from typing import List
import pygame
import sys
import numpy as np
import imageio
from config import HEIGHT, WIDTH
from tqdm import tqdm
from ember import Ember
from wind import WindField


IS_DEVELOPMENT = False
EMBER_COUNT = 600
FRAME_RATE = 60
RUNNING_TIME_IN_SECONDS = 300
WIND_STRENTH = 300

delta_time = 1 / 150


embers: List[Ember] = [Ember() for _ in range(EMBER_COUNT)]
wind_field = WindField(WIDTH, HEIGHT)


pygame.init()

flags = pygame.HWSURFACE | pygame.HWACCEL
screen = pygame.display.set_mode((WIDTH, HEIGHT), flags)
pygame.display.set_caption("Animated Video")

clock = pygame.time.Clock()

video_file = "animation.mp4"
video_writer = imageio.get_writer(video_file, fps=FRAME_RATE)

for frame in tqdm(range(RUNNING_TIME_IN_SECONDS * FRAME_RATE)):
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

    # screen.fill((0, 0, 0))
    # Add a semi-transparent rectangle to create the trail effect
    fade_surface = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    fade_surface.fill(
        (0, 0, 0, 10)
    )  # Adjust the last parameter for the opacity (0-255)
    screen.blit(fade_surface, (0, 0))

    wind_field.update(time=frame / 150)
    # wind_field.draw(screen)

    for ember in embers:
        wind = wind_field.get_wind(ember.position)
        velocity_x = wind[0] * WIND_STRENTH
        velocity_y = wind[1] * WIND_STRENTH

        ember.update((velocity_x, velocity_y), delta_time=delta_time)
        ember.draw(screen)

    pygame.display.flip()

    if not IS_DEVELOPMENT:
        current_frame = pygame.surfarray.array3d(screen)
        current_frame = np.rot90(current_frame)
        # current_frame = np.flip(current_frame, axis=0)
        video_writer.append_data(current_frame.copy())

    if IS_DEVELOPMENT:
        clock.tick(FRAME_RATE)

if not IS_DEVELOPMENT:
    video_writer.close()