#include "gles2_debayer.h"

#include "camera.h"
#include "dcp.h"
#include "gl_util.h"
#include "matrix.h"
#include <stdlib.h>

#define VERTEX_ATTRIBUTE 0
#define TEX_COORD_ATTRIBUTE 1

char* pixel_format_precision(uint32_t pixfmt){
	// Large bit formats should use high precision to remove artifacting
	switch (pixfmt) {
                case V4L2_PIX_FMT_SBGGR10:
                case V4L2_PIX_FMT_SGBRG10:
                case V4L2_PIX_FMT_SGRBG10:
                case V4L2_PIX_FMT_SRGGB10:
                case V4L2_PIX_FMT_SBGGR12:
                case V4L2_PIX_FMT_SGBRG12:
                case V4L2_PIX_FMT_SGRBG12:
                case V4L2_PIX_FMT_SRGGB12:
                case V4L2_PIX_FMT_SBGGR10P:
                case V4L2_PIX_FMT_SGBRG10P:
                case V4L2_PIX_FMT_SGRBG10P:
                case V4L2_PIX_FMT_SRGGB10P:
			return "#define PRECISION_HIGH\n";
        }
	return "#define PRECISION_MEDIUM\n";
}

GLES2Debayer *
gles2_debayer_new(int format)
{
        // Cannot run on format 0 (Undefined)
        assert(format != 0);

        uint32_t pixfmt = libmegapixels_format_to_v4l_pixfmt(format);

        int shader = 0;
        int texture_mode = 0;

        switch (pixfmt) {
                case V4L2_PIX_FMT_SBGGR8:
                case V4L2_PIX_FMT_SGBRG8:
                case V4L2_PIX_FMT_SGRBG8:
                case V4L2_PIX_FMT_SRGGB8:
                case V4L2_PIX_FMT_SBGGR10:
                case V4L2_PIX_FMT_SGBRG10:
                case V4L2_PIX_FMT_SGRBG10:
                case V4L2_PIX_FMT_SRGGB10:
                case V4L2_PIX_FMT_SBGGR12:
                case V4L2_PIX_FMT_SGBRG12:
                case V4L2_PIX_FMT_SGRBG12:
                case V4L2_PIX_FMT_SRGGB12:
                        shader = SHADER_DEBAYER;
                        texture_mode = TEXTURE_BAYER;
                        break;
                case V4L2_PIX_FMT_SBGGR10P:
                case V4L2_PIX_FMT_SGBRG10P:
                case V4L2_PIX_FMT_SGRBG10P:
                case V4L2_PIX_FMT_SRGGB10P:
                        shader = SHADER_DEBAYER_PACKED;
                        texture_mode = TEXTURE_BAYER;
                        break;
                case V4L2_PIX_FMT_YUYV:
                case V4L2_PIX_FMT_YVYU:
                case V4L2_PIX_FMT_UYVY:
                case V4L2_PIX_FMT_VYUY:
                        shader = SHADER_YUV;
                        texture_mode = TEXTURE_YUV;
                        break;
        }

        if (shader == 0) {
                return NULL;
        }

        GLuint frame_buffer;
        glGenFramebuffers(1, &frame_buffer);
        check_gl();

        char format_def[64];
        snprintf(format_def,
                 64,
                 "#define CFA_%s\n#define BITS_%d\n#define FMT_%s\n",
                 libmegapixels_format_cfa(format),
                 libmegapixels_format_bits_per_pixel(format),
                 libmegapixels_format_name(format));


	const GLchar *vert_def[1] = {pixel_format_precision(pixfmt)};
        const GLchar *def[1] = { format_def };

        char shader_vertex[64];
        char shader_fragment[64];
        snprintf(shader_vertex,
                 64,
                 "/org/postmarketos/Megapixels/%s.vert",
                 gles2_shader_to_vertex(shader));
        snprintf(shader_fragment,
                 64,
                 "/org/postmarketos/Megapixels/%s.frag",
                 gles2_shader_to_fragment(shader));

        GLuint shaders[] = {
                gl_util_load_shader(shader_vertex, GL_VERTEX_SHADER, vert_def, 1),
                gl_util_load_shader(shader_fragment, GL_FRAGMENT_SHADER, def, 1),
        };
        printf("Using shader %s and %s\n", shader_vertex, shader_fragment);
        GLuint program = gl_util_link_program(shaders, 2);
        glBindAttribLocation(program, VERTEX_ATTRIBUTE, "vert");
        glBindAttribLocation(program, TEX_COORD_ATTRIBUTE, "tex_coord");
        check_gl();

        GLES2Debayer *self = malloc(sizeof(GLES2Debayer));
        self->format = format;
        self->shader = shader;
        self->texture_mode = texture_mode;

        self->frame_buffer = frame_buffer;
        self->program = program;

        self->uniform_transform = glGetUniformLocation(self->program, "transform");
        self->uniform_pixel_size = glGetUniformLocation(self->program, "pixel_size");
        self->uniform_padding_ratio =
                glGetUniformLocation(self->program, "padding_ratio");
        self->uniform_texture = glGetUniformLocation(self->program, "texture");
        self->uniform_color_matrix =
                glGetUniformLocation(self->program, "color_matrix");
        self->uniform_inv_gamma = glGetUniformLocation(self->program, "inv_gamma");
        self->uniform_blacklevel = glGetUniformLocation(self->program, "blacklevel");
        if (libmegapixels_format_bits_per_pixel(self->format) == 10)
                self->uniform_row_length =
                        glGetUniformLocation(self->program, "row_length");
        check_gl();

        self->quad = gl_util_new_quad();

        return self;
}

void
gles2_debayer_free(GLES2Debayer *self)
{
        glDeleteFramebuffers(1, &self->frame_buffer);

        glDeleteProgram(self->program);

        free(self);
}

void
gles2_debayer_use(GLES2Debayer *self)
{
        assert(self != NULL);
        glUseProgram(self->program);
        check_gl();

        gl_util_bind_quad(self->quad);
}

void
gles2_debayer_set_shading(GLES2Debayer *self,
                          float red,
                          float blue,
                          float blacklevel)
{
        if (self->forward_matrix[0]) {
                float wb[9] = {
                        // clang-format off
                        red, 0,   0,
                        0,   1,   0,
                        0,   0, blue,
                        // clang-format on
                };
                float colormat[9];
                float xyz[9];
                float xyzd65[9];
                multiply_matrices(wb, self->forward_matrix, xyz);
                multiply_matrices(xyz, XYZD50_to_D65, xyzd65);
                multiply_matrices(xyzd65, XYZD65_to_sRGB, colormat);
                glUniformMatrix3fv(
                        self->uniform_color_matrix, 1, GL_FALSE, colormat);
        } else if (self->shader == SHADER_YUV) {
                glUniformMatrix3fv(
                        self->uniform_color_matrix, 1, GL_FALSE, YUV_to_RGB);
        } else {
                glUniformMatrix3fv(
                        self->uniform_color_matrix, 1, GL_FALSE, IDENTITY);
        }
        check_gl();

        glUniform1f(self->uniform_blacklevel, blacklevel);
        check_gl();
}

void
gles2_debayer_configure(GLES2Debayer *self,
                        const uint32_t dst_width,
                        const uint32_t dst_height,
                        const uint32_t src_width,
                        const uint32_t src_height,
                        const uint32_t rotation,
                        const bool mirrored,
                        struct MPCameraCalibration calibration)
{
        glViewport(0, 0, (int)dst_width, (int)dst_height);
        check_gl();

        /* Rotation matrix for orientation correction */
        GLfloat rotation_list[4] = { 0, -1, 0, 1 };
        int rotation_index = 4 - (int)rotation / 90;

        GLfloat sin_rot = rotation_list[rotation_index];
        GLfloat cos_rot = rotation_list[(rotation_index + 1) % 4];
        GLfloat scale_x = mirrored ? 1 : -1;
        GLfloat matrix[9] = {
                // clang-format off
		cos_rot * scale_x,  sin_rot, 0,
		-sin_rot * scale_x, cos_rot, 0,
		0,                        0, 1,
                // clang-format on
        };
        glUniformMatrix3fv(self->uniform_transform, 1, GL_FALSE, matrix);
        check_gl();

        GLfloat pixel_size_x = 1.0f / src_width;
        GLfloat pixel_size_y = 1.0f / src_height;
        glUniform2f(self->uniform_pixel_size, pixel_size_x, pixel_size_y);
        check_gl();

        /* Color calibration curves and matrices */
        float gamma = 1.0f;
        for (int i = 2; i < calibration.tone_curve_length; i += 2) {
                float g = calibration.tone_curve[i + 1] / calibration.tone_curve[i];
                if (g > gamma) {
                        gamma = g;
                }
        }
        glUniform1f(self->uniform_inv_gamma, 1.0f / gamma);
        check_gl();

        for (int i = 0; i < 9; i++) {
                self->forward_matrix[i] = calibration.forward_matrix_1[i];
        }

        gles2_debayer_set_shading(self, 1.0f, 1.0f, 0.0f);

        GLuint row_length =
                libmegapixels_mode_width_to_bytes(self->format, src_width);
        if (libmegapixels_format_bits_per_pixel(self->format) == 10) {
                assert(src_width % 4 == 0);
                glUniform1f(self->uniform_row_length, row_length);
                check_gl();
        }

        GLuint padding_bytes =
                libmegapixels_mode_width_to_padding(self->format, src_width);
        GLfloat padding_ratio = (float)row_length / (row_length + padding_bytes);
        glUniform1f(self->uniform_padding_ratio, padding_ratio);
}

void
gles2_debayer_process(GLES2Debayer *self, GLuint dst_id, GLuint source_id)
{
        glBindFramebuffer(GL_FRAMEBUFFER, self->frame_buffer);
        glBindTexture(GL_TEXTURE_2D, dst_id);
        glFramebufferTexture2D(
                GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, dst_id, 0);
        check_gl();

        assert(glCheckFramebufferStatus(GL_FRAMEBUFFER) == GL_FRAMEBUFFER_COMPLETE);

        glActiveTexture(GL_TEXTURE0);
        glBindTexture(GL_TEXTURE_2D, source_id);
        glUniform1i(self->uniform_texture, 0);
        check_gl();

        gl_util_draw_quad(self->quad);
}

/* Returns which fragment file to use for a particular shader */
const char*
gles2_shader_to_fragment(int shader)
{
        switch (shader) {
                case SHADER_DEBAYER:
                        return "debayer";
                case SHADER_DEBAYER_PACKED:
                        return "debayer_packed";
                case SHADER_YUV:
                        return "yuv";
                default:
                        return "solid";
        }
}

/* Returns which vertex file to use for a particular shader */
const char*
gles2_shader_to_vertex(int shader)
{
        switch (shader) {
                case SHADER_DEBAYER:
                case SHADER_DEBAYER_PACKED:
                        return "debayer";
                case SHADER_YUV:
                        return "yuv";
                default:
                        return "solid";
        }
}
