"""Cut from each target pixel file the cadences around its counted flares, keeping the columns
localize.py reads, so a re-run needs a few megabytes instead of the whole files.

python make_cutouts.py <counted_flares.csv> <tpf dir> <out dir> <sources.csv>
"""
import hashlib
import os
import sys

import numpy as np
import pandas as pd
from astropy.io import fits

CAD_D = 120.0 / 86400.0
PAD = 60 * CAD_D   # localize.py reads up to 45 cadences either side

flares, src, out, listing = sys.argv[1:5]
os.makedirs(out, exist_ok=True)
rows = []
for lc, g in pd.read_csv(flares).groupby('file'):
    name = lc.replace('_lc.fits', '_tp.fits')
    path = os.path.join(src, name)
    raw = open(path, 'rb').read()
    with fits.open(path, memmap=False) as h:
        d = h[1].data
        t = np.asarray(d['TIME'], float)
        keep = np.zeros(len(t), bool)
        for _, f in g.iterrows():
            keep |= (t >= f.t0 - PAD) & (t <= f.t1 + PAD)
        cols = fits.ColDefs([c for c in h[1].columns if c.name in ('TIME', 'FLUX', 'QUALITY')])
        table = fits.BinTableHDU.from_columns(cols, nrows=int(keep.sum()), header=h[1].header)
        for c in ('TIME', 'FLUX', 'QUALITY'):
            table.data[c] = d[c][keep]
        fits.HDUList([fits.PrimaryHDU(header=h[0].header), table, h[2].copy()]).writeto(os.path.join(out, name), overwrite=True)
    rows.append({'file': name, 'source_sha256': 'sha256:' + hashlib.sha256(raw).hexdigest(), 'source_bytes': len(raw),
                 'cadences_kept': int(keep.sum())})
pd.DataFrame(rows).to_csv(listing, index=False)
print(len(rows), 'cutouts')