from nipype.interfaces.ants import N4BiasFieldCorrection
from src.utils import splitext 
from nipype.interfaces.ants.segmentation import BrainExtraction
from src.arg_parser import icbm_default_template, file_dir
from nipype.interfaces.base import CommandLine, CommandLineInputSpec
from nipype.interfaces.base import (TraitedSpec, File, traits, InputMultiPath, 
        BaseInterface, OutputMultiPath, BaseInterfaceInputSpec, isdefined)
from nipype.interfaces.ants import registration, segmentation
from nipype.interfaces.ants.segmentation import Atropos
from nipype.interfaces.ants import Registration, ApplyTransforms
from src.utils import copyCommand
from src.ants import APPIANRegistration, APPIANApplyTransforms
from nipype.interfaces.utility import Rename
import nipype.pipeline.engine as pe
import nipype.interfaces.utility as niu
import nipype.interfaces.utility as util
import src.initialization as init
import nipype.interfaces.io as nio
import nipype.interfaces.minc as minc
import src.ants_nibabel as nib
import numpy as np
import ntpath
import os


global icbm_default_csf  
global icbm_default_gm
global icbm_default_wm 

icbm_default_csf=file_dir+os.sep+"/atlas/MNI152/mni_icbm152_csf_tal_nlin_asym_09c.nii.gz"
icbm_default_gm=file_dir+os.sep+"/atlas/MNI152/mni_icbm152_gm_tal_nlin_asym_09c.nii.gz"
icbm_default_wm=file_dir+os.sep+"/atlas/MNI152/mni_icbm152_wm_tal_nlin_asym_09c.nii.gz"
icbm_default_brain=file_dir+os.sep+"/atlas/MNI152/mni_icbm152_t1_tal_nlin_asym_09c_mask.nii.gz"

def get_workflow(name, opts):
    workflow = pe.Workflow(name=name)
    in_fields = ['mri']
    if opts.user_brainmask :
        in_fields += ['brain_mask_space_stx']

    if opts.user_mri_stx :
        in_fields += ['tfm_mri_stx', 'tfm_stx_mri']

    label_types = [ opts.results_label_type]
    label_imgs = [ opts.results_label_img  ]
    stages = ['results']

    if opts.pvc_method != None :
        stages += ['pvc']
        label_imgs += [ opts.pvc_label_img  ]
        label_types += [ opts.pvc_label_type ]

    if opts.quant_method != None :
        stages += ['quant']
        label_imgs += [ opts.results_label_img  ]
        label_types += [ opts.results_label_type ]


    inputnode = pe.Node(niu.IdentityInterface(fields=in_fields), name="inputnode")

    out_fields=['tfm_stx_mri', 'tfm_mri_stx', 'brain_mask_space_stx', 'brain_mask_space_mri', 'mri_space_stx', 'mri_space_nat', 'template_space_mri' ]
    for stage, label_type in zip(stages, label_types):
        if 'internal_cls' == label_type :
            out_fields += [ stage+'_label_img']

    outputnode = pe.Node(niu.IdentityInterface(fields=out_fields), name='outputnode')

    ##########################################
    # T1 spatial (+ intensity) normalization #
    ##########################################
    if opts.n4_bspline_fitting_distance != 0 :
        n4 =  pe.Node(N4BiasFieldCorrection(), "mri_intensity_normalized" )
        workflow.connect(inputnode, 'mri', n4, 'input_image')
        n4.inputs.dimension = 3 
        n4.inputs.bspline_fitting_distance = opts.n4_bspline_fitting_distance 
        n4.inputs.shrink_factor = opts.n4_shrink_factor
        n4.inputs.n_iterations = opts.n4_n_iterations
        n4.inputs.convergence_threshold = opts.n4_convergence_threshold
    else :
        n4 = pe.Node(niu.IdentityInterface(fields=["output_image"]), name='mri_no_intensity_normalization')
        workflow.connect(inputnode, 'mri', n4, 'output_image')

    # If the use does not provide a transformation file to align the MRI to the stereotaxic template, then use ants to do the registration
    if opts.user_mri_stx == '':
        mri2template = pe.Node(interface=APPIANRegistration(), name="mri_spatial_normalized")
        mri2template.inputs.moving_image_space="T1w"
        mri2template.inputs.fixed_image_space="stx"
        mri2template.inputs.fixed_image_mask = icbm_default_brain
        mri2template.inputs.fixed_image = opts.template
        workflow.connect(n4, 'output_image', mri2template, 'moving_image')
        if opts.user_ants_command != None :
            mri2template.inputs.user_ants_command = opts.user_ants_command
        if opts.normalization_type :
            mri2template.inputs.normalization_type = opts.normalization_type

        mri_stx_file = 'warped_image'
        mri_stx_node = mri2template
        template_nat_file = 'inverse_warped_image'

        tfm_node= mri2template
        tfm_inv_node= mri2template
        if opts.normalization_type == 'nl' :
            tfm_file='composite_transform'
            tfm_inv_file='inverse_composite_transform'
        elif opts.normalization_type == 'affine' or opts.normalization_type == 'rigid' :
            tfm_file='out_matrix'
            tfm_inv_file='out_matrix_inverse'
        else :
            print("Error: --normalization-type should be either rigid, lin, or nl")
            exit(1)
    else :
        transform_mri = pe.Node(interface=APPIANApplyTransforms(), name="transform_mri"  )
        transform_mri.inputs.target_space='t1'
        workflow.connect(inputnode, 'mri', transform_mri, 'input_image')
        workflow.connect(inputnode, 'tfm_mri_stx', transform_mri, 'transform_1')
        transform_mri.inputs.reference_image = opts.template
        transform_mri.inputs.create_inverse_image = True

        mri_stx_node = transform_mri
        mri_stx_file = 'output_image'
        template_nat_file = 'output_image_inverse'
        tfm_node = inputnode
        tfm_file = 'tfm_mri_stx'
        tfm_inv_node=inputnode
        tfm_inv_file='tfm_stx_mri'

    #
    # T1 in native space will be part of the APPIAN target directory
    # and hence it won't be necessary to link to the T1 in the source directory.
    #
    copy_mri_nat = pe.Node(interface=copyCommand(), name="mri_nat"  )
    workflow.connect(inputnode, 'mri', copy_mri_nat, 'input_file')

    ###################################
    # Segment T1 in Stereotaxic space #
    ###################################
    atropos = None

    if opts.ants_atropos_priors == [] and opts.template == icbm_default_template :
        opts.ants_atropos_priors =  file_dir+os.sep+"/atlas/MNI152/mni_icbm152_%02d_tal_nlin_asym_09c.nii.gz"
    if opts.ants_atropos_priors == [] :
        print("Warning : user did not provide alternative priors for template. This will affect your T1 MRI segmentation. Check this segmentation visually to make sure it is what you want ")

    for stage, label_type, img in zip(stages, label_types, label_imgs) :
        if  atropos == None :
            atropos = pe.Node(interface=Atropos(), name='atropos')
            atropos.inputs.dimension=3
            atropos.inputs.number_of_tissue_classes=3 #len(opts.ants_atropos_priors)
            atropos.inputs.initialization = 'PriorProbabilityImages'
            atropos.inputs.prior_weighting = opts.ants_atropos_prior_weighting
            atropos.inputs.prior_image = opts.ants_atropos_priors
            atropos.inputs.likelihood_model = 'Gaussian'
            atropos.inputs.posterior_formulation = 'Socrates'
            atropos.inputs.use_mixture_model_proportions = True
            atropos.inputs.args="-v 1"
            workflow.connect(mri_stx_node, mri_stx_file,  atropos, 'intensity_images' )
            atropos.inputs.mask_image = icbm_default_brain

            seg = pe.Node(interface=fix_nifti(), name='segmentation_ants' ) 

            #workflow.connect(brain_mask_node, brain_mask_file,  seg, 'mask_image' )

            workflow.connect(atropos, 'classified_image',  seg, 'input_image' )
        print(stage, img) 
        if 'antsAtropos' == img :
            workflow.connect(seg, 'classified_image', outputnode, stage+'_label_img')

    ####################
    # T1 Brain masking #
    ####################
    if not opts.user_brainmask :
        #  if opts.brain_extraction_method == 'beast':
       #      #Brain Mask MNI-Space
       #      mriMNI_brain_mask = pe.Node(interface=beast(), name="mri_stx_brain_mask")
       #      mriMNI_brain_mask.inputs.library_dir  = library_dir
       #      mriMNI_brain_mask.inputs.template  = library_dir+"/margin_mask.mnc"
       #      mriMNI_brain_mask.inputs.configuration = mriMNI_brain_mask.inputs.library_dir+os.sep+"default.2mm.conf"
       #      mriMNI_brain_mask.inputs.same_resolution = True
       #      mriMNI_brain_mask.inputs.median = True
       #      mriMNI_brain_mask.inputs.fill = True
       #      mriMNI_brain_mask.inputs.median = True

       #      workflow.connect(mri_stx_node, mri_stx_file, mriMNI_brain_mask, "in_file" )

       #      brain_mask_node = mriMNI_brain_mask
       #      brain_mask_file = 'out_file'
       #  else : 
            #mriMNI_brain_mask = pe.Node(interface=BrainExtraction(), name="mri_stx_brain_mask")
            #mriMNI_brain_mask.inputs.dimension = 3
            #mriMNI_brain_mask.inputs.brain_template = opts.template
            #template_base, template_ext = splitext(opts.template)
            #mriMNI_brain_mask.inputs.brain_probability_mask =template_base+'_variant-brain_pseg'+template_ext

        mriMNI_brain_mask = pe.Node(interface=SegmentationToBrainMask(), name="mri_stx_brain_mask")
        #workflow.connect(mri_stx_node, mri_stx_file, mriMNI_brain_mask, "anatomical_image" )
        workflow.connect(seg, 'classified_image', mriMNI_brain_mask, "seg_file" )

        brain_mask_node = mriMNI_brain_mask
        brain_mask_file = 'output_image'
    else :			
        brain_mask_node = inputnode
        brain_mask_file = 'brain_mask_space_stx'

    #
    # Transform brain mask from stereotaxic to T1 native space
    #
    transform_brain_mask = pe.Node(interface=APPIANApplyTransforms(),name="transform_brain_mask")
    transform_brain_mask.inputs.interpolation = 'NearestNeighbor'
    transform_brain_mask.inputs.target_space = 't1'
    workflow.connect(brain_mask_node, brain_mask_file, transform_brain_mask, 'input_image')
    workflow.connect(tfm_node, tfm_inv_file, transform_brain_mask, 'transform_1')
    workflow.connect(copy_mri_nat,'output_file', transform_brain_mask,'reference_image') 

    ###############################
    # Pass results to output node #
    ###############################
    workflow.connect(brain_mask_node, brain_mask_file, outputnode, 'brain_mask_space_stx')
    workflow.connect(tfm_node, tfm_file, outputnode, 'tfm_mri_stx' )
    workflow.connect(tfm_node, tfm_inv_file, outputnode, 'tfm_stx_mri' )
    workflow.connect(transform_brain_mask, 'output_image', outputnode, 'brain_mask_space_mri')
    #workflow.connect(mri_stx_node, mri_stx_file, outputnode, 'mri_space_stx')
    workflow.connect(copy_mri_nat, 'output_file', outputnode, 'mri_space_nat')
    
    workflow.connect(tfm_node, template_nat_file, outputnode, 'template_space_mri')
    return(workflow)


class SegmentationToBrainMaskOutput(TraitedSpec):
    output_image = File(argstr="%s",  desc="Brain Mask")

class SegmentationToBrainMaskInput(CommandLineInputSpec):
    output_image= File(argstr="%s",  desc="Brain Mask", position=-1)
    seg_file = File(argstr="%s",  desc="Segmentation", position=-1)


class SegmentationToBrainMask(BaseInterface):
    input_spec =  SegmentationToBrainMaskInput
    output_spec = SegmentationToBrainMaskOutput


    def _run_interface(self, runtime) :
        if not isdefined(self.inputs.output_image):
            self.inputs.output_image = self._gen_output(self.inputs.seg_file)

        img = nib.load(self.inputs.seg_file)
        data = img.get_data()
        data[ data > 1 ] = 1
        out = nib.Nifti1Image(data.astype(np.int16), img.get_affine())
        out.to_filename (self.inputs.output_image)

        return runtime

    def _gen_output(self, basefile):
        fname = ntpath.basename(basefile)
        fname_list = splitext(fname) # [0]= base filename; [1] =extension
        dname = os.getcwd()
        return dname+ os.sep+fname_list[0] + "_brain_mask" + fname_list[1]

    def _list_outputs(self):
        if not isdefined(self.inputs.output_image):
            self.inputs.output_image = self._gen_output(self.inputs.seg_file)
        outputs = self.output_spec().get()
        outputs["output_image"] = self.inputs.output_image
        return outputs



class fix_niftiOutput(TraitedSpec):
    classified_image = File(argstr="%s",  desc="Brain Mask")

class fix_niftiInput(CommandLineInputSpec):
    input_image = File(argstr="%s",  desc="Brain Mask", position=-1)
    classified_image = File(argstr="%s",  desc="Brain Mask", position=-1)

class fix_nifti(BaseInterface):
    input_spec = fix_niftiInput 
    output_spec = fix_niftiOutput 

    def _run_interface(self, runtime) :
        if not isdefined(self.inputs.classified_image):
            self.inputs.classified_image = os.getcwd() + os.sep + os.path.basename(self.inputs.input_image)
        print('Fixing nifti file')
        print(self.inputs.classified_image)
        nib.load(self.inputs.input_image).to_filename(self.inputs.classified_image)

        return runtime

    def _list_outputs(self):
        outputs = self.output_spec().get()
        outputs["classified_image"] = self.inputs.classified_image
        return outputs


