#!/usr/bin/env fslpython

import sys,os,glob,subprocess,re,shutil
import argparse, textwrap

import numpy as np
import pandas as pd
import scipy.sparse as sps

# Image stuff
import nibabel as nib
from nibabel import cifti2
from fsl.data.image import Image
from fsl.data.cifti import cifti2_axes
from fsl.data.cifti import Cifti

FSLDIR = os.getenv('FSLDIR')
FSLbin = os.path.join(FSLDIR, 'bin')

# normalise blueprint
def normalise(M):
    D       = np.sum(M,axis=1)
    D[D==0] = 1
    M       = M / D[:,None]
    return M

# load fdt matrix and convert to dense
def load_fdt(fdt_path, norm=True):
    print('Reading tractography matrix. This may take a few minutes...')
    # mat         = pd.read_csv(fdt_path, header=None, dtype={0:np.int32, 1:np.int32, 2:np.float32}, delim_whitespace=True).to_numpy()
    mat         = pd.read_csv(fdt_path, header=None, dtype={0:np.int32, 1:np.int32, 2:np.float32}, sep='\s+').to_numpy()
    data        = mat[:-1, -1]
    rows        = np.array(mat[:-1, 0]-1, dtype=int)
    cols        = np.array(mat[:-1, 1]-1, dtype=int)
    nrows,ncols = int(mat[-1, 0]), int(mat[-1,1])
    M           = sps.csc_matrix((data, (rows,cols)), shape=(nrows,ncols)).toarray()
    if norm == True:
        M       = normalise(M)
    return M

# function to find which hemisphere/cifti structure we're working with
def get_structure(p):
    f = open(p, 'r')
    text = f.read()
    text = text.split('Cortex',1)[1]
    cstruct = text.split(']]></Value>',1)[0]
    return cstruct

# some useful functions
def errchk(errflag):
    if errflag:
        print("Exit without doing anything..")
        quit()

class MyParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n' % message)
        self.print_help()
        sys.exit(2)

class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


parser = MyParser(prog='XTRACT Blueprint',
                  description='xtract_blueprint: for generating connectivity blueprints from xtract output',
                  epilog='')

# Compulsory arguments:
parser.add_argument("-xtract_folder", metavar='<folder>', help="Path to xtract folder", required=True)
parser.add_argument("-ptx_folder", metavar='<list>', help="the ptx_folder(s)- if >1, comma separated", required=True)
parser.add_argument("-seed_path", metavar='<list>', help="the seed mask(s) - if >1, comma separated", required=True)
parser.add_argument("-tracts", metavar='<list>', help="Comma separated list of tracts to include", required=True)
parser.add_argument("-roi_path", metavar='<list>', help="using the medial wall as a mask(s)? - if >1, comma separated")
parser.add_argument("-savetxt", action="store_true", default=False, help="save as txt?")
parser.add_argument("-thr", metavar='<float>', default=0.001, type=float, help="Threshold applied to XTRACT tracts prior to blueprint calculation (default = 0.001)")
# parser.add_argument("-distnorm", action="store_true", default=False, help="distance normalisation?")
parser.add_argument("-prefix", metavar='<string>', help="prefix for filename")
argsa = parser.parse_args()

######################
####### Some preparation stuff
######################
CIFTI_STRUCTURES = ['CEREBELLUM', 'ACCUMBENS_LEFT', 'ACCUMBENS_RIGHT', 'AMYGDALA_LEFT', 'AMYGDALA_RIGHT',
                    'BRAIN_STEM', 'CAUDATE_LEFT', 'CAUDATE_RIGHT', 'CEREBELLUM_LEFT', 'CEREBELLUM_RIGHT',
                    'DIENCEPHALON_VENTRAL_LEFT', 'DIENCEPHALON_VENTRAL_RIGHT', 'HIPPOCAMPUS_LEFT',
                    'HIPPOCAMPUS_RIGHT', 'OTHER', 'OTHER_GREY_MATTER', 'PALLIDUM_LEFT', 'PALLIDUM_RIGHT',
                    'PUTAMEN_LEFT', 'PUTAMEN_RIGHT', 'THALAMUS_LEFT', 'THALAMUS_RIGHT']

# split the arguments here
xtract_folder = argsa.xtract_folder
ptx_folder = argsa.ptx_folder.split(",")
seed_path = argsa.seed_path.split(",")
tracts = argsa.tracts.split(",")
ntracts = len(tracts)

if argsa.prefix is None:
    prefix = ''

# determine if volume or surface seed: if volume, force savetxt
# 0 is surface, 1 is volume, 2 is surface plus subcortical
if '.nii' in seed_path[0]:
    seed_type = 1
    seed_geom = "voxels"
elif '.gii' in seed_path[0] or '.asc' in seed_path[0]:
    seed_type = 0
    seed_geom = "vertices"

if seed_type == 0:
    if len(ptx_folder) == 2:
        print('Building whole-brain connectivity blueprint')
    elif len(ptx_folder) == 3:
        seed_type = 2
        seed_geom = "vertices+voxels"
        # load subseed list with number of voxels per structure
        subseed_list = pd.read_csv(os.path.join(ptx_folder[2], "VoxNum_Subcort_ROIs"), sep=' ', names=['structure', 'nvox'])
        print(f'Building whole-brain plus subcortical connectivity blueprint with {subseed_list.shape[0]} subcortical structure(s)')
    else:
        print('Building single hemisphere connectivity blueprint')
else:
    print('Building connectivity blueprint using volume seed')

######################
####### Load and prepare data for cortical blueprint
######################
# mask and lut are equal across hemispheres, so only need 1
maskfile = os.path.join(ptx_folder[0], 'lookup_tractspace_fdt_matrix2')
mask     = Image(maskfile)
lut      = Image(os.path.join(ptx_folder[0], 'lookup_tractspace_fdt_matrix2.nii.gz'))
lut      = lut.data[mask.data>0]-1

out_folder = os.path.dirname(os.path.dirname(ptx_folder[0]))

# Get num voxels in mask
nvoxels = np.sum(mask.data>0)

print(f'Generating blueprint with {ntracts} tracts')
# Collect tracts
tracts_mat = np.zeros((nvoxels, ntracts))
print('Reading tracts...')
trm=[]
for idx,t in enumerate(tracts):
    t = os.path.join(xtract_folder, f'{t}.nii.gz')
    if os.path.exists(t):
        tract = Image(t)
        tracts_mat[lut,idx] = tract[mask.data>0]
    else:
        print(f'Could not find {tracts[idx]}. Skipping... (will remove from final output)')
        trm.append(int(idx))

# remove missing tract columns
tracts_mat = np.delete(tracts_mat, [trm], 1)
for i in trm:
    tracts.pop(i)

ntracts = len(tracts)

# Open matrix2 file
M = []
for p in ptx_folder[0:2]:
    if os.path.exists(os.path.join(p, "fdt_matrix2.dot.gz")):
        foo = subprocess.run(['gunzip', os.path.join(p, "fdt_matrix2.dot.gz")])
    M.append(load_fdt(os.path.join(p, "fdt_matrix2.dot")))
    foo = subprocess.run(['gzip', os.path.join(p, "fdt_matrix2.dot"), "--fast"])

if len(ptx_folder) >= 2:
    print('Stacking hemispheres...')
    Mdims = [M[0].shape[0], M[1].shape[0]]
    M = np.concatenate((M[0], M[1]))
else:
    Mdims = M[0].shape[0]
    M = M[0]

######################
####### Create cortical blueprint
######################
print('Calculating blueprint...')
BP = M@tracts_mat # the cortical blueprint


######################
####### Create subcortical blueprint
######################
if seed_type == 2:
    print('Preparing subcortical blueprint...')
    nsubvox = np.sum(subseed_list.nvox)
    # now load all_structures_all_tracts and convert to blueprint format
    subcort_bp = nib.load(os.path.join(ptx_folder[2], "all_structures_all_tracts.nii.gz"))
    subcort_bp_img = subcort_bp.get_fdata(dtype=np.float32)
    subcort_bp_flat = np.zeros((nsubvox, ntracts))
    coords = np.zeros((nsubvox, 3))
    start_idx = 0
    for ind, row in subseed_list.iterrows():
        str_mask = nib.load(row.structure).get_fdata(dtype=np.float32)
        getind = np.argwhere(str_mask > 0)
        for i, t in enumerate(tracts):
            subcort_bp_flat[start_idx:start_idx+row.nvox, i] = subcort_bp_img[[g[0] for g in getind], [g[1] for g in getind], [g[2] for g in getind], i]
        
        coords[start_idx:start_idx+row.nvox, :] = getind
        start_idx += row.nvox
    
    ######################
    ####### Concatenate blueprints
    ######################
    BP = np.concatenate((BP, subcort_bp_flat))

######################
####### Nroamlise blueprint
######################
BP = normalise(BP)

print(f'Blueprint dimensions: {BP.shape[0]} greyordinates, {BP.shape[1]} tracts')

######################
####### Save blueprint
######################
# if surface seeding, load ROIs to build CIFTI/add in empty medial wall
if seed_type == 0:
    # load the seed ROIs
    seed = []
    for p in seed_path:
        temp = nib.load(p).darrays[0].data != 0
        seed.append(temp[:,0])

    # Convert to full cortex structure here (i.e. add in empty medial wall)
    if argsa.roi_path is not None:
        roi_path = argsa.roi_path.split(",")
        roi = []
        for p in roi_path:
            temp = nib.load(p).darrays[0].data
            if np.unique(temp).shape[0] > 2:
                print('Warning!! Medial wall mask is not binary.')
                print('Binarising...')
                temp = (temp > 0)*1
            
            roi.append(temp)
        
        # and stack
        if len(ptx_folder) == 2:
            roi = np.concatenate((roi[0],roi[1]))
        else:
            roi = roi[0]

        # and stack
        if len(ptx_folder) == 2:
            seed_temp = np.concatenate((seed[0],seed[1]))
        else:
            seed_temp = seed[0]

        full_BP = np.zeros([np.shape(seed_temp)[0], np.shape(BP)[1]])
        full_BP[roi == 1, :] = BP
        BP = full_BP
    if argsa.savetxt:
        if len(ptx_folder) == 1:
            cstruct = get_structure(seed_path[0])
            if cstruct == 'Left':
                side='L'
            else:
                side='R'

            new_fname = os.path.join(out_folder, f'{prefix}BP.{side}.txt')
            np.savetxt(new_fname, BP)
            new_fname_tord = os.path.join(out_folder, f'{prefix}tract_order.{side}.txt')
            np.savetxt(new_fname_tord, tracts, fmt="%s")
        elif len(ptx_folder) == 2:
            new_fname = os.path.join(out_folder, f'{prefix}BP.LR.txt')
            np.savetxt(new_fname, BP)
            new_fname_tord = os.path.join(out_folder, f'{prefix}tract_order.LR.txt')
            np.savetxt(new_fname_tord, tracts, fmt="%s")
    else:
        # set up cifti brain model axes
        if len(ptx_folder) == 1:
            cstruct = get_structure(seed_path[0])
            if cstruct == 'Left':
                side='L'
            else:
                side='R'
            bm        = cifti2_axes.BrainModelAxis.from_mask(seed[0], name=f'Cortex{cstruct}')
            new_fname = os.path.join(out_folder, f'{prefix}BP.{side}.dscalar.nii')
        elif len(ptx_folder) == 2:
            bm_l      = cifti2_axes.BrainModelAxis.from_mask(seed[0], name=f'CortexLeft')
            bm_r      = cifti2_axes.BrainModelAxis.from_mask(seed[1], name=f'CortexRight')
            bm        = bm_l + bm_r
            new_fname = os.path.join(out_folder, f'{prefix}BP.LR.dscalar.nii')

        # save cifti
        sc        = cifti2_axes.ScalarAxis(tracts)
        hdr       = cifti2.Cifti2Header.from_axes((sc, bm))
        img       = cifti2.Cifti2Image(BP.T, hdr)
        img.update_headers()
        img.nifti_header.set_intent(3006) # the dscalar intent code
        nib.save(img, new_fname)
# else, if volume seeding, default save as nifti
elif seed_type == 1:
    # for nifti output, save tract_order always
    new_fname_tord = os.path.join(out_folder, f'{prefix}tract_order.txt')
    np.savetxt(new_fname_tord, tracts, fmt="%s")

    # if savetxt save blueprint as txt
    if argsa.savetxt:
        new_fname = os.path.join(out_folder, f'{prefix}BP.LR.txt')
        np.savetxt(new_fname, BP)
    else:
        # read the volume seed and use coords file to fill voxel data
        vol_seed = nib.load(seed_path[0])
        vol_seed_img = vol_seed.get_fdata()
        coords = np.loadtxt(os.path.join(ptx_folder[0], 'coords_for_fdt_matrix2'), dtype='int')

        # create empty blueprint nifti and set each volume as a tract
        BP_nii = np.zeros((vol_seed_img.shape[0], vol_seed_img.shape[1], vol_seed_img.shape[2], ntracts))
        for idx,t in enumerate(tracts):
            BP_nii[coords[:,0], coords[:,1], coords[:,2], idx] = BP[:, idx]

        new_fname = os.path.join(out_folder, f'{prefix}BP.LR.nii.gz')

        # prepare nifti file, change header and substitute data from vol_seed for output
        vol_seed.header['dim'][0] = 4
        vol_seed.header['dim'][4] = ntracts
        img = nib.Nifti1Image(BP_nii, vol_seed.get_sform(), header=vol_seed.header)
        nib.save(img, new_fname)
# else, if surface+subcortical seeding, save as full CIFTI structure
elif seed_type == 2:
    # load the seed ROIs
    seed = []
    for p in seed_path[0:2]:
        temp = nib.load(p).darrays[0].data != 0
        seed.append(temp[:,0])

    # Convert to full cortex structure here (i.e. add in empty medial wall)
    if argsa.roi_path is not None:
        roi_path = argsa.roi_path.split(",")
        roi = []
        for p in roi_path:
            temp = nib.load(p).darrays[0].data
            if np.unique(temp).shape[0] > 2:
                print('Warning!! Medial wall mask is not binary.')
                print('Binarising...')
                temp = (temp > 0)*1
            
            roi.append(temp)

        roi = np.concatenate((roi[0],roi[1]))
        # and append the subcortical seed
        roi = np.concatenate((roi, np.ones([nsubvox,])))
        # and stack
        seed_temp = np.concatenate((seed[0], seed[1], np.ones([nsubvox,], dtype=bool)))
        full_BP = np.zeros([np.shape(seed_temp)[0], np.shape(BP)[1]])
        full_BP[roi == 1, :] = BP
        BP = full_BP

    if argsa.savetxt:
        new_fname = os.path.join(out_folder, f'{prefix}BP.LR.txt')
        np.savetxt(new_fname, BP)
        new_fname_tord = os.path.join(out_folder, f'{prefix}tract_order.LR.txt')
        np.savetxt(new_fname_tord, tracts, fmt="%s")
    else:
        # set up cifti brain model axes
        bm_l      = cifti2_axes.BrainModelAxis.from_mask(seed[0], name=f'CortexLeft')
        bm_r      = cifti2_axes.BrainModelAxis.from_mask(seed[1], name=f'CortexRight')
        bm        = bm_l + bm_r
        ref_img = nib.load(subseed_list.structure[0])
        start_idx = 0
        for ind, row in subseed_list.iterrows():
            struct_name = row.structure.split('CIFTI_STRUCTURE_')[1].replace('.nii.gz', '').replace('.nii', '')
            struct_name = struct_name.upper()
            if struct_name not in CIFTI_STRUCTURES:
                struct_name = 'OTHER'
            bm += nib.cifti2.BrainModelAxis(
                name=struct_name,
                voxel=coords[start_idx : start_idx + row.nvox, :3],
                affine=ref_img.affine,
                volume_shape=ref_img.shape,
            )
            start_idx += row.nvox

        # save cifti
        new_fname = os.path.join(out_folder, f'{prefix}BP.LR.dscalar.nii')
        sc        = cifti2_axes.ScalarAxis(tracts)
        hdr       = cifti2.Cifti2Header.from_axes((sc, bm))
        img       = cifti2.Cifti2Image(BP.T, hdr)
        img.update_headers()
        img.nifti_header.set_intent(3006) # the dscalar intent code
        nib.save(img, new_fname)

print(f'Saved: {new_fname}')
print('')
quit()
