"""A nonsymmetric 3D variational benchmark of the local transition model.

One periodic potential is varied in all three Cartesian directions. The
source is a synthetic density contrast, not the Sun or an observed galaxy.
The effective energy is specified; its quantum origin is not derived here.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
import time

from dark_transition_tweak import BASE, Tweak, digest, fit
import numpy as np
from scipy.fft import fftn, ifftn, fftfreq
from scipy.interpolate import PchipInterpolator
from scipy.signal import resample
from scipy.sparse.linalg import LinearOperator, cg

OUT = BASE / "outputs/dark_gravity_3d"
OUT.mkdir(parents=True,exist_ok=True)
BOX = 8.


def rms(values):
    return float(np.sqrt(np.mean(values*values)))


def resize(values, n):
    for axis in range(3):
        values = resample(values,n,axis=axis)
    return np.real(values)


class EnergyLaw:
    def __init__(self, amplitude=-.45, position=4., newton=False):
        self.newton = newton
        if newton:
            self.response_error = 0.
            return
        self.law = Tweak(amplitude,position)
        x = np.r_[0.,np.geomspace(1e-10,1e4,24001)]
        mu, _ = self.law.evaluate(x)
        # Interpolate U'(g)=g*mu(g), then derive both U and U'' from that
        # same interpolant. This preserves the discrete variational identity.
        force = x*mu
        force[0] = 0.
        assert np.all(np.diff(force)>0)
        self.first = PchipInterpolator(x,force,extrapolate=False)
        self.second = self.first.derivative()
        self.primitive = self.first.antiderivative()
        self.offset = float(self.primitive(0.))
        probe = np.geomspace(.001,100,12347)
        self.response_error = float(abs(self.first(probe)/probe-self.law.evaluate(probe)[0]).max())
        assert self.response_error < 2e-7

    def evaluate(self, g):
        if self.newton:
            return .5*g*g, np.ones_like(g), np.ones_like(g)
        assert np.max(g)<1e4
        first = self.first(g)
        mu = np.divide(first,g,out=np.zeros_like(g),where=g>0)
        radial = self.second(g)
        energy = self.primitive(g)-self.offset
        assert np.all(mu>=0) and np.all(radial>=0)
        assert np.isfinite(energy).all()
        return energy,mu,radial


class Grid:
    def __init__(self,n,external,source=None):
        assert n%2==1, "Odd grids avoid a Nyquist derivative null space"
        self.n,self.external = n,np.asarray(external)[:,None,None,None]
        k = 2*np.pi*fftfreq(n,d=BOX/n)
        self.k = [k[:,None,None],k[None,:,None],k[None,None,:]]
        self.k2 = sum(w*w for w in self.k)
        self.inverse_k2 = np.divide(1.,self.k2,out=np.zeros_like(self.k2),where=self.k2>0)
        self.source = density(n) if source is None else source
        self.source = self.source-self.source.mean()

    def gradient(self,potential):
        spectrum = fftn(potential)
        return np.array([ifftn(1j*wave*spectrum).real for wave in self.k])

    def divergence(self,vector):
        return sum(ifftn(1j*wave*fftn(component)).real for wave,component in zip(self.k,vector))

    def newton_potential(self):
        return -ifftn(fftn(self.source)*self.inverse_k2).real

    def energy_gradient(self,potential,law):
        gvec = self.external+self.gradient(potential)
        g = np.sqrt(np.sum(gvec*gvec,axis=0))
        local_energy,mu,radial = law.evaluate(g)
        energy = float(np.mean(local_energy+self.source*potential))
        gradient = -self.divergence(mu*gvec)+self.source
        gradient -= gradient.mean()
        unit = np.divide(gvec,g,out=np.zeros_like(gvec),where=g>0)
        return energy,gradient,(mu,radial,unit,g)

    def hessian(self,variation,coefficients):
        mu,radial,unit,_ = coefficients
        dv = self.gradient(variation)
        parallel = np.sum(unit*dv,axis=0)
        action = -self.divergence(mu*dv+(radial-mu)*unit*parallel)
        return action-action.mean()


def density(n):
    coordinate = np.arange(n)*BOX/n-BOX/2
    blobs = [((-1.3,.5,-.2),(.65,.8,.55),18.),
             ((1.1,-1.2,.9),(.5,.65,.9),9.)]
    value = np.zeros((n,n,n))
    for center,width,amplitude in blobs:
        factors = [sum(np.exp(-.5*((coordinate-c+shift*BOX)/s)**2) for shift in [-1,0,1])
                   for c,s in zip(center,width)]
        value += amplitude*factors[0][:,None,None]*factors[1][None,:,None]*factors[2][None,None,:]
    return value-value.mean()


def solve(grid,law,initial=None):
    shape = (grid.n,)*3
    size = grid.n**3
    potential = np.zeros(shape) if initial is None else initial.copy()
    potential -= potential.mean()
    source_norm = rms(grid.source)
    history = []
    start = time.monotonic()
    for iteration in range(35):
        energy,gradient,coefficients = grid.energy_gradient(potential,law)
        residual = rms(gradient)/source_norm
        history.append({"iteration":iteration,"energy":energy,"relative_residual":residual})
        if residual<2e-9:
            break
        hessian = LinearOperator((size,size),matvec=lambda v:grid.hessian(v.reshape(shape),coefficients).ravel(),dtype=float)
        mu,radial,_,_ = coefficients
        coefficient = float(np.mean((2*mu+radial)/3))
        def precondition(v):
            return (ifftn(fftn(v.reshape(shape))*grid.inverse_k2).real/coefficient).ravel()
        pre = LinearOperator((size,size),matvec=precondition,dtype=float)
        count = [0]
        def callback(_):
            count[0] += 1
        delta,status = cg(hessian,-gradient.ravel(),M=pre,rtol=1e-8,atol=0.,maxiter=600,callback=callback)
        assert status==0, (grid.n,iteration,status)
        delta = delta.reshape(shape)
        delta -= delta.mean()
        slope = float(np.mean(gradient*delta))
        assert slope<0
        step = 1.
        for _ in range(30):
            proposed = potential+step*delta
            trial,_,_ = grid.energy_gradient(proposed,law)
            if trial <= energy+1e-4*step*slope+1e-14*max(1.,abs(energy)):
                potential = proposed
                break
            step *= .5
        else:
            raise RuntimeError("3D Newton line search failed")
        history[-1].update({"cg_iterations":count[0],"step":step})
    else:
        raise RuntimeError(f"3D nonlinear solve did not converge: {residual}")
    mu,radial,_,g = coefficients
    record = {"n":grid.n,"seconds":time.monotonic()-start,"iterations":iteration,
              "relative_residual":residual,"energy":energy,
              "minimum_acceleration":float(g.min()),"maximum_acceleration":float(g.max()),
              "minimum_mu":float(mu.min()),"minimum_radial_stiffness":float(radial.min()),
              "history":history}
    return potential,record


def controls(law,external):
    n = 9
    grid = Grid(n,external)
    axis = np.arange(n)*2*np.pi/n
    x,y,z = np.meshgrid(axis,axis,axis,indexing="ij")
    potential = .08*np.sin(x+2*y)+.05*np.cos(y-z)+.02*np.sin(2*z-x)
    u = np.sin(x+y+z)+.3*np.cos(2*y-z)
    v = np.cos(2*x+y-z)+.2*np.sin(x-2*z)
    energy,gradient,coefficients = grid.energy_gradient(potential,law)
    errors = []
    for h in [1e-4,3e-5]:
        plus,gplus,_ = grid.energy_gradient(potential+h*v,law)
        minus,gminus,_ = grid.energy_gradient(potential-h*v,law)
        fd = (plus-minus)/(2*h)
        expected = float(np.mean(gradient*v))
        hessian_error = rms((gplus-gminus)/(2*h)-grid.hessian(v,coefficients))/rms(grid.hessian(v,coefficients))
        errors.append({"step":h,"energy_gradient_absolute_error":abs(fd-expected),
                       "hessian_relative_error":hessian_error})
    symmetry = abs(float(np.mean(u*grid.hessian(v,coefficients)-v*grid.hessian(u,coefficients))))
    positive = float(np.mean(v*grid.hessian(v,coefficients)))
    # Cyclic permutation is a rotation of the cube and all vector components.
    order = (1,2,0)
    rotated = Grid(n,np.asarray(external)[list(order)],grid.source.transpose(order))
    er,gr,_ = rotated.energy_gradient(potential.transpose(order),law)
    rotation_error = rms(gr-gradient.transpose(order))
    phi_newton,newton_record = solve(grid,EnergyLaw(newton=True))
    newton_error = rms(phi_newton-grid.newton_potential())/rms(grid.newton_potential())
    assert errors[-1]["energy_gradient_absolute_error"]<2e-8
    assert errors[-1]["hessian_relative_error"]<2e-5
    assert symmetry<1e-11 and positive>0
    assert rotation_error<1e-11 and abs(er-energy)<1e-12
    assert newton_error<1e-10
    return {"directional_checks":errors,"hessian_symmetry_error":symmetry,
            "tested_direction_curvature":positive,"cyclic_rotation_gradient_error":rotation_error,
            "newton_exact_poisson_relative_error":newton_error,
            "interpolated_response_max_error":law.response_error}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pilot",action="store_true")
    args = parser.parse_args()
    candidate = fit(-.45,4.)
    direction = np.array([2.,3.,5.])/np.sqrt(38.)
    external = direction*2.32e-10/candidate["galaxy_scale_m_s2"]
    law = EnergyLaw()
    check = controls(law,external)
    print("3D variational, Hessian, rotation and exact Newtonian controls passed.",flush=True)
    previous = None
    records = []
    for n in ([17] if args.pilot else [17,25,33,49]):
        grid = Grid(n,external)
        initial = None if previous is None else resize(previous,n)
        potential,record = solve(grid,law,initial)
        if previous is not None:
            restricted = resize(potential,len(previous))
            previous_grid = Grid(len(previous),external)
            record["refinement_potential_relative_L2"] = rms(restricted-previous)/rms(restricted)
            record["refinement_gradient_relative_L2"] = rms(previous_grid.gradient(restricted-previous))/rms(previous_grid.gradient(restricted))
        path = OUT/f"field_n{n}.npz"
        np.savez_compressed(path,potential=potential,source=grid.source,external=external,
                            potential_gradient=external[:,None,None,None]+grid.gradient(potential),
                            acceleration=-(external[:,None,None,None]+grid.gradient(potential)))
        record["field_sha256"] = digest(path)
        records.append(record)
        previous = potential
        print(json.dumps({k:v for k,v in record.items() if k!="history"}),flush=True)
    comparison = None
    if not args.pilot:
        baseline,baseline_record = solve(grid,EnergyLaw(amplitude=0.),potential)
        np.savez_compressed(OUT/"baseline_n49.npz",potential=baseline,
                            potential_gradient=external[:,None,None,None]+grid.gradient(baseline),
                            acceleration=-(external[:,None,None,None]+grid.gradient(baseline)))
        comparison = {"baseline_solve":baseline_record,
                      "gradient_change_from_original_local_relative_L2":rms(grid.gradient(potential-baseline))/rms(grid.gradient(baseline)),
                      "baseline_sha256":digest(OUT/"baseline_n49.npz")}
        assert records[-1]["refinement_gradient_relative_L2"]<.02
    report = {"scope":"Nonsymmetric Cartesian 3D benchmark of one local effective potential; not the quantum nonlocal theory",
              "boundary":"Periodic cube; zero-mean potential and zero-mean synthetic density contrast",
              "energy":"mean[U(|external+grad phi|)+source*phi]; U'=g*mu_new(g)",
              "equation":"div[mu_new(|external+grad phi|)*(external+grad phi)]=source",
              "method":"Odd-grid Fourier derivatives; variational Newton-CG with a Poisson preconditioner",
              "box_length":BOX,"candidate":candidate,"external_vector":external.tolist(),
              "external_vector_convention":"Background potential gradient; physical acceleration has the opposite sign",
              "controls":check,"grids":records,"comparison":comparison,
              "source_sha256":digest(__file__),
              "dependencies_sha256":{"dark_transition_tweak.py":digest(BASE/"dark_transition_tweak.py")},
              "limitations":["Synthetic periodic source with compensating background; no fit to a galaxy",
                             "Finite-grid static equilibrium; no dynamical or relativistic validation",
                             "New energy is phenomenological; quantum-interface geometry and common-potential assembly remain underived",
                             "The 3D benchmark is not used to extract a solar Cassini quadrupole"]}
    (OUT/("pilot.json" if args.pilot else "summary.json")).write_text(json.dumps(report,indent=2,allow_nan=False)+"\n")


if __name__ == "__main__":
    main()
