/*
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

#version 460
#pragma shader_stage(compute)

#extension GL_EXT_scalar_block_layout : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_GOOGLE_include_directive : require

#define PB_UNALIGNED
#include "common.glsl"

layout (constant_id = 0) const int max_mbs_per_slice = 8;
layout (constant_id = 1) const int chroma_factor = 0;
layout (constant_id = 2) const int alpha_bits = 0;
layout (constant_id = 3) const int num_planes = 0;
layout (constant_id = 4) const int slices_per_picture = 0;
layout (constant_id = 5) const int max_quant = 0;

struct SliceData {
    uint32_t mbs_per_slice;
    int16_t coeffs[4][8 * 256];
};

struct SliceScore {
    ivec4 bits[16];
    ivec4 score[16];
    int total_bits[16];
    int total_score[16];
    int overquant;
    int buf_start;
    int quant;
};

layout(push_constant, scalar) uniform EncodeSliceInfo {
    u8buf bytestream;
    u8vec2buf seek_table;
};

layout (set = 0, binding = 0, scalar) readonly buffer SliceBuffer {
    SliceData slices[];
};
layout (set = 0, binding = 1, scalar) readonly buffer SliceScores {
    SliceScore scores[];
};
layout (set = 0, binding = 2, scalar) uniform ProresDataTables {
    int16_t qmat[128][64]; int16_t qmat_chroma[128][64];
};

#define CFACTOR_Y444 3

void encode_vlc_codeword(inout PutBitContext pb, uint codebook, int val)
{
    /* number of prefix bits to switch between Rice and expGolomb */
    uint switch_bits = (codebook & 3) + 1;
    uint rice_order  =  codebook >> 5;       /* rice code order */
    uint exp_order   = (codebook >> 2) & 7;  /* exp golomb code order */

    uint switch_val  = switch_bits << rice_order;

    if (val >= switch_val) {
        val -= int(switch_val - (1 << exp_order));
        int exponent = findMSB(val);

        put_bits(pb, exponent - exp_order + switch_bits, 0);
        put_bits(pb, exponent + 1, val);
    } else {
        int exponent = val >> rice_order;
        if (exponent != 0)
            put_bits(pb, exponent, 0);
        put_bits(pb, 1, 1);
        if (rice_order != 0)
            put_bits(pb, rice_order, zero_extend(val, rice_order));
    }
}

#define GET_SIGN(x)  ((x) >> 31)
#define MAKE_CODE(x) (((x) * 2) ^ GET_SIGN(x))

#define FIRST_DC_CB 0xB8 // rice_order = 5, exp_golomb_order = 6, switch_bits = 0

void encode_dcs(inout PutBitContext pb, bool is_chroma, int q)
{
    const uint8_t dc_codebook[7] = { U8(0x04), U8(0x28), U8(0x28), U8(0x4D), U8(0x4D), U8(0x70), U8(0x70) };

    uint slice = gl_GlobalInvocationID.x;
    uint plane = gl_GlobalInvocationID.y;
    uint blocks_per_mb = is_chroma && chroma_factor != CFACTOR_Y444 ? 2 : 4;
    uint blocks_per_slice = slices[slice].mbs_per_slice * blocks_per_mb;
    int codebook = 5;
    int scale = is_chroma ? qmat_chroma[q][0] : qmat[q][0];
    int coeff = slices[slice].coeffs[plane][0];
    int prev_dc = (coeff - 0x4000) / scale;
    encode_vlc_codeword(pb, FIRST_DC_CB, MAKE_CODE(prev_dc));
    int sign = 0;
    for (int i = 1; i < blocks_per_slice; i++) {
        coeff = slices[slice].coeffs[plane][i];
        int dc = (coeff - 0x4000) / scale;
        int delta = dc - prev_dc;
        int new_sign = GET_SIGN(delta);
        delta = (delta ^ sign) - sign;
        int code = MAKE_CODE(delta);
        encode_vlc_codeword(pb, dc_codebook[codebook], code);
        codebook = min(code, 6);
        sign = new_sign;
        prev_dc = dc;
    }
}

void encode_acs(inout PutBitContext pb, bool is_chroma, int q)
{
    const uint8_t run_to_cb[16] = { U8(0x06), U8(0x06), U8(0x05), U8(0x05), U8(0x04), U8(0x29),
                                    U8(0x29), U8(0x29), U8(0x29), U8(0x28), U8(0x28), U8(0x28),
                                    U8(0x28), U8(0x28), U8(0x28), U8(0x4C) };

    const uint8_t level_to_cb[10] = { U8(0x04), U8(0x0A), U8(0x05), U8(0x06), U8(0x04), U8(0x28),
                                      U8(0x28), U8(0x28), U8(0x28), U8(0x4C) };

    uint slice = gl_GlobalInvocationID.x;
    uint plane = gl_GlobalInvocationID.y;
    uint blocks_per_mb = is_chroma && chroma_factor != CFACTOR_Y444 ? 2 : 4;
    uint blocks_per_slice = slices[slice].mbs_per_slice * blocks_per_mb;
    int prev_run = 4;
    int prev_level = 2;
    int run = 0;

    for (uint i = 1; i < 64; i++) {
        int quant = is_chroma ? qmat_chroma[q][i] : qmat[q][i];
        for (uint j = 0; j < blocks_per_slice; j++) {
            uint idx = i * blocks_per_slice + j;
            int coeff = slices[slice].coeffs[plane][idx];
            int level = coeff / quant;
            if (level != 0) {
                int abs_level = abs(level);
                encode_vlc_codeword(pb, run_to_cb[prev_run], run);
                encode_vlc_codeword(pb, level_to_cb[prev_level], abs_level - 1);
                put_bits(pb, 1, zero_extend(GET_SIGN(level), 1));
                prev_run = min(run, 15);
                prev_level = min(abs_level, 9);
                run = 0;
            } else {
                run++;
            }
        }
    }
}

void encode_slice_plane(inout PutBitContext pb, int q)
{
    uint plane = gl_GlobalInvocationID.y;
    bool is_chroma = plane == 1 || plane == 2;
    encode_dcs(pb, is_chroma, q);
    encode_acs(pb, is_chroma, q);
}

void put_alpha_diff(inout PutBitContext pb, int cur, int prev)
{
    const int dbits = (alpha_bits == 8) ? 4 : 7;
    const int dsize = 1 << dbits - 1;
    int diff = cur - prev;

    diff = zero_extend(diff, alpha_bits);
    if (diff >= (1 << alpha_bits) - dsize)
        diff -= 1 << alpha_bits;
    if (diff < -dsize || diff > dsize || diff == 0) {
        put_bits(pb, 1, 1);
        put_bits(pb, alpha_bits, diff);
    } else {
        put_bits(pb, 1, 0);
        put_bits(pb, dbits - 1, abs(diff) - 1);
        put_bits(pb, 1, int(diff < 0));
    }
}

void put_alpha_run(inout PutBitContext pb, int run)
{
    if (run != 0) {
        put_bits(pb, 1, 0);
        if (run < 0x10)
            put_bits(pb, 4, run);
        else
            put_bits(pb, 15, run);
    } else {
        put_bits(pb, 1, 1);
    }
}

void encode_alpha_plane(inout PutBitContext pb)
{
    uint slice = gl_GlobalInvocationID.x;
    const int mask = (1 << alpha_bits) - 1;
    const int num_coeffs = int(slices[slice].mbs_per_slice) * 256;
    int prev = mask, cur;
    int idx = 0;
    int run = 0;

    cur = slices[slice].coeffs[3][idx++];
    put_alpha_diff(pb, cur, prev);
    prev = cur;
    do {
        cur = slices[slice].coeffs[3][idx++];
        if (cur != prev) {
            put_alpha_run(pb, run);
            put_alpha_diff(pb, cur, prev);
            prev = cur;
            run  = 0;
        } else {
            run++;
        }
    } while (idx < num_coeffs);
    put_alpha_run(pb, run);
}

u8vec2 byteswap16(int value)
{
    return unpack8(uint16_t(value)).yx;
}

void main()
{
    uint slice = gl_GlobalInvocationID.x;
    if (slice >= slices_per_picture)
        return;

    uint plane = gl_GlobalInvocationID.y;
    int q = scores[slice].quant;
    int q_idx = min(q, max_quant + 1);
    ivec4 bits = scores[slice].bits[q_idx];
    int slice_hdr_size = 2 * num_planes;
    int slice_size = slice_hdr_size + ((bits.x + bits.y + bits.z + bits.w) / 8);
    int buf_start = scores[slice].buf_start;
    u8buf buf = OFFBUF(u8buf, bytestream, buf_start);

    /* Write slice header */
    if (plane == 0) {
        buf[0].v = uint8_t(slice_hdr_size * 8);
        buf[1].v = uint8_t(q);
        u8vec2buf slice_hdr = OFFBUF(u8vec2buf, buf, 2);
        for (int i = 0; i < num_planes - 1; i++) {
            slice_hdr[i].v = byteswap16(bits[i] / 8);
        }
        seek_table[slice].v = byteswap16(slice_size);
    }

    int plane_offset = 0;
    for (int i = 0; i < plane; ++i)
        plane_offset += bits[i] / 8;

    /* Encode slice plane */
    PutBitContext pb;
    init_put_bits(pb, OFFBUF(u8buf, buf, slice_hdr_size + plane_offset), 0);
    if (plane == 3)
        encode_alpha_plane(pb);
    else
        encode_slice_plane(pb, q);
    flush_put_bits(pb);
}