"""Weights and derivatives of spline orders 0 to 7."""
import torch
from enum import Enum
from .jit_utils import square, cube, pow4, pow5, pow6, pow7


class InterpolationType(Enum):
    nearest = zeroth = 0
    linear = first = 1
    quadratic = second = 2
    cubic = third = 3
    fourth = 4
    fifth = 5
    sixth = 6
    seventh = 7


@torch.jit.script
class Spline:

    def __init__(self, order: int = 1):
        self.order = order

    def weight(self, x):
        w = self.fastweight(x)
        zero = torch.zeros([1], dtype=x.dtype, device=x.device)
        w = torch.where(x.abs() >= (self.order + 1)/2, zero, w)
        return w

    def fastweight(self, x):
        if self.order == 0:
            return torch.ones(x.shape, dtype=x.dtype, device=x.device)
        x = x.abs()
        if self.order == 1:
            return 1 - x
        if self.order == 2:
            x_low = 0.75 - square(x)
            x_up = 0.5 * square(1.5 - x)
            return torch.where(x < 0.5, x_low, x_up)
        if self.order == 3:
            x_low = (x * x * (x - 2.) * 3. + 4.) / 6.
            x_up = cube(2. - x) / 6.
            return torch.where(x < 1., x_low, x_up)
        if self.order == 4:
            x_low = square(x)
            x_low = x_low * (x_low * 0.25 - 0.625) + 115. / 192.
            x_mid = x * (x * (x * (5. - x) / 6. - 1.25) + 5./24.) + 55./96.
            x_up = pow4(x - 2.5) / 24.
            return torch.where(x < 0.5, x_low, torch.where(x < 1.5, x_mid, x_up))
        if self.order == 5:
            x_low = square(x)
            x_low = x_low * (x_low * (0.25 - x / 12.) - 0.5) + 0.55
            x_mid = x * (x * (x * (x * (x / 24. - 0.375) + 1.25) - 1.75) + 0.625) + 0.425
            x_up = pow5(3 - x) / 120.
            return torch.where(x < 1., x_low, torch.where(x < 2., x_mid, x_up))
        if self.order == 6:
            x_low = square(x)
            x_low = x_low * (x_low * (7./48. - x_low/36.) - 77./192.) + 5887./11520.
            x_mid_low = (x * (x * (x * (x * (x * (x / 48. - 7./48.) + 0.328125)
                         - 35./288.) - 91./256.) - 7./768.) + 7861./15360.)
            x_mid_up = (x * (x * (x * (x * (x * (7./60. - x / 120.) - 0.65625)
                        + 133./72.) - 2.5703125) + 1267./960.) + 1379./7680.)
            x_up = pow6(x - 3.5) / 720.
            return torch.where(x < .5, x_low,
                               torch.where(x < 1.5, x_mid_low,
                                           torch.where(x < 2.5, x_mid_up, x_up)))
        if self.order == 7:
            x_low = square(x)
            x_low = (x_low * (x_low * (x_low * (x / 144. - 1./36.)
                     + 1./9.) - 1./3.) + 151./315.)
            x_mid_low = (x * (x * (x * (x * (x * (x * (0.05 - x/240.) - 7./30.)
                         + 0.5) - 7./18.) - 0.1) - 7./90.) + 103./210.)
            x_mid_up = (x * (x * (x * (x * (x * (x * (x / 720. - 1./36.)
                        + 7./30.) - 19./18.) + 49./18.) - 23./6.) + 217./90.)
                        - 139./630.)
            x_up = pow7(4 - x) / 5040.
            return torch.where(x < 1., x_low,
                               torch.where(x < 2., x_mid_low,
                                           torch.where(x < 3., x_mid_up, x_up)))
        raise NotImplementedError

    def grad(self, x):
        if self.order == 0:
            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
        g = self.fastgrad(x)
        zero = torch.zeros([1], dtype=x.dtype, device=x.device)
        g = torch.where(x.abs() >= (self.order + 1)/2, zero, g)
        return g

    def fastgrad(self, x):
        if self.order == 0:
            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
        return self._fastgrad(x.abs()).mul(x.sign())

    def _fastgrad(self, x):
        if self.order == 1:
            return torch.ones(x.shape, dtype=x.dtype, device=x.device)
        if self.order == 2:
            return torch.where(x < 0.5, -2*x, x - 1.5)
        if self.order == 3:
            g_low = x * (x * 1.5 - 2)
            g_up = -0.5 * square(2 - x)
            return torch.where(x < 1, g_low, g_up)
        if self.order == 4:
            g_low = x * (square(x) - 1.25)
            g_mid = x * (x * (x * (-2./3.) + 2.5) - 2.5) + 5./24.
            g_up = cube(2. * x - 5.) / 48.
            return torch.where(x < 0.5, g_low,
                               torch.where(x < 1.5, g_mid, g_up))
        if self.order == 5:
            g_low = x * (x * (x * (x * (-5./12.) + 1.)) - 1.)
            g_mid = x * (x * (x * (x * (5./24.) - 1.5) + 3.75) - 3.5) + 0.625
            g_up = pow4(x - 3.) / (-24.)
            return torch.where(x < 1, g_low,
                               torch.where(x < 2, g_mid, g_up))
        if self.order == 6:
            g_low = square(x)
            g_low = x * (g_low * (7./12.) - square(g_low) / 6. - 77./96.)
            g_mid_low = (x * (x * (x * (x * (x * 0.125 - 35./48.) + 1.3125)
                         - 35./96.) - 0.7109375) - 7./768.)
            g_mid_up = (x * (x * (x * (x * (x / (-20.) + 7./12.) - 2.625)
                        + 133./24.) - 5.140625) + 1267./960.)
            g_up = pow5(2*x - 7) / 3840.
            return torch.where(x < 0.5, g_low,
                               torch.where(x < 1.5, g_mid_low,
                                           torch.where(x < 2.5, g_mid_up,
                                                       g_up)))
        if self.order == 7:
            g_low = square(x)
            g_low = x * (g_low * (g_low * (x * (7./144.) - 1./6.) + 4./9.) - 2./3.)
            g_mid_low = (x * (x * (x * (x * (x * (x * (-7./240.) + 3./10.)
                         - 7./6.) + 2.) - 7./6.) - 1./5.) - 7./90.)
            g_mid_up = (x * (x * (x * (x * (x * (x * (7./720.) - 1./6.)
                        + 7./6.) - 38./9.) + 49./6.) - 23./3.) + 217./90.)
            g_up = pow6(x - 4) / (-720.)
            return torch.where(x < 1, g_low,
                               torch.where(x < 2, g_mid_low,
                                           torch.where(x < 3, g_mid_up, g_up)))
        raise NotImplementedError

    def hess(self, x):
        if self.order == 0:
            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
        h = self.fasthess(x)
        zero = torch.zeros([1], dtype=x.dtype, device=x.device)
        h = torch.where(x.abs() >= (self.order + 1)/2, zero, h)
        return h

    def fasthess(self, x):
        if self.order in (0, 1):
            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
        x = x.abs()
        if self.order == 2:
            one = torch.ones([1], dtype=x.dtype, device=x.device)
            return torch.where(x < 0.5, -2 * one, one)
        if self.order == 3:
            return torch.where(x < 1, 3. * x - 2., 2. - x)
        if self.order == 4:
            return torch.where(x < 0.5, 3. * square(x) - 1.25,
                               torch.where(x < 1.5, x * (-2. * x + 5.) - 2.5,
                                           square(2. * x - 5.) / 8.))
        if self.order == 5:
            h_low = square(x)
            h_low = - h_low * (x * (5./3.) - 3.) - 1.
            h_mid = x * (x * (x * (5./6.) - 9./2.) + 15./2.) - 7./2.
            h_up = 9./2. - x * (x * (x/6. - 3./2.) + 9./2.)
            return torch.where(x < 1, h_low,
                               torch.where(x < 2, h_mid, h_up))
        if self.order == 6:
            h_low = square(x)
            h_low = - h_low * (h_low * (5./6) - 7./4.) - 77./96.
            h_mid_low = (x * (x * (x * (x * (5./8.) - 35./12.) + 63./16.)
                         - 35./48.) - 91./128.)
            h_mid_up = -(x * (x * (x * (x/4. - 7./3.) + 63./8.) - 133./12.)
                         + 329./64.)
            h_up = (x * (x * (x * (x/24. - 7./12.) + 49./16.) - 343./48.)
                    + 2401./384.)
            return torch.where(x < 0.5, h_low,
                               torch.where(x < 1.5, h_mid_low,
                                           torch.where(x < 2.5, h_mid_up,
                                                       h_up)))
        if self.order == 7:
            h_low = square(x)
            h_low = h_low * (h_low*(x * (7./24.) - 5./6.) + 4./3.) - 2./3.
            h_mid_low = - (x * (x * (x * (x * (x * (7./40.) - 3./2.) + 14./3.)
                           - 6.) + 7./3.) + 1./5.)
            h_mid_up = (x * (x * (x * (x * (x * (7./120.) - 5./6.) + 14./3.)
                        - 38./3.) + 49./3.) - 23./3.)
            h_up = - (x * (x * (x * (x * (x/120. - 1./6.) + 4./3.) - 16./3.)
                      + 32./3.) - 128./15.)
            return torch.where(x < 1, h_low,
                               torch.where(x < 2, h_mid_low,
                                           torch.where(x < 3, h_mid_up,
                                                       h_up)))
        raise NotImplementedError