#!/usr/bin/env fslpython

import sys,os,glob,subprocess,shutil,tempfile
import pandas as pd
import numpy as np
import nibabel as nib
import argparse, textwrap
from tqdm import tqdm

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

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

def imgtest(fname):
    r = subprocess.run([f'{os.path.join(FSLbin, "imtest")} {fname}'], capture_output=True, text=True, shell=True)
    return int(r.stdout)

def applywarp(fin, out, ref, warp):
    r = subprocess.run([os.path.join(FSLbin, "applywarp"), '-i', fin, '-o', out, '-r', ref, '-w', warp, '-d', 'float', '--interp=spline'])

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


splash = r"""
__  _______ ____      _    ____ _____    _        _
\ \/ /_   _|  _ \    / \  / ___|_   _|__| |_ __ _| |_ ___
 \  /  | | | |_) |  / _ \| |     | |/ __| __/ _  | __/ __|
 /  \  | | |  _ <  / ___ \ |___  | |\__ \ || (_| | |_\__ \\
/_/\_\ |_| |_| \_\/_/   \_\____| |_||___/\__\__ _|\__|___/

 """
print(splash)

parser = MyParser(prog='XTRACT Stats',
                  description='xtract_stats: summary tract-wise measures',
                  formatter_class=argparse.RawDescriptionHelpFormatter,
                  epilog=textwrap.dedent('''Example usage:
                        xtract_stats -xtract /data/xtract -d /data/Diffusion/dti_ -warp std2diff.nii.gz
                                         '''))

required = parser.add_argument_group('Required arguments')
optional = parser.add_argument_group('Optional arguments')

# Compulsory arguments:
required.add_argument("-xtract", metavar='<folder>', help="Path to XTRACT output folder", required=True)
required.add_argument("-d", metavar='<folder_basename>', help="Path to microstructure folder and basename of data (e.g. /home/DTI/dti_)", required=True)

# Optional arguments:
optional.add_argument("-warp", metavar='<nifti>', nargs=2, help="If tract are not in the same space as the data in '-d', provide the reference image in diffusion space (e.g. /home/DTI/dti_FA) and the xtract space to reference space warp field")
optional.add_argument("-out", metavar='<path>', help="Output filepath (Default <XTRACT_dir>/stats.csv)")
optional.add_argument("-meas", metavar='<list>', default='vol,prob,length,FA,MD', help="Comma separated list of features to extract (Default = vol,prob,length,FA,MD - assumes DTI folder has been provided)\nvol = tract volume, prob = tract probability, length = tract length\nAdditional metrics must follow file naming conventions. e.g. for dti_L1 use 'L1'")
optional.add_argument("-tract_list", metavar='<list>', help="Comma separated list of tracts to include (default = all found under -xtract <folder>)")
optional.add_argument("-thr", metavar='<float>', default=0.001, type=float, help="Threshold applied to tracts (default = 0.001).")
argsa = parser.parse_args()

xtract = argsa.xtract
thr = str(argsa.thr)
meas_dir = argsa.d
meas = argsa.meas.split(',')

# Check compulsory arguments
errflag = 0
if os.path.isdir(xtract) is False:
    print(f'xtract folder {xtract} not found')
    errflag=1
if os.path.isdir(os.path.dirname(meas_dir)) is False:
    print(f'Data folder {meas_dir} not found')
    errflag=1
errchk(errflag)
    
# check optional arguments
if argsa.warp is not None:
    ref = argsa.warp[0]
    warp = argsa.warp[1]
    if imgtest(ref) == 0:
        print(f'Reference image {ref} not found.')
        errflag=1
    if imgtest(warp) == 0:
        print(f'Warp field {warp} not found.')
        errflag=1
else:
    warp = 'native'
errchk(errflag)

if argsa.out == None:
    out = os.path.join(xtract, 'stats.csv')
else:
    out = argsa.out

if os.path.isfile(out) is True:
    print('Warning: Output file already exsits. Will be overwritten.')

# get tract list
if argsa.tract_list is not None:
    tracts = argsa.tract_list.split(',')
else:
    print('Getting stats for all tracts under "-xtract"')
    tracts = glob.glob(os.path.join(xtract, 'tracts', '*', 'densityNorm.nii.gz'))
    tracts = [os.path.basename(os.path.dirname(t)) for t in tracts]
    tracts.sort()

# report tracts and measures being included
print(f'Tracts: {tracts}')
print(f'\nGetting summary stats for: {meas}')

# Location of temp folder for warped tracts
temp_dir = tempfile.TemporaryDirectory()

# expand meas_list to contain mean, median, standard deviation
# for 'vol' volume, only get -V
# for 'length' and 'prob' get mean, median and std
# for all other measures get mean, median and std and the probability-weighted stats
meas_hdr = []
for m in meas:
    if m == 'vol':
        # just get tract volume
        meas_hdr.append('volume')
    elif m in ['prob', 'length']:
        # get mean, median and std
        meas_hdr.append(f'{m}_mean')
        meas_hdr.append(f'{m}_median')
        meas_hdr.append(f'{m}_std')
    else:
        # get mean, median and std from binary mask and weighted by tract probability
        meas_hdr.append(f'{m}_mean')
        meas_hdr.append(f'{m}_median')
        meas_hdr.append(f'{m}_std')
        meas_hdr.append(f'{m}_WP_mean')
        meas_hdr.append(f'{m}_WP_median')
        meas_hdr.append(f'{m}_WP_std')

stats = pd.DataFrame(index=tracts, columns=meas_hdr)
print('\nProcessing...')
progression_bar = tqdm(tracts)
for t in (progression_bar):
    progression_bar.set_description(f'{t}')
    tpath = os.path.join(xtract, 'tracts', t, 'densityNorm.nii.gz')
    tlengthpath = os.path.join(xtract, 'tracts', t, 'density_lengths.nii.gz')
    # deal with warping
    if warp != 'native':
        progression_bar.set_description(f'{t}: warping')
        tpath = os.path.join(temp_dir.name, 'densityNorm_diffspace.nii.gz')
        applywarp(os.path.join(xtract, 'tracts', t, 'densityNorm.nii.gz'), tpath, ref, warp)
        # only warp tract_lengths if using it
        if 'length' in meas:
            tlengthpath = os.path.join(temp_dir.name, 'density_lengths_diffspace.nii.gz')
            applywarp(os.path.join(xtract, 'tracts', t, 'density_lengths.nii.gz'), tlengthpath, ref, warp)
    # binarise prob map
    tbinpath = tpath.replace('.nii.gz', '_bin.nii.gz')
    cmout = subprocess.run([os.path.join(FSLbin, "fslmaths"), tpath, '-thr', thr, '-bin', tbinpath])
    # get stats
    for m in meas:
        progression_bar.set_description(f'{t}: getting metrics')
        if m == 'vol':
            # tract volume (mm)
            cmout = subprocess.run([os.path.join(FSLbin, "fslstats"), tpath, '-l', thr, '-V'], capture_output=True, text=True)
            stats.at[t, 'volume'] =  float(cmout.stdout.split(' ')[1])
        elif m == 'prob':
            # median, mean, and standard deviation of probability
            cmout = subprocess.run([os.path.join(FSLbin, "fslstats"), tpath, '-l', thr, '-M', '-P', '50', '-S'], capture_output=True, text=True)
            # cmout.stdout.split: 0 is mean, 1 is median, 2 is std
            stats.at[t, 'prob_mean'] =  cmout.stdout.split(' ')[0]
            stats.at[t, 'prob_median'] =  cmout.stdout.split(' ')[1]
            stats.at[t, 'prob_std'] =  cmout.stdout.split(' ')[2]
        elif m == 'length':
            # median, mean, and standard deviation of tract lengths
            cmout = subprocess.run([os.path.join(FSLbin, "fslstats"), tlengthpath, '-l', thr, '-M', '-P', '50', '-S'], capture_output=True, text=True)
            stats.at[t, 'length_mean'] =  cmout.stdout.split(' ')[0]
            stats.at[t, 'length_median'] =  cmout.stdout.split(' ')[1]
            stats.at[t, 'length_std'] =  cmout.stdout.split(' ')[2]
        else:
            mpath = meas_dir + m
            # get mean, median and std from binary mask
            cmout = subprocess.run([os.path.join(FSLbin, "fslstats"), mpath, '-k', tbinpath, '-M', '-P', '50', '-S'], capture_output=True, text=True)
            stats.at[t, f'{m}_mean'] =  cmout.stdout.split(' ')[0]
            stats.at[t, f'{m}_median'] =  cmout.stdout.split(' ')[1]
            stats.at[t, f'{m}_std'] =  cmout.stdout.split(' ')[2]
            # get mean, median and std from binary mask weighted by tract probability
            mpath_temp = os.path.join(temp_dir.name, f'{m}_{t}')
            cmout = subprocess.run([os.path.join(FSLbin, "fslmaths"), mpath, '-mul', tpath, mpath_temp])
            cmout = subprocess.run([os.path.join(FSLbin, "fslstats"), mpath_temp, '-k', tbinpath, '-M', '-P', '50', '-S'], capture_output=True, text=True)
            stats.at[t, f'{m}_WP_mean'] =  cmout.stdout.split(' ')[0]
            stats.at[t, f'{m}_WP_median'] =  cmout.stdout.split(' ')[1]
            stats.at[t, f'{m}_WP_std'] =  cmout.stdout.split(' ')[2]
  

# save output
stats.to_csv(out, index_label='tract')
print(f'\nOutput saved to {out}')

temp_dir.cleanup()
quit()