"""
Continuation GRPO training on PROCEDURAL scenarios.

Loads the existing LoRA adapter from Round 1 (trained on 3 fixed benchmarks)
and continues training on procedurally generated scenarios to fix:
1. Memorization of 3 fixed action sequences → procedural variety
2. Low token entropy / overconfidence → higher temperature + diverse prompts
3. No generalization eval → held-out procedural eval at the end
4. ~6% format failures → increased format reward weight

This is a SHORT continuation (~30-40 min on a10g) not a full retrain.

Run with:
  hf jobs uv run --no-project --flavor a10g-small --timeout 75m \
    --secrets HF_TOKEN --with-requirements requirements-job.txt \
    train_continuation.py
"""

import os, sys, gc, json, re, pathlib, subprocess, time
from collections import defaultdict
import numpy as np

# ── Clone the Space repo ──
SPACE_REPO = "https://huggingface.co/spaces/A-HK/agentic-security-lab"
REPO_ROOT  = pathlib.Path("/app/repo")
if not (REPO_ROOT / "models.py").exists():
    print(f"Cloning {SPACE_REPO} ...")
    REPO_ROOT.mkdir(parents=True, exist_ok=True)
    subprocess.run(["git", "clone", "--depth", "1", SPACE_REPO, str(REPO_ROOT)], check=True)
    print("Clone done.")
else:
    print("Repo already present.")

sys.path.insert(0, str(REPO_ROOT))
os.environ["AGENTIC_SECURITY_LAB_ALLOW_ENV_BASE_FALLBACK"] = "true"
ARTIFACT_DIR = REPO_ROOT / "artifacts"
ARTIFACT_DIR.mkdir(exist_ok=True)

HF_TOKEN = os.environ.get("HF_TOKEN", "")
if HF_TOKEN:
    from huggingface_hub import login
    login(token=HF_TOKEN)

# ── Config ──
BASE_ADAPTER    = "A-HK/security-incident-responder-grpo"   # Round 1 checkpoint
BASE_MODEL      = "Qwen/Qwen2.5-3B-Instruct"
HUB_MODEL_ID    = "A-HK/security-incident-responder-grpo"   # Push back to same repo
OUTPUT_DIR      = str(ARTIFACT_DIR / "grpo_v2_checkpoint")

# Shorter run: 30 procedural episodes × 4 gens × 1 epoch
NUM_EPISODES         = 30
NUM_GENERATIONS      = 4
MAX_COMPLETION_LENGTH = 512
LEARNING_RATE        = 5e-7      # Lower LR for continuation
GRAD_ACCUM           = 4
NUM_TRAIN_EPOCHS     = 1
SAVE_STEPS           = 999999
LOGGING_STEPS        = 1
TEMPERATURE          = 1.0       # Higher than before (was 0.9) for diversity
BATCH_SIZE           = 1

PUSH_TO_HUB = bool(HF_TOKEN)

print(f"=== CONTINUATION TRAINING (Procedural Scenarios) ===")
print(f"Base adapter: {BASE_ADAPTER}")
print(f"Episodes: {NUM_EPISODES} (procedural) | Gens: {NUM_GENERATIONS} | Epochs: {NUM_TRAIN_EPOCHS}")
print(f"Temperature: {TEMPERATURE} | LR: {LEARNING_RATE}")

# ── Stub out missing optional deps ──
import types as _types
class _Stub(_types.ModuleType):
    def __getattr__(self, name):
        if name.startswith("__") and name.endswith("__"): raise AttributeError(name)
        return _Stub(f"{self.__name__}.{name}")
    def __call__(self, *a, **k): return None
for _s in ["mergekit", "mergekit.config", "mergekit.merge_methods",
           "llm_blender", "llm_blender.blender", "llm_blender.pair_ranker"]:
    sys.modules.setdefault(_s, _Stub(_s))

# ── Imports ──
import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainerCallback
from peft import PeftModel
from trl import GRPOTrainer, GRPOConfig
from server.agentic_security_lab_environment import AgenticSecurityLabEnvironment
from models import AgenticSecurityLabAction
from training.procedural_scenarios import generate_procedural_scenario

# ── Tool call parsing (same as round 1) ──
_TOOL_RE = re.compile(r'\{\s*"name"\s*:\s*"(\w+)"\s*,\s*"arguments"\s*:\s*(\{[^{}]*\})\s*\}', re.DOTALL)
_QWEN_TOOL_RE = re.compile(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', re.DOTALL)

def _completion_to_str(c):
    if isinstance(c, str): return c
    if isinstance(c, list): return "\n".join(i.get("content","") if isinstance(i,dict) else str(i) for i in c)
    return str(c)

def _extract_tool_calls(text):
    calls = []
    for m in _TOOL_RE.finditer(text):
        try: calls.append((m.group(1), json.loads(m.group(2))))
        except: pass
    for m in _QWEN_TOOL_RE.finditer(text):
        try:
            obj = json.loads(m.group(1)); n = obj.get("name","")
            if n: calls.append((n, obj.get("arguments",{})))
        except: pass
    return calls

# ── Environment replay using procedural scenarios ──
def _replay_on_scenario(calls, scenario):
    """Replay tool calls on a fresh env initialized with a specific scenario."""
    env = AgenticSecurityLabEnvironment("easy")
    env.reset(task_name="easy", mode="benchmark")
    # Override with procedural scenario data
    env._scenario = scenario
    state = env._state
    state.packages = scenario["packages"]
    state.dependents = scenario["dependents"]
    state.secrets = {k: dict(v) for k, v in scenario["secrets"].items()}
    state.max_steps = scenario["max_steps"]
    state.exfiltration_step = scenario["exfiltration_step"]
    state.pending_hidden_iocs = list(scenario.get("hidden_iocs", []))

    n_invalid = 0
    for name, args in calls:
        s = env._state
        if s.incident_contained or s.attacker_succeeded or s.step_count >= s.max_steps:
            break
        pkg = args.get("package", "")
        sec = args.get("secret", "")
        team = args.get("team", "")
        try:
            if name == "scan_logs" and pkg:
                env.step(AgenticSecurityLabAction(command="scan_logs", parameters={"package": pkg}))
            elif name == "inspect_package" and pkg:
                env.step(AgenticSecurityLabAction(command="inspect_package", parameters={"package": pkg}))
            elif name == "check_dependents" and pkg:
                env.step(AgenticSecurityLabAction(command="check_dependents", parameters={"package": pkg}))
            elif name in ("quarantine_package", "quarantine") and pkg:
                env.step(AgenticSecurityLabAction(command="quarantine", parameters={"package": pkg}))
            elif name == "rotate_secret" and sec:
                env.step(AgenticSecurityLabAction(command="rotate_secret", parameters={"secret": sec}))
            elif name in ("notify_team", "notify") and team:
                env.step(AgenticSecurityLabAction(command="notify", parameters={"team": team}))
            elif name in ("conclude_incident", "conclude"):
                env.step(AgenticSecurityLabAction(command="conclude", parameters={}))
            else:
                n_invalid += 1
        except:
            n_invalid += 1
    return env, len(calls), n_invalid

def compute_trajectory_reward(env):
    """Same reward function as round 1."""
    state = env._state
    scenario = env._scenario
    req = scenario.get("required_actions", {})
    req_q = req.get("quarantine", [])
    req_r = req.get("rotate_secret", [])
    req_n = req.get("notify", [])

    q = len(set(state.quarantined) & set(req_q)) / max(1, len(req_q))
    r = len(set(state.rotated_secrets) & set(req_r)) / max(1, len(req_r))
    n = len(set(state.notified_teams) & set(req_n)) / max(1, len(req_n))

    benchmark = q * 1.2 + r * 1.0 + n * 0.8
    concluded = state.incident_contained
    comp_bonus = 0.5 if concluded else -0.3
    eff = max(0.0, 0.5 * (1.0 - state.step_count / max(1, state.max_steps))) if concluded and state.step_count > 0 else 0.0
    actions = [t.get("command", "") for t in state.trajectory_log]
    div = min(0.5, len(set(actions)) * 0.1) if len(actions) >= 2 else 0.0
    inv_bonus = 0.2 if len(set(state.inspected) | set(state.scanned_logs)) > 0 else 0.0

    sigs = [f"{t.get('command', '')}:{json.dumps(t.get('params', {}), sort_keys=True)}" for t in state.trajectory_log]
    rep_pen = min(0.5, (len(sigs) - len(set(sigs))) * 0.1)
    inv_pen = state.invalid_action_count * 0.15
    fp_pen = state.false_positive_count * 0.3
    atk_pen = 1.0 if state.attacker_succeeded else 0.0
    pre_pen = 1.0 if (concluded and q == 0 and r == 0) else 0.0

    total = max(-1.0, min(5.0,
        benchmark + comp_bonus + eff + div + inv_bonus
        - rep_pen - inv_pen - fp_pen - atk_pen - pre_pen))
    return total, {
        "benchmark": benchmark, "q_ratio": q, "r_ratio": r, "n_ratio": n,
        "concluded": concluded, "steps": state.step_count,
        "efficiency": eff, "diversity": div
    }

# ── Reward functions ──
CURRENT_DIFFICULTY = "easy"

def trajectory_reward_func(completions, difficulty=None, scenario_seed=None, **kwargs):
    """Replay completions on PROCEDURAL scenarios (not fixed benchmarks)."""
    rewards = []
    diffs = difficulty if difficulty is not None else [CURRENT_DIFFICULTY] * len(completions)
    seeds = scenario_seed if scenario_seed is not None else [None] * len(completions)
    for comp, diff, seed in zip(completions, diffs, seeds):
        text = _completion_to_str(comp)
        calls = _extract_tool_calls(text)
        if not calls:
            rewards.append(-0.5)
            continue
        scenario = generate_procedural_scenario(difficulty=diff, seed=seed)
        env, _, ni = _replay_on_scenario(calls, scenario)
        rv, info = compute_trajectory_reward(env)
        if not env._state.incident_contained:
            rv -= 0.5
        rv -= ni * 0.1
        rv = max(-1.0, min(5.0, rv))
        rewards.append(float(rv))
    return rewards

def format_reward_func(completions, **kwargs):
    """Stronger format penalty than round 1."""
    rewards = []
    for c in completions:
        calls = _extract_tool_calls(_completion_to_str(c))
        rewards.append(min(1.0, len(calls) * 0.15) if calls else -0.5)
    return rewards

# ── System prompt ──
SYSTEM_PROMPT = (
    "You are an expert security incident responder. Respond with a sequence of tool calls "
    "to investigate, contain, and remediate the incident.\n\n"
    "Available tools (call each as a JSON object on its own line):\n"
    '{"name": "scan_logs", "arguments": {"package": "<pkg@ver>"}}\n'
    '{"name": "inspect_package", "arguments": {"package": "<pkg@ver>"}}\n'
    '{"name": "check_dependents", "arguments": {"package": "<pkg@ver>"}}\n'
    '{"name": "quarantine_package", "arguments": {"package": "<pkg@ver>"}}\n'
    '{"name": "rotate_secret", "arguments": {"secret": "<SECRET_NAME>"}}\n'
    '{"name": "notify_team", "arguments": {"team": "<team-name>"}}\n'
    '{"name": "conclude_incident", "arguments": {}}\n\n'
    "IMPORTANT: Replace placeholders with actual values from the incident description.\n"
    "Recommended order: scan/inspect suspicious packages -> quarantine malicious ones -> "
    "check_dependents -> rotate compromised secrets -> notify affected teams -> conclude.\n"
    "Output ONLY the tool call JSON objects, one per line. Be efficient."
)

# ── Build PROCEDURAL dataset ──
def build_procedural_dataset(num_episodes, difficulty_dist=None):
    """Each episode gets a UNIQUE procedural scenario with a stored seed."""
    if difficulty_dist is None:
        difficulty_dist = {"easy": 0.4, "medium": 0.35, "hard": 0.25}
    import random
    rng = random.Random(12345)
    rows = []
    for i in range(num_episodes):
        diff = rng.choices(list(difficulty_dist.keys()), list(difficulty_dist.values()))[0]
        seed = rng.randint(100000, 999999)
        scenario = generate_procedural_scenario(difficulty=diff, seed=seed)

        pkgs = list(scenario["packages"].keys())
        secs = list(scenario["secrets"].keys())
        teams = set()
        for ts in scenario["dependents"].values():
            teams.update(ts)

        urgency = rng.choice(["URGENT", "CRITICAL", "HIGH PRIORITY", "IMMEDIATE"])
        act = rng.choice(["Investigate", "Analyze", "Examine", "Assess"])
        desc = (
            f"[{urgency}] {scenario['description']}\n\n"
            f"Packages in scope: {pkgs}\n"
            f"Known secrets at risk: {secs}\n"
            f"Known affected teams: {sorted(teams)}\n"
            f"Budget: {scenario['max_steps']} steps, exfiltration in {scenario['exfiltration_step']} steps.\n\n"
            f"{act} all packages, quarantine the malicious ones, rotate all secrets, "
            f"notify all affected teams, then call conclude_incident."
        )
        rows.append({
            "prompt": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": desc}
            ],
            "difficulty": diff,
            "scenario_seed": seed,
        })
    return Dataset.from_list(rows)

# ── Load model with existing adapter ──
print("\n=== Loading base model + Round 1 adapter ===")
gpu_name = torch.cuda.get_device_name(0).lower() if torch.cuda.is_available() else ""
use_bf16 = any(x in gpu_name for x in ["a100", "a10", "h100", "l4", "l40"])
print(f"GPU: {gpu_name} | bf16={use_bf16}")

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16 if use_bf16 else torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, quantization_config=bnb_config, device_map="auto",
    torch_dtype=torch.bfloat16 if use_bf16 else torch.float16,
    attn_implementation="eager"
)
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Load the Round 1 LoRA adapter (is_trainable=True is MANDATORY for continued training)
print(f"Loading adapter from {BASE_ADAPTER}...")
model = PeftModel.from_pretrained(base_model, BASE_ADAPTER, is_trainable=True)
model.print_trainable_parameters()

# ── Build dataset ──
train_dataset = build_procedural_dataset(NUM_EPISODES)
print(f"\nDataset: {NUM_EPISODES} procedural episodes")
from collections import Counter
diff_counts = Counter(train_dataset["difficulty"])
print(f"Difficulty distribution: {dict(diff_counts)}")

# ── GRPO Config (NO peft_config — we already have a PeftModel) ──
grpo_config = GRPOConfig(
    output_dir=OUTPUT_DIR,
    num_generations=NUM_GENERATIONS,
    generation_batch_size=NUM_GENERATIONS,
    max_completion_length=MAX_COMPLETION_LENGTH,
    temperature=TEMPERATURE,
    beta=0.0,
    epsilon=0.2,
    epsilon_high=0.28,
    scale_rewards="group",
    loss_type="dapo",
    learning_rate=LEARNING_RATE,
    num_train_epochs=NUM_TRAIN_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    gradient_checkpointing=True,
    bf16=use_bf16,
    fp16=False,
    logging_steps=LOGGING_STEPS,
    logging_first_step=True,
    log_completions=True,
    disable_tqdm=True,
    save_strategy="steps",
    save_steps=SAVE_STEPS,
    save_total_limit=2,
    push_to_hub=PUSH_TO_HUB,
    hub_model_id=HUB_MODEL_ID,
    hub_strategy="end",
    warmup_steps=2,
    reward_weights=[1.0, 0.3],
    report_to=[],
)

# ── Metrics callback ──
class MetricsCallback(TrainerCallback):
    def __init__(self):
        self.step_count = 0
        self.rewards = []
        self.losses = []
        self.metrics_file = ARTIFACT_DIR / "grpo_v2_metrics.jsonl"
        if self.metrics_file.exists():
            self.metrics_file.unlink()

    def _log(self, data):
        with self.metrics_file.open("a") as f:
            f.write(json.dumps(data) + "\n")

    def on_log(self, args, state, control, logs=None, **kwargs):
        if not logs:
            return
        self.step_count += 1
        reward = logs.get("reward")
        loss = logs.get("loss") or logs.get("train_loss")
        reward_std = logs.get("reward_std", 0)
        frac_zero = logs.get("frac_reward_zero_std", 0)
        entropy = logs.get("entropy", 0)
        if reward is not None:
            self.rewards.append(reward)
        if loss is not None:
            self.losses.append(loss)
        self._log({
            "step": self.step_count, "reward": reward,
            "reward_std": reward_std, "frac_reward_zero_std": frac_zero,
            "entropy": entropy, "loss": loss,
            "lr": logs.get("learning_rate", 0), "phase": "procedural_continuation"
        })
        r_str = f"{reward:.3f}" if reward is not None else "None"
        l_str = f"{loss:.6f}" if loss is not None else "None"
        recent = self.rewards[-5:] if len(self.rewards) >= 5 else self.rewards
        avg = np.mean(recent) if recent else 0
        print(f"  step {self.step_count}: reward={r_str} avg5={avg:.3f} "
              f"std={reward_std:.4f} frac0={frac_zero:.2f} entropy={entropy:.3f} loss={l_str}")


# ── Held-out eval on PROCEDURAL scenarios ──
def evaluate_procedural(model, tokenizer, n_episodes=15, label="procedural_eval"):
    """Evaluate on NEVER-SEEN procedural scenarios."""
    print(f"\n{'='*60}\nEvaluating: {label} ({n_episodes} procedural episodes)\n{'='*60}")
    model.eval()
    results = []
    import random
    eval_rng = random.Random(99999)

    for i in range(n_episodes):
        diff = ["easy", "easy", "medium", "medium", "hard"][i % 5]
        seed = eval_rng.randint(2000000, 2999999)
        scenario = generate_procedural_scenario(difficulty=diff, seed=seed)

        pkgs = list(scenario["packages"].keys())
        secs = list(scenario["secrets"].keys())
        teams = set()
        for ts in scenario["dependents"].values():
            teams.update(ts)

        desc = (
            f"{scenario['description']}\n\n"
            f"Packages in scope: {pkgs}\n"
            f"Known secrets at risk: {secs}\n"
            f"Known affected teams: {sorted(teams)}\n"
            f"Budget: {scenario['max_steps']} steps.\n\n"
            f"Investigate, quarantine malicious, rotate secrets, notify teams, conclude."
        )
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": desc}
        ]
        prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(prompt_text, return_tensors="pt", truncation=True, max_length=2048).to(model.device)

        with torch.no_grad():
            output = model.generate(
                **inputs, max_new_tokens=MAX_COMPLETION_LENGTH,
                temperature=0.7, do_sample=True,
                pad_token_id=tokenizer.pad_token_id
            )
        generated = tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        calls = _extract_tool_calls(generated)

        if calls:
            env, _, _ = _replay_on_scenario(calls, scenario)
            rv, info = compute_trajectory_reward(env)
            if not env._state.incident_contained:
                rv -= 0.5
        else:
            rv = -0.5
            info = {"benchmark": 0, "q_ratio": 0, "r_ratio": 0, "n_ratio": 0,
                    "concluded": False, "steps": 0, "efficiency": 0, "diversity": 0}

        results.append({
            "episode": i, "difficulty": diff, "seed": seed,
            "reward": rv, "info": info,
            "n_calls": len(calls), "text": generated[:300]
        })
        print(f"  Ep {i} ({diff}): r={rv:.3f} bench={info['benchmark']:.2f} "
              f"calls={len(calls)} q={info['q_ratio']:.1f} r={info['r_ratio']:.1f} n={info['n_ratio']:.1f}")

    avg = np.mean([r["reward"] for r in results])
    by_diff = {}
    for r in results:
        by_diff.setdefault(r["difficulty"], []).append(r["reward"])
    print(f"\n{label} overall avg: {avg:.3f}")
    for d in ["easy", "medium", "hard"]:
        if d in by_diff:
            print(f"  {d}: avg={np.mean(by_diff[d]):.3f} (n={len(by_diff[d])})")

    format_failures = sum(1 for r in results if r["n_calls"] == 0)
    print(f"  Format failures: {format_failures}/{n_episodes} ({format_failures/n_episodes*100:.1f}%)")
    model.train()
    return results, avg

# ── Evaluate BEFORE continuation (on procedural) ──
print("\n=== Pre-continuation eval on procedural scenarios ===")
pre_results, pre_avg = evaluate_procedural(model, tokenizer, n_episodes=15, label="PRE-CONTINUATION (procedural)")

# ── Train ──
gc.collect()
torch.cuda.empty_cache()
model.train()

metrics_cb = MetricsCallback()
trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[trajectory_reward_func, format_reward_func],
    train_dataset=train_dataset,
    args=grpo_config,
    callbacks=[metrics_cb],
)

print(f"\n{'='*60}")
print(f"GRPO Continuation Training | Procedural Scenarios")
print(f"{'='*60}")
train_result = trainer.train()
print(f"Done: {train_result.global_step} steps, loss={train_result.training_loss:.4f}")

# ── Evaluate AFTER continuation ──
post_results, post_avg = evaluate_procedural(model, tokenizer, n_episodes=15, label="POST-CONTINUATION (procedural)")

# ── Also eval on fixed benchmarks to check we didn't regress ──
def _replay_on_fixed_benchmark(calls, difficulty):
    from scenarios import get_scenario
    scenario = get_scenario(difficulty)
    return _replay_on_scenario(calls, scenario)

def evaluate_benchmark(model, tokenizer, n_episodes=5, label="benchmark"):
    print(f"\n{'='*60}\nEvaluating: {label} (fixed benchmarks)\n{'='*60}")
    model.eval()
    results = []
    for i in range(n_episodes):
        diff = ["easy", "easy", "easy", "medium", "hard"][i % 5]
        from scenarios import get_scenario
        scenario = get_scenario(diff)
        pkgs = list(scenario["packages"].keys())
        secs = list(scenario["secrets"].keys())
        teams = set()
        for ts in scenario["dependents"].values():
            teams.update(ts)
        desc = (
            f"{scenario['description']}\n\nPackages in scope: {pkgs}\n"
            f"Known secrets at risk: {secs}\nKnown affected teams: {sorted(teams)}\n"
            f"Budget: {scenario['max_steps']} steps.\n\n"
            f"Investigate, quarantine malicious, rotate secrets, notify teams, conclude."
        )
        messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": desc}]
        prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(prompt_text, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
        with torch.no_grad():
            output = model.generate(**inputs, max_new_tokens=MAX_COMPLETION_LENGTH, temperature=0.7, do_sample=True, pad_token_id=tokenizer.pad_token_id)
        generated = tokenizer.decode(output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        calls = _extract_tool_calls(generated)
        if calls:
            env, _, _ = _replay_on_fixed_benchmark(calls, diff)
            rv, info = compute_trajectory_reward(env)
            if not env._state.incident_contained: rv -= 0.5
        else:
            rv = -0.5
            info = {"benchmark": 0, "q_ratio": 0, "r_ratio": 0, "n_ratio": 0, "concluded": False, "steps": 0}
        results.append({"episode": i, "difficulty": diff, "reward": rv, "info": info, "n_calls": len(calls)})
        print(f"  Ep {i} ({diff}): r={rv:.3f} bench={info['benchmark']:.2f} calls={len(calls)}")
    avg = np.mean([r["reward"] for r in results])
    print(f"\n{label} avg: {avg:.3f}")
    model.train()
    return results, avg

bench_results, bench_avg = evaluate_benchmark(model, tokenizer, n_episodes=5, label="POST-CONTINUATION (benchmark)")

# ── Generate plots ──
try:
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    metrics_path = ARTIFACT_DIR / "grpo_v2_metrics.jsonl"
    if metrics_path.exists():
        rows = [json.loads(l) for l in metrics_path.read_text().splitlines() if l.strip()]
        valid = [r for r in rows if r.get("reward") is not None]
        steps = [r["step"] for r in valid]
        rewards = [float(r["reward"]) for r in valid]
        reward_stds = [float(r.get("reward_std") or 0) for r in valid]

        fig, ax = plt.subplots(figsize=(10, 5))
        ax.plot(steps, rewards, alpha=0.4, color="blue", label="Per-step reward")
        w = min(5, max(1, len(rewards) // 3))
        if len(rewards) >= w:
            roll = np.convolve(rewards, np.ones(w) / w, mode='valid')
            ax.plot(steps[w-1:], roll, color="blue", lw=2, label=f"Rolling avg (w={w})")
        ax.axhline(pre_avg, color="red", ls="--", lw=2, label=f"Pre (procedural): {pre_avg:.3f}")
        ax.axhline(post_avg, color="green", ls="--", lw=2, label=f"Post (procedural): {post_avg:.3f}")
        ax.set_xlabel("Training Step"); ax.set_ylabel("Episode Reward")
        ax.set_title("Continuation Training on Procedural Scenarios"); ax.legend(); ax.grid(alpha=0.3)
        fig.tight_layout(); fig.savefig(str(ARTIFACT_DIR / "continuation_reward_curve.png"), dpi=150); plt.close(fig)

        fig, ax = plt.subplots(figsize=(10, 4))
        ax.plot(steps, reward_stds, color="purple", lw=1.5)
        ax.axhline(0, color="red", ls="--", alpha=0.5)
        ax.set_xlabel("Step"); ax.set_ylabel("Reward STD")
        ax.set_title("Reward STD During Continuation (procedural scenarios)"); ax.grid(alpha=0.3)
        fig.tight_layout(); fig.savefig(str(ARTIFACT_DIR / "continuation_reward_std.png"), dpi=150); plt.close(fig)

        fig, ax = plt.subplots(figsize=(8, 5))
        pre_by_diff, post_by_diff = {}, {}
        for r in pre_results: pre_by_diff.setdefault(r["difficulty"], []).append(r["reward"])
        for r in post_results: post_by_diff.setdefault(r["difficulty"], []).append(r["reward"])
        diffs = ["easy", "medium", "hard"]; x = np.arange(len(diffs)); width = 0.35
        pre_means = [np.mean(pre_by_diff.get(d, [0])) for d in diffs]
        post_means = [np.mean(post_by_diff.get(d, [0])) for d in diffs]
        ax.bar(x - width/2, pre_means, width, label=f"Before: {pre_avg:.2f}", color="lightcoral")
        ax.bar(x + width/2, post_means, width, label=f"After: {post_avg:.2f}", color="lightgreen")
        ax.set_xlabel("Difficulty"); ax.set_ylabel("Avg Reward (procedural)")
        ax.set_title("Generalization to Unseen Procedural Scenarios")
        ax.set_xticks(x); ax.set_xticklabels(diffs); ax.legend(); ax.grid(alpha=0.3, axis="y")
        fig.tight_layout(); fig.savefig(str(ARTIFACT_DIR / "generalization_comparison.png"), dpi=150); plt.close(fig)
    print("All plots saved.")
except Exception as e:
    print(f"Plot error: {e}"); import traceback; traceback.print_exc()

# ── Push to Hub ──
model.save_pretrained(str(ARTIFACT_DIR / "lora_v2_adapter"))
tokenizer.save_pretrained(str(ARTIFACT_DIR / "lora_v2_adapter"))
if PUSH_TO_HUB:
    from huggingface_hub import HfApi; api = HfApi(token=HF_TOKEN)
    model.push_to_hub(HUB_MODEL_ID, token=HF_TOKEN)
    tokenizer.push_to_hub(HUB_MODEL_ID, token=HF_TOKEN)
    for fname in ["continuation_reward_curve.png", "continuation_reward_std.png",
                   "generalization_comparison.png", "grpo_v2_metrics.jsonl"]:
        fpath = ARTIFACT_DIR / fname
        if fpath.exists():
            api.upload_file(path_or_fileobj=str(fpath), path_in_repo=f"artifacts/{fname}",
                            repo_id=HUB_MODEL_ID, repo_type="model", token=HF_TOKEN)
            print(f"  uploaded artifacts/{fname}")
    print(f"\n-> https://huggingface.co/{HUB_MODEL_ID}")

# ── Summary ──
print(f"\n{'='*60}")
print(f"CONTINUATION TRAINING SUMMARY")
print(f"{'='*60}")
print(f"Pre-continuation  (procedural, 15 ep): {pre_avg:.3f}")
print(f"Post-continuation (procedural, 15 ep): {post_avg:.3f}")
print(f"Post-continuation (fixed bench, 5 ep): {bench_avg:.3f}")
imp = post_avg - pre_avg
print(f"Procedural improvement: {imp:+.3f} ({imp/max(0.001,abs(pre_avg))*100:+.1f}%)")
if metrics_cb.rewards:
    v2_rows = [json.loads(l) for l in (ARTIFACT_DIR / "grpo_v2_metrics.jsonl").read_text().splitlines() if l.strip()]
    zero_std = sum(1 for r in v2_rows if r.get("reward_std") == 0 and r.get("reward") is not None)
    total = len([r for r in v2_rows if r.get("reward") is not None])
    print(f"Zero-std steps: {zero_std}/{total} ({zero_std/max(1,total)*100:.1f}%) — Round 1 was 30.6%")
    ff_pre = sum(1 for r in pre_results if r["n_calls"] == 0)
    ff_post = sum(1 for r in post_results if r["n_calls"] == 0)
    print(f"Format failures: pre={ff_pre}/15 post={ff_post}/15")
    print(f"Rewards: min={min(metrics_cb.rewards):.3f} max={max(metrics_cb.rewards):.3f} avg={np.mean(metrics_cb.rewards):.3f}")