"""Small vocabulary and deterministic candidate retrieval, without site rules."""
from collections import Counter
import re
import torch

KINDS = ('C', 'T', 'O')
ROLES = {'button', 'link', 'textbox', 'searchbox', 'combobox', 'checkbox', 'radio', 'listbox', 'spinbutton'}


def words(text):
    return re.findall(r'\w+', text.casefold(), flags=re.UNICODE)


def fit_vocab(rows, limit=4096):
    counts = Counter()
    for row in rows:
        counts.update(words(row['goal']))
        for element in row['elements']:
            counts.update(words(element['role'] + ' ' + element['name']))
    return {'<pad>':0, '<unk>':1, **{token:i+2 for i, (token, _) in enumerate(counts.most_common(limit-2))}}


def lexical(goal, element):
    goal_words, name_words = set(words(goal)), set(words(element['name']))
    overlap = len(goal_words & name_words)
    return [overlap / max(1,len(name_words)), overlap / max(1,len(goal_words)),
            float(element['name'].casefold() in goal.casefold()),
            float(element.get('visible', True)), float(element.get('enabled', True))]


def candidates(goal, elements, limit=40):
    eligible = [i for i, e in enumerate(elements) if e.get('visible', True) and e.get('enabled', True)
                and not e.get('sensitive', False) and e['role'] in ROLES]
    return sorted(eligible, key=lambda i: (-lexical(goal,elements[i])[0], i))[:limit]


def encode(rows, vocab, max_tokens=24, max_candidates=40):
    # Training padding uses batch size; inference uses only the surviving candidates.
    maps = [candidates(row['goal'],row['elements'],max_candidates) for row in rows]
    count = max(1, max(map(len,maps),default=0))
    goal = torch.zeros((len(rows),max_tokens),dtype=torch.long)
    element = torch.zeros((len(rows),count,max_tokens),dtype=torch.long)
    features = torch.zeros((len(rows),count,5))
    mask = torch.zeros((len(rows),count),dtype=torch.bool)
    targets = torch.full((len(rows),),-100,dtype=torch.long)
    actions = torch.full((len(rows),),-100,dtype=torch.long)
    def tokens(text):
        return [vocab.get(token,1) for token in words(text)[:max_tokens]]
    for batch,row in enumerate(rows):
        ids = tokens(row['goal'])
        goal[batch,:len(ids)] = torch.tensor(ids,dtype=torch.long)
        if row.get('action') in KINDS:
            actions[batch] = KINDS.index(row['action'])
        for j,index in enumerate(maps[batch]):
            e = row['elements'][index]
            ids = tokens(e['role'] + ' ' + e['name'])
            element[batch,j,:len(ids)] = torch.tensor(ids,dtype=torch.long)
            features[batch,j] = torch.tensor(lexical(row['goal'],e))
            mask[batch,j] = True
            if row.get('target') == index:
                targets[batch] = j
    return (goal,element,features,mask),actions,targets,maps