"""Handle CMS suppression and standardize both datasets.

Inpatient: no missing values per Phase 1 — keep as-is after rename/dtype.
Outpatient: separate flows for cost-analysis rows vs coverage-count rows.
"""
import pandas as pd
from .load import load_inpatient, load_outpatient
from .config import DATA_PROC


def clean_inpatient() -> pd.DataFrame:
    """Inpatient: just dedupe + light standardization. No missing values expected."""
    df = load_inpatient()
    df = df.drop_duplicates(subset=["provider_ccn", "drg_code"])
    df["provider_name"] = df["provider_name"].str.strip()
    df["provider_city"] = df["provider_city"].str.strip()
    df["state"] = df["state"].str.strip().str.upper()
    df["zip5"] = df["zip5"].str.zfill(5)
    return df.reset_index(drop=True)


def clean_outpatient(drop_suppressed: bool = True) -> pd.DataFrame:
    """Outpatient: handle suppression.

    - avg_allowed_amount is the primary cost measure. CMS blanks it when
      services ≤ 10. We surface the suppression flag explicitly so downstream
      code can choose to drop or keep.
    - outlier_services / avg_outlier_amount blanks are NOT zero — they mean
      "suppressed because count < 11". Left as NaN; an explicit
      `outlier_suppressed` flag aids any sensitivity analysis.

    drop_suppressed=True drops rows with blank primary cost fields — that's
    what the cost-distribution analyses want. The full (incl. suppressed) frame
    can be retrieved with drop_suppressed=False.
    """
    df = load_outpatient()
    df = df.drop_duplicates(subset=["provider_ccn", "apc_code"])
    df["provider_name"] = df["provider_name"].str.strip()
    df["provider_city"] = df["provider_city"].str.strip()
    df["state"] = df["state"].str.strip().str.upper()
    df["zip5"] = df["zip5"].str.zfill(5)

    # Explicit suppression flags
    df["cost_suppressed"] = df["avg_allowed_amount"].isna()
    df["outlier_suppressed"] = df["outlier_services"].isna()

    if drop_suppressed:
        df = df[~df["cost_suppressed"]].copy()
    return df.reset_index(drop=True)


def write_meta_summary(inp: pd.DataFrame, out_full: pd.DataFrame, out_cost: pd.DataFrame) -> dict:
    """Compact summary used by the site's 'approach' page."""
    summary = {
        "inpatient": {
            "rows": int(len(inp)),
            "providers": int(inp["provider_ccn"].nunique()),
            "services": int(inp["drg_code"].nunique()),
            "states": int(inp["state"].nunique()),
            "missing_values": 0,
            "suppression": "none",
        },
        "outpatient_full": {
            "rows": int(len(out_full)),
            "providers": int(out_full["provider_ccn"].nunique()),
            "services": int(out_full["apc_code"].nunique()),
            "states": int(out_full["state"].nunique()),
            "cost_suppressed_rows": int(out_full["cost_suppressed"].sum()),
            "outlier_suppressed_rows": int(out_full["outlier_suppressed"].sum()),
            "cost_suppressed_pct": round(out_full["cost_suppressed"].mean() * 100, 2),
        },
        "outpatient_cost_subset": {
            "rows": int(len(out_cost)),
            "providers": int(out_cost["provider_ccn"].nunique()),
            "services": int(out_cost["apc_code"].nunique()),
            "note": "rows with avg_allowed_amount not suppressed (services > 10)",
        },
        "provider_overlap": {
            "inpatient_only": int(set(inp["provider_ccn"]) - set(out_full["provider_ccn"]) and 1 or 0)
                or len(set(inp["provider_ccn"]) - set(out_full["provider_ccn"])),
            "outpatient_only": len(set(out_full["provider_ccn"]) - set(inp["provider_ccn"])),
            "both": len(set(inp["provider_ccn"]) & set(out_full["provider_ccn"])),
        },
    }
    return summary


if __name__ == "__main__":  # pragma: no cover
    inp = clean_inpatient()
    out_cost = clean_outpatient(drop_suppressed=True)
    out_full = clean_outpatient(drop_suppressed=False)
    inp.to_parquet(DATA_PROC / "inpatient_clean.parquet", index=False)
    out_cost.to_parquet(DATA_PROC / "outpatient_clean.parquet", index=False)
    out_full.to_parquet(DATA_PROC / "outpatient_full.parquet", index=False)
    print(f"inpatient_clean: {inp.shape}")
    print(f"outpatient_clean (cost): {out_cost.shape}")
    print(f"outpatient_full: {out_full.shape}")