import pygame
import math


class Entity(object):
    def __call__(cls, *args, **kwargs):
        obj = type.__call__(cls, *args, **kwargs)
        obj.__post_init__()
        return obj

    def __init__(self, env):
        self.shape = None
        self.env = env
        self.pos = (env.random(), env.random())
        self.last_pos = None
        self.speed = (0, 0)
        self.acc = (0, 0)
        self.drag = 0
        self.col = (255, 255, 255)
        self.solid = False
        self.movable = False  # False = Non movable, True = Movable, x>1: lighter movable
        self.elasticity = 1
        self.collision_changes_speed = self.env.controll_type == 'ACC'
        self.collision_elasticity = self.env.default_collision_elasticity
        self._crash_list = []
        self._coll_add_pushback = 0
        self.crash_conservation_of_energy = True
        self.draw_path = False
        self.draw_path_col = [int(c/5) for c in self.col]
        self.draw_path_width = 2
        self.draw_path_harm = False
        self.draw_path_harm_col = [c for c in self.draw_path_col]
        self.draw_path_harm_col[0] += int(255/3)

    def __post_init__(self):
        pass

    def physics_step(self):
        x, y = self.pos
        vx, vy = self.speed
        ax, ay = self.acc
        vx, vy = vx+ax*self.env.acc_fac,  vy+ay*self.env.acc_fac
        x, y = x+vx*self.env.speed_fac, y+vy*self.env.speed_fac
        if not self.env.torus_topology:
            if x > 1 or x < 0:
                x, y, vx, vy = self.calc_void_collision(x < 0, x, y, vx, vy)
            if y > 1 or y < 0:
                x, y, vx, vy = self.calc_void_collision(
                    2 + (x < 0), x, y, vx, vy)
        else:
            x = x % 1
            y = y % 1
        self.speed = vx/(1+self.drag), vy/(1+self.drag)
        self.pos = x, y

    def controll_step(self):
        pass

    def step(self):
        self.controll_step()
        self.physics_step()
        self._crash_list = []

    def draw(self):
        self._draw_path()

    def _draw_path(self):
        if self.draw_path and self.last_pos:
            col = self.draw_path_col
            if self.draw_path_harm:
                if self.env.gotHarm:
                    col = self.draw_path_harm_col
            pygame.draw.line(self.env.path_overlay, col,
                             (self.last_pos[0]*self.env.width, self.last_pos[1]*self.env.height), (self.pos[0]*self.env.width, self.pos[1]*self.env.height), self.draw_path_width)
            pygame.draw.circle(self.env.path_overlay, col,
                               (self.pos[0]*self.env.width, self.pos[1]*self.env.height), max(0, self.draw_path_width/2-1))
        self.last_pos = self.pos[0], self.pos[1]

    def on_collision(self, other, depth):
        if self.solid and other.solid:
            if self.movable:
                self.on_crash(other, depth)

    def on_crash(self, other, depth):
        if other in self._crash_list:
            return
        self._crash_list.append(other)
        force_dir = self._get_crash_force_dir(other)
        #print(force_dir, depth)
        force_dir_len = math.sqrt(force_dir[0]**2+force_dir[1]**2)
        if force_dir_len == 0:
            return
        force_dir = force_dir[0]/force_dir_len, force_dir[1]/force_dir_len
        if not self.env.torus_topology:
            if self.env.agent.pos[0] > 0.99 or self.env.agent.pos[0] < 0.01:
                force_dir = force_dir[0], force_dir[1] * 2
            if self.env.agent.pos[1] > 0.99 or self.env.agent.pos[1] < 0.01:
                force_dir = force_dir[0] * 2, force_dir[1]
        depth *= 1.0*self.movable/(self.movable + other.movable)/2
        depth /= other.elasticity
        force_vec = force_dir[0]*depth/self.env.width, \
            force_dir[1]*depth/self.env.height
        if sum([abs(f) for f in force_vec]) > 0.005:
            self.pos = self.pos[0] + force_vec[0], self.pos[1] + force_vec[1]
            if self._coll_add_pushback:
                self.pos = self.pos[0] - self.env.inp[0]*self._coll_add_pushback * \
                    self.env.speed_fac, self.pos[1] - self.env.inp[1] * \
                    self._coll_add_pushback*self.env.speed_fac
        if self.collision_changes_speed:
            oldspeed = math.sqrt(self.speed[0]**2+self.speed[1]**2)
            self.speed = self.speed[0] + \
                force_vec[0]*self.collision_elasticity/self.env.speed_fac, self.speed[1] + \
                force_vec[1]*self.collision_elasticity/self.env.speed_fac
            newspeed = math.sqrt(self.speed[0]**2+self.speed[1]**2)
            if self.crash_conservation_of_energy and newspeed > oldspeed*1.1:
                self.speed = self.speed[0]/newspeed*1.1 * \
                    oldspeed, self.speed[1]/newspeed*oldspeed*1.1

    def _get_crash_force_dir(self, other):
        if 1 == 1:  # linter hack
            raise Exception(
                '[!] No collision-logic implemented for shape"'+str(self.shape)+'"')

    def on_collect(self, other):
        pass

    def on_collected(self):
        pass

    def calc_void_collision(self, dir, x, y, vx, vy):
        if dir < 2:
            x = min(max(x, 0), 1)
            vx = -vx*self.collision_elasticity*0.5*self.collision_changes_speed
        else:
            y = min(max(y, 0), 1)
            vy = -vy*self.collision_elasticity*0.5*self.collision_changes_speed
        return x, y, vx, vy

    def kill(self):
        self.env.kill_entity(self)


class CircularEntity(Entity):
    def __init__(self, env):
        super().__init__(env)
        self.shape = 'circle'
        self.radius = 10

    def draw(self):
        super().draw()
        x, y = self.pos
        pygame.draw.circle(self.env.surface, self.col,
                           (x*self.env.width, y*self.env.height), self.radius, width=0)

    def _get_crash_force_dir(self, other):
        if other.shape == 'circle':
            return self.pos[0] - other.pos[0], self.pos[1] - other.pos[1]
        elif other.shape == 'rect':
            pad = 0
            edge_size = min(self.radius, min(
                other.width/3, other.height/3)) + 1

            x, y = self.pos
            x, y = x*self.env.height, y*self.env.width
            left, top = x - self.radius + pad, y - self.radius + pad
            right, bottom = x + self.radius - pad, y + self.radius - pad
            lrcenter, tbcenter = x, y

            ox, oy = other.pos
            ox, oy = ox*self.env.height, oy*self.env.width
            oleft, otop = ox + pad, oy + pad
            oright, obottom = ox + other.width - pad, oy + other.height - pad
            olrcenter, otbcenter = ox + other.width/2, oy + other.height/2

            lr, tb = 0, 0

            if otop < bottom and obottom > bottom:
                # col from top
                tb = otop - bottom
                #print('t', tb)
            elif top < obottom and top > otop:
                # col from bottom
                tb = - top + obottom
                #print('b', tb)

            if right > oleft and right < oright:
                # col from left
                lr = oleft - right
                #print('l', lr)
            elif left < oright and left > oleft:
                # col from right
                lr = - left + oright
                #print('r', lr)

            if lr != 0 and tb != 0:
                if abs(abs(tb) - abs(lr)) < edge_size:
                    if abs(tb) < abs(lr):
                        return lr/5, tb
                    else:
                        return lr, tb/5
                if abs(tb) < abs(lr):
                    return 0, tb
                else:
                    return lr, 0

            return 0, 0
        else:
            raise Exception(
                '[!] Shape "circle" does not know how to collide with shape "'+str(other.shape)+'"')


class RectangularEntity(Entity):
    def __init__(self, env):
        super().__init__(env)
        self.shape = 'rect'
        self.width = 10
        self.height = 10

    def draw(self):
        super().draw()
        x, y = self.pos
        rect = pygame.Rect(x*self.env.width, y *
                           self.env.width, self.width, self.height)
        pygame.draw.rect(self.env.surface, self.col,
                         rect, width=0)

    def _get_crash_force_dir(self, other):
        raise Exception(
            '[!] Collisions in this direction not implemented for shape "rectangle"')


class Agent(CircularEntity):
    def __init__(self, env):
        super(Agent, self).__init__(env)
        self.pos = (0.5, 0.5)
        self.col = (0, 0, 255)
        self.drag = self.env.agent_drag
        self.controll_type = self.env.controll_type
        self.solid = True
        self.movable = True

    def controll_step(self):
        self._read_input()
        self.env.check_collisions_for(self)

    def _read_input(self):
        if self.controll_type == 'SPEED':
            self.speed = self.env.inp[0] - 0.5, self.env.inp[1] - 0.5
        elif self.controll_type == 'ACC':
            self.acc = self.env.inp[0] - 0.5, self.env.inp[1] - 0.5
        else:
            raise Exception('Unsupported controll_type')


class Enemy(Entity):
    def __init__(self, env):
        super(Enemy, self).__init__(env)
        self.col = (255, 0, 0)
        self.damage = 100
        self.radiateDamage = True

    def on_collision(self, other, depth):
        super().on_collision(other, depth)
        if isinstance(other, Agent):
            self.env.new_reward -= self.damage


class Barrier(Enemy):
    def __init__(self, env):
        super(Barrier, self).__init__(env)
        self.solid = True
        self.movable = False


class CircleBarrier(Barrier, CircularEntity):
    def __init__(self, env):
        super(CircleBarrier, self).__init__(env)


class RectBarrier(Barrier, RectangularEntity):
    def __init__(self, env):
        super().__init__(env)


class Chaser(Enemy, CircularEntity):
    def __init__(self, env):
        super(Chaser, self).__init__(env)
        self.target = self.env.agent
        self.arrow_fak = 100
        self.lookahead = 0

    def _get_arrow(self):
        tx, ty = self.target.pos
        x, y = self.pos
        fx, fy = x + self.speed[0]*self.lookahead*self.env.speed_fac, y + \
            self.speed[1]*self.lookahead*self.env.speed_fac
        dx, dy = (tx-fx)*self.arrow_fak, (ty-fy)*self.arrow_fak
        return self.env._limit_to_unit_circle((dx, dy))


class WalkingChaser(Chaser):
    def __init__(self, env):
        super(WalkingChaser, self).__init__(env)
        self.col = (255, 0, 0)
        self.chase_speed = 0.45

    def controll_step(self):
        arrow = self._get_arrow()
        self.speed = arrow[0] * self.chase_speed, arrow[1] * self.chase_speed


class FlyingChaser(Chaser):
    def __init__(self, env):
        super(FlyingChaser, self).__init__(env)
        self.col = (255, 0, 0)
        self.chase_acc = 0.5
        self.arrow_fak = 5
        self.lookahead = 8 + env.random()*2

    def controll_step(self):
        arrow = self._get_arrow()
        self.acc = arrow[0] * self.chase_acc, arrow[1] * self.chase_acc


class Collectable(CircularEntity):
    def __init__(self, env):
        super(Collectable, self).__init__(env)
        self.avaible = True
        self.enforce_not_on_barrier = False
        self.reward = 10
        self.collectors = []

    def on_collision(self, other, depth):
        super().on_collision(other, depth)
        if isinstance(other, Barrier):
            self.on_barrier_collision()
        else:
            for Col in self.collectors:
                if isinstance(other, Col):
                    other.on_collect(self)
                    self.on_collected()

    def on_collected(self):
        self.env.new_reward += self.reward

    def on_barrier_collision(self):
        if self.enforce_not_on_barrier:
            self.pos = (self.env.random(), self.env.random())
            self.env.check_collisions_for(self)


class Reward(Collectable):
    def __init__(self, env):
        super(Reward, self).__init__(env)
        self.col = (0, 255, 0)
        self.reward = 10
        self.collectors = [Agent]


class OnceReward(Reward):
    def __init__(self, env):
        super(OnceReward, self).__init__(env)
        self.reward = 500

    def on_collected(self):
        # Force rerender of value func (even in static envs)
        self.env._invalidate_value_map()

        self.env.new_abs_reward += self.reward
        self.kill()


class TeleportingReward(OnceReward):
    def __init__(self, env):
        super(TeleportingReward, self).__init__(env)
        self.enforce_not_on_barrier = True
        self.env.check_collisions_for(self)

    def on_collected(self):
        # Force rerender of value func (even in static envs)
        self.env._invalidate_value_map()

        self.env.new_abs_reward += self.reward
        self.pos = (self.env.random(), self.env.random())
        self.env.check_collisions_for(self)


class LoopReward(OnceReward):
    def __init__(self, env):
        super().__init__(env)
        self.loop = [[0.25, 0.5], [0.75, 0.5]]
        self.state = 0
        self.jump_to_state()
        self.barrier_physics = False

    def jump_to_state(self):
        # Force rerender of value func (even in static envs)
        self.env._invalidate_value_map()

        pos_vec = [v for v in self.loop[self.state]]
        if len(pos_vec) == 4:
            pos_vec = pos_vec[0] + pos_vec[2] * \
                (self.env.random()-0.5), pos_vec[1] + \
                pos_vec[3]*(self.env.random()-0.5)
        self.pos = pos_vec

    def next_state(self):
        self.state = (self.state + 1) % len(self.loop)

    def jump_next(self):
        self.next_state()
        self.jump_to_state()

    def on_collected(self):
        self.env.new_abs_reward += self.reward
        self.jump_next()

    def physics_step(self):
        if self.barrier_physics:
            self.env.check_collisions_for(self)
        super().physics_step()


class TimeoutReward(OnceReward):
    def __init__(self, env):
        super(TimeoutReward, self).__init__(env)
        self.enforce_not_on_barrier = True
        self.env.check_collisions_for(self)
        self.timeout = 10

    def set_avaible(self, value):
        self.avaible = value
        if self.avaible:
            self.col = (0, 255, 0)
        else:
            self.col = (50, 100, 50)

    def on_collected(self):
        if self.avaible:
            # Force rerender of value func (even in static envs)
            self.env._invalidate_value_map()

            self.env.new_abs_reward += self.reward
            self.set_avaible(False)
            self.env.timers.append((self.timeout, self.set_avaible, True))


class Ball(CircularEntity):
    def __init__(self, env):
        super(Ball, self).__init__(env)
        self.col = (255, 128, 0)
        self.drag = 0.0025
        self.solid = True
        self.movable = 10
        self.elasticity = 1
        self.collision_changes_speed = True
        self.wall_reflect_damping = 0.1

    def calc_void_collision(self, dir, x, y, vx, vy):
        if dir < 2:
            x = min(max(x, 0), 1)
            vx = -(vx/(1+self.wall_reflect_damping))
        else:
            y = min(max(y, 0), 1)
            vy = -(vy/(1+self.wall_reflect_damping))
        return x, y, vx, vy

    def physics_step(self):
        self.env.check_collisions_for(self)
        super().physics_step()


class Goal(Collectable):
    def __init__(self, env):
        super(Goal, self).__init__(env)
        self.col = (0, 200, 0)
        self.reward = 500
        self.radius = 20
        self.collectors = [Ball]


class TeleportingGoal(Goal):
    def __init__(self, env):
        super(TeleportingGoal, self).__init__(env)
        self.enforce_not_on_barrier = True
        self.env.check_collisions_for(self)

    def on_collected(self):
        # Force rerender of value func (even in static envs)
        self.env._invalidate_value_map()

        self.env.new_abs_reward += self.reward
        self.pos = (self.env.random(), self.env.random())
        self.env.check_collisions_for(self)


class FootballPlayer():
    def __init__(self, env, target):
        super(FootballPlayer, self).__init__(env)
        self.col = (200, 0, 100)
        self.target = target
        self.solid = True
        self.movable = 1
        self.elasticity = 1


class WalkingFootballPlayer(FootballPlayer, WalkingChaser):
    def __init__(self, env, target):
        super(WalkingFootballPlayer, self).__init__(env, target)
        self.target = target


class FlyingFootballPlayer(FootballPlayer, FlyingChaser):
    def __init__(self, env, target):
        super(FlyingFootballPlayer, self).__init__(env, target)

# Not a real entity. Is used in the config of RayObserver to reference the outer boundary of the environment.


class Void():
    def __init__(self, env):
        self.col = (50, 50, 50)
    pass