import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

def visualize_results(simulator, length, width, thickness, stress, deformation):
    """
    Generates 3D and 2D visualizations for simulation results.
    Parameters:
        simulator (str): Name of the simulator (e.g., 'Python-Based Solver').
        length (float): Length of the object.
        width (float): Width of the object.
        thickness (float): Thickness of the object.
        stress (float): Stress value.
        deformation (float): Deformation value.
    Returns:
        str: Path to the 3D and 2D visualization images.
    """
    # Generate 2D bar chart
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.bar(["Stress", "Deformation"], [stress, deformation], color=["red", "blue"])
    ax.set_title(f"Simulation Results ({simulator})")
    ax.set_ylabel("Magnitude")
    ax.grid(True, linestyle="--", alpha=0.6)
    
    # Save 2D chart
    output_2d_path = "results_2d.png"
    plt.savefig(output_2d_path)
    plt.close(fig)
    
    # Generate 3D bar chart
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')
    
    x = np.array([1, 2])
    y = np.array([0, 0])
    z = np.array([0, 0])
    dx = np.array([1, 1])
    dy = np.array([1, 1])
    dz = np.array([stress, deformation])
    
    ax.bar3d(x, y, z, dx, dy, dz, color=["red", "blue"], alpha=0.6)
    ax.set_xlabel('Parameters')
    ax.set_ylabel('Category')
    ax.set_zlabel('Magnitude')
    ax.set_xticks([1, 2])
    ax.set_xticklabels(['Stress', 'Deformation'])
    ax.set_title(f"3D Simulation Results ({simulator})")

    # Save 3D chart
    output_3d_path = "results_3d.png"
    plt.savefig(output_3d_path)
    plt.close(fig)

    return output_2d_path, output_3d_path

def visualize_end_product(simulation_type, length, width, thickness, deformation):
    """
    Visualize the end product based on the parameters with 3D visualization.
    Parameters:
        simulation_type (str): 'plate' or 'beam'.
        length (float): Length of the product (mm).
        width (float): Width of the product (mm).
        thickness (float): Thickness of the product (mm).
        deformation (float): Deformation (mm).
    Returns:
        str: Path to the saved image of the 3D visualization.
    """
    fig = plt.figure(figsize=(8, 6))
    ax = fig.add_subplot(111, projection='3d')
    
    # Create a grid of coordinates for the surface
    x = np.linspace(0, length, 10)
    y = np.linspace(0, width, 10)
    x, y = np.meshgrid(x, y)
    z = np.zeros_like(x)  # Base plane
    
    # Add deformation to the surface
    z_deformed = z + np.random.normal(0, deformation / 10, size=z.shape)  # Adding deformation visually
    
    if simulation_type == "plate":
        ax.plot_surface(x, y, z_deformed, color='lightblue', alpha=0.8)
        ax.text(length / 2, width / 2, np.max(z_deformed) + 2, f"Deformation: {deformation:.2f} mm", color="red", fontsize=10, ha="center")
    elif simulation_type == "beam":
        ax.plot_surface(x, y, z_deformed, color='lightgreen', alpha=0.8)
        ax.text(length / 2, thickness / 2, np.max(z_deformed) + 2, f"Deflection: {deformation:.2f} mm", color="red", fontsize=10, ha="center")
    else:
        raise ValueError("Invalid simulation type.")
    
    # Add dimensions
    ax.set_xlabel("Length (mm)")
    ax.set_ylabel("Width (mm)")
    ax.set_zlabel("Deformation (mm)")
    ax.set_title(f"3D Visualization of {simulation_type.capitalize()} End Product")
    
    # Save the 3D visualization
    output_3d_path = f"{simulation_type}_end_product_3d.png"
    plt.savefig(output_3d_path)
    plt.close(fig)

    return output_3d_path

# Example usage:
simulate_results_2d_path, simulate_results_3d_path = visualize_results('Python-Based Solver', 200, 100, 20, 350, 0.02)
end_product_3d_path = visualize_end_product('plate', 200, 100, 20, 0.02)

print(f"2D Simulation Visualization saved at: {simulate_results_2d_path}")
print(f"3D Simulation Visualization saved at: {simulate_results_3d_path}")
print(f"3D End Product Visualization saved at: {end_product_3d_path}")