#!/usr/bin/env python3
"""Offline independent dispersion roots, temporal ODE and bounded linear grid. NumPy2.0.2/SciPy1.13.1."""
import argparse
import hashlib
import importlib.util
import json
from pathlib import Path
import numpy as np
from scipy import __version__ as scipy_version
from scipy.linalg import eigvals
from scipy.integrate import solve_ivp

HERE=Path(__file__).resolve().parent
ROOT=HERE.parents[2] if len(HERE.parents)>2 else HERE
REPOSITORY_INPUT=ROOT/'src/features/labs/protoplanetary-dust/data/streaming-input.json'

def load(name):
    spec=importlib.util.spec_from_file_location(name,HERE/(name+'.py'));module=importlib.util.module_from_spec(spec);spec.loader.exec_module(module);return module
linear=load('linear-system');grid=load('grid-patch')

def pairs(values):return [[float(value.real),float(value.imag)]for value in values]

def analyze(case,settings):
    stokes,loading,kx,kz=[case[key]for key in ['stokes','loading','kx','kz']]
    matrix,basis=linear.incompressible_matrix(stokes,loading,kx,kz)
    value,vector,roots=linear.dominant(matrix)
    operator,mass=linear.youdin_goodman_pencil(stokes,loading,kx,kz)
    pencil=eigvals(operator,mass);finite=pencil[np.isfinite(pencil)]
    root_difference=max(min(abs(root-other)for other in finite)for root in roots)
    full_value,_,full_roots=linear.dominant(linear.compressible_matrix(stokes,loading,kx,kz))
    times=np.linspace(0,4,81)
    solution=solve_ivp(lambda time,state:matrix@state,(0,4),vector*1e-6,t_eval=times,method='DOP853',rtol=1e-11,atol=1e-15)
    if not solution.success:raise RuntimeError(solution.message)
    amplitude=np.abs(solution.y[-1]);ode_growth=float(np.polyfit(times,np.log(amplitude),1)[0])
    gas,dust=linear.equilibrium(stokes,loading)
    runs=[]
    for cells in settings['resolutions']:
        stencil_factor=np.sin(2*np.pi/cells)/(2*np.pi/cells)
        discrete_matrix,discrete_basis=linear.incompressible_matrix(stokes,loading,kx*stencil_factor,kz*stencil_factor)
        discrete_value,discrete_vector,_=linear.dominant(discrete_matrix)
        run=grid.evolve_grid(stokes,loading,kx,kz,gas,dust,discrete_vector,discrete_basis,cells,settings['end_time'],settings['step'])
        run.update({'discrete_expected_growth':float(discrete_value.real),'effective_wavenumber_factor':float(stencil_factor),'seed':'discrete-grid eigenvector, no supplied growth rate in integration'})
        runs.append(run)
    factor=np.sin(2*np.pi/32)/(2*np.pi/32)
    _,discrete_vector,_=linear.dominant(linear.incompressible_matrix(stokes,loading,kx*factor,kz*factor)[0])
    refined=grid.evolve_grid(stokes,loading,kx,kz,gas,dust,discrete_vector,basis,32,settings['end_time'],settings['step']/2)
    controls={}
    for name,options in [('no_feedback',{'loading':0}),('no_vertical_mode',{'kz':0}),('no_pressure_drift',{'pressure_speed':0}),('toy_diffusion',{'diffusion':.001})]:
        params={'stokes':stokes,'loading':loading,'kx':kx,'kz':kz};params.update(options)
        controls[name]=float(linear.dominant(linear.incompressible_matrix(**params)[0])[0].real)
    fg=1/(1+loading);fp=1-fg;wave_squared=kx*kx+kz*kz
    cubic=np.roots([1,stokes*(1j*fp*kx*kx/wave_squared+2*fg*fg*kx),-kz*kz/wave_squared,2*fg*kx*kz*kz/wave_squared*(fp-fg)*stokes])
    return {'name':case['name'],'settings':case,'sigma_incompressible':[float(value.real),float(value.imag)],'six_roots':pairs(roots),'direct_yg_pencil_finite_roots':pairs(finite),'max_independent_root_difference':float(root_difference),'sigma_full_weak_compressible':[float(full_value.real),float(full_value.imag)],'eight_roots':pairs(full_roots),'published_full_growth':case['published_full_growth'],'published_yj_growth':case['published_yj_growth'],'published_full_absolute_difference':abs(float(full_value.real)-case['published_full_growth']),'seed_relative_eigenvector':pairs(vector),'ode_growth':ode_growth,'ode_history':[{'time':float(time),'density_amplitude':float(amp)}for time,amp in zip(times,amplitude)],'grid_runs':runs,'halved_step_growth_32':refined['measured_growth'],'controls':controls,'terminal_velocity_cubic_growth':float(max(cubic.imag)),'terminal_velocity_cubic_domain_measure':float(stokes*np.sqrt(wave_squared))}

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input',type=Path,default=REPOSITORY_INPUT if REPOSITORY_INPUT.exists() else HERE/'streaming-input.json')
    parser.add_argument('--out',type=Path,default=ROOT/'docs/astronomy/ac14/streaming-replay.json' if REPOSITORY_INPUT.exists() else HERE/'streaming-replay-rerun.json')
    args=parser.parse_args();raw=args.input.read_bytes();data=json.loads(raw)
    result={'schema':'ac14-independent-linear-replay-v1','kind':data['kind'],'input_sha256':hashlib.sha256(raw).hexdigest(),'numpy_version':np.__version__,'scipy_version':scipy_version,'scope':'linear unstratified periodic patch only; no nonlinear saturation, vertical settling, self-gravity or planet production','conventions':data['units'],'results':[analyze(case,data['grid'])for case in data['cases']]}
    args.out.parent.mkdir(parents=True,exist_ok=True);args.out.write_text(json.dumps(result,indent=2,allow_nan=False)+'\n')
    for entry in result['results']:print(entry['name'],entry['sigma_incompressible'],'full8',entry['sigma_full_weak_compressible'],'ODE',entry['ode_growth'],'grid',[run['measured_growth']for run in entry['grid_runs']])

if __name__=='__main__':main()