import core.modules
import core.modules.module_registry
from core.modules.vistrails_module import Module, ModuleError
from core.bundles import py_import
scipy = py_import('scipy', {'linux-ubuntu':'python-scipy',
                            'linux-fedora':'scipy'})
import scipy.interpolate

# Numpy package imports
import Array
import ArrayAccess
import ArrayOperations
import ArrayUtilities
import ArrayConvert
import ArrayIO
import Imaging

# Scipy package imports
import Matrix
import MatrixUtilities
import DSP
import Filters
import EnsembleOrdering

# Statistics package imports
import StatFunctions
import StatDistributions
import ContinuousDistributions

# NumSciPy plotting
import ArrayPlot

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

    clslst = []
    
    #########################################################################################
    #  Numpy Registry
    reg.add_module(Array.NDArray, name="Numpy Array", namespace=Array.NDArray.my_namespace)
    reg.add_output_port(Array.NDArray, "self", (Array.NDArray, 'self'))
    reg.add_module(Matrix.Matrix, name="Scipy Matrix", namespace=Matrix.Matrix.my_namespace)
    
    
    #########################################################################################
    #  Array Access registry
    accessclass = [ArrayAccess.GetShape,
                   ArrayAccess.GetReals,
                   ArrayAccess.GetImaginaries,
                   ArrayAccess.GetMax,
                   ArrayAccess.GetMean,
                   ArrayAccess.GetMin,
                   ArrayAccess.GetDiagonal,
                   ArrayAccess.GetArrayAsType,
                   ArrayAccess.GetConjugate,
                   ArrayAccess.GetFlattenedArray,
                   ArrayAccess.GetField,
                   ArrayAccess.ToScalar,
                   ArrayAccess.GetMemoryFootprint,
                   ArrayAccess.GetArrayRank,
                   ArrayAccess.GetNonZeroEntries,
                   ArrayAccess.GetArraySize,
                   ArrayAccess.GetTranspose,
                   ArrayAccess.GetRowRange,
                   ArrayAccess.GetColumnRange,
                   ArrayAccess.GetRows]

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


    #########################################################################################
    #  Array Operations registry
    opclasses = [ArrayOperations.ArrayReshape,
                 ArrayOperations.ArrayCumulativeSum,
                 ArrayOperations.ArraySort,
                 ArrayOperations.ArrayCumulativeProduct,
                 ArrayOperations.ArrayFill,
                 ArrayOperations.ArrayAbs,
                 ArrayOperations.ArrayResize,
                 ArrayOperations.ArrayExtractRegion,
                 ArrayOperations.ArrayRavel,
                 ArrayOperations.ArrayRound,
                 ArrayOperations.ArrayGetSigma,
                 ArrayOperations.ArrayGradient,
                 ArrayOperations.ArraySum,
                 ArrayOperations.ArrayElementMultiply,
                 ArrayOperations.ArraySetElement,
                 ArrayOperations.ArrayVariance,
                 ArrayOperations.ArrayTrace,
                 ArrayOperations.ArraySwapAxes,
                 ArrayOperations.ArraySqueeze,
                 ArrayOperations.ArrayScalarMultiply,
                 ArrayOperations.ArrayAdd,
                 ArrayOperations.ArrayScalarAdd,
                 ArrayOperations.ArrayLog10,
                 ArrayOperations.ArrayAtan2,
                 ArrayOperations.ArraySqrt,
                 ArrayOperations.ArrayThreshold,
                 ArrayOperations.ArrayWindow,
                 ArrayOperations.ArrayNormalize,
                 ArrayOperations.ArrayName,
                 ArrayOperations.ArrayFlipLR,
                 ArrayOperations.ArrayFlipUD]

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

    #########################################################################################
    #  Array Convert registry
    convertclasses = [ArrayConvert.ArrayDumpToFile,
                      ArrayConvert.ArrayDumpToString,
                      ArrayConvert.ArrayToFile,
                      ArrayConvert.ArrayToString,
                      ArrayConvert.ArrayToMatrix,
                      ArrayConvert.ArrayToVTKImageData]
    
    for cls in convertclasses:
        cls.register(reg, basic)

    #########################################################################################
    #  Array IO registry
    ioclasses = [ArrayIO.ReadRAW,
                 ArrayIO.ReadNHDR,
                 ArrayIO.WriteRAW,
                 ArrayIO.WriteNHDR,
                 ArrayIO.ReadStatisticalSummary]

    try:
        from core.bundles import py_import
        mpl_dict = {'linux-ubuntu': 'python-matplotlib',
                    'linux-fedora': 'python-matplotlib'}
        matplotlib = py_import('matplotlib', mpl_dict)
        matplotlib.use('Qt4Agg')
        pylab = py_import('pylab', mpl_dict)
        ioclasses.append(ArrayIO.ReadPNG)
        ioclasses.append(ArrayIO.WritePNG)
    except:
        pass
    
    for cls in ioclasses:
        cls.register(reg, basic)

    #########################################################################################
    #  Array Imaging registry
    ar_util = [ArrayUtilities.ArrayToVTKScalars,
               ArrayUtilities.ArrayToTimeVaryingVTKScalars,
               ArrayUtilities.ArrayToTimeVaryingVTKVectors,
               ArrayUtilities.ArrayToVTKVectors,
               ArrayUtilities.VTKDataSetToPointArray]

    clslst.extend(ar_util)
    
    for cls in ar_util:
        cls.register(reg, basic)        

    #########################################################################################
    #  Array Imaging registry
    imclasses = [Imaging.ExtractRGBAChannel,
                 Imaging.GaussianGradientMagnitude,
                 Imaging.JointHistogram,
                 Imaging.GaussianSmooth,
                 Imaging.MedianFilter,
                 Imaging.ImageDifference,
                 Imaging.ImageNormalize,
                 Imaging.SobelGradientMagnitude]

    for cls in imclasses:
        cls.register(reg, basic)
        
    #########################################################################################
    #  Scipy Registry
    matrixclasses = [Matrix.MatrixMultiply,
                     Matrix.MatrixConjugate,
                     Matrix.MatrixToArray]

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

    #########################################################################################
    #  Scipy Matrix Utilities Registry
    matrixutils = [MatrixUtilities.MatlabReader,
                   MatrixUtilities.MatlabWriter]

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

    #########################################################################################
    #  Scipy DSP Registry
    dspclasses = [DSP.FFT,
                  DSP.FFTN,
                  DSP.SignalSmoothing,
                  DSP.ShortTimeFourierTransform,
                  DSP.SignalGenerator]
#                  DSP.FrequencyPhaseLocking,
#                  DSP.SingleTrialPhaseLocking

    try:
        import smt
        import Stockwell
        dspclasses.append(Stockwell.StockwellTransform)
        dspclasses.append(Stockwell.MultiTaperStockwellTransform)
        dspclasses.append(Stockwell.FastStockwell3D)
        dspclasses.append(Stockwell.Stockwell2D)
        dspclasses.append(Stockwell.PyStockwellTransform)
        dspclasses.append(Stockwell.PointBasedStockwell)
        dspclasses.append(Stockwell.ScaleSpaceHistogram)
        dspclasses.append(Stockwell.PyScaleSpaceHistogram)
        dspclasses.append(Stockwell.MaximalScaleVolume)
        dspclasses.append(Stockwell.ScaleVolumes)
        dspclasses.append(Stockwell.IsotropicScaleVolumes)
        dspclasses.append(Stockwell.FrequencyPhaseLocking)
    except:
        pass

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

    #########################################################################################
    #  Scipy Interpolation Registry
    interpclasses = []
    try:
        import ArrayInterpolate
        interpclasses.append(ArrayInterpolate.RBFInterpolate)
        interpclasses.append(ArrayInterpolate.BSplineInterpolate)
    except:
        pass

    interpclasses.append(ArrayInterpolate.BSplineResample)
    for cls in interpclasses:
        cls.register(reg, basic)

    #########################################################################################
    #  Scipy Windows Registry
    winclasses = [Filters.HanningWindow,
                  Filters.TriangularWindow,
                  Filters.BlackmanWindow,
                  Filters.BlackmanHarrisWindow,
                  Filters.ParzenWindow,
                  Filters.HammingWindow,
                  Filters.KaiserWindow,
                  Filters.BartlettHannWindow,
                  Filters.GaussianWindow,
                  Filters.BoxcarWindow,
                  Filters.BohmanWindow,
                  Filters.BartlettWindow,
                  Filters.NuttallBlackmanHarrisWindow]    

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

    #########################################################################################
    #  Scipy Signal Ensembles Registry
    ensembles = [EnsembleOrdering.OrderByIndexes,
                 EnsembleOrdering.OrderByCorrelation,
                 EnsembleOrdering.OrderByProgressiveCorrelation,
                 EnsembleOrdering.ComputeDistance]

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

    #########################################################################################    
    #  NumSciPy Plotting Registry
    pl = [ArrayPlot.LinePlot,
          ArrayPlot.ScatterPlot,
          ArrayPlot.BarChart,
          ArrayPlot.Histogram,
          ArrayPlot.ArrayImage]

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

    clslst.extend(pl)

    #########################################################################################    
    #  Scipy Stats package registry
    stats_cls = [StatFunctions.RelatedTTest]

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

    clslst.extend(stats_cls)
    
    #########################################################################################    
    #  Scipy Stats distributions package registry
    dist_cls = [ContinuousDistributions.ContinuousRandomVariable,
                ContinuousDistributions.FindPDF,
                ContinuousDistributions.FindCDF,
                ContinuousDistributions.FindMoments,
                ContinuousDistributions.NormalDistribution]

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

    clslst.extend(dist_cls)
    
    #########################################################################################    
    #  Register all ports.
    for cls in clslst:
        cls.register_ports(reg, basic)