#!/usr/bin/env python3
"""
Stress verification + density vs stress hold-time A/B test.
Pre-committed thresholds:
  - Stress change: |delta_sxy| > 1e-8 (must exceed zero by 8 orders)
  - Return threshold: L2 distance < 0.01 * baseline_variance (must drop below 1% of baseline noise)
  - Setup time: >=2510 cycles (>=10 khra periods at canonical w_khra=0.025, period=2pi/0.025≈251)
  - Bootstrap: 100 resamples, 95% CI
"""
import zmq, json, time, struct, numpy as np
from collections import deque

CTX = zmq.Context()

# ========== Helpers ==========
def send_cmd(cmd_pub, d):
    """Send JSON command to port 5557."""
    cmd_pub.send_string(json.dumps(d))
    time.sleep(0.05)

def recv_ack(ack_sub, timeout=4.0):
    """Receive ack messages until timeout."""
    ack_sub.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
    acks = []
    while True:
        try:
            acks.append(json.loads(ack_sub.recv_string()))
        except zmq.ZMQError:
            break
    ack_sub.setsockopt(zmq.RCVTIMEO, -1)  # back to blocking
    return acks

def recv_stress_snapshot(stress_sub, timeout=5.0):
    """Receive one stress snapshot from port 5560 (8-byte header + 3*4MB)."""
    stress_sub.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
    try:
        data = stress_sub.recv()
        stress_sub.setsockopt(zmq.RCVTIMEO, -1)
        cycle = struct.unpack_from('<I', data, 0)[0]
        w = struct.unpack_from('<H', data, 4)[0]
        h = struct.unpack_from('<H', data, 6)[0]
        n = w * h
        sxx = np.frombuffer(data, dtype=np.float32, count=n, offset=8)
        syy = np.frombuffer(data, dtype=np.float32, count=n, offset=8 + 4*n)
        sxy = np.frombuffer(data, dtype=np.float32, count=n, offset=8 + 8*n)
        return cycle, sxx.reshape(h, w), syy.reshape(h, w), sxy.reshape(h, w)
    except zmq.ZMQError:
        stress_sub.setsockopt(zmq.RCVTIMEO, -1)
        return None

def recv_coarse(coarse_sub, timeout=5.0, max_frames=500):
    """Receive coarse frames from port 5561 (16-byte KGCF header + 6144 floats).
    Returns list of (cycle, 32x32x6 array)."""
    coarse_sub.setsockopt(zmq.RCVTIMEO, int(timeout * 1000))
    frames = []
    while len(frames) < max_frames:
        try:
            data = coarse_sub.recv()
            if len(data) < 16 + 4*6144:
                continue
            # Parse KGCF header
            magic = data[:4]
            if magic != 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)
            frames.append((cycle, field))
        except zmq.ZMQError:
            break
    coarse_sub.setsockopt(zmq.RCVTIMEO, -1)
    return frames

# ========== Connect ==========
cmd_pub = CTX.socket(zmq.PUB)
cmd_pub.connect('tcp://127.0.0.1:5557')
time.sleep(0.3)

ack_sub = CTX.socket(zmq.SUB)
ack_sub.connect('tcp://127.0.0.1:5559')
ack_sub.setsockopt_string(zmq.SUBSCRIBE, '')

stress_sub = CTX.socket(zmq.SUB)
stress_sub.connect('tcp://127.0.0.1:5560')
stress_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, '')

# Drain any queued messages
time.sleep(2)
while True:
    try:
        coarse_sub.recv(flags=zmq.NOBLOCK)
    except zmq.ZMQError:
        break
while True:
    try:
        stress_sub.recv(flags=zmq.NOBLOCK)
    except zmq.ZMQError:
        break
while True:
    try:
        ack_sub.recv(flags=zmq.NOBLOCK)
    except zmq.ZMQError:
        break

# ================================
# PART 1: STRESS VERIFICATION
# ================================
print("=" * 60)
print("PART 1: STRESS PERTURBATION VERIFICATION")
print("=" * 60)

# Request stress snapshot before perturbation
print("Requesting baseline stress snapshot...")
send_cmd(cmd_pub, {"cmd": "stress_snapshot_now"})
time.sleep(1)
before = recv_stress_snapshot(stress_sub, timeout=5.0)
if before is None:
    print("FAIL: Could not receive baseline stress snapshot")
    exit(1)
bcycle, bsxx, bsyy, bsxy = before
print(f"Baseline stress snapshot at cycle {bcycle}")

# Compute stats at the injection site (patch around 512,512)
cx, cy = 512, 512
patch_r = 5  # 11x11 patch
y0, y1 = max(0, cy-patch_r), min(1024, cy+patch_r+1)
x0, x1 = max(0, cx-patch_r), min(1024, cx+patch_r+1)
before_sxy_mean = np.mean(bsxy[y0:y1, x0:x1])
before_sxy_std = np.std(bsxy[y0:y1, x0:x1])
print(f"Before perturb: sxy mean={before_sxy_mean:.6e}, std={before_sxy_std:.6e} at ({cx},{cy})")

# Apply perturb_stress
print("\nApplying perturb_stress...")
send_cmd(cmd_pub, {"cmd": "perturb_stress", "x": 512.0, "y": 512.0, "sigma": 16.0, "strength": 0.1})
time.sleep(1)
acks = recv_ack(ack_sub, timeout=3.0)
for a in acks:
    print(f"  ACK: {a}")

# Request stress snapshot after perturbation
print("Requesting post-perturb stress snapshot...")
send_cmd(cmd_pub, {"cmd": "stress_snapshot_now"})
time.sleep(1)
after = recv_stress_snapshot(stress_sub, timeout=5.0)
if after is None:
    print("FAIL: Could not receive post-perturb stress snapshot")
    exit(1)
acycle, asxx, asyy, asxy = after
print(f"Post-perturb stress snapshot at cycle {acycle}")

after_sxy_mean = np.mean(asxy[y0:y1, x0:x1])
after_sxy_std = np.std(asxy[y0:y1, x0:x1])
delta_sxy = after_sxy_mean - before_sxy_mean
threshold_sxy = 1e-8

print(f"\nAfter perturb:  sxy mean={after_sxy_mean:.6e}, std={after_sxy_std:.6e}")
print(f"Delta sxy mean: {delta_sxy:.6e}")
print(f"Pre-committed threshold: |delta| > {threshold_sxy:.1e}")
if abs(delta_sxy) > threshold_sxy:
    print("STRESS VERIFICATION: PASS -- f_neq actually changed!")
else:
    print("STRESS VERIFICATION: FAIL -- no detectable change in sxy")
    print("(This means the perturbation was equilibrated away before the snapshot)")

# ================================
# PART 2: HOLD-TIME A/B TEST
# ================================
print("\n" + "=" * 60)
print("PART 2: HOLD-TIME A/B -- DENSITY vs STRESS")
print("=" * 60)

SETTLE_CYCLES = 2600   # >=10 khra periods at default w_khra=0.025
MEASURE_CYCLES = 2600

# Drain coarse channel, let system settle
print(f"Waiting {SETTLE_CYCLES//100*10} sec for settle + baseline collection...")
time.sleep(SETTLE_CYCLES // 100 * 10 / 1000 + 5)  # rough wait

def capture_baseline(n_frames=300):
    """Capture baseline coarse frames for comparison."""
    frames = []
    while len(frames) < n_frames:
        try:
            data = coarse_sub.recv(flags=zmq.NOBLOCK)
            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)
                frames.append((cycle, field.astype(np.float64)))
        except zmq.ZMQError:
            time.sleep(0.02)
    return frames

def run_hold_time_test(perturb_cmd, label):
    """Run a hold-time measurement for a given perturbation.
    Returns (hold_time_cycles, bootstrap_ci_low, bootstrap_ci_high)."""
    print(f"\n--- {label} ---")
    
    # Capture baseline
    print("Capturing baseline...")
    baseline = capture_baseline(n_frames=250)
    if len(baseline) < 200:
        print(f"FAIL: only captured {len(baseline)} baseline frames")
        return None
    baseline_cycle = baseline[-1][0]
    print(f"Baseline: {len(baseline)} frames, last cycle {baseline_cycle}")
    
    # Compute baseline variance for threshold
    # Use the mean field over baseline frames
    baseline_fields = np.stack([f[1] for f in baseline])
    baseline_mean = np.mean(baseline_fields, axis=0)
    baseline_var = np.var(baseline_fields, axis=0)
    global_baseline_var = np.mean(baseline_var)
    threshold_d = 0.01 * global_baseline_var  # pre-committed: 1% of baseline noise
    print(f"Baseline mean variance: {global_baseline_var:.6e}, return threshold: {threshold_d:.6e}")
    
    # Apply perturbation
    print(f"Applying perturbation at cycle ~{baseline_cycle}...")
    send_cmd(cmd_pub, perturb_cmd)
    time.sleep(0.2)
    inject_cycle = baseline_cycle + 5  # estimate
    print(f"Estimated injection cycle: {inject_cycle}")
    
    # Collect post-perturbation frames (>=2510 cycles = ~260 frames at 10-cycle tick)
    # The coarse stream emits every 10 cycles, so 260 frames = 2600 cycles
    print(f"Collecting post-perturb frames (target ~260 frames)...")
    post_frames = capture_baseline(n_frames=280)
    print(f"Captured {len(post_frames)} post-perturb frames")
    
    if len(post_frames) < 50:
        print("FAIL: too few post-perturb frames")
        return None
    
    # Compute L2 distance from baseline mean at each frame
    distances = []
    cycles = []
    for cycle, field in post_frames:
        d = np.mean((field - baseline_mean) ** 2)
        distances.append(d)
        cycles.append(cycle)
    
    distances = np.array(distances)
    cycles = np.array(cycles)
    
    # Find hold-time: first frame where distance drops below threshold
    # Start from frame 20 (skip initial transient)
    hold_idx = None
    for i in range(20, len(distances)):
        if distances[i] < threshold_d:
            # Confirm: next 5 frames also below threshold
            if i + 5 < len(distances) and np.all(distances[i:i+5] < threshold_d):
                hold_idx = i
                break
    
    if hold_idx is not None:
        hold_cycle = cycles[hold_idx] - inject_cycle
        hold_cycles_float = float(hold_cycle)
    else:
        print("WARNING: never crossed threshold -- perturbation may persist beyond window")
        hold_cycle = len(post_frames) * 10  # rough max
        hold_cycles_float = float(hold_cycle)
    
    print(f"Hold-time (raw): {hold_cycles_float:.0f} cycles")
    print(f"Distance range: [{distances[0]:.6e}, {distances[-1]:.6e}]")
    
    # Bootstrap CI
    np.random.seed(42)
    boot_samples = []
    for _ in range(200):
        # Resample baseline frames with replacement
        idx = np.random.choice(len(baseline_fields), len(baseline_fields), replace=True)
        boot_baseline_mean = np.mean(baseline_fields[idx], axis=0)
        boot_baseline_var = np.mean(np.var(baseline_fields[idx], axis=0))
        boot_threshold = 0.01 * boot_baseline_var
        
        # Compute distances against boot baseline
        boot_distances = []
        for cycle, field in post_frames:
            d = np.mean((field - boot_baseline_mean) ** 2)
            boot_distances.append(d)
        boot_distances = np.array(boot_distances)
        
        # Find hold-time with boot threshold
        boot_hold_idx = None
        for i in range(20, len(boot_distances)):
            if boot_distances[i] < boot_threshold and i + 5 < len(boot_distances):
                if np.all(boot_distances[i:i+5] < boot_threshold):
                    boot_hold_idx = i
                    break
        if boot_hold_idx is not None:
            boot_samples.append(float(cycles[boot_hold_idx] - inject_cycle))
        else:
            boot_samples.append(float(len(post_frames) * 10))
    
    ci_low = np.percentile(boot_samples, 2.5)
    ci_high = np.percentile(boot_samples, 97.5)
    
    print(f"Hold-time: {hold_cycles_float:.0f} cycles, 95% CI: [{ci_low:.0f}, {ci_high:.0f}]")
    return hold_cycles_float, ci_low, ci_high

# Run Arm A: density
result_a = run_hold_time_test(
    {"cmd": "inject_density", "x": 512.0, "y": 512.0, "sigma": 16.0, "strength": 0.1},
    "Arm A: inject_density"
)

# Wait for system to settle back
print(f"\nWaiting {SETTLE_CYCLES//100*10/1000:.0f}s for field to settle...")
time.sleep(15)

# Run Arm B: stress
result_b = run_hold_time_test(
    {"cmd": "perturb_stress", "x": 512.0, "y": 512.0, "sigma": 16.0, "strength": 0.1},
    "Arm B: perturb_stress"
)

# ================================
print("\n" + "=" * 60)
print("RESULTS SUMMARY")
print("=" * 60)
print(f"\nSTRESS VERIFICATION: {'PASS' if abs(delta_sxy) > threshold_sxy else 'FAIL'}")
print(f"  Before sxy at site: {before_sxy_mean:.6e}")
print(f"  After  sxy at site: {after_sxy_mean:.6e}")
print(f"  Delta: {delta_sxy:.6e} (threshold: {threshold_sxy:.1e})")

print(f"\nHOLD-TIME A/B:")
if result_a:
    ht, lo, hi = result_a
    print(f"  Arm A (density):  {ht:.0f} cycles, 95% CI [{lo:.0f}, {hi:.0f}]")
else:
    print(f"  Arm A (density):  FAILED")

if result_b:
    ht, lo, hi = result_b
    print(f"  Arm B (stress):   {ht:.0f} cycles, 95% CI [{lo:.0f}, {hi:.0f}]")
else:
    print(f"  Arm B (stress):   FAILED")

# Cleanup
cmd_pub.close()
ack_sub.close()
stress_sub.close()
coarse_sub.close()
CTX.term()
print("\nDone.")