import random
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, TYPE_CHECKING

import torch
from numpy import inf
import numpy as np
from torch import Tensor, cat
from torch.nn import Module, CrossEntropyLoss
from torch.utils.data import random_split, DataLoader
from torch.optim import Optimizer, SGD, AdamW

from avalanche.benchmarks.utils import AvalancheConcatDataset, \
	AvalancheDataset, AvalancheSubset
from avalanche.benchmarks.utils.data_loader import \
	ReplayDataLoader
from avalanche.models import FeatureExtractorBackbone
from avalanche.training.plugins.strategy_plugin import StrategyPlugin

if TYPE_CHECKING:
	from avalanche.training.strategies import BaseStrategy

from avalanche.models.utils import avalanche_forward
from avalanche.benchmarks.utils.data_loader import TaskBalancedDataLoader, _seq_collate_mbatches_fn
from avalanche.models.dynamic_optimizers import reset_optimizer
import copy
import numpy as np
from avalanche.models.utils import LabelSmoothingCrossEntropy



def _get_class_error_based_indicies(strategy, rehearsal_plugin):

	nr_indices_to_return = int(len(strategy.experience.dataset)*rehearsal_plugin.buffer_data_ratio)
	if len(rehearsal_plugin.memory_dataset)<= nr_indices_to_return:
		return None

	assert strategy.val_cls_acc_dict is not None

	idx_rehersal_prob = torch.zeros((len(rehearsal_plugin.memory_dataset)))
	for cls_i, cls_acc_i in strategy.val_cls_acc_dict.items():
		if isinstance(cls_acc_i.result(), dict) and 0 in cls_acc_i.result().keys() : 
			cls_acc = cls_acc_i.result()[0]
			cls_rehersal_prob= 1-cls_acc
		else: 
			cls_acc = 0

		cls_rehersal_prob= 1-cls_acc
		total_seen_for_cls = len(idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]])

		if total_seen_for_cls!=0:
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = cls_rehersal_prob/total_seen_for_cls
		else:
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = 0

	idxs = idx_rehersal_prob.multinomial(num_samples=nr_indices_to_return, replacement=False)
	return idxs


def _get_class_error_temp_based_indicies(strategy, rehearsal_plugin):

	nr_indices_to_return = int(len(strategy.experience.dataset)*rehearsal_plugin.buffer_data_ratio)
	if len(rehearsal_plugin.memory_dataset)<= nr_indices_to_return:
		return None

	assert strategy.val_cls_acc_dict is not None

	idx_rehersal_prob = torch.zeros((len(rehearsal_plugin.memory_dataset)))

	for cls_i, cls_acc_i in strategy.val_cls_acc_dict.items():
		if isinstance(cls_acc_i.result(), dict) and 0 in cls_acc_i.result().keys() : 
			cls_acc = cls_acc_i.result()[0]
			cls_rehersal_prob= 1-cls_acc
		else:  
			cls_acc = 0
		cls_rehersal_prob= torch.exp(torch.tensor((1-cls_acc)/rehearsal_plugin.temperature))
		total_seen_for_cls = len(rehearsal_plugin.cls_idx_dict[cls_i])
		if total_seen_for_cls!=0:
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = cls_rehersal_prob/total_seen_for_cls
		else:
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = 0

	idxs = idx_rehersal_prob.multinomial(num_samples=nr_indices_to_return, replacement=False)
	return idxs


def _get_class_random_indicies(strategy, rehearsal_plugin):
	nr_indices_to_return = int(len(strategy.experience.dataset)*rehearsal_plugin.buffer_data_ratio)
	if len(rehearsal_plugin.memory_dataset)<= nr_indices_to_return:
		return None
	else:
		idxs = random.sample(range(len(rehearsal_plugin.memory_dataset)), nr_indices_to_return)
		idxs = torch.tensor(idxs)
		return idxs

def _get_frequency_based_indicies(strategy, rehearsal_plugin):

	nr_indices_to_return = int(len(strategy.experience.dataset)*rehearsal_plugin.buffer_data_ratio)
	if len(rehearsal_plugin.memory_dataset)<= nr_indices_to_return:
		return None

	idx_rehersal_prob = torch.zeros((len(rehearsal_plugin.memory_dataset)))
	total_data_len = len(rehearsal_plugin.memory_dataset)
	summed_inv_cls_freq = 0
	for cls_i, c_idxs  in rehearsal_plugin.cls_idx_dict.items():
		if c_idxs.shape[0] != 0:
			inv_cls_freq = 1/c_idxs.shape[0]
		else:
                    inv_cls_freq = 0
		summed_inv_cls_freq += inv_cls_freq
		total_seen_for_cls = len(idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]])
		if total_seen_for_cls != 0 :
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = inv_cls_freq/total_seen_for_cls
		else:
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = 0

	idx_rehersal_prob = idx_rehersal_prob/summed_inv_cls_freq
	idxs = idx_rehersal_prob.multinomial(num_samples=nr_indices_to_return, replacement=False)
	return idxs


def _get_frequency_temp_based_indicies(strategy, rehearsal_plugin):

	nr_indices_to_return = int(len(strategy.experience.dataset)*rehearsal_plugin.buffer_data_ratio)
	if len(rehearsal_plugin.memory_dataset)<= nr_indices_to_return:
		return None

	idx_rehersal_prob = torch.zeros((len(rehearsal_plugin.memory_dataset)))
	total_data_len = len(rehearsal_plugin.memory_dataset)
	summed_inv_cls_freq = 0
	for cls_i, c_idxs  in rehearsal_plugin.cls_idx_dict.items():
		if c_idxs.shape[0] != 0:
			inv_cls_freq = 1/c_idxs.shape[0]
			inv_cls_freq = torch.exp(torch.tensor(inv_cls_freq/rehearsal_plugin.temperature))
		else:
			inv_cls_freq = 0
		summed_inv_cls_freq += inv_cls_freq
		total_seen_for_cls = len(idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]])
		if total_seen_for_cls != 0 :
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = inv_cls_freq/total_seen_for_cls
		else:
			idx_rehersal_prob[rehearsal_plugin.cls_idx_dict[cls_i]] = 0

	idx_rehersal_prob = idx_rehersal_prob/summed_inv_cls_freq
	idxs = idx_rehersal_prob.multinomial(num_samples=nr_indices_to_return, replacement=False)
	return idxs



global rehearsal_selection_strategy_dict
rehearsal_selection_strategy_dict = {'ce':_get_class_error_based_indicies,
									 'cet': _get_class_error_temp_based_indicies,
									 'cft': _get_frequency_temp_based_indicies,
									 'cf': _get_frequency_based_indicies,
									 'rr': _get_class_random_indicies }

class ClassImbalanceMemoryRehersalPlugin(StrategyPlugin):
	# This is the parent class of plugins that handle a limited memory and the filling of that memory. 
	# The memory_dataset variable stores the data that is saved. 
	# The variable memory_size is an integer variable defining how many instances can be stored
	# in the memory. 
	# The rehearsal_selection_strategy has to be a string code defining how the data from the memory should 
	# selceted for the rehearsal set. The global rehearsal_selection_strategy_dict selects the correct function
	# for returning the the indicies for the memry picked for rehearsal. 

	# 

	def __init__(self, buffer_data_ratio, memory_size, rehearsal_selection_strategy, temperature):
		super().__init__()
	
		self.memory_dataset = None
		self.cls_idx_dict = {}
		self.buffer_data_ratio = buffer_data_ratio
		self.memory_size = memory_size
		if temperature is not None and rehearsal_selection_strategy=='ce':
			rehearsal_selection_strategy = 'cet'
		if temperature is not None and rehearsal_selection_strategy=='cf':
			rehearsal_selection_strategy = 'cft'
		self.temperature = temperature
		self.rehearsal_selection_strategy = rehearsal_selection_strategy_dict[rehearsal_selection_strategy]


	def train_dataset_adaptation(self, strategy: 'BaseStrategy', **kwargs):
		
		indices = None
		if self.memory_dataset is not None:

			indices = self._get_indices(strategy, **kwargs)
			if indices is None:
				self.exp_rehersal_set = self.memory_dataset # until r experiences have been seen _get_indicies functions return None
			else:
				self.exp_rehersal_set = AvalancheSubset(self.memory_dataset, indices=indices)

			strategy.adapted_dataset = AvalancheConcatDataset([strategy.experience.dataset, self.exp_rehersal_set], paths=[strategy.experience.dataset.paths, self.exp_rehersal_set.paths])
	
		else:
			strategy.adapted_dataset = AvalancheConcatDataset([strategy.experience.dataset], paths=strategy.experience.dataset.paths) # in first experience


		strategy.rehearsal_indicies_picked[strategy.experience.current_experience] = indices


class ReservoirMemoryRehearsalPlugin(ClassImbalanceMemoryRehersalPlugin):

	# This plugin is an adaptation of the reservoir sampling algorithm from a stream of data for batches. The original algorithm 
	# is designed to sample uniformly from a stream, meaning at the end of the probability of each stream element to be in the memory 
	# is exactly 1/nr_elements_seen_so_far.  
	# This is a batch-wise adaptation of this algorithm to avoid looping over every element in the experience. 
	# This memory filling strategy can be seen as a baseline and compared to the class balancing reservoir sampling method described below.

	def __init__(self, buffer_data_ratio, memory_size, nr_classes, rehearsal_selection_strategy, temperature):
		super().__init__(buffer_data_ratio, memory_size, rehearsal_selection_strategy, temperature )
		
		self.nr_classes = nr_classes
		self.total_stream_length = 0


	def after_training_exp(self, strategy: 'BaseStrategy', **kwargs):
		if self.memory_dataset is None:
			self.memory_dataset = strategy.experience.dataset
			self.total_stream_length = len(strategy.experience.dataset)


		elif len(self.memory_dataset)+len(strategy.experience.dataset) <= self.memory_size:
			combined_paths = self.memory_dataset.paths+strategy.experience.dataset.paths

			self.memory_dataset = AvalancheConcatDataset(
				[self.memory_dataset, strategy.experience.dataset], paths=combined_paths)
			self.total_stream_length += len(strategy.experience.dataset)

		else:

			if len(self.memory_dataset)==self.memory_size:
				experience_data = strategy.experience.dataset

			elif len(self.memory_dataset)<self.memory_size:

				experience_fill_memory = AvalancheSubset(strategy.experience.dataset, indices=range(self.memory_size-len(self.memory_dataset)), paths=strategy.experience.dataset.paths)
				experience_data = AvalancheSubset(strategy.experience.dataset, indices=range(self.memory_size-len(self.memory_dataset), len(strategy.experience.dataset)), paths=strategy.experience.dataset.paths)
				memory_dataset_paths = self.memory_dataset.paths+experience_fill_memory.paths
				self.memory_dataset = AvalancheConcatDataset([self.memory_dataset, experience_fill_memory ], paths=memory_dataset_paths)

				self.total_stream_length += len(experience_fill_memory)#


			stream_positions = range(self.total_stream_length, self.total_stream_length+len(experience_data))

			self.total_stream_length = self.total_stream_length+len(experience_data)
			random_M = [random.randint(0, i) for i in stream_positions]
			random_M = torch.tensor(random_M)

			exp_idx_less_than_mem_size = (random_M<self.memory_size).nonzero().flatten()

			mem_idx_to_replace = np.random.choice(range(self.memory_size), size=exp_idx_less_than_mem_size.shape[0], replace=True)

			unique_mem_to_replace, index_exp_to_replace = np.unique(mem_idx_to_replace[::-1], return_index=True)#

			index_exp_to_replace = -index_exp_to_replace +(len(exp_idx_less_than_mem_size)-1) 


			mem_idx_to_retain = np.array(range(self.memory_size))[torch.logical_not(torch.isin(torch.arange(self.memory_size), torch.tensor(unique_mem_to_replace)))]
			exp_subset = AvalancheSubset(experience_data, indices=index_exp_to_replace, paths=experience_data.paths)
			mem_subset = AvalancheSubset(self.memory_dataset, indices=mem_idx_to_retain, paths=self.memory_dataset.paths)

			new_memory_paths = exp_subset.paths+ mem_subset.paths
			self.memory_dataset = AvalancheConcatDataset([mem_subset, exp_subset],  paths=new_memory_paths)

		all_memory_targets = torch.tensor(self.memory_dataset.targets)
		strategy.memory_dataset_paths.append(self.memory_dataset.paths)
		strategy.memory_dataset_targets.append(all_memory_targets)
		for i in range(strategy.model.num_classes):
			self.cls_idx_dict[i] = (all_memory_targets==i).nonzero().flatten()


	def _get_indices(self, strategy):
		return self.rehearsal_selection_strategy(strategy, self)


class ClassBalancingReservoirMemoryRehersalPlugin(ClassImbalanceMemoryRehersalPlugin):

	# This memory filling strategy was based on the Class-Balancing Reservoir Sampling Algorithm defined in the paper 
	# 'Online Continual Learning from Imbalanced Data' by Chrysakis et al. The basic idea is to limit the replacement of old instances 
	# in the memory to classes that are 'full'. Please see the paper for more tetails. This again is a batchwise adaptation of the 
	# original algorithm.

	def __init__(self, buffer_data_ratio, memory_size, nr_classes, rehearsal_selection_strategy, temperature):
		super().__init__(buffer_data_ratio, memory_size, rehearsal_selection_strategy, temperature )

		self.full_classes = torch.tensor([])
		self.instances_in_stream_count_dict = {}
		self.instances_in_memory_count_dict = {}
		self.mem_to_stream_ratios_dict = {}
		self.nr_classes = nr_classes

		for i in range(self.nr_classes):
			self.instances_in_memory_count_dict[i] = 0
			self.instances_in_stream_count_dict[i] = 0 
			self.mem_to_stream_ratios_dict[i] = np.Inf 



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

		for i in range(strategy.model.num_classes):

			nr_cls_instances = (torch.tensor(strategy.experience.dataset.targets)==i).nonzero().flatten().shape[0]

			if nr_cls_instances !=0:

				self.instances_in_stream_count_dict[i] += nr_cls_instances
				self.mem_to_stream_ratios_dict[i]=self.instances_in_memory_count_dict[i]/self.instances_in_stream_count_dict[i]

			
		if self.memory_dataset is None:
			self.memory_dataset = strategy.experience.dataset
		elif len(self.memory_dataset)+len(strategy.experience.dataset) <= self.memory_size:
			combined_paths = self.memory_dataset.paths+strategy.experience.dataset.paths

			self.memory_dataset = AvalancheConcatDataset(
				[self.memory_dataset, strategy.experience.dataset], paths=combined_paths)


		else:
			if len(self.memory_dataset)==self.memory_size:

				experience_data = strategy.experience.dataset
			elif len(self.memory_dataset)<self.memory_size:

				experience_fill_memory = AvalancheSubset(strategy.experience.dataset, indices=range(self.memory_size-len(self.memory_dataset)), paths=strategy.experience.dataset.paths)
				experience_data = AvalancheSubset(strategy.experience.dataset, indices=range(self.memory_size-len(self.memory_dataset), len(strategy.experience.dataset)), paths=strategy.experience.dataset.paths)

				memory_dataset_paths = self.memory_dataset.paths+experience_fill_memory.paths
				self.memory_dataset = AvalancheConcatDataset([self.memory_dataset, experience_fill_memory ], paths=memory_dataset_paths)


			all_memory_targets = torch.tensor(self.memory_dataset.targets)
			u, c = torch.unique(all_memory_targets, return_counts=True)
			most_freq_cls_in_mem = u[c==torch.max(c)]
			self.full_classes = torch.unique(torch.cat([self.full_classes, most_freq_cls_in_mem]))

			
			
			mem_subset_in_full_classes_idxs = (torch.isin(all_memory_targets, self.full_classes)).nonzero().flatten()
			mem_subset_NOT_in_full_classes_idxs = (torch.logical_not(torch.isin(all_memory_targets, self.full_classes))).nonzero().flatten()


			exp_subset_in_NOT_full_classes_idxs = (torch.logical_not(torch.isin(torch.tensor(experience_data.targets), self.full_classes))).nonzero().flatten()

			idxs_from_large_classes_retained = random.sample(mem_subset_in_full_classes_idxs.tolist(), mem_subset_in_full_classes_idxs.shape[0]-exp_subset_in_NOT_full_classes_idxs.shape[0] )
			large_classes_retained_dataset = AvalancheSubset(self.memory_dataset, indices=idxs_from_large_classes_retained, paths=self.memory_dataset.paths)
			NOT_large_classes_memory_subset = AvalancheSubset(self.memory_dataset, indices=mem_subset_NOT_in_full_classes_idxs, paths=self.memory_dataset.paths)

			new_memory = AvalancheSubset(experience_data, indices= exp_subset_in_NOT_full_classes_idxs, paths=experience_data.paths)
			new_memory_paths =new_memory.paths+NOT_large_classes_memory_subset.paths
			new_memory = AvalancheConcatDataset([new_memory, NOT_large_classes_memory_subset], paths=new_memory_paths)

			for j, cls_i in enumerate(self.full_classes):

				cls_i_in_mem_retained_idxs = (torch.tensor(large_classes_retained_dataset.targets)==cls_i).nonzero().flatten()#

				cls_i_exp_subset_idxs =  (torch.tensor(experience_data.targets)==cls_i).nonzero().flatten()


				uniform_samples = torch.rand(cls_i_exp_subset_idxs.shape[0])

				exp_idxs_for_replacing_mem = (uniform_samples<=self.mem_to_stream_ratios_dict[int(cls_i.item())]).nonzero().flatten()
				mem_idx_to_replace = np.random.choice(cls_i_in_mem_retained_idxs, size=exp_idxs_for_replacing_mem.shape[0], replace=True)
				unique_mem_to_replace, index_exp_to_replace = np.unique(mem_idx_to_replace[::-1], return_index=True)

				index_exp_to_replace = -index_exp_to_replace +(len(exp_idxs_for_replacing_mem)-1) 
				index_exp_to_replace = cls_i_exp_subset_idxs[index_exp_to_replace]

				mem_idx_to_retain = cls_i_in_mem_retained_idxs[torch.logical_not(torch.isin(cls_i_in_mem_retained_idxs, torch.tensor(unique_mem_to_replace)))]

				cls_i_exp_subset = AvalancheSubset(experience_data, indices=index_exp_to_replace, paths=experience_data.paths)
				cls_i_mem_subset = AvalancheSubset(large_classes_retained_dataset, indices=mem_idx_to_retain, paths=large_classes_retained_dataset.paths)

				cls_i_paths = cls_i_exp_subset.paths+cls_i_mem_subset.paths
				cls_i_data = AvalancheConcatDataset([cls_i_exp_subset, cls_i_mem_subset], paths=cls_i_paths )

				new_memory_paths = new_memory.paths+ cls_i_data.paths

				new_memory = AvalancheConcatDataset([new_memory, cls_i_data],  paths=new_memory_paths)			
			self.memory_dataset = new_memory
		all_memory_targets = torch.tensor(self.memory_dataset.targets)
		strategy.memory_dataset_paths.append(self.memory_dataset.paths)
		strategy.memory_dataset_targets.append(all_memory_targets)
		print('full classes: ', self.full_classes)
		for i in range(strategy.model.num_classes):
			self.cls_idx_dict[i] = (all_memory_targets==i).nonzero().flatten()
			self.instances_in_memory_count_dict[i] = self.cls_idx_dict[i].shape[0]



	def _get_indices(self, strategy):
		return self.rehearsal_selection_strategy(strategy, self)


