import streamlit as st
import tempfile
from OCC.Extend.DataExchange import read_step_file
from OCC.Core.BRep import BRep_Tool
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.TopAbs import TopAbs_FACE
from OCC.Core.TopoDS import topods_Face
from OCC.Core.BRepTools import breptools_UVBounds
from OCC.Core.gp import gp_Pnt

def extract_faces_points(shape):
    BRepMesh_IncrementalMesh(shape, 0.1)
    exp = TopExp_Explorer(shape, TopAbs_FACE)
    points = []

    while exp.More():
        face = topods_Face(exp.Current())
        u_min, u_max, v_min, v_max = breptools_UVBounds(face)
        for u in [u_min, (u_min + u_max)/2, u_max]:
            for v in [v_min, (v_min + v_max)/2, v_max]:
                pnt = gp_Pnt()
                BRep_Tool().Surface(face).D0(u, v, pnt)
                points.append((pnt.X(), pnt.Y(), pnt.Z()))
        exp.Next()

    return points

def compare_step_files(file1, file2):
    shape1 = read_step_file(file1)
    shape2 = read_step_file(file2)

    points1 = extract_faces_points(shape1)
    points2 = extract_faces_points(shape2)

    diffs = []
    for p1, p2 in zip(points1, points2):
        diff = ((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2 + (p1[2]-p2[2])**2) ** 0.5
        diffs.append(diff)

    max_dev = max(diffs)
    avg_dev = sum(diffs) / len(diffs)

    return max_dev, avg_dev, diffs

st.title("🔍 STEP File Comparison Tool")
st.markdown("Upload two STEP files with the same origin and compare their 3D geometry.")

uploaded_file1 = st.file_uploader("📤 Upload STEP File 1", type=["step", "stp"])
uploaded_file2 = st.file_uploader("📤 Upload STEP File 2", type=["step", "stp"])

if uploaded_file1 and uploaded_file2:
    with tempfile.NamedTemporaryFile(delete=False) as tmp1:
        tmp1.write(uploaded_file1.read())
        path1 = tmp1.name

    with tempfile.NamedTemporaryFile(delete=False) as tmp2:
        tmp2.write(uploaded_file2.read())
        path2 = tmp2.name

    st.info("🔧 Processing and comparing files...")
    max_dev, avg_dev, all_diffs = compare_step_files(path1, path2)

    st.success(f"📏 Max Deviation: {max_dev:.3f} mm")
    st.info(f"📉 Average Deviation: {avg_dev:.3f} mm")

    st.line_chart(all_diffs[:100])