# -*- coding: utf-8 -*-
"""Provides interfaces to various commands provided by dcmstack
"""
from __future__ import (print_function, division, unicode_literals,
                        absolute_import)

import os
from os import path as op
import string
import errno
from glob import glob

import nibabel as nb
import imghdr

from .base import (TraitedSpec, DynamicTraitedSpec, InputMultiPath, File,
                   Directory, traits, BaseInterface, isdefined, Undefined)
from ..utils import NUMPY_MMAP

have_dcmstack = True
try:
    import dicom
    import dcmstack
    from dcmstack.dcmmeta import NiftiWrapper
except ImportError:
    have_dcmstack = False


def sanitize_path_comp(path_comp):
    result = []
    for char in path_comp:
        if char not in string.letters + string.digits + '-_.':
            result.append('_')
        else:
            result.append(char)
    return ''.join(result)


class NiftiGeneratorBaseInputSpec(TraitedSpec):
    out_format = traits.Str(desc="String which can be formatted with "
                            "meta data to create the output filename(s)")
    out_ext = traits.Str(
        '.nii.gz', usedefault=True, desc="Determines output file type")
    out_path = Directory(
        desc='output path, current working directory if not set')


class NiftiGeneratorBase(BaseInterface):
    '''Base class for interfaces that produce Nifti files, potentially with
    embedded meta data.'''

    def _get_out_path(self, meta, idx=None):
        '''Return the output path for the gernerated Nifti.'''
        if self.inputs.out_format:
            out_fmt = self.inputs.out_format
        else:
            # If no out_format is specified, use a sane default that will work
            # with the provided meta data.
            out_fmt = []
            if idx is not None:
                out_fmt.append('%03d' % idx)
            if 'SeriesNumber' in meta:
                out_fmt.append('%(SeriesNumber)03d')
            if 'ProtocolName' in meta:
                out_fmt.append('%(ProtocolName)s')
            elif 'SeriesDescription' in meta:
                out_fmt.append('%(SeriesDescription)s')
            else:
                out_fmt.append('sequence')
            out_fmt = '-'.join(out_fmt)
        out_fn = (out_fmt % meta) + self.inputs.out_ext
        out_fn = sanitize_path_comp(out_fn)

        out_path = os.getcwd()
        if isdefined(self.inputs.out_path):
            out_path = op.abspath(self.inputs.out_path)

            # now, mkdir -p $out_path
            try:
                os.makedirs(out_path)
            except OSError as exc:  # Python >2.5
                if exc.errno == errno.EEXIST and op.isdir(out_path):
                    pass
                else:
                    raise

        return op.join(out_path, out_fn)


class DcmStackInputSpec(NiftiGeneratorBaseInputSpec):
    dicom_files = traits.Either(
        InputMultiPath(File(exists=True)),
        Directory(exists=True),
        traits.Str(),
        mandatory=True)
    embed_meta = traits.Bool(desc="Embed DICOM meta data into result")
    exclude_regexes = traits.List(desc="Meta data to exclude, suplementing "
                                  "any default exclude filters")
    include_regexes = traits.List(desc="Meta data to include, overriding any "
                                  "exclude filters")
    force_read = traits.Bool(
        True,
        usedefault=True,
        desc=('Force reading files without DICM marker'))


class DcmStackOutputSpec(TraitedSpec):
    out_file = File(exists=True)


class DcmStack(NiftiGeneratorBase):
    '''Create one Nifti file from a set of DICOM files. Can optionally embed
    meta data.

    Example
    -------

    >>> from nipype.interfaces.dcmstack import DcmStack
    >>> stacker = DcmStack()
    >>> stacker.inputs.dicom_files = 'path/to/series/'
    >>> stacker.run() # doctest: +SKIP
    >>> result.outputs.out_file # doctest: +SKIP
    '/path/to/cwd/sequence.nii.gz'
    '''
    input_spec = DcmStackInputSpec
    output_spec = DcmStackOutputSpec

    def _get_filelist(self, trait_input):
        if isinstance(trait_input, (str, bytes)):
            if op.isdir(trait_input):
                return glob(op.join(trait_input, '*.dcm'))
            else:
                return glob(trait_input)

        return trait_input

    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        include_regexes = dcmstack.default_key_incl_res
        if isdefined(self.inputs.include_regexes):
            include_regexes += self.inputs.include_regexes
        exclude_regexes = dcmstack.default_key_excl_res
        if isdefined(self.inputs.exclude_regexes):
            exclude_regexes += self.inputs.exclude_regexes
        meta_filter = dcmstack.make_key_regex_filter(exclude_regexes,
                                                     include_regexes)
        stack = dcmstack.DicomStack(meta_filter=meta_filter)
        for src_path in src_paths:
            if not imghdr.what(src_path) == "gif":
                src_dcm = dicom.read_file(
                    src_path, force=self.inputs.force_read)
                stack.add_dcm(src_dcm)
        nii = stack.to_nifti(embed_meta=True)
        nw = NiftiWrapper(nii)
        self.out_path = \
            self._get_out_path(nw.meta_ext.get_class_dict(('global', 'const')))
        if not self.inputs.embed_meta:
            nw.remove_extension()
        nb.save(nii, self.out_path)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_file"] = self.out_path
        return outputs


class GroupAndStackOutputSpec(TraitedSpec):
    out_list = traits.List(desc="List of output nifti files")


class GroupAndStack(DcmStack):
    '''Create (potentially) multiple Nifti files for a set of DICOM files.
    '''
    input_spec = DcmStackInputSpec
    output_spec = GroupAndStackOutputSpec

    def _run_interface(self, runtime):
        src_paths = self._get_filelist(self.inputs.dicom_files)
        stacks = dcmstack.parse_and_stack(src_paths)

        self.out_list = []
        for key, stack in list(stacks.items()):
            nw = NiftiWrapper(stack.to_nifti(embed_meta=True))
            const_meta = nw.meta_ext.get_class_dict(('global', 'const'))
            out_path = self._get_out_path(const_meta)
            if not self.inputs.embed_meta:
                nw.remove_extension()
            nb.save(nw.nii_img, out_path)
            self.out_list.append(out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs["out_list"] = self.out_list
        return outputs


class LookupMetaInputSpec(TraitedSpec):
    in_file = File(mandatory=True, exists=True, desc='The input Nifti file')
    meta_keys = traits.Either(
        traits.List(),
        traits.Dict(),
        mandatory=True,
        desc=("List of meta data keys to lookup, or a "
              "dict where keys specify the meta data "
              "keys to lookup and the values specify "
              "the output names"))


class LookupMeta(BaseInterface):
    '''Lookup meta data values from a Nifti with embedded meta data.

    Example
    -------

    >>> from nipype.interfaces import dcmstack
    >>> lookup = dcmstack.LookupMeta()
    >>> lookup.inputs.in_file = 'functional.nii'
    >>> lookup.inputs.meta_keys = {'RepetitionTime' : 'TR', \
                                   'EchoTime' : 'TE'}
    >>> result = lookup.run() # doctest: +SKIP
    >>> result.outputs.TR # doctest: +SKIP
    9500.0
    >>> result.outputs.TE # doctest: +SKIP
    95.0
    '''
    input_spec = LookupMetaInputSpec
    output_spec = DynamicTraitedSpec

    def _make_name_map(self):
        if isinstance(self.inputs.meta_keys, list):
            self._meta_keys = {}
            for key in self.inputs.meta_keys:
                self._meta_keys[key] = key
        else:
            self._meta_keys = self.inputs.meta_keys

    def _outputs(self):
        self._make_name_map()
        outputs = super(LookupMeta, self)._outputs()
        undefined_traits = {}
        for out_name in list(self._meta_keys.values()):
            outputs.add_trait(out_name, traits.Any)
            undefined_traits[out_name] = Undefined
        outputs.trait_set(trait_change_notify=False, **undefined_traits)
        # Not sure why this is needed
        for out_name in list(self._meta_keys.values()):
            _ = getattr(outputs, out_name)
        return outputs

    def _run_interface(self, runtime):
        # If the 'meta_keys' input is a list, covert it to a dict
        self._make_name_map()
        nw = NiftiWrapper.from_filename(self.inputs.in_file)
        self.result = {}
        for meta_key, out_name in list(self._meta_keys.items()):
            self.result[out_name] = nw.meta_ext.get_values(meta_key)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs.update(self.result)
        return outputs


class CopyMetaInputSpec(TraitedSpec):
    src_file = File(mandatory=True, exists=True)
    dest_file = File(mandatory=True, exists=True)
    include_classes = traits.List(desc="List of specific meta data "
                                  "classifications to include. If not "
                                  "specified include everything.")
    exclude_classes = traits.List(desc="List of meta data "
                                  "classifications to exclude")


class CopyMetaOutputSpec(TraitedSpec):
    dest_file = File(exists=True)


class CopyMeta(BaseInterface):
    '''Copy meta data from one Nifti file to another. Useful for preserving
    meta data after some processing steps.'''
    input_spec = CopyMetaInputSpec
    output_spec = CopyMetaOutputSpec

    def _run_interface(self, runtime):
        src_nii = nb.load(self.inputs.src_file)
        src = NiftiWrapper(src_nii, make_empty=True)
        dest_nii = nb.load(self.inputs.dest_file)
        dest = NiftiWrapper(dest_nii, make_empty=True)
        classes = src.meta_ext.get_valid_classes()
        if self.inputs.include_classes:
            classes = [
                cls for cls in classes if cls in self.inputs.include_classes
            ]
        if self.inputs.exclude_classes:
            classes = [
                cls for cls in classes
                if cls not in self.inputs.exclude_classes
            ]

        for cls in classes:
            src_dict = src.meta_ext.get_class_dict(cls)
            dest_dict = dest.meta_ext.get_class_dict(cls)
            dest_dict.update(src_dict)
        # Update the shape and slice dimension to reflect the meta extension
        # update.
        dest.meta_ext.slice_dim = src.meta_ext.slice_dim
        dest.meta_ext.shape = src.meta_ext.shape

        self.out_path = op.join(os.getcwd(), op.basename(
            self.inputs.dest_file))
        dest.to_filename(self.out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['dest_file'] = self.out_path
        return outputs


class MergeNiftiInputSpec(NiftiGeneratorBaseInputSpec):
    in_files = traits.List(mandatory=True, desc="List of Nifti files to merge")
    sort_order = traits.Either(
        traits.Str(),
        traits.List(),
        desc="One or more meta data keys to "
        "sort files by.")
    merge_dim = traits.Int(desc="Dimension to merge along. If not "
                           "specified, the last singular or "
                           "non-existant dimension is used.")


class MergeNiftiOutputSpec(TraitedSpec):
    out_file = File(exists=True, desc="Merged Nifti file")


def make_key_func(meta_keys, index=None):
    def key_func(src_nii):
        result = [src_nii.get_meta(key, index) for key in meta_keys]
        return result

    return key_func


class MergeNifti(NiftiGeneratorBase):
    '''Merge multiple Nifti files into one. Merges together meta data
    extensions as well.'''
    input_spec = MergeNiftiInputSpec
    output_spec = MergeNiftiOutputSpec

    def _run_interface(self, runtime):
        niis = [nb.load(fn, mmap=NUMPY_MMAP) for fn in self.inputs.in_files]
        nws = [NiftiWrapper(nii, make_empty=True) for nii in niis]
        if self.inputs.sort_order:
            sort_order = self.inputs.sort_order
            if isinstance(sort_order, (str, bytes)):
                sort_order = [sort_order]
            nws.sort(key=make_key_func(sort_order))
        if self.inputs.merge_dim == traits.Undefined:
            merge_dim = None
        else:
            merge_dim = self.inputs.merge_dim
        merged = NiftiWrapper.from_sequence(nws, merge_dim)
        const_meta = merged.meta_ext.get_class_dict(('global', 'const'))
        self.out_path = self._get_out_path(const_meta)
        nb.save(merged.nii_img, self.out_path)
        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_file'] = self.out_path
        return outputs


class SplitNiftiInputSpec(NiftiGeneratorBaseInputSpec):
    in_file = File(exists=True, mandatory=True, desc="Nifti file to split")
    split_dim = traits.Int(desc="Dimension to split along. If not "
                           "specified, the last dimension is used.")


class SplitNiftiOutputSpec(TraitedSpec):
    out_list = traits.List(File(exists=True), desc="Split Nifti files")


class SplitNifti(NiftiGeneratorBase):
    '''
    Split one Nifti file into many along the specified dimension. Each
    result has an updated meta data extension as well.
    '''
    input_spec = SplitNiftiInputSpec
    output_spec = SplitNiftiOutputSpec

    def _run_interface(self, runtime):
        self.out_list = []
        nii = nb.load(self.inputs.in_file)
        nw = NiftiWrapper(nii, make_empty=True)
        split_dim = None
        if self.inputs.split_dim == traits.Undefined:
            split_dim = None
        else:
            split_dim = self.inputs.split_dim
        for split_idx, split_nw in enumerate(nw.split(split_dim)):
            const_meta = split_nw.meta_ext.get_class_dict(('global', 'const'))
            out_path = self._get_out_path(const_meta, idx=split_idx)
            nb.save(split_nw.nii_img, out_path)
            self.out_list.append(out_path)

        return runtime

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['out_list'] = self.out_list
        return outputs
