#!/usr/bin/env python3

r'''Studies the 2d surface b-spline interpolation

The expressions for the interpolated surface I need aren't entirely clear, so I
want to study this somewhat. This script does that.

I want to use b-splines to drive the generic camera models. These provide
local support in the control points at the expense of not interpolating the
control points. I don't NEED an interpolating spline, so this is just fine.
Let's assume the knots lie at integer coordinates.

I want a function that takes in

- the control points in a neighborhood of the spline segment
- the query point x, scaled to [0,1] in the spline segment

This is all nontrivial for some reason, and there're several implementations
floating around with slightly different inputs. I have

- sample_segment_cubic()
  From a generic calibration research library:
  https://github.com/puzzlepaint/camera_calibration/blob/master/applications/camera_calibration/scripts/derive_jacobians.py
  in the EvalUniformCubicBSpline() function. This does what I want (query
  point in [0,1], control points around it, one segment at a time), but the
  source of the expression isn't given. I'd like to re-derive it, and then
  possibly extend it

- splev_local()
  wraps scipy.interpolate.splev()

- splev_translated()
  Followed sources of scipy.interpolate.splev() to the core fortran functions
  in fpbspl.f and splev.f. I then ran these through f2c, pythonified it, and
  simplified it. The result is sympy-able

- splev_wikipedia()
  Sample implementation of De Boor's algorithm from wikipedia:
  https://en.wikipedia.org/wiki/De_Boor%27s_algorithm
from EvalUniformCubicBSpline() in
camera_calibration/applications/camera_calibration/scripts/derive_jacobians.py.
translated to [0,1] from [3,4]

'''

import sys
import numpy as np
import numpysane as nps
import gnuplotlib as gp



skip_plots = False



import scipy.interpolate
from scipy.interpolate import _fitpack

def sample_segment_cubic(x, a,b,c,d):
    A =  (-x**3 + 3*x**2 - 3*x + 1)/6
    B = (3 * x**3/2 - 3*x**2 + 2)/3
    C = (-3 * x**3 + 3*x**2 + 3*x + 1)/6
    D = (x * x * x) / 6

    return A*a + B*b + C*c + D*d

def splev_local(x, t, c, k, der=0, ext=0):
    y = scipy.interpolate.splev(x, (t,c,k), der, ext)
    return y.reshape(x.shape)

def splev_translated(x, t, c, k, l):

    # assumes that t[l] <= x <= t[l+1]

    # print(f"l = {l}")
    # print((t[l], x, t[l+1]))

    l += 1 # l is now a fortran-style index

    hh  = [0] * 19
    h__ = [0] * 19

    h__[-1 + 1] = 1
    i__1 = k
    for j in range(1,i__1+1):
        i__2 = j
        for i__ in range(1,i__2+1):
            hh[-1 + i__] = h__[-1 + i__]
        h__[-1 + 1] = 0
        i__2 = j
        for i__ in range(1,i__2+1):
            li = l + i__
            lj = li - j
            if t[-1 + li] != t[-1 + lj]:
                h__[-1 + i__]    += (t[-1 + li] - x) * hh[-1 + i__] / (t[-1 + li] - t[-1 + lj])
                h__[-1 + i__ + 1] = (x - t[-1 + lj]) * hh[-1 + i__] / (t[-1 + li] - t[-1 + lj])
            else:
                h__[-1 + i__ + 1] = 0

    sp = 0
    ll = l - (k+1)
    i__2 = (k+1)
    for j in range(1,i__2+1):
        ll += 1
        sp += c[-1 + ll] * h__[-1 + j]
    return sp

# from https://en.wikipedia.org/wiki/De_Boor%27s_algorithm
def splev_wikipedia(k: int, x: int, t, c, p: int):
    """Evaluates S(x).

    Arguments
    ---------
    k: Index of knot interval that contains x.
    x: Position.
    t: Array of knot positions, needs to be padded as described above.
    c: Array of control points.
       We will look at c[k-p .. k]
    p: Degree of B-spline.
    """
    # make sure I never reference c out of bounds
    if k-p < 0: raise Exception("c referenced out of min bounds")
    if k >= len(c): raise Exception("c referenced out of max bounds")
    d = [c[j + k - p] for j in range(0, p+1)]

    for r in range(1, p+1):
        for j in range(p, r-1, -1):
            alpha = (x - t[j+k-p]) / (t[j+1+k-r] - t[j+k-p])
            d[j] = (1 - alpha) * d[j-1] + alpha * d[j]

    return d[p]


##############################################################################
# First I confirm that the functions from numpy and from wikipedia produce the
# same results. I don't care about edge cases for now. I query an arbitrary
# point

N = 30
t = np.arange( N, dtype=int)
k = 3
c = np.random.rand(len(t))

x  = 4.3
# may have trouble exactly AT the knots. the edge logic isn't quite what I want
l = np.searchsorted(np.array(t), x)-1
y0 = splev_local(np.array((x,)),t,c,k)[0]
y1 = splev_translated(x, t, c, k, l)
y2 = splev_wikipedia(l, x, t, c, k)

print(f"These should all match: {y0} {y1} {y2}")
print(f"  err1 = {y1-y0}")
print(f"  err2 = {y2-y0}")



##############################################################################
# OK, good. I want to make sure that the spline roughly follows the curve
# defined by the control points. There should be no x offset or anything of that
# sort
N = 30
t = np.arange( N, dtype=int)
k = 3
c = np.random.rand(len(t))
Npad = 10
x = np.linspace(Npad, N-Npad, 1000)

########### sample_segment_cubic()
@nps.broadcast_define(((),), ())
def sample_cubic(x, cp):
    i = int(x//1)
    q = x-i
    try:    return sample_segment_cubic(q, *cp[i-1:i+3])
    except: return 0

y = sample_cubic(x, c)
c2 = c.copy()
c2[int(N//2)] *= 1.1
y2 = sample_cubic(x, c2)

if not skip_plots:
    plot1 = gp.gnuplotlib(title = 'Cubic splines: response to control point perturbation',
                          # hardcopy = '/tmp/cubic-spline-perturbations.svg',
                          # terminal = 'svg size 800,600 noenhanced solid dynamic font ",12"'
                          )
    plot1.plot( (x, nps.cat(y,y2),
                 dict(_with = np.array(('lines lc "blue"',
                                        'lines lc "sea-green"')),
                      legend = np.array(('Spline: baseline',
                                         'Spline: tweaked one control point')))),
                (t[:len(c)], nps.cat(c,c2),
                 dict(_with = np.array(('points pt 1 ps 1 lc "blue"',
                                        'points pt 2 ps 1 lc "sea-green"')),
                      legend= np.array(('Control points: baseline',
                                        'Control points: tweaked one control point')))),
                (x, y-y2,
                 dict(_with  = 'lines lc "red"',
                      legend = 'Difference',
                      y2     = 1)),
                _xrange = (10.5,19.5),
                y2max   = 0.01,
                ylabel   = 'Spline value',
                y2label  = 'Difference due to perturbation',)

########### sample_wikipedia()
@nps.broadcast_define(((),), ())
def sample_wikipedia(x, t, c, k):
    return splev_wikipedia(np.searchsorted(np.array(t), x)-1,x,t,c,k)

@nps.broadcast_define(((),), ())
def sample_wikipedia_integer_knots(x, c, k):
    t = np.arange(len(c) + k, dtype=int)
    l = int(x//1)
    offset = int((k+1)//2)
    return splev_wikipedia(l,x,t,c[offset:],k)

if 1:
    y = sample_wikipedia_integer_knots(x,c,k)
else:
    offset = int((k+1)//2)
    y = sample_wikipedia(x,t, c[offset:], k)

if not skip_plots:
    plot2 = gp.gnuplotlib(title = 'sample_wikipedia')
    plot2.plot( (x, y, dict(_with='lines')),
                (t[:len(c)], c,  dict(_with='linespoints pt 7 ps 2')) )
    print("these two plots should look the same: we're using two implementation of the same algorithm to interpolate the same data")
    plot2.wait()
    plot1.wait()


# Now I use sympy to get the polynomial coefficients from sample_wikipedia.
# These should match the ones in sample_segment_cubic()

import sympy
c = sympy.symbols(f'c:{len(c)}')
x = sympy.symbols('x')

# Which interval we're in. Arbitrary. In the middle somewhere
l = int(N//2)
print("Should match A,B,C,D coefficients in sample_segment_cubic()")
s = splev_wikipedia(l,
                    # I want 'x' to be [0,1] within the interval, but this
                    # function wants x in the whole domain
                    x+l,
                    np.arange( N, dtype=int), c, k).expand()
print(s.coeff(c[12]))
print(s.coeff(c[13]))
print(s.coeff(c[14]))
print(s.coeff(c[15]))

print("Should also match A,B,C,D coefficients in sample_segment_cubic()")
s = splev_translated(# I want 'x' to be [0,1] within the interval, but this
                     # function wants x in the whole domain
                     x+l,
                     np.arange( N, dtype=int), c, k, l).expand()
print(s.coeff(c[12]))
print(s.coeff(c[13]))
print(s.coeff(c[14]))
print(s.coeff(c[15]))


#########################################################
# Great. More questions. Here I have a cubic spline (k==3). And to evaluate the
# spline value I need to have 4 control point values available, 2 on either
# side. Piecewise-linear splines are too rough, but quadratic splines could work
# (k==2). What do the expressions look like? How many neighbors do I need? Here
# the control point values c represent the function value between adjacent
# knots, so each segment uses 3 neighboring control points, and is defined in a
# region [-0.5..0.5] off the center control point
print("======================== k = 2")
N = 30
t = np.arange( N, dtype=int)
k = 2
# c[0,1,2] corresponds to is t[-0.5 0.5 1.5 2.5 ...
c = np.random.rand(len(t))

x  = 4.3
# may have trouble exactly AT the knots. the edge logic isn't quite what I want
l = np.searchsorted(np.array(t), x)-1
y0 = splev_local(np.array((x,)),t,c,k)[0]
y1 = splev_translated(x, t, c, k, l)
y2 = splev_wikipedia(l, x, t, c, k)
print(f"These should all match: {y0} {y1} {y2}")
print(f"  err1 = {y1-y0}")
print(f"  err2 = {y2-y0}")


##############################################################################
# OK, good. I want to make sure that the spline roughly follows the curve
# defined by the control points. There should be no x offset or anything of that
# sort
if not skip_plots:
    N = 30
    t = np.arange( N, dtype=int)
    k = 2
    c = np.random.rand(len(t))
    Npad = 10
    x = np.linspace(Npad, N-Npad, 1000)
    offset = int((k+1)//2)

    y = sample_wikipedia(x,t-0.5, c[offset:], k)

    xm = (x[1:] + x[:-1]) / 2.
    d = np.diff(y) / np.diff(x)
    plot1 = gp.gnuplotlib(title = 'k==2; sample_wikipedia')
    plot1.plot( (x, y, dict(_with='lines', legend='spline')),
                (xm, d, dict(_with='lines', y2=1, legend='diff')),
                (t[:len(c)], c,  dict(_with='linespoints pt 7 ps 2', legend='control points')))

    @nps.broadcast_define(((),), ())
    def sample_splev_translated(x, t, c, k):
        l = np.searchsorted(np.array(t), x)-1
        return splev_translated(x,t,c,k,l)
    y = sample_splev_translated(x,t-0.5, c[offset:], k)
    xm = (x[1:] + x[:-1]) / 2.
    d = np.diff(y) / np.diff(x)
    plot2 = gp.gnuplotlib(title = 'k==2; splev_translated')
    plot2.plot( (x, y, dict(_with='lines', legend='spline')),
                (xm, d, dict(_with='lines', y2=1, legend='diff')),
                (t[:len(c)], c,  dict(_with='linespoints pt 7 ps 2', legend='control points')))


    # These are the functions I'm going to use. Derived by the sympy steps
    # immediately after this
    def sample_segment_quadratic(x, a,b,c):
        A = (4*x**2 - 4*x + 1)/8
        B = (3 - 4*x**2)/4
        C = (4*x**2 + 4*x + 1)/8
        return A*a + B*b + C*c
    @nps.broadcast_define(((),), ())
    def sample_quadratic(x, cp):
        i = int((x+0.5)//1)
        q = x-i
        try:    return sample_segment_quadratic(q, *cp[i-1:i+2])
        except: return 0
    y = sample_quadratic(x, c)
    xm = (x[1:] + x[:-1]) / 2.
    d = np.diff(y) / np.diff(x)
    plot3 = gp.gnuplotlib(title = 'k==2; sample_quadratic')
    plot3.plot( (x, y, dict(_with='lines', legend='spline')),
                (xm, d, dict(_with='lines', y2=1, legend='diff')),
                (t[:len(c)], c,  dict(_with='linespoints pt 7 ps 2', legend='control points')))

    plot3.wait()
    plot2.wait()
    plot1.wait()

    print("these 3 plots should look the same: we're using different implementation of the same algorithm to interpolate the same data")


# ##################################
# # OK, these match. Let's get the expression of the polynomial in a segment. This
# # was used to construct sample_segment_quadratic() above
# c = sympy.symbols(f'c:{len(c)}')
# x = sympy.symbols('x')
# l = int(N//2)
# print("A,B,C for k==2 using splev_wikipedia()")
# s = splev_wikipedia(l,
#                     # I want 'x' to be [-0.5..0.5] within the interval, but this
#                     # function wants x in the whole domain
#                     x+l,
#                     np.arange( N, dtype=int) - sympy.Rational(1,2), c, k).expand()
# print(s)
# print(s.coeff(c[13]))
# print(s.coeff(c[14]))
# print(s.coeff(c[15]))
# # I see this:
# #   c13*x**2/2 - c13*x/2 + c13/8 - c14*x**2 + 3*c14/4 + c15*x**2/2 + c15*x/2 + c15/8
# #   x**2/2 - x/2 + 1/8
# #   3/4 - x**2
# #   x**2/2 + x/2 + 1/8


############## compare cubic, quadratic
# I now have nice expressions for quadratic and cubic interpolations. Both are
# C1 continuous, but the cubic interpolation is C2 continuous too. How much do I
# care? Let's at least compare them visually
N = 30
t = np.arange( N, dtype=int)
c = np.random.rand(len(t))
Npad = 10
x = np.linspace(Npad, N-Npad, 1000)

y2 = sample_quadratic(x,c)
y3 = sample_cubic(    x,c)

if not skip_plots:
    gp.plot( (x, nps.cat(y2,y3), dict(_with='lines',
                                      legend=np.array(('quadratic',
                                                       'cubic')))),
             (t[:len(c)], c, dict(_with='linespoints pt 7 ps 2',
                                  legend='control points')),
             title = "Comparing quadratic and cubic b-spline interpolation",
             wait = True)
    # Visually, neither looks better than the other. The quadratic spline follows
    # the control points closer. Looks like it's trying to match the slope at the
    # halfway point. Maybe that's bad, and the quadratic curve is too wiggly? We'll
    # see

####################################################################
# Great. Final set of questions: how do you we make a 2D spline surface? The
# generic calibration research library
# (https://github.com/puzzlepaint/camera_calibration) Does a set of 1d
# interpolations in one dimension, and then interpolates the 4 interpolated
# values along the other dimension. Questions:
#
# Does order matter? I can do x and then y or y then x
#
# And is there a better way? scipy has 2d b-spline interpolation routines. Do
# they do something better?
#
# ############### x-y or y-x?
import sympy
from sympy.abc import x,y
cp = sympy.symbols('cp:4(:4)')

# x then y
xs = [sample_segment_cubic( x,
                            cp[i*4 + 0],
                            cp[i*4 + 1],
                            cp[i*4 + 2],
                            cp[i*4 + 3] ) for i in range(4)]
yxs = sample_segment_cubic(y, *xs)

# y then x
ys = [sample_segment_cubic( y,
                            cp[0*4 + i],
                            cp[1*4 + i],
                            cp[2*4 + i],
                            cp[3*4 + i] ) for i in range(4)]
xys = sample_segment_cubic(x, *ys)

print(f"Bicubic interpolation. x-then-y and y-then-x difference: {(xys - yxs).expand()}")

########### Alright. Apparently either order is ok. Does scipy do something
########### different for 2d interpolation? I compare the 1d-then-1d
########### interpolation above to fitpack results:

from fpbisp import fpbisp_

N = 30
t = np.arange( N, dtype=int)
k = 3

cp = np.array(sympy.symbols('cp:40(:40)'), dtype=np.object).reshape(40,40)

lx = 3
ly = 5

z   = fpbisp_(t, t, k, cp, x+lx, lx, y+ly, ly)
err = z - yxs

print(f"Bicubic interpolation. difference(1d-then-1d, FITPACK): {err.expand()}")

# Apparently chaining 1D interpolations produces identical results to what FITPACK is doing
