"""Frozen-parameter local AQUAL solar test of LQ8. No ephemeris refit.

--run builds the Rust solver, verifies response tables, and executes the
predeclared eight-parameter first suite. Additional numerical controls are
explicit cases, not a reoptimization of the galactic scales.
"""
from __future__ import annotations
import argparse
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
from pathlib import Path
import subprocess
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from dark_band_report import Response

BASE=Path(__file__).resolve().parent
OUT=BASE/"outputs"
MANIFEST=BASE/"rust/Cargo.toml"
BINARY=BASE/"rust/target/release/lesage-research"


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


def responses():
    bands=np.genfromtxt(OUT/"dark_lattice_bands.csv",delimiter=",",names=True)
    controls=np.genfromtxt(OUT/"dark_lattice_controls.csv",delimiter=",",names=True)
    return {float(c["wilson"]):Response(bands[np.isclose(bands["wilson"],c["wilson"])],c)
            for c in controls[:-1]}


def label(w):
    return f"w{round(1000*w):03d}"


def table_path(w,refined=False):
    return OUT/f"dark_solar_table_{label(w)}{'_fine' if refined else ''}.csv"


def write_tables(response_map):
    checks=[]
    for w,response in response_map.items():
        for refined in [False,True]:
            x=np.exp(np.linspace(-25,25,20001 if refined else 10001))
            mu,radial=response.evaluate(x)
            assert np.all(mu>0) and np.all(radial>0)
            path=table_path(w,refined)
            np.savetxt(path,np.column_stack((x,mu)),delimiter=",",header="x,mu",comments="",fmt="%.17e")
        control=OUT/f"dark_solar_interpolation_{label(w)}.csv"
        subprocess.run([str(BINARY),"table-response",str(table_path(w)),str(control)],check=True)
        data=np.genfromtxt(control,delimiter=",",names=True)
        exact,_=response.evaluate(data["x"])
        absolute=float(np.max(abs(data["mu"]-exact)))
        # The declared deep-asymptotic switch at x=1e-4 is checked in
        # absolute error; the main transition receives a relative check too.
        main=data["x"]>=1e-3
        relative=float(np.max(abs(data["mu"][main]/exact[main]-1)))
        assert absolute<1e-5 and relative<1e-5,(w,absolute,relative)
        h=.01
        plus,_=response.evaluate(data["x"]*np.exp(h))
        minus,_=response.evaluate(data["x"]*np.exp(-h))
        slope=(np.log(plus)-np.log(minus))/(2*h)
        slope_error=float(np.max(abs(data["log_slope"][main]-slope[main])))
        assert slope_error<2e-4,(w,slope_error)
        checks.append({"wilson":w,"absolute_mu_error":absolute,"relative_mu_error_x_ge_1e_minus3":relative,
                       "log_slope_error_x_ge_1e_minus3":slope_error,"table_sha256":digest(table_path(w))})
    return checks


def cases(fits):
    first=[{"name":label(f["wilson"])+"_96","wilson":f["wilson"],"scale":f["g_dark_m_s2"],
            "nr":96,"nu":72,"rmin":.003,"rmax":40.,"external":2e-10,"floor":1e-9,"fine_table":False}
           for f in fits]
    extra=[]
    # Chosen after the first eight cases: smallest Q2, best RAR, largest
    # overshoot. These are numerical controls, not new fitted parameters.
    for w in [.1,.45,.5]:
        base=next(c for c in first if np.isclose(c["wilson"],w))
        for nr,nu in [(128,96),(192,144)]:
            extra.append(dict(base,name=label(w)+f"_{nr}",nr=nr,nu=nu))
        if w==.5:
            extra.extend(dict(base,name=label(w)+f"_{nr}",nr=nr,nu=nu)
                         for nr,nu in [(256,192),(384,288)])
    for w in [.1,.45]:
        base=next(c for c in extra if c["name"]==label(w)+"_128")
        for radius in [20.,80.]:
            extra.append(dict(base,name=label(w)+f"_box{round(radius)}",rmax=radius))
        for ext in [1.5e-10,2.5e-10]:
            extra.append(dict(base,name=label(w)+f"_ext{round(ext/1e-11)}",external=ext))
    base=next(c for c in extra if c["name"]=="w450_128")
    extra.extend([dict(base,name="w450_inner",rmin=.001),
                  dict(base,name="w450_floor",floor=1e-7),
                  dict(base,name="w450_fine_table",fine_table=True)])
    return first+extra


def run_case(case):
    prefix=OUT/f"dark_solar_{case['name']}"
    command=list(map(str,[BINARY,"table",case["nr"],case["nu"],case["rmin"],case["rmax"],
                          case["external"],prefix,case["floor"],table_path(case["wilson"],case["fine_table"]),case["scale"]]))
    result=subprocess.run(command,text=True,capture_output=True,check=False)
    prefix.with_suffix(".log").write_text(result.stdout+"\n"+result.stderr)
    return {"name":case["name"],"returncode":result.returncode}


def inspect_case(case):
    path=OUT/f"dark_solar_{case['name']}.json"
    result=json.loads(path.read_text())
    assert result["model"]=="Table"
    assert result["model_scale_m_s2"]==case["scale"]
    assert (result["nr"],result["nu"])==(case["nr"],case["nu"])
    assert (result["rmin"],result["rmax"])==(case["rmin"],case["rmax"])
    assert result["external_acceleration_m_s2"]==case["external"]
    assert result["mu_floor"]==case["floor"]
    return dict(case=case,result=result,table_sha256=digest(table_path(case["wilson"],case["fine_table"])))


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run",action="store_true")
    parser.add_argument("--only",nargs="+",help="With --run, execute only these declared cases; inspect the entire suite afterwards")
    args=parser.parse_args()
    if args.only is not None and not args.run:
        parser.error("--only requires --run")
    reference=json.loads((OUT/"dark_summary.json").read_text())
    suite=cases(reference["RAR"]["fits"])
    interpolation_path=OUT/"dark_solar_interpolation_checks.json"
    runs=[]
    if args.run:
        subprocess.run(["cargo","test","--release","--manifest-path",str(MANIFEST),"--bin","lesage-research"],check=True)
        subprocess.run(["cargo","build","--release","--manifest-path",str(MANIFEST),"--bin","lesage-research"],check=True)
        interpolation=write_tables(responses())
        interpolation_path.write_text(json.dumps(interpolation,indent=2)+"\n")
        print("Response interpolation checks passed; starting frozen-scale solar cases.",flush=True)
        selected=[c for c in suite if args.only is None or c["name"] in args.only]
        if args.only is not None:
            assert {c["name"] for c in selected}==set(args.only)
        with ThreadPoolExecutor(max_workers=4) as pool:
            for result in pool.map(run_case,selected):
                runs.append(result)
                print(result,flush=True)
        for model,nr,nu in [("newton",96,72),("tanh",128,96)]:
            subprocess.run(list(map(str,[BINARY,model,nr,nu,.003,40,2e-10,
                                         OUT/f"dark_solar_{model}_regression"])),
                           check=True,capture_output=True,text=True)
    interpolation=json.loads(interpolation_path.read_text())
    results={case["name"]:inspect_case(case) for case in suite}
    response_map=responses()
    external_response=[]
    for fit in reference["RAR"]["fits"]:
        w=fit["wilson"]
        eta=2e-10/fit["g_dark_m_s2"]
        mu,radial=response_map[w].evaluate(eta)
        external_response.append({"wilson":w,"eta":eta,"mu_external":float(mu),
                                  "log_slope_external":float(radial/mu-1),"spherical_boost_proxy":float(1/mu)})
    q=lambda name:results[name]["result"]["Q2_s_minus2"][1]
    numerical={"primary_case":"w450_192","primary_Q2_s_minus2":q("w450_192"),
               "smallest_family_case":"w100_192","smallest_refined_Q2_s_minus2":q("w100_192"),
               "largest_refined_case":"w500_384","largest_refined_Q2_s_minus2":q("w500_384"),
               "relative_grid_change":{"w100_128_to_192":abs(q("w100_128")/q("w100_192")-1),
                                       "w450_128_to_192":abs(q("w450_128")/q("w450_192")-1),
                                       "w500_256_to_384":abs(q("w500_256")/q("w500_384")-1)},
               "primary_response_table_change_s_minus2":abs(q("w450_128")-q("w450_fine_table")),
               "primary_box_change_s_minus2":abs(q("w450_box20")-q("w450_box80")),
               "primary_inner_boundary_change_s_minus2":abs(q("w450_128")-q("w450_inner")),
               "primary_floor_change_s_minus2":abs(q("w450_128")-q("w450_floor"))}
    assert max(numerical["relative_grid_change"].values())<.005
    newton=json.loads((OUT/"dark_solar_newton_regression.json").read_text())
    tanh=json.loads((OUT/"dark_solar_tanh_regression.json").read_text())
    old_tanh=json.loads((OUT/"solar_tanh_128.json").read_text())
    regression={"newton_max_abs_Q2_s_minus2":max(map(abs,newton["Q2_s_minus2"])),
                "tanh_max_abs_change_from_archive_s_minus2":float(np.max(abs(np.array(tanh["Q2_s_minus2"])-old_tanh["Q2_s_minus2"])))}
    assert newton["converged"] and tanh["converged"]
    assert regression["newton_max_abs_Q2_s_minus2"]<1e-35
    assert regression["tanh_max_abs_change_from_archive_s_minus2"]<1e-35
    report={"scope":"LQ8 stationary local AQUAL test at frozen galactic scales; no full quantum metric dynamics or ephemeris refit",
            "source_sha256":digest(BASE/"rust/src/main.rs"),"report_source_sha256":digest(BASE/"dark_solar_report.py"),
            "galactic_reference_sha256":digest(OUT/"dark_summary.json"),"entire_suite_rerun_this_invocation":args.run and args.only is None,
            "runs_this_invocation":runs,"numerical_controls":numerical,"external_response":external_response,
            "analytic_response_regression_checks":regression,
            "response_interpolation":interpolation,"cases":results,
            "cassini":{"central_s_minus2":1.6e-27,"one_sigma_s_minus2":1.8e-27,
                       "source":"https://arxiv.org/abs/2602.17884v2"},
            "limitations":["No independent second nonlinear field solver", "No joint likelihood or formal significance",
                           "Galactic scales come from the spherical full-sample RAR proxy", "Nonlocal and metric corrections are not computed in this solar test"]}
    (OUT/"dark_solar_summary.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")
    for name,result in results.items():
        data=result["result"]
        print(name,"converged=",data["converged"],"Q2 / 1e-27 =",np.array(data["Q2_s_minus2"])/1e-27)
    assert all(v["result"]["converged"] and v["result"]["nonlinear_residual"]<2e-9 for v in results.values())

    fig,axes=plt.subplots(1,2,figsize=(10.5,4.2),layout="constrained")
    fits=reference["RAR"]["fits"]
    wilson=np.array([f["wilson"] for f in fits])
    axes[0].plot(wilson,[f["rms_dex"] for f in fits],"o-",color="#305b7c",label="LQ8, une échelle par r_W")
    axes[0].axhline(reference["RAR"]["empirical_RAR_one_scale_rms_dex"],color="gray",ls="--",label="Référence empirique")
    q=[results[label(w)+"_96"]["result"]["Q2_s_minus2"][1]/1e-27 for w in wilson]
    axes[1].plot(wilson,q,"o-",color="#a64f36",label="Grille 96 × 72")
    for w in [.1,.45,.5]:
        axes[1].plot(w,results[label(w)+"_192"]["result"]["Q2_s_minus2"][1]/1e-27,"x",color="black",
                     label="Contrôles 192 × 144" if w==.1 else None)
    axes[1].axhspan(-.2,3.4,color="#77a987",alpha=.3,label="Cassini : intervalle publié à 1σ")
    axes[0].set(xlabel="Paramètre de réseau r_W",ylabel="RMS RAR (dex)",title="Comparaison galactique descriptive")
    axes[1].set(xlabel="Même paramètre r_W",ylabel="Q₂ (10⁻²⁷ s⁻²)",title="Test solaire à échelles RAR figées")
    for ax in axes:ax.grid(alpha=.2);ax.legend(fontsize=8)
    fig.supxlabel("M=1, Ω=20 ; réalisation locale supposée ; aucune réoptimisation sur Cassini ni significativité jointe.",fontsize=9)
    for suffix in ("png","pdf"):fig.savefig(OUT/f"dark_solar_comparison.{suffix}",dpi=170)
    plt.close(fig)


if __name__=="__main__":main()
