import time
import itertools
import numpy as np

import qiskit
from qiskit import *
from qiskit.quantum_info import Statevector
from qiskit.algorithms.optimizers import COBYLA, L_BFGS_B, SPSA
from qiskit.opflow.state_fns import StateFn, CircuitStateFn
from qiskit.providers.aer import StatevectorSimulator, AerSimulator

from qiskit.opflow import CircuitSampler

from qiskit.ignis.mitigation.measurement import CompleteMeasFitter # you will need to pip install qiskit-ignis
from qiskit.ignis.mitigation.measurement import complete_meas_cal

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import BoundaryNorm
cmap = plt.get_cmap("plasma") #'viridis'

from modules.utils import *

from qae import *

import datetime

print(qiskit.__version__, np.__version__)

#plt.rc('text', usetex=True)
plt.rc('font', family='serif',size=16)
#plt.rc('font', family='serif')

temp = np.load("data/QAEAnsatz_scaling_losses.npz", allow_pickle=True)
losses0 = temp["losses"]
Ls = temp["Ls"]

# It took 51 hrs to obtain these results.
# However, you can confirm the validity by just executing the circuit once for the optimized parameters (as demonstrated down below)
temp2 = np.load("data/QAEAnsatz_scaling_MPS_script_L-32_tracked-losses_results.npz",allow_pickle=True)
losses32 = temp2["losses"]

Ls = list(Ls)
Ls += [32]
Ls

fig, ax = plt.subplots(figsize=(6,5))
for i in range(4):
    ax.plot(np.array(losses0[i]),label=f"L = {Ls[i]}")
ax.plot(losses32[1],label=f"L = {Ls[-1]}")
ax.legend()

left, bottom, width, height = [0.4,0.61, 0.25, 0.325]
ax2 = fig.add_axes([left, bottom, width, height])
for i in range(4):
    ax2.plot(np.array(losses0[i]),label=f"L = {Ls[i]}")
ax2.plot(losses32[1],label=f"L = {Ls[-1]}")
ax2.set_yscale("log")

#ax.set_xlim(0,405)
#ax.set_yscale("log")
ax.set_ylabel("cost", fontsize=16)
ax.set_xlabel("iteration", fontsize=16)
plt.tight_layout()
plt.savefig("plots/QAEAnsatz_Ising-scaling.pdf")
plt.savefig("plots/QAEAnsatz_Ising-scaling.png")

#IBMQ.load_account() # this then automatically loads your saved account
#provider = IBMQ.get_provider(hub='ibm-q-research')

#device = provider.backend.ibmq_rome # 6 bogota ; 4 rome

### Real device execution:
#backend = device 

### Simulation with noise profile from real device
#backend = qiskit.providers.aer.AerSimulator.from_backend(device)

### Simulation without noise
backend = qiskit.providers.aer.AerSimulator()
#backend = StatevectorSimulator()

### Preliminaries
L = 5
num_trash = 2
anti = 1 # 1 for ferromagnetic Ising model, -1 for antiferromagnet

filename = "data/QAEAnsatz_scaling"

gz = 0
gx = 0.3

ED_state, ED_E, ham = ising_groundstate(L, anti, np.float32(gx), np.float32(gz))
Sen = ED_E
Smag = ED_state.T.conj()@Mag(L,anti)@ED_state

def prepare_circuit(thetas, L=6, num_trash=2, init_state=None, measurement=True,vqe=False):
    # QAE ansatz
    QAE_circ = QAEAnsatz(num_qubits = L, num_trash_qubits= num_trash, trash_qubits_idxs = list(range(num_trash)), measure_trash=measurement).assign_parameters(thetas)
    # initialize state vector
    initcirc = QuantumCircuit(QuantumRegister(L,"q"),ClassicalRegister(num_trash, 'c'))
    initcirc.initialize(init_state, initcirc.qubits)
    # compose circuits
    fullcirc = initcirc.compose(QAE_circ)
    return fullcirc

circ = prepare_circuit(thetas = np.random.rand(2*L+2), L = L, init_state = ED_state)
circ.draw("mpl")

### Execute circuit
count = 0
def run_circuit(thetas, L, num_trash, init_state, vqe=False, shots=100, meas_fitter = None):
    #global count
    #count += 1
    #print(count, "thetas: ", thetas)
    circ = prepare_circuit(thetas, L, num_trash, init_state, vqe=vqe)
    tcirc = qiskit.transpile(circ, backend)
    # Execute the circuit 
    job_sim = backend.run(tcirc, shots=shots,seed_simulator=333, seed_transpiler=444) # , seed_simulator=123, seed_transpiler=234 fix seed to make it reproducible
    result = job_sim.result()
    # Results without mitigation
    counts = result.get_counts()
    if meas_fitter != None:
        # Get the filter object
        meas_filter = meas_fitter.filter

        # Results with mitigation
        mitigated_results = meas_filter.apply(result)
        counts = mitigated_results.get_counts(0)
    return counts

res = run_circuit(thetas = np.random.rand(num_trash*L+num_trash), L = L, num_trash = num_trash, init_state = ED_state, shots=1000)

res

def count_ones(string):
    return np.sum([int(_) for _ in string])

count_ones("01010111")

[_ for _ in res]

[_ for _ in res if _ != "0" * num_trash]

np.sum([res[_]*count_ones(_) for _ in res if _ != "0" * num_trash]) # all measurement results except "000"

### Optimize circuit
def cost_function_single(thetas, L, num_trash, init_state, shots=1000, vqe=True, param_encoding=False, x=0, meas_fitter=None):
    """ Optimizes circuit """
    if param_encoding: thetas = feature_encoding(thetas, x) 
    out = run_circuit(thetas, L, num_trash, init_state, vqe=vqe, shots=shots, meas_fitter=meas_fitter)
    cost = np.sum([out[_]*count_ones(_) for _ in out if _ != "0" * num_trash]) # all measurement results except "000"
    return cost/shots

def cost_function(thetas, L, num_trash, init_states, shots=1000, vqe=True, param_encoding=False, x=0, meas_fitter=None):
    """ Optimizes circuit """
    cost = 0.
    for init_state in init_states:
        cost += cost_function_single(thetas, L, num_trash, init_state, shots, vqe, param_encoding, meas_fitter=meas_fitter)
    return cost/len(init_states)

def optimize(init_states, L=6, num_trash=2, thetas=None, shots=1000, max_iter=400, vqe=True, param_encoding=False, x=0, pick_optimizer = None,
            meas_fitter=None):
    if thetas is None:
        n_params = (num_trash*L+num_trasg)*2 if param_encoding else (num_trash*L+num_trash)
        thetas = np.random.uniform(0, 2*np.pi, n_params) # initial parameters without feature encoding
        
    #print("Initial cost: {:.3f}".format(cost_function(thetas, L, num_trash, init_states, shots, vqe, param_encoding, x)))
    
    counts, values, accepted = [], [], []
    def store_intermediate_result(eval_count, parameters, mean, std, ac):
        # counts.append(eval_count)
        values.append(mean)
        accepted.append(ac)

    # Initialize optimizer
    if pick_optimizer == "cobyla":
        optimizer = COBYLA(maxiter=max_iter, tol=0.0001)
    if pick_optimizer == "adam" or pick_optimizer == "ADAM":
        optimizer = qiskit.algorithms.optimizers.ADAM(maxiter=max_iter)
    # optimizer = L_BFGS_B(maxfun=300, maxiter=max_iter)#, factr=10, iprint=- 1, epsilon=1e-08)
    if pick_optimizer == "spsa" or pick_optimizer == None:
        optimizer = SPSA(maxiter=max_iter,
                         #blocking=True,
                         callback=store_intermediate_result,
                         #learning_rate=0.3,
                         #perturbation=0.1
                         ) # recommended from qiskit (first iteraction takes quite long)
                           # to reduce time figure out optimal learning rate and perturbation in advance
    start_time = time.time()
    ret = optimizer.optimize(
                            num_vars=len(thetas),
                            objective_function=(lambda thetas: cost_function(thetas, L, num_trash, init_states, shots, vqe, param_encoding, x, meas_fitter=meas_fitter)),
                            initial_point=thetas
                            )
    print("Time: {:.5f} sec".format(time.time()-start_time))
    print(ret)
    return ret[0], values, accepted



thetas_opt_mitigated, losses, accepted = optimize([ED_state], max_iter=120, L=5, meas_fitter=None) #, pick_optimizer="adam")

plt.plot(losses)

Ls = [3,4,8,16] #note that L=16 may take up few hours, it is faster with MPS further below
max_iter = [400] * len(Ls)
num_trashs = np.log(Ls)/np.log(2)
num_trashs = np.array(num_trashs, dtype="int")
gz = 0
gx = 0.3

losses = [None] * len(Ls); accepted = [None] * len(Ls); thetas_opt= [None] * len(Ls)

# may have to re-run some sizes when it gets stuck in local minima
for j,(L,num_trash) in enumerate(zip(Ls,num_trashs)):
    print(L,num_trash)
    ED_state, ED_E, ham = ising_groundstate(L, anti, np.float32(gx), np.float32(gz))
    thetas_opt[j], losses[j], accepted[j] = optimize([ED_state], max_iter=max_iter[j], L=L, num_trash=num_trash, meas_fitter=None) #, pick_optimizer="adam")
    plt.plot(losses[j])
    plt.show()

np.savez(filename + "_losses", losses=losses, thetas_opt = thetas_opt, Ls=Ls, max_iter=max_iter, num_trashs=num_trashs)

temp = np.load(filename + "_losses.npz", allow_pickle=True)
losses0 = temp["losses"]

fig, ax = plt.subplots(figsize=(6,5))
for i in range(4):
    ax.plot(np.array(losses0[i]) + 1e-4,label=f"L = {Ls[i]}")
ax.legend()
ax.set_xlim(0,400)
ax.set_yscale("log")
ax.set_yticks([1,1e-1,1e-2,1e-3,1e-4])
ax.set_yticklabels(["$10^{}$".format(i) for i in range(4)] + ["0.0"])
ax.set_ylabel("cost", fontsize=16)
ax.set_xlabel("iterations", fontsize=16)

min_losses = [np.min(l) for l in losses0]

min_losses

fig, ax = plt.subplots(figsize=(6,5))
ax.plot(Ls, min_losses,"x--")
#ax.set_yscale("log")
#ax.set_xscale("log")

def DMRG_Ising(L, J, g, chi_max=30, bc_MPS='finite'):
    print("finite DMRG, transverse field Ising model")
    print("L={L:d}, g={g:.2f}".format(L=L, g=g))
    model_params = dict(L=L, J=J, g=g, bc_MPS=bc_MPS, conserve=None)
    M = TFIChain(model_params)
    product_state = ["up"] * M.lat.N_sites
    psi = MPS.from_product_state(M.lat.mps_sites(), product_state, bc=M.lat.bc_MPS)
    dmrg_params = {
        'mixer': None,  # setting this to True helps to escape local minima
        'max_E_err': 1.e-10,
        'trunc_params': {
            'chi_max': chi_max,
            'svd_min': 1.e-10
        },
        'combine': True
    }
    info = dmrg.run(psi, M, dmrg_params)  # the main work...
    E = info['E']
    print("E = {E:.13f}".format(E=E))
    print("final bond dimensions: ", psi.chi)
    mag_x = np.sum(psi.expectation_value("Sigmax"))
    mag_z = np.sum(psi.expectation_value("Sigmaz"))
    print("magnetization in X = {mag_x:.5f}".format(mag_x=mag_x))
    print("magnetization in Z = {mag_z:.5f}".format(mag_z=mag_z))
    #if L < 20:  # compare to exact result
    #    from tfi_exact import finite_gs_energy
    #    E_exact = finite_gs_energy(L, 1., g)
    #    print("Exact diagonalization: E = {E:.13f}".format(E=E_exact))
    #    print("relative error: ", abs((E - E_exact) / E_exact))
    return E, psi #, M

### Preliminaries
qiskit_chi = 100
L = 8
num_trash = int(np.log(L)/np.log(2))
anti = 1 # 1 for ferromagnetic Ising model, -1 for antiferromagnet
g = 0.05
J=anti
filename = "data/QAEAnsatz_scaling_MPS_script_L-32_tracked-losses"

backend = qiskit.providers.aer.AerSimulator(method="matrix_product_state",
                                            precision="single",
                                            matrix_product_state_max_bond_dimension = qiskit_chi,
                                            matrix_product_state_truncation_threshold = 1e-10,
                                            #mps_sample_measure_algorithm = "mps_apply_measure", #alt: "mps_probabilities" 
                                           )

from qiskit.utils import algorithm_globals
algorithm_globals.random_seed = 333

def qiskit_state(psi0):
    # G is only the local tensor (not multiplied by any singular values) - see https://tenpy.readthedocs.io/en/latest/reference/tenpy.networks.mps.html
    A_list = [psi0.get_B(i, form="G").to_ndarray().transpose([1,0,2]) for i in range(L)]
    for i,A in enumerate(A_list):
        A_list[i] = (A[0], A[1])
    S_list = [psi0.get_SR(i) for i in range(L-1)] # skip trivial last bond; hast to be of size L-1
    return (A_list, S_list)

def prepare_circuit(thetas, L=6, num_trash=2, init_state=None, measurement=True,vqe=False):
    # QAE ansatz
    QAE_circ = QAEAnsatz(num_qubits = L, num_trash_qubits= num_trash, trash_qubits_idxs = list(range(L//2-1,L//2-1+num_trash)), measure_trash=measurement).assign_parameters(thetas)
    # initialize state vector
    initcirc = QuantumCircuit(QuantumRegister(L,"q"),ClassicalRegister(num_trash, 'c'))
    if init_state != None:
        initcirc.set_matrix_product_state(qiskit_state(init_state))
    # compose circuits
    fullcirc = initcirc.compose(QAE_circ)
    return fullcirc

### Execute circuit
count = 0
def run_circuit(thetas, L, num_trash, init_state, vqe=False, shots=100, meas_fitter = None):
    #global count
    #count += 1
    #print(count, "thetas: ", thetas)
    #print(L, num_trash)
    circ = prepare_circuit(thetas, L, num_trash, init_state, vqe=vqe)
    #circ.draw("mpl")
    #tcirc = qiskit.transpile(circ, backend)
    # Execute the circuit 
    job_sim = backend.run(circ, shots=shots, seed_simulator=333, seed_transpiler=444) #fix seed to make it reproducible
    result = job_sim.result()
    # Results without mitigation
    counts = result.get_counts()
    if meas_fitter != None:
        # Get the filter object
        meas_filter = meas_fitter.filter

        # Results with mitigation
        mitigated_results = meas_filter.apply(result)
        counts = mitigated_results.get_counts(0)
    return counts

def count_ones(string):
    return np.sum([int(_) for _ in string])

### Optimize circuit
def cost_function_single(thetas, L, num_trash, init_state, shots=1000, vqe=True, param_encoding=False, x=0, meas_fitter=None):
    """ Optimizes circuit """
    if param_encoding: thetas = feature_encoding(thetas, x) 
    out = run_circuit(thetas, L, num_trash, init_state, vqe=vqe, shots=shots, meas_fitter=meas_fitter)
    cost = np.sum([out[_]*count_ones(_) for _ in out if _ != "0" * num_trash]) # all measurement results except "000"
    return cost/shots

def cost_function(thetas, L, num_trash, init_states, shots=1000, vqe=True, param_encoding=False, x=0, meas_fitter=None):
    """ Optimizes circuit """
    cost = 0.
    for init_state in init_states:
        cost += cost_function_single(thetas, L, num_trash, init_state, shots, vqe, param_encoding, meas_fitter=meas_fitter)
    return cost/len(init_states)

def optimize(init_states, L=6, num_trash=2, thetas=None, shots=1000, max_iter=400, vqe=True, param_encoding=False, x=0, pick_optimizer = None,
            meas_fitter=None):
    if thetas is None:
        n_params = (num_trash*L+num_trasg)*2 if param_encoding else (num_trash*L+num_trash)
        thetas = np.random.uniform(0, 2*np.pi, n_params) # initial parameters without feature encoding
        
    #print("Initial cost: {:.3f}".format(cost_function(thetas, L, num_trash, init_states, shots, vqe, param_encoding, x)))
    
    counts, values, accepted = [], [], []
    def store_intermediate_result(eval_count, parameters, mean, std, ac):
        # counts.append(eval_count)
        values.append(mean)
        accepted.append(ac)

    # Initialize optimizer
    if pick_optimizer == "cobyla":
        optimizer = COBYLA(maxiter=max_iter, tol=0.0001)
    if pick_optimizer == "adam" or pick_optimizer == "ADAM":
        optimizer = qiskit.algorithms.optimizers.ADAM(maxiter=max_iter)
    # optimizer = L_BFGS_B(maxfun=300, maxiter=max_iter)#, factr=10, iprint=- 1, epsilon=1e-08)
    if pick_optimizer == "spsa" or pick_optimizer == None:
        optimizer = SPSA(maxiter=max_iter,
                         #blocking=True,
                         callback=store_intermediate_result,
                         #learning_rate=0.3,
                         #perturbation=0.1
                         ) # recommended from qiskit (first iteraction takes quite long)
                           # to reduce time figure out optimal learning rate and perturbation in advance
    start_time = time.time()
    ret = optimizer.optimize(
                            num_vars=len(thetas),
                            objective_function=(lambda thetas: cost_function(thetas, L, num_trash, init_states, shots, vqe, param_encoding, x, meas_fitter=meas_fitter)),
                            initial_point=thetas
                            )
    print("Time: {:.5f} sec".format(time.time()-start_time))
    print(ret)
    return ret[0], values, accepted

temp2 = np.load("data/QAEAnsatz_scaling/QAEAnsatz_scaling_MPS_script_L-32_results.npz",allow_pickle=True) # 500 iterations, gx = 0.3, chi=100 (but resulting bond dimension of state much smaller)

psi0 = temp2["psi0"].item()
thetas_opt = temp2["thetas_opt"][0]
L = psi0.L

loss_final = cost_function_single(thetas_opt, L, int(np.log(L)/np.log(2)), psi0, shots=1000)

print(f"Result for L = {L}: loss_opt = {loss_final} with optimal parameters {thetas_opt}")

temp2 = np.load("data/QAEAnsatz_scaling_MPS_script_L-32_tracked-losses_results.npz",allow_pickle=True) # 500 iterations, gx = 0.3, chi=100 (but resulting bond dimension of state much smaller)

psi0 = temp2["psis"][1]
thetas_opt = temp2["thetas_opt"][1]
L = psi0.L

loss_final = cost_function_single(thetas_opt, L, int(np.log(L)/np.log(2)), psi0, shots=1000)

print(f"Result for L = {L}: loss_opt = {loss_final} with optimal parameters {thetas_opt}")

