# Generated by Claude Code -- 2026-02-08
"""Model 1: Naive Baseline -- Orbital Shell Density Prior.

Predicts collision risk based solely on the altitude band of the conjunction,
using historical base rates. This establishes that altitude alone is predictive
(LEO is more crowded) but insufficient for actionable conjunction assessment.
"""

import json
import numpy as np
from pathlib import Path
from collections import defaultdict


class OrbitalShellBaseline:
    """
    Altitude-band collision rate baseline.

    For any conjunction event, predict the average risk and miss distance
    for that altitude regime. Bins events into 50km altitude bands.
    """

    def __init__(self, bin_width_km: float = 50.0):
        self.bin_width = bin_width_km
        self.bins: dict[int, dict] = {}
        self.global_stats: dict = {}

    def _altitude_to_bin(self, alt_km: float) -> int:
        return int(round(alt_km / self.bin_width) * self.bin_width)

    def fit(self, altitudes: np.ndarray, y_risk: np.ndarray, y_miss_log: np.ndarray):
        """
        Fit baseline from altitude array and labels.

        Args:
            altitudes: altitude in km for each event
            y_risk: binary risk labels
            y_miss_log: log1p(miss_distance_km) targets
        """
        # Global fallback stats
        self.global_stats = {
            "mean_risk": float(np.mean(y_risk)),
            "mean_miss_log": float(np.mean(y_miss_log)),
            "count": int(len(y_risk)),
        }

        # Per-bin statistics
        bin_data = defaultdict(lambda: {"risks": [], "misses": []})

        for alt, risk, miss in zip(altitudes, y_risk, y_miss_log):
            b = self._altitude_to_bin(alt)
            bin_data[b]["risks"].append(risk)
            bin_data[b]["misses"].append(miss)

        self.bins = {}
        for b, data in bin_data.items():
            self.bins[b] = {
                "mean_risk": float(np.mean(data["risks"])),
                "mean_miss_log": float(np.mean(data["misses"])),
                "count": len(data["risks"]),
                "risk_rate": float(np.sum(data["risks"]) / len(data["risks"])),
            }

        print(f"Baseline fit: {len(self.bins)} altitude bins, "
              f"global risk rate = {self.global_stats['mean_risk']:.4f}")

    def predict(self, altitudes: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """
        Predict risk probability and log miss distance for each altitude.

        Returns: (risk_probs, miss_log_preds)
        """
        risk_preds = []
        miss_preds = []

        for alt in altitudes:
            b = self._altitude_to_bin(alt)
            if b in self.bins:
                risk_preds.append(self.bins[b]["risk_rate"])
                miss_preds.append(self.bins[b]["mean_miss_log"])
            else:
                risk_preds.append(self.global_stats["mean_risk"])
                miss_preds.append(self.global_stats["mean_miss_log"])

        return np.array(risk_preds), np.array(miss_preds)

    def save(self, path: Path):
        """Save model to JSON."""
        data = {
            "bin_width": self.bin_width,
            "bins": {str(k): v for k, v in self.bins.items()},
            "global_stats": self.global_stats,
        }
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
        print(f"Baseline saved to {path}")

    @classmethod
    def load(cls, path: Path) -> "OrbitalShellBaseline":
        """Load model from JSON."""
        with open(path) as f:
            data = json.load(f)
        model = cls(bin_width_km=data["bin_width"])
        model.bins = {int(k): v for k, v in data["bins"].items()}
        model.global_stats = data["global_stats"]
        return model