# src/models/ensemble.py

import numpy as np
import pandas as pd
import lightgbm as lgb
import xgboost as xgb
from catboost import CatBoostClassifier
from sklearn.linear_model import LogisticRegression
import joblib
import os

class SuperLearnerEnsemble:
    """
    Ensemble of XGBoost, LightGBM, CatBoost with a logistic regression meta-learner.
    Now stores feature names to prevent shape mismatches.
    """

    def __init__(self):
        self.xgb_model = None
        self.lgb_model = None
        self.cat_model = None
        self.meta_learner = LogisticRegression(max_iter=1000)
        self.is_fitted = False
        self.feature_names = None  # NEW: store feature names

    def fit(self, X_train, y_train, X_val=None, y_val=None, feature_names=None):
        """
        Train base models and meta-learner.

        Args:
            X_train: Training features (numpy array or DataFrame)
            y_train: Training targets
            X_val: Validation features for meta-learner training
            y_val: Validation targets for meta-learner training
            feature_names: List of feature names (optional, extracted if X is DataFrame)
        """
        # Extract feature names if available
        if feature_names is not None:
            self.feature_names = feature_names
        elif hasattr(X_train, 'columns'):
            self.feature_names = list(X_train.columns)
        else:
            self.feature_names = [f"feature_{i}" for i in range(X_train.shape[1])]

        # Convert to numpy if DataFrame
        if hasattr(X_train, 'values'):
            X_train = X_train.values
        if X_val is not None and hasattr(X_val, 'values'):
            X_val = X_val.values

        # Initialize base models
        self.xgb_model = xgb.XGBClassifier(
            n_estimators=200, max_depth=6, learning_rate=0.05,
            objective='binary:logistic', random_state=42,
            use_label_encoder=False, eval_metric='logloss'
        )
        self.lgb_model = lgb.LGBMClassifier(
            n_estimators=200, max_depth=6, learning_rate=0.05,
            random_state=42, verbose=-1
        )
        self.cat_model = CatBoostClassifier(
            n_estimators=200, depth=6, learning_rate=0.05,
            random_state=42, verbose=0
        )

        # Train base models
        self.xgb_model.fit(X_train, y_train)
        self.lgb_model.fit(X_train, y_train)
        self.cat_model.fit(X_train, y_train)

        # Meta-learner training
        if X_val is not None and y_val is not None:
            meta_X = np.column_stack([
                self.xgb_model.predict_proba(X_val)[:, 1],
                self.lgb_model.predict_proba(X_val)[:, 1],
                self.cat_model.predict_proba(X_val)[:, 1]
            ])
            self.meta_learner.fit(meta_X, y_val)
        else:
            split = int(len(X_train) * 0.8)
            X_train_base, X_meta = X_train[:split], X_train[split:]
            y_train_base, y_meta = y_train[:split], y_train[split:]

            self.xgb_model.fit(X_train_base, y_train_base)
            self.lgb_model.fit(X_train_base, y_train_base)
            self.cat_model.fit(X_train_base, y_train_base)

            meta_X = np.column_stack([
                self.xgb_model.predict_proba(X_meta)[:, 1],
                self.lgb_model.predict_proba(X_meta)[:, 1],
                self.cat_model.predict_proba(X_meta)[:, 1]
            ])
            self.meta_learner.fit(meta_X, y_meta)

        self.is_fitted = True
        return self

    def predict_proba(self, X):
        if not self.is_fitted:
            raise ValueError("Model not fitted.")

        # Convert to numpy if DataFrame
        if hasattr(X, 'values'):
            X = X.values

        xgb_proba = self.xgb_model.predict_proba(X)[:, 1]
        lgb_proba = self.lgb_model.predict_proba(X)[:, 1]
        cat_proba = self.cat_model.predict_proba(X)[:, 1]

        meta_X = np.column_stack([xgb_proba, lgb_proba, cat_proba])
        return self.meta_learner.predict_proba(meta_X)[:, 1]

    def predict(self, X):
        proba = self.predict_proba(X)
        return (proba >= 0.5).astype(int)

    def save(self, path_prefix: str):
        os.makedirs(os.path.dirname(path_prefix), exist_ok=True)
        joblib.dump(self.xgb_model, f"{path_prefix}_xgb.joblib")
        joblib.dump(self.lgb_model, f"{path_prefix}_lgb.joblib")
        joblib.dump(self.cat_model, f"{path_prefix}_cat.joblib")
        joblib.dump(self.meta_learner, f"{path_prefix}_meta.joblib")
        # Save feature names
        joblib.dump(self.feature_names, f"{path_prefix}_features.joblib")
        print(f"✅ Ensemble saved with prefix '{path_prefix}'")

    def load(self, path_prefix: str):
        self.xgb_model = joblib.load(f"{path_prefix}_xgb.joblib")
        self.lgb_model = joblib.load(f"{path_prefix}_lgb.joblib")
        self.cat_model = joblib.load(f"{path_prefix}_cat.joblib")
        self.meta_learner = joblib.load(f"{path_prefix}_meta.joblib")
        # Load feature names if available
        try:
            self.feature_names = joblib.load(f"{path_prefix}_features.joblib")
        except:
            self.feature_names = None
            print("⚠️ No feature names file found.")
        self.is_fitted = True
        print(f"✅ Ensemble loaded from '{path_prefix}_*.joblib'")
        return self