<commit_before>"""Correlation matrix View: show correlation matrix."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import numpy as np
import numpy.random as rdn
from galry import (Manager, DefaultPaintManager, PlotInteractionManager,
    Visual, PlotPaintManager,
    GalryWidget, QtGui, QtCore, QtOpenGL, enforce_dtype, RectanglesVisual,
    TextVisual, TextureVisual)
from galry.tools import hsv_to_rgb

from kwiklib.dataio.selection import get_indices
from kwiklib.dataio.tools import get_array
from kwiklib.utils.colors import COLORMAP
from kwiklib.utils import logger as log
from klustaviewa.views.common import HighlightManager, KlustaViewaBindings, KlustaView

# -----------------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------------
def colormap(x, col0=None, col1=None):
    """Colorize a 2D grayscale array.

    Arguments:
      * x:an NxM array with values in [0,1].
      * col0=None: a tuple (H, S, V) corresponding to color 0. By default, a
        rainbow color gradient is used.
      * col1=None: a tuple (H, S, V) corresponding to color 1.

    Returns:
      * y: an NxMx3 array with a rainbow color palette.

    """
    # record values to be removed
    removed = x == -1

    x[np.isnan(x)] = 0.
    x -= x.min()
    x *= (1. / x.max())
    # Set the maximum values above which the max color should be used.
    max = .1
    x = np.clip(x / max, 0., 1.)
    # Gamma correction. Doesn't look very good.
    # x = x ** .2

    shape = x.shape

    if col0 is None:
        col0 = (.67, .91, .65)
    if col1 is None:
        col1 = (0., 1., 1.)

    col0 = np.array(col0).reshape((1, 1, -1))
    col1 = np.array(col1).reshape((1, 1, -1))

    col0 = np.tile(col0, x.shape + (1,))
    col1 = np.tile(col1, x.shape + (1,))

    x = np.tile(x.reshape(shape + (1,)), (1, 1, 3))

    y = hsv_to_rgb(col0 + (col1 - col0) * x)

    # value of -1 = black
    y[removed,:] = 0
    # Remove diagonal.
    n = y.shape[0]
    y[xrange(n), xrange(n), :] = 0

    return y

# -----------------------------------------------------------------------------
# Data manager
# -----------------------------------------------------------------------------
class SimilarityMatrixDataManager(Manager):
    def set_data(self, similarity_matrix=None,
        cluster_colors_full=None,
        clusters_hidden=[],
        ):

        if similarity_matrix is None:
            similarity_matrix = np.zeros(0)
            cluster_colors_full = np.zeros(0)

        if similarity_matrix.size == 0:
            similarity_matrix = -np.ones((2, 2))
        elif similarity_matrix.shape[0] == 1:
            similarity_matrix = -np.ones((2, 2))

        self.similarity_matrix = similarity_matrix
        n = similarity_matrix.shape[0]

        self.texture = colormap(similarity_matrix)
        # similarity_matrix axes are originally (x, y) from the lower left corner
        # but when displayed, they are (i, j) from the upper left corner
        # so we need to transpose
        self.texture = np.swapaxes(self.texture, 0, 1)[::-1, :, :]

        # Hide some clusters.
        tex0 = self.texture.copy()
        for clu in clusters_hidden:
            # Inversion of axes in the y axis
            self.texture[- clu - 1, :, :] = tex0[- clu - 1, :, :] * .5
            self.texture[:, clu, :] = tex0[:, clu, :] * .5

        self.clusters_unique = get_indices(cluster_colors_full)
        self.cluster_colors = cluster_colors_full
        self.nclusters = len(self.clusters_unique)

# -----------------------------------------------------------------------------
# Visuals
# -----------------------------------------------------------------------------
class SimilarityMatrixPaintManager(PlotPaintManager):
    def initialize(self):
        self.add_visual(TextureVisual,
            texture=self.data_manager.texture,
            name='similarity_matrix')

        self.add_visual(RectanglesVisual, coordinates=np.zeros((2, 4)),
            color=(1., 1., 1., .5), autonormalizable=False, name='square')

        self.add_visual(TextVisual, text='0', name='clusterinfo', fontsize=16,
            background_transparent=False,
            posoffset=(150., -60.),
            color=(1., 1., 1., 1.),
            letter_spacing=350.,
            depth=-1,
            visible=False)

    def update(self):
        self.set_data(
            texture=self.data_manager.texture, visual='similarity_matrix')

# -----------------------------------------------------------------------------
# Interaction
# -----------------------------------------------------------------------------
class SimilarityMatrixInfoManager(Manager):
    def initialize(self):
        pass

    def get_closest_cluster(self, xd, yd):
        nclu = self.data_manager.nclusters

        cx = int((xd + 1) / 2. * nclu)
        cy = int((yd + 1) / 2. * nclu)

        cx_rel = np.clip(cx, 0, nclu - 1)
        cy_rel = np.clip(cy, 0, nclu - 1)

        return cx_rel, cy_rel

    def show_closest_cluster(self, xd, yd):

        if len(self.data_manager.clusters_unique) == 0:
            return

        cx_rel, cy_rel = self.get_closest_cluster(xd, yd)

        cx = self.data_manager.clusters_unique[cx_rel]
        cy = self.data_manager.clusters_unique[cy_rel]

        if ((cx_rel >= self.data_manager.similarity_matrix.shape[0]) or
            (cy_rel >= self.data_manager.similarity_matrix.shape[1])):
            return

        val = self.data_manager.similarity_matrix[cx_rel, cy_rel]

        text = "%d/%d:%.3f" % (cx, cy, val)

        x, y = self.interaction_manager.get_processor('navigation').get_window_coordinates(xd, yd)

        posoffset=(-np.sign(x) * 140., -np.sign(y) * 40.)

        self.paint_manager.set_data(coordinates=(xd, yd),
            text=text,
            visible=True,
            posoffset=posoffset,
            visual='clusterinfo')

class SimilarityMatrixInteractionManager(PlotInteractionManager):
    def initialize(self):
        # self.register('ShowClosestCluster', self.show_closest_cluster)
        self.register('SelectPair', self.select_pair)
        self.register('AddPair', self.add_pair)
        self.register('MoveSquare', self.move_square)
        self.register(None, self.hide_closest_cluster)

        self.square_coordinates = np.zeros((2, 4))
        self.clusters_selected = np.array([])

    def hide_closest_cluster(self, parameter):
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def select_pair(self, parameter, add=False):
        if self.data_manager.nclusters == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        cx_rel, cy_rel = self.info_manager.get_closest_cluster(xd, yd)
        cx = self.data_manager.clusters_unique[cx_rel]
        cy = self.data_manager.clusters_unique[cy_rel]
        if cx != cy:
            clusters = np.array([cx, cy])
        else:
            clusters = np.array([cx])

        if add:
            clusters = np.array(sorted(set(self.clusters_selected).union(
                clusters)))
            # clusters_new = clusters
            # clusters = self.clusters_selected.extend([cluster
                # for cluster in clusters_new
                    # if cluster not in self.clusters_selected])

        self.clusters_selected = clusters

        # Emit signal.
        log.debug("Selected clusters {0:s}.".format(str(clusters)))
        self.parent.clustersSelected.emit(clusters)

    def add_pair(self, parameter):
        self.select_pair(parameter, True)

    def show_cluster(self, clu0, clu1):
        if self.data_manager.nclusters <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_unique) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.info_manager.show_closest_cluster((x0 + x1) / 2, (y0 + y1) / 2)

    def show_closest_cluster(self, parameter):
        if self.data_manager.nclusters <= 1:
            return

        nclu = self.data_manager.nclusters

        if nclu == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        self.info_manager.show_closest_cluster(xd, yd)

    def show_square(self, clu0, clu1):
        if self.data_manager.nclusters <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_unique) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.square_coordinates[0,:] = (x0, -1, x1, 1)
        self.square_coordinates[1,:] = (-1, y0, 1, y1)
        self.paint_manager.set_data(coordinates=self.square_coordinates,
            visible=True,
            visual='square')

    def move_square(self, parameter):
        if self.data_manager.nclusters <= 1:
            return

        self.show_closest_cluster(parameter)

        # data coordinates
        x, y = parameter
        nav = self.get_processor('navigation')
        x, y = nav.get_data_coordinates(x, y)

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i = np.clip(int((x + 1) / 2. * n), 0, n - 1)
        j = np.clip(int((y + 1) / 2. * n), 0, n - 1)

        clu0 = self.data_manager.clusters_unique[i]
        clu1 = self.data_manager.clusters_unique[j]
        self.show_square(clu0, clu1)

class SimilarityMatrixBindings(KlustaViewaBindings):
    def get_base_cursor(self):
        return 'ArrowCursor'

    def set_selectcluster(self):
        self.set('RightClick', 'SelectPair', #key_modifier='Shift',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))
        self.set('RightClick', 'AddPair', key_modifier='Control',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))

    def set_move(self):
        self.set('Move', 'MoveSquare',
            param_getter=lambda p: p['mouse_position'])
        self.set('Move', 'MoveSquare', key_modifier='Control',
            param_getter=lambda p: p['mouse_position'])

    def initialize(self):
        # super(SimilarityMatrixBindings, self).initialize()
        # self.set_clusterinfo()
        self.set_selectcluster()
        self.set_move()

# -----------------------------------------------------------------------------
# Top-level module
# -----------------------------------------------------------------------------
class SimilarityMatrixView(KlustaView):

    # Raise the list of highlighted spike absolute indices.
    clustersSelected = QtCore.pyqtSignal(np.ndarray)

    def initialize(self):
        self.set_bindings(SimilarityMatrixBindings)
        self.set_companion_classes(
            paint_manager=SimilarityMatrixPaintManager,
            info_manager=SimilarityMatrixInfoManager,
            interaction_manager=SimilarityMatrixInteractionManager,
            data_manager=SimilarityMatrixDataManager,)

    def set_data(self, *args, **kwargs):
        # if kwargs.get('similarity_matrix', None) is None:
            # return
        self.data_manager.set_data(*args, **kwargs)

        # update?
        if self.initialized:
            self.paint_manager.update()
            self.updateGL()

    def clear(self):
        self.set_data()
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def show_selection(self, clu0, clu1):
        self.interaction_manager.show_square(clu0, clu1)
        self.interaction_manager.show_cluster(clu0, clu1)
        self.updateGL()

    # def sizeHint(self):
        # return QtCore.QSize(300, 400)
        <commit_msg>WIP: removing hidden clusters from similarity matrix. TODO: update interaction and text in matrix view.<commit_after>"""Correlation matrix View: show correlation matrix."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import numpy as np
import numpy.random as rdn
from galry import (Manager, DefaultPaintManager, PlotInteractionManager,
    Visual, PlotPaintManager,
    GalryWidget, QtGui, QtCore, QtOpenGL, enforce_dtype, RectanglesVisual,
    TextVisual, TextureVisual)
from galry.tools import hsv_to_rgb

from kwiklib.dataio.selection import get_indices
from kwiklib.dataio.tools import get_array
from kwiklib.utils.colors import COLORMAP
from kwiklib.utils import logger as log
from klustaviewa.views.common import HighlightManager, KlustaViewaBindings, KlustaView

# -----------------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------------
def colormap(x, col0=None, col1=None):
    """Colorize a 2D grayscale array.

    Arguments:
      * x:an NxM array with values in [0,1].
      * col0=None: a tuple (H, S, V) corresponding to color 0. By default, a
        rainbow color gradient is used.
      * col1=None: a tuple (H, S, V) corresponding to color 1.

    Returns:
      * y: an NxMx3 array with a rainbow color palette.

    """
    # record values to be removed
    removed = x == -1

    x[np.isnan(x)] = 0.
    x -= x.min()
    x *= (1. / x.max())
    # Set the maximum values above which the max color should be used.
    max = .1
    x = np.clip(x / max, 0., 1.)
    # Gamma correction. Doesn't look very good.
    # x = x ** .2

    shape = x.shape

    if col0 is None:
        col0 = (.67, .91, .65)
    if col1 is None:
        col1 = (0., 1., 1.)

    col0 = np.array(col0).reshape((1, 1, -1))
    col1 = np.array(col1).reshape((1, 1, -1))

    col0 = np.tile(col0, x.shape + (1,))
    col1 = np.tile(col1, x.shape + (1,))

    x = np.tile(x.reshape(shape + (1,)), (1, 1, 3))

    y = hsv_to_rgb(col0 + (col1 - col0) * x)

    # value of -1 = black
    y[removed,:] = 0
    # Remove diagonal.
    n = y.shape[0]
    y[xrange(n), xrange(n), :] = 0

    return y

# -----------------------------------------------------------------------------
# Data manager
# -----------------------------------------------------------------------------
class SimilarityMatrixDataManager(Manager):
    def set_data(self, similarity_matrix=None,
        cluster_colors_full=None,
        clusters_hidden=[],  # WARNING: relative indexing
        ):
        if similarity_matrix is None:
            similarity_matrix = np.zeros(0)
            cluster_colors_full = np.zeros(0)

        if similarity_matrix.size == 0:
            similarity_matrix = -np.ones((2, 2))
        elif similarity_matrix.shape[0] == 1:
            similarity_matrix = -np.ones((2, 2))

        self.similarity_matrix = similarity_matrix
        n = similarity_matrix.shape[0]

        self.texture = colormap(similarity_matrix)
        # similarity_matrix axes are originally (x, y) from the lower left corner
        # but when displayed, they are (i, j) from the upper left corner
        # so we need to transpose
        self.texture = np.swapaxes(self.texture, 0, 1)

        self.clusters_unique = get_indices(cluster_colors_full)

        self.cluster_colors = cluster_colors_full
        self.nclusters = len(self.clusters_unique)

        # Remove hidden clusters.
        indices = np.array(sorted(set(range(self.nclusters)) - set(clusters_hidden)),
                                dtype=np.int32)

        if len(indices >= 2):
            tex0 = self.texture.copy()
            k = self.nclusters
            self.texture = tex0[[[_] for _ in (indices)],indices,:]

        self.texture = self.texture[::-1,:,:]

# -----------------------------------------------------------------------------
# Visuals
# -----------------------------------------------------------------------------
class SimilarityMatrixPaintManager(PlotPaintManager):
    def initialize(self):
        self.add_visual(TextureVisual,
            texture=self.data_manager.texture,
            name='similarity_matrix')

        self.add_visual(RectanglesVisual, coordinates=np.zeros((2, 4)),
            color=(1., 1., 1., .5), autonormalizable=False, name='square')

        self.add_visual(TextVisual, text='0', name='clusterinfo', fontsize=16,
            background_transparent=False,
            posoffset=(150., -60.),
            color=(1., 1., 1., 1.),
            letter_spacing=350.,
            depth=-1,
            visible=False)

    def update(self):
        self.set_data(
            texture=self.data_manager.texture, visual='similarity_matrix')

# -----------------------------------------------------------------------------
# Interaction
# -----------------------------------------------------------------------------
class SimilarityMatrixInfoManager(Manager):
    def initialize(self):
        pass

    def get_closest_cluster(self, xd, yd):
        nclu = self.data_manager.nclusters

        cx = int((xd + 1) / 2. * nclu)
        cy = int((yd + 1) / 2. * nclu)

        cx_rel = np.clip(cx, 0, nclu - 1)
        cy_rel = np.clip(cy, 0, nclu - 1)

        return cx_rel, cy_rel

    def show_closest_cluster(self, xd, yd):

        if len(self.data_manager.clusters_unique) == 0:
            return

        cx_rel, cy_rel = self.get_closest_cluster(xd, yd)

        cx = self.data_manager.clusters_unique[cx_rel]
        cy = self.data_manager.clusters_unique[cy_rel]

        if ((cx_rel >= self.data_manager.similarity_matrix.shape[0]) or
            (cy_rel >= self.data_manager.similarity_matrix.shape[1])):
            return

        val = self.data_manager.similarity_matrix[cx_rel, cy_rel]

        text = "%d/%d:%.3f" % (cx, cy, val)

        x, y = self.interaction_manager.get_processor('navigation').get_window_coordinates(xd, yd)

        posoffset=(-np.sign(x) * 140., -np.sign(y) * 40.)

        self.paint_manager.set_data(coordinates=(xd, yd),
            text=text,
            visible=True,
            posoffset=posoffset,
            visual='clusterinfo')

class SimilarityMatrixInteractionManager(PlotInteractionManager):
    def initialize(self):
        # self.register('ShowClosestCluster', self.show_closest_cluster)
        self.register('SelectPair', self.select_pair)
        self.register('AddPair', self.add_pair)
        self.register('MoveSquare', self.move_square)
        self.register(None, self.hide_closest_cluster)

        self.square_coordinates = np.zeros((2, 4))
        self.clusters_selected = np.array([])

    def hide_closest_cluster(self, parameter):
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def select_pair(self, parameter, add=False):
        if self.data_manager.nclusters == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        cx_rel, cy_rel = self.info_manager.get_closest_cluster(xd, yd)
        cx = self.data_manager.clusters_unique[cx_rel]
        cy = self.data_manager.clusters_unique[cy_rel]
        if cx != cy:
            clusters = np.array([cx, cy])
        else:
            clusters = np.array([cx])

        if add:
            clusters = np.array(sorted(set(self.clusters_selected).union(
                clusters)))
            # clusters_new = clusters
            # clusters = self.clusters_selected.extend([cluster
                # for cluster in clusters_new
                    # if cluster not in self.clusters_selected])

        self.clusters_selected = clusters

        # Emit signal.
        log.debug("Selected clusters {0:s}.".format(str(clusters)))
        self.parent.clustersSelected.emit(clusters)

    def add_pair(self, parameter):
        self.select_pair(parameter, True)

    def show_cluster(self, clu0, clu1):
        if self.data_manager.nclusters <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_unique) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.info_manager.show_closest_cluster((x0 + x1) / 2, (y0 + y1) / 2)

    def show_closest_cluster(self, parameter):
        if self.data_manager.nclusters <= 1:
            return

        nclu = self.data_manager.nclusters

        if nclu == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        self.info_manager.show_closest_cluster(xd, yd)

    def show_square(self, clu0, clu1):
        if self.data_manager.nclusters <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_unique) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.square_coordinates[0,:] = (x0, -1, x1, 1)
        self.square_coordinates[1,:] = (-1, y0, 1, y1)
        self.paint_manager.set_data(coordinates=self.square_coordinates,
            visible=True,
            visual='square')

    def move_square(self, parameter):
        if self.data_manager.nclusters <= 1:
            return

        self.show_closest_cluster(parameter)

        # data coordinates
        x, y = parameter
        nav = self.get_processor('navigation')
        x, y = nav.get_data_coordinates(x, y)

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i = np.clip(int((x + 1) / 2. * n), 0, n - 1)
        j = np.clip(int((y + 1) / 2. * n), 0, n - 1)

        clu0 = self.data_manager.clusters_unique[i]
        clu1 = self.data_manager.clusters_unique[j]
        self.show_square(clu0, clu1)

class SimilarityMatrixBindings(KlustaViewaBindings):
    def get_base_cursor(self):
        return 'ArrowCursor'

    def set_selectcluster(self):
        self.set('RightClick', 'SelectPair', #key_modifier='Shift',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))
        self.set('RightClick', 'AddPair', key_modifier='Control',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))

    def set_move(self):
        self.set('Move', 'MoveSquare',
            param_getter=lambda p: p['mouse_position'])
        self.set('Move', 'MoveSquare', key_modifier='Control',
            param_getter=lambda p: p['mouse_position'])

    def initialize(self):
        # super(SimilarityMatrixBindings, self).initialize()
        # self.set_clusterinfo()
        self.set_selectcluster()
        self.set_move()

# -----------------------------------------------------------------------------
# Top-level module
# -----------------------------------------------------------------------------
class SimilarityMatrixView(KlustaView):

    # Raise the list of highlighted spike absolute indices.
    clustersSelected = QtCore.pyqtSignal(np.ndarray)

    def initialize(self):
        self.set_bindings(SimilarityMatrixBindings)
        self.set_companion_classes(
            paint_manager=SimilarityMatrixPaintManager,
            info_manager=SimilarityMatrixInfoManager,
            interaction_manager=SimilarityMatrixInteractionManager,
            data_manager=SimilarityMatrixDataManager,)

    def set_data(self, *args, **kwargs):
        # if kwargs.get('similarity_matrix', None) is None:
            # return
        self.data_manager.set_data(*args, **kwargs)

        # update?
        if self.initialized:
            self.paint_manager.update()
            self.updateGL()

    def clear(self):
        self.set_data()
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def show_selection(self, clu0, clu1):
        self.interaction_manager.show_square(clu0, clu1)
        self.interaction_manager.show_cluster(clu0, clu1)
        self.updateGL()

    # def sizeHint(self):
        # return QtCore.QSize(300, 400)
        <|endoftext|>


<commit_before>"""Correlation matrix View: show correlation matrix."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import numpy as np
import numpy.random as rdn
from galry import (Manager, DefaultPaintManager, PlotInteractionManager,
    Visual, PlotPaintManager,
    GalryWidget, QtGui, QtCore, QtOpenGL, enforce_dtype, RectanglesVisual,
    TextVisual, TextureVisual)
from galry.tools import hsv_to_rgb

from kwiklib.dataio.selection import get_indices
from kwiklib.dataio.tools import get_array
from kwiklib.utils.colors import COLORMAP
from kwiklib.utils import logger as log
from klustaviewa.views.common import HighlightManager, KlustaViewaBindings, KlustaView

# -----------------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------------
def colormap(x, col0=None, col1=None):
    """Colorize a 2D grayscale array.

    Arguments:
      * x:an NxM array with values in [0,1].
      * col0=None: a tuple (H, S, V) corresponding to color 0. By default, a
        rainbow color gradient is used.
      * col1=None: a tuple (H, S, V) corresponding to color 1.

    Returns:
      * y: an NxMx3 array with a rainbow color palette.

    """
    # record values to be removed
    removed = x == -1

    x[np.isnan(x)] = 0.
    x -= x.min()
    x *= (1. / x.max())
    # Set the maximum values above which the max color should be used.
    max = .1
    x = np.clip(x / max, 0., 1.)
    # Gamma correction. Doesn't look very good.
    # x = x ** .2

    shape = x.shape

    if col0 is None:
        col0 = (.67, .91, .65)
    if col1 is None:
        col1 = (0., 1., 1.)

    col0 = np.array(col0).reshape((1, 1, -1))
    col1 = np.array(col1).reshape((1, 1, -1))

    col0 = np.tile(col0, x.shape + (1,))
    col1 = np.tile(col1, x.shape + (1,))

    x = np.tile(x.reshape(shape + (1,)), (1, 1, 3))

    y = hsv_to_rgb(col0 + (col1 - col0) * x)

    # value of -1 = black
    y[removed,:] = 0
    # Remove diagonal.
    n = y.shape[0]
    y[xrange(n), xrange(n), :] = 0

    return y

# -----------------------------------------------------------------------------
# Data manager
# -----------------------------------------------------------------------------
class SimilarityMatrixDataManager(Manager):
    def set_data(self, similarity_matrix=None,
        cluster_colors_full=None,
        clusters_hidden=[],  # WARNING: relative indexing
        ):
        if similarity_matrix is None:
            similarity_matrix = np.zeros(0)
            cluster_colors_full = np.zeros(0)

        if similarity_matrix.size == 0:
            similarity_matrix = -np.ones((2, 2))
        elif similarity_matrix.shape[0] == 1:
            similarity_matrix = -np.ones((2, 2))

        self.similarity_matrix = similarity_matrix
        n = similarity_matrix.shape[0]

        self.texture = colormap(similarity_matrix)
        # similarity_matrix axes are originally (x, y) from the lower left corner
        # but when displayed, they are (i, j) from the upper left corner
        # so we need to transpose
        self.texture = np.swapaxes(self.texture, 0, 1)

        self.clusters_unique = get_indices(cluster_colors_full)

        self.cluster_colors = cluster_colors_full
        self.nclusters = len(self.clusters_unique)

        # Remove hidden clusters.
        indices = np.array(sorted(set(range(self.nclusters)) - set(clusters_hidden)),
                                dtype=np.int32)

        if len(indices) >= 2:
            tex0 = self.texture.copy()
            self.texture = tex0[[[_] for _ in (indices)],indices,:]
        else:
            indices = range(self.nclusters)

        self.clusters_displayed = self.clusters_unique[indices]
        self.nclusters_displayed = len(indices)

        self.texture = self.texture[::-1,:,:]

# -----------------------------------------------------------------------------
# Visuals
# -----------------------------------------------------------------------------
class SimilarityMatrixPaintManager(PlotPaintManager):
    def initialize(self):
        self.add_visual(TextureVisual,
            texture=self.data_manager.texture,
            name='similarity_matrix')

        self.add_visual(RectanglesVisual, coordinates=np.zeros((2, 4)),
            color=(1., 1., 1., .5), autonormalizable=False, name='square')

        self.add_visual(TextVisual, text='0', name='clusterinfo', fontsize=16,
            background_transparent=False,
            posoffset=(150., -60.),
            color=(1., 1., 1., 1.),
            letter_spacing=350.,
            depth=-1,
            visible=False)

    def update(self):
        self.set_data(
            texture=self.data_manager.texture, visual='similarity_matrix')

# -----------------------------------------------------------------------------
# Interaction
# -----------------------------------------------------------------------------
class SimilarityMatrixInfoManager(Manager):
    def initialize(self):
        pass

    def get_closest_cluster(self, xd, yd):
        nclu = self.data_manager.nclusters_displayed

        cx = int((xd + 1) / 2. * nclu)
        cy = int((yd + 1) / 2. * nclu)

        cx_rel = np.clip(cx, 0, nclu - 1)
        cy_rel = np.clip(cy, 0, nclu - 1)

        return cx_rel, cy_rel

    def show_closest_cluster(self, xd, yd):

        if len(self.data_manager.clusters_displayed) == 0:
            return

        cx_rel, cy_rel = self.get_closest_cluster(xd, yd)

        cx = self.data_manager.clusters_displayed[cx_rel]
        cy = self.data_manager.clusters_displayed[cy_rel]

        if ((cx_rel >= self.data_manager.similarity_matrix.shape[0]) or
            (cy_rel >= self.data_manager.similarity_matrix.shape[1])):
            return

        val = self.data_manager.similarity_matrix[cx_rel, cy_rel]

        text = "%d/%d:%.3f" % (cx, cy, val)

        x, y = self.interaction_manager.get_processor('navigation').get_window_coordinates(xd, yd)

        posoffset=(-np.sign(x) * 140., -np.sign(y) * 40.)

        self.paint_manager.set_data(coordinates=(xd, yd),
            text=text,
            visible=True,
            posoffset=posoffset,
            visual='clusterinfo')

class SimilarityMatrixInteractionManager(PlotInteractionManager):
    def initialize(self):
        # self.register('ShowClosestCluster', self.show_closest_cluster)
        self.register('SelectPair', self.select_pair)
        self.register('AddPair', self.add_pair)
        self.register('MoveSquare', self.move_square)
        self.register(None, self.hide_closest_cluster)

        self.square_coordinates = np.zeros((2, 4))
        self.clusters_selected = np.array([])

    def hide_closest_cluster(self, parameter):
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def select_pair(self, parameter, add=False):
        if self.data_manager.nclusters_displayed == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        cx_rel, cy_rel = self.info_manager.get_closest_cluster(xd, yd)
        cx = self.data_manager.clusters_displayed[cx_rel]
        cy = self.data_manager.clusters_displayed[cy_rel]
        if cx != cy:
            clusters = np.array([cx, cy])
        else:
            clusters = np.array([cx])

        if add:
            clusters = np.array(sorted(set(self.clusters_selected).union(
                clusters)))
            # clusters_new = clusters
            # clusters = self.clusters_selected.extend([cluster
                # for cluster in clusters_new
                    # if cluster not in self.clusters_selected])

        self.clusters_selected = clusters

        # Emit signal.
        log.debug("Selected clusters {0:s}.".format(str(clusters)))
        self.parent.clustersSelected.emit(clusters)

    def add_pair(self, parameter):
        self.select_pair(parameter, True)

    def show_cluster(self, clu0, clu1):
        if self.data_manager.nclusters_displayed <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_displayed) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.info_manager.show_closest_cluster((x0 + x1) / 2, (y0 + y1) / 2)

    def show_closest_cluster(self, parameter):
        if self.data_manager.nclusters_displayed <= 1:
            return

        nclu = self.data_manager.nclusters_displayed

        if nclu == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        self.info_manager.show_closest_cluster(xd, yd)

    def show_square(self, clu0, clu1):
        if self.data_manager.nclusters_displayed <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_displayed) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.square_coordinates[0,:] = (x0, -1, x1, 1)
        self.square_coordinates[1,:] = (-1, y0, 1, y1)
        self.paint_manager.set_data(coordinates=self.square_coordinates,
            visible=True,
            visual='square')

    def move_square(self, parameter):
        if self.data_manager.nclusters_displayed <= 1:
            return

        self.show_closest_cluster(parameter)

        # data coordinates
        x, y = parameter
        nav = self.get_processor('navigation')
        x, y = nav.get_data_coordinates(x, y)

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i = np.clip(int((x + 1) / 2. * n), 0, n - 1)
        j = np.clip(int((y + 1) / 2. * n), 0, n - 1)

        clu0 = self.data_manager.clusters_displayed[i]
        clu1 = self.data_manager.clusters_displayed[j]
        self.show_square(clu0, clu1)

class SimilarityMatrixBindings(KlustaViewaBindings):
    def get_base_cursor(self):
        return 'ArrowCursor'

    def set_selectcluster(self):
        self.set('RightClick', 'SelectPair', #key_modifier='Shift',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))
        self.set('RightClick', 'AddPair', key_modifier='Control',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))

    def set_move(self):
        self.set('Move', 'MoveSquare',
            param_getter=lambda p: p['mouse_position'])
        self.set('Move', 'MoveSquare', key_modifier='Control',
            param_getter=lambda p: p['mouse_position'])

    def initialize(self):
        # super(SimilarityMatrixBindings, self).initialize()
        # self.set_clusterinfo()
        self.set_selectcluster()
        self.set_move()

# -----------------------------------------------------------------------------
# Top-level module
# -----------------------------------------------------------------------------
class SimilarityMatrixView(KlustaView):

    # Raise the list of highlighted spike absolute indices.
    clustersSelected = QtCore.pyqtSignal(np.ndarray)

    def initialize(self):
        self.set_bindings(SimilarityMatrixBindings)
        self.set_companion_classes(
            paint_manager=SimilarityMatrixPaintManager,
            info_manager=SimilarityMatrixInfoManager,
            interaction_manager=SimilarityMatrixInteractionManager,
            data_manager=SimilarityMatrixDataManager,)

    def set_data(self, *args, **kwargs):
        # if kwargs.get('similarity_matrix', None) is None:
            # return
        self.data_manager.set_data(*args, **kwargs)

        # update?
        if self.initialized:
            self.paint_manager.update()
            self.updateGL()

    def clear(self):
        self.set_data()
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def show_selection(self, clu0, clu1):
        self.interaction_manager.show_square(clu0, clu1)
        self.interaction_manager.show_cluster(clu0, clu1)
        self.updateGL()

    # def sizeHint(self):
        # return QtCore.QSize(300, 400)
        <commit_msg>Fixed bug in displaying cluster in similarity view matrix (due to index discrepancy after removing hidden clusters).<commit_after>"""Correlation matrix View: show correlation matrix."""

# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import numpy as np
import numpy.random as rdn
from galry import (Manager, DefaultPaintManager, PlotInteractionManager,
    Visual, PlotPaintManager,
    GalryWidget, QtGui, QtCore, QtOpenGL, enforce_dtype, RectanglesVisual,
    TextVisual, TextureVisual)
from galry.tools import hsv_to_rgb

from kwiklib.dataio.selection import get_indices
from kwiklib.dataio.tools import get_array
from kwiklib.utils.colors import COLORMAP
from kwiklib.utils import logger as log
from klustaviewa.views.common import HighlightManager, KlustaViewaBindings, KlustaView

# -----------------------------------------------------------------------------
# Utility functions
# -----------------------------------------------------------------------------
def colormap(x, col0=None, col1=None):
    """Colorize a 2D grayscale array.

    Arguments:
      * x:an NxM array with values in [0,1].
      * col0=None: a tuple (H, S, V) corresponding to color 0. By default, a
        rainbow color gradient is used.
      * col1=None: a tuple (H, S, V) corresponding to color 1.

    Returns:
      * y: an NxMx3 array with a rainbow color palette.

    """
    # record values to be removed
    removed = x == -1

    x[np.isnan(x)] = 0.
    x -= x.min()
    x *= (1. / x.max())
    # Set the maximum values above which the max color should be used.
    max = .1
    x = np.clip(x / max, 0., 1.)
    # Gamma correction. Doesn't look very good.
    # x = x ** .2

    shape = x.shape

    if col0 is None:
        col0 = (.67, .91, .65)
    if col1 is None:
        col1 = (0., 1., 1.)

    col0 = np.array(col0).reshape((1, 1, -1))
    col1 = np.array(col1).reshape((1, 1, -1))

    col0 = np.tile(col0, x.shape + (1,))
    col1 = np.tile(col1, x.shape + (1,))

    x = np.tile(x.reshape(shape + (1,)), (1, 1, 3))

    y = hsv_to_rgb(col0 + (col1 - col0) * x)

    # value of -1 = black
    y[removed,:] = 0
    # Remove diagonal.
    n = y.shape[0]
    y[xrange(n), xrange(n), :] = 0

    return y

# -----------------------------------------------------------------------------
# Data manager
# -----------------------------------------------------------------------------
class SimilarityMatrixDataManager(Manager):
    def set_data(self, similarity_matrix=None,
        cluster_colors_full=None,
        clusters_hidden=[],  # WARNING: relative indexing
        ):
        if similarity_matrix is None:
            similarity_matrix = np.zeros(0)
            cluster_colors_full = np.zeros(0)

        if similarity_matrix.size == 0:
            similarity_matrix = -np.ones((2, 2))
        elif similarity_matrix.shape[0] == 1:
            similarity_matrix = -np.ones((2, 2))

        self.similarity_matrix = similarity_matrix
        n = similarity_matrix.shape[0]

        self.texture = colormap(similarity_matrix)
        # similarity_matrix axes are originally (x, y) from the lower left corner
        # but when displayed, they are (i, j) from the upper left corner
        # so we need to transpose
        self.texture = np.swapaxes(self.texture, 0, 1)

        self.clusters_unique = get_indices(cluster_colors_full)

        self.cluster_colors = cluster_colors_full
        self.nclusters = len(self.clusters_unique)

        # Remove hidden clusters.
        indices = np.array(sorted(set(range(self.nclusters)) - set(clusters_hidden)),
                                dtype=np.int32)
        self.indices = indices

        if len(indices) >= 2:
            tex0 = self.texture.copy()
            self.texture = tex0[[[_] for _ in (indices)],indices,:]
        else:
            indices = range(self.nclusters)

        self.clusters_displayed = self.clusters_unique[indices]
        self.nclusters_displayed = len(indices)

        self.texture = self.texture[::-1,:,:]

# -----------------------------------------------------------------------------
# Visuals
# -----------------------------------------------------------------------------
class SimilarityMatrixPaintManager(PlotPaintManager):
    def initialize(self):
        self.add_visual(TextureVisual,
            texture=self.data_manager.texture,
            name='similarity_matrix')

        self.add_visual(RectanglesVisual, coordinates=np.zeros((2, 4)),
            color=(1., 1., 1., .5), autonormalizable=False, name='square')

        self.add_visual(TextVisual, text='0', name='clusterinfo', fontsize=16,
            background_transparent=False,
            posoffset=(150., -60.),
            color=(1., 1., 1., 1.),
            letter_spacing=350.,
            depth=-1,
            visible=False)

    def update(self):
        self.set_data(
            texture=self.data_manager.texture, visual='similarity_matrix')

# -----------------------------------------------------------------------------
# Interaction
# -----------------------------------------------------------------------------
class SimilarityMatrixInfoManager(Manager):
    def initialize(self):
        pass

    def get_closest_cluster(self, xd, yd):
        nclu = self.data_manager.nclusters_displayed

        cx = int((xd + 1) / 2. * nclu)
        cy = int((yd + 1) / 2. * nclu)

        cx_rel = np.clip(cx, 0, nclu - 1)
        cy_rel = np.clip(cy, 0, nclu - 1)

        return cx_rel, cy_rel

    def show_closest_cluster(self, xd, yd):

        if len(self.data_manager.clusters_displayed) == 0:
            return

        cx_rel, cy_rel = self.get_closest_cluster(xd, yd)

        cx = self.data_manager.clusters_displayed[cx_rel]
        cy = self.data_manager.clusters_displayed[cy_rel]

        ind = self.data_manager.indices
        matrix = self.data_manager.similarity_matrix
        matrix = matrix[ind,:][:,ind]

        if ((cx_rel >= matrix.shape[0]) or
            (cy_rel >= matrix.shape[1])):
            return

        val = matrix[cx_rel, cy_rel]

        text = "%d/%d:%.3f" % (cx, cy, val)

        x, y = self.interaction_manager.get_processor('navigation').get_window_coordinates(xd, yd)

        posoffset=(-np.sign(x) * 140., -np.sign(y) * 40.)

        self.paint_manager.set_data(coordinates=(xd, yd),
            text=text,
            visible=True,
            posoffset=posoffset,
            visual='clusterinfo')

class SimilarityMatrixInteractionManager(PlotInteractionManager):
    def initialize(self):
        # self.register('ShowClosestCluster', self.show_closest_cluster)
        self.register('SelectPair', self.select_pair)
        self.register('AddPair', self.add_pair)
        self.register('MoveSquare', self.move_square)
        self.register(None, self.hide_closest_cluster)

        self.square_coordinates = np.zeros((2, 4))
        self.clusters_selected = np.array([])

    def hide_closest_cluster(self, parameter):
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def select_pair(self, parameter, add=False):
        if self.data_manager.nclusters_displayed == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        cx_rel, cy_rel = self.info_manager.get_closest_cluster(xd, yd)
        cx = self.data_manager.clusters_displayed[cx_rel]
        cy = self.data_manager.clusters_displayed[cy_rel]
        if cx != cy:
            clusters = np.array([cx, cy])
        else:
            clusters = np.array([cx])

        if add:
            clusters = np.array(sorted(set(self.clusters_selected).union(
                clusters)))
            # clusters_new = clusters
            # clusters = self.clusters_selected.extend([cluster
                # for cluster in clusters_new
                    # if cluster not in self.clusters_selected])

        self.clusters_selected = clusters

        # Emit signal.
        log.debug("Selected clusters {0:s}.".format(str(clusters)))
        self.parent.clustersSelected.emit(clusters)

    def add_pair(self, parameter):
        self.select_pair(parameter, True)

    def show_cluster(self, clu0, clu1):
        if self.data_manager.nclusters_displayed <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_displayed) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.info_manager.show_closest_cluster((x0 + x1) / 2, (y0 + y1) / 2)

    def show_closest_cluster(self, parameter):
        if self.data_manager.nclusters_displayed <= 1:
            return

        nclu = self.data_manager.nclusters_displayed

        if nclu == 0:
            return

        nav = self.get_processor('navigation')

        # window coordinates
        x, y = parameter
        # data coordinates
        xd, yd = nav.get_data_coordinates(x, y)

        self.info_manager.show_closest_cluster(xd, yd)

    def show_square(self, clu0, clu1):
        if self.data_manager.nclusters_displayed <= 1:
            return

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i, j = np.digitize([clu0, clu1], self.data_manager.clusters_displayed) - 1
        x0, y0 = i * dx * 2 - 1, j * dx * 2 - 1,
        x1, y1 = (i + 1) * dx * 2 - 1, (j + 1) * dx * 2 - 1
        self.square_coordinates[0,:] = (x0, -1, x1, 1)
        self.square_coordinates[1,:] = (-1, y0, 1, y1)
        self.paint_manager.set_data(coordinates=self.square_coordinates,
            visible=True,
            visual='square')

    def move_square(self, parameter):
        if self.data_manager.nclusters_displayed <= 1:
            return

        self.show_closest_cluster(parameter)

        # data coordinates
        x, y = parameter
        nav = self.get_processor('navigation')
        x, y = nav.get_data_coordinates(x, y)

        n = self.data_manager.texture.shape[0]
        dx = 1 / float(n)
        i = np.clip(int((x + 1) / 2. * n), 0, n - 1)
        j = np.clip(int((y + 1) / 2. * n), 0, n - 1)

        clu0 = self.data_manager.clusters_displayed[i]
        clu1 = self.data_manager.clusters_displayed[j]
        self.show_square(clu0, clu1)

class SimilarityMatrixBindings(KlustaViewaBindings):
    def get_base_cursor(self):
        return 'ArrowCursor'

    def set_selectcluster(self):
        self.set('RightClick', 'SelectPair', #key_modifier='Shift',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))
        self.set('RightClick', 'AddPair', key_modifier='Control',
            param_getter=lambda p:
            (p['mouse_position'][0], p['mouse_position'][1]))

    def set_move(self):
        self.set('Move', 'MoveSquare',
            param_getter=lambda p: p['mouse_position'])
        self.set('Move', 'MoveSquare', key_modifier='Control',
            param_getter=lambda p: p['mouse_position'])

    def initialize(self):
        # super(SimilarityMatrixBindings, self).initialize()
        # self.set_clusterinfo()
        self.set_selectcluster()
        self.set_move()

# -----------------------------------------------------------------------------
# Top-level module
# -----------------------------------------------------------------------------
class SimilarityMatrixView(KlustaView):

    # Raise the list of highlighted spike absolute indices.
    clustersSelected = QtCore.pyqtSignal(np.ndarray)

    def initialize(self):
        self.set_bindings(SimilarityMatrixBindings)
        self.set_companion_classes(
            paint_manager=SimilarityMatrixPaintManager,
            info_manager=SimilarityMatrixInfoManager,
            interaction_manager=SimilarityMatrixInteractionManager,
            data_manager=SimilarityMatrixDataManager,)

    def set_data(self, *args, **kwargs):
        # if kwargs.get('similarity_matrix', None) is None:
            # return
        self.data_manager.set_data(*args, **kwargs)

        # update?
        if self.initialized:
            self.paint_manager.update()
            self.updateGL()

    def clear(self):
        self.set_data()
        self.paint_manager.set_data(visible=False, visual='clusterinfo')
        self.paint_manager.set_data(visible=False, visual='square')

    def show_selection(self, clu0, clu1):
        self.interaction_manager.show_square(clu0, clu1)
        self.interaction_manager.show_cluster(clu0, clu1)
        self.updateGL()

    # def sizeHint(self):
        # return QtCore.QSize(300, 400)
        <|endoftext|>


<commit_before># Copyright (C) 2009 - TODAY Renato Lima - Akretion
# Copyright (C) 2019 - TODAY Raphaël Valyi - Akretion
# License AGPL-3 - See http://www.gnu.org/licenses/agpl-3.0.html

from lxml import etree

from odoo import api, fields, models

from odoo.addons.l10n_br_fiscal.constants.fiscal import FISCAL_OUT

INVOICE_TO_OPERATION = {
    'out_invoice': 'out',
    'in_invoice': 'in',
    'out_refund': 'in',
    'in_refund': 'out',
}

REFUND_TO_OPERATION = {
    'out_invoice': 'in',
    'in_invoice': 'out',
    'out_refund': 'out',
    'in_refund': 'in',
}

FISCAL_TYPE_REFUND = {
    'out': ['purchase_return', 'in_return'],
    'in': ['sale_return', 'out_return'],
}

SHADOWED_FIELDS = [
    'partner_id',
    'company_id',
    'date',
    'currency_id',
    'partner_shipping_id',
]

class AccountInvoice(models.Model):
    _name = 'account.invoice'
    _inherit = [_name, 'l10n_br_fiscal.document.mixin.methods']
    _inherits = {'l10n_br_fiscal.document': 'fiscal_document_id'}
    _order = 'date_invoice DESC, number DESC'

    # initial account.invoice inherits on fiscal.document that are
    # disable with active=False in their fiscal_document table.
    # To make these invoices still visible, we set active=True
    # in the invoice table.
    active = fields.Boolean(
        string='Active',
        default=True,
    )

    cnpj_cpf = fields.Char(
        string='CNPJ/CPF',
        related='partner_id.cnpj_cpf',
    )

    legal_name = fields.Char(
        string='Adapted Legal Name',
        related='partner_id.legal_name',
    )

    ie = fields.Char(
        string='Adapted State Tax Number',
        related='partner_id.inscr_est',
    )

    financial_move_line_ids = fields.Many2many(
        comodel_name='account.move.line',
        string='Financial Move Lines',
        store=True,
        compute='_compute_financial',
    )

    fiscal_number = fields.Char(
        string='Fiscal Number',
        related='fiscal_document_id.number',
    )

    # this default should be overwritten to False in a module pretending to
    # create fiscal documents from the invoices. But this default here
    # allows to install the l10n_br_account module without creating issues
    # with the existing Odoo invoice (demo or not).
    fiscal_document_id = fields.Many2one(
        comodel_name='l10n_br_fiscal.document',
        string='Fiscal Document',
        required=True,
        copy=False,
        ondelete='cascade',
    )

    @api.multi
    @api.depends('move_id.line_ids')
    def _compute_financial(self):
        for invoice in self:
            lines = invoice.move_id.line_ids.filtered(
                lambda l: l.account_id == invoice.account_id and
                l.account_id.internal_type in ('receivable', 'payable'))
            invoice.financial_move_line_ids = lines.sorted()

    @api.model
    def _shadowed_fields(self):
        """Returns the list of shadowed fields that are synced
        from the parent."""
        return SHADOWED_FIELDS

    @api.multi
    def _prepare_shadowed_fields_dict(self, default=False):
        self.ensure_one()
        vals = self._convert_to_write(self.read(self._shadowed_fields())[0])
        if default:  # in case you want to use new rather than write later
            return {'default_%s' % (k,): vals[k] for k in vals.keys()}
        return vals

    @api.model
    def fields_view_get(self, view_id=None, view_type="form",
                        toolbar=False, submenu=False):

        order_view = super().fields_view_get(
            view_id, view_type, toolbar, submenu
        )

        if view_type == 'form':
            view = self.env['ir.ui.view']

            sub_form_view = self.env['account.invoice.line'].fields_view_get(
                view_id=self.env.ref('l10n_br_account.invoice_line_form').id,
                view_type='form')['arch']

            sub_form_node = etree.fromstring(
                self.env['account.invoice.line'].fiscal_form_view(
                    sub_form_view))

            sub_arch, sub_fields = view.postprocess_and_fields(
                'account.invoice.line', sub_form_node, None)

            order_view['fields']['invoice_line_ids']['views']['form'] = {}

            order_view['fields']['invoice_line_ids']['views']['form'][
                'fields'] = sub_fields
            order_view['fields']['invoice_line_ids']['views']['form'][
                'arch'] = sub_arch

        return order_view

    @api.model
    def default_get(self, fields_list):
        defaults = super().default_get(fields_list)
        invoice_type = self.env.context.get('type', 'out_invoice')
        defaults['fiscal_operation_type'] = INVOICE_TO_OPERATION[invoice_type]
        if defaults['fiscal_operation_type'] == FISCAL_OUT:
            defaults['issuer'] = 'company'
        else:
            defaults['issuer'] = 'partner'
        return defaults

    @api.model
    def create(self, values):
        dummy_doc = self.env.ref('l10n_br_fiscal.fiscal_document_dummy')
        if not values.get('document_type_id'):
            values.update({'fiscal_document_id': dummy_doc.id})
        invoice = super().create(values)
        if invoice.fiscal_document_id != dummy_doc:
            shadowed_fiscal_vals = invoice._prepare_shadowed_fields_dict()
            invoice.fiscal_document_id.write(shadowed_fiscal_vals)
        return invoice

    @api.multi
    def write(self, values):
        dummy_doc = self.env.ref('l10n_br_fiscal.fiscal_document_dummy')
        result = super().write(values)
        for invoice in self:
            if invoice.fiscal_document_id != dummy_doc:
                shadowed_fiscal_vals = invoice._prepare_shadowed_fields_dict()
                invoice.fiscal_document_id.write(shadowed_fiscal_vals)
        return result

    @api.multi
    def unlink(self):
        """Allows delete a draft or cancelled invoices"""
        self.filtered(lambda i: i.state in ('draft', 'cancel')).write(
            {'move_name': False}
        )
        return super().unlink()

    @api.one
    @api.depends(
        'invoice_line_ids.price_total',
        'tax_line_ids.amount',
        'tax_line_ids.amount_rounding',
        'currency_id',
        'company_id',
        'date_invoice',
        'type')
    def _compute_amount(self):
        for inv_line in self.invoice_line_ids:
            if inv_line.cfop_id:
                if inv_line.cfop_id.finance_move:
                    self.amount_untaxed += inv_line.price_subtotal
                    self.amount_tax += inv_line.price_tax
            else:
                self.amount_untaxed += inv_line.price_subtotal
                self.amount_tax += inv_line.price_tax

        self.amount_total = self.amount_untaxed + self.amount_tax
        amount_total_company_signed = self.amount_total
        amount_untaxed_signed = self.amount_untaxed
        if (self.currency_id and self.company_id and
                self.currency_id != self.company_id.currency_id):
            currency_id = self.currency_id
            amount_total_company_signed = currency_id._convert(
                self.amount_total, self.company_id.currency_id,
                self.company_id, self.date_invoice or fields.Date.today())
            amount_untaxed_signed = currency_id._convert(
                self.amount_untaxed, self.company_id.currency_id,
                self.company_id, self.date_invoice or fields.Date.today())
        sign = self.type in ['in_refund', 'out_refund'] and -1 or 1
        self.amount_total_company_signed = amount_total_company_signed * sign
        self.amount_total_signed = self.amount_total * sign
        self.amount_untaxed_signed = amount_untaxed_signed * sign

    @api.model
    def invoice_line_move_line_get(self):
        move_lines_dict = super().invoice_line_move_line_get()
        new_mv_lines_dict = []
        for line in move_lines_dict:
            invoice_line = self.invoice_line_ids.browse(line.get('invl_id'))
            line['price'] = invoice_line.price_total
            if invoice_line.cfop_id:
                if invoice_line.cfop_id.finance_move:
                    new_mv_lines_dict.append(line)
            else:
                new_mv_lines_dict.append(line)

        return new_mv_lines_dict

    @api.model
    def tax_line_move_line_get(self):
        tax_lines_dict = super().tax_line_move_line_get()
        # new_tax_lines_dict = []
        # for tax in tax_lines_dict:
        #     new_tax_lines_dict.append(tax)
        #
        #     new_tax = tax.copy()
        #     new_tax['type'] = 'src'
        #
        #     new_tax_lines_dict.append(new_tax)
        return tax_lines_dict

    @api.multi
    def finalize_invoice_move_lines(self, move_lines):
        lines = super().finalize_invoice_move_lines(move_lines)
        financial_lines = [
            l for l in lines if l[2]['account_id'] == self.account_id.id]

        count = 1
        for l in financial_lines:
            if l[2]['debit'] or l[2]['credit']:
                l[2]['name'] = '{}/{}-{}'.format(
                    self.fiscal_number, count, len(financial_lines))
                count += 1
        return lines

    @api.multi
    def get_taxes_values(self):
        # uncomment these lines
        # dummy_doc = self.env.ref('l10n_br_fiscal.fiscal_document_dummy')
        # if self.fiscal_document_id == dummy_doc:
        #     return super().get_taxes_values()

        tax_grouped = {}
        round_curr = self.currency_id.round
        for line in self.invoice_line_ids:
            if not line.account_id or line.display_type:
                continue

            taxes = line.invoice_line_tax_ids.compute_all(
                price_unit=line.price_unit,
                currency=line.invoice_id.currency_id,
                quantity=line.quantity,
                product=line.product_id,
                partner=line.invoice_id.partner_id,
                fiscal_taxes=line.fiscal_tax_ids,
                operation_line=line.fiscal_operation_line_id,
                ncm=line.ncm_id,
                nbs=line.nbs_id,
                nbm=line.nbm_id,
                cest=line.cest_id,
                discount_value=line.discount_value,
                insurance_value=line.insurance_value,
                other_costs_value=line.other_costs_value,
                freight_value=line.freight_value,
                fiscal_price=line.fiscal_price,
                fiscal_quantity=line.fiscal_quantity,
                uot=line.uot_id,
                icmssn_range=line.icmssn_range_id)['taxes']

            for tax in taxes:
                if tax.get('amount', 0.0) != 0.0:
                    val = self._prepare_tax_line_vals(line, tax)
                    key = self.env['account.tax'].browse(
                        tax['id']).get_grouping_key(val)

                    if key not in tax_grouped:
                        tax_grouped[key] = val
                        tax_grouped[key]['base'] = round_curr(val['base'])
                    else:
                        tax_grouped[key]['amount'] += val['amount']
                        tax_grouped[key]['base'] += round_curr(val['base'])
        return tax_grouped

    @api.multi
    def action_move_create(self):
        for inv in self.filtred(lambda i: i.fiscal._document_id !=
                self.env.ref('l10n_br_fiscal.fiscal_document_dummy')):
            inv.fiscal_document_id.action_button_confirm()
        return super().action_move_create()
<commit_msg>[FIX] Invoice's total amount compute<commit_after># Copyright (C) 2009 - TODAY Renato Lima - Akretion
# Copyright (C) 2019 - TODAY Raphaël Valyi - Akretion
# License AGPL-3 - See http://www.gnu.org/licenses/agpl-3.0.html

from lxml import etree

from odoo import api, fields, models

from odoo.addons.l10n_br_fiscal.constants.fiscal import FISCAL_OUT

INVOICE_TO_OPERATION = {
    'out_invoice': 'out',
    'in_invoice': 'in',
    'out_refund': 'in',
    'in_refund': 'out',
}

REFUND_TO_OPERATION = {
    'out_invoice': 'in',
    'in_invoice': 'out',
    'out_refund': 'out',
    'in_refund': 'in',
}

FISCAL_TYPE_REFUND = {
    'out': ['purchase_return', 'in_return'],
    'in': ['sale_return', 'out_return'],
}

SHADOWED_FIELDS = [
    'partner_id',
    'company_id',
    'date',
    'currency_id',
    'partner_shipping_id',
]

class AccountInvoice(models.Model):
    _name = 'account.invoice'
    _inherit = [_name, 'l10n_br_fiscal.document.mixin.methods']
    _inherits = {'l10n_br_fiscal.document': 'fiscal_document_id'}
    _order = 'date_invoice DESC, number DESC'

    # initial account.invoice inherits on fiscal.document that are
    # disable with active=False in their fiscal_document table.
    # To make these invoices still visible, we set active=True
    # in the invoice table.
    active = fields.Boolean(
        string='Active',
        default=True,
    )

    cnpj_cpf = fields.Char(
        string='CNPJ/CPF',
        related='partner_id.cnpj_cpf',
    )

    legal_name = fields.Char(
        string='Adapted Legal Name',
        related='partner_id.legal_name',
    )

    ie = fields.Char(
        string='Adapted State Tax Number',
        related='partner_id.inscr_est',
    )

    financial_move_line_ids = fields.Many2many(
        comodel_name='account.move.line',
        string='Financial Move Lines',
        store=True,
        compute='_compute_financial',
    )

    fiscal_number = fields.Char(
        string='Fiscal Number',
        related='fiscal_document_id.number',
    )

    # this default should be overwritten to False in a module pretending to
    # create fiscal documents from the invoices. But this default here
    # allows to install the l10n_br_account module without creating issues
    # with the existing Odoo invoice (demo or not).
    fiscal_document_id = fields.Many2one(
        comodel_name='l10n_br_fiscal.document',
        string='Fiscal Document',
        required=True,
        copy=False,
        ondelete='cascade',
    )

    @api.multi
    @api.depends('move_id.line_ids')
    def _compute_financial(self):
        for invoice in self:
            lines = invoice.move_id.line_ids.filtered(
                lambda l: l.account_id == invoice.account_id and
                l.account_id.internal_type in ('receivable', 'payable'))
            invoice.financial_move_line_ids = lines.sorted()

    @api.model
    def _shadowed_fields(self):
        """Returns the list of shadowed fields that are synced
        from the parent."""
        return SHADOWED_FIELDS

    @api.multi
    def _prepare_shadowed_fields_dict(self, default=False):
        self.ensure_one()
        vals = self._convert_to_write(self.read(self._shadowed_fields())[0])
        if default:  # in case you want to use new rather than write later
            return {'default_%s' % (k,): vals[k] for k in vals.keys()}
        return vals

    @api.model
    def fields_view_get(self, view_id=None, view_type="form",
                        toolbar=False, submenu=False):

        order_view = super().fields_view_get(
            view_id, view_type, toolbar, submenu
        )

        if view_type == 'form':
            view = self.env['ir.ui.view']

            sub_form_view = self.env['account.invoice.line'].fields_view_get(
                view_id=self.env.ref('l10n_br_account.invoice_line_form').id,
                view_type='form')['arch']

            sub_form_node = etree.fromstring(
                self.env['account.invoice.line'].fiscal_form_view(
                    sub_form_view))

            sub_arch, sub_fields = view.postprocess_and_fields(
                'account.invoice.line', sub_form_node, None)

            order_view['fields']['invoice_line_ids']['views']['form'] = {}

            order_view['fields']['invoice_line_ids']['views']['form'][
                'fields'] = sub_fields
            order_view['fields']['invoice_line_ids']['views']['form'][
                'arch'] = sub_arch

        return order_view

    @api.model
    def default_get(self, fields_list):
        defaults = super().default_get(fields_list)
        invoice_type = self.env.context.get('type', 'out_invoice')
        defaults['fiscal_operation_type'] = INVOICE_TO_OPERATION[invoice_type]
        if defaults['fiscal_operation_type'] == FISCAL_OUT:
            defaults['issuer'] = 'company'
        else:
            defaults['issuer'] = 'partner'
        return defaults

    @api.model
    def create(self, values):
        dummy_doc = self.env.ref('l10n_br_fiscal.fiscal_document_dummy')
        if not values.get('document_type_id'):
            values.update({'fiscal_document_id': dummy_doc.id})
        invoice = super().create(values)
        if invoice.fiscal_document_id != dummy_doc:
            shadowed_fiscal_vals = invoice._prepare_shadowed_fields_dict()
            invoice.fiscal_document_id.write(shadowed_fiscal_vals)
        return invoice

    @api.multi
    def write(self, values):
        dummy_doc = self.env.ref('l10n_br_fiscal.fiscal_document_dummy')
        result = super().write(values)
        for invoice in self:
            if invoice.fiscal_document_id != dummy_doc:
                shadowed_fiscal_vals = invoice._prepare_shadowed_fields_dict()
                invoice.fiscal_document_id.write(shadowed_fiscal_vals)
        return result

    @api.multi
    def unlink(self):
        """Allows delete a draft or cancelled invoices"""
        self.filtered(lambda i: i.state in ('draft', 'cancel')).write(
            {'move_name': False}
        )
        return super().unlink()

    @api.one
    @api.depends(
        'invoice_line_ids.price_total',
        'tax_line_ids.amount',
        'tax_line_ids.amount_rounding',
        'currency_id',
        'company_id',
        'date_invoice',
        'type')
    def _compute_amount(self):
        self.fiscal_document_id._compute_amount()
        for inv_line in self.invoice_line_ids:
            if inv_line.cfop_id:
                if inv_line.cfop_id.finance_move:
                    self.amount_untaxed += inv_line.price_subtotal
                    self.amount_tax += inv_line.price_tax
            else:
                self.amount_untaxed += inv_line.price_subtotal
                self.amount_tax += inv_line.price_tax

        self.amount_total = (
            self.amount_untaxed + self.amount_tax -
            self.amount_tax_withholding)

        self.amount_total = self.amount_untaxed + self.amount_tax
        amount_total_company_signed = self.amount_total
        amount_untaxed_signed = self.amount_untaxed
        if (self.currency_id and self.company_id and
                self.currency_id != self.company_id.currency_id):
            currency_id = self.currency_id
            amount_total_company_signed = currency_id._convert(
                self.amount_total, self.company_id.currency_id,
                self.company_id, self.date_invoice or fields.Date.today())
            amount_untaxed_signed = currency_id._convert(
                self.amount_untaxed, self.company_id.currency_id,
                self.company_id, self.date_invoice or fields.Date.today())
        sign = self.type in ['in_refund', 'out_refund'] and -1 or 1
        self.amount_total_company_signed = amount_total_company_signed * sign
        self.amount_total_signed = self.amount_total * sign
        self.amount_untaxed_signed = amount_untaxed_signed * sign

    @api.model
    def invoice_line_move_line_get(self):
        move_lines_dict = super().invoice_line_move_line_get()
        new_mv_lines_dict = []
        for line in move_lines_dict:
            invoice_line = self.invoice_line_ids.browse(line.get('invl_id'))
            line['price'] = invoice_line.price_total
            if invoice_line.cfop_id:
                if invoice_line.cfop_id.finance_move:
                    new_mv_lines_dict.append(line)
            else:
                new_mv_lines_dict.append(line)

        return new_mv_lines_dict

    @api.model
    def tax_line_move_line_get(self):
        tax_lines_dict = super().tax_line_move_line_get()
        # new_tax_lines_dict = []
        # for tax in tax_lines_dict:
        #     new_tax_lines_dict.append(tax)
        #
        #     new_tax = tax.copy()
        #     new_tax['type'] = 'src'
        #
        #     new_tax_lines_dict.append(new_tax)
        return tax_lines_dict

    @api.multi
    def finalize_invoice_move_lines(self, move_lines):
        lines = super().finalize_invoice_move_lines(move_lines)
        financial_lines = [
            l for l in lines if l[2]['account_id'] == self.account_id.id]

        count = 1
        for l in financial_lines:
            if l[2]['debit'] or l[2]['credit']:
                l[2]['name'] = '{}/{}-{}'.format(
                    self.fiscal_number, count, len(financial_lines))
                count += 1
        return lines

    @api.multi
    def get_taxes_values(self):
        # uncomment these lines
        # dummy_doc = self.env.ref('l10n_br_fiscal.fiscal_document_dummy')
        # if self.fiscal_document_id == dummy_doc:
        #     return super().get_taxes_values()

        tax_grouped = {}
        round_curr = self.currency_id.round
        for line in self.invoice_line_ids:
            if not line.account_id or line.display_type:
                continue

            taxes = line.invoice_line_tax_ids.compute_all(
                price_unit=line.price_unit,
                currency=line.invoice_id.currency_id,
                quantity=line.quantity,
                product=line.product_id,
                partner=line.invoice_id.partner_id,
                fiscal_taxes=line.fiscal_tax_ids,
                operation_line=line.fiscal_operation_line_id,
                ncm=line.ncm_id,
                nbs=line.nbs_id,
                nbm=line.nbm_id,
                cest=line.cest_id,
                discount_value=line.discount_value,
                insurance_value=line.insurance_value,
                other_costs_value=line.other_costs_value,
                freight_value=line.freight_value,
                fiscal_price=line.fiscal_price,
                fiscal_quantity=line.fiscal_quantity,
                uot=line.uot_id,
                icmssn_range=line.icmssn_range_id)['taxes']

            for tax in taxes:
                if tax.get('amount', 0.0) != 0.0:
                    val = self._prepare_tax_line_vals(line, tax)
                    key = self.env['account.tax'].browse(
                        tax['id']).get_grouping_key(val)

                    if key not in tax_grouped:
                        tax_grouped[key] = val
                        tax_grouped[key]['base'] = round_curr(val['base'])
                    else:
                        tax_grouped[key]['amount'] += val['amount']
                        tax_grouped[key]['base'] += round_curr(val['base'])
        return tax_grouped

    @api.multi
    def action_move_create(self):
        for inv in self.filtred(lambda i: i.fiscal._document_id !=
                self.env.ref('l10n_br_fiscal.fiscal_document_dummy')):
            inv.fiscal_document_id.action_button_confirm()
        return super().action_move_create()
<|endoftext|>


<commit_before>"""Implementation of trie."""

from collections import OrderedDict

class Trie(object):
    """Class representation of trie.

    insert(string): Insert a new word into the trie.
    contains(self, string): Return true if the string is in trie else false.
    size(self): Return the size of the trie.
    remove(self, string): Remove the string from the tree if it exists.
    """

    def __init__(self):
        self.root = OrderedDict()
        self._size = 0

    def insert(self, string):
        """Insert a new word into the trie."""
        if not isinstance(string, str):
            raise TypeError("Value input must be a string")
        if self.contains(string):
            return
        check = self.root
        for letter in string:
            if letter not in check:
                check.setdefault(letter, OrderedDict())
            check = check[letter]
        check['$'] = 'END'
        self._size += 1

    def contains(self, string):
        """Return true if the string is in trie else false."""
        check = self.root
        for letter in string:
            if letter not in check:
                return False
            check = check[letter]
        return "$" in check.keys()

    @property
    def size(self):
        """Return the size of the trie."""
        return self._size

    def remove(self, string):
        """Remove the string from the tree if it exists."""
        if not self.contains(string):
            raise ValueError('The trie does not contain this string.')
        check = self.root
        list_deleted = []
        prev = None
        for letter in string:
            check = check[letter]
            if prev and len(prev.keys()) > 1:
                list_deleted.insert(0, letter)
            prev = check
        del check['$']
        if len(check.keys()) > 0:
            self._size -= 1
            return
        check = self.root
        prev = self.root
        for letter in string:
            if len(check[letter].keys()) > 1:
                prev = check[letter]
            check = check[letter]
        if len(list_deleted) > 0:
            del prev[list_deleted[0]]
        else:
            del self.root[string[0]]
        self._size -= 1

    def _traversal_start(self, start=None):
        """Return the node in trie corresponding to the input string if exists."""
        if start:
            try:
                check = self.root
                for letter in start:
                    check = check[letter]
            except KeyError:
                raise KeyError('The input string is not in the trie.')
            return check
        return self.root

    def traversal(self, start=None):
        """Perform a depth traversal of the string."""
        check = self._traversal_start(start)
        if isinstance(check, OrderedDict):
            for letter in check.keys():
                if letter != "$":
                    yield letter
                if start:
                    newstart = start + letter
                else:
                    newstart = letter
                for item in self.traversal(newstart):
                    if item != "$":
                        yield item
trie = Trie()
trie.insert('alpha')
trie.insert('alpaca')
trie.insert('boy')
gen = trie.traversal('alp')
print (list(gen))<commit_msg>remove some garbage<commit_after>"""Implementation of trie."""

from collections import OrderedDict

class Trie(object):
    """Class representation of trie.

    insert(string): Insert a new word into the trie.
    contains(self, string): Return true if the string is in trie else false.
    size(self): Return the size of the trie.
    remove(self, string): Remove the string from the tree if it exists.
    """

    def __init__(self):
        self.root = OrderedDict()
        self._size = 0

    def insert(self, string):
        """Insert a new word into the trie."""
        if not isinstance(string, str):
            raise TypeError("Value input must be a string")
        if self.contains(string):
            return
        check = self.root
        for letter in string:
            if letter not in check:
                check.setdefault(letter, OrderedDict())
            check = check[letter]
        check['$'] = 'END'
        self._size += 1

    def contains(self, string):
        """Return true if the string is in trie else false."""
        check = self.root
        for letter in string:
            if letter not in check:
                return False
            check = check[letter]
        return "$" in check.keys()

    @property
    def size(self):
        """Return the size of the trie."""
        return self._size

    def remove(self, string):
        """Remove the string from the tree if it exists."""
        if not self.contains(string):
            raise ValueError('The trie does not contain this string.')
        check = self.root
        list_deleted = []
        prev = None
        for letter in string:
            check = check[letter]
            if prev and len(prev.keys()) > 1:
                list_deleted.insert(0, letter)
            prev = check
        del check['$']
        if len(check.keys()) > 0:
            self._size -= 1
            return
        check = self.root
        prev = self.root
        for letter in string:
            if len(check[letter].keys()) > 1:
                prev = check[letter]
            check = check[letter]
        if len(list_deleted) > 0:
            del prev[list_deleted[0]]
        else:
            del self.root[string[0]]
        self._size -= 1

    def _traversal_start(self, start=None):
        """Return the node in trie corresponding to the input string if exists."""
        if start:
            try:
                check = self.root
                for letter in start:
                    check = check[letter]
            except KeyError:
                raise KeyError('The input string is not in the trie.')
            return check
        return self.root

    def traversal(self, start=None):
        """Perform a depth traversal of the string."""
        check = self._traversal_start(start)
        if isinstance(check, OrderedDict):
            for letter in check.keys():
                if letter != "$":
                    yield letter
                if start:
                    newstart = start + letter
                else:
                    newstart = letter
                for item in self.traversal(newstart):
                    if item != "$":
                        yield item
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('countries_plus', '__first__'),
    ]

    operations = [
        migrations.CreateModel(
            name='CultureCode',
            fields=[
                ('code', models.CharField(max_length=10, serialize=False, primary_key=True)),
                ('country', models.ForeignKey(to='countries_plus.Country')),
            ],
            options={
                'ordering': ['code'],
                'verbose_name': 'CultureCode',
                'verbose_name_plural': 'CultureCodes',
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Language',
            fields=[
                ('iso_639_1', models.CharField(max_length=2, serialize=False, primary_key=True)),
                ('iso_639_2T', models.CharField(unique=True, max_length=3, blank=True)),
                ('iso_639_2B', models.CharField(unique=True, max_length=3, blank=True)),
                ('iso_639_3', models.CharField(max_length=3, blank=True)),
                ('iso_639_6', models.CharField(max_length=4, blank=True)),
                ('name_en', models.CharField(max_length=100)),
                ('name_native', models.CharField(max_length=100)),
                ('name_other', models.CharField(max_length=50, blank=True)),
                ('family', models.CharField(max_length=50)),
                ('notes', models.CharField(max_length=100, blank=True)),
                ('countries_spoken', models.ManyToManyField(to='countries_plus.Country', null=True, blank=True)),
            ],
            options={
                'ordering': ['name_en'],
                'verbose_name': 'Language',
                'verbose_name_plural': 'Languages',
            },
            bases=(models.Model,),
        ),
        migrations.AddField(
            model_name='culturecode',
            name='language',
            field=models.ForeignKey(to='languages_plus.Language'),
            preserve_default=True,
        ),
        migrations.CreateModel(
            name='LangCountry',
            fields=[
            ],
            options={
                'proxy': True,
            },
            bases=('countries_plus.country',),
        ),
    ]
<commit_msg>Adding `on_delete` to migration<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('countries_plus', '__first__'),
    ]

    operations = [
        migrations.CreateModel(
            name='CultureCode',
            fields=[
                ('code', models.CharField(max_length=10, serialize=False, primary_key=True)),
                ('country', models.ForeignKey(to='countries_plus.Country', on_delete=models.CASCADE)),
            ],
            options={
                'ordering': ['code'],
                'verbose_name': 'CultureCode',
                'verbose_name_plural': 'CultureCodes',
            },
            bases=(models.Model,),
        ),
        migrations.CreateModel(
            name='Language',
            fields=[
                ('iso_639_1', models.CharField(max_length=2, serialize=False, primary_key=True)),
                ('iso_639_2T', models.CharField(unique=True, max_length=3, blank=True)),
                ('iso_639_2B', models.CharField(unique=True, max_length=3, blank=True)),
                ('iso_639_3', models.CharField(max_length=3, blank=True)),
                ('iso_639_6', models.CharField(max_length=4, blank=True)),
                ('name_en', models.CharField(max_length=100)),
                ('name_native', models.CharField(max_length=100)),
                ('name_other', models.CharField(max_length=50, blank=True)),
                ('family', models.CharField(max_length=50)),
                ('notes', models.CharField(max_length=100, blank=True)),
                ('countries_spoken', models.ManyToManyField(to='countries_plus.Country', null=True, blank=True)),
            ],
            options={
                'ordering': ['name_en'],
                'verbose_name': 'Language',
                'verbose_name_plural': 'Languages',
            },
            bases=(models.Model,),
        ),
        migrations.AddField(
            model_name='culturecode',
            name='language',
            field=models.ForeignKey(to='languages_plus.Language', on_delete=models.CASCADE),
            preserve_default=True,
        ),
        migrations.CreateModel(
            name='LangCountry',
            fields=[
            ],
            options={
                'proxy': True,
            },
            bases=('countries_plus.country',),
        ),
    ]
<|endoftext|>


<commit_before># coding: utf-8
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import os
import logging
import argparse
import json
from laniakea.core.common import Focus
from .manager import AzureManager

logger = logging.getLogger('Laniakea')

class AzureCommandLine(object):
    """
    Sub command-line interface for the Microsoft Azure provider.
    """
    HOME = os.path.dirname(os.path.abspath(__file__))
    VERSION = 0.1

    @classmethod
    def add_arguments(cls, subparsers, dirs):
        parser = subparsers.add_parser(
            'azure',
            help='Microsoft Azure',
            formatter_class=lambda prog:
                argparse.ArgumentDefaultsHelpFormatter(prog, max_help_position=30, width=100))

        m = parser.add_argument_group('Mandatory Azure Parameters')

        m.add_argument('-region',
                       type=str,
                       metavar='name',
                       help='Azure region.')
        m.add_argument('-count',
                       type=int,
                       metavar='n',
                       help='Number of instances to launch.')
        m.add_argument('-create',
                       action='store_true',
                       help='Create an instance pool.')
        m.add_argument('-delete',
                       action='store_true',
                       help='Delete an instance pool.')
        m.add_argument('-storage-name',
                       type=str,
                       metavar='name',
                       help='Storage name.')
        m.add_argument('-group-name',
                       type=str,
                       metavar='name',
                       help='Group name to be deleted.')
        m.add_argument('-azure',
                       metavar='path',
                       type=argparse.FileType(),
                       default=os.path.join(dirs.user_config_dir, 'azure.json'),
                       help='Deployment template for Windows Azure')

        u = parser.add_argument_group('UserData Parameters')
        u.add_argument('-template',
                       metavar='path',
                       type=argparse.FileType(),
                       default=os.path.join(dirs.user_config_dir, 'userdata', 'azure', 'template.json'),
                       help='Deployment template for Windows Azure')

        o = parser.add_argument_group('Optional Parameters')
        o.add_argument('-version',
                       action='version',
                       version='%(prog)s {}'.format(cls.VERSION),
                       help=argparse.SUPPRESS)

    @classmethod
    def main(self, args, settings):
        logger.info('Using Azure configuration from %s', Focus.info(args.azure.name))
        try:
            settings = json.loads(args.azure.read())
        except ValueError as msg:
            logger.error('Unable to parse %s: %s', args.settings.name, msg)
            return 1

        logger.info('Using Azure instance template from %s', Focus.info(args.template.name))
        try:
            template = json.loads(args.template.read())
        except ValueError as msg:
            logger.error('Unable to parse %s: %s', args.template.name, msg)
            return 1

        try:
            cluster = AzureManager(settings)
        except Exception as e:
            logger.error(e)
            return 1

        logger.info('Creating instance parameters.')
        try:
            parameters = cluster.create_parameters(args)
        except ValueError as msg:
            logger.error('Unable to create parameters.')
            return 1

        try:
            cluster.connect()
        except Exception as e:
            logger.error(e)
            return 1

        if args.create:
            try:
                cluster.create(parameters, args.storage_name, template)
            except ValueError as msg:
                return 1

        if args.delete:
                try:
                    cluster.terminate(args.group_name)
                except ValueError as msg:
                    return 1

<commit_msg>Update __init__.py<commit_after># coding: utf-8
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import os
import logging
import argparse
import json
from laniakea.core.common import Focus
from .manager import AzureManager

logger = logging.getLogger('Laniakea')

class AzureCommandLine(object):
    """
    Sub command-line interface for the Microsoft Azure provider.
    """
    HOME = os.path.dirname(os.path.abspath(__file__))
    VERSION = 0.1

    @classmethod
    def add_arguments(cls, subparsers, dirs):
        parser = subparsers.add_parser(
            'azure',
            help='Microsoft Azure',
            formatter_class=lambda prog:
                argparse.ArgumentDefaultsHelpFormatter(prog, max_help_position=30, width=100))

        m = parser.add_argument_group('Mandatory Azure Parameters')

        m.add_argument('-region',
                       type=str,
                       metavar='name',
                       help='Azure region.')
        m.add_argument('-count',
                       type=int,
                       metavar='n',
                       help='Number of instances to launch.')
        m.add_argument('-create',
                       action='store_true',
                       help='Create an instance pool.')
        m.add_argument('-delete',
                       action='store_true',
                       help='Delete an instance pool.')
        m.add_argument('-group-name',
                       type=str,
                       metavar='name',
                       help='Group name to be deleted.')
        m.add_argument('-azure',
                       metavar='path',
                       type=argparse.FileType(),
                       default=os.path.join(dirs.user_config_dir, 'azure.json'),
                       help='Deployment template for Windows Azure')

        u = parser.add_argument_group('UserData Parameters')
        u.add_argument('-template',
                       metavar='path',
                       type=argparse.FileType(),
                       default=os.path.join(dirs.user_config_dir, 'userdata', 'azure', 'template.json'),
                       help='Deployment template for Windows Azure')

        o = parser.add_argument_group('Optional Parameters')
        o.add_argument('-version',
                       action='version',
                       version='%(prog)s {}'.format(cls.VERSION),
                       help=argparse.SUPPRESS)

    @classmethod
    def main(self, args, settings):
        logger.info('Using Azure configuration from %s', Focus.info(args.azure.name))
        try:
            settings = json.loads(args.azure.read())
        except ValueError as msg:
            logger.error('Unable to parse %s: %s', args.settings.name, msg)
            return 1

        logger.info('Using Azure instance template from %s', Focus.info(args.template.name))
        try:
            template = json.loads(args.template.read())
        except ValueError as msg:
            logger.error('Unable to parse %s: %s', args.template.name, msg)
            return 1

        try:
            cluster = AzureManager(settings)
        except Exception as e:
            logger.error(e)
            return 1

        logger.info('Creating instance parameters.')
        try:
            parameters = cluster.create_parameters(args)
        except ValueError as msg:
            logger.error('Unable to create parameters.')
            return 1

        try:
            cluster.connect()
        except Exception as e:
            logger.error(e)
            return 1

        if args.create:
            try:
                cluster.create(parameters, args.storage_name, template)
            except ValueError as msg:
                return 1

        if args.delete:
                try:
                    cluster.terminate(args.group_name)
                except ValueError as msg:
                    return 1

<|endoftext|>