utils/tle_dataset.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TLE daily-grid dataset v2: outlier-cleaned + space-weather channels.

Builds on the v1 daily-grid cumulative-phase formulation and adds the two
data-quality / physics levers from the cm-tle-pred benchmark analysis:

  (1) Robust outlier removal per satellite (tle_clean.clean_records) before
      resampling -- their single biggest accuracy lever.
  (2) Space-weather input channels (F10.7, F10.7-81d, Ap) -- exogenous drivers
      of atmospheric drag, i.e. the secular decay of mean motion.

Channels (CHANNEL_NAMES):
  orbital, PREDICTED + in loss (indices 0..5):
    bstar_slog, mean_motion, eccentricity, inclination_deg,
    draan_deg_per_day, dargp_deg_per_day
  solar, INPUT-ONLY context, excluded from loss (indices 6..8):
    f107, f107_81, ap
Absolute angles [MA, RAAN, argp] are kept as aux (anchor/truth), never predicted.

Angles are reconstructed by anchoring at the last observed TLE and integrating
the predicted rates (reconstruct_track) -- absolute phase is never predicted.
"""

from __future__ import annotations

import hashlib
import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional

import numpy as np
import torch
from torch.utils.data import Dataset
from tqdm import tqdm

_HERE = Path(__file__).resolve().parent
_CODE_UTILS = _HERE.parents[1] / "code" / "utils"
for p in (str(_HERE), str(_CODE_UTILS)):
    if p not in sys.path:
        sys.path.insert(0, p)

from tle_future_dataset import (  # noqa: E402
    collect_txt_files, iter_tle_pairs_from_txt, parse_tle_pair,
    parse_date_to_unix, signed_log1p, extract_year_from_filename,
)
from tle_clean import clean_records, cumulative_mean_anomaly  # noqa: E402
from space_weather import SpaceWeather, load_space_weather, SOLAR_FEATURES, N_SOLAR  # noqa: E402

# v3: persistence-RESIDUAL (drift) targets.
#
# v2 predicted absolute mean_motion / bstar. Because these barely change over the
# forecast window, the loss-optimal prediction was ≈persistence and the model
# could not beat SGP4 propagation at short/mid horizons. v3 makes the slowly
# DRIFTING magnitudes into per-day DELTA channels, so the training target IS the
# drift (zero-mean, small) — the learning signal we care about is amplified.
# Absolute mean_motion / bstar at the anchor live in aux, and the trajectory is
# reconstructed by anchor + cumulative predicted delta (like the angle rates).
#
# Eccentricity / inclination drift even less and are well predicted as absolutes,
# so they stay absolute. RAAN/argp remain secular rate channels (as in v2).
ORBITAL_NAMES = [
    "d_bstar_slog_per_day",   # 0  daily drift of signed_log1p(bstar)
    "d_mean_motion_per_day",  # 1  daily drift of mean motion (orbital decay rate)
    "eccentricity",           # 2  absolute
    "inclination_deg",        # 3  absolute
    "draan_deg_per_day",      # 4  RAAN secular rate
    "dargp_deg_per_day",      # 5  argp secular rate
]
N_ORBITAL = len(ORBITAL_NAMES)

# v4 = cm-tle-pred-style processing. The cm-tle-pred benchmark feeds, on top of the
# raw elements, a set of engineered physical features. We add the reproducible ones
# as INPUT-ONLY channels (excluded from the loss), mirroring their feature names:
#   - SAT_RX..VZ      : Cartesian ECI state from the day's elements (two-body
#                       osculating r,v -- same role as their SGP4 SAT_R*/V*)
#   - SEMIMAJOR_AXIS / PERIOD / APOAPSIS / PERIAPSIS : derived orbit geometry
#   - *_COS / *_SIN   : cyclical encodings of the angular elements (MA/RAAN/argp)
# What is PREDICTED is unchanged (the 6 orbital drift/rate channels) -- this only
# enriches Toto's context, exactly like the solar channels do.
PHYS_NAMES = [
    "sat_rx", "sat_ry", "sat_rz",      # 6..8   ECI position (km)
    "sat_vx", "sat_vy", "sat_vz",      # 9..11  ECI velocity (km/s)
    "semimajor_axis", "period_min",    # 12..13 (km), (minutes/rev)
    "apoapsis_alt", "periapsis_alt",   # 14..15 altitude (km)
    "ma_cos", "ma_sin",                # 16..17 mean anomaly cyclical
    "raan_cos", "raan_sin",            # 18..19 RAAN cyclical
    "argp_cos", "argp_sin",            # 20..21 argp cyclical
]
N_PHYS = len(PHYS_NAMES)

CHANNEL_NAMES = ORBITAL_NAMES + PHYS_NAMES + SOLAR_FEATURES   # 6 + 16 + 3 = 25
N_CHANNELS = len(CHANNEL_NAMES)
# aux carries the absolute anchors/truth NOT directly predicted:
# [mean_anomaly, RAAN, argp, mean_motion, bstar_slog]
N_AUX = 5
AUX_MA, AUX_RAAN, AUX_ARGP, AUX_MM, AUX_BSTAR = 0, 1, 2, 3, 4
# orbital channels are prediction targets; PHYS + solar are input-only context.
LOSS_CHANNEL_MASK = np.array([1] * N_ORBITAL + [0] * (N_PHYS + N_SOLAR), dtype=bool)
# Position-critical drift channels: d_bstar_slog (0) and d_mean_motion (1). SGP4
# along-track error is governed by mean motion and drag, so these are upweighted in
# the training loss (--drift-loss-weight) to align the objective with position.
DRIFT_CHANNELS = (0, 1)
SECONDS_PER_DAY = 86400.0
MU_EARTH = 398600.4418      # km^3/s^2  (WGS-72 gravitational parameter)
R_EARTH = 6378.137          # km        (equatorial radius, for apo/peri altitude)
LEO_MIN_MEAN_MOTION = 11.25  # rev/day; cm-tle-pred LEO cutoff
CHANNEL_VERSION = "v5cmtle_feats"


def signed_expm1(x: float) -> float:
    return math.copysign(math.expm1(abs(x)), x)


def _clip_ecc(x):
    return float(min(max(x, 0.0), 0.9999999))


def _clip_inc(x):
    return float(min(max(x, 0.0), 180.0))


def elements_from_feat_aux(feat_row, aux_row) -> Dict[str, float]:
    """Full physical elements for an OBSERVED grid row (truth/anchor).

    ecc/inc come from the absolute channels; bstar/mean_motion/angles from aux.
    """
    f = np.asarray(feat_row, dtype=np.float64)
    a = np.asarray(aux_row, dtype=np.float64)
    return {
        "bstar": signed_expm1(float(a[AUX_BSTAR])),
        "mean_motion_rev_per_day": float(max(a[AUX_MM], 1e-6)),
        "eccentricity": _clip_ecc(f[2]),
        "inclination_deg": _clip_inc(f[3]),
        "mean_anomaly_deg": float(a[AUX_MA] % 360.0),
        "raan_deg": float(a[AUX_RAAN] % 360.0),
        "argp_deg": float(a[AUX_ARGP] % 360.0),
    }


def reconstruct_track(anchor_aux, pred_seq, grid_step_days=1.0):
    """Absolute elements per forecast day, anchored at the last observed TLE.

    bstar(log) and mean_motion are integrated from their anchor (aux) by summing
    the predicted per-day DELTAS; mean-anomaly phase is integrated trapezoidally
    from the (reconstructed) mean motion; RAAN/argp from predicted daily rates;
    ecc/inc are read from the absolute channels. Nothing absolute-phase is predicted.
    """
    a = np.asarray(anchor_aux, dtype=np.float64)
    pred = np.asarray(pred_seq, dtype=np.float64)
    ma, raan, argp = a[AUX_MA], a[AUX_RAAN], a[AUX_ARGP]
    mm = float(a[AUX_MM]); bstar_slog = float(a[AUX_BSTAR])
    prev_mm = mm
    out = []
    for h in range(pred.shape[0]):
        bstar_slog += float(pred[h, 0]) * grid_step_days   # cumulative d_bstar
        mm += float(pred[h, 1]) * grid_step_days           # cumulative d_mean_motion
        mm_h = max(mm, 1e-6)
        ma += 360.0 * 0.5 * (prev_mm + mm_h) * grid_step_days
        prev_mm = mm_h
        raan += float(pred[h, 4]) * grid_step_days
        argp += float(pred[h, 5]) * grid_step_days
        out.append({
            "bstar": signed_expm1(bstar_slog),
            "mean_motion_rev_per_day": mm_h,
            "eccentricity": _clip_ecc(pred[h, 2]),
            "inclination_deg": _clip_inc(pred[h, 3]),
            "mean_anomaly_deg": ma % 360.0,
            "raan_deg": raan % 360.0,
            "argp_deg": argp % 360.0,
        })
    return out


# -----------------------------
# cm-tle-pred physical features (input-only channels)
# -----------------------------

def _kepler_rv(a, e, inc_r, raan_r, argp_r, M_r):
    """Vectorized two-body osculating ECI state from classical elements.

    Inputs are numpy arrays (one element per grid day); returns rx,ry,rz (km) and
    vx,vy,vz (km/s). This plays the same role as cm-tle-pred's SGP4 SAT_R*/V*
    Cartesian features but is fully vectorized over the daily grid (no per-day
    sgp4init), which matters for the all-years (~23 GB) build.
    """
    E = np.array(M_r, dtype=np.float64, copy=True)
    for _ in range(12):                      # Newton on Kepler's equation
        E = E - (E - e * np.sin(E) - M_r) / np.maximum(1.0 - e * np.cos(E), 1e-9)
    cosE, sinE = np.cos(E), np.sin(E)
    r = a * (1.0 - e * cosE)
    nu = np.arctan2(np.sqrt(np.maximum(1.0 - e * e, 0.0)) * sinE, cosE - e)
    cosnu, sinnu = np.cos(nu), np.sin(nu)
    p = np.maximum(a * (1.0 - e * e), 1e-6)
    rp_x, rp_y = r * cosnu, r * sinnu
    vfac = np.sqrt(MU_EARTH / p)
    vp_x, vp_y = -vfac * sinnu, vfac * (e + cosnu)
    co, so = np.cos(raan_r), np.sin(raan_r)
    ci, si = np.cos(inc_r), np.sin(inc_r)
    cw, sw = np.cos(argp_r), np.sin(argp_r)
    # perifocal -> ECI rotation R3(-raan) R1(-inc) R3(-argp)
    R11 = co * cw - so * sw * ci
    R12 = -co * sw - so * cw * ci
    R21 = so * cw + co * sw * ci
    R22 = -so * sw + co * cw * ci
    R31, R32 = sw * si, cw * si
    rx = R11 * rp_x + R12 * rp_y
    ry = R21 * rp_x + R22 * rp_y
    rz = R31 * rp_x + R32 * rp_y
    vx = R11 * vp_x + R12 * vp_y
    vy = R21 * vp_x + R22 * vp_y
    vz = R31 * vp_x + R32 * vp_y
    return rx, ry, rz, vx, vy, vz


def physics_features(g_mm, g_ecc, g_inc_deg, g_raan_deg, g_argp_deg, g_ma_deg):
    """(T, N_PHYS) cm-tle-pred input features from the daily-grid elements."""
    mm = np.maximum(np.asarray(g_mm, dtype=np.float64), 1e-6)
    n_rad_s = mm * 2.0 * math.pi / SECONDS_PER_DAY
    a = (MU_EARTH / (n_rad_s ** 2)) ** (1.0 / 3.0)             # semimajor axis (km)
    e = np.clip(np.asarray(g_ecc, dtype=np.float64), 0.0, 0.999999)
    inc_r = np.radians(g_inc_deg)
    raan_r = np.radians(np.asarray(g_raan_deg) % 360.0)
    argp_r = np.radians(np.asarray(g_argp_deg) % 360.0)
    ma_r = np.radians(np.asarray(g_ma_deg) % 360.0)
    rx, ry, rz, vx, vy, vz = _kepler_rv(a, e, inc_r, raan_r, argp_r, ma_r)
    period_min = 1440.0 / mm
    apo = a * (1.0 + e) - R_EARTH                              # apoapsis altitude (km)
    peri = a * (1.0 - e) - R_EARTH                            # periapsis altitude (km)
    return np.stack([
        rx, ry, rz, vx, vy, vz,
        a, period_min, apo, peri,
        np.cos(ma_r), np.sin(ma_r),
        np.cos(raan_r), np.sin(raan_r),
        np.cos(argp_r), np.sin(argp_r),
    ], axis=1)


def sat_split_of(norad_id, train_frac=0.70, valid_frac=0.15):
    """cm-tle-pred-style SATELLITE-level split (70/15/15): every record of a given
    satellite lands in the same split, deterministically by NORAD id."""
    h = int(hashlib.md5(str(int(norad_id)).encode()).hexdigest(), 16) % 1000
    if h < int(train_frac * 1000):
        return "train"
    if h < int((train_frac + valid_frac) * 1000):
        return "valid"
    return "test"


# -----------------------------
# Parse + clean (with progress)
# -----------------------------

def _is_physically_possible(r) -> bool:
    """cm-tle-pred 'data integrity' check: drop physically impossible records."""
    return (0.0 <= r.eccentricity < 1.0
            and 0.0 < r.inclination_deg < 180.0
            and 0.1 < r.mean_motion_rev_per_day < 20.0)


def load_and_clean(
    input_dir, years: Optional[Iterable[int]] = None, clean: bool = True,
    min_records_per_object: int = 2, drop_first_n: int = 5, leo_only: bool = True,
) -> Dict[int, List]:
    """Parse raw TLE txt (tqdm over files), group by NORAD, dedup, outlier-clean.

    v4 adds the cm-tle-pred 'data integrity' steps:
      - drop physically impossible records (impossible ecc/inc/mean_motion),
      - discard the first ``drop_first_n`` TLEs per satellite (least reliable),
      - keep only LEO objects (median mean_motion >= LEO_MIN_MEAN_MOTION) when
        ``leo_only`` is set.
    """
    files = collect_txt_files(input_dir, years=years)
    by_norad: Dict[int, List] = {}
    for path in tqdm(files, desc="parse files", unit="file"):
        year = extract_year_from_filename(path)
        ridx = 0
        for obj, l1, l2, ln in iter_tle_pairs_from_txt(path):
            ridx += 1
            rec = parse_tle_pair(l1, l2, source_file=str(path), source_year=year,
                                 source_line_number=ln, object_name_raw=obj, record_index=ridx)
            if rec is not None and _is_physically_possible(rec):
                by_norad.setdefault(rec.norad_id, []).append(rec)

    out: Dict[int, List] = {}
    removed_total = 0
    n_nonleo = 0
    for norad, recs in tqdm(by_norad.items(), desc="dedup+clean", unit="sat"):
        recs = sorted(recs, key=lambda r: r.epoch_unix)
        deduped, seen = [], set()
        for r in recs:
            key = round(r.epoch_unix, 3)
            if key in seen:
                continue
            seen.add(key)
            deduped.append(r)
        if drop_first_n > 0 and len(deduped) > drop_first_n:   # least-reliable early TLEs
            deduped = deduped[drop_first_n:]
        if leo_only:
            med_mm = float(np.median([r.mean_motion_rev_per_day for r in deduped])) if deduped else 0.0
            if med_mm < LEO_MIN_MEAN_MOTION:
                n_nonleo += 1
                continue
        if clean and len(deduped) >= 5:
            deduped, nrem = clean_records(deduped)
            removed_total += nrem
        if len(deduped) >= min_records_per_object:
            out[norad] = deduped
    print(f"[load] {len(out)} satellites, removed {removed_total} outlier records, "
          f"dropped {n_nonleo} non-LEO objects (leo_only={leo_only})")
    return out


# -----------------------------
# Daily-grid series + solar
# -----------------------------

@dataclass
class DailySeries:
    norad_id: int
    grid_epochs: np.ndarray   # (T,)
    feats: np.ndarray         # (T, N_CHANNELS)
    mask: np.ndarray          # (T,) orbital coverage
    aux: np.ndarray           # (T, N_AUX) absolute angles


def _build_one(recs, sw: Optional[SpaceWeather], grid_step_days=1.0) -> Optional[DailySeries]:
    if len(recs) < 2:
        return None
    recs = sorted(recs, key=lambda r: r.epoch_unix)
    ep = np.array([r.epoch_unix for r in recs], dtype=np.float64)
    bstar_log = np.array([signed_log1p(r.bstar) for r in recs])
    mm = np.array([r.mean_motion_rev_per_day for r in recs])
    ecc = np.array([r.eccentricity for r in recs])
    inc = np.array([r.inclination_deg for r in recs])
    raan_uw = np.degrees(np.unwrap(np.radians([r.raan_deg for r in recs])))
    argp_uw = np.degrees(np.unwrap(np.radians([r.argp_deg for r in recs])))
    phiM = cumulative_mean_anomaly(recs)

    step = grid_step_days * SECONDS_PER_DAY
    t0 = math.floor(ep[0] / step) * step
    t1 = math.ceil(ep[-1] / step) * step
    grid = np.arange(t0, t1 + 0.5 * step, step, dtype=np.float64)
    if grid.size < 2:
        return None

    def itp(y):
        return np.interp(grid, ep, y)

    g_raan, g_argp = itp(raan_uw), itp(argp_uw)
    g_mm, g_bstar = itp(mm), itp(bstar_log)
    g_ecc, g_inc, g_ma = itp(ecc), itp(inc), itp(phiM)   # g_ma = cumulative MA (deg)
    draan = np.gradient(g_raan) / grid_step_days
    dargp = np.gradient(g_argp) / grid_step_days
    d_bstar = np.gradient(g_bstar) / grid_step_days      # daily drift (persistence residual)
    d_mm = np.gradient(g_mm) / grid_step_days

    orbital = np.stack([d_bstar, d_mm, g_ecc, g_inc, draan, dargp], axis=1)
    # cm-tle-pred input-only physical features (Cartesian r/v, SMA/period/apo/peri,
    # cyclical angle encodings) derived from the day's elements
    phys = physics_features(g_mm, g_ecc, g_inc, g_raan, g_argp, g_ma)
    if sw is not None:
        solar = sw.for_epochs(grid)                      # (T, N_SOLAR)
    else:
        solar = np.zeros((grid.size, N_SOLAR), dtype=np.float32)
    feats = np.concatenate([orbital, phys, solar], axis=1).astype(np.float32)

    # aux = [MA, RAAN, argp, mean_motion_abs, bstar_slog_abs]
    aux = np.stack(
        [g_ma % 360.0, g_raan % 360.0, g_argp % 360.0, g_mm, g_bstar], axis=1
    ).astype(np.float32)

    idx = np.clip(np.searchsorted(ep, grid), 1, len(ep) - 1)
    gap = np.minimum(np.abs(grid - ep[idx]), np.abs(grid - ep[idx - 1]))
    mask = gap <= (2.0 * step)
    return DailySeries(recs[0].norad_id, grid, feats, mask, aux)


def _cache_key(input_dir, years, grid_step_days, clean, with_solar, min_grid_points, leo_only) -> str:
    h = hashlib.md5(str(Path(input_dir).resolve()).encode()).hexdigest()[:6]
    tag = Path(input_dir).resolve().name
    ys = "all" if years is None else f"{min(years)}-{max(years)}_{len(list(years))}"
    flags = f"c{int(clean)}s{int(with_solar)}m{int(min_grid_points)}L{int(leo_only)}"
    return f"{tag}_{h}_{ys}_g{grid_step_days:g}_{CHANNEL_VERSION}_{flags}"


def _load_cache_npz(cache_path, verbose) -> Dict[int, DailySeries]:
    d = np.load(cache_path, allow_pickle=True)
    out = {int(n): DailySeries(int(n), d[f"e_{int(n)}"], d[f"f_{int(n)}"],
                               d[f"m_{int(n)}"], d[f"a_{int(n)}"]) for n in d["norads"]}
    if verbose:
        print(f"[cache] {len(out)} satellites loaded from {cache_path}")
    return out


def build_daily_series(
    input_dir, years=None, grid_step_days=1.0, min_grid_points=64,
    clean=True, sw_csv=None, cache_dir=None, cache_file=None, rebuild=False,
    leo_only=True, verbose=True,
) -> Dict[int, DailySeries]:
    # Explicit prebuilt cache (e.g. the full 2005-2024 superset): load it directly
    # and let the dataset/eval filter by window + time-split. Skips parsing & solar.
    if cache_file is not None:
        cp = Path(cache_file)
        if not cp.exists():
            raise FileNotFoundError(f"--cache-file not found: {cp}")
        if verbose:
            print(f"[cache] using explicit cache {cp}")
        return _load_cache_npz(cp, verbose)

    sw = load_space_weather(sw_csv)
    with_solar = sw is not None
    cache_path = None
    if cache_dir is not None:
        cache_dir = Path(cache_dir); cache_dir.mkdir(parents=True, exist_ok=True)
        cache_path = cache_dir / f"tle_v4_{_cache_key(input_dir, years, grid_step_days, clean, with_solar, min_grid_points, leo_only)}.npz"
        if cache_path.exists() and not rebuild:
            if verbose:
                print(f"[cache] loading {cache_path}")
            d = np.load(cache_path, allow_pickle=True)
            out = {int(n): DailySeries(int(n), d[f"e_{int(n)}"], d[f"f_{int(n)}"],
                                       d[f"m_{int(n)}"], d[f"a_{int(n)}"]) for n in d["norads"]}
            if verbose:
                print(f"[cache] {len(out)} satellites loaded")
            return out

    by_norad = load_and_clean(input_dir, years=years, clean=clean, leo_only=leo_only)
    out: Dict[int, DailySeries] = {}
    for norad, recs in tqdm(by_norad.items(), desc="daily grid+solar", unit="sat"):
        ser = _build_one(recs, sw, grid_step_days)
        if ser is None or int(ser.mask.sum()) < min_grid_points:
            continue
        out[norad] = ser

    if cache_path is not None:
        kw: Dict[str, Any] = {"norads": np.asarray(sorted(out.keys()))}
        for n, s in out.items():
            kw[f"e_{n}"] = s.grid_epochs; kw[f"f_{n}"] = s.feats
            kw[f"m_{n}"] = s.mask; kw[f"a_{n}"] = s.aux
        np.savez(cache_path, **kw)
        if verbose:
            print(f"[cache] saved {len(out)} satellites -> {cache_path}")
    return out


# -----------------------------
# Windowed dataset
# -----------------------------

@dataclass
class DayWindow:
    norad_id: int
    start: int
    end_epoch: float


class TLEDatasetV2(Dataset):
    def __init__(
        self, input_dir, patch_size, years=None, window_patches=4, stride_patches=2,
        min_valid_steps=None, min_grid_points=None, split="train", train_until="2022-01-01",
        valid_until="2023-01-01", clean=True, sw_csv=None, cache_dir=None,
        cache_file=None, rebuild_cache=False, max_satellites=None,
        grid_step_days=1.0, leo_only=True, split_mode="time", verbose=True,
    ):
        if split not in {"train", "valid", "test", "all"}:
            raise ValueError("split must be train|valid|test|all")
        if split_mode not in {"time", "satellite"}:
            raise ValueError("split_mode must be time|satellite")
        self.patch_size = int(patch_size)
        self.window_length = int(window_patches) * self.patch_size
        self.stride = max(1, int(stride_patches) * self.patch_size)
        self.min_obs = int(min_valid_steps or self.window_length)
        self.split = split
        self.split_mode = split_mode    # 'time' (epoch cutoffs) or 'satellite' (cm-tle-pred 70/15/15)
        self.train_until_unix = parse_date_to_unix(train_until)
        self.valid_until_unix = parse_date_to_unix(valid_until)

        # cache-build satellite filter: keep a satellite if it has at least this many
        # observed grid days. Decoupled from window length so the cache can hold the
        # full population (per-window length is enforced later in _build_index).
        cache_min = int(min_grid_points) if min_grid_points is not None else self.window_length
        self.series = build_daily_series(
            input_dir, years=years, grid_step_days=grid_step_days,
            min_grid_points=cache_min, clean=clean, sw_csv=sw_csv,
            cache_dir=cache_dir, cache_file=cache_file, rebuild=rebuild_cache,
            leo_only=leo_only, verbose=verbose,
        )
        if max_satellites is not None:
            keep = sorted(self.series.keys())[:max_satellites]
            self.series = {k: self.series[k] for k in keep}
        self.index: List[DayWindow] = self._build_index()
        if verbose:
            print(json.dumps({
                "dataset": "TLEDatasetV2", "split": split, "split_mode": split_mode,
                "num_satellites": len(self.series), "num_windows": len(self.index),
                "window_days": self.window_length, "patch_size": self.patch_size,
                "n_channels": N_CHANNELS, "orbital": N_ORBITAL,
                "phys": N_PHYS, "solar": N_SOLAR,
            }, indent=2))

    def _split_of(self, epoch):
        if self.train_until_unix is None or self.valid_until_unix is None:
            return "all"
        if epoch < self.train_until_unix:
            return "train"
        if epoch < self.valid_until_unix:
            return "valid"
        return "test"

    def _build_index(self):
        out, W = [], self.window_length
        for norad, s in self.series.items():
            T = s.feats.shape[0]
            if T < W:
                continue
            # satellite-level split: the whole satellite belongs to one split
            if self.split != "all" and self.split_mode == "satellite" \
                    and sat_split_of(norad) != self.split:
                continue
            for st in range(0, T - W + 1, self.stride):
                if int(s.mask[st:st + W].sum()) < self.min_obs:
                    continue
                end_epoch = float(s.grid_epochs[st + W - 1])
                if self.split != "all" and self.split_mode == "time" \
                        and self._split_of(end_epoch) != self.split:
                    continue
                out.append(DayWindow(norad, st, end_epoch))
        return out

    def __len__(self):
        return len(self.index)

    def __getitem__(self, i):
        w = self.index[i]
        s = self.series[w.norad_id]
        W = self.window_length
        feats = s.feats[w.start:w.start + W]
        mask = s.mask[w.start:w.start + W]
        epochs = s.grid_epochs[w.start:w.start + W]
        aux = s.aux[w.start:w.start + W]

        # orbital channels use coverage mask; solar channels are always observed
        chan_mask = np.zeros((N_CHANNELS, W), dtype=bool)
        chan_mask[:N_ORBITAL] = mask[None, :]
        chan_mask[N_ORBITAL:] = True

        return {
            "target": torch.from_numpy(feats.T.copy()),          # (C, W)
            "target_mask": torch.from_numpy(chan_mask),           # (C, W)
            "series_ids": torch.zeros(N_CHANNELS, dtype=torch.long),
            "meta": {"norad_id": w.norad_id, "length": int(mask.sum()),
                     "epochs": epochs.copy(), "feats": feats.copy(), "aux": aux.copy()},
        }


def series_collate_fn(batch):
    return {
        "target": torch.stack([b["target"] for b in batch], 0),
        "target_mask": torch.stack([b["target_mask"] for b in batch], 0),
        "series_ids": torch.stack([b["series_ids"] for b in batch], 0),
        "meta": [b["meta"] for b in batch],
    }


def main():
    import argparse
    from torch.utils.data import DataLoader
    ap = argparse.ArgumentParser()
    ap.add_argument("--input-dir", default="/home/irteam/data-vol1/models/OrbitGPT/data/TLEs")
    ap.add_argument("--cache-dir", default="/home/irteam/data-vol1/models/OrbitGPT/v2/cache")
    ap.add_argument("--sw-csv", default="/home/irteam/data-vol1/models/OrbitGPT/v2/data/SW-All.csv")
    ap.add_argument("--start-year", type=int, default=2020)
    ap.add_argument("--end-year", type=int, default=2020)
    ap.add_argument("--patch-size", type=int, default=32)
    ap.add_argument("--window-patches", type=int, default=3)
    ap.add_argument("--stride-patches", type=int, default=1)
    ap.add_argument("--min-grid-points", type=int, default=64,
                    help="cache-build filter: keep satellites with >= this many observed "
                         "grid days. Low value (e.g. 64) keeps the full population.")
    ap.add_argument("--no-clean", action="store_true")
    ap.add_argument("--no-leo", action="store_true", help="disable LEO-only filter")
    ap.add_argument("--split", default="all")
    ap.add_argument("--split-mode", default="time", choices=["time", "satellite"])
    args = ap.parse_args()

    ds = TLEDatasetV2(
        input_dir=args.input_dir, cache_dir=args.cache_dir, sw_csv=args.sw_csv,
        years=range(args.start_year, args.end_year + 1), patch_size=args.patch_size,
        window_patches=args.window_patches, stride_patches=args.stride_patches,
        min_grid_points=args.min_grid_points, clean=not args.no_clean,
        leo_only=not args.no_leo, split=args.split, split_mode=args.split_mode,
    )
    if len(ds) == 0:
        print("No windows."); return
    b = next(iter(DataLoader(ds, batch_size=4, shuffle=True, collate_fn=series_collate_fn)))
    print("target", tuple(b["target"].shape), "obs frac", round(float(b["target_mask"].float().mean()), 3))
    print("channels:", N_CHANNELS, "= orbital", N_ORBITAL, "+ phys", N_PHYS, "+ solar", N_SOLAR)
    m0 = b["meta"][0]
    print("phys[t0] (rx,ry,rz,vx,vy,vz,...):",
          [round(float(x), 3) for x in m0["feats"][0][N_ORBITAL:N_ORBITAL + N_PHYS]])
    print("solar[t0] (f107,f107_81,ap):", m0["feats"][0][N_ORBITAL + N_PHYS:].tolist())
    track = reconstruct_track(m0["aux"][0], m0["feats"][1:6])
    print("reconstructed[+5d]:", json.dumps(track[-1], indent=2))
    print("true[+5d]:", json.dumps(elements_from_feat_aux(m0["feats"][5], m0["aux"][5]), indent=2))


if __name__ == "__main__":
    main()