import torch
from sklearn.datasets import load_iris
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import torch.nn.functional as fnn
import torch.optim as optim

class IrisDataset(Dataset):
    def __init__(self):
        data = load_iris()
        self.data = torch.tensor(data.data, dtype=torch.float32)
        self.targets = torch.tensor(data.target, dtype=torch.long)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.targets[idx]

dataset = IrisDataset()
dataloader = DataLoader(dataset, batch_size=16, shuffle=True)

class INet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 16)
        self.fc4 = nn.Linear(16, 8)
        self.fc5 = nn.Linear(8, 4)

    def forward(self, x):
        x = self.fc1(x)
        x = fnn.relu(x)
        x = self.fc2(x)
        x = fnn.relu(x)
        x = self.fc3(x)
        x = fnn.relu(x)
        x = self.fc4(x)
        x = fnn.relu(x)
        x = self.fc5(x)
        return x

model = INet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

for epoch in range(30):
    for X, y in dataloader:
        preds = model(X)
        loss = criterion(preds, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch}, Loss: {loss.item()}")

torch.save(model.state_dict(), "/models/inet.pth")