import sys
sys.path.append("./")
sys.path.append("..")

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from src.data.data import ReadPesSpec, PesChannelSelector
from src.data.data_preproc import SpecPreprocessing

from src.models.find_components.Find_Component import FindPCAcomps

from src.models.fit_methods.model import Model

from src.utils.utils import create_experiment_dirs, save_train_test_h5, save_rec_data_h5, save_checkpoint, save_checkpoint_fft

from sklearn.model_selection import train_test_split
import numpy as np

run_directory = "/gpfs/exfel/exp/SA3/202121/p002935/raw"
run_number = "r0015" 
spec_ofset = -2
n_pca_comps_pes = 400
n_pca_comps_spec = 20
pes_pca_preprocessing = True
spec_pca_preprocessing = True
spec_gc = True
spec_fft = False
use_data_subset = True
data_subset_start, data_subset_end = 1000, 6000     # Min Max 
channel_names = "all_chs"                           # "default", "all_chs", or list with channel names ex  ["channel_1_D", "channel_2_B"]
test_size = 0.05 
model_type = "bfgs_pca_eps"                                 # "bfgs", "bfgs_pca", "bfgs_pca_eps"

exp_dir = "test3_pulseen_short_test_eps_" + run_number
general_comment = exp_dir

data_info = {"run_directory": run_directory,
            "run_number": run_number,
            "spec_ofset": spec_ofset,
            "exp_dir": exp_dir,
            "n_pca_comps_pes": n_pca_comps_pes,
            "n_pca_comps_spec": n_pca_comps_spec,
            "pes_pca_preprocessing": pes_pca_preprocessing,
            "spec_pca_preprocessing": spec_pca_preprocessing,
            "spec_gc": spec_gc,
            "spec_fft": spec_fft,
            "use_data_subset": use_data_subset,
            "data_subset_start": data_subset_start,
            "data_subset_end": data_subset_end,
            "channel_names": channel_names,
            "test_size": test_size,
            "model_type": model_type,
            "general_comment": general_comment

            }


experiment_dir, summary_dir, checkpoint_dir, output_dir, test_dir = create_experiment_dirs(exp_dir)

import datetime
now = datetime.datetime.now()
print ("Start Read Pes Spec Data : ")
print (now.strftime("%Y-%m-%d %H:%M:%S"))

# Read Pes Spec Data
RPS = ReadPesSpec(data_info)
data_dict = RPS.get_data()

now = datetime.datetime.now()
print ("End Read Pes Spec Data : ")
print (now.strftime("%Y-%m-%d %H:%M:%S"))

# Model
model_instance = Model(model_type=model_type, data_info=data_info)
X_train, X_test, Y_train, Y_test = model_instance.preprocess(data_dict, data_info)              # Model Preprocess

now = datetime.datetime.now()
print ("End Model Preprocessor : ")
print (now.strftime("%Y-%m-%d %H:%M:%S"))

model_instance.save_preproc()
print("################")
print(X_train.shape)
print("################")
model_instance.fit_eval(X_train, Y_train, X_test, Y_test)           # Model FIT

now = datetime.datetime.now()
print ("End Model Train : ")
print (now.strftime("%Y-%m-%d %H:%M:%S"))

model_instance.save_latest_ckp()

result = model_instance.predict(input_value=X_test)

model_instance.save_prediction()




print("FINISHED ALL")