import pygame
import sys
import math
import json
import numpy as np

class MapGenerator:
    def __init__(self, width, height, difficulty, goals):
        self.surface = pygame.Surface((width, height))
        self.width = width
        self.height = height
        self.difficulty = difficulty
        self.goals = []
        self.num_goals = goals if isinstance(goals, int) else len(goals) if goals else 12
        self.curve_pts = []  # Added to store curve points

    def update(self):
        self.surface.fill((255, 255, 255))

        if self.difficulty == 1:
            border = 50
            num_ctrl = 4
            x_coords = [i * (self.width/(num_ctrl-1)) for i in range(num_ctrl)]
            x_coords[0], x_coords[-1] = border*2, self.width-border*2
            y_coords = np.random.randint(border*4, self.height-border*4, num_ctrl)
            ctrl_pts = list(zip(x_coords, y_coords))
            curve_pts = self.compute_catmull_rom(ctrl_pts, 800)
            road_width = 100
        
        if self.difficulty == 2:
            # Define constants
            border = 150
            cx = self.width / 2
            cy = self.height / 2

            r = min(self.width / 2 - border, self.height / 2 - border)
            a = r
            b = r

            # Generate control points for open loop (exclude last few points to create gap)
            N = 25
            gap_points = 3  # Remove last 3 control points to create visible gap
            angles = [2 * np.pi * i / N for i in range(N - gap_points)]
            control_pts = []

            # Perlin-like noise parameters for soft, organic wobble
            octaves = 1
            frequencies = [2.5] # Multiple octaves for perlin-like effect
            amplitudes = [16]  # Decreasing amplitudes for each octave
            
            # Random phases for each octave
            phi_x = [[np.random.uniform(0, 2 * np.pi) for _ in range(octaves)]]
            phi_y = [[np.random.uniform(0, 2 * np.pi) for _ in range(octaves)]]

            # Generate noisy control points with perlin-like noise
            for i, angle in enumerate(angles):
                x = cx + a * np.cos(angle)
                y = cy + b * np.sin(angle)
                
                # Apply multi-octave noise (perlin-like)
                dx = 0
                dy = 0
                for j in range(octaves):
                    # Use both angle and position index for variation
                    noise_input = frequencies[j] * angle + phi_x[0][j]
                    dx += amplitudes[j] * np.sin(noise_input)
                    
                    noise_input_y = frequencies[j] * angle + phi_y[0][j]
                    dy += amplitudes[j] * np.cos(noise_input_y)  # Use cos for y to decorrelate
                
                x += dx
                y += dy
                control_pts.append((x, y))

            # Randomly invert the list of control points
            if np.random.rand() < 0.5:
                control_pts = control_pts[::-1]

            # Compute the curve points with an open loop
            curve_pts = self.compute_catmull_rom(control_pts, 800)
            
            # Remove additional points at the end to ensure visible gap
            if len(curve_pts) > 100:
                curve_pts = curve_pts[:-50]  # Remove last 50 curve points for clear gap
            road_width = 100
            
        if self.difficulty == 3:
            border = 50
            num_ctrl = 5
            x_coords = [i * (self.width/(num_ctrl-1)) for i in range(num_ctrl)]
            x_coords[0], x_coords[-1] = border*2, self.width-border*2
            
            # Generate y coordinates with alternating high/low pattern
            y_coords = []
            
            # Define high and low ranges
            high_range = (self.height//2 + 2*border, self.height - int(border*5.5))
            low_range = (int(border*5.5), self.height//2 - 2*border)
            
            # Generate first point randomly
            first_y = np.random.randint(int(border*5.5), self.height-int(border*5.5))
            y_coords.append(first_y)
            
            # Determine if first point is high or low
            is_first_high = first_y > self.height // 2
            
            # Generate remaining points alternating between high and low
            for i in range(1, num_ctrl):
                if is_first_high:
                    # If previous point was high, make this one low
                    if (i - 1) % 2 == 0:  # Previous was high (first point)
                        y = np.random.randint(low_range[0], low_range[1])
                    else:  # Previous was low
                        y = np.random.randint(high_range[0], high_range[1])
                else:
                    # If previous point was low, make this one high
                    if (i - 1) % 2 == 0:  # Previous was low (first point)
                        y = np.random.randint(high_range[0], high_range[1])
                    else:  # Previous was high
                        y = np.random.randint(low_range[0], low_range[1])
                
                y_coords.append(y)
                
            ctrl_pts = list(zip(x_coords, y_coords))
            
            # 50% chance to flip the control points order
            if np.random.rand() < 0.5:
                ctrl_pts = ctrl_pts[::-1]
            
            curve_pts = self.compute_catmull_rom(ctrl_pts, 800)
            road_width = 100

        self.curve_pts = curve_pts  # Store curve points

        # Create goals     
        goals_pts = self.create_goals(curve_pts)
        self.goals = goals_pts  

        # Draw roads as connected segments
        self.draw_road_segments(curve_pts, road_width)

        # Draw goals
        for goal_index in goals_pts: 
            pygame.draw.line(self.surface, (0,0,255), goal_index[0], goal_index[1], width=2)   

        # Get car start
        if len(curve_pts) >= 2:
            start_pos = curve_pts[0]
            xs, ys = start_pos
            xe, ye = curve_pts[1]
            x = xe - xs
            y = ye - ys
            start_dir = (x, y)
        else:
            start_pos = (100, self.height//2)
            start_dir = (1, 0)
            
        return pygame.math.Vector2(start_pos), pygame.math.Vector2(start_dir)

    def draw_road_segments(self, points, road_width):
        if len(points) < 2:
            return
        for i in range(len(points) - 1):
            start_pos = (int(points[i][0]), int(points[i][1]))
            end_pos = (int(points[i+1][0]), int(points[i+1][1]))
            pygame.draw.line(self.surface, (0, 0, 0), start_pos, end_pos, width=road_width)
        cap_radius = road_width // 2
        for point in points:
            pygame.draw.circle(self.surface, (0, 0, 0), (int(point[0]), int(point[1])), cap_radius)

    def compute_catmull_rom(self, control_points, steps_per_segment=20):
        try:
            is_loop = control_points[0] == control_points[-1]
            if hasattr(is_loop, '__len__') and len(is_loop) > 1:
                is_loop = np.allclose(control_points[0], control_points[-1])
        except:
            first = np.array(control_points[0])
            last = np.array(control_points[-1])
            is_loop = np.allclose(first, last)
        
        pts = control_points[:-1] if is_loop else control_points
        n = len(pts)
        points = []

        def get(idx):
            if is_loop:
                return pts[idx % n]
            else:
                return pts[min(max(idx, 0), n - 1)]

        for i in range(n if is_loop else n - 1):
            p0 = get(i - 1)
            p1 = get(i)
            p2 = get(i + 1)
            p3 = get(i + 2)
            for j in range(steps_per_segment):
                t = j / steps_per_segment
                t2 = t * t
                t3 = t2 * t
                f1 = -0.5 * t3 + t2 - 0.5 * t
                f2 = 1.5 * t3 - 2.5 * t2 + 1.0
                f3 = -1.5 * t3 + 2.0 * t2 + 0.5 * t
                f4 = 0.5 * t3 - 0.5 * t2
                x = f1 * p0[0] + f2 * p1[0] + f3 * p2[0] + f4 * p3[0]
                y = f1 * p0[1] + f2 * p1[1] + f3 * p2[1] + f4 * p3[1]
                points.append((x, y))
        if not is_loop:
            points.append(tuple(pts[-1]))
        return points
    
    def resample_by_distance(self, points, desired_spacing):
        if len(points) < 2:
            return points
        dists = [0.0]
        for i in range(1, len(points)):
            prev, curr = np.array(points[i-1]), np.array(points[i])
            dist = float(np.linalg.norm(curr - prev))
            dists.append(dists[-1] + dist)
        total_len = dists[-1]
        if total_len == 0:
            return points
        n_points = int(total_len // desired_spacing)
        new_pts = []
        for k in range(n_points+1):
            target = k * desired_spacing
            for j in range(1, len(dists)):
                if dists[j] >= target:
                    ratio = (target - dists[j-1])/(dists[j] - dists[j-1])
                    A, B = np.array(points[j-1]), np.array(points[j])
                    P = A + ratio*(B - A)
                    new_pts.append(tuple(P))
                    break
        return new_pts if new_pts else points

    def create_goals(self, pts, goal_number=None):
        n_pts = len(pts)
        if goal_number is not None:
            num_goals = goal_number
        else:
            num_goals = self.num_goals
        if n_pts == 0 or num_goals == 0:
            return []
        actual_goals = min(num_goals, n_pts // 10)
        if actual_goals == 0:
            actual_goals = 1
        lines = []
        last_dir = None
        for k in range(actual_goals):
            if actual_goals == 1:
                idx = n_pts // 2
            else:
                idx = int((k * (n_pts - 1)) / (actual_goals - 1))
            idx = max(0, min(idx, n_pts - 1))
            idx2 = idx
            look_ahead = 5
            if idx + look_ahead < n_pts:
                idx2 = idx + look_ahead
            elif idx >= look_ahead:
                idx2 = idx - look_ahead
            elif idx + 1 < n_pts:
                idx2 = idx + 1
            elif idx > 0:
                idx2 = idx - 1
            x, y = pts[idx]
            if idx2 != idx:
                x1, y1 = pts[idx2]
                dir_vec = pygame.math.Vector2(x1 - x, y1 - y)
                if idx2 < idx:
                    dir_vec *= -1
            else:
                dir_vec = last_dir if last_dir else pygame.math.Vector2(0, -1)
            if dir_vec.length_squared() < 1e-8:
                dir_vec = pygame.math.Vector2(0, -1)
            else:
                dir_vec = dir_vec.normalize()
                if last_dir and dir_vec.dot(last_dir) < -0.5:
                    dir_vec *= -1
            last_dir = dir_vec
            normal = dir_vec.rotate(90)
            p0 = pygame.math.Vector2(x, y) + normal * 50
            p1 = pygame.math.Vector2(x, y) - normal * 50
            lines.append((p0, p1))
        return lines

    def draw(self, surface):
        surface.blit(self.surface, (0,0))

class Car:
    def __init__(self, image, pos, max_velocity=275, turn_speed=77.5):
        self.original_image = image.convert_alpha()
        self.image = self.original_image
        self.pos = pygame.math.Vector2(pos)
        self.velocity = 0.0
        self.max_velocity = max_velocity
        self.direction = pygame.math.Vector2(1, 0)
        self.turn_speed = turn_speed
        self.angle = 0.0
        self.is_braking = False
        self.acceleration = 55.0
        self.deceleration = 16.0
        self.brake_deceleration = 55.0
        self.min_velocity_threshold = 5.0

    def update(self, dt, keys=None, steer_factor=1.0):
        if keys is None:
            pressed = pygame.key.get_pressed()
            keys = {
                'w': pressed[pygame.K_w],
                'a': pressed[pygame.K_a],
                's': pressed[pygame.K_s],
                'd': pressed[pygame.K_d]
            }
        self.is_braking = False
        if keys.get('w', False):
            acceleration_multiplier = 1.0
            if self.velocity > 50:
                acceleration_multiplier = 1.2
            elif self.velocity < 20:
                acceleration_multiplier = 1.5
            if self.velocity < self.max_velocity:
                self.velocity += self.acceleration * acceleration_multiplier * dt
                if self.velocity > self.max_velocity:
                    self.velocity = self.max_velocity
        if keys.get('s', False):
            self.is_braking = True
            if self.velocity > 0:
                brake_multiplier = 1.0 + (self.velocity / self.max_velocity) * 0.5
                self.velocity -= self.brake_deceleration * brake_multiplier * dt
                if self.velocity < 0:
                    self.velocity = 0
        else:
            if not keys.get('w', False) and self.velocity > 0:
                decel_rate = self.deceleration
                if self.velocity > 100:
                    decel_rate *= 0.7
                self.velocity -= decel_rate * dt
                if self.velocity < self.min_velocity_threshold:
                    self.velocity = max(0, self.velocity - 2.0 * dt)
        if self.velocity > 0:
            self.pos += self.direction.normalize() * self.velocity * dt
        if self.velocity > 10:
            angular_velocity = 0
            speed_factor = min(1.0, self.velocity / 200.0)
            effective_turn_speed = self.turn_speed * (0.5 + 0.5 * speed_factor)
            if keys.get('a', False):
                angular_velocity = -effective_turn_speed * steer_factor
            elif keys.get('d', False):
                angular_velocity = effective_turn_speed * steer_factor
            if angular_velocity != 0:
                angle_change = angular_velocity * dt
                target_direction = self.direction.rotate(angle_change)
                base_drift = 0.15
                velocity_factor = min(1.0, self.velocity / self.max_velocity)
                drift_factor = base_drift * (0.5 + 0.5 * velocity_factor)
                self.direction = (self.direction * drift_factor + 
                                target_direction * (1 - drift_factor)).normalize()
                self.angle = -math.degrees(math.atan2(self.direction.y, self.direction.x))
        self.image = pygame.transform.rotate(self.original_image, self.angle)

    def draw(self, surface):
        rect = self.image.get_rect(center=self.pos)
        surface.blit(self.image, rect.topleft)
    
    def reset(self, pos, direction):
        self.pos = pygame.math.Vector2(pos)
        self.direction = pygame.math.Vector2(direction).normalize()
        self.velocity = 0.0
        self.angle = -math.degrees(math.atan2(self.direction.y, self.direction.x))

class RayCasting:
    def __init__(self, car_pos, car_dir, track_surf,
                 ray_angles, ray_offsets, max_distance=1600, step_pixels=1, draw_rays=True):
        self.car_pos = pygame.math.Vector2(car_pos)
        self.car_dir = pygame.math.Vector2(car_dir).normalize()
        self.track_surf = track_surf
        self.ray_angles = list(ray_angles)
        self.ray_offsets = list(ray_offsets)
        self.max_distance = max_distance
        self.step_pixels = step_pixels
        self.last_collisions = [self.car_pos.copy()] * len(self.ray_angles)
        self.ray_distances = [0.0] * len(self.ray_angles)
        self.draw_rays = draw_rays

    def cast_ray(self, start_pos, direction):
        pos = pygame.math.Vector2(start_pos)
        dir_norm = pygame.math.Vector2(direction).normalize()
        inv = max(abs(dir_norm.x), abs(dir_norm.y))
        step_size = (1.0 / inv) if inv != 0 else 1.0
        traveled = 0.0
        w, h = self.track_surf.get_size()
        while traveled < self.max_distance:
            ix, iy = int(pos.x), int(pos.y)
            if 0 <= ix < w and 0 <= iy < h:
                color = self.track_surf.get_at((ix, iy))
                if color[:3] == (255, 255, 255):
                    return pos, traveled
            pos += dir_norm * (step_size * self.step_pixels)
            traveled += step_size * self.step_pixels
        return pos, traveled
    
    def draw_ray(self, start, end, color=(255, 255, 0)):
        pygame.draw.line(self.track_surf, color,
                         (int(start.x), int(start.y)),
                         (int(end.x), int(end.y)),
                         width=2)

    def update(self, car_pos, car_dir):
        self.car_pos = pygame.math.Vector2(car_pos)
        self.car_dir = pygame.math.Vector2(car_dir).normalize()
        self.last_collisions = []
        self.ray_distances = []
        any_origin_in_white = False
        for angle, offset in zip(self.ray_angles, self.ray_offsets):
            ray_dir = self.car_dir.rotate(angle).normalize()
            origin = self.car_pos + ray_dir * offset
            ix, iy = int(origin.x), int(origin.y)
            w, h = self.track_surf.get_size()
            if 0 <= ix < w and 0 <= iy < h:
                if self.track_surf.get_at((ix, iy))[:3] == (255, 255, 255):
                    any_origin_in_white = True
                    self.last_collisions.append(origin)
                    self.ray_distances.append(0.0)
                    continue
            hit_pt, distance = self.cast_ray(origin, ray_dir)
            self.last_collisions.append(hit_pt)
            self.ray_distances.append(distance)
            if self.draw_rays:
                self.draw_ray(origin, hit_pt)
        return not any_origin_in_white

    def get_ray_distances(self):
        return self.ray_distances.copy()
    
    def has_zero_length_ray(self):
        return any(distance <= 0.0 for distance in self.ray_distances)

if __name__ == "__main__":
    pygame.init()
    screen = pygame.display.set_mode((1600, 900))
    clock = pygame.time.Clock()
    car_surf = pygame.image.load("assets/car_transparent.png").convert_alpha()
    w, h = car_surf.get_size()
    car_surf = pygame.transform.scale(car_surf, (w//18, h//18))
    car = Car(car_surf, pos=(400, 300))
    map = MapGenerator(1600, 900, difficulty=3, goals=22)
    startpos, startdir = map.update()
    car.direction = startdir
    car.pos = startpos
    car.velocity = 0
    ray_angles = [-60, -30, -15, 0, 15, 30, 60]
    ray_offsets = [12, 22, 27, 26, 27, 22, 12]
    rays = RayCasting(car.pos, car.direction, screen, ray_angles, ray_offsets)
    car.draw(screen)
    frames = 0
    while True:
        frames += 1 
        dt = clock.tick(90) / 1000
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
        rays.update(car.pos, car.direction)
        if frames % 3000 == 0 or rays.has_zero_length_ray():
            dt = 0
            startpos, startdir = map.update()
            car.direction = startdir
            car.pos = startpos
            car.velocity = 0 
        screen.fill((255, 255, 255))
        map.draw(screen)
        car.update(dt)
        car.draw(screen)
        rays.update(car.pos, car.direction)
        pygame.display.flip()