"""Reproduce numerical summaries and figures from saved CloudScat batch curves."""
from pathlib import Path
import sys, csv, json, math
ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT.parents[2]/'work'/'calculation_python'))
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import timing_metrics as tm

OUT=ROOT/'analysis'; OUT.mkdir(exist_ok=True)
DT=2e-6; A=math.pi*(.025/2)**2; ETA=.20
def rows(path):
    with open(path,newline='') as f: return list(csv.DictReader(f))
def curve(path):
    return np.loadtxt(path,delimiter=',',skiprows=1)
def dumpcsv(path,rs):
    with open(path,'w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(rs[0])); w.writeheader();w.writerows(rs)
def main():
    raw=rows(ROOT/'results'/'batch_summary.csv')
    paths={r['case_id']:ROOT/'results' for r in raw}
    if (ROOT/'full_air_grid'/'summary.csv').exists():
        for f in rows(ROOT/'full_air_grid'/'summary.csv'):
            r=dict(raw[0]); r.update(case_id=f['case_id'],group='full_air_tau_grid',batch=f['batch'],N=f['N'],tau_over=f['tau_over'],**{'beta_m-1':str(float(f['tau_over'])/2000),'direct_fluence_m-2':f['direct_fluence'],'domainradius_m':'400000','domaintop_m':'120000'})
            raw.append(r); paths[r['case_id']]=ROOT/'full_air_grid'
    if (ROOT/'boundary_checks'/'summary.csv').exists():
        for f in rows(ROOT/'boundary_checks'/'summary.csv'):
            if f['case_id'] not in ('top_full_air','above_full_air'):continue
            r=dict(raw[0]);r.update(case_id=f['case_id'],group='full_air_top',batch=f['batch'],N=f['N'],zsrc_m='18000',ztop_m='18000' if f['case_id']=='top_full_air' else '17000',zbottom_m='12000' if f['case_id']=='top_full_air' else '11000',tau_over='0',**{'beta_m-1':'.005','direct_fluence_m-2':str(math.exp(-.054509901035)/(4*math.pi*532000**2))})
            raw.append(r);paths[r['case_id']]=ROOT/'boundary_checks'
    ids=list(dict.fromkeys(r['case_id'] for r in raw))
    summaries=[]; timing=[]; curves={}; allbatch={}
    for id in ids:
        rr=[r for r in raw if r['case_id']==id]
        if len(rr)!=4: continue
        r=rr[0]; hh=np.array([curve(paths[id]/f'{id}_batch{i}.csv')[:,1] for i in range(1,5)])
        h=hh.mean(axis=0); curves[id]=h; allbatch[id]=hh
        ff=hh.sum(axis=1)*DT; F=float(ff.mean()); se=float(ff.std(ddof=1)/2)
        z=float(r['zsrc_m']); vacuum=1/(4*math.pi*(550000-z)**2)
        sm=dict(case_id=id,group=r['group'],source_altitude_km=z/1000,depth_km=max(0,(float(r['ztop_m'])-z)/1000),tau_over=float(r['tau_over']),beta_m_inv=float(r['beta_m-1']),g=float(r['g']),N=sum(int(x['N']) for x in rr),fluence_m2_per_emitted_photon=F,fluence_se=se,total_to_vacuum=F/vacuum,total_to_vacuum_se=se/vacuum,direct_to_vacuum=float(r['direct_fluence_m-2'])/vacuum,direct_fraction_of_received=float(r['direct_fluence_m-2'])/F,aperture_fraction_25mm=F*A,photoelectrons_per_1e18_emitted=F*A*ETA*1e18)
        sm.update(tm.impulse_metrics(h,DT)); summaries.append(sm)
        for rise in (30.,80.):
            for fc in (0.,10000.,50000.):
                m=tm.metrics(h,DT,rise,fc if fc else None)
                t,y=tm.propagate_waveform(h,DT,rise,fc if fc else None)
                m['fraction_signal_first_100us']=float(y[t<=100e-6].sum()/y.sum())
                m['fraction_signal_first_200us']=float(y[t<=200e-6].sum()/y.sum())
                if fc==50000.:
                    bs=[tm.metrics(hi,DT,rise,fc)['rise_10_90_us'] for hi in hh]
                    m['rise_mc_se_us']=float(np.std(bs,ddof=1)/2)
                else:m['rise_mc_se_us']=None
                timing.append({**m,'case_id':id,'intrinsic_rise_us':rise,'fc_hz':fc})
    dumpcsv(OUT/'fluence_and_delay_summary.csv',summaries)
    dumpcsv(OUT/'pulse_timing_summary.csv',timing)
    data=dict(summary=summaries,timing=timing,assumptions=dict(dt_s=DT,record_s=.012,aperture_m2=A,efficiency=ETA))
    with open(OUT/'results.json','w') as f: json.dump(data,f,indent=2)
    plt.rcParams.update({'font.family':'DejaVu Sans','font.size':10,'axes.spines.top':False,'axes.spines.right':False,'axes.grid':True,'grid.alpha':.2,'figure.dpi':140})
    main=[r for r in summaries if r['group']=='full_air_tau_grid']; x=np.array([r['tau_over'] for r in main])
    fig,ax=plt.subplots(figsize=(8.2,4.2),layout='constrained')
    ax.errorbar(x,[r['total_to_vacuum'] for r in main],yerr=[r['total_to_vacuum_se'] for r in main],marker='o',capsize=3,label='All accepted arrivals, within 12 ms (mean +/- 1 SE)')
    ax.plot(x,[r['direct_to_vacuum'] for r in main],'--',color='#b35c20',label='Direct, unscattered contribution')
    ax.set_xscale('symlog',linthresh=.1);ax.set_xlim(0,120);ax.set_xlabel('Vertical overlying cloud extinction optical depth at 337.1 nm')
    ax.set_ylabel('Received fluence / isotropic vacuum fluence');ax.set_ylim(bottom=0);ax.legend(fontsize=8)
    ax.set_title('Nadir redistribution: 16 km source, 18 km cloud top, 550 km receiver')
    fig.savefig(OUT/'fluence_vs_optical_depth.png');plt.close(fig)
    fig,ax=plt.subplots(figsize=(8.2,4.3),layout='constrained')
    for rise,color in ((30.,'#147c9c'),(80.,'#9a3b71')):
        tr=[next(t for t in timing if t['case_id']==r['case_id'] and t['intrinsic_rise_us']==rise and t['fc_hz']==50000) for r in main]
        ax.plot(x,[t['rise_10_90_us'] for t in tr],'-o',color=color,label=f'{rise:g} us intrinsic rise + 50 kHz electronics')
    ax.axhline(200,ls='--',color='#b35c20',label='200 us reference')
    ax.set_xscale('symlog',linthresh=.1);ax.set_xlim(0,120);ax.set_xlabel('Overlying cloud optical depth');ax.set_ylabel('Observed 10-90% rising-edge time (us)');ax.legend(fontsize=8)
    ax.set_title('Finite pulse: assumed 3 ms exponential decay, no detector noise')
    fig.savefig(OUT/'rise_vs_optical_depth.png');plt.close(fig)
    fig,axes=plt.subplots(1,2,figsize=(9,3.7),layout='constrained')
    for id,label in [('full_tau_0p0','No cloud'),('full_tau_10p0','Overlying tau=10'),('full_tau_30p0','Overlying tau=30'),('full_tau_100p0','Overlying tau=100'),('top_full_air','At cloud top')]:
        if id not in curves:continue
        h=curves[id];mass=h*DT;cdf=np.cumsum(mass)/sum(mass)
        axes[0].plot(np.arange(len(h))*DT*1e6,cdf,label=label)
        t,y=tm.propagate_waveform(h,DT,30.,50000.)
        axes[1].plot(t*1e6,y/max(y),label=label)
    axes[0].set(xlim=(0,1500),ylim=(0,1.02),xlabel='Excess arrival delay (us)',ylabel='Fraction of received photons');axes[0].set_title('Transport impulse response')
    axes[1].set(xlim=(0,1500),ylim=(0,1.05),xlabel='Time since source onset (us)',ylabel='Waveform / its own peak');axes[1].set_title('30 us source rise + 50 kHz electronics')
    axes[0].legend(fontsize=7);fig.savefig(OUT/'arrival_and_pulse_shapes.png');plt.close(fig)
    # Sample-phase and pulse-front counts are conditional, noiseless examples.
    ss=[]
    for id in ('full_tau_0p0','top_full_air','above_full_air','full_tau_10p0','full_tau_100p0'):
        if id not in curves:continue
        for rise in (30.,80.):
            t,y=tm.propagate_waveform(curves[id],DT,rise,10000.)
            for fs in (10000.,50000.,100000.,200000.):
                ss.append(dict(case_id=id,intrinsic_rise_us=rise,fc_hz=10000.,fs_hz=fs,**tm.sampling_phase_metrics(y,DT,fs,fc_hz=10000.)))
    dumpcsv(OUT/'noiseless_sampling_phase.csv',ss)
    print('Completed cases:',len(summaries),'packets:',sum(r['N'] for r in summaries))
    for r in main:print('MAIN',r)
    for r in summaries:
        if r['group'] in ('same_tau10','top_above','full_air_top'):print('OTHER',r)
    for t in timing:
        if t['case_id'] in ('full_tau_0p0','full_tau_10p0','full_tau_30p0','full_tau_100p0','top_full_air','above_full_air') and t['fc_hz']==50000:print('TIMING',t['case_id'],t['intrinsic_rise_us'],t['rise_10_90_us'],t['onset_to_peak_us'],t['rise_mc_se_us'])

if __name__=='__main__':main()