import numpy as np
from scipy.interpolate import CubicSpline
from scipy import integrate

def arc_length(cubic_spline,parameterization:np.ndarray):
    """
    Parameters
    ----------
    cubic_spline : scipy.interpolate._cubic.CubicSpline
        Cubic spline object
    parameterization : numpy.ndarray
        Array of size N. Represents the parameterization values of the cubic spline.
        Must be monotonically increasing.
    Returns
    -------
    arclength : numpy.ndarray
        Array of size N.  
    errs :  numpy.ndarray
        DESCRIPTION.
    """
    df = lambda t : np.sqrt(np.sum(cubic_spline.derivative()(t)**2)) # arc-length
    segments = len(parameterization)-1
    sub_arclength = np.zeros((segments,))
    errs = np.zeros((segments,))
    for i in range(segments):
        subf,suberr = integrate.quad(df,parameterization[i],parameterization[i+1])
        sub_arclength[i] = subf
        errs[i] = suberr
    arclength = np.hstack((0,sub_arclength.cumsum()))
    return arclength,errs

def alp_spline(xyz:np.ndarray):
    ppdist = np.sqrt(np.sum(np.diff(xyz)**2,axis=0)) # point to point distance
    
    non_zero_indices = ppdist>0
    ppdist = ppdist[non_zero_indices]
    
    xyz = xyz[:,np.hstack((True,non_zero_indices))]
    
    linear_dist = np.hstack((0,ppdist.cumsum()))
    ldp_cs = CubicSpline(linear_dist,np.transpose(xyz))
    
    arclength,errs = arc_length(ldp_cs,linear_dist)
    alp_cs = CubicSpline(arclength,np.transpose(xyz))
    return linear_dist,ldp_cs,arclength,alp_cs

