"""
    Code taken from ColBERT (new_api branch), which is distributed under Apache-compatible MIT license.
    We changed how the model is initialized to make it compatible with FlexNeuART saving
    and loading of models. There are two options to create the model:
    1. by specifying the COLBERT checkpoint generated by original COLBERT code
    2. by specifying the config
    In both cases, one needs to provide the type of the Transformer/BERT model (bert_flavor)

    https://github.com/stanford-futuredata/ColBERT/tree/new_api/colbert
"""
import torch

from flexneuart import if_none

from .hf_colbert import HF_ColBERT
from .config import ColBERTConfig


class BaseColBERT(torch.nn.Module):
    """
    Shallow module that wraps the ColBERT parameters, custom configuration, and underlying tokenizer.
    Unlike the original code that provided direct instantiation and saving of the model/colbert_config/tokenizer package,
    this code delegates initialization to upstream code.

    Like HF, evaluation mode is the default.
    """

    def __init__(self,
                 bert_flavor : str,
                 colbert_config: ColBERTConfig=None,
                 colbert_checkpoint: str=None):
        super().__init__()
        self.raw_tokenizer = HF_ColBERT.raw_tokenizer_from_pretrained(bert_flavor)

        if colbert_checkpoint is None:
            assert colbert_config is not None, "config must be specified when checkpoint isn't"
            self.colbert_config = colbert_config
            self.model: torch.nn.Module = HF_ColBERT.from_pretrained(bert_flavor, colbert_config=self.colbert_config)
        else:
            assert colbert_config is None, "config should NOT be specified when checkpoint is"

            self.colbert_config = ColBERTConfig.from_existing(ColBERTConfig.load_from_checkpoint(colbert_checkpoint))
            self.model = HF_ColBERT.from_pretrained(colbert_checkpoint, colbert_config=self.colbert_config)

        self.eval()

    @property
    def device(self):
        return self.model.device

    @property
    def bert(self):
        return self.model.bert

    @property
    def linear(self):
        return self.model.linear
    
    @property
    def score_scaler(self):
        return self.model.score_scaler

    def save(self, path):
        assert not path.endswith('.dnn'), f"{path}: We reserve *.dnn names for the deprecated checkpoint format."

        self.model.save_pretrained(path)
        self.raw_tokenizer.save_pretrained(path)

        self.colbert_config.save_for_checkpoint(path)

