"""ENGIN pipeline correctness suite, version 2.

This filename is intentionally new to avoid stale downloads of the original
``test_benchmark_grouped.py`` artifact.
"""

from __future__ import annotations

import unittest
from pathlib import Path

import numpy as np
import pandas as pd

from benchmark_grouped import FREQ_COLS, SpectrumFeatures, make_splits, validate_data
from robustness_grouped import mask_spectral_cells
from severity_benchmark import OrdinalLogistic, OrdinalRidge, SeverityFeatures


REPO_DIR = Path(__file__).resolve().parent
SUITE_VERSION = "ENGIN_PIPELINE_TESTS_V2_SEVERITY_ENABLED"


class SuiteVersionTests(unittest.TestCase):
    def test_v2_severity_suite_marker(self) -> None:
        self.assertEqual(SUITE_VERSION, "ENGIN_PIPELINE_TESTS_V2_SEVERITY_ENABLED")


class SpectrumFeatureTests(unittest.TestCase):
    def make_frame(self) -> pd.DataFrame:
        rows = []
        for engine_id, offsets in (
            ("a", [0.0, 2.0, 4.0]),
            ("b", [10.0, 14.0, 18.0]),
        ):
            for cylinder, offset in enumerate(offsets, start=1):
                row = {"engine_id": engine_id, "cylinder": cylinder}
                row.update(
                    {
                        column: float(frequency + offset)
                        for frequency, column in enumerate(FREQ_COLS)
                    }
                )
                rows.append(row)
        return pd.DataFrame(rows)

    def test_relative_feature_median_is_zero_within_each_engine(self) -> None:
        frame = self.make_frame()
        transformed = SpectrumFeatures("relative").fit_transform(frame)
        relative = pd.DataFrame(transformed, columns=FREQ_COLS)
        relative["engine_id"] = frame["engine_id"].to_numpy()
        medians = relative.groupby("engine_id")[FREQ_COLS].median().to_numpy()
        self.assertTrue(np.allclose(medians, 0.0))

    def test_combined_features_have_expected_width_and_no_nan(self) -> None:
        frame = self.make_frame()
        frame.loc[0, "mV_10"] = np.nan
        transformed = SpectrumFeatures("combined").fit_transform(frame)
        self.assertEqual(transformed.shape, (len(frame), 2 * len(FREQ_COLS)))
        self.assertFalse(np.isnan(transformed).any())
        self.assertAlmostEqual(transformed[0, 10], 10.0)


class MissingMaskTests(unittest.TestCase):
    def make_frame(self, rows: int = 20) -> pd.DataFrame:
        frame = pd.DataFrame(
            {
                "engine_id": [f"engine_{row // 4}" for row in range(rows)],
                "cylinder": [(row % 4) + 1 for row in range(rows)],
                "label": ["ok"] * rows,
            }
        )
        for frequency, column in enumerate(FREQ_COLS):
            frame[column] = np.arange(rows, dtype=float) + frequency
        return frame

    def test_mask_is_exact_deterministic_and_spectral_only(self) -> None:
        frame = self.make_frame()
        masked_a, manifest_a = mask_spectral_cells(frame, 0.05, 2026)
        masked_b, manifest_b = mask_spectral_cells(frame, 0.05, 2026)
        expected = round(0.05 * len(frame) * len(FREQ_COLS))
        self.assertEqual(masked_a[FREQ_COLS].isna().to_numpy().sum(), expected)
        pd.testing.assert_frame_equal(manifest_a, manifest_b)
        pd.testing.assert_frame_equal(masked_a, masked_b)
        pd.testing.assert_frame_equal(
            frame[["engine_id", "cylinder", "label"]],
            masked_a[["engine_id", "cylinder", "label"]],
        )


class GroupSplitTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.frame = pd.read_csv(REPO_DIR / "val.csv").reset_index(drop=True)

    def test_every_engine_is_held_out_once_without_overlap(self) -> None:
        validate_data(self.frame, n_splits=5)
        splits = make_splits(
            self.frame,
            splitter_name="stratified-group",
            n_splits=5,
            random_state=42,
        )
        validation_engines = []
        for train_idx, valid_idx in splits:
            train_engines = set(self.frame.iloc[train_idx]["engine_id"])
            valid_engines = set(self.frame.iloc[valid_idx]["engine_id"])
            self.assertTrue(train_engines.isdisjoint(valid_engines))
            validation_engines.extend(valid_engines)
        self.assertEqual(len(validation_engines), 40)
        self.assertEqual(len(set(validation_engines)), 40)


class SeverityTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.frame = pd.read_csv(REPO_DIR / "val.csv").reset_index(drop=True)
        cls.frame["diagnostic_label"] = cls.frame["label"]

    def test_severity_features_are_finite_and_engine_complete(self) -> None:
        features = SeverityFeatures(
            "all", include_fault_type=True
        ).fit_transform(self.frame)
        self.assertEqual(features.shape[0], len(self.frame))
        self.assertGreater(features.shape[1], 4 * len(FREQ_COLS))
        self.assertTrue(np.isfinite(features).all())

    def test_ordinal_estimators_emit_only_allowed_severities(self) -> None:
        fault = self.frame["label"].isin(
            ["zakoksowany", "lejacy", "pompa", "iglica"]
        )
        features = SeverityFeatures("deviation").fit_transform(self.frame)[
            fault.to_numpy()
        ]
        severity = self.frame.loc[fault, "severity"].to_numpy()
        for estimator in (OrdinalRidge(), OrdinalLogistic(random_state=42)):
            predicted = estimator.fit(features, severity).predict(features[:10])
            self.assertTrue(
                set(predicted).issubset({"male", "srednie", "duze"})
            )


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