import numpy as np
import qiskit
from qiskit import *
from qiskit import Aer
import pandas as pd
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 itertools
import mitiq
import argparse
import cma
import os
import sys
from qiskit import IBMQ
import pickle
import random
import re
from pprint import pprint

#! ここからmainの実行処理
IBMQ.load_account()
# provider = IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q-utokyo', group='internal', project='hirashi-jst')

print("provider:", provider)

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

#TODO 外部実装
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))
            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)

#! convert 完了
def TrotterEvolveCircuit(dt, nt, init):
    """
    Implements trotter evolution of the Heisenberg hamiltonian using the circuit from https://arxiv.org/pdf/1906.06343.pdf #! 要チェック
    :param tf: time to evolve to #! dt * nt = tf ???
    :param nt: number of trotter steps to use
    :param init: initial state for the trotter evolution. Should be another Qiskit circuit
    
    外部変数: L
    """

    # def get_angles(a, b, c):
    #     return (np.pi/2 - 2*c, 2*a - np.pi/2, np.pi/2 - 2*b)
    def get_angles(a): 
        #! 角度計算, aはalpha, return値タプルの0はtheta, 1はphi, 2はlambd = theta
        return (np.pi/2 - 2*a, 2*a - np.pi/2, np.pi/2 - 2*a)

    def N(cir, qb0, qb1):
        #! fig 4を実装: thetaとphiとlambdはglobal変数
        #! cnotのdepthは3
        cir.rz(-np.pi/2, qb1)
        cir.cnot(qb1, qb0)
        cir.rz(theta, qb0)
        cir.ry(phi, qb1)
        cir.cnot(qb0, qb1)
        cir.ry(lambd, qb1)
        cir.cnot(qb1, qb0)
        cir.rz(np.pi/2, qb0)
        return cir

    #! dtはtrotter step size ← step sizeとは？？？ (default: 0.25)
    theta, phi, lambd = get_angles(-dt/4) #! why divided by 4??? 少なくとも時間間隔ではある
    circ = init

    for i in range(nt): #! ntはTrotterステップ数 (ここではcnotが深さnt * 3かかる)
        # even (odd indices)
        if (L % 2 == 0): #! Lはsystem size
            # UEven
            for i in range(1, L-1, 2): # L for periodic bdy conditions
                circ = N(circ, i, (i+1)%L)
            # UOdd
            for i in range(0, L-1, 2):
                circ = N(circ, i, (i+1)%L)
        else:
            # UEven
            for i in range(1, L, 2):
                circ = N(circ, i, (i+1)%L)
            # UOdd
            for i in range(0, L-1, 2):
                circ = N(circ, i, (i+1)%L)
            # UBdy
            # circ = N(circ, L-1, 0)

    return circ

#! convert完了
def AnsatzCircuit(params: list, p: int) -> QuantumCircuit:
    """
    Implements HVA ansatz using circuits from https://arxiv.org/pdf/1906.06343.pdf #! 要チェック
    #! HVA := Hamiltonian Variational Ansatz
    :param params: parameters to parameterize circuit
    :param p: depth of the ansatz
    
    外部変数: L, p
    """
    circ = QuantumCircuit(L) #! L = system size

    def get_angles(a): #! 回転角度の計算 (肩に乗せるやつ)
        return (np.pi/2 - 2*a, 2*a - np.pi/2, np.pi/2 - 2*a)

    def N(cir, angles, qb0, qb1):
        #! angles = (theta, phi, lambd)
        cir.rz(-np.pi/2, qb1)
        cir.cnot(qb1, qb0)
        cir.rz(angles[0], qb0)
        cir.ry(angles[1], qb1)
        cir.cnot(qb0, qb1)
        cir.ry(angles[2], qb1)
        cir.cnot(qb1, qb0)
        cir.rz(np.pi/2, qb0)
        return cir

    for i in range(p):
        if (L % 2 == 0):
            for j in range(1, L-1, 2): # L for periodic bdy conditions #! periodicなので、Lで割って、0とn-1にまたがる回路が存在する
                circ = N(circ, get_angles(-params[((L-1)*i)+j]/4), j, (j+1)%L)
            for j in range(0, L-1, 2):
                circ = N(circ, get_angles(-params[((L-1)*i)+j]/4), j, (j+1)%L)
        else:
            for j in range(1, L, 2):
                circ = N(circ, get_angles(-params[((L-1)*i)+j]/4), j, (j+1)%L)
            for j in range(0, L-1, 2):
                circ = N(circ, get_angles(-params[((L-1)*i)+j]/4), j, (j+1)%L)
            # circ = N(circ, get_angles(-params[(L*i)+L-1]/4), L-1, 0) # boundary
    return circ

#TODO reverse_bitを適宜挟む
def ReorderBasis(circ):
    """
    #! changing the big endian to little endian
    #! unnecessary function: equal to reverse_bit() method
    Reorders basis so that 0th qubit is on the left side of the tensor product
    :param circ: circuit to reorder, can also be a vector
    """
    if (isinstance(circ, qiskit.circuit.quantumcircuit.QuantumCircuit)):
        for i in range(L//2):
            circ.swap(i, L-i-1)
        return circ
    else:
        perm = np.eye(2**L)
        for i in range(1, 2**L//2):
            perm[:, [i, 2**L-i-1]] = perm[:, [2**L-i-1, i]]
        return perm @ circ

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

#TODO 
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")

    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()

