"""Data cleaning workflow for the encounter-level Phase 3 base table."""

from __future__ import annotations

import logging
from typing import Any

import pandas as pd

from src.config import IngestionPaths
from src.utils.io import write_json

LOGGER = logging.getLogger(__name__)

COVID_ERA_START = pd.Timestamp("2020-01-01")
DATE_COLUMNS = [
    "birth_datetime",
    "visit_start_date",
    "visit_end_date",
    "measurement_date",
    "flu_last_administered",
    "tdap_last_administered",
    "mmr_last_administered",
    "polio_last_administered",
]
STRING_COLUMNS = [
    "race_source_value",
    "ethnicity_source_value",
    "gender_source_value",
    "visit_type",
    "condition",
    "observation_source",
    "deceased",
]
NUMERIC_COLUMNS = [
    "age_at_visit_years",
    "body_height_cm",
    "bmi",
    "body_temperature_c",
    "body_weight_kg",
    "systolic",
    "diastolic",
    "heart_rate_bpm",
    "oxygen_saturation_percent",
    "respiratory_rate_per_minute",
]


def _deduplicate_condition_terms(value: Any) -> tuple[str | pd.NA, bool, int]:
    if pd.isna(value):
        return pd.NA, False, 0
    terms = [term.strip() for term in str(value).split(":") if term.strip()]
    if not terms:
        return pd.NA, False, 0
    seen: set[str] = set()
    ordered_terms: list[str] = []
    repeated = False
    for term in terms:
        if term in seen:
            repeated = True
            continue
        seen.add(term)
        ordered_terms.append(term)
    return "|".join(ordered_terms), repeated, len(ordered_terms)


def _write_cleaned_csv(df: pd.DataFrame, path: pd.Timestamp | Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(path, index=False)


def clean_dataset(paths: IngestionPaths | None = None) -> dict[str, Any]:
    """Create the cleaned base table and machine-readable cleaning report."""
    active_paths = paths or IngestionPaths()
    df = pd.read_csv(active_paths.raw_csv)
    df = df.rename(columns={"measurement_Date": "measurement_date"})

    for column in STRING_COLUMNS + [
        "measurement_date",
        "birth_datetime",
        "visit_start_date",
        "visit_end_date",
        "flu_last_administered",
        "tdap_last_administered",
        "mmr_last_administered",
        "polio_last_administered",
    ]:
        df[column] = df[column].astype("string").str.strip()
        df[column] = df[column].replace({"": pd.NA, "nan": pd.NA})

    for column in DATE_COLUMNS:
        df[column] = pd.to_datetime(df[column], errors="coerce")

    for column in NUMERIC_COLUMNS:
        df[column] = pd.to_numeric(df[column], errors="coerce")

    condition_cleaned = df["condition"].apply(_deduplicate_condition_terms)
    df["condition_terms_cleaned"] = condition_cleaned.map(lambda item: item[0])
    df["condition_terms_repeated"] = condition_cleaned.map(lambda item: item[1])
    df["condition_term_count"] = condition_cleaned.map(lambda item: item[2]).astype(
        "Int64"
    )

    df["has_suspected_covid_observation"] = (
        df["observation_source"] == "Suspected COVID-19"
    ).fillna(False).astype(bool)
    symptom_pattern = (
        r"pneumonia|dyspnea|respiratory distress|wheezing|cough|sore throat"
    )
    df["is_candidate_covid_like_visit"] = (
        df["has_suspected_covid_observation"]
        | df["condition"].str.contains(symptom_pattern, case=False, na=False)
    ).fillna(False).astype(bool)
    df["is_pre_covid_era_visit"] = (
        (df["visit_start_date"] < COVID_ERA_START).fillna(False).astype(bool)
    )
    df["is_pre_covid_era_measurement"] = (
        (df["measurement_date"] < COVID_ERA_START).fillna(False).astype(bool)
    )
    df["measurement_after_visit_start"] = (
        df["measurement_date"].notna()
        & df["visit_start_date"].notna()
        & (df["measurement_date"] > df["visit_start_date"])
    ).fillna(False).astype(bool)

    for column in DATE_COLUMNS:
        df[column] = df[column].dt.strftime("%Y-%m-%d").astype("string")

    _write_cleaned_csv(df, active_paths.cleaned_csv)

    report = {
        "source_file": str(active_paths.raw_csv),
        "cleaned_file": str(active_paths.cleaned_csv),
        "row_count": int(len(df)),
        "column_count": int(df.shape[1]),
        "cleaned_columns": list(df.columns),
        "renamed_columns": {"measurement_Date": "measurement_date"},
        "missing_observation_source_count": int(df["observation_source"].isna().sum()),
        "pre_covid_visit_count": int(df["is_pre_covid_era_visit"].sum()),
        "pre_covid_measurement_count": int(df["is_pre_covid_era_measurement"].sum()),
        "measurement_after_visit_start_count": int(
            df["measurement_after_visit_start"].sum()
        ),
        "candidate_covid_like_visit_count": int(df["is_candidate_covid_like_visit"].sum()),
        "condition_terms_repeated_count": int(df["condition_terms_repeated"].sum()),
        "row_count_preserved": int(len(df)) == 156030,
    }
    write_json(active_paths.cleaning_report, report)
    LOGGER.info("Wrote cleaned dataset to %s", active_paths.cleaned_csv)
    LOGGER.info("Wrote cleaning report to %s", active_paths.cleaning_report)
    return report


if __name__ == "__main__":
    from src.logging_config import configure_logging

    configure_logging()
    clean_dataset()
