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

This tool produces conservative diagnostics that are useful in benchmarks and
artifact evidence fields.

Supported modes
- --mode circle_mean_motion  : reproduces the mean-motion demo (circle)
- --mode torus_translation   : synthetic k-torus translation with nonresonant omega

Outputs JSON to stdout.
"""

from __future__ import annotations

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

import numpy as np

TAU = 2.0 * math.pi


def mod_tau(x: np.ndarray) -> np.ndarray:
    return np.mod(x, TAU)


def circ_increment(theta: np.ndarray) -> np.ndarray:
    """Return forward increments in (-pi, pi] for a phase series."""
    d = np.diff(theta)
    d = (d + math.pi) % TAU - math.pi
    return d


def best_rational_approx(x: float, qmax: int) -> dict:
    """Brute best p/q for x in [0,1) with q<=qmax."""
    best = (0, 1, abs(x))
    for q in range(1, qmax + 1):
        p = int(round(x * q))
        err = abs(x - p / q)
        if err < best[2]:
            best = (p, q, err)
    p, q, err = best
    return {"p": int(p), "q": int(q), "error": float(err)}


def low_mode_fourier_magnitudes(theta: np.ndarray, kmax: int) -> List[float]:
    mags: List[float] = []
    for k in range(1, kmax + 1):
        z = np.exp(1j * k * theta)
        mags.append(float(abs(np.mean(z))))
    return mags


def histogram_stats(theta: np.ndarray, bins: int) -> dict:
    counts, _ = np.histogram(theta, bins=bins, range=(0.0, TAU))
    n = int(np.sum(counts))
    expected = n / float(bins)
    # chi^2 statistic against uniform (no p-value without scipy)
    chi2 = float(np.sum((counts - expected) ** 2 / max(expected, 1e-12)))
    df = int(bins - 1)
    chi2_over_df = float(chi2 / max(df, 1))
    max_bin_dev = float(np.max(np.abs(counts - expected)) / max(expected, 1e-12))
    return {
        "bins": int(bins),
        "n": int(n),
        "chi2": chi2,
        "df": df,
        "chi2_over_df": chi2_over_df,
        "max_relative_bin_deviation": max_bin_dev,
    }


def joint_fourier_magnitudes(thetas: np.ndarray, mmax: int) -> dict:
    """For k-torus samples thetas shape (N,k), compute joint low-mode magnitudes.

    We compute |E exp(i m·theta)| for integer vectors m with ||m||_inf<=mmax,
    excluding the zero vector.
    """
    n, k = thetas.shape
    out = []
    # iterate all m vectors in [-mmax, mmax]^k
    # keep small list only
    def rec(idx: int, cur: List[int]):
        if idx == k:
            if all(v == 0 for v in cur):
                return
            m = np.array(cur, dtype=float)
            phase = np.dot(thetas, m)
            val = float(abs(np.mean(np.exp(1j * phase))))
            out.append({"m": [int(v) for v in cur], "magnitude": val})
            return
        for v in range(-mmax, mmax + 1):
            cur.append(v)
            rec(idx + 1, cur)
            cur.pop()

    rec(0, [])
    # sort by magnitude descending for readability
    out.sort(key=lambda d: d["magnitude"], reverse=True)
    return {"mmax": int(mmax), "count": int(len(out)), "modes": out}


def small_integer_relation_scan(omega: np.ndarray, lmax: int) -> dict:
    """Search for a small integer relation a·omega ≈ 0 mod 2pi.

    This is a conservative resonance screen: if we find an a with small residual,
    we report it.
    """
    k = int(omega.shape[0])
    best = {"a": [0] * k, "residual_over_tau": 1.0}

    def rec(idx: int, cur: List[int]):
        nonlocal best
        if idx == k:
            if all(v == 0 for v in cur):
                return
            a = np.array(cur, dtype=float)
            r = float(np.dot(a, omega)) / TAU
            # residual to nearest integer
            residual = abs(r - round(r))
            if residual < best["residual_over_tau"]:
                best = {"a": [int(v) for v in cur], "residual_over_tau": float(residual)}
            return
        for v in range(-lmax, lmax + 1):
            cur.append(v)
            rec(idx + 1, cur)
            cur.pop()

    rec(0, [])
    best["lmax"] = int(lmax)
    return best


@dataclass
class Thresholds:
    qmax: int = 200
    nonresonance_min_error: float = 1e-4
    kmax: int = 10
    max_low_mode_mag: float = 0.03
    bins: int = 60
    max_chi2_over_df: float = 2.0
    max_bin_dev: float = 0.35
    joint_mmax: int = 2
    relation_lmax: int = 6
    relation_min_residual: float = 1e-4


def run_circle_mean_motion(N: int) -> Tuple[np.ndarray, dict]:
    D = 86400.0
    omega_E = TAU / 86400.0
    omega_sun = TAU / (365.2422 * D)
    omega_moon = TAU / (27.32166 * D)

    n = np.arange(N, dtype=float)
    t = n * D
    alpha_sun = mod_tau(omega_sun * t)
    alpha_moon = mod_tau(omega_moon * t)
    phi = mod_tau(alpha_moon - alpha_sun)

    rho = float(mod_tau((omega_moon - omega_sun) * D))
    context = {
        "model": "mean_motion_demo",
        "N": int(N),
        "step_seconds": float(D),
        "rho_radians": float(rho),
        "rho_over_tau": float(rho / TAU),
    }
    return phi, context


def run_torus_translation(N: int, k: int) -> Tuple[np.ndarray, dict]:
    # deterministic nonresonant omega using square roots
    base = [math.sqrt(2.0), math.sqrt(3.0), math.sqrt(5.0), math.sqrt(7.0), math.sqrt(11.0)]
    if k > len(base):
        raise ValueError("k too large for built-in omega generator")
    omega = np.array(base[:k], dtype=float)
    omega = (omega / np.max(omega)) * (0.37 * TAU)

    n = np.arange(N, dtype=float)
    thetas = (n[:, None] * omega[None, :]) % TAU

    context = {
        "model": "torus_translation",
        "N": int(N),
        "k": int(k),
        "omega_radians": [float(x) for x in omega.tolist()],
        "omega_over_tau": [float(x / TAU) for x in omega.tolist()],
    }
    return thetas, context


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", required=True, choices=["circle_mean_motion", "torus_translation"])
    ap.add_argument("--N", type=int, default=4000)
    ap.add_argument("--k", type=int, default=2)

    ap.add_argument("--qmax", type=int, default=200)
    ap.add_argument("--nonres_min_error", type=float, default=1e-4)
    ap.add_argument("--kmax", type=int, default=10)
    ap.add_argument("--max_low_mode_mag", type=float, default=0.03)
    ap.add_argument("--bins", type=int, default=60)
    ap.add_argument("--max_chi2_over_df", type=float, default=2.0)
    ap.add_argument("--max_bin_dev", type=float, default=0.35)

    ap.add_argument("--joint_mmax", type=int, default=2)
    ap.add_argument("--relation_lmax", type=int, default=6)
    ap.add_argument("--relation_min_residual", type=float, default=1e-4)

    args = ap.parse_args()

    thr = Thresholds(
        qmax=args.qmax,
        nonresonance_min_error=args.nonres_min_error,
        kmax=args.kmax,
        max_low_mode_mag=args.max_low_mode_mag,
        bins=args.bins,
        max_chi2_over_df=args.max_chi2_over_df,
        max_bin_dev=args.max_bin_dev,
        joint_mmax=args.joint_mmax,
        relation_lmax=args.relation_lmax,
        relation_min_residual=args.relation_min_residual,
    )

    if args.mode == "circle_mean_motion":
        theta, context = run_circle_mean_motion(args.N)
        x = float(context["rho_over_tau"]) % 1.0
        rat = best_rational_approx(x, thr.qmax)
        nonresonance_ok = rat["error"] >= thr.nonresonance_min_error

        fourier = low_mode_fourier_magnitudes(theta, thr.kmax)
        low_mode_ok = max(fourier) <= thr.max_low_mode_mag

        hist = histogram_stats(theta, thr.bins)
        hist_ok = (hist["chi2_over_df"] <= thr.max_chi2_over_df) and (hist["max_relative_bin_deviation"] <= thr.max_bin_dev)

        out = {
            "kind": "RNF_EMPIRICAL_CHECKS",
            "mode": args.mode,
            "context": context,
            "checks": {
                "nonresonance": {
                    "qmax": thr.qmax,
                    "best_rational_approx": rat,
                    "min_error_threshold": thr.nonresonance_min_error,
                    "meets_threshold": bool(nonresonance_ok),
                },
                "low_mode_uniformity": {
                    "kmax": thr.kmax,
                    "mode_magnitudes": fourier,
                    "max_allowed": thr.max_low_mode_mag,
                    "meets_threshold": bool(low_mode_ok),
                },
                "histogram_uniformity": {
                    **hist,
                    "max_chi2_over_df": thr.max_chi2_over_df,
                    "max_relative_bin_deviation_allowed": thr.max_bin_dev,
                    "meets_threshold": bool(hist_ok),
                },
            },
            "conclusion": {
                "evidence_supports_rnf_equidistribution": bool(nonresonance_ok and low_mode_ok and hist_ok),
                "note": "These are empirical diagnostics used as evidence in benchmarks. They are not proofs of RNF.",
            },
        }
        print(json.dumps(out, indent=2) + "\n")
        return

    # torus translation
    thetas, context = run_torus_translation(args.N, args.k)
    omega_over_tau = np.array(context["omega_over_tau"], dtype=float)

    # per-coordinate rational approximations
    rat_list = [best_rational_approx(float(x % 1.0), thr.qmax) for x in omega_over_tau.tolist()]
    nonresonance_ok = all(r["error"] >= thr.nonresonance_min_error for r in rat_list)

    # small relation scan
    rel = small_integer_relation_scan(np.array(context["omega_radians"], dtype=float), thr.relation_lmax)
    relation_ok = rel["residual_over_tau"] >= thr.relation_min_residual

    # per-coordinate low mode
    per_coord = []
    coord_ok = True
    for j in range(thetas.shape[1]):
        mags = low_mode_fourier_magnitudes(thetas[:, j], thr.kmax)
        per_coord.append({"coord": int(j), "mode_magnitudes": mags})
        coord_ok = coord_ok and (max(mags) <= thr.max_low_mode_mag)

    # joint modes
    joint = joint_fourier_magnitudes(thetas, thr.joint_mmax)
    joint_max = float(max(m["magnitude"] for m in joint["modes"]))
    joint_ok = joint_max <= thr.max_low_mode_mag

    out = {
        "kind": "RNF_EMPIRICAL_CHECKS",
        "mode": args.mode,
        "context": context,
        "checks": {
            "nonresonance_per_coordinate": {
                "qmax": thr.qmax,
                "best_rational_approx": rat_list,
                "min_error_threshold": thr.nonresonance_min_error,
                "meets_threshold": bool(nonresonance_ok),
            },
            "small_integer_relation_scan": {
                **rel,
                "min_residual_threshold": thr.relation_min_residual,
                "meets_threshold": bool(relation_ok),
            },
            "low_mode_uniformity_per_coordinate": {
                "kmax": thr.kmax,
                "max_allowed": thr.max_low_mode_mag,
                "coords": per_coord,
                "meets_threshold": bool(coord_ok),
            },
            "joint_low_mode_uniformity": {
                "mmax": thr.joint_mmax,
                "max_joint_mode_magnitude": joint_max,
                "max_allowed": thr.max_low_mode_mag,
                "top_joint_modes": joint["modes"][: min(20, len(joint["modes"]))],
                "meets_threshold": bool(joint_ok),
            },
        },
        "conclusion": {
            "evidence_supports_rnf_equidistribution": bool(nonresonance_ok and relation_ok and coord_ok and joint_ok),
            "note": "These are empirical diagnostics used as evidence in benchmarks. They are not proofs of RNF.",
        },
    }
    print(json.dumps(out, indent=2) + "\n")


if __name__ == "__main__":
    main()
