"""Step 6: Preflossing experiment (Engelken-style separate phases).

Phase 1 — Pure flossing: only optimize L_floss, no task loss.
  engelken: L = (1/k) Σ λ_i²    (push all top-k toward 0, two-sided)
  cf:       L = (1/k) Σ max(0, λ_i)²  (only push positive λ toward 0)

Phase 2 — Pure task training: standard HRM ACT loss, no flossing.

Baseline mode (--prefloss-steps 0): skip phase 1, go straight to task training.

Key fix vs step3: lyap_act_steps defaults to halt_max_steps (16 for HRM).
"""
from __future__ import annotations
import sys, os, yaml, json, math, time, argparse
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F

HRM_DIR = Path("/home/yurenh2/rrm/hrm")
sys.path.insert(0, str(HRM_DIR))

from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1
from models.losses import ACTLossHead
from adam_atan2 import AdamATan2


def load_model(ckpt_root: Path, ckpt_name: str, device: str):
    cfg = yaml.safe_load((ckpt_root / "all_config.yaml").read_text())
    arch_cfg = dict(cfg["arch"])
    train_meta = json.loads((Path(cfg["data_path"]) / "train" / "dataset.json").read_text())
    arch_cfg.update(batch_size=cfg["global_batch_size"], seq_len=train_meta["seq_len"],
                    vocab_size=train_meta["vocab_size"],
                    num_puzzle_identifiers=train_meta["num_puzzle_identifiers"], causal=False)
    base = HierarchicalReasoningModel_ACTV1(arch_cfg)
    head = ACTLossHead(base, loss_type=arch_cfg["loss"]["loss_type"])
    sd = torch.load(ckpt_root / ckpt_name, map_location="cpu", weights_only=True)
    stripped = {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
    missing, unexpected = head.load_state_dict(stripped, strict=False)
    print(f"[load {ckpt_name}] missing={len(missing)} unexpected={len(unexpected)}")
    head.to(device)
    return head, base, cfg, train_meta


def jvp_train(f, x, v):
    return torch.autograd.functional.jvp(f, x, v=v, create_graph=True, strict=False)


def compute_joint_lyap_spec(base, batch, k_lyap, lyap_act_steps, device, seed):
    inner = base.inner
    cfg = inner.config
    B = batch["inputs"].shape[0]
    seq_full = cfg.seq_len + inner.puzzle_emb_len
    hidden = cfg.hidden_size
    D = seq_full * hidden

    z_H = inner.H_init.unsqueeze(0).expand(B, seq_full, hidden).clone().to(inner.forward_dtype)
    z_L = inner.L_init.unsqueeze(0).expand(B, seq_full, hidden).clone().to(inner.forward_dtype)
    seq_info = dict(cos_sin=inner.rotary_emb() if hasattr(inner, "rotary_emb") else None)
    input_embeddings = inner._input_embeddings(batch["inputs"], batch["puzzle_identifiers"])

    g = torch.Generator(device=device).manual_seed(seed)
    Q0 = torch.randn(B, 2*D, k_lyap, device=device, dtype=torch.float32, generator=g)
    Q, _ = torch.linalg.qr(Q0)
    log_R_sum = torch.zeros(B, k_lyap, device=device, dtype=torch.float32)
    n_steps = 0

    n_act = min(lyap_act_steps, cfg.halt_max_steps)
    for _act in range(n_act):
        for _h in range(cfg.H_cycles):
            for _l in range(cfg.L_cycles):
                v_H_j = Q[:, :D, :]
                v_L_j = Q[:, D:, :]
                v_comb = v_H_j + v_L_j
                new_v_L_cols = []
                f_L = lambda z: inner.L_level(z, z_H + input_embeddings, **seq_info)
                for i in range(k_lyap):
                    v_i = v_comb[:, :, i].reshape(B, seq_full, hidden).to(inner.forward_dtype)
                    z_L_new, Dv = jvp_train(f_L, z_L, v_i)
                    new_v_L_cols.append(Dv.reshape(B, D).to(torch.float32))
                new_v_L = torch.stack(new_v_L_cols, dim=-1)
                Q = torch.cat([v_H_j, new_v_L], dim=1)
                z_L = z_L_new
                Q, R = torch.linalg.qr(Q)
                log_R_sum = log_R_sum + R.diagonal(dim1=-2, dim2=-1).abs().clamp_min(1e-30).log()
                n_steps += 1
            v_H_j = Q[:, :D, :]
            v_L_j = Q[:, D:, :]
            v_comb = v_H_j + v_L_j
            new_v_H_cols = []
            f_H = lambda z: inner.H_level(z, z_L, **seq_info)
            for i in range(k_lyap):
                v_i = v_comb[:, :, i].reshape(B, seq_full, hidden).to(inner.forward_dtype)
                z_H_new, Dv = jvp_train(f_H, z_H, v_i)
                new_v_H_cols.append(Dv.reshape(B, D).to(torch.float32))
            new_v_H = torch.stack(new_v_H_cols, dim=-1)
            Q = torch.cat([new_v_H, v_L_j], dim=1)
            z_H = z_H_new
            Q, R = torch.linalg.qr(Q)
            log_R_sum = log_R_sum + R.diagonal(dim1=-2, dim2=-1).abs().clamp_min(1e-30).log()
            n_steps += 1

    return log_R_sum / max(n_steps, 1)


def load_train_batches(data_path: Path, batch_size: int, n_iters: int, seed: int = 0):
    rng = np.random.default_rng(seed)
    inputs = np.load(data_path / "train" / "all__inputs.npy")
    labels = np.load(data_path / "train" / "all__labels.npy")
    pid    = np.load(data_path / "train" / "all__puzzle_identifiers.npy")
    N = len(inputs)
    for _ in range(n_iters):
        idx = rng.choice(N, size=batch_size, replace=False)
        yield {
            "inputs": torch.from_numpy(inputs[idx].astype(np.int32)),
            "labels": torch.from_numpy(labels[idx].astype(np.int32)),
            "puzzle_identifiers": torch.from_numpy(pid[idx].astype(np.int32)),
        }


def evaluate(head, base, data_path, n_samples, batch_size, device, seed=42):
    rng = np.random.default_rng(seed)
    inputs = np.load(data_path / "test" / "all__inputs.npy")
    labels = np.load(data_path / "test" / "all__labels.npy")
    pid    = np.load(data_path / "test" / "all__puzzle_identifiers.npy")
    idx_all = rng.choice(len(inputs), size=n_samples, replace=False)
    head.eval()
    correct = 0; token_correct = 0; token_total = 0
    for s in range(0, n_samples, batch_size):
        e = min(s + batch_size, n_samples)
        idx = idx_all[s:e]
        batch = {
            "inputs": torch.from_numpy(inputs[idx].astype(np.int32)).to(device),
            "labels": torch.from_numpy(labels[idx].astype(np.int32)).to(device),
            "puzzle_identifiers": torch.from_numpy(pid[idx].astype(np.int32)).to(device),
        }
        with torch.no_grad():
            with torch.device(device):
                carry = base.initial_carry(batch)
            for _ in range(base.config.halt_max_steps):
                carry, outputs = base(carry=carry, batch=batch)
            preds = outputs["logits"].argmax(dim=-1)
            mask = batch["labels"] > 0
            exact = ((preds == batch["labels"]) | ~mask).all(dim=-1).float()
            correct += exact.sum().item()
            token_correct += ((preds == batch["labels"]) & mask).sum().item()
            token_total += mask.sum().item()
    return correct / n_samples, token_correct / max(token_total, 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt-root", required=True)
    ap.add_argument("--ckpt-name", default="step_13020")
    ap.add_argument("--prefloss-steps", type=int, default=500,
                    help="Phase 1: pure flossing steps. 0 = skip (baseline).")
    ap.add_argument("--train-steps", type=int, default=3000,
                    help="Phase 2: pure task training steps.")
    ap.add_argument("--floss-mode", choices=["engelken", "cf", "volume_cf"], default="engelken",
                    help="engelken: Σλ_i² (two-sided). cf: Σmax(0,λ_i)² (one-sided hinge). "
                         "volume_cf: max(0, mean_i λ_i)² over the measured top-k spectrum.")
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--floss-lr", type=float, default=1e-4,
                    help="LR for flossing phase (Engelken uses higher LR)")
    ap.add_argument("--train-lr", type=float, default=1e-5,
                    help="LR for task training phase")
    ap.add_argument("--k-lyap", type=int, default=2)
    ap.add_argument("--lyap-act-steps", type=int, default=16,
                    help="ACT steps for Lyapunov computation (default=halt_max_steps)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--eval-every", type=int, default=100)
    ap.add_argument("--eval-n", type=int, default=512)
    ap.add_argument("--eval-batch-size", type=int, default=32)
    ap.add_argument("--out", default="step6_log.json")
    args = ap.parse_args()

    device = "cuda"
    head, base, cfg, train_meta = load_model(Path(args.ckpt_root), args.ckpt_name, device)
    data_path = Path(cfg["data_path"])

    print(f"\n=== Initial eval (loaded {args.ckpt_name}) ===")
    acc0, tacc0 = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device)
    print(f"  initial: exact_acc = {acc0:.4f}  token_acc = {tacc0:.4f}")

    log = {"args": vars(args), "initial_acc": acc0, "initial_tok_acc": tacc0,
           "phase1_steps": [], "phase1_evals": [], "phase2_steps": [], "phase2_evals": []}

    global_step = 0

    # ========== PHASE 1: Pure flossing ==========
    if args.prefloss_steps > 0:
        print(f"\n=== Phase 1: Pure {args.floss_mode} flossing ({args.prefloss_steps} steps, lr={args.floss_lr}) ===")
        floss_optim = AdamATan2(head.parameters(), lr=args.floss_lr, betas=(0.9, 0.95), weight_decay=0.0)
        floss_iter = load_train_batches(data_path, args.batch_size, args.prefloss_steps, seed=args.seed)
        t0 = time.time()

        for step, batch in enumerate(floss_iter):
            batch = {k: v.to(device) for k, v in batch.items()}
            head.train()
            base.inner.puzzle_emb.eval()
            for p in base.inner.puzzle_emb.parameters():
                p.requires_grad_(False)

            lyap_spec = compute_joint_lyap_spec(
                base, batch, k_lyap=args.k_lyap,
                lyap_act_steps=args.lyap_act_steps, device=device,
                seed=args.seed + step,
            )

            if args.floss_mode == "engelken":
                floss_loss = (lyap_spec ** 2).mean()
            elif args.floss_mode == "volume_cf":
                floss_loss = (lyap_spec.mean(dim=1).clamp_min(0.0) ** 2).mean()
            else:
                floss_loss = (lyap_spec.clamp_min(0.0) ** 2).mean()

            floss_optim.zero_grad(set_to_none=True)
            floss_loss.backward()
            torch.nn.utils.clip_grad_norm_([p for p in head.parameters() if p.requires_grad], 1.0)
            floss_optim.step()

            lyap1 = lyap_spec[:, 0].detach()
            rec = {
                "step": step, "floss_loss": float(floss_loss.item()),
                "lyap1_mean": float(lyap1.mean().item()),
                "lyap1_max": float(lyap1.max().item()),
                "lyap_all_mean": float(lyap_spec.detach().mean().item()),
                "lyap_volume_mean": float(lyap_spec.detach().mean(dim=1).mean().item()),
                "lyap_volume_max": float(lyap_spec.detach().mean(dim=1).max().item()),
            }
            log["phase1_steps"].append(rec)
            if step % 10 == 0 or step == args.prefloss_steps - 1:
                print(f"  P1[{step:>4}/{args.prefloss_steps}] dt={time.time()-t0:.1f}s "
                      f"floss={rec['floss_loss']:.6f}  "
                      f"λ1={rec['lyap1_mean']:+.4f}  max={rec['lyap1_max']:+.4f}  "
                      f"λ_all={rec['lyap_all_mean']:+.4f}", flush=True)

            if (step + 1) % args.eval_every == 0:
                acc, tacc = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device)
                print(f"    >> P1 EVAL @ step {step+1}: exact_acc={acc:.4f}  (Δ={acc-acc0:+.4f})", flush=True)
                log["phase1_evals"].append({"step": step + 1, "acc": acc, "tok_acc": tacc})

        acc_p1, tacc_p1 = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device)
        print(f"  Phase 1 final: exact_acc={acc_p1:.4f}  (Δ from init: {acc_p1-acc0:+.4f})")
        log["phase1_final_acc"] = acc_p1
        log["phase1_evals"].append({"step": args.prefloss_steps, "acc": acc_p1, "tok_acc": tacc_p1})
        global_step = args.prefloss_steps
    else:
        print("\n=== Phase 1 skipped (baseline mode) ===")
        log["phase1_final_acc"] = acc0

    # ========== PHASE 2: Pure task training ==========
    print(f"\n=== Phase 2: Pure task training ({args.train_steps} steps, lr={args.train_lr}) ===")
    train_optim = AdamATan2(head.parameters(), lr=args.train_lr, betas=(0.9, 0.95),
                            weight_decay=cfg["weight_decay"])
    train_iter = load_train_batches(data_path, args.batch_size, args.train_steps,
                                    seed=args.seed + 10000)
    t0 = time.time()
    acc_ref = log.get("phase1_final_acc", acc0)

    for step, batch in enumerate(train_iter):
        batch = {k: v.to(device) for k, v in batch.items()}
        head.train()
        base.inner.puzzle_emb.eval()
        for p in base.inner.puzzle_emb.parameters():
            p.requires_grad_(False)

        with torch.device(device):
            carry = base.initial_carry(batch)
        sup_loss_sum = 0.0
        n_loss = 0
        for _ in range(base.config.halt_max_steps):
            carry, l, metrics, _, all_finish = head(return_keys=[], carry=carry, batch=batch)
            sup_loss_sum = sup_loss_sum + l
            n_loss += 1
            if all_finish:
                break
        sup_loss = sup_loss_sum / max(n_loss, 1) / args.batch_size

        train_optim.zero_grad(set_to_none=True)
        sup_loss.backward()
        torch.nn.utils.clip_grad_norm_([p for p in head.parameters() if p.requires_grad], 1.0)
        train_optim.step()

        rec = {"step": step, "sup_loss": float(sup_loss.item())}
        log["phase2_steps"].append(rec)
        if step % 50 == 0 or step == args.train_steps - 1:
            print(f"  P2[{step:>4}/{args.train_steps}] dt={time.time()-t0:.1f}s "
                  f"sup={rec['sup_loss']:.4f}", flush=True)

        if (step + 1) % args.eval_every == 0:
            acc, tacc = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device)
            print(f"    >> P2 EVAL @ step {step+1}: exact_acc={acc:.4f}  "
                  f"(Δ from init: {acc-acc0:+.4f}, Δ from P1: {acc-acc_ref:+.4f})", flush=True)
            log["phase2_evals"].append({"step": global_step + step + 1, "acc": acc, "tok_acc": tacc})

    acc_f, tacc_f = evaluate(head, base, data_path, args.eval_n, args.eval_batch_size, device)
    print(f"\n=== Final eval ===")
    print(f"  initial: {acc0:.4f}  phase1_end: {log.get('phase1_final_acc', acc0):.4f}  "
          f"final: {acc_f:.4f}  (total Δ: {acc_f-acc0:+.4f})")
    log["final_acc"] = acc_f
    log["final_tok_acc"] = tacc_f
    log["phase2_evals"].append({"step": global_step + args.train_steps, "acc": acc_f, "tok_acc": tacc_f})

    Path(args.out).write_text(json.dumps(log, indent=2))
    print(f"log → {args.out}")


if __name__ == "__main__":
    main()