import sys
sys.path.insert(1, '..') # sets the import path to the parent folder

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

from tenpy.networks.mps import MPS
from tenpy.models.hubbard import BoseHubbardChain
from tenpy.algorithms import dmrg
from tenpy.linalg import np_conserved

def DMRG_EBH(L, V, t_list, chi_max=30, bc_MPS='infinite'):

    model_params = dict(n_max=1, filling=0.5, bc_MPS=bc_MPS, t=t_list,
                        L=L, V=V, mu=0, conserve='N')
    M = BoseHubbardChain(model_params)
        
    vector=[]
    for i in range(M.lat.N_sites):
        if i%2:
            vector.append(1)
        else:
            vector.append(0)

    psi = MPS.from_product_state(M.lat.mps_sites(), vector, bc=M.lat.bc_MPS)    
        
    dmrg_params = {                                                                                             
        'mixer': True,                                                                                          
        'trunc_params': {                                                                                       
        'chi_max': chi_max,                                                                                                                                                                    
        },                                                                                                      
        'max_E_err': 1.e-16,                                                                                    
        #'verbose': 0
    }
    
    info = dmrg.run(psi, M, dmrg_params)
    
    return info['E'], psi

print(qiskit.__version__, np.__version__)

#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(method="matrix_product_state")
#backend = StatevectorSimulator()

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

filename = "data/QAEAnsatz_scaling_MPS"

V = 1
deltat=1
chi = 2**3
print("bond dimension : ", chi)

t_list = np.ones(L-1)
for i in range(len(t_list)):
    t_list[i] -= deltat*(-1)**i
E0, psi0 = DMRG_EBH(L, V, t_list, chi_max=chi, bc_MPS='finite')

thetas = np.random.rand(num_trash*(L+1))
QAE_circ = QAEAnsatz(num_qubits = L, num_trash_qubits= num_trash, trash_qubits_idxs = list(range(num_trash)), measure_trash=False).assign_parameters(thetas)

QAE_circ.save_matrix_product_state(label="my_mps")

result = backend.run(QAE_circ).result()

#result.data(0)["my_mps"]

[(_[0].shape, _[1].shape) for _ in result.data(0)["my_mps"][0]]

[_.shape for _ in result.data(0)["my_mps"][1]]

# 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])

[(_[0].shape, _[1].shape) for _ in A_list]

S_list = [psi0.get_SR(i) for i in range(L-1)] # skip trivial last bond; hast to be of size L-1

state = (A_list, S_list)

initcirc = QuantumCircuit(QuantumRegister(L,"q")) #,ClassicalRegister(num_trash, 'c')

initcirc.set_matrix_product_state(state)
initcirc.draw()

plt.plot(psi0.expectation_value("N"))

i=0

qi = qiskit.utils.QuantumInstance(backend=backend, # , seed_simulator=seed, seed_transpiler=seed
                                  shots = 1000,
                                  #coupling_map=coupling_map,
                                  #noise_model=noise_model, # comment out on real device execution
                                  #measurement_error_mitigation_cls= CompleteMeasFitter, 
                                  #cals_matrix_refresh_period=30  #How often to refresh the calibration matrix in measurement mitigation. in minutes
                                 )

QZ = np.zeros(L)

for i in range(L):
    meas_outcome = ~StateFn(QNKron(L,Z,I,i)) @ StateFn(initcirc)
    QZ[i] = CircuitSampler(qi).convert(meas_outcome).eval()

plt.plot(QZ)

