import collections
import itertools


def merge_mutation_sets(
    mutations_a,
    mutations_b):
  """Constructs all possible merges of mutations sets `mutations_a` and `mutations_b`.

  Merging two mutation sets involves combining all the mutations together into
  a combined mutation set. A valid mutation set has mutations to unique
  positions. For a given position P, if there are mutations in the both of the
  two sets, then a valid merge can include either the mutation to position P
  from set A or B. Thus, for M positions with overlap, this function returns
  2^{M} merged mutation sets.

  >>> m1 = [(0, 2), (1, 3)]
  >>> m2 = [(0, 4), (3, 5)]
  >>> merge_mutation_sets(m1, m2)
  [((0, 2), (1, 3), (3, 5)),
   ((0, 4), (1, 3), (3, 5))]

  Args:
    mutations_a: An iterable of tuples encoding mutations (position, mutation).
    mutations_b: See `mutations_a`.

  Returns:
    All possible merges of mutation sets `mutations_a` and `mutations_b`.
  """
  assert all(len(m) == 2 for m in mutations_a)
  assert all(len(m) == 2 for m in mutations_b)
  position_idx = 0

  grouped_mutations = collections.defaultdict(list)
  for m in list(mutations_a) + list(mutations_b):
    position = m[position_idx]
    grouped_mutations[position].append(m)
  singletons = []
  collisions = []
  for position, muts in grouped_mutations.items():
    if len(muts) == 1:
      singletons.append(muts[0])
    else:
      collisions.append(muts)

  to_return = []
  for c in itertools.product(*collisions):
    mutations = c + tuple(singletons)
    sorted_mutations = tuple(sorted(mutations, key=lambda m: m[position_idx]))
    to_return.append(sorted_mutations)
  to_return = sorted(to_return)
  return to_return