import numpy as np
from dipy.reconst.dsi import DiffusionSpectrumModel
from dipy.data import get_data
from dipy.core.gradients import gradient_table
from numpy.testing import (assert_almost_equal,
                           run_module_suite)
from dipy.reconst.canal import ShoreModel, SHOREmatrix
from dipy.sims.voxel import MultiTensor, all_tensor_evecs, multi_tensor_odf, single_tensor_odf, multi_tensor_rtop, multi_tensor_msd, multi_tensor_pdf
from dipy.data import fetch_isbi2013_2shell, read_isbi2013_2shell
from dipy.data import fetch_taiwan_ntu_dsi, read_taiwan_ntu_dsi
from dipy.data import get_sphere


def test_canal_metrics():
    fetch_taiwan_ntu_dsi()
    img, gtab = read_taiwan_ntu_dsi()
    # fetch_isbi2013_2shell()
    # img, gtab = read_isbi2013_2shell()

    mevals = np.array(([0.0015, 0.0003, 0.0003],
                       [0.0015, 0.0003, 0.0003]))
    angl = [(0, 0), (60, 0)]
    S, sticks = MultiTensor(gtab, mevals, S0=100, angles=angl,
                            fractions=[50, 50], snr=None)
    S = S / S[0, None].astype(np.float)

    asm = ShoreModel(gtab)
    asmfit = asm.fit(S)
    radialOrder = 8
    zeta = 800
    lambdaN = 1e-12
    lambdaL = 1e-12
    Cshore = asmfit.l2estimation(radialOrder=radialOrder, zeta=zeta,
                                 lambdaN=lambdaN, lambdaL=lambdaL)

    Cmat = SHOREmatrix(radialOrder, zeta, gtab)
    S_reconst = np.dot(Cmat, Cshore)
    nmse_signal = np.sqrt(np.sum((S - S_reconst) ** 2)) / (S.sum())
    assert_almost_equal(nmse_signal, 0.0, 4)

    mevecs2 = np.zeros((2, 3, 3))
    angl = np.array(angl)
    for i in range(2):
        mevecs2[i] = all_tensor_evecs(sticks[i]).T

    sphere = get_sphere('symmetric724')
    v = sphere.vertices
    radius = 10e-3
    pdf_shore = asmfit.pdf_iso(v * radius)
    pdf_mt = multi_tensor_pdf(
        v * radius, [.5, .5], mevals=mevals, mevecs=mevecs2)
    nmse_pdf = np.sqrt(np.sum((pdf_mt - pdf_shore) ** 2)) / (pdf_mt.sum())
    assert_almost_equal(nmse_pdf, 0.0, 2)

    rtop_shore_signal = asmfit.rtop_signal()
    rtop_shore_pdf = asmfit.rtop_pdf()
    assert_almost_equal(rtop_shore_signal, rtop_shore_pdf, 9)
    #rtop_mt = multi_tensor_rtop([.5, .5], mevals=mevals)
    #err_rtop = np.abs(rtop_mt - rtop_shore_pdf) / rtop_mt
    #assert_almost_equal(err_rtop, 0.0, 1)
    msd_mt = multi_tensor_msd([.5, .5], mevals=mevals)
    msd_shore = asmfit.msd()
    err_msd = np.abs(msd_mt - msd_shore) / msd_mt
    assert_almost_equal(err_msd, 0, 1)

if __name__ == '__main__':
    run_module_suite()

add a first draft for the test of canal metrics