/* -----------------------------------------------------------------------------

The copyright in this software is being made available under the Clear BSD

License, included below. No patent rights, trademark rights and/or 

other Intellectual Property Rights other than the copyrights concerning 

the Software are granted under this license.

The Clear BSD License

Copyright (c) 2019-2026, Fraunhofer-Gesellschaft zur Förderung der angewandten Forschung e.V. & The VVenC Authors.

All rights reserved.

Redistribution and use in source and binary forms, with or without modification,

are permitted (subject to the limitations in the disclaimer below) provided that

the following conditions are met:

     * Redistributions of source code must retain the above copyright notice,

     this list of conditions and the following disclaimer.

     * Redistributions in binary form must reproduce the above copyright

     notice, this list of conditions and the following disclaimer in the

     documentation and/or other materials provided with the distribution.

     * Neither the name of the copyright holder nor the names of its

     contributors may be used to endorse or promote products derived from this

     software without specific prior written permission.

NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY

THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND

CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT

LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A

PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR

CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,

EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,

PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR

BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER

IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)

ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE

POSSIBILITY OF SUCH DAMAGE.

------------------------------------------------------------------------------------------- */

/** \file     EncSearch.cpp

 *  \brief    encoder intra search class

 */

#include "IntraSearch.h"

#include "EncPicture.h"

#include "CommonLib/CommonDef.h"

#include "CommonLib/Rom.h"

#include "CommonLib/Picture.h"

#include "CommonLib/UnitTools.h"

#include "CommonLib/dtrace_next.h"

#include "CommonLib/dtrace_buffer.h"

#include <math.h>

#include "vvenc/vvencCfg.h"

//! \ingroup EncoderLib

//! \{

namespace vvenc {

#define PLTCtx(c) SubCtx( Ctx::Palette, c )

IntraSearch::IntraSearch()

  : m_pSaveCS       (nullptr)

  , m_pcEncCfg      (nullptr)

  , m_pcTrQuant     (nullptr)

  , m_pcRdCost      (nullptr)

  , m_CABACEstimator(nullptr)

  , m_CtxCache      (nullptr)

{

}

void IntraSearch::init(const VVEncCfg &encCfg, TrQuant *pTrQuant, RdCost *pRdCost, SortedPelUnitBufs<SORTED_BUFS> *pSortedPelUnitBufs, XUCache &unitCache )

{

  IntraPrediction::init( encCfg.m_internChromaFormat, encCfg.m_internalBitDepth[ CH_L ] );

  m_pcEncCfg          = &encCfg;

  m_pcTrQuant         = pTrQuant;

  m_pcRdCost          = pRdCost;

  m_SortedPelUnitBufs = pSortedPelUnitBufs;

  const ChromaFormat chrFormat = encCfg.m_internChromaFormat;

  const int maxCUSize          = encCfg.m_CTUSize;

  Area area = Area( 0, 0, maxCUSize, maxCUSize );

  m_pTempCS = new CodingStructure( unitCache, nullptr );

  m_pBestCS = new CodingStructure( unitCache, nullptr );

  m_pTempCS->createForSearch( chrFormat, area );

  m_pBestCS->createForSearch( chrFormat, area );

  const int uiNumSaveLayersToAllocate = 3;

  m_pSaveCS = new CodingStructure*[uiNumSaveLayersToAllocate];

  for( int layer = 0; layer < uiNumSaveLayersToAllocate; layer++ )

  {

    m_pSaveCS[ layer ] = new CodingStructure( unitCache, nullptr );

    m_pSaveCS[ layer ]->createForSearch( chrFormat, Area( 0, 0, maxCUSize, maxCUSize ) );

    m_pSaveCS[ layer ]->initStructData();

  }

  CompArea chromaArea( COMP_Cb, chrFormat, area, true );

  for( int i = 0; i < 5; i++ )

  {

    m_orgResiCb[i].create( chromaArea );

    m_orgResiCr[i].create( chromaArea );

  }

}

void IntraSearch::destroy()

{

  if ( m_pSaveCS )

  {

    const int uiNumSaveLayersToAllocate = 3;

    for( int layer = 0; layer < uiNumSaveLayersToAllocate; layer++ )

    {

      if ( m_pSaveCS[ layer ] ) { m_pSaveCS[ layer ]->destroy(); delete m_pSaveCS[ layer ]; }

    }

    delete[] m_pSaveCS;

    m_pSaveCS = nullptr;

  }

  if( m_pTempCS )

  {

    m_pTempCS->destroy();

    delete m_pTempCS; m_pTempCS = nullptr;

  }

  if( m_pBestCS )

  {

    m_pBestCS->destroy();

    delete m_pBestCS; m_pBestCS = nullptr;

  }

}

IntraSearch::~IntraSearch()

{

  destroy();

}

void IntraSearch::setCtuEncRsrc( CABACWriter* cabacEstimator, CtxCache *ctxCache )

{

  m_CABACEstimator = cabacEstimator;

  m_CtxCache       = ctxCache;

}

//////////////////////////////////////////////////////////////////////////

// INTRA PREDICTION

//////////////////////////////////////////////////////////////////////////

static constexpr double COST_UNKNOWN = -65536.0;

double IntraSearch::xFindInterCUCost( CodingUnit &cu )

{

  if( CU::isConsIntra(cu) && !cu.slice->isIntra() )

  {

    //search corresponding inter CU cost

    for( int i = 0; i < m_numCuInSCIPU; i++ )

    {

      if( cu.lumaPos() == m_cuAreaInSCIPU[i].pos() && cu.lumaSize() == m_cuAreaInSCIPU[i].size() )

      {

        return m_cuCostInSCIPU[i];

      }

    }

  }

  return COST_UNKNOWN;

}

void IntraSearch::xEstimateLumaRdModeList(int& numModesForFullRD,

  static_vector<ModeInfo, FAST_UDI_MAX_RDMODE_NUM>& RdModeList,

  static_vector<ModeInfo, FAST_UDI_MAX_RDMODE_NUM>& HadModeList,

  static_vector<double, FAST_UDI_MAX_RDMODE_NUM>& CandCostList,

  static_vector<double, FAST_UDI_MAX_RDMODE_NUM>& CandHadList, CodingUnit& cu, bool testMip )

{

  PROFILER_SCOPE_AND_STAGE_EXT( 1, _TPROF, P_INTRA_EST_RD_CAND, cu.cs, CH_L );

  const uint16_t intra_ctx_size = Ctx::IntraLumaMpmFlag.size() + Ctx::IntraLumaPlanarFlag.size() + Ctx::MultiRefLineIdx.size() + Ctx::ISPMode.size() + Ctx::MipFlag.size();

  const TempCtx  ctxStartIntraCtx(m_CtxCache, SubCtx(CtxSet(Ctx::IntraLumaMpmFlag(), intra_ctx_size), m_CABACEstimator->getCtx()));

  const double   sqrtLambdaForFirstPass = m_pcRdCost->getMotionLambda() * FRAC_BITS_SCALE;

  const int numModesAvailable = NUM_LUMA_MODE; // total number of Intra modes

  CHECK(numModesForFullRD >= numModesAvailable, "Too many modes for full RD search");

  const SPS& sps     = *cu.cs->sps;

  const bool fastMip = sps.MIP && m_pcEncCfg->m_useFastMIP;

  // this should always be true

  CHECK( !cu.Y().valid(), "CU is not valid" );

  const CompArea& area = cu.Y();

  const UnitArea localUnitArea(area.chromaFormat, Area(0, 0, area.width, area.height));

  if( testMip)

  {

    numModesForFullRD += fastMip ? numModesForFullRD - std::min( m_pcEncCfg->m_useFastMIP, numModesForFullRD )

                                 : numModesForFullRD;

    m_SortedPelUnitBufs->prepare( localUnitArea, numModesForFullRD + 1 );

  }

  else

  {

    m_SortedPelUnitBufs->prepare( localUnitArea, numModesForFullRD );

  }

  CPelBuf piOrg   = cu.cs->getOrgBuf(COMP_Y);

  PelBuf piPred  = m_SortedPelUnitBufs->getTestBuf(COMP_Y);

  DistParam distParam    = m_pcRdCost->setDistParam( piOrg, piPred, sps.bitDepths[ CH_L ], DF_HAD_2SAD); // Use HAD (SATD) cost

  const int numHadCand = (testMip ? 2 : 1) * 3;

  //*** Derive (regular) candidates using Hadamard

  cu.mipFlag = false;

  cu.multiRefIdx = 0;

  //===== init pattern for luma prediction =====

  initIntraPatternChType(cu, cu.Y(), true);

  bool satdChecked[NUM_INTRA_MODE] = { false };

  const unsigned* mpmLst = m_searchMpmLst;   // derived for this CU in estIntraPredLumaQT

  const int decMsk = ( 1 << m_pcEncCfg->m_IntraEstDecBit ) - 1;

  m_parentCandList.resize( 0 );

  m_parentCandList.reserve( ( numModesAvailable >> m_pcEncCfg->m_IntraEstDecBit ) + 2 );

  for( unsigned mode = 0; mode < numModesAvailable; mode++ )

  {

    // Skip checking extended Angular modes in the first round of SATD

    if( mode > DC_IDX && ( mode & decMsk ) )

    {

      continue;

    }

    m_parentCandList.push_back( ModeInfo( false, false, 0, NOT_INTRA_SUBPARTITIONS, mode ) );

  }

  for( int decDst = 1 << m_pcEncCfg->m_IntraEstDecBit; decDst > 0; decDst >>= 1 )

  {

    for( unsigned idx = 0; idx < m_parentCandList.size(); idx++ )

    {

      int modeParent = m_parentCandList[idx].modeId;

      int off = decDst & decMsk;

      int inc = decDst << 1;

#if 1 // INTRA_AS_IN_VTM

      if( off != 0 && ( modeParent <= ( DC_IDX + 1 ) || modeParent >= ( NUM_LUMA_MODE - 1 ) ) )

      {

        continue;

      }

#endif

      for( int mode = modeParent - off; mode < modeParent + off + 1; mode += inc )

      {

        if( satdChecked[mode] || mode < 0 || mode >= NUM_LUMA_MODE )

        {

          continue;

        }

        cu.intraDir[0] = mode;

        initPredIntraParams( cu, cu.Y(), sps );

        distParam.cur.buf = piPred.buf = m_SortedPelUnitBufs->getTestBuf().Y().buf;

        predIntraAng( COMP_Y, piPred, cu );

        // Use the min between SAD and HAD as the cost criterion

        // SAD is scaled by 2 to align with the scaling of HAD

        Distortion minSadHad = distParam.distFunc( distParam );

        uint64_t fracModeBits = xFracModeBitsIntraLuma( cu, mpmLst );

        //restore ctx

        m_CABACEstimator->getCtx() = SubCtx( CtxSet( Ctx::IntraLumaMpmFlag(), intra_ctx_size ), ctxStartIntraCtx );

        double cost = ( double ) minSadHad + ( double ) fracModeBits * sqrtLambdaForFirstPass;

        DTRACE( g_trace_ctx, D_INTRA_COST, "IntraHAD: %u, %llu, %f (%d)\n", minSadHad, fracModeBits, cost, mode );

        int insertPos = -1;

        updateCandList( ModeInfo( false, false, 0, NOT_INTRA_SUBPARTITIONS, mode ), cost, RdModeList, CandCostList, numModesForFullRD, &insertPos );

        updateCandList( ModeInfo( false, false, 0, NOT_INTRA_SUBPARTITIONS, mode ), ( double ) minSadHad, HadModeList, CandHadList, numHadCand );

        m_SortedPelUnitBufs->insert( insertPos, ( int ) RdModeList.size() );

        satdChecked[mode] = true;

      }

    }

    m_parentCandList.resize( RdModeList.size() );

    std::copy( RdModeList.cbegin(), RdModeList.cend(), m_parentCandList.begin() );

  }

  const bool isFirstLineOfCtu = (((cu.block(COMP_Y).y)&((cu.cs->sps)->CTUSize - 1)) == 0);

  if( m_pcEncCfg->m_MRL && ! isFirstLineOfCtu )

  {

    cu.multiRefIdx = 1;

    unsigned  multiRefMPM [NUM_MOST_PROBABLE_MODES];

    CU::getIntraMPMs(cu, multiRefMPM);

    for (int mRefNum = 1; mRefNum < MRL_NUM_REF_LINES; mRefNum++)

    {

      int multiRefIdx = MULTI_REF_LINE_IDX[mRefNum];

      cu.multiRefIdx = multiRefIdx;

      initIntraPatternChType(cu, cu.Y(), true);

      for (int x = 1; x < NUM_MOST_PROBABLE_MODES; x++)

      {

        cu.intraDir[0] = multiRefMPM[x];

        initPredIntraParams(cu, cu.Y(), sps);

        distParam.cur.buf = piPred.buf = m_SortedPelUnitBufs->getTestBuf().Y().buf;

        predIntraAng(COMP_Y, piPred, cu);

        // Use the min between SAD and SATD as the cost criterion

        // SAD is scaled by 2 to align with the scaling of HAD

        Distortion minSadHad = distParam.distFunc(distParam);

        // NB xFracModeBitsIntra will not affect the mode for chroma that may have already been pre-estimated.

        uint64_t fracModeBits = xFracModeBitsIntraLuma( cu, mpmLst );

        //restore ctx

        m_CABACEstimator->getCtx() = SubCtx(CtxSet(Ctx::IntraLumaMpmFlag(), intra_ctx_size), ctxStartIntraCtx);

        double cost = (double) minSadHad + (double) fracModeBits * sqrtLambdaForFirstPass;

//        DTRACE(g_trace_ctx, D_INTRA_COST, "IntraMRL: %u, %llu, %f (%d)\n", minSadHad, fracModeBits, cost, cu.intraDir[0]);

        int insertPos = -1;

        updateCandList( ModeInfo( false, false, multiRefIdx, NOT_INTRA_SUBPARTITIONS, cu.intraDir[0] ), cost, RdModeList,  CandCostList, numModesForFullRD, &insertPos );

        updateCandList( ModeInfo( false, false, multiRefIdx, NOT_INTRA_SUBPARTITIONS, cu.intraDir[0] ), (double)minSadHad, HadModeList, CandHadList,  numHadCand );

        m_SortedPelUnitBufs->insert(insertPos, (int)RdModeList.size());

      }

    }

    cu.multiRefIdx = 0;

  }

  if (testMip)

  {

    cu.mipFlag = true;

    cu.multiRefIdx = 0;

    double mipHadCost[MAX_NUM_MIP_MODE] = { MAX_DOUBLE };

    initIntraPatternChType(cu, cu.Y());

    initIntraMip( cu );

    const int transpOff    = getNumModesMip( cu.Y() );

    const int numModesFull = (transpOff << 1);

    for( uint32_t uiModeFull = 0; uiModeFull < numModesFull; uiModeFull++ )

    {

      const bool     isTransposed = (uiModeFull >= transpOff ? true : false);

      const uint32_t uiMode       = (isTransposed ? uiModeFull - transpOff : uiModeFull);

      cu.mipTransposedFlag = isTransposed;

      cu.intraDir[CH_L] = uiMode;

      distParam.cur.buf = piPred.buf = m_SortedPelUnitBufs->getTestBuf().Y().buf;

      predIntraMip(piPred, cu);

      // Use the min between SAD and HAD as the cost criterion

      // SAD is scaled by 2 to align with the scaling of HAD

      Distortion minSadHad = distParam.distFunc(distParam);

      uint64_t fracModeBits = xFracModeBitsIntraLuma( cu, mpmLst );

      //restore ctx

      m_CABACEstimator->getCtx() = SubCtx(CtxSet(Ctx::IntraLumaMpmFlag(), intra_ctx_size), ctxStartIntraCtx);

      double cost = double(minSadHad) + double(fracModeBits) * sqrtLambdaForFirstPass;

      mipHadCost[uiModeFull] = cost;

      DTRACE(g_trace_ctx, D_INTRA_COST, "IntraMIP: %u, %llu, %f (%d)\n", minSadHad, fracModeBits, cost, uiModeFull);

      int insertPos = -1;

      updateCandList( ModeInfo( true, isTransposed, 0, NOT_INTRA_SUBPARTITIONS, cu.intraDir[0] ), cost, RdModeList,  CandCostList, numModesForFullRD+1, &insertPos );

      updateCandList( ModeInfo( true, isTransposed, 0, NOT_INTRA_SUBPARTITIONS, cu.intraDir[0] ), 0.8*(double)minSadHad, HadModeList, CandHadList,  numHadCand );

      m_SortedPelUnitBufs->insert(insertPos, (int)RdModeList.size());

    }

    const double thresholdHadCost = 1.0 + 1.4 / sqrt((double)(cu.lwidth()*cu.lheight()));

    xReduceHadCandList(RdModeList, CandCostList, *m_SortedPelUnitBufs, numModesForFullRD, thresholdHadCost, mipHadCost, cu, fastMip);

  }

  if( m_pcEncCfg->m_bFastUDIUseMPMEnabled )

  {

    const int numMPMs = NUM_MOST_PROBABLE_MODES;

    unsigned  intraMpms[numMPMs];

    cu.multiRefIdx = 0;

    const int numCand = CU::getIntraMPMs( cu, intraMpms );

    ModeInfo mostProbableMode(false, false, 0, NOT_INTRA_SUBPARTITIONS, 0);

    for( int j = 0; j < numCand; j++ )

    {

      bool mostProbableModeIncluded = false;

      mostProbableMode.modeId = intraMpms[j];

      for( int i = 0; i < numModesForFullRD; i++ )

      {

        mostProbableModeIncluded |= ( mostProbableMode == RdModeList[i] );

      }

      if( !mostProbableModeIncluded )

      {

        numModesForFullRD++;

        RdModeList.push_back( mostProbableMode );

        CandCostList.push_back(0);

      }

    }

  }

}

bool IntraSearch::estIntraPredLumaQT(CodingUnit &cu, Partitioner &partitioner, double bestCost)

{

  CU::getIntraMPMs( cu, m_searchMpmLst );

  m_searchMpmArea  = cu.Y();

  m_searchMpmValid = true;

  const bool validReturn = xEstIntraPredLumaQT( cu, partitioner, bestCost );

  m_searchMpmValid = false;

  return validReturn;

}

bool IntraSearch::xEstIntraPredLumaQT(CodingUnit &cu, Partitioner &partitioner, double bestCost)

{

  CodingStructure       &cs           = *cu.cs;

  const int             width         = partitioner.currArea().lwidth();

  const int             height        = partitioner.currArea().lheight();

  //===== loop over partitions =====

  const TempCtx ctxStart           ( m_CtxCache, m_CABACEstimator->getCtx() );

  // variables for saving fast intra modes scan results across multiple LFNST passes

  double costInterCU = xFindInterCUCost( cu );

  bool validReturn = false;

  //===== determine set of modes to be tested (using prediction signal only) =====

  int numModesAvailable = NUM_LUMA_MODE; // total number of Intra modes

  static_vector<ModeInfo, FAST_UDI_MAX_RDMODE_NUM> RdModeList;

  static_vector<ModeInfo, FAST_UDI_MAX_RDMODE_NUM> HadModeList;

  static_vector<double, FAST_UDI_MAX_RDMODE_NUM> CandCostList;

  static_vector<double, FAST_UDI_MAX_RDMODE_NUM> CandHadList;

  int numModesForFullRD = g_aucIntraModeNumFast_UseMPM_2D[Log2(width) - MIN_CU_LOG2][Log2(height) - MIN_CU_LOG2];

  if (m_pcEncCfg->m_numIntraModesFullRD > 0)

    numModesForFullRD=m_pcEncCfg->m_numIntraModesFullRD;

#if INTRA_FULL_SEARCH

  numModesForFullRD = numModesAvailable;

#endif

  const SPS& sps = *cu.cs->sps;

  const bool mipAllowed = sps.MIP && cu.lwidth() <= sps.getMaxTbSize() && cu.lheight() <= sps.getMaxTbSize() && ((cu.lfnstIdx == 0) || allowLfnstWithMip(cu.lumaSize()));

  const int SizeThr     = 8 >> std::max( 0, m_pcEncCfg->m_useFastMIP - 1 );

  const bool testMip    = mipAllowed && ( cu.lwidth() <= ( SizeThr * cu.lheight() ) && cu.lheight() <= ( SizeThr * cu.lwidth() ) ) && ( cu.lwidth() <= MIP_MAX_WIDTH && cu.lheight() <= MIP_MAX_HEIGHT );

  bool testISP = sps.ISP && CU::canUseISP(width, height, cu.cs->sps->getMaxTbSize());

  if (testISP)

  {

    int numTotalPartsHor = (int)width >> floorLog2(CU::getISPSplitDim(width, height, TU_1D_VERT_SPLIT));

    int numTotalPartsVer = (int)height >> floorLog2(CU::getISPSplitDim(width, height, TU_1D_HORZ_SPLIT));

    m_ispTestedModes[0].init(numTotalPartsHor, numTotalPartsVer, 0);

    // the total number of subpartitions is modified to take into account the cases where LFNST cannot be combined with

    // ISP due to size restrictions

    numTotalPartsHor = sps.LFNST && CU::canUseLfnstWithISP(cu.Y(), HOR_INTRA_SUBPARTITIONS) ? numTotalPartsHor : 0;

    numTotalPartsVer = sps.LFNST && CU::canUseLfnstWithISP(cu.Y(), VER_INTRA_SUBPARTITIONS) ? numTotalPartsVer : 0;

    for (int j = 1; j < NUM_LFNST_NUM_PER_SET; j++)

    {

      m_ispTestedModes[j].init(numTotalPartsHor, numTotalPartsVer, 0);

    }

    testISP = m_ispTestedModes[0].numTotalParts[0];

  }

  else

  {

    m_ispTestedModes[0].init(0, 0, 0);

  }

  xEstimateLumaRdModeList(numModesForFullRD, RdModeList, HadModeList, CandCostList, CandHadList, cu, testMip);

  CHECK( (size_t)numModesForFullRD != RdModeList.size(), "Inconsistent state!" );

  // after this point, don't use numModesForFullRD

  if( m_pcEncCfg->m_usePbIntraFast && !cs.slice->isIntra() && RdModeList.size() < numModesAvailable )

  {

    double pbintraRatio = m_pcEncCfg->m_usePbIntraFast == 1 && ( cs.area.lwidth() >= 16 && cs.area.lheight() >= 16 ) ? 1.2 : PBINTRA_RATIO;

    int maxSize = -1;

    ModeInfo bestMipMode;

    int bestMipIdx = -1;

    for( int idx = 0; idx < RdModeList.size(); idx++ )

    {

      if( RdModeList[idx].mipFlg )

      {

        bestMipMode = RdModeList[idx];

        bestMipIdx = idx;

        break;

      }

    }

    const int numHadCand = 3;

    for (int k = numHadCand - 1; k >= 0; k--)

    {

      if (CandHadList.size() < (k + 1) || CandHadList[k] > cs.interHad * pbintraRatio) { maxSize = k; }

    }

    if (maxSize > 0)

    {

      RdModeList.resize(std::min<size_t>(RdModeList.size(), maxSize));

      if( bestMipIdx >= 0 )

      {

        if( RdModeList.size() <= bestMipIdx )

        {

          RdModeList.push_back(bestMipMode);

          m_SortedPelUnitBufs->swap( maxSize, bestMipIdx );

        }

      }

    }

    if (maxSize == 0)

    {

      cs.dist = MAX_DISTORTION;

      cs.interHad = 0;

      return false;

    }

  }

  //===== check modes (using r-d costs) =====

  ModeInfo bestPUMode;

  CodingStructure *csTemp = m_pTempCS;

  CodingStructure *csBest = m_pBestCS;

  csTemp->slice   = csBest->slice   = cs.slice;

  csTemp->picture = csBest->picture = cs.picture;

  csTemp->compactResize( cu );

  csBest->compactResize( cu );

  csTemp->initStructData();

  csBest->initStructData();

  int   bestLfnstIdx  = 0;

  const bool useBDPCM = cs.picture->useBDPCM;

  int   NumBDPCMCand  = (useBDPCM && sps.BDPCM && CU::bdpcmAllowed(cu, ComponentID(partitioner.chType))) ? 2 : 0;

  int   bestbdpcmMode = 0;

  int   bestISP       = 0;

  int   bestMrl       = 0;

  bool  bestMip       = 0;

  int   EndMode       = (int)RdModeList.size();

  bool  useISPlfnst   = testISP && sps.LFNST;

  bool  noLFNST_ts    = false;

  double bestCostIsp[2] = { MAX_DOUBLE, MAX_DOUBLE };

  bool disableMTS = false;

  bool disableLFNST = false;

  bool disableDCT2test = false;

  if (m_pcEncCfg->m_FastIntraTools)

  {

    int speedIntra = 0;

    xSpeedUpIntra(bestCost, EndMode, speedIntra, cu);

    disableMTS = (speedIntra >> 2 ) & 0x1;

    disableLFNST = (speedIntra >> 1) & 0x1;

    disableDCT2test = speedIntra>>3;

    if (disableLFNST)

    {

      noLFNST_ts = true;

      useISPlfnst = false;

    }

    if (speedIntra & 0x1)

    {

      testISP = false;

    }

  }

  for (int mode_cur = 0; mode_cur < EndMode + NumBDPCMCand; mode_cur++)

  {

    int mode = mode_cur;

    if (mode_cur >= EndMode)

    {

      mode = mode_cur - EndMode ? -1 : -2;

      testISP = false;

    }

    // set CU/PU to luma prediction mode

    ModeInfo testMode;

    int noISP = 0;

    int endISP = testISP ? 2 : 0;

    bool noLFNST = false || noLFNST_ts;

    if (mode && useISPlfnst)

    {

      noLFNST |= (bestCostIsp[0] > (bestCostIsp[1] * 1.4));

      if (mode > 2)

      {

        endISP = 0;

        testISP = false;

      }

    }

    if (testISP)

    {

      xSpeedUpISP(1, testISP, mode, noISP, endISP, cu, RdModeList, bestPUMode, bestISP, bestLfnstIdx);

    }

    int startISP = 0;

    if (disableDCT2test && mode && bestISP)

    {

      startISP = endISP ? 1 : 0;

    }

    for (int ispM = startISP; ispM <= endISP; ispM++)

    {

      if (ispM && (ispM == noISP))

      {

        continue;

      }

      if (mode < 0)

      {

        cu.bdpcmM[CH_L] = -mode;

        testMode = ModeInfo(false, false, 0, NOT_INTRA_SUBPARTITIONS, cu.bdpcmM[CH_L] == 2 ? VER_IDX : HOR_IDX);

      }

      else

      {

        testMode = RdModeList[mode];

        cu.bdpcmM[CH_L] = 0;

      }

      cu.ispMode = ispM;

      cu.mipFlag = testMode.mipFlg;

      cu.mipTransposedFlag = testMode.mipTrFlg;

      cu.multiRefIdx = testMode.mRefId;

      cu.intraDir[CH_L] = testMode.modeId;

      if (cu.ispMode && xSpeedUpISP(0, testISP, mode, noISP, endISP, cu, RdModeList, bestPUMode, bestISP, 0) )

      {

        continue;

      }

      if (m_pcEncCfg->m_FastIntraTools && (cu.ispMode || sps.LFNST || sps.MTS))

      {

        m_ispTestedModes[0].intraWasTested = true;

      }

      CHECK(cu.mipFlag && cu.multiRefIdx, "Error: combination of MIP and MRL not supported");

      CHECK(cu.multiRefIdx && (cu.intraDir[0] == PLANAR_IDX), "Error: combination of MRL and Planar mode not supported");

      CHECK(cu.ispMode && cu.mipFlag, "Error: combination of ISP and MIP not supported");

      CHECK(cu.ispMode && cu.multiRefIdx, "Error: combination of ISP and MRL not supported");

      // determine residual for partition

      cs.initSubStructure(*csTemp, partitioner.chType, cs.area, true);

      int doISP = (((cu.ispMode == 0) && noLFNST) || (useISPlfnst && mode && cu.ispMode && (bestLfnstIdx == 0)) || disableLFNST) ? -mode : mode;

      xIntraCodingLumaQT(*csTemp, partitioner, m_SortedPelUnitBufs->getBufFromSortedList(mode), bestCost, doISP, disableMTS);

      DTRACE(g_trace_ctx, D_INTRA_COST, "IntraCost T [x=%d,y=%d,w=%d,h=%d] %f (%d,%d,%d,%d,%d,%d) \n", cu.blocks[0].x,

        cu.blocks[0].y, width, height, csTemp->cost, testMode.modeId, testMode.ispMod,

        cu.multiRefIdx, cu.mipFlag, cu.lfnstIdx, cu.mtsFlag);

      if (cu.ispMode && !csTemp->cus[0]->firstTU->cbf[COMP_Y])

      {

        csTemp->cost = MAX_DOUBLE;

        csTemp->costDbOffset = 0;

      }

      if (useISPlfnst)

      {

        int n = (cu.ispMode == 0) ? 0 : 1;

        bestCostIsp[n] = csTemp->cost < bestCostIsp[n] ? csTemp->cost : bestCostIsp[n];

      }

      // check r-d cost

      if (csTemp->cost < csBest->cost)

      {

        validReturn   = true;

        std::swap(csTemp, csBest);

        bestPUMode    = testMode;

        bestLfnstIdx  = csBest->cus[0]->lfnstIdx;

        bestISP       = csBest->cus[0]->ispMode;

        bestMip       = csBest->cus[0]->mipFlag;

        bestMrl       = csBest->cus[0]->multiRefIdx;

        bestbdpcmMode = cu.bdpcmM[CH_L];

        m_ispTestedModes[bestLfnstIdx].bestSplitSoFar = ISPType(bestISP);

        if (csBest->cost < bestCost)

        {

          bestCost = csBest->cost;

        }

        if ((csBest->getTU(partitioner.chType)->mtsIdx[COMP_Y] == MTS_SKIP) && ( floorLog2(csBest->getTU(partitioner.chType)->blocks[COMP_Y].area()) >= 6 ))

        {

          noLFNST_ts = 1;

        }

      }

      // reset context models

      m_CABACEstimator->getCtx() = ctxStart;

      csTemp->releaseIntermediateData();

      if (m_pcEncCfg->m_fastLocalDualTreeMode && CU::isConsIntra(cu) && !cu.slice->isIntra() && csBest->cost != MAX_DOUBLE && costInterCU != COST_UNKNOWN && mode >= 0)

      {

        if( (m_pcEncCfg->m_fastLocalDualTreeMode == 2) || (csBest->cost > costInterCU * 1.5))

        {

          //Note: only try one intra mode, which is especially useful to reduce EncT for LDB case (around 4%)

          EndMode = 0;

          break;

        }

      }

    }

  } // Mode loop

  if (m_pcEncCfg->m_FastIntraTools && (sps.ISP|| sps.LFNST || sps.MTS))

  {

    int bestMode = csBest->getTU(partitioner.chType)->mtsIdx[COMP_Y] ? 4 : 0;

    bestMode |= bestLfnstIdx ? 2 : 0;

    bestMode |= bestISP ? 1 : 0;

    m_ispTestedModes[0].bestIntraMode = bestMode;

  }

  cu.ispMode = bestISP;

  if( validReturn )

  {

    cs.useSubStructure( *csBest, partitioner.chType, TREE_D, cu.singleChan( CH_L ), true );

    //=== update PU data ====

    cu.lfnstIdx           = bestLfnstIdx;

    cu.mipTransposedFlag  = bestPUMode.mipTrFlg;

    cu.intraDir[CH_L]     = bestPUMode.modeId;

    cu.bdpcmM[CH_L]       = bestbdpcmMode;

    cu.mipFlag            = bestMip;

    cu.multiRefIdx        = bestMrl;

  }

  else

  {

    THROW("fix this");

  }

  csBest->releaseIntermediateData();

  return validReturn;

}

void IntraSearch::estIntraPredChromaQT( CodingUnit& cu, Partitioner& partitioner, const double maxCostAllowed )

{

  PROFILER_SCOPE_AND_STAGE_EXT( 0, _TPROF, P_INTRA_CHROMA, cu.cs, CH_C );

  const TempCtx ctxStart( m_CtxCache, m_CABACEstimator->getCtx() );

  CodingStructure &cs   = *cu.cs;

  bool lumaUsesISP      = !CU::isSepTree(cu) && cu.ispMode;

  PartSplit ispType     = lumaUsesISP ? CU::getISPType(cu, COMP_Y) : TU_NO_ISP;

  double bestCostSoFar  = maxCostAllowed;

  const uint32_t numberValidComponents = getNumberValidComponents( cu.chromaFormat );

  const bool useBDPCM   = cs.picture->useBDPCM;

  uint32_t   uiBestMode = 0;

  Distortion uiBestDist = 0;

  double     dBestCost  = MAX_DOUBLE;

  //----- init mode list ----

  {

    uint32_t  uiMinMode = 0;

    uint32_t  uiMaxMode = NUM_CHROMA_MODE;

    const int reducedModeNumber = uiMaxMode >> (m_pcEncCfg->m_reduceIntraChromaModesFullRD ? 1 : 2);

    //----- check chroma modes -----

    uint32_t chromaCandModes[ NUM_CHROMA_MODE ];

    CU::getIntraChromaCandModes( cu, chromaCandModes );

    // create a temporary CS

    CodingStructure &saveCS = *m_pSaveCS[0];

    saveCS.pcv      = cs.pcv;

    saveCS.picture  = cs.picture;

    saveCS.area.repositionTo( cs.area );

    saveCS.clearTUs();

    if( !CU::isSepTree(cu) && cu.ispMode )

    {

      saveCS.clearCUs();

    }

    if( CU::isSepTree(cu) )

    {

      if( partitioner.canSplit( TU_MAX_TR_SPLIT, cs ) )

      {

        partitioner.splitCurrArea( TU_MAX_TR_SPLIT, cs );

        do

        {

          cs.addTU( CS::getArea( cs, partitioner.currArea(), partitioner.chType, partitioner.treeType ), partitioner.chType, &cu ).depth = partitioner.currTrDepth;

        } while( partitioner.nextPart( cs ) );

        partitioner.exitCurrSplit();

      }

      else

        cs.addTU( CS::getArea( cs, partitioner.currArea(), partitioner.chType, partitioner.treeType ), partitioner.chType, &cu );

    }

    // create a store for the TUs

    std::vector<TransformUnit*> orgTUs;

    for( const auto &ptu : cs.tus )

    {

      // for split TUs in HEVC, add the TUs without Chroma parts for correct setting of Cbfs

      if (lumaUsesISP || cu.contains(*ptu, CH_C))

      {

        saveCS.addTU( *ptu, partitioner.chType, nullptr );

        orgTUs.push_back( ptu );

      }

    }

    // SATD pre-selecting.

    int     satdModeList  [NUM_CHROMA_MODE] = { 0 };

    int64_t satdSortedCost[NUM_CHROMA_MODE] = { 0 };

    bool    modeDisable[NUM_INTRA_MODE + 1] = { false }; // use intra mode idx to check whether enable

    CodingStructure& cs = *(cu.cs);

    CompArea areaCb = cu.Cb();

    CompArea areaCr = cu.Cr();

    CPelBuf orgCb  = cs.getOrgBuf (COMP_Cb);

    PelBuf predCb  = cs.getPredBuf(COMP_Cb);

    CPelBuf orgCr  = cs.getOrgBuf (COMP_Cr);

    PelBuf predCr  = cs.getPredBuf(COMP_Cr);

    DistParam distParamSadCb  = m_pcRdCost->setDistParam( orgCb, predCb, cu.cs->sps->bitDepths[ CH_C ], DF_SAD);

    DistParam distParamSatdCb = m_pcRdCost->setDistParam( orgCb, predCb, cu.cs->sps->bitDepths[ CH_C ], DF_HAD);

    DistParam distParamSadCr  = m_pcRdCost->setDistParam( orgCr, predCr, cu.cs->sps->bitDepths[ CH_C ], DF_SAD);

    DistParam distParamSatdCr = m_pcRdCost->setDistParam( orgCr, predCr, cu.cs->sps->bitDepths[ CH_C ], DF_HAD);

    cu.intraDir[1] = MDLM_L_IDX; // temporary assigned, just to indicate this is a MDLM mode. for luma down-sampling operation.

    initIntraPatternChType(cu, cu.Cb());

    initIntraPatternChType(cu, cu.Cr());

    loadLMLumaRecPels(cu, cu.Cb());

    for (int idx = uiMinMode; idx < uiMaxMode; idx++)

    {

      int mode = chromaCandModes[idx];

      satdModeList[idx] = mode;

      if (CU::isLMCMode(mode) && ( !CU::isLMCModeEnabled(cu, mode) || cu.slice->lmChromaCheckDisable ) )

      {

        continue;

      }

      if ((mode == LM_CHROMA_IDX) || (mode == PLANAR_IDX) || (mode == DM_CHROMA_IDX)) // only pre-check regular modes and MDLM modes, not including DM ,Planar, and LM

      {

        continue;

      }

      cu.intraDir[1]    = mode; // temporary assigned, for SATD checking.

      const bool isLMCMode = CU::isLMCMode(mode);

      if( isLMCMode )

      {

        predIntraChromaLM(COMP_Cb, predCb, cu, areaCb, mode);

      }

      else

      {

        initPredIntraParams(cu, cu.Cb(), *cs.sps);

        predIntraAng(COMP_Cb, predCb, cu);

      }

      int64_t sadCb = distParamSadCb.distFunc(distParamSadCb) * 2;

      int64_t satdCb = distParamSatdCb.distFunc(distParamSatdCb);

      int64_t sad = std::min(sadCb, satdCb);

      if( isLMCMode )

      {

        predIntraChromaLM(COMP_Cr, predCr, cu, areaCr, mode);

      }

      else

      {

        initPredIntraParams(cu, cu.Cr(), *cs.sps);

        predIntraAng(COMP_Cr, predCr, cu);

      }

      int64_t sadCr = distParamSadCr.distFunc(distParamSadCr) * 2;

      int64_t satdCr = distParamSatdCr.distFunc(distParamSatdCr);

      sad += std::min(sadCr, satdCr);

      satdSortedCost[idx] = sad;

    }

    // sort the mode based on the cost from small to large.

    for (int i = uiMinMode; i <= uiMaxMode - 1; i++)

    {

      for (int j = i + 1; j <= uiMaxMode - 1; j++)

      {

        if (satdSortedCost[j] < satdSortedCost[i])

        {

          std::swap( satdModeList[i],   satdModeList[j]);

          std::swap( satdSortedCost[i], satdSortedCost[j]);

        }

      }

    }

    for (int i = 0; i < reducedModeNumber; i++)

    {

      modeDisable[satdModeList[uiMaxMode - 1 - i]] = true; // disable the last reducedModeNumber modes

    }

    int bestLfnstIdx = 0;

    // save the dist

    Distortion baseDist = cs.dist;

    int32_t bestbdpcmMode = 0;

    uint32_t numbdpcmModes = ( useBDPCM && CU::bdpcmAllowed(cu, COMP_Cb)

        && ((partitioner.chType == CH_C) || (cu.ispMode == 0 && cu.lfnstIdx == 0 && cu.firstTU->mtsIdx[COMP_Y] == MTS_SKIP))) ? 2 : 0;

    for (int mode_cur = uiMinMode; mode_cur < (int)(uiMaxMode + numbdpcmModes); mode_cur++)

    {

      int mode = mode_cur;

      if (mode_cur >= uiMaxMode)

      {

        mode = mode_cur > uiMaxMode ? -1 : -2; //set bdpcm mode

        if ((mode == -1) && (saveCS.tus[0]->mtsIdx[COMP_Cb] != MTS_SKIP) && (saveCS.tus[0]->mtsIdx[COMP_Cr] != MTS_SKIP))

        {

          continue;

        }

      }

      int chromaIntraMode;

      if (mode < 0)

      {

        cu.bdpcmM[CH_C] = -mode;

        chromaIntraMode = cu.bdpcmM[CH_C] == 2 ? chromaCandModes[1] : chromaCandModes[2];

      }

      else

      {

        cu.bdpcmM[CH_C] = 0;

        chromaIntraMode = chromaCandModes[mode];

        if (CU::isLMCMode(chromaIntraMode) && ( !CU::isLMCModeEnabled(cu, chromaIntraMode) || cu.slice->lmChromaCheckDisable ) )

        {

          continue;

        }

        if (modeDisable[chromaIntraMode] && CU::isLMCModeEnabled(cu, chromaIntraMode)) // when CCLM is disable, then MDLM is disable. not use satd checking

        {

          continue;

        }

      }

      cs.dist = baseDist;

      //----- restore context models -----

      m_CABACEstimator->getCtx() = ctxStart;

      //----- chroma coding -----

      cu.intraDir[1] = chromaIntraMode;

      m_ispTestedModes[0].IspType = ispType;

      m_ispTestedModes[0].subTuCounter = -1;

      xIntraChromaCodingQT( cs, partitioner );

      if (lumaUsesISP && cs.dist == MAX_UINT)

      {

        continue;

      }

      if (cs.sps->transformSkip)

      {

        m_CABACEstimator->getCtx() = ctxStart;

      }

      m_ispTestedModes[0].IspType = ispType;

      m_ispTestedModes[0].subTuCounter = -1;

      uint64_t fracBits   = xGetIntraFracBitsQT( cs, partitioner, false );

      Distortion uiDist = cs.dist;

      double    dCost   = m_pcRdCost->calcRdCost( fracBits, uiDist - baseDist );

      //----- compare -----

      if( dCost < dBestCost )

      {

        if (lumaUsesISP && (dCost < bestCostSoFar))

        {

          bestCostSoFar = dCost;

        }

        for( uint32_t i = getFirstComponentOfChannel( CH_C ); i < numberValidComponents; i++ )

        {

          const CompArea& area = cu.blocks[i];

          saveCS.getRecoBuf     ( area ).copyFrom( cs.getRecoBuf   ( area ) );

          cs.picture->getRecoBuf( area ).copyFrom( cs.getRecoBuf   ( area ) );

          for( uint32_t j = 0; j < saveCS.tus.size(); j++ )

          {

            saveCS.tus[j]->copyComponentFrom( *orgTUs[j], area.compID );

          }

        }

        dBestCost    = dCost;

        uiBestDist   = uiDist;

        uiBestMode   = chromaIntraMode;

        bestLfnstIdx = cu.lfnstIdx;

        bestbdpcmMode = cu.bdpcmM[CH_C];

      }

    }

    cu.lfnstIdx = bestLfnstIdx;

    cu.bdpcmM[CH_C]= bestbdpcmMode;

    for( uint32_t i = getFirstComponentOfChannel( CH_C ); i < numberValidComponents; i++ )

    {

      const CompArea& area = cu.blocks[i];

      cs.getRecoBuf         ( area ).copyFrom( saveCS.getRecoBuf( area ) );

      cs.picture->getRecoBuf( area ).copyFrom( cs.getRecoBuf    ( area ) );

      for( uint32_t j = 0; j < saveCS.tus.size(); j++ )

      {

        orgTUs[ j ]->copyComponentFrom( *saveCS.tus[ j ], area.compID );

      }

    }

  }

  cu.intraDir[1] = uiBestMode;

  cs.dist        = uiBestDist;

  //----- restore context models -----

  m_CABACEstimator->getCtx() = ctxStart;

  if (lumaUsesISP && bestCostSoFar >= maxCostAllowed)

  {

    cu.ispMode = 0;

  }

}

void IntraSearch::saveCuAreaCostInSCIPU( Area area, double cost )

{

  if( m_numCuInSCIPU < NUM_INTER_CU_INFO_SAVE )

  {

    m_cuAreaInSCIPU[m_numCuInSCIPU] = area;

    m_cuCostInSCIPU[m_numCuInSCIPU] = cost;

    m_numCuInSCIPU++;

  }

}

void IntraSearch::initCuAreaCostInSCIPU()

{

  for( int i = 0; i < NUM_INTER_CU_INFO_SAVE; i++ )

  {

    m_cuAreaInSCIPU[i] = Area();

    m_cuCostInSCIPU[i] = 0;

  }