#!/usr/bin/env python3
"""
Phase 4 Runner & Validator: Zero-Proxy Layer-47 Dual-Orthogonal Whitened Fisher Discriminant (WFLD)
====================================================================================================
Implements and verifies Phase 4 v3 (`PLAN.md` §2.2, §2.4, §2.6, §5, and `audits/opus_strata_metal_opts_v3.md` §5):
1. Keeps `blk.0..46.ffn_down_shexp.weight` (47 tensors) 100% bit-identical to Phase 3 (`K256-q2_0-healed-metal.gguf`).
2. Solves the exact zero-proxy Layer-47 Dual-Orthogonal Whitened Fisher Linear Discriminant (`scripts/solve_phase4_wfld_l47.cpp`)
   on `blk.47.ffn_down_shexp.weight` (`Q8_0` type 8) over:
   - 42 natural multi-turn assistant slop contexts across all 7 banned phrases (28 train + 14 strictly held-out = 33.3% held-out),
     strictly enforcing prefix token matching (`tuple[:-1]`) with zero unprefixed token penalties.
   - 207 legitimate collateral occurrences of `" than"`, `" what"`, `" end"`, `" away"`, `" version"`, `" actually"` across 48 unique
     natural sentences (`data/m5_val_prompts/phase4_collateral_1024tok.txt`: 65 train_anchor + 142 heldout_eval with 0 8-gram overlap).
   - 1,022 real multi-domain calibration tokens (`models/K256_reap_l2_sum/cal_cache/cal_p3_1024`).
"""

import argparse
import json
import os
import struct
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from scripts.phase3_lora_and_transplant import classify_qwen4exp_tensor
from scripts.prepare_phase3_sft_data import apply_speech_act_surgery
from scripts.slice_qwen4exp_gguf import compute_tensor_nbytes, parse_gguf_header


def run_self_test(root: Path) -> None:
    meta_path = root / "data" / "m5_val_prompts" / "phase4_ftpo_sft_meta.json"
    col_path  = root / "data" / "m5_val_prompts" / "phase4_collateral_1024tok.txt"
    ftpo_path = root / "data" / "m5_val_prompts" / "phase4_ftpo_sft_1024tok.txt"
    val_path  = root / "results" / "phase4_colab_validation.json"
    bin_path  = root / "results" / "phase4_healed_persona_ftpo_writers.bin"
    cpp_path  = root / "scripts" / "solve_phase4_wfld_l47.cpp"

    for p in (meta_path, col_path, ftpo_path, val_path, bin_path, cpp_path):
        assert p.exists(), f"Missing required Phase 4 artifact: {p}"

    meta = json.loads(meta_path.read_text(encoding="utf-8"))
    b_json = json.loads((root / "results" / "banned_llm_isms.json").read_text(encoding="utf-8"))
    banned_tuples = {tuple(x) for x in b_json["ftpo_qwen_token_id_tuples"]}

    ftpo_toks = [int(x) for x in ftpo_path.read_text(encoding="utf-8").split()]
    col_toks  = [int(x) for x in col_path.read_text(encoding="utf-8").split()]
    assert len(ftpo_toks) == 1024, f"Expected 1024 ftpo tokens, got {len(ftpo_toks)}"
    assert len(col_toks) >= 1600, f"Expected >= 1600 collateral tokens, got {len(col_toks)}"

    # Verify strict prefix match on all 42 FTPO positions (28 train + 14 heldout = 33.3% heldout)
    ftpo_pos = meta["ftpo_positions"]
    assert len(ftpo_pos) == 42, f"Expected 42 FTPO positions, got {len(ftpo_pos)}"
    n_tr = sum(1 for x in ftpo_pos if x["split"] == "train")
    n_ho = sum(1 for x in ftpo_pos if x["split"] == "heldout")
    assert n_tr == 28 and n_ho == 14, f"Expected 28 train + 14 heldout, got {n_tr} + {n_ho}"
    for x in ftpo_pos:
        pos = x["t_pos"]
        tup = tuple(x["tuple"])
        assert tup in banned_tuples, f"Unknown tuple {tup}"
        plen = len(tup) - 1
        assert tuple(ftpo_toks[pos + 1 - plen : pos + 2]) == tup, f"Prefix mismatch at {pos}"

    # Verify >= 200 collateral positions, 0 banned prefixes, and 0 8-gram overlap between train_anchor and heldout_eval
    col_pos = meta["collateral_positions"]
    assert len(col_pos) >= 200, f"Expected >= 200 collateral positions, got {len(col_pos)}"
    train_8g = {tuple(col_toks[i : i + 8]) for i in range(0, 480 - 8)}
    held_8g  = {tuple(col_toks[i : i + 8]) for i in range(480, len(col_toks) - 8)}
    assert len(train_8g & held_8g) == 0, "8-gram overlap detected in collateral trace!"

    # Verify bundle format and that 47/48 tensors are bit-identical to Phase 3 GGUF
    healed_gguf = root / "models" / "K256_reap_l2_sum" / "K256-q2_0-healed-metal.gguf"
    p3_writers = root / "results" / "phase3_healed_writers_1024.bin"
    if healed_gguf.exists():
        fd = os.open(str(healed_gguf), os.O_RDONLY)
        try:
            hdr = parse_gguf_header(fd, healed_gguf.stat().st_size)
            tmap = {t["name"]: (t["ggml_type"], compute_tensor_nbytes(t["dims"], t["ggml_type"]), t["abs_start"]) for t in hdr["tensors"]}
            del hdr
            p3_map = {}
            if p3_writers.exists():
                with open(p3_writers, "rb") as pf:
                    (pnt,) = struct.unpack("<I", pf.read(4))
                    for _ in range(pnt):
                        (pnl,) = struct.unpack("<I", pf.read(4))
                        pnm = pf.read(pnl).decode("utf-8")
                        pgt, pnb = struct.unpack("<II", pf.read(8))
                        p3_map[pnm] = pf.read(pnb)
            with open(bin_path, "rb") as f:
                (nt,) = struct.unpack("<I", f.read(4))
                assert nt == 48
                for l in range(48):
                    (nl,) = struct.unpack("<I", f.read(4))
                    nm = f.read(nl).decode("utf-8")
                    gt, nb = struct.unpack("<II", f.read(8))
                    raw = f.read(nb)
                    assert classify_qwen4exp_tensor(nm, scope="swift_writers_plus_router") == "trainable_dense_router_shared"
                    eg, enb, eoff = tmap[nm]
                    assert gt == eg and nb == enb
                    deployed = os.pread(fd, enb, eoff)
                    if p3_map:
                        if l < 47:
                            assert raw == p3_map[nm], f"Layer {l} modified vs Phase 3!"
                        else:
                            assert raw != p3_map[nm], "Layer 47 not updated vs Phase 3!"
                    assert raw == deployed or l == 47
        finally:
            os.close(fd)

    # Verify JSON bars (a)-(g)
    vj = json.loads(val_path.read_text(encoding="utf-8"))
    ho_sep = vj["ftpo_natural_sentences_summary"]["heldout_separate_doc_split"]
    assert ho_sep["count"] == 21
    assert ho_sep["prob_reduction_pct"] >= 69.0
    assert ho_sep["after_mean_prob_im_end"] <= ho_sep["before_mean_prob_im_end"]
    fg = vj["special_token_and_chat_doc_stability_bars_f_g"]
    assert fg["bar_f_argmax_flips_to_special_tokens"]["heldout_chat_doc_total_flips"] == 0
    assert fg["bar_f_argmax_flips_to_special_tokens"]["sft_chat_doc_total_flips"] == 0
    assert fg["bar_g_heldout_chat_doc_stability"]["ordinary_token_mean_delta_logprob_nats"] >= -0.01
    assert fg["bar_g_heldout_chat_doc_stability"]["nonslop_kld_nats"] <= 0.005
    col_ho = vj["collateral_legitimate_target_benchmark"]["heldout_eval_split"]
    assert col_ho["loss_gt_0_1_pct"] <= 10.0, f"loss_gt_0_1_pct {col_ho['loss_gt_0_1_pct']} > 10.0"
    for w, wstats in col_ho["per_word_breakdown"].items():
        assert wstats["mean_delta_logprob"] >= -0.05, f"Word {w!r} mean {wstats['mean_delta_logprob']} < -0.05" 
    re_v = vj["real_strata_metal_engine_verification"]["kld_vs_phase3_healed_m0"]
    for dom in ("ops_tools", "chat_json", "ja"):
        assert re_v[dom]["mean_kld"] <= 0.01, f"{dom} KLD {re_v[dom]['mean_kld']} > 0.01"


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--self-test", action="store_true")
    args = ap.parse_args()
    root = Path(__file__).resolve().parent.parent
    if args.self_test:
        run_self_test(root)
        print("scripts/run_phase4_colab_ftpo.py self-test: PASS")