// Copyright (c) 2022-2024 Advanced Micro Devices, Inc. All rights reserved.

#include <string.h>
#include <math.h>
#include "color_bg.h"
#include "vpe_priv.h"

struct csc_vector {
    float x;
    float y;
    float z;
};

struct csc_table {
    struct csc_vector rgb_offset; // RGB offset
    struct csc_vector red_coef;   // RED coefficient
    struct csc_vector green_coef; // GREEN coefficient
    struct csc_vector blue_coef;  // BLUE coefficient
};


const double bt_709_rgb_xyz_matrix[] = {
    0.135676572958501,   0.117645247657296,   0.059378179384203,
    0.069958232931727,   0.235290495314592,   0.023751271753681,
    0.006359839357430,   0.039215082552432,   0.312725078090138
};

const double bt_601_rgb_xyz_matrix[] = {
    0.129468377303939,   0.120169907240092,   0.063061715455969,
    0.069871822671967,   0.230648692928563,   0.028479484399470,
    0.006165160823997,   0.036826261896157,   0.315308577279846
};

const double bt_2020_rgb_xyz_matrix[] = {
    0.209559197891125,   0.047578961279863,   0.055561840829013,
    0.086428369751707,   0.223061365529709,   0.019510264718585,
    0.000000000000000,   0.009235916013150,   0.349064083986850
};

const double bt_709_xyz_rgb_matrix[] = {
    9.850972467794900,    -4.672897196261683,    -1.515534225814599,
   -2.946029289607537,     5.702028879962675,     0.126307165371354,
    0.169088388136759,    -0.619990756501448,     3.212679374598414
};

const double bt_601_xyz_rgb_matrix[] = {
    10.656544932293809,   -5.288117709127149,    -1.653672548215019,
   -3.249384680406732,     6.011485965740993,     0.106904010143450,
    0.171144655726832,    -0.598710197023623,     3.191344462670923
};

const double bt_2020_xyz_rgb_matrix[] = {
    5.217784765870115,    -1.081066212086299,    -0.770110277731489,
   -2.026396206177778,     4.913316828677627,     0.047928710680581,
    0.053616587979668,    -0.130001864005497,     2.863535322904176
};

static struct csc_table bgcolor_to_rgbfull_table[COLOR_SPACE_MAX] = {
    [COLOR_SPACE_YCBCR601] =
        {
            {0.0f, -0.5f, -0.5f},
            {1.0f, 0.0f, 1.402f},
            {1.0f, -0.344136286f, -0.714136286f},
            {1.0f, 1.772f, 0.0f},
        },
    [COLOR_SPACE_YCBCR709] =
        {
            {0.0f, -0.5f, -0.5f},
            {1.0f, 0.0f, 1.5748f},
            {1.0f, -0.187324273f, -0.468124273f},
            {1.0f, 1.8556f, 0.0f},
        },
    [COLOR_SPACE_YCBCR601_LIMITED] =
        {
            {-0.0625f, -0.5f, -0.5f},
            {1.164383562f, 0.0f, 1.596026786f},
            {1.164383562f, -0.39176229f, -0.812967647f},
            {1.164383562f, 2.017232143f, 0.0f},
        },
    [COLOR_SPACE_YCBCR709_LIMITED] =
        {
            {-0.0625f, -0.5f, -0.5f},
            {1.164383562f, 0.0f, 1.792741071f},
            {1.164383562f, -0.213248614f, -0.532909329f},
            {1.164383562f, 2.112401786f, 0.0f},
        },
    [COLOR_SPACE_2020_YCBCR] =
        {
            {0.0f, -512.f / 1023.f, -512.f / 1023.f},
            {1.0f, 0.0f, 1.4746f},
            {1.0f, -0.164553127f, -0.571353127f},
            {1.0f, 1.8814f, 0.0f},
        },
    [COLOR_SPACE_2020_YCBCR_LIMITED] =
        {
            {-0.0625f, -0.5f, -0.5f},
            {1.167808219f, 0.0f, 1.683611384f},
            {1.167808219f, -0.187877063f, -0.652337331f},
            {1.167808219f, 2.148071652f, 0.0f},
        },
    [COLOR_SPACE_SRGB_LIMITED] =
        {
            {-0.0626221f, -0.0626221f, -0.0626221f},
            {1.167783652f, 0.0f, 0.0f},
            {0.0f, 1.167783652f, 0.0f},
            {0.0f, 0.0, 1.167783652f},
        },
    [COLOR_SPACE_2020_RGB_LIMITEDRANGE] = {
        {-0.0626221f, -0.0626221f, -0.0626221f},
        {1.167783652f, 0.0f, 0.0f},
        {0.0f, 1.167783652f, 0.0f},
        {0.0f, 0.0, 1.167783652f},
    }};

static double clip_double(double x)
{
    if (x < 0.0)
        return 0.0;
    else if (x > 1.0)
        return 1.0;
    else
        return x;
}

static float clip_float(float x)
{
    if (x < 0.0f)
        return 0.0f;
    else if (x > 1.0f)
        return 1.0f;
    else
        return x;
}

static void color_multiply_matrices_double(double *mResult, double *M1,
    double *M2, unsigned int Rows1, unsigned int Cols1, unsigned int Cols2)
{
    unsigned int i, j, k;

    for (i = 0; i < Rows1; i++) {
        for (j = 0; j < Cols2; j++) {
            mResult[(i * Cols2) + j] = 0.0;
            for (k = 0; k < Cols1; k++)
                mResult[(i * Cols2) + j] = mResult[(i * Cols2) + j] +
                    M1[(i * Cols1) + k] * M2[(k * Cols2) + j];
        }
    }
}

static void set_gamut_remap_matrix(double* res, enum color_space src_cs, enum color_space dst_cs) {

    double rgb_to_xyz[9] = { 0.0 };
    double xyz_to_rgb[9] = { 0.0 };

    switch (src_cs)
    {
    case COLOR_SPACE_SRGB:
    case COLOR_SPACE_SRGB_LIMITED:
    case COLOR_SPACE_MSREF_SCRGB:
    case COLOR_SPACE_YCBCR709_LIMITED:
    case COLOR_SPACE_YCBCR709:
    case COLOR_SPACE_YCBCR_JFIF:
    case COLOR_SPACE_RGB_JFIF:
        memcpy(rgb_to_xyz, bt_709_rgb_xyz_matrix, 9 * sizeof(double));
        break;
    case COLOR_SPACE_YCBCR601:
    case COLOR_SPACE_YCBCR601_LIMITED:
        memcpy(rgb_to_xyz, bt_601_rgb_xyz_matrix, 9 * sizeof(double));
        break;
    case COLOR_SPACE_2020_RGB_FULLRANGE:
    case COLOR_SPACE_2020_RGB_LIMITEDRANGE:
    case COLOR_SPACE_2020_YCBCR:
    case COLOR_SPACE_2020_YCBCR_LIMITED:
        memcpy(rgb_to_xyz, bt_2020_rgb_xyz_matrix, 9 * sizeof(double));
        break;
    default:
        VPE_ASSERT(0);
        break;
    }

    switch (dst_cs)
    {
    case COLOR_SPACE_SRGB:
    case COLOR_SPACE_SRGB_LIMITED:
    case COLOR_SPACE_MSREF_SCRGB:
    case COLOR_SPACE_YCBCR709_LIMITED:
    case COLOR_SPACE_YCBCR709:
    case COLOR_SPACE_YCBCR_JFIF:
    case COLOR_SPACE_RGB_JFIF:
        memcpy(xyz_to_rgb, bt_709_xyz_rgb_matrix, 9 * sizeof(double));
        break;
    case COLOR_SPACE_YCBCR601:
    case COLOR_SPACE_YCBCR601_LIMITED:
        memcpy(xyz_to_rgb, bt_601_xyz_rgb_matrix, 9 * sizeof(double));
        break;
    case COLOR_SPACE_2020_RGB_FULLRANGE:
    case COLOR_SPACE_2020_RGB_LIMITEDRANGE:
    case COLOR_SPACE_2020_YCBCR:
    case COLOR_SPACE_2020_YCBCR_LIMITED:
        memcpy(xyz_to_rgb, bt_2020_xyz_rgb_matrix, 9 * sizeof(double));
        break;
    default:
        VPE_ASSERT(0);
        break;
    }

    color_multiply_matrices_double(res, xyz_to_rgb, rgb_to_xyz, 3, 3, 3);

}

bool vpe_bg_csc(struct vpe_color *bg_color, enum color_space cs)
{
    struct csc_table *entry             = &bgcolor_to_rgbfull_table[cs];
    float             csc_final[3]      = {0};
    float             csc_mm[3][4]      = {0};
    bool              output_is_clipped = false;

    memcpy(&csc_mm[0][0], &entry->red_coef, sizeof(struct csc_vector));
    memcpy(&csc_mm[1][0], &entry->green_coef, sizeof(struct csc_vector));
    memcpy(&csc_mm[2][0], &entry->blue_coef, sizeof(struct csc_vector));

    csc_mm[0][3] = entry->rgb_offset.x * csc_mm[0][0] + entry->rgb_offset.y * csc_mm[0][1] +
                   entry->rgb_offset.z * csc_mm[0][2];

    csc_mm[1][3] = entry->rgb_offset.x * csc_mm[1][0] + entry->rgb_offset.y * csc_mm[1][1] +
                   entry->rgb_offset.z * csc_mm[1][2];

    csc_mm[2][3] = entry->rgb_offset.x * csc_mm[2][0] + entry->rgb_offset.y * csc_mm[2][1] +
                   entry->rgb_offset.z * csc_mm[2][2];

    csc_final[0] = csc_mm[0][0] * bg_color->ycbcra.y + csc_mm[0][1] * bg_color->ycbcra.cb +
                   csc_mm[0][2] * bg_color->ycbcra.cr + csc_mm[0][3];

    csc_final[1] = csc_mm[1][0] * bg_color->ycbcra.y + csc_mm[1][1] * bg_color->ycbcra.cb +
                   csc_mm[1][2] * bg_color->ycbcra.cr + csc_mm[1][3];

    csc_final[2] = csc_mm[2][0] * bg_color->ycbcra.y + csc_mm[2][1] * bg_color->ycbcra.cb +
                   csc_mm[2][2] * bg_color->ycbcra.cr + csc_mm[2][3];

    // switch to RGB components
    bg_color->rgba.a = bg_color->ycbcra.a;
    bg_color->rgba.r = clip_float(csc_final[0]);
    bg_color->rgba.g = clip_float(csc_final[1]);
    bg_color->rgba.b = clip_float(csc_final[2]);
    if ((bg_color->rgba.r != csc_final[0]) || (bg_color->rgba.g != csc_final[1]) ||
        (bg_color->rgba.b != csc_final[2])) {
        output_is_clipped = true;
    }
    bg_color->is_ycbcr = false;
    return output_is_clipped;
}

bool vpe_is_global_bg_blend_applied(struct stream_ctx *stream_ctx)
{

    return (stream_ctx->stream.blend_info.blending)  &&
        (stream_ctx->stream.blend_info.global_alpha) &&
        (stream_ctx->stream.blend_info.global_alpha_value != 1.0);
}

struct gamma_coefs {
    float a0;
    float a1;
    float a2;
    float a3;
    float user_gamma;
    float user_contrast;
    float user_brightness;
};

// srgb, 709, G24
static const int32_t numerator01[] = {31308, 180000, 0};
static const int32_t numerator02[] = {12920, 4500, 0};
static const int32_t numerator03[] = {55, 99, 0};
static const int32_t numerator04[] = {55, 99, 0};
static const int32_t numerator05[] = {2400, 2222, 2400};

static bool build_coefficients(struct gamma_coefs *coefficients, enum color_transfer_func type)
{
    uint32_t index = 0;
    bool     ret   = true;

    if (type == TRANSFER_FUNC_SRGB)
        index = 0;
    else if (type == TRANSFER_FUNC_BT709)
        index = 1;
    else if (type == TRANSFER_FUNC_BT1886)
        index = 2;
    else {
        ret = false;
        goto release;
    }

    coefficients->a0         = (float)numerator01[index] / 10000000.0f;
    coefficients->a1         = (float)numerator02[index] / 1000.0f;
    coefficients->a2         = (float)numerator03[index] / 1000.0f;
    coefficients->a3         = (float)numerator04[index] / 1000.0f;
    coefficients->user_gamma = (float)numerator05[index] / 1000.0f;

release:
    return ret;
}

static double translate_to_linear_space(
    double arg, double a0, double a1, double a2, double a3, double gamma)
{
    double linear;
    double base;

    a0 *= a1;
    if (arg <= -a0) {
        base   = (a2 - arg) / (1.0 + a3);
        linear = -pow(base, gamma);
    } else if ((-a0 <= arg) && (arg <= a0))
        linear = arg / a1;
    else {
        base   = (a2 + arg) / (1.0 + a3);
        linear = pow(base, gamma);
    }

    return linear;
}

// for 709 & sRGB
static void compute_degam(enum color_transfer_func tf, double inY, double *outX, bool clip)
{
    double             ret;
    struct gamma_coefs coefs = {0};

    build_coefficients(&coefs, tf);

    ret = translate_to_linear_space(inY, (double)coefs.a0, (double)coefs.a1, (double)coefs.a2,
        (double)coefs.a3, (double)coefs.user_gamma);

    if (clip) {
        ret = clip_double(ret);
    }
    *outX = ret;
}

static double get_maximum_fp(double a, double b)
{
    if (a > b)
        return a;
    return b;
}

static void compute_depq(double inY, double *outX, bool clip)
{
    double M1 = 0.159301758;
    double M2 = 78.84375;
    double C1 = 0.8359375;
    double C2 = 18.8515625;
    double C3 = 18.6875;

    double nPowM2;
    double base;
    double one      = 1.0;
    double zero     = 0.0;
    bool   negative = false;
    double ret;

    if (inY < zero) {
        inY      = -inY;
        negative = true;
    }
    nPowM2 = pow(inY, one / M2);
    base   = get_maximum_fp(nPowM2 - C1, zero) / (C2 - C3 * nPowM2);
    ret    = pow(base, one / M1);
    if (clip) {
        ret = clip_double(ret);
    }
    if (negative)
        ret = -ret;

    *outX = ret;
}

bool vpe_is_limited_cs(enum color_space cs)
{
    bool is_limited = false;

    switch (cs)
    {
    case COLOR_SPACE_RGB601:
    case COLOR_SPACE_SRGB:
    case COLOR_SPACE_2020_RGB_FULLRANGE:
    case COLOR_SPACE_MSREF_SCRGB:
    case COLOR_SPACE_YCBCR601:
    case COLOR_SPACE_YCBCR709:
    case COLOR_SPACE_YCBCR_JFIF:
    case COLOR_SPACE_RGB_JFIF:
    case COLOR_SPACE_2020_YCBCR:
        is_limited = false;
        break;
    case COLOR_SPACE_RGB601_LIMITED:
    case COLOR_SPACE_SRGB_LIMITED:
    case COLOR_SPACE_YCBCR601_LIMITED:
    case COLOR_SPACE_YCBCR709_LIMITED:
    case COLOR_SPACE_2020_RGB_LIMITEDRANGE:
    case COLOR_SPACE_2020_YCBCR_LIMITED:
        is_limited = true;
        break;
    default:
        VPE_ASSERT(0);
        is_limited = false;
        break;
    }
    return is_limited;
}

void vpe_bg_degam(struct transfer_func *output_tf, struct vpe_color *bg_color)
{

    double degam_r = (double)bg_color->rgba.r;
    double degam_g = (double)bg_color->rgba.g;
    double degam_b = (double)bg_color->rgba.b;

    // de-gam
    switch (output_tf->tf) {

    case TRANSFER_FUNC_PQ2084:
        compute_depq((double)bg_color->rgba.r, &degam_r, true);
        compute_depq((double)bg_color->rgba.g, &degam_g, true);
        compute_depq((double)bg_color->rgba.b, &degam_b, true);
        break;
    case TRANSFER_FUNC_SRGB:
    case TRANSFER_FUNC_BT709:
    case TRANSFER_FUNC_BT1886:
        compute_degam(output_tf->tf, (double)bg_color->rgba.r, &degam_r, true);
        compute_degam(output_tf->tf, (double)bg_color->rgba.g, &degam_g, true);
        compute_degam(output_tf->tf, (double)bg_color->rgba.b, &degam_b, true);
        break;
    case TRANSFER_FUNC_LINEAR:
        break;
    default:
        VPE_ASSERT(0);
        break;
    }
    bg_color->rgba.r = (float)degam_r;
    bg_color->rgba.g = (float)degam_g;
    bg_color->rgba.b = (float)degam_b;
}

void vpe_bg_inverse_gamut_remap(
    enum color_space output_cs, struct transfer_func *output_tf, struct vpe_color *bg_color)
{

        double bg_rgb[3] = { 0.0 };
        double final_bg_rgb[3] = { 0.0 };
        double matrix[9] = { 0.0 };
        bg_rgb[0] = (double)bg_color->rgba.r;
        bg_rgb[1] = (double)bg_color->rgba.g;
        bg_rgb[2] = (double)bg_color->rgba.b;

        switch (output_tf->tf) {
        case TRANSFER_FUNC_LINEAR:
            /* Since linear output uses Bt709, and this conversion is only needed
             * when the tone mapping is enabled on (Bt2020) input, it is needed to
             * apply the reverse of Bt2020 -> Bt709 on the background color to
             * cancel out the effect of Bt2020 -> Bt709 on the background color.
             */
            set_gamut_remap_matrix(matrix, COLOR_SPACE_SRGB, COLOR_SPACE_2020_RGB_FULLRANGE);
            color_multiply_matrices_double(final_bg_rgb, matrix, bg_rgb, 3, 3, 1);

            bg_color->rgba.r = (float)clip_double(final_bg_rgb[0]);
            bg_color->rgba.g = (float)clip_double(final_bg_rgb[1]);
            bg_color->rgba.b = (float)clip_double(final_bg_rgb[2]);

            break;
        case TRANSFER_FUNC_PQ2084:
        case TRANSFER_FUNC_SRGB:
        case TRANSFER_FUNC_BT709:
        case TRANSFER_FUNC_BT1886:
            break;
        default:
            VPE_ASSERT(0);
            break;
        }

}

void vpe_inverse_output_csc(enum color_space output_cs, struct vpe_color *bg_color)
{
    enum color_space bgcolor_cs = COLOR_SPACE_YCBCR709;

    switch (output_cs) {
        // output is ycbr cs, follow output's setting
    case COLOR_SPACE_YCBCR601:
    case COLOR_SPACE_YCBCR709:
    case COLOR_SPACE_YCBCR601_LIMITED:
    case COLOR_SPACE_YCBCR709_LIMITED:
    case COLOR_SPACE_2020_YCBCR:
    case COLOR_SPACE_2020_YCBCR_LIMITED:
        bgcolor_cs = output_cs;
        break;
        // output is RGB cs, follow output's range
        // but need yuv to rgb csc
    case COLOR_SPACE_SRGB_LIMITED:
    case COLOR_SPACE_RGB601_LIMITED:
        bgcolor_cs = COLOR_SPACE_YCBCR709_LIMITED;
        break;
    case COLOR_SPACE_2020_RGB_LIMITEDRANGE:
        bgcolor_cs = COLOR_SPACE_2020_YCBCR_LIMITED;
        break;
    case COLOR_SPACE_SRGB:
    case COLOR_SPACE_MSREF_SCRGB:
    case COLOR_SPACE_RGB601:
        bgcolor_cs = COLOR_SPACE_YCBCR709;
        break;
    case COLOR_SPACE_2020_RGB_FULLRANGE:
        bgcolor_cs = COLOR_SPACE_2020_YCBCR;
        break;
    default:
        // should revise the newly added CS
        // and set corresponding bgcolor_cs accordingly
        VPE_ASSERT(0);
        bgcolor_cs = COLOR_SPACE_YCBCR709;
        break;
    }

    // input is [0-0xffff]
    // convert bg color to RGB full range for use inside pipe
    vpe_bg_csc(bg_color, bgcolor_cs);
}

