"""
Dialectical Transition Operator — interactive demo.

Given a (question, answer, 3 critiques), the model returns a (revised answer, 3 fresh
critiques). Iterating it deepens the answer.

Model = Qwen3-8B + two stacked LoRA adapters (a frozen "SFT-voice" + a trainable "RL"
adapter trained with GRPO / set-VPO under a readability-and-addressability-aware judge).
"""
import spaces          # MUST be imported before anything initializes CUDA
import os
import re
import random

from threading import Thread

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from peft import PeftModel
from huggingface_hub import snapshot_download
import gradio as gr

from questions import QUESTIONS

MODEL_ID = "Qwen/Qwen3-8B"
REPO = "andreiski/dialectical-transition-cot"
SFT_SUB = "mission-diversity-sft/adapter-final"   # frozen SFT-voice adapter
RL_SUB = "mission-diversity-rl/adapter"            # trained RL adapter
HF_TOKEN = os.environ.get("HF_TOKEN")

# Pre-download weights at startup (pure file fetch, no CUDA) so the lazy GPU-load is quick.
print("warming weight cache ...")
snapshot_download(MODEL_ID, token=HF_TOKEN)
snapshot_download(REPO, token=HF_TOKEN, allow_patterns=[SFT_SUB + "/*", RL_SUB + "/*"])

TOK = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
if TOK.pad_token is None:
    TOK.pad_token = TOK.eos_token
MODEL = None

def _ensure_model():
    """Load base + stacked adapters INSIDE the GPU context (cached after first call).
    ZeroGPU has no GPU at import time, so model/CUDA setup must happen here."""
    global MODEL
    if MODEL is not None:
        return MODEL
    base = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, torch_dtype=torch.bfloat16, token=HF_TOKEN).to("cuda")
    m = PeftModel.from_pretrained(base, REPO, subfolder=SFT_SUB, adapter_name="sftvoice", token=HF_TOKEN)
    m.load_adapter(REPO, subfolder=RL_SUB, adapter_name="rl", token=HF_TOKEN)
    m.base_model.set_adapter(["sftvoice", "rl"])   # BOTH active/stacked — matches training
    m.eval()
    MODEL = m
    return MODEL

# ------------------------------------------------------------------ format (verbatim)
def render_prompt(question, answer, critiques):
    lines = [f"Question: {question}", f"Answer: {answer}", "Critiques of answer:"]
    for i, c in enumerate(critiques):
        lines.append(f"[c{i + 1}] {c}")
    return "\n".join(lines)

_CLEAN_RE = re.compile(r"<\|[^>]*\|>")
_REVISED_RE = re.compile(r"(?msi)^[ \t]*Revised answer[ \t]*:[ \t]*(.*?)"
                         r"(?=^[ \t]*Critiques of revised answer[ \t]*:|\Z)")
_CRITHDR_RE = re.compile(r"(?mi)^[ \t]*Critiques of revised answer[ \t]*:[ \t]*")
_CRIT_RE = re.compile(r"(?ms)^[ \t]*\[c(\d+)\][ \t]*(.*?)(?=^[ \t]*\[c\d+\]|\Z)")

def _clean(s):
    return _CLEAN_RE.sub("", s or "").strip()

_REVHDR_RE = re.compile(r"(?i)revised answer\s*:\s*")
_CRITHDR2_RE = re.compile(r"(?i)critiques of revised answer\s*:")
_CRITN_RE = re.compile(r"\[c([123])\]\s*(.*?)(?=\[c[123]\]|\Z)", re.S)

def partial_parse(text):
    """Best-effort incremental parse for streaming: returns (thinking, revised, [c1,c2,c3])."""
    op = text.find("<think>")
    start = op + len("<think>") if op != -1 else 0
    cl = text.find("</think>")
    if cl != -1:
        thinking = text[start:cl]
        body = text[cl + len("</think>"):]
    else:
        thinking = text[start:]
        body = ""
    revised = ""
    mr = _REVHDR_RE.search(body)
    if mr:
        after = body[mr.end():]
        mc = _CRITHDR2_RE.search(after)
        revised = after[:mc.start()] if mc else after
    crits = ["", "", ""]
    mh = _CRITHDR2_RE.search(body)
    if mh:
        cb = body[mh.end():]
        for m in _CRITN_RE.finditer(cb):
            crits[int(m.group(1)) - 1] = m.group(2)
    return _clean(thinking), _clean(revised), [_clean(c) for c in crits]

# ------------------------------------------------------------------------- seeds
GENERIC_ANSWERS = [
    "Who knows?", "It depends.", "That's hard to say.", "There isn't really a clear answer.",
    "It's complicated.", "Honestly, it could go either way.", "I'm not sure.",
]
GENERIC_CRITIQUES = [
    "This doesn't actually answer the question.",
    "It gives the reader nothing useful to act on.",
    "It's vague and noncommittal.",
    "It dodges the question entirely.",
    "It offers no reasoning or explanation.",
    "It's too generic to be helpful.",
    "It never engages with what was actually asked.",
    "It leaves the reader exactly where they started.",
]

def random_seed():
    a = random.choice(GENERIC_ANSWERS)
    cs = random.sample(GENERIC_CRITIQUES, 3)
    return a, cs[0], cs[1], cs[2]

def random_question():
    return random.choice(QUESTIONS)

# ---------------------------------------------------------------------- generation
def _run_stream(question, ans, c1, c2, c3):
    model = _ensure_model()
    user = render_prompt(question, ans, [c1, c2, c3])
    prompt = TOK.apply_chat_template([{"role": "user", "content": user}],
                                     tokenize=False, add_generation_prompt=True)
    inputs = TOK(prompt, return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(TOK, skip_prompt=True, skip_special_tokens=False)
    kw = dict(**inputs, max_new_tokens=1280, do_sample=True, temperature=1.0,
              top_p=1.0, top_k=0, repetition_penalty=1.0,
              pad_token_id=TOK.eos_token_id, streamer=streamer)
    Thread(target=model.generate, kwargs=kw).start()
    full = ""
    for chunk in streamer:
        full += chunk
        yield full

# --------------------------------------------------------------------------- UI html
_IDLE = """
<div style="display:flex;align-items:center;justify-content:center;height:100%;min-height:430px;">
  <div style="writing-mode:vertical-rl;transform:rotate(180deg);
              background:#334155;color:#f8fafc;border-radius:6px;
              display:flex;align-items:center;justify-content:center;
              padding:10px 16px;min-height:430px;
              font-weight:600;font-size:.92em;letter-spacing:3px;text-transform:uppercase;">
    transition&nbsp;operator
  </div>
</div>"""

_BUSY = """
<style>@keyframes spin{to{transform:rotate(360deg)}}</style>
<div style="display:flex;align-items:center;justify-content:center;height:100%;min-height:430px;">
  <div style="background:#334155;border-radius:6px;min-height:430px;padding:14px 16px;
              display:flex;flex-direction:column;align-items:center;justify-content:center;gap:16px;">
    <div style="width:30px;height:30px;border:4px solid #64748b;border-top-color:#f8fafc;border-radius:50%;animation:spin .8s linear infinite;"></div>
    <div style="writing-mode:vertical-rl;transform:rotate(180deg);color:#cbd5e1;font-size:.82em;letter-spacing:2px;text-transform:uppercase;">thinking</div>
  </div>
</div>"""

@spaces.GPU(duration=60)
def ui_generate(question, ans, c1, c2, c3):
    if not (question or "").strip():
        yield _IDLE, "", "Enter a question first.", "", "", ""
        return
    yield _BUSY, "", "", "", "", ""
    cot, rev, cs = "", "", ["", "", ""]
    for full in _run_stream(question, ans or "", c1 or "", c2 or "", c3 or ""):
        cot, rev, cs = partial_parse(full)
        box = _IDLE if rev else _BUSY   # spinner during <think>, then stream answer + critiques
        yield box, cot, rev, cs[0], cs[1], cs[2]
    msg = rev if rev else "Generation was malformed — press Generate to try again."
    yield _IDLE, cot, msg, cs[0], cs[1], cs[2]

# ------------------------------------------------------------------------------- app
_THEME = gr.themes.Default(font=[gr.themes.GoogleFont("Inter")])
with gr.Blocks(title="Dialectical Transition Operator", theme=_THEME) as demo:
    gr.Markdown("# Dialectical Transition Operator")
    gr.Markdown(
        "Give the model a **question**, a starting **answer**, and **three critiques** of that answer. "
        "It returns a **revised answer** and **three fresh critiques**. "
        "**Copy the output back to the input and generate again** to iteratively deepen the answer.")

    with gr.Row():
        question = gr.Textbox(label="Question (X)", lines=2, placeholder="Ask anything…", scale=8)
        randq_btn = gr.Button("Random question", scale=1, min_width=130)

    with gr.Row(equal_height=True):
        with gr.Column(scale=6):
            gr.Markdown("#### Input")
            in_ans = gr.Textbox(label="Answer", lines=5)
            gr.Markdown("Critiques")
            in_c1 = gr.Textbox(show_label=False, lines=2)
            in_c2 = gr.Textbox(show_label=False, lines=2)
            in_c3 = gr.Textbox(show_label=False, lines=2)
        with gr.Column(scale=1, min_width=66):
            model_box = gr.HTML(_IDLE)
        with gr.Column(scale=6):
            gr.Markdown("#### Output")
            out_ans = gr.Textbox(label="Revised answer", lines=5)
            gr.Markdown("Critiques")
            out_c1 = gr.Textbox(show_label=False, lines=2)
            out_c2 = gr.Textbox(show_label=False, lines=2)
            out_c3 = gr.Textbox(show_label=False, lines=2)

    with gr.Row():
        seed_btn = gr.Button("Random input")
        gen_btn = gr.Button("Generate", variant="primary")
        copy_btn = gr.Button("Copy output → input")

    with gr.Accordion("Chain of thought (the model's live reasoning)", open=True):
        cot_box = gr.Textbox(show_label=False, lines=8)

    gr.Examples(
        examples=[
            "How does Hurricane Florence lose strength?",
            "Explain Kant like I'm five but a philosophy professor would approve.",
            "I signed a lease last week and found someone already living in the basement — what can I do?",
            "Write a Python script that alerts me if NVIDIA stock moves 10% within 5 days.",
        ],
        inputs=question)

    randq_btn.click(random_question, outputs=question)
    seed_btn.click(random_seed, outputs=[in_ans, in_c1, in_c2, in_c3])
    gen_btn.click(ui_generate, inputs=[question, in_ans, in_c1, in_c2, in_c3],
                  outputs=[model_box, cot_box, out_ans, out_c1, out_c2, out_c3])
    copy_btn.click(lambda a, b, c, d: (a, b, c, d, "", "", "", ""),
                   inputs=[out_ans, out_c1, out_c2, out_c3],
                   outputs=[in_ans, in_c1, in_c2, in_c3, out_ans, out_c1, out_c2, out_c3])
    demo.load(random_seed, outputs=[in_ans, in_c1, in_c2, in_c3])

if __name__ == "__main__":
    demo.queue(max_size=20).launch()