# Copyright 2026 Gaofeng Fan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""SaveConfig for managing simulation output saves."""

from dataclasses import dataclass, field
from typing import List, Set, Optional
from pathlib import Path
import yaml


@dataclass
class SaveConfig:
    """
    Defines what signals to save for a circuit/subcircuit.

    SaveConfig uses relative paths within the circuit it's defined for.
    When instantiated in a hierarchy, use with_prefix() to apply the
    hierarchical path.

    Example:
        # Define saves for an OTA
        ota_saves = (SaveConfig("ota_debug")
            .voltage("out", "tail", "bias")
            .current("M1:d", "M2:d")
            .op("M1:gm", "M1:vth"))

        # In testbench, apply prefix
        tb.save(ota_saves.with_prefix("X_LDO.X_OTA"))
    """
    name: Optional[str] = None
    _voltage: List[str] = field(default_factory=list)
    _current: List[str] = field(default_factory=list)
    _op: List[str] = field(default_factory=list)
    _tags: dict = field(default_factory=dict)

    def voltage(self, *nets, tag: Optional[str] = None) -> 'SaveConfig':
        """Save voltage at specified nets."""
        for net in nets:
            self._voltage.append(net)
            if tag:
                self._tags.setdefault(tag, []).append(('v', net))
        return self

    def current(self, *terminals, tag: Optional[str] = None) -> 'SaveConfig':
        """Save current through terminals (e.g., 'M1:d')."""
        for term in terminals:
            self._current.append(term)
            if tag:
                self._tags.setdefault(tag, []).append(('i', term))
        return self

    def op(self, *params, tag: Optional[str] = None) -> 'SaveConfig':
        """Save operating point parameters (e.g., 'M1:gm', 'M1:vth')."""
        for param in params:
            self._op.append(param)
            if tag:
                self._tags.setdefault(tag, []).append(('op', param))
        return self

    def with_prefix(self, prefix: str) -> 'SaveConfig':
        """
        Return new SaveConfig with hierarchical prefix applied.

        Args:
            prefix: Hierarchical path like "X_LDO.X_OTA"

        Returns:
            New SaveConfig with prefixed paths
        """
        prefixed = SaveConfig(self.name)
        prefixed._voltage = [f"{prefix}.{s}" for s in self._voltage]
        prefixed._current = [f"{prefix}.{c}" for c in self._current]
        prefixed._op = [f"{prefix}.{p}" for p in self._op]
        # Prefix tags too
        for tag, items in self._tags.items():
            prefixed._tags[tag] = [
                (sig_type, f"{prefix}.{sig}") for sig_type, sig in items
            ]
        return prefixed

    def merge(self, other: 'SaveConfig') -> 'SaveConfig':
        """Combine with another SaveConfig."""
        self._voltage.extend(other._voltage)
        self._current.extend(other._current)
        self._op.extend(other._op)
        for tag, items in other._tags.items():
            self._tags.setdefault(tag, []).extend(items)
        return self

    def only_tags(self, *tags) -> 'SaveConfig':
        """Return new SaveConfig with only specified tags."""
        filtered = SaveConfig(self.name)
        for tag in tags:
            for sig_type, sig in self._tags.get(tag, []):
                if sig_type == 'v':
                    filtered._voltage.append(sig)
                elif sig_type == 'i':
                    filtered._current.append(sig)
                elif sig_type == 'op':
                    filtered._op.append(sig)
        return filtered

    def without_tags(self, *tags) -> 'SaveConfig':
        """Return new SaveConfig excluding specified tags."""
        exclude = set()
        for tag in tags:
            for _, sig in self._tags.get(tag, []):
                exclude.add(sig)

        filtered = SaveConfig(self.name)
        filtered._voltage = [v for v in self._voltage if v not in exclude]
        filtered._current = [c for c in self._current if c not in exclude]
        filtered._op = [o for o in self._op if o not in exclude]
        return filtered

    def to_spectre(self) -> str:
        """Generate Spectre save statements."""
        lines = []
        if self._voltage:
            lines.append(f"save {' '.join(self._voltage)}")
        if self._current:
            lines.append(f"save {' '.join(self._current)}")
        if self._op:
            lines.append(f"save {' '.join(self._op)}")
        return '\n'.join(lines)

    def to_ngspice(self) -> str:
        """Generate ngspice save statements."""
        lines = []
        for sig in self._voltage:
            lines.append(f".save v({sig})")
        for cur in self._current:
            lines.append(f".save i({cur})")
        return '\n'.join(lines)

    @classmethod
    def from_yaml(cls, path: str) -> 'SaveConfig':
        """Load SaveConfig from YAML file."""
        with open(path) as f:
            data = yaml.safe_load(f)
        config = cls(data.get('name'))
        config._voltage = data.get('voltage', [])
        config._current = data.get('current', [])
        config._op = data.get('op', [])
        return config

    def to_yaml(self, path: str):
        """Save SaveConfig to YAML file."""
        data = {
            'name': self.name,
            'voltage': self._voltage,
            'current': self._current,
            'op': self._op,
        }
        with open(path, 'w') as f:
            yaml.dump(data, f, default_flow_style=False)

    @property
    def signals(self) -> List[str]:
        """All voltage signals."""
        return self._voltage.copy()

    @property
    def currents(self) -> List[str]:
        """All current signals."""
        return self._current.copy()

    @property
    def op_params(self) -> List[str]:
        """All operating point parameters."""
        return self._op.copy()

    def __len__(self) -> int:
        """Total number of saved signals."""
        return len(self._voltage) + len(self._current) + len(self._op)


class SaveOptions:
    """
    Testbench-level save control.

    Provides modes for controlling what gets saved:
    - "all": Save all public nodes (saveOptions save=allpub)
    - "selected": Only save explicitly specified signals
    - "none": Disable all saves

    Example:
        tb.save_options.mode = "selected"
        tb.save_options.nestlvl = 2
    """
    def __init__(self):
        self.mode: str = "selected"  # "all", "selected", "none"
        self.nestlvl: Optional[int] = None
        self.pwr: str = "none"  # "all", "none"

    def to_spectre(self) -> str:
        parts = ["saveOptions options"]
        if self.mode == "all":
            parts.append("save=allpub")
        elif self.mode == "selected":
            parts.append("save=selected")

        if self.nestlvl is not None:
            parts.append(f"nestlvl={self.nestlvl}")

        parts.append(f"pwr={self.pwr}")

        return " ".join(parts)
