from __future__ import division, print_function, absolute_import

import numpy as np
import nose
from nose.tools import (assert_true, assert_false, assert_equal,
                        assert_almost_equal)
from numpy.testing import assert_array_equal, assert_array_almost_equal
from dipy.tracking import metrics as tm
from dipy.tracking import distances as pf


def test_LSCv2():
    xyz1 = np.array([[1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype='float32')
    xyz2 = np.array([[1, 0, 0], [1, 2, 0], [1, 3, 0]], dtype='float32')
    xyz3 = np.array([[1.1, 0, 0], [1, 2, 0], [1, 3, 0]], dtype='float32')
    xyz4 = np.array([[1, 0, 0], [2.1, 0, 0], [3, 0, 0]], dtype='float32')

    xyz5 = np.array([[100, 0, 0], [200, 0, 0], [300, 0, 0]], dtype='float32')
    xyz6 = np.array([[0, 20, 0], [0, 40, 0], [300, 50, 0]], dtype='float32')

    T = [xyz1, xyz2, xyz3, xyz4, xyz5, xyz6]
    C = pf.local_skeleton_clustering(T, 0.2)

    # print C
    # print len(C)

    C2 = pf.local_skeleton_clustering_3pts(T, 0.2)

    # print C2
    # print len(C2)

    # """

    for i in range(40):
        xyz = np.random.rand(3, 3).astype('f4')
        T.append(xyz)

    from time import time
    t1 = time()
    C3 = pf.local_skeleton_clustering(T, .5)
    t2 = time()
    print(t2-t1)
    print(len(C3))

    t1 = time()
    C4 = pf.local_skeleton_clustering_3pts(T, .5)
    t2 = time()
    print(t2-t1)
    print(len(C4))

    for c in C3:
        assert_equal(np.sum(C3[c]['hidden']-C4[c]['hidden']), 0)

    T2 = []
    for i in range(10**4):
        xyz = np.random.rand(10, 3).astype('f4')
        T2.append(xyz)
    t1 = time()
    C5 = pf.local_skeleton_clustering(T2, .5)
    t2 = time()
    print(t2-t1)
    print(len(C5))

    from dipy.data import get_data
    from nibabel import trackvis as tv
    try:
        from dipy.viz import fvtk
    except ImportError as e:
        raise nose.plugins.skip.SkipTest(
            'Fails to import dipy.viz due to %s' % str(e))

    streams, hdr = tv.read(get_data('fornix'))
    T3 = [tm.downsample(s[0], 6) for s in streams]

    print('lenT3', len(T3))

    C = pf.local_skeleton_clustering(T3, 10.)

    print('lenC', len(C))

    """

    r = fvtk.ren()
    colors = np.zeros((len(C), 3))
    for c in C:
        color = np.random.rand(3)
        for i in C[c]['indices']:
            fvtk.add(r, fvtk.line(T3[i], color))
        colors[c] = color
    fvtk.show(r)
    fvtk.clear(r)
    skeleton = []

    def width(w):
        if w<1:
            return 1
        else:
            return w

    for c in C:

        bundle = [T3[i] for i in C[c]['indices']]
        si,s = pf.most_similar_track_mam(bundle, 'avg')
        skeleton.append(bundle[si])
        fvtk.label(r,text = str(len(bundle)), pos=(bundle[si][-1]),
                   scale=(2, 2, 2))
        fvtk.add(r, fvtk.line(skeleton, colors, opacity=1,
                              linewidth = width(len(bundle)/10.)))

    fvtk.show(r)

    """


def test_bundles_distances_mam():
    xyz1A = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]],
                     dtype='float32')
    xyz2A = np.array([[0, 1, 1], [1, 0, 1], [2, 3, -2]], dtype='float32')
    xyz1B = np.array([[-1, 0, 0], [2, 0, 0], [2, 3, 0], [3, 0, 0]],
                     dtype='float32')
    tracksA = [xyz1A, xyz2A]
    tracksB = [xyz1B, xyz1A, xyz2A]
    for metric in ('avg', 'min', 'max'):
        DM2 = pf.bundles_distances_mam(tracksA, tracksB, metric=metric)


def test_bundles_distances_mdf():
    xyz1A = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype='float32')
    xyz2A = np.array([[0, 1, 1], [1, 0, 1], [2, 3, -2]], dtype='float32')
    xyz3A = np.array([[0, 0, 0], [1, 0, 0], [3, 0, 0]], dtype='float32')
    xyz1B = np.array([[-1, 0, 0], [2, 0, 0], [2, 3, 0]], dtype='float32')

    tracksA = [xyz1A, xyz2A]
    tracksB = [xyz1B, xyz1A, xyz2A]
    DM2 = pf.bundles_distances_mdf(tracksA, tracksB)

    tracksA = [xyz1A, xyz1A]
    tracksB = [xyz1A, xyz1A]

    DM2 = pf.bundles_distances_mdf(tracksA, tracksB)
    assert_array_almost_equal(DM2, np.zeros((2, 2)))

    tracksA = [xyz1A, xyz3A]
    tracksB = [xyz2A]

    DM2 = pf.bundles_distances_mdf(tracksA, tracksB)
    print(DM2)

    # assert_array_almost_equal(DM2,np.zeros((2,2)))
    DM = np.zeros(DM2.shape)
    for (a, ta) in enumerate(tracksA):
        for (b, tb) in enumerate(tracksB):
            md = np.sum(np.sqrt(np.sum((ta-tb)**2, axis=1)))/3.
            md2 = np.sum(np.sqrt(np.sum((ta-tb[::-1])**2, axis=1)))/3.
            DM[a, b] = np.min((md, md2))
    print(DM)

    print('--------------')
    for t in tracksA:
        print(t)
    print('--------------')
    for t in tracksB:
        print(t)

    assert_array_almost_equal(DM, DM2, 4)


def test_mam_distances():
    xyz1 = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]])
    xyz2 = np.array([[0, 1, 1], [1, 0, 1], [2, 3, -2]])
    # dm=array([[ 2,  2, 17], [ 3,  1, 14], [6,  2, 13], [11,  5, 14]])
    # this is the distance matrix between points of xyz1
    # and points of xyz2
    xyz1 = xyz1.astype('float32')
    xyz2 = xyz2.astype('float32')
    zd2 = pf.mam_distances(xyz1, xyz2)
    assert_almost_equal(zd2[0], 1.76135602742)


def test_approx_ei_traj():

    segs = 100
    t = np.linspace(0, 1.75*2*np.pi, segs)
    x = t
    y = 5*np.sin(5*t)
    z = np.zeros(x.shape)
    xyz = np.vstack((x, y, z)).T
    xyza = pf.approx_polygon_track(xyz)
    assert_equal(len(xyza), 27)


def test_approx_mdl_traj():

    t = np.linspace(0, 1.75*2*np.pi, 100)
    x = np.sin(t)
    y = np.cos(t)
    z = t
    xyz = np.vstack((x, y, z)).T
    xyza1 = pf.approximate_mdl_trajectory(xyz, alpha=1.)
    xyza2 = pf.approximate_mdl_trajectory(xyz, alpha=2.)
    assert_equal(len(xyza1), 10)
    assert_equal(len(xyza2), 8)
    assert_array_almost_equal(
        xyza1,
        np.array([[  0.00000000e+00,   1.00000000e+00,   0.00000000e+00],
                  [  9.39692621e-01,   3.42020143e-01,   1.22173048e+00],
                  [  6.42787610e-01,  -7.66044443e-01,   2.44346095e+00],
                  [ -5.00000000e-01,  -8.66025404e-01,   3.66519143e+00],
                  [ -9.84807753e-01,   1.73648178e-01,   4.88692191e+00],
                  [ -1.73648178e-01,   9.84807753e-01,   6.10865238e+00],
                  [  8.66025404e-01,   5.00000000e-01,   7.33038286e+00],
                  [  7.66044443e-01,  -6.42787610e-01,   8.55211333e+00],
                  [ -3.42020143e-01,  -9.39692621e-01,   9.77384381e+00],
                  [ -1.00000000e+00,  -4.28626380e-16,   1.09955743e+01]]))

    assert_array_almost_equal(
        xyza2,
        np.array([[  0.00000000e+00,   1.00000000e+00,   0.00000000e+00],
                  [  9.95471923e-01,  -9.50560433e-02,   1.66599610e+00],
                  [ -1.89251244e-01,  -9.81928697e-01,   3.33199221e+00],
                  [ -9.59492974e-01,   2.81732557e-01,   4.99798831e+00],
                  [  3.71662456e-01,   9.28367933e-01,   6.66398442e+00],
                  [  8.88835449e-01,  -4.58226522e-01,   8.32998052e+00],
                  [ -5.40640817e-01,  -8.41253533e-01,   9.99597663e+00],
                  [ -1.00000000e+00,  -4.28626380e-16,   1.09955743e+01]]))


def test_point_track_sq_distance():

    t = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]], dtype='f4')
    p = np.array([-1, -1., -1], dtype='f4')
    assert_equal( pf.point_track_sq_distance_check(t, p, .2**2), False)
    pf.point_track_sq_distance_check(t, p, 2**2), True
    t = np.array([[0, 0, 0], [1, 0, 0], [2, 2, 0]], dtype='f4')
    p = np.array([.5, 0, 0], dtype='f4')
    assert_equal( pf.point_track_sq_distance_check(t, p, .2**2), True)
    p = np.array([.5, 1, 0], dtype='f4')
    assert_equal( pf.point_track_sq_distance_check(t, p, .2**2), False)


def test_track_roi_intersection_check():
    roi = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0]], dtype='f4')
    t = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]], dtype='f4')
    assert_equal( pf.track_roi_intersection_check(t, roi, 1), True)
    t = np.array([[0, 0, 0], [1, 0, 0], [2, 2, 2]], dtype='f4')
    assert_equal(pf.track_roi_intersection_check(t, roi, 1), True)
    t = np.array([[1, 1, 0], [1, 0, 0], [1, -1, 0]], dtype='f4')
    assert_equal( pf.track_roi_intersection_check(t, roi, 1), True)
    t = np.array([[4, 0, 0], [4, 1, 1], [4, 2, 0]], dtype='f4')
    assert_equal(pf.track_roi_intersection_check(t, roi, 1), False)


def test_minimum_distance():
    xyz1 = np.array([[1, 0, 0], [2, 0, 0]], dtype='float32')
    xyz2 = np.array([[3, 0, 0], [4, 0, 0]], dtype='float32')
    assert_equal(pf.minimum_closest_distance(xyz1, xyz2), 1.0)


def test_most_similar_mam():
    xyz1 = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]],
                    dtype='float32')
    xyz2 = np.array([[0, 1, 1], [1, 0, 1], [2, 3, -2]],
                    dtype='float32')
    xyz3 = np.array([[-1, 0, 0], [2, 0, 0], [2, 3, 0], [3, 0, 0]],
                    dtype='float32')
    tracks = [xyz1, xyz2, xyz3]
    for metric in ('avg', 'min', 'max'):
        # pf should be much faster and the results equivalent
        si2, s2 = pf.most_similar_track_mam(tracks, metric=metric)


def test_cut_plane():
    dt = np.dtype(np.float32)
    refx = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0], [3, 0, 0]], dtype=dt)
    bundlex = [np.array([[0.5, 1, 0], [1.5, 2, 0], [2.5, 3, 0]], dtype=dt),
               np.array([[0.5, 2, 0], [1.5, 3, 0], [2.5, 4, 0]], dtype=dt),
               np.array([[0.5, 1, 1], [1.5, 2, 2], [2.5, 3, 3]], dtype=dt),
               np.array([[-0.5, 2, -1], [-1.5, 3, -2], [-2.5, 4, -3]],
                        dtype=dt)]
    expected_hit0 = [
        [ 1.        ,  1.5       ,  0.        ,  0.70710683,  0.        ],
        [ 1.        ,  2.5       ,  0.        ,  0.70710677,  1.        ],
        [ 1.        ,  1.5       ,  1.5       ,  0.81649661,  2.        ]]
    expected_hit1 = [
        [ 2.        ,  2.5       ,  0.        ,  0.70710677,  0.        ],
        [ 2.        ,  3.5       ,  0.        ,  0.70710677,  1.        ],
        [ 2.        ,  2.5       ,  2.5       ,  0.81649655,  2.        ]]
    hitx = pf.cut_plane(bundlex, refx)
    assert_array_almost_equal(hitx[0], expected_hit0)
    assert_array_almost_equal(hitx[1], expected_hit1)
    # check that algorithm allows types other than float32
    bundlex[0] = np.asarray(bundlex[0], dtype=np.float64)
    hitx = pf.cut_plane(bundlex, refx)
    assert_array_almost_equal(hitx[0], expected_hit0)
    assert_array_almost_equal(hitx[1], expected_hit1)
    refx = np.asarray(refx, dtype=np.float64)
    hitx = pf.cut_plane(bundlex, refx)
    assert_array_almost_equal(hitx[0], expected_hit0)
    assert_array_almost_equal(hitx[1], expected_hit1)
