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

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)
from sentence_transformers import CrossEncoder

import os, sys
import torch
import random
import time
import faiss
import logging
import argparse
import numpy as np
import json
import csv
import datetime
from tqdm import tqdm, trange
from collections import defaultdict

csv.field_size_limit(sys.maxsize)
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)
        self.negdict = 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]
        ret_val['neg_ids'] = self.negdict[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]
        all_neg_pids = [x["neg_ids"] for x in batch]
        return data, qids, all_rel_pids, all_neg_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, qids, all_rel_pids, all_neg_ids, lambda_cut, corpus, queries, cross_encoder):
    loss = 0
    mrr = 0
    train_batch_size = len(batch_neighbors)
    loss_function = torch.nn.MSELoss(reduction="mean")
    for qid, qembedding, retrieve_pids, cur_rel_pids, cur_neg_ids in zip(
        qids, query_embeddings, batch_neighbors, all_rel_pids, all_neg_ids):

        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)])

        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]
        pair_diff = rel_scores - irrel_scores
        
        query, docs = queries[qid], [corpus[doc_id] for doc_id in retrieve_pids.tolist()]
        scores_ce = cross_encoder.predict(
            [(query, x) for x in docs], 
            show_progress_bar=False, 
            convert_to_numpy=False, 
            convert_to_tensor=True,
            batch_size=128,
        )
        rel_ce_scores = scores_ce[target_labels]
        irrel_ce_scores = scores_ce[~target_labels]
        labels = rel_ce_scores - irrel_ce_scores
        cur_loss = loss_function(pair_diff, labels)
        loss += cur_loss

    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)
    corpus, queries = {}, {}
    
    def fix_nulls(s):
        for line in s:
            yield line.replace('\0', '')

    reader = csv.reader(fix_nulls(open(os.path.join(args.data_path, "collection.tsv"), encoding="utf-8")), delimiter="\t", quoting=csv.QUOTE_MINIMAL)
    for id, row in enumerate(reader):
        corpus[int(row[0])] = row[1]
        
    
    reader = csv.reader(fix_nulls(open(os.path.join(args.data_path, "queries.tsv"), encoding="utf-8")), delimiter="\t", quoting=csv.QUOTE_MINIMAL)
    for id, row in enumerate(reader):
        queries[int(row[0])] = row[1]

    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=args.model_device, max_length=350)

    
    for epoch_idx, _ in enumerate(train_iterator):
        epoch_iterator = tqdm(train_dataloader, desc="Iteration")
        for step, (batch, qids, all_rel_poffsets, all_neg_ids) 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, qids, all_rel_poffsets, all_neg_ids, args.lambda_cut, corpus, queries, cross_encoder)   
            # 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
            
            if step == 45000:
                save_model(model, args.model_save_dir, f'step-{step}', args)
                tokenizer.save_pretrained(os.path.join(args.model_save_dir, f'step-{step}'))
                model.config.save_pretrained(os.path.join(args.model_save_dir, f'step-{step}'))
                faiss.write_index(opq_index,
                    os.path.join(args.model_save_dir, f'step-{step}', 
                        os.path.basename(args.init_index_path)))
        
        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=6, type=int)

    parser.add_argument("--data_path", default=None, type=str, required=True, help="path to original JPQ collection and query path")
    parser.add_argument("--cross_encoder", type=str, default="cross-encoder/ms-marco-MiniLM-L-6-v2")
    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_gpl(args):

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

    logger.warning("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 GPL 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 GPL 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_gpl(args)
