import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

country = "Nepal"

df_confirmed = pd.read_csv("confirmed_cases.csv")

df_confirmed_country= df_confirmed[df_confirmed["Country/Region"] == country]
df_confirmed_country= pd.DataFrame(df_confirmed_country[df_confirmed_country.columns[1:]].sum(), columns=["confirmed"])
df_confirmed_country.index = pd.to_datetime(df_confirmed_country.index,format='%m/%d/%y')

df_confirmed_country.plot(figsize=(12,12),title="COVID confirmed cases")
# Use data untill 14 days before as training
x= len(df_confirmed_country)-14

train = df_confirmed_country.iloc[:x]
test = df_confirmed_country.iloc[x:]

#scale or normalize data as the data is too skewed
from sklearn.preprocessing import MinMaxScaler
scaler = MinMaxScaler()
scaler.fit(train)

train_scaled = scaler.transform(train)
test_scaled = scaler.transform(test)

# Use Timeseries generator to generate data in sequences
from keras.preprocessing.sequence import TimeseriesGenerator

#Sequence size has an impact on prediction, especially since COVID is unpredictable
seq_size = 7 ## number of steps (lookback)
n_features = 1 ## number of features. This dataset is univariate so it is 1
train_generator = TimeseriesGenerator(train_scaled,train_scaled, length = seq_size, batch_size=1)
print("Total number of samples in the original training data= ", len(train)) #271
print("Total number of samples in the generated data =",len(train_generator)) #264 with seq_size=7

#Check data shape from generator
x,y= train_generator[10]

#Also generate test data
test_generator = TimeseriesGenerator(test_scaled, test_scaled, length=seq_size, batch_size=1)
print("Total number of samples in the original training data=", len(test)) #14 as we are using last 14 days for test
print("Total number of samples in the generated data = ", len(test_generator)) #7
# Check data shape from generator
x,y = test_generator[0]

from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout, Activation

#Define Model
model = Sequential()
model.add(LSTM(150, activation ='relu',return_sequences=True,input_shape=(seq_size,n_features)))
model.add(LSTM(64,activation='relu'))
model.add(Dense(64))
model.add(Dense(1))
model.compile(optimizer='adam',loss='mean_squared_error')

model.summary()
print('Train...')

########################

history = model.fit_generator(train_generator, validation_data=test_generator, epochs=50, steps_per_epoch=10)

#plot the training and validation accuracy and loss at each epoch
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1,len(loss) +1)
plt.plot(epochs,loss,'y', label='Training loss')
plt.plot(epochs, val_loss, 'r', label='Validation Loss')
plt.title('Training and validation loss')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

#forecast
prediction = [] # Empty list to populate later with prediction

current_batch = train_scaled[-seq_size:] #Final data points in train
current_batch = current_batch.reshape(1,seq_size, n_features) #Reshape

## Predict future, beyond test dates
future = 7
for i in range(len(test) + future):
    current_pred = model.predict(current_batch)[0]
    prediction.append(current_pred)
    current_batch = np.append(current_batch[:,1:,:],[[current_pred]],axis=1)

### Inverse transform to before scaling so we get actual numbers
rescaled_prediction = scaler.inverse_transform(prediction)

time_series_array = test.index # Get dates for test data

# Add new dates for the forecast period
for k in range(0,future):
    time_series_array = time_series_array.append(time_series_array[-1:] + pd.DateOffset(1))

# Create a dataframe to capture the forecast data
df_forecast = pd.DataFrame(columns=["actual_confirmed","predicted"],index=time_series_array)

df_forecast.loc[:,"predicted"] = rescaled_prediction[:,0]
df_forecast.loc[:,"actual_confirmed"] = test["confirmed"]

#Plot
df_forecast.plot(title="Prediction for next 7 days")



