# -*- coding: utf-8 -*-
"""
Created on Fri Aug  5 11:44:45 2022

@author: Eva Brenner
"""

import cirq
import matplotlib.pyplot as plt
import numpy as np
import collections

# Choose a unitary matrix U 
phi = 0.546 # 0 <= phi < 1
unitary = np.array([
    [0.0, 1.0, 0.0,  0.0],
    [0.0,  0.0, 0.0,  1.0],
    [0.0,  0.0, np.exp(complex(0,2*np.pi*phi)),  0.0],
    [1.0,  0.0, 0.0, 0.0]
]) 

# Build the corresponding gate
U = cirq.MatrixGate(unitary)
controlled_U = cirq.ControlledGate(U)

# Get some qubits
t = 4  # accuracy of approximation
n = 2  # to represent eigenvector
approx = cirq.LineQubit.range(t)
eigenvector = cirq.LineQubit.range(t,t+n)


circ = cirq.Circuit()

"""
TODO: 
    1) Präparieren Sie den Eigenvektor der oben definierten Matrix U zum 
Eigenwert exp(2*pi*i*phi) im Register eigenvector.

    2) Implementieren Sie einen Schaltkreis für die Phasenschätzung unter Verwendung
der Klasse ControlledGate in cirq.
"""

# Run the circuit
simulator = cirq.Simulator()
result = simulator.run(circ)
print("\nMeasurement:")
print(result)

# Convert result into float
def approx_period(bits):
    # Calculates an approximation to the period given a measurement result
    return np.sum(2 ** np.arange(start=t-1,stop=-1,step=-1) * bits) / 2**t

approx_result = approx_period(result.measurements['m'])
print(approx_result)


# Repeat 100 times
samples = simulator.run(circ, repetitions=100)
approx_results = np.sum(2 ** np.arange(start=t-1,stop=-1,step=-1) * 
                        samples.measurements['m'],axis=1) / 2**t
counts = samples.histogram(key="m")
# Print Measurements 
print("\nMeasurements:")
print(counts)

# Count how many times the measurements occured 
counter = collections.Counter(approx_results)

# Print the counter and, optionally, the measurements
print("\nThe measurements correspond to the following approximations to the period:")
print(counter)
#print(approx_results)

# Plot the approximations to the period
custom_histogram = samples.histogram(key ='m', fold_func=approx_period)
cirq.plot_state_histogram(custom_histogram, plt.subplot(), 
                          xlabel="Approximation to the period", 
                          title="Result Period Histogram")
plt.show()