import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np


def generate_picture_clp(all_group_vars, picture_path, DEBUG=False):
    try:
        # Get x_wpraw, y_wpraw, and z_wpraw values
        vars = all_group_vars['wp_raw']
        x_wpraw = float(vars['x_wpraw'])
        y_wpraw = float(vars['y_wpraw'])
        z_wpraw = float(vars['z_wpraw'])
        x_cube = x_wpraw*2
        y_cube = y_wpraw*2
        z_cube = z_wpraw
        clp_offset = float(all_group_vars['clp']['clp_offset'])

        # ERROR HANDLING
        if x_cube <= 0 or y_cube <= 0 or z_cube <= 0:
            raise ValueError("x_wpraw, y_wpraw, and z_wpraw must be positive values.")

        # Define the vertices of the cube based on wpraw dimensions
        vertices = (np.array([[0, 0, 0], [x_cube, 0, 0], [x_cube, y_cube, 0], [0, y_cube, 0],
                             [0, 0, z_cube], [x_cube, 0, z_cube], [x_cube, y_cube, z_cube], [0, y_cube, z_cube]])
                    + np.array([0, -clp_offset, 0]))

        # Define the 6 faces of the cube
        faces = [[vertices[j] for j in [0, 1, 2, 3]],  # Bottom face
                 [vertices[j] for j in [4, 5, 6, 7]],  # Top face
                 [vertices[j] for j in [0, 3, 7, 4]],  # Left face
                 [vertices[j] for j in [1, 2, 6, 5]],  # Right face
                 [vertices[j] for j in [0, 1, 5, 4]],  # Front face
                 [vertices[j] for j in [2, 3, 7, 6]]]  # Back face

        # Plot the cube in 3D
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        # Plot the cube with the specified faces
        ax.add_collection3d(Poly3DCollection(faces, facecolors='k', linewidths=1, edgecolors='k', alpha=.15))

        # Set the aspect ratio for each axis
        ax.set_box_aspect([x_cube, y_cube, z_cube])

        # Hide the axis ticks and labels
        ax.set_xticks([])  # Hide X axis ticks
        ax.set_yticks([])  # Hide Y axis ticks
        ax.set_zticks([])  # Hide Z axis ticks

        # Add a ball (sphere) at the center of the cube
        ball_radius = min(x_wpraw, y_wpraw, z_wpraw) * 0.1  # Size relative to cube dimensions

        u = np.linspace(0, 2 * np.pi, 20)
        v = np.linspace(0, np.pi, 20)
        x_sphere = x_wpraw + ball_radius * np.outer(np.cos(u), np.sin(v))
        y_sphere = y_wpraw - clp_offset + ball_radius * np.outer(np.sin(u), np.sin(v))
        z_sphere = z_wpraw + ball_radius * np.outer(np.ones(np.size(u)), np.cos(v))
        ax.plot_surface(x_sphere, y_sphere, z_sphere, color='red', alpha=1)

        # Add the origin sparrow (a small 3D marker at the origin)
        max_size = max(x_cube, y_cube, z_cube) * 0.2  # Scale the arrows relative to the smallest dimension
        ax.quiver(0, -clp_offset, 0, max_size, 0, 0, color='red', linewidth=2, label='X-axis')  # Red arrow for X-axis
        ax.quiver(0, -clp_offset, 0, 0, max_size, 0, color='green', linewidth=2, label='Y-axis')  # Green arrow for Y-axis
        ax.quiver(0, -clp_offset, 0, 0, 0, max_size, color='blue', linewidth=2, label='Z-axis')  # Blue arrow for Z-axis

        # Add two black blocks (greifer fingers) on the Y-faces in the middle of the cube height
        max_size2 = max(x_cube, y_cube, z_cube) * 0.4  # Scale the arrows relative to the smallest dimension
        ax.bar3d(-max_size2, -(2.0-1)/2 * y_cube, -z_wpraw * 0.0,
                 max_size2, 2.0 * y_cube, z_wpraw * 0.2,
                 color='black', alpha=.10)
        ax.bar3d(-max_size2, -(2.0-1)/2 * y_cube, -z_wpraw * 0.0,
                 max_size2+3, 2.0 * y_cube, -z_wpraw * 0.2,
                 color='black', alpha=.10)

        ax.bar3d(x_cube, -(2-1)/2 * y_cube, -z_wpraw * 0.0,
                 max_size2, 2.0 * y_cube, z_wpraw * 0.2,
                 color='black', alpha=.10)
        ax.bar3d(x_cube-3, -(2-1)/2 * y_cube, -z_wpraw * 0.0,
                 max_size2+3, 2.0 * y_cube, -z_wpraw * 0.2,
                 color='black', alpha=.10)

        # Add the origin sparrow (a small 3D marker at the origin)
        max_size3 = max(x_cube, y_cube, z_cube) * 0.3  # Scale the arrows relative to the smallest dimension
        ax.quiver(x_wpraw, y_wpraw, 0, max_size3, 0, 0, color='red', linewidth=3, label='X-axis')  # Red arrow for X-axis
        ax.quiver(x_wpraw, y_wpraw, 0, 0, max_size3, 0, color='green', linewidth=3, label='Y-axis')  # Green arrow for Y-axis
        ax.quiver(x_wpraw, y_wpraw, 0, 0, 0, max_size3, color='blue', linewidth=3, label='Z-axis')  # Blue arrow for Z-axis

        # Set viewing angle
        ax.view_init(elev=30, azim=-150)  # Set elevation to 24° and azimuth to -151°
        # ax.view_init(elev=90, azim=180)  # Set elevation to 24° and azimuth to -151°

        # Set limits for the axes
        # ax.set_xlim([0, x_cube])
        # ax.set_ylim([0, y_cube])
        # ax.set_zlim([0, z_cube])

        # Set equal aspect ratio to make grid cells look square
        ax.set_aspect('equal')

        if not DEBUG:
            # Save the figure to the specified picture path
            fig.savefig(picture_path, bbox_inches='tight', pad_inches=0.1, dpi=300, transparent=True)
        else:
            plt.show()

        # Close the figure to free memory
        plt.close()

    except ValueError as ve:
        # Handle value error (e.g., invalid input)
        print(f"Invalid Input: Error: {ve}")

    except Exception as e:
        # Catch any other exceptions and show a generic error message
        print(f"An error occurred while generating the 3D cube: {e}")


# Example use-case
if __name__ == "__main__":
    all_group_vars = {
        'wp_raw': {
            'x_wpraw': '10',   # Example values for testing
            'y_wpraw': '15',
            'z_wpraw': '20'
        },
        'clp': {
            'clp_offset': '4'
        }
    }
    generate_picture_clp(all_group_vars, 'output_image.png', DEBUG=True)