"""Independent static force reference and fixed-slice diagnostics; no trajectory engine."""
import json,math
from pathlib import Path
import numpy as np
from scipy.optimize import minimize_scalar
KPC=3.0857e19;MYR=3.156e13;MSUN=1.989e30
G=6.67259e-11*MSUN*MYR**2/KPC**3
MB=1.3e10*(1+10**2)**1.5/10**3
MH=3e11*(1+.5**2)**1.5/.5**3
MD=6.5e10/(1-1/math.sqrt(17))

def potential(radius,z):
 return -G*(MB/math.sqrt(radius*radius+z*z+.4**2)+MH/math.sqrt(radius*radius+z*z+40**2)+MD/math.sqrt(radius*radius+(4+math.sqrt(z*z+.25**2))**2))

def restoring(radius,z):
 """Positive magnitude -g_z for positive height, in kpc/Myr^2."""
 height=math.sqrt(z*z+.25**2)
 return G*(MB*z/(radius*radius+z*z+.4**2)**1.5+MH*z/(radius*radius+z*z+40**2)**1.5+MD*(4+height)*z/(height*(radius*radius+(4+height)**2)**1.5))

def maximum_restoring(radius,samples=512):
 """Resolve every local maximum bracket found on a log height grid; check refinement."""
 heights=np.r_[0.,np.geomspace(1e-5,1e4,samples)]
 values=np.array([restoring(radius,z)for z in heights]);candidates=[]
 for index in range(1,len(heights)-1):
  if values[index]>=values[index-1]and values[index]>=values[index+1]:
   result=minimize_scalar(lambda z:-restoring(radius,z),bounds=(heights[index-1],heights[index+1]),method='bounded',options={'xatol':1e-12})
   if not result.success:raise RuntimeError('Static maximum bracket did not converge')
   candidates.append((float(-result.fun),float(result.x)))
 if not candidates:raise RuntimeError('No resolved restoring maximum')
 peak,height=max(candidates)
 # For all larger positive z, the acceleration is bounded by GM_total/z^2.
 if G*(MB+MH+MD)/heights[-1]**2>=peak:raise RuntimeError('Unbounded height search')
 return peak,height

def analyse(inputs):
 common=inputs['common'];radii=np.array(common['radiiKpc']);mass=np.array(common['massSolar']);area=np.array(common['areaKpc2'])
 conversion=MSUN/(KPC*MYR**2)
 criterion=[];max_refinement=0.
 for radius,m,a in zip(radii,mass,area):
  coarse=maximum_restoring(float(radius),256);fine=maximum_restoring(float(radius),512)
  error=abs(coarse[0]-fine[0])/fine[0];max_refinement=max(max_refinement,error)
  criterion.append({'radiusKpc':float(radius),'heightAtMaximumKpc':fine[1],'restoringKpcPerMyr2':fine[0],'initialColumnPressurePa':m/a*fine[0]*conversion,'nominalPeakColumnPressurePa':m/a*math.exp(-500/2000)*fine[0]*conversion})
 reports=[]
 for item in inputs['slices']:
  fields=np.array(item['fields']).reshape(64,13);radius=np.hypot(fields[:,0],fields[:,1]);z=fields[:,2];m=fields[:,9]
  inside=(radius<=16)&(np.abs(z)<=1)
  energy=np.sum(fields[:,3:6]**2,axis=1)/2+np.array([potential(float(r),float(h))for r,h in zip(radius,z)])
  resident=np.where(inside)[0];order=resident[np.argsort(radius[resident])]
  disk_mass=float(m[inside].sum());r90=float(radius[order[np.searchsorted(np.cumsum(m[order]),.9*disk_mass)]])
  reports.append({'q':item['q'],'frame':item['frame'],'modelMyr':item['timeMyr'],'outsideFraction':float(m[~inside].sum()/m.sum()),'unboundFraction':float(m[energy>0].sum()/m.sum()),'diskR90Kpc':r90,'diskSFRSolarPerYear':disk_mass/2e9,'globalSFRSolarPerYear':float(m.sum()/2e9),'maximumZKpc':float(z.max())})
 return {'schema':'ac42-static-reference-v1','criterion':criterion,'slices':reports,'max256to512PeakRelativeChange':max_refinement,'interpretation':'Static fixed-R maximum restoring acceleration with initial and nominal t=0 consumed columns. This is a quasi-static reference, not a finite-pulse stripping or escape boundary; trajectories move radially and consume gas. No trajectory or SSP is integrated.'}

def main():
 import argparse
 parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--inputs',type=Path,default=Path(__file__).parent/'inputs/fixed-slices.json');parser.add_argument('--out',type=Path,default=Path('recomputed.json'));args=parser.parse_args()
 result=analyse(json.loads(args.inputs.read_text()));
 if result['max256to512PeakRelativeChange']>1e-7:raise RuntimeError('Static maximum refinement gate failed')
 args.out.write_text(json.dumps(result,indent=2)+'\n')
 print(json.dumps({'status':'PASS_STATIC_REFERENCE_AND_FIVE_FIXED_SLICES','maxRefinement':result['max256to512PeakRelativeChange'],'out':str(args.out)}))
if __name__=='__main__':main()