"""
Stages DICOM series directories for import into CTP.
"""

from nipype.caching import Memory
from nipype.interfaces.base import (traits,
    BaseInterfaceInputSpec, TraitedSpec, BaseInterface,
    InputMultiPath, OutputMultiPath, Directory)
from . import GroupDicom, FixDicom, Copy
from ..staging.staging_helpers import match_session_hierarchy

class StageCTPInputSpec(BaseInterfaceInputSpec):
    collection = traits.Str(desc='The image collection', mandatory=True)

    subject_dirs = InputMultiPath(Directory(exists=True), mandatory=True,
        desc='The input subject directories to stage')
    
    dest = Directory(desc='The output directory (default working directory)')


class StageCTPOutputSpec(TraitedSpec):
    series_dirs = OutputMultiPath(Directory(exists=True),
        desc='The output series directories')


class StageCTP(BaseInterface):
    input_spec = StageCTPInputSpec
    
    output_spec = StageCTPOutputSpec

    def _run_interface(self, runtime):
        opts = dict(dest=self.inputs.dest)
        self.series_dirs = _stage(*self.inputs.subject_dirs)
        return runtime

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


    def _stage(self, dest, *subject_dirs):
        """
        Stages the given directories for import into CTP.

        @param subject_dirs: the AIRC source subject directories to stage
        @return: the staged series directories
        """
        # Create a new memory context.
        mem = Memory('.')
        # The destination directory.
        dest = self.inputs.dest or os.getcwd()
        # The FixDicom nipype node factory.
        fixer = mem.cache(FixDicom)
        # The Move nipype node factory.
        mover = mem.cache(Move)
        for d in self.input.subject_dirs:
            fixed = fixer(source=d, dest=dest, collection=self.inputs.collection)
            for f in fixed.outputs.out_files:
                # The subject, session and series directory components.
                sbj, sess, ser = match_session_hierarchy(f)
                # The copy target directory.
                tgt = os.path.join(self.ctp_dir, sbj, sess, ser)
                # Move the file to the target series directory.
                mover(in_file=f, dest=tgt)
                # Add the move target series directory to the staged set.
                staged.add(tgt)
        return staged

@@ -55,16 +55,26 @@
 P.%0A%22%22%22%0A%0A
+import os%0A
 from nip
@@ -265,34 +265,36 @@
 ort 
-GroupDicom, FixDicom, Copy
+FixDicom, MapCTP, XNATUpload
 %0Afro
@@ -328,37 +328,36 @@
  import match_se
-ssion
+ries
 _hierarchy%0A%0Aclas
@@ -480,22 +480,21 @@
 )%0A%0A    s
-ubject
+eries
 _dirs = 
@@ -573,22 +573,21 @@
  input s
-ubject
+eries
  directo
@@ -1050,16 +1050,21 @@
 _dirs = 
+self.
 _stage(*
@@ -1076,22 +1076,21 @@
 inputs.s
-ubject
+eries
 _dirs)%0A 
@@ -1255,17 +1255,16 @@
 utputs%0A%0A
-%0A
     def 
@@ -1284,22 +1284,21 @@
 dest, *s
-ubject
+eries
 _dirs):%0A
@@ -1384,22 +1384,21 @@
 @param s
-ubject
+eries
 _dirs: t
@@ -1747,32 +1747,99 @@
 The 
-Move nipype node fa
+XNAT uploader.%0A        xnat = mem.Cache(XNATUpload)%0A        # The staged series dire
 ctor
-y
+ies
 .%0A  
@@ -1848,31 +1848,90 @@
     
-mover = mem.cache(Move)
+staged = %5B%5D%0A        # The subject =%3E DICOM file dictionary.%0A        sbj_files = %7B%7D
 %0A   
@@ -1958,16 +1958,16 @@
 nput
+s
 .s
-ubject
+eries
 _dir
@@ -1985,200 +1985,173 @@
     
-fixed = fixer(source=d, dest=dest, collection=self.inputs.collection)%0A            for f in fixed.outputs.out_files:%0A                # The subject, session and series directory components.%0A    
+# Fix the input DICOM files.%0A            fix_result = fixer(series_dir=d, collection=self.inputs.collection)%0A            # Add the fixed result to the return value.%0A
     
@@ -2187,13 +2187,12 @@
 h_se
-ssion
+ries
 _hie
@@ -2198,17 +2198,17 @@
 erarchy(
-f
+d
 )%0A      
@@ -2217,63 +2217,22 @@
     
-    # The copy 
+s
 ta
-r
 ge
-t directory.%0A                tgt = 
+d.append(
 os.p
@@ -2244,20 +2244,12 @@
 oin(
-self.ctp_dir
+dest
 , sb
@@ -2261,20 +2261,18 @@
 ss, ser)
+)
 %0A
-   
         
@@ -2279,99 +2279,156 @@
     
- # Move the file to the target series directory.%0A                mover(in_file=f, dest=tgt)
+# Pick one file for each subject to map.%0A            if not sbj_files.has_key(sbj):%0A                sbj_files%5Bsbj%5D = fix_result.outputs.out_files%5B0%5D
 %0A   
@@ -2424,32 +2424,33 @@
 iles%5B0%5D%0A        
+%0A
         # Add th
@@ -2447,96 +2447,321 @@
   # 
-Add
+Map
  the 
-move target series directory to the staged set.%0A                staged.add(tg
+subject ids.%0A        sbj_dirs = %5Bmatch_series_hierarchy(d)%5B0%5D for d in staged%5D%0A        map_result = mem.cache(MapCTP)(collection=self.inputs.collection, subject_files=sbj_files)%0A        # Copy the id map to the destination.%0A        mem.cache(Copy)(in_file=map_result.outputs.out_file, dest=des
 t)%0A
+        %0A
     
@@ -2778,8 +2778,17 @@
  staged%0A
+        %0A

Call on series_dir rather than subject_dir.