#!/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")