"""
General utility functions.
"""

import torch

def interpolate(x: torch.Tensor, xp: torch.Tensor, fp: torch.Tensor) -> torch.Tensor:
    """One-dimensional linear interpolation for monotonically increasing sample
    points.

    Returns the one-dimensional piecewise linear interpolant to a function with
    given discrete data points :math:`(xp, fp)`, evaluated at :math:`x`.

    From https://github.com/pytorch/pytorch/issues/50334#issuecomment-1247611276.
  
    Parameters
    ----------
    x : torch.Tensor
        the :math:`x`-coordinates at which to evaluate the interpolated values.
    xp : torch.Tensor
        the :math:`x`-coordinates of the data points, must be increasing.
    fp : torch.Tensor
        the :math:`y`-coordinates of the data points, same length as `xp`.

    Returns
    -------
    torch.Tensor
        the interpolated values, same size as `x`
    """
    m = (fp[:,1:] - fp[:,:-1]) / (xp[:,1:] - xp[:,:-1])  #slope
    b = fp[:, :-1] - (m.mul(xp[:, :-1]) )

    indicies = torch.sum(torch.ge(x[:, :, None], xp[:, None, :]), -1) - 1  #torch.ge:  x[i] >= xp[i] ? true: false
    indicies = torch.clamp(indicies, 0, m.shape[-1] - 1)

    line_idx = torch.linspace(0, indicies.shape[0], 1, device=indicies.device).to(torch.long)
    line_idx = line_idx.expand(indicies.shape)
    # idx = torch.cat([line_idx, indicies] , 0)
    return m[line_idx, indicies].mul(x) + b[line_idx, indicies]
