"""A constrained metric candidate retaining the fitted LQ8 transition.

The kinetic-conformal action is imported, and a small infrared detuning is
postulated. This is a local effective candidate, not a quantum derivation or
a successful cosmology. All historical calculations remain untouched.
"""
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
from scipy.interpolate import PchipInterpolator
from scipy.optimize import minimize_scalar

from dark_transition_tweak import BASE, BINARY, DATA, LOG_BAR, LOG_OBS, REFERENCE, Tweak, digest
from dark_joint_solar import response_offsets

OUT = BASE / "outputs/dark_constrained_metric"
AMPLITUDE = -.45
POSITION = 4.25
IR_DETUNING = 1e-6


class ConstrainedResponse:
    def __init__(self, detuning=IR_DETUNING):
        assert detuning >= 0
        self.detuning = detuning
        self.raw = Tweak(AMPLITUDE,POSITION)

    def evaluate(self, x):
        mu, radial = self.raw.evaluate(x)
        return (mu+self.detuning)/(1+self.detuning),(radial+self.detuning)/(1+self.detuning)

    def offsets(self, x):
        dm, dr = response_offsets(self.raw,x)
        return dm/(1+self.detuning),dr/(1+self.detuning)

    def lapse_coefficients(self, x):
        dm, dr = self.offsets(x)
        return -2*dm,-2*dr

    def record(self):
        return {"amplitude":AMPLITUDE,"position":POSITION,"IR_detuning":self.detuning,
                "mu_at_zero":self.detuning/(1+self.detuning),
                "lambda_H":1/3,"alpha_K_high_acceleration":0.,"b_K":0.}


def fit(law):
    x = np.geomspace(1e-11,1e11,30001)
    mu, radial = law.evaluate(x)
    assert np.all(mu>0) and np.all(radial>0) and np.all(np.diff(x*mu)>0)
    inverse = PchipInterpolator(np.log10(x*mu),np.log10(x),extrapolate=False)
    optimum = minimize_scalar(lambda z:float(np.mean((z+inverse(LOG_BAR-z)-LOG_OBS)**2)),
                              bounds=(-12,-8),method="bounded",options={"xatol":1e-11})
    assert optimum.success
    return {**law.record(),"galaxy_scale_m_s2":10**float(optimum.x),
            "rar_rms_dex":float(np.sqrt(optimum.fun)),"rar_rows":len(DATA),
            "scope":"One scale fitted to the existing full-sample spherical RAR proxy; response selected after known solar constraints",
            "intermediate_MOND_a0_m_s2":10**float(optimum.x)*(1+law.detuning)/REFERENCE.slope,
            "IR_crossover_acceleration_m_s2":10**float(optimum.x)*law.detuning/REFERENCE.slope}


def solve(law, record, nr=256, external=2.32e-10, rmin=.0003, rmax=40., table_nodes=30001):
    OUT.mkdir(parents=True,exist_ok=True)
    x = np.geomspace(1e-11,1e11,table_nodes)
    mu,_ = law.evaluate(x)
    spec = {"law":law.record(),"scale":record["galaxy_scale_m_s2"],"nr":nr,"nu":nr*3//4,
            "external":external,"rmin":rmin,"rmax":rmax,
            "floor":law.detuning/(1+law.detuning),"table_nodes":table_nodes,
            "source_sha256":digest(__file__),"rust_sha256":digest(BASE/"rust/src/main.rs")}
    code = hashlib.sha256(json.dumps(spec,sort_keys=True).encode()).hexdigest()[:16]
    prefix = OUT/f"constrained_{code}"
    table = OUT/f"constrained_{code}_table.csv"
    np.savetxt(table,np.column_stack((x,mu)),delimiter=",",header="x,mu",comments="",fmt="%.17e")
    if not prefix.with_suffix(".json").exists():
        command = list(map(str,[BINARY,"table",nr,spec["nu"],rmin,rmax,external,prefix,
                                spec["floor"],table,spec["scale"]]))
        done = subprocess.run(command,capture_output=True,text=True)
        prefix.with_suffix(".log").write_text(done.stdout+"\n"+done.stderr)
        if done.returncode:
            return {"spec":spec,"status":"solver failure","returncode":done.returncode}
    result = json.loads(prefix.with_suffix(".json").read_text())
    assert result["converged"] and result["nonlinear_residual"]<2e-9
    q2 = result["Q2_s_minus2"][1]
    return {"spec":spec,"status":"converged","result":result,
            "inside_quoted_Q2_two_sigma":bool(-2e-27<=q2<=5.2e-27),
            "field":str(prefix.with_suffix(".csv").relative_to(BASE)),
            "table":str(table.relative_to(BASE)),"field_sha256":digest(prefix.with_suffix(".csv")),
            "table_sha256":digest(table)}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--stage",choices=["pilot","fields","controls","ir"],default="pilot")
    args = parser.parse_args()
    law = ConstrainedResponse()
    record = fit(law)
    if args.stage == "pilot":
        jobs = [(law,record,{"nr":192})]
    elif args.stage == "fields":
        jobs = [(law,record,{"nr":384,"external":float(g)})
                for g in np.linspace(2.,2.64,17)*1e-10]
    elif args.stage == "controls":
        jobs = [(law,record,kw) for kw in [
            {"nr":256},{"nr":512},
            {"nr":384,"rmin":.0001},{"nr":384,"rmax":80.},
            {"nr":384,"table_nodes":60001}]]
    else:
        laws = [ConstrainedResponse(delta) for delta in [0.,1e-8,1e-5,1e-4]]
        jobs = [(item,fit(item),{"nr":256}) for item in laws]
    def run(job):
        item, fit_record, kwargs = job
        result = solve(item,fit_record,**kwargs)
        print(json.dumps({"law":item.record(),"run":kwargs,
              "Q2":result.get("result",{}).get("Q2_s_minus2"),"status":result["status"]}),flush=True)
        return {"candidate":fit_record,"solar":result}
    results = []
    OUT.mkdir(parents=True,exist_ok=True)
    with ThreadPoolExecutor(max_workers=2) as pool:
        for row in pool.map(run,jobs):
            results.append(row)
            (OUT/f"{args.stage}_checkpoint.json").write_text(json.dumps(results,indent=2,allow_nan=False)+"\n")
    report = {"source_sha256":digest(__file__),"stage":args.stage,"candidate":record,"cases":results,
              "input_sha256":{name:digest(BASE/name) for name in ["dark_transition_tweak.py","dark_joint_solar.py",
                  "dark_band_report.py","dark_solar_report.py","research.py","outputs/dark_lattice_bands.csv",
                  "outputs/dark_lattice_controls.csv","OVERNIGHT_SOLAR_PROTOCOL.md"]}}
    (OUT/f"{args.stage}_summary.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")


if __name__ == "__main__":
    main()
