2 files changed, 23 insertions, 19 deletions

diff --git a/quantum/visualizer/visualizer.c b/quantum/visualizer/visualizer.c
index 98cd7ba55..5fbd12031 100644
--- a/quantum/visualizer/visualizer.c
+++ b/quantum/visualizer/visualizer.c
@@ -154,6 +154,14 @@ void stop_all_keyframe_animations(void) {
     }
 }
 
+static uint8_t get_num_running_animations(void) {
+    uint8_t count = 0;
+    for (int i=0;i<MAX_SIMULTANEOUS_ANIMATIONS;i++) {
+        count += animations[i] ? 1 : 0;
+    }
+    return count;
+}
+
 static bool update_keyframe_animation(keyframe_animation_t* animation, visualizer_state_t* state, systemticks_t delta, systemticks_t* sleep_time) {
     // TODO: Clean up this messy code
     dprintf("Animation frame%d, left %d, delta %d\n", animation->current_frame,
@@ -228,14 +236,6 @@ bool keyframe_no_operation(keyframe_animation_t* animation, visualizer_state_t*
     return false;
 }
 
-bool enable_visualization(keyframe_animation_t* animation, visualizer_state_t* state) {
-    (void)animation;
-    (void)state;
-    dprint("User visualizer inited\n");
-    visualizer_enabled = true;
-    return false;
-}
-
 // TODO: Optimize the stack size, this is probably way too big
 static DECLARE_THREAD_STACK(visualizerThreadStack, 1024);
 static DECLARE_THREAD_FUNCTION(visualizerThread, arg) {
@@ -276,13 +276,15 @@ static DECLARE_THREAD_FUNCTION(visualizerThread, arg) {
 
     systemticks_t sleep_time = TIME_INFINITE;
     systemticks_t current_time = gfxSystemTicks();
+    bool force_update = true;
 
     while(true) {
         systemticks_t new_time = gfxSystemTicks();
         systemticks_t delta = new_time - current_time;
         current_time = new_time;
         bool enabled = visualizer_enabled;
-        if (!same_status(&state.status, &current_status)) {
+        if (force_update || !same_status(&state.status, &current_status)) {
+            force_update = false;
             if (visualizer_enabled) {
                 if (current_status.suspended) {
                     stop_all_keyframe_animations();
@@ -320,10 +322,10 @@ static DECLARE_THREAD_FUNCTION(visualizerThread, arg) {
 #ifdef EMULATOR
         draw_emulator();
 #endif
-        // The animation can enable the visualizer
-        // And we might need to update the state when that happens
-        // so don't sleep
-        if (enabled != visualizer_enabled) {
+        // Enable the visualizer when the startup or the suspend animation has finished
+        if (!visualizer_enabled && state.status.suspended == false && get_num_running_animations() == 0) {
+            visualizer_enabled = true;
+            force_update = true;
             sleep_time = 0;
         }
 
diff --git a/quantum/visualizer/visualizer.h b/quantum/visualizer/visualizer.h
index f37ce8416..463934849 100644
--- a/quantum/visualizer/visualizer.h
+++ b/quantum/visualizer/visualizer.h
@@ -130,20 +130,22 @@ void run_next_keyframe(keyframe_animation_t* animation, visualizer_state_t* stat
 // Does nothing, useful for adding delays
 bool keyframe_no_operation(keyframe_animation_t* animation, visualizer_state_t* state);
 
-// Call this once, when the initial animation has finished, alternatively you can call it
-// directly from the initalize_user_visualizer function (the animation can be null)
-bool enable_visualization(keyframe_animation_t* animation, visualizer_state_t* state);
-
 // The master can set userdata which will be transferred to the slave
 #ifdef VISUALIZER_USER_DATA_SIZE
 void visualizer_set_user_data(void* user_data);
 #endif
 
 // These functions have to be implemented by the user
-void initialize_user_visualizer(visualizer_state_t* state);
+// Called regularly each time the state has changed (but not every scan loop)
 void update_user_visualizer_state(visualizer_state_t* state, visualizer_keyboard_status_t* prev_status);
+// Called when the computer goes to suspend, will also stop calling update_user_visualizer_state
 void user_visualizer_suspend(visualizer_state_t* state);
+// You have to start at least one animation as a response to the following two functions
+// When the animation has finished the visualizer will resume normal operation and start calling the
+// update_user_visualizer_state again
+// Called when the keyboard boots up
+void initialize_user_visualizer(visualizer_state_t* state);
+// Called when the computer resumes from a suspend
 void user_visualizer_resume(visualizer_state_t* state);
 
-
 #endif /* VISUALIZER_H */