from pathlib import Path

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


# --------------------------------------------------
# 1. Load the data
# --------------------------------------------------
file_name = "workspace minisci collision.csv"
csv_path = Path(file_name)

# The file has no column names, so header=None is important.
df_csv = pd.read_csv(
    csv_path,
)
rate_res = df_csv.iloc[:, :4]
try:
    rate_res_sort = rate_res.sort_values('z').reset_index(drop=True)
except:
    rate_res_sort = rate_res

DECIMALS = 4

try:
    x_unique = np.round(rate_res_sort['x'], DECIMALS).unique()
    y_unique = np.round(rate_res_sort['y'], DECIMALS).unique()
    z_unique = np.round(rate_res_sort['z'], DECIMALS).unique()
except:
    x_unique = np.round(rate_res_sort.iloc[:,0], DECIMALS).unique()
    y_unique = np.round(rate_res_sort.iloc[:, 1], DECIMALS).unique()
    z_unique = np.round(rate_res_sort.iloc[:, 2], DECIMALS).unique()

nx, ny, nz = len(x_unique), len(y_unique), len(z_unique)


ik_rates = rate_res_sort.to_numpy()

# --------------------------------------------------
# 2. Create an output directory
# --------------------------------------------------
output_dir = Path(file_name.split(".")[0])
output_dir.mkdir(exist_ok=True)


# --------------------------------------------------
# 3. Use the same colour scale for every z-plane
# --------------------------------------------------
df = rate_res_sort
value_min = df["ik_success_rate"].min()
value_max = df["ik_success_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_success_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",
    )

    highlight_levels = [0.6, 0.7]
    # Only plot if the levels are within the data range (optional)
    if value_min <= 0.6 <= value_max or value_min <= 0.7 <= value_max:
        lines = ax.contour(X, Y, ik_grid, levels=highlight_levels,
                           colors='red', linewidths=2, linestyles='solid')
        # Optionally label the lines
        ax.clabel(lines, inline=True, fontsize=10, fmt='%1.1f')


    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()}")