version = '0.2'
identifier = 'edu.utah.sci.vistrails.itk'
name = 'ITK'

import core.bundles.utils
import core.requirements
from core.modules.vistrails_module import Module, ModuleError

# Ugly, but Carlos doesnt know any better
if core.bundles.utils.guess_system() == 'linux-ubuntu':
    import sys
    sys.path.append('/usr/local/lib/VisTrailsITK')

try:
    from core.bundles import py_import
    itk = py_import('itk', {'linux-ubuntu': 'vistrails-itk'})
except ImportError:
    raise core.requirements.MissingRequirement("ITK and WrapITK")

import core.modules
import core.modules.module_registry

# ITK Package imports
from PixelType import *
from FeatureExtractionFilters import *
from ITK import *
from Image import Image
from IntensityFilters import *
from SegmentationFilters import *
from SelectionFilters import *
from SmoothingFilters import *
from ThresholdFilters import *
from GradientFilters import *
from NeighborhoodFilters import *
from ImageReader import *

def initialize(*args, **keywords):
    reg = core.modules.module_registry
    basic = core.modules.basic_modules

########################################################################################
# Misc.
    Index2D.register(reg,basic)
    Index3D.register(reg,basic)
    Size.register(reg,basic)
    Region.register(reg,basic)
    PixelType.register(reg,basic)
    Filter.register(reg,basic)
    Kernel.register(reg,basic)

    Image.register(reg,basic)

########################################################################################
# Pixel Types
    pixeltypes = [PixelTypeFloat,
                  PixelTypeUnsignedChar,
                  PixelTypeUnsignedShort,
                  PixelTypeRGB]

    for cls in pixeltypes:
        cls.register(reg,basic)


########################################################################################
# Feature Extraction Filters
    featurefilters = [GradientMagnitudeRecursiveGaussianImageFilter,
                      DanielssonDistanceMapImageFilter,
                      SobelEdgeDetectionImageFilter]

    for cls in featurefilters:
        cls.register(reg,basic)

########################################################################################
# Intensity Filters
    intensityfilters = [RescaleIntensityImageFilter,
                        SigmoidImageFilter,
                        ThresholdImageFilter,
                        ShiftScaleImageFilter,
                        NormalizeImageFilter]

    for cls in intensityfilters:
        cls.register(reg,basic)

########################################################################################
# Segmentation Filters
    segmentationfilters = [IsolatedWatershedImageFilter,
                           ConnectedThresholdImageFilter,
                           ConfidenceConnectedImageFilter,
                           IsolatedConnectedImageFilter]

    for cls in segmentationfilters:
        cls.register(reg,basic)

########################################################################################
# Selection Filters
    selectionfilters = [RegionOfInterestImageFilter,
                        CastImageFilter,
                        ExtractImageFilter]

    for cls in selectionfilters:
        cls.register(reg,basic)

########################################################################################
# Smoothing Filters
    smoothingfilters = [CurvatureAnisotropicDiffusionFilter,
                        RecursiveGaussianImageFilter,
                        DiscreteGaussianImageFilter,
                        GradientAnisotropicDiffusionImageFilter,
                        MinMaxCurvatureFlowImageFilter,
                        BinomialBlurImageFilter,
                        BilateralImageFilter,
                        CurvatureFlowImageFilter]

    for cls in smoothingfilters:
        cls.register(reg,basic)

########################################################################################
# Threshold Filters
    thresholdfilters = [BinaryThresholdImageFilter]

    for cls in thresholdfilters:
        cls.register(reg,basic)

########################################################################################
# Gradient Filters
    gradientfilters = [GradientMagnitudeImageFilter]

    for cls in gradientfilters:
        cls.register(reg,basic)

########################################################################################
# Neighborhood Filters
    neighborhoodfilters = [MeanImageFilter,
                           MedianImageFilter,
                           BinaryErodeImageFilter]

    for cls in neighborhoodfilters:
        cls.register(reg,basic)

########################################################################################
# Image Reader
    imagereader = [ImageReader,
                   ImageToFile,
                   GDCMReader,
                   DICOMReader]


    for cls in imagereader:
        cls.register(reg,basic)
