"""Trainable action classifier + candidate pointer; non-autoregressive baseline."""
import torch
from torch import nn
from huggingface_hub import PyTorchModelHubMixin


class PointerPolicy(nn.Module, PyTorchModelHubMixin, library_name='baim', tags=['browser-agent','cpu']):
    def __init__(self, vocab_size=4096, width=64, encoder='mean', lexical_features=True):
        super().__init__()
        self.encoder_kind = encoder
        self.lexical_features = lexical_features
        self.embedding = nn.Embedding(vocab_size,width,padding_idx=0)
        if encoder == 'gru':
            self.encoder = nn.GRU(width,width,batch_first=True)
        elif encoder == 'transformer':
            self.position = nn.Embedding(24,width)
            self.encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(width,4,width*2,
                dropout=0.0,batch_first=True),num_layers=1,enable_nested_tensor=False)
        elif encoder != 'mean':
            raise ValueError('unknown encoder')
        self.action = nn.Sequential(nn.Linear(width,width),nn.ReLU(),nn.Linear(width,3))
        self.pointer = nn.Sequential(nn.Linear(width*4+5,width),nn.ReLU(),nn.Linear(width,1))

    def embed(self, ids):
        mask = ids.ne(0)
        x = self.embedding(ids)
        if self.encoder_kind == 'gru':
            x,_ = self.encoder(x)
        elif self.encoder_kind == 'transformer':
            x = x + self.position(torch.arange(ids.shape[-1],device=ids.device))
            # Empty padded candidates need one unmasked token to avoid NaNs.
            safe = mask.clone()
            safe[:,0] = True
            x = self.encoder(x,src_key_padding_mask=~safe)
        return (x*mask.unsqueeze(-1)).sum(1)/mask.sum(1,keepdim=True).clamp_min(1)

    def forward(self, goal, elements, features, mask):
        g = self.embed(goal)
        batch,count,length = elements.shape
        e = self.embed(elements.reshape(batch*count,length)).reshape(batch,count,-1)
        expanded = g.unsqueeze(1).expand_as(e)
        lexical = features if self.lexical_features else torch.zeros_like(features)
        pairs = torch.cat([expanded,e,expanded*e,torch.abs(expanded-e),lexical],dim=-1)
        pointers = self.pointer(pairs).squeeze(-1).masked_fill(~mask,-1e4)
        return self.action(g),pointers