"""
Code borrowed from https://gist.github.com/alper111/8233cdb0414b4cb5853f2f730ab95a49#file-vgg_perceptual_loss-py-L5
"""
import torch
import torchvision
from models.vggface import VGGFaceFeats


def cos_loss(fi, ft):
    return 1 - torch.nn.functional.cosine_similarity(fi, ft).mean()


class VGGPerceptualLoss(torch.nn.Module):
    def __init__(self, resize=False):
        super(VGGPerceptualLoss, self).__init__()
        blocks = []
        blocks.append(torchvision.models.vgg16(pretrained=True).features[:4].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[4:9].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[9:16].eval())
        blocks.append(torchvision.models.vgg16(pretrained=True).features[16:23].eval())
        for bl in blocks:
            for p in bl:
                p.requires_grad = False
        self.blocks = torch.nn.ModuleList(blocks)
        self.transform = torch.nn.functional.interpolate
        self.mean = torch.nn.Parameter(torch.tensor([0.485, 0.456, 0.406]).view(1,3,1,1))
        self.std = torch.nn.Parameter(torch.tensor([0.229, 0.224, 0.225]).view(1,3,1,1))
        self.resize = resize

    def forward(self, input, target, max_layer=4, cos_dist: bool = False):
        target = (target + 1) * 0.5
        input = (input + 1) * 0.5

        if input.shape[1] != 3:
            input = input.repeat(1, 3, 1, 1)
            target = target.repeat(1, 3, 1, 1)
        input = (input-self.mean) / self.std
        target = (target-self.mean) / self.std
        if self.resize:
            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)
        x = input
        y = target
        loss = 0.0
        loss_func = cos_loss if cos_dist else torch.nn.functional.l1_loss
        for bi, block in enumerate(self.blocks[:max_layer]):
            x = block(x)
            y = block(y)
            loss += loss_func(x, y.detach())
        return loss


class VGGFacePerceptualLoss(torch.nn.Module):
    def __init__(self, weight_path: str = "checkpoint/vgg_face_dag.pt", resize: bool = False):
        super().__init__()
        self.vgg = VGGFaceFeats()
        self.vgg.load_state_dict(torch.load(weight_path))

        mean = torch.tensor(self.vgg.meta["mean"]).view(1, 3, 1, 1) / 255.0
        self.register_buffer("mean", mean)

        self.transform = torch.nn.functional.interpolate
        self.resize = resize

    def forward(self, input, target, max_layer: int = 4, cos_dist: bool = False):
        target = (target + 1) * 0.5
        input = (input + 1) * 0.5

        # preprocessing
        if input.shape[1] != 3:
            input = input.repeat(1, 3, 1, 1)
            target = target.repeat(1, 3, 1, 1)
        input = input - self.mean
        target = target - self.mean
        if self.resize:
            input = self.transform(input, mode='bilinear', size=(224, 224), align_corners=False)
            target = self.transform(target, mode='bilinear', size=(224, 224), align_corners=False)

        input_feats = self.vgg(input)
        target_feats = self.vgg(target)

        loss_func = cos_loss if cos_dist else torch.nn.functional.l1_loss
        # calc perceptual loss
        loss = 0.0
        for fi, ft in zip(input_feats[:max_layer], target_feats[:max_layer]):
            loss = loss + loss_func(fi, ft.detach())
        return loss


class PerceptualLoss(torch.nn.Module):
    def __init__(
            self, lambda_vggface: float = 0.025 / 0.15, lambda_vgg: float = 1,  eps: float = 1e-8, cos_dist: bool = False
    ):
        super().__init__()
        self.register_buffer("lambda_vggface", torch.tensor(lambda_vggface))
        self.register_buffer("lambda_vgg", torch.tensor(lambda_vgg))
        self.cos_dist = cos_dist

        if lambda_vgg > eps:
            self.vgg = VGGPerceptualLoss()
        if lambda_vggface > eps:
            self.vggface = VGGFacePerceptualLoss()

    def forward(self, input, target, eps=1e-8, use_vggface: bool = True, use_vgg=True, max_vgg_layer=4):
        loss = 0.0
        if self.lambda_vgg > eps and use_vgg:
            loss = loss + self.lambda_vgg * self.vgg(input, target, max_layer=max_vgg_layer)
        if self.lambda_vggface > eps and use_vggface:
            loss = loss + self.lambda_vggface * self.vggface(input, target, cos_dist=self.cos_dist)
        return loss