"""
AdamOps Ensemble Models Module

Provides voting, stacking, blending, and weighted averaging ensembles.
"""

from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
import pandas as pd
from sklearn.ensemble import VotingClassifier, VotingRegressor, StackingClassifier, StackingRegressor
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.model_selection import cross_val_predict

from adamops.utils.logging import get_logger
from adamops.models.modelops import CLASSIFICATION_MODELS, REGRESSION_MODELS

logger = get_logger(__name__)


class VotingEnsemble:
    """Voting ensemble for classification or regression."""
    
    def __init__(self, estimators: List[Tuple[str, Any]], voting: str = "soft",
                 weights: Optional[List[float]] = None, task: str = "classification"):
        self.estimators = estimators
        self.voting = voting
        self.weights = weights
        self.task = task
        
        if task == "classification":
            self.model = VotingClassifier(estimators, voting=voting, weights=weights)
        else:
            self.model = VotingRegressor(estimators, weights=weights)
    
    def fit(self, X, y):
        self.model.fit(X, y)
        return self
    
    def predict(self, X):
        return self.model.predict(X)
    
    def predict_proba(self, X):
        if self.task == "classification" and hasattr(self.model, "predict_proba"):
            return self.model.predict_proba(X)
        raise ValueError("Not available for regression")


class StackingEnsemble:
    """Stacking ensemble with meta-learner."""
    
    def __init__(self, estimators: List[Tuple[str, Any]], 
                 final_estimator: Optional[Any] = None,
                 task: str = "classification", cv: int = 5):
        self.estimators = estimators
        self.task = task
        self.cv = cv
        
        if final_estimator is None:
            final_estimator = LogisticRegression() if task == "classification" else Ridge()
        
        if task == "classification":
            self.model = StackingClassifier(
                estimators, final_estimator=final_estimator, cv=cv
            )
        else:
            self.model = StackingRegressor(
                estimators, final_estimator=final_estimator, cv=cv
            )
    
    def fit(self, X, y):
        self.model.fit(X, y)
        return self
    
    def predict(self, X):
        return self.model.predict(X)


class BlendingEnsemble:
    """Blending ensemble (holdout-based stacking)."""
    
    def __init__(self, estimators: List[Tuple[str, Any]],
                 final_estimator: Optional[Any] = None,
                 task: str = "classification", blend_ratio: float = 0.2):
        self.estimators = estimators
        self.task = task
        self.blend_ratio = blend_ratio
        self.final_estimator = final_estimator or (
            LogisticRegression() if task == "classification" else Ridge()
        )
        self.fitted_estimators = []
    
    def fit(self, X, y):
        from sklearn.model_selection import train_test_split
        
        # Split for blending
        X_train, X_blend, y_train, y_blend = train_test_split(
            X, y, test_size=self.blend_ratio, random_state=42
        )
        
        # Fit base models and get blend predictions
        blend_features = []
        self.fitted_estimators = []
        
        for name, estimator in self.estimators:
            estimator.fit(X_train, y_train)
            self.fitted_estimators.append((name, estimator))
            
            if self.task == "classification" and hasattr(estimator, "predict_proba"):
                preds = estimator.predict_proba(X_blend)[:, 1]
            else:
                preds = estimator.predict(X_blend)
            blend_features.append(preds)
        
        # Stack blend predictions
        blend_X = np.column_stack(blend_features)
        
        # Fit meta-learner
        self.final_estimator.fit(blend_X, y_blend)
        
        return self
    
    def predict(self, X):
        # Get predictions from base models
        features = []
        for name, estimator in self.fitted_estimators:
            if self.task == "classification" and hasattr(estimator, "predict_proba"):
                preds = estimator.predict_proba(X)[:, 1]
            else:
                preds = estimator.predict(X)
            features.append(preds)
        
        meta_X = np.column_stack(features)
        return self.final_estimator.predict(meta_X)


class WeightedAverageEnsemble:
    """Weighted average ensemble."""
    
    def __init__(self, estimators: List[Tuple[str, Any]], 
                 weights: Optional[List[float]] = None,
                 task: str = "classification"):
        self.estimators = estimators
        self.weights = weights or [1.0 / len(estimators)] * len(estimators)
        self.task = task
        self.fitted_estimators = []
    
    def fit(self, X, y):
        self.fitted_estimators = []
        for name, estimator in self.estimators:
            estimator.fit(X, y)
            self.fitted_estimators.append((name, estimator))
        return self
    
    def predict(self, X):
        predictions = []
        for (name, estimator), weight in zip(self.fitted_estimators, self.weights):
            pred = estimator.predict(X)
            predictions.append(pred * weight)
        
        weighted_sum = sum(predictions)
        
        if self.task == "classification":
            return (weighted_sum > 0.5).astype(int)
        return weighted_sum
    
    def predict_proba(self, X):
        if self.task != "classification":
            raise ValueError("Not available for regression")
        
        probas = []
        for (name, estimator), weight in zip(self.fitted_estimators, self.weights):
            if hasattr(estimator, "predict_proba"):
                probas.append(estimator.predict_proba(X) * weight)
        
        return sum(probas)


def create_voting_ensemble(
    algorithms: List[str], task: str = "classification",
    voting: str = "soft", weights: Optional[List[float]] = None
) -> VotingEnsemble:
    """Create voting ensemble from algorithm names."""
    models = CLASSIFICATION_MODELS if task == "classification" else REGRESSION_MODELS
    estimators = [(alg, models[alg]()) for alg in algorithms if alg in models]
    return VotingEnsemble(estimators, voting=voting, weights=weights, task=task)


def create_stacking_ensemble(
    algorithms: List[str], task: str = "classification",
    final_estimator: Optional[Any] = None, cv: int = 5
) -> StackingEnsemble:
    """Create stacking ensemble from algorithm names."""
    models = CLASSIFICATION_MODELS if task == "classification" else REGRESSION_MODELS
    estimators = [(alg, models[alg]()) for alg in algorithms if alg in models]
    return StackingEnsemble(estimators, final_estimator, task, cv)


def auto_ensemble(
    X, y, task: str = "classification", top_n: int = 3, cv: int = 5
) -> Tuple[Any, Dict]:
    """
    Automatically select and create best ensemble.
    
    Returns:
        (ensemble, results): Best ensemble and evaluation results.
    """
    from adamops.models.modelops import compare_models
    
    # Compare base models
    comparison = compare_models(X, y, task, cv=cv)
    top_algorithms = comparison.head(top_n)["algorithm"].tolist()
    
    logger.info(f"Selected top {top_n} algorithms: {top_algorithms}")
    
    # Create and evaluate ensembles
    results = {}
    
    # Voting
    voting = create_voting_ensemble(top_algorithms, task)
    voting.fit(X, y)
    results["voting"] = voting
    
    # Stacking
    stacking = create_stacking_ensemble(top_algorithms, task, cv=cv)
    stacking.fit(X, y)
    results["stacking"] = stacking
    
    return stacking, {"algorithms": top_algorithms, "ensembles": list(results.keys())}
