# -*- coding: utf-8 -*-
from core.toad.generictask import GenericTask
from lib.images import Images

__author__ = "Mathieu Desrosiers, Arnaud Bore"
__copyright__ = "Copyright (C) 2016, TOAD"
__credits__ = ["Mathieu Desrosiers", "Arnaud Bore"]


class TensorMrtrix(GenericTask):


    def __init__(self, subject):
        GenericTask.__init__(self, subject, 'upsampling', 'registration', 'masking', 'qa')


    def implement(self):

        dwi = self.getUpsamplingImage('dwi', 'upsample')
        bFile = self.getUpsamplingImage('grad',  None, 'b')
        mask = self.getRegistrationImage('mask', 'resample')

        iterWLS = self.get('iter')  # Number of iteration for tensor estimations

        tensorsMrtrix = self.__produceTensors(dwi, bFile, iterWLS, mask)
        self.__produceMetrics(tensorsMrtrix, mask, dwi)

    # convert diffusion-weighted images to tensor images.
    def __produceTensors(self, source, encodingFile, iterWLS, mask=None):
        self.info("Starting DWI2Tensor from mrtrix using weighted linear least squares estimator.")

        tmp = self.buildName(source, "tmp")
        target = self.buildName(source, "tensor")
        cmd = "dwi2tensor {} {} -iter {} -grad {} -nthreads {} -quiet ".format(source, tmp, iterWLS, encodingFile, self.getNTreadsMrtrix())
        if mask is not None:
            cmd += "-mask {}".format(mask)

        self.launchCommand(cmd)
        return self.rename(tmp, target)


    def __produceMetrics(self, source, mask, target):
        self.info("Launch tensor2metric from mrtrix.\n")
        adc = self.buildName(target, "adc")
        fa = self.buildName(target, "fa")
        vector = self.buildName(target, "vector")
        adImage = self.buildName(target, "ad")
        rdImage = self.buildName(target, "rd")
        mdImage = self.buildName(target, "md")
        value2 = self.buildName(target, "l2")
        value3 = self.buildName(target, "l3")
        modulate = self.get('modulate')

        cmd1 = "tensor2metric {} -adc {} -fa {} -num 1 -vector {} -value {} -modulate {} -nthreads {} -quiet "\
            .format(source, adc, fa, vector, adImage , modulate, self.getNTreadsMrtrix())
        cmd2 = "tensor2metric {} -num 2 -value {} -modulate {} -nthreads {} -quiet "\
            .format(source, value2, modulate, self.getNTreadsMrtrix())
        cmd3 = "tensor2metric {} -num 3 -value {} -modulate {} -nthreads {} -quiet "\
            .format(source, value3, modulate, self.getNTreadsMrtrix())

        for cmd in [cmd1, cmd2, cmd3]:
            if mask is not None:
                cmd += "-mask {} ".format(mask)
            self.launchCommand(cmd)

        cmd = "mrmath {} {} mean {} -nthreads {} -quiet ".format(value2, value3, rdImage, self.getNTreadsMrtrix())
        self.launchCommand(cmd)

        cmd = "mrmath {} {} {} mean {} -nthreads {} -quiet ".format(adImage, value2, value3, mdImage, self.getNTreadsMrtrix())
        self.launchCommand(cmd)


    def isIgnore(self):
        



    def meetRequirement(self):
        return Images((self.getUpsamplingImage('dwi', 'upsample'), "upsampled diffusion"),
                  (self.getUpsamplingImage('grad', None, 'b'), "gradient encoding b file"),
                  (self.getRegistrationImage('mask', 'resample'), 'brain  mask'))



    def isDirty(self):

        return Images((self.getImage("dwi", "tensor"), "mrtrix tensor"),
                     (self.getImage('dwi', 'adc'), "mean apparent diffusion coefficient (ADC)"),
                     (self.getImage('dwi', 'vector'), "selected eigenvector(s)"),
                     (self.getImage('dwi', 'fa'), "fractional anisotropy"),
                     (self.getImage('dwi', 'ad'), "selected eigenvalue(s) AD" ),
                     (self.getImage('dwi', 'rd'), "selected eigenvalue(s) RD"),
                     (self.getImage('dwi', 'md'), "mean diffusivity"))

    def qaSupplier(self):
        """Create and supply images for the report generated by qa task

        """
        qaImages = Images()
        softwareName = 'mrtrix'

        #Set information
        information = "Estimation using WLS with {} iteration(s)".format(self.get('iter'))
        qaImages.setInformation(information)

        #Get images
        mask = self.getRegistrationImage('mask', 'resample')

        #Build qa images
        tags = (
            ('fa', 0.7, 'Fractional anisotropy'),
            ('ad', 0.005, 'Axial Diffusivity'),
            ('md', 0.005, 'Mean Diffusivity'),
            ('rd', 0.005, 'Radial Diffusivity'),
            )

        for postfix, vmax, description in tags:
            image = self.getImage('dwi', postfix)
            if image:
                imageQa = self.plot3dVolume(
                        image, fov=mask, vmax=vmax,
                        colorbar=True, postfix=softwareName)
                qaImages.append((imageQa, description))

        return qaImages

return self.get("ignore")