from dataclasses import dataclass
from itertools import combinations

import pygame as pg

from tools import debug
from collider.system import intersect 
from collider.types import *
from transform import Transform

@dataclass
class RigidBody:
    transform: Transform
    collider: BaseCollider
    velocity: pg.Vector2
    angular_velocity: float = 0
    mass: float = 1.0
    restitution: float = 0.5
    coef_friction: float = 0.6

    def apply_impulse(self, j: pg.Vector2, point: pg.Vector2) -> None:
        moment_arm = point - self.transform.global_position
        self.velocity += j * self.inv_mass
        self.angular_velocity += moment_arm.cross(j) * self.inv_moment_of_inertia

    @property
    def inv_mass(self) -> float:
        return 0.0 if self.mass == 0.0 else 1/self.mass
    
    @property
    def moment_of_inertia(self) -> float:
        return 0.0 if self.mass == 0.0 else self.collider.moment_of_inertia(self.mass)
    
    @property
    def inv_moment_of_inertia(self) -> float:
        return 0.0 if self.mass == 0.0 else 1/self.moment_of_inertia

class PhysicsSystem:

    def __init__(self, gravity: int = 250):
        self.bodies: list[RigidBody] = []
        self.gravity = pg.Vector2(0,gravity)
    
    def add_body(self, body: RigidBody) -> None:
        self.bodies.append(body)

    def update(self, dt: float):
        g = self.gravity * dt
        for body in self.bodies:
            if body.mass != 0.0:
                body.velocity += g
            body.transform.position += dt * body.velocity
            body.transform.rotation += dt * body.angular_velocity
            debug.draw_collider(body.collider, body.transform)

        for a, b in combinations(self.bodies, 2):
            if collision := intersect(a.collider, b.collider, a.transform, b.transform):
                self.resolve_collision(a, b, collision)
        
    def resolve_collision(self, a: RigidBody, b: RigidBody, collision: ColliderContact) -> None:
        SLACK=0.2
        correction = collision.penetration / (a.inv_mass + b.inv_mass) * SLACK * collision.normal
        
        if a.mass != 0.0:
            a.transform.position += correction
        if b.mass != 0.0:
            b.transform.position -= correction

        restitution = a.restitution * b.restitution
        friction = a.coef_friction * b.coef_friction
        tangent = collision.normal.rotate(90)
        v_rel_linear = pg.Vector2(a.velocity - b.velocity)
        a_w = a.angular_velocity
        b_w = b.angular_velocity

        for point in collision.points:
            r_a = point - a.transform.position
            r_b = point - b.transform.position

            w_cross_r_a = pg.Vector2(-r_a.y * a_w, r_a.x * a_w) #omega is not encoded as a vector, so we pretend it is
            w_cross_r_b = pg.Vector2(-r_b.y * b_w, r_b.x * b_w)
            v_rel = v_rel_linear + w_cross_r_a - w_cross_r_b
            v_rel_tangent = v_rel.dot(tangent)
            
            collision_impulse = -(1+restitution)*(v_rel.dot(collision.normal))\
                / (a.inv_mass + b.inv_mass + (r_a.cross(collision.normal)**2 * a.inv_moment_of_inertia)\
                + (r_b.cross(collision.normal)**2 * b.inv_moment_of_inertia)) / len(collision.points)
            friction_impulse = (-v_rel_tangent / (a.inv_mass + b.inv_mass +\
                (r_a.cross(tangent))**2 * a.inv_moment_of_inertia +\
                (r_b.cross(tangent)**2  * b.inv_moment_of_inertia))) / len(collision.points)
            friction_impulse = pg.math.clamp(friction_impulse, -abs(collision_impulse)*friction, abs(collision_impulse)*friction)
            a.apply_impulse(collision.normal*collision_impulse + friction_impulse * tangent, point)
            b.apply_impulse(-1 * collision.normal*collision_impulse - friction_impulse * tangent, point)