#! /usr/bin/python3
import numpy as np
import torch
from torch import nn
import pandas as pd

from sklearn.metrics import accuracy_score, precision_score, recall_score
from sklearn.preprocessing import LabelEncoder
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self, input_dim):
        super(Model, self).__init__()
        self.layer1 = nn.Linear(input_dim, 50)
        self.layer2 = nn.Linear(50, 20)
        self.layer3 = nn.Linear(20, 2)
        
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = F.softmax(self.layer3(x))
        return x


test_df = pd.read_csv('testing_data.csv')

X = test_df[['Pclass', 'Sex', 'Age','SibSp', 'Fare']]
Y = test_df[['Survived']]
Y = np.ravel(Y) 
encoder = LabelEncoder()
encoder.fit(Y)
Y = encoder.transform(Y)

model = Model(X.shape[1])
model.load_state_dict(torch.load('model.pt'))


x_test = torch.tensor(X.values, dtype=torch.float32)
pred = model(x_test)
pred = pred.detach().numpy()
acc = accuracy_score(Y, np.argmax(pred, axis=1))
prec = precision_score(Y, np.argmax(pred, axis=1))
recall = recall_score(Y, np.argmax(pred, axis=1))
print ("The accuracy is", acc)
print ("The precission score is ", prec)
print ("The recall score is ", recall)

file = open('metrics.txt', 'a')
file.write(str(acc)  + '\t' + str(prec) + '\t' + str(recall))
file.close()


np.savetxt('prediction.tsv', pred, delimiter='\t')