geMoldInsight/src/core/mold_generator.py

from typing import Dict, List, Any, Tuple, Optional
import numpy as np
from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_MakeFace
from OCC.Core.gp import gp_Pln, gp_Dir, gp_Pnt
from OCC.Core.TopoDS import TopoDS_Face
from OCC.Core.BRepAdaptor import BRepAdaptor_Surface
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.TopAbs import TopAbs_FACE
from OCC.Core.Bnd import Bnd_Box
from OCC.Core.BRepBndLib import brepbndlib_Add

from models.schemas import create_mold_cavity_data, create_mold_key_info
from utils.logger import get_logger
from core.base_mold_generator import BaseMoldGenerator
from core.side_action_designer import SideActionDesigner

logger = get_logger(__name__)


class MoldCavityGenerator(BaseMoldGenerator):
    """模具型腔生成器 - 基于产品模型生成Cavity和Core"""

    def __init__(self, shrinkage_rate: float = 0.005, draft_angle: float = 2.0,
                 material_density: float = 1.05):
        super().__init__(shrinkage_rate, draft_angle, material_density)

        self.material_densities = {
            "ABS": 1.05,
            "PP": 0.90,
            "PC": 1.20,
            "PE": 0.95,
            "PS": 1.05,
            "PA": 1.14,
            "POM": 1.42,
            "PMMA": 1.18
        }

        self.parting_line_tolerance = 0.1
        self.max_draft_angle = 5.0
        self.side_action_designer = SideActionDesigner()

    def set_material(self, material: str):
        """设置产品材料"""
        if material in self.material_densities:
            self.material_density = self.material_densities[material]
            logger.info(f"材料设置为 {material}, 密度: {self.material_density} g/cm³")
        else:
            logger.warning(f"未知材料 {material}, 使用默认密度 {self.material_density} g/cm³")

    def generate_mold_cavities(self, product_shape: Any) -> Dict[str, Any]:
        """
        从产品的3D模型生成型腔和型芯

        Returns:
            {
                "cavity": cavity_shape,
                "core": core_shape,
                "parting_surface": parting_surface,
                "parting_line": parting_line
            }
        """
        logger.info("开始生成模具型腔...")

        try:
            analysis = self._analyze_product_geometry(product_shape)

            parting_result = self._detect_primary_parting(product_shape, analysis)
            parting_surface = parting_result["surface"]
            parting_line = self.optimize_parting_line(parting_result["line"])
            parting_direction = parting_result["direction"]

            side_action_result = self.side_action_designer.analyze_and_design(
                shape=product_shape,
                parting_direction=parting_direction,
                mold_size=self._calculate_mold_size(analysis),
                parting_surface=parting_surface,
            )
            undercut_regions = self._build_undercut_regions(
                side_action_result.get("undercut_analysis", {})
            )

            scaled_shape = self._apply_shrinkage_compensation(product_shape)

            drafted_shape = self._apply_draft_angles(scaled_shape, parting_surface)

            cavity, core = self._split_cavity_core(drafted_shape, parting_surface)

            logger.info("模具型腔生成完成")

            return {
                "cavity": cavity,
                "core": core,
                "parting_surface": parting_surface,
                "parting_line": parting_line,
                "analysis": analysis,
                "undercut_regions": undercut_regions,
                "side_actions": side_action_result,
            }

        except Exception as e:
            logger.error(f"模具型腔生成失败: {e}")
            raise

    def generate_detailed_cavity_json(self, cavity_data: Dict) -> Dict[str, Any]:
        """
        生成详细的型腔三维JSON数据

        Returns:
            包含完整几何信息的JSON结构
        """
        cavity = cavity_data["cavity"]
        core = cavity_data["core"]
        parting_surface = cavity_data["parting_surface"]
        analysis = cavity_data["analysis"]

        cavity_geometry = self._extract_shape_geometry(cavity, "cavity")
        core_geometry = self._extract_shape_geometry(core, "core")

        parting_geometry = self._extract_parting_surface_geometry(
            parting_surface
        )

        detailed_json = {
            "metadata": {
                "version": "2.0",
                "generated_at": str(np.datetime64('now')),
                "shrinkage_rate": self.shrinkage_rate,
                "draft_angle": self.draft_angle,
                "unit": "mm"
            },
            "product_analysis": {
                "bounding_box": analysis.get("bounding_box", {}),
                "volume": analysis.get("volume", 0),
                "surface_area": analysis.get("surface_area", 0),
                "center_of_mass": analysis.get("center_of_mass", [0, 0, 0])
            },
            "mold_cavities": {
                "cavity": cavity_geometry,
                "core": core_geometry
            },
            "parting_surface": parting_geometry,
            "quality_checks": {
                "undercut_regions": cavity_data.get("undercut_regions", []),
                "side_actions": cavity_data.get("side_actions", {}),
            },
            "manufacturing_info": {
                "estimated_mold_size": self._calculate_mold_size(analysis),
                "estimated_clamping_force": self._calculate_clamping_force(analysis),
                "recommended_material": self._get_recommended_material()
            }
        }

        return detailed_json

    def generate_cavity_key_info(self, cavity_data: Dict) -> Dict[str, Any]:
        """
        生成模具型腔的关键信息

        Returns:
            关键参数摘要
        """
        analysis = cavity_data["analysis"]

        key_info = {
            "mold_parameters": {
                "shrinkage_rate": f"{self.shrinkage_rate * 100:.2f}%",
                "draft_angle": f"{self.draft_angle}°",
                "parting_line_length": self._calculate_parting_line_length(
                    cavity_data["parting_line"]
                ),
                "cavity_depth": analysis.get("bounding_box", {}).get("dimensions", [0, 0, 0])[2]
            },
            "geometric_characteristics": {
                "product_volume": f"{analysis.get('volume', 0) / 1000:.2f} cm³",
                "product_weight": self._calculate_product_weight(analysis),
                "wall_thickness_range": self._estimate_wall_thickness(analysis),
                "complexity_score": self._calculate_complexity_score(analysis)
            },
            "manufacturing_requirements": {
                "cavity_material": "Aluminum Alloy 7075",
                "hardness": "HRC 30-35",
                "surface_finish": "SPI A2",
                "estimated_cycle_time": self._estimate_cycle_time(analysis),
                "recommended_injection_pressure": "80-120 MPa"
            },
            "quality_considerations": {
                "undercut_count": len(cavity_data.get("undercut_regions", [])),
                "side_action_summary": cavity_data.get("side_actions", {}).get("summary", {}),
                "potential_weld_lines": self._identify_weld_line_risk(analysis),
                "sink_mark_areas": self._identify_sink_mark_risk(analysis),
                "warpage_risk": self._assess_warpage_risk(analysis)
            }
        }

        return key_info

    # ==================== 内部方法 ====================

    def _detect_parting_surface(self, shape: Any, analysis: Dict) -> Tuple[Any, List]:
        """
        检测分型面和分型线

        优先级：
        1. AI 模型检测（如果已设置）
        2. 基于法向量分析的几何方法
        3. 简化方法（基于边界框）
        """
        try:
            parting_result = self._detect_primary_parting(shape, analysis)
            logger.info(
                f"使用 {parting_result['method']} 方法检测分型面，"
                f"置信度={parting_result['confidence']:.3f}"
            )
            return parting_result["surface"], self.optimize_parting_line(parting_result["line"])

        except Exception as e:
            logger.warning(f"法向量分析失败，使用简化方法：{e}")

        logger.info("使用简化方法检测分型面")
        return self._simple_parting_surface(shape, analysis)

    def _detect_primary_parting(self, shape: Any, analysis: Dict) -> Dict[str, Any]:
        """检测主分型面（AI优先 → 几何法向量 → 简化回退）"""
        if self.ai_parting_detector is not None:
            try:
                ai_result = self.ai_parting_detector.detect(shape, analysis)
                if ai_result is not None:
                    surface, line = self._create_parting_surface_from_ai(ai_result, analysis, shape)
                    return {
                        "surface": surface,
                        "line": line,
                        "direction": ai_result.get("normal", [0, 0, 1]),
                        "method": ai_result.get("method", "ai"),
                        "confidence": ai_result.get("confidence", 0.8),
                    }
            except Exception as e:
                logger.warning(f"AI 分型面检测失败: {e}")

        try:
            normal_dir = self._analyze_face_normals(shape)
            parting_plane = self._create_optimal_parting_plane(shape, analysis, normal_dir)
            dims = analysis.get("bounding_box", {}).get("dimensions", [100, 100, 100])
            span = max(dims) * 1.5 + 30
            parting_surface = BRepBuilderAPI_MakeFace(
                parting_plane, -span, span, -span, span
            ).Face()
            parting_surface = self.extend_parting_surface(parting_surface, shape, extension=30.0)
            parting_line = self._calculate_parting_line(shape, parting_surface)
            return {
                "surface": parting_surface,
                "line": parting_line,
                "direction": [float(normal_dir.X()), float(normal_dir.Y()), float(normal_dir.Z())],
                "method": "face_normal_analysis",
                "confidence": 0.85,
            }
        except Exception as e:
            logger.warning(f"法向量分析失败，使用简化方法：{e}")

        surface, line = self._simple_parting_surface(shape, analysis)
        return {
            "surface": surface,
            "line": line,
            "direction": [0, 0, 1],
            "method": "simple",
            "confidence": 0.6,
        }

    def _build_undercut_regions(self, undercut_analysis: Dict[str, Any]) -> List[Dict[str, Any]]:
        """将侧向机构分析结果转换为兼容旧结构的倒扣区域列表。"""
        undercut_faces = undercut_analysis.get("undercut_faces", [])
        regions = []

        for face in undercut_faces:
            regions.append({
                "type": "negative_draft",
                "location": face.get("center", [0, 0, 0]),
                "severity": face.get("severity", "medium"),
                "area": face.get("area", 0),
                "is_outer": face.get("is_outer", False),
                "face_index": face.get("face_index"),
            })

        logger.info(f"转换得到 {len(regions)} 个兼容倒扣区域")
        return regions

    def _analyze_face_normals(self, shape: Any) -> gp_Dir:
        """
        分析产品表面的法向量分布，找出最优分型方向

        原理：
        - 统计所有面的法向量
        - 选择法向量变化最小的方向作为分型方向
        - 避免倒扣（undercut）区域
        """
        face_normals = []
        explorer = TopExp_Explorer(shape, TopAbs_FACE)

        while explorer.More():
            face = TopoDS_Face(explorer.Current())
            surface = BRepAdaptor_Surface(face)

            try:
                if surface.GetType() == 0:
                    normal = surface.Plane().Position().Direction()
                else:
                    bbox = Bnd_Box()
                    brepbndlib_Add(face, bbox)
                    normal = gp_Dir(0, 0, 1)

                face_normals.append(normal)
            except Exception as e:
                logger.debug(f"面法向量计算失败：{e}")

            explorer.Next()

        if not face_normals:
            return gp_Dir(0, 0, 1)

        avg_x = sum(n.X() for n in face_normals) / len(face_normals)
        avg_y = sum(n.Y() for n in face_normals) / len(face_normals)
        avg_z = sum(n.Z() for n in face_normals) / len(face_normals)

        length = np.sqrt(avg_x**2 + avg_y**2 + avg_z**2)
        if length > 0.001:
            return gp_Dir(avg_x/length, avg_y/length, avg_z/length)
        else:
            return gp_Dir(0, 0, 1)

    def _create_optimal_parting_plane(self, shape: Any, analysis: Dict,
                                      direction: gp_Dir) -> gp_Pln:
        """
        创建最优分型面

        Args:
            shape: 产品形状
            analysis: 几何分析结果
            direction: 分型方向（法向量）

        Returns:
            gp_Pln: 分型面方程
        """
        bbox = analysis["bounding_box"]
        center = bbox["center"]

        parting_plane = gp_Pln(
            gp_Pnt(center[0], center[1], center[2]),
            direction
        )

        logger.info(f"创建分型面：原点=({center[0]:.2f}, {center[1]:.2f}, {center[2]:.2f}), "
                   f"法向量=({direction.X():.3f}, {direction.Y():.3f}, {direction.Z():.3f})")

        return parting_plane

    def _simple_parting_surface(self, shape: Any, analysis: Dict) -> Tuple[Any, List]:
        """简化的分型面检测（回退方案）"""
        bbox = analysis["bounding_box"]
        center_z = bbox["center"][2]

        parting_plane = gp_Pln(
            gp_Pnt(0, 0, center_z),
            gp_Dir(0, 0, 1)
        )
        parting_surface = BRepBuilderAPI_MakeFace(
            parting_plane,
            bbox["min"][0] - 10, bbox["max"][0] + 10,
            bbox["min"][1] - 10, bbox["max"][1] + 10
        ).Face()

        parting_line = self._simple_parting_line(shape)

        return parting_surface, parting_line

    def _create_parting_surface_from_ai(self, ai_result: Dict,
                                        analysis: Dict, shape: Any = None) -> Tuple[Any, List]:
        """
        从 AI 模型结果创建分型面（预留接口）

        Args:
            ai_result: AI 模型输出，应包含：
                - origin: [x, y, z] 平面原点
                - normal: [nx, ny, nz] 法向量
            analysis: 几何分析结果
            shape: 产品形状（用于计算分型线）

        Returns:
            (parting_surface, parting_line)
        """
        origin = ai_result.get("origin", [0, 0, 0])
        normal = ai_result.get("normal", [0, 0, 1])

        parting_plane = gp_Pln(
            gp_Pnt(origin[0], origin[1], origin[2]),
            gp_Dir(normal[0], normal[1], normal[2])
        )
        parting_surface = BRepBuilderAPI_MakeFace(parting_plane).Face()

        if "parting_line" in ai_result:
            parting_line = ai_result["parting_line"]
        elif shape is not None:
            parting_line = self._calculate_parting_line(shape, parting_surface)
        else:
            parting_line = []

        logger.info(f"从 AI 结果创建分型面：原点={origin}, 法向量={normal}")
        return parting_surface, parting_line

    def _extract_parting_surface_geometry(self, surface: Any) -> Dict[str, Any]:
        """提取分型面几何数据"""
        metadata = self._extract_plane_metadata(surface)

        return {
            "type": "plane",
            "normal": metadata["normal"],
            "origin": metadata["origin"],
            "bounds": metadata["bounds"],
        }

    def _calculate_mold_size(self, analysis: Dict) -> Dict[str, float]:
        """估算模具尺寸"""
        product_bbox = analysis["bounding_box"]["dimensions"]
        margin = 30

        return {
            "length": product_bbox[0] + 2 * margin,
            "width": product_bbox[1] + 2 * margin,
            "height": product_bbox[2] + 2 * margin + 100,
            "margin": margin
        }

    def _calculate_clamping_force(self, analysis: Dict) -> str:
        """估算锁模力"""
        volume_cm3 = analysis.get("volume", 0) / 1000

        if volume_cm3 < 10:
            return "50-100 吨"
        elif volume_cm3 < 100:
            return "150-300 吨"
        elif volume_cm3 < 500:
            return "400-600 吨"
        else:
            return "800+ 吨"

    def _get_recommended_material(self) -> str:
        """推荐模具材料"""
        return "Aluminum Alloy 7075 (铝合金模具)"

    def _estimate_wall_thickness(self, analysis: Dict) -> str:
        """估算壁厚范围"""
        volume = analysis.get("volume", 0)
        surface_area = analysis.get("surface_area", 0)

        if surface_area > 0 and volume > 0:
            avg_thickness = (volume / surface_area) * 0.6
            return f"{avg_thickness * 0.7:.2f} - {avg_thickness * 1.3:.2f} mm"
        elif volume > 0:
            bbox_dims = analysis.get("bounding_box", {}).get("dimensions", [1, 1, 1])
            bbox_volume = bbox_dims[0] * bbox_dims[1] * bbox_dims[2]
            if bbox_volume > 0:
                efficiency = volume / bbox_volume
                avg_thickness = (bbox_dims[0] + bbox_dims[1]) / 2 * efficiency
                return f"{avg_thickness * 0.7:.2f} - {avg_thickness * 1.3:.2f} mm"

        return "2.0 - 4.0 mm (默认)"

    def _calculate_complexity_score(self, analysis: Dict) -> float:
        """计算复杂度评分（0-10）"""
        volume = analysis.get("volume", 0)
        surface_area = analysis.get("surface_area", 0)

        if surface_area > 0 and volume > 0:
            thickness_ratio = (volume / surface_area) * 0.6
            complexity = min(thickness_ratio / 5.0, 10.0)
            return round(complexity, 1)
        elif volume > 0:
            bbox_dims = analysis.get("bounding_box", {}).get("dimensions", [100, 100, 100])
            bbox_volume = bbox_dims[0] * bbox_dims[1] * bbox_dims[2]
            if bbox_volume > 0:
                volume_ratio = volume / bbox_volume
                complexity = (1.0 - volume_ratio) * 10
                return round(min(max(complexity, 0), 10), 1)

        return 5.0

    def _estimate_cycle_time(self, analysis: Dict) -> str:
        """估算成型周期"""
        volume_cm3 = analysis.get("volume", 0) / 1000

        if volume_cm3 < 10:
            return "15-25 秒"
        elif volume_cm3 < 50:
            return "25-40 秒"
        elif volume_cm3 < 200:
            return "40-60 秒"
        else:
            return "60-90 秒"

    def _identify_weld_line_risk(self, analysis: Dict) -> str:
        """识别熔接痕风险"""
        complexity = self._calculate_complexity_score(analysis)

        if complexity > 7:
            return "高 - 建议优化浇口位置"
        elif complexity > 4:
            return "中 - 需仿真验证"
        else:
            return "低"

    def _identify_sink_mark_risk(self, analysis: Dict) -> str:
        """识别缩痕风险"""
        return "中 - 建议壁厚均匀性检查"