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

__all__ = ['rpy2', 'AddStringsMixin', 'CorrCramerVModel', 'CorrTetraModel', 'CorrModel', 'CorrSpearmansRhoModel',
           'CorrPearsonModel', 'CorrPBiModel']

# Cell
import ipywidgets as ipyw
import traitlets
from traitlets import observe, Unicode, Bool, Int, Float
from traittypes import Array
import numpy as np
from fastcore.test import test_eq
from IPython.display import display, Javascript

# Cell
from .ratio import RatioModel
from ..mvc import Model
from ..items import Collection
from ..traits import LinkAllMixin
from ..superpower.corr import CorrCramerV, CorrTetra, CorrSpearmansRho, CorrPBi, CorrPearson
from ..log import HANDLERS
from ..rpy import RFunction, R_CONSOLE, RMixins
rpy2 = R_CONSOLE.rpy2

# Cell
class AddStringsMixin():
    ''' Like numpy.char.add but it works on dtype np.object'''

    def add_strings(self, first, second):
        if isinstance(first,str):
            first = [first] * len(second)
        if isinstance(second,str):
            second = [second] * len(first)
        if len(first) != len(second):
            raise ValueError('If adding arrays, they must have matching lengths')
        return np.array([[first[i] + second[j] for j in range(len(second))] for i in range(len(first))], dtype=np.object)

# Cell
class CorrCramerVModel(RatioModel, AddStringsMixin):

    pVal = Float()
    rcomp = Bool()
    mata = Array(dtype=np.dtype(int))
    matb = Array(dtype=np.dtype(int))
    matrix = Array(dtype=np.dtype(int))
    w = Float()
    hasColors = Bool()
    tail = Unicode()

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.cramerv = CorrCramerV()

    def setEstimates(self):
        super().setEstimates() #ratio model
        self.pVal = 0.05
        self.w = 0.3
        self.rcomp = False
        self.hasColors = False
        self.tail = 'two'

    def changeColor(self, i, j, color):
        self.colors[i][j] = color
        if np.all(np.array(self.colors) == 'none'):
            self.hasColors = False
        else:
            self.hasColors = True

    def runEffectSize(self):
        if self.hasColors:
            self.validateColors()
            kwargs = self.getTraits(['matrix', 'pVal', 'tail'])
            matrix = kwargs.pop('matrix')
            kwargs['n_row'], kwargs['n_col'] = matrix.shape
            kwargs['ns_cells'] = matrix.flatten()
            kwargs['alternative'] = kwargs.pop('tail')
            if kwargs['alternative'] == 'two':
                kwargs['alternative'] = 'two sided'
            return self.cramerv.run(**kwargs)

    def run(self):
        if self.hasColors:
            return self.runEffectSize()
        else:
            kwargs = self.getTraits(['ns_cells', 'pVal', 'tail'])
            kwargs['n_row'], kwargs['n_col'] = kwargs['ns_cells'].shape
            kwargs['ns_cells'] = kwargs['ns_cells'].flatten()
            kwargs['alternative'] = kwargs.pop('tail')
            if kwargs['alternative'] == 'two':
                kwargs['alternative'] = 'two sided'
            return self.cramerv.run(**kwargs)

    def validateColors(self):
        self.mata = []
        self.matb = []
        self.strings = []
        colors = np.array(self.colors)
        ns_cells = np.array(self.ns_cells)
        strings = self.add_strings(self.rowLevelNames, self.colLevelNames)
        first = None
        switched = False
        done = False
        i = 0
        if self.rcomp:
            length = self.n_col
            stack = np.vstack
        else:
            length = self.n_row
            stack = np.column_stack
        while (i < length):
            if self.rcomp:
                color_vec = colors[i]
                vector = ns_cells[i]
                string = strings[i]
            else:
                color_vec = colors[:,i]
                vector = ns_cells[:,i]
                string = strings[:,i]
            if not first:
                red = color_vec == 'red'
                green = color_vec == 'green'
                if np.any(red) or np.any(green):
                    if np.any(red):
                        first = 'red'
                        second = 'green'
                        self.pattern = red
                    elif np.any(green):
                        first = 'green'
                        second = 'red'
                        self.pattern = green
                    self.mata = vector[self.pattern]
                    self.strings.append(np.array2string(string[self.pattern], separator=',', max_line_width=10000, formatter={'object':lambda x: x}))
                    if not self.rcomp:
                        self.mata = np.transpose([self.mata])
            elif not switched: # if we haven't switched to the second color yet
                vec = color_vec == first
                if np.any(vec): # if this is still a vector with the first color
                    assert np.array_equal(vec, self.pattern), 'Selected self.colors should be the same for every selected %s' % string
                    self.mata = stack((self.mata, vector[self.pattern]))
                    self.strings.append(np.array2string(string[self.pattern], separator=',', max_line_width=10000, formatter={'object':lambda x: x}))
                else: # if this vector is no longer the first color
                    vec = color_vec == second
                    if not np.any(vec):
                        raise Exception('Both colors must be present.')
                    else:
                        assert np.array_equal(vec, self.pattern), 'Selected self.colors should be the same for every selected %s' % string
                        self.matb = vector[self.pattern]
                        self.strings.append(np.array2string(string[self.pattern], separator=',', max_line_width=10000, formatter={'object':lambda x: x}))
                        if not self.rcomp:
                            self.matb = np.transpose([self.matb])
                        switched = True
            elif not done:
                vec = color_vec == second
                if np.any(vec):
                    assert np.array_equal(vec, self.pattern), 'Selected self.colors should be the same for every selected %s' % string
                    self.matb = stack((self.matb, vector[self.pattern]))
                    self.strings.append(np.array2string(string[self.pattern], separator=',', max_line_width=10000, formatter={'object':lambda x: x}))
                else:
                    done = True
            else:
                if np.any(color_vec != 'none'):
                    raise Exception('Misplaced color selections.')
            i += 1
        self.matrix = stack((self.mata, self.matb))

# Cell
class CorrTetraModel(CorrCramerVModel, RMixins):

    tail = Unicode()
    pVal = Float()

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.function = CorrTetra()

    def setLanding(self):
        super().setLanding()
        self.tail = 'two'
        self.pVal = 0.05

    def run(self):
        #corrTetra(tail = "two", conttab = matrix(rep(10, 36), nrow = 6, byrow = TRUE), obs = TRUE, pVal = 0.05)
        kwargs = self.getTraits(['ns_cells', 'pVal', 'tail', 'n_row'])
        kwargs['conttab'] = self.matrix_func(kwargs.pop('ns_cells').flatten(), nrow = kwargs['n_row'], byrow = True)
        kwargs['obs'] = True
        if kwargs['tail'] != 'two':
            kwargs['tail'] = 'one'
        del kwargs['n_row']
        print(kwargs)
        return self.function.run(**kwargs)

# Cell
class CorrModel(Model):

    r = Float()
    N = Int()
    alpha = Float()
    tail = Unicode()

    def setLanding(self):
        pass

    def setEstimates(self):
        self.alpha = 0.05
        self.tail = 'two'
        self.r = 0.5
        self.N = 100

class CorrSpearmansRhoModel(CorrModel):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.function = CorrSpearmansRho()

    def run(self, **kwargs):
        kwargs = self.getTraits()
        kwargs['alternative'] = kwargs.pop('tail')
        if kwargs['alternative'] == 'two':
            kwargs['alternative'] = 'two sided'
        kwargs['pVal'] = kwargs.pop('alpha')
        return self.function.run(**kwargs)

# Cell
class CorrPearsonModel(CorrModel):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.function = CorrPearson()

    def run(self, **kwargs):
        kwargs = self.getTraits()
        kwargs['alternative'] = kwargs.pop('tail')
        if kwargs['alternative'] == 'two':
            kwargs['alternative'] = 'two sided'
        kwargs['pVal'] = kwargs.pop('alpha')
        return self.function.run(**kwargs)

# Cell
class CorrPBiModel(CorrModel):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.function = CorrPBi()

    def run(self, **kwargs):
        kwargs = self.getTraits()
        if kwargs['tail'] != 'two':
            kwargs['tail'] = 'one'
        kwargs['pVal'] = kwargs.pop('alpha')
        print(kwargs)
        return self.function.run(**kwargs)