import numpy as np
import time

from controller import ControllerBase
from casadi_opt import OpenLoopSolver


class MPCController(ControllerBase):
    def __init__(self, N=20, T=1.0):
        super().__init__()
        self.t = None

        self.ols = OpenLoopSolver(N=N, T=T)
        self.ols.setup()
        self.control_rate = self.ols.T / self.ols.N

        self.mstep = 2

        # integrator
        self.omega_max = 5.0
        self.control_scaling = 0.5

    def set_target_position(self, target_pos):
        super(MPCController, self).set_target_position(target_pos)
        self.t = time.time()

    def compute_control(self, state):
        x_pred = np.array(state[1:])

        error_pos = np.linalg.norm(x_pred[0:2] - self.target_pos[0:2])
        angles_unwrapped = np.unwrap([x_pred[2], self.target_pos[2]])  # unwrap angle to avoid jump in data
        error_ang = np.abs(angles_unwrapped[0] - angles_unwrapped[1])

        if error_pos > 0.05 or error_ang > 0.2:
        #if error_pos > 0.05:
            # solve mpc open loop problem
            res = self.ols.solve(x_pred, self.target_pos)

            us1 = res[0] * self.control_scaling
            us2 = res[1] * self.control_scaling

            dt_mpc = time.time() - self.t
            if dt_mpc < self.control_rate:  # wait until next control can be applied
                time.sleep(self.control_rate - dt_mpc)
        else:
            us1 = [0] * self.mstep
            us2 = [0] * self.mstep

        return us1, us2

    def apply_control(self, control):
        if self.robot is not None:
            for i in range(0, self.mstep):  # option to use multistep mpc if len(range) > 1
                u1 = control[0][i]
                u2 = control[1][i]
                self.robot.send_cmd(u1, u2)
                if i < self.mstep:
                    time.sleep(self.control_rate)
                self.t = time.time()  # save time the most recent control was applied
        else:
            raise Exception("error: controller cannot apply control!\n"
                            "       there is no robot attached to the controller!")