import pandas as pd
import torch
from torch import nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch import optim
from sklearn import metrics




class DatasetIris(Dataset):
	def __init__(self, path):
                df = pd.read_csv(path)
                self.data = torch.tensor(df.to_numpy()[:, :4].astype(float),dtype=torch.float32)
                tmp = torch.from_numpy(df["variety"].replace({"Virginica": 0, "Versicolor": 1, "Setosa": 2}).to_numpy(dtype=int))
                self.target = nn.functional.one_hot(tmp, num_classes=3).to(dtype=torch.float32) # one hot encoding (i maps to [0,..,1,...0]) 
	def __len__(self):
		return len(self.data)
	def __getitem__(self, idx):
		return self.data[idx], self.target[idx]

def predict(model, X):
    model.eval()
    with torch.no_grad():
        l = model(X)
        p = torch.softmax(l, dim=1)
        predictions = torch.argmax(p, dim=1)
    return predictions

# Construction du jeu de données

training_data = DatasetIris("iris_train.csv") 
train_dataloader = DataLoader(training_data, batch_size=16, shuffle=True)

test_data = DatasetIris("iris_test.csv") 
test_dataloader = DataLoader(training_data, batch_size=16, shuffle=True)


# Construction du modèle

m = nn.Sequential(nn.Linear(4,8),
                  nn.ReLU(),
                  nn.Linear(8,3),
                  )

# Pas de softmax, la fonction d'erreur le caclulera pour nous. 

o = optim.Adam(m.parameters(), lr=0.03)
loss_fn = nn.CrossEntropyLoss()

for e in range(2000):
        for X, y in train_dataloader:
                o.zero_grad()
                loss = loss_fn(m(X), y)
                loss.backward()
                o.step()

target_true = torch.argmax(test_data.target, dim=1)
target_predicted = predict(m,test_data.data)
print(target_predicted)
print(target_true)
accuracy = metrics.accuracy_score(target_true, target_predicted)
print(f"Test accuracy {accuracy}")

target_true = torch.argmax(training_data.target, dim=1)
target_predicted = predict(m,training_data.data)
accuracy = metrics.accuracy_score(target_true, target_predicted)
print(f"Train accuracy {accuracy}")