"""
Market price data preprocessing.

Aggregates daily Agmarknet mandi prices into monthly district-crop price
series, which is the granularity every downstream model (SARIMAX market
model, Random Forest recommendation model) expects.

Usage:
    python -m ml_pipeline.preprocessing.clean_market_data \
        --input data/raw/agmarknet/prices.csv \
        --output data/processed/market_clean.csv
"""

import argparse
import logging
import sys
from pathlib import Path

import pandas as pd

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


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)
    # The raw CSV export from data.gov.in encodes spaces in original XML
    # field names as "_x0020_" (e.g. "Min_x0020_Price"). Strip that out
    # before the generic lowercase/space-to-underscore pass below, or
    # "min_x0020_price" never matches any alias.
    df.columns = [
        c.strip().lower().replace("_x0020_", "_").replace(" ", "_")
        for c in df.columns
    ]

    # Confirmed against a real data.gov.in export of resource
    # 9ef84268-d588-465a-a308-a864a43d0070 ("Current Daily Price of Various
    # Commodities from Various Markets (Mandi)") — columns after the
    # cleanup above are: state, district, market, commodity, variety,
    # grade, arrival_date, min_price, max_price, modal_price.
    column_aliases = {
        "state": ["state"],
        "district": ["district", "district_name"],
        "commodity": ["commodity", "commodity_name"],
        "date": ["date", "price_date", "arrival_date"],
        "modal_price": ["modal_price", "modalprice"],
        "min_price": ["min_price"],
        "max_price": ["max_price"],
    }

    for canonical, aliases in column_aliases.items():
        found = next((a for a in aliases if a in df.columns), None)
        if found and found != canonical:
            df = df.rename(columns={found: canonical})

    required = {"district", "commodity", "date", "modal_price"}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(
            f"Input is missing required columns {missing} after alias mapping. "
            f"Found: {list(df.columns)}. Update column_aliases in this script "
            f"if the source uses different field names."
        )

    return df


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

    df["date"] = pd.to_datetime(df["date"], errors="coerce", dayfirst=True)
    bad_dates = df["date"].isnull().sum()
    if bad_dates > 0:
        logger.warning(f"Dropping {bad_dates} rows with unparseable dates")
        df = df.dropna(subset=["date"])

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

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

    df["year"] = df["date"].dt.year
    df["month"] = df["date"].dt.month

    monthly = (
        df.groupby(["district", "crop", "year", "month"])
        .agg(
            price_per_quintal=("modal_price", "mean"),
            price_std=("modal_price", "std"),
            sample_count=("modal_price", "count"),
        )
        .reset_index()
    )
    monthly["price_std"] = monthly["price_std"].fillna(0)
    monthly["price_volatility"] = (
        monthly["price_std"] / monthly["price_per_quintal"].replace(0, pd.NA)
    ).fillna(0)

    low_sample = (monthly["sample_count"] < 3).sum()
    if low_sample > 0:
        logger.warning(
            f"{low_sample} district-crop-month combinations have fewer than 3 "
            f"daily price observations — their monthly average is less reliable. "
            f"Not dropped, but flagged in 'sample_count' for downstream awareness."
        )

    return monthly


def main() -> None:
    parser = argparse.ArgumentParser(description="Clean and aggregate raw market price data")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument(
        "--state", default=None,
        help="Filter to a single state before cleaning, e.g. 'Tamil Nadu'. "
             "Use this when the raw file was downloaded without a state filter applied.",
    )
    parser.add_argument(
        "--aggregate-to-state", action="store_true",
        help="Roll district-level prices up to state-level (replace 'district' "
             "with 'state' as the grouping key before aggregating). Use this "
             "when your climate and/or crop data is only available at state "
             "level — every source must share the same join granularity for "
             "build_feature_matrix.py's merge to work.",
    )
    args = parser.parse_args()

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

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

    if args.state:
        if "state" not in raw_df.columns:
            raise ValueError(
                "--state filter was requested but no 'state' column exists in "
                "the input file after column alias mapping."
            )
        before = len(raw_df)
        raw_df = raw_df[raw_df["state"].astype(str).str.strip().str.lower() == args.state.strip().lower()]
        logger.info(f"Filtered to state='{args.state}': {before} -> {len(raw_df)} rows")
        if len(raw_df) == 0:
            raise ValueError(
                f"No rows matched state='{args.state}'. Check the exact spelling "
                f"used in this dataset (e.g. some exports use 'Keralam' not 'Kerala')."
            )

    if args.aggregate_to_state:
        if "state" not in raw_df.columns:
            raise ValueError(
                "--aggregate-to-state was requested but no 'state' column exists "
                "in the input file after column alias mapping."
            )
        logger.info(
            "Rolling up to state-level: 'district' values will be replaced with "
            "'state' before aggregation, matching state-level climate/crop data."
        )
        raw_df = raw_df.drop(columns=["district"]).rename(columns={"state": "district"})

    logger.info("Cleaning and aggregating to monthly district-crop prices")
    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)} monthly price records to {output_path}")


if __name__ == "__main__":
    main()