"""Continue the fixed-parameter Cassini audit without claiming a nonlocal Sun.

New grid refinements test the unmodified local model. Two deliberately
large local response floors are counterfactual sensitivity tests, NOT
derived corrections to LQ8 or new candidates fitted to Cassini.

The spatial budget uses the scalar Poisson source, not an arbitrary
partition of a polarization flux: adding a divergence-free flux would
otherwise change the apparent budget of a selected spatial region.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
import os
from pathlib import Path
import subprocess

BASE = Path(__file__).resolve().parent
OUT = BASE / "outputs/dark_cassini_followup"
OUT.mkdir(exist_ok=True, parents=True)
os.environ.setdefault("MPLCONFIGDIR", str(OUT / ".mplconfig"))
import numpy as np
from scipy.interpolate import CubicSpline, RectBivariateSpline
from scipy.optimize import brentq
from dark_solar_report import responses
from dark_solar_green_check import normalization_check

BINARY = BASE / "rust/target/release/lesage-research"
SCALE = 1.8295491247141643e-10
EXTERNAL = 2e-10
GM = 6.67430e-11 * 1.98847e30  # Same normalization as the existing solver.
R_SCALE = np.sqrt(GM / SCALE)
FACTOR = SCALE / R_SCALE


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def gauss(n, a, b):
    t, w = np.polynomial.legendre.leggauss(n)
    return (a+b)/2+(b-a)*t/2, w*(b-a)/2


def composite_gauss(n, edges):
    """Respect each spline cell instead of integrating through Hessian knots."""
    t, w = np.polynomial.legendre.leggauss(n)
    half = np.diff(edges)/2
    center = (edges[:-1]+edges[1:])/2
    return (center[:,None]+half[:,None]*t).ravel(), (half[:,None]*w).ravel()


def declared_cases():
    return [{"name":f"w450_{nr}_floor{floor:g}".replace(".", "p"), "nr":nr, "nu":nr*3//4,
             "rmin":.003, "rmax":40., "scale":SCALE, "external":EXTERNAL,
             "floor":floor,
             "interpretation":("Unmodified local LQ8; numerical floor only" if floor==1e-9
                               else "Counterfactual local regularization; not the LQ8 nonlocal Hamiltonian")}
            for floor in [1e-9, .3, .6] for nr in [256, 384]]


def run_case(case):
    prefix = OUT / case["name"]
    command = list(map(str, [BINARY, "table", case["nr"], case["nu"],
                            case["rmin"], case["rmax"], EXTERNAL, prefix,
                            case["floor"], BASE/"outputs/dark_solar_table_w450.csv", SCALE]))
    result = subprocess.run(command, capture_output=True, text=True)
    prefix.with_suffix(".log").write_text(result.stdout+"\n"+result.stderr)
    if result.returncode:
        raise RuntimeError(f"Solver failed: {case['name']}: {result.stderr[-1000:]}")
    record = json.loads(prefix.with_suffix(".json").read_text())
    assert record["converged"] and record["nonlinear_residual"] < 2e-9
    assert record["model_scale_m_s2"] == SCALE
    return case["name"], record


def load_field(path):
    data = np.genfromtxt(path, delimiter=",", names=True)
    radius, angular = np.unique(data["r"]), np.unique(data["u"])
    correction = data["correction"].reshape(len(radius), len(angular))
    full_angular = np.r_[-1., angular, 1.]
    endpoint_extension = CubicSpline(angular, correction, axis=1)(full_angular)
    spline = RectBivariateSpline(np.log(radius), full_angular, endpoint_extension,
                                 kx=3, ky=3, s=0)
    return radius, angular, spline


def field_derivatives(spline, y, u, eta):
    r = np.exp(y)[:, None]
    u = u[None, :]
    s = np.sqrt(1-u*u)
    axis = u[0]
    wy = spline(y, axis, dx=1)
    wu = spline(y, axis, dy=1)
    wyy = spline(y, axis, dx=2)
    wyu = spline(y, axis, dx=1, dy=1)
    wuu = spline(y, axis, dy=2)
    gr = 1/r**2-eta*u+wy/r
    gu = -eta*r+wu
    gt = -s*gu/r
    g = np.hypot(gr, gt)
    hrr = -2/r**3+(wyy-wy)/r**2
    hrt = -s*(wyu-wu)/r**2
    htt = 1/r**3+(wy+(1-u*u)*wuu-u*wu)/r**2
    lap = (wyy+wy+(1-u*u)*wuu-2*u*wu)/r**2
    along_gradient = (gr*gr*hrr+2*gr*gt*hrt+gt*gt*htt)/g**2
    return {"r":r, "u":u, "gr":gr, "gu":gu, "gt":gt,
            "g":g, "laplacian":lap, "along_gradient":along_gradient}


def source_budget(path, response, order, floor=1e-9):
    radius, angular, spline = load_field(path)
    y, wy = composite_gauss(order, np.log(radius))
    u, wu = composite_gauss(order, np.r_[-1.,angular,1.])
    f = field_derivatives(spline, y, u, EXTERNAL/SCALE)
    mu, radial = response.evaluate(f["g"])
    active = mu > floor
    mu = np.maximum(mu, floor)
    log_slope = np.where(active, radial/mu-1, 0.)
    # Vacuum local AQUAL: mu*lap(Phi)+(dmu/dlogg)*e_g.H.e_g=0.
    source_pde = -log_slope*f["along_gradient"]
    source_laplacian = f["laplacian"]
    source_flux_divergence = (1-mu)*source_laplacian-mu*log_slope*f["along_gradient"]
    p2 = (3*u*u-1)/2
    measure = 1.5*FACTOR*wy[:,None]*wu[None,:]*p2[None,:]
    source_terms = {
        "pde_source":measure*source_pde,
        "laplacian_source":measure*source_laplacian,
        "flux_divergence_source":measure*source_flux_divergence,
    }
    # Full-domain flux identity, an independent extraction using first derivatives.
    integrand = (mu-1)*(p2*f["gr"]-u*(1-u*u)*f["gu"]/f["r"])
    q_flux = -4.5*FACTOR*np.einsum("i,ij,j->", wy/np.exp(y), integrand, wu)
    totals = {name:float(terms.sum()) for name, terms in source_terms.items()}
    assert all(np.isfinite(terms).all() for terms in source_terms.values())
    budgets = {}
    for name, terms in source_terms.items():
        total = totals[name]
        abs_total = float(abs(terms).sum())
        by_g = []
        for limit in [.03, .1, .3, .5, .75, 1., 1.5, 2., 3., 5., 10.]:
            mask = f["g"] < limit
            signed = float(terms[mask].sum())
            absolute = float(abs(terms[mask]).sum())
            by_g.append({"g_over_gdark_below":limit, "signed_Q2":signed,
                         "absolute_contribution":absolute,
                         "signed_fraction_of_net":signed/total,
                         "absolute_fraction_of_net":absolute/abs(total)})
        by_r = []
        for low, high in [(0., .1), (.1,.3), (.3,.6), (.6,1.),
                          (1.,1.5), (1.5,2.), (2.,3.), (3.,5.), (5.,10.), (10.,40.)]:
            mask = (np.exp(y)>=low)&(np.exp(y)<high)
            signed = float(terms[mask].sum())
            by_r.append({"r_over_Rdark_low":low, "r_over_Rdark_high":high,
                         "signed_Q2":signed, "fraction_of_net":signed/total,
                         "absolute_contribution":float(abs(terms[mask]).sum())})
        budgets[name] = {"Q2":total, "absolute_integral":abs_total,
                         "by_acceleration_threshold":by_g, "by_radial_shell":by_r}
    return {"quadrature_nodes_per_spline_cell":order, "total_nodes":[len(y),len(u)],
            "domain_r_over_Rdark":[float(radius[0]),float(radius[-1])],
            "minimum_sampled_g_over_gdark":float(f["g"].min()),
            "green_flux_Q2":float(q_flux), "source_Q2":totals, "budgets":budgets}


def manufactured_derivative_check():
    # A regular l=2 harmonic plus monopole and uniform field has zero Laplacian.
    y = np.linspace(np.log(.2), np.log(3), 240)
    u = np.linspace(-1., 1., 120)
    amplitude = .02
    w = amplitude*np.exp(2*y[:,None])*(3*u[None,:]**2-1)/2
    spline = RectBivariateSpline(y,u,w,s=0)
    yy = np.linspace(np.log(.3),np.log(2.5),45)
    uu = np.linspace(-.9,.9,40)
    f = field_derivatives(spline, yy, uu, 1.1)
    lap_error = float(abs(f["laplacian"]).max())
    assert lap_error < 2e-5
    # A non-harmonic l=2 potential checks the nonzero source and its sign.
    w4 = amplitude*np.exp(4*y[:,None])*(3*u[None,:]**2-1)/2
    spline4 = RectBivariateSpline(y,u,w4,s=0)
    f4 = field_derivatives(spline4,yy,uu,1.1)
    expected = 14*amplitude*np.exp(2*yy[:,None])*(3*uu[None,:]**2-1)/2
    source_relative = float(abs(f4["laplacian"]-expected).max()/abs(expected).max())
    assert source_relative < .001
    return {"regular_quadrupole_laplacian_max_error":lap_error,
            "nonharmonic_source_relative_error":source_relative,
            "green_normalization":normalization_check()}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", action="store_true")
    parser.add_argument("--pilot", action="store_true", help="Only inspect the existing 192x144 field; do not solve")
    args = parser.parse_args()
    controls = manufactured_derivative_check()
    response = responses()[.45]
    previous = json.loads((BASE/"outputs/dark_solar_summary.json").read_text())
    assert digest(BASE/"rust/src/main.rs") == previous["source_sha256"]
    cases = declared_cases()
    if args.run:
        subprocess.run(["cargo","build","--release","--offline","--manifest-path",
                        str(BASE/"rust/Cargo.toml"),"--bin","lesage-research"],check=True)
        with ThreadPoolExecutor(max_workers=2) as pool:
            for name, record in pool.map(run_case, cases):
                print(f"{name}: Q2={record['Q2_s_minus2'][1]:.10e}; "
                      f"residual={record['nonlinear_residual']:.3g}; "
                      f"{record['seconds']:.1f}s",flush=True)
    records = []
    for case in ([] if args.pilot else cases):
        result = json.loads((OUT/case["name"]).with_suffix(".json").read_text())
        assert result["converged"] and result["nonlinear_residual"]<2e-9
        assert (result["nr"],result["nu"],result["mu_floor"])==(case["nr"],case["nu"],case["floor"])
        assert result["model_scale_m_s2"]==SCALE and result["external_acceleration_m_s2"]==EXTERNAL
        records.append({"case":case, "result":result,
                        "field_sha256":digest((OUT/case["name"]).with_suffix(".csv"))})
    fields = [("reference_192",BASE/"outputs/dark_solar_w450_192.csv")]
    if not args.pilot:
        fields += [(f"refined_{n}",OUT/f"w450_{n}_floor1e-09.csv") for n in [256,384]]
    budgets = []
    for name,path in fields:
        low = source_budget(path,response,2)
        high = source_budget(path,response,4)
        entry = {"name":name,"field_sha256":digest(path),"coarse":low,"fine":high}
        budgets.append(entry)
        print(name, "flux", high["green_flux_Q2"], "source routes", high["source_Q2"], flush=True)
        print("  |g| budgets", [(r["g_over_gdark_below"],r["absolute_fraction_of_net"])
              for r in high["budgets"]["pde_source"]["by_acceleration_threshold"][:6]],flush=True)
    floor_effects = []
    if not args.pilot:
        for floor in [.3,.6]:
            threshold = brentq(lambda x:float(response.evaluate(x)[0])-floor,1e-4,EXTERNAL/SCALE)
            for nr in [256,384]:
                base = next(r["result"] for r in records if r["case"]["nr"]==nr and r["case"]["floor"]==1e-9)
                modified = next(r["result"] for r in records if r["case"]["nr"]==nr and r["case"]["floor"]==floor)
                floor_effects.append({"nr":nr,"floor":floor,"modified_below_g_over_gdark":threshold,
                                     "Q2":modified["Q2_s_minus2"][1],
                                     "relative_change":modified["Q2_s_minus2"][1]/base["Q2_s_minus2"][1]-1})
    selected = (previous["cases"]["w450_192"]["result"]["Q2_s_minus2"][1] if args.pilot
                else next(r["result"]["Q2_s_minus2"][1] for r in records if r["case"]["nr"]==384 and r["case"]["floor"]==1e-9))
    measurement = previous["cassini"]
    upper = measurement["central_s_minus2"]+2*measurement["one_sigma_s_minus2"]
    numerical = {}
    if not args.pilot:
        base256 = next(r["result"]["Q2_s_minus2"][1] for r in records
                       if r["case"]["nr"]==256 and r["case"]["floor"]==1e-9)
        numerical["relative_Q2_change_256_to_384"] = abs(base256/selected-1)
        numerical["relative_Q2_change_192_to_384"] = abs(previous["cases"]["w450_192"]["result"]["Q2_s_minus2"][1]/selected-1)
        numerical["relative_green_flux_difference_at_384"] = abs(budgets[-1]["fine"]["green_flux_Q2"]/selected-1)
        numerical["maximum_relative_source_method_difference_from_flux"] = max(
            abs(value/b["fine"]["green_flux_Q2"]-1)
            for b in budgets for value in b["fine"]["source_Q2"].values())
        numerical["maximum_relative_source_quadrature_change"] = max(
            abs(b["coarse"]["source_Q2"][key]/value-1)
            for b in budgets for key,value in b["fine"]["source_Q2"].items())
        assert numerical["relative_Q2_change_256_to_384"] < .001
        assert numerical["relative_green_flux_difference_at_384"] < .001
        assert numerical["maximum_relative_source_method_difference_from_flux"] < .01
        assert numerical["maximum_relative_source_quadrature_change"] < .001
    report = {
        "scope":"Refined local Cassini test, Poisson-source budget, and counterfactual local-floor sensitivity; no nonlocal solar solution",
        "controls":controls, "galaxy_parameters_held_fixed":{"M":1,"Omega":20,"wilson":.45,"g_dark_m_s2":SCALE},
        "solar_scale_Rdark_m":float(R_SCALE),
        "solar_GM_convention":"6.67430e-11 * 1.98847e30, retained from the existing solver; distinct from the nominal IAU scale illustration",
        "cases":records, "source_budgets":budgets, "floor_diagnostics":floor_effects,
        "numerical_controls":numerical,
        "cassini":measurement,
        "comparison":{"selected_local_Q2":selected,"observational_upper_two_sigma":upper,
                      "fractional_reduction_to_upper_two_sigma":1-upper/selected,
                      "interpretation":"Comparison with the quoted observational band; no combined likelihood or complete theoretical uncertainty"},
        "source_sha256":digest(__file__),
        "input_sha256":{name:digest(BASE/name) for name in [
            "rust/src/main.rs","outputs/dark_summary.json","outputs/dark_solar_summary.json",
            "outputs/dark_lattice_bands.csv","outputs/dark_lattice_controls.csv",
            "outputs/dark_solar_table_w450.csv"]},
        "dependencies_sha256":{name:digest(BASE/name) for name in [
            "dark_solar_report.py","dark_solar_green_check.py","dark_band_report.py"]},
        "limitations":[
            "The new gravity solutions remain local AQUAL, not the nonlocal eight-band field equation",
            "Large mu floors are counterfactual diagnostics and alter the microscopic local law",
            "Poisson-source budgets depend on interpolated second derivatives and finite-domain boundaries",
            "The source budget is not a bound on the backreaction from a nonlocal modification",
            "The very small saddle region is not resolved to microscopic length scales",
            "The separate Cassini relativistic light-propagation/PPN test has no LQ8 metric prediction",
        ],
    }
    path = OUT/("pilot.json" if args.pilot else "summary.json")
    path.write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")
    print(json.dumps({"comparison":report["comparison"],"floor_diagnostics":floor_effects},indent=2))


if __name__ == "__main__":
    main()
