import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from matplotlib.pyplot import imshow
import numpy as np
from matplotlib.pyplot import imshow
import matplotlib.pyplot as ppl

def plotdigit(image):
    img = np.reshape(image, (-250, 250))
    imshow(img, cmap='Greys', vmin=0, vmax=255)
    ppl.show()

train_images = np.load('neural_network\\image-dataset.npy')
print(train_images.shape)

train_labels = np.load('neural_network\\label-dataset.npy')

train_images = train_images / 255
train_labels = train_labels / 255

train_images = [torch.tensor(image, dtype=torch.float32) for image in train_images]
print(train_images[0].shape)

train_labels = [torch.tensor(label, dtype=torch.long) for label in train_labels]


input_dim = 100*100*3
output_dim = 2

model = nn.Sequential(
    nn.Linear(input_dim, output_dim),
    nn.LogSoftmax()
)

def train(model, n_iter):

    criterion = nn.NLLLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    for epoch in range(n_iter):
        for image, label in zip(train_images, train_labels):
            optimizer.zero_grad()

            output = model(image)
            loss = criterion(output.unsqueeze(0), label.unsqueeze(0))
            loss.backward()
            optimizer.step()

        print(f'epoch: {epoch:03}')


train(model, 100)