import numpy as np
from utils import expm2x2
from numba import njit, prange


def calc_chevron_analityc(fL, ffd, theta, DriveAmplitude, N, Tmax):
    """
    Chevron pattern of RING (magnetic drive case).

    Parameters
    ----------
    fL : float
        Larmor frequency (Hz).
    ffd : array-like
        Array of drive frequencies (Hz).
    theta : float
        Angle parameter.
    DriveAmplitude : float
        Driving amplitude (already given, unlike in EDSR case).
    N : int
        Number of time samples.
    Tmax : float
        Maximum time.

    Returns
    -------
    chevron : ndarray, shape (len(ffd), N)
        Rabi oscillation probabilities for each drive frequency.
    tt : ndarray, shape (N,)
        Time array used.
    """
    ffd = np.asarray(ffd)
    Nd = len(ffd)
    chevron = np.zeros((Nd, N), dtype=float)
    

    tt = np.linspace(0, Tmax, N)

    for j, fd in enumerate(ffd):
        Omega = ((DriveAmplitude*2*np.pi)**2 * np.sin(theta)) / (2 * fd * 2 * np.pi)
        Delta = fL * 2 * np.pi - ((DriveAmplitude*2*np.pi)**2 * np.cos(theta)) / (2 * fd * 2 * np.pi)
        WR = np.sqrt(Omega**2 + Delta**2)
        Rabi = (np.sin(WR * tt / 2)**2) * (Omega**2 / WR**2)
        chevron[j, :] = Rabi

    return chevron, tt



@njit(parallel=True)
def trotter_Tmax_circular_field(psi0, fL, fd, theta, DriveAmplitude, N, Tmax):
    """
    Trotter solver for 2x2 Hamiltonians with circular magnetic field,
    parallelized with numba.

    Parameters
    ----------
    psi0 : np.ndarray, shape (2,)
        Initial state vector.
    fL : float
        Larmor frequency (Hz).
    fd : float
        Drive frequency (Hz).
    theta : float
        Angle parameter (radians).
    DriveAmplitude : float
        Drive amplitude (same units as frequency).
    N : int
        Number of timesteps for Tmax.
    Tmax : float
        Maximum evolution time.
    """
    ts = np.linspace(0, Tmax, N)
    dt = Tmax / N

    X = np.array([[0, 1], [1, 0]], dtype=np.complex128)
    Y = np.array([[0, -1j], [1j, 0]], dtype=np.complex128)
    Z = np.array([[1, 0], [0, -1]], dtype=np.complex128)
    nvec = np.array([np.sin(theta), 0, np.cos(theta)], dtype=np.complex128)
    wd = 2 * np.pi * fd   # drive frequency
    wL = 2 * np.pi * fL   # Larmor frequency
    Omega = 2 * np.pi * DriveAmplitude 
    Hstat = 0.5 * wL * (nvec[0]*X + nvec[1]*Y + nvec[2]*Z) #hbar is set to 1
    psipsi = np.zeros((2, N), dtype=np.complex128)
    psi = psi0.copy()
    for k in range(N):
        H = Hstat + 0.5 * Omega * (np.cos(wd * ts[k]) * X + np.sin(wd * ts[k]) * Y)
        U = expm2x2(H, dt) # hbar is set to 1
        psi = U @ psi
        psipsi[:, k] = psi
    return psipsi

@njit
def trotter_cyclic_circular_field(psi0, fL, fd, theta, DriveAmplitude, Npercycle, Ncycle):
    """
    Calculate the time evolution of a two-level system under a periodic time-dependent Hamiltonian using the Trotter-Suzuki decomposition.

    Parameters
    ----------
    psi0 : ndarray, shape (2,)
        Initial state vector.
    fL : float
        Larmor frequency (Hz).
    fd : float
        Drive frequency (Hz).
    theta : float
        Angle between the drive plane normal vector and the static field (radians).
    DriveAmplitude : float
        Drive amplitude (Hz).
    Npercycle : int
        Number of timesteps per cycle.
    Ncycle : int
        Number of cycles.
    """
    T = 1 / fd
    dt = T / Npercycle

    X = np.array([[0, 1], [1, 0]], dtype=np.complex128)
    Y = np.array([[0, -1j], [1j, 0]], dtype=np.complex128)
    Z = np.array([[1, 0], [0, -1]], dtype=np.complex128)
    nvec = np.array([np.sin(theta), 0, np.cos(theta)], dtype=np.complex128)
    wd = 2 * np.pi * fd   # drive frequency
    wL = 2 * np.pi * fL   # Larmor frequency
    Omega = 2 * np.pi * DriveAmplitude 
    Hstat = 0.5 * wL * (nvec[0]*X + nvec[1]*Y + nvec[2]*Z)


    # store evolution operators at each step (similar to UUt in MATLAB)
    UUt = np.zeros((2 * Npercycle, 2), dtype=np.complex128)

    H0 = Hstat + 0.5 * Omega * (np.cos(0) * X + np.sin(0) * Y)
    # first step
    UUt[0:2, :] = np.eye(2, dtype=np.complex128)

    # iterative propagation
    for k in range(2, Npercycle + 1):
        t = (k - 1) * dt
        # hbar = 1
        H = Hstat + 0.5 * Omega * (np.cos(wd * t) * X + np.sin(wd * t) * Y)
        Udt = expm2x2(H, dt) # hbar = 1
        UUt[2*k-2:2*k, :] = Udt @ UUt[2*(k-1)-2:2*(k-1), :]

    # one full-cycle propagator
    UT = UUt[2*Npercycle-2:2*Npercycle, :]

    # results: (2, Npercycle*Ncycle)
    psipsi = np.zeros((2, Npercycle * Ncycle), dtype=np.complex128)
    

    for c in range(Ncycle):
        block = UUt @ psi0   # shape (2*Npercycle,)
        for k in range(Npercycle):
            psipsi[0, c*Npercycle+k] = block[2*k]
            psipsi[1, c*Npercycle+k] = block[2*k+1]
        psi0 = UT @ psi0

    return psipsi
    
    

@njit(parallel=True)
def calc_chevron_ESR(fL, ffd, theta, DriveAmplitude, N, Tmax):
    """
    Calculate ESR chevron (time evolution of excitation probability)
    using non-perturbative Trotter evolution.
    
    Parameters
    ----------
    fL : float
        Larmor frequency (Hz).
    ffd : array_like
        Array of drive frequencies (Hz).
    theta : float
        Angle parameter (radians).
    DriveAmplitude : float
        Drive amplitude (Hz).
    N : int
        Number of timesteps for Tmax.
    Tmax : float
        Maximum evolution time.

    Returns
    -------
    chevron : np.ndarray
        Array of shape (len(ffd), N) with excitation probabilities.
    """
    # initial ground state vector
    ground = np.array([np.sin(theta/2), -np.cos(theta/2)], dtype=np.complex128)
    Nd = len(ffd)
    chevron = np.zeros((Nd, N), dtype=float)


    for j in prange(Nd):

        # propagate
        psipsi = trotter_Tmax_circular_field(ground, fL, ffd[j], theta, DriveAmplitude, N, Tmax)

        # projection onto ground state
        ppg = np.abs(np.conj(ground) @ psipsi)**2
        chevron[j, :] = 1 - ppg[:N]

    return chevron, np.linspace(0, Tmax, N)

@njit(parallel=True)
def calc_pemax_ESR(ffL, ffd, theta, DriveAmplitude, Npercycle,
                   Ncycle=1000):
    """
    Calculate (and optionally plot) the maximal excited-state probability
    for ESR under circular magnetic drive, non-perturbative.

    Parameters
    ----------
    ffL : array-like
        Larmor frequencies (Hz).
    ffd : array-like
        Drive frequencies (Hz).
    theta : float
        Angle parameter (radians).
    DriveAmplitude : float
        Drive amplitude (Hz or angular freq units).
    Npercycle : int
        Trotter steps per cycle.
    scale : {'linear', 'log'}, optional
        Axis scaling for plot.
    Ncycle : int, optional
        Number of drive cycles for time evolution (default 1000).

    Returns
    -------
    pemax : ndarray, shape (len(ffL), len(ffd))
        Maximal excited state probability.
    """

    ffL = np.asarray(ffL)
    ffd = np.asarray(ffd)
    NL, Nd = len(ffL), len(ffd)

    pemax = np.zeros((NL, Nd))

    # Ground state vector
    ground = np.array([np.sin(theta/2), -np.cos(theta/2)], dtype=np.complex128)

    for j in prange(Nd):
        fd = ffd[j]
        for l in prange(NL):
            fL = ffL[l]

            # Propagate over Ncycle periods
            psipsi = trotter_cyclic_circular_field(ground, fL, fd, theta, DriveAmplitude, Npercycle, Ncycle)

            # Ground state overlap over time
            overlaps = np.abs(np.conjugate(ground) @ psipsi)**2

            # Maximal excited state probability
            pemax[l, j] = 1 - np.min(overlaps)

    return pemax