kicad-mcp/kicad_mcp/utils/model3d_analyzer.py

"""
3D Model Analysis utilities for KiCad PCB files.

Provides functionality to analyze 3D models, visualizations, and mechanical constraints
from KiCad PCB files including component placement, clearances, and board dimensions.
"""

from dataclasses import dataclass
import logging
import re
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class Component3D:
    """Represents a 3D component with position and model information."""
    reference: str
    position: tuple[float, float, float]  # X, Y, Z coordinates in mm
    rotation: tuple[float, float, float]  # Rotation around X, Y, Z axes
    model_path: str | None
    model_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
    model_offset: tuple[float, float, float] = (0.0, 0.0, 0.0)
    footprint: str | None = None
    value: str | None = None


@dataclass
class BoardDimensions:
    """PCB board physical dimensions and constraints."""
    width: float  # mm
    height: float  # mm
    thickness: float  # mm
    outline_points: list[tuple[float, float]]  # Board outline coordinates
    holes: list[tuple[float, float, float]]  # Hole positions and diameters
    keepout_areas: list[dict[str, Any]]  # Keepout zones


@dataclass
class MechanicalAnalysis:
    """Results of mechanical/3D analysis."""
    board_dimensions: BoardDimensions
    components: list[Component3D]
    clearance_violations: list[dict[str, Any]]
    height_analysis: dict[str, float]  # min, max, average heights
    mechanical_constraints: list[str]  # Constraint violations or warnings


class Model3DAnalyzer:
    """Analyzer for 3D models and mechanical aspects of KiCad PCBs."""

    def __init__(self, pcb_file_path: str):
        """Initialize with PCB file path."""
        self.pcb_file_path = pcb_file_path
        self.pcb_data = None
        self._load_pcb_data()

    def _load_pcb_data(self) -> None:
        """Load and parse PCB file data."""
        try:
            with open(self.pcb_file_path, encoding='utf-8') as f:
                content = f.read()
                # Parse S-expression format (simplified)
                self.pcb_data = content
        except Exception as e:
            logger.error(f"Failed to load PCB file {self.pcb_file_path}: {e}")
            self.pcb_data = None

    def extract_3d_components(self) -> list[Component3D]:
        """Extract 3D component information from PCB data."""
        components = []

        if not self.pcb_data:
            return components

        # Parse footprint modules with 3D models
        footprint_pattern = r'\(footprint\s+"([^"]+)"[^)]*\(at\s+([\d.-]+)\s+([\d.-]+)(?:\s+([\d.-]+))?\)'
        model_pattern = r'\(model\s+"([^"]+)"[^)]*\(at\s+\(xyz\s+([\d.-]+)\s+([\d.-]+)\s+([\d.-]+)\)\)[^)]*\(scale\s+\(xyz\s+([\d.-]+)\s+([\d.-]+)\s+([\d.-]+)\)\)'
        reference_pattern = r'\(fp_text\s+reference\s+"([^"]+)"'
        value_pattern = r'\(fp_text\s+value\s+"([^"]+)"'

        # Find all footprints
        for footprint_match in re.finditer(footprint_pattern, self.pcb_data, re.MULTILINE):
            footprint_name = footprint_match.group(1)
            x_pos = float(footprint_match.group(2))
            y_pos = float(footprint_match.group(3))
            rotation = float(footprint_match.group(4)) if footprint_match.group(4) else 0.0

            # Extract the footprint section
            start_pos = footprint_match.start()
            footprint_section = self._extract_footprint_section(start_pos)

            # Find reference and value within this footprint
            ref_match = re.search(reference_pattern, footprint_section)
            val_match = re.search(value_pattern, footprint_section)

            reference = ref_match.group(1) if ref_match else "Unknown"
            value = val_match.group(1) if val_match else ""

            # Find 3D model within this footprint
            model_match = re.search(model_pattern, footprint_section)

            if model_match:
                model_path = model_match.group(1)
                model_x = float(model_match.group(2))
                model_y = float(model_match.group(3))
                model_z = float(model_match.group(4))
                scale_x = float(model_match.group(5))
                scale_y = float(model_match.group(6))
                scale_z = float(model_match.group(7))

                component = Component3D(
                    reference=reference,
                    position=(x_pos, y_pos, 0.0),  # Z will be calculated from model
                    rotation=(0.0, 0.0, rotation),
                    model_path=model_path,
                    model_scale=(scale_x, scale_y, scale_z),
                    model_offset=(model_x, model_y, model_z),
                    footprint=footprint_name,
                    value=value
                )
                components.append(component)

        logger.info(f"Extracted {len(components)} 3D components from PCB")
        return components

    def _extract_footprint_section(self, start_pos: int) -> str:
        """Extract a complete footprint section from PCB data."""
        if not self.pcb_data:
            return ""

        # Find the matching closing parenthesis
        level = 0
        i = start_pos
        while i < len(self.pcb_data):
            if self.pcb_data[i] == '(':
                level += 1
            elif self.pcb_data[i] == ')':
                level -= 1
                if level == 0:
                    return self.pcb_data[start_pos:i+1]
            i += 1

        return self.pcb_data[start_pos:start_pos + 10000]  # Fallback

    def analyze_board_dimensions(self) -> BoardDimensions:
        """Analyze board physical dimensions and constraints."""
        if not self.pcb_data:
            return BoardDimensions(0, 0, 1.6, [], [], [])

        # Extract board outline (Edge.Cuts layer)
        edge_pattern = r'\(gr_line\s+\(start\s+([\d.-]+)\s+([\d.-]+)\)\s+\(end\s+([\d.-]+)\s+([\d.-]+)\)\s+\(stroke[^)]*\)\s+\(layer\s+"Edge\.Cuts"\)'

        outline_points = []
        for match in re.finditer(edge_pattern, self.pcb_data):
            start_x, start_y = float(match.group(1)), float(match.group(2))
            end_x, end_y = float(match.group(3)), float(match.group(4))
            outline_points.extend([(start_x, start_y), (end_x, end_y)])

        # Calculate board dimensions
        if outline_points:
            x_coords = [p[0] for p in outline_points]
            y_coords = [p[1] for p in outline_points]
            width = max(x_coords) - min(x_coords)
            height = max(y_coords) - min(y_coords)
        else:
            width = height = 0

        # Extract board thickness from stackup (if available) or default to 1.6mm
        thickness = 1.6
        thickness_pattern = r'\(thickness\s+([\d.]+)\)'
        thickness_match = re.search(thickness_pattern, self.pcb_data)
        if thickness_match:
            thickness = float(thickness_match.group(1))

        # Find holes
        holes = []
        hole_pattern = r'\(pad[^)]*\(type\s+thru_hole\)[^)]*\(at\s+([\d.-]+)\s+([\d.-]+)\)[^)]*\(size\s+([\d.-]+)'
        for match in re.finditer(hole_pattern, self.pcb_data):
            x, y, diameter = float(match.group(1)), float(match.group(2)), float(match.group(3))
            holes.append((x, y, diameter))

        return BoardDimensions(
            width=width,
            height=height,
            thickness=thickness,
            outline_points=list(set(outline_points)),  # Remove duplicates
            holes=holes,
            keepout_areas=[]  # TODO: Extract keepout zones
        )

    def analyze_component_heights(self, components: list[Component3D]) -> dict[str, float]:
        """Analyze component height distribution."""
        heights = []

        for component in components:
            if component.model_path:
                # Estimate height from model scale and type
                estimated_height = self._estimate_component_height(component)
                heights.append(estimated_height)

        if not heights:
            return {"min": 0, "max": 0, "average": 0, "count": 0}

        return {
            "min": min(heights),
            "max": max(heights),
            "average": sum(heights) / len(heights),
            "count": len(heights)
        }

    def _estimate_component_height(self, component: Component3D) -> float:
        """Estimate component height based on footprint and model."""
        # Component height estimation based on common footprint patterns
        footprint_heights = {
            # SMD packages
            "0402": 0.6,
            "0603": 0.95,
            "0805": 1.35,
            "1206": 1.7,

            # IC packages
            "SOIC": 2.65,
            "QFP": 1.75,
            "BGA": 1.5,
            "TQFP": 1.4,

            # Through-hole
            "DIP": 4.0,
            "TO-220": 4.5,
            "TO-92": 4.5,
        }

        # Check footprint name for height hints
        footprint = component.footprint or ""
        for pattern, height in footprint_heights.items():
            if pattern in footprint.upper():
                return height * component.model_scale[2]  # Apply Z scaling

        # Default height based on model scale
        return 2.0 * component.model_scale[2]

    def check_clearance_violations(self, components: list[Component3D],
                                 board_dims: BoardDimensions) -> list[dict[str, Any]]:
        """Check for 3D clearance violations between components."""
        violations = []

        # Component-to-component clearance
        for i, comp1 in enumerate(components):
            for j, comp2 in enumerate(components[i+1:], i+1):
                distance = self._calculate_3d_distance(comp1, comp2)
                min_clearance = self._get_minimum_clearance(comp1, comp2)

                if distance < min_clearance:
                    violations.append({
                        "type": "component_clearance",
                        "component1": comp1.reference,
                        "component2": comp2.reference,
                        "distance": distance,
                        "required_clearance": min_clearance,
                        "severity": "warning" if distance > min_clearance * 0.8 else "error"
                    })

        # Board edge clearance
        for component in components:
            edge_distance = self._distance_to_board_edge(component, board_dims)
            min_edge_clearance = 0.5  # 0.5mm minimum edge clearance

            if edge_distance < min_edge_clearance:
                violations.append({
                    "type": "board_edge_clearance",
                    "component": component.reference,
                    "distance": edge_distance,
                    "required_clearance": min_edge_clearance,
                    "severity": "warning"
                })

        return violations

    def _calculate_3d_distance(self, comp1: Component3D, comp2: Component3D) -> float:
        """Calculate 3D distance between two components."""
        dx = comp1.position[0] - comp2.position[0]
        dy = comp1.position[1] - comp2.position[1]
        dz = comp1.position[2] - comp2.position[2]
        return (dx*dx + dy*dy + dz*dz) ** 0.5

    def _get_minimum_clearance(self, comp1: Component3D, comp2: Component3D) -> float:
        """Get minimum required clearance between components."""
        # Base clearance rules (can be made more sophisticated)
        base_clearance = 0.2  # 0.2mm base clearance

        # Larger clearance for high-power components
        if any(keyword in (comp1.value or "") + (comp2.value or "")
               for keyword in ["POWER", "REGULATOR", "MOSFET"]):
            return base_clearance + 1.0

        return base_clearance

    def _distance_to_board_edge(self, component: Component3D,
                               board_dims: BoardDimensions) -> float:
        """Calculate minimum distance from component to board edge."""
        if not board_dims.outline_points:
            return float('inf')

        # Simplified calculation - distance to bounding rectangle
        x_coords = [p[0] for p in board_dims.outline_points]
        y_coords = [p[1] for p in board_dims.outline_points]

        min_x, max_x = min(x_coords), max(x_coords)
        min_y, max_y = min(y_coords), max(y_coords)

        comp_x, comp_y = component.position[0], component.position[1]

        # Distance to each edge
        distances = [
            comp_x - min_x,  # Left edge
            max_x - comp_x,  # Right edge
            comp_y - min_y,  # Bottom edge
            max_y - comp_y   # Top edge
        ]

        return min(distances)

    def generate_3d_visualization_data(self) -> dict[str, Any]:
        """Generate data structure for 3D visualization."""
        components = self.extract_3d_components()
        board_dims = self.analyze_board_dimensions()
        height_analysis = self.analyze_component_heights(components)
        clearance_violations = self.check_clearance_violations(components, board_dims)

        return {
            "board_dimensions": {
                "width": board_dims.width,
                "height": board_dims.height,
                "thickness": board_dims.thickness,
                "outline": board_dims.outline_points,
                "holes": board_dims.holes
            },
            "components": [
                {
                    "reference": comp.reference,
                    "position": comp.position,
                    "rotation": comp.rotation,
                    "model_path": comp.model_path,
                    "footprint": comp.footprint,
                    "value": comp.value,
                    "estimated_height": self._estimate_component_height(comp)
                }
                for comp in components
            ],
            "height_analysis": height_analysis,
            "clearance_violations": clearance_violations,
            "stats": {
                "total_components": len(components),
                "components_with_3d_models": len([c for c in components if c.model_path]),
                "violation_count": len(clearance_violations)
            }
        }

    def perform_mechanical_analysis(self) -> MechanicalAnalysis:
        """Perform comprehensive mechanical analysis."""
        components = self.extract_3d_components()
        board_dims = self.analyze_board_dimensions()
        height_analysis = self.analyze_component_heights(components)
        clearance_violations = self.check_clearance_violations(components, board_dims)

        # Generate mechanical constraints and warnings
        constraints = []

        if height_analysis["max"] > 10.0:  # 10mm height limit example
            constraints.append(f"Board height {height_analysis['max']:.1f}mm exceeds 10mm limit")

        if board_dims.width > 100 or board_dims.height > 100:
            constraints.append(f"Board dimensions {board_dims.width:.1f}x{board_dims.height:.1f}mm are large")

        if len(clearance_violations) > 0:
            constraints.append(f"{len(clearance_violations)} clearance violations found")

        return MechanicalAnalysis(
            board_dimensions=board_dims,
            components=components,
            clearance_violations=clearance_violations,
            height_analysis=height_analysis,
            mechanical_constraints=constraints
        )


def analyze_pcb_3d_models(pcb_file_path: str) -> dict[str, Any]:
    """Convenience function to analyze 3D models in a PCB file."""
    try:
        analyzer = Model3DAnalyzer(pcb_file_path)
        return analyzer.generate_3d_visualization_data()
    except Exception as e:
        logger.error(f"Failed to analyze 3D models in {pcb_file_path}: {e}")
        return {"error": str(e)}


def get_mechanical_constraints(pcb_file_path: str) -> MechanicalAnalysis:
    """Get mechanical analysis and constraints for a PCB."""
    analyzer = Model3DAnalyzer(pcb_file_path)
    return analyzer.perform_mechanical_analysis()