/*
 * FFv1 codec
 *
 * Copyright (c) 2024 Lynne <dev@lynne.ee>
 *
 * 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
 */

#pragma shader_stage(compute)
#extension GL_GOOGLE_include_directive : require
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_EXT_maximal_reconvergence : require

#define ENCODE
#define SB_QUALI
#include "common.glsl"
#include "ffv1_common.glsl"

layout (set = 1, binding = 1) uniform uimage2D src[];

ivec3 load_components(ivec2 pos)
{
    ivec3 pix = ivec3(imageLoad(src[0], pos));
    if (planar_rgb)
        for (int i = 1; i < 3; i++)
            pix[i] = int(imageLoad(src[i], pos)[0]);

    return ivec3(pix[fmt_lut[0]], pix[fmt_lut[1]], pix[fmt_lut[2]]);
}

#define NUM_CHECKS 15
const ivec2 rct_y_coeff[NUM_CHECKS] = {
    ivec2(0, 0), //      4G

    ivec2(0, 1), //      3G +  B
    ivec2(1, 0), //  R + 3G
    ivec2(1, 1), //  R + 2G + B

    ivec2(0, 2), //      2G + 2B
    ivec2(2, 0), // 2R + 2G
    ivec2(2, 2), // 2R      + 2B

    ivec2(0, 3), //      1G + 3B
    ivec2(3, 0), // 3R + 1G

    ivec2(0, 4), //           4B
    ivec2(4, 0), // 4R

    ivec2(1, 2), //  R +  G + 2B
    ivec2(2, 1), // 2R +  G +  B

    ivec2(3, 1), // 3R      +  B
    ivec2(1, 3), //  R      + 3B
};

shared ivec3 pix_buf[gl_WorkGroupSize.x + 1][gl_WorkGroupSize.y + 1] = { };

ivec3 transform_sample(ivec3 pix, ivec2 rct_coef)
{
    pix.b -= pix.g;
    pix.r -= pix.g;
    pix.g += (pix.b*rct_coef.g + pix.r*rct_coef.r) >> 2;
    pix.b += rct_offset;
    pix.r += rct_offset;
    return pix;
}

uint get_dist(ivec3 cur)
{
    ivec3 LL = pix_buf[gl_LocalInvocationID.x + 0][gl_LocalInvocationID.y + 1];
    ivec3 TL = pix_buf[gl_LocalInvocationID.x + 0][gl_LocalInvocationID.y + 0];
    ivec3 TT = pix_buf[gl_LocalInvocationID.x + 1][gl_LocalInvocationID.y + 0];

    ivec3 pred = ivec3(predict(LL.r, ivec2(TL.r, TT.r)),
                       predict(LL.g, ivec2(TL.g, TT.g)),
                       predict(LL.b, ivec2(TL.b, TT.b)));

    uvec3 c = abs(pred - cur);
    return mid_pred(c.r, c.g, c.b);
}

shared uint score_cols[gl_WorkGroupSize.y] = { };
shared uint score_mode[16] = { };

/* One scoring step: publish this lane's tx_pix to shared, then read the
 * neighbours' to compute the prediction error. `valid` selects whether
 * this lane has a real pixel; invalid lanes write zero into pix_buf so
 * the cache stays well-defined while still participating in the barrier. */
void process(ivec2 pos, bool valid, int i)
{
    ivec3 pix = valid ? load_components(pos) : ivec3(0);
    ivec3 tx_pix = transform_sample(pix, rct_y_coeff[i]);
    pix_buf[gl_LocalInvocationID.x + 1][gl_LocalInvocationID.y + 1] =
        valid ? tx_pix : ivec3(0);
    barrier();

    uint dist = 0u;
    if (valid)
        dist = get_dist(tx_pix);

    uint sum = subgroupAdd(dist);
    if (subgroupElect())
        atomicAdd(score_mode[i], sum);
}

void coeff_search(uint slice_idx)
{
    uvec2 img_size = imageSize(src[0]);
    uint sxs = slice_coord(img_size.x, gl_WorkGroupID.x + 0,
                           gl_NumWorkGroups.x, 0);
    uint sxe = slice_coord(img_size.x, gl_WorkGroupID.x + 1,
                           gl_NumWorkGroups.x, 0);
    uint sys = slice_coord(img_size.y, gl_WorkGroupID.y + 0,
                           gl_NumWorkGroups.y, 0);
    uint sye = slice_coord(img_size.y, gl_WorkGroupID.y + 1,
                           gl_NumWorkGroups.y, 0);

    /* Uniform iteration: every lane in the workgroup runs the same number
     * of tile iterations so that the in-process barrier is always reached
     * by all lanes. Lanes outside the slice extents pass valid=false. */
    uint sw = sxe - sxs;
    uint sh = sye - sys;
    uint n_xi = (sw + gl_WorkGroupSize.x - 1u) / gl_WorkGroupSize.x;
    uint n_yi = (sh + gl_WorkGroupSize.y - 1u) / gl_WorkGroupSize.y;

    for (uint yi = 0u; yi < n_yi; yi++) {
        uint y = sys + yi*gl_WorkGroupSize.y + gl_LocalInvocationID.y;
        for (uint xi = 0u; xi < n_xi; xi++) {
            uint x = sxs + xi*gl_WorkGroupSize.x + gl_LocalInvocationID.x;
            bool valid = (x < sxe) && (y < sye);
            for (int i = 0; i < NUM_CHECKS; i++) {
                process(ivec2(x, y), valid, i);
                barrier();
            }
        }
    }

    /* All lanes must have finished accumulating into score_mode before
     * lane (0,0) inspects it for the argmin. */
    barrier();

    if (gl_LocalInvocationID.x == 0 && gl_LocalInvocationID.y == 0) {
        uint min_score = 0xFFFFFFFF;
        uint min_idx = 3;
        for (int i = 0; i < NUM_CHECKS; i++) {
            if (score_mode[i] < min_score) {
                min_score = score_mode[i];
                min_idx = i;
            }
        }
        slice_ctx[slice_idx].slice_rct_coef = rct_y_coeff[min_idx];
    }
}

void main(void) [[maximally_reconverges]]
{
    if (force_pcm)
        return;

    const uint slice_idx = gl_WorkGroupID.y*gl_NumWorkGroups.x + gl_WorkGroupID.x;
    coeff_search(slice_idx);
}
