#!/usr/bin/env python
# coding: utf-8

# # Differential testing
# 
# Test multiple different methods for implementation of finite difference differentiation.
# This includes using `findiff`

# In[1]:


import findiff as fd
import numpy as np
import sympy as sp
from termcolor import colored

import discretisedfield as df
import discretisedfield.util as dfu


# Create testing functions

# In[2]:


def test_function(fun, verbose=1):
    print(colored(f"Testing {fun.__name__}...", "blue"))
    p1 = (0.0, 0.0, 0.0)
    p2 = (10.0, 10.0, 10.0)

    def f(pos):
        x, y, z = pos
        return (x**4, y**4, z**4)

    total = 0
    success = 0

    for i in range(1, 7):
        for order in range(1, 5):
            total += 1
            print(f"Test {i} cell(s) in z - order {order}", end=" ")
            mesh = df.Mesh(p1=p1, p2=p2, n=(10, 10, i))
            field = df.Field(mesh, dim=3, value=f)
            try:
                out = fun(field, "z", order=order)
                if np.allclose(out.array, 0):
                    print(colored("Success with assumption", "yellow"))
                else:
                    print(colored("Success", "green"))
                success += 1
            except Exception as e:
                if verbose:
                    print(colored("Failed", "red"), end="    ")
                    print(e)
                else:
                    print(colored("Failed", "red"))

    print(colored(f"Success rate: {success}/{total}", "blue"))


def test_value(fun, verbose=1):
    # Test results against sympy
    print(colored(f"Testing {fun.__name__}...", "blue"))
    p1 = (0, 0, 0)
    p2 = (10, 10, 10)

    x = sp.symbols("x")
    fx = sp.sqrt(x)

    lam_f = sp.lambdify(x, fx, "numpy")

    def value_fun(point):
        x, y, z = point
        return lam_f(x)

    total = 0
    success = 0

    for i in range(1, 5):
        for order in range(1, 5):
            value = float(fx.diff(x, order).subs(x, 5))
            total += 1
            print(f"Test {i} cell(s) in x - order {order}", end=" ")
            mesh = df.Mesh(p1=p1, p2=p2, n=(i, 1, 1))
            field = df.Field(mesh, dim=1, value=value_fun)
            try:
                fun_value = fun(field, "x", order=order)((5, 5, 5))
                if np.allclose(fun_value, value, rtol=0.05):
                    if verbose:
                        print(colored("Success", "green"), end="    ")
                        print(f"Expected {value}, got {fun_value}")
                    else:
                        print(colored("Success", "green"))
                    success += 1
                else:
                    if verbose:
                        print(colored("Failed", "red"), end="    ")
                        print(f"Expected {value}, got {fun_value}")
                    else:
                        print(colored("Failed", "red"))
            except:
                print("")

    print(colored(f"Success rate: {success}/{total}", "blue"))


def test_accuracy(fun, verbose=1):
    print(colored(f"Testing {fun.__name__}...", "blue"))
    # Test results against sympy
    p1 = (0, 0, 0)
    p2 = (10, 10, 10)
    n = (30, 1, 1)
    mesh = df.Mesh(p1=p1, p2=p2, n=n)

    x = sp.symbols("x")
    fx = x**6

    lam_f = sp.lambdify(x, fx, "numpy")

    def value_fun(point):
        x, y, z = point
        return lam_f(x)

    field = df.Field(mesh, nvdim=1, value=value_fun)

    total = 0
    success = 0

    for order in [1, 2, 3, 4, 5]:
        value = float(fx.diff(x, order).integrate((x, 0, 10)) / 10)
        for acc in [2, 4, 6, 8, 10]:
            print(f"Test {order=}, {acc=}", end=" ")
            total += 1
            try:
                fun_value = fun(field, "x", order=order).mean()
                if np.allclose(fun_value, value, rtol=0.05):
                    if verbose:
                        print(colored("Success", "green"), end="    ")
                        print(f"Expected {value}, got {fun_value}")
                    else:
                        print(colored("Success", "green"))
                    success += 1
                else:
                    if verbose:
                        print(colored("Failed", "red"), end="    ")
                        print(f"Expected {value}, got {fun_value}")
                    else:
                        print(colored("Failed", "red"))
            except:
                print("")
    print(colored(f"Success rate: {success}/{total}", "blue"))


# ### Original df.Field differential

# In[3]:


def diff_original(self, direction, order=1):
    direction = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.
    if self.mesh.n[direction] <= 1:
        return self.zero

    padded_array = self.array

    if order not in (1, 2):
        msg = f"Derivative of the n={order} order is not implemented."
        raise NotImplementedError(msg)

    elif order == 1:
        if self.nvdim == 1:
            derivative_array = np.gradient(
                padded_array[..., 0], self.mesh.cell[direction], axis=direction
            )[..., np.newaxis]
        else:
            derivative_array = np.gradient(
                padded_array, self.mesh.cell[direction], axis=direction
            )

    elif order == 2:
        derivative_array = np.zeros_like(padded_array)
        for i in range(padded_array.shape[direction]):
            if i == 0:
                i1, i2, i3 = i + 2, i + 1, i
            elif i == padded_array.shape[direction] - 1:
                i1, i2, i3 = i, i - 1, i - 2
            else:
                i1, i2, i3 = i + 1, i, i - 1
            index1 = dfu.assemble_index(slice(None), 4, {direction: i1})
            index2 = dfu.assemble_index(slice(None), 4, {direction: i2})
            index3 = dfu.assemble_index(slice(None), 4, {direction: i3})
            index = dfu.assemble_index(slice(None), 4, {direction: i})
            derivative_array[index] = (
                padded_array[index1] - 2 * padded_array[index2] + padded_array[index3]
            ) / self.mesh.cell[direction] ** 2

    return self.__class__(
        self.mesh, dim=self.nvdim, value=derivative_array, vdims=self.vdims
    )


# In[4]:


test_function(diff_original)


# ### Edit Original df.Field differential
# Accounts for `order=2` and number of cells `2`.

# In[5]:


def diff_original_edit(self, direction, order=1):
    direction_idx = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.
    if self.mesh.n[direction_idx] <= order:
        return self.zero

    padded_array = self.array

    if order not in (1, 2):
        msg = f"Derivative of the n={order} order is not implemented."
        raise NotImplementedError(msg)

    elif order == 1:
        if self.nvdim == 1:
            derivative_array = np.gradient(
                padded_array[..., 0], self.mesh.cell[direction_idx], axis=direction_idx
            )[..., np.newaxis]
        else:
            derivative_array = np.gradient(
                padded_array, self.mesh.cell[direction_idx], axis=direction_idx
            )

    elif order == 2:
        derivative_array = np.zeros_like(padded_array)
        for i in range(padded_array.shape[direction_idx]):
            if i == 0:
                i1, i2, i3 = i + 2, i + 1, i
            elif i == padded_array.shape[direction_idx] - 1:
                i1, i2, i3 = i, i - 1, i - 2
            else:
                i1, i2, i3 = i + 1, i, i - 1
            index1 = dfu.assemble_index(slice(None), 4, {direction_idx: i1})
            index2 = dfu.assemble_index(slice(None), 4, {direction_idx: i2})
            index3 = dfu.assemble_index(slice(None), 4, {direction_idx: i3})
            index = dfu.assemble_index(slice(None), 4, {direction_idx: i})
            derivative_array[index] = (
                padded_array[index1] - 2 * padded_array[index2] + padded_array[index3]
            ) / self.mesh.cell[direction_idx] ** 2

    return self.__class__(
        self.mesh, dim=self.nvdim, value=derivative_array, vdims=self.vdims
    )


# In[6]:


test_function(diff_original_edit)


# ### Edit Fast Original df.Field differential
# Accounts for `order=2` and number of cells `2` and faster implementation of `order=2`.
# 
# Pad array so the same stencil can be used everywhere
# \begin{align}
#     f(1) + f(-1) - 2f(0) &= f(2) + f(0) - 2 f(1) \\ 
#     f(-1) &= f(2) - 3 f(1) + 3f(0)
# \end{align}

# In[7]:


def diff_original_fast(self, direction, order=1):
    direction_idx = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.
    if self.mesh.n[direction_idx] <= order:
        return self.zero

    padded_array = self.array

    if order not in (1, 2):
        msg = f"Derivative of the n={order} order is not implemented."
        raise NotImplementedError(msg)

    elif order == 1:
        if self.nvdim == 1:
            derivative_array = np.gradient(
                padded_array[..., 0], self.mesh.cell[direction_idx], axis=direction_idx
            )[..., np.newaxis]
        else:
            derivative_array = np.gradient(
                padded_array, self.mesh.cell[direction_idx], axis=direction_idx
            )

    elif order == 2:

        def pad_fun(vector, pad_width, iaxis, kwargs):
            if iaxis == direction_idx:
                vector[0] = vector[3] - 3 * vector[2] + 3 * vector[1]
                vector[-1] = vector[-4] - 3 * vector[-3] + 3 * vector[-2]

        pad_width = [(0, 0)] * 4
        pad_width[direction_idx] = (1, 1)
        padded_array = np.pad(padded_array, pad_width, pad_fun)
        # padded_array = self.pad({'x': (1, 1)}, mode='constant').array
        # padded_array[0] = padded_array[3] - 3*padded_array[2] + 3*padded_array[1]
        # padded_array[-1] = padded_array[-4] - 3*padded_array[-3] + 3*padded_array[-2]
        # derivative_array = np.empty_like(padded_array)
        index_p1 = slice(2, None)
        index_0 = slice(1, -1)
        index_m1 = slice(None, -2)
        derivative_array = (
            padded_array[index_p1] - 2 * padded_array[index_0] + padded_array[index_m1]
        ) / self.mesh.cell[direction_idx] ** 2

        # derivative_array[0] = derivative_array[1]
        # derivative_array[-1] = derivative_array[-2]

    return self.__class__(
        self.mesh, dim=self.nvdim, value=derivative_array, vdims=self.vdims
    )


# In[8]:


get_ipython().run_cell_magic('timeit', '', 'diff_original_fast(field, "x", order=2)\n')


# In[ ]:


get_ipython().run_cell_magic('timeit', '', 'diff_original_edit(field, "x", order=2)\n')


# In[ ]:


test_function(diff_original_fast)


# ### Findiff implementation

# In[9]:


def diff_a(self, direction, order=1, acc=None):
    direction_idx = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.

    # ASSUME that needs more cells than curvature!
    if self.mesh.n[direction_idx] <= order:
        return self.zero

    # Create FinDiff object.
    diff_fd = fd.FinDiff(direction_idx, self.mesh.cell[direction_idx], order, acc=acc)

    derivative_array = diff_fd(self.array)

    return self.__class__(
        self.mesh,
        nvdim=self.nvdim,
        value=derivative_array,
        vdims=self.vdims,
    )


# In[10]:


test_function(diff_a)


# ### Findiff implementation
# Allowing for mixed accuracy default behaviour

# In[11]:


def fd_diff_edit_b(array, h, order, dim, acc=None):
    """Edit of findiff.diff.diff to allow differnt accuracies"""
    if acc is None:
        coefs = {}
        coefs["center"] = fd.coefficients(order, acc=2)["center"]
        coefs["forward"] = fd.coefficients(order, offsets=range(0, order + 1))
        coefs["backward"] = fd.coefficients(order, offsets=range(-order, 1))
    else:
        coefs = fd.coefficients(order, acc=acc)

    try:
        npts = array.shape[dim]
    except AttributeError as err:
        raise ValueError(
            "FinDiff objects can only be applied to arrays or evaluated(!) functions"
            " returning arrays"
        ) from err
    yd = np.zeros_like(array)
    scheme = "center"
    weights = coefs[scheme]["coefficients"]
    offsets = coefs[scheme]["offsets"]

    num_bndry_points = len(weights) // 2
    ref_slice = slice(num_bndry_points, npts - num_bndry_points, 1)
    off_slices = [
        fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
        for k in range(len(offsets))
    ]

    fd.diff.Diff._apply_to_array(..., yd, array, weights, off_slices, ref_slice, dim)

    scheme = "forward"
    weights = coefs[scheme]["coefficients"]
    offsets = coefs[scheme]["offsets"]

    ref_slice = slice(0, num_bndry_points, 1)
    off_slices = [
        fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
        for k in range(len(offsets))
    ]

    fd.diff.Diff._apply_to_array(..., yd, array, weights, off_slices, ref_slice, dim)

    scheme = "backward"
    weights = coefs[scheme]["coefficients"]
    offsets = coefs[scheme]["offsets"]

    ref_slice = slice(npts - num_bndry_points, npts, 1)
    off_slices = [
        fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
        for k in range(len(offsets))
    ]

    fd.diff.Diff._apply_to_array(..., yd, array, weights, off_slices, ref_slice, dim)

    h_inv = 1.0 / h**order

    return yd * h_inv


# In[12]:


def diff_b(self, direction, order=1, acc=None):
    direction_idx = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.

    # ASSUME that needs more cells than curvature!
    if self.mesh.n[direction_idx] <= order:
        return self.zero

    # Create FinDiff object.
    derivative_array = fd_diff_edit_b(
        self.array, self.mesh.cell[direction_idx], order, direction_idx, acc=acc
    )

    return self.__class__(
        self.mesh,
        nvdim=self.nvdim,
        value=derivative_array,
        vdims=self.vdims,
    )


# In[13]:


test_function(diff_b)


# ### Findiff implementation
# Generalise defaults to work for all orders

# In[14]:


def fd_diff_edit_c(array, h, order, dim, acc=None):
    """Edit of findiff.diff.diff to allow differnt accuracies"""
    if acc is None:
        coefs = {}
        coefs["center"] = fd.coefficients(order, acc=2)["center"]
        coefs["forward"] = fd.coefficients(order, offsets=range(0, order + 1))
        coefs["backward"] = fd.coefficients(order, offsets=range(-order, 1))
    else:
        coefs = fd.coefficients(order, acc=acc)

    try:
        npts = array.shape[dim]
    except AttributeError as err:
        raise ValueError(
            "FinDiff objects can only be applied to arrays or evaluated(!) functions"
            " returning arrays"
        ) from err

    yd = np.zeros_like(array)

    len_needed = len(coefs["center"]["coefficients"]) // 2 - 1 + order
    if npts <= len_needed:
        for i in range(npts):
            coefs = fd.coefficients(order, offsets=(np.arange(npts) - i).tolist())
            weights = coefs["coefficients"]
            offsets = coefs["offsets"]

            ref_slice = slice(i, i + 1, 1)
            off_slices = [
                fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
                for k in range(len(offsets))
            ]
            fd.diff.Diff._apply_to_array(
                ..., yd, array, weights, off_slices, ref_slice, dim
            )

    else:
        scheme = "center"
        weights = coefs[scheme]["coefficients"]
        offsets = coefs[scheme]["offsets"]

        num_bndry_points = len(weights) // 2

        ref_slice = slice(num_bndry_points, npts - num_bndry_points, 1)
        off_slices = [
            fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
            for k in range(len(offsets))
        ]

        fd.diff.Diff._apply_to_array(
            ..., yd, array, weights, off_slices, ref_slice, dim
        )

        scheme = "forward"
        weights = coefs[scheme]["coefficients"]
        offsets = coefs[scheme]["offsets"]

        ref_slice = slice(0, num_bndry_points, 1)
        off_slices = [
            fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
            for k in range(len(offsets))
        ]

        fd.diff.Diff._apply_to_array(
            ..., yd, array, weights, off_slices, ref_slice, dim
        )

        scheme = "backward"
        weights = coefs[scheme]["coefficients"]
        offsets = coefs[scheme]["offsets"]

        ref_slice = slice(npts - num_bndry_points, npts, 1)
        off_slices = [
            fd.diff.Diff._shift_slice(..., ref_slice, offsets[k], npts)
            for k in range(len(offsets))
        ]

        fd.diff.Diff._apply_to_array(
            ..., yd, array, weights, off_slices, ref_slice, dim
        )

    h_inv = 1.0 / h**order

    return yd * h_inv


# In[15]:


def diff_c(self, direction, order=1, acc=None):
    direction_idx = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.

    # ASSUME that needs more cells than curvature!
    if self.mesh.n[direction_idx] <= order:
        return self.zero

    # Create FinDiff object.
    derivative_array = fd_diff_edit_c(
        self.array, self.mesh.cell[direction_idx], order, direction_idx, acc=acc
    )

    return self.__class__(
        self.mesh,
        nvdim=self.nvdim,
        value=derivative_array,
        vdims=self.vdims,
    )


# In[16]:


test_function(diff_c)


# ### Mixed implementation

# In[17]:


def diff_d(self, direction, order=1, acc=None):
    direction_idx = dfu.axesdict[direction]

    # If there are no neighbouring cells in the specified direction, zero
    # field is returned.
    if self.mesh.n[direction_idx] <= order:
        return self.zero

    padded_array = self.array

    # Only use our implimentation if we have to
    if order in (1, 2) and acc is None:
        if order == 1:
            if self.nvdim == 1:
                derivative_array = np.gradient(
                    padded_array[..., 0],
                    self.mesh.cell[direction_idx],
                    axis=direction_idx,
                )[..., np.newaxis]
            else:
                derivative_array = np.gradient(
                    padded_array, self.mesh.cell[direction_idx], axis=direction_idx
                )

        elif order == 2:
            derivative_array = np.zeros_like(padded_array)
            for i in range(padded_array.shape[direction_idx]):
                if i == 0:
                    i1, i2, i3 = i + 2, i + 1, i
                elif i == padded_array.shape[direction_idx] - 1:
                    i1, i2, i3 = i, i - 1, i - 2
                else:
                    i1, i2, i3 = i + 1, i, i - 1
                index1 = dfu.assemble_index(slice(None), 4, {direction_idx: i1})
                index2 = dfu.assemble_index(slice(None), 4, {direction_idx: i2})
                index3 = dfu.assemble_index(slice(None), 4, {direction_idx: i3})
                index = dfu.assemble_index(slice(None), 4, {direction_idx: i})
                derivative_array[index] = (
                    padded_array[index1]
                    - 2 * padded_array[index2]
                    + padded_array[index3]
                ) / self.mesh.cell[direction_idx] ** 2
    else:
        if acc is None:
            acc = 2
        coeffs = fd.coefficients(order, acc=acc)
        stencil_len_back = len(coeffs["backward"]["offsets"]) - 1
        stencil_max_cent = max(coeffs["center"]["offsets"]) - 1

        len_needed = stencil_len_back + stencil_max_cent
        if self.mesh.n[direction_idx] <= len_needed:
            raise ValueError(
                f"The minimum size of the mesh in the direction {direction} "
                f"is {len_needed+1} for an order order {order} and acc {acc}."
            )
        diff_fd = fd.FinDiff(
            direction_idx, self.mesh.cell[direction_idx], order, acc=acc
        )
        derivative_array = diff_fd(padded_array)

    return self.__class__(
        self.mesh, dim=self.nvdim, value=derivative_array, vdims=self.vdims
    )


# In[18]:


test_function(diff_d)


# In[19]:


test_value(diff_d)


# ### ndimage implementation

# In[20]:


from scipy import ndimage


# In[37]:


p1 = (0.0, 0.0, 0.0)
p2 = (10.0, 10.0, 10.0)


def f(pos):
    x, y, z = pos
    return (x**4, y**4, z**4)


mesh = df.Mesh(p1=p1, p2=p2, n=(10, 10, 20), bc="xyz")
field = df.Field(mesh, dim=3, value=f)


# In[62]:


field.diff("x").array[0, 5, 5]


# In[60]:


k = np.zeros((3, 1, 1, 1))
k[0, 0, 0, 0] = 0.5
k[2, 0, 0, 0] = -0.5


# In[63]:


(
    ndimage.convolve(field.array, k, mode="wrap", origin=(0, 0, 0, 0))
    / field.mesh.cell[0]
)[0, 5, 5]