# 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 math import sqrt, pi
from nose import tools as nt
import os
import numpy as np
from numpy.testing import assert_allclose
import neurom as nm
from neurom.geom import convex_hull
from neurom.fst import _neuritefunc as _nf
from neurom.fst.sectionfunc import section_volume

from utils import _close

_PWD = os.path.dirname(os.path.abspath(__file__))
H5_PATH = os.path.join(_PWD, '../../../test_data/h5/v1/')
DATA_PATH = os.path.join(H5_PATH, 'Neuron.h5')
SWC_PATH = os.path.join(_PWD, '../../../test_data/swc')
SIMPLE = nm.load_neuron(os.path.join(SWC_PATH, 'simple.swc'))
NRN = nm.load_neuron(DATA_PATH)


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(os.path.join(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(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))

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_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)))
