"""
TODO: DEPRECATE THIS ONE IN FAVOR OF lib_q_computer.py
"""

import random

import numpy as np

# |0> and |1>
_0 = np.array([[1],
               [0]])
_1 = np.array([[0],
               [1]])

# |+> and |->
_p = np.array([[1 / np.sqrt(2)],
               [1 / np.sqrt(2)]])
_m = np.array([[1 / np.sqrt(2)],
               [-1 / np.sqrt(2)]])

# Gates - Identity, Pauli X and Hadamard
I = np.array([[1, 0],
              [0, 1]])
X = np.array([[0, 1],
              [1, 0]])
H = np.array([[1 / np.sqrt(2), 1 / np.sqrt(2)],
              [1 / np.sqrt(2), -1 / np.sqrt(2)]])


def measure_probability(qbit):
    """
    In a qbit [a, b] normalized: |a|^2 + |b|^2 = 1
    Probability of 0 is |a|^2 and 1 with prob |b|^2

    :returns: tuple of probabilities to measure 0 or 1"""
    return np.abs(qbit[0][0]) ** 2, np.abs(qbit[1][0]) ** 2


def measure(qbit):
    """
    This gets a random choice of either 0 and 1 with weights
    based on the probabilities of the qbit

    :returns: classical bit based on qbit probabilities"""
    return random.choices([0, 1], measure_probability(qbit))[0]


def run_qbit_tests():
    # asserts are sets of tests to check if mathz workz

    # Identity: verify that I|0> == |0> and I|1> == |0>
    assert np.array_equal(I.dot(_0), _0)
    assert np.array_equal(I.dot(_1), _1)

    # Pauli X: verify that X|0> == |1> and X|1> == |0>
    assert np.array_equal(X.dot(_0), _1)
    assert np.array_equal(X.dot(_1), _0)

    # measure probabilities in sigma_x of |0> and |1>
    # using allclose since dealing with floats
    assert np.allclose(measure_probability(_0), (1.0, 0.0))
    assert np.allclose(measure_probability(_1), (0.0, 1.0))

    # applying Hadamard puts the qbit in orthogonal +/- basis
    assert np.array_equal(H.dot(_0), _p)
    assert np.array_equal(H.dot(_1), _m)

    # measure probabilities in sigma_x of |+> and |->
    # using allclose since dealing with floats
    assert np.allclose(measure_probability(_p), (0.5, 0.5))
    assert np.allclose(measure_probability(_m), (0.5, 0.5))