################################################################################
# Copyright (c) 2021 ContinualAI.                                              #
# Copyrights licensed under the MIT License.                                   #
# See the accompanying LICENSE file for terms.                                 #
#                                                                              #
# Date: 30-12-2020                                                             #
# Author(s): Lorenzo Pellegrini                                                #
# E-mail: contact@continualai.org                                              #
# Website: www.continualai.org                                                 #
################################################################################

from typing import List, Union, Dict

import torch
from torch import Tensor
from avalanche.evaluation import Metric, PluginMetric, GenericPluginMetric
from avalanche.evaluation.metrics.mean import Mean
from avalanche.evaluation.metric_utils import phase_and_task
from collections import defaultdict
import numpy as np

def break_tie_mode(mb_output, nr_modes):

    # this funciton returns the top nr_modes prediction for a sequence
    # the mini batch outputs for each image in the sequence are averaged and 
    # the averaged output is that used in a standard way to predict the label.
    # this is an alternative approach to using the maximum softmaxed output from
    # all images as a method to aggregate the predictions

    mb_output = torch.nn.Softmax(dim =1)(mb_output)
    avg_mb_output = torch.mean(mb_output, dim=0, keepdim=True)
    _, pos = torch.topk(avg_mb_output, k=nr_modes)

    return pos


class Top1AccTransfer(GenericPluginMetric[float]): 

    # here the accuracy of each experience in the training stream is collected
    # tracking the top1 accuracy for each experience through the course of the training allows
    # for better understanding of how the performance on an experienc decreases with more, 
    # new experiences being used for training. 


    def __init__(self, nr_train_exp, reset_at, emit_at, mode):
        self.nr_train_exp = nr_train_exp
        self.train_exp_acc_dict = {}
        for i in range(self.nr_train_exp):
            self.train_exp_acc_dict[i] = Accuracy()
        
        super(Top1AccTransfer, self).__init__(
            self.train_exp_acc_dict, reset_at='stream', emit_at='stream',
            mode='eval')

    def reset(self, strategy=None) -> None:


        if strategy.experience.origin_stream.name =='train':
            for key, val in self.train_exp_acc_dict.items():
                val.reset(device=strategy.device)


    def result(self, strategy=None) -> float:
        if strategy.experience.origin_stream.name =='train':
            result_dict = {}

            for key, val in self.train_exp_acc_dict.items():
                result_dict[key] = val.result()
            
            # flag to pass to _package_results in metric_definitions.py
            # to indicate this is a dictionary.
            result_dict[-1] = [-1] 

            return result_dict
        

    def update(self, strategy):
        if strategy.experience.origin_stream.name =='train':
            task_labels = strategy.experience.task_labels
            exp_nr = strategy.experience.current_experience
            self.train_exp_acc_dict[exp_nr].update(strategy.preds, strategy.mb_y, task_labels[0], device=strategy.device)


    def __str__(self):
        return " Top1AccTransfer"+'_'+self._emit_at.capitalize()


class ClasswiseTop_nAcc(GenericPluginMetric[float]): 
    # collects the accuracies for each class in the cls_acc_dict
    # also the strategy.val_cls_acc_dict variable is updated, 
    # which is used in the class error basedrehearsal plugin


    def __init__(self, nr_classes, n, reset_at, emit_at, mode):
        self.nr_classes = nr_classes
        self.cls_acc_dict = {}
        self.n = n
        for i in range(self.nr_classes):
            self.cls_acc_dict[i] = Accuracy()
        super(ClasswiseTop_nAcc, self).__init__(
            self.cls_acc_dict, reset_at='stream', emit_at='stream',
            mode=mode)

    def reset(self, strategy=None) -> None:

        for key, val in self.cls_acc_dict.items():
            val.reset(device=strategy.device)


    def result(self, strategy=None) -> float:
        result_dict={}      
        for key, val in self.cls_acc_dict.items():
            result_dict[key] = val.result()

        #making clear distiction here such that test classification restults are ensured not to be used in rehersal set building 
        if strategy.experience.origin_stream.name =='validation' and self.n==1:
            print('found validation stream name')
            strategy.val_cls_acc_dict = self.cls_acc_dict

        if strategy.experience.origin_stream.name =='test' and self.n==1:
            print('found validation stream name')
            strategy.test_cls_acc_dict = self.cls_acc_dict

        if strategy.experience.origin_stream.name =='train' and self.n==1:
            print('found validation stream name')
            strategy.train_cls_acc_dict = self.cls_acc_dict

        result_dict[-1] = [-1]
        return result_dict

    def update(self, strategy):

        task_labels = strategy.experience.task_labels
        for key, val in self.cls_acc_dict.items():
            true_key_cls = torch.eq(strategy.mb_y,key)
            cls_true_y = strategy.mb_y[true_key_cls]
            predicted_y = strategy.preds[true_key_cls]
            cls_mb_output = strategy.mb_output[true_key_cls,:]
            sum_tensor = torch.zeros_like(cls_true_y, dtype=bool, device=strategy.device)
            _, top_n = torch.topk(cls_mb_output ,k=self.n, dim=1)
            for i in range(self.n):
                pred_overlap = top_n[:,i] == cls_true_y
                sum_tensor = torch.logical_or(sum_tensor, pred_overlap)
            if len(cls_true_y)>0:
                val.update(torch.ones_like(cls_true_y, device=strategy.device), sum_tensor*1, task_labels[0], device=strategy.device)


    def __str__(self):
        return "ClasswiseTop_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()


class ClassTop_nAvgAcc(GenericPluginMetric[float]):

    # with this metric the average of the class accuracies (view as binary decision, correct class or not)


    def __init__(self, nr_classes,n, reset_at, emit_at, mode):
        self.nr_classes = nr_classes
        self.n = n
        self.cls_acc_dict = {}
        self.emit_at = emit_at
        for i in range(self.nr_classes):
            self.cls_acc_dict[i] = Accuracy()
        self._cls_avg_accuracy = Accuracy()
        super(ClassTop_nAvgAcc, self).__init__(
            self._cls_avg_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._cls_avg_accuracy.reset(device=strategy.device)
        else:
            self._cls_avg_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        for key, val in self.cls_acc_dict.items():
            val.reset(device=strategy.device)

    def result(self, strategy=None) -> float:

        cls_avg_acc = 0
        for key, val in self.cls_acc_dict.items():
            cls_returned_results = val.result()
            if len(cls_returned_results.keys())==1:
                cls_acc = cls_returned_results[0]
                cls_avg_acc += cls_acc

        if self.emit_at == 'stream':
            result = cls_avg_acc/self.nr_classes
            
        else:
            result = float(cls_avg_acc/torch.sum(strategy.classes_used_for_training))
        if self.n==1 and  not isinstance(self, SeasonClassTop_nAcc):
            strategy.cls_top1_avg_acc = result
           
        return result

    def update(self, strategy):

        task_labels = strategy.experience.task_labels
        for key, val in self.cls_acc_dict.items():
            true_key_cls = torch.eq(strategy.mb_y,key)
            cls_true_y = strategy.mb_y[true_key_cls]
            predicted_y = strategy.preds[true_key_cls]
            cls_mb_output = strategy.mb_output[true_key_cls,:]
            sum_tensor = torch.zeros_like(cls_true_y, dtype=bool, device=strategy.device)
            _, top_n = torch.topk(cls_mb_output ,k=self.n, dim=1)
            for i in range(self.n):
                pred_overlap = top_n[:,i] == cls_true_y
                sum_tensor = torch.logical_or(sum_tensor, pred_overlap)

            if len(cls_true_y)>0:
                val.update(torch.ones_like(cls_true_y, device=strategy.device), sum_tensor*1, task_labels[0], device=strategy.device)


    def __str__(self):
        return "ClassTop_"+str(self.n)+"AvgAcc"+'_'+self._emit_at.capitalize()

class SeasonClassTop_nAcc(ClassTop_nAvgAcc):    


    # based on the files in the winter_files_list, that should countain all the file names from train, test and validation data, 
    # that are taken in the winter season, the class avg accuracy is computed wither for the 'winter' or 'summer' season.
    # this accuracy requires a variable that is set in the SeqDataPlugin
    # all season wise accuracy compute the accuracy on either all winter data or all summer data depending on the 
    # class variable 'season'. If both seasons should be tracked, there should be an instance for the summer and one 
    # for the winter data


    def __init__(self, nr_classes, n, winter_files_list, season, reset_at, emit_at, mode):
        self.winter_files_list = winter_files_list
        self.season = season

        super(SeasonClassTop_nAcc, self).__init__( 
        nr_classes=nr_classes, n=n, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def update(self, strategy):
        # the strategy.paths variable is updated in the _unpack_minibatch function of the SeqDataPlugin
        in_season = np.isin(np.array(strategy.paths), self.winter_files_list)
        
        if self.season == 'summer':
            in_season = np.logical_not(in_season)

        self.mb_y = strategy.mb_y[in_season]
        self.preds = strategy.preds[in_season]
        self.mb_output = strategy.mb_output[in_season,:]
        self.seq_codes = strategy.seq_codes[in_season]

        task_labels = strategy.experience.task_labels
        for key, val in self.cls_acc_dict.items():
            true_key_cls = torch.eq(self.mb_y,key)
            cls_true_y = self.mb_y[true_key_cls]
            predicted_y = self.preds[true_key_cls]
            cls_mb_output = self.mb_output[true_key_cls,:]
            sum_tensor = torch.zeros_like(cls_true_y, dtype=bool, device=strategy.device)
            _, top_n = torch.topk(cls_mb_output ,k=self.n, dim=1)
            for i in range(self.n):
                pred_overlap = top_n[:,i] == cls_true_y
                sum_tensor = torch.logical_or(sum_tensor, pred_overlap)
               
            if len(cls_true_y)>0:
                val.update(torch.ones_like(cls_true_y, device=strategy.device), sum_tensor*1, task_labels[0], device=strategy.device)



    def __str__(self):
        return self.season.capitalize()+"ClassTop_"+str(self.n)+"AvgAcc"+'_'+self._emit_at.capitalize()

class SeqTop_nAcc(GenericPluginMetric[float]):
    # Sequence wise top-n accuracy

    def __init__(self, n, reset_at, emit_at, mode):

        self._seq_wise_accuracy = Accuracy()
        self.n = n
        
        super(SeqTop_nAcc, self).__init__(
            self._seq_wise_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self._seq_wise_accuracy.reset(device=strategy.device)
        else:
            self._seq_wise_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None


    def result(self, strategy=None) -> float:

        task_labels = strategy.experience.task_labels
        unique_seq_codes = torch.unique(self.seq_codes)
        true_y = torch.zeros_like(unique_seq_codes, device=strategy.device)
        predicted_y = torch.zeros_like(unique_seq_codes, device=strategy.device)
        for i, seq_code in enumerate(unique_seq_codes):

            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]
            seq_mb_output_subset = self.mb_output[true_key_cls, :]
            top_n_predictions = break_tie_mode(seq_mb_output_subset, self.n)

            if seq_true_y not in top_n_predictions:
                predicted_y[i] = torch.tensor(1, device=strategy.device)
                
        self._seq_wise_accuracy.update(predicted_y, true_y, task_labels[0], device=strategy.device)


        result = self._seq_wise_accuracy.result()
        if self.n== 1 and not isinstance(self, SeasonSeqTop_nAcc):
            strategy.seq_top1_acc = result[0] 

        if result[0] == torch.tensor(0.):
            result[0] = 0.0
        return result

    def update(self, strategy):

        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])



    def __str__(self):
        return "SeqTop_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()

class SeqMaxAcc(GenericPluginMetric[float]):
    # sequence wise accuracy but using the class with the maximum node output as the prediction for the sequence
    # this metric assumes, that there is only one species per sequence

    def __init__(self, reset_at, emit_at, mode):

        self._seq_wise_accuracy = Accuracy()

        
        super(SeqMaxAcc, self).__init__(
            self._seq_wise_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self._seq_wise_accuracy.reset(device=strategy.device)
        else:
            self._seq_wise_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None


    def result(self, strategy=None) -> float:

        task_labels = strategy.experience.task_labels
        unique_seq_codes = torch.unique(self.seq_codes)
        true_y = torch.ones_like(unique_seq_codes, device=strategy.device)
        predicted_y = torch.zeros_like(unique_seq_codes, device=strategy.device)
        for i, seq_code in enumerate(unique_seq_codes):

            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]

            seq_mb_output_subset = self.mb_output[true_key_cls, :]
            seq_mb_output_subset = torch.nn.Softmax(dim =1)(seq_mb_output_subset)
            max_pos = torch.argmax(seq_mb_output_subset)
        
            max_cls = max_pos%(seq_mb_output_subset.shape[1])
       
            #seq_mb_output_max,_ = torch.max(self.mb_output, axis=0)
            #max_pred_in_seq = torch.argmax(seq_mb_output_max)

            if seq_true_y == max_cls:
             
                predicted_y[i] = torch.tensor(1, device=strategy.device)


        self._seq_wise_accuracy.update(predicted_y, true_y, task_labels[0], device=strategy.device)
        result = self._seq_wise_accuracy.result()
        strategy.seq_max_acc = result[0] 

        if result[0] == torch.tensor(0.):
            result[0] = 0.0
        return result

    def update(self, strategy):

        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])



    def __str__(self):
        return "SeqMAXAcc"+'_'+self._emit_at.capitalize()

class SeasonSeqTop_nAcc(SeqTop_nAcc):
    # Season-wise sequence accuracy
    # all season wise accuracy compute the accuracy on either all winter data or all summer data depending on the 
    # class variable 'season'. If both seasons should be tracked, there should be an instance for the summer and one 
    # for the winter data

    def __init__(self, num, winter_files_list, season, reset_at, emit_at, mode):
        self.winter_files_list = winter_files_list
        self.season = season

        super(SeasonSeqTop_nAcc, self).__init__( 
        num, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def update(self, strategy):

        in_season = np.isin(np.array(strategy.paths), self.winter_files_list)
        
        if self.season == 'summer':
            in_season = np.logical_not(in_season)

        if self.mb_y is None: 
            self.mb_y = strategy.mb_y[in_season]
            self.preds = strategy.preds[in_season]
            self.mb_output = strategy.mb_output[in_season,:]
            self.seq_codes = strategy.seq_codes[in_season]

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y[in_season]])
            self.preds = torch.cat([self.preds, strategy.preds[in_season]])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output[in_season,:]])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes[in_season]])

    def __str__(self):
        return self.season.capitalize()+"SeqTop_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()

class SeqClassTop_nAvgAcc(GenericPluginMetric[float]):
    # Here the the average of the sequence-wise class accuracies is computed 

    def __init__(self, nr_classes, n, reset_at, emit_at, mode):
        self.nr_classes = nr_classes
        self.cls_acc_dict = {}
        self.n = n 
        self.emit_at = emit_at
        for i in range(self.nr_classes):
            self.cls_acc_dict[i] = Accuracy()
        self._cls_avg_accuracy = Accuracy()
        super(SeqClassTop_nAvgAcc, self).__init__(
            self._cls_avg_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._cls_avg_accuracy.reset(device=strategy.device)
        else:
            self._cls_avg_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        for key, val in self.cls_acc_dict.items():
            val.reset(device=strategy.device)
        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None

    def result(self, strategy=None) -> float:
        task_labels = strategy.experience.task_labels

        unique_seq_codes = torch.unique(self.seq_codes)
        for i, seq_code in enumerate(unique_seq_codes):
            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]
            #seq_preds_subset = self.preds[true_key_cls]
            seq_mb_output_subset = self.mb_output[true_key_cls, :]
            seq_predicted_y = break_tie_mode(seq_mb_output_subset, self.n)

            gt_in_pred = torch.tensor([(seq_true_y in seq_predicted_y)*1], device=strategy.device)
            self.cls_acc_dict[seq_true_y.item()].update(torch.tensor([1], device=strategy.device), gt_in_pred, task_labels[0], device=strategy.device)



        cls_avg_acc = 0
        for key, val in self.cls_acc_dict.items():
            cls_returned_results = val.result()
            if len(cls_returned_results.keys())==1:
                cls_acc = cls_returned_results[0]
                cls_avg_acc += cls_acc

        if self.emit_at == 'stream':
            result = cls_avg_acc/self.nr_classes
        else:
            result = float(cls_avg_acc/torch.sum(strategy.classes_used_for_training))
        if self.n ==1 and not isinstance(self, SeasonSeqClassTop_nAcc):
            strategy.seq_cls_top1_avg_acc =  result

        return result

    def update(self, strategy):
        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])

    def __str__(self):
        return "SeqClassTop_"+str(self.n)+"AvgAcc"+'_'+self._emit_at.capitalize()


class SeqClassAvgMaxAcc(GenericPluginMetric[float]):
    # ere the the average of the sequence-wise class accuracies is computed, 
    # but with the maximum method for aggregating single image classifications/outputs into one prediction

    def __init__(self, nr_classes,reset_at, emit_at, mode):
        self.nr_classes = nr_classes
        self.cls_acc_dict = {}
        self.emit_at = emit_at
        for i in range(self.nr_classes):
            self.cls_acc_dict[i] = Accuracy()
        self._cls_avg_accuracy = Accuracy()
        super(SeqClassAvgMaxAcc, self).__init__(
            self._cls_avg_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._cls_avg_accuracy.reset(device=strategy.device)
        else:
            self._cls_avg_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        for key, val in self.cls_acc_dict.items():
            val.reset(device=strategy.device)
        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None

    def result(self, strategy=None) -> float:
        task_labels = strategy.experience.task_labels

        unique_seq_codes = torch.unique(self.seq_codes)

        for i, seq_code in enumerate(unique_seq_codes):
            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]
            seq_preds_subset = self.preds[true_key_cls]
            seq_mb_output_subset = self.mb_output[true_key_cls, :]
            seq_mb_output_subset = torch.nn.Softmax(dim =1)(seq_mb_output_subset)

            max_pos= torch.argmax(seq_mb_output_subset)
            max_pred_in_seq = max_pos%(seq_mb_output_subset.shape[1])


            gt_in_pred = torch.tensor([(seq_true_y == max_pred_in_seq)*1], device=strategy.device)
            self.cls_acc_dict[seq_true_y.item()].update(torch.tensor([1], device=strategy.device), gt_in_pred, task_labels[0], device=strategy.device)



        cls_avg_acc = 0
        for key, val in self.cls_acc_dict.items():
            cls_returned_results = val.result()
            if len(cls_returned_results.keys())==1:
                cls_acc = cls_returned_results[0]
                cls_avg_acc += cls_acc

        if self.emit_at == 'stream':
            result = cls_avg_acc/self.nr_classes
        else:
            result = float(cls_avg_acc/torch.sum(strategy.classes_used_for_training))
      
        strategy.seq_cls_max_avg_acc =  result

        return result

    def update(self, strategy):
        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])

    def __str__(self):
        return "SeqClassAvgMAXAcc"+'_'+self._emit_at.capitalize()



class SeasonSeqClassTop_nAcc(SeqClassTop_nAvgAcc):
    # season-wise sequence wise class average accuracy
    # all season wise accuracy compute the accuracy on either all winter data or all summer data depending on the 
    # class variable 'season'. If both seasons should be tracked, there should be an instance for the summer and one 
    # for the winter data


    def __init__(self, nr_classes, n, winter_files_list, season, reset_at, emit_at, mode):
        self.winter_files_list = winter_files_list
        self.season = season

        super(SeasonSeqClassTop_nAcc, self).__init__( nr_classes=nr_classes,
        n=n, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def update(self, strategy):

        in_season = np.isin(np.array(strategy.paths), self.winter_files_list)
        
        if self.season == 'summer':
            in_season = np.logical_not(in_season)

        if self.mb_y is None: 

            self.mb_y = strategy.mb_y[in_season]
            self.preds = strategy.preds[in_season]
            self.mb_output = strategy.mb_output[in_season,:]
            self.seq_codes = strategy.seq_codes[in_season]

        else:

            self.mb_y = torch.cat([self.mb_y, strategy.mb_y[in_season]])
            self.preds = torch.cat([self.preds, strategy.preds[in_season]])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output[in_season,:]])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes[in_season]])

    def __str__(self):
        return self.season.capitalize()+"SeqClassTop_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()


class SeqClasswiseTop_nAcc(GenericPluginMetric[float]):
      # sequence wise accuracy for each class

    def __init__(self, nr_classes, n, reset_at, emit_at, mode):
        self.nr_classes = nr_classes
        self.cls_acc_dict = {}
        self.n = n 
        self.emit_at = emit_at
        for i in range(self.nr_classes):
            self.cls_acc_dict[i] = Accuracy()
        self._cls_avg_accuracy = Accuracy()
        super(SeqClasswiseTop_nAcc, self).__init__(
            self._cls_avg_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._cls_avg_accuracy.reset(device=strategy.device)
        else:
            self._cls_avg_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        for key, val in self.cls_acc_dict.items():
            val.reset(device=strategy.device)
        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None

    def result(self, strategy=None) -> float:
        task_labels = strategy.experience.task_labels
        unique_seq_codes = torch.unique(self.seq_codes)


        for i, seq_code in enumerate(unique_seq_codes):
            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]
            seq_mb_output_subset = self.mb_output[true_key_cls, :]
            seq_predicted_y = break_tie_mode(seq_mb_output_subset, self.n)
            gt_in_pred = torch.tensor([(seq_true_y in seq_predicted_y)*1], device=strategy.device)
            self.cls_acc_dict[seq_true_y.item()].update(torch.tensor([1], device=strategy.device), gt_in_pred, task_labels[0], device=strategy.device)


        result_dict={}      
        for key, val in self.cls_acc_dict.items():
            result_dict[key] = val.result()

        result_dict[-1] = [-1]
        return result_dict


    def update(self, strategy):

        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:

            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])

      
    def __str__(self):
        return "SeqClasswiseTop_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()

class SeqAnyAcc(GenericPluginMetric[float]):
    # was the correct prediction made for any of the images in a sequence

    def __init__(self,  reset_at, emit_at, mode):

        self._seq_any_accuracy = Accuracy()
        super(SeqAnyAcc, self).__init__(
            self._seq_any_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self._seq_any_accuracy.reset(device=strategy.device)
        else:
            self._seq_any_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)
        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None

    def result(self, strategy=None) -> float:

        task_labels = strategy.experience.task_labels
        unique_seq_codes = torch.unique(self.seq_codes)
        true_y = torch.zeros_like(unique_seq_codes, device=strategy.device)
        predicted_y = torch.zeros_like(unique_seq_codes, device=strategy.device)

        for i, seq_code in enumerate(unique_seq_codes):
            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]
            seq_preds_subset = self.preds[true_key_cls]
            seq_predicted_y = torch.mode(seq_preds_subset)[0]
            if seq_true_y not in  seq_preds_subset:
                predicted_y[i] = torch.tensor(1, device=strategy.device)
                
        self._seq_any_accuracy.update(predicted_y, true_y, task_labels[0], device=strategy.device)



        if self._emit_at == 'stream' or strategy is None:
            result = self._seq_any_accuracy.result()
        else:
            result = self._seq_any_accuracy.result()
        strategy.seq_any_acc = result[0]

        if result[0] == torch.tensor(0.):
            result[0] = 0.0
        return result

    def update(self, strategy):

        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])

      
    def __str__(self):
        return "SeqAnyAcc"+'_'+self._emit_at.capitalize()




class Top_nAcc(GenericPluginMetric[float]):
    # standard top-n accuracy

    def __init__(self, n, reset_at, emit_at, mode):

        self._top_n_accuracy = Accuracy()
        self.n = n
        super(Top_nAcc, self).__init__(
            self._top_n_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self._top_n_accuracy.reset(device=strategy.device)
        else:
            self._top_n_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

    def result(self, strategy=None) -> float:
        if self._emit_at == 'stream' or strategy is None:
            result =  self._top_n_accuracy.result()
        else:
            result = self._top_n_accuracy.result(phase_and_task(strategy)[1])

        if self.n ==1 and not isinstance(self, SeasonTop_nAcc):
            strategy.top1_acc = result[0]
        if result[0] == torch.tensor(0.):
            result[0] = 0.0
        return result
    def update(self, strategy):

        task_labels = strategy.experience.task_labels
        sum_tensor = torch.zeros_like(strategy.mb_y, dtype=bool, device=strategy.device)

        _, top_n = torch.topk(strategy.mb_output,k=self.n, dim=1)
        for i in range(self.n):
            pred_overlap = top_n[:,i] == strategy.mb_y
            sum_tensor = torch.logical_or(sum_tensor, pred_overlap)
                
        self._top_n_accuracy.update(torch.ones_like(strategy.mb_y), sum_tensor*1, task_labels[0], device=strategy.device)


    def __str__(self):
        return "Top_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()



class SeasonTop_nAcc(Top_nAcc):
    # seasons-wise top accuracy. 
    # all season wise accuracy compute the accuracy on either all winter data or all summer data depending on the 
    # class variable 'season'. If both seasons should be tracked, there should be an instance for the summer and one 
    # for the winter data

    def __init__(self, n, winter_files_list, season, reset_at, emit_at, mode):
        self.winter_files_list = winter_files_list
        self.season = season

        super(SeasonTop_nAcc, self).__init__(
        n=n, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def update(self, strategy):

        in_season = np.isin(np.array(strategy.paths), self.winter_files_list)
        
        if self.season == 'summer':
            in_season = np.logical_not(in_season)

        self.mb_y = strategy.mb_y[in_season]
        self.preds = strategy.preds[in_season]
        self.mb_output = strategy.mb_output[in_season,:]
        self.seq_codes = strategy.seq_codes[in_season]

        task_labels = strategy.experience.task_labels
        sum_tensor = torch.zeros_like(self.mb_y, dtype=bool, device=strategy.device)

        _, top_n = torch.topk(self.mb_output,k=self.n, dim=1)
        for i in range(self.n):
            pred_overlap = top_n[:,i] == self.mb_y
            sum_tensor = torch.logical_or(sum_tensor, pred_overlap)
                
        self._top_n_accuracy.update(torch.ones_like(self.mb_y), sum_tensor*1, task_labels[0], device=strategy.device)



    def __str__(self):
        return self.season.capitalize()+"Top_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()



class BinaryAnimalAcc(GenericPluginMetric[float]):
    # At some stage there where three non-animal classes in the bavarian forrest data.
    # This metric basically lumps all animal classes together an when an image with an animal in it is 
    # classified as an image with any of the species classes, that is counted as a correctr prediction. 
    # Simirlary, all non-animal classes are lumped together
    # This was done as part of a study investigating the best practices for handling non-animal images that
    # will occure in a real lif dataset.

    def __init__(self, non_animal_class, reset_at, emit_at, mode):
        self.non_animal_class = non_animal_class
        self._binary_accuracy = Accuracy()
        super(BinaryAnimalAcc, self).__init__(
            self._binary_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._binary_accuracy.reset(device=strategy.device)
        else:
            self._binary_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)


    def result(self, strategy=None) -> float:
        if self._emit_at == 'stream' or strategy is None:
            result = self._binary_accuracy.result()
        else:
            result = self._binary_accuracy.result(phase_and_task(strategy)[1])

        strategy.bin_animal_acc = result[0]
        return result

    def update(self, strategy):

        task_labels = strategy.experience.task_labels

        sum_tensor = torch.ones_like(strategy.mb_y, dtype=bool, device=strategy.device)
        true_ghost = torch.logical_and(torch.isin(strategy.preds,self.non_animal_class),torch.isin(strategy.mb_y,self.non_animal_class))
        true_animal = torch.logical_and(torch.logical_not(torch.isin(strategy.preds,self.non_animal_class)),torch.logical_not(torch.isin(strategy.mb_y,self.non_animal_class)))

        correctly_predicted_animal = torch.logical_or(true_animal, true_ghost)*1

        self._binary_accuracy.update(sum_tensor, correctly_predicted_animal, task_labels[0], device=strategy.device)

    def __str__(self):
        return "BinaryAnimalAcc"+'_'+self._emit_at.capitalize()

class QuatrupleAnimalAcc(GenericPluginMetric[float]):
    # At some stage there where three non-animal classes in the bavarian forrest data.
    # This metric basically lumps all animal classes together an when an image with an animal in it is 
    # classified as an image with any of the species classes, that is counted as a correctr prediction. 

    def __init__(self, non_animal_class, reset_at, emit_at, mode):
        self.non_animal_class = non_animal_class
        self.quadruple_acc = Accuracy()
        super(QuatrupleAnimalAcc, self).__init__(
            self.quadruple_acc, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self.quadruple_acc.reset(device=strategy.device)
        else:
            self.quadruple_acc.reset(phase_and_task(strategy)[1], device=strategy.device)


    def result(self, strategy=None) -> float:
        if self._emit_at == 'stream' or strategy is None:
            result = self.quadruple_acc.result()
        else:
            result = self.quadruple_acc.result(phase_and_task(strategy)[1])
        strategy.quadruple_acc = result[0]
        return result

    def update(self, strategy):
        task_labels = strategy.experience.task_labels

        sum_tensor = torch.zeros_like(strategy.mb_y, dtype=bool, device=strategy.device)
        fake_gt = torch.ones_like(strategy.mb_y, dtype=bool, device=strategy.device)
        for i in self.non_animal_class:

            true_non_animal = torch.logical_and(torch.isin(strategy.preds,torch.tensor(i)),torch.isin(strategy.mb_y,torch.tensor(i)))
            sum_tensor = torch.logical_or(sum_tensor, true_non_animal)

        true_animal = torch.logical_and(torch.logical_not(torch.isin(strategy.preds,self.non_animal_class)),torch.logical_not(torch.isin(strategy.mb_y,self.non_animal_class)))

        correctly_predicted= torch.logical_or(true_animal, sum_tensor)*1
        self.quadruple_acc.update(fake_gt*1, correctly_predicted, task_labels[0], device=strategy.device)

    def __str__(self):
        return "QuatrupleAnimalAcc"+'_'+self._emit_at.capitalize()


class SeqBinaryAnimalAcc(GenericPluginMetric[float]):
    # sequence-wise version of the binary animal accuracy

    def __init__(self, non_animal_class, reset_at, emit_at, mode):
        self.non_animal_class = non_animal_class
        self._seq_binary_accuracy = Accuracy()
        super(SeqBinaryAnimalAcc, self).__init__(
            self._seq_binary_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self._seq_binary_accuracy.reset( device=strategy.device)
        else:
            self._seq_binary_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None


    def result(self, strategy=None) -> float:

        task_labels = strategy.experience.task_labels
        unique_seq_codes = torch.unique(self.seq_codes)
        true_y = torch.zeros_like(unique_seq_codes, device=strategy.device)

        for i, seq_code in enumerate(unique_seq_codes):
            true_key_cls = torch.eq(self.seq_codes,seq_code)
            seq_true_y = self.mb_y[true_key_cls][0]
            seq_mb_output_subset = self.mb_output[true_key_cls, :]
            #seq_preds_subset = self.preds[true_key_cls]
            top_n_predictions = break_tie_mode(seq_mb_output_subset, 1)

            if (torch.isin(top_n_predictions, self.non_animal_class) and torch.isin(seq_true_y, self.non_animal_class)) or (torch.logical_not(torch.isin(top_n_predictions, self.non_animal_class)) and torch.logical_not(torch.isin(seq_true_y, self.non_animal_class))):
                self._seq_binary_accuracy.update(torch.tensor([1], device=strategy.device), torch.tensor([1], device=strategy.device), task_labels[0], device=strategy.device)
            else:
                self._seq_binary_accuracy.update(torch.tensor([1], device=strategy.device), torch.tensor([0], device=strategy.device), task_labels[0], device=strategy.device)

        if self._emit_at == 'stream' or strategy is None:

            result =  self._seq_binary_accuracy.result()
        else:
            result = self._seq_binary_accuracy.result(phase_and_task(strategy)[1])

        strategy.seq_bin_animal_acc = result[0]
        return result

    def update(self, strategy):
        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])


    def __str__(self):
        return "SeqBinaryAnimalAcc"+'_'+self._emit_at.capitalize()



class IgnoreNonAnimalOutputsTop_nAcc(GenericPluginMetric[float]):
    # Here the top-n accuracy of all images with animals are computed for settings where there are non-animal
    # classes in the datastream. This too was done as a way of analysing how to handle images without animals in them
    # wheter a two step approach, where non-animal images are filtered in the first step performs better than a one step approach 


    def __init__(self, non_animal_class, n, nr_classes, reset_at, emit_at, mode):
        self.non_animal_class = non_animal_class
        self.n = n
        self.nr_classes = nr_classes
        self._ignore_non_animal_accuracy = Accuracy()
        super(IgnoreNonAnimalOutputsTop_nAcc, self).__init__(
            self._ignore_non_animal_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._ignore_non_animal_accuracy.reset(device=strategy.device)
        else:
            self._ignore_non_animal_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

    def result(self, strategy=None) -> float:
        if self._emit_at == 'stream' or strategy is None:
            result = self._ignore_non_animal_accuracy.result()
        else:
            result = self._ignore_non_animal_accuracy.result(phase_and_task(strategy)[1])

        strategy.ignore_non_animal_acc = result[0]
        return result

    def update(self, strategy):

        task_labels = strategy.experience.task_labels

        gt_animal = torch.logical_not(torch.isin(strategy.mb_y, self.non_animal_class))

        animal_mby = strategy.mb_y[gt_animal] 
        animal_mb_output = strategy.mb_output[gt_animal,:]

        species_columns = torch.range(0,self.nr_classes-1, dtype=int, device=strategy.device)
        species_columns = torch.logical_not(torch.isin(species_columns, self.non_animal_class))
        animal_mb_output = animal_mb_output[:,species_columns]

        sum_tensor = torch.zeros_like(animal_mby, dtype=bool, device=strategy.device)

        _, top_n = torch.topk(animal_mb_output,k=self.n, dim=1)
        for i in range(self.n):
            pred_overlap = top_n[:,i] == animal_mby
            sum_tensor = torch.logical_or(sum_tensor, pred_overlap)

        self._ignore_non_animal_accuracy.update(torch.ones_like(animal_mby), sum_tensor*1, task_labels[0], device=strategy.device)

    def __str__(self):
        return "IgnoreNonAnimalOutputsTop_"+str(self.n)+"Acc"+'_'+self._emit_at.capitalize()


class SeqIgnoreNonAnimalOutputsAcc(GenericPluginMetric[float]):
    # sequence wise accuracy of all images with animals

    def __init__(self, non_animal_class, nr_classes, reset_at, emit_at, mode):
        self.non_animal_class = non_animal_class
        self.nr_classes = nr_classes
        self._seq_ignore_non_animal_accuracy = Accuracy()
        super(SeqIgnoreNonAnimalOutputsAcc, self).__init__(
            self._seq_ignore_non_animal_accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:
        if self._reset_at == 'stream' or strategy is None:
            self._seq_ignore_non_animal_accuracy.reset(device=strategy.device)
        else:
            self._seq_ignore_non_animal_accuracy.reset(phase_and_task(strategy)[1], device=strategy.device)

        self.mb_y = None
        self.pred = None
        self.mb_output = None
        self.seq_codes = None


    def result(self, strategy=None) -> float:
        task_labels = strategy.experience.task_labels
        unique_seq_codes = torch.unique(strategy.seq_codes)
        true_y = torch.tensor([0], device=strategy.device)

        for i, seq_code in enumerate(unique_seq_codes):

            true_key_cls = torch.eq(strategy.seq_codes,seq_code)

            seq_true_y = strategy.mb_y[true_key_cls][0]
            if seq_true_y not in self.non_animal_class:
                predicted_y = torch.tensor([0], device=strategy.device)
                seq_mb_output_subset = strategy.mb_output[true_key_cls, :]
                species_columns = torch.range(0,self.nr_classes-1, dtype=int, device=strategy.device)
                species_columns = torch.logical_not(torch.isin(species_columns, self.non_animal_class))
                seq_mb_output_subset = seq_mb_output_subset[:,species_columns]
                top_n_predictions = break_tie_mode(seq_mb_output_subset, 1)

                if seq_true_y not in top_n_predictions:
                    predicted_y= torch.tensor([1], device=strategy.device)

                self._seq_ignore_non_animal_accuracy.update(predicted_y, true_y, task_labels[0], device=strategy.device)


        if self._emit_at == 'stream' or strategy is None:
            result =  self._seq_ignore_non_animal_accuracy.result()
        else:
            result =  self._seq_ignore_non_animal_accuracy.result(phase_and_task(strategy)[1])
        strategy.seq_ignore_non_animal_acc = result[0]
        return result

    def update(self, strategy):
        if self.mb_y is None: 
            self.mb_y = strategy.mb_y
            self.preds = strategy.preds
            self.mb_output = strategy.mb_output
            self.seq_codes = strategy.seq_codes

        else:
            self.mb_y = torch.cat([self.mb_y, strategy.mb_y])
            self.preds = torch.cat([self.preds, strategy.preds])
            self.mb_output = torch.cat([self.mb_output, strategy.mb_output])
            self.seq_codes = torch.cat([self.seq_codes, strategy.seq_codes])

    def __str__(self):
        return "SeqIgnoreNonAnimalOutputsAcc"+'_'+self._emit_at.capitalize()

class ClassIgnoreNonAnimalOutputsTop_nAvgAcc(GenericPluginMetric[float]):
    #class average accuracy of all images with animals
    def __init__(self, nr_classes,n, non_animal_class, reset_at, emit_at, mode):
        self.nr_classes = nr_classes
        self.n = n
        self.non_animal_class = non_animal_class
        self.cls_ignore_non_animal_acc_dict = {}
        self.emit_at = emit_at

        for i in range(self.nr_classes):
            if i not in self.non_animal_class:
                self.cls_ignore_non_animal_acc_dict[i] = Accuracy()
        self._cls_ignore_non_animal_acc= Accuracy()
        super(ClassIgnoreNonAnimalOutputsTop_nAvgAcc, self).__init__(
            self._cls_ignore_non_animal_acc, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._cls_ignore_non_animal_acc.reset(device=strategy.device)
        else:
            self._cls_ignore_non_animal_acc.reset(phase_and_task(strategy)[1], device=strategy.device)

        for key, val in self.cls_ignore_non_animal_acc_dict.items():
            val.reset(device=strategy.device)

    def result(self, strategy=None) -> float:

        cls_avg_acc = 0
        for key, val in self.cls_ignore_non_animal_acc_dict.items():
            cls_returned_results = val.result()
            if len(cls_returned_results.keys())==1:
                cls_acc = cls_returned_results[0]
                cls_avg_acc += cls_acc

        if self.emit_at == 'stream':
            result = cls_avg_acc/float(len(self.cls_ignore_non_animal_acc_dict.keys()))
        else:
            result = float(cls_avg_acc/torch.sum(torch.logical_not(torch.isin(strategy.classes_used_for_training, self.non_animal_class))*1))
        strategy.cls_ignore_non_animal_acc = result    
        return result

    def update(self, strategy):

        task_labels = strategy.experience.task_labels
        for key, val in self.cls_ignore_non_animal_acc_dict.items():
            true_key_cls = torch.eq(strategy.mb_y,key)
            cls_true_y = strategy.mb_y[true_key_cls]
            predicted_y = strategy.preds[true_key_cls]
            cls_mb_output = strategy.mb_output[true_key_cls,:]
            species_columns = torch.range(0,self.nr_classes-1, dtype=int, device=strategy.device)
            species_columns = torch.logical_not(torch.isin(species_columns, self.non_animal_class))
            cls_mb_output = cls_mb_output[:,species_columns]
            sum_tensor = torch.zeros_like(cls_true_y, dtype=bool, device=strategy.device)
            _, top_n = torch.topk(cls_mb_output ,k=self.n, dim=1)
            for i in range(self.n):
                pred_overlap = top_n[:,i] == cls_true_y
                sum_tensor = torch.logical_or(sum_tensor, pred_overlap)

            if len(cls_true_y)>0:
                val.update(torch.ones_like(cls_true_y, device=strategy.device), sum_tensor*1, task_labels[0], device=strategy.device)


    def __str__(self):
        return "ClassIgnoreNonAnimalOutputsTop_"+str(self.n)+"AvgAcc"+'_'+self._emit_at.capitalize()


class Accuracy(Metric[float]):
    """
    The Accuracy metric. This is a standalone metric.

    The metric keeps a dictionary of <task_label, accuracy value> pairs.
    and update the values through a running average over multiple
    <prediction, target> pairs of Tensors, provided incrementally.
    The "prediction" and "target" tensors may contain plain labels or
    one-hot/logit vectors.

    Each time `result` is called, this metric emits the average accuracy
    across all predictions made since the last `reset`.

    The reset method will bring the metric to its initial state. By default
    this metric in its initial state will return an accuracy value of 0.
    """

    def __init__(self):
        """
        Creates an instance of the standalone Accuracy metric.

        By default this metric in its initial state will return an accuracy
        value of 0. The metric can be updated by using the `update` method
        while the running accuracy can be retrieved using the `result` method.
        """
        super().__init__()
        self._mean_accuracy = defaultdict(Mean)
        """
        The mean utility that will be used to store the running accuracy
        for each task label.
        """

    @torch.no_grad()
    def update(self, predicted_y: Tensor, true_y: Tensor,
               task_labels: Union[float, Tensor], device='cpu') -> None:
        """
        Update the running accuracy given the true and predicted labels.
        Parameter `task_labels` is used to decide how to update the inner
        dictionary: if Float, only the dictionary value related to that task
        is updated. If Tensor, all the dictionary elements belonging to the
        task labels will be updated.

        :param predicted_y: The model prediction. Both labels and logit vectors
            are supported.
        :param true_y: The ground truth. Both labels and one-hot vectors
            are supported.
        :param task_labels: the int task label associated to the current
            experience or the task labels vector showing the task label
            for each pattern.

        :return: None.
        """
        if len(true_y) != len(predicted_y):
            raise ValueError('Size mismatch for true_y and predicted_y tensors')

        if isinstance(task_labels, Tensor) and len(task_labels) != len(true_y):
            raise ValueError('Size mismatch for true_y and task_labels tensors')

        true_y = torch.as_tensor(true_y)
        predicted_y = torch.as_tensor(predicted_y)

        # Check if logits or labels
        if len(predicted_y.shape) > 1:
            # Logits -> transform to labels
            predicted_y = torch.max(predicted_y, 1)[1]

        if len(true_y.shape) > 1:
            # Logits -> transform to labels
            true_y = torch.max(true_y, 1)[1]
        _mean_accuracy_list = [list(self._mean_accuracy.values())[k].result() for k in self._mean_accuracy.keys()]

        if isinstance(task_labels, int):
            check_device1 = torch.sum(torch.eq(predicted_y, true_y))

            true_positives = torch.sum(torch.eq(predicted_y, true_y))
            total_patterns = torch.tensor(len(true_y))

            self._mean_accuracy[task_labels].update(
                true_positives / total_patterns, total_patterns)

        elif isinstance(task_labels, Tensor):
            for pred, true, t in zip(predicted_y, true_y, task_labels):
                true_positives = (pred == true).float().item()
                self._mean_accuracy[t.item()].update(
                    true_positives, 1, device=device)

        else:
            raise ValueError(f"Task label type: {type(task_labels)}, "
                             f"expected int/float or Tensor")

    def result(self, task_label=None) -> Dict[int, float]:
        """
        Retrieves the running accuracy.

        Calling this method will not change the internal state of the metric.

        :param task_label: if None, return the entire dictionary of accuracies
            for each task. Otherwise return the dictionary
            `{task_label: accuracy}`.
        :return: A dict of running accuracies for each task label,
            where each value is a float value between 0 and 1.
        """
        assert(task_label is None or isinstance(task_label, int))
        if task_label is None:
            return {k: v.result() for k, v in self._mean_accuracy.items()}
        else:
            return {task_label: self._mean_accuracy[task_label].result()}

    def reset(self, task_label=None, device='cpu') -> None:
        """
        Resets the metric.
        :param task_label: if None, reset the entire dictionary.
            Otherwise, reset the value associated to `task_label`.

        :return: None.
        """
        assert(task_label is None or isinstance(task_label, int))
        if task_label is None:
            self._mean_accuracy = defaultdict(Mean)
        else:
            self._mean_accuracy[task_label].reset()


class AccuracyPluginMetric(GenericPluginMetric[float]):
    """
    Base class for all accuracies plugin metrics
    """
    def __init__(self, reset_at, emit_at, mode):
        self._accuracy = Accuracy()

        super(AccuracyPluginMetric, self).__init__(
            self._accuracy, reset_at=reset_at, emit_at=emit_at,
            mode=mode)

    def reset(self, strategy=None) -> None:

        if self._reset_at == 'stream' or strategy is None:
            self._metric.reset(device=strategy.device)
        else:
            self._metric.reset(phase_and_task(strategy)[1], device=strategy.device)

        _mean_accuracy_list = [list(self._metric._mean_accuracy.values())[k].result() for k in self._metric._mean_accuracy.keys()]


    def result(self, strategy=None) -> float:
        if self._emit_at == 'stream' or strategy is None:
            return self._metric.result()
        else:
            return self._metric.result(phase_and_task(strategy)[1])

    def update(self, strategy):
        # task labels defined for each experience
        task_labels = strategy.experience.task_labels
        if len(task_labels) > 1:
            # task labels defined for each pattern
            task_labels = strategy.mb_task_id
        else:
            task_labels = task_labels[0]

        self._accuracy.update(strategy.mb_output, strategy.mb_y, task_labels, device=strategy.device)


class MinibatchAccuracy(AccuracyPluginMetric):
    """
    The minibatch plugin accuracy metric.
    This metric only works at training time.

    This metric computes the average accuracy over patterns
    from a single minibatch.
    It reports the result after each iteration.

    If a more coarse-grained logging is needed, consider using
    :class:`EpochAccuracy` instead.
    """
    def __init__(self):
        """
        Creates an instance of the MinibatchAccuracy metric.
        """
        super(MinibatchAccuracy, self).__init__(
            reset_at='iteration', emit_at='iteration', mode='train')

    def __str__(self):
        return "Top1_Acc_MB"


class EpochAccuracy(AccuracyPluginMetric):
    """
    The average accuracy over a single training epoch.
    This plugin metric only works at training time.

    The accuracy will be logged after each training epoch by computing
    the number of correctly predicted patterns during the epoch divided by
    the overall number of patterns encountered in that epoch.
    """

    def __init__(self):
        """
        Creates an instance of the EpochAccuracy metric.
        """

        super(EpochAccuracy, self).__init__(
            reset_at='epoch', emit_at='epoch', mode='train')

    def __str__(self):
        return "Top1_Acc_Epoch"


class RunningEpochAccuracy(AccuracyPluginMetric):
    """
    The average accuracy across all minibatches up to the current
    epoch iteration.
    This plugin metric only works at training time.

    At each iteration, this metric logs the accuracy averaged over all patterns
    seen so far in the current epoch.
    The metric resets its state after each training epoch.
    """

    def __init__(self):
        """
        Creates an instance of the RunningEpochAccuracy metric.
        """

        super(RunningEpochAccuracy, self).__init__(
            reset_at='epoch', emit_at='iteration', mode='train')

    def __str__(self):
        return "Top1_RunningAcc_Epoch"


class ExperienceAccuracy(AccuracyPluginMetric):
    """
    At the end of each experience, this plugin metric reports
    the average accuracy over all patterns seen in that experience.
    This metric only works at eval time.
    """

    def __init__(self):
        """
        Creates an instance of ExperienceAccuracy metric
        """
        super(ExperienceAccuracy, self).__init__(
            reset_at='experience', emit_at='experience', mode='eval')

    def __str__(self):
        return "Top1_Acc_Exp"


class StreamAccuracy(AccuracyPluginMetric):
    """
    At the end of the entire stream of experiences, this plugin metric
    reports the average accuracy over all patterns seen in all experiences.
    This metric only works at eval time.
    """

    def __init__(self):
        """
        Creates an instance of StreamAccuracy metric
        """
        super(StreamAccuracy, self).__init__(
            reset_at='stream', emit_at='stream', mode='eval')

    def __str__(self):
        return "Top1_Acc_Stream"


class TrainedExperienceAccuracy(AccuracyPluginMetric):
    """
    At the end of each experience, this plugin metric reports the average
    accuracy for only the experiences that the model has been trained on so far.

    This metric only works at eval time.
    """

    def __init__(self):
        """
        Creates an instance of TrainedExperienceAccuracy metric by first 
        constructing AccuracyPluginMetric
        """
        super(TrainedExperienceAccuracy, self).__init__(
            reset_at='stream', emit_at='stream', mode='eval')
        self._current_experience = 0

    def after_training_exp(self, strategy) -> None:
        self._current_experience = strategy.experience.current_experience
        # Reset average after learning from a new experience 
        AccuracyPluginMetric.reset(self, strategy)
        return AccuracyPluginMetric.after_training_exp(self, strategy)
        
    def update(self, strategy):
        """
        Only update the accuracy with results from experiences that have been 
        trained on
        """
        if strategy.experience.current_experience <= self._current_experience:
            AccuracyPluginMetric.update(self, strategy)

    def __str__(self):
        return "Accuracy_On_Trained_Experiences"


def accuracy_metrics(*, 
                     minibatch=False,
                     epoch=False,
                     epoch_running=False,
                     experience=False,
                     stream=False,
                     trained_experience=False) -> List[PluginMetric]:
    """
    Helper method that can be used to obtain the desired set of
    plugin metrics.

    :param minibatch: If True, will return a metric able to log
        the minibatch accuracy at training time.
    :param epoch: If True, will return a metric able to log
        the epoch accuracy at training time.
    :param epoch_running: If True, will return a metric able to log
        the running epoch accuracy at training time.
    :param experience: If True, will return a metric able to log
        the accuracy on each evaluation experience.
    :param stream: If True, will return a metric able to log
        the accuracy averaged over the entire evaluation stream of experiences.
    :param trained_experience: If True, will return a metric able to log
        the average evaluation accuracy only for experiences that the
        model has been trained on         

    :return: A list of plugin metrics.
    """

    metrics = []
    if minibatch:
        metrics.append(MinibatchAccuracy())

    if epoch:
        metrics.append(EpochAccuracy())

    if epoch_running:
        metrics.append(RunningEpochAccuracy())

    if experience:
        metrics.append(ExperienceAccuracy())

    if stream:
        metrics.append(StreamAccuracy())

    if trained_experience:
        metrics.append(TrainedExperienceAccuracy())

    return metrics


__all__ = [
    'Accuracy',
    'Top1AccTransfer',
    'ClasswiseTop_nAcc',
    'ClassIgnoreNonAnimalOutputsTop_nAvgAcc',
    'SeqIgnoreNonAnimalOutputsAcc',
    'IgnoreNonAnimalOutputsTop_nAcc',
    'SeqTop_nAcc',
    'SeqMaxAcc',
    'SeqClassAvgMaxAcc',
    'SeasonSeqTop_nAcc',
    'Top_nAcc',  
    'SeasonTop_nAcc',
    'ClassTop_nAvgAcc',
    'SeasonClassTop_nAcc',
    'SeqBinaryAnimalAcc',
    'BinaryAnimalAcc',
    'QuatrupleAnimalAcc',
    'SeqClassTop_nAvgAcc',
    'SeasonSeqClassTop_nAcc',
    'SeqClasswiseTop_nAcc',
    'SeqAnyAcc',
    'MinibatchAccuracy',
    'EpochAccuracy',
    'RunningEpochAccuracy',
    'ExperienceAccuracy',
    'StreamAccuracy',
    'TrainedExperienceAccuracy',
    'accuracy_metrics'
]
