"""
model_train.py
Generates synthetic Indian crop data, trains a Random Forest model,
and saves model + encoders to disk.
Run this once before starting app.py
"""

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score, mean_absolute_error
from sklearn.preprocessing import LabelEncoder
import joblib
import os

np.random.seed(42)
N = 5000

# ── Categorical options ───────────────────────────────────────────────────────
CROPS    = ["Wheat", "Rice", "Maize", "Soybean", "Cotton"]
SOILS    = ["Loamy", "Clayey", "Sandy", "Black", "Red"]
SEASONS  = ["Kharif", "Rabi", "Zaid"]
DISEASE  = ["None", "Mild", "Moderate", "Severe"]
IRRIGATE = ["Rainfed", "Canal", "Drip", "Sprinkler"]

# ── Base yields (kg/acre) per crop ────────────────────────────────────────────
BASE_YIELD = {
    "Wheat":   950,
    "Rice":    1100,
    "Maize":   800,
    "Soybean": 600,
    "Cotton":  500,
}

# ── Synthetic data generation ─────────────────────────────────────────────────
def generate_data(n=N):
    crop      = np.random.choice(CROPS,    n)
    soil      = np.random.choice(SOILS,    n)
    season    = np.random.choice(SEASONS,  n)
    disease   = np.random.choice(DISEASE,  n)
    irrigate  = np.random.choice(IRRIGATE, n)

    rainfall    = np.random.uniform(100, 1200, n)
    temperature = np.random.uniform(10,  45,   n)
    fertilizer  = np.random.uniform(0,   150,  n)
    area        = np.random.uniform(1,   50,   n)

    # ── Yield simulation ──────────────────────────────────────────────────────
    base = np.array([BASE_YIELD[c] for c in crop], dtype=float)

    # rainfall: sweet-spot around 600 mm
    rain_factor = 1 - 0.0003 * (rainfall - 600) ** 2 / 1000
    rain_factor = np.clip(rain_factor, 0.5, 1.1)

    # temperature: optimal ~28 °C
    temp_factor = 1 - 0.002 * (temperature - 28) ** 2
    temp_factor = np.clip(temp_factor, 0.6, 1.1)

    # fertilizer: log benefit
    fert_factor = 1 + 0.003 * fertilizer - 0.00001 * fertilizer ** 2
    fert_factor = np.clip(fert_factor, 0.8, 1.3)

    # disease penalty
    disease_penalty = {"None": 1.0, "Mild": 0.85, "Moderate": 0.65, "Severe": 0.40}
    dis_factor = np.array([disease_penalty[d] for d in disease])

    # irrigation boost
    irr_boost = {"Rainfed": 1.0, "Canal": 1.1, "Drip": 1.2, "Sprinkler": 1.15}
    irr_factor = np.array([irr_boost[i] for i in irrigate])

    # soil-crop compatibility bonus
    compat = {
        ("Wheat",   "Loamy"): 1.1, ("Wheat",  "Black"):  1.05,
        ("Rice",    "Clayey"):1.15,("Rice",   "Black"):  1.1,
        ("Maize",   "Loamy"): 1.1, ("Soybean","Black"):  1.1,
        ("Cotton",  "Black"): 1.2, ("Cotton", "Red"):    1.0,
    }
    soil_factor = np.array([compat.get((c, s), 1.0) for c, s in zip(crop, soil)])

    yield_kg = (base
                * rain_factor
                * temp_factor
                * fert_factor
                * dis_factor
                * irr_factor
                * soil_factor
                + np.random.normal(0, 40, n))      # noise
    yield_kg = np.clip(yield_kg, 50, 2500)

    return pd.DataFrame({
        "crop": crop, "soil": soil, "season": season,
        "disease": disease, "irrigation": irrigate,
        "rainfall": rainfall, "temperature": temperature,
        "fertilizer": fertilizer, "area": area,
        "yield_kg_per_acre": yield_kg,
    })

# ── Encode + train ────────────────────────────────────────────────────────────
def train():
    print("Generating synthetic data...")
    df = generate_data()

    cat_cols = ["crop", "soil", "season", "disease", "irrigation"]
    encoders = {}
    for col in cat_cols:
        le = LabelEncoder()
        df[col + "_enc"] = le.fit_transform(df[col])
        encoders[col] = le

    features = [c + "_enc" for c in cat_cols] + [
        "rainfall", "temperature", "fertilizer", "area"
    ]
    X = df[features]
    y = df["yield_kg_per_acre"]

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    print("Training Random Forest (200 trees)...")
    model = RandomForestRegressor(
        n_estimators=200,
        max_depth=12,
        min_samples_split=4,
        random_state=42,
        n_jobs=-1,
    )
    model.fit(X_train, y_train)

    preds = model.predict(X_test)
    r2  = r2_score(y_test, preds)
    mae = mean_absolute_error(y_test, preds)
    print(f"\n✅  R² = {r2:.3f}   MAE = {mae:.1f} kg/acre")

    os.makedirs("model", exist_ok=True)
    joblib.dump(model,    "model/yield_model.pkl")
    joblib.dump(encoders, "model/encoders.pkl")
    joblib.dump(features, "model/features.pkl")
    print("Model saved → model/yield_model.pkl")
    return r2, mae

if __name__ == "__main__":
    train()