# EMerge is an open source Python based FEM EM simulation module.
# Copyright (C) 2025  Andrés Martínez Mera.
# Copyright (C) 2025  Robert Fennis.

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program; if not, see
# <https://www.gnu.org/licenses/>.

import os
from typing import Literal
import numpy as np
from datetime import datetime
from loguru import logger

_F_UNIT = {'HZ': 1,
           'KHZ': 1e3,
           'MHZ': 1e6,
           'GHZ': 1e9,
}
def generate_touchstone(filename: str,
                        freq: np.ndarray,
                        Smat: np.ndarray,
                        data_format: Literal['RI','MA','DB'],
                        custom_comments: list[str] = None,
                        funit: Literal['HZ','KHZ','MHZ','GHZ'] = 'GHZ') -> None:
    """
    Export S-parameter data to a Touchstone file

    Parameters
    ----------
    filename : str
        Base filename (may include path). If no extension is given, one
        will be added of the form '.sNp' where N = number of ports.
    freq : np.ndarray
        1-D array of M frequency points, in Hz.
    Smat : np.ndarray
        3-D array of S-parameters with shape (M, N, N).
    data_format : {'RI','MA','DB'}
        RI = real&imag, MA = magnitude&angle (deg), DB = dB&angle (deg).
    custom_comments : list[str], optional
        List of custom comment strings to add to the touchstone file header.
        Each string will be prefixed with "! " automatically.
    """
    # --- validations ---
    if Smat.ndim != 3:
        raise ValueError(f"Smat must be 3-D with shape (M,N,N), got ndim={Smat.ndim}")
    M, Nports, N2 = Smat.shape
    if Nports != N2:
        raise ValueError(f"Smat must be square in its last two dims, got {Nports}×{N2}")
    if freq.ndim != 1 or freq.size != M:
        raise ValueError(f"freq must be 1-D of length {M}, got shape {freq.shape}")

    # --- determine output filename & extension ---
    base, ext = os.path.splitext(filename)
    if ext == '':
        ext = f".s{Nports}p"
    filename_out = base + ext

    # --- write the Touchstone file ---
    with open(filename_out, 'w') as f:
        # Write header
        f.write(f"! Generated by EMerge - {Nports}-port S-parameters\n")
        timestamp_str = datetime.now().strftime("%Y-%m-%d %H:%M %Z")
        f.write(f"! File generated: {timestamp_str}\n")
        
        # Write custom comments if provided
        if custom_comments:
            for comment in custom_comments:
                # Ensure comment starts with "! " if not already present
                if not comment.startswith("! "):
                    comment = "! " + comment
                f.write(f"{comment}\n")

        f.write(f"# {funit} S {data_format.upper()} R 50\n")

        # Write data
        for i in range(M):
            freq_ghz = freq[i] / _F_UNIT[funit]  # Convert Hz to GHz
            
            # Build all S-parameter values for this frequency
            s_values = []
            for row in range(Nports):
                for col in range(Nports):
                    s_val = Smat[i, row, col]
                    
                    if data_format == 'RI':
                        # Real and Imaginary
                        val1 = s_val.real
                        val2 = s_val.imag
                    elif data_format == 'MA':
                        # Magnitude and Angle (degrees)
                        val1 = np.abs(s_val)
                        val2 = np.angle(s_val) * 180.0 / np.pi
                    elif data_format == 'DB':
                        # dB and Angle (degrees)
                        val1 = 20.0 * np.log10(np.abs(s_val)) if np.abs(s_val) > 0 else -300.0
                        val2 = np.angle(s_val) * 180.0 / np.pi
                    
                    s_values.extend([val1, val2])
            
            # Write frequency and S-parameters according to touchstone format
            if Nports == 1:
                # s1p: freq S11_real S11_imag (3 values per line)
                f.write(f"{freq_ghz:12.6e}  {s_values[0]:12.6e}  {s_values[1]:12.6e}\n")
            
            elif Nports == 2:
                # s2p: freq S11_r S11_i S21_r S21_i S12_r S12_i S22_r S22_i (9 values per line)
                f.write(f"{freq_ghz:12.6e}")
                for val in s_values:
                    f.write(f"  {val:12.6e}")
                f.write("\n")
            
            else:
                # s3p, s4p, etc: freq on first line, then S-parameters on subsequent lines
                f.write(f"{freq_ghz:12.6e}")
                
                if (Nports == 3):
                    values_per_line = 6  # 3-port S-parameter
                else: 
                    values_per_line = 8  # 4-port S-parameter pairs
                for j in range(0, len(s_values), values_per_line):
                    f.write("   ")  # Indent continuation lines
                    if (j != 0):
                        f.write(' ' * 12)  # Extra indentation
                    end_idx = min(j + values_per_line, len(s_values))
                    for val in s_values[j:end_idx]:
                        f.write(f" {val:12.6e}")
                    f.write("\n")
    
    logger.info(f"Touchstone file written to '{filename_out}'")