# -*- coding: utf-8 -*- 
""" 
@File : myTransfer.py.py 
@Author: csc
@Date : 2022/7/24
"""
import os

# os.environ['CUDA_VISIBLE_DEVICES'] = '5'

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

import torchvision
import torchvision.transforms as transforms
import torchvision.models as models

import random
import shutil
from glob import glob
from tqdm import tqdm
from PIL import Image
import imagehash
import copy

from utils import *
from models import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = 'cpu'


class MetaNet(nn.Module):
    def __init__(self, param_dict):
        super(MetaNet, self).__init__()
        self.param_num = len(param_dict)
        self.hidden = nn.Linear(1920, 128 * self.param_num)
        self.fc_dict = {}
        for i, (name, params) in enumerate(param_dict.items()):
            self.fc_dict[name] = i
            setattr(self, 'fc{}'.format(i + 1), nn.Linear(128, params))

    # ONNX 要求输出 tensor 或者 list，不能是 dict
    def forward(self, mean_std_features):
        hidden = F.relu(self.hidden(mean_std_features))
        filters = {}
        for name, i in self.fc_dict.items():
            fc = getattr(self, 'fc{}'.format(i + 1))
            filters[name] = fc(hidden[:, i * 128:(i + 1) * 128])
        return list(filters.values())

    def forward2(self, mean_std_features):
        hidden = F.relu(self.hidden(mean_std_features))
        filters = {}
        for name, i in self.fc_dict.items():
            fc = getattr(self, 'fc{}'.format(i + 1))
            filters[name] = fc(hidden[:, i * 128:(i + 1) * 128])
        return filters


backbone = models.vgg19(pretrained=True)
backbone = VGG19(backbone.features[:30]).to(device).eval()

base = 32

# 可视化
width = 256
data_transform = transforms.Compose([
    transforms.RandomResizedCrop(width, scale=(256 / 480, 1), ratio=(1, 1)),
    transforms.ToTensor(),
    tensor_normalizer
])

style_weight = 3e5
content_weight = 1
tv_weight = 1e-6
batch_size = 8

content_dataset = torchvision.datasets.ImageFolder('./models/COCO2014_1000/', transform=data_transform)
content_data_loader = torch.utils.data.DataLoader(content_dataset, batch_size=batch_size, shuffle=True)

# style buffer
buffer = {}


def style_transfer(style_img, content_img):
    transform_net = TransformNet(base, residuals='resnext').to(device)
    metanet = MetaNet(transform_net.get_param_dict()).to(device)
    # 每次都加载一次来实现类似深拷贝的效果
    transform_net.load_state_dict(torch.load('./models/metanet_base32_style300000.0_tv1e-06_tagnohvd_transform_net.pth'))
    metanet.load_state_dict(torch.load('./models/metanet_base32_style300000.0_tv1e-06_tagnohvd.pth'))

    trainable_params = {}
    trainable_param_shapes = {}
    for model in [backbone, transform_net, metanet]:
        for name, param in model.named_parameters():
            if param.requires_grad:
                trainable_params[name] = param
                trainable_param_shapes[name] = param.shape

    optimizer = optim.Adam(trainable_params.values(), 1e-3)

    # hash
    key = imagehash.phash(Image.fromarray(style_img), hash_size=8, highfreq_factor=4)

    style_img = image_to_tensor(style_img, 256)
    content_img = image_to_tensor(content_img, 256)
    style_image = style_img.to(device)
    content_image = content_img.to(device)

    if key in buffer:
        _transform_net = buffer[key]
        transformed_images = _transform_net(content_image)
    else:
        style_features = backbone(style_image)
        style_mean_std = mean_std(style_features)

        n_batch = 20
        with tqdm(enumerate(content_data_loader), total=n_batch) as pbar:
            for batch, (content_images, _) in pbar:
                x = content_images.cpu().numpy()
                if (x.min(-1).min(-1) == x.max(-1).max(-1)).any():
                    continue

                optimizer.zero_grad()

                # 使用风格图像生成风格模型
                weights = metanet.forward2(mean_std(style_features))
                transform_net.set_weights(weights, 0)

                # 使用风格模型预测风格迁移图像
                content_images = content_images.to(device)
                transformed_images = transform_net(content_images)

                # 使用 vgg19 计算特征
                content_features = backbone(content_images)
                transformed_features = backbone(transformed_images)
                transformed_mean_std = mean_std(transformed_features)

                # content loss
                content_loss = content_weight * F.mse_loss(transformed_features[2], content_features[2])

                # style loss
                style_loss = style_weight * F.mse_loss(transformed_mean_std,
                                                       style_mean_std.expand_as(transformed_mean_std))

                # total variation loss
                y = transformed_images
                tv_loss = tv_weight * (torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) +
                                       torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :])))

                # 求和
                loss = content_loss + style_loss + tv_loss

                loss.backward()
                optimizer.step()

                if batch > n_batch:
                    break

        buffer[key] = transform_net
        transformed_images = transform_net(content_image)

    transformed_images_vis = torch.cat([x for x in transformed_images], dim=-1)

    return recover_image(transformed_images_vis)