#!/usr/bin/env python3
"""
Differential hold-time A/B: perturbed vs unperturbed CONTROL from same checkpoint.
Save to "." (daemon working dir), find checkpoint by glob.
Cancels mass-leak drift — both runs share the same leak, difference isolates perturbation.

Mandatory sanity gate: Arm A (density, dead channel) MUST decay faster than Arm B (stress).
If density == stress, the measurement is still broken — STOP.
"""
import zmq, json, time, struct, glob, os
import numpy as np

# Daemon's working dir is /mnt/d/resonance-engine-active
CKPT_DIR = "/mnt/d/resonance-engine-active"

def main():
    ctx = zmq.Context()

    cp = ctx.socket(zmq.PUB); cp.connect("tcp://127.0.0.1:5557")
    ack_sub = ctx.socket(zmq.SUB); ack_sub.connect("tcp://127.0.0.1:5559")
    ack_sub.setsockopt_string(zmq.SUBSCRIBE, "")
    coarse_sub = ctx.socket(zmq.SUB); coarse_sub.connect("tcp://127.0.0.1:5561")
    coarse_sub.setsockopt_string(zmq.SUBSCRIBE, "")
    time.sleep(1.0)

    # Drain
    for s in [coarse_sub, ack_sub]:
        while True:
            try: s.recv(flags=zmq.NOBLOCK)
            except zmq.ZMQError: break

    def send_cmd(d):
        cp.send_string(json.dumps(d))
        time.sleep(0.05)

    def wait_ack(cmd_name, timeout=8.0):
        ack_sub.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
        start = time.time()
        while time.time() - start < timeout:
            try:
                ack = json.loads(ack_sub.recv_string())
                if ack.get("ack") == cmd_name:
                    ack_sub.setsockopt(zmq.RCVTIMEO, -1)
                    return ack
            except zmq.ZMQError:
                pass
        ack_sub.setsockopt(zmq.RCVTIMEO, -1)
        return None

    def drain_coarse():
        while True:
            try: coarse_sub.recv(flags=zmq.NOBLOCK)
            except zmq.ZMQError: break

    def collect_coarse(n_frames):
        coarse_sub.setsockopt(zmq.RCVTIMEO, 90000)
        frames = []
        while len(frames) < n_frames:
            try:
                data = coarse_sub.recv()
                if len(data) >= 16 + 4 * 6144 and data[:4] == b"KGCF":
                    cycle = struct.unpack_from("<I", data, 4)[0]
                    vals = np.frombuffer(data, dtype=np.float32, count=6144, offset=16)
                    field = vals.reshape(32, 32, 6).astype(np.float64)
                    frames.append((cycle, field))
            except zmq.ZMQError:
                break
        coarse_sub.setsockopt(zmq.RCVTIMEO, -1)
        return frames

    # ========== SETTLE ==========
    print(f"Settling for 2600 cycles...")
    drain_coarse()
    time.sleep(2)
    skip = collect_coarse(260)
    print(f"Settle complete. Last cycle: {skip[-1][0]}")

    # ========== SAVE CHECKPOINT ==========
    # Daemon treats path as a directory; save to "." = daemon's cwd = CKPT_DIR
    # Remove any old test checkpoint written by this script
    for f in glob.glob(f"{CKPT_DIR}/ckpt_*_test.bin"):
        os.remove(f)

    send_cmd({"cmd": "save_state", "path": "."})
    ack = wait_ack("save_state", timeout=12.0)
    if ack is None or ack.get("status") != "ok":
        print(f"FAIL: save_state failed. ACK: {ack}")
        exit(1)
    save_cycle = ack.get("cycle", 0)
    print(f"Checkpoint saved at cycle {save_cycle}")

    # Find the actual checkpoint file
    time.sleep(0.5)
    ckpt_files = sorted(glob.glob(f"{CKPT_DIR}/ckpt_*.bin"))
    # Filter to the one just created (closest to save_cycle)
    ckpt_path = None
    for f in ckpt_files:
        # Parse cycle from filename
        basename = os.path.basename(f)
        parts = basename.replace(".bin", "").split("_c")
        if len(parts) == 2:
            try:
                fcycle = int(parts[1])
                if abs(fcycle - save_cycle) < 100:  # within 100 cycles
                    ckpt_path = basename
                    break
            except ValueError:
                pass
    if ckpt_path is None:
        # Fallback: use latest
        ckpt_path = os.path.basename(ckpt_files[-1])
    print(f"Using checkpoint: {ckpt_path}")

    # ========== CONTROL RUN ==========
    print("\n=== CONTROL RUN (no perturbation) ===")
    drain_coarse()
    time.sleep(1)
    control_frames = collect_coarse(260)
    print(f"Control: {len(control_frames)} frames, cycles {control_frames[0][0]} - {control_frames[-1][0]}")

    # ========== LOAD CHECKPOINT ==========
    print(f"\nLoading checkpoint for Arm A...")
    send_cmd({"cmd": "load_state", "path": ckpt_path})
    ack = wait_ack("load_state", timeout=15.0)
    if ack is None or ack.get("status") != "ok":
        print(f"FAIL: load_state failed. ACK: {ack}")
        exit(1)
    drain_coarse()
    time.sleep(2)
    skip = collect_coarse(60)
    print(f"Post-load settle: last cycle {skip[-1][0]}")

    # ========== PERTURBED RUN (Arm A: density) ==========
    print("\n=== ARM A: inject_density ===")
    drain_coarse()
    time.sleep(0.5)
    send_cmd({"cmd": "inject_density", "x": 512.0, "y": 512.0, "sigma": 16.0, "strength": 0.1})
    time.sleep(0.5)
    perturb_a = collect_coarse(260)
    print(f"Arm A: {len(perturb_a)} frames")

    # ========== LOAD CHECKPOINT again for Arm B ==========
    print(f"\nLoading checkpoint for Arm B...")
    send_cmd({"cmd": "load_state", "path": ckpt_path})
    ack = wait_ack("load_state", timeout=15.0)
    if ack is None or ack.get("status") != "ok":
        print(f"FAIL: load_state failed for Arm B")
        exit(1)
    drain_coarse()
    time.sleep(2)
    skip = collect_coarse(60)
    print(f"Post-load settle: last cycle {skip[-1][0]}")

    # ========== PERTURBED RUN (Arm B: stress) ==========
    print("\n=== ARM B: perturb_stress ===")
    drain_coarse()
    time.sleep(0.5)
    send_cmd({"cmd": "perturb_stress", "x": 512.0, "y": 512.0, "sigma": 16.0, "strength": 0.1})
    time.sleep(0.5)
    perturb_b = collect_coarse(260)
    print(f"Arm B: {len(perturb_b)} frames")

    # ========== ANALYSIS ==========
    print("\n" + "=" * 60)
    print("DIFFERENTIAL HOLD-TIME ANALYSIS")
    print("=" * 60)

    n = min(len(control_frames), len(perturb_a), len(perturb_b))
    control_arrays = np.stack([f[1] for f in control_frames[:n]])
    diffs = np.diff(control_arrays, axis=0)
    baseline_noise = np.mean(diffs ** 2)
    threshold = 0.05 * baseline_noise
    print(f"Baseline noise: {baseline_noise:.6e}")
    print(f"Return threshold (5% of baseline noise): {threshold:.6e}")

    def compute_hold_time(perturb_frames, label):
        distances = []
        for i in range(n):
            d = np.mean((perturb_frames[i][1] - control_frames[i][1]) ** 2)
            distances.append(d)
        distances = np.array(distances)

        print(f"\n{label}:")
        print(f"  Distance range: [{distances[0]:.6e}, {distances[-1]:.6e}]")
        print(f"  Frames above threshold: {np.sum(distances >= threshold)}/{n}")

        hold_idx = None
        for i in range(10, n - 5):
            if distances[i] < threshold and np.all(distances[i:i + 5] < threshold):
                hold_idx = i
                break

        if hold_idx is not None:
            hold_cycles = hold_idx * 10
            print(f"  Returned at frame {hold_idx} (~{hold_cycles} cycles)")
        else:
            hold_cycles = n * 10
            print(f"  Never returned within window ({hold_cycles} cycles)")

        # Bootstrap CI
        np.random.seed(42)
        boot_samples = []
        for _ in range(100):
            idx = np.random.choice(n, n, replace=True)
            boot_d = np.array([np.mean((perturb_frames[j][1] - control_frames[j][1]) ** 2) for j in idx if j < n])
            boot_hold = None
            for j in range(10, n - 5):
                if boot_d[j] < threshold and np.all(boot_d[j:j + 5] < threshold):
                    boot_hold = j
                    break
            boot_samples.append(float(boot_hold * 10) if boot_hold is not None else float(n * 10))

        ci_low = np.percentile(boot_samples, 2.5)
        ci_high = np.percentile(boot_samples, 97.5)
        print(f"  Hold-time: {hold_cycles} cycles, 95% CI: [{ci_low:.0f}, {ci_high:.0f}]")
        return hold_cycles, ci_low, ci_high

    ht_a, lo_a, hi_a = compute_hold_time(perturb_a, "Arm A: inject_density")
    ht_b, lo_b, hi_b = compute_hold_time(perturb_b, "Arm B: perturb_stress")

    # ========== SANITY GATE ==========
    print("\n" + "=" * 60)
    print("SANITY GATE: Does density decay faster than stress?")
    print("=" * 60)
    print(f"  Arm A (density): {ht_a} cycles, CI [{lo_a:.0f}, {hi_a:.0f}]")
    print(f"  Arm B (stress):  {ht_b} cycles, CI [{lo_b:.0f}, {hi_b:.0f}]")

    if ht_a < ht_b:
        print(f"  GATE PASS: density ({ht_a}) < stress ({ht_b}) — instrument is working.")
    elif ht_a > ht_b:
        print(f"  ANOMALOUS: density ({ht_a}) > stress ({ht_b}) — unexpected, investigate.")
    else:
        print(f"  GATE FAIL: density ({ht_a}) == stress ({ht_b}) — instrument STILL broken.")
        print(f"  STOP. Do NOT proceed to sweep engine until this is fixed.")

    cp.close(); ack_sub.close(); coarse_sub.close(); ctx.term()
    print("\nDone.")

if __name__ == "__main__":
    main()