import json, math
from pathlib import Path
import numpy as np

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

# Analytic continuous-time latent flow: a circular orbit with fixed radius.
# The speed parameter makes clock-rate perturbations explicit.
def flow(z0, t, omega=1.0):
    z0 = np.asarray(z0, dtype=float)
    r = np.linalg.norm(z0)
    phase = math.atan2(z0[1], z0[0])
    a = phase + omega * np.asarray(t)
    return r * np.stack([np.cos(a), np.sin(a)], axis=-1)

def discrepancy(x, y, t, alpha, omega_x=1.0, omega_y=1.0):
    return np.linalg.norm(flow(x, t, omega_x) - flow(y, alpha * t, omega_y), axis=1).max()

def optimize_affine_warp(x, y, t, grid, omega_x=1.0, omega_y=1.0):
    vals = np.array([discrepancy(x, y, t, a, omega_x, omega_y) for a in grid])
    i = int(vals.argmin())
    return float(vals[i]), float(grid[i])

def make_point(r, phase):
    return np.array([r * np.cos(phase), r * np.sin(phase)])

def empirical_delta(ds, q=0.10):
    # Largest empirical threshold with approximately q false-closeness.
    return float(np.quantile(np.asarray(ds), q, method='linear'))

def main():
    # A moderate horizon avoids trivial saturation of a max distance at 2r.
    t = np.linspace(0, 2.0, 81)
    warp_grid = np.linspace(0.55, 1.05, 1001)
    rows = {}

    # Prediction 1: if y runs at omega_y, the optimal warp slope is
    # alpha*=omega_x/omega_y and removes clock mismatch when geometry agrees.
    mismatch = np.linspace(0.0, 0.35, 8)
    p1 = []
    x = make_point(1.0, 0.0)
    for dm in mismatch:
        omega_y = 1.0 + dm
        y = x.copy()
        ordinary = discrepancy(x, y, t, 1.0, omega_y=omega_y)
        warped, alpha = optimize_affine_warp(x, y, t, warp_grid, omega_y=omega_y)
        predicted_alpha = 1.0 / omega_y
        p1.append({'mismatch': float(dm), 'ordinary_D': float(ordinary),
                   'warped_D': warped, 'alpha_star': alpha,
                   'predicted_alpha': predicted_alpha})

    # Prediction 2: distinct invariant radii cannot be removed by a time warp;
    # D >= |r_x-r_y| and should increase with the radial gap.
    gaps = np.linspace(0.02, 0.40, 8)
    p2 = []
    for g in gaps:
        x = make_point(1.0, 0.2)
        y = make_point(1.0 + g, 0.2)
        D, alpha = optimize_affine_warp(x, y, t, warp_grid)
        p2.append({'radius_gap': float(g), 'D': D, 'lower_bound': float(g), 'alpha_star': alpha})

    # Prediction 3: L_exp=(max(0,m-D))^2 activates exactly for D<m;
    # therefore activation probability follows the empirical CDF and crosses
    # q near delta_emp(q).
    negatives = []
    for _ in range(300):
        x = make_point(rng.uniform(0.85, 1.15), rng.uniform(-math.pi, math.pi))
        y = make_point(rng.uniform(0.65, 1.35), rng.uniform(-math.pi, math.pi))
        negatives.append(optimize_affine_warp(x, y, t, warp_grid)[0])
    negatives = np.asarray(negatives)
    q = 0.10
    delta = empirical_delta(negatives, q)
    margins = np.array([0.05, 0.15, 0.25, 0.35, 0.50])
    activation = [float(np.mean(negatives < m)) for m in margins]

    # Tiny controlled comparison on clock-perturbed copies: ordinary pointwise
    # separation treats speed change as a difference, whereas the warp-aware D
    # largely removes it. This is the intended benefit, not a trained model.
    ordinary_clock, warped_clock = [], []
    for _ in range(100):
        r = rng.uniform(.9, 1.1)
        phase = rng.uniform(-math.pi, math.pi)
        dm = rng.uniform(0.0, .35)
        x = make_point(r, phase)
        y = x.copy()
        ordinary_clock.append(discrepancy(x, y, t, 1.0, omega_y=1.0 + dm))
        warped_clock.append(optimize_affine_warp(x, y, t, warp_grid, omega_y=1.0 + dm)[0])

    alpha_err = [abs(x['alpha_star'] - x['predicted_alpha']) for x in p1]
    result = {
      'seed': SEED,
      'prediction_1_clock_mismatch': p1,
      'prediction_2_radius_gap': p2,
      'prediction_3_margin_transition': {
          'delta_emp_q0.10': delta, 'margins': margins.tolist(),
          'activation_rates': activation, 'observed_rate_at_delta': float(np.mean(negatives < delta))},
      'mini_comparison': {
          'clock_pairs_ordinary_mean_D': float(np.mean(ordinary_clock)),
          'clock_pairs_warped_mean_D': float(np.mean(warped_clock)),
          'relative_reduction': float(1 - np.mean(warped_clock) / np.mean(ordinary_clock))},
      'checks': {
        'clock_warp_reduces_at_max_mismatch': bool(p1[-1]['warped_D'] < p1[-1]['ordinary_D']),
        'alpha_matches_inverse_speed_max_error': float(max(alpha_err)),
        'radius_lower_bound': bool(all(x['D'] + 1e-8 >= x['lower_bound'] for x in p2)),
        'radius_monotone': bool(np.all(np.diff([x['D'] for x in p2]) > 0)),
        'activation_monotone': bool(np.all(np.diff(activation) >= 0)),
        'delta_is_quantile': bool(abs(np.mean(negatives < delta) - q) < 0.02)}
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()