import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D, proj3d
from lokky.pionmath import SSolver  # import the SSolver class from the lokky package

def main():
    # Number of points
    n_points = 10

    # Parameters for SSolver
    params = {
        "kp": np.ones((n_points, 6)),
        "ki": np.zeros((n_points, 6)),
        "kd": np.ones((n_points, 6)) * 0,
        "attraction_weight": 1.0,
        "cohesion_weight": 1.0,
        "alignment_weight": 1.0,
        "repulsion_weight": 1.0,
        "unstable_weight": 1.0,
        "noise_weight": 1.0,
        "safety_radius": 1.0,
        "max_acceleration": 1.0,
        "max_speed": 1.0,
        "unstable_radius": 2,
    }
    solver = SSolver(params)

    # Generate random positions and velocities
    np.random.seed(42)
    positions = np.random.rand(n_points, 3) * np.array([10, 10, 0])
    velocities = np.zeros((n_points, 3))

    # Form the state_matrix (positions + velocities)
    state_matrix = np.hstack([positions, velocities])

    # Form the target_matrix (target positions with an offset)
    target_positions = positions + np.random.rand(n_points, 3) * 1.0
    target_matrix = np.hstack([target_positions, np.zeros((n_points, 3))])

    # Time step
    dt = 0.1

    # Compute control signals
    control_signals = solver.solve(state_matrix, target_matrix, dt)

    # Create a figure with 4 subplots: Top, Side, Front, and 3D view
    fig = plt.figure(figsize=(12, 10))
    ax_top = fig.add_subplot(221)      # Top view (XY)
    ax_front = fig.add_subplot(222)    # Front view (YZ)
    ax_side = fig.add_subplot(223)     # Side view (XZ)
    ax_3d = fig.add_subplot(224, projection="3d")  # 3D view

    # Disable 3D navigation
    ax_3d.set_navigate(False)
    ax_3d.mouse_init = lambda: None

    # Set up 2D axes
    ax_top.set_title("Top view (XY)")
    ax_top.set_xlabel("X")
    ax_top.set_ylabel("Y")
    ax_top.set_xlim(-2, 12)
    ax_top.set_ylim(-2, 12)
    ax_top.grid(True)

    ax_side.set_title("Side view (XZ)")
    ax_side.set_xlabel("X")
    ax_side.set_ylabel("Z")
    ax_side.set_xlim(-2, 12)
    ax_side.set_ylim(-2, 12)
    ax_side.grid(True)

    ax_front.set_title("Front view (YZ)")
    ax_front.set_xlabel("Y")
    ax_front.set_ylabel("Z")
    ax_front.set_xlim(-2, 12)
    ax_front.set_ylim(-2, 12)
    ax_front.grid(True)

    # Set up 3D axes
    ax_3d.set_title("3D view")
    ax_3d.set_xlabel("X")
    ax_3d.set_ylabel("Y")
    ax_3d.set_zlabel("Z")
    ax_3d.set_xlim(-2, 12)
    ax_3d.set_ylim(-2, 12)
    ax_3d.set_zlim(-2, 12)
    ax_3d.grid(True)
    ax_3d.view_init(elev=30, azim=-60)

    # Create scatter plots for all views
    state_scatter_top = ax_top.scatter(positions[:, 0], positions[:, 1], color="black", label="State")
    target_scatter_top = ax_top.scatter(target_positions[:, 0], target_positions[:, 1],
                                        color="blue", marker="x", label="Target")

    state_scatter_side = ax_side.scatter(positions[:, 0], positions[:, 2], color="black", label="State")
    target_scatter_side = ax_side.scatter(target_positions[:, 0], target_positions[:, 2],
                                          color="blue", marker="x", label="Target")

    state_scatter_front = ax_front.scatter(positions[:, 1], positions[:, 2], color="black", label="State")
    target_scatter_front = ax_front.scatter(target_positions[:, 1], target_positions[:, 2],
                                            color="blue", marker="x", label="Target")

    state_scatter_3d = ax_3d.scatter(positions[:, 0], positions[:, 1], positions[:, 2],
                                     color="black", label="State")
    target_scatter_3d = ax_3d.scatter(target_positions[:, 0], target_positions[:, 1],
                                      target_positions[:, 2], color="blue", marker="x", label="Target")

    # Lists for arrows in 2D views
    error_arrows_top, ctrl_arrows_top = [], []
    error_arrows_side, ctrl_arrows_side = [], []
    error_arrows_front, ctrl_arrows_front = [], []

    # Lists for quiver arrows in the 3D view
    error_quivers_3d, ctrl_quivers_3d = [], []

    # Local variables for the selected point and active view
    selected_index = None
    selected_view = None  # can be "top", "side", "front", or "3d"
    last_y = None       # for vertical offsets in 3D

    # Functions to draw arrows and update scatter plots
    def draw_arrows_2d():
        nonlocal error_arrows_top, ctrl_arrows_top, error_arrows_side, ctrl_arrows_side, error_arrows_front, ctrl_arrows_front
        # Remove old arrows
        for arr in error_arrows_top + ctrl_arrows_top:
            arr.remove()
        for arr in error_arrows_side + ctrl_arrows_side:
            arr.remove()
        for arr in error_arrows_front + ctrl_arrows_front:
            arr.remove()
        error_arrows_top.clear()
        ctrl_arrows_top.clear()
        error_arrows_side.clear()
        ctrl_arrows_side.clear()
        error_arrows_front.clear()
        ctrl_arrows_front.clear()

        for i in range(n_points):
            err = target_positions[i] - positions[i]
            ctrl = control_signals[i, :3]
            # Top view (XY)
            a = ax_top.arrow(positions[i, 0], positions[i, 1],
                             err[0], err[1],
                             color="orange", width=0.05, head_width=0.3,
                             alpha=0.6, linestyle="dashed")
            error_arrows_top.append(a)
            a = ax_top.arrow(positions[i, 0], positions[i, 1],
                             ctrl[0], ctrl[1],
                             color="red", width=0.05, head_width=0.3,
                             alpha=0.8)
            ctrl_arrows_top.append(a)
            # Side view (XZ)
            a = ax_side.arrow(positions[i, 0], positions[i, 2],
                              err[0], err[2],
                              color="orange", width=0.05, head_width=0.3,
                              alpha=0.6, linestyle="dashed")
            error_arrows_side.append(a)
            a = ax_side.arrow(positions[i, 0], positions[i, 2],
                              ctrl[0], ctrl[2],
                              color="red", width=0.05, head_width=0.3,
                              alpha=0.8)
            ctrl_arrows_side.append(a)
            # Front view (YZ)
            a = ax_front.arrow(positions[i, 1], positions[i, 2],
                               err[1], err[2],
                               color="orange", width=0.05, head_width=0.3,
                               alpha=0.6, linestyle="dashed")
            error_arrows_front.append(a)
            a = ax_front.arrow(positions[i, 1], positions[i, 2],
                               ctrl[1], ctrl[2],
                               color="red", width=0.05, head_width=0.3,
                               alpha=0.8)
            ctrl_arrows_front.append(a)
        fig.canvas.draw_idle()

    def draw_arrows_3d():
        nonlocal error_quivers_3d, ctrl_quivers_3d
        for q in error_quivers_3d + ctrl_quivers_3d:
            q.remove()
        error_quivers_3d.clear()
        ctrl_quivers_3d.clear()
        for i in range(n_points):
            err = target_positions[i] - positions[i]
            ctrl = control_signals[i, :3]
            q = ax_3d.quiver(positions[i, 0], positions[i, 1], positions[i, 2],
                             err[0], err[1], err[2],
                             color="orange", arrow_length_ratio=0.1)
            error_quivers_3d.append(q)
            q = ax_3d.quiver(positions[i, 0], positions[i, 1], positions[i, 2],
                             ctrl[0], ctrl[1], ctrl[2],
                             color="red", arrow_length_ratio=0.1)
            ctrl_quivers_3d.append(q)
        fig.canvas.draw_idle()

    def update_scatter():
        state_scatter_top.set_offsets(positions[:, [0, 1]])
        target_scatter_top.set_offsets(target_positions[:, [0, 1]])

        state_scatter_side.set_offsets(positions[:, [0, 2]])
        target_scatter_side.set_offsets(target_positions[:, [0, 2]])

        state_scatter_front.set_offsets(positions[:, [1, 2]])
        target_scatter_front.set_offsets(target_positions[:, [1, 2]])

        # For 3D scatter, update data manually:
        state_scatter_3d._offsets3d = (positions[:, 0], positions[:, 1], positions[:, 2])
        target_scatter_3d._offsets3d = (target_positions[:, 0], target_positions[:, 1], target_positions[:, 2])
        fig.canvas.draw_idle()

    # Helper functions for 3D projection
    def get_2d_coords(ax, x, y, z):
        """Projects 3D coordinates to 2D screen coordinates."""
        x2, y2, _ = proj3d.proj_transform(x, y, z, ax.get_proj())
        return ax.transData.transform((x2, y2))

    def get_3d_from_2d(ax, x_pixel, y_pixel, z_fixed):
        """Approximate conversion of 2D coordinates to 3D (on a plane with fixed z)."""
        inv = ax.transData.inverted()
        x2d, y2d = inv.transform((x_pixel, y_pixel))
        return x2d, y2d

    # Mouse event handlers
    def on_press(event):
        nonlocal selected_index, selected_view, last_y
        if event.inaxes not in [ax_top, ax_side, ax_front, ax_3d]:
            return
        click = np.array([event.xdata, event.ydata])
        distances = []
        if event.inaxes == ax_top:
            for pos in positions:
                distances.append(np.hypot(pos[0] - click[0], pos[1] - click[1]))
            current_view = "top"
        elif event.inaxes == ax_side:
            for pos in positions:
                distances.append(np.hypot(pos[0] - click[0], pos[2] - click[1]))
            current_view = "side"
        elif event.inaxes == ax_front:
            for pos in positions:
                distances.append(np.hypot(pos[1] - click[0], pos[2] - click[1]))
            current_view = "front"
        elif event.inaxes == ax_3d:
            click_screen = np.array([event.x, event.y])
            for pos in positions:
                proj = get_2d_coords(ax_3d, pos[0], pos[1], pos[2])
                distances.append(np.hypot(proj[0] - click_screen[0], proj[1] - click_screen[1]))
            current_view = "3d"
        distances = np.array(distances)
        threshold = 15 if current_view == "3d" else 0.5
        if distances.min() < threshold:
            selected_index = int(np.argmin(distances))
            selected_view = current_view
            last_y = event.y

    def on_motion(event):
        nonlocal selected_index, positions, state_matrix, control_signals, last_y
        if selected_index is None or event.inaxes is None:
            return

        if event.inaxes == ax_top and selected_view == "top":
            positions[selected_index, 0] = event.xdata
            positions[selected_index, 1] = event.ydata
            state_matrix[selected_index, :2] = [event.xdata, event.ydata]
        elif event.inaxes == ax_side and selected_view == "side":
            positions[selected_index, 0] = event.xdata
            positions[selected_index, 2] = event.ydata
            state_matrix[selected_index, 0] = event.xdata
            state_matrix[selected_index, 2] = event.ydata
        elif event.inaxes == ax_front and selected_view == "front":
            positions[selected_index, 1] = event.xdata
            positions[selected_index, 2] = event.ydata
            state_matrix[selected_index, 1] = event.xdata
            state_matrix[selected_index, 2] = event.ydata
        elif event.inaxes == ax_3d and selected_view == "3d":
            if event.key is not None and ("shift" in event.key.lower() or "control" in event.key.lower()):
                if last_y is not None and event.y is not None:
                    dy = event.y - last_y
                    scale = 0.05
                    modifier = 1 if "shift" in event.key.lower() else -1
                    positions[selected_index, 2] += dy * scale * modifier
                    state_matrix[selected_index, 2] = positions[selected_index, 2]
            else:
                new_x, new_y = get_3d_from_2d(ax_3d, event.x, event.y, positions[selected_index, 2])
                positions[selected_index, 0] = new_x
                positions[selected_index, 1] = new_y
                state_matrix[selected_index, :2] = [new_x, new_y]
        control_signals[:] = solver.solve(state_matrix, target_matrix, dt)
        update_scatter()
        draw_arrows_2d()
        draw_arrows_3d()
        last_y = event.y

    def on_release(event):
        nonlocal selected_index, selected_view
        selected_index = None
        selected_view = None

    # Connect event handlers to the axes
    for ax in [ax_top, ax_side, ax_front, ax_3d]:
        ax.figure.canvas.mpl_connect("button_press_event", on_press)
        ax.figure.canvas.mpl_connect("motion_notify_event", on_motion)
        ax.figure.canvas.mpl_connect("button_release_event", on_release)

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()
