import gradio as gr
import time
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM, pipeline

# ── Device ────────────────────────────────────────────────────────────────────
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🖥️ Device: {DEVICE}")

# ── Reasoner — loads locally, no API calls ────────────────────────────────────
REASONER_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
print(f"🔄 Loading reasoner: {REASONER_MODEL}")
reasoner_tokenizer = AutoTokenizer.from_pretrained(REASONER_MODEL, trust_remote_code=True)
reasoner_model = AutoModelForCausalLM.from_pretrained(
    REASONER_MODEL,
    trust_remote_code=True,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    device_map="auto",
)
reasoner_pipe = pipeline(
    "text-generation",
    model=reasoner_model,
    tokenizer=reasoner_tokenizer,
    max_new_tokens=600,
    temperature=0.3,
    do_sample=True,
)
print("✅ Reasoner loaded!")

# ── Step Probe — fine-tuned on PRM800K ───────────────────────────────────────
PROBE_MODEL = "realArceus/twt-probe"
print(f"🔄 Loading probe: {PROBE_MODEL}")
probe_tokenizer = AutoTokenizer.from_pretrained(PROBE_MODEL, trust_remote_code=True)
probe_model = AutoModelForSequenceClassification.from_pretrained(
    PROBE_MODEL,
    trust_remote_code=True,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
).to(DEVICE)
probe_model.eval()
print("✅ Probe loaded!")

SYSTEM_PROMPT = """You are a careful step-by-step reasoner. When given a problem, solve it by thinking through exactly numbered steps.
Format EVERY step as:
Step 1: <your reasoning>
Step 2: <your reasoning>
...
Final Answer: <answer>
Be deliberate. Show your full working. Each step should be one clear thought."""

# ── Step Probe inference ──────────────────────────────────────────────────────
def probe_step(problem: str, steps_so_far: list) -> tuple:
    context = " ".join(steps_so_far[:-1]) if len(steps_so_far) > 1 else ""
    current = steps_so_far[-1] if steps_so_far else ""
    input_text = (
        f"[PROBLEM] {problem.strip()} "
        f"[STEPS SO FAR] {context} "
        f"[CURRENT STEP] {current}"
    )
    inputs = probe_tokenizer(
        input_text, return_tensors="pt",
        truncation=True, max_length=512,
    ).to(DEVICE)
    with torch.no_grad():
        logits = probe_model(**inputs).logits
        probs  = torch.softmax(logits, dim=-1)[0]
    ok_conf   = probs[1].item()
    fail_conf = probs[0].item()
    label = "OK" if ok_conf >= 0.5 else "FAIL"
    conf  = round(ok_conf if label == "OK" else fail_conf, 2)
    return label, conf

# ── Core pipeline ─────────────────────────────────────────────────────────────
@spaces.GPU
def run_twt(problem: str):
    if not problem.strip():
        yield ("", "<div class='msg warn'>⚠️ Enter a problem to analyze.</div>", "<div class='msg empty'>Waiting...</div>")
        return

    cot_html = ""

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": problem.strip()}
    ]
    prompt = reasoner_tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    try:
        result = reasoner_pipe(prompt)
        full_response = result[0]["generated_text"][len(prompt):]
    except Exception as e:
        yield ("", f"<div class='msg error'>❌ Reasoner error: {str(e)}</div>", "")
        return

    lines        = full_response.strip().split("\n")
    parsed_steps = [l.strip() for l in lines if l.strip()]
    displayed_steps = []
    first_fail_seen = False

    for i, step in enumerate(parsed_steps):
        displayed_steps.append(step)
        is_final = step.lower().startswith("final answer")

        if not is_final:
            label, conf = probe_step(problem, displayed_steps)
        else:
            label, conf = "FINAL", 1.0

        if is_final:
            icon, badge_class, badge_text = "🏁", "badge-final", "ANSWER"
        elif label == "OK":
            icon, badge_class, badge_text = "🟢", "badge-ok", f"OK · {int(conf*100)}%"
        else:
            icon, badge_class, badge_text = "🔴", "badge-fail", f"FAULT · {int(conf*100)}%"
            if not first_fail_seen:
                first_fail_seen = True

        step_card = f"""
        <div class='step-card {"step-fault" if label == "FAIL" else ""}'>
            <span class='step-icon'>{icon}</span>
            <span class='step-text'>{step}</span>
            <span class='badge {badge_class}'>{badge_text}</span>
        </div>"""
        cot_html += step_card

        yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen))
        time.sleep(0.1)

    yield (full_response, cot_html, build_trace(problem, displayed_steps, first_fail_seen, done=True))


def build_trace(problem, displayed, first_fail, done=False):
    faults, rows = 0, ""
    for i, s in enumerate(displayed):
        is_final = s.lower().startswith("final answer")
        if not is_final:
            lbl, conf = probe_step(problem, displayed[:i+1])
            if lbl == "FAIL": faults += 1
            dot   = f"<span class='dot dot-{'ok' if lbl == 'OK' else 'fail'}'></span>"
            rows += f"<div class='trace-row'>{dot} Step {i+1} <span class='trace-conf'>&mdash; {int(conf*100)}%</span></div>"
        else:
            rows += "<div class='trace-row'><span class='dot dot-final'></span> Final Answer</div>"

    health    = max(0, 100 - faults * 20)
    bar_color = "#10b981" if health > 60 else "#e11d48"
    summary   = f"""
    <div class='trace-summary'>
        <div class='trace-label'>Reasoning Health</div>
        <div class='trace-bar-bg'><div class='trace-bar' style='width:{health}%;background:{bar_color}'></div></div>
        <div class='trace-pct'>{health}%</div>
    </div>""" if done else ""
    return f"{rows}{summary}"


# ── CSS (Compact Enterprise Corporate Navy Theme) ──────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500;600&display=swap');

*, *::before, *::after { box-sizing: border-box; }
body, .gradio-container, .gradio-container * {
    font-family: 'Inter', -apple-system, sans-serif !important;
    color: #1e293b !important;
}
.gradio-container {
    background-color: #f4f6f8 !important;
    max-width: 1280px !important;
    padding-top: 1rem !important;
}
.twt-header {
    text-align: center;
    padding: 0 1rem 1rem;
    margin-bottom: 1.5rem;
    border-bottom: 1px solid #cbd5e1;
}
.twt-title {
    font-size: 2rem;
    font-weight: 800;
    letter-spacing: -0.03em;
    font-family: 'JetBrains Mono', monospace !important;
    margin-bottom: 0.25rem;
    color: #0f172a !important;
}
.twt-title span { color: #0033a0 !important; }
.twt-sub {
    font-size: 0.85rem;
    color: #475569 !important;
    font-weight: 500;
    letter-spacing: 0.02em;
}
.input-label, .panel-label {
    font-size: 0.7rem;
    font-weight: 700;
    color: #334155 !important;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    margin-bottom: 0.5rem;
    display: block;
}
.panel-label {
    padding-bottom: 0.4rem;
    border-bottom: 2px solid #cbd5e1;
}
textarea {
    background: #ffffff !important;
    border: 1px solid #94a3b8 !important;
    border-radius: 6px !important;
    color: #0f172a !important;
    font-size: 0.85rem !important;
    line-height: 1.5 !important;
    padding: 0.75rem !important;
    box-shadow: inset 0 1px 2px rgba(15, 23, 42, 0.05) !important;
    transition: all 0.2s ease !important;
}
textarea:focus {
    border-color: #0033a0 !important;
    box-shadow: 0 0 0 3px rgba(0, 51, 160, 0.15) !important;
    outline: none !important;
}
textarea::placeholder { color: #94a3b8 !important; }
button.primary {
    background: #0033a0 !important;
    color: #ffffff !important;
    font-weight: 600 !important;
    border: none !important;
    border-radius: 6px !important;
    padding: 0.6rem 1.25rem !important;
    font-size: 0.85rem !important;
    box-shadow: 0 2px 4px -1px rgba(0, 51, 160, 0.2) !important;
    transition: all 0.15s ease-in-out !important;
    cursor: pointer !important;
    width: 100% !important;
}
button.primary:hover {
    background: #002266 !important;
    transform: translateY(-1px) !important;
    box-shadow: 0 4px 6px -2px rgba(0, 51, 160, 0.3) !important;
}
button.primary:active {
    transform: translateY(0) !important;
    box-shadow: 0 1px 2px rgba(0, 51, 160, 0.2) !important;
}
.step-card {
    display: flex;
    align-items: flex-start;
    gap: 0.75rem;
    padding: 0.75rem 1rem;
    margin-bottom: 0.5rem;
    background: #ffffff;
    border: 1px solid #cbd5e1;
    border-radius: 6px;
    line-height: 1.5;
    animation: fadeUp 0.2s ease-out both;
    box-shadow: 0 1px 2px rgba(15, 23, 42, 0.03);
    transition: all 0.15s ease;
}
.step-card:hover {
    box-shadow: 0 2px 6px rgba(15, 23, 42, 0.06);
    border-color: #94a3b8;
}
.step-card.step-fault {
    border-color: #fca5a5;
    background: #fff1f2;
}
.step-card.step-fault:hover { border-color: #f87171; }
@keyframes fadeUp {
    from { opacity: 0; transform: translateY(6px); }
    to   { opacity: 1; transform: translateY(0); }
}
.step-icon { font-size: 1rem; margin-top: 1px; }
.step-text { flex: 1; color: #1e293b !important; font-size: 0.85rem; }
.badge {
    flex-shrink: 0;
    font-size: 0.6rem;
    font-weight: 700;
    font-family: 'JetBrains Mono', monospace !important;
    letter-spacing: 0.02em;
    padding: 0.15rem 0.4rem;
    border-radius: 4px;
    margin-top: 2px;
    text-transform: uppercase;
}
.badge-ok    { background: #f0fdf4; color: #15803d !important; border: 1px solid #86efac; }
.badge-fail  { background: #fff1f2; color: #be123c !important; border: 1px solid #fda4af; }
.badge-final { background: #f0f4ff; color: #0033a0 !important; border: 1px solid #bfdbfe; }
.trace-row {
    display: flex;
    align-items: center;
    gap: 0.5rem;
    font-size: 0.75rem;
    color: #334155 !important;
    font-family: 'JetBrains Mono', monospace !important;
    padding: 0.35rem 0;
    border-bottom: 1px solid #e2e8f0;
}
.trace-conf { color: #64748b !important; }
.dot { width: 6px; height: 6px; border-radius: 50%; flex-shrink: 0; }
.dot-ok    { background: #10b981; box-shadow: 0 0 0 2px #d1fae5; }
.dot-fail  { background: #e11d48; box-shadow: 0 0 0 2px #ffe4e6; }
.dot-final { background: #0033a0; box-shadow: 0 0 0 2px #dbeafe; }
.trace-summary {
    margin-top: 1rem;
    padding-top: 0.75rem;
    border-top: 1px solid #cbd5e1;
}
.trace-label {
    font-size: 0.65rem;
    font-weight: 700;
    color: #475569 !important;
    text-transform: uppercase;
    letter-spacing: 0.05em;
    margin-bottom: 0.4rem;
}
.trace-bar-bg {
    background: #cbd5e1;
    border-radius: 999px;
    height: 4px;
    overflow: hidden;
}
.trace-bar {
    height: 100%;
    border-radius: 999px;
    transition: width 0.8s cubic-bezier(0.16, 1, 0.3, 1);
}
.trace-pct {
    font-size: 0.75rem;
    color: #0f172a !important;
    margin-top: 0.4rem;
    font-family: 'JetBrains Mono', monospace !important;
    font-weight: 600;
}
.msg { padding: 0.75rem; border-radius: 6px; font-size: 0.8rem; font-weight: 500; }
.msg.warn  { background: #fefce8; color: #a16207 !important; border: 1px solid #fde047; }
.msg.error { background: #fff1f2; color: #be123c !important; border: 1px solid #fca5a5; }
.msg.empty {
    color: #64748b !important;
    font-size: 0.8rem;
    font-family: 'JetBrains Mono', monospace !important;
    padding: 1.5rem 1rem;
    text-align: center;
    background: #ffffff;
    border: 1px dashed #94a3b8;
    border-radius: 6px;
}
.gr-box, .gr-form { background: transparent !important; border: none !important; }
"""

HEADER = """
<div class='twt-header'>
    <div class='twt-title'>Think<span>While</span>Thinking</div>
    <div class='twt-sub'>Real-time reasoning failure detection · Step-level process supervision</div>
</div>
"""

EXAMPLES = [
    "If a bat and a ball cost $1.10 in total, and the bat costs $1 more than the ball, how much does the ball cost?",
    "A farmer has 17 sheep. All but 9 die. How many sheep are left?",
    "What is 15% of 80? Then add that to 25% of 60.",
    "If you have a 3-gallon jug and a 5-gallon jug, how do you measure exactly 4 gallons?",
]

with gr.Blocks(css=CSS, title="ThinkWhileThinking") as demo:
    gr.HTML(HEADER)
    with gr.Row():
        with gr.Column(scale=2):
            gr.HTML("<div class='input-label'>Problem</div>")
            problem_input = gr.Textbox(placeholder="Enter a math, logic, or reasoning problem...", lines=5, show_label=False)
            run_btn = gr.Button("▶  Analyze Reasoning", variant="primary")
            gr.Examples(examples=EXAMPLES, inputs=problem_input, label="Try an example")
        with gr.Column(scale=3):
            with gr.Row():
                with gr.Column(scale=3):
                    gr.HTML("<div class='panel-label'>Chain of Thought · Step Scores</div>")
                    cot_output = gr.HTML(value="<div class='msg empty'>// awaiting problem input</div>")
                with gr.Column(scale=1):
                    gr.HTML("<div class='panel-label'>Step Trace</div>")
                    trace_output = gr.HTML(value="<div class='msg empty'>// trace</div>")
    raw_output = gr.Textbox(visible=False)
    run_btn.click(fn=run_twt, inputs=[problem_input], outputs=[raw_output, cot_output, trace_output])
    problem_input.submit(fn=run_twt, inputs=[problem_input], outputs=[raw_output, cot_output, trace_output])

demo.launch()