import pygame
import pygame.gfxdraw
import gas_sim as gs
import numpy as np

##### CONFIG #####
window_size = (1280, 720)
target_fps = 60

particle_radius_ws = 0.025
particle_compressibility = 0.1
particle_friction = 0.1


cursor_radius_ws = 0.15

substeps = 10
##################

# ws - world space
# ss - screen space

window_aspect = window_size[0] / window_size[1]
cursor_radius_ss = cursor_radius_ws * window_size[1]
particle_radius_ss = particle_radius_ws * window_size[1]


# pygame init
pygame.init()
window = pygame.display.set_mode(window_size)
clock = pygame.time.Clock()


# gas sim init
gas_sim = gs.GasSim(particle_radius_ws, particle_compressibility, particle_friction, substeps)


# add walls
gas_sim.add_wall((0, 0), (window_aspect, 0))
gas_sim.add_wall((0, 1), (window_aspect, 1))
gas_sim.add_wall((0, 0), (0, 1))
gas_sim.add_wall((1, 0), (1, 1))


not_loaded_particles = True
dt = target_fps / 1000



# main loop
while True:
    # events
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            quit()


    # on mouse down
    if pygame.mouse.get_pressed()[0]:
        
        mouse_pos_ss = pygame.mouse.get_pos()

        mouse_pos_ws = (mouse_pos_ss[0] / window_size[1], mouse_pos_ss[1] / window_size[1])

        # if not colliding with any particles
        if not gas_sim.cursor_particle_collide_test(mouse_pos_ws, particle_radius_ws):
            gas_sim.add_particle(mouse_pos_ws)

    # on s key press
    if pygame.key.get_pressed()[pygame.K_s]:
        # save particles to file
        particles = tuple(pos for pos, _ in gas_sim.particles)
        np.save("particles.npy", particles)

        print("Saved particles to file")
    
    # on l key press
    if not_loaded_particles and pygame.key.get_pressed()[pygame.K_l]:
        # load particles from file
        particles = np.load("particles.npy")
        for particle in particles:
            gas_sim.add_particle(particle)

        print("Loaded particles from file")

        not_loaded_particles = False



    # update particles
    #gas_sim.update(dt)
    



    window.fill((0, 0, 0))

    # draw cursor
    cursor_pos = pygame.mouse.get_pos()
    pygame.gfxdraw.aacircle(window, cursor_pos[0], cursor_pos[1], int(cursor_radius_ss), (255, 255, 255))

    # draw particles
    for particle_pos, _ in gas_sim.particles:
        particle_pos_ss = (int(particle_pos[0] * window_size[1]), int(particle_pos[1] * window_size[1]))
        pygame.gfxdraw.aacircle(window, particle_pos_ss[0], particle_pos_ss[1], int(particle_radius_ss), (255, 255, 255))


    # update
    pygame.display.update()
    dt = clock.tick(60) / 1000