"""This is a utility for converting (hierarchical) annotated images into annotated image snippets"""
import json
import os
import sys
from argparse import ArgumentParser

import cv2

from circuit_recognizer.annotations import Annotation
from circuit_recognizer.utils import get_annotation_source_image, get_image_path

parser = ArgumentParser()
parser.add_argument("annotations", nargs="+")
parser.add_argument("--outdir", default=".")
parser.add_argument("--padding", default=0, type=int)
parser.add_argument(
    "--parents", default="", help="Comma separated list of parent types"
)
parser.add_argument(
    "--absolute",
    default=False,
    action="store_true",
    help="Use absolute file path in annotations",
)
args = parser.parse_args()


def pad_dims(x, y, width, height, img_shape, padding=0):
    (img_height, img_width, _) = img_shape
    left_pad = min(x, padding)
    top_pad = min(y, padding)
    right_pad = min(img_width - (x + width), padding)
    bottom_pad = min(img_height - (y + height), padding)
    padded_width = width + right_pad + left_pad
    padded_height = height + top_pad + bottom_pad
    padded_x = x - left_pad
    padded_y = y - top_pad
    return padded_x, padded_y, padded_width, padded_height


def select_annotations(annotations, parent_types):
    for type in parent_types:
        annotations = [a for a in annotations if a.label() == type]

    return annotations


for annotations_path in args.annotations:
    if os.path.isdir(annotations_path):
        print(f"{annotations_path} is a directory. Skipping...", file=sys.stderr)
        continue

    with open(annotations_path, "r") as f:
        annotations = Annotation.from_dicts(json.load(f))

    image_name = os.path.basename(annotations_path).replace(".json", "")
    os.makedirs(os.path.join(args.outdir, "images"), exist_ok=True)
    os.makedirs(os.path.join(args.outdir, "annotations"), exist_ok=True)

    parent_types = args.parents.split(",")
    annotations = select_annotations(annotations, parent_types)

    if len(annotations) > 0:
        image = get_annotation_source_image(annotations[0])
        for annotation in annotations:
            x, y, width, height = pad_dims(
                *annotation.bounding_box().astype(int), image.shape, args.padding
            )

            snippet = image[y : y + height, x : x + width]
            filename = f"{image_name}.{annotation.id}.{annotation.label()}"
            image_path = f"{args.outdir}/images/{filename}.png"
            if args.absolute:
                image_path = os.path.abspath(image_path)
            cv2.imwrite(image_path, snippet)
            children = []
            for child in annotation.children:
                px, py, pw, ph = child.bounding_box().astype(int)
                overflow = max(
                    x - px,
                    y - py,
                    (px + pw) - (x + width - 1),
                    (py + ph) - (y + height - 1),
                )

                px = max(px - x, 0)
                py = max(py - y, 0)

                # check if the child annotation is off the snippet. Print a warning
                if overflow > 0:
                    print(
                        f"Annotation {child.id} is overflowing. Needs at least {overflow + args.padding} padding to fit.",
                        file=sys.stderr,
                    )
                    px = min(px, width)
                    py = min(py, height)
                    pw = min(pw, width - px - 1)
                    ph = min(ph, height - py - 1)

                rect = ((px, py), (px + pw, py + ph))
                annotation = Annotation.from_component(child.label(), rect, image_path)
                children.append(annotation)

            children_path = f"{args.outdir}/annotations/{filename}.json"
            with open(children_path, "w") as f:
                json.dump(Annotation.to_dicts(children), f)
