# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, transpile, Aer, IBMQ, execute
from qiskit.tools.jupyter import *
from qiskit.visualization import *

from qiskit.providers.aer import QasmSimulator
from numpy import pi
from qiskit.circuit import Parameter, ParameterVector


import numpy as np
from numpy import pi, sin, cos, exp
import cmath

degree = 1
scaling = 1
coeffs = [0.15 + 0.15j]*degree
coeff0 = 0.1

#Generate a Fourier series.
def function(x):
    res = coeff0
    for i,c in enumerate(coeffs):
        exponent = complex(0, scaling*(i+1)*x)
        conj_c = np.conjugate(c)
        res += c * np.exp(exponent) + conj_c * np.exp(-exponent)
    return np.real(res)

#Original data points
x = np.linspace(-6,6,70)
y_0 = np.array([function(x_) for x_ in x])

from qiskit.circuit.library import TwoLocal
from qiskit.opflow import CircuitStateFn, CircuitOp
import matplotlib.pyplot as plt

n = 2
qc = QuantumCircuit(n)
num_parameters = 8

#Data encoding circuit
def M(cir, x):
    for w in range(n):
        cir.rx(scaling*x, w)
    return cir

#Trainable circuit block (Defining the ansatz)
def U(theta):
    encode = TwoLocal(num_qubits = 2, reps = 1, rotation_blocks=['ry','rz'], 
               entanglement_blocks='cx', entanglement='full', insert_barriers=True)
    return encode.assign_parameters(theta)
    
#Defining the Quantum Model. Here we are mapping the data to the fourier series.
def Quantum_model(cir, weights, x=None):
    weights = weights.reshape(2,num_parameters)
    cir = cir.compose(U(weights[0]))
    cir = cir.compose(M(cir, x))        
    cir = cir.compose(U(weights[1]))
    cir = CircuitStateFn(cir)
    weights = weights.reshape(1,2*num_parameters)
    
    circuit = QuantumCircuit(2)
    circuit.z(0)
    op = CircuitOp(circuit)
    
    return (cir.adjoint().compose(op).compose(cir).eval().real)*scaling

scaling = 1
weight = 2*np.pi*np.random.random(size = (1,2*num_parameters))

random_y = [Quantum_model(qc, weight, x = x_) for x_ in x]

plt.plot(x,random_y, c='blue')
plt.plot(x,y_0, c = 'red')

plt.show()

def square_loss(targets, predictions):
    s = 0
    for t, p in zip(targets, predictions):
        s = s + (t - p) ** 2
    loss = s / 2*len(targets)
    return loss

def cost(weights):
    y = y_0
    predictions = [Quantum_model(qc,weights, x=x_) for x_ in x]
    return square_loss(y, predictions)

import scipy.optimize as op
weight2 = op.fmin(lambda w: cost(w), weight, xtol=0.0001, ftol=0.0001, maxiter=100,
                  maxfun=None, full_output=0, disp=1, retall=0, callback=None)

new_weights = np.array(weight2)
trained_y = [Quantum_model(qc, new_weights, x = x_) for x_ in x]

plt.plot(x, trained_y, c='blue')
plt.plot(x,y_0, c = 'red')

plt.show()

#Computes the first 2*K+1 Fourier coefficients of a 2*pi periodic function.

def fourier_coefficients(f, K):
    n_coeffs = 2*K+1
    t = np.linspace(0, 2*np.pi, n_coeffs, endpoint=False)
    y = np.fft.rfft(f(t)) / t.size
    return y

n_coeffs = len(new_weights)
n_samples = 100

coeffs = []
def f(x):
    return np.array([Quantum_model(qc, new_weights, x=i) for i in x])

coeffs_sample = fourier_coefficients(f, n_coeffs)
coeffs.append(coeffs_sample)

coeffs = np.array(coeffs)
coeffs_real = np.real(coeffs)
coeffs_imag = np.imag(coeffs)

n_coeffs = len(coeffs_real[0])

fig, ax = plt.subplots(1, n_coeffs, figsize=(15,4))

for idx, ax_ in enumerate(ax):
    ax_.set_title(r"$c_{}$".format(idx))
    ax_.scatter(coeffs_real[:, idx], coeffs_imag[:, idx], s=35, facecolor='white', edgecolor='red')
    ax_.set_aspect("equal")
    ax_.set_ylim(-1, 1)
    ax_.set_xlim(-1, 1)


plt.tight_layout(pad=0.5)
plt.show()