#!/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


[((1581, 1636), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(1)', 'out_channels': '(4)', 'kernel_size': '(3)'}), '(in_channels=1, out_channels=4, kernel_size=3)\n', (1590, 1636), True, 'import torch.nn as nn\n'), ((1724, 1741), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(4)'], {}), '(4)\n', (1738, 1741), True, 'import torch.nn as nn\n'), ((1812, 1856), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv1.weight'], {}), '(self.conv1.weight)\n', (1837, 1856), True, 'import torch.nn.init as torch_init\n'), ((1878, 1915), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(3)', 'stride': '(1)'}), '(kernel_size=3, stride=1)\n', (1890, 1915), True, 'import torch.nn as nn\n'), ((2185, 2240), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(4)', 'out_channels': '(8)', 'kernel_size': '(3)'}), '(in_channels=4, out_channels=8, kernel_size=3)\n', (2194, 2240), True, 'import torch.nn as nn\n'), ((2269, 2286), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(8)'], {}), '(8)\n', (2283, 2286), True, 'import torch.nn as nn\n'), ((2295, 2339), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv2.weight'], {}), '(self.conv2.weight)\n', (2320, 2339), True, 'import torch.nn.init as torch_init\n'), ((2378, 2415), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(3)', 'stride': '(1)'}), '(kernel_size=3, stride=1)\n', (2390, 2415), True, 'import torch.nn as nn\n'), ((2536, 2592), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(8)', 'out_channels': '(16)', 'kernel_size': '(3)'}), '(in_channels=8, out_channels=16, kernel_size=3)\n', (2545, 2592), True, 'import torch.nn as nn\n'), ((2621, 2639), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(16)'], {}), '(16)\n', (2635, 2639), True, 'import torch.nn as nn\n'), ((2648, 2692), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv3.weight'], {}), '(self.conv3.weight)\n', (2673, 2692), True, 'import torch.nn.init as torch_init\n'), ((2811, 2868), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(16)', 'out_channels': '(16)', 'kernel_size': '(3)'}), '(in_channels=16, out_channels=16, kernel_size=3)\n', (2820, 2868), True, 'import torch.nn as nn\n'), ((2897, 2915), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(16)'], {}), '(16)\n', (2911, 2915), True, 'import torch.nn as nn\n'), ((2924, 2968), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv4.weight'], {}), '(self.conv4.weight)\n', (2949, 2968), True, 'import torch.nn.init as torch_init\n'), ((2990, 3027), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(3)', 'stride': '(1)'}), '(kernel_size=3, stride=1)\n', (3002, 3027), True, 'import torch.nn as nn\n'), ((3145, 3201), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(16)', 'out_channels': '(8)', 'kernel_size': '(3)'}), '(in_channels=16, out_channels=8, kernel_size=3)\n', (3154, 3201), True, 'import torch.nn as nn\n'), ((3230, 3247), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(8)'], {}), '(8)\n', (3244, 3247), True, 'import torch.nn as nn\n'), ((3256, 3300), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv5.weight'], {}), '(self.conv5.weight)\n', (3281, 3300), True, 'import torch.nn.init as torch_init\n'), ((3323, 3378), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(8)', 'out_channels': '(8)', 'kernel_size': '(3)'}), '(in_channels=8, out_channels=8, kernel_size=3)\n', (3332, 3378), True, 'import torch.nn as nn\n'), ((3407, 3424), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(8)'], {}), '(8)\n', (3421, 3424), True, 'import torch.nn as nn\n'), ((3433, 3477), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv6.weight'], {}), '(self.conv6.weight)\n', (3458, 3477), True, 'import torch.nn.init as torch_init\n'), ((3499, 3536), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(3)', 'stride': '(1)'}), '(kernel_size=3, stride=1)\n', (3511, 3536), True, 'import torch.nn as nn\n'), ((3647, 3702), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(8)', 'out_channels': '(8)', 'kernel_size': '(3)'}), '(in_channels=8, out_channels=8, kernel_size=3)\n', (3656, 3702), True, 'import torch.nn as nn\n'), ((3731, 3748), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(8)'], {}), '(8)\n', (3745, 3748), True, 'import torch.nn as nn\n'), ((3757, 3801), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv7.weight'], {}), '(self.conv7.weight)\n', (3782, 3801), True, 'import torch.nn.init as torch_init\n'), ((3824, 3879), 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': '(8)', 'out_channels': '(8)', 'kernel_size': '(3)'}), '(in_channels=8, out_channels=8, kernel_size=3)\n', (3833, 3879), True, 'import torch.nn as nn\n'), ((3908, 3925), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(8)'], {}), '(8)\n', (3922, 3925), True, 'import torch.nn as nn\n'), ((3934, 3978), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.conv8.weight'], {}), '(self.conv8.weight)\n', (3959, 3978), True, 'import torch.nn.init as torch_init\n'), ((4000, 4037), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(4)', 'stride': '(4)'}), '(kernel_size=4, stride=4)\n', (4012, 4037), True, 'import torch.nn as nn\n'), ((4121, 4175), 'torch.nn.Linear', 'nn.Linear', ([], {'in_features': '(122 * 122 * 8)', 'out_features': '(512)'}), '(in_features=122 * 122 * 8, out_features=512)\n', (4130, 4175), True, 'import torch.nn as nn\n'), ((4199, 4218), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['(512)'], {}), '(512)\n', (4213, 4218), True, 'import torch.nn as nn\n'), ((4227, 4269), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.fc1.weight'], {}), '(self.fc1.weight)\n', (4252, 4269), True, 'import torch.nn.init as torch_init\n'), ((4309, 4353), 'torch.nn.Linear', 'nn.Linear', ([], {'in_features': '(512)', 'out_features': '(128)'}), '(in_features=512, out_features=128)\n', (4318, 4353), True, 'import torch.nn as nn\n'), ((4380, 4399), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', (['(128)'], {}), '(128)\n', (4394, 4399), True, 'import torch.nn as nn\n'), ((4408, 4450), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.fc2.weight'], {}), '(self.fc2.weight)\n', (4433, 4450), True, 'import torch.nn.init as torch_init\n'), ((4490, 4533), 'torch.nn.Linear', 'nn.Linear', ([], {'in_features': '(128)', 'out_features': '(14)'}), '(in_features=128, out_features=14)\n', (4499, 4533), True, 'import torch.nn as nn\n'), ((4542, 4584), 'torch.nn.init.xavier_normal_', 'torch_init.xavier_normal_', (['self.fc3.weight'], {}), '(self.fc3.weight)\n', (4567, 4584), True, 'import torch.nn.init as torch_init\n')]