#!/usr/bin/env python3
"""
TASK: gixx_amp=0 Discriminator — full experiment.
Obey TASK_GIXX_ZERO_DISCRIMINATOR.md protocol exactly.
"""
import zmq, json, time, struct, numpy as np, os, sys
from datetime import datetime, timezone

ANALYSIS_DIR = "/mnt/d/resonance-engine-active/analysis"
HANDOFF_FILE = "/mnt/d/resonance-engine-active/TASK_GIXX_ZERO_DISCRIMINATOR.md"
CHRONICLE_FILE = "/mnt/d/resonance-engine-active/analysis/observer_chronicle.jsonl"

def chronicle(kind, data, note=""):
    """Append to chronicle JSONL."""
    entry = {
        "ts": datetime.now(timezone.utc).isoformat(),
        "actor": "cline-agent",
        "kind": kind,
        "data": data,
        "note": note
    }
    with open(CHRONICLE_FILE, "a", encoding="utf-8") as f:
        f.write(json.dumps(entry) + "\n")
    print(f"  [CHRONICLE] {kind}: {note}")

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

    # --- Sockets ---
    # Telemetry SUB (port 5556)
    tel_sub = ctx.socket(zmq.SUB)
    tel_sub.connect("tcp://127.0.0.1:5556")
    tel_sub.setsockopt_string(zmq.SUBSCRIBE, "")

    # Command PUB (port 5557)
    cmd_pub = ctx.socket(zmq.PUB)
    cmd_pub.connect("tcp://127.0.0.1:5557")

    # Ack SUB (port 5559)
    ack_sub = ctx.socket(zmq.SUB)
    ack_sub.connect("tcp://127.0.0.1:5559")
    ack_sub.setsockopt_string(zmq.SUBSCRIBE, "")

    # Coarse SUB (port 5561)
    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)  # Allow ZMQ connections to establish

    # --- Helpers ---
    def drain_socket(sock, timeout=0.5):
        sock.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
        while True:
            try:
                sock.recv(flags=zmq.NOBLOCK)
            except zmq.ZMQError:
                break
        sock.setsockopt(zmq.RCVTIMEO, -1)

    def recv_telemetry(timeout=5.0):
        """Receive one telemetry JSON frame. Returns dict or None."""
        tel_sub.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
        try:
            data = tel_sub.recv_string()
            tel_sub.setsockopt(zmq.RCVTIMEO, -1)
            return json.loads(data)
        except zmq.ZMQError:
            tel_sub.setsockopt(zmq.RCVTIMEO, -1)
            return None

    def recv_coarse_frames(n_frames, timeout=120.0):
        """Receive n_frames from coarse stream (5561). Returns list of (cycle, 32x32x6 array)."""
        coarse_sub.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
        frames = []
        while len(frames) < n_frames:
            try:
                data = coarse_sub.recv()
                if len(data) < 16 + 4 * 6144:
                    continue
                if data[:4] != b"KGCF":
                    continue
                cycle = struct.unpack_from("<I", data, 4)[0]
                tiles = struct.unpack_from("<H", data, 8)[0]
                ch = struct.unpack_from("<H", data, 10)[0]
                if tiles != 32 or ch != 6:
                    continue
                vals = np.frombuffer(data, dtype=np.float32, count=6144, offset=16)
                field = vals.reshape(32, 32, 6)  # tile-major, channels per tile
                frames.append((cycle, field))
            except zmq.ZMQError:
                break
        coarse_sub.setsockopt(zmq.RCVTIMEO, -1)
        return frames

    def send_cmd_verified(cmd_dict, timeout=8.0):
        """Send command, return (accepted, ack_dict). accepted=True if status=ok."""
        cmd_pub.send_string(json.dumps(cmd_dict))
        time.sleep(0.3)
        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_dict.get("cmd", ""):
                    ack_sub.setsockopt(zmq.RCVTIMEO, -1)
                    status = ack.get("status", "ok")
                    return (status == "ok", ack)
            except zmq.ZMQError:
                pass
        ack_sub.setsockopt(zmq.RCVTIMEO, -1)
        return (False, None)

    def save_npz(fname, cycles, frames, note=""):
        """Save npz: cycles array + frames array (n_frames, 32, 32, 6)."""
        path = os.path.join(ANALYSIS_DIR, fname)
        np.savez_compressed(path, cycles=np.array(cycles, dtype=np.int32),
                           frames=np.stack(frames).astype(np.float32))
        print(f"  Saved {path}: {len(cycles)} frames, cycles {cycles[0]}-{cycles[-1]}")
        if note:
            chronicle("capture", {"file": fname, "n_frames": len(cycles),
                       "cycle_range": [int(cycles[0]), int(cycles[-1])]}, note)

    # Drain all
    drain_socket(tel_sub)
    drain_socket(ack_sub)
    drain_socket(coarse_sub)

    # ==========================================
    # STEP 1: Verify telemetry + canonical params
    # ==========================================
    print("=" * 60)
    print("STEP 1: Verify daemon live + canonical params")
    print("=" * 60)
    tel = None
    for _ in range(5):
        tel = recv_telemetry(timeout=5.0)
        if tel:
            break
        time.sleep(1)
    if tel is None:
        print("FATAL: No telemetry received.")
        chronicle("fatal", {}, "No telemetry — daemon may not be running")
        return 1

    print(f"  Cycle: {tel['cycle']}, coherence: {tel['coherence']:.4f}, alpha: {tel.get('alpha','?')}")
    gixx_amp_val = tel.get("gixx_amp", None)
    if gixx_amp_val is None:
        print("FATAL: gixx_amp not in telemetry.")
        return 1
    print(f"  gixx_amp: {gixx_amp_val} (expect 0.008)")
    if abs(gixx_amp_val - 0.008) > 0.001:
        print(f"WARNING: gixx_amp={gixx_amp_val} not canonical. Proceeding anyway — handoff says compiled-in defaults.")
    chronicle("telemetry_check", {"cycle": tel["cycle"], "coherence": tel["coherence"],
               "gixx_amp": gixx_amp_val}, "Daemon live, params verified")

    # ==========================================
    # STEP 2: Settle ≥ 2,000 cycles
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 2: Settle to cycle ≥ 2,000")
    print("=" * 60)
    current_cycle = tel["cycle"]
    if current_cycle < 2000:
        wait_cycles = 2000 - current_cycle
        wait_sec = wait_cycles / 70.0 + 2  # ~70 cycles/s
        print(f"  At cycle {current_cycle}, waiting ~{wait_sec:.0f}s to reach 2000...")
        time.sleep(wait_sec)
    # Confirm
    tel = recv_telemetry(timeout=5.0)
    print(f"  Settled at cycle {tel['cycle']}, coh={tel['coherence']:.4f}")
    chronicle("settle", {"cycle": tel["cycle"]}, "Settled ≥ 2,000 cycles")

    # ==========================================
    # STEP 3: Baseline capture B0 (~200 frames)
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 3: Baseline capture B0 (~200 frames)")
    print("=" * 60)
    drain_socket(coarse_sub)
    time.sleep(1)
    b0_frames = recv_coarse_frames(200, timeout=60.0)
    if len(b0_frames) < 180:
        print(f"FATAL: only {len(b0_frames)} baseline frames")
        return 1
    b0_cycles = [f[0] for f in b0_frames]
    b0_fields = [f[1] for f in b0_frames]
    save_npz("gixx0_B0.npz", b0_cycles, b0_fields, "Baseline at canonical gixx_amp=0.008")

    # ==========================================
    # STEP 4: Set gixx_amp = 0.0
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 4: Set gixx_amp = 0.0")
    print("=" * 60)
    accepted, ack = send_cmd_verified({"cmd": "set_param", "param": "gixx_amp", "value": 0.0}, timeout=8.0)
    if not accepted:
        print(f"FATAL: set_param gixx_amp=0.0 failed or rejected. ACK: {ack}")
        chronicle("set_param_failed", {"param": "gixx_amp", "value": 0.0, "ack": ack}, "Set gixx_amp=0.0 FAILED")
        return 1
    set_cycle = ack.get("cycle", 0)
    print(f"  gixx_amp=0.0 accepted at cycle {set_cycle}")
    chronicle("set_param", {"param": "gixx_amp", "value": 0.0, "cycle": set_cycle, "ack": ack}, "gixx_amp set to 0.0")

    # ==========================================
    # STEP 5: Hold ≥10,000 cycles, capture W1/W2/W3
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 5: Hold ≥10,000 cycles, capture W1, W2, W3")
    print("=" * 60)

    # W1: +500 cycles after set
    wait_sec = 500 / 70.0 + 2
    print(f"  Waiting {wait_sec:.0f}s for W1 (+500 cycles)...")
    time.sleep(wait_sec)
    drain_socket(coarse_sub)
    time.sleep(0.5)
    w1_frames = recv_coarse_frames(150, timeout=60.0)
    if len(w1_frames) < 100:
        print(f"WARNING: only {len(w1_frames)} W1 frames")
    else:
        w1c = [f[0] for f in w1_frames]
        w1f = [f[1] for f in w1_frames]
        save_npz("gixx0_W1.npz", w1c, w1f, f"W1: +500 after gixx_amp=0.0, cycles {w1c[0]}-{w1c[-1]}")
        # Health check
        tel = recv_telemetry(timeout=5.0)
        if tel:
            print(f"  W1 health: cycle={tel['cycle']} coh={tel['coherence']:.4f} asym={tel['asymmetry']:.2f}")

    # W2: +5,000 cycles after set
    elapsed = (w1_frames[-1][0] if w1_frames else set_cycle + 500) - set_cycle
    wait_more = max(0, (5000 - elapsed) / 70.0 + 2)
    print(f"\n  Waiting {wait_more:.0f}s for W2 (+5,000 cycles)...")
    time.sleep(wait_more)
    drain_socket(coarse_sub)
    time.sleep(0.5)
    w2_frames = recv_coarse_frames(150, timeout=60.0)
    if len(w2_frames) < 100:
        print(f"WARNING: only {len(w2_frames)} W2 frames")
    else:
        w2c = [f[0] for f in w2_frames]
        w2f = [f[1] for f in w2_frames]
        save_npz("gixx0_W2.npz", w2c, w2f, f"W2: +5,000 after gixx_amp=0.0, cycles {w2c[0]}-{w2c[-1]}")
        tel = recv_telemetry(timeout=5.0)
        if tel:
            print(f"  W2 health: cycle={tel['cycle']} coh={tel['coherence']:.4f} asym={tel['asymmetry']:.2f}")

    # W3: +9,000 cycles after set
    elapsed = (w2_frames[-1][0] if w2_frames else set_cycle + 5000) - set_cycle
    wait_more = max(0, (9000 - elapsed) / 70.0 + 2)
    print(f"\n  Waiting {wait_more:.0f}s for W3 (+9,000 cycles)...")
    time.sleep(wait_more)
    drain_socket(coarse_sub)
    time.sleep(0.5)
    w3_frames = recv_coarse_frames(150, timeout=60.0)
    if len(w3_frames) < 100:
        print(f"WARNING: only {len(w3_frames)} W3 frames")
    else:
        w3c = [f[0] for f in w3_frames]
        w3f = [f[1] for f in w3_frames]
        save_npz("gixx0_W3.npz", w3c, w3f, f"W3: +9,000 after gixx_amp=0.0, cycles {w3c[0]}-{w3c[-1]}")
        tel = recv_telemetry(timeout=5.0)
        if tel:
            print(f"  W3 health: cycle={tel['cycle']} coh={tel['coherence']:.4f} asym={tel['asymmetry']:.2f}")

    chronicle("hold_complete", {"set_cycle": set_cycle, "final_cycle": w3c[-1] if w3_frames else 0},
              "≥10,000-cycle hold at gixx_amp=0.0 complete")

    # ==========================================
    # STEP 6: Restore gixx_amp = 0.008, capture recovery
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 6: Restore gixx_amp = 0.008")
    print("=" * 60)
    accepted, ack = send_cmd_verified({"cmd": "set_param", "param": "gixx_amp", "value": 0.008}, timeout=8.0)
    if not accepted:
        print(f"FATAL: restore gixx_amp=0.008 failed. ACK: {ack}")
        chronicle("set_param_failed", {"param": "gixx_amp", "value": 0.008, "ack": ack}, "Restore FAILED")
        return 1
    print(f"  gixx_amp=0.008 restored at cycle {ack.get('cycle',0)}")
    chronicle("set_param", {"param": "gixx_amp", "value": 0.008, "cycle": ack.get("cycle",0), "ack": ack}, "gixx_amp restored")

    # Recovery window
    print("  Collecting recovery window (~100 frames)...")
    time.sleep(2)
    drain_socket(coarse_sub)
    time.sleep(1)
    r_frames = recv_coarse_frames(100, timeout=60.0)
    if len(r_frames) >= 80:
        rc = [f[0] for f in r_frames]
        rf = [f[1] for f in r_frames]
        save_npz("gixx0_R.npz", rc, rf, f"Recovery after restore gixx_amp=0.008, cycles {rc[0]}-{rc[-1]}")
    tel = recv_telemetry(timeout=5.0)
    if tel:
        print(f"  Recovery health: cycle={tel['cycle']} coh={tel['coherence']:.4f}")
        if tel['coherence'] > 0.7:
            print("  Coherence recovering toward canonical band.")

    # ==========================================
    # STEP 7: Spectral analysis
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 7: Spectral analysis (periodogram, Hann window)")
    print("=" * 60)

    def spectral_analysis(frames, label, channel=4):
        """
        channel=4 is sxy (index 5 if 0-based: rho=0, ux=1, uy=2, sxx=3, syy=4, sxy=5)
        Wait - the spec says [rho, ux, uy, sxx, syy, sxy] per tile.
        Channel 4 = syy, channel 5 = sxy. Let me use channel 5 for sxy.
        """
        if len(frames) < 20:
            return {"error": "too few frames"}

        fields = np.stack(frames)  # (n_frames, 32, 32, 6)
        # Global spatial mean of |sxy|
        sxy = fields[:, :, :, 5]  # channel 5 = sxy
        signal = np.mean(np.abs(sxy), axis=(1, 2))  # (n_frames,)

        # Also ux for secondary
        ux = fields[:, :, :, 1]  # channel 1 = ux
        ux_signal = np.mean(np.abs(ux), axis=(1, 2))

        n = len(signal)
        # Frame spacing = 10 cycles
        # Periodogram with Hann window
        window = np.hanning(n)
        signal_dt = signal - np.mean(signal)
        ux_dt = ux_signal - np.mean(ux_signal)

        # Compute FFT
        fft = np.fft.rfft(signal_dt * window)
        fft_ux = np.fft.rfft(ux_dt * window)

        power = np.abs(fft) ** 2
        power_ux = np.abs(fft_ux) ** 2

        # Frequency bins in cycles per FRAME. Convert to cycles: period = (n * 10) / k
        # Actually: bin k corresponds to period = (n * frame_spacing) / k
        # frame_spacing = 10 cycles
        freqs = np.fft.rfftfreq(n, d=1.0)  # cycles/frame
        periods = np.where(np.arange(len(power)) > 0,
                          n * 10.0 / np.arange(1, len(power) + 1), np.inf)

        # Top-4 spectral peaks (excluding DC)
        # Sort by power, skip DC (bin 0)
        sorted_idx = np.argsort(power[1:])[::-1] + 1
        top4_sxy = [(int(periods[i]), float(power[i])) for i in sorted_idx[:4] if i < len(periods)]
        sorted_idx_ux = np.argsort(power_ux[1:])[::-1] + 1
        top4_ux = [(int(periods[i]), float(power_ux[i])) for i in sorted_idx_ux[:4] if i < len(periods)]

        # Fraction of total AC power in 100-150 cycle band
        ac_total = np.sum(power[1:])
        band_mask = (periods >= 100) & (periods <= 150)
        # Apply mask skipping DC
        valid_mask = np.zeros(len(power), dtype=bool)
        valid_mask[1:] = band_mask[1:]
        band_power = np.sum(power[valid_mask])
        band_fraction = band_power / ac_total if ac_total > 0 else 0.0

        # top-2 spectral peak in 90-150?
        top2_in_band = any(90 <= p <= 150 for p, _ in top4_sxy[:2])

        # Also 57-63 band for ux
        ux_ac_total = np.sum(power_ux[1:])
        ux_band_57_63 = np.sum(power_ux[(periods >= 57) & (periods <= 63)])
        ux_57_63_frac = ux_band_57_63 / ux_ac_total if ux_ac_total > 0 else 0.0

        return {
            "label": label,
            "n_frames": n,
            "top4_sxy": top4_sxy,
            "top4_ux": top4_ux,
            "band_fraction_100_150": band_fraction,
            "top2_in_90_150": top2_in_band,
            "ac_total": float(ac_total),
            "band_power": float(band_power),
            "ux_57_63_frac": ux_57_63_frac,
        }

    # Analyze B0 and W3 (the decision windows)
    b0_analysis = spectral_analysis(b0_fields, "B0 (baseline, gixx_amp=0.008)")
    w3_analysis = spectral_analysis(w3f if w3_frames else b0_fields, "W3 (gixx_amp=0.0, +9k cycles)")

    print(f"\n  B0 (baseline):")
    print(f"    Top-4 sxy periods: {b0_analysis.get('top4_sxy', [])}")
    print(f"    100-150 band fraction: {b0_analysis['band_fraction_100_150']:.4f}")
    print(f"    Top-2 in 90-150: {b0_analysis['top2_in_90_150']}")

    print(f"\n  W3 (gixx_amp=0.0):")
    print(f"    Top-4 sxy periods: {w3_analysis.get('top4_sxy', [])}")
    print(f"    100-150 band fraction: {w3_analysis['band_fraction_100_150']:.4f}")
    print(f"    Top-2 in 90-150: {w3_analysis['top2_in_90_150']}")

    # Secondary: ux 57-63 band
    print(f"\n  B0 ux 57-63 band frac: {b0_analysis['ux_57_63_frac']:.4f}")
    print(f"  W3 ux 57-63 band frac: {w3_analysis['ux_57_63_frac']:.4f}")

    # ==========================================
    # STEP 8: Decision rule
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 8: Pre-committed decision rule")
    print("=" * 60)

    b0_band_frac = b0_analysis["band_fraction_100_150"]
    w3_band_frac = w3_analysis["band_fraction_100_150"]
    w3_top2_in_band = w3_analysis["top2_in_90_150"]

    verdict = "AMBIGUOUS"
    verdict_reason = "none of the gates fired"

    if w3_band_frac >= 0.5 * b0_band_frac and w3_top2_in_band:
        verdict = "SURVIVES"
        verdict_reason = f"W3 band-fraction ({w3_band_frac:.4f}) ≥ 0.5 × B0 ({0.5*b0_band_frac:.4f}) AND top-2 peak in 90-150"
    elif w3_band_frac < 0.1 * b0_band_frac:
        verdict = "DIES"
        verdict_reason = f"W3 band-fraction ({w3_band_frac:.4f}) < 0.1 × B0 ({0.1*b0_band_frac:.4f})"

    print(f"  B0 100-150 band fraction: {b0_band_frac:.4f}")
    print(f"  W3 100-150 band fraction: {w3_band_frac:.4f}")
    print(f"  W3 top-2 in 90-150: {w3_top2_in_band}")
    print(f"\n  VERDICT: **{verdict}**")
    print(f"  Reason: {verdict_reason}")

    chronicle("verdict", {
        "b0_band_frac": b0_band_frac,
        "w3_band_frac": w3_band_frac,
        "w3_top2_in_90_150": w3_top2_in_band,
        "verdict": verdict,
        "reason": verdict_reason
    }, f"gixx_amp=0 discriminator: {verdict}")

    # ==========================================
    # STEP 9: Write RESULTS to handoff file
    # ==========================================
    print("\n" + "=" * 60)
    print("STEP 9: Writing RESULTS to handoff file")
    print("=" * 60)

    results_text = f"""
## RESULTS (auto-written by cline-agent, {datetime.now(timezone.utc).isoformat()})

### Execution
- Daemon launched, canonical params confirmed (gixx_amp=0.008)
- Baseline B0: {len(b0_frames)} frames, cycles {b0_cycles[0]}-{b0_cycles[-1]}
- gixx_amp=0.0 set at cycle {set_cycle}, ack-verified
- W1 (+500): {len(w1_frames) if w1_frames else 0} frames
- W2 (+5k): {len(w2_frames) if w2_frames else 0} frames
- W3 (+9k): {len(w3_frames) if w3_frames else 0} frames, cycles {w3c[0] if w3_frames else '?'}-{w3c[-1] if w3_frames else '?'}
- gixx_amp=0.008 restored, ack-verified
- Recovery R: {len(r_frames) if r_frames else 0} frames

### Spectral Analysis

**B0 (baseline, gixx_amp=0.008):**
- Top-4 sxy periods (cycles): {b0_analysis.get('top4_sxy', [])}
- 100-150 cycle band fraction: {b0_band_frac:.4f}
- ux 57-63 band fraction: {b0_analysis['ux_57_63_frac']:.4f}

**W3 (gixx_amp=0.0, +9,000 cycles):**
- Top-4 sxy periods (cycles): {w3_analysis.get('top4_sxy', [])}
- 100-150 cycle band fraction: {w3_band_frac:.4f}
- Top-2 in 90-150: {w3_top2_in_band}
- ux 57-63 band fraction: {w3_analysis['ux_57_63_frac']:.4f}

### Verdict

**`{verdict}`** — {verdict_reason}

- SURVIVES gate: band-fraction ≥ 0.5×B0 (need {0.5*b0_band_frac:.4f}, got {w3_band_frac:.4f}) AND top-2 in 90-150 (need True, got {w3_top2_in_band})
- DIES gate: band-fraction < 0.1×B0 (need <{0.1*b0_band_frac:.4f}, got {w3_band_frac:.4f})

### Recovery
- Coherence after restore: {tel['coherence']:.4f} (canonical band ~0.737-0.740)

### Artifacts
- `analysis/gixx0_B0.npz` — baseline
- `analysis/gixx0_W1.npz` — +500 cycles after gixx_amp=0.0
- `analysis/gixx0_W2.npz` — +5,000 cycles
- `analysis/gixx0_W3.npz` — +9,000 cycles
- `analysis/gixx0_R.npz` — recovery after restore
"""

    with open(HANDOFF_FILE, "a", encoding="utf-8") as f:
        f.write(results_text)
    print("  RESULTS appended to TASK_GIXX_ZERO_DISCRIMINATOR.md")

    # Cleanup
    cmd_pub.close()
    tel_sub.close()
    ack_sub.close()
    coarse_sub.close()
    ctx.term()

    print("\n" + "=" * 60)
    print(f"EXPERIMENT COMPLETE. Verdict: {verdict}")
    print("=" * 60)
    return 0

if __name__ == "__main__":
    sys.exit(main())