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
from torch.utils.data import random_split, DataLoader

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.benchmarks.utils.data_loader import TaskBalancedDataLoader, _seq_collate_mbatches_fn
import matplotlib.pyplot as plt

import pickle as pkl


class SeqDataPlugin(StrategyPlugin):
    """Plugin that handles SeqPathDataset where sequence information is relevant. 
    the class ant seq wise accuracies are logged using the funcitons implemented in this function. 
    """

    def __init__(self):
        super().__init__()
        

    def before_train_dataset_adaptation(self, strategy: "BaseStrategy"):
        # counting unique class occurances in train stream to produce help evaluate the results 
        unique, counts = torch.unique(torch.as_tensor(strategy.experience.dataset.targets), return_counts=True)

        for i, cls_i in enumerate(unique):
            strategy.unique_train_cls_dict[cls_i.item()] += counts[i]

    def before_eval_dataset_adaptation(self, strategy: "BaseStrategy"):
        # counting class occurances in validation and test data to help evaluate the results 
        if strategy.last_eval:

            if strategy.experience.origin_stream.name =='test':
                unique, counts = torch.unique(torch.as_tensor(strategy.experience.dataset.targets), return_counts=True)
                if strategy.current_train_exp_seen == 0: 
                    for i, cls_i in enumerate(unique):
                        strategy.total_test_cls_dict[cls_i.item()] +=counts[i]

            if strategy.experience.origin_stream.name =='validation':
                unique, counts = torch.unique(torch.as_tensor(strategy.experience.dataset.targets), return_counts=True)
                if strategy.current_train_exp_seen == 0: 
                    for i, cls_i in enumerate(unique):
                        strategy.total_validation_cls_dict[cls_i.item()] +=counts[i]

    def after_train_dataset_adaptation(self, strategy: "BaseStrategy"):
        # counting class occurances in train data after the rehearsal set has been selected to help evaluate the results. 
        unique, counts =torch.unique(torch.as_tensor(strategy.adapted_dataset.targets), return_counts=True)

        for i, cls_i in enumerate(unique):
            strategy.total_train_cls_dict[cls_i.item()] +=counts[i]


    def make_train_dataloader(self, strategy: "BaseStrategy", num_workers=0, shuffle=True,
                              pin_memory=False, **kwargs):
        kwargs['oversample_small_groups'] =True
        strategy.dataloader = TaskBalancedDataLoader(
                strategy.adapted_dataset,
                collate_mbatches=_seq_collate_mbatches_fn,
                num_workers=num_workers,
                batch_size=strategy.train_mb_size,
                shuffle=shuffle,
                pin_memory=pin_memory, **kwargs)

    def before_eval(self, strategy: "BaseStrategy"):

        if strategy.last_eval:
            strategy.all_preds = torch.tensor([], dtype=int, device=strategy.device)
            strategy.all_targets = torch.tensor([], dtype=int, device=strategy.device)



    def make_eval_dataloader(self, strategy: "BaseStrategy", num_workers=0, shuffle=True,
                              pin_memory=False, **kwargs):
        kwargs['oversample_small_groups'] =True
        strategy.dataloader = TaskBalancedDataLoader(
                strategy.adapted_dataset,
                collate_mbatches=_seq_collate_mbatches_fn,
                num_workers=num_workers,
                batch_size=strategy.train_mb_size,
                shuffle=shuffle,
                pin_memory=pin_memory, **kwargs)

    def _unpack_minibatch(self, strategy: "BaseStrategy",  **kwargs):

        # Callback function to handle the third seq_id field that the Dataloader of the SeqPathsDataset returns

        strategy.paths = strategy.mbatch[1][0]
        strategy.seq_codes = torch.tensor(strategy.mbatch[2][0]).to(strategy.device)
        strategy.mbatch = strategy.mbatch[0]

        assert len(strategy.mbatch) >= 3
        for i in range(len(strategy.mbatch)):

            strategy.mbatch[i] = strategy.mbatch[i].to(strategy.device)



    def after_forward(self, strategy:"BaseStrategy", **kwargs):
        _, strategy.preds = torch.max(strategy.mb_output, dim=1)


    def after_eval_forward(self, strategy:"BaseStrategy", **kwargs):
        # keeps track of predictions and targets during evaluation phase to produce comprehensive confusion matricies after evaluation.
        _, strategy.preds = torch.max(strategy.mb_output, dim=1)

        if strategy.last_eval:
            strategy.all_targets = torch.cat([strategy.all_targets, strategy.mb_y])
            strategy.all_preds = torch.cat([strategy.all_preds, strategy.preds])

    


