from agent.Base_Agent import Base_Agent
from behaviors.custom.Sprint.Env import Env
from math_ops.Math_Ops import Math_Ops as M
from math_ops.Neural_Network import run_mlp
import numpy as np
import pickle


class Sprint:

    def __init__(self, base_agent: Base_Agent) -> None:
        self.behavior = base_agent.behavior
        self.path_manager = base_agent.path_manager
        self.world = base_agent.world
        self.description = "RL dribble"
        self.auto_head = True
        self.env = Env(base_agent, 0.9 if self.world.robot.type == 3 else 1.2)

        with open(M.get_active_directory([
                                             "/behaviors/custom/Sprint/Sprint_R0.pkl",
                                             "/behaviors/custom/Sprint/Sprint_R1.pkl",
                                             "/behaviors/custom/Sprint/Sprint_R2.pkl",
                                             "/behaviors/custom/Sprint/Sprint_R3.pkl",
                                             "/behaviors/custom/Sprint/Sprint_R4.pkl"
                                         ][self.world.robot.type]), 'rb') as f:
            self.model = pickle.load(f)

    def define_approach_orientation(self):

        w = self.world
        b = w.ball_abs_pos[:2]
        me = w.robot.loc_head_position[:2]

        self.approach_orientation = None

        MARGIN = 0.8  # safety margin (if ball is near the field limits by this amount, the approach orientation is considered)
        M90 = 90 / MARGIN  # auxiliary variable
        DEV = 25  # when standing on top of sidelines or endlines, the approach direction deviates from that line by this amount
        MDEV = (90 + DEV) / MARGIN  # auxiliary variable

        a1 = -180  # angle range start (counterclockwise rotation)
        a2 = 180  # angle range end   (counterclockwise rotation)

        if b[1] < -10 + MARGIN:
            if b[0] < -15 + MARGIN:
                a1 = DEV - M90 * (b[1] + 10)
                a2 = 90 - DEV + M90 * (b[0] + 15)
            elif b[0] > 15 - MARGIN:
                a1 = 90 + DEV - M90 * (15 - b[0])
                a2 = 180 - DEV + M90 * (b[1] + 10)
            else:
                a1 = DEV - MDEV * (b[1] + 10)
                a2 = 180 - DEV + MDEV * (b[1] + 10)
        elif b[1] > 10 - MARGIN:
            if b[0] < -15 + MARGIN:
                a1 = -90 + DEV - M90 * (b[0] + 15)
                a2 = -DEV + M90 * (10 - b[1])
            elif b[0] > 15 - MARGIN:
                a1 = 180 + DEV - M90 * (10 - b[1])
                a2 = 270 - DEV + M90 * (15 - b[0])
            else:
                a1 = -180 + DEV - MDEV * (10 - b[1])
                a2 = -DEV + MDEV * (10 - b[1])
        elif b[0] < -15 + MARGIN:
            a1 = -90 + DEV - MDEV * (b[0] + 15)
            a2 = 90 - DEV + MDEV * (b[0] + 15)
        elif b[0] > 15 - MARGIN and abs(b[1]) > 1.2:
            a1 = 90 + DEV - MDEV * (15 - b[0])
            a2 = 270 - DEV + MDEV * (15 - b[0])

        cad = M.vector_angle(b - me)  # current approach direction

        a1 = M.normalize_deg(a1)
        a2 = M.normalize_deg(a2)

        if a1 < a2:
            if a1 <= cad <= a2:
                return  # current approach orientation is within accepted range
        else:
            if a1 <= cad or cad <= a2:
                return  # current approach orientation is within accepted range

        a1_diff = abs(M.normalize_deg(a1 - cad))
        a2_diff = abs(M.normalize_deg(a2 - cad))

        self.approach_orientation = a1 if a1_diff < a2_diff else a2  # fixed normalized orientation

    def execute(self, reset, target, orientation, is_orientation_absolute, speed=1, stop=False):
        '''
        Parameters
        ----------
        target
        orientation : float
            absolute or relative orientation of torso (relative to imu_torso_orientation), in degrees
            set to None to dribble towards the opponent's goal (is_orientation_absolute is ignored)
        is_orientation_absolute : bool
            True if orientation is relative to the field, False if relative to the robot's torso
        speed : float
            speed from 0 to 1 (scale is not linear)
        stop : bool
            return True immediately if walking, wind down if dribbling, and return True when possible
        '''
        w = self.world
        r = self.world.robot
        me = r.loc_head_position[:2]
        b = w.ball_abs_pos[:2]
        b_rel = w.ball_rel_torso_cart_pos[:2]
        b_dist = np.linalg.norm(b - me)
        behavior = self.behavior
        reset_dribble = False
        lost_ball = (w.ball_last_seen <= w.time_local_ms - w.VISUALSTEP_MS) or np.linalg.norm(b_rel) > 0.4

        if reset:
            self.phase = 0
            if behavior.previous_behavior == "Push_RL" and 0 < b_rel[0] < 0.25 and abs(b_rel[1]) < 0.07:
                self.phase = 1
                reset_dribble = True

        if self.phase == 0:  # prepare to sprint
            self.phase = 1
            reset_dribble = True

            if stop:
                return True

        if self.phase == 1 and (stop):  # go back to walking
            self.phase += 1
        elif self.phase == 1:  # dribble
            # ------------------------ 1. Define dribble parameters
            self.env.dribble_speed = speed
            if target is not None:
                raw_target = target - r.loc_head_position[:2]
                walk_rel_target = M.rotate_2d_vec(raw_target, -r.imu_torso_orientation)
                orientation = M.vector_angle(walk_rel_target)
            # Relative orientation values are decreased to avoid overshoot
            if orientation is None:
                if b[0] < 0:  # dribble to the sides
                    if b[1] > 0:
                        dribble_target = (15, 5)
                    else:
                        dribble_target = (15, -5)
                else:
                    dribble_target = None  # go to goal
                self.env.dribble_rel_orientation = \
                    self.path_manager.get_dribble_path(optional_2d_target=dribble_target)[1]
            elif is_orientation_absolute:
                self.env.dribble_rel_orientation = M.normalize_deg(orientation - r.imu_torso_orientation)
            else:  # sprint use this
                self.env.dribble_rel_orientation = float(orientation)  # copy if numpy float

            # ------------------------ 2. Execute behavior
            obs = self.env.observe(reset_dribble, virtual_ball=True)
            action = run_mlp(obs, self.model)
            self.env.execute(action)

        # wind down dribbling, and then reset phase
        if self.phase > 1:
            WIND_DOWN_STEPS = 60
            # ------------------------ 1. Define dribble wind down parameters
            self.env.dribble_speed = 1 - self.phase / WIND_DOWN_STEPS
            self.env.dribble_rel_orientation = 0

            # ------------------------ 2. Execute behavior
            obs = self.env.observe(reset_dribble, virtual_ball=True)
            action = run_mlp(obs, self.model)
            self.env.execute(action)

            # ------------------------ 3. Reset behavior
            self.phase += 1
            if self.phase >= WIND_DOWN_STEPS - 5:
                self.phase = 0
                return True

        return False

    def is_ready(self):
        ''' Returns True if this behavior is ready to start/continue under current game/robot conditions '''
        return True
