import FreeCAD, Fem, FemGui
from FreeCAD import Gui
import os
import sys
sys.path.append(r"C:\Program Files\FreeCAD 1.0\bin")  # Path to FreeCAD's Python library
sys.path.append(r"C:\Program Files\FreeCAD 1.0\lib")  # Additional FreeCAD modules

try:
    import FreeCAD, Fem, FemGui
    print("FreeCAD modules imported successfully.")
except ImportError as e:
    print(f"Error importing FreeCAD modules: {e}")


def run_fea(step_file, output_dir):
    """
    Automates FEA analysis in FreeCAD using the FEM Workbench.

    Args:
        step_file (str): Path to the STEP file.
        output_dir (str): Directory to save output files (images and analysis data).
    """
    # Create a new FreeCAD document
    doc = FreeCAD.newDocument("FEM_Analysis")
    
    # Import STEP file
    obj = doc.addObject("Part::Feature", "ImportedPart")
    obj.Shape = FreeCAD.importShape(step_file)
    doc.recompute()

    # Assign material properties
    mat = doc.addObject("App::MaterialObjectPython", "Material")
    mat.Material = {
        "Name": "Steel",
        "YoungsModulus": 210e9,  # Pa (N/m²)
        "PoissonRatio": 0.3,
    }
    doc.addObject("Fem::ConstraintMaterial", "MaterialConstraint").References = [(obj, "Face1")]  # Adjust face as needed
    doc.recompute()

    # Create and configure the mesh
    mesh = doc.addObject("Fem::FemMeshShapeNetgenObject", "FEMMesh")
    mesh.Shape = obj.Shape
    mesh.MaxSize = 5  # Set mesh size (adjust as required)
    doc.recompute()

    # Apply boundary conditions
    fixed_constraint = doc.addObject("Fem::ConstraintFixed", "FixedConstraint")
    fixed_constraint.References = [(obj, "Face1")]  # Adjust face as needed
    doc.recompute()

    # Apply force
    force_constraint = doc.addObject("Fem::ConstraintForce", "ForceConstraint")
    force_constraint.References = [(obj, "Face2")]  # Adjust face as needed
    force_constraint.Force = 1000  # N
    doc.recompute()

    # Add solver (CalculiX)
    solver = doc.addObject("Fem::SolverCalculixCcxTools", "Solver")
    solver.InputFile = os.path.join(output_dir, "input.inp")
    solver.WorkingDir = output_dir
    doc.recompute()

    # Write input file and run the solver
    solver.write_input_file()
    solver.run()

    # Load and process results
    results = solver.load_results()
    results_obj = doc.addObject("Fem::FemResultObject", "Results")
    results_obj.Mesh = results.Mesh
    results_obj.NodeCount = len(results.Mesh.Nodes)
    doc.recompute()

    # Save 2D and 3D views as images
    two_d_image_path = os.path.join(output_dir, "2D_View.png")
    three_d_image_path = os.path.join(output_dir, "3D_View.png")

    Gui.ActiveDocument.ActiveView.saveImage(two_d_image_path, 1024, 768, "White")
    Gui.ActiveDocument.ActiveView.viewIsometric()
    Gui.ActiveDocument.ActiveView.saveImage(three_d_image_path, 1024, 768, "White")

    # Save the FreeCAD project file
    doc.saveAs(os.path.join(output_dir, "FEM_Analysis.FCStd"))

    print(f"Analysis completed successfully.")
    print(f"2D View saved to: {two_d_image_path}")
    print(f"3D View saved to: {three_d_image_path}")

# Example Usage
if __name__ == "__main__":
    step_file = "path_to_your_step_file.stp"  # Replace with your STEP file path
    output_dir = "path_to_output_directory"  # Replace with your desired output directory

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    run_fea(step_file, output_dir)

# Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# ANSYS APDL Script for 2D and 3D Views")

    with gr.Row():
        step_file = gr.File(label="Upload STEP File")
        step_output = gr.Text(label="Upload Status")
        step_file.upload(upload_step_file, inputs=step_file, outputs=step_output)

    with gr.Row():
        apdl_view = gr.Textbox(apdl_script, lines=20, label="APDL Script")
        modify_button = gr.Button("Save Changes")
        modify_button.click(modify_apdl_script, inputs=apdl_view, outputs=step_output)

    with gr.Row():
        execute_button = gr.Button("Run Script")
        execution_output = gr.Text(label="Execution Status")
        log_output = gr.Textbox(label="Output Log", lines=10, interactive=False)
        execute_button.click(execute_script, outputs=[execution_output])
        execute_button.click(read_output_log, outputs=log_output)

    with gr.Row():
        gr.Markdown("## Generated Images")
        view_2d = gr.Image(label="2D View", interactive=False)
        view_3d = gr.Image(label="3D View", interactive=False)
        refresh_button = gr.Button("Refresh Images")
        refresh_button.click(get_generated_images, outputs=[view_2d, view_3d])

app.launch()