import numpy as np
from numpy.testing import assert_array_equal
from nose.tools import assert_raises
from dipy.io.bvectxt import orientation_from_string, reorient_vectors, \
    orientation_to_string, reorient_vectors

def test_orientation_from_to_string():
    ras = np.array(((0,1), (1,1), (2,1)))
    lps = np.array(((0,-1), (1,-1), (2,1)))
    asl = np.array(((1,1), (2,1), (0,-1)))
    assert_array_equal(orientation_from_string('ras'), ras)
    assert_array_equal(orientation_from_string('lps'), lps)
    assert_array_equal(orientation_from_string('asl'), asl)
    assert_raises(ValueError, orientation_from_string, 'aasl')

    assert orientation_to_string(ras) == 'ras'
    assert orientation_to_string(lps) == 'lps'
    assert orientation_to_string(asl) == 'asl'

def test_reorient_vectors():
    bvec = np.arange(12).reshape((3,4))
    assert_array_equal(reorient_vectors(bvec, 'ras', 'ras'), bvec)
    assert_array_equal(reorient_vectors(bvec, 'ras', 'lpi'), -bvec)
    result = bvec[[1,2,0]]
    assert_array_equal(reorient_vectors(bvec, 'ras', 'asr'), result)
    bvec = result
    result = bvec[[1,0,2]]*[[-1],[1],[-1]]
    assert_array_equal(reorient_vectors(bvec, 'asr', 'ial'), result)
    result = bvec[[1,0,2]]*[[-1],[1],[1]]
    assert_array_equal(reorient_vectors(bvec, 'asr', 'iar'), result)
    assert_raises(ValueError, reorient_vectors, bvec, 'ras', 'ra')

def test_reorient_vectors():
    bvec = np.arange(12).reshape((3,4))
    assert_array_equal(reorient_vectors(bvec, 'ras', 'ras'), bvec)
    assert_array_equal(reorient_vectors(bvec, 'ras', 'lpi'), -bvec)
    result = bvec[[1,2,0]]
    assert_array_equal(reorient_vectors(bvec, 'ras', 'asr'), result)
    bvec = result
    result = bvec[[1,0,2]]*[[-1],[1],[-1]]
    assert_array_equal(reorient_vectors(bvec, 'asr', 'ial'), result)
    result = bvec[[1,0,2]]*[[-1],[1],[1]]
    assert_array_equal(reorient_vectors(bvec, 'asr', 'iar'), result)
    assert_raises(ValueError, reorient_vectors, bvec, 'ras', 'ra')

    bvec = np.arange(12).reshape((3,4))
    bvec = bvec.T
    assert_array_equal(reorient_vectors(bvec, 'ras', 'ras', axis=1), bvec)
    assert_array_equal(reorient_vectors(bvec, 'ras', 'lpi', axis=1), -bvec)
    result = bvec[:, [1,2,0]]
    assert_array_equal(reorient_vectors(bvec, 'ras', 'asr', axis=1), result)
    bvec = result
    result = bvec[:, [1,0,2]]*[-1, 1, -1]
    assert_array_equal(reorient_vectors(bvec, 'asr', 'ial', axis=1), result)
    result = bvec[:, [1,0,2]]*[-1, 1, 1]
    assert_array_equal(reorient_vectors(bvec, 'asr', 'iar', axis=1), result)
