#!/usr/bin/env python3

import time

import numpy as np

import pinocchio as pin

from manager import QueueClient
from task import Task, TaskParameters


class Boss(QueueClient):
    def run(self, qs, vqs, aqs):
        # create some tasks
        for i, [q, v, a] in enumerate(zip(qs, vqs, aqs)):
            self.tasks.put(Task(i, q, v, a))

        # Wait for the results
        results = []
        for _ in qs:
            result = self.results.get()
            print(f'got result {result.identifier} of processed in {result.time:.3f}s : ')
            results.append(result)
        print('Got all {tasks} results !')


if __name__ == '__main__':

    # Allocate data to be proceeded.
    NDATA = 20
    qs = np.array([pin.randomConfiguration(TaskParameters.robot.model) for _ in range(NDATA)])
    vqs = np.random.rand(NDATA, TaskParameters.robot.model.nv) * 2 - 1
    aqs = np.random.rand(NDATA, TaskParameters.robot.model.nv) * 2 - 1

    start = time.perf_counter()
    Boss().run(qs, vqs, aqs)
    total = time.perf_counter() - start

    print(f'Completed in {total:.3f} secs')