import copy
import random
import torch
import os
import numpy as np
from tqdm import tqdm
from os.path import join as pjoin
from torch.utils.data import Dataset, DataLoader
from utils.utils import MotionNormalizer

class EvaluationDatasetInterHuman(Dataset):
    """
    Evaluation Dataset of InterHuman.
    Motions are generated by the trained model to later be compared with the ground truth.
    But this manner is too slow, should make it suitable for multiple batch sizes. 
    """
    def __init__(self, model, dataset, device, 
                 batch_size=32, num_workers=9):
        """
        Initialization of the dataset and generation of the motions.
            :param model: Model to generate the motions.
            :param dataset: Ground truth dataset.
            :param device: Device to run the model.
            :param mm_num_samples: Number of samples to generate for the MultiModality metric.
            :param mm_num_repeats: Number of repeats for each sample in the MultiModality metric.
        """

        # Configuration variables
        self.normalizer = MotionNormalizer()
        self.model = model.to(device)
        self.model.eval()
        dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True)
        self.max_length = dataset.max_length
        self.extended = dataset.extended

        # Indexes of the motions to generate
        idxs = list(range(len(dataset)))
        random.shuffle(idxs)

        # Data structures
        generated_motions = []

        with torch.no_grad():
            for i, data in enumerate(tqdm(dataloader)):

                # Get the data from the data loader
                if self.extended:
                    name, text, motion1, motion2, motion_lens, text_individual1, text_individual2 = data
                else:
                    name, text, motion1, motion2, motion_lens = data
                
                # If data into MM list, duplicate them mm_num_repeats times
                batch = {}

                batch["text"] = list(text)
                batch["motions"] = torch.stack([motion1, motion2], dim=2)
                if self.extended:
                    batch["text_individual1"] = list(text_individual1)
                    batch["text_individual2"] = list(text_individual2)

                batch["motion_lens"] = motion_lens

                # Predict the motions with the model to be evaluated
                batch = self.model.forward_test(batch)
                motions_output = batch["output"].reshape(batch["output"].shape[0], batch["output"].shape[1], 2, -1) # (B, T, 2, D)
                motions_output = self.normalizer.backward(motions_output.cpu().detach().numpy()) # (B, T, 2, D)

                # Padding the motions to the max_length
                B,T = motions_output.shape[0], motions_output.shape[1]
                if T < self.max_length:
                    padding_len = self.max_length - T
                    D = motions_output.shape[-1]
                    padding_zeros = np.zeros((B, padding_len, 2, D))
                    motions_output = np.concatenate((motions_output, padding_zeros), axis=1)
                assert motions_output.shape[1] == self.max_length

                # Save the generated motions
                if self.extended:
                    sub_dicts = []
                    for i in range(B):
                        sub_dict = {'motion1': motions_output[i, :,0],
                                    'motion2': motions_output[i, :,1],
                                    'motion_lens': motion_lens[i],
                                    'text': text[i],
                                    'text_individual1': text_individual1[0],
                                    'text_individual2': text_individual2[0]}
                        sub_dicts.append(sub_dict)
                    generated_motions.extend(sub_dicts)

                else:
                    sub_dicts = []
                    for i in range(B):
                        sub_dict = {'motion1': motions_output[i, :,0], # motions_output in [b, l, 2, d]
                                    'motion2': motions_output[i, :,1],
                                    'motion_lens': motion_lens[i],
                                    'text': text[i]}
                        sub_dicts.append(sub_dict)
                    generated_motions.extend(sub_dicts)

        self.generated_motions = generated_motions

    def __len__(self):
        """
        Get the length of the dataset.
        """
        return len(self.generated_motions)

    def __getitem__(self, item):
        """
        Get the item of the dataset.
            :param item: Index of the item.
        """
        data = self.generated_motions[item]

        # Return the data, if dataset extended also return individual descriptions
        if self.extended:
            motion1, motion2, motion_lens, text, text_individual1, text_individual2 = data['motion1'], data['motion2'], data['motion_lens'], data['text'], data['text_individual1'], data['text_individual2']
            return "generated", text, motion1, motion2, motion_lens, text_individual1, text_individual2
        else:
            motion1, motion2, motion_lens, text = data['motion1'], data['motion2'], data['motion_lens'], data['text']
            return "generated", text, motion1, motion2, motion_lens


class MMGeneratedDatasetInterHuman(Dataset):
    """
    Dataset for the MultiModality metric.
    """
    def __init__(self, model, dataset, device, mm_num_samples=100, mm_num_repeats=30):
        """
        Initialization of the dataset and generation of the motions.
            :param model: Model to generate the motions.
            :param dataset: Ground truth dataset.
            :param device: Device to run the model.
            :param mm_num_samples: Number of samples to generate for the MultiModality metric.
            :param mm_num_repeats: Number of repeats for each sample in the MultiModality metric.
        """

        # Configuration variables
        self.normalizer = MotionNormalizer()
        self.model = model.to(device)
        self.model.eval()
        dataloader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=True)
        self.max_length = dataset.max_length
        self.extended = dataset.extended

        # Indexes of the motions to generate
        idxs = list(range(len(dataset)))
        random.shuffle(idxs)
        mm_idxs = idxs[:mm_num_samples]

        # Data structures
        mm_generated_motions = []

        with torch.no_grad():
            for i, data in tqdm(enumerate(dataloader)):
                # Get the data from the data loader
                if self.extended:
                    name, text, motion1, motion2, motion_lens, text_individual1, text_individual2 = data
                else:
                    name, text, motion1, motion2, motion_lens = data
                
                # If data into MM list, duplicate them mm_num_repeats times
                batch = {}
                if i in mm_idxs:
                    batch["text"] = list(text) * mm_num_repeats
                    if self.extended:
                        batch["text_individual1"] = list(text_individual1) * mm_num_repeats
                        batch["text_individual2"] = list(text_individual2) * mm_num_repeats
                else:
                    continue
                    batch["text"] = list(text)
                    if self.extended:
                        batch["text_individual1"] = list(text_individual1)
                        batch["text_individual2"] = list(text_individual2)

                batch["motion_lens"] = torch.repeat_interleave(motion_lens, mm_num_repeats, dim=0)
                batch["motions"] = torch.stack([motion1, motion2], dim=2)
                batch["motions"] = torch.cat([batch["motions"]] * mm_num_repeats, dim=0)

                # Predict the motions with the model to be evaluated
                batch = self.model.forward_test(batch)
                motions_output = batch["output"].reshape(batch["output"].shape[0], batch["output"].shape[1], 2, -1)
                motions_output = self.normalizer.backward(motions_output.cpu().detach().numpy())

                # Padding the motions to the max_length
                B,T = motions_output.shape[0], motions_output.shape[1]
                if T < self.max_length:
                    padding_len = self.max_length - T
                    D = motions_output.shape[-1]
                    padding_zeros = np.zeros((B, padding_len, 2, D))
                    motions_output = np.concatenate((motions_output, padding_zeros), axis=1)
                assert motions_output.shape[1] == self.max_length

                # Save the generated motions
                if self.extended:
                    if i in mm_idxs:
                        mm_sub_dict = {'mm_motions': motions_output,
                                       'motion_lens': motion_lens[0],
                                       'text': text[0],
                                       'text_individual1': text_individual1[0],
                                       'text_individual2': text_individual2[0]}
                        mm_generated_motions.append(mm_sub_dict)
                else:
                    if i in mm_idxs:
                        mm_sub_dict = {'mm_motions': motions_output, # [sbmm, l, 2, d]
                                       'motion_lens': motion_lens[0],
                                       'text': text[0]}
                        mm_generated_motions.append(mm_sub_dict)
        self.mm_generated_motions = mm_generated_motions

        
    def __len__(self):
        """
        Get the length of the dataset.
        """
        return len(self.mm_generated_motions)
    
    def __getitem__(self, item):
        data = self.mm_generated_motions[item]
        mm_motions = data['mm_motions']
        motion_lens = data['motion_lens']
        mm_motions1 = mm_motions[:,:,0]
        mm_motions2 = mm_motions[:,:,1]
        text = data['text']
        motion_lens = np.array([motion_lens]*mm_motions1.shape[0])

        # If dataset extended also return individual descriptions
        if self.extended:
            text_individual1 = data['text_individual1']
            text_individual2 = data['text_individual2']
            return "mm_generated", text, mm_motions1, mm_motions2, motion_lens, text_individual1, text_individual2
        else:
            return "mm_generated", text, mm_motions1, mm_motions2, motion_lens






class EvaluationDatasetInterHumanOld(Dataset):
    """
    Evaluation Dataset of InterHuman.
    Motions are generated by the trained model to later be compared with the ground truth.
    """
    def __init__(self, model, dataset, device, mm_num_samples, mm_num_repeats):
        """
        Initialization of the dataset and generation of the motions.
            :param model: Model to generate the motions.
            :param dataset: Ground truth dataset.
            :param device: Device to run the model.
            :param mm_num_samples: Number of samples to generate for the MultiModality metric.
            :param mm_num_repeats: Number of repeats for each sample in the MultiModality metric.
        """

        # Configuration variables
        self.normalizer = MotionNormalizer()
        self.model = model.to(device)
        self.model.eval()
        dataloader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=True)
        self.max_length = dataset.max_length
        self.extended = dataset.extended

        # Indexes of the motions to generate
        idxs = list(range(len(dataset)))
        random.shuffle(idxs)
        mm_idxs = idxs[:mm_num_samples]

        # Data structures
        generated_motions = []
        mm_generated_motions = []

        with torch.no_grad():
            for i, data in tqdm(enumerate(dataloader)):

                # Get the data from the data loader
                if self.extended:
                    name, text, motion1, motion2, motion_lens, text_individual1, text_individual2 = data
                else:
                    name, text, motion1, motion2, motion_lens = data
                
                # If data into MM list, duplicate them mm_num_repeats times
                batch = {}
                if i in mm_idxs:
                    batch["text"] = list(text) * mm_num_repeats
                    if self.extended:
                        batch["text_individual1"] = list(text_individual1) * mm_num_repeats
                        batch["text_individual2"] = list(text_individual2) * mm_num_repeats
                else:
                    batch["text"] = list(text)
                    if self.extended:
                        batch["text_individual1"] = list(text_individual1)
                        batch["text_individual2"] = list(text_individual2)

                batch["motion_lens"] = motion_lens

                # Predict the motions with the model to be evaluated
                batch = self.model.forward_test(batch)
                motions_output = batch["output"].reshape(batch["output"].shape[0], batch["output"].shape[1], 2, -1)
                motions_output = self.normalizer.backward(motions_output.cpu().detach().numpy())

                # Padding the motions to the max_length
                B,T = motions_output.shape[0], motions_output.shape[1]
                if T < self.max_length:
                    padding_len = self.max_length - T
                    D = motions_output.shape[-1]
                    padding_zeros = np.zeros((B, padding_len, 2, D))
                    motions_output = np.concatenate((motions_output, padding_zeros), axis=1)
                assert motions_output.shape[1] == self.max_length

                # Save the generated motions
                if self.extended:
                    sub_dict = {'motion1': motions_output[0, :,0],
                                'motion2': motions_output[0, :,1],
                                'motion_lens': motion_lens[0],
                                'text': text[0],
                                'text_individual1': text_individual1[0],
                                'text_individual2': text_individual2[0]}
                    generated_motions.append(sub_dict)
                    if i in mm_idxs:
                        mm_sub_dict = {'mm_motions': motions_output,
                                       'motion_lens': motion_lens[0],
                                       'text': text[0],
                                       'text_individual1': text_individual1[0],
                                       'text_individual2': text_individual2[0]}
                        mm_generated_motions.append(mm_sub_dict)
                else:
                    sub_dict = {'motion1': motions_output[0, :,0],
                                'motion2': motions_output[0, :,1],
                                'motion_lens': motion_lens[0],
                                'text': text[0]}
                    generated_motions.append(sub_dict)
                    if i in mm_idxs:
                        mm_sub_dict = {'mm_motions': motions_output,
                                       'motion_lens': motion_lens[0],
                                       'text': text[0]}
                        mm_generated_motions.append(mm_sub_dict)


        self.generated_motions = generated_motions
        self.mm_generated_motions = mm_generated_motions

    def __len__(self):
        """
        Get the length of the dataset.
        """
        return len(self.generated_motions)

    def __getitem__(self, item):
        """
        Get the item of the dataset.
            :param item: Index of the item.
        """
        data = self.generated_motions[item]

        # Return the data, if dataset extended also return individual descriptions
        if self.extended:
            motion1, motion2, motion_lens, text, text_individual1, text_individual2 = data['motion1'], data['motion2'], data['motion_lens'], data['text'], data['text_individual1'], data['text_individual2']
            return "generated", text, motion1, motion2, motion_lens, text_individual1, text_individual2
        else:
            motion1, motion2, motion_lens, text = data['motion1'], data['motion2'], data['motion_lens'], data['text']
            return "generated", text, motion1, motion2, motion_lens


class MMGeneratedDatasetInterHumanOld(Dataset):
    """
    Dataset for the MultiModality metric.
    """
    def __init__(self, motion_dataset):
        """
        Initialization of the dataset.
            :param motion_dataset: EvaluationDataset of the generated motions.
        """
        self.dataset = motion_dataset.mm_generated_motions
        self.extended = motion_dataset.extended

    def __len__(self):
        """
        Get the length of the dataset.
        """
        return len(self.dataset)

    def __getitem__(self, item):
        """
        Get the item of the dataset.
            :param item: Index of the item.
        """
        data = self.dataset[item]
        mm_motions = data['mm_motions']
        motion_lens = data['motion_lens']
        mm_motions1 = mm_motions[:,:,0]
        mm_motions2 = mm_motions[:,:,1]
        text = data['text']
        motion_lens = np.array([motion_lens]*mm_motions1.shape[0])

        # If dataset extended also return individual descriptions
        if self.extended:
            text_individual1 = data['text_individual1']
            text_individual2 = data['text_individual2']
            return "mm_generated", text, mm_motions1, mm_motions2, motion_lens, text_individual1, text_individual2
        else:
            return "mm_generated", text, mm_motions1, mm_motions2, motion_lens



class EvaluationDatasetDualMDM(Dataset):
    """
    Evaluation Dataset of DualMDM.
    Motions are generated by the trained model to later be compared with the ground truth.
    The dataset is generated by combining interaction descriptions from InterHuman with individual descriptions from HumanML3D
    """
    def __init__(self, model, dataset, device, num_repeats):
        """
        Initialization of the dataset and generation of the motions.
            :param model: Model to generate the motions.
            :param dataset: Ground truth dataset.
            :param device: Device to run the model.
            :param num_repeats: Number of repeats for each sample.
        """
        self.normalizer = MotionNormalizer()
        self.model = model.to(device)
        self.model.eval()
        dataloader = DataLoader(dataset, batch_size=1, num_workers=0, shuffle=True)
        self.max_length = dataset.max_length

        # Paths to the individual descriptions from HumanML3D
        self.individual_text_path = "data/HumanML3D/texts"
        self.individual_text_files = os.listdir(self.individual_text_path)

        # Data structures
        generated_motions = []

        # Model parameters
        composition_weight_func = self.model.decoder.cfg_composition_weight_func
        composition_weight_value = self.model.decoder.cfg_composition_weight_value

        with torch.no_grad():
            for i, data in tqdm(enumerate(dataloader)):

                # Get the data from the gt InterHuman dataset
                name, text, motion1, motion2, motion_lens, text_individual1, text_individual2 = data

                batch = {}
                batch["motion_lens"] = motion_lens.repeat(num_repeats * 2)
                batch["text"] = list(text) * (num_repeats * 2)
                batch["text_individual1"] = list(text_individual1) * num_repeats
                batch["text_individual2"] = list(text_individual2) * num_repeats

                # Modify individual textual description to be from HumanML3D
                for j in range(num_repeats):
                    # Select 2 random files from HumanML3D for the individual descriptions
                    files = random.sample(self.individual_text_files, 2)

                    with open(pjoin(self.individual_text_path, files[0]), "r") as f:
                        hml3d_texts_individual1 = f.readlines()
                        hml3d_text_individual1 = random.choice(hml3d_texts_individual1)
                        hml3d_text_individual1 = hml3d_text_individual1.strip().split("#")[0]
                        batch["text_individual1"].append(hml3d_text_individual1)

                    with open(pjoin(self.individual_text_path, files[1]), "r") as f:
                        hml3d_texts_individual2 = f.readlines()
                        hml3d_text_individual2 = random.choice(hml3d_texts_individual2)
                        hml3d_text_individual2 = hml3d_text_individual2.strip().split("#")[0]
                        batch["text_individual2"].append(hml3d_text_individual2)


                # Generate interactions using the base interaction model (in2IN)
                batch_interaction = copy.deepcopy(batch)
                batch_interaction["motion_lens"] = batch_interaction["motion_lens"][:num_repeats]
                batch_interaction["text"] =  batch_interaction["text"][:num_repeats]
                batch_interaction["text_individual1"] = batch_interaction["text_individual1"][:num_repeats]
                batch_interaction["text_individual2"] = batch_interaction["text_individual2"][:num_repeats]
                self.model.decoder.cfg_composition_weight_func = "const"
                self.model.decoder.cfg_composition_weight_value = 0
                batch_interaction = self.model.forward_test(batch_interaction)
                motions_output_interaction = batch_interaction["output"].reshape(batch_interaction["output"].shape[0], batch_interaction["output"].shape[1], 2, -1)
                motions_output_interaction = self.normalizer.backward(motions_output_interaction.cpu().detach().numpy())

                # Generate interactions using the combined interaction model with the interaction (DualMDM)
                batch_individual = copy.deepcopy(batch)
                batch_individual["motion_lens"] = batch_individual["motion_lens"][num_repeats:]
                batch_individual["text"] =  batch_individual["text"][num_repeats:]
                batch_individual["text_individual1"] = batch_individual["text_individual1"][num_repeats:]
                batch_individual["text_individual2"] = batch_individual["text_individual2"][num_repeats:]
                self.model.decoder.cfg_composition_weight_func = composition_weight_func
                self.model.decoder.cfg_composition_weight_value = composition_weight_value
                batch_individual = self.model.forward_test(batch_individual)
                motions_output_individual = batch_individual["output"].reshape(batch_individual["output"].shape[0], batch_individual["output"].shape[1], 2, -1)
                motions_output_individual = self.normalizer.backward(motions_output_individual.cpu().detach().numpy())

                motions_output = np.concatenate((motions_output_interaction, motions_output_individual), axis=0)
                B,T = motions_output.shape[0], motions_output.shape[1]

                # Padding all the generated motions to the biggest length of the dataset
                if T < self.max_length:
                    padding_len = self.max_length - T
                    D = motions_output.shape[-1]
                    padding_zeros = np.zeros((B, padding_len, 2, D))
                    motions_output = np.concatenate((motions_output, padding_zeros), axis=1)
                assert motions_output.shape[1] == self.max_length

                # Save the generated motions
                sub_dict = {'generated_motions': motions_output,
                            'motion1': motion1,
                            'motion2': motion2,
                            'motion_lens': batch["motion_lens"],
                            'text': batch["text"],
                            'text_individual1': batch["text_individual1"],
                            'text_individual2': batch["text_individual2"]}
            
                generated_motions.append(sub_dict)


        self.generated_motions = generated_motions

    def __len__(self):
        """
        Get the length of the dataset.
        """
        return len(self.generated_motions)

    def __getitem__(self, item):
        """
        Get the item of the dataset.
            :param item: Index of the item.
        """
        data = self.generated_motions[item]

        motion_lens = data['motion_lens']
        text = data['text']
        text_individual1 = data['text_individual1']
        text_individual2 = data['text_individual2']

        generated_motions = data['generated_motions']
        generated_motions1 = generated_motions[:, :, 0, :]
        generated_motions2 = generated_motions[:, :, 1, :]
        motion1 = data['motion1']
        motion2 = data['motion2']

        return generated_motions1, generated_motions2, motion1, motion2, motion_lens, text, text_individual1, text_individual2