# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/models.reg.ipynb (unless otherwise specified).

__all__ = ['CovariateRegressionLandingModel', 'RegressionLandingModel', 'RegressionEstimatesModel',
           'CovariateRegressionModel', 'NormalRegressionModel', 'LogisticRegressionModel', 'PoissonRegressionModel',
           'NegativeBinomialRegressionModel', 'PoissonZInfRegressionModel', 'MultinomialRegressionModel']

# Cell
from ..mvc import Model
from ..rpy import PyFunction
from powerpy.regression import Regression_Model

# Cell
from traitlets import Bool, Int, Unicode, List, Float
from traittypes import Array
import numpy as np
import string
from IPython.display import display, Javascript

# Cell
class CovariateRegressionLandingModel(Model):

    predictors = Int()
    num_variables = Int()
    predLabels = List()
    _predLabels = List()
    yLabel = Unicode()
    N = Int()

    def setLanding(self):
        self._predLabels = ['X%i' % i for i in range(1, 21)]
        self.yLabel = 'Y'
        self.setPredictors(4)
        self.N = 100

    def setPredictors(self, num):
        self.predictors = num
        self.num_variables = num + 1
        self.predLabels = [self._predLabels[i] for i in range(num)]

    def setLabel(self, index, label):
        if index < self.predictors:
            self.predLabels[index] = label
            self._predLabels[index] = label

    def setYLabel(self, label):
        self.yLabel = label

# Cell
class RegressionLandingModel(CovariateRegressionLandingModel):

    standardized = Bool()
    model_family = Unicode()
    varTypes = List()
    _varTypes = List()
    numLevels = List()
    levels = List()
    yLabel = Unicode()

    def __init__(self):
        self.setLanding()

    def setLanding(self):
        self._varTypes = ['Continuous' for i in range(20)]
        super().setLanding()
        self.setPredictors(4)
        self.standardized = False
        self.numLevels = [1 for i in range(self.predictors)]
        self.levels = [[''] for i in range(self.predictors)]
        self.model_family = 'Gaussian'

    def setPredictors(self, num):
        if num < self.predictors:
            self.levels = self.levels[:num]
            self.numLevels = self.numLevels[:num]
        elif num > self.predictors:
            diff = num - self.predictors
            self.numLevels += [1] * diff
            self.levels += [[''] for i in range(diff)]
        self.varTypes = [self._varTypes[i] for i in range(num)]
        self.predictors = num
        self.num_variables = num + 1
        self.predLabels = [self._predLabels[i] for i in range(num)]

    def setVarType(self, index, varType):
        self.varTypes[index] = varType
        self._varTypes[index] = varType
        if varType == 'Continuous':
            self.setNumLevels(index, 1)
        else:
            self.setNumLevels(index, 2)

    def setNumLevels(self, i, num):
        if i >= self.predictors:
            raise Exception('Index %i out of range when there are %i predictors' % (i, self.predictors))
        if self.varTypes[i] == 'Continuous' and num > 1:
            raise Exception('Cannot set number of levels for a continuous variable at index %i' % i)
        if num == 1:
            self.levels[i] = ['']
        elif num < self.numLevels[i]:
            self.levels[i] = self.levels[i][:num]
        else:
            diff = num - self.numLevels[i]
            if self.levels[i]:
                self.levels[i] = self.levels[i] + [''] * diff
            else:
                self.levels[i] = [''] * diff
        self.numLevels[i] = num

    def setLevel(self, i, j, level):
        assert i < self.predictors, 'Index of predictor out of bounds'
        assert j < self.numLevels[i], 'Index of level name out of bounds'
        self.levels[i][j] = level

# Cell
class RegressionEstimatesModel(RegressionLandingModel):

    fixedVals = List()
    sds = Array(dtype=np.dtype(float))
    parameters = List()
    measure = Unicode()
    num_variables = Int()
    parameters = List()
    boxLabels = List()
    musd_spec = Bool()
    labels = List()
    descriptions = List()
    yParam = List()
    yDescription = List()
    equation = List()
    distributions = List()
    dist_0 = Unicode()
    N = Int()


    def setEstimates(self):
        self.measure = 'sd' # or 'var'
        self.yParam = [0, 0]
        self.equation.append('$' + self.yLabel + ' = $')
        self.distributions = []
        self.dist_0 = 'Gaussian'
        self.parameters.append([0, 1])
        self.equation = []
        self.N = 100

        for i in range(self.predictors):
            if self.varTypes[i] == 'Categorical':
                self.distributions.append('Multinomial')
                self.parameters.append([])
                self.labels.append([])
                self.fixedVals.append([])
                self.equation.append([])
                remainder = 1
                for j in range(self.numLevels[i]): # create level name if empty string
                    level = string.ascii_uppercase[j] if self.levels[i][j] == '' else self.levels[i][j]
                    self.labels[i].append(self.predLabels[i] + ':' + level)
                    if j == self.numLevels[i] - 1: # on last proportion
                        self.parameters[i + 1].append(remainder)
                    else:
                        proportion = round(1/self.numLevels[i], 2) # equal part of 100%
                        self.parameters[i + 1].append(proportion)
                        remainder = round(remainder - proportion, 2)
                    if self.standardized:
                        self.descriptions.append([r'$$\beta_%i$$' % (i + 1), 'Proportion'])
                    else:
                        self.descriptions.append(['b', 'Proportion'])
                    self.fixedVals[i].append(.1)
                    self.equation[i].append(['$' + str(.1) + '$', '$* ' + level + ' +$'])

            else:
                self.distributions.append('Gaussian')
                label = 'X' + str(self.num_variables + 1) if self.predLabels[i] == '' else self.predLabels[i]
                self.labels.append(label)
                self.parameters.append([0, 1])
                self.fixedVals.append(.1)
                self.equation.append(['$' + str(.1) + '$', '$* ' + label + ' +$'])
                if self.standardized:
                    self.descriptions.append([r'$$\beta_%i$$' % (i + 1), None])
                else:
                    self.descriptions.append(['b', 'SD'])

        # remove last '+' from equation
        if isinstance(self.equation[-1], list):
            self.equation[-1][-1] = self.equation[-1][-1][:-2] + '$'
        else:
            self.equation[-1] = self.equation[-1][:-2] + '$'

    def setN(self, N):
        self.N = N

    def setFixedVal(self, i, value):
        if isinstance(i, tuple):
            j, k = i
            self.fixedVals[j][k] = value
            self.equation[j][k][0] = '$' + str(value) + '$'
        else:
            self.fixedVals[i] = value
            self.equation[i][0] = '$' + str(value) + '$'

    def setSD(self, i, sd):
        if self.standardized:
            raise Exception('SD\'s are always 1 when standardized')
        if self.varTypes[i] == 'Categorical':
            raise Exception('Categorical variables do not have SDs')
        self.parameters[i+1][1] = sd

    def setProportion(self, i, j, prop):
        self.parameters[i+1][j] = prop

    def setStandardized(self, std):
        self.standardized = std

    def run(self, **kwargs):
        for i, levels in enumerate(self.numLevels):
            if levels > 1 and int(sum((self.parameters[i+1]))) != 1:
                if 'alert' in kwargs.keys() and kwargs['alert']:
                    display(Javascript('alert("Proportions of categorical variables must sum to 1.");'))
                else:
                    raise Exception('Proportions of categorical variables must sum to 1.')
                return None

        distributions = [self.dist_0] + self.distributions

        attrNames = ['num_variables', 'parameters', 'fixedVals']
        argNames = ['N', 'sims', 'pVal']
        modelAttrs = ['standardized', 'distributions', 'model_family']

        self.function = PyFunction('regression', Regression_Model(self.num_variables, parameters=self.parameters, fixedVals=self.fixedVals,
                                          standardized = self.standardized, distributions=distributions,
                                          model_family = self.model_family,
                                                   ), argNames, attrNames, modelAttrs)
        return self.function.run(N=self.N, **kwargs)

# Cell
class CovariateRegressionModel(CovariateRegressionLandingModel):

    labels = List()
    N = Int()
    matrix = Array()
    standardized = Bool()


    def setEstimates(self):
        self.N = 100
        self.standardized = True
        self.matrix = np.zeros((self.num_variables, self.num_variables), float)
        np.fill_diagonal(self.matrix, 1)

    def setN(self, N):
        self.N = N

    def setMatrix(self, i, j, value):
        # symmetric matrix
        self.matrix[i][j] = value
        self.matrix[j][i] = value

    def run(self, **kwargs):
        attrNames = ['num_variables', 'matrix']
        argNames = ['N', 'sims', 'pVal']

        self.function = PyFunction('regression', Regression_Model(self.num_variables, matrix=self.matrix,
                                                   ), argNames, attrNames, None)
        return self.function.run(**kwargs, N=self.N)

# Cell
class NormalRegressionModel(RegressionEstimatesModel):

    sd = Float()

    def setEstimates(self):
        super().setEstimates()
        self.yDescription = ['SD', None]
        self.setYValue1(1)
        self.model_family = 'gaussian'
        self.dist_0 = 'Gaussian'

    def setYValue1(self, sd):
        self.yParam = [0, sd]
        self.parameters[0] = [0, sd]
        self.sd = sd

    def getYValue1(self):
        return self.sd

# Cell
class LogisticRegressionModel(RegressionEstimatesModel):

    proportion = Float()

    def setEstimates(self):
        super().setEstimates()
        self.yDescription = ['Propotion', None]
        self.setYValue1(.5)
        self.model_family = 'Binomial'
        self.dist_0 = 'Bernoulli'

    def setYValue1(self, prop):
        self.yParam = [prop * (1 - prop)]
        self.parameters[0] = [prop * (1 - prop)]
        self.proportion = prop

    def getYValue1(self):
        return self.proportion

# Cell
class PoissonRegressionModel(RegressionEstimatesModel):

    mean = Float()

    def setEstimates(self):
        super().setEstimates()
        self.yDescription = ['Rate/Mean', None]
        self.setYValue1(1)
        self.model_family = 'Poisson'
        self.dist_0 = 'Poisson'

    def setYValue1(self, mean):
        self.yParam[1] = mean
        self.parameters[0][1] = mean

    def getYValue1(self):
        return self.mean

# Cell
class NegativeBinomialRegressionModel(RegressionEstimatesModel):

    mean = Float()
    sd = Float()

    def setEstimates(self):
        super().setEstimates()
        self.yParam = [0, 0]
        self.yDescription = ['Mean', 'SD']
        self.setYValue1(1)
        self.setYValue2(1)
        self.model_family = 'NegativeBinomial'
        self.dist_0 = 'NegativeBinomial'

    def setYValue1(self, mean):
        self.yParam[0] = mean
        self.parameters[0][0] = mean
        self.mean = mean

    def setYValue2(self, sd):
        self.yParam[1] = sd
        self.parameters[0][1] = sd
        self.sd = sd

    def getYValue1(self):
        return self.mean

    def getYValue2(self):
        return self.sd

# Cell
class PoissonZInfRegressionModel(RegressionEstimatesModel):

    mean = Float()
    proportion = Float()

    def setEstimates(self):
        super().setEstimates()
        self.yParam = [0, 0]
        self.yDescription = ['Rate', 'Zero-Inflation Rate']
        self.setYValue1(1)
        self.setYValue2(.5)
        self.model_family = 'Poisson_zinf'
        self.dist_0 = 'Poisson_zinf'

    def setYValue1(self, mean):
        self.yParam[0] = mean
        self.parameters[0][0] = mean
        self.mean = mean

    def getYValue1(self):
        return self.mean

    def setYValue2(self, prop):
        sd = prop * (1 - prop)
        self.yParam[1] = sd
        self.parameters[0][1] = sd
        self.proportion = prop

    def getYValue2(self):
        return self.proportion

# Cell
class MultinomialRegressionModel(RegressionEstimatesModel):

    mean = Float()
    sd = Float()

    def setEstimates(self):
        super().setEstimates()
        self.yParam = [0, 0]
        self.yDescription = ['Mean', 'SD']
        self.setYValue1(1)
        self.setYValue2(.5)
        self.model_family = 'Multinomial'
        self.dist_0 = 'Multinomial'

    def setYValue1(self, mean):
        self.yParam[0] = mean
        self.parameters[0][0] = mean
        self.mean = mean

    def setYValue2(self, sd):
        self.yParam[1] = sd
        self.parameters[0][1] = sd
        self.sd = sd

    def getYValue1(self):
        return self.mean

    def getYValue2(self):
        return self.sd