#!/usr/bin/env python

g_program_name = __file__.split('/')[-1]  # = 'interpolate_mt.py'
g_version_str = '0.1.0'
g_date_str = '2019-8-22'

g_usage_str = """
Usage:

   """ + g_program_name + """ Ndesired [scale] < coords_orig.raw > coords.raw

Example:

   """ + g_program_name + """ 30117 3.0 < coords_orig.raw > coords.raw

"""


import sys
from math import *
import numpy as np
import bisect



## Tri Diagonal Matrix Algorithm(a.k.a Thomas algorithm) solver stolen from:
# https://gist.github.com/ofan666/1875903

def TDMAsolver(a, b, c, d):
    '''
    TDMA solver, a b c d can be NumPy array type or Python list type.
    Solves a (non-cyclic) system of equations of the form:

      a_i*x_{i-1} + b_i*x_i + c_i*x_{i+1} = d_i
      where a_1=0, and c_n=0

    refer to http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm

    '''

    nf = len(a)     # number of equations
    ac, bc, cc, dc = list(map(np.array, (a, b, c, d)))     # copy the array
    for it in range(1, nf):
        mc = ac[it]/bc[it-1]
        bc[it] = bc[it] - mc*cc[it-1] 
        dc[it] = dc[it] - mc*dc[it-1]

    xc = ac
    xc[-1] = dc[-1]/bc[-1]

    for il in range(nf-2, -1, -1):
        xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]

    del bc, cc, dc  # delete variables from memory

    return xc



def CalcNaturalCubicSplineCoeffs(r, spline_exponent_alpha):
    N = len(r)
    assert(N >= 4)
    D = len(r[0])
    tcontrol = np.zeros(N)

    # e_d2rdt2[d][i] is the second derivative of the spline at the ith 
    #               control point (and in the dth direction)
    #
    # Once we have figured out e_d2rdt2[d][i], we can calculate the spline
    # anywhere using:
    #   SplineEval(t, h_i, 
    #              e_{i+1} / (6*h_i), 
    #              e_i / (6*h_{i+1}),
    #              r_{i+1}/h_i - e_{i+1}*h_i,
    #              r_i/h_i - e_i*h_i)

    e_d2rdt2 = np.zeros((D,N))

    # We want to solve this system of equations for e_1, e_2, ..., e_{n-2}:
    #    (note: e_i is shorthand for e_d2rdt2[d][i])
    #
    # h_{i-1}*e_{i-1} + u_i*e_i + h_{i+1} * e_{i+1}  =  v_i
    #   where h_i, u_i and v_i are shorthand for:
    #
    # h_i = change in "time" parameter for the ith interval, which is usually:
    #     = |r_{i+1} - r_i |^alpha    (alpha varies depending on settings)
    #   and
    # u_i = 2 * (h_{i-1} + h_i )
    # v_i = 6 * (b_i     - b_{i-1})
    # b_i =  (r_i - r_{i-1}) / h_i
    #
    # ...subject to the constraints that the curve is straight at the endpoints:
    # e_0 = 0       <-- first control point  (indexing begins at 0)
    # e_{n-1} = 0   <-- this is the last control point  (indexing begins at 0)

    h_dt = np.zeros(N-1)   # h_dt[i] is the i'th time interval 
                         # in the parameterization
    b_drdt = np.zeros((N-1,D))  # b_drdt is a discrete version of the derivative

    t_total=0.0
    for i in range(0, N-1):
        tcontrol[i] = t_total
        sqr_distance_i_ip1 = 0.0
        for d in range(0, D):
            sqr_distance_i_ip1 += (r[i+1][d] - r[i][d])**2
        h_dt[i] = sqr_distance_i_ip1**(0.5*spline_exponent_alpha)
        for d in range(0, D):
            b_drdt[i][d] = (r[i+1][d] - r[i][d]) / h_dt[i]
        t_total += h_dt[i]
    tcontrol[N-1] = t_total

    a_coeff = np.zeros(N)
    b_coeff = np.zeros(N)
    c_coeff = np.zeros(N)
    d_coeff = np.zeros((D,N))

    for i in range(1, N-1):
        # h_dt[i] is the difference in "time" in the parametric curve 
        # between pairs of control points.  If spline_exponenent_alpha is 0
        # then the time interval between control points is uniform.  (Typically,
        # "spline_exponent_alpha" is 0.5 for centripital Catmull-Rom splines.)
        a_coeff[i]    =      h_dt[i-1]
        b_coeff[i]    = 2.0*(h_dt[i-1] + h_dt[i])
        c_coeff[i]    =      h_dt[i]
        for d in range(0, D):
            d_coeff[d][i] = 6.0*(b_drdt[i][d] - b_drdt[i-1][d])

    a_coeff[0] = 0.0
    b_coeff[0] = 1.0
    c_coeff[0] = 0.0
    for d in range(0, D):
        d_coeff[d][0] = 0.0
    a_coeff[N-1] = 0.0
    b_coeff[N-1] = 1.0
    c_coeff[N-1] = 0.0
    for d in range(0, D):
        d_coeff[d][N-1] = 0.0

    for d in range(0, D):
        e_d2rdt2[d] = TDMAsolver(a_coeff, b_coeff, c_coeff, d_coeff[d])

    # alternately, if that fails, try the matrix inverter that comes with numpy:
    #M = np.zeros((N,N))
    #for i in range(0,N):
    #    if i-1>=0:
    #        M[i][i-1] = a_coeff[i]
    #    M[i][i]   = b_coeff[i]
    #    if i+1 < N:
    #        M[i][i+1] = c_coeff[i]
    #for d in range(0, D):
    #    e_d2rdt2[d] = np.linalg.solve(M, d_coeff[d])

    e_d2rdt2 = e_d2rdt2.transpose()
    c3a = np.zeros((N, D))
    c3b = np.zeros((N, D))
    c1a = np.zeros((N, D))
    c1b = np.zeros((N, D))
    # faster to precompute these coefficients in advance:
    for i in range(0, N-1):
        # c3a = e_{i+1} / (6*h_i)
        # c3b =     e_i / (6*h_{i+1})
        # c1a = r_{i+1}/h_i - e_{i+1}*h_i
        # c1b =  r_i / h_i - e_i*h_i)
        c3a[i] = e_d2rdt2[i+1] / (6*h_dt[i])
        c3b[i] = e_d2rdt2[i]   / (6*h_dt[i])
        c1a[i] = r[i+1]/h_dt[i]  -  e_d2rdt2[i+1]*h_dt[i]/6.0
        c1b[i] = r[i]/h_dt[i]    -  e_d2rdt2[i]*h_dt[i]/6.0

    # Return these spline coefficients to the caller.
    # We can use these to quickly evaluate the spline repetatively later on

    return c3a, c3b, c1a, c1b, tcontrol





def SplineEval(t, t_interval, c3a, c3b, c1a, c1b):
    ta = t
    tb = t_interval - t
    return (c3a*ta*ta + c1a)*ta + (c3b*tb*tb + c1b)*tb



def SplineEvalD1(t, t_interval, c3a, c3b, c1a, c1b):
    ta = t
    tb = t_interval - t
    return c3a*ta*ta*3.0 + c1a - (c3b*tb*tb*3.0 + c1b)



def SplineEvalD2(t, t_interval, c3a, c3b, c1a, c1b):
    ta = t
    tb = t_interval - t
    return c3a*ta*6.0 + c3b*tb*6.0



def SplineInterpEval(t, c3a, c3b, c1a, c1b, tcontrol):
    i = bisect.bisect(tcontrol, t) - 1
    if i < 0: i = 0
    if i >= len(tcontrol)-1: i = len(tcontrol)-2
    return SplineEval(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
                      c3a[i], c3b[i], c1a[i], c1b[i])


def SplineInterpEvalD1(t, c3a, c3b, c1a, c1b, tcontrol):
    i = bisect.bisect(tcontrol, t) - 1
    if i < 0: i = 0
    if i >= len(tcontrol)-1: i = len(tcontrol)-2
    return SplineEvalD1(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
                        c3a[i], c3b[i], c1a[i], c1b[i])



def SplineInterpEvalD2(t, c3a, c3b, c1a, c1b, tcontrol):
    i = bisect.bisect(tcontrol, t) - 1
    if i < 0: i = 0
    if i >= len(tcontrol)-1: i = len(tcontrol)-2
    return SplineEvalD2(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
                        c3a[i], c3b[i], c1a[i], c1b[i])



def SplineCurvature2D(t, t_interval, c3a, c3b, c1a, c1b):
    # first derivatives
    x1,y1 = SplineEvalD1(t, t_interval, c3a, c3b, c1a, c1b)
    # second derivatives
    x2,y2 = SplineEvalD2(t, t_interval, c3a, c3b, c1a, c1b)
    denom = pow((x1*x1 + y1*y1), 1.5)
    curvature = (x1*y2 - x2*y1) / denom
    return curvature



def SplineInterpCurvature2D(t, c3a, c3b, c1a, c1b, tcontrol):
    i = bisect.bisect(tcontrol, t) - 1
    if i < 0: i = 0
    if i >= len(tcontrol)-1: i = len(tcontrol)-2
    return SplineCurvature2D(t-tcontrol[i], tcontrol[i+1]-tcontrol[i],
                             c3a[i], c3b[i], c1a[i], c1b[i])



class InputError(Exception):
    """ A generic exception object containing a string for error reporting.
        (Raising this exception implies that the caller has provided
         a faulty input file or argument.)

    """

    def __init__(self, err_msg):
        self.err_msg = err_msg

    def __str__(self):
        return self.err_msg

    def __repr__(self):
        return str(self)


def main():

    try:

        #######  Main Code Below: #######
        sys.stderr.write(g_program_name+' v'+g_version_str+' '+g_date_str+'\n')

        spline_exponent_alpha = 0.5
        use_linear_interpolation = False

        # Parse the argument list:
        if len(sys.argv) <= 1:
            raise InputError("Missing arguments\n"+g_usage_str+"\n")
                             

        n_new = int(sys.argv[1])

        if len(sys.argv) > 2:
            scale = float(sys.argv[2])
        else:
            scale = 1.0

        coords_orig = []

        lines = sys.stdin.readlines()

        for line in lines:
            tokens = line.split()
            if (len(tokens) > 0):
                coords_orig.append(list(map(float, tokens)))
                g_dim = len(tokens)

        n_orig = len(coords_orig)

        x_orig = np.array(coords_orig)

        if n_orig < 2:
            raise InputError("Input file contains less than two lines of coordinates.")

        if n_new < 2:
            raise InputError("Output file will contain less than two lines of coordinates.")

        #coords_new = [[0.0 for d in range(0, g_dim)] for i in range(0, n_new)]
        x_new = np.zeros((n_new, g_dim))

        if use_linear_interpolation:
            for i_new in range(0, n_new):
                I_orig = (i_new) * (float(n_orig-1) / float(n_new-1))
                i_orig = int(floor(I_orig))
                i_remainder = I_orig - i_orig

                if (i_new < n_new-1):
                    for d in range(0, g_dim):
                        x_new[i_new][d] = scale*(x_orig[i_orig][d]
                                                 +
                                                 i_remainder*(x_orig[i_orig+1][d]-
                                                              x_orig[i_orig][d]))
                else:
                    for d in range(0, g_dim):
                        x_new[i_new][d] = scale*x_orig[n_orig-1][d]

        else:
            # Figure out the spline parameters: c3a, c3b, c1a, c1b, and tcontrol:
            c3a, c3b, c1a, c1b, tcontrol = \
                 CalcNaturalCubicSplineCoeffs(x_orig, spline_exponent_alpha)
            tmin = 0.0
            tmax = tcontrol[-1]
            for i in range(0, n_new):
                t = tmin + i*((tmax - tmin) / (n_new-1))
                x_new[i] = SplineInterpEval(t, c3a, c3b, c1a, c1b, tcontrol)

        # print the coordates
        for i in range(0, n_new):
            for d in range(0, g_dim-1):
                sys.stdout.write(str(x_new[i][d]*scale) + ' ')
            sys.stdout.write(str(x_new[i][g_dim-1]*scale) + "\n")

    except (ValueError, InputError) as err:
        sys.stderr.write('\n' + 'Error:\n\n' + str(err) + '\n')
        sys.exit(1)


if __name__ == '__main__':
    main()
