import functools
import io
import json
import math
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # disable the tokenizer parallelism warning
import random
import re
import string
import subprocess
import sys
import yaml

import numpy as np

from collections import defaultdict
from copy import deepcopy
from dataclasses import dataclass
from functools import partial
from pydub import AudioSegment
from tqdm import tqdm

import torch
import torchvision
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset, get_worker_info
from torch.utils.data.distributed import DistributedSampler


from transformers import AutoTokenizer

import librosa
import soundfile as sf

EMOTION_MAP_DICT = {
    'amused':       'amused'      , 
    'anger':        'angry'       , 'angry':        'angry'       , 
    'anxious':      'anxious'     , 
    'apologetic':   'apologetic'  , 
    'assertive':    'assertive'   ,
    'calm':         'calm'        , 
    'concerned':    'concerned'   , 
    'contempt':     'contempt'    , 
    'disgust':      'disgusted'   , 'disgusted':    'disgusted'   , 
    'encouraging':  'encouraging' , 
    'excited':      'excited'     , 
    'fear':         'fearful'     , 'fearful':      'fearful'     , 
    'frustated':    'frustated'   ,
    'happy':        'happy'       , 'joy':          'happy'       , 
    'neutral':      'neutral'     , 
    'sad':          'sad'         , 'sadness':      'sad'         , 
    'sleepy':       'sleepy'      , 
    'surprise':     'surprised'   , 'surprised':    'surprised'   ,
    'pleasantly surprised': 'pleasantly surprised' ,
}


def int16_to_float32(x):
    return (x / 32767.0).astype(np.float32)


def float32_to_int16(x):
    x = np.clip(x, a_min=-1., a_max=1.)
    return (x * 32767.).astype(np.int16)


class DataCollator:
    def __init__(self, tokenizer, clap_config):

        self.tokenizer = tokenizer
        self.clap_config = clap_config
        self.max_num_window = clap_config["max_num_window"]

    def __call__(self, batch):

        filenames, audio_clips, audio_embed_masks, input_ids, attention_masks = zip(*batch)

        num_windows_all = [sum(audio_embed_mask) for audio_embed_mask in audio_embed_masks]
        max_window_batch = int(max(num_windows_all))

        if max_window_batch > self.max_num_window:
            max_window_batch = self.max_num_window

        padded_audio_clips = []
        padded_audio_embed_masks = []
        for audio_clip, audio_embed_mask in zip(audio_clips,audio_embed_masks):
            this_audio_clip_clips = [clip for clip in audio_clip]
            num_windows = len(this_audio_clip_clips)
            if num_windows < max_window_batch:
                for _ in range(max_window_batch - num_windows):
                    this_audio_clip_clips.append(torch.zeros_like(this_audio_clip_clips[-1]))
                audio_clip = torch.cat(this_audio_clip_clips)
                audio_embed_mask = torch.zeros(max_window_batch)
                audio_embed_mask[:num_windows] = 1
            elif num_windows < max_window_batch:
                audio_clip = this_audio_clip_clips[:max_window_batch]
                audio_clip = torch.cat(this_audio_clip_clips)
                audio_embed_mask = audio_embed_mask[:max_window_batch]
            else:
                audio_clip = torch.cat(this_audio_clip_clips)
                
            padded_audio_clips.append(audio_clip)
            padded_audio_embed_masks.append(audio_embed_mask)

        audio_clips = torch.cat([x.unsqueeze(0) for x in padded_audio_clips], dim=0)
        audio_embed_mask = torch.cat([x.unsqueeze(0) for x in padded_audio_embed_masks], dim=0)

        max_length = max([ids.shape[1] for ids in input_ids])

        padded_input_ids = []
        padded_attention_masks = []
        for ids, mask in zip(input_ids, attention_masks):
            if ids.shape[1] < max_length:
                padded_input_ids.append(
                    torch.cat([ids, torch.LongTensor([self.tokenizer.pad_token_id] * (max_length - ids.shape[1])).unsqueeze(0)], dim=1)
                )
                padded_attention_masks.append(
                    torch.cat([mask, torch.LongTensor([0] * (max_length - mask.shape[1])).unsqueeze(0)], dim=1)
                )
            else:
                padded_input_ids.append(ids)
                padded_attention_masks.append(mask)
        
        padded_input_ids = torch.cat(padded_input_ids, dim=0)
        padded_attention_masks = torch.cat(padded_attention_masks, dim=0).bool()
        
        out_dict = dict(
            filenames=filenames,
            audio_clips=audio_clips,
            audio_embed_mask=audio_embed_mask,
            input_ids=padded_input_ids,
            attention_mask=padded_attention_masks
        )
        return out_dict


class AudioTextData(torch.utils.data.Dataset):
    def __init__(
        self,
        dataset_file_root: str,
        data_root: str,
        clap_config: dict,
        dataset_blending_global_weight: float,
        dataset_blending_config: dict,
        dataset_blending_output: str,
        tokenizer,
        max_tokens: int,
        split: str = 'train',
        valid_dataset_config: dict = {},
        valid_dataset_name: str = '',
        epoch: int = 0,
        force_reblend: bool = False,
        sr = 16000,
        **kwargs
    ):
        self.dataset_file_root = dataset_file_root
        self.data_root = data_root
        self.clap_config = clap_config
        self.dataset_blending_global_weight = dataset_blending_global_weight
        self.dataset_blending_config = dataset_blending_config
        self.sr = sr
        
        self.split = split
        self.epoch = epoch
        self.force_reblend = force_reblend

        assert self.split in ['train', 'val', 'test']

        if self.split == 'train':
            self.data = self.blend_dataset(dataset_blending_config, dataset_blending_output)

        elif self.split in ['val', 'test']:
            self.valid_data = self.validation_dataset(valid_dataset_config, valid_dataset_name)
        
        self.tokenizer = tokenizer
        self.tokenizer.padding_side = "right"
        self.max_tokens = max_tokens

    @staticmethod
    def shuffle_dict_fixed_rand(dic, seed=0):
        print('randomly shuffling key-value pairs')
        
        local_random = np.random.default_rng(seed)
        original_keys = list(dic.keys())
        shuffled_keys = deepcopy(original_keys)
        local_random.shuffle(shuffled_keys)
        shuffling_mapping = {x: y for (x, y) in zip(original_keys, shuffled_keys)}

        shuffled_dic = {}
        for idx in original_keys:
            shuffled_idx = shuffling_mapping[idx]
            shuffled_dic[idx] = dic[shuffled_idx]
        return shuffled_dic

    @staticmethod
    def is_broken_file(audiopath):
        BROKEN_FILES = [
            "/lustre/fsw/portfolios/adlr/users/zkong/datasets/FMA/fma_large/023/023431.mp3",
            "/lustre/fsw/portfolios/adlr/users/zkong/datasets/FMA/fma_large/033/033690.mp3",
            "/lustre/fsw/portfolios/adlr/users/zkong/datasets/FMA/fma_large/119/119217.mp3",
            "/lustre/fsw/portfolios/adlr/users/zkong/datasets/FMA/fma_large/119/119222.mp3",
            "/lustre/fsw/portfolios/adlr/users/zkong/datasets/FMA/fma_large/119/119219.mp3",
            "/lustre/fsw/portfolios/adlr/users/zkong/datasets/GTZAN/gtzan/data/genres/jazz/jazz.00054.wav"
        ]
        return audiopath in BROKEN_FILES

    def _read_dataset_file(self, dataset_file):
        print("reading", dataset_file)
        with open(dataset_file) as f:
            contents = f.read()
        contents = json.loads(contents)

        if contents['split_path'] is not None:
            abs_path = contents['split_path']

        """
        for normal data
        contents['data'] = {idx: {
                'name': rel_path/name, 
                'prompt': prompt, 
                'output': output, 
                [optional] 'audio_start': audio_start,
                'task': task,
            }}
        """

        if 'interleaved' not in dataset_file:
            for idx in contents["data"]:
                contents["data"][idx]['task'] = contents["flamingo_task"]
                contents["data"][idx]['name'] = os.path.join(
                    abs_path, contents["data"][idx]['name']
                )
            return contents
    
    def blend_dataset(self, dataset_blending_config, dataset_blending_output):
        if os.path.exists(dataset_blending_output) and not self.force_reblend:
            print("loading blended dataset file from:", dataset_blending_output)
            with open(dataset_blending_output) as f:
                contents = f.read()
            self_data = json.loads(contents)
        
        else:
            if not self.force_reblend:
                print("no blended dataset file found; reading all dataset files")
            else:
                print("force reblending dataset at epoch {}; reading all dataset files".format(self.epoch))

            all_data = {}
            for dataset_name in dataset_blending_config:
                dataset_file = os.path.join(self.dataset_file_root, '{}.json'.format(dataset_name))
                contents = self._read_dataset_file(dataset_file)
                contents['data'] = self.shuffle_dict_fixed_rand(
                    contents['data'], 
                    seed=sum(list(map(ord, dataset_name)))
                )

                weight_global = float(self.dataset_blending_global_weight)
                weight_dataset = float(dataset_blending_config[dataset_name]["weight"])
                weight = weight_global * weight_dataset

                all_data[dataset_name] = {
                    "contents": contents,
                    "weight": weight
                }

            self_data = {
                "dataset_path": self.data_root,
                "split_path": None,
                "total_num": 0,
                "data": {}  # {id: {'name': rel_path/name or [rel_path/names], 'prompt': prompt or [prompts], 'output': output or [outputs], 'task': task, 'interleaved': interleave_method}}
            }

            for dataset_name in all_data:
                print('blending {}'.format(dataset_name))

                contents = all_data[dataset_name]["contents"]
                shuffled_contents_data = contents['data']
                weight = all_data[dataset_name]["weight"]
                assert type(weight) == float and weight > 0.0

                dataset_total_num = contents['total_num']
                start_idx = int(self.epoch * dataset_total_num * weight)
                end_idx = int((self.epoch + 1) * dataset_total_num * weight)

                for idx in range(start_idx, end_idx):
                    if idx > 0 and idx % dataset_total_num == 0:
                        print('force shuffling at new epoch {} for dataset {}'.format(idx // dataset_total_num, dataset_name))
                        shuffled_contents_data = self.shuffle_dict_fixed_rand(
                            contents['data'], 
                            seed=sum(list(map(ord, '{}-epoch-{}'.format(dataset_name, idx // dataset_total_num))))
                        )

                    key = str(idx % dataset_total_num)
                    item = shuffled_contents_data[key]

                    found_broken = False
                    if type(item['name']) is str:
                        audiopath = item['name']
                        if self.is_broken_file(audiopath):
                            print('cannot read {}'.format(audiopath))
                            found_broken = True

                    if found_broken:
                        continue 
                    
                    self_data['data'][self_data['total_num']] = item
                    self_data['total_num'] += 1 

            if not self.force_reblend:
                print('writing blended dataset file to:', dataset_blending_output)
                with open(dataset_blending_output, 'w') as json_file:
                    json.dump(self_data, json_file)
            else:
                print('writing reblended dataset file to:', dataset_blending_output.replace('.json', '-reblended.json'))
                with open(dataset_blending_output.replace('.json', '-reblended.json'), 'w') as json_file:
                    json.dump(self_data, json_file)

        return self_data

    def get_num_windows(self, T, sr):
        clap_config = self.clap_config
        window_length  = int(float(clap_config["window_length"]) * sr)
        window_overlap = int(float(clap_config["window_overlap"]) * sr)
        max_num_window = int(clap_config["max_num_window"])

        num_windows = 1
        if T <= window_length:
            num_windows = 1
            full_length = window_length
        elif T >= (max_num_window * window_length - (max_num_window - 1) * window_overlap):
            num_windows = max_num_window
            full_length = (max_num_window * window_length - (max_num_window - 1) * window_overlap)
        else:
            num_windows = 1 + int(np.ceil((T - window_length) / float(window_length - window_overlap)))
            full_length = num_windows * window_length - (num_windows - 1) * window_overlap
        
        return num_windows, full_length

    def load_audio(self, file_path, target_sr=16000, duration=30.0, start=0.0):
        if file_path.endswith('.mp3'):
            audio = AudioSegment.from_file(file_path)
            if len(audio) > (start + duration) * 1000:
                audio = audio[start * 1000:(start + duration) * 1000]

            if audio.frame_rate != target_sr:
                audio = audio.set_frame_rate(target_sr)

            if audio.channels > 1:
                audio = audio.set_channels(1)
            
            data = np.array(audio.get_array_of_samples())
            if audio.sample_width == 2:
                data = data.astype(np.float32) / np.iinfo(np.int16).max
            elif audio.sample_width == 4:
                data = data.astype(np.float32) / np.iinfo(np.int32).max
            else:
                raise ValueError("Unsupported bit depth: {}".format(audio.sample_width))

        else:
            with sf.SoundFile(file_path) as audio:
                original_sr = audio.samplerate
                channels = audio.channels

                max_frames = int((start + duration) * original_sr)

                audio.seek(int(start * original_sr))
                frames_to_read = min(max_frames, len(audio))
                data = audio.read(frames_to_read)

                if data.max() > 1 or data.min() < -1:
                    data = data / max(abs(data.max()), abs(data.min()))
            
            if original_sr != target_sr:
                if channels == 1:
                    data = librosa.resample(data.flatten(), orig_sr=original_sr, target_sr=target_sr)
                else:
                    data = librosa.resample(data.T, orig_sr=original_sr, target_sr=target_sr)[0]
            else:
                if channels != 1:
                    data = data.T[0]
        
        if data.min() >= 0:
            data = 2 * data / abs(data.max()) - 1.0
        else:
            data = data / max(abs(data.max()), abs(data.min()))
        
        assert len(data.shape) == 1, data.shape
        return data

    def compute_sliding_window(self, audio_file, audio_start=0.0, audio="sound"):
        if type(audio_start) == str:
            audio_start = float(audio_start)

        if audio == "sound":
            encoder_config = self.clap_config
        else:
            raise NotImplementedError

        if encoder_config["method"] == 'nvclap-large':
            sr = 16000
        else:
            raise NotImplementedError

        window_length  = int(float(encoder_config["window_length"]) * sr)
        window_overlap = int(float(encoder_config["window_overlap"]) * sr)
        max_num_window = int(encoder_config["max_num_window"])
        duration = max_num_window * (encoder_config["window_length"] - encoder_config["window_overlap"]) + encoder_config["window_overlap"]

        audio_data = self.load_audio(os.path.join(self.data_root, audio_file), sr, duration, audio_start) # already cuts to max duration
        T = len(audio_data)
        num_windows, full_length = self.get_num_windows(T, sr)

        # pads to the nearest multiple of window_length
        if full_length > T:
            audio_data = np.append(audio_data, np.zeros(full_length - T))

        audio_data = audio_data.reshape(1, -1)
        audio_data_tensor = torch.from_numpy(int16_to_float32(float32_to_int16(audio_data))).float()

        audio_clips = []
        audio_embed_mask = torch.ones(num_windows)
        for i in range(num_windows):
            start = i * (window_length - window_overlap)
            audio_data_tensor_this = audio_data_tensor[:, start:start+window_length]
            audio_clips.append(audio_data_tensor_this)            
        
        return audio_clips, audio_embed_mask

    def validation_dataset(self, valid_dataset_config, valid_dataset_name):
        dataset_file = os.path.join(self.dataset_file_root, '{}.json'.format(valid_dataset_name))
        contents = self._read_dataset_file(dataset_file)

        contents['data'] = self.shuffle_dict_fixed_rand(
            contents['data'], 
            seed=sum(list(map(ord, valid_dataset_name)))
        )

        return contents

    def preprocess_string_for_eval(self, x):
        x = x.rstrip().lstrip()
        x = x.lower()
        return x

    def _actual_getitem(self, i):
        if self.split == 'train':
            try:
                item = self.data['data'][str(i)]
            except:
                item = self.data['data'][i]

            if type(item['name']) is str:
                audio_file = item['name']
                audio_start = 0 if 'audio_start' not in item else float(item['audio_start'])
            else:
                raise Exception(f"The item has a {type(item['name'])}. Only single path as a string is supported")

            # compute window for long audios
            audio_clips, audio_embed_mask = self.compute_sliding_window(audio_file, audio_start, audio="sound")
        
            # make the text prompt
            text_prompt = str(item['prompt']).lower()
            text_output = str(item['output']).lower()

            sample = f"<audio>{text_prompt.strip()}{self.tokenizer.sep_token}{text_output.strip()}<|endofchunk|>{self.tokenizer.eos_token}"

            text = self.tokenizer(
                sample,
                max_length=self.max_tokens,
                padding="longest",
                truncation="only_first",
                return_tensors="pt"
            )
        
        elif self.split in ['val', 'test']:
            try:
                item = self.valid_data['data'][str(i)]
            except:
                item = self.valid_data['data'][i]

            if type(item['name']) is str:
                audio_file = os.path.join(self.data_root, item['name'])
                audio_start = 0 if 'audio_start' not in item else float(item['audio_start'])
            else:
                raise Exception(f"The item has a {type(item['name'])}. Only single path as a string is supported")

            # compute window for long audios
            audio_clips, audio_embed_mask = self.compute_sliding_window(audio_file, audio_start, audio="sound")
        
            # make the text prompt
            text_prompt = self.preprocess_string_for_eval(str(item['prompt']).lower())
            text_output = self.preprocess_string_for_eval(str(item['output']).lower())

            sample = f"<audio>{text_prompt.strip()}{self.tokenizer.sep_token}{text_output.strip()}<|endofchunk|>{self.tokenizer.eos_token}"

            text = self.tokenizer(
                sample,
                max_length=self.max_tokens,
                padding="longest",
                truncation="only_first",
                return_tensors="pt"
            )
            
        # audio_clips_clap, audio_embed_mask_clap, audio_clips_speech, audio_embed_mask_speech, audio_clips_music, audio_embed_mask_music,
        return (item['name'], audio_clips, audio_embed_mask, text["input_ids"], text["attention_mask"])

    def __getitem__(self, i):
        try: 
            return self._actual_getitem(i)
        except Exception as e:
            print('batch {} failed with reason {}'.format(i, e))
            try:
                return self._actual_getitem((i-42)%99)
            except:
                return self._actual_getitem((i-84)%99)

    def __len__(self):
        if self.split == 'train':
            return len(list(self.data['data'].keys()))

        elif self.split == 'val':
            return min(len(list(self.valid_data['data'].keys())), 64)

        elif self.split == 'test':
            return len(list(self.valid_data['data'].keys()))


@dataclass
class DataInfo:
    dataset: Dataset
    dataloader: DataLoader
    sampler: DistributedSampler = None

    def set_epoch(self, epoch):
        if self.sampler is not None and isinstance(self.sampler, DistributedSampler):
            self.sampler.set_epoch(epoch)


def get_audiotext_dataloader(data_config, clap_config, text_tokenizer, batch_size, split='train', epoch=0, force_reblend=False):
    assert split in ['train', 'val', 'test']

    data_collator = DataCollator(text_tokenizer, clap_config)
    dataloader_shuffle = False

    if split == 'train':
        trainset = AudioTextData(
            **data_config, 
            clap_config=clap_config,
            tokenizer=text_tokenizer, 
            split=split,
            epoch=epoch,
            force_reblend=force_reblend
        )
        sampler = DistributedSampler(trainset, shuffle=True)
        trainloader = DataLoader(
            trainset, 
            sampler=sampler, 
            batch_size=batch_size, 
            shuffle=dataloader_shuffle, 
            collate_fn=data_collator, 
            num_workers=data_config["num_workers"]
        )
        return DataInfo(dataset=trainset, dataloader=trainloader, sampler=sampler)
    
    elif split in ['val', 'test']:
        all_DataInfo = {}
        for valid_dataset_name in list(data_config["valid_dataset_config"].keys()):
            valid_dataset_name = valid_dataset_name.strip()
            validset = AudioTextData(
                **data_config, 
                clap_config=clap_config,
                tokenizer=text_tokenizer, 
                split=split, 
                valid_dataset_name=valid_dataset_name
            )
            if split == 'val':
                # distributed sampler
                all_DataInfo[valid_dataset_name] = DataInfo(
                    dataset=validset,
                    dataloader=DataLoader(
                        validset, 
                        sampler=DistributedSampler(validset, shuffle=False),
                        batch_size=batch_size, 
                        shuffle=dataloader_shuffle, 
                        collate_fn=data_collator, 
                        num_workers=data_config["num_workers"]
                ))
            else:
                # single GPU
                all_DataInfo[valid_dataset_name] = DataInfo(
                    dataset=validset,
                    dataloader=DataLoader(
                        validset, 
                        batch_size=batch_size, 
                        shuffle=dataloader_shuffle, 
                        collate_fn=data_collator, 
                        num_workers=data_config["num_workers"]
                ))

        return all_DataInfo
    

def main():
    import time
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, default='../configs/config.yaml', help='yaml config path')
    args = parser.parse_args()

    config = yaml.load(open(args.config), Loader=yaml.FullLoader)

    data_config = config['data_config']
    clap_config = config['clap_config']
    whisper_config = config["whisper_config"]
    mert_config = config["mert_config"]

    tokenizer_path = "facebook/opt-1.3b"
    cache_dir = '/lustre/fsw/portfolios/adlr/users/sreyang/.cache'
    text_tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_path,
        local_files_only=False,
        trust_remote_code=True,
        cache_dir=cache_dir,
    )
    text_tokenizer.add_special_tokens(
        {"additional_special_tokens": ["<audio>", "<|endofchunk|>"]}
    )
    if text_tokenizer.pad_token is None:
        text_tokenizer.add_special_tokens({"pad_token": "<|PAD_TOKEN|>"})
    if text_tokenizer.sep_token is None:
        text_tokenizer.add_special_tokens({"sep_token": "<SEP>"})

    trainset = AudioTextData(
        **data_config,
        clap_config=clap_config, tokenizer=text_tokenizer,
        epoch=66, force_reblend=True
    )

    data_collator = DataCollator(text_tokenizer)
    dataloader = DataLoader(trainset, batch_size=16, shuffle=True, collate_fn=data_collator, num_workers=4)

    for step, batch in enumerate(dataloader):
        filenames = batch["filenames"]
        audio_clips = batch["audio_clips"]
        audio_embed_mask = batch["audio_embed_mask"]
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]

        print(
            'batch {}:'.format(step+1), 
            audio_clips.shape, audio_embed_mask.shape, 
            input_ids.shape, attention_mask.shape
        )

        print('filenames', filenames)
        print('audio_embed_mask', audio_embed_mask)
        print('input_ids', input_ids)
        for input_id in input_ids:
            print('-' * 50)
            print(text_tokenizer.decode(input_id))
        print('attention_mask', attention_mask)

        if step == 20:
            break


if __name__ == "__main__":
    main()