# -*- coding: utf-8 -*-
"""Change directory to provide relative paths for doctests
   >>> import os
   >>> filepath = os.path.dirname( os.path.realpath( __file__ ) )
   >>> datadir = os.path.realpath(os.path.join(filepath, '../../testing/data'))
   >>> os.chdir(datadir)
"""
from __future__ import print_function, division, unicode_literals, absolute_import
from multiprocessing import (Pool, cpu_count)
import os.path as op
from builtins import range

import nibabel as nb

from ... import logging
from ..base import (traits, TraitedSpec, BaseInterfaceInputSpec,
                    File, InputMultiPath, isdefined)
from .base import DipyBaseInterface
IFLOGGER = logging.getLogger('interface')


class SimulateMultiTensorInputSpec(BaseInterfaceInputSpec):
    in_dirs = InputMultiPath(File(exists=True), mandatory=True,
                             desc='list of fibers (principal directions)')
    in_frac = InputMultiPath(File(exists=True), mandatory=True,
                             desc=('volume fraction of each fiber'))
    in_vfms = InputMultiPath(File(exists=True), mandatory=True,
                             desc=('volume fractions of isotropic '
                                   'compartiments'))
    in_mask = File(exists=True, desc='mask to simulate data')

    diff_iso = traits.List(
        [3000e-6, 960e-6, 680e-6], traits.Float, usedefault=True,
        desc='Diffusivity of isotropic compartments')
    diff_sf = traits.Tuple(
        (1700e-6, 200e-6, 200e-6),
        traits.Float, traits.Float, traits.Float, usedefault=True,
        desc='Single fiber tensor')

    n_proc = traits.Int(0, usedefault=True, desc='number of processes')
    baseline = File(exists=True, mandatory=True, desc='baseline T2 signal')
    gradients = File(exists=True, desc='gradients file')
    in_bvec = File(exists=True, desc='input bvecs file')
    in_bval = File(exists=True, desc='input bvals file')
    num_dirs = traits.Int(32, usedefault=True,
                          desc=('number of gradient directions (when table '
                                'is automatically generated)'))
    bvalues = traits.List(traits.Int, value=[1000, 3000], usedefault=True,
                          desc=('list of b-values (when table '
                                'is automatically generated)'))
    out_file = File('sim_dwi.nii.gz', usedefault=True,
                    desc='output file with fractions to be simluated')
    out_mask = File('sim_msk.nii.gz', usedefault=True,
                    desc='file with the mask simulated')
    out_bvec = File('bvec.sim', usedefault=True, desc='simulated b vectors')
    out_bval = File('bval.sim', usedefault=True, desc='simulated b values')
    snr = traits.Int(0, usedefault=True, desc='signal-to-noise ratio (dB)')


class SimulateMultiTensorOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc='simulated DWIs')
    out_mask = File(exists=True, desc='mask file')
    out_bvec = File(exists=True, desc='simulated b vectors')
    out_bval = File(exists=True, desc='simulated b values')


class SimulateMultiTensor(DipyBaseInterface):

    """
    Interface to MultiTensor model simulator in dipy
    http://nipy.org/dipy/examples_built/simulate_multi_tensor.html

    Example
    -------

    >>> import nipype.interfaces.dipy as dipy
    >>> sim = dipy.SimulateMultiTensor()
    >>> sim.inputs.in_dirs = ['fdir00.nii', 'fdir01.nii']
    >>> sim.inputs.in_frac = ['ffra00.nii', 'ffra01.nii']
    >>> sim.inputs.in_vfms = ['tpm_00.nii.gz', 'tpm_01.nii.gz',
    ...                       'tpm_02.nii.gz']
    >>> sim.inputs.baseline = 'b0.nii'
    >>> sim.inputs.in_bvec = 'bvecs'
    >>> sim.inputs.in_bval = 'bvals'
    >>> sim.run()                                   # doctest: +SKIP
    """
    input_spec = SimulateMultiTensorInputSpec
    output_spec = SimulateMultiTensorOutputSpec

    def _run_interface(self, runtime):
        from dipy.core.gradients import gradient_table

        # Gradient table
        if isdefined(self.inputs.in_bval) and isdefined(self.inputs.in_bvec):
            # Load the gradient strengths and directions
            bvals = np.loadtxt(self.inputs.in_bval)
            bvecs = np.loadtxt(self.inputs.in_bvec).T
            gtab = gradient_table(bvals, bvecs)
        else:
            gtab = _generate_gradients(self.inputs.num_dirs,
                                       self.inputs.bvalues)
        ndirs = len(gtab.bvals)
        np.savetxt(op.abspath(self.inputs.out_bvec), gtab.bvecs.T)
        np.savetxt(op.abspath(self.inputs.out_bval), gtab.bvals)

        # Load the baseline b0 signal
        b0_im = nb.load(self.inputs.baseline)
        hdr = b0_im.header
        shape = b0_im.shape
        aff = b0_im.affine

        # Check and load sticks and their volume fractions
        nsticks = len(self.inputs.in_dirs)
        if len(self.inputs.in_frac) != nsticks:
            raise RuntimeError(('Number of sticks and their volume fractions'
                                ' must match.'))

        # Volume fractions of isotropic compartments
        nballs = len(self.inputs.in_vfms)
        vfs = np.squeeze(nb.concat_images(
            [nb.load(f) for f in self.inputs.in_vfms]).get_data())
        if nballs == 1:
            vfs = vfs[..., np.newaxis]
        total_vf = np.sum(vfs, axis=3)

        # Generate a mask
        if isdefined(self.inputs.in_mask):
            msk = nb.load(self.inputs.in_mask).get_data()
            msk[msk > 0.0] = 1.0
            msk[msk < 1.0] = 0.0
        else:
            msk = np.zeros(shape)
            msk[total_vf > 0.0] = 1.0

        msk = np.clip(msk, 0.0, 1.0)
        nvox = len(msk[msk > 0])

        # Fiber fractions
        ffsim = nb.concat_images([nb.load(f) for f in self.inputs.in_frac])
        ffs = np.nan_to_num(np.squeeze(ffsim.get_data()))  # fiber fractions
        ffs = np.clip(ffs, 0., 1.)
        if nsticks == 1:
            ffs = ffs[..., np.newaxis]

        for i in range(nsticks):
            ffs[..., i] *= msk

        total_ff = np.sum(ffs, axis=3)

        # Fix incongruencies in fiber fractions
        for i in range(1, nsticks):
            if np.any(total_ff > 1.0):
                errors = np.zeros_like(total_ff)
                errors[total_ff > 1.0] = total_ff[total_ff > 1.0] - 1.0
                ffs[..., i] -= errors
                ffs[ffs < 0.0] = 0.0
            total_ff = np.sum(ffs, axis=3)

        for i in range(vfs.shape[-1]):
            vfs[..., i] -= total_ff
        vfs = np.clip(vfs, 0., 1.)

        fractions = np.concatenate((ffs, vfs), axis=3)

        nb.Nifti1Image(fractions, aff, None).to_filename('fractions.nii.gz')
        nb.Nifti1Image(np.sum(fractions, axis=3), aff, None).to_filename(
            'total_vf.nii.gz')

        mhdr = hdr.copy()
        mhdr.set_data_dtype(np.uint8)
        mhdr.set_xyzt_units('mm', 'sec')
        nb.Nifti1Image(msk, aff, mhdr).to_filename(
            op.abspath(self.inputs.out_mask))

        # Initialize stack of args
        fracs = fractions[msk > 0]

        # Stack directions
        dirs = None
        for i in range(nsticks):
            f = self.inputs.in_dirs[i]
            fd = np.nan_to_num(nb.load(f).get_data())
            w = np.linalg.norm(fd, axis=3)[..., np.newaxis]
            w[w < np.finfo(float).eps] = 1.0
            fd /= w
            if dirs is None:
                dirs = fd[msk > 0].copy()
            else:
                dirs = np.hstack((dirs, fd[msk > 0]))

        # Add random directions for isotropic components
        for d in range(nballs):
            fd = np.random.randn(nvox, 3)
            w = np.linalg.norm(fd, axis=1)
            fd[w < np.finfo(float).eps, ...] = np.array([1., 0., 0.])
            w[w < np.finfo(float).eps] = 1.0
            fd /= w[..., np.newaxis]
            dirs = np.hstack((dirs, fd))

        sf_evals = list(self.inputs.diff_sf)
        ba_evals = list(self.inputs.diff_iso)

        mevals = [sf_evals] * nsticks + \
            [[ba_evals[d]] * 3 for d in range(nballs)]

        b0 = b0_im.get_data()[msk > 0]
        args = []
        for i in range(nvox):
            args.append(
                {'fractions': fracs[i, ...].tolist(),
                 'sticks': [tuple(dirs[i, j:j + 3])
                            for j in range(nsticks + nballs)],
                 'gradients': gtab,
                 'mevals': mevals,
                 'S0': b0[i],
                 'snr': self.inputs.snr
                 })

        n_proc = self.inputs.n_proc
        if n_proc == 0:
            n_proc = cpu_count()

        try:
            pool = Pool(processes=n_proc, maxtasksperchild=50)
        except TypeError:
            pool = Pool(processes=n_proc)

        # Simulate sticks using dipy
        IFLOGGER.info(('Starting simulation of %d voxels, %d diffusion'
                       ' directions.') % (len(args), ndirs))
        result = np.array(pool.map(_compute_voxel, args))
        if np.shape(result)[1] != ndirs:
            raise RuntimeError(('Computed directions do not match number'
                                'of b-values.'))

        signal = np.zeros((shape[0], shape[1], shape[2], ndirs))
        signal[msk > 0] = result

        simhdr = hdr.copy()
        simhdr.set_data_dtype(np.float32)
        simhdr.set_xyzt_units('mm', 'sec')
        nb.Nifti1Image(signal.astype(np.float32), aff,
                       simhdr).to_filename(op.abspath(self.inputs.out_file))

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_file'] = op.abspath(self.inputs.out_file)
        outputs['out_mask'] = op.abspath(self.inputs.out_mask)
        outputs['out_bvec'] = op.abspath(self.inputs.out_bvec)
        outputs['out_bval'] = op.abspath(self.inputs.out_bval)

        return outputs


def _compute_voxel(args):
    """
    Simulate DW signal for one voxel. Uses the multi-tensor model and
    three isotropic compartments.

    Apparent diffusivity tensors are taken from [Alexander2002]_
    and [Pierpaoli1996]_.

    .. [Alexander2002] Alexander et al., Detection and modeling of non-Gaussian
      apparent diffusion coefficient profiles in human brain data, MRM
      48(2):331-340, 2002, doi: `10.1002/mrm.10209
      <http://dx.doi.org/10.1002/mrm.10209>`_.
    .. [Pierpaoli1996] Pierpaoli et al., Diffusion tensor MR imaging
      of the human brain, Radiology 201:637-648. 1996.
    """
    from dipy.sims.voxel import multi_tensor

    ffs = args['fractions']
    gtab = args['gradients']
    signal = np.zeros_like(gtab.bvals, dtype=np.float32)

    # Simulate dwi signal
    sf_vf = np.sum(ffs)
    if sf_vf > 0.0:
        ffs = ((np.array(ffs) / sf_vf) * 100)
        snr = args['snr'] if args['snr'] > 0 else None

        try:
            signal, _ = multi_tensor(
                gtab, args['mevals'], S0=args['S0'],
                angles=args['sticks'], fractions=ffs, snr=snr)
        except Exception as e:
            pass
            # IFLOGGER.warn('Exception simulating dwi signal: %s' % e)

    return signal.tolist()


def _generate_gradients(ndirs=64, values=[1000, 3000], nb0s=1):
    """
    Automatically generate a `gradient table
    <http://nipy.org/dipy/examples_built/gradients_spheres.html#example-gradients-spheres>`_

    """
    import numpy as np
    from dipy.core.sphere import (disperse_charges, Sphere, HemiSphere)
    from dipy.core.gradients import gradient_table

    theta = np.pi * np.random.rand(ndirs)
    phi = 2 * np.pi * np.random.rand(ndirs)
    hsph_initial = HemiSphere(theta=theta, phi=phi)
    hsph_updated, potential = disperse_charges(hsph_initial, 5000)

    values = np.atleast_1d(values).tolist()
    vertices = hsph_updated.vertices
    bvecs = vertices.copy()
    bvals = np.ones(vertices.shape[0]) * values[0]

    for v in values[1:]:
        bvecs = np.vstack((bvecs, vertices))
        bvals = np.hstack((bvals, v * np.ones(vertices.shape[0])))

    for i in range(0, nb0s):
        bvals = bvals.tolist()
        bvals.insert(0, 0)

        bvecs = bvecs.tolist()
        bvecs.insert(0, np.zeros(3))

    return gradient_table(bvals, bvecs)
