#!/usr/bin/env python3
"""Offline Parker critical-wind and Jeans benchmark, standard-library Python3.10+.
Numerical synthetic boundary conditions, not measured planet profiles. No Lab imports.
"""
import argparse
import hashlib
import json
import math
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2] if len(HERE.parents)>=3 else HERE
REPOSITORY_INPUT = ROOT/'src/features/labs/atmospheric-escape/data/escape-input.json'
DEFAULT_INPUT = REPOSITORY_INPUT if REPOSITORY_INPUT.exists() else HERE/'escape-input.json'


def bisect(function, lower, upper, iterations=100):
    left = function(lower)
    if left*function(upper)>0:raise ValueError('Root is not bracketed')
    for _ in range(iterations):
        middle = (lower+upper)/2
        value = function(middle)
        if value*left>0:lower,left=middle,value
        else:upper=middle
    return (lower+upper)/2


def transonic_log_mach(radius_over_sonic, maximum_step=.0015):
    """Independent regular ODE integration, not the browser's implicit root solver."""
    target = radius_over_sonic
    if abs(target-1)<1e-4:
        delta=target-1
        return delta-delta*delta+13*delta**3/12
    direction=1 if target>1 else -1
    epsilon=direction*1e-5
    x=1+epsilon
    value=epsilon-epsilon**2+13*epsilon**3/12
    def derivative(position,log_mach):
        return 2*(1/position-1/position**2)/math.expm1(2*log_mach)
    while abs(target-x)>1e-13:
        step=direction*min(maximum_step*min(x,1),abs(target-x))
        a=derivative(x,value)
        b=derivative(x+step/2,value+step*a/2)
        c=derivative(x+step/2,value+step*b/2)
        d=derivative(x+step,value+step*c)
        value+=step*(a+2*b+2*c+d)/6
        x+=step
    return value


def integrate_tail(lambda_value):
    """Direct Simpson integral of the flux-weighted Maxwellian energy tail."""
    intervals=1200
    step=50/intervals
    total=0
    for index in range(intervals+1):
        energy=lambda_value+index*step
        weight=1 if index in (0,intervals) else 4 if index%2 else 2
        total+=weight*energy*math.exp(-energy)
    return total*step/3


def analyze(case,constants):
    settings=case['settings']
    gm=settings['planetMassEarth']*constants['earth_gm_m3_s2']
    base=settings['baseRadiusEarth']*constants['earth_radius_m']
    sound=math.sqrt(constants['boltzmann_j_per_k']*settings['temperatureKelvin']/constants['particle_mass_kg'])
    density=10**settings['logBaseNumberDensity']
    sonic=bisect(lambda radius:2*sound**2/radius-gm/radius**2,.01*base,100*base)
    base_log=transonic_log_mach(base/sonic)
    # Integrate each target independently from the regular critical expansion.
    radii=[math.exp(index/120*math.log(constants['outer_radius_ratio']))for index in range(121)]
    if 1<=sonic/base<=constants['outer_radius_ratio']:radii.append(sonic/base)
    profile=[]
    for ratio in sorted(radii):
        log_mach=transonic_log_mach(ratio*base/sonic)
        velocity=sound*math.exp(log_mach)
        number=density*math.exp(base_log-log_mach-2*math.log(ratio))
        mass_density=number*constants['particle_mass_kg']
        profile.append({'radius_ratio':ratio,'radius_m':ratio*base,'velocity_m_s':velocity,'mach':math.exp(log_mach),'number_density_m3':number,'mass_density_kg_m3':mass_density,'mass_flux_kg_s':4*math.pi*(ratio*base)**2*mass_density*velocity})
    wind_rate=profile[0]['mass_flux_kg_s']
    sonic_number=density*math.exp(base_log)*(base/sonic)**2
    sonic_kn=3/(math.sqrt(2)*sonic_number*constants['collision_cross_section_m2']*sonic)
    available=settings['heatingEfficiency']*math.pi*base**2*10**settings['logXuvFlux']
    outer=profile[-1]
    mechanical=.5*(outer['velocity_m_s']**2-profile[0]['velocity_m_s']**2)+gm*(1/base-1/outer['radius_m'])
    energy_limited=available/(gm/base)
    lambda_base=gm/(base*sound**2)
    def log_kn(ratio):
        radius=ratio*base
        log_number=math.log(density)+lambda_base*(1/ratio-1)
        height=sound**2*radius**2/gm
        return -math.log(math.sqrt(2)*constants['collision_cross_section_m2']*height)-log_number
    exobase=None
    if sonic>base and log_kn(1)<0 and log_kn(sonic/base)>0:
        ratio=bisect(log_kn,1,sonic/base)
        number=density*math.exp(lambda_base*(1/ratio-1))
        binding=lambda_base/ratio
        numerical_tail=integrate_tail(binding)
        exact_tail=(1+binding)*math.exp(-binding)
        exobase={'radius_ratio':ratio,'lambda':binding,'number_density_m3':number,'flux_weighted_escape_fraction':exact_tail,'numerical_tail_integral':numerical_tail,'relative_tail_error':abs(numerical_tail/exact_tail-1),'mass_loss_kg_s':4*math.pi*(ratio*base)**2*constants['particle_mass_kg']*number*sound/math.sqrt(2*math.pi)*numerical_tail}
    fluxes=[point['mass_flux_kg_s']for point in profile]
    bernoulli=[.5*point['velocity_m_s']**2+sound**2*math.log(point['number_density_m3']/density)-gm/point['radius_m']for point in profile]
    return {'name':case['name'],'settings':settings,'critical_radius_m':sonic,'analytic_critical_radius_m':gm/(2*sound**2),'sound_speed_m_s':sound,'lambda_base':lambda_base,'sonic_radius_ratio':sonic/base,'sonic_knudsen':sonic_kn,'wind_mass_loss_kg_s':wind_rate,'energy_limited_mass_loss_kg_s':energy_limited,'available_power_w':available,'required_power_to_outer_radius_w':wind_rate*mechanical,'required_to_available_power':wind_rate*mechanical/available,'relative_mass_flux_spread':(max(fluxes)-min(fluxes))/wind_rate,'bernoulli_spread_over_sound_speed_squared':(max(bernoulli)-min(bernoulli))/sound**2,'validity':{'subsonic_base':sonic>base,'collisional_sonic_region':sonic_kn<.1,'power_to_outer_radius_sufficient':wind_rate*mechanical<=available,'strong_binding_jeans':exobase is not None and exobase['lambda']>=6},'exobase':exobase,'profile':profile}


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--input',type=Path,default=DEFAULT_INPUT)
    parser.add_argument('--out',type=Path,default=ROOT/'docs/astronomy/ac13/wind-replay.json' if REPOSITORY_INPUT.exists() else HERE/'wind-replay-rerun.json')
    args=parser.parse_args()
    raw=args.input.read_bytes();data=json.loads(raw)
    refinement=[{'radius_over_sonic':x,'absolute_log_mach_difference':abs(transonic_log_mach(x)-transonic_log_mach(x,.00075))}for x in [.03,.1,.5,.9,1.1,2,12]]
    failures={'naive_critical_slope':'0/0 at r_s and u=c; regular accelerating slope is c/r_s','buried_base':'formal supersonic solution at the reservoir does not attach to a subsonic atmosphere','collisionless_sonic':'fluid transport fails before the sonic point; a Jeans illustration uses its separately located strong-binding exobase','insufficient_power':'XUV alone cannot maintain the prescribed-temperature mechanical energy increase; no temperature is silently adjusted'}
    settings=data['cases'][0]['settings'];constants=data['constants']
    gm=settings['planetMassEarth']*constants['earth_gm_m3_s2'];base=settings['baseRadiusEarth']*constants['earth_radius_m']
    sound=math.sqrt(constants['boltzmann_j_per_k']*settings['temperatureKelvin']/constants['particle_mass_kg']);sonic=gm/(2*sound**2)
    log_base=transonic_log_mach(base/sonic);velocity_base=sound*math.exp(log_base)
    boundary_comparison=[]
    for ratio in [4,12,24,48,96]:
        log_outer=transonic_log_mach(ratio*base/sonic)
        mechanical=.5*sound**2*(math.exp(2*log_outer)-math.exp(2*log_base))+gm/base*(1-1/ratio)
        expansion_heat=sound**2*(2*math.log(ratio)+log_outer-log_base)
        boundary_comparison.append({'outer_radius_ratio':ratio,'mechanical_heat_j_kg':mechanical,'independent_isothermal_expansion_heat_j_kg':expansion_heat,'relative_first_law_difference':abs(mechanical/expansion_heat-1)})
    result={'outer_boundary_comparison':boundary_comparison,'step_refinement':refinement,'failed_regimes':failures,'schema':'ac13-independent-ode-v1','kind':data['kind'],'input_sha256':hashlib.sha256(raw).hexdigest(),'method':'bracketed critical radius; independent RK4 log-Mach ODE from regular sonic series; flux-weighted Jeans tail Simpson quadrature','scope':'prescribed-temperature single-component synthetic atmosphere; no heating/chemistry/evolution closure; not an observational reduction','constants':data['constants'],'results':[analyze(case,data['constants'])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')
    print(json.dumps([{key:value for key,value in entry.items() if key not in ['profile','settings']}for entry in result['results']],indent=2))


if __name__=='__main__':main()