import numpy as np
import pandas as pd
import itertools
import cma
import os
import sys
import argparse
import pickle
import random
import re
from pprint import pprint

import qiskit
from qiskit import *
from qiskit import Aer
from qiskit import IBMQ
from qiskit.providers.aer.noise.noise_model import NoiseModel
from qiskit.test.mock import *
from qiskit.providers.aer import AerSimulator, QasmSimulator
from qiskit.ignis.mitigation.measurement import complete_meas_cal, CompleteMeasFitter
import mitiq

IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q-utokyo', group='internal', project='hirashi-jst')
provider = IBMQ.get_provider(hub='', group='internal', project='hirashi-jst')
print("provider:", provider)

L = 3
p = 2
dt = 1.0
tf = 20
shots = 8192

def TwirlCircuit(circ: str) -> QuantumCircuit:
    """
    そのまま使う: 修正は後回し
    """
    #! qasm ベタ書き
    def apply_pauli(num: int, qb: int) -> str:
        if (num == 0):
            return f'id q[{qb}];\n'
        elif (num == 1):
            return f'x q[{qb}];\n'
        elif (num == 2):
            return f'y q[{qb}];\n'
        else:
            return f'z q[{qb}];\n'

    paulis = [(i,j) for i in range(0,4) for j in range(0,4)]
    paulis.remove((0,0))
    paulis_map = [(0, 1), (3, 2), (3, 3), (1, 1), (1, 0), (2, 3), (2, 2), (2, 1), (2, 0), (1, 3), (1, 2), (3, 0), (3, 1), (0, 2), (0, 3)]

    new_circ = ''
    ops = circ.qasm().splitlines(True) #! 生のqasmコードを持ってきてる: オペレータに分解
    for op in ops:
        if (op[:2] == 'cx'): # can add for cz, etc.
            num = random.randrange(len(paulis)) #! permute paulis
            qbs = re.findall('q\[(.)\]', op)
            new_circ += apply_pauli(paulis[num][0], qbs[0])
            new_circ += apply_pauli(paulis[num][1], qbs[1])
            new_circ += op
            new_circ += apply_pauli(paulis_map[num][0], qbs[0])
            new_circ += apply_pauli(paulis_map[num][1], qbs[1])
        else:
            new_circ += op
    return qiskit.circuit.QuantumCircuit.from_qasm_str(new_circ)

# とりあえずOK
def evolve(alpha: float, q0: Union[int, QuantumRegister], q1: Union[int, QuantumRegister]) -> QuantumCircuit:
    """
    The implementation of Fig. 4 in https://arxiv.org/abs/2112.12654 
    """
    qc = QuantumCircuit(2)
    qc.rz(-np.pi / 2, q1)
    qc.cnot(q1, q0)
    qc.rz(alpha - np.pi / 2, q0)
    qc.ry(np.pi / 2 - alpha, q1)
    qc.cnot(q0, q1)
    qc.ry(alpha - np.pi / 2, q1)
    qc.cnot(q1, q0)
    qc.rz(np.pi / 2, q0)
    return qc

# とりあえずOK
def make_ansatz_circuit(num_qubits: int, ansatz_depth: int, parameters: np.array) -> QuantumCircuit:
    """
    Prepare ansatz circuit
    code reference: https://gitlab.com/QANED/heis_dynamics
    method reference: https://arxiv.org/abs/1906.06343
    
    == AnsatzCircuit(param, p) -> QuantumCircuit
    
    Args:
        parameters: 1d array (for 2d parameters on circuit)
    """
    qc = QuantumCircuit(num_qubits)
    for l in range(ansatz_depth):
        if num_qubits & 1:
            for i in range(0, num_qubits, 2):  # linear condition
                qc.compose(evolve(parameters[l * ansatz_depth + i] / 4), [i, i + 1], inplace = True) #! we do not have to divide the angle by 4
            for i in range(1, num_qubits - 1, 2):  # linear condition
                # ! we do not have to divide the angle by 4
                qc.compose(evolve(parameters[l * ansatz_depth + i] / 4), [i, i + 1], inplace=True)
        else:
            for i in range(0, num_qubits - 1, 2):  # linear condition
                # ! we do not have to divide the angle by 4
                qc.compose(evolve(parameters[l * ansatz_depth + i]), [i, i + 1], inplace=True)
            for i in range(1, num_qubits - 1, 2):  # linear condition
                # ! we do not have to divide the angle by 4
                qc.compose(evolve(parameters[l * ansatz_depth + i]), [i, i + 1], inplace=True)
    return qc

# とりあえずOK
def make_trotter_circuit(num_qubits: int, time_interval: float) -> QuantumCircuit:
    """ 
    Prepare Trotter circuit
    code reference: https://gitlab.com/QANED/heis_dynamics
    method reference: https://arxiv.org/abs/1906.06343
    
    == TrotterEvolveCircuit(dt, nt, init) -> QuantumCircuit
    """
    qc = QuantumCircuit(num_qubits)
    for n in range(trotter_steps): #! time_interval の符号に注意
        if num_qubits & 1:
            for i in range(0, num_qubits, 2):  # linear condition
                qc.compose(evolve(time_interval / 4), [i, i + 1], inplace = True) #! we do not have to divide the angle by 4
            for i in range(1, num_qubits - 1, 2):  # linear condition
                qc.compose(evolve(time_interval / 4), [i, i + 1], inplace = True)  # ! we do not have to divide the angle by 4
        else:
            for i in range(0, num_qubits - 1, 2):  # linear condition
                qc.compose(evolve(time_interval / 4), [i, i + 1], inplace = True) #! we do not have to divide the angle by 4
            for i in range(1, num_qubits - 1, 2):  # linear condition
                qc.compose(evolve(time_interval / 4), [i, i + 1], inplace = True)  # ! we do not have to divide the angle by 4
    return qc

#TODO VTCとは別実装?→ no, 同じ実装に。
def SimulateAndReorder(circ):
    """
    #! execution wrapper
    Executes a circuit using the statevector simulator and reorders basis to match with standard
    """
    backend = Aer.get_backend('statevector_simulator')
    return execute(circ.reverse_bits(), backend).result().get_statevector()

def Simulate(circ):
    """
    #! execution wrapper
    Executes a circuit using the statevector simulator. Doesn't reorder -- which is needed for intermediate steps in the VTC
    """
    backend = Aer.get_backend('statevector_simulator')
    return execute(circ, backend).result().get_statevector()

#TODO
def LoschmidtEchoExecutor(circuits, backend, shots, filter):
    """
    #! 回路を実行
    Returns the expectation value to be mitigated.
    :param circuit: Circuit to run. #! ここでのcircuitsは
    :param backend: backend to run the circuit  on
    :param shots: Number of times to execute the circuit to compute the expectation value.
    :param fitter: measurement error mitigator
    """
    # circuits = [TwirlCircuit(circ) for circ in circuits]
    scale_factors = [1.0, 2.0, 3.0] #! ZNEのノイズスケーリングパラメタ
    folded_circuits = [] #! ZNE用の回路
    for circuit in circuits:
        folded_circuits.append([mitiq.zne.scaling.fold_gates_at_random(circuit, scale) for scale in scale_factors]) #! ここでmitiqを使用
    folded_circuits = list(itertools.chain(*folded_circuits)) #! folded_circuitsを平坦化
    folded_circuits = [TwirlCircuit(circ) for circ in folded_circuits] #! 後からPauli Twirlingを施す!

    print("length of circuit in job", len(folded_circuits))
    
    #! jobを投げる
    job = qiskit.execute(
        experiments=folded_circuits,
        backend=backend,
        optimization_level=0,
        shots=shots
    )
    print("casted job")

    #! fidelity測定用(VTCをしないなら、ここはtomographyで良い)
    c = ['1','1','0'] #! これをpermutationする
    # c = [str((1 + (-1)**(i+1)) // 2) for i in range(L)]
    c = ''.join(c)[::-1] #! endianを反転 (big endianへ)
    res = job.result()
    if (filter is not None): #! QREM
        res = filter.apply(res)

    print("retrieved job")
        
    all_counts = [job.result().get_counts(i) for i in range(len(folded_circuits))]
    expectation_values = []
    for counts in all_counts:
        total_allowed_shots = [counts.get(''.join(p)) for p in set(itertools.permutations(c))] #! ここでcをpermutationしている
        total_allowed_shots = sum([0 if x is None else x for x in total_allowed_shots])
        if counts.get(c) is None:
            expectation_values.append(0)
        else:
            expectation_values.append(counts.get(c)/total_allowed_shots)
    # expectation_values = [counts.get(c) / shots for counts in all_counts]

    zero_noise_values = []
    if isinstance(backend, qiskit.providers.aer.backends.qasm_simulator.QasmSimulator): # exact_sim
        for i in range(len(circuits)):
            zero_noise_values.append(np.mean(expectation_values[i*len(scale_factors):(i+1)*len(scale_factors)]))
    else: #device_sim, real_device
        fac = mitiq.zne.inference.LinearFactory(scale_factors)
        for i in range(len(circuits)):
            zero_noise_values.append(fac.extrapolate(scale_factors, 
            expectation_values[i*len(scale_factors):(i+1)*len(scale_factors)]))

    print("zero_noise_values")
    pprint(zero_noise_values)
    print()
    return zero_noise_values

#TODO 
def LoschmidtEchoCircuit(params, U_v, U_trot, init, p):
    """
    #! 回路を作成
    Cost function using the Loschmidt Echo. Just using statevectors currently -- can rewrite using shots
    :param params: parameters new variational circuit that represents U_trot U_v | init >. Need dagger for cost function
    :param U_v: variational circuit that stores the state before the trotter step
    :param U_trot: trotter step
    :param init: initial state
    :param p: number of ansatz steps
    """
    U_v_prime = AnsatzCircuit(params, p)
    circ = init + U_v + U_trot + U_v_prime.inverse()
    circ.measure_all()
    return circ

def LoschmidtEcho(params, U_v, U_trot, init, p, backend, shots, filter):
    """
    #! 実行パート
    """
    circs = []
    for param in params:
        circs.append(LoschmidtEchoCircuit(param, U_v, U_trot, init, p)) #! 回路を作成
    print("length of circuits without zne:", len(circs))
    res = LoschmidtEchoExecutor(circs, backend, shots, filter) #! 回路を実行
    return abs(1 - np.array(res))

def LoschmidtEchoExact(params, U_v, U_trot, init, p):
    """
    #! unused function
    """
    U_v_prime = AnsatzCircuit(params, p)
    circ = init + U_v + U_trot + U_v_prime.inverse()

    circ_vec = Simulate(circ)
    init_vec = Simulate(init)
    return 1 - abs(np.conj(circ_vec) @ init_vec)**2

def CMAES(U_v, U_trot, init, p, backend, shots, filter):
    """
    #! 実行 + 最適化パート
    """
    init_params = np.random.uniform(0, 2*np.pi, (L-1)*p)
    es = cma.CMAEvolutionStrategy(init_params, np.pi/2)
    es.opts.set({'ftarget':5e-3, 'maxiter':1000})
    # es = pickle.load(open(f'./results_{L}/optimizer_dump', 'rb'))
    while not es.stop(): #! 最適化パート
        # solutions = es.ask(25) # ! 25 = number of returned solutions
        solutions = es.ask(10)
        print("solutions")
        pprint(solutions)
        es.tell(solutions, LoschmidtEcho(solutions, U_v, U_trot, init, p, backend, shots, filter)) #! 実行パート
        # es.tell(solutions, LoschmidtEchoExact(solutions, U_v, U_trot, init, p)) #! 実行パート
        es.disp()
        open(f'./results_{L}/optimizer_dump', 'wb').write(es.pickle_dumps())
    return es.result_pretty()

def VTC(tf, dt, p, init, backend, shots, filter):
    """
    #! tf: 総経過時間
    #! dt: trotter step size: 時間間隔
    #! p: ansatzのステップ数
    """

    VTCParamList = [np.zeros((L-1)*p)] #! デフォルトのパラメタ(初期値)
    VTCStepList = [SimulateAndReorder(init.copy())] #! type: List[Statevector]
    # TrotterFixStepList = [init]
    TimeStep = [0]

    if (os.path.exists(f'./results_{L}/VTD_params_{tf}_{L}_{p}_{dt}_{shots}.csv')): #! 2巡目からこっち
        VTCParamList = pd.read_csv(f'./results_{L}/VTD_params_{tf}_{L}_{p}_{dt}_{shots}.csv', index_col=0)
        VTCStepList = pd.read_csv(f'./results_{L}/VTD_results_{tf}_{L}_{p}_{dt}_{shots}.csv', index_col=0)
        temp = VTCParamList.iloc[-1]
        print(temp, "th time interval")
        U_v = AnsatzCircuit(temp, p)
    else: #! 最初はこっちに入る
        VTCParamList = pd.DataFrame(np.array(VTCParamList), index=np.array(TimeStep))
        VTCStepList = pd.DataFrame(np.array(VTCStepList), index=np.array(TimeStep))
        print("0 th time interval")
        print()
        U_v = QuantumCircuit(L)

    ts = VTCParamList.index 
    #! 時間間隔
    U_trot = TrotterEvolveCircuit(dt, p, QuantumCircuit(L)) #! Trotter分解のunitaryを作る

    print()
    print("start CMAES")
    print()
    res = CMAES(U_v, U_trot, init, p, backend, shots, filter) #! ここでプロセスを実行！！！！
    print()
    print("res")
    pprint(res)

    #! 新しいループ結果を追加し、tsを更新
    res = res.xbest  # ! best solution evaluated
    print("res.xbest")
    pprint(res)
    VTCParamList.loc[ts[-1]+(dt*p)] = np.array(res) 
    VTCStepList.loc[ts[-1]+(dt*p)] = np.array(SimulateAndReorder(init + AnsatzCircuit(res, p)))
    ts = VTCParamList.index

    # VTCParamList = pd.DataFrame(np.array(VTCParamList), index=np.array(TimeStep))
    # VTCStepList = pd.DataFrame(np.array(VTCStepList), index=np.array(TimeStep))

    #! csvファイルを更新
    VTCParamList.to_csv(f'./results_{L}/VTD_params_{tf}_{L}_{p}_{dt}_{shots}.csv')
    VTCStepList.to_csv(f'./results_{L}/VTD_results_{tf}_{L}_{p}_{dt}_{shots}.csv')

    if (ts[-1] >= tf):
        return
    else:
        print("next step")
        VTC(tf, dt, p, init, backend, shots, filter)

#! ここからQREM回路
qr = QuantumRegister(L)
meas_calibs, state_labels = complete_meas_cal(qr=qr, circlabel='mcal')

# device_backend = FakeJakarta()
# device_sim = AerSimulator.from_backend(device_backend)
real_device = provider.get_backend('ibmq_jakarta')
noise_model = NoiseModel.from_backend(real_device)

device_sim = QasmSimulator(method='statevector', noise_model=noise_model)
exact_sim = Aer.get_backend('qasm_simulator') # QasmSimulator(method='statevector')

t_qc = transpile(meas_calibs)
qobj = assemble(t_qc, shots=8192)
# cal_results = real_device.run(qobj, shots=8192).result()
cal_results = device_sim.run(qobj, shots=8192).result()
meas_fitter = CompleteMeasFitter(cal_results, state_labels, circlabel='mcal')

print("qrem done")

# np.around(meas_fitter.cal_matrix, decimals=2)

init = QuantumCircuit(L)
# c = [str((1 + (-1)**(i+1)) // 2) for i in range(L)]
c = ['1','1','0'] #! なぜinitial stateが110なの？？？？？？？ もしかしてopen science prizeを意識???
#! けどループでこのプログラムが実行されるたびにここが|110>だとおかしくないか？
for q in range(len(c)):
    if (c[q] == '1'):
        init.x(q)
#! ここまでQREM回路

nt = int(np.ceil(tf / (dt * p)))
# f = open(f'./results_{L}/logging.txt', 'a')
# sys.stdout = f
#! tf: シミュレーションの(経過)時間
#! dt: trotter分解のステップ数
#! p: ansatzのステップ数 (論文中のL)
# VTC(tf, dt, p, init, real_device, shots, meas_fitter.filter) #! mainの処理

print("vtc start!!!! \n\n\n")
VTC(tf, dt, p, init, device_sim, shots, meas_fitter.filter) #! mainの処理
# f.close()

