"""
PHASE 3: ML Model Training
Goal: Predict environmental impact category (Low / Medium / High Impact)
Algorithm: RandomForestClassifier
HF Spaces compatible — uses BASE_DIR for all file paths.
"""

import os
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import classification_report, accuracy_score
import joblib
import warnings
warnings.filterwarnings("ignore")

# ── HF-safe absolute paths ──────────────────────────────────────────────────
BASE_DIR          = os.path.dirname(os.path.abspath(__file__))
MODEL_PATH        = os.path.join(BASE_DIR, "models", "eco_classifier.pkl")
ENCODER_PATH      = os.path.join(BASE_DIR, "models", "category_encoder.pkl")
LABEL_ENCODER_PATH= os.path.join(BASE_DIR, "models", "label_encoder.pkl")
DATA_PATH         = os.path.join(BASE_DIR, "dataset", "products.csv")
# ────────────────────────────────────────────────────────────────────────────


def preprocess_data(df: pd.DataFrame):
    df = df.copy()
    cat_encoder = LabelEncoder()
    df["category_encoded"] = cat_encoder.fit_transform(df["category"])

    feature_cols = [
        "category_encoded",
        "deforestation_risk",
        "pollution_level",
        "biodiversity_impact",
    ]
    X = df[feature_cols].values

    label_encoder = LabelEncoder()
    y = label_encoder.fit_transform(df["eco_label"])

    return X, y, cat_encoder, label_encoder, feature_cols


def train_model(data_path: str = DATA_PATH):
    os.makedirs(os.path.join(BASE_DIR, "models"), exist_ok=True)

    print("📂 Loading dataset...")
    df = pd.read_csv(data_path)
    print(f"   Dataset shape: {df.shape}")
    print(f"   Label distribution:\n{df['eco_label'].value_counts()}\n")

    print("⚙️  Preprocessing data...")
    X, y, cat_encoder, label_encoder, feature_cols = preprocess_data(df)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    print(f"   Train: {len(X_train)} | Test: {len(X_test)}\n")

    print("🌲 Training RandomForest model...")
    model = RandomForestClassifier(
        n_estimators=200,
        max_depth=10,
        min_samples_split=2,
        random_state=42,
        class_weight="balanced",
    )
    model.fit(X_train, y_train)

    y_pred   = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    report   = classification_report(
        y_test, y_pred,
        target_names=label_encoder.classes_,
        output_dict=True,
    )
    report_str = classification_report(y_test, y_pred, target_names=label_encoder.classes_)

    print(f"✅ Accuracy: {accuracy:.2%}\n{report_str}")

    joblib.dump(model,         MODEL_PATH)
    joblib.dump(cat_encoder,   ENCODER_PATH)
    joblib.dump(label_encoder, LABEL_ENCODER_PATH)
    print(f"💾 Saved → {MODEL_PATH}")

    return {
        "accuracy": accuracy,
        "report":   report,
        "feature_importances": dict(zip(feature_cols, model.feature_importances_)),
        "classes": list(label_encoder.classes_),
    }


def load_model():
    if not os.path.exists(MODEL_PATH):
        print("Model not found — training now...")
        train_model()
    model         = joblib.load(MODEL_PATH)
    cat_encoder   = joblib.load(ENCODER_PATH)
    label_encoder = joblib.load(LABEL_ENCODER_PATH)
    return model, cat_encoder, label_encoder


def predict_impact(
    category: str,
    deforestation_risk: float,
    pollution_level: float,
    biodiversity_impact: float,
) -> dict:
    model, cat_encoder, label_encoder = load_model()

    known_cats = list(cat_encoder.classes_)
    if category not in known_cats:
        category = known_cats[0]

    cat_encoded = cat_encoder.transform([category])[0]
    X           = np.array([[cat_encoded, deforestation_risk, pollution_level, biodiversity_impact]])
    pred_idx    = model.predict(X)[0]
    proba       = model.predict_proba(X)[0]

    predicted_label  = label_encoder.inverse_transform([pred_idx])[0]
    confidence_dict  = {
        cls: round(float(p), 3)
        for cls, p in zip(label_encoder.classes_, proba)
    }

    return {
        "predicted_impact":  predicted_label,
        "confidence":        round(float(proba[pred_idx]), 3),
        "all_probabilities": confidence_dict,
        "known_categories":  known_cats,
    }


if __name__ == "__main__":
    print("=== EcoVision ML Model Training ===\n")
    results = train_model()
    print(f"\nFinal Accuracy: {results['accuracy']:.2%}")