#!/usr/bin/env python3
"""Recompute photon-axion survival along saved cluster sightlines.

Sandbox Physics AC79 method tutorial. Standard library only.
Usage:  python3 -I recompute_sightline.py sightline-inputs.json sightline-outputs.json

Conventions: natural Heaviside-Lorentz units (hbar = c = 1), basis (A_x, A_y, a),
i d(psi)/dz = -M psi, so each cell multiplies the state by exp(i M L).
Delta_pl = -omega_pl^2/(2E), Delta_a = -m_a^2/(2E), Delta_ag = g B_perp / 2.
This script does not use the Lab code. It diagonalizes each 2x2 block itself.
"""
import cmath
import json
import math
import platform
import random
import sys

# CODATA 2018 values (SI); every conversion below is derived from them.
HBAR_EV_S = 6.582119569e-16
HBARC_EV_M = 1.973269804e-7
ELEMENTARY_CHARGE_C = 1.602176634e-19
VACUUM_PERMITTIVITY = 8.8541878128e-12
VACUUM_PERMEABILITY = 1.25663706212e-6
ELECTRON_MASS_KG = 9.1093837015e-31
METRES_PER_KPC = 3.0856775814913673e19

KPC_IN_INVERSE_EV = METRES_PER_KPC / HBARC_EV_M
TESLA_IN_EV2 = math.sqrt((1 / VACUUM_PERMEABILITY / ELEMENTARY_CHARGE_C) * HBARC_EV_M ** 3)
MICROGAUSS_IN_EV2 = TESLA_IN_EV2 * 1e-10


def plasma_frequency_ev(density_per_cm3):
    density_per_m3 = density_per_cm3 * 1e6
    omega = math.sqrt(density_per_m3 * ELEMENTARY_CHARGE_C ** 2 / (VACUUM_PERMITTIVITY * ELECTRON_MASS_KG))
    return HBAR_EV_S * omega


def cell_terms(case, domain, energy_kev, plasma_sign=-1.0):
    """Return (Delta_pl, Delta_a, Delta_ag) in kpc^-1."""
    energy_ev = energy_kev * 1e3
    plasma = plasma_sign * plasma_frequency_ev(case["electronDensityPerCubicCentimetre"]) ** 2 / (2 * energy_ev) * KPC_IN_INVERSE_EV
    axion = -case["axionMassEv"] ** 2 / (2 * energy_ev) * KPC_IN_INVERSE_EV
    field_ev2 = case["fieldMicrogauss"] * domain["strengthScale"] * MICROGAUSS_IN_EV2
    mixing = case["couplingPerGev"] * 1e-9 * field_ev2 / 2 * KPC_IN_INVERSE_EV
    return plasma, axion, mixing


def evolve_cell(state, plasma, axion, mixing, length, angle):
    """Rotate into (perpendicular, parallel), evolve with the 2x2 eigenbasis, rotate back."""
    ax, ay, a = state
    c, s = math.cos(angle), math.sin(angle)
    parallel = c * ax + s * ay
    perpendicular = (-s * ax + c * ay) * cmath.exp(1j * plasma * length)
    theta = 0.5 * math.atan2(2 * mixing, plasma - axion)
    mean = 0.5 * (plasma + axion)
    half_split = 0.5 * math.hypot(plasma - axion, 2 * mixing)
    ct, st = math.cos(theta), math.sin(theta)
    upper = (ct * parallel + st * a) * cmath.exp(1j * (mean + half_split) * length)
    lower = (-st * parallel + ct * a) * cmath.exp(1j * (mean - half_split) * length)
    parallel, a = ct * upper - st * lower, st * upper + ct * lower
    return [c * parallel - s * perpendicular, s * parallel + c * perpendicular, a]


def observe(case, domains, energy_kev, plasma_sign=-1.0):
    """Unpolarized start: average the two pure photon polarizations."""
    finals = []
    for start in ([1 + 0j, 0j, 0j], [0j, 1 + 0j, 0j]):
        state = start
        for domain in domains:
            plasma, axion, mixing = cell_terms(case, domain, energy_kev, plasma_sign)
            state = evolve_cell(state, plasma, axion, mixing, case["domainLengthKiloparsec"], domain["angle"])
        finals.append(state)
    rho = [[0.5 * sum(f[i] * f[j].conjugate() for f in finals) for j in range(3)] for i in range(3)]
    xx, yy, aa, xy = rho[0][0].real, rho[1][1].real, rho[2][2].real, rho[0][1].real
    photon = xx + yy
    return {
        "photonSurvival": photon,
        "axionFraction": aa,
        "linearPolarization": math.hypot(xx - yy, 2 * xy) / photon if photon > 1e-15 else 0.0,
        "normalizationError": abs(photon + aa - 1),
    }


def parallel_probability(plasma, axion, mixing, length):
    split = math.hypot(plasma - axion, 2 * mixing)
    return 0.0 if split == 0 else (2 * mixing / split) ** 2 * math.sin(split * length / 2) ** 2


def ideal_reference(case, energy_kev):
    domains = case["domains"]
    if case["fieldMode"] == "uniform":
        plasma, axion, mixing = cell_terms(case, domains[0], energy_kev)
        total = case["domainLengthKiloparsec"] * len(domains)
        return {"kind": "uniform-analytic", "survival": 1 - parallel_probability(plasma, axion, mixing, total) / 2}
    product = 1.0
    for domain in domains:
        product *= 1 - 1.5 * parallel_probability(*cell_terms(case, domain, energy_kev), case["domainLengthKiloparsec"])
    return {"kind": "random-angle-mean", "survival": 1 - (1 - product) / 3}


def random_domains(generator, count):
    return [{"angle": 2 * math.pi * generator.random(), "strengthScale": 1.0} for _ in range(count)]


def monte_carlo_check(case, energy_kev, realizations, seed):
    """Mean over fresh random angles (Mersenne Twister, independent of the Lab generator)."""
    generator = random.Random(seed)
    values = [observe(case, random_domains(generator, len(case["domains"])), energy_kev)["axionFraction"] for _ in range(realizations)]
    mean = sum(values) / len(values)
    spread = math.sqrt(sum((v - mean) ** 2 for v in values) / (len(values) - 1))
    expected = 1 - ideal_reference(case, energy_kev)["survival"]
    return {"energyKev": energy_kev, "realizations": realizations, "seed": seed, "meanAxionFraction": mean,
            "standardError": spread / math.sqrt(len(values)), "ensembleFormula": expected,
            "pullInStandardErrors": (mean - expected) / (spread / math.sqrt(len(values)))}


def log_grid(count, low=0.5, high=10.0):
    return [low * (high / low) ** (i / (count - 1)) for i in range(count)]


def failure_cases(cases):
    by_id = {case["id"]: case for case in cases}
    resonance, tangled = by_id["plasma-resonance"], by_id["tangled-cluster"]
    grid = log_grid(40)
    right = [observe(resonance, resonance["domains"], e)["photonSurvival"] for e in grid]
    flipped = [observe(resonance, resonance["domains"], e, plasma_sign=+1.0)["photonSurvival"] for e in grid]
    coarse_grid, fine_grid = log_grid(12), log_grid(400)
    coarse = [observe(tangled, tangled["domains"], e)["photonSurvival"] for e in coarse_grid]
    interpolation_error = 0.0
    for e in fine_grid:
        k = min(10, int(math.log(e / 0.5) / math.log(20.0) * 11))
        t = math.log(e / coarse_grid[k]) / math.log(coarse_grid[k + 1] / coarse_grid[k])
        guess = coarse[k] + t * (coarse[k + 1] - coarse[k])
        interpolation_error = max(interpolation_error, abs(guess - observe(tangled, tangled["domains"], e)["photonSurvival"]))
    # Forward Euler on one strong cell: the update matrix is not unitary, so probability drifts.
    plasma, axion, mixing = cell_terms(tangled, {"angle": 0.0, "strengthScale": 1.0}, 3.0)
    parallel, a, steps = 1 + 0j, 0j, 16
    h = tangled["domainLengthKiloparsec"] / steps
    for _ in range(steps):
        parallel, a = parallel + 1j * h * (plasma * parallel + mixing * a), a + 1j * h * (mixing * parallel + axion * a)
    deviations = [abs(observe(tangled, tangled["domains"], e)["photonSurvival"] - ideal_reference(tangled, e)["survival"]) for e in grid]
    return [
        {"id": "plasma-sign-flip", "what": "Use Delta_pl = +omega_pl^2/(2E) at the resonance mass.",
         "correctSpread": max(right) - min(right), "flippedSpread": max(flipped) - min(flipped),
         "lesson": "With the right sign the resonance survival is flat in energy; the flipped sign destroys the resonance."},
        {"id": "coarse-energy-grid", "what": "Sample the tangled spectrum at 12 instead of 400 log-spaced energies.",
         "largestInterpolationError": interpolation_error,
         "lesson": "Interpolating a coarse grid misses narrow ripples of a single realization."},
        {"id": "forward-euler", "what": "Replace the exact cell transfer with 16 forward-Euler steps in one 8 kpc cell at 3 keV.",
         "totalProbability": abs(parallel) ** 2 + abs(a) ** 2,
         "lesson": "Euler steps are not unitary; the total probability no longer stays one."},
        {"id": "one-sightline-is-not-the-mean", "what": "Largest gap between the saved realization and the ensemble mean, 0.5 to 10 keV.",
         "maximumGap": max(deviations),
         "lesson": "A single field realization scatters around the mean; do not compare it to the mean as if they were the same curve."},
    ]


def main(input_path, output_path):
    with open(input_path, encoding="utf-8") as handle:
        inputs = json.load(handle)
    results = []
    for case in inputs["cases"]:
        rows = []
        for energy in inputs["energiesKev"]:
            row = {"energyKev": energy, **observe(case, case["domains"], energy)}
            row["idealReference"] = ideal_reference(case, energy)
            rows.append(row)
        results.append({"id": case["id"], "rows": rows})
    tangled = next(case for case in inputs["cases"] if case["id"] == "tangled-cluster")
    output = {
        "format": "sandboxphysics.ac79.tutorial-outputs.v1",
        "kind": "offline simulation of saved model inputs; not an observation",
        "generatedBy": {"script": "recompute_sightline.py", "python": platform.python_version(), "libraries": "standard library only"},
        "units": inputs["units"],
        "derivedConstants": {"kpcInInverseEv": KPC_IN_INVERSE_EV, "teslaInEv2": TESLA_IN_EV2,
                             "plasmaFrequencyEvAtOnePerCubicCentimetre": plasma_frequency_ev(1.0)},
        "results": results,
        "monteCarlo": monte_carlo_check(tangled, 3.0, 2000, 79),
        "failureCases": failure_cases(inputs["cases"]),
    }
    with open(output_path, "w", encoding="utf-8") as handle:
        json.dump(output, handle, indent=2)
        handle.write("\n")


if __name__ == "__main__":
    if len(sys.argv) != 3:
        sys.exit("usage: recompute_sightline.py INPUTS.json OUTPUTS.json")
    main(sys.argv[1], sys.argv[2])
