"""Independent quadrupole extraction from the archived AQUAL field.

This is a Green-identity check of one field solution, not a second nonlinear
PDE solver. It integrates F=(mu-1) grad(phi) over the resolved spatial domain.
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
from scipy.interpolate import CubicSpline, RectBivariateSpline
from dark_solar_report import responses

BASE=Path(__file__).resolve().parent
OUT=BASE/"outputs"


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


def normalization_check():
    r,wr=gauss(128,0.,10.)
    u,wu=gauss(64,-1.,1.)
    p2=(3*u*u-1)/2
    radial=r**4*np.exp(-r*r)
    flux_integral=-4.5*np.dot(wr,radial/r**2)*np.dot(wu,p2*p2)
    # -div F for F_r=r^4 exp(-r^2) P2, separately integrated as a source.
    source=-(6*r**3-2*r**5)*np.exp(-r*r)
    source_integral=1.5*np.dot(wr,source/r)*np.dot(wu,p2*p2)
    exact=-9*np.sqrt(np.pi)/20
    assert abs(flux_integral-exact)<1e-12
    assert abs(source_integral-exact)<1e-12
    return {"exact":exact,"flux_integral":flux_integral,"source_integral":source_integral}


def green(case,response,order):
    path=OUT/f"dark_solar_{case['name']}.csv"
    data=np.genfromtxt(path,delimiter=",",names=True)
    r=np.unique(data["r"])
    u=np.unique(data["u"])
    w=data["correction"].reshape(len(r),len(u))
    # Extend only the half-cell at each polar endpoint, then interpolate
    # the small correction. The large harmonic monopole is analytic.
    full_u=np.r_[-1.,u,1.]
    full_w=CubicSpline(u,w,axis=1)(full_u)
    spline=RectBivariateSpline(np.log(r),full_u,full_w,kx=3,ky=3,s=0)
    y,wy=gauss(order,np.log(r[0]),np.log(r[-1]))
    angular,wa=gauss(order,-1.,1.)
    radial=np.exp(y)[:,None]
    eta=case["external"]/case["scale"]
    gr=1/radial**2-eta*angular+spline(y,angular,dx=1)/radial
    du=-eta*radial+spline(y,angular,dy=1)
    gt=-np.sqrt(1-angular**2)*du/radial
    mu,_=response.evaluate(np.hypot(gr,gt))
    mu=np.maximum(mu,case["floor"])
    p2=(3*angular**2-1)/2
    integrand=(mu-1)*(p2*gr-angular*(1-angular**2)*du/radial)
    assert np.all(np.isfinite(integrand)) and np.all(np.isfinite(wy))
    # Explicit contractions avoid platform BLAS floating-status warnings;
    # no warning is silenced and all physical intermediate arrays are checked.
    q_dimensionless=-4.5*np.einsum("i,ij,j->",wy/np.exp(y),integrand,wa)
    length=np.sqrt(6.67430e-11*1.98847e30/case["scale"])
    return float(q_dimensionless*case["scale"]/length)


if __name__=="__main__":
    normalization=normalization_check()
    summary=json.loads((OUT/"dark_solar_summary.json").read_text())
    response_map=responses()
    results=[]
    for name in ["w100_96","w100_128","w100_192","w450_96","w450_128","w450_192","w500_96","w500_128","w500_192","w500_256","w500_384"]:
        item=summary["cases"][name]
        case=item["case"]
        low=green(case,response_map[case["wilson"]],320)
        high=green(case,response_map[case["wilson"]],640)
        direct=item["result"]["Q2_s_minus2"][1]
        result={"name":name,"green_Q2_s_minus2":high,"local_fit_Q2_s_minus2":direct,
                "quadrature_change_s_minus2":abs(high-low),"relative_difference":high/direct-1}
        results.append(result)
        print(result,flush=True)
        assert abs(high/direct-1)<.01 and abs(high-low)/abs(high)<.005
    report={"scope":"Independent global Green extraction of the same archived field; not a second nonlinear field solver",
            "normalization_check":normalization,"results":results,
            "limitations":["Field interpolated from finite-volume output", "Polar endpoints extrapolated by half a cell",
                           "Integral restricted to resolved domain; omitted boundaries require controls", "No rigorous field error bound"]}
    (OUT/"dark_solar_green_checks.json").write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")
