maybe optimize the GreyWorld?

Color/GreyWorld.sv

@@ -23,19 +23,21 @@ module GreyWorld #(
     // Gain: red = 0.803881, green = 0.885894, blue = 1.594308
     input wire enable,
     input wire [8:0] flame_rate,
-    input wire [31:0] white_gain[3]
+    input wire [15:0] white_gain[3]
 );
 
-  reg [2:0] state, nextState;
+  reg [2:0] state, nextState, calState;
   localparam reg [2:0] READ_DATA = 0;
   localparam reg [2:0] CALC_DATA = 1;
   localparam reg [2:0] SEND_DATA = 2;
 
   reg [8:0] cnt_flame;
-  reg [31:0] red_total, green_total, blue_total, r_white_gain[3];
+  reg [31:0] red_total, green_total, blue_total;
+  reg [39:0] r_white_gain[3];
   reg [31:0] data_cal[3], data_cache[3];
   reg [31:0] cnt_pexels;
-  wire [39:0] average;
+  reg [39:0] average;
+  reg isCal;
 
   always @(posedge clk) begin
     if (reset) state <= READ_DATA;
@@ -45,7 +47,7 @@ module GreyWorld #(
   always @(*) begin
     case (state)
       READ_DATA: nextState = in_en ? CALC_DATA : READ_DATA;
-      CALC_DATA: nextState = SEND_DATA;
+      CALC_DATA: nextState = isCal ? CALC_DATA : SEND_DATA;
       SEND_DATA: nextState = in_receive ? READ_DATA : SEND_DATA;
       default:   nextState = READ_DATA;
     endcase
@@ -54,7 +56,8 @@ module GreyWorld #(
   assign out_ready   = (!in_en && state == READ_DATA && !reset) ? 1 : 0;
   assign out_receive = (in_en && state == READ_DATA && !reset) ? 1 : 0;
 
-  assign average = (({8'b0, red_total } + {8'b0, green_total } + {8'b0, blue_total }) << 8) / 3 ;
+  // assign average = (({8'b0, red_total } + {8'b0, green_total } + {8'b0, blue_total }) << 8) / 3 ;
+  // assign isCal = cnt_pexels >= IM_SIZE && cnt_flame == flame_rate;
 
   always @(posedge clk) begin
     if (reset) begin
@@ -63,6 +66,8 @@ module GreyWorld #(
       blue_total <= 0;
       cnt_flame <= flame_rate;
       cnt_pexels <= 0;
+      calState <= 4;
+
       r_white_gain[0] <= white_gain[0];
       r_white_gain[1] <= white_gain[1];
       r_white_gain[2] <= white_gain[2];
@@ -94,20 +99,41 @@ module GreyWorld #(
               cnt_pexels <= 0;
               if (cnt_flame < flame_rate) cnt_flame <= cnt_flame + 1;
               else cnt_flame <= 0;
+              if (cnt_pexels >= IM_SIZE && cnt_flame == flame_rate) isCal <= 1;
             end
           end
         end
 
         CALC_DATA: begin
-          if (cnt_pexels >= IM_SIZE && cnt_flame == flame_rate) begin
+          if (calState == 0) begin
-            r_white_gain[0] <= { ( average / {8'b0, red_total } ) }[31:0];
+            average <= {8'b0, red_total } + {8'b0, green_total };
-            r_white_gain[1] <= { ( average / {8'b0, green_total } ) }[31:0];
+            calState <= 1;
-            r_white_gain[2] <= { ( average / {8'b0, blue_total } ) }[31:0];
           end
-
+          else if (calState == 1) begin
+            average <= average + {8'b0, blue_total};
+            calState <= 2;
+          end
+          else if (calState == 2) begin
+            average <= average << 8;
+            calState <= 3;
+          end
+          else if (calState == 3) begin
+            r_white_gain[0] <= average / {8'b0, red_total};
+            r_white_gain[1] <= average / {8'b0, green_total};
+            r_white_gain[2] <= average / {8'b0, blue_total};
+            isCal <= 0;
+            calState <= 4;
+          end
+          else if (calState == 4) begin
             data_cal[0] <= (data_cache[0] * r_white_gain[0][31:0]);
             data_cal[1] <= (data_cache[1] * r_white_gain[1][31:0]);
             data_cal[2] <= (data_cache[2] * r_white_gain[2][31:0]);
+            calState <= 5;
+          end
+          else begin
+            calState <= isCal ? 0 : 4;
+          end
+
         end
 
         SEND_DATA: begin