# Copyright 2020 The Magenta Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Beam search library."""

import collections
import copy
import heapq

# A beam entry containing a) the current sequence, b) a "state" containing any
# information needed to extend the sequence, and c) a score for the current
# sequence e.g. log-likelihood.
BeamEntry = collections.namedtuple('BeamEntry', ['sequence', 'state', 'score'])


def _generate_branches(beam_entries, generate_step_fn, branch_factor,
                       num_steps):
  """Performs a single iteration of branch generation for beam search.

  This method generates `branch_factor` branches for each sequence in the beam,
  where each branch extends the event sequence by `num_steps` steps (via calls
  to `generate_step_fn`). The resulting beam is returned.

  Args:
    beam_entries: A list of BeamEntry tuples, the current beam.
    generate_step_fn: A function that takes three parameters: a list of
        sequences, a list of states, and a list of scores, all of the same size.
        The function should generate a single step for each of the sequences and
        return the extended sequences, updated states, and updated (total)
        scores, as three lists. The function may modify the sequences, states,
        and scores in place, but should also return the modified values.
    branch_factor: The integer branch factor to use.
    num_steps: The integer number of steps to take per branch.

  Returns:
    The updated beam, with `branch_factor` times as many BeamEntry tuples.
  """
  if branch_factor > 1:
    branched_entries = beam_entries * branch_factor
    all_sequences = [copy.deepcopy(entry.sequence)
                     for entry in branched_entries]
    all_states = [copy.deepcopy(entry.state) for entry in branched_entries]
    all_scores = [entry.score for entry in branched_entries]
  else:
    # No need to make copies if there's no branching.
    all_sequences = [entry.sequence for entry in beam_entries]
    all_states = [entry.state for entry in beam_entries]
    all_scores = [entry.score for entry in beam_entries]

  for _ in range(num_steps):
    all_sequences, all_states, all_scores = generate_step_fn(
        all_sequences, all_states, all_scores)

  return [BeamEntry(sequence, state, score)
          for sequence, state, score
          in zip(all_sequences, all_states, all_scores)]


def _prune_branches(beam_entries, k):
  """Prune all but the `k` sequences with highest score from the beam."""
  indices = heapq.nlargest(k, range(len(beam_entries)),
                           key=lambda i: beam_entries[i].score)
  return [beam_entries[i] for i in indices]


def beam_search(initial_sequence, initial_state, generate_step_fn, num_steps,
                beam_size, branch_factor, steps_per_iteration):
  """Generates a sequence using beam search.

  Initially, the beam is filled with `beam_size` copies of the initial sequence.

  Each iteration, the beam is pruned to contain only the `beam_size` event
  sequences with highest score. Then `branch_factor` new event sequences are
  generated for each sequence in the beam. These new sequences are formed by
  extending each sequence in the beam by `steps_per_iteration` steps. So between
  a branching and a pruning phase, there will be `beam_size` * `branch_factor`
  active event sequences.

  After the final iteration, the single sequence in the beam with highest
  likelihood will be returned.

  The `generate_step_fn` function operates on lists of sequences + states +
  scores rather than single sequences. This is to allow for the possibility of
  batching.

  Args:
    initial_sequence: The initial sequence, a Python list-like object.
    initial_state: The state corresponding to the initial sequence, with any
        auxiliary information needed for extending the sequence.
    generate_step_fn: A function that takes three parameters: a list of
        sequences, a list of states, and a list of scores, all of the same size.
        The function should generate a single step for each of the sequences and
        return the extended sequences, updated states, and updated (total)
        scores, as three lists.
    num_steps: The integer length in steps of the final sequence, after
        generation.
    beam_size: The integer beam size to use.
    branch_factor: The integer branch factor to use.
    steps_per_iteration: The integer number of steps to take per iteration.

  Returns:
    A tuple containing a) the highest-scoring sequence as computed by the beam
    search, b) the state corresponding to this sequence, and c) the score of
    this sequence.
  """
  sequences = [copy.deepcopy(initial_sequence) for _ in range(beam_size)]
  states = [copy.deepcopy(initial_state) for _ in range(beam_size)]
  scores = [0] * beam_size

  beam_entries = [BeamEntry(sequence, state, score)
                  for sequence, state, score
                  in zip(sequences, states, scores)]

  # Choose the number of steps for the first iteration such that subsequent
  # iterations can all take the same number of steps.
  first_iteration_num_steps = (num_steps - 1) % steps_per_iteration + 1

  beam_entries = _generate_branches(
      beam_entries, generate_step_fn, branch_factor, first_iteration_num_steps)

  num_iterations = (num_steps -
                    first_iteration_num_steps) // steps_per_iteration

  for _ in range(num_iterations):
    beam_entries = _prune_branches(beam_entries, k=beam_size)
    beam_entries = _generate_branches(
        beam_entries, generate_step_fn, branch_factor, steps_per_iteration)

  # Prune to the single best beam entry.
  beam_entry = _prune_branches(beam_entries, k=1)[0]

  return beam_entry.sequence, beam_entry.state, beam_entry.score
