#!/usr/bin/env python
"""
Computes the number of BVALs and B0s for the input scan, and plots the DTI
directions (BVECs) as a 3D scatterplot. This scatterplot will be output in the
same directory as the output.csv file.

Usage:
    qc-dti <nifti> <bvec> <bval> <output_prefix>

Arguments:
    <nifti>           Input DTI nifti scan
    <bvec>            Input bvec file
    <bval>            Input bval file
    <output_prefix>   prefix of all output files

Outputs:
    <output_prefix>_stats.csv
    <output_prefix>_directions.png
    <output_prefix>_qascripts_dti.csv

[1] Roalf et al. The impact of quality assurance assessment on diffusion tensor
imaging outcomes in a large-scale population-based cohort. Neuroimage. 2016.
"""

import os,sys
import logging
import qcmon as qc
import numpy as np
import matplotlib
matplotlib.use('Agg')   # Force matplotlib to not use any Xwindows backend
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from qcmon.utilities import make_tmpdir

logging.basicConfig(level=logging.WARN, format="[%(name)s] %(levelname)s: %(message)s")
logger = logging.getLogger(os.path.basename(__file__))

def absmax(x):
    """
    Finds the maximum absoloute value of the input vector x.
    """
    return(np.max((np.abs(np.max(x)), np.abs(np.min(x)))))

def float_to_int(filename):
    '''
    Convert space delimited text file from float to int
    '''

    with open(filename,'r') as f:
        raw_file = f.read().strip().split(' ')

    int_list = [str(int(float(b))) for b in raw_file]
    write_str = ' '.join(int_list)

    with open(filename,'w') as f:
        f.write(write_str)

@make_tmpdir
def main(nifti, bvec, bval, output_prefix, tmpdir=None):

    # add qascript to path
    qascriptsPath = os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir, 'assets/qascripts_version2')
    os.environ['PATH'] += os.pathsep + qascriptsPath

    # run metrics from qascripts toolchain
    qc.utilities.run('cp {nifti} {t}/dti.nii.gz'.format(nifti=nifti, t=tmpdir))
    qc.utilities.run('cp {bval} {t}/dti.bval'.format(bval=bval, t=tmpdir))
    qc.utilities.run('cp {bvec} {t}/dti.bvec'.format(bvec=bvec, t=tmpdir))
    float_to_int(os.path.join(tmpdir,'dti.bval'))
    qc.utilities.run('qa_dti_v2.sh {t}/dti.nii.gz {t}/dti.bval {t}/dti.bvec {t}/qc_dti.csv'.format(t=tmpdir))
    qc.utilities.run('mv {t}/qc_dti.csv {output_prefix}_qascripts_dti.csv'.format(t=tmpdir, output_prefix=output_prefix))

    # load in data for remaining QC
    bvec = np.genfromtxt(bvec)
    bval = np.genfromtxt(bval)

    # HOTFIX: Account for small deviations in observed bvals
    bval[np.abs(bval) < 10] = 0
    idxDirs = np.where(bval > 0)[0]
    nDirs = len(idxDirs)
    nB0s = len(bval) - len(idxDirs)

    # 3d scatterplot of directions
    xlim = absmax(bvec[0, idxDirs])
    ylim = absmax(bvec[1, idxDirs])
    zlim = absmax(bvec[2, idxDirs])

    fig = plt.figure(figsize=(5, 5))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(bvec[0, idxDirs], bvec[1, idxDirs], bvec[2, idxDirs])
    ax.set_xlim3d([-xlim, xlim])
    ax.set_ylim3d([-ylim, ylim])
    ax.set_zlim3d([-zlim, zlim])
    plt.savefig(os.path.join(output_prefix + '_directions.png'))

    with open(output_prefix + '_stats.csv', 'w') as f:
        f.write('#ndirs,nb0s\n{},{}'.format(nDirs, nB0s))

if __name__ == '__main__':
    if len(sys.argv) == 5:
        main(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
    else:
        print(__doc__)
