import gradio as gr
import matplotlib.pyplot as plt
import plotly.graph_objects as go
import tempfile
import numpy as np
import os


# Function to simulate G-code generation
def generate_gcode(length, diameter, step_file):
    try:
        # Mock G-code generation as a string
        gcode = f"(Generated G-code for Toolpath)\n"
        gcode += f"G21 ; Set units to mm\n"
        gcode += f"G00 X0 Y0 Z0 ; Move to start\n"
        gcode += f"G01 X{length} Y0 Z0 F1000 ; Move along X-axis\n"
        gcode += f"G01 X{length} Y{diameter} Z0 ; Move to end point\n"
        gcode += f"G00 X0 Y0 Z0 ; Return to origin\n"
        gcode += f"M30 ; End of program"

        # Save the G-code to a file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".nc")
        with open(temp_file.name, "w") as f:
            f.write(gcode)
        
        print("G-code file generated successfully")
        return temp_file.name  # Return file path
    
    except Exception as e:
        print(f"Error in G-code Generation: {str(e)}")
        return None


# Function to generate 2D Visualization with Graphs
def generate_2d_plot(length, diameter):
    try:
        # Create 2D Rectangle and Overlay a Graph
        fig, ax = plt.subplots()
        ax.plot([0, length, length, 0, 0], [0, 0, diameter, diameter, 0], color='blue', linewidth=2, label="Toolpath")
        
        # Example graph (sinusoidal curve)
        x = np.linspace(0, length, 100)
        y = diameter / 2 * np.sin(2 * np.pi * x / length)
        ax.plot(x, y, color='red', linestyle='--', label="Graph Overlay")
        
        ax.set_title("2D Visualization with Graphs")
        ax.set_xlabel("Length (mm)")
        ax.set_ylabel("Diameter (mm)")
        ax.legend()

        # Save the figure to a temporary file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        plt.savefig(temp_file.name, format='png')
        plt.close(fig)
        return temp_file.name  # Return the file path
    
    except Exception as e:
        print(f"Error in 2D Plot: {str(e)}")
        return None


# Function to generate 3D Visualization
def generate_3d_plot(length, diameter):
    try:
        t = np.linspace(0, 2 * np.pi, 50)
        z = np.linspace(0, length, 20)
        T, Z = np.meshgrid(t, z)
        X = diameter * np.cos(T)
        Y = diameter * np.sin(T)

        fig = go.Figure(data=[go.Surface(x=X, y=Y, z=Z, colorscale='Viridis')])
        fig.update_layout(
            title="3D Visualization of Toolpath",
            scene=dict(
                xaxis_title="X (mm)",
                yaxis_title="Y (mm)",
                zaxis_title="Z (mm)"
            )
        )
        return fig
    
    except Exception as e:
        print(f"Error in 3D Plot: {str(e)}")
        return None


# Main function to handle all actions
def app_interface(length, diameter, step_file):
    # Process STEP file (Placeholder logic)
    step_file_path = None
    if step_file is not None:
        step_file_path = step_file.name
        print(f"Uploaded STEP file: {step_file_path}")

    # Generate G-code
    gcode_file = generate_gcode(length, diameter, step_file_path)
    
    # Generate Visualizations
    file_path_2d = generate_2d_plot(length, diameter)
    fig_3d = generate_3d_plot(length, diameter)
    
    return file_path_2d, fig_3d, gcode_file


# Define Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## CNC Toolpath Generator with 2D/3D Visualization and G-code Generation")

    with gr.Row():
        length_input = gr.Number(label="Length (mm)", value=100)
        diameter_input = gr.Number(label="Diameter (mm)", value=50)
        step_file_input = gr.File(label="Upload STEP File")

    submit_button = gr.Button("Submit")

    with gr.Row():
        output_2d = gr.Image(label="2D Visualization with Graphs")
        output_3d = gr.Plot(label="3D Visualization")

    gcode_output = gr.File(label="G-code File")

    submit_button.click(
        app_interface,
        inputs=[length_input, diameter_input, step_file_input],
        outputs=[output_2d, output_3d, gcode_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()