import jax
import jax.numpy as jnp

from space.continuous.structs import *




def create_rays(agent_pos:tuple, ray_span = RAY_SPAN, ray_length = RAY_MAX_LENGTH):
    """
    Generate rays for the  agent given the agent position, the span of the rays, the length of the rays and the number of rays
    Args:
        - agent_pos: The position of the agent: (x, y, angle)
        - ray_span: The span of the rays i.e. how wide the agent can see times 0.5
        - ray_length: The length of the rays
        - RAY_RESOLUTION: The number of rays is global
    Returns:
        The rays generated by the foraer agent at the given position

    """
    x, y, angle = agent_pos

    ray_angles = jnp.linspace(angle - ray_span, angle + ray_span, RAY_RESOLUTION)
    cos_ray_angles = jnp.cos(ray_angles)
    sin_ray_angles = jnp.sin(ray_angles)
    ray_directions = jax.vmap(Point)(cos_ray_angles, sin_ray_angles)
    
    ray_origin = Point(x, y)
    
    rays = jax.vmap(Ray, in_axes=(None, 0, None))(ray_origin, ray_directions, ray_length)
    return rays

jit_create_rays = jax.jit(create_rays) #jitting is needed as rays are often generated at each instance of a loop


def create_space(x_min, x_max, y_min, y_max, torous, wall_array):
    """
    Create a space with the given parameters
    Args:
        - x_min: The minimum x value of the space
        - x_max: The maximum x value of the space
        - y_min: The minimum y value of the space
        - y_max: The maximum y value of the space
        - torous: Whether the space is torous or not
        - wall_array: The array of walls in the space in the form [[[wall1_begin(x,y)], [wall1_end(x,y)]], [[wall2_begin(x,y)], [wall2_end(x,y)], ...]]]
    Returns:
        A new space where wall is in the form of ([P1(x1,y1), P1(x2,y2), ..] [P2(x1,y1), P2(x2,y2), ..] ...)
    """
    x_wall_begins = jnp.array([wall_array[i][0][0] for i in range(len(wall_array))])
    y_wall_begins = jnp.array([wall_array[i][0][1] for i in range(len(wall_array))])
    x_wall_ends = jnp.array([wall_array[i][1][0] for i in range(len(wall_array))])
    y_wall_ends = jnp.array([wall_array[i][1][1] for i in range(len(wall_array))])
    wall_points_begins = jax.vmap(Point)(x_wall_begins, y_wall_begins)
    wall_points_ends = jax.vmap(Point)(x_wall_ends, y_wall_ends)
    walls = jax.vmap(Line)(wall_points_begins, wall_points_ends)
    return Space(x_min = x_min, x_max = x_max, y_min = y_min, y_max = y_max, torous = torous, walls = walls)


def is_walls_ray_collision(walls: Line, old_position:Point, new_position:Point):
    """
    Code to check if the particle collides with any of the walls
    code adapted from https://www.geeksforgeeks.org/check-if-two-given-line-segments-intersect/
    Args:
        - walls: The walls in the space
        - old_position: The old position of the particle
        - new_position: The new position of the particle
    Returns:
        - True if the particle collides with any of the walls, False otherwise
    """

    def is_wall_ray(wall, old_position, new_position):

        def is_on_segment(p, q, r):
            cond = jnp.logical_and(q.x <= jnp.maximum(p.x, r.x), 
                                   jnp.logical_and(q.x >= jnp.minimum(p.x, r.x), 
                                    jnp.logical_and(q.y <= jnp.maximum(p.y, r.y), q.y >= jnp.minimum(p.y, r.y))))
            return cond
    
        def get_orientation(p, q, r):
            val = (q.y - p.y) * (r.x - q.x) - (q.x - p.x) * (r.y - q.y)
            is_zero = jnp.isclose(val, 0.0, atol=1e-6)
            return jax.lax.select(is_zero, 0, jax.lax.select(val > 0, 1, 2))
        
        def is_intersect(wall, old_position, new_position):
            p1 = wall.p1
            q1 = wall.p2
            p2 = old_position
            q2 = new_position
            o1 = get_orientation(p1, q1, p2)
            o2 = get_orientation(p1, q1, q2)
            o3 = get_orientation(p2, q2, p1)
            o4 = get_orientation(p2, q2, q1)
            
            c1 = jnp.logical_and(o1 != o2, o3 != o4)
            c2 = jnp.logical_and(o1 == 0, is_on_segment(p1, p2, q1))
            c3 = jnp.logical_and(o2 == 0, is_on_segment(p1, q2, q1))
            c4 = jnp.logical_and(o3 == 0, is_on_segment(p2, p1, q2))
            c5 = jnp.logical_and(o4 == 0, is_on_segment(p2, q1, q2))
            return jnp.logical_or(jnp.logical_or(c1, c2), jnp.logical_or(c3, jnp.logical_or(c4, c5)))
        
        return is_intersect(wall, old_position, new_position)
    
    return jnp.any(jax.vmap(is_wall_ray, in_axes=(0, None, None))(walls, old_position, new_position)) # any rertuns true if any wall intersects with the particle

jit_is_walls_ray_collision = jax.jit(is_walls_ray_collision)



def get_ray_wall_collision(ray:Ray, wall:Line):
    """
    Step:
    - Get the distance of the collision of the ray with the walls along the ray
    - Def get_orientation(p, q, r): gets the orientation of the points p, q, r
    - Def get_line_intercept(p1, q1, p2, q2): gets the distance of the intersection point of the lines p1q1 and p2q2
    - Def is_on_segment(p, q, r): checks if the point q is on the segment pr
    Args:
        - ray: Ray, The ray to check for collision
        - walls: Line, The walls in the space
    Returns:
        The distance of the collision of the ray with the walls along the ray
    """

    def get_orientation(p, q, r):
        val = (q.y - p.y) * (r.x - q.x) - (q.x - p.x) * (r.y - q.y)
        is_zero = jnp.isclose(val, 0.0, atol=1e-6)
        return jax.lax.select(is_zero, 0, jax.lax.select(val > 0, 1, 2))
    
    def get_line_intercept(p1, q1, p2, q2):
        # converting points to arrays for cross product
        # p1 -> q1 is the ray, p2 -> q2 is the wall
        # algorithm adapted from https://stackoverflow.com/questions/20677795/how-do-i-compute-the-intersection-point-of-two-lines
        p1 = jnp.array([p1.x, p1.y])
        q1 = jnp.array([q1.x, q1.y])
        p2 = jnp.array([p2.x, p2.y])
        q2 = jnp.array([q2.x, q2.y])

        d1 = jnp.cross(q1 - p1, p2 - p1) # determinant of the ray with the (wall begin - ray begin)
        d2 = jnp.cross(q1 - p1, q2 - p1) # determinant of the ray with the (wall end - ray begin)
        denom = jax.lax.cond(d1 == d2, lambda _: 1e-6, lambda _: d1 - d2, None) # to avoid division by zero
        r = d1 / denom
        intersection_pt = p2 + r*(q2 - p2) # intersection point on the wall
        distance = jnp.linalg.norm(p1 - intersection_pt) # distance between the intersection point and the ray begin

        return distance
    
    def is_on_segment(p, q, r):
        # check if the point q is on the segment pr
        cond = jnp.logical_and(q.x <= jnp.maximum(p.x, r.x), 
                               jnp.logical_and(q.x >= jnp.minimum(p.x, r.x), 
                               jnp.logical_and(q.y <= jnp.maximum(p.y, r.y), q.y >= jnp.minimum(p.y, r.y))))
        return cond
    
    p1 = ray.origin
    q1 = Point(ray.origin.x + ray.direction.x*ray.length, ray.origin.y + ray.direction.y*ray.length)
    
    p2 = wall.p1
    q2 = wall.p2

    o_p1_q1_p2 = get_orientation(p1, q1, p2)
    o_p1_q1_q2 = get_orientation(p1, q1, q2)
    o_p2_q2_p1 = get_orientation(p2, q2, p1)
    o_p2_q2_q1 = get_orientation(p2, q2, q1)

    # the usual case
    c1 = jnp.logical_and(o_p1_q1_p2 != o_p1_q1_q2, o_p2_q2_p1 != o_p2_q2_q1) # if the orientations are different
    d1 = jax.lax.cond(c1, lambda _: get_line_intercept(p1, q1, p2, q2), lambda _: ray.length, None) # get the distance of the intersection point from the ray begin

    # if wall begin is on the ray
    c2 = jnp.logical_and(o_p1_q1_p2 == 0, is_on_segment(p1, p2, q1))
    d2 = jax.lax.cond(c2, lambda _: jnp.linalg.norm(jnp.array([p1.x, p1.y]) - jnp.array([p2.x, p2.y])), lambda _: ray.length, None)

    # if wall end is on the ray
    c3 = jnp.logical_and(o_p1_q1_q2 == 0, is_on_segment(p1, q2, q1))
    d3 = jax.lax.cond(c3, lambda _: jnp.linalg.norm(jnp.array([p1.x, p1.y]) - jnp.array([q2.x, q2.y])), lambda _: ray.length, None)

    # if ray begin is on the wall
    c4 = jnp.logical_and(o_p2_q2_p1 == 0, is_on_segment(p2, p1, q2))
    d4 = jax.lax.cond(c4, lambda _: 0.0, lambda _: ray.length, None)

    # return the minimum distance
    return jnp.min(jnp.array([d1, d2, d3, d4]))

jit_get_ray_wall_collision = jax.jit(get_ray_wall_collision)



def get_ray_circle_collision(ray:Ray, circle:Circle):
    """
    ray casting algorithm to check for collision between a ray and a circle, adapted from https://www.youtube.com/watch?v=ebzlMOw79Yw&ab_channel=MagellanicMath
    Args:
        - ray: Ray, The ray to check for collision
        - circle: Circle, The circle to check for collision
    Returns:
        The distance of the collision of the ray with the circle along the ray
    """
    circle_center = jnp.reshape(jnp.array([circle.center.x, circle.center.y]), (2,))
    ray_origin = jnp.reshape(jnp.array([ray.origin.x, ray.origin.y]), (2,))
    ray_direction = jnp.reshape(jnp.array([ray.direction.x, ray.direction.y]), (2,))

    s = ray_origin - circle_center
    b = jnp.dot(s, ray_direction)
    c = jnp.dot(s, s) - circle.radius**2
    h = b**2 - c
    h = jax.lax.cond(h < 0, lambda _: -1.0, lambda _: jnp.sqrt(h), None)
    t = jax.lax.cond(h >= 0, lambda _: -b - h, lambda _: ray.length, None)
    t = jax.lax.cond(t < 0, lambda _: ray.length, lambda _: t, None)
    return jnp.minimum(t, ray.length)

jit_get_ray_circle_collision = jax.jit(get_ray_circle_collision)

if __name__ == "__main__":
    agent_pos = (0.0, 0.0, 0.0)
    rays = jit_create_rays(agent_pos)
    

    wall_array = jnp.array([[[1.0,-1.0],[1.0,1.0]],[[1.0,1.0],[-1.0,1.0]],[[-1.0,1.0],[-1.0,-1.0]]]) # 3 sided wall
    space = create_space(-2.0, 2.0, -2.0, 2.0, True, wall_array)


    old_point = Point(0.0, 0.0)
    new_point = Point(2.0, 1.0)
    print(jit_is_walls_ray_collision(space.walls, old_point, new_point))

    old_point = Point(0.0, 0.0)
    new_point = Point(0.0, -2.0)
    print(jit_is_walls_ray_collision(space.walls, old_point, new_point))

    '''
    a_wall = jax.tree_util.tree_map(lambda x: x[0], space.walls)
    print(" a_wall start", a_wall.p1)
    print(" a_wall end", a_wall.p2)
    ray_wall_collisions = jax.vmap(jit_get_ray_wall_collision, in_axes=(0, None))(rays, a_wall)
    print(ray_wall_collisions)
    '''
    wall = Line(p1 = Point(1.0, -1.0), p2 = Point(1.0, 1.0))
    ray = Ray(origin = Point(0.0, 0.0), direction = Point(jnp.cos(0.0), jnp.sin(0.0)), length = 40.0)
    print(jit_get_ray_wall_collision(ray, wall))

    circle = Circle(center = Point(0.0, 0.0), radius = 1.0)
