"""Correctness tests for the :class:`grader.Grader` module."""
from __future__ import annotations

import pytest

from agents import HeuristicAgent, RandomAgent
from grader import SCENARIOS, Grader


def test_scenarios_contain_required_keys() -> None:
    assert "normal_day" in SCENARIOS
    assert "surge" in SCENARIOS
    assert "mass_casualty" in SCENARIOS
    for cfg in SCENARIOS.values():
        assert "arrival_rate" in cfg
        assert "severity_dist" in cfg


def test_grade_output_shape_and_bounds() -> None:
    grader = Grader(n_episodes_per_scenario=1)
    agent = RandomAgent(seed=0)
    scores = grader.grade(agent)

    expected_keys = {
        "overall_score",
        "survival_rate",
        "avg_wait_time",
        "resource_utilization",
        "critical_patient_survival",
        "invalid_action_rate",
        "scenario_scores",
    }
    assert expected_keys.issubset(scores.keys())

    # Composite score lives on [0, 100]
    assert 0.0 <= scores["overall_score"] <= 100.0
    assert 0.0 <= scores["survival_rate"] <= 1.0
    assert 0.0 <= scores["critical_patient_survival"] <= 1.0
    assert 0.0 <= scores["invalid_action_rate"] <= 1.0

    for name in SCENARIOS.keys():
        assert name in scores["scenario_scores"]
        assert 0.0 <= scores["scenario_scores"][name] <= 100.0


def test_heuristic_beats_random_overall() -> None:
    """Heuristic should outscore random on overall composite.

    Run with fewer episodes for speed; we only need the direction to be right.
    """
    grader = Grader(n_episodes_per_scenario=2, base_seed=17)

    random_agent = RandomAgent(seed=0)
    heuristic_agent = HeuristicAgent()

    random_scores = grader.grade(random_agent)
    heuristic_scores = grader.grade(heuristic_agent)

    # Heuristic should have a higher overall score AND fewer invalid actions.
    assert heuristic_scores["overall_score"] > random_scores["overall_score"]
    assert heuristic_scores["invalid_action_rate"] <= random_scores["invalid_action_rate"]


def test_grader_is_reproducible() -> None:
    grader_a = Grader(n_episodes_per_scenario=1, base_seed=99)
    grader_b = Grader(n_episodes_per_scenario=1, base_seed=99)

    scores_a = grader_a.grade(HeuristicAgent())
    scores_b = grader_b.grade(HeuristicAgent())

    assert scores_a == scores_b