# Importing useful packages
import sys
import time
import numpy as np
import pandas as pd
import seaborn as sns
from tabulate import tabulate
import matplotlib.pyplot as plt

# Importing QISKit
from qiskit import Aer, execute, BasicAer
from qiskit.aqua.components.initial_states import Zero
from qiskit.aqua.algorithms.adaptive import VQE
from qiskit.aqua import QuantumInstance, aqua_globals
from qiskit.aqua.operators import WeightedPauliOperator
from qiskit.aqua.algorithms import ExactEigensolver
from qiskit.aqua.components.initial_states import Custom
from qiskit.aqua.components.variational_forms import RY, RYRZ
from qiskit.aqua.operators.op_converter import to_weighted_pauli_operator
from qiskit.aqua.components.optimizers import COBYLA, SPSA, L_BFGS_B, SLSQP
from qiskit.aqua.operators import (TPBGroupedWeightedPauliOperator, WeightedPauliOperator, MatrixOperator)

import qiskit.tools.jupyter 
%qiskit_version_table

def run_VQE(**kwargs):
    """Runs the Variational Quantum Eigensolver (VQE)
    
    Args:
        file_path (str): path to txt file containing Hamiltonian
        simulator (str): [statevector_simulator, qasm_simulator] The type of simulator 
            you wish to use (statevector is noiseless, while qasm contains a noise model)
        
    Returns:
        return_dict (dict): Dictionary with the following Keys:
            file_name (str): Name of file used for simulation
            
            result_df (DataFrame): Pandas DataFrame containing the 
                convergence count (df index), convergence_vals, and percent_error
                
            reference_val (float): GS Energy found using the Exact Eigen Solver
            
            algos (QISKit VQE Object): Object returned by the VQE algorithm class
            
            algo_results (dict): Results from VQE run (one algorithim result dict
                for each optimizer) 
    """
    file_path = kwargs['file_path']
    init_state_name = kwargs['init_state_name']
    vf_name = kwargs['vf_name']
    depth = kwargs['depth']
    shots = kwargs['shots']
    simulator = kwargs['simulator']
    
    # Loading matrix representation of Hamiltonian txt file
    H = np.loadtxt(file_path)
    
    # Converting Hamiltonian to Matrix Operator
    qubit_op = MatrixOperator(matrix=H)
    
    # Converting to Pauli Operator
    qubit_op = to_weighted_pauli_operator(qubit_op)
    start = time.time()
    
    num_qubits = qubit_op.num_qubits
    num_paulis = len(qubit_op.paulis)
    num_qubits = qubit_op.num_qubits
    
    print(f"""{num_paulis} Pauli factors \n{round(time.time()-start)} s to process""")
    
    # Solving system exactly for reference value
    ee = ExactEigensolver(qubit_op)
    result = ee.run()
    ref = result['energy']

    # Setting initial state, variational form, and backend
    init_state = init_state_name(num_qubits)
    var_form = vf_name(num_qubits,
                      depth=depth,
                      entanglement='linear',
                      initial_state=init_state)
    
    backend = BasicAer.get_backend(simulator)

    # Don't use SPSA if using a noiseless simulator
    if simulator == 'statevector_simulator':
        optimizers = [COBYLA, L_BFGS_B, SLSQP]  
    else:
        optimizers = [SPSA]
    
    # Initializing empty lists & dicts for storage
    dfs = []
    algos = {}
    algo_results = {}
    for optimizer in optimizers:
        # For reproducibility
        aqua_globals.random_seed = 250
        print(f'\rOptimizer: {optimizer.__name__}          ', end='')

        counts = []
        values = []
        params = []
        def store_intermediate_result(eval_count, parameters, mean, std):
            counts.append(eval_count)
            values.append(mean)
            params.append(parameters)
        
        # Running VQE 
        algo = VQE(qubit_op, var_form, optimizer(), callback=store_intermediate_result)
        quantum_instance = QuantumInstance(backend=backend, shots=shots)  
        algo_result = algo.run(quantum_instance)
        
        # Appending each optimizer data frame to one list
        dfs.append(pd.DataFrame(values, columns=['convergence_vals'], index=counts))
        algos[optimizer.__name__] = algo
        algo_results[optimizer.__name__] = algo_result

    print('\rOptimization complete')
    
    # Formatting return data frame
    results = pd.concat([df for df in dfs], keys=[optimizer.__name__ for optimizer in optimizers])
    results['energy_diff'] = results['convergence_vals'].apply(lambda x: abs(x - ref))
    results['percent_error'] = results['convergence_vals'].apply(lambda x: (abs(x - ref)/ref)*100)
    
    return_dict = {}
    return_dict['file_name'] = file_path
    return_dict['result_df'] = results
    return_dict['reference_val'] = ref
    return_dict['algos'] = algos
    return_dict['algo_results'] = algo_results
    
    return return_dict

def print_table(**result_dict):
    """Prints results of VQE simulation for each optimizer
    
    Args:
        result_dict (dict): Dictionary returned from run_VQE function
        
    Returns: None
    """
    
    algo_results = result_dict['algo_results']
    ref_val = result_dict['reference_val']
    file_name = result_dict['file_name']
    fn = file_name.split('/')[-1].replace(".txt", "")
    
    print(f'\nHamiltonian: {fn}')
    print(f'Reference Value: {ref_val}')
    
    optimizers = algo_results.keys()
    vals = []
    for opt in optimizers:
        vals.append([opt, 
                     algo_results[opt]['energy'], 
                     abs((algo_results[opt]['energy'] - ref_val)/ref_val)*100])
    
    sorted_vals = sorted(vals, key=lambda x: x[2])
    print(tabulate(sorted_vals, 
                   tablefmt="fancy_grid", 
                   headers=['Optimizer', 'VQE Energy', '% Error'], 
                   numalign="right"))

def visualize_results(**result_dict):
    """Graphs optimizer energy convergence and convergence of most effective
        optimizer.
        
    Args:
        result_dict (dict): Dictionary returned from run_VQE function
        
    Returns: None
    """ 
    result_df = result_dict['result_df']
    algo_results = result_dict['algo_results']
    ref_val = result_dict['reference_val']
    
    # Optimizer Convergence
    plt.figure(figsize=(10, 5))
    for opt in result_df.index.unique(0):
        result_df['energy_diff'][opt].plot(logy=True, label=opt)
        
    plt.title('Optimizer Energy Convergence')
    plt.xlabel('Eval count')
    plt.ylabel('Energy difference from solution reference value')
    plt.legend()
    sns.despine()
    plt.show();
    
    # Most effective optimizer
    most_eff = sorted([(opt, 
                        result_df['percent_error'][opt].iloc[-1]) for opt in result_df.index.unique(0)], 
                      key=lambda x: x[1])
    fig, ax = plt.subplots(figsize=(10, 5))
    result_df['convergence_vals'][most_eff[0][0]].plot(label=f"{most_eff[0][0]} = {algo_results[most_eff[0][0]]['energy']:.6f}")

    ax.hlines(y=ref_val, 
              xmin=0, 
              xmax=len(result_df['convergence_vals'][most_eff[0][0]]), 
              colors='r', 
              label=f'Exact = {ref_val:.6f}')

    plt.title('Convergence', size=24)
    plt.xlabel('Optimization Steps', size=18)
    plt.ylabel('Energy', size=18)
    plt.legend()
    sns.despine()
    plt.show();
        
def plt_state_vector(**result_dict):
    """Plots state vector for each optimizer
    
    Args:
       result_dict (dict): Dictionary returned from run_VQE function 
    
    Returns: None
    """
    algos = result_dict['algos']
    algo_results = result_dict['algo_results']
    optimizers = result_dict['algos'].keys()
    
    for optimizer in optimizers:
        circ = algos[optimizer].construct_circuit(algo_results[optimizer]['opt_params'])
    
        e0wf = execute(circ[0],
                       Aer.get_backend('statevector_simulator'),
                       shots=1).result().get_statevector(circ[0])
        plt.figure(figsize=(10,5))
        plt.plot(np.flip(e0wf.real))
        plt.title(f'({optimizer}) State Vector', size=24)
        sns.despine()
        plt.show();

h_osc_params = {
    'file_path': r'JA_harmonic_oscillator_q=4.txt',
    'init_state_name': Zero,
    'vf_name': RY,
    'depth': 3,
    'shots': 1000,
    'simulator': 'statevector_simulator'
}

harmonic_osc = run_VQE(**h_osc_params)
print_table(**harmonic_osc)

for key in harmonic_osc.keys():
    print(f'{key}: {harmonic_osc[key]}')
    print('\n')

visualize_results(**harmonic_osc)
plt_state_vector(**harmonic_osc)

anh_gs = (3/8) * (6 / (1/2)**2) ** (1/3)
print(f'Upper bound for the ground state energy of the anharmonic oscillator: {anh_gs:.2f} eV')
print(f'Percent error: {abs((anh_gs-1.06)/1.06)*100:.2f}%')

anh_osc_params = {
    'file_path': r'JA_adapted_anharmonic_oscillator_q=4.txt',
    'init_state_name': Zero,
    'vf_name': RY,
    'depth': 5,
    'shots': 5000,
    'simulator': 'statevector_simulator'
}

anharmonic_osc = run_VQE(**anh_osc_params)
print_table(**anharmonic_osc)

visualize_results(**anharmonic_osc)
plt_state_vector(**anharmonic_osc)

%qiskit_copyright