# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import numpy as np
from torch.utils.data import Dataset
import torch
import random
import cv2
from ..utils.image_processor import ImageProcessor, load_fixed_mask
from ..utils.audio import melspectrogram
from decord import AudioReader, VideoReader, cpu


class UNetDataset(Dataset):
    def __init__(self, train_data_dir: str, config):
        if config.data.train_fileslist != "":
            with open(config.data.train_fileslist) as file:
                self.video_paths = [line.rstrip() for line in file]
        elif train_data_dir != "":
            self.video_paths = []
            for file in os.listdir(train_data_dir):
                if file.endswith(".mp4"):
                    self.video_paths.append(os.path.join(train_data_dir, file))
        else:
            raise ValueError("data_dir and fileslist cannot be both empty")

        self.resolution = config.data.resolution
        self.num_frames = config.data.num_frames

        if self.num_frames == 16:
            self.mel_window_length = 52
        elif self.num_frames == 5:
            self.mel_window_length = 16
        else:
            raise NotImplementedError("Only support 16 and 5 frames now")

        self.audio_sample_rate = config.data.audio_sample_rate
        self.video_fps = config.data.video_fps
        self.mask = config.data.mask
        self.mask_image = load_fixed_mask(self.resolution)
        self.load_audio_data = config.model.add_audio_layer and config.run.use_syncnet
        self.audio_mel_cache_dir = config.data.audio_mel_cache_dir
        os.makedirs(self.audio_mel_cache_dir, exist_ok=True)

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

    def read_audio(self, video_path: str):
        ar = AudioReader(video_path, ctx=cpu(self.worker_id), sample_rate=self.audio_sample_rate)
        original_mel = melspectrogram(ar[:].asnumpy().squeeze(0))
        return torch.from_numpy(original_mel)

    def crop_audio_window(self, original_mel, start_index):
        start_idx = int(80.0 * (start_index / float(self.video_fps)))
        end_idx = start_idx + self.mel_window_length
        return original_mel[:, start_idx:end_idx].unsqueeze(0)

    def get_frames(self, video_reader: VideoReader):
        total_num_frames = len(video_reader)

        start_idx = random.randint(self.num_frames // 2, total_num_frames - self.num_frames - self.num_frames // 2)
        frames_index = np.arange(start_idx, start_idx + self.num_frames, dtype=int)

        while True:
            wrong_start_idx = random.randint(0, total_num_frames - self.num_frames)
            if wrong_start_idx > start_idx - self.num_frames and wrong_start_idx < start_idx + self.num_frames:
                continue
            wrong_frames_index = np.arange(wrong_start_idx, wrong_start_idx + self.num_frames, dtype=int)
            break

        frames = video_reader.get_batch(frames_index).asnumpy()
        wrong_frames = video_reader.get_batch(wrong_frames_index).asnumpy()

        return frames, wrong_frames, start_idx

    def worker_init_fn(self, worker_id):
        # Initialize the face mesh object in each worker process,
        # because the face mesh object cannot be called in subprocesses
        self.worker_id = worker_id
        setattr(
            self,
            f"image_processor_{worker_id}",
            ImageProcessor(self.resolution, self.mask, mask_image=self.mask_image),
        )

    def __getitem__(self, idx):
        image_processor = getattr(self, f"image_processor_{self.worker_id}")
        while True:
            try:
                idx = random.randint(0, len(self) - 1)

                # Get video file path
                video_path = self.video_paths[idx]

                vr = VideoReader(video_path, ctx=cpu(self.worker_id))

                if len(vr) < 3 * self.num_frames:
                    continue

                continuous_frames, ref_frames, start_idx = self.get_frames(vr)

                if self.load_audio_data:
                    mel_cache_path = os.path.join(
                        self.audio_mel_cache_dir, os.path.basename(video_path).replace(".mp4", "_mel.pt")
                    )

                    if os.path.isfile(mel_cache_path):
                        try:
                            original_mel = torch.load(mel_cache_path)
                        except Exception as e:
                            print(f"{type(e).__name__} - {e} - {mel_cache_path}")
                            os.remove(mel_cache_path)
                            original_mel = self.read_audio(video_path)
                            torch.save(original_mel, mel_cache_path)
                    else:
                        original_mel = self.read_audio(video_path)
                        torch.save(original_mel, mel_cache_path)

                    mel = self.crop_audio_window(original_mel, start_idx)

                    if mel.shape[-1] != self.mel_window_length:
                        continue
                else:
                    mel = []

                gt, masked_gt, mask = image_processor.prepare_masks_and_masked_images(
                    continuous_frames, affine_transform=False
                )

                if self.mask == "fix_mask":
                    ref, _, _ = image_processor.prepare_masks_and_masked_images(ref_frames, affine_transform=False)
                else:
                    ref = image_processor.process_images(ref_frames)
                vr.seek(0)  # avoid memory leak
                break

            except Exception as e:  # Handle the exception of face not detcted
                print(f"{type(e).__name__} - {e} - {video_path}")
                if "vr" in locals():
                    vr.seek(0)  # avoid memory leak

        sample = dict(
            gt=gt,
            masked_gt=masked_gt,
            ref=ref,
            mel=mel,
            mask=mask,
            video_path=video_path,
            start_idx=start_idx,
        )

        return sample