from qiskit import *
import qiskit.tools.jupyter
import numpy as np
IBMQ.load_account()

from qiskit.aqua.components.variational_forms import VariationalForm
from qiskit.aqua.components.initial_states import Zero, VarFormBased
from qiskit.aqua.operators import MatrixOperator, TPBGroupedWeightedPauliOperator
from qiskit.aqua.operators.op_converter import to_tpb_grouped_weighted_pauli_operator
from qiskit.aqua.components.optimizers import SPSA
from qiskit.aqua.algorithms.adaptive import VQE

from scipy.sparse import *
from scipy import *

num_qubits = 3

num_var_params = 6
var_form_depth = 1
initial_params = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

class ConcatenatedSchwingerAnsatz(VariationalForm):
    """ Variational Layer from Martin's Paper """
    
    CONFIGURATION = {
        'name': 'Schwinger Ansatz',
        'description': 'Variational Form for Schwinger VQE',
        'input_schema': {
            '$schema': 'http://json-schema.org/draft-07/schema#',
            'id': 'schwinger_schema',
            'type': 'object',
            'properties': {
                'depth': {
                    'type': 'integer',
                    'default': 3,
                    'minimum': 1
                },
                'entangler_map': {
                    'type': ['array', 'null'],
                    'default': None
                },
                'entanglement_gate': {
                    'type': 'string',
                    'default': 'cx',
                    'enum': ['cx']
                },
            },
            'additionalProperties': False
        },
        'depends': [
            {
                'pluggable_type': 'initial_state',
                'default': {
                    'name': 'ZERO',
                }
            },
        ],
    }
    
    def __init__(self, depth=3, entangler_map=None, initial_state=None):
        """
        Constructor.
        Args:
            depth (int) : number of rotation layers
            entangler_map (list[list]): describe the connectivity of qubits, each list describes
                             [source, target], or None for full entanglement.
                             Note that the order is the list is the order of
                             applying the two-qubit gate.
            initial_state (InitialState): an initial state object
        """
        self.validate(locals())
        super().__init__()
        #_support_parameterized_circuit = True
        # hard-coded, first pass through FIXME
        num_qubits = 3
        num_parameters = 6
        entanglement='full'
        entanglement_gate='cx'
        skip_unentangled_qubits = False
        self._num_qubits = num_qubits 
        self._depth = depth
        if entangler_map is None:
            self._entangler_map = VariationalForm.get_entangler_map(entanglement, num_qubits)
        else:
            self._entangler_map = VariationalForm.validate_entangler_map(entangler_map, num_qubits)
        # determine the entangled qubits
        all_qubits = []
        for src, targ in self._entangler_map:
            all_qubits.extend([src, targ])
        self._entangled_qubits = sorted(list(set(all_qubits)))
        self._initial_state = initial_state
        self._entanglement_gate = entanglement_gate
        self._skip_unentangled_qubits = skip_unentangled_qubits
        # for the first layer
        self._num_parameters = num_parameters
        # for repeated block
        self._num_parameters += num_parameters * (depth-1)
        self._bounds = [(-np.pi, np.pi)] * self._num_parameters
    
    def construct_circuit(self, parameters, q=None):
        """
        Construct the variational form, given its parameters.
        
        Args:
            parameters (numpy.ndarray): circuit parameters
            q (QuantumRegister): Quantum Register for the circuit.
            
        Returns:
            QuantumCircuit: a quantum circuit with given `parameters`
            
        Raises:
            ValueError: the number of parameters is incorrect.
        """
        if len(parameters) != self._num_parameters:
            raise ValueError('The number of parameters has to be {}'.format(self._num_parameters))
        
        if q is None:
            q = QuantumRegister(self._num_qubits, name='q')
        if self._initial_state is not None:
            circuit = self._initial_state.construct_circuit('circuit', q)
        else:
            circuit = QuantumCircuit(q)
        #
        param_idx = 0
        
        def angle(theta=0., phi=0., lam=0.):
            return [theta, phi, lam]

        for _ in range(self._depth):
            # get variational parameters for this layer
            t0_1 = parameters[param_idx] # theta-angle values
            t1_1 = parameters[param_idx+1]
            t2_1 = parameters[param_idx+2]
            t0_2 = parameters[param_idx] # theta-angle values
            t1_2 = parameters[param_idx+1]
            t2_3 = parameters[param_idx+2]
            t1 = [angle(t_i) for t_i in [t0_1, t1_1, t2_1]]
            t2 = [angle(t_i) for t_i in [t0_1, t1_1, t2_1]]
            # FIRST
            circuit.u3(*(t1[1]), 1)  # u3(*angles, q)
            circuit.u3(*(t1[0]), 2)
            circuit.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
            circuit.u3(*(t1[0]), 1)
            circuit.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
            circuit.u3(*(t1[2]), 2)
            circuit.barrier(range(3))
            # SECOND
            circuit.cx(0,2)
            #
            circuit.u3(*(t2[1]), 1)  # u3(*angles, q)
            circuit.u3(*(t2[0]), 2)
            circuit.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
            circuit.u3(*(t2[0]), 1)
            circuit.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
            circuit.u3(*(t1[2]), 2)
            #
            circuit.cx(0,2)
            circuit.barrier(range(3))
        return circuit  

schwinger_form = ConcatenatedSchwingerAnsatz(var_form_depth)
vqs_circ = schwinger_form.construct_circuit(initial_params)
vqs_circ.draw(output='latex', scale=0.5)

#qiskit.compiler.transpile(qc, basis_gates=['u3', 'cx'], optimization_level=3).draw(output='latex')

# even space
# [0.13028114, -1.18048464, 0.17737505]
# [ 0.12398951, -1.18068356,  0.15356187]
t0_1 = 0.13028114 # theta-angle values
t1_1 = -1.18048464
t2_1 = 0.17737505
# odd space
# [-0.30306561, -1.02142291,  0.15013751
# [ 2.8238186 , -5.22031494,  0.16478762]
t0_2 = 2.8238186
t1_2 = -5.22031494
t2_2 = 0.16478762
def angle(theta=0., phi=0., lam=0.):
    return [theta, phi, lam]
t1 = [angle(t_i) for t_i in [t0_1, t1_1, t2_1]]
t2 = [angle(t_i) for t_i in [t0_2, t1_2, t2_2]]
# create registers
q = QuantumRegister(3, name="q")
c = ClassicalRegister(3, name="c")
regs = [q, c]
circ = QuantumCircuit(*regs, name="super+")
# FIRST
circ.u3(*(t1[1]), 1)  # u3(*angles, q)
circ.u3(*(t1[0]), 2)
circ.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
circ.u3(*(t1[0]), 1)
circ.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
circ.u3(*(t1[2]), 2)
circ.barrier(range(3))
# SECOND
circ.cx(0,2)
#
circ.u3(*(t2[1]), 1)  # u3(*angles, q)
circ.u3(*(t2[0]), 2)
circ.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
circ.u3(*(t2[0]), 1)
circ.cx(1, 2)        # entangling qubit cx(ctrl, tgt)
circ.u3(*(t1[2]), 2)
#
circ.cx(0,2)
circ.barrier(range(3))
# MEAS
meas = QuantumCircuit(*regs, name="meas")
meas.measure(q, c)#  map qubits to clbits
# Full circuit
qc = circ + meas
qc.draw(output='latex', scale=0.5, filename='my_circuit.png', plot_barriers=True, fold=40)

fiducial_state = Zero(num_qubits)
fid_circ = fiducial_state.construct_circuit()
fid_circ.draw(output='latex', scale=0.5)

var_form_wavefunction = VarFormBased(schwinger_form, initial_params)
var_form_wavefunction.construct_circuit().draw(output='latex', scale=0.5)

# these correspond to L=2 spatial sites. 
# They have slightly different truncations so that the dimensions of the matrices are the same.

h = np.array([ 
    [-0.20036, 1.19976, 0., 0., 0., 0., 0., 0.],            #1
    [1.19976, 1., 0.848358, 0., 0., 0., 0., 0.],            #2
    [0., 0.848358, 2.20036, 0.848358, 0., 0., 0., 0.],      #3
    [0., 0., 0.848358, 3., 0., 0., 0., 0.],                 #4
    [0., 0., 0., 0., 1., 0.848358, 0., 0.],                 #5
    [0., 0., 0., 0., 0.848358, 2.20036, -0.848358, 0.],     #6
    [0., 0., 0., 0., 0., -0.848358, 3., -0.848358],         #7
    [0., 0., 0., 0., 0., 0., -0.848358, 3.79964]            #8
])
# trying different ordering?
h = np.array([ 
    [-0.20036, 1.19976, 0., 0., 0., 0., 0., 0.],            #1
    [0., 0., 0., 0., 1., 0.848358, 0., 0.],                 #5
    [0., 0.848358, 2.20036, 0.848358, 0., 0., 0., 0.],      #3
    [0., 0., 0., 0., 0., -0.848358, 3., -0.848358],         #7
    [1.19976, 1., 0.848358, 0., 0., 0., 0., 0.],            #2
    [0., 0., 0.848358, 3., 0., 0., 0., 0.],                 #4
    [0., 0., 0., 0., 0.848358, 2.20036, -0.848358, 0.],     #6
    [0., 0., 0., 0., 0., 0., -0.848358, 3.79964]            #8
])

H_mat = csr_matrix(h)
H_MatrixOperator = MatrixOperator(H_mat)
H_Operator = to_tpb_grouped_weighted_pauli_operator(H_MatrixOperator, 
                                                    TPBGroupedWeightedPauliOperator.sorted_grouping)
print(H_Operator.print_details())

operator_mode=None
input_circuit=var_form_wavefunction.construct_circuit()
backend=None
qr=None
cr=None
use_simulator_operator_mode=False
use_simulator_snapshot_mode=False
wave_function=fid_circ
statevector_mode=False
input_circuit.draw(output='latex')

H_circs = H_Operator.construct_evaluation_circuit(input_circuit, # wavefunction
                                        statevector_mode,
                                        qr,
                                        cr,
                                        use_simulator_snapshot_mode)

H_circs[0].draw(output='latex', scale=0.5)

max_trials = 500
optimizer = qiskit.aqua.components.optimizers.SPSA(max_trials)
print(optimizer.setting) 

# a callback that can access the intermediate data 
# during the optimization.
# Internally, four arguments are provided as follows 
# the index of evaluation, parameters of variational form,
# evaluated mean, evaluated standard deviation.
def simple_callback(eval_count, parameter_set, mean, std):
    if eval_count % 1 == 0:
        print('Energy evaluation %s returned %4f +/- %4f' % (eval_count, np.real(mean), np.real(std)))

# classVQE(operator, var_form, optimizer, 
# operator_mode=None, initial_point=None, 
# max_evals_grouped=1, 
# aux_operators=None, callback=None, 
# auto_conversion=True)
operator_mode = None
initial_point = initial_params
max_evals_grouped=1              # max number of evaluations performed simultaneously
aux_operators = None
auto_conversion=False
sim = VQE(H_Operator, schwinger_form, optimizer,
          # operator_mode, not needed in the most recent version
          initial_point, 
          max_evals_grouped, 
          aux_operators, 
          simple_callback, 
          auto_conversion)

backend = None
use_simulator_operator_mode=False
# is backend from AerProvider,
# if True and mode is paulis, single circuit is generated.
use_simulator_snapshot_mode=True
statevector_mode=False

sim_circs = sim.construct_circuit(initial_params, 
                                  statevector_mode,
                                  use_simulator_snapshot_mode)

backend_sim = Aer.get_backend('qasm_simulator')

from qiskit.aqua import QuantumInstance
nshots = 1024
my_quantum_instance = QuantumInstance(backend_sim, nshots)

print(my_quantum_instance)

print(sim.print_settings())

res = sim.run(my_quantum_instance)

res

sim.get_optimal_vector()

sim.get_optimal_cost()

energy_eval = sim.get_optimal_cost()

energy_eval_err = 0.011

sim._ret['opt_params']

sim._energy_evaluation(sim._ret['opt_params'])
sim._energy_evaluation(sim._ret['opt_params'])
sim._energy_evaluation(sim._ret['opt_params'])
sim._energy_evaluation(sim._ret['opt_params'])
sim._energy_evaluation(sim._ret['opt_params'])

sim.quantum_instance.backend

sim.quantum_instance.run_config.shots

sim.quantum_instance.set_config(shots=1024, memory=True)



sim.quantum_instance.run_config

from qiskit.aqua.utils.run_circuits import find_regs_by_name
opt_circ = sim.get_optimal_circuit()
c = ClassicalRegister(opt_circ.width(), name='c')
q = find_regs_by_name(opt_circ, 'q')
opt_circ.add_register(c)
opt_circ.barrier(q)
opt_circ.measure(q, c)
opt_circ.draw(output='latex', filename='example_gate.png', scale=1.0)

# takes in results from a single circuit 
# with memory=True, and estimates
# the uncertanities on the probability of each outcome 
# (currently with bootstrap)
# returns outcome:(pba, std) dictionary
def compute_stats(res, invert_qubit_order = False):
    # check that there is a single circuit in the results
    # (ambiguous otherwise)
    assert(len(res.results)==1)
    # generates all possible outcomes given \
    # the number of qubits 
    def _generate_bitstrings(nqubits, invert_qubit_order):
        # (recursive, modifies all_strings in place)
        def _generate_bitstrings_rec(nqubits, all_strings, a_string, irecur): 
            # base
            if irecur == nqubits: 
                all_strings.append(''.join([bit for bit in a_string]))
                return 
            # append 0
            _generate_bitstrings_rec(nqubits, all_strings, a_string + ['0'], irecur + 1)
            # append 1
            _generate_bitstrings_rec(nqubits, all_strings, a_string + ['1'], irecur + 1)
        all_strings = []
        _generate_bitstrings_rec(nqubits, all_strings, [], 0)
        if (invert_qubit_order): 
        # pesky Qiskit messes up qubit ordering... this may or may not be necessary to translate the results.
            all_strings = [''.join(reversed(bitstring)) for bitstring in all_strings]
        return all_strings
    # bootstrap specific?
    # given ensembles of outcomes and a particular outcome
    # calculates the statistics of that outcome:
    # returns outcome:(pba, std) estimates for the outcome
    def _calc_outcome_stats(ensembles, nshots, outcome):
        cts = np.count_nonzero(ensembles==outcome, axis=0)
        pba = np.mean(cts)/nshots
        std = np.std(cts/nshots, ddof = 1) # use unbiased estimator
        return (pba, std)
    nqubits = int(np.log2(len(res.get_counts(0)))) # pull number of qubits from results
    outcomes = _generate_bitstrings(nqubits, invert_qubit_order)
    mem = res.get_memory(0)
    nshots = sum(list(res.get_counts(0).values()))
    nens = nshots # choose number of ensemles = number of samples
    nsam = nshots
    ensembles = np.random.choice(mem, (nens, nsam))
    stats = map(lambda outcome: (outcome, _calc_outcome_stats(ensembles, nshots, outcome)), outcomes)
    return dict(stats)

sim.quantum_instance.set_config(shots=8*1024, memory=True)
opt_res = sim.quantum_instance.execute(opt_circ)
sim.quantum_instance.set_config(shots=8*1024, memory=False) # reset to no memory! saves resources?

from qiskit.visualization import plot_histogram
plot_histogram(opt_res.get_counts(opt_circ))

transp_circ = transpile(opt_circ, sim.quantum_instance.backend, optimization_level=3)
transp_circ.draw(output='latex', scale=0.5)

opt_circ.decompose().draw(output='latex', scale=0.5)

measured_coeffs =compute_stats(opt_res, False)

measured_coeffs

# compare with Mathematica's expectations
mathematica_res = -1.01162 # Even
mathematica_res2 = 0.486204 # Odd

# odd coefficients
even_coeffs = [0.669069, 0.305919, 0.023941, 0.00107068]
even_coeffs_at_theta = [0.664109, 0.309821, 0.0236388, 0.00243079]
odd_coeffs = [0.7041621406145364, 0.2582831083289985, 0.03524433188695736, 0.002310419169507747]

import matplotlib.pyplot as plt
outcomes = list(measured_coeffs.keys())
outcomes_bars = np.arange(len(outcomes))
avgs = [measured_coeffs[outcome][0] for outcome in outcomes]
stds = [measured_coeffs[outcome][1] for outcome in outcomes]
fig, ax = plt.subplots()
ax.bar(outcomes_bars, avgs, yerr=stds, 
       align='center', alpha=0.5, ecolor='black',
       edgecolor='k',
       error_kw=dict(lw=0.5, capsize=10, capthick=2.0),
       label='IBM sim')
ax.set_ylabel(r'$|c_i|^2$', fontsize=16)
ax.set_xlabel(r'outcome', fontsize=16)
ax.set_xticks(outcomes_bars)
ax.set_xticklabels(outcomes)
ax.set_title('ibmqx2: $E \\pm \\Delta E = %.3f \\pm %.3f$\n($E = %.3f$ exact),\n nshots=%d' 
             % (energy_eval, energy_eval_err, mathematica_res, 
                sim.quantum_instance.run_config.shots), fontsize=16)
for icoeff in range(len(even_coeffs)):
    ax.axhline(y = even_coeffs[icoeff], color ='red', linewidth = 0.5, 
               label=r'exact $|c_{%d}|^2 = %.3f$' % (icoeff+1, even_coeffs[icoeff]))
    ax.axhline(y = even_coeffs_at_theta[icoeff], color ='green', linewidth = 0.5, 
               label=r'exact $|c(\theta)_{%d}|^2 = %.3f$''\n'r'vs measured $%.3f \pm %.3f$' % (icoeff+1, even_coeffs_at_theta[icoeff], *list(measured_coeffs.values())[icoeff]))
ax.set_yscale('log')
# plt.tight_layout()
plt.legend(frameon=False, fontsize=16, bbox_to_anchor=(1.04,1), loc="upper left")
plt.savefig("ibmqx2_log.png", dpi=300)

