# -*- coding: utf-8 -*-
"""
@brief Computing gradient

@author: Anna Petrasova
"""

import numpy as np
import unittest


def gradient(array, step):
    """Computing gradient (second order approximation)
    using central differencing scheme (plus forward and backward
    difference of second order approx.)"""
    size_z, size_y, size_x = array.shape

    gradient_array_x = np.zeros(array.shape)
    for depth in range(size_z):
        for row in range(size_y):
            col_array = array[depth, row, :]
            gradient_array_x[depth, row, 0] = (-3 * col_array[0] +
                                               4 * col_array[1] -
                                               col_array[2]) / (2 * step[0])
            gradient_array_x[depth, row, size_x - 1] = (col_array[size_x - 3] -
                                                        4 * col_array[size_x - 2] +
                                                        3 * col_array[size_x - 1]) / (2 * step[0])
            for col in range(1, size_x - 1):
                gradient_array_x[depth, row, col] = (col_array[col + 1] -
                                                     col_array[col - 1]) / (2 * step[0])

    gradient_array_y = np.zeros(array.shape)
    for depth in range(size_z):
        for col in range(size_x):
            row_array = array[depth, :, col]
            gradient_array_y[depth, 0, col] = (-3 * row_array[0] +
                                               4 * row_array[1] -
                                               row_array[2]) / (2 * step[1])
            gradient_array_y[depth, size_y - 1, col] = (row_array[size_y - 3] -
                                                        4 * row_array[size_y - 2] +
                                                        3 * row_array[size_y - 1]) / (2 * step[1])
            for row in range(1, size_y - 1):
                gradient_array_y[depth, row, col] = (row_array[row + 1] -
                                                     row_array[row - 1]) / (2 * step[1])

    gradient_array_z = np.zeros(array.shape)
    for row in range(size_y):
        for col in range(size_x):
            depth_array = array[:, row, col]
            gradient_array_z[0, row, col] = (-3 * depth_array[0] +
                                             4 * depth_array[1] -
                                             depth_array[2]) / (2 * step[2])
            gradient_array_z[size_z - 1, row, col] = (depth_array[size_z - 3] -
                                                      4 * depth_array[size_z - 2] +
                                                      3 * depth_array[size_z - 1]) / (2 * step[2])
            for depth in range(1, size_z - 1):
                gradient_array_z[depth, row, col] = (depth_array[depth + 1] -
                                                     depth_array[depth - 1]) / (2 * step[2])

    return gradient_array_x, -gradient_array_y, gradient_array_z


class TestGradient(unittest.TestCase):

    def test_gradient(self):
        array = np.random.rand(10, 15, 3)
        step = [2, 0.6, 1]
        # numpy_gradient is function gradient which must be taken from numpy master branch
        # https://github.com/numpy/numpy/blob/master/numpy/lib/function_base.py
        # because 1.8 uses first order approx on edges, not second
        numpy_gradient_z, numpy_gradient_y, numpy_gradient_x = numpy_gradient(array, step[2], step[1], step[0])
        gradient_x, gradient_y, gradient_z = gradient(array, step)

        self.assertEqual(True, np.allclose(gradient_z, numpy_gradient_z, atol=1e-8))
        self.assertEqual(True, np.allclose(gradient_y, -numpy_gradient_y, atol=1e-8))
        self.assertEqual(True, np.allclose(gradient_x, numpy_gradient_x, atol=1e-8))


if __name__ == '__main__':
    unittest.main()
