"""
  This script provides an example to wrap TencentPretrain for feature extraction.
"""
import sys
import os
import torch
import torch.nn as nn
import argparse
import numpy as np

tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)

from tencentpretrain.embeddings import *
from tencentpretrain.encoders import *
from tencentpretrain.targets import *
from tencentpretrain.utils.constants import *
from tencentpretrain.utils import *
from tencentpretrain.utils.config import load_hyperparam
from tencentpretrain.utils.misc import pooling
from tencentpretrain.model_loader import load_model
from tencentpretrain.opts import infer_opts, tokenizer_opts


def batch_loader(batch_size, src, seg):
    instances_num = src.size(0)
    for i in range(instances_num // batch_size):
        src_batch = src[i * batch_size : (i + 1) * batch_size]
        seg_batch = seg[i * batch_size : (i + 1) * batch_size]
        yield src_batch, seg_batch
    if instances_num > instances_num // batch_size * batch_size:
        src_batch = src[instances_num // batch_size * batch_size:]
        seg_batch = seg[instances_num // batch_size * batch_size:]
        yield src_batch, seg_batch


def read_dataset(args, path):
    dataset = []
    PAD_ID = args.tokenizer.vocab.get(PAD_TOKEN)
    with open(path, mode="r", encoding="utf-8") as f:
        for line in f:
            src = args.tokenizer.convert_tokens_to_ids(args.tokenizer.tokenize(line))
            if len(src) == 0:
                continue
            src = args.tokenizer.convert_tokens_to_ids([CLS_TOKEN]) + src + args.tokenizer.convert_tokens_to_ids([SEP_TOKEN])
            seg = [1] * len(src)

            if len(src) > args.seq_length:
                src = src[:args.seq_length]
                seg = seg[:args.seq_length]
            while len(src) < args.seq_length:
                src.append(PAD_ID)
                seg.append(PAD_ID)
            dataset.append((src, seg))
    return dataset


class FeatureExtractor(torch.nn.Module):    
    def __init__(self, args):
        super(FeatureExtractor, self).__init__()
        self.embedding = Embedding(args)
        for embedding_name in args.embedding:
            tmp_emb = str2embedding[embedding_name](args, len(args.tokenizer.vocab))
            self.embedding.update(tmp_emb, embedding_name)
        self.encoder = str2encoder[args.encoder](args)
        self.pooling_type = args.pooling

    def forward(self, src, seg):
        emb = self.embedding(src, seg)
        output = self.encoder(emb, seg)
        output = pooling(output, seg, self.pooling_type)

        return output        


class WhiteningHandle(torch.nn.Module):
    """
    Whitening operation.
    @ref: https://github.com/bojone/BERT-whitening/blob/main/demo.py
    """
    def __init__(self, args, vecs):
        super(WhiteningHandle, self).__init__()
        self.kernel, self.bias = self._compute_kernel_bias(vecs)

    def forward(self, vecs, n_components=None, normal=True, pt=True):
        vecs = self._format_vecs_to_np(vecs)
        vecs = self._transform(vecs, n_components)
        vecs = self._normalize(vecs) if normal else vecs
        vecs = torch.tensor(vecs) if pt else vecs
        return vecs

    def _compute_kernel_bias(self, vecs):
        vecs = self._format_vecs_to_np(vecs)
        mu = vecs.mean(axis=0, keepdims=True)
        cov = np.cov(vecs.T)
        u, s, vh = np.linalg.svd(cov)
        W = np.dot(u, np.diag(1 / np.sqrt(s)))
        return W, -mu

    def _transform(self, vecs, n_components):
        w = self.kernel[:, :n_components] \
                if isinstance(n_components, int) else self.kernel
        return (vecs + self.bias).dot(w)

    def _normalize(self, vecs):
        return vecs / (vecs**2).sum(axis=1, keepdims=True)**0.5

    def _format_vecs_to_np(self, vecs):
        vecs_np = []
        for vec in vecs:
            if isinstance(vec, list):
                vec = np.array(vec)
            elif torch.is_tensor(vec):
                vec = vec.detach().numpy()
            elif isinstance(vec, np.ndarray):
                vec = vec
            else:
                raise Exception('Unknown vec type.')
            vecs_np.append(vec)
        vecs_np = np.array(vecs_np)
        return vecs_np


if __name__ == '__main__':
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    infer_opts(parser)

    parser.add_argument("--whitening_size", type=int, default=None, help="Output vector size after whitening.")

    tokenizer_opts(parser)

    args = parser.parse_args()
    args = load_hyperparam(args)

    args.tokenizer = str2tokenizer[args.tokenizer](args)

    # Build feature extractor model.
    model = FeatureExtractor(args)
    model = load_model(model, args.load_model_path)

    # For simplicity, we use DataParallel wrapper to use multiple GPUs.
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    if torch.cuda.device_count() > 1:
        print("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
        model = nn.DataParallel(model)
    model.eval()

    dataset = read_dataset(args, args.test_path)

    src = torch.LongTensor([sample[0] for sample in dataset])
    seg = torch.LongTensor([sample[1] for sample in dataset])

    feature_vectors = []
    for i, (src_batch, seg_batch) in enumerate(batch_loader(args.batch_size, src, seg)):
        src_batch = src_batch.to(device)
        seg_batch = seg_batch.to(device)
        output = model(src_batch, seg_batch)
        feature_vectors.append(output.cpu().detach())
    feature_vectors = torch.cat(feature_vectors, 0)

    # Vector whitening.
    if args.whitening_size is not None:
        whitening = WhiteningHandle(args, feature_vectors)
        feature_vectors = whitening(feature_vectors, args.whitening_size, pt=True)

    print("The size of feature vectors (sentences_num * vector size): {}".format(feature_vectors.shape))
    torch.save(feature_vectors, args.prediction_path)