
import argparse
from helper import load_image, save_image_file, inverse_normalize
import clip
import open_clip
import math
import torch
from sklearn.decomposition import TruncatedSVD

def zero_gradients(x):
    if isinstance(x, torch.Tensor):
        if x.grad is not None:
            x.grad.detach_()
            x.grad.zero_()
    elif isinstance(x, collections.abc.Iterable):
        for elem in x:
            zero_gradients(elem)

def compute_jacobian(inputs, output):
	"""
	:param inputs: Batch X Size (e.g. Depth X Width X Height)
	:param output: Batch X Classes
	:return: jacobian: Batch X Classes X Size
	"""
	inputs = inputs.requires_grad_(True)
	num_classes = output.size()[1]

	jacobian = torch.zeros(num_classes, *inputs.size())
	grad_output = torch.zeros(*output.size())
	if inputs.is_cuda:
		grad_output = grad_output.cuda()
		jacobian = jacobian.cuda()

	for i in range(num_classes):
		zero_gradients(inputs)
		grad_output.zero_()
		grad_output[:, i] = 1
		output.backward(grad_output, retain_graph = True)
		jacobian[i] = inputs.grad

	return torch.transpose(jacobian, dim0=0, dim1=1)



def f(x):
    return model2.encode_image(x)


def main():
	
    parser = argparse.ArgumentParser(description="J")
    parser.add_argument("--current_image_path", type=str, required=True, help="File path of the current image")

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    model, _, preprocess = open_clip.create_model_and_transforms('ViT-H-14', pretrained='laion2b_s32b_b79k', device=device)

    optimizer = EmbeddingOptimizer(model, args.learning_rate)

    x = load_image(args.current_image_path, preprocess, device)
    jacs = torch.autograd.functional.jacobian(f, x)
    
    
    svd1= TruncatedSVD(n_components=1024, n_iter=10, random_state=42)
    svd1.fit(jacb_matrix)
    print(svd1.singular_values_)
    np.savetxt("./zshot/res/singular_vals_church_to_tench.txt", svd1.singular_values_)


if __name__ == "__main__":
    main()
