import copy
import pathlib

import cv2
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from pipeline import pipeline_conf

import iris


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

    print("Loading pipeline...")
    iris_pipeline = iris.IRISPipeline(
        env=iris.IRISPipeline.DEBUGGING_ENVIRONMENT, config=pipeline_conf
    )
    iris_visualizer = iris.visualisation.IRISVisualizer()
    print("Pipeline loaded")

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

        print(f"Processing {image_path}...")
        try:
            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(
                iris.IRImage(
                    img_data=image_pixels, image_id="image_id", eye_side="right"
                )
            )

            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")
            plt.close()

            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")
            plt.close()

            (path / "normalized").mkdir(parents=True, exist_ok=True)
            data = iris_pipeline.call_trace["normalization"].normalized_image
            plt.imsave(
                path / "normalized" / f"{image_path.stem}.png", data, cmap="Greys_r"
            )

            cm = copy.copy(mpl.colormaps["binary"])
            cm.set_bad("red")

            data = output["iris_template"].iris_codes[0][:, :, 0]
            mask = np.invert(output["iris_template"].mask_codes[0][:, :, 0])
            data = np.ma.masked_array(data, mask=mask)

            (path / "template").mkdir(parents=True, exist_ok=True)
            plt.imsave(path / "template" / f"{image_path.stem}.png", data, cmap=cm)

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

            print(f"Finished processing: {image_path}")

        except:
            print(
                f"Failed to process: {image_path}. Most likely not centered, out of focus, or eye not opened enough"
            )

    print("Done")
