import unittest
import sys
import math

from qiskit import QuantumCircuit, Aer, transpile, assemble, execute
from qiskit.visualization import plot_histogram, plot_bloch_multivector
from qiskit.circuit.library import Barrier, HGate, XGate
from numpy.random import randint
import numpy as np
np.seterr(divide='ignore', invalid='ignore')

import hypothesis.strategies as st
from hypothesis import given, settings

def generate_binary(len):
    return randint(2, size=len)

def encode_message(bits, bases, messageLen):
    message = []
    for i in range(messageLen):
        qc = QuantumCircuit(1,1)
        if bases[i] == 0: # Prepare qubit in Z-basis
            if bits[i] == 0:
                pass 
            else:
                qc.x(0)
        else: # Prepare qubit in X-basis
            if bits[i] == 0:
                qc.h(0)
            else:
                qc.x(0)
                qc.h(0)
        qc.barrier()
        message.append(qc)
    return message

def measure_message(message, bases, messageLen):
    measurements = []
    for q in range(messageLen):
        if bases[q] == 0: # measuring in Z-basis
            message[q].measure(0,0)
        if bases[q] == 1: # measuring in X-basis
            message[q].h(0)
            message[q].measure(0,0)
        aer_sim = Aer.get_backend('aer_simulator')
        qobj = assemble(message[q], shots=1, memory=True)
        result = aer_sim.run(qobj).result()
        measured_bit = int(result.get_memory()[0])
        measurements.append(measured_bit)
    return measurements

def remove_garbage(a_bases, b_bases, bits, messageLen):
    good_bits = []
    for q in range(messageLen):
        if a_bases[q] == b_bases[q]:
            # If both used the same basis, add
            # this to the list of 'good' bits
            good_bits.append(bits[q])
    return good_bits

def sample_bits(bits, selection):
    sample = []
    for i in selection:
        # use np.mod to make sure the
        # bit we sample is always in 
        # the list range
        i = np.mod(i, len(bits))
        # pop(i) removes the element of the
        # list at index 'i'
        sample.append(bits.pop(i))
    return sample

np.random.seed(seed=10)
messageLen = 100

## Step 1
# Alice generates bits
alice_bits = generate_binary(messageLen)

## Step 2
# Create an array to tell us which qubits
# are encoded in which bases
alice_bases = generate_binary(messageLen)
message = encode_message(alice_bits, alice_bases, messageLen)

## Step 3
# Decide which basis to measure in:
bob_bases = generate_binary(messageLen)
bob_results = measure_message(message, bob_bases, messageLen)

## Step 4
alice_key = remove_garbage(alice_bases, bob_bases, alice_bits, messageLen)
bob_key = remove_garbage(alice_bases, bob_bases, bob_results, messageLen)

## Step 5
sample_size = 15
bit_selection = generate_binary(sample_size)

print(alice_key)
print(bob_key)

bob_sample = sample_bits(bob_key, bit_selection)
print("  bob_sample = " + str(bob_sample))
alice_sample = sample_bits(alice_key, bit_selection)
print("alice_sample = "+ str(alice_sample))

bob_sample == alice_sample

n = 100

# Step 1
alice_bits = generate_binary(messageLen)
alice_bases = generate_binary(messageLen)

# Step 2
message = encode_message(alice_bits, alice_bases, messageLen)

# Interception!
eve_bases = generate_binary(messageLen)
intercepted_message = measure_message(message, eve_bases, messageLen)

# Step 3
bob_bases = generate_binary(messageLen)
bob_results = measure_message(message, bob_bases, messageLen)

# Step 4
bob_key = remove_garbage(alice_bases, bob_bases, bob_results, messageLen)
alice_key = remove_garbage(alice_bases, bob_bases, alice_bits, messageLen)

# Step 5
sample_size = 1 # Change this to something lower and see if 
                 # Eve can intercept the message without Alice
                 # and Bob finding out
        
bit_selection = generate_binary(sample_size)
bob_sample = sample_bits(bob_key, bit_selection)
alice_sample = sample_bits(alice_key, bit_selection)

if bob_sample != alice_sample:
    print("Eve's interference was detected.")
else:
    print("Eve went undetected!")

###########################################################################
## Define composite strategies to generate lists of ints in equal length ##
###########################################################################
@st.composite
def single_list(draw):
    arrayLengths = draw(st.integers(min_value=1, max_value=100))
    fixed_length_list = st.lists(st.integers(min_value=0, max_value=1), min_size=arrayLengths, max_size=arrayLengths)
    return (draw(fixed_length_list))

@st.composite
def pair_of_lists(draw):
    arrayLengths = draw(st.integers(min_value=1, max_value=100))
    fixed_length_list = st.lists(st.integers(min_value=0, max_value=1), min_size=arrayLengths, max_size=arrayLengths)
    return (draw(fixed_length_list), draw(fixed_length_list))

@st.composite
def trio_of_lists(draw):
    arrayLengths = draw(st.integers(min_value=1, max_value=100))
    fixed_length_list = st.lists(st.integers(min_value=0, max_value=1), min_size=arrayLengths, max_size=arrayLengths)
    return (draw(fixed_length_list), draw(fixed_length_list), draw(fixed_length_list))

@st.composite
def long_trio_of_lists(draw):
    arrayLengths = draw(st.integers(min_value=100, max_value=110))
    fixed_length_list = st.lists(st.integers(min_value=0, max_value=1), min_size=arrayLengths, max_size=arrayLengths)
    return (draw(fixed_length_list), draw(fixed_length_list), draw(fixed_length_list))

##########################
## test generate binary ##
##########################
@given(testLength = st.integers(min_value=0, max_value=10000))
def test_created_message_is_binary(testLength):
    binArr = generate_binary(testLength)
    for i in binArr:
        assert (i == 1 or i == 0) 

@given(testLength = st.integers(min_value=1, max_value=10000))
def test_created_message_equal_length_to_int_passed_in(testLength):
    binArr = generate_binary(testLength)
    assert(len(binArr) == testLength) 

############################
## encoding message tests ##
############################
@given(pair_of_lists())
@settings(deadline=None)
def test_encode_message_equal_length_to_base(lists):
    alice_bits, alice_bases = lists
    circuitArr = encode_message(alice_bits, alice_bases, len(alice_bits))
    assert(len(circuitArr) ==  len(alice_bits))

############################
## decoding message tests ##
############################
@given(lists = trio_of_lists())
@settings(deadline=None)
def test_decode_message_length_equals_base_length(lists):
    alice_bits, alice_bases, bob_base = lists
    encoded_message = encode_message(alice_bits, alice_bases, len(bob_base))
    msmtArr = measure_message(encoded_message, bob_base, len(bob_base))
    assert len(msmtArr) == len(bob_base)

@given(lists = trio_of_lists())
@settings(deadline=None)
def test_decode_message_is_binary(lists):
    alice_bits, alice_bases, bob_base = lists
    encoded_message = encode_message(alice_bits, alice_bases, len(bob_base))
    msmtArr = measure_message(encoded_message, bob_base, len(bob_base))
    for i in msmtArr:
        assert (i == 1 or i == 0) 
        
@given(lists = pair_of_lists())
@settings(deadline=None)
def test_decode_with_same_base_returns_original_bits(lists):
    alice_bits, alice_bases = lists
    encoded_message = encode_message(alice_bits, alice_bases, len(alice_bits))
    decodeWithSameBases = measure_message(encoded_message, alice_bases, len(alice_bases))
    assert(np.array_equal(np.array(alice_bits), np.array(decodeWithSameBases)))        
        
@given(lists = pair_of_lists())
@settings(deadline=None)
def test_decode_with_same_bases_return_same_array(lists):
    alice_bits, alice_bases = lists
    encoded_message = encode_message(alice_bits, alice_bases, len(alice_bits))
    encoded_message2 = encode_message(alice_bits, alice_bases, len(alice_bits))
    decodeWithSameBases = measure_message(encoded_message, alice_bases, len(alice_bases))
    decodeWithSameBases2 = measure_message(encoded_message2, alice_bases, len(alice_bases))
    assert(np.array_equal(np.array(decodeWithSameBases), np.array(decodeWithSameBases2)))
    
    
@given(lists = long_trio_of_lists())
@settings(deadline=None)
def test_decoding_runs_likely_different(lists):
    alice_bits, alice_bases, bob_base = lists
    encoded_message = encode_message(alice_bits, alice_bases, len(bob_base))
    msmtArr = measure_message(encoded_message, alice_bases, len(alice_bases))
    msmtArrRun2 = measure_message(encoded_message, bob_bases, len(bob_base))
    assert(not np.array_equal(np.array(msmtArr), np.array(msmtArrRun2)))

##############################
## remove garbage/key tests ##
##############################
@given(lists = trio_of_lists())
@settings(deadline=None)
def test_key_smaller_or_equal_len_to_original_bits(lists):
    alice_bits, alice_bases, bob_base = lists
    assert len(remove_garbage(alice_bits, alice_bases, bob_base, len(bob_base))) <= len(bob_base)

@given(lists = trio_of_lists())
@settings(deadline=None)
def test_check_keys_equal(lists):
    alice_bits, alice_bases, bob_bases = lists
    message = encode_message(alice_bits, alice_bases, len(bob_bases))
    bob_results = measure_message(message, bob_bases, len(bob_bases))
    alice_key = remove_garbage(alice_bases, bob_bases, alice_bits, len(bob_bases))
    bob_key = remove_garbage(alice_bases, bob_bases, bob_results, len(bob_bases)) 
    assert(np.array_equal(np.array(alice_key), np.array(bob_key)))

@given(lists = trio_of_lists())
@settings(deadline=None)
def test_key_is_binary(lists):
    alice_bits, alice_bases, bob_bases = lists
    alice_key = remove_garbage(alice_bases, bob_bases, alice_bits, len(bob_bases))
    for i in alice_key:
        assert (i == 1 or i == 0) 

@given(pair_of_lists())
@settings(deadline=None, max_examples=10)
def test_encode_message_applies_correct_phase(lists):
    backend = Aer.get_backend('aer_simulator')
    alice_bits, alice_bases = lists
    circuitArr = encode_message(alice_bits, alice_bases, len(alice_bits))
    for idx, base in enumerate(alice_bases):
        if (base == 0):
            if(alice_bits[idx] == 0):
                print("no assert")
            else:
                print("no assert")
        else:
            if(alice_bits[idx] == 0):
                assertPhase(backend, circuitArr[idx], [0], [0], 100000, 0.0001)
            else:
                assertPhase(backend, circuitArr[idx], [0], [180], 100000, 0.0001)

if __name__ == "__main__":
    # test_created_message_is_binary()
    # test_created_message_equal_length_to_int_passed_in()
    # test_encode_message_equal_length_to_base()
    # test_decode_message_length_equals_base_length()
    # test_decode_message_is_binary()
    # test_decode_with_same_base_returns_original_bits()
    # test_decode_with_same_bases_return_same_array()
    #test_decoding_runs_likely_different()
    # test_key_smaller_or_equal_len_to_original_bits()
    # test_check_keys_equal()
    # test_key_is_binary()
    #test_encode_message_applies_correct_phase()