'''ResNet using PSG in PyTorch.

For Pre-activation ResNet, see 'preact_resnet.py'.

Reference:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
    Deep Residual Learning for Image Recognition. arXiv:1512.03385
'''
from numpy.lib.arraysetops import isin
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np

from models.masked_psg_seed_conv import PredictiveSeedConv2d
from masked_layers import layers


# Fixed
NUM_BITS = 32
NUM_BITS_WEIGHT = 32
NUM_BITS_GRAD = None

BIPRECISION = False
PREDICTIVE_FORWARD = False

WRITER = None
WRITER_PREFIX_COUNTER = 0

# Tunable
PREDICTIVE_BACKWARD = True

MSB_BITS = 4
MSB_BITS_WEIGHT = 4
MSB_BITS_GRAD = 8

THRESHOLD = 0.0
SPARSIFY = False
SIGN = True



def conv1x1(in_planes, out_planes, stride=1, input_signed=True, predictive_forward=True, writer_prefix=""):
    "1x1 convolution with no padding"
    predictive_forward = PREDICTIVE_FORWARD and predictive_forward
    return PredictiveSeedConv2d(
        in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=False,
        num_bits=NUM_BITS, num_bits_weight=NUM_BITS_WEIGHT, num_bits_grad=NUM_BITS_GRAD,
        biprecision=BIPRECISION, input_signed=input_signed,
        predictive_forward=predictive_forward, predictive_backward=PREDICTIVE_BACKWARD,
        msb_bits=MSB_BITS, msb_bits_weight=MSB_BITS_WEIGHT, msb_bits_grad=MSB_BITS_GRAD,
        threshold=THRESHOLD, sparsify=SPARSIFY, sign=SIGN,
        writer=WRITER, writer_prefix=writer_prefix)


def conv3x3(in_planes, out_planes, stride=1, input_signed=False, predictive_forward=True, writer_prefix=""):
    "3x3 convolution with padding"
    predictive_forward = PREDICTIVE_FORWARD and predictive_forward
    return PredictiveSeedConv2d(
        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False,
        num_bits=NUM_BITS, num_bits_weight=NUM_BITS_WEIGHT, num_bits_grad=NUM_BITS_GRAD,
        biprecision=BIPRECISION, input_signed=input_signed,
        predictive_forward=predictive_forward, predictive_backward=PREDICTIVE_BACKWARD,
        msb_bits=MSB_BITS, msb_bits_weight=MSB_BITS_WEIGHT, msb_bits_grad=MSB_BITS_GRAD,
        threshold=THRESHOLD, sparsify=SPARSIFY, sign=SIGN,
        writer=WRITER, writer_prefix=writer_prefix)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_planes, planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes, stride=1, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                # nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                conv1x1(in_planes, self.expansion*planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_planes, planes, stride=1):
        super(Bottleneck, self).__init__()
        # self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
        self.conv1 = conv1x1(in_planes, planes, stride=1, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn1 = nn.BatchNorm2d(planes)
        # self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.conv2 = conv3x3(planes, planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn2 = nn.BatchNorm2d(planes)
        # self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
        self.conv3 = conv1x1(planes, self.expansion*planes, stride=1, input_signed=False, predictive_forward=False, writer_prefix=None)
        self.bn3 = nn.BatchNorm2d(self.expansion*planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                # nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                conv1x1(in_planes, self.expansion*planes, stride=stride, input_signed=False, predictive_forward=False, writer_prefix=None),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, in_planes=64, num_classes=10, init_method='standard'):
        super(ResNet, self).__init__()
        self.in_planes = in_planes

        self.conv1 = conv3x3(3, self.in_planes, stride=1, input_signed=True, predictive_forward=False, writer_prefix=None)
        self.bn1 = nn.BatchNorm2d(self.in_planes)
        if self.in_planes ==  64:
            # self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
            self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
            self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
            self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
            self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
            self.linear = nn.Linear(512*block.expansion, num_classes)
            #self.linear = layers.Linear(512*block.expansion, num_classes)
        elif self.in_planes == 16:
            self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
            self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
            self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
            self.layer4 = None
            self.linear = nn.Linear(64, num_classes)

        self.reset_conv_parameters(init_method)
        print('conv weights reset to {}'.format(init_method))

    def reset_parameters(self, module, init_method="kaiming_uniform") -> None:
        if init_method == "kaiming_constant_signed":
            fan = nn.init._calculate_correct_fan(module.weight, "fan_in")
            gain = nn.init.calculate_gain("relu")
            std = gain / math.sqrt(fan)
            with torch.no_grad():
                module.weight.data = module.weight.data.sign() * std
        elif init_method == "kaiming_constant_unsigned":
            fan = nn.init._calculate_correct_fan(module.weight, "fan_in")
            gain = nn.init.calculate_gain("relu")
            std = gain / math.sqrt(fan)
            with torch.no_grad():
                module.weight.data = torch.ones_like(module.weight.data) * std
        elif init_method == "kaiming_normal":
            nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu")
        elif init_method == "kaiming_uniform":
            nn.init.kaiming_uniform_(module.weight, mode="fan_in", nonlinearity="relu")
        elif init_method == "kaiming_laplace":
            fan = nn.init._calculate_correct_fan(module.weight, "fan_in")
            gain = nn.init.calculate_gain("relu")
            scale = gain / math.sqrt(2.0 * fan)
            with torch.no_grad():
                new_weight = np.random.laplace(loc=0.0, scale=scale, size=module.weight.shape)
                module.weight.data = module.weight.data.new_tensor(torch.from_numpy(new_weight).clone().detach())
        elif init_method == "xavier_normal":
            nn.init.xavier_normal_(module.weight)
        elif init_method == "xavier_constant":
            fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(module.weight)
            std = math.sqrt(2.0 / float(fan_in + fan_out))
            with torch.no_grad():
                module.weight.data = module.weight.data.sign() * std
        elif init_method == "standard":
            nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
        else:
            raise ValueError(f"{init_method} is not an initialization option!")

    def reset_conv_parameters(self, init_method="standard") -> None:
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                self.reset_parameters(m, init_method)

    def get_bop_params(self):
        bop_params = []
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                bop_params += list(m.parameters())
        return bop_params

    def get_non_bop_params(self):
        non_bop_params = []
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.BatchNorm2d,)):
                non_bop_params += list(m.parameters())
        return non_bop_params

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        if self.layer4 is not None:
            out = self.layer4(out)
        # out = F.avg_pool2d(out, 4)
        out = F.avg_pool2d(out, out.size()[3])
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


def PsgSeedResNet20(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(BasicBlock, [3,3,3], in_planes=16, num_classes=num_classes, init_method=init_method)

def PsgSeedResNet18(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(BasicBlock, [2,2,2,2], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet34(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(BasicBlock, [3,4,6,3], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet50(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(Bottleneck, [3,4,6,3], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet101(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(Bottleneck, [3,4,23,3], num_classes=num_classes, init_method=init_method)

def PsgSeedResNet152(
    num_classes=10,
    init_method='standard',
    predictive_backward=True,
    msb_bits=4,
    msb_bits_weight=4,
    msb_bits_grad=8,
    threshold=0.0,
    sparsify=False,
    sign=True
):
    global PREDICTIVE_BACKWARD, MSB_BITS, MSB_BITS_WEIGHT, MSB_BITS_GRAD, THRESHOLD, SPARSIFY, SIGN
    PREDICTIVE_BACKWARD = predictive_backward
    MSB_BITS = msb_bits
    MSB_BITS_WEIGHT = msb_bits_weight
    MSB_BITS_GRAD = msb_bits_grad
    THRESHOLD = threshold
    SPARSIFY = sparsify
    SIGN = sign
    return ResNet(Bottleneck, [3,8,36,3], num_classes=num_classes, init_method=init_method)


def test():
    net = ResNet18()
    y = net(torch.randn(1,3,32,32))
    print(y.size())

# test()

[((47, 11, 54, 51), 'models.masked_psg_seed_conv.PredictiveSeedConv2d', 'PredictiveSeedConv2d', (), '', False, 'from models.masked_psg_seed_conv import PredictiveSeedConv2d\n'), ((60, 11, 67, 51), 'models.masked_psg_seed_conv.PredictiveSeedConv2d', 'PredictiveSeedConv2d', (), '', False, 'from models.masked_psg_seed_conv import PredictiveSeedConv2d\n'), ((76, 19, 76, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(76, 34, 76, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((78, 19, 78, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(78, 34, 78, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((80, 24, 80, 39), 'torch.nn.Sequential', 'nn.Sequential', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((92, 14, 92, 25), 'torch.nn.functional.relu', 'F.relu', ({(92, 21, 92, 24): 'out'}, {}), '(out)', True, 'import torch.nn.functional as F\n'), ((103, 19, 103, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(103, 34, 103, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((106, 19, 106, 41), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(106, 34, 106, 40): 'planes'}, {}), '(planes)', True, 'import torch.nn as nn\n'), ((109, 19, 109, 56), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(109, 34, 109, 55): 'self.expansion * planes'}, {}), '(self.expansion * planes)', True, 'import torch.nn as nn\n'), ((111, 24, 111, 39), 'torch.nn.Sequential', 'nn.Sequential', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((124, 14, 124, 25), 'torch.nn.functional.relu', 'F.relu', ({(124, 21, 124, 24): 'out'}, {}), '(out)', True, 'import torch.nn.functional as F\n'), ((134, 19, 134, 49), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(134, 34, 134, 48): 'self.in_planes'}, {}), '(self.in_planes)', True, 'import torch.nn as nn\n'), ((214, 15, 214, 37), 'torch.nn.Sequential', 'nn.Sequential', ({(214, 29, 214, 36): '*layers'}, {}), '(*layers)', True, 'import torch.nn as nn\n'), ((359, 12, 359, 34), 'torch.randn', 'torch.randn', ({(359, 24, 359, 25): '1', (359, 26, 359, 27): '3', (359, 28, 359, 30): '32', (359, 31, 359, 33): '32'}, {}), '(1, 3, 32, 32)', False, 'import torch\n'), ((141, 26, 141, 69), 'torch.nn.Linear', 'nn.Linear', ({(141, 36, 141, 55): '512 * block.expansion', (141, 57, 141, 68): 'num_classes'}, {}), '(512 * block.expansion, num_classes)', True, 'import torch.nn as nn\n'), ((155, 18, 155, 73), 'torch.nn.init._calculate_correct_fan', 'nn.init._calculate_correct_fan', ({(155, 49, 155, 62): 'module.weight', (155, 64, 155, 72): '"""fan_in"""'}, {}), "(module.weight, 'fan_in')", True, 'import torch.nn as nn\n'), ((156, 19, 156, 49), 'torch.nn.init.calculate_gain', 'nn.init.calculate_gain', ({(156, 42, 156, 48): '"""relu"""'}, {}), "('relu')", True, 'import torch.nn as nn\n'), ((85, 16, 85, 53), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(85, 31, 85, 52): 'self.expansion * planes'}, {}), '(self.expansion * planes)', True, 'import torch.nn as nn\n'), ((116, 16, 116, 53), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(116, 31, 116, 52): 'self.expansion * planes'}, {}), '(self.expansion * planes)', True, 'import torch.nn as nn\n'), ((148, 26, 148, 52), 'torch.nn.Linear', 'nn.Linear', ({(148, 36, 148, 38): '64', (148, 40, 148, 51): 'num_classes'}, {}), '(64, num_classes)', True, 'import torch.nn as nn\n'), ((157, 25, 157, 39), 'math.sqrt', 'math.sqrt', ({(157, 35, 157, 38): 'fan'}, {}), '(fan)', False, 'import math\n'), ((158, 17, 158, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((161, 18, 161, 73), 'torch.nn.init._calculate_correct_fan', 'nn.init._calculate_correct_fan', ({(161, 49, 161, 62): 'module.weight', (161, 64, 161, 72): '"""fan_in"""'}, {}), "(module.weight, 'fan_in')", True, 'import torch.nn as nn\n'), ((162, 19, 162, 49), 'torch.nn.init.calculate_gain', 'nn.init.calculate_gain', ({(162, 42, 162, 48): '"""relu"""'}, {}), "('relu')", True, 'import torch.nn as nn\n'), ((163, 25, 163, 39), 'math.sqrt', 'math.sqrt', ({(163, 35, 163, 38): 'fan'}, {}), '(fan)', False, 'import math\n'), ((164, 17, 164, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((167, 12, 167, 86), 'torch.nn.init.kaiming_normal_', 'nn.init.kaiming_normal_', (), '', True, 'import torch.nn as nn\n'), ((165, 37, 165, 72), 'torch.ones_like', 'torch.ones_like', ({(165, 53, 165, 71): 'module.weight.data'}, {}), '(module.weight.data)', False, 'import torch\n'), ((169, 12, 169, 87), 'torch.nn.init.kaiming_uniform_', 'nn.init.kaiming_uniform_', (), '', True, 'import torch.nn as nn\n'), ((171, 18, 171, 73), 'torch.nn.init._calculate_correct_fan', 'nn.init._calculate_correct_fan', ({(171, 49, 171, 62): 'module.weight', (171, 64, 171, 72): '"""fan_in"""'}, {}), "(module.weight, 'fan_in')", True, 'import torch.nn as nn\n'), ((172, 19, 172, 49), 'torch.nn.init.calculate_gain', 'nn.init.calculate_gain', ({(172, 42, 172, 48): '"""relu"""'}, {}), "('relu')", True, 'import torch.nn as nn\n'), ((173, 27, 173, 47), 'math.sqrt', 'math.sqrt', ({(173, 37, 173, 46): '(2.0 * fan)'}, {}), '(2.0 * fan)', False, 'import math\n'), ((174, 17, 174, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((175, 29, 175, 94), 'numpy.random.laplace', 'np.random.laplace', (), '', True, 'import numpy as np\n'), ((178, 12, 178, 49), 'torch.nn.init.xavier_normal_', 'nn.init.xavier_normal_', ({(178, 35, 178, 48): 'module.weight'}, {}), '(module.weight)', True, 'import torch.nn as nn\n'), ((180, 30, 180, 82), 'torch.nn.init._calculate_fan_in_and_fan_out', 'nn.init._calculate_fan_in_and_fan_out', ({(180, 68, 180, 81): 'module.weight'}, {}), '(module.weight)', True, 'import torch.nn as nn\n'), ((182, 17, 182, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((185, 54, 185, 66), 'math.sqrt', 'math.sqrt', ({(185, 64, 185, 65): '(5)'}, {}), '(5)', False, 'import math\n'), ((176, 67, 176, 95), 'torch.from_numpy', 'torch.from_numpy', ({(176, 84, 176, 94): 'new_weight'}, {}), '(new_weight)', False, 'import torch\n')]