# Copyright (C) 2019-2021 Intel Corporation
#
# SPDX-License-Identifier: MIT

import argparse
import logging as log
import os
import os.path as osp

from datumaro.components.project import Project
from datumaro.util.command_targets import (
    ImageTarget, ProjectTarget, SourceTarget, TargetKinds, is_project_path,
    target_selector,
)
from datumaro.util.image import load_image, save_image

from ..util import MultilineFormatter
from ..util.project import load_project


def build_parser(parser_ctor=argparse.ArgumentParser):
    parser = parser_ctor(help="Run Explainable AI algorithm",
        description="""
        Runs an explainable AI algorithm for a model.|n
        |n
        This tool is supposed to help an AI developer to debug
        a model and a dataset. Basically, it executes inference and
        tries to find problems in the trained model - determine decision
        boundaries and belief intervals for the classifier.|n
        |n
        Currently, the only available algorithm is
        RISE (https://arxiv.org/pdf/1806.07421.pdf), which runs
        inference and then re-runs a model multiple times
        on each image to produce a heatmap of activations for
        each output of the first inference. As a result, we obtain
        few heatmaps, which shows, how image pixels affected
        the inference result. This algorithm doesn't require any special
        information about the model, but it requires the model to
        return all the outputs and confidences. Check the User Manual
        for usage examples.|n
        Supported scenarios:|n
        - RISE for classification|n
        - RISE for Object Detection|n
        |n
        Examples:|n
        - Run RISE on an image, display results:|n
        |s|s%(prog)s -t path/to/image.jpg -m mymodel rise --max-samples 50
        """, formatter_class=MultilineFormatter)

    parser.add_argument('-m', '--model', required=True,
        help="Model to use for inference")
    parser.add_argument('-t', '--target', default=None,
        help="Inference target - image, source, project "
             "(default: current project)")
    parser.add_argument('-o', '--output-dir', dest='save_dir', default=None,
        help="Directory to save output (default: display only)")

    method_sp = parser.add_subparsers(dest='algorithm')

    rise_parser = method_sp.add_parser('rise',
        description="""
        RISE: Randomized Input Sampling for
        Explanation of Black-box Models algorithm|n
        |n
        See explanations at: https://arxiv.org/pdf/1806.07421.pdf
        """,
        formatter_class=MultilineFormatter)
    rise_parser.add_argument('-s', '--max-samples', default=None, type=int,
        help="Number of algorithm iterations (default: mask size ^ 2)")
    rise_parser.add_argument('--mw', '--mask-width',
        dest='mask_width', default=7, type=int,
        help="Mask width (default: %(default)s)")
    rise_parser.add_argument('--mh', '--mask-height',
        dest='mask_height', default=7, type=int,
        help="Mask height (default: %(default)s)")
    rise_parser.add_argument('--prob', default=0.5, type=float,
        help="Mask pixel inclusion probability (default: %(default)s)")
    rise_parser.add_argument('--iou', '--iou-thresh',
        dest='iou_thresh', default=0.9, type=float,
        help="IoU match threshold for detections (default: %(default)s)")
    rise_parser.add_argument('--nms', '--nms-iou-thresh',
        dest='nms_iou_thresh', default=0.0, type=float,
        help="IoU match threshold in Non-maxima suppression (default: no NMS)")
    rise_parser.add_argument('--conf', '--det-conf-thresh',
        dest='det_conf_thresh', default=0.0, type=float,
        help="Confidence threshold for detections (default: include all)")
    rise_parser.add_argument('-b', '--batch-size', default=1, type=int,
        help="Inference batch size (default: %(default)s)")
    rise_parser.add_argument('--display', action='store_true',
        help="Visualize results during computations")

    parser.add_argument('-p', '--project', dest='project_dir', default='.',
        help="Directory of the project to operate on (default: current dir)")
    parser.set_defaults(command=explain_command)

    return parser

def explain_command(args):
    project_path = args.project_dir
    if is_project_path(project_path):
        project = Project.load(project_path)
    else:
        project = None
    args.target = target_selector(
        ProjectTarget(is_default=True, project=project),
        SourceTarget(project=project),
        ImageTarget()
    )(args.target)
    if args.target[0] == TargetKinds.project:
        if is_project_path(args.target[1]):
            args.project_dir = osp.dirname(osp.abspath(args.target[1]))


    from matplotlib import cm
    import cv2

    project = load_project(args.project_dir)

    model = project.make_executable_model(args.model)

    if str(args.algorithm).lower() != 'rise':
        raise NotImplementedError()

    from datumaro.components.algorithms.rise import RISE
    rise = RISE(model,
        max_samples=args.max_samples,
        mask_width=args.mask_width,
        mask_height=args.mask_height,
        prob=args.prob,
        iou_thresh=args.iou_thresh,
        nms_thresh=args.nms_iou_thresh,
        det_conf_thresh=args.det_conf_thresh,
        batch_size=args.batch_size)

    if args.target[0] == TargetKinds.image:
        image_path = args.target[1]
        image = load_image(image_path)

        log.info("Running inference explanation for '%s'" % image_path)
        heatmap_iter = rise.apply(image, progressive=args.display)

        image = image / 255.0
        file_name = osp.splitext(osp.basename(image_path))[0]
        if args.display:
            for i, heatmaps in enumerate(heatmap_iter):
                for j, heatmap in enumerate(heatmaps):
                    hm_painted = cm.jet(heatmap)[:, :, 2::-1]
                    disp = (image + hm_painted) / 2
                    cv2.imshow('heatmap-%s' % j, hm_painted)
                    cv2.imshow(file_name + '-heatmap-%s' % j, disp)
                cv2.waitKey(10)
                print("Iter", i, "of", args.max_samples, end='\r')
        else:
            heatmaps = next(heatmap_iter)

        if args.save_dir is not None:
            log.info("Saving inference heatmaps at '%s'" % args.save_dir)
            os.makedirs(args.save_dir, exist_ok=True)

            for j, heatmap in enumerate(heatmaps):
                save_path = osp.join(args.save_dir,
                    file_name + '-heatmap-%s.png' % j)
                save_image(save_path, heatmap * 255.0)
        else:
            for j, heatmap in enumerate(heatmaps):
                disp = (image + cm.jet(heatmap)[:, :, 2::-1]) / 2
                cv2.imshow(file_name + '-heatmap-%s' % j, disp)
            cv2.waitKey(0)
    elif args.target[0] == TargetKinds.source or \
         args.target[0] == TargetKinds.project:
        if args.target[0] == TargetKinds.source:
            source_name = args.target[1]
            dataset = project.make_source_project(source_name).make_dataset()
            log.info("Running inference explanation for '%s'" % source_name)
        else:
            project_name = project.config.project_name
            dataset = project.make_dataset()
            log.info("Running inference explanation for '%s'" % project_name)

        for item in dataset:
            image = item.image.data
            if image is None:
                log.warning(
                    "Dataset item %s does not have image data. Skipping." % \
                    (item.id))
                continue

            heatmap_iter = rise.apply(image)

            image = image / 255.0
            heatmaps = next(heatmap_iter)

            if args.save_dir is not None:
                log.info("Saving inference heatmaps to '%s'" % args.save_dir)
                os.makedirs(args.save_dir, exist_ok=True)

                for j, heatmap in enumerate(heatmaps):
                    save_image(osp.join(args.save_dir,
                            item.id + '-heatmap-%s.png' % j),
                        heatmap * 255.0, create_dir=True)

            if not args.save_dir or args.display:
                for j, heatmap in enumerate(heatmaps):
                    disp = (image + cm.jet(heatmap)[:, :, 2::-1]) / 2
                    cv2.imshow(item.id + '-heatmap-%s' % j, disp)
                cv2.waitKey(0)
    else:
        raise NotImplementedError()

    return 0
