# AUTOGENERATED! DO NOT EDIT! File to edit: annotation_viewer.ipynb (unless otherwise specified).

__all__ = ['WINDOW_NAME', 'show_annotated_images', 'configure_logging']


# Cell
import sys
import argparse
import logging
import cv2
import numpy as np
from .core import list_subclasses, parse_known_args_with_help
from aiforce import annotation as annotation_package
from .annotation.core import AnnotationAdapter, SubsetType
from .image.opencv_tools import fit_to_max_size
from .annotation.core import RegionShape


# Cell
WINDOW_NAME = 'Annotation'


# Cell
def show_annotated_images(annotation_adapter, max_width=0, max_height=0):
    """
    Show images with corresponding annotations.
    Images are shown one at a time with switching by using the arrow left/right keys.
    `annotation_adapter`: The annotation adapter to use
    `max_width`: The maximum width to scale the image for visibility.
    `max_height`: The maximum height to scale the image for visibility.
    """
    categories = annotation_adapter.read_categories()
    annotations = annotation_adapter.read_annotations(SubsetType.TRAINVAL)
    len_annotations = len(annotations)

    if len_annotations == 0:
        logging.error("No Annotations found")
        return

    index = 0
    annotation_keys = list(annotations.keys())

    logging.info("Keys to use:")
    logging.info("n = Next Image")
    logging.info("b = Previous Image")
    logging.info("q = Quit")

    logging.info("Annotations to view: {}".format(len_annotations))

    while True:
        annotation_id = annotation_keys[index]
        annotation = annotations[annotation_id]
        logging.info("View Image {}/{}: {}".format(index + 1, len_annotations, annotation.file_path))
        img = cv2.imread(annotation.file_path, cv2.IMREAD_COLOR)

        if img is None:
            logging.info("Image not found at {}".format(annotation.file_path))
            img = np.zeros(shape=(1, 1, 3))

        if annotation.regions:
            logging.info("Found {} regions".format(len(annotation.regions)))
            for region_index, region in enumerate(annotation.regions):
                points = list(zip(region.points_x, region.points_y))
                logging.info("Found {} of category {} with {} points: {}".format(region.shape,
                                                                                 ','.join(region.labels),
                                                                                 len(points), points))
                if region.shape == RegionShape.CIRCLE:
                    img = cv2.circle(img, points[0], region.radius_x, (0, 255, 255), 2)
                elif region.shape == RegionShape.ELLIPSE:
                    img = cv2.ellipse(img, points[0], (region.radius_x, region.radius_y), 0, 0, 360,
                                      (0, 255, 255), 2)
                elif region.shape == RegionShape.POINT:
                    img = cv2.circle(img, points[0], 1, (0, 255, 255), 2)
                elif region.shape == RegionShape.POLYGON:
                    pts = np.array(points, np.int32)
                    pts = pts.reshape((-1, 1, 2))
                    img = cv2.polylines(img, [pts], True, (0, 255, 255), 2)
                elif region.shape == RegionShape.RECTANGLE:
                    img = cv2.rectangle(img, points[0], points[1], (0, 255, 255), 2)

        if max_width and max_height:
            img = fit_to_max_size(img, max_width, max_height)

        cv2.imshow(WINDOW_NAME, img)
        cv2.setWindowTitle(WINDOW_NAME, "Image {}/{}".format(index + 1, len_annotations))

        k = cv2.waitKey(0)
        if k == ord('q'):    # 'q' key to stop
            break
        elif k == ord('b'):
            index = max(0, index - 1)
        elif k == ord('n'):
            index = min(len_annotations - 1, index + 1)

    cv2.destroyAllWindows()


# Cell
def configure_logging(logging_level=logging.INFO):
    """
    Configures logging for the system.

    :param logging_level: The logging level to use.
    """
    logging.basicConfig(level=logging_level)


# Cell
if __name__ == '__main__' and '__file__' in globals():
    # for direct shell execution
    configure_logging()

    # read annotation adapters to use
    adapters = list_subclasses(annotation_package, AnnotationAdapter)

    parser = argparse.ArgumentParser()
    parser.add_argument("-a",
                        "--annotation",
                        help="The annotation adapter to read the annotations.",
                        type=str,
                        choices=adapters.keys(),
                        required=True)
    parser.add_argument("--max-width",
                        help="The maximum width to scale the image for visibility.",
                        type=int,
                        default=0)
    parser.add_argument("--max-height",
                        help="The maximum height to scale the image for visibility.",
                        type=int,
                        default=0)

    argv = sys.argv
    args, argv = parse_known_args_with_help(parser, argv)

    adapter_class = adapters[args.annotation]

    # parse the annotation arguments
    annotation_parser = getattr(adapter_class, 'argparse')()
    annotation_args, argv = parse_known_args_with_help(annotation_parser, argv)

    show_annotated_images(adapter_class(**vars(annotation_args)), args.max_width, args.max_height)
