"""Tests for M3: Scoring model, Sinkhorn solver, and scheduling policies."""

from __future__ import annotations

import math

import numpy as np

from radix_edge.scheduler.scoring_model import SchedulerModel, JobFeatures, Stats
from radix_edge.scheduler.sinkhorn import SinkhornSolver
from radix_edge.scheduler.policies import order_jobs
from radix_edge.models import Job, JobStatus


# --- Scoring Model ---

class TestScoringModel:
    def test_score_single_gpu_type(self, db):
        model = SchedulerModel()
        features = JobFeatures(job_type="training")
        result = model.score_job(features, ["A100"])

        assert result.chosen_gpu == "A100"
        assert result.priority_score == 50.0  # Only one GPU type -> default
        assert "mu" in result.terms
        assert "sigma" in result.terms
        assert "ig" in result.terms
        assert "cost" in result.terms

    def test_score_multiple_gpu_types(self, db):
        model = SchedulerModel()

        # Seed stats: A100 is faster than V100 for training
        model.stats[("training", "A100")] = Stats(5.0, 0.5, 100, 0.0)
        model.stats[("training", "V100")] = Stats(15.0, 0.5, 100, 0.0)

        features = JobFeatures(job_type="training")
        result = model.score_job(features, ["A100", "V100"])

        # A100 should be chosen (lower cost)
        assert result.chosen_gpu == "A100"
        assert result.priority_score > 50.0  # A100 clearly better

    def test_exploration_prefers_high_uncertainty(self, db):
        model = SchedulerModel()
        model.beta_exploration = 2.0  # High exploration weight
        model.exploration_cap = 1.0  # Disable cap so exploration can happen

        # A100: well-known (low variance), V100: unknown (high variance)
        model.stats[("training", "A100")] = Stats(8.0, 0.1, 100, 0.0)
        model.stats[("training", "V100")] = Stats(10.0, 50.0, 2, 0.0)

        features = JobFeatures(job_type="training")
        result = model.score_job(features, ["A100", "V100"])

        # With high beta and no cap, V100's high IG should make it attractive
        assert result.chosen_gpu == "V100"

    def test_observe_updates_stats(self, db):
        model = SchedulerModel()

        model.observe("training", "A100", 10.0)
        mu, var = model.get_stats("training", "A100")
        assert mu == 10.0
        assert var == 1.0  # Initial variance

        model.observe("training", "A100", 12.0)
        mu, var = model.get_stats("training", "A100")
        # EMA: new_mean = 0.9*10 + 0.1*12 = 10.2
        assert abs(mu - 10.2) < 0.01

    def test_default_prior(self, db):
        model = SchedulerModel()
        mu, var = model.get_stats("unknown_type", "unknown_gpu")
        assert mu == 10.0
        assert var == 25.0

    def test_information_gain(self, db):
        model = SchedulerModel()
        # High variance -> high IG
        ig_high = model.information_gain(100.0)
        ig_low = model.information_gain(0.01)
        assert ig_high > ig_low

    def test_exploration_cap(self, db):
        model = SchedulerModel()
        model.exploration_cap = 0.0  # No exploration allowed
        model.exploration_decisions = 100
        model.total_decisions = 100

        # V100 has extremely high variance but higher mean
        model.stats[("training", "A100")] = Stats(5.0, 0.01, 1000, 0.0)
        model.stats[("training", "V100")] = Stats(20.0, 1000.0, 1, 0.0)

        features = JobFeatures(job_type="training")
        result = model.score_job(features, ["A100", "V100"])

        # With cap = 0, should exploit (choose A100 with lower mean)
        assert result.chosen_gpu == "A100"

    def test_checkpoint_and_reload(self, db):
        model = SchedulerModel()
        model.observe("training", "A100", 10.0)
        model.observe("inference", "V100", 5.0)
        model.checkpoint()

        # Create new model instance (should load from DB)
        model2 = SchedulerModel()
        mu, var = model2.get_stats("training", "A100")
        assert mu == 10.0

    def test_metrics(self, db):
        model = SchedulerModel()
        model.observe("training", "A100", 10.0)
        features = JobFeatures(job_type="training")
        model.score_job(features, ["A100"])

        metrics = model.get_metrics()
        assert metrics["decisions_total"] == 1
        assert metrics["stats_count"] == 1
        assert "exploration_ratio" in metrics


# --- Sinkhorn Solver ---

class TestSinkhorn:
    def test_basic_2x2(self):
        solver = SinkhornSolver(epsilon=0.1)
        costs = np.array([[1.0, 2.0], [2.0, 1.0]])
        supply = np.array([1.0, 1.0])
        demand = np.array([1.0, 1.0])

        assignment = solver.solve(costs, supply, demand)

        # Check marginal constraints
        assert np.allclose(assignment.sum(axis=1), supply, atol=1e-3)
        assert np.allclose(assignment.sum(axis=0), demand, atol=1e-3)

        # Job 0 prefers GPU 0, Job 1 prefers GPU 1
        assert assignment[0, 0] > assignment[0, 1]
        assert assignment[1, 1] > assignment[1, 0]

    def test_job_assignment_interface(self):
        solver = SinkhornSolver()
        job_costs = [
            [1.0, 2.0],  # Job 0 prefers GPU 0
            [2.0, 1.0],  # Job 1 prefers GPU 1
            [1.5, 1.5],  # Job 2 indifferent
        ]
        gpu_capacities = [2, 2]

        assignments = solver.assign_jobs_to_gpus(job_costs, gpu_capacities)
        assert len(assignments) == 3

        for a in assignments:
            assert len(a) > 0
            assert abs(sum(a.values()) - 1.0) < 0.1

    def test_empty_input(self):
        solver = SinkhornSolver()
        assert solver.assign_jobs_to_gpus([], []) == []
        assert solver.assign_jobs_to_gpus([], [1, 2]) == []

    def test_uniform_costs(self):
        solver = SinkhornSolver(epsilon=0.1)
        costs = np.array([[1.0, 1.0], [1.0, 1.0]])
        supply = np.array([1.0, 1.0])
        demand = np.array([1.0, 1.0])

        assignment = solver.solve(costs, supply, demand)

        # Uniform costs -> roughly equal assignment
        assert np.allclose(assignment, 0.5 * np.ones((2, 2)), atol=0.1)


# --- Scheduling Policies ---

class TestPolicies:
    def _make_job(self, job_id, priority=5, user_id="u1", created_at="2026-01-01T00:00:00Z",
                  timeout=3600, job_type="generic"):
        return Job(
            job_id=job_id, user_id=user_id, image="img",
            priority=priority, created_at=created_at,
            timeout_seconds=timeout, job_type=job_type,
        )

    def test_fifo(self):
        jobs = [
            self._make_job("j3", created_at="2026-01-03T00:00:00Z"),
            self._make_job("j1", created_at="2026-01-01T00:00:00Z"),
            self._make_job("j2", created_at="2026-01-02T00:00:00Z"),
        ]
        ordered = order_jobs(jobs, "fifo")
        assert [j.job_id for j in ordered] == ["j1", "j2", "j3"]

    def test_priority(self):
        jobs = [
            self._make_job("j1", priority=3),
            self._make_job("j2", priority=8),
            self._make_job("j3", priority=5),
        ]
        ordered = order_jobs(jobs, "priority")
        assert [j.job_id for j in ordered] == ["j2", "j3", "j1"]

    def test_priority_tiebreak_by_time(self):
        jobs = [
            self._make_job("j1", priority=5, created_at="2026-01-02T00:00:00Z"),
            self._make_job("j2", priority=5, created_at="2026-01-01T00:00:00Z"),
        ]
        ordered = order_jobs(jobs, "priority")
        # Same priority -> FIFO tiebreak
        assert [j.job_id for j in ordered] == ["j2", "j1"]

    def test_sjf(self):
        jobs = [
            self._make_job("j1", timeout=7200),
            self._make_job("j2", timeout=600),
            self._make_job("j3", timeout=3600),
        ]
        ordered = order_jobs(jobs, "sjf")
        assert [j.job_id for j in ordered] == ["j2", "j3", "j1"]

    def test_fairshare_without_db(self):
        """Without DB, falls back to priority ordering."""
        jobs = [
            self._make_job("j1", priority=3),
            self._make_job("j2", priority=8),
        ]
        ordered = order_jobs(jobs, "fairshare")
        assert [j.job_id for j in ordered] == ["j2", "j1"]

    def test_fairshare_with_db(self, db):
        """With DB, users with less usage get priority."""
        # User-a has lots of recent usage, user-b has none
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, runtime_seconds, num_gpus, completed_at)
            VALUES ('old-1', 'user-a', 'img', 'succeeded', 3600, 1, datetime('now'))"""
        )
        db.commit()

        jobs = [
            self._make_job("j1", priority=5, user_id="user-a"),
            self._make_job("j2", priority=5, user_id="user-b"),
        ]
        ordered = order_jobs(jobs, "fairshare", db)
        # user-b should come first (lower usage)
        assert ordered[0].user_id == "user-b"

    def test_unknown_policy_fallback(self):
        jobs = [self._make_job("j1")]
        ordered = order_jobs(jobs, "nonexistent_policy")
        assert len(ordered) == 1

    def test_empty_list(self):
        assert order_jobs([], "fifo") == []
