import matplotlib.pyplot as plt
import numpy as np


def generate_picture_grp(all_group_vars, picture_path, DEBUG=False):
    try:
        # Extract variables for the gripper
        vars = all_group_vars['grp']
        h_offset = float(vars['h_offset'])
        h_step = float(vars['h_step'])
        w_offset = float(vars['w_offset'])
        w_step = float(vars['w_step'])
        steps_num = int(vars['steps_num'])

        # ERROR HANDLING
        if h_step <= 0 or w_step <= 0 or steps_num <= 0:
            raise ValueError("h_step, w_step, and steps_num must be positive values.")

        # Initialize figure
        fig, ax = plt.subplots(figsize=(6, 6))  # Square figure
        ax.set_aspect('equal')

        # Draw the gripper finger
        def draw_gripper(x_sign=1):
            max_width = x_sign * (w_offset + w_step * (steps_num - 1))
            current_height = h_offset
            current_width = x_sign * w_offset

            # Draw the first block
            rect = plt.Rectangle((x_sign, 0), max_width - x_sign, h_offset, facecolor='black', linewidth=0)
            ax.add_patch(rect)

            # Draw the steps
            for _ in range(steps_num):
                rect = plt.Rectangle((current_width, current_height), max_width - current_width, h_step, facecolor='black', linewidth=0)
                ax.add_patch(rect)
                current_height += h_step
                current_width += x_sign * w_step

        # Draw original gripper finger and x flipped
        draw_gripper()
        draw_gripper(x_sign=-1)

        # Plot the origin point at (0, 0)
        ax.plot(0, 0, marker='o', markersize=10, color='red')  # Red point at origin

        # Flip the y-axis (invert it)
        ax.invert_yaxis()

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

        # Hide major ticks and labels (only gridlines should be visible)
        ax.set_xticks([])
        ax.set_yticks([])

        # Add a white border around the plot to ensure the outer gridlines are visible
        fig.patch.set_facecolor('white')
        plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)  # Adjust padding for the white border

        if not DEBUG:
            # Save the figure to the specified picture path
            plt.gca().set_aspect('equal', adjustable='box')
            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 gripper finger: {e}")


if __name__ == "__main__":
    vars = {
        'h_offset': 10,
        'h_step': 5,
        'w_offset': 20,
        'w_step': 10,
        'steps_num': 5
    }
    picture_path = "cfg_picture/TEST_gripper_finger.jpg"
    generate_picture_grp(vars, picture_path, DEBUG=True)