import pygame
import math
from enum import Enum

# import pprint
# from pygame import gfxdraw

import numpy as np

SIZE = (1000, 1000)

LEFT_COLOR = "#606858"
RIGHT_COLOR = "#F7F9F6"

BLOCKS_SIZE = 10

PADDLE_LENGTH = 100

HALF_SPEED = 250


def draw_circle(surface, x, y, radius, color):
    pygame.draw.circle(surface, color, (x, y), radius)
    # AA Drawing
    # gfxdraw.aacircle(surface, x, y, radius, color)
    # gfxdraw.filled_circle(surface, x, y, radius, color)


def draw_rect(surface, topL, topT, width, height, color, corner_radius=0):
    pygame.draw.rect(
        surface,
        color,
        pygame.Rect(topL, topT, width, height),
        border_radius=corner_radius,
    )


def angle_arccos(v1, v2):
    y_diff = v2[1] - v1[1]
    x_diff = v2[0] - v1[0]
    angle = math.atan2(y_diff, x_diff)
    return angle


def intersects(rect, r, center):
    circle_distance_x = abs(center[0] - rect.centerx)
    circle_distance_y = abs(center[1] - rect.centery)
    if circle_distance_x > rect.w / 2.0 + r or circle_distance_y > rect.h / 2.0 + r:
        return False
    if circle_distance_x <= rect.w / 2.0 or circle_distance_y <= rect.h / 2.0:
        return True
    corner_x = circle_distance_x - rect.w / 2.0
    corner_y = circle_distance_y - rect.h / 2.0
    corner_distance_sq = corner_x**2.0 + corner_y**2.0
    return corner_distance_sq <= r**2.0


class TileType(Enum):
    PLAYER = 1
    CPU = 2


class ObjectType(Enum):
    PLAYER = 1
    CPU = 2


class Tile:

    def __init__(self, positionx, positiony, sidelength, color, tileType) -> None:
        self.x = positionx
        self.y = positiony
        self.length = sidelength
        self.rect = pygame.Rect(positionx, positiony, sidelength, sidelength)
        self.color = color
        self.tileType = tileType


class Paddle:

    def __init__(
        self, startx, starty, length, depth, color, radius, initial_vel=300
    ) -> None:
        self.x = startx
        self.y = starty
        self.length = length
        self.depth = depth
        self.vy = initial_vel
        self.color = color
        self.radius = radius
        self.elasticity = 0.08

    def draw(self, screen):
        draw_rect(
            screen,
            self.x - (self.depth / 2),
            self.y - (self.length / 2),
            self.depth,
            self.length,
            self.color,
            corner_radius=self.radius,
        )

    def move(self, dt):
        newy = self.y + (self.vy * dt)
        if (newy - self.length / 2 <= 0) or (newy + self.length / 2 >= SIZE[0]):
            return
        self.y = newy

    def set_vel(self, vel):
        self.vy = vel


class Ball:

    def __init__(self, startx, starty, size, color, charType) -> None:
        self.x = startx
        self.y = starty
        self.radius = size
        self.color = pygame.Color(color)
        self.mass = 0.25
        self.vx = 0
        self.vy = 0
        self.charType = charType
        self.speed = HALF_SPEED

    def draw(self, screen):
        draw_circle(screen, self.x, self.y, self.radius, self.color)

    def set_speed(self, speed):
        self.speed = speed

    def move(self, grid, object_grid, dt):
        # Pretend to move to new location
        newx = self.x + (self.vx * dt)
        newy = self.y + (self.vy * dt)

        # If we are at the top/bottom edge of the board, flip Y velocity
        if newy + self.radius >= SIZE[0] or newy - self.radius <= 0:
            self.vy = self.vy * -1.0

        # Back wall
        # TODO: Make this a life or lose condition
        if newx + self.radius >= SIZE[0] or newx - self.radius <= 0:
            self.vx = self.vx * -1.0

        grid_x = math.floor(newx / 100)
        grid_y = math.floor(newy / 100)

        candidates = []
        if (
            grid[grid_y][grid_x] == TileType.CPU and self.charType == ObjectType.PLAYER
        ) or (
            grid[grid_y][grid_x] == TileType.PLAYER and self.charType == ObjectType.CPU
        ):
            candidates.append(object_grid[grid_y][grid_x])
        if grid_x + 1 < BLOCKS_SIZE:
            if self.vx > 0:
                if (
                    grid[grid_y][grid_x + 1] == TileType.CPU
                    and self.charType == ObjectType.PLAYER
                ) or (
                    grid[grid_y][grid_x + 1] == TileType.PLAYER
                    and self.charType == ObjectType.CPU
                ):
                    candidates.append(object_grid[grid_y][grid_x + 1])
        if grid_x - 1 >= 0:
            if self.vx <= 0:
                if (
                    grid[grid_y][grid_x - 1] == TileType.CPU
                    and self.charType == ObjectType.PLAYER
                ) or (
                    grid[grid_y][grid_x - 1] == TileType.PLAYER
                    and self.charType == ObjectType.CPU
                ):
                    candidates.append(object_grid[grid_y][grid_x - 1])
        if grid_y + 1 < BLOCKS_SIZE:
            if self.vy > 0:
                if (
                    grid[grid_y + 1][grid_x] == TileType.CPU
                    and self.charType == ObjectType.PLAYER
                ) or (
                    grid[grid_y + 1][grid_x] == TileType.PLAYER
                    and self.charType == ObjectType.CPU
                ):
                    candidates.append(object_grid[grid_y + 1][grid_x])
        if grid_y - 1 >= 0:
            if self.vy <= 0:
                if (
                    grid[grid_y - 1][grid_x] == TileType.CPU
                    and self.charType == ObjectType.PLAYER
                ) or (
                    grid[grid_y - 1][grid_x] == TileType.PLAYER
                    and self.charType == ObjectType.CPU
                ):
                    candidates.append(object_grid[grid_y - 1][grid_x])

        for candidate in candidates:
            if intersects(candidate.rect, self.radius, (self.x, self.y)):

                if self.x >= candidate.x and self.x <= (candidate.x + candidate.length):
                    if (candidate.y + candidate.length) < self.y:
                        self.vy = self.vy * -1
                    if self.y < candidate.y:
                        self.vy = self.vy * -1

                if candidate.x > self.x:
                    self.vx = self.vx * -1.0

                if (candidate.x + candidate.length) < self.x:
                    self.vx = self.vx * -1.0

                if candidate.tileType == TileType.PLAYER:
                    candidate.tileType = TileType.CPU
                    candidate.color = RIGHT_COLOR
                else:
                    candidate.tileType = TileType.PLAYER
                    candidate.color = LEFT_COLOR

                tgrid_x = int(candidate.x / 100)
                tgrid_y = int(candidate.y / 100)
                if grid[tgrid_y][tgrid_x] == TileType.PLAYER:
                    grid[tgrid_y][tgrid_x] = TileType.CPU
                else:
                    grid[tgrid_y][tgrid_x] = TileType.PLAYER
                break

        self.x = self.x + (self.vx * dt)
        self.y = self.y + (self.vy * dt)

        # Fix speed if necessary:
        if abs(self.vx) < 3:
            if self.vx < 0:
                self.vx = -5
            else:
                self.vx = 5

        # Normalize to speed
        mag = math.sqrt((self.vx * self.vx) + (self.vy * self.vy))
        if mag != 0:
            self.vx = (self.vx / mag) * self.speed
            self.vy = (self.vy / mag) * self.speed

    def perform_spawn_hit(self):
        forceX = np.random.uniform(low=-50.0, high=50.0)
        forceY = np.random.uniform(low=-50.0, high=50.0)

        # Impulse dv = (f * dt)/m
        # Assuming dt = 1
        self.vx = (forceX) / self.mass
        self.vy = (forceY) / self.mass


def ball_paddle_collide(ball: Ball, paddle: Paddle):
    pRect = pygame.Rect(
        paddle.x - paddle.depth / 2,
        paddle.y - paddle.length / 2,
        paddle.depth,
        paddle.length,
    )

    collide = intersects(pRect, ball.radius, (ball.x, ball.y))

    if collide:
        ball.vx = (ball.vx * -1)

        # Depending on where the ball hit on the paddle 
        # relative to the center, add more or less vy
        mag = (ball.y - paddle.y) / (paddle.length * 1.25)

        ball.vy = ball.vy + (mag * 50)
        return True

    return False


pygame.init()
screen = pygame.display.set_mode(SIZE)
clock = pygame.time.Clock()

running = True
dt = 0

ball_pos_left = pygame.Vector2(screen.get_width() / 4, screen.get_height() / 2)
ball_pos_right = pygame.Vector2((screen.get_width() / 4) * 3, screen.get_height() / 2)
left_ball = Ball(ball_pos_left.x, ball_pos_left.y, 10, RIGHT_COLOR, ObjectType.PLAYER)
right_ball = Ball(ball_pos_right.x, ball_pos_right.y, 10, LEFT_COLOR, ObjectType.CPU)

p1 = Paddle(20, screen.get_height() / 2, PADDLE_LENGTH, 20, RIGHT_COLOR, 5, 300)
p2 = Paddle(
    screen.get_width() - 20,
    screen.get_height() / 2,
    PADDLE_LENGTH,
    20,
    LEFT_COLOR,
    5,
    300,
)

p1score = 0
p2score = 0

grid = [[] * BLOCKS_SIZE]

LENGTH = BLOCKS_SIZE

A = TileType.PLAYER
B = TileType.CPU

# Initialize the 2D array with zeros
grid = [[0 for _ in range(LENGTH)] for _ in range(LENGTH)]

# Determine the midpoint of the array
midpoint = LENGTH // 2

# Fill the first half of the array with value A
for i in range(midpoint):
    for j in range(LENGTH):
        grid[j][i] = A

# Fill the second half of the grid with value B
for i in range(midpoint, LENGTH):
    for j in range(LENGTH):
        grid[j][i] = B


object_grid = [[0 for _ in range(LENGTH)] for _ in range(LENGTH)]

for i in range(LENGTH):
    for j in range(LENGTH):
        if grid[j][i] == TileType.PLAYER:
            object_grid[j][i] = Tile(i * 100, j * 100, 100, LEFT_COLOR, TileType.PLAYER)
        else:
            object_grid[j][i] = Tile(i * 100, j * 100, 100, RIGHT_COLOR, TileType.CPU)


BACKGROUND_RIGHT = pygame.Rect(
    screen.get_width() / 2, 0, screen.get_width() / 2, screen.get_height()
)

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    # Render grid
    for i in range(LENGTH):
        for j in range(LENGTH):
            t = object_grid[j][i]
            draw_rect(screen, t.rect.x, t.rect.y, t.rect.width, t.rect.height, t.color)

    # Draw ball
    left_ball.draw(screen)
    right_ball.draw(screen)
    p1.draw(screen)
    p2.draw(screen)

    # input handling
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    p1.set_vel(0)
    p2.set_vel(0)
    keys = pygame.key.get_pressed()
    if keys[pygame.K_w]:
        p1.set_vel(-300.0)
    if keys[pygame.K_s]:
        p1.set_vel(300.0)
    if keys[pygame.K_i]:
        p2.set_vel(-300.0)
    if keys[pygame.K_k]:
        p2.set_vel(300.0)
    if keys[pygame.K_h]:
        left_ball.perform_spawn_hit()
        right_ball.perform_spawn_hit()
    if keys[pygame.K_q]:
        running = False

    # Movement and physics
    left_ball.move(grid, object_grid, dt)
    right_ball.move(grid, object_grid, dt)
    p1.move(dt)
    p2.move(dt)

    ball_paddle_collide(left_ball, p1)
    ball_paddle_collide(right_ball, p2)

    # Compute Score

    p1score = 0
    p2score = 0
    for i in grid:
        for j in i:
            if j == TileType.PLAYER:
                p1score += 1
            else:
                p2score += 1

    # if your score is higher, your ball is faster
    left_ball.set_speed(HALF_SPEED * (p1score/50.0))
    right_ball.set_speed(HALF_SPEED * (p2score/50.0))


    pygame.display.flip()
    dt = clock.tick(60) / 1000
    print("\x1b[1A\x1b[2K", end="")
    print(f"scores -> \t {p1score} ({left_ball.speed}, ({left_ball.vx}, {left_ball.vy}) \t {p2score} ({right_ball.speed}, {right_ball.vx})")

pygame.quit()