import math
import logging

from learning_racer.config.config import ConfigReader

logger = logging.getLogger(__name__)

config = ConfigReader()


def reward(action, e_i, done):
    if done:
        norm_throttle = (action[1] - config.agent_min_throttle()) / (
                config.agent_max_throttle() - config.agent_min_throttle())
        return config.reward_reward_crash() - (config.reward_crash_reward_weight() * norm_throttle), done
    throttle_reward = config.reward_throttle_reward_weight() * (action[1] / config.agent_max_throttle())
    return 1 + throttle_reward, done


# For gym_donkey
def reward_sim_speed(self, done):
    """'pos': (self.x, self.y, self.z), 'cte': self.cte,
                   "speed": self.speed, "hit": self.hit"""
    if self.determine_crash():
        return config.reward_reward_crash() - config.reward_crash_reward_weight() * (self.speed / 18.0)
    throttle_reward = config.reward_throttle_reward_weight() * (self.speed / 18.0)
    return 1 + throttle_reward - config.reward_cte_reward_weight() * math.fabs(self.cte + config.reward_cte_offset()) / self.max_cte

def reward_sim_throttle(self, done):
    if self.determine_crash():
        norm_throttle = (self.throttle - config.agent_min_throttle()) / (
                config.agent_max_throttle() - config.agent_min_throttle())
        return config.reward_reward_crash() - config.reward_crash_reward_weight() * norm_throttle
    throttle_reward = config.reward_throttle_reward_weight() * (self.throttle / config.agent_max_throttle())
    return 1 + throttle_reward - config.reward_cte_reward_weight() * math.fabs(self.cte  + config.reward_cte_offset()) / self.max_cte

def reward_sim_vd(self, done):
    if self.determine_crash():
        return config.reward_reward_crash()
    return (1.0 - (math.fabs(self.cte + config.reward_cte_offset()) / self.max_cte)) * self.speed

def reward_sim_v(self, done):
    if self.determine_crash() :
        return config.reward_reward_crash()
    return self.speed

first_lap = -1 

def reward_sim_lap_fix_offset(self, done):
    if self.determine_crash():
        return config.reward_reward_crash()

    reward = (1.0 - (math.fabs(self.cte + config.reward_cte_offset()) / self.max_cte)) * self.speed
    if len(self.lap_times) > 0 and self.reward_param["lap_num"] != len(self.lap_times):
        print('[reward_sim_lap_fix_offset()] laptime is added to reward')
        self.reward_param["lap_num"] = len(self.lap_times)
        reward += math.exp(2*(config.reward_fix_offset() - self.lap_times[-1]))
    
    return reward

def reward_sim_lap_first_lap(self, done):
    global first_lap
    if self.determine_crash():
        first_lap = -1 
        return config.reward_reward_crash()

    if first_lap == -1 and len(self.lap_times) == 1:
        first_lap = self.lap_times[0]

    reward = (1.0 - (math.fabs(self.cte + config.reward_cte_offset()) / self.max_cte)) * self.speed
        
    if len(self.lap_times) > 0 and self.reward_param["lap_num"] != len(self.lap_times):
        self.reward_param["lap_num"] = len(self.lap_times)
        reward += math.exp(2*(first_lap - self.lap_times[-1]))
    
    return reward


# For gym_donkey
hit_counter = 0
pass_counter = 0

def determine_crash_sim(self):
    global hit_counter, pass_counter

    # we have a few initial frames on start that are sometimes very large CTE when it's behind
    # the path just slightly. We ignore those.
    if math.fabs(self.cte + config.reward_cte_offset()) > 5 * self.max_cte:
        pass_counter += 1
        logger.info(f"Game Over: pass")
        if pass_counter > 30:
            pass_counter = 0
            return True
        pass
    elif math.fabs(self.cte + config.reward_cte_offset()) > self.max_cte:
        logger.info(f"Game Over: cte {self.cte + config.reward_cte_offset()}")
        return True
    elif self.hit != "none":
        hit_counter += 1
        logger.info(f"Game Over: hit({self.hit}), counter({hit_counter})")
        return True
    elif self.missed_checkpoint:
        logger.info(f"Game Over: missed checkpoint")
        return True
    elif self.dq:
        logger.info(f"Game Over: disqualified")
        return True
    return False
    
def episode_over_sim(self):
    if self.determine_crash():
        self.over = True

    if config.agent_max_lapnum_per_episode() == -1:
       return
    if self.starting_line_num >=  config.agent_max_lapnum_per_episode():
        logger.info(f"Game Clear: Collision with starting line")
        self.over = True

