mcltspice/src/mcp_ltspice/drc.py

"""Design Rule Checks for LTspice schematics."""

from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path

from .schematic import Schematic, parse_schematic


class Severity(Enum):
    ERROR = "error"  # Will likely cause simulation failure
    WARNING = "warning"  # May cause unexpected results
    INFO = "info"  # Suggestion for improvement


@dataclass
class DRCViolation:
    """A single design rule violation."""

    rule: str  # Short rule identifier
    severity: Severity
    message: str  # Human-readable description
    component: str | None = None  # Related component name
    location: tuple[int, int] | None = None  # (x, y) if applicable


@dataclass
class DRCResult:
    """Results of a design rule check."""

    violations: list[DRCViolation] = field(default_factory=list)
    checks_run: int = 0

    @property
    def errors(self) -> list[DRCViolation]:
        return [v for v in self.violations if v.severity == Severity.ERROR]

    @property
    def warnings(self) -> list[DRCViolation]:
        return [v for v in self.violations if v.severity == Severity.WARNING]

    @property
    def passed(self) -> bool:
        return len(self.errors) == 0

    def summary(self) -> str:
        """Human-readable summary."""
        total = len(self.violations)
        err_count = len(self.errors)
        warn_count = len(self.warnings)
        info_count = total - err_count - warn_count

        if total == 0:
            return f"DRC passed: {self.checks_run} checks run, no violations found."

        parts = []
        if err_count:
            parts.append(f"{err_count} error{'s' if err_count != 1 else ''}")
        if warn_count:
            parts.append(f"{warn_count} warning{'s' if warn_count != 1 else ''}")
        if info_count:
            parts.append(f"{info_count} info")

        status = "FAILED" if err_count else "passed with warnings"
        return f"DRC {status}: {self.checks_run} checks run, {', '.join(parts)}."

    def to_dict(self) -> dict:
        """Convert to JSON-serializable dict."""
        return {
            "passed": self.passed,
            "checks_run": self.checks_run,
            "summary": self.summary(),
            "error_count": len(self.errors),
            "warning_count": len(self.warnings),
            "violations": [
                {
                    "rule": v.rule,
                    "severity": v.severity.value,
                    "message": v.message,
                    "component": v.component,
                    "location": list(v.location) if v.location else None,
                }
                for v in self.violations
            ],
        }


def run_drc(schematic_path: Path | str) -> DRCResult:
    """Run all design rule checks on a schematic.

    Args:
        schematic_path: Path to .asc file

    Returns:
        DRCResult with all violations found
    """
    sch = parse_schematic(schematic_path)
    result = DRCResult()

    _check_ground(sch, result)
    _check_floating_nodes(sch, result)
    _check_simulation_directive(sch, result)
    _check_voltage_source_loops(sch, result)
    _check_component_values(sch, result)
    _check_duplicate_names(sch, result)
    _check_unconnected_components(sch, result)

    return result


def _check_ground(sch: Schematic, result: DRCResult) -> None:
    """Check that at least one ground (node '0') exists."""
    result.checks_run += 1
    has_ground = any(f.name == "0" for f in sch.flags)
    if not has_ground:
        result.violations.append(
            DRCViolation(
                rule="NO_GROUND",
                severity=Severity.ERROR,
                message=(
                    "No ground node found. Every circuit needs at least one ground (0) connection."
                ),
            )
        )


def _check_floating_nodes(sch: Schematic, result: DRCResult) -> None:
    """Check for nodes with only one wire connection (likely floating).

    Build a connectivity map from wires and flags.
    A node is a unique (x,y) point where wires meet or flags are placed.
    If a point only has one wire endpoint and no flag, it might be floating.

    This is approximate -- we cannot fully determine connectivity without
    knowing component pin locations, but we can flag obvious cases.
    """
    result.checks_run += 1

    # Count how many wire endpoints touch each point
    point_count: dict[tuple[int, int], int] = defaultdict(int)
    for wire in sch.wires:
        point_count[(wire.x1, wire.y1)] += 1
        point_count[(wire.x2, wire.y2)] += 1

    # Flags also represent connections at a point
    flag_points = {(f.x, f.y) for f in sch.flags}

    # Component positions also represent connection points (approximately)
    component_points = {(c.x, c.y) for c in sch.components}

    for point, count in point_count.items():
        if count == 1 and point not in flag_points and point not in component_points:
            result.violations.append(
                DRCViolation(
                    rule="FLOATING_NODE",
                    severity=Severity.WARNING,
                    message=(
                        f"Possible floating node at ({point[0]}, {point[1]}). "
                        f"Wire endpoint has only one connection."
                    ),
                    location=point,
                )
            )


def _check_simulation_directive(sch: Schematic, result: DRCResult) -> None:
    """Check that at least one simulation directive exists."""
    result.checks_run += 1
    directives = sch.get_spice_directives()
    sim_types = [".tran", ".ac", ".dc", ".op", ".noise", ".tf"]
    has_sim = any(any(d.lower().startswith(s) for s in sim_types) for d in directives)
    if not has_sim:
        result.violations.append(
            DRCViolation(
                rule="NO_SIM_DIRECTIVE",
                severity=Severity.ERROR,
                message=("No simulation directive found. Add .tran, .ac, .dc, .op, etc."),
            )
        )


def _check_voltage_source_loops(sch: Schematic, result: DRCResult) -> None:
    """Check for voltage sources connected in parallel (short circuit).

    This is a simplified check -- look for voltage sources that share both
    pin nodes via wire connectivity. Two voltage sources whose positions
    connect to the same pair of nets would create a loop or conflict.
    """
    result.checks_run += 1

    # Build a union-find structure from wire connectivity so we can
    # determine which points belong to the same electrical net.
    parent: dict[tuple[int, int], tuple[int, int]] = {}

    def find(p: tuple[int, int]) -> tuple[int, int]:
        if p not in parent:
            parent[p] = p
        while parent[p] != p:
            parent[p] = parent[parent[p]]
            p = parent[p]
        return p

    def union(a: tuple[int, int], b: tuple[int, int]) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    # Each wire merges its two endpoints into the same net
    for wire in sch.wires:
        union((wire.x1, wire.y1), (wire.x2, wire.y2))

    # Also merge flag locations (named nets connect distant points with
    # the same name)
    flag_groups: dict[str, list[tuple[int, int]]] = defaultdict(list)
    for flag in sch.flags:
        flag_groups[flag.name].append((flag.x, flag.y))
    for pts in flag_groups.values():
        for pt in pts[1:]:
            union(pts[0], pt)

    # Find voltage sources and approximate their pin positions.
    # LTspice voltage sources have pins at the component origin and
    # offset along the component axis. Standard pin spacing is 64 units
    # vertically for a non-rotated voltage source (pin+ at top, pin- at
    # bottom relative to the symbol origin).
    voltage_sources = [c for c in sch.components if "voltage" in c.symbol.lower()]

    if len(voltage_sources) < 2:
        return

    # Estimate pin positions per voltage source based on rotation.
    # Default (R0): positive pin at (x, y-16), negative at (x, y+16)
    # We use a coarse offset; the exact value depends on the symbol but
    # 16 is a common half-pin-spacing in LTspice grid units.
    pin_offset = 16

    def _pin_positions(comp):
        """Return approximate (positive_pin, negative_pin) coordinates."""
        x, y = comp.x, comp.y
        rot = comp.rotation
        if rot == 0:
            return (x, y - pin_offset), (x, y + pin_offset)
        elif rot == 90:
            return (x + pin_offset, y), (x - pin_offset, y)
        elif rot == 180:
            return (x, y + pin_offset), (x, y - pin_offset)
        elif rot == 270:
            return (x - pin_offset, y), (x + pin_offset, y)
        return (x, y - pin_offset), (x, y + pin_offset)

    def _nearest_net(pin: tuple[int, int]) -> tuple[int, int]:
        """Find the nearest wire/flag point to a pin and return its net root.

        If the pin is directly on a known point, use it. Otherwise search
        within a small radius for the closest connected point.
        """
        if pin in parent:
            return find(pin)
        # Search nearby points (LTspice grid snap is typically 16 units)
        best = None
        best_dist = float("inf")
        for pt in parent:
            dx = pin[0] - pt[0]
            dy = pin[1] - pt[1]
            dist = dx * dx + dy * dy
            if dist < best_dist:
                best_dist = dist
                best = pt
        if best is not None and best_dist <= 32 * 32:
            return find(best)
        return pin  # isolated -- return pin itself as its own net

    # For each pair of voltage sources, check if they share both nets
    for i in range(len(voltage_sources)):
        pin_a_pos, pin_a_neg = _pin_positions(voltage_sources[i])
        net_a_pos = _nearest_net(pin_a_pos)
        net_a_neg = _nearest_net(pin_a_neg)
        for j in range(i + 1, len(voltage_sources)):
            pin_b_pos, pin_b_neg = _pin_positions(voltage_sources[j])
            net_b_pos = _nearest_net(pin_b_pos)
            net_b_neg = _nearest_net(pin_b_neg)

            # Parallel if both nets match (in either polarity)
            parallel = (net_a_pos == net_b_pos and net_a_neg == net_b_neg) or (
                net_a_pos == net_b_neg and net_a_neg == net_b_pos
            )
            if parallel:
                name_i = voltage_sources[i].name
                name_j = voltage_sources[j].name
                result.violations.append(
                    DRCViolation(
                        rule="VSOURCE_LOOP",
                        severity=Severity.ERROR,
                        message=(
                            f"Voltage sources '{name_i}' and '{name_j}' "
                            f"appear to be connected in parallel, which "
                            f"creates a short circuit / voltage conflict."
                        ),
                        component=name_i,
                    )
                )


def _check_component_values(sch: Schematic, result: DRCResult) -> None:
    """Check that components have values where expected.

    Resistors, capacitors, inductors should have values.
    Voltage/current sources should have values.
    Skip if the value looks like a parameter expression (e.g., "{R1}").
    """
    result.checks_run += 1

    # Map symbol substrings to human-readable type names
    value_required = {
        "res": "Resistor",
        "cap": "Capacitor",
        "ind": "Inductor",
        "voltage": "Voltage source",
        "current": "Current source",
    }

    for comp in sch.components:
        symbol_lower = comp.symbol.lower()
        matched_type = None
        for pattern, label in value_required.items():
            if pattern in symbol_lower:
                matched_type = label
                break

        if matched_type is None:
            continue

        val = comp.value
        if not val or not val.strip():
            result.violations.append(
                DRCViolation(
                    rule="MISSING_VALUE",
                    severity=Severity.WARNING,
                    message=(f"{matched_type} '{comp.name}' has no value set."),
                    component=comp.name,
                    location=(comp.x, comp.y),
                )
            )
        elif val.strip().startswith("{") and val.strip().endswith("}"):
            # Parameter expression -- valid, skip
            pass


def _check_duplicate_names(sch: Schematic, result: DRCResult) -> None:
    """Check for duplicate component instance names."""
    result.checks_run += 1
    seen: dict[str, bool] = {}
    for comp in sch.components:
        if not comp.name:
            continue
        if comp.name in seen:
            result.violations.append(
                DRCViolation(
                    rule="DUPLICATE_NAME",
                    severity=Severity.ERROR,
                    message=f"Duplicate component name '{comp.name}'.",
                    component=comp.name,
                    location=(comp.x, comp.y),
                )
            )
        else:
            seen[comp.name] = True


def _check_unconnected_components(sch: Schematic, result: DRCResult) -> None:
    """Check for components that don't seem to be connected to anything.

    A component at (x, y) should have wire endpoints near its pins.
    This is approximate without knowing exact pin positions, so we check
    whether any wire endpoint falls within a reasonable distance (16
    units -- one LTspice grid step) of the component's origin.
    """
    result.checks_run += 1

    proximity = 16  # LTspice grid spacing

    # Collect all wire endpoints into a set for fast lookup
    wire_points: set[tuple[int, int]] = set()
    for wire in sch.wires:
        wire_points.add((wire.x1, wire.y1))
        wire_points.add((wire.x2, wire.y2))

    # Also include flag positions (flags connect to nets)
    flag_points: set[tuple[int, int]] = set()
    for flag in sch.flags:
        flag_points.add((flag.x, flag.y))

    all_connection_points = wire_points | flag_points

    for comp in sch.components:
        if not comp.name:
            continue

        # Check if any connection point is within proximity of the
        # component origin. We scan a small bounding box rather than
        # iterating all points.
        connected = False
        for pt in all_connection_points:
            dx = abs(pt[0] - comp.x)
            dy = abs(pt[1] - comp.y)
            if dx <= proximity and dy <= proximity:
                connected = True
                break

        if not connected:
            result.violations.append(
                DRCViolation(
                    rule="UNCONNECTED_COMPONENT",
                    severity=Severity.WARNING,
                    message=(
                        f"Component '{comp.name}' ({comp.symbol}) at "
                        f"({comp.x}, {comp.y}) has no nearby wire "
                        f"connections."
                    ),
                    component=comp.name,
                    location=(comp.x, comp.y),
                )
            )