"""
Climate data preprocessing.

Transforms raw IMD-format CSVs (wide format: one column per month) into
a clean long-format table (one row per district-year-month) ready for
feature engineering.

Usage:
    python -m ml_pipeline.preprocessing.clean_climate_data \
        --input data/raw/imd/rainfall.csv \
        --output data/processed/climate_clean.csv
"""

import argparse
import logging
import sys
from pathlib import Path

import numpy as np
import pandas as pd

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

MONTH_COLUMNS = [
    "jan", "feb", "mar", "apr", "may", "jun",
    "jul", "aug", "sep", "oct", "nov", "dec",
]
MONTH_NUMBER = {name: i + 1 for i, name in enumerate(MONTH_COLUMNS)}


def load_and_validate(input_path: Path) -> pd.DataFrame:
    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_path}")

    df = pd.read_csv(input_path)
    df.columns = [c.strip().lower() for c in df.columns]

    # IMD's commonly-distributed open rainfall dataset is published at
    # METEOROLOGICAL SUBDIVISION level (e.g. "Tamil Nadu" as one region),
    # not district level — column name is typically "subdivision". True
    # district-level rainfall requires the licensed IMD product (see
    # docs/data_sourcing.md). As a deliberate stopgap to prove the
    # pipeline end-to-end, we treat "subdivision" as "district" here,
    # at the cost of state-level rather than district-level granularity.
    if "district" not in df.columns:
        if "subdivision" in df.columns:
            logger.warning(
                "No 'district' column found, but 'subdivision' is present — "
                "using subdivision as a stopgap district value. Climate "
                "granularity will be state/subdivision-level, not "
                "district-level, until real district data is sourced."
            )
            df = df.rename(columns={"subdivision": "district"})
        else:
            raise ValueError(
                "Input has neither a 'district' nor a 'subdivision' column — "
                f"cannot proceed. Found columns: {list(df.columns)}"
            )

    required_base = {"district", "year"}
    missing = required_base - set(df.columns)
    if missing:
        raise ValueError(
            f"Input is missing required columns {missing}. "
            f"Found columns: {list(df.columns)}"
        )

    present_month_cols = [c for c in MONTH_COLUMNS if c in df.columns]
    if len(present_month_cols) < 12:
        missing_months = set(MONTH_COLUMNS) - set(present_month_cols)
        logger.warning(
            f"Input is missing month columns {missing_months} — "
            f"these months will be absent from the output, not interpolated."
        )

    return df


def wide_to_long(df: pd.DataFrame) -> pd.DataFrame:
    present_month_cols = [c for c in MONTH_COLUMNS if c in df.columns]

    long_df = df.melt(
        id_vars=["district", "year"],
        value_vars=present_month_cols,
        var_name="month_name",
        value_name="rainfall_mm",
    )
    long_df["month"] = long_df["month_name"].map(MONTH_NUMBER)
    long_df = long_df.drop(columns=["month_name"])
    return long_df


def clean(df: pd.DataFrame) -> pd.DataFrame:
    df["district"] = df["district"].astype(str).str.strip().str.lower()
    df["rainfall_mm"] = pd.to_numeric(df["rainfall_mm"], errors="coerce")

    negative_count = (df["rainfall_mm"] < 0).sum()
    if negative_count > 0:
        logger.warning(f"Found {negative_count} negative rainfall values — clamping to 0")
        df.loc[df["rainfall_mm"] < 0, "rainfall_mm"] = 0

    extreme_count = (df["rainfall_mm"] > 3000).sum()
    if extreme_count > 0:
        logger.warning(
            f"Found {extreme_count} rainfall values above 3000mm in a single month — "
            f"these are physically implausible for monthly totals and likely indicate "
            f"a unit error in the source file. Review before proceeding; not auto-corrected."
        )

    before_fill = df["rainfall_mm"].isnull().sum()
    df["rainfall_mm"] = df.groupby(["district", "month"])["rainfall_mm"].transform(
        lambda x: x.fillna(x.mean())
    )
    after_fill = df["rainfall_mm"].isnull().sum()
    logger.info(
        f"Filled {before_fill - after_fill} missing rainfall values using "
        f"district-month historical mean. {after_fill} remain null "
        f"(district+month combinations with no data at all)."
    )

    df = df.sort_values(["district", "year", "month"]).reset_index(drop=True)
    return df


def main() -> None:
    parser = argparse.ArgumentParser(description="Clean raw IMD climate data")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument(
        "--filter-value", default=None,
        help="Filter to a single district/subdivision value before cleaning, "
             "e.g. 'Tamil Nadu'. Use when the raw file covers all of India.",
    )
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output)

    logger.info(f"Loading raw climate data from {input_path}")
    raw_df = load_and_validate(input_path)

    if args.filter_value:
        before = len(raw_df)
        raw_df = raw_df[
            raw_df["district"].astype(str).str.strip().str.lower() == args.filter_value.strip().lower()
        ]
        logger.info(f"Filtered to '{args.filter_value}': {before} -> {len(raw_df)} rows")
        if len(raw_df) == 0:
            raise ValueError(
                f"No rows matched '{args.filter_value}'. Check the exact "
                f"spelling used in this dataset's district/subdivision column."
            )

    logger.info("Converting wide format to long format")
    long_df = wide_to_long(raw_df)

    logger.info("Cleaning: type coercion, outlier flagging, missing value imputation")
    clean_df = clean(long_df)

    if clean_df["rainfall_mm"].isnull().mean() > 0.3:
        logger.error(
            "More than 30% of rainfall values are null after cleaning. "
            "This dataset is not safe to use for training as-is. Investigate "
            "the source file before proceeding."
        )
        sys.exit(1)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    clean_df.to_csv(output_path, index=False)
    logger.info(f"Saved {len(clean_df)} cleaned rows to {output_path}")


if __name__ == "__main__":
    main()