pyNgSpice/pyngspice/runner.py

"""
SpiceRunner interface and implementations for pyTesla integration.

Provides a common interface for running SPICE simulations with different
backends (embedded ngspice, subprocess ngspice, etc.). pyTesla uses
this interface to swap between LTspice and ngspice seamlessly.

Usage:
    from pyngspice import NgspiceRunner

    runner = NgspiceRunner(working_directory="./output")
    raw_file, log_file = runner.run("circuit.net")

    # Or with auto-detection:
    from pyngspice.runner import get_runner
    runner = get_runner("./output")
"""

import os
import shutil
import subprocess
import sys
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional, Tuple

from .netlist import preprocess_netlist


def _build_probe_mapping(processed_netlist_path: str) -> dict:
    """Build a mapping from probe trace names to original capacitor trace names.

    Reads the processed netlist to find V_probe_* lines and builds:
        {'i(v_probe_c_mmc)': 'i(c_mmc)', ...}

    All names are lowercased to match ngspice's raw file convention.

    Args:
        processed_netlist_path: Path to the preprocessed netlist file

    Returns:
        Dict mapping probe trace names to original capacitor trace names
    """
    mapping = {}
    try:
        with open(processed_netlist_path, 'r') as f:
            for line in f:
                tokens = line.strip().split()
                if not tokens:
                    continue
                name = tokens[0]
                if name.upper().startswith('V_PROBE_'):
                    # V_probe_C_mmc -> C_mmc
                    cap_name = name[8:]  # strip 'V_probe_'
                    old_trace = f'i({name.lower()})'
                    new_trace = f'i({cap_name.lower()})'
                    mapping[old_trace] = new_trace
    except (OSError, IOError):
        pass
    return mapping


def _postprocess_raw_file(raw_file: str, processed_netlist_path: str) -> None:
    """Rename capacitor probe traces in raw file header.

    After simulation, the raw file contains traces like i(v_probe_c_mmc).
    This function renames them back to i(c_mmc) so pyTesla sees the
    expected capacitor current names.

    Only modifies the ASCII header; the binary data section is untouched.

    Args:
        raw_file: Path to the .raw file to post-process
        processed_netlist_path: Path to the preprocessed netlist (to find probes)
    """
    mapping = _build_probe_mapping(processed_netlist_path)
    if not mapping:
        return

    try:
        with open(raw_file, 'rb') as f:
            content = f.read()
    except (OSError, IOError):
        return

    # Find header/data boundary
    header_end = -1
    for marker in [b'Binary:\n', b'Binary:\r\n', b'Values:\n', b'Values:\r\n']:
        pos = content.find(marker)
        if pos >= 0:
            header_end = pos + len(marker)
            break

    if header_end < 0:
        return  # Can't find boundary, skip

    header = content[:header_end].decode('ascii', errors='replace')
    data = content[header_end:]

    # Apply renames in header
    for old_name, new_name in mapping.items():
        header = header.replace(old_name, new_name)

    with open(raw_file, 'wb') as f:
        f.write(header.encode('ascii'))
        f.write(data)


class SimulationError(Exception):
    """Raised when a SPICE simulation fails."""

    def __init__(self, message: str, log_content: str = None):
        super().__init__(message)
        self.log_content = log_content


class SpiceRunner(ABC):
    """Abstract base class for SPICE simulator runners.

    This interface is used by pyTesla to run simulations with different
    SPICE backends (LTspice, ngspice embedded, ngspice subprocess).
    All implementations produce .raw and .log files that can be parsed
    with PyLTSpice's RawRead or pyngspice's RawRead.
    """

    def __init__(self, working_directory: str):
        """Initialize with working directory for netlist and output files.

        Args:
            working_directory: Directory where output files will be written.
                Created if it doesn't exist.
        """
        self.working_directory = os.path.abspath(working_directory)
        os.makedirs(self.working_directory, exist_ok=True)

    @abstractmethod
    def run(self, netlist_path: str, timeout: int = None) -> Tuple[str, str]:
        """Execute SPICE simulation.

        Args:
            netlist_path: Absolute or relative path to .net/.cir file
            timeout: Optional timeout in seconds

        Returns:
            (raw_file_path, log_file_path) - absolute paths to output files

        Raises:
            SimulationError: If simulation fails or times out
            FileNotFoundError: If netlist file doesn't exist
        """
        ...

    @staticmethod
    @abstractmethod
    def detect() -> bool:
        """Return True if this SPICE engine is available."""
        ...

    @staticmethod
    @abstractmethod
    def get_executable_path() -> Optional[str]:
        """Auto-detect and return path to executable, or None."""
        ...

    def _preprocess_netlist(self, netlist_path: str) -> str:
        """Pre-process netlist for ngspice compatibility.

        If the netlist needs translation (Rser= on inductors, etc.),
        writes a processed copy to the working directory and returns
        its path. Otherwise returns the original path.

        Args:
            netlist_path: Path to the original netlist

        Returns:
            Path to the netlist to actually simulate (may be original or processed copy)
        """
        netlist_path = os.path.abspath(netlist_path)
        if not os.path.isfile(netlist_path):
            raise FileNotFoundError(f"Netlist not found: {netlist_path}")

        with open(netlist_path, 'r') as f:
            original = f.read()

        processed = preprocess_netlist(original)
        if processed == original:
            return netlist_path

        # Write processed netlist to working directory
        stem = Path(netlist_path).stem
        out_path = os.path.join(self.working_directory, f"{stem}_ngspice.net")
        with open(out_path, 'w') as f:
            f.write(processed)
        return out_path


class NgspiceRunner(SpiceRunner):
    """NgspiceRunner using embedded ngspice via pybind11 C++ extension.

    This is the primary runner. It uses the statically-linked ngspice
    library through pybind11 bindings — no external ngspice installation
    is needed.

    The embedded approach is faster than subprocess invocation since there's
    no process spawn overhead and the simulator is already loaded in memory.
    """

    def __init__(self, working_directory: str = "."):
        super().__init__(working_directory)
        from . import _cpp_available
        if not _cpp_available:
            from . import _cpp_error
            raise ImportError(
                f"pyngspice C++ extension not available: {_cpp_error}\n"
                f"Use SubprocessRunner as a fallback, or rebuild with: pip install -e ."
            )
        from ._pyngspice import SimRunner as _CppSimRunner
        self._runner = _CppSimRunner(output_folder=self.working_directory)

    def run(self, netlist_path: str, timeout: int = None) -> Tuple[str, str]:
        """Run simulation using embedded ngspice.

        Args:
            netlist_path: Path to .net/.cir netlist file
            timeout: Optional timeout in seconds (not yet implemented for embedded mode)

        Returns:
            (raw_file_path, log_file_path) as absolute paths

        Raises:
            SimulationError: If simulation fails
        """
        processed_path = self._preprocess_netlist(netlist_path)

        try:
            raw_file, log_file = self._runner.run_now(processed_path)
        except Exception as e:
            raise SimulationError(f"Simulation failed: {e}") from e

        # Normalize to absolute paths
        raw_file = os.path.abspath(raw_file)
        log_file = os.path.abspath(log_file)

        if not os.path.isfile(raw_file):
            raise SimulationError(
                f"Simulation completed but raw file not found: {raw_file}"
            )

        # Post-process: rename capacitor probe traces in raw file header
        _postprocess_raw_file(raw_file, processed_path)

        return raw_file, log_file

    @staticmethod
    def detect() -> bool:
        """Return True if the embedded C++ extension is available."""
        try:
            from . import _cpp_available
            return _cpp_available
        except ImportError:
            return False

    @staticmethod
    def get_executable_path() -> Optional[str]:
        """Return None — embedded mode has no separate executable."""
        return None


class SubprocessRunner(SpiceRunner):
    """Fallback runner using ngspice as a subprocess.

    Invokes ngspice in batch mode via the command line. Requires ngspice
    to be installed and available on PATH (or specified explicitly).

    This is useful when:
    - The C++ extension is not compiled
    - You need to use a specific ngspice build
    - Debugging simulation issues with ngspice's own output
    """

    def __init__(self, working_directory: str = ".", executable: str = None):
        """Initialize with optional explicit path to ngspice executable.

        Args:
            working_directory: Directory for output files
            executable: Path to ngspice executable. If None, auto-detects.
        """
        super().__init__(working_directory)
        self._executable = executable or self._find_ngspice()
        if self._executable is None:
            raise FileNotFoundError(
                "ngspice executable not found. Install ngspice or specify "
                "the path explicitly via the 'executable' parameter."
            )

    def run(self, netlist_path: str, timeout: int = None) -> Tuple[str, str]:
        """Run simulation using ngspice subprocess.

        Invokes: ngspice -b -r <output.raw> -o <output.log> <netlist>

        Args:
            netlist_path: Path to .net/.cir netlist file
            timeout: Optional timeout in seconds

        Returns:
            (raw_file_path, log_file_path) as absolute paths

        Raises:
            SimulationError: If simulation fails or times out
        """
        processed_path = self._preprocess_netlist(netlist_path)

        stem = Path(processed_path).stem
        raw_file = os.path.join(self.working_directory, f"{stem}.raw")
        log_file = os.path.join(self.working_directory, f"{stem}.log")

        cmd = [
            self._executable,
            '-b',                    # Batch mode (no GUI)
            '-r', raw_file,          # Raw output file
            '-o', log_file,          # Log output file
            processed_path,          # Input netlist
        ]

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd=self.working_directory,
            )
        except subprocess.TimeoutExpired:
            raise SimulationError(
                f"Simulation timed out after {timeout} seconds"
            )
        except FileNotFoundError:
            raise SimulationError(
                f"ngspice executable not found: {self._executable}"
            )

        # Check for errors
        if result.returncode != 0:
            log_content = None
            if os.path.isfile(log_file):
                with open(log_file, 'r') as f:
                    log_content = f.read()
            raise SimulationError(
                f"ngspice exited with code {result.returncode}: {result.stderr}",
                log_content=log_content,
            )

        if not os.path.isfile(raw_file):
            raise SimulationError(
                f"Simulation completed but raw file not found: {raw_file}\n"
                f"stderr: {result.stderr}"
            )

        return os.path.abspath(raw_file), os.path.abspath(log_file)

    @staticmethod
    def detect() -> bool:
        """Return True if ngspice executable is found on PATH."""
        return SubprocessRunner._find_ngspice() is not None

    @staticmethod
    def get_executable_path() -> Optional[str]:
        """Return path to ngspice executable, or None if not found."""
        return SubprocessRunner._find_ngspice()

    @staticmethod
    def _find_ngspice() -> Optional[str]:
        """Search for ngspice executable.

        Checks:
        1. PATH (via shutil.which)
        2. Common install locations on each platform
        """
        # Check PATH first
        found = shutil.which("ngspice")
        if found:
            return found

        # Check common install locations
        if sys.platform == 'win32':
            candidates = [
                r"C:\Program Files\ngspice\bin\ngspice.exe",
                r"C:\Program Files (x86)\ngspice\bin\ngspice.exe",
                os.path.expanduser(r"~\ngspice\bin\ngspice.exe"),
            ]
        elif sys.platform == 'darwin':
            candidates = [
                "/usr/local/bin/ngspice",
                "/opt/homebrew/bin/ngspice",
            ]
        else:  # Linux
            candidates = [
                "/usr/bin/ngspice",
                "/usr/local/bin/ngspice",
                "/snap/bin/ngspice",
            ]

        for path in candidates:
            if os.path.isfile(path):
                return path

        return None


def get_runner(working_directory: str = ".",
               backend: str = "auto") -> SpiceRunner:
    """Factory function to get the best available SpiceRunner.

    Args:
        working_directory: Directory for simulation output files
        backend: One of "auto", "embedded", "subprocess".
            - "auto": try embedded first, fall back to subprocess
            - "embedded": use NgspiceRunner (requires C++ extension)
            - "subprocess": use SubprocessRunner (requires ngspice on PATH)

    Returns:
        A SpiceRunner instance

    Raises:
        RuntimeError: If no suitable backend is found
    """
    if backend == "embedded":
        return NgspiceRunner(working_directory)
    elif backend == "subprocess":
        return SubprocessRunner(working_directory)
    elif backend == "auto":
        # Try embedded first (faster, no external dependency)
        if NgspiceRunner.detect():
            return NgspiceRunner(working_directory)
        # Fall back to subprocess
        if SubprocessRunner.detect():
            return SubprocessRunner(working_directory)
        raise RuntimeError(
            "No ngspice backend available. Either:\n"
            "  1. Build the C++ extension: pip install -e .\n"
            "  2. Install ngspice and add it to PATH"
        )
    else:
        raise ValueError(f"Unknown backend: {backend!r}. Use 'auto', 'embedded', or 'subprocess'.")