# SPDX-FileCopyrightText: 2023 Jeff Epler for Adafruit Industries
# SPDX-License-License: MIT

# Version 11
# Modified by Warmer_Summers

# Hardware:
# - Adafruit RP2040 Scorpio
# - NeoKey FeatherWing

# Code Function:
# This code integrates a NeoKey RP2040 FeatherWing with NeoPixel animations using CircuitPython.
# It allows the user to enter and exit a brightness control menu by holding down both buttons,
# and cycle through different animations using the left and right keys.
# Additionally, a reboot sequence is initiated by holding both keys for 15 seconds.

# Service Menus and Button Functionality:
# - Press both keys simultaneously for 5 seconds to enter the brightness menu.
# - In the brightness menu, the left key decreases brightness and the right key increases brightness in increments of 0.1.
# - NeoPixels behind the keys indicate the current brightness setting.
# - The main LED strands are turned off during brightness adjustment.
# - Press both keys simultaneously to exit the brightness menu.
# - Press the left key to move to the previous animation.
# - Press the right key to move to the next animation.
# - Hold both keys for 15 seconds to initiate a reboot sequence, indicated by four green flashes.

# Optimisations:
# - Added garbage collection to manage memory usage and avoid memory allocation errors.
# - Optimised the creation of custom animations (DAHX Heat Chase and DAHX Heat Rainbow) to use less memory.
# - Ensured consistent application of the colour map across the entire length of the LED strips.
# - Improved key press handling to prevent sluggish behaviour.
# - Memory logging is performed only once at startup and after switching animations.
# - Introduced dynamic instantiation of animations using lambda functions to better manage memory.

import time
import board
import microcontroller
import gc
from digitalio import DigitalInOut, Pull
from adafruit_debouncer import Debouncer
from adafruit_led_animation.animation.chase import Chase
from adafruit_led_animation.animation.comet import Comet
from adafruit_led_animation.animation.rainbow import Rainbow
from adafruit_led_animation.animation.sparkle import Sparkle
from adafruit_led_animation.animation.pulse import Pulse
from adafruit_led_animation.animation.blink import Blink
from adafruit_led_animation.animation import Animation
from adafruit_led_animation.color import (RED, YELLOW, ORANGE, GREEN, TEAL, CYAN, BLUE, PURPLE, MAGENTA, WHITE, BLACK, GOLD, PINK, AQUA, JADE, AMBER, OLD_LACE)
from adafruit_led_animation.group import AnimationGroup
from adafruit_neopxl8 import NeoPxl8
from adafruit_pixelmap import PixelMap
import neopixel

print("- DexThermic HX Brightness Control and Animations -")
time.sleep(1)  # Sleep for a bit to avoid a race condition on some systems

# ----- Key setup ----- #
switch_a_in = DigitalInOut(board.D5)
switch_b_in = DigitalInOut(board.D6)
switch_a_in.pull = Pull.UP
switch_b_in.pull = Pull.UP
switch_a = Debouncer(switch_a_in)
switch_b = Debouncer(switch_b_in)

# ----- LED Setup ----- #
strand_length = 250  # Set strand length to 250 LEDs
pixel_brightness = 0.2  # Set initial brightness to 0.2
num_strands = 8

num_pixels = num_strands * strand_length
pixels = NeoPxl8(board.NEOPIXEL0, num_pixels, auto_write=False, brightness=pixel_brightness)

def strand(n):
    return PixelMap(
        pixels,
        tuple(range(n * strand_length, (n + 1) * strand_length)),
        individual_pixels=True,
    )

strips = [strand(i) for i in range(num_strands)]

# ----- Key Pixel Setup ----- #
pixel_pin = board.D9
key_pixels = neopixel.NeoPixel(pixel_pin, 2, brightness=0.9)  # Set brightness to 0.9 for key pixels

# Define the colour relationship
left_key_color = BLUE
right_key_color = RED
strip_color = GOLD

brightness_menu_active = False
brightness_value = 0.2  # Initial brightness value
brightness_step = 0.1
hold_start_time = None
keys_released = False
flash_start_time = None
flash_state = False
cooldown_end_time = None

def clear_all_strips():
    for strip in strips:
        strip.fill(0)
    pixels.show()

def fill_all_strips(color):
    for strip in strips:
        strip.fill(color)
    pixels.show()

# Function to create a gradient from a list of colours with adjusted distribution
def create_color_gradient(colors, length, distribution=None):
    if distribution is None:
        distribution = [1] * len(colors)
    gradient = []
    total = sum(distribution)
    n = len(colors)
    for i in range(length):
        pos = (i / (length - 1)) * total
        left_color_index = 0
        while pos > distribution[left_color_index]:
            pos -= distribution[left_color_index]
            left_color_index += 1
        right_color_index = min(left_color_index + 1, n - 1)
        blend = pos / distribution[left_color_index]
        left_color = colors[left_color_index]
        right_color = colors[right_color_index]
        r = int(left_color[0] * (1 - blend) + right_color[0] * blend)
        g = int(left_color[1] * (1 - blend) + right_color[1] * blend)
        b = int(left_color[2] * (1 - blend) + right_color[2] * blend)
        gradient.append((r, g, b))
    return gradient

# Create a gradient for the strand length with custom distribution
gradient_colors = [RED, ORANGE, YELLOW, TEAL, BLUE]
distribution = [1, 1, 1, 1, 2]  # Give more space to blue
gradient = create_color_gradient(gradient_colors, strand_length, distribution)

# Custom Animation class to use gradient colours
class DAHXHeatChase(Animation):
    def __init__(self, pixel_object, speed, spacing, size, colors, name=None):
        super().__init__(pixel_object, speed, color=colors[0], name=name)
        self.colors = colors
        self.size = size
        self.spacing = spacing
        self.current_pixel = 0

    def draw(self):
        self.pixel_object.fill((0, 0, 0))  # Clear the strip
        for i in range(0, len(self.pixel_object), self.spacing + self.size):
            for j in range(self.size):
                pixel_index = (self.current_pixel + i + j) % len(self.pixel_object)
                color_index = (pixel_index) % len(self.colors)
                self.pixel_object[pixel_index] = self.colors[color_index]
        self.current_pixel = (self.current_pixel + 1) % (self.spacing + self.size)

class DAHXHeatRainbow(Animation):
    def __init__(self, pixel_object, speed, colors, name=None):
        super().__init__(pixel_object, speed, color=colors[0], name=name)
        self.colors = colors

    def draw(self):
        for i in range(len(self.pixel_object)):
            color_index = i % len(self.colors)
            self.pixel_object[i] = self.colors[color_index]

# Animation setup
animation_definitions = [
    ("Rainbow", lambda: AnimationGroup(
        Rainbow(strips[0], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[1], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[2], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[3], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[4], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[5], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[6], speed=0.1, period=2, precompute_rainbow=True),
        Rainbow(strips[7], speed=0.1, period=2, precompute_rainbow=True),
    )),
    ("Chase", lambda: AnimationGroup(
        Chase(strips[0], speed=0.05, color=BLUE, size=5, spacing=5),
        Chase(strips[1], speed=0.1, color=RED, size=5, spacing=5),
        Chase(strips[2], speed=0.1, color=BLUE, size=5, spacing=5),
        Chase(strips[3], speed=0.1, color=RED, size=5, spacing=5),
        Chase(strips[4], speed=0.1, color=BLUE, size=5, spacing=5),
        Chase(strips[5], speed=0.1, color=RED, size=5, spacing=5),
        Chase(strips[6], speed=0.1, color=BLUE, size=5, spacing=5),
        Chase(strips[7], speed=0.1, color=RED, size=5, spacing=5),
    )),
    ("Comet", lambda: AnimationGroup(
        Comet(strips[0], speed=0.02, color=RED, tail_length=65, bounce=False),
        Comet(strips[1], speed=0.05, color=YELLOW, tail_length=45, bounce=False),
        Comet(strips[2], speed=0.02, color=ORANGE, tail_length=55, bounce=False),
        Comet(strips[3], speed=0.05, color=RED, tail_length=65, bounce=False),
        Comet(strips[4], speed=0.02, color=YELLOW, tail_length=45, bounce=False),
        Comet(strips[5], speed=0.05, color=ORANGE, tail_length=55, bounce=False),
        Comet(strips[6], speed=0.02, color=RED, tail_length=65, bounce=False),
        Comet(strips[7], speed=0.05, color=YELLOW, tail_length=55, bounce=False),
    )),
    ("Sparkle", lambda: AnimationGroup(
        Sparkle(strips[0], speed=0.1, color=TEAL, num_sparkles=10),
        Sparkle(strips[1], speed=0.1, color=BLUE, num_sparkles=10),
        Sparkle(strips[2], speed=0.1, color=YELLOW, num_sparkles=10),
        Sparkle(strips[3], speed=0.1, color=TEAL, num_sparkles=10),
        Sparkle(strips[4], speed=0.1, color=BLUE, num_sparkles=10),
        Sparkle(strips[5], speed=0.1, color=YELLOW, num_sparkles=10),
        Sparkle(strips[6], speed=0.1, color=TEAL, num_sparkles=10),
        Sparkle(strips[7], speed=0.1, color=BLUE, num_sparkles=10),
    )),
    ("Pulse", lambda: AnimationGroup(
        Pulse(strips[0], speed=0.1, color=GOLD, period=4),
        Pulse(strips[1], speed=0.1, color=AMBER, period=3),
        Pulse(strips[2], speed=0.1, color=YELLOW, period=5),
        Pulse(strips[3], speed=0.1, color=GOLD, period=4),
        Pulse(strips[4], speed=0.1, color=AMBER, period=3),
        Pulse(strips[5], speed=0.1, color=YELLOW, period=5),
        Pulse(strips[6], speed=0.1, color=GOLD, period=4),
        Pulse(strips[7], speed=0.1, color=AMBER, period=3),
    )),
    ("Blink", lambda: AnimationGroup(
        Blink(strips[0], speed=0.05, color=PURPLE),
        Blink(strips[1], speed=0.1, color=MAGENTA),
        Blink(strips[2], speed=0.08, color=PINK),
        Blink(strips[3], speed=0.1, color=PURPLE),
        Blink(strips[4], speed=0.05, color=MAGENTA),
        Blink(strips[5], speed=0.08, color=PINK),
        Blink(strips[6], speed=0.05, color=PURPLE),
        Blink(strips[7], speed=0.1, color=MAGENTA),
    )),
    ("DAHX Heat Chase", lambda: AnimationGroup(
        DAHXHeatChase(strips[0], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[1], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[2], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[3], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[4], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[5], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[6], speed=0.1, spacing=5, size=5, colors=gradient),
        DAHXHeatChase(strips[7], speed=0.1, spacing=5, size=5, colors=gradient),
    )),
    ("DAHX Heat Rainbow", lambda: AnimationGroup(
        DAHXHeatRainbow(strips[0], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[1], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[2], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[3], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[4], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[5], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[6], speed=0.1, colors=gradient),
        DAHXHeatRainbow(strips[7], speed=0.1, colors=gradient),
    )),
]

animation_index = 0
current_animation = animation_definitions[animation_index][1]()  # Instantiate initial animation

# Set initial colours
fill_all_strips(strip_color)
key_pixels[0] = left_key_color
key_pixels[1] = right_key_color

# Function to log memory usage
def log_memory_usage():
    free_memory = gc.mem_free()
    allocated_memory = gc.mem_alloc()
    print(f"Memory Info - Free: {free_memory} Allocated: {allocated_memory}")

log_memory_usage()

while True:
    switch_a.update()  # Debouncer checks for changes in switch state
    switch_b.update()

    current_time = time.monotonic()

    # Check if we're in the cooldown period
    if cooldown_end_time and current_time < cooldown_end_time:
        continue
    else:
        cooldown_end_time = None

    # Check if both buttons are pressed simultaneously
    if not switch_a.value and not switch_b.value:
        if hold_start_time is None:
            hold_start_time = current_time
            flash_start_time = current_time
            print("Both keys pressed, starting hold timer")
        elif hold_start_time and current_time - hold_start_time >= 15:
            print("Initiating reboot sequence")
            for _ in range(4):
                key_pixels.fill(GREEN)
                time.sleep(0.25)
                key_pixels.fill((0, 0, 0))
                time.sleep(0.25)
            microcontroller.reset()
        elif hold_start_time and current_time - hold_start_time >= 5 and not brightness_menu_active:
            brightness_menu_active = True
            hold_start_time = None
            keys_released = False
            flash_start_time = None
            flash_state = False
            print("Entering brightness menu")
            key_pixels.fill(WHITE)
            clear_all_strips()  # Clear the main LED strands
        elif not brightness_menu_active and current_time - flash_start_time >= 1:
            flash_start_time = current_time
            flash_state = not flash_state
            key_pixels[0] = WHITE if flash_state else left_key_color
            key_pixels[1] = WHITE if flash_state else right_key_color
    else:
        if hold_start_time is not None and not brightness_menu_active:
            hold_start_time = None
            flash_start_time = None
            print("Keys released before 5 seconds, resetting hold timer")

    if brightness_menu_active:
        if keys_released:
            if not switch_a.value and not switch_b.value:
                brightness_menu_active = False
                key_pixels[0] = left_key_color
                key_pixels[1] = right_key_color
                pixels.brightness = brightness_value
                fill_all_strips(strip_color)
                cooldown_end_time = current_time + 2  # Set cooldown period for 2 seconds
                print("Exiting brightness menu")
            elif switch_a.fell:
                brightness_value = max(0.1, brightness_value - brightness_step)
                brightness_value = round(brightness_value, 1)
                key_pixels.brightness = brightness_value
                key_pixels.fill(WHITE)
                key_pixels.show()
                print(f"Decreasing brightness: {brightness_value}")
            elif switch_b.fell:
                brightness_value = min(1.0, brightness_value + brightness_step)
                brightness_value = round(brightness_value, 1)
                key_pixels.brightness = brightness_value
                key_pixels.fill(WHITE)
                key_pixels.show()
                print(f"Increasing brightness: {brightness_value}")
        else:
            keys_released = switch_a.rose or switch_b.rose
            if keys_released:
                print("Keys released in brightness menu, now ready for brightness adjustment")
    else:
        # Animation control
        if switch_a.fell and (switch_a.value or switch_b.value):
            animation_index = (animation_index - 1) % len(animation_definitions)
            print(f"Switching to animation: {animation_definitions[animation_index][0]}")
            log_memory_usage()
            gc.collect()  # Force garbage collection
            current_animation = animation_definitions[animation_index][1]()  # Instantiate animation
        if switch_b.fell and (switch_a.value or switch_b.value):
            animation_index = (animation_index + 1) % len(animation_definitions)
            print(f"Switching to animation: {animation_definitions[animation_index][0]}")
            log_memory_usage()
            gc.collect()  # Force garbage collection
            current_animation = animation_definitions[animation_index][1]()  # Instantiate animation

        current_animation.animate()
        pixels.show()