/*****************************************************************************

* Copyright (C) 2013-2020 MulticoreWare, Inc

*

* Author: Gopu Govindaswamy <gopu@multicorewareinc.com>

*         Min Chen <chenm003@163.com>

*

* This program is free software; you can redistribute it and/or modify

* it under the terms of the GNU General Public License as published by

* the Free Software Foundation; either version 2 of the License, or

* (at your option) any later version.

*

* This program 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 General Public License for more details.

*

* You should have received a copy of the GNU General Public License

* along with this program; if not, write to the Free Software

* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02111, USA.

*

* This program is also available under a commercial proprietary license.

* For more information, contact us at license @ x265.com.

*****************************************************************************/

#include "common.h"

#include "deblock.h"

#include "framedata.h"

#include "picyuv.h"

#include "slice.h"

#include "mv.h"

using namespace X265_NS;

#define DEBLOCK_SMALLEST_BLOCK  8

#define DEFAULT_INTRA_TC_OFFSET 2

void Deblock::deblockCTU(const CUData* ctu, const CUGeom& cuGeom, int32_t dir)

{

    uint8_t blockStrength[MAX_NUM_PARTITIONS];

    memset(blockStrength, 0, sizeof(uint8_t) * cuGeom.numPartitions);

    deblockCU(ctu, cuGeom, dir, blockStrength);

}

static inline uint8_t bsCuEdge(const CUData* cu, uint32_t absPartIdx, int32_t dir)

{

    if (dir == Deblock::EDGE_VER)

    {

        if (cu->m_cuPelX + g_zscanToPelX[absPartIdx] > 0)

        {

            uint32_t    tempPartIdx;

            const CUData* tempCU = cu->getPULeft(tempPartIdx, absPartIdx);

            return tempCU ? 2 : 0;

        }

    }

    else

    {

        if (cu->m_cuPelY + g_zscanToPelY[absPartIdx] > 0)

        {

            uint32_t    tempPartIdx;

            const CUData* tempCU = cu->getPUAbove(tempPartIdx, absPartIdx);

            return tempCU ? 2 : 0;

        }

    }

    return 0;

}

/* Deblocking filter process in CU-based (the same function as conventional's)

 * param Edge the direction of the edge in block boundary (horizonta/vertical), which is added newly */

void Deblock::deblockCU(const CUData* cu, const CUGeom& cuGeom, const int32_t dir, uint8_t blockStrength[])

{

    uint32_t absPartIdx = cuGeom.absPartIdx;

    uint32_t depth = cuGeom.depth;

    if (cu->m_predMode[absPartIdx] == MODE_NONE)

        return;

    if (cu->m_cuDepth[absPartIdx] > depth)

    {

        for (uint32_t subPartIdx = 0; subPartIdx < 4; subPartIdx++)

        {

            const CUGeom& childGeom = *(&cuGeom + cuGeom.childOffset + subPartIdx);

            if (childGeom.flags & CUGeom::PRESENT)

                deblockCU(cu, childGeom, dir, blockStrength);

        }

        return;

    }

    uint32_t numUnits = 1 << (cuGeom.log2CUSize - LOG2_UNIT_SIZE);

    setEdgefilterPU(cu, absPartIdx, dir, blockStrength, numUnits);

    setEdgefilterTU(cu, absPartIdx, 0, dir, blockStrength);

    setEdgefilterMultiple(absPartIdx, dir, 0, bsCuEdge(cu, absPartIdx, dir), blockStrength, numUnits);

    uint32_t numParts = cuGeom.numPartitions;

    for (uint32_t partIdx = absPartIdx; partIdx < absPartIdx + numParts; partIdx++)

    {

        uint32_t bsCheck = !(partIdx & (1 << dir));

        if (bsCheck && blockStrength[partIdx])

            blockStrength[partIdx] = getBoundaryStrength(cu, dir, partIdx, blockStrength);

    }

    const uint32_t partIdxIncr = DEBLOCK_SMALLEST_BLOCK >> LOG2_UNIT_SIZE;

    uint32_t shiftFactor = (dir == EDGE_VER) ? cu->m_hChromaShift : cu->m_vChromaShift;

    uint32_t chromaMask = ((DEBLOCK_SMALLEST_BLOCK << shiftFactor) >> LOG2_UNIT_SIZE) - 1;

    uint32_t e0 = (dir == EDGE_VER ? g_zscanToPelX[absPartIdx] : g_zscanToPelY[absPartIdx]) >> LOG2_UNIT_SIZE;

    for (uint32_t e = 0; e < numUnits; e += partIdxIncr)

    {

        edgeFilterLuma(cu, absPartIdx, depth, dir, e, blockStrength);

        if (!((e0 + e) & chromaMask) && cu->m_chromaFormat != X265_CSP_I400)

            edgeFilterChroma(cu, absPartIdx, depth, dir, e, blockStrength);

    }

}

static inline uint32_t calcBsIdx(uint32_t absPartIdx, int32_t dir, int32_t edgeIdx, int32_t baseUnitIdx)

{

    if (dir)

        return g_rasterToZscan[g_zscanToRaster[absPartIdx] + (edgeIdx << LOG2_RASTER_SIZE) + baseUnitIdx];

    else

        return g_rasterToZscan[g_zscanToRaster[absPartIdx] + (baseUnitIdx << LOG2_RASTER_SIZE) + edgeIdx];

}

void Deblock::setEdgefilterMultiple(uint32_t scanIdx, int32_t dir, int32_t edgeIdx, uint8_t value, uint8_t blockStrength[], uint32_t numUnits)

{

    X265_CHECK(numUnits > 0, "numUnits edge filter check\n");

    for (uint32_t i = 0; i < numUnits; i++)

    {

        const uint32_t bsidx = calcBsIdx(scanIdx, dir, edgeIdx, i);

        blockStrength[bsidx] = value;

    }

}

void Deblock::setEdgefilterTU(const CUData* cu, uint32_t absPartIdx, uint32_t tuDepth, int32_t dir, uint8_t blockStrength[])

{

    uint32_t log2TrSize = cu->m_log2CUSize[absPartIdx] - tuDepth;

    if (cu->m_tuDepth[absPartIdx] > tuDepth)

    {

        uint32_t qNumParts = 1 << (log2TrSize - LOG2_UNIT_SIZE - 1) * 2;

        for (uint32_t qIdx = 0; qIdx < 4; ++qIdx, absPartIdx += qNumParts)

            setEdgefilterTU(cu, absPartIdx, tuDepth + 1, dir, blockStrength);

        return;

    }

    uint32_t numUnits = 1 << (log2TrSize - LOG2_UNIT_SIZE);

    setEdgefilterMultiple(absPartIdx, dir, 0, 2, blockStrength, numUnits);

}

void Deblock::setEdgefilterPU(const CUData* cu, uint32_t absPartIdx, int32_t dir, uint8_t blockStrength[], uint32_t numUnits)

{

    const uint32_t hNumUnits = numUnits >> 1;

    const uint32_t qNumUnits = numUnits >> 2;

    switch (cu->m_partSize[absPartIdx])

    {

    case SIZE_2NxN:

        if (EDGE_HOR == dir)

            setEdgefilterMultiple(absPartIdx, dir, hNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_Nx2N:

        if (EDGE_VER == dir)

            setEdgefilterMultiple(absPartIdx, dir, hNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_NxN:

        setEdgefilterMultiple(absPartIdx, dir, hNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_2NxnU:

        if (EDGE_HOR == dir)

            setEdgefilterMultiple(absPartIdx, dir, qNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_nLx2N:

        if (EDGE_VER == dir)

            setEdgefilterMultiple(absPartIdx, dir, qNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_2NxnD:

        if (EDGE_HOR == dir)

            setEdgefilterMultiple(absPartIdx, dir, numUnits - qNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_nRx2N:

        if (EDGE_VER == dir)

            setEdgefilterMultiple(absPartIdx, dir, numUnits - qNumUnits, 1, blockStrength, numUnits);

        break;

    case SIZE_2Nx2N:

    default:

        break;

    }

}

uint8_t Deblock::getBoundaryStrength(const CUData* cuQ, int32_t dir, uint32_t partQ, const uint8_t blockStrength[])

{

    // Calculate block index

    uint32_t partP;

    const CUData* cuP = (dir == EDGE_VER ? cuQ->getPULeft(partP, partQ) : cuQ->getPUAbove(partP, partQ));

    // Set BS for Intra MB : BS = 2

    if (cuP->isIntra(partP) || cuQ->isIntra(partQ))

        return 2;

    // Set BS for not Intra MB : BS = 1 or 0

    if (blockStrength[partQ] > 1 &&

        (cuQ->getCbf(partQ, TEXT_LUMA, cuQ->m_tuDepth[partQ]) ||

         cuP->getCbf(partP, TEXT_LUMA, cuP->m_tuDepth[partP])))

        return 1;

    static const MV zeroMv(0, 0);

    const Slice* const sliceQ = cuQ->m_slice;

    const Slice* const sliceP = cuP->m_slice;

    const Frame* refP0 = (cuP->m_refIdx[0][partP] >= 0) ? sliceP->m_refFrameList[0][cuP->m_refIdx[0][partP]] : NULL;

    const Frame* refQ0 = (cuQ->m_refIdx[0][partQ] >= 0) ? sliceQ->m_refFrameList[0][cuQ->m_refIdx[0][partQ]] : NULL;

    const MV& mvP0 = refP0 ? cuP->m_mv[0][partP] : zeroMv;

    const MV& mvQ0 = refQ0 ? cuQ->m_mv[0][partQ] : zeroMv;

    if (sliceQ->isInterP() && sliceP->isInterP())

    {

        return ((refP0 != refQ0) ||

                (abs(mvQ0.x - mvP0.x) >= 4) || (abs(mvQ0.y - mvP0.y) >= 4)) ? 1 : 0;

    }

    // (sliceQ->isInterB() || sliceP->isInterB())

    const Frame* refP1 = (cuP->m_refIdx[1][partP] >= 0) ? sliceP->m_refFrameList[1][cuP->m_refIdx[1][partP]] : NULL;

    const Frame* refQ1 = (cuQ->m_refIdx[1][partQ] >= 0) ? sliceQ->m_refFrameList[1][cuQ->m_refIdx[1][partQ]] : NULL;

    const MV& mvP1 = refP1 ? cuP->m_mv[1][partP] : zeroMv;

    const MV& mvQ1 = refQ1 ? cuQ->m_mv[1][partQ] : zeroMv;

    if (((refP0 == refQ0) && (refP1 == refQ1)) || ((refP0 == refQ1) && (refP1 == refQ0)))

    {

        if (refP0 != refP1) // Different L0 & L1

        {

            if (refP0 == refQ0)

                return ((abs(mvQ0.x - mvP0.x) >= 4) || (abs(mvQ0.y - mvP0.y) >= 4) ||

                        (abs(mvQ1.x - mvP1.x) >= 4) || (abs(mvQ1.y - mvP1.y) >= 4)) ? 1 : 0;

            else

                return ((abs(mvQ1.x - mvP0.x) >= 4) || (abs(mvQ1.y - mvP0.y) >= 4) ||

                        (abs(mvQ0.x - mvP1.x) >= 4) || (abs(mvQ0.y - mvP1.y) >= 4)) ? 1 : 0;

        }

        else // Same L0 & L1

        {

            return (((abs(mvQ0.x - mvP0.x) >= 4) || (abs(mvQ0.y - mvP0.y) >= 4) ||

                     (abs(mvQ1.x - mvP1.x) >= 4) || (abs(mvQ1.y - mvP1.y) >= 4)) &&

                    ((abs(mvQ1.x - mvP0.x) >= 4) || (abs(mvQ1.y - mvP0.y) >= 4) ||

                     (abs(mvQ0.x - mvP1.x) >= 4) || (abs(mvQ0.y - mvP1.y) >= 4))) ? 1 : 0;

        }

    }

    // for all different Ref_Idx

    return 1;

}

static inline int32_t calcDP(pixel* src, intptr_t offset)

{

    return abs(static_cast<int32_t>(src[-offset * 3]) - 2 * src[-offset * 2] + src[-offset]);

}

static inline int32_t calcDQ(pixel* src, intptr_t offset)

{

    return abs(static_cast<int32_t>(src[0]) - 2 * src[offset] + src[offset * 2]);

}

static inline bool useStrongFiltering(intptr_t offset, int32_t beta, int32_t tc, pixel* src)

{

    int16_t m4     = (int16_t)src[0];

    int16_t m3     = (int16_t)src[-offset];

    int16_t m7     = (int16_t)src[offset * 3];

    int16_t m0     = (int16_t)src[-offset * 4];

    int32_t strong = abs(m0 - m3) + abs(m7 - m4);

    return (strong < (beta >> 3)) && (abs(m3 - m4) < ((tc * 5 + 1) >> 1));

}

/* Deblocking for the luminance component with strong or weak filter

 * \param src     pointer to picture data

 * \param offset  offset value for picture data

 * \param tc      tc value

 * \param maskP   indicator to enable filtering on partP

 * \param maskQ   indicator to enable filtering on partQ

 * \param maskP1  decision weak filter/no filter for partP

 * \param maskQ1  decision weak filter/no filter for partQ */

static inline void pelFilterLuma(pixel* src, intptr_t srcStep, intptr_t offset, int32_t tc, int32_t maskP, int32_t maskQ,

                                 int32_t maskP1, int32_t maskQ1)

{

    int32_t thrCut = tc * 10;

    int32_t tc2 = tc >> 1;

    maskP1 &= maskP;

    maskQ1 &= maskQ;

    for (int32_t i = 0; i < UNIT_SIZE; i++, src += srcStep)

    {

        int16_t m4  = (int16_t)src[0];

        int16_t m3  = (int16_t)src[-offset];

        int16_t m5  = (int16_t)src[offset];

        int16_t m2  = (int16_t)src[-offset * 2];

        int32_t delta = (9 * (m4 - m3) - 3 * (m5 - m2) + 8) >> 4;

        if (abs(delta) < thrCut)

        {

            delta = x265_clip3(-tc, tc, delta);

            src[-offset] = x265_clip(m3 + (delta & maskP));

            src[0] = x265_clip(m4 - (delta & maskQ));

            if (maskP1)

            {

                int16_t m1  = (int16_t)src[-offset * 3];

                int32_t delta1 = x265_clip3(-tc2, tc2, ((((m1 + m3 + 1) >> 1) - m2 + delta) >> 1));

                src[-offset * 2] = x265_clip(m2 + delta1);

            }

            if (maskQ1)

            {

                int16_t m6  = (int16_t)src[offset * 2];

                int32_t delta2 = x265_clip3(-tc2, tc2, ((((m6 + m4 + 1) >> 1) - m5 - delta) >> 1));

                src[offset] = x265_clip(m5 + delta2);

            }

        }

    }

}

void Deblock::edgeFilterLuma(const CUData* cuQ, uint32_t absPartIdx, uint32_t depth, int32_t dir, int32_t edge, const uint8_t blockStrength[])

{

    PicYuv* reconPic = cuQ->m_encData->m_reconPic[0];

    pixel* src = reconPic->getLumaAddr(cuQ->m_cuAddr, absPartIdx);

    intptr_t stride = reconPic->m_stride;

    const PPS* pps = cuQ->m_slice->m_pps;

    intptr_t offset, srcStep;

    int32_t maskP = -1;

    int32_t maskQ = -1;

    int32_t betaOffset = pps->deblockingFilterBetaOffsetDiv2 << 1;

    int32_t tcOffset = pps->deblockingFilterTcOffsetDiv2 << 1;

    bool bCheckNoFilter = pps->bTransquantBypassEnabled;

    if (dir == EDGE_VER)

    {

        offset = 1;

        srcStep = stride;

        src += (edge << LOG2_UNIT_SIZE);

    }

    else // (dir == EDGE_HOR)

    {

        offset = stride;

        srcStep = 1;

        src += (edge << LOG2_UNIT_SIZE) * stride;

    }

    uint32_t numUnits = cuQ->m_slice->m_sps->numPartInCUSize >> depth;

    for (uint32_t idx = 0; idx < numUnits; idx++)

    {

        uint32_t partQ = calcBsIdx(absPartIdx, dir, edge, idx);

        uint32_t bs = blockStrength[partQ];

        if (!bs)

            continue;

        // Derive neighboring PU index

        uint32_t partP;

        const CUData* cuP = (dir == EDGE_VER ? cuQ->getPULeft(partP, partQ) : cuQ->getPUAbove(partP, partQ));

        if (bCheckNoFilter)

        {

            // check if each of PUs is lossless coded

            maskP = cuP->m_tqBypass[partP] - 1;

            maskQ = cuQ->m_tqBypass[partQ] - 1;

            if (!(maskP | maskQ))

                continue;

        }

        int32_t qpQ = cuQ->m_qp[partQ];

        int32_t qpP = cuP->m_qp[partP];

        int32_t qp  = (qpP + qpQ + 1) >> 1;

        int32_t indexB = x265_clip3(0, QP_MAX_SPEC, qp + betaOffset);

        const int32_t bitdepthShift = X265_DEPTH - 8;

        int32_t beta = s_betaTable[indexB] << bitdepthShift;

        intptr_t unitOffset = idx * srcStep << LOG2_UNIT_SIZE;

        int32_t dp0 = calcDP(src + unitOffset              , offset);

        int32_t dq0 = calcDQ(src + unitOffset              , offset);

        int32_t dp3 = calcDP(src + unitOffset + srcStep * 3, offset);

        int32_t dq3 = calcDQ(src + unitOffset + srcStep * 3, offset);

        int32_t d0 = dp0 + dq0;

        int32_t d3 = dp3 + dq3;

        int32_t d =  d0 + d3;

        if (d >= beta)

            continue;

        int32_t indexTC = x265_clip3(0, QP_MAX_SPEC + DEFAULT_INTRA_TC_OFFSET, int32_t(qp + DEFAULT_INTRA_TC_OFFSET * (bs - 1) + tcOffset));

        int32_t tc = s_tcTable[indexTC] << bitdepthShift;

        bool sw = (2 * d0 < (beta >> 2) &&

                   2 * d3 < (beta >> 2) &&

                   useStrongFiltering(offset, beta, tc, src + unitOffset              ) &&

                   useStrongFiltering(offset, beta, tc, src + unitOffset + srcStep * 3));

        if (sw)

        {

            int32_t tc2 = 2 * tc;

            int32_t tcP = (tc2 & maskP);

            int32_t tcQ = (tc2 & maskQ);

            primitives.pelFilterLumaStrong[dir](src + unitOffset, srcStep, offset, tcP, tcQ);

        }

        else

        {

            int32_t sideThreshold = (beta + (beta >> 1)) >> 3;

            int32_t dp = dp0 + dp3;

            int32_t dq = dq0 + dq3;

            int32_t maskP1 = (dp < sideThreshold ? -1 : 0);

            int32_t maskQ1 = (dq < sideThreshold ? -1 : 0);

            pelFilterLuma(src + unitOffset, srcStep, offset, tc, maskP, maskQ, maskP1, maskQ1);

        }

    }

}

void Deblock::edgeFilterChroma(const CUData* cuQ, uint32_t absPartIdx, uint32_t depth, int32_t dir, int32_t edge, const uint8_t blockStrength[])

{

    int32_t chFmt = cuQ->m_chromaFormat, chromaShift;

    intptr_t offset, srcStep;

    const PPS* pps = cuQ->m_slice->m_pps;

    int32_t maskP = -1;

    int32_t maskQ = -1;

    int32_t tcOffset = pps->deblockingFilterTcOffsetDiv2 << 1;

    X265_CHECK(((dir == EDGE_VER)

                ? ((g_zscanToPelX[absPartIdx] + edge * UNIT_SIZE) >> cuQ->m_hChromaShift)

                : ((g_zscanToPelY[absPartIdx] + edge * UNIT_SIZE) >> cuQ->m_vChromaShift)) % DEBLOCK_SMALLEST_BLOCK == 0,

               "invalid edge\n");

    PicYuv* reconPic = cuQ->m_encData->m_reconPic[0];

    intptr_t stride = reconPic->m_strideC;

    intptr_t srcOffset = reconPic->getChromaAddrOffset(cuQ->m_cuAddr, absPartIdx);

    bool bCheckNoFilter = pps->bTransquantBypassEnabled;

    if (dir == EDGE_VER)

    {

        chromaShift = cuQ->m_vChromaShift;

        srcOffset += (edge << (LOG2_UNIT_SIZE - cuQ->m_hChromaShift));

        offset     = 1;

        srcStep    = stride;

    }

    else // (dir == EDGE_HOR)

    {

        chromaShift = cuQ->m_hChromaShift;

        srcOffset += edge * stride << (LOG2_UNIT_SIZE - cuQ->m_vChromaShift);

        offset     = stride;

        srcStep    = 1;

    }

    pixel* srcChroma[2];

    srcChroma[0] = reconPic->m_picOrg[1] + srcOffset;

    srcChroma[1] = reconPic->m_picOrg[2] + srcOffset;

    uint32_t numUnits = cuQ->m_slice->m_sps->numPartInCUSize >> (depth + chromaShift);

    for (uint32_t idx = 0; idx < numUnits; idx++)

    {

        uint32_t partQ = calcBsIdx(absPartIdx, dir, edge, idx << chromaShift);

        uint32_t bs = blockStrength[partQ];

        if (bs <= 1)

            continue;

        // Derive neighboring PU index

        uint32_t partP;

        const CUData* cuP = (dir == EDGE_VER ? cuQ->getPULeft(partP, partQ) : cuQ->getPUAbove(partP, partQ));

        if (bCheckNoFilter)

        {

            // check if each of PUs is lossless coded

            maskP = (cuP->m_tqBypass[partP] ? 0 : -1);

            maskQ = (cuQ->m_tqBypass[partQ] ? 0 : -1);

            if (!(maskP | maskQ))

                continue;

        }

        int32_t qpQ = cuQ->m_qp[partQ];

        int32_t qpP = cuP->m_qp[partP];

        int32_t qpA = (qpP + qpQ + 1) >> 1;

        intptr_t unitOffset = idx * srcStep << LOG2_UNIT_SIZE;

        for (uint32_t chromaIdx = 0; chromaIdx < 2; chromaIdx++)

        {

            int32_t qp = qpA + pps->chromaQpOffset[chromaIdx];

            if (qp >= 30)

                qp = chFmt == X265_CSP_I420 ? g_chromaScale[qp] : X265_MIN(qp, QP_MAX_SPEC);

            int32_t indexTC = x265_clip3(0, QP_MAX_SPEC + DEFAULT_INTRA_TC_OFFSET, int32_t(qp + DEFAULT_INTRA_TC_OFFSET + tcOffset));

            const int32_t bitdepthShift = X265_DEPTH - 8;

            int32_t tc = s_tcTable[indexTC] << bitdepthShift;

            pixel* srcC = srcChroma[chromaIdx];

            primitives.pelFilterChroma[dir](srcC + unitOffset, srcStep, offset, tc, maskP, maskQ);

        }

    }

}

const uint8_t Deblock::s_tcTable[54] =

{

    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2,

    2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 5, 5, 6, 6, 7, 8, 9, 10, 11, 13, 14, 16, 18, 20, 22, 24

};

const uint8_t Deblock::s_betaTable[52] =

{

    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17,

    18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64

};