# -*- coding: utf-8 -*-
"""328lab10.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1BAGGzKp3Q2Ay2nO-wKy0sgTBDZ8l6Ori
"""

import os
import time

import torch.nn.functional as F
import torch
from torch import nn
from torchvision import models
import numpy as np

from google.colab import drive
# drive.mount('/content/drive')

root_dir = '/content/drive/My Drive/Colab Notebooks/A10'


class ImagesData(object):
    def __init__(self, subset='train', batch_size=64, shuffle=True):
        if subset == 'train':
            images = np.load(os.path.join(root_dir, 'TextureImagesDataset',
                                          'train_images.npy'))
            masks = np.load(os.path.join(root_dir, 'TextureImagesDataset',
                                         'train_masks.npy'))
        elif subset == 'test':
            images = np.load(os.path.join(root_dir, 'TextureImagesDataset',
                                          'test_images.npy'))
            masks = np.load(os.path.join(root_dir, 'TextureImagesDataset',
                                         'test_masks.npy'))
        else:
            raise NotImplementedError
        self._images = images
        self.images = self._images
        self._masks = masks
        self.masks = self._masks
        self.batch_size = batch_size
        self.num_samples = len(self.images)
        self.shuffle = shuffle
        if self.shuffle:
            self.shuffle_samples()
        self.next_batch_pointer = 0

    def shuffle_samples(self):
        image_indices = np.random.permutation(np.arange(self.num_samples))
        self.images = self._images[image_indices]
        self.masks = self._masks[image_indices]

    def get_next_batch(self):
        num_samples_left = self.num_samples - self.next_batch_pointer
        if num_samples_left >= self.batch_size:
            x_batch = self.images[self.next_batch_pointer:self.next_batch_pointer + self.batch_size]
            y_batch = self.masks[self.next_batch_pointer:self.next_batch_pointer + self.batch_size]
            self.next_batch_pointer += self.batch_size
        else:
            x_partial_batch_1 = self.images[self.next_batch_pointer:self.num_samples]
            y_partial_batch_1 = self.masks[self.next_batch_pointer:self.num_samples]
            if self.shuffle:
                self.shuffle_samples()
            x_partial_batch_2 = self.images[0:self.batch_size - num_samples_left]
            y_partial_batch_2 = self.masks[0:self.batch_size - num_samples_left]
            x_batch = np.vstack((x_partial_batch_1, x_partial_batch_2))
            y_batch = np.vstack((y_partial_batch_1, y_partial_batch_2))
            self.next_batch_pointer = self.batch_size - num_samples_left
        return x_batch, y_batch

class ClassificationNet(nn.Module):
    def __init__(self, num_classes=10, num_digits=2):
        super(ClassificationNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
        self.fc1 = nn.Linear(64*4*4, 1024)
        self.dropout = nn.Dropout(0.1)
        self.fc2 = nn.Linear(1024, num_classes*num_digits)

    def forward(self, out):
        # print(out.size())
        out = self.conv1(out)
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv3(out))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv4(out))
        out = F.max_pool2d(out, 2)
        # print(out.size())
        out = out.view(out.size()[0], -1)
        out = F.relu(self.dropout(self.fc1(out)))
        out = self.fc2(out)
        out = torch.reshape(out, tuple(out.size()[:-1])+(2, 10))
        print(out)
        return out


class LocalizationNet(nn.Module):
    def __init__(self, num_bbox=2, num_digits=2):
        super(LocalizationNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
        self.fc1 = nn.Linear(64*4*4, 1024)
        self.dropout = nn.Dropout(0.1)
        self.fc2 = nn.Linear(1024, num_bbox*num_digits)

    def forward(self, out):
        out = F.relu(self.conv1(out))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv3(out))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv4(out))
        out = F.max_pool2d(out, 2)
        # print(out.size())
        out = out.view(out.size()[0], -1)
        out = F.relu(self.dropout(self.fc1(out)))
        out = self.fc2(out)
        out = torch.reshape(out, tuple(out.size()[:-1])+(2, 2))
        return out


#!pip install torchsummary

from torchsummary import summary
summary(LocalizationNet().to('cuda'), input_size=(1, 64, 64))
