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

__all__ = ['CramerVLandingModel', 'AddStringsMixin', 'CramerVModel', 'TetraLandingModel', 'TetraModel',
           'CorrLandingModel', 'CorrModel', 'SpearmansRhoModel', 'PearsonModel', 'PBiModel']

# 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 RatioLandingModel, RatioEstimatesModel
from ..mvc import Model
from ..items import Collection
from ..traits import LinkAllMixin
from ..rpy import RMixins
from ..superpower.corr import CorrCramerV, CorrTetra, CorrSpearmansRho, CorrPBi, CorrPearson

# Cell
class CramerVLandingModel(RatioLandingModel):
    pass

# 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 CramerVModel(RatioEstimatesModel, 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 TetraLandingModel(RatioLandingModel):
    pass

# Cell
class TetraModel(CramerVModel, 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 CorrLandingModel(Model):

    def setLanding(self):
        pass

class CorrModel(CorrLandingModel):

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

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

class SpearmansRhoModel(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 PearsonModel(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 PBiModel(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)