"""Parse Touchstone (.s1p, .s2p, .snp) S-parameter files for LTspice."""

import re
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

# Frequency unit multipliers to Hz
_FREQ_MULTIPLIERS: dict[str, float] = {
    "HZ": 1.0,
    "KHZ": 1e3,
    "MHZ": 1e6,
    "GHZ": 1e9,
}


@dataclass
class TouchstoneData:
    """Parsed contents of a Touchstone file.

    All frequencies are stored in Hz regardless of the original file's
    unit.  The ``data`` array holds complex-valued parameters with shape
    ``(n_freq, n_ports, n_ports)``.
    """

    filename: str
    n_ports: int
    freq_unit: str
    parameter_type: str  # S, Y, Z, H, G
    format_type: str  # MA, DB, RI
    reference_impedance: float
    frequencies: np.ndarray  # shape (n_freq,), always in Hz
    data: np.ndarray  # shape (n_freq, n_ports, n_ports), complex128
    comments: list[str] = field(default_factory=list)


def _to_complex(v1: float, v2: float, fmt: str) -> complex:
    """Convert a value pair to complex according to the format type."""
    if fmt == "RI":
        return complex(v1, v2)
    elif fmt == "MA":
        mag = v1
        angle_rad = np.deg2rad(v2)
        return complex(mag * np.cos(angle_rad), mag * np.sin(angle_rad))
    elif fmt == "DB":
        mag = 10.0 ** (v1 / 20.0)
        angle_rad = np.deg2rad(v2)
        return complex(mag * np.cos(angle_rad), mag * np.sin(angle_rad))
    else:
        raise ValueError(f"Unknown format type: {fmt!r}")


def _detect_ports(path: Path) -> int:
    """Detect port count from the file extension (.s<n>p)."""
    suffix = path.suffix.lower()
    m = re.match(r"\.s(\d+)p$", suffix)
    if not m:
        raise ValueError(
            f"Cannot determine port count from extension {suffix!r}. "
            "Expected .s1p, .s2p, .s3p, .s4p, etc."
        )
    return int(m.group(1))


def parse_touchstone(path: str | Path) -> TouchstoneData:
    """Parse a Touchstone file into a TouchstoneData object.

    Handles .s1p through .s4p (and beyond), all three format types
    (MA, DB, RI), all frequency units, and continuation lines used
    by files with more than two ports.

    Args:
        path: Path to the Touchstone file.

    Returns:
        A TouchstoneData with frequencies converted to Hz and
        parameters stored as complex128.
    """
    path = Path(path)
    n_ports = _detect_ports(path)

    # Defaults per Touchstone spec
    freq_unit = "GHZ"
    param_type = "S"
    fmt = "MA"
    ref_impedance = 50.0

    comments: list[str] = []
    data_lines: list[str] = []
    option_found = False

    with path.open() as fh:
        for raw_line in fh:
            line = raw_line.strip()

            # Comment lines
            if line.startswith("!"):
                comments.append(line[1:].strip())
                continue

            # Option line (only the first one is used)
            if line.startswith("#"):
                if not option_found:
                    option_found = True
                    tokens = line[1:].split()
                    # Parse tokens case-insensitively
                    i = 0
                    while i < len(tokens):
                        tok = tokens[i].upper()
                        if tok in _FREQ_MULTIPLIERS:
                            freq_unit = tok
                        elif tok in ("S", "Y", "Z", "H", "G"):
                            param_type = tok
                        elif tok in ("MA", "DB", "RI"):
                            fmt = tok
                        elif tok == "R" and i + 1 < len(tokens):
                            i += 1
                            ref_impedance = float(tokens[i])
                        i += 1
                continue

            # Skip blank lines
            if not line:
                continue

            # Inline comments after data (some files use !)
            if "!" in line:
                line = line[: line.index("!")]

            data_lines.append(line)

    # Compute the expected number of value pairs per frequency point.
    # Each frequency point has n_ports * n_ports parameters, each
    # consisting of two floats.
    values_per_freq = n_ports * n_ports * 2  # pairs * 2 values each

    # Flatten all data tokens
    all_tokens: list[float] = []
    for dl in data_lines:
        all_tokens.extend(float(t) for t in dl.split())

    # Each frequency row starts with the frequency value, followed by
    # values_per_freq data values.  For n_ports <= 2, everything fits
    # on one line.  For n_ports > 2, continuation lines are used and
    # don't repeat the frequency.
    stride = 1 + values_per_freq  # freq + data
    if len(all_tokens) % stride != 0:
        raise ValueError(
            f"Token count {len(all_tokens)} is not a multiple of "
            f"expected stride {stride} for a {n_ports}-port file."
        )

    n_freq = len(all_tokens) // stride
    freq_mult = _FREQ_MULTIPLIERS[freq_unit]

    frequencies = np.empty(n_freq, dtype=np.float64)
    data = np.empty((n_freq, n_ports, n_ports), dtype=np.complex128)

    for k in range(n_freq):
        offset = k * stride
        frequencies[k] = all_tokens[offset] * freq_mult

        # Data values come in pairs: (v1, v2) per parameter
        idx = offset + 1
        for row in range(n_ports):
            for col in range(n_ports):
                v1 = all_tokens[idx]
                v2 = all_tokens[idx + 1]
                data[k, row, col] = _to_complex(v1, v2, fmt)
                idx += 2

    return TouchstoneData(
        filename=path.name,
        n_ports=n_ports,
        freq_unit=freq_unit,
        parameter_type=param_type,
        format_type=fmt,
        reference_impedance=ref_impedance,
        frequencies=frequencies,
        data=data,
        comments=comments,
    )


def get_s_parameter(data: TouchstoneData, i: int, j: int) -> tuple[np.ndarray, np.ndarray]:
    """Extract a single S-parameter across all frequencies.

    Args:
        data: Parsed Touchstone data.
        i: Row index (1-based, as in S(i,j)).
        j: Column index (1-based).

    Returns:
        Tuple of (frequencies_hz, complex_values) where both are 1-D
        numpy arrays.
    """
    if i < 1 or i > data.n_ports:
        raise IndexError(f"Row index {i} out of range for {data.n_ports}-port data")
    if j < 1 or j > data.n_ports:
        raise IndexError(f"Column index {j} out of range for {data.n_ports}-port data")
    return data.frequencies.copy(), data.data[:, i - 1, j - 1].copy()


def s_param_to_db(complex_values: np.ndarray) -> np.ndarray:
    """Convert complex S-parameter values to decibels.

    Computes 20 * log10(|S|), flooring magnitudes at -300 dB to avoid
    log-of-zero warnings.

    Args:
        complex_values: Array of complex S-parameter values.

    Returns:
        Magnitude in dB as a real-valued numpy array.
    """
    magnitude = np.abs(complex_values)
    return 20.0 * np.log10(np.maximum(magnitude, 1e-15))


def generate_ltspice_subcircuit(touchstone_data: TouchstoneData, name: str) -> str:
    """Generate an LTspice-compatible subcircuit wrapping S-parameter data.

    LTspice can reference Touchstone files from within a subcircuit
    using the ``.net`` directive.  This function produces a ``.sub``
    file body that instantiates the S-parameter block.

    Args:
        touchstone_data: Parsed Touchstone data.
        name: Subcircuit name (used in .SUBCKT and filename references).

    Returns:
        A string containing the complete subcircuit definition.
    """
    td = touchstone_data
    n = td.n_ports

    # Build port list: port1, port2, ..., portN, plus a reference node
    port_names = [f"port{k}" for k in range(1, n + 1)]
    port_list = " ".join(port_names)

    lines: list[str] = []
    lines.append(f"* LTspice subcircuit for {td.filename}")
    lines.append(f"* {n}-port {td.parameter_type}-parameters, Z0={td.reference_impedance} Ohm")
    lines.append(
        f"* Frequency range: {td.frequencies[0]:.6g} Hz "
        f"to {td.frequencies[-1]:.6g} Hz "
        f"({len(td.frequencies)} points)"
    )
    lines.append(f".SUBCKT {name} {port_list} ref")
    lines.append(f".net {td.filename} {port_list} ref")
    lines.append(f".ends {name}")

    return "\n".join(lines) + "\n"