import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from numba import njit, prange

mpl.rcParams["font.family"] = "serif"
mpl.rcParams["mathtext.fontset"] = "stix"
mpl.rcParams.update({'font.size': 24})

@njit
def expm2x2(H, dt):
    """
    Compute exp(-i H dt) for a 2x2 Hermitian matrix H.
    Uses closed form with Pauli decomposition.
    """
    # Eigen-decomposition of 2x2 Hermitian
    tr = (H[0,0] + H[1,1]) / 2.0
    disc = np.sqrt(((H[0,0]-H[1,1])/2.0)**2 + (abs(H[0,1]))**2).real

    lam1 = tr + disc
    lam2 = tr - disc
    
    # exp(-i λ t)
    e1 = np.exp(-1j*lam1*dt)
    e2 = np.exp(-1j*lam2*dt)
    
    # Projectors
    if disc > 1e-14:
        P1 = (H - lam2*np.eye(2, dtype=np.complex128)) / (lam1 - lam2)
    else:
        P1 = np.eye(2, dtype=np.complex128) / 2.0
    P2 = np.eye(2, dtype=np.complex128) - P1
    
    return e1*P1 + e2*P2


@njit(parallel=True)
def reshape_columnwise(block, new_shape):
    """
    Reshape a 1D array into a 2D array in column-major order using
    parallel processing with numba.
    
    Parameters
    ----------
    block : 1D ndarray
        Input array to be reshaped.
    new_shape : tuple
        Desired shape of the output array (rows, columns).
    
    Returns
    -------
    out : 2D ndarray
        Reshaped array.
    """
    N, M = new_shape
    out = np.empty(new_shape, dtype=block.dtype)
    total = N * M
    for idx in prange(total):
        i = idx % N         # row index
        j = idx // N        # column index
        out[i, j] = block[idx]
    return out

def plot_chevron(tt, ffd, chevron, fdlim=None, savefile=None):
    """
    Plot a chevron pattern as a heatmap.

    Parameters
    ----------
    tt : array_like
        Time array (seconds).
    ffd : array_like
        Drive frequencies (Hz).
    chevron : 2D array
        Excited state probability vs time and drive frequency.
    fdlim : tuple or list
        Limits for y-axis in Hz, e.g. (fmin, fmax).
    savefile : str or None
        If provided, save the figure to this file.
    """
    fig, ax = plt.subplots()

    # Create heatmap
    cmap = plt.get_cmap("Blues_r")  # "sky" doesn't exist, this is similar
    im = ax.imshow(
        chevron,
        aspect="auto",
        origin="lower",
        extent=[tt[0]*1e9, tt[-1]*1e9, ffd[0]*1e-9, ffd[-1]*1e-9],
        cmap=cmap,
        vmin=0,
        vmax=1,
        interpolation="none" 
    )

    # Colorbar
    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label(r"$p_\mathrm{e}$")

    # Labels and scaling
    ax.set_xlabel(r"$t$ (ns)")
    ax.set_ylabel(r"$f_\mathrm{d}$ (GHz)")
    if fdlim is None:
        fdlim = (ffd[0], ffd[-1])
    ax.set_ylim(fdlim[0]*1e-9, fdlim[1]*1e-9)
    ax.xaxis.set_major_locator(MaxNLocator(nbins=5))  # ≈ 5 ticks
    ax.yaxis.set_major_locator(MaxNLocator(nbins=6))  # ≈ 5 ticks

    plt.tight_layout()
    if savefile is not None:
        plt.savefig(savefile, dpi=500)


    plt.show()


def plot_pemax_ring(ffL, ffd, pemax, scale='linear', savefile=None):
    """ Plot the maximal excited-state probability as a heatmap.
    Parameters
    ----------
    ffL : array-like
        Larmor frequencies (Hz).
    ffd : array-like    
        Drive frequencies (Hz).
    pemax : 2D array
        Maximal excited-state probability vs Larmor and drive frequencies.
    scale : str
        Scale for axes, 'linear' or 'log'.
    savefile : str or None
        If provided, save the figure to this file.
    """

    fig, ax = plt.subplots()
    cmap = plt.get_cmap("Blues_r")
    im = ax.imshow(
        pemax.T,  # transpose so axes match MATLAB's imagesc
        extent=[ffL[0]*1e-6, ffL[-1]*1e-6, ffd[0]*1e-9, ffd[-1]*1e-9],
        origin='lower',
        aspect='auto',
        cmap=cmap,
        vmin=0, vmax=1
    )
    ax.set_xlabel(r"$f_\mathrm{L}$ (MHz)")
    ax.set_ylabel(r"$f_\mathrm{d}$ (GHz)")
    ax.set_xscale(scale)
    ax.set_yscale(scale)
    ax.xaxis.set_major_locator(MaxNLocator(nbins=4))  # ≈ 4 ticks
    ax.yaxis.set_major_locator(MaxNLocator(nbins=6))  # ≈ 4 ticks

    plt.colorbar(im, ax=ax, label=r"$p_\mathrm{e,max}$")

    plt.tight_layout()
    if savefile is not None:
        plt.savefig(savefile, dpi=500)
    plt.show()