import json
import os
from typing import Dict, Optional

import pandas as pd


BASE_DIR = os.path.dirname(os.path.abspath(__file__))


def load_config() -> dict:
    cfg_path = os.path.join(BASE_DIR, "config.json")
    if not os.path.exists(cfg_path):
        raise FileNotFoundError(
            f"config.json file not found at {cfg_path}. Please create this file with the required configuration."
        )
    with open(cfg_path, "r") as f:
        return json.load(f)


class QualityCheck:
    """Run post-ProcessData quality checks on processed outputs.

    Typical usage:
        qc = geopops.QualityCheck()
        # detailed dict is available at qc.results
    """

    def __init__(self, config_dict: Optional[dict] = None, output_dir: Optional[str] = None, auto_run: bool = True):
        self.config = config_dict if config_dict is not None else load_config()
        if output_dir is not None:
            self.output_dir = output_dir
        else:
            self.output_dir = self.config.get("path", BASE_DIR)
        self.processed_dir = os.path.join(self.output_dir, "processed")
        self.results: Dict[str, object] = {}

        if auto_run:
            self.run_all()

    def _require_file(self, filename: str) -> str:
        path = os.path.join(self.processed_dir, filename)
        if not os.path.exists(path):
            raise FileNotFoundError(
                f"Required processed file not found: {path}. "
                "Run geopops.ProcessData() before geopops.QualityCheck()."
            )
        return path

    @staticmethod
    def _normalize_code(series: pd.Series, width: int = 7) -> pd.Series:
        return series.astype(str).str.strip().str.zfill(width)

    @staticmethod
    def _is_missing_code(series: pd.Series) -> pd.Series:
        s = series.astype(str).str.strip()
        return s.eq("") | s.eq("nan") | s.eq("None") | s.eq("<NA>")

    def check_st_puma_overlap(self) -> Dict[str, object]:
        cbg_path = self._require_file("cbg_geo.csv")
        samp_path = self._require_file("samp_geo.csv")

        cbg_geo = pd.read_csv(cbg_path, dtype=str)
        samp_geo = pd.read_csv(samp_path, dtype=str)

        for col in ["st_puma"]:
            if col not in cbg_geo.columns or col not in samp_geo.columns:
                raise KeyError(f"Expected column '{col}' in both cbg_geo.csv and samp_geo.csv.")

        cbg_missing_mask = self._is_missing_code(cbg_geo["st_puma"])
        cbg_non_missing = cbg_geo.loc[~cbg_missing_mask, "st_puma"]
        cbg_set = set(self._normalize_code(cbg_non_missing).unique())
        samp_set = set(self._normalize_code(samp_geo["st_puma"]).unique())
        inter = cbg_set & samp_set

        by_state = {}
        for s in sorted({x[:2] for x in cbg_set | samp_set}):
            cbg_s = {x for x in cbg_set if x.startswith(s)}
            samp_s = {x for x in samp_set if x.startswith(s)}
            inter_s = cbg_s & samp_s
            if cbg_s or samp_s:
                by_state[s] = {
                    "cbg_unique_st_puma": len(cbg_s),
                    "samp_unique_st_puma": len(samp_s),
                    "intersection": len(inter_s),
                }

        coverage = (len(inter) / len(cbg_set)) if cbg_set else 0.0
        return {
            "cbg_rows_total": int(len(cbg_geo)),
            "cbg_rows_missing_st_puma": int(cbg_missing_mask.sum()),
            "cbg_rows_non_missing_st_puma": int((~cbg_missing_mask).sum()),
            "cbg_unique_st_puma": len(cbg_set),
            "samp_unique_st_puma": len(samp_set),
            "intersection": len(inter),
            "cbg_covered_by_samp_fraction": round(float(coverage), 4),
            "cbg_not_in_samp_examples": sorted(list(cbg_set - samp_set))[:20],
            "samp_not_in_cbg_examples": sorted(list(samp_set - cbg_set))[:20],
            "by_state": by_state,
        }

    def check_samp_geo_missingness(self) -> Dict[str, object]:
        samp_path = self._require_file("samp_geo.csv")
        samp_geo = pd.read_csv(samp_path, dtype=str)

        checks = ["cbsa", "county", "R", "U"]
        for col in checks:
            if col not in samp_geo.columns:
                raise KeyError(f"Expected column '{col}' in samp_geo.csv.")

        n = max(len(samp_geo), 1)
        missing = {}
        stripped = {}
        for col in checks:
            s = samp_geo[col].astype(str).str.strip()
            stripped[col] = s
            is_missing = s.eq("") | s.eq("nan") | s.eq("None")
            missing[col] = {
                "missing_count": int(is_missing.sum()),
                "missing_fraction": round(float(is_missing.sum() / n), 6),
            }

        cbsa_missing = stripped["cbsa"].eq("") | stripped["cbsa"].eq("nan") | stripped["cbsa"].eq("None")
        county_missing = stripped["county"].eq("") | stripped["county"].eq("nan") | stripped["county"].eq("None")
        r_missing = stripped["R"].eq("") | stripped["R"].eq("nan") | stripped["R"].eq("None")
        u_missing = stripped["U"].eq("") | stripped["U"].eq("nan") | stripped["U"].eq("None")

        blank_cbsa_with_geo = int((cbsa_missing & ~(county_missing & r_missing & u_missing)).sum())
        blank_cbsa_all_geo_blank = int((cbsa_missing & county_missing & r_missing & u_missing).sum())

        return {
            "missingness": missing,
            "blank_cbsa_with_county_or_RU_present": blank_cbsa_with_geo,
            "blank_cbsa_with_county_and_RU_also_blank": blank_cbsa_all_geo_blank,
        }

    def run_all(self) -> Dict[str, object]:
        if not os.path.isdir(self.processed_dir):
            raise FileNotFoundError(
                f"Processed directory not found: {self.processed_dir}. "
                "Run geopops.ProcessData() before geopops.QualityCheck()."
            )

        overlap = self.check_st_puma_overlap()
        missing = self.check_samp_geo_missingness()

        self.results = {
            "output_dir": self.output_dir,
            "processed_dir": self.processed_dir,
            "main_year": self.config.get("main_year"),
            "st_puma_overlap": overlap,
            "samp_geo_missingness": missing,
        }

        print("QualityCheck summary")
        print(f"- output_dir: {self.output_dir}")
        print(f"- main_year: {self.results['main_year']}")
        print(
            "- st_puma overlap: "
            f"{overlap['intersection']} / {overlap['cbg_unique_st_puma']} "
            f"({overlap['cbg_covered_by_samp_fraction']:.2%} of target st_puma codes covered)"
        )
        print(
            "- samp_geo missing fractions (cbsa/county/R/U): "
            f"{missing['missingness']['cbsa']['missing_fraction']:.3%}, "
            f"{missing['missingness']['county']['missing_fraction']:.3%}, "
            f"{missing['missingness']['R']['missing_fraction']:.3%}, "
            f"{missing['missingness']['U']['missing_fraction']:.3%}"
        )
        print(
            "- rows with blank cbsa but county/R/U present: "
            f"{missing['blank_cbsa_with_county_or_RU_present']}"
        )
        print(
            "- rows with blank cbsa and county/R/U also blank: "
            f"{missing['blank_cbsa_with_county_and_RU_also_blank']}"
        )

        return self.results

