block_size = 512

# results in the amplitude spectra
fft = ac.RFFT(source=ts, block_size=block_size, window='Rectangular')

spec = next(fft.result(num=1))[0]  # return a single snapshot
time_block = next(ts.result(num=block_size))[:, 0]

print("Parseval's theorem:")
print('signal energy in time domain', np.round(np.sum(time_block**2)))
print(
    'signal energy in frequency domain (single-sided)',
    np.round((1 / block_size * np.sum(spec * spec.conjugate())).real),
)

# It can be seen, that the energy is not preserved when using the default 'none' scaling. Therefore,
# we set the scaling to 'energy' to compensate for the energy loss due to truncation to a
# single-sided spectrum.

fft.scaling = 'energy'
spec = next(fft.result(num=1))[0]
print(
    'signal energy in frequency domain (single-sided) with energy scaling',
    np.round((1 / block_size * np.sum(spec * spec.conjugate() / 2)).real),
)

fft.scaling = 'none'
ifft = ac.IRFFT(source=fft)
time_block_rec = next(ifft.result(num=block_size))[:, 0]
print('reconstructed signal energy in time domain', np.round(np.sum(time_block_rec**2)))

import matplotlib.pyplot as plt

fft.scaling = 'amplitude'
spectrogram = ac.tools.return_result(fft)

# plot the power spectrogram
plt.figure()
plt.imshow(np.abs(spectrogram.T), origin='lower', aspect='auto', extent=(0, t_in_s, 0, sample_freq / 2), vmax=2.0)
plt.xlabel('Time / s')
plt.ylabel('Frequency / Hz')
plt.colorbar(label='Power Spectrum')
plt.show()

plt.figure()
for block_size in [4096, 256]:
    fft.block_size = block_size
    spec = next(fft.result(num=1))[0]
    plt.plot(fft.freqs, np.abs(spec), label=f'block_size = {block_size}')
plt.xlim(0, sample_freq / 2)
plt.xlabel('Frequency / Hz')
plt.ylabel('Amplitude Spectrum')
plt.legend()
plt.grid()
plt.show()

fft.scaling = 'none'
ap = ac.AutoPowerSpectra(source=fft, scaling='psd')  # results in the power spectrum
avg = ac.Average(source=ap, naverage=64)  # results in the time averaged power spectrum

for block_size in [256, 1024]:
    plt.figure()
    fft.block_size = block_size
    for navg in [1, 100]:
        avg.naverage = navg
        spectrum = next(avg.result(num=1))
        levels = ac.L_p(spectrum)
        plt.plot(fft.freqs, levels[0], label=f'navg = {navg}, block_size = {block_size}')
    plt.xlabel('Frequency / Hz')
    plt.ylabel('PSD / dB')
    plt.grid()
    plt.legend()
    plt.semilogx()
    plt.show()

fft.block_size = 128

mixed_signal = noise.signal()[:, np.newaxis]
ts = ac.TimeSamples(data=np.concatenate([mixed_signal, 0.5 * mixed_signal], axis=1), sample_freq=sample_freq)

# compare with PowerSpectra
ps1 = ac.PowerSpectra(source=ts, block_size=fft.block_size, cached=False)
csm_comp = ps1.csm[:, :, :]

fft.source = ts  # we need to change the source of the fft object to the new TimeSamples object
ps = ac.CrossPowerSpectra(source=fft)
avg = ac.Average(source=ps, naverage=int(ps1.num_blocks))
csm = next(avg.result(num=1))

# reshape the cross-spectral matrix to a 3D array of shape (numfreq, num_channels, num_channels)
csm = csm.reshape(fft.numfreqs, ts.num_channels, ts.num_channels)

plt.figure()
colors = ['r', 'b']
for i in range(ts.num_channels):
    c = colors[i]
    auto_pow_spectrum = ac.L_p(csm_comp[:, i, i].real)
    plt.plot(fft.freqs, auto_pow_spectrum, label=f'PowerSpectra ({i})', color=c, linestyle='-', linewidth=0.8)
    auto_pow_spectrum = ac.L_p(csm[:, i, i].real)
    plt.plot(fft.freqs, auto_pow_spectrum, label=f'CrossPowerSpectra ({i})', color=c, linestyle='--')
plt.xlabel('Frequency / Hz')
plt.ylabel('Power Spectrum / dB')
plt.grid()
plt.legend()
plt.semilogx()
plt.show()