import sys
_module = sys.modules[__name__]
del sys
data_generator = _module
infer = _module
m_resunet = _module
metrics = _module
process_image = _module
resunet = _module
run = _module
unet = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import torch


import torch.nn as nn


class Squeeze_Excitation(nn.Module):

    def __init__(self, channel, r=8):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.net = nn.Sequential(nn.Linear(channel, channel // r, bias=False), nn.ReLU(inplace=True), nn.Linear(channel // r, channel, bias=False), nn.Sigmoid())

    def forward(self, inputs):
        b, c, _, _ = inputs.shape
        x = self.pool(inputs).view(b, c)
        x = self.net(x).view(b, c, 1, 1)
        x = inputs * x
        return x


class Stem_Block(nn.Module):

    def __init__(self, in_c, out_c, stride):
        super().__init__()
        self.c1 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=3, stride=stride, padding=1), nn.BatchNorm2d(out_c), nn.ReLU(), nn.Conv2d(out_c, out_c, kernel_size=3, padding=1))
        self.c2 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=1, stride=stride, padding=0), nn.BatchNorm2d(out_c))
        self.attn = Squeeze_Excitation(out_c)

    def forward(self, inputs):
        x = self.c1(inputs)
        s = self.c2(inputs)
        y = self.attn(x + s)
        return y


class ResNet_Block(nn.Module):

    def __init__(self, in_c, out_c, stride):
        super().__init__()
        self.c1 = nn.Sequential(nn.BatchNorm2d(in_c), nn.ReLU(), nn.Conv2d(in_c, out_c, kernel_size=3, padding=1, stride=stride), nn.BatchNorm2d(out_c), nn.ReLU(), nn.Conv2d(out_c, out_c, kernel_size=3, padding=1))
        self.c2 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=1, stride=stride, padding=0), nn.BatchNorm2d(out_c))
        self.attn = Squeeze_Excitation(out_c)

    def forward(self, inputs):
        x = self.c1(inputs)
        s = self.c2(inputs)
        y = self.attn(x + s)
        return y


class ASPP(nn.Module):

    def __init__(self, in_c, out_c, rate=[1, 6, 12, 18]):
        super().__init__()
        self.c1 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=3, dilation=rate[0], padding=rate[0]), nn.BatchNorm2d(out_c))
        self.c2 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=3, dilation=rate[1], padding=rate[1]), nn.BatchNorm2d(out_c))
        self.c3 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=3, dilation=rate[2], padding=rate[2]), nn.BatchNorm2d(out_c))
        self.c4 = nn.Sequential(nn.Conv2d(in_c, out_c, kernel_size=3, dilation=rate[3], padding=rate[3]), nn.BatchNorm2d(out_c))
        self.c5 = nn.Conv2d(out_c, out_c, kernel_size=1, padding=0)

    def forward(self, inputs):
        x1 = self.c1(inputs)
        x2 = self.c2(inputs)
        x3 = self.c3(inputs)
        x4 = self.c4(inputs)
        x = x1 + x2 + x3 + x4
        y = self.c5(x)
        return y


class Attention_Block(nn.Module):

    def __init__(self, in_c):
        super().__init__()
        out_c = in_c[1]
        self.g_conv = nn.Sequential(nn.BatchNorm2d(in_c[0]), nn.ReLU(), nn.Conv2d(in_c[0], out_c, kernel_size=3, padding=1), nn.MaxPool2d((2, 2)))
        self.x_conv = nn.Sequential(nn.BatchNorm2d(in_c[1]), nn.ReLU(), nn.Conv2d(in_c[1], out_c, kernel_size=3, padding=1))
        self.gc_conv = nn.Sequential(nn.BatchNorm2d(in_c[1]), nn.ReLU(), nn.Conv2d(out_c, out_c, kernel_size=3, padding=1))

    def forward(self, g, x):
        g_pool = self.g_conv(g)
        x_conv = self.x_conv(x)
        gc_sum = g_pool + x_conv
        gc_conv = self.gc_conv(gc_sum)
        y = gc_conv * x
        return y


class Decoder_Block(nn.Module):

    def __init__(self, in_c, out_c):
        super().__init__()
        self.a1 = Attention_Block(in_c)
        self.up = nn.Upsample(scale_factor=2, mode='nearest')
        self.r1 = ResNet_Block(in_c[0] + in_c[1], out_c, stride=1)

    def forward(self, g, x):
        d = self.a1(g, x)
        d = self.up(d)
        d = torch.cat([d, g], axis=1)
        d = self.r1(d)
        return d


class build_resunetplusplus(nn.Module):

    def __init__(self):
        super().__init__()
        self.c1 = Stem_Block(3, 16, stride=1)
        self.c2 = ResNet_Block(16, 32, stride=2)
        self.c3 = ResNet_Block(32, 64, stride=2)
        self.c4 = ResNet_Block(64, 128, stride=2)
        self.b1 = ASPP(128, 256)
        self.d1 = Decoder_Block([64, 256], 128)
        self.d2 = Decoder_Block([32, 128], 64)
        self.d3 = Decoder_Block([16, 64], 32)
        self.aspp = ASPP(32, 16)
        self.output = nn.Conv2d(16, 1, kernel_size=1, padding=0)

    def forward(self, inputs):
        c1 = self.c1(inputs)
        c2 = self.c2(c1)
        c3 = self.c3(c2)
        c4 = self.c4(c3)
        b1 = self.b1(c4)
        d1 = self.d1(c3, b1)
        d2 = self.d2(c2, d1)
        d3 = self.d3(c1, d2)
        output = self.aspp(d3)
        output = self.output(output)
        return output


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (ASPP,
     lambda: ([], {'in_c': 4, 'out_c': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (Attention_Block,
     lambda: ([], {'in_c': [4, 4]}),
     lambda: ([torch.rand([4, 4, 2, 2]), torch.rand([4, 4, 4, 4])], {}),
     True),
    (ResNet_Block,
     lambda: ([], {'in_c': 4, 'out_c': 4, 'stride': 1}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (Squeeze_Excitation,
     lambda: ([], {'channel': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (Stem_Block,
     lambda: ([], {'in_c': 4, 'out_c': 4, 'stride': 1}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (build_resunetplusplus,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
]

class Test_DebeshJha_ResUNetPlusPlus(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

    def test_004(self):
        self._check(*TESTCASES[4])

    def test_005(self):
        self._check(*TESTCASES[5])

