from _context import sparse
from sparse import util

import torch
from torch import nn
from torch.autograd import Variable
import torch.nn.functional as F
import torch.distributions as dist

import numpy as np

from argparse import ArgumentParser
from torch.utils.tensorboard import SummaryWriter

import torchvision
import torchvision.transforms as transforms
from torchvision.transforms import ToTensor
from torch.utils.data import TensorDataset, DataLoader
from torchvision.datasets import MNIST, CIFAR10, CIFAR100

import random, tqdm, sys, math, os

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt

from util import d


### Wide resnet. Ported from fast.ai lib

def _bn(ni, init_zero=False):
    """
    Batchnorm layer with (optional) 0 initialization

    :param ni:
    :param init_zero:
    :return:
    """
    m = nn.BatchNorm2d(ni)
    m.weight.data.fill_(0 if init_zero else 1)
    m.bias.data.zero_()
    return m

def bn_relu_conv(ni, nf, ks, stride, init_zero=False):
    """
    Batchnorm/relu/convolution block
    """

    bn_initzero = _bn(ni, init_zero=init_zero)

    return nn.Sequential(
        bn_initzero,
        nn.ReLU(inplace=True),
        nn.Conv2d(in_channels=ni, out_channels=nf, kernel_size=ks, stride=stride, padding=ks//2)
    )


class WideBlock(nn.Module):
    """
    Wide ResNet block
    """
    def __init__(self, ni, nf, stride, drop_p=0.0, mult=0.2):
        super().__init__()

        self.mult = mult
        self.bn = nn.BatchNorm2d(ni)

        self.conv1 = nn.Conv2d(in_channels=ni, out_channels=nf, kernel_size=3, stride=stride, padding=1)
        self.conv2 = bn_relu_conv(nf, nf, 3, 1)

        self.drop = nn.Dropout(drop_p) if drop_p else None
        self.shortcut = nn.Conv2d(in_channels=ni, out_channels=nf, kernel_size=1, stride=stride) if ni != nf else util.Lambda(lambda x : x)

    def forward(self, x):

        x2 = F.relu(self.bn(x))

        r = self.shortcut(x2)
        x = self.conv1(x2)

        if self.drop:
            x = self.drop(x)
        x = self.conv2(x) * self.mult

        return x + r

def _make_group(N, ni, nf, block, stride, drop_p):
    return [block(ni if i == 0 else nf, nf, stride if i == 0 else 1, drop_p) for i in range(N)]

class WideResNet(nn.Module):
    """
    Wide ResNet with `num_groups` and a width of `k`.
    """
    def __init__(self, num_classes:int, num_groups:int=3, N:int=3, k:int=6, drop_p:float=0.0, start_nf:int=16, n_in_channels:int=3, **kwargs):
        super().__init__()

        n_channels = [start_nf]
        for i in range(num_groups): n_channels.append(start_nf*(2**i)*k)

        layers = [nn.Conv2d(n_in_channels, n_channels[0], 3, 1, padding=1)]  # conv1
        for i in range(num_groups):
            layers += _make_group(N, n_channels[i], n_channels[i+1], WideBlock, (1 if i==0 else 2), drop_p)

        layers += [
            nn.BatchNorm2d(n_channels[num_groups]), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1), util.Flatten(),
            nn.Linear(n_channels[num_groups], num_classes)
        ]

        self.features = nn.Sequential(*layers)

    def forward(self, x): return self.features(x)

class Cutout(object):
    """
    Cutout transform. From https://github.com/uoguelph-mlrg/Cutout/blob/master/util/cutout.py

    Randomly mask out one or more patches from an image.
    Args:
        n_holes (int): Number of patches to cut out of each image.
        length (int): The length (in pixels) of each square patch.
    """
    def __init__(self, n_holes, length):
        self.n_holes = n_holes
        self.length = length

    def __call__(self, img):
        """
        Args:
            img (Tensor): Tensor image of size (C, H, W).
        Returns:
            Tensor: Image with n_holes of dimension length x length cut out of it.
        """

        h = img.size(1)
        w = img.size(2)

        mask = np.ones((h, w), np.float32)

        for n in range(self.n_holes):
            y = np.random.randint(h)
            x = np.random.randint(w)

            y1 = np.clip(y - self.length // 2, 0, h)
            y2 = np.clip(y + self.length // 2, 0, h)
            x1 = np.clip(x - self.length // 2, 0, w)
            x2 = np.clip(x + self.length // 2, 0, w)

            mask[y1: y2, x1: x2] = 0.

        mask = torch.from_numpy(mask)
        mask = mask.expand_as(img)
        img = img * mask

        return img



class MobBlock(nn.Module):
    """
    Inverted residual block (as in mobilenetv2)
    """
    def __init__(self,  c, wide=None, kernel=3):
        super().__init__()

        wide = 6 * c if wide is None else wide
        padding = int(math.floor(kernel/2))

        self.convs = nn.Sequential(
            nn.Conv2d(c, wide, kernel_size=1),
            #nn.BatchNorm2d(wide),
            nn.ReLU6(),
            nn.Conv2d(wide, wide, kernel_size=kernel, padding=padding, groups=wide),
            #nn.BatchNorm2d(wide),
            nn.ReLU6(),
            nn.Conv2d(wide, c, kernel_size=1),
            nn.BatchNorm2d(c), nn.ReLU6()
        )

    def forward(self, x):

        return self.convs(x) + x

class Convolution(nn.Module):
    """
    A "convolution" with a learned sparsity structure (as opposed to a fixed k x k grid)

    A sparse transformation is used to select k pixels from the input. Their channel
    vectors are concatenated and linearly transformed to the desired number of output
    channels (this is done for every output pixel).

    The pattern of input pixels, relative to the current output pixels is determined
    adaptively.
    """
    def __init__(self, cin, cout, k, gadditional, radditional, region, min_sigma=0.05, sigma_scale=0.05,
                 mmult=1.0, admode='full', edges='sigmoid', bias=True, constraints='free'):
        """
        :param k: Number of connections to the input in total
        :param gadditional:
        :param radditional:
        :param pooling factor: The output resolution is reduced by this factor.
        :param region:
        """

        super().__init__()

        # cin, _, _ = in_size
        # cout, hout, wout = cout, hin, win

        self.gadditional, self.radditional = gadditional, radditional

        self.region = (region, region)

        self.min_sigma, self.sigma_scale = min_sigma, sigma_scale
        self.admode, self.mmult = admode, mmult
        self.edges = edges

        if admode == 'none':
            self.mmult *= 0.1
        if edges == 'modulo' or edges == 'clamp':
            self.mmult *= 5.0

        self.k = k
        self.unify = nn.Linear(k*cin, cout, bias=bias)
        self.constraints = constraints

        # Set number of output parameters per pixel
        if self.constraints == 'free':
            self.nparms = k * 3
        elif self.constraints == 'affine':
            self.nparms = 4
        elif self.constraints == 'scale':
            self.nparms = 2
        else:
            raise Exception(f'{self.constraints}')

        if admode == 'none':
            self.params = nn.Parameter(torch.randn(size=(self.nparms, )))
        else:
            if admode == 'full':
                adin = 2 + cin
            elif admode == 'coords':
                adin = 2
            elif admode == 'inputs':
                adin = cin
            else:
                raise Exception(f'adaptivity mode {admode} not recognized')

            # network that generates the coordinates and sigmas
            hidden = cin * 4
            self.toparams = nn.Sequential(
                nn.Linear(adin, hidden), nn.ReLU(),
                nn.Linear(hidden, self.nparms) # two means, one sigma
            )

        self.register_buffer('mvalues', torch.ones((k,)))

        self.smp = True

        if self.constraints == 'affine' or self.constraints == 'scale':
            kp = int(math.sqrt(k))
            self.register_buffer('base', util.coordinates((kp, kp)) - 0.5)

    def sample(self, smp):
        self.smp = smp

    def hyper(self, x):

        b, c, h, w = x.size()
        k = self.k

        coords = util.coordinates((h, w), cuda=x.is_cuda)

        # the coordinates of the current pixels in parameters space
        # - the index tuples are described relative to these
        hw = torch.tensor((h, w), device=d(x), dtype=torch.float)
        mids = coords[None, :, :, :].expand(b, 2, h, w) * (hw - 1)[None, :, None, None]
        mids = mids.permute(0, 2, 3, 1)
        if self.edges == 'sigmoid':
            mids = util.inv(mids, mx=hw[None, None, None, :])
        mids = mids[:, :, :, None, :].expand(b, h, w, k, 2)

        # add coords to channels
        if self.admode == 'none':
            params = self.params[None, None, None, :].expand(b, h, w, self.nparms)
        else:
            if self.admode == 'full':
                crds = coords[None, :, :, :].expand(b, 2, h, w)
                x = torch.cat([x, crds], dim=1)
            elif self.admode == 'coords':
                x = coords[None, :, :, :].expand(b, 2, h, w)
            elif self.admode == 'inputs':
                pass
            else:
                raise Exception(f'adaptivity mode {self.admode} not recognized')

            x = x.permute(0, 2, 3, 1)
            params = self.toparams(x)

        if self.constraints == 'affine':
            # Take the three output values (per pixel) and use them to transform a grid of
            # inputs (rotated and scale in two dimensions)
            MULT = 10.0

            kp = int(math.sqrt(k))

            assert params.size() == (b, h, w, 4)

            scalex = params[:, :, :, 0:1, None, None]
            scaley = params[:, :, :, 1:2, None, None]
            rotate = params[:, :, :, 2:3, None, None]
            sigmas = params[:, :, :, 3:4, None, None]

            coords = self.base.view(1, 1, 1, 2, kp, kp).expand(b, h, w, 2, kp, kp).contiguous() * MULT

            # scale
            coords[:, :, :, 0:1, :, :] *= scalex
            coords[:, :, :, 1:2, :, :] *= scaley

            # rotate
            rotx = coords[:, :, :, 0:1, :, :] * rotate.cos() - coords[:, :, :, 1:2, :, :] * rotate.sin()
            roty = coords[:, :, :, 0:1, :, :] * rotate.sin() + coords[:, :, :, 1:2, :, :] * rotate.cos()

            sigmas = sigmas.expand(b, h, w, 1, kp, kp)
            params = torch.cat([rotx, roty, sigmas], dim=3)

            assert params.size() == (b, h, w, 3, kp, kp)

            params = params.reshape(b, h, w, 3, k).permute(0, 1, 2, 4, 3)
        elif self.constraints == 'scale':
            # Take the output values (per pixel) and use them to transform a grid of
            # inputs (uniformly scaled)
            MULT = 10.0

            kp = int(math.sqrt(k))

            assert params.size() == (b, h, w, 2)

            scales = params[:, :, :, 0:1, None, None]
            sigmas = params[:, :, :, 1:2, None, None]

            coords = self.base.view(1, 1, 1, 2, kp, kp).expand(b, h, w, 2, kp, kp).contiguous() * MULT

            # scale
            coords[:, :, :, :, :, :] *= scales

            sigmas = sigmas.expand(b, h, w, 1, kp, kp)
            params = torch.cat([coords, sigmas], dim=3)

            assert params.size() == (b, h, w, 3, kp, kp)

            params = params.reshape(b, h, w, 3, k).permute(0, 1, 2, 4, 3)
        elif self.constraints == 'free':
            params = params.reshape(b, h, w, k, 3)
        else:
            raise Exception(f'{self.constraints}')

        assert params.size() == (b, h, w, k, 3) # k index tuples per output pixel

        means  = params[:, :, :, :, :2]
        sigmas = params[:, :, :, :, 2]
        values = self.mvalues[None, None, None, :].expand(b, h, w, k)

        means = mids + self.mmult * means

        s = (h, w)
        means  = sparse.transform_means(means, s, method=self.edges)
        sigmas = sparse.transform_sigmas(sigmas, s, min_sigma=self.min_sigma) * self.sigma_scale

        return means, sigmas, values

    def forward(self, x):

        b, c, h, w = x.size()
        k = self.k
        s = (h, w)

        assert h > 2 and w > 2, 'Input resolution must be larger than 2x2 for the sparse convolution to work.'

        means, sigmas, mvalues = self.hyper(x)

        # This is a bit confusing, but k is the chunk dimension here. This is because the sparse operation
        # only selects in the k separate input pixels, it doens not sum/merge them.
        # In other words, we add a separate tuple dimension.
        means   = means  [:, :, :, :, None, :]
        sigmas  = sigmas [:, :, :, :, None, :]
        mvalues = mvalues[:, :, :, :, None]

        if self.smp:
            # sample integer indices and values
            indices = sparse.ngenerate(means, self.gadditional, self.radditional, rng=s, relative_range=self.region, cuda=x.is_cuda)

            vs = (4 + self.radditional + self.gadditional)
            assert indices.size() == (b, h, w, k, vs, 2), f'{indices.size()}, {(b, h, w, k, vs, 2)}'

            indices = indices.view(b, h, w, k, vs, 2)
            indfl = indices.float()

            # Mask for duplicate indices
            dups = util.nduplicates(indices).to(torch.bool)

            # compute (unnormalized) densities under the given MVNs (proportions)
            props = sparse.densities(indfl, means, sigmas).clone() # (b, h, w, k, vs, 1)
            assert props.size() == (b, h, w, k, vs, 1)

            props[dups, :] = 0
            props = props / props.sum(dim=4, keepdim=True)  # normalize over all points of a given index tuple

            # weight the values by the proportions
            weights = mvalues[:, :, :, :, None, :].expand_as(props)
            # - add a dim for the MVNs

            weights = props * weights # TODO mult by one, this can be removed
            weights = weights.sum(dim=5)  # - sum out the MVNs

            assert indices.size() == (b, h, w, k, vs, 2)
            assert weights.size() == (b, h, w, k, vs)

        else:
            vs = 1
            indices = means.floor().to(torch.long).detach()

        l = h * w * k * vs
        indices = indices.view(b*l, 2)

        br = torch.arange(b, device=d(x), dtype=torch.long)[:, None].expand(b, l).contiguous().view(-1)
        features = x[br, :, indices[:, 0], indices[:, 1]]
        assert features.size() == (b*l, c)

        if self.smp:
            features = features.view(b, h, w, k, vs, c)
            features = features * weights[:, :, :, :, :, None]
            features = features.sum(dim=4)
        else:
            features = features.view(b, h, w, k, c)

        # features now contains the selected input pixels (or weighted sum thereover): k inputs per output pixel
        assert features.size() == (b, h, w, k, c), f'Was {features.size()}, expected {(b, h, w, k, c)}.'

        features = features.view(b, h, w, k * c)

        return self.unify(features).permute(0, 3, 1, 2) # (b, c_out, h, w)

    def plot(self, inputs, numpixels=5, ims=None):

        ims = inputs if ims is None else ims

        b, c, h, w = inputs.size()
        b, cims, hims, wims = ims.size()

        k = self.k

        # choose 5 random pixels, for which we'll plot the input pixels.
        choices = torch.randint(low=0, high=h*w, size=(numpixels,))

        perrow = 5

        rows = int(math.ceil(b/perrow))

        means, sigmas, _ = self.hyper(inputs)

        inputs = inputs.data

        plt.figure(figsize=(perrow * 3, rows*3))

        # scale up to image coordinates
        scale = torch.tensor((hims/h, wims/w), device=d(inputs))
        means = means * scale + (scale/2)

        for current in range(b):

            # select subset of means, sigmas
            smeans = means[current, :, :, :, :].view(h*w, k, 2)
            ssigmas = sigmas[current, :, :, :].view(h*w, k, 2)

            color = (torch.arange(numpixels, dtype=torch.float)[:, None].expand(numpixels, k)/numpixels) * 2.0 - 1.0

            smeans = smeans[choices, :, :]
            ssigmas = ssigmas[choices, :]

            ax = plt.subplot(rows, perrow, current+1)

            im = np.transpose(ims[current, :, :, :].cpu().numpy(), (1, 2, 0))
            im = np.squeeze(im)

            ax.imshow(im, interpolation='nearest', extent=(-0.5, wims-0.5, -0.5, hims-0.5), cmap='gray_r')

            util.plot(smeans.reshape(1, -1, 2), ssigmas.reshape(1, -1, 2), color.reshape(1, -1), axes=ax, flip_y=hims, tanh=False)