#!python

import flammkuchen as fk
from glob import glob
import pandas as pd
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser(description="Summarize TLS results.")
parser.add_argument("indir", type=str)
parser.add_argument("-param", choices=['TOI','power_gls','Prot_gls','amp_gls','SDE_tls','Porb_tls'], default='SDE_tls', type=str)
args = parser.parse_args()

indir = args.indir
files = glob(indir + "/*.h5")
param = args.param
assert len(files) > 0, "no *.h5 files found!"

# print(files)
ss = []
errors = []
print("Reading *.tls files...")
for file in tqdm(files):
    try:
        data = fk.load(file)
        d = {
            "TOI": data.get("toiid"),
            "TIC": data.get("ticid"),
            "gaiaid": data.get("gaiaid"),
            "power_gls": data.get("power_gls")[0],
            "SDE_tls": data.get("SDE"),
            "Porb_tls": data.get("period"),
            "Prot_gls": data.get("Prot_gls")[0],
            "amp_gls": data.get("amp_gls")[0],
            # "Porb_tfop": data.get("period_tfop")[0],
            "depth": (1 - data.get("depth")) * 1e3,
            # "depth_tfop": data.get("depth_tfop")[0],
            "simbad_object": data.get("simbad_obj"),
            "filename": file,
        }
        s = pd.Series(d)
        ss.append(s)
    except Exception as e:
        errors.append(data.get("ticid"))
print(f"Sorting by {param}...")
# sort by sde
df = (
    pd.DataFrame(ss)
    .sort_values(
        by=param, ascending=False
    )
    .reset_index(drop=True)
    .drop_duplicates()
)
fp = indir + "_tls.csv"
df.to_csv(fp)
print(f"Saved: {fp}")
print("Error:", errors)
