"spanspan""Interactive Labeler

This module includes an InteractiveLabeler.
"""

import asyncio, signal
from trunklucator.webserver.server import WebServer
from trunklucator.webserver.aiothread import AioThread

import trunklucator.protocol.dto as dto
import trunklucator.const.msg as const_msg
import trunklucator.const.task_types as const_ttype

from trunklucator.streams import helper as streams
from trunklucator.mabandit.epsilon_greedy import EpsilonGreedy

from typing import List


TIMEOUT = 9000 #Do we really need it?
INDEX_KEY = 'idx'

def signal_handler(loop):
    loop.remove_signal_handler(signal.SIGTERM)
    #is_working = False

class Trunklucator:

    """Interactive Labeler

    InteractiveLabeler is a Labeler object that shows the feature through image
    using matplotlib and lets human label each feature through command line
    interface.

    Parameters
    ----------
    label_name: list
        Let the label space be from 0 to len(label_name)-1, this list
        corresponds to each label's name.

    """

    def __init__(self, *args, **kwargs):
        self.aiothread = None
        self.loop = None
        self.ws = None
        self.bandit = None
        self.args = args
        self.kwargs = kwargs


    def start_thread(self):
        self.aiothread = AioThread()
        self.loop = self.aiothread.get_loop()
        #self.loop.add_signal_handler(signal.SIGTERM, signal_handler, self.loop)
        self.ws = WebServer(loop=self.loop, *self.args, **self.kwargs)
        self.aiothread.set_server(self.ws)
        self.aiothread.start()
        self.aiothread.event.wait()

    def stop_thread(self):
        self.aiothread.finalize()

    def __enter__(self):
        self.start_thread()
        return self

    def __exit__(self, *args):
        self.stop_thread()
        return False

    def ask(self, X):
        #create task object
        task_data = dto.Data(dto.get_id(), X)
        coro = self.aiothread.server.add_task(task_data)
        future = self.aiothread.add_task(coro)
        #Make sure you wait for loop to start. Calling future.cancel() in main thread will cancel asyncio coroutine in background thread.
        try:
            result = future.result(TIMEOUT)
            return result
        except asyncio.TimeoutError:
            print('The coroutine took too long, cancelling the task')
            future.cancel()
        except Exception as exc:
            print('The coroutine raised an exception: {!r}'.format(exc))


    def update(self, X):
        #create task object
        coro = self.aiothread.server.publish_update(dto.Update(X))
        _ = self.aiothread.add_task(coro)

    def process_iterables(self, iterables, labels: List[str], index_key=INDEX_KEY):
        if not self.bandit:
            self.bandit = EpsilonGreedy(0.4)

        self.bandit.initialize(len(labels))
        arm: int = None
        for X in streams.read_iterables(iterables, index_name=index_key):
            if not arm:
                arm = self.bandit.select_arm()
            if X[index_key] != arm:
                continue
            result = self.ask(X)
            if int(result.y) >= 0:
                self.bandit.update(arm, result.y)
                update = [{"value": v, "count": c, "label": l.name} for v, c, l in zip(self.bandit.values, self.bandit.counts, labels)]
                self.update(update)

            X[ANSWER_KEY] = result.y
            X['arm'] = arm # Do we need this?
            arm = None

            yield X

    def process_iterable(self, stream):
        for X in stream:
            result = self.ask(X)
            X[ANSWER_KEY] = result.y
            yield res
