from rgbmatrix import RGBMatrix, RGBMatrixOptions, graphics
import itertools as it
from time import sleep
import random
import colorsys

from PIL import Image
from PIL import ImageDraw

mode = 0

options = RGBMatrixOptions()
options.rows = 32
options.chain_length = 6
options.disable_hardware_pulsing = 0
options.gpio_slowdown = 1
options.brightness = 100
options.pwm_bits = 11

FIRE_WIDTH = 192
FIRE_HEIGHT = 32

COLORS = [
    {"r": 7, "g": 7, "b": 7},
    {"r": 31, "g": 7, "b": 7},
    {"r": 47, "g": 15, "b": 7},
    {"r": 71, "g": 15, "b": 7},
    {"r": 87, "g": 23, "b": 7},
    {"r": 103, "g": 31, "b": 7},
    {"r": 119, "g": 31, "b": 7},
    {"r": 143, "g": 39, "b": 7},
    {"r": 159, "g": 47, "b": 7},
    {"r": 175, "g": 63, "b": 7},
    {"r": 191, "g": 71, "b": 7},
    {"r": 199, "g": 71, "b": 7},
    {"r": 223, "g": 79, "b": 7},
    {"r": 223, "g": 87, "b": 7},
    {"r": 223, "g": 87, "b": 7},
    {"r": 215, "g": 95, "b": 7},
    {"r": 215, "g": 95, "b": 7},
    {"r": 215, "g": 103, "b": 15},
    {"r": 207, "g": 111, "b": 15},
    {"r": 207, "g": 119, "b": 15},
    {"r": 207, "g": 127, "b": 15},
    {"r": 207, "g": 135, "b": 23},
    {"r": 199, "g": 135, "b": 23},
    {"r": 199, "g": 143, "b": 23},
    {"r": 199, "g": 151, "b": 31},
    {"r": 191, "g": 159, "b": 31},
    {"r": 191, "g": 159, "b": 31},
    {"r": 191, "g": 167, "b": 39},
    {"r": 191, "g": 167, "b": 39},
    {"r": 191, "g": 175, "b": 47},
    {"r": 183, "g": 175, "b": 47},
    {"r": 183, "g": 183, "b": 47},
    {"r": 183, "g": 183, "b": 55},
    {"r": 207, "g": 207, "b": 111},
    {"r": 223, "g": 223, "b": 159},
    {"r": 239, "g": 239, "b": 199},
    {"r": 255, "g": 255, "b": 255},
]

matrix = RGBMatrix(options=options)
print("Matrix initialized\n")


intensities = []
offscreen_canvas = matrix.CreateFrameCanvas()
image = Image.new("RGB", (192, 32))

# Fill up intensities
for i in range(0, FIRE_WIDTH * FIRE_HEIGHT):
    intensities.append(0)

# Create fire source
for column in range(0, FIRE_WIDTH):
    overflow_pixel_index = FIRE_WIDTH * FIRE_HEIGHT
    index = (overflow_pixel_index - FIRE_WIDTH) + column
    intensities[index] = 36 # max intensity


def update_pixel_intensity(pixel_index):
    below_pixel_index = pixel_index + FIRE_WIDTH

    # overflow ?
    if below_pixel_index >= FIRE_WIDTH * FIRE_HEIGHT:
        return

    decay = random.randint(0, 2)
    below_pixel_intensity = intensities[below_pixel_index]
    new_intensity = (
        below_pixel_intensity - decay if (below_pixel_intensity - decay) >= 0 else 0
    )
    intensities[pixel_index - decay] = new_intensity


def propagate():
    for column in range(0, FIRE_WIDTH):
        for row in range(0, FIRE_HEIGHT):
            index = column + (FIRE_WIDTH * row)
            update_pixel_intensity(index)


def burn():
    propagate()
    draw = ImageDraw.Draw(image)
    for col in range(0, FIRE_WIDTH):
        for row in range(0, FIRE_HEIGHT):
            index = col + (FIRE_WIDTH * row)
            color = COLORS[intensities[index]]
            draw.point((col, row), (color["r"], color["g"], color["b"]))
    matrix.Clear()
    matrix.SetImage(image, 0, 0)
    # diretly setting pixel is as slow
    # matrix.SetPixel(i % FIRE_WIDTH, i / FIRE_WIDTH, color["r"], color["g"], color["b"])


while True:
    burn()
    # sleep(.01)