import pathlib

import cv2
import matplotlib.pyplot as plt

import iris


def process(path: str) -> None:
    path = pathlib.Path(path)

    iris_pipeline = iris.IRISPipeline(env=iris.IRISPipeline.DEBUGGING_ENVIRONMENT)
    iris_visualizer = iris.visualisation.IRISVisualizer()

    outputs = []
    for image_path in (path / "input").iterdir():
        if image_path.suffix != ".png":
            continue

        print(f"Processing {image_path}...")

        eye_side = "right" if "left" not in image_path.stem else "left"

        image_pixels = cv2.imread(str(image_path), cv2.IMREAD_GRAYSCALE)
        img = iris.IRImage(
            img_data=image_pixels,
            eye_side=eye_side,
        )

        output = iris_pipeline(image_pixels, eye_side=eye_side)

        iris_visualizer.plot_segmentation_map(
            ir_image=img,
            segmap=iris_pipeline.call_trace["segmentation"],
        )
        (path / "segmentation").mkdir(parents=True, exist_ok=True)
        plt.savefig(path / "segmentation" / f"{image_path.stem}.png")

        iris_visualizer.plot_all_geometry(
            ir_image=img,
            geometry_polygons=iris_pipeline.call_trace["geometry_estimation"],
            eye_orientation=iris_pipeline.call_trace["eye_orientation"],
            eye_center=iris_pipeline.call_trace["eye_center_estimation"],
        )
        (path / "geometry").mkdir(parents=True, exist_ok=True)
        plt.savefig(path / "geometry" / f"{image_path.stem}.png")

        iris_visualizer.plot_normalized_iris(
            normalized_iris=iris_pipeline.call_trace["normalization"],
        )
        (path / "normalized").mkdir(parents=True, exist_ok=True)
        plt.savefig(path / "normalized" / f"{image_path.stem}.png")

        iris_visualizer.plot_iris_template(
            iris_template=iris_pipeline.call_trace["encoder"],
        )
        (path / "template").mkdir(parents=True, exist_ok=True)
        plt.savefig(path / "template" / f"{image_path.stem}.png")

        outputs.append((image_path, img, output))
