import proveit
# Automation is not needed when building an expression:
proveit.defaults.automation = False # This will speed things up.
proveit.defaults.inline_pngs = False # Makes files smaller.
%load_expr # Load the stored expression as 'stored_expr'
# import Expression classes needed to build the expression
from proveit import Conditional, ExprRange, Lambda, U, Variable, VertExprArray, s, t
from proveit.linear_algebra import MatrixMult, ScalarMult
from proveit.logic import And, Equals, InSet
from proveit.numbers import Add, Exp, Interval, IntervalCO, Mod, Mult, Real, Round, e, four, frac, greater, i, one, pi, two, zero
from proveit.physics.quantum import I, NumKet, Z, ket_plus, var_ket_u
from proveit.physics.quantum.QPE import QPE, phase, two_pow_t
from proveit.physics.quantum.circuits import Gate, Input, Measure, MultiQubitElem, Output, Qcircuit
from proveit.statistics import Prob
# build up the expression from sub-expressions
sub_expr1 = Variable("_a", latex_format = r"{_{-}a}")
sub_expr2 = Add(t, one)
sub_expr3 = Add(t, s)
sub_expr4 = Interval(sub_expr2, sub_expr3)
sub_expr5 = MultiQubitElem(element = Gate(operation = QPE(U, t), part = sub_expr1), targets = Interval(one, sub_expr3))
expr = Lambda(phase, Conditional(greater(Prob(Qcircuit(vert_expr_array = VertExprArray([ExprRange(sub_expr1, Input(state = ket_plus), one, t), ExprRange(sub_expr1, MultiQubitElem(element = Input(state = var_ket_u, part = sub_expr1), targets = sub_expr4), one, s)], [ExprRange(sub_expr1, sub_expr5, one, t), ExprRange(sub_expr1, sub_expr5, sub_expr2, sub_expr3)], [ExprRange(sub_expr1, Measure(basis = Z), one, t), ExprRange(sub_expr1, Gate(operation = I).with_implicit_representation(), one, s)], [ExprRange(sub_expr1, MultiQubitElem(element = Output(state = NumKet(Mod(Round(Mult(two_pow_t, phase)), two_pow_t), t), part = sub_expr1), targets = Interval(one, t)), one, t), ExprRange(sub_expr1, MultiQubitElem(element = Output(state = var_ket_u, part = sub_expr1), targets = sub_expr4), one, s)]))), frac(four, Exp(pi, two))), And(InSet(phase, Real), InSet(phase, IntervalCO(zero, one)), Equals(MatrixMult(U, var_ket_u), ScalarMult(Exp(e, Mult(two, pi, i, phase)), var_ket_u)))))
# check that the built expression is the same as the stored expression
assert expr == stored_expr
assert expr._style_id == stored_expr._style_id
print("Passed sanity check: expr matches stored_expr")
# Show the LaTeX representation of the expression for convenience if you need it.
print(stored_expr.latex())
stored_expr.style_options()
# display the expression information
stored_expr.expr_info()