# Copyright (C) 1995-2021, Christopher Torrence and Gilbert P.Compo
# Python version of the code is written by Evgeniya Predybaylo in 2014
# edited by Michael von Papen (FZ Juelich, INM-6), 2018, to include
# analysis at arbitrary frequencies
#
#   This software may be used, copied, or redistributed as long as it is not
#   sold and this copyright notice is reproduced on each copy made. This
#   routine is provided as is without any express or implied warranties
#   whatsoever.
#
# Notice: Please acknowledge the use of the above software in any publications:
#            Wavelet software was provided by C. Torrence and G. Compo,
#      and is available at URL: http://paos.colorado.edu/research/wavelets/''.
#
# Reference: Torrence, C. and G. P. Compo, 1998: A Practical Guide to
#            Wavelet Analysis. <I>Bull. Amer. Meteor. Soc.</I>, 79, 61-78.
#
# Please send a copy of such publications to either C. Torrence or G. Compo:
#  Dr. Christopher Torrence               Dr. Gilbert P. Compo
#  Research Systems, Inc.                 Climate Diagnostics Center
#  4990 Pearl East Circle                 325 Broadway R/CDC1
#  Boulder, CO 80301, USA                 Boulder, CO 80305-3328, USA
#  E-mail: chris[AT]rsinc[DOT]com         E-mail: compo[AT]colorado[DOT]edu
#
# ----------------------------------------------------------------------------


# # WAVELET  1D Wavelet transform with optional significance testing
#   wave, period, scale, coi = wavelet(Y, dt, pad, dj, s0, J1, mother, param)
#
#   Computes the wavelet transform of the vector Y (length N),
#   with sampling rate DT.
#
#   By default, the Morlet wavelet (k0=6) is used.
#   The wavelet basis is normalized to have total energy=1 at all scales.
#
# INPUTS:
#
#    Y = the time series of length N.
#    DT = amount of time between each Y value, i.e. the sampling time.
#
# OUTPUTS:
#
#    WAVE is the WAVELET transform of Y. This is a complex array
#    of dimensions (N,J1+1). FLOAT(WAVE) gives the WAVELET amplitude,
#    ATAN(IMAGINARY(WAVE),FLOAT(WAVE) gives the WAVELET phase.
#    The WAVELET power spectrum is ABS(WAVE)**2.
#    Its units are sigma**2 (the time series variance).
#
# OPTIONAL INPUTS:
#
# *** Note *** if none of the optional variables is set up, then the program
#   uses default values of -1.
#
#    PAD = if set to 1 (default is 0), pad time series with zeroes to get
#         N up to the next higher power of 2. This prevents wraparound
#         from the end of the time series to the beginning, and also
#         speeds up the FFT's used to do the wavelet transform.
#         This will not eliminate all edge effects (see COI below).
#
#    DJ = the spacing between discrete scales. Default is 0.25.
#         A smaller # will give better scale resolution, but be slower to plot.
#
#    S0 = the smallest scale of the wavelet.  Default is 2*DT.
#
#    J1 = the # of scales minus one. Scales range from S0 up to S0*2**(J1*DJ),
#        to give a total of (J1+1) scales. Default is J1 = (LOG2(N DT/S0))/DJ.
#
#    MOTHER = the mother wavelet function.
#             The choices are 'MORLET', 'PAUL', or 'DOG'
#
#    PARAM = the mother wavelet parameter.
#            For 'MORLET' this is k0 (wavenumber), default is 6.
#            For 'PAUL' this is m (order), default is 4.
#            For 'DOG' this is m (m-th derivative), default is 2.
#
#
# OPTIONAL OUTPUTS:
#
#    PERIOD = the vector of "Fourier" periods (in time units) that corresponds
#           to the SCALEs.
#
#    SCALE = the vector of scale indices, given by S0*2**(j*DJ), j=0...J1
#            where J1+1 is the total # of scales.
#
#    COI = if specified, then return the Cone-of-Influence, which is a vector
#        of N points that contains the maximum period of useful information
#        at that particular time.
#        Periods greater than this are subject to edge effects.

import numpy as np

from scipy.optimize import fminbound
from scipy.special._ufuncs import gamma, gammainc
from scipy.signal import find_peaks

__author__ = "Evgeniya Predybaylo, Michael von Papen, Daniel Gonzalez-Duque"


def wavelet(Y, dt, pad=0, dj=-1, s0=-1, J1=-1, mother=-1, param=-1, freq=None):
    n1 = len(Y)

    if s0 == -1:
        s0 = 2 * dt
    if dj == -1:
        dj = 1.0 / 4.0
    if J1 == -1:
        J1 = np.fix((np.log(n1 * dt / s0) / np.log(2)) / dj)
    if mother == -1:
        mother = "MORLET"

    # construct time series to analyze, pad if necessary
    x = Y - np.mean(Y)
    if pad == 1:
        # power of 2 nearest to N
        base2 = np.fix(np.log(n1) / np.log(2) + 0.4999)
        nzeroes = (2 ** (base2 + 1) - n1).astype(np.int64)
        x = np.concatenate((x, np.zeros(nzeroes)))

    n = len(x)

    # construct wavenumber array used in transform [Eqn(5)]
    kplus = np.arange(1, int(n / 2) + 1).astype(np.int64)
    kplus = kplus * 2 * np.pi / (n * dt)
    kminus = np.arange(1, int((n - 1) / 2) + 1).astype(np.int64)
    kminus = np.sort((-kminus * 2 * np.pi / (n * dt)))
    k = np.concatenate(([0.0], kplus, kminus))

    # compute FFT of the (padded) time series
    f = np.fft.fft(x)  # [Eqn(3)]

    # construct SCALE array & empty PERIOD & WAVE arrays
    if mother.upper() == "MORLET":
        empir = [2.0, -1, -1, -1]
        empir[1:] = [0.776, 2.32, 0.60]
        if param == -1:
            param = 6.0
            empir[1:] = [0.776, 2.32, 0.60]
        fourier_factor = 4 * np.pi / (param + np.sqrt(2 + param**2))
    elif mother.upper() == "PAUL":
        empir = [2, -1, -1, -1]
        empir[1:] = [1.132, 1.17, 1.5]
        if param == -1:
            param = 4.0
            empir[1:] = [1.132, 1.17, 1.5]
        fourier_factor = 4 * np.pi / (2 * param + 1)
    elif mother.upper() == "DOG":
        empir = [1.0, -1, -1, -1]
        empir[1:] = [3.541, 1.43, 1.4]
        if param == -1:
            param = 2.0
            empir[1:] = [3.541, 1.43, 1.4]
        elif param == 6:  # --------------------------------------DOG
            empir[1:] = [1.966, 1.37, 0.97]
        fourier_factor = 2 * np.pi * np.sqrt(2.0 / (2 * param + 1))
    else:
        fourier_factor = np.nan

    if freq is None:
        j = np.arange(0, J1 + 1)
        scale = s0 * 2.0 ** (j * dj)
        freq = 1.0 / (fourier_factor * scale)
        period = 1.0 / freq
    else:
        scale = 1.0 / (fourier_factor * freq)
        period = 1.0 / freq
    # define the wavelet array
    wave = np.zeros(shape=(len(scale), n), dtype=complex)

    # loop through all scales and compute transform
    for a1 in range(0, len(scale)):
        daughter, fourier_factor, coi, _ = wave_bases(mother, k, scale[a1], param)
        wave[a1, :] = np.fft.ifft(f * daughter)  # wavelet transform[Eqn(4)]

    # COI [Sec.3g]
    coi = (
        coi
        * dt
        * np.concatenate(
            (
                np.insert(
                    np.arange(int((n1 + 1) / 2) - 1).astype(np.int64), [0], [1e-5]
                ),
                np.insert(
                    np.flipud(np.arange(0, int(n1 / 2) - 1).astype(np.int64)),
                    [-1],
                    [1e-5],
                ),
            )
        )
    )
    wave = wave[:, :n1]  # get rid of padding before returning

    # Save parameters in dictionary
    parameters = {
        "pad": pad,
        "s0": s0,
        "dj": dj,
        "J1": J1,
        "mother": mother,
        "C_delta": empir[1],
    }

    return wave, period, scale, coi, parameters


# --------------------------------------------------------------------------
# WAVE_BASES  1D Wavelet functions Morlet, Paul, or DOG
#
#  DAUGHTER,FOURIER_FACTOR,COI,DOFMIN = wave_bases(MOTHER,K,SCALE,PARAM)
#
#   Computes the wavelet function as a function of Fourier frequency,
#   used for the wavelet transform in Fourier space.
#   (This program is called automatically by WAVELET)
#
# INPUTS:
#
#    MOTHER = a string, equal to 'MORLET' or 'PAUL' or 'DOG'
#    K = a vector, the Fourier frequencies at which to calculate the wavelet
#    SCALE = a number, the wavelet scale
#    PARAM = the nondimensional parameter for the wavelet function
#
# OUTPUTS:
#
#    DAUGHTER = a vector, the wavelet function
#    FOURIER_FACTOR = the ratio of Fourier period to scale
#    COI = a number, the cone-of-influence size at the scale
#    DOFMIN = a number, degrees of freedom for each point in the wavelet power
#             (either 2 for Morlet and Paul, or 1 for the DOG)


def wave_bases(mother, k, scale, param):
    n = len(k)
    kplus = np.array(k > 0.0, dtype=float)

    if mother == "MORLET":  # -----------------------------------  Morlet

        if param == -1:
            param = 6.0

        k0 = np.copy(param)
        # calc psi_0(s omega) from Table 1
        expnt = -((scale * k - k0) ** 2) / 2.0 * kplus
        norm = np.sqrt(scale * k[1]) * (np.pi ** (-0.25)) * np.sqrt(n)
        daughter = norm * np.exp(expnt)
        daughter = daughter * kplus  # Heaviside step function
        # Scale-->Fourier [Sec.3h]
        fourier_factor = (4 * np.pi) / (k0 + np.sqrt(2 + k0**2))
        coi = fourier_factor / np.sqrt(2)  # Cone-of-influence [Sec.3g]
        dofmin = 2  # Degrees of freedom
    elif mother == "PAUL":  # --------------------------------  Paul
        if param == -1:
            param = 4.0
        m = param
        # calc psi_0(s omega) from Table 1
        expnt = -scale * k * kplus
        norm_bottom = np.sqrt(m * np.prod(np.arange(1, (2 * m)).astype(np.int64)))
        norm = np.sqrt(scale * k[1]) * (2**m / norm_bottom) * np.sqrt(n)
        daughter = norm * ((scale * k) ** m) * np.exp(expnt) * kplus
        fourier_factor = 4 * np.pi / (2 * m + 1)
        coi = fourier_factor * np.sqrt(2)
        dofmin = 2
    elif mother == "DOG":  # --------------------------------  DOG
        if param == -1:
            param = 2.0
        m = param
        # calc psi_0(s omega) from Table 1
        expnt = -((scale * k) ** 2) / 2.0
        norm = np.sqrt(scale * k[1] / gamma(m + 0.5)) * np.sqrt(n)
        daughter = -norm * (1j**m) * ((scale * k) ** m) * np.exp(expnt)
        fourier_factor = 2 * np.pi * np.sqrt(2.0 / (2 * m + 1))
        coi = fourier_factor / np.sqrt(2)
        dofmin = 1
    else:
        print("Mother must be one of MORLET, PAUL, DOG")

    return daughter, fourier_factor, coi, dofmin


# --------------------------------------------------------------------------
# WAVE_SIGNIF  Significance testing for the 1D Wavelet transform WAVELET
#
#   SIGNIF = wave_signif(Y,DT,SCALE,SIGTEST,LAG1,SIGLVL,DOF,MOTHER,PARAM)
#
# INPUTS:
#
#    Y = the time series, or, the VARIANCE of the time series.
#        (If this is a single number, it is assumed to be the variance...)
#    DT = amount of time between each Y value, i.e. the sampling time.
#    SCALE = the vector of scale indices, from previous call to WAVELET.
#
#
# OUTPUTS:
#
#    SIGNIF = significance levels as a function of SCALE
#    FFT_THEOR = output theoretical red-noise spectrum as fn of PERIOD
#
#
# OPTIONAL INPUTS:
#    SIGTEST = 0, 1, or 2.    If omitted, then assume 0.
#
#         If 0 (the default), then just do a regular chi-square test,
#             i.e. Eqn (18) from Torrence & Compo.
#         If 1, then do a "time-average" test, i.e. Eqn (23).
#             In this case, DOF should be set to NA, the number
#             of local wavelet spectra that were averaged together.
#             For the Global Wavelet Spectrum, this would be NA=N,
#             where N is the number of points in your time series.
#         If 2, then do a "scale-average" test, i.e. Eqns (25)-(28).
#             In this case, DOF should be set to a
#             two-element vector [S1,S2], which gives the scale
#             range that was averaged together.
#             e.g. if one scale-averaged scales between 2 and 8,
#             then DOF=[2,8].
#
#    LAG1 = LAG 1 Autocorrelation, used for SIGNIF levels. Default is 0.0
#
#    SIGLVL = significance level to use. Default is 0.95
#
#    DOF = degrees-of-freedom for signif test.
#         IF SIGTEST=0, then (automatically) DOF = 2 (or 1 for MOTHER='DOG')
#         IF SIGTEST=1, then DOF = NA, the number of times averaged together.
#         IF SIGTEST=2, then DOF = [S1,S2], the range of scales averaged.
#
#       Note: IF SIGTEST=1, then DOF can be a vector (same length as SCALEs),
#            in which case NA is assumed to vary with SCALE.
#            This allows one to average different numbers of times
#            together at different scales, or to take into account
#            things like the Cone of Influence.
#            See discussion following Eqn (23) in Torrence & Compo.
#
#    GWS = global wavelet spectrum, a vector of the same length as scale.
#          If input then this is used as the theoretical background spectrum,
#          rather than white or red noise.


def wave_signif(
    Y,
    dt,
    scale,
    sigtest=0,
    lag1=0.0,
    siglvl=0.95,
    dof=None,
    mother="MORLET",
    param=None,
    gws=None,
):
    n1 = len(np.atleast_1d(Y))
    J1 = len(scale) - 1
    dj = np.log2(scale[1] / scale[0])

    if n1 == 1:
        variance = Y
    else:
        variance = np.std(Y) ** 2

    # get the appropriate parameters [see Table(2)]
    if mother == "MORLET":  # ----------------------------------  Morlet
        empir = [2.0, -1, -1, -1]
        if param is None:
            param = 6.0
            empir[1:] = [0.776, 2.32, 0.60]
        k0 = param
        # Scale-->Fourier [Sec.3h]
        fourier_factor = (4 * np.pi) / (k0 + np.sqrt(2 + k0**2))
    elif mother == "PAUL":
        empir = [2, -1, -1, -1]
        if param is None:
            param = 4
            empir[1:] = [1.132, 1.17, 1.5]
        m = param
        fourier_factor = (4 * np.pi) / (2 * m + 1)
    elif mother == "DOG":  # -------------------------------------Paul
        empir = [1.0, -1, -1, -1]
        if param is None:
            param = 2.0
            empir[1:] = [3.541, 1.43, 1.4]
        elif param == 6:  # --------------------------------------DOG
            empir[1:] = [1.966, 1.37, 0.97]
        m = param
        fourier_factor = 2 * np.pi * np.sqrt(2.0 / (2 * m + 1))
    else:
        print("Mother must be one of MORLET, PAUL, DOG")

    period = scale * fourier_factor
    dofmin = empir[0]  # Degrees of freedom with no smoothing
    Cdelta = empir[1]  # reconstruction factor
    gamma_fac = empir[2]  # time-decorrelation factor
    dj0 = empir[3]  # scale-decorrelation factor

    freq = dt / period  # normalized frequency

    if gws is not None:  # use global-wavelet as background spectrum
        fft_theor = gws
    else:
        # [Eqn(16)]
        fft_theor = (1 - lag1**2) / (1 - 2 * lag1 * np.cos(freq * 2 * np.pi) + lag1**2)
        fft_theor = variance * fft_theor  # include time-series variance

    signif = fft_theor
    if dof is None:
        dof = dofmin

    if sigtest == 0:  # no smoothing, DOF=dofmin [Sec.4]
        dof = dofmin
        chisquare = chisquare_inv(siglvl, dof) / dof
        signif = fft_theor * chisquare  # [Eqn(18)]
    elif sigtest == 1:  # time-averaged significance
        if len(np.atleast_1d(dof)) == 1:
            dof = np.zeros(J1) + dof
        dof[dof < 1] = 1
        # [Eqn(23)]
        dof = dofmin * np.sqrt(1 + (dof * dt / gamma_fac / scale) ** 2)
        dof[dof < dofmin] = dofmin  # minimum DOF is dofmin
        for a1 in range(0, J1 + 1):
            chisquare = chisquare_inv(siglvl, dof[a1]) / dof[a1]
            signif[a1] = fft_theor[a1] * chisquare
    elif sigtest == 2:  # time-averaged significance
        if len(dof) != 2:
            print("ERROR: DOF must be set to [S1,S2]," " the range of scale-averages")
        if Cdelta == -1:
            print(
                "ERROR: Cdelta & dj0 not defined"
                " for " + mother + " with param = " + str(param)
            )

        s1 = dof[0]
        s2 = dof[1]
        avg = np.logical_and(scale >= 2, scale < 8)  # scales between S1 & S2
        navg = np.sum(np.array(np.logical_and(scale >= 2, scale < 8), dtype=int))
        if navg == 0:
            print("ERROR: No valid scales between " + s1 + " and " + s2)
        Savg = 1.0 / np.sum(1.0 / scale[avg])  # [Eqn(25)]
        Smid = np.exp((np.log(s1) + np.log(s2)) / 2.0)  # power-of-two midpoint
        dof = (dofmin * navg * Savg / Smid) * np.sqrt(
            1 + (navg * dj / dj0) ** 2
        )  # [Eqn(28)]
        fft_theor = Savg * np.sum(fft_theor[avg] / scale[avg])  # [Eqn(27)]
        chisquare = chisquare_inv(siglvl, dof) / dof
        signif = (dj * dt / Cdelta / Savg) * fft_theor * chisquare  # [Eqn(26)]
    else:
        print("ERROR: sigtest must be either 0, 1, or 2")

    return signif


# --------------------------------------------------------------------------
# CHISQUARE_INV  Inverse of chi-square cumulative distribution function (cdf).
#
#   X = chisquare_inv(P,V) returns the inverse of chi-square cdf with V
#   degrees of freedom at fraction P.
#   This means that P*100 percent of the distribution lies between 0 and X.
#
#   To check, the answer should satisfy:   P==gammainc(X/2,V/2)

# Uses FMIN and CHISQUARE_SOLVE


def chisquare_inv(P, V):

    if (1 - P) < 1e-4:
        print("P must be < 0.9999")

    if P == 0.95 and V == 2:  # this is a no-brainer
        X = 5.9915
        return X

    MINN = 0.01  # hopefully this is small enough
    MAXX = 1  # actually starts at 10 (see while loop below)
    X = 1
    TOLERANCE = 1e-4  # this should be accurate enough

    while (X + TOLERANCE) >= MAXX:  # should only need to loop thru once
        MAXX = MAXX * 10.0
        # this calculates value for X, NORMALIZED by V
        X = fminbound(chisquare_solve, MINN, MAXX, args=(P, V), xtol=TOLERANCE)
        MINN = MAXX

    X = X * V  # put back in the goofy V factor

    return X  # end of code


# --------------------------------------------------------------------------
# CHISQUARE_SOLVE  Internal function used by CHISQUARE_INV
#
#   PDIFF=chisquare_solve(XGUESS,P,V)  Given XGUESS, a percentile P,
#   and degrees-of-freedom V, return the difference between
#   calculated percentile and P.

# Uses GAMMAINC
#
# Written January 1998 by C. Torrence

# extra factor of V is necessary because X is Normalized


def chisquare_solve(XGUESS, P, V):

    PGUESS = gammainc(V / 2, V * XGUESS / 2)  # incomplete Gamma function

    PDIFF = np.abs(PGUESS - P)  # error in calculated P

    TOL = 1e-4
    if PGUESS >= 1 - TOL:  # if P is very close to 1 (i.e. a bad guess)
        PDIFF = XGUESS  # then just assign some big number like XGUESS

    return PDIFF


def iwavelet(Y, dt, scale, scale_indices=None, mother="MORLET", param=None):
    """
    Description:
    ------------

        Inverse wavelet transform of Y, using [Eqn(11)] in Torrence and Compo
        (1998).

        Created by Daniel Gonzalez-Duque on 07/02/2023

    Args:
    ------

    :param Y: np.ndarray,
        Wavelet transform of a time series.
    :param dt: float,
        Time interval between samples.
    :param scale: np.ndarray,
        Vector of scale indices.
    :param scale_indices: np.ndarray,
        Vector of scale indices where the reconstruction will be applied.
    :param mother: str,
        Wavelet type.
    :param param: np.ndarray
        Wavelet parameters.
    :return:
    """
    if isinstance(Y.dtype, np.complex128):
        wave = Y.real
    else:
        wave = Y
    # Calculate dj from the scales
    dj = np.log2(scale[1] / scale[0])
    # Select the appropiate parameters [see Table(2)]
    if mother == "MORLET":  # ----------------------------------  Morlet
        empir = [2.0, -1, -1, -1, -1]
        if param is None:
            param = 6.0
            empir[1:] = [0.776, 2.32, 0.60, np.pi ** (1 / 4)]
        k0 = param
        # Scale-->Fourier [Sec.3h]
        fourier_factor = (4 * np.pi) / (k0 + np.sqrt(2 + k0**2))
    elif mother == "PAUL":
        empir = [2, -1, -1, -1, -1, 1.079]
        if param is None:
            param = 4
            empir[1:] = [1.132, 1.17, 1.5, 0.867]
        m = param
        fourier_factor = (4 * np.pi) / (2 * m + 1)
    elif mother == "DOG":  # -------------------------------------Paul
        empir = [1.0, -1, -1, -1, -1]
        if param is None:
            param = 2.0
            empir[1:] = [3.541, 1.43, 1.4, 0.867]
        elif param == 6:  # --------------------------------------DOG
            empir[1:] = [1.966, 1.37, 0.97, 0.884]
        m = param
        fourier_factor = 2 * np.pi * np.sqrt(2.0 / (2 * m + 1))
    else:
        raise ValueError("Mother must be one of MORLET, PAUL, DOG")
    c_delta = empir[1]  # reconstruction factor
    wave_0 = empir[4]  # wavelet zeroth moment

    # Filtering over scales
    if scale_indices is None:
        scale_indices = np.arange(len(scale)).astype(np.int64)

    new_wave = wave[scale_indices, :]
    new_scale = scale[scale_indices]

    # Reconstruction with Torrence and Compo (1998) (Eq[11])
    factor = (dj * np.sqrt(dt)) / (c_delta * wave_0)
    r_sum = np.transpose(np.sum(new_wave.T / new_scale**0.5, axis=-1))
    c_r = r_sum * factor

    return c_r


def calculate_global_wavelet_spectrum(wave):
    """
    Description:
    ------------

        Calculates the global wavelet spectrum (GWS) from the wavelet transform
        using equation (22) in Torrence and Compo (1998). This function also
        finds the index for the peaks in the global wavelet spectrum.

        Created by Daniel Gonzalez-Duque on 11/19/2023.

        References:
        
        Torrence, C., and G. P. Compo, 1998: A Practical Guide to Wavelet
        Analysis. Bull. Amer. Meteor. Soc., 79, 61–78,
        https://doi.org/10.1175/1520-0477(1998)079<0061:APGTWA>2.0.CO;2

    Args:
    -----
        :param wave: np.ndarray,
            Wavelet transform.
        :return: gws: np.ndarray,
            Global wavelet spectrum.
            peaks: np.ndarray,
            Indices for the peaks in the GWS.
    """
    gws = np.nanmean(np.abs(wave) ** 2, axis=1)
    # find peaks in the gws
    peaks, _ = find_peaks(gws)
    return gws, peaks


def calculate_scale_averaged_wavelet_power(wave, scales, ds, dj, c_delta):
    """
    Description:
    ------------

        Calculates the scale-averaged wavelet power (SAWP) from the wavelet
        transform using equation (9) in Zolezzi and Güneralp (2016).

        Created by Daniel Gonzalez-Duque on 11/19/2023.

        References:
        
        Zolezzi, G., & Güneralp, I. (2016). Continuous wavelet
        characterization of the wavelengths and regularity of
        meandering rivers. Geomorphology, 252, 98–111.
        https://doi.org/10.1016/j.geomorph.2015.07.029

    Args:
    -----
        :param wave: np.ndarray,
            Wavelet transform.
        :param scales: np.ndarray,
            Vector of scales.
        :param dt: float,
            Time or spatial resolution of the data.
        :param dj: float,
            Scale increment.
        :param c_delta: float,
            Reconstruction factor.
        :return:
            Scale-averaged wavelet power.
    """
    power = np.abs(wave) ** 2
    scales_p = np.tile(scales, (power.shape[1], 1)).T
    sawp = (dj * ds) / c_delta * np.nansum(power / scales_p, axis=0)

    return sawp
