diff --git a/lib_q_computer_math.py b/lib_q_computer_math.py
index 689d9e1..ee8e28c 100644
--- a/lib_q_computer_math.py
+++ b/lib_q_computer_math.py
@@ -428,8 +428,8 @@ class HermitianOperator(LinearOperator, HermitianMatrix):
         return str(self.m)
 
 
-# TODO - How to add a CNOT gate to the Quantum Processor?
-#        Imagine if I have to act on a 3-qubit computer and CNOT(q1, q3)
+# How to add a CNOT gate to the Quantum Processor?
+# Imagine if I have to act on a 3-qubit computer and CNOT(q1, q3)
 #
 # Decomposed CNOT :
 # reverse engineered from
@@ -454,8 +454,8 @@ class TwoQubitPartial(object):
         return str("-{}-".format(self.rpr))
 
 
-C_ = TwoQubitPartial("C")
-x_ = TwoQubitPartial("x")
+C_partial = TwoQubitPartial("C")
+x_partial = TwoQubitPartial("x")
 
 
 class TwoQubitOperator(UnitaryOperator):
@@ -468,13 +468,16 @@ class TwoQubitOperator(UnitaryOperator):
         self.A_p = A_p
         self.B_p = B_p
 
+        A.operator = self
+        B.operator = self
+
     def verify_step(self, step):
         if not (step.count(self.A) == 1 and step.count(self.B) == 1):
-            raise RuntimeError("Both CONTROL and TARGET need to be defined in the same step exactly once")
+            raise RuntimeError("Both {} and {} need to be defined in the same step exactly once".format(
+                self.A, self.B
+            ))
 
     def compose(self, step, state):
-        # TODO: Hacky way to do CNOT
-        #       Should generalize for a 2-Qubit gate
         # _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
         outer_0, outer_1 = [], []
         for s in step:
@@ -555,11 +558,12 @@ Z = UnitaryOperator([[1, 0],
                     name="Z")
 
 
-# TODO: These are rotations that are specified commonly e.g. in
+# These are rotations that are specified commonly e.g. in
 #       https://www.quantum-inspire.com/kbase/rotation-operators/
 #       http://www.vcpc.univie.ac.at/~ian/hotlist/qc/talks/bloch-sphere-rotations.pdf
 #       and elsewhere DO NOT equate X, Y and Z for theta=np.pi
-#       ??????????????????????????????????????????????????????
+# However - they are correct up to a global phase which is all that matters for measurement purposes
+#
 
 def Rx(theta):
     return UnitaryOperator([[np.cos(theta / 2), -1j * np.sin(theta / 2)],
@@ -578,11 +582,6 @@ def Rz(theta):
                             [0, np.power(np.e, 1j * theta / 2)]],
                            name="Rz")
 
-
-# def unitary(alpha, vector:Vector, theta):
-#     return np.exp(1j * alpha)
-
-
 H = UnitaryOperator([[1 / np.sqrt(2), 1 / np.sqrt(2)],
                      [1 / np.sqrt(2), -1 / np.sqrt(2)], ],
                     name="H")
@@ -591,16 +590,7 @@ CNOT = TwoQubitOperator([[1, 0, 0, 0],
                          [0, 1, 0, 0],
                          [0, 0, 0, 1],
                          [0, 0, 1, 0], ],
-                        TwoQubitPartial("C"),
-                        TwoQubitPartial("x"),
-                        I,
-                        X)
-
-C, x = CNOT.A, CNOT.B
-
-
-# TODO: End Hacky way to define 2-qbit gate
-###########################################################
+                         C_partial, x_partial, I, X)
 
 
 def assert_raises(exception, msg, callable, *args, **kwargs):
@@ -664,7 +654,7 @@ class MeasurementOpeartor(HermitianOperator):
 
     def get_prob(self, state: State):
         state_ct = state.conjugate_transpose()
-        return np.asscalar(state_ct.m.dot(self.m.dot(state.m)))
+        return state_ct.m.dot(self.m.dot(state.m)).item()
 
 
 def test_measurement_ops():
@@ -939,24 +929,12 @@ class QuantumCircuit(object):
             self.add_row(row_data)
 
     def compose_quantum_state(self, step):
-        if C in step or x in step:
-            if not (step.count(C) == 1 and step.count(x) == 1):
-                raise RuntimeError("Both CONTROL and TARGET need to be defined in the same step exactly once")
-            # TODO: Hacky way to do CNOT
-            #       Should generalize for a 2-Qubit gate
-            # _0.x(_0) * Matrix(I.m) + _1.x(_1) * Matrix(X.m)
-            outer_0, outer_1 = [], []
-            for s in step:
-                if s == C:
-                    outer_0.append(_0.x(_0))
-                    outer_1.append(_1.x(_1))
-                elif s == x:
-                    outer_0.append(Matrix(I.m))
-                    outer_1.append(Matrix(X.m))
-                else:
-                    outer_0.append(Matrix(s.m))
-                    outer_1.append(Matrix(s.m))
-            return reduce((lambda x, y: x * y), outer_0) + reduce((lambda x, y: x * y), outer_1)
+        partials = [op for op in step if isinstance(op, TwoQubitPartial)]
+        # TODO: No more than 1 TwoQubitGate **OR** UnitaryOperator can be used in a step
+        for partial in partials:
+            two_qubit_op = partial.operator
+            two_qubit_op.verify_step()
+            return two_qubit_op.compose(step)
         return reduce((lambda x, y: x * y), step)
 
     def step(self):
@@ -1077,9 +1055,9 @@ class QuantumProcessor(object):
 def test_quantum_processor():
     # Produce Bell state between 0 and 2 qubit
     qc = QuantumCircuit(3)
-    qc.add_row([H, C])
+    qc.add_row([H, C_partial])
     qc.add_row([_, _])
-    qc.add_row([_, x])
+    qc.add_row([_, x_partial])
     qc.print()
     qp = QuantumProcessor(qc)
     qp.get_sample(100)