#!/usr/bin/python3

import re
import torch
import pickle

from PyRNN.Data import Data
from PyRNN.RNNTagger import RNNTagger

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Tagger:
    ''' Tagger interface for the software "Versanalyse" '''

    def __init__(self, parfilename):
        ''' Einlesen des Taggermodelles '''
        self.data = Data(parfilename+'.io')
        with open(parfilename+'.hyper', 'rb') as file:
            hyper_params = pickle.load(file)
        self.model = RNNTagger(*hyper_params)
        self.model.load_state_dict(torch.load(parfilename+'.rnn', 
                                              map_location=torch.device('cpu'),
                                              weights_only=False))
        self.model = self.model.to(device)
        self.model.eval()

    def __call__(self, all_tokens):
        ''' Annotation of a sentence with POS tags '''
        all_tags = []
        with torch.no_grad():
            for i in range(0, len(all_tokens), 1000):
                tokens = all_tokens[i:i+1000]
                # map words to numbers and create Torch variables
                fwd_charIDs, bwd_charIDs = self.data.words2charIDvec(tokens)
                fwd_charIDs = torch.LongTensor(fwd_charIDs).to(device)
                bwd_charIDs = torch.LongTensor(bwd_charIDs).to(device)
                
                tagscores = self.model(fwd_charIDs, bwd_charIDs)
                tagIDs = tagscores.argmax(-1).tolist()
                tags = self.data.IDs2tags(tagIDs)
                
                # drop morphosyntactic features
                tags = [tag.split('.')[0] for tag in tags]
                
                # Replace the tag VAFIN with VVFIN unless an infinite verb
                # occurs immediately before or soon after it
                for i, tag in enumerate(tags):
                    # if tag == 'VAFIN':
                    #     start = max(i-1, 0)
                    #     end   = min(i+6, len(tags))
                    #     for k in range(start, end):
                    #         if re.match(r'V[VAM][^F]', tags[k]):
                    #             break
                    #     else:  # no other verb has been found
                    #         tag = 'VVFIN'
                    all_tags.append(tag)
            
            return all_tags
