import pygame, sys, random, math

pygame.init()
screen = pygame.display.set_mode((800,600))
pygame.display.set_caption("Realistic Rotational Physics Derived from Collisions")
clock = pygame.time.Clock()
font = pygame.font.SysFont(None, 20)

# Physics constants
GRAVITY = 500.0         # pixels/s^2
RESTITUTION = 0.8       # bounciness coefficient
DAMPING = 0.99          # linear damping per frame
ANGULAR_DAMPING = 0.99  # angular damping per frame

# Friction coefficient for collisions
FRICTION_COEFF = 0.5

# Modes
MODE_EDIT = "edit"
MODE_SIMULATE = "simulate"
mode = MODE_EDIT

# Global lists and undo/redo stacks
shapes = []
undo_stack = []
redo_stack = []

# State for drawing a new shape
drawing_new_shape = False
current_drawing_points = []  # absolute coordinates

# State for dragging a shape
dragging_shape = None
drag_offset = (0,0)
drag_start_offset = None

# Collision mesh grid square size (for visual oriented boxes)
square_size = 10

# Debug flag and counters
debug_mode = False
collision_checks = 0      # number of SAT axis tests
impulse_resolutions_count = 0  # count of impulse resolutions performed

def rotate_point(x, y, angle):
    """Rotate a point (x,y) by angle (in radians)."""
    cos_a = math.cos(angle)
    sin_a = math.sin(angle)
    return (x*cos_a - y*sin_a, x*sin_a + y*cos_a)

def compute_moment_of_inertia(local_poly, mass):
    """Approximate moment of inertia (using maximum distance from the local center)."""
    cx = sum(p[0] for p in local_poly)/len(local_poly)
    cy = sum(p[1] for p in local_poly)/len(local_poly)
    max_dist_sq = max((p[0]-cx)**2 + (p[1]-cy)**2 for p in local_poly)
    return mass * max_dist_sq / 2

def point_in_polygon(x, y, poly):
    """Ray–casting algorithm to test if (x,y) is inside poly."""
    num = len(poly)
    j = num - 1
    inside = False
    for i in range(num):
        xi, yi = poly[i]
        xj, yj = poly[j]
        if ((yi > y) != (yj > y)) and (x < (xj-xi)*(y-yi)/(yj-yi+1e-9)+xi):
            inside = not inside
        j = i
    return inside

def greedy_mesh(grid, rows, cols, square_size):
    """Combine adjacent True cells into larger rectangles (in local coordinates)."""
    used = [[False for _ in range(cols)] for _ in range(rows)]
    rects = []
    for row in range(rows):
        for col in range(cols):
            if grid[row][col] and not used[row][col]:
                width = 0
                while col+width < cols and grid[row][col+width] and not used[row][col+width]:
                    width += 1
                height = 0
                valid = True
                while valid and row+height < rows:
                    for i in range(width):
                        if not grid[row+height][col+i] or used[row+height][col+i]:
                            valid = False
                            break
                    if valid:
                        height += 1
                for r in range(row, row+height):
                    for c in range(col, col+width):
                        used[r][c] = True
                rects.append(pygame.Rect(col*square_size, row*square_size, width*square_size, height*square_size))
    return rects

def compute_local_collision_boxes(local_poly, square_size):
    """Compute grid cells inside local_poly and merge them using greedy meshing."""
    xs = [p[0] for p in local_poly]
    ys = [p[1] for p in local_poly]
    max_x = int(max(xs))
    max_y = int(max(ys))
    cols = (max_x // square_size) + 1
    rows = (max_y // square_size) + 1
    grid = [[False for _ in range(cols)] for _ in range(rows)]
    for row in range(rows):
        for col in range(cols):
            cx = col*square_size + square_size/2
            cy = row*square_size + square_size/2
            if point_in_polygon(cx, cy, local_poly):
                grid[row][col] = True
    rects = greedy_mesh(grid, rows, cols, square_size)
    return rects

def sat_collision_polygon(polyA, polyB):
    """
    Uses the Separating Axis Theorem (SAT) to check collision between two convex polygons.
    Returns (colliding, mtv_normal, penetration_depth).
    Increments the global collision_checks count for each axis test.
    """
    global collision_checks
    min_overlap = float('inf')
    mtv_axis = None
    for polygon in (polyA, polyB):
        for i in range(len(polygon)):
            p1 = polygon[i]
            p2 = polygon[(i+1)%len(polygon)]
            edge = (p2[0]-p1[0], p2[1]-p1[1])
            axis = (-edge[1], edge[0])
            length = math.hypot(axis[0], axis[1])
            if length == 0:
                continue
            axis = (axis[0]/length, axis[1]/length)
            collision_checks += 1
            minA, maxA = float('inf'), float('-inf')
            for p in polyA:
                proj = p[0]*axis[0] + p[1]*axis[1]
                minA = min(minA, proj)
                maxA = max(maxA, proj)
            minB, maxB = float('inf'), float('-inf')
            for p in polyB:
                proj = p[0]*axis[0] + p[1]*axis[1]
                minB = min(minB, proj)
                maxB = max(maxB, proj)
            overlap = min(maxA, maxB) - max(minA, minB)
            if overlap < 0:
                return (False, None, 0)
            if overlap < min_overlap:
                min_overlap = overlap
                mtv_axis = axis
                centerA = (sum(p[0] for p in polyA)/len(polyA), sum(p[1] for p in polyA)/len(polyA))
                centerB = (sum(p[0] for p in polyB)/len(polyB), sum(p[1] for p in polyB)/len(polyB))
                d = (centerB[0]-centerA[0], centerB[1]-centerA[1])
                if d[0]*mtv_axis[0] + d[1]*mtv_axis[1] < 0:
                    mtv_axis = (-mtv_axis[0], -mtv_axis[1])
    return (True, mtv_axis, min_overlap)

class Shape:
    def __init__(self, local_poly, offset):
        """
        local_poly: list of points (x,y) in local coordinates (with minimum at (0,0)).
        offset: (x,y) position in the window.
        """
        self.local_poly = local_poly
        self.offset = list(offset)
        self.collision_rects = compute_local_collision_boxes(self.local_poly, square_size)
        self.velocity = [0.0, 0.0]
        self.angular_velocity = 0.0  # starts at zero; rotation will be derived from collisions
        self.angle = 0.0
        self.color = (random.randint(50,255), random.randint(50,255), random.randint(50,255))
        area = abs(0.5 * sum(local_poly[i][0]*local_poly[(i+1)%len(local_poly)][1] - local_poly[(i+1)%len(local_poly)][0]*local_poly[i][1]
                           for i in range(len(local_poly))))
        self.mass = area if area > 0 else 1.0
        self.inertia = compute_moment_of_inertia(local_poly, self.mass)
        self.local_center = (sum(p[0] for p in local_poly)/len(local_poly),
                             sum(p[1] for p in local_poly)/len(local_poly))
        self.colliding_boxes = set()

    def get_absolute_polygon(self):
        """Return the rotated and translated polygon."""
        abs_poly = []
        for p in self.local_poly:
            rel = (p[0]-self.local_center[0], p[1]-self.local_center[1])
            rot = rotate_point(rel[0], rel[1], self.angle)
            abs_poly.append((rot[0] + self.offset[0] + self.local_center[0],
                             rot[1] + self.offset[1] + self.local_center[1]))
        return abs_poly

    def get_absolute_collision_polys(self):
        """Return oriented collision boxes (for visualization)."""
        polys = []
        for rect in self.collision_rects:
            corners = [(rect.x, rect.y),
                       (rect.x+rect.width, rect.y),
                       (rect.x+rect.width, rect.y+rect.height),
                       (rect.x, rect.y+rect.height)]
            transformed = []
            for pt in corners:
                rel = (pt[0]-self.local_center[0], pt[1]-self.local_center[1])
                rot = rotate_point(rel[0], rel[1], self.angle)
                transformed.append((rot[0] + self.offset[0] + self.local_center[0],
                                    rot[1] + self.offset[1] + self.local_center[1]))
            polys.append(transformed)
        return polys

    def get_bounding_rect(self):
        poly = self.get_absolute_polygon()
        xs = [p[0] for p in poly]
        ys = [p[1] for p in poly]
        return pygame.Rect(min(xs), min(ys), max(xs)-min(xs), max(ys)-min(ys))
    
    def get_center(self):
        """Return the global center (local center translated)."""
        return (self.offset[0] + self.local_center[0],
                self.offset[1] + self.local_center[1])
    
    def update(self, dt):
        # Apply gravity.
        self.velocity[1] += GRAVITY * dt
        # Update position.
        self.offset[0] += self.velocity[0] * dt
        self.offset[1] += self.velocity[1] * dt
        # Update rotation.
        self.angle += self.angular_velocity * dt
        # Bounce off window edges.
        br = self.get_bounding_rect()
        if br.left < 0:
            self.offset[0] -= br.left
            self.velocity[0] = -RESTITUTION * self.velocity[0]
        if br.right > screen.get_width():
            self.offset[0] -= (br.right - screen.get_width())
            self.velocity[0] = -RESTITUTION * self.velocity[0]
        if br.top < 0:
            self.offset[1] -= br.top
            self.velocity[1] = -RESTITUTION * self.velocity[1]
        if br.bottom > screen.get_height():
            self.offset[1] -= (br.bottom - screen.get_height())
            self.velocity[1] = -RESTITUTION * self.velocity[1]
        self.velocity[0] *= DAMPING
        self.velocity[1] *= DAMPING
        self.angular_velocity *= ANGULAR_DAMPING

    def draw(self, surface):
        pygame.draw.polygon(surface, self.color, self.get_absolute_polygon(), 2)
        polys = self.get_absolute_collision_polys()
        for idx, poly in enumerate(polys):
            col = (255,0,0) if idx in self.colliding_boxes else (0,0,255)
            pygame.draw.polygon(surface, col, poly, 1)

def resolve_collision_realistic(shapeA, shapeB, normal, penetration):
    """
    Applies positional correction and impulse–based collision resolution (with friction)
    so that rotation is derived from collision impulses.
    """
    global impulse_resolutions_count
    percent = 0.2
    slop = 0.01
    inv_massA = 1/shapeA.mass
    inv_massB = 1/shapeB.mass
    correction = (max(penetration - slop, 0) / (inv_massA + inv_massB)) * percent
    shapeA.offset[0] -= correction * inv_massA * normal[0]
    shapeA.offset[1] -= correction * inv_massA * normal[1]
    shapeB.offset[0] += correction * inv_massB * normal[0]
    shapeB.offset[1] += correction * inv_massB * normal[1]
    
    centerA = shapeA.get_center()
    centerB = shapeB.get_center()
    contact_point = ((centerA[0]+centerB[0])/2, (centerA[1]+centerB[1])/2)
    rA = (contact_point[0]-centerA[0], contact_point[1]-centerA[1])
    rB = (contact_point[0]-centerB[0], contact_point[1]-centerB[1])
    velA_contact = (shapeA.velocity[0] + -shapeA.angular_velocity * rA[1],
                    shapeA.velocity[1] + shapeA.angular_velocity * rA[0])
    velB_contact = (shapeB.velocity[0] + -shapeB.angular_velocity * rB[1],
                    shapeB.velocity[1] + shapeB.angular_velocity * rB[0])
    rv = (velB_contact[0]-velA_contact[0], velB_contact[1]-velA_contact[1])
    vel_along_normal = rv[0]*normal[0] + rv[1]*normal[1]
    if vel_along_normal > 0:
        return
    e = RESTITUTION
    rA_cross_n = rA[0]*normal[1] - rA[1]*normal[0]
    rB_cross_n = rB[0]*normal[1] - rB[1]*normal[0]
    inv_inertiaA = 1/shapeA.inertia if shapeA.inertia != 0 else 0
    inv_inertiaB = 1/shapeB.inertia if shapeB.inertia != 0 else 0
    denom = inv_massA + inv_massB + (rA_cross_n**2)*inv_inertiaA + (rB_cross_n**2)*inv_inertiaB
    if denom == 0:
        return
    j = -(1+e)*vel_along_normal/denom
    impulse = (j*normal[0], j*normal[1])
    shapeA.velocity[0] -= impulse[0]*inv_massA
    shapeA.velocity[1] -= impulse[1]*inv_massA
    shapeB.velocity[0] += impulse[0]*inv_massB
    shapeB.velocity[1] += impulse[1]*inv_massB
    shapeA.angular_velocity -= rA_cross_n * j * inv_inertiaA
    shapeB.angular_velocity += rB_cross_n * j * inv_inertiaB

    tangent = (rv[0]-vel_along_normal*normal[0],
               rv[1]-vel_along_normal*normal[1])
    t_length = math.hypot(tangent[0], tangent[1])
    if t_length != 0:
        tangent = (tangent[0]/t_length, tangent[1]/t_length)
        jt = - (rv[0]*tangent[0] + rv[1]*tangent[1]) / denom
        jt = max(-j*FRICTION_COEFF, min(jt, j*FRICTION_COEFF))
        friction_impulse = (jt*tangent[0], jt*tangent[1])
        shapeA.velocity[0] -= friction_impulse[0]*inv_massA
        shapeA.velocity[1] -= friction_impulse[1]*inv_massA
        shapeB.velocity[0] += friction_impulse[0]*inv_massB
        shapeB.velocity[1] += friction_impulse[1]*inv_massB
        shapeA.angular_velocity -= (rA[0]*tangent[1] - rA[1]*tangent[0]) * jt * inv_inertiaA
        shapeB.angular_velocity += (rB[0]*tangent[1] - rB[1]*tangent[0]) * jt * inv_inertiaB
    global impulse_resolutions_count
    impulse_resolutions_count += 1

class Button:
    def __init__(self, rect, text):
        self.rect = pygame.Rect(rect)
        self.text = text
    def draw(self, surface):
        pygame.draw.rect(surface, (100,100,100), self.rect)
        txt = font.render(self.text, True, (255,255,255))
        txt_rect = txt.get_rect(center=self.rect.center)
        surface.blit(txt, txt_rect)
    def is_clicked(self, pos):
        return self.rect.collidepoint(pos)

new_shape_btn = Button((10,10,100,30), "New Shape")
play_btn = Button((120,10,100,30), "Play")
undo_btn = Button((230,10,100,30), "Undo")
redo_btn = Button((340,10,100,30), "Redo")
edit_btn = Button((450,10,100,30), "Edit")

running = True
while running:
    dt = clock.tick(60)/1000.0
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_d:
                debug_mode = not debug_mode
        if mode == MODE_EDIT:
            if event.type == pygame.MOUSEBUTTONDOWN:
                pos = pygame.mouse.get_pos()
                if new_shape_btn.is_clicked(pos):
                    drawing_new_shape = True
                    current_drawing_points = []
                    dragging_shape = None
                elif play_btn.is_clicked(pos):
                    mode = MODE_SIMULATE
                    # Launch shapes with random linear velocities but zero angular velocity.
                    for shape in shapes:
                        shape.velocity = [random.choice([-200, -150, 150, 200]),
                                          random.choice([-200, -150, 150, 200])]
                        shape.angular_velocity = 0.0
                elif undo_btn.is_clicked(pos):
                    if undo_stack:
                        action = undo_stack.pop()
                        if action[0] == "add":
                            shape = action[1]
                            if shape in shapes:
                                shapes.remove(shape)
                                redo_stack.append(("add", shape))
                        elif action[0] == "move":
                            shape, old_offset, new_offset = action[1], action[2], action[3]
                            shape.offset = list(old_offset)
                            redo_stack.append(("move", shape, new_offset, old_offset))
                elif redo_btn.is_clicked(pos):
                    if redo_stack:
                        action = redo_stack.pop()
                        if action[0] == "add":
                            shape = action[1]
                            shapes.append(shape)
                            undo_stack.append(("add", shape))
                        elif action[0] == "move":
                            shape, old_offset, new_offset = action[1], action[2], action[3]
                            shape.offset = list(new_offset)
                            undo_stack.append(("move", shape, old_offset, new_offset))
                elif edit_btn.is_clicked(pos):
                    pass
                else:
                    if drawing_new_shape:
                        current_drawing_points.append(pos)
                    else:
                        for shape in reversed(shapes):
                            if point_in_polygon(pos[0], pos[1], shape.get_absolute_polygon()):
                                dragging_shape = shape
                                drag_offset = (pos[0]-shape.offset[0], pos[1]-shape.offset[1])
                                drag_start_offset = shape.offset.copy()
                                break
            elif event.type == pygame.MOUSEBUTTONUP:
                if dragging_shape:
                    if drag_start_offset != dragging_shape.offset:
                        undo_stack.append(("move", dragging_shape, drag_start_offset, dragging_shape.offset.copy()))
                        redo_stack.clear()
                    dragging_shape = None
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN and drawing_new_shape and len(current_drawing_points) >= 3:
                    xs = [p[0] for p in current_drawing_points]
                    ys = [p[1] for p in current_drawing_points]
                    min_x, min_y = min(xs), min(ys)
                    local_poly = [(p[0]-min_x, p[1]-min_y) for p in current_drawing_points]
                    new_shape = Shape(local_poly, (min_x, min_y))
                    shapes.append(new_shape)
                    undo_stack.append(("add", new_shape))
                    redo_stack.clear()
                    drawing_new_shape = False
                elif event.key == pygame.K_z:
                    if undo_stack:
                        action = undo_stack.pop()
                        if action[0] == "add":
                            shape = action[1]
                            if shape in shapes:
                                shapes.remove(shape)
                                redo_stack.append(("add", shape))
                        elif action[0] == "move":
                            shape, old_offset, new_offset = action[1], action[2], action[3]
                            shape.offset = list(old_offset)
                            redo_stack.append(("move", shape, new_offset, old_offset))
                elif event.key == pygame.K_y:
                    if redo_stack:
                        action = redo_stack.pop()
                        if action[0] == "add":
                            shape = action[1]
                            shapes.append(shape)
                            undo_stack.append(("add", shape))
                        elif action[0] == "move":
                            shape, old_offset, new_offset = action[1], action[2], action[3]
                            shape.offset = list(new_offset)
                            undo_stack.append(("move", shape, old_offset, new_offset))
        elif mode == MODE_SIMULATE:
            if event.type == pygame.MOUSEBUTTONDOWN:
                pos = pygame.mouse.get_pos()
                if edit_btn.is_clicked(pos):
                    mode = MODE_EDIT
                    for shape in shapes:
                        shape.velocity = [0.0, 0.0]
                        shape.angular_velocity = 0.0
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False
    if mode == MODE_EDIT and dragging_shape:
        mpos = pygame.mouse.get_pos()
        dragging_shape.offset[0] = mpos[0]-drag_offset[0]
        dragging_shape.offset[1] = mpos[1]-drag_offset[1]
    if mode == MODE_SIMULATE:
        for shape in shapes:
            shape.update(dt)
            shape.colliding_boxes = set()
        collision_checks = 0
        impulse_resolutions_count = 0
        for i in range(len(shapes)):
            for j in range(i+1, len(shapes)):
                shapeA = shapes[i]
                shapeB = shapes[j]
                polyA = shapeA.get_absolute_polygon()
                polyB = shapeB.get_absolute_polygon()
                colliding, normal, penetration = sat_collision_polygon(polyA, polyB)
                if colliding:
                    resolve_collision_realistic(shapeA, shapeB, normal, penetration)
    screen.fill((30,30,30))
    new_shape_btn.draw(screen)
    play_btn.draw(screen)
    undo_btn.draw(screen)
    redo_btn.draw(screen)
    edit_btn.draw(screen)
    if drawing_new_shape:
        if len(current_drawing_points) > 1:
            pygame.draw.lines(screen, (255,255,255), False, current_drawing_points, 2)
        for p in current_drawing_points:
            pygame.draw.circle(screen, (255,0,0), p, 4)
    for shape in shapes:
        shape.draw(screen)
    if mode == MODE_SIMULATE:
        sim_text = font.render("Simulation Mode - Click Edit to return", True, (255,255,255))
        screen.blit(sim_text, (10,50))
    else:
        edit_text = font.render("Edit Mode - Draw/Drag shapes. Enter to finish shape.", True, (255,255,255))
        screen.blit(edit_text, (10,50))
    if debug_mode:
        debug_lines = []
        debug_lines.append(f"Mode: {mode}")
        debug_lines.append(f"Shapes: {len(shapes)}")
        total_boxes = sum(len(s.get_absolute_collision_polys()) for s in shapes)
        debug_lines.append(f"Oriented Collision Boxes: {total_boxes}")
        debug_lines.append(f"SAT Axis Tests: {collision_checks}")
        debug_lines.append(f"Impulse Resolutions: {impulse_resolutions_count}")
        debug_lines.append(f"FPS: {int(clock.get_fps())}")
        for i, shape in enumerate(shapes):
            debug_lines.append(f"Shape {i}: angle={math.degrees(shape.angle):.1f}°, ang_vel={shape.angular_velocity:.2f}")
        for idx, line in enumerate(debug_lines):
            txt = font.render(line, True, (255,255,0))
            screen.blit(txt, (10, 90+idx*18))
    pygame.display.flip()
pygame.quit()
sys.exit()