import unittest

import cirq
from cirq.ops import H, X, I
import random
import matplotlib.pyplot as plt
import numpy as np
from numpy.random import randint

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):
        qubits = cirq.LineQubit.range(1)
        qc = cirq.Circuit()
        if bases[i] == 0: # Prepare qubit in Z-basis
            if bits[i] == 0:
                qc.append(cirq.I(qubits[0]))
            else:
                qc.append(cirq.X(qubits[0]))
        else: # Prepare qubit in X-basis
            if bits[i] == 0:
                qc.append(cirq.H(qubits[0]))
            else:
                qc.append(cirq.X(qubits[0]))
                qc.append(cirq.H(qubits[0]))
        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
            if (not message[q].has_measurements()):
                for qubit in message[q].all_qubits():
                    message[q].append(cirq.measure(qubit))
        if bases[q] == 1: # measuring in X-basis
            if (not message[q].has_measurements()):
                for qubit in message[q].all_qubits():
                    message[q].append(cirq.H(qubit))
                    message[q].append(cirq.measure(qubit))
        simulator = cirq.Simulator()
        measured_bit = simulator.run(message[q])
        measurements.append((measured_bit.data.iat[0,0]))   
    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=0)
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)

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 = 15 # 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))

@given(pair_of_lists())
@settings(deadline=None)
def test_encode_message_are_circuits(lists):
    alice_bits, alice_bases = lists
    circuitArr = encode_message(alice_bits, alice_bases, len(alice_bits))
    for i in circuitArr:
        assert(isinstance(i, cirq.Circuit))

@given(pair_of_lists())
@settings(deadline=None)
def test_encode_message_circuits_are_not_longer_than_2(lists):
    alice_bits, alice_bases = lists
    circuitArr = encode_message(alice_bits, alice_bases, len(alice_bits))
    for i in circuitArr:
        assert(not(sum(1 for _ in i.all_operations()) > 2))

@given(pair_of_lists())
@settings(deadline=None)
def test_encode_message_circuits_use_only_H_X_I(lists):
    alice_bits, alice_bases = lists
    circuitArr = encode_message(alice_bits, alice_bases, len(alice_bits))
    for i in circuitArr:
        for gate in i.all_operations():
            assert(isinstance(gate.gate, cirq.ops.identity.IdentityGate)
                   or isinstance(gate.gate, cirq.ops.pauli_gates._PauliX)
                   or isinstance(gate.gate, cirq.ops.common_gates.HPowGate)) 

############################
## 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) 

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_encode_message_are_circuits()
    #test_encode_message_circuits_are_not_longer_than_2()
    #test_encode_message_circuits_use_only_H_X_I()
    #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()

