import gradio as gr
import os
import cv2
import numpy as np
from OCC.Extend.DataExchange import read_step_file
from OCC.Core.BRepBnd import Bnd_Box
from OCC.Core.gp import gp_Pnt
from OCC.Core.BRep import BRep_Builder
from OCC.Core.TopAbs import TopAbs_SOLID
from OCC.Core.TopoDS import topods_Solid
from OCC.Core.STEPControl import STEPControl_Reader

# Function to process image and extract dimensions (example: simple edge detection)
def process_image_for_dimensions(image_file_path):
    """
    Process CAD image to extract dimensions. For simplicity, we assume
    a 2D image where we can detect the outline or significant objects.
    """
    # Load image
    image = cv2.imread(image_file_path, cv2.IMREAD_GRAYSCALE)
    _, thresholded = cv2.threshold(image, 128, 255, cv2.THRESH_BINARY)
    
    # Find contours of the image
    contours, _ = cv2.findContours(thresholded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        raise ValueError("No contours found in the image.")
    
    # Get the bounding box of the largest contour (assuming it's the object of interest)
    contour = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(contour)

    # Example: assume extracted dimensions from bounding box
    width = w  # Width of the bounding box
    height = h  # Height of the bounding box
    depth = 50  # Default depth, you can modify this based on input or assumptions
    
    # Return the extracted dimensions
    return (width, height, depth)

# Function to process STEP file and extract dimensions
def process_step_file(step_file_path):
    """
    Process a STEP file to extract the 3D bounding box dimensions of the geometry.
    """
    # Read the STEP file
    reader = STEPControl_Reader()
    status = reader.ReadFile(step_file_path)
    
    if status != 1:
        raise ValueError(f"Failed to load STEP file: {step_file_path}")
    
    # Get the shape from the STEP file
    shape = reader.Shape()
    
    # Calculate bounding box
    bbox = Bnd_Box()
    builder = BRep_Builder()
    builder.Add(bbox, shape)
    
    # Get the bounds of the box
    xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()
    
    # Return dimensions (width, height, depth)
    width = xmax - xmin
    height = ymax - ymin
    depth = zmax - zmin
    
    return (width, height, depth)

# Function to generate APDL script from extracted properties
def generate_apdl_script(cad_file, element_type, material_ex, material_nu):
    """
    Generate an APDL script based on CAD file properties and user input.
    """
    file_extension = os.path.splitext(cad_file.name)[1].lower()
    
    # Extract dimensions based on the file type
    if file_extension in ['.png', '.jpg', '.jpeg', '.bmp']:  # For image files
        cad_dimensions = process_image_for_dimensions(cad_file.name)
    elif file_extension in ['.step', '.stp']:  # For STEP files
        cad_dimensions = process_step_file(cad_file.name)
    else:
        raise ValueError(f"Unsupported CAD file format: {file_extension}")
    
    # File name for the generated APDL script
    file_name = "generated_apdl_script.inp"
    
    # Generate the APDL script
    with open(file_name, 'w') as apdl_file:
        apdl_file.write("! APDL Script Generated from CAD File\n")
        apdl_file.write("/PREP7\n")
        apdl_file.write(f"ET,1,{element_type}\n")  # Define element type
        apdl_file.write(f"BLOCK,0,{cad_dimensions[0]},0,{cad_dimensions[1]},0,{cad_dimensions[2]}\n")  # Create a block
        apdl_file.write(f"MP,EX,1,{material_ex}\n")  # Define material property (Young's modulus)
        apdl_file.write(f"MP,NUXY,1,{material_nu}\n")  # Define material property (Poisson's ratio)
        apdl_file.write("MAT,1\n")
        apdl_file.write("ESEL,S,ALL\n")  # Select all elements
        apdl_file.write("ESIZE,10\n")  # Element size
        apdl_file.write("MSHAPE,1,3D\n")  # Define mesh shape
        apdl_file.write("MESH,ALL\n")  # Generate the mesh
        apdl_file.write("/SOLU\n")  # Switch to solution processor
        apdl_file.write("ANTYPE,STATIC\n")  # Define analysis type
        apdl_file.write("D,1,UX,0\n")  # Apply boundary condition (UX=0)
        apdl_file.write("D,1,UY,0\n")  # Apply boundary condition (UY=0)
        apdl_file.write("D,1,UZ,0\n")  # Apply boundary condition (UZ=0)
        apdl_file.write("FINISH\n")  # Finish pre-processing
        apdl_file.write("/POST1\n")  # Post-processing
        apdl_file.write("SET,LAST\n")  # Use the last result set
        apdl_file.write("PLDISP,ALL\n")  # Plot displacement
        apdl_file.write("FINISH\n")  # Finish post-processing
    
    return file_name

# Gradio Interface for CAD File Input (Image or STEP)
def generate_script_ui_from_cad(cad_file, element_type, material_ex, material_nu):
    """
    Interface function to generate an APDL script from CAD file (image or STEP).
    """
    try:
        file_path = generate_apdl_script(cad_file, element_type, material_ex, material_nu)
        return file_path
    except ValueError as e:
        return str(e)

# Gradio Input and Output for CAD File Upload
iface = gr.Interface(
    fn=generate_script_ui_from_cad,
    inputs=[
        gr.File(label="Upload CAD File (Image or STEP)"),
        gr.Textbox(label="Element Type (e.g., SOLID186)", value="SOLID186"),
        gr.Number(label="Material EX (Young's Modulus)", value=200e9),
        gr.Number(label="Material NU (Poisson's Ratio)", value=0.3),
    ],
    outputs=gr.File(label="Generated APDL Script"),
    title="APDL Script Generator from CAD File",
    description="Upload a CAD file (2D image or STEP format) and provide material properties to generate an APDL script."
)

# Launch the Interface
iface.launch()