#!/usr/bin/env python3
"""Lattice dispersion relation on the RS cubic ledger lattice.

This mirrors the Lean module
IndisputableMonolith/Foundation/SimplicialLedger/LorentzEmergence.lean,
whose dispersion is the standard cos-Laplacian on Z^3:

    omega^2(k) = (2/a^2) * sum_{i=1}^{3} [1 - cos(k_i * a)]

where a is the voxel spacing.  Lean proves (0 sorry, THEOREM tier):
  - dispersion_nonneg:  omega^2 >= 0
  - dispersion_upper_bound_by_isotropic:  omega^2 <= |k|^2
  - isotropic_envelope_rotation_invariant:  the continuum envelope |k|^2
    is rotationally invariant (Lorentz emergence at leading order)

At finite a the lattice corrections are SUBLUMINAL:

    omega^2 = |k|^2 - (a^2/12) * sum k_i^4 + (a^4/360) * sum k_i^6 - ...

so the group velocity deficit is delta = |v_g/c - 1| ~ (k a)^2 / 8: QUADRATIC
in energy over the lattice scale.  This order-2 suppression is what the Lean
carries (Physics/LorentzViolationBoundFromRS.lean: lvOrderOfMagnitude = 2,
a DEF/MODEL, no numeric coefficient).  There is NO Lean theorem that massless
modes travel at exactly c at finite k; earlier versions of this script cited
a "No-Dispersion Theorem" that does not exist in the library.  The bounds
comparison below is therefore a MODEL estimate scored against published
experimental limits, not a theorem.

This script computes:
  - Exact and Taylor-expanded omega(k) (cos form, matching the Lean)
  - Group velocity v_g and phase velocity v_p
  - Lorentz violation parameter delta = |v_g/c - 1|
  - A MODEL-tier comparison to experimental LV bounds (quadratic suppression
    evaluated from the exact lattice v_g at k*a = E/E_Planck)
"""

from __future__ import annotations

import argparse
import json
import math
from typing import Any, Dict, List, Optional, Tuple


PHI = (1.0 + math.sqrt(5.0)) / 2.0
PLANCK_LENGTH_M = 1.616255e-35
PLANCK_ENERGY_GEV = 1.220890e19
PLANCK_ENERGY_EV = PLANCK_ENERGY_GEV * 1e9

EXPERIMENTAL_BOUNDS: List[Dict[str, Any]] = [
    {
        "name": "Fermi-LAT GRB 090510 (photon)",
        "sector": "massless",
        "bound_linear": 1e-20,
        "bound_quadratic": None,
        "energy_GeV": 31.0,
        "reference": "Nature 462, 331 (2009)",
    },
    {
        "name": "MAGIC Mrk 501 (photon)",
        "sector": "massless",
        "bound_linear": 3e-19,
        "bound_quadratic": None,
        "energy_GeV": 1.2e3,
        "reference": "Phys. Lett. B 668, 253 (2008)",
    },
    {
        "name": "Hughes-Drever (electron)",
        "sector": "massive",
        "bound_linear": 1e-25,
        "bound_quadratic": None,
        "energy_GeV": 0.511e-3,
        "reference": "Phys. Rev. Lett. 4, 342 (1960)",
    },
    {
        "name": "IceCube (neutrino)",
        "sector": "massive",
        "bound_linear": 1e-23,
        "bound_quadratic": None,
        "energy_GeV": 100.0,
        "reference": "Nature Phys. 14, 961 (2018)",
    },
    {
        "name": "Clock comparison (proton)",
        "sector": "massive",
        "bound_linear": 1e-27,
        "bound_quadratic": None,
        "energy_GeV": 0.938,
        "reference": "Phys. Rev. Lett. 85, 5038 (2000)",
    },
]


def jcost(x: float) -> float:
    """J(x) = (x + 1/x)/2 - 1, the unique RS cost function."""
    return 0.5 * (x + 1.0 / x) - 1.0


def omega_squared_exact(kx: float, ky: float, kz: float, a: float) -> float:
    """Exact lattice dispersion: (2/a^2) * sum_i [1 - cos(k_i * a)].

    Matches LorentzEmergence.axis_dispersion / dispersion exactly.
    """
    s = 0.0
    for ki in (kx, ky, kz):
        s += 1.0 - math.cos(ki * a)
    return 2.0 * s / (a * a)


def omega_squared_taylor(kx: float, ky: float, kz: float, a: float,
                         order: int = 4) -> Tuple[float, Dict[str, float]]:
    """Taylor expansion of omega^2 in powers of a^2.

    Returns (omega^2_approx, dict of correction terms by order).
    """
    k_components = [kx, ky, kz]
    k2 = sum(ki ** 2 for ki in k_components)
    k4 = sum(ki ** 4 for ki in k_components)
    k6 = sum(ki ** 6 for ki in k_components)
    k8 = sum(ki ** 8 for ki in k_components)

    terms = {"O(0)_continuum": k2}
    result = k2

    # cos expansion: 2(1 - cos x)/a^2 = k^2 - a^2 k^4/12 + a^4 k^6/360 - ...
    if order >= 2:
        c2 = -(a ** 2 / 12.0) * k4
        terms["O(a^2)_quartic"] = c2
        result += c2

    if order >= 4:
        c4 = (a ** 4 / 360.0) * k6
        terms["O(a^4)_sextic"] = c4
        result += c4

    if order >= 6:
        c6 = -(a ** 6 / 20160.0) * k8
        terms["O(a^6)_octic"] = c6
        result += c6

    return result, terms


def group_velocity_exact(kx: float, ky: float, kz: float, a: float,
                         direction: Optional[Tuple[float, float, float]] = None
                         ) -> float:
    """Group velocity |d omega / dk| along a direction (default: k-hat).

    From omega^2 = (2/a^2) sum (1 - cos(k_i a)):  d(omega^2)/dk_i = (2/a) sin(k_i a),
    so v_g_i = sin(k_i * a) / (a * omega).
    """
    om2 = omega_squared_exact(kx, ky, kz, a)
    if om2 <= 0.0:
        return 0.0
    omega = math.sqrt(om2)

    grad = [math.sin(ki * a) / a for ki in (kx, ky, kz)]

    if direction is None:
        k_mag = math.sqrt(kx ** 2 + ky ** 2 + kz ** 2)
        if k_mag < 1e-30:
            return 0.0
        direction = (kx / k_mag, ky / k_mag, kz / k_mag)

    vg_dot = sum(g * d for g, d in zip(grad, direction))
    return abs(vg_dot) / omega


def phase_velocity_exact(kx: float, ky: float, kz: float, a: float) -> float:
    """Phase velocity omega / |k|."""
    k_mag = math.sqrt(kx ** 2 + ky ** 2 + kz ** 2)
    if k_mag < 1e-30:
        return 1.0
    om2 = omega_squared_exact(kx, ky, kz, a)
    if om2 <= 0.0:
        return 0.0
    return math.sqrt(om2) / k_mag


def lv_parameter(energy_gev: float) -> float:
    """MODEL-tier lattice LV estimate: the exact group-velocity deficit at
    k*a = E/E_Planck (identifying the lattice spacing with the Planck length).

    Leading order this is (E/E_P)^2 / 8.  Quadratic suppression is the only
    part carried by Lean (lvOrderOfMagnitude = 2); the coefficient and the
    a = Planck-length identification are modeling choices, not theorems.
    """
    ka = energy_gev / PLANCK_ENERGY_GEV
    if ka < 1e-14:
        # below double-precision resolution of 1 - v_g: use the leading term
        return ka * ka / 8.0
    vg = group_velocity_exact(ka, 0.0, 0.0, 1.0)
    return abs(vg - 1.0)


def compute_dispersion(kx: float, ky: float, kz: float, a: float = 1.0,
                       taylor_order: int = 4) -> Dict[str, Any]:
    """Full dispersion analysis for a given wavevector and lattice spacing."""
    k_mag = math.sqrt(kx ** 2 + ky ** 2 + kz ** 2)

    om2_exact = omega_squared_exact(kx, ky, kz, a)
    omega_exact = math.sqrt(max(om2_exact, 0.0))

    om2_taylor, terms = omega_squared_taylor(kx, ky, kz, a, order=taylor_order)
    omega_taylor = math.sqrt(max(om2_taylor, 0.0))

    vg = group_velocity_exact(kx, ky, kz, a)
    vp = phase_velocity_exact(kx, ky, kz, a)

    delta_vg = abs(vg - 1.0) if k_mag > 1e-30 else 0.0
    delta_vp = abs(vp - 1.0) if k_mag > 1e-30 else 0.0

    correction_fraction = 0.0
    if k_mag > 1e-30:
        om2_continuum = k_mag ** 2
        correction_fraction = abs(om2_exact - om2_continuum) / om2_continuum

    return {
        "kx": kx,
        "ky": ky,
        "kz": kz,
        "k_magnitude": k_mag,
        "lattice_spacing_a": a,
        "omega_squared_exact": om2_exact,
        "omega_exact": omega_exact,
        "omega_squared_taylor": om2_taylor,
        "omega_taylor": omega_taylor,
        "taylor_terms": terms,
        "group_velocity": vg,
        "phase_velocity": vp,
        "delta_v_group": delta_vg,
        "delta_v_phase": delta_vp,
        "correction_fraction": correction_fraction,
    }


def compute_lv_bounds_table() -> List[Dict[str, Any]]:
    """RS predictions vs experimental LV bounds."""
    rows = []
    for exp in EXPERIMENTAL_BOUNDS:
        e_gev = exp["energy_GeV"]
        rs_prediction = lv_parameter(e_gev)
        prediction_note = f"~(E/E_P)^2/8 = {rs_prediction:.2e}  [MODEL]"

        bound = exp["bound_linear"]
        safe = rs_prediction < bound

        rows.append({
            "experiment": exp["name"],
            "sector": exp["sector"],
            "energy_GeV": e_gev,
            "experimental_bound": bound,
            "rs_prediction": rs_prediction,
            "prediction_note": prediction_note,
            "within_bound": safe,
            "reference": exp["reference"],
        })
    return rows


def format_dispersion_report(result: Dict[str, Any]) -> str:
    lines = [
        "--- J-cost Lattice Dispersion ---",
        f"k = ({result['kx']:.6g}, {result['ky']:.6g}, {result['kz']:.6g})",
        f"|k| = {result['k_magnitude']:.6g}",
        f"a   = {result['lattice_spacing_a']:.6g}",
        "",
        f"omega^2 (exact)  = {result['omega_squared_exact']:.10g}",
        f"omega   (exact)  = {result['omega_exact']:.10g}",
        f"omega^2 (Taylor) = {result['omega_squared_taylor']:.10g}",
        "",
        "Taylor decomposition:",
    ]
    for name, val in result["taylor_terms"].items():
        lines.append(f"  {name:24s} = {val:+.6e}")
    lines += [
        "",
        f"group velocity v_g = {result['group_velocity']:.10g}",
        f"phase velocity v_p = {result['phase_velocity']:.10g}",
        f"|v_g - c| / c      = {result['delta_v_group']:.6e}",
        f"|v_p - c| / c      = {result['delta_v_phase']:.6e}",
        f"correction fraction = {result['correction_fraction']:.6e}",
    ]
    return "\n".join(lines)


def format_bounds_table(rows: List[Dict[str, Any]]) -> str:
    lines = [
        "",
        "--- RS vs Experimental Lorentz Violation Bounds ---",
        "",
        f"{'Experiment':<35s} {'Sector':<10s} {'E (GeV)':<12s} "
        f"{'Expt Bound':<14s} {'RS Pred':<14s} {'OK?':<5s}",
        "-" * 95,
    ]
    for r in rows:
        pred_str = f"{r['rs_prediction']:.2e}"
        ok_str = "YES" if r["within_bound"] else "NO"
        lines.append(
            f"{r['experiment']:<35s} {r['sector']:<10s} {r['energy_GeV']:<12.4g} "
            f"{r['experimental_bound']:<14.2e} {pred_str:<14s} {ok_str:<5s}"
        )

    lines += [
        "",
        "RS lattice MODEL prediction (all sectors): subluminal quadratic suppression,",
        f"delta ~ (E/E_Planck)^2 / 8 with E_Planck = {PLANCK_ENERGY_GEV:.4e} GeV.",
        "Lean carries the ORDER (lvOrderOfMagnitude = 2, LorentzViolationBoundFromRS.lean);",
        "the coefficient 1/8 and the a = Planck-length identification are MODEL choices.",
        "There is no Lean theorem of exact c at finite k for any sector.",
        "Any detection of LINEAR Lorentz violation O(E/E_P) falsifies this lattice picture.",
    ]
    return "\n".join(lines)


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description="J-cost lattice dispersion relation and Lorentz violation bounds."
    )
    p.add_argument("--kx", type=float, default=0.1,
                   help="x-component of wavevector (RS-native units, default 0.1)")
    p.add_argument("--ky", type=float, default=0.0, help="y-component of wavevector")
    p.add_argument("--kz", type=float, default=0.0, help="z-component of wavevector")
    p.add_argument("--a", type=float, default=1.0,
                   help="Lattice spacing in voxel units (default 1.0 = Planck scale)")
    p.add_argument("--taylor-order", type=int, default=4, choices=[2, 4, 6],
                   help="Taylor expansion order (default 4)")
    p.add_argument("--energy-gev", type=float, default=None,
                   help="Particle energy in GeV (overrides kx/ky/kz for isotropic k)")
    p.add_argument("--bounds", action="store_true",
                   help="Print experimental LV bounds comparison table")
    p.add_argument("--scan", action="store_true",
                   help="Scan k from 0.01 to pi/a in 20 steps along x-axis")
    p.add_argument("--json", action="store_true", help="Emit JSON output")
    return p


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.energy_gev is not None:
        # identify k*a = E/E_Planck (a = Planck length, hbar = c = 1)
        k_val = args.energy_gev / PLANCK_ENERGY_GEV / args.a
        args.kx = k_val
        args.ky = 0.0
        args.kz = 0.0

    if args.scan:
        k_max = math.pi / args.a
        steps = 20
        scan_results = []
        for i in range(1, steps + 1):
            k = k_max * i / steps
            r = compute_dispersion(k, 0.0, 0.0, a=args.a,
                                   taylor_order=args.taylor_order)
            scan_results.append(r)

        if args.json:
            print(json.dumps(scan_results, indent=2))
        else:
            print(f"{'|k|':>10s}  {'omega':>12s}  {'v_g':>12s}  "
                  f"{'v_p':>12s}  {'|dv_g/c|':>12s}  {'correction':>12s}")
            print("-" * 78)
            for r in scan_results:
                print(f"{r['k_magnitude']:10.6f}  {r['omega_exact']:12.8f}  "
                      f"{r['group_velocity']:12.8f}  {r['phase_velocity']:12.8f}  "
                      f"{r['delta_v_group']:12.6e}  {r['correction_fraction']:12.6e}")
        return

    result = compute_dispersion(args.kx, args.ky, args.kz, a=args.a,
                                taylor_order=args.taylor_order)

    if args.json:
        output = {"dispersion": result}
        if args.bounds:
            output["lv_bounds"] = compute_lv_bounds_table()
        print(json.dumps(output, indent=2))
    else:
        print(format_dispersion_report(result))
        if args.bounds:
            print(format_bounds_table(compute_lv_bounds_table()))


if __name__ == "__main__":
    main()
