// -----------------------------------------------------------------------------------------
// NVEnc by rigaya
// -----------------------------------------------------------------------------------------
//
// The MIT License
//
// Copyright (c) 2014-2016 rigaya
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
//
// ------------------------------------------------------------------------------------------

#include <map>
#include <array>
#include <cmath>
#include "convert_csp.h"
#include "NVEncFilterAfs.h"
#include "NVEncParam.h"
#pragma warning (push)
#pragma warning (disable: 4819)
#include <cuda_runtime.h>
#include <cuda_runtime_api.h>
#include <device_launch_parameters.h>
#include <vector_types.h>
#pragma warning (pop)
#include "rgy_cuda_util_kernel.h"

#define MERGE_BLOCK_INT_X  (32) //work groupサイズ(x) = スレッド数/work group
#define MERGE_BLOCK_Y       (8) //work groupサイズ(y) = スレッド数/work group
#define MERGE_BLOCK_LOOP_Y  (1) //work groupのy方向反復数

#define u8x4(x)  (uint32_t)(((uint32_t)(x)) | (((uint32_t)(x)) <<  8) | (((uint32_t)(x)) << 16) | (((uint32_t)(x)) << 24))

template<typename Type>
__global__ void kernel_afs_merge_scan(
    Type *__restrict__ ptr_dst,
    int *__restrict__ ptr_count,
    const Type *__restrict__ src_p0,
    const Type *__restrict__ src_p1,
    const int si_w_type, const int pitch_type, const int height,
    const int tb_order,
    const uint32_t scan_left, const uint32_t scan_top, const uint32_t scan_width, const uint32_t scan_height) {
    //int lx = threadIdx.x; //スレッド数=MERGE_BLOCK_INT_X
    int ly = threadIdx.y; //スレッド数=MERGE_BLOCK_Y
    int imgx = blockIdx.x * MERGE_BLOCK_INT_X /*blockDim.x*/ + threadIdx.x;

    int stripe_count = 0;
    const int field_select = ((ly + tb_order) & 1);

    if (imgx < si_w_type) {
        const int mask_lsft = (field_select) ? 1 : 2;
        int imgy = blockIdx.y * MERGE_BLOCK_LOOP_Y * MERGE_BLOCK_Y + ly;
        src_p0  += imgx + pitch_type * imgy;
        src_p1  += imgx + pitch_type * imgy;
        ptr_dst += imgx + pitch_type * imgy;
        #pragma unroll
        for (int iloop = 0; iloop < MERGE_BLOCK_LOOP_Y; iloop++,
            imgy    += MERGE_BLOCK_Y,
            ptr_dst += MERGE_BLOCK_Y * pitch_type,
            src_p0  += MERGE_BLOCK_Y * pitch_type,
            src_p1  += MERGE_BLOCK_Y * pitch_type
            ) {
            if (imgy < height) {
                const int offsetm = (imgy == 0       ) ? 0 : -pitch_type;
                const int offsetp = (imgy >= height-1) ? 0 :  pitch_type;
                Type p0m = src_p0[offsetm];
                Type p0c = src_p0[0];
                Type p0p = src_p0[offsetp];

                Type p1m = src_p1[offsetm];
                Type p1c = src_p1[0];
                Type p1p = src_p1[offsetp];

                Type m4 = (p0m | p0p | u8x4(0xf3)) & p0c;
                Type m5 = (p1m | p1p | u8x4(0xf3)) & p1c;

                Type m6 = (m4 & m5 & u8x4(0x44)) | (~(p0c) & u8x4(0x33));

                ptr_dst[0] = m6;

                if ((((uint32_t)imgx - scan_left) < scan_width) && (((uint32_t)imgy - scan_top) < scan_height)) {
                    //(m6 & MASK) == 0 を各バイトについて加算する
                    //MASK = 0x50 or 0x60 (ly, tb_order依存)
                    //まず、0x10 or 0x20 のビットを 0x40に集約する
                    Type check = m6 | ((m6) << mask_lsft);
                    //0x40のビットが1なら0、0なら1としたいので、xorでチェック
                    //他のビットはいらないので、自分とxorして消す
                    stripe_count += __popc(check ^ (check | u8x4(0x40)));
                }
            }
        }
    }
    //motion countの総和演算
    // 32               16              0
    //  |  count_latter ||  count_first |
    static_assert(MERGE_BLOCK_INT_X * sizeof(int) * MERGE_BLOCK_Y * MERGE_BLOCK_LOOP_Y < (1<<(sizeof(short)*8-1)), "reduce block size for proper reduction in 16bit.");
    int stripe_count_01 = (int)(field_select ? (uint32_t)stripe_count << 16 : (uint32_t)stripe_count);

    __shared__ int shared[MERGE_BLOCK_INT_X * MERGE_BLOCK_Y / WARP_SIZE]; //int単位でアクセスする
    stripe_count_01 = block_sum<decltype(stripe_count_01), MERGE_BLOCK_INT_X, MERGE_BLOCK_Y>(stripe_count_01, (int *)shared);

    const int lid = threadIdx.y * MERGE_BLOCK_INT_X + threadIdx.x;
    if (lid == 0) {
        const int gid = blockIdx.y * gridDim.x + blockIdx.x;
        ptr_count[gid] = stripe_count_01;
    }
}

template<typename Type>
RGY_ERR run_merge_scan(uint8_t *dst,
    uint8_t *sp0, uint8_t *sp1,
    const int srcWidth, const int srcPitch, const int srcHeight,
    CUMemBufPair *count_stripe, const VppAfs *pAfsPrm, cudaStream_t stream) {
    auto sts = RGY_ERR_NONE;

    dim3 blockSize(MERGE_BLOCK_INT_X, MERGE_BLOCK_Y);
    dim3 gridSize(divCeil(srcWidth, blockSize.x * sizeof(Type)), divCeil(srcHeight, blockSize.y * MERGE_BLOCK_LOOP_Y));

    const auto grid_count = gridSize.x * gridSize.y;
    if (count_stripe->nSize < grid_count) {
        count_stripe->clear();
        sts = count_stripe->alloc(grid_count * sizeof(int));
        if (sts != RGY_ERR_NONE) {
            return sts;
        }
    }
    const uint32_t scan_left   = pAfsPrm->clip.left / sizeof(Type);
    const uint32_t scan_width  = (srcWidth - pAfsPrm->clip.left - pAfsPrm->clip.right) / sizeof(Type);
    const uint32_t scan_top    = pAfsPrm->clip.top;
    const uint32_t scan_height = srcHeight - pAfsPrm->clip.top - pAfsPrm->clip.bottom;

    kernel_afs_merge_scan<<<gridSize, blockSize, 0, stream>>>(
        (Type *)dst, (int *)count_stripe->ptrDevice, (Type *)sp0, (Type *)sp1,
        divCeil(srcWidth, sizeof(Type)), divCeil(srcPitch, sizeof(Type)), srcHeight,
        pAfsPrm->tb_order ? 1 : 0,
        scan_left, scan_top, scan_width, scan_height);
    return err_to_rgy(cudaGetLastError());
}

RGY_ERR NVEncFilterAfs::merge_scan(AFS_STRIPE_DATA *sp, AFS_SCAN_DATA *sp0, AFS_SCAN_DATA *sp1, CUMemBufPair *count_motion, const NVEncFilterParamAfs *pAfsParam, cudaStream_t stream) {
    auto sts = run_merge_scan<uint32_t>(
        sp->map.frame.ptr[0], sp0->map.frame.ptr[0], sp1->map.frame.ptr[0],
        sp1->map.frame.width, sp1->map.frame.pitch[0], sp1->map.frame.height,
        count_motion, &pAfsParam->afs, stream);
    if (sts != RGY_ERR_NONE) {
        return sts;
    }
    return RGY_ERR_NONE;
}
