# SPDX-License-Identifier: MIT
"""YS-EN-0003 · build OGLE-SLC: a StarLightCurves look-alike that keeps the period.

StarLightCurves stores period-folded, resampled light curves and drops the
period. We rebuild the same kind of dataset from the OGLE-IV Collection of
Variable Stars in the Large Magellanic Cloud, keep each star's catalogue
period, and match StarLightCurves' size and class mix exactly:

  1,329 classical Cepheids (fundamental mode + first overtone)
  2,580 eclipsing binaries
  5,327 RR Lyrae (RRab + RRc)
  -> 9,236 curves, 1,000 for training (stratified), 8,236 for test

Each curve: OGLE-IV I-band photometry, folded on the catalogue period with the
catalogue epoch at phase 0, kept in magnitudes (as StarLightCurves is:
its eclipses are peaks), smoothed onto 1,024 phase
points by circular Gaussian kernel regression (bandwidth 0.01 in phase), then
z-normalised. Stars with fewer than 100 good epochs are skipped.

  python build_ogle.py                  # -> ogle_slc.npz
  python build_ogle.py --max-epochs 250 # -> ogle_slc_e250.npz: the same stars, each
        # cut to a random 250 of its epochs, closer to OGLE-II's 200-300

Inputs (download first with get_ogle.sh, ~1.3 GB):
  https://www.astrouw.edu.pl/ogle/ogle4/OCVS/lmc/{cep,rrlyr,ecl}/phot.tar.gz
  and the parameter tables cepF.dat, cep1O.dat, RRab.dat, RRc.dat, ecl.dat
Output: ~/ys-work/ogle/ogle_slc.npz (X, y, period, ids, train mask)
"""
from __future__ import annotations

import argparse
import os
import tarfile
from pathlib import Path

import numpy as np

SRC = Path(os.environ.get("YS_WORK", Path.home() / "ys-work")) / "ogle"
L = 1024
BW = 0.01
COUNTS = {0: 1329, 1: 2580, 2: 5327}  # Cepheid, eclipsing binary, RR Lyrae (StarLightCurves order)
TABLES = {0: ["cep_cepF.dat", "cep_cep1O.dat"], 1: ["ecl_ecl.dat"], 2: ["rrlyr_RRab.dat", "rrlyr_RRc.dat"]}
TAR = {0: "cep_phot.tar.gz", 1: "ecl_phot.tar.gz", 2: "rrlyr_phot.tar.gz"}


def read_table(name):
    rows = []
    for line in open(SRC / name):
        t = line.split()
        try:
            # columns: ID, <I>, <V>, P [d], (sigma P), epoch ... (epoch is col 5 for ecl, col 6 for pulsators)
            p = float(t[3])
            ep = float(t[4]) if name.startswith("ecl") else float(t[5])
        except (IndexError, ValueError):
            continue  # no I-band epoch (32 bright cepF stars) or no period (22 ecl rows)
        rows.append((t[0], p, ep))
    return rows


def fold(t, mag, period, epoch):
    """Folded, smoothed, z-normalised brightness on L phase points."""
    # bring times and epoch to HJD-2450000 (ecl epochs are full HJD; photometry files may be either)
    epoch = epoch - 2450000 if epoch > 2400000 else epoch
    t = t - 2450000 if t[0] > 2400000 else t
    ph = ((t - epoch) / period) % 1.0
    b = mag  # magnitudes, like StarLightCurves (fainter = larger)
    grid = (np.arange(L) + 0.5) / L
    d = grid[:, None] - ph[None, :]
    d = (d + 0.5) % 1.0 - 0.5  # circular distance
    w = np.exp(-0.5 * (d / BW) ** 2)
    f = (w @ b) / np.maximum(w.sum(1), 1e-12)
    return (f - f.mean()) / max(f.std(), 1e-9)


def main(seed=0, max_epochs=0):
    rng = np.random.default_rng(seed)
    erng = np.random.default_rng(seed + 1)  # epoch thinning has its own stream, so star choice is unchanged
    X, y, P, ids = [], [], [], []
    for c, tables in TABLES.items():
        cat = [r for t in tables for r in read_table(t)]
        order = rng.permutation(len(cat))
        want = COUNTS[c]
        lookup = {r[0]: r for r in cat}
        chosen = [cat[i][0] for i in order]
        # One sequential pass over the archive (a .tar.gz cannot seek backwards
        # cheaply), keeping the photometry of the first 1.5x `want` stars in our
        # random order; then take, in that order, the first `want` with >= 100 epochs.
        cand = set(chosen[: int(want * 1.5)])
        phot = {}
        with tarfile.open(SRC / TAR[c], "r|gz") as tf:
            for m in tf:
                if "/I/" in m.name and m.name.endswith(".dat") and Path(m.name).stem in cand:
                    phot[Path(m.name).stem] = np.loadtxt(tf.extractfile(m), ndmin=2)
        got = 0
        for sid in chosen:
            if got == want:
                break
            if sid not in cand:
                raise SystemExit(f"class {c}: candidate set exhausted, widen it")
            a = phot.get(sid)
            if a is None or len(a) < 100:
                continue
            if max_epochs and len(a) > max_epochs:
                a = a[np.sort(erng.choice(len(a), max_epochs, replace=False))]
            _, p, ep = lookup[sid]
            X.append(fold(a[:, 0], a[:, 1], p, ep).astype(np.float32))
            y.append(c)
            P.append(p)
            ids.append(sid)
            got += 1
        print(f"class {c}: {got} curves (wanted {want})", flush=True)
    X, y, P, ids = np.array(X), np.array(y), np.array(P), np.array(ids)
    # stratified 1,000-curve training set in StarLightCurves' proportions (152 / 275 / 573)
    train = np.zeros(len(y), bool)
    for c, k in {0: 152, 1: 275, 2: 573}.items():
        idx = np.where(y == c)[0]
        train[rng.choice(idx, k, replace=False)] = True
    out = SRC / (f"ogle_slc_e{max_epochs}.npz" if max_epochs else "ogle_slc.npz")
    np.savez_compressed(out, X=X, y=y, period=P, ids=ids, train=train)
    print("saved", X.shape, "train", train.sum(), "test", (~train).sum())


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-epochs", type=int, default=0)
    main(max_epochs=ap.parse_args().max_epochs)