import numpy as np
from scipy.signal import savgol_filter
import matplotlib.pyplot as plt

def min_max_delta_to_array(min_, max_, delta_):
    return np.arange(min_, max_ + delta_ / 2, delta_)

def volt_p_to_dBm(vp, Z0=50):
    """
    Convert peak voltage `vp` to dBm, assuming impedance `Z0`
    """
    P = vp ** 2 / (2 * Z0)
    return watt_to_dBm(P)

def volt_pp_to_dBm(vpp, Z0=50):
    """
    Convert peak-to-peak voltage `vpp` to dBm, assuming impedance `Z0`
    """
    return volt_p_to_dBm(vpp / 2)

def watt_to_dBm(p):
    """
    Convert power in units of Watt into dBm
    """
    return 10 * np.log10(p / 1e-3)

def dBm_to_watt(p_dBm):
    """
    Convert power in units of dBm into Watt
    """
    return 1e-3 * 10 ** (p_dBm / 10)

def dBm_to_volt_p(p_dBm, Z0=50):
    """
    Convert power in units of dBm into peak voltage
    """
    return np.sqrt(2 * Z0 * dBm_to_watt(p_dBm))

def dBm_to_volt_pp(p_dBm, Z0=50):
    """
    Convert power in units of dBm into peak-to-peak voltage
    """
    return 2 * dBm_to_volt_p(p_dBm, Z0=50)

def pow_ratio_to_dB(r):
    """
    Convert power ratio `r` into dB number
    """
    return 10 * np.log10(r)

def dB_to_amp_ratio(x):
    return 10 ** (x / 20)

def savgol_filter_complex(sig, *args, **kwargs):
    r = savgol_filter(sig.real, *args, **kwargs)
    i = savgol_filter(sig.imag, *args, **kwargs)
    return r + 1j * i

def simple_resonance_idx_finder(sig):
    """
    Simple
    """
    q_idx = 0

    if sig.dtype == complex:
        abs_sig_filt = np.abs(savgol_filter_complex(sig, 11, 3))
        q_idx = 0
        avg_, max_, min_ = np.mean(abs_sig_filt), np.max(abs_sig_filt), np.min(abs_sig_filt)
        if (max_ - avg_) > (avg_ - min_): # check if maximum signal has larger separation from average than the minimum signal
            q_idx = np.argmax(abs_sig_filt)
        else:
            q_idx = np.argmin(abs_sig_filt)
    elif sig.dtype == float:
        sig_filt = savgol_filter(sig, 11, 3)
        avg_, max_, min_ = np.mean(sig_filt), np.max(sig_filt), np.min(sig_filt)
        if (max_ - avg_) > (avg_ - min_): # check if maximum signal has larger separation from average than the minimum signal
            q_idx = np.argmax(sig_filt)
        else:
            q_idx = np.argmin(sig_filt)
    return q_idx

def plot_fft(times, sig):
    
    dt = times[1] - times[0]
    N = len(times)
    
    freqs = np.fft.fftshift(np.fft.fftfreq(N, d=dt))
    sig_f = np.fft.fftshift(np.fft.fft(sig) * np.sqrt(dt / N))
    
    fig = plt.figure()
    plt.plot(freqs / 1e6, np.abs(sig_f) ** 2)
    plt.xlabel("Frequency (MHz)")
    plt.ylabel("Power Spectral Density ($\mathrm{V}^2/\mathrm{Hz}$)")
    
    return fig