import time
import uuid

import crocoddyl
import numpy as np


def makeDDP():
    rcost_x = 0.1
    rcost_u = 0.2
    tcost_x = 3.0
    horizon = 100

    model = crocoddyl.ActionModelUnicycle()
    model.costWeights = np.array([rcost_x, rcost_u])
    modelT = crocoddyl.ActionModelUnicycle()
    modelT.costWeights = np.array([tcost_x, 0])
    state = np.zeros(3)
    problem = crocoddyl.ShootingProblem(state, [model] * horizon, modelT)
    return crocoddyl.SolverDDP(problem)


class TaskParameters:
    """Common information to all tasks, to be allocated by each minion and the boss."""

    ddp = makeDDP()


class Result:
    def __init__(self, ddp, converged=True):
        self.converged = converged
        self.x0 = ddp.problem.x0.copy()
        self.xs = np.array(ddp.xs)
        self.us = np.array(ddp.xs)
        self.v = ddp.cost
        self.Vx = ddp.Vx[0].copy()
        self.Vxx = ddp.Vxx[0].copy()
        self.iter = ddp.iter


class Task:
    def __init__(self, x0s, identifier=None):
        """
        Set a task with a list of initial state.
        Identifier is a uniq identifier of the task (for documentation only)
        that is generated randomly if none is provided.
        """

        self.identifier = identifier if identifier is not None else uuid.uuid1()
        self.time = 0
        self.x0s = x0s
        self.results = []

    def work(self):
        start = time.perf_counter()

        ddp = TaskParameters.ddp
        for x0 in self.x0s:
            ddp.problem.x0 = x0.copy()
            res = ddp.solve()
            self.results.append(Result(ddp, res))

        self.time = time.perf_counter() - start


if __name__ == "__main__":
    task = Task(np.random.rand(10, 3))
    task.work()