#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Space-weather (solar / geomagnetic activity) features for TLE drag modelling.

Atmospheric density -- and therefore drag and the secular decay of mean motion
-- is driven mostly by solar EUV (tracked by the F10.7 cm radio flux) and
geomagnetic activity (Ap index). Feeding these as extra input channels gives the
model the exogenous information it needs to predict how an orbit decays, which is
exactly where a learned model can beat "hold the last mean motion constant" SGP4
propagation at multi-day horizons.

Data source (download once, no auth):
    https://celestrak.org/SpaceData/SW-All.csv
Save it to v2/data/SW-All.csv (or pass --sw-csv). The CSV is daily from 1957.

Columns used: DATE, F10.7_OBS, F10.7_OBS_CENTER81, AP_AVG.
"""

from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

SOLAR_FEATURES = ["f107", "f107_81", "ap"]
N_SOLAR = len(SOLAR_FEATURES)
SW_URL = "https://celestrak.org/SpaceData/SW-All.csv"


class SpaceWeather:
    """Daily F10.7 / Ap lookup, aligned to arbitrary unix epochs."""

    def __init__(self, day_unix: np.ndarray, table: np.ndarray):
        self.day_unix = day_unix      # (D,) sorted unix seconds at 00:00 UTC
        self.table = table            # (D, N_SOLAR) float32

    @classmethod
    def from_csv(cls, csv_path: str | Path) -> "SpaceWeather":
        df = pd.read_csv(csv_path)
        cols = {c.upper(): c for c in df.columns}

        def col(*names):
            for nm in names:
                if nm in cols:
                    return df[cols[nm]]
            raise KeyError(f"none of {names} in SW csv columns {list(df.columns)}")

        dates = pd.to_datetime(col("DATE"))
        f107 = pd.to_numeric(col("F10.7_OBS", "F10.7_ADJ"), errors="coerce")
        f107_81 = pd.to_numeric(col("F10.7_OBS_CENTER81", "F10.7_ADJ_CENTER81",
                                    "F10.7_OBS_LAST81"), errors="coerce")
        ap = pd.to_numeric(col("AP_AVG"), errors="coerce")

        tab = pd.DataFrame({"f107": f107, "f107_81": f107_81, "ap": ap})
        tab = tab.ffill().bfill()  # fill predicted/missing tail+head
        day_unix = np.array(
            [d.replace(tzinfo=timezone.utc).timestamp() for d in dates.dt.to_pydatetime()],
            dtype=np.float64,
        )
        order = np.argsort(day_unix)
        return cls(day_unix[order], tab.to_numpy(dtype=np.float32)[order])

    def for_epochs(self, epochs_unix: np.ndarray) -> np.ndarray:
        """Return (len(epochs), N_SOLAR) by nearest-preceding-day lookup."""
        idx = np.searchsorted(self.day_unix, epochs_unix, side="right") - 1
        idx = np.clip(idx, 0, len(self.day_unix) - 1)
        return self.table[idx]


def load_space_weather(csv_path: Optional[str | Path]) -> Optional[SpaceWeather]:
    if csv_path is None:
        return None
    p = Path(csv_path)
    if not p.exists():
        print(f"[space_weather] WARNING: {p} not found -> solar channels will be ZERO.\n"
              f"               download once: {SW_URL}")
        return None
    sw = SpaceWeather.from_csv(p)
    print(f"[space_weather] loaded {len(sw.day_unix)} daily records from {p}")
    return sw