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

__all__ = ['CovariateRegressionLandingView', 'CovariateRegressionLanding', 'CovariateRegressionView',
           'CovariateRegression', 'RegressionLandingView', 'RegressionLanding', 'StandardizedRegressionLanding',
           'RegressionView', 'Regression']

# Cell
from ..widgets import BoundedIntText, BoundedFloatText, RadioButtons, Dropdown, HTMLMath, FloatText, Text, Grid, Label, AlignedLabel
from ..widget_containers import  HBox, Box, VBox
from ..mvc import Controller

# Cell
import ipywidgets as ipyw
import traitlets
from functools import partial
import string

# Cell
class CovariateRegressionLandingView(VBox):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.maxPredictors = 20
        self.maxLevels=6
        layout = {'description_width': 'initial', 'max_width': '130px', 'margin': '2px 2px 40px 2px'}
        self.numPredictors = BoundedIntText(description='Predictors', layout=ipyw.Layout(**layout), min=1, max=self.maxPredictors)

        self.yLabel = Text(placeholder='Y', layout=ipyw.Layout(width='90px'))
        first = HBox(children=(self.yLabel, ), layout=ipyw.Layout(margin='1px 15px 1px 15px'))
        self.labels = [Text(placeholder='X' + str(i+1), layout=ipyw.Layout(width='90px')) for i in range(self.maxPredictors)]

        self.variables = []
        self.levelsBoxes = []
        for i in range(self.maxPredictors):
            topBox = HBox(layout=ipyw.Layout(margin='0 15px 0 15px', flex='0 1 auto'))
            topBox.children = (self.labels[i],)
            self.variables.append(VBox(children=(topBox,)))

        flexLayout= {'flex_flow': 'column wrap', 'width': '100%', 'height': '340px'}
        self.flexbox = Box(layout=flexLayout, children = tuple(self.variables))
        self.children += (self.numPredictors, first, self.flexbox)

# Cell
class CovariateRegressionLanding(CovariateRegressionLandingView, Controller):

    def __init__(self, model, **kwargs):
        super().__init__(**kwargs)
        self.model = model
        self.numPredictors.observe(self.observePredictors, 'value')
        for i in range(self.maxPredictors):
            self.labels[i].observe(partial(self.observeLabels, i), 'value')
        self.setLanding()

    def setLanding(self):
        self.numPredictors.value = self.model.predictors
        self.yLabel.value = self.model.yLabel
        for i, label in enumerate(self.model._predLabels):
            self.labels[i].value = label

    def observePredictors(self, change):
        self.flexbox.children = tuple(self.variables[:change['new']])
        self.model.setPredictors(change['new'])

    def observeLabels(self, index, change):
        self.model.setLabel(index, change['new'])

    def observeYLabel(self, change):
        self.model.setYLabel(change['new'])

# Cell
class CovariateRegressionView(VBox):

    def __init__(self, num_variables, **kwargs):
        super().__init__(**kwargs)

        n = num_variables + 1

        grid_lo = {'max_width': '%ipx' % (n*70), 'grid_gap':'3px 5px'}
        child_lo = {'width': 'auto', 'height': 'auto', 'min_width': '0', 'grid_gap': '5px 5px'}

        self.N = BoundedIntText(description='N',min=1, max=999999, style={'description_width': 'initial'}, layout=ipyw.Layout(width='100px', margin='2px 2px 30px 55px'))
        self.grid = Grid(n, n, layout=ipyw.Layout(**grid_lo))
        self.colLabels = [AlignedLabel(value='X%i' % i, justification='flex-end', layout=ipyw.Layout(**child_lo)) for i in range(n-1)]
        self.rowLabels = [AlignedLabel(value='X%i' % i, alignment='flex-start', layout=ipyw.Layout(**child_lo)) for i in range(n-1)]
        self.vars = []
        self.covariates = []
        for i in range(0, n-1):
            self.grid[i, 0] = self.colLabels[i]
        for i in range(1, n):
            self.grid[n-1, i] = self.rowLabels[i-1]
        for i in range(0, n-1):
            for j in range(0, i+1):
                var = FloatText(step=0.01, layout=ipyw.Layout(**child_lo))
                self.covariates.append(var)
                self.grid[i, j+1] = var
        self.children += (self.N, self.grid, )

# Cell
class CovariateRegression(CovariateRegressionView):

    def __init__(self, model):
        super().__init__(num_variables = model.num_variables)
        self.model = model
        self.setEstimates()
        for i in range(self.model.num_variables):
            for j in range(i+1):
                self.grid[i, j+1].observe(partial(self.observeMatrix, i, j), 'value')
        self.N.observe(self.observeN, 'value')

    def setEstimates(self):
        self.N.value = self.model.N
        self.rowLabels[0].value = self.model.yLabel
        self.colLabels[0].value = self.model.yLabel
        for i in range(self.model.predictors):
            self.rowLabels[i+1].value = self.model.predLabels[i]
            self.colLabels[i+1].value = self.model.predLabels[i]
        for i in range(self.model.num_variables):
            for j in range(i+1):
                self.grid[i, j+1].value = self.model.matrix[i][j]
                if self.model.standardized and i == j:
                    self.grid[i, j+1].disabled = True

    def observeMatrix(self, i, j, change):
        self.model.setMatrix(i, j, change['new'])

    def observeN(self, change):
        self.model.N = change['new']

# Cell
class RegressionLandingView(VBox):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.maxPredictors = 20
        self.maxLevels=6
        layout = {'description_width': 'initial', 'max_width': '130px', 'margin': '2px 2px 40px 2px'}
        self.numPredictors = BoundedIntText(description='Predictors', layout=ipyw.Layout(**layout), min=1, max=self.maxPredictors)

        self.yLabel = Text(placeholder='Y', layout=ipyw.Layout(width='90px'))
        options = ('Gaussian','Logistic', 'Poisson', 'Zero-inflated Poisson', 'Negative Binomial', 'Multinomial')
        first = HBox(children=(self.yLabel, ), layout=ipyw.Layout(margin='1px 15px 1px 15px'))

        options = ('Continuous', 'Categorical')
        self.varTypes = [Dropdown(options=options, layout=ipyw.Layout(width='110px')) for i in range(self.maxPredictors)]
        self.labels = [Text(placeholder='X' + str(i+1), layout=ipyw.Layout(width='90px')) for i in range(self.maxPredictors)]
        self.numLevels = [BoundedIntText(description='Levels', layout=ipyw.Layout(**layout), min=2, max=self.maxLevels) for i in range(self.maxPredictors)]
        self.levels = [[Text(placeholder=string.ascii_uppercase[i],layout=ipyw.Layout(width='75px')) for i in range(self.maxLevels)] for i in range(self.maxPredictors)]

        self.variables = []
        self.levelsBoxes = []
        for i in range(self.maxPredictors):
            topBox = HBox(layout=ipyw.Layout(margin='0 15px 0 15px', flex='0 1 auto'))
            topBox.children = (self.labels[i], self.varTypes[i])
            levelsBox = HBox(layout=ipyw.Layout(margin='15px 15px 0 15px', flex='0 1 auto'))
            levelsBox.children = tuple([self.numLevels[i]] + self.levels[i])
            self.levelsBoxes.append(levelsBox)
            self.variables.append(VBox(children=(topBox,levelsBox)))
            for j in range(2, self.maxLevels):
                self.levels[i][j].hide()
            levelsBox.hide()

        flexLayout= {'flex_flow': 'column wrap', 'width': '100%', 'height': '340px'}
        self.flexbox = Box(layout=flexLayout, children = tuple(self.variables))
        self.children += (self.numPredictors, first, self.flexbox)

# Cell
class RegressionLanding(RegressionLandingView, Controller):

    def __init__(self, model, **kwargs):
        super().__init__(**kwargs)
        self.model = model
        self.numPredictors.observe(self.observePredictors, 'value')
        for i in range(self.maxPredictors):
            self.labels[i].observe(partial(self.observeLabels, i), 'value')
            self.varTypes[i].observe(partial(self.observeVarType, i), 'value')
            self.numLevels[i].observe(partial(self.observeNumLevels, i), 'value')
            for j, intText in enumerate(self.levels[i]):
                intText.observe(partial(self.observeLevels, i, j), 'value')
        self.setLanding()

    def setLanding(self):
        self.numPredictors.value = self.model.predictors
        self.yLabel.value = self.model.yLabel
        for i, label in enumerate(self.model._predLabels):
            self.labels[i].value = label
        for i, varType in enumerate(self.model._varTypes):
            self.varTypes[i].value = varType

    def observePredictors(self, change):
        self.flexbox.children = tuple(self.variables[:change['new']])
        self.model.setPredictors(change['new'])
        if change['new'] > change['old']:
            for i in range(change['old'], change['new']):
                if self.varTypes[i] == 'Categorical':
                    self.model.setNumLevels(i, self.numLevels[i].value)
                    for j in range(self.numLevels[i].value):
                        self.model.setLevel(i, j, self.levels[i][j].value)

    def observeLabels(self, index, change):
        self.model.setLabel(index, change['new'])

    def observeLevels(self, i, j, change):
        self.model.setLevel(i, j, change['new'])

    def observeYLabel(self, change):
        self.model.setYLabel(change['new'])

    def observeNumLevels(self, index, change):
        self.model.setNumLevels(index, change['new'])
        for i in range(self.maxLevels):
            if i < change['new']:
                self.levels[index][i].show()
                self.model.levels[index][i] = self.levels[index][i].value
            else:
                self.levels[index][i].hide()

    def populateLevelsBox(self, index):
        self.numLevels[index].value = self.model.numLevels[index]
        for i in range(self.maxLevels):
            if i < self.model.numLevels[index]:
                self.levels[index][i].show()
                self.levels[index][i].value = self.model.levels[index][i]
            else:
                self.levels[index][i].hide()

    def observeVarType(self, index, change):
        self.model.setVarType(index, change['new'])
        if change['new'] == 'Categorical':
            self.levelsBoxes[index].show()
            self.populateLevelsBox(index)
        else:
            self.levelsBoxes[index].hide()

# Cell
class StandardizedRegressionLanding(RegressionLanding):

    def __init__(self, model, **kwargs):
        super().__init__(model, **kwargs)
        # perhaps not the most obvious place to do this, but it works
        self.model.standardized = True
        for dropdown in self.varTypes:
            dropdown.options = ['Continuous']

# Cell
class RegressionView(VBox):

    def __init__(self, num_levels=[1]*20, standardized=True,  **kwargs):
        super().__init__(**kwargs)

        # calculate the total number of boxes needed
        self.total = sum(num_levels)

        options = [('Standard Deviation', 'sd'), ('Variance', 'var')]
        style = {'description_width': '20px'}
        fvLayout = {'width': '90px', 'margin': '2px 20px 2px 2px'}
        sdLayout = {'max_width': '140px', 'margin': '2px 20px 2px 2px'}
        trdLayout = {'max_width': '200px', 'margin': '2px 20px 2px 2px'}
        lLayout = {'margin': '1px 20px 1px 1px', 'width': '140px'}
        lStyle = {'description_width': '10px'}
        sdStyle = {'description_width': 'initial'}
        fbLayout= {'flex_flow': 'row wrap', 'justify_content': 'flex-start', 'align_content': 'space-around','max_height':'400px', 'margin': '0 0 50px 0'}

        nBox = HBox(layout=ipyw.Layout(margin='0 24px 22px 0', flex='0 1 auto'))
        self.N = BoundedIntText(description='N', max=999999, style=sdStyle, layout=ipyw.Layout(max_width='100px'))
        nBox.children = (self.N, )

        # Y
        yBox = HBox(layout=ipyw.Layout(margin='0 24px 22px 0', flex='0 1 auto'))
        self.yLabel = AlignedLabel(value='Label', justification='flex-end', layout=lLayout, style=lStyle)
        self.yValues = [BoundedFloatText(step=0.01, description='Third', style=sdStyle, layout=ipyw.Layout(**trdLayout)) for i in range(2)]
        yBox.children = tuple([self.yLabel] + self.yValues)

        self.varNames = [AlignedLabel(value='A realy really long label', justification='flex-end', layout=lLayout, style=lStyle) for i in range(self.total)]
        if standardized:
            fixedValLabel = r'\beta'
        else:
            fixedValLabel = 'b'
        self.fixedVals = [BoundedFloatText(step=0.01, description=r'$$%s_{%s}$$' % (fixedValLabel, str(i+1)), style=style, layout=ipyw.Layout(**fvLayout)) for i in range(self.total)]
        self.thirdFloat = [BoundedFloatText(step=0.01, description='Third', style=sdStyle, layout=ipyw.Layout(**sdLayout)) for i in range(self.total)]

        self.equation =  [HTMLMath(value = str(i)) for i in range((self.total + 1) * 2 - 1)]

        self.flexbox = Box(layout=fbLayout)
        for i in range(self.total):
            hbox = HBox(layout=ipyw.Layout(margin='0 24px 0 0', width='435px'))
            hbox.children = (self.varNames[i], self.fixedVals[i], self.thirdFloat[i])
            self.flexbox.children += (hbox, )
        self.children = (nBox, yBox, self.flexbox, HBox(children = tuple(self.equation)))

# Cell
class Regression(RegressionView):

    def __init__(self, model, **kwargs):
        self.model = model
        super().__init__(num_levels = model.numLevels, standardized=model.standardized, **kwargs)

        viewIndex = 0
        for i in range(self.model.predictors):
            if self.model.numLevels[i] == 1:
                self.equation[viewIndex * 2 + 1].value = self.model.equation[i][0]
                self.equation[(viewIndex + 1) * 2].value = self.model.equation[i][1]
                self.varNames[viewIndex].value = self.model.labels[i]
                self.fixedVals[viewIndex].description = self.model.descriptions[viewIndex][0]
                self.fixedVals[viewIndex].value = self.model.fixedVals[i]
                self.fixedVals[viewIndex].observe(partial(self.observeFixedVals, i, viewIndex), 'value')
                if self.model.descriptions[viewIndex][1]:
                    self.thirdFloat[viewIndex].description = self.model.descriptions[viewIndex][1]
                    self.thirdFloat[viewIndex].min = 0
                    self.thirdFloat[viewIndex].max = 999999
                    self.thirdFloat[viewIndex].value = self.model.parameters[i+1][1]
                    self.thirdFloat[viewIndex].observe(partial(self.observeSD, i), 'value')
                else:
                    self.thirdFloat[viewIndex].hide()
                viewIndex += 1
            else: # if categorical
                for j in range(self.model.numLevels[i]):
                    self.equation[viewIndex * 2 + 1].value = self.model.equation[i][j][0]
                    self.equation[(viewIndex + 1) * 2].value = self.model.equation[i][j][1]
                    self.varNames[viewIndex].value = self.model.labels[i][j]
                    self.fixedVals[viewIndex].description = self.model.descriptions[viewIndex][0]
                    self.fixedVals[viewIndex].value = self.model.fixedVals[i][j]
                    self.fixedVals[viewIndex].observe(partial(self.observeFixedVals, (i, j), viewIndex), 'value')
                    if self.model.descriptions[viewIndex][1]:
                        self.thirdFloat[viewIndex].description = self.model.descriptions[viewIndex][1]
                        self.thirdFloat[viewIndex].min = 0
                        self.thirdFloat[viewIndex].max = 1
                        self.thirdFloat[viewIndex].value = self.model.parameters[i+1][j]
                        self.thirdFloat[viewIndex].observe(partial(self.observeProportion, i, j), 'value')
                    else:
                        self.thirdFloat[viewIndex].hide()
                    viewIndex += 1

        self.yLabel.value = self.model.yLabel
        for i, description in enumerate(self.model.yDescription):
            if description:
                self.yValues[i].description = self.model.yDescription[i]
                self.yValues[i].value = self.model.getYValue1()
            else:
                self.yValues[i].hide()
        self.yValues[0].observe(self.observeYValue1, 'value')
        self.yValues[1].observe(self.observeYValue2, 'value')

        self.equation[0].value = '$%s = $' % self.model.yLabel

        self.N.value = self.model.N
        self.N.observe(self.observeN, 'value')

    def observeFixedVals(self, i, viewIndex, change):
        self.model.setFixedVal(i, change['new'])
        if isinstance(i, tuple):
            j, k= i
            self.equation[viewIndex* 2 + 1].value = self.model.equation[j][k][0]
        else:
            self.equation[viewIndex* 2 + 1].value = self.model.equation[i][0]

    def observeProportion(self, i, j, change):
        self.model.setProportion(i, j, change['new'])

    def observeSD(self, i, change):
        self.model.setSD(i, change['new'])

    def observeYValue1(self, change):
        self.model.setYValue1(change['new'])

    def observeYValue2(self, change):
        self.model.setYValue2(change['new'])

    def observeN(self, change):
        self.model.N = change['new']