"""
Crop yield data preprocessing.

Cleans raw crop production data (ICRISAT VDSA or APY government format)
into a standardized long-format table: one row per district-year-crop.

Usage:
    python -m ml_pipeline.preprocessing.clean_crop_data \
        --input data/raw/icrisat/yield.csv \
        --output data/processed/crop_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__)

SEASON_TO_MONTH = {"kharif": 6, "rabi": 11, "summer": 3, "whole year": 6, "autumn": 9, "winter": 11}


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().replace(" ", "_") for c in df.columns]

    # Some public crop production datasets (e.g. the common Kaggle-style
    # "Crop_Year, Season, State, Area, Production, Yield" export) are
    # STATE-level only, with no district column at all. The rest of this
    # pipeline (feature engineering, model training, live inference) is
    # built around district-level granularity. As a deliberate stopgap —
    # see docs/data_sourcing.md — we synthesize "district" as a copy of
    # "state" when no real district column exists. This lets the full
    # pipeline run end-to-end today to prove correctness, at the cost of
    # state-level precision instead of district-level. Replace with a
    # true district-level source (e.g. a direct aps.dac.gov.in district
    # export) before relying on this for real predictions.
    if "district" not in df.columns:
        if "state" not in df.columns:
            raise ValueError(
                "Input has neither a 'district' nor a 'state' column — cannot "
                f"proceed. Found columns: {list(df.columns)}"
            )
        logger.warning(
            "No district column found in input — using 'state' as a stopgap "
            "district value. Model granularity will be state-level, not "
            "district-level, until real district data is sourced."
        )
        df["district"] = df["state"]

    if "crop_year" in df.columns and "year" not in df.columns:
        df = df.rename(columns={"crop_year": "year"})

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

    # Different sources name the yield column differently — normalize to yield_ton_ha
    yield_aliases = ["yield_ton_ha", "yield_kg_ha", "yield", "productivity"]
    found_yield_col = next((c for c in yield_aliases if c in df.columns), None)
    if found_yield_col is None:
        raise ValueError(
            f"No recognizable yield column found. Looked for {yield_aliases}, "
            f"found columns: {list(df.columns)}. Rename the source column or "
            f"add the new alias to this script."
        )
    if found_yield_col != "yield_ton_ha":
        df = df.rename(columns={found_yield_col: "yield_ton_ha"})
        if found_yield_col == "yield_kg_ha":
            logger.info("Converting yield_kg_ha to yield_ton_ha (dividing by 1000)")
            df["yield_ton_ha"] = df["yield_ton_ha"] / 1000

    if "area_ha" not in df.columns:
        area_aliases = ["area", "area_hectares"]
        found_area_col = next((c for c in area_aliases if c in df.columns), None)
        if found_area_col:
            df = df.rename(columns={found_area_col: "area_ha"})
        else:
            logger.warning("No area column found — area_ha will be absent from output")

    return df


def remove_outliers_per_crop(df: pd.DataFrame, std_threshold: float = 3.0) -> pd.DataFrame:
    """
    Drops rows more than std_threshold standard deviations from their
    crop's mean yield. Implemented via boolean masking rather than
    groupby().apply() returning filtered sub-frames — recent pandas
    versions (2.2+) can silently drop the grouping column itself when
    apply() returns a row-filtered subset of the original frame, which
    would corrupt 'crop' out of the output entirely.
    """
    group_stats = df.groupby("crop")["yield_ton_ha"].agg(["mean", "std"]).reset_index()
    group_stats.columns = ["crop", "_group_mean", "_group_std"]

    merged = df.merge(group_stats, on="crop", how="left")
    # A crop with only one sample has std=NaN — keep those rows rather
    # than dropping them, since there's no meaningful outlier test for n=1.
    keep_mask = (
        merged["_group_std"].isna()
        | (merged["_group_std"] == 0)
        | (np.abs(merged["yield_ton_ha"] - merged["_group_mean"]) < std_threshold * merged["_group_std"])
    )

    before = len(df)
    result = merged[keep_mask].drop(columns=["_group_mean", "_group_std"]).reset_index(drop=True)
    after = len(result)
    logger.info(f"Outlier removal: dropped {before - after} rows ({(before - after) / before:.1%})")
    return result


def clean(df: pd.DataFrame) -> pd.DataFrame:
    df["district"] = df["district"].astype(str).str.strip().str.lower()
    df["crop"] = df["crop"].astype(str).str.strip().str.lower()

    if "season" in df.columns:
        df["season"] = df["season"].astype(str).str.strip().str.lower()
        df["month"] = df["season"].map(SEASON_TO_MONTH)
        unmapped = df["month"].isnull().sum()
        if unmapped > 0:
            logger.warning(
                f"{unmapped} rows have an unrecognized season value and got month=NaN. "
                f"Recognized seasons: {list(SEASON_TO_MONTH.keys())}"
            )
    else:
        logger.warning("No 'season' column present — month will be absent, defaulting to 6 (kharif)")
        df["month"] = 6

    df["yield_ton_ha"] = pd.to_numeric(df["yield_ton_ha"], errors="coerce")

    before = len(df)
    df = df[df["yield_ton_ha"] > 0]
    after = len(df)
    logger.info(f"Dropped {before - after} rows with zero, negative, or non-numeric yield")

    df = remove_outliers_per_crop(df)
    df = df.dropna(subset=["yield_ton_ha", "month"]).reset_index(drop=True)
    return df


def main() -> None:
    parser = argparse.ArgumentParser(description="Clean raw crop yield data")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

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

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

    logger.info("Cleaning: type coercion, season mapping, outlier removal")
    clean_df = clean(raw_df)

    if len(clean_df) == 0:
        logger.error("Zero rows remain after cleaning. Check the source file and column mappings.")
        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}")
    logger.info(f"Crops present: {sorted(clean_df['crop'].unique())}")


if __name__ == "__main__":
    main()