"""Tests for the raw data ingestion and validation workflow."""

from __future__ import annotations

import unittest

import pandas as pd

from src.config import IngestionPaths
from src.data.clean import clean_dataset
from src.data.validate import validate_raw_dataset


class DataValidationTestCase(unittest.TestCase):
    """Validate the current raw dataset assumptions from Phase 2."""

    @classmethod
    def setUpClass(cls) -> None:
        cls.paths = IngestionPaths()
        cls.report = validate_raw_dataset(cls.paths)
        cls.cleaning_report = clean_dataset(cls.paths)
        cls.cleaned_df = pd.read_csv(cls.paths.cleaned_csv)

    def test_row_and_column_counts_match_expected_source(self) -> None:
        self.assertEqual(self.report["row_count"], 156030)
        self.assertEqual(self.report["column_count"], 27)

    def test_required_columns_exist(self) -> None:
        self.assertTrue(all(self.report["required_columns_present"].values()))

    def test_visit_occurrence_id_is_unique(self) -> None:
        self.assertEqual(self.report["duplicate_visit_occurrence_id_count"], 0)
        self.assertTrue(self.report["encounter_identifier_assumption_validated"])

    def test_visit_dates_are_not_reversed(self) -> None:
        self.assertEqual(self.report["visit_end_before_start_count"], 0)

    def test_header_mismatch_is_detected_and_documented(self) -> None:
        self.assertFalse(self.report["header_matches_metadata"])
        issue_codes = {issue["code"] for issue in self.report["issues"]}
        self.assertIn("header_metadata_mismatch", issue_codes)

    def test_broad_cohort_scope_warning_is_present(self) -> None:
        issue_codes = {issue["code"] for issue in self.report["issues"]}
        self.assertIn("broad_cohort_scope", issue_codes)
        self.assertGreater(self.report["pre_2020_visit_count"], 0)

    def test_cleaned_dataset_is_created_with_expected_shape(self) -> None:
        self.assertEqual(len(self.cleaned_df), 156030)
        self.assertNotIn("measurement_Date", self.cleaned_df.columns)
        self.assertIn("measurement_date", self.cleaned_df.columns)

    def test_cleaned_dataset_has_quality_flags(self) -> None:
        for column in [
            "is_pre_covid_era_visit",
            "is_pre_covid_era_measurement",
            "measurement_after_visit_start",
            "has_suspected_covid_observation",
            "is_candidate_covid_like_visit",
            "condition_terms_repeated",
            "condition_terms_cleaned",
            "condition_term_count",
        ]:
            self.assertIn(column, self.cleaned_df.columns)

    def test_cleaning_report_matches_known_counts(self) -> None:
        self.assertTrue(self.cleaning_report["row_count_preserved"])
        self.assertEqual(self.cleaning_report["pre_covid_visit_count"], 69324)
        self.assertEqual(self.cleaning_report["measurement_after_visit_start_count"], 2790)
        self.assertEqual(self.cleaning_report["condition_terms_repeated_count"], 2263)


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