from core.modules.vistrails_module import Module, ModuleError
from Array import *
import numpy
import scipy
import scipy.stats

class ArrayStatsModule(object):
    namespace = 'scipy|statistics|functions'

class RelatedTTest(ArrayStatsModule, Module):
    namespace = ArrayStatsModule.namespace + '|tests'
    __doc__ = scipy.stats.ttest_rel.__doc__
    
    def compute(self):
        array_a = self.getInputFromPort("Array 1").get_array()
        array_b = self.getInputFromPort("Array 2").get_array()

        axis_ = self.forceGetInputFromPort("Axis", 0)
        ravel_ = self.forceGetInputFromPort("Ravel", False)
        if ravel_:
            axis_ = None
            array_a = array_a.ravel()
            array_b = array_b.ravel()
            
        (t, p) = scipy.stats.ttest_rel(array_a, array_b, axis=axis_)

        if type(t) == type(0.0):
            t = numpy.array([t])
        if type(p) == type(0.0):
            p = numpy.array([p])

        t_out = NDArray()
        t_out.set_array(t)
        p_out = NDArray()
        p_out.set_array(p)
        self.setResult("T Array", t_out)
        self.setResult("P Array", p_out)

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name="Related T-Test", namespace=cls.namespace)

    @classmethod
    def register_ports(cls, reg, basic):
        reg.add_input_port(cls, "Array 1", (NDArray, 'Input Array 1'))
        reg.add_input_port(cls, "Array 2", (NDArray, 'Input Array 2'))
        reg.add_input_port(cls, "Axis", (basic.Integer, 'Axis to Calculate On'), True)
        reg.add_input_port(cls, "Ravel", (basic.Boolean, 'Ravel arrays before test'), True)
        reg.add_output_port(cls, "T Array", (NDArray, 'Output T Values'))
        reg.add_output_port(cls, "P Array", (NDArray, 'Output P Values'))