#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
STEP -> OCC triangulate -> PyTorch3D orthographic 3 views -> SVG (silhouette + visible/hidden feature edges)
Then combine 3 SVGs into a single SVG and a PNG.

Views:
  - top  : from +Z towards origin (view direction -Z)
  - left : from -X towards origin (view direction +X)
  - front: from +Y towards origin (view direction -Y)

Outputs (default in --out_dir):
  - top.svg, left.svg, front.svg
  - combined_views.svg
  - combined_views.png
"""

import os, sys, math, argparse, tempfile
import numpy as np
import torch
from collections import defaultdict

# ---------- PyTorch3D ----------
from pytorch3d.structures import Meshes
from pytorch3d.renderer import (
    RasterizationSettings, MeshRasterizer,
    MeshRenderer, BlendParams, SoftSilhouetteShader,
    look_at_view_transform, OrthographicCameras
)
from pytorch3d.renderer.mesh.rasterizer import Fragments

# ---------- SVG / Image ----------
import svgwrite
from skimage import measure

# ---------- pythonocc (OCC) ----------
from OCC.Core.STEPControl import STEPControl_Reader
from OCC.Core.TopAbs import TopAbs_FACE
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.BRep import BRep_Tool
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.ShapeFix import ShapeFix_Shape

# ---------- Combine SVG ----------
import svgutils.transform as st
from svgutils.compose import Unit
import cairosvg


# ---------------- Utils ----------------

def autodevice():
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def read_step_solid(path: str):
    reader = STEPControl_Reader()
    stat = reader.ReadFile(path)
    if stat != 1:
        raise RuntimeError(f"STEP 读取失败: {path}")
    reader.TransferRoots()
    shape = reader.OneShape()
    fixer = ShapeFix_Shape(shape)
    fixer.Perform()
    return fixer.Shape()


def triangulate(shape, deflection=0.5, angle=0.5):
    """OCC 三角化：返回 (V(np.float32 N×3), F(np.int64 M×3))"""
    BRepMesh_IncrementalMesh(shape, deflection, False, angle, True)
    verts, faces, v_off = [], [], 0
    exp = TopExp_Explorer(shape, TopAbs_FACE)
    while exp.More():
        face = exp.Current()
        loc = TopLoc_Location()
        tri = BRep_Tool.Triangulation(face, loc)
        if tri is not None:
            nb_nodes = tri.NbNodes()
            has_nodes_arr = hasattr(tri, "Nodes")
            for i in range(1, nb_nodes + 1):
                p = tri.Nodes().Value(i) if has_nodes_arr else tri.Node(i)
                p = p.Transformed(loc.Transformation())
                verts.append([p.X(), p.Y(), p.Z()])
            nb_tris = tri.NbTriangles()
            has_tris_arr = hasattr(tri, "Triangles")
            for i in range(1, nb_tris + 1):
                t = tri.Triangles().Value(i) if has_tris_arr else tri.Triangle(i)
                a, b, c = t.Get()
                faces.append([v_off + a - 1, v_off + b - 1, v_off + c - 1])
            v_off += nb_nodes
        exp.Next()
    if not verts or not faces:
        raise RuntimeError("三角化为空，尝试减小 --defl")
    return np.asarray(verts, np.float32), np.asarray(faces, np.int64)


def normalize_mesh_np(verts: np.ndarray, unit_scale=1.0):
    c = verts.mean(axis=0, keepdims=True)
    v0 = verts - c
    s = np.max(np.abs(v0))
    s = max(s, 1e-12)
    return v0 / s * unit_scale


def build_mesh(V_np, F_np, device):
    V = torch.from_numpy(V_np).to(device)
    F = torch.from_numpy(F_np).to(device)
    return Meshes(verts=[V], faces=[F])


def compute_face_normals(verts: torch.Tensor, faces: torch.Tensor):
    v0 = verts[faces[:, 0]];
    v1 = verts[faces[:, 1]];
    v2 = verts[faces[:, 2]]
    n = torch.cross(v1 - v0, v2 - v0, dim=1)
    return torch.nn.functional.normalize(n, dim=1, eps=1e-12)


def extract_feature_edges(faces: torch.Tensor, face_normals: torch.Tensor, angle_deg: float):
    """二面角>=angle_deg 视为特征棱；边界边必取。"""
    thr = math.cos(math.radians(max(0.0, angle_deg)))
    e2f = defaultdict(list)
    F = faces.shape[0]
    for f in range(F):
        i, j, k = faces[f].tolist()
        for a, b in ((i, j), (j, k), (k, i)):
            e = (a, b) if a < b else (b, a)
            e2f[e].append(f)
    feat = []
    for e, fl in e2f.items():
        if len(fl) == 1:
            feat.append(e)
        elif len(fl) == 2:
            n0, n1 = face_normals[fl[0]], face_normals[fl[1]]
            cosv = torch.clamp((n0 * n1).sum(), -1.0, 1.0).item()
            if cosv <= thr:
                feat.append(e)
    return feat


def raster_fragments(mesh: Meshes, cameras, image_size=1600, faces_per_pixel=1):
    rs = RasterizationSettings(
        image_size=image_size,
        blur_radius=0.0,
        faces_per_pixel=faces_per_pixel,
        cull_backfaces=True
    )
    rast = MeshRasterizer(cameras=cameras, raster_settings=rs)
    with torch.no_grad():
        frags: Fragments = rast(mesh, cameras=cameras)
    return frags, frags.zbuf[0, ..., 0]


def render_silhouette(mesh: Meshes, cameras, image_size=1600):
    rs = RasterizationSettings(
        image_size=image_size,
        blur_radius=1e-6,
        faces_per_pixel=50,
        cull_backfaces=True
    )
    renderer = MeshRenderer(
        rasterizer=MeshRasterizer(cameras=cameras, raster_settings=rs),
        shader=SoftSilhouetteShader(blend_params=BlendParams(sigma=1e-4, gamma=1e-4))
    )
    with torch.no_grad():
        img = renderer(mesh, cameras=cameras)  # (1,H,W,4)
    return np.clip(img[0, ..., 3].detach().cpu().numpy(), 0.0, 1.0)


def project_points_to_screen(cameras, points_world: torch.Tensor, image_size: int):
    with torch.no_grad():
        scr = cameras.transform_points_screen(
            points_world[None, ...],
            image_size=((image_size, image_size),)
        )
    return scr[0]  # (N,3)


def edge_visibility_split(edges_idx, verts_world, cameras, zbuf, image_size, eps=1e-4):
    H = W = image_size
    vis, hid = [], []
    scr = project_points_to_screen(cameras, verts_world, image_size)
    zmin = zbuf.detach().cpu().numpy()
    for i, j in edges_idx:
        p0, p1 = scr[i], scr[j]
        m = 0.5 * (p0 + p1)
        x = int(np.clip(m[0].item(), 0, W - 1))
        y = int(np.clip(m[1].item(), 0, H - 1))
        z_proj = m[2].item()
        z_ref = zmin[y, x]
        seg = (p0[0].item(), p0[1].item(), p1[0].item(), p1[1].item())
        if np.isfinite(z_ref) and (z_proj <= z_ref + eps):
            vis.append(seg)
        else:
            hid.append(seg)
    return vis, hid


def trace_silhouettes(alpha: np.ndarray, threshold=0.5, step=1):
    contours = measure.find_contours(alpha, level=threshold)
    polys = []
    for cnt in contours:
        if len(cnt) < 8:
            continue
        cnt = cnt[::max(1, step)]
        polys.append(np.stack([cnt[:, 1], cnt[:, 0]], axis=1))  # (x,y)
    return polys


def svg_from_view(out_svg, W, H, silhouettes, edges_vis, edges_hid, stroke=2.0, margin=10, fill="none"):
    dwg = svgwrite.Drawing(out_svg, size=(W + 2 * margin, H + 2 * margin))
    dwg.add(dwg.rect(insert=(0, 0), size=(W + 2 * margin, H + 2 * margin), fill="none"))

    # silhouette fill (optional)
    for poly in silhouettes:
        if len(poly) < 3:
            continue
        path = [f"M {poly[0, 0] + margin:.2f} {poly[0, 1] + margin:.2f}"]
        for i in range(1, len(poly)):
            path.append(f"L {poly[i, 0] + margin:.2f} {poly[i, 1] + margin:.2f}")
        path.append("Z")
        dwg.add(dwg.path(" ".join(path), fill=fill, stroke="none"))

    # visible edges
    for (x0, y0, x1, y1) in edges_vis:
        dwg.add(dwg.line(
            (x0 + margin, y0 + margin),
            (x1 + margin, y1 + margin),
            stroke="black",
            stroke_width=stroke
        ))

    # hidden edges
    for (x0, y0, x1, y1) in edges_hid:
        dwg.add(dwg.line(
            (x0 + margin, y0 + margin),
            (x1 + margin, y1 + margin),
            stroke="black",
            stroke_width=max(1.0, 0.75 * stroke),
            stroke_dasharray=[6, 6],
            opacity=0.85
        ))

    dwg.save()


# ---------------- Cameras ----------------

def make_camera(view: str, device, dist=2.7):
    """
    look_at_view_transform spherical params for engineering views.
    """
    if view == "top":
        # from +Z to origin (view -Z)
        R, T = look_at_view_transform(dist=dist, elev=90.0, azim=180.0, device=device)
    elif view == "left":
        # from -X to origin (view +X)
        R, T = look_at_view_transform(dist=dist, elev=0.0, azim=270.0, device=device)
    elif view == "front":
        # from +Y to origin (view -Y)
        R, T = look_at_view_transform(dist=dist, elev=0.0, azim=180.0, device=device)
    else:
        raise ValueError("view must be one of ['top','left','front']")
    return OrthographicCameras(R=R, T=T, device=device)


# ---------------- Combine 3 SVGs ----------------

def combine_svgs(view1_path, view2_path, view3_path, output_svg, output_image):
    """
    横向拼接 3 张 SVG，并额外导出 PNG。
    """
    out_dir = os.path.dirname(output_svg)
    if out_dir:
        os.makedirs(out_dir, exist_ok=True)

    view1 = st.fromfile(view1_path)
    view2 = st.fromfile(view2_path)
    view3 = st.fromfile(view3_path)

    r1 = view1.getroot()
    r2 = view2.getroot()
    r3 = view3.getroot()

    w1, h1 = [float(x.replace("px", "")) for x in view1.get_size()]
    w2, h2 = [float(x.replace("px", "")) for x in view2.get_size()]
    w3, h3 = [float(x.replace("px", "")) for x in view3.get_size()]

    max_w = max(w1, w2, w3)
    max_h = max(h1, h2, h3)

    combined_width = max_w * 3 + 40
    combined_height = max_h + 20

    combined = st.SVGFigure(Unit(combined_width), Unit(combined_height))

    x_offset = 10
    y_offset = 10

    r1.moveto(x_offset, y_offset)
    x_offset += max_w + 20
    r2.moveto(x_offset, y_offset)
    x_offset += max_w + 20
    r3.moveto(x_offset, y_offset)

    combined.append([r1, r2, r3])
    combined.save(output_svg)

    # SVG -> PNG
    cairosvg.svg2png(url=output_svg, write_to=output_image)


# ---------------- Main ----------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--step_path", type=str, required=False, help="输入 STEP 文件", default="output/sample_clean.step")
    ap.add_argument("--out_dir", type=str, required=False, help="输出目录", default="output")

    ap.add_argument("--res", type=int, default=1400, help="渲染分辨率（正方形）")
    ap.add_argument("--defl", type=float, default=1.5, help="三角化线性偏差")
    ap.add_argument("--ang", type=float, default=65.0, help="二面角阈值(°)")
    ap.add_argument("--stroke", type=float, default=1.3, help="SVG 线宽(px)")
    ap.add_argument("--scale", type=float, default=1.0, help="归一化后模型最大边长")

    ap.add_argument("--no_combine", action="store_true", help="只导出三视图 SVG，不合成")
    ap.add_argument("--combined_svg", type=str, default=None, help="合成 SVG 输出路径（默认 out_dir/combined_views.svg）")
    ap.add_argument("--combined_png", type=str, default=None, help="合成 PNG 输出路径（默认 out_dir/combined_views.png）")

    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    device = autodevice()
    print(f"[Device] {device}")

    # STEP -> mesh
    shape = read_step_solid(args.step_path)
    V_np, F_np = triangulate(shape, deflection=args.defl)
    V_np = normalize_mesh_np(V_np, unit_scale=args.scale)

    mesh = build_mesh(V_np, F_np, device)
    verts = mesh.verts_packed()
    faces = mesh.faces_packed()
    fn = compute_face_normals(verts, faces)
    feat_edges = extract_feature_edges(faces, fn, angle_deg=args.ang)

    view_paths = {}
    for name in ["top", "left", "front"]:
        print(f"[View] {name}")
        cam = make_camera(name, device)
        alpha = render_silhouette(mesh, cam, image_size=args.res)
        _, zbuf = raster_fragments(mesh, cam, image_size=args.res, faces_per_pixel=1)

        e_vis, e_hid = edge_visibility_split(feat_edges, verts, cam, zbuf, args.res, eps=1e-4)
        polys = trace_silhouettes(alpha, threshold=0.5, step=1)

        out_svg = os.path.join(args.out_dir, f"{name}.svg")
        svg_from_view(out_svg, args.res, args.res, polys, e_vis, e_hid,
                      stroke=args.stroke, margin=10, fill="none")
        view_paths[name] = out_svg
        print(f"  -> {out_svg} (edges: vis={len(e_vis)}, hid={len(e_hid)}, polys={len(polys)})")

    if args.no_combine:
        print("[Done] 3 views exported (no combine).")
        return

    combined_svg = args.combined_svg or os.path.join(args.out_dir, "combined_views.svg")
    combined_png = args.combined_png or os.path.join(args.out_dir, "combined_views.png")

    # 注意：这里组合顺序按你第二段示例：front / top / left
    combine_svgs(
        view1_path=view_paths["front"],
        view2_path=view_paths["top"],
        view3_path=view_paths["left"],
        output_svg=combined_svg,
        output_image=combined_png
    )
    print(f"[Combined] SVG: {combined_svg}")
    print(f"[Combined] PNG: {combined_png}")
    print("[Done] All outputs exported.")


if __name__ == "__main__":
    main()