"""FATHOM Streamlit Demo — DEM-03.

2-panel layout:
  Left:   Recursion tree visualization (D3 or plotly sunburst)
  Right:  W&B training curves embed (live iframe)

Run: streamlit run space_demo/app.py
"""
from __future__ import annotations

import json
import os

import streamlit as st

st.set_page_config(
    page_title="FATHOM — RL-Trained Recursive Language Model",
    page_icon="🧠",
    layout="wide",
)

st.markdown(
    f"**Live env:** [Pratham-math/fathom-env]({os.environ.get('FATHOM_SPACE_URL','https://Pratham-math-fathom-env.hf.space')}) "
    f"· **Trained model:** [Pratham-math/fathom-1.5b-grpo](https://huggingface.co/Pratham-math/fathom-1.5b-grpo) "
    f"· **W&B:** [run sy1tqun0](https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/sy1tqun0)"
)

# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------
st.markdown("## 🧠 FATHOM Demo")
st.markdown(
    "_First RL-Trained Recursive Language Model — Meta × PyTorch × HF Hackathon_"
)
st.divider()

# ---------------------------------------------------------------------------
# Reward composition badge (DEM-03 C.4: visible-to-judges reward overview)
# ---------------------------------------------------------------------------
with st.expander("Reward composition (4 components, deterministic verifier)", expanded=True):
    cb1, cb2 = st.columns([1, 2], gap="medium")
    with cb1:
        st.markdown("**Format gate** (multiplier)")
        st.success("`<answer>...</answer>` required \u2014 if missing, soft bonus + cap applies")
        st.caption("Source: `rewards/compose.py` \u2014 audited in REWARD_AUDIT.md")
    with cb2:
        try:
            import plotly.graph_objects as go  # type: ignore
            labels = ["correctness", "token_budget", "recursion_efficiency"]
            weights = [0.70, 0.15, 0.15]
            colors = ["#2ca02c", "#1f77b4", "#ff7f0e"]
            fig0 = go.Figure(go.Pie(
                labels=labels, values=weights, marker=dict(colors=colors),
                hole=0.4, textinfo="label+percent",
            ))
            fig0.update_layout(margin=dict(l=10, r=10, t=10, b=10), height=200, showlegend=False)
            st.plotly_chart(fig0, use_container_width=True)
        except ImportError:
            st.metric("correctness", 0.70)
            st.metric("token_budget", 0.15)
            st.metric("recursion_efficiency", 0.15)

st.divider()

# ---------------------------------------------------------------------------
# Sidebar: controls
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("Controls")
    env_url = st.text_input(
        "Env server URL",
        value=os.environ.get("FATHOM_SPACE_URL", "https://Pratham-math-fathom-env.hf.space"),
        key="env_url",
    )
    st.divider()
    st.caption("Source: [github.com/Pratham-math/fathom](https://github.com/Pratham-math/fathom)")

@st.cache_data(ttl=3600)
def fetch_trace(url: str) -> dict:
    import httpx
    try:
        # Provide a quick dummy trace by interacting with the env.
        # Note: the real recursive model is not loaded in Streamlit,
        # so this just grabs the reset observation and submits a dummy answer.
        r = httpx.post(f"{url}/reset", json={"seed": 42}, timeout=10)
        r.raise_for_status()
        obs = r.json()
        q = obs.get("question", "Question")
        
        s = httpx.post(f"{url}/step", json={"action_type": "answer", "answer": "simulated"}, timeout=10)
        s.raise_for_status()
        
        return {
            "name": f"Live init: {q[:30]}...",
            "children": [
                {"name": "REPL: simulated step"},
                {"name": "→ <answer>simulated</answer>"}
            ]
        }
    except Exception as e:
        return {
            "name": "[Example Trace] root: 200K doc",
            "children": [
                {
                    "name": "llm(chunk_0-50K)",
                    "children": [{"name": "REPL: grep → 'azure'"}],
                },
                {"name": "REPL: count_tokens → 200K"},
                {"name": "→ <answer>azure</answer>"},
            ],
        }

# ---------------------------------------------------------------------------
# 2 columns
# ---------------------------------------------------------------------------
col_tree, col_wb = st.columns([1.5, 1.5], gap="medium")

# ── Column 1: Recursion tree ────────────────────────────────────────────────
with col_tree:
    st.subheader("Recursion Tree")
    st.caption("Sample episode trace")

    sample_tree = fetch_trace(env_url)
    tree_json = json.dumps(sample_tree)
    d3_html = f"""
    <html>
    <head>
    <script src="https://cdn.jsdelivr.net/npm/d3@7"></script>
    <style>
      body {{ font-family: monospace; font-size: 12px; }}
      .node circle {{ fill: #6366f1; stroke: #312e81; stroke-width: 1.5px; }}
      .node text {{ fill: #1e1b4b; }}
      .link {{ fill: none; stroke: #a5b4fc; stroke-width: 1.5px; }}
    </style>
    </head>
    <body>
    <div id="tree"></div>
    <script>
    const data = {tree_json};
    const width = 400, height = 240;
    const svg = d3.select("#tree").append("svg").attr("width", width).attr("height", height);
    const g = svg.append("g").attr("transform", "translate(40,20)");
    const tree = d3.tree().size([height-40, width-120]);
    const root = d3.hierarchy(data);
    tree(root);
    g.selectAll(".link").data(root.links()).enter().append("path")
      .attr("class","link")
      .attr("d", d3.linkHorizontal().x(d=>d.y).y(d=>d.x));
    const node = g.selectAll(".node").data(root.descendants()).enter()
      .append("g").attr("class","node")
      .attr("transform", d=>`translate(${{d.y}},${{d.x}})`);
    node.append("circle").attr("r", 5);
    node.append("text").attr("dy","0.35em").attr("x", d=>d.children?-8:8)
      .attr("text-anchor", d=>d.children?"end":"start")
      .text(d=>d.data.name.slice(0,35));
    </script>
    </body></html>
    """
    st.components.v1.html(d3_html, height=280)

# ── Column 3: W&B training curves ────────────────────────────────────────────
with col_wb:
    st.subheader("Training Curves")
    wb_url = os.environ.get("WANDB_RUN_URL", "https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/sy1tqun0")
    if wb_url:
        st.components.v1.iframe(wb_url, height=260, scrolling=True)
    else:
        st.info("Set `WANDB_RUN_URL` env var to embed live training curves.")

st.divider()
st.caption("FATHOM Demo Space")