import json
import numpy as np
from chern_gap_monitor import monitor_field, qwz_field, fixed_point_parity


def main():
    np.random.seed(7)
    # Prediction 1: the spherical formula gives the exact octant area pi/2.
    e1, e2, e3 = np.eye(3)
    from chern_gap_monitor import spherical_area
    area = float(spherical_area(e1, e2, e3))

    K = 80
    masses = [-3.0, -1.5, -0.5, 0.5, 1.5, 3.0]
    rows = []
    for mass in masses:
        gap, ch = monitor_field(qwz_field(K, mass))
        rows.append({"mass": mass, "gap": gap, "chern": ch,
                     "rounded_chern": int(np.rint(ch)),
                     "parity_proxy": fixed_point_parity(mass)})

    # Prediction 2: sector is stable away from the three gap-closing masses;
    # prediction 3: at a sampled critical momentum the gap is |delta mass|.
    critical = []
    for mc in (-2.0, 0.0, 2.0):
        deltas = np.array([-0.20, -0.10, -0.05, 0.05, 0.10, 0.20])
        gaps = np.array([monitor_field(qwz_field(K, mc+d))[0] for d in deltas])
        slope = float(np.dot(np.abs(deltas), gaps) / np.dot(deltas, deltas))
        critical.append({"critical_mass": mc, "min_gap_at_transition": float(monitor_field(qwz_field(K, mc))[0]),
                         "deltas": deltas.tolist(), "gaps": gaps.tolist(),
                         "fit_slope_gap_vs_abs_delta": slope,
                         "predicted_slope": 1.0})

    # Finite-batch/EMA sanity: additive mean noise should be reduced by EMA.
    true_field = qwz_field(40, 1.0)
    rng = np.random.default_rng(11)
    raw_errors, ema_errors = [], []
    ema = np.zeros_like(true_field)
    alpha = 0.15
    for _ in range(100):
        observed = true_field + rng.normal(0, 0.20, true_field.shape)
        ema = (1-alpha)*ema + alpha*observed
        raw_errors.append(np.mean((observed-true_field)**2))
        ema_errors.append(np.mean((ema-true_field)**2))
    result = {
        "octant_area": area, "octant_area_error": abs(area-np.pi/2),
        "sector_rows": rows, "critical_sweeps": critical,
        "ema_noise_mse": float(np.mean(ema_errors[20:])),
        "raw_noise_mse": float(np.mean(raw_errors[20:])),
        "ema_reduction_fraction": float(1-np.mean(ema_errors[20:])/np.mean(raw_errors[20:]))
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()