import glob
import json
import os
import argparse

import numpy as np
import torch
import tqdm
from PIL import Image
import matplotlib.pyplot as plt
from detectron2.evaluation import SemSegEvaluator


import sys
sys.path.insert(0, "./")

import data.register_coco_panoptic_annos_semseg

try:
    np.int
except:
    np.int = np.int32
    np.float = np.float32


class SemSegEvaluatorCustom(SemSegEvaluator):
    def __init__(
        self,
        dataset_name,
        distributed=True,
        output_dir=None,
        palette=None,
        pred_dir=None,
        dist_type=None,
    ):
        """
        Args:
            dataset_name (str): name of the dataset to be evaluated.
            distributed (bool): if True, will collect results from all ranks for evaluation.
                Otherwise, will evaluate the results in the current process.
            output_dir (str): an output directory to dump results.
        """
        super().__init__(
            dataset_name=dataset_name,
            distributed=distributed,
            output_dir=output_dir,
        )

        # update source names
        print("num gt:", len(self.input_file_to_gt_file))
        self.input_file_to_gt_file_custom = {}
        for src_file, tgt_file in self.input_file_to_gt_file.items():
            assert os.path.basename(src_file).replace('.jpg', '.png') == os.path.basename(tgt_file)
            src_file_custom = os.path.join(pred_dir, os.path.basename(tgt_file))  # output is saved as png
            self.input_file_to_gt_file_custom[src_file_custom] = tgt_file

        color_to_idx = {}
        for cls_idx, color in enumerate(palette):
            color = tuple(color)
            # in coco, foreground index starts from 0
            color_to_idx[color] = cls_idx
        self.color_to_idx = color_to_idx
        self.palette = torch.tensor(palette, dtype=torch.float, device="cuda")  # (num_cls, 3)
        self.pred_dir = pred_dir
        self.dist_type = dist_type

    def process(self, inputs, outputs):
        """
        Args:
            inputs: the inputs to a model.
                It is a list of dicts. Each dict corresponds to an image and
                contains keys like "height", "width", "file_name".
            outputs: the outputs of a model. It is either list of semantic segmentation predictions
                (Tensor [H, W]) or list of dicts with key "sem_seg" that contains semantic
                segmentation prediction in the same format.
        """
        print("processing")
        for input in tqdm.tqdm(inputs):
            # output = output["sem_seg"].argmax(dim=0).to(self._cpu_device)  # chw --> hw
            output = input["file_name"]
            output = Image.open(output)
            output = np.array(output)  # (h, w, 3)

            # use custom input_file_to_gt_file mapping
            gt_filename = self.input_file_to_gt_file_custom[input["file_name"]]
            gt = self.sem_seg_loading_fn(gt_filename, dtype=np.int)

            gt[gt == self._ignore_label] = self._num_classes

            pred, dist = self.post_process_segm_output(output)

            self._conf_matrix += np.bincount(
                (self._num_classes + 1) * pred.reshape(-1) + gt.reshape(-1),
                minlength=self._conf_matrix.size,
            ).reshape(self._conf_matrix.shape)

            if self._compute_boundary_iou:
                b_gt = self._mask_to_boundary(gt.astype(np.uint8))
                b_pred = self._mask_to_boundary(pred.astype(np.uint8))

                self._b_conf_matrix += np.bincount(
                    (self._num_classes + 1) * b_pred.reshape(-1) + b_gt.reshape(-1),
                    minlength=self._conf_matrix.size,
                ).reshape(self._conf_matrix.shape)

            self._predictions.extend(self.encode_json_sem_seg(pred, input["file_name"]))

    def post_process_segm_output(self, segm):
        """
        Post-processing to turn output segm image to class index map

        Args:
            segm: (H, W, 3)

        Returns:
            class_map: (H, W)
        """
        segm = torch.from_numpy(segm).float().to(self.palette.device)  # (h, w, 3)
        h, w, k = segm.shape[0], segm.shape[1], self.palette.shape[0]

        if self.dist_type == 'abs':
            dist = torch.abs(segm.view(h, w, 1, 3) - self.palette.view(1, 1, k, 3))  # (h, w, k)
        elif self.dist_type == 'square':
            dist = torch.pow(segm.view(h, w, 1, 3) - self.palette.view(1, 1, k, 3), 2)  # (h, w, k)
        elif self.dist_type == 'mean':
            dist_abs = torch.abs(segm.view(h, w, 1, 3) - self.palette.view(1, 1, k, 3))  # (h, w, k)
            dist_square = torch.pow(segm.view(h, w, 1, 3) - self.palette.view(1, 1, k, 3), 2)  # (h, w, k)
            dist = (dist_abs + dist_square) / 2.
        else:
            raise NotImplementedError

        dist = torch.sum(dist, dim=-1)
        pred = dist.argmin(dim=-1).cpu()  # (h, w)
        pred = np.array(pred, dtype=np.int)

        return pred, dist.cpu().numpy()


def get_args_parser():
    parser = argparse.ArgumentParser('COCO semantic segmentation', add_help=False)
    parser.add_argument('--pred_dir', type=str, help='dir to ckpt', default=None)
    parser.add_argument('--dist_type', type=str, help='color type',
                        default='abs', choices=['abs', 'square', 'mean'])
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args_parser()
    dataset_name = 'coco_2017_val_panoptic_with_sem_seg'
    pred_dir = args.pred_dir

    output_folder = os.path.join(pred_dir, 'eval_pano_semseg')

    # load cat info
    panoptic_coco_categories = 'data/panoptic_coco_categories.json'
    with open(panoptic_coco_categories, 'r') as f:
        categories_list = json.load(f)
    categories = {category['id']: category for category in categories_list}
    catid2colorid = {category['id']: idx for idx, category in enumerate(categories_list)}

    # define colors (dict of cat_id to color mapper)
    num_colors = len(categories)
    from data.coco_semseg.gen_color_coco_panoptic_segm import define_colors_by_mean_sep
    PALETTE_DICT = define_colors_by_mean_sep(num_colors=num_colors)
    PALETTE = [v for k, v in PALETTE_DICT.items()]

    evaluator = SemSegEvaluatorCustom(
        dataset_name,
        distributed=True,
        output_dir=output_folder,
        palette=PALETTE,
        pred_dir=pred_dir,
        dist_type=args.dist_type,
    )

    inputs = []
    outputs = []
    prediction_list = glob.glob(os.path.join(pred_dir, "*.png"))
    print("num_pred: ", len(prediction_list))
    print("loading predictions")
    for file_name in prediction_list:
        # keys in input: "file_name", keys in output: "sem_seg"
        input_dict = {"file_name": file_name}
        output_dict = {"sem_seg": file_name}
        inputs.append(input_dict)
        outputs.append(output_dict)

    evaluator.reset()
    evaluator.process(inputs, outputs)
    results = evaluator.evaluate()
    print(results)

    copy_paste_results = {}
    for key in ['mIoU', 'fwIoU', 'mACC', 'pACC']:
        copy_paste_results[key] = results['sem_seg'][key]
    print(copy_paste_results)

    result_file = os.path.join(output_folder, "results.txt")
    print("writing to {}".format(result_file))
    with open(result_file, 'w') as f:
        print(results, file=f)
        print(copy_paste_results, file=f)
