"""Tests for district label normalization at the case-loader boundary.

The processed/cases_by_district_daily.parquet carries both bare ("武昌") and
区-suffixed ("武昌区") spellings of each district (26 labels = 13 districts × 2
spellings), which double-counts in any roll-up. data.case_loader normalizes
these to the canonical 13 区-suffixed names and re-aggregates. These tests pin
that behavior.
"""
import sys
from pathlib import Path

import pandas as pd
import pytest

# Ensure the backend package root is importable at collection time (mirrors the
# sys.path handling other modules rely on once the app is imported).
BACKEND_ROOT = Path(__file__).parent.parent
if str(BACKEND_ROOT) not in sys.path:
    sys.path.insert(0, str(BACKEND_ROOT))

from data.case_loader import (  # noqa: E402
    CANONICAL_DISTRICTS,
    normalize_district,
    load_cases_by_district_daily,
)

PROJECT_ROOT = Path(__file__).parent.parent.parent
RAW_PARQUET = PROJECT_ROOT / "processed" / "cases_by_district_daily.parquet"


def test_normalize_district_known_bare_forms():
    """Every known bare form maps to its canonical 区-suffixed name."""
    cases = {
        "武昌": "武昌区", "汉阳": "汉阳区", "江岸": "江岸区", "硚口": "硚口区",
        "青山": "青山区", "洪山": "洪山区", "东西湖": "东西湖区", "汉南": "汉南区",
        "蔡甸": "蔡甸区", "江夏": "江夏区", "黄陂": "黄陂区", "新洲": "新洲区",
        "江汉": "江汉区",
    }
    for bare, canonical in cases.items():
        assert normalize_district(bare) == canonical


def test_normalize_district_already_suffixed_passes_through():
    for d in CANONICAL_DISTRICTS:
        assert normalize_district(d) == d


def test_canonical_set_is_exactly_thirteen():
    assert len(CANONICAL_DISTRICTS) == 13
    assert len(set(CANONICAL_DISTRICTS)) == 13


@pytest.mark.skipif(not RAW_PARQUET.exists(), reason="case parquet not present")
def test_loader_collapses_to_thirteen_canonical_districts():
    df = load_cases_by_district_daily()
    districts = set(df["district"].unique())

    # (a) exactly 13 unique districts, all canonical
    assert len(districts) == 13, f"expected 13 districts, got {len(districts)}: {sorted(districts)}"
    assert districts == set(CANONICAL_DISTRICTS)

    # (b) no bare / unsuffixed duplicates remain
    for name in districts:
        assert name.endswith(("区", "县", "市")), f"unsuffixed district leaked: {name}"


@pytest.mark.skipif(not RAW_PARQUET.exists(), reason="case parquet not present")
def test_loader_preserves_totals_no_rows_dropped_or_double_counted():
    """Sum integrity: normalized total == raw parquet total."""
    raw = pd.read_parquet(RAW_PARQUET)
    normalized = load_cases_by_district_daily()

    assert int(normalized["total_cases"].sum()) == int(raw["total_cases"].sum())
    assert int(normalized["outpatient_count"].sum()) == int(raw["outpatient_count"].sum())
    assert int(normalized["inpatient_count"].sum()) == int(raw["inpatient_count"].sum())


@pytest.mark.skipif(not RAW_PARQUET.exists(), reason="case parquet not present")
def test_raw_parquet_actually_has_dirty_labels():
    """Sanity: the raw file really has the 26-label problem we are fixing."""
    raw = pd.read_parquet(RAW_PARQUET)
    assert raw["district"].nunique() > 13