# We will use this file to create a dataloader for the real and fake dataset
import os
import json
import torch
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import numpy as np
import pandas as pd
import cv2

import cv2
import numpy as np
import matplotlib.pyplot as plt
import pywt

class Extracted_Frames_Dataset(Dataset):
    def __init__(self, root_dir, split = "train", transform = None, extend = 'None', multi_modal = "dct"):
        """
        Args:   
        returns:
            """
        AssertionError(split in ["train", "val", "test"]), "Split must be one of (train, val, test)"
        self.multi_modal = multi_modal
        self.root_dir = root_dir
        self.split = split
        self.transform = transform
        if extend == 'faceswap':
            self.dataset = pd.read_csv(os.path.join(root_dir, f"faceswap_extended_{self.split}.csv"))
        elif extend == 'fsgan':
            self.dataset = pd.read_csv(os.path.join(root_dir, f"fsgan_extended_{self.split}.csv"))
        else:
            self.dataset = pd.read_csv(os.path.join(root_dir, f"{self.split}.csv"))
                

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

    def __getitem__(self, idx):
        sample_input = self.get_sample_input(idx)
        return sample_input
    

    def get_sample_input(self, idx):
        rgb_image = self.get_rgb_image(idx)
        label = self.get_label(idx) 
        if self.multi_modal == "dct":
            dct_image = self.get_dct_image(idx)
            sample_input = {"rgb_image": rgb_image, "dct_image": dct_image, "label": label}

        # dct_image = self.get_dct_image(idx)
        elif self.multi_modal == "fft":
            fft_image = self.get_fft_image(idx)
            sample_input = {"rgb_image": rgb_image, "dct_image": fft_image, "label": label}
        elif self.multi_modal == "hh":
            hh_image = self.get_hh_image(idx)
            sample_input = {"rgb_image": rgb_image, "dct_image": hh_image, "label": label}
        else:
            AssertionError("multi_modal must be one of (dct:discrete cosine transform, fft: fast forier transform, hh)")

        return sample_input

    
    def get_fft_image(self, idx):
        gray_image_path = self.dataset.iloc[idx, 0]
        gray_image = cv2.imread(gray_image_path, cv2.IMREAD_GRAYSCALE)
        fft_image = self.compute_fft(gray_image)
        if self.transform:
            fft_image = self.transform(fft_image)
        
        return fft_image

    
    def compute_fft(self, image):
        f = np.fft.fft2(image)
        fshift = np.fft.fftshift(f)
        magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1)  # Add 1 to avoid log(0)
        return magnitude_spectrum


    def get_hh_image(self, idx):
        gray_image_path = self.dataset.iloc[idx, 0]
        gray_image = cv2.imread(gray_image_path, cv2.IMREAD_GRAYSCALE)
        hh_image = self.compute_hh(gray_image)
        if self.transform:
            hh_image = self.transform(hh_image)
        return hh_image
    
    def compute_hh(self, image):
        coeffs2 = pywt.dwt2(image, 'haar')
        LL, (LH, HL, HH) = coeffs2
        return HH
        
    def get_rgb_image(self, idx):
        rgb_image_path = self.dataset.iloc[idx, 0]
        rgb_image = Image.open(rgb_image_path)
        if self.transform:
            rgb_image = self.transform(rgb_image)
        return rgb_image
    
    def get_dct_image(self, idx):
        rgb_image_path = self.dataset.iloc[idx, 0]
        rgb_image = cv2.imread(rgb_image_path)
        dct_image = self.compute_dct_color(rgb_image)
        if self.transform:
            dct_image = self.transform(dct_image)
        
        return dct_image
    
    def get_label(self, idx):
        return self.dataset.iloc[idx, 1]
    

    def compute_dct_color(self, image):
        image_float = np.float32(image)
        dct_image = np.zeros_like(image_float)
        for i in range(3):  
            dct_image[:, :, i] = cv2.dct(image_float[:, :, i])
        return dct_image