#!/usr/bin/env python3
"""eval-A: paper-era pipeline simplified — CoT reasoning + GPT-4o judge step-based belief.

For each question (no system prompt = "no-prompt" config matching paper's +0.0207 entry):
  1) Run CoT: model produces full reasoning trace.
  2) Split trace into halves; ask GPT-4o judge: "given this partial reasoning, what
     probability is the reasoner converging to?" for mid + final.
  3) Compute MS slope from (mid → final) deltas, Brier from final vs truth.
"""
from __future__ import annotations
import asyncio, json, os, re, sys, time
from pathlib import Path
import aiohttp
import numpy as np
from scipy import stats

QUESTIONS_PATH = "/home/ubuntu/Martingale-Training/data/questions/paper_subset_437.json"
SGLANG_URL = "http://127.0.0.1:30000/generate"
OPENAI_URL = "https://api.openai.com/v1/chat/completions"
JUDGE_MODEL = "gpt-4o"
OPENAI_KEY = os.environ.get("OPENAI_API_KEY", "")
assert OPENAI_KEY, "set OPENAI_API_KEY"

OUT_PATH = Path("/home/ubuntu/runs/r1_distill_eval_A_paper_subset_437.json")
OUT_PATH.parent.mkdir(parents=True, exist_ok=True)

CONCURRENCY_REASONER = 16
CONCURRENCY_JUDGE = 16
MAX_NEW_TOKENS_TRACE = 2500
TEMPERATURE = 0.6

NUM_RE = re.compile(r"(?<![\d.])(?:0?\.\d+|1(?:\.0+)?|0(?:\.0+)?)(?!\d)")


def chat_prompt(text):
    return "<｜begin▁of▁sentence｜><｜User｜>" + text + "<｜Assistant｜>"


def cot_prompt(qtext, opt):
    return chat_prompt(
        f"Forecasting question:\n\n{qtext}\n\n"
        f"Think step by step about whether the resolution will be \"{opt}\". "
        f"Walk through your reasoning carefully. At the very end state your final probability "
        f"(between 0 and 1) on its own line as: P = X"
    )


def split_trace(trace):
    """Split trace into (mid_partial, full) for step-based belief extraction."""
    # Mid = first half of the actual reasoning (post-<think> if present)
    if "</think>" in trace:
        thought, rest = trace.split("</think>", 1)
        full = thought + "</think>" + rest
        mid_thought = thought[: len(thought) // 2]
        mid = mid_thought + "[...reasoning truncated here for mid-step probe...]"
    else:
        full = trace
        mid = trace[: len(trace) // 2] + "[...reasoning truncated here for mid-step probe...]"
    return mid, full


def judge_prompt(qtext, opt, trace_excerpt, label):
    return (
        f"You are an impartial judge extracting probabilities from a reasoner's chain of thought.\n\n"
        f"Forecasting question:\n{qtext}\n\n"
        f"The reasoner is estimating the probability that the resolution is \"{opt}\".\n\n"
        f"--- Reasoner's {label} reasoning ---\n{trace_excerpt}\n--- End of reasoning ---\n\n"
        f"Based ONLY on this {label} reasoning, what probability is the reasoner converging toward "
        f"for \"{opt}\"? Output ONLY a single number between 0 and 1 on its own line, "
        f"with no other text."
    )


def extract_prob(text):
    matches = NUM_RE.findall(text or "")
    if matches:
        try:
            v = float(matches[-1])
            if 0 <= v <= 1:
                return v
        except ValueError:
            pass
    return None


async def call_sglang(session, prompt, max_new):
    payload = {"text": prompt, "sampling_params": {"temperature": TEMPERATURE, "max_new_tokens": max_new, "top_p": 0.95}}
    async with session.post(SGLANG_URL, json=payload, timeout=aiohttp.ClientTimeout(total=900)) as r:
        d = await r.json()
        return d.get("text", "") if isinstance(d, dict) else str(d)


async def call_judge(session, prompt):
    payload = {
        "model": JUDGE_MODEL,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.0,
        "max_tokens": 30,
    }
    headers = {"Authorization": f"Bearer {OPENAI_KEY}", "Content-Type": "application/json"}
    async with session.post(OPENAI_URL, json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=120)) as r:
        d = await r.json()
        try:
            return d["choices"][0]["message"]["content"]
        except Exception:
            return ""


async def reason_step(session, sem_r, q, idx, total):
    qtext = (q.get("description") or "").replace("\n\n", "\n")
    title = q.get("title") or ""
    body = (qtext + " " + title).strip() if qtext else title
    outcomes = q.get("outcomes", ["yes", "no"])
    opt = str(outcomes[0]).strip()
    async with sem_r:
        try:
            trace = await call_sglang(session, cot_prompt(body, opt), MAX_NEW_TOKENS_TRACE)
        except Exception as e:
            return {"id": q.get("id"), "error_reason": str(e)}
        if idx % 25 == 0:
            print(f"  [reason {idx}/{total}] {q.get('id')} trace_len={len(trace)}", flush=True)
        return {
            "id": q.get("id"),
            "title": q.get("title"),
            "resolution": q.get("resolution"),
            "outcomes": q.get("outcomes"),
            "first_option": opt,
            "qtext": body,
            "trace": trace,
        }


async def judge_step(session, sem_j, rec, idx, total):
    if "trace" not in rec:
        return rec
    mid, full = split_trace(rec["trace"])
    async with sem_j:
        try:
            mid_text, final_text = await asyncio.gather(
                call_judge(session, judge_prompt(rec["qtext"], rec["first_option"], mid, "partial")),
                call_judge(session, judge_prompt(rec["qtext"], rec["first_option"], full, "full")),
            )
        except Exception as e:
            rec["error_judge"] = str(e)
            return rec
        rec["mid_raw"] = mid_text
        rec["final_raw"] = final_text
        rec["prior"] = extract_prob(mid_text)  # paper convention: mid = prior
        rec["posterior"] = extract_prob(final_text)  # final = posterior
        if idx % 25 == 0:
            print(f"  [judge  {idx}/{total}] {rec['id']} mid={rec['prior']} final={rec['posterior']}", flush=True)
        return rec


async def main():
    questions = json.load(open(QUESTIONS_PATH))
    print(f"loaded {len(questions)} questions", flush=True)
    sem_r = asyncio.Semaphore(CONCURRENCY_REASONER)
    sem_j = asyncio.Semaphore(CONCURRENCY_JUDGE)
    t0 = time.time()
    async with aiohttp.ClientSession() as session:
        print("=== stage 1: reasoning ===", flush=True)
        reason_tasks = [reason_step(session, sem_r, q, i, len(questions)) for i, q in enumerate(questions)]
        recs = await asyncio.gather(*reason_tasks)
        print(f"  reasoning done in {time.time()-t0:.1f}s", flush=True)
        print("=== stage 2: judge ===", flush=True)
        t1 = time.time()
        judge_tasks = [judge_step(session, sem_j, r, i, len(recs)) for i, r in enumerate(recs)]
        recs = await asyncio.gather(*judge_tasks)
        print(f"  judge done in {time.time()-t1:.1f}s", flush=True)

    json.dump(recs, open(OUT_PATH, "w"))
    print(f"\nDONE total {time.time()-t0:.1f}s; wrote {OUT_PATH}", flush=True)

    usable = [r for r in recs if r.get("prior") is not None and r.get("posterior") is not None]
    print(f"USABLE: {len(usable)}/{len(recs)}", flush=True)
    if not usable:
        return
    priors = np.array([r["prior"] for r in usable])
    posteriors = np.array([r["posterior"] for r in usable])
    deltas = posteriors - priors
    pc = priors - 0.5

    slope, intercept, r_value, p_value, _ = stats.linregress(pc, deltas)
    signed_lin_r2 = float(np.sign(slope)) * (r_value ** 2)
    print(f"\nMartingale slope: {slope:+.4f}  (signed linR^2: {signed_lin_r2:+.4f}, p={p_value:.3e})", flush=True)
    print(f"mean |delta|: {np.mean(np.abs(deltas)):.4f}", flush=True)

    y_resolve = np.array([1 if (r["resolution"] == r["first_option"]) else 0 for r in usable])
    bp = float(np.mean((y_resolve - priors) ** 2))
    bq = float(np.mean((y_resolve - posteriors) ** 2))
    print(f"Brier(mid): {bp:.4f}  |  Brier(final): {bq:.4f}  |  Delta: {bq-bp:+.4f}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())