# Import additional libraries for 3D visualization
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

def visualize_3d_results(length, width, stress, deformation):
    """
    Generate a 3D visualization of the simulation results.
    """
    x = np.linspace(0, length, 100)
    y = np.linspace(0, width, 100)
    X, Y = np.meshgrid(x, y)
    Z = np.sin((X / length) * np.pi) * np.cos((Y / width) * np.pi) * deformation

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(X, Y, Z, cmap='viridis')
    ax.set_title("3D Simulation Visualization")
    ax.set_xlabel("Length (mm)")
    ax.set_ylabel("Width (mm)")
    ax.set_zlabel("Deformation (mm)")
    output_path = "3d_visualization.png"
    plt.savefig(output_path)  # Save the visualization to a file
    plt.close()
    return output_path

def simulation_workflow(tool_type, use_case, include_hole, include_force, include_load, thickness, length, width, hole_diameter, force, load, elastic_modulus):
    """
    Main simulation workflow.
    """
    # Handle optional parameters
    force = force if include_force else 0
    load = load if include_load else 0

    # Generate APDL script dynamically
    if use_case == "plate":
        hole_diameter = hole_diameter if include_hole else 0
        apdl_path = generate_plate_apdl(thickness, length, width, hole_diameter, force)
    elif use_case == "beam":
        apdl_path = generate_beam_apdl(length, width, thickness, load)
    else:
        return "Invalid use case selected.", None, None, None, None

    # Read the generated APDL program
    with open(apdl_path, "r") as file:
        apdl_program = file.read()

    # Run simulation using Python-based solver
    stress, deformation = run_python_simulation(apdl_path, use_case, thickness, length, width, force, load, elastic_modulus)

    # Explanation for negligible deformation
    if deformation < 1e-6:
        deformation_note = "\nNote: Deformation is negligible due to the high rigidity of the material."
    else:
        deformation_note = ""

    # Generate 2D simulation visualization
    graph_path, _ = visualize_results("Python-Based Solver", length, width, thickness, stress, deformation)
    
    # Generate end-product visualization
    product_path = visualize_end_product(use_case, length, width, thickness, deformation)

    # Generate 3D visualization
    visualization_3d_path = visualize_3d_results(length, width, stress, deformation)

    return (
        f"{tool_type} Simulation\nStress: {stress:.2f} MPa\nDeformation: {deformation:.6f} mm{deformation_note}",
        graph_path,
        product_path,
        apdl_program,
        visualization_3d_path
    )

# Update Gradio interface
with gr.Blocks() as interface:
    gr.Markdown("# Punch and Die Simulation Tool")
    
    with gr.Row():
        tool_type = gr.Radio(["Punch", "Die"], label="Select Tool Type")
        use_case = gr.Radio(["plate", "beam"], label="Select Use Case")
    
    with gr.Row():
        include_hole = gr.Checkbox(label="Include Hole for Plate Simulation")
        include_force = gr.Checkbox(label="Include Force")
        include_load = gr.Checkbox(label="Include Load")
    
    with gr.Row():
        thickness = gr.Slider(10, 50, step=1, label="Thickness (mm)")
        length = gr.Slider(100, 500, step=10, label="Length (mm)")
        width = gr.Slider(50, 200, step=10, label="Width (mm)")
    
    with gr.Row():
        hole_diameter = gr.Slider(5, 25, step=1, label="Hole Diameter (mm)")
        force = gr.Slider(1000, 10000, step=500, label="Force (N)")
        load = gr.Slider(1000, 20000, step=1000, label="Load (N)")
        elastic_modulus = gr.Slider(5e10, 3e11, step=1e10, label="Elastic Modulus (Pa)")
    
    submit_button = gr.Button("Run Simulation")
    
    with gr.Row():
        result_text = gr.Textbox(label="Simulation Results")
        result_graph = gr.Image(label="2D Simulation Visualization")
        result_product = gr.Image(label="End Product Visualization")
        result_3d = gr.Image(label="3D Simulation Visualization")
        apdl_output = gr.Code(language="python", label="Generated APDL Program")
    
    # Link the button to the simulation function
    submit_button.click(
        simulation_workflow,
        inputs=[
            tool_type, use_case, include_hole, include_force, include_load,
            thickness, length, width, hole_diameter, force, load, elastic_modulus
        ],
        outputs=[result_text, result_graph, result_product, apdl_output, result_3d]
    )

# Launch interface
interface.launch()