#!/usr/bin/env python3
"""S-Parameter Analysis Demonstration.

This demo showcases Oscura's S-parameter analysis capabilities for
characterizing interconnects, cables, and transmission line components.

**Features Demonstrated**:
- Touchstone file parsing (.s1p, .s2p, .snp)
- Return loss (S11) calculation
- Insertion loss (S21) calculation
- S-to-ABCD parameter conversion
- Network cascading
- Frequency-domain interpolation

**S-Parameters**:
S-parameters (Scattering parameters) characterize the electrical
behavior of linear networks at high frequencies:
- S11: Input reflection coefficient (return loss)
- S21: Forward transmission coefficient (insertion loss)
- S12: Reverse transmission coefficient (isolation)
- S22: Output reflection coefficient

**Key Formulas**:
- Return Loss (dB) = -20 * log10(|S11|)
- Insertion Loss (dB) = -20 * log10(|S21|)
- VSWR = (1 + |S11|) / (1 - |S11|)

**Standards**:
- Touchstone 2.0 File Format Specification
- IEEE 370-2020: Electrical Characterization of Interconnects

Usage:
    python sparams_demo.py
    python sparams_demo.py --verbose

Author: Oscura Development Team
Date: 2026-01-16
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np

# Add parent for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from demonstrations.common import BaseDemo, ValidationSuite, print_info, print_result
from demonstrations.common.base_demo import run_demo_main
from demonstrations.common.formatting import GREEN, RED, RESET, YELLOW, print_subheader

# Oscura imports
from oscura.analyzers.signal_integrity.sparams import (
    SParameterData,
    abcd_to_s,
    insertion_loss,
    return_loss,
    s_to_abcd,
)


class SParametersDemo(BaseDemo):
    """S-Parameter Analysis Demonstration.

    This demo creates synthetic S-parameter data representing typical
    interconnect characteristics and demonstrates analysis capabilities.
    """

    def __init__(self, **kwargs):
        """Initialize demo."""
        super().__init__(
            name="S-Parameters Demo",
            description="Demonstrates S-parameter analysis for interconnect characterization",
            **kwargs,
        )

        # Frequency range
        self.f_start = 100e6  # 100 MHz
        self.f_stop = 10e9  # 10 GHz
        self.n_points = 201

        # Cable parameters
        self.cable_length_m = 1.0
        self.characteristic_z = 50.0
        self.attenuation_db_per_m_ghz = 0.5  # dB/m at 1 GHz

        self.s_params = None

    def _generate_cable_s_params(self) -> SParameterData:
        """Generate S-parameters for a lossy transmission line cable.

        Uses simple transmission line model with frequency-dependent loss.

        Returns:
            SParameterData for 2-port cable model.
        """
        frequencies = np.linspace(self.f_start, self.f_stop, self.n_points)

        # S-matrix: n_freq x n_ports x n_ports
        s_matrix = np.zeros((self.n_points, 2, 2), dtype=np.complex128)

        for i, freq in enumerate(frequencies):
            # Frequency-dependent loss (sqrt(f) dependence is typical for skin effect)
            freq_ghz = freq / 1e9
            loss_db = self.attenuation_db_per_m_ghz * self.cable_length_m * np.sqrt(freq_ghz)
            loss_linear = 10 ** (-loss_db / 20)

            # Phase shift through cable
            wavelength = 3e8 / freq / 1.5  # Assume velocity factor of 0.67
            phase_rad = 2 * np.pi * self.cable_length_m / wavelength

            # Transmission coefficient (S21 = S12)
            s21 = loss_linear * np.exp(-1j * phase_rad)

            # Small reflection due to impedance mismatch (assume 1% mismatch)
            gamma = 0.01 + 0.005 * np.sin(2 * phase_rad)  # Standing wave effect
            s11 = gamma * np.exp(-1j * 0.1)  # Small phase

            # Symmetric 2-port
            s_matrix[i, 0, 0] = s11  # S11
            s_matrix[i, 0, 1] = s21  # S12
            s_matrix[i, 1, 0] = s21  # S21
            s_matrix[i, 1, 1] = s11  # S22

        return SParameterData(
            frequencies=frequencies,
            s_matrix=s_matrix,
            n_ports=2,
            z0=self.characteristic_z,
            format="ri",
            source_file="synthetic_cable.s2p",
            comments=["Synthetic 1m lossy cable model", "Generated by Oscura demo"],
        )

    def generate_test_data(self) -> dict:
        """Generate synthetic S-parameter data.

        Loads from file if available (--data-file override or default NPZ),
        otherwise generates synthetic cable S-parameters.
        """
        # Try loading data from file
        file_to_load = None

        # 1. Check CLI override
        if self.data_file and self.data_file.exists():
            file_to_load = self.data_file
            print_info(f"Loading S-parameter data from CLI override: {self.data_file}")
        # 2. Check default generated data
        elif default_file := self.find_default_data_file("sparams.npz"):
            file_to_load = default_file
            print_info(f"Loading S-parameter data from default file: {default_file.name}")

        # Load from file if found
        if file_to_load:
            try:
                data = np.load(file_to_load)
                frequencies = data["frequencies"]
                s_matrix = data["s_matrix"]
                z0 = float(data["z0"])
                n_ports = int(data["n_ports"])

                # Load cable parameters if available
                if "cable_length_m" in data:
                    self.cable_length_m = float(data["cable_length_m"])
                if "characteristic_z" in data:
                    self.characteristic_z = float(data["characteristic_z"])
                if "attenuation_db_per_m_ghz" in data:
                    self.attenuation_db_per_m_ghz = float(data["attenuation_db_per_m_ghz"])

                # Update frequency range from loaded data
                self.f_start = frequencies[0]
                self.f_stop = frequencies[-1]
                self.n_points = len(frequencies)

                # Create S-parameter data object
                self.s_params = SParameterData(
                    frequencies=frequencies,
                    s_matrix=s_matrix,
                    n_ports=n_ports,
                    z0=z0,
                    format="ri",
                    source_file=str(file_to_load),
                    comments=["Loaded from NPZ file"],
                )

                print_result("Data loaded from file", file_to_load.name)
                print_result("Frequency points", self.n_points)
                print_result("Number of ports", self.s_params.n_ports)
                print_result("Reference impedance", f"{self.s_params.z0} Ohms")
                print_result(
                    "Frequency range",
                    f"{self.f_start / 1e6:.0f} MHz - {self.f_stop / 1e9:.0f} GHz",
                )
                return
            except Exception as e:
                print_info(f"Failed to load data from file: {e}, falling back to synthetic")
                file_to_load = None

        # Generate synthetic data if not loaded
        print_info("Generating synthetic cable S-parameters...")

        print_info(f"  Cable length: {self.cable_length_m} m")
        print_info(f"  Characteristic impedance: {self.characteristic_z} Ohms")
        print_info(f"  Attenuation: {self.attenuation_db_per_m_ghz} dB/m/sqrt(GHz)")
        print_info(f"  Frequency range: {self.f_start / 1e6:.0f} MHz - {self.f_stop / 1e9:.0f} GHz")

        self.s_params = self._generate_cable_s_params()

        print_result("Frequency points", self.n_points)
        print_result("Number of ports", self.s_params.n_ports)
        print_result("Reference impedance", f"{self.s_params.z0} Ohms")

        return {}

    def run_demonstration(self, data: dict) -> dict:
        """Analyze S-parameters and calculate key metrics."""
        # Store flags for validation
        self.results["s11_present"] = self.s_params is not None
        self.results["s21_present"] = self.s_params is not None and self.s_params.n_ports >= 2

        # ===== Return Loss Analysis =====
        print_subheader("Return Loss Analysis (S11)")

        rl_all = return_loss(self.s_params)

        print_result("Min return loss", f"{np.min(rl_all):.1f} dB")
        print_result("Max return loss", f"{np.max(rl_all):.1f} dB")
        print_result("Mean return loss", f"{np.mean(rl_all):.1f} dB")

        self.results["min_return_loss_db"] = np.min(rl_all)
        self.results["max_return_loss_db"] = np.max(rl_all)

        # Sample at specific frequencies
        test_freqs = [1e9, 5e9, 10e9]
        print_info("Return loss at key frequencies:")
        for freq in test_freqs:
            rl = return_loss(self.s_params, frequency=freq)
            print_info(f"  {freq / 1e9:.0f} GHz: {rl:.1f} dB")

        # Evaluate return loss quality
        min_rl = np.min(rl_all)
        if min_rl > 20:
            rating = f"{GREEN}Excellent (RL > 20 dB){RESET}"
        elif min_rl > 15:
            rating = f"{GREEN}Good (RL > 15 dB){RESET}"
        elif min_rl > 10:
            rating = f"{YELLOW}Marginal (RL > 10 dB){RESET}"
        else:
            rating = f"{RED}Poor (RL < 10 dB){RESET}"
        print_info(f"Return loss quality: {rating}")

        # ===== Insertion Loss Analysis =====
        print_subheader("Insertion Loss Analysis (S21)")

        il_all = insertion_loss(self.s_params)

        print_result("IL at 100 MHz", f"{il_all[0]:.2f} dB")
        print_result("IL at 10 GHz", f"{il_all[-1]:.2f} dB")
        print_result("Total IL variation", f"{np.ptp(il_all):.2f} dB")

        self.results["il_100mhz_db"] = il_all[0]
        self.results["il_10ghz_db"] = il_all[-1]
        self.results["il_variation_db"] = np.ptp(il_all)

        # Sample at specific frequencies
        print_info("Insertion loss at key frequencies:")
        for freq in test_freqs:
            il = insertion_loss(self.s_params, frequency=freq)
            print_info(f"  {freq / 1e9:.0f} GHz: {il:.2f} dB")

        # Evaluate insertion loss
        il_max = np.max(il_all)
        if il_max < 3:
            rating = f"{GREEN}Excellent (IL < 3 dB){RESET}"
        elif il_max < 6:
            rating = f"{GREEN}Good (IL < 6 dB){RESET}"
        elif il_max < 10:
            rating = f"{YELLOW}Marginal (IL < 10 dB){RESET}"
        else:
            rating = f"{RED}Poor (IL >= 10 dB){RESET}"
        print_info(f"Insertion loss quality: {rating}")

        # ===== VSWR Calculation =====
        print_subheader("VSWR Analysis")

        s11_all = np.abs(self.s_params.get_s(1, 1))
        vswr = (1 + s11_all) / (1 - s11_all)

        print_result("Min VSWR", f"{np.min(vswr):.2f}:1")
        print_result("Max VSWR", f"{np.max(vswr):.2f}:1")
        print_result("Mean VSWR", f"{np.mean(vswr):.2f}:1")

        self.results["max_vswr"] = np.max(vswr)

        # VSWR at key frequencies
        print_info("VSWR at key frequencies:")
        for _i, freq in enumerate([1e9, 5e9, 10e9]):
            idx = int((freq - self.f_start) / (self.f_stop - self.f_start) * (self.n_points - 1))
            idx = max(0, min(idx, self.n_points - 1))
            print_info(f"  {freq / 1e9:.0f} GHz: {vswr[idx]:.2f}:1")

        # ===== ABCD Parameter Conversion =====
        print_subheader("ABCD Parameter Analysis")

        abcd = s_to_abcd(self.s_params)

        # Show ABCD at 1 GHz
        idx_1ghz = int((1e9 - self.f_start) / (self.f_stop - self.f_start) * (self.n_points - 1))
        abcd_1ghz = abcd[idx_1ghz]

        print_info("ABCD matrix at 1 GHz:")
        print_info(f"  A = {abcd_1ghz[0, 0]:.4f}")
        print_info(f"  B = {abcd_1ghz[0, 1]:.4f}")
        print_info(f"  C = {abcd_1ghz[1, 0]:.6f}")
        print_info(f"  D = {abcd_1ghz[1, 1]:.4f}")

        # Verify round-trip conversion
        s_roundtrip = abcd_to_s(abcd, z0=self.s_params.z0)
        error = np.max(np.abs(s_roundtrip - self.s_params.s_matrix))
        print_result("ABCD round-trip error", f"{error:.2e}")

        self.results["abcd_roundtrip_error"] = error

        if error < 1e-10:
            print_info(f"  {GREEN}Conversion verified (error < 1e-10){RESET}")
        else:
            print_info(f"  {YELLOW}Some conversion error detected{RESET}")

        # ===== Network Cascading Example =====
        print_subheader("Network Cascading")

        print_info("Cascading two identical cables:")

        # ABCD cascade: multiply matrices
        abcd_cascade = np.zeros_like(abcd)
        for i in range(len(abcd)):
            abcd_cascade[i] = abcd[i] @ abcd[i]

        # Convert back to S-parameters
        s_cascade = abcd_to_s(abcd_cascade, z0=self.s_params.z0)

        # Create cascaded S-parameter data
        s_params_cascade = SParameterData(
            frequencies=self.s_params.frequencies,
            s_matrix=s_cascade,
            n_ports=2,
            z0=self.s_params.z0,
        )

        # Compare insertion loss
        il_single = insertion_loss(self.s_params)
        il_cascade = insertion_loss(s_params_cascade)

        print_result("Single cable IL at 10 GHz", f"{il_single[-1]:.2f} dB")
        print_result("Cascaded IL at 10 GHz", f"{il_cascade[-1]:.2f} dB")
        print_result("IL ratio", f"{il_cascade[-1] / il_single[-1]:.2f}x")

        # Should be approximately 2x (6 dB more)
        il_ratio = il_cascade[-1] / max(il_single[-1], 0.001)
        self.results["cascade_il_ratio"] = il_ratio
        self.results["cascade_loss_db"] = il_cascade[-1]

        if 1.8 < il_ratio < 2.2:
            print_info(f"  {GREEN}Cascade IL is ~2x single (expected){RESET}")
        else:
            print_info(f"  {YELLOW}Cascade ratio: {il_ratio:.2f}x{RESET}")

        # ===== Frequency Response Summary =====
        print_subheader("Summary")

        print_info("Parameter        100 MHz    1 GHz      10 GHz")
        print_info("-" * 50)

        # Get values at key frequencies
        freqs = [0.1e9, 1e9, 10e9]
        rl_vals = [return_loss(self.s_params, f) for f in freqs]
        il_vals = [insertion_loss(self.s_params, f) for f in freqs]

        print_info(
            f"Return Loss (dB)  {rl_vals[0]:>6.1f}    {rl_vals[1]:>6.1f}    {rl_vals[2]:>6.1f}"
        )
        print_info(
            f"Insertion Loss    {il_vals[0]:>6.2f}    {il_vals[1]:>6.2f}    {il_vals[2]:>6.2f}"
        )

        # Store for validation
        self.results["return_loss_1ghz"] = rl_vals[1]
        self.results["insertion_loss_1ghz"] = il_vals[1]

        return self.results

    def validate(self, results: dict) -> bool:
        """Validate S-parameter demo results."""
        suite = ValidationSuite()

        # Check S-parameters were generated
        s_params_present = results.get("s11_present", False) and results.get("s21_present", False)
        suite.add_check(
            "S-parameters generated", s_params_present, "S11 and S21 parameters present"
        )

        # Check return loss is reasonable
        min_return_loss = results.get("min_return_loss_db", 0)
        suite.add_check(
            "Minimum return loss",
            min_return_loss >= 5,
            f"Got {min_return_loss:.2f} dB (≥5 dB expected)",
        )

        # Check insertion loss is reasonable
        il_10ghz = results.get("il_10ghz_db", 100)
        suite.add_check("IL at 10 GHz", il_10ghz < 20, f"Got {il_10ghz:.2f} dB (<20 dB expected)")

        # Check VSWR is reasonable
        max_vswr = results.get("max_vswr", 100)
        suite.add_check("Maximum VSWR", max_vswr < 2.0, f"Got {max_vswr:.2f} (<2:1 expected)")

        # Check ABCD conversion accuracy
        abcd_error = results.get("abcd_roundtrip_error", 1)
        suite.add_check(
            "ABCD round-trip error", abcd_error < 1e-6, f"Got {abcd_error:.2e} (<1e-6 expected)"
        )

        # Check cascade behavior (ratio should be ~2.0x)
        cascade_ratio = results.get("cascade_il_ratio", 0)
        suite.add_check(
            "Cascade loss ratio",
            1.8 < cascade_ratio < 2.2,
            f"Got {cascade_ratio:.2f}x (should be ~2.0x for identical cascaded networks)",
        )

        suite.report()
        return suite.all_passed()


if __name__ == "__main__":
    sys.exit(run_demo_main(SParametersDemo))
