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

__author__ = 'desmat'

class TensorMrtrix(GenericTask):


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


    def implement(self):

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


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


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

        tmp = self.buildName(source, "tmp")
        target = self.buildName(source, "tensor")
        cmd = "dwi2tensor {} {} -grad {} -nthreads {} -quiet ".format(source, tmp, 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):
        return self.get("ignore")


    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'

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

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

        for postfix, description in tags:
            image = self.getImage('dwi', postfix)
            if image:
                qaImage = self.buildName(image, softwareName, 'png')
                self.slicerPng(image, qaImage, boundaries=mask)
                qaImages.extend(Images((qaImage, description)))

        return qaImages
