import numpy as np
import sys


class AccelerationSimulator:
    def __init__(self, weather_info=3.0, max_speed=2.0):
        self.old_speed = np.zeros(2,)  # speed values for last steps, in m/s
        self.computed_new_speed = np.zeros(2,)  # speed values for current steps, in m/s
        self.accelerate_value_new = np.zeros(2,)  # accelerate values for current step, in m/s^2
        self.goal_dis = np.zeros(2,)  # distance that agent needs to move, in meters
        self.weight_of_agent = 20.0  # the weight for agent, in kg, shouldn't smaller than or equal to 0.0
        self.gravity_coefficient = 9.80665  # gravitational constant, in m/s^2, shouldn't be easily changed
        self.coefficient_of_friction = 0.0  # coefficient of friction, changes with weather
        self.weather_info = weather_info  # weather, 0.0 present sunny, default weather is sunny
        self.gain = 0.05  # gain factor to adjust the friction, positive, maximum 1.0
        self.none_zero = np.array([1e-32, 1e-32])
        # to prevent the speed change to be zero, which makes the incorrect force disintegration
        self.max_speed = max_speed  # maximum speed that agent is allowed to go
        '''''''''
        list of weather info
        0.0 presents sunny
        1.0 presents rain
        2.0 presents snow without freezing land
        3.0 presents snow with freezing land 
        '''''''''
    def _get_coefficient_of_friction(self):
        if self.weather_info == 0.0:  # sunny
            self.coefficient_of_friction = 0.6 * self.gain
        elif self.weather_info == 1.0:  # rain
            self.coefficient_of_friction = 0.4 * self.gain
        elif self.weather_info == 2.0:  # snow without freezing land
            self.coefficient_of_friction = 0.28 * self.gain
        elif self.weather_info == 3.0:  # snow with freezing land
            self.coefficient_of_friction = 0.18 * self.gain
        else:
            print("the weather information is not correct, please checkout your weather setting")
            sys.exit()

    def _make_vector_force(self, force):
        acceleration = np.linalg.norm(self.accelerate_value_new)
        force_x = (self.accelerate_value_new[0] / (acceleration + 1e-32)) * force
        force_y = (self.accelerate_value_new[1] / (acceleration + 1e-32)) * force
        # 1e-32 is to prevent division by zero errors
        force_vector = np.array([force_x, force_y])
        return force_vector

    def compute_actual_movement_and_speed(self, action_time, goal_dis, last_step, reset):
        if reset:
            self.old_speed = np.array([0.0, 0.0])  # while env has reset, the history speed should be back to zero
        self._get_coefficient_of_friction()  # get weather factor that effects the friction of agent
        self.computed_new_speed = goal_dis / action_time
        # According to Newton's Laws of Kinematics, that is the agent's speeds are supposed to go without friction
        self.goal_dis = goal_dis
        speed_change = self.computed_new_speed - self.old_speed + self.none_zero

        # to prevent the speed change to be zero, which makes the incorrect force disintegration
        self.accelerate_value_new = speed_change / (action_time + 1e-32)
        # 1e-32 is to prevent division by zero errors
        agent_force = self.weight_of_agent * np.linalg.norm(self.accelerate_value_new)
        friction = self.weight_of_agent * self.gravity_coefficient * self.coefficient_of_friction
        # 117.6N with these settings
        if friction >= agent_force and np.linalg.norm(self.old_speed) == 0.0:
            friction = agent_force
        # 1e-32 is to prevent division by zero errors
        friction_vector = self._make_vector_force(friction)  # compute the vector friction
        agent_force_vector = self._make_vector_force(agent_force)
        # Friction will prevent agent from accelerating to the expected acceleration within delta t
        # compute the effect of friction on both X and Y direction
        accelerate_value_true = np.zeros(2,)  # Acceleration buffer for pure acceleration/deceleration
        accelerate_value_true_0 = np.zeros(2, )  # Both acceleration and deceleration acceleration buffer
        accelerate_value_true_1 = np.zeros(2, )  # Both acceleration and deceleration acceleration buffer
        true_movement = np.zeros(2,)
        true_speed = np.zeros(2,)
        true_average_speed = np.zeros(2,)

        # we compute the actual movement and speed that agent could do in this dynamic env based on
        # Newton's Laws of Kinematics
        for i in range(2):
            if self.computed_new_speed[i] > self.old_speed[i] >= 0.0 or \
                    self.computed_new_speed[i] < self.old_speed[i] <= 0.0:
                true_force = agent_force_vector[i] - friction_vector[i]
                accelerate_value_true[i] = true_force / self.weight_of_agent
                # if true force < 0, agent doesn't have enough force to speed up, start to speed down
                true_speed[i] = self.old_speed[i] + accelerate_value_true[i] * action_time
                # true speed should be limited under the maximum speed
                if np.abs(true_speed[i]) >= self.max_speed:
                    if true_speed[i] >= 0:
                        true_speed[i] = self.max_speed
                    else:
                        true_speed[i] = -1.0 * self.max_speed
                true_average_speed[i] = np.average([true_speed[i], self.old_speed[i]])
                true_movement[i] = true_average_speed[i] * action_time
            elif self.old_speed[i] > self.computed_new_speed[i] >= 0.0 or \
                    self.old_speed[i] < self.computed_new_speed[i] <= 0.0:
                true_force = agent_force_vector[i] + friction_vector[i]
                accelerate_value_true[i] = true_force / self.weight_of_agent
                # friction can help agent speed down
                true_speed[i] = self.old_speed[i] + accelerate_value_true[i] * action_time
                # true speed should be limited under the maximum speed
                if np.abs(true_speed[i]) >= self.max_speed:
                    if true_speed[i] >= 0:
                        true_speed[i] = self.max_speed
                    else:
                        true_speed[i] = -1.0 * self.max_speed
                true_average_speed[i] = np.average([true_speed[i], self.old_speed[i]])
                true_movement[i] = true_average_speed[i] * action_time
            elif self.computed_new_speed[i] / (self.old_speed[i] + 1e-32) < 0.0 \
                    and np.linalg.norm(self.old_speed) != 0.0:
                # 1e-32 is to prevent division by zero errors
                true_force_0 = agent_force_vector[i] + friction_vector[i]
                true_force_1 = agent_force_vector[i] - friction_vector[i]
                # agent needs to go slow down first, then speed up
                accelerate_value_true_0[i] = true_force_0 / self.weight_of_agent
                time_0 = np.abs(self.old_speed[i] / (accelerate_value_true_0[i] + 1e-32))
                if time_0 >= action_time:
                    true_speed[i] = self.old_speed[i] + accelerate_value_true_0[i] * action_time
                    # true speed should be limited under the maximum speed
                    if np.abs(true_speed[i]) >= self.max_speed:
                        if true_speed[i] >= 0:
                            true_speed[i] = self.max_speed
                        else:
                            true_speed[i] = -1.0 * self.max_speed
                    true_movement[i] = (true_speed[i] / 2.0) * action_time
                    true_force = true_force_0
                    accelerate_value_true[i] = true_force / self.weight_of_agent
                else:
                    speed_true_0 = self.old_speed[i] + accelerate_value_true_0[i] * time_0
                    if np.abs(speed_true_0) >= self.max_speed:
                        if speed_true_0 >= 0:
                            speed_true_0 = self.max_speed
                        else:
                            speed_true_0 = -1.0 * self.max_speed
                    movement_0 = (speed_true_0/2.0) * time_0
                    true_force = true_force_1
                    accelerate_value_true_1[i] = true_force / self.weight_of_agent
                    time_1 = np.abs(self.computed_new_speed[i] / (accelerate_value_true_1[i] + 1e-32))
                    if time_1 >= action_time - time_0:
                        time_1 = action_time - time_0
                    speed_true_1 = 0.0 + accelerate_value_true_1[i] * time_1
                    if np.abs(speed_true_1) >= self.max_speed:
                        if speed_true_1 >= 0:
                            speed_true_1 = self.max_speed
                        else:
                            speed_true_1 = -1.0 * self.max_speed
                    movement_1 = (speed_true_1/2.0) * time_1
                    true_movement[i] = movement_0 + movement_1
                    true_speed[i] = speed_true_1
                    accelerate_value_true[i] = accelerate_value_true_1[i]
            else:  # the old and new speed are equaled
                true_force = agent_force_vector[i] - friction_vector[i]
                accelerate_value_true[i] = true_force / self.weight_of_agent
                true_speed[i] = self.old_speed[i] + accelerate_value_true[i] * action_time
                # true speed should be limited under the maximum speed
                if np.abs(true_speed[i]) >= self.max_speed:
                    if true_speed[i] >= 0:
                        true_speed[i] = self.max_speed
                    else:
                        true_speed[i] = -1.0 * self.max_speed
                true_average_speed[i] = np.average([true_speed[i], self.old_speed[i]])
                true_movement[i] = true_average_speed[i] * action_time
        self.old_speed = true_speed  # real speed will be recorded
        if last_step:
            # there is no more force provided by agent, so the agent will be stopped
            # when its velocity is zero due to friction
            last_speed = np.linalg.norm(self.old_speed)
            if last_speed != 0:
                friction = self.weight_of_agent * self.gravity_coefficient * self.coefficient_of_friction
                accelerate_value_last = np.abs(0.0 - friction / self.weight_of_agent)
                accelerate_value_last_vector = self._make_vector_force(accelerate_value_last)
                for i in range(len(accelerate_value_last_vector)):
                    if accelerate_value_last_vector[i] / (self.old_speed[i] + 1e-32) < 0.0:
                        accelerate_value_last_vector[i] *= -1.0
                # make a correction of the direction of acceleration if needed
                time_last = last_speed / accelerate_value_last
                true_movement_last = np.zeros(2,)
                for i in range(len(true_movement)):
                    true_movement_last[i] = self.old_speed[i] * time_last + 0.5 * accelerate_value_last_vector[i] * \
                                            time_last**2
                true_speed = np.zeros(2,)
                true_movement += true_movement_last
            else:
                true_movement = true_movement
                true_speed = np.zeros(2,)
        ###############################################################################################################
        return true_movement, true_speed
