#!/usr/bin/env python3

import os.path as osp

import numpy as np
from PIL import Image

import mmcv
from mmcv.utils import print_log

from .builder import DATASETS
from .custom import CustomJointDataset, OTFCustomJointDataset

from pyEdgeEval.datasets.cityscapes_attributes import (
    CITYSCAPES_inst_labelIds,
    CITYSCAPES_labelIds,
    CITYSCAPES_label2trainId,
)

CITYSCAPES_CLASSES = (
    "road",
    "sidewalk",
    "building",
    "wall",
    "fence",
    "pole",
    "traffic light",
    "traffic sign",
    "vegetation",
    "terrain",
    "sky",
    "person",
    "rider",
    "car",
    "truck",
    "bus",
    "train",
    "motorcycle",
    "bicycle",
)
CITYSCAPES_PALETTE = [
    [128, 64, 128],
    [244, 35, 232],
    [70, 70, 70],
    [102, 102, 156],
    [190, 153, 153],
    [153, 153, 153],
    [250, 170, 30],
    [220, 220, 0],
    [107, 142, 35],
    [152, 251, 152],
    [70, 130, 180],
    [220, 20, 60],
    [255, 0, 0],
    [0, 0, 142],
    [0, 0, 70],
    [0, 60, 100],
    [0, 80, 100],
    [0, 0, 230],
    [119, 11, 32],
]
ignore_ids = [2, 3]


@DATASETS.register_module()
class JointCityscapesDataset(CustomJointDataset):
    CLASSES = CITYSCAPES_CLASSES
    PALETTE = CITYSCAPES_PALETTE

    def __init__(
        self,
        img_suffix="_leftImg8bit.png",
        seg_map_suffix="_gtFine_labelTrainIds.png",
        edge_map_suffix="_gtProc_raw_edge.png",
        inst_sensitive=True,
        thin=False,
        test_mode=False,
        **kwargs,
    ):

        if test_mode:
            if inst_sensitive:
                if thin:
                    edge_map_suffix = "_gtProc_thin_isedge.png"
                else:
                    edge_map_suffix = "_gtProc_raw_isedge.png"
            else:
                if thin:
                    edge_map_suffix = "_gtProc_thin_edge.png"
                else:
                    edge_map_suffix = "_gtProc_raw_edge.png"
        else:
            if inst_sensitive:
                edge_map_suffix = "_gtProc_isedge.png"
            else:
                edge_map_suffix = "_gtProc_edge.png"

        super(JointCityscapesDataset, self).__init__(
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            edge_map_suffix=edge_map_suffix,
            inst_sensitive=inst_sensitive,
            **kwargs,
        )

    @staticmethod
    def _convert_to_label_id(result):
        """Convert trainId to id for cityscapes."""
        if isinstance(result, str):
            result = np.load(result)
        import cityscapesscripts.helpers.labels as CSLabels

        result_copy = result.copy()
        for trainId, label in CSLabels.trainId2label.items():
            result_copy[result == trainId] = label.id

        return result_copy

    def results2img(self, results, imgfile_prefix, to_label_id, indices=None):
        """Write the segmentation results to images.

        Args:
            results (list[ndarray]): Testing results of the
                dataset.
            imgfile_prefix (str): The filename prefix of the png files.
                If the prefix is "somepath/xxx",
                the png files will be named "somepath/xxx.png".
            to_label_id (bool): whether convert output to label_id for
                submission.
            indices (list[int], optional): Indices of input results,
                if not set, all the indices of the dataset will be used.
                Default: None.

        Returns:
            list[str: str]: result txt files which contains corresponding
            semantic segmentation images.
        """
        if indices is None:
            indices = list(range(len(self)))

        mmcv.mkdir_or_exist(imgfile_prefix)
        result_files = []
        for result, idx in zip(results, indices):
            if to_label_id:
                result = self._convert_to_label_id(result)
            filename = self.img_infos[idx]["filename"]
            basename = osp.splitext(osp.basename(filename))[0]

            png_filename = osp.join(imgfile_prefix, f"{basename}.png")

            output = Image.fromarray(result.astype(np.uint8)).convert("P")
            import cityscapesscripts.helpers.labels as CSLabels

            palette = np.zeros((len(CSLabels.id2label), 3), dtype=np.uint8)
            for label_id, label in CSLabels.id2label.items():
                palette[label_id] = label.color

            output.putpalette(palette)
            output.save(png_filename)
            result_files.append(png_filename)

        return result_files

    def format_results(self, results, imgfile_prefix, to_label_id=True, indices=None):
        """Format the results into dir (standard format for Cityscapes
        evaluation).

        Args:
            results (list): Testing results of the dataset.
            imgfile_prefix (str): The prefix of images files. It
                includes the file path and the prefix of filename, e.g.,
                "a/b/prefix".
            to_label_id (bool): whether convert output to label_id for
                submission. Default: False
            indices (list[int], optional): Indices of input results,
                if not set, all the indices of the dataset will be used.
                Default: None.

        Returns:
            tuple: (result_files, tmp_dir), result_files is a list containing
                the image paths, tmp_dir is the temporal directory created
                for saving json/png files when img_prefix is not specified.
        """
        if indices is None:
            indices = list(range(len(self)))

        assert isinstance(results, list), "results must be a list."
        assert isinstance(indices, list), "indices must be a list."

        result_files = self.results2img(results, imgfile_prefix, to_label_id, indices)

        return result_files

    def evaluate(self, results, metric="mIoU", logger=None, imgfile_prefix=None):
        """Evaluation in Cityscapes/default protocol.

        Args:
            results (list): Testing results of the dataset.
            metric (str | list[str]): Metrics to be evaluated.
            logger (logging.Logger | None | str): Logger used for printing
                related information during evaluation. Default: None.
            imgfile_prefix (str | None): The prefix of output image file,
                for cityscapes evaluation only. It includes the file path and
                the prefix of filename, e.g., "a/b/prefix".
                If results are evaluated with cityscapes protocol, it would be
                the prefix of output png files. The output files would be
                png images under folder "a/b/prefix/xxx.png", where "xxx" is
                the image name of cityscapes. If not specified, a temp file
                will be created for evaluation.
                Default: None.

        Returns:
            dict[str, float]: Cityscapes/default metrics.
        """

        eval_results = dict()
        metrics = metric.copy() if isinstance(metric, list) else [metric]
        if "cityscapes" in metrics:
            eval_results.update(
                self._evaluate_cityscapes(results, logger, imgfile_prefix)
            )
            metrics.remove("cityscapes")
        if len(metrics) > 0:
            eval_results.update(
                super(JointCityscapesDataset, self).evaluate(results, metrics, logger)
            )

        return eval_results

    def _evaluate_cityscapes(self, results, logger, imgfile_prefix):
        """Evaluation in Cityscapes protocol.

        Args:
            results (list): Testing results of the dataset.
            logger (logging.Logger | str | None): Logger used for printing
                related information during evaluation. Default: None.
            imgfile_prefix (str | None): The prefix of output image file

        Returns:
            dict[str: float]: Cityscapes evaluation results.
        """
        try:
            import cityscapesscripts.evaluation.evalPixelLevelSemanticLabeling as CSEval  # noqa
        except ImportError:
            raise ImportError(
                'Please run "pip install cityscapesscripts" to '
                "install cityscapesscripts first."
            )
        msg = "Evaluating in Cityscapes style"
        if logger is None:
            msg = "\n" + msg
        print_log(msg, logger=logger)

        result_dir = imgfile_prefix

        eval_results = dict()
        print_log(f"Evaluating results under {result_dir} ...", logger=logger)

        CSEval.args.evalInstLevelScore = True
        CSEval.args.predictionPath = osp.abspath(result_dir)
        CSEval.args.evalPixelAccuracy = True
        CSEval.args.JSONOutput = False

        seg_map_list = []
        pred_list = []

        # when evaluating with official cityscapesscripts,
        # **_gtFine_labelIds.png is used
        for seg_map in mmcv.scandir(
            self.ann_dir, "_gtFine_labelIds.png", recursive=True
        ):
            seg_map_list.append(osp.join(self.other_dir, seg_map))
            pred_list.append(CSEval.getPrediction(CSEval.args, seg_map))

        eval_results.update(
            CSEval.evaluateImgLists(pred_list, seg_map_list, CSEval.args)
        )

        return eval_results


@DATASETS.register_module()
class OTFJointCityscapesDataset(OTFCustomJointDataset):
    """OTF edge generation for Cityscapes dataset

    - should only be used for training the model
    """

    CLASSES = CITYSCAPES_CLASSES
    PALETTE = CITYSCAPES_PALETTE

    def __init__(
        self,
        img_suffix="_leftImg8bit.png",
        seg_map_suffix="_gtFine_labelIds.png",
        inst_map_suffix="_gtFine_instanceIds.png",
        inst_sensitive=True,  # default instance sensitive
        labelIds=CITYSCAPES_labelIds,
        inst_labelIds=CITYSCAPES_inst_labelIds,
        ignore_indicies=[2, 3],
        label2trainId=CITYSCAPES_label2trainId,
        **kwargs,
    ):
        super(OTFJointCityscapesDataset, self).__init__(
            img_suffix=img_suffix,
            seg_map_suffix=seg_map_suffix,
            inst_map_suffix=inst_map_suffix,
            inst_sensitive=inst_sensitive,
            labelIds=labelIds,
            inst_labelIds=inst_labelIds,
            ignore_indices=ignore_indicies,
            label2trainId=label2trainId,
            **kwargs,
        )

    def format_results(self, **kwargs):
        raise ValueError("ERR: Should not use OTF for test set")

    def evaluate(self, **kwargs):
        raise ValueError("ERR: Should not use OTF for evaluation!")
