#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Script to compute powder average (mean diffusion weighted image) from set of
diffusion images.

By default will output an average image calculated from all images with
non-zero bvalue.

specify --bvalue to output an image for a single shell

Script currently does not take into account the diffusion gradient directions
being averaged.
"""

import argparse
import logging

import nibabel as nib

import numpy as np

from dipy.io.gradients import read_bvals_bvecs

from scilpy.io.image import (get_data_as_mask, assert_same_resolution)
from scilpy.io.utils import (add_overwrite_arg, assert_inputs_exist,
                             assert_outputs_exist, add_verbose_arg)

from scilpy.utils.bvec_bval_tools import get_shell_indices

logger = logging.getLogger("Compute_Powder_Average")
logger.setLevel(logging.INFO)


def _build_arg_parser():
    p = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter)
    p.add_argument('in_dwi',
                   help='Path of the input diffusion volume.')
    p.add_argument('in_bval',
                   help='Path of the bvals file, in FSL format.')
    p.add_argument('out_avg',
                   help='Path of the output file.')

    add_overwrite_arg(p)

    p.add_argument('--mask', metavar='file',
                   help='Path to a binary mask.\nOnly data inside the'
                   ' mask will be used for powder avg. '
                   '(Default: %(default)s)')

    p.add_argument('--b0_thr', type=int, default='50',
                   help='Exclude b0 volumes from powder average with'
                   ' bvalue less than specified threshold.\n'
                   '(Default: remove volumes with bvalue < %(default)s')

    p.add_argument('--shells', nargs='+', type=int, default=None,
                   help='bvalue (shells) to include in powder average'
                   ' passed as a list \n(e.g. --shells 1000 2000). '
                   'If not specified will include all volumes with'
                   ' a non-zero bvalue.')

    p.add_argument('--shell_thr', type=int, default='50',
                   help='Include volumes with bvalue +- the specified'
                   ' threshold.\n(Default: [%(default)s]')

    add_verbose_arg(p)

    return p


def main():
    parser = _build_arg_parser()
    args = parser.parse_args()
    inputs = [args.in_dwi, args.in_bval]
    if args.mask:
        inputs.append(args.mask)

    assert_inputs_exist(parser, inputs)
    assert_outputs_exist(parser, args, args.out_avg)

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    img = nib.load(args.in_dwi)
    data = img.get_fdata(dtype=np.float32)
    affine = img.affine
    if args.mask is None:
        mask = None
    else:
        mask_img = nib.load(args.mask)
        assert_same_resolution((img, mask_img))
        mask = get_data_as_mask(mask_img, dtype='uint8')

    # Read bvals (bvecs not needed at this point)
    logging.info('Performing powder average')
    bvals, _ = read_bvals_bvecs(args.in_bval, None)

    # Select diffusion volumes to average
    if not (args.shells):
        # If no shell given, average all diffusion weighted images
        pwd_avg_idx = np.squeeze(np.where(bvals > 0 + args.b0_thr))
        logging.debug('Calculating powder average from all diffusion'
                      '-weighted volumes, {} volumes '
                      'included.'.format(len(pwd_avg_idx)))
    else:
        pwd_avg_idx = []
        logging.debug('Calculating powder average from {} '
                      'shells {}'.format(len(args.shells), args.shells))

        for shell in args.shells:
            pwd_avg_idx = np.int64(
                np.concatenate((pwd_avg_idx,
                                get_shell_indices(bvals,
                                                  shell,
                                                  tol=args.shell_thr))))
            logging.debug('{} b{} volumes detected and included'.format(
                len(pwd_avg_idx), shell))

        # remove b0 indices
        b0_idx = get_shell_indices(bvals, 0, args.b0_thr)
        logging.debug('{} b0 volumes detected and not included'.format(
            len(b0_idx)))
        for val in b0_idx:
            pwd_avg_idx = pwd_avg_idx[pwd_avg_idx != val]

    if len(pwd_avg_idx) == 0:
        raise ValueError('No shells selected for powder average, ensure '
                         'shell, shell_thr and b0_thr are set '
                         'appropriately')

    powder_avg = np.squeeze(np.mean(data[:, :, :, pwd_avg_idx], axis=3))

    if args.mask:
        powder_avg = powder_avg * mask

    powder_avg_img = nib.Nifti1Image(powder_avg.astype(np.float32), affine)
    nib.save(powder_avg_img, args.out_avg)


if __name__ == "__main__":
    main()
