from abc import ABCMeta, abstractmethod
from nipype.interfaces.base import (
    traits, DynamicTraitedSpec, Undefined, File, Directory,
    BaseInterface)
from arcana.node import Node
from arcana.dataset import (
    Dataset, DatasetSpec, FieldSpec, BaseField, BaseDataset)
from arcana.exception import ArcanaError
from arcana.utils import PATH_SUFFIX, FIELD_SUFFIX

PATH_TRAIT = traits.Either(File(exists=True), Directory(exists=True))
FIELD_TRAIT = traits.Either(traits.Int, traits.Float, traits.Str)
MULTIPLICITIES = ('per_session', 'per_subject', 'per_visit', 'per_project')


class Archive(object):
    """
    Abstract base class for all Archive systems, DaRIS, XNAT and local file
    system. Sets out the interface that all Archive classes should implement.
    """

    __metaclass__ = ABCMeta

    @abstractmethod
    def source(self, inputs, name=None, study_name=None, **kwargs):
        """
        Returns a NiPype node that gets the input data from the archive
        system. The input spec of the node's interface should inherit from
        ArchiveSourceInputSpec

        Parameters
        ----------
        project_id : str
            The ID of the project to return the sessions for
        inputs : list(Dataset|Field)
            An iterable of arcana.Dataset or arcana.Field
            objects, which specify the datasets to extract from the
            archive system
        name : str
            Name of the NiPype node
        study_name: str
            Prefix used to distinguish datasets generated by a particular
            study. Used for derived datasets only
        """
        if name is None:
            name = "{}_source".format(self.type)
        inputs = list(inputs)  # protected against iterators
        datasets = [i for i in inputs if isinstance(i, BaseDataset)]
        fields = [i for i in inputs if isinstance(i, BaseField)]
        return Node(self.Source(study_name, datasets, fields, **kwargs),
                    name=name)

    @abstractmethod
    def sink(self, outputs, frequency='per_session', name=None,
             study_name=None, **kwargs):
        """
        Returns a NiPype node that puts the output data back to the archive
        system. The input spec of the node's interface should inherit from
        ArchiveSinkInputSpec

        Parameters
        ----------
        project_id : str
            The ID of the project to return the sessions for
        outputs : List(BaseFile|Field) | list(
            An iterable of arcana.Dataset arcana.Field objects,
            which specify the datasets to put into the archive system
        name : str
            Name of the NiPype node
        study_name: str
            Prefix used to distinguish datasets generated by a particular
            study. Used for derived datasets only

        """
        if name is None:
            name = "{}_{}_sink".format(self.type, frequency)
        outputs = list(outputs)  # protected against iterators
        if frequency.startswith('per_session'):
            sink_class = self.Sink
        elif frequency.startswith('per_subject'):
            sink_class = self.SubjectSink
        elif frequency.startswith('per_visit'):
            sink_class = self.VisitSink
        elif frequency.startswith('per_project'):
            sink_class = self.ProjectSink
        else:
            raise ArcanaError(
                "Unrecognised frequency '{}' can be one of '{}'"
                .format(frequency,
                        "', '".join(Dataset.MULTIPLICITY_OPTIONS)))
        datasets = [o for o in outputs if isinstance(o, BaseDataset)]
        fields = [o for o in outputs if isinstance(o, BaseField)]
        return Node(sink_class(study_name, datasets, fields, **kwargs),
                    name=name)

    def __ne__(self, other):
        return not (self == other)


class BaseArchiveNode(BaseInterface):
    """
    Parameters
    ----------
    infields : list of str
        Indicates the input fields to be dynamically created

    outfields: list of str
        Indicates output fields to be dynamically created

    See class examples for usage

    """

    def __init__(self, study_name, datasets, fields):
        super(BaseArchiveNode, self).__init__()
        self._study_name = study_name
        self._datasets = datasets
        self._fields = fields

    def __eq__(self, other):
        try:
            return (self.study_name == other.study_name and
                    self.datasets == other.datasets and
                    self.fields == other.fields)
        except AttributeError:
            return False

    def __repr__(self):
        return "{}(study_name='{}', datasets={}, fields={})".format(
            type(self).__name__, self.study_name, self.datasets,
            self.fields)

    def __ne__(self, other):
        return not self == other

    def _run_interface(self, runtime, *args, **kwargs):  # @UnusedVariable
        return runtime

    @property
    def study_name(self):
        return self._study_name

    @property
    def datasets(self):
        return self._datasets

    @property
    def fields(self):
        return self._fields

    @classmethod
    def _add_trait(cls, spec, name, trait_type):
        spec.add_trait(name, trait_type)
        spec.trait_set(trait_change_notify=False, **{name: Undefined})
        # Access the trait (not sure why but this is done in add_traits
        # so I have also done it here
        getattr(spec, name)

    def prefix_study_name(self, name, is_spec=True):
        """Prepend study name if defined"""
        if is_spec:
            name = self.study_name + '_' + name
        return name


class ArchiveSourceInputSpec(DynamicTraitedSpec):
    """
    Base class for archive source input specifications. Provides a common
    interface for 'run_pipeline' when using the archive source to extract
    primary and preprocessed datasets from the archive system
    """
    subject_id = traits.Str(mandatory=True, desc="The subject ID")
    visit_id = traits.Str(mandatory=True, usedefult=True,
                            desc="The visit or derived group ID")


class ArchiveSource(BaseArchiveNode):
    """
    Parameters
    ----------
    datasets: list
        List of all datasets to be extracted from the archive
    fields: list
        List of all the fields that are to be extracted from the archive
    study_name: str
        Prefix prepended onto derived dataset "names"
    """

    output_spec = DynamicTraitedSpec
    _always_run = True

    def _outputs(self):
        outputs = super(ArchiveSource, self)._outputs()
        # Add output datasets
        for dataset in self.datasets:
            assert isinstance(dataset, BaseDataset)
            self._add_trait(outputs, dataset.name + PATH_SUFFIX,
                            PATH_TRAIT)
        # Add output fields
        for field in self.fields:
            assert isinstance(field, BaseField)
            self._add_trait(outputs, field.name + FIELD_SUFFIX,
                            field.dtype)
        return outputs


class BaseArchiveSinkSpec(DynamicTraitedSpec):
    pass


class ArchiveSinkInputSpec(BaseArchiveSinkSpec):

    subject_id = traits.Str(mandatory=True, desc="The subject ID"),
    visit_id = traits.Str(mandatory=False,
                            desc="The session or derived group ID")


class ArchiveSubjectSinkInputSpec(BaseArchiveSinkSpec):

    subject_id = traits.Str(mandatory=True, desc="The subject ID")


class ArchiveVisitSinkInputSpec(BaseArchiveSinkSpec):

    visit_id = traits.Str(mandatory=True, desc="The visit ID")


class ArchiveProjectSinkInputSpec(BaseArchiveSinkSpec):
    pass


class BaseArchiveSinkOutputSpec(DynamicTraitedSpec):

    out_files = traits.List(PATH_TRAIT, desc='Output datasets')

    out_fields = traits.List(
        traits.Tuple(traits.Str, FIELD_TRAIT), desc='Output fields')


class ArchiveSinkOutputSpec(BaseArchiveSinkOutputSpec):

    subject_id = traits.Str(desc="The subject ID")
    visit_id = traits.Str(desc="The visit ID")


class ArchiveSubjectSinkOutputSpec(BaseArchiveSinkOutputSpec):

    subject_id = traits.Str(desc="The subject ID")


class ArchiveVisitSinkOutputSpec(BaseArchiveSinkOutputSpec):

    visit_id = traits.Str(desc="The visit ID")


class ArchiveProjectSinkOutputSpec(BaseArchiveSinkOutputSpec):

    project_id = traits.Str(desc="The project ID")


class BaseArchiveSink(BaseArchiveNode):

    def __init__(self, study_name, datasets, fields):
        super(BaseArchiveSink, self).__init__(study_name, datasets,
                                              fields)
        # Add input datasets
        for dataset in datasets:
            assert isinstance(dataset, DatasetSpec)
            self._add_trait(self.inputs, dataset.name + PATH_SUFFIX,
                            PATH_TRAIT)
        # Add input fields
        for field in fields:
            assert isinstance(field, FieldSpec)
            self._add_trait(self.inputs, field.name + FIELD_SUFFIX,
                            field.dtype)


class ArchiveSink(BaseArchiveSink):

    input_spec = ArchiveSinkInputSpec
    output_spec = ArchiveSinkOutputSpec

    frequency = 'per_session'

    def _base_outputs(self):
        outputs = self.output_spec().get()
        outputs['subject_id'] = self.inputs.subject_id
        outputs['visit_id'] = self.inputs.visit_id
        return outputs


class ArchiveSubjectSink(BaseArchiveSink):

    input_spec = ArchiveSubjectSinkInputSpec
    output_spec = ArchiveSubjectSinkOutputSpec

    frequency = 'per_subject'

    def _base_outputs(self):
        outputs = self.output_spec().get()
        outputs['subject_id'] = self.inputs.subject_id
        return outputs


class ArchiveVisitSink(BaseArchiveSink):

    input_spec = ArchiveVisitSinkInputSpec
    output_spec = ArchiveVisitSinkOutputSpec

    frequency = 'per_visit'

    def _base_outputs(self):
        outputs = self.output_spec().get()
        outputs['visit_id'] = self.inputs.visit_id
        return outputs


class ArchiveProjectSink(BaseArchiveSink):

    input_spec = ArchiveProjectSinkInputSpec
    output_spec = ArchiveProjectSinkOutputSpec

    frequency = 'per_project'

    def _base_outputs(self):
        outputs = self.output_spec().get()
        return outputs
