import collections
import itertools


def convert_call_variants_outputs_to_probs_dict(
    canonical_variant, call_variants_outputs, alt_alleles_to_remove):
  """Converts a list of CallVariantsOutput to an internal allele probs dict.

  Args:
    canonical_variant: variants_pb2.Variant.
    call_variants_outputs: list of CallVariantsOutput.
    alt_alleles_to_remove: set of strings. Alleles to remove.

  Returns:
    Dictionary of {(allele1, allele2): list of probabilities},
    where allele1 and allele2 are strings.
  """
  flattened_dict = collections.defaultdict(lambda: [])
  if not call_variants_outputs:
    return flattened_dict

  for call_variants_output in call_variants_outputs:
    allele_set1 = frozenset([canonical_variant.reference_bases])
    allele_set2 = frozenset(
        canonical_variant.alternate_bases[index]
        for index in call_variants_output.alt_allele_indices.indices)
    if alt_alleles_to_remove.intersection(allele_set2):
      continue
    p11, p12, p22 = call_variants_output.genotype_probabilities
    for (set1, set2, p) in [(allele_set1, allele_set1, p11), (allele_set1,
                                                              allele_set2, p12),
                            (allele_set2, allele_set2, p22)]:
      for indices in itertools.product(set1, set2):
        flattened_dict[indices].append(p)
  return flattened_dict