import math
import numpy as np

class Register:

    def __init__(self, num_qubits, state):
        """Initialise quantum register

        state - column vector representing initial quantum state. Example:
            [
                (0.0, 0.0i) # 0 0
                (0.0, 0.0i) # 0 1
                (1.0, 0.0i) # 1 0
                (0.0, 0.0i) # 1 1
            ]
        shows a "collapsed quantum state" because 1 0 is the only state set to 1

        A register can be in a measured state superposition of states e.g.
        if a Hadamard gate is applied to a 1-qubit system then the qubit
        would be in a superposition of being both 0 and 1

        The state of the register must always follow this equation:
        Let a, b and c be complex numbers representing the quantum states of a
        2-qubit system: |a|^2 + |b|^2 + |c|^2 + |d|^2 = 1
        """

        num_bases = len(state)
        if(num_bases != pow(2, num_qubits)):
            raise Exception("Invalid number of bases vectors")

        eps = 0.0001
        total = complex(0, 0)
        for s in state:
            total = total + pow(np.abs(s), 2)
        if(not (abs(total - (complex(1, 0))) < eps)):
            raise Exception("Quantum state is invalid")

        self.state = state
        self.num_qubits = num_qubits

    # TODO
    # TODO make me static
    # TODO make me private but still somehow testable
    def generate_bases(self, num_qubits):
        """ Generate bases vectors

        Example: [000, 001, 010, 011, 100, 101, 110, 111] for a 3-qubit system
        """
        bases = []
        num_vectors = pow(2, num_qubits)
        base_idx = 0
        for idx in xrange(num_vectors):
            base = []
            current_base = idx
            for idy in xrange(num_qubits):
                base.append(current_base % 2)
                current_base = current_base / 2
            base.reverse()
            bases.append(base)
        return bases

    def measure(self):
        """Perform quantum measurement

        Collapses from quantum state to classical state
        qubit_num -- the qubits to measure
        """
        # TODO use a {}?
        # TODO to do this, loop through each bases and do (1/value)^2
        current_state = np.squeeze(np.asarray(self.state))
        probabilities = {}
        for idx, basis in enumerate(current_state):
            probabilities[idx] = pow(basis, 2)

@@ -1959,16 +1959,589 @@
  bases%0A%0A
+    def filter_bases(self, num_qubits, qubit_nums):%0A        %22%22%22Filter bases%0A%0A        Example:%0A            - Generate bases: %5B000, 001, 010, 011, 100, 101, 110, 111%5D%0A            - Filter to qubits 0 and 2: %5B00, 01, 00, 01, 10, 11, 10, 11%5D%0A        %22%22%22%0A        bases = self.generate_bases(num_qubits)%0A        for base_num, base in enumerate(bases):%0A            bases_before_filter = bases%5Bbase_num%5D%0A            bases%5Bbase_num%5D = %5B%5D%0A            for qubit_position in qubit_nums:%0A                bases%5Bbase_num%5D.append(bases_before_filter%5Bqubit_position%5D)%0A        return bases%0A%0A
     def 

add filter bases