"""
Narada GRPO Training Script
Equivalent to narada_grpo.ipynb but runs as a single process.

Required env vars:
    HF_TOKEN       — HuggingFace write token
    HF_PUSH_REPO   — where to push adapter, e.g. "KrishVenky/narada-detective-lora"

Optional env vars:
    ENV_URL        — Narada environment (default: HF Space)
    BASE_MODEL     — base model ID (default: Qwen/Qwen3-1.7B)
    LORA_RANK      — LoRA rank (default: 16)
"""

from __future__ import annotations

import asyncio
import json
import math
import os
import re
import sys
import textwrap
import threading
import time
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any, Dict, List, Optional

# Must come before transformers/torchao: torchao 0.17+ calls register_constant
# which was added in PyTorch 2.7. Stub it so the import chain works on 2.6.
import torch
if not hasattr(torch.utils._pytree, "register_constant"):
    torch.utils._pytree.register_constant = lambda cls: cls

from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments

import nest_asyncio
nest_asyncio.apply()

import websockets
from datasets import Dataset
from unsloth import FastLanguageModel

# ── Config ────────────────────────────────────────────────────────────────────

HF_TOKEN     = os.environ["HF_TOKEN"]
HF_PUSH_REPO = os.environ["HF_PUSH_REPO"]
BASE_MODEL   = os.environ.get("BASE_MODEL", "Qwen/Qwen3-1.7B")
ENV_URL      = os.environ.get("ENV_URL", "https://krishvenky-narada-env.hf.space")
LORA_RANK    = int(os.environ.get("LORA_RANK", "16"))
ADAPTER_NAME = "narada-detective-lora"

os.environ["HF_TOKEN"] = HF_TOKEN

CURRICULUM = [
    {"task": "monogenic",          "steps": 80},
    {"task": "oligogenic",         "steps": 60},
    {"task": "phenotype_mismatch", "steps": 60},
]
EVAL_SEEDS       = [42, 7, 999, 1337, 2024]
N_SEEDS_PER_TASK = 40
MAX_SEQ_LEN      = 2048
MINI_BATCH_SIZE  = 2
GRAD_ACCUM       = 4
LR               = 5e-6
WARMUP_STEPS     = 20

# ── System prompt ─────────────────────────────────────────────────────────────

SYSTEM_PROMPT = textwrap.dedent("""
You are an expert clinical geneticist. Generate a DIAGNOSTIC PLAN: 3-5 JSON action blocks
to navigate the gene-disease knowledge graph and identify the causal variant.

Output each action block on its own line. End with flag_causal when you are confident.

ACTIONS (one JSON per line, no other text):
  {"action_type": "hop",          "node_id": "<id>",       "reasoning": "<one sentence>"}
  {"action_type": "flag_causal",  "variant_id": "VAR:...", "reasoning": "<one sentence>"}
  {"action_type": "backtrack",                              "reasoning": "<one sentence>"}
  {"action_type": "summarise_trail",                        "reasoning": "<one sentence>"}

STRATEGY:
  1. Navigate phenotype -> disease -> gene -> variant chains.
  2. BRCA1/TP53 is a DECOY if phenotypes are cardiac/neurological -- skip it.
  3. Oligogenic: flag ALL causal variants, not just the first one.
  4. Flag before step 8 for a timing bonus.
  5. ABSENT PHENOTYPES are strong rule-out signals -- use them.
""").strip()

# ── Environment helpers ───────────────────────────────────────────────────────

def format_obs(obs: Dict[str, Any]) -> str:
    lines = [
        f"STEP {obs['step']}/{obs['max_steps']} | Task: {obs['task_type']}",
        "",
        "PATIENT PHENOTYPES (present):",
    ]
    for hid, name in zip(obs["patient_phenotypes"], obs["phenotype_names"]):
        lines.append(f"  + {hid} -- {name}")

    absent_ids   = obs.get("phenotypes_absent") or []
    absent_names = obs.get("phenotype_absent_names") or []
    if absent_ids:
        lines += ["", "ABSENT PHENOTYPES (rule-out signal):"]
        for hid, name in zip(absent_ids, absent_names):
            lines.append(f"  - {hid} -- {name}")

    n = obs["current_node"]
    lines += [
        "",
        f"CURRENT NODE: [{n['type'].upper()}] {n['name']} ({n['id']})",
        f"  Neighbors ({len(n['connected_node_ids'])}): {', '.join(n['connected_node_ids'][:8])}",
    ]
    if obs.get("trail"):
        trail = [f"{t['name']}({t['id']})" for t in obs["trail"][-4:]]
        lines.append(f"  Trail: {' -> '.join(trail)}")

    lines += ["", "CANDIDATE VARIANTS:"]
    for v in obs["candidate_variants"]:
        lines.append(
            f"  {v['id']} | {v['gene']} | {v['variant_type']} "
            f"| path={v['pathogenicity_score']:.2f} | {v['clinical_significance']}"
        )
    lines.append(f"\nStep reward: {obs['step_reward']:+.4f} | Cumulative: {obs['cumulative_reward']:.4f}")
    lines.append("Generate your diagnostic plan (3-5 JSON action blocks):")
    return "\n".join(lines)


def parse_all_actions(text: str) -> List[Dict[str, Any]]:
    actions = []
    for m in re.finditer(r'\{[^{}]*"action_type"[^{}]*\}', text, re.DOTALL):
        try:
            d = json.loads(m.group(0))
            atype = str(d.get("action_type", "")).lower()
            if atype not in ("hop", "flag_causal", "backtrack", "summarise_trail", "request_lab"):
                continue
            actions.append({
                "action_type": atype,
                "node_id":     str(d["node_id"])    if d.get("node_id")    else None,
                "variant_id":  str(d["variant_id"]) if d.get("variant_id") else None,
                "reasoning":   str(d.get("reasoning", ""))[:200],
            })
            if atype == "flag_causal":
                break
        except Exception:
            continue
    return actions or [{"action_type": "summarise_trail", "reasoning": "fallback"}]


def parse_action(text: str) -> Dict[str, Any]:
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m:
        return {"action_type": "summarise_trail", "reasoning": "fallback"}
    try:
        d = json.loads(m.group(0))
        atype = str(d.get("action_type", "summarise_trail")).lower()
        if atype not in ("hop", "flag_causal", "backtrack", "summarise_trail", "request_lab"):
            atype = "summarise_trail"
        return {
            "action_type": atype,
            "node_id":    str(d["node_id"])    if d.get("node_id")    else None,
            "variant_id": str(d["variant_id"]) if d.get("variant_id") else None,
            "reasoning":  str(d.get("reasoning", ""))[:200],
        }
    except Exception:
        return {"action_type": "summarise_trail", "reasoning": "parse error"}


async def run_episode_async(
    task_type: str,
    actions: List[Dict[str, Any]],
    seed: Optional[int] = None,
) -> float:
    ws_url = ENV_URL.replace("https://", "wss://").replace("http://", "ws://") + "/ws"
    async with websockets.connect(ws_url, open_timeout=30, ping_interval=20) as ws:
        reset_msg: Dict[str, Any] = {"type": "reset", "task_type": task_type}
        if seed is not None:
            reset_msg["seed"] = seed
        await ws.send(json.dumps(reset_msg))
        raw = json.loads(await ws.recv())
        if raw.get("type") == "error":
            return 0.1

        obs = raw["data"]["observation"]
        last_reward = 0.1

        for action in actions:
            if obs.get("done"):
                break
            await ws.send(json.dumps({"type": "step", "action": action}))
            raw = json.loads(await ws.recv())
            if raw.get("type") == "error":
                break
            data = raw["data"]
            obs  = data["observation"]
            last_reward = data["reward"]
            if obs.get("done"):
                return float(last_reward)

        return float(last_reward)


async def collect_episode_async(task_type: str, seed: Optional[int] = None) -> List[Dict]:
    ws_url = ENV_URL.replace("https://", "wss://").replace("http://", "ws://") + "/ws"
    steps: List[Dict] = []
    async with websockets.connect(ws_url, open_timeout=30, ping_interval=20) as ws:
        reset_msg: Dict[str, Any] = {"type": "reset", "task_type": task_type}
        if seed is not None:
            reset_msg["seed"] = seed
        await ws.send(json.dumps(reset_msg))
        raw = json.loads(await ws.recv())
        if raw.get("type") == "error":
            return steps
        obs = raw["data"]["observation"]
        steps.append({"prompt": format_obs(obs), "obs": obs, "task_type": task_type, "seed": seed})
    return steps


def run_episode(task_type: str, actions: List[Dict], seed: Optional[int] = None) -> float:
    return asyncio.get_event_loop().run_until_complete(run_episode_async(task_type, actions, seed))


def collect_episode(task_type: str, seed: Optional[int] = None) -> List[Dict]:
    return asyncio.get_event_loop().run_until_complete(collect_episode_async(task_type, seed))


# ── Reward tracker callback ───────────────────────────────────────────────────

class RewardTracker(TrainerCallback):
    def __init__(self, phase: str, log: List[Dict]):
        self.phase = phase
        self.log   = log

    def on_log(self, args: TrainingArguments, state: TrainerState,
               control: TrainerControl, logs=None, **kwargs):
        if logs is None:
            return
        step   = state.global_step
        reward = logs.get("reward", logs.get("train/reward", None))
        if reward is not None:
            self.log.append({"phase": self.phase, "step": step, "reward": float(reward)})


# ── Reward function ───────────────────────────────────────────────────────────

def clamp(v: float, lo: float = 0.01, hi: float = 0.99) -> float:
    return max(lo, min(hi, v)) if math.isfinite(v) else 0.1


async def _eval_one_async(text: str, task: str, seed: Any) -> float:
    actions = parse_all_actions(text)
    try:
        return clamp(await run_episode_async(task, actions, seed=seed))
    except Exception:
        return 0.1


def narada_reward(completions, prompts, task_type=None, seed=None, **kwargs):
    n     = len(completions)
    tasks = task_type if task_type is not None else ["monogenic"] * n
    seeds = seed       if seed      is not None else [None] * n

    texts = []
    for c in completions:
        if isinstance(c, list):
            texts.append(c[-1]["content"] if c else "")
        else:
            texts.append(str(c))

    async def _batch():
        return list(await asyncio.gather(*[
            _eval_one_async(t, task, s)
            for t, task, s in zip(texts, tasks, seeds)
        ]))

    return asyncio.get_event_loop().run_until_complete(_batch())


# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> None:
    # ── Load model ────────────────────────────────────────────────────────────
    os.environ["UNSLOTH_DISABLE_STATISTICS"] = "1"

    print(f"Loading {BASE_MODEL}...", flush=True)
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name    = BASE_MODEL,
        max_seq_length= MAX_SEQ_LEN,
        dtype         = None,
        load_in_4bit  = True,
        token         = HF_TOKEN,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r              = LORA_RANK,
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                          "gate_proj", "up_proj", "down_proj"],
        lora_alpha     = LORA_RANK * 2,
        lora_dropout   = 0.0,
        bias           = "none",
        use_gradient_checkpointing = "unsloth",
        random_state   = 42,
    )
    print(f"Model loaded. Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}", flush=True)

    # Disable Qwen3 thinking mode
    _orig = tokenizer.apply_chat_template
    def _no_think(*args, **kwargs):
        kwargs["enable_thinking"] = False
        return _orig(*args, **kwargs)
    tokenizer.apply_chat_template = _no_think
    print("Thinking mode disabled.", flush=True)

    # ── Build dataset (parallel collection) ──────────────────────────────────
    import random
    random.seed(42)
    train_seeds = random.sample(range(1, 10000), N_SEEDS_PER_TASK * 3)

    async def _collect_one(task: str, seed: int) -> Optional[Dict]:
        steps = await collect_episode_async(task, seed=seed)
        if not steps:
            return None
        return {
            "prompt":    [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": steps[0]["prompt"]},
            ],
            "task_type": task,
            "seed":      seed,
        }

    async def _collect_all() -> List[Dict]:
        pairs = []
        for i, phase in enumerate(CURRICULUM):
            task = phase["task"]
            for seed in train_seeds[i * N_SEEDS_PER_TASK : (i + 1) * N_SEEDS_PER_TASK]:
                pairs.append((task, seed))
        results = await asyncio.gather(*[_collect_one(t, s) for t, s in pairs])
        return [r for r in results if r is not None]

    print(f"Collecting {N_SEEDS_PER_TASK * len(CURRICULUM)} prompts in parallel...", flush=True)
    all_prompts = asyncio.get_event_loop().run_until_complete(_collect_all())
    dataset = Dataset.from_list(all_prompts)
    print(f"Dataset: {len(dataset)} prompts", flush=True)

    # ── GRPO config ───────────────────────────────────────────────────────────
    import shutil
    from trl import GRPOConfig, GRPOTrainer

    cache_path = "/tmp/unsloth_compiled_cache"
    if os.path.exists(cache_path):
        shutil.rmtree(cache_path)

    grpo_config = GRPOConfig(
        num_generations             = 8,
        temperature                 = 1.1,
        top_p                       = 0.95,
        learning_rate               = LR,
        per_device_train_batch_size = MINI_BATCH_SIZE,
        gradient_accumulation_steps = GRAD_ACCUM,
        warmup_steps                = WARMUP_STEPS,
        max_grad_norm               = 0.1,
        optim                       = "adamw_8bit",
        max_prompt_length           = 1200,
        max_completion_length       = 800,
        logging_steps               = 5,
        output_dir                  = f"/tmp/{ADAPTER_NAME}",
        report_to                   = "none",
    )

    # ── Zero-shot baseline (before any training) ─────────────────────────────
    print(f"\n{'='*60}", flush=True)
    print("BASELINE (zero-shot, untrained LoRA weights)", flush=True)
    print(f"{'='*60}", flush=True)
    FastLanguageModel.for_inference(model)
    baseline_results: Dict[str, float] = {}
    for phase in CURRICULUM:
        task = phase["task"]
        scores = []
        for es in EVAL_SEEDS:
            steps = collect_episode(task, seed=es)
            if not steps:
                continue
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": steps[0]["prompt"]},
            ]
            inputs = tokenizer.apply_chat_template(
                messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
            ).to("cuda")
            with torch.no_grad():
                out = model.generate(inputs, max_new_tokens=200, temperature=0.3)
            completion = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)
            scores.append(run_episode(task, [parse_action(completion)], seed=es))
        avg = sum(scores) / len(scores) if scores else 0.0
        baseline_results[task] = avg
        print(f"Baseline {task}: {avg:.4f}  (n={len(scores)})", flush=True)
    FastLanguageModel.for_training(model)

    # ── Curriculum training ───────────────────────────────────────────────────
    eval_results: Dict[str, float] = {}
    reward_log:   List[Dict]       = []
    step_offset = 0

    for phase in CURRICULUM:
        task    = phase["task"]
        n_steps = phase["steps"]

        phase_data = dataset.filter(lambda x: x["task_type"] == task)
        if len(phase_data) == 0:
            print(f"Skipping {task} — no data.", flush=True)
            continue

        print(f"\n{'='*60}", flush=True)
        print(f"Phase: {task}  |  {len(phase_data)} prompts  |  {n_steps} steps", flush=True)
        print(f"{'='*60}", flush=True)

        grpo_config.max_steps = n_steps

        tracker = RewardTracker(task, reward_log)
        tracker._step_offset = step_offset

        trainer = GRPOTrainer(
            model            = model,
            processing_class = tokenizer,
            reward_funcs     = narada_reward,
            args             = grpo_config,
            train_dataset    = phase_data,
            callbacks        = [tracker],
        )

        t0 = time.time()
        trainer.train()
        step_offset += n_steps
        print(f"Phase {task} done in {(time.time()-t0)/60:.1f} min", flush=True)

        # Eval
        FastLanguageModel.for_inference(model)
        scores = []
        for es in EVAL_SEEDS:
            steps = collect_episode(task, seed=es)
            if not steps:
                continue
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": steps[0]["prompt"]},
            ]
            inputs = tokenizer.apply_chat_template(
                messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
            ).to("cuda")
            with torch.no_grad():
                out = model.generate(inputs, max_new_tokens=200, temperature=0.3)
            completion = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)
            scores.append(run_episode(task, [parse_action(completion)], seed=es))

        avg = sum(scores) / len(scores) if scores else 0.0
        eval_results[task] = avg
        print(f"Eval {task}: {avg:.4f}  (n={len(scores)})", flush=True)
        FastLanguageModel.for_training(model)

    # ── Save & push adapter ───────────────────────────────────────────────────
    model.save_pretrained(ADAPTER_NAME)
    tokenizer.save_pretrained(ADAPTER_NAME)
    print(f"\nAdapter saved locally to ./{ADAPTER_NAME}", flush=True)

    model.push_to_hub(HF_PUSH_REPO, token=HF_TOKEN)
    tokenizer.push_to_hub(HF_PUSH_REPO, token=HF_TOKEN)
    print(f"Pushed to https://huggingface.co/{HF_PUSH_REPO}", flush=True)

    # ── Generate training curves ──────────────────────────────────────────────
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        colors = {"monogenic": "#2196F3", "oligogenic": "#FF9800", "phenotype_mismatch": "#E91E63"}

        # ── Fig 1: reward over all steps (curriculum) ─────────────────────────
        fig, ax = plt.subplots(figsize=(10, 5))
        global_step = 0
        for phase in CURRICULUM:
            task = phase["task"]
            pts  = [(e["step"] + global_step, e["reward"])
                    for e in reward_log if e["phase"] == task]
            if pts:
                xs, ys = zip(*pts)
                ax.plot(xs, ys, "o-", color=colors[task], label=task, linewidth=1.5, markersize=3)
            global_step += phase["steps"]
        ax.set_xlabel("Training step")
        ax.set_ylabel("Mean reward")
        ax.set_title("Narada GRPO — reward curve (curriculum order)")
        ax.legend()
        ax.grid(True, alpha=0.3)
        fig.tight_layout()
        fig.savefig("training_curve.png", dpi=150)
        plt.close(fig)

        # ── Fig 2: per-task before/after bar chart ─────────────────────────────
        tasks_list = [p["task"] for p in CURRICULUM]
        x          = range(len(tasks_list))
        w          = 0.35
        fig2, ax2  = plt.subplots(figsize=(9, 5))
        ax2.bar([i - w/2 for i in x],
                [baseline_results.get(t, 0) for t in tasks_list],
                w, label="Zero-shot baseline", color="#90A4AE")
        ax2.bar([i + w/2 for i in x],
                [eval_results.get(t, 0) for t in tasks_list],
                w, label="After GRPO", color="#43A047")
        ax2.set_xticks(list(x))
        ax2.set_xticklabels(tasks_list)
        ax2.set_ylabel("Avg reward (5 eval seeds)")
        ax2.set_title("Narada — zero-shot vs GRPO-trained (Qwen3-1.7B)")
        ax2.set_ylim(0, 1.0)
        ax2.legend()
        ax2.grid(True, axis="y", alpha=0.3)
        fig2.tight_layout()
        fig2.savefig("before_after.png", dpi=150)
        plt.close(fig2)
        print("Plots saved: training_curve.png, before_after.png", flush=True)

        # Upload plots to the adapter repo
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)
        for fname in ("training_curve.png", "before_after.png"):
            api.upload_file(
                path_or_fileobj=fname,
                path_in_repo=fname,
                repo_id=HF_PUSH_REPO,
                repo_type="model",
                commit_message=f"add {fname}",
            )
        print(f"Plots uploaded to https://huggingface.co/{HF_PUSH_REPO}", flush=True)
    except Exception as e:
        print(f"Plot generation skipped: {e}", flush=True)

    print("\n=== TRAINING COMPLETE ===", flush=True)
    print(f"{'Task':<25} {'Baseline':>10} {'Trained':>10} {'Delta':>10}", flush=True)
    print("-" * 57, flush=True)
    for task in [p["task"] for p in CURRICULUM]:
        b = baseline_results.get(task, 0.0)
        t = eval_results.get(task, 0.0)
        print(f"  {task:<23} {b:>10.4f} {t:>10.4f} {t-b:>+10.4f}", flush=True)


class _HealthHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.end_headers()
        self.wfile.write(b"Training in progress...")
    def log_message(self, *args):
        pass


if __name__ == "__main__":
    threading.Thread(
        target=lambda: HTTPServer(("0.0.0.0", 7860), _HealthHandler).serve_forever(),
        daemon=True,
    ).start()
    main()