import numpy as np
import pandas as pd
import tensorflow as tf 
import sys
import wget
from tensorflow import keras 
from sklearn.metrics import r2_score, mean_squared_error
from math import sqrt
from sklearn.model_selection import train_test_split
from sklearn import preprocessing
# Importing the dataset

df = pd.read_csv('train.csv').dropna()
dataset = df.iloc[:, 3:-3]
sys.stdout=open("train_output.txt","w")
print(dataset.head())
dataset = df.groupby(by=["country"], dropna=True).sum()
X = dataset.loc[:,dataset.columns != "daily_vaccinations"]
y = dataset.loc[:,dataset.columns == "daily_vaccinations"]

# Splitting the dataset into the Training set and Test set
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 42)

# Feature Scaling
model = keras.Sequential([
    keras.layers.Dense(512,input_dim = X_train.shape[1],kernel_initializer='normal', activation='relu'),
    keras.layers.Dense(512,kernel_initializer='normal', activation='relu'),
    keras.layers.Dense(256,kernel_initializer='normal', activation='relu'),
    keras.layers.Dense(256,kernel_initializer='normal', activation='relu'),
    keras.layers.Dense(128,kernel_initializer='normal', activation='relu'),
    keras.layers.Dense(1,kernel_initializer='normal', activation='linear'),
])

model.compile(loss='mean_absolute_error', optimizer='adam', metrics=['mean_absolute_error'])

model.fit(X_train, y_train, epochs=50, validation_split = 0.3)

prediction = model.predict(X_test)
print(prediction)
sys.stdout.close()