import torch


def nearest_neighbor_grad(
    averaged_grad,
    embedding_matrix,
    trigger_token_ids,
    tree,
    step_size,
    increase_loss=False,
    num_candidates=1,
):
    """
    Takes a small step in the direction of the averaged_grad and finds the nearest
    vector in the embedding matrix using a kd-tree.
    """
    new_trigger_token_ids = [
        [None] * num_candidates for _ in range(len(trigger_token_ids))
    ]
    averaged_grad = averaged_grad.cpu()
    embedding_matrix = embedding_matrix.cpu()
    if increase_loss:  # reverse the sign
        step_size *= -1
    for token_pos, trigger_token_id in enumerate(trigger_token_ids):
        # take a step in the direction of the gradient
        trigger_token_embed = (
            torch.nn.functional.embedding(
                torch.LongTensor([trigger_token_id]), embedding_matrix
            )
            .detach()
            .cpu()
            .numpy()[0]
        )
        stepped_trigger_token_embed = (
            trigger_token_embed
            + averaged_grad[token_pos].detach().cpu().numpy() * step_size
        )
        # look in the k-d tree for the nearest embedding
        _, neighbors = tree.query([stepped_trigger_token_embed], k=num_candidates)
        for candidate_number, neighbor in enumerate(neighbors[0]):
            new_trigger_token_ids[token_pos][candidate_number] = neighbor
    return new_trigger_token_ids