"""Predeclared paired injection-recovery pilot. All input arrays are public TESS data."""
import argparse, csv, hashlib, json, math
from pathlib import Path
import numpy as np
from scipy.signal import fftconvolve, savgol_filter
from scipy.ndimage import maximum_filter1d
from astropy.io import fits

DT=10/(60*24)
HALF=216
GRID=np.arange(-HALF,HALF+1)*DT
WINDOW=len(GRID)
DURATIONS=np.array([3,4.5,6,9,12,18])/24
LAMBDAS=[0.25,0.5,1.0]
METHODS=['baseline','gap_only']+['control_'+str(v) for v in LAMBDAS]
SENTINEL=-1e100
G=6.67430e-11; MSUN=1.98847e30; RSUN=6.957e8; C=299792458.

def seed_for(tic):
    return int(hashlib.sha256(('sj-inject-v1:'+str(tic)).encode()).hexdigest()[:16],16)%2**32

def chord(t,duration):
    x=2*t/duration; inside=np.abs(x)<1
    out=np.zeros_like(t);out[inside]=.4+.6*np.sqrt(np.maximum(0,1-x[inside]**2))
    return out

def flare(t,width):
    x=t/width; y=np.zeros_like(t); rise=(x>=-1)&(x<0); decay=x>=0
    q=x[rise];y[rise]=1+1.941*q-.175*q**2-2.246*q**3-1.125*q**4
    y[decay]=.689*np.exp(-1.600*x[decay])+.303*np.exp(-.2783*x[decay])
    return np.maximum(0,y)

def orthonormal_kernel(y):
    z=y-y.mean();z=z-GRID*np.dot(GRID,z)/np.dot(GRID,GRID)
    norm=np.linalg.norm(z)
    assert norm>0
    return z/norm

LENS=[orthonormal_kernel(chord(GRID,d)) for d in DURATIONS]
ALT=[orthonormal_kernel(flare(GRID,w/24)) for w in [.25,.5,1,2,4,8]]
impulse=np.zeros(WINDOW);impulse[HALF]=1;ALT.append(orthonormal_kernel(impulse))
for tau in [.5,2,6,12]:
    pulse=np.zeros(WINDOW);mask=GRID>=0;pulse[mask]=np.exp(-GRID[mask]/(tau/24))
    ALT.extend([orthonormal_kernel(pulse),orthonormal_kernel(-pulse)])

def read_curve(path):
    with fits.open(path,memmap=False) as h:
        d=h[1].data
        # Retain raw columns until after injection; filtering is deterministic and signal-independent.
        return {k:np.asarray(d[k],dtype=np.float64).copy() for k in ['TIME','PDCSAP_FLUX','QUALITY']}

def preprocess(curve,addition=None):
    t=curve['TIME'];f=curve['PDCSAP_FLUX'].copy()
    if addition is not None:f=f+addition
    mask=np.isfinite(t)&np.isfinite(f)&(curve['QUALITY']==0)
    t=t[mask];f=f[mask];origin=t.min()
    ix=np.floor((t-origin)/DT+1e-8).astype(int);n=int(ix.max()+1)
    counts=np.bincount(ix,minlength=n);observed=counts>0
    binned=np.divide(np.bincount(ix,weights=f,minlength=n),counts,out=np.zeros(n),where=observed)
    binned=binned/np.median(binned[observed])-1
    grid=origin+(np.arange(n)+.5)*DT
    filled=np.interp(np.arange(n),np.flatnonzero(observed),binned[observed])
    trend=savgol_filter(filled,WINDOW,2,mode='interp')
    y=np.where(observed,binned-trend,0.)
    adjacent=observed[1:]&observed[:-1];dif=np.diff(y)[adjacent]
    sigma=np.median(np.abs(dif-np.median(dif)))/.67448975/np.sqrt(2)
    if not np.isfinite(sigma) or sigma<=0:raise ValueError('Degenerate noise estimate')
    coverage=fftconvolve(observed.astype(float),np.ones(WINDOW)/WINDOW,mode='same')
    endpoint=np.arange(n);within=(endpoint>=HALF)&(endpoint<n-HALF)
    return grid,y,sigma,coverage,within

def positive_score(y,sigma,kernels):
    scores=np.vstack([fftconvolve(y,k[::-1],mode='same')/sigma for k in kernels])
    return np.maximum(0,scores.max(axis=0))**2

def scan(curve,addition=None):
    t,y,sigma,cov,within=preprocess(curve,addition)
    lens=positive_score(y,sigma,LENS)
    competitor=positive_score(y,sigma,ALT)
    competitor=maximum_filter1d(competitor,size=73,mode='constant',cval=0) # +/- six hours
    base=within&(cov>=.8-1e-10);strict=within&(cov>=.95-1e-10)
    scores={'baseline':np.where(base,lens,SENTINEL),'gap_only':np.where(strict,lens,SENTINEL)}
    for lam in LAMBDAS:scores['control_'+str(lam)]=np.where(strict,lens-lam*competitor,SENTINEL)
    return t,scores,{'sigma':float(sigma),'baseline_search_bins':int(base.sum()),'controlled_search_bins':int(strict.sum())}

def physical_pulse(times,center,mass,period,radius,b):
    host_mass=radius**1.25
    a=(G*(host_mass+mass)*MSUN*(period*86400/(2*np.pi))**2)**(1/3)
    speed=2*np.pi*a/(period*86400)
    re2=4*G*mass*MSUN*a/C**2
    wd_radius=.01125*RSUN*np.sqrt((1.454/mass)**(2/3)-(mass/1.454)**(2/3))
    prefactor=(2*re2-wd_radius**2)/(radius*RSUN)**2
    half=radius*RSUN*np.sqrt(1-b*b)/speed/86400
    # Five midpoint samples across the 120-second integration.
    signal=np.zeros(len(times))
    for offset in (np.arange(5)+.5)/5-.5:
        chordx=(times+offset*120/86400-center)*86400*speed/(radius*RSUN)
        r2=b*b+chordx*chordx;inside=np.isfinite(r2)&(r2<1)
        signal[inside]+=prefactor*(.4+.6*np.sqrt(1-r2[inside]))/.8/5
    peak=prefactor*(.4+.6*np.sqrt(1-b*b))/.8
    return signal,dict(duration_hours=float(2*half*24),peak_ppm=float(peak*1e6),impact_parameter=float(b),einstein_to_source=float(np.sqrt(re2)/(radius*RSUN)))

def threshold(background,alpha):
    return float(np.quantile(background,1-alpha,method='linear'))

def recovery(injected,local_null,cut):
    return (injected>cut)&(local_null<=cut)

def evaluate_arrays(background,inj,local,method,alpha=.1):
    cut=threshold(background[:,method],alpha)
    rec=recovery(inj[:,:,method],local[:,:,method],cut)
    return dict(threshold=cut,background_count=int(np.sum(background[:,method]>cut)),background_fraction=float(np.mean(background[:,method]>cut)),recovered=int(rec.sum()),recovery_pct=float(100*rec.mean()),raw_recovery_pct=float(100*np.mean(inj[:,:,method]>cut)))

def summarize(records,rows,output):
    n=len(records);background=np.array([[r['background'][m] for m in METHODS] for r in records])
    inj=np.array([[[q['scores'][m] for m in METHODS] for q in rows if q['tic']==r['tic']] for r in records])
    local=np.array([[[q['local_null'][m] for m in METHODS] for q in rows if q['tic']==r['tic']] for r in records])
    assert inj.shape==(n,9,len(METHODS))
    dev=np.array([r['split']=='development' for r in records]);ev=~dev
    dev_metrics={m:evaluate_arrays(background[dev],inj[dev],local[dev],j) for j,m in enumerate(METHODS)}
    selected=max(range(2,len(METHODS)),key=lambda j:(dev_metrics[METHODS[j]]['recovery_pct'],-LAMBDAS[j-2]))
    matched={m:evaluate_arrays(background[ev],inj[ev],local[ev],j) for j,m in enumerate(METHODS)}
    fixed={}
    for j,m in enumerate(METHODS):
        cut=dev_metrics[m]['threshold'];r=recovery(inj[ev,:,j],local[ev,:,j],cut)
        fixed[m]=dict(threshold=cut,background_count=int(np.sum(background[ev,j]>cut)),background_fraction=float(np.mean(background[ev,j]>cut)),recovered=int(r.sum()),recovery_pct=float(r.mean()*100))
    # The selected control is conditioned on the independent development split.
    rng=np.random.default_rng(20261023);bkg=background[ev];si=inj[ev];sn=local[ev];boot=[]
    for _ in range(2000):
        idx=rng.integers(0,len(bkg),len(bkg));r=[]
        for j in [0,selected]:
            cut=threshold(bkg[idx,j],.1);r.append(recovery(si[idx,:,j],sn[idx,:,j],cut).mean())
        boot.append(100*(r[1]-r[0]))
    interval=np.quantile(boot,[.025,.975]).tolist()
    operating={}
    for alpha in [.05,.1,.2]:
        operating[str(alpha)]={m:evaluate_arrays(background[ev],inj[ev],local[ev],j,alpha) for j,m in [(0,'baseline'),(selected,METHODS[selected])]}
    by_parameter={}
    evrows=[q for q in rows if q['split']=='evaluation']
    for param in ['mass_msun','period_days']:
        by_parameter[param]=[]
        for val in sorted(set(r[param] for r in evrows)):
            subset=[q for q in evrows if q[param]==val];d={'value':val,'injections':len(subset)}
            for j in [0,selected]:
                method=METHODS[j];cut=matched[method]['threshold'];d[method]=100*np.mean([q['scores'][method]>cut and q['local_null'][method]<=cut for q in subset])
            by_parameter[param].append(d)
    result=dict(targets=n,development_targets=int(dev.sum()),evaluation_targets=int(ev.sum()),injections=len(rows),evaluation_injections=int(ev.sum()*9),selected_method=METHODS[selected],selected_lambda=LAMBDAS[selected-2],development=dev_metrics,evaluation_matched=matched,evaluation_fixed=fixed,difference_pp=matched[METHODS[selected]]['recovery_pct']-matched['baseline']['recovery_pct'],difference_ci95_pp=interval,bootstrap_draws=2000,operating_points=operating,by_parameter=by_parameter,injection_duration_hours_range=[min(q['duration_hours'] for q in rows),max(q['duration_hours'] for q in rows)],injection_peak_ppm_range=[min(q['peak_ppm'] for q in rows),max(q['peak_ppm'] for q in rows)],einstein_ratio_range=[min(q['einstein_to_source'] for q in rows),max(q['einstein_to_source'] for q in rows)])
    result['baseline_recovered']=matched['baseline']['recovered']
    result['control_recovered']=matched[METHODS[selected]]['recovered']
    result['baseline_background_count']=matched['baseline']['background_count']
    result['control_background_count']=matched[METHODS[selected]]['background_count']
    # Presentation-only rounded summary, with integer counts retained above.
    result['display']={'baseline_pct':round(matched['baseline']['recovery_pct'],1),'control_pct':round(matched[METHODS[selected]]['recovery_pct'],1),'difference_pp':round(result['difference_pp'],1),'ci_low_pp':round(interval[0],1),'ci_high_pp':round(interval[1],1),'background_pct':round(100*matched['baseline']['background_fraction'],1),'fixed_baseline_background_pct':round(100*fixed['baseline']['background_fraction'],1),'fixed_control_background_pct':round(100*fixed[METHODS[selected]]['background_fraction'],1),'fixed_baseline_recovery_pct':round(fixed['baseline']['recovery_pct'],1),'fixed_control_recovery_pct':round(fixed[METHODS[selected]]['recovery_pct'],1)}
    (output/'R1.json').write_text(json.dumps(result,indent=2,allow_nan=False)+'\n')
    (output/'bootstrap_difference_pp.json').write_text(json.dumps(boot))
    return result

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--fits-dir',type=Path,default=Path('data/fits'));args=parser.parse_args()
    sample=json.loads(Path('data/sample.json').read_text());out=Path('results');out.mkdir(exist_ok=True)
    records=[];rows=[]
    for index,s in enumerate(sample):
        path=args.fits_dir/s['name'];raw=path.read_bytes()
        assert 'sha256:'+hashlib.sha256(raw).hexdigest()==s['sha256'],s['tic']
        curve=read_curve(path);t,null,meta=scan(curve)
        records.append(dict(tic=s['tic'],split=s['split'],background={m:float(v.max()) for m,v in null.items()},**meta))
        eligible=np.isfinite(curve['TIME'])&np.isfinite(curve['PDCSAP_FLUX'])&(curve['QUALITY']==0)
        valid_time=curve['TIME'][eligible];centers=valid_time[(valid_time>=valid_time.min()+1.5)&(valid_time<=valid_time.max()-1.5)]
        rng=np.random.default_rng(seed_for(s['tic']));scale=np.median(curve['PDCSAP_FLUX'][eligible])
        for mass in [.3,.6,1.]:
            for period in [100,300,1000]:
                center=float(rng.choice(centers));b=float(rng.uniform(0,.8))
                signal,phys=physical_pulse(curve['TIME'],center,mass,period,s['radius'],b)
                ti,score,_=scan(curve,scale*signal)
                assert np.array_equal(t,ti)
                near=np.abs(t-center)<=max(20/(60*24),phys['duration_hours']/24/4)
                row=dict(tic=s['tic'],split=s['split'],mass_msun=mass,period_days=period,center_btjd=center,**phys,scores={m:float(v[near].max()) for m,v in score.items()},local_null={m:float(v[near].max()) for m,v in null.items()})
                rows.append(row)
        print('processed',index+1,'/',len(sample),'TIC',s['tic'],flush=True)
    (out/'target_scores.json').write_text(json.dumps(records,indent=2,allow_nan=False))
    (out/'injections.json').write_text(json.dumps(rows,indent=2,allow_nan=False))
    result=summarize(records,rows,out)
    write_csv(rows,out)
    plot_results(result,out)
    print(json.dumps(result['display'],indent=2))

def write_csv(rows,out):
    flat=[]
    for row in rows:
        q={k:v for k,v in row.items() if k not in ['scores','local_null']}
        for group in ['scores','local_null']:
            for m,v in row[group].items():q[group+'_'+m]=v
        flat.append(q)
    with (out/'injections.csv').open('w') as f:
        writer=csv.DictWriter(f,fieldnames=list(flat[0]));writer.writeheader();writer.writerows(flat)

def plot_results(r,out):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    fig,axs=plt.subplots(1,2,figsize=(10,4.2),layout='constrained')
    chosen=r['selected_method'];colors=['#3c556b','#c26b37']
    for method,label,color in zip(['baseline',chosen],['Lensing template','Flare + artifact controls'],colors):
        a=r['by_parameter']['mass_msun'];axs[0].plot([x['value'] for x in a],[x[method] for x in a],marker='o',label=label,color=color)
        a=r['by_parameter']['period_days'];axs[1].plot([x['value'] for x in a],[x[method] for x in a],marker='o',color=color)
    axs[0].set_xlabel('Injected white-dwarf mass (solar masses)');axs[1].set_xlabel('Injected orbital period (days)');axs[1].set_xscale('log')
    for ax in axs:ax.set_ylim(0,100);ax.set_ylabel('New localized recovery (%)');ax.grid(alpha=.2)
    axs[0].legend(fontsize=8)
    fig.suptitle('Held-out TESS pilot at 10% empirical background triggers\n'+str(r['evaluation_targets'])+' stars; '+str(r['evaluation_injections'])+' paired injections',fontsize=11)
    fig.savefig(out/'recovery.png',dpi=160);plt.close(fig)

if __name__=='__main__':main()