
from collections import namedtuple

import torch
import torch.nn as nn
import torch.nn.functional as F

__all__ = ['Inception3', 'inception_v3']

_InceptionOutputs = namedtuple('InceptionOutputs', ['logits', 'aux_logits'])


class BasicConv2d(nn.Module):

  def __init__(self, in_channels, out_channels, **kwargs):
    super(BasicConv2d, self).__init__()
    self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
    self.bn = nn.BatchNorm2d(out_channels, eps=0.001)

  def forward(self, x):
    x = self.conv(x)
    x = self.bn(x)
    return F.relu(x, inplace=True)


class InceptionAux(nn.Module):

  def __init__(self, in_channels, num_classes):
    super(InceptionAux, self).__init__()
    self.conv0 = BasicConv2d(in_channels, 128, kernel_size=1)
    self.conv1 = BasicConv2d(128, 512, kernel_size=5)
    self.conv1.stddev = 0.01
    self.fc = nn.Linear(512, num_classes)
    self.fc.stddev = 0.001

  def forward(self, x):
    x = F.avg_pool2d(x, kernel_size=3, stride=3)
    x = self.conv0(x)
    # x = x.reshape(x.shape[0], -1)
    x = self.conv1(x)
    x = F.adaptive_avg_pool2d(x, (1, 1))
    x = torch.flatten(x, 1)
    x = self.fc(x)
    return x


class InceptionD(nn.Module):

  def __init__(self, in_channels):
    super(InceptionD, self).__init__()
    self.branch3x3_1 = BasicConv2d(in_channels, 32, kernel_size=1)
    self.branch3x3_2 = BasicConv2d(32,64, kernel_size=3, stride=2)

    self.branch7x7x3_1 = BasicConv2d(in_channels, 32, kernel_size=1)
    self.branch7x7x3_2 = BasicConv2d(32,64, kernel_size=(1, 7), padding=(0, 3))
    self.branch7x7x3_3 = BasicConv2d(64, 128, kernel_size=(7, 1), padding=(3, 0))
    self.branch7x7x3_4 = BasicConv2d(128,192, kernel_size=3, stride=2)

  def forward(self, x):
    branch3x3 = self.branch3x3_1(x)
    branch3x3 = self.branch3x3_2(branch3x3)

    branch7x7x3 = self.branch7x7x3_1(x)
    branch7x7x3 = self.branch7x7x3_2(branch7x7x3)
    branch7x7x3 = self.branch7x7x3_3(branch7x7x3)
    branch7x7x3 = self.branch7x7x3_4(branch7x7x3)

    branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
    outputs = [branch3x3, branch7x7x3, branch_pool]
    return torch.cat(outputs, 1)

class InceptionC(nn.Module):

  def __init__(self, in_channels, channels_7x7):
    super(InceptionC, self).__init__()
    self.branch1x1 = BasicConv2d(in_channels, 128, kernel_size=1)

    c7 = channels_7x7
    self.branch7x7_1 = BasicConv2d(in_channels, c7, kernel_size=1)
    self.branch7x7_2 = BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3))
    self.branch7x7_3 = BasicConv2d(c7, 128, kernel_size=(7, 1), padding=(3, 0))

    self.branch7x7dbl_1 = BasicConv2d(in_channels, c7, kernel_size=1)
    self.branch7x7dbl_2 = BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0))
    self.branch7x7dbl_3 = BasicConv2d(c7, c7, kernel_size=(1, 7), padding=(0, 3))
    self.branch7x7dbl_4 = BasicConv2d(c7, c7, kernel_size=(7, 1), padding=(3, 0))
    self.branch7x7dbl_5 = BasicConv2d(c7, 128, kernel_size=(1, 7), padding=(0, 3))

    self.branch_pool = BasicConv2d(in_channels, 128, kernel_size=1)

  def forward(self, x):
    branch1x1 = self.branch1x1(x)

    branch7x7 = self.branch7x7_1(x)
    branch7x7 = self.branch7x7_2(branch7x7)
    branch7x7 = self.branch7x7_3(branch7x7)

    branch7x7dbl = self.branch7x7dbl_1(x)
    branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
    branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
    branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
    branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)

    branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
    branch_pool = self.branch_pool(branch_pool)

    outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
    return torch.cat(outputs, 1)

class InceptionB(nn.Module):

  def __init__(self, in_channels):
    super(InceptionB, self).__init__()
    self.branch3x3 = BasicConv2d(in_channels, 32, kernel_size=3, stride=2)

    self.branch3x3dbl_1 = BasicConv2d(in_channels, 32, kernel_size=1)
    self.branch3x3dbl_2 = BasicConv2d(32, 64, kernel_size=3, padding=1)
    self.branch3x3dbl_3 = BasicConv2d(64, 96, kernel_size=3, stride=2)

  def forward(self, x):
    branch3x3 = self.branch3x3(x)

    branch3x3dbl = self.branch3x3dbl_1(x)
    branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
    branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)

    branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)

    outputs = [branch3x3, branch3x3dbl, branch_pool]
    return torch.cat(outputs, 1)


class InceptionA(nn.Module):

  def __init__(self, in_channels, pool_features):
    super(InceptionA, self).__init__()
    self.branch1x1 = BasicConv2d(in_channels, 8, kernel_size=1)

    self.branch5x5_1 = BasicConv2d(in_channels, 8, kernel_size=1)
    self.branch5x5_2 = BasicConv2d(8, 16, kernel_size=5, padding=2)

    self.branch3x3dbl_1 = BasicConv2d(in_channels, 8, kernel_size=1)
    self.branch3x3dbl_2 = BasicConv2d(8, 16, kernel_size=3, padding=1)
    self.branch3x3dbl_3 = BasicConv2d(16, 32, kernel_size=3, padding=1)

    self.branch_pool = BasicConv2d(in_channels, pool_features, kernel_size=1)

  def forward(self, x):
    branch1x1 = self.branch1x1(x)

    branch5x5 = self.branch5x5_1(x)
    branch5x5 = self.branch5x5_2(branch5x5)

    branch3x3dbl = self.branch3x3dbl_1(x)
    branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
    branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)

    branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
    branch_pool = self.branch_pool(branch_pool)

    outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
    return torch.cat(outputs, 1)

class Inception3(nn.Module):

  def __init__(self, num_classes=10, aux_logits=True, training=True, transform_input=True):
    super(Inception3, self).__init__()
    self.aux_logits = aux_logits
    self.training = training
    self.transform_input = transform_input
    self.Conv2d_4a_3x3 = BasicConv2d(3, 32, kernel_size=3,padding=1)
    self.Mixed_5b = InceptionA(32, pool_features=8)
    self.Mixed_5c = InceptionA(64, pool_features=72)
    self.Mixed_6a = InceptionB(128)
    self.Mixed_6b = InceptionC(256, channels_7x7=64)
    if aux_logits:
      self.AuxLogits = InceptionAux(512, num_classes)
    self.Mixed_7a = InceptionD(512)
    self.fc = nn.Linear(768, num_classes)

    for m in self.modules():
      if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        import scipy.stats as stats
        stddev = m.stddev if hasattr(m, 'stddev') else 0.1
        X = stats.truncnorm(-2, 2, scale=stddev)
        values = torch.as_tensor(X.rvs(m.weight.numel()), dtype=m.weight.dtype)
        values = values.view(m.weight.size())
        with torch.no_grad():
          m.weight.copy_(values)
      elif isinstance(m, nn.BatchNorm2d):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)

  def forward(self, x):
    global aux
    # print(x.shape)
    x = self.Conv2d_4a_3x3(x)
    x = self.Mixed_5b(x)
    x = self.Mixed_5c(x)
    x = self.Mixed_6a(x)
    x = self.Mixed_6b(x)
    if self.training and self.aux_logits:
      aux = self.AuxLogits(x)
    x = self.Mixed_7a(x)
    x = F.adaptive_avg_pool2d(x, (1, 1))
    x = F.dropout(x, training=self.training)
    x = torch.flatten(x, 1)
    x = self.fc(x)
    # print(x.shape)
    if self.training and self.aux_logits:
      return _InceptionOutputs(x, aux)
    return x


def inception_v3(pretrained=False, **kwargs):
  
  if pretrained:
    if 'transform_input' not in kwargs:
      kwargs['transform_input'] = True
    if 'aux_logits' in kwargs:
      original_aux_logits = kwargs['aux_logits']
      kwargs['aux_logits'] = True
    else:
      original_aux_logits = True
    model = Inception3(**kwargs)
    if not original_aux_logits:
      model.aux_logits = False
    return model

  return Inception3(**kwargs)

