from pathlib import Path
import json,sys,hashlib,warnings
import numpy as np
from scipy.optimize import minimize_scalar,least_squares
from astropy.io import fits
import argparse
parser=argparse.ArgumentParser(description='Recompute the exact fixed SDSS coadd; needs astropy, NumPy and SciPy.')
parser.add_argument('fits',type=Path);parser.add_argument('--out',type=Path,default=Path('observed-output'));args=parser.parse_args()
root=Path(__file__).resolve().parent;out=args.out;out.mkdir(parents=True,exist_ok=False)
identity=json.loads((root/'input-manifest.json').read_text())['sdss']
registration=dict(sourceFITS=dict(sha256=identity['sha256']),identity=dict(specObjID=identity['identity']['specObjID'],redshiftCatalogue=2.694561),continuumRestWindowsAngstrom=[[1270,1290],[1315,1325],[1340,1360]],independentFitRelativeGate=1e-7,caveats='Fixed continuum and diagonal pixel noise; no LSF, covariance, metals/DLA/BAL rejection, gas-temperature or dark-matter inference.')
p=args.fits;assert p.stat().st_size==identity['bytes'] and hashlib.sha256(p.read_bytes()).hexdigest()==identity['sha256']
with warnings.catch_warnings(record=True)as caught:
 warnings.simplefilter('always')
 with fits.open(p,memmap=False,checksum=True)as hdus:
  table=hdus[1].data;catalogue=hdus[2].data[0]
  assert str(catalogue['SPECOBJID']).strip()==registration['identity']['specObjID']
  assert str(catalogue['CLASS']).strip()=='QSO'and int(catalogue['ZWARNING'])==0
  z=float(catalogue['Z']);assert abs(z-registration['identity']['redshiftCatalogue'])<1e-6
  arrays={name:np.array(table[name],copy=True)for name in ['flux','loglam','ivar','and_mask','or_mask','wdisp']}
assert not caught,[str(w.message)for w in caught]
wavelength=10**arrays['loglam'].astype(float);rest=wavelength/(1+z);flux=arrays['flux'].astype(float);ivar=arrays['ivar'].astype(float)
assert np.all(np.diff(wavelength)>0)
valid=np.isfinite(flux)&np.isfinite(rest)&np.isfinite(ivar)&(ivar>0)&(arrays['and_mask']==0)
continuum_window=np.zeros(len(rest),dtype=bool)
for lower,upper in registration['continuumRestWindowsAngstrom']:continuum_window|=(rest>=lower)&(rest<=upper)
selected=valid&continuum_window;assert selected.sum()>=30
x=rest[selected]/1280;y=flux[selected];weight=ivar[selected]
def amplitude(slope):
 basis=x**slope
 return np.dot(weight*basis,y)/np.dot(weight*basis,basis)
def objective(slope):return np.sum(weight*(y-amplitude(slope)*x**slope)**2)
profiled=minimize_scalar(objective,bounds=(-5,3),method='bounded',options={'xatol':1e-12})
assert profiled.success and -4.999<profiled.x<2.999
A=float(amplitude(profiled.x));assert A>0
independent=least_squares(lambda parameters:np.sqrt(weight)*(parameters[0]*x**parameters[1]-y),[A,profiled.x],bounds=([0,-5],[np.inf,3]),xtol=1e-13,ftol=1e-13,gtol=1e-13,max_nfev=2000)
assert independent.success
relative=np.max(np.abs(independent.x-[A,profiled.x])/np.maximum(1,np.abs([A,profiled.x])))
assert relative<registration['independentFitRelativeGate'],relative
continuum=A*(rest/1280)**float(profiled.x)
forest_window=(rest>=1040)&(rest<=1180);baseline=valid&forest_window;assert baseline.sum()>=50
rows=[]
def statistics(mask,factor=1.,tilt=0.):
 C=continuum[mask]*factor*(rest[mask]/1280)**tilt
 F=flux[mask]/C;noise=1/(ivar[mask]*C*C)
 mean=float(np.mean(F));variance=float(np.var(F,ddof=1));noise_mean=float(np.mean(noise))
 assert abs(variance-(np.sum(F*F)-len(F)*mean*mean)/(len(F)-1))<1e-12
 histogram,edges=np.histogram(F,bins=np.linspace(-.5,1.5,21))
 return dict(n=int(mask.sum()),meanTransmission=mean,sampleVariance=variance,conditionalDiagonalNoiseVariance=noise_mean,noiseSubtractedVariance=variance-noise_mean,effectiveOpticalDepth=float(-np.log(mean))if mean>0 else None,pdfEdges=edges.tolist(),pdfCounts=histogram.tolist(),belowPDF=int(np.sum(F<edges[0])),abovePDF=int(np.sum(F>edges[-1])),belowZero=int(np.sum(F<0)),aboveOne=int(np.sum(F>1)),minimum=float(np.min(F)),maximum=float(np.max(F)))
base=statistics(baseline)
for factor in [.95,1.,1.05]:
 row=statistics(baseline,factor=factor)
 assert abs(row['meanTransmission']-base['meanTransmission']/factor)<1e-12
 assert abs(row['sampleVariance']-base['sampleVariance']/factor**2)<1e-12
 if base['effectiveOpticalDepth']is not None:assert abs(row['effectiveOpticalDepth']-base['effectiveOpticalDepth']-np.log(factor))<1e-12
 rows.append(dict(normalizationFactor=factor,slopeOffset=0,statistics=row))
for tilt in [-.3,.3]:rows.append(dict(normalizationFactor=1,slopeOffset=tilt,statistics=statistics(baseline,tilt=tilt)))
conservative=baseline&(arrays['or_mask']==0)
mask_sensitivity=statistics(conservative)if conservative.sum()>=50 else dict(status='NOT_ENOUGH_PIXELS',n=int(conservative.sum()))
receipt=dict(status='PASS_BOUNDED_OBSERVED_FLUX_EXERCISE',sourceSHA256=registration['sourceFITS']['sha256'],redshift=z,coaddPixels=len(rest),continuumPixels=int(selected.sum()),forestPixels=int(baseline.sum()),excludedForestPixels=int(forest_window.sum()-baseline.sum()),continuumAmplitude=A,continuumSlope=float(profiled.x),independentFit=independent.x.tolist(),independentFitRelativeDifference=float(relative),continuumWeightedResidualSum=float(profiled.fun),continuumResidualDegreesOfFreedom=int(selected.sum())-2,baseline=base,sensitivities=rows,orMaskZeroSensitivity=mask_sensitivity,astropyWarnings=0,scope=registration['caveats'])
(out/'observed-statistics.json').write_text(json.dumps(receipt,indent=2)+'\n')
with (out/'observed-pixels.csv').open('w')as stream:
 stream.write('pixel,wavelength_vacuum_angstrom,rest_angstrom,flux_1e-17,ivar_reciprocal_flux_squared,and_mask,or_mask,continuum_1e-17,accepted_forest,accepted_continuum\n')
 for index in range(len(rest)):
  if 1000<=rest[index]<=1400:stream.write(','.join(map(str,[index,wavelength[index],rest[index],flux[index],ivar[index],int(arrays['and_mask'][index]),int(arrays['or_mask'][index]),continuum[index],int(baseline[index]),int(selected[index])]))+'\n')
print(json.dumps({key:value for key,value in receipt.items()if key not in ['sensitivities','baseline','orMaskZeroSensitivity']}))

expected=json.loads((root/'expected-observed.json').read_text())
for key in ['coaddPixels','continuumPixels','forestPixels','excludedForestPixels']:assert receipt[key]==expected[key],key
for key in ['continuumAmplitude','continuumSlope','continuumWeightedResidualSum']:assert abs(receipt[key]-expected[key])<1e-7*max(1,abs(expected[key])),key
for key in ['meanTransmission','sampleVariance','conditionalDiagonalNoiseVariance','noiseSubtractedVariance','effectiveOpticalDepth']:assert abs(base[key]-expected['baseline'][key])<1e-7,key
for key in ['belowZero','aboveOne','belowPDF','abovePDF','pdfCounts']:assert base[key]==expected['baseline'][key],key
print('PASS_FIXED_OBSERVED_RECONSTRUCTION')