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

df = pd.read_csv('/home/lgv/cmvr/cmvr-es/data/ik_psi_sweep.csv')

joints = ['q1','q2','q3','q4','q5','q6','q7']
fig, axes = plt.subplots(len(joints), 1, figsize=(10, 2.6*len(joints)), sharex=True)
x = df['psi'].to_numpy()

for ax, col in zip(axes, joints):
    y = df[col].to_numpy()

    # 过滤 NaN
    mask = ~np.isnan(x) & ~np.isnan(y)
    xv, yv = x[mask], y[mask]
    if len(xv) == 0:
        ax.set_title(f'{col}: no data')
        continue

    # 散点图（不画连线）
    ax.scatter(xv, yv, s=10, alpha=0.8)  # s 调整点大小

    # 极值与对应 psi
    i_min = int(np.argmin(yv))
    i_max = int(np.argmax(yv))
    y_min, psi_min = float(yv[i_min]), float(xv[i_min])
    y_max, psi_max = float(yv[i_max]), float(xv[i_max])

    # 水平红虚线（值）
    ax.axhline(y_min, color='red', linestyle='--', linewidth=1)
    ax.axhline(y_max, color='red', linestyle='--', linewidth=1)
    # 垂直红虚线（psi）
    ax.axvline(psi_min, color='red', linestyle='--', linewidth=1, alpha=0.7)
    ax.axvline(psi_max, color='red', linestyle='--', linewidth=1, alpha=0.7)

    # 极值点标记
    ax.plot(psi_min, y_min, 'ro', markersize=4)
    ax.plot(psi_max, y_max, 'ro', markersize=4)

    # 标注：值 + psi
    ax.annotate(f"min={y_min:.4f}\nψ={psi_min:.4f}",
                xy=(psi_min, y_min), xytext=(8, -10),
                textcoords='offset points', color='red',
                ha='left', va='top', fontsize=9,
                bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none"))
    ax.annotate(f"max={y_max:.4f}\nψ={psi_max:.4f}",
                xy=(psi_max, y_max), xytext=(8, 10),
                textcoords='offset points', color='red',
                ha='left', va='bottom', fontsize=9,
                bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none"))

    ax.set_ylabel(f"{col} (rad)")
    ax.grid(True, alpha=0.35)

axes[-1].set_xlabel('psi (rad)')
fig.suptitle('IK joints vs arm-angle psi', y=0.995)
plt.tight_layout()
plt.show()