#!/usr/bin/env python3
"""Quick stress perturbation verification."""
import zmq, json, time, struct
import numpy as np

ctx = zmq.Context()
cp = ctx.socket(zmq.PUB)
cp.connect('tcp://127.0.0.1:5557')
asub = ctx.socket(zmq.SUB)
asub.connect('tcp://127.0.0.1:5559')
asub.setsockopt_string(zmq.SUBSCRIBE, '')
ssub = ctx.socket(zmq.SUB)
ssub.connect('tcp://127.0.0.1:5560')
ssub.setsockopt_string(zmq.SUBSCRIBE, '')
time.sleep(0.5)

# Drain
while True:
    try:
        ssub.recv(flags=zmq.NOBLOCK)
    except zmq.ZMQError:
        break
while True:
    try:
        asub.recv(flags=zmq.NOBLOCK)
    except zmq.ZMQError:
        break

# Baseline stress snapshot
cp.send_string(json.dumps({'cmd': 'stress_snapshot_now'}))
time.sleep(0.8)
d0 = ssub.recv()
sxy0 = np.frombuffer(d0, dtype=np.float32, count=1048576, offset=8 + 8 * 1048576)
sxy0 = sxy0.reshape(1024, 1024)
bx = np.mean(sxy0[507:518, 507:518])
print(f'Baseline sxy at [507:518,507:518]: {bx:.8e}')

# Apply perturb_stress
cp.send_string(json.dumps({'cmd': 'perturb_stress', 'x': 512, 'y': 512, 'sigma': 16, 'strength': 0.1}))
time.sleep(0.3)
acks = []
while True:
    try:
        acks.append(json.loads(asub.recv_string()))
    except zmq.ZMQError:
        break
for a in acks:
    s = json.dumps(a)
    print(f'ACK: {s[:150]}')

# Post stress snapshot
time.sleep(0.2)
cp.send_string(json.dumps({'cmd': 'stress_snapshot_now'}))
time.sleep(0.8)
d1 = ssub.recv()
sxy1 = np.frombuffer(d1, dtype=np.float32, count=1048576, offset=8 + 8 * 1048576)
sxy1 = sxy1.reshape(1024, 1024)
ax = np.mean(sxy1[507:518, 507:518])
delta = ax - bx
print(f'Post    sxy at [507:518,507:518]: {ax:.8e}')
print(f'Delta: {delta:.8e}')
print(f'Pre-committed threshold: |delta| > 1e-8')
passed = abs(delta) > 1e-8
print(f'VERIFICATION: {"PASS" if passed else "FAIL"}')
if passed:
    print('f_neq actually changed -- stress perturbation confirmed.')
else:
    print('f_neq did NOT change -- perturbation was equilibrated away before snapshot.')

cp.close()
asub.close()
ssub.close()
ctx.term()