%matplotlib inline

import pandas as pd
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, LeakyReLU, Dropout, ReLU, GRU, TimeDistributed, Conv2D, MaxPooling2D, Flatten
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.callbacks import EarlyStopping
from jlab import load_test_data, get_test_detector_plane

"""
Explanation: 04 - Feed Forward Approach

After getting 4th place out of 9, I decided to try a simpler strategy
Input: one detector plane, output: the next detector plane
Each track has 24 training samples
Let's try it out and see what happens
End of explanation
"""


data = pd.read_csv('MLchallenge2_training.csv')

"""
Explanation: Load up and prep the datasets
End of explanation
"""


def extract_X_y(df):
    N_TRACKS = len(df)

    X_train_single = np.ndarray(shape=(N_TRACKS * 22, 7))
    y_train_single = np.ndarray(shape=(N_TRACKS * 22, 5))
    for ix in range(0, N_TRACKS):
        alldetvals = df.iloc[ix]
        for twix in range(2, N_DETECTORS-1):
            detvals = list(alldetvals[[f"{kin}{twix}" for kin in ["x", "y", "z", "px", "py", "pz"]]])
            detvals.append(twix)
            detvals = np.array(detvals)
            X_train_single[(22*ix + twix - 2)] = detvals
            nextdetvals = np.array(alldetvals[[f"{kin}{twix + 1}" for kin in ["x", "y", "px", "py", "pz"]]])
            y_train_single[(22*ix + twix - 2)] = nextdetvals
            
    return X_train_single, y_train_single

X_train, y_train = extract_X_y(data.loc[0:5000])

X_test, y_test = extract_X_y(data.loc[5000:6000])

X_train[50]

y_train[50]

from sklearn.ensemble import RandomForestRegressor

rfr = RandomForestRegressor(n_estimators=1000)

rfr.fit(X_train, y_train)

y_pred = rfr.predict(X_test)

from sklearn.metrics import mean_squared_error

mse = mean_squared_error(y_test, y_pred)

mse

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import GRU, Dense, LeakyReLU, Dropout
from tensorflow.keras.callbacks import EarlyStopping
import joblib

def lrelu(x):
    return LeakyReLU()(x)


def dense_model(dense_units=100, dropout_rate=0.25):
    """Model definition.
    
    Parameters
    ----------
    dense_units : int
    dropout_rate : float
    
    Returns
    -------
    tensorflow.keras.models.Sequential
    
    """
    
    model = Sequential()
    
    model.add(Dense(500, activation=lrelu,
                  input_dim=7))
    model.add(Dropout(dropout_rate))
    model.add(Dense(100, activation=lrelu))
    model.add(Dropout(dropout_rate))
    model.add(Dense(5))
    
    model.compile(loss='mse', optimizer='adam')
    
    return model

model = dense_model()

model.summary()

es = EarlyStopping(monitor='val_loss', mode='min',
                   patience=5, restore_best_weights=True)
history = model.fit(
    x=X_train,
    y=y_train,
    validation_data=(X_test, y_test),
    callbacks=[es],
    epochs=50,
)

model.save("dense_model.h5")
joblib.dump(history.history, "dense_model.history")

import matplotlib.pyplot as plt

# Plot training & validation loss values
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.show()

"""
Explanation: Construct the training data and targets

For each track
For each detector between 1-24
Extract x, y, z, px, py, pz for detectors 1-24
Add detector ID
Append to X_train
Append the next detector's x, y, px, py, pz to y_train



Training Dataset
End of explanation
"""


