from __future__ import annotations

import numpy as np
from dataclasses import dataclass, field
from random import uniform
import typing
import time
import turtle as t

if typing.TYPE_CHECKING:
    from nptyping import NDArray


HALF_WIDTH, HALF_HEIGHT = 200, 200
INITIAL_VELOCITY_COEF = 1000
FORCE_LAW = lambda r: 600000 / (r ** 2)
DT, N, VELOCITY_CUTOFF, ACCELERATION_CUTOFF = 0.01, 10, 1500, 10000


@dataclass
class Particle:
    coordinates: NDArray[(2,), float] = field(default_factory=lambda: np.array(
        [uniform(-HALF_WIDTH, HALF_WIDTH), uniform(-HALF_HEIGHT, HALF_HEIGHT)]
        ))
    velocity: NDArray[(2,)] = field(
            default_factory=lambda: INITIAL_VELOCITY_COEF * np.array([uniform(-1, 1), uniform(-1, 1)]))

    def adjust_velocity(self, other_particles: typing.List[Particle], dt: float = DT):
        a = np.array([0., 0.])
        for particle in other_particles:
            if list(particle.coordinates) == list(self.coordinates):
                continue
            r = self.coordinates - particle.coordinates
            a += (-1) * r * FORCE_LAW(np.abs(r))
        if np.linalg.norm(a) > ACCELERATION_CUTOFF:
            a *= ACCELERATION_CUTOFF / np.linalg.norm(a)
        self.velocity += a * dt
        if np.linalg.norm(self.velocity) > VELOCITY_CUTOFF:
            self.velocity *= VELOCITY_CUTOFF / np.linalg.norm(self.velocity)
        if abs(self.coordinates[0]) >= HALF_WIDTH:
            self.velocity[0] *= -1
            self.coordinates[0] = np.sign(self.coordinates[0]) * min(abs(self.coordinates[0]), HALF_WIDTH - 2)
        if abs(self.coordinates[1]) >= HALF_HEIGHT:
            self.velocity[1] *= -1
            self.coordinates[1] = np.sign(self.coordinates[1]) * min(abs(self.coordinates[1]), HALF_HEIGHT - 2)

    def move(self, dt: float = DT):
        self.coordinates += self.velocity * dt


@dataclass
class Gas:
    particles: typing.List[Particle]
    dt: float = DT

    def __init__(self, number: int = N) -> Gas:
        self.particles = [Particle() for _ in range(N)]

    def step(self):
        for particle in self.particles:
            particle.adjust_velocity(self.particles, self.dt)
        for particle in self.particles:
            particle.move(self.dt)

    def visualize_with_turtles(self, turtles: typing.List[t.Turtle]):
        for turtle, particle in zip(turtles, self.particles):
            turtle.goto(*particle.coordinates)

    def run_with_turtles(self, turtles: typing.List[t.Turtle]):
        while True:
            last_time: float = time.time()
            self.visualize_with_turtles(turtles)
            self.step()
            time.sleep(max(0, self.dt - (time.time() - last_time)))

    @classmethod
    def create_turtles_and_run(cls, number: int = N):
        self = cls()
        turtles = [t.Turtle(shape='circle') for i in range(number)]
        for turtle in turtles:
            turtle.penup()
            turtle.speed(0)
        master_turtle = t.Turtle()
        master_turtle.penup()
        master_turtle.speed(0)
        delta = 5
        master_turtle.goto(-HALF_WIDTH - delta, -HALF_HEIGHT - delta)
        master_turtle.pendown()
        master_turtle.goto(HALF_WIDTH + delta, -HALF_HEIGHT - delta)
        master_turtle.goto(HALF_WIDTH + delta, HALF_HEIGHT + delta)
        master_turtle.goto(-HALF_WIDTH - delta, HALF_HEIGHT + delta)
        master_turtle.goto(-HALF_WIDTH - delta, -HALF_HEIGHT - delta)
        master_turtle.up()
        master_turtle.goto(1000, 1000)
        del master_turtle
        self.run_with_turtles(turtles)


if __name__ == '__main__':
    Gas.create_turtles_and_run()