geMoldInsight/scripts/test_mold_splitting.py

"""
测试改进后的分模算法

功能：
1. 测试法向量分析
2. 测试真实分型线计算
3. 验证 AI 接口可用性
"""
import sys
from pathlib import Path

# 添加项目根目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from OCC.Core.STEPControl import STEPControl_Reader
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeBox
from core.mold_generator import MoldCavityGenerator
from utils.logger import get_logger

logger = get_logger(__name__)


def test_simple_shape():
    """测试简单形状（长方体）的分模"""
    print("\n" + "="*60)
    print("测试 1: 简单长方体分模")
    print("="*60)
    
    # 创建简单长方体
    box = BRepPrimAPI_MakeBox(100, 80, 50).Shape()
    
    # 创建模具生成器
    generator = MoldCavityGenerator(
        shrinkage_rate=0.005,
        draft_angle=2.0
    )
    
    # 生成模具型腔
    try:
        result = generator.generate_mold_cavities(box)
        
        print(f"✓ 分模成功")
        print(f"  - 分型面法向量：{result['analysis']['bounding_box']['dimensions']}")
        print(f"  - 分型线点数：{len(result['parting_line'])}")
        print(f"  - 产品体积：{result['analysis']['volume']:.2f} mm³")
        print(f"  - 产品重量：{generator._calculate_product_weight(result['analysis'])}")
        
        # 检查分型线是否合理
        if len(result['parting_line']) > 4:
            print(f"  ✓ 分型线使用真实几何计算（{len(result['parting_line'])} 个点）")
        else:
            print(f"  ⚠ 分型线使用简化矩形（4 个点）")
        
        return True
        
    except Exception as e:
        print(f"✗ 测试失败：{e}")
        import traceback
        traceback.print_exc()
        return False


def test_step_file(step_path: str):
    """测试 STEP 文件的分模"""
    print("\n" + "="*60)
    print(f"测试 2: STEP 文件分模 - {step_path}")
    print("="*60)
    
    # 读取 STEP 文件
    step_reader = STEPControl_Reader()
    status = step_reader.ReadFile(step_path)
    
    if status != 1:
        print(f"✗ STEP 文件读取失败")
        return False
    
    step_reader.TransferRoots()
    shape = step_reader.OneShape()
    
    # 创建模具生成器
    generator = MoldCavityGenerator(
        shrinkage_rate=0.005,
        draft_angle=2.0,
        material_density=1.05  # ABS
    )
    
    # 生成模具型腔
    try:
        result = generator.generate_mold_cavities(shape)
        
        print(f"✓ 分模成功")
        print(f"  - 边界框：{result['analysis']['bounding_box']['dimensions']}")
        print(f"  - 体积：{result['analysis']['volume']:.2f} mm³")
        print(f"  - 表面积：{result['analysis']['surface_area']:.2f} mm²")
        print(f"  - 分型线点数：{len(result['parting_line'])}")
        
        # 计算分型线长度
        parting_line_length = generator._calculate_parting_line_length(result['parting_line'])
        print(f"  - 分型线长度：{parting_line_length:.2f} mm")
        
        # 生成详细 JSON
        detailed_json = generator.generate_detailed_cavity_json(result)
        print(f"  ✓ 生成详细型腔数据")
        print(f"    - 型腔顶点数：{detailed_json['mold_cavities']['cavity']['vertex_count']}")
        print(f"    - 型芯顶点数：{detailed_json['mold_cavities']['core']['vertex_count']}")
        
        return True
        
    except Exception as e:
        print(f"✗ 测试失败：{e}")
        import traceback
        traceback.print_exc()
        return False


def test_ai_interface():
    """测试 AI 模型接口"""
    print("\n" + "="*60)
    print("测试 3: AI 模型接口")
    print("="*60)
    
    from core.ai_mold_assistant import AIPartingSurfaceDetector, AIDraftAnalyzer
    
    # 创建 AI 模型（示例）
    parting_detector = AIPartingSurfaceDetector()
    draft_analyzer = AIDraftAnalyzer()
    
    # 创建模具生成器
    generator = MoldCavityGenerator()
    
    # 设置 AI 模型
    generator.set_ai_model(
        parting_detector=parting_detector,
        draft_analyzer=draft_analyzer
    )
    
    print(f"✓ AI 模型接口已设置")
    print(f"  - 分型面检测器：{generator.ai_parting_detector is not None}")
    print(f"  - 拔模分析器：{generator.ai_draft_analyzer is not None}")
    
    # 测试简单形状
    box = BRepPrimAPI_MakeBox(50, 40, 30).Shape()
    
    try:
        result = generator.generate_mold_cavities(box)
        print(f"✓ 使用 AI 接口分模成功（AI 模型会回退到几何方法）")
        return True
    except Exception as e:
        print(f"✗ 测试失败：{e}")
        return False


def test_parting_line_calculation():
    """测试分型线计算算法"""
    print("\n" + "="*60)
    print("测试 4: 分型线计算算法")
    print("="*60)
    
    from OCC.Core.gp import gp_Pln, gp_Pnt, gp_Dir
    from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_MakeFace
    
    # 创建测试形状
    box = BRepPrimAPI_MakeBox(100, 80, 60).Shape()
    
    # 创建分型面（Z=30）
    parting_plane = gp_Pln(gp_Pnt(0, 0, 30), gp_Dir(0, 0, 1))
    parting_surface = BRepBuilderAPI_MakeFace(parting_plane).Face()
    
    # 创建模具生成器
    generator = MoldCavityGenerator()
    
    # 计算分型线
    parting_line = generator._calculate_parting_line(box, parting_surface)
    
    print(f"✓ 分型线计算完成")
    print(f"  - 点数：{len(parting_line)}")
    print(f"  - 长度：{generator._calculate_parting_line_length(parting_line):.2f} mm")
    
    # 打印前几个点
    if len(parting_line) > 0:
        print(f"  - 示例点：{parting_line[0]}")
    
    return True


def main():
    """运行所有测试"""
    print("\n" + "="*60)
    print("分模算法改进测试")
    print("="*60)
    
    results = []
    
    # 测试 1: 简单形状
    results.append(("简单长方体", test_simple_shape()))
    
    # 测试 2: STEP 文件（如果有）
    test_files = [
        "uploads/test.stp",
        "uploads/box.stp",
        "test.stp"
    ]
    
    for test_file in test_files:
        if Path(test_file).exists():
            results.append((f"STEP 文件 ({test_file})", test_step_file(test_file)))
            break
    else:
        print("\n⚠ 跳过 STEP 文件测试（未找到测试文件）")
    
    # 测试 3: AI 接口
    results.append(("AI 模型接口", test_ai_interface()))
    
    # 测试 4: 分型线算法
    results.append(("分型线计算", test_parting_line_calculation()))
    
    # 汇总结果
    print("\n" + "="*60)
    print("测试结果汇总")
    print("="*60)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for name, result in results:
        status = "✓ 通过" if result else "✗ 失败"
        print(f"{status}: {name}")
    
    print(f"\n总计：{passed}/{total} 测试通过")
    
    if passed == total:
        print("\n🎉 所有测试通过！")
        return 0
    else:
        print(f"\n⚠ {total - passed} 个测试失败")
        return 1


if __name__ == "__main__":
    sys.exit(main())