"""Search many TESS light curves: download each from MAST's public bucket, search it, delete it.

python run_search.py <sample_lcs.csv> <sample_stars.csv> <out dir> <which: dev|main|audit|list:FILE>
                     [--workers N] [--threads N] [--keep-cands] [--inject] [--lcdir DIR]

With --lcdir the light curves are read from that directory instead (what verifiers do with the
files data/external.json points at). Results go to <out>/lcs.csv, cands.csv, dips.csv, inj.csv;
a rerun skips light curves already done.
"""
import argparse
import concurrent.futures as cf
import csv
import hashlib
import os
import re
import sys
import time

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import sflare as sf  # noqa: E402
import inject as ij  # noqa: E402

LC_COLS = ['file', 'tic', 'sector', 'camera', 'ccd', 'sha256', 'bytes', 'n_good', 'days_good', 'sigma_ppm', 'cdpp1',
           'crowdsap', 'rvar', 'p_rot', 'ls_power', 'ls_amp', 'n_cands', 'n_dips', 'k_erg_per_s']
CAND_COLS = ['file', 'tic', 'sector', 'camera', 'kind', 'i0', 'i1', 'ipk', 't0', 't1', 'tpk', 'n', 'n_core', 'amp', 'z_pk',
             'ed', 'snr', 'energy', 'fwhm', 'amp_fit', 'tpk_fit', 'ed_fit', 'chi2_flare', 'chi2_gauss', 'chi2_null',
             'dchi2', 'npts', 'flagged', 'dt_desat', 'dt_edge0', 'dt_edge1', 'bkg_ratio', 'bkg_z', 'sap_ratio', 'cen_dx', 'cen_dy', 'cen_z', 'cen_offset', 'cen_offset_z',
             'rise', 'decay', 'vetted', 'reject']
INJ_COLS = ['file', 'tic', 'sector', 'model', 'log_e', 'n_inj'] + [f'n_rec_{k}' for k in ij.THRESHOLD_NAMES] + ['n_rec_any', 'med_ratio']


def s3_url(f):
    m = re.search(r'-s(\d{4})-(\d{16})-', f)
    sec, tic = m.group(1), m.group(2)
    return (f'https://stpubdata.s3.amazonaws.com/tess/public/tid/s{sec}/'
            f'{tic[0:4]}/{tic[4:8]}/{tic[8:12]}/{tic[12:16]}/{f}')


def fetch(f, tmp):
    import requests
    p = os.path.join(tmp, f)
    for attempt in range(5):
        try:
            r = requests.get(s3_url(f), timeout=180)
            if r.status_code == 200:
                with open(p + '.part', 'wb') as fh:
                    fh.write(r.content)
                os.replace(p + '.part', p)
                return p
        except Exception:
            pass
        time.sleep(2 + 3 * attempt)
    return None


def process(path, star, resp, do_inject, seed):
    with open(path, 'rb') as fh:
        raw = fh.read()
    sha = hashlib.sha256(raw).hexdigest()
    lc, meta = sf.read_lc(path)
    fname = os.path.basename(path)
    k = sf.energy_per_ed(star['Teff'], star['Rad'], resp)
    row = dict(file=fname, tic=meta['tic'], sector=meta['sector'], camera=meta['camera'], ccd=meta['ccd'],
               sha256=sha, bytes=len(raw), n_good=len(lc['t']), days_good=len(lc['t']) * sf.CADENCE_S / sf.DAY_S,
               sigma_ppm=np.nan, cdpp1=meta['cdpp1'], crowdsap=meta['crowdsap'], rvar=np.nan, p_rot=np.nan,
               ls_power=np.nan, ls_amp=np.nan, n_cands=0, n_dips=0, k_erg_per_s=k)
    cands, inj = [], []
    if len(lc['t']) < 1000:
        return row, cands, inj
    trend, sigma, r, found = sf.search(lc)
    row['sigma_ppm'] = float(np.nanmedian(sigma) * 1e6)
    tn = trend / np.nanmedian(trend)
    row['rvar'] = float(np.nanpercentile(tn, 95) - np.nanpercentile(tn, 5))
    row['p_rot'], row['ls_power'], row['ls_amp'] = sf.rotation(lc['t'], trend)
    for kind, events, sign in (('flare', found, 1.0), ('dip', sf.candidates(lc['t'], -r, sigma), -1.0)):
        for c in events:
            c.update(sf.vet(lc, sign * r, sigma, c, meta['crowdsap']))
            c.update(file=fname, tic=meta['tic'], sector=meta['sector'], camera=meta['camera'], kind=kind,
                     energy=c['ed'] * k)
            cands.append(c)
    row['n_cands'] = sum(1 for c in cands if c['kind'] == 'flare')
    row['n_dips'] = sum(1 for c in cands if c['kind'] == 'dip')
    if do_inject:
        models = [('base', 1.0)]
        # Flares half and two and a half times as long, in one light curve in five
        if seed % 5 == 0:
            models += [('short', 0.5), ('long', 2.5)]
        for k_m, (name, scale) in enumerate(models):
            rng = np.random.default_rng(seed + k_m)
            for rec in ij.inject_recover(lc, k, rng, t_half_34=ij.T_HALF_34_D * scale, crowdsap=meta['crowdsap']):
                rec.update(file=fname, tic=meta['tic'], sector=meta['sector'], model=name)
                inj.append(rec)
    return row, cands, inj


def worker(args):
    path, star, resp, do_inject, seed, delete = args
    try:
        out = process(path, star, resp, do_inject, seed)
    except Exception as e:  # report and move on; the light curve is listed as failed
        out = ({'file': os.path.basename(path), 'error': repr(e)[:200]}, [], [])
    if delete:
        try:
            os.remove(path)
        except OSError:
            pass
    return out


def seed_of(fname):
    return int(hashlib.sha256(fname.encode()).hexdigest()[:8], 16)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('lcs'); ap.add_argument('stars'); ap.add_argument('out'); ap.add_argument('which')
    ap.add_argument('--workers', type=int, default=10); ap.add_argument('--threads', type=int, default=48)
    ap.add_argument('--inject', action='store_true'); ap.add_argument('--lcdir'); ap.add_argument('--resp', required=True)
    ap.add_argument('--keep', help='directory to keep light curves with flare candidates in')
    ap.add_argument('--limit', type=int)
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    stars = pd.read_csv(a.stars).set_index('tic')
    lcs = pd.read_csv(a.lcs)
    if a.which == 'dev':
        lcs = lcs[lcs.tic.isin(stars.index[stars.dev])]
    elif a.which == 'main':
        lcs = lcs[lcs.tic.isin(stars.index[~stars.dev])]
    elif a.which.startswith('list:'):
        lcs = lcs[lcs.file.isin(set(open(a.which[5:]).read().split()))]
    if a.limit:
        lcs = lcs.head(a.limit)
    done = set()
    lp = os.path.join(a.out, 'lcs.csv')
    if os.path.exists(lp):
        done = set(pd.read_csv(lp, usecols=['file']).file)
    todo = [f for f in lcs.file if f not in done]
    print(f'{len(lcs)} light curves, {len(todo)} to do', flush=True)
    tmp = os.path.join(a.out, 'tmp')
    os.makedirs(tmp, exist_ok=True)
    files = {name: open(os.path.join(a.out, name + '.csv'), 'a', newline='') for name in ('lcs', 'cands', 'inj', 'failed')}
    writers = {'lcs': csv.DictWriter(files['lcs'], LC_COLS, extrasaction='ignore'),
               'cands': csv.DictWriter(files['cands'], CAND_COLS, extrasaction='ignore'),
               'inj': csv.DictWriter(files['inj'], INJ_COLS, extrasaction='ignore'),
               'failed': csv.DictWriter(files['failed'], ['file', 'error'], extrasaction='ignore')}
    for name, w in writers.items():
        if files[name].tell() == 0:
            w.writeheader()
    tic_of = dict(zip(lcs.file, lcs.tic))
    t_start = time.time()
    n_done = 0
    with cf.ThreadPoolExecutor(a.threads) as dl, cf.ProcessPoolExecutor(a.workers) as pool:
        pending_dl, pending_cpu = {}, set()
        it = iter(todo)

        def feed():
            while len(pending_dl) + len(pending_cpu) < 3 * a.threads:
                try:
                    f = next(it)
                except StopIteration:
                    return
                if a.lcdir:
                    fut = cf.Future(); fut.set_result(os.path.join(a.lcdir, f))
                else:
                    fut = dl.submit(fetch, f, tmp)
                pending_dl[fut] = f

        feed()
        while pending_dl or pending_cpu:
            ready = [fu for fu in list(pending_dl) if fu.done()]
            for fu in ready:
                f = pending_dl.pop(fu)
                p = fu.result()
                if p is None:
                    writers['failed'].writerow({'file': f, 'error': 'download failed'})
                    continue
                star = stars.loc[tic_of[f]]
                pending_cpu.add(pool.submit(worker, (p, {'Teff': float(star.Teff), 'Rad': float(star.Rad)}, a.resp,
                                                     a.inject, seed_of(f), not a.lcdir and not a.keep)))
            done_cpu = [fu for fu in pending_cpu if fu.done()]
            for fu in done_cpu:
                pending_cpu.discard(fu)
                row, cands, inj = fu.result()
                if 'error' in row:
                    writers['failed'].writerow(row)
                    continue
                writers['lcs'].writerow(row)
                for c in cands:
                    writers['cands'].writerow(c)
                for r in inj:
                    writers['inj'].writerow(r)
                if a.keep and not a.lcdir:
                    src = os.path.join(tmp, row['file'])
                    if any(c['kind'] == 'flare' for c in cands):
                        os.makedirs(a.keep, exist_ok=True)
                        os.replace(src, os.path.join(a.keep, row['file']))
                    elif os.path.exists(src):
                        os.remove(src)
                n_done += 1
                if n_done % 500 == 0:
                    for fh in files.values():
                        fh.flush()
                    el = time.time() - t_start
                    print(f'{n_done} done, {el / 60:.1f} min, {n_done / el:.1f}/s', flush=True)
            feed()
            if not ready and not done_cpu:
                time.sleep(0.05)
    for fh in files.values():
        fh.close()
    print(f'finished {n_done} in {(time.time() - t_start) / 60:.1f} min', flush=True)


if __name__ == '__main__':
    main()