import json
from pathlib import Path

from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
import torch

def _get_model_config(cfg_path: Path):
    with open(cfg_path, 'r') as mconfig:
        cfg = json.load(mconfig)
        return cfg

class SentenceClassPredictor:
    def __init__(self, model_path: str) -> None:
        self.model_path = Path(model_path)
        self.config_path = Path.joinpath(self.model_path , 'config.json')
        self.model_config = _get_model_config(self.config_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_path)
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_config['_name_or_path'])
        self.id2label = self.model_config['id2label']
    
    def predict(self, sentence: str) -> str:
        tokenized = self.tokenizer(sentence, truncation=True, padding=True, return_tensors='pt')
        input_ids = tokenized['input_ids']
        result = torch.argmax(self.model(input_ids).logits, 1)
        return self.id2label[str(result.item())]


class TokenClassPredictor:
    def __init__(self, model_path: str) -> None:
        self.model_path = Path(model_path)
        self.config_path = Path.joinpath(self.model_path , 'config.json')
        self.model_config = _get_model_config(self.config_path)
        self.id2tag = self.model_config['id2tag']
        self.model = AutoModelForTokenClassification.from_pretrained(self.model_config['base_model'],
                                                                     num_labels=len(self.model_config['unique_tags']))
        self.model.load_state_dict(torch.load(Path.joinpath(self.model_path, 'pytorch_model.bin'), map_location=torch.device('cpu')))
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_config['base_model'])

    def predict(self, sentence: str) -> any:
        tokenized = self.tokenizer(sentence, 
                                   is_split_into_words=False,
                                   return_special_tokens_mask=True,
                                   return_tensors='pt')
        input_ids = tokenized['input_ids']
        outputs = self.model(input_ids).logits

        predicted_labels = [l for (i, l) in enumerate(torch.argmax(outputs, dim=2).tolist()[0])
            if tokenized["special_tokens_mask"].tolist()[0][i] == 0]
        merged_tokens = [l for (i, l) in 
            enumerate(self.tokenizer.convert_ids_to_tokens(input_ids.tolist()[0]))
            if tokenized["special_tokens_mask"].tolist()[0][i] == 0]

        return merged_tokens, [self.id2tag[str(i)] for i in predicted_labels]