"""
Rectangular Waveguide Design Class

High-level interface for rectangular waveguide S-parameter analysis.
"""

import subprocess
import tempfile
import os
from typing import Optional, List, Tuple
import numpy as np

import pyedgefem as em


class RectWaveguideDesign:
    """
    Rectangular waveguide S-parameter analysis.

    This class provides a high-level interface for computing S-parameters
    of rectangular waveguide sections. It handles mesh generation, port
    setup, and simulation automatically.

    Standard waveguide designations:
        WR-90 (X-band):   a=22.86mm, b=10.16mm, fc=6.56 GHz
        WR-62 (Ku-band):  a=15.80mm, b=7.90mm,  fc=9.49 GHz
        WR-42 (Ka-band):  a=10.67mm, b=4.32mm,  fc=14.1 GHz

    Example:
        wg = RectWaveguideDesign(a=22.86e-3, b=10.16e-3, length=50e-3)
        wg.generate_mesh(density=10)
        S = wg.sparams_at_freq(10e9)
        print(f"|S21| = {abs(S[1,0]):.4f}")

    Attributes:
        a (float): Waveguide width in meters
        b (float): Waveguide height in meters
        length (float): Waveguide length in meters
        mesh: The FEM mesh (after generate_mesh is called)
        ports: Wave port definitions
    """

    def __init__(
        self,
        a: float,
        b: float,
        length: float,
        *,
        pec_tag: int = 1,
        port1_tag: int = 2,
        port2_tag: int = 3,
        volume_tag: int = 100,
    ):
        """
        Initialize rectangular waveguide design.

        Args:
            a: Waveguide width (broadwall dimension) in meters
            b: Waveguide height (narrowwall dimension) in meters
            length: Waveguide length in meters
            pec_tag: Physical tag for PEC walls (default: 1)
            port1_tag: Physical tag for port 1 surface (default: 2)
            port2_tag: Physical tag for port 2 surface (default: 3)
            volume_tag: Physical tag for waveguide volume (default: 100)
        """
        self.a = a
        self.b = b
        self.length = length

        self._pec_tag = pec_tag
        self._port1_tag = port1_tag
        self._port2_tag = port2_tag
        self._volume_tag = volume_tag

        self._mesh = None
        self._mesh_path = None
        self._bc = None
        self._ports = None
        self._last_freq = None
        self._last_S = None

        # Compute cutoff frequency
        c0 = 299792458.0
        self._fc = c0 / (2 * a)  # TE10 cutoff

    @property
    def mesh(self):
        """The FEM mesh (None until generate_mesh is called)."""
        return self._mesh

    @property
    def ports(self):
        """Wave port definitions (None until setup_ports is called)."""
        return self._ports

    @property
    def cutoff_frequency(self) -> float:
        """TE10 cutoff frequency in Hz."""
        return self._fc

    def _check_propagating(self, freq: float):
        """Check if frequency is above cutoff."""
        if freq <= self._fc:
            raise ValueError(
                f"Frequency {freq/1e9:.3f} GHz is below TE10 cutoff "
                f"{self._fc/1e9:.3f} GHz"
            )

    def generate_mesh(
        self,
        density: float = 10.0,
        output_path: Optional[str] = None,
        gmsh_options: Optional[List[str]] = None,
    ) -> None:
        """
        Generate FEM mesh using Gmsh.

        Args:
            density: Elements per wavelength at 1.5× cutoff frequency (default: 10)
            output_path: Path for mesh file. If None, uses temp directory.
            gmsh_options: Additional Gmsh command-line options

        Raises:
            RuntimeError: If Gmsh is not available or mesh generation fails
        """
        # Compute mesh size based on wavelength at 1.5× fc
        c0 = 299792458.0
        freq_design = 1.5 * self._fc
        wavelength = c0 / freq_design
        lc = wavelength / density

        # Generate .geo file content
        geo_content = f"""// Rectangular waveguide mesh
// Generated by EdgeFEM RectWaveguideDesign
// a = {self.a*1000:.3f} mm, b = {self.b*1000:.3f} mm, length = {self.length*1000:.3f} mm

lc = {lc};

// Define points
Point(1) = {{0, 0, 0, lc}};
Point(2) = {{{self.a}, 0, 0, lc}};
Point(3) = {{{self.a}, {self.b}, 0, lc}};
Point(4) = {{0, {self.b}, 0, lc}};
Point(5) = {{0, 0, {self.length}, lc}};
Point(6) = {{{self.a}, 0, {self.length}, lc}};
Point(7) = {{{self.a}, {self.b}, {self.length}, lc}};
Point(8) = {{0, {self.b}, {self.length}, lc}};

// Define lines
Line(1) = {{1, 2}};
Line(2) = {{2, 3}};
Line(3) = {{3, 4}};
Line(4) = {{4, 1}};
Line(5) = {{5, 6}};
Line(6) = {{6, 7}};
Line(7) = {{7, 8}};
Line(8) = {{8, 5}};
Line(9) = {{1, 5}};
Line(10) = {{2, 6}};
Line(11) = {{3, 7}};
Line(12) = {{4, 8}};

// Define surfaces
Curve Loop(1) = {{1, 2, 3, 4}};   // Port 1 (z=0)
Plane Surface(1) = {{1}};

Curve Loop(2) = {{5, 6, 7, 8}};   // Port 2 (z=length)
Plane Surface(2) = {{2}};

Curve Loop(3) = {{1, 10, -5, -9}};  // Bottom wall
Plane Surface(3) = {{3}};

Curve Loop(4) = {{2, 11, -6, -10}}; // Right wall
Plane Surface(4) = {{4}};

Curve Loop(5) = {{3, 12, -7, -11}}; // Top wall
Plane Surface(5) = {{5}};

Curve Loop(6) = {{4, 9, -8, -12}};  // Left wall
Plane Surface(6) = {{6}};

// Define surface loop and volume
Surface Loop(1) = {{1, 2, 3, 4, 5, 6}};
Volume(1) = {{1}};

// Physical groups
Physical Surface("{self._pec_tag}") = {{3, 4, 5, 6}};  // PEC walls
Physical Surface("{self._port1_tag}") = {{1}};         // Port 1
Physical Surface("{self._port2_tag}") = {{2}};         // Port 2
Physical Volume("{self._volume_tag}") = {{1}};         // Waveguide volume

// Mesh settings
Mesh.Algorithm = 6;  // Frontal-Delaunay
Mesh.Algorithm3D = 1;  // Delaunay
"""

        # Write .geo file and run Gmsh
        if output_path is None:
            fd, geo_path = tempfile.mkstemp(suffix=".geo")
            os.close(fd)
            msh_path = geo_path.replace(".geo", ".msh")
        else:
            if output_path.endswith(".msh"):
                msh_path = output_path
                geo_path = output_path.replace(".msh", ".geo")
            else:
                geo_path = output_path + ".geo"
                msh_path = output_path + ".msh"

        with open(geo_path, "w") as f:
            f.write(geo_content)

        # Run Gmsh
        cmd = ["gmsh", geo_path, "-3", "-format", "msh2", "-o", msh_path]
        if gmsh_options:
            cmd.extend(gmsh_options)

        try:
            result = subprocess.run(
                cmd, capture_output=True, text=True, check=True
            )
        except FileNotFoundError:
            raise RuntimeError(
                "Gmsh not found. Please install Gmsh and ensure it's in PATH."
            )
        except subprocess.CalledProcessError as e:
            raise RuntimeError(f"Gmsh failed: {e.stderr}")

        # Load mesh
        self._mesh = em.load_gmsh(msh_path)
        self._mesh_path = msh_path
        self._bc = None
        self._ports = None

        # Validate mesh
        report = em.validate_mesh(self._mesh)
        if not report.is_valid():
            print(
                f"Warning: Mesh has {report.num_degenerate_tets} degenerate "
                f"and {report.num_inverted_tets} inverted elements"
            )

    def load_mesh(self, mesh_path: str) -> None:
        """
        Load a pre-existing mesh file.

        Args:
            mesh_path: Path to Gmsh v2 .msh file
        """
        self._mesh = em.load_gmsh(mesh_path)
        self._mesh_path = mesh_path
        self._bc = None
        self._ports = None

    def _setup_ports(self, freq: float) -> None:
        """Set up wave ports for the given frequency."""
        if self._mesh is None:
            raise RuntimeError("Mesh not generated. Call generate_mesh() first.")

        # Build BC
        self._bc = em.build_edge_pec(self._mesh, self._pec_tag)

        # Create port dimensions
        port_dim = em.RectWaveguidePort()
        port_dim.a = self.a
        port_dim.b = self.b

        # Compute TE10 mode
        mode = em.solve_te10_mode(port_dim, freq)

        # Build port 1
        surface1 = em.extract_surface_mesh(self._mesh, self._port1_tag)
        em.populate_te10_field(surface1, port_dim, mode)
        port1 = em.build_wave_port(self._mesh, surface1, mode)

        # Build port 2
        surface2 = em.extract_surface_mesh(self._mesh, self._port2_tag)
        em.populate_te10_field(surface2, port_dim, mode)
        port2 = em.build_wave_port(self._mesh, surface2, mode)

        self._ports = [port1, port2]

    def sparams_at_freq(
        self,
        freq: float,
        *,
        use_eigenmode: bool = True,
        port_abc_scale: float = 0.5,
    ) -> np.ndarray:
        """
        Compute S-parameters at a single frequency.

        Args:
            freq: Frequency in Hz
            use_eigenmode: Use eigenmode-based S-parameter extraction (default: True)
            port_abc_scale: ABC scaling factor (default: 0.5, optimal for accuracy)

        Returns:
            2x2 complex S-parameter matrix as numpy array

        Raises:
            ValueError: If frequency is below cutoff
            RuntimeError: If mesh not generated
        """
        self._check_propagating(freq)
        self._setup_ports(freq)

        # Set up Maxwell parameters
        params = em.MaxwellParams()
        params.omega = 2 * np.pi * freq
        params.port_abc_scale = port_abc_scale

        # Compute S-parameters
        if use_eigenmode:
            S = em.calculate_sparams_eigenmode(
                self._mesh, params, self._bc, self._ports
            )
        else:
            S = em.calculate_sparams(self._mesh, params, self._bc, self._ports)

        self._last_freq = freq
        self._last_S = S

        # Check passivity
        em.check_passivity(S)

        return np.array(S)

    def frequency_sweep(
        self,
        f_start: float,
        f_stop: float,
        n_points: int = 11,
        *,
        use_eigenmode: bool = True,
        port_abc_scale: float = 0.5,
        verbose: bool = False,
    ) -> Tuple[np.ndarray, List[np.ndarray]]:
        """
        Compute S-parameters over a frequency range.

        Args:
            f_start: Start frequency in Hz
            f_stop: Stop frequency in Hz
            n_points: Number of frequency points (default: 11)
            use_eigenmode: Use eigenmode-based extraction (default: True)
            port_abc_scale: ABC scaling factor (default: 0.5)
            verbose: Print progress (default: False)

        Returns:
            Tuple of (frequencies, S_matrices) where:
                frequencies: 1D array of frequencies in Hz
                S_matrices: List of 2x2 complex S-parameter matrices

        Raises:
            ValueError: If frequency range includes below-cutoff frequencies
        """
        freqs = np.linspace(f_start, f_stop, n_points)
        S_list = []

        for i, f in enumerate(freqs):
            if verbose:
                print(f"  {i+1}/{n_points}: {f/1e9:.3f} GHz")
            S = self.sparams_at_freq(
                f, use_eigenmode=use_eigenmode, port_abc_scale=port_abc_scale
            )
            S_list.append(S)

        return freqs, S_list

    def export_touchstone(
        self,
        path: str,
        freqs: Optional[np.ndarray] = None,
        S_matrices: Optional[List[np.ndarray]] = None,
    ) -> None:
        """
        Export S-parameters to Touchstone format.

        Args:
            path: Output file path (will add .s2p extension if needed)
            freqs: Frequencies in Hz. If None, uses last simulation frequency.
            S_matrices: List of S-matrices. If None, uses last simulation result.

        Raises:
            RuntimeError: If no simulation has been run and no data provided
        """
        if freqs is None:
            if self._last_freq is None:
                raise RuntimeError("No simulation data. Run sparams_at_freq() first.")
            freqs = np.array([self._last_freq])
            S_matrices = [self._last_S]

        if not path.endswith(".s2p"):
            path = path + ".s2p"

        # Convert to Eigen matrices for the API
        S_eigen = [np.array(S, dtype=complex) for S in S_matrices]
        em.write_touchstone_nport(path, list(freqs), S_eigen)

    def export_fields_vtk(
        self,
        path: str,
        freq: Optional[float] = None,
        *,
        active_port: int = 0,
    ) -> None:
        """
        Export E-field solution to VTK format for ParaView visualization.

        Args:
            path: Output file path (should end in .vtu)
            freq: Frequency in Hz. If None, uses last simulation frequency.
            active_port: Port to excite (0 or 1, default: 0)

        Raises:
            RuntimeError: If mesh not generated
        """
        if self._mesh is None:
            raise RuntimeError("Mesh not generated. Call generate_mesh() first.")

        if freq is None:
            if self._last_freq is None:
                raise ValueError("No frequency specified and no prior simulation.")
            freq = self._last_freq
        else:
            self._setup_ports(freq)

        # Set up parameters
        params = em.MaxwellParams()
        params.omega = 2 * np.pi * freq
        params.port_abc_scale = 0.5

        # Assemble and solve
        assembly = em.assemble_maxwell(
            self._mesh, params, self._bc, self._ports, active_port
        )
        result = em.solve_linear(assembly.A, assembly.b)

        # Export
        em.export_fields_vtk(self._mesh, result.x, path)

    def __repr__(self) -> str:
        return (
            f"RectWaveguideDesign(a={self.a*1000:.2f}mm, "
            f"b={self.b*1000:.2f}mm, length={self.length*1000:.2f}mm, "
            f"fc={self._fc/1e9:.3f}GHz)"
        )
