#!/usr/bin/env python3
"""Freeze deterministic WASP-126 windows from one hash-pinned original TESS LC FITS."""
import argparse
import csv
import hashlib
import json
from pathlib import Path
from astropy.io import fits
from astropy import __version__ as astropy_version
import numpy as np

SOURCE_URL = 'https://archive.stsci.edu/missions/tess/tid/s0001/0000/0000/2515/5310/tess2018206045859-s0001-0000000025155310-0120-s_lc.fits'
SOURCE_HASH = '12ddd623b04bdd354f27b5c89287b3a25a800d65cdf2733a1d0d04ab98e2d790'
COLUMNS = ['TIME','TIMECORR','CADENCENO','QUALITY','SAP_FLUX','SAP_FLUX_ERR','SAP_BKG','SAP_BKG_ERR','PDCSAP_FLUX','PDCSAP_FLUX_ERR','MOM_CENTR1','MOM_CENTR1_ERR','MOM_CENTR2','MOM_CENTR2_ERR','POS_CORR1','POS_CORR2']


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('source', type=Path)
    parser.add_argument('--outdir', type=Path, default=Path('.'))
    args = parser.parse_args()
    raw = args.source.read_bytes()
    if hashlib.sha256(raw).hexdigest() != SOURCE_HASH:
        raise ValueError('Original FITS hash does not match the pinned science input')
    args.outdir.mkdir(parents=True, exist_ok=True)
    with fits.open(args.source) as source:
        table, header = source[1].data, source[1].header
        windows = [{'index': index, 'center_btjd': 1327.52 + index * 3.2888,
                    'half_width_days': .28} for index in range(5)]
        rows = []
        for window in windows:
            selected = np.flatnonzero(abs(table['TIME'] - window['center_btjd']) <= window['half_width_days'])
            # Preserve every original row between explicit time-window endpoints, including null times/flags.
            for original in range(int(selected.min()), int(selected.max()) + 1):
                row = [original, window['index']]
                for key in COLUMNS:
                    value = table[key][original]
                    row.append('' if not np.isfinite(value) else int(value) if key in ['CADENCENO','QUALITY'] else float(value))
                rows.append(row)
        with (args.outdir / 'wasp126-tess-windows.csv').open('w', newline='') as stream:
            writer = csv.writer(stream, lineterminator='\n')
            writer.writerow(['source_row','window'] + [key.lower() for key in COLUMNS])
            writer.writerows(rows)
        products = []
        for index, filename in [(0,'tess-primary.header'),(1,'tess-lightcurve.header')]:
            info = source[index].fileinfo()
            data = raw[info['hdrLoc']:info['datLoc']]
            (args.outdir / filename).write_bytes(data)
            products.append({'file':filename,'source_byte_start':info['hdrLoc'],'bytes':len(data),'sha256':hashlib.sha256(data).hexdigest(),'scope':'exact original FITS header blocks only'})
        csv_raw = (args.outdir / 'wasp126-tess-windows.csv').read_bytes()
        products.append({'file':'wasp126-tess-windows.csv','bytes':len(csv_raw),'sha256':hashlib.sha256(csv_raw).hexdigest(),'rows':len(rows)})
        provenance = {'kind':'observed public TESS SPOC WASP-126 calibrated light-curve windows; not synthetic',
                      'freezer_versions':{'astropy':astropy_version,'numpy':np.__version__},'source_url':SOURCE_URL,'source_bytes':len(raw),'source_sha256':SOURCE_HASH,
                      'object':'TIC25155310 / WASP-126','sector':1,'data_release':int(source[0].header['DATA_REL']),
                      'processing_version':source[0].header['PROCVER'],'source_rows':len(table),
                      'selection':'five predetermined +/-0.28-day windows centered on1327.52+n*3.2888 BTJD; preserve all source rows between window endpoints; no flux/quality selection',
                      'windows':windows,'columns':COLUMNS,
                      'units':{key:source[1].columns[key].unit for key in COLUMNS},
                      'header':{key:header[key] for key in ['TIMESYS','TIMEREF','BJDREFI','BJDREFF','TIMEUNIT','TIMEPIXR','TIMEDEL','INT_TIME','READTIME','FRAMETIM','NUM_FRM','DEADC','CROWDSAP','FLFRCSAP']},
                      'handling':'TIME already barycentric; do not add TIMECORR. PDCSAP already corrects known crowding/aperture losses. SAP_BKG is retained pipeline-subtracted background, not another subtraction.',
                      'exposure_approximation':'nominal120s cadence-span uniform window; effective live time/onboard rejection are not exactly that kernel',
                      'permissions':'NASA TESS public science data via MAST; retain mission/SPOC/MAST acknowledgements and source data-use terms; no new license asserted for original observations',
                      'products':products}
        (args.outdir / 'observed-provenance.json').write_text(json.dumps(provenance, indent=2, allow_nan=False)+'\n')
    print('Frozen original rows:',len(rows),'bytes:',len(csv_raw))


if __name__ == '__main__':
    main()