# 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.

"""
SimulationBatch for organizing PVT sweeps and Monte Carlo simulations.

This module generates simulation infrastructure but does NOT execute.
It creates:
- Directory structures
- Netlists for each run
- Runner scripts (Python)
- Commands for execution
"""

from typing import List, Dict, Callable, Any, Optional
from pathlib import Path
from dataclasses import dataclass

from .testbench import Testbench
from .spectre import generate_spectre
from .command import SpectreCommand
from .devices import vsource # For docstring example


@dataclass
class SimulationRun:
    """Represents a single simulation run."""
    name: str
    directory: Path
    testbench: Testbench
    netlist: Optional[str] = None
    command: Optional[str] = None


class SimulationBatch:
    """
    Organizes multiple simulations for PVT sweeps or Monte Carlo.

    Does NOT execute simulations - generates scripts and commands.

    Example (PVT sweep):
        batch = SimulationBatch("ldo_pvt", "/sim/ldo_pvt")
        batch.pvt_sweep(make_tb_ldo, corners=[
            {"process": "tt", "voltage": 1.8, "temp": 27},
            {"process": "ff", "voltage": 1.98, "temp": -40},
        ])
        batch.generate()
        batch.write_runner("run_pvt.py")

    Example (Monte Carlo):
        batch = SimulationBatch("ldo_mc", "/sim/ldo_mc")
        batch.monte_carlo(make_tb_ldo, num_runs=500, seed_start=1)
        batch.generate()
        batch.write_runner("run_mc.py")
    """

    def __init__(self, name: str, root_dir: str):
        """
        Initialize batch simulation.

        Args:
            name: Batch name (used in output files)
            root_dir: Root directory for all simulation runs
        """
        self.name = name
        self.root_dir = Path(root_dir)
        self.runs: List[SimulationRun] = []
        self._command_opts: Dict[str, Any] = {}

    def add_run(self, name: str, testbench: Testbench) -> 'SimulationBatch':
        """
        Add a single simulation run.

        Args:
            name: Name for this run (used as subdirectory)
            testbench: Testbench instance
        """
        run_dir = self.root_dir / name
        self.runs.append(SimulationRun(
            name=name,
            directory=run_dir,
            testbench=testbench
        ))
        return self

    def pvt_sweep(self, testbench_factory: Callable[..., Testbench],
                  corners: List[Dict[str, Any]]) -> 'SimulationBatch':
        """
        Add PVT corner sweep.

        Args:
            testbench_factory: Function that creates testbench given corner params
            corners: List of corner dictionaries with process/voltage/temp

        Example:
            def make_tb(process="tt", voltage=1.8, temp=27):
                tb = Testbench(f"tb_{process}")
                gnd = tb.gnd()
                vdd_net = tb.net("vdd")
                tb.set_temp(temp)

                tb.instantiate(vsource, instance_name="I_Vdd", p=vdd_net, n=gnd, dc=voltage)
                # ... build circuit ...
                return tb

            batch.pvt_sweep(make_tb, corners=[
                {"process": "tt", "voltage": 1.8, "temp": 27},
                {"process": "ff", "voltage": 1.98, "temp": -40},
            ])
        """
        for corner in corners:
            # Generate corner name
            p = corner.get('process', 'typ')
            v = corner.get('voltage', 1.8)
            t = corner.get('temp', 27)
            name = f"{p}_v{v}_t{t}"

            # Create testbench with corner parameters
            tb = testbench_factory(**corner)
            self.add_run(name, tb)

        return self

    def monte_carlo(self, testbench_factory: Callable[..., Testbench],
                    num_runs: int, seed_start: int = 1) -> 'SimulationBatch':
        """
        Add Monte Carlo simulation runs.

        Args:
            testbench_factory: Function that creates testbench given seed
            num_runs: Number of Monte Carlo runs
            seed_start: Starting seed value
        """
        for i in range(num_runs):
            seed = seed_start + i
            name = f"mc_{seed:04d}"
            tb = testbench_factory(seed=seed)
            self.add_run(name, tb)

        return self

    def command_options(self, **kwargs) -> 'SimulationBatch':
        """
        Set default command options for all runs.

        Args:
            **kwargs: Options passed to SpectreCommand
                accuracy: "conservative", "moderate", "liberal"
                threads: Number of threads
                timeout: Timeout in seconds
        """
        self._command_opts.update(kwargs)
        return self

    def generate(self, create_dirs: bool = True) -> List[SimulationRun]:
        """
        Generate netlists and commands for all runs.

        Args:
            create_dirs: Whether to create directories (default True)

        Returns:
            List of SimulationRun with netlists and commands populated
        """
        for run in self.runs:
            # Create directory
            if create_dirs:
                run.directory.mkdir(parents=True, exist_ok=True)

            # Generate netlist
            run.netlist = generate_spectre(run.testbench)
            netlist_path = run.directory / "input.scs"

            # Write netlist
            if create_dirs:
                netlist_path.write_text(run.netlist)

            # Build command
            cmd = SpectreCommand(str(netlist_path))

            # Apply default options
            if 'accuracy' in self._command_opts:
                cmd.accuracy(self._command_opts['accuracy'])
            if 'threads' in self._command_opts:
                cmd.threads(self._command_opts['threads'])
            if 'timeout' in self._command_opts:
                cmd.timeout(self._command_opts['timeout'])

            run.command = cmd.build()

        return self.runs

    def write_runner(self, output_file: str = None,
                     config_at_top: bool = True) -> str:
        """
        Write Python runner script.

        Args:
            output_file: Output file path (default: {name}_runner.py)
            config_at_top: Include configuration section at top

        Returns:
            Path to generated runner script
        """
        if output_file is None:
            output_file = f"{self.name}_runner.py"

        lines = [
            "#!/usr/bin/env python3",
            f'"""',
            f'{self.name} Simulation Runner',
            f'Generated by analogpy',
            f'Total runs: {len(self.runs)}',
            f'"""',
            "",
        ]

        if config_at_top:
            lines.extend([
                "# " + "=" * 60,
                "# CONFIGURATION - Modify these as needed",
                "# " + "=" * 60,
                "#",
                "# Common adjustments:",
                "#   --accuracy conservative  # For convergence issues",
                "#   --accuracy liberal       # For faster simulation",
                "#   --threads 32             # More parallelism",
                "#",
                "DEFAULTS = {",
                f'    "accuracy": "{self._command_opts.get("accuracy", "liberal")}",',
                f'    "threads": {self._command_opts.get("threads", 16)},',
                f'    "timeout": {self._command_opts.get("timeout", 900)},',
                "}",
                "# " + "=" * 60,
                "",
            ])

        lines.extend([
            "import argparse",
            "from pathlib import Path",
            "",
            "# Simulation runs",
            "RUNS = [",
        ])

        for run in self.runs:
            lines.append(f'    {{"name": "{run.name}", "dir": "{run.directory}"}},')

        lines.extend([
            "]",
            "",
            "def get_commands(config):",
            '    """Generate spectre commands with current config."""',
            "    commands = []",
            "    for run in RUNS:",
            '        netlist = Path(run["dir"]) / "input.scs"',
            "        cmd_parts = [",
            '            "spectre",',
            "            str(netlist),",
            '            "+escchars",',
            '            f"+log ../psf/spectre.out",',
            '            "-format psfxl",',
            '            "-raw ../psf",',
            '            f\'++aps={config["accuracy"]}\',',
            '            f\'+mt={config["threads"]}\',',
            '            f\'+lqtimeout {config["timeout"]}\',',
            "        ]",
            '        commands.append(" ".join(cmd_parts))',
            "    return commands",
            "",
            "def print_commands(config):",
            '    """Print commands (pipe to tmux-ssh or parallel)."""',
            "    for cmd in get_commands(config):",
            "        print(cmd)",
            "",
            "def main():",
            '    parser = argparse.ArgumentParser(description=f"{__doc__}")',
            '    parser.add_argument("action", choices=["commands", "list", "help"],',
            '                        nargs="?", default="help")',
            '    parser.add_argument("-a", "--accuracy",',
            '                        choices=["conservative", "moderate", "liberal"],',
            '                        default=DEFAULTS["accuracy"])',
            '    parser.add_argument("-t", "--threads", type=int,',
            '                        default=DEFAULTS["threads"])',
            '    parser.add_argument("--timeout", type=int,',
            '                        default=DEFAULTS["timeout"])',
            "    args = parser.parse_args()",
            "",
            '    config = {',
            '        "accuracy": args.accuracy,',
            '        "threads": args.threads,',
            '        "timeout": args.timeout,',
            '    }',
            "",
            '    if args.action == "commands":',
            "        print_commands(config)",
            '    elif args.action == "list":',
            '        print(f"Runs ({len(RUNS)}):")',
            "        for run in RUNS:",
            '            print(f"  {run[\'name\']}: {run[\'dir\']}")',
            "    else:",
            '        print("Usage:")',
            '        print("  python {script} commands    # Print spectre commands")',
            '        print("  python {script} list        # List all runs")',
            '        print("")',
            '        print("Current configuration:")',
            "        for k, v in config.items():",
            '            print(f"  {k}: {v}")',
            '        print("")',
            '        print("Override with: --accuracy, --threads, --timeout")',
            '        print("")',
            '        print("Example:")',
            '        print("  python {script} commands | parallel -j 8 tmux-ssh {{}}")',
            "",
            'if __name__ == "__main__":',
            "    main()",
        ])

        script_content = "\n".join(lines)

        with open(output_file, 'w') as f:
            f.write(script_content)

        return output_file

    def get_commands(self) -> List[str]:
        """Get list of all commands."""
        return [run.command for run in self.runs if run.command]

    def __len__(self) -> int:
        return len(self.runs)
