from itertools import product

import numpy as np
from PIL import Image, ImageOps

PALETTE = [
    (0x07, 0x07, 0x07),
    (0x1F, 0x07, 0x07),
    (0x2F, 0x0F, 0x07),
    (0x47, 0x0F, 0x07),
    (0x57, 0x17, 0x07),
    (0x67, 0x1F, 0x07),
    (0x77, 0x1F, 0x07),
    (0x8F, 0x27, 0x07),
    (0x9F, 0x2F, 0x07),
    (0xAF, 0x3F, 0x07),
    (0xBF, 0x47, 0x07),
    (0xC7, 0x47, 0x07),
    (0xDF, 0x4F, 0x07),
    (0xDF, 0x57, 0x07),
    (0xDF, 0x57, 0x07),
    (0xD7, 0x5F, 0x07),
    (0xD7, 0x5F, 0x07),
    (0xD7, 0x67, 0x0F),
    (0xCF, 0x6F, 0x0F),
    (0xCF, 0x77, 0x0F),
    (0xCF, 0x7F, 0x0F),
    (0xCF, 0x87, 0x17),
    (0xC7, 0x87, 0x17),
    (0xC7, 0x8F, 0x17),
    (0xC7, 0x97, 0x1F),
    (0xBF, 0x9F, 0x1F),
    (0xBF, 0x9F, 0x1F),
    (0xBF, 0xA7, 0x27),
    (0xBF, 0xA7, 0x27),
    (0xBF, 0xAF, 0x2F),
    (0xB7, 0xAF, 0x2F),
    (0xB7, 0xB7, 0x2F),
    (0xB7, 0xB7, 0x37),
    (0xCF, 0xCF, 0x6F),
    (0xDF, 0xDF, 0x9F),
    (0xEF, 0xEF, 0xC7),
    (0xFF, 0xFF, 0xFF),
]


def update_fire(fire: np.ndarray) -> None:
    """Update DOOM fire.

    Args:
        fire: DOOM fire array.
    """
    height, width = fire.shape
    for y, x in product(range(height - 1), range(width)):
        new_x = (x + np.random.randint(-1, 2)) % width
        fire[y + 1][new_x] = max(fire[y][x] - np.random.randint(0, 3), 0)


if __name__ == "__main__":
    # Initial state of the fire
    width, height = 64, 64
    doom_fire = np.zeros((height, width))
    doom_fire[0] = len(PALETTE) - 1

    frames = 24
    for _ in range(frames):
        update_fire(doom_fire)

    # Save as PNG
    img_array = np.array(PALETTE, dtype=np.uint8)[doom_fire.astype(np.uint8)]
    img = Image.fromarray(img_array[::-1])
    img = ImageOps.scale(img, 4, Image.Resampling.NEAREST)
    img.save("doom_fire.png")