#include "step3.h"

static const char *TAG = "step3";

static std::random_device rd;
static std::mt19937 gen(rd());
static std::uniform_int_distribution<> puzzle_dist(0, 7);
static std::uniform_int_distribution<> led_picker_dist(0, 20);
static std::uniform_int_distribution<> led_color_dist(0, 5);

static const int INDICATOR_RED[6] = {30, 0, 0, 25, 15, 10};
static const int INDICATOR_GREEN[6] = {0, 0, 30, 5, 0, 10};
static const int INDICATOR_BLUE[6] = {0, 30, 0, 0, 15, 10};

void step3(void) {
    int solved_puzzles = 0;
    while (solved_puzzles < 3) {
        lcd_set_cursor(&lcd, 1, 1);

        int puzzle = puzzle_dist(gen);
        // int puzzle = 8;
        switch (puzzle) {
            case 0: {
                lcd_print(&lcd, "Clear");
                set_module_time(15000);
                start_module_timer();

                std::uniform_int_distribution<> green_indicators_dist(2, 15);

                // create all green indicators
                uint8_t green_indicators = green_indicators_dist(gen);
                std::set<int> indicators;

                // ESP_LOGI(TAG, "green indicators: %i", green_indicators);

                // ***CHANGE TO PUZZLE 5 METHOD***
                while (indicators.size() < green_indicators) {
                    int led = led_picker_dist(gen);
                    indicators.insert(led);
                    // ESP_LOGI(TAG, "added green led at %i", led);
                }

                for (std::set<int>::iterator it = indicators.begin(); it != indicators.end(); it++) {
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, *it, 0, 10, 0));
                }

                // add non-green indicators
                const int NON_GREEN_INDICATOR_RED[3] = {30, 0, 15};
                const int NON_GREEN_INDICATOR_BLUE[3] = {0, 30, 15};
                std::uniform_int_distribution<> non_green_indicator_color_dist(0, 2);

                std::uniform_int_distribution<> remaining_indicators_dist(0, 20 - green_indicators);
                int non_green_indicators = remaining_indicators_dist(gen);
                std::set<int> remaining_indicators;

                // ESP_LOGI(TAG, "non-green indicators: %i", non_green_indicators);

                while (remaining_indicators.size() < non_green_indicators) {
                    int led_attempt = led_picker_dist(gen);
                    if (indicators.find(led_attempt) == indicators.end()) {
                        remaining_indicators.insert(led_attempt);
                        // ESP_LOGI(TAG, "added non-green led at %i", led_attempt);
                    }
                }

                for (std::set<int>::iterator it = remaining_indicators.begin(); it != remaining_indicators.end(); it++) {
                    int color = non_green_indicator_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, *it, NON_GREEN_INDICATOR_RED[color], 0, NON_GREEN_INDICATOR_BLUE[color]));
                }

                ESP_ERROR_CHECK(led_strip_refresh(leds));

                // wait for time
                while (1) {
                    if (get_module_time() <= 0) {
                        uint8_t state = get_switch_state();
                        uint8_t flipped_state = 0;
                        flipped_state |= (state & 0x01) << 3;
                        flipped_state |= (state & 0x02) << 1;
                        flipped_state |= (state & 0x04) >> 1;
                        flipped_state |= (state & 0x08) >> 3;

                        if (flipped_state == green_indicators) {
                            solved_puzzles++;
                        } else {
                            strike("Incorrect switch combination! (step 3, puzzle 0)");
                        }
                        break;
                    }
                    vTaskDelay(pdMS_TO_TICKS(10));
                }

                break;
            }
            case 1: {
                lcd_print(&lcd, "Blank");
                set_module_time(20000);
                start_module_timer();
                
                std::uniform_int_distribution<> on_indicators_dist(16, 21);

                uint8_t indicators_num = on_indicators_dist(gen);
                std::set<int> indicators;

                while (indicators.size() < indicators_num) {
                    indicators.insert(led_picker_dist(gen));
                }

                for (std::set<int>::iterator it = indicators.begin(); it != indicators.end(); it++) {
                    int color = led_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, *it, INDICATOR_RED[color], INDICATOR_GREEN[color], INDICATOR_BLUE[color]));
                }
                ESP_ERROR_CHECK(led_strip_refresh(leds));
                // ESP_LOGI(TAG, "puzzle 1 LEDs set (%i leds)", indicators_num);

                if (indicators_num < 18) {
                    while (1) {
                        if (get_module_time() <= 0) {
                            if (get_switch_state() == 0b1111) {
                                solved_puzzles++;
                                // ESP_LOGI(TAG, "puzzle 1 solved (switches)!");
                            } else {
                                strike("Switch state was changed! (step 3, puzzle 1)");
                            }
                            break;
                        }
                        vTaskDelay(pdMS_TO_TICKS(10));
                    }
                } else {
                    uint8_t starting_switch_state = get_switch_state();
                    KeypadKey key;
                    std::string pressed_keys;
                    while (1) {
                        if (get_pressed_keypad(&key)) {
                            pressed_keys += char_of_keypad_key(key);
                            // ESP_LOGI(TAG, "key %c pressed", char_of_keypad_key(key));
                        }
                        if (get_module_time() <= 0) {
                            if (starting_switch_state == get_switch_state()) {
                                if (pressed_keys == "ADCB") {
                                    solved_puzzles++;
                                    // ESP_LOGI(TAG, "puzzle 1 solved (keypad)!");
                                }
                            } else {
                                strike("Switch state was changed! (step 3, puzzle 1)");
                            }
                            break;
                        }
                        vTaskDelay(pdMS_TO_TICKS(10));
                    }
                }



                break;
            }
            case 2: {
                set_module_time(20000);
                start_module_timer();

                std::uniform_int_distribution<> lit_led_dist(0, 1);
                bool rfid_lit = lit_led_dist(gen);
                // ESP_LOGI(TAG, "rfid lit: %i", rfid_lit);
                bool lcd_lit = lit_led_dist(gen);
                bool speaker_lit = lit_led_dist(gen);
                bool keypad_lit = lit_led_dist(gen);
                bool tft_lit = lit_led_dist(gen);

                if (rfid_lit) {
                    int color = led_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, 10, INDICATOR_RED[color], INDICATOR_GREEN[color], INDICATOR_BLUE[color]));
                }
                if (lcd_lit) {
                    int color = led_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, 12, INDICATOR_RED[color], INDICATOR_GREEN[color], INDICATOR_BLUE[color]));
                }
                if (speaker_lit) {
                    int color = led_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, 9, INDICATOR_RED[color], INDICATOR_GREEN[color], INDICATOR_BLUE[color]));
                }
                if (keypad_lit) {
                    int color = led_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, 11, INDICATOR_RED[color], INDICATOR_GREEN[color], INDICATOR_BLUE[color]));
                }
                if (tft_lit) {
                    int color = led_color_dist(gen);
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, 6, INDICATOR_RED[color], INDICATOR_GREEN[color], INDICATOR_BLUE[color]));
                }
                ESP_ERROR_CHECK(led_strip_refresh(leds));

                int green_button_pressed = 0;
                // ESP_LOGI(TAG, "green pressed: %i", green_button_pressed);
                int blue_button_pressed = 0;
                int fingerprint_sensor_pressed = 0;
                std::string keypad_string;

                ButtonKey button;
                KeypadKey key;
                SwitchKey switch1 = s1;
                SwitchKey switch2 = s2;

                while (1) {
                    if (get_pressed_button(&button)) {
                        uint8_t button_state = get_button_state();
                        // ESP_LOGI(TAG, "Button pressed! (%i)", button_state());

                        if ((button_state & 0b1) == 0b1) {
                            green_button_pressed++;
                            if ((green_button_pressed > 1) || !rfid_lit) {
                                strike("Green button pressed too many times! (step 3, puzzle 2)");
                                break;
                            }
                        }
                        if ((button_state & 0b1000) == 0b1000) {
                            blue_button_pressed++;
                            if ((blue_button_pressed > 1) || !lcd_lit) {
                                strike("Blue button pressed too many times! (step 3, puzzle 2)");
                                break;
                            }
                        }
                    }
                    if (get_touch_pressed()) {
                        fingerprint_sensor_pressed++;
                        if ((fingerprint_sensor_pressed > 2) || !speaker_lit) {
                                strike("Fingerprint sensor pressed too many times! (step 3, puzzle 2)");
                                break;
                            }
                    }
                    if (get_pressed_keypad(&key)) {
                        bool wrong = false;
                        keypad_string += char_of_keypad_key(key);
                        switch (keypad_string.length()) {
                            case 1: {
                                if (keypad_string != "1") {
                                    strike("1 was not pressed on Keypad! (step 3, puzzle 2)");
                                    wrong = true;
                                }
                                break;
                            }
                            case 2: {
                                if (keypad_string != "12") {
                                    strike("12 was not pressed on Keypad! (step 3, puzzle 2)");
                                    wrong = true;
                                }
                                break;
                            }
                            default: {
                                strike("keypad length was more than 2! (step 3, puzzle 2)");
                                wrong = true;
                                break;
                            }
                        }
                        if (wrong) {
                            break;
                        }
                    }
                    if (get_flipped_up_switch(&switch1) || get_flipped_up_switch(&switch2)) {
                        if (!tft_lit) {
                            strike("Switch 1 or 2 were flipped up! (step 3, puzzle 2)");
                            break;
                        }
                    }
                    if (get_module_time() <= 0) {
                        bool rfid_correct = !(rfid_lit && (green_button_pressed != 1));
                        bool lcd_correct = !(lcd_lit && (blue_button_pressed != 1));
                        bool speaker_correct = !(speaker_lit && (fingerprint_sensor_pressed != 2));
                        bool keypad_correct = !(keypad_lit && (keypad_string != "12"));
                        bool tft_correct = !(tft_lit && ((get_switch_state() & 0b11) != 0b11));

                        if (rfid_correct && lcd_correct && speaker_correct && keypad_correct && tft_correct) {
                            solved_puzzles++;
                        } else {
                            strike("Final state was not correct! (step 3, puzzle 2)");
                        }
                        break;
                    }
                    vTaskDelay(pdMS_TO_TICKS(10));
                }

                break;
            }
            case 3: {
                lcd_print(&lcd, "Nothing");
                set_module_time(10000);
                start_module_timer();

                const int COLOR_RED[4] = {0, 20, 10, 0};
                const int COLOR_GREEN[4] = {20, 0, 10, 0};
                const int COLOR_BLUE[4] = {0, 0, 0, 20};

                static std::uniform_int_distribution<> color_dist(0, 3);

                int tft_color = color_dist(gen);
                int speaker_color = color_dist(gen);
                int s3_color = color_dist(gen);

                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 6, COLOR_RED[tft_color], COLOR_GREEN[tft_color], COLOR_BLUE[tft_color]));
                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 9, COLOR_RED[speaker_color], COLOR_GREEN[speaker_color], COLOR_BLUE[speaker_color]));
                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 14, COLOR_RED[s3_color], COLOR_GREEN[s3_color], COLOR_BLUE[s3_color]));

                ESP_ERROR_CHECK(led_strip_refresh(leds));

                int buttons_pressed = 0;

                ButtonKey button;

                while (1) {
                    if (get_pressed_button(&button)) {
                        buttons_pressed++;
                        uint8_t button_state = get_button_state();

                        if (buttons_pressed == 1) {
                            if ((button_state >> tft_color) != 0b1) {
                                strike("Wrong button pressed! (step 3, puzzle 3)");
                                break;
                            }
                        } else if (buttons_pressed == 2) {
                            if ((button_state >> speaker_color) != 0b1) {
                                strike("Wrong button pressed! (step 3, puzzle 3)");
                                break;
                            }
                        } else if (buttons_pressed == 3) {
                            if ((button_state >> s3_color) != 0b1) {
                                strike("Wrong button pressed! (step 3, puzzle 3)");
                            } else {
                                solved_puzzles++;
                            }
                            break;
                        }

                    }
                    if (get_module_time() <= 0) {
                        strike("Ran out of time! (step 3, puzzle 3)");
                        break;
                    }
                    vTaskDelay(pdMS_TO_TICKS(10));
                }



                break;
            }
            case 4: {
                lcd_print(&lcd, "Blink");
                set_module_time(25000);
                start_module_timer();

                // buttons
                const int BUTTON_COLOR_RED[4] = {0, 20, 10, 0};
                const int BUTTON_COLOR_GREEN[4] = {10, 0, 5, 0};
                const int BUTTON_COLOR_BLUE[4] = {0, 0, 0, 10};

                static std::uniform_int_distribution<> button_color_dist(0, 3);

                int b1_color = button_color_dist(gen);
                int b2_color = button_color_dist(gen);
                int b3_color = button_color_dist(gen);
                int b4_color = button_color_dist(gen);

                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 20, BUTTON_COLOR_RED[b1_color], BUTTON_COLOR_GREEN[b1_color], BUTTON_COLOR_BLUE[b1_color]));
                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 19, BUTTON_COLOR_RED[b2_color], BUTTON_COLOR_GREEN[b2_color], BUTTON_COLOR_BLUE[b2_color]));
                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 18, BUTTON_COLOR_RED[b3_color], BUTTON_COLOR_GREEN[b3_color], BUTTON_COLOR_BLUE[b3_color]));
                ESP_ERROR_CHECK(led_strip_set_pixel(leds, 17, BUTTON_COLOR_RED[b4_color], BUTTON_COLOR_GREEN[b4_color], BUTTON_COLOR_BLUE[b4_color]));

                // switches
                const int SWITCH_COLOR_RED[3] = {20, 0, 10};
                const int SWITCH_COLOR_GREEN[3] = {0, 10, 5};

                static std::uniform_int_distribution<> switch_color_dist(0, 2);

                int switch_colors[4] = {switch_color_dist(gen), switch_color_dist(gen), switch_color_dist(gen), switch_color_dist(gen)};

                // int s1_color = switch_color_dist(gen);
                // int s2_color = switch_color_dist(gen);
                // int s3_color = switch_color_dist(gen);
                // int s4_color = switch_color_dist(gen);

                for (int i = 0; i < 4; i++) {
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, (16 - i), SWITCH_COLOR_RED[switch_colors[i]], SWITCH_COLOR_GREEN[switch_colors[i]], 0));
                }

                // ESP_ERROR_CHECK(led_strip_set_pixel(leds, 16, SWITCH_COLOR_RED[s1_color], SWITCH_COLOR_GREEN[s1_color], 0));
                // ESP_ERROR_CHECK(led_strip_set_pixel(leds, 15, SWITCH_COLOR_RED[s2_color], SWITCH_COLOR_GREEN[s2_color], 0));
                // ESP_ERROR_CHECK(led_strip_set_pixel(leds, 14, SWITCH_COLOR_RED[s3_color], SWITCH_COLOR_GREEN[s3_color], 0));
                // ESP_ERROR_CHECK(led_strip_set_pixel(leds, 13, SWITCH_COLOR_RED[s4_color], SWITCH_COLOR_GREEN[s4_color], 0));

                ESP_ERROR_CHECK(led_strip_refresh(leds));

                ButtonKey button;
                uint8_t starting_switch_state = get_switch_state();

                while (1) {
                    if (get_pressed_button(&button)) {
                        uint8_t button_state = get_button_state();


                        // change button colors
                        if ((button_state & 0b1) == 0b1) {
                            b1_color++;
                            if (b1_color > 3) {
                               b1_color = 0; 
                            }
                            ESP_ERROR_CHECK(led_strip_set_pixel(leds, 20, BUTTON_COLOR_RED[b1_color], BUTTON_COLOR_GREEN[b1_color], BUTTON_COLOR_BLUE[b1_color]));
                        }
                        if ((button_state & 0b10) == 0b10) {
                            b2_color++;
                            if (b2_color > 3) {
                               b2_color = 0; 
                            }
                            ESP_ERROR_CHECK(led_strip_set_pixel(leds, 19, BUTTON_COLOR_RED[b2_color], BUTTON_COLOR_GREEN[b2_color], BUTTON_COLOR_BLUE[b2_color]));
                        }
                        if ((button_state & 0b100) == 0b100) {
                            b3_color++;
                            if (b3_color > 3) {
                               b3_color = 0; 
                            }
                            ESP_ERROR_CHECK(led_strip_set_pixel(leds, 18, BUTTON_COLOR_RED[b3_color], BUTTON_COLOR_GREEN[b3_color], BUTTON_COLOR_BLUE[b3_color]));
                        }
                        if ((button_state & 0b1000) == 0b1000) {
                            b4_color++;
                            if (b4_color > 3) {
                               b4_color = 0; 
                            }
                            ESP_ERROR_CHECK(led_strip_set_pixel(leds, 17, BUTTON_COLOR_RED[b4_color], BUTTON_COLOR_GREEN[b4_color], BUTTON_COLOR_BLUE[b4_color]));
                        }
                        ESP_ERROR_CHECK(led_strip_refresh(leds));
                    }
                    if (get_module_time() <= 0) {
                        // check button state
                        if (b1_color == 0 && b2_color == 1 && b3_color == 2 && b4_color == 3) {
                            // check switch state
                            uint8_t switch_state = get_switch_state();
                            bool correct = true;
                            ESP_LOGI(TAG, "starting switch state: %i, current switch state: %i", starting_switch_state, switch_state);
                            


                            for (int i = 0; i < 4; i++) {
                                if (switch_colors[i] == 2) {
                                    // ESP_LOGI(TAG, "color yellow triggered: %i", (((switch_state ^ starting_switch_state) >> i) & 0b1));
                                    if ((((switch_state ^ starting_switch_state) >> i) & 0b1) != 1) {
                                        correct = false;
                                        break;
                                    }
                                } else {
                                    // ESP_LOGI(TAG, "color green or red triggered: %i != %i", ((switch_state >> i) & 0b1), (switch_colors[i] & 0b1));
                                    if (((switch_state >> i) & 0b1) != (switch_colors[i] & 0b1)) {
                                        correct = false;
                                        break;
                                    }
                                }
                            }

                            if (correct) {
                                solved_puzzles++;
                            } else {
                                strike("Wrong switch state! (step 3, puzzle 4)");
                            }
                        } else {
                            strike("Wrong button state! (step 3, puzzle 4)");
                        }
                        break;
                    }
                    vTaskDelay(pdMS_TO_TICKS(10));
                }


                break;
            }
            case 5: {
                lcd_print(&lcd, "Ummm");
                set_module_time(20000);
                start_module_timer();

                std::uniform_int_distribution<> indicator_number_dist(0, 5);

                uint8_t green_indicators = indicator_number_dist(gen);
                uint8_t red_indicators = indicator_number_dist(gen);
                uint8_t yellow_indicators = indicator_number_dist(gen);
                uint8_t blue_indicators = indicator_number_dist(gen);
                // ESP_LOGI(TAG, "Green: %i, Red: %i, Yellow: %i, Blue: %i", green_indicators, red_indicators, yellow_indicators, blue_indicators);

                std::set<uint8_t> indicators;

                while (indicators.size() < green_indicators) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 0, 10, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 15, 0, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators + yellow_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 10, 5, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators + yellow_indicators + blue_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 0, 0, 10));
                    }
                }

                ESP_ERROR_CHECK(led_strip_refresh(leds));

                uint8_t green_pressed = 0;
                uint8_t red_pressed = 0;
                uint8_t yellow_pressed = 0;
                uint8_t blue_pressed = 0;

                ButtonKey button;

                while (1) {
                    if (get_pressed_button(&button)) {
                        uint8_t button_state = get_button_state();

                        if ((button_state & 0b1) == 0b1) {
                            green_pressed++; 
                            if (green_pressed > green_indicators) {
                                strike("Green button pressed too many times! (step 3, puzzle 5)");
                                break;
                            }
                        }
                        if ((button_state & 0b10) == 0b10) {
                            red_pressed++;
                            if (red_pressed > red_indicators) {
                                strike("Red button pressed too many times! (step 3, puzzle 5)");
                                break;
                            }
                        }
                        if ((button_state & 0b100) == 0b100) {
                            yellow_pressed++;
                            if (yellow_pressed > yellow_indicators) {
                                strike("Yellow button pressed too many times! (step 3, puzzle 5)");
                                break;
                            }
                        }
                        if ((button_state & 0b1000) == 0b1000) {
                            blue_pressed++;
                            if (blue_pressed > blue_indicators) {
                                strike("Blue button pressed too many times! (step 3, puzzle 5)");
                            }
                        }
                    }
                    if (get_module_time() <= 0) {
                        // check for correct button presses
                        if (green_pressed == green_indicators && red_pressed == red_indicators && yellow_pressed == yellow_indicators && blue_pressed == blue_indicators) {
                            uint8_t switch_state = get_switch_state();

                            // ESP_LOGI(TAG, "%i, %i, %i, %i", (~green_indicators & 0b1), ((~red_indicators & 0b1) << 1), ((~yellow_indicators & 0b1) << 2), ((~blue_indicators & 0b1) << 3));
                            // ESP_LOGI(TAG, "correct: %i, inputted: %i", ((~green_indicators & 0b1) | (~red_indicators & 0b1) << 1) | (~yellow_indicators & 0b1) << 2) | (~blue_indicators & 0b1) << 3)), switch_state);
                            // check for correct switch states
                            if (((~green_indicators & 0b1) + ((~red_indicators & 0b1) << 1) + ((~yellow_indicators & 0b1) << 2) + ((~blue_indicators & 0b1) << 3)) == switch_state) {
                                solved_puzzles++;
                            } else {
                                strike("Wrong switch state! (step 3, puzzle 5)");
                            }
                        } else {
                            strike("Wrong button state! (step 3, puzzle 5)");
                        }
                        break;
                    }
                    vTaskDelay(pdMS_TO_TICKS(10));
                }



                break;
            }
            case 6: {
                lcd_print(&lcd, "Plank");
                set_module_time(15000);
                start_module_timer();

                // red, purple, blue, white, green, yellow
                const uint8_t COLORS_RED[6] = {20, 10, 0, 5, 0, 10};
                const uint8_t COLORS_GREEN[6] = {0, 0, 0, 5, 10, 5};
                const uint8_t COLORS_BLUE[6] = {0, 10, 10, 5, 0, 0};

                int button_colors[4] = {led_color_dist(gen), led_color_dist(gen), led_color_dist(gen), led_color_dist(gen)};

                for (int i = 0; i < 4; i++) {
                    ESP_ERROR_CHECK(led_strip_set_pixel(leds, (20 - i), COLORS_RED[button_colors[i]], COLORS_GREEN[button_colors[i]], COLORS_BLUE[button_colors[i]]));
                }
                ESP_ERROR_CHECK(led_strip_refresh(leds));

                bool buttons_cycling[4] = {true, true, true, true};
                TickType_t lastCycleTime = xTaskGetTickCount();

                ButtonKey button;
                while (1) {
                    if (get_pressed_button(&button)) {
                        uint8_t button_state = get_button_state();
                        bool correct = true;

                        if (button_state & 0b1) {
                            if (button_colors[0] != 4) {
                                correct = false;
                            } else {
                                buttons_cycling[0] = false;
                            }
                        }
                        if (button_state & 0b10) {
                            if (button_colors[1] != 0) {
                                correct = false;
                            } else {
                                buttons_cycling[1] = false;
                            }
                        }
                        if (button_state & 0b100) {
                            if (button_colors[2] != 5) {
                                correct = false;
                            } else {
                                buttons_cycling[2] = false;
                            }
                        }
                        if (button_state & 0b1000) {
                            if (button_colors[3] != 2) {
                                correct = false;
                            } else {
                                buttons_cycling[3] = false;
                            }
                        }

                        if (!correct) {
                            strike("Paused buttons at the wrong time! (step 3, puzzle 6)");
                            break;
                        }
                    }
                    if ((xTaskGetTickCount() - lastCycleTime) >= pdMS_TO_TICKS(500)) {
                        ESP_LOGI(TAG, "Cycling LEDs");
                        for (int i = 0; i < 4; i++) {
                            if (buttons_cycling[i]) {
                                button_colors[i]++;
                                if (button_colors[i] > 5) {
                                    button_colors[i] = 0;
                                }
                                ESP_ERROR_CHECK(led_strip_set_pixel(leds, (20 - i), COLORS_RED[button_colors[i]], COLORS_GREEN[button_colors[i]], COLORS_BLUE[button_colors[i]]));
                            }
                        }
                        ESP_ERROR_CHECK(led_strip_refresh(leds));

                        lastCycleTime = xTaskGetTickCount();
                    }
                    if (get_module_time() <= 0) {
                        bool failed = false;
                        for (int i = 0; i < sizeof(buttons_cycling); i++) {
                            if (buttons_cycling[i] == true) {
                                failed = true;
                            }
                        }
                        if (failed) {
                            strike("Ran out of time! (step 3, puzzle 6)");
                        } else {
                            solved_puzzles++;
                        }
                        break;
                    }


                    vTaskDelay(pdMS_TO_TICKS(10));
                }

                // *** ADD PART 2 ***


                break;
            }
            case 7: {
                lcd_print(&lcd, "What");
                set_module_time(20000);
                start_module_timer();

                std::uniform_int_distribution<> math_number_dist(1, 9);
                std::uniform_int_distribution<> math_operation_dist(0, 3);

                std::vector<float> math_numbers;
                std::vector<int> math_operations;

                // ESP_LOGI(TAG, "math_numbers: %f, %f, %f, %f", math_numbers[0], math_numbers[1], math_numbers[2], math_numbers[3]);
                // ESP_LOGI(TAG, "math_operations: %i, %i, %i", math_operations[0], math_operations[1], math_operations[2]);

                std::map<int, char> operation_map = {
                    {0, '+'},
                    {1, '-'},
                    {2, '*'},
                    {3, '/'},
                };

                int expression_answer = -1;
                std::string display_expression;

                while (expression_answer < 0) {
                    math_numbers = {static_cast<float>(math_number_dist(gen)), static_cast<float>(math_number_dist(gen)), static_cast<float>(math_number_dist(gen)), static_cast<float>(math_number_dist(gen))};
                    math_operations = {math_operation_dist(gen), math_operation_dist(gen), math_operation_dist(gen)};

                    display_expression = std::to_string(static_cast<int>(math_numbers[0]));
                    for (int i = 0; i < 3; i++) {
                        display_expression += operation_map[math_operations[i]];
                        display_expression += std::to_string(static_cast<int>(math_numbers[i + 1]));
                    }

                    for (int j = 0; j < 3; j++) {
                        bool found = false;
                        for (int i = 0; i < math_operations.size(); i++) {
                            if (math_operations[i] == 2) {
                                // ESP_LOGI(TAG, "i = %i, condensing %f * %f to %f", i, math_numbers[i], math_numbers[i + 1], (math_numbers[i] * math_numbers[i+1]));
                                math_numbers[i] = math_numbers[i] * math_numbers[i + 1];
                                math_numbers.erase(math_numbers.begin() + i + 1);
                                math_operations.erase(math_operations.begin() + i);
                                found = true;
                                break;
                            } else if (math_operations[i] == 3) {
                                // ESP_LOGI(TAG, "i = %i, condensing %f / %f to %f", i, math_numbers[i], math_numbers[i + 1], (math_numbers[i] / math_numbers[i+1]));
                                math_numbers[i] = math_numbers[i] / math_numbers[i + 1];
                                math_numbers.erase(math_numbers.begin() + i + 1);
                                math_operations.erase(math_operations.begin() + i);
                                found = true;
                                break;
                            }
                        }
                        if (found) {
                            continue;
                        }
                        for (int i = 0; i < math_operations.size(); i++) {
                            if (math_operations[i] == 0) {
                                // ESP_LOGI(TAG, "i = %i, condensing %f + %f to %f", i, math_numbers[i], math_numbers[i + 1], (math_numbers[i] + math_numbers[i+1]));
                                math_numbers[i] = math_numbers[i] + math_numbers[i + 1];
                                math_numbers.erase(math_numbers.begin() + i + 1);
                                math_operations.erase(math_operations.begin() + i);
                                found = true;
                                break;
                            } else if (math_operations[i] == 1) {
                                // ESP_LOGI(TAG, "i = %i, condensing %f - %f to %f", i, math_numbers[i], math_numbers[i + 1], (math_numbers[i] - math_numbers[i+1]));
                                math_numbers[i] = math_numbers[i] - math_numbers[i + 1];
                                math_numbers.erase(math_numbers.begin() + i + 1);
                                math_operations.erase(math_operations.begin() + i);
                                found = true;
                                break;
                            }
                        }
                    }
                    expression_answer = static_cast<int>(std::floor(math_numbers[0]));
                }

                // display expression
                lcd_set_cursor(&lcd, 1, 2);
                lcd_print(&lcd, display_expression.c_str());

                // ESP_LOGI(TAG, "Display expression: %s", display_expression.c_str());
                // ESP_LOGI(TAG, "Solved expression answer: %i", static_cast<int>(expression_answer));


                // set LEDs
                std::uniform_int_distribution<> indicator_number_dist(0, 5);

                uint8_t green_indicators = indicator_number_dist(gen);
                uint8_t red_indicators = indicator_number_dist(gen);
                uint8_t blue_indicators = indicator_number_dist(gen);

                while (((expression_answer + (blue_indicators * 3) - red_indicators) * (green_indicators ^ 2)) < 0) {
                    green_indicators = indicator_number_dist(gen);
                    red_indicators = indicator_number_dist(gen);
                    blue_indicators = indicator_number_dist(gen);
                }


                // ESP_LOGI(TAG, "Green: %i, Red: %i, Blue: %i", green_indicators, red_indicators, blue_indicators);

                std::set<uint8_t> indicators;

                while (indicators.size() < green_indicators) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 0, 10, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 15, 0, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators + blue_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 0, 0, 10));
                    }
                }

                ESP_ERROR_CHECK(led_strip_refresh(leds));

                // *** ALSO IMPLEMENT LAST STEPS ***

                std::string answer_string = std::to_string(expression_answer);
                std::string entered_string;

                // ESP_LOGI(TAG, "Solved full answer: %s", answer_string.c_str());

                KeypadKey key;

                while (1) {
                    if (get_pressed_keypad(&key)) {
                        if (key == KeypadKey::star) {
                            // clear
                            entered_string = "";
                        } else if (key == KeypadKey::pound) {
                            // submit
                            if (entered_string != answer_string) {
                                strike("Entered answer is not correct! (step 3, puzzle 7)");
                            } else {
                                solved_puzzles++;
                            }
                            break;
                        } else {
                            entered_string += char_of_keypad_key(key);
                        }

                        lcd_clear(&lcd);
                        lcd_set_cursor(&lcd, 1, 1);
                        lcd_print(&lcd, "What");
                        lcd_set_cursor(&lcd, 1, 2);
                        lcd_print(&lcd, display_expression.c_str());
                        lcd_set_cursor(&lcd, 1, 3);
                        lcd_print(&lcd, entered_string.c_str());
                    }
                    if (get_module_time() <= 0) {
                        strike("Ran out of time! (step 3, puzzle 7)");
                        break;
                    }


                    vTaskDelay(pdMS_TO_TICKS(10));
                }

                break;
            }
            case 8: {
                lcd_print(&lcd, "Plink");
                set_module_time(10000);
                start_module_timer();

                std::uniform_int_distribution<> indicator_number_dist(0, 4);

                uint8_t green_indicators = indicator_number_dist(gen);
                uint8_t red_indicators = indicator_number_dist(gen);
                uint8_t yellow_indicators = indicator_number_dist(gen);
                uint8_t blue_indicators = indicator_number_dist(gen);
                uint8_t purple_indicators = indicator_number_dist(gen);
                // ESP_LOGI(TAG, "Green: %i, Red: %i, Yellow: %i, Blue: %i", green_indicators, red_indicators, yellow_indicators, blue_indicators);

                std::set<uint8_t> indicators;

                while (indicators.size() < green_indicators) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 0, 10, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 15, 0, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators + yellow_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 10, 5, 0));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators + yellow_indicators + blue_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 0, 0, 10));
                    }
                }

                while (indicators.size() < (green_indicators + red_indicators + yellow_indicators + blue_indicators + purple_indicators)) {
                    uint8_t led = led_picker_dist(gen);
                    if (indicators.insert(led).second) {
                        ESP_ERROR_CHECK(led_strip_set_pixel(leds, led, 10, 0, 5));
                    }
                }

                ESP_ERROR_CHECK(led_strip_refresh(leds));

                std::uniform_int_distribution<> answer_color_dist(0, 4);

                int answer_color = answer_color_dist(gen);
                std::string color_string;
                std::string answer_num;

                switch (answer_color) {
                    case 0: {
                        answer_num = std::to_string(green_indicators);
                        color_string = "Green";
                        break;
                    }
                    case 1: {
                        answer_num = std::to_string(red_indicators);
                        color_string = "Red";
                        break;
                    }
                    case 2: {
                        answer_num = std::to_string(yellow_indicators);
                        color_string = "Yellow";
                        break;
                    }
                    case 3: {
                        answer_num = std::to_string(blue_indicators);
                        color_string = "Blue";
                        break;
                    }
                    case 4: {
                        answer_num = std::to_string(purple_indicators);
                        color_string = "Purple";
                        break;
                    }
                }
                // ESP_LOGI(TAG, "color string: %s", color_string.c_str());

                lcd_set_cursor(&lcd, 1, 2);
                lcd_print(&lcd, color_string.c_str());

                std::string entered_string;

                KeypadKey key;
                while (1) {
                    if (get_pressed_keypad(&key)) {
                        if (key == KeypadKey::star) {
                            // clear
                            entered_string = "";
                        } else if (key == KeypadKey::pound) {
                            // submit
                            if (entered_string != answer_num) {
                                strike("Entered answer is not correct! (step 3, puzzle 7)");
                            } else {
                                solved_puzzles++;
                            }
                            break;
                        } else {
                            entered_string += char_of_keypad_key(key);
                        }

                        lcd_clear(&lcd);
                        lcd_set_cursor(&lcd, 1, 1);
                        lcd_print(&lcd, "Plink");
                        lcd_set_cursor(&lcd, 1, 2);
                        lcd_print(&lcd, color_string.c_str());
                        lcd_set_cursor(&lcd, 1, 3);
                        lcd_print(&lcd, entered_string.c_str());
                    }
                    if (get_module_time() <= 0) {
                        strike("Ran out of time! (step 3, puzzle 7)");
                        break;
                    }


                    vTaskDelay(pdMS_TO_TICKS(10));
                }

                break;
            }


        }
        vTaskDelay(pdMS_TO_TICKS(1000));
        clean_bomb();
    }
}