import torch
import torch.nn as nn
from torch.optim import SGD

import MinkowskiEngine as ME

from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck

from examples.common import data_loader
from examples.resnet import ResNetBase


class MinkUNetBase(ResNetBase):
    BLOCK = None
    PLANES = None
    DILATIONS = (1, 1, 1, 1, 1, 1, 1, 1)
    LAYERS = (2, 2, 2, 2, 2, 2, 2, 2)
    INIT_DIM = 32
    OUT_TENSOR_STRIDE = 1

    # To use the model, must call initialize_coords before forward pass.
    # Once data is processed, call clear to reset the model before calling
    # initialize_coords
    def __init__(self, in_channels, out_channels, D=3):
        ResNetBase.__init__(self, in_channels, out_channels, D)

    def network_initialization(self, in_channels, out_channels, D):
        # Output of the first conv concated to conv6
        self.inplanes = self.INIT_DIM
        self.conv0p1s1 = ME.MinkowskiConvolution(
            in_channels, self.inplanes, kernel_size=5, dimension=D)

        self.bn0 = ME.MinkowskiBatchNorm(self.inplanes)

        self.conv1p1s2 = ME.MinkowskiConvolution(
            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
        self.bn1 = ME.MinkowskiBatchNorm(self.inplanes)

        self.block1 = self._make_layer(self.BLOCK, self.PLANES[0],
                                       self.LAYERS[0])

        self.conv2p2s2 = ME.MinkowskiConvolution(
            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
        self.bn2 = ME.MinkowskiBatchNorm(self.inplanes)

        self.block2 = self._make_layer(self.BLOCK, self.PLANES[1],
                                       self.LAYERS[1])

        self.conv3p4s2 = ME.MinkowskiConvolution(
            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)

        self.bn3 = ME.MinkowskiBatchNorm(self.inplanes)
        self.block3 = self._make_layer(self.BLOCK, self.PLANES[2],
                                       self.LAYERS[2])

        self.conv4p8s2 = ME.MinkowskiConvolution(
            self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)
        self.bn4 = ME.MinkowskiBatchNorm(self.inplanes)
        self.block4 = self._make_layer(self.BLOCK, self.PLANES[3],
                                       self.LAYERS[3])

        self.convtr4p16s2 = ME.MinkowskiConvolutionTranspose(
            self.inplanes, self.PLANES[4], kernel_size=2, stride=2, dimension=D)
        self.bntr4 = ME.MinkowskiBatchNorm(self.PLANES[4])

        self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
        self.block5 = self._make_layer(self.BLOCK, self.PLANES[4],
                                       self.LAYERS[4])
        self.convtr5p8s2 = ME.MinkowskiConvolutionTranspose(
            self.inplanes, self.PLANES[5], kernel_size=2, stride=2, dimension=D)
        self.bntr5 = ME.MinkowskiBatchNorm(self.PLANES[5])

        self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
        self.block6 = self._make_layer(self.BLOCK, self.PLANES[5],
                                       self.LAYERS[5])
        self.convtr6p4s2 = ME.MinkowskiConvolutionTranspose(
            self.inplanes, self.PLANES[6], kernel_size=2, stride=2, dimension=D)
        self.bntr6 = ME.MinkowskiBatchNorm(self.PLANES[6])

        self.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion
        self.block7 = self._make_layer(self.BLOCK, self.PLANES[6],
                                       self.LAYERS[6])
        self.convtr7p2s2 = ME.MinkowskiConvolutionTranspose(
            self.inplanes, self.PLANES[7], kernel_size=2, stride=2, dimension=D)
        self.bntr7 = ME.MinkowskiBatchNorm(self.PLANES[7])

        self.inplanes = self.PLANES[7] + self.INIT_DIM
        self.block8 = self._make_layer(self.BLOCK, self.PLANES[7],
                                       self.LAYERS[7])

        self.final = ME.MinkowskiConvolution(
            self.PLANES[7],
            out_channels,
            kernel_size=1,
            has_bias=True,
            dimension=D)
        self.relu = ME.MinkowskiReLU(inplace=True)

    def forward(self, x):
        out = self.conv0p1s1(x)
        out = self.bn0(out)
        out_p1 = self.relu(out)

        out = self.conv1p1s2(out_p1)
        out = self.bn1(out)
        out = self.relu(out)
        out_b1p2 = self.block1(out)

        out = self.conv2p2s2(out_b1p2)
        out = self.bn2(out)
        out = self.relu(out)
        out_b2p4 = self.block2(out)

        out = self.conv3p4s2(out_b2p4)
        out = self.bn3(out)
        out = self.relu(out)
        out_b3p8 = self.block3(out)

        # tensor_stride=16
        out = self.conv4p8s2(out_b3p8)
        out = self.bn4(out)
        out = self.relu(out)
        out = self.block4(out)

        # tensor_stride=8
        out = self.convtr4p16s2(out)
        out = self.bntr4(out)
        out = self.relu(out)

        out = ME.cat((out, out_b3p8))
        out = self.block5(out)

        # tensor_stride=4
        out = self.convtr5p8s2(out)
        out = self.bntr5(out)
        out = self.relu(out)

        out = ME.cat((out, out_b2p4))
        out = self.block6(out)

        # tensor_stride=2
        out = self.convtr6p4s2(out)
        out = self.bntr6(out)
        out = self.relu(out)

        out = ME.cat((out, out_b1p2))
        out = self.block7(out)

        # tensor_stride=1
        out = self.convtr7p2s2(out)
        out = self.bntr7(out)
        out = self.relu(out)

        out = ME.cat((out, out_p1))
        out = self.block8(out)

        return self.final(out)


class MinkUNet14(MinkUNetBase):
    BLOCK = BasicBlock
    LAYERS = (1, 1, 1, 1, 1, 1, 1, 1)


class MinkUNet18(MinkUNetBase):
    BLOCK = BasicBlock
    LAYERS = (2, 2, 2, 2, 2, 2, 2, 2)


class MinkUNet34(MinkUNetBase):
    BLOCK = BasicBlock
    LAYERS = (2, 3, 4, 6, 2, 2, 2, 2)


class MinkUNet50(MinkUNetBase):
    BLOCK = Bottleneck
    LAYERS = (2, 3, 4, 6, 2, 2, 2, 2)


class MinkUNet101(MinkUNetBase):
    BLOCK = Bottleneck
    LAYERS = (2, 3, 4, 23, 2, 2, 2, 2)


class MinkUNet14A(MinkUNet14):
    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)


class MinkUNet14B(MinkUNet14):
    PLANES = (32, 64, 128, 256, 128, 128, 128, 128)


class MinkUNet14C(MinkUNet14):
    PLANES = (32, 64, 128, 256, 192, 192, 128, 128)


class MinkUNet14D(MinkUNet14):
    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)


class MinkUNet18A(MinkUNet18):
    PLANES = (32, 64, 128, 256, 128, 128, 96, 96)


class MinkUNet18B(MinkUNet18):
    PLANES = (32, 64, 128, 256, 128, 128, 128, 128)


class MinkUNet18D(MinkUNet18):
    PLANES = (32, 64, 128, 256, 384, 384, 384, 384)


class MinkUNet34A(MinkUNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 64, 64)


class MinkUNet34B(MinkUNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 64, 32)


class MinkUNet34C(MinkUNet34):
    PLANES = (32, 64, 128, 256, 256, 128, 96, 96)


if __name__ == '__main__':
    # loss and network
    criterion = nn.CrossEntropyLoss()
    net = MinkUNet14A(in_channels=3, out_channels=5, D=2)
    print(net)

    # a data loader must return a tuple of coords, features, and labels.
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = net.to(device)
    optimizer = SGD(net.parameters(), lr=1e-2)

    for i in range(10):
        optimizer.zero_grad()

        # Get new data
        coords, feat, label = data_loader(is_classification=False)
        input = ME.SparseTensor(feat, coords=coords).to(device)
        label = label.to(device)

        # Forward
        output = net(input)

        # Loss
        loss = criterion(output.F, label)
        print('Iteration: ', i, ', Loss: ', loss.item())

        # Gradient
        loss.backward()
        optimizer.step()

    # Saving and loading a network
    torch.save(net.state_dict(), 'test.pth')
    net.load_state_dict(torch.load('test.pth'))

[((227, 16, 227, 37), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((25, 8, 25, 63), 'examples.resnet.ResNetBase.__init__', 'ResNetBase.__init__', ({(25, 28, 25, 32): 'self', (25, 34, 25, 45): 'in_channels', (25, 47, 25, 59): 'out_channels', (25, 61, 25, 62): 'D'}, {}), '(self, in_channels, out_channels, D)', False, 'from examples.resnet import ResNetBase\n'), ((30, 25, 31, 67), 'MinkowskiEngine.MinkowskiConvolution', 'ME.MinkowskiConvolution', (), '', True, 'import MinkowskiEngine as ME\n'), ((33, 19, 33, 55), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(33, 41, 33, 54): 'self.inplanes'}, {}), '(self.inplanes)', True, 'import MinkowskiEngine as ME\n'), ((35, 25, 36, 79), 'MinkowskiEngine.MinkowskiConvolution', 'ME.MinkowskiConvolution', (), '', True, 'import MinkowskiEngine as ME\n'), ((37, 19, 37, 55), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(37, 41, 37, 54): 'self.inplanes'}, {}), '(self.inplanes)', True, 'import MinkowskiEngine as ME\n'), ((42, 25, 43, 79), 'MinkowskiEngine.MinkowskiConvolution', 'ME.MinkowskiConvolution', (), '', True, 'import MinkowskiEngine as ME\n'), ((44, 19, 44, 55), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(44, 41, 44, 54): 'self.inplanes'}, {}), '(self.inplanes)', True, 'import MinkowskiEngine as ME\n'), ((49, 25, 50, 79), 'MinkowskiEngine.MinkowskiConvolution', 'ME.MinkowskiConvolution', (), '', True, 'import MinkowskiEngine as ME\n'), ((52, 19, 52, 55), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(52, 41, 52, 54): 'self.inplanes'}, {}), '(self.inplanes)', True, 'import MinkowskiEngine as ME\n'), ((56, 25, 57, 79), 'MinkowskiEngine.MinkowskiConvolution', 'ME.MinkowskiConvolution', (), '', True, 'import MinkowskiEngine as ME\n'), ((58, 19, 58, 55), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(58, 41, 58, 54): 'self.inplanes'}, {}), '(self.inplanes)', True, 'import MinkowskiEngine as ME\n'), ((62, 28, 63, 80), 'MinkowskiEngine.MinkowskiConvolutionTranspose', 'ME.MinkowskiConvolutionTranspose', (), '', True, 'import MinkowskiEngine as ME\n'), ((64, 21, 64, 58), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(64, 43, 64, 57): 'self.PLANES[4]'}, {}), '(self.PLANES[4])', True, 'import MinkowskiEngine as ME\n'), ((69, 27, 70, 80), 'MinkowskiEngine.MinkowskiConvolutionTranspose', 'ME.MinkowskiConvolutionTranspose', (), '', True, 'import MinkowskiEngine as ME\n'), ((71, 21, 71, 58), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(71, 43, 71, 57): 'self.PLANES[5]'}, {}), '(self.PLANES[5])', True, 'import MinkowskiEngine as ME\n'), ((76, 27, 77, 80), 'MinkowskiEngine.MinkowskiConvolutionTranspose', 'ME.MinkowskiConvolutionTranspose', (), '', True, 'import MinkowskiEngine as ME\n'), ((78, 21, 78, 58), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(78, 43, 78, 57): 'self.PLANES[6]'}, {}), '(self.PLANES[6])', True, 'import MinkowskiEngine as ME\n'), ((83, 27, 84, 80), 'MinkowskiEngine.MinkowskiConvolutionTranspose', 'ME.MinkowskiConvolutionTranspose', (), '', True, 'import MinkowskiEngine as ME\n'), ((85, 21, 85, 58), 'MinkowskiEngine.MinkowskiBatchNorm', 'ME.MinkowskiBatchNorm', ({(85, 43, 85, 57): 'self.PLANES[7]'}, {}), '(self.PLANES[7])', True, 'import MinkowskiEngine as ME\n'), ((91, 21, 96, 24), 'MinkowskiEngine.MinkowskiConvolution', 'ME.MinkowskiConvolution', (), '', True, 'import MinkowskiEngine as ME\n'), ((97, 20, 97, 50), 'MinkowskiEngine.MinkowskiReLU', 'ME.MinkowskiReLU', (), '', True, 'import MinkowskiEngine as ME\n'), ((130, 14, 130, 37), 'MinkowskiEngine.cat', 'ME.cat', ({(130, 21, 130, 36): '(out, out_b3p8)'}, {}), '((out, out_b3p8))', True, 'import MinkowskiEngine as ME\n'), ((138, 14, 138, 37), 'MinkowskiEngine.cat', 'ME.cat', ({(138, 21, 138, 36): '(out, out_b2p4)'}, {}), '((out, out_b2p4))', True, 'import MinkowskiEngine as ME\n'), ((146, 14, 146, 37), 'MinkowskiEngine.cat', 'ME.cat', ({(146, 21, 146, 36): '(out, out_b1p2)'}, {}), '((out, out_b1p2))', True, 'import MinkowskiEngine as ME\n'), ((154, 14, 154, 35), 'MinkowskiEngine.cat', 'ME.cat', ({(154, 21, 154, 34): '(out, out_p1)'}, {}), '((out, out_p1))', True, 'import MinkowskiEngine as ME\n'), ((241, 30, 241, 66), 'examples.common.data_loader', 'data_loader', (), '', False, 'from examples.common import data_loader\n'), ((258, 24, 258, 46), 'torch.load', 'torch.load', ({(258, 35, 258, 45): '"""test.pth"""'}, {}), "('test.pth')", False, 'import torch\n'), ((232, 36, 232, 61), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((242, 16, 242, 52), 'MinkowskiEngine.SparseTensor', 'ME.SparseTensor', (), '', True, 'import MinkowskiEngine as ME\n')]