# Copyright (C) 2023-2024 Intel Corporation
#
# SPDX-License-Identifier: MIT

import logging as log
import os
import os.path as osp
from textwrap import wrap
from typing import Dict, List, Set, Tuple
from unittest import TestCase

from attr import attrib, attrs
from tabulate import tabulate

from datumaro.cli.util.project import generate_next_file_name
from datumaro.components.annotation import AnnotationType, LabelCategories
from datumaro.components.annotations.matcher import LineMatcher, PointsMatcher, match_segments_pair
from datumaro.components.dataset import Dataset
from datumaro.components.operations import (
    compute_ann_statistics,
    compute_image_statistics,
    match_items_by_id,
    match_items_by_image_hash,
)
from datumaro.components.shift_analyzer import ShiftAnalyzer
from datumaro.util import dump_json_file, filter_dict, find
from datumaro.util.annotation_util import find_instances, max_bbox
from datumaro.util.attrs_util import default_if_none


@attrs
class DistanceComparator:
    iou_threshold = attrib(converter=float, default=0.5)

    def match_annotations(self, item_a, item_b):
        return {t: self._match_ann_type(t, item_a, item_b) for t in AnnotationType}

    def _match_ann_type(self, t, *args):
        # pylint: disable=no-value-for-parameter
        if t == AnnotationType.label:
            return self.match_labels(*args)
        elif t == AnnotationType.bbox:
            return self.match_boxes(*args)
        elif t == AnnotationType.polygon:
            return self.match_polygons(*args)
        elif t == AnnotationType.mask:
            return self.match_masks(*args)
        elif t == AnnotationType.points:
            return self.match_points(*args)
        elif t == AnnotationType.polyline:
            return self.match_lines(*args)
        # pylint: enable=no-value-for-parameter
        else:
            raise NotImplementedError("Unexpected annotation type %s" % t)

    @staticmethod
    def _get_ann_type(t, item):
        return [a for a in item.annotations if a.type == t]

    def match_labels(self, item_a, item_b):
        a_labels = set(a.label for a in self._get_ann_type(AnnotationType.label, item_a))
        b_labels = set(a.label for a in self._get_ann_type(AnnotationType.label, item_b))

        matches = a_labels & b_labels
        a_unmatched = a_labels - b_labels
        b_unmatched = b_labels - a_labels
        return matches, a_unmatched, b_unmatched

    def _match_segments(self, t, item_a, item_b):
        a_boxes = self._get_ann_type(t, item_a)
        b_boxes = self._get_ann_type(t, item_b)
        return match_segments_pair(a_boxes, b_boxes, dist_thresh=self.iou_threshold)

    def match_polygons(self, item_a, item_b):
        return self._match_segments(AnnotationType.polygon, item_a, item_b)

    def match_masks(self, item_a, item_b):
        return self._match_segments(AnnotationType.mask, item_a, item_b)

    def match_boxes(self, item_a, item_b):
        return self._match_segments(AnnotationType.bbox, item_a, item_b)

    def match_points(self, item_a, item_b):
        a_points = self._get_ann_type(AnnotationType.points, item_a)
        b_points = self._get_ann_type(AnnotationType.points, item_b)

        instance_map = {}
        for s in [item_a.annotations, item_b.annotations]:
            s_instances = find_instances(s)
            for inst in s_instances:
                inst_bbox = max_bbox(inst)
                for ann in inst:
                    instance_map[id(ann)] = [inst, inst_bbox]
        matcher = PointsMatcher(instance_map=instance_map)

        return match_segments_pair(
            a_points, b_points, dist_thresh=self.iou_threshold, distance=matcher.distance
        )

    def match_lines(self, item_a, item_b):
        a_lines = self._get_ann_type(AnnotationType.polyline, item_a)
        b_lines = self._get_ann_type(AnnotationType.polyline, item_b)

        matcher = LineMatcher()

        return match_segments_pair(
            a_lines, b_lines, dist_thresh=self.iou_threshold, distance=matcher.distance
        )


@attrs
class EqualityComparator:
    match_images: bool = attrib(kw_only=True, default=False)
    ignored_fields = attrib(kw_only=True, factory=set, validator=default_if_none(set))
    ignored_attrs = attrib(kw_only=True, factory=set, validator=default_if_none(set))
    ignored_item_attrs = attrib(kw_only=True, factory=set, validator=default_if_none(set))
    all = attrib(kw_only=True, default=False)

    _test: TestCase = attrib(init=False)
    errors: list = attrib(init=False)

    def __attrs_post_init__(self):
        self._test = TestCase()
        self._test.maxDiff = None

    def _match_items(self, a, b):
        if self.match_images:
            return match_items_by_image_hash(a, b)
        else:
            return match_items_by_id(a, b)

    def _compare_categories(self, a, b):
        test = self._test
        errors = self.errors

        try:
            test.assertEqual(sorted(a, key=lambda t: t.value), sorted(b, key=lambda t: t.value))
        except AssertionError as e:
            errors.append({"type": "categories", "message": str(e)})

        if AnnotationType.label in a:
            try:
                test.assertEqual(
                    a[AnnotationType.label].items,
                    b[AnnotationType.label].items,
                )
            except AssertionError as e:
                errors.append({"type": "labels", "message": str(e)})
        if AnnotationType.mask in a:
            try:
                test.assertEqual(
                    a[AnnotationType.mask].colormap,
                    b[AnnotationType.mask].colormap,
                )
            except AssertionError as e:
                errors.append({"type": "colormap", "message": str(e)})
        if AnnotationType.points in a:
            try:
                test.assertEqual(
                    a[AnnotationType.points].items,
                    b[AnnotationType.points].items,
                )
            except AssertionError as e:
                errors.append({"type": "points", "message": str(e)})

    def _compare_annotations(self, a, b):
        ignored_fields = self.ignored_fields
        ignored_attrs = self.ignored_attrs

        a_fields = {k: None for k in a.as_dict() if k in ignored_fields}
        b_fields = {k: None for k in b.as_dict() if k in ignored_fields}
        if "attributes" not in ignored_fields:
            a_fields["attributes"] = filter_dict(a.attributes, ignored_attrs)
            b_fields["attributes"] = filter_dict(b.attributes, ignored_attrs)

        result = a.wrap(**a_fields) == b.wrap(**b_fields)

        return result

    def _compare_items(self, item_a, item_b):
        test = self._test

        a_id = (item_a.id, item_a.subset)
        b_id = (item_b.id, item_b.subset)

        matched = []
        unmatched = []
        errors = []

        try:
            test.assertEqual(
                filter_dict(item_a.attributes, self.ignored_item_attrs),
                filter_dict(item_b.attributes, self.ignored_item_attrs),
            )
        except AssertionError as e:
            errors.append({"type": "item_attr", "a_item": a_id, "b_item": b_id, "message": str(e)})

        b_annotations = item_b.annotations[:]
        for ann_a in item_a.annotations:
            ann_b_candidates = [x for x in item_b.annotations if x.type == ann_a.type]

            ann_b = find(
                enumerate(self._compare_annotations(ann_a, x) for x in ann_b_candidates),
                lambda x: x[1],
            )
            if ann_b is None:
                unmatched.append(
                    {
                        "item": a_id,
                        "source": "a",
                        "ann": str(ann_a),
                    }
                )
                continue
            else:
                ann_b = ann_b_candidates[ann_b[0]]

            b_annotations.remove(ann_b)  # avoid repeats
            matched.append({"a_item": a_id, "b_item": b_id, "a": str(ann_a), "b": str(ann_b)})

        for ann_b in b_annotations:
            unmatched.append({"item": b_id, "source": "b", "ann": str(ann_b)})

        return matched, unmatched, errors

    @staticmethod
    def _print_output(output: dict):
        print("Found:")
        print("The first project has %s unmatched items" % len(output.get("a_extra_items", [])))
        print("The second project has %s unmatched items" % len(output.get("b_extra_items", [])))
        print("%s item conflicts" % len(output.get("errors", [])))
        print("%s matching annotations" % len(output.get("matches", [])))
        print("%s mismatching annotations" % len(output.get("mismatches", [])))

    def compare_datasets(self, a, b):
        self.errors = []
        errors = self.errors

        self._compare_categories(a.categories(), b.categories())

        matched = []
        unmatched = []

        matches, a_unmatched, b_unmatched = self._match_items(a, b)

        if a.categories().get(AnnotationType.label) != b.categories().get(AnnotationType.label):
            output = {
                "mismatches": unmatched,
                "a_extra_items": sorted(a_unmatched),
                "b_extra_items": sorted(b_unmatched),
                "errors": errors,
            }
            if self.all:
                output["matches"] = matched

            self._print_output(output)
            return output

        _dist = lambda s: len(s[1]) + len(s[2])
        for a_ids, b_ids in matches:
            # build distance matrix
            match_status = {}  # (a_id, b_id): [matched, unmatched, errors]
            a_matches = {a_id: None for a_id in a_ids}
            b_matches = {b_id: None for b_id in b_ids}

            for a_id in a_ids:
                item_a = a.get(*a_id)
                candidates = {}

                for b_id in b_ids:
                    item_b = b.get(*b_id)

                    i_m, i_um, i_err = self._compare_items(item_a, item_b)
                    candidates[b_id] = [i_m, i_um, i_err]

                    if len(i_um) == 0:
                        a_matches[a_id] = b_id
                        b_matches[b_id] = a_id
                        matched.extend(i_m)
                        errors.extend(i_err)
                        break

                match_status[a_id] = candidates

            # assign
            for a_id in a_ids:
                if len(b_ids) == 0:
                    break

                # find the closest, ignore already assigned
                matched_b = a_matches[a_id]
                if matched_b is not None:
                    continue
                min_dist = -1
                for b_id in b_ids:
                    if b_matches[b_id] is not None:
                        continue
                    d = _dist(match_status[a_id][b_id])
                    if d < min_dist and 0 <= min_dist:
                        continue
                    min_dist = d
                    matched_b = b_id

                if matched_b is None:
                    continue
                a_matches[a_id] = matched_b
                b_matches[matched_b] = a_id

                m = match_status[a_id][matched_b]
                matched.extend(m[0])
                unmatched.extend(m[1])
                errors.extend(m[2])

            a_unmatched |= set(a_id for a_id, m in a_matches.items() if not m)
            b_unmatched |= set(b_id for b_id, m in b_matches.items() if not m)

        output = {
            "mismatches": unmatched,
            "a_extra_items": sorted(a_unmatched),
            "b_extra_items": sorted(b_unmatched),
            "errors": errors,
        }
        if self.all:
            output["matches"] = matched
        self._print_output(output)
        return output

    @staticmethod
    def save_compare_report(
        output: Dict,
        report_dir: str,
    ) -> None:
        """Saves the comparison report to JSON and text files.

        Args:
            output: A dictionary containing the comparison data.
            report_dir: A string representing the directory to save the report files.
        """
        os.makedirs(report_dir, exist_ok=True)
        output_file = osp.join(
            report_dir,
            generate_next_file_name("equality_compare", ext=".json", basedir=report_dir),
        )

        log.info(f"Saving compare json to {output_file}")
        dump_json_file(output_file, output, indent=True)


@attrs
class TableComparator:
    """
    Class for comparing datasets and generating comparison report table.
    """

    @staticmethod
    def _extract_labels(dataset: Dataset) -> Set[str]:
        """Extracts labels from the dataset.

        Args:
            dataset: An instance of a Dataset class.

        Returns:
            A set of labels present in the dataset.
        """
        label_cat = dataset.categories().get(AnnotationType.label, LabelCategories())
        return set(c.name for c in label_cat)

    @staticmethod
    def _compute_statistics(dataset: Dataset) -> Tuple[Dict, Dict]:
        """Computes image and annotation statistics of the dataset.

        Args:
            dataset: An instance of a Dataset class.

        Returns:
            A tuple containing image statistics and annotation statistics.
        """
        image_stats = compute_image_statistics(dataset)
        ann_stats = compute_ann_statistics(dataset)
        return image_stats, ann_stats

    def _analyze_dataset(self, dataset: Dataset) -> Tuple[str, Set[str], Dict, Dict]:
        """Analyzes the dataset to get labels, format, and statistics.

        Args:
            dataset: An instance of a Dataset class.

        Returns:
            A tuple containing Dataset format, set of label names, image statistics,
            and annotation statistics.
        """
        dataset_format = dataset.format
        dataset_labels = self._extract_labels(dataset)
        image_stats, ann_stats = self._compute_statistics(dataset)
        return dataset_format, dataset_labels, image_stats, ann_stats

    @staticmethod
    def _create_table(headers: List[str], rows: List[List[str]]) -> str:
        """Creates a table with the given headers and rows using the tabulate module.

        Args:
            headers: A list containing table headers.
            rows: A list containing table rows.

        Returns:
            A string representation of the table.
        """

        def wrapfunc(item):
            """Wrap a item consisted of text, returning a list of wrapped lines."""
            max_len = 35
            return "\n".join(wrap(item, max_len))

        wrapped_rows = []
        for row in rows:
            new_row = [wrapfunc(item) for item in row]
            wrapped_rows.append(new_row)

        return tabulate(wrapped_rows, headers, tablefmt="grid")

    @staticmethod
    def _create_dict(rows: List[List[str]]) -> Dict[str, List[str]]:
        """Creates a dictionary from the rows of the table.

        Args:
            rows: A list containing table rows.

        Returns:
            A dictionary where the key is the first element of a row and the value is
            the rest of the row.
        """
        data_dict = {row[0]: row[1:] for row in rows[1:]}
        return data_dict

    def _create_high_level_comparison_table(
        self, first_info: Tuple, second_info: Tuple
    ) -> Tuple[str, Dict]:
        """Generates a high-level comparison table.

        Args:
            first_info: A tuple containing information about the first dataset.
            second_info: A tuple containing information about the second dataset.

        Returns:
            A tuple containing the table as a string and a dictionary representing the data
            of the table.
        """
        first_format, first_labels, first_image_stats, first_ann_stats = first_info
        second_format, second_labels, second_image_stats, second_ann_stats = second_info

        headers = ["Field", "First", "Second"]

        rows = [
            ["Format", first_format, second_format],
            ["Number of classes", str(len(first_labels)), str(len(second_labels))],
            [
                "Common classes",
                ", ".join(sorted(list(first_labels.intersection(second_labels)))),
                ", ".join(sorted(list(second_labels.intersection(first_labels)))),
            ],
            ["Classes", ", ".join(sorted(first_labels)), ", ".join(sorted(second_labels))],
            [
                "Images count",
                str(first_image_stats["dataset"]["images count"]),
                str(second_image_stats["dataset"]["images count"]),
            ],
            [
                "Unique images count",
                str(first_image_stats["dataset"]["unique images count"]),
                str(second_image_stats["dataset"]["unique images count"]),
            ],
            [
                "Repeated images count",
                str(first_image_stats["dataset"]["repeated images count"]),
                str(second_image_stats["dataset"]["repeated images count"]),
            ],
            [
                "Annotations count",
                str(first_ann_stats["annotations count"]),
                str(second_ann_stats["annotations count"]),
            ],
            [
                "Unannotated images count",
                str(first_ann_stats["unannotated images count"]),
                str(second_ann_stats["unannotated images count"]),
            ],
        ]

        table = self._create_table(headers, rows)
        data_dict = self._create_dict(rows)

        return table, data_dict

    def _create_mid_level_comparison_table(
        self, first_info: Tuple, second_info: Tuple
    ) -> Tuple[str, Dict]:
        """Generates a mid-level comparison table.

        Args:
            first_info: A tuple containing information about the first dataset.
            second_info: A tuple containing information about the second dataset.

        Returns:
            A tuple containing the table as a string and a dictionary representing the data
            of the table.
        """
        _, _, first_image_stats, first_ann_stats = first_info
        _, _, second_image_stats, second_ann_stats = second_info

        headers = ["Field", "First", "Second"]

        rows = []

        first_subsets = sorted(list(first_image_stats["subsets"].keys()))
        second_subsets = sorted(list(second_image_stats["subsets"].keys()))

        subset_names = first_subsets.copy()
        subset_names.extend(item for item in second_subsets if item not in first_subsets)

        for subset_name in subset_names:
            first_subset_data = first_image_stats["subsets"].get(subset_name, {})
            second_subset_data = second_image_stats["subsets"].get(subset_name, {})
            mean_str_first = (
                ", ".join(f"{val:6.2f}" for val in first_subset_data.get("image mean (RGB)", []))
                if "image mean (RGB)" in first_subset_data
                else ""
            )
            std_str_first = (
                ", ".join(f"{val:6.2f}" for val in first_subset_data.get("image std (RGB)", []))
                if "image std" in first_subset_data
                else ""
            )
            mean_str_second = (
                ", ".join(f"{val:6.2f}" for val in second_subset_data.get("image mean (RGB)", []))
                if "image mean (RGB)" in second_subset_data
                else ""
            )
            std_str_second = (
                ", ".join(f"{val:6.2f}" for val in second_subset_data.get("image std", []))
                if "image std (RGB)" in second_subset_data
                else ""
            )
            rows.append([f"{subset_name} - Image Mean (RGB)", mean_str_first, mean_str_second])
            rows.append([f"{subset_name} - Image Std (RGB)", std_str_first, std_str_second])

        first_labels = sorted(list(first_ann_stats["annotations"]["labels"]["distribution"].keys()))
        second_labels = sorted(
            list(second_ann_stats["annotations"]["labels"]["distribution"].keys())
        )

        label_names = first_labels.copy()
        label_names.extend(item for item in second_labels if item not in first_labels)

        for label_name in label_names:
            count_dist_first = first_ann_stats["annotations"]["labels"]["distribution"].get(
                label_name, [0, 0.0]
            )
            count_dist_second = second_ann_stats["annotations"]["labels"]["distribution"].get(
                label_name, [0, 0.0]
            )
            count_first, dist_first = count_dist_first if count_dist_first[0] != 0 else ["", ""]
            count_second, dist_second = count_dist_second if count_dist_second[0] != 0 else ["", ""]
            rows.append(
                [
                    f"Label - {label_name}",
                    f"imgs: {count_first}, percent: {dist_first:.4f}" if count_first != "" else "",
                    f"imgs: {count_second}, percent: {dist_second:.4f}"
                    if count_second != ""
                    else "",
                ]
            )

        table = self._create_table(headers, rows)
        data_dict = self._create_dict(rows)

        return table, data_dict

    def _create_low_level_comparison_table(
        self, first_dataset: Dataset, second_dataset: Dataset
    ) -> Tuple[str, Dict]:
        """Generates a low-level comparison table.

        Args:
            first_dataset: The first dataset to compare.
            second_dataset: The second dataset to compare.

        Returns:
            A tuple containing the table as a string and a dictionary representing the data
            of the table.
        """
        shift_analyzer = ShiftAnalyzer()
        cov_shift = shift_analyzer.compute_covariate_shift([first_dataset, second_dataset])
        label_shift = shift_analyzer.compute_label_shift([first_dataset, second_dataset])

        headers = ["Field", "Value"]

        rows = [
            ["Covariate shift", str(cov_shift)],
            ["Label shift", str(label_shift)],
        ]

        table = self._create_table(headers, rows)
        data_dict = self._create_dict(rows)

        return table, data_dict

    def compare_datasets(
        self, first: Dataset, second: Dataset, mode: str = "all"
    ) -> Tuple[str, str, str, Dict]:
        """Compares two datasets and generates comparison reports.

        Args:
            first: The first dataset to compare.
            second: The second dataset to compare.

        Returns:
            A tuple containing high-level table, mid-level table, low-level table, and a
            dictionary representation of the comparison.
        """
        first_info = self._analyze_dataset(first)
        second_info = self._analyze_dataset(second)

        high_level_table, high_level_dict = None, {}
        mid_level_table, mid_level_dict = None, {}
        low_level_table, low_level_dict = None, {}

        if mode in ["high", "all"]:
            high_level_table, high_level_dict = self._create_high_level_comparison_table(
                first_info, second_info
            )
        if mode in ["mid", "all"]:
            mid_level_table, mid_level_dict = self._create_mid_level_comparison_table(
                first_info, second_info
            )
        if mode in ["low", "all"]:
            low_level_table, low_level_dict = self._create_low_level_comparison_table(first, second)

        comparison_dict = dict(
            high_level=high_level_dict, mid_level=mid_level_dict, low_level=low_level_dict
        )

        print(f"High-level comparison:\n{high_level_table}\n")
        print(f"Mid-level comparison:\n{mid_level_table}\n")
        print(f"Low-level comparison:\n{low_level_table}\n")

        return high_level_table, mid_level_table, low_level_table, comparison_dict

    @staticmethod
    def save_compare_report(
        high_level_table: str,
        mid_level_table: str,
        low_level_table: str,
        comparison_dict: Dict,
        report_dir: str,
    ) -> None:
        """Saves the comparison report to JSON and text files.

        Args:
            high_level_table: High-level comparison table as a string.
            mid_level_table: Mid-level comparison table as a string.
            low_level_table: Low-level comparison table as a string.
            comparison_dict: A dictionary containing the comparison data.
            report_dir: A string representing the directory to save the report files.
        """
        os.makedirs(report_dir, exist_ok=True)
        json_output_file = osp.join(
            report_dir, generate_next_file_name("table_compare", ext=".json", basedir=report_dir)
        )
        txt_output_file = osp.join(
            report_dir, generate_next_file_name("table_compare", ext=".txt", basedir=report_dir)
        )

        log.info(f"Saving compare json to {json_output_file}")
        log.info(f"Saving compare table to {txt_output_file}")

        dump_json_file(json_output_file, comparison_dict, indent=True)
        with open(txt_output_file, "w") as f:
            f.write(f"High-level Comparison:\n{high_level_table}\n\n")
            f.write(f"Mid-level Comparison:\n{mid_level_table}\n\n")
            f.write(f"Low-level Comparison:\n{low_level_table}\n\n")
