`timescale 1ns / 1ps

module SaturationCorrection #(
    parameter reg [4:0] COLOR_DEPTH = 8
) (
    input wire clk,
    input wire reset,

    input wire in_en,
    input wire [COLOR_DEPTH - 1 : 0] in_data[3],
    output wire out_ready,
    output wire out_receive,

    output reg out_en,
    output reg [COLOR_DEPTH - 1 : 0] out_data[3],
    input wire in_ready,
    input wire in_receive,

    input wire enable,
    input wire signed [31:0] saturation_inc
);
  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 signed [31:0] data_cal[3], data_cache[3];
  // wire signed [31:0] max, min, delta, value, light, saturation, alpha;
  reg signed [31:0] max, min, delta, value, light, saturation, alpha;

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

  always @(*) begin
    case (state)
      READ_DATA: nextState = in_en ? 3 : READ_DATA;
      3: nextState = CALC_DATA;
      CALC_DATA: nextState = (calState >= 5 || !enable) ? SEND_DATA : CALC_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 max = data_cache[0] > data_cache[1]?
  //              (data_cache[0] > data_cache[2] ? data_cache[0] : data_cache[2]):
  //              (data_cache[1] > data_cache[2] ? data_cache[1] : data_cache[2]);
  // assign min = data_cache[0] < data_cache[1]?
  //              (data_cache[0] < data_cache[2] ? data_cache[0] : data_cache[2]):
  //              (data_cache[1] < data_cache[2] ? data_cache[1] : data_cache[2]);
  // assign delta = max - min;
  // assign value = max + min;
  // assign light = value >>> 1;
  // // // assign saturation = (light <= 128) ? (delta <<< 8) / value : (delta <<< 8) / (512 - value);
  // assign saturation = (delta <<< 8) / max;
  // assign alpha = (saturation_inc[31] == 0)
  //                 ? ((saturation_inc + saturation >= 256)
  //                   ? (65536 / saturation) - 256 : (65536 / (256 - saturation_inc)) - 256)
  //                   : (saturation_inc);

  always @(posedge clk) begin
    if (reset) begin
      calState <= 0;
      out_en <= 0;
      out_data[0] <= 0;
      out_data[1] <= 0;
      out_data[2] <= 0;

      min <= 0;
      max <= 0;
      delta <= 0;
      value <= 0;
      light <= 0;
      saturation <= 0;
      alpha <= 0;

      data_cal[0] <= 0;
      data_cal[1] <= 0;
      data_cal[2] <= 0;
      data_cache[0] <= 0;
      data_cache[1] <= 0;
      data_cache[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]};
          end
        end

        CALC_DATA: begin
          if (enable) begin
            if (calState == 0) begin
              max <= data_cache[0] > data_cache[1]?
               (data_cache[0] > data_cache[2] ? data_cache[0] : data_cache[2]):
               (data_cache[1] > data_cache[2] ? data_cache[1] : data_cache[2]);
              min <= data_cache[0] < data_cache[1]?
               (data_cache[0] < data_cache[2] ? data_cache[0] : data_cache[2]):
               (data_cache[1] < data_cache[2] ? data_cache[1] : data_cache[2]);

              calState <= 1;
            end else if (calState == 1) begin
              delta <= max - min;
              value <= max + min;

              calState <= 2;
            end else if (calState == 2) begin
              light <= value >>> 1;
              saturation <= (delta <<< 8) / max;

              calState <= 3;
            end else if (calState == 3) begin
              alpha <= (saturation_inc[31] == 0) ? ((saturation_inc + saturation >= 255)
                    ? (65536 / saturation) - 255 : (65536 / (256 - saturation_inc)) - 255)
                    : (saturation_inc);

              calState <= 4;
            end else if (calState == 4) begin
              if (saturation_inc[31] == 0) begin
                data_cal[0] <= (data_cache[0] << 8) + ((data_cache[0] - light) * alpha);
                data_cal[1] <= (data_cache[1] << 8) + ((data_cache[1] - light) * alpha);
                data_cal[2] <= (data_cache[2] << 8) + ((data_cache[2] - light) * alpha);
              end else begin
                data_cal[0] <= (light << 8) + (data_cache[0] - light) * (256 + alpha);
                data_cal[1] <= (light << 8) + (data_cache[1] - light) * (256 + alpha);
                data_cal[2] <= (light << 8) + (data_cache[2] - light) * (256 + alpha);
              end

              calState <= 5;
            end else begin
              calState <= 0;
            end
          end
        end

        SEND_DATA: begin
          if (in_ready && !in_receive) begin
            out_en <= 1;
            if (enable && delta != 0) 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
