"""Graph-PrefLexOR: Graph-Native Reinforcement Learning for Scientific Hypothesis Generation.

A chat demo that loads the lamm-mit/Graph-Preflexor-8b_12292025 model, streams its
structured reasoning output, parses the <graph_json> block, and renders the relational
graph as an interactive Mermaid diagram alongside the conversation.
"""

import json
import re
import spaces  # MUST come before torch / transformers
import torch
from threading import Thread
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
    GenerationConfig,
)
import gradio as gr

MODEL_ID = "lamm-mit/Graph-Preflexor-8b_12292025"

# ---------------------------------------------------------------------------
# Model loading (module scope, eager .to("cuda") per ZeroGPU rules)
# ---------------------------------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()

# ---------------------------------------------------------------------------
# Graph JSON → Mermaid conversion
# ---------------------------------------------------------------------------
GRAPH_JSON_RE = re.compile(
    r"<graph_json>\s*(.*?)\s*</graph_json>", re.DOTALL
)


def extract_graph_json(text: str):
    """Extract and parse the first <graph_json> block from model output."""
    m = GRAPH_JSON_RE.search(text)
    if not m:
        return None
    raw = m.group(1).strip()
    try:
        return json.loads(raw)
    except Exception:
        i1 = raw.find("{")
        i2 = raw.rfind("}")
        if i1 != -1 and i2 > i1:
            try:
                return json.loads(raw[i1 : i2 + 1])
            except Exception:
                return None
    return None


def _safe_id(s: str) -> str:
    """Sanitize a node id for use as a Mermaid node identifier."""
    out = re.sub(r"[^A-Za-z0-9_]", "_", str(s))
    if out and out[0].isdigit():
        out = "_" + out
    return out or "_"


def graph_to_mermaid(graph_obj: dict) -> str:
    """Convert a {nodes, edges} graph dict to a Mermaid flowchart definition."""
    if not graph_obj:
        return ""
    nodes = graph_obj.get("nodes", []) or []
    edges = graph_obj.get("edges", []) or []

    lines = ["flowchart LR"]
    seen = set()

    for n in nodes:
        if not isinstance(n, dict):
            continue
        nid = n.get("id")
        if not nid:
            continue
        safe = _safe_id(nid)
        if safe in seen:
            continue
        seen.add(safe)
        ntype = n.get("type", "")
        label = nid
        if ntype:
            label = f"{nid} ({ntype})"
        lines.append(f'  {safe}["{label}"]')

    for e in edges:
        if not isinstance(e, dict):
            continue
        src = e.get("source")
        tgt = e.get("target")
        if not src or not tgt:
            continue
        rel = e.get("relation", "")
        s_safe = _safe_id(src)
        t_safe = _safe_id(tgt)
        if s_safe not in seen:
            seen.add(s_safe)
            lines.append(f'  {s_safe}["{src}"]')
        if t_safe not in seen:
            seen.add(t_safe)
            lines.append(f'  {t_safe}["{tgt}"]')
        if rel:
            lines.append(f"  {s_safe} -->|{rel}| {t_safe}")
        else:
            lines.append(f"  {s_safe} --> {t_safe}")

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def _estimate_duration(message, history, max_new_tokens, *args, **kwargs):
    """Callable duration estimator: base + proportional to max_new_tokens."""
    return min(300, 30 + int(max_new_tokens * 0.004))


def _content_to_text(content) -> str:
    """Normalize a Gradio chat message 'content' to a plain string.

    Gradio 6 chatbots deliver message content as a list of parts
    (e.g. [{"type": "text", "text": "..."}]) rather than a bare string.
    The model's chat template only keeps content when it is a string, so
    list-shaped content must be flattened before it reaches the model or
    the user's text is silently dropped.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for part in content:
            if isinstance(part, dict):
                if part.get("type", "text") == "text":
                    parts.append(part.get("text", ""))
            elif isinstance(part, str):
                parts.append(part)
        return "".join(parts)
    if content is None:
        return ""
    return str(content)


@spaces.GPU(duration=_estimate_duration)
def generate(
    message: str,
    history: list,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
):
    """Stream a graph-native reasoning response from the model.

    Args:
        message: The user's scientific question or prompt.
        history: Chat history as a list of message dicts with 'role' and 'content'.
        max_new_tokens: Maximum tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling threshold.
    """
    # Build the messages list from history (list of {role, content} dicts).
    # Gradio may hand us content as a list of parts, so flatten to plain text.
    messages = []
    for msg in history:
        if isinstance(msg, dict) and msg.get("content"):
            text = _content_to_text(msg["content"])
            if text:
                messages.append({"role": msg["role"], "content": text})
    messages.append({"role": "user", "content": _content_to_text(message)})

    # Apply chat template with thinking enabled
    try:
        prompt_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=True,
        )
    except TypeError:
        prompt_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

    inputs = tokenizer(prompt_text, return_tensors="pt").to("cuda")

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=False,
    )

    gen_config = GenerationConfig(
        max_new_tokens=max_new_tokens,
        do_sample=temperature > 0,
        temperature=max(temperature, 0.01),
        top_p=top_p,
    )

    generation_kwargs = dict(
        **inputs,
        generation_config=gen_config,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    full_text = ""
    for chunk in streamer:
        full_text += chunk
        yield full_text, ""

    thread.join()

    graph_obj = extract_graph_json(full_text)
    mermaid_code = graph_to_mermaid(graph_obj) if graph_obj else ""

    # The graph panel is a custom gr.HTML component whose value is the raw
    # Mermaid graph definition; the component renders it into an actual diagram.
    yield full_text, mermaid_code


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
#graph-panel { min-height: 300px; }
"""

EXAMPLES = [
    "What are the key mechanical properties of spider silk and how do they arise from its molecular structure?",
    "Explain the relationship between hierarchical structures and material toughness in biological materials.",
    "How do proteins fold and why is this important for their function?",
    "Give me a short introduction to materiomics.",
    "Propose a novel hypothesis for self-healing biopolymer composites.",
]

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🧠 Graph-PrefLexOR: Graph-Native Scientific Reasoning
        Ask a scientific question and the model will reason through it using
        structured graph-native thinking — brainstorming, building a knowledge
        graph, extracting patterns, and synthesizing a final answer.
        The extracted relational graph is visualized as an interactive Mermaid
        diagram.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=520)
            with gr.Row():
                msg_input = gr.Textbox(
                    show_label=False,
                    placeholder="Ask a scientific question…",
                    container=False,
                    scale=4,
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)
            with gr.Accordion("Advanced settings", open=False):
                with gr.Row():
                    max_tokens = gr.Slider(
                        512, 16384, value=8192, step=512,
                        label="Max new tokens",
                    )
                    temperature = gr.Slider(
                        0.0, 1.5, value=0.2, step=0.05,
                        label="Temperature",
                    )
                    top_p = gr.Slider(
                        0.1, 1.0, value=0.95, step=0.05,
                        label="Top-p",
                    )
            with gr.Row():
                stop_btn = gr.Button("Stop", variant="stop")
                clear_btn = gr.Button("Clear")

        with gr.Column(scale=2):
            gr.Markdown("### 📊 Relational Graph")
            # Custom gr.HTML component (see Gradio "Custom HTML Components" guide):
            # mermaid.js is loaded via `head`, and the component's `value` holds the
            # raw Mermaid graph definition. `html_template` re-renders on every value
            # update; the ${...} expression calls mermaid.render() with a unique id
            # (timestamp) each time so the diagram is re-rendered as an actual SVG
            # diagram whenever the underlying definition changes.
            graph_output = gr.HTML(
                value="",
                elem_id="graph-panel",
                head="""<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>""",
                html_template="""
                <div style="display:flex;justify-content:center;padding:8px;overflow:auto;min-height:120px;">
                  ${
                    (() => {
                      const def = (value || '').toString().trim();
                      if (!def) {
                        return "<p style='color:#888;text-align:center;padding:40px;'>The extracted knowledge graph will appear here after generation.</p>";
                      }
                      // Unique id per render forces mermaid to re-render fresh.
                      const uid = 'mermaid-svg-' + Date.now() + '-' + Math.floor(Math.random() * 1e6);
                      const holderId = 'mermaid-holder-' + uid;
                      // Render asynchronously and inject the resulting SVG once ready.
                      setTimeout(() => {
                        try {
                          if (typeof mermaid === 'undefined') return;
                          if (!window.__mermaidInit) {
                            mermaid.initialize({ startOnLoad: false, theme: 'default', securityLevel: 'loose' });
                            window.__mermaidInit = true;
                          }
                          mermaid.render(uid, def).then(({ svg }) => {
                            const holder = document.getElementById(holderId);
                            if (holder) holder.innerHTML = svg;
                          }).catch((e) => {
                            const holder = document.getElementById(holderId);
                            if (holder) holder.innerHTML = "<pre style='color:#c00;white-space:pre-wrap;'>Mermaid render error: " + (e && e.message ? e.message : e) + "</pre>";
                          });
                        } catch (e) { console.warn('Mermaid render:', e); }
                      }, 0);
                      return "<div id='" + holderId + "' style='width:100%;'>Rendering diagram…</div>";
                    })()
                  }
                </div>
                """,
            )

    def run_example(message):
        """Wrapper for gr.Examples: run generate with defaults, return chat history + graph."""
        for text, graph_html in generate(message, [], 8192, 0.2, 0.95):
            pass
        history = [
            {"role": "user", "content": message},
            {"role": "assistant", "content": text},
        ]
        return history, graph_html

    gr.Examples(
        examples=EXAMPLES,
        inputs=msg_input,
        outputs=[chatbot, graph_output],
        fn=run_example,
        cache_examples=False,
        run_on_click=True,
    )

    # Wire events — use messages format for Gradio 6
    def user_submit(message, history):
        """Add user message to chat history and clear input."""
        if not message.strip():
            return gr.skip(), history
        new_history = history + [{"role": "user", "content": message}]
        return "", new_history

    def bot_respond(history, max_tok, temp, tp):
        """Run generation and stream the response into the chatbot."""
        if not history or history[-1].get("role") != "user":
            yield history, gr.skip()
            return
        message = _content_to_text(history[-1]["content"])
        history_for_gen = history[:-1]
        for text, graph_html in generate(message, history_for_gen, max_tok, temp, tp):
            updated = history + [{"role": "assistant", "content": text}]
            yield updated, graph_html

    submit_events = [send_btn.click, msg_input.submit]
    cancel_targets = []
    for evt in submit_events:
        click_event = evt(
            user_submit,
            [msg_input, chatbot],
            [msg_input, chatbot],
        ).then(
            bot_respond,
            [chatbot, max_tokens, temperature, top_p],
            [chatbot, graph_output],
        )
        cancel_targets.append(click_event)

    clear_btn.click(
        lambda: ([], ""),
        None,
        [chatbot, graph_output],
    )
    stop_btn.click(None, None, None, cancels=cancel_targets)


demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)