#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Mar 24 10:38:09 2020

@author: ckielasjensen
"""

try:
    from openGJK_cython import pygjk
except:
    print('[!] Warning: OpenGJK library not installed.')

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from numba import jit, njit
import numpy as np
from scipy.special import binom

from BeBOT.polynomial.base import Base
from BeBOT.polynomial.rationalbernstein import RationalBernstein

# TODO
# * Add mapping function from [t0, tf] --> [0, 1] (t to tau)
# * Change cache values to reflect tf-t0 and use a hashtable lookup (in Base
#   class)
class Bernstein(Base):
    """Bernstein polynomial class for trajectory generation.

    Allows the user to construct a Bernstein (Bezier) polynomial of arbitrary
    dimension and degree.

    NOTE: cpts must be a float type.

    :param cpts: Control points used to define the Bernstein polynomial. The
        degree of the polynomial is equal to the number of columns -1. The
        dimension of the polynomial is equal to the number of rows.
    :type cpts: numpy.ndarray or None
    :param t0: Initial time of the Bernstein polynomial.
    :type t0: float
    :param tf: Final time of the Bernstein polynomial.
    :type tf: float

    """

    def __init__(self, cpts=None, t0=0.0, tf=1.0):
        super().__init__(cpts=cpts, t0=t0, tf=tf)
        # self.tau = np.linspace(t0, tf, 1001)

    def __add__(self, curve):
        return self.add(curve)

    def __sub__(self, curve):
        return self.sub(curve)

    def __mul__(self, curve):
        return self.mul(curve)

    def __truediv__(self, curve):
        return self.div(curve)

    def __pow__(self, power):
        raise Exception('Power is not yet implemented.')

    def __repr__(self):
        return 'Bernstein({}, {}, {})'.format(self.cpts, self.t0, self.tf)

    def __call__(self, t):
        """Calling the object returns the values of the curve at the t values

        Note that unlike the curve property, this will NOT cache the computed
        values. This is meant to be a convenience function to quickly peek at
        the values of the curve.

        :param t: Single value or numpy array of values at which to compute the
            curve.
        :type t: float or numpy.ndarray
        """
        tau = np.atleast_1d(t)
        curve = np.empty((self.dim, tau.size))
        for i, pts in enumerate(self.cpts):
            curve[i] = deCasteljauCurve(pts, tau, self.t0, self.tf)

        return curve

    @property
    def tau(self):
        if self._tau is None:
            self._tau = np.linspace(self.t0, self.tf, 1001)
        return self._tau

    @tau.setter
    def tau(self, val):
        self._tau = val

    @property
    def x(self):
        """Convenience function to return only the X dimension of the curve

        Returns a Bernstein object whose control points are the 0th row of the
        original object's control points.
        """
        return Bernstein(self.cpts[0], t0=self.t0, tf=self.tf)

    @property
    def y(self):
        """Convenience function to return only the Y dimension of the curve

        Returns a Bernstein object whose control points are the 1st row of the
        original object's control points. If the original object is less than
        2 dimensions, this returns None.
        """
        if self.dim > 1:
            return Bernstein(self.cpts[1], t0=self.t0, tf=self.tf)
        else:
            return None

    @property
    def z(self):
        """Convenience function to return only the Z dimension of the curve

        Returns a Bernstein object whose control points are the 2nd row of the
        original object's control points. If the original object is less than
        3 dimensions, this returns None.
        """
        if self.dim > 2:
            return Bernstein(self.cpts[2], t0=self.t0, tf=self.tf)
        else:
            return None

    @property
    def curve(self):
        """Returns the curve computed at each value of Tau.

        This function will use the De Casteljau algorithm to find the value of
        the curve at each value found within Tau \in [t0, tf]. Note that by
        default Tau is computed by np.linspace(t0, tf, 1001). This function
        will also cache the curve for future use to avoid computing the curve
        multiple times.

        The computation of the curve uses the De Casteljau algorithm rather
        than the definition of Bernstein polynomials due to roundoff errors for
        high order curves. While Bernstein polynomials exhibit numerical
        stability, Python will not properly handle values to the power of large
        numbers.
        """
        if self._curve is None:
            self._curve = np.zeros([self.dim, len(self.tau)])
            for i, pts in enumerate(self.cpts):
                self._curve[i] = deCasteljauCurve(pts, self.tau, self.t0,
                                                  self.tf)

        return self._curve

    def copy(self):
        """Creates an exact, deep copy of the current Bernstein object

        :return: Deep copy of Bernstein object
        :rtype: Bernstein
        """
        return Bernstein(self.cpts, self.t0, self.tf)

    def plot(self, axisHandle=None, showCpts=True, **kwargs):
        """Plots the Bernstein polynomial in 1D or 2D

        Note: Currently only supports plotting in 1D or 2D.

        This is a convenience function to quickly plot Bernstein polynomials.
        It is simply a wrapper around pyplot's plot function. For advanced
        plots, the user should use the data from the Bernstein object's curve
        attribute

        :param axisHandle: Handle to the figure axis. If it is None, a new
            figure will be plotted.
        :type axisHandle: matplotlib.axes._subplots.AxesSubplot or None
        :param showCpts: Flag that decides whether to show the control points
            in the plot. Default is True.
        :type showCpts: bool
        :param **kwargs: Keyword arguments passed into the plot command. Note
            that the arguments are only passed into the plot command that
            plots the curve and not the command that plots the control points.
        :type **kwargs: dict
        :return: Axis object where the curve was plotted.
        :rtype: matplotlib.axes._subplots.AxesSubplot
        """
        if axisHandle is None:
            fig, ax = plt.subplots()
        else:
            ax = axisHandle

        cpts = np.asarray(self.cpts)

        if self.dim == 1:
            ax.plot(self.tau, self.curve[0], **kwargs)
            if showCpts:
                ax.plot(np.linspace(self.t0, self.tf, self.deg+1),
                        self.cpts.squeeze(), '.--')
        elif self.dim == 2 or (self.cpts[2, :] == 0).all():
            ax.plot(self.curve[0], self.curve[1], **kwargs)
            if showCpts:
                ax.plot(cpts[0], cpts[1], '.--')
        else:
            # Check whether ax is 3D
            if not hasattr(ax, 'get_zlim'):
                parent = ax.get_figure()
                ax.remove()
                ax = parent.add_subplot(111, projection='3d')
            ax.plot(self.curve[0], self.curve[1], self.curve[2], **kwargs)
            if showCpts:
                ax.plot(self.cpts[0], self.cpts[1], self.cpts[2], '.--')

        return ax

    def add(self, other):
        """Adds two Bernstein polynomials

        This function will automatically check to make sure that the initial
        and final times are aligned.

        Paper Reference: Property 7: Arithmetic Operations

        :param other: Other Bernstein polynomial to be added
        :type other: Bernstein
        :return: Sum of the two Bernstein polynomials, or None if their times
            do not overlap.
        :rtype: Bernstein
        """
        if self.t0 == other.t0 and self.tf == other.tf:
            cpts = self.cpts + other.cpts
            t0 = self.t0
            tf = self.tf
        else:
            c1, c2 = _temporalAlignment(self, other)
            cpts = c1.cpts + c2.cpts
            t0 = c1.t0
            tf = c1.tf

        if t0 >= tf:
            return None
        else:
            return Bernstein(cpts, t0=t0, tf=tf)

    def sub(self, other):
        """Subtracts two Bernstein polynomials

        This function will automatically check to make sure that the initial
        and final times are aligned.

        Paper Reference: Property 7: Arithmetic Operations

        :param other: Bernstein polynomial to subtract from the original
        :type other: Bernstein
        :return: Original curve - Other curve, or None if their times do not
            overlap.
        :rtype: Bernstein
        """
        if self.t0 == other.t0 and self.tf == other.tf:
            cpts = self.cpts - other.cpts
            t0 = self.t0
            tf = self.tf
        else:
            c1, c2 = _temporalAlignment(self, other)
            cpts = c1.cpts - c2.cpts
            t0 = c1.t0
            tf = c1.tf

        if t0 >= tf:
            return None
        else:
            return Bernstein(cpts, t0=t0, tf=tf)

    def mul(self, multiplicand):
        """Computes the product of two Bernstein polynomials.

        Paper Reference: Property 7: Arithmetic Operations

        Source: Section 5.1 of "The Bernstein Polynomial Basis: A Centennial
        Retrospective" by Farouki.

        :param multiplicand: Multiplicand
        :type multiplicand: Bernstein
        :return: Product of the two curve
        :rtype: Bernstein
        """
        if not isinstance(multiplicand, Bernstein):
            msg = 'The multiplicand must be a {} object, not a {}'.format(
                    Bernstein, type(multiplicand))
            raise TypeError(msg)

        dim = self.dim
        if multiplicand.dim != dim:
            msg = ('The dimension of both Bernstein polynomials must be the'
                   'same. '
                   'The first dimension is {} and the second is {}'.format(
                           dim, multiplicand.dim))
            raise ValueError(msg)

        a = np.array(self.cpts, ndmin=2)
        b = np.array(multiplicand.cpts, ndmin=2)
        m = self.deg
        n = multiplicand.deg

        c = np.empty((dim, m+n+1))

        try:
            coefMat = Bernstein.prodMatCache[m][n]
        except KeyError:
            coefMat = bezProductCoefficients(m, n)
            Bernstein.prodMatCache[m][n] = coefMat

        for d in range(dim):
            c[d] = multiplyBezCurves(a[d], b[d], coefMat)

#   This code uses Farouki's method for multiplication but does not simplify
#   the problem using matrices.
#        for d in range(dim):
#            for k in np.arange(0, m+n+1):
#                summation = 0
#                for j in np.arange(max(0, k-n), min(m, k)+1):
#                    summation += binom(m, j)  \
#                                 * binom(n, k-j)  \
#                                 * a[d, j]  \
#                                 * b[d, k-j]
#                c[d, k] = summation / binom(m+n, k)

        newCurve = self.copy()
        newCurve.cpts = c

        return newCurve

    def div(self, denominator):
        """Divides one Bernstein polynomial by another.

        The division of two Bernstein polynomials results in a rational
        Bernstein polynomial.

        Paper Reference: Property 7: Arithmetic Operations

        :param denominator: Denominator of the division
        :type denominator: Bernstein
        :return: Rational Bernstein polynomial representing the division of
            the two polynomials.
        :rtype: RationalBernstein
        """
        if not isinstance(denominator, Bernstein):
            msg = ('The denominator must be a Bernstein object, not a {}. '
                   'Or the module has been reloaded.').format(
                           type(denominator))
            raise TypeError(msg)

        cpts = np.empty((self.dim, self.deg+1))
        for i in range(self.dim):
            for j in range(self.deg+1):
                if self.cpts[i, j] == 0:
                    cpts[i, j] = 0
                elif denominator.cpts[i, j] == 0:
                    cpts[i, j] = np.inf
                else:
                    cpts[i, j] = self.cpts[i, j] / denominator.cpts[i, j]

        weights = denominator.cpts

        return RationalBernstein(cpts.astype(np.float64),
                                 weights.astype(np.float64),
                                 t0=self.t0, tf=self.tf)

    def elev(self, R=1):
        """Elevates the degree of the Bernstein polynomial.

        Elevates the degree of the Bernstein polynomial by R (default is 1)
        and returns a new, higher degree Bernstein object.

        :param R: Number of degrees to elevate the curve
        :type R: int
        :return: Elevated Bernstein polynomial
        :rtype: Bernstein
        """
        if R < 1:
            return self

        try:
            elevMat = Bernstein.elevMatCache[self.deg][R]
        except KeyError:
            elevMat = elevMatrix(self.deg, R)
            Bernstein.elevMatCache[self.deg][R] = elevMat

        elevPts = []
        for cpts in self.cpts:
            elevPts.append([np.dot(cpts, elevMat)])

        elevPts = np.concatenate(elevPts)

        curveElev = self.copy()
        curveElev.cpts = elevPts

        return curveElev

    def diff(self):
        """Calculates the derivative of the Bernstein polynomial

        Note that this does not affect the object. Instead it returns the
        derivative.

        :return: Derivative of the Bernstein polynomial
        :rtype: Bernstein
        """
        try:
            Dm = Bernstein.diffMatCache[self.deg][self.tf-self.t0]
        except KeyError:
            Dm = diffMatrix(self.deg, self.t0, self.tf)
            Bernstein.diffMatCache[self.deg][self.tf-self.t0] = Dm

        cptsDot = []
        for i in range(self.dim):
            cptsDot.append(np.dot(self.cpts[i, :], Dm))

        curveDot = self.copy()
        curveDot.cpts = cptsDot

        return curveDot.elev()

    def integrate(self):
        """Calculates the area under the curve in each dimension

        :return: Area under the curve in each dimension.
        :rtype: numpy.ndarray
        """
        areas = np.empty(self.dim)
        for d in range(self.dim):
            areas[d] = self.tf * sum(self.cpts[d]) / (self.deg+1)

        return areas

    def split(self, tDiv):
        """Splits the curve into two curves at point tDiv

        Note that the two returned curves will have the SAME tf value as the
        original curve. This may result in slightly unexpected behavior for a
        1D curve when plotting since both slices of the original curve will
        also be plotted on [0, tf]. The behavior should work as expected when
        plotting in 2D though.

        Paper Reference: Property 5: The de Casteljau Algorithm

        :param tDiv: Point at which to split the curve
        :type tDiv: float
        :return: Tuple of curves. One before the split point and one after.
        :rtype: tuple(Bernstein, Bernstein)
        """
        # c1 = self.copy()
        # c2 = self.copy()

        cpts1 = []
        cpts2 = []

        if np.isnan(tDiv):
            print('[!] Warning, tDiv is {}, changing to 0.'.format(tDiv))
            tDiv = 0

        for d in range(self.dim):
            left, right = deCasteljauSplit(self.cpts[d, :], tDiv - self.t0,
                                           self.tf - self.t0)
            cpts1.append([left])
            cpts2.append([right[::-1]])

        # c1.cpts = cpts1
        # c1.t0 = self.t0
        # c1.tf = tDiv
        # c2.cpts = cpts2
        # c2.t0 = tDiv
        # c2.tf = self.tf

        c1 = Bernstein(cpts=np.concatenate(cpts1), t0=self.t0, tf=tDiv)
        c2 = Bernstein(cpts=np.concatenate(cpts2), t0=tDiv, tf=self.tf)

        return c1, c2

    def min(self, dim=0, globMin=np.inf, tol=1e-6):
        """Return the minimum value of the Bernstein polynomialin a single dimension.

        Finds the minimum value of the Bernstein polynomial. This is done by
        first checking the first and last control points since the first and
        last point lie on the curve. If the first or last control point is not
        the minimum value, the curve is split at the lowest control point. The
        new minimum value is then defined as the lowest control point of the
        two new curves. This continues until the difference between the new
        minimum and old minimum values is within the desired tolerance.

        :param dim: Which dimension to return the minimum of.
        :type dim: int
        :param tol: Tolerance of the minimum value.
        :type tol: float
        :param maxIter: Maximum number of iterations to search for the minimum.
        :type maxIter: int
        :return: Minimum value of the Bernstein polynomial. None if maximum
            iterations is met.
        :rtype: float or None
        """
        ub = self.cpts[dim, (0, -1)].min()
        if ub < globMin:
            globMin = ub

        minIdx = self.cpts[dim, :].argmin()
        lb = self.cpts[dim, minIdx]

        # Prune if the global min is less than the lower bound
        if globMin < lb:
            return globMin

        # If we are within the desired tolerance, return
        if ub - lb < tol:
            return globMin

        # Otherwise split and continue
        else:
            tdiv = (minIdx/self.deg)*(self.tf - self.t0) + self.t0
            c1, c2 = self.split(tdiv)
            c1min = c1.min(dim=dim, globMin=globMin, tol=tol)
            c2min = c2.min(dim=dim, globMin=globMin, tol=tol)

            return min(c1min, c2min)

    def max(self, dim=0, globMax=-np.inf, tol=1e-6):
        """Return the maximum value of the Bernstein polynomial in a single dimension.

        Finds the maximum value of the Bernstein polynomial. This is done by
        first checking the first and last control points since the first and
        last point lie on the curve. If the first or last control point is not
        the maximum value, the curve is split at the highest control point. The
        new maximum value is then defined as the highest control point of the
        two new curves. This continues until the difference between the new
        maximum and old maximum values is within the desired tolerance.

        :param dim: Which dimension to return the maximum of.
        :type dim: int
        :param tol: Tolerance of the maximum value.
        :type tol: float
        :param maxIter: Maximum number of iterations to search for the minimum.
        :type maxIter: int
        :return: Maximum value of the Bernstein polynomial. None if maximum
            iterations is met.
        :rtype: float or None
        """
        lb = self.cpts[dim, (0, -1)].max()
        if lb > globMax:
            globMax = lb

        maxIdx = self.cpts[dim, :].argmax()
        ub = self.cpts[dim, maxIdx]

        if ub - lb < tol:
            return globMax

        else:
            tdiv = (maxIdx/self.deg)*(self.tf - self.t0) + self.t0
            c1, c2 = self.split(tdiv)
            c1max = c1.max(dim=dim, globMax=globMax, tol=tol)
            c2max = c2.max(dim=dim, globMax=globMax, tol=tol)

            newMax = max(c1max, c2max)

        return newMax

    def minDist(self, otherCurve):
        """

        Parameters
        ----------
        otherCurve : Bernstein
            Bernstein object of the curve being compared.

        Raises
        ------
        ValueError
            Both curves must be either 2D or 3D only.

        Returns
        -------
        (float, float, float)
            (minimum distance, t1, t2) where t1 and t2 are the t values corresponding to the point at which the
            minimum distance occurs.

        """
        if (self.dim < 2 or self.dim > 3 or
                otherCurve.dim < 2 or otherCurve.dim > 3):
            err = ('Both curves must be either 2D or 3D, not {}D and {}D.'
                   ).format(self.dim, otherCurve.dim)
            raise ValueError(err)

        cpts1 = self.cpts.copy()
        cpts2 = otherCurve.cpts.copy()

        if cpts1.shape[0] < 3:
            cpts1 = np.concatenate([cpts1, np.zeros((1, cpts1.shape[1]))])

        if cpts2.shape[0] < 3:
            cpts2 = np.concatenate([cpts2, np.zeros((1, cpts2.shape[1]))])

        alpha, t1, t2 = _minDist(cpts1, cpts2)
        return alpha, self.t0 + t1*(self.tf-self.t0), otherCurve.t0 + t2*(otherCurve.tf - otherCurve.t0)

    def minDist2Poly(self, poly):
        """
        """
        return _minDist2Poly(self, poly)

    def collCheck(self, otherCurve):
        """
        """
        if (self.dim < 2 or self.dim > 3 or
                otherCurve.dim < 2 or otherCurve.dim > 3):
            err = ('Both curves must be either 2D or 3D, not {}D and {}D.'
                   ).format(self.dim, otherCurve.dim)
            raise ValueError(err)

        cpts1 = self.cpts.copy()
        cpts2 = otherCurve.cpts.copy()

        if cpts1.shape[0] < 3:
            cpts1 = np.concatenate([cpts1, np.zeros((1, cpts1.shape[1]))])

        if cpts2.shape[0] < 3:
            cpts2 = np.concatenate([cpts2, np.zeros((1, cpts2.shape[1]))])

        res = _collCheckBez2Bez(cpts1, cpts2)

        if res == -1:
            res = 1

        return res

    def collCheck2Poly(self, poly):
        """
        """
        return _collCheckBez2Poly(self, poly)

    def normSquare(self):
        """Calculates the norm squared of the Bernstein polynomial

        Returns a Bernstein object for the norm squared result of the current
        Bernstein polynomial.

        :return: Norm squared of the Bernstein polynomial
        :rtype: Bernstein
        """
        try:
            prodM = Bernstein.prodMatCache[self.deg][self.deg]
        except KeyError:
            prodM = prodMatrix(self.deg).T
            Bernstein.prodMatCache[self.deg][self.deg] = prodM

        normCpts = _normSquare(self.cpts, 1, self.dim, prodM.T)/2

        newCurve = self.copy()
        newCurve.cpts = normCpts

        return newCurve
#        return Bezier(_normSquare(self.cpts, 1, self.dim, prodM.T),
#                      tau=self.tau, tf=self.tf)

def _temporalAlignment(c1, c2):
    """Returns two curves that are temporally aligned (i.e. same t0 and tf)

    :param c1: First curve to be aligned
    :type c1: Bernstein
    :param c2: Second curve to be aligned
    :type c2: Bernstein
    :return: Two curves whose control points are defined for the intersection
        of time of c1 and c2. (i.e. t0 = max(c1.t0, c2.t0),
        tf = min(c1.tf, c2.tf))
    :rtype: (Bernstein, Bernstein)
    """
    newC1 = c1.copy()
    newC2 = c2.copy()

    if c1.t0 < c2.t0:
        t0 = c2.t0
        _, newC1 = newC1.split(t0)
    elif c1.t0 > c2.t0:
        t0 = c1.t0
        _, newC2 = newC2.split(t0)
    else:
        t0 = c1.t0

    if c1.tf < c2.tf:
        tf = c1.tf
        newC2, _ = newC2.split(tf)
    elif c1.tf > c2.tf:
        tf = c2.tf
        newC1, _ = newC1.split(tf)
    else:
        tf = c1.tf

    newC1.t0 = t0
    newC2.t0 = t0
    newC1.tf = tf
    newC2.tf = tf

    return newC1, newC2

@njit(cache=True)
def deCasteljauCurve(cpts, tau, t0=0.0, tf=1.0):
    """Returns a Bernstein polynomial using the de Casteljau algorithm

    Uses the de Casteljau algorithm to generate the Bernstein polynomial
    defined by the provided control points. Note that the datatypes are
    important due to the nature of the numba library.

    Paper Reference: Property 5: The de Casteljau Algorithm

    :param cpts: Control points defining the 1D Bernstein polynomial.
    :type cpts: numpy.ndarray(dtype=numpy.float64)
    :param tau: Values at which to evaluate Bernstein polynomial.
    :type tau: numpy.ndarray(dtype=numpy.float64)
    :param t0: Initial time of the curve. Default is 0.0.
    :type t0: float
    :param tf: Final time of the curve. Default is 1.0.
    :type tf: float
    :return: Numpy array of length tau of the Bernstein polynomial evaluated
        at each value of tau.
    :rtype: numpy.ndarray(dtype=numpy.float64)
    """
    T = (tau - t0) / (tf - t0)

    curveLen = T.size
    curve = np.empty(curveLen)
    curveIdx = 0

    for t in T:
        newCpts = cpts.copy()
        while newCpts.size > 1:
            cptsTemp = np.empty(newCpts.size-1)
            for i in range(cptsTemp.size):
                cptsTemp[i] = (1-t)*newCpts[i] + t*newCpts[i+1]
            newCpts = cptsTemp.copy()
        curve[curveIdx] = newCpts[0]
        curveIdx += 1

    return curve

@njit(cache=True)
def deCasteljauSplit(cpts, tDiv, tf=1.0):
    """Uses the de Casteljau algorithm to split the curve at tDiv

    This function is similar to the de Casteljau curve function but instead of
    drawing a curve, it returns two sets of control points which define the
    curve to the left and to the right of the split point, tDiv.

    Paper Reference: Property 5: The de Casteljau Algorithm

    :param cpts: Control points defining the 1D Bernstein polynomial.
    :type cpts: numpy.ndarray(dtype=numpy.float64)
    :param tDiv: Point at which to divide the curve.
    :type tDiv: float
    :param tf: Final tau value for the 1D curve. Default is 1.0. Note that the
        Bernstein polynomial is defined on the range of [0, tf].
    :type tf: float
    :return: Returns a tuple of numpy arrays. The zeroth element is the
        control points defining the curve to the left of tDiv. The first
        element is the control points defining the curve to the right of tDiv.
    :rtype: tuple(numpy.ndarray, numpy.ndarray)
    """
    tDiv = tDiv/tf
    cptsLeft = np.zeros(cpts.size)
    cptsRight = np.zeros(cpts.size)
    idx = 0

    newCpts = cpts.copy()
    cptsTemp = cpts.copy()
    while newCpts.size > 1:
        cptsLeft[idx] = cptsTemp[0]
        cptsRight[idx] = cptsTemp[-1]
        idx += 1

        cptsTemp = np.empty(newCpts.size-1)
        for i in range(cptsTemp.size):
            cptsTemp[i] = (1-tDiv)*newCpts[i] + tDiv*newCpts[i+1]

        newCpts = cptsTemp.copy()

    cptsLeft[-1] = cptsRight[-1] = newCpts[0]

    return cptsLeft, cptsRight

def bezierCurve(cpts, tau, tf=1.0):
    """Computes the values of a 1D Bezier curve defined by the control points.

    Creates a 1 dimensional Bezier curve using the designated control points
    and values of tau.

    Effectively evaluates the following expression:
        T*M*P
    Where
    T is the power basis vector [1 t t^2 t^3 ... t^N]
    M is the binomial matrix (more information found in buildBezMatrix)
    P is a vector of Bezier weights (i.e. control points)

    :param cpts: Single row matrix of N+1 control points for
        a one dimensional Bezier curve.
    :type cpts: numpy.ndarray
    :param tau: Values at which to evaluate Bezier curve. Should
        typically only be on the range of [0,tf] but it should work if it's not
        on that range.
    :type tau: numpy.ndarray
    :return: Numpy array of length tau of the Bezier curve evaluated at each
        value of tau.
    :rtype: numpy.ndarray
    """
    cpts = np.array(cpts)
    tau = np.array(tau, dtype=np.float64, ndmin=1)/float(tf)
    tauLen = tau.size
    n = cpts.size-1
    curve = np.empty(tauLen)

    coeffs = buildBezMatrix(n)
    for i, t in enumerate(tau):
        powerBasis = np.power(t, range(n+1))
        curve[i] = np.dot(powerBasis, np.dot(coeffs, cpts.T))

    return curve

@jit(cache=True, forceobj=True)
def buildBezMatrix(n):
    """Builds a matrix of coefficients of the power basis to a Bernstein
    polynomial.

    The coefficients matrix allows us to represent a Bezier curve of arbitrary
    degree as a matrix. This is useful for computational efficiency and ease of
    calculations.

    Taken from the source:
    "Since the power basis {1, t, t^2, t^3, ...} forms a basis for the space of
     polynomials of degree less than or equal to n, any Bernstein polynomial of
     degree n can be written in terms of the power basis."

    Source:
        http://graphics.cs.ucdavis.edu/education/CAGDNotes/CAGDNotes/
        Bernstein-Polynomials/Bernstein-Polynomials.html#conversion

    :param n: Degree of the Bezier curve
    :type n: int
    :return: Power basis matrix of coefficients for a Bezier curve
    :rtype: numpy.ndarray
    """
    bezMatrix = np.zeros((n+1, n+1))

    for k in np.arange(0, n+1):
        for i in np.arange(k, n+1):
            bezMatrix[i, k] = (-1)**(i-k) * binom(n, i) * binom(i, k)

    return bezMatrix

@njit(cache=True)
def diffMatrix(n, t0=0.0, tf=1.0):
    """Generates the differentiation matrix to find the derivative

    Takes the derivative of the control points for a Bezier curve. The
    resulting control points can be used to construct a Bezier curve that is
    the derivative of the original.

    Paper Reference: Property 3: Derivatives

    :param n: Degree of the Bezier curve
    :type n: int
    :param tf: Final time for the Bezier trajectory
    :type tf: float
    :return: Differentiation matrix for a Bezier curve of degree n
    :rtype: numpy.ndarray
    """
    val = n / (tf - t0)
    Dm = np.zeros((n+1, n))
    for i in range(n):
        Dm[i, i] = -val
        Dm[i+1, i] = val

    return Dm

#@jit(cache=True, forceobj=True)
def elevMatrix(N, R=1):
    """Creates an elevation matrix for a Bezier curve.

    Creates a matrix to elevate a Bezier curve of degree N to degree N+R.
    The elevation is performed as such:
        B_(N)*T = B_(N+1) where * is the dot product.

    :param N: Degree of the Bezier curve being elevated
    :type N: int
    :param R: Number of degrees to elevate the Bezier curve
    :type R: int
    :return: Elevation matrix to raise a Bezier curve of degree N by R degrees
    :rtype: numpy.ndarray
    """
    T = np.zeros((N+1, N+R+1))
    for i in range(N+R+1):
        den = binom(N+R, i)
        for j in range(N+1):
            T[j, i] = binom(N, j) * binom(R, i-j) / den

    return T

#@jit(cache=True, forceobj=True)
def prodMatrix(N):
    """Produces a product matrix for obtaining the norm of a Bezier curve

    This function produces a matrix which can be used to compute ||x dot x||^2
    i.e. xaug = x'*x;
    xaug = reshape(xaug',[length(x)^2,1]);
    y = Prod_T*xaug;
    or simply norm_square(x)
    prodM is the coefficient of bezier multiplication.

    Code ported over from Venanzio Cichella's MATLAB Prod_Matrix function.

    :param N: Degree of the Bezier curve
    :type N: int
    :return: Product matrix
    :rtype: numpy.ndarray
    """
    T = np.zeros((2*N+1, (N+1)**2))

    for j in np.arange(2*N+1):
        den = binom(2*N, j)
        for i in np.arange(max(0, j-N), min(N, j)+1):
            if N >= i and N >= j-i and 2*N >= j and j-i >= 0:
                T[j, N*i+j] = binom(N, i)*binom(N, j-i) / den

    return T

# TODO:
#    Change this function name to prodM.
#    Clean up and slightly change _normSquare to accommodate this change
#@jit(cache=True, forceobj=True)
def bezProductCoefficients(m, n=None):
    """Produces a product matrix for obtaining the product of two Bezier curves

    This function computes the matrix of coefficients for multiplying two
    Bezier curves. This function exists so that the coefficients matrix can be
    computed ahead of time when performing many multiplications.

    :param m: Degree of the first Bezier curve
    :type m: int
    :param n: Degree of the second Bezier curve
    :type n: int
    :return: Product matrix
    :rtype: numpy.ndarray
    """

    if n is None:
        n = m

    coefMat = np.zeros(((m+1)*(n+1), m+n+1))

    for k in range(m+n+1):
        den = binom(m+n, k)
        for j in range(max(0, k-n), min(m, k)+1):
            coefMat[m*j+k, k] = binom(m, j)*binom(n, k-j)/den

    return coefMat

@jit(cache=True, forceobj=True)
def multiplyBezCurves(multiplier, multiplicand, coefMat=None):
    """Multiplies two Bezier curves together

    The product of two Bezier curves can be computed directly from their
    control points. This function specifically uses matrix multiplication to
    increase the speed of the multiplication.

    Note that this function is made for 1D control points.

    One can pass in a matrix of product coefficients to compute the product
    about 10x faster. It is recommended that the coefficient matrix is
    precomputed and saved in memory when performing multiplication many times.
    The function bezProductCoefficients will produce this matrix.

    :param multiplier: Control points of the multiplier curve. Single dimension
    :type multiplier: numpy.ndarray
    :param multiplicand: Control points of the multiplicand curve.
    :type multiplicand: numpy.ndarray
    :param coefMat: Precomputed coefficient matrix from bezProductCoefficients
    :type coefMat: numpy.ndarray or None
    :return: Product of two Bezier curves
    :rtype: numpy.ndarray
    """
    multiplier = np.atleast_2d(multiplier)
    multiplicand = np.atleast_2d(multiplicand)
    m = multiplier.shape[1] - 1
    n = multiplicand.shape[1] - 1

    augMat = np.dot(multiplier.T, multiplicand)
    newMat = augMat.reshape((1, -1))

    if coefMat is None:
        coefMat = bezProductCoefficients(m, n)

    return np.dot(newMat, coefMat)

@jit(cache=True, forceobj=True)
def splitCurveMat(deg, z, coefMat=None):
    """Creates matrices Q and Qp that are used to compute control points for a
        split curve.

    :param deg: Degree of the Bezier curve
    :type deg: int
    :param z: Point on the curve at which to split the curve [0,1]
    :type z: float
    :param coefMat: Matrix of Binomial coefficients for a Bezier curve. Passing
        in a precomputed matrix will significantly increase the speed of the
        function.
    :type coefMat: numpy.ndarray or None
    :return: Returns a tuple of matrices. The zeroth element is the Q matrix
        for computing control points before the point z and the first element
        is the Q matrix for computing after the point z.
    :rtype: tuple(numpy.ndarray, numpy.ndarray)
    """
    powMat = np.diag(np.power(z, range(deg+1)))

    if coefMat is None:
        coefMat = buildBezMatrix(deg)

    # Q = M^-1 * Z * M
    Q = np.linalg.inv(coefMat).dot(powMat).dot(coefMat)

    # Qp is just Q but rolled and flipped
    Qp = np.empty((deg+1, deg+1))
    for i, row in enumerate(Q):
        Qp[deg-i, :] = np.roll(row, deg-i)

    return Q, Qp

def _minDist(cpts1, cpts2, count=0, alpha=np.inf, eps=1e-9, t1_l=0., t1_h=1., t2_l=0., t2_h=1.):
    """
    Find the minimum distance between two Bernstein polynomials given their control points.

    TODO Documentation

    Parameters
    ----------
    cpts1 : TYPE
        DESCRIPTION.
    cpts2 : TYPE
        DESCRIPTION.
    cnt : TYPE, optional
        DESCRIPTION. The default is 0.
    alpha : TYPE, optional
        DESCRIPTION. The default is np.inf.
    eps : TYPE, optional
        DESCRIPTION. The default is 1e-9.
    t1_l : TYPE, optional
        DESCRIPTION. The default is 0.
    t1_h : TYPE, optional
        DESCRIPTION. The default is 1.
    t2_l : TYPE, optional
        DESCRIPTION. The default is 0.
    t2_h : TYPE, optional
        DESCRIPTION. The default is 1.

    Returns
    -------
    alpha : TYPE
        DESCRIPTION.

    """
    t1len = t1_h - t1_l
    t2len = t2_h - t2_l

    lb = pygjk(cpts1.T, cpts2.T)
    if lb > alpha:
        return np.inf, -1, -1

    ub, t1, t2 = _upperbound(cpts1, cpts2)
    tdiv1 = t1_l + t1len*t1
    tdiv2 = t2_l + t2len*t2

    if ub < alpha:
        alpha = ub

    if alpha - lb < eps:
        return alpha, tdiv1, tdiv2

    else:
        c1L = []
        c1R = []
        for row in cpts1:
            left, right = deCasteljauSplit(row, 0.5)
            c1L.append(left)
            c1R.append(right[::-1])
        c1L = np.array(c1L, dtype=float)
        c1R = np.array(c1R, dtype=float)

        c2L = []
        c2R = []
        for row in cpts2:
            left, right = deCasteljauSplit(row, 0.5)
            c2L.append(left)
            c2R.append(right[::-1])
        c2L = np.array(c2L, dtype=float)
        c2R = np.array(c2R, dtype=float)

        # c1L, c1R = deCasteljauSplit(cpts1, 0.5)
        # c1R = c1R[::-1]
        # c2L, c2R = deCasteljauSplit(cpts2, 0.5)
        # c2R = c2R[::-1]

        count += 1
        alphaNew, t1New, t2New = _minDist(c1L, c2L, count=count, alpha=alpha,
                                          t1_l=t1_l, t1_h=t1_l+0.5*t1len,
                                          t2_l=t2_l, t2_h=t2_l+0.5*t2len)
        if alphaNew < alpha:
            alpha = alphaNew
            tdiv1 = t1New
            tdiv2 = t2New

        alphaNew, t1New, t2New = _minDist(c1L, c2R, count=count, alpha=alpha,
                                          t1_l=t1_l, t1_h=t1_l+0.5*t1len,
                                          t2_l=t2_l+0.5*t2len, t2_h=t2_h)
        if alphaNew < alpha:
            alpha = alphaNew
            tdiv1 = t1New
            tdiv2 = t2New

        alphaNew, t1New, t2New = _minDist(c1R, c2L, count=count, alpha=alpha,
                                          t1_l=t1_l+0.5*t1len, t1_h=t1_h,
                                          t2_l=t2_l, t2_h=t2_l+0.5*t2len)
        if alphaNew < alpha:
            alpha = alphaNew
            tdiv1 = t1New
            tdiv2 = t2New

        alphaNew, t1New, t2New = _minDist(c1R, c2R, count=count, alpha=alpha,
                                          t1_l=t1_l+0.5*t1len, t1_h=t1_h,
                                          t2_l=t2_l+0.5*t2len, t2_h=t2_h)
        if alphaNew < alpha:
            alpha = alphaNew
            tdiv1 = t1New
            tdiv2 = t2New

        return alpha, tdiv1, tdiv2

# def _colCheck(cpts1, cpts2, count=0, alpha=np.inf, maxIter=10, eps=1e-9):
def _collCheckBez2Bez(cpts1, cpts2, count=0, alpha=np.inf, maxIter=10, eps=1e-9):
    """
    Find the minimum distance between two Bernstein polynomials given their control points.

    TODO Documentation

    Parameters
    ----------
    cpts1 : TYPE
        DESCRIPTION.
    cpts2 : TYPE
        DESCRIPTION.
    cnt : TYPE, optional
        DESCRIPTION. The default is 0.
    alpha : TYPE, optional
        DESCRIPTION. The default is np.inf.
    maxIter : TYPE, optional
        DESCRIPTION. The default is 10.

    Returns
    -------
    alpha : TYPE
        DESCRIPTION.

    """
    lb = pygjk(cpts1.T, cpts2.T)
    if lb > eps:
        # No collision
        return 0

    elif count > maxIter:
        # Possible collision
        return -1

    else:
        c1L = []
        c1R = []
        for row in cpts1:
            left, right = deCasteljauSplit(row, 0.5)
            c1L.append(left)
            c1R.append(right[::-1])
        c1L = np.array(c1L, dtype=float)
        c1R = np.array(c1R, dtype=float)

        c2L = []
        c2R = []
        for row in cpts2:
            left, right = deCasteljauSplit(row, 0.5)
            c2L.append(left)
            c2R.append(right[::-1])
        c2L = np.array(c2L, dtype=float)
        c2R = np.array(c2R, dtype=float)

        # c1L, c1R = deCasteljauSplit(cpts1, 0.5)
        # c1R = c1R[::-1]
        # c2L, c2R = deCasteljauSplit(cpts2, 0.5)
        # c2R = c2R[::-1]

        count += 1
        alphaNew = _collCheckBez2Bez(c1L, c2L, count=count, alpha=alpha, maxIter=maxIter, eps=eps)  # ,
                                          # t1_l=t1_l, t1_h=t1_l+0.5*t1len,
                                          # t2_l=t2_l, t2_h=t2_l+0.5*t2len)
        if alphaNew < alpha:
            alpha = alphaNew

        alphaNew = _collCheckBez2Bez(c1L, c2R, count=count, alpha=alpha, maxIter=maxIter, eps=eps)  # ,
                                          # t1_l=t1_l, t1_h=t1_l+0.5*t1len,
                                          # t2_l=t2_l+0.5*t2len, t2_h=t2_h)
        if alphaNew < alpha:
            alpha = alphaNew

        alphaNew = _collCheckBez2Bez(c1R, c2L, count=count, alpha=alpha, maxIter=maxIter, eps=eps)  # ,
                                          # t1_l=t1_l+0.5*t1len, t1_h=t1_h,
                                          # t2_l=t2_l, t2_h=t2_l+0.5*t2len)
        if alphaNew < alpha:
            alpha = alphaNew

        alphaNew = _collCheckBez2Bez(c1R, c2R, count=count, alpha=alpha, maxIter=maxIter, eps=eps)  # ,
                                          # t1_l=t1_l+0.5*t1len, t1_h=t1_h,
                                          # t2_l=t2_l+0.5*t2len, t2_h=t2_h)
        if alphaNew < alpha:
            alpha = alphaNew

        return alpha

def _minDist2Poly(c1, poly2, cnt=0, alpha=np.inf, eps=1e-6, t1_l=0, t1_h=1):
    """Similar to _minDist but finds the distance between a curve and a polygon
    """
    x1 = c1.cpts[0, :]
    y1 = c1.cpts[1, :]

    if c1.dim == 3:
        z1 = c1.cpts[2, :]
    else:
        z1 = [0]*x1.size

    c1 = Bernstein([x1, y1, z1])

    poly1 = np.array(tuple(zip(x1, y1, z1)))
    poly2 = poly2.astype(float)

    cnt += 1
    if cnt > 1000:
        return (-1, -1, -1)

    flag, info = gjkNew(poly1, poly2)
    if flag > 0:
        closest1 = info[0]
        closest2 = info[1]
        lb = info[2]

        # Check to see if the closest point on the shape is a control point
        p1idx = np.where((poly1 == closest1).all(axis=1))[0]
        if p1idx.size > 0:
            t1 = p1idx[0]/c1.deg

        # If the closest point is not a control point, find t by weighting all
        # the control points by their distance from the closest point
        else:
            eucDist1 = np.linalg.norm(closest1-c1.cpts.T, axis=1)
            N = eucDist1.size
            W = np.empty(N)
            for i in range(N):
                W[i] = 1 / (1 +
                            (eucDist1[i]/eucDist1[:i]).sum() +
                            (eucDist1[i]/eucDist1[i+1:]).sum())

            t1 = (W*range(N)/N).sum()

        # Upper bound check
        ub, t1local = _upperboundPoly(c1.cpts, closest2)
        if ub <= alpha:
            alpha = ub
            newT1 = (1-t1local)*t1_l + t1local*t1_h
        else:
            newT1 = -1

    else:
        t1 = 0.5
        newT1 = -1
        closest2 = -1
        lb = eps**3
        ub = np.inf

    t1len = t1_h - t1_l

    retval = (alpha, newT1, closest2)

    if lb >= alpha*(1-eps):
        return retval

    else:
        c3, c4 = c1.split(t1)

        newAlpha, newT1, closest2 = _minDist2Poly(c3, poly2, cnt=cnt,
                                                  alpha=retval[0],
                                                  t1_l=t1_l,
                                                  t1_h=t1_l+t1*t1len)

        if newAlpha < retval[0]:
            retval = (newAlpha, newT1, closest2)

        newAlpha, newT1, closest2 = _minDist2Poly(c4, poly2, cnt=cnt,
                                                  alpha=retval[0],
                                                  t1_l=t1_l+t1*t1len,
                                                  t1_h=t1_h)

        if newAlpha < retval[0]:
            retval = (newAlpha, newT1, closest2)

    return retval

@njit(cache=True)
def _upperbound(c1, c2):
    """
    Find the upper bound on the minimum distance between two Bernstein polynomials.

    Parameters
    ----------
    c1 : np.array
        Control points of the first Bernstein polynomial.
    c2 : np.array
        Control points of the second Bernstein polynomial.

    Returns
    -------
    minimum distance : float
        Upper bound on the minimum distance.
    t1 : float
        t value at which the minimum distance occurs on the first Bernstein polynomial.
    t2 : float
        t value at which the minimum distance occurs on the second Bernstein polynomial.

    """
    distances = np.empty(4)
    tvals = np.array([(0., 0.),
                      (0., 1.),
                      (1., 0.),
                      (1., 1.)])

    distances[0] = _norm(c1[:, 0] - c2[:, 0])
    distances[1] = _norm(c1[:, 0] - c2[:, -1])
    distances[2] = _norm(c1[:, -1] - c2[:, 0])
    distances[3] = _norm(c1[:, -1] - c2[:, -1])

    t1, t2 = tvals[distances.argmin()]

    return distances.min(), t1, t2

#@njit(cache=True)
#def _upperboundPoly(c1, poly):
#    """
#    """
#    distances = np.empty(2*poly.shape[0])
#    tvals = np.array((0., 1.))
#
#    for i in range(poly.shape[0]):
#        distances[2*i] = _norm(c1[:, 0] - poly[i])
#        distances[2*i+1] = _norm(c1[:, -1] - poly[i])
#
#    t1 = tvals[distances.argmin() % 2]
#
#    return distances.min(), t1

@njit(cache=True)
def _upperboundPoly(c1, pt):
    """
    """
    distances = np.empty(2)
    tvals = np.array((0., 1.))

    distances[0] = _norm(c1[:, 0] - pt)
    distances[1] = _norm(c1[:, -1] - pt)

    t1 = tvals[distances.argmin()]

    return distances.min(), t1

@njit(cache=True)
def _norm(x):
    """
    """
    summation = 0.
    for val in x:
        summation += val*val

    return np.sqrt(summation)

# def _collCheckBez2Bez(c1, c2, cnt=0, alpha=np.inf, eps=1e-9):
#     """
#     Source: Computation of the minimum distance between two Bezier
#     curves/surfaces
#     """
#     x1 = c1.cpts[0, :]
#     y1 = c1.cpts[1, :]
#     x2 = c2.cpts[0, :]
#     y2 = c2.cpts[1, :]

#     if c1.dim == 3:
#         z1 = c1.cpts[2, :]
#     else:
#         z1 = [0]*x1.size

#     if c2.dim == 3:
#         z2 = c2.cpts[2, :]
#     else:
#         z2 = [0]*x1.size

#     c1 = Bernstein([x1, y1, z1])
#     c2 = Bernstein([x2, y2, z2])

#     poly1 = np.array(tuple(zip(x1, y1, z1)))
#     poly2 = np.array(tuple(zip(x2, y2, z2)))

#     cnt += 1
#     if cnt > 100:
#         return -1

#     ub, t1local, t2local = _upperbound(c1.cpts, c2.cpts)

#     flag, info = gjkNew(poly1, poly2)
#     if flag > 0:
#         return 1
#     else:
#         t1 = 0.5
#         t2 = 0.5
#         lb = 0

#     if ub <= alpha:
#         alpha = ub

#     if lb >= alpha*(1-eps):
#         return alpha

#     else:
#         c3, c4 = c1.split(t1)
#         c5, c6 = c2.split(t2)
#         alpha = min(alpha, _collCheckBez2Bez(c3, c5, cnt=cnt, alpha=alpha))
#         alpha = min(alpha, _collCheckBez2Bez(c3, c6, cnt=cnt, alpha=alpha))
#         alpha = min(alpha, _collCheckBez2Bez(c4, c5, cnt=cnt, alpha=alpha))
#         alpha = min(alpha, _collCheckBez2Bez(c4, c6, cnt=cnt, alpha=alpha))

#     return alpha

def _collCheckBez2Poly(c1, poly2, cnt=0, alpha=np.inf, eps=1e-9):
    """
    Source: Computation of the minimum distance between two Bezier
    curves/surfaces
    """
    x1 = c1.cpts[0, :]
    y1 = c1.cpts[1, :]

    if c1.dim == 3:
        z1 = c1.cpts[2, :]
    else:
        z1 = [0]*x1.size

    c1 = Bernstein([x1, y1, z1])

    poly1 = np.array(tuple(zip(x1, y1, z1)))

    cnt += 1
    if cnt > 100:
        return -1

    flag, info = gjkNew(poly1, poly2)
    if flag > 0:
        return 1
    else:
#        t1 = 0.5
#        lb = eps

        c3, c4 = c1.split(0.5)
        if (_collCheckBez2Poly(c3, poly2, cnt=cnt) == 1 and
            _collCheckBez2Poly(c4, poly2, cnt=cnt) == 1):
            return 1

    return 0

#    ub = np.inf
#
#    flag1, info1 = gjkNew(poly2, np.atleast_2d(poly1[0]))
#    flag2, info2 = gjkNew(poly2, np.atleast_2d(poly1[-1]))
#
#    if flag1 > 0:
#        ub = info1[2]
#
#    if flag2 > 0 and info2[2] < ub:
#        ub = info2[2]
#
#    if ub <= alpha:
#        alpha = ub
#
#    if lb >= alpha*(1-eps):
#        return round(alpha, int(round(np.log(eps)/np.log(10))))
#
#    else:
#        c3, c4 = c1.split(t1)
#        alpha = min(alpha, _collCheckBez2Poly(c3, poly2, cnt=cnt, alpha=alpha))
#        alpha = min(alpha, _collCheckBez2Poly(c4, poly2, cnt=cnt, alpha=alpha))
#
#    return alpha

# TODO
#   * Find a fast implementation for the calculation of binomial coefficients
#     that doesn't break for large numbers. Try fastBinom for 70 choose 20 and
#     it returns 0 which it shouldn't.
# @numba.jit(nopython=True)
# def fastBinom(n, k):
#    """Quickly computes the binomial coefficients.
#
#    A fast way to calculate binomial coefficients by Andrew Dalke.
#    See http://stackoverflow.com/questions/3025162/
#        statistics-combinations-in-python
#
#    :param n: n portion of "n choose k"
#    :type n: int
#    :param k: k portion of "n choose k"
#    :type k: int
#    :return: binomial coefficient of n choose k
#    :rtype: int
#    """
#    if 0 <= k <= n:
#        ntok = 1
#        ktok = 1
#        for t in range(1, min(k, n-k) + 1):
#            ntok *= n
#            ktok *= t
#            n -= 1
#        return ntok // ktok
#    else:
#        return 0
#
#
# @numba.jit(nopython=True)
# def binomialCoefficient(n, k):
#    # since C(n, k) = C(n, n - k)
#    if(k > n - k):
#        k = n - k
#    # initialize result
#    res = 1
#    # Calculate value of
#    # [n * (n-1) *---* (n-k + 1)] / [k * (k-1) *----* 1]
#    for i in range(k):
#        res = res * (n - i)
#        res = res // (i + 1)
#    return res

@njit(cache=True)
def _normSquare(x, Nveh, Ndim, prodM):
    """Compute the control points of the square of the norm of a vector

    normSquare(x, Nveh, Ndim, prodM)

    INPUT: Ndim*Nveh by N matrix x = [x1,...,x_Nveh)], x_i in R^Ndim
    OUTPUT: control points of ||x_i||^2 ... Nveh by N matrix

    Code ported over from Venanzio Cichella's MATLAB norm_square function.
    NOTE: This only works on 1D or 2D matricies. It will fail for 3 or more.
    """
#    x = np.array(x)
#    if x.ndim == 1:
#        x = x[None]

    m, N = x.shape

    xsquare = np.zeros((m, prodM.shape[0]))

    for i in range(m):
#        xaug = np.dot(x[i, None].T, x[i, None])
        xaug = np.dot(x.T, x)
        xnew = xaug.reshape((N**2, 1))
        xsquare[i, :] = np.dot(prodM, xnew).T[0]

    S = np.zeros((Nveh, Nveh*Ndim))

    for i in range(Nveh):
        for j in range(Ndim):
            S[i, Ndim*i+j] = 1

    return np.dot(S, xsquare)

if __name__ == '__main__':
    plt.close('all')
    cpts1 = np.array([[0, 2, 4, 6, 8, 10],
                      [1, 0, 2, 3, 10, 3],
                      [0, 0, 0, 0, 0, 0]], dtype=float)
    cpts2 = np.array([[1, 3, 6, 8, 10, 12],
                      [6, 9, 10, 11, 8, 8],
                      [0, 0, 0, 0, 0, 0]], dtype=float)

    # Bernstein polynomials
    c1 = Bernstein(cpts1, t0=10, tf=20)
    c2 = Bernstein(cpts2, t0=10, tf=20)

    alpha, t1, t2 = _minDist(cpts1, cpts2)

    approxMin = np.inf
    approxT1 = 0
    approxT2 = 0
    for i, val in enumerate(c1.curve.T):
        for j, val2 in enumerate(c2.curve[:, i:].T):
            dist = np.linalg.norm(val - val2)
            if dist < approxMin:
                approxMin = dist
                approxT1 = i / 1001
                approxT2 = i+j / 1001

    print(f'Alpha: {alpha}, t1: {t1}, t2: {t2}')
    print(f'Approx: {approxMin}, t1: {approxT1}, t2: {approxT2}')

    ax = c1.plot()
    c2.plot(ax)
    pt1 = c1(10+10*t1)
    pt2 = c2(10+10*t2)
    ax.plot([pt1[0], pt2[0]], [pt1[1], pt2[1]])
    plt.show()


# Write a Python program to count the number of characters (characterfrequency) in a string.
# Sample String : google.com

def no_of_each_char(str):
    dict={}
    for n in str:
        keys=dict.keys()
        print(keys)
        if n in keys:
            dict[n] +=1
        else:
            dict[n] =1
    return dict
print(no_of_each_char("google.com"))


import json
import re
from typing import List, Optional, Any

import bs4
import numpy as np
import pandas as pd
import dask.dataframe as dd
from dask.diagnostics import ProgressBar
from sqlalchemy.engine.base import Engine

from scrc.preprocessors.extractors.abstract_extractor import AbstractExtractor
from scrc.utils.main_utils import clean_text, get_config
from scrc.utils.sql_select_utils import join_decision_on_parameter, where_decisionid_in_list, where_string_spider

class Cleaner(AbstractExtractor):
    """
    Cleans the different courts from unnecessary strings to improve further processing.

    IMPORTANT: Make sure you have enough free RAM to run this (especially to clean the bigger spiders like BGer)

    For the biggest courts by size, special cleaning regexes (for pdfs)
    and cleaning functions (for htmls) have been developed for cleaner data.
    The regexes and functions have been tested on the basis of 3 random court decisions per court.

    Additional information to specific courts:
        - CH_BVGer:
            - sometimes contains spaces inside words: e.g. U r t e i l or B u n d e s v e r w a l t u n g s g e r i c h t
                => hard to fix because one does not know where the next word starts
        - ZH_Verwaltungsgericht:
            - contains an interesting summary at the top. But unfortunately it is quite hard to extract
    """

    def __init__(self, config: dict):
        super().__init__(config,
                         function_name='cleaning_functions',
                         col_name='text',
                         col_type='text')
        self.cleaning_regexes = self.load_cleaning_regexes(config['files']['cleaning_regexes'])
        self.processed_file_path = self.progress_dir / "spiders_cleaned.txt"

        self.logger_info = {
            'start': 'Started cleaning decisions',
            'finished': 'Finished cleaning decisions',
            'start_spider': 'Started cleaning decisions for spider',
            'finish_spider': 'Finished cleaning decisions for spider',
            'saving': 'Saving chunk of cleaned decisions',
            'processing_one': 'Cleaning the decision',
            'no_functions': 'Not cleaning the decision.'
        }

    def get_coverage(self, spider: str):
        self.logger.info('no coverage function implemented')

    def load_cleaning_regexes(self, file_name):
        """loads the cleaning regexes used for pdf files"""
        regex_file = self.spider_specific_dir / file_name  # mainly used for pdf spiders
        with open(regex_file) as f:
            cleaning_regexes = json.load(f)
        return cleaning_regexes

    def clean(self, decision_ids: Optional[List] = None):
        """cleans all the raw court rulings with the defined regexes (for pdfs) and functions (for htmls)"""
        self.logger.info("Started cleaning raw court rulings")
        processed_file_path = self.progress_dir / "spiders_cleaned.txt"

        if self.rebuild_entire_database:
            processed_file_path.unlink()
        if decision_ids is not None:
            self.decision_ids = decision_ids
        spider_list, message = self.compute_remaining_spiders(processed_file_path)
        self.logger.info(message)

    def select_df(self, engine: Engine, spider: str):
        """Returns the `where` clause of the select statement for the entries to be processed by extractor"""
        only_given_decision_ids_string = f" AND {where_decisionid_in_list(self.decision_ids)}" if self.decision_ids is not None else ""

        tables = f"file {join_decision_on_parameter('file_id', 'file.file_id')} LEFT JOIN section ON section.decision_id = decision.decision_id"
        columns = f"decision.decision_id,file.file_id,file_name,html_raw,pdf_raw,'{spider}' as spider, section_id"
        where_string = f"file.file_id IN {where_string_spider('file_id', spider)} {only_given_decision_ids_string}"
        return self.select(engine, tables, columns, where=where_string, chunksize=self.chunksize)

    def get_required_data(self, series: pd.DataFrame) -> Any:
        return series['spider']

    def check_condition_before_process(self, spider: str, data: Any, namespace: dict) -> bool:
        """Override if data has to conform to a certain condition before processing.
        e.g. data is required to be present for analysis"""
        return True

    def save_data_to_database(self, df: pd.DataFrame, engine: Engine):
        df['section_type_id'] = 1 # Set the Section type to full text
        self.update(engine, df, 'section', ['section_text'], self.output_dir, index_name='section_id')

    def process_one_spider(self, engine, spider):
        """Cleans one spider csv file"""
        self.logger.info(f"Started cleaning {spider}")

        #self.start_progress(engine, spider, engine)
        dfs = self.select_df(engine, spider)  # stream dfs from the db
        for df in dfs:
            # according to docs you should aim for a partition size of 100MB
            df['section_text'] = ''
            ddf = dd.from_pandas(df, npartitions=self.num_cpus)
            # apply cleaning function to each row
            ddf = ddf.apply(self.process_one_df_row, axis='columns', meta=ddf)
            with ProgressBar():
                df = ddf.compute(scheduler='processes')
            self.save_data_to_database(df, engine)
        #    self.log_progress(self.chunksize)
        #self.log_coverage(engine, spider, lang)
        self.logger.info(f"{self.logger_info['finish_spider']} {spider}")

    def process_one_df_row(self, series):
        """Cleans one row of a raw df"""
        self.logger.debug(
            f"{self.logger_info['processing_one']} {series['file_name']}")
        namespace = series.reindex(['file_id', 'pdf_url', 'html_url']).to_dict()
        spider = self.get_required_data(series)

        html_clean, pdf_clean = '', ''

        html_raw = series['html_raw']
        if pd.notna(html_raw) and html_raw not in [None, '']:
            # Parses the html string with bs4 and returns the body content
            soup = bs4.BeautifulSoup(html_raw, "html.parser").find('body')
            assert soup
            html_clean = self.clean_html(spider, soup, namespace)

        pdf_raw = series['pdf_raw']
        if pd.notna(pdf_raw) and pdf_raw not in [None, '']:
            pdf_clean = self.clean_pdf(spider, pdf_raw, namespace)

        # Combine columns into one easy to use text column (prioritize html content if both exist)
        series['section_text'] = np.where(html_clean != '', html_clean, pdf_clean)
        if series['section_text'] == '':
            # series['text'] = np.nan  # set to nan, so that it can be removed afterwards
            self.logger.warning(f"No raw text available for court decision {series['file_id']}")

        # otherwise we get an error when we want to save it into the db
        series.section_text = str(series.section_text)

        return series

    def clean_pdf(self, spider: str, text: str, namespace: dict) -> str:
        """Cleans first the text first with court specific regexes and then with general ones"""
        cleaned_text = self.clean_with_regexes(spider, text, namespace)
        return clean_text(cleaned_text)

    def clean_html(self, spider: str, soup: bs4.BeautifulSoup, namespace: dict) -> str:
        """Cleans first the text first with court specific regexes and then with general ones"""
        cleaned_text = self.clean_with_functions(spider, soup, namespace)
        return clean_text(cleaned_text)

    def clean_with_functions(self, spider: str, soup: bs4.BeautifulSoup, namespace: dict) -> str:
        """Cleans html documents with cleaning functions"""
        if not hasattr(self.processing_functions, spider):
            self.logger.debug(f"There are no special functions for spider {spider}. Just performing default cleaning.")
        else:
            cleaning_function = getattr(self.processing_functions,
                                        spider)  # retrieve cleaning function by spider
            # invoke cleaning function with soup and namespace
            soup = cleaning_function(soup, namespace)

        # we cannot just remove tables because sometimes the content of the entire court decision is inside a table (GL_Omni)
        # for table in soup.find_all("table"):
        #    table.decompose()  # remove all tables from content because they are not useful in raw text
        return soup.get_text()

    def clean_with_regexes(self, spider: str, text: str, namespace: dict) -> str:
        """Cleans pdf documents with cleaning regexes"""
        if spider not in self.cleaning_regexes or not self.cleaning_regexes[spider]:
            self.logger.debug(f"There are no special regexes for spider {spider}. Just performing default cleaning.")
            return text
        else:
            regexes = self.cleaning_regexes[spider]
            assert regexes  # this should not be empty
            for regex in regexes:
                # these strings can be used in the patterns and will be replaced by the variable content
                # see: https://stackoverflow.com/questions/44757222/transform-string-to-f-string
                pattern = r'' + regex['pattern']
                for key in namespace.keys():
                    # IMPORTANT: regex quantifiers (e.g. {4}) could be replaced by string format. => check first
                    if key in pattern:  # only format when the key is in the pattern.
                        # add the namespace to the pattern
                        pattern.format(**namespace)
                # perform the replacement
                cleaned_text = re.sub(pattern, regex['replacement'], text)
            return cleaned_text

if __name__ == '__main__':
    config = get_config()

    cleaner = Cleaner(config)
    cleaner.start()


import fnmatch
import locale
import logging
import os
import stat
from io import BytesIO, StringIO
from typing import (
    TYPE_CHECKING,
    Callable,
    Dict,
    Iterable,
    List,
    Mapping,
    Optional,
    Tuple,
    Union,
)

from funcy import cached_property

from dvc.path_info import PathInfo
from dvc.progress import Tqdm
from dvc.scm.base import SCMError
from dvc.utils import relpath

from ..objects import GitObject
from .base import BaseGitBackend

if TYPE_CHECKING:
    from dvc.types import StrPath

    from ..objects import GitCommit

logger = logging.getLogger(__name__)

class DulwichObject(GitObject):
    def __init__(self, repo, name, mode, sha):
        self.repo = repo
        self._name = name
        self._mode = mode
        self.sha = sha

    def open(self, mode: str = "r", encoding: str = None):
        if not encoding:
            encoding = locale.getpreferredencoding(False)
        # NOTE: we didn't load the object before as Dulwich will also try to
        # load the contents of it into memory, which will slow down Trie
        # building considerably.
        obj = self.repo[self.sha]
        data = obj.as_raw_string()
        if mode == "rb":
            return BytesIO(data)
        return StringIO(data.decode(encoding))

    @property
    def name(self) -> str:
        return self._name

    @property
    def mode(self) -> int:
        return self._mode

    def scandir(self) -> Iterable["DulwichObject"]:
        tree = self.repo[self.sha]
        for entry in tree.iteritems():  # noqa: B301
            yield DulwichObject(
                self.repo, entry.path.decode(), entry.mode, entry.sha
            )

class DulwichBackend(BaseGitBackend):  # pylint:disable=abstract-method
    """Dulwich Git backend."""

    # Dulwich progress will return messages equivalent to git CLI,
    # our pbars should just display the messages as formatted by dulwich
    BAR_FMT_NOTOTAL = "{desc}{bar:b}|{postfix[info]} [{elapsed}]"

    def __init__(  # pylint:disable=W0231
        self, root_dir=os.curdir, search_parent_directories=True
    ):
        from dulwich.errors import NotGitRepository
        from dulwich.repo import Repo

        try:
            if search_parent_directories:
                self.repo = Repo.discover(start=root_dir)
            else:
                self.repo = Repo(root_dir)
        except NotGitRepository as exc:
            raise SCMError(f"{root_dir} is not a git repository") from exc

        self._submodules: Dict[str, "PathInfo"] = self._find_submodules()
        self._stashes: dict = {}

    def _find_submodules(self) -> Dict[str, "PathInfo"]:
        """Return dict mapping submodule names to submodule paths.

        Submodule paths will be relative to Git repo root.
        """
        from dulwich.config import ConfigFile, parse_submodules

        submodules: Dict[str, "PathInfo"] = {}
        config_path = os.path.join(self.root_dir, ".gitmodules")
        if os.path.isfile(config_path):
            config = ConfigFile.from_path(config_path)
            for path, _url, section in parse_submodules(config):
                submodules[os.fsdecode(section)] = PathInfo(os.fsdecode(path))
        return submodules

    def close(self):
        self.repo.close()

    @property
    def root_dir(self) -> str:
        return self.repo.path

    @staticmethod
    def clone(
        url: str,
        to_path: str,
        rev: Optional[str] = None,
        shallow_branch: Optional[str] = None,
    ):
        raise NotImplementedError

    @property
    def dir(self) -> str:
        return self.repo.commondir()

    def add(self, paths: Union[str, Iterable[str]], update=False):
        from dvc.utils.fs import walk_files

        assert paths or update

        if isinstance(paths, str):
            paths = [paths]

        if update and not paths:
            self.repo.stage(list(self.repo.open_index()))
            return

        files: List[bytes] = []
        for path in paths:
            if not os.path.isabs(path) and self._submodules:
                # NOTE: If path is inside a submodule, Dulwich expects the
                # staged paths to be relative to the submodule root (not the
                # parent git repo root). We append path to root_dir here so
                # that the result of relpath(path, root_dir) is actually the
                # path relative to the submodule root.
                path_info = PathInfo(path).relative_to(self.root_dir)
                for sm_path in self._submodules.values():
                    if path_info.isin(sm_path):
                        path = os.path.join(
                            self.root_dir, path_info.relative_to(sm_path)
                        )
                        break
            if os.path.isdir(path):
                files.extend(
                    os.fsencode(relpath(fpath, self.root_dir))
                    for fpath in walk_files(path)
                )
            else:
                files.append(os.fsencode(relpath(path, self.root_dir)))

        # NOTE: this doesn't check gitignore, same as GitPythonBackend.add
        if update:
            index = self.repo.open_index()
            if os.name == "nt":
                # NOTE: we need git/unix separator to compare against index
                # paths but repo.stage() expects to be called with OS paths
                self.repo.stage(
                    [
                        fname
                        for fname in files
                        if fname.replace(b"\\", b"/") in index
                    ]
                )
            else:
                self.repo.stage([fname for fname in files if fname in index])
        else:
            self.repo.stage(files)

    def commit(self, msg: str, no_verify: bool = False):
        from dulwich.errors import CommitError
        from dulwich.porcelain import commit
        from dulwich.repo import InvalidUserIdentity

        try:
            commit(self.root_dir, message=msg, no_verify=no_verify)
        except CommitError as exc:
            raise SCMError("Git commit failed") from exc
        except InvalidUserIdentity as exc:
            raise SCMError(
                "Git username and email must be configured"
            ) from exc

    def checkout(
        self,
        branch: str,
        create_new: Optional[bool] = False,
        force: bool = False,
        **kwargs,
    ):
        raise NotImplementedError

    def pull(self, **kwargs):
        raise NotImplementedError

    def push(self):
        raise NotImplementedError

    def branch(self, branch: str):
        from dulwich.porcelain import Error, branch_create

        try:
            branch_create(self.root_dir, branch)
        except Error as exc:
            raise SCMError(f"Failed to create branch '{branch}'") from exc

    def tag(self, tag: str):
        raise NotImplementedError

    def untracked_files(self) -> Iterable[str]:
        _staged, _unstaged, untracked = self.status()
        return untracked

    def is_tracked(self, path: str) -> bool:
        rel = PathInfo(path).relative_to(self.root_dir).as_posix().encode()
        rel_dir = rel + b"/"
        for path in self.repo.open_index():
            if path == rel or path.startswith(rel_dir):
                return True
        return False

    def is_dirty(self, untracked_files: bool = False) -> bool:
        staged, unstaged, untracked = self.status()
        return bool(staged or unstaged or (untracked_files and untracked))

    def active_branch(self) -> str:
        raise NotImplementedError

    def list_branches(self) -> Iterable[str]:
        raise NotImplementedError

    def list_tags(self) -> Iterable[str]:
        raise NotImplementedError

    def list_all_commits(self) -> Iterable[str]:
        raise NotImplementedError

    def get_tree_obj(self, rev: str, **kwargs) -> DulwichObject:
        from dulwich.objectspec import parse_tree

        tree = parse_tree(self.repo, rev)
        return DulwichObject(self.repo, ".", stat.S_IFDIR, tree.id)

    def get_rev(self) -> str:
        rev = self.get_ref("HEAD")
        if rev:
            return rev
        raise SCMError("Empty git repo")

    def resolve_rev(self, rev: str) -> str:
        raise NotImplementedError

    def resolve_commit(self, rev: str) -> "GitCommit":
        raise NotImplementedError

    def _get_stash(self, ref: str):
        from dulwich.stash import Stash as DulwichStash

        if ref not in self._stashes:
            self._stashes[ref] = DulwichStash(self.repo, ref=os.fsencode(ref))
        return self._stashes[ref]

    @cached_property
    def ignore_manager(self):
        from dulwich.ignore import IgnoreFilterManager

        return IgnoreFilterManager.from_repo(self.repo)

    def is_ignored(self, path: "StrPath") -> bool:
        # `is_ignored` returns `false` if excluded in `.gitignore` and
        # `None` if it's not mentioned at all. `True` if it is ignored.
        relative_path = relpath(path, self.root_dir)
        # if checking a directory, a trailing slash must be included
        if str(path)[-1] == os.sep:
            relative_path += os.sep
        return bool(self.ignore_manager.is_ignored(relative_path))

    def set_ref(
        self,
        name: str,
        new_ref: str,
        old_ref: Optional[str] = None,
        message: Optional[str] = None,
        symbolic: Optional[bool] = False,
    ):
        name_b = os.fsencode(name)
        new_ref_b = os.fsencode(new_ref)
        old_ref_b = os.fsencode(old_ref) if old_ref else None
        message_b = message.encode("utf-8") if message else None
        if symbolic:
            return self.repo.refs.set_symbolic_ref(
                name_b, new_ref_b, message=message_b
            )
        if not self.repo.refs.set_if_equals(
            name_b, old_ref_b, new_ref_b, message=message_b
        ):
            raise SCMError(f"Failed to set '{name}'")

    def get_ref(self, name, follow: bool = True) -> Optional[str]:
        from dulwich.refs import parse_symref_value

        name_b = os.fsencode(name)
        if follow:
            try:
                ref = self.repo.refs[name_b]
            except KeyError:
                ref = None
        else:
            ref = self.repo.refs.read_ref(name_b)
            try:
                if ref:
                    ref = parse_symref_value(ref)
            except ValueError:
                pass
        if ref:
            return os.fsdecode(ref)
        return None

    def remove_ref(self, name: str, old_ref: Optional[str] = None):
        name_b = name.encode("utf-8")
        old_ref_b = old_ref.encode("utf-8") if old_ref else None
        if not self.repo.refs.remove_if_equals(name_b, old_ref_b):
            raise SCMError(f"Failed to remove '{name}'")

    def iter_refs(self, base: Optional[str] = None):
        base_b = os.fsencode(base) if base else None
        for key in self.repo.refs.keys(base=base_b):
            if base:
                if base.endswith("/"):
                    base = base[:-1]
                yield "/".join([base, os.fsdecode(key)])
            else:
                yield os.fsdecode(key)

    def iter_remote_refs(self, url: str, base: Optional[str] = None):
        from dulwich.client import get_transport_and_path
        from dulwich.porcelain import get_remote_repo

        try:
            _remote, location = get_remote_repo(self.repo, url)
            client, path = get_transport_and_path(location)
        except Exception as exc:
            raise SCMError(
                f"'{url}' is not a valid Git remote or URL"
            ) from exc

        if base:
            yield from (
                os.fsdecode(ref)
                for ref in client.get_refs(path)
                if ref.startswith(os.fsencode(base))
            )
        else:
            yield from (os.fsdecode(ref) for ref in client.get_refs(path))

    def get_refs_containing(self, rev: str, pattern: Optional[str] = None):
        raise NotImplementedError

    def push_refspec(
        self,
        url: str,
        src: Optional[str],
        dest: str,
        force: bool = False,
        on_diverged: Optional[Callable[[str, str], bool]] = None,
    ):
        from dulwich.client import get_transport_and_path
        from dulwich.errors import NotGitRepository, SendPackError
        from dulwich.porcelain import (
            DivergedBranches,
            check_diverged,
            get_remote_repo,
        )

        dest_refs, values = self._push_dest_refs(src, dest)

        try:
            _remote, location = get_remote_repo(self.repo, url)
            client, path = get_transport_and_path(location)
        except Exception as exc:
            raise SCMError(
                f"'{url}' is not a valid Git remote or URL"
            ) from exc

        def update_refs(refs):
            new_refs = {}
            for ref, value in zip(dest_refs, values):
                if ref in refs:
                    local_sha = self.repo.refs[ref]
                    remote_sha = refs[ref]
                    try:
                        check_diverged(self.repo, remote_sha, local_sha)
                    except DivergedBranches:
                        if not force:
                            overwrite = False
                            if on_diverged:
                                overwrite = on_diverged(
                                    os.fsdecode(ref), os.fsdecode(remote_sha)
                                )
                            if not overwrite:
                                continue
                new_refs[ref] = value
            return new_refs

        try:
            with Tqdm(
                desc="Pushing git refs", bar_format=self.BAR_FMT_NOTOTAL
            ) as pbar:

                def progress(msg_b):
                    msg = msg_b.decode("ascii").strip()
                    pbar.update_msg(msg)
                    pbar.refresh()
                    logger.trace(msg)

                client.send_pack(
                    path,
                    update_refs,
                    self.repo.object_store.generate_pack_data,
                    progress=progress,
                )
        except (NotGitRepository, SendPackError) as exc:
            raise SCMError("Git failed to push '{src}' to '{url}'") from exc

    def _push_dest_refs(
        self, src: Optional[str], dest: str
    ) -> Tuple[Iterable[bytes], Iterable[bytes]]:
        from dulwich.objects import ZERO_SHA

        if src is not None and src.endswith("/"):
            src_b = os.fsencode(src)
            keys = self.repo.refs.subkeys(src_b)
            values = [self.repo.refs[b"".join([src_b, key])] for key in keys]
            dest_refs = [b"".join([os.fsencode(dest), key]) for key in keys]
        else:
            if src is None:
                values = [ZERO_SHA]
            else:
                values = [self.repo.refs[os.fsencode(src)]]
            dest_refs = [os.fsencode(dest)]
        return dest_refs, values

    def fetch_refspecs(
        self,
        url: str,
        refspecs: Iterable[str],
        force: Optional[bool] = False,
        on_diverged: Optional[Callable[[str, str], bool]] = None,
    ):
        from dulwich.client import get_transport_and_path
        from dulwich.objectspec import parse_reftuples
        from dulwich.porcelain import (
            DivergedBranches,
            check_diverged,
            get_remote_repo,
        )

        fetch_refs = []

        def determine_wants(remote_refs):
            fetch_refs.extend(
                parse_reftuples(
                    remote_refs,
                    self.repo.refs,
                    [os.fsencode(refspec) for refspec in refspecs],
                    force=force,
                )
            )
            return [
                remote_refs[lh]
                for (lh, _, _) in fetch_refs
                if remote_refs[lh] not in self.repo.object_store
            ]

        try:
            _remote, location = get_remote_repo(self.repo, url)
            client, path = get_transport_and_path(location)
        except Exception as exc:
            raise SCMError(
                f"'{url}' is not a valid Git remote or URL"
            ) from exc

        with Tqdm(
            desc="Fetching git refs", bar_format=self.BAR_FMT_NOTOTAL
        ) as pbar:

            def progress(msg_b):
                msg = msg_b.decode("ascii").strip()
                pbar.update_msg(msg)
                pbar.refresh()
                logger.trace(msg)

            fetch_result = client.fetch(
                path,
                self.repo,
                progress=progress,
                determine_wants=determine_wants,
            )
        for (lh, rh, _) in fetch_refs:
            try:
                if rh in self.repo.refs:
                    check_diverged(
                        self.repo, self.repo.refs[rh], fetch_result.refs[lh]
                    )
            except DivergedBranches:
                if not force:
                    overwrite = False
                    if on_diverged:
                        overwrite = on_diverged(
                            os.fsdecode(rh), os.fsdecode(fetch_result.refs[lh])
                        )
                    if not overwrite:
                        continue
            self.repo.refs[rh] = fetch_result.refs[lh]

    def _stash_iter(self, ref: str):
        stash = self._get_stash(ref)
        yield from stash.stashes()

    def _stash_push(
        self,
        ref: str,
        message: Optional[str] = None,
        include_untracked: Optional[bool] = False,
    ) -> Tuple[Optional[str], bool]:
        from dulwich.repo import InvalidUserIdentity

        from dvc.scm.git import Stash

        if include_untracked or ref == Stash.DEFAULT_STASH:
            # dulwich stash.push does not support include_untracked and does
            # not touch working tree
            raise NotImplementedError

        stash = self._get_stash(ref)
        message_b = message.encode("utf-8") if message else None
        try:
            rev = stash.push(message=message_b)
        except InvalidUserIdentity as exc:
            raise SCMError(
                "Git username and email must be configured"
            ) from exc
        return os.fsdecode(rev), True

    def _stash_apply(self, rev: str):
        raise NotImplementedError

    def _stash_drop(self, ref: str, index: int):
        from dvc.scm.git import Stash

        if ref == Stash.DEFAULT_STASH:
            raise NotImplementedError

        stash = self._get_stash(ref)
        try:
            stash.drop(index)
        except ValueError as exc:
            raise SCMError("Failed to drop stash entry") from exc

    def describe(
        self,
        rev: str,
        base: Optional[str] = None,
        match: Optional[str] = None,
        exclude: Optional[str] = None,
    ) -> Optional[str]:
        if not base:
            base = "refs/tags"
        for ref in self.iter_refs(base=base):
            if (match and not fnmatch.fnmatch(ref, match)) or (
                exclude and fnmatch.fnmatch(ref, exclude)
            ):
                continue
            if self.get_ref(ref, follow=False) == rev:
                return ref
        return None

    def diff(self, rev_a: str, rev_b: str, binary=False) -> str:
        from dulwich.patch import write_tree_diff

        commit_a = self.repo[os.fsencode(rev_a)]
        commit_b = self.repo[os.fsencode(rev_b)]

        buf = BytesIO()
        write_tree_diff(
            buf, self.repo.object_store, commit_a.tree, commit_b.tree
        )
        return buf.getvalue().decode("utf-8")

    def reset(self, hard: bool = False, paths: Iterable[str] = None):
        raise NotImplementedError

    def checkout_index(
        self,
        paths: Optional[Iterable[str]] = None,
        force: bool = False,
        ours: bool = False,
        theirs: bool = False,
    ):
        raise NotImplementedError

    def status(
        self, ignored: bool = False
    ) -> Tuple[Mapping[str, Iterable[str]], Iterable[str], Iterable[str]]:
        from dulwich.porcelain import status as git_status

        staged, unstaged, untracked = git_status(
            self.root_dir, ignored=ignored
        )
        return (
            {
                status: [os.fsdecode(name) for name in paths]
                for status, paths in staged.items()
                if paths
            },
            [os.fsdecode(name) for name in unstaged],
            [os.fsdecode(name) for name in untracked],
        )

    def _reset(self) -> None:
        self.__dict__.pop("ignore_manager", None)

    def merge(
        self,
        rev: str,
        commit: bool = True,
        msg: Optional[str] = None,
        squash: bool = False,
    ) -> Optional[str]:
        raise NotImplementedError


from textwrap import dedent
from nmodl.initial import initial_blk
from nmodl.comment import comments

def test_initial():
    fun = dedent('''
    INITIAL {
        tadj = q10^((celsius - temp)/(10 (degC))) : make all threads calculate tadj at initialization

        trates(v+vshift)
        m = minf
        h = hinf
    }
    ''')
    assert(initial_blk.ignore(comments).parseString(fun).asList() ==
           ['INITIAL',
            ['tadj', '=', ['q10', '^', [['celsius', '-', 'temp'], '/', '10', 'degC']]],
            'trates',
            [['v', '+', 'vshift']],
            ['m', '=', 'minf'],
            ['h', '=', 'hinf']])


from pathlib import Path
import numpy as np
from skimage.io import imread
from skimage.transform import resize
from skimage.color import rgb2gray
from utils import (
    resize_to_block_shape,
    identify_biggest_connected_component,
    get_XYLAB,
    WindowXY,
)
from scipy.sparse import csr_matrix
from sklearn.neighbors import (
    KDTree,
    RadiusNeighborsTransformer,
)
from sklearn.metrics import pairwise_distances
from frequency import compute_dct_centers, compute_dct_coefficients

def enforce_connectivity(img, cluster_centers, labels):
    img_flat = img.reshape(-1, img.shape[-1])
    unique_labels = np.unique(labels)
    labels_img = labels.reshape(*img.shape[:2])
    new_labels = labels_img.copy()
    for my_label in unique_labels:
        mask_label = labels_img == my_label
        non_maximal_components = identify_biggest_connected_component(mask_label)
        new_labels[non_maximal_components] = -1
    img_grid = img_flat[:, 3:]
    mask_outliers = new_labels.reshape(-1) == -1

    outliers = img_grid[mask_outliers]
    non_outliers = img_grid[~mask_outliers]
    cluster_grid = cluster_centers[:, 3:]
    kd_tree_non_outliers = KDTree(img_grid)

    nearest_neighbor_ind = kd_tree_non_outliers.query(
        outliers, 1, return_distance=False
    )
    new_labels.reshape(-1)[mask_outliers] = labels[nearest_neighbor_ind.reshape(-1)]
    return new_labels

class BasePixeliser:
    """
    Base class for the Pixelsiers. We are going to follow skimge convention.
    Each Pixelisser will have a `pixelate` method which returns an array of the same @ dimensions than the original image.
    Each pixel of this 2D array will be an integer representing the cluster it belongs to.
    """

    def __init__(self):
        pass

    def pixelate(
        self,
    ):
        pass

    def pixelate_file(self, img_file: Path, *args, **kwargs):
        return self.pixelate(imread(img_file, *args, **kwargs))

class BlockPixeliser(BasePixeliser):
    def __init__(self, patch_size: int = 21):
        super().__init__()
        self.patch_size = patch_size

    def pixelate(self, img: np.ndarray):
        """Pixelate an image by just dividing it into blocks."""

        n_height, n_width = resize_to_block_shape(img.shape, self.patch_size)
        labels = np.arange(n_height * n_width).reshape(n_height, n_width)

        heigh_block, width_block = (
            n_height // self.patch_size,
            n_width // self.patch_size,
        )
        labels = np.arange(heigh_block * width_block).reshape(heigh_block, width_block)
        labels = resize(labels, img.shape[:2], order=0, anti_aliasing=False)

        return labels

class SLICPixeliser(BasePixeliser):
    def __init__(
        self,
        m_spatial: int = 30,
        step: int = 50,
        max_iter: int = 10,
        threshold_error: float = 1.0,
        clever_init: bool = True,
        connectivity: bool = True,
    ):
        super().__init__()
        self.m_spatial = m_spatial
        self.step = step
        self.max_iter = max_iter
        self.threshold_error = threshold_error
        self.clever_init = clever_init
        self.connectivity = connectivity

    def init_centers(self, img: np.ndarray):
        """init centers as a list of array coordinates
        also gives them in XYLAB space
        init can be clever (see class init method)"""
        h, w, _ = img.shape
        centers_X = self.step // 2 + np.arange(0, w - self.step // 2, self.step)
        centers_Y = self.step // 2 + np.arange(0, h - self.step // 2, self.step)
        centers_X, centers_Y = np.meshgrid(centers_X, centers_Y)
        centers_XY = np.array(
            list(zip(np.nditer(centers_X), np.nditer(centers_Y))),
        )

        # clever center initialisation
        if self.clever_init:
            img_gray = rgb2gray(img)
            grad_x, grad_y = np.gradient(img_gray)
            grad_img = (
                grad_x ** 2 + grad_y ** 2
            )  # no need to np.sqrt because only relative order matters
            for k, (X_center, Y_center) in enumerate(centers_XY):
                small_window = WindowXY(X_center, Y_center, 1, w, h)
                # WARNING should not be called if center is on the edges of the image :
                # the small window should be entirely in the image
                sub_grad = grad_img[
                    small_window.y_min : small_window.y_max + 1,
                    small_window.x_min : small_window.x_max + 1,
                ]
                Y_jitter, X_jitter = np.unravel_index(
                    np.argmin(sub_grad, axis=None), sub_grad.shape
                )
                centers_XY[k] = np.array(
                    [
                        min(X_center + X_jitter - 1, w),
                        min(Y_center + Y_jitter - 1, h),
                    ]
                )

        centers_XY[:, [0, 1]] = centers_XY[:, [1, 0]]
        return centers_XY

    def pixelate(self, img: np.ndarray):
        img_XYLAB = get_XYLAB(img)

        h, w, c = img_XYLAB.shape
        initial_centers = self.init_centers(img)

        centers_XY = np.ravel_multi_index(initial_centers.T, dims=(h, w))

        img_flat_XYLAB = img_XYLAB.reshape(-1, c)
        centers_XYLAB = img_flat_XYLAB[centers_XY].copy()

        current_error = np.inf
        current_iteration = 0

        radius_neighbors_transformer = RadiusNeighborsTransformer(
            mode="connectivity", metric="chebyshev", radius=self.step
        )
        radius_neighbors_transformer.fit(img_flat_XYLAB[:, :2])
        while (current_iteration < self.max_iter) and (
            current_error > self.threshold_error
        ):
            new_centers_XYLAB, labels = self.one_slic_iteration(
                img_flat_XYLAB, centers_XYLAB, radius_neighbors_transformer
            )

            current_error = np.mean(
                (new_centers_XYLAB[:, :2] - centers_XYLAB[:, :2]) ** 2
            )
            centers_XYLAB = new_centers_XYLAB
            current_iteration += 1

        if self.connectivity:
            labels = self.enforce_connectivity(img_XYLAB, centers_XYLAB, labels)
        return labels.reshape(h, w)

    def one_slic_iteration(
        self,
        img_flat_XYLAB: np.ndarray,
        centers_XYLAB: np.ndarray,
        radius_neighbors_transformer: RadiusNeighborsTransformer,
    ):
        labels = self.get_closest_center(
            img_flat_XYLAB,
            centers_XYLAB,
            radius_neighbors_transformer=radius_neighbors_transformer,
        )

        centers_XYLAB = self.compute_centers(img_flat_XYLAB, labels)
        return centers_XYLAB, labels

    def get_closest_center(
        self,
        img_flat_XYLAB: np.ndarray,
        centers_XYLAB: np.ndarray,
        radius_neighbors_transformer: RadiusNeighborsTransformer,
    ):

        centers_LAB = centers_XYLAB[:, 2:]
        centers_XY = centers_XYLAB[:, :2]
        img_LAB = img_flat_XYLAB[:, 2:]
        img_XY = img_flat_XYLAB[:, :2]

        connectivity = radius_neighbors_transformer.transform(centers_XY)

        data_distance_matrix = []

        for i, row in enumerate(connectivity):
            close_pixels_indices = row.indices
            close_pixels = img_flat_XYLAB[close_pixels_indices]
            current_center_XY = centers_XY[i : i + 1]
            current_center_LAB = centers_LAB[i : i + 1]
            distance_XY = pairwise_distances(
                current_center_XY, close_pixels[:, :2], metric="sqeuclidean"
            )
            distance_LAB = pairwise_distances(
                current_center_LAB, close_pixels[:, 2:], metric="sqeuclidean"
            )
            distance = distance_LAB + (self.m_spatial / self.step) ** 2 * distance_XY
            data_distance_matrix.extend(np.sqrt(distance.reshape(-1)))

        distance_matrix = csr_matrix(
            (data_distance_matrix, connectivity.indices, connectivity.indptr),
            connectivity.shape,
        )
        mask_distance_matrix = ~connectivity.todense(order="F").astype(bool)

        distance_matrix_ma = np.ma.array(
            distance_matrix.todense(order="F"), mask=mask_distance_matrix
        )

        closest_center_labels = distance_matrix_ma.argmin(axis=0, fill_value=np.inf)

        isolated_pixels_mask = ~mask_distance_matrix.any(axis=0)
        if isolated_pixels_mask.any():

            kd_tree_centers_XY = KDTree(centers_XY)
            indices = kd_tree_centers_XY.query(
                img_XY[isolated_pixels_mask], return_distance=False
            )
            closest_center_labels[isolated_pixels_mask] = indices

        return closest_center_labels

    def compute_centers(self, img_flat_XYLAB: np.ndarray, labels: np.ndarray):

        unique_labels = np.unique(labels)
        centers_XYLAB = np.zeros((len(unique_labels), img_flat_XYLAB.shape[-1]))
        for i, current_label in enumerate(unique_labels):
            mask_label = labels == current_label
            centers_XYLAB[i] = np.mean(img_flat_XYLAB[mask_label], axis=0)
        return centers_XYLAB.astype(np.int32)

    def enforce_connectivity(
        self, img_XYLAB: np.ndarray, centers_XYLAB: np.ndarray, labels: np.ndarray
    ):
        unique_labels = np.unique(labels)
        labels_img = labels.reshape(img_XYLAB.shape[:2])
        new_labels = labels_img.copy()
        img_flat_XYLAB = img_XYLAB.reshape(-1, img_XYLAB.shape[-1])

        maximal_components_mask = np.zeros_like(new_labels, dtype=bool)
        for my_label in unique_labels:
            mask_label = labels_img == my_label
            maximal_components_mask = np.logical_or(
                maximal_components_mask,
                identify_biggest_connected_component(mask_label),
            )

        new_labels[~maximal_components_mask] = -1

        img_XY = img_flat_XYLAB[:, :2]
        mask_outliers = ~maximal_components_mask.reshape(-1)
        centers_XY = centers_XYLAB[:, :2]
        if mask_outliers.any():
            outliers = img_XY[mask_outliers]
            kd_tree_clusters = KDTree(centers_XY)
            indices = kd_tree_clusters.query(outliers, return_distance=False)
            new_labels[~maximal_components_mask] = indices.squeeze()
            return new_labels
        else:
            return labels

class EnergeticSLICPixeliser(SLICPixeliser):
    def __init__(
        self,
        m_spatial: float = 30.0,
        m_frequency: float = 50.0,
        step: int = 50,
        max_iter: int = 10,
        threshold_error: float = 1.0,
        clever_init: bool = True,
        connectivity: bool = True,
        reduce_func=np.median,
        patch_size=9,
        log_transform=True,
        dft=False,
    ):
        super().__init__(
            m_spatial=m_spatial,
            step=step,
            max_iter=max_iter,
            threshold_error=threshold_error,
            clever_init=clever_init,
            connectivity=connectivity,
        )

        self.m_frequency = m_frequency
        self.patch_size = patch_size
        self.log_transform = log_transform
        self.reduce_func = reduce_func
        self.dft = dft

    def pixelate(self, img: np.ndarray):
        img_XYLAB = get_XYLAB(img)

        h, w, c = img_XYLAB.shape
        initial_centers = self.init_centers(img)

        centers_XY = np.ravel_multi_index(initial_centers.T, dims=(h, w))

        img_flat_XYLAB = img_XYLAB.reshape(-1, c)
        centers_XYLAB = img_flat_XYLAB[centers_XY].copy()

        radius_neighbors_transformer = RadiusNeighborsTransformer(
            mode="connectivity", metric="chebyshev", radius=self.step
        )
        radius_neighbors_transformer.fit(img_flat_XYLAB[:, :2])

        kdtree_centers = KDTree(centers_XYLAB[:, :2], metric="chebyshev")

        initial_labels = kdtree_centers.query(
            img_flat_XYLAB[:, :2], return_distance=False, k=1
        ).reshape(-1)

        img_dct = compute_dct_coefficients(
            img,
            patch_size=self.patch_size,
            log_transform=self.log_transform,
            as_channels=True,
            dft=self.dft,
        )

        img_XYLABF = np.concatenate([img_XYLAB, img_dct], axis=-1)

        img_flat_XYLABF = img_XYLABF.reshape(-1, img_XYLABF.shape[-1])

        centers_dct = compute_dct_centers(
            img_flat_XYLABF[:, 5:], initial_labels, reduce_func=self.reduce_func
        ).reshape(-1, self.patch_size ** 2)
        centers_XYLABF = np.concatenate([centers_XYLAB, centers_dct], axis=-1)

        current_error = np.inf
        current_iteration = 0
        while (current_iteration < self.max_iter) and (
            current_error > self.threshold_error
        ):
            new_centers_XYLABF, labels = self.one_slic_iteration(
                img_flat_XYLABF, centers_XYLABF, radius_neighbors_transformer
            )

            current_error = np.mean(
                (new_centers_XYLABF[:, :2] - centers_XYLABF[:, :2]) ** 2
            )
            centers_XYLABF = new_centers_XYLABF
            current_iteration += 1

        if self.connectivity:
            labels = self.enforce_connectivity(img_XYLABF, centers_XYLABF, labels)
        return labels.reshape(h, w)

    def one_slic_iteration(
        self,
        img_flat_XYLABF: np.ndarray,
        centers_XYLABF: np.ndarray,
        radius_neighbors_transformer: RadiusNeighborsTransformer,
    ):
        labels = self.get_closest_center(
            img_flat_XYLABF,
            centers_XYLABF,
            radius_neighbors_transformer=radius_neighbors_transformer,
        )

        centers_XYLABF = self.compute_centers(img_flat_XYLABF, labels)
        return centers_XYLABF, labels

    def get_closest_center(
        self,
        img_flat_XYLABF: np.ndarray,
        centers_XYLABF: np.ndarray,
        radius_neighbors_transformer: RadiusNeighborsTransformer,
    ):

        img_XY = img_flat_XYLABF[:, :2]

        connectivity = radius_neighbors_transformer.transform(centers_XYLABF[:, :2])

        data_distance_matrix = []

        for i, row in enumerate(connectivity):

            close_pixels_indices = row.indices
            close_pixels = img_flat_XYLABF[close_pixels_indices]
            current_center = centers_XYLABF[i : i + 1]

            distance_XY = pairwise_distances(
                current_center[:, :2], close_pixels[:, :2], metric="sqeuclidean"
            )
            distance_LAB = pairwise_distances(
                current_center[:, 2:5], close_pixels[:, 2:5], metric="sqeuclidean"
            )

            distance_F = pairwise_distances(
                current_center[:, 5:], close_pixels[:, 5:], metric="sqeuclidean"
            )
            distance = (
                distance_LAB
                + (self.m_spatial / self.step) ** 2 * distance_XY
                + (self.m_frequency / self.step) ** 2 * distance_F
            )
            data_distance_matrix.extend(np.sqrt(distance.reshape(-1)))

        distance_matrix = csr_matrix(
            (data_distance_matrix, connectivity.indices, connectivity.indptr),
            connectivity.shape,
        )
        mask_distance_matrix = ~connectivity.todense(order="F").astype(bool)

        distance_matrix_ma = np.ma.array(
            distance_matrix.todense(order="F"), mask=mask_distance_matrix
        )

        closest_center_labels = distance_matrix_ma.argmin(axis=0, fill_value=np.inf)

        isolated_pixels_mask = ~mask_distance_matrix.any(axis=0)
        if isolated_pixels_mask.any():

            kd_tree_centers_XY = KDTree(centers_XYLABF[:, :2])
            indices = kd_tree_centers_XY.query(
                img_XY[isolated_pixels_mask], return_distance=False
            )
            closest_center_labels[isolated_pixels_mask] = indices

        return closest_center_labels