#!/usr/bin/env python3
"""Empirical NDB/MXC checks (evidence, not proofs).

This tool simulates a canonical noise-to-diffusion model on the circle:

  phi_{n+1} = phi_n + rho + xi_n  (mod 2pi)

with xi_n Gaussian(0, sigma^2). It measures the decay of low Fourier modes
from a non-uniform initial density as a practical diagnostic.

Outputs JSON to stdout.
"""

from __future__ import annotations

import argparse
import json
import math
from typing import List, Tuple

import numpy as np

TAU = 2.0 * math.pi


def simulate_random_walk(N: int, rho: float, sigma: float, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    phi = np.zeros(N, dtype=float)
    # non-uniform start: concentrated at 0
    phi[0] = 0.0
    noise = rng.normal(0.0, sigma, size=N - 1)
    for i in range(1, N):
        phi[i] = (phi[i - 1] + rho + noise[i - 1]) % TAU
    return phi


def mode_magnitudes_over_time(phi: np.ndarray, modes: List[int], window: int) -> dict:
    """Compute sliding-window estimates of |E exp(i k phi)|."""
    n = int(phi.shape[0])
    if window <= 10 or window >= n:
        raise ValueError("window must be between 10 and N-1")

    centers = []
    series = {int(k): [] for k in modes}
    for start in range(0, n - window + 1, window):
        seg = phi[start : start + window]
        centers.append(int(start + window // 2))
        for k in modes:
            val = float(abs(np.mean(np.exp(1j * k * seg))))
            series[int(k)].append(val)

    return {"centers": centers, "modes": series, "window": int(window)}


def fit_exponential_decay(x: List[int], y: List[float]) -> Tuple[float, float]:
    """Fit log(y) ~ a - lambda * t, return (lambda, r2) with robust guards."""
    xx = np.array(x, dtype=float)
    yy = np.array(y, dtype=float)
    # avoid zeros
    yy = np.clip(yy, 1e-12, None)
    ly = np.log(yy)

    A = np.vstack([np.ones_like(xx), -xx]).T
    coef, *_ = np.linalg.lstsq(A, ly, rcond=None)
    a, lam = float(coef[0]), float(coef[1])

    pred = A @ np.array([a, lam])
    ss_res = float(np.sum((ly - pred) ** 2))
    ss_tot = float(np.sum((ly - np.mean(ly)) ** 2))
    r2 = 1.0 if ss_tot == 0.0 else float(1.0 - ss_res / ss_tot)
    return lam, r2


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--N", type=int, default=8000)
    ap.add_argument("--rho_over_tau", type=float, default=0.137)
    ap.add_argument("--sigma", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--window", type=int, default=400)
    ap.add_argument("--modes", type=str, default="1,2,3")

    ap.add_argument("--min_fit_r2", type=float, default=0.85)
    ap.add_argument("--min_lambda", type=float, default=1e-4)

    args = ap.parse_args()

    modes = [int(x.strip()) for x in args.modes.split(",") if x.strip()]
    rho = float(args.rho_over_tau) * TAU

    phi = simulate_random_walk(args.N, rho, float(args.sigma), int(args.seed))
    mm = mode_magnitudes_over_time(phi, modes, int(args.window))

    # fit on mode 1 by default
    k0 = modes[0] if modes else 1
    y0 = mm["modes"][int(k0)]
    x0 = mm["centers"]
    lam, r2 = fit_exponential_decay(x0, y0)

    meets_threshold = bool((r2 >= float(args.min_fit_r2)) and (lam >= float(args.min_lambda)))

    out = {
        "kind": "NDB_MXC_EMPIRICAL_CHECKS",
        "mode": "random_walk_circle",
        "context": {
            "N": int(args.N),
            "rho_over_tau": float(args.rho_over_tau),
            "sigma": float(args.sigma),
            "seed": int(args.seed),
            "window": int(args.window),
            "modes": modes,
        },
        "measured": {
            "mode_magnitudes": mm,
            "fit": {
                "fit_mode": int(k0),
                "lambda_per_step": float(lam),
                "r2": float(r2),
            },
        },
        "thresholds": {
            "min_fit_r2": float(args.min_fit_r2),
            "min_lambda": float(args.min_lambda),
        },
        "conclusion": {
            "evidence_supports_diffusive_decay": meets_threshold,
            "note": "This is an empirical diagnostic for a declared NDB model. It is not a proof of MXC for a given physical system.",
        },
    }
    print(json.dumps(out, indent=2) + "\n")


if __name__ == "__main__":
    main()
