from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict, List, Any
import pandas as pd
import numpy as np
import pickle
import json
import shap
import os
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI(title="Fisherman Response Prediction API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load configuration and models
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

try:
    with open(os.path.join(BASE_DIR, "feature_cols_eng.json"), "r") as f:
        FEATURE_COLS = json.load(f)
    
    with open(os.path.join(BASE_DIR, "target_cols_A.json"), "r") as f:
        TARGET_COLS_A = json.load(f)
        
    with open(os.path.join(BASE_DIR, "target_cols_B.json"), "r") as f:
        TARGET_COLS_B = json.load(f)
        
    with open(os.path.join(BASE_DIR, "modelA_catboost.pkl"), "rb") as f:
        model_A = pickle.load(f)
        
    with open(os.path.join(BASE_DIR, "modelB_xgboost.pkl"), "rb") as f:
        model_B = pickle.load(f)
        
    # Load background data for SHAP
    df_train = pd.read_csv(os.path.join(BASE_DIR, "X_train_eng.csv"))
    # CatBoost explainer uses a sample
    bg_data_A = df_train.sample(50, random_state=42)
    bg_data_B = df_train.sample(50, random_state=42)
    
except Exception as e:
    print(f"Error loading models/data: {e}")
    # Still allow app to start, but it will fail on /predict

class PredictRequest(BaseModel):
    features: Dict[str, float]

def get_shap_factors(model, feature_names, X_input, bg_data, target_cols, top_n=5):
    """
    Menghitung SHAP values dinamis untuk model ClassifierChain.
    Karena base estimatornya dilatih dengan augmented features,
    kita juga harus memberikan padding nol untuk background dan test set.
    """
    n_base = len(feature_names)
    X_input_np = X_input.values
    
    # Rata-ratakan nilai absolute SHAP lintas semua base_estimator/label
    total_shap_abs = np.zeros(n_base)
    
    try:
        for i, estimator in enumerate(model.estimators_):
            # Augment background and input data with zeros representing previous labels
            bg_aug = np.hstack([bg_data.values, np.zeros((len(bg_data), i))]) if i > 0 else bg_data.values
            X_aug = np.hstack([X_input_np, np.zeros((1, i))]) if i > 0 else X_input_np
            
            explainer = shap.TreeExplainer(estimator, bg_aug)
            sv = explainer.shap_values(X_aug)
            
            # xgboost tree explainer might return a list [sv_class0, sv_class1], we want class 1
            if isinstance(sv, list):
                sv = sv[1]
                
            # Keep only the original base features (discard augmented feature SHAP values)
            base_sv = sv[0, :n_base]
            
            # Kita gabungkan pengaruh absolute-nya
            total_shap_abs += np.abs(base_sv)
            
        # Dapatkan index Top N
        top_idx = np.argsort(total_shap_abs)[-top_n:][::-1]
        
        top_factors = [
            {"feature": feature_names[i], "importance": float(total_shap_abs[i])}
            for i in top_idx if total_shap_abs[i] > 0
        ]
        return top_factors
    except Exception as e:
        print(f"SHAP error: {e}")
        return []

@app.post("/predict")
def predict(request: PredictRequest):
    # Ensure all required features are present
    input_dict = request.features
    
    # Fill missing features with 0
    row = {col: input_dict.get(col, 0.0) for col in FEATURE_COLS}
    
    # Create DataFrame (1 row)
    X_df = pd.DataFrame([row])
    
    # --- MODEL A PREDICTION ---
    probas_A = model_A.predict_proba(X_df)[0]
    best_idx_A = np.argmax(probas_A)
    cat_A = TARGET_COLS_A[best_idx_A]
    
    prob_dict_A = {TARGET_COLS_A[i]: float(probas_A[i]) for i in range(len(TARGET_COLS_A))}
    
    # --- MODEL B PREDICTION ---
    probas_B = model_B.predict_proba(X_df)[0]
    
    # Get Top 5 Actions
    top_indices_B = np.argsort(probas_B)[-5:][::-1]
    top_5_actions = [
        {"action": TARGET_COLS_B[i].replace("_", " "), "probability": float(probas_B[i])}
        for i in top_indices_B
    ]
    
    # --- EXPLAINABILITY (SHAP) ---
    top_factors = get_shap_factors(model_A, FEATURE_COLS, X_df, bg_data_A, TARGET_COLS_A, top_n=5)
    
    # --- PERSONALIZED MESSAGE ---
    lik_trust = input_dict.get("Kepercayaan LIK (1-5)", 3.0)
    trust_level = "Tinggi" if lik_trust >= 4.0 else "Sedang" if lik_trust >= 3.0 else "Rendah"
    
    message = (
        f"Peringatan dinamis. Anda memiliki kepercayaan LIK yang {trust_level} ({lik_trust}). "
        "Disarankan untuk memperhatikan tanda alam yang biasa Anda gunakan. "
    )
    if top_5_actions:
         message += f"Sangat disarankan untuk: {top_5_actions[0]['action']}."
    
    return {
        "model_A": {
            "predicted_category": cat_A.replace("_", " "),
            "probabilities": prob_dict_A
        },
        "model_B": {
            "top5_actions": top_5_actions
        },
        "explanation": {
            "top_factors": top_factors
        },
        "message": message
    }