from typing import Dict
import math
import itertools


def match_fidelity(target_fidelity: float, cuts: Dict[frozenset, int]):
    """Determines the "width" of each cut (what values to evaluate over) to
    match the target fidelity as closely as possible.

  Args:
    target_fidelity: the requested fidelity. The final fidelity must be
      greater than this, but should otherwise be as close as possible.
    cuts: a dict of (cut_qubits, bond_dim) for all cuts.

  Returns:
    The final fidelity approximation and a list of tuples
      (cut qubits, cut width, cut_dimension).
  """
    if (target_fidelity == 1.0 or not cuts):
        return (1.0,
                zip(list(cuts.keys()), list(cuts.values()),
                    list(cuts.values())))

    value_sets = []
    denominator = 1.0
    for cut, dim in cuts.items():
        # Using min() prevents float math from producing dim+1.
        min_width = min(dim, math.ceil(dim * target_fidelity))
        value_sets.append(list(range(min_width, dim + 1)))
        denominator *= dim

    combinations = itertools.product(*value_sets)
    closest_set = []
    closest_fidelity = 1.0
    for width_set in combinations:
        fidelity = 1.0
        for width in width_set:
            fidelity *= width
        fidelity /= denominator
        if fidelity >= target_fidelity and fidelity < closest_fidelity:
            closest_set = width_set
            closest_fidelity = fidelity
            if fidelity == target_fidelity:
                break

    return (closest_fidelity,
            zip(list(cuts.keys()), closest_set, list(cuts.values())))