from typing import List, Optional

import numpy as np
import pandas as pd

import torch
from torch import nn
from torch.optim import SGD

from lightwood.mixer.base import BaseMixer
from lightwood.ensemble.mean_ensemble import MeanEnsemble
from lightwood.api.types import PredictionArguments
from lightwood.data.encoded_ds import EncodedDs
from lightwood.helpers.log import log


class StackedEnsemble(MeanEnsemble):
    """
    This ensemble will learn an optimal weight vector via Stochastic Gradient Descent on the validation dataset and the respective mixer predictions.

    Starting weights for the vector are uniformly set.

    Note this mixer is still in experimental phase. Some features in the roadmap are:
      - support for handling faulty mixers
      - support for custom initial vector weights
      - early stopping
      - arbitrarily complex secondary model

    """  # noqa
    def __init__(self, target, mixers: List[BaseMixer], data: EncodedDs, dtype_dict: dict,
                 args: PredictionArguments, fit: Optional[bool] = True, **kwargs) -> None:
        super().__init__(target, mixers, data, dtype_dict, fit=False)

        self.target_cols = [target]
        self.softmax = nn.Softmax(dim=0)
        self.mixer_weights = nn.Parameter(torch.full((len(mixers),), 1 / len(mixers)))
        self.criterion = nn.MSELoss()
        self.optimizer = SGD([self.mixer_weights], lr=0.01)
        self.agg_dim = 1

        if fit and len(mixers) > 1:
            all_preds = torch.tensor(self.predict(data, args)).squeeze().reshape(-1, len(mixers))
            actual = torch.tensor(data.data_frame[self.target_cols].values)

            def _eval_loss():
                self.optimizer.zero_grad()
                weighted = torch.sum(all_preds * self.softmax(self.mixer_weights), dim=self.agg_dim)
                loss = self.criterion(weighted, actual)
                loss.backward()
                return loss

            self.optimizer.step(_eval_loss)
            self.mixer_weights = self.softmax(self.mixer_weights)
            log.info(f'Optimal stacking weights: {self.mixer_weights.detach().tolist()}')
        self.prepared = True

    def predict(self, ds: EncodedDs, args: PredictionArguments) -> List:
        outputs = []
        for mixer in self.mixers:
            output = mixer(ds, args=args)['prediction'].tolist()
            output = np.expand_dims(np.array(output), self.agg_dim)
            outputs.append(output)
        return outputs

    def __call__(self, ds: EncodedDs, args: PredictionArguments) -> pd.DataFrame:
        assert self.prepared
        mixer_weights = torch.tensor(args.mixer_weights) if args.mixer_weights else self.mixer_weights
        output = pd.DataFrame()
        predictions = torch.tensor(self.predict(ds, args)).squeeze().reshape(-1, len(self.mixers))
        predictions = (predictions * mixer_weights).sum(axis=1)
        output['prediction'] = predictions.detach().numpy().tolist()
        return output

    def set_weights(self, weights: List):
        if len(weights) != len(self.mixers):
            raise Exception(f"Expected weight vector to have {len(self.mixers)} entries, got {len(weights)} instead.")

        self.mixer_weights = torch.tensor(weights)
