# Copyright (c) 2015, Ecole Polytechnique Federale de Lausanne, Blue Brain Project
# All rights reserved.
#
# This file is part of NeuroM <https://github.com/BlueBrain/NeuroM>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     1. Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#     2. Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#     3. Neither the name of the copyright holder nor the names of
#        its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Test neurom._neuritefunc functionality."""

from pathlib import Path
from math import pi, sqrt

import numpy as np
from nose import tools as nt
from numpy.testing import assert_allclose, assert_array_almost_equal

import neurom as nm

# NOTE: The 'bf' alias is used in the fst/tests modules
# Do NOT change it.
# TODO: If other neurom.features are imported,
# the should use the aliasing used in fst/tests module files
from neurom.features import neuritefunc as _nf
from neurom.features import sectionfunc as sectionfunc
from neurom.geom import convex_hull
from neurom.features.tests.utils import _close

DATA_PATH = Path(__file__).parent.parent.parent.parent / 'test_data'
H5_PATH = DATA_PATH / 'h5/v1'
SWC_PATH = DATA_PATH / 'swc'
SIMPLE = nm.load_neuron(Path(SWC_PATH, 'simple.swc'))
NRN = nm.load_neuron(Path(H5_PATH, 'Neuron.h5'))


def test_principal_direction_extents():
    principal_dir = list(_nf.principal_direction_extents(SIMPLE))
    assert_allclose(principal_dir,
                    (14.736052694538641, 12.105102672688004))

    # test with a realistic neuron
    nrn = nm.load_neuron(Path(H5_PATH, 'bio_neuron-000.h5'))

    p_ref = [1672.9694359427331, 142.43704397865031, 226.45895382204986,
             415.50612748523838, 429.83008974193206, 165.95410536922873,
             346.83281498399697]

    p = _nf.principal_direction_extents(nrn)
    _close(np.array(p), np.array(p_ref))


def test_n_bifurcation_points():
    nt.assert_equal(_nf.n_bifurcation_points(SIMPLE.neurites[0]), 1)
    nt.assert_equal(_nf.n_bifurcation_points(SIMPLE.neurites[1]), 1)
    nt.assert_equal(_nf.n_bifurcation_points(SIMPLE.neurites), 2)


def test_n_forking_points():
    nt.assert_equal(_nf.n_forking_points(SIMPLE.neurites[0]), 1)
    nt.assert_equal(_nf.n_forking_points(SIMPLE.neurites[1]), 1)
    nt.assert_equal(_nf.n_forking_points(SIMPLE.neurites), 2)


def test_n_leaves():
    nt.assert_equal(_nf.n_leaves(SIMPLE.neurites[0]), 2)
    nt.assert_equal(_nf.n_leaves(SIMPLE.neurites[1]), 2)
    nt.assert_equal(_nf.n_leaves(SIMPLE.neurites), 4)


def test_total_area_per_neurite():
    def surface(r0, r1, h):
        return pi * (r0 + r1) * sqrt((r0 - r1) ** 2 + h ** 2)

    basal_area = surface(1, 1, 5) + surface(1, 0, 5) + surface(1, 0, 6)
    ret = _nf.total_area_per_neurite(SIMPLE,
                                     neurite_type=nm.BASAL_DENDRITE)
    nt.assert_almost_equal(ret[0], basal_area)

    axon_area = surface(1, 1, 4) + surface(1, 0, 5) + surface(1, 0, 6)
    ret = _nf.total_area_per_neurite(SIMPLE, neurite_type=nm.AXON)
    nt.assert_almost_equal(ret[0], axon_area)

    ret = _nf.total_area_per_neurite(SIMPLE)
    nt.ok_(np.allclose(ret, [basal_area, axon_area]))


def test_total_volume_per_neurite():
    vol = _nf.total_volume_per_neurite(NRN)
    nt.eq_(len(vol), 4)

    # calculate the volumes by hand and compare
    vol2 = [sum(sectionfunc.section_volume(s) for s in n.iter_sections())
            for n in NRN.neurites
            ]
    nt.eq_(vol, vol2)

    # regression test
    ref_vol = [271.94122143951864, 281.24754646913954,
               274.98039928781355, 276.73860261723024]
    nt.ok_(np.allclose(vol, ref_vol))


def test_neurite_volume_density():

    vol = np.array(_nf.total_volume_per_neurite(NRN))
    hull_vol = np.array([convex_hull(n).volume for n in nm.iter_neurites(NRN)])

    vol_density = _nf.neurite_volume_density(NRN)
    nt.eq_(len(vol_density), 4)
    nt.ok_(np.allclose(vol_density, vol / hull_vol))

    ref_density = [0.43756606998299519, 0.52464681266899216,
                   0.24068543213643726, 0.26289304906104355]
    assert_allclose(vol_density, ref_density)


def test_terminal_path_length_per_neurite():
    terminal_distances = _nf.terminal_path_lengths_per_neurite(SIMPLE)
    assert_allclose(terminal_distances,
                    (5 + 5., 5 + 6., 4. + 6., 4. + 5))
    terminal_distances = _nf.terminal_path_lengths_per_neurite(SIMPLE,
                                                               neurite_type=nm.AXON)
    assert_allclose(terminal_distances,
                    (4. + 6., 4. + 5.))


def test_total_length_per_neurite():
    total_lengths = _nf.total_length_per_neurite(SIMPLE)
    assert_allclose(total_lengths,
                    (5. + 5. + 6., 4. + 5. + 6.))


def test_n_segments():
    n_segments = _nf.n_segments(SIMPLE)
    nt.eq_(n_segments, 6)


def test_n_neurites():
    n_neurites = _nf.n_neurites(SIMPLE)
    nt.eq_(n_neurites, 2)


def test_n_sections():
    n_sections = _nf.n_sections(SIMPLE)
    nt.eq_(n_sections, 6)


def test_neurite_volumes():
    # note: cannot use SIMPLE since it lies in a plane
    total_volumes = _nf.total_volume_per_neurite(NRN)
    assert_allclose(total_volumes,
                    [271.94122143951864, 281.24754646913954,
                     274.98039928781355, 276.73860261723024]
                    )


def test_section_path_lengths():
    path_lengths = list(_nf.section_path_lengths(SIMPLE))
    assert_allclose(path_lengths,
                    (5., 10., 11.,  # type 3, basal dendrite
                     4., 10., 9.))  # type 2, axon


def test_section_term_lengths():
    term_lengths = list(_nf.section_term_lengths(SIMPLE))
    assert_allclose(term_lengths,
                    (5., 6., 6., 5.))


def test_section_bif_lengths():
    bif_lengths = list(_nf.section_bif_lengths(SIMPLE))
    assert_allclose(bif_lengths,
                    (5.,  4.))


def test_section_end_distances():
    end_dist = list(_nf.section_end_distances(SIMPLE))
    assert_allclose(end_dist,
                    [5.0, 5.0, 6.0, 4.0, 6.0, 5.0])


def test_section_partition_pairs():
    part_pairs = list(_nf.partition_pairs(SIMPLE))
    assert_allclose(part_pairs,
                    [(1.0, 1.0), (1.0, 1.0)])


def test_section_bif_radial_distances():
    bif_rads = list(_nf.section_bif_radial_distances(SIMPLE))
    assert_allclose(bif_rads,
                    [5.,  4.])
    trm_rads = list(_nf.section_bif_radial_distances(NRN, neurite_type=nm.AXON))
    assert_allclose(trm_rads,
                    [8.842008561870646,
                     16.7440421479104,
                     23.070306480850533,
                     30.181121708042546,
                     36.62766031035137,
                     43.967487830324885,
                     51.91971040624528,
                     59.427722328770955,
                     66.25222507299583,
                     74.05119754074926])


def test_section_term_radial_distances():
    trm_rads = list(_nf.section_term_radial_distances(SIMPLE))
    assert_allclose(trm_rads,
                    [7.0710678118654755, 7.810249675906654, 7.211102550927978, 6.4031242374328485])
    trm_rads = list(_nf.section_term_radial_distances(NRN, neurite_type=nm.APICAL_DENDRITE))
    assert_allclose(trm_rads,
                    [16.22099879395879,
                     25.992977561564082,
                     33.31600613822663,
                     42.721314797308175,
                     52.379508081911546,
                     59.44327819128149,
                     67.07832724133213,
                     79.97743930553612,
                     87.10434825508366,
                     97.25246040544428,
                     99.58945832481642])


def test_number_of_sections_per_neurite():
    sections = _nf.number_of_sections_per_neurite(SIMPLE)
    assert_allclose(sections,
                    (3, 3))


def test_section_branch_orders():
    branch_orders = list(_nf.section_branch_orders(SIMPLE))
    assert_allclose(branch_orders,
                    (0, 1, 1,  # type 3, basal dendrite
                     0, 1, 1))  # type 2, axon


def test_section_bif_branch_orders():
    bif_branch_orders = list(_nf.section_bif_branch_orders(SIMPLE))
    assert_allclose(bif_branch_orders,
                    (0,  # type 3, basal dendrite
                     0))  # type 2, axon


def test_section_term_branch_orders():
    term_branch_orders = list(_nf.section_term_branch_orders(SIMPLE))
    assert_allclose(term_branch_orders,
                    (1, 1,  # type 3, basal dendrite
                     1, 1))  # type 2, axon


def test_section_radial_distances():
    radial_distances = _nf.section_radial_distances(SIMPLE)
    assert_allclose(radial_distances,
                    (5.0, sqrt(5**2 + 5**2), sqrt(6**2 + 5**2),   # type 3, basal dendrite
                     4.0, sqrt(6**2 + 4**2), sqrt(5**2 + 4**2)))  # type 2, axon


def test_local_bifurcation_angles():
    local_bif_angles = list(_nf.local_bifurcation_angles(SIMPLE))
    assert_allclose(local_bif_angles,
                    (pi, pi))


def test_remote_bifurcation_angles():
    remote_bif_angles = list(_nf.remote_bifurcation_angles(SIMPLE))
    assert_allclose(remote_bif_angles,
                    (pi, pi))


def test_partition():
    partition = list(_nf.bifurcation_partitions(SIMPLE))
    assert_allclose(partition,
                    (1.0, 1.0))


def test_partition_asymmetry():
    partition = list(_nf.partition_asymmetries(SIMPLE))
    assert_allclose(partition,
                    (0.0, 0.0))

    partition = list(_nf.partition_asymmetries(SIMPLE, variant='length'))
    assert_allclose(partition,
                    (0.0625, 0.06666666666666667))

    nt.assert_raises(ValueError, _nf.partition_asymmetries, SIMPLE, variant='unvalid-variant')


def test_segment_lengths():
    segment_lengths = _nf.segment_lengths(SIMPLE)
    assert_allclose(segment_lengths,
                    (5.0, 5.0, 6.0,   # type 3, basal dendrite
                     4.0, 6.0, 5.0))  # type 2, axon


def test_segment_areas():
    result = _nf.segment_areas(SIMPLE)
    assert_allclose(result,
                    [31.415927,
                     16.019042,
                     19.109562,
                     25.132741,
                     19.109562,
                     16.019042])


def test_segment_volumes():
    expected = [
        15.70796327,
        5.23598776,
        6.28318531,
        12.56637061,
        6.28318531,
        5.23598776,
    ]
    result = _nf.segment_volumes(SIMPLE)
    assert_allclose(result, expected)


def test_segment_midpoints():
    midpoints = np.array(_nf.segment_midpoints(SIMPLE))
    assert_allclose(midpoints,
                    np.array([[0.,  (5. + 0) / 2,  0.],  # trunk type 2
                              [-2.5,  5.,  0.],
                              [3.,  5.,  0.],
                              [0., (-4. + 0) / 2.,  0.],  # trunk type 3
                              [3., -4.,  0.],
                              [-2.5, -4.,  0.]]))


def test_segment_radial_distances():
    """midpoints on segments."""
    radial_distances = _nf.segment_radial_distances(SIMPLE)
    assert_allclose(radial_distances,
                    [2.5, sqrt(2.5**2 + 5**2), sqrt(3**2 + 5**2), 2.0, 5.0, sqrt(2.5**2 + 4**2)])


def test_segment_path_lengths():
    pathlengths = _nf.segment_path_lengths(SIMPLE)
    assert_allclose(pathlengths, [5., 10., 11., 4., 10., 9.])

    pathlengths = _nf.segment_path_lengths(NRN)[:5]
    assert_array_almost_equal(pathlengths, [0.1, 1.332525, 2.530149, 3.267878, 4.471462])


def test_principal_direction_extents():
    principal_dir = list(_nf.principal_direction_extents(SIMPLE))
    assert_allclose(principal_dir,
                    (14.736052694538641, 12.105102672688004))


def test_section_taper_rates():
    assert_allclose(list(_nf.section_taper_rates(NRN.neurites[0]))[:10],
                    [0.06776235492169848,
                     0.0588716599404923,
                     0.03791571485186163,
                     0.04674653812192691,
                     -0.026399800285566058,
                     -0.026547582897720887,
                     -0.045038414440432537,
                     0.02083822978267914,
                     -0.0027721371791201038,
                     0.0803069042861474],
                    atol=1e-4)
