# -*- coding: utf-8 -*-
# coding: utf-8
import numpy as np
import sys
sys.path.append("../../Methods")
from DiET import Params_DiET, Make_Spec, get_model

Station_number = 2
Input_data = np.ones((Station_number, Params_DiET.frequency_size, Params_DiET.time_size, Params_DiET.num_channels)) ### Spectrogram with all pixel values of 1 is identified as noise. 
Input_data [0, 25:73, :, :] = 0.2 ### Remain only the power in the 2-8 Hz

##### Loading model #####
model = get_model()
#model.summary()
pred_result = model.predict(Input_data, batch_size = Station_number, verbose = 0)
print('Probability', pred_result)
print('EQ', pred_result[:,0])
print('T', pred_result[:,1])
print('N', pred_result[:,2])

##### Save user results #####
save = 'of'
if save == 'on':
    header_info = ['EQ', 'T', 'N']
    Save_df = pd.DataFrame(data = pred_result)
    Save_df.to_csv('../../Results/User_results/DiET_prediction.csv', index = False, header = header_info)

quit()
