from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


# --------------------------------------------------
# 1. Load the data
# --------------------------------------------------
csv_path = Path("workspace_nocollisiondetection.csv")

# The file has no column names, so header=None is important.
df = pd.read_csv(
    csv_path,
    header=None,
    names=["x", "y", "z", "ik_rate"],
)

print(df.head())
print("z planes:", np.sort(df["z"].unique()))


# --------------------------------------------------
# 2. Create an output directory
# --------------------------------------------------
output_dir = Path("contour_plots")
output_dir.mkdir(exist_ok=True)


# --------------------------------------------------
# 3. Use the same colour scale for every z-plane
# --------------------------------------------------
value_min = df["ik_rate"].min()
value_max = df["ik_rate"].max()

# More levels give a smoother-looking contour plot.
levels = np.linspace(value_min, value_max, 51)


# --------------------------------------------------
# 4. Draw one contour plot for each z-plane
# --------------------------------------------------
for z_value, plane in df.groupby("z", sort=True):

    # Rows become y-coordinates, columns become x-coordinates.
    grid = plane.pivot(index="y", columns="x", values="ik_rate")

    x = grid.columns.to_numpy()
    y = grid.index.to_numpy()
    ik_grid = grid.to_numpy()

    X, Y = np.meshgrid(x, y)

    fig, ax = plt.subplots(figsize=(7, 6))

    contour = ax.contourf(
        X,
        Y,
        ik_grid,
        levels=levels,
        cmap="viridis",
        extend="both",
    )

    colorbar = fig.colorbar(contour, ax=ax)
    colorbar.set_label("IK rate")

    ax.set_title(f"IK rate at z = {z_value:.2f}")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_aspect("equal")

    fig.tight_layout()

    output_path = output_dir / f"ik_contour_z_{z_value:.2f}.png"
    fig.savefig(output_path, dpi=200, bbox_inches="tight")
    plt.close(fig)

print(f"Plots saved to: {output_dir.resolve()}")