import math
import numpy as np
import random
import imageio
from PIL import Image, ImageDraw
import io
import gradio as gr

# Parameters
SPACE_SIZE = 600
WORLD_UNITS = 20
SCALE = SPACE_SIZE / WORLD_UNITS
OBSTACLE_RADIUS = 0.5
ROBOT_RADIUS = 0.3
STEP_SIZE = 0.3
TURN_ANGLE = 30
SCAN_ANGLE = 60
SCAN_RESOLUTION = 30
SCAN_RANGE_DEFAULT = 6.0
OBSTACLE_COUNT = 10

# Utilities
def to_canvas_coords(x, y):
    cx = int((x + WORLD_UNITS / 2) * SCALE)
    cy = int((WORLD_UNITS / 2 - y) * SCALE)
    return cx, cy

def generate_obstacles(count):
    obs = []
    for _ in range(count):
        while True:
            x = random.uniform(-WORLD_UNITS / 2 + OBSTACLE_RADIUS, WORLD_UNITS / 2 - OBSTACLE_RADIUS)
            y = random.uniform(-WORLD_UNITS / 2 + OBSTACLE_RADIUS, WORLD_UNITS / 2 - OBSTACLE_RADIUS)
            if math.hypot(x, y) > 2:
                obs.append((x, y))
                break
    return obs

def check_collision(x, y, obstacles):
    if not (-WORLD_UNITS / 2 + ROBOT_RADIUS <= x <= WORLD_UNITS / 2 - ROBOT_RADIUS):
        return True
    if not (-WORLD_UNITS / 2 + ROBOT_RADIUS <= y <= WORLD_UNITS / 2 - ROBOT_RADIUS):
        return True
    for ox, oy in obstacles:
        if math.hypot(ox - x, oy - y) <= OBSTACLE_RADIUS + ROBOT_RADIUS:
            return True
    return False

def scan_environment(robot_pose, obstacles, scan_range):
    points = []
    angles = np.linspace(-SCAN_ANGLE / 2, SCAN_ANGLE / 2, SCAN_RESOLUTION) + robot_pose['angle']
    for a in angles:
        rad = math.radians(a)
        for d in np.linspace(0, scan_range, 100):
            x = robot_pose['x'] + d * math.cos(rad)
            y = robot_pose['y'] + d * math.sin(rad)
            if any(math.hypot(ox - x, oy - y) <= OBSTACLE_RADIUS for ox, oy in obstacles):
                points.append((x, y))
                break
    return points

def move_robot(robot_pose, obstacles):
    rad = math.radians(robot_pose['angle'])
    for d in np.linspace(0, SCAN_RANGE_DEFAULT, 30):
        x = robot_pose['x'] + d * math.cos(rad)
        y = robot_pose['y'] + d * math.sin(rad)
        if check_collision(x, y, obstacles):
            robot_pose['angle'] -= TURN_ANGLE
            return robot_pose, False
    new_x = robot_pose['x'] + STEP_SIZE * math.cos(rad)
    new_y = robot_pose['y'] + STEP_SIZE * math.sin(rad)
    if not check_collision(new_x, new_y, obstacles):
        robot_pose['x'] = new_x
        robot_pose['y'] = new_y
        return robot_pose, True
    else:
        robot_pose['angle'] -= TURN_ANGLE
        return robot_pose, False

def draw_environment(robot_pose, trajectory, obstacles):
    img = Image.new("RGB", (SPACE_SIZE, SPACE_SIZE), "white")
    draw = ImageDraw.Draw(img)
    
    # Grid
    for i in range(int(WORLD_UNITS) + 1):
        x = int(i * SCALE)
        draw.line([(x, 0), (x, SPACE_SIZE)], fill=(230, 230, 230))
        draw.line([(0, x), (SPACE_SIZE, x)], fill=(230, 230, 230))
    
    # Obstacles
    for ox, oy in obstacles:
        cx, cy = to_canvas_coords(ox, oy)
        r = int(OBSTACLE_RADIUS * SCALE)
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill="gray")

    # Trajectory
    if len(trajectory) > 1:
        points = [to_canvas_coords(x, y) for x, y in trajectory]
        draw.line(points, fill="blue", width=2)

    # Robot
    rx, ry = to_canvas_coords(robot_pose['x'], robot_pose['y'])
    r = int(ROBOT_RADIUS * SCALE)
    draw.ellipse([rx - r, ry - r, rx + r, ry + r], fill="red")

    # Heading
    dx = math.cos(math.radians(robot_pose['angle'])) * ROBOT_RADIUS * 2 * SCALE
    dy = math.sin(math.radians(robot_pose['angle'])) * ROBOT_RADIUS * 2 * SCALE
    draw.line([rx, ry, rx + dx, ry - dy], fill="darkred", width=2)
    return img

def draw_slam_map(robot_pose, slam_points, scan_range):
    img = Image.new("RGB", (SPACE_SIZE, SPACE_SIZE), "white")
    draw = ImageDraw.Draw(img)

    # Grid
    for i in range(int(WORLD_UNITS) + 1):
        x = int(i * SCALE)
        draw.line([(x, 0), (x, SPACE_SIZE)], fill=(230, 230, 230))
        draw.line([(0, x), (SPACE_SIZE, x)], fill=(230, 230, 230))

    # Laser scan lines
    angles = np.linspace(-SCAN_ANGLE / 2, SCAN_ANGLE / 2, SCAN_RESOLUTION) + robot_pose['angle']
    for a in angles:
        rad = math.radians(a)
        dist = scan_range
        for px, py in slam_points:
            vx, vy = px - robot_pose['x'], py - robot_pose['y']
            point_dist = math.hypot(vx, vy)
            point_angle = math.degrees(math.atan2(vy, vx))
            angle_diff = abs((point_angle - a + 180) % 360 - 180)
            if angle_diff < SCAN_ANGLE / SCAN_RESOLUTION:
                if point_dist < dist:
                    dist = point_dist
        x_end = robot_pose['x'] + dist * math.cos(rad)
        y_end = robot_pose['y'] + dist * math.sin(rad)
        cx_start, cy_start = to_canvas_coords(robot_pose['x'], robot_pose['y'])
        cx_end, cy_end = to_canvas_coords(x_end, y_end)
        draw.line([cx_start, cy_start, cx_end, cy_end], fill="green")

    # SLAM points
    for px, py in slam_points:
        cx, cy = to_canvas_coords(px, py)
        draw.ellipse([cx - 2, cy - 2, cx + 2, cy + 2], fill="black")

    # Robot
    rx, ry = to_canvas_coords(robot_pose['x'], robot_pose['y'])
    r = int(ROBOT_RADIUS * SCALE)
    draw.ellipse([rx - r, ry - r, rx + r, ry + r], fill="red")

    # Heading
    dx = math.cos(math.radians(robot_pose['angle'])) * ROBOT_RADIUS * 2 * SCALE
    dy = math.sin(math.radians(robot_pose['angle'])) * ROBOT_RADIUS * 2 * SCALE
    draw.line([rx, ry, rx + dx, ry - dy], fill="red", width=2)
    return img

# Simulation State
class SimulationState:
    def __init__(self):
        self.robot_pose = {'x': 0.0, 'y': 0.0, 'angle': 0.0}
        self.trajectory = []
        self.slam_points = []
        self.obstacles = generate_obstacles(OBSTACLE_COUNT)
        self.recording = False
        self.video_frames = []

state = SimulationState()

# Handlers
def reset_simulation(obstacle_count):
    state.robot_pose = {'x': 0.0, 'y': 0.0, 'angle': 0.0}
    state.trajectory = []
    state.slam_points = []
    state.obstacles = generate_obstacles(obstacle_count)
    state.recording = False
    state.video_frames = []
    env_img = draw_environment(state.robot_pose, state.trajectory, state.obstacles)
    slam_img = draw_slam_map(state.robot_pose, state.slam_points, SCAN_RANGE_DEFAULT)
    return env_img, slam_img, "Simulation reset"

def step_simulation(scan_range):
    pose_before = state.robot_pose.copy()
    new_pose, moved = move_robot(state.robot_pose, state.obstacles)
    if moved:
        state.trajectory.append((new_pose['x'], new_pose['y']))

    new_points = scan_environment(state.robot_pose, state.obstacles, scan_range)
    for pt in new_points:
        if all(math.hypot(pt[0] - p[0], pt[1] - p[1]) > 0.1 for p in state.slam_points):
            state.slam_points.append(pt)

    env_img = draw_environment(state.robot_pose, state.trajectory, state.obstacles)
    slam_img = draw_slam_map(state.robot_pose, state.slam_points, scan_range)

    if state.recording:
        combined = Image.new('RGB', (SPACE_SIZE * 2, SPACE_SIZE))
        combined.paste(env_img, (0, 0))
        combined.paste(slam_img, (SPACE_SIZE, 0))
        state.video_frames.append(np.array(combined))

    return env_img, slam_img, f"Robot moved: {moved}, Position: ({state.robot_pose['x']:.2f}, {state.robot_pose['y']:.2f})"

def toggle_recording(record):
    state.recording = record
    return f"Recording {'started' if record else 'stopped'}."

def export_video():
    if not state.video_frames:
        return None, "No frames recorded."
    buffer = io.BytesIO()
    imageio.mimsave(buffer, state.video_frames, format="GIF", duration=0.1)
    buffer.seek(0)
    return buffer, "Video exported."


# Default obstacle count
OBSTACLE_COUNT = 10

with gr.Blocks() as demo:
    gr.Markdown("# 2D SLAM Simulation")
    with gr.Row():
        env_img = gr.Image(label="Environment View", height=SPACE_SIZE, width=SPACE_SIZE)
        slam_img = gr.Image(label="SLAM Map View", height=SPACE_SIZE, width=SPACE_SIZE)
    with gr.Row():
        scan_range_slider = gr.Slider(1.0, 10.0, value=SCAN_RANGE_DEFAULT, step=0.1, label="Laser Scan Range")
        obstacle_count_slider = gr.Slider(0, 30, value=OBSTACLE_COUNT, step=1, label="Obstacle Count")
    with gr.Row():
        step_btn = gr.Button("Step")
        reset_btn = gr.Button("Reset")
    record_toggle = gr.Checkbox(label="Record Video")
    status_text = gr.Textbox(value="Simulation ready.", interactive=False)

    reset_btn.click(reset_simulation, inputs=obstacle_count_slider, outputs=[env_img, slam_img, status_text])
    step_btn.click(step_simulation, inputs=scan_range_slider, outputs=[env_img, slam_img, status_text])
    record_toggle.change(toggle_recording, inputs=record_toggle, outputs=None)

demo.launch()