import nbimporter
from typing import Dict, Tuple, List
import numpy as np
from tqdm import tqdm

QUBITS_NUM = 4 
N = 16
K = 4
NUM_SHOTS = 1024
NUM_ITERATIONS = 50

w = 0.5

approximated_energies = []

from qiskit import Aer
from qiskit.utils import QuantumInstance, algorithm_globals

seed = 50
algorithm_globals.random_seed = seed

simulator_backend = Aer.get_backend('qasm_simulator')

from scipy.optimize import minimize

from utiles import *

input_states = get_first_k_eigenvectors_from_n_computational_basis(K, N)

from ansatz_circuit_item2 import get_full_variational_quantum_circuit

init_circuit_params = {
    "thetas": np.random.uniform(low=0, high=2*np.pi, size=8),
    "phis": np.random.uniform(low=0, high=2*np.pi, size=4),
    "D1": 2,
    "D2": 6
}

def prepare_circuit_params(thetas) -> Dict:
     return {
    "thetas": thetas[4:],
    "phis": thetas[:4],
    "D1": 2,
    "D2": 6
     }

def get_ansatz_state(circuit_params, input_state):
    circuit_params_with_input_state = {**circuit_params, "input_state": input_state}
    return get_full_variational_quantum_circuit(**circuit_params_with_input_state)

def transfrom_hamiltonian_into_pauli_strings(hamiltonian) -> List:
    pauli_operators = hamiltonian.to_pauli_op().settings['oplist']
    pauli_coeffs = list(map(lambda pauli_operator: pauli_operator.coeff, pauli_operators))
    pauli_strings = list(map(lambda pauli_operator: pauli_operator.primitive, pauli_operators))
    return pauli_coeffs, pauli_strings

from qiskit.circuit.library.standard_gates import HGate, SGate
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister

reducing_to_pauli_z_mapping = {
    'I': 'I',
    'Z': 'Z',
    'X': 'Z',
    'Y': 'Z'
} 

def reduce_pauli_matrixes_into_sigma_z(pauli_string) -> str:
    reduced_pauli_string = ""
    for matrix_index in range(QUBITS_NUM):
        pauli_matrix = str(pauli_string[matrix_index])
        reduced_pauli_matrix = reducing_to_pauli_z_mapping[pauli_matrix]
        reduced_pauli_string = reduced_pauli_matrix + reduced_pauli_string
    
    return reduced_pauli_string

def add_layer_of_gates_for_reducing_paulis_to_sigma_z(pauli_string, quantum_circuit):
    quantum_registers = QuantumRegister(QUBITS_NUM, name="qubit")
    additional_circuit_layer = QuantumCircuit(quantum_registers)
    
    for quantum_register_index, pauli_matrix in enumerate(pauli_string):
        if pauli_matrix == "X":
            additional_circuit_layer.append(HGate(), [quantum_registers[quantum_register_index]])
        if pauli_string == "Y":
            additional_circuit_layer.append(HGate(), [quantum_registers[quantum_register_index]])
            additional_circuit_layer.append(SGate(), [quantum_registers[quantum_register_index]])
                
    extended_quantum_circuit = quantum_circuit.compose(additional_circuit_layer)
    return extended_quantum_circuit

def get_probability_distribution(counts: Dict) -> Dict:
    proba_distribution = {state: (count / NUM_SHOTS) for state, count in counts.items()}
    return proba_distribution

def calculate_probabilities_of_measurments_in_computational_basis(quantum_state_circuit) -> Dict:
    quantum_state_circuit.measure_all()
    
    transpiled_quantum_state_circuit = transpile(quantum_state_circuit, simulator_backend) 
    Qobj = assemble(transpiled_quantum_state_circuit)
    result = simulator_backend.run(Qobj).result()
    counts = result.get_counts(quantum_state_circuit)
    
    return get_probability_distribution(counts)

def sort_probas_dict_by_qubits_string_keys(proba_distribution: Dict) -> Dict:
    return dict(sorted(proba_distribution.items()))

def reset_power_of_minus_1(power_of_minus_1):
    power_of_minus_1 = 0
    return power_of_minus_1

def convert_pauli_string_into_str(pauli_string) -> str:
    return str(pauli_string)

def calculate_expectation_value_of_pauli_string_by_measurments_probas(pauli_string, ansatz_circuit):
    pauli_string_expectation_value = 0
    power_of_minus_1 = 0
    
    pauli_string_str = convert_pauli_string_into_str(pauli_string)
    extended_ansatz_circuit = add_layer_of_gates_for_reducing_paulis_to_sigma_z(pauli_string_str, ansatz_circuit)
    probas_distribution = calculate_probabilities_of_measurments_in_computational_basis(extended_ansatz_circuit)
    
    reduced_pauli_string = reduce_pauli_matrixes_into_sigma_z(pauli_string)
    sorted_probas_distribuition = sort_probas_dict_by_qubits_string_keys(probas_distribution)
    for qubits_string, proba in sorted_probas_distribuition.items():
        for string_index in range(QUBITS_NUM):
            if(str(qubits_string[string_index])=="1" and str(pauli_string[string_index])=="Z"):
                power_of_minus_1 += 1
            
        pauli_string_expectation_value += pow(-1, power_of_minus_1)*proba
        power_of_minus_1 = reset_power_of_minus_1(power_of_minus_1)
        
    return pauli_string_expectation_value

def get_expectation_value(ansatz_circuit, pauli_coeffs, pauli_strings):
    total_expection_value = 0
    
    for pauli_coeff, pauli_string in zip(pauli_coeffs, pauli_strings):
        total_expection_value += pauli_coeff*calculate_expectation_value_of_pauli_string_by_measurments_probas(
                                                                                    pauli_string, ansatz_circuit)
    
    return total_expection_value

from qiskit import assemble, transpile

def cost_function(thetas, hamiltonian):
    circuit_params = prepare_circuit_params(thetas)
    computational_eigenvectors = get_first_k_eigenvectors_from_n_computational_basis(K, N)
    
    pauli_coeffs, pauli_strings = transfrom_hamiltonian_into_pauli_strings(hamiltonian)
    k_ansatz_state = get_ansatz_state(circuit_params, computational_eigenvectors[K-1])
    
    approximated_energey = get_expectation_value(k_ansatz_state, pauli_coeffs, pauli_strings)
    insert_approximated_energy_to_list_of_all_approximated_energies(approximated_energey)
    
    L_w = w*approximated_energey
    for j in range(K-1):
        ansatz_state = get_ansatz_state(circuit_params, computational_eigenvectors[j])
        L_w += get_expectation_value(ansatz_state, pauli_coeffs, pauli_strings)
        
    return L_w

def get_optimal_thetas_of_ansatz_circuit_for_hamiltonian(hamiltonian):
    initial_thetas = np.random.uniform(low=0, high=2*np.pi, size=12)
    optimizer_result = minimize(cost_function,
                                x0=initial_thetas,
                                args=(hamiltonian),
                                method="COBYLA",
                                options={"disp": True, "maxiter":NUM_ITERATIONS})
    optimal_thetas = prepare_circuit_params(optimizer_result.x)
    
    return optimal_thetas

def get_approximated_eigenvalue_of_hamiltonian(hamiltonian):
    optimal_thetas = get_optimal_thetas_of_ansatz_circuit_for_hamiltonian(hamiltonian)
    print(optimal_thetas)
    computational_eigenvectors = get_first_k_eigenvectors_from_n_computational_basis(K, N)
    optimal_ansatz_state = get_ansatz_state(optimal_thetas, computational_eigenvectors[K-1])
    
    pauli_coeffs, pauli_strings = transfrom_hamiltonian_into_pauli_strings(hamiltonian)
    approximated_eigenvalue = get_expectation_value(optimal_ansatz_state, pauli_coeffs, pauli_strings)

    return approximated_eigenvalue

from numpy import linalg as LA

def get_approximation_error(exact_eigenvalue, approximated_eigenvalue):
    return abs(abs(exact_eigenvalue)-abs(approximated_eigenvalue))/abs(exact_eigenvalue)

def get_k_exact_eigenvalue_of_hamiltonian(hamiltonian, k):
    eigen_values = LA.eig(hamiltonian.to_matrix())[0]
    print(sorted(eigen_values, reverse=True))
    
    return sorted(eigen_values,reverse=True)[k-1]

def compare_exact_and_approximated_eigenvalue(hamiltonian, approximated_eigenvalue):
    exact_eigenvalue = get_k_exact_eigenvalue_of_hamiltonian(hamiltonian, K)
    print("Exact Eigenvalue:")
    print(exact_eigenvalue)
    print("Approximated K Eigenvalues:")
    print(approximated_eigenvalue)

    print("Approximation Error")
    print(get_approximation_error(exact_eigenvalue, approximated_eigenvalue))

def insert_approximated_energy_to_list_of_all_approximated_energies(energy):
    approximated_energies.append(energy)

import matplotlib.pyplot as plt

def plot_convergence_of_optimization_process(approximated_energies, exact_eigenvalue, margin):
    plt.title("convergence of optimization process to the exact eigenvalue")
    plt.margins(0, margin)
    plt.plot(approximated_energies)
    plt.axhline(y = exact_eigenvalue, color = 'r', linestyle = '-')
    plt.xlabel("# of iterations")
    plt.ylabel("Energy")
    
def plot_fidelity():
    plt.plot(LiH_approximated_energies)
    plt.xlabel("# of iterations")
    plt.ylabel("Energy")

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

LiH_molecule_4_qubits = -7.49894690201071*(I^I^I^I) + \
-0.0029329964409502266*(X^X^Y^Y) + \
0.0029329964409502266*(X^Y^Y^X) + \
0.01291078027311749*(X^Z^X^I) + \
-0.0013743761078958677*(X^Z^X^Z) + \
0.011536413200774975*(X^I^X^I) + \
0.0029329964409502266*(Y^X^X^Y) + \
-0.0029329964409502266*(Y^Y^X^X) + \
0.01291078027311749*(Y^Z^Y^I) + \
-0.0013743761078958677*(Y^Z^Y^Z) + \
0.011536413200774975*(Y^I^Y^I) + \
0.16199475388004184*(Z^I^I^I) + \
0.011536413200774975*(Z^X^Z^X) + \
0.011536413200774975*(Z^Y^Z^Y) + \
0.12444770133137588*(Z^Z^I^I) + \
0.054130445793298836*(Z^I^Z^I) + \
0.05706344223424907*(Z^I^I^Z) + \
0.012910780273117487*(I^X^Z^X) + \
-0.0013743761078958677*(I^X^I^X) + \
0.012910780273117487*(I^Y^Z^Y) + \
-0.0013743761078958677*(I^Y^I^Y) + \
0.16199475388004186*(I^Z^I^I) + \
0.05706344223424907*(I^Z^Z^I) + \
0.054130445793298836*(I^Z^I^Z) + \
-0.013243698330265966*(I^I^Z^I) + \
0.08479609543670981*(I^I^Z^Z) + \
-0.013243698330265952*(I^I^I^Z)

%%time
LiH_approximated_eigenvalue = get_approximated_eigenvalue_of_hamiltonian(LiH_molecule_4_qubits)

compare_exact_and_approximated_eigenvalue(LiH_molecule_4_qubits, LiH_approximated_eigenvalue)

print(approximated_energies)
approximated_energies = []

LiH_approximated_energies = [-7.6274591415582105, -7.607576828257072, -7.632200935003245, -7.609855196623901, -7.6059250397280795, -7.854597415729809, -7.628956794318834, -7.6784378048763715, -7.640342773091634, -7.665298199108273, -7.4589086431520615, -7.64400728822478, -7.63919074876535, -7.5712982052834015, -7.644638404908841, -7.646954760531615, -7.755301596699023, -7.631831246253383, -7.631061799605215, -7.616310788526333, -7.611576222526171, -7.620105483420993, -7.6417862171136335, -7.589952292685995, -7.610418644012381, -7.674320046442037, -7.61956694546109, -7.627683802497983, -7.623678295984975, -7.617228428549745, -7.633104604686701, -7.629086274819348, -7.595694407363554, -7.617836676264696, -7.617389773769819, -7.634299221424122, -7.598467942904893, -7.630191480144404, -7.63698619980098, -7.623626950424333, -7.612007795510998, -7.6261074907011634, -7.6148982305668405, -7.623709634670492, -7.61204075447679, -7.6106110013026385, -7.620711346718306, -7.618187492707977, -7.612462844967995, -7.623640830688807]
plot_convergence_of_optimization_process(LiH_approximated_energies, exact_eigenvalue=-7.7140566916607005,margin=1)

H2_molecule_Hamiltonian_4_qubits =  -0.8105479805373279 * (I^I^I^I) \
                                    + 0.1721839326191554 * (I^I^I^Z) \
                                    - 0.22575349222402372 * (I^I^Z^I) \
                                    + 0.17218393261915543 * (I^Z^I^I) \
                                    - 0.2257534922240237 * (Z^I^I^I) \
                                    + 0.12091263261776627 * (I^I^Z^Z) \
                                    + 0.16892753870087907 * (I^Z^I^Z) \
                                    + 0.045232799946057826 * (Y^Y^Y^Y) \
                                    + 0.045232799946057826 * (X^X^Y^Y) \
                                    + 0.045232799946057826 * (Y^Y^X^X) \
                                    + 0.045232799946057826 * (X^X^X^X) \
                                    + 0.1661454325638241 * (Z^I^I^Z) \
                                    + 0.1661454325638241 * (I^Z^Z^I) \
                                    + 0.17464343068300453 * (Z^I^Z^I) \
                                    + 0.12091263261776627 * (Z^Z^I^I)

%%time
H2_approximated_eigenvalue = get_approximated_eigenvalue_of_hamiltonian(H2_molecule_Hamiltonian_4_qubits)

compare_exact_and_approximated_eigenvalue(H2_molecule_Hamiltonian_4_qubits, H2_approximated_eigenvalue)

print(approximated_energies)
approximated_energies = []

H2_approximated_energies = [-1.009976657006965, -0.9238284763706393, -0.9884320852467877, -0.9805917857984984, -1.0167176344158073, -0.803009921527331, -0.4598306468944341, -0.5114890962692614, -0.8589176393087158, -0.9346855734783651, -0.8412382644084648, -0.6224404303456501, -0.653317087612677, -0.7607656559263039, -0.7930275875404521, -0.4281158503125878, -0.5788815212789364, -0.7176165386431208, -0.8701142070097521, -0.8688095826524974, -0.9151572025822381, -0.6635735611840398, -0.6037067550586409, -0.9285839833437391, -0.5058550482270441, -0.9291800596309192, -0.49550483773798376, -0.7464398349372543, -0.9433354608506894, -0.8895971844704103, -0.9048816527272305, -0.8978302385287464, -0.6949574562429168, -0.8304208103551916, -0.4692230966986817, -0.8462497998117152, -0.7185585023635032, -0.8077110062007753, -0.5192275804166049, -0.5640061209546482, -0.8011222125790844, -0.9513753098344495, -0.8402620344123755, -0.6086751350074949, -0.7369530952207817, -0.7603612456563129, -0.7950958666448285, -0.7999579057635281, -0.7706575601258369, -0.7909119439657882]
plot_convergence_of_optimization_process(H2_approximated_energies, exact_eigenvalue=-1.2445845498133272, margin=5)

transverse_ising_4_qubits = 0.0 * (I^I^I^I) \
    + 0.8398088405253477 * (X^I^I^I) \
    + 0.7989496312070936 * (I^X^I^I) \
    + 0.38189710487113193 * (Z^Z^I^I) \
    + 0.057753122422666725 * (I^I^X^I) \
    + 0.5633292636970458 * (Z^I^Z^I) \
    + 0.3152740621483513 * (I^Z^Z^I) \
    + 0.07209487981989715 * (I^I^I^X) \
    + 0.17892334004292654 * (Z^I^I^Z) \
    + 0.2273896497668042 * (I^Z^I^Z) \
    + 0.09762902934216211 * (I^I^Z^Z)

%%time
TI_approximated_eigen_value = get_approximated_eigenvalue_of_hamiltonian(transverse_ising_4_qubits)

compare_exact_and_approximated_eigenvalue(transverse_ising_4_qubits, TI_approximated_eigen_value)

print(approximated_energies)
approximated_energies = []

TI_approximated_energies = [1.0877911566875362, 1.3339897771771856, 1.1497195956414945, 1.3463775026022828, 1.1824868775928117, 1.6423958284840114, 1.5842407965410792, 1.6377219801098, 1.5250316270225506, 1.626129653972439, 1.3560460449434424, 1.6695319025733455, 1.6842571288891142, 1.5407570120203158, 1.5187677867233031, 1.4564666271629088, 1.7323725293893963, 1.792109013228278, 1.7481282811675565, 1.8139670143156743, 1.6629383661129478, 1.68337823955209, 1.4848944286809986, 1.4845001982657569, 1.6129182843066505, 1.6625400070359255, 1.6104761870181472, 1.7152744787413363, 1.6678413737920048, 1.6869395320619605, 1.6590487040095492, 1.6723344557176432, 1.627252139576561, 1.6998615622327233, 1.6227578573441954, 1.6709380388303952, 1.747587311276882, 1.7146592410359045, 1.7219455615683754, 1.6503664867345986, 1.7349602602365404, 1.767898541795154, 1.7803996165862503, 1.6955748077929296, 1.7174656756704858, 1.715728620581076, 1.7002352954018423, 1.7436412571902764, 1.6757447332820925, 1.7032804122037015]
plot_convergence_of_optimization_process(TI_approximated_energies, exact_eigenvalue=-1.7583827504312988, margin=3)

