import argparse

import torch
from sklearn.decomposition import FastICA
import numpy as np
import random

def indipendent_components_decomposition(W, n_components):
      fast_ica = FastICA(n_components=n_components)
      fast_ica.fit(W)
      W_ = fast_ica.components_
      norm = np.linalg.norm(W_, axis = 1).reshape(-1, n_components)
      W_nomralize = W_ / norm.T
      indipendent_components = torch.from_numpy(W_nomralize.T).float()
      return indipendent_components

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Extract factor/eigenvectors of latent spaces using closed form factorization"
    )

    parser.add_argument(
        "-n", "--number_of_component", type=int, default=8, help="index of eigenvector"
    )

    parser.add_argument(
        "--out", type=str, default="factor.pt", help="name of the result factor file"
    )
    parser.add_argument("ckpt", type=str, help="name of the model checkpoint")
    parser.add_argument('--full_model', default=False, action='store_true')

    args = parser.parse_args()

    if args.full_model:
      state_dict = torch.load(args.ckpt).state_dict()
    else:
      state_dict = torch.load(args.ckpt)["g_ema"]
      
    modulate = {
        k: v
        for k, v in state_dict.items()
        if "modulation" in k and "to_rgbs" not in k and "weight" in k
    }

    weight_mat = []
    for k, v in modulate.items():
        weight_mat.append(v)

    W = torch.cat(weight_mat, 0)

    np.random.seed(0)
    random.seed(0)
    eigvec = indipendent_components_decomposition(W, args.number_of_component).to("cpu")

    torch.save({"ckpt": args.ckpt, "eigvec": eigvec}, args.out)

