#!/usr/bin/env python3

import argparse
import csv
from dataclasses import dataclass
from pathlib import Path

import cv2


TOLERANCE_PIXELS = 20


@dataclass
class Annotation:
    class_name: str
    x: float
    y: float


@dataclass
class GroundTruth:
    image_name: str
    image_width: int
    image_height: int
    annotations: list[Annotation]


def read_ground_truth(csv_path: str) -> GroundTruth:
    """
    Read a CSV with columns:

        class_name, x, y, image_name, image_width, image_height
    """

    annotations = []
    image_names = set()
    image_widths = set()
    image_heights = set()

    with open(csv_path, "r", newline="", encoding="utf-8") as file:
        reader = csv.reader(file)

        for line_number, row in enumerate(reader, start=1):
            if not row or all(not value.strip() for value in row):
                continue

            if len(row) < 6:
                raise ValueError(
                    f"{csv_path}, line {line_number}: "
                    f"expected 6 columns, found {len(row)}"
                )

            class_name = row[0].strip()
            image_name = row[3].strip()

            try:
                x = float(row[1])
                y = float(row[2])
                image_width = int(float(row[4]))
                image_height = int(float(row[5]))
            except ValueError as error:
                raise ValueError(
                    f"{csv_path}, line {line_number}: invalid numeric value"
                ) from error

            annotations.append(
                Annotation(
                    class_name=class_name,
                    x=x,
                    y=y,
                )
            )

            image_names.add(image_name)
            image_widths.add(image_width)
            image_heights.add(image_height)

    if not annotations:
        raise ValueError(f"{csv_path} contains no annotations")

    if len(image_names) != 1:
        raise ValueError(
            f"{csv_path} must contain annotations for exactly one image. "
            f"Found: {sorted(image_names)}"
        )

    if len(image_widths) != 1 or len(image_heights) != 1:
        raise ValueError(
            f"{csv_path} contains inconsistent image dimensions"
        )

    return GroundTruth(
        image_name=next(iter(image_names)),
        image_width=next(iter(image_widths)),
        image_height=next(iter(image_heights)),
        annotations=annotations,
    )


def squared_distance(a: Annotation, b: Annotation) -> float:
    return (a.x - b.x) ** 2 + (a.y - b.y) ** 2


def match_annotations(
    annotations_a: list[Annotation],
    annotations_b: list[Annotation],
    tolerance: float,
):
    """
    Match points one-to-one using nearest spatial locations.

    A point is considered spatially matched when its Euclidean distance
    from the other point is no greater than the specified tolerance.
    """

    possible_matches = []

    for index_a, annotation_a in enumerate(annotations_a):
        for index_b, annotation_b in enumerate(annotations_b):
            distance_squared = squared_distance(annotation_a, annotation_b)

            if distance_squared <= tolerance * tolerance:
                possible_matches.append(
                    (distance_squared, index_a, index_b)
                )

    # Match closest pairs first.
    possible_matches.sort(key=lambda item: item[0])

    matched_a = set()
    matched_b = set()
    matches = []

    for _, index_a, index_b in possible_matches:
        if index_a in matched_a or index_b in matched_b:
            continue

        matched_a.add(index_a)
        matched_b.add(index_b)

        matches.append(
            (annotations_a[index_a], annotations_b[index_b])
        )

    unmatched_a = [
        annotation
        for index, annotation in enumerate(annotations_a)
        if index not in matched_a
    ]

    unmatched_b = [
        annotation
        for index, annotation in enumerate(annotations_b)
        if index not in matched_b
    ]

    return matches, unmatched_a, unmatched_b


def draw_unmatched(
    image,
    annotation: Annotation,
    label: str,
):
    """
    Draw a red marker for a point found in only one CSV.
    """

    point = (round(annotation.x), round(annotation.y))

    # Red circle and cross.
    cv2.circle(image, point, 14, (0, 0, 255), 3)

    cv2.drawMarker(
        image,
        point,
        (0, 0, 255),
        markerType=cv2.MARKER_CROSS,
        markerSize=28,
        thickness=2,
    )

    text = f"{label}: {annotation.class_name}"

    cv2.putText(
        image,
        text,
        (point[0] + 12, point[1] - 12),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6,
        (0, 0, 255),
        2,
        cv2.LINE_AA,
    )


def draw_class_mismatch(
    image,
    annotation_a: Annotation,
    annotation_b: Annotation,
):
    """
    Draw a black marker when locations match but class names differ.
    """

    point_a = (round(annotation_a.x), round(annotation_a.y))
    point_b = (round(annotation_b.x), round(annotation_b.y))

    # Black line connecting the two locations.
    cv2.line(
        image,
        point_a,
        point_b,
        (0, 0, 0),
        2,
    )

    # Black circles at both locations.
    cv2.circle(image, point_a, 14, (0, 0, 0), 3)
    cv2.circle(image, point_b, 14, (0, 0, 0), 3)

    midpoint = (
        (point_a[0] + point_b[0]) // 2,
        (point_a[1] + point_b[1]) // 2,
    )

    text = f"{annotation_a.class_name} -> {annotation_b.class_name}"

    cv2.putText(
        image,
        text,
        (midpoint[0] + 12, midpoint[1] - 12),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6,
        (0, 0, 0),
        2,
        cv2.LINE_AA,
    )


def draw_legend(image):
    """
    Draw a legend in the upper-left corner.
    """

    # Red legend entry for unmatched points.
    cv2.circle(image, (25, 30), 10, (0, 0, 255), 3)

    cv2.putText(
        image,
        "Unmatched point",
        (45, 36),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6,
        (0, 0, 255),
        2,
        cv2.LINE_AA,
    )

    # Black legend entry for class mismatches.
    cv2.circle(image, (25, 65), 10, (0, 0, 0), 3)

    cv2.putText(
        image,
        "Class mismatch",
        (45, 71),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.6,
        (0, 0, 0),
        2,
        cv2.LINE_AA,
    )


def main():
    parser = argparse.ArgumentParser(
        description="Compare two image ground-truth CSV files."
    )

    parser.add_argument(
        "csv_a",
        help="First ground-truth CSV file",
    )

    parser.add_argument(
        "csv_b",
        help="Second ground-truth CSV file",
    )

    parser.add_argument(
        "-t",
        "--tolerance",
        type=float,
        default=TOLERANCE_PIXELS,
        help=f"Maximum point-matching distance in pixels "
             f"(default: {TOLERANCE_PIXELS})",
    )

    args = parser.parse_args()

    ground_truth_a = read_ground_truth(args.csv_a)
    ground_truth_b = read_ground_truth(args.csv_b)

    if ground_truth_a.image_name != ground_truth_b.image_name:
        raise ValueError(
            "The two CSV files refer to different images:\n"
            f"  CSV A: {ground_truth_a.image_name}\n"
            f"  CSV B: {ground_truth_b.image_name}"
        )

    image_name = ground_truth_a.image_name

    # Load the original image using the name from the CSV files.
    original_image = cv2.imread(image_name)

    if original_image is None:
        raise FileNotFoundError(
            f"Could not load the original image: {image_name}\n"
            "The image must be in the current directory, or the CSV must "
            "contain a path that can be opened by OpenCV."
        )

    # Create a copy of the original image for drawing.
    output_image = original_image.copy()

    matches, unmatched_a, unmatched_b = match_annotations(
        ground_truth_a.annotations,
        ground_truth_b.annotations,
        args.tolerance,
    )

    # Draw red markers for spatially unmatched points.
    for annotation in unmatched_a:
        draw_unmatched(output_image, annotation, "A")

    for annotation in unmatched_b:
        draw_unmatched(output_image, annotation, "B")

    # Draw black markers for class-name mismatches.
    class_mismatch_count = 0

    for annotation_a, annotation_b in matches:
        if annotation_a.class_name != annotation_b.class_name:
            draw_class_mismatch(
                output_image,
                annotation_a,
                annotation_b,
            )
            class_mismatch_count += 1

    draw_legend(output_image)

    # Use the second CSV filename, replacing its extension with .png.
    # The output is written to the current directory.
    output_filename = Path(args.csv_b).stem + ".png"

    if not cv2.imwrite(output_filename, output_image):
        raise IOError(
            f"Could not write output image: {output_filename}"
        )

    print(f"Original image: {image_name}")
    print(f"Output image:   {output_filename}")
    print(f"Tolerance:      {args.tolerance} pixels")
    print(f"Unmatched points in CSV A: {len(unmatched_a)}")
    print(f"Unmatched points in CSV B: {len(unmatched_b)}")
    print(f"Class mismatches:          {class_mismatch_count}")


if __name__ == "__main__":
    main()