from BasicTools import *

def measure(dm, operator, **kwargs):
    """Measures a density matrix dm with an operator, the operator is decomposed
    into projectors, then the probabilities are calculated for each projector.

    A random draw takes place and the final density matrix is returned with
    the energy measured."""
    ops, e_vals = proj_meas_decomd(operator)

    if not np.allclose(np.sum([a.H * a for a in ops], 0), I(len(ops[0]))) and 'force' not in kwargs:
        raise Exception("Operators do not sum to the identity matrix.")

    probabilities = []
    for op in ops:
        probabilities.append(trace(op.H * op * vec_to_dm(dm)))

    draw = np.random.rand()
    bins = np.cumsum(probabilities)

    measured = np.digitize(draw, bins)  # identify which measurement has happened

    final_dm = ops[measured] * vec_to_dm(dm) * ops[measured].H
    final_dm /= probabilities[measured]

    return final_dm, e_vals[measured]
