# Installation of the requirements
#!python -m pip install -r requirements.txt

'''
(C) Renata Wong 2023

Qiskit code for testing fidelity of randomised classical shadow on the ground state energy of molecules.

Procedure:
1. Choose a variational ansatz with initial parameters selected at random.
2. Generate a set of random basis change operators.
3. Apply the random operators to change bases in the ansatz.
4. Measure the ansatz in the Pauli Z basis and store the results as a shadow.
5. Obtain the expectation value of the molecular Hamiltonian from the shadow.
6. Optimize for minimum Hamiltonian expectation value. 
7. Feed the calculated angles/parameters back to the ansatz.
8. Repeat steps 3-7 till the optimization is completed. 
9. Output the minimized expectation value of the molecular Hamiltonian and the mean-square-root-error. 

Note: Below we perform calculations on the molecular Hamiltonian of H_2.
To perform calculations on other molecules, you will need to specify their geometry, charge and spin 
to replace the values in the driver. 

Note: predicting_quantum_properties module comes from https://github.com/hsinyuan-huang/predicting-quantum-properties
'''

from qiskit import QuantumCircuit, execute
from qiskit_aer import QasmSimulator

from qiskit_nature.units import DistanceUnit
from qiskit_nature.second_q.drivers import PySCFDriver
from qiskit_nature.second_q.mappers import BravyiKitaevMapper

from qiskit.algorithms.minimum_eigensolvers import NumPyMinimumEigensolver
from qiskit_nature.second_q.algorithms import GroundStateEigensolver

from qiskit.circuit.library import EfficientSU2
from qiskit.algorithms.optimizers import SLSQP, COBYLA, SPSA

import time
import numpy as np
from functools import partial
import matplotlib.pyplot as plt
from collections import Counter

from predicting_quantum_properties.data_acquisition_shadow import randomized_classical_shadow
from predicting_quantum_properties.prediction_shadow import estimate_exp

# handle deprecation issues
import qiskit_nature
qiskit_nature.settings.use_pauli_sum_op = False
import h5py
H5PY_DEFAULT_READONLY=1





# classically obtained ground state energy
EXPECTED_EIGENVALUE = -1.86

# specifying the geometry of the molecule in question
driver = PySCFDriver(
    atom="H 0 0 0; H 0 0 0.735",
    basis="sto3g",
    charge=0,
    spin=0,
    unit=DistanceUnit.ANGSTROM,
)


problem = driver.run()
hamiltonian = problem.hamiltonian

# electronic Hamiltonian of the system
second_q_op = hamiltonian.second_q_op()

# The Bravyi-Kitaev repserentation of the Fermionic Hamiltonian
mapper = BravyiKitaevMapper()
bkencoded_hamiltonian = mapper.map(second_q_op)

print(bkencoded_hamiltonian)

def process_hamiltonian(hamiltonian, derandomize = False):

    hamiltonian_observables = []
    hamiltonian_coefficients = []

    for observable in hamiltonian.paulis:
        op_list = []
        for op_index, pauli_op in enumerate(observable):
            pauli_op = str(pauli_op)
            if pauli_op == 'I' or pauli_op == 'X' or pauli_op == 'Y' or pauli_op == 'Z':
                op_list.append((pauli_op, op_index))
        hamiltonian_observables.append(op_list)
    
    hamiltonian_coefficients = hamiltonian.coeffs.real
    
    system_size = len(hamiltonian_observables[0])

    # removing all occurrences of Pauli-I, for all-Pauli-I there is an empty list left
    # these observables are needed for estimate_exp()
    observables_xyze = []

    for observable in hamiltonian_observables:
        XYZE = []
        for pauli in observable:
            if pauli[0] != 'I':
                XYZE.append(pauli)
        observables_xyze.append(XYZE)   
    
    # derandomisation procedure requires that coefficients are non-negative
    if derandomize == True:
        absolute_coefficients = [abs(coeffcient) for coeffcient in hamiltonian_coefficients]

        # removing the empty list as well
        # these observables are needed for derandomisation procedure
        observables_xyz = []
        for idx, observable in enumerate(observables_xyze):
            if observable:
                observables_xyz.append(observable)
            else:
                absolute_coefficients.pop(idx)
        
        return system_size, observables_xyze, observables_xyz, hamiltonian_coefficients, absolute_coefficients
    
    return system_size, observables_xyze, hamiltonian_coefficients



# process the Hamiltonian to obtain properly formatted data
hamiltonian_data = process_hamiltonian(bkencoded_hamiltonian, derandomize = False)
system_size, observables_xyze, hamiltonian_coefficients = hamiltonian_data

#print('HAMILTONIAN\n', observables_xyze)

'''
VARIATIONAL ANSATZ
Note that for molecules other than H_2 you may need to specify a different number of reps.
'''

reps = 1   
ansatz = EfficientSU2(system_size, su2_gates=['rx', 'ry'], entanglement='circular', reps=reps, skip_final_rotation_layer=False)

    
ansatz.decompose().draw('mpl')

'''
COST FUNCTION
'''


def basis_change_circuit(pauli_op):

    #   Generating circuit with just the basis change operators
    #
    #       pauli_op: n-qubit Pauli operator

    basis_change = QuantumCircuit(ansatz.num_qubits, ansatz.num_qubits)
    
    for idx, op in enumerate(pauli_op):
        if op == 'X':
            basis_change.h(idx)
        elif op == 'Y':
            basis_change.h(idx)
            basis_change.p(-np.pi/2, idx)

    return basis_change



def min_cost():
    
    return min(cost_history)


def log_cost(cost):

    global cost_history
    cost_history.append(cost)
    
    


def objective_function(operators, params):
    
    backend = QasmSimulator(method='statevector', shots=1)
    
    pauli_op_dict = Counter(tuple(x) for x in operators)    
    
    shadow = []
    for pauli_op in pauli_op_dict:
        
        qc = ansatz.bind_parameters(params)
        qc = qc.compose(basis_change_circuit(pauli_op))                                   
        qc.measure(reversed(range(system_size)), range(system_size))
        result = execute(qc, backend, shots=pauli_op_dict[pauli_op]).result()
        counts = result.get_counts()    
        
        # generate a separate circuit for each basis change Pauli operator
        output_str = list(list(counts.keys())[list(counts.values()).index(1)])
        output = [int(i) for i in output_str]
        eigenvals = [x+1 if x == 0 else x-2 for x in output]
        snapshot = [(op, eigenval) for op, eigenval in zip(pauli_op, eigenvals)] 
                
        shadow.append(snapshot)
    
    
    cost = 0.0
    for term, weight in zip(observables_xyze, hamiltonian_coefficients):
        sum_product, match_count = estimate_exp(shadow, term)
        if match_count != 0:
            exp_val = sum_product / match_count
            cost += weight * exp_val
            
    log_cost(cost)
    
    return cost


'''
RUNNING EXPERIMENTS
''' 

start_time = time.time()
rmse_errors = []

print('NUMBER OF OPERATORS | RANDOMISED OPERATORS | AVERAGE RMSE ERROR\n')
measurement_range = [50, 100, 200, 600, 1000, 1600]

for num_operators in measurement_range:   
    
    basis_change = randomized_classical_shadow(num_operators, system_size)

    tuples = (tuple(pauli) for pauli in basis_change)
    counts = Counter(tuples)

    optimizer = SPSA(maxiter=2000) 
    cost_function = partial(objective_function, basis_change)
    
    expectation_values = []
    num_experiments = 5


    for iteration in range(num_experiments):
        
        cost_history = []
        
        params = np.random.rand(ansatz.num_parameters)
        result = optimizer.minimize(fun=cost_function, x0=params)
        minimal_cost = min_cost()
        expectation_values.append(minimal_cost)
                
        print("EXPERIMENT {}: GROUND STATE ENERGY FOUND = {}".format(iteration, minimal_cost))


    rmse_randomised_cs = np.sqrt(np.sum([(expectation_values[i] - EXPECTED_EIGENVALUE)**2 
                                           for i in range(num_experiments)])/num_experiments)

    rmse_errors.append(rmse_randomised_cs)
    
    print('{} | {} | {}'.format(num_operators, counts, rmse_randomised_cs))


elapsed_time = time.time() - start_time
print("Execution time = ", time.strftime("%H:%M:%S", time.gmtime(elapsed_time)))

          
points = measurement_range
num_points = len(measurement_range)

plt.plot([i for i in points], [rmse_errors[i] for i in range(num_points)], 'r', label='SPSA(maxiter=2000)')
plt.xlabel('Number of measurements')
plt.ylabel('Average RMSE error')
plt.legend(loc=1)

