import os
import ujson
import torch
import random

from collections import defaultdict, OrderedDict

from transformers import BertConfig

from colbert.parameters import DEVICE
from colbert.modeling.colbert import ColBERT
from colbert.modeling.prefix import PrefixColBERT
from colbert.utils.utils import print_message, load_checkpoint


def load_model(args, do_print=True):
    if args.prefix:
        config = BertConfig.from_pretrained('bert-base-uncased', cache_dir=".cache")
        config.pre_seq_len = args.pre_seq_len
        config.prefix_hidden_size = args.prefix_hidden_size
        config.prefix_mlp = args.prefix_mlp
        colbert = PrefixColBERT.from_pretrained('bert-base-uncased', config=config,
                                          query_maxlen=args.query_maxlen,
                                          doc_maxlen=args.doc_maxlen,
                                          dim=args.dim,
                                          similarity_metric=args.similarity,
                                          mask_punctuation=args.mask_punctuation)
    else:
        config = BertConfig.from_pretrained('bert-base-uncased', cache_dir=".cache")
        colbert = ColBERT.from_pretrained('bert-base-uncased',
                                        config=config,
                                        query_maxlen=args.query_maxlen,
                                        doc_maxlen=args.doc_maxlen,
                                        dim=args.dim,
                                        similarity_metric=args.similarity,
                                        mask_punctuation=args.mask_punctuation)
    colbert = colbert.to(DEVICE)

    print_message("#> Loading model checkpoint.", condition=do_print)

    checkpoint = load_checkpoint(args.checkpoint, colbert, do_print=do_print)

    colbert.eval()

    return colbert, checkpoint
