%load_ext pymatbridge

# put some system-specific variables in the namespace ('le' dict)
%run ~/set_localenv_vars.py

nb_name = 'making_a_region_mapping_for_aal'

# this is where I look for stuf
outdir = '%s/notebooks/%s' %(le['data_dir'],nb_name)

# this is where stuff actually lives
outdir_loc = '%s/%s' %('/mnt/titania-hecuba/mcintosh_lab/john/Data/notebooks', nb_name)

# symlink the two
import os
if not os.path.isdir(outdir_loc): os.makedirs(outdir_loc)
if not os.path.isdir(outdir): os.system('ln -s %s %s' %(outdir_loc,outdir))
                                        

# root path to tvb-library, tvb-data, etc. libraries. 
tvb_folder = le['code_dir'] + '/libraries_of_mine/github'


# stuff for workdocs-cloudfiles

aws_key = 'drowssaperucesyreva'
aws_secret = '?teytidettopsuoyevah'  

# Generic imports

import os,sys,glob,h5py,itertools,multiprocessing,\
       numpy as np,pandas as pd
from datetime import datetime
from time import time
from scipy import optimize
from scipy.spatial import cKDTree

# Visualization stuff

%matplotlib inline
from matplotlib import pyplot as plt
from IPython.display import Image,display as d
import seaborn as sns



# Neuroimaging stuff

import nibabel as nib
from nilearn.plotting import plot_roi
from dipy.tracking.utils import apply_affine

# TVB stuff

nbso,nbse = sys.stdout,sys.stderr # hack part 1/2 to keep output printing properly
sys.path.append(tvb_folder +'/tvb-library')
sys.path.append(tvb_folder + '/tvb-data')
from tvb.simulator.lab import *
sys.stdout,sys.stderr = nbso,nbse  # ...hack part 2/2
from tvb.simulator.plot.tools import plot_surface_mpl,plot_surface_mpl_mv


from copy import copy # seems that this needs to be done last otherwise it gets 
                      # replaced with numpy copy function
    
    
# aws api and workdocs-cloudfiles stuff

sys.path.append(le['ipynb_workdocs_dir'])
from new_utils import nb_fig,cloudfiles_nb

def get_roi_xyzs(roi_file,match_type='com'):
  """
  Get xyz centroids for a nifti image volume with labelled voxels
  """
    
  roi_img = nib.load(roi_file)
  roi_dat = roi_img.get_data()  
  roi_aff = roi_img.get_affine()
    
  roi_list = np.unique(roi_dat)
  if match_type == 'com':
    roi_all_inds = np.indices(roi_img.shape)
    roi_coords = [apply_affine(roi_aff, roi_all_inds[:,roi_dat==r].T) for r in roi_list]
    roi_xyzs = np.array([np.mean(rc,axis=0) for rc in roi_coords])
        
  elif match_type == 'xyzcuts': 
            
    roi_xyzs = np.array([find_xyz_cut_coords(roi_img,roi_dat==r,
                                             activation_threshold=0) for r in roi_list])

  return roi_xyzs,roi_list
    
    
    
def roi_to_surface(roi_xyzs,vtx_xyzs,faces):
  """
  Find nearest surface vertex and faces for a set of xyz locations
  """
    
  # average across the coords for each vertex comprising the face
  faces_xyz = np.mean(vtx_xyzs[faces],axis=1)
  # same as  fac_xyzs2 = np.array([np.mean(xyz[f],axis=0) for f in fac])    
    
  face_nearestroi_inds = cKDTree(roi_xyzs).query(faces_xyz, k=1)[1]
  #(gives the closest index in A for each point in B)
    
  vtx_nearestroi_inds = cKDTree(roi_xyzs).query(vtx_xyzs, k=1)[1]
    
  return vtx_nearestroi_inds,face_nearestroi_inds


def project_to_cs(roi_vals,vtx_nearestroi_inds,face_nearestroi_inds):
  """
  Take a set of roi values and a set of roi-vertex mappings and return 
  a of mapped vertex values 
  """
    
  vtx_vals = np.array([roi_vals[v] for v in vtx_nearestroi_inds])
  face_vals = np.array([roi_vals[v] for v in face_nearestroi_inds])
      
  return vtx_vals,face_vals
  

cnb = cloudfiles_nb('aws', [aws_key,aws_secret])
cnb.initialize_folder(nb_name)

os.chdir(outdir)

# Define some variables


tpl_dir = le['code_dir'] + '/libraries_of_others/misc/WFU_PickAtlas_3.0.5b/wfu_pickatlas/MNI_atlas_templates'
aal_txt_file = '%s/aal_MNI_V4.txt' %tpl_dir
aal_nii_file = '%s/aal_MNI_V4.nii' %tpl_dir

df_aal_txt = pd.read_csv(aal_txt_file, sep='\t', skiprows=[0],names=['region'])
df_aal_txt.index.names = ['id']
_tmp = pd.DataFrame(['None'],columns=['region']); _tmp.index.names = ['id']
df_aal_txt = df_aal_txt.append(_tmp).sort_index()


roi_xyzs,roi_list = get_roi_xyzs(aal_nii_file)

df_roi_xyz = pd.DataFrame(roi_xyz.T,['x', 'y', 'z'],
                          columns=df_aal_txt)

df_aal_txt_xyz = df_aal_txt.copy()
df_aal_txt_xyz['x'] = roi_xyz[:,0]
df_aal_txt_xyz['y'] = roi_xyz[:,1]
df_aal_txt_xyz['z'] = roi_xyz[:,2]


hemis = []
for r in df_aal_txt['region']: 
  if '_R' in r: hemis.append('R')
  elif '_L' in r: hemis.append('L')
  else: hemis.append('N')
                
hemis = np.array(hemis)

df_aal_txt_xyz['hemi'] = hemis

# Sort by hemisphere and y axis value
sortreg = df_aal_txt_xyz.sort_values(['hemi', 'y'])['region']

# (probably a simpler way of doing this...)
df_aal_txt_xyz = df_aal_txt_xyz.reset_index().set_index('region')\
                 .ix[sortreg].reset_index().set_index(['id', 'hemi'])

_rxyz = df_aal_txt_xyz[['x','y','z']].values

nr,nc = _rxyz.shape

eucldists = np.array([np.abs((_rxyz - np.tile(r,[nr,1]))).sum(axis=1) for r in _rxyz])

df_eucldists = pd.DataFrame(eucldists,columns=df_aal_txt_xyz['region'],
                            index=df_aal_txt_xyz['region'])

df_aal_txt_xyz_lh = df_aal_txt_xyz.query("hemi=='L'")
df_aal_txt_xyz_rh = df_aal_txt_xyz.query("hemi=='R'")

data_lh = df_eucldists.iloc[:,0].ix[df_aal_txt_xyz_lh['region']]
data_rh = df_eucldists.iloc[:,0].ix[df_aal_txt_xyz_rh['region']]

lhp_file = '/software/freesurfer/subjects/fsaverage/surf/lh.pial'
lhp_vtx,lhp_tri = nib.freesurfer.read_geometry(lhp_file)

rhp_file = '/software/freesurfer/subjects/fsaverage/surf/rh.pial'
rhp_vtx,rhp_tri = nib.freesurfer.read_geometry(rhp_file)

rlh = df_aal_txt_xyz_lh[['x', 'y', 'z']].values
rrh = df_aal_txt_xyz_rh[['x', 'y', 'z']].values

vtx_nrst_lh,tri_nrst_lh = roi_to_surface(rlh,lhp_vtx,lhp_tri)
vtx_nrst_rh,tri_nrst_rh = roi_to_surface(rrh,rhp_vtx,rhp_tri)

# that gave us correspondences in order of the roi coordinates list entered. 
# now get the corresponding id values

vtx_nrst_lh_id = roi_xyz_lh.index.get_level_values('id')[vtx_nrst_lh]
vtx_nrst_rh_id = roi_xyz_rh.index.get_level_values('id')[vtx_nrst_rh]

data_lh_vtx,data_lh_tri = project_to_cs(data_lh,vtx_nrst_lh,tri_nrst_lh)
data_rh_vtx,data_rh_tri = project_to_cs(data_rh,vtx_nrst_rh,tri_nrst_rh)

%%matlab -i outdir,lhp_file,rhp_file,data_lh_vtx,data_rh_vtx

fig = figure; 

[lhp_vtx_ml,lhp_tri_ml] = read_surf(lhp_file);
[rhp_vtx_ml,rhp_tri_ml] = read_surf(rhp_file);

viewfromleft = [-90,10];
viewfromright = [90,10];
viewfromtop = [0,90]; 
viewfrombottom = [0,-90];

subplot(2,2,1)
patch('Vertices',lhp_vtx_ml,'Faces',lhp_tri_ml+1,'CData',data_lh_vtx,...
      'FaceColor','interp','EdgeColor','none');
view(viewfromleft);
axis image; axis off; camlight; lighting gouraud; 
colormap('jet'); caxis auto;

subplot(2,2,2)
patch('Vertices',rhp_vtx_ml,'Faces',rhp_tri_ml+1,'CData',data_rh_vtx,...
      'FaceColor','interp','EdgeColor','none');
view(viewfromright);
axis image; axis off; camlight; lighting gouraud; 
colormap('jet'); caxis auto;

subplot(2,2,3)
patch('Vertices',lhp_vtx_ml,'Faces',lhp_tri_ml+1,'CData',data_lh_vtx,...
      'FaceColor','interp','EdgeColor','none');
view(viewfromright);
axis image; axis off; camlight; lighting gouraud; 
colormap('jet'); caxis auto;

subplot(2,2,4)
patch('Vertices',rhp_vtx_ml,'Faces',rhp_tri_ml+1,'CData',data_rh_vtx,...
      'FaceColor','interp','EdgeColor','none');
view(viewfromleft);
axis image; axis off; camlight; lighting gouraud; 
colormap('jet'); caxis auto;

print(fig, '-dpng',[outdir '/aal_euclidean_distances_patch.png']);
                    
close(fig)