from pycrtools import *
import matplotlib.pyplot as plt

plt.figure()

filename_lofar = LOFARSOFT + "/data/lofar/VHECR_example.h5"

datafile = open(filename_lofar)

fftdata = datafile.empty("FFT_DATA")

nblocks = datafile["MAXIMUM_READ_LENGTH"] / datafile["BLOCKSIZE"]

avspectrum = hArray(float, dimensions=fftdata, name="Average spectrum")

# Calculate spectral power
for block in range(nblocks):
    datafile["BLOCK"] = block
    fftdata.read(datafile, "FFT_DATA")
    hSpectralPower(avspectrum[...], fftdata[...])

# Create 1st plot
frequencies = datafile["FREQUENCY_DATA"].setUnit("M", "")
plt.subplot(1, 2, 1)
plt.title("Average spectrum for two antennas")
plt.xlabel(frequencies.getKey("name") + " [" + frequencies.getUnit() + "]")
plt.ylabel(avspectrum.getKey("name") + " [" + avspectrum.getUnit() + "]")
for i in range(2):
    plt.semilogy(frequencies.vec(), avspectrum[i].vec())

# Prepare block
datafile["BLOCK"] = 2
datafile["BLOCKSIZE"] = 2 ** 16
timeall = datafile["TIME_DATA"]
timeall.setUnit("m", "s")
fxall = datafile["TIMESERIES_DATA"]

# Create 2nd plot
plt.subplot(1, 2, 2)
plt.title("Time series of antenna 0")
plt.plot(timeall.vec(), fxall[0].vec())
plt.xlabel(timeall.getKey("name") + " [" + timeall.getUnit() + "]")
plt.ylabel("Electric Field [ADC counts]")