import torch
from torch import nn
from torchvision.models.vgg import vgg16
import model

class GeneratorLoss(nn.Module):
    def __init__(self, inchannel=3):
        super(GeneratorLoss, self).__init__()
        vgg = vgg16(pretrained=True)
        layer1 = nn.Conv2d(inchannel, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        loss_network = nn.Sequential(layer1, *list(vgg.features)[1:31]).eval()
        for param in loss_network.parameters():
            param.requires_grad = False
        self.loss_network = loss_network
        self.mse_loss = nn.MSELoss()
        self.tv_loss = TVLoss()

    def forward(self, out_labels, out_images, target_images):
        # Adversarial Loss
        adversarial_loss = torch.mean(1 - out_labels)
        # Perception Loss
        perception_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))
        # Image Loss
        image_loss = self.mse_loss(out_images, target_images)
        # TV Loss
        tv_loss = self.tv_loss(out_images)
        return image_loss + 0.001 * adversarial_loss + 0.006 * perception_loss + 2e-8 * tv_loss


class TVLoss(nn.Module):
    def __init__(self, tv_loss_weight=1):
        super(TVLoss, self).__init__()
        self.tv_loss_weight = tv_loss_weight

    def forward(self, x):
        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h = self.tensor_size(x[:, :, 1:, :])
        count_w = self.tensor_size(x[:, :, :, 1:])
        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()
        return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size

    @staticmethod
    def tensor_size(t):
        return t.size()[1] * t.size()[2] * t.size()[3]

class PerHourglass(nn.Module):
    '''Hourglass model from Newell et al ECCV 2016'''
    def __init__(self):
        super(PerHourglass, self).__init__()
        block = model.Bottleneck
        num_stacks = 1
        num_blocks = 1
        num_classes = 6
        self.inplanes = 64
        self.num_feats = 128
        self.num_stacks = num_stacks

        self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=True)
        self.bn1 = nn.BatchNorm2d(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self._make_residual(block, self.inplanes, 1)
        self.layer2 = self._make_residual(block, self.inplanes, 1)
        self.layer3 = self._make_residual(block, self.num_feats, 1)
        self.maxpool = nn.MaxPool2d(2, stride=2)

        # build hourglass modules
        ch = self.num_feats*block.expansion
        hg, res, fc, score, fc_, score_ = [], [], [], [], [], []
        for i in range(num_stacks):
            hg.append(model.Hourglass(block, num_blocks, self.num_feats, 4))
            res.append(self._make_residual(block, self.num_feats, num_blocks))
            fc.append(self._make_fc(ch, ch))
            score.append(nn.Conv2d(ch, num_classes, kernel_size=1, bias=True))
            if i < num_stacks-1:
                fc_.append(nn.Conv2d(ch, ch, kernel_size=1, bias=True))
                score_.append(nn.Conv2d(num_classes, ch, kernel_size=1, bias=True))
        self.hg = nn.ModuleList(hg)
        self.res = nn.ModuleList(res)
        self.fc = nn.ModuleList(fc)
        self.score = nn.ModuleList(score)
        self.fc_ = nn.ModuleList(fc_)
        self.score_ = nn.ModuleList(score_)

    def _make_residual(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=True),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def _make_fc(self, inplanes, outplanes):
        bn = nn.BatchNorm2d(inplanes)
        conv = nn.Conv2d(inplanes, outplanes, kernel_size=1, bias=True)
        return nn.Sequential(
                conv,
                bn,
                self.relu,
            )

    def forward(self, x):
        out = []
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.maxpool(x)
        x = self.layer2(x)
        x = self.layer3(x)

        for i in range(self.num_stacks):
            y = self.hg[i](x)
            y = self.res[i](y)
            y = self.fc[i](y)
            score = self.score[i](y)
            out.append(score)
            if i < self.num_stacks-1:
                fc_ = self.fc_[i](y)
                score_ = self.score_[i](score)
                x = x + fc_ + score_

        return out



class PercepLoss(nn.Module):
    def __init__(self):
        super(PercepLoss, self).__init__()
        block = model.Bottleneck
        self.inplanes = 64
        self.num_feats = 128

        self.loss_net = PerHourglass()

        model_path = r'hg_best.pth.tar'
        best_model = torch.load(model_path)
        full_state_dict = best_model['state_dict']

        my_state = self.loss_net.state_dict()

        need_state_dict = dict()
        for k in full_state_dict.keys():
            if k.replace('module.','') in my_state.keys():
                need_state_dict[k.replace('module.','')] = full_state_dict[k]

        # print(my_state.keys())
        # print()
        # print()
        # print()
        # print([k.replace('module.','') for k in full_state_dict.keys()])

        self.loss_net.load_state_dict(need_state_dict)
        self.loss_net.cuda()
        # need_state_dict = {k:v for k, v in full_state_dict.keys(),full_state_dict.values() if k in my_state}
        self.mse_loss = nn.MSELoss()

    def _make_residual(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=True),
            )
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x, y):
        keypoint_loss = self.mse_loss(self.loss_net(x)[0], self.loss_net(y)[0])
        return keypoint_loss


if __name__ == "__main__":
    g_loss = PercepLoss()
    print(g_loss)

import torch
from torch import nn
from torchvision.models.vgg import vgg16
import model

class GeneratorLoss(nn.Module):
    def __init__(self, inchannel=3):
        super(GeneratorLoss, self).__init__()
        vgg = vgg16(pretrained=True)
        layer1 = nn.Conv2d(inchannel, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        loss_network = nn.Sequential(layer1, *list(vgg.features)[1:31]).eval()
        for param in loss_network.parameters():
            param.requires_grad = False
        self.loss_network = loss_network
        self.mse_loss = nn.MSELoss()
        self.tv_loss = TVLoss()

    def forward(self, out_labels, out_images, target_images):
        # Adversarial Loss
        adversarial_loss = torch.mean(1 - out_labels)
        # Perception Loss
        perception_loss = self.mse_loss(self.loss_network(out_images), self.loss_network(target_images))
        # Image Loss
        image_loss = self.mse_loss(out_images, target_images)
        # TV Loss
        tv_loss = self.tv_loss(out_images)
        return image_loss + 0.001 * adversarial_loss + 0.006 * perception_loss + 2e-8 * tv_loss


class TVLoss(nn.Module):
    def __init__(self, tv_loss_weight=1):
        super(TVLoss, self).__init__()
        self.tv_loss_weight = tv_loss_weight

    def forward(self, x):
        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h = self.tensor_size(x[:, :, 1:, :])
        count_w = self.tensor_size(x[:, :, :, 1:])
        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()
        return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size

    @staticmethod
    def tensor_size(t):
        return t.size()[1] * t.size()[2] * t.size()[3]

class PerHourglass(nn.Module):
    '''Hourglass model from Newell et al ECCV 2016'''
    def __init__(self):
        super(PerHourglass, self).__init__()
        block = model.Bottleneck
        num_stacks = 1
        num_blocks = 1
        num_classes = 6
        self.inplanes = 64
        self.num_feats = 128
        self.num_stacks = num_stacks

        self.conv1 = nn.Conv2d(1, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=True)
        self.bn1 = nn.BatchNorm2d(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self._make_residual(block, self.inplanes, 1)
        self.layer2 = self._make_residual(block, self.inplanes, 1)
        self.layer3 = self._make_residual(block, self.num_feats, 1)
        self.maxpool = nn.MaxPool2d(2, stride=2)

        # build hourglass modules
        ch = self.num_feats*block.expansion
        hg, res, fc, score, fc_, score_ = [], [], [], [], [], []
        for i in range(num_stacks):
            hg.append(model.Hourglass(block, num_blocks, self.num_feats, 4))
            res.append(self._make_residual(block, self.num_feats, num_blocks))
            fc.append(self._make_fc(ch, ch))
            score.append(nn.Conv2d(ch, num_classes, kernel_size=1, bias=True))
            if i < num_stacks-1:
                fc_.append(nn.Conv2d(ch, ch, kernel_size=1, bias=True))
                score_.append(nn.Conv2d(num_classes, ch, kernel_size=1, bias=True))
        self.hg = nn.ModuleList(hg)
        self.res = nn.ModuleList(res)
        self.fc = nn.ModuleList(fc)
        self.score = nn.ModuleList(score)
        self.fc_ = nn.ModuleList(fc_)
        self.score_ = nn.ModuleList(score_)

    def _make_residual(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=True),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def _make_fc(self, inplanes, outplanes):
        bn = nn.BatchNorm2d(inplanes)
        conv = nn.Conv2d(inplanes, outplanes, kernel_size=1, bias=True)
        return nn.Sequential(
                conv,
                bn,
                self.relu,
            )

    def forward(self, x):
        out = []
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.maxpool(x)
        x = self.layer2(x)
        x = self.layer3(x)

        for i in range(self.num_stacks):
            y = self.hg[i](x)
            y = self.res[i](y)
            y = self.fc[i](y)
            score = self.score[i](y)
            out.append(score)
            if i < self.num_stacks-1:
                fc_ = self.fc_[i](y)
                score_ = self.score_[i](score)
                x = x + fc_ + score_

        return out



class PercepLoss(nn.Module):
    def __init__(self):
        super(PercepLoss, self).__init__()
        block = model.Bottleneck
        self.inplanes = 64
        self.num_feats = 128

        self.loss_net = PerHourglass()

        model_path = r'hg_best.pth.tar'
        best_model = torch.load(model_path)
        full_state_dict = best_model['state_dict']

        my_state = self.loss_net.state_dict()

        need_state_dict = dict()
        for k in full_state_dict.keys():
            if k.replace('module.','') in my_state.keys():
                need_state_dict[k.replace('module.','')] = full_state_dict[k]

        # print(my_state.keys())
        # print()
        # print()
        # print()
        # print([k.replace('module.','') for k in full_state_dict.keys()])

        self.loss_net.load_state_dict(need_state_dict)
        self.loss_net.cuda()
        # need_state_dict = {k:v for k, v in full_state_dict.keys(),full_state_dict.values() if k in my_state}
        self.mse_loss = nn.MSELoss()

    def _make_residual(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=True),
            )
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x, y):
        keypoint_loss = self.mse_loss(self.loss_net(x)[0], self.loss_net(y)[0])
        return keypoint_loss


if __name__ == "__main__":
    g_loss = PercepLoss()
    print(g_loss)

# Adversarial Loss # Perception Loss # Image Loss # TV Loss Hourglass model from Newell et al ECCV 2016 # build hourglass modules # print(my_state.keys()) # print() # print() # print() # print([k.replace('module.','') for k in full_state_dict.keys()]) # need_state_dict = {k:v for k, v in full_state_dict.keys(),full_state_dict.values() if k in my_state}