#固定风格固定内容的普通迁移
#coding:utf-8

from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.optim as optim
from torchvision import models
from torchvision import transforms
import time
import hiddenlayer as h1
from skimage.io import imread



def load_image(img_path,max_size=400,shape=None):

    #读取图像并保证图像的高和宽在默认情况下都小于400

    image=Image.open(img_path).convert('RGB')
    #如果图像尺寸过大，就对图像进行尺寸转换
    if max(image.size)>max_size:
        size=max_size
    else:
        size=max(image.size)
    #如果指定了图像的尺寸，就将图像转化为shape指定的尺寸
    if shape is not None:
        size=shape
    #使用transforms将图像转化为张量，并进行标准化
    in_transform=transforms.Compose(
        [transforms.Resize(size),
         transforms.ToTensor(),
         transforms.Normalize((0.485,0.456,0.406),
                              (0.229,0.224,0.225))])
    image=in_transform(image)[:3,:,:].unsqueeze(dim=0)
    return image

def im_convert(tensor):

    #将[1,c,h,w]维度的张量转化为[h,w,c]的数组
    #因为张量进行了标准化，所以要进行标准化逆变换

    image=tensor.data.numpy().squeeze()
    image=image.transpose(1,2,0)
    image=image*np.array((0.229,0.224,0.225))+np.array((0.485,0.456,0.406))
    image=image.clip(0,1)
    return image

def get_features(image,model,layers=None):
    if layers is None:
        layers={'0':'conv1_1',
                '5':'conv2_1',
                '10': 'conv3_1',
                '19': 'conv4_1',
                '21': 'conv4_2',
                '28': 'conv5_1',}
    features={}
    x=image
    for name,layer in model._modules.items():
        x=layer(x)
        if name in layers:
            features[layers[name]]=x
    return features

def gram_matrix(tensor):
    _,d,h,w=tensor.size()
    tensor=tensor.view(d,h*w)
    gram=torch.mm(tensor,tensor.t())
    return gram

def gray(content_dir,content_gray):
    a = np.asarray(Image.open(content_dir).convert('L')).astype('float')
    depath = 10  # (0-100)
    grad = np.gradient(a)  # 取图像灰度的梯度值
    grad_x, grad_y = grad  # 分别取横纵图像梯度值
    grad_x = grad_x * depath / 100.
    grad_y = grad_y * depath / 100.
    A = np.sqrt(grad_x ** 2 + grad_y ** 2 + 1.)
    uni_x = grad_x / A
    uni_y = grad_y / A
    uni_z = 1. / A

    vec_el = np.pi / 2.2  # 光源的俯视角度，弧度值
    vec_az = np.pi / 4  # 光源的方位角度，弧度值
    dx = np.cos(vec_el) * np.cos(vec_az)  # 光源对 x 轴的影响
    dy = np.cos(vec_el) * np.sin(vec_az)  # 光源对 y 轴的影响
    dz = np.sin(vec_el)  # 光源对 z 轴的影响

    b = 255 * (dx * uni_x + dy * uni_y + dz * uni_z)  # 光源归一化
    b = b.clip(0, 255)

    im = Image.fromarray(b.astype('uint8'))  # 重构图像
    im.save(content_gray)
    return

def main(content_dir,style_dir,result_dir):
    vgg19=models.vgg19(pretrained=False)#读取已经预训练的VGG19网络,已下载在缓存中，pretrained设置为False
    vgg=vgg19.features
    #是否将VGG19的特征提取网络权重冻结，训练时不更新
    for param in vgg.parameters():
        param.requires_grad_(False)

    dot_index = content_dir.rfind(".")
    content_gray = content_dir[:dot_index] + "_gray" + content_dir[dot_index:]
    result5000_dir = result_dir + "5000.bmp"

    gray(content_dir,content_gray)

    pre_content=load_image(content_dir,max_size=400)
    content=load_image(content_dir,max_size=400)
    #print("content shape:",content.shape)
    style=load_image(style_dir,shape=content.shape[-2:])
    #print("style shape:",style.shape)

    fig,(ax1,ax2,ax3)=plt.subplots(1,3,figsize=(16,5))
    ax1.imshow(im_convert(pre_content))
    ax1.set_title("pre_content")
    ax2.imshow(im_convert(content))
    ax2.set_title("content")
    ax3.imshow(im_convert(style))
    ax3.set_title("style")
    plt.show()

    content_features=get_features(content,vgg)
    style_features=get_features(style,vgg)
    style_grams={layer:gram_matrix(style_features[layer]) for layer in style_features}
    target=content.clone().requires_grad_(True)

    style_weights={'conv1_1':1.,
                   'conv2_1':0.75,
                   'conv3_1':0.2,
                   'conv4_1':0.2,
                   'conv5_1':0.2}
    alpha=1
    beta=1e6
    content_weight=alpha
    style_weight=beta

    show_every=1000#每迭代1000次输出一个结果
    #将损失保存
    total_loss_all=[]
    content_loss_all=[]
    style_loss_all=[]
    #使用Adam优化器
    optimizer=optim.Adam([target],lr=0.0003)
    steps=10000#优化时迭代的次数
    t0=time.time()#计算需要的时间
    for ii in range(1,steps+1):
        #获取目标图像特征
        target_features=get_features(target,vgg)
        #计算内容损失
        content_loss=torch.mean((target_features['conv4_2']-content_features['conv4_2'])**2)
        #计算风格损失，并且初始化为0
        style_loss=0
        #将每个层的gram matrix损失相加
        for layer in style_weights:
            target_feature=target_features[layer]
            target_gram=gram_matrix(target_feature)
            _,d,h,w=target_feature.shape
            #获取风格图像在每层的风格的gram matrix
            style_gram=style_grams[layer]
            #计算要生成图像的风格和风格图像的风格之间的差异，每层都有一个权重
            layer_style_loss=style_weights[layer]*torch.mean((target_gram-style_gram)**2)
            #累加计算风格差异损失
            style_loss += layer_style_loss/(d*h*w)
        #计算一次迭代的总的损失，即内容损失和风格损失的加权和
        total_loss=content_weight*content_loss+style_weight*style_loss
        #保存三种损失大小
        content_loss_all.append(content_loss.item())
        style_loss_all.append(style_loss.item())
        total_loss_all.append(total_loss.item())
        #更新需要生成的目标图像
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        #输出每show_every次迭代后的生成图像
        if ii%show_every==0:
            print('Total loss:',total_loss.item())
            print('Use time:',(time.time()-t0)/3600,"hour")
            newim=im_convert(target)
            #显示图像
            #plt.imshow(newim)
            #plt.title("Iteration:"+str(ii)+"times")
            #plt.show()
            result=Image.fromarray((newim*255).astype(np.uint8))
            result.save(result_dir+str(ii)+".bmp")
    #损失可视化
    plt.figure(figsize=(12,4))
    plt.subplot(1,2,1)
    plt.plot(total_loss_all,"r",label="total_loss")
    plt.legend()
    plt.title("total loss")
    plt.subplot(1, 2, 2)
    plt.plot(content_loss_all, "g-", label="content_loss")
    plt.plot(style_loss_all, "b-", label="style_loss")
    plt.legend()
    plt.title("content and style loss")
    plt.show()

    content = load_image(content_dir, max_size=400)
    style = load_image(style_dir, shape=content.shape[-2:])
    result=load_image(result5000_dir,max_size=400)
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12, 5))
    ax1.imshow(im_convert(content))
    ax1.set_title("content")
    ax2.imshow(im_convert(style))
    ax2.set_title("style")
    ax3.imshow(im_convert(result))
    ax3.set_title("result")
    plt.show()