# -*- coding: utf-8 -*-
# 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)

"""
from __future__ import (print_function, division, unicode_literals,
                        absolute_import)
from builtins import open

import os.path as op

from ... import logging
from ..base import (BaseInterface, BaseInterfaceInputSpec, isdefined,
                    TraitedSpec, File, traits)
iflogger = logging.getLogger('nipype.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'  # doctest: +SKIP
    >>> tfm.inputs.reference_image = 'fixed1.nii'  # doctest: +SKIP
    >>> 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.header.get_zooms()) == 4:
                im = nb.func.four_to_three(im)[0]

            size = ' '.join(["%01d" % s for s in im.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.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.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.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 val is not None and 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
