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

class ContinuousDistribution(object):
    namespace = 'scipy|statistics|distributions|continuous'

class ContinuousRandomVariable(ContinuousDistribution, Module):
    name = "Continuous Distribution"

    def __init__(self):
        Module.__init__(self)
        self.rv = None

    def value(self):
        return self.rv

    def value(self, distribution):
        self.rv = distribution

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)

    @classmethod
    def register_ports(cls, reg, basic):
        pass

class FindPDF(ContinuousDistribution, Module):
    __doc__ = scipy.stats.rv_continuous.pdf.__doc__
    name = 'Calculate PDF'
    namespace = ContinuousDistribution.namespace + '|functions'
    
    def compute(self):
        distribution = self.getInputFromPort("Distribution")
        lo_val = self.forceGetInputFromPort("Low Value", 0.)
        hi_val = self.forceGetInputFromPort("High Value", 10.)
        x = numpy.linspace(lo_val, hi_val)
        out_ar = distribution.pdf(x)
        out = NDArray()
        out.set_array(out_ar)

        self.setResult("PDF Array", out)
        self.setResult("Distribution", distribution)

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)

    @classmethod
    def register_ports(cls, reg, basic):
        reg.add_input_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))
        reg.add_input_port(cls, "Low Value", (basic.Float, 'Low Value of PDF'))
        reg.add_input_port(cls, "High Value", (basic.Float, 'High Value of PDF'))
        reg.add_output_port(cls, "PDF Array", (NDArray, 'PDF Array'))
        reg.add_output_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))

class FindCDF(ContinuousDistribution, Module):
    __doc__ = scipy.stats.rv_continuous.cdf.__doc__
    name = 'Calculate CDF'
    namespace = ContinuousDistribution.namespace + '|functions'

    def compute(self):
        distribution = self.getInputFromPort("Distribution")
        lo_val = self.forceGetInputFromPort("Low Value", -5.)
        hi_val = self.forceGetInputFromPort("High Value", 5.)
        x = numpy.linspace(lo_val, hi_val)
        out_ar = distribution.pdf(x)
        out = NDArray()
        out.set_array(out_ar)

        self.setResult("CDF Array", out)
        self.setResult("Distribution", distribution)

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)

    @classmethod
    def register_ports(cls, reg, basic):
        reg.add_input_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))
        reg.add_input_port(cls, "Low Value", (basic.Float, 'Low Value of CDF'))
        reg.add_input_port(cls, "High Value", (basic.Float, 'High Value of CDF'))
        reg.add_output_port(cls, "CDF Array", (NDArray, 'CDF Array'))
        reg.add_output_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))

class FindMoments(ContinuousDistribution, Module):
    __doc__ = scipy.stats.rv_continuous.stats.__doc__
    name = 'Find Moments'
    namespace = ContinuousDistribution.namespace + '|functions'

    def compute(self):
        distribution = self.getInputFromPort("Distribution")
        array_like = self.forceGetInputListFromPort("Additional Arrays")
        ar_list = []
        if array_like != None:
            ar_list = [ar.get_array() for ar in array_like]
        ar_list.append('moments="mvsk"')
        (mean, variance, skew, kurtosis) = getattr(distribution, 'stats')(ar_list)

        self.setResult("Mean", mean)
        self.setResult("Variance", variance)
        self.setResult("Skew", skew)
        self.setResult("Kurtosis", kurtosis)
        self.setResult("Distribution", distribution)

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)

    @classmethod
    def register_ports(cls, reg, basic):
        reg.add_input_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))
        reg.add_input_port(cls, "Additional Arrays", (NDArray, 'Array Like'), True)
        reg.add_output_port(cls, "Mean", (NDArray, 'Mean Value'))
        reg.add_output_port(cls, "Variance", (NDArray, 'Variance Value'))
        reg.add_output_port(cls, "Skew", (NDArray, 'Skew Value'))
        reg.add_output_port(cls, "Kurtosis", (NDArray, 'Kurtosis Value'))
        reg.add_output_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))

class NormalDistribution(ContinuousDistribution, Module):
    __doc__ = scipy.stats.norm.__doc__
    name = "Normal Distribution"

    def compute(self):
        mean = self.getInputFromPort("Mean")
        stddev = self.getInputFromPort("Std Dev")
        num_pts = self.getInputFromPort("Size")

        rv = ContinuousRandomVariable()
        print "1"
        dist = scipy.stats.norm(mean, stddev, num_pts)
        print "2"
        rv.value(dist)
        print "3"
        out = NDArray()
        out_ar = dist.rvs(mean, stddev, num_pts)
        print "4"
        out.set_array(out_ar)

        self.setResult("Distribution Array", out)
        self.setResult("Distribution", rv)

    @classmethod
    def register(cls, reg, basic):
        reg.add_module(cls, name=cls.name, namespace=cls.namespace)

    @classmethod
    def register_ports(cls, reg, basic):
        reg.add_input_port(cls, "Mean", (basic.Float, 'Mean Value'))
        reg.add_input_port(cls, "Std Dev", (basic.Float, 'Std. Deviation'))
        reg.add_input_port(cls, "Size", (basic.Integer, 'Size of Distribution'))
        reg.add_output_port(cls, "Distribution Array", (NDArray, 'Distribution Array'))
        reg.add_output_port(cls, "Distribution", (ContinuousRandomVariable, 'Distribution'))