`timescale 1ns / 1ps

// 三通道图像合成一个RGB图像
module WhiteBalance #(
    parameter reg [4:0] IN_DEPTH = 12,  // 输入图像的色深
    parameter reg [4:0] OUT_DEPTH = 8,  // 输出图像的色深
    parameter reg [8:0] BUFF_SIZE = 32,
    parameter reg [31:0] RAM_BEGIN_ADDR = 0,
    parameter reg [31:0] IM_SIZE = 1920 * 1080
) (
    input wire clk,
    input wire reset,

    input wire in_en,
    input wire [15: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 wire out_en,
    output wire [OUT_DEPTH - 1:0] out_data[3],

    input wire enable,
    input wire [8:0] flame_rate,
    input wire [31:0] white_gain[3]
);

  wire ram_read_en, ram_write_en, ram_read_ready, ram_write_ready;
  wire [31:0] ram_read_data, ram_read_addr, ram_write_data, ram_write_addr;

  reg [8:0] cnt_flame;
  reg [31:0] cnt_in_pic, cnt_out_pic, cnt_pixels;
  reg [31:0] red_total, green_total, blue_total;
  reg [31:0] r_white_gain[3];

  wire is_finish_cal;
  always @(*) begin
    if (cnt_in_pic >= IM_SIZE) is_finish_cal = 1;
    else if (cnt_out_pic >= IM_SIZE) is_finish_cal = 0;
    else is_finish_cal = is_finish_cal;
  end

  assign ram_read_addr = cnt_out_pic;
  assign ram_write_addr = cnt_in_pic;
  assign cnt_pixels = (is_finish_cal)
                      ? IM_SIZE - cnt_out_pic + cnt_in_pic
                      : cnt_in_pic - cnt_out_pic;

  assign out_ready = (!in_en && !reset && cnt_pixels < IM_SIZE) ? 1 : 0;
  assign out_receive = (in_en && !reset && cnt_in_pic < IM_SIZE) ? 1 : 0;

  DiffWidthSyncFIFO #(
      .DATA_WIDTH (8),
      .DATA_DEPTH (12),
      .READ_DEPTH (4),
      .WRITE_DEPTH(3)
  ) in_fifo (
      .clk(clk),
      .reset(reset),
      .read_ready(ram_write_ready),
      .read_en(ram_write_en),
      .read_data(ram_write_data),
      .write_ready(in_ready),
      .write_en(in_en),
      .write_data(in_data)
  );

  DiffWidthSyncFIFO #(
      .DATA_WIDTH (8),
      .DATA_DEPTH (12),
      .READ_DEPTH (3),
      .WRITE_DEPTH(4)
  ) out_fifo (
      .clk(clk),
      .reset(reset),
      .read_ready(out_ready),
      .read_en(out_en),
      .read_data(out_data),
      .write_ready(ram_read_ready),
      .write_en(ram_read_en),
      .write_data(ram_read_data)
  );

  SDRAM inst_RAM (
      .clk(clk),
      .read_en(ram_read_en),
      .read_addr(ram_read_addr),
      .read_data(ram_read_data),
      .read_ready(ram_read_ready),
      .write_en(ram_write_en),
      .write_addr(ram_write_addr),
      .write_data(ram_write_data),
      .write_ready(ram_write_ready)
  );

  // calculate white gain
  always @(posedge clk) begin
    if (reset) begin
      red_total <= 0;
      green_total <= 0;
      blue_total <= 0;
      cnt_flame <= 0;
      cnt_in_pic <= 0;
      r_white_gain[0] <= white_gain[0];
      r_white_gain[1] <= white_gain[1];
      r_white_gain[2] <= white_gain[2];
    end else begin
      if (in_en && cnt_pixels < IM_SIZE) begin
        red_total   <= red_total + in_data[0];
        green_total <= green_total + in_data[1];
        blue_total  <= blue_total + in_data[2];
        if (cnt_in_pic < IM_SIZE) cnt_in_pic <= cnt_in_pic + 1;
        else cnt_in_pic <= 0;
      end else begin
        
      end
    end
  end

  // calculate out data
  always @(posedge clk) begin
    if (reset) begin
      cnt_out_pic <= 0;
    end else begin
      if (out_ready && is_finish_cal) begin
        out_en <= 1;
      end else begin
        out_en <= 0;
      end
    end
  end

endmodule