<reponame>Animadversio/Visual_Neuro_InSilico_Exp<gh_stars>1-10
# -*- coding: utf-8 -*-
"""
Created on Wed Jul 22 19:08:54 2020
Modified Oct. 14th

@author: <NAME>

Find important Nuisanced + Class transformations in Noise + Class space for a BigGAN evolved image. 
"""

# backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_integrated\2020-06-01-09-46-37"
# Put the backup folder and the thread to analyze here 
#backup_dir = r"C:\Users\Poncelab-ML2a\Documents\monkeylogic2\generate_BigGAN\2020-07-22-10-14-22"
# backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_BigGAN\2020-08-06-10-18-55"#2020-08-04-09-54-25"#
backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_BigGAN\2020-11-06-11-53-09"
threadid = 1

score_rank_avg = False  # If True, it will try to read "scores_record.mat", from the backup folder and read "scores_record"
                        # Else, it will use the unweighted mean code of the last generation as the center vector. 
                        # Need to run the BigGAN postHoc Analysis to save the `scores_record` mat and use this flag

exact_distance = True   # Control if exact distance search is used or approximate heuristic rule is used.
target_distance = [0.16, 0.24, 0.32, 0.40] # 0.08, 
#target_distance = [0.09, 0.18, 0.27, 0.36, 0.45]  # if exact_distance is True it will search for images with these
                                                  # distance to reference image along each eigenvector.
                                                  
#%% Prepare the generator model and perceptual loss networks
from time import time
import os
from os.path import join
import sys
import cv2 # for video writing
if os.environ['COMPUTERNAME'] == 'PONCELAB-ML2B':
    Python_dir = r"C:\Users\<NAME>\Documents\Python"
elif os.environ['COMPUTERNAME'] == 'PONCELAB-ML2A':
    Python_dir = r"C:\Users\Poncelab-ML2a\Documents\Python"
elif os.environ['COMPUTERNAME'] == 'DESKTOP-MENSD6S':
    Python_dir = r"E:\Github_Projects"
elif os.environ['COMPUTERNAME'] == 'DESKTOP-9DDE2RH':
    Python_dir = r"D:\Github"

sys.path.append(join(Python_dir, "Visual_Neuro_InSilico_Exp"))
sys.path.append(join(Python_dir, "PerceptualSimilarity"))
import torch
import numpy as np
from pytorch_pretrained_biggan import BigGAN, truncated_noise_sample
from GAN_utils import upconvGAN
from hessian_eigenthings.lanczos import lanczos
from GAN_hvp_operator import GANHVPOperator, GANForwardHVPOperator, GANForwardMetricHVPOperator, \
    compute_hessian_eigenthings, get_full_hessian
from skimage.io import imsave, imsave
from torchvision.utils import make_grid
from build_montages import build_montages
from torchvision.transforms import ToPILImage, ToTensor
from IPython.display import clear_output
from hessian_eigenthings.utils import progress_bar
from tqdm import tqdm
T00 = time()
import lpips
# import models # from PerceptualSimilarity folder
# ImDist = models.PerceptualLoss(model='net-lin', net='squeeze', use_gpu=1, gpu_ids=[0])
ImDist = lpips.LPIPS(net='squeeze', )
# model_vgg = models.PerceptualLoss(model='net-lin', net='vgg', use_gpu=1, gpu_ids=[0])
ImDist.cuda()
for param in ImDist.parameters():
    param.requires_grad_(False)
#%%
# BGAN = BigGAN.from_pretrained("biggan-deep-256")
# for param in BGAN.parameters():
#     param.requires_grad_(False)
# embed_mat = BGAN.embeddings.parameters().__next__().data
# BGAN.cuda()
from GAN_utils import BigGAN_wrapper, loadBigGAN
BGAN = loadBigGAN()
G = BigGAN_wrapper(BGAN)
#%%
def LExpMap(refvect, tangvect, ticks=11, lims=(-1,1)):
    refvect, tangvect = refvect.reshape(1, -1), tangvect.reshape(1, -1)
    steps = np.linspace(lims[0], lims[1], ticks)[:, np.newaxis]
    interp_vects = steps @ tangvect + refvect
    return interp_vects

def SExpMap(refvect, tangvect, ticks=11, lims=(-1,1)):
    refvect, tangvect = refvect.reshape(1, -1), tangvect.reshape(1, -1)
    steps = np.linspace(lims[0], lims[1], ticks)[:, np.newaxis] * np.pi / 2
    interp_vects = steps @ tangvect + refvect
    return interp_vects


#%% Test code for hessian eigendecomposition
#t0 = time()
#feat = torch.randn((1, 4096), dtype=torch.float32).requires_grad_(False).cuda()
#eigenvals, eigenvecs = compute_hessian_eigenthings(G, feat, ImDist,
#    num_eigenthings=300, mode="lanczos", use_gpu=True,)
#print(time() - t0,"\n")  # 81.02 s 
#%% Load the codes from the Backup folder 
import os
from  scipy.io import loadmat
import re
def load_codes_mat(backup_dir, threadnum=None, savefile=False):
    """ load all the code mat file in the experiment folder and summarize it into nparrays
    threadnum: can select one thread of the code if it's a parallel evolution. Usually, 0 or 1. 
        None for all threads. 
    """
    # make sure enough codes for requested size
    if "codes_all.npz" in os.listdir(backup_dir):
        # if the summary table exist, just read from it!
        with np.load(join(backup_dir, "codes_all.npz")) as data:
            codes_all = data["codes_all"]
            generations = data["generations"]
        return codes_all, generations
    if threadnum is None:
        codes_fns = sorted([fn for fn in os.listdir(backup_dir) if "_code.mat" in fn])
    else:
        codes_fns = sorted([fn for fn in os.listdir(backup_dir) if "thread%03d_code.mat"%(threadnum) in fn])
    codes_all = []
    img_ids = []
    for i, fn in enumerate(codes_fns[:]):
        matdata = loadmat(join(backup_dir, fn))
        codes_all.append(matdata["codes"])
        img_ids.extend(list(matdata["ids"]))

    codes_all = np.concatenate(tuple(codes_all), axis=0)
    img_ids = np.concatenate(tuple(img_ids), axis=0)
    img_ids = [img_ids[i][0] for i in range(len(img_ids))]
    generations = [int(re.findall("gen(\d+)", img_id)[0]) if 'gen' in img_id else -1 for img_id in img_ids]
    if savefile:
        np.savez(join(backup_dir, "codes_all.npz"), codes_all=codes_all, generations=generations)
    return codes_all, generations
#%% Use Rank weight like CMAES
def rankweight(popsize):
    weights_pad =np.zeros(popsize)
    mu = popsize/2
    weights = np.log(mu + 1 / 2) - (np.log(np.arange(1, 1 + np.floor(mu))))
    weights = weights / sum(weights)
    mu = int(mu)
    weights_pad[:mu] = weights
    return weights_pad
#%% Compute Image distance using the ImDist
def Hess_img_distmat(ImDist, img_all, nrow=11):
    """
    img_all: a torch 4D array of shape [N, 3, 256, 256] on cpu()
    nrow: specify how many images are generated in one axis.
        It will arrange the images in `img_all` in a matrix and calculate the distance to the center image in each row.
    """
    distmat = torch.zeros(img_all.shape[0]).view(-1,nrow)
    nTot = distmat.shape[0]
    for irow in range(nTot):
        rstr = irow * nrow
        rend = nrow + rstr
        rmid = (nrow-1)//2 + rstr
        with torch.no_grad():
            dists = ImDist(img_all[rstr:rend,:],img_all[rmid,:]).squeeze()
        distmat[irow, :]=dists.cpu()
    return distmat
#%% Utility functions to create movies from a row of interpolation.
def subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=11):
    targ_val_all = list(targ_val[::-1]) + [0] + list(targ_val)
    subsamp_ticks = []
    targ_ticks = []
    for i in range(len(targ_val_all)-1):
        seg = np.linspace(xticks_row[i], xticks_row[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
        tseg = np.linspace(targ_val_all[i], targ_val_all[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
        subsamp_ticks.extend(list(seg))
        targ_ticks.extend(list(tseg))
    subsamp_ticks.append(xticks_row[-1])
    targ_ticks.append(targ_val_all[-1])
    # codes_row = ref_vect + torch.tensor(subsamp_ticks).unsqueeze(1).float().cuda() @ tan_vec.cuda()
    codes_row = ref_vect.cpu().numpy() + np.array([subsamp_ticks]).T @ tan_vec.cpu().numpy()
    imgs = G.render(codes_row, B=8)
    return imgs, subsamp_ticks, targ_ticks, codes_row

def LERP_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=11):
    targ_val_all = list(targ_val[::-1]) + [0] + list(targ_val)
    subsamp_ticks = []
    targ_ticks = []
    seg = np.linspace(xticks_row[0], 0, int(abs((targ_val_all[0] - 0))/unit * density) , endpoint=False)
    tseg = np.linspace(targ_val_all[0], 0, int(abs((targ_val_all[0] - 0))/unit * density) , endpoint=False)
    subsamp_ticks.extend(list(seg))
    targ_ticks.extend(list(tseg))
    seg = np.linspace(0, xticks_row[-1], int(abs((targ_val_all[-1] - 0))/unit * density) , endpoint=False)
    tseg = np.linspace(0, targ_val_all[-1], int(abs((targ_val_all[-1] - 0))/unit * density) , endpoint=False)
    subsamp_ticks.extend(list(seg))
    targ_ticks.extend(list(tseg))
#    for i in range(len(targ_val_all)-1):
#        seg = np.linspace(xticks_row[i], xticks_row[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
#        tseg = np.linspace(targ_val_all[i], targ_val_all[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
#        subsamp_ticks.extend(list(seg))
#        targ_ticks.extend(list(tseg))
#    subsamp_ticks.append(xticks_row[-1])
#    targ_ticks.append(targ_val_all[-1])
    # codes_row = ref_vect + torch.tensor(subsamp_ticks).unsqueeze(1).float().cuda() @ tan_vec.cuda()
    codes_row = ref_vect.cpu().numpy() + np.array([subsamp_ticks]).T @ tan_vec.cpu().numpy()
    imgs = G.render(codes_row, B=8)
    return imgs, subsamp_ticks, targ_ticks, codes_row

def createSinuMovie(imgs, movdir="", savenm="eig", fps=20, writeframe=False):
    out = cv2.VideoWriter(join(movdir, "%s.avi"%savenm), cv2.VideoWriter_fourcc(*'XVID'), fps, imgs[0].shape[0:2])
    fN = len(imgs)
    centi = fN // 2
    for fi in [*range(centi, fN)] + [*range(fN - 1, -1, -1)] + [*range(0, centi+1)]:
        out.write((255.0*imgs[fi][:,:,::-1]).astype('uint8'))
    out.release()
    if writeframe:
        os.makedirs(join(movdir, 'frames'), exist_ok=True)
        for i, fi in enumerate([*range(centi, fN)] + [*range(fN - 1, -1, centi)]):
            imsave(join(movdir, 'frames', savenm+"_F%03d.jpg"%i), (255.0*imgs[fi]).astype('uint8'), )
#%% Load up the codes
from sklearn.decomposition import PCA 
import numpy as np
import matplotlib.pylab as plt
from imageio import imwrite

newimg_dir = join(backup_dir,"Hess_imgs")
summary_dir = join(backup_dir,"Hess_imgs","summary")
movie_dir = join(backup_dir,"Hess_imgs","Movie")
os.makedirs(newimg_dir,exist_ok=True)
os.makedirs(summary_dir,exist_ok=True)
os.makedirs(movie_dir,exist_ok=True)

print("Loading the codes from experiment folder %s", backup_dir)
evo_codes_all, generations = load_codes_mat(backup_dir, threadnum=threadid) 
generations = np.array(generations) 
print("Shape of codes", evo_codes_all.shape) 
# Use penultimate generation to generate the center
final_gen_codes = evo_codes_all[generations==max(generations)-1, :]
final_gen_norms = np.linalg.norm(final_gen_codes, axis=1)
final_gen_norm = final_gen_norms.mean()
print("Average norm of the last generation samples %.2f" % final_gen_norm)
#%% If there is score of images, load them up here. And compute the weighted average code
if score_rank_avg:
    try:
        scores_record = loadmat(join(backup_dir, "scores_record.mat"))["scores_record"]
        scores_thread = scores_record[:,threadid]
        assert len(scores_thread)==max(generations)-min(generations) # block of scores is the number of block of codes -1  
        final_gen_scores = scores_thread[-1].squeeze()
        assert len(final_gen_scores)==final_gen_codes.shape[0]
        print("Loading scores successful, use the Score Rank Weighted mean as center code.")
        sort_idx = np.argsort( - final_gen_scores)
        weights = rankweight(len(final_gen_scores))
        w_avg_code = weights[np.newaxis,:] @ final_gen_codes
    except Exception as e:
        score_rank_avg = False
        print(e)
        print("Loading scores not successful, use the unweighted mean instead.")

#%% PCA of the Existing code (adopted from Manifold experiment)
sphere_norm = final_gen_norm
print("Set sphere norm to the last generations norm!")
#% Do PCA and find the major trend of evolution
print("Computing PCs")
code_pca = PCA(n_components=50)
PC_Proj_codes = code_pca.fit_transform(evo_codes_all)
PC_vectors = code_pca.components_
if PC_Proj_codes[-1, 0] < 0:  # decide which is the positive direction for PC1
    inv_PC1 = True
    PC1_sign = -1
else:
    inv_PC1 = False
    PC1_sign = 1

PC1_vect = PC1_sign * PC_vectors[0,:] 
#%% Prepare the center vector to use in Hessian computation. 
#   Use this vector as the reference vector (center) in the Hessian computation
#   Before Aug. 6th, it's designed to use one code / image from the last generation and explore around it. 
#       this can be dangerous, sometimes one image will lose the major feature that we want in the population. (lose the 2 balls)
#   From Aug. 6th on, we decided to use the mean code from the last generation, which has a higher probability of 
#   From the setting of the bhv2 files find the fixed noise vectors
space_data = loadmat(join(backup_dir, "space_opts.mat"))["space_opts"]
evolspace = space_data[0,threadid]["name"][0]
print("Evolution happens in %s space, load the fixed code in `space_opts`" % evolspace)
if evolspace == "BigGAN_class":
    ref_noise_vec = space_data[0,threadid]['fix_noise_vec']
    #% Choose the mean final generation codes as our reference
    ref_class_vec = final_gen_codes.mean(axis=0, keepdims=True)  # final_gen_codes[0:1, :]
    if score_rank_avg:
        ref_class_vec = w_avg_code
if evolspace == "BigGAN_noise":
    ref_class_vec = space_data[0,threadid]['fix_class_vec']
    ref_noise_vec = final_gen_codes.mean(axis=0, keepdims=True)  # final_gen_codes[0:1, :]
    if score_rank_avg:
        ref_noise_vec = w_avg_code
elif evolspace == "BigGAN":
    ref_vec = final_gen_codes.mean(axis=0, keepdims=True)  # final_gen_codes[0:1, :]
    if score_rank_avg:
        ref_vec = w_avg_code
    ref_noise_vec = ref_vec[:, :128]
    ref_class_vec = ref_vec[:, 128:]
## View image correspond to the reference code
ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
refimg = G.visualize(ref_vect).cpu()
centimg = ToPILImage()(refimg[0,:,:,:])
centimg.show(title="Center Reference Image")
#%% Visualize the Final Generation  together  with the center reference image. 
VisFinalGen = False
if VisFinalGen:
    #% If you want to regenerate the images from last generation here.
    print("Review the last generation codes w.r.t. the center code for manifold.")
    imgs_final = G.visualize_batch_np(np.concatenate((ref_noise_vec.repeat(25,axis=0), final_gen_codes[:,:]), axis=1))
    ToPILImage()(make_grid(imgs_final,nrow=5)).show()
    #G.visualize(torch.from_numpy(np.concatenate((ref_noise_vec.repeat(5,axis=0), final_gen_codes[:5,:]), axis=1)).float().cuda()).cpu()
    #ToPILImage()(make_grid(imgs.cpu())).show()
#%% Compute Hessian decomposition and get the vectors
Hess_method = "BP"  # "BackwardIter" "ForwardIter"
Hess_all = False # Set to False to reducerrrrrr computation time. 
t0 = time()
if Hess_method == "BP":
    print("Computing Hessian Decomposition Through auto-grad and full eigen decomposition.")
    classvec = torch.from_numpy(ref_class_vec).float().cuda()  # embed_mat[:, class_id:class_id+1].cuda().T
    noisevec = torch.from_numpy(ref_noise_vec).float().cuda()
    ref_vect = torch.cat((noisevec, classvec, ), dim=1).detach().clone()
    mov_vect = ref_vect.detach().clone().requires_grad_(True)
    #%
    imgs1 = G.visualize(ref_vect)
    if Hess_all:
        imgs2 = G.visualize(mov_vect)
        dsim = ImDist(imgs1, imgs2)
        H = get_full_hessian(dsim, mov_vect)  # 77sec to compute a Hessian. # 114sec on ML2a
        # ToPILImage()(imgs[0,:,:,:].cpu()).show()
        eigvals, eigvects = np.linalg.eigh(H)  # 75 ms
    #%
    noisevec.requires_grad_(True)
    classvec.requires_grad_(False)
    mov_vect = torch.cat((noisevec, classvec, ), dim=1)
    imgs2 = G.visualize(mov_vect)
    dsim = ImDist(imgs1, imgs2)
    H_nois = get_full_hessian(dsim, noisevec)  # 39.3 sec to compute a Hessian.# 59 sec on ML2a
    eigvals_nois, eigvects_nois = np.linalg.eigh(H_nois)  # 75 ms
    #%
    noisevec.requires_grad_(False)
    classvec.requires_grad_(True)
    mov_vect = torch.cat((noisevec, classvec, ), dim=1)
    imgs2 = G.visualize(mov_vect)
    dsim = ImDist(imgs1, imgs2)
    H_clas = get_full_hessian(dsim, classvec)  # 39.3 sec to compute a Hessian.
    eigvals_clas, eigvects_clas = np.linalg.eigh(H_clas)  # 75 ms
    classvec.requires_grad_(False)
    del dsim, imgs2, imgs1
    torch.cuda.empty_cache()
    if Hess_all:
        np.savez(join(summary_dir, "Hess_mat.npz"), H=H, eigvals=eigvals, eigvects=eigvects, 
             H_clas=H_clas, eigvals_clas=eigvals_clas, eigvects_clas=eigvects_clas, 
             H_nois=H_nois, eigvals_nois=eigvals_nois, eigvects_nois=eigvects_nois, 
             vect=ref_vect.cpu().numpy(), noisevec=noisevec.cpu().numpy(), classvec=classvec.cpu().numpy())
    else:
        np.savez(join(summary_dir, "Hess_mat.npz"), #H=H, eigvals=eigvals, eigvects=eigvects, 
             H_clas=H_clas, eigvals_clas=eigvals_clas, eigvects_clas=eigvects_clas, 
             H_nois=H_nois, eigvals_nois=eigvals_nois, eigvects_nois=eigvects_nois, 
             vect=ref_vect.cpu().numpy(), noisevec=noisevec.cpu().numpy(), classvec=classvec.cpu().numpy())
elif Hess_method == "BackwardIter":
    print("Computing Hessian Decomposition Through Lanczos decomposition on Backward HVP operator.")
    feat = torch.from_numpy(sphere_norm * PC1_vect).float().requires_grad_(False).cuda()
    eigenvals, eigenvecs = compute_hessian_eigenthings(G, feat, ImDist,
        num_eigenthings=128, mode="lanczos", use_gpu=True)
    eigenvals = eigenvals[::-1]
    eigenvecs = eigenvecs[::-1, :]

elif Hess_method == "ForwardIter": 
    print("Computing Hessian Decomposition Through Lanczos decomposition on Forward HVP operator.")
    pass

print("%.2f sec"% (time() - t0))  # 31.75 secs for 300 eig, 87.52 secs for 800 eigs. 
#%% Visualize spectrum
plt.figure(figsize=[8,5])
plt.subplot(1,2,1)
if Hess_all: plt.plot(eigvals[::-1], label="all")
plt.plot(eigvals_clas[::-1], label="class")
plt.plot(eigvals_nois[::-1], label="noise")
plt.ylabel("eigval");plt.legend()
plt.subplot(1,2,2)
if Hess_all: plt.plot(np.log10(eigvals[::-1]), label="all")
plt.plot(np.log10(eigvals_clas[::-1]), label="class")
plt.plot(np.log10(eigvals_nois[::-1]), label="noise")
plt.ylabel("log(eigval)");plt.legend()
plt.savefig(join(summary_dir, "spectrum.jpg"))
#%% Optional: Angle with PC1 vector
if evolspace == "BigGAN_class":
    innerprod2PC1 = PC1_vect @ eigvects_clas.T
elif evolspace == "BigGAN_noise":
    innerprod2PC1 = PC1_vect @ eigvects_nois.T
elif evolspace == "BigGAN":
    innerprod2PC1 = PC1_vect @ eigvects.T
print("Eigen vector: Innerproduct max %.3E min %.3E std %.3E"% (innerprod2PC1.max(), innerprod2PC1.min(), innerprod2PC1.std()))
print("EigenDecomposition of Hessian of Image Similarity Metric\nEigen value: Class space max %.3E min %.3E std %.3E; Noise space max %.3E min %.3E std %.3E"%
      (eigvals_clas.max(), eigvals_clas.min(), eigvals_clas.std(), eigvals_nois.max(), eigvals_nois.min(), eigvals_nois.std(), )) 
if Hess_all: 
    print("EigenDecomposition of Hessian of Image Similarity Metric\nEigen value: All: max %.3E min %.3E std %.3E"%
      (eigvals.max(), eigvals.min(), eigvals.std(),))

#%% Do interpolation along each axes
#%%
#classvec = torch.from_numpy(ref_class_vec).float().cuda()  # embed_mat[:, class_id:class_id+1].cuda().T
#noisevec = torch.from_numpy(ref_noise_vec).float().cuda()
if not exact_distance:
    #% Interpolation in the class space, but inversely scale the step size w.r.t. eigenvalue
    codes_all = []
    img_names = []
    scale = 5
    expon = 2.5
    eiglist_class = [0, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40,]
    for eigi in eiglist_class:#range(20):  # eigvects.shape[1] # 60, 80
        interp_class = LExpMap(classvec.cpu().numpy(), eigvects_clas[:, -eigi-1], 11, (-scale * eigvals_clas[-eigi-1] ** (-1/expon), scale * eigvals_clas[-eigi-1] ** (-1/expon)))
        interp_codes = np.hstack((noisevec.cpu().numpy().repeat(11, axis=0), interp_class, ))
        codes_all.append(interp_codes.copy())
        img_names.extend("class_eig%d_exp%.1f_lin%.1f.jpg"%(eigi, expon, dist) for dist in np.linspace(-scale, scale, 11))

    codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
    img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
    imggrid = make_grid(img_all, nrow=11)
    PILimg2 = ToPILImage()(imggrid)#.show()
    PILimg2.save(join(summary_dir, "eigvect_clas_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    npimgs = img_all.permute([2,3,1,0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
#%
    nrow = 11
    distmat = Hess_img_distmat(ImDist, img_all, nrow=11)
    plt.figure(figsize=[6, 6])
    plt.matshow(distmat, fignum=0)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nclass space exponent %.1f Scale%d "%(expon, scale, ))
    plt.savefig(join(summary_dir, "distmat_eigvect_clas_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    plt.show()
    #% Interpolation in the noise space
    codes_all = []
    img_names = []
    scale = 6
    expon = 3
    eiglist_noise = [0, 1, 2, 3, 4, 6, 10, 15, 20, 40]
    for eigi in eiglist_noise:#range(20):#eigvects_nois.shape[1]
    #    interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-4.5, 4.5))
        interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-scale * eigvals_nois[-eigi-1] ** (-1/expon), scale * eigvals_nois[-eigi-1] ** (-1/expon)))
        interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), ))
        codes_all.append(interp_codes.copy())
        img_names.extend("noise_eig%d_exp%.1f_lin%.1f.jpg"%(eigi, expon, dist) for dist in np.linspace(-scale, scale, 11))

    codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
    img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
    imggrid = make_grid(img_all, nrow=11)
    PILimg3 = ToPILImage()(imggrid)#.show()
    PILimg3.save(join(summary_dir, "eigvect_nois_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    npimgs = img_all.permute([2,3,1,0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
    #
    nrow = 11
    distmat = Hess_img_distmat(ImDist, img_all, nrow=11)
    # Show the image distance from center reference to each image around it.
    plt.figure(figsize=[6,6])
    plt.matshow(distmat, fignum=0)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nnoise space  Scale%d exponent %.1f"%(scale, expon))
    plt.savefig(join(summary_dir, "distmat_eigvect_nois_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    plt.show()

else:  # exact_distance by line search
    from ImDist_Line_Searcher import find_level_step
    targ_val = np.array(target_distance)
    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
    refimg = G.visualize(ref_vect)
    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda()
    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda()

    space = "noise"
    imgall = None
    xtick_col = []
    dsim_col = []
    vecs_col = []
    img_names = []
    tick_labels = list(-targ_val[::-1]) + [0] + list(targ_val)  # -0.5, -0.4 ...  0.4, 0.5
    t0 = time()
    # eiglist_noise = [0, 1, 2, 3, 4, 5, 6, 8, 10, 20, 30, 40, ]#[0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 20, 30, 40, 50, 60, 70, 80]
    eiglist_noise = [0, 1, 4, 8, 10, 20, ]# [0, 1, 2, 4, 6, 8, 10, 20, ]
    for eigid in tqdm(eiglist_noise):  # range(128):  # #
        if space == "class":
            tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
        elif space == "noise":
            tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
        xtar_pos, ytar_pos, stepimgs_pos = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=True, maxdist=30)
        xtar_neg, ytar_neg, stepimgs_neg = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=False, maxdist=30)
        imgrow = torch.cat((torch.flip(stepimgs_neg, (0,)), refimg, stepimgs_pos)).cpu()
        xticks_row = xtar_neg[::-1] + [0.0] + xtar_pos
        dsim_row = list(ytar_neg[::-1]) + [0.0] + list(ytar_pos)
        vecs_row = torch.tensor(xticks_row).float().cuda().view(-1, 1) @ tan_vec + ref_vect

        xtick_col.append(xticks_row)
        dsim_col.append(dsim_row)
        vecs_col.append(vecs_row.cpu().numpy())
        img_names.extend("noise_eig%d_lin%.2f.jpg" % (eigid, dist) for dist in tick_labels)  # dsim_row)
        imgall = imgrow if imgall is None else torch.cat((imgall, imgrow))
        # Subsample the exploration and visualize the images on the line. ANd then render it into a movie. 
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)

    mtg1 = ToPILImage()(make_grid(imgall, nrow=2*len(target_distance)+1).cpu())  # 20sec for 13 rows not bad
    mtg1.show()
    mtg1.save(join(summary_dir, "noise_space_all_var.jpg"))
    npimgs = imgall.permute([2, 3, 1, 0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]),
                                                  np.uint8(npimgs[:, :, :, imgi] * 255))
    print(time() - t0)
    # %
    xtick_arr = np.array(xtick_col)
    dsim_arr = np.array(dsim_col)
    vecs_arr = np.array(vecs_col)
    np.savez(join(summary_dir, "noise_ImDist_root_data.npz"), xtick_arr=xtick_arr, dsim_arr=dsim_arr, vecs_arr=vecs_arr,
             targ_val=targ_val, eiglist=eiglist_noise)
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(xtick_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("L2 deviation from center")
    plt.xticks(range(len(eiglist_noise)), eiglist_noise)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Distance Travel Along Given Eigen vector to achieve certain Image Distance")
    plt.savefig(join(summary_dir, "noise_code_deviation.jpg"))
    plt.show()
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(dsim_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("Image Distance")
    plt.xticks(range(len(eiglist_noise)), eiglist_noise)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Achieved Image Distance Along Each Axis")
    plt.savefig(join(summary_dir, "noise_space_dist_curv.jpg"))
    plt.show()
    # %
    plt.figure()
    plt.matshow(dsim_arr, fignum=0)
    plt.xticks(range(len(tick_labels)), tick_labels)
    plt.yticks(range(len(eiglist_noise)), eiglist_noise)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nnoise space")
    plt.savefig(join(summary_dir, "noise_space_distmat.jpg"))
    plt.show()
    #%
    space = "class"
    imgall = None
    xtick_col = []
    dsim_col = []
    vecs_col = []
    img_names = []
    tick_labels = list(-targ_val[::-1]) + [0] + list(targ_val)
    t0 = time()
    #eiglist_class = [0, 1, 2, 3, 6, 9, 13, 17, 21, 25, 30, 40, 60, ] #[0, 1, 2, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40, 50, 60, 70, 80]
    eiglist_class = [0, 1, 2, 6, 17, 30, ]#[0, 1, 2, 3, 6, 9, 17, 21, 30, 60, ]
    for eigid in tqdm(eiglist_class):  # [0,1,2,3,4,5,6,7,8,10,20,30,
        # 40]:#
        if space == "class":
            tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
        elif space == "noise":
            tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
        xtar_pos, ytar_pos, stepimgs_pos = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=True, maxdist=30)
        xtar_neg, ytar_neg, stepimgs_neg = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=False, maxdist=30)
        imgrow = torch.cat((torch.flip(stepimgs_neg, (0,)), refimg, stepimgs_pos)).cpu()
        xticks_row = xtar_neg[::-1] + [0.0] + xtar_pos
        dsim_row = list(ytar_neg[::-1]) + [0.0] + list(ytar_pos)
        vecs_row = torch.tensor(xticks_row).float().cuda().view(-1, 1) @ tan_vec + ref_vect
        xtick_col.append(xticks_row)
        dsim_col.append(dsim_row)
        vecs_col.append(vecs_row.cpu().numpy())
        img_names.extend(
            "class_eig%d_lin%.2f.jpg" % (eigid, dist) for dist in tick_labels)  # np.linspace(-0.4, 0.4,11))
        imgall = imgrow if imgall is None else torch.cat((imgall, imgrow))
        # Subsample the exploration and visualize the images on the line. ANd then render it into a movie. 
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)

    mtg2 = ToPILImage()(make_grid(imgall, nrow=2*len(target_distance)+1).cpu())  # 20sec for 13 rows not bad
    mtg2.show()
    mtg2.save(join(summary_dir, "class_space_all_var.jpg"))
    npimgs = imgall.permute([2, 3, 1, 0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]),
                                                  np.uint8(npimgs[:, :, :, imgi] * 255))
    print(time() - t0)
    # %
    xtick_arr = np.array(xtick_col)
    dsim_arr = np.array(dsim_col)
    vecs_arr = np.array(vecs_col)
    np.savez(join(summary_dir, "class_ImDist_root_data.npz"), xtick_arr=xtick_arr, dsim_arr=dsim_arr, vecs_arr=vecs_arr,
             targ_val=targ_val, eiglist=eiglist_class)
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(xtick_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("L2 deviation from center")
    plt.xticks(range(len(eiglist_class)), eiglist_class)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Distance Travel Along Given Eigen vector to achieve certain Image Distance")
    plt.savefig(join(summary_dir, "class_code_deviation.jpg"))
    plt.show()
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(dsim_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("Image Distance")
    plt.xticks(range(len(eiglist_class)), eiglist_class)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Achieved Image Distance Along Each Axis")
    plt.savefig(join(summary_dir, "class_space_dist_curv.jpg"))
    plt.show()
    # %
    plt.figure()
    plt.matshow(dsim_arr, fignum=0)
    plt.xticks(range(len(tick_labels)), tick_labels)
    plt.yticks(range(len(eiglist_class)), eiglist_class)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nnoise space")
    plt.savefig(join(summary_dir, "class_space_distmat.jpg"))
    plt.show()
#%%

#movie_dir = "."
if exact_distance:
    space = "noise"
#    targ_val = np.array(target_distance)
    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda()
    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda()
    
    data = np.load(join(summary_dir, "noise_ImDist_root_data.npz"))
    xtick_arr = data["xtick_arr"]
    eiglist_noise = data["eiglist"]
    for i, eigid in enumerate(eiglist_noise):
        tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
        xticks_row = xtick_arr[i, :]
        imgs, _, _, _ = LERP_img_row(ref_vect, tan_vec, target_distance, xticks_row, unit=0.08, density=5)
        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort_LERP" % (space, eigid), fps=30, writeframe=True)
    
    space = "class"
    data = np.load(join(summary_dir, "class_ImDist_root_data.npz"))
    xtick_arr = data["xtick_arr"]
    eiglist_class = data["eiglist"]
    for i, eigid in enumerate(eiglist_class):
        tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
        xticks_row = xtick_arr[i, :]
        imgs, _, _, _ = LERP_img_row(ref_vect, tan_vec, target_distance, xticks_row, unit=0.08, density=5)
        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort_LERP" % (space, eigid), fps=30, writeframe=True)
    

#%% Regenerate the movies with different fps
#summary_dir = os.path.join(backup_dir, "Hess_imgs", "summary")
#movie_dir = "."
#data = np.load(join(summary_dir, "Hess_mat.npz")) #H=H, eigvals=eigvals, eigvects=eigvects, 
#eigvects_clas = data["eigvects_clas"]
#eigvects_nois = data["eigvects_nois"]
#ref_noise_vec = data["noisevec"]
#ref_class_vec = data["classvec"]

#if exact_distance:
#    space = "noise"
##    targ_val = np.array(target_distance)
#    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
#    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda()
#    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda()
#    
#    data = np.load(join(summary_dir, "noise_ImDist_root_data.npz"))
#    xtick_arr = data["xtick_arr"]
#    eiglist_noise = data["eiglist"]
#    for i, eigid in enumerate(eiglist_noise):
#        tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
#        xticks_row = xtick_arr[i, :]
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)
#    
#    space = "class"
#    data = np.load(join(summary_dir, "class_ImDist_root_data.npz"))
#    xtick_arr = data["xtick_arr"]
#    eiglist_class = data["eiglist"]
#    for i, eigid in enumerate(eiglist_class):
#        tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
#        xticks_row = xtick_arr[i, :]
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)
#%%
##% Interpolation in the class space
#codes_all = []
#img_names = []
#for eigi in [0, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40,]:#range(20):  # eigvects.shape[1] # 60, 80
#    interp_class = LExpMap(classvec.cpu().numpy(), eigvects_clas[:, -eigi-1], 11, (-2.5, 2.5))
#    interp_codes = np.hstack((noisevec.cpu().numpy().repeat(11, axis=0), interp_class, ))
#    codes_all.append(interp_codes.copy())
#    img_names.extend("class_eig%d_lin%.1f.jpg"%(eigi, dist) for dist in np.linspace(-2.5, 2.5, 11))
#codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
#img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
#imggrid = make_grid(img_all, nrow=11)
#PILimg2 = ToPILImage()(imggrid)#.show()
#PILimg2.save(join(summary_dir, "eigvect_clas_interp.jpg"))
#npimgs = img_all.permute([2,3,1,0]).numpy()
#for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
##%%
##% Interpolation in the noise space
#codes_all = []
#img_names = []
#for eigi in [0, 1, 2, 3, 4, 6, 10, 15, 20, 40]:#range(20):#eigvects_nois.shape[1]
#    interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-4.5, 4.5))
#    interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), ))
#    codes_all.append(interp_codes.copy())
#    img_names.extend("noise_eig%d_lin%.1f.jpg"%(eigi, dist) for dist in np.linspace(-4.5, 4.5, 11))
#
#codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
#img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
#imggrid = make_grid(img_all, nrow=11)
#PILimg3 = ToPILImage()(imggrid)#.show()
#PILimg3.save(join(summary_dir, "eigvect_nois_interp.jpg"))
#npimgs = img_all.permute([2,3,1,0]).numpy()
#for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
#print("Spent %.1f sec from start" % (time() - T00))
#%% Interpolation in the noise space
#codes_all = []
#img_names = []
#for eigi in range(20):#eigvects_nois.shape[1]
#    interp_noise = SExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-1, 1))
#    interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), ))
#    codes_all.append(interp_codes.copy())
#    img_names.extend("noise_eig%d_sphang%d.jpg"%(eigi, angle) for angle in np.linspace(-90,90,11))
#
#codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
#img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
#imggrid = make_grid(img_all, nrow=11)
#PILimg4 = ToPILImage()(imggrid)#.show()
#PILimg4.save(join(summary_dir, "eigvect_nois_sph_interp.jpg"))
#npimgs = img_all.permute([2,3,1,0]).numpy()
#for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))

#%% 
#vec_norm = 220# sphere_norm
#eig_id = 0
#perturb_vect = eigenvecs[eig_id,:] # PC_vectors[1,:]    
#ang_step = 180 / 10
#theta_arr_deg = ang_step * np.arange(-5, 6)
#theta_arr = ang_step * np.arange(-5, 6) / 180 * np.pi
#codes_arc = np.array([np.cos(theta_arr), 
#                      np.sin(theta_arr) ]).T @ np.array([PC1_vect, perturb_vect])
#norms = np.linalg.norm(codes_arc, axis=1)
#codes_arc = codes_arc / norms[:,np.newaxis] * vec_norm
#imgs = G.visualize(torch.from_numpy(codes_arc).float().cuda())
#
#npimgs = imgs.detach().cpu().permute([2, 3, 1, 0]).numpy()
#for i in range(npimgs.shape[3]):
#    angle = theta_arr_deg[i]
#    imwrite(join(newimg_dir, "norm%d_eig%d_ang%d.jpg"%(vec_norm, eig_id, angle)), npimgs[:,:,:,i])
#
#img_list = [npimgs[:,:,:,i] for i in range(npimgs.shape[3])]
#mtg1 = build_montages(img_list, [256, 256], [11, 1])[0]
##imwrite(join(backup_dir, "norm%d_eig%d.jpg"%(vec_norm, eig_id)),mtg1)
#imwrite(join(newimg_dir, "norm%d_eig%d.jpg"%(vec_norm, eig_id)),mtg1)
#%%

# -*- coding: utf-8 -*-
"""
Created on Wed Jul 22 19:08:54 2020
Modified Oct. 14th

@author: <NAME>

Find important Nuisanced + Class transformations in Noise + Class space for a BigGAN evolved image. 
"""

# backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_integrated\2020-06-01-09-46-37"
# Put the backup folder and the thread to analyze here 
#backup_dir = r"C:\Users\Poncelab-ML2a\Documents\monkeylogic2\generate_BigGAN\2020-07-22-10-14-22"
# backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_BigGAN\2020-08-06-10-18-55"#2020-08-04-09-54-25"#
backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_BigGAN\2020-11-06-11-53-09"
threadid = 1

score_rank_avg = False  # If True, it will try to read "scores_record.mat", from the backup folder and read "scores_record"
                        # Else, it will use the unweighted mean code of the last generation as the center vector. 
                        # Need to run the BigGAN postHoc Analysis to save the `scores_record` mat and use this flag

exact_distance = True   # Control if exact distance search is used or approximate heuristic rule is used.
target_distance = [0.16, 0.24, 0.32, 0.40] # 0.08, 
#target_distance = [0.09, 0.18, 0.27, 0.36, 0.45]  # if exact_distance is True it will search for images with these
                                                  # distance to reference image along each eigenvector.
                                                  
#%% Prepare the generator model and perceptual loss networks
from time import time
import os
from os.path import join
import sys
import cv2 # for video writing
if os.environ['COMPUTERNAME'] == 'PONCELAB-ML2B':
    Python_dir = r"C:\Users\<NAME>\Documents\Python"
elif os.environ['COMPUTERNAME'] == 'PONCELAB-ML2A':
    Python_dir = r"C:\Users\Poncelab-ML2a\Documents\Python"
elif os.environ['COMPUTERNAME'] == 'DESKTOP-MENSD6S':
    Python_dir = r"E:\Github_Projects"
elif os.environ['COMPUTERNAME'] == 'DESKTOP-9DDE2RH':
    Python_dir = r"D:\Github"

sys.path.append(join(Python_dir, "Visual_Neuro_InSilico_Exp"))
sys.path.append(join(Python_dir, "PerceptualSimilarity"))
import torch
import numpy as np
from pytorch_pretrained_biggan import BigGAN, truncated_noise_sample
from GAN_utils import upconvGAN
from hessian_eigenthings.lanczos import lanczos
from GAN_hvp_operator import GANHVPOperator, GANForwardHVPOperator, GANForwardMetricHVPOperator, \
    compute_hessian_eigenthings, get_full_hessian
from skimage.io import imsave, imsave
from torchvision.utils import make_grid
from build_montages import build_montages
from torchvision.transforms import ToPILImage, ToTensor
from IPython.display import clear_output
from hessian_eigenthings.utils import progress_bar
from tqdm import tqdm
T00 = time()
import lpips
# import models # from PerceptualSimilarity folder
# ImDist = models.PerceptualLoss(model='net-lin', net='squeeze', use_gpu=1, gpu_ids=[0])
ImDist = lpips.LPIPS(net='squeeze', )
# model_vgg = models.PerceptualLoss(model='net-lin', net='vgg', use_gpu=1, gpu_ids=[0])
ImDist.cuda()
for param in ImDist.parameters():
    param.requires_grad_(False)
#%%
# BGAN = BigGAN.from_pretrained("biggan-deep-256")
# for param in BGAN.parameters():
#     param.requires_grad_(False)
# embed_mat = BGAN.embeddings.parameters().__next__().data
# BGAN.cuda()
from GAN_utils import BigGAN_wrapper, loadBigGAN
BGAN = loadBigGAN()
G = BigGAN_wrapper(BGAN)
#%%
def LExpMap(refvect, tangvect, ticks=11, lims=(-1,1)):
    refvect, tangvect = refvect.reshape(1, -1), tangvect.reshape(1, -1)
    steps = np.linspace(lims[0], lims[1], ticks)[:, np.newaxis]
    interp_vects = steps @ tangvect + refvect
    return interp_vects

def SExpMap(refvect, tangvect, ticks=11, lims=(-1,1)):
    refvect, tangvect = refvect.reshape(1, -1), tangvect.reshape(1, -1)
    steps = np.linspace(lims[0], lims[1], ticks)[:, np.newaxis] * np.pi / 2
    interp_vects = steps @ tangvect + refvect
    return interp_vects


#%% Test code for hessian eigendecomposition
#t0 = time()
#feat = torch.randn((1, 4096), dtype=torch.float32).requires_grad_(False).cuda()
#eigenvals, eigenvecs = compute_hessian_eigenthings(G, feat, ImDist,
#    num_eigenthings=300, mode="lanczos", use_gpu=True,)
#print(time() - t0,"\n")  # 81.02 s 
#%% Load the codes from the Backup folder 
import os
from  scipy.io import loadmat
import re
def load_codes_mat(backup_dir, threadnum=None, savefile=False):
    """ load all the code mat file in the experiment folder and summarize it into nparrays
    threadnum: can select one thread of the code if it's a parallel evolution. Usually, 0 or 1. 
        None for all threads. 
    """
    # make sure enough codes for requested size
    if "codes_all.npz" in os.listdir(backup_dir):
        # if the summary table exist, just read from it!
        with np.load(join(backup_dir, "codes_all.npz")) as data:
            codes_all = data["codes_all"]
            generations = data["generations"]
        return codes_all, generations
    if threadnum is None:
        codes_fns = sorted([fn for fn in os.listdir(backup_dir) if "_code.mat" in fn])
    else:
        codes_fns = sorted([fn for fn in os.listdir(backup_dir) if "thread%03d_code.mat"%(threadnum) in fn])
    codes_all = []
    img_ids = []
    for i, fn in enumerate(codes_fns[:]):
        matdata = loadmat(join(backup_dir, fn))
        codes_all.append(matdata["codes"])
        img_ids.extend(list(matdata["ids"]))

    codes_all = np.concatenate(tuple(codes_all), axis=0)
    img_ids = np.concatenate(tuple(img_ids), axis=0)
    img_ids = [img_ids[i][0] for i in range(len(img_ids))]
    generations = [int(re.findall("gen(\d+)", img_id)[0]) if 'gen' in img_id else -1 for img_id in img_ids]
    if savefile:
        np.savez(join(backup_dir, "codes_all.npz"), codes_all=codes_all, generations=generations)
    return codes_all, generations
#%% Use Rank weight like CMAES
def rankweight(popsize):
    weights_pad =np.zeros(popsize)
    mu = popsize/2
    weights = np.log(mu + 1 / 2) - (np.log(np.arange(1, 1 + np.floor(mu))))
    weights = weights / sum(weights)
    mu = int(mu)
    weights_pad[:mu] = weights
    return weights_pad
#%% Compute Image distance using the ImDist
def Hess_img_distmat(ImDist, img_all, nrow=11):
    """
    img_all: a torch 4D array of shape [N, 3, 256, 256] on cpu()
    nrow: specify how many images are generated in one axis.
        It will arrange the images in `img_all` in a matrix and calculate the distance to the center image in each row.
    """
    distmat = torch.zeros(img_all.shape[0]).view(-1,nrow)
    nTot = distmat.shape[0]
    for irow in range(nTot):
        rstr = irow * nrow
        rend = nrow + rstr
        rmid = (nrow-1)//2 + rstr
        with torch.no_grad():
            dists = ImDist(img_all[rstr:rend,:],img_all[rmid,:]).squeeze()
        distmat[irow, :]=dists.cpu()
    return distmat
#%% Utility functions to create movies from a row of interpolation.
def subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=11):
    targ_val_all = list(targ_val[::-1]) + [0] + list(targ_val)
    subsamp_ticks = []
    targ_ticks = []
    for i in range(len(targ_val_all)-1):
        seg = np.linspace(xticks_row[i], xticks_row[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
        tseg = np.linspace(targ_val_all[i], targ_val_all[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
        subsamp_ticks.extend(list(seg))
        targ_ticks.extend(list(tseg))
    subsamp_ticks.append(xticks_row[-1])
    targ_ticks.append(targ_val_all[-1])
    # codes_row = ref_vect + torch.tensor(subsamp_ticks).unsqueeze(1).float().cuda() @ tan_vec.cuda()
    codes_row = ref_vect.cpu().numpy() + np.array([subsamp_ticks]).T @ tan_vec.cpu().numpy()
    imgs = G.render(codes_row, B=8)
    return imgs, subsamp_ticks, targ_ticks, codes_row

def LERP_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=11):
    targ_val_all = list(targ_val[::-1]) + [0] + list(targ_val)
    subsamp_ticks = []
    targ_ticks = []
    seg = np.linspace(xticks_row[0], 0, int(abs((targ_val_all[0] - 0))/unit * density) , endpoint=False)
    tseg = np.linspace(targ_val_all[0], 0, int(abs((targ_val_all[0] - 0))/unit * density) , endpoint=False)
    subsamp_ticks.extend(list(seg))
    targ_ticks.extend(list(tseg))
    seg = np.linspace(0, xticks_row[-1], int(abs((targ_val_all[-1] - 0))/unit * density) , endpoint=False)
    tseg = np.linspace(0, targ_val_all[-1], int(abs((targ_val_all[-1] - 0))/unit * density) , endpoint=False)
    subsamp_ticks.extend(list(seg))
    targ_ticks.extend(list(tseg))
#    for i in range(len(targ_val_all)-1):
#        seg = np.linspace(xticks_row[i], xticks_row[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
#        tseg = np.linspace(targ_val_all[i], targ_val_all[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False)
#        subsamp_ticks.extend(list(seg))
#        targ_ticks.extend(list(tseg))
#    subsamp_ticks.append(xticks_row[-1])
#    targ_ticks.append(targ_val_all[-1])
    # codes_row = ref_vect + torch.tensor(subsamp_ticks).unsqueeze(1).float().cuda() @ tan_vec.cuda()
    codes_row = ref_vect.cpu().numpy() + np.array([subsamp_ticks]).T @ tan_vec.cpu().numpy()
    imgs = G.render(codes_row, B=8)
    return imgs, subsamp_ticks, targ_ticks, codes_row

def createSinuMovie(imgs, movdir="", savenm="eig", fps=20, writeframe=False):
    out = cv2.VideoWriter(join(movdir, "%s.avi"%savenm), cv2.VideoWriter_fourcc(*'XVID'), fps, imgs[0].shape[0:2])
    fN = len(imgs)
    centi = fN // 2
    for fi in [*range(centi, fN)] + [*range(fN - 1, -1, -1)] + [*range(0, centi+1)]:
        out.write((255.0*imgs[fi][:,:,::-1]).astype('uint8'))
    out.release()
    if writeframe:
        os.makedirs(join(movdir, 'frames'), exist_ok=True)
        for i, fi in enumerate([*range(centi, fN)] + [*range(fN - 1, -1, centi)]):
            imsave(join(movdir, 'frames', savenm+"_F%03d.jpg"%i), (255.0*imgs[fi]).astype('uint8'), )
#%% Load up the codes
from sklearn.decomposition import PCA 
import numpy as np
import matplotlib.pylab as plt
from imageio import imwrite

newimg_dir = join(backup_dir,"Hess_imgs")
summary_dir = join(backup_dir,"Hess_imgs","summary")
movie_dir = join(backup_dir,"Hess_imgs","Movie")
os.makedirs(newimg_dir,exist_ok=True)
os.makedirs(summary_dir,exist_ok=True)
os.makedirs(movie_dir,exist_ok=True)

print("Loading the codes from experiment folder %s", backup_dir)
evo_codes_all, generations = load_codes_mat(backup_dir, threadnum=threadid) 
generations = np.array(generations) 
print("Shape of codes", evo_codes_all.shape) 
# Use penultimate generation to generate the center
final_gen_codes = evo_codes_all[generations==max(generations)-1, :]
final_gen_norms = np.linalg.norm(final_gen_codes, axis=1)
final_gen_norm = final_gen_norms.mean()
print("Average norm of the last generation samples %.2f" % final_gen_norm)
#%% If there is score of images, load them up here. And compute the weighted average code
if score_rank_avg:
    try:
        scores_record = loadmat(join(backup_dir, "scores_record.mat"))["scores_record"]
        scores_thread = scores_record[:,threadid]
        assert len(scores_thread)==max(generations)-min(generations) # block of scores is the number of block of codes -1  
        final_gen_scores = scores_thread[-1].squeeze()
        assert len(final_gen_scores)==final_gen_codes.shape[0]
        print("Loading scores successful, use the Score Rank Weighted mean as center code.")
        sort_idx = np.argsort( - final_gen_scores)
        weights = rankweight(len(final_gen_scores))
        w_avg_code = weights[np.newaxis,:] @ final_gen_codes
    except Exception as e:
        score_rank_avg = False
        print(e)
        print("Loading scores not successful, use the unweighted mean instead.")

#%% PCA of the Existing code (adopted from Manifold experiment)
sphere_norm = final_gen_norm
print("Set sphere norm to the last generations norm!")
#% Do PCA and find the major trend of evolution
print("Computing PCs")
code_pca = PCA(n_components=50)
PC_Proj_codes = code_pca.fit_transform(evo_codes_all)
PC_vectors = code_pca.components_
if PC_Proj_codes[-1, 0] < 0:  # decide which is the positive direction for PC1
    inv_PC1 = True
    PC1_sign = -1
else:
    inv_PC1 = False
    PC1_sign = 1

PC1_vect = PC1_sign * PC_vectors[0,:] 
#%% Prepare the center vector to use in Hessian computation. 
#   Use this vector as the reference vector (center) in the Hessian computation
#   Before Aug. 6th, it's designed to use one code / image from the last generation and explore around it. 
#       this can be dangerous, sometimes one image will lose the major feature that we want in the population. (lose the 2 balls)
#   From Aug. 6th on, we decided to use the mean code from the last generation, which has a higher probability of 
#   From the setting of the bhv2 files find the fixed noise vectors
space_data = loadmat(join(backup_dir, "space_opts.mat"))["space_opts"]
evolspace = space_data[0,threadid]["name"][0]
print("Evolution happens in %s space, load the fixed code in `space_opts`" % evolspace)
if evolspace == "BigGAN_class":
    ref_noise_vec = space_data[0,threadid]['fix_noise_vec']
    #% Choose the mean final generation codes as our reference
    ref_class_vec = final_gen_codes.mean(axis=0, keepdims=True)  # final_gen_codes[0:1, :]
    if score_rank_avg:
        ref_class_vec = w_avg_code
if evolspace == "BigGAN_noise":
    ref_class_vec = space_data[0,threadid]['fix_class_vec']
    ref_noise_vec = final_gen_codes.mean(axis=0, keepdims=True)  # final_gen_codes[0:1, :]
    if score_rank_avg:
        ref_noise_vec = w_avg_code
elif evolspace == "BigGAN":
    ref_vec = final_gen_codes.mean(axis=0, keepdims=True)  # final_gen_codes[0:1, :]
    if score_rank_avg:
        ref_vec = w_avg_code
    ref_noise_vec = ref_vec[:, :128]
    ref_class_vec = ref_vec[:, 128:]
## View image correspond to the reference code
ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
refimg = G.visualize(ref_vect).cpu()
centimg = ToPILImage()(refimg[0,:,:,:])
centimg.show(title="Center Reference Image")
#%% Visualize the Final Generation  together  with the center reference image. 
VisFinalGen = False
if VisFinalGen:
    #% If you want to regenerate the images from last generation here.
    print("Review the last generation codes w.r.t. the center code for manifold.")
    imgs_final = G.visualize_batch_np(np.concatenate((ref_noise_vec.repeat(25,axis=0), final_gen_codes[:,:]), axis=1))
    ToPILImage()(make_grid(imgs_final,nrow=5)).show()
    #G.visualize(torch.from_numpy(np.concatenate((ref_noise_vec.repeat(5,axis=0), final_gen_codes[:5,:]), axis=1)).float().cuda()).cpu()
    #ToPILImage()(make_grid(imgs.cpu())).show()
#%% Compute Hessian decomposition and get the vectors
Hess_method = "BP"  # "BackwardIter" "ForwardIter"
Hess_all = False # Set to False to reducerrrrrr computation time. 
t0 = time()
if Hess_method == "BP":
    print("Computing Hessian Decomposition Through auto-grad and full eigen decomposition.")
    classvec = torch.from_numpy(ref_class_vec).float().cuda()  # embed_mat[:, class_id:class_id+1].cuda().T
    noisevec = torch.from_numpy(ref_noise_vec).float().cuda()
    ref_vect = torch.cat((noisevec, classvec, ), dim=1).detach().clone()
    mov_vect = ref_vect.detach().clone().requires_grad_(True)
    #%
    imgs1 = G.visualize(ref_vect)
    if Hess_all:
        imgs2 = G.visualize(mov_vect)
        dsim = ImDist(imgs1, imgs2)
        H = get_full_hessian(dsim, mov_vect)  # 77sec to compute a Hessian. # 114sec on ML2a
        # ToPILImage()(imgs[0,:,:,:].cpu()).show()
        eigvals, eigvects = np.linalg.eigh(H)  # 75 ms
    #%
    noisevec.requires_grad_(True)
    classvec.requires_grad_(False)
    mov_vect = torch.cat((noisevec, classvec, ), dim=1)
    imgs2 = G.visualize(mov_vect)
    dsim = ImDist(imgs1, imgs2)
    H_nois = get_full_hessian(dsim, noisevec)  # 39.3 sec to compute a Hessian.# 59 sec on ML2a
    eigvals_nois, eigvects_nois = np.linalg.eigh(H_nois)  # 75 ms
    #%
    noisevec.requires_grad_(False)
    classvec.requires_grad_(True)
    mov_vect = torch.cat((noisevec, classvec, ), dim=1)
    imgs2 = G.visualize(mov_vect)
    dsim = ImDist(imgs1, imgs2)
    H_clas = get_full_hessian(dsim, classvec)  # 39.3 sec to compute a Hessian.
    eigvals_clas, eigvects_clas = np.linalg.eigh(H_clas)  # 75 ms
    classvec.requires_grad_(False)
    del dsim, imgs2, imgs1
    torch.cuda.empty_cache()
    if Hess_all:
        np.savez(join(summary_dir, "Hess_mat.npz"), H=H, eigvals=eigvals, eigvects=eigvects, 
             H_clas=H_clas, eigvals_clas=eigvals_clas, eigvects_clas=eigvects_clas, 
             H_nois=H_nois, eigvals_nois=eigvals_nois, eigvects_nois=eigvects_nois, 
             vect=ref_vect.cpu().numpy(), noisevec=noisevec.cpu().numpy(), classvec=classvec.cpu().numpy())
    else:
        np.savez(join(summary_dir, "Hess_mat.npz"), #H=H, eigvals=eigvals, eigvects=eigvects, 
             H_clas=H_clas, eigvals_clas=eigvals_clas, eigvects_clas=eigvects_clas, 
             H_nois=H_nois, eigvals_nois=eigvals_nois, eigvects_nois=eigvects_nois, 
             vect=ref_vect.cpu().numpy(), noisevec=noisevec.cpu().numpy(), classvec=classvec.cpu().numpy())
elif Hess_method == "BackwardIter":
    print("Computing Hessian Decomposition Through Lanczos decomposition on Backward HVP operator.")
    feat = torch.from_numpy(sphere_norm * PC1_vect).float().requires_grad_(False).cuda()
    eigenvals, eigenvecs = compute_hessian_eigenthings(G, feat, ImDist,
        num_eigenthings=128, mode="lanczos", use_gpu=True)
    eigenvals = eigenvals[::-1]
    eigenvecs = eigenvecs[::-1, :]

elif Hess_method == "ForwardIter": 
    print("Computing Hessian Decomposition Through Lanczos decomposition on Forward HVP operator.")
    pass

print("%.2f sec"% (time() - t0))  # 31.75 secs for 300 eig, 87.52 secs for 800 eigs. 
#%% Visualize spectrum
plt.figure(figsize=[8,5])
plt.subplot(1,2,1)
if Hess_all: plt.plot(eigvals[::-1], label="all")
plt.plot(eigvals_clas[::-1], label="class")
plt.plot(eigvals_nois[::-1], label="noise")
plt.ylabel("eigval");plt.legend()
plt.subplot(1,2,2)
if Hess_all: plt.plot(np.log10(eigvals[::-1]), label="all")
plt.plot(np.log10(eigvals_clas[::-1]), label="class")
plt.plot(np.log10(eigvals_nois[::-1]), label="noise")
plt.ylabel("log(eigval)");plt.legend()
plt.savefig(join(summary_dir, "spectrum.jpg"))
#%% Optional: Angle with PC1 vector
if evolspace == "BigGAN_class":
    innerprod2PC1 = PC1_vect @ eigvects_clas.T
elif evolspace == "BigGAN_noise":
    innerprod2PC1 = PC1_vect @ eigvects_nois.T
elif evolspace == "BigGAN":
    innerprod2PC1 = PC1_vect @ eigvects.T
print("Eigen vector: Innerproduct max %.3E min %.3E std %.3E"% (innerprod2PC1.max(), innerprod2PC1.min(), innerprod2PC1.std()))
print("EigenDecomposition of Hessian of Image Similarity Metric\nEigen value: Class space max %.3E min %.3E std %.3E; Noise space max %.3E min %.3E std %.3E"%
      (eigvals_clas.max(), eigvals_clas.min(), eigvals_clas.std(), eigvals_nois.max(), eigvals_nois.min(), eigvals_nois.std(), )) 
if Hess_all: 
    print("EigenDecomposition of Hessian of Image Similarity Metric\nEigen value: All: max %.3E min %.3E std %.3E"%
      (eigvals.max(), eigvals.min(), eigvals.std(),))

#%% Do interpolation along each axes
#%%
#classvec = torch.from_numpy(ref_class_vec).float().cuda()  # embed_mat[:, class_id:class_id+1].cuda().T
#noisevec = torch.from_numpy(ref_noise_vec).float().cuda()
if not exact_distance:
    #% Interpolation in the class space, but inversely scale the step size w.r.t. eigenvalue
    codes_all = []
    img_names = []
    scale = 5
    expon = 2.5
    eiglist_class = [0, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40,]
    for eigi in eiglist_class:#range(20):  # eigvects.shape[1] # 60, 80
        interp_class = LExpMap(classvec.cpu().numpy(), eigvects_clas[:, -eigi-1], 11, (-scale * eigvals_clas[-eigi-1] ** (-1/expon), scale * eigvals_clas[-eigi-1] ** (-1/expon)))
        interp_codes = np.hstack((noisevec.cpu().numpy().repeat(11, axis=0), interp_class, ))
        codes_all.append(interp_codes.copy())
        img_names.extend("class_eig%d_exp%.1f_lin%.1f.jpg"%(eigi, expon, dist) for dist in np.linspace(-scale, scale, 11))

    codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
    img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
    imggrid = make_grid(img_all, nrow=11)
    PILimg2 = ToPILImage()(imggrid)#.show()
    PILimg2.save(join(summary_dir, "eigvect_clas_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    npimgs = img_all.permute([2,3,1,0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
#%
    nrow = 11
    distmat = Hess_img_distmat(ImDist, img_all, nrow=11)
    plt.figure(figsize=[6, 6])
    plt.matshow(distmat, fignum=0)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nclass space exponent %.1f Scale%d "%(expon, scale, ))
    plt.savefig(join(summary_dir, "distmat_eigvect_clas_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    plt.show()
    #% Interpolation in the noise space
    codes_all = []
    img_names = []
    scale = 6
    expon = 3
    eiglist_noise = [0, 1, 2, 3, 4, 6, 10, 15, 20, 40]
    for eigi in eiglist_noise:#range(20):#eigvects_nois.shape[1]
    #    interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-4.5, 4.5))
        interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-scale * eigvals_nois[-eigi-1] ** (-1/expon), scale * eigvals_nois[-eigi-1] ** (-1/expon)))
        interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), ))
        codes_all.append(interp_codes.copy())
        img_names.extend("noise_eig%d_exp%.1f_lin%.1f.jpg"%(eigi, expon, dist) for dist in np.linspace(-scale, scale, 11))

    codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
    img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
    imggrid = make_grid(img_all, nrow=11)
    PILimg3 = ToPILImage()(imggrid)#.show()
    PILimg3.save(join(summary_dir, "eigvect_nois_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    npimgs = img_all.permute([2,3,1,0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
    #
    nrow = 11
    distmat = Hess_img_distmat(ImDist, img_all, nrow=11)
    # Show the image distance from center reference to each image around it.
    plt.figure(figsize=[6,6])
    plt.matshow(distmat, fignum=0)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nnoise space  Scale%d exponent %.1f"%(scale, expon))
    plt.savefig(join(summary_dir, "distmat_eigvect_nois_interp_exp%.1f_d%d.jpg"%(expon, scale)))
    plt.show()

else:  # exact_distance by line search
    from ImDist_Line_Searcher import find_level_step
    targ_val = np.array(target_distance)
    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
    refimg = G.visualize(ref_vect)
    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda()
    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda()

    space = "noise"
    imgall = None
    xtick_col = []
    dsim_col = []
    vecs_col = []
    img_names = []
    tick_labels = list(-targ_val[::-1]) + [0] + list(targ_val)  # -0.5, -0.4 ...  0.4, 0.5
    t0 = time()
    # eiglist_noise = [0, 1, 2, 3, 4, 5, 6, 8, 10, 20, 30, 40, ]#[0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 20, 30, 40, 50, 60, 70, 80]
    eiglist_noise = [0, 1, 4, 8, 10, 20, ]# [0, 1, 2, 4, 6, 8, 10, 20, ]
    for eigid in tqdm(eiglist_noise):  # range(128):  # #
        if space == "class":
            tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
        elif space == "noise":
            tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
        xtar_pos, ytar_pos, stepimgs_pos = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=True, maxdist=30)
        xtar_neg, ytar_neg, stepimgs_neg = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=False, maxdist=30)
        imgrow = torch.cat((torch.flip(stepimgs_neg, (0,)), refimg, stepimgs_pos)).cpu()
        xticks_row = xtar_neg[::-1] + [0.0] + xtar_pos
        dsim_row = list(ytar_neg[::-1]) + [0.0] + list(ytar_pos)
        vecs_row = torch.tensor(xticks_row).float().cuda().view(-1, 1) @ tan_vec + ref_vect

        xtick_col.append(xticks_row)
        dsim_col.append(dsim_row)
        vecs_col.append(vecs_row.cpu().numpy())
        img_names.extend("noise_eig%d_lin%.2f.jpg" % (eigid, dist) for dist in tick_labels)  # dsim_row)
        imgall = imgrow if imgall is None else torch.cat((imgall, imgrow))
        # Subsample the exploration and visualize the images on the line. ANd then render it into a movie. 
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)

    mtg1 = ToPILImage()(make_grid(imgall, nrow=2*len(target_distance)+1).cpu())  # 20sec for 13 rows not bad
    mtg1.show()
    mtg1.save(join(summary_dir, "noise_space_all_var.jpg"))
    npimgs = imgall.permute([2, 3, 1, 0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]),
                                                  np.uint8(npimgs[:, :, :, imgi] * 255))
    print(time() - t0)
    # %
    xtick_arr = np.array(xtick_col)
    dsim_arr = np.array(dsim_col)
    vecs_arr = np.array(vecs_col)
    np.savez(join(summary_dir, "noise_ImDist_root_data.npz"), xtick_arr=xtick_arr, dsim_arr=dsim_arr, vecs_arr=vecs_arr,
             targ_val=targ_val, eiglist=eiglist_noise)
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(xtick_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("L2 deviation from center")
    plt.xticks(range(len(eiglist_noise)), eiglist_noise)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Distance Travel Along Given Eigen vector to achieve certain Image Distance")
    plt.savefig(join(summary_dir, "noise_code_deviation.jpg"))
    plt.show()
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(dsim_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("Image Distance")
    plt.xticks(range(len(eiglist_noise)), eiglist_noise)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Achieved Image Distance Along Each Axis")
    plt.savefig(join(summary_dir, "noise_space_dist_curv.jpg"))
    plt.show()
    # %
    plt.figure()
    plt.matshow(dsim_arr, fignum=0)
    plt.xticks(range(len(tick_labels)), tick_labels)
    plt.yticks(range(len(eiglist_noise)), eiglist_noise)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nnoise space")
    plt.savefig(join(summary_dir, "noise_space_distmat.jpg"))
    plt.show()
    #%
    space = "class"
    imgall = None
    xtick_col = []
    dsim_col = []
    vecs_col = []
    img_names = []
    tick_labels = list(-targ_val[::-1]) + [0] + list(targ_val)
    t0 = time()
    #eiglist_class = [0, 1, 2, 3, 6, 9, 13, 17, 21, 25, 30, 40, 60, ] #[0, 1, 2, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40, 50, 60, 70, 80]
    eiglist_class = [0, 1, 2, 6, 17, 30, ]#[0, 1, 2, 3, 6, 9, 17, 21, 30, 60, ]
    for eigid in tqdm(eiglist_class):  # [0,1,2,3,4,5,6,7,8,10,20,30,
        # 40]:#
        if space == "class":
            tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
        elif space == "noise":
            tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
        xtar_pos, ytar_pos, stepimgs_pos = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=True, maxdist=30)
        xtar_neg, ytar_neg, stepimgs_neg = find_level_step(BGAN, ImDist, targ_val, ref_vect, tan_vec, refimg, iter=20,
                                                           pos=False, maxdist=30)
        imgrow = torch.cat((torch.flip(stepimgs_neg, (0,)), refimg, stepimgs_pos)).cpu()
        xticks_row = xtar_neg[::-1] + [0.0] + xtar_pos
        dsim_row = list(ytar_neg[::-1]) + [0.0] + list(ytar_pos)
        vecs_row = torch.tensor(xticks_row).float().cuda().view(-1, 1) @ tan_vec + ref_vect
        xtick_col.append(xticks_row)
        dsim_col.append(dsim_row)
        vecs_col.append(vecs_row.cpu().numpy())
        img_names.extend(
            "class_eig%d_lin%.2f.jpg" % (eigid, dist) for dist in tick_labels)  # np.linspace(-0.4, 0.4,11))
        imgall = imgrow if imgall is None else torch.cat((imgall, imgrow))
        # Subsample the exploration and visualize the images on the line. ANd then render it into a movie. 
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)

    mtg2 = ToPILImage()(make_grid(imgall, nrow=2*len(target_distance)+1).cpu())  # 20sec for 13 rows not bad
    mtg2.show()
    mtg2.save(join(summary_dir, "class_space_all_var.jpg"))
    npimgs = imgall.permute([2, 3, 1, 0]).numpy()
    for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]),
                                                  np.uint8(npimgs[:, :, :, imgi] * 255))
    print(time() - t0)
    # %
    xtick_arr = np.array(xtick_col)
    dsim_arr = np.array(dsim_col)
    vecs_arr = np.array(vecs_col)
    np.savez(join(summary_dir, "class_ImDist_root_data.npz"), xtick_arr=xtick_arr, dsim_arr=dsim_arr, vecs_arr=vecs_arr,
             targ_val=targ_val, eiglist=eiglist_class)
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(xtick_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("L2 deviation from center")
    plt.xticks(range(len(eiglist_class)), eiglist_class)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Distance Travel Along Given Eigen vector to achieve certain Image Distance")
    plt.savefig(join(summary_dir, "class_code_deviation.jpg"))
    plt.show()
    # %
    plt.figure(figsize=[10, 7])
    plt.plot(dsim_arr)
    plt.xlabel("Eigenvalue index")
    plt.ylabel("Image Distance")
    plt.xticks(range(len(eiglist_class)), eiglist_class)
    plt.legend(["Neg%.2f" % d for d in targ_val[::-1]] + ["orig"] + ["Pos%.2f" % d for d in targ_val])
    plt.title("Achieved Image Distance Along Each Axis")
    plt.savefig(join(summary_dir, "class_space_dist_curv.jpg"))
    plt.show()
    # %
    plt.figure()
    plt.matshow(dsim_arr, fignum=0)
    plt.xticks(range(len(tick_labels)), tick_labels)
    plt.yticks(range(len(eiglist_class)), eiglist_class)
    plt.colorbar()
    plt.title("Perceptual distance metric along each row\nnoise space")
    plt.savefig(join(summary_dir, "class_space_distmat.jpg"))
    plt.show()
#%%

#movie_dir = "."
if exact_distance:
    space = "noise"
#    targ_val = np.array(target_distance)
    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda()
    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda()
    
    data = np.load(join(summary_dir, "noise_ImDist_root_data.npz"))
    xtick_arr = data["xtick_arr"]
    eiglist_noise = data["eiglist"]
    for i, eigid in enumerate(eiglist_noise):
        tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
        xticks_row = xtick_arr[i, :]
        imgs, _, _, _ = LERP_img_row(ref_vect, tan_vec, target_distance, xticks_row, unit=0.08, density=5)
        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort_LERP" % (space, eigid), fps=30, writeframe=True)
    
    space = "class"
    data = np.load(join(summary_dir, "class_ImDist_root_data.npz"))
    xtick_arr = data["xtick_arr"]
    eiglist_class = data["eiglist"]
    for i, eigid in enumerate(eiglist_class):
        tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
        xticks_row = xtick_arr[i, :]
        imgs, _, _, _ = LERP_img_row(ref_vect, tan_vec, target_distance, xticks_row, unit=0.08, density=5)
        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort_LERP" % (space, eigid), fps=30, writeframe=True)
    

#%% Regenerate the movies with different fps
#summary_dir = os.path.join(backup_dir, "Hess_imgs", "summary")
#movie_dir = "."
#data = np.load(join(summary_dir, "Hess_mat.npz")) #H=H, eigvals=eigvals, eigvects=eigvects, 
#eigvects_clas = data["eigvects_clas"]
#eigvects_nois = data["eigvects_nois"]
#ref_noise_vec = data["noisevec"]
#ref_class_vec = data["classvec"]

#if exact_distance:
#    space = "noise"
##    targ_val = np.array(target_distance)
#    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda()
#    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda()
#    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda()
#    
#    data = np.load(join(summary_dir, "noise_ImDist_root_data.npz"))
#    xtick_arr = data["xtick_arr"]
#    eiglist_noise = data["eiglist"]
#    for i, eigid in enumerate(eiglist_noise):
#        tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1)
#        xticks_row = xtick_arr[i, :]
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)
#    
#    space = "class"
#    data = np.load(join(summary_dir, "class_ImDist_root_data.npz"))
#    xtick_arr = data["xtick_arr"]
#    eiglist_class = data["eiglist"]
#    for i, eigid in enumerate(eiglist_class):
#        tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1)
#        xticks_row = xtick_arr[i, :]
#        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5)
#        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30)
#%%
##% Interpolation in the class space
#codes_all = []
#img_names = []
#for eigi in [0, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40,]:#range(20):  # eigvects.shape[1] # 60, 80
#    interp_class = LExpMap(classvec.cpu().numpy(), eigvects_clas[:, -eigi-1], 11, (-2.5, 2.5))
#    interp_codes = np.hstack((noisevec.cpu().numpy().repeat(11, axis=0), interp_class, ))
#    codes_all.append(interp_codes.copy())
#    img_names.extend("class_eig%d_lin%.1f.jpg"%(eigi, dist) for dist in np.linspace(-2.5, 2.5, 11))
#codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
#img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
#imggrid = make_grid(img_all, nrow=11)
#PILimg2 = ToPILImage()(imggrid)#.show()
#PILimg2.save(join(summary_dir, "eigvect_clas_interp.jpg"))
#npimgs = img_all.permute([2,3,1,0]).numpy()
#for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
##%%
##% Interpolation in the noise space
#codes_all = []
#img_names = []
#for eigi in [0, 1, 2, 3, 4, 6, 10, 15, 20, 40]:#range(20):#eigvects_nois.shape[1]
#    interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-4.5, 4.5))
#    interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), ))
#    codes_all.append(interp_codes.copy())
#    img_names.extend("noise_eig%d_lin%.1f.jpg"%(eigi, dist) for dist in np.linspace(-4.5, 4.5, 11))
#
#codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
#img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
#imggrid = make_grid(img_all, nrow=11)
#PILimg3 = ToPILImage()(imggrid)#.show()
#PILimg3.save(join(summary_dir, "eigvect_nois_interp.jpg"))
#npimgs = img_all.permute([2,3,1,0]).numpy()
#for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))
#print("Spent %.1f sec from start" % (time() - T00))
#%% Interpolation in the noise space
#codes_all = []
#img_names = []
#for eigi in range(20):#eigvects_nois.shape[1]
#    interp_noise = SExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-1, 1))
#    interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), ))
#    codes_all.append(interp_codes.copy())
#    img_names.extend("noise_eig%d_sphang%d.jpg"%(eigi, angle) for angle in np.linspace(-90,90,11))
#
#codes_all_arr = np.concatenate(tuple(codes_all), axis=0)
#img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10)
#imggrid = make_grid(img_all, nrow=11)
#PILimg4 = ToPILImage()(imggrid)#.show()
#PILimg4.save(join(summary_dir, "eigvect_nois_sph_interp.jpg"))
#npimgs = img_all.permute([2,3,1,0]).numpy()
#for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255))

#%% 
#vec_norm = 220# sphere_norm
#eig_id = 0
#perturb_vect = eigenvecs[eig_id,:] # PC_vectors[1,:]    
#ang_step = 180 / 10
#theta_arr_deg = ang_step * np.arange(-5, 6)
#theta_arr = ang_step * np.arange(-5, 6) / 180 * np.pi
#codes_arc = np.array([np.cos(theta_arr), 
#                      np.sin(theta_arr) ]).T @ np.array([PC1_vect, perturb_vect])
#norms = np.linalg.norm(codes_arc, axis=1)
#codes_arc = codes_arc / norms[:,np.newaxis] * vec_norm
#imgs = G.visualize(torch.from_numpy(codes_arc).float().cuda())
#
#npimgs = imgs.detach().cpu().permute([2, 3, 1, 0]).numpy()
#for i in range(npimgs.shape[3]):
#    angle = theta_arr_deg[i]
#    imwrite(join(newimg_dir, "norm%d_eig%d_ang%d.jpg"%(vec_norm, eig_id, angle)), npimgs[:,:,:,i])
#
#img_list = [npimgs[:,:,:,i] for i in range(npimgs.shape[3])]
#mtg1 = build_montages(img_list, [256, 256], [11, 1])[0]
##imwrite(join(backup_dir, "norm%d_eig%d.jpg"%(vec_norm, eig_id)),mtg1)
#imwrite(join(newimg_dir, "norm%d_eig%d.jpg"%(vec_norm, eig_id)),mtg1)
#%%
# -*- coding: utf-8 -*- Created on Wed Jul 22 19:08:54 2020 Modified Oct. 14th  @author: <NAME>  Find important Nuisanced + Class transformations in Noise + Class space for a BigGAN evolved image. # backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_integrated\2020-06-01-09-46-37" # Put the backup folder and the thread to analyze here #backup_dir = r"C:\Users\Poncelab-ML2a\Documents\monkeylogic2\generate_BigGAN\2020-07-22-10-14-22" # backup_dir = r"C:\Users\Ponce lab\Documents\ml2a-monk\generate_BigGAN\2020-08-06-10-18-55"#2020-08-04-09-54-25"# # If True, it will try to read "scores_record.mat", from the backup folder and read "scores_record" # Else, it will use the unweighted mean code of the last generation as the center vector. # Need to run the BigGAN postHoc Analysis to save the `scores_record` mat and use this flag # Control if exact distance search is used or approximate heuristic rule is used. # 0.08, #target_distance = [0.09, 0.18, 0.27, 0.36, 0.45]  # if exact_distance is True it will search for images with these # distance to reference image along each eigenvector. #%% Prepare the generator model and perceptual loss networks # for video writing # import models # from PerceptualSimilarity folder # ImDist = models.PerceptualLoss(model='net-lin', net='squeeze', use_gpu=1, gpu_ids=[0]) # model_vgg = models.PerceptualLoss(model='net-lin', net='vgg', use_gpu=1, gpu_ids=[0]) #%% # BGAN = BigGAN.from_pretrained("biggan-deep-256") # for param in BGAN.parameters(): #     param.requires_grad_(False) # embed_mat = BGAN.embeddings.parameters().__next__().data # BGAN.cuda() #%% #%% Test code for hessian eigendecomposition #t0 = time() #feat = torch.randn((1, 4096), dtype=torch.float32).requires_grad_(False).cuda() #eigenvals, eigenvecs = compute_hessian_eigenthings(G, feat, ImDist, #    num_eigenthings=300, mode="lanczos", use_gpu=True,) #print(time() - t0,"\n")  # 81.02 s #%% Load the codes from the Backup folder load all the code mat file in the experiment folder and summarize it into nparrays     threadnum: can select one thread of the code if it's a parallel evolution. Usually, 0 or 1.          None for all threads. # make sure enough codes for requested size # if the summary table exist, just read from it! #%% Use Rank weight like CMAES #%% Compute Image distance using the ImDist img_all: a torch 4D array of shape [N, 3, 256, 256] on cpu()     nrow: specify how many images are generated in one axis.         It will arrange the images in `img_all` in a matrix and calculate the distance to the center image in each row. #%% Utility functions to create movies from a row of interpolation. # codes_row = ref_vect + torch.tensor(subsamp_ticks).unsqueeze(1).float().cuda() @ tan_vec.cuda() #    for i in range(len(targ_val_all)-1): #        seg = np.linspace(xticks_row[i], xticks_row[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False) #        tseg = np.linspace(targ_val_all[i], targ_val_all[i+1], int(abs((targ_val_all[i] - targ_val_all[i+1]))/unit * density) , endpoint=False) #        subsamp_ticks.extend(list(seg)) #        targ_ticks.extend(list(tseg)) #    subsamp_ticks.append(xticks_row[-1]) #    targ_ticks.append(targ_val_all[-1]) # codes_row = ref_vect + torch.tensor(subsamp_ticks).unsqueeze(1).float().cuda() @ tan_vec.cuda() #%% Load up the codes # Use penultimate generation to generate the center #%% If there is score of images, load them up here. And compute the weighted average code # block of scores is the number of block of codes -1 #%% PCA of the Existing code (adopted from Manifold experiment) #% Do PCA and find the major trend of evolution # decide which is the positive direction for PC1 #%% Prepare the center vector to use in Hessian computation. #   Use this vector as the reference vector (center) in the Hessian computation #   Before Aug. 6th, it's designed to use one code / image from the last generation and explore around it. #       this can be dangerous, sometimes one image will lose the major feature that we want in the population. (lose the 2 balls) #   From Aug. 6th on, we decided to use the mean code from the last generation, which has a higher probability of #   From the setting of the bhv2 files find the fixed noise vectors #% Choose the mean final generation codes as our reference # final_gen_codes[0:1, :] # final_gen_codes[0:1, :] # final_gen_codes[0:1, :] ## View image correspond to the reference code #%% Visualize the Final Generation  together  with the center reference image. #% If you want to regenerate the images from last generation here. #G.visualize(torch.from_numpy(np.concatenate((ref_noise_vec.repeat(5,axis=0), final_gen_codes[:5,:]), axis=1)).float().cuda()).cpu() #ToPILImage()(make_grid(imgs.cpu())).show() #%% Compute Hessian decomposition and get the vectors # "BackwardIter" "ForwardIter" # Set to False to reducerrrrrr computation time. # embed_mat[:, class_id:class_id+1].cuda().T #% # 77sec to compute a Hessian. # 114sec on ML2a # ToPILImage()(imgs[0,:,:,:].cpu()).show() # 75 ms #% # 39.3 sec to compute a Hessian.# 59 sec on ML2a # 75 ms #% # 39.3 sec to compute a Hessian. # 75 ms #H=H, eigvals=eigvals, eigvects=eigvects, # 31.75 secs for 300 eig, 87.52 secs for 800 eigs. #%% Visualize spectrum #%% Optional: Angle with PC1 vector #%% Do interpolation along each axes #%% #classvec = torch.from_numpy(ref_class_vec).float().cuda()  # embed_mat[:, class_id:class_id+1].cuda().T #noisevec = torch.from_numpy(ref_noise_vec).float().cuda() #% Interpolation in the class space, but inversely scale the step size w.r.t. eigenvalue #range(20):  # eigvects.shape[1] # 60, 80 #.show() #% #% Interpolation in the noise space #range(20):#eigvects_nois.shape[1] #    interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-4.5, 4.5)) #.show() # # Show the image distance from center reference to each image around it. # exact_distance by line search # -0.5, -0.4 ...  0.4, 0.5 # eiglist_noise = [0, 1, 2, 3, 4, 5, 6, 8, 10, 20, 30, 40, ]#[0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 20, 30, 40, 50, 60, 70, 80] # [0, 1, 2, 4, 6, 8, 10, 20, ] # range(128):  # # # dsim_row) # Subsample the exploration and visualize the images on the line. ANd then render it into a movie. #        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5) #        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30) # 20sec for 13 rows not bad # % # % # % # % #% #eiglist_class = [0, 1, 2, 3, 6, 9, 13, 17, 21, 25, 30, 40, 60, ] #[0, 1, 2, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40, 50, 60, 70, 80] #[0, 1, 2, 3, 6, 9, 17, 21, 30, 60, ] # [0,1,2,3,4,5,6,7,8,10,20,30, # 40]:# # np.linspace(-0.4, 0.4,11)) # Subsample the exploration and visualize the images on the line. ANd then render it into a movie. #        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5) #        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30) # 20sec for 13 rows not bad # % # % # % # % #%% #movie_dir = "." #    targ_val = np.array(target_distance) #%% Regenerate the movies with different fps #summary_dir = os.path.join(backup_dir, "Hess_imgs", "summary") #movie_dir = "." #data = np.load(join(summary_dir, "Hess_mat.npz")) #H=H, eigvals=eigvals, eigvects=eigvects, #eigvects_clas = data["eigvects_clas"] #eigvects_nois = data["eigvects_nois"] #ref_noise_vec = data["noisevec"] #ref_class_vec = data["classvec"] #if exact_distance: #    space = "noise" ##    targ_val = np.array(target_distance) #    ref_vect = torch.from_numpy(np.concatenate((ref_noise_vec, ref_class_vec), axis=1)).float().cuda() #    evc_clas_tsr = torch.from_numpy(eigvects_clas[:, ::-1].copy()).float().cuda() #    evc_nois_tsr = torch.from_numpy(eigvects_nois[:, ::-1].copy()).float().cuda() # #    data = np.load(join(summary_dir, "noise_ImDist_root_data.npz")) #    xtick_arr = data["xtick_arr"] #    eiglist_noise = data["eiglist"] #    for i, eigid in enumerate(eiglist_noise): #        tan_vec = torch.cat((evc_nois_tsr[:, eigid:eigid + 1].t(), torch.zeros(1, 128).cuda()), dim=1) #        xticks_row = xtick_arr[i, :] #        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5) #        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30) # #    space = "class" #    data = np.load(join(summary_dir, "class_ImDist_root_data.npz")) #    xtick_arr = data["xtick_arr"] #    eiglist_class = data["eiglist"] #    for i, eigid in enumerate(eiglist_class): #        tan_vec = torch.cat((torch.zeros(1, 128).cuda(), evc_clas_tsr[:, eigid:eigid + 1].t()), dim=1) #        xticks_row = xtick_arr[i, :] #        imgs, _, _, _ = subsampled_img_row(ref_vect, tan_vec, targ_val, xticks_row, unit=0.08, density=5) #        createSinuMovie(imgs, movie_dir, savenm="%s_eig%d_shortshort" % (space, eigid), fps=30) #%% ##% Interpolation in the class space #codes_all = [] #img_names = [] #for eigi in [0, 3, 6, 9, 11, 13, 15, 17, 19, 21, 25, 40,]:#range(20):  # eigvects.shape[1] # 60, 80 #    interp_class = LExpMap(classvec.cpu().numpy(), eigvects_clas[:, -eigi-1], 11, (-2.5, 2.5)) #    interp_codes = np.hstack((noisevec.cpu().numpy().repeat(11, axis=0), interp_class, )) #    codes_all.append(interp_codes.copy()) #    img_names.extend("class_eig%d_lin%.1f.jpg"%(eigi, dist) for dist in np.linspace(-2.5, 2.5, 11)) #codes_all_arr = np.concatenate(tuple(codes_all), axis=0) #img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10) #imggrid = make_grid(img_all, nrow=11) #PILimg2 = ToPILImage()(imggrid)#.show() #PILimg2.save(join(summary_dir, "eigvect_clas_interp.jpg")) #npimgs = img_all.permute([2,3,1,0]).numpy() #for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255)) ##%% ##% Interpolation in the noise space #codes_all = [] #img_names = [] #for eigi in [0, 1, 2, 3, 4, 6, 10, 15, 20, 40]:#range(20):#eigvects_nois.shape[1] #    interp_noise = LExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-4.5, 4.5)) #    interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), )) #    codes_all.append(interp_codes.copy()) #    img_names.extend("noise_eig%d_lin%.1f.jpg"%(eigi, dist) for dist in np.linspace(-4.5, 4.5, 11)) # #codes_all_arr = np.concatenate(tuple(codes_all), axis=0) #img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10) #imggrid = make_grid(img_all, nrow=11) #PILimg3 = ToPILImage()(imggrid)#.show() #PILimg3.save(join(summary_dir, "eigvect_nois_interp.jpg")) #npimgs = img_all.permute([2,3,1,0]).numpy() #for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255)) #print("Spent %.1f sec from start" % (time() - T00)) #%% Interpolation in the noise space #codes_all = [] #img_names = [] #for eigi in range(20):#eigvects_nois.shape[1] #    interp_noise = SExpMap(noisevec.cpu().numpy(), eigvects_nois[:, -eigi-1], 11, (-1, 1)) #    interp_codes = np.hstack((interp_noise, classvec.cpu().numpy().repeat(11, axis=0), )) #    codes_all.append(interp_codes.copy()) #    img_names.extend("noise_eig%d_sphang%d.jpg"%(eigi, angle) for angle in np.linspace(-90,90,11)) # #codes_all_arr = np.concatenate(tuple(codes_all), axis=0) #img_all = G.visualize_batch_np(codes_all_arr, truncation=0.7, B=10) #imggrid = make_grid(img_all, nrow=11) #PILimg4 = ToPILImage()(imggrid)#.show() #PILimg4.save(join(summary_dir, "eigvect_nois_sph_interp.jpg")) #npimgs = img_all.permute([2,3,1,0]).numpy() #for imgi in range(npimgs.shape[-1]):  imwrite(join(newimg_dir, img_names[imgi]), np.uint8(npimgs[:,:,:,imgi]*255)) #%% #vec_norm = 220# sphere_norm #eig_id = 0 #perturb_vect = eigenvecs[eig_id,:] # PC_vectors[1,:] #ang_step = 180 / 10 #theta_arr_deg = ang_step * np.arange(-5, 6) #theta_arr = ang_step * np.arange(-5, 6) / 180 * np.pi #codes_arc = np.array([np.cos(theta_arr), #                      np.sin(theta_arr) ]).T @ np.array([PC1_vect, perturb_vect]) #norms = np.linalg.norm(codes_arc, axis=1) #codes_arc = codes_arc / norms[:,np.newaxis] * vec_norm #imgs = G.visualize(torch.from_numpy(codes_arc).float().cuda()) # #npimgs = imgs.detach().cpu().permute([2, 3, 1, 0]).numpy() #for i in range(npimgs.shape[3]): #    angle = theta_arr_deg[i] #    imwrite(join(newimg_dir, "norm%d_eig%d_ang%d.jpg"%(vec_norm, eig_id, angle)), npimgs[:,:,:,i]) # #img_list = [npimgs[:,:,:,i] for i in range(npimgs.shape[3])] #mtg1 = build_montages(img_list, [256, 256], [11, 1])[0] ##imwrite(join(backup_dir, "norm%d_eig%d.jpg"%(vec_norm, eig_id)),mtg1) #imwrite(join(newimg_dir, "norm%d_eig%d.jpg"%(vec_norm, eig_id)),mtg1) #%%