"""Probability calibration for WeatherNext forecasts.

Converts raw model outputs into well-calibrated probabilities
suitable for prediction market trading.
"""

import json
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional, List, Tuple

import numpy as np


class ProbabilityCalibrator:
    """
    Calibrate raw model probabilities using historical performance.

    Methods:
    - Platt scaling (logistic regression on historical outcomes)
    - Isotonic regression (non-parametric)
    - Ensemble (combine multiple calibration methods)
    """

    def __init__(self, calibration_file: str = "data/calibration_history.json"):
        self.calibration_file = Path(__file__).parent.parent / calibration_file
        self.history: List[Dict] = self._load_history()
        self.platt_params: Dict[str, Tuple[float, float]] = {}
        self._fit()

    def _load_history(self) -> List[Dict]:
        """Load historical forecast vs outcome data."""
        if self.calibration_file.exists():
            try:
                with open(self.calibration_file) as f:
                    return json.load(f)
            except Exception:
                return []
        return []

    def save_history(self):
        """Save calibration history."""
        self.calibration_file.parent.mkdir(parents=True, exist_ok=True)
        with open(self.calibration_file, "w") as f:
            json.dump(self.history, f, indent=2)

    def record_outcome(
        self,
        date: str,
        variable: str,
        predicted_probability: float,
        actual_outcome: bool,
    ):
        """Record a prediction-outcome pair for future calibration."""
        self.history.append({
            "date": date,
            "variable": variable,
            "predicted_probability": predicted_probability,
            "actual_outcome": actual_outcome,
            "recorded_at": datetime.now().isoformat(),
        })
        self.save_history()
        self._fit()  # Re-fit on new data

    def _fit(self):
        """Fit Platt scaling parameters from history."""
        by_variable: Dict[str, List[Tuple[float, int]]] = {}
        for record in self.history:
            var = record["variable"]
            if var not in by_variable:
                by_variable[var] = []
            by_variable[var].append((
                record["predicted_probability"] / 100.0,
                1 if record["actual_outcome"] else 0,
            ))

        for var, data in by_variable.items():
            if len(data) >= 5:
                try:
                    from sklearn.linear_model import LogisticRegression
                    X = np.array([[d[0]] for d in data])
                    y = np.array([d[1] for d in data])
                    lr = LogisticRegression()
                    lr.fit(X, y)
                    self.platt_params[var] = (lr.coef_[0][0], lr.intercept_[0])
                except ImportError:
                    # Fallback: simple linear correction
                    self._simple_fit(var, data)
                except Exception:
                    self._simple_fit(var, data)
            elif len(data) >= 2:
                self._simple_fit(var, data)

    def _simple_fit(self, var: str, data: List[Tuple[float, int]]):
        """Simple linear calibration for small datasets."""
        probs = np.array([d[0] for d in data])
        outcomes = np.array([d[1] for d in data])
        mean_prob = probs.mean()
        mean_outcome = outcomes.mean()

        slope = 1.0
        intercept = mean_outcome - mean_prob

        self.platt_params[var] = (slope, intercept)

    def calibrate(self, variable: str, raw_probability: float) -> float:
        """Calibrate a raw probability (0-100) to a calibrated one."""
        x = raw_probability / 100.0

        if variable in self.platt_params:
            a, b = self.platt_params[variable]
            calibrated = 1.0 / (1.0 + np.exp(-(a * x + b)))
            return float(np.clip(calibrated * 100.0, 0.5, 99.5))

        return float(np.clip(raw_probability, 0.5, 99.5))

    def ensemble_calibrate(
        self, variable: str, raw_probability: float
    ) -> Tuple[float, float]:
        """
        Return (calibrated_probability, confidence_interval_width).
        Confidence width shrinks with more historical data.
        """
        cal_prob = self.calibrate(variable, raw_probability)

        n_obs = sum(1 for h in self.history if h["variable"] == variable)
        if n_obs < 5:
            ci_width = 15.0
        elif n_obs < 20:
            ci_width = 10.0
        elif n_obs < 50:
            ci_width = 5.0
        else:
            ci_width = 3.0

        return cal_prob, ci_width

    def get_calibration_stats(self, variable: str) -> Dict:
        """Get calibration statistics for a variable."""
        relevant = [h for h in self.history if h["variable"] == variable]
        if not relevant:
            return {"n_observations": 0, "brier_score": None, "calibration_error": None}

        preds = np.array([h["predicted_probability"] / 100.0 for h in relevant])
        outcomes = np.array([1 if h["actual_outcome"] else 0 for h in relevant])

        brier = float(np.mean((preds - outcomes) ** 2))
        cal_error = float(np.abs(preds.mean() - outcomes.mean()))

        return {
            "n_observations": len(relevant),
            "brier_score": brier,
            "calibration_error": cal_error,
            "mean_prediction": float(preds.mean() * 100),
            "mean_outcome": float(outcomes.mean() * 100),
        }