# -*- coding: utf-8 -*-
"""The ants visualisation module provides basic functions based on ITK.
"""
from __future__ import (print_function, division, unicode_literals,
                        absolute_import)

import os

from ..base import TraitedSpec, File, traits
from .base import ANTSCommand, ANTSCommandInputSpec


class ConvertScalarImageToRGBInputSpec(ANTSCommandInputSpec):
    dimension = traits.Enum(
        3,
        2,
        argstr='%d',
        usedefault=True,
        desc='image dimension (2 or 3)',
        mandatory=True,
        position=0)
    input_image = File(
        argstr='%s',
        exists=True,
        desc='Main input is a 3-D grayscale image.',
        mandatory=True,
        position=1)
    output_image = traits.Str(
        'rgb.nii.gz',
        argstr='%s',
        usedefault=True,
        desc='rgb output image',
        position=2)
    mask_image = File(
        'none',
        argstr='%s',
        exists=True,
        desc='mask image',
        position=3,
        usedefault=True)
    colormap = traits.Str(
        argstr='%s',
        usedefault=True,
        desc=('Possible colormaps: grey, red, green, '
              'blue, copper, jet, hsv, spring, summer, '
              'autumn, winter, hot, cool, overunder, custom '),
        mandatory=True,
        position=4)
    custom_color_map_file = traits.Str(
        'none',
        argstr='%s',
        usedefault=True,
        desc='custom color map file',
        position=5)
    minimum_input = traits.Int(
        argstr='%d', desc='minimum input', mandatory=True, position=6)
    maximum_input = traits.Int(
        argstr='%d', desc='maximum input', mandatory=True, position=7)
    minimum_RGB_output = traits.Int(
        0, usedefault=True, argstr='%d', desc='', position=8)
    maximum_RGB_output = traits.Int(
        255, usedefault=True, argstr='%d', desc='', position=9)


class ConvertScalarImageToRGBOutputSpec(TraitedSpec):
    output_image = File(exists=True, desc='converted RGB image')


class ConvertScalarImageToRGB(ANTSCommand):
    """
    Examples
    --------
    >>> from nipype.interfaces.ants.visualization import ConvertScalarImageToRGB
    >>> converter = ConvertScalarImageToRGB()
    >>> converter.inputs.dimension = 3
    >>> converter.inputs.input_image = 'T1.nii.gz'
    >>> converter.inputs.colormap = 'jet'
    >>> converter.inputs.minimum_input = 0
    >>> converter.inputs.maximum_input = 6
    >>> converter.cmdline
    'ConvertScalarImageToRGB 3 T1.nii.gz rgb.nii.gz none jet none 0 6 0 255'
    """
    _cmd = 'ConvertScalarImageToRGB'
    input_spec = ConvertScalarImageToRGBInputSpec
    output_spec = ConvertScalarImageToRGBOutputSpec

    def _format_arg(self, opt, spec, val):
        return super(ConvertScalarImageToRGB, self)._format_arg(opt, spec, val)

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['output_image'] = os.path.join(os.getcwd(),
                                               self.inputs.output_image)
        return outputs


class CreateTiledMosaicInputSpec(ANTSCommandInputSpec):
    input_image = File(
        argstr='-i %s',
        exists=True,
        desc='Main input is a 3-D grayscale image.',
        mandatory=True)
    rgb_image = File(
        argstr='-r %s',
        exists=True,
        desc=('An optional Rgb image can be added as an overlay.'
              'It must have the same image'
              'geometry as the input grayscale image.'),
        mandatory=True)
    mask_image = File(
        argstr='-x %s',
        exists=True,
        desc='Specifies the ROI of the RGB voxels used.')
    alpha_value = traits.Float(
        argstr='-a %.2f',
        desc=('If an Rgb image is provided, render the overlay '
              'using the specified alpha parameter.'))
    output_image = traits.Str(
        'output.png',
        argstr='-o %s',
        desc='The output consists of the tiled mosaic image.',
        usedefault=True)
    tile_geometry = traits.Str(
        argstr='-t %s',
        desc=(
            'The tile geometry specifies the number of rows and columns'
            'in the output image. For example, if the user specifies "5x10", '
            'then 5 rows by 10 columns of slices are rendered. If R < 0 and C > '
            '0 (or vice versa), the negative value is selected'
            'based on direction.'))
    direction = traits.Int(
        argstr='-d %d',
        desc=('Specifies the direction of '
              'the slices. If no direction is specified, the '
              'direction with the coarsest spacing is chosen.'))
    pad_or_crop = traits.Str(
        argstr='-p %s',
        desc='argument passed to -p flag:'
        '[padVoxelWidth,<constantValue=0>]'
        '[lowerPadding[0]xlowerPadding[1],upperPadding[0]xupperPadding[1],'
        'constantValue]'
        'The user can specify whether to pad or crop a specified '
        'voxel-width boundary of each individual slice. For this '
        'program, cropping is simply padding with negative voxel-widths.'
        'If one pads (+), the user can also specify a constant pad '
        'value (default = 0). If a mask is specified, the user can use '
        'the mask to define the region, by using the keyword "mask"'
        ' plus an offset, e.g. "-p mask+3".')
    slices = traits.Str(
        argstr='-s %s',
        desc=('Number of slices to increment Slice1xSlice2xSlice3'
              '[numberOfSlicesToIncrement,<minSlice=0>,<maxSlice=lastSlice>]'))
    flip_slice = traits.Str(argstr='-f %s', desc='flipXxflipY')
    permute_axes = traits.Bool(argstr='-g', desc='doPermute')


class CreateTiledMosaicOutputSpec(TraitedSpec):
    output_image = File(exists=True, desc='image file')


class CreateTiledMosaic(ANTSCommand):
    """The program CreateTiledMosaic in conjunction with ConvertScalarImageToRGB
    provides useful functionality for common image analysis tasks. The basic
    usage of CreateTiledMosaic is to tile a 3-D image volume slice-wise into
    a 2-D image.

    Examples
    --------

    >>> from nipype.interfaces.ants.visualization import CreateTiledMosaic
    >>> mosaic_slicer = CreateTiledMosaic()
    >>> mosaic_slicer.inputs.input_image = 'T1.nii.gz'
    >>> mosaic_slicer.inputs.rgb_image = 'rgb.nii.gz'
    >>> mosaic_slicer.inputs.mask_image = 'mask.nii.gz'
    >>> mosaic_slicer.inputs.output_image = 'output.png'
    >>> mosaic_slicer.inputs.alpha_value = 0.5
    >>> mosaic_slicer.inputs.direction = 2
    >>> mosaic_slicer.inputs.pad_or_crop = '[ -15x -50 , -15x -30 ,0]'
    >>> mosaic_slicer.inputs.slices = '[2 ,100 ,160]'
    >>> mosaic_slicer.cmdline
    'CreateTiledMosaic -a 0.50 -d 2 -i T1.nii.gz -x mask.nii.gz -o output.png -p [ -15x -50 , -15x -30 ,0] \
-r rgb.nii.gz -s [2 ,100 ,160]'
    """

    _cmd = 'CreateTiledMosaic'
    input_spec = CreateTiledMosaicInputSpec
    output_spec = CreateTiledMosaicOutputSpec

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['output_image'] = os.path.join(os.getcwd(),
                                               self.inputs.output_image)
        return outputs
