import ipywidgets as widgets
from ipywidgets import Button, Layout, ButtonStyle
from IPython.display import display

class Entator:
    current_annotations = []
    curent_class = 'no'
    index = 0
    
    clases2color = {}
    colors = ['lightgray','red','blue','yellow','green','orange','lightblue']
    
    labels_buttons = []
    text_buttons = []
    
    output = widgets.Output()
    
    annotation_label = widgets.Label('Annotations')
    sample_label = widgets.Label('Text for annotation')
    
    def __init__(self, labels, inputs, targets=None):
        
        self.inputs = inputs
        if targets is None:
            targets = [None] * len(inputs)
        self.targets = targets
        self.sample_label.value =f'Text for annotation: 1/{len(inputs)}'
        for label, color in zip(labels, self.colors):
            if self.curent_class == 'no' and label != 'no':
                self.curent_class = label
                self.annotation_label.value = f'Current Label: {self.curent_class}'
            self.clases2color[label] = color
            button = self.addButton(label, color)
            def on_button_clicked(b):
                with self.output:
                    self.curent_class = b.description
                    self.annotation_label.value = f'Current Label: {self.curent_class}'

            button.on_click(on_button_clicked)
            self.labels_buttons.append(button)
        
    def addSample(self, x, y):        
        '''add text to the output
        x - list of strings reprezenting the text
        y - list of strings reprezenting the classes
        '''
        if y is None:
            y = ['no']*len(x)
        self.current_annotations = y
        nr = 0
        self.text_buttons = []
        for text, label in zip(x,y):
            button = self.addButton(text, self.clases2color[label], nr)
            def on_button_clicked(b):
                with self.output:
                    label = self.curent_class
                    index = int(b.tooltip)
                    if self.current_annotations[index] == label:
                        # if click second time, add default label
                        label = 'no'
                    self.current_annotations[index] = label
                    color = self.clases2color[label]        
                    b.style=ButtonStyle(button_color=color)    

            button.on_click(on_button_clicked)
            self.text_buttons.append(button)
            nr+=1
    
    def splitLines(self):
        """split words of the sample in different lines"""
        lines = []
        line = []
        number_letters = 0
        for b in self.text_buttons:
            
            if number_letters+len(b.description)>100:
                lines.append(widgets.HBox(line))
                line = []
                number_letters = 0
            line.append(b)
            number_letters += len(b.description)
        if len(line)>0:
            lines.append(widgets.HBox(line))
        return widgets.VBox(lines)
    
    def addControlButtons(self):
        controls = []
        for text in ['back','next','skip']:
            button = self.addButton(text, 'gray')
            def on_button_clicked(b):
                with self.output:
                    if b.description == 'next':
                        if self.index< len(self.inputs):
                            self.targets[self.index] = self.current_annotations 
                        self.index+=1
                    if b.description == 'skip':
                        self.index+=1
                    if b.description == 'back':
                        if self.index< len(self.inputs):
                            self.targets[self.index] = self.current_annotations 
                        self.index-=1
                    
                    self.index = max(0, self.index)
                    self.index = min(len(self.inputs), self.index)
                    
                    if self.index< len(self.inputs):
                        self.sample_label.value = f'Text for annotation: {self.index+1}/{len(self.inputs)}'
                        self.addSample(self.inputs[self.index], self.targets[self.index])
                        self.text_section.children =  self.splitLines().children
                    else:
                        self.text_section.children = [widgets.Label('No more examples')]
            
            button.on_click(on_button_clicked)
            controls.append(button)
        return widgets.HBox(controls)
    
    def run(self):
        self.addSample(self.inputs[self.index], self.targets[self.index])
        self.text_section = self.splitLines()
        return display(widgets.VBox([self.annotation_label,
                                     widgets.HBox(self.labels_buttons),
                                     self.sample_label,
                                     self.text_section, 
                                     widgets.Label('Controls'),
                                     self.addControlButtons() ]), self.output)
        
    def addButton(self, text, color, nr=None):
        """
        text - 
        color - button color
        nr - index of the word in the sample
        """
        if nr is None:
            nr = text
        button = widgets.Button(
            description=text,
            disabled=False,
            button_style='',
            tooltip=str(nr),
            layout=Layout(margin='2px', padding='1px', width="auto" if nr!=text else '100px'),
            style=ButtonStyle(button_color=color)
        )
        return button

