# -*- coding: utf-8 -*-
"""
@author RuoyuLiu
@Email lry199801@163.com
@这段说自己是author其实不太合适......这些网络架构有的我是根据人家论文里面的网络图仿真的，有的（比如ResNet-18这种烂大街网
  络）就直接从网上找来改了改参数就放进来了，其实与其说自己是author不如说自己是一个mover。感谢各位作者的无私开源，之后我
  会整理这些网络的思路来源，贴在下面。
"""
import torch
from torch import nn
from torch.nn import init
import torch.nn.functional as F


class HResNet(nn.Module):
    def __init__(self, num_of_bands, num_of_class, patch_size):
        super(HResNet, self).__init__()
        self.num_of_bands = num_of_bands
        self.num_of_class = num_of_class
        self.conv0 = nn.Conv2d(self.num_of_bands, 64, kernel_size=(3, 3), padding=(0, 0))
        self.bn1 = nn.BatchNorm2d(64)
        self.relu11 = nn.ReLU()
        self.conv11 = nn.Conv2d(64, 64, kernel_size=(3, 3), padding=(1, 1))
        self.relu12 = nn.ReLU()
        self.conv12 = nn.Conv2d(64, 64, kernel_size=(3, 3), padding=(1, 1))
        self.bn2 = nn.BatchNorm2d(64)
        self.relu21 = nn.ReLU()
        self.conv21 = nn.Conv2d(64, 64, kernel_size=(3, 3), padding=(1, 1))
        self.relu22 = nn.ReLU()
        self.conv22 = nn.Conv2d(64, 64, kernel_size=(3, 3), padding=(1, 1))
        self.avg_pool = nn.AvgPool2d((patch_size - 2, patch_size - 2))
        # self.dense = nn.Linear(64 * patch_size - 2, patch_size - 2)
        self.dense = nn.Linear(64, num_of_class)

    def forward(self, x):
        x1 = self.conv0(x)
        x0 = x1
        x1 = self.bn1(x1)
        x1 = self.relu11(x1)
        x1 = self.conv11(x1)
        x1 = self.relu12(x1)
        x1 = self.conv12(x1)
        x1 = x0 + x1
        x2 = self.bn2(x1)
        x2 = self.relu21(x2)
        x2 = self.conv21(x2)
        x2 = self.relu22(x2)
        x2 = self.conv22(x2)
        res = x1 + x2
        # res = x0 + x1
        res = self.avg_pool(res)
        res = res.contiguous().view(res.size(0), -1)
        res = self.dense(res)
        return res


class ResidualBlock(nn.Module):
    def __init__(self, inchannel, outchannel, stride=(1, 1)):
        super(ResidualBlock, self).__init__()
        self.left = nn.Sequential(
            nn.Conv2d(inchannel, outchannel, kernel_size=(3, 3), stride=stride, padding=(1, 1), bias=False),
            nn.BatchNorm2d(outchannel),
            nn.ReLU(inplace=True),
            nn.Conv2d(outchannel, outchannel, kernel_size=(3, 3), stride=stride, padding=(1, 1), bias=False),
            nn.BatchNorm2d(outchannel)
        )
        self.shortcut = nn.Sequential()
        if stride != 1 or inchannel != outchannel:
            self.shortcut = nn.Sequential(
                nn.Conv2d(inchannel, outchannel, kernel_size=(1, 1), stride=stride, bias=False),
                nn.BatchNorm2d(outchannel)
            )

    def forward(self, x):
        out = self.left(x)
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, ResidualBlock, num_of_bands, num_of_class):
        super(ResNet, self).__init__()
        self.inchannel = 64
        self.num_of_bands = num_of_bands
        self.num_of_class = num_of_class
        self.conv1 = nn.Sequential(
            nn.Conv2d(self.num_of_bands, 64, kernel_size=(3, 3), padding=(0, 0), bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(),
        )
        self.layer1 = self.make_layer(ResidualBlock, 64, 2, stride=1)
        self.layer2 = self.make_layer(ResidualBlock, 128, 2, stride=1)
        self.layer3 = self.make_layer(ResidualBlock, 256, 2, stride=1)
        self.layer4 = self.make_layer(ResidualBlock, 512, 2, stride=1)
        self.fc = nn.Linear(512*7*7, self.num_of_class)

    def make_layer(self, block, channels, num_blocks, stride):
        strides = [stride] + [1] * (num_blocks - 1)  # strides=[1,1]
        layers = []
        for stride in strides:
            layers.append(block(self.inchannel, channels, stride))
            self.inchannel = channels
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv1(x)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        # out = F.avg_pool2d(out, 4)
        out = out.contiguous().view(out.size(0), -1)
        out = self.fc(out)
        return out


def ResNet18(num_of_bands, num_of_class):
    return ResNet(ResidualBlock, num_of_bands, num_of_class)
