from agent.Base_Agent import Base_Agent
from behaviors.custom.Stop.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 Stop:

    def __init__(self, base_agent: Base_Agent) -> None:
        self.last_executed = 0
        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/Stop/best_model.zip.pkl",
                                             "/behaviors/custom/Stop/best_model.zip.pkl",
                                             "/behaviors/custom/Stop/best_model.zip.pkl",
                                             "/behaviors/custom/Stop/best_model.zip.pkl",
                                             "/behaviors/custom/Stop/best_model.zip.pkl"
                                         ][self.world.robot.type]), 'rb') as f:
            self.model = pickle.load(f)

    def execute(self, reset):
        '''
        Parameters
        ----------
        target_2d : array_like
            2D target in absolute or relative coordinates (use is_target_absolute to specify)
        is_target_absolute : bool
            True if target_2d is in absolute coordinates, False if relative to robot's torso
        orientation : float
            absolute or relative orientation of torso, in degrees
            set to None to go towards the target (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
        distance : float
            distance to final target [0,0.5] (influences walk speed when approaching the final target)
            set to None to consider target_2d the final target
        '''
        r = self.world.robot

        # ------------------------ 0. Override reset (since some behaviors use this as a sub-behavior)
        if reset and self.world.time_local_ms - self.last_executed == 20:
            reset = False
        self.last_executed = self.world.time_local_ms

        # ------------------------ 1. Define walk parameters

        self.env.walk_rel_target = (0, 0)
        self.env.walk_distance = 0  # MAX_LINEAR_DIST = 0.5
        self.env.walk_rel_orientation = 0

        # ------------------------ 2. Execute behavior

        obs = self.env.observe(reset)
        action = run_mlp(obs, self.model)
        if self.env.execute(action):
            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
