#!/usr/bin/env python
# encoding: utf-8

import argparse
import os
import numpy as np
from sklearn.datasets import load_svmlight_file
import scipy.sparse as sp
import pickle
from sklearn.preprocessing import normalize
from tqdm import tqdm

import torch
from transformers import *


def main(args):
    if args.embed_type == "text-emb":
        label_text_list = [line.strip() for line in open("./{}/mapping/label_map.txt".format(args.dataset), "r")]
        n_label = len(label_text_list)

        # xlnet-large-cased tokenizer
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        # tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
        # model = XLNetModel.from_pretrained('xlnet-base-cased')
        # tokenizer = BertTokenizer.from_pretrained('bert-large-cased-whole-word-masking')
        # model = BertModel.from_pretrained('bert-large-cased-whole-word-masking')
        tokenizer = RobertaTokenizer.from_pretrained("roberta-large")
        model = RobertaModel.from_pretrained("roberta-large")
        model = model.to(device)

        # get label embedding
        label_embedding = []
        for idx in tqdm(range(n_label)):
            inputs = torch.tensor([tokenizer.encode(label_text_list[idx])])
            inputs = inputs.to(device)
            with torch.no_grad():
                last_hidden_states = model(inputs)[0]  # [1, seq_len, hidden_dim]
                seq_embedding = last_hidden_states.mean(dim=1)
            label_embedding.append(seq_embedding)
        label_embedding = torch.cat(label_embedding, dim=0)
        label_embedding = label_embedding.cpu().numpy()
        label_embedding = sp.csr_matrix(label_embedding)
        label_embedding = normalize(label_embedding, axis=1, norm="l2")

    elif args.embed_type == "pifa":
        # load TF-IDF and label matrix
        X = sp.load_npz("./{}/X.trn.npz".format(args.dataset))
        Y = sp.load_npz("./{}/Y.trn.npz".format(args.dataset))
        assert Y.getformat() == "csr"
        print("X", type(X), X.shape)
        print("Y", type(Y), Y.shape)
        # create label embedding
        Y_avg = normalize(Y, axis=1, norm="l2")
        label_embedding = sp.csr_matrix(Y_avg.T.dot(X))
        label_embedding = normalize(label_embedding, axis=1, norm="l2")

    elif args.embed_type == "pifa-neural":
        # load neural embedding from matcher
        X_trn = np.load(args.trn_embedding_npy)
        Y_trn = sp.load_npz("./{}/Y.trn.npz".format(args.dataset))
        print("X_trn", type(X_trn), X_trn.shape)
        print("Y_trn", type(Y_trn), Y_trn.shape)
        Y_avg = normalize(Y_trn, axis=1, norm="l2")
        label_embedding = sp.csr_matrix(Y_avg.T.dot(X_trn))
        label_embedding = normalize(label_embedding, axis=1, norm="l2")

    else:
        raise NotImplementedError("unknown embed_type {}".format(args.embed_type))

    # save label embedding
    print("label_embedding", type(label_embedding), label_embedding.shape)
    label_embedding_path = "{}/L.{}.npz".format(args.dataset, args.embed_type)
    sp.save_npz(label_embedding_path, label_embedding)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-d", "--dataset", type=str, required=True, help="dataset name: [ Eurlex-4K | Wiki10-31K | AmazonCat-13K | Wiki-500K ]"
    )
    parser.add_argument(
        "-e", "--embed-type", type=str, required=True, help="label embedding type: [ pifa | pifa-tst | pifa-neural | text-emb"
    )
    parser.add_argument("-x1", "--trn_embedding_npy", type=str, default=None, help="train embedding extracted from neural matcher")
    parser.add_argument("-x2", "--tst_embedding_npy", type=str, default=None, help="test embedding extracted from neural matcher")
    args = parser.parse_args()
    print(args)
    main(args)
