"""
analysis.py
-----------
Pure numerical analysis functions for extracting resonance frequencies
and linewidths from S-parameter networks.

None of these functions build networks or do optimization — they only
inspect existing rf.Network objects.  This makes them easy to test
independently and reuse across models.
"""
from __future__ import annotations
import logging

import numpy as np
import skrf as rf
from scipy.signal import find_peaks

logger = logging.getLogger(__name__)


def resonance(ntwk: rf.Network, m: int = 0, n: int = 0, use_max: bool = False) -> float:
    """
    Estimate resonance frequency from the real part of S[m,n].

    By default returns the frequency of the minimum of Re(S), which
    corresponds to the resonance dip in S11/S22.  Set use_max=True
    for transmission peaks.

    Parameters
    ----------
    ntwk : rf.Network
    m, n : int
        S-parameter indices (0-based).
    use_max : bool
        If True, find frequency of maximum Re(S) instead of minimum.

    Returns
    -------
    float
        Resonance frequency in Hz.
    """
    f    = ntwk.frequency.f
    re_s = np.real(ntwk.s[:, m, n])
    idx  = np.argmax(re_s) if use_max else np.argmin(re_s)
    return float(f[idx])


def resonance_from_s_max(network: rf.Network, m: int = 0, n: int = 0) -> float:
    """
    Resonance frequency from the dominant peak in |S[m,n]|.

    Finds all peaks in the magnitude and returns the frequency of the
    largest one.  Useful for transmission parameters (S12/S21) where the
    resonance appears as a peak rather than a dip.

    Parameters
    ----------
    network : rf.Network
    m, n : int
        S-parameter indices (0-based).

    Returns
    -------
    float
        Dominant resonance frequency in GHz.

    Raises
    ------
    ValueError
        If no peaks are found in |S[m,n]|.
    """
    f     = network.frequency.f / 1e9
    s_mn  = network.s[:, m, n]
    peaks, _ = find_peaks(np.abs(s_mn))

    if len(peaks) == 0:
        raise ValueError("No peaks found in |S[{},{}]|.".format(m, n))

    peak_idx = peaks[np.argmax(np.abs(s_mn)[peaks])]
    return float(f[peak_idx])

def resonances_from_s_max(network: rf.Network, m: int = 0, n: int = 0) -> float:
    """
    Resonance frequency from the dominant peak in |S[m,n]|.

    Finds all peaks in the magnitude and returns the frequency of the
    largest one.  Useful for transmission parameters (S12/S21) where the
    resonance appears as a peak rather than a dip.

    Parameters
    ----------
    network : rf.Network
    m, n : int
        S-parameter indices (0-based).

    Returns
    -------
    float
        Dominant resonance frequency in GHz.

    Raises
    ------
    ValueError
        If no peaks are found in |S[m,n]|.
    """
    f     = network.frequency.f / 1e9
    s_mn  = network.s[:, m, n]
    peaks, _ = find_peaks(np.abs(s_mn))
    freqs = f[peaks]
    logger.debug("Resonance frequencies (GHz): %s", freqs)

    if len(peaks) == 0:
        raise ValueError("No peaks found in |S[{},{}]|.".format(m, n))
    return freqs


def resonances_from_s(network: rf.Network, m: int = 0, n: int = 0) -> np.ndarray:
    """
    Find all resonance frequencies from peaks in Re(S[m,n]).

    Parameters
    ----------
    network : rf.Network
    m, n : int
        S-parameter indices (0-based).

    Returns
    -------
    np.ndarray
        Array of resonance frequencies in GHz.
    """
    f    = network.frequency.f / 1e9
    s_mn = network.s[:, m, n]
    peaks, _ = find_peaks(np.real(s_mn))
    freqs = f[peaks]
    logger.debug("Resonance frequencies (GHz): %s", freqs)
    return freqs


def fwhm_from_trace_db(
    ntwk: rf.Network,
    m: int = 0,
    n: int = 0,
    kind: str = "dip",
    smooth: int = None,
) -> float:
    """
    Full-Width at Half Maximum (FWHM) from S-parameter magnitude in dB.

    Uses crossing-interpolation to find the two frequencies where the
    magnitude crosses the half-depth level, giving a more accurate result
    than a pure index-based approach.

    Parameters
    ----------
    ntwk : rf.Network
    m, n : int
        S-parameter indices (0-based).
    kind : {"dip", "peak"}
        "dip"  — resonance appears as a downward dip (typical for S11/S22).
        "peak" — resonance appears as an upward peak (typical for S12/S21).
    smooth : int or None
        Optional moving-average window length for noisy traces.
        Use an odd integer; None disables smoothing.

    Returns
    -------
    float
        FWHM linewidth in Hz.

    Raises
    ------
    ValueError
        If two crossings at the half-depth level cannot be found.
    """
    f      = np.asarray(ntwk.frequency.f, dtype=float)
    s      = ntwk.s[:, m, n]
    mag_db = 20 * np.log10(np.maximum(np.abs(s), 1e-300))

    if smooth is not None and smooth > 1:
        kernel = np.ones(int(smooth)) / int(smooth)
        mag_db = np.convolve(mag_db, kernel, mode="same")

    if kind == "dip":
        i0       = np.argmin(mag_db)
        y0       = mag_db[i0]
        baseline = np.median(np.r_[
            mag_db[:max(3, i0 // 8)],
            mag_db[min(len(mag_db) - 3, i0 + len(mag_db) // 8):],
        ])
        level = 0.5 * (baseline + y0)
        above = mag_db > level
    elif kind == "peak":
        i0       = np.argmax(mag_db)
        y0       = mag_db[i0]
        baseline = np.median(np.r_[
            mag_db[:max(3, i0 // 8)],
            mag_db[min(len(mag_db) - 3, i0 + len(mag_db) // 8):],
        ])
        level = 0.5 * (baseline + y0)
        above = mag_db < level
    else:
        raise ValueError("kind must be 'dip' or 'peak', got '{}'.".format(kind))

    flips = np.where(above[:-1] != above[1:])[0]
    if flips.size < 2:
        raise ValueError("Could not find two crossings for FWHM.")

    left_flips  = flips[flips < i0]
    right_flips = flips[flips >= i0]
    if left_flips.size == 0 or right_flips.size == 0:
        raise ValueError("Could not bracket resonance with crossings.")

    iL = left_flips[-1]
    iR = right_flips[0]

    def _interp_crossing(i: int) -> float:
        x0, x1 = f[i], f[i + 1]
        y0_, y1_ = mag_db[i], mag_db[i + 1]
        if y1_ == y0_:
            return 0.5 * (x0 + x1)
        return x0 + (level - y0_) * (x1 - x0) / (y1_ - y0_)

    return float(_interp_crossing(iR) - _interp_crossing(iL))


def fwhm_from_res11(ntwk: rf.Network) -> float:
    """
    Linewidth from zero crossings of Re(S11).

    Re(S11) crosses zero at the two half-power frequencies of the resonance.
    Uses linear interpolation between samples for accuracy.

    Parameters
    ----------
    ntwk : rf.Network

    Returns
    -------
    float
        Linewidth (distance between the two zero crossings) in Hz.

    Raises
    ------
    ValueError
        If fewer than two zero crossings are found.
    """
    f     = ntwk.frequency.f
    re_s11 = np.real(ntwk.s[:, 0, 0])

    sign_changes = np.where(np.diff(np.sign(re_s11)))[0]
    if len(sign_changes) < 2:
        raise ValueError("Could not determine linewidth: fewer than 2 zero crossings found.")

    zeros = []
    for idx in sign_changes[:2]:
        f1, f2 = f[idx], f[idx + 1]
        y1, y2 = re_s11[idx], re_s11[idx + 1]
        f_zero = f1 - y1 * (f2 - f1) / (y2 - y1)
        zeros.append(f_zero)

    return float(zeros[1] - zeros[0])


def resonance_from_res11(ntwk: rf.Network) -> float:
    """
    Resonance frequency from the minimum of Re(S11).

    Thin convenience wrapper around `resonance` for the common single-port
    S11 case.

    Parameters
    ----------
    ntwk : rf.Network

    Returns
    -------
    float
        Resonance frequency in Hz.
    """
    return resonance(ntwk, m=0, n=0, use_max=False)


def find_resonant_frequency(network: rf.Network) -> complex:
    freqs = network.f
    S21 = network.s[:, 1, 0]

    mag = np.abs(S21)
    peak_idx = np.argmax(mag)
    omega_0 = 2 * np.pi * freqs[peak_idx]

    half_pow = mag[peak_idx] / np.sqrt(2)
    left  = np.where(mag[:peak_idx] < half_pow)[0]
    right = np.where(mag[peak_idx:] < half_pow)[0]
    alpha = (2 * np.pi * freqs[peak_idx + right[0]] - 2 * np.pi * freqs[left[-1]]) / 2

    pole = complex(-alpha, omega_0)
    f0 = pole.imag / (2 * np.pi)
    return f0


def find_resonant_frequency(network: rf.Network) -> complex:
    freqs = network.f
    S21 = network.s[:, 1, 0]

    mag = np.abs(S21)
    peak_idx = np.argmax(mag)
    omega_0 = 2 * np.pi * freqs[peak_idx]

    half_pow = mag[peak_idx] / np.sqrt(2)
    left  = np.where(mag[:peak_idx] < half_pow)[0]
    right = np.where(mag[peak_idx:] < half_pow)[0]
    alpha = (2 * np.pi * freqs[peak_idx + right[0]] - 2 * np.pi * freqs[left[-1]]) / 2

    pole = complex(-alpha, omega_0)
    f0 = pole.imag / (2 * np.pi)
    return f0
