import pandas as pd
import torch
import torch.nn as nn
import itertools

from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import f1_score as f1
from sklearn.metrics import confusion_matrix
from sklearn.metrics import *

from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader

test_file = pd.read_csv('test.data').drop('Unnamed: 0', axis=1)
df_test = test_file.dropna()

X_test = df_test.drop('class', axis=1)
Y_test = df_test['class']

scaler = MinMaxScaler()
X_test =  scaler.fit_transform(X_test)

x_test_tensor = torch.tensor(X_test).float()
y_test_tensor = torch.tensor(Y_test.values).float()

test_ds = TensorDataset(x_test_tensor, y_test_tensor.unsqueeze(1))
test_dl = DataLoader(test_ds, batch_size=64)

model = torch.jit.load('model_scripted.pt')
model.eval()

y_pred_list = []
model.eval()

with torch.no_grad():
    for xb_test,yb_test  in test_dl:
        y_test_pred = model(xb_test)
        y_pred_tag = torch.round(y_test_pred)
        y_pred_list.append(y_pred_tag.detach().numpy())
    
y_pred_list = [a.squeeze().tolist() for a in y_pred_list]

ytest_pred = list(itertools.chain.from_iterable(y_pred_list))

y_true_test = Y_test.values.ravel()
conf_matrix = confusion_matrix(y_true_test ,ytest_pred)
print("Confusion Matrix of the Test Set")
print("-----------")
print(conf_matrix)
print("Precision of the MLP :\t"+str(precision_score(y_true_test,ytest_pred)))
print("Accuracy of the MLP :\t"+str(accuracy_score(y_true_test,ytest_pred)))
print("Recall of the MLP    :\t"+str(recall_score(y_true_test,ytest_pred)))
print("F1 Score of the Model :\t"+str(f1_score(y_true_test,ytest_pred)))


print(*map(int, ytest_pred))
print(*y_true_test)

with open('predicted_values.txt', 'w') as f:
    f.write(" ".join(map(str, map(int, ytest_pred))))