// Tencent is pleased to support the open source community by making ncnn available.
//
// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
//
// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// https://opensource.org/licenses/BSD-3-Clause
//
// Unless required by applicable law or agreed to in writing, software distributed
// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, either express or implied. See the License for the
// specific language governing permissions and limitations under the License.

#version 450

#if NCNN_fp16_storage
#extension GL_EXT_shader_16bit_storage: require
struct sfpvec8 { f16vec4 abcd; f16vec4 efgh; };
#endif
#if NCNN_fp16_arithmetic
#extension GL_EXT_shader_explicit_arithmetic_types_float16: require
#endif

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

layout (constant_id = shape_constant_id_offset + 4) const int outcstep = 0;

layout (constant_id = shape_constant_id_offset + 5) const int block_x = 0;
layout (constant_id = shape_constant_id_offset + 6) const int block_y = 0;

#if NCNN_image_shader
layout (binding = 0) uniform unfp sampler3D bottom_blob;
layout (binding = 1, imfmtc4) writeonly uniform unfp image3D bottom_tm_blob;
#else
layout (binding = 0) readonly buffer bottom_blob { sfpvec8 bottom_blob_data[]; };
layout (binding = 1) writeonly buffer bottom_tm_blob { sfpvec8 bottom_tm_blob_data[]; };
#endif

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

    int outcstep;

    int block_x;
    int block_y;
} p;

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

    if (gx >= p.block_x || gy >= p.block_y || gz >= psc(c))
        return;

    // load 4x4
#if NCNN_image_shader
    int sx = gx * 2;
    int sy = gy * 2;

    afpvec8 v00 = image3d_ld8(bottom_blob, ivec3(sx + 0, sy + 0, gz));
    afpvec8 v01 = image3d_ld8(bottom_blob, ivec3(sx + 1, sy + 0, gz));
    afpvec8 v02 = image3d_ld8(bottom_blob, ivec3(sx + 2, sy + 0, gz));
    afpvec8 v03 = image3d_ld8(bottom_blob, ivec3(sx + 3, sy + 0, gz));

    afpvec8 v10 = image3d_ld8(bottom_blob, ivec3(sx + 0, sy + 1, gz));
    afpvec8 v11 = image3d_ld8(bottom_blob, ivec3(sx + 1, sy + 1, gz));
    afpvec8 v12 = image3d_ld8(bottom_blob, ivec3(sx + 2, sy + 1, gz));
    afpvec8 v13 = image3d_ld8(bottom_blob, ivec3(sx + 3, sy + 1, gz));

    afpvec8 v20 = image3d_ld8(bottom_blob, ivec3(sx + 0, sy + 2, gz));
    afpvec8 v21 = image3d_ld8(bottom_blob, ivec3(sx + 1, sy + 2, gz));
    afpvec8 v22 = image3d_ld8(bottom_blob, ivec3(sx + 2, sy + 2, gz));
    afpvec8 v23 = image3d_ld8(bottom_blob, ivec3(sx + 3, sy + 2, gz));

    afpvec8 v30 = image3d_ld8(bottom_blob, ivec3(sx + 0, sy + 3, gz));
    afpvec8 v31 = image3d_ld8(bottom_blob, ivec3(sx + 1, sy + 3, gz));
    afpvec8 v32 = image3d_ld8(bottom_blob, ivec3(sx + 2, sy + 3, gz));
    afpvec8 v33 = image3d_ld8(bottom_blob, ivec3(sx + 3, sy + 3, gz));
#else
    int v_offset_0 = gz * psc(cstep) + gy * 2 * psc(w) + gx * 2;
    ivec4 v_offset = v_offset_0 + ivec4(0, 1, 2, 3) * psc(w);

    afpvec8 v00 = buffer_ld8(bottom_blob_data, v_offset.r + 0);
    afpvec8 v01 = buffer_ld8(bottom_blob_data, v_offset.r + 1);
    afpvec8 v02 = buffer_ld8(bottom_blob_data, v_offset.r + 2);
    afpvec8 v03 = buffer_ld8(bottom_blob_data, v_offset.r + 3);

    afpvec8 v10 = buffer_ld8(bottom_blob_data, v_offset.g + 0);
    afpvec8 v11 = buffer_ld8(bottom_blob_data, v_offset.g + 1);
    afpvec8 v12 = buffer_ld8(bottom_blob_data, v_offset.g + 2);
    afpvec8 v13 = buffer_ld8(bottom_blob_data, v_offset.g + 3);

    afpvec8 v20 = buffer_ld8(bottom_blob_data, v_offset.b + 0);
    afpvec8 v21 = buffer_ld8(bottom_blob_data, v_offset.b + 1);
    afpvec8 v22 = buffer_ld8(bottom_blob_data, v_offset.b + 2);
    afpvec8 v23 = buffer_ld8(bottom_blob_data, v_offset.b + 3);

    afpvec8 v30 = buffer_ld8(bottom_blob_data, v_offset.a + 0);
    afpvec8 v31 = buffer_ld8(bottom_blob_data, v_offset.a + 1);
    afpvec8 v32 = buffer_ld8(bottom_blob_data, v_offset.a + 2);
    afpvec8 v33 = buffer_ld8(bottom_blob_data, v_offset.a + 3);
#endif

    // const float itm[4][4] = {
    //     {1.0f,  0.0f, -1.0f,  0.0f},
    //     {0.0f,  1.0f,  1.0f,  0.0f},
    //     {0.0f, -1.0f,  1.0f,  0.0f},
    //     {0.0f, -1.0f,  0.0f,  1.0f}
    // };

    // implicit transpose
    afpvec8 m00 = v00 - v02;
    afpvec8 m01 = v10 - v12;
    afpvec8 m02 = v20 - v22;
    afpvec8 m03 = v30 - v32;

    afpvec8 m10 = v02 + v01;
    afpvec8 m11 = v12 + v11;
    afpvec8 m12 = v22 + v21;
    afpvec8 m13 = v32 + v31;

    afpvec8 m20 = v02 - v01;
    afpvec8 m21 = v12 - v11;
    afpvec8 m22 = v22 - v21;
    afpvec8 m23 = v32 - v31;

    afpvec8 m30 = v03 - v01;
    afpvec8 m31 = v13 - v11;
    afpvec8 m32 = v23 - v21;
    afpvec8 m33 = v33 - v31;

    v00 = m00 - m02;
    v10 = m10 - m12;
    v20 = m20 - m22;
    v30 = m30 - m32;

    v01 = m02 + m01;
    v11 = m12 + m11;
    v21 = m22 + m21;
    v31 = m32 + m31;

    v02 = m02 - m01;
    v12 = m12 - m11;
    v22 = m22 - m21;
    v32 = m32 - m31;

    v03 = m03 - m01;
    v13 = m13 - m11;
    v23 = m23 - m21;
    v33 = m33 - m31;

    // store 16
#if NCNN_image_shader
    int y = gy * p.block_x + gx;

    image3d_st8(bottom_tm_blob, ivec3(0, y, gz), v00);
    image3d_st8(bottom_tm_blob, ivec3(1, y, gz), v01);
    image3d_st8(bottom_tm_blob, ivec3(2, y, gz), v02);
    image3d_st8(bottom_tm_blob, ivec3(3, y, gz), v03);
    image3d_st8(bottom_tm_blob, ivec3(4, y, gz), v10);
    image3d_st8(bottom_tm_blob, ivec3(5, y, gz), v11);
    image3d_st8(bottom_tm_blob, ivec3(6, y, gz), v12);
    image3d_st8(bottom_tm_blob, ivec3(7, y, gz), v13);
    image3d_st8(bottom_tm_blob, ivec3(8, y, gz), v20);
    image3d_st8(bottom_tm_blob, ivec3(9, y, gz), v21);
    image3d_st8(bottom_tm_blob, ivec3(10, y, gz), v22);
    image3d_st8(bottom_tm_blob, ivec3(11, y, gz), v23);
    image3d_st8(bottom_tm_blob, ivec3(12, y, gz), v30);
    image3d_st8(bottom_tm_blob, ivec3(13, y, gz), v31);
    image3d_st8(bottom_tm_blob, ivec3(14, y, gz), v32);
    image3d_st8(bottom_tm_blob, ivec3(15, y, gz), v33);
#else
    int v_tm_offset = gz * psc(outcstep) + (gy * p.block_x + gx) * 16;

    buffer_st8(bottom_tm_blob_data, v_tm_offset + 0, v00);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 1, v01);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 2, v02);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 3, v03);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 4, v10);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 5, v11);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 6, v12);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 7, v13);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 8, v20);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 9, v21);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 10, v22);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 11, v23);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 12, v30);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 13, v31);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 14, v32);
    buffer_st8(bottom_tm_blob_data, v_tm_offset + 15, v33);
#endif
}
