"""Explore a universal variational correction to the LQ8 transition.

The correction is a new phenomenological term, not derived from eight bands.
Its amplitude and position are scanned; each acceleration scale is fitted
only to the existing full-sample spherical RAR proxy before the solar solve.
Cassini then informs family selection, so it is not an independent validation.
"""
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_transition_tweak"
OUT.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str(BASE / "outputs/dark_cassini_followup/.mplconfig"))
import numpy as np
from scipy.interpolate import PchipInterpolator
from scipy.optimize import minimize_scalar
from dark_solar_report import responses
from research import load_rar

BINARY = BASE / "rust/target/release/lesage-research"
REFERENCE = responses()[.45]
DATA, RAR_HASH = load_rar()
LOG_BAR, _, LOG_OBS, _ = DATA.T


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


def tag(amplitude, position):
    return f"A{amplitude:g}_T{position:g}".replace("-", "m").replace(".", "p")


class Tweak:
    def __init__(self, amplitude, position):
        self.amplitude, self.position = amplitude, position

    def added_energy(self, x):
        """Delta V / (K_b * g_dark**2); its derivative defines the response."""
        t = np.asarray(x, dtype=float) / self.position
        return (27/16)*self.amplitude*self.position**2*t**4/(1+t*t)**2

    def correction(self, x):
        t = np.asarray(x, dtype=float) / self.position
        t2 = t*t
        delta_mu = self.amplitude*(27/4)*t2/(1+t2)**3
        delta_radial = delta_mu*3*(1-t2)/(1+t2)
        return delta_mu, delta_radial

    def evaluate(self, x):
        mu, radial = REFERENCE.evaluate(x)
        delta_mu, delta_radial = self.correction(x)
        return mu+delta_mu, radial+delta_radial


def controls():
    prior = json.loads((BASE / "outputs/dark_solar_summary.json").read_text())
    assert digest(BASE / "rust/src/main.rs") == prior["source_sha256"]
    x = np.geomspace(.03, 30, 90)
    law = Tweak(.3, 1.7)
    h = 1e-4*x
    derivative = (law.added_energy(x+h)-law.added_energy(x-h))/(2*h)
    mu, radial = law.correction(x)
    first = float(np.max(abs(derivative/x-mu)))
    second = float(np.max(abs(((x+h)*law.correction(x+h)[0]-(x-h)*law.correction(x-h)[0])/(2*h)-radial)))
    assert first < 1e-8 and second < 1e-8
    zero = Tweak(0., 1.)
    expected = REFERENCE.evaluate(x)
    assert all(np.array_equal(a,b) for a,b in zip(zero.evaluate(x), expected))
    return {"energy_derivative_max_error":first, "radial_derivative_max_error":second,
            "zero_amplitude_matches_original":True}


def fit(amplitude, position):
    law = Tweak(amplitude, position)
    x = np.exp(np.linspace(-25,25,20001))
    mu, radial = law.evaluate(x)
    info = {"name":tag(amplitude, position), "amplitude":amplitude, "position":position,
            "min_mu":float(mu.min()), "min_radial":float(radial.min()),
            "max_mu":float(mu.max())}
    if np.any(mu <= 0) or np.any(radial <= 0) or np.any(np.diff(x*mu) <= 0):
        return {**info, "status":"rejected: static ellipticity/invertibility"}
    inverse = PchipInterpolator(np.log10(x*mu), np.log10(x), extrapolate=False)
    def objective(logscale):
        residual = logscale+inverse(LOG_BAR-logscale)-LOG_OBS
        return float(np.mean(residual**2))
    result = minimize_scalar(objective, bounds=(-12,-8), method="bounded", options={"xatol":1e-10})
    assert result.success and np.isfinite(result.fun)
    scale = 10**float(result.x)
    residual = float(result.x)+inverse(LOG_BAR-float(result.x))-LOG_OBS
    bins = []
    for low, high in [(-15,-11),(-11,-10.5),(-10.5,-10),(-10,-9.5),(-9.5,-7)]:
        mask = (LOG_BAR >= low)&(LOG_BAR < high)
        if mask.any():
            bins.append({"log_gbar_low":low,"log_gbar_high":high,"rows":int(mask.sum()),
                         "mean_residual_dex":float(residual[mask].mean()),
                         "rms_dex":float(np.sqrt(np.mean(residual[mask]**2)))})
    return {**info, "status":"admissible", "galaxy_scale_m_s2":scale,
            "deep_a0_m_s2":scale/REFERENCE.slope, "rar_rms_dex":float(np.sqrt(result.fun)),
            "rar_bins":bins, "scale_fit_scope":"One scale from full-sample spherical RAR; no Cassini term in scale fit"}


def solve(record, nr=96, external=2e-10, rmin=.003, rmax=40., table_nodes=10001):
    law = Tweak(record["amplitude"], record["position"])
    table = OUT / f"{record['name']}_table{table_nodes}.csv"
    x = np.exp(np.linspace(-25,25,table_nodes))
    mu, radial = law.evaluate(x)
    assert np.all(mu > 0) and np.all(radial > 0)
    np.savetxt(table, np.column_stack((x,mu)), delimiter=",", header="x,mu", comments="", fmt="%.17e")
    spec = {"name":record["name"], "nr":nr, "nu":nr*3//4, "rmin":rmin, "rmax":rmax,
            "external":external, "scale":record["galaxy_scale_m_s2"], "floor":1e-9,
            "table_sha256":digest(table), "solver_source_sha256":digest(BASE/"rust/src/main.rs")}
    code = hashlib.sha256(json.dumps(spec,sort_keys=True).encode()).hexdigest()[:12]
    prefix = OUT / f"{record['name']}_n{nr}_{code}"
    path = prefix.with_suffix(".json")
    if not path.exists():
        command = list(map(str,[BINARY,"table",spec["nr"],spec["nu"],rmin,rmax,external,prefix,
                                spec["floor"],table,spec["scale"]]))
        process = subprocess.run(command, capture_output=True, text=True)
        prefix.with_suffix(".log").write_text(process.stdout+"\n"+process.stderr)
        if process.returncode:
            return {"spec":spec, "status":"solver failed", "returncode":process.returncode,
                    "log":str(prefix.with_suffix(".log").relative_to(BASE))}
    result = json.loads(path.read_text())
    assert result["converged"] and result["nonlinear_residual"] < 2e-9
    for key, expected in [("nr",nr),("nu",spec["nu"]),("rmin",rmin),("rmax",rmax),
                          ("external_acceleration_m_s2",external),("model_scale_m_s2",spec["scale"])]:
        assert result[key] == expected
    q2 = result["Q2_s_minus2"][1]
    return {"spec":spec, "status":"converged", "result":result,
            "inside_quoted_cassini_two_sigma":bool(-2e-27 <= q2 <= 5.2e-27),
            "field":str(prefix.with_suffix(".csv").relative_to(BASE)),
            "field_sha256":digest(prefix.with_suffix(".csv")),
            "table":str(table.relative_to(BASE))}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scan", action="store_true")
    parser.add_argument("--amplitude", type=float, default=0.)
    parser.add_argument("--position", type=float, default=1.)
    parser.add_argument("--nr", type=int, default=96)
    parser.add_argument("--external", type=float, default=2e-10)
    parser.add_argument("--rmin", type=float, default=.003)
    parser.add_argument("--rmax", type=float, default=40.)
    parser.add_argument("--table-nodes", type=int, default=10001)
    args = parser.parse_args()
    check = controls()
    subprocess.run(["cargo","build","--release","--offline","--manifest-path",str(BASE/"rust/Cargo.toml"),
                    "--bin","lesage-research"], check=True)
    candidates = ([(0.,1.)]+[(amplitude,position) for position in [.5,1.,2.,4.]
                            for amplitude in [-.2,-.1,.1,.25,.5,1.]] if args.scan
                  else [(args.amplitude,args.position)])
    records = [fit(amplitude,position) for amplitude,position in candidates]
    admissible = [record for record in records if record["status"] == "admissible"]
    def run(record):
        field = solve(record,args.nr,args.external,args.rmin,args.rmax,args.table_nodes)
        record["solar"] = field
        q = field.get("result",{}).get("Q2_s_minus2",[None,None,None])[1]
        print(f"{record['name']}: RAR={record['rar_rms_dex']:.7f} dex; Q2={q}; "
              f"Cassini band={field.get('inside_quoted_cassini_two_sigma')}",flush=True)
    with ThreadPoolExecutor(max_workers=2) as pool:
        list(pool.map(run,admissible))
    report = {"scope":"Post-Cassini exploratory transition correction; no microscopic derivation or independent validation",
              "formula":"Delta V/(Kb*g_dark^2)=(27/16)*A*T^2*(x/T)^4/(1+(x/T)^2)^2; x=|grad Phi|/g_dark",
              "new_parameters":["A: amplitude","T: position"], "fixed_band_parameters":{"M":1,"Omega":20,"rW":.45},
              "controls":check, "rar_rows":len(DATA), "rar_sha256":RAR_HASH,
              "cases":records, "source_sha256":digest(__file__),
              "input_sha256":{name:digest(BASE/name) for name in ["rust/src/main.rs","dark_band_report.py",
                      "dark_solar_report.py","research.py","outputs/dark_lattice_bands.csv","outputs/dark_lattice_controls.csv"]},
              "cassini":{"central_s_minus2":1.6e-27,"one_sigma_s_minus2":1.8e-27,
                          "source":"https://arxiv.org/abs/2602.17884v2"},
              "limitations":["Shape family is selected after seeing Cassini; being in its band is a fit compatibility statement",
                             "Full-sample spherical RAR proxy lacks galaxy-level holdout and resolved disk fields",
                             "The added energy term is phenomenological; the original eight-band spectrum is unchanged",
                             "Static ellipticity does not establish dynamical stability or a metric completion",
                             "This model has no PPN gamma prediction"]}
    label = "scan" if args.scan else f"{tag(args.amplitude,args.position)}_n{args.nr}_ext{args.external:g}_r{args.rmin:g}_{args.rmax:g}_t{args.table_nodes}"
    (OUT/(label+"_summary.json")).write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")
    print("Saved",label,"with",len(records),"candidates",flush=True)


if __name__ == "__main__":
    main()
