import torch


def get_attention(model, tokenizer, text, include_queries_and_keys=False):

    """Compute representation of the attention to pass to the d3 visualization

    Args:
      model: pytorch_transformers model
      tokenizer: pytorch_transformers tokenizer
      text: Input text
      include_queries_and_keys: Indicates whether to include queries/keys in results

    Returns:
      Dictionary of attn representations with the structure:
      {
        'left_text': list of source tokens, to be displayed on the left of the vis
        'right_text': list of target tokens, to be displayed on the right of the vis
        'attn': list of attention matrices, one for each layer. Each has shape (num_heads, source_seq_len, target_seq_len)
        'queries' (optional): list of query vector arrays, one for each layer. Each has shape (num_heads, source_seq_len, vector_size)
        'keys' (optional): list of key vector arrays, one for each layer. Each has shape (num_heads, target_seq_len, vector_size)
      }
    """

    # Prepare inputs to model
    token_ids = tokenizer.encode(text)
    tokens = [tokenizer.decode([t]).strip() for t in token_ids]
    tokens_tensor = torch.tensor(token_ids).unsqueeze(0)

    # Call model to get attention data
    model.eval()
    _, _, attn_data_list = model(tokens_tensor)

    # Format attention data for visualization
    all_attns = []
    all_queries = []
    all_keys = []
    for layer, attn_data in enumerate(attn_data_list):
        attn = attn_data['attn'][0]  # assume batch_size=1; output shape = (num_heads, seq_len, seq_len)
        all_attns.append(attn.tolist())
        if include_queries_and_keys:
            queries = attn_data['queries'][0]  # assume batch_size=1; output shape = (num_heads, seq_len, vector_size)
            all_queries.append(queries.tolist())
            keys = attn_data['keys'][0]  # assume batch_size=1; output shape = (num_heads, seq_len, vector_size)
            all_keys.append(keys.tolist())
    results = {
        'attn': all_attns,
        'left_text': tokens,
        'right_text': tokens
    }
    if include_queries_and_keys:
        results.update({
            'queries': all_queries,
            'keys': all_keys,
        })
    return {'all': results}