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


class Kick_RL:

    def __init__(self, base_agent: Base_Agent) -> None:
        self.base_agent = base_agent
        self.world = base_agent.world
        self.description = 'Walk to ball and kick in a given direction'
        self.auto_head = True
        self.env = Env(self.base_agent,self.world )
        self.ik = base_agent.inv_kinematics
        self.behavior = base_agent.behavior

        if self.world.robot.type == 1:
            model_file = "/behaviors/custom/Kick_RL/kick_all.pkl"
            self.kick_time = (8, 8)
            self.kick_target_bias = (14, 7)
            self.kick_cap_spd = 4
            self.kick_dist = 7
            self.kick_init_pos = 0.3
            with open(M.get_active_directory(model_file), 'rb') as f:
                self.model = pickle.load(f)
        else:
            model_file = "/behaviors/custom/Kick_RL/kick_r2.pkl"
            self.kick_time = (11, 9)
            self.kick_target_bias = (7, 0)
            self.kick_cap_spd = 7.03
            self.kick_dist = 11
            self.kick_init_pos = 0.26 if self.world.robot.type == 0 else 0.26
            with open(M.get_active_directory(model_file), 'rb') as f:
                self.model = pickle.load(f)

    # WARNING: Decompyle incomplete

    def get_kick_pos_rot(self, direction, ball_2d, cap_power):
        ''' Returns kick position and desired robot orientation (with correction bias) '''
        direction += self.kick_target_bias[int(cap_power)]
        dir_rad = direction * math.pi / 180
        target_unit_vec = np.array([
            math.cos(dir_rad),
            math.sin(dir_rad)])
        left_unit_vec = np.array([
            -target_unit_vec[1],
            target_unit_vec[0]])
        return (ball_2d - self.kick_init_pos * target_unit_vec) + 0.045 * left_unit_vec, direction

    def execute(self, reset, direction, cap_power=(False,)):
        r = self.world.robot
        b = self.world.ball_abs_pos[:2]
        me = r.loc_torso_position[:2]
        if reset:
            self.phase = 0
        kick_dir = direction + self.kick_target_bias[int(cap_power)] if self.phase == 0 else self.env.kick_dir
        dir_rad = kick_dir * math.pi / 180
        target_unit_vec = np.array([
            math.cos(dir_rad),
            math.sin(dir_rad)])
        left_unit_vec = np.array([
            -target_unit_vec[1],
            target_unit_vec[0]])
        ball_rel_hip_center = self.ik.torso_to_hip_transform(self.world.ball_rel_torso_cart_pos)
        ball_dist = np.linalg.norm(ball_rel_hip_center)
        rot_with_bias = kick_dir + 8
        if self.phase == 0:
            target = (b - self.kick_init_pos * target_unit_vec) + 0.045 * left_unit_vec
            if np.linalg.norm(target - me) < 0.1:
                self.phase = 1
                self.env.kick_dir = kick_dir
            if reset:
                pass
            reset_walk = self.behavior.previous_behavior != 'Walk'
            self.behavior.execute_sub_behavior('Walk', reset_walk, target, True, rot_with_bias, True, None)
        elif self.phase == 1:
            target = (b - 0.1 * target_unit_vec) + 0.045 * left_unit_vec
            if ball_dist < 0.27:
                self.phase = 2
            else:
                self.behavior.execute_sub_behavior('Walk', False, target, True, rot_with_bias, True, None)
        if self.phase == 2:
            walk_rl = self.behavior.get_custom_behavior_object('Walk').env
            if walk_rl.step_generator.state_is_left_active and walk_rl.step_generator.state_current_ts == 1:
                self.phase = 3
            else:
                self.behavior.execute_sub_behavior('Walk', False, (0, 0), True, 0, True, None)
        if self.phase == 3:
            obs = self.env.observe(True)
            action = run_mlp(obs, self.model)
            self.env.execute(action)
            self.phase = 4
            return False
        if self.phase > 3:
            obs = self.env.observe(False)
            action = run_mlp(obs, self.model)
            self.env.execute(action)
            self.phase += 1
            print(self.phase)
            if self.phase >= self.kick_time[int(cap_power)]:
                print("True")
                return True
            return False

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