import numpy as np
import pickle
from tensorflow.keras.models import load_model # type: ignore
from tensorflow.keras.preprocessing.sequence import pad_sequences # type: ignore
import re


class NERSRLRepository:
    def __init__(self):
        # Load model and artifacts
        self.model = load_model("app/lstm_model/ner_srl/lstm_ner_srl_model.keras")

        with open("app/lstm_model/ner_srl/word2idx.pkl", "rb") as f:
            self.word2idx = pickle.load(f)
        with open("app/lstm_model/ner_srl/tag2idx_ner.pkl", "rb") as f:
            self.tag2idx_ner = pickle.load(f)
        with open("app/lstm_model/ner_srl/tag2idx_srl.pkl", "rb") as f:
            self.tag2idx_srl = pickle.load(f)

        self.idx2tag_ner = {i: t for t, i in self.tag2idx_ner.items()}
        self.idx2tag_srl = {i: t for t, i in self.tag2idx_srl.items()}

        self.PAD_WORD_ID = self.word2idx["PAD"]
        self.MAXLEN = self.model.input_shape[1]

    def _preprocess_tokens(self, tokens: list[str]) -> np.ndarray:
        seq = [self.word2idx.get(tok.lower(), self.word2idx["UNK"]) for tok in tokens]
        return pad_sequences(
            [seq], maxlen=self.MAXLEN, padding="post", value=self.PAD_WORD_ID
        )

    def predict_sentence(self, sentence: str) -> dict:
        tokens = re.findall(r"\d{1,2}\.\d{2}|\w+|[^\w\s]", sentence.lower())
        print(tokens)
        seq_padded = self._preprocess_tokens(tokens)

        pred_ner_prob, pred_srl_prob = self.model.predict(seq_padded, verbose=0)
        pred_ner = pred_ner_prob.argmax(-1)[0][: len(tokens)]
        pred_srl = pred_srl_prob.argmax(-1)[0][: len(tokens)]

        return {
            "tokens": tokens,
            "labels_ner": [self.idx2tag_ner[int(i)] for i in pred_ner],
            "labels_srl": [self.idx2tag_srl[int(i)] for i in pred_srl],
        }

    def labeling_token(self, tokens: list[str]) -> dict:
        sentence = " ".join(tokens)
        return self.predict_sentence(sentence)
