"""
Feature engineering pipeline.

Produces data/processed/final_feature_matrix.csv — the single training
table every model in ml_pipeline/training/ reads from. Also produces
data/processed/climate_clean.csv (already created by clean_climate_data.py,
read directly by the live API's FeatureBuilder for historical averages).

This script must be re-run any time the raw source data is refreshed
(new IMD/ICRISAT/Agmarknet export). Models trained on a stale feature
matrix will silently drift from what the live system computes at
inference time — keep this pipeline run as part of your retraining
checklist, not a one-time setup step.

Usage:
    python -m ml_pipeline.feature_engineering.build_feature_matrix \
        --climate data/processed/climate_clean.csv \
        --crop data/processed/crop_clean.csv \
        --market data/processed/market_clean.csv \
        --crop-lookup data/lookup/crops.json \
        --output data/processed/final_feature_matrix.csv
"""

import argparse
import json
import logging
import sys
from pathlib import Path

import numpy as np
import pandas as pd
from scipy import stats

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

SEASON_MAP = {
    6: 0, 7: 0, 8: 0, 9: 0,
    10: 1, 11: 1, 12: 1, 1: 1,
    2: 2, 3: 2, 4: 2, 5: 2,
}

HISTORICAL_CUTOFF_YEAR = 2013  # years up to and including this are used to compute
                                 # "historical" baselines, avoiding leakage of a year's
                                 # own data into its own anomaly score where possible


def add_climate_features(climate: pd.DataFrame) -> pd.DataFrame:
    climate = climate.sort_values(["district", "year", "month"]).reset_index(drop=True)

    hist = climate[climate["year"] <= HISTORICAL_CUTOFF_YEAR].groupby(
        ["district", "month"]
    ).agg(
        hist_mean_rainfall=("rainfall_mm", "mean"),
        hist_std_rainfall=("rainfall_mm", "std"),
    ).reset_index()

    climate = pd.merge(climate, hist, on=["district", "month"], how="left")

    fallback_mean = climate["rainfall_mm"].mean()
    climate["hist_mean_rainfall"] = climate["hist_mean_rainfall"].fillna(fallback_mean)
    climate["hist_std_rainfall"] = climate["hist_std_rainfall"].fillna(climate["rainfall_mm"].std())
    climate["hist_std_rainfall"] = climate["hist_std_rainfall"].replace(0, 1)  # avoid div-by-zero

    climate["rainfall_anomaly"] = (
        (climate["rainfall_mm"] - climate["hist_mean_rainfall"]) / climate["hist_std_rainfall"]
    )

    climate["rainfall_lag1"] = climate.groupby("district")["rainfall_mm"].shift(1)
    climate["rainfall_ma3"] = (
        climate.groupby("district")["rainfall_mm"]
        .transform(lambda x: x.rolling(3, min_periods=1).mean())
    )
    climate["rainfall_lag1"] = climate["rainfall_lag1"].fillna(climate["rainfall_mm"])

    climate["season_index"] = climate["month"].map(SEASON_MAP)

    # temp_c / humidity may not exist depending on which IMD export was
    # used — synthesize reasonable defaults rather than crash, but log it
    # clearly so it is not mistaken for real measured data downstream.
    if "temp_c" not in climate.columns:
        logger.warning("No temp_c column in climate data — using a flat 30.0C default for all rows")
        climate["temp_c"] = 30.0
    if "humidity" not in climate.columns:
        logger.warning("No humidity column in climate data — using a flat 70.0 default for all rows")
        climate["humidity"] = 70.0

    return climate


def load_crop_properties(lookup_path: Path) -> pd.DataFrame:
    # utf-8-sig transparently strips a UTF-8 byte order mark if present,
    # and behaves identically to plain utf-8 if no BOM exists. Windows
    # tools (e.g. PowerShell's `Set-Content -Encoding utf8`) commonly
    # write a BOM that plain utf-8 + json.load() rejects outright.
    with open(lookup_path, encoding="utf-8-sig") as f:
        crops_dict = json.load(f)
    rows = []
    for crop, props in crops_dict.items():
        row = {"crop": crop}
        row.update(props)
        rows.append(row)
    return pd.DataFrame(rows)


def _compute_slope(series: pd.Series) -> float:
    values = series.dropna().values
    if len(values) < 3:
        return 0.0
    x = np.arange(len(values))
    slope, _, _, _, _ = stats.linregress(x, values)
    return float(slope)


def add_crop_features(crop: pd.DataFrame, crop_props: pd.DataFrame) -> pd.DataFrame:
    crop = pd.merge(crop, crop_props, on="crop", how="left")

    missing_props = crop[crop["water_req_mm"].isnull()]["crop"].unique()
    if len(missing_props) > 0:
        logger.warning(
            f"Crops present in yield data but missing from crops.json lookup: "
            f"{list(missing_props)}. Add them to data/lookup/crops.json or these "
            f"rows will have null agronomic features and get dropped before training."
        )

    crop = crop.sort_values(["district", "crop", "year"])
    trend_slopes = (
        crop.groupby(["district", "crop"])["yield_ton_ha"]
        .apply(lambda s: s.rolling(5, min_periods=3).apply(_compute_slope, raw=False))
        .reset_index(level=[0, 1], drop=True)
    )
    crop["yield_trend_slope"] = trend_slopes
    crop["yield_trend_slope"] = crop["yield_trend_slope"].fillna(0.0)

    return crop


def add_market_features(market: pd.DataFrame) -> pd.DataFrame:
    market = market.sort_values(["district", "crop", "year", "month"])
    market["price_lag1"] = market.groupby(["district", "crop"])["price_per_quintal"].shift(1)

    demand_trend = (
        market.groupby(["district", "crop"])["price_per_quintal"]
        .apply(lambda s: s.rolling(6, min_periods=3).apply(_compute_slope, raw=False))
        .reset_index(level=[0, 1], drop=True)
    )
    market["demand_trend"] = demand_trend.fillna(0.0)
    return market
def label_feasibility(row: pd.Series) -> str:
    temp_ok = row["min_temp_c"] <= row["temp_c"] <= row["max_temp_c"]
    rain_ok = row["rainfall_mm"] >= (row["water_req_mm"] / 12)
    anomaly_bad = abs(row["rainfall_anomaly"]) > 1.5
    has_yield = row.get("yield_ton_ha", 0) > 0
    extreme_drought = row["rainfall_mm"] < (row["water_req_mm"] / 12) * 0.10

    # not_suitable: temperature out of range OR extreme drought
    if not temp_ok or extreme_drought:
        return "not_suitable"

    # suitable: temp ok + rain ok + no anomaly + proven historical yield
    if temp_ok and rain_ok and not anomaly_bad and has_yield:
        return "suitable"

    # risky: temp ok + rain borderline OR no historical yield in this district
    if temp_ok and not anomaly_bad and has_yield:
        return "risky"

    # not_suitable: no yield history + rain fails + not drought tolerant
    if not rain_ok and not row["drought_tolerant"] and not has_yield:
        return "not_suitable"

    return "risky"

def label_trend(slope: float) -> str:
    if slope > 0.1:
        return "upward"
    elif slope < -0.1:
        return "downward"
    elif abs(slope) > 0.05:
        return "volatile"
    return "stable"


def label_suitability(row: pd.Series) -> str:
    if row["feasibility_label"] == "suitable" and row.get("demand_trend", 0) > 0:
        return "high"
    elif row["feasibility_label"] == "not_suitable":
        return "low"
    return "medium"


def build(climate: pd.DataFrame, crop: pd.DataFrame, market: pd.DataFrame, crop_props: pd.DataFrame) -> pd.DataFrame:
    climate = add_climate_features(climate)
    crop = add_crop_features(crop, crop_props)
    market = add_market_features(market)

    merged = pd.merge(crop, climate, on=["district", "year", "month"], how="inner")
    logger.info(f"After climate+crop merge: {len(merged)} rows")

    if len(merged) == 0:
        raise ValueError(
            "Zero rows after merging climate and crop data. Likely cause: "
            "district name spelling mismatch, or year/month ranges that "
            "don't overlap between the two sources. Inspect both inputs."
        )

    merged = pd.merge(merged, market, on=["district", "crop", "year", "month"], how="left")
    logger.info(f"After market merge: {len(merged)} rows ({merged['price_per_quintal'].isnull().sum()} missing market price)")

    merged["price_per_quintal"] = merged.groupby("crop")["price_per_quintal"].transform(
        lambda x: x.fillna(x.median())
    )
    merged["price_volatility"] = merged["price_volatility"].fillna(0.1)
    merged["demand_trend"] = merged["demand_trend"].fillna(0.0)

    before_drop = len(merged)
    required_for_labels = ["temp_c", "min_temp_c", "max_temp_c", "rainfall_mm", "water_req_mm"]
    merged = merged.dropna(subset=required_for_labels)
    logger.info(f"Dropped {before_drop - len(merged)} rows missing required label-input fields")

    merged["feasibility_label"] = merged.apply(label_feasibility, axis=1)
    merged["trend_label"] = merged["yield_trend_slope"].apply(label_trend)
    merged["suitability_label"] = merged.apply(label_suitability, axis=1)

    return merged


def main() -> None:
    parser = argparse.ArgumentParser(description="Build the final feature matrix for model training")
    parser.add_argument("--climate", required=True)
    parser.add_argument("--crop", required=True)
    parser.add_argument("--market", required=True)
    parser.add_argument("--crop-lookup", required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    logger.info("Loading inputs...")
    climate_df = pd.read_csv(args.climate)
    crop_df = pd.read_csv(args.crop)
    market_df = pd.read_csv(args.market)
    crop_props = load_crop_properties(Path(args.crop_lookup))

    logger.info("Building feature matrix...")
    final_df = build(climate_df, crop_df, market_df, crop_props)

    if len(final_df) < 100:
        logger.error(
            f"Only {len(final_df)} rows in the final feature matrix. This is "
            f"too small to train reliable models — review the merge keys and "
            f"source data coverage before proceeding to training."
        )
        sys.exit(1)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    final_df.to_csv(output_path, index=False)

    logger.info(f"Saved {len(final_df)} rows to {output_path}")
    logger.info(f"Feasibility label distribution:\n{final_df['feasibility_label'].value_counts()}")
    logger.info(f"Trend label distribution:\n{final_df['trend_label'].value_counts()}")
    logger.info(f"Suitability label distribution:\n{final_df['suitability_label'].value_counts()}")


if __name__ == "__main__":
    main()