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

#include "bmp.hpp"

enum BayerPattern { GRBG, RGGB, BGGR, GBRG };
const int IN_WIDTH = 1920;
const int IN_HEIGHT = 1080;
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 = RGGB;
const int COLOR_DEPTH = 10;

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

int main() {
    std::ifstream in_image;
    in_image.open("./test.RAW", 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 data_cache[9] = {0};
            for (int i = 0; i < 3; i++)
                for (int j = 0; j < 3; j++) {
                    data_cache[i * 3 + j] = image[y + i][x + j];
                }

            //  data       case 0      case 1      case 2      case 3
            //  0 1 2      G R G       R G R       B G B       G B G
            //  3 4 5      B G B       G B G       G R G       R G R
            //  6 7 8 	   G R G       R G R       B G B       G B G
            switch (raw_type) {
                case 0:  // Missing B, R on G
                    red = (data_cache[1] + data_cache[7]) / 2;
                    blue = (data_cache[3] + data_cache[5]) / 2;
                    green = data_cache[4];
                    break;

                case 1:  // Missing G, R on B
                    green = (data_cache[1] + data_cache[3] + data_cache[5] +
                             data_cache[7]) /
                            4;
                    red = (data_cache[0] + data_cache[2] + data_cache[6] +
                           data_cache[8]) /
                          4;
                    blue = data_cache[4];
                    break;

                case 2:  // Missing G, B on R
                    green = (data_cache[1] + data_cache[3] + data_cache[5] +
                             data_cache[7]) /
                            4;
                    blue = (data_cache[0] + data_cache[2] + data_cache[6] +
                            data_cache[8]) /
                           4;
                    red = data_cache[4];
                    break;

                case 3:  // Missing B, R on G
                    blue = (data_cache[1] + data_cache[7]) / 2;
                    red = (data_cache[3] + data_cache[5]) / 2;
                    green = data_cache[4];
                    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("test.bmp", data, OUT_WIDTH, OUT_HEIGHT);
    delete[] data;

    return 0;
}