import numpy as np
from pathlib import Path
from PIL import Image

import torch
from torch.utils.data import Dataset
import torchvision.transforms as T

from transformers import CLIPImageProcessor

import sys
sys.path.append("/path/to/FollowYourEmoji")
from media_pipe import FaceMeshDetector, FaceMeshAlign
from media_pipe.draw_util import FaceMeshVisualizer


def val_collate_fn(samples):
    return {
        'ref_frame': [sample['ref_frame'] for sample in samples],
        'clip_image': [sample['clip_image'] for sample in samples],
        'motions': [sample['motions'] for sample in samples],
        'file_name': [sample['file_name'] for sample in samples],
        'lmk_name': [sample['lmk_name'] for sample in samples],
    }


class ValDataset(Dataset):
    def __init__(self, input_path, lmk_path, resolution_w=512, resolution_h=512):
        
        print(f'Loading dataset from {input_path} and {lmk_path}')
        
        all_img_paths = self._get_path_files(Path(input_path), file_suffix=['.jpg', '.jpeg', '.png', '.webp'])
        all_lmk_paths = self._get_path_files(Path(lmk_path), file_suffix=['.npy'])
        
        print(f'Found {len(all_img_paths)} image files and {len(all_lmk_paths)} lmk files')
        print(f"ALL IMG PATH: {all_img_paths}")
        print(f"ALL LKM PATH: {all_lmk_paths}")
        self.all_paths = []
        for lmk_path in all_lmk_paths:
            for img_path in all_img_paths:
                self.all_paths.append((img_path, lmk_path))
        
        self.W = resolution_w
        self.H = resolution_h
        self.to_tensor = T.ToTensor()

        self.detector = FaceMeshDetector()
        self.aligner = FaceMeshAlign()

        self.clip_image_processor = CLIPImageProcessor()
        self.vis = FaceMeshVisualizer(forehead_edge=False, iris_edge=False, iris_point=True)

    def __len__(self):
        return len(self.all_paths)

    def _get_path_files(self, path, file_suffix):
        all_paths = []
        if path.is_file():
            if path.suffix.lower() in file_suffix:
                all_paths = [path]
            else:
                raise ValueError('Path is not valid image file.')
        elif path.is_dir():
            all_paths = sorted(
                [
                    f
                    for f in path.iterdir()
                    if f.is_file() and f.suffix.lower() in file_suffix
                ]
            )
            if len(all_paths) == 0:
                raise ValueError('Folder does not contain any images.')
        else:
            raise ValueError

        return all_paths

    def get_align_motion(self, ref_lmk, temp_lmks):
        motions = self.aligner(ref_lmk, temp_lmks)
        motions = [self.to_tensor(motion) for motion in motions]
        motions = torch.stack(motions).permute((1,0,2,3))
        return motions

    def __getitem__(self, index):
        img_path, lmk_path = self.all_paths[index]
        W, H = self.W, self.H

        image = Image.open(img_path).convert('RGB')

        # resize and center crop
        scale = min(W / image.size[0], H / image.size[1])
        ref_image = image.resize(
            (int(image.size[0] * scale), int(image.size[1] * scale)))
        w, h = ref_image.size[0], ref_image.size[1]
        ref_image = ref_image.crop((w//2-W//2, h//2-H//2, w//2+W//2, h//2+H//2))
        ref_image = np.array(ref_image)

        # reference image lmk
        ref_lmk_image, ref_lmk = self.detector(ref_image)

        # clip image
        clip_image = Image.fromarray(np.array(ref_image))
        clip_image = self.clip_image_processor(images=clip_image, return_tensors="pt").pixel_values[0]

        # reference image
        ref_image = self.to_tensor(ref_image).unsqueeze(1)
        ref_image = ref_image * 2.0 - 1.0

        # motion sequence
        temp_lmks = np.load(lmk_path, allow_pickle=True)
        # landmark align and draw motions
        if ref_lmk is not None:
            motions = self.get_align_motion(ref_lmk, temp_lmks)
        else:
            motions = [
                self.vis.draw_landmarks((H, W), lmk['lmks'].astype(np.float32), normed=True)
                for lmk in temp_lmks
            ]
            motions = [self.to_tensor(motion) for motion in motions]
            motions = torch.stack(motions).permute((1,0,2,3))

        example = dict()
        example["file_name"] = str(img_path.stem).split('/')[-1]
        example["lmk_name"] = str(lmk_path.stem).split('/')[-1]
        example["motions"] = motions # value in [0, 1]
        example["ref_frame"] = ref_image # value in [-1, 1]
        example["ref_lmk_image"] = ref_lmk_image
        example["clip_image"] = clip_image

        return example