# Equatorial SpreadF Occurrence prediction - a single station approach using keras functional API
#The Four input parameters (DOY, Vzmax, Time of Vzmax and F10.7) and two output parameter is either occurrence status (1 or 0) or TOC
#TOC (Time of commencement of ESF)
#import the reqiured modules and function
import os
import sklearn
from sklearn.metrics import mean_absolute_error
from sklearn.tree import DecisionTreeRegressor
import numpy as np
from numpy import loadtxt # the function to load data
from sklearn.model_selection import train_test_split
# load data to be used below here
# input variables are DOY, Vzmax, Time of Vzmax and F10.7
#target variables are (occurrence status 1 or 0, and TOC)  
FileName="ffVzDatapreparedF107.txt"
data_all = loadtxt(FileName)
#format of data_all [1] doy [2] Vz,max  [3] Time of Vz,max [4] F10.7  [5] ESF indicator and [6] TOC 
#divide data into train and test data 
X=data_all[:,0:4] # The number of input parameters can be varied here as discussed in the manuscript 
Y=data_all[:,4]  # Y here is a singel column of ESF indicator on the column 5 of the input data.
#F107=data_all[:,3:4] # This line and the next can be jointly used with line 19 to vary the number of input
#X=np.append(Xx, F107, axis=1 )
#for TOC change 4 to 5 on line 23
# on the next line the data is split into train and data for both input and output
train_X, val_X, train_y, val_y = train_test_split(X,Y, random_state=0)
# X the input data doy, Vz,max, Time of Vz,max and F10.7
# below is a function to test the perfomance of the generated model - the input are either 
#[1] ESF indicator [2] ESF indicator [3] predicted ESF indicator  [4] predicted ESF indicator - for ESF prediction
# OR
#[1] TOC [2] TOC [3] predicted TOC  [4] predicted TOC -  for TOC prediction
#to make the input consistent with that of Joint model, columns were repeated
#the output are percentage perfomance for ESF, percentage perfomance for ESF - repeated output
#the code can also estimate %perfomance for TOC within + or - 30 minutes as implemented for Joint model.
def getPerfScore(val_output):
    nTotal=0
    nTrue=0
    n30min=0
    nTime=0
    dataLen=val_output[:,1]
    for n in range(0,len( dataLen)):
        nTotal=nTotal+1
        if val_output[n][0]==val_output[n][2]:        
        #if val_output[n][1]==val_output[n][3]:
            nTrue=nTrue+1        
        if val_output[n][1]==val_output[n][3]:
            nTime=nTime+1 
#        if (val_output[n][3]>=val_output[n][1]-0.5 and val_output[n][3]<=val_output[n][1]+0.5):
#            n30min=n30min+1 
#    perf_score=(nTime/nTotal)*100 
    perf_score=[(nTrue/nTotal)*100, (nTime/nTotal)*100] # The current form of this function is for ESF prediction
    return perf_score
#model  
model =DecisionTreeRegressor(max_leaf_nodes=400, random_state=0) 
model.fit(train_X, train_y)
#prediction
y_preds=model.predict(val_X)
# the code below will concatenate two rows 
fData1=np.vstack([val_y, val_y]) # repeated column
fData2=np.vstack([y_preds, y_preds]) # repeated column
#concatenate output data to reflect the input into function 'getPerfScore'
ffData=np.concatenate([np.transpose(fData1),np.transpose(fData2)], axis=1)
print(getPerfScore(ffData))
#save the data used as input into into function 'getPerfScore' 
np.savetxt('outESFModel.out', ffData, delimiter=',')   #  is an 
#The format of '' is [1] Observed ESF Occurrence [2] Observed ESF Occurrence 
#[3]Predicted ESF Occurrence [4] Predicted ESF Occurrence