from qiskit.circuit.library.standard_gates import RXGate, RZGate, CXGate, CZGate
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister

def anzats_circ1(thetas, D2, in_state):
    qr = QuantumRegister(4, name="q")
    qc = QuantumCircuit(qr)
    qc.initialize(in_state)
        
    for d in range(D2):
        qc.append(RXGate(thetas[0]), [qr[0]])
        qc.append(RXGate(thetas[1]), [qr[1]])
        qc.append(RXGate(thetas[2]), [qr[2]])
        qc.append(RXGate(thetas[3]), [qr[3]])
        
        qc.append(RZGate(thetas[4]), [qr[0]])
        qc.append(RZGate(thetas[5]), [qr[1]])
        qc.append(RZGate(thetas[6]), [qr[2]])
        qc.append(RZGate(thetas[7]), [qr[3]])
        
        qc.append(CZGate(), [qr[0], qr[1]])
        qc.append(CZGate(), [qr[1], qr[2]])
        qc.append(CZGate(), [qr[2], qr[3]])
        qc.barrier(qr)
    
    qc.append(RXGate(thetas[0]), [qr[0]])
    qc.append(RXGate(thetas[1]), [qr[1]])
    qc.append(RXGate(thetas[2]), [qr[2]])
    qc.append(RXGate(thetas[3]), [qr[3]])

    qc.append(RZGate(thetas[4]), [qr[0]])
    qc.append(RZGate(thetas[5]), [qr[1]])
    qc.append(RZGate(thetas[6]), [qr[2]])
    qc.append(RZGate(thetas[7]), [qr[3]])
    
    qc.measure_all()
    
    return qc

def anzats_circ1_uninitialized(thetas, D2):
    qr = QuantumRegister(4, name="q")
    qc = QuantumCircuit(qr)
        
    for d in range(D2):
        qc.append(RXGate(thetas[0]), [qr[0]])
        qc.append(RXGate(thetas[1]), [qr[1]])
        qc.append(RXGate(thetas[2]), [qr[2]])
        qc.append(RXGate(thetas[3]), [qr[3]])
        
        qc.append(RZGate(thetas[4]), [qr[0]])
        qc.append(RZGate(thetas[5]), [qr[1]])
        qc.append(RZGate(thetas[6]), [qr[2]])
        qc.append(RZGate(thetas[7]), [qr[3]])
        
        qc.append(CZGate(), [qr[0], qr[1]])
        qc.append(CZGate(), [qr[1], qr[2]])
        qc.append(CZGate(), [qr[2], qr[3]])
        qc.barrier(qr)
    
    qc.append(RXGate(thetas[0]), [qr[0]])
    qc.append(RXGate(thetas[1]), [qr[1]])
    qc.append(RXGate(thetas[2]), [qr[2]])
    qc.append(RXGate(thetas[3]), [qr[3]])

    qc.append(RZGate(thetas[4]), [qr[0]])
    qc.append(RZGate(thetas[5]), [qr[1]])
    qc.append(RZGate(thetas[6]), [qr[2]])
    qc.append(RZGate(thetas[7]), [qr[3]])
    
    qc.measure_all()
    
    return qc

import numpy as np
def get_k_basis(k, n):
    full_basis = np.identity(n)
    return full_basis[:k]


from qiskit.opflow import X, Z, I, H, Y

# for convinience let all coefficients ai and Jij be 0.5
H_transverse_ising = 0.5*((I^I^I^X) + (I^I^X^I) + (I^X^I^I) + (X^I^I^I) + \
                          (Z^Z^I^I) + (Z^I^Z^I) + (Z^I^I^Z) + (I^Z^Z^I) + \
                          (I^Z^I^Z) + (I^I^Z^Z))

# TODO: change this to a 4 qubits Hamiltonian
H2_molecule_Hamiltonian = -0.5053051899926562*(I^I) + \
                            -0.3277380754984016*(Z^I) + \
                            0.15567463610622564*(Z^Z) + \
                            -0.3277380754984016*(I^Z)


import matplotlib
from qiskit import assemble, Aer
from qiskit import *
from qiskit.visualization import plot_histogram
import math

sim = Aer.get_backend('aer_simulator')

def convert_string_to_index(string):
    return int(string, 2)

# TODO: this isn't the right way to calculate the result state, should be changed
def convert_result_to_state(result):
    total_meas = 0
    
    for k in result.keys():
        total_meas += result[k]
        amount_of_qubits = len(k)

    state = np.zeros(2**amount_of_qubits)
    
    for k in result.keys():
        index = convert_string_to_index(k)
        state[index] = result[k]/total_meas
        
    return state

# TODO: finish these functions
def calc_ising_avg(qc):
    avg = 0
    for i in range(4):
        qc_tmp = qc.copy()
        qc_tmp.append(H, [i])

def calc_molecular_avg(qc):
    pass

def calc_target_func1(thetas, basis, D2, Ham):
    target_func = 0
    for j in basis:
        qc = anzats_circ1(thetas, D2, j)
        print(qc.draw())
        qobj = assemble(qc)
        result = sim.run(qobj).result().get_counts()
        state = convert_result_to_state(result)
        dagger_state = np.matrix(state)
        state = dagger_state.getH()
        state = np.array(state)
        dagger_state = np.array(dagger_state)
        
#         TODO: calculate this with estimation with the Ham pauli strings
#         if Ham == "Ising Model":
#             product = calc_ising_avg(qc)
#         else             product = calc_molecular_avg(qc)
        product = np.matmul(Ham, state)
        product = np.matmul(dagger_state, product)
                
        target_func += product
        
    return target_func



# JUST TO RUN AND CHECK
k = 4
basis = get_k_basis(k,16)
D2 = 1
Ham = H_transverse_ising.to_matrix()
thetas = np.zeros(8)
calc_target_func1(thetas, basis, D2, Ham)


def objective_func1(thetas):
    k = 4
    basis = get_k_basis(k,16)
    D2 = 1
    Ham = H_transverse_ising.to_matrix()
#     Ham = H_transverse_ising
    target_func = calc_target_func1(thetas, basis,D2, Ham)
    return target_func[0][0]


from qiskit.algorithms.optimizers import L_BFGS_B

bfgs_optimizer = L_BFGS_B(maxiter=60)

point, value, nfev = bfgs_optimizer.optimize(8,objective_func1,initial_point=np.zeros(8))
print(point)
print("---point---")
print(value)
print("---value---")

thetas_opt = point


def anzats_circ2(phis, D1, in_state):
    qr = QuantumRegister(4, name="q")
    cr = ClassicalRegister(4)
    qc = QuantumCircuit(qr, cr)
    qc.initialize(in_state)
    
    for d in range(D1):
        qc.append(RXGate(phis[0]), [qr[2]])
        qc.append(RXGate(phis[1]), [qr[3]])
        
        qc.append(RZGate(phis[2]), [qr[2]])
        qc.append(RZGate(phis[3]), [qr[3]])
        
        qc.append(CZGate(), [qr[2], qr[3]])
        qc.barrier(qr)
    return qc

def calc_target_func2(thetas_opt, phis, in_state, D1, D2, Ham):
    target_func = 0

    qc2 = anzats_circ2(phis, D1, in_state)
    qc1 = anzats_circ1_uninitialized(thetas_opt, D2)
    qc = qc2.compose(qc1)
    print(qc.draw())
    qobj = assemble(qc)
    result = sim.run(qobj).result().get_counts()
    state = convert_result_to_state(result)
    dagger_state = np.matrix(state)
    state = dagger_state.getH()
    state = np.array(state)
    dagger_state = np.array(dagger_state)
        
#         TODO: calculate this with estimation with the Ham pauli strings
#         if Ham == "Ising Model":
#             product = calc_ising_avg(qc)
#         else             product = calc_molecular_avg(qc)
    product = np.matmul(Ham, state)
    product = np.matmul(dagger_state, product)

    target_func += product
        
    return target_func

i = np.random.randint(0,k)
def objective_func2(phis):
    in_state = basis[i]
    print("in state")
    print(in_state)
    D1 = 1
    D2 = 1
    Ham = H_transverse_ising.to_matrix()
#     Ham = H_transverse_ising
    target_func2 = calc_target_func2(thetas_opt, phis, in_state, D1, D2, Ham)
    print("target func:")
    print(target_func2[0][0])
    return target_func2[0][0]

def objective_func2_neg(phis):
    return -1*objective_func2(phis)

#to see wraping function is working well
phis = np.array([np.pi, np.pi, np.pi, np.pi])
exp =objective_func2(phis)
print("exp:")
print(exp)
exp_neg =objective_func2_neg(phis)
print("exp_neg:")
print(exp_neg)

point, value, nfev = bfgs_optimizer.optimize(4,objective_func2_neg,initial_point=np.array([2, 2, 2, 2]))
print(point)
print("---point---")
print(value)
print("---value---")

value = -value
print(value)



