import matplotlib.pyplot as plt
from functions import *


def generate_picture_wp(all_group_vars, picture_path, DEBUG=False):
    vars = all_group_vars['wp']
    WPR_X, WPR_Y, WPR_Z, WPF_X, WPF_Y, WPF_Z, WP_CLP_OFFSET_X, WP_GRP_OFFSET_Z = (float(vars[k]) for k in
                                                                                  ['WPR_X', 'WPR_Y', 'WPR_Z', 'WPF_X',
                                                                                   'WPF_Y', 'WPF_Z', 'WP_CLP_OFFSET_X',
                                                                                   'WP_GRP_OFFSET_Z'])

    if min(WPR_X, WPR_Y, WPR_Z, WPF_X, WPF_Y, WPF_Z) <= 0:
        raise ValueError("Dimensions must be positive.")

    fig, ax = plt.subplots(subplot_kw={'projection': '3d'})

    # Define raw workpiece dimensions
    x_cube, y_cube, z_cube = 2 * WPR_X, 2 * WPR_Y, WPR_Z
    vertices_raw = np.array([[0, 0, 0], [x_cube, 0, 0], [x_cube, y_cube, 0], [0, y_cube, 0],
                             [0, 0, WPR_Z], [x_cube, 0, WPR_Z], [x_cube, y_cube, WPR_Z], [0, y_cube, WPR_Z]])
    faces_raw = [vertices_raw[[0, 1, 2, 3]], vertices_raw[[4, 5, 6, 7]], vertices_raw[[0, 3, 7, 4]],
                 vertices_raw[[1, 2, 6, 5]], vertices_raw[[0, 1, 5, 4]], vertices_raw[[2, 3, 7, 6]]]

    # Set axis properties
    ax.set_box_aspect([x_cube, y_cube, z_cube])
    ax.set(xticks=[], yticks=[], zticks=[])

    # Draw WP coordinate system
    axis_len = max(x_cube, y_cube, z_cube) * 0.2
    colors = ['r', 'g', 'b']
    origin_x = 0
    origin_y = 0
    origin_z = 0
    for i, vec in enumerate(np.eye(3)):
        ax.quiver(origin_x, origin_y, origin_z, *(axis_len * vec), color=colors[i], linewidth=3)
    ax.text(origin_x, origin_y, origin_z - axis_len * 0.3, 'WP', color='black', fontsize=9, weight='bold', ha='center',
            va='top')

    # Draw WP
    ax.add_collection3d(Poly3DCollection(faces_raw, facecolors='k', linewidths=1, edgecolors='k', alpha=.15))

    # Draw gripping point
    u, v = np.mgrid[0:2 * np.pi:20j, 0:np.pi:20j]
    ball_radius = min(WPR_X, WPR_Y, WPR_Z) * 0.1
    ax.plot_surface(
        WPR_X + ball_radius * np.cos(u) * np.sin(v),
        WPR_Y + ball_radius * np.sin(u) * np.sin(v),
        WPR_Z + ball_radius * np.cos(v),
        color='r'
    )

    # Draw gripper
    gripper_thickness = max(x_cube, y_cube, z_cube) * 0.1
    gripper_wide = WPR_X * 0.4
    GRP_ALPHA = .40
    bars = [
        ((x_cube - gripper_wide) * 0.5, -gripper_thickness, WPR_Z - WP_GRP_OFFSET_Z),
        ((x_cube - gripper_wide) * 0.5, y_cube, WPR_Z - WP_GRP_OFFSET_Z),
    ]
    for (x, y, z) in bars:
        ax.bar3d(
            x, y, z,
            gripper_wide, gripper_thickness, 5 * WPR_Z + WP_GRP_OFFSET_Z,
            color='r', alpha=GRP_ALPHA, edgecolor='k', linewidth=0.5
        )

    # Define finished workpiece dimensions
    x_cube2, y_cube2, z_cube2 = 2 * WPF_X, 2 * WPF_Y, WPF_Z

    # Draw gripping point for finished workpiece
    u, v = np.mgrid[0:2 * np.pi:20j, 0:np.pi:20j]
    ball_radius = min(WPF_X, WPF_Y, WPF_Z) * 0.08
    ax.plot_surface(
        WPF_X + ball_radius * np.cos(u) * np.sin(v),
        WPF_Y + ball_radius * np.sin(u) * np.sin(v),
        WPF_Z + ball_radius * np.cos(v),
        color='g'
    )

    # Draw gripper for finished workpiece
    gripper_thickness2 = max(x_cube, y_cube, z_cube) * 0.05
    gripper_wide2 = WPR_X * 0.4
    GRP_ALPHA2 = .25
    yyy2 = [-gripper_thickness2, y_cube2] if WPF_Y < WPR_Y else [-gripper_thickness2 + (WPF_Y - WPR_Y) * 2, y_cube]
    bars = [
        ((x_cube2 - gripper_wide2) * 0.5, yyy2[0], WPF_Z - WP_GRP_OFFSET_Z),
        ((x_cube2 - gripper_wide2) * 0.5, yyy2[1], WPF_Z - WP_GRP_OFFSET_Z),
    ]
    for (x, y, z) in bars:
        ax.bar3d(
            x, y, z,
            gripper_wide2, gripper_thickness2, 5 * WPF_Z + WP_GRP_OFFSET_Z,
            color='g', alpha=GRP_ALPHA2, edgecolor='k', linewidth=0.5
        )

    # Draw CLP coordinate system
    axis_len = max(x_cube, y_cube, z_cube) * 0.15
    colors = ['r', 'g', 'b']
    origin_x = WPR_X - WP_CLP_OFFSET_X
    origin_y = WPR_Y
    origin_z = 0
    for i, vec in enumerate(np.eye(3)):
        ax.quiver(origin_x, origin_y, origin_z, *(axis_len * vec), color=colors[i], linewidth=3)
    ax.text(origin_x, origin_y, origin_z - axis_len * 0.3, 'CLP', color='blue', fontsize=9, weight='bold', ha='center',
            va='top')

    # Define clamping block dimensions
    clp_height = z_cube * 0.2
    clp_width = 0.6 * x_cube
    clp_deep = y_cube * 0.3
    clp_supportdeep = 3
    base_x = WPR_X - clp_width / 2
    alpha = 0.35
    color = 'blue'
    linewidth = 1

    # Draw clamping blocks
    x = base_x - WP_CLP_OFFSET_X
    dx = clp_width
    y = 0
    dy1 = clp_deep
    dy2 = clp_supportdeep
    z = 0
    dz1 = z_cube * 0.2
    dz2 = z_cube * 0.6
    plot_clamping(ax, x, y, z, dx, dy1, dy2, dz1, dz2, color=color,
                  alpha=alpha, linewidth=linewidth)
    y = y_cube
    plot_clamping(ax, x, y, z, dx, dy1, dy2, dz1, dz2, color=color,
                  alpha=alpha, linewidth=linewidth, inverty=True)

    # Set view properties
    z_trans = WPR_Z*0.2
    zoom_factor = 1.1  # Adjust this factor to zoom out more or less
    ax.view_init(elev=30, azim=-160)
    ax.set_xlim([-x_cube * zoom_factor * 0.1, x_cube * zoom_factor])
    ax.set_ylim([-y_cube * zoom_factor * 0.1, y_cube * zoom_factor])
    ax.set_zlim([-WPR_Z * zoom_factor * 0.1 - z_trans, WPR_Z * zoom_factor - z_trans])
    ax.set_aspect('auto')
    ax.grid(False)  # Remove grid lines
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])
    for axis in [ax.xaxis, ax.yaxis, ax.zaxis]:
        axis.set_pane_color((1, 1, 1, 0))  # Hide background panes
        axis.line.set_color((1, 1, 1, 0))  # Hide axis lines

    # Save or display the plot
    if not DEBUG:
        fig.savefig(picture_path, bbox_inches='tight', dpi=300, transparent=True)
    else:
        plt.show()
    plt.close()


if __name__ == "__main__":
    vars = {'WPR_X': 20,
            'WPR_Y': 15,
            'WPR_Z': 20,
            'WPF_X': 10,
            'WPF_Y': 18,
            'WPF_Z': 15,
            'WP_CLP_OFFSET_X': 3,
            'WP_GRP_OFFSET_Z': 5}
    generate_picture_wp({'wp': vars}, "cfg_picture/TEST_wpraw.jpg", DEBUG=True)