import numpy as np
import random
import torch
from pathlib import Path
from torch.utils import data as data

from basicsr.data.transforms import augment, paired_random_crop
from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
from basicsr.utils.flow_util import dequantize_flow
from basicsr.utils.registry import DATASET_REGISTRY


@DATASET_REGISTRY.register()
class REDSDataset(data.Dataset):
    """REDS dataset for training.

    The keys are generated from a meta info txt file.
    basicsr/data/meta_info/meta_info_REDS_GT.txt

    Each line contains:
    1. subfolder (clip) name; 2. frame number; 3. image shape, separated by
    a white space.
    Examples:
    000 100 (720,1280,3)
    001 100 (720,1280,3)
    ...

    Key examples: "000/00000000"
    GT (gt): Ground-Truth;
    LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.

    Args:
        opt (dict): Config for train dataset. It contains the following keys:
        dataroot_gt (str): Data root path for gt.
        dataroot_lq (str): Data root path for lq.
        dataroot_flow (str, optional): Data root path for flow.
        meta_info_file (str): Path for meta information file.
        val_partition (str): Validation partition types. 'REDS4' or 'official'.
        io_backend (dict): IO backend type and other kwarg.
        num_frame (int): Window size for input frames.
        gt_size (int): Cropped patched size for gt patches.
        interval_list (list): Interval list for temporal augmentation.
        random_reverse (bool): Random reverse input frames.
        use_hflip (bool): Use horizontal flips.
        use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
        scale (bool): Scale, which will be added automatically.
    """

    def __init__(self, opt):
        super(REDSDataset, self).__init__()
        self.opt = opt
        self.gt_root, self.lq_root = Path(opt["dataroot_gt"]), Path(opt["dataroot_lq"])
        self.flow_root = (
            Path(opt["dataroot_flow"]) if opt["dataroot_flow"] is not None else None
        )
        assert (
            opt["num_frame"] % 2 == 1
        ), f'num_frame should be odd number, but got {opt["num_frame"]}'
        self.num_frame = opt["num_frame"]
        self.num_half_frames = opt["num_frame"] // 2

        self.keys = []
        with open(opt["meta_info_file"], "r") as fin:
            for line in fin:
                folder, frame_num, _ = line.split(" ")
                self.keys.extend([f"{folder}/{i:08d}" for i in range(int(frame_num))])

        # remove the video clips used in validation
        if opt["val_partition"] == "REDS4":
            val_partition = ["000", "011", "015", "020"]
        elif opt["val_partition"] == "official":
            val_partition = [f"{v:03d}" for v in range(240, 270)]
        else:
            raise ValueError(
                f'Wrong validation partition {opt["val_partition"]}.'
                f"Supported ones are ['official', 'REDS4']."
            )
        self.keys = [v for v in self.keys if v.split("/")[0] not in val_partition]

        # file client (io backend)
        self.file_client = None
        self.io_backend_opt = opt["io_backend"]
        self.is_lmdb = False
        if self.io_backend_opt["type"] == "lmdb":
            self.is_lmdb = True
            if self.flow_root is not None:
                self.io_backend_opt["db_paths"] = [
                    self.lq_root,
                    self.gt_root,
                    self.flow_root,
                ]
                self.io_backend_opt["client_keys"] = ["lq", "gt", "flow"]
            else:
                self.io_backend_opt["db_paths"] = [self.lq_root, self.gt_root]
                self.io_backend_opt["client_keys"] = ["lq", "gt"]

        # temporal augmentation configs
        self.interval_list = opt["interval_list"]
        self.random_reverse = opt["random_reverse"]
        interval_str = ",".join(str(x) for x in opt["interval_list"])
        logger = get_root_logger()
        logger.info(
            f"Temporal augmentation interval list: [{interval_str}]; "
            f"random reverse is {self.random_reverse}."
        )

    def __getitem__(self, index):
        if self.file_client is None:
            self.file_client = FileClient(
                self.io_backend_opt.pop("type"), **self.io_backend_opt
            )

        scale = self.opt["scale"]
        gt_size = self.opt["gt_size"]
        key = self.keys[index]
        clip_name, frame_name = key.split("/")  # key example: 000/00000000
        center_frame_idx = int(frame_name)

        # determine the neighboring frames
        interval = random.choice(self.interval_list)

        # ensure not exceeding the borders
        start_frame_idx = center_frame_idx - self.num_half_frames * interval
        end_frame_idx = center_frame_idx + self.num_half_frames * interval
        # each clip has 100 frames starting from 0 to 99
        while (start_frame_idx < 0) or (end_frame_idx > 99):
            center_frame_idx = random.randint(0, 99)
            start_frame_idx = center_frame_idx - self.num_half_frames * interval
            end_frame_idx = center_frame_idx + self.num_half_frames * interval
        frame_name = f"{center_frame_idx:08d}"
        neighbor_list = list(range(start_frame_idx, end_frame_idx + 1, interval))
        # random reverse
        if self.random_reverse and random.random() < 0.5:
            neighbor_list.reverse()

        assert (
            len(neighbor_list) == self.num_frame
        ), f"Wrong length of neighbor list: {len(neighbor_list)}"

        # get the GT frame (as the center frame)
        if self.is_lmdb:
            img_gt_path = f"{clip_name}/{frame_name}"
        else:
            img_gt_path = self.gt_root / clip_name / f"{frame_name}.png"
        img_bytes = self.file_client.get(img_gt_path, "gt")
        img_gt = imfrombytes(img_bytes, float32=True)

        # get the neighboring LQ frames
        img_lqs = []
        for neighbor in neighbor_list:
            if self.is_lmdb:
                img_lq_path = f"{clip_name}/{neighbor:08d}"
            else:
                img_lq_path = self.lq_root / clip_name / f"{neighbor:08d}.png"
            img_bytes = self.file_client.get(img_lq_path, "lq")
            img_lq = imfrombytes(img_bytes, float32=True)
            img_lqs.append(img_lq)

        # get flows
        if self.flow_root is not None:
            img_flows = []
            # read previous flows
            for i in range(self.num_half_frames, 0, -1):
                if self.is_lmdb:
                    flow_path = f"{clip_name}/{frame_name}_p{i}"
                else:
                    flow_path = self.flow_root / clip_name / f"{frame_name}_p{i}.png"
                img_bytes = self.file_client.get(flow_path, "flow")
                cat_flow = imfrombytes(
                    img_bytes, flag="grayscale", float32=False
                )  # uint8, [0, 255]
                dx, dy = np.split(cat_flow, 2, axis=0)
                flow = dequantize_flow(
                    dx, dy, max_val=20, denorm=False
                )  # we use max_val 20 here.
                img_flows.append(flow)
            # read next flows
            for i in range(1, self.num_half_frames + 1):
                if self.is_lmdb:
                    flow_path = f"{clip_name}/{frame_name}_n{i}"
                else:
                    flow_path = self.flow_root / clip_name / f"{frame_name}_n{i}.png"
                img_bytes = self.file_client.get(flow_path, "flow")
                cat_flow = imfrombytes(
                    img_bytes, flag="grayscale", float32=False
                )  # uint8, [0, 255]
                dx, dy = np.split(cat_flow, 2, axis=0)
                flow = dequantize_flow(
                    dx, dy, max_val=20, denorm=False
                )  # we use max_val 20 here.
                img_flows.append(flow)

            # for random crop, here, img_flows and img_lqs have the same
            # spatial size
            img_lqs.extend(img_flows)

        # randomly crop
        img_gt, img_lqs = paired_random_crop(
            img_gt, img_lqs, gt_size, scale, img_gt_path
        )
        if self.flow_root is not None:
            img_lqs, img_flows = img_lqs[: self.num_frame], img_lqs[self.num_frame :]

        # augmentation - flip, rotate
        img_lqs.append(img_gt)
        if self.flow_root is not None:
            img_results, img_flows = augment(
                img_lqs, self.opt["use_hflip"], self.opt["use_rot"], img_flows
            )
        else:
            img_results = augment(img_lqs, self.opt["use_hflip"], self.opt["use_rot"])

        img_results = img2tensor(img_results)
        img_lqs = torch.stack(img_results[0:-1], dim=0)
        img_gt = img_results[-1]

        if self.flow_root is not None:
            img_flows = img2tensor(img_flows)
            # add the zero center flow
            img_flows.insert(self.num_half_frames, torch.zeros_like(img_flows[0]))
            img_flows = torch.stack(img_flows, dim=0)

        # img_lqs: (t, c, h, w)
        # img_flows: (t, 2, h, w)
        # img_gt: (c, h, w)
        # key: str
        if self.flow_root is not None:
            return {"lq": img_lqs, "flow": img_flows, "gt": img_gt, "key": key}
        else:
            return {"lq": img_lqs, "gt": img_gt, "key": key}

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


@DATASET_REGISTRY.register()
class REDSRecurrentDataset(data.Dataset):
    """REDS dataset for training recurrent networks.

    The keys are generated from a meta info txt file.
    basicsr/data/meta_info/meta_info_REDS_GT.txt

    Each line contains:
    1. subfolder (clip) name; 2. frame number; 3. image shape, separated by
    a white space.
    Examples:
    000 100 (720,1280,3)
    001 100 (720,1280,3)
    ...

    Key examples: "000/00000000"
    GT (gt): Ground-Truth;
    LQ (lq): Low-Quality, e.g., low-resolution/blurry/noisy/compressed frames.

    Args:
        opt (dict): Config for train dataset. It contains the following keys:
        dataroot_gt (str): Data root path for gt.
        dataroot_lq (str): Data root path for lq.
        dataroot_flow (str, optional): Data root path for flow.
        meta_info_file (str): Path for meta information file.
        val_partition (str): Validation partition types. 'REDS4' or 'official'.
        io_backend (dict): IO backend type and other kwarg.
        num_frame (int): Window size for input frames.
        gt_size (int): Cropped patched size for gt patches.
        interval_list (list): Interval list for temporal augmentation.
        random_reverse (bool): Random reverse input frames.
        use_hflip (bool): Use horizontal flips.
        use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
        scale (bool): Scale, which will be added automatically.
    """

    def __init__(self, opt):
        super(REDSRecurrentDataset, self).__init__()
        self.opt = opt
        self.gt_root, self.lq_root = Path(opt["dataroot_gt"]), Path(opt["dataroot_lq"])
        self.num_frame = opt["num_frame"]

        self.keys = []
        with open(opt["meta_info_file"], "r") as fin:
            for line in fin:
                folder, frame_num, _ = line.split(" ")
                self.keys.extend([f"{folder}/{i:08d}" for i in range(int(frame_num))])

        # remove the video clips used in validation
        if opt["val_partition"] == "REDS4":
            val_partition = ["000", "011", "015", "020"]
        elif opt["val_partition"] == "official":
            val_partition = [f"{v:03d}" for v in range(240, 270)]
        else:
            raise ValueError(
                f'Wrong validation partition {opt["val_partition"]}.'
                f"Supported ones are ['official', 'REDS4']."
            )
        if opt["test_mode"]:
            self.keys = [v for v in self.keys if v.split("/")[0] in val_partition]
        else:
            self.keys = [v for v in self.keys if v.split("/")[0] not in val_partition]

        # file client (io backend)
        self.file_client = None
        self.io_backend_opt = opt["io_backend"]
        self.is_lmdb = False
        if self.io_backend_opt["type"] == "lmdb":
            self.is_lmdb = True
            if hasattr(self, "flow_root") and self.flow_root is not None:
                self.io_backend_opt["db_paths"] = [
                    self.lq_root,
                    self.gt_root,
                    self.flow_root,
                ]
                self.io_backend_opt["client_keys"] = ["lq", "gt", "flow"]
            else:
                self.io_backend_opt["db_paths"] = [self.lq_root, self.gt_root]
                self.io_backend_opt["client_keys"] = ["lq", "gt"]

        # temporal augmentation configs
        self.interval_list = opt.get("interval_list", [1])
        self.random_reverse = opt.get("random_reverse", False)
        interval_str = ",".join(str(x) for x in self.interval_list)
        logger = get_root_logger()
        logger.info(
            f"Temporal augmentation interval list: [{interval_str}]; "
            f"random reverse is {self.random_reverse}."
        )

    def __getitem__(self, index):
        if self.file_client is None:
            self.file_client = FileClient(
                self.io_backend_opt.pop("type"), **self.io_backend_opt
            )

        scale = self.opt["scale"]
        gt_size = self.opt["gt_size"]
        key = self.keys[index]
        clip_name, frame_name = key.split("/")  # key example: 000/00000000

        # determine the neighboring frames
        interval = random.choice(self.interval_list)

        # ensure not exceeding the borders
        start_frame_idx = int(frame_name)
        if start_frame_idx > 100 - self.num_frame * interval:
            start_frame_idx = random.randint(0, 100 - self.num_frame * interval)
        end_frame_idx = start_frame_idx + self.num_frame * interval

        neighbor_list = list(range(start_frame_idx, end_frame_idx, interval))

        # random reverse
        if self.random_reverse and random.random() < 0.5:
            neighbor_list.reverse()

        # get the neighboring LQ and GT frames
        img_lqs = []
        img_gts = []
        for neighbor in neighbor_list:
            if self.is_lmdb:
                img_lq_path = f"{clip_name}/{neighbor:08d}"
                img_gt_path = f"{clip_name}/{neighbor:08d}"
            else:
                img_lq_path = self.lq_root / clip_name / f"{neighbor:08d}.png"
                img_gt_path = self.gt_root / clip_name / f"{neighbor:08d}.png"

            # get LQ
            img_bytes = self.file_client.get(img_lq_path, "lq")
            img_lq = imfrombytes(img_bytes, float32=True)
            img_lqs.append(img_lq)

            # get GT
            img_bytes = self.file_client.get(img_gt_path, "gt")
            img_gt = imfrombytes(img_bytes, float32=True)
            img_gts.append(img_gt)

        # randomly crop
        img_gts, img_lqs = paired_random_crop(
            img_gts, img_lqs, gt_size, scale, img_gt_path
        )

        # augmentation - flip, rotate
        img_lqs.extend(img_gts)
        img_results = augment(img_lqs, self.opt["use_hflip"], self.opt["use_rot"])

        img_results = img2tensor(img_results)
        img_gts = torch.stack(img_results[len(img_lqs) // 2 :], dim=0)
        img_lqs = torch.stack(img_results[: len(img_lqs) // 2], dim=0)

        # img_lqs: (t, c, h, w)
        # img_gts: (t, c, h, w)
        # key: str
        return {"lq": img_lqs, "gt": img_gts, "key": key}

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