#!/usr/bin/env python
# coding: utf-8

# In[ ]:


################################################################################
# CSE 253: Programming Assignment 3
# Winter 2019
# Code author: Jenny Hamer (+ modifications by Tejash Desai)
#
# Filename: baseline_cnn.py
#
# Description:
#
# This file contains the starter code for the baseline architecture you will use
# to get a little practice with PyTorch and compare the results of with your
# improved architecture.
#
# Be sure to fill in the code in the areas marked #TODO.
################################################################################


# PyTorch and neural network imports
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as func
import torch.nn.init as torch_init
import torch.optim as optim

# Data utils and dataloader
import torchvision
from torchvision import transforms, utils
from xray_dataloader_zscored import ChestXrayDataset, create_split_loaders

import matplotlib.pyplot as plt
import numpy as np
import os



class Arch2CNN(nn.Module):
    """
<<<<<<< HEAD
    conv1 -> maxpool -> conv2 -> maxpool -> conv3 -> conv4 ->maxpool -> conv5 -> conv6 -> maxpool -> conv7 -> conv8 -> maxpool -> fc1 -> fc2 -> fc3 (outputs)
=======
    conv1 -> conv2 -> maxpool -> conv3 -> conv4 -> conv5 -> maxpool -> fc1 -> fc2 -> fc3 (outputs)
>>>>>>> 6652e3cfb72835ac4a7c802c9a703b59d5f63ae6
    """

    def __init__(self):
        super(Arch2CNN, self).__init__()

        # conv1: 1 input channel, 4 output channels, [3x3] kernel size
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=4, kernel_size=3)

        # Add batch-normalization to the outputs of conv1
        self.conv1_normed = nn.BatchNorm2d(4)

        # Initialized weights using the Xavier-Normal method
        torch_init.xavier_normal_(self.conv1.weight)
        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=1)

        #TODO: Fill in the remaining initializations replacing each '_' with
        # the necessary value based on the provided specs for each layer

        #TODO: conv2: 4 input channels, 8 output channels, [3x3] kernel, initialization: xavier
        self.conv2 = nn.Conv2d(in_channels=4, out_channels=8, kernel_size=3)
        self.conv2_normed = nn.BatchNorm2d(8)
        torch_init.xavier_normal_(self.conv2.weight)
        #Maxpool
        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=1)


        #TODO: conv3: X input channels, 12 output channels, [8x8] kernel, initialization: xavier
        self.conv3 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3)
        self.conv3_normed = nn.BatchNorm2d(16)
        torch_init.xavier_normal_(self.conv3.weight)
        #TODO: conv4: X input channels, 10 output channels, [6x6] kernel, initialization: xavier
        self.conv4 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3)
        self.conv4_normed = nn.BatchNorm2d(16)
        torch_init.xavier_normal_(self.conv4.weight)
        self.pool3 = nn.MaxPool2d(kernel_size=3, stride=1)
        #TODO: conv5: X input channels, 8 output channels, [5x5] kernel, initialization: xavier
        self.conv5 = nn.Conv2d(in_channels=16, out_channels=8, kernel_size=3)
        self.conv5_normed = nn.BatchNorm2d(8)
        torch_init.xavier_normal_(self.conv5.weight)

        self.conv6 = nn.Conv2d(in_channels=8, out_channels=8, kernel_size=3)
        self.conv6_normed = nn.BatchNorm2d(8)
        torch_init.xavier_normal_(self.conv6.weight)
        self.pool4 = nn.MaxPool2d(kernel_size=3, stride=1)

        #TODO: Apply max-pooling with a [3x3] kernel using tiling (*NO SLIDING WINDOW*)
        self.conv7 = nn.Conv2d(in_channels=8, out_channels=8, kernel_size=3)
        self.conv7_normed = nn.BatchNorm2d(8)
        torch_init.xavier_normal_(self.conv7.weight)

        self.conv8 = nn.Conv2d(in_channels=8, out_channels=8, kernel_size=3)
        self.conv8_normed = nn.BatchNorm2d(8)
        torch_init.xavier_normal_(self.conv8.weight)
        self.pool5 = nn.MaxPool2d(kernel_size=4, stride=4)


        # Define 2 fully connected layers:
        #TODO: fc1
        self.fc1 = nn.Linear(in_features=122*122*8, out_features=512)

        self.fc1_normed = nn.BatchNorm1d(512)
        torch_init.xavier_normal_(self.fc1.weight)

        #TODO: fc2
        self.fc2 = nn.Linear(in_features=512, out_features=128)
        self.fc2_normed = nn.BatchNorm1d(128)
        torch_init.xavier_normal_(self.fc2.weight)

        #TODO: fc3
        self.fc3 = nn.Linear(in_features=128, out_features=14)
        torch_init.xavier_normal_(self.fc3.weight)

        #TODO: Output layer: what should out_features be?
        self.out_features = 14


    def forward(self, batch):
        """Pass the batch of images through each layer of the network, applying
        non-linearities after each layer.

        Note that this function *needs* to be called "forward" for PyTorch to
        automagically perform the forward pass.

        Params:
        -------
        - batch: (Tensor) An input batch of images

        Returns:
        --------
        - logits: (Variable) The output of the network
        """

        # Apply first convolution, followed by ReLU non-linearity;
        # use batch-normalization on its outputs
        batch = func.rrelu(self.conv1_normed(self.conv1(batch)))

        batch = self.pool1(batch)
        # Apply conv2 and conv3 similarly
        batch = func.rrelu(self.conv2_normed(self.conv2(batch)))
        batch = self.pool2(batch)

        batch = func.rrelu(self.conv3_normed(self.conv3(batch)))
        batch = func.rrelu(self.conv4_normed(self.conv4(batch)))
        batch = self.pool3(batch)
        batch = func.rrelu(self.conv5_normed(self.conv5(batch)))
        batch = func.rrelu(self.conv6_normed(self.conv6(batch)))
        # Pass the output of conv3 to the pooling layer
        batch = self.pool4(batch)
        batch = func.rrelu(self.conv7_normed(self.conv7(batch)))
        batch = func.rrelu(self.conv8_normed(self.conv8(batch)))
        # Pass the output of conv3 to the pooling layer
        batch = self.pool5(batch)

        # Reshape the output of the conv3 to pass to fully-connected layer
        batch = batch.view(-1, self.num_flat_features(batch))

        # Connect the reshaped features of the pooled conv3 to fc1
        batch =  func.rrelu(self.fc1_normed(self.fc1(batch)))
        batch =  func.rrelu(self.fc2_normed(self.fc2(batch)))

        # Connect fc1 to fc2 - this layer is slightly different than the rest (why?)
        batch = self.fc3(batch)


        # Return the class predictions
        #TODO: apply an activition function to 'batch'
        #batch = func.sigmoid(batch)
        return batch

    def num_flat_features(self, inputs):

        # Get the dimensions of the layers excluding the inputs
        size = inputs.size()[1:]
        # Track the number of features
        num_features = 1

        for s in size:
            num_features *= s

        return num_features


[((56, 21, 56, 76), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((59, 28, 59, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(59, 43, 59, 44): '4'}, {}), '(4)', True, 'import torch.nn as nn\n'), ((62, 8, 62, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(62, 34, 62, 51): 'self.conv1.weight'}, {}), '(self.conv1.weight)', True, 'import torch.nn.init as torch_init\n'), ((63, 21, 63, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((69, 21, 69, 76), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((70, 28, 70, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(70, 43, 70, 44): '8'}, {}), '(8)', True, 'import torch.nn as nn\n'), ((71, 8, 71, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(71, 34, 71, 51): 'self.conv2.weight'}, {}), '(self.conv2.weight)', True, 'import torch.nn.init as torch_init\n'), ((73, 21, 73, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((77, 21, 77, 77), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((78, 28, 78, 46), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(78, 43, 78, 45): '16'}, {}), '(16)', True, 'import torch.nn as nn\n'), ((79, 8, 79, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(79, 34, 79, 51): 'self.conv3.weight'}, {}), '(self.conv3.weight)', True, 'import torch.nn.init as torch_init\n'), ((81, 21, 81, 78), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((82, 28, 82, 46), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(82, 43, 82, 45): '16'}, {}), '(16)', True, 'import torch.nn as nn\n'), ((83, 8, 83, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(83, 34, 83, 51): 'self.conv4.weight'}, {}), '(self.conv4.weight)', True, 'import torch.nn.init as torch_init\n'), ((84, 21, 84, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((86, 21, 86, 77), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((87, 28, 87, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(87, 43, 87, 44): '8'}, {}), '(8)', True, 'import torch.nn as nn\n'), ((88, 8, 88, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(88, 34, 88, 51): 'self.conv5.weight'}, {}), '(self.conv5.weight)', True, 'import torch.nn.init as torch_init\n'), ((90, 21, 90, 76), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((91, 28, 91, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(91, 43, 91, 44): '8'}, {}), '(8)', True, 'import torch.nn as nn\n'), ((92, 8, 92, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(92, 34, 92, 51): 'self.conv6.weight'}, {}), '(self.conv6.weight)', True, 'import torch.nn.init as torch_init\n'), ((93, 21, 93, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((96, 21, 96, 76), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((97, 28, 97, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(97, 43, 97, 44): '8'}, {}), '(8)', True, 'import torch.nn as nn\n'), ((98, 8, 98, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(98, 34, 98, 51): 'self.conv7.weight'}, {}), '(self.conv7.weight)', True, 'import torch.nn.init as torch_init\n'), ((100, 21, 100, 76), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((101, 28, 101, 45), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(101, 43, 101, 44): '8'}, {}), '(8)', True, 'import torch.nn as nn\n'), ((102, 8, 102, 52), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(102, 34, 102, 51): 'self.conv8.weight'}, {}), '(self.conv8.weight)', True, 'import torch.nn.init as torch_init\n'), ((103, 21, 103, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((108, 19, 108, 69), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((110, 26, 110, 45), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', ({(110, 41, 110, 44): '512'}, {}), '(512)', True, 'import torch.nn as nn\n'), ((111, 8, 111, 50), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(111, 34, 111, 49): 'self.fc1.weight'}, {}), '(self.fc1.weight)', True, 'import torch.nn.init as torch_init\n'), ((114, 19, 114, 63), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((115, 26, 115, 45), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', ({(115, 41, 115, 44): '128'}, {}), '(128)', True, 'import torch.nn as nn\n'), ((116, 8, 116, 50), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(116, 34, 116, 49): 'self.fc2.weight'}, {}), '(self.fc2.weight)', True, 'import torch.nn.init as torch_init\n'), ((119, 19, 119, 62), 'torch.nn.Linear', 'nn.Linear', (), '', True, 'import torch.nn as nn\n'), ((120, 8, 120, 50), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', ({(120, 34, 120, 49): 'self.fc3.weight'}, {}), '(self.fc3.weight)', True, 'import torch.nn.init as torch_init\n')]