#!/usr/bin/env python
# coding: utf-8

# # Exercise 9.2 - Solution
# Large arrays of radio antennas can be used to measure cosmic rays by recording the electromagnetic radiation generated in the atmosphere.
# These radio signals are strongly contaminated by galactic noise as well as signals from human origin. Since these signals appear to be similar to the background, the discovery of cosmic-ray events can be challenging.
# ## Identification of signals
# In this exercise, we design an RNN to classify if the recorded radio signals contain a cosmic-ray event or only noise.
# 
# The signal-to-noise ratio (SNR) of a measured trace $S(t)$ is defined as follows:
# 
# $$\mathrm{SNR}=\frac{S^{\mathrm{signal}}(t)_\mathrm{max}}{\mathrm{RMS}[S(t)]},$$  
# where $S^{\mathrm{signal}}(t)_\mathrm{max}$ denotes the maximum amplitude of the (true) signal.
# 
# Typical cosmic-ray observatories enable a precise reconstruction at an SNR of roughly 3.
# 
# We choose a challenging setup in this task and try to identify cosmic-ray events in signal traces with an SNR of 2.  
# Training RNNs can be computationally demanding, thus, we recommend to use a GPU for this task.

# In[2]:


import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt

layers = keras.layers


# ### Load and prepare dataset
# In this task, we use a simulation of cosmic-ray-induced air showers that are measured by radion antennas.  
# For more information, see https://arxiv.org/abs/1901.04079.  
# The task is to design an RNN which is able to identify if the measured signal traces (shortened to 500 time steps) contains a signal or not.

# In[2]:


import os
import gdown
url = "https://drive.google.com/u/0/uc?export=download&confirm=HgGH&id=1R-qfxO1jVh88TC9Gnm9JGMomSRg0Zpkx"
output = 'radio_data.npz'

if os.path.exists(output) == False:
    gdown.download(url, output, quiet=True)

f = np.load(output)
n_train = 40000

x_train, x_test = f["traces"][:n_train], f["traces"][n_train:]  # measured traces (signal + colored noise)
signals = f["signals"]  # signal part (only available for cosmic-ray events)

labels = (signals.std(axis=-1) != 0).astype(float)  # define training label (1=cosmic event, 0=noise)
y_train, y_test = labels[:n_train], labels[n_train:]


# ## Plot example signal traces
# Left: signal trace containing a cosmic-ray event. The underlying cosmic-ray signal is shown in red, the backgrounds + signal is shown in blue.
# Right: background noise.

# In[3]:


from matplotlib import pyplot as plt
fs = 180e6  # Sampling frequency of antenna setup 180 MHz
t = np.arange(500) / fs * 1e6
idx = np.random.randint(0, labels.sum()-1)
idx2 =  np.random.randint(0, n_train - labels.sum())

plt.figure(1, (12, 4))
plt.subplot(1, 2, 1)
plt.plot(t, np.real(f["traces"][labels.astype(bool)][idx]), linewidth = 1, color="b", label="Measured trace")
plt.plot(t, np.real(signals[labels.astype(bool)][idx]), linewidth = 1, color="r", label="CR signal")
plt.ylabel('Amplitude / mV')
plt.xlabel('Time / $\mu \mathrm{s}$')
plt.legend()
plt.title("Cosmic-ray event")
plt.subplot(1, 2, 2)

plt.plot(t, np.real(x_train[~y_train.astype(bool)][idx2]), linewidth = 1, color="b", label="Measured trace")
plt.ylabel('Amplitude / mV')
plt.xlabel('Time / $\mu \mathrm{s}$')
plt.legend()
plt.title("Noise event")

plt.grid(True)
plt.tight_layout()


# In[4]:


sigma = x_train.std()
x_train /= sigma
x_test /= sigma


# ### Define RNN model
# In the following, design a cosmic-ray model to identify cosmic-ray events using an RNN-based classifier.

# In[5]:


model = keras.models.Sequential()
model.add(layers.Bidirectional(layers.LSTM(32, return_sequences=True), input_shape=(500,1)))
model.add(layers.Bidirectional(layers.LSTM(64, return_sequences=True)))
model.add(layers.Bidirectional(layers.LSTM(10, return_sequences=True)))
model.add(layers.Flatten())
model.add(layers.Dropout(0.3))
model.add(layers.Dense(1, activation="sigmoid"))

model.summary()


# #### Pre-processing of data and RNN training

# In[6]:


model.compile(
    loss='binary_crossentropy',
    optimizer=keras.optimizers.Adam(1e-3, decay=0.00008),
    metrics=['accuracy'])


results = model.fit(x_train[...,np.newaxis], y_train,
                    batch_size=128,
                    epochs=100,
                    verbose=1,
                    validation_split=0.1,
                    callbacks = [keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=5, verbose=1, min_lr=1e-5),
                                 keras.callbacks.EarlyStopping(patience=15, verbose=1)]
                    )


# In[7]:


model.evaluate(x_test[...,np.newaxis], y_test)


# ### Plot loss and accuracy

# In[8]:


plt.figure(1, (12, 4))
plt.subplot(1, 2, 1)
plt.plot(results.history['loss'])
plt.plot(results.history['val_loss'])
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper right')

plt.subplot(1, 2, 2)
plt.plot(results.history['accuracy'])
plt.plot(results.history['val_accuracy'])
plt.ylabel('accuracy')
plt.xlabel('epoch')
plt.legend(['train', 'val'], loc='upper left')
plt.tight_layout()