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

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

class ContinuousRandomVariable(Distributions, Module):
    name = "Continuous Distribution"
    namespace = 'scipy|statistics|distributions|continuous'

    def __init__(self):
        self.rv = None

    @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 DiscreteRandomVariable(Distributions, Module):
    name = "Discrete Distribution"
    namespace = Distributions.namespace + '|discrete'

    def __init__(self):
        self.rv = None

    @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(Distributions, Module):
    __doc__ = scipy.stats.rv_continuous.pdf.__doc__
    name = 'Calculate PDF'
    namespace = Distributions.namespace + '|continuous'
    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", (RandomVariable, '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", (RandomVariable, 'Distribution'))

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

    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", (RandomVariable, '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", (RandomVariable, 'Distribution'))

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

    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", (RandomVariable, '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", (RandomVariable, 'Distribution'))

class NormalVariates(Distributions, Module):
    __doc__ = scipy.stats.norm.__doc__
    namespace = Distributions.namespace + '|continuous'
    name = 'Normal Variates'
    
    def compute(self):
        mean = self.getInputFromPort("Mean")
        stddev = self.getInputFromPort("Std Dev")
        num_pts = self.getInputFromPort("Size")

        rv = RandomVariable()
        rv.rv = scipy.stats.norm(mean, stddev, num_pts)
        out_ar = rv.rv(mean, stddev, num_pts)
        out = NDArray()
        out.set_array(out_ar)

        self.setResult("Distribution Array", out)
        self.setResult("Random Variable", 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", (RandomVariable, 'Distribution'))