"""
automl_lib.ensemble
~~~~~~~~~~~~~~~~~~~
Model ensembling via voting or averaging.
"""

from __future__ import annotations

from typing import Any
import numpy as np
import torch
import torch.nn as nn

from .utils import get_logger

logger = get_logger(__name__)

class EnsembleModel(nn.Module):
    """A wrapper for multiple models to perform ensembling."""

    def __init__(self, models: list[nn.Module], task: str = "classification", weights: list[float] | None = None):
        super().__init__()
        self.models = nn.ModuleList([m for m in models if not getattr(m, "_is_sklearn", False)])
        self.sklearn_models = [m for m in models if getattr(m, "_is_sklearn", False)]
        self.task = task
        self.weights = weights if weights else [1.0 / len(models)] * len(models)

    def forward(self, x):
        """Average forward pass for PyTorch models."""
        if not self.models:
            return None
        
        outputs = []
        for i, model in enumerate(self.models):
            outputs.append(model(x) * self.weights[i])
        
        return torch.stack(outputs).sum(dim=0)

    def predict(self, x: np.ndarray) -> np.ndarray:
        """Prediction logic for ensemble (mix of torch and sklearn)."""
        all_preds = []
        
        # PyTorch predictions
        self.eval()
        with torch.no_grad():
            if self.models:
                # Assuming x is compatible or can be converted
                x_tensor = torch.tensor(x, dtype=torch.float32)
                logits = self.forward(x_tensor)
                if self.task == "classification":
                    all_preds.append(torch.softmax(logits, dim=1).cpu().numpy())
                else:
                    all_preds.append(logits.cpu().numpy())
        
        # Sklearn predictions
        for i, model in enumerate(self.sklearn_models):
            if self.task == "classification":
                # Use predict_proba for better averaging
                if hasattr(model.estimator, "predict_proba"):
                    all_preds.append(model.estimator.predict_proba(x))
                elif hasattr(model, "predict_proba"):
                    all_preds.append(model.predict_proba(x))
            else:
                pred = model.estimator.predict(x) if hasattr(model, "estimator") else model.predict(x)
                all_preds.append(pred.reshape(-1, 1))

        if not all_preds:
            return np.array([])

        # Weighted average of probabilities/values
        # This is strictly a 'soft voting' or 'averaging' ensemble
        ensemble_out = np.zeros_like(all_preds[0])
        total_weight = 0
        for i, preds in enumerate(all_preds):
            w = self.weights[i] if i < len(self.weights) else 1.0
            ensemble_out += preds * w
            total_weight += w
        
        ensemble_out /= total_weight

        if self.task == "classification":
            return np.argmax(ensemble_out, axis=1)
        return ensemble_out.flatten()

def create_voting_ensemble(models: list[Any], task: str) -> EnsembleModel:
    """Helper to create an ensemble from a list of models."""
    logger.info(f"Creating ensemble from {len(models)} models using soft voting/averaging.")
    return EnsembleModel(models, task=task)
