// Copyright 2020 Tencent
// SPDX-License-Identifier: BSD-3-Clause

#version 450

layout (constant_id = 0) const int resize_type = 0;
layout (constant_id = 1) const int align_corner = 0;

#define shape_constant_id_offset 2
layout (constant_id = shape_constant_id_offset + 0) const int dims = 0;
layout (constant_id = shape_constant_id_offset + 1) const int w = 0;
layout (constant_id = shape_constant_id_offset + 2) const int h = 0;
layout (constant_id = shape_constant_id_offset + 3) const int c = 0;
layout (constant_id = shape_constant_id_offset + 4) const int cstep = 0;

layout (constant_id = shape_constant_id_offset + 5) const int outdims = 0;
layout (constant_id = shape_constant_id_offset + 6) const int outw = 0;
layout (constant_id = shape_constant_id_offset + 7) const int outh = 0;
layout (constant_id = shape_constant_id_offset + 8) const int outc = 0;
layout (constant_id = shape_constant_id_offset + 9) const int outcstep = 0;

layout (binding = 0) readonly buffer bottom_blob { sfpvec8 bottom_blob_data[]; };
layout (binding = 1) writeonly buffer top_blob { sfpvec8 top_blob_data[]; };

layout (push_constant) uniform parameter
{
    int dims;
    int w;
    int h;
    int c;
    int cstep;

    int outdims;
    int outw;
    int outh;
    int outc;
    int outcstep;

    float scale_x;
    float scale_y;
} p;

void main()
{
    int gx = int(gl_GlobalInvocationID.x);
    int gy = int(gl_GlobalInvocationID.y);
    int gz = int(gl_GlobalInvocationID.z);

    if (gx >= psc(outw) || gy >= psc(outh) || gz >= psc(outc))
        return;

    if (psc(dims) == 1)
    {
        const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;

        buffer_cp8(top_blob_data, gi, bottom_blob_data, gz);

        return;
    }

    if (psc(dims) == 2)
    {
        if (resize_type == 1) // nearest
        {
            int sx = min(int(floor(afp(gx) * afp(p.scale_x))), psc(w) - 1);

            int v_offset = gz * psc(cstep) + gy * psc(w) + sx;

            const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;

            buffer_cp8(top_blob_data, gi, bottom_blob_data, v_offset);
        }
        if (resize_type == 2) // bilinear
        {
            afp fx;
            if (align_corner == 1)
            {
                fx = afp(gx) * afp(p.scale_x);
            }
            else
            {
                fx = (afp(gx) + afp(0.5f)) * afp(p.scale_x) - afp(0.5f);
            }

            int sx = int(floor(fx));

            fx -= afp(sx);

            int sx_max = psc(w) - 2;

            if (sx < 0)
            {
                sx = 0;
                fx = afp(0.f);
            }
            else if (sx > sx_max)
            {
                sx = sx_max;
                fx = afp(1.f);
            }

            int v_offset_0 = gz * psc(cstep) + gy * psc(w) + sx;

            afpvec8 a0 = buffer_ld8(bottom_blob_data, v_offset_0);
            afpvec8 a1 = buffer_ld8(bottom_blob_data, v_offset_0 + 1);

            afpvec8 res;
            res[0] = a0[0] * (afp(1.f) - fx) + a1[0] * fx;
            res[1] = a0[1] * (afp(1.f) - fx) + a1[1] * fx;

            const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;

            buffer_st8(top_blob_data, gi, res);
        }

        return;
    }

    if (resize_type == 1) // nearest
    {
        afpvec2 gxy = afpvec2(gx, gy);
        ivec2 sxy_max = ivec2(psc(w) - 1, psc(h) - 1);
        ivec2 sxy = min(ivec2(floor(gxy * afpvec2(p.scale_x, p.scale_y))), sxy_max);

        int sx = sxy.r;
        int sy = sxy.g;

        int v_offset = gz * psc(cstep) + sy * psc(w) + sx;

        const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;

        buffer_cp8(top_blob_data, gi, bottom_blob_data, v_offset);
    }
    if (resize_type == 2) // bilinear
    {
        afpvec2 gxy = afpvec2(gx, gy);
        afpvec2 fxy;
        if (align_corner == 1)
        {
            fxy = gxy * afpvec2(p.scale_x, p.scale_y);
        }
        else
        {
            fxy = (gxy + afp(0.5f)) * afpvec2(p.scale_x, p.scale_y) - afp(0.5f);
        }

        ivec2 sxy = ivec2(floor(fxy));

        fxy -= afpvec2(sxy);

        ivec2 sxy_max = ivec2(psc(w) - 2, psc(h) - 2);

        bvec2 underflow = lessThan(sxy, ivec2(0));
        bvec2 overflow = greaterThan(sxy, sxy_max);

        sxy = clamp(sxy, ivec2(0), sxy_max);

        fxy = mix(fxy, afpvec2(0.f), underflow);
        fxy = mix(fxy, afpvec2(1.f), overflow);

        int sx = sxy.r;
        int sy = sxy.g;

        int v_offset_0 = gz * psc(cstep) + sy * psc(w) + sx;
        int v_offset_1 = gz * psc(cstep) + (sy + 1) * psc(w) + sx;

        afpvec8 a0 = buffer_ld8(bottom_blob_data, v_offset_0);
        afpvec8 a1 = buffer_ld8(bottom_blob_data, v_offset_0 + 1);
        afpvec8 b0 = buffer_ld8(bottom_blob_data, v_offset_1);
        afpvec8 b1 = buffer_ld8(bottom_blob_data, v_offset_1 + 1);

        afp fx = fxy.r;
        afp fy = fxy.g;

        afpvec8 a;
        afpvec8 b;
        a[0] = a0[0] * (afp(1.f) - fx) + a1[0] * fx;
        a[1] = a0[1] * (afp(1.f) - fx) + a1[1] * fx;
        b[0] = b0[0] * (afp(1.f) - fx) + b1[0] * fx;
        b[1] = b0[1] * (afp(1.f) - fx) + b1[1] * fx;

        afpvec8 res;
        res[0] = a[0] * (afp(1.f) - fy) + b[0] * fy;
        res[1] = a[1] * (afp(1.f) - fy) + b[1] * fy;

        const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;

        buffer_st8(top_blob_data, gi, res);
    }
}
