
import numpy as np

_SQ2 = np.sqrt(2.0)

def rx_mat(theta):
    c, s = np.cos(theta/2), np.sin(theta/2)
    return np.array([[c, -1j*s],[-1j*s, c]], dtype=complex)

def ry_mat(theta):
    c, s = np.cos(theta/2), np.sin(theta/2)
    return np.array([[c, -s],[s, c]], dtype=complex)

def rz_mat(theta):
    return np.array([[np.exp(-1j*theta/2), 0],
                     [0, np.exp(1j*theta/2)]], dtype=complex)

def u1_mat(lam):
    return np.array([[1,0],[0,np.exp(1j*lam)]], dtype=complex)

def u2_mat(phi, lam):
    return np.array([[1, -np.exp(1j*lam)],
                     [np.exp(1j*phi), np.exp(1j*(phi+lam))]], dtype=complex) / _SQ2

def u3_mat(theta, phi, lam):
    c, s = np.cos(theta/2), np.sin(theta/2)
    return np.array([[c, -np.exp(1j*lam)*s],
                     [np.exp(1j*phi)*s, np.exp(1j*(phi+lam))*c]], dtype=complex)

def apply_1q(state, gate, qubit, n):
    """Fast single-qubit gate — tensor reshape, no Kronecker."""
    psi = state.reshape([2]*n)
    psi = np.tensordot(gate, psi, axes=[[1], [qubit]])
    order = list(range(1, qubit+1)) + [0] + list(range(qubit+1, n))
    return np.transpose(psi, order).reshape(-1)

def apply_2q(state, gate, q0, q1, n):
    """Fast two-qubit gate — einsum contraction."""
    psi = state.reshape([2]*n)
    G = gate.reshape(2, 2, 2, 2)
    in_idx  = list(range(n))
    out_idx = list(range(n))
    ni, nj  = n, n+1
    out_idx[q0] = ni
    out_idx[q1] = nj
    return np.einsum(G, [ni, nj, q0, q1], psi, in_idx, out_idx).reshape(-1)

def apply_3q(state, gate, q0, q1, q2, n):
    """Fast three-qubit gate — einsum contraction."""
    psi = state.reshape([2]*n)
    G   = gate.reshape(2,2,2,2,2,2)
    in_idx  = list(range(n))
    out_idx = list(range(n))
    ni, nj, nk = n, n+1, n+2
    out_idx[q0]=ni; out_idx[q1]=nj; out_idx[q2]=nk
    return np.einsum(G, [ni,nj,nk,q0,q1,q2], psi, in_idx, out_idx).reshape(-1)


def rx_mat(t):
    import numpy as np
    c, s = np.cos(t/2), np.sin(t/2)
    return np.array([[c, -1j*s], [-1j*s, c]], dtype=np.complex64)

def ry_mat(t):
    import numpy as np
    c, s = np.cos(t/2), np.sin(t/2)
    return np.array([[c, -s], [s, c]], dtype=np.complex64)

def rz_mat(t):
    import numpy as np
    return np.array([[np.exp(-1j*t/2), 0], [0, np.exp(1j*t/2)]], dtype=np.complex64)

def u1_mat(lam):
    import numpy as np
    return np.array([[1, 0], [0, np.exp(1j*lam)]], dtype=np.complex64)

def u2_mat(phi, lam):
    import numpy as np
    return np.array([[1, -np.exp(1j*lam)],
                     [np.exp(1j*phi), np.exp(1j*(phi+lam))]],
                    dtype=np.complex64) / np.sqrt(2)

def u3_mat(theta, phi, lam):
    import numpy as np
    c, s = np.cos(theta/2), np.sin(theta/2)
    return np.array([[c, -np.exp(1j*lam)*s],
                     [np.exp(1j*phi)*s, np.exp(1j*(phi+lam))*c]],
                    dtype=np.complex64)

def apply_1q(sv, mat, qubit, n):
    import numpy as np
    psi = sv.reshape([2]*n)
    psi = np.tensordot(mat, psi, axes=[[1], [qubit]])
    order = list(range(1, qubit+1)) + [0] + list(range(qubit+1, n))
    return np.transpose(psi, order).reshape(-1)

def apply_2q(sv, mat, q0, q1, n):
    import numpy as np
    psi = sv.reshape([2]*n)
    G   = mat.reshape(2, 2, 2, 2)
    ii  = list(range(n))
    oi  = list(range(n))
    ni, nj = n, n+1
    oi[q0] = ni; oi[q1] = nj
    return np.einsum(G, [ni, nj, q0, q1], psi, ii, oi).reshape(-1)

def apply_3q(sv, mat, q0, q1, q2, n):
    import numpy as np
    psi = sv.reshape([2]*n)
    G   = mat.reshape(2, 2, 2, 2, 2, 2)
    ii  = list(range(n))
    oi  = list(range(n))
    na, nb, nc = n, n+1, n+2
    oi[q0] = na; oi[q1] = nb; oi[q2] = nc
    return np.einsum(G, [na, nb, nc, q0, q1, q2], psi, ii, oi).reshape(-1)
