# This file includes code adapted from:
# https://github.com/Liyan06/MiniCheck/blob/b58b9fa69acbd1015ec970fa65dd752413a053d2/minicheck/inference.py
#
# The original work is licensed under the Apache License, Version 2.0.
# See the LICENSE file in the original repository for details.


from abc import ABC, abstractmethod
from itertools import batched
from typing import Dict, Generator, override

import torch
from nltk.tokenize import word_tokenize
from pydantic import BaseModel
from transformers import (
    AutoConfig,
    AutoModelForSeq2SeqLM,
    AutoModelForSequenceClassification,
    AutoTokenizer,
)

from srs2std.util import sent_tokenize_with_newlines

DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DEFAULT_BATCH_SIZE = 4


def generate_chunks(text: str, chunk_size: int) -> Generator[str, None, None]:
    sentences = sent_tokenize_with_newlines(text) or [""]
    current_chunks = []
    current_word_count = 0
    for sentence in sentences:
        sentence_word_count = len(word_tokenize(sentence))
        if current_word_count + sentence_word_count > chunk_size:
            chunk = " ".join(current_chunks).replace(" \n ", "\n").strip()
            if chunk:
                yield chunk
            current_chunks = [sentence]
            current_word_count = sentence_word_count
        else:
            current_chunks.append(sentence)
            current_word_count += sentence_word_count
    if current_chunks:
        chunk = " ".join(current_chunks).replace(" \n ", "\n").strip()
        if chunk:
            yield chunk


class Reference(BaseModel):
    chunk: str
    score: float


class Scorer(ABC):
    @property
    @abstractmethod
    def model_path(self) -> str:
        pass

    @property
    @abstractmethod
    def max_model_len(self) -> int:
        pass

    @property
    def device(self) -> str:
        return DEFAULT_DEVICE

    @property
    def batch_size(self) -> int:
        return DEFAULT_BATCH_SIZE

    def __init__(self):
        self.model_name = self.model_path.split("/")[-1]
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_path, use_fast=True, revision="main"
        )

        # reserve some space (hard coded, 200) for the claim to be checked
        self.default_chunk_size = self.max_model_len - 200
        assert self.default_chunk_size > 0

    def score(self, doc: str, claim: str) -> Reference:
        doc_chunk_gen = generate_chunks(doc, self.default_chunk_size)
        max_supported_prob = 0.0
        max_prob_chunk = ""
        for doc_batch in batched(doc_chunk_gen, self.batch_size):
            doc_batch = list(doc_batch)
            batch_input = self._batch_tokenize(doc_batch, claim)
            batch_input = {k: v.to(self.device) for k, v in batch_input.items()}
            with torch.no_grad():
                probs = self._infer_batch(batch_input)
                current_max_prob = probs.max().item()
            if current_max_prob > max_supported_prob:
                max_supported_prob = current_max_prob
                max_prob_index = int(probs.argmax().item())
                max_prob_chunk = doc_batch[max_prob_index]
        return Reference(chunk=max_prob_chunk, score=max_supported_prob)

    @abstractmethod
    def _infer_batch(self, batch_input: Dict[str, torch.Tensor]) -> torch.Tensor:
        pass

    @abstractmethod
    def _concat_doc_and_claim(self, doc, claim):
        pass

    def _batch_tokenize(self, doc_batch, claim):
        return self.tokenizer(
            [self._concat_doc_and_claim(doc, claim) for doc in doc_batch],
            max_length=self.max_model_len,
            truncation=True,
            padding=True,
            return_tensors="pt",
        )


class BERT(Scorer):
    @property
    @abstractmethod
    def label2id(self) -> Dict:
        pass

    def __init__(self):
        super().__init__()
        config = AutoConfig.from_pretrained(
            self.model_path,
            label2id=self.label2id,
            finetuning_task="text-classification",
            revision="main",
        )
        config.num_labels = len(set(self.label2id.values()))
        config.problem_type = "single_label_classification"
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_path,
            config=config,
            revision="main",
            ignore_mismatched_sizes=False,
        )
        self.supported_label = self.label2id.get("supported", 0)
        self.model.to(self.device)
        self.model.eval()

    @override
    def _infer_batch(self, batch_input):
        logits = self.model(**batch_input).logits.cpu()
        label_probs = torch.nn.functional.softmax(logits, dim=1)
        return label_probs[:, self.supported_label]

    @override
    def _concat_doc_and_claim(self, doc, claim):
        eos_token = self.tokenizer.eos_token
        return eos_token.join([doc, claim])


class T5(Scorer):
    @property
    @abstractmethod
    def pos_token_id(self) -> int:
        pass

    @property
    @abstractmethod
    def neg_token_id(self) -> int:
        pass

    def __init__(self):
        super().__init__()
        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_path)
        self.model.to(self.device)
        self.model.eval()

    @override
    def _infer_batch(self, batch_input):
        decoder_input_ids = torch.zeros(
            (batch_input["input_ids"].size(0), 1), dtype=torch.long
        ).to(self.device)
        outputs = self.model(
            input_ids=batch_input["input_ids"],
            attention_mask=batch_input["attention_mask"],
            decoder_input_ids=decoder_input_ids,
        )
        logits = outputs.logits.squeeze(1)
        label_logits = logits[
            :, torch.tensor([self.neg_token_id, self.pos_token_id])
        ].cpu()
        label_probs = torch.nn.functional.softmax(label_logits, dim=-1)
        return label_probs[:, 1]
