import numpy as np
from numba import njit, prange
import scipy.sparse as sp
from utils import reshape_columnwise

def discretize_potential_sparse(Vgrid, spin_multiplicity=2):
    """
    Discretize a potential operator as COO triplets (indices + values).
    The potential is diagonal in real space.

    Parameters
    ----------
    Vgrid : array_like
        Potential values on the discretization grid.
    spin_multiplicity : int, optional
        Spin multiplicity (default 2 for spin-1/2).

    Returns
    -------
    H : scipy.sparse.coo_matrix
    """
    Vgrid = np.asarray(Vgrid).T.flatten()

    # repeat each potential value spin_multiplicity times
    values = np.repeat(Vgrid, spin_multiplicity)

    return sp.diags(values)


sx = np.array([[0,1],[1,0]])
sy = np.array([[0,-1j],[1j,0]])
sz = np.array([[1,0],[0,-1]])
I2 = np.eye(2)

@njit
def onsite_potential(V, omega_L, n):
    HZ = 0.5 * omega_L * (n[0]*sx + n[1]*sy + n[2]*sz)
    return V * I2 + HZ

@njit
def hopping_x(t, tR):
    return -t * I2 - 1j * tR * sy

@njit
def hopping_y(t, tR):
    return -t * I2 + 1j * tR * sx

@njit(parallel=True)
def build_hamiltonian_parallel(Nx, Ny, a, meff, alpha, omega_L, n, Vgrid):
    """
    Build the full Hamiltonian for a 2D grid with Rashba SOI and onsite potential   
    in sparse COO format using parallel loops.

    Parameters
    ----------
    Nx, Ny : int
        Number of grid points in x and y.
    a : float
        Lattice constant (m).
    meff : float
        Effective mass (in units of rad/m^2 * s).
    alpha : float
        Rashba SOI parameter (m/s).
    omega_L : float
        Larmor frequency (Hz).
    n : array_like
        Unit vector for magnetic field direction.
    Vgrid : array_like
        Onsite potential values on the Nx by Ny grid.
    
    Returns
    -------
    rows, cols, vals : arrays
        COO-format triplets for sparse matrix construction.
    """

    t = 1 / (2.0 * meff * a**2) # hbar = 1, units of frequency
    tR = alpha / (2.0 * a) # hbar = 1, units of frequency

    # Each site has:
    # - onsite: 2x2 = 4 entries
    # - x hopping: 2x2 = 4 (if ix+1<Nx) * 2
    # - y hopping: 2x2 = 4 (if iy+1<Ny) * 2
    max_entries_per_site = 4 + 8 + 8
    total_nnz = max_entries_per_site * Nx * Ny
    rows = np.zeros(total_nnz, dtype=np.int32)
    cols = np.zeros(total_nnz, dtype=np.int32)
    vals = np.zeros(total_nnz, dtype=np.complex128)

    def idx(ix, iy, spin):
        return 2*(iy*Nx + ix) + spin

    def fill_site(ix, iy):
        local_rows = np.zeros(max_entries_per_site, dtype=np.int32)
        local_cols = np.zeros(max_entries_per_site, dtype=np.int32)
        local_vals = np.zeros(max_entries_per_site, dtype=np.complex128)
        count = 0

        V = Vgrid[iy, ix]
        Hons = onsite_potential(V, omega_L, n)

        # onsite
        for s1 in range(2):
            for s2 in range(2):
                local_rows[count] = idx(ix,iy,s1)
                local_cols[count] = idx(ix,iy,s2)
                local_vals[count] = Hons[s1,s2]
                count += 1

        # x hopping
        if ix+1 < Nx:
            Hx = hopping_x(t, tR)
            for s1 in range(2):
                for s2 in range(2):
                    local_rows[count] = idx(ix,iy,s1)
                    local_cols[count] = idx(ix+1,iy,s2)
                    local_vals[count] = Hx[s1,s2]
                    count += 1
                    local_rows[count] = idx(ix+1,iy,s2)
                    local_cols[count] = idx(ix,iy,s1)
                    local_vals[count] = np.conj(Hx[s1,s2])
                    count += 1

        # y hopping
        if iy+1 < Ny:
            Hy = hopping_y(t, tR)
            for s1 in range(2):
                for s2 in range(2):
                    local_rows[count] = idx(ix,iy,s1)
                    local_cols[count] = idx(ix,iy+1,s2)
                    local_vals[count] = Hy[s1,s2]
                    count += 1
                    local_rows[count] = idx(ix,iy+1,s2)
                    local_cols[count] = idx(ix,iy,s1)
                    local_vals[count] = np.conj(Hy[s1,s2])
                    count += 1

        return local_rows[:count], local_cols[:count], local_vals[:count]


    for ix in prange(Nx):
        for iy in range(Ny):
            r, c, v = fill_site(ix, iy)
            offset = (ix * Ny + iy) * max_entries_per_site
            count = len(r)
            rows[offset:offset+count] = r
            cols[offset:offset+count] = c
            vals[offset:offset+count] = v

    return rows, cols, vals

@njit(parallel=True)
def build_electric_field_potential(Nx: np.int64, 
                                   Ny: np.int64, 
                                   a: np.float64, 
                                   E0: np.float64,
                                   angle: np.float64):
    """
    Build a 2D grid potential for a uniform electric field

    Parameters
    ----------
    Nx, Ny : number of grid points in x and y
    a      : grid spacing
    E0     : electric field amplitude (Hz/(m C))
    angle  : angle of the electric field (radians)
    Returns
    -------
    Vgrid : 2D ndarray
        Potential values on the grid.
    """
    e = 1.6e-19  # C
    Vgrid = np.zeros((Ny, Nx), dtype=np.float64)
    Ex = E0 * np.cos(angle)
    Ey = -E0 * np.sin(angle)
    cx = (Nx - 1) / 2
    cy = (Ny - 1) / 2

    for idx in prange(Nx * Ny):
        ix = idx % Nx
        iy = idx // Nx
        x = (ix - cx) * a
        y = (iy - cy) * a
        Vgrid[iy, ix] = - e * (Ex * x + Ey * y)

    return Vgrid


@njit(parallel=True)
def build_harmonic_potential(Nx, Ny, a, m_eff, w0):
    """
    Build a 2D harmonic potential V = 0.5 * m_eff * w0^2 * (x^2 + y^2)
    centered on the grid.
    
    Nx, Ny : number of grid points in x and y
    a      : grid spacing
    m_eff  : effective mass (in units of s/m^2, i.e., hbar = 1)
    w0     : harmonic frequency
    """
    Vgrid = np.zeros((Ny, Nx), dtype=np.float64)
    
    cx = (Nx - 1) / 2
    cy = (Ny - 1) / 2
    
    for idx in prange(Nx * Ny):
        ix = idx % Nx
        iy = idx // Nx
        x = (ix - cx) * a
        y = (iy - cy) * a
        Vgrid[iy, ix] = 0.5 * m_eff * w0**2 * (x**2 + y**2)
    
    return Vgrid


def trotter_rotating_electric_field(psi0, Hstatic, fd, Ed, a, Nx, Ny, Npercycle, Ncycle, Nkeep):
    """
    Calculate the time evolution of a tight-binding system under a rotating electric field
    using the Trotter-Suzuki decomposition in a reduced, comoving basis.

    Parameters
    ----------
    psi0 : ndarray, shape (Nx * Ny * 2,)
        Initial state vector.
    Hstatic : sparse matrix
        Static part of the Hamiltonian in rad/s.
    fd : float
        Drive frequency (Hz).
    Ed : float
        Electric field amplitude (V/m).
    a : float
        Lattice constant (m).
    Nx : int
        Number of grid points in x.
    Ny : int
        Number of grid points in y.
    Npercycle : int
        Number of timesteps per cycle.
    Ncycle : int
        Number of cycles.
    Nkeep : int
        Number of energy levelss to keep.
    """
    T = 1 / fd
    dt = T / Npercycle

    Egrid = build_electric_field_potential(Nx, Ny, a, Ed, 0)  # field along y
    H0 = Hstatic + discretize_potential_sparse(Egrid, spin_multiplicity=2)
    E0, V0 = sp.linalg.eigsh(H0, k=Nkeep, which='SA')  # Smallest real parts
    order = np.argsort(E0)
    E0 = np.array(E0[order])
    V0 = np.array(V0[:, order])
    for k in range(Nkeep):
        j = np.argmax(np.abs(V0[:, k]))
        V0[:, k] = V0[:, k] / V0[j, k] * np.abs(V0[j, k])

    # store evolution operators at each step (similar to UUt in MATLAB)
    UUt = np.zeros((Nkeep * Npercycle, Nkeep), dtype=np.complex128)
    # first step
    # UUt[0:Nkeep, :] = np.eye(Nkeep, dtype=np.complex128)

    Utemp = V0 @ np.diag(np.exp(-1j *E0 * dt)) # time evolution in reduced basis, psi0 is already in reduced basis

    # iterative propagation, Utemp is the evolution operator in the full space at time t, projected to the reduced basis
    # Utemp has shape (Nx*Ny*2, Nkeep), but we store the time evolution only in the comoving, reduced basis (Nkeep, Nkeep)
    for k in range(1, Npercycle):
        t = k * dt
        Egrid = build_electric_field_potential(Nx, Ny, a, Ed, fd * 2 * np.pi * t)  # electric field
        Ht = Hstatic + discretize_potential_sparse(Egrid, spin_multiplicity=2)
        E, V = sp.linalg.eigsh(Ht, k=Nkeep, which='SA')  # Smallest real parts
        
        order = np.argsort(E)
        E = np.array(E[order])
        V = np.array(V[:, order])
        
        Utemp = np.conj(V.T) @ Utemp  # change basis to the comoving basis, Utemp has shape (Nkeep, Nkeep)
        # for l in range(Nkeep):
        #     Utemp[:,l] = Utemp[:,l]/Utemp[l,l] * np.abs(Utemp[l,l])
        
        UUt[Nkeep*(k-1):Nkeep*k, :] = Utemp
        Utemp = V @ np.diag(np.exp(-1j * E * dt)) @ Utemp  # propagate and switch back to full basis, Utemp has shape (Nx*Ny*2, Nkeep)

    # End of cycle
    UT = V0.conj().T @ Utemp  # one full-cycle propagator in the reduced basis, shape (Nkeep, Nkeep)

    UUt[Nkeep*(Npercycle-1):Nkeep*Npercycle, :] = UT
    # results: (Nkeep, Npercycle*Ncycle)
    psipsi = np.zeros((Nkeep, Npercycle * Ncycle + 1), dtype=np.complex128)
    psipsi[:, 0] = psi0 # initial state in the reduced basis

    for c in range(Ncycle):
        block = UUt @ psi0   # (Nkeep*Npercycle,)
        psipsi[:, c*Npercycle+1:(c+1)*Npercycle+1] = reshape_columnwise(block, (Nkeep, Npercycle))
        psi0 = UT @ psi0

    return psipsi