import cma
import tqdm
from pytorch_pretrained_biggan import BigGAN, truncated_noise_sample, BigGANConfig
import torch
import torchvision
from torchvision.transforms import ToPILImage
from torchvision.utils import make_grid
from torch.optim import SGD, Adam
import os
import sys
import numpy as np
import matplotlib.pylab as plt
import pandas as pd
from os.path import join
from imageio import imread
from scipy.linalg import block_diag
#%%
def get_BigGAN(version="biggan-deep-256"):
    cache_path = "/scratch/binxu/torch/"
    cfg = BigGANConfig.from_json_file(join(cache_path, "%s-config.json" % version))
    BGAN = BigGAN(cfg)
    BGAN.load_state_dict(torch.load(join(cache_path, "%s-pytorch_model.bin" % version)))
    return BGAN
from argparse import ArgumentParser
parser = ArgumentParser()
parser.add_argument("--ImDist", type=str, default="squeeze", help="Network model to use for Image distance computation")
parser.add_argument("--dataset", type=str, default="ImageNet", help="")
parser.add_argument("--imgidx", type=int, nargs='+', default=[200, 220], help="Network model to use for Image distance "
                                                                   "computation")
parser.add_argument("--cmasteps", type=int, default=10, help="")
parser.add_argument("--gradsteps", type=int, default=10, help="")
parser.add_argument("--finalgradsteps", type=int, default=500, help="")
parser.add_argument("--CMApostGrad", type=bool, default=True, help="")
parser.add_argument("--basis", type=str, default="all", help="")
args = parser.parse_args()
#%%
if sys.platform == "linux":
    BGAN = get_BigGAN()
    sys.path.append(r"/home/binxu/PerceptualSimilarity")
    Hpath = r"/scratch/binxu/GAN_hessian/BigGAN/summary/H_avg_1000cls.npz"
    if args.dataset == "BigGAN_rnd":
        imgfolder = r"/scratch/binxu/Datasets/BigGAN_rnd"
        savedir = r"/scratch/binxu/GAN_invert/BasinCMA/BigGAN_rnd"
    elif args.dataset == "ImageNet":
        imgfolder = r"/scratch/binxu/Datasets/ImageTranslation/GAN_real/B/train"
        savedir = r"/scratch/binxu/GAN_invert/BasinCMA/ImageNet"
else:
    BGAN = BigGAN.from_pretrained("biggan-deep-256")
    sys.path.append(r"D:\Github\PerceptualSimilarity")
    sys.path.append(r"E:\Github_Projects\PerceptualSimilarity")
    Hpath = r"E:\OneDrive - Washington University in St. Louis\Hessian_summary\BigGAN\H_avg_1000cls.npz"
    if args.dataset == "BigGAN_rnd":
        imgfolder = r"E:\Cluster_Backup\Datasets\BigGAN_rnd"
        savedir = r"E:\OneDrive - Washington University in St. Louis\BigGAN_invert\BasinCMA\BigGAN_rnd"
    elif args.dataset == "ImageNet":
        imgfolder = r"E:\Cluster_Backup\Datasets\ImageTranslation\GAN_real\B\train"
        savedir = r"E:\OneDrive - Washington University in St. Louis\BigGAN_invert\BasinCMA\ImageNet"
    # savedir = r"E:\Cluster_Backup\BigGAN_invert\ImageNet"

#%%
import models  # from PerceptualSimilarity folder
ImDist = models.PerceptualLoss(model='net-lin', net=args.ImDist, use_gpu=1, gpu_ids=[0])
for param in ImDist.parameters():
    param.requires_grad_(False)
def L1loss(target, img):
    return (img - target).abs().sum(axis=1).mean(axis=1)

# alpha = 5 # relative weight
BGAN.cuda()
BGAN.eval()
for param in BGAN.parameters():
    param.requires_grad_(False)
#%% Load the precomputed Hessian and Form them as basis matrix
data = np.load(Hpath)
evc_clas = torch.from_numpy(data['eigvects_clas_avg'])#.cuda()
evc_nois = torch.from_numpy(data['eigvects_nois_avg'])#.cuda()
evc_all = torch.from_numpy(data['eigvects_avg']).cuda()
evc_sep = torch.from_numpy(block_diag(data['eigvects_nois_avg'], data['eigvects_clas_avg'])).cuda()
evc_none = torch.eye(256).cuda()
#%%
"""Main function to do inversion."""
def BasinCMA(target_tsr, cmasteps=30, gradsteps=40, finalgrad=500, batch_size=4,
             basis="all", CMApostAdam=False, RND=None, savedir=savedir, imgnm=""):
    Record = {"L1cma": [],"dsimcma": [], "L1Adam": [], "dsimAdam": [], "L1refine":[], "dsimrefine":[], "classnorm":[], "noisenorm":[]}
    # choose the basis vector to use in Adam optimization
    basisvec = {"all": evc_all, "sep": evc_sep, "none": evc_none}[basis]
    fixnoise = truncated_noise_sample(1, 128)
    optim_noise = cma.CMAEvolutionStrategy(fixnoise, 0.4)#0.4)  # 0.2
    optim_class = cma.CMAEvolutionStrategy(128 * [0.0], 0.2)#0.12)  # 0.06
    # noise_vec = torch.from_numpy(fixnoise)
    RND = np.random.randint(1E6) if RND is None else RND
    # Outer Loop: CMA optimization of initial points
    for i in tqdm.trange(cmasteps, desc="CMA steps"):
        class_codes = optim_class.ask()
        noise_codes = optim_noise.ask()
        # TODO: boundary handling by projection in code space
        # Evaluate the cma proposed codes `latent_code` at first.
        codes_tsr = torch.from_numpy(np.array(class_codes)).float()
        noise_tsr = torch.from_numpy(np.array(noise_codes)).float()
        latent_code = torch.cat((noise_tsr, codes_tsr), dim=1).cuda()  # this initialize inner loop
        with torch.no_grad():
            imgs = BGAN.generator(latent_code, 0.7)
            imgs = (imgs + 1.0) / 2.0
            dsims = ImDist(imgs, target_tsr).squeeze()
            L1dsim = (imgs - target_tsr).abs().mean([1, 2, 3])
        scores = dsims.detach().cpu().numpy()
        L1score = L1dsim.detach().cpu().numpy()
        print("step %d pre-ADAM dsim %.3f L1 %.3f (norm %.2f noise norm %.2f)" % (
            i, scores.mean(), L1score.mean(), codes_tsr.norm(dim=1).mean(), noise_tsr.norm(dim=1).mean()))
        Record["L1cma"].append(L1score)
        Record["dsimcma"].append(scores)
        # Inner loop: ADAM optimization from the cma proposed codes `latent_code` batch by batch
        codes_post = np.zeros_like(np.hstack((noise_codes, class_codes)))
        scores_post = np.zeros_like(scores)
        L1score_post = np.zeros_like(L1score)
        if gradsteps > 0:
            csr = 0
            while csr < len(class_codes):  # pbar = tqdm.tqdm(total=len(codes), initial=csr, desc="batchs")
                csr_end = min(len(class_codes), csr + batch_size)
                # codes_batch = codes_tsr[csr:csr_end, :].detach().clone().requires_grad_(True)
                coef_batch = (latent_code[csr:csr_end, :] @ basisvec).detach().clone().requires_grad_(True)
                optim = Adam([coef_batch], lr=0.05, )
                for step in range(gradsteps):  # tqdm.trange(gradsteps, desc="ADAM steps"):
                    optim.zero_grad()
                    latent_batch = coef_batch @ basisvec.T
                    imgs = BGAN.generator(latent_batch, 0.7)
                    imgs = (imgs + 1.0) / 2.0
                    dsims = ImDist(imgs, target_tsr).squeeze()
                    L1dsim = (imgs - target_tsr).abs().mean([1, 2, 3])
                    loss = (dsims + L1dsim).sum()
                    loss.backward()
                    optim.step()
                    if (step + 1) % 10 == 0:
                        print("step %d dsim %.3f L1 %.3f" % (step, dsims.mean().item(), L1dsim.mean().item(),))
                code_batch = (coef_batch @ evc_all.T).detach().cpu().numpy()
                scores_batch = dsims.detach().cpu().numpy()
                L1score_batch = L1dsim.detach().cpu().numpy()
                codes_post[csr:csr_end, :] = code_batch
                scores_post[csr:csr_end] = scores_batch
                L1score_post[csr:csr_end] = L1score_batch
                csr = csr_end
            # Use the ADAM optimized scores as utility for `latent_code` and do cma update
            print("step %d post-ADAM dsim %.3f L1 %.3f (norm %.2f, norm %.2f)" % (
                i, scores_post.mean(), L1score_post.mean(),
                np.linalg.norm(codes_post[:, 128:], axis=1).mean(),
                np.linalg.norm(codes_post[:, :128], axis=1).mean()))
        else:  # if no grad step is performed
            scores_post = scores
            L1score_post = L1score
            codes_post = np.hstack((noise_codes, class_codes))
        Record["L1Adam"].append(L1score_post)
        Record["dsimAdam"].append(scores_post)
        if CMApostAdam:
            optim_class.tell(codes_post[:, :128], scores_post + L1score_post)
            optim_noise.tell(codes_post[:, 128:], scores_post + L1score_post)
        else:
            optim_class.tell(class_codes, scores_post + L1score_post)
            optim_noise.tell(noise_codes, scores_post + L1score_post)

    # Sort the scores and find the codes with the least scores to be final refined
    idx = np.argsort((L1score_post + scores_post))
    codes_batch = torch.from_numpy(codes_post[idx[:4]]).float().cuda()
    fit_img = BGAN.generator(codes_batch, 0.7)
    fit_img = (fit_img + 1) / 2.0
    CMAfitimg = ToPILImage()(make_grid(torch.cat((fit_img, target_tsr)).cpu()))
    CMAfitimg.save(join(savedir, "%s_CMA_final%06d.jpg" % (imgnm, RND)))
    CMAfitimg.show()
    # Linear Reparametrization using basisvec
    coef_batch = (codes_batch @ basisvec).detach().clone().requires_grad_(True)
    optim = Adam([coef_batch], lr=0.05, )
    scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=200, gamma=0.5)
    for step in range(finalgrad):  # tqdm.trange(gradsteps, desc="ADAM steps"):
        optim.zero_grad()
        # latent_code = torch.cat((noise_vec.repeat(codes_batch.shape[0], 1), codes_batch), dim=1).cuda()
        latent_code = coef_batch @ basisvec.T  #evc_all.T
        imgs = BGAN.generator(latent_code, 0.7)
        imgs = (imgs + 1.0) / 2.0
        dsims = ImDist(imgs, target_tsr).squeeze()
        L1dsim = (imgs - target_tsr).abs().mean([1, 2, 3])
        loss = (dsims + L1dsim).sum()
        loss.backward()
        optim.step()
        scheduler.step()
        Record["L1refine"].append(L1dsim.detach().cpu().numpy())
        Record["dsimrefine"].append(dsims.detach().cpu().numpy())
        Record["classnorm"].append(latent_code[:, 128:].norm(dim=1).detach().cpu().numpy())
        Record["noisenorm"].append(latent_code[:, :128].norm(dim=1).detach().cpu().numpy())
        if (step + 1) % 10 == 0:
            print("step %d dsim %.3f L1 %.3f (norm %.2f)" % (
                step, dsims.mean().item(), L1dsim.mean().item(), latent_code.norm(dim=1).mean().item()))
    scores_final = dsims.detach().cpu().numpy()
    L1score_final = L1dsim.detach().cpu().numpy()
    finalimg = ToPILImage()(make_grid(torch.cat((imgs, target_tsr)).cpu()))
    finalimg.save(join(savedir, "%srefinefinal%06d.jpg" % (imgnm, RND)))
    finalimg.show()

    fig = visualize_optim(Record, titlestr="cmasteps %d gradsteps %d refinesteps %d Hbasis %s, CMApostAdam %d"%(cmasteps, gradsteps, finalgrad, basis, CMApostAdam))
    fig.savefig(join(savedir, "%straj_H%s%s_%d_dsim_%.3f_L1_%.3f.jpg" % (imgnm, basis, "_postAdam" if CMApostAdam else "",
                                                                        RND, scores_final.min(), L1score_final.min())))
    np.savez(join(savedir, "%soptim_data_%d.npz" % (imgnm, RND)), codes=latent_code.cpu().detach().numpy(),
             Record=Record, dsims=scores_final, L1dsims=L1score_final)

    return imgs.cpu().detach().numpy(), latent_code.cpu().detach().numpy(), scores_final, L1score_final, Record

def visualize_optim(Record, titlestr="", savestr=""):
    fig, ax = plt.subplots()
    L1_cma_tr = np.array(Record["L1cma"])
    dsim_cma_tr = np.array(Record["dsimcma"])
    L1_adam_tr = np.array(Record["L1Adam"])
    dsim_adam_tr = np.array(Record["dsimAdam"])
    dsim_tr = np.array(Record["dsimrefine"])
    L1_tr = np.array(Record["L1refine"])
    nos_norm = np.array(Record["noisenorm"])
    cls_norm = np.array(Record["classnorm"])
    cma_steps, popsize = L1_cma_tr.shape
    xticks_arr = np.arange(-cma_steps*10, 0, 10)[:,np.newaxis].repeat(popsize, 1)
    ax.scatter(xticks_arr, L1_cma_tr, color="blue", s=15)
    ax.scatter(xticks_arr, dsim_cma_tr, color="green", s=15)
    ax.scatter(xticks_arr+0.2, L1_adam_tr, color="blue", s=15)
    ax.scatter(xticks_arr+0.2, dsim_adam_tr, color="green", s=15)
    ax.plot(dsim_tr, label="dsim", color="green", alpha=0.7)
    ax.plot(L1_tr, label="L1", color="blue", alpha=0.7)
    ax.set_ylabel("Image Dissimilarity", color="blue", fontsize=14)
    plt.legend()
    ax2 = ax.twinx()
    ax2.plot(nos_norm, color="orange", label="noise", alpha=0.7)
    ax2.plot(cls_norm, color="magenta", label="class", alpha=0.7)
    ax2.set_ylabel("L2 Norm", color="red", fontsize=14)
    plt.legend()
    plt.title(titlestr)
    plt.show()
    return fig
#%%
# from imageio import imread
# target = imread("block042_thread000_gen_gen041_001030.bmp")
# target_tsr = torch.from_numpy(target / 255.0).permute([2, 0, 1]).unsqueeze(0)
# target_tsr = target_tsr.float().cuda()
#%%
# savedir = r"E:\OneDrive - Washington University in St. Louis\BigGAN_invert\BasinCMA\ImageNet"
# imgfolder = r"E:\Datasets\ImageNet\resize"
# Append the name of parameters to the folder
expdir = join(savedir, "CMA%dAdam%dFinal%d%s_%s"%(args.cmasteps, args.gradsteps, args.finalgradsteps,
                                                  "_postAdam" if args.CMApostGrad else "", args.basis))
os.makedirs(expdir, exist_ok=True)
idmin, idmax = args.imgidx
for imgid in tqdm.trange(idmin, idmax):
    print("Processing image %d" %imgid)
    if args.dataset == "ImageNet":
        imgnm = "val_crop_%08d" % imgid
        img = imread(join(imgfolder, "val_crop_%08d.JPEG"%imgid))
    elif args.dataset == "BigGAN_rnd":
        imgnm = "BigGAN_rnd_%04d" % imgid
        img = imread(join(imgfolder, "rnd_%04d.jpg" % imgid))
    target_tsr = torch.from_numpy(img / 255.0).permute([2, 0, 1]).unsqueeze(0)
    target_tsr = target_tsr.float().cuda()
    imgs_np, codes_np, dsims, L1dsims, Record = BasinCMA(target_tsr, cmasteps=args.cmasteps, gradsteps=args.gradsteps, finalgrad=args.finalgradsteps,
                                                         batch_size=4, basis=args.basis, CMApostAdam=args.CMApostGrad,
                                                         imgnm=imgnm, savedir=expdir)
    # imgs_np, codes_np, dsims, L1dsims, Record = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500,
    #                                                      batch_size=4, basis="all", CMApostAdam=True)

#%%
# imgfolder = r"E:\Cluster_Backup\Datasets\BigGAN_rnd"
# G = BigGAN_wrapper(BGAN)
# #%%
# from imageio import imwrite
# imgnum = 200
# noise_samps = truncated_noise_sample(250, 128)
# class_samps = np.concatenate((0.05 * np.random.randn(50, 128),  0.1 * np.random.randn(50, 128),
#                 0.15 * np.random.randn(50, 128),  0.2 * np.random.randn(50, 128),  0.25 * np.random.randn(50, 128)),
#                              axis=0)
# codes_all_arr = np.hstack((noise_samps, class_samps))
# #%%
# img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=15)
# npimgs = img_all.permute([2,3,1,0]).numpy()
# for imgi in range(npimgs.shape[-1]):  imwrite(join(imgfolder, "rnd_%04d.jpg"%imgi), np.uint8(npimgs[:,:,:,imgi]*255))
# np.savez(join(imgfolder, "codes_all.npz"), codes_all=codes_all_arr)
#%%
# L1_cma_tr = np.array(Record["L1cma"])
# dsim_cma_tr = np.array(Record["dsimcma"])
# L1_adam_tr = np.array(Record["L1Adam"])
# dsim_adam_tr = np.array(Record["dsimAdam"])
# dsim_tr = np.array(Record["dsimrefine"])
# L1_tr = np.array(Record["L1refine"])
# nos_norm = np.array(Record["noisenorm"])
# cls_norm = np.array(Record["classnorm"])
#
# cma_steps, popsize = L1_cma_tr.shape
# xticks_arr = np.arange(-cma_steps, 0)[:, np.newaxis].repeat(popsize,1)
# plt.scatter(xticks_arr, L1_cma_tr, color="blue")
# plt.scatter(xticks_arr, dsim_cma_tr, color="green")
# plt.scatter(xticks_arr+0.2, L1_adam_tr, color="blue")
# plt.scatter(xticks_arr+0.2, dsim_adam_tr, color="green")
# plt.show()


#%%
# imgs_np, codes_np, dsims, L1dsims, Record = BasinCMA(target_tsr, cmasteps=50, gradsteps=0, finalgrad=400, batch_size=4, basis="all")
# #%%
# imgs_np2, codes_np2, dsims2, L1dsims2, Record2 = BasinCMA(target_tsr, cmasteps=50, gradsteps=0, finalgrad=400, batch_size=4, basis="all")
# #%%
# imgs_np3, codes_np3, dsims3, L1dsims3, Record3 = BasinCMA(target_tsr, cmasteps=10, gradsteps=20, finalgrad=400, batch_size=4, basis="all")
# #%%
# imgs_np4, codes_np4, dsims4, L1dsims4, Record4 = BasinCMA(target_tsr, cmasteps=50, gradsteps=0, finalgrad=300, batch_size=4, basis="all")
#
# #%%
# RND = np.random.randint(1E6)
# imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=5, gradsteps=5, finalgrad=400, batch_size=4, basis="all", CMApostAdam=True,  RND=RND)
# fig = visualize_optim(Record5, titlestr="cmasteps5 gradsteps5 Hbasis all, CMApostAdam")
# fig.savefig(join(savedir, "traj_%d_dsim_%.3f_L1_%.3f.jpg"%(RND, dsims5.min(), L1dsims5.min())))
# np.savez(join(savedir, "optim_data_%d.npz"%RND), codes=codes_np5, Record=Record5, dsims=dsims5, L1dsims=L1dsims5)
# #%%
# RND = np.random.randint(1E6)
# imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="all", CMApostAdam=True,  RND=RND)
# fig = visualize_optim(Record5, titlestr="cmasteps5 gradsteps5 Hbasis all, CMApostAdam")
# fig.savefig(join(savedir, "traj_Hall_%d_dsim_%.3f_L1_%.3f.jpg"%(RND, dsims5.min(), L1dsims5.min())))
# np.savez(join(savedir, "optim_data_%d.npz"%RND), codes=codes_np5, Record=Record5, dsims=dsims5, L1dsims=L1dsims5)
# #%%
# RND = np.random.randint(1E6)
# imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="all", CMApostAdam=True,  RND=RND)
# fig = visualize_optim(Record5, titlestr="cmasteps5 gradsteps5 Hbasis all, CMApostAdam")
# fig.savefig(join(savedir, "traj_Hall_%d_dsim_%.3f_L1_%.3f.jpg"%(RND, dsims5.min(), L1dsims5.min())))
# np.savez(join(savedir, "optim_data_%d.npz"%RND), codes=codes_np5, Record=Record5, dsims=dsims5, L1dsims=L1dsims5)
# #%% This is pretty gfood
# imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=400, batch_size=4, basis="all", CMApostAdam=True)
# #%%
# imgs_np6, codes_np6, dsims6, L1dsims6, Record6 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="none")
# #%%
#
# from time import time
# T0 = time()
# for trial in range(5):
#     imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="all", CMApostAdam=True)
#
#     imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="sep", CMApostAdam=True)
#
#     imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="none", CMApostAdam=True)
#
#     imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="all", CMApostAdam=False)
#
#     imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="sep", CMApostAdam=False)
#
#     imgs_np5, codes_np5, dsims5, L1dsims5, Record5 = BasinCMA(target_tsr, cmasteps=10, gradsteps=10, finalgrad=500, batch_size=4, basis="none", CMApostAdam=False)
#     print(time() - T0) # 5 hours
