import numpy as np
import math
import time

from controller import ControllerBase


class PIDController(ControllerBase):
    def __init__(self):
        super().__init__()
        self.t = None

        self.P_angle = 0.4
        self.I_angle = 0.35
        self.D_angle = 0.1

        self.P_pos = 0.50
        self.I_pos = 0.3
        self.D_pos = 0.1

        self.mode = 'combined'

        self.e_angle_old = 0.0
        self.e_pos_old = 0.0

        self.i = 0.0
        self.i_angle = 0.0
        self.i_pos = 0.0

    def set_target_position(self, target_pos):
        super(PIDController, self).set_target_position(target_pos)
        self.mode = 'combined'
        self.e_angle_old = 0.0
        self.e_pos_old = 0.0

        self.i = 0.0
        self.i_angle = 0.0
        self.i_pos = 0.0

    def compute_control(self, state):
        # measure state
        x_pred = state[1:]

        if self.t is None:
            dt = 0.1
        else:
            dt = time.time() - self.t   # time since last control was applied

        if self.mode == 'angle':
            # compute angle such that robot faces to target point
            target_angle = self.target_pos[2]

            angles_unwrapped = np.unwrap([x_pred[2], target_angle])  # unwrap angle to avoid jump in data

            e_angle = angles_unwrapped[0] - angles_unwrapped[1]   # angle difference
            p = self.P_angle * e_angle
            self.i += self.I_angle * dt * (e_angle + self.e_angle_old)/2.0  # trapezoidal rule
            d = self.D_angle * (e_angle - self.e_angle_old)/dt

            u1 = p + self.i + d
            u2 = - u1

            self.e_angle_old = e_angle

        elif self.mode == 'combined':
            # compute angle such that robot faces to target point
            v = self.target_pos[0:2] - x_pred[0:2]
            target_angle = math.atan2(v[1], v[0])

            angles_unwrapped = np.unwrap([x_pred[2], target_angle])  # unwrap angle to avoid jump in data

            e_angle = angles_unwrapped[0] - angles_unwrapped[1]  # angle difference
            e_pos = np.linalg.norm(v)

            if e_pos < 0.05:
                self.mode = 'angle'
                self.e_angle_old = 0
                self.e_pos_old = 0
                self.i_angle = 0
                self.i_pos = 0
                u1 = 0
                u2 = 0
            else:
                forward = abs(e_angle) < np.pi/2.0

                if not forward:
                    if e_angle > np.pi/2.0:
                        e_angle -= np.pi
                    else:
                        e_angle += np.pi

                p_angle = self.P_angle * e_angle
                self.i_angle += self.I_angle * dt * (e_angle + self.e_angle_old) / 2.0  # trapezoidal rule
                d_angle = self.D_angle * (e_angle - self.e_angle_old) / dt

                p_pos = self.P_pos * e_pos
                self.i_pos += self.I_pos * dt * (e_pos + self.e_pos_old) / 2.0  # trapezoidal rule
                d_pos = self.D_pos * (e_pos - self.e_pos_old) / dt

                if forward:
                    u1 = p_angle + p_pos + self.i_angle + self.i_pos + d_angle + d_pos
                    u2 = - p_angle + p_pos - self.i_angle + self.i_pos - d_angle + d_pos
                else:
                    u1 = p_angle - p_pos + self.i_angle - self.i_pos + d_angle - d_pos
                    u2 = - p_angle - p_pos - self.i_angle - self.i_pos - d_angle - d_pos

                self.e_pos_old = e_pos
                self.e_angle_old = e_angle
        else:
            u1 = 0.0
            u2 = 0.0
        return u1, u2

    def apply_control(self, control):
        super(PIDController, self).apply_control(control)
        self.t = time.time()  # save time when the most recent control was applied
        time.sleep(self.control_rate)