"""
A collection of neural network code. The first part of the script includes
blocks, which are the building blocks of our models. The second part includes
the actual Pytorch models.
"""
import torch
import torchvision.transforms as transforms


class ConvBlock(torch.nn.Module):
    """
    A ConvBlock represents a convolution. It's not just a convolution however,
    as some common operations (dropout, activation, batchnorm, 2x2 pooling)
    can be set and run in the order mentioned.
    """

    def __init__(
        self,
        dim,
        n_out,
        kernel_size=3,
        stride=1,
        padding=1,
        batchnorm=False,
        dropout=0,
        activation=True,
    ):
        """ A convolution operation """
        super(ConvBlock, self).__init__()
        n_in = int(dim[0])
        self.conv2d = torch.nn.Conv2d(
            n_in, n_out, kernel_size=kernel_size, stride=stride, padding=padding
        )
        self.batchnorm = torch.nn.BatchNorm2d(n_out) if batchnorm else None
        self.activation = torch.nn.ReLU(inplace=True) if activation else None
        self.dropout = torch.nn.Dropout2d(dropout) if dropout else None
        dim[0] = n_out
        dim[1:] = 1 + (dim[1:] + padding * 2 - kernel_size) // stride
        self.n_params = n_out * (n_in * kernel_size * kernel_size + (3 if batchnorm else 1))
        print(
            "Conv2d in %4i out %4i h %4i w %4i k %i s %i params %9i"
            % (n_in, *dim, kernel_size, stride, self.n_params)
        )

    def forward(self, batch):
        """ Forward the 4D batch """
        out = self.conv2d(batch)
        if self.activation:
            out = self.activation(out)
        if self.batchnorm:
            out = self.batchnorm(out)
        if self.dropout:
            out = self.dropout(out)
        return out


class LinearBlock(torch.nn.Module):
    """
    A LinearBlock represents a fully connected layer. It's not just this, as
    some common operations (dropout, activation, batchnorm) can be set and run
    in the order mentioned.
    """

    def __init__(self, dim, n_out, batchnorm=False, dropout=0.0, activation=True):
        """ A fully connected operation """
        super(LinearBlock, self).__init__()
        n_in = int(dim[0])
        self.linear = torch.nn.Linear(n_in, n_out)
        dim[0] = n_out if type(n_out) in (int, float) else n_out[0]
        self.batchnorm = torch.nn.BatchNorm1d(dim[0]) if batchnorm else None
        self.dropout = torch.nn.Dropout(dropout) if dropout > 0.0 else None
        self.activation = torch.nn.ReLU(inplace=True) if activation else None
        self.n_params = n_out * (n_in + (3 if batchnorm else 1))
        print(
            "Linear in %4i out %4i                       params %9i" % (n_in, n_out, self.n_params)
        )

    def forward(self, batch):
        """ Forward the 2D batch """
        out = self.linear(batch)
        if self.activation:
            out = self.activation(out)
        if self.batchnorm:
            out = self.batchnorm(out)
        if self.dropout:
            out = self.dropout(out)
        return out


class PoolBlock(torch.nn.Module):
    """
    A PoolBlock is a pooling operation that happens on a matrix, often between
    convolutional layers, on each channel individually. By default only two are
    supported: max and avg.
    """

    def __init__(self, dim, pool="max", size=None, stride=None):
        """ A pooling operation """
        super(PoolBlock, self).__init__()

        stride = size if stride is None else stride
        if size:
            dim[1:] //= stride
        else:
            size = [int(x) for x in dim[1:]]
            dim[1:] = 1
        if pool == "max":
            self.pool = torch.nn.MaxPool2d(size, stride=stride, padding=0)
        elif pool == "avg":
            self.pool = torch.nn.AvgPool2d(size, stride=stride, padding=0)
        self.n_params = 0

    def forward(self, batch):
        """ Forward the 4D batch """
        out = self.pool(batch)
        return out


class ViewBlock(torch.nn.Module):
    """
    A ViewBlock restructures the shape of our activation maps so they're
    represented as 1D instead of 3D.
    """

    def __init__(self, dim, shape=-1):
        """ A reshape operation """
        super(ViewBlock, self).__init__()
        self.shape = shape
        if self.shape == -1:
            dim[0] = dim[0] * dim[1] * dim[2]
            dim[-2] = 0
            dim[-1] = 0
        else:
            dim[:] = shape

        self.n_params = 0
        print("View             d %4i h %4i w %4i" % (*dim,))

    def forward(self, batch):
        """ Forward the 4D batch into a 2D batch """
        return batch.view(batch.size(0), self.shape)


class Tiny(torch.nn.Module):
    """ A small and quick model """

    def __init__(self, in_dim, n_status, n_out):
        """
        Args:
            in_dim (list): The input size of each example
            n_status (int): Number of status inputs to add
            n_out (int): Number of values to predict
        """
        super(Tiny, self).__init__()
        self.n_status = n_status
        dim = in_dim.copy()
        self.feat = torch.nn.Sequential(
            ConvBlock(dim, 16),
            PoolBlock(dim, "max", 2),
            ConvBlock(dim, 32),
            PoolBlock(dim, "max", 2),
            ConvBlock(dim, 48),
            PoolBlock(dim, "max", 2),
            ConvBlock(dim, 64),
            PoolBlock(dim, "max", 2),
        )
        self.view = ViewBlock(dim)
        dim[0] += n_status
        self.head = torch.nn.Sequential(LinearBlock(dim, n_out, activation=False))
        self.n_params = sum([x.n_params for x in self.feat]) + sum([x.n_params for x in self.head])
        print("Tiny                                          params %9i" % self.n_params)

    def forward(self, batch, status):
        """
        Args:
            batch (4D tensor): A batch of camera input.
            status (1D tensor): Status inputs indicating things like speed.
        """
        out = self.feat(batch)
        out = self.view(out)
        if self.n_status:
            out = torch.cat((out, status), 1)
        out = self.head(out)
        return out


class StarTree(torch.nn.Module):
    """
    A medium-sized model that uses layers with few activation maps to
    efficiently increase the number of layers, and therefore nonlinearities.
    """

    def __init__(self, in_dim, n_status, n_out):
        """
        Args:
            in_dim (list): The input size of each example
            n_status (int): Number of status inputs to add
            n_out (int): Number of values to predict
        """
        super(StarTree, self).__init__()
        self.n_status = n_status
        dim = in_dim.copy()
        self.feat = torch.nn.Sequential(
            ConvBlock(dim, 64, dropout=0.25),
            ConvBlock(dim, 16),
            ConvBlock(dim, 32),
            PoolBlock(dim, "max", 2),
            ConvBlock(dim, 24),
            ConvBlock(dim, 48),
            PoolBlock(dim, "max", 2),
            ConvBlock(dim, 32),
            ConvBlock(dim, 64),
            PoolBlock(dim, "max", 2),
            ConvBlock(dim, 40),
            ConvBlock(dim, 80, dropout=0.25),
            PoolBlock(dim, "max", 2),
        )
        self.view = ViewBlock(dim)
        dim[0] += n_status
        self.head = torch.nn.Sequential(
            LinearBlock(dim, 50), LinearBlock(dim, n_out, activation=False),
        )
        self.n_params = sum([x.n_params for x in self.feat]) + sum([x.n_params for x in self.head])
        print("StarTree                                  params %9i" % self.n_params)

    def forward(self, batch, status):
        """
        Args:
            batch (4D tensor): A batch of camera input.
            status (1D tensor): Status inputs indicating things like speed.
        """
        out = self.feat(batch)
        out = self.view(out)
        if self.n_status:
            out = torch.cat((out, status), 1)
        out = self.head(out)
        return out


def train_epoch(device, model, optimizer, criterion, loader):
    """ Run the optimzer over all batches in an epoch """
    model.train()
    epoch_loss = 0
    batch_index = 0
    for batch_index, (examples, statuses, labels) in enumerate(loader):
        optimizer.zero_grad()
        guesses = model(examples.to(device), statuses.to(device))
        loss = criterion(guesses, labels.to(device))
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    return epoch_loss / (batch_index + 1)


def test_epoch(device, model, criterion, loader):
    """ Run the evaluator over all batches in an epoch """
    model.eval()
    epoch_loss = 0
    batch_index = 0
    with torch.no_grad():
        for batch_index, (examples, statuses, labels) in enumerate(loader):
            guesses = model(examples.to(device), statuses.to(device))
            loss = criterion(guesses, labels.to(device))
            epoch_loss += loss.item()
    return epoch_loss / (batch_index + 1)


def compose_transforms(transform_config):
    """ Apply all image transforms """
    transform_list = []
    for perturb_config in transform_config:
        if perturb_config["name"] == "colorjitter":
            transform = transforms.ColorJitter(
                brightness=perturb_config["brightness"],
                contrast=perturb_config["contrast"],
                saturation=perturb_config["saturation"],
                hue=perturb_config["hue"],
            )
            transform_list.append(transform)
    transform_list.append(transforms.ToTensor())
    return transforms.Compose(transform_list)

    
