import itertools
import matplotlib.pyplot as plt
import numpy as np
import dimod
%matplotlib inline
np.set_printoptions(precision=3, suppress=True)

n_spins = 10
n_samples = 1000
h = {v: np.random.uniform(-2, 2) for v in range(n_spins)}
J = {}
for u, v in itertools.combinations(h, 2):
    if np.random.random() < .05:
        J[(u, v)] = np.random.uniform(-1, 1)
model = dimod.BinaryQuadraticModel(h, J, 0.0, dimod.SPIN)
sampler = dimod.SimulatedAnnealingSampler()        

temperature_0 = 1
response = sampler.sample(model, beta_range=[1/temperature_0, 1/temperature_0], num_reads=n_samples)
energies_0 = [solution.energy for solution in response.data()]
temperature_1 = 10
response = sampler.sample(model, beta_range=[1/temperature_1, 1/temperature_1], num_reads=n_samples)
energies_1 = [solution.energy for solution in response.data()]
temperature_2 = 100
response = sampler.sample(model, beta_range=[1/temperature_2, 1/temperature_2], num_reads=n_samples)
energies_2 = [solution.energy for solution in response.data()]

def plot_probabilities(energy_samples, temperatures):
    fig, ax = plt.subplots()
    for i, (energies, T) in enumerate(zip(energy_samples, temperatures)):
        probabilities = np.exp(-np.array(sorted(energies))/T)
        Z = probabilities.sum()
        probabilities /= Z
        ax.plot(energies, probabilities, linewidth=3, label = "$T_" + str(i+1)+"$")
    minimum_energy = min([min(energies) for energies in energy_samples])
    maximum_energy = max([max(energies) for energies in energy_samples])
    ax.set_xlim(minimum_energy, maximum_energy)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_xlabel('Energy')
    ax.set_ylabel('Probability')
    ax.legend()
    plt.show()

plot_probabilities([energies_0, energies_1, energies_2], 
                   [temperature_0, temperature_1, temperature_2])

import itertools
import numpy as np
from functools import reduce
from qiskit import BasicAer, QuantumRegister, QuantumCircuit, ClassicalRegister
from qiskit import execute
from qiskit.quantum_info import Pauli
from qiskit_aqua import get_aer_backend, QuantumInstance
from qiskit_aqua.operator import Operator
from qiskit_aqua.components.optimizers import COBYLA
from qiskit_aqua.algorithms import VQE
from qiskit_aqua.algorithms.adaptive.qaoa.varform import QAOAVarForm

n_qubits = 2
n_system = n_qubits * 2
β = 0

weights = np.array([[0,1],[0,0]])
p = 1

def pauli_z(qubit, coeff):
    eye = np.eye((n_system))
    return Operator([[coeff, Pauli(eye[qubit], np.zeros(n_system))]])

def product_pauli_z(q1, q2, coeff):
    eye = np.eye((n_system))
    return Operator([[coeff, Pauli(eye[q1], np.zeros(n_system)) * Pauli(eye[q2], np.zeros(n_system))]])

def ising_hamiltonian(weights):
    H = reduce(lambda x,y:x+y,
            [product_pauli_z(i,j, -weights[i,j])
             for (i,j) in itertools.product(range(n_qubits), range(n_qubits))])
    H.to_matrix()
    return H

Hc = ising_hamiltonian(weights)

qr = QuantumRegister(n_qubits * 2)
cr = ClassicalRegister(n_qubits)
backend = BasicAer.get_backend('qasm_simulator')
circuit_init = QuantumCircuit(qr)

α = 2 * np.arctan(np.exp(- β/2))
for i in range(n_qubits):
    circuit_init.rx(α, qr[n_qubits+i])
    circuit_init.cx(qr[n_qubits+i], qr[i])

class InitialState:
    def __init__(self, β, n_qubits):
        self.β = β
        self.n_qubits = n_qubits
    
    def construct_circuit(self, mode, qr):
        circuit_init = QuantumCircuit(qr)

        alpha = 2 * np.arctan(np.exp(- self.β / 2))
        for i in range(n_qubits):
            circuit_init.rx(alpha, qr[n_qubits+i])
            circuit_init.cx(qr[n_qubits+i], qr[i])
            
        return circuit_init

class MyQAOA(VQE):
    def __init__(self, operator, optimizer, init_state, p=1, operator_mode='matrix', initial_point=None,
                 batch_mode=False, aux_operators=None):
            self.validate(locals())
            var_form = QAOAVarForm(operator, p, init_state)
            super().__init__(operator, var_form, optimizer,
                             operator_mode=operator_mode, initial_point=initial_point)
        
class MyQAOAVarForm(QAOAVarForm):
    def construct_circuit(self, angles, quantum_register, classical_register):
        if not len(angles) == self.num_parameters:
            raise ValueError('Incorrect number of angles: expecting {}, but {} given.'.format(
                self.num_parameters, len(angles)
            ))
        q = quantum_register
        circuit = QuantumCircuit(quantum_register, classical_register)
        if self._initial_state:
            circuit += self._initial_state.construct_circuit('circuit', q)
        else:
            circuit.u2(0, np.pi, q)
        for idx in range(self._p):
            β, γ = angles[idx], angles[idx + self._p]
            circuit += self._cost_operator.evolve(None, γ, 'circuit', 1, quantum_registers=q)
            circuit += self._mixer_operator.evolve(None, β, 'circuit', 1, quantum_registers=q)
        return circuit

initial_state = InitialState(β, n_qubits)

def get_thermal_state(weights, p):
    Hc = ising_hamiltonian(weights)
    print("Begin QAOA...")
    
    optimizer = COBYLA()
    qaoa = MyQAOA(Hc, optimizer, initial_state, p, "matrix")
    backend = get_aer_backend('statevector_simulator')
    quantum_instance = QuantumInstance(backend, shots=100)
    result = qaoa.run(quantum_instance)
    print("Results of QAOA", result)
    
    return result
result = get_thermal_state(weights, 1)

var_form = MyQAOAVarForm(Hc, p, initial_state)
thermal_state = var_form.construct_circuit(result['opt_params'], qr, cr)
for i in range(n_qubits):
    thermal_state.measure(qr[i], cr[i])
    job = execute(thermal_state, backend, shots=2000)
results = job.result().get_counts(thermal_state)

def get_energy(spin_configuration):
    x = spin_configuration.reshape(-1, 1)
    return np.sum([[-weights[i,j] * x[i] * x[j] for j in range(n_qubits)] for i in range(n_qubits)])

list_spin_configs = np.array(np.concatenate([[list(spin_config)] * results[spin_config] for spin_config in results]), dtype="int")
list_spin_configs[list_spin_configs == 0] = -1
list_energy = np.array([get_energy(spin_config) for spin_config in list_spin_configs])
hist = plt.hist(list_energy, density=True)

β = 5
initial_state = InitialState(β, n_qubits)
result = get_thermal_state(weights, 1)

var_form = MyQAOAVarForm(Hc, p, initial_state)
thermal_state = var_form.construct_circuit(result['opt_params'], qr, cr)
for i in range(n_qubits):
    thermal_state.measure(qr[i], cr[i])
    job = execute(thermal_state, backend, shots=2000)
results = job.result().get_counts(thermal_state)

list_spin_configs = np.array(np.concatenate([[list(spin_config)] * results[spin_config] for spin_config in results]), dtype="int")
list_spin_configs[list_spin_configs == 0] = -1
list_energy = np.array([get_energy(spin_config) for spin_config in list_spin_configs])
plt.hist(list_energy, density=True)