import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
from ansys.mapdl.core import launch_mapdl

# Function to generate APDL script
def generate_apdl_script(youngs_modulus, poissons_ratio, length, width, thickness, element_size, pressure):
    script = f"""
    /PREP7
    MP,EX,1,{youngs_modulus}
    MP,PRXY,1,{poissons_ratio}
    BLOCK,0,{length},0,{width},0,{thickness}
    ESIZE,{element_size}
    ET,1,SOLID185
    VMESH,ALL
    NSEL,S,LOC,X,0
    D,ALL,ALL
    ALLSEL,ALL
    NSEL,S,LOC,Z,{thickness}
    SF,ALL,PRES,{pressure}
    ALLSEL,ALL
    /SOLU
    ANTYPE,STATIC
    SOLVE
    FINISH
    /POST1
    SET,1,1
    PRNSOL,U,Z
    PRNSOL,S,EQV
    FINISH
    """
    return script

# Function to run simulation and return results
def run_simulation(apdl_script):
    mapdl = launch_mapdl()
    mapdl.input_strings(apdl_script)
    mapdl.post1()
    mapdl.set(1, 1)
    nodal_displacements = mapdl.post_processing.nodal_displacement('UZ')
    nodal_stress = mapdl.post_processing.nodal_stress(0)
    mapdl.exit()
    return nodal_displacements, nodal_stress

# Function to visualize results
def visualize_results(displacements, stress):
    fig, ax = plt.subplots(2, 1, figsize=(10, 8))
    
    ax[0].plot(displacements, label='Nodal Displacement (UZ)')
    ax[0].set_xlabel('Node Number')
    ax[0].set_ylabel('Displacement (mm)')
    ax[0].set_title('Nodal Displacement in Z Direction')
    ax[0].legend()
    ax[0].grid(True)
    
    ax[1].plot(stress, label='Nodal Stress (EQV)')
    ax[1].set_xlabel('Node Number')
    ax[1].set_ylabel('Stress (MPa)')
    ax[1].set_title('Equivalent Nodal Stress')
    ax[1].legend()
    ax[1].grid(True)
    
    plt.tight_layout()
    plt.show()

# Function to evaluate results against design criteria
def evaluate_results(displacements, stress, displacement_threshold=0.5, stress_threshold=250):
    max_displacement = np.max(np.abs(displacements))
    max_stress = np.max(np.abs(stress))
    displacement_ok = max_displacement <= displacement_threshold
    stress_ok = max_stress <= stress_threshold
    return displacement_ok and stress_ok, max_displacement, max_stress

# Gradio interface functions
def project_requirements(youngs_modulus, poissons_ratio, length, width, thickness, element_size, pressure):
    params = {
        'youngs_modulus': youngs_modulus,
        'poissons_ratio': poissons_ratio,
        'length': length,
        'width': width,
        'thickness': thickness,
        'element_size': element_size,
        'pressure': pressure
    }
    return params

def apdl_program(params):
    apdl_script = generate_apdl_script(**params)
    return apdl_script

def simulation_and_visualization(apdl_script):
    displacements, stress = run_simulation(apdl_script)
    visualize_results(displacements, stress)
    design_acceptable, max_disp, max_strs = evaluate_results(displacements, stress)
    result_message = (
        f"Max Displacement: {max_disp:.2f} mm\n"
        f"Max Stress: {max_strs:.2f} MPa\n"
        f"Design {'meets' if design_acceptable else 'does not meet'} all criteria."
    )
    return result_message

# Gradio Blocks for multi-tab interface
with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.TabItem("Project Requirements"):
            youngs_modulus = gr.Number(label="Young's Modulus (MPa)", value=210000)
            poissons_ratio = gr.Number(label="Poisson's Ratio", value=0.3)
            length = gr.Number(label="Length (mm)", value=100)
            width = gr.Number(label="Width (mm)", value=50)
            thickness = gr.Number(label="Thickness (mm)", value=10)
            element_size = gr.Number(label="Element Size (mm)", value=5)
            pressure = gr.Number(label="Applied Pressure (MPa)", value=10)
            params_button = gr.Button("Submit")
            params_output = gr.JSON()
            params_button.click(
                project_requirements,
                inputs=[youngs_modulus, poissons_ratio, length, width, thickness, element_size, pressure],
                outputs=params_output
            )
        with gr.TabItem("APDL Program Generation"):
            apdl_output = gr.Textbox(label="Generated APDL Script", lines=20)
            params_output.change(
                apdl_program,
                inputs=params_output,
                outputs=apdl_output
            )
        with gr.TabItem("Simulation and Visualization"):
            result_output = gr.Textbox(label="Simulation Results", lines=10)
            simulate_button = gr.Button("Run Simulation")
            simulate_button.click(
                simulation_and_visualization,
                inputs=apdl_output,
                outputs=result_output
            )

demo.launch()