import copy
import itertools
from typing import TYPE_CHECKING, Optional, List
import torch
from torch.optim import Optimizer

from avalanche.benchmarks.utils import AvalancheConcatDataset, \
    AvalancheTensorDataset, AvalancheSubset
from math import ceil

from avalanche.models import TrainEvalModel, NCMClassifier
from avalanche.training.plugins import EvaluationPlugin
from avalanche.training.plugins.evaluation import default_logger
from avalanche.training.losses import ICaRLLossPlugin
from avalanche.training.plugins.strategy_plugin import StrategyPlugin
from torch.nn import Module
from torch.utils.data import DataLoader
from avalanche.training.strategies import BaseStrategy


class ICaRL(BaseStrategy):
    def __init__(self, feature_extractor: Module, classifier: Module,
                 optimizer: Optimizer, memory_size, buffer_transform,
                 fixed_memory, criterion=ICaRLLossPlugin(),
                 train_mb_size: int = 1, train_epochs: int = 1,
                 eval_mb_size: int = None, device=None,
                 plugins: Optional[List[StrategyPlugin]] = None,
                 evaluator: EvaluationPlugin = default_logger, eval_every=-1):
        """ iCaRL Strategy.
        This strategy does not use task identities.

        :param feature_extractor: The feature extractor.
        :param classifier: The differentiable classifier that takes as input
            the output of the feature extractor.
        :param optimizer: The optimizer to use.
        :param memory_size: The nuber of patterns saved in the memory.
        :param buffer_transform: transform applied on buffer elements already
            modified by test_transform (if specified) before being used for
             replay
        :param fixed_memory: If True a memory of size memory_size is
            allocated and partitioned between samples from the observed
            experiences. If False every time a new class is observed
            memory_size samples of that class are added to the memory.
        :param train_mb_size: The train minibatch size. Defaults to 1.
        :param train_epochs: The number of training epochs. Defaults to 1.
        :param eval_mb_size: The eval minibatch size. Defaults to 1.
        :param device: The device to use. Defaults to None (cpu).
        :param plugins: Plugins to be added. Defaults to None.
        :param evaluator: (optional) instance of EvaluationPlugin for logging
            and metric computations.
        :param eval_every: the frequency of the calls to `eval` inside the
            training loop.
                if -1: no evaluation during training.
                if  0: calls `eval` after the final epoch of each training
                    experience.
                if >0: calls `eval` every `eval_every` epochs and at the end
                    of all the epochs for a single experience.
        """
        model = TrainEvalModel(feature_extractor,
                               train_classifier=classifier,
                               eval_classifier=NCMClassifier())

        icarl = _ICaRLPlugin(memory_size, buffer_transform, fixed_memory)

        if plugins is None:
            plugins = [icarl]
        else:
            plugins += [icarl]

        if isinstance(criterion, StrategyPlugin):
            plugins += [criterion]

        super().__init__(
            model, optimizer, criterion=criterion,
            train_mb_size=train_mb_size, train_epochs=train_epochs,
            eval_mb_size=eval_mb_size, device=device, plugins=plugins,
            evaluator=evaluator, eval_every=eval_every)


class _ICaRLPlugin(StrategyPlugin):
    """
        iCaRL Plugin.
        iCaRL uses nearest class exemplar classification to prevent
        forgetting to occur at the classification layer. The feature extractor
        is continually learned using replay and distillation. The exemplars
        used for replay and classification are selected through herding.
        This plugin does not use task identities.
        """

    def __init__(self, memory_size, buffer_transform=None, fixed_memory=True):
        """
        :param memory_size: amount of patterns saved in the memory.
        :param buffer_transform: transform applied on buffer elements already
            modified by test_transform (if specified) before being used for
             replay
        :param fixed_memory: If True a memory of size memory_size is
            allocated and partitioned between samples from the observed
            experiences. If False every time a new class is observed
            memory_size samples of that class are added to the memory.
        """
        super().__init__()

        self.memory_size = memory_size
        self.buffer_transform = buffer_transform
        self.fixed_memory = fixed_memory

        self.x_memory = []
        self.y_memory = []
        self.order = []

        self.old_model = None
        self.observed_classes = []
        self.class_means = None
        self.embedding_size = None
        self.output_size = None
        self.input_size = None

    def after_train_dataset_adaptation(self, strategy: 'BaseStrategy',
                                       **kwargs):
        if strategy.training_exp_counter != 0:
            memory = AvalancheTensorDataset(
                torch.cat(self.x_memory).cpu(),
                list(itertools.chain.from_iterable(self.y_memory)),
                transform=self.buffer_transform, target_transform=None)

            strategy.adapted_dataset = \
                AvalancheConcatDataset((strategy.adapted_dataset, memory))

    def before_training_exp(self, strategy: 'BaseStrategy', **kwargs):
        tid = strategy.training_exp_counter
        benchmark = strategy.experience.benchmark
        nb_cl = benchmark.n_classes_per_exp[tid]

        self.observed_classes.extend(
            benchmark.classes_order[tid * nb_cl:(tid + 1) * nb_cl])

    def before_forward(self, strategy: 'BaseStrategy', **kwargs):
        if self.input_size is None:
            with torch.no_grad():
                self.input_size = strategy.mb_x.shape[1:]
                self.output_size = strategy.model(strategy.mb_x).shape[1]
                self.embedding_size = strategy.model.feature_extractor(
                    strategy.mb_x).shape[1]

    def after_training_exp(self, strategy: 'BaseStrategy', **kwargs):
        strategy.model.eval()

        self.construct_exemplar_set(strategy)
        self.reduce_exemplar_set(strategy)
        self.compute_class_means(strategy)

    def compute_class_means(self, strategy):
        if self.class_means is None:
            n_classes = sum(strategy.experience.benchmark.n_classes_per_exp)
            self.class_means = torch.zeros(
                (self.embedding_size, n_classes)).to(strategy.device)

        for i, class_samples in enumerate(self.x_memory):
            label = self.y_memory[i][0]
            class_samples = class_samples.to(strategy.device)

            with torch.no_grad():
                mapped_prototypes = strategy.model.feature_extractor(
                    class_samples).detach()
            D = mapped_prototypes.T
            D = D / torch.norm(D, dim=0)

            if len(class_samples.shape) == 4:
                class_samples = torch.flip(class_samples, [3])

            with torch.no_grad():
                mapped_prototypes2 = strategy.model.feature_extractor(
                    class_samples).detach()

            D2 = mapped_prototypes2.T
            D2 = D2 / torch.norm(D2, dim=0)

            div = torch.ones(class_samples.shape[0], device=strategy.device)
            div = div / class_samples.shape[0]

            m1 = torch.mm(D, div.unsqueeze(1)).squeeze(1)
            m2 = torch.mm(D2, div.unsqueeze(1)).squeeze(1)
            self.class_means[:, label] = (m1 + m2) / 2
            self.class_means[:, label] /= torch.norm(self.class_means[:, label])

            strategy.model.eval_classifier.class_means = self.class_means

    def construct_exemplar_set(self, strategy):
        tid = strategy.training_exp_counter
        nb_cl = strategy.experience.benchmark.n_classes_per_exp

        if self.fixed_memory:
            nb_protos_cl = int(ceil(
                self.memory_size / len(self.observed_classes)))
        else:
            nb_protos_cl = self.memory_size
        new_classes = self.observed_classes[tid * nb_cl[tid]:
                                            (tid + 1) * nb_cl[tid]]

        dataset = strategy.experience.dataset
        targets = torch.tensor(dataset.targets)
        for iter_dico in range(nb_cl[tid]):
            cd = AvalancheSubset(dataset,
                                 torch.where(targets == new_classes[iter_dico])
                                 [0])

            class_patterns, _, _ = next(iter(
                DataLoader(cd.eval(), batch_size=len(cd))))
            class_patterns = class_patterns.to(strategy.device)

            with torch.no_grad():
                mapped_prototypes = strategy.model.feature_extractor(
                    class_patterns).detach()
            D = mapped_prototypes.T
            D = D / torch.norm(D, dim=0)

            mu = torch.mean(D, dim=1)
            order = torch.zeros(class_patterns.shape[0])
            w_t = mu

            i, added, selected = 0, 0, []
            while not added == nb_protos_cl and i < 1000:
                tmp_t = torch.mm(w_t.unsqueeze(0), D)
                ind_max = torch.argmax(tmp_t)

                if ind_max not in selected:
                    order[ind_max] = 1 + added
                    added += 1
                    selected.append(ind_max.item())

                w_t = w_t + mu - D[:, ind_max]
                i += 1

            pick = (order > 0) * (order < nb_protos_cl + 1) * 1.
            self.x_memory.append(class_patterns[torch.where(pick == 1)[0]])
            self.y_memory.append(
                [new_classes[iter_dico]] * len(torch.where(pick == 1)[0]))
            self.order.append(order[torch.where(pick == 1)[0]])

    def reduce_exemplar_set(self, strategy):
        tid = strategy.training_exp_counter
        nb_cl = strategy.experience.benchmark.n_classes_per_exp

        if self.fixed_memory:
            nb_protos_cl = int(ceil(
                self.memory_size / len(self.observed_classes)))
        else:
            nb_protos_cl = self.memory_size

        for i in range(len(self.x_memory) - nb_cl[tid]):
            pick = (self.order[i] < nb_protos_cl + 1) * 1.
            self.x_memory[i] = self.x_memory[i][torch.where(pick == 1)[0]]
            self.y_memory[i] = self.y_memory[i][:len(torch.where(pick == 1)[0])]
            self.order[i] = self.order[i][torch.where(pick == 1)[0]]
