"""Simple model ensemble."""

from __future__ import annotations

import numpy as np


class RankEnsemble:
    def __init__(self, models: list, weights: list[float] | None = None):
        self.models = models
        self.weights = weights or [1.0 / len(models)] * len(models)

    def predict(self, X) -> np.ndarray:
        preds = np.zeros(len(X))
        for model, w in zip(self.models, self.weights):
            if hasattr(model, "predict"):
                preds += w * np.asarray(model.predict(X))
            else:
                import torch
                model.eval()
                with torch.no_grad():
                    xt = torch.tensor(X, dtype=torch.float32)
                    preds += w * model(xt).numpy()
        return preds