# 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 csv
import io
import logging
import re
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, NamedTuple
import pickle
import random
import os
import gc
from psutil import Process
import sys

import spacy
import numpy as np
import torch
from fairseq.data import (
    ConcatDataset,
    Dictionary,
    FairseqDataset,
    ResamplingDataset,
    data_utils as fairseq_data_utils,
)
from fairseq.data.audio.audio_utils import (
    get_fbank,
    get_waveform,
    read_from_stored_zip,
    is_npy_data,
    is_sf_audio_data,
    parse_path,
    FEATURE_OR_SF_AUDIO_FILE_EXTENSIONS,
)
from fairseq.data.audio.feature_transforms import CompositeAudioFeatureTransform
from fairseq.data.audio.speech_to_text_dataset import (
    S2TDataConfig,
    SpeechToTextDataset,
)
from fairseq.data.audio.speech_to_text_dataset import (
    get_features_or_waveform,
    get_features_or_waveform_from_stored_zip,
    get_features_from_npy_or_audio,
    _collate_frames
)

from fairseq import utils


logger = logging.getLogger(__name__)


class SpeechToTextDatasetItem(NamedTuple):
    index: int
    source: torch.Tensor
    target: Optional[torch.Tensor] = None


class StrASTDataset(SpeechToTextDataset):
    LANG_TAG_TEMPLATE = "<lang:{}>"

    def __init__(
        self,
        split: str,
        is_train_split: bool,
        cfg: S2TDataConfig,
        audio_paths: List[str],
        n_frames: List[int],
        src_texts: Optional[List[str]] = None,
        tgt_texts: Optional[List[str]] = None,
        speakers: Optional[List[str]] = None,
        src_langs: Optional[List[str]] = None,
        tgt_langs: Optional[List[str]] = None,
        ids: Optional[List[str]] = None,
        tgt_dict: Optional[Dictionary] = None,
        pre_tokenizer=None,
        bpe_tokenizer=None,
        frame_min=None,
        frame_max=None,
        whereToAugment=None,
    ):
        self.split, self.is_train_split = split, is_train_split
        self.cfg = cfg
        self.audio_paths, self.n_frames = audio_paths, n_frames
        self.n_samples = len(audio_paths)
        assert len(n_frames) == self.n_samples > 0
        assert src_texts is None or len(src_texts) == self.n_samples
        assert tgt_texts is None or len(tgt_texts) == self.n_samples
        assert speakers is None or len(speakers) == self.n_samples
        assert src_langs is None or len(src_langs) == self.n_samples
        assert tgt_langs is None or len(tgt_langs) == self.n_samples
        assert ids is None or len(ids) == self.n_samples
        assert (tgt_dict is None and tgt_texts is None) or (
            tgt_dict is not None and tgt_texts is not None
        )
        self.src_texts, self.tgt_texts = src_texts, tgt_texts
        self.src_langs, self.tgt_langs = src_langs, tgt_langs
        self.tgt_dict = tgt_dict
        self.check_tgt_lang_tag()
        self.ids = ids
        self.shuffle = cfg.shuffle if is_train_split else False

        self.feature_transforms = CompositeAudioFeatureTransform.from_config_dict(
            self.cfg.get_feature_transforms(split, is_train_split)
        )

        self.pre_tokenizer = pre_tokenizer
        self.bpe_tokenizer = bpe_tokenizer

        self.tgt_lens = self.get_tgt_lens_and_check_oov()

        self.audioPath_prefix = '/'.join(self.audio_paths[0].split('/')[:-1])
        self.frame_min = frame_min
        self.frame_max = frame_max
        self.whereToAugment = whereToAugment

        self.normFbankBeforeAugmentation = self.cfg.normFbankBeforeAugmentation
        self.updateFbankWithSilence = self.cfg.updateFbankWithSilence

        self.audioDict = None
        if self.cfg.audioDict:
            with open(self.cfg.audioDict, 'rb') as f:
                self.audioDict = pickle.load(f)
                logger.info('Load audioDict')

        logger.info('audioPath_prefix: {}'.format(self.audioPath_prefix))
        logger.info(self.__repr__())

    def get_tgt_lens_and_check_oov(self):
        if self.tgt_texts is None:
            return [0 for _ in range(self.n_samples)]
        tgt_lens = []
        n_tokens, n_oov_tokens = 0, 0
        for i in range(self.n_samples):
            tokenized = self.get_tokenized_tgt_text(i).split(" ")
            oov_tokens = [
                t
                for t in tokenized
                if self.tgt_dict.index(t) == self.tgt_dict.unk_index
            ]
            n_tokens += len(tokenized)
            n_oov_tokens += len(oov_tokens)
            tgt_lens.append(len(tokenized))
        logger.info(f"'{self.split}' has {n_oov_tokens / n_tokens * 100:.2f}% OOV")
        return tgt_lens

    def __repr__(self):
        return (
            self.__class__.__name__
            + f'(split="{self.split}", n_samples={self.n_samples}, '
            f"prepend_tgt_lang_tag={self.cfg.prepend_tgt_lang_tag}, "
            f"audioDictPath={self.cfg.audioDict}, "
            f"normFbankBeforeAugmentation={self.normFbankBeforeAugmentation}, "
            f"updateFbankWithSilence={self.updateFbankWithSilence}, "
            f"shuffle={self.shuffle}, "
            f"transforms={self.feature_transforms})"
        )

    @classmethod
    def is_lang_tag(cls, token):
        pattern = cls.LANG_TAG_TEMPLATE.replace("{}", "(.*)")
        return re.match(pattern, token)

    def check_tgt_lang_tag(self):
        if self.cfg.prepend_tgt_lang_tag:
            assert self.tgt_langs is not None and self.tgt_dict is not None
            tgt_lang_tags = [
                self.LANG_TAG_TEMPLATE.format(t) for t in set(self.tgt_langs)
            ]
            assert all(t in self.tgt_dict for t in tgt_lang_tags)

    @classmethod
    def tokenize(cls, tokenizer, text: str):
        return text if tokenizer is None else tokenizer.encode(text)

    #def get_tokenized_tgt_text(self, index: int):
    #    text = self.tokenize(self.pre_tokenizer, self.tgt_texts[index])
    #    text = self.tokenize(self.bpe_tokenizer, text)
    #    return text

    def get_tokenized_tgt_text(self, index: int, tgt_text: str = None):
        if tgt_text is None:
            tgt_text = self.tgt_texts[index]

        text = self.tokenize(self.pre_tokenizer, tgt_text)
        text = self.tokenize(self.bpe_tokenizer, text)
        return text

    @classmethod
    def get_lang_tag_idx(cls, lang: str, dictionary: Dictionary):
        lang_tag_idx = dictionary.index(cls.LANG_TAG_TEMPLATE.format(lang))
        assert lang_tag_idx != dictionary.unk()
        return lang_tag_idx

    def _getArrayFromZipFile(self, zipPath, offset, file_size):
        with open(zipPath, 'rb') as f:
            f.seek(offset)
            data = f.read(file_size)

        dataByte = io.BytesIO(data)
        fbankFromZip = np.load(dataByte)

        return fbankFromZip

    def _updateFbank_phrase(self, source, frame_min, whereToAugment):
        """
        Concatenate the original fbank segment and the sampled fbank segment
        """

        frame_min = [int(f) for f in frame_min.split('-')]

        if self.normFbankBeforeAugmentation is True:
            source = self._utteranceCMVN(source)

        audioPath = list(whereToAugment.values())[0]
        zipPath, offset, filesize, frameLoc = audioPath.split(':')
        zipPath = os.path.join(self.audioPath_prefix, zipPath)

        fbank = self._getArrayFromZipFile(
            zipPath, int(offset), int(filesize))

        if self.normFbankBeforeAugmentation is True:
            fbank = self._utteranceCMVN(fbank)

        idxMin, idxMax = frameLoc.split('-')
        fbank_suffix = fbank[int(idxMin):(int(idxMax)+1)].copy()
        
        loc = list(whereToAugment.keys())[0]
        source_prefix = source[:frame_min[loc],:].copy()

        return np.vstack((source_prefix, fbank_suffix))

    def _updateFbankWithSilence(self, source, frame_min, frame_max, srcInfo):
        frame_min = [int(f) for f in frame_min.split('-')]
        frame_max = [int(f) for f in frame_max.split('-')]

        if self.normFbankBeforeAugmentation is True:
            source = self._utteranceCMVN(source)

        collect_fbank_min_pos, collect_fbank_max_pos = [], []
        collect_sampled_fbanks = []

        for pos, fbank in srcInfo.items():
            collect_fbank_min_pos.append(frame_min[pos])
            collect_fbank_max_pos.append(frame_max[pos])
            collect_sampled_fbanks.append(fbank)

        #Update: First insert the remained fbanks
        collect_fbank_max_pos.insert(0, 0)
        collect_fbank_min_pos.append(source.shape[0])
        assert len(collect_fbank_max_pos) == len(collect_fbank_min_pos)

        collect_sampled_fbanks.append(np.array([])) #maintain the same length
        assert len(collect_fbank_max_pos) == len(collect_sampled_fbanks)

        fbank_updated = []
        for idx in range(len(collect_fbank_max_pos)):
            maxIdx = collect_fbank_max_pos[idx]
            minIdx = collect_fbank_min_pos[idx]

            fbank_updated.append(source[maxIdx:(minIdx), :])

            #Handle the augmented fbanks
            sampledFbank = collect_sampled_fbanks[idx]
            if isinstance(sampledFbank, list):
                assert len(sampledFbank) == 0
                continue
            else:
                assert isinstance(sampledFbank, np.ndarray)
                if collect_sampled_fbanks[idx].shape[0] == 0:
                    #for the maintain the same length stuff
                    continue

                fbank_updated.append(collect_sampled_fbanks[idx])

        return np.concatenate(fbank_updated)

    def _utteranceCMVN(self, spectrogram):
        """
        Channelwise normalization. To test the effect of normalization before
        or after augmentation on the final model performance.
        
        Copied from utterance_cmvn.py

        Return:
         spectrogram (np.array): channel-wise normalized filterbanks
        """
        mean = spectrogram.mean(axis=0)
        square_sums = (spectrogram ** 2).sum(axis=0)

        spectrogram = np.subtract(spectrogram, mean)
        
        var = square_sums / spectrogram.shape[0] - mean ** 2
        std = np.sqrt(np.maximum(var, 1e-10))

        spectrogram = np.divide(spectrogram, std)

        return spectrogram

    def __getitem__(self, index: int) -> SpeechToTextDatasetItem:
        source = get_features_or_waveform(
            self.audio_paths[index],
            need_waveform=self.cfg.use_audio_input,
            use_sample_rate=self.cfg.use_sample_rate,
        )
        if self.feature_transforms is not None:
            assert not self.cfg.use_audio_input

            whToAug = eval(self.whereToAugment[index])

            if (len(whToAug) != 0):
                source = self._updateFbank_phrase(
                    source=source, frame_min=self.frame_min[index],
                    whereToAugment=whToAug)

                if self.normFbankBeforeAugmentation is False:
                    source = self._utteranceCMVN(source)

            else:
                source = self._utteranceCMVN(source)

            ## Note: have to remove utteranceCMVN in the config.
            source = self.feature_transforms(source)
        source = torch.from_numpy(source).float()

        target = None
        if self.tgt_texts is not None:
            #tokenized = self.get_tokenized_tgt_text(index)
            tokenized = self.get_tokenized_tgt_text(
                index, tgt_text=self.tgt_texts[index])

            target = self.tgt_dict.encode_line(
                tokenized, add_if_not_exist=False, append_eos=True
            ).long()
            if self.cfg.prepend_tgt_lang_tag:
                lang_tag_idx = self.get_lang_tag_idx(
                    self.tgt_langs[index], self.tgt_dict
                )
                target = torch.cat((torch.LongTensor([lang_tag_idx]), target), 0)

        return SpeechToTextDatasetItem(index=index, source=source, target=target)

    def __len__(self):
        return self.n_samples

    def collater(
        self, samples: List[SpeechToTextDatasetItem], return_order: bool = False
    ) -> Dict:
        if len(samples) == 0:
            return {}
        indices = torch.tensor([x.index for x in samples], dtype=torch.long)
        frames = _collate_frames([x.source for x in samples], self.cfg.use_audio_input)
        # sort samples by descending number of frames
        n_frames = torch.tensor([x.source.size()[0] for x in samples], dtype=torch.long)
        n_frames, order = n_frames.sort(descending=True)
        indices = indices.index_select(0, order)
        frames = frames.index_select(0, order)

        target, target_lengths = None, None
        prev_output_tokens = None
        ntokens = None
        if self.tgt_texts is not None:
            target = fairseq_data_utils.collate_tokens(
                [x.target for x in samples],
                self.tgt_dict.pad(),
                self.tgt_dict.eos(),
                left_pad=False,
                move_eos_to_beginning=False,
            )
            target = target.index_select(0, order)
            target_lengths = torch.tensor(
                [x.target.size()[0] for x in samples], dtype=torch.long
            ).index_select(0, order)
            prev_output_tokens = fairseq_data_utils.collate_tokens(
                [x.target for x in samples],
                self.tgt_dict.pad(),
                self.tgt_dict.eos(),
                left_pad=False,
                move_eos_to_beginning=True,
            )
            prev_output_tokens = prev_output_tokens.index_select(0, order)
            ntokens = sum(x.target.size()[0] for x in samples)

        net_input = {
            "src_tokens": frames,
            "src_lengths": n_frames,
            "prev_output_tokens": prev_output_tokens,
        }
        out = {
            "id": indices,
            "net_input": net_input,
            "target": target,
            "target_lengths": target_lengths,
            "ntokens": ntokens,
            "nsentences": len(samples),
        }
        if return_order:
            out["order"] = order
        return out

    def num_tokens(self, index):
        return self.n_frames[index]

    def size(self, index):
        return self.n_frames[index], self.tgt_lens[index]

    @property
    def sizes(self):
        return np.array(self.n_frames)

    @property
    def can_reuse_epoch_itr_across_epochs(self):
        return True

    def ordered_indices(self):
        if self.shuffle:
            order = [np.random.permutation(len(self))]
        else:
            order = [np.arange(len(self))]
        # first by descending order of # of frames then by original/random order
        order.append([-n for n in self.n_frames])
        return np.lexsort(order)
        

    def prefetch(self, indices):
        raise False


class StrASTDatasetCreator(object):
    # mandatory columns
    KEY_ID, KEY_AUDIO, KEY_N_FRAMES = "id", "audio", "n_frames"
    KEY_TGT_TEXT = "tgt_text"
    # optional columns
    KEY_SPEAKER, KEY_SRC_TEXT = "speaker", "src_text"
    KEY_SRC_LANG, KEY_TGT_LANG = "src_lang", "tgt_lang"
    # default values
    DEFAULT_SPEAKER = DEFAULT_SRC_TEXT = DEFAULT_LANG = ""
    # columns for aligned info
    KEY_FRAME_MIN, KEY_FRAME_MAX = 'frame_min', 'frame_max'
    KEY_WHERETOAUGMENT = 'whereToAugment'

    @classmethod
    def _from_list(
        cls,
        split_name: str,
        is_train_split,
        samples: List[Dict],
        cfg: S2TDataConfig,
        tgt_dict,
        pre_tokenizer,
        bpe_tokenizer,
    ) -> StrASTDataset:
        audio_root = Path(cfg.audio_root)
        ids = [s[cls.KEY_ID] for s in samples]
        audio_paths = [(audio_root / s[cls.KEY_AUDIO]).as_posix() for s in samples]
        n_frames = [int(s[cls.KEY_N_FRAMES]) for s in samples]
        tgt_texts = [s[cls.KEY_TGT_TEXT] for s in samples]
        src_texts = [s.get(cls.KEY_SRC_TEXT, cls.DEFAULT_SRC_TEXT) for s in samples]
        speakers = [s.get(cls.KEY_SPEAKER, cls.DEFAULT_SPEAKER) for s in samples]
        src_langs = [s.get(cls.KEY_SRC_LANG, cls.DEFAULT_LANG) for s in samples]
        tgt_langs = [s.get(cls.KEY_TGT_LANG, cls.DEFAULT_LANG) for s in samples]

        frame_min = [s[cls.KEY_FRAME_MIN] for s in samples]
        frame_max = [s[cls.KEY_FRAME_MAX] for s in samples]
        whereToAugment = [s[cls.KEY_WHERETOAUGMENT] for s in samples]

        return StrASTDataset(
            split_name,
            is_train_split,
            cfg,
            audio_paths,
            n_frames,
            src_texts=src_texts,
            tgt_texts=tgt_texts,
            speakers=speakers,
            src_langs=src_langs,
            tgt_langs=tgt_langs,
            ids=ids,
            tgt_dict=tgt_dict,
            pre_tokenizer=pre_tokenizer,
            bpe_tokenizer=bpe_tokenizer,
            frame_min=frame_min,
            frame_max=frame_max,
            whereToAugment=whereToAugment,
        )

    @classmethod
    def get_size_ratios(
        cls, datasets: List[StrASTDataset], alpha: float = 1.0
    ) -> List[float]:
        """Size ratios for temperature-based sampling
        (https://arxiv.org/abs/1907.05019)"""

        id_to_lp, lp_to_sz = {}, defaultdict(int)
        for ds in datasets:
            lang_pairs = {f"{s}->{t}" for s, t in zip(ds.src_langs, ds.tgt_langs)}
            assert len(lang_pairs) == 1
            lang_pair = list(lang_pairs)[0]
            id_to_lp[ds.split] = lang_pair
            lp_to_sz[lang_pair] += sum(ds.n_frames)

        sz_sum = sum(v for v in lp_to_sz.values())
        lp_to_prob = {k: v / sz_sum for k, v in lp_to_sz.items()}
        lp_to_tgt_prob = {k: v ** alpha for k, v in lp_to_prob.items()}
        prob_sum = sum(v for v in lp_to_tgt_prob.values())
        lp_to_tgt_prob = {k: v / prob_sum for k, v in lp_to_tgt_prob.items()}
        lp_to_sz_ratio = {
            k: (lp_to_tgt_prob[k] * sz_sum) / v for k, v in lp_to_sz.items()
        }
        size_ratio = [lp_to_sz_ratio[id_to_lp[ds.split]] for ds in datasets]

        p_formatted = {
            k: f"{lp_to_prob[k]:.3f}->{lp_to_tgt_prob[k]:.3f}" for k in lp_to_sz
        }
        logger.info(f"sampling probability balancing: {p_formatted}")
        sr_formatted = {ds.split: f"{r:.3f}" for ds, r in zip(datasets, size_ratio)}
        logger.info(f"balanced sampling size ratio: {sr_formatted}")
        return size_ratio

    @classmethod
    def _load_samples_from_tsv(cls, root: str, split: str):
        tsv_path = Path(root) / f"{split}.tsv"
        if not tsv_path.is_file():
            raise FileNotFoundError(f"Dataset not found: {tsv_path}")
        with open(tsv_path) as f:
            reader = csv.DictReader(
                f,
                delimiter="\t",
                quotechar=None,
                doublequote=False,
                lineterminator="\n",
                quoting=csv.QUOTE_NONE,
            )
            samples = [dict(e) for e in reader]
        if len(samples) == 0:
            raise ValueError(f"Empty manifest: {tsv_path}")
        return samples

    @classmethod
    def _from_tsv(
        cls,
        root: str,
        cfg: S2TDataConfig,
        split: str,
        tgt_dict,
        is_train_split: bool,
        pre_tokenizer,
        bpe_tokenizer,
    ) -> StrASTDataset:
        samples = cls._load_samples_from_tsv(root, split)
        return cls._from_list(
            split, is_train_split, samples, cfg, tgt_dict, pre_tokenizer, bpe_tokenizer
        )

    @classmethod
    def from_tsv(
        cls,
        root: str,
        cfg: S2TDataConfig,
        splits: str,
        tgt_dict,
        pre_tokenizer,
        bpe_tokenizer,
        is_train_split: bool,
        epoch: int,
        seed: int,
    ) -> StrASTDataset:
        datasets = [
            cls._from_tsv(
                root, cfg, split, tgt_dict, is_train_split, pre_tokenizer, bpe_tokenizer
            )
            for split in splits.split(",")
        ]

        if is_train_split and len(datasets) > 1 and cfg.sampling_alpha != 1.0:
            # temperature-based sampling
            size_ratios = cls.get_size_ratios(datasets, alpha=cfg.sampling_alpha)
            datasets = [
                ResamplingDataset(
                    d, size_ratio=r, seed=seed, epoch=epoch, replace=(r >= 1.0)
                )
                for r, d in zip(size_ratios, datasets)
            ]

        return ConcatDataset(datasets) if len(datasets) > 1 else datasets[0]
