"""Solar study of a metric model with an auxiliary momentum-trace constraint.

The action, constraint and extremely weak-field return to Newtonian response
are new phenomenological assumptions. The previously fitted galaxy transition
and its acceleration scale are retained. This file does not fit tracking data.
"""
from __future__ import annotations

import argparse
from concurrent.futures import ThreadPoolExecutor
import hashlib
import json
import subprocess

import numpy as np
from scipy.interpolate import PchipInterpolator, CubicHermiteSpline

from dark_transition_tweak import BASE, DATA, LOG_BAR, LOG_OBS, Tweak, digest
from dark_joint_solar import response_offsets
from build_cmc_solver import build

OUT=BASE/"outputs/dark_cmc"
IR_EPSILON=1e-9
IR_CUTOFF=.01


class CMCResponse:
    def __init__(self, epsilon=IR_EPSILON, cutoff=IR_CUTOFF):
        assert epsilon>0 and cutoff>0
        self.epsilon,self.cutoff=epsilon,cutoff
        self.raw=Tweak(-.45,4.25)
        # Repair the inherited low/core numerical splice at x=1e-4.
        # Join the force, including its derivative, well below all RAR points.
        self.join_interval=np.array([5e-5,2e-4])
        mu,radial=self.raw.evaluate(self.join_interval)
        self.force_join=CubicHermiteSpline(self.join_interval,
                                          self.join_interval*mu,radial)

    def smooth_previous(self,x):
        x=np.asarray(x,dtype=float)
        mu,radial=self.raw.evaluate(x)
        inside=(x>=self.join_interval[0]) & (x<=self.join_interval[1])
        safe=np.clip(x,*self.join_interval)
        return (np.where(inside,self.force_join(safe)/safe,mu),
                np.where(inside,self.force_join(safe,1),radial))

    def correction(self,x):
        x=np.asarray(x,dtype=float)
        z=(x/self.cutoff)**2
        extra=self.epsilon/(x+self.epsilon)/(1+z)
        radial=extra*(self.epsilon/(x+self.epsilon)-2*z/(1+z))
        return extra,radial

    def evaluate(self,x):
        mu,radial=self.smooth_previous(x)
        dm,dr=self.correction(x)
        return mu+dm,radial+dr

    def offsets(self,x):
        dm,dr=response_offsets(self.raw,x)
        mu_smooth,r_smooth=self.smooth_previous(x)
        x=np.asarray(x,dtype=float)
        dm=np.where(x<=100,mu_smooth-1,dm)
        dr=np.where(x<=100,r_smooth-1,dr)
        extra,radial=self.correction(x)
        return dm+extra,dr+radial

    def lapse_coefficients(self,x):
        mu,radial=self.evaluate(x)
        at,ar=2*(1-mu),2*(1-radial)
        x=np.asarray(x)
        dm,dr=self.offsets(x)
        return np.where(x>100,-2*dm,at),np.where(x>100,-2*dr,ar)

    def record(self):
        return {"amplitude":-.45,"position":4.25,"IR_epsilon":self.epsilon,"IR_cutoff":self.cutoff,
                "mu_at_zero":1.,"radial_stiffness_at_zero":1.,
                "numerical_force_splice_interval":self.join_interval.tolist(),
                "response":"mu_C(x)=mu_previous_smoothed(x)+epsilon/[(x+epsilon)*(1+(x/ell)^2)]"}


def candidate(law):
    frozen=json.loads((BASE/"outputs/dark_transition_tweak/publication_summary.json").read_text())["candidate"]
    scale=frozen["galaxy_scale_m_s2"]
    # Use the same grid as the historical fit to isolate the new response.
    x=np.exp(np.linspace(-25,25,20001))
    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)
    logscale=np.log10(scale)
    residual=logscale+inverse(LOG_BAR-logscale)-LOG_OBS
    old_mu,_=law.raw.evaluate(x)
    old_inverse=PchipInterpolator(np.log10(x*old_mu),np.log10(x),extrapolate=False)
    old_residual=logscale+old_inverse(LOG_BAR-logscale)-LOG_OBS
    predicted_x=10**inverse(LOG_BAR-logscale)
    return {**law.record(),"galaxy_scale_m_s2":scale,"rar_rms_dex":float(np.sqrt(np.mean(residual**2))),
            "previous_rar_rms_same_grid_dex":float(np.sqrt(np.mean(old_residual**2))),
            "maximum_RAR_prediction_change_dex":float(np.max(abs(residual-old_residual))),
            "minimum_RAR_predicted_x":float(predicted_x.min()),"rar_rows":len(DATA),
            "scale_refitted":False,"scope":"Frozen previous galaxy scale; full-sample spherical RAR proxy",
            "IR_parameter_acceleration_m_s2":law.epsilon*scale,
            "IR_cutoff_acceleration_m_s2":law.cutoff*scale}


def solve(law,record,binary,build_record,nr=384,external=2.32e-10,rmin=.0003,rmax=40.,table_nodes=60001):
    OUT.mkdir(parents=True,exist_ok=True)
    x=np.geomspace(1e-14,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":0.,"physical_mu_at_zero":1.,
          "table_nodes":table_nodes,"source_sha256":digest(__file__),
          "solver_source_sha256":build_record["generated_source_sha256"]}
    code=hashlib.sha256(json.dumps(spec,sort_keys=True).encode()).hexdigest()[:16]
    prefix=OUT/f"solar_{code}"
    table=OUT/f"solar_{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,0.,table,spec["scale"],1.]))
        done=subprocess.run(command,capture_output=True,text=True)
        prefix.with_suffix(".log").write_text(done.stdout+"\n"+done.stderr)
        if done.returncode:
            raise RuntimeError(f"Solar solver failed; inspect {prefix.with_suffix('.log')}")
    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()
    binary,build_record=build()
    law=CMCResponse()
    record=candidate(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":120001},
            {"nr":512,"rmin":.00001}]]
    else:
        variants=[CMCResponse(e,l) for e,l in [(1e-10,.01),(1e-8,.01),(1e-9,.003),(1e-9,.03)]]
        jobs=[(variant,candidate(variant),{"nr":256}) for variant in variants]
    def run(job):
        item,fit_record,kwargs=job
        result=solve(item,fit_record,binary,build_record,**kwargs)
        print(json.dumps({"law":item.record(),"run":kwargs,"Q2":result["result"]["Q2_s_minus2"]}),flush=True)
        return {"candidate":fit_record,"solar":result}
    results=[]
    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,
            "solver":build_record,"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",
                "outputs/dark_transition_tweak/publication_summary.json","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()
