"""Create the experiment-vs-simulation parabola comparison figure."""

from pathlib import Path

import cv2
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import AutoMinorLocator
from mpl_toolkits.axes_grid1.inset_locator import inset_axes


EXPERIMENT_ANALYSIS = Path("output_parabola_fits_experiment/frame_0267.npz")
SIMULATION_ANALYSIS = Path("output_parabola_fits_simulation/frame_0061.npz")
RENDERED_OUTPUT_DIR = Path("rendered_output")
OUTPUT_PNG = RENDERED_OUTPUT_DIR / "simulation_experiment_shape_comparison.png"
OUTPUT_PDF = RENDERED_OUTPUT_DIR / "simulation_experiment_shape_comparison.pdf"

EXPERIMENT_COLOR = "#8400FF"
SIMULATION_COLOR = "black"
STANDARD_FONT_SIZE = 20


def load_analysis(path: Path):
    """Read one saved parabola-fit analysis file."""
    with np.load(path, allow_pickle=False) as data:
        return {
            "orig_bgr": data["overlay_bgr"],
            "pts_left": data["pts_left"],
            "pts_right": data["pts_right"],
            "coef_left": data["coef_left"],
            "coef_right": data["coef_right"],
            "base_y": data["base_y"],
            "upper_y": data["upper_y"],
            "mid_img": data["mid_img"],
        }


def curve_x(coef, y_values):
    """Evaluate the fitted x=f(y) parabola."""
    return np.polyval(coef, y_values)


def configure_matplotlib():
    """Use publication-size text for the generated figure."""
    plt.rcParams["font.size"] = STANDARD_FONT_SIZE
    plt.rcParams["axes.titlesize"] = STANDARD_FONT_SIZE
    plt.rcParams["axes.labelsize"] = STANDARD_FONT_SIZE
    plt.rcParams["xtick.labelsize"] = STANDARD_FONT_SIZE
    plt.rcParams["ytick.labelsize"] = STANDARD_FONT_SIZE
    plt.rcParams["legend.fontsize"] = STANDARD_FONT_SIZE


def scaled_right_edge(data, y_scaled):
    """Scale the right edge by the initial diameter and height."""
    y_unscaled = data["upper_y"] + (data["base_y"] - data["upper_y"]) * y_scaled
    x_right_unscaled = curve_x(data["coef_right"], y_unscaled) + data["mid_img"]
    x_right_scaled = (x_right_unscaled - data["mid_img"]) / (
        x_right_unscaled[0] - data["mid_img"]
    )
    return x_right_unscaled, x_right_scaled


def add_inset(ax, data, color, label, loc, anchor, alpha=1.0):
    """Add an image inset with the fitted perimeter points overlaid."""
    height, width = data["orig_bgr"].shape[:2]
    inset = inset_axes(
        ax,
        width="40%",
        height="40%",
        loc=loc,
        bbox_to_anchor=anchor,
        bbox_transform=ax.transAxes,
        borderpad=0,
    )

    rgb = cv2.cvtColor(data["orig_bgr"], cv2.COLOR_BGR2RGB)
    inset.imshow(rgb, extent=[0, width, 0, height], alpha=alpha)
    inset.set_xlim(0, width)
    inset.set_ylim(0, height)
    inset.set_aspect("equal", "box")
    inset.axis("off")
    inset.plot(data["pts_right"][:, 0], height - data["pts_right"][:, 1], "D", color=color, ms=7)
    inset.text(
        0.02,
        1.02,
        label,
        transform=inset.transAxes,
        ha="left",
        va="bottom",
        color=color,
        fontweight="bold",
    )
    inset.tick_params(left=False, bottom=False)

    for spine in inset.spines.values():
        spine.set_visible(True)
        spine.set_linewidth(0.6)
        spine.set_color("black")


def main():
    RENDERED_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    experiment = load_analysis(EXPERIMENT_ANALYSIS)
    simulation = load_analysis(SIMULATION_ANALYSIS)
    configure_matplotlib()

    y_scaled = np.linspace(0.0, 1.0, 50)
    x_experiment_unscaled, x_experiment_scaled = scaled_right_edge(experiment, y_scaled)
    x_simulation_unscaled, x_simulation_scaled = scaled_right_edge(simulation, y_scaled)

    fig, ax = plt.subplots(figsize=(10.2, 7.8))
    ax.set_xlim(0, 2)
    ax.set_ylim(-1.0, 1.1)
    ax.set_ylim(ax.get_ylim()[::-1])
    ax.set_xlabel("Width scaled to diameter / -", fontsize=24)
    ax.set_ylabel("Scaled height", fontsize=24)
    ax.set_yticks([])
    ax.xaxis.set_minor_locator(AutoMinorLocator(2))
    ax.tick_params(direction="in", labelsize=24)
    ax.tick_params(which="minor", direction="in")
    ax.tick_params(axis="x", pad=12)
    ax.xaxis.labelpad = 16

    experiment_width_scale = x_experiment_unscaled[0] - experiment["mid_img"]
    simulation_width_scale = x_simulation_unscaled[0] - simulation["mid_img"]

    ax.plot(
        (experiment["pts_right"][:, 0] - experiment["mid_img"]) / experiment_width_scale,
        (experiment["pts_right"][:, 1] - experiment["upper_y"])
        / (experiment["base_y"] - experiment["upper_y"]),
        "D",
        color=EXPERIMENT_COLOR,
        ms=7,
        label="Perimeter points (Exp.)",
        zorder=1,
    )
    ax.plot(
        (simulation["pts_right"][:, 0] - simulation["mid_img"]) / simulation_width_scale,
        (simulation["pts_right"][:, 1] - simulation["upper_y"])
        / (simulation["base_y"] - simulation["upper_y"]),
        "D",
        color=SIMULATION_COLOR,
        ms=7,
        label="Perimeter points (Sim.)",
        zorder=2,
    )
    ax.plot(
        x_experiment_scaled,
        y_scaled,
        color=EXPERIMENT_COLOR,
        lw=4,
        label="Parabola (Exp.)",
        zorder=3,
    )
    ax.plot(
        x_simulation_scaled,
        y_scaled,
        color=SIMULATION_COLOR,
        lw=4,
        label="Parabola (Sim.)",
        zorder=4,
    )
    ax.hlines(0.0, 0, 1.0, color="black", lw=2, linestyle="-", label="Diameter", zorder=5)
    ax.hlines(
        1.0,
        0,
        max([x_experiment_scaled[-1], x_simulation_scaled[-1]]),
        color="black",
        lw=2,
        linestyle="--",
        label="Baseline",
        zorder=6,
    )

    ax.legend(loc="upper left", frameon=False, framealpha=1.0)
    add_inset(
        ax,
        experiment,
        EXPERIMENT_COLOR,
        "b) Strain: 0.32",
        loc="lower right",
        anchor=(0.0, 0.05, 1.0, 1.0),
    )
    add_inset(
        ax,
        simulation,
        SIMULATION_COLOR,
        "a) Strain: 0.31",
        loc="upper right",
        anchor=(0.0, -0.05, 1.0, 0.98),
        alpha=0.5,
    )

    fig.savefig(OUTPUT_PNG, dpi=900, bbox_inches="tight", pil_kwargs={"quality": 95})
    fig.savefig(OUTPUT_PDF, bbox_inches="tight")


if __name__ == "__main__":
    main()
