from .dataset import TextTokenIdsCache, SequenceDataset, pack_tensor_2D
from .models.backbones import DistilBertDot, RobertaDot, BertDot
from .models.backbones.roberta_tokenizer import RobertaTokenizer

from tqdm.autonotebook import tqdm, trange
from collections import defaultdict
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader, RandomSampler
from transformers import (AdamW, get_linear_schedule_with_warmup,
    RobertaConfig, AutoConfig, AutoTokenizer)

import os
import torch
import random
import time
import faiss
import logging
import argparse
import numpy as np

logger = logging.getLogger(__name__)

def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def save_model(model, output_dir, save_name, args, optimizer=None):
    save_dir = os.path.join(output_dir, save_name)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    model_to_save = model.module if hasattr(model, 'module') else model  
    model_to_save.save_pretrained(save_dir)
    torch.save(args, os.path.join(save_dir, 'training_args.bin'))
    if optimizer is not None:
        torch.save(optimizer.state_dict(), os.path.join(save_dir, "optimizer.bin"))


class TrainQueryDataset(SequenceDataset):
    def __init__(self, queryids_cache, 
            rel_file, max_query_length):
        SequenceDataset.__init__(self, queryids_cache, max_query_length)
        self.reldict = defaultdict(list)
        for line in tqdm(open(rel_file), desc=os.path.split(rel_file)[1]):
            qid, _, pid, _ = line.split()
            qid, pid = int(qid), int(pid)
            self.reldict[qid].append((pid))

    def __getitem__(self, item):
        ret_val = super().__getitem__(item)
        ret_val['rel_ids'] = self.reldict[item]
        return ret_val


def get_collate_function(max_seq_length):
    cnt = 0
    def collate_function(batch):
        nonlocal cnt
        length = None
        if cnt < 10:
            length = max_seq_length
            cnt += 1

        input_ids = [x["input_ids"] for x in batch]
        attention_mask = [x["attention_mask"] for x in batch]
        data = {
            "input_ids": pack_tensor_2D(input_ids, default=1, 
                dtype=torch.int64, length=length),
            "attention_mask": pack_tensor_2D(attention_mask, default=0, 
                dtype=torch.int64, length=length),
        }
        qids = [x['id'] for x in batch]
        all_rel_pids = [x["rel_ids"] for x in batch]
        return data, all_rel_pids
    return collate_function  
    

def get_doc_embeds(psg_ids, pq_codes, centroids):
    M = centroids.shape[0]
    first_indices = torch.arange(M).to(psg_ids.device)
    first_indices = first_indices.expand(len(psg_ids), M).reshape(-1)
    second_indices = pq_codes[psg_ids].reshape(-1)
    embeddings = centroids[first_indices, second_indices].reshape(len(psg_ids), -1)
    return embeddings


def compute_loss(query_embeddings, pq_codes, centroids, 
        batch_neighbors, all_rel_pids, lambda_cut):
    loss = 0
    mrr = 0
    train_batch_size = len(batch_neighbors)
    for qembedding, retrieve_pids, cur_rel_pids in zip(
        query_embeddings, batch_neighbors, all_rel_pids):
        cur_rel_pids = list(set(cur_rel_pids))
        cur_rel_pids = torch.LongTensor(cur_rel_pids).to(batch_neighbors.device)
        target_labels = (retrieve_pids[:, None]==cur_rel_pids).any(-1)
        retrieved_rel_pids = retrieve_pids[target_labels]
        if len(retrieved_rel_pids) < len(cur_rel_pids):
            not_retrieved_rel_pids = cur_rel_pids[
                (cur_rel_pids[:, None]!=retrieved_rel_pids).all(-1)]
            assert len(not_retrieved_rel_pids) > 0
            retrieve_pids = torch.hstack([retrieve_pids, not_retrieved_rel_pids])
            target_labels = torch.hstack([target_labels, 
                torch.tensor([True]*len(not_retrieved_rel_pids)).to(target_labels.device)])
        position_matrix = (1+torch.arange(len(target_labels))).to(target_labels.device)
        if len(retrieved_rel_pids) > 0:
            first_rel_pos = position_matrix[target_labels][0].item()
            if first_rel_pos <= 10:
                mrr += 1/first_rel_pos 
        psg_embeddings = get_doc_embeds(retrieve_pids, pq_codes, centroids)
        cur_top_scores = (qembedding.reshape(1, -1) * psg_embeddings).sum(-1)
        rel_scores = cur_top_scores[target_labels]
        irrel_scores = cur_top_scores[~target_labels]
        rel_poses = position_matrix[target_labels]
        rel_weights = 1/rel_poses
        irrel_poses = position_matrix[~target_labels]
        irrel_weights = 1/irrel_poses
        if lambda_cut > 0:
            rel_pos_masks = (rel_poses <= lambda_cut).float()
            rel_weights *= rel_pos_masks
            irrel_pos_masks = (irrel_poses <= lambda_cut).float()
            irrel_weights *= irrel_pos_masks
        # relnum, irrelnum
        pair_diff = irrel_scores.reshape(1, -1) - rel_scores.reshape(-1, 1)
        pair_weights = torch.abs(irrel_weights.reshape(1, -1) - rel_weights.reshape(-1, 1))
        cur_loss = torch.log(1 + torch.exp(pair_diff)) * pair_weights
        loss += cur_loss.mean()
    mrr /= train_batch_size
    loss /= train_batch_size
    return loss, mrr


def train(args, model, pq_codes, centroid_embeds, opq_transform, opq_index, tokenizer):
    """ Train the model """
    ivf_index = faiss.downcast_index(opq_index.index)
    if args.gpu_search:
        res = faiss.StandardGpuResources()
        res.setTempMemory(128*1024*1024)
        co = faiss.GpuClonerOptions()
        co.useFloat16 = len(pq_codes) >= 56     
        gpu_index = faiss.index_cpu_to_gpu(res, 0, opq_index, co)

    tb_writer = SummaryWriter(args.log_dir)

    train_dataset = TrainQueryDataset(
        TextTokenIdsCache(args.preprocess_dir, "train-query"),
        os.path.join(args.preprocess_dir, "train-qrel.tsv"),
        args.max_seq_length
    )

    train_sampler = RandomSampler(train_dataset) 
    collate_fn = get_collate_function(args.max_seq_length)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, 
        batch_size=args.train_batch_size, collate_fn=collate_fn)

    t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
        {'params': [centroid_embeds], 'weight_decay': args.centroid_weight_decay, 'lr': args.centroid_lr},
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=args.lr, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps,
        num_training_steps=t_total)
    
    # multi-gpu training (should be after apex fp16 initialization)
    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)

    # Train!
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info("  Total train batch size (w. accumulation) = %d",
                   args.train_batch_size * args.gradient_accumulation_steps)
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    tr_loss, logging_loss = 0.0, 0.0
    tr_mrr, logging_mrr = 0.0, 0.0
    optimizer.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch")
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)  
    
    for epoch_idx, _ in enumerate(train_iterator):
        epoch_iterator = tqdm(train_dataloader, desc="Iteration")
        for step, (batch, all_rel_poffsets) in enumerate(epoch_iterator):

            batch = {k:v.to(args.model_device) for k, v in batch.items()}
            model.train()            
            query_embeddings = model(
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"])
            
            if args.gpu_search:
                batch_neighbors = gpu_index.search(
                    query_embeddings.detach().cpu().numpy(), args.loss_neg_topK)[1]
            else:
                batch_neighbors = opq_index.search(
                    query_embeddings.detach().cpu().numpy(), args.loss_neg_topK)[1]
            batch_neighbors = torch.tensor(batch_neighbors).to(model.device)
            loss, mrr = compute_loss(
                    query_embeddings @ opq_transform.T, 
                    pq_codes, centroid_embeds,
                    batch_neighbors, all_rel_poffsets, args.lambda_cut)   
            tr_mrr += mrr
            loss /= args.gradient_accumulation_steps
            loss.backward()

            tr_loss += loss.item()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            # if args.n_gpu > 1:
            #     loss = loss.mean() # mean() to average on multi-gpu parallel (not distributed) training
            if (step + 1) % args.gradient_accumulation_steps == 0:
                optimizer.step()
                scheduler.step()  # Update learning rate schedule
                optimizer.zero_grad()
                # model.zero_grad()
                global_step += 1
                faiss.copy_array_to_vector(
                    centroid_embeds.detach().cpu().numpy().ravel(), 
                    ivf_index.pq.centroids)
                
                if args.gpu_search:
                    gpu_index = None
                    gpu_index = faiss.index_cpu_to_gpu(res, 0, opq_index, co)
                
                if args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
                    cur_loss =  (tr_loss - logging_loss)/args.logging_steps
                    tb_writer.add_scalar('train/all_loss', cur_loss, global_step)
                    logging_loss = tr_loss

                    cur_mrr =  (tr_mrr - logging_mrr)/(
                        args.logging_steps * args.gradient_accumulation_steps)
                    tb_writer.add_scalar('train/mrr_10', cur_mrr, global_step)
                    logging_mrr = tr_mrr
        
        save_model(model, args.model_save_dir, f'epoch-{epoch_idx+1}', args)
        tokenizer.save_pretrained(os.path.join(args.model_save_dir, f'epoch-{epoch_idx+1}'))
        model.config.save_pretrained(os.path.join(args.model_save_dir, f'epoch-{epoch_idx+1}'))
        faiss.write_index(opq_index,
            os.path.join(args.model_save_dir, f'epoch-{epoch_idx+1}', 
                os.path.basename(args.init_index_path)))


def run_parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--preprocess_dir", type=str, required=True)
    parser.add_argument("--model_save_dir", type=str, required=True)
    parser.add_argument("--log_dir", type=str, required=True)
    parser.add_argument("--init_index_path", type=str, required=True)
    parser.add_argument("--init_model_path", type=str, required=True)
    parser.add_argument("--init_backbone", type=str, required=True, choices=["roberta", "distilbert", "bert"])
    
    parser.add_argument("--lambda_cut", type=int, required=True)
    parser.add_argument("--centroid_lr", type=float, required=True)
    parser.add_argument("--gpu_search", action="store_true")
    parser.add_argument("--centroid_weight_decay", type=float, default=0)

    parser.add_argument("--loss_neg_topK", type=int, default=200)
    parser.add_argument("--max_seq_length", type=int, default=64)
    parser.add_argument("--train_batch_size", type=int, default=32)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--warmup_steps", default=2000, type=int)
    parser.add_argument("--no_cuda", action='store_true')
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument("--logging_steps", type=int, default=100)

    parser.add_argument("--lr", default=5e-6, type=float)
    parser.add_argument("--weight_decay", default=0.01, type=float)
    parser.add_argument("--adam_epsilon", default=1e-8, type=float)
    parser.add_argument("--max_grad_norm", default=1.0, type=float)
    parser.add_argument("--num_train_epochs", default=1, type=int)

    parser.add_argument("--threads", type=int, default=1)
    args = parser.parse_args()
    faiss.omp_set_num_threads(args.threads)

    os.makedirs(args.model_save_dir, exist_ok=True)
    return args


def train_genq(args):
    
    # Setup CUDA, GPU 
    args.model_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    args.n_gpu = 1

    logger.info("Model Device: %s, n_gpu: %s", args.model_device, args.n_gpu)

    # Set seed
    set_seed(args)

    if args.init_backbone == "distilbert":
        logger.info("loading Model for GenQ training: {}".format(args.init_model_path)) 
        config = AutoConfig.from_pretrained(args.init_model_path)
        config.return_dict = False
        config.gradient_checkpointing = args.gpu_search # to save cuda memory
        model = DistilBertDot.from_pretrained(args.init_model_path, config=config)
        tokenizer = AutoTokenizer.from_pretrained(args.init_model_path)
    
    elif args.init_backbone == "bert":
        logger.info("loading Model for GenQ training: {}".format(args.init_model_path))
        config = AutoConfig.from_pretrained(args.init_model_path)
        config.return_dict = False
        config.gradient_checkpointing = args.gpu_search # to save cuda memory
        model = BertDot.from_pretrained(args.init_model_path, config=config)
        tokenizer = AutoTokenizer.from_pretrained(args.init_model_path)
    
    elif args.init_backbone == "roberta":
        logger.info("loading Model for GenQ training: {}".format(args.init_model_path))
        config = RobertaConfig.from_pretrained(args.init_model_path)
        config.return_dict = False
        config.gradient_checkpointing = args.gpu_search # to save cuda memory
        model = RobertaDot.from_pretrained(args.init_model_path, config=config)
        tokenizer = RobertaTokenizer.from_pretrained("roberta-base", do_lower_case=True)
    
    model.to(args.model_device)

    opq_index = faiss.read_index(args.init_index_path)
    vt = faiss.downcast_VectorTransform(opq_index.chain.at(0))            
    assert isinstance(vt, faiss.LinearTransform)
    opq_transform = faiss.vector_to_array(vt.A).reshape(vt.d_out, vt.d_in)
    opq_transform = torch.FloatTensor(opq_transform).to(args.model_device)

    ivf_index = faiss.downcast_index(opq_index.index)
    invlists = faiss.extract_index_ivf(ivf_index).invlists
    ls = invlists.list_size(0)
    pq_codes = faiss.rev_swig_ptr(invlists.get_codes(0), ls * invlists.code_size)
    pq_codes = pq_codes.reshape(-1, invlists.code_size)
    pq_codes = torch.LongTensor(pq_codes).to(args.model_device)

    centroid_embeds = faiss.vector_to_array(ivf_index.pq.centroids)
    centroid_embeds = centroid_embeds.reshape(ivf_index.pq.M, ivf_index.pq.ksub, ivf_index.pq.dsub)
    coarse_quantizer = faiss.downcast_index(ivf_index.quantizer)
    coarse_embeds = faiss.vector_to_array(coarse_quantizer.xb)
    centroid_embeds += coarse_embeds.reshape(ivf_index.pq.M, -1, ivf_index.pq.dsub)
    faiss.copy_array_to_vector(centroid_embeds.ravel(), ivf_index.pq.centroids)
    coarse_embeds[:] = 0   
    faiss.copy_array_to_vector(coarse_embeds.ravel(), coarse_quantizer.xb)

    centroid_embeds = torch.FloatTensor(centroid_embeds).to(args.model_device)
    centroid_embeds.requires_grad = True
    
    train(args, model, pq_codes, centroid_embeds, opq_transform, opq_index, tokenizer)
    

if __name__ == "__main__":
    args = run_parse_args()
    logger.info(args)
    train_genq(args)
