"""Information-theoretic GPU scheduling model with Bayesian learning.

Ported from agents/scheduler-agent/app/model.py — adapted to use shared db.py.
"""

from __future__ import annotations

import math
import threading
import time
import logging
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Optional

from radix_edge.config import get_config
from radix_edge.db import get_db

logger = logging.getLogger("radix_edge.scheduler")


@dataclass
class JobFeatures:
    """Features extracted from a job request."""
    job_type: str
    gpu_mem_gb: float = 0.0
    num_gpus: int = 1
    tenant: str = "default"


@dataclass
class Stats:
    """Statistics for a (job_type, gpu_type) pair."""
    mean_runtime: float
    variance: float
    count: int
    last_updated: float


@dataclass
class ScoringResult:
    """Result of scoring a job for GPU assignment."""
    chosen_gpu: str
    priority_score: float
    gpu_indices: list[int] = field(default_factory=list)
    avoid_co_locate_with: list[str] = field(default_factory=list)
    terms: dict[str, float] = field(default_factory=dict)


class SchedulerModel:
    """Information-theoretic GPU scheduler model.

    Scores jobs for GPU assignment using: cost(j,g) = μ + λ√σ² - β·IG + γ·penalty
    Learns runtime distributions via Bayesian EMA updates.
    """

    def __init__(self):
        config = get_config()
        self.lambda_uncertainty = config.lambda_uncertainty
        self.beta_exploration = config.beta_exploration
        self.gamma_interference = config.gamma_interference
        self.tau_squared = config.tau_squared
        self.exploration_cap = config.exploration_cap
        self.enable_interference = config.enable_interference
        self.retention_limit = config.retention_observations
        self.alpha = 0.1  # EMA learning rate

        self.lock = threading.RLock()

        # In-memory caches (backed by SQLite)
        self.stats: dict[tuple[str, str], Stats] = {}
        self.interference: dict[str, dict[str, float]] = defaultdict(lambda: defaultdict(float))

        # Exploration tracking
        self.exploration_decisions = 0
        self.total_decisions = 0

        self._load_from_db()

    def _load_from_db(self):
        """Load stats and interference from shared SQLite."""
        with self.lock:
            try:
                with get_db() as db:
                    for row in db.execute("SELECT * FROM scheduler_stats").fetchall():
                        key = (row["job_type"], row["gpu_type"])
                        self.stats[key] = Stats(
                            row["mean_runtime"], row["variance"],
                            row["count"], row["last_updated"],
                        )

                    for row in db.execute("SELECT * FROM scheduler_interference").fetchall():
                        self.interference[row["job_type"]][row["co_job_type"]] = row["slowdown_factor"]
            except Exception as e:
                logger.warning("Could not load scheduler stats: %s", e)

    def checkpoint(self):
        """Save current state to shared SQLite."""
        with self.lock:
            try:
                with get_db() as db:
                    db.execute("DELETE FROM scheduler_stats")
                    for (job_type, gpu_type), s in self.stats.items():
                        db.execute(
                            "INSERT INTO scheduler_stats VALUES (?, ?, ?, ?, ?, ?)",
                            (job_type, gpu_type, s.mean_runtime, s.variance, s.count, s.last_updated),
                        )

                    db.execute("DELETE FROM scheduler_interference")
                    for jt, co_jobs in self.interference.items():
                        for co_jt, slowdown in co_jobs.items():
                            if slowdown > 0:
                                db.execute(
                                    "INSERT INTO scheduler_interference VALUES (?, ?, ?)",
                                    (jt, co_jt, slowdown),
                                )
            except Exception as e:
                logger.warning("Could not checkpoint scheduler stats: %s", e)

    def get_stats(self, job_type: str, gpu_type: str) -> tuple[float, float]:
        """Get (mean, variance) for a (job_type, gpu_type) pair."""
        with self.lock:
            key = (job_type, gpu_type)
            if key in self.stats:
                s = self.stats[key]
                return s.mean_runtime, s.variance
            # Uninformative prior: high mean, high variance
            return 10.0, 25.0

    def information_gain(self, variance: float) -> float:
        """IG = 0.5 * log(1 + σ²/τ²)."""
        return 0.5 * math.log(1 + variance / self.tau_squared)

    def interference_penalty(self, job_type: str, colocated_types: list[str]) -> float:
        """Penalty for job collocation based on learned interference."""
        if not self.enable_interference or not colocated_types:
            return 0.0
        penalty = sum(
            self.interference[job_type].get(co, 0.0) for co in colocated_types
        )
        return penalty / len(colocated_types)

    def score_job(self, features: JobFeatures, gpu_types: list[str],
                  colocated_types: Optional[list[str]] = None) -> ScoringResult:
        """Score a job for GPU assignment using information-theoretic objective.

        cost(j,g) = μ + λ√σ² - β·IG + γ·penalty
        Chooses the GPU type with minimum cost.
        """
        if colocated_types is None:
            colocated_types = []

        records = []
        for gpu_type in gpu_types:
            mu, sigma2 = self.get_stats(features.job_type, gpu_type)
            ig = self.information_gain(sigma2)
            penalty = self.interference_penalty(features.job_type, colocated_types)

            cost = (
                mu
                + self.lambda_uncertainty * math.sqrt(sigma2)
                - self.beta_exploration * ig
                + self.gamma_interference * penalty
            )
            records.append((gpu_type, mu, sigma2, ig, penalty, cost))

        # Choose minimum cost
        chosen = min(records, key=lambda r: r[-1])
        chosen_gpu, mu, sigma2, ig, penalty, cost = chosen

        # Exploration cap enforcement
        is_exploration = sigma2 > 1.0
        with self.lock:
            self.total_decisions += 1
            if is_exploration:
                self.exploration_decisions += 1
                ratio = self.exploration_decisions / self.total_decisions
                if ratio > self.exploration_cap:
                    exploit = min(records, key=lambda r: r[1])  # min by μ
                    chosen_gpu, mu, sigma2, ig, penalty, cost = exploit

        # Normalize cost to priority score [0, 100]
        costs = [r[-1] for r in records]
        if len(costs) > 1:
            max_c, min_c = max(costs), min(costs)
            if max_c > min_c:
                priority_score = 100 * (max_c - cost) / (max_c - min_c)
            else:
                priority_score = 50.0
        else:
            priority_score = 50.0

        return ScoringResult(
            chosen_gpu=chosen_gpu,
            priority_score=min(100.0, max(0.0, priority_score)),
            terms={
                "mu": mu,
                "sigma": math.sqrt(sigma2),
                "ig": ig,
                "penalty": penalty,
                "cost": cost,
            },
        )

    def observe(self, job_type: str, gpu_type: str, runtime: float,
                colocated_types: Optional[list[str]] = None):
        """Update model with observed runtime (Bayesian EMA)."""
        if colocated_types is None:
            colocated_types = []

        with self.lock:
            key = (job_type, gpu_type)
            now = time.time()

            if key in self.stats:
                s = self.stats[key]
                new_mean = (1 - self.alpha) * s.mean_runtime + self.alpha * runtime
                error = runtime - s.mean_runtime
                new_var = (1 - self.alpha) * s.variance + self.alpha * (error ** 2)
                self.stats[key] = Stats(new_mean, max(0.01, new_var), s.count + 1, now)
            else:
                self.stats[key] = Stats(runtime, 1.0, 1, now)

            # Interference learning
            if self.enable_interference and colocated_types:
                solo_mean, _ = self.get_stats(job_type, gpu_type)
                if solo_mean > 0:
                    slowdown = max(0.0, (runtime - solo_mean) / solo_mean)
                    for co in colocated_types:
                        cur = self.interference[job_type][co]
                        self.interference[job_type][co] = 0.9 * cur + 0.1 * slowdown

            self._apply_retention()

    def _apply_retention(self):
        """Limit memory by removing oldest entries."""
        if len(self.stats) <= self.retention_limit:
            return
        sorted_stats = sorted(self.stats.items(), key=lambda x: x[1].last_updated)
        to_remove = len(self.stats) - self.retention_limit
        for i in range(to_remove):
            del self.stats[sorted_stats[i][0]]

    def get_metrics(self) -> dict[str, float]:
        """Get model metrics for monitoring."""
        with self.lock:
            ratio = self.exploration_decisions / max(1, self.total_decisions)
            avg_unc = 0.0
            if self.stats:
                avg_unc = sum(math.sqrt(s.variance) for s in self.stats.values()) / len(self.stats)
            return {
                "decisions_total": self.total_decisions,
                "exploration_ratio": ratio,
                "avg_uncertainty": avg_unc,
                "stats_count": len(self.stats),
                "interference_pairs": sum(len(v) for v in self.interference.values()),
            }
