# coding=utf-8
import os
# import moxing as mox
import numpy as np
import cv2 as cv
import mindspore
from mindspore import nn, Tensor, Parameter, context
from mindspore import load_checkpoint, load_param_into_net, save_checkpoint
from vgg import Vgg19
from method import create_dataset, load_image, gram_matrix, Optim_Loss, save_img, load_parameters
# 设置为动态图模式
import argparse
# parser = argparse.ArgumentParser(description='Training img transfer')
# parser.add_argument('--device_target', type=str, default="Ascend", choices=['Ascend', 'GPU', 'CPU'])
# parser.add_argument('--data_url', required=True, default=None)
# parser.add_argument('--train_url', required=True, default=None)
# # parser.add_argument("model_url", type=str, default=None, help="pretrained checkpoint file path of vgg.")
# args = parser.parse_known_args()[0]
# context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)

# if not os.path.exists("./img_data"):
#     os.mkdir("./img_data")
#     mox.file.copy_parallel(args.data_url, "./img_data")
#
# if not os.path.exists('./output_img'):
#     os.mkdir("./output_img")

vgg = Vgg19()
param_dict = load_parameters("vgg19.ckpt")
print("ckpt load success")
load_param_into_net(vgg, param_dict)
print("param load success")
vgg.set_train(False)   # 评估模式


def load_img_ckpt(path):
    return load_parameters(path)


def optimize_picture(content_path, style_path, iter_num, learning_rate, output_path, mode):
    if mode == 0:
        optimize_img = load_image(content_path, 400)
        content_maps = vgg(Tensor(optimize_img))
        optimizing_param = Parameter(Tensor(optimize_img, mindspore.float32), "optimizing_img")
    else:
        optimize_img = load_img_ckpt(content_path)["img"]
        content_maps = vgg(Tensor(optimize_img))
        optimizing_param = Parameter(Tensor(optimize_img, mindspore.float32), "optimizing_img")
    style_img = load_image(style_path, 400)
    # 特征图提取（6层）
    content_maps = content_maps[4].squeeze(axis=0)    # 内容图片vgg19提取的特征
    style_maps = list(vgg(Tensor(style_img)))    # 风格图片的特征
    style_maps.pop(4)      # 剔除掉不需要的map
    style_gram = list(gram_matrix(x) for x in style_maps)
    # 要优化的图片

    num_of_iterations = iter_num + 1    # 迭代的次数
    loss_opt = Optim_Loss(vgg, [content_maps, style_gram, optimizing_param])
    cosine_decay_lr = nn.CosineDecayLR(0.05, 2.0, 2000)
    optimizer_n = nn.Adam([optimizing_param], learning_rate=cosine_decay_lr)  # 损失优化方法
    # optimizer = nn.Adam([optimizing_param], learning_rate=learning_rate)   # 损失优化方法
    train_net = nn.TrainOneStepCell(loss_opt, optimizer_n)
    train_net.set_train()
    for i in range(num_of_iterations):
        if i % 50 == 0:
            save_img(i, optimizer_n, output_path)
        loss = train_net()
        print(str(content_path[-13:-4]) + "_to_" + str(style_path[-8:-4]) + ":iteration_", str(i), ":  loss", loss)
    final = optimizer_n.parameters
    save_checkpoint([{"name": "img", "data": final[0]}], output_path+"/img_1.ckpt")


optimize_picture("output_img/iter_19/img.ckpt",
                 "transfer_img/trans_6/sty_fam_1.jpg", 2000, 0.05, "output_img/iter_19", 1)

# img = cv.imread("transfer_img/trans_5/style_img.jpg")[:, :, ::-1]
# print(img.shape)