from pathlib import Path

import einops
import numpy as np
import pandas as pd
import torch
import yaml

import priceforecast.models.nn.MAE_utils.models as models
from priceforecast.models.base import ForecastApiRequest, ForecastApiResponse
from priceforecast.models.nn.base_nn import BaseNN, WEIGHTS_KEY


class MAE(BaseNN):
    """Class that implements a very simple neural network model."""

    def __init__(self, config_path: Path):
        super().__init__(config_path)

        with open(self.config['model_info'], 'r') as f:
            self.model_info = yaml.safe_load(f)

        self.model = self.load_model()

        self.past_data = np.zeros(self.model_info['past_window_hrs'])
        self.last_time_step = 0

        self.forecast_window = self.config['forecast_window']

        self.check_model_compatibility()

    def check_model_compatibility(self):

        if self.forecast_window % self.model_info['model_config']['dim_tokens'] != 0:
            raise Exception('The forecast window must be divisible by the token dimension.')
        elif 'PerceiverIO' in self.model_info['model_name'] and self.forecast_window > self.model_info[
            'forecast_window_hrs']:
            raise Exception(f"The forecast window must not be larger than {self.model_info['forecast_window_hrs']}.")

    def load_model(self):
        # load model
        if not self.config.get(WEIGHTS_KEY):
            raise ValueError("MAE model requires weights path")
        state = torch.load(self.config[WEIGHTS_KEY], map_location=torch.device('cpu'), weights_only=True)[
            'model_state']
        # we need to remove the first word from the keys
        state = {k[7:]: v for k, v in state.items()}
        model = getattr(models, self.model_info['model_name'])(**self.model_info['model_config'])
        model.load_state_dict(state)
        model.eval()
        print('Model loaded')
        return model

    def forecast(self, request: ForecastApiRequest) -> ForecastApiResponse:
        """Return `ForecastApiResponse` on given `request`"""
        past_tokens, past_idxs, future_idxs = self.create_dataset(request=request)

        # forward pass
        if 'PerceiverIO' in self.model_info['model_name']:
            forecasts = self.model(past_tokens, past_idxs, future_idxs)
        elif 'PerceiverAR' in self.model_info['model_name']:
            forecasts = None
            for i in range(future_idxs.shape[1]):
                next_prediction = self.model(past_tokens)[..., -1:, :]
                forecasts = next_prediction if forecasts is None else torch.cat((forecasts, next_prediction), dim=-2)
                if 'prob_ensemble' in self.model_info['model_name']:
                    next_prediction = next_prediction.mean(dim=1)
                past_tokens = torch.cat((past_tokens[..., 1:, :], next_prediction), dim=-2)

        # de-tokenise
        forecasts = forecasts.detach().cpu().numpy()
        forecasts = einops.rearrange(forecasts, 'b t td -> b (t td)') if self.model_info[
                                                                             'loss'] != 'crps' else einops.rearrange(
            forecasts, 'batch ensemble tokens tokendim -> batch ensemble (tokens tokendim)')

        # de-normalise
        if self.model_info['normalized']:
            forecasts = forecasts * self.model_info['normalization_parameters']['std'] + \
                        self.model_info['normalization_parameters']['mean']

        forecasts = forecasts.squeeze(axis=0)

        if len(forecasts.shape) > 1:
            means = dict(zip(self._get_requested_time_steps(request), forecasts.mean(axis=0)))
            variances = dict(zip(self._get_requested_time_steps(request), forecasts.var(axis=0)))
        else:
            means = dict(zip(self._get_requested_time_steps(request), forecasts))
            variances = {k: 0 for k in self._get_requested_time_steps(request)}

        return self.cast_response(request, means=[means], variances=[variances])

    def create_dataset(self, request: ForecastApiRequest) -> (torch.Tensor, torch.Tensor, torch.Tensor):
        """
        """
        past_targets = self._get_sorted_by_key(request.pastTargets)
        past_targets_keys = np.array(list(past_targets.keys()))
        mask = past_targets_keys > self.last_time_step
        past_targets_values = np.array(list(past_targets.values()))[mask]
        len_input = len(past_targets_values)
        if len_input > 0:
            self.last_time_step = past_targets_keys[-1]
            self.past_data = np.roll(self.past_data, - len_input)
            self.past_data[-len_input:] = past_targets_values

        # normalize
        if self.model_info['normalized']:
            data = (self.past_data - self.model_info['normalization_parameters']['mean']) / \
                   self.model_info['normalization_parameters']['std']

        # transform to tensors and add batch dimension
        data = torch.tensor(data.astype(np.float32)).unsqueeze(dim=0)

        # tokenize
        past_tokens = einops.rearrange(data, 'b (t td) -> b t td',
                                       td=self.model_info['model_config']['dim_tokens'])

        past_idxs = torch.arange(past_tokens.shape[1]).expand(past_tokens.shape[0], -1)

        future_idxs = torch.arange(past_tokens.shape[1], past_tokens.shape[1] + (
                self.forecast_window // self.model_info['model_config']['dim_tokens'])).expand(
            past_tokens.shape[0], -1)

        return past_tokens, past_idxs, future_idxs

    def train_nn(self, data: dict[str, pd.DataFrame]) -> None:
        """
        Trains a NN to predict targets based on features.

        Args:
            data: dictionary of scenarios to be used for training
        """
        raise ValueError("MAE Model is used as zero shot model and cannot be trained")
