# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import os

import numpy as np

from fairseq.data import (
    ConcatSentencesDataset,
    data_utils,
    Dictionary,
    IdDataset,
    NestedDictionaryDataset,
    NumSamplesDataset,
    NumelDataset,
    OffsetTokensDataset,
    PrependTokenDataset,
    RawLabelDataset,
    RightPadDataset,
    SortDataset,
    StripTokenDataset,
    TruncateDataset,
    RawAudioDataset
)

from . import FairseqTask, register_task


@register_task('emotion_prediction_old')
class EmotionPredictionTask_Old(FairseqTask):
    """
    Sentence (or sentence pair) prediction (classification or regression) task.

    Args:
        dictionary (Dictionary): the dictionary for the input of the task
    """

    @staticmethod
    def add_args(parser):
        """Add task-specific arguments to the parser."""
        parser.add_argument('--data_T', metavar='FILE',
                            help='file prefix for Text data')

        parser.add_argument('--data_A', metavar='FILE',
                            help='file prefix for Audio data')
        parser.add_argument('--sample-rate', default=16000, type=int,
                            help='target sample rate. audio files will be up/down sampled to this rate')

        parser.add_argument('--max-sample-size', default=None, type=int,
                            help='max sample size to crop to for batching. default = min sample length')
        parser.add_argument('--min-sample-size', default=None, type=int,
                            help='min sample size to crop to for batching. default = same as --max-sample-size')

        parser.add_argument('--num-classes', type=int, default=-1,
                            help='number of classes')
        parser.add_argument('--init-token', type=int, default=None,
                            help='add token at the beginning of each batch item')
        parser.add_argument('--separator-token', type=int, default=None,
                            help='add separator token between inputs')
        parser.add_argument('--regression-target', action='store_true', default=False)
        parser.add_argument('--no-shuffle', action='store_true', default=False)
        parser.add_argument('--truncate-sequence', action='store_true', default=False,
                            help='Truncate sequence to max_sequence_length')

    def __init__(self, args, data_dictionary, label_dictionary):
        super().__init__(args)
        self.dictionary = data_dictionary
        self.label_dictionary = label_dictionary

    @classmethod
    def load_dictionary(cls, args, filename, source=True):
        """Load the dictionary from the filename

        Args:
            filename (str): the filename
        """
        dictionary = Dictionary.load(filename)
        dictionary.add_symbol('<mask>')
        return dictionary

    @classmethod
    def setup_task(cls, args, **kwargs): # to_d0 :Might need to modify this also to merge audio and video

        #Check how they make the dictionary file. Is that dataset oriented or unique ?
      
        assert args.num_classes > 0, 'Must set --num-classes'

        args.tokens_per_sample = args.max_positions
        #print(args.data)
        #exit()

        #Getting the statisitc with the created dictionary <fairseq.data.dictionary.Dictionary object at 0x7f6fb59ac550>
        # load data dictionary
        data_dict = cls.load_dictionary( #Go to the text dictionary
            args,
            os.path.join(args.data_T, 'input0', 'dict.txt'),
            source=True,
        )
    
        print('| [input] dictionary: {} types'.format(len(data_dict)))  #There are 50264 symbols 

        label_dict = None
        if not args.regression_target:

            # load label dictionary
            label_dict = cls.load_dictionary(
                args,
                os.path.join(args.data_T, 'label', 'dict.txt'),
                source=False,
            )
            print('| [label] dictionary: {} types'.format(len(label_dict))) #symbols that can be found it the label section
        else:
            label_dict = data_dict

        
        return EmotionPredictionTask_Old(args, data_dict, label_dict)  #Initalization of an instance of our class

    def load_dataset(self, split, combine=False, **kwargs): 
        '''
        #to_do : Merge audio and video data loaders
        '''
        """Load a given dataset split (e.g., train, valid, test)."""

     
        #wave2vec_dataloader needs to send the 
        ################################################
        manifest = os.path.join(self.args.data_A, '{}.tsv'.format(split+'_a'))
        
        self.datasets[split+'_audio'] = RawAudioDataset(manifest,
                                               sample_rate=self.args.sample_rate,
                                               max_sample_size=self.args.max_sample_size,
                                               min_sample_size=self.args.min_sample_size)
        ##################################################

        def get_path(type, split):
            return os.path.join(self.args.data_T, type, split)

        def make_dataset(type, dictionary):
            split_path = get_path(type, split)

            dataset = data_utils.load_indexed_dataset(
                split_path,
                self.source_dictionary,
                self.args.dataset_impl,
                combine=combine,
            )



            return dataset

        input0 = make_dataset('input0', self.source_dictionary)
        assert input0 is not None, 'could not find dataset: {}'.format(get_path(type, split))
        input1 = make_dataset('input1', self.source_dictionary)
  

        if self.args.init_token is not None:
            input0 = PrependTokenDataset(input0, self.args.init_token)

        if input1 is None:
            src_tokens = input0
        else:
            if self.args.separator_token is not None:
                input1 = PrependTokenDataset(input1, self.args.separator_token)

            src_tokens = ConcatSentencesDataset(input0, input1)

      
        with data_utils.numpy_seed(self.args.seed):
            shuffle = np.random.permutation(len(src_tokens))


        if self.args.truncate_sequence:
            src_tokens = TruncateDataset(src_tokens, self.args.max_positions) #Tranucate the input seq


        dataset = {
            'id': IdDataset(),
            'net_input': {
                'src_tokens': RightPadDataset(
                    src_tokens,
                    pad_idx=self.source_dictionary.pad(),
                ),
                'src_lengths': NumelDataset(src_tokens, reduce=False),
            },
            'nsentences': NumSamplesDataset(),
            'ntokens': NumelDataset(src_tokens, reduce=True),
        }

        if not self.args.regression_target:
            label_dataset = make_dataset('label', self.target_dictionary)
            if label_dataset is not None:
                dataset.update(
                    target=OffsetTokensDataset(
                        StripTokenDataset(
                            label_dataset,
                            id_to_strip=self.target_dictionary.eos(),
                        ),
                        offset=-self.target_dictionary.nspecial,
                    )
                )
             
        else:
            label_path = "{0}.label".format(get_path('label', split))
            if os.path.exists(label_path):
                dataset.update(
                    target=RawLabelDataset([
                        float(x.strip()) for x in open(label_path).readlines()
                    ])
                )

        nested_dataset = NestedDictionaryDataset(
            dataset,
            sizes=[src_tokens.sizes],
        )

        if self.args.no_shuffle:
            dataset = nested_dataset
        else:
            dataset = SortDataset(
                nested_dataset,
                # shuffle
                sort_order=[shuffle],
            )

        print("| Loaded {0} with #samples: {1}".format(split, len(dataset)))

        self.datasets[split] = dataset

        return self.datasets[split] , self.datasets[split+'_audio']

    def build_model(self, args):
        
        from fairseq import models
      
        model = models.build_model(args, self) #/home/gsir059/Documents/PhD/MulFie/fairseq/fairseq/models/__init__.py'
        print("Model Initialization Done")
        model.register_classification_head(
            'emotion_classification_head_old',
            num_classes=self.args.num_classes,
        )
        #print(model,"sentece_prediction.py")
      
        return model

    def max_positions(self):
        return self.args.max_positions

    @property
    def source_dictionary(self):
        return self.dictionary

    @property
    def target_dictionary(self):
        return self.label_dictionary
