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


class LogisticRegression(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LogisticRegression, self).__init__()
        self.linear = torch.nn.Linear(input_dim, output_dim)

    def forward(self, x):
        outputs = self.linear(x)
        return outputs


def test_model(model, test_loader):
    total_correct = 0
    pred_y = []
    target_y = []
    for batch_x, batch_y in test_loader:
        logits = model(batch_x)
        pred = torch.argmax(logits, dim=1)

        correct = torch.sum(pred == batch_y)
        total_correct += correct

        pred_y.append(pred)
        target_y.append(batch_y)
    pred_y = torch.concat(pred_y, dim=0)
    target_y = torch.concat(target_y, dim=0)
    rst = torch.sum(pred_y == target_y) / len(pred_y)
    #     print(pred_y == target_y)
    #     print(target_y)
    return rst


batch_size = 16
learning_rate = 0.001
epochs = 1000

from sklearn import datasets

iris = datasets.load_iris()
X = iris.data
Y = iris.target

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2, random_state=0)

X_train, y_train = torch.Tensor(X_train), torch.LongTensor(y_train)
train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

X_test, y_test = torch.Tensor(X_test), torch.LongTensor(y_test)
test_dataset = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

model = LogisticRegression(X_train.shape[1], len(set(y_train.numpy())))
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = torch.nn.CrossEntropyLoss()

for i in range(epochs):
    for batch_x, batch_y in train_loader:
        logits = model(batch_x)
        loss = criterion(logits, batch_y)

        model.zero_grad()
        loss.backward()
        optimizer.step()
    if i % 100 == 0:
        rst = test_model(model, test_loader)
        print(f'logistic regression: acc = {rst.item():.4f}')
rst = test_model(model, test_loader)
print(f'logistic regression: acc = {rst.item():.4f}')