#include <fstream>
#include <iostream>
#include <vector>

#include "transform/bmp.hpp"

enum BayerPattern { GRBG, RGGB, BGGR, GBRG };
const int IN_WIDTH = 1936;
const int IN_HEIGHT = 1088;
const int IN_SIZE = IN_WIDTH * IN_HEIGHT;
const int OUT_WIDTH = 1280;
const int OUT_HEIGHT = 720;
const int OUT_SIZE = OUT_WIDTH * OUT_HEIGHT;
const BayerPattern RAW_TYPE = GBRG;
const int COLOR_DEPTH = 12;

// const float red_gain = 1.2f;  // Adjust these values as necessary
// const float green_gain = 0.5f;
// const float blue_gain = 0.95f;

int main() {
    std::ifstream in_image;
    in_image.open("./test.bin", std::ios::in | std::ios::binary);
    auto image = std::vector<std::vector<uint16_t>>(
        IN_HEIGHT, std::vector<uint16_t>(IN_WIDTH, 0));
    for (int y = 0; y < IN_HEIGHT; y++)
        for (int x = 0; x < IN_WIDTH; x++) {
            uint8_t buf[2] = {0};
            in_image.read((char*)buf, sizeof(buf));

            image[y][x] = buf[0] + ((uint16_t)buf[1] << 8);
        }

    BayerPattern raw_type = RAW_TYPE;
    uint8_t* data =
        new uint8_t[OUT_WIDTH * OUT_HEIGHT * 3];  // RGB24格式像素数据
    for (int32_t y = 0; y < OUT_HEIGHT; ++y) {
        for (int32_t x = 0; x < OUT_WIDTH; ++x) {
            int32_t index = (y * OUT_WIDTH + x) * 3;

            uint16_t red = 0, green = 0, blue = 0;
            uint16_t cache[25] = {0};
            for (int i = 0; i < 5; i++)
                for (int j = 0; j < 5; j++) {
                    cache[i * 5 + j] = image[y + i][x + j];
                }
            //       data        case GRBG   case RGGB   case BGGR   case GBRG
            //  00 01 02 03 04   G R G R G   R G R G R   B G B G B   G B G B G
            //  05 06 07 08 09   B G B G B   G B G B G   G R G R G   R G R G R
            //  10 11 12 13 14   G R G R G   R G R G R   B G B G B   G B G B G
            //  15 16 17 18 19   B G B G B   G B G B G   G R G R G   R G R G R
            //  20 21 22 23 24   G R G R G   R G R G R   B G B G B   G B G B G

            switch (raw_type) {
                case GRBG:
                    green = cache[12];
                    blue =
                        (cache[7] + cache[17]) / 2 + cache[12] -
                        ((cache[2] + cache[6] + cache[8] + cache[12]) / 4 +
                         (cache[22] + cache[16] + cache[18] + cache[12]) / 4) /
                            2;
                    red = (cache[11] + cache[13]) / 2 + cache[12] -
                          ((cache[10] + cache[6] + cache[16] + cache[12]) / 4 +
                           (cache[14] + cache[8] + cache[18] + cache[12]) / 4) /
                              2;
                    break;

                case RGGB:
                    red = cache[12];
                    green = (cache[7] + cache[11] + cache[13] + cache[17]) / 4;
                    blue =
                        (cache[6] + cache[8] + cache[16] + cache[18]) / 4 +
                        (cache[7] + cache[11] + cache[13] + cache[17]) / 4 -
                        ((cache[1] + cache[5] + cache[7] + cache[11]) / 4 +
                         (cache[3] + cache[7] + cache[9] + cache[13]) / 4 +
                         (cache[11] + cache[15] + cache[17] + cache[18]) / 4 +
                         (cache[13] + cache[17] + cache[19] + cache[23]) / 4) /
                            4;
                    break;

                case BGGR:
                    blue = cache[12];
                    green = (cache[7] + cache[11] + cache[13] + cache[17]) / 4;
                    red =
                        (cache[6] + cache[8] + cache[16] + cache[18]) / 4 +
                        (cache[7] + cache[11] + cache[13] + cache[17]) / 4 -
                        ((cache[1] + cache[5] + cache[7] + cache[11]) / 4 +
                         (cache[3] + cache[7] + cache[9] + cache[13]) / 4 +
                         (cache[11] + cache[15] + cache[17] + cache[18]) / 4 +
                         (cache[13] + cache[17] + cache[19] + cache[23]) / 4) /
                            4;
                    break;
                case GBRG:
                    green = cache[12];
                    red =
                        (cache[7] + cache[17]) / 2 + cache[12] -
                        ((cache[2] + cache[6] + cache[8] + cache[12]) / 4 +
                         (cache[22] + cache[16] + cache[18] + cache[12]) / 4) /
                            2;
                    blue =
                        (cache[11] + cache[13]) / 2 + cache[12] -
                        ((cache[10] + cache[6] + cache[16] + cache[12]) / 4 +
                         (cache[14] + cache[8] + cache[18] + cache[12]) / 4) /
                            2;
                    break;
            }
            printf(
                "x=%4d, y=%4d, red=0x%03x, green=0x%03x, blue=0x%03x, "
                "raw_type=%d\n",
                x, y, red, green, blue, raw_type);

            switch (raw_type) {
                case 0:
                    raw_type = RGGB;
                    break;
                case 1:
                    raw_type = GRBG;
                    break;
                case 2:
                    raw_type = GBRG;
                    break;
                case 3:
                    raw_type = BGGR;
                    break;
            }

            data[index + 0] = red >> (COLOR_DEPTH - 8);    // R
            data[index + 1] = green >> (COLOR_DEPTH - 8);  // G
            data[index + 2] = blue >> (COLOR_DEPTH - 8);   // B
        }

        if (y % 2) {
            raw_type = RAW_TYPE;
        } else {
            switch (RAW_TYPE) {
                case 0:
                    raw_type = BGGR;
                    break;
                case 1:
                    raw_type = GBRG;
                    break;
                case 2:
                    raw_type = GRBG;
                    break;
                case 3:
                    raw_type = RGGB;
                    break;
            }
        }
    }

    // for (int i = 0; i < OUT_WIDTH * OUT_HEIGHT * 3; i += 3) {
    //     data[i + 0] = std::min(255, static_cast<int>(data[i + 0] *
    //     red_gain)); data[i + 1] = std::min(255, static_cast<int>(data[i + 1]
    //     * green_gain)); data[i + 2] = std::min(255, static_cast<int>(data[i +
    //     2] * blue_gain));
    // }

    write_bmp("test1.bmp", data, OUT_WIDTH, OUT_HEIGHT);
    delete[] data;

    return 0;
}
