from obspy import Stream
from scipy.fft import rfft, rfftfreq
from scipy.interpolate import interp1d
import numpy as np


def calculate_amp_spectra(station_stream: Stream):
    """
    Function returns frequencies and corresponding mean amplitude spectra (for all available traces in stream).
    All spectra are interpolated
    :param station_stream:
    :return:
    """

    # Fourier's transformation of signal for each trace (channel) in the stream
    freq_bins = None
    psd = 0
    for tr in station_stream:
        fft_signal = rfft(tr.data)
        freq_bins = rfftfreq(len(tr.data), d=tr.stats.delta)
        amp_spectrum = abs(tr.stats.delta * fft_signal)  # Multiplication by sampling step
        amp_spectrum = amp_spectrum * np.sqrt(2)  # Multiplication by square root to calibrate energy
        psd += amp_spectrum ** 2  # for calculation of root-mean-square
    # Calculation of root-mean-square for all traces in the stream
    mean_amp_spectra = psd ** 0.5

    # Interpolation
    fi = np.logspace(np.log10(freq_bins[1]), np.log10(freq_bins[-1]), len(freq_bins))
    interpolation_function = interp1d(freq_bins, mean_amp_spectra, kind='linear', fill_value='extrapolate')

    return fi, interpolation_function(fi)