#!/usr/bin/env python3
"""Toy verification of a contractive Floquet return map.

State is (theta, r), where theta is phase and r is transverse deviation from
an ideal periodic orbit.  One return period applies
    r' = q*r + delta
    theta' = theta + shear*r'  (mod 2*pi)
The orbit distance is |r|; pointwise Euclidean distance on the unit circle is
computed against the phase-aligned target point.
"""
import json
import math
from pathlib import Path
import numpy as np

SEED = 2094
rng = np.random.default_rng(SEED)


def step(z, q, defect=0.0, shear=0.0):
    theta, r = z
    rn = q * r + defect
    return np.array([theta + shear * rn, rn], dtype=float)


def simulate(q, r0, n, defect=0.0, shear=0.0):
    z = np.array([0.0, r0], dtype=float)
    zs = [z.copy()]
    for _ in range(n):
        z = step(z, q, defect, shear)
        zs.append(z.copy())
    return np.asarray(zs)


def orbital_error(z):
    return np.abs(z[:, 1])


def pointwise_error(z):
    # Target orbit Gamma(theta)=(cos(theta), sin(theta)); learned point uses
    # the model phase and radius 1+r. This intentionally retains phase drift.
    theta, r = z[:, 0], z[:, 1]
    target = np.stack([np.cos(theta), np.sin(theta)], axis=1)
    model = np.stack([(1.0 + r) * np.cos(theta), (1.0 + r) * np.sin(theta)], axis=1)
    # Compare to phase-zero target: pointwise error exposes phase mismatch.
    target0 = np.array([1.0, 0.0])
    return np.linalg.norm(model - target0[None, :], axis=1)


def predicted_r(q, r0, n, delta):
    if abs(1.0 - q) < 1e-12:
        return r0 + n * delta
    return q**n * r0 + delta * (1.0 - q**n) / (1.0 - q)


def empirical_q(q, n_pairs=1000):
    # For the affine map, this should equal |q| exactly, independently of
    # phase shear when transverse distance is measured in the return coordinate.
    rs1 = rng.normal(0, 0.2, n_pairs)
    rs2 = rng.normal(0, 0.2, n_pairs)
    d0 = np.abs(rs1 - rs2)
    d1 = np.abs((q * rs1) - (q * rs2))
    return float(np.max(d1 / d0))


def main():
    report = {"seed": SEED, "predictions": {}, "comparison": {}}

    # Prediction 1: transition at q=1; contraction/expansion rate is log(q).
    qs = [0.80, 0.95, 1.00, 1.05, 1.20]
    n = 30
    r0 = 0.1
    threshold_rows = []
    for q in qs:
        z = simulate(q, r0, n)
        observed = abs(z[-1, 1] / r0)
        expected = abs(q) ** n
        slope = math.log(observed) / n if observed > 0 else float("-inf")
        threshold_rows.append({"q": q, "observed_gain": observed,
                               "predicted_gain": expected,
                               "observed_log_rate": slope,
                               "predicted_log_rate": math.log(abs(q))})
    report["predictions"]["threshold_q_equals_1"] = threshold_rows

    # Prediction 2: constant defect gives floor delta/(1-q), linearly in delta.
    q = 0.8
    deltas = [0.001, 0.003, 0.01, 0.03]
    n = 200
    defect_rows = []
    for delta in deltas:
        z = simulate(q, 0.0, n, defect=delta)
        observed = abs(z[-1, 1])
        expected = delta / (1 - q)
        defect_rows.append({"delta": delta, "observed_floor": observed,
                            "predicted_floor": expected,
                            "relative_error": abs(observed-expected)/expected})
    report["predictions"]["defect_floor"] = defect_rows

    # Prediction 3: phase shear makes pointwise error grow, while orbital error
    # follows q^n and remains bounded. Compare shear=0 and a sheared orbit.
    q, r0, shear, n = 0.85, 0.15, 2.5, 80
    z = simulate(q, r0, n, shear=shear)
    orb = orbital_error(z)
    point = pointwise_error(z)
    report["predictions"]["phase_invariant_vs_pointwise"] = {
        "q": q, "shear": shear, "initial_transverse": r0,
        "orbital_error_at_80": float(orb[-1]),
        "orbital_predicted_at_80": float(abs(q**n * r0)),
        "pointwise_error_at_80": float(point[-1]),
        "pointwise_peak": float(np.max(point)),
        "orbital_max": float(np.max(orb)),
        "interpretation": "phase shear leaves transverse contraction intact but increases phase-sensitive error"
    }

    # Secondary mini-experiment: unconstrained return map versus contractive map.
    # Small random per-cycle model defect makes this a deployment-like test.
    cycles = 1000
    trials = 200
    sigma = 0.004
    configs = {"baseline_unconstrained": 1.02, "contractive_idea": 0.80}
    for name, q in configs.items():
        orbital_end, point_end, qhat = [], [], []
        for _ in range(trials):
            z = np.array([0.0, 0.10])
            for _ in range(cycles):
                d = float(rng.normal(0.0, sigma))
                z = step(z, q, defect=d, shear=1.5)
            orbital_end.append(abs(z[1]))
            point_end.append(float(pointwise_error(z[None, :])[0]))
            qhat.append(empirical_q(q, 100))
        report["comparison"][name] = {
            "q": q, "cycles": cycles, "defect_sigma": sigma,
            "median_orbital_error": float(np.median(orbital_end)),
            "p95_orbital_error": float(np.quantile(orbital_end, .95)),
            "median_pointwise_error": float(np.median(point_end)),
            "fraction_orbital_inside_0.05": float(np.mean(np.asarray(orbital_end) < .05)),
            "qhat_max_over_trials": float(np.max(qhat))
        }

    # Overall numerical checks used to decide whether the mechanism manifested.
    checks = {
        "threshold_relative_errors": [abs(x["observed_gain"]-x["predicted_gain"])/max(x["predicted_gain"],1e-12) for x in threshold_rows],
        "max_defect_relative_error": max(x["relative_error"] for x in defect_rows),
        "empirical_q_contractive": report["comparison"]["contractive_idea"]["qhat_max_over_trials"] < 0.8000001,
        "empirical_q_baseline_noncontractive": report["comparison"]["baseline_unconstrained"]["qhat_max_over_trials"] > 1.0,
    }
    report["checks"] = checks
    out = Path("results.json")
    out.write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()