import json

import torch

from model.request_models import InputElement, InputNREType
from model.response_models import ResponseNREModel

from typing import List, Dict

from pathlib import Path

from src.OpenNRE import opennre
from src.helpers import DownloadModel


class ResponsesNRE:
    responses: Dict[int, ResponseNREModel] = {}

    @staticmethod
    def check_id(request_id: int) -> bool:
        if ResponsesNRE.responses.get(request_id):
            return True
        return False

    @staticmethod
    def get_response(request_id: int) -> ResponseNREModel:
        if ResponsesNRE.responses.get(request_id):
            return ResponsesNRE.responses[request_id]


class NREResponse:
    def __init__(self, name="./data/model.ckpt.pth.tar", BERT='HooshvareLab/bert-base-parsbert-uncased',
                 device: str = "gpu"):

        if name == "./data/mbertEntail-2GBEn.bin":
            model_path = Path(name)
            if not model_path.is_file():
                DownloadModel().download_nre_model(name)
        sentence_encoder = opennre.encoder.BERTEntityEncoder(
            max_length=128,
            pretrain_path=BERT
        )
        rel2id = json.load(open('./data/Perlex_rel2id.json'))
        self.model = opennre.model.SoftmaxNN(sentence_encoder, len(rel2id), rel2id)
        checkpoint = torch.load(name, map_location=torch.device('cpu'))
        self.model.load_state_dict(checkpoint['state_dict'])

    def predict_json(self, input_json: InputNREType, request_id: int) -> ResponseNREModel:
        ResponsesNRE.responses[request_id] = ResponseNREModel(progression="in going")
        result = []
        for element in input_json:
            element.h.pos = eval(element.h.pos)
            element.t.pos = eval(element.t.pos)
            result.append(self.model.infer(element.dict()))

        ResponsesNRE.responses[request_id] = ResponseNREModel(progression="done", result=result)
        print(request_id, " done ")
        return ResponsesNRE.responses[request_id]
