import sys
import os
from pathlib import Path

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


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