"""Independent scalar finite-cylinder Monte Carlo. NumPy only.
Original implementation for this study; not CloudScat source.
See README.md and QA.json for validation scope and numerical limits.
"""
from dataclasses import dataclass,asdict
from pathlib import Path
import argparse,hashlib,json,math,time
import numpy as np
C=299792458.

def sha(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()

class Phase:
    def __init__(self,kind,mu=None,f=None,g=None):
        self.kind=kind
        if kind=="table":
            self.mu=np.asarray(mu,float); self.f=np.asarray(f,float)
            self.raw_integral=float(np.trapezoid(self.f,self.mu))
            self.f=self.f/self.raw_integral
            self.dx=np.diff(self.mu); self.slope=np.diff(self.f)/self.dx
            self.mass=.5*(self.f[:-1]+self.f[1:])*self.dx
            self.cdf=np.r_[0.,np.cumsum(self.mass)];self.cdf[-1]=1.
            # Exact first moment of each piecewise-linear segment.
            x=self.mu[:-1];d=self.dx;y=self.f[:-1];s=self.slope
            self.g=float(np.sum(x*y*d+(x*s+y)*d*d/2+s*d**3/3))
        else:
            self.g=float(g)
    def pdf_mu(self,x):
        x=np.clip(x,-1.,1.)
        if self.kind=="table":return np.interp(x,self.mu,self.f)
        g=self.g
        return (1-g*g)/(2*(1+g*g-2*g*x)**1.5)
    def pdf_sr(self,x):return self.pdf_mu(x)/(2*math.pi)
    def sample_mu(self,rng,n):
        return self.inverse(rng.random(n))
    def inverse(self,u):
        u=np.asarray(u)
        if self.kind!="table":
            g=self.g
            if abs(g)<1e-12:return 2*u-1
            return np.clip((1+g*g-((1-g*g)/(1-g+2*g*u))**2)/(2*g),-1,1)
        k=np.minimum(np.searchsorted(self.cdf,u,side="right")-1,len(self.dx)-1)
        k=np.maximum(k,0);v=u-self.cdf[k];y=self.f[k];s=self.slope[k]
        disc=np.maximum(0,y*y+2*s*v)
        den=y+np.sqrt(disc)
        d=np.divide(2*v,den,out=np.zeros_like(v,dtype=float),where=den>0)
        return np.clip(self.mu[k]+np.clip(d,0,self.dx[k]),-1,1)

def load_phases(lut,reff=10.,wavelength=337.1):
    z=np.load(lut)
    i=int(np.argmin(abs(z["reff_um"]-reff)))
    if abs(z["reff_um"][i]-reff)>1e-9:raise ValueError("Exact requested reff unavailable")
    a,b=z["wavelength_nm"];v=(wavelength-a)/(b-a)
    if not 0<=v<=1:raise ValueError("Not a bracketed interpolation")
    mu=np.cos(np.deg2rad(z["theta_deg"]))[::-1].copy()
    f=((1-v)*z["phase"][i,0,0]+v*z["phase"][i,1,0])[::-1].astype(float)
    tab=Phase("table",mu,f)
    meta={"input_file_name":Path(lut).name,"input_sha256":sha(lut),
      "original_full_nc_sha256":"2b37e27cffbbf0dd7e048df16c5fb612cd7a83e3c657c851bb50e758e11e870e",
      "reff_um":reff,"wavelength_nm":wavelength,"bracket_nm":[float(a),float(b)],
      "linear_weight_upper":float(v),"native_angle_nodes":len(mu),
      "raw_P11_integral_dmu":tab.raw_integral,
      "normalization":"f(mu)=P11/integral(P11 dmu); p_per_sr=f/(2pi)",
      "representation":"Linear interpolation of P11 in wavelength, then piecewise-linear probability density in mu; each inverse CDF segment solved quadratically.",
      "derived_g":tab.g,"scalar_SSA_assumed":1.,
      "input_source_url":"https://docs.hygeos.com/public.php/webdav/ic_baum_ghm_sol.nc",
      "registry_url":"https://raw.githubusercontent.com/hygeos/smartg/master/smartg/auxdata.py",
      "raw_LUT_redistributed":False,
      "license_status":"Original LUT is not included in release; permission for redistribution not established.",
      "physical_limits":"Randomly oriented severely rough bulk ice mixture; processed-LUT lineage and fresh-anvil PSD applicability remain uncertain. No vector polarization propagation or measured337nm absorption."}
    return {"ghm":tab,"hg":Phase("hg",g=tab.g)},meta

@dataclass
class Geometry:
    base_m:float=14000.
    top_m:float=17000.
    radius_m:float=30000.
    observer_z_m:float=550000.
    full_fov_deg:float=60.
    @property
    def thickness(self):return self.top_m-self.base_m

def intervals(pos,direction,geo):
    """Forward ray interval inside convex cylinder; no physical epsilon."""
    p=np.asarray(pos);u=np.asarray(direction);n=len(p)
    uz=u[:,2];z=p[:,2]
    moving=abs(uz)>1e-15
    az=np.full(n,-np.inf);bz=np.full(n,np.inf)
    t0=np.divide(geo.base_m-z,uz,out=np.zeros(n),where=moving)
    t1=np.divide(geo.top_m-z,uz,out=np.zeros(n),where=moving)
    az[moving]=np.minimum(t0,t1)[moving];bz[moving]=np.maximum(t0,t1)[moving]
    outz=(~moving)&((z<geo.base_m)|(z>geo.top_m));az[outz]=np.inf;bz[outz]=-np.inf
    aa=u[:,0]**2+u[:,1]**2
    bb=p[:,0]*u[:,0]+p[:,1]*u[:,1]
    cc=p[:,0]**2+p[:,1]**2-geo.radius_m**2
    disc=bb*bb-aa*cc
    radial=aa>1e-28
    ar=np.full(n,-np.inf);br=np.full(n,np.inf)
    good=radial&(disc>=0)
    sq=np.sqrt(np.maximum(disc,0))
    ar[good]=(-bb[good]-sq[good])/aa[good]
    br[good]=(-bb[good]+sq[good])/aa[good]
    bad=(radial&(disc<0))|((~radial)&(cc>0))
    ar[bad]=np.inf;br[bad]=-np.inf
    lo=np.maximum(0,np.maximum(az,ar));hi=np.minimum(bz,br)
    hit=hi>lo
    return lo,hi,hit

def segment_inside(pos,direction,length,geo):
    lo,hi,hit=intervals(pos,direction,geo)
    return np.where(hit,np.maximum(0,np.minimum(hi,length)-lo),0)

def rotate(u,mu,phi):
    n=len(u)
    a=np.sqrt(u[:,0]**2+u[:,1]**2)
    e=np.zeros((n,3));nz=a>1e-12
    e[nz,0]=-u[nz,1]/a[nz];e[nz,1]=u[nz,0]/a[nz];e[~nz,0]=1.
    f=np.cross(u,e)
    transverse=np.sqrt(np.maximum(0,1-mu*mu))
    out=mu[:,None]*u+transverse[:,None]*(np.cos(phi)[:,None]*e+np.sin(phi)[:,None]*f)
    # Roundoff only; physical directions should remain unit vectors.
    return out/np.linalg.norm(out,axis=1)[:,None]

def launch(rng,n,phase,importance):
    u=rng.uniform(-1,1,n);phi=rng.uniform(0,2*math.pi,n)
    if importance:
        choose=rng.random(n)<importance
        u[choose]=phase.sample_mu(rng,int(choose.sum()))
        weight=1/(1-importance+importance*4*math.pi*phase.pdf_sr(u))
    else:weight=np.ones(n)
    st=np.sqrt(np.maximum(0,1-u*u))
    return np.column_stack((st*np.cos(phi),st*np.sin(phi),u)),weight

def simulate(phase,source_m,total_tau,n,seed,geo=None,dt_s=2e-6,max_delay_s=.012,
             max_scatter=20000,importance=.5):
    """Per-emitted-photon plane-aperture fluence; vacuum exterior, SSA1.
    Initial physical source remains isotropic. Optional source-only proposal
    q=(1-a)/(4pi)+a*p_phase(mu_to_nadir) uses exact likelihood weights.
    Collision directions ALWAYS use the physical phase, identically evaluated
    in the next-event estimator. No biased walk or implicit absorption.
    """
    geo=geo or Geometry();start=time.perf_counter();rng=np.random.default_rng(seed)
    beta=total_tau/geo.thickness
    r0=geo.observer_z_m-source_m
    vacuum=1/(4*math.pi*r0*r0)
    src=np.array([[0.,0.,source_m]]);vertical=np.array([[0.,0.,1.]])
    direct_length=float(segment_inside(src,vertical,np.array([r0]),geo)[0])
    direct=math.exp(-beta*direct_length)
    nbins=int(round(max_delay_s/dt_s));hist=np.zeros(nbins+1)
    imp=importance if geo.base_m<source_m<geo.top_m else 0.
    dirs,weights=launch(rng,n,phase,imp)
    launch_sum=float(weights.sum())
    positions=np.tile(src,(n,1));lo,hi,hit=intervals(positions,dirs,geo)
    if beta==0:hit[:]=False
    escaped_initial=float(weights[~hit].sum())
    escaped_weight=escaped_initial;escaped_n=int((~hit).sum())
    escaped_up=float(weights[(~hit)&(dirs[:,2]>0)].sum())
    boundary_weights={"top":0.,"base":0.,"side":0.,"miss_or_vacuum":escaped_initial}
    ids=np.flatnonzero(hit)
    pos=positions[hit]+lo[hit,None]*dirs[hit]
    u=dirs[hit];w=weights[hit];path=lo[hit].copy()
    nee=np.zeros(n);collisions=0;iteration=0;min_delay=np.inf;negative=0
    first_hist=np.zeros(nbins+1)
    cone_weights={3.:0.,5.:0.,10.:0.}
    def cone_tally(positions,directions,ww):
        up=directions[:,2]>0
        if not np.any(up):return
        pp=positions[up];dd=directions[up];wt=ww[up]
        travel=(geo.observer_z_m-pp[:,2])/dd[:,2]
        xx=pp[:,0]+travel*dd[:,0];yy=pp[:,1]+travel*dd[:,1]
        rr=xx*xx+yy*yy
        for deg in cone_weights:
            radius=r0*math.tan(math.radians(deg))
            cone_weights[deg]+=float(wt[rr<=radius*radius].sum())
    max_dir_error=0.;weighted_path_exit=0.;exit_max_path=0.
    while len(ids) and iteration<max_scatter:
        iteration+=1
        _,exit_distance,ok=intervals(pos,u,geo)
        if not np.all(ok):
            raise RuntimeError("Interior ray lost cylinder intersection")
        free=rng.exponential(1/beta,len(ids))
        collide=free<exit_distance
        leaving=~collide
        if np.any(leaving):
            ww=w[leaving];endpos=pos[leaving]+exit_distance[leaving,None]*u[leaving]
            if iteration>1:cone_tally(endpos,u[leaving],ww)
            plen=path[leaving]+exit_distance[leaving]
            escaped_weight+=float(ww.sum());escaped_n+=int(leaving.sum())
            escaped_up+=float(ww[u[leaving,2]>0].sum())
            weighted_path_exit+=float(np.sum(ww*plen));exit_max_path=max(exit_max_path,float(plen.max()))
            top=abs(endpos[:,2]-geo.top_m)<1e-5
            base=(~top)&(abs(endpos[:,2]-geo.base_m)<1e-5)
            side=~(top|base)
            for label,mask in [("top",top),("base",base),("side",side)]:
                boundary_weights[label]+=float(ww[mask].sum())
        if not np.any(collide):
            ids=ids[:0];break
        ids=ids[collide];pos=pos[collide]+free[collide,None]*u[collide]
        path=path[collide]+free[collide];u=u[collide];w=w[collide]
        collisions+=len(ids)
        toward=-pos.copy();toward[:,2]+=geo.observer_z_m
        distance=np.linalg.norm(toward,axis=1);direction=toward/distance[:,None]
        cosine=np.einsum("ij,ij->i",u,direction)
        optical_length=segment_inside(pos,direction,distance,geo)
        aperture_cos=direction[:,2]
        accepted=aperture_cos>=math.cos(math.radians(geo.full_fov_deg/2))
        # Dimensionless relative to isotropic vacuum aperture fluence.
        contribution=w*4*math.pi*r0*r0*phase.pdf_sr(cosine)*aperture_cos/(distance*distance)*np.exp(-beta*optical_length)*accepted
        delay=(path+distance-r0)/C
        min_delay=min(min_delay,float(delay.min()))
        negative+=int(np.sum(delay < -1e-12))
        if negative:raise RuntimeError("Noncausal next-event contribution beyond roundoff")
        delay=np.maximum(delay,0)
        bins=np.minimum((delay/dt_s).astype(np.int64),nbins)
        part=np.bincount(bins,weights=contribution,minlength=nbins+1)
        hist+=part
        if iteration==1:first_hist+=part
        nee[ids]+=contribution
        mu=phase.sample_mu(rng,len(ids));phi=rng.uniform(0,2*math.pi,len(ids))
        u=rotate(u,mu,phi)
        max_dir_error=max(max_dir_error,float(np.max(abs(np.linalg.norm(u,axis=1)-1))))
    truncated_n=len(ids);truncated_weight=float(w.sum()) if truncated_n else 0.
    residual=launch_sum-escaped_weight-truncated_weight
    info={"seed":int(seed),"n_packets":n,"total_tau":float(total_tau),
      "overlying_tau":beta*direct_length,"source_altitude_m":source_m,
      "phase":phase.kind,"g":phase.g,"initial_importance_fraction":imp,
      "vacuum_transfer_m_inverse2":vacuum,
      "direct_relative_to_vacuum":direct,
      "scattered_relative_to_vacuum":float(hist.sum()/n),
      "total_relative_to_vacuum":direct+float(hist.sum()/n),
      "weighted_launch_mean":launch_sum/n,
      "weighted_escape_fraction":escaped_weight/n,"weighted_upward_escape_fraction":escaped_up/n,
      "weighted_truncation_fraction":truncated_weight/n,
      "energy_balance_residual_per_packet":residual/n,
      "escaped_packet_count":escaped_n,"truncated_packet_count":truncated_n,
      "boundary_escape_weights_per_packet":{k:v/n for k,v in boundary_weights.items()},
      "mean_collisions_per_packet":collisions/n,"max_scatter_iterations":iteration,
      "weighted_insideplus_entry_path_exit_mean_m":weighted_path_exit/n,
      "maximum_entryplus_cloud_path_m":exit_max_path,
      "min_scattered_delay_s":None if not np.isfinite(min_delay) else min_delay,
      "negative_delay_beyond_tolerance_count":negative,
      "max_direction_norm_error":max_dir_error,
      "response_overflow_relative_to_vacuum":float(hist[-1]/n),
      "first_scatter_relative_to_vacuum":float(first_hist.sum()/n),
      "finite_receiver_scattered_gain":{str(k):4*v/n/math.tan(math.radians(k))**2 for k,v in cone_weights.items()},
      "single_packet_standard_error_relative":float(np.std(nee,ddof=1)/math.sqrt(n)),
      "elapsed_s":time.perf_counter()-start}
    return info,hist/n

def phase_qa(phase,seed=821013,n=300000):
    rng=np.random.default_rng(seed);mu=phase.sample_mu(rng,n)
    # Independent Gauss integration of piecewise-linear input.
    if phase.kind=="table":
        x,w=np.polynomial.legendre.leggauss(8)
        xx=(phase.mu[1:]+phase.mu[:-1])[:,None]/2+np.diff(phase.mu)[:,None]*x/2
        ww=np.diff(phase.mu)[:,None]*w/2
        f=phase.pdf_mu(xx)
        integrals=[float(np.sum(ww*f)),float(np.sum(ww*f*xx)),float(np.sum(ww*f*(3*xx*xx-1)/2))]
    else:
        integrals=[1.,phase.g,phase.g**2]
    sample=[float(mu.mean()),float(np.mean((3*mu*mu-1)/2))]
    se=[float(mu.std(ddof=1)/math.sqrt(n)),float(np.std((3*mu*mu-1)/2,ddof=1)/math.sqrt(n))]
    u=np.linspace(.000001,.999999,10000);inv=phase.inverse(u)
    if phase.kind=="table":
        k=np.minimum(np.searchsorted(phase.mu,inv,side="right")-1,len(phase.dx)-1);k=np.maximum(0,k)
        dx=inv-phase.mu[k]
        back=phase.cdf[k]+phase.f[k]*dx+.5*phase.slope[k]*dx*dx
        inverse_error=float(np.max(abs(back-u)))
    else:
        g=phase.g
        back=(1-g*g)/(2*g)*(1/np.sqrt(1+g*g-2*g*inv)-1/(1+g))
        inverse_error=float(np.max(abs(back-u)))
    return {"normalization":integrals[0],"g_integral":integrals[1],"g_parameter":phase.g,
      "second_Legendre_moment":integrals[2],"sample_size":n,"sample_moments":sample,
      "sample_standard_errors":se,"sample_z_scores":[(sample[i]-integrals[i+1])/se[i]for i in range(2)],
      "inverse_CDF_max_abs_error":inverse_error,
      "sample_in_bounds":bool(np.all((mu>=-1)&(mu<=1))),
      "pass":bool(abs(integrals[0]-1)<1e-10 and inverse_error<1e-9 and all(abs((sample[i]-integrals[i+1])/se[i])<5 for i in range(2)))}

if __name__=="__main__":
    ap=argparse.ArgumentParser();ap.add_argument("--lut",required=True);ap.add_argument("--out",required=True)
    ap.add_argument("--pilot",action="store_true");args=ap.parse_args()
    out=Path(args.out);out.mkdir(exist_ok=True)
    phases,meta=load_phases(args.lut)
    q={"phase_input":meta,"phases":{k:phase_qa(v)for k,v in phases.items()}}
    (out/"phase_QA.json").write_text(json.dumps(q,indent=2),encoding="utf-8")
    cases=[]
    for name in ["ghm","hg"]:
        for altitude,tau in [(15000.,1.),(15000.,45.),(17000.,10.),(18000.,10.)]:
            info,h=simulate(phases[name],altitude,tau,5000,810022+len(cases))
            cases.append(info);print(json.dumps(info),flush=True)
    (out/"pilot_results.json").write_text(json.dumps(cases,indent=2),encoding="utf-8")