`timescale 1ns / 1ps

// 三通道图像合成一个RGB图像
module GreyWorld #(
    parameter reg [4:0] COLOR_DEPTH = 8,
    parameter reg [31:0] IM_SIZE = 1920 * 1080
) (
    input wire clk,
    input wire reset,

    input wire in_en,
    input wire [7:0] in_data[3],  // 0:R 1:G 2:B
    output wire out_ready,
    output wire out_receive,

    // 输出相关
    input wire in_ready,
    input wire in_receive,
    output reg out_en,
    output reg [COLOR_DEPTH - 1:0] out_data[3],

    // 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]
);

  reg [2:0] state, nextState;
  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] data_cal[3], data_cache[3];
  reg  [31:0] cnt_pexels;
  wire [39:0] average;

  always @(posedge clk) begin
    if (reset) state <= READ_DATA;
    else state <= nextState;
  end

  always @(*) begin
    case (state)
      READ_DATA: nextState = in_en ? CALC_DATA : READ_DATA;
      CALC_DATA: nextState = SEND_DATA;
      SEND_DATA: nextState = in_receive ? READ_DATA : SEND_DATA;
      default:   nextState = READ_DATA;
    endcase
  end

  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 ;

  always @(posedge clk) begin
    if (reset) begin
      red_total <= 0;
      green_total <= 0;
      blue_total <= 0;
      cnt_flame <= flame_rate;
      cnt_pexels <= 0;
      r_white_gain[0] <= white_gain[0];
      r_white_gain[1] <= white_gain[1];
      r_white_gain[2] <= white_gain[2];
      data_cache[0] <= 0;
      data_cache[1] <= 0;
      data_cache[2] <= 0;

      out_en <= 0;
      out_data[0] <= 0;
      out_data[1] <= 0;
      out_data[2] <= 0;
    end else begin
      case (state)
        READ_DATA: begin
          if (in_en) begin
            data_cache[0] <= {24'b0, in_data[0]};
            data_cache[1] <= {24'b0, in_data[1]};
            data_cache[2] <= {24'b0, in_data[2]};

            if (cnt_flame == flame_rate) begin
              red_total   <= red_total + {24'b0, in_data[0]};
              green_total <= green_total + {24'b0, in_data[1]};
              blue_total  <= blue_total + {24'b0, in_data[2]};
            end

            if (cnt_pexels <= IM_SIZE) begin
              cnt_pexels <= cnt_pexels + 1;
            end else begin
              cnt_pexels <= 0;
              if (cnt_flame < flame_rate) cnt_flame <= cnt_flame + 1;
              else cnt_flame <= 0;
            end
          end
        end

        CALC_DATA: begin
          if (cnt_pexels >= IM_SIZE && cnt_flame == flame_rate) begin
            r_white_gain[0] <= { ( average / {8'b0, red_total } ) }[31:0];
            r_white_gain[1] <= { ( average / {8'b0, green_total } ) }[31:0];
            r_white_gain[2] <= { ( average / {8'b0, blue_total } ) }[31:0];
          end

          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]);
        end

        SEND_DATA: begin
          if (in_ready) begin
            out_en <= 1;
            if (enable) begin
              out_data[0] <= (|data_cal[0][31:16]) ? 255 : (data_cal[0] > 0 ? data_cal[0][15:8] : 0);
              out_data[1] <= (|data_cal[1][31:16]) ? 255 : (data_cal[1] > 0 ? data_cal[1][15:8] : 0);
              out_data[2] <= (|data_cal[2][31:16]) ? 255 : (data_cal[2] > 0 ? data_cal[2][15:8] : 0);
            end else begin
              out_data[0] <= data_cache[0][7:0];
              out_data[1] <= data_cache[1][7:0];
              out_data[2] <= data_cache[2][7:0];
            end
          end else out_en <= 0;
        end
        default: ;
      endcase
    end
  end



endmodule