# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

import torch
from torch.utils.data import DataLoader
import torch.nn.functional as F
from sentence_transformers import models, losses, util, SentenceTransformer, LoggingHandler
from sentence_transformers.cross_encoder import CrossEncoder
from sentence_transformers.cross_encoder.evaluation import CECorrelationEvaluator
from sentence_transformers.readers import InputExample
from datetime import datetime
import argparse
import logging 
import sys
import random
import tqdm
import math
import os
import numpy as np
import pandas as pd

# import from local codes
from data import load_data
from eval import eval_encoder

from sentence_transformers_ext.bi_encoder_eval import (
    EmbeddingSimilarityEvaluator, 
    EmbeddingSimilarityEvaluatorEnsemble,
    EmbeddingSimilarityEvaluatorAUC,
    EmbeddingSimilarityEvaluatorAUCEnsemble
)
from sentence_transformers_ext.cross_encoder_eval import (
    CECorrelationEvaluatorEnsemble, 
    CECorrelationEvaluatorAUC,
    CECorrelationEvaluatorAUCEnsemble
)

parser = argparse.ArgumentParser()
parser.add_argument("--model_name_or_path", type=str, 
        default="princeton-nlp/unsup-simcse-bert-base-uncased",
        help="Transformers' model name or path")
parser.add_argument("--task", type=str, default='sts', 
        help='{sts|sickr|sts_sickr|qqp|qnli|mrpc|snli|custom}')
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--cycle", type=int, default=3)
parser.add_argument("--bi_encoder_pooling_mode", type=str, default='cls', 
        help="{cls|mean}")
parser.add_argument("--num_epochs_cross_encoder", type=int, default=1)
parser.add_argument("--num_epochs_bi_encoder", type=int, default=10)
parser.add_argument("--batch_size_cross_encoder", type=int, default=32)
parser.add_argument("--batch_size_bi_encoder", type=int, default=128)
parser.add_argument("--init_with_new_models", action="store_true")
parser.add_argument("--random_seed", type=int, default=2021)
#parser.add_argument("--use_raw_data_from_all_tasks", action="store_true")
parser.add_argument("--add_snli_data", type=int, default=0)
parser.add_argument("--custom_corpus_path", type=str, default=None)
parser.add_argument("--num_training_pairs", type=int, default=None)
parser.add_argument("--save_all_predictions", action="store_true")
parser.add_argument("--quick_test", action="store_true")


args = parser.parse_args()
print (args)

torch.manual_seed(args.random_seed)

#### Just some code to print debug information to stdout
logging.basicConfig(format="%(asctime)s - %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S",
                    level=logging.INFO,
                    handlers=[LoggingHandler()])


### read datasets
all_pairs, all_test, dev_samples = load_data(args.task, fpath=args.custom_corpus_path)

# load_pairs from other tasks
"""
if args.use_raw_data_from_all_tasks:
    all_pairs_qqp, _, _ = load_data("qqp")
    all_pairs_qnli, _, _ = load_data("qnli")
    all_pairs_mrpc, _, _ = load_data("mrpc")
    all_pairs = all_pairs + all_pairs_qqp + all_pairs_qnli + all_pairs_mrpc
"""

if args.add_snli_data != 0:
    random.seed(args.random_seed)
    all_pairs_snli, _, _ = load_data("snli")
    all_pairs_snli_sampled = random.sample(all_pairs_snli, args.add_snli_data)
    all_pairs = all_pairs + all_pairs_snli_sampled

if args.quick_test:
    all_pairs = all_pairs[:1000] # for quick testing

# randomly select training pairs for control study
if args.num_training_pairs is not None:
    print ("before sampling |raw sentence pairs|:", len(all_pairs))
    if args.num_training_pairs == -1:
        # use all
        pass
    else:    
        random.seed(args.random_seed)
        all_pairs = random.sample(all_pairs, args.num_training_pairs)

print ("|raw sentence pairs|:", len(all_pairs))
print ("|dev set|:", len(dev_samples))
for key in all_test:
    print ("|test set: %s|" % key, len(all_test[key]))

model_name = args.model_name_or_path #"princeton-nlp/unsup-simcse-bert-base-uncased"
simcse2base = {
    "princeton-nlp/unsup-simcse-roberta-base": "roberta-base", 
    "princeton-nlp/unsup-simcse-roberta-large": "roberta-large", 
    "princeton-nlp/unsup-simcse-bert-base-uncased": "bert-base-uncased",
    "princeton-nlp/unsup-simcse-bert-large-uncased": "bert-large-uncased"}
batch_size_cross_encoder = args.batch_size_cross_encoder
batch_size_bi_encoder = args.batch_size_bi_encoder
num_epochs_cross_encoder = args.num_epochs_cross_encoder 
num_epochs_bi_encoder = args.num_epochs_bi_encoder
max_seq_length = 32
total_cycle = args.cycle
device=args.device

logging.info ("########## load base model and evaluate ##########")

###### Bi-encoder (sentence-transformers) ######
logging.info(f"Loading bi-encoder model: {model_name}")
# Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for mapping tokens to embeddings
word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)

# Apply mean pooling to get one fixed sized sentence vector
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode=args.bi_encoder_pooling_mode)

bi_encoder = SentenceTransformer(modules=[word_embedding_model, pooling_model], device=device)

# eval bi-encoder
logging.info ("Evaluate bi-encoder...")
eval_encoder(all_test, bi_encoder, task=args.task, enc_type="bi")

start_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

bi_encoder_dev_scores = []
cross_encoder_dev_scores = []
bi_encoder_path = model_name
all_predictions_bi_encoder = {}
all_predictions_cross_encoder = {}


for cycle in range(total_cycle):
    cycle += 1
    logging.info (f"########## cycle {cycle:.0f} starts ##########")

    ###### label data with bi-encoder ######
    # label sentence pairs with bi-encoder
    logging.info ("Label sentence pairs with bi-encoder...")

    # Two lists of sentences
    sents1 = [p[0] for p in all_pairs]
    sents2 = [p[1] for p in all_pairs]

    #Compute embedding for both lists
    embeddings1 = bi_encoder.encode(sents1, convert_to_tensor=True)
    embeddings2 = bi_encoder.encode(sents2, convert_to_tensor=True)

    #Compute cosine-similarits
    cosine_scores = F.cosine_similarity(embeddings1, embeddings2)

    # save the predictions
    all_predictions_bi_encoder["bi_encoder_cycle_"+str(cycle-1)] = cosine_scores.cpu().numpy()

    # form (self-labelled) train set
    train_samples = []

    for i in range(len(sents1)):
        if args.task in ["qnli"]:
            train_samples.append(InputExample(texts=[sents1[i], sents2[i]], label=cosine_scores[i]))
        else:
            train_samples.append(InputExample(texts=[sents1[i], sents2[i]], label=cosine_scores[i]))
            train_samples.append(InputExample(texts=[sents2[i], sents1[i]], label=cosine_scores[i]))

    del bi_encoder, embeddings1, embeddings2, cosine_scores
    torch.cuda.empty_cache()

    ###### Cross-encoder learning ######
    if args.init_with_new_models:
        bi_encoder_path = simcse2base[model_name] #model_name # always use new model (PLM)
    logging.info(f"Loading cross-encoder model: {bi_encoder_path}")
    # Use Huggingface/transformers model (like BERT, RoBERTa, XLNet, XLM-R) for cross-encoder model
    cross_encoder = CrossEncoder(bi_encoder_path, num_labels=1, device=device, max_length=64)

    # We wrap gold_samples (which is a List[InputExample]) into a pytorch DataLoader
    train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=batch_size_cross_encoder)

    # We add an evaluator, which evaluates the performance during training
    if args.task in ["sts", "sickr", "sts_sickr", "custom"]:
        evaluator = CECorrelationEvaluator.from_input_examples(dev_samples, name='dev')
    else:
        evaluator = CECorrelationEvaluatorAUC.from_input_examples(dev_samples, name='dev')

    # Configure the training
    warmup_steps = math.ceil(len(train_dataloader) * num_epochs_cross_encoder * 0.1) #10% of train data for warm-up
    logging.info(f"Warmup-steps: {warmup_steps}")

    cross_encoder_path = f"output/cross-encoder/" \
        f"{args.task}_cycle{cycle}_{model_name.replace('/', '-')}-{start_time}" 

    # Train the cross-encoder model
    cross_encoder.fit(
            train_dataloader=train_dataloader,
            evaluator=evaluator,
            evaluation_steps=200,
            use_amp=True,
            epochs=num_epochs_cross_encoder,
            warmup_steps=warmup_steps,
            output_path=cross_encoder_path)



    cross_encoder = CrossEncoder(cross_encoder_path, max_length=64, device=device)
    #cross_encoder = CrossEncoder(cross_encoder_path, device=device)

    # eval cross-encoder
    dev_score = evaluator(cross_encoder)
    cross_encoder_dev_scores.append(dev_score)
    logging.info (f"***** dev's spearman's rho: {dev_score:.4f} *****")

    ###### label data with cross-encoder ######
    # label sentence pairs with cross-encoder
    logging.info ("Label sentence pairs with cross-encoder...")
    silver_scores = cross_encoder.predict(all_pairs)
    silver_samples = list(InputExample(texts=[data[0], data[1]], label=score) for \
        data, score in zip(all_pairs, silver_scores))

    del cross_encoder
    torch.cuda.empty_cache()

    # save the predictions
    all_predictions_cross_encoder["cross_encoder_cycle_"+str(cycle)] = silver_scores

    ###### Bi-encoder learning ######
    if args.init_with_new_models:
        cross_encoder_path = model_name # always use new model (SimCSE)
    logging.info(f"Loading bi-encoder model: {cross_encoder_path}") 
    word_embedding_model = models.Transformer(cross_encoder_path, max_seq_length=32)
    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode=args.bi_encoder_pooling_mode)
    bi_encoder = SentenceTransformer(modules=[word_embedding_model, pooling_model], device=device)

    train_dataloader = DataLoader(silver_samples, shuffle=True, batch_size=batch_size_bi_encoder)
    train_loss = losses.CosineSimilarityLoss(model=bi_encoder)

    if args.task in ["sts", "sickr", "sts_sickr", "custom"]:
        evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_samples, name='dev')
    else:
        evaluator = EmbeddingSimilarityEvaluatorAUC.from_input_examples(dev_samples, name='dev')

    # Configure the training.
    #warmup_steps = math.ceil(len(train_dataloader) * num_epochs  * 0.1) #10% of train data for warm-up
    #logging.info(f"Warmup-steps: {warmup_steps}")

    bi_encoder_path = f"output/bi-encoder/" \
        f"{args.task}_cycle{cycle}_{model_name.replace('/', '-')}-{start_time}" 
    
    # Train the bi-encoder model
    bi_encoder.fit(
          train_objectives=[(train_dataloader, train_loss)],
          evaluator=evaluator,
          epochs=num_epochs_bi_encoder,
          evaluation_steps=200,
          warmup_steps=0,
          output_path=bi_encoder_path,
          optimizer_params= {"lr": 5e-5},
          use_amp=True,
          )

    bi_encoder = SentenceTransformer(bi_encoder_path, device=device)

    # eval bi-encoder
    dev_score = evaluator(bi_encoder)
    bi_encoder_dev_scores.append(dev_score)
    logging.info (f"***** dev's spearman's rho: {dev_score:.4f} *****")

logging.info (bi_encoder_dev_scores)
logging.info (cross_encoder_dev_scores)

# best bi-encoder
best_cycle_bi_encoder = np.argmax(bi_encoder_dev_scores)+1
best_cycle_bi_encoder_path = f"output/bi-encoder/"\
    f"{args.task}_cycle{best_cycle_bi_encoder}_{model_name.replace('/', '-')}-{start_time}" 
# eval bi-encoder
logging.info (f"Evaluate best bi-encoder (from cycle {best_cycle_bi_encoder})...")
bi_encoder = SentenceTransformer(best_cycle_bi_encoder_path, device=device)
logging.info (best_cycle_bi_encoder_path)
eval_encoder(all_test, bi_encoder, task=args.task, enc_type="bi")


# best cross-encoder
best_cycle_cross_encoder = np.argmax(cross_encoder_dev_scores)+1
best_cycle_cross_encoder_path = f"output/cross-encoder/"\
    f"{args.task}_cycle{best_cycle_cross_encoder}_{model_name.replace('/', '-')}-{start_time}" 
# eval cross-encoder
logging.info (f"Evaluate best cross-encoder (from cycle {best_cycle_cross_encoder})...")
logging.info (best_cycle_cross_encoder_path)
#cross_encoder = CrossEncoder(best_cycle_cross_encoder_path, max_length=64, device=device)
cross_encoder = CrossEncoder(best_cycle_cross_encoder_path, device=device)
eval_encoder(all_test, cross_encoder, task=args.task, enc_type="cross")

# save all predictions
best_cycle_bi_encoder_csv_path = f"output/bi-encoder/" \
    f"{args.task}_cycle{best_cycle_bi_encoder}_{model_name.replace('/', '-')}-{start_time}_all_preds.csv" 
best_cycle_cross_encoder_csv_path = f"output/cross-encoder/" \
    f"{args.task}_cycle{best_cycle_cross_encoder}_{model_name.replace('/', '-')}-{start_time}_all_preds.csv"

if args.save_all_predictions:
    bi_pred_df = pd.DataFrame(np.array(list(all_predictions_bi_encoder.values())).T, columns=list(all_predictions_bi_encoder.keys()))
    cross_pred_df = pd.DataFrame(np.array(list(all_predictions_cross_encoder.values())).T, columns=list(all_predictions_cross_encoder.keys()))
    bi_pred_df.to_csv(best_cycle_bi_encoder_csv_path, index=False) 
    cross_pred_df.to_csv(best_cycle_cross_encoder_csv_path, index=False) 
    
logging.info ("\n")
print (args)
logging.info ("\n")
logging.info ("***** END *****")
