# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: keras
#     language: python
#     name: keras
# ---

# ### 1. Import modules

# +
import os
import os.path
import re
import pickle
import random
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import keras
from keras import models
from keras import layers
import keras.backend as K
from keras.models import load_model 
from keras.utils import to_categorical
from sklearn.preprocessing import MinMaxScaler
from sklearn.utils import shuffle
from sklearn.utils import class_weight

import tensorflow as tf
# -

# ### 2. Build model

# +
N_XICs = 170  # specify the number of XICs for the RSM data here
N_RTs = 12    # specify the RT cycles for the RSM data here

def build_model(n_features, n_rt = 10):
    model = models.Sequential()
    model.add(layers.LSTM(128, dropout = 0.4, recurrent_dropout = 0.3, 
                          return_sequences = True, input_shape = (n_rt, n_features)))
    model.add(layers.LSTM(64, dropout = 0.4, recurrent_dropout = 0.3, activation = "relu"))
    model.add(layers.Dense(16, activation = "relu"))
    model.add(layers.Dense(1, activation = "sigmoid"))
    model.compile(optimizer = "adam", loss = "binary_crossentropy", metrics = ["accuracy"])
    return model

model = build_model(N_XICs, n_rt = N_RTs)
# -

# ### 3. Import your own training data

# * The training data should consist of two numpy array.
# * X: a numpy array with three dimensions, (Number of peptides, N_RTs, N_XICs).
# * y: a 1-d numpy array containing 1 as real peptide and 0 as decoy.

X, y = pickle.load(open("my_data.pkl", "rb"))

# ### 4. Assign class weight

# * This step is signicant especially for unbalanced-labeled data.

my_class_weight = class_weight.compute_class_weight("balanced", np.array([0, 1]), y)

# ### 5. Train and validation

history = model.fit(X, np.array(y), epochs = 100, batch_size = 1024, 
                    validation_split = 0.3, class_weight = {0 : my_class_weight[0], 1 : my_class_weight[1]})

# ### 6. Choose the best epoch and retrain the model

# +
best_epoch = 11 # specify the best epoch here

model1 = build_model(N_XICs, n_rt = N_RTs)
model1.fit(X, np.array(y), epochs = best_epoch, batch_size = 1024, class_weight = {0 : my_class_weight[0], 1 : my_class_weight[1]})
# -

# ### 7. Save the model

# +
model1.save("./models/My_own_model.h5")

intermediate_layer_model = models.Model(inputs = model1.input,outputs = model1.get_layer("dense_2").output)
intermediate_layer_model.save("./models/My_own_model.rep.h5")
# -

# ### 8. Modify the model names in the `dream_score.py` script

# From:
# ```python
# BM_model_file = os.path.join(os.path.dirname(sys.argv[0]), "models/M2_170_12_11epochs_BM.h5")
# ```
# To:
# ```python
# BM_model_file = os.path.join(os.path.dirname(sys.argv[0]), "models/My_own_model.h5")
# ```
# &
#
# From:
# ```python
# RM_model_file = os.path.join(os.path.dirname(sys.argv[0]), "models/M2_170_12_11epochs_RM.h5")
# ```
# To:
# ```python
# RM_model_file = os.path.join(os.path.dirname(sys.argv[0]), "models/My_own_model.rep.h5")
# ```

# ### 9. Run DreamDIA-XMBD with your own parameters

# * Run `DreamDIA-XMBD-v2.0.2/DreamDIA.py dreamscore --help` to see the parameters.

StarcoderdataJupyterScriptsDedupFiltered