from typing import Tuple

import numpy as np

from pypulseq.opts import Opts


def traj_to_grad(
    k: np.ndarray, raster_time: float = None
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Convert k-space trajectory `k` into gradient waveform in compliance with `raster_time` gradient raster time.

    Parameters
    ----------
    k : numpy.ndarray
        K-space trajectory to be converted into gradient waveform.
    raster_time : float, default=Opts().grad_raster_time
        Gradient raster time.

    Returns
    -------
    g : numpy.ndarray
        Gradient waveform.
    sr : numpy.ndarray
        Slew rate.
    """
    if raster_time == None:
        raster_time = Opts.default.grad_raster_time
    
    # Compute finite difference for gradients in Hz/m
    g = (k[...,1:] - k[...,:-1]) / raster_time
    # Compute the slew rate
    sr0 = (g[...,1:] - g[...,:-1]) / raster_time

    # Gradient is now sampled between k-space points whilst the slew rate is between gradient points
    sr = np.zeros(sr0.shape[:-1] + (sr0.shape[-1] + 1,))
    sr[...,0] = sr0[...,0]
    sr[...,1:-1] = 0.5 * (sr0[...,:-1] + sr0[...,1:])
    sr[...,-1] = sr0[...,-1]

    return g, sr
