from torch.utils.data import Dataset
import os
import torch
import random
import numpy as np

class ReactDataset(Dataset):
    def __init__(
            self,
            root_dir,
            split="train",
            clip_length=750,
            random_crop=True
    ):
        self.clip_length = clip_length
        self.random_crop = random_crop
        base = os.path.join(root_dir, split)
        self.audio_dir = os.path.join(base, "audio-features")
        self.dmm_dir = os.path.join(base, "coefficients")
        self.emo_dir = os.path.join(base, "facial-attributes")

        speaker_emo_root = os.path.join(self.emo_dir, "speaker")
        self.items = []
        for root, _, files in os.walk(speaker_emo_root):
            session = os.path.basename(root)
            for f in files:
                if not f.lower().endswith(".npy"):
                    continue
                self.items.append(os.path.join(session, f))
        self.items.sort()
        assert len(self.items) > 0, "No .npy files found in " + speaker_emo_root

        mean_face_path = os.path.join("external_tools", "mean_face.npy")
        std_face_path = os.path.join("external_tools", "std_face.npy")

        self.mean_face = np.load(mean_face_path).astype(np.float32)
        self.std_face = np.load(std_face_path).astype(np.float32)

        self.mean_face = torch.from_numpy(self.mean_face).unsqueeze(0)
        self.std_face = torch.from_numpy(self.std_face).unsqueeze(0)

    def __len__(self):
        return len(self.items)
    
    def load_data(self, modality_dir, speaker_listener, current, target_dim=None):
        path = os.path.join(modality_dir, speaker_listener, current)
        t = np.load(path).astype(np.float32)
        t = torch.from_numpy(t).squeeze()

        if target_dim is not None:
            assert t.dim() == target_dim

        return t

    def __getitem__(self, index):
        current = self.items[index] #e.g "session0/123.npy"

        sp_audio = self.load_data(self.audio_dir, "speaker", current, target_dim=2)
        sp_dmm = self.load_data(self.dmm_dir, "speaker", current, target_dim=2)
        sp_emo = self.load_data(self.emo_dir, "speaker", current, target_dim=2)
        li_emo = self.load_data(self.emo_dir, "listener", current, target_dim=2)

        T = min(sp_audio.size(0), sp_dmm.size(0), sp_emo.size(0), li_emo.size(0))
        sp_audio = sp_audio[:T]
        sp_dmm = sp_dmm[:T]
        sp_emo = sp_emo[:T]
        li_emo = li_emo[:T]

        sp_dmm = (sp_dmm - self.mean_face) / self.std_face

        if T > self.clip_length:
            if self.random_crop:
                start = random.randint(0, T - self.clip_length)
            else:
                start = (T - self.clip_length) // 2
            sp_audio = sp_audio[start: start + self.clip_length]
            sp_dmm = sp_dmm[start: start + self.clip_length]
            sp_emo = sp_emo[start: start + self.clip_length]
            li_emo = li_emo[start: start + self.clip_length]
            motion_length = self.clip_length
        else:
            pad = self.clip_length - T
            sp_audio = torch.cat([sp_audio, torch.zeros(pad, 768)], dim=0)
            sp_dmm = torch.cat([sp_dmm, torch.zeros(pad, 58)], dim=0)
            sp_emo = torch.cat([sp_emo, torch.zeros(pad, 25)], dim=0)
            li_emo = torch.cat([li_emo, torch.zeros(pad, 25)], dim=0)
            motion_length = T
        
        assert sp_audio.size() == (750, 768)
        assert sp_dmm.size() == (750, 58)
        assert sp_emo.size() == (750, 25)
        assert li_emo.size() == (750, 25)

        return {
            "speaker_audio": sp_audio,
            "speaker_3dmm": sp_dmm,
            "speaker_emotion": sp_emo,
            "listener_emotion": li_emo,
            "motion_length": torch.tensor(motion_length, dtype=torch.long)
        }