import torch

from lm_eval.api.model import LM
from lm_eval.api.registry import register_model

import sys
sys.path.insert(0, "runtime")

from litgpt import Config
from litgpt.model import GPT


@register_model("obsidian_multiscreen")
class ObsidianMultiscreenLM(LM):

    def __init__(
        self,
        checkpoint,
        device="cuda",
        dtype="bfloat16",
        backend="triton",
        **kwargs
    ):

        self.device = device

        if backend:
            import os
            os.environ["MULTISCREEN_BACKEND"] = backend

        config = Config.from_file(
            f"{checkpoint}/model_config.yaml"
        )

        self.model = GPT(config)

        state = torch.load(
            f"{checkpoint}/lit_model.pth",
            map_location="cpu"
        )

        self.model.load_state_dict(
            state["model"]
        )

        self.model.to(device)

        if dtype == "bfloat16":
            self.model.to(torch.bfloat16)

        self.model.eval()

        self.vocab_size = config.padded_vocab_size
        self.max_length = config.block_size


    @property
    def eot_token_id(self):
        return 0


    @property
    def max_length(self):
        return self._max_length


    @max_length.setter
    def max_length(self, x):
        self._max_length=x


    def tok_encode(self, string):
        return self.tokenizer.encode(string)


    def loglikelihood(self, requests):

        results=[]

        for request in requests:

            context, continuation = request.args

            text=context+continuation

            ids=torch.tensor(
                [self.tokenizer.encode(text)],
                device=self.device
            )

            with torch.no_grad():
                logits=self.model(ids)

            log_probs=torch.log_softmax(
                logits,
                dim=-1
            )

            cont_ids=self.tokenizer.encode(
                continuation
            )

            score=0

            for i,tok in enumerate(cont_ids):
                score += log_probs[
                    0,
                    -(len(cont_ids)-i+1),
                    tok
                ]

            results.append(
                (float(score), True)
            )

        return results