import torch
import torch.nn.functional as F
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
from torch import optim
import numpy as np
import os
import pickle


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 32, kernel_size=6, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(32, 32, kernel_size=6, stride=1, padding=1)
        self.conv3 = torch.nn.Conv2d(32, 10, kernel_size=6, stride=1, padding=1)

    def forward(self, x):
        batch_size = x.shape[0]

        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.adaptive_avg_pool2d(x, 1)
        return x.view(batch_size, 10)


def load_model():
    directory = os.path.abspath(os.path.dirname(__file__))

    model = Model()
    model.load_state_dict(torch.load(directory + '/weights.pth', map_location='cpu'))

    return model


def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)


def get_validation_loss(model, loss_func, valid_dl):
    model.eval()
    with torch.no_grad():
        losses, nums = zip(
            *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]
        )
    return np.sum(np.multiply(losses, nums)) / np.sum(nums)


def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    best_loss = get_validation_loss(model, loss_func, valid_dl)

    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)

        val_loss = get_validation_loss(model, loss_func, valid_dl)
        if val_loss < best_loss:
            best_loss = val_loss
            torch.save(model.state_dict(), 'weights.pth')

        print(epoch, val_loss)


def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )


if __name__ == "__main__":
    batch_size = 50

    with open('trn.pkl', 'rb') as f:
        trn = pickle.load(f)
    with open('val.pkl', 'rb') as f:
        val = pickle.load(f)

    x_train = trn['data'].transpose((0, 3, 1, 2)).astype('f4')
    y_train = trn['labels'].astype('i8')

    x_valid = val['data'].transpose((0, 3, 1, 2)).astype('f4')
    y_valid = val['labels'].astype('i8')

    x_train, y_train, x_valid, y_valid = map(torch.tensor, (x_train, y_train, x_valid, y_valid))

    train_ds = TensorDataset(x_train, y_train)
    valid_ds = TensorDataset(x_valid, y_valid)
    train_dl, valid_dl = get_data(train_ds, valid_ds, batch_size)

    model = Model()
    model.load_state_dict(torch.load('weights.pth', map_location='cpu'))
    loss_func = F.cross_entropy
    epochs = 10

    opt = optim.Adamax(model.parameters(), lr=0.0001)
    fit(epochs, model, loss_func, opt, train_dl, valid_dl)

import torch
import torch.nn.functional as F
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
from torch import optim
import numpy as np
import os
import pickle


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 32, kernel_size=6, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(32, 32, kernel_size=6, stride=1, padding=1)
        self.conv3 = torch.nn.Conv2d(32, 10, kernel_size=6, stride=1, padding=1)

    def forward(self, x):
        batch_size = x.shape[0]

        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = F.adaptive_avg_pool2d(x, 1)
        return x.view(batch_size, 10)


def load_model():
    directory = os.path.abspath(os.path.dirname(__file__))

    model = Model()
    model.load_state_dict(torch.load(directory + '/weights.pth', map_location='cpu'))

    return model


def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)


def get_validation_loss(model, loss_func, valid_dl):
    model.eval()
    with torch.no_grad():
        losses, nums = zip(
            *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]
        )
    return np.sum(np.multiply(losses, nums)) / np.sum(nums)


def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    best_loss = get_validation_loss(model, loss_func, valid_dl)

    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)

        val_loss = get_validation_loss(model, loss_func, valid_dl)
        if val_loss < best_loss:
            best_loss = val_loss
            torch.save(model.state_dict(), 'weights.pth')

        print(epoch, val_loss)


def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )


if __name__ == "__main__":
    batch_size = 50

    with open('trn.pkl', 'rb') as f:
        trn = pickle.load(f)
    with open('val.pkl', 'rb') as f:
        val = pickle.load(f)

    x_train = trn['data'].transpose((0, 3, 1, 2)).astype('f4')
    y_train = trn['labels'].astype('i8')

    x_valid = val['data'].transpose((0, 3, 1, 2)).astype('f4')
    y_valid = val['labels'].astype('i8')

    x_train, y_train, x_valid, y_valid = map(torch.tensor, (x_train, y_train, x_valid, y_valid))

    train_ds = TensorDataset(x_train, y_train)
    valid_ds = TensorDataset(x_valid, y_valid)
    train_dl, valid_dl = get_data(train_ds, valid_ds, batch_size)

    model = Model()
    model.load_state_dict(torch.load('weights.pth', map_location='cpu'))
    loss_func = F.cross_entropy
    epochs = 10

    opt = optim.Adamax(model.parameters(), lr=0.0001)
    fit(epochs, model, loss_func, opt, train_dl, valid_dl)
