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

__all__ = ['Ratios', 'RatioArray', 'RatioLandingModel', 'RatioEstimatesModel']

# Cell
import numpy as np
from traitlets import List, Unicode, HasTraits, TraitType, observe, Int, Float
from traittypes import Array

# Cell
from ..rpy import DTYPES
from ..mvc import Model
from ..superpower.chisq import ChiSqInd

# Cell
class Ratios(np.ndarray):

    '''    Note that this doesn't support view casting and new-from-template. It just
    allows us to have access
    '''
    def __new__(subtype, array=None, size=3, decimals=3, **kwargs):
        if array is not None:
            array = np.array(array, dtype=np.dtype(float))
            size = array.size
        else:
            size = size
        precision = 10**decimals
        fill_value = round(int(1/size * precision)/precision, decimals)
        total = int(fill_value * precision) * size
        if array is not None:
            arrayTotal = int(np.sum(array) * precision)
            if not total == arrayTotal:
                raise Exception('Array does not contain valid sum %.3f' % round(total/precision, decimals))
            obj = array.view(subtype)
        else:
            obj = np.full(size, fill_value, dtype=np.dtype(float)).view(subtype)
        obj.precision = precision
        obj.decimals = decimals
        obj.total = total
        obj.step = 1/obj.precision
        obj.n = size
        obj.j = 0
        obj.k = 0
        return obj

    def __array_finalize__(self, obj):
        if obj is None: return
        self.precision = getattr(obj, 'precision', None)
        self.decimals = getattr(obj, 'decimals', None)
        self.total = getattr(obj, 'total', None)
        self.step = getattr(obj, 'step', None)
        self.n = getattr(obj, 'n', None)
        self.j = getattr(obj, 'j', None)
        self.k = getattr(obj, 'k', None)

    def __setitem__(self, i, value):
        diff = int(round(value - self[i], 3) * self.precision)
        if diff > 0:
            self.inc(i, diff)
        elif diff < 0:
            self.dec(i, abs(diff))

    def __add__(self, b):
        print('add: %s'% str(b))

    def __iadd__(self, b):
        print('addi: %s'% str(b))

    def __sadd__(self, b):
        print('adds: %s'% str(b))

    def inc(self, i, diff):
        change = 0
        nums = np.copy(self) * self.precision
        nums = nums.astype(int)
        if diff + nums[i] > self.total:
            diff = self.total - nums[i]
        while diff > 0:
            if self.j != i:
                if nums[self.j] - 1 >= 0:
                    nums[self.j] = nums[self.j] - 1
                    nums[i] = nums[i] + 1
                    diff -= 1
                    change += 1
            self._inc_j()
        for i, val in enumerate(np.around(nums/self.precision, 3)):
            self.itemset(i, val)

    def dec(self, i, diff):
        nums = np.copy(self) * self.precision
        nums = nums.astype(int)
        if nums[i] - diff < 0:
            diff = nums[i]
        while diff > 0:
            if self.k != i:
                if nums[self.k] + 1 <= self.total:
                    nums[self.k] = nums[self.k] + 1
                    nums[i] = nums[i] - 1
                    diff -= 1
            self._inc_k()
        for i, val in enumerate(np.around(nums/self.precision, 3)):
            self.itemset(i, val)

    def _inc_j(self):
        self.j = (self.j + 1) % self.n

    def _inc_k(self):
        self.k = (self.k + 1) % self.n

# Cell
class RatioArray(Array):

    info_text = 'a Ratio object'

    def __init__(self, default_value=None, size=3, decimals=3, **kwargs):
        super().__init__(**kwargs)
        if default_value:
            self.default_value = default_value
        else:
            self.default_value = Ratios(size=size, decimals=decimals)

    def validate(self, obj, value):
        if isinstance(value, Ratios):
            return value
        self.error(obj, value)

# Cell
class RatioLandingModel(Model):

    '''landing'''
    rowMax = Int(6)
    colMax = Int(6)
    n_row = Int()
    n_col = Int()
    rowName = Unicode()
    colName = Unicode()
    rowLevels = List()
    colLevels = List()
    rowLevelNames = List()
    colLevelNames = List()

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

    ''' landing '''
    def setLanding(self):
        with self.hold_trait_notifications():
            self.n_row = 6
            self.n_col = 6
            self.rowName = 'Row'
            self.colName = 'Col'
            self.rowLevels = [str(i+1) for i in range(6)]
            self.colLevels = [str(i+1) for i in range(6)]

    @observe('n_row')
    def observe_n_row(self, change):
        self.rowLevelNames = [self.rowName + self.rowLevels[i] for i in range(change['new'])]

    @observe('n_col')
    def observe_n_col(self, change):
        self.colLevelNames = [self.colName + self.colLevels[i] for i in range(change['new'])]

    @observe('rowName')
    def observeRowName(self, change):
        self.rowLevelNames = [change['new'] + self.rowLevels[i] for i in range(self.n_row)]

    @observe('colName')
    def observeColName(self, change):
        self.colLevelNames = [change['new'] + self.colLevels[i] for i in range(self.n_col)]

    @observe('rowLevels')
    def observeRowLevels(self, change):
        self.rowLevelNames = [self.rowName + change['new'][i] for i in range(self.n_row)]

    @observe('colLevels')
    def observeColLevels(self, change):
        self.colLevelNames = [self.colName + change['new'][i] for i in range(self.n_col)]

    def changeRowLevel(self, i, value):
        self.rowLevels[i] = value
        if self.n > i:
            self.rowLevelNames[i] = self.name + value

# Cell
class RatioEstimatesModel(RatioLandingModel):

    decimals = Int(3)
    rowRatios = RatioArray()
    colRatios = RatioArray()
    total = Int()
    N = Int()
    colors = List()
    ns_cells = Array(dtype=np.dtype(int))
    percents = Array(dtype=np.dtype(float))
    rowTotals = Array(dtype=np.dtype(int))
    colTotals = Array(dtype=np.dtype(int))


    def setEstimates(self):
        with self.hold_trait_notifications():
            self.rowRatios = Ratios(size=self.n_row, decimals=self.decimals)
            self.colRatios = Ratios(size=self.n_col, decimals=self.decimals)
            self.colors = [['none'] * self.n_col for i in range(self.n_row)]
            self.total = self.n_row * self.n_col * 10
            percents = np.outer(self.rowRatios, self.colRatios)
            self.percents = np.around(percents, self.decimals)
            ns_cells = self.percents * self.total
            self.ns_cells = ns_cells.astype(np.dtype(int))
            self.rowTotals = np.sum(self.ns_cells, axis=1)
            self.colTotals = np.sum(self.ns_cells, axis=0)
            self.N = int(np.sum(self.ns_cells))

    @observe('total')
    def observeTotal(self, change):
        self.total = change['new']
        self.onChange()

    def incRowRatio(self, i, diff):
        self.rowRatios[i] += diff
        self.onRatioChange()

    def incColRatio(self, i, diff):
        self.colRatios[i] += diff
        self.onRatioChange()

    def onRatioChange(self):
        percents = np.outer(self.rowRatios, self.colRatios)
        self.percents = np.around(percents, self.decimals)
        self.onChange()

    def onChange(self):
        ns_cells = self.percents * self.total
        self.ns_cells = ns_cells.astype(np.dtype(int))
        self.rowTotals = np.sum(self.ns_cells, axis=1)
        self.colTotals = np.sum(self.ns_cells, axis=0)
        self.N = int(np.sum(self.ns_cells))