import unittest

import numpy as np

from quread.severity_mapper import SeverityConfig, SeverityThresholds, SeverityWeights, compute_severity_rows, severity_rows_to_csv


class SeverityMapperTest(unittest.TestCase):
    def _metrics(self):
        return {
            "activity_count": np.array([1.0, 4.0, 2.0], dtype=float),
            "activity_norm": np.array([0.25, 1.0, 0.5], dtype=float),
            "gate_error": np.array([0.02, 0.75, 0.30], dtype=float),
            "readout_error": np.array([0.03, 0.40, 0.08], dtype=float),
            "coherence_health": np.array([0.95, 0.30, 0.70], dtype=float),
            "decoherence_risk": np.array([0.10, 0.80, 0.35], dtype=float),
            "fidelity": np.array([0.98, 0.50, 0.90], dtype=float),
            "state_fidelity": np.array([0.99, 0.55, 0.92], dtype=float),
            "process_fidelity": np.array([0.98, 0.60, 0.91], dtype=float),
            "composite_risk": np.array([0.18, 0.86, 0.41], dtype=float),
        }

    def test_linear_mode_uses_selected_heatmap_metric(self):
        rows = compute_severity_rows(
            self._metrics(),
            cfg=SeverityConfig(mode="linear", source_metric="gate_error"),
            qubit_coords={0: (0, 0), 1: (0, 1), 2: (1, 0)},
        )
        self.assertEqual(rows[0]["qubit"], 1)
        self.assertEqual(rows[0]["severity_band"], "CRITICAL")
        self.assertAlmostEqual(rows[0]["severity_score"], 0.75, places=6)
        self.assertEqual(rows[0]["source_metric"], "gate_error")

    def test_bucket_mode_quantizes_to_three_levels(self):
        rows = compute_severity_rows(
            self._metrics(),
            cfg=SeverityConfig(
                mode="bucket",
                source_metric="composite_risk",
                thresholds=SeverityThresholds(warning=0.4, critical=0.7),
            ),
        )
        values = {int(r["qubit"]): float(r["severity_score"]) for r in rows}
        self.assertEqual(values[0], 0.0)
        self.assertEqual(values[1], 1.0)
        self.assertEqual(values[2], 0.5)

    def test_percentile_mode_uses_current_session_ranking(self):
        rows = compute_severity_rows(
            self._metrics(),
            cfg=SeverityConfig(
                mode="percentile",
                source_metric="composite_risk",
                thresholds=SeverityThresholds(percentile_warning=50.0, percentile_critical=90.0),
            ),
        )
        by_q = {int(r["qubit"]): r for r in rows}
        self.assertEqual(by_q[1]["severity_band"], "CRITICAL")
        self.assertEqual(by_q[1]["severity_rank"], 1)
        self.assertGreater(by_q[1]["severity_percentile"], by_q[2]["severity_percentile"])

    def test_weighted_mode_uses_separate_weight_model(self):
        rows = compute_severity_rows(
            self._metrics(),
            cfg=SeverityConfig(
                mode="weighted",
                weights=SeverityWeights(activity=0.0, gate_error=1.0, readout_error=0.0, decoherence=0.0, fidelity=0.0),
            ),
        )
        self.assertEqual(rows[0]["qubit"], 1)
        self.assertEqual(rows[0]["source_metric"], "weighted_blend")
        self.assertAlmostEqual(rows[0]["severity_score"], 0.75, places=6)

    def test_pnr_cost_mode_emits_richer_outputs(self):
        rows = compute_severity_rows(
            self._metrics(),
            cfg=SeverityConfig(mode="pnr_cost"),
        )
        top = rows[0]
        self.assertEqual(top["severity_band"], "CRITICAL")
        self.assertGreater(top["pnr_cost"], 0.0)
        self.assertGreater(top["timing_derate"], 1.0)
        self.assertLess(top["density_cap"], 100.0)
        self.assertIn(top["route_priority"], {"LOW", "MEDIUM", "HIGH"})

    def test_severity_rows_to_csv_contains_headers(self):
        csv_text = severity_rows_to_csv(
            compute_severity_rows(
                self._metrics(),
                cfg=SeverityConfig(mode="linear", source_metric="composite_risk"),
                qubit_coords={0: (0, 0), 1: (0, 1), 2: (1, 0)},
            )
        )
        self.assertIn("severity_mode", csv_text.splitlines()[0])
        self.assertIn("pnr_cost", csv_text.splitlines()[0])
        self.assertIn("\n1,0,1,", csv_text)


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