"""
ClimaIQ Engine — Climate-Adjusted Credit Risk Model
Core prediction module for ClimaIQ Kisan
Author: Krishna Dahale
"""

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
import joblib
import os
import warnings
warnings.filterwarnings("ignore")

# ─── Feature Definitions ───────────────────────────────────────────────────────

TRADITIONAL_FEATURES = [
    "Age", "Land_Size_Acres", "Annual_Income_Lakhs",
    "Loan_Amount_Lakhs", "Debt_to_Income_Ratio",
    "Previous_Defaults", "Land_Productivity"
]

CLIMATE_FEATURES = [
    "Rainfall_Deficit_Pct", "SPI", "Consecutive_Drought_Years"
]

ENHANCED_FEATURES = [
    "State_Maharashtra", "Crop_Water_Intensive", "Crop_Drought_Interaction"
]

ALL_FEATURES = TRADITIONAL_FEATURES + CLIMATE_FEATURES + ENHANCED_FEATURES

CROP_WATER_MAP = {
    "Rice": 1.0,
    "Sugarcane": 1.0,
    "Cotton": 0.5,
    "Wheat": 0.2,
    "Millets": 0.0
}

FEATURE_DISPLAY_NAMES = {
    "Crop_Drought_Interaction": "Crop × Drought Severity",
    "Rainfall_Deficit_Pct": "Rainfall Deficit",
    "SPI": "Drought Severity (SPI)",
    "Consecutive_Drought_Years": "Consecutive Drought Years",
    "Previous_Defaults": "Previous Default History",
    "Debt_to_Income_Ratio": "Debt-to-Income Ratio",
    "State_Maharashtra": "Geographic Risk (Maharashtra)",
    "Crop_Water_Intensive": "Crop Water Dependency",
    "Annual_Income_Lakhs": "Annual Income",
    "Loan_Amount_Lakhs": "Loan Amount",
    "Land_Size_Acres": "Land Size",
    "Age": "Borrower Age",
    "Land_Productivity": "Land Productivity"
}

# ─── Data Generation ───────────────────────────────────────────────────────────

def generate_training_data(n_loans=5000, seed=42):
    """Generate synthetic agricultural loan dataset."""
    np.random.seed(seed)

    df = pd.DataFrame({
        "Loan_ID": range(1, n_loans + 1),
        "State": np.random.choice(["Maharashtra", "Punjab"], n_loans, p=[0.6, 0.4]),
        "Age": np.random.normal(45, 12, n_loans).clip(25, 70),
        "Land_Size_Acres": np.random.exponential(3, n_loans).clip(0.5, 20),
        "Annual_Income_Lakhs": np.random.lognormal(2.5, 0.8, n_loans).clip(1, 15),
        "Loan_Amount_Lakhs": np.random.uniform(0.5, 8, n_loans),
        "Previous_Defaults": np.random.choice([0, 1], n_loans, p=[0.85, 0.15]),
        "Crop_Type": np.random.choice(
            ["Rice", "Cotton", "Wheat", "Sugarcane", "Millets"],
            n_loans,
            p=[0.25, 0.20, 0.25, 0.15, 0.15],
        ),
    })

    df["Rainfall_Deficit_Pct"] = np.random.normal(-12, 15, n_loans).clip(-60, 20)
    df["SPI"] = np.random.normal(-0.3, 1.0, n_loans).clip(-3, 2)
    df["Consecutive_Drought_Years"] = np.random.choice([0, 1, 2, 3], n_loans, p=[0.5, 0.3, 0.15, 0.05])
    df["Debt_to_Income_Ratio"] = df["Loan_Amount_Lakhs"] / df["Annual_Income_Lakhs"]
    df["Land_Productivity"] = df["Annual_Income_Lakhs"] / df["Land_Size_Acres"]

    def generate_default(row):
        p = 0.08
        if row["Previous_Defaults"] == 1: p *= 4
        if row["Debt_to_Income_Ratio"] > 0.5: p *= 2
        elif row["Debt_to_Income_Ratio"] > 0.3: p *= 1.5
        if row["Rainfall_Deficit_Pct"] < -25: p *= 2.5
        elif row["Rainfall_Deficit_Pct"] < -10: p *= 1.4
        if row["SPI"] < -1.5: p *= 2
        elif row["SPI"] < -1.0: p *= 1.3
        if row["Consecutive_Drought_Years"] >= 2: p *= 2.5
        if row["Crop_Type"] in ["Rice", "Sugarcane"] and row["Rainfall_Deficit_Pct"] < -15: p *= 1.6
        if row["Age"] < 30 or row["Age"] > 60: p *= 1.2
        if row["State"] == "Maharashtra": p *= 1.3
        return 1 if np.random.rand() < min(p, 0.95) else 0

    df["Default"] = df.apply(generate_default, axis=1)

    # Feature engineering
    df["State_Maharashtra"] = (df["State"] == "Maharashtra").astype(int)
    df["Crop_Water_Intensive"] = df["Crop_Type"].map(CROP_WATER_MAP)
    df["Crop_Drought_Interaction"] = (
        df["Crop_Water_Intensive"] * df["Rainfall_Deficit_Pct"].clip(upper=0).abs()
    )

    return df


# ─── Model Training ────────────────────────────────────────────────────────────

def train_model():
    """Train the climate-adjusted credit risk model. Returns model and scaler."""
    df = generate_training_data()

    X = df[ALL_FEATURES]
    y = df["Default"]

    X_train, _, y_train, _ = train_test_split(
        X, y, test_size=0.3, stratify=y, random_state=42
    )

    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)

    model = LogisticRegression(max_iter=1000, class_weight="balanced")
    model.fit(X_train_scaled, y_train)

    return model, scaler


# ─── Save / Load ───────────────────────────────────────────────────────────────

def save_model(model, scaler, model_path="climaiq_model.pkl", scaler_path="climaiq_scaler.pkl"):
    joblib.dump(model, model_path)
    joblib.dump(scaler, scaler_path)


def load_model(model_path="climaiq_model.pkl", scaler_path="climaiq_scaler.pkl"):
    """Load saved model and scaler. Trains fresh if files not found."""
    if os.path.exists(model_path) and os.path.exists(scaler_path):
        model = joblib.load(model_path)
        scaler = joblib.load(scaler_path)
    else:
        model, scaler = train_model()
        save_model(model, scaler, model_path, scaler_path)
    return model, scaler


# ─── Feature Engineering for Single Input ─────────────────────────────────────

def engineer_single_input(data: dict) -> pd.DataFrame:
    """
    Convert raw farmer input dict into model-ready feature DataFrame.

    Expected keys:
        age, land_size_acres, annual_income_lakhs, loan_amount_lakhs,
        previous_defaults, crop_type, state,
        rainfall_deficit_pct, spi, consecutive_drought_years
    """
    debt_to_income = data["loan_amount_lakhs"] / data["annual_income_lakhs"]
    land_productivity = data["annual_income_lakhs"] / data["land_size_acres"]
    state_mh = 1 if data["state"] == "Maharashtra" else 0
    crop_water = CROP_WATER_MAP.get(data["crop_type"], 0.5)
    crop_drought = crop_water * abs(min(data["rainfall_deficit_pct"], 0))

    features = {
        "Age": data["age"],
        "Land_Size_Acres": data["land_size_acres"],
        "Annual_Income_Lakhs": data["annual_income_lakhs"],
        "Loan_Amount_Lakhs": data["loan_amount_lakhs"],
        "Debt_to_Income_Ratio": debt_to_income,
        "Previous_Defaults": data["previous_defaults"],
        "Land_Productivity": land_productivity,
        "Rainfall_Deficit_Pct": data["rainfall_deficit_pct"],
        "SPI": data["spi"],
        "Consecutive_Drought_Years": data["consecutive_drought_years"],
        "State_Maharashtra": state_mh,
        "Crop_Water_Intensive": crop_water,
        "Crop_Drought_Interaction": crop_drought,
    }

    return pd.DataFrame([features])[ALL_FEATURES]


# ─── Credit Score Calculation ──────────────────────────────────────────────────

def compute_credit_score(model, scaler, feature_row: pd.DataFrame,
                          base_score=650, pdo=50) -> int:
    """Convert default probability into 300–850 credit score."""
    coef = model.coef_[0]
    factor = pdo / np.log(2)
    offset = base_score - factor * np.log(20)

    z = scaler.transform(feature_row)[0]
    score = offset + np.sum(z * (-factor * coef))
    return int(np.clip(score, 300, 850))


def get_risk_band(score: int) -> str:
    if score >= 700:
        return "Very Low Risk"
    elif score >= 650:
        return "Low Risk"
    elif score >= 600:
        return "Medium Risk"
    else:
        return "High Risk"


# ─── Risk Driver Extraction ────────────────────────────────────────────────────

def get_top_risk_drivers(model, scaler, feature_row: pd.DataFrame, top_n=3) -> list:
    """
    Return the top N features driving default risk for this borrower.
    Each driver is a dict: {feature, display_name, contribution, direction}
    """
    z = scaler.transform(feature_row)[0]
    coef = model.coef_[0]

    # Contribution = standardized value × coefficient
    # Positive contribution = increases default probability = increases risk
    contributions = z * coef

    driver_df = pd.DataFrame({
        "feature": ALL_FEATURES,
        "contribution": contributions
    }).sort_values("contribution", ascending=False)

    # Top risk-increasing drivers only
    top_drivers = driver_df.head(top_n)

    result = []
    for _, row in top_drivers.iterrows():
        result.append({
            "feature": row["feature"],
            "display_name": FEATURE_DISPLAY_NAMES.get(row["feature"], row["feature"]),
            "contribution": round(row["contribution"], 4),
            "direction": "increases risk" if row["contribution"] > 0 else "reduces risk"
        })

    return result


# ─── Main Prediction Function ──────────────────────────────────────────────────

def predict_single(data: dict, model, scaler) -> dict:
    """
    Full prediction pipeline for a single farmer/borrower.

    Returns:
        default_probability  : float (0–100)
        credit_score         : int (300–850)
        risk_band            : str
        top_risk_drivers     : list of dicts
        recommended_action   : str
    """
    feature_row = engineer_single_input(data)

    scaled = scaler.transform(feature_row)
    default_prob = model.predict_proba(scaled)[0][1] * 100

    credit_score = compute_credit_score(model, scaler, feature_row)
    risk_band = get_risk_band(credit_score)
    top_drivers = get_top_risk_drivers(model, scaler, feature_row)

    # Recommended action based on risk band
    action_map = {
        "Very Low Risk": "Auto-Approve",
        "Low Risk": "Approve",
        "Medium Risk": "Manual Review",
        "High Risk": "Decline or Offer Climate-Linked Premium Rate"
    }

    return {
        "default_probability": round(default_prob, 2),
        "credit_score": credit_score,
        "risk_band": risk_band,
        "top_risk_drivers": top_drivers,
        "recommended_action": action_map[risk_band]
    }


# ─── Stress Test ───────────────────────────────────────────────────────────────

def run_stress_test(model, scaler, base_portfolio: list) -> list:
    """
    Run 4 climate scenarios across a list of farmer input dicts.
    Returns list of scenario result dicts.
    """
    scenarios = {
        "Normal Monsoon":     {"rainfall_deficit_pct": 0,   "spi": 0,    "consecutive_drought_years": 0},
        "Moderate Drought":   {"rainfall_deficit_pct": -20, "spi": -1.2, "consecutive_drought_years": 1},
        "Severe Drought":     {"rainfall_deficit_pct": -35, "spi": -1.8, "consecutive_drought_years": 1},
        "Back-to-Back Drought": {"rainfall_deficit_pct": -25, "spi": -1.5, "consecutive_drought_years": 2},
    }

    avg_loan = 3.0
    recovery = 0.30
    results = []

    for scenario_name, overrides in scenarios.items():
        probs = []
        for farmer in base_portfolio:
            stressed = {**farmer, **overrides}
            result = predict_single(stressed, model, scaler)
            probs.append(result["default_probability"] / 100)

        avg_default_pct = np.mean(probs) * 100
        total_loss = np.mean(probs) * avg_loan * (1 - recovery) * len(base_portfolio)

        results.append({
            "scenario": scenario_name,
            "avg_default_pct": round(avg_default_pct, 2),
            "total_loss_lakhs": round(total_loss, 2),
            "portfolio_size": len(base_portfolio)
        })

    return results


# ─── Kaggle / research export ─────────────────────────────────────────────────

KAGGLE_DATASET_COLUMNS = [
    "loan_id",
    "state",
    "age",
    "land_size_acres",
    "annual_income_lakhs",
    "loan_amount_lakhs",
    "crop_type",
    "previous_defaults",
    "rainfall_deficit_pct",
    "spi",
    "consecutive_drought_years",
    "debt_to_income_ratio",
    "land_productivity",
    "crop_water_intensity",
    "crop_drought_interaction",
    "state_maharashtra",
    "default",
    "climaiq_score",
    "risk_band",
]


def _simplify_risk_band(label: str) -> str:
    """Kaggle-friendly labels: High / Medium / Low / Very Low."""
    return label.replace(" Risk", "").strip()


def _training_row_to_farmer_dict(row: pd.Series) -> dict:
    """Map a generate_training_data row to the dict expected by predict_single."""
    return {
        "age": int(round(float(row["Age"]))),
        "land_size_acres": float(row["Land_Size_Acres"]),
        "annual_income_lakhs": float(row["Annual_Income_Lakhs"]),
        "loan_amount_lakhs": float(row["Loan_Amount_Lakhs"]),
        "previous_defaults": int(row["Previous_Defaults"]),
        "crop_type": str(row["Crop_Type"]),
        "state": str(row["State"]),
        "rainfall_deficit_pct": float(row["Rainfall_Deficit_Pct"]),
        "spi": float(row["SPI"]),
        "consecutive_drought_years": int(row["Consecutive_Drought_Years"]),
    }


def build_kaggle_export_dataframe(
    n_rows: int = 1000,
    seed: int = 42,
    model=None,
    scaler=None,
) -> pd.DataFrame:
    """
    Build a tabular dataset aligned with the ClimaIQ engine: raw inputs, engineered
    features, synthetic default target (causal logic in generate_training_data), and
    ClimaIQ credit score / risk band from the trained logistic model.

    Intended for Kaggle / research (e.g. *ClimaIQ — Climate-Adjusted Agricultural
    Credit Risk Dataset (India)*).
    """
    if model is None or scaler is None:
        model, scaler = load_model()

    src = generate_training_data(n_rows, seed=seed)
    records = []

    for _, row in src.iterrows():
        farmer = _training_row_to_farmer_dict(row)
        pred = predict_single(farmer, model, scaler)
        records.append(
            {
                "loan_id": int(row["Loan_ID"]),
                "state": str(row["State"]),
                "age": farmer["age"],
                "land_size_acres": round(farmer["land_size_acres"], 3),
                "annual_income_lakhs": round(farmer["annual_income_lakhs"], 3),
                "loan_amount_lakhs": round(farmer["loan_amount_lakhs"], 3),
                "crop_type": farmer["crop_type"],
                "previous_defaults": farmer["previous_defaults"],
                "rainfall_deficit_pct": round(farmer["rainfall_deficit_pct"], 3),
                "spi": round(farmer["spi"], 3),
                "consecutive_drought_years": farmer["consecutive_drought_years"],
                "debt_to_income_ratio": round(float(row["Debt_to_Income_Ratio"]), 4),
                "land_productivity": round(float(row["Land_Productivity"]), 4),
                "crop_water_intensity": round(float(row["Crop_Water_Intensive"]), 4),
                "crop_drought_interaction": round(float(row["Crop_Drought_Interaction"]), 4),
                "state_maharashtra": int(row["State_Maharashtra"]),
                "default": int(row["Default"]),
                "climaiq_score": int(pred["credit_score"]),
                "risk_band": _simplify_risk_band(str(pred["risk_band"])),
            }
        )

    out = pd.DataFrame(records)
    return out[KAGGLE_DATASET_COLUMNS]


# ─── Quick Test ────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("Training ClimaIQ model...")
    model, scaler = load_model()
    print("Model ready.\n")

    sample_farmer = {
        "age": 42,
        "land_size_acres": 3.5,
        "annual_income_lakhs": 3.0,
        "loan_amount_lakhs": 2.0,
        "previous_defaults": 0,
        "crop_type": "Cotton",
        "state": "Maharashtra",
        "rainfall_deficit_pct": -35.0,
        "spi": -1.8,
        "consecutive_drought_years": 1
    }

    result = predict_single(sample_farmer, model, scaler)

    print(f"Credit Score     : {result['credit_score']}")
    print(f"Risk Band        : {result['risk_band']}")
    print(f"Default Prob     : {result['default_probability']}%")
    print(f"Recommended      : {result['recommended_action']}")
    print(f"\nTop Risk Drivers:")
    for d in result["top_risk_drivers"]:
        print(f"  - {d['display_name']} ({d['direction']})")