"""The stars and light curves the study searches, from the TIC, Gaia DR3, and TESS's own lists.

Matches the selection of Vasilyev et al. (2024) as far as TESS allows: effective temperature
5000-6500 K, absolute G magnitude 4-6, RUWE below 1.4, no known eclipsing binary, no known
rotation period shorter than 20 days. Writes data/sample_stars.csv and data/sample_lcs.csv.
"""
import sys
import numpy as np
import pandas as pd

META = sys.argv[1] if len(sys.argv) > 1 else 'meta'
OUT = sys.argv[2] if len(sys.argv) > 2 else 'data'
TMAG_MAX = 10.5
A_G_PER_EBV = 2.74  # Casagrande & VandenBerg (2018), for solar-type stars


def main():
    p = pd.read_csv(f'{META}/parent_bright.csv.gz', low_memory=False, dtype={'GAIA': str, 'dr3_source_id': str})
    rot = pd.read_csv(f'{META}/gaia_rotmod.csv', dtype={'source_id': str})
    ebs = pd.read_csv(f'{META}/prsa2022_tessebs.tsv', sep='\t', comment='#', dtype=str)
    eb_tics = set(int(v) for v in ebs.iloc[:, 0] if str(v).strip().isdigit())
    lcs = pd.read_csv(f'{META}/lc_files_meta.csv.gz')

    s = p[p.dr3_source_id.notna()].copy()
    ebv = s['E(B-V)'].fillna(0.0).clip(lower=0.0)
    s['A_G'] = A_G_PER_EBV * ebv
    s['MG'] = s.phot_g_mean_mag + 5 * np.log10(s.parallax / 100.0) - s.A_G
    s['prot_gaia'] = s.dr3_source_id.map(dict(zip(rot.source_id, rot.best_rotation_period)))
    tmag = lcs.groupby('tic').Tmag.median()
    s['tmag_list'] = s.tic.map(tmag)

    steps = [
        ('Tmag below 10.5 in the TESS target lists', s.tmag_list < TMAG_MAX),
        ('TIC Teff 5000-6500 K', (s.Teff > 5000) & (s.Teff < 6500)),
        ('parallax over error above 10', s.parallax / s.parallax_error > 10),
        ('absolute G 4-6', (s.MG > 4) & (s.MG < 6)),
        ('RUWE below 1.4', s.ruwe < 1.4),
        ('TIC radius known', s.Rad.notna() & (s.Rad > 0)),
        ('not a Gaia DR3 eclipsing binary', s.best_class_name.fillna('') != 'ECL'),
        ('not a TESS eclipsing binary (Prsa et al. 2022)', ~s.tic.isin(eb_tics)),
        ('no Gaia DR3 rotation period under 20 days', ~(s.prot_gaia < 20)),
        ('Gaia DR3 vbroad below 15 km/s or none', ~(s.vbroad >= 15)),
    ]
    keep = np.ones(len(s), bool)
    log = []
    for name, cond in steps:
        keep &= cond.fillna(False).values
        log.append((name, int(keep.sum())))
    s = s[keep].copy()
    s['dev'] = (s.tic % 20 == 0)
    cols = ['tic', 'dr3_source_id', 'GAIA', 'tmag_list', 'Tmag', 'Teff', 'e_Teff', 'E_Teff', 'logg', 'Rad', 'e_Rad', 'E_Rad',
            'Mass', 'phot_g_mean_mag', 'bp_rp', 'parallax', 'parallax_error', 'A_G', 'MG', 'ruwe', 'non_single_star',
            'rv_chisq_pvalue', 'rv_amplitude_robust', 'vbroad', 'activityindex_espcs', 'activityindex_espcs_uncertainty',
            'best_class_name', 'prot_gaia', 'Rcont', 'ra', 'dec', 'dev']
    s = s[cols].sort_values('tic')
    l = lcs[lcs.tic.isin(set(s.tic))][['sector', 'tic', 'file', 'Camera', 'CCD']].sort_values(['tic', 'sector'])
    s.to_csv(f'{OUT}/sample_stars.csv', index=False)
    l.to_csv(f'{OUT}/sample_lcs.csv', index=False)
    for name, n in log:
        print(f'{n:7d}  {name}')
    print(f'{len(s)} stars, {len(l)} light curves; dev {int(s.dev.sum())} stars, '
          f'{int(l.tic.isin(set(s.tic[s.dev])).sum())} light curves')


if __name__ == '__main__':
    main()