import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torch.autograd import Variable
from torchvision import datasets, transforms
import os
import shutil
import random
import numpy as np
import matplotlib.pyplot as plt

'''
Implementation of the logistic regression classifier which solves the problem of recognition of cats and dogs.

The dataset can be found here:
    https://www.kaggle.com/c/dogs-vs-cats
'''

'''
Splits train and test data into to the following folder structure:
---/train
------/cat
------/dog
---/test
------/cat
------/dog

Randomly selects 5000 images as test examples (2500 for each class).
'''


def split_dogs_and_cats(source):
    cats_test = './test/cat'
    dogs_test = './test/dog'
    cats_train = './train/cat'
    dogs_train = './train/dog'

    files = os.listdir(source)
    os.makedirs(cats_test)
    os.makedirs(dogs_test)
    os.makedirs(cats_train)
    os.makedirs(dogs_train)

    cats_test_index = [i for i in range(12500)]
    dogs_test_index = [i for i in range(12500)]
    random.shuffle(cats_test_index)
    random.shuffle(dogs_test_index)
    cats_test_index = cats_test_index[:2500]
    dogs_test_index = dogs_test_index[:2500]

    for file in files:
        srcname = os.path.join(source, file)
        tag, number, _ = file.split('.')
        number = int(number)
        if tag == 'cat':
            if number in cats_test_index:
                dst = cats_test
            else:
                dst = cats_train
        else:
            if number in dogs_test_index:
                dst = dogs_test
            else:
                dst = dogs_train

        dstname = os.path.join(dst, file)
        shutil.move(srcname, dstname)


split_dogs_and_cats('./dataset')

batch_size = 32
image_size = 128


#Normalize the data.


transformation = transforms.Compose([
    transforms.Resize((image_size, image_size)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])


#Loads the data.


train_data = datasets.ImageFolder(root='./train', transform=transformation)
test_data = datasets.ImageFolder(root='./test', transform=transformation)

torchvision.utils.make_grid(train_data[0][0])

train_data_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=4)
test_data_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=True, num_workers=4)


#Let's check that we loaded the data correctly.



def imshow(inp, title=None):
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)


#1 - dog, 0 - cat.


inputs, classes = next(iter(train_data_loader))

out = torchvision.utils.make_grid(inputs)

imshow(out, title=classes)

'''
Class that represents logistic regression.
'''


class NeuralNetwork(nn.Module):

    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.layer = nn.Linear(3 * image_size * image_size, 1)

    def forward(self, x):
        return F.sigmoid(self.layer(x)).squeeze()


network = NeuralNetwork()
criterion = nn.BCELoss()

'''
Main procedure.
Firstly, we train our classifier then test.
'''


def run(learning_rate):
    optimizer = optim.Adam(network.parameters(), lr=learning_rate)
    for epoch in range(1, 3):
        for batch_idx, (data, target) in enumerate(train_data_loader):
            data, target = Variable(data), Variable(target)
            data = data.view(-1, image_size * image_size * 3)
            output = network(data)
            cost = criterion(output, target.float())
            optimizer.zero_grad()
            cost.backward()
            optimizer.step()

            if batch_idx % 10 == 0:
                print('Train Epoch: {} [{}/{} ({:.0f}%)]tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data), len(train_data_loader.dataset),
                           100. * batch_idx / len(train_data_loader), cost.item()))

    '''
    Test set evaluation.
    '''

    network.eval()
    test_loss = 0
    correct = 0
    for data, target in test_data_loader:
        data, target = Variable(data), Variable(target)
        data = data.view(-1, image_size * image_size * 3)
        output = network(data)
        test_loss += criterion(output, target.float()).item()
        pred = output.ge(0.5)
        correct += torch.eq(pred, target.byte()).sum()

    test_loss /= len(test_data_loader.dataset)

    print('nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)n'.format(
        test_loss, correct, len(test_data_loader.dataset),
        100. * correct / len(test_data_loader.dataset)))

    '''
    Train set evaluation.
    '''

    train_loss = 0
    correct = 0
    for data, target in train_data_loader:
        data, target = Variable(data), Variable(target)
        data = data.view(-1, image_size * image_size * 3)
        output = network(data)
        train_loss += criterion(output, target.float()).item()
        pred = output.ge(0.5)
        correct += torch.eq(pred, target.byte()).sum()

    train_loss /= len(train_data_loader.dataset)

    print('nTrain set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)n'.format(
        train_loss, correct, len(train_data_loader.dataset),
        100. * correct / len(train_data_loader.dataset)))


run(0.0001)