"""Data-loading functions in order to create a Dataset and DataLoaders.
Transforms for loaders are in transforms.py
"""

import json
import os
from pathlib import Path

import numpy as np
import torch
import yaml
from imageio import imread
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

from climategan.transforms import get_transforms
from climategan.tutils import get_normalized_depth_t
from climategan.utils import env_to_path, is_image_file

classes_dict = {
    "s": {  # unity
        0: [0, 0, 255, 255],  # Water
        1: [55, 55, 55, 255],  # Ground
        2: [0, 255, 255, 255],  # Building
        3: [255, 212, 0, 255],  # Traffic items
        4: [0, 255, 0, 255],  # Vegetation
        5: [255, 97, 0, 255],  # Terrain
        6: [255, 0, 0, 255],  # Car
        7: [60, 180, 60, 255],  # Trees
        8: [255, 0, 255, 255],  # Person
        9: [0, 0, 0, 255],  # Sky
        10: [255, 255, 255, 255],  # Default
    },
    "r": {  # deeplab v2
        0: [0, 0, 255, 255],  # Water
        1: [55, 55, 55, 255],  # Ground
        2: [0, 255, 255, 255],  # Building
        3: [255, 212, 0, 255],  # Traffic items
        4: [0, 255, 0, 255],  # Vegetation
        5: [255, 97, 0, 255],  # Terrain
        6: [255, 0, 0, 255],  # Car
        7: [60, 180, 60, 255],  # Trees
        8: [220, 20, 60, 255],  # Person
        9: [8, 19, 49, 255],  # Sky
        10: [0, 80, 100, 255],  # Default
    },
    "kitti": {
        0: [210, 0, 200],  # Terrain
        1: [90, 200, 255],  # Sky
        2: [0, 199, 0],  # Tree
        3: [90, 240, 0],  # Vegetation
        4: [140, 140, 140],  # Building
        5: [100, 60, 100],  # Road
        6: [250, 100, 255],  # GuardRail
        7: [255, 255, 0],  # TrafficSign
        8: [200, 200, 0],  # TrafficLight
        9: [255, 130, 0],  # Pole
        10: [80, 80, 80],  # Misc
        11: [160, 60, 60],  # Truck
        12: [255, 127, 80],  # Car
        13: [0, 139, 139],  # Van
        14: [0, 0, 0],  # Undefined
    },
    "flood": {
        0: [255, 0, 0],  # Cannot flood
        1: [0, 0, 255],  # Must flood
        2: [0, 0, 0],  # May flood
    },
}

kitti_mapping = {
    0: 5,  # Terrain -> Terrain
    1: 9,  # Sky -> Sky
    2: 7,  # Tree -> Trees
    3: 4,  # Vegetation ->  Vegetation
    4: 2,  # Building -> Building
    5: 1,  # Road -> Ground
    6: 3,  # GuardRail -> Traffic items
    7: 3,  # TrafficSign ->  Traffic items
    8: 3,  # TrafficLight -> Traffic items
    9: 3,  # Pole -> Traffic items
    10: 10,  # Misc -> default
    11: 6,  # Truck -> Car
    12: 6,  # Car -> Car
    13: 6,  # Van -> Car
    14: 10,  # Undefined -> Default
}


def encode_exact_segmap(seg, classes_dict, default_value=14):
    """
    When the mapping (rgb -> label) is known to be exact (no approximative rgb values)
    maps rgb image to segmap labels

    Args:
        seg (np.ndarray): H x W x 3 RGB image
        classes_dict (dict): Mapping {class: rgb value}
        default_value (int, optional): Value for unknown label. Defaults to 14.

    Returns:
        np.ndarray: Segmap as labels, not RGB
    """
    out = np.ones((seg.shape[0], seg.shape[1])) * default_value
    for cindex, cvalue in classes_dict.items():
        out[np.where((seg == cvalue).all(-1))] = cindex
    return out


def merge_labels(labels, mapping, default_value=14):
    """
    Maps labels from a source domain to labels of a target domain,
    typically kitti -> climategan

    Args:
        labels (np.ndarray): input segmap labels
        mapping (dict): source_label -> target_label
        default_value (int, optional): Unknown label. Defaults to 14.

    Returns:
        np.ndarray: Adapted labels
    """
    out = np.ones_like(labels) * default_value
    for source, target in mapping.items():
        out[labels == source] = target
    return out


def process_kitti_seg(path, kitti_classes, merge_map, default=14):
    """
    Processes a path to produce a 1 x 1 x H x W torch segmap

    %timeit process_kitti_seg(path, classes_dict, mapping, default=14)
    326 ms ± 118 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

    Args:
        path (str | pathlib.Path): Segmap RBG path
        kitti_classes (dict): Kitti map label -> rgb
        merge_map (dict): map kitti_label -> climategan_label
        default (int, optional): Unknown kitti label. Defaults to 14.

    Returns:
        torch.Tensor: 1 x 1 x H x W torch segmap
    """
    seg = imread(path)
    labels = encode_exact_segmap(seg, kitti_classes, default_value=default)
    merged = merge_labels(labels, merge_map, default_value=default)
    return torch.tensor(merged).unsqueeze(0).unsqueeze(0)


def decode_segmap_merged_labels(tensor, domain, is_target, nc=11):
    """Creates a label colormap for classes used in Unity segmentation benchmark.
    Arguments:
        tensor -- segmented image of size (1) x (nc) x (H) x (W)
            if prediction, or size (1) x (1) x (H) x (W) if target
    Returns:
        RGB tensor of size (1) x (3) x (H) x (W)
    #"""

    if is_target:  # Target is size 1 x 1 x H x W
        idx = tensor.squeeze(0).squeeze(0)
    else:  # Prediction is size 1 x nc x H x W
        idx = torch.argmax(tensor.squeeze(0), dim=0)

    indexer = torch.tensor(list(classes_dict[domain].values()))[:, :3]
    return indexer[idx.long()].permute(2, 0, 1).to(torch.float32).unsqueeze(0)


def decode_segmap_cityscapes_labels(image, nc=19):
    """Creates a label colormap used in CITYSCAPES segmentation benchmark.
    Arguments:
        image {array} -- segmented image
        (array of image size containing class at each pixel)
    Returns:
        array of size 3*nc -- A colormap for visualizing segmentation results.
    """
    colormap = np.zeros((19, 3), dtype=np.uint8)
    colormap[0] = [128, 64, 128]
    colormap[1] = [244, 35, 232]
    colormap[2] = [70, 70, 70]
    colormap[3] = [102, 102, 156]
    colormap[4] = [190, 153, 153]
    colormap[5] = [153, 153, 153]
    colormap[6] = [250, 170, 30]
    colormap[7] = [220, 220, 0]
    colormap[8] = [107, 142, 35]
    colormap[9] = [152, 251, 152]
    colormap[10] = [70, 130, 180]
    colormap[11] = [220, 20, 60]
    colormap[12] = [255, 0, 0]
    colormap[13] = [0, 0, 142]
    colormap[14] = [0, 0, 70]
    colormap[15] = [0, 60, 100]
    colormap[16] = [0, 80, 100]
    colormap[17] = [0, 0, 230]
    colormap[18] = [119, 11, 32]

    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for col in range(nc):
        idx = image == col
        r[idx] = colormap[col, 0]
        g[idx] = colormap[col, 1]
        b[idx] = colormap[col, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb


def find_closest_class(pixel, dict_classes):
    """Takes a pixel as input and finds the closest known pixel value corresponding
    to a class in dict_classes

    Arguments:
        pixel -- tuple pixel (R,G,B,A)
    Returns:
        tuple pixel (R,G,B,A) corresponding to a key (a class) in dict_classes
    """
    min_dist = float("inf")
    closest_pixel = None
    for pixel_value in dict_classes.keys():
        dist = np.sqrt(np.sum(np.square(np.subtract(pixel, pixel_value))))
        if dist < min_dist:
            min_dist = dist
            closest_pixel = pixel_value
    return closest_pixel


def encode_segmap(arr, domain):
    """Change a segmentation RGBA array to a segmentation array
                            with each pixel being the index of the class
    Arguments:
        numpy array -- segmented image of size (H) x (W) x (4 RGBA values)
    Returns:
        numpy array of size (1) x (H) x (W) with each pixel being the index of the class
    """
    new_arr = np.zeros((1, arr.shape[0], arr.shape[1]))
    dict_classes = {
        tuple(rgba_value): class_id
        for (class_id, rgba_value) in classes_dict[domain].items()
    }
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            pixel_rgba = tuple(arr[i, j, :])
            if pixel_rgba in dict_classes.keys():
                new_arr[0, i, j] = dict_classes[pixel_rgba]
            else:
                pixel_rgba_closest = find_closest_class(pixel_rgba, dict_classes)
                new_arr[0, i, j] = dict_classes[pixel_rgba_closest]
    return new_arr


def encode_mask_label(arr, domain):
    """Change a segmentation RGBA array to a segmentation array
                            with each pixel being the index of the class
    Arguments:
        numpy array -- segmented image of size (H) x (W) x (3 RGB values)
    Returns:
        numpy array of size (1) x (H) x (W) with each pixel being the index of the class
    """
    diff = np.zeros((len(classes_dict[domain].keys()), arr.shape[0], arr.shape[1]))
    for cindex, cvalue in classes_dict[domain].items():
        diff[cindex, :, :] = np.sqrt(
            np.sum(
                np.square(arr - np.tile(cvalue, (arr.shape[0], arr.shape[1], 1))),
                axis=2,
            )
        )
    return np.expand_dims(np.argmin(diff, axis=0), axis=0)


def transform_segmap_image_to_tensor(path, domain):
    """
    Transforms a segmentation image to a tensor of size (1) x (1) x (H) x (W)
    with each pixel being the index of the class
    """
    arr = np.array(Image.open(path).convert("RGBA"))
    arr = encode_segmap(arr, domain)
    arr = torch.from_numpy(arr).float()
    arr = arr.unsqueeze(0)
    return arr


def save_segmap_tensors(path_to_json, path_to_dir, domain):
    """
    Loads the segmentation images mentionned in a json file, transforms them to
    tensors and save the tensors in the wanted directory

    Args:
        path_to_json: complete path to the json file where to find the original data
        path_to_dir: path to the directory where to save the tensors as tensor_name.pt
        domain: domain of the images ("r" or "s")

    e.g:
        save_tensors(
            "/network/tmp1/ccai/data/climategan/seg/train_s.json",
            "/network/tmp1/ccai/data/munit_dataset/simdata/Unity11K_res640/Seg_tensors/",
            "s",
        )
    """
    ims_list = None
    if path_to_json:
        path_to_json = Path(path_to_json).resolve()
        with open(path_to_json, "r") as f:
            ims_list = yaml.safe_load(f)

    assert ims_list is not None

    for im_dict in ims_list:
        for task_name, path in im_dict.items():
            if task_name == "s":
                file_name = os.path.splitext(path)[0]  # remove extension
                file_name = file_name.rsplit("/", 1)[-1]  # keep only the file_name
                tensor = transform_segmap_image_to_tensor(path, domain)
                torch.save(tensor, path_to_dir + file_name + ".pt")


def pil_image_loader(path, task):
    if Path(path).suffix == ".npy":
        arr = np.load(path).astype(np.uint8)
    elif is_image_file(path):
        # arr = imread(path).astype(np.uint8)
        arr = np.array(Image.open(path).convert("RGB"))
    else:
        raise ValueError("Unknown data type {}".format(path))

    # Convert from RGBA to RGB for images
    if len(arr.shape) == 3 and arr.shape[-1] == 4:
        arr = arr[:, :, 0:3]

    if task == "m":
        arr[arr != 0] = 1
        # Make sure mask is single-channel
        if len(arr.shape) >= 3:
            arr = arr[:, :, 0]

    # assert len(arr.shape) == 3, (path, task, arr.shape)

    return Image.fromarray(arr)


def tensor_loader(path, task, domain, opts):
    """load data as tensors
    Args:
        path (str): path to data
        task (str)
        domain (str)
    Returns:
        [Tensor]: 1 x C x H x W
    """
    if task == "s":
        if domain == "kitti":
            return process_kitti_seg(
                path, classes_dict["kitti"], kitti_mapping, default=14
            )
        return torch.load(path)
    elif task == "d":
        if Path(path).suffix == ".npy":
            arr = np.load(path)
        else:
            arr = imread(path)  # .astype(np.uint8) /!\ kitti is np.uint16
        tensor = torch.from_numpy(arr.astype(np.float32))
        tensor = get_normalized_depth_t(
            tensor,
            domain,
            normalize="d" in opts.train.pseudo.tasks,
            log=opts.gen.d.classify.enable,
        )
        tensor = tensor.unsqueeze(0)
        return tensor

    elif Path(path).suffix == ".npy":
        arr = np.load(path).astype(np.float32)
    elif is_image_file(path):
        arr = imread(path).astype(np.float32)
    else:
        raise ValueError("Unknown data type {}".format(path))

    # Convert from RGBA to RGB for images
    if len(arr.shape) == 3 and arr.shape[-1] == 4:
        arr = arr[:, :, 0:3]

    if task == "x":
        arr -= arr.min()
        arr /= arr.max()
        arr = np.moveaxis(arr, 2, 0)
    elif task == "s":
        arr = np.moveaxis(arr, 2, 0)
    elif task == "m":
        if arr.max() > 127:
            arr = (arr > 127).astype(arr.dtype)
        # Make sure mask is single-channel
        if len(arr.shape) >= 3:
            arr = arr[:, :, 0]
        arr = np.expand_dims(arr, 0)

    return torch.from_numpy(arr).unsqueeze(0)


class OmniListDataset(Dataset):
    def __init__(self, mode, domain, opts, transform=None):

        self.opts = opts
        self.domain = domain
        self.mode = mode
        self.tasks = set(opts.tasks)
        self.tasks.add("x")
        if "p" in self.tasks:
            self.tasks.add("m")

        file_list_path = Path(opts.data.files[mode][domain])
        if "/" not in str(file_list_path):
            file_list_path = Path(opts.data.files.base) / Path(
                opts.data.files[mode][domain]
            )

        if file_list_path.suffix == ".json":
            self.samples_paths = self.json_load(file_list_path)
        elif file_list_path.suffix in {".yaml", ".yml"}:
            self.samples_paths = self.yaml_load(file_list_path)
        else:
            raise ValueError("Unknown file list type in {}".format(file_list_path))

        if opts.data.max_samples and opts.data.max_samples != -1:
            assert isinstance(opts.data.max_samples, int)
            self.samples_paths = self.samples_paths[: opts.data.max_samples]

        self.filter_samples()
        if opts.data.check_samples:
            print(f"Checking samples ({mode}, {domain})")
            self.check_samples()
        self.file_list_path = str(file_list_path)
        self.transform = transform

    def filter_samples(self):
        """
        Filter out data which is not required for the model's tasks
        as defined in opts.tasks
        """
        self.samples_paths = [
            {k: v for k, v in s.items() if k in self.tasks} for s in self.samples_paths
        ]

    def __getitem__(self, i):
        """Return an item in the dataset with fields:
        {
            data: transform({
                domains: values
            }),
            paths: [{task: path}],
            domain: [domain],
            mode: [train|val]
        }
        Args:
            i (int): index of item to retrieve
        Returns:
            dict: dataset item where tensors of data are in item["data"] which is a dict
                  {task: tensor}
        """
        paths = self.samples_paths[i]

        # always apply transforms,
        # if no transform is specified, ToTensor and Normalize will be applied

        item = {
            "data": self.transform(
                {
                    task: tensor_loader(
                        env_to_path(path),
                        task,
                        self.domain,
                        self.opts,
                    )
                    for task, path in paths.items()
                }
            ),
            "paths": paths,
            "domain": self.domain if self.domain != "kitti" else "s",
            "mode": self.mode,
        }

        return item

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

    def json_load(self, file_path):
        with open(file_path, "r") as f:
            return json.load(f)

    def yaml_load(self, file_path):
        with open(file_path, "r") as f:
            return yaml.safe_load(f)

    def check_samples(self):
        """Checks that every file listed in samples_paths actually
        exist on the file-system
        """
        for s in self.samples_paths:
            for k, v in s.items():
                assert Path(v).exists(), f"{k} {v} does not exist"


def get_loader(mode, domain, opts):
    if (
        domain != "kitti"
        or not opts.train.kitti.pretrain
        or not opts.train.kitti.batch_size
    ):
        batch_size = opts.data.loaders.get("batch_size", 4)
    else:
        batch_size = opts.train.kitti.get("batch_size", 4)

    return DataLoader(
        OmniListDataset(
            mode,
            domain,
            opts,
            transform=transforms.Compose(get_transforms(opts, mode, domain)),
        ),
        batch_size=batch_size,
        shuffle=True,
        num_workers=opts.data.loaders.get("num_workers", 8),
        pin_memory=True,  # faster transfer to gpu
        drop_last=True,  # avoids batchnorm pbs if last batch has size 1
    )


def get_all_loaders(opts):
    loaders = {}
    for mode in ["train", "val"]:
        loaders[mode] = {}
        for domain in opts.domains:
            if mode in opts.data.files:
                if domain in opts.data.files[mode]:
                    loaders[mode][domain] = get_loader(mode, domain, opts)
    return loaders