"""Segmentation task interface module.

This module contains a class that defines the objectives of models/trainers for segmentation tasks.
"""
import logging
from typing import Optional  # noqa: F401

import numpy as np
import PIL.Image
import torch
import tqdm

import thelper.concepts
import thelper.utils
from thelper.ifaces import ClassNamesHandler
from thelper.tasks.utils import Task

logger = logging.getLogger(__name__)


@thelper.concepts.segmentation
class Segmentation(Task, ClassNamesHandler):
    """Interface for pixel-level labeling/classification (segmentation) tasks.

    This specialization requests that when given an input tensor, the trained model should
    provide prediction scores for each predefined label (or class), for each element of the
    input tensor. The label names are used here to help categorize samples, and to assure that
    two tasks are only identical when their label counts and ordering match.

    Attributes:
        class_names: map of class name-value pairs to predict for each pixel.
        input_key: the key used to fetch input tensors from a sample dictionary.
        label_map_key: the key used to fetch label (class) maps from a sample dictionary.
        meta_keys: the list of extra keys provided by the data parser inside each sample.
        dontcare: value of the 'dontcare' label (if any) used in the class map.
        color_map: map of class name-color pairs to use when displaying results.

    .. seealso::
        | :class:`thelper.tasks.utils.Task`
        | :class:`thelper.train.segm.ImageSegmTrainer`
    """

    def __init__(self, class_names, input_key, label_map_key,
                 meta_keys=None, dontcare=None, color_map=None):
        """Receives and stores the class (or label) names to predict, the input tensor key,
        the groundtruth label (class) map key, the extra (meta) keys produced by the dataset
        parser(s), the 'dontcare' label value that might be present in gt maps (if any), and
        the color map used to swap label indices for colors when displaying results.

        The class names can be provided as a list of strings, as a path to a json file that
        contains such a list, or as a map of predefined name-value pairs to use in gt maps.
        This list/map must contain at least two elements. All other arguments are used as-is
        to index dictionaries, and must therefore be key-compatible types.
        """
        super(Segmentation, self).__init__(input_key, label_map_key, meta_keys)
        if "dontcare" in class_names:
            logger.warning("found reserved 'dontcare' label in input classes; it will be removed from the internal list")
            if isinstance(class_names, list):
                del class_names[class_names.index("dontcare")]
            elif isinstance(class_names, dict):
                del class_names["dontcare"]
        ClassNamesHandler.__init__(self, class_names=class_names)
        self.dontcare = dontcare
        self.color_map = color_map

    @property
    def dontcare(self):
        """Returns the 'dontcare' label value used in loss functions (can be ``None``)."""
        return self._dontcare

    @dontcare.setter
    def dontcare(self, dontcare):
        """Sets the 'dontcare' label value for this segmentation task (can be ``None``)."""
        if dontcare is not None:
            assert isinstance(dontcare, int), "'dontcare' value should be integer (index)"
            assert dontcare not in self.class_indices.values(), "found 'dontcare' value tied to another class label"
        self._dontcare = dontcare

    @property
    def color_map(self):
        """Returns the color map used to swap label indices for colors when displaying results."""
        return self._color_map

    @color_map.setter
    def color_map(self, color_map):
        """Sets the color map used to swap label indices for colors when displaying results."""
        if color_map is not None:
            assert isinstance(color_map, dict), "color map should be given as dictionary"
            self._color_map = {}
            assert all([isinstance(k, int) for k in color_map]) or all([isinstance(k, str) for k in color_map]), \
                "color map keys should be only class names or only class indices"
            for key, val in color_map.items():
                if isinstance(key, str):
                    if key == "dontcare" and self.dontcare is not None:
                        key = self.dontcare
                    else:
                        assert key in self.class_indices, f"could not find color map key '{key}' in class names"
                        key = self.class_indices[key]
                assert key in self.class_indices.values() or key == self.dontcare, f"unrecognized class index '{key}'"
                if isinstance(val, (list, tuple)):
                    val = np.asarray(val)
                assert isinstance(val, np.ndarray) and val.size == 3, "color values should be given as triplets"
                self._color_map[key] = val
            if self.dontcare is not None and self.dontcare not in self._color_map:
                self._color_map[self.dontcare] = np.asarray([0, 0, 0])  # use black as default 'dontcare' color
        else:
            self._color_map = {}

    def get_class_sizes(self, samples):
        """Given a list of samples, returns a map of element counts for each class label."""
        assert samples is not None and samples, "provided invalid sample list"
        elem_counts = {class_name: 0 for class_name in self.class_names}
        if self.dontcare is not None:
            elem_counts["dontcare"] = 0
        warned_unknown_value_flag = False
        for sample_idx, sample in tqdm.tqdm(enumerate(samples), desc="cumulating label counts", total=len(samples)):
            if self.gt_key is None or self.gt_key not in sample:
                continue
            else:
                labels = sample[self.gt_key]
                if isinstance(labels, torch.Tensor):
                    labels = labels.cpu().numpy()
                if isinstance(labels, PIL.Image.Image):
                    labels = np.array(labels)
                assert isinstance(labels, np.ndarray), "unsupported label map type ('%s')" % str(type(labels))
                # here, we assume labels are given as some integer type that corresponds to class name indices
                curr_elem_counts = {cname: np.count_nonzero(labels == cval) for cname, cval in self.class_indices.items()}
                dontcare_elem_count = 0 if self.dontcare is None else np.count_nonzero(labels == self.dontcare)
                if (sum(curr_elem_counts.values()) + dontcare_elem_count) != labels.size and not warned_unknown_value_flag:
                    logger.warning("some label maps contain values that are unknown (i.e. with no proper class mapping)")
                    warned_unknown_value_flag = True
                for class_name in self.class_names:
                    elem_counts[class_name] += curr_elem_counts[class_name]
                if dontcare_elem_count > 0:
                    elem_counts["dontcare"] += dontcare_elem_count
        return elem_counts

    def check_compat(self, task, exact=False):
        # type: (Segmentation, Optional[bool]) -> bool
        """Returns whether the current task is compatible with the provided one or not.

        This is useful for sanity-checking, and to see if the inputs/outputs of two models
        are compatible. If ``exact = True``, all fields will be checked for exact (perfect)
        compatibility (in this case, matching meta keys and class maps).
        """
        if isinstance(task, Segmentation):
            # if both tasks are related to segmentation: gt keys, class names, and dc must match
            return self.input_key == task.input_key and self.dontcare == task.dontcare and \
                (self.gt_key is None or task.gt_key is None or self.gt_key == task.gt_key) and \
                all([cls in self.class_names for cls in task.class_names]) and \
                (not exact or (self.class_names == task.class_names and
                               set(self.meta_keys) == set(task.meta_keys) and
                               self.color_map.keys() == task.color_map.keys() and
                               all([np.array_equal(self.color_map[k], task.color_map[k]) for k in self.color_map]) and
                               self.gt_key == task.gt_key))
        elif type(task) == Task:
            # if 'task' simply has no gt, compatibility rests on input key only
            return not exact and self.input_key == task.input_key and task.gt_key is None
        return False

    def get_compat(self, task):
        """Returns a task instance compatible with the current task and the given one."""
        assert isinstance(task, Segmentation) or type(task) == Task, \
            f"cannot create compatible task from types '{type(task)}' and '{type(self)}'"
        if isinstance(task, Segmentation):
            assert self.input_key == task.input_key, "input key mismatch, cannot create compatible task"
            assert self.gt_key is None or task.gt_key is None or self.gt_key == task.gt_key, \
                "gt key mismatch, cannot create compatible task"
            assert self.dontcare == task.dontcare, "dontcare value mismatch, cannot create compatible task"
            meta_keys = list(set(self.meta_keys + task.meta_keys))
            # cannot use set for class names, order needs to stay intact!
            class_indices = {cname: cval for cname, cval in task.class_indices.items() if cname not in self.class_indices}
            class_indices = {**self.class_indices, **class_indices}
            color_map = {cname: cval for cname, cval in task.color_map.items() if cname not in self.color_map}
            color_map = {**self.color_map, **color_map}
            return Segmentation(class_names=class_indices, input_key=self.input_key, label_map_key=self.gt_key,
                                meta_keys=meta_keys, dontcare=self.dontcare, color_map=color_map)
        elif type(task) == Task:
            assert self.check_compat(task), f"cannot create compatible task between:\n\t{str(self)}\n\t{str(task)}"
            meta_keys = list(set(self.meta_keys + task.meta_keys))
            return Segmentation(class_names=self.class_indices, input_key=self.input_key, label_map_key=self.gt_key,
                                meta_keys=meta_keys, dontcare=self.dontcare, color_map=self.color_map)

    def __repr__(self):
        """Creates a print-friendly representation of a segmentation task."""
        color_map = {k: v.tolist() for k, v in self.color_map.items()}
        return self.__class__.__module__ + "." + self.__class__.__qualname__ + \
            f"(class_names={repr(self.class_indices)}, input_key={repr(self.input_key)}, " + \
            f"label_map_key={repr(self.gt_key)}, meta_keys={repr(self.meta_keys)}, " + \
            f"dontcare={repr(self.dontcare)}, color_map={repr(color_map)})"
