# This code is licensed under the MIT License (see LICENSE file for details)

import itertools
from PyQt5 import Qt
from ris_widget.qwidgets import annotator

from .. import process_data

class StageField(annotator.AnnotationField):
    FIRST_COLOR = (255, 255, 255)
    LAST_COLOR = (184, 184, 184)
    COLOR_CYCLE = itertools.cycle([(184, 255, 184), (255, 255, 184), (184, 184, 255), (255, 184, 184), (255, 184, 255)])

    def __init__(self, name='stage', stages=['egg', 'larva', 'adult', 'dead'], shortcuts=None):
        """Annotate the life-stage of a worm.

        If timestamp annotations are already present (i.e. as generated by
        process_data.update_annotations()), 'age' and 'ghost_age' annotations
        will also be produced as possible.

        Parameters:
            name: annotation name
            stages: list of life stages to annotate, generally starting with egg
                and ending with dead.
            shortcuts: shortcut keys to select the different stages; if not
                specified, the first letter of each stage will be used.
        """
        self.stages = stages
        self.stage_indices = {stage: i for i, stage in enumerate(stages)}
        if shortcuts is None:
            # take the first letter of each as the shortcut
            shortcuts = [stage[0] for stage in stages]
        self.shortcuts = shortcuts
        self.colors = {stages[0]: self.FIRST_COLOR, stages[-1]: self.LAST_COLOR}
        self.colors.update(zip(stages[1:-1], self.COLOR_CYCLE))
        super().__init__(name)

    def init_widget(self):
        self.widget = Qt.QGroupBox(self.name)
        layout = Qt.QHBoxLayout()
        self.widget.setLayout(layout)
        self.buttons = []
        self.button_group = Qt.QButtonGroup()
        for stage, key in zip(self.stages, self.shortcuts):
            button = Qt.QPushButton(stage)
            self.buttons.append(button)
            button.setCheckable(True)
            self.button_group.addButton(button)
            callback = self._make_stage_callback(button, stage)
            button.clicked.connect(callback)
            layout.addWidget(button)
            Qt.QShortcut(key, self.widget, callback, context=Qt.Qt.ApplicationShortcut)

    def _make_stage_callback(self, button, stage):
        def callback():
            self.set_stage(stage)
        return callback

    def set_stage(self, stage):
        self.update_annotation(stage)
        # now fix up pages before this page to comply with newly set stage
        fb_i = self.flipbook.pages.index(self.page)
        youngest_stage_i = self.stage_indices[stage]
        if fb_i > 0: # Exclude update of previous pages if this is the first image
            for page in self.flipbook.pages[fb_i-1::-1]:
                page_stage = self.get_annotation(page)
                if page_stage is None:
                    # if previous page has no annotation, use the annotation BEFORE
                    # the one clicked (if possible)
                    page.annotations[self.name] = self.stages[max(0, youngest_stage_i-1)]
                else:
                    page_stage_i = self.stage_indices[page_stage]
                    if page_stage_i < youngest_stage_i:
                        youngest_stage_i = page_stage_i
                    else:
                        page.annotations[self.name] = self.stages[youngest_stage_i]

        # pages after this will be brought into compliance by update_widget
        self.update_widget(stage)

    def update_widget(self, value):
        if value is None:
            # to deselect all buttons, need to make group temporarily non-exclusive
            self.button_group.setExclusive(False)
            for button in self.buttons:
                button.setChecked(False)
            self.button_group.setExclusive(True)
        elif value not in self.stages:
            raise ValueError(f'Value {value} not in list of stages.')
        else:
            self.buttons[self.stage_indices[value]].setChecked(True)
        # TODO: why do we need to manually call repaint() after setting the checked property here?
        # TODO: also, calling update() instead of repaint() can cause other parts of a floated widget
        # to get corrupt. WTF?
        for button in self.buttons:
            button.repaint()

        # now ensure that all pages follow correct stage ordering, and set
        # the page colors
        oldest_stage_i = -1
        for page in self.flipbook.pages:
            page_stage = self.get_annotation(page)
            page_stage_i = self.stage_indices.get(page_stage, -1) # will be < all others if stage is None
            if page_stage_i > oldest_stage_i:
                oldest_stage_i = page_stage_i
            elif page_stage_i < oldest_stage_i:
                # NB: if both are -1, then we haven't actually seen any annotations yet,
                # so we shouldn't do anything. Hence the test above for strict inequality.
                page_stage = page.annotations[self.name] = self.stages[oldest_stage_i]

            if page_stage is None:
                page.color = None
            else:
                page.color = self.colors[page_stage]
        self.update_ages()

    def update_ages(self):
        timepoint_annotations = []
        for page in self.flipbook.pages:
            if not hasattr(page, 'annotations'):
                return
            timepoint_annotations.append(page.annotations)
        if not hasattr(self, 'experiment_annotator'):
            return
        process_data._update_ages(timepoint_annotations, self.experiment_annotator.position.annotations, self.name, self.stages[0])
