models

import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm
from utils import *


class Network1(nn.Module):
    def __init__(self):
        super().__init__()

        # input 28 x 28 x 1
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3)
        # self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3)
        self.fc1 = nn.Linear(in_features=5 * 5 * 64, out_features=50)
        self.out = nn.Linear(in_features=50, out_features=10)

    def forward(self, t):
        x = t
        # conv1 layer
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)  # output ->  28 | 26 | 13

        # conv2 layer
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)  # output ->  13 | 11 | 5

        x = x.reshape(-1, 5 * 5 * 64)

        # fc1 layer
        x = self.fc1(x)
        x = F.relu(x)  # output ->  5 x 5 x 64

        # out layer
        x = self.out(x)
        x = F.softmax(x, dim=1)  # output ->  1 x 10

        return x


class Network2(nn.Module):
    def __init__(self):
        super().__init__()

        # input 28 # output 24 # receptive_field = 5
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
        # input 24 # output 20 # receptive_field = 9
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
        # input 12x20x20, output 120
        # input 10*512
        self.fc1 = nn.Linear(in_features=12 * 20 * 20, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        self.out = nn.Linear(in_features=60, out_features=10)

    def forward(self, t):
        return t


class Network3(nn.Module):
    def __init__(self):
        super(Network3, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, padding=1)  # input -? Output? RF
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv4 = nn.Conv2d(128, 256, 3, padding=1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.conv5 = nn.Conv2d(256, 512, 3)
        self.conv6 = nn.Conv2d(512, 1024, 3)
        self.conv7 = nn.Conv2d(1024, 10, 3)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv2(F.relu(self.conv1(x)))))
        x = self.pool2(F.relu(self.conv4(F.relu(self.conv3(x)))))
        x = F.relu(self.conv6(F.relu(self.conv5(x))))
        x = F.relu(self.conv7(x))
        x = x.view(-1, 10)
        return F.log_softmax(x)


def train(model, device, train_loader, optimizer, criterion):

    model.train()
    pbar = tqdm(train_loader)

    # Data to plot accuracy and loss graphs
    train_losses = []
    train_acc = []
    train_loss = 0
    correct = 0
    processed = 0

    for batch_idx, (data, target) in enumerate(pbar):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()

        # Predict
        pred = model(data)

        # Calculate loss
        loss = criterion(pred, target)
        train_loss += loss.item()

        # Backpropagation
        loss.backward()
        optimizer.step()

        correct += get_num_correct(pred, target)
        processed += len(data)

        pbar.set_description(
            desc=f'Train: Loss={loss.item(): 0.4f} Batch_id={batch_idx} Accuracy={100 * correct / processed:0.2f}')

    train_acc.append(100 * correct / processed)
    train_losses.append(train_loss / len(train_loader))

def test(model, device, test_loader, criterion):
    model.eval()

    test_loss = 0
    correct = 0
    test_losses = []
    test_acc = []

    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)

            output = model(data)
            test_loss += criterion(output, target).item()  # sum up batch loss

            correct += get_num_correct(output, target)

    test_loss /= len(test_loader.dataset)
    test_acc.append(100. * correct / len(test_loader.dataset))
    test_losses.append(test_loss)

    print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))