import json
from glob import glob
from os.path import join
from dataset import AbstractDataset

SPLIT = ["train", "val", "test"]
LABEL_MAP = {"REAL": 0, "FAKE": 1}


class DFDC(AbstractDataset):
    """
    Deepfake Detection Challenge organized by Facebook
    """

    def __init__(self, cfg, seed=2022, transforms=None, transform=None, target_transform=None):
        # pre-check
        if cfg['split'] not in SPLIT:
            raise ValueError(f"split should be one of {SPLIT}, but found {cfg['split']}.")
        super(DFDC, self).__init__(cfg, seed, transforms, transform, target_transform)
        print(f"Loading data from 'DFDC' of split '{cfg['split']}'"
              f"\nPlease wait patiently...")
        self.categories = ['original', 'fake']
        self.root = cfg['root']
        self.num_real = 0
        self.num_fake = 0
        if self.split == "test":
            self.__load_test_data()
        elif self.split == "train":
            self.__load_train_data()
        assert len(self.images) == len(self.targets), "Length of images and targets not the same!"
        print(f"Data from 'DFDC' loaded.")
        print(f"Real: {self.num_real}, Fake: {self.num_fake}.")
        print(f"Dataset contains {len(self.images)} images\n")

    def __load_test_data(self):
        label_path = join(self.root, "test", "labels.csv")
        with open(label_path, encoding="utf-8") as file:
            content = file.readlines()
        for _ in content:
            if ".mp4" in _:
                key = _.split(".")[0]
                label = _.split(",")[1].strip()
                label = int(label)
                imgs = glob(join(self.root, "test", "images", key, "*.png"))
                num = len(imgs)
                self.images.extend(imgs)
                self.targets.extend([label] * num)
                if label == 0:
                    self.num_real += num
                elif label == 1:
                    self.num_fake += num

    def __load_train_data(self):
        train_folds = glob(join(self.root, "dfdc_train_part_*"))
        for fold in train_folds:
            fold_imgs = list()
            fold_tgts = list()
            metadata_path = join(fold, "metadata.json")
            try:
                with open(metadata_path, "r", encoding="utf-8") as file:
                    metadata = json.loads(file.readline())
                for k, v in metadata.items():
                    index = k.split(".")[0]
                    label = LABEL_MAP[v["label"]]
                    imgs = glob(join(fold, "images", index, "*.png"))
                    fold_imgs.extend(imgs)
                    fold_tgts.extend([label] * len(imgs))
                    if label == 0:
                        self.num_real += len(imgs)
                    elif label == 1:
                        self.num_fake += len(imgs)
                self.images.extend(fold_imgs)
                self.targets.extend(fold_tgts)
            except FileNotFoundError:
                continue


if __name__ == '__main__':
    import yaml

    config_path = "../config/dataset/dfdc.yml"
    with open(config_path) as config_file:
        config = yaml.load(config_file, Loader=yaml.FullLoader)
    config = config["train_cfg"]
    # config = config["test_cfg"]


    def run_dataset():
        dataset = DFDC(config)
        print(f"dataset: {len(dataset)}")
        for i, _ in enumerate(dataset):
            path, target = _
            print(f"path: {path}, target: {target}")
            if i >= 9:
                break


    def run_dataloader(display_samples=False):
        from torch.utils import data
        import matplotlib.pyplot as plt

        dataset = DFDC(config)
        dataloader = data.DataLoader(dataset, batch_size=8, shuffle=True)
        print(f"dataset: {len(dataset)}")
        for i, _ in enumerate(dataloader):
            path, targets = _
            image = dataloader.dataset.load_item(path)
            print(f"image: {image.shape}, target: {targets}")
            if display_samples:
                plt.figure()
                img = image[0].permute([1, 2, 0]).numpy()
                plt.imshow(img)
                # plt.savefig("./img_" + str(i) + ".png")
                plt.show()
            if i >= 9:
                break


    ###########################
    # run the functions below #
    ###########################

    # run_dataset()
    run_dataloader(False)