# coding: utf-8
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:

"""
Generic interfaces to manipulate registration parameters files, including
transform files (to configure warpings)

"""

import os.path as op
import re

from nipype.interfaces.base import (BaseInterface, BaseInterfaceInputSpec, isdefined,
                                    TraitedSpec, File, traits, InputMultiPath)


from nipype import logging
logger = logging.getLogger('interface')


class EditTransformInputSpec(BaseInterfaceInputSpec):
    transform_file = File(exists=True, mandatory=True,
                          desc='transform-parameter file, only 1')
    reference_image = File(exists=True,
                           desc=('set a new reference image to change the '
                                 'target coordinate system.'))
    interpolation = traits.Enum('cubic','linear','nearest', usedefault=True,
                                argstr='FinalBSplineInterpolationOrder',
                                desc='set a new interpolator for transformation')

    output_type = traits.Enum('float', 'unsigned char', 'unsigned short','short',
                              'unsigned long','long','double',
                              argstr='ResultImagePixelType',
                              desc='set a new output pixel type for resampled images')
    output_format = traits.Enum('nii.gz', 'nii', 'mhd', 'hdr', 'vtk',
                                argstr='ResultImageFormat',
                                desc='set a new image format for resampled images')
    output_file = File(desc='the filename for the resulting transform file')

class EditTransformOutputSpec(TraitedSpec):
    output_file = File(exists=True, desc='output transform file')


class EditTransform(BaseInterface):
    """
    Manipulates an existing transform file generated with elastix

    Example
    -------

    >>> from nipype.interfaces.elastix import EditTransform
    >>> tfm = EditTransform()
    >>> tfm.inputs.transform_file = 'TransformParameters.0.txt'
    >>> tfm.inputs.reference_image = 'fixed1.nii'
    >>> tfm.inputs.output_type = 'unsigned char'
    >>> tfm.run() # doctest: +SKIP


    """

    input_spec = EditTransformInputSpec
    output_spec = EditTransformOutputSpec
    _out_file = ''
    _pattern = '\((?P<entry>%s\s\"?)([-\.\s\w]+)(\"?\))'

    _interp = { 'nearest': 0, 'linear': 1, 'cubic': 3 }

    def _run_interface(self, runtime):
        import re
        import nibabel as nb
        import numpy as np

        contents = ''

        with open(self.inputs.transform_file, 'r') as f:
            contents = f.read()

        if isdefined(self.inputs.output_type):
            p = re.compile((self._pattern % 'ResultImagePixelType').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % self.inputs.output_type
            contents = p.sub(rep, contents)

        if isdefined(self.inputs.output_format):
            p = re.compile((self._pattern % 'ResultImageFormat').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % self.inputs.output_format
            contents = p.sub(rep, contents)

        if isdefined(self.inputs.interpolation):
            p = re.compile((self._pattern % 'FinalBSplineInterpolationOrder').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % self._interp[self.inputs.interpolation]
            contents = p.sub(rep, contents)

        if isdefined(self.inputs.reference_image):
            im = nb.load(self.inputs.reference_image)

            if len(im.get_header().get_zooms()) == 4:
                im = nb.func.four_to_three(im)[0]

            size = ' '.join(["%01d" % s for s in im.get_shape() ])
            p = re.compile((self._pattern % 'Size').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % size
            contents = p.sub(rep, contents)

            index = ' '.join(["0" for s in im.get_shape() ])
            p = re.compile((self._pattern % 'Index').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % index
            contents = p.sub(rep, contents)

            spacing = ' '.join(["%0.4f" % f for f in im.get_header().get_zooms()])
            p = re.compile((self._pattern % 'Spacing').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % spacing
            contents = p.sub(rep, contents)

            itkmat = np.eye(4)
            itkmat[0,0] = -1
            itkmat[1,1] = -1

            affine = np.dot( itkmat, im.get_affine() )
            dirs = ' '.join(['%0.4f' % f for f in affine[0:3,0:3].reshape(-1)])
            orig = ' '.join(['%0.4f' % f for f in affine[0:3,3].reshape(-1)])

            #p = re.compile((self._pattern % 'Direction').decode('string-escape'))
            #rep = '(\g<entry>%s\g<3>' % dirs
            #contents = p.sub(rep, contents)

            p = re.compile((self._pattern % 'Origin').decode('string-escape'))
            rep = '(\g<entry>%s\g<3>' % orig
            contents = p.sub(rep, contents)


        with open(self._get_outfile(), 'w') as of:
            of.write(contents)

        return runtime

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs['output_file'] = getattr(self, '_out_file')
        return outputs

    def _get_outfile(self):
        val = getattr(self,'_out_file')
        if not val is None and not val=='':
            return val

        if isdefined(self.inputs.output_file):
            setattr(self,'_out_file',self.inputs.output_file)
            return self.inputs.output_file

        out_file = op.abspath(op.basename(self.inputs.transform_file))
        setattr(self, '_out_file', out_file)
        return out_file
