import time

import numpy as np

import nibabel as nib

from .. import recspeed as rp
from .. import gqi as gq
from .. import dti as dt
from ...core import meshes
from ...data import get_data, get_sphere

from nose.tools import assert_true, assert_false, assert_equal, assert_raises
from numpy.testing import assert_array_equal, assert_array_almost_equal


def test_gqiodfmask():

    fimg,fbvals,fbvecs=get_data('small_64D')    
    bvals=np.load(fbvals)
    gradients=np.load(fbvecs)
    data=nib.load(fimg).get_data()
    mask =np.random.random(data.shape[:3])
    mask[mask >.5]=1
    mask[mask<=.5]=0    
    gqs = gq.GeneralizedQSampling(data,bvals,gradients,mask=mask)
    mask=np.zeros(data.shape[:3])
    gqs = gq.GeneralizedQSampling(data,bvals,gradients,mask=mask)

    assert_equal(np.sum(np.isnan(gqs.QA)),5000) #all voxels should be NULL


def test_gqiodf():

    #read bvals,gradients and data
    fimg,fbvals,fbvecs=get_data('small_64D')
    bvals=np.load(fbvals)
    gradients=np.load(fbvecs)
    data=nib.load(fimg).get_data()

    #print(bvals.shape)
    #print(gradients.shape)
    #print(data.shape)
    t1=time.clock()
    
    gqs = gq.GeneralizedQSampling(data,bvals,gradients)
    ten = dt.Tensor(data,bvals,gradients,thresh=50)
    
    fa=ten.fa()

    x,y,z,a,b=ten.evecs.shape
    evecs=ten.evecs
    xyz=x*y*z
    evecs = evecs.reshape(xyz,3,3)
    #vs = np.sign(evecs[:,2,:])
    #print vs.shape
    #print np.hstack((vs,vs,vs)).reshape(1000,3,3).shape
    #evecs = np.hstack((vs,vs,vs)).reshape(1000,3,3)
    #print evecs.shape
    evals=ten.evals
    evals = evals.reshape(xyz,3)
    #print evals.shape

    t2=time.clock()
    #print('GQS in %d' %(t2-t1))
        
    odf_vertices, odf_faces = get_sphere('symmetric362')

    #Yeh et.al, IEEE TMI, 2010
    #calculate the odf using GQI

    scaling=np.sqrt(bvals*0.01506) # 0.01506 = 6*D where D is the free
    #water diffusion coefficient 
    #l_values sqrt(6 D tau) D free water
    #diffusion coefficiet and tau included in the b-value

    tmp=np.tile(scaling,(3,1))
    b_vector=gradients.T*tmp
    Lambda = 1.2 # smoothing parameter - diffusion sampling length
    
    q2odf_params=np.sinc(np.dot(b_vector.T, odf_vertices.T) * Lambda/np.pi)
    #implements equation no. 9 from Yeh et.al.

    S=data.copy()

    x,y,z,g=S.shape
    S=S.reshape(x*y*z,g)
    QA = np.zeros((x*y*z,5))
    IN = np.zeros((x*y*z,5))

    fwd = 0
    
    #Calculate Quantitative Anisotropy and find the peaks and the indices
    #for every voxel

    summary = {}

    summary['vertices'] = odf_vertices
    v = odf_vertices.shape[0]
    summary['faces'] = odf_faces
    f = odf_faces.shape[0]

    '''
    If e = number_of_edges
    the Euler formula says f-e+v = 2 for a mesh on a sphere
    Here, assuming we have a healthy triangulation
    every face is a triangle, all 3 of whose edges should belong to
    exactly two faces = so 2*e = 3*f
    to avoid division we test whether 2*f - 3*f + 2*v == 4
    or equivalently 2*v - f == 4
    '''

    assert_equal(2*v-f, 4,'Euler test fails')
    
    for (i,s) in enumerate(S):
        #print 'Volume %d' % i
        istr = str(i)
        summary[istr] = {}
        
        odf = Q2odf(s,q2odf_params)
        peaks,inds=rp.peak_finding(odf,odf_faces)
        fwd=max(np.max(odf),fwd)
        peaks = peaks - np.min(odf)
        l=min(len(peaks),5)
        QA[i][:l] = peaks[:l]
        IN[i][:l] = inds[:l]

        summary[istr]['odf'] = odf
        summary[istr]['peaks'] = peaks
        summary[istr]['inds'] = inds
        summary[istr]['evecs'] = evecs[i,:,:]
        summary[istr]['evals'] = evals[i,:]
   
    QA/=fwd
    QA=QA.reshape(x,y,z,5)    
    IN=IN.reshape(x,y,z,5)

    peaks_1 = [i for i in range(1000) if len(summary[str(i)]['inds'])==1]
    peaks_2 = [i for i in range(1000) if len(summary[str(i)]['inds'])==2]
    peaks_3 = [i for i in range(1000) if len(summary[str(i)]['inds'])==3]

    # correct numbers of voxels with respectively 1,2,3 ODF/QA peaks
    assert_array_equal((len(peaks_1),len(peaks_2),len(peaks_3)), (790,196,14),
                             'error in numbers of QA/ODF peaks')

    # correct indices of odf directions for voxels 0,10,44
    # with respectively 1,2,3 ODF/QA peaks
    assert_array_equal(summary['0']['inds'],[116],
                             'wrong peak indices for voxel 0')
    assert_array_equal(summary['10']['inds'],[105, 78],
                             'wrong peak indices for voxel 10')
    assert_array_equal(summary['44']['inds'],[95, 84, 108],
                             'wrong peak indices for voxel 44')

    assert_equal(np.argmax(summary['0']['odf']), 116)
    assert_equal(np.argmax(summary['10']['odf']), 105)


def upper_hemi_map(v):
    '''
    maps a 3-vector into the z-upper hemisphere
    '''
    return np.sign(v[2])*v


def equatorial_maximum(vertices, odf, pole, width):
    eqvert = meshes.equatorial_vertices(vertices, pole, width)
    '''
    need to test for whether eqvert is empty or not
    '''
    if len(eqvert) == 0:
        print 'empty equatorial band at pole', pole, 'with width', width
        return Null, Null

    eqvals = [odf[i] for i in eqvert]
    eqargmax = np.argmax(eqvals)
    eqvertmax = eqvert[eqargmax]
    eqvalmax = eqvals[eqargmax]
    return eqvertmax, eqvalmax


def patch_vertices(vertices,pole, width):
    '''
    find 'vertices' within the cone of 'width' around 'pole'
    '''
    return [i for i,v in enumerate(vertices) if np.dot(v,pole) > 1- width]


def patch_maximum(vertices, odf, pole, width):
    eqvert = patch_vertices(vertices, pole, width)
    '''
    need to test for whether eqvert is empty or not
    '''
    if len(eqvert) == 0:
        print 'empty cone around pole', pole, 'with width', width
        return Null, Null

    eqvals = [odf[i] for i in eqvert]
    eqargmax = np.argmax(eqvals)
    eqvertmax = eqvert[eqargmax]
    eqvalmax = eqvals[eqargmax]
    return eqvertmax, eqvalmax


def triple_odf_maxima(vertices, odf, width):

    indmax1 = np.argmax([odf[i] for i,v in enumerate(vertices)])
    odfmax1 = odf[indmax1]
    indmax2, odfmax2 = equatorial_maximum(vertices, odf, vertices[indmax1], width)
    cross12 = np.cross(vertices[indmax1],vertices[indmax2])
    indmax3, odfmax3 = patch_maximum(vertices, odf, cross12, width)
    return [(indmax1, odfmax1),(indmax2, odfmax2),(indmax3, odfmax3)]


def test_gqi_small():

    #read bvals,gradients and data
    fimg,fbvals,fbvecs=get_data('small_64D')    
    bvals=np.load(fbvals)
    gradients=np.load(fbvecs)
    data=nib.load(fimg).get_data()

    print(bvals.shape)
    print(gradients.shape)
    print(data.shape)

    t1=time.clock()
    
    gqs = gq.GeneralizedQSampling(data,bvals,gradients)

    t2=time.clock()
    print('GQS in %d' %(t2-t1))
    odf_vertices, odf_faces = get_sphere('symmetric362')

    #Yeh et.al, IEEE TMI, 2010
    #calculate the odf using GQI

    scaling=np.sqrt(bvals*0.01506) # 0.01506 = 6*D where D is the free
    #water diffusion coefficient 
    #l_values sqrt(6 D tau) D free water
    #diffusion coefficiet and tau included in the b-value

    tmp=np.tile(scaling,(3,1))
    b_vector=gradients.T*tmp
    Lambda = 1.2 # smoothing parameter - diffusion sampling length
    
    q2odf_params=np.sinc(np.dot(b_vector.T, odf_vertices.T) * Lambda/np.pi)
    #implements equation no. 9 from Yeh et.al.
    S=data.copy()
    x,y,z,g=S.shape
    S=S.reshape(x*y*z,g)
    QA = np.zeros((x*y*z,5))
    IN = np.zeros((x*y*z,5))
    fwd = 0
    
    #Calculate Quantitative Anisotropy and find the peaks and the indices
    #for every voxel

    for (i,s) in enumerate(S):

        odf = Q2odf(s,q2odf_params)
        peaks,inds=rp.peak_finding(odf,odf_faces)
        fwd=max(np.max(odf),fwd)
        peaks = peaks - np.min(odf)
        l=min(len(peaks),5)
        QA[i][:l] = peaks[:l]
        IN[i][:l] = inds[:l]

    QA/=fwd
    QA=QA.reshape(x,y,z,5)    
    IN=IN.reshape(x,y,z,5)
    
    print('Old %d secs' %(time.clock() - t2))    
    assert_equal((gqs.QA-QA).max(),0.,'Frank QA different than dipy QA')
    assert_equal((gqs.QA.shape),QA.shape, 'Frank QA shape is different') 


def Q2odf(s,q2odf_params):
    odf=np.dot(s,q2odf_params)
    return odf


def peak_finding(odf,odf_faces):
    #proton density already include from the scaling b_table[0][0] and s[0]
    #find local maxima
    peak=odf.copy()
    # where the smallest odf values in the vertices of a face remove the
    # two smallest vertices 

    for face in odf_faces:
        i, j, k = face
        check=np.array([odf[i],odf[j],odf[k]])
        zeroing=check.argsort()
        peak[face[zeroing[0]]]=0
        peak[face[zeroing[1]]]=0

    #for later testing expecting peak.max 794595.94774980657 and
    #np.where(peak>0) (array([166, 347]),)
    #we just need the first half of peak
    peak=peak[0:len(peak)/2]
    #find local maxima and give fiber orientation (inds) and magnitute
    #peaks in a descending order
    inds=np.where(peak>0)[0]
    pinds=np.argsort(peak[inds])
    peaks=peak[inds[pinds]][::-1]
    return peaks, inds[pinds][::-1]


if __name__ == "__main__":
    T=test_gqiodf()








