"""Probability calibration tracking for prediction market strategies.

Tracks estimated probabilities vs actual outcomes to measure and improve
prediction accuracy. Uses lightweight SQLite storage.
"""

from __future__ import annotations

import math
import sqlite3
import time
from dataclasses import dataclass, field


@dataclass(frozen=True, slots=True)
class CalibrationBucket:
    """One bucket in a calibration curve."""
    bin_start: float
    bin_end: float
    predicted_mean: float
    actual_frequency: float
    count: int


@dataclass(frozen=True, slots=True)
class CalibrationReport:
    """Summary of calibration performance."""
    brier_score: float
    log_loss: float
    n_predictions: int
    n_resolved: int
    buckets: list[CalibrationBucket] = field(default_factory=list)

    @property
    def overconfidence(self) -> float:
        """Positive = overconfident (predicted higher than actual)."""
        if not self.buckets:
            return 0.0
        total_pred = 0.0
        total_actual = 0.0
        total_count = 0
        for b in self.buckets:
            total_pred += b.predicted_mean * b.count
            total_actual += b.actual_frequency * b.count
            total_count += b.count
        if total_count == 0:
            return 0.0
        return (total_pred - total_actual) / total_count


class CalibrationTracker:
    """Track probability calibration with SQLite persistence.

    Usage:
        tracker = CalibrationTracker("calibration.db")
        tracker.log_prediction("mkt_1", 0.75)
        # ... later, when market resolves ...
        tracker.resolve_market("mkt_1", True)
        report = tracker.report()
        print(f"Brier: {report.brier_score:.4f}")
    """

    def __init__(self, db_path: str = ":memory:") -> None:
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA synchronous=NORMAL")
        self._create_tables()

    def _create_tables(self) -> None:
        self._conn.executescript("""
            CREATE TABLE IF NOT EXISTS predictions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                market_id TEXT NOT NULL,
                probability REAL NOT NULL,
                timestamp REAL NOT NULL
            );
            CREATE TABLE IF NOT EXISTS resolutions (
                market_id TEXT PRIMARY KEY,
                outcome INTEGER NOT NULL,
                resolved_at REAL NOT NULL
            );
            CREATE INDEX IF NOT EXISTS idx_pred_market ON predictions(market_id);
        """)

    def log_prediction(
        self,
        market_id: str,
        probability: float,
        timestamp: float | None = None,
    ) -> None:
        """Log a probability prediction for a market."""
        ts = timestamp if timestamp is not None else time.time()
        self._conn.execute(
            "INSERT INTO predictions (market_id, probability, timestamp) VALUES (?, ?, ?)",
            (market_id, probability, ts),
        )
        self._conn.commit()

    def resolve_market(self, market_id: str, outcome: bool) -> None:
        """Record the resolution of a market (True = YES won)."""
        self._conn.execute(
            "INSERT OR REPLACE INTO resolutions (market_id, outcome, resolved_at) VALUES (?, ?, ?)",
            (market_id, int(outcome), time.time()),
        )
        self._conn.commit()

    def brier_score(self) -> float:
        """Compute Brier score across all resolved markets.

        Uses the last prediction per market (closest to resolution).
        Lower is better. Perfect = 0.0, random = 0.25.
        """
        rows = self._conn.execute("""
            SELECT p.probability, r.outcome
            FROM resolutions r
            JOIN predictions p ON p.market_id = r.market_id
            WHERE p.id = (
                SELECT id FROM predictions
                WHERE market_id = r.market_id
                ORDER BY timestamp DESC LIMIT 1
            )
        """).fetchall()
        if not rows:
            return 0.0
        total = 0.0
        for prob, outcome in rows:
            total += (prob - outcome) ** 2
        return total / len(rows)

    def log_loss(self) -> float:
        """Compute log loss (cross-entropy) across resolved markets.

        Uses the last prediction per market. Lower is better.
        """
        rows = self._conn.execute("""
            SELECT p.probability, r.outcome
            FROM resolutions r
            JOIN predictions p ON p.market_id = r.market_id
            WHERE p.id = (
                SELECT id FROM predictions
                WHERE market_id = r.market_id
                ORDER BY timestamp DESC LIMIT 1
            )
        """).fetchall()
        if not rows:
            return 0.0
        eps = 1e-15
        total = 0.0
        for prob, outcome in rows:
            p = max(eps, min(1.0 - eps, prob))
            if outcome:
                total -= math.log(p)
            else:
                total -= math.log(1.0 - p)
        return total / len(rows)

    def calibration_curve(self, n_bins: int = 10) -> list[CalibrationBucket]:
        """Generate calibration curve: predicted probability vs actual frequency.

        Divides [0, 1] into n_bins equal-width bins. For each bin, computes
        the mean predicted probability and the actual resolution frequency.
        """
        rows = self._conn.execute("""
            SELECT p.probability, r.outcome
            FROM resolutions r
            JOIN predictions p ON p.market_id = r.market_id
            WHERE p.id = (
                SELECT id FROM predictions
                WHERE market_id = r.market_id
                ORDER BY timestamp DESC LIMIT 1
            )
        """).fetchall()
        if not rows:
            return []

        bin_width = 1.0 / n_bins
        bins: list[list[tuple[float, int]]] = [[] for _ in range(n_bins)]

        for prob, outcome in rows:
            if not (0.0 <= prob <= 1.0):
                continue
            idx = min(int(prob / bin_width), n_bins - 1)
            bins[idx].append((prob, outcome))

        buckets = []
        for i, bin_data in enumerate(bins):
            if not bin_data:
                continue
            bin_start = i * bin_width
            bin_end = (i + 1) * bin_width
            predicted_mean = sum(p for p, _ in bin_data) / len(bin_data)
            actual_freq = sum(o for _, o in bin_data) / len(bin_data)
            buckets.append(CalibrationBucket(
                bin_start=bin_start,
                bin_end=bin_end,
                predicted_mean=predicted_mean,
                actual_frequency=actual_freq,
                count=len(bin_data),
            ))
        return buckets

    def suggest_adjustment(self, probability: float, n_bins: int = 10) -> float:
        """Suggest an adjusted probability based on historical calibration.

        If historical data shows overconfidence at this probability level,
        the adjustment pulls the estimate toward the empirical frequency.
        Returns the original probability if insufficient data.
        """
        buckets = self.calibration_curve(n_bins)
        if not buckets:
            return probability

        for b in buckets:
            if b.bin_start <= probability < b.bin_end or (
                probability == 1.0 and b.bin_end == 1.0
            ):
                if b.count < 5:
                    return probability
                # Blend: 70% original, 30% empirical adjustment
                bias = b.predicted_mean - b.actual_frequency
                return max(0.01, min(0.99, probability - 0.3 * bias))

        return probability

    def report(self, n_bins: int = 10) -> CalibrationReport:
        """Generate a full calibration report."""
        n_pred = self._conn.execute("SELECT COUNT(*) FROM predictions").fetchone()[0]
        n_res = self._conn.execute("SELECT COUNT(*) FROM resolutions").fetchone()[0]
        return CalibrationReport(
            brier_score=self.brier_score(),
            log_loss=self.log_loss(),
            n_predictions=n_pred,
            n_resolved=n_res,
            buckets=self.calibration_curve(n_bins),
        )

    def clear(self) -> None:
        """Clear all calibration data."""
        self._conn.executescript("""
            DELETE FROM predictions;
            DELETE FROM resolutions;
        """)

    def close(self) -> None:
        """Close the database connection."""
        self._conn.close()
