import numpy as np
import scipy.sparse as sp
import math
import random
import matplotlib.pyplot as plt
from scipy.special import jn
from scipy.sparse import identity, csr_matrix, kron, diags, eye
from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.circuit.library import MCXGate, MCPhaseGate, RXGate, CRXGate, QFTGate, StatePreparation, PauliEvolutionGate, RZGate
from qiskit.quantum_info import SparsePauliOp, Statevector, Operator, Pauli
from scipy.linalg import expm
# from tools import *
from qiskit.qasm3 import dumps  # QASM 3 exporter
from qiskit.qasm3 import loads
from qiskit.circuit.library import QFT
from qiskit.primitives import StatevectorEstimator
from qiskit import transpile
from qiskit_addon_aqc_tensor.simulation import tensornetwork_from_circuit, apply_circuit_to_state, compute_overlap
from qiskit_aer import AerSimulator
from qiskit_ionq import IonQProvider
import os
my_api_key = os.getenv("IONQ_API_KEY")
from qiskit_ibm_runtime import QiskitRuntimeService, EstimatorV2 as Estimator




# provider = IonQProvider()

api_token = "SgUkiDq1r2bVEadyiUfvtuxQ03Qci6UW"
provider = IonQProvider(api_token)
ionq_backend = provider.get_backend("simulator")




simulator_settings = AerSimulator(
    method="matrix_product_state",
    matrix_product_state_max_bond_dimension=100,
)

def Wj(j, theta, lam, name='Wj', xgate=False):    
    if not xgate:
        name = f' $W_{j}$ '
    qc=QuantumCircuit(j, name=name)

    if j > 1:
        qc.cx(j-1, range(j-1))
    if lam != 0:
        qc.p(lam, j-1)
    qc.h(j-1)
    if xgate:
        qc.x(range(j-1))

    # the multicontrolled rz gate 
    # it will be decomposed in qiskit
    if j > 1:
        qc.mcrz(theta, range(j-1), j-1)
    else:
        qc.rz(theta, j-1)
    
    if xgate:
        qc.x(range(j-1))
    qc.h(j-1)
    if lam != 0:
        qc.p(-lam, j-1)
    if j > 1:
        qc.cx(j-1, range(j-1))

    return qc

def Wj_block(j, n, ctrl_state, theta, lam, name='Wj_block', xgate=False):    
    if not xgate:
        name = f' $W_{j}_block$ '
    qc=QuantumCircuit(n + j, name=name)

    if j > 1:
        qc.cx(n + j-1, range(n, n+j-1))
    if lam != 0:
        qc.p(lam, n + j -1)
    qc.h(n + j -1)
    
    if xgate and j>1:
        if isinstance(xgate, (list, tuple)):  # selective application
            for idx, flag in enumerate(xgate):
                if flag:   # only apply where flag == 1
                    qc.x(n + idx)
        elif xgate is True:  # apply to all
            qc.x(range(n, n+j-1))

    # the multicontrolled rz gate 
    # it will be decomposed in qiskit
    if j > 1:
        mcrz = RZGate(theta).control(len(ctrl_state) + j-1, ctrl_state = "1"*(j-1)+ctrl_state)
        qc.append(mcrz, range(0, n + j))
    else:
        mcrz = RZGate(theta).control(len(ctrl_state), ctrl_state = ctrl_state)
        qc.append(mcrz, range(0, n+j))
    
    if xgate and j>1:
        if isinstance(xgate, (list, tuple)):  # selective application
            for idx, flag in enumerate(xgate):
                if flag:   # only apply where flag == 1
                    qc.x(n + idx)
        elif xgate is True:  # apply to all
            qc.x(range(n, n+j-1))
            
    qc.h(n+ j-1)
    if lam != 0:
        qc.p(-lam, n + j-1)
    if j > 1:
        qc.cx(n + j-1, range(n, n +j-1))

    return qc.to_gate(label=name)

def V1(nx, dt, name = "V1"):
    n = int(np.ceil(np.log2(nx)))

    derivatives = QuantumRegister(2*n)
    blocks = QuantumRegister(2)

    qc = QuantumCircuit(derivatives, blocks)
    
    W1 = Wj_block(2, n, "0"*n, -dt , 0, xgate=True)
    qc.append(W1, list(derivatives[0:n])+list(blocks[:]))

    # qc.barrier()

    W2 = Wj_block(3, n-1, "1"*(n-1), dt , 0, xgate=[0,1])
    qc.append(W2, list(derivatives[1:n])+[derivatives[0]]+list(blocks[:]))

    # qc.barrier()

    W3 = Wj_block(1, n+1, "0"*(n+1), dt , 0, xgate=False)
    qc.append(W3, list(derivatives[n:2*n])+list(blocks[:]))

    # qc.barrier()

    W4 = Wj_block(2, n, "0"+"1"*(n-1), -dt , 0, xgate=False)
    qc.append(W4, list(derivatives[n+1:2*n]) + [blocks[0]] + [derivatives[n]] + [blocks[1]])

    return qc

def V2(nx, dt, name = "V2"):
    n = int(np.ceil(np.log2(nx)))

    derivatives = QuantumRegister(2*n)
    blocks = QuantumRegister(2)

    qc = QuantumCircuit(derivatives, blocks)
    
    W1 = Wj_block(2, 0, "", -2*dt , -np.pi/2, xgate=True)
    qc.append(W1, list(blocks[:]))

    # qc.barrier()

    for j in range(1, n+1):
        W2 = Wj_block(2+j, 0, "", 2*dt , -np.pi/2, xgate=[1]*(j-1)+[0,1])
        qc.append(W2, list(derivatives[0:j])+list(blocks[:]))

    # qc.barrier()

    W3 = Wj_block(2, n, "0"*n, -dt , -np.pi/2, xgate=True)
    qc.append(W3, list(derivatives[0:n])+list(blocks[:]))

    # qc.barrier()

    W4 = Wj_block(2, n, "1"*n, 2*dt , -np.pi/2, xgate=True)
    qc.append(W4, list(derivatives[0:n])+list(blocks[:]))

    # qc.barrier()

    W5 = Wj_block(3, n-1, "1"*(n-1), dt , -np.pi/2, xgate=[0,1])
    qc.append(W5, list(derivatives[1:n])+[derivatives[0]]+list(blocks[:]))

    # qc.barrier()

    W6 = Wj_block(1, 1, "0", 2*dt , -np.pi/2, xgate=False)
    qc.append(W6, list(blocks[:]))

    # qc.barrier()

    for j in range(1, n+1):
        W7 = Wj_block(1+j, 1, "0", -2*dt , -np.pi/2, xgate=[1]*(j-1))
        qc.append(W7, [blocks[0]]+list(derivatives[n:n+j])+[blocks[1]])

        # qc.barrier()

    W8 = Wj_block(1, n+1, "0"*(n+1), dt , -np.pi/2, xgate=False)
    qc.append(W8, list(derivatives[n:2*n])+list(blocks[:]))

    # qc.barrier()

    W9 = Wj_block(1, n+1, "0"+"1"*(n), -2*dt , -np.pi/2, xgate=False)
    qc.append(W9, list(derivatives[n:2*n])+list(blocks[:]))

    # qc.barrier()

    W10 = Wj_block(2, n, "0"+"1"*(n-1), -dt , -np.pi/2, xgate=False)
    qc.append(W10, list(derivatives[n+1:2*n]) + [blocks[0]] + [derivatives[n]] + [blocks[1]])

    # qc.barrier()
    
    return qc

def schro(nx, na, R, dt, initial_state, steps):

    nq = int(np.ceil(np.log2(nx)))

    # warped phase transformation
    dp = 2 * R * np.pi / 2**na
    p = np.arange(- R * np.pi, R * np.pi, step=dp)
    fp = np.exp(-np.abs(p))
    norm1 = np.linalg.norm(fp[2**(na-1):]) # norm of p>=0

    # construct quantum circuit
    system = QuantumRegister(2*nq+2, name='system')
    ancilla = QuantumRegister(na, name='ancilla')
    qc = QuantumCircuit(system, ancilla)

    # initialization
    prep = StatePreparation(initial_state)
    anc_prep = StatePreparation(fp / np.linalg.norm(fp))

    qc.append(prep, system)
    # qc.append(anc_prep, ancilla)
    qc.initialize(fp / np.linalg.norm(fp), ancilla)


    # QFT
    qc.append(QFTGate(na), ancilla)
    qc.x(ancilla[-1])

    A1 = V1(nx, dt, name = "V1").to_gate()
    A2 = V2(nx, dt, name = "V2")

    
    # Hamiltonian simulation for Nt steps
    for i in range(steps):
        # circuit for one step
        for j in range(na):
            # repeat controlled H1 for 2**j times
            qc.append(A1.control().repeat(2**j), [ancilla[j]] + system[:])

        # qc.append(A1.inverse().control(ctrl_state = "0").repeat(2**(na-1)), [ancilla[na-1]] + system[:])
        qc.append(A1.inverse().repeat(2**(na-1)), system[:])
        qc.append(A2, system[:])

    # rearrange eta
    qc.x(ancilla[-1])
    qc.append(QFTGate(na).inverse(), ancilla)

    return qc



def circ_for_magnitude(field, x, y, nx, na, R, dt, initial_state, steps):

    qc  = schro(nx, na, R, dt, initial_state, steps)
    naimark = QuantumRegister(1, name='Naimark')
    qc.add_register(naimark)

    if field == 'Ez':
        index = nx * y + x
    elif field == 'Hx':
        index = 2*nx*nx + nx * y + x
    else:
        index = 3*nx*nx + nx * y + x

    index_bin = format(index, f'0{qc.num_qubits-2}b')
    ctrl_state = '1' + index_bin
    ctrl_qubits = qc.qubits[:-1]
    qc.mcx(ctrl_qubits, naimark[0], ctrl_state=ctrl_state)

    return qc

def circuits_for_sign(field, x, y, nx, na, dt, R, initial_state, steps, xref, yref, field_ref = 'Ez'):    
    qc  = schro(nx, na, R, dt, initial_state, steps)

    naimark = QuantumRegister(1, name='Naimark')
    qc.add_register(naimark)

    if field == 'Ez':
        index = nx * y + x
    elif field == 'Hx':
        index = 2*nx*nx + nx * y + x
    else:
        index = 3*nx*nx + nx * y + x

    if field_ref == 'Ez':
        index_ref = nx * yref + xref
    elif field_ref == 'Hx':
        index_ref = 2*nx*nx + nx * yref + xref
    else:
        index_ref = 3*nx*nx + nx * yref + xref

    index_bin = [(index >> i) & 1 for i in range(qc.num_qubits-2)]
    index_ref_bin = [(index_ref >> i) & 1 for i in range(qc.num_qubits-2)]
    index_bin.append(1)
    index_ref_bin.append(1)

    #Convert reference bitstring to 00000 
    for i, bit in enumerate(index_ref_bin):
        if bit == 1:
            qc.x(i)

    d_bits = [b ^ r for b, r in zip(index_ref_bin, index_bin)]
    control = d_bits.index(1)

    #Convert the other bitstring to 0001000
    for target, bit in enumerate(d_bits):
        if bit == 1 and target != control:
            qc.cx(control, target)
    qc.h(control)

    ctrl_state_sum = '0'*(qc.num_qubits-1)
    ctrl_state_diff = '0'*(qc.num_qubits-1-control-1)+'1'+'0'*(control)
    
    qcdiff = qc.copy()

    ctrl_qubits = qc.qubits[:-1]

    qc.mcx(ctrl_qubits, naimark[0], ctrl_state=ctrl_state_sum)
    qcdiff.mcx(ctrl_qubits, naimark[0], ctrl_state=ctrl_state_diff)

    return qc, qcdiff

def get_absolute_field_value(qc, nq, na, offset, norm):

    pauli_label = 'Z'+'I'*(2*nq+2+na)
    observable  = SparsePauliOp(Pauli(pauli_label))
    ########################################################################################
    estimator = StatevectorEstimator() 

    # === Run Estimator (no parameters needed) ===
    pub = (qc, observable)
    job = estimator.run([pub])
    result = job.result()[0]
    z_exp = result.data.evs.item()
    #########################################################################################
    # === Compute projector expectation ===
    pi_expect = (1 - z_exp) / 2

    Absolute_value = norm*np.sqrt(pi_expect)-offset
    
    return Absolute_value

def get_relative_sign(qc, qcdiff, nq, na):

    pauli_label = 'Z'+'I'*(2*nq+2+na)
    observable  = SparsePauliOp(Pauli(pauli_label))
    ########################################################################################
    estimator = StatevectorEstimator() 

    # === Run Estimator ===
    pub = (qc, observable)
    job = estimator.run([pub])
    result = job.result()[0]
    z_exp = result.data.evs.item()

    pub_diff = (qcdiff, observable)
    job_diff = estimator.run([pub_diff])
    result_diff = job_diff.result()[0]
    z_exp_diff = result_diff.data.evs.item()
    #########################################################################################
    # === Compute projector expectation ===
    pi_expect_sum = (1 - z_exp) / 2
    pi_expect_diff = (1 - z_exp_diff) / 2

    relative_sign = 'same' if pi_expect_sum >= pi_expect_diff else 'different'
    
    return relative_sign

def Eref_value(nx, nq, R, dt, na, steps, xref, yref, field_ref = 'Ez'):
    if steps < 31:
        offset = 1
    else :
        offset = 0.15
    deltastate = np.zeros(4*nx*nx)
    # deltastate[nx*nx//2+nx//2:nx*nx//2+nx//2+1] = 1
    deltastate[nx*yref+xref] = 1
    deltastate[0:nx*nx] = deltastate[0:nx*nx] + offset
    norm1 = np.linalg.norm(deltastate)
    initial_state = deltastate/norm1

    dp = 2 * R * np.pi / 2**na
    p = np.arange(- R * np.pi, R * np.pi, step=dp)
    fp = np.exp(-np.abs(p))
    norm2 = np.linalg.norm(fp)
    norm = norm1 * norm2

    qc = circ_for_magnitude(field_ref, xref, yref, nx, na, R, dt, initial_state, steps)

    Ezref = get_absolute_field_value(qc, nq, na, offset, norm)

    return Ezref


def transpile_circ(circ, basis_gates=None):
    """
    Transpile the circuit to the specified basis gates.
    """
    if basis_gates is None:
        basis_gates = ['z', 'y', 'x', 'sdg', 's', 'h', 'rz', 'ry', 'rx', 'ecr', 'cz', 'cx']
    
    transpiled_circ = transpile(circ, basis_gates=basis_gates)
    return transpiled_circ

def compute_fidelity(circ1, circ2):
    
    circ_1 = tensornetwork_from_circuit(transpile_circ(circ1), simulator_settings)
    circ_2 = tensornetwork_from_circuit(transpile_circ(circ2), simulator_settings)
    fidelity = abs(compute_overlap(circ_1, circ_2))**2
    
    return fidelity

# def create_impulse_state(grid_dims, impulse_pos):
#     """
#     Creates an initial state vector with a single delta impulse at a specified grid position.

#     The 2D grid is flattened into a 1D vector in row-major order, and this
#     vector is then padded to match the full simulation state space size (4x).

#     Args:
#         grid_dims (tuple): A tuple (width, height) defining the simulation grid dimensions.
#                            For your original code, this would be (nx, nx).
#         impulse_pos (tuple): A tuple (x, y) for the position of the impulse.
#                              Coordinates are 0-indexed.

#     Returns:
#         numpy.ndarray: The full, padded initial state vector with a single 1.
    
#     Raises:
#         ValueError: If the impulse position is outside the grid dimensions.
#     """
#     grid_width, grid_height = grid_dims
#     impulse_x, impulse_y = impulse_pos

#     # --- Input Validation ---
#     # Ensure the requested impulse position is actually on the grid.
#     if not (0 <= impulse_x < grid_width and 0 <= impulse_y < grid_height):
#         raise ValueError(f"Impulse position ({impulse_x}, {impulse_y}) is outside the "
#                          f"grid dimensions ({grid_width}x{grid_height}).")

#     # --- 1. Calculate the 1D Array Index ---
#     # Convert the (x, y) coordinate to a single index in a flattened 1D array.
#     # The formula for row-major order is: index = y_coord * width + x_coord
#     flat_index = impulse_y * grid_width + impulse_x

#     # --- 2. Create the Full, Padded State Vector ---
#     grid_size = grid_width * grid_height
#     total_size = 4 * grid_size  # The simulation space is 4x the grid size.
#     initial_state = np.zeros(total_size)

#     # --- 3. Set the Delta Impulse ---
#     initial_state[flat_index] = 1

#     return initial_state