tests/test_solver.py

"""
Tests for core simulation components:
- DC power flow solver
- Environment lifecycle (reset, step, terminate)
- Grading system (scoring, bounds, reproducibility)
- Baseline heuristic policy
"""

import copy
import unittest

import numpy as np

from src.physics import DCSolver, IslandedException
from src.environment import OpenGridEnv
from src.tasks import TASKS
from src.models import GridAction, BusAdjustment
from src.grader import RobustnessGrader, compute_analytical_ceiling
from src.baseline import heuristic_policy


def task(task_id: str):
    """Get a deep-copied task config to prevent cross-test contamination."""
    return copy.deepcopy(TASKS[task_id])


class TestDCSolver(unittest.TestCase):
    def setUp(self):
        self.num_buses = 3
        self.lines = [
            {'id': 'L01', 'from': 0, 'to': 1, 'susceptance': 100, 'connected': True},
            {'id': 'L12', 'from': 1, 'to': 2, 'susceptance': 50, 'connected': True},
            {'id': 'L02', 'from': 0, 'to': 2, 'susceptance': 100, 'connected': True}
        ]
        self.solver = DCSolver(self.num_buses)
        self.solver.update_grid(self.lines)

    def test_power_flow_balance(self):
        """Slack bus should absorb any generation/load imbalance."""
        p_inj = np.array([0.0, 50.0, -100.0])
        theta, flows, slack_inj = self.solver.solve(p_inj)

        # Check that flows are computed
        self.assertIn('L01', flows)
        self.assertIn('L02', flows)

    def test_islanding_detection(self):
        """Disconnecting lines to island bus 2 should raise IslandedException."""
        with self.assertRaises(IslandedException):
            broken_lines = [
                {'id': 'L01', 'from': 0, 'to': 1, 'susceptance': 100, 'connected': True},
                {'id': 'L12', 'from': 1, 'to': 2, 'susceptance': 50, 'connected': False},
                {'id': 'L02', 'from': 0, 'to': 2, 'susceptance': 100, 'connected': False}
            ]
            self.solver.update_grid(broken_lines)

    def test_slack_injection_returned(self):
        """solve() should return slack bus injection as third element."""
        p_inj = np.array([0.0, 50.0, -100.0])
        result = self.solver.solve(p_inj)
        self.assertEqual(len(result), 3)
        theta, flows, slack_inj = result
        # Slack should inject ~50 MW to cover the deficit
        self.assertAlmostEqual(slack_inj, 50.0, places=0)

    def test_solve_before_update_raises(self):
        """Calling solve() on a fresh solver should raise RuntimeError."""
        fresh = DCSolver(3)
        with self.assertRaises(RuntimeError):
            fresh.solve(np.array([0.0, 10.0, -10.0]))

    def test_invalid_bus_index_raises(self):
        """Lines referencing out-of-range bus IDs should raise ValueError."""
        bad_lines = [
            {'id': 'L_bad', 'from': 0, 'to': 99, 'susceptance': 50, 'connected': True},
        ]
        solver = DCSolver(3)
        with self.assertRaises(ValueError):
            solver.update_grid(bad_lines)


class TestEnvironment(unittest.TestCase):
    def test_reset_returns_observation(self):
        """reset() should return a valid GridObservation."""
        env = OpenGridEnv(task("task_easy"))
        obs = env.reset()
        self.assertEqual(obs.timestep, 0)
        self.assertGreater(len(obs.buses), 0, "Observation should have buses")
        self.assertGreater(len(obs.lines), 0, "Observation should have lines")

    def test_step_returns_tuple(self):
        """step() should return (obs, reward, done, info)."""
        env = OpenGridEnv(task("task_easy"))
        env.reset()
        obs, reward, done, info = env.step(GridAction())
        self.assertEqual(obs.timestep, 1)
        self.assertIsInstance(reward.value, float)
        self.assertIsInstance(done, bool)

    def test_reproducibility(self):
        """Running the same task twice should produce identical initial observations."""
        env1 = OpenGridEnv(task("task_easy"))
        obs1 = env1.reset()

        env2 = OpenGridEnv(task("task_easy"))
        obs2 = env2.reset()

        self.assertEqual(obs1.grid_frequency, obs2.grid_frequency)
        self.assertEqual(len(obs1.buses), len(obs2.buses))

    def test_episode_terminates(self):
        """Episode should end after max_steps."""
        config = task("task_easy")
        env = OpenGridEnv(config)
        env.reset()
        done = False
        steps = 0
        while not done and steps < 100:
            _, _, done, _ = env.step(GridAction())
            steps += 1
        self.assertTrue(done, "Episode should terminate")
        self.assertLessEqual(steps, config["max_steps"])

    def test_frequency_reasonable(self):
        """Frequency should stay in a reasonable range for do-nothing agent."""
        env = OpenGridEnv(task("task_easy"))
        obs = env.reset()
        for _ in range(10):
            obs, _, done, _ = env.step(GridAction())
            if done:
                break
            self.assertGreater(obs.grid_frequency, 40.0,
                               "Frequency below reasonable minimum")
            self.assertLess(obs.grid_frequency, 60.0,
                            "Frequency above reasonable maximum")


class TestGrader(unittest.TestCase):
    def test_grader_score_range(self):
        """Grader should return score strictly in (0, 1) — never 0.0 or 1.0."""
        grader = RobustnessGrader(task("task_easy"))
        result = grader.evaluate_policy(heuristic_policy, n_episodes=1)
        self.assertGreater(result["score"], 0.0)
        self.assertLess(result["score"], 1.0)

    def test_grader_all_tasks(self):
        """Grader should work on all registered tasks."""
        for task_id, config in TASKS.items():
            grader = RobustnessGrader(copy.deepcopy(config))
            result = grader.evaluate_policy(heuristic_policy, n_episodes=1)
            self.assertIn("score", result, f"Missing 'score' for {task_id}")
            self.assertIn("avg_raw_reward", result,
                          f"Missing 'avg_raw_reward' for {task_id}")


class TestBaseline(unittest.TestCase):
    def test_heuristic_returns_valid_action(self):
        """Heuristic policy should return a valid GridAction."""
        env = OpenGridEnv(task("task_easy"))
        obs = env.reset()
        action = heuristic_policy(obs)
        self.assertIsInstance(action, GridAction)


class TestReproducibility(unittest.TestCase):
    def test_floor_deterministic(self):
        """Two calls to _estimate_bounds should produce identical floors (seeded RNG)."""
        grader1 = RobustnessGrader(task("task_easy"))
        grader1._estimate_bounds(n_samples=3)

        grader2 = RobustnessGrader(task("task_easy"))
        grader2._estimate_bounds(n_samples=3)

        self.assertEqual(grader1.reward_floor, grader2.reward_floor,
                         "Floor should be deterministic with same seed")

    def test_ceiling_is_analytical(self):
        """Ceiling should be max_steps * 1.2, not an empirical estimate."""
        config = task("task_easy")
        grader = RobustnessGrader(config)
        bounds = grader.get_bounds()
        expected_ceiling = compute_analytical_ceiling(config["max_steps"])
        self.assertEqual(bounds["reward_ceiling"], expected_ceiling,
                         "Ceiling should match analytical formula")

    def test_heuristic_score_below_one(self):
        """With analytical ceiling, heuristic should score < 1.0 (not degenerate)."""
        grader = RobustnessGrader(task("task_easy"))
        result = grader.evaluate_policy(heuristic_policy, n_episodes=1)
        self.assertLess(result["score"], 1.0)
        self.assertGreater(result["score"], 0.0)


if __name__ == '__main__':
    unittest.main()