#include "TComAdaptiveColorSpace.h"
#if QC_SCC_ADAPTIVE_COLORSPACE
#include "TComRom.h"
#include "TComPic.h"

Void TComAdaptiveColorSpace::getNoCorrImageIntra(TComDataCU* pcCU, TComYuv* pcOrg, TComYuv* pcNoCorrOrg, UInt uiAbsPartIdx, UInt uiLog2TrSize, Int leftAvailable, Int aboveAvailable)
{
  Pel pixelY, pixelU, pixelV;
  const Int *v = getV_YCgCo();
  UInt blockWidth = 1<<uiLog2TrSize;
  UInt blockHeight = 1<<uiLog2TrSize;
  Int uvoffset = (1<<(g_bitDepth[CHANNEL_TYPE_LUMA] + 2 - 1));
  Bool bLL     = pcCU->getSlice()->getSPS()->getUseLossless();
  if(bLL)
    uvoffset = uvoffset >> 1;

  if (pcOrg) {
    Pel* pOrg[3] = {pcOrg->getAddr(COMPONENT_Y, uiAbsPartIdx), pcOrg->getAddr(COMPONENT_Cb, uiAbsPartIdx), pcOrg->getAddr(COMPONENT_Cr, uiAbsPartIdx)};
    UInt uiOrgStride = pcOrg->getStride(COMPONENT_Y);
    Pel* pNoCorrOrg[3] = {pcNoCorrOrg->getAddr(COMPONENT_Y, uiAbsPartIdx), pcNoCorrOrg->getAddr(COMPONENT_Cb, uiAbsPartIdx), pcNoCorrOrg->getAddr(COMPONENT_Cr, uiAbsPartIdx)};

    for (Int y = 0; y < blockHeight; y++)
    {
      for (Int x = 0; x < blockWidth; x++)
      {
        if(bLL)
        {
          pixelV = pOrg[2][x] - pOrg[1][x];
          pixelU = pOrg[0][x] - (pOrg[1][x] + (pixelV>>1));
          pixelY = pOrg[1][x] + (pixelV>>1) + (pixelU>>1);
        }
        else
        {
          pixelY = pOrg[0][x]*v[0]+ pOrg[1][x]*v[1] + pOrg[2][x]*v[2];
          pixelU = pOrg[0][x]*v[3]+ pOrg[1][x]*v[4] + pOrg[2][x]*v[5];
          pixelV = pOrg[0][x]*v[6]+ pOrg[1][x]*v[7] + pOrg[2][x]*v[8];
        }

        pixelU +=uvoffset;
        pixelV +=uvoffset;

        pNoCorrOrg[0][x] = pixelY;          
        pNoCorrOrg[1][x] = pixelU;
        pNoCorrOrg[2][x] = pixelV;
      }
      pOrg[0] += uiOrgStride;
      pOrg[1] += uiOrgStride;
      pOrg[2] += uiOrgStride;
      pNoCorrOrg[0] += uiOrgStride;
      pNoCorrOrg[1] += uiOrgStride;
      pNoCorrOrg[2] += uiOrgStride;
    }
  }

  UInt uiZOrderIdx = pcCU->getZorderIdxInCU()+uiAbsPartIdx;
  Pel* piRef[3] = {pcCU->getPic()->getPicYuvRec()->getAddr(COMPONENT_Y, pcCU->getAddr(), uiZOrderIdx), pcCU->getPic()->getPicYuvRec()->getAddr(COMPONENT_Cb, pcCU->getAddr(), uiZOrderIdx),
    pcCU->getPic()->getPicYuvRec()->getAddr(COMPONENT_Cr, pcCU->getAddr(), uiZOrderIdx)};
  Int iRefStride = static_cast<Int>(pcCU->getPic()->getStride(COMPONENT_Y));

  if (aboveAvailable)
  {
    for (Int x = 0; x < blockWidth*2; x++)
    {
      if(bLL)
      {
        pixelV = piRef[2][-iRefStride+x] -  piRef[1][-iRefStride+x];
        pixelU = piRef[0][-iRefStride+x] - (piRef[1][-iRefStride+x] + (pixelV>>1));
        pixelY = piRef[1][-iRefStride+x] + (pixelV>>1) + (pixelU>>1);
      }
      else
      {
        pixelY = piRef[0][-iRefStride+x]*v[0] + piRef[1][-iRefStride+x]*v[1] + piRef[2][-iRefStride+x]*v[2];
        pixelU = piRef[0][-iRefStride+x]*v[3] + piRef[1][-iRefStride+x]*v[4] + piRef[2][-iRefStride+x]*v[5];
        pixelV = piRef[0][-iRefStride+x]*v[6] + piRef[1][-iRefStride+x]*v[7] + piRef[2][-iRefStride+x]*v[8];
      }

      piRef[0][-iRefStride + x] = pixelY;
      piRef[1][-iRefStride + x] = pixelU + uvoffset; 
      piRef[2][-iRefStride + x] = pixelV + uvoffset; 
    }
  }
  if (leftAvailable)
  {
    for (Int y = (aboveAvailable ? -1 : 0); y < static_cast<Int>(blockHeight)*2; y++)
    {
      if(bLL)
      {
        pixelV = piRef[2][y*iRefStride-1] -  piRef[1][y*iRefStride-1];
        pixelU = piRef[0][y*iRefStride-1] - (piRef[1][y*iRefStride-1] + (pixelV>>1));
        pixelY = piRef[1][y*iRefStride-1] + (pixelV>>1) + (pixelU>>1);
      }
      else
      {
        pixelY = piRef[0][y*iRefStride-1]*v[0]+  piRef[1][y*iRefStride-1]*v[1] + piRef[2][y*iRefStride-1]*v[2];
        pixelU = piRef[0][y*iRefStride-1]*v[3]+  piRef[1][y*iRefStride-1]*v[4] + piRef[2][y*iRefStride-1]*v[5];
        pixelV = piRef[0][y*iRefStride-1]*v[6]+  piRef[1][y*iRefStride-1]*v[7] + piRef[2][y*iRefStride-1]*v[8];   
      }

      piRef[0][y*iRefStride-1] = pixelY;
      piRef[1][y*iRefStride-1]   = pixelU + uvoffset; 
      piRef[2][y*iRefStride-1]   = pixelV + uvoffset; 
    }
  }
}

Void TComAdaptiveColorSpace::getNoCorrResiInter(TComDataCU* pcCU, TComYuv* pcPred, TComYuv* pcResi, TComYuv* pcNoCorrResi)
{
  const Int *v = getV_YCgCo();

  Bool bLL     = pcCU->getSlice()->getSPS()->getUseLossless();
  Pel pixelY, pixelU, pixelV;

  for (Int iPartIdx = 0; iPartIdx < pcCU->getNumPartInter(); iPartIdx++ )
  {
    Int iWidth;
    Int iHeight;
    UInt uiPartAddr;
    pcCU->getPartIndexAndSize( iPartIdx, uiPartAddr, iWidth, iHeight );

    if (pcResi)
    {
      UInt uiResiStride = pcResi->getStride(COMPONENT_Y);
      Pel* pResi[3] = {pcResi->getAddr(COMPONENT_Y, uiPartAddr), pcResi->getAddr(COMPONENT_Cb, uiPartAddr), pcResi->getAddr(COMPONENT_Cr, uiPartAddr)};
      Pel* pNoCorrResi[3] = {pcNoCorrResi->getAddr(COMPONENT_Y, uiPartAddr), pcNoCorrResi->getAddr(COMPONENT_Cb, uiPartAddr), pcNoCorrResi->getAddr(COMPONENT_Cr, uiPartAddr)};

      for (Int y = 0; y < iHeight; y++) {
        for (Int x = 0; x < iWidth; x++) {
          if(bLL)
          {
            pixelV = pResi[2][x] -  pResi[1][x]; 
            pixelU = pResi[0][x] - (pResi[1][x] + (pixelV>>1));
            pixelY = pResi[1][x] + (pixelV>>1) + (pixelU>>1);
          }
          else
          {
            pixelY = pResi[0][x] *v[0] + pResi[1][x]*v[1] + pResi[2][x]*v[2]; 
            pixelU = pResi[0][x] *v[3] + pResi[1][x]*v[4] + pResi[2][x]*v[5];
            pixelV = pResi[0][x] *v[6] + pResi[1][x]*v[7] + pResi[2][x]*v[8];
          }

          pNoCorrResi[0][x] = pixelY;
          pNoCorrResi[1][x] = pixelU;
          pNoCorrResi[2][x] = pixelV;

        }
        pResi[0] += uiResiStride;
        pResi[1] += uiResiStride;
        pResi[2] += uiResiStride;
        pNoCorrResi[0] += uiResiStride;
        pNoCorrResi[1] += uiResiStride;
        pNoCorrResi[2] += uiResiStride;
      }
    }
  }
}

Int* TComAdaptiveColorSpace::getV_YCgCo()
{
  m_v[0][0] = (1<<BIT_V) >> 1;
  m_v[0][1] = (1<<BIT_V) >> 2;
  m_v[0][2] = (1<<BIT_V) >> 2;

  m_v[0][3] = (1<<BIT_V) >> 1;
  m_v[0][4] = ((1<<BIT_V) >> 2)*(-1);
  m_v[0][5] = ((1<<BIT_V) >> 2)*(-1);

  m_v[0][6] = 0;
  m_v[0][7] = ((1<<BIT_V) >> 1)*(-1);
  m_v[0][8] = (1<<BIT_V) >> 1;

  return m_v[0];
}
Void TComAdaptiveColorSpace::reconstructInterResi(TComDataCU* pcCU, TComYuv* pcResi, TComYuv* pcDstResi, Int iAbsPartIdx, Int iWidth, Int iHeight)
{
  Pel decYint, decUint, decVint;
  UInt uiStride = pcResi->getStride(COMPONENT_Y);
  const Int offset = 1 << (BIT_V - 1);
  Bool bLL     = pcCU->getSlice()->getSPS()->getUseLossless();

  if (iAbsPartIdx < 0)
  {
    Pel* pResi[3] = {pcResi->getAddr( COMPONENT_Y, 0 ), pcResi->getAddr( COMPONENT_Cb, 0 ), pcResi->getAddr( COMPONENT_Cr, 0 )};
    Pel* pDstResi[3] = {pcDstResi->getAddr( COMPONENT_Y, 0 ), pcDstResi->getAddr( COMPONENT_Cb, 0 ), pcDstResi->getAddr( COMPONENT_Cr, 0 )};
    iWidth  = pcCU->getWidth( 0 ) ;
    iHeight = iWidth              ;
    for (Int iY = 0; iY < iHeight; iY++)
    {
      for (Int iX = 0; iX < iWidth; iX++)
      {
        if(bLL)
        {
          Pel     tmp = pResi[0][iX] - (pResi[1][iX]>>1);
          decYint = pResi[1][iX] + tmp;
          decUint = tmp - (pResi[2][iX] >> 1);
          decVint = pResi[2][iX] + decUint;
        }
        else
        {
          decYint = ( pResi[0][iX] + pResi[1][iX]                + offset) >> BIT_V;
          decUint = ( pResi[0][iX] - pResi[1][iX] - pResi[2][iX] + offset) >> BIT_V;
          decVint = ( pResi[0][iX] - pResi[1][iX] + pResi[2][iX] + offset) >> BIT_V;
        }

        pDstResi[0][iX] = decYint;
        pDstResi[1][iX] = decUint;
        pDstResi[2][iX] = decVint;
      }
      pResi[0] += uiStride;
      pResi[1] += uiStride;
      pResi[2] += uiStride;
      pDstResi[0] += uiStride;
      pDstResi[1] += uiStride;
      pDstResi[2] += uiStride;
    }
  }
  else
  {
    PartSize ePartSize = pcCU->getPartitionSize(iAbsPartIdx);
    Int iHorTh = pcCU->getWidth(0);
    Int iVerTh = pcCU->getHeight(0);
    switch(ePartSize) {
    case SIZE_2NxN:
      iVerTh = pcCU->getHeight(0)>>1;
      break;
    case SIZE_Nx2N:
      iHorTh = pcCU->getWidth(0)>>1;
      break;
    case SIZE_NxN:
      iVerTh = pcCU->getHeight(0)>>1;
      iHorTh = pcCU->getWidth(0)>>1;
      break;
    case SIZE_2NxnU:
      iVerTh = pcCU->getHeight(0)>>2;
      break;
    case SIZE_2NxnD:
      iVerTh = (pcCU->getHeight(0)*3)>>2;
      break;
    case SIZE_nLx2N:
      iHorTh = pcCU->getWidth(0)>>2;
      break;
    case SIZE_nRx2N:
      iHorTh = (pcCU->getWidth(0)*3)>>2;
      break;
    case SIZE_2Nx2N:
    default:
      break;
    }
    UInt uiHorPos = (g_auiZscanToRaster[ iAbsPartIdx ]%pcCU->getPic()->getNumPartInWidth()) * pcCU->getPic()->getMinCUWidth();
    UInt uiVerPos = (g_auiZscanToRaster[ iAbsPartIdx ]/pcCU->getPic()->getNumPartInWidth()) * pcCU->getPic()->getMinCUHeight();

    iHorTh -= uiHorPos;
    iVerTh -= uiVerPos;

    Pel* pResi[3] = {pcResi->getAddr( COMPONENT_Y, iAbsPartIdx ), pcResi->getAddr( COMPONENT_Cb, iAbsPartIdx ), pcResi->getAddr( COMPONENT_Cr, iAbsPartIdx )};
    Pel* pDstResi[3] = {pcDstResi->getAddr( COMPONENT_Y, iAbsPartIdx ), pcDstResi->getAddr( COMPONENT_Cb, iAbsPartIdx ), pcDstResi->getAddr( COMPONENT_Cr, iAbsPartIdx )};
    for (Int iY = 0; iY < iHeight; iY++)
    {
      for (Int iX = 0; iX < iWidth; iX++)
      {
        if(bLL)
        {
          Pel     tmp = pResi[0][iX] - (pResi[1][iX]>>1);
          decYint = pResi[1][iX] + tmp;
          decUint = tmp - (pResi[2][iX] >> 1);
          decVint = pResi[2][iX] + decUint;
        }
        else
        {
          decYint = ( pResi[0][iX] + pResi[1][iX]                + offset) >> BIT_V ;
          decUint = ( pResi[0][iX] - pResi[1][iX] - pResi[2][iX] + offset) >> BIT_V  ;
          decVint = ( pResi[0][iX] - pResi[1][iX] + pResi[2][iX] + offset) >> BIT_V  ;
        }

        pDstResi[0][iX] = decYint;
        pDstResi[1][iX] = decUint;
        pDstResi[2][iX] = decVint;
      }
      pResi[0] += uiStride;
      pResi[1] += uiStride;
      pResi[2] += uiStride;
      pDstResi[0] += uiStride;
      pDstResi[1] += uiStride;
      pDstResi[2] += uiStride;
    }

  }
}

Void TComAdaptiveColorSpace::getNoCorrImage(TComDataCU* pcCU, TComYuv* pcOrigYuv, TComYuv* pcNoCorrYuv, Pel NeighbourPixel[][MAX_CU_SIZE*4+1])
{
  UInt uiPartTmp;
  Int aboveAvailable = (pcCU->getPUAbove( uiPartTmp, pcCU->getZorderIdxInCU(), true, false ) != NULL) ? 2 : 0;
  Int leftAvailable = (pcCU->getPULeft( uiPartTmp, pcCU->getZorderIdxInCU(), true, false ) != NULL) ? 2 : 0;
  UInt uiZOrder          = pcCU->getZorderIdxInCU();
  Pel* pRecIPred[3]     = { pcCU->getPic()->getPicYuvRec()->getAddr( COMPONENT_Y, pcCU->getAddr(), uiZOrder ), pcCU->getPic()->getPicYuvRec()->getAddr( COMPONENT_Cb, pcCU->getAddr(), uiZOrder ),
    pcCU->getPic()->getPicYuvRec()->getAddr( COMPONENT_Cr, pcCU->getAddr(), uiZOrder )};
  UInt uiRecIPredStride  = pcCU->getPic()->getPicYuvRec()->getStride  (COMPONENT_Y);
  UInt uiRecIPredStrideC = pcCU->getPic()->getPicYuvRec()->getStride  (COMPONENT_Cb);
  const ChromaFormat chFmt = pcCU->getPic()->getChromaFormat();
  Int iWidth = pcCU->getWidth(0);
  Int iHeight = pcCU->getHeight(0);
  Int iWidthC = (chFmt == CHROMA_444) ? iWidth : iWidth>>1;
  Int iHeightC = (chFmt != CHROMA_420) ? iHeight : iHeight>>1;
  if (aboveAvailable)
  {
    for (Int iX = (leftAvailable ? 0 : 1); iX < iWidth*2+1; iX++)
    {
      NeighbourPixel[0][iHeight*2+iX] = pRecIPred[0][-static_cast<Int>(uiRecIPredStride)-1+iX];
    }
    for (Int iX = (leftAvailable ? 0 : 1); iX < iWidthC*2+1; iX++)
    {
      NeighbourPixel[1][iHeightC*2+iX] = pRecIPred[1][-static_cast<Int>(uiRecIPredStrideC)-1+iX];
      NeighbourPixel[2][iHeightC*2+iX] = pRecIPred[2][-static_cast<Int>(uiRecIPredStrideC)-1+iX];
    }
  }
  if (leftAvailable)
  {
    for (Int iY = 0; iY < iHeight*2; iY++)
    {
      NeighbourPixel[0][iHeight*2-iY-1] = pRecIPred[0][iY*static_cast<Int>(uiRecIPredStride)-1];
    }
    for (Int iY = 0; iY < iHeightC*2; iY++)
    {
      NeighbourPixel[1][iHeightC*2-iY-1] = pRecIPred[1][iY*static_cast<Int>(uiRecIPredStrideC)-1];
      NeighbourPixel[2][iHeightC*2-iY-1] = pRecIPred[2][iY*static_cast<Int>(uiRecIPredStrideC)-1];
    }
  }
  UInt uiLog2TrSize = g_aucConvertToBit[pcCU->getSlice()->getSPS()->getMaxCUWidth() >> pcCU->getDepth(0)]+2;
  getNoCorrImageIntra(pcCU, pcOrigYuv, pcNoCorrYuv, 0, uiLog2TrSize, leftAvailable, aboveAvailable);
}

Void TComAdaptiveColorSpace::reconstructPixel(TComDataCU* pcCU, TComYuv* pcRecoYuvBest, Pel NeighbourPixel[][MAX_CU_SIZE*4+1])
{
  UInt uiPartTmp;
  Int aboveAvailable = (pcCU->getPUAbove( uiPartTmp, pcCU->getZorderIdxInCU(), true, false ) != NULL) ? 2 : 0;
  Int leftAvailable = (pcCU->getPULeft( uiPartTmp, pcCU->getZorderIdxInCU(), true, false ) != NULL) ? 2 : 0;
  Int iWidth  = pcCU->getWidth(0);
  Int iHeight = pcCU->getHeight(0);

  UInt uiZOrder          = pcCU->getZorderIdxInCU();
  Pel* pRecIPred[3]     = { pcCU->getPic()->getPicYuvRec()->getAddr( COMPONENT_Y, pcCU->getAddr(), uiZOrder ), pcCU->getPic()->getPicYuvRec()->getAddr( COMPONENT_Cb, pcCU->getAddr(), uiZOrder ),
    pcCU->getPic()->getPicYuvRec()->getAddr( COMPONENT_Cr, pcCU->getAddr(), uiZOrder )};
  UInt uiRecIPredStride  = pcCU->getPic()->getPicYuvRec()->getStride  (COMPONENT_Y);
  UInt uiRecIPredStrideC = pcCU->getPic()->getPicYuvRec()->getStride  (COMPONENT_Cb);
  if (aboveAvailable)
  {
    for (Int iX = (leftAvailable ? 0 : 1); iX < iWidth*2+1; iX++)
    {
      pRecIPred[0][-static_cast<Int>(uiRecIPredStride)-1+iX] = NeighbourPixel[0][iHeight*2+iX];
    }
    for (Int iX = (leftAvailable ? 0 : 1); iX < iWidth*2+1; iX++)
    {
      pRecIPred[1][-static_cast<Int>(uiRecIPredStrideC)-1+iX] = NeighbourPixel[1][iHeight*2+iX];
      pRecIPred[2][-static_cast<Int>(uiRecIPredStrideC)-1+iX] = NeighbourPixel[2][iHeight*2+iX];
    }
  }
  if (leftAvailable)
  {
    for (Int iY = 0; iY < iHeight*2; iY++)
    {
      pRecIPred[0][iY*static_cast<Int>(uiRecIPredStride)-1] = NeighbourPixel[0][iHeight*2-iY-1];
    }
    for (Int iY = 0; iY < iHeight*2; iY++)
    {
      pRecIPred[1][iY*static_cast<Int>(uiRecIPredStrideC)-1] = NeighbourPixel[1][iHeight*2-iY-1];
      pRecIPred[2][iY*static_cast<Int>(uiRecIPredStrideC)-1] = NeighbourPixel[2][iHeight*2-iY-1];
    }
  }
  if (pcRecoYuvBest)
  {
    Pel* pReco[3] = {pcRecoYuvBest->getAddr(COMPONENT_Y), pcRecoYuvBest->getAddr(COMPONENT_Cb), pcRecoYuvBest->getAddr(COMPONENT_Cr)};
    UInt uiRecoStride = pcRecoYuvBest->getStride(COMPONENT_Y);

    Bool bLL     = pcCU->getSlice()->getSPS()->getUseLossless();
    Pel decYint, decUint, decVint;
    Int uvoffset = (1<<(g_bitDepth[CHANNEL_TYPE_LUMA] + 2 - 1));
    const Int offset = (1 << (BIT_V - 1)); 
    if( bLL )
      uvoffset = uvoffset >> 1;

    for (Int iY = 0; iY < iHeight; iY++)
    {
      for (Int iX = 0; iX < iWidth; iX++)
      {
        if(bLL)
        {
          Pel     tmp = pReco[0][iX] - ((pReco[1][iX]-uvoffset)>>1);
          decYint = (pReco[1][iX] -uvoffset) + tmp;
          decUint = tmp - ( (pReco[2][iX]-uvoffset) >> 1);
          decVint = (pReco[2][iX]-uvoffset) + decUint;
        }
        else
        {
          decYint = ( pReco[0][iX] + (pReco[1][iX] - uvoffset)                             + offset) >> BIT_V;
          decUint = ( pReco[0][iX] - (pReco[1][iX] - uvoffset) - (pReco[2][iX] - uvoffset) + offset) >> BIT_V;
          decVint = ( pReco[0][iX] - (pReco[1][iX] - uvoffset) + (pReco[2][iX] - uvoffset) + offset) >> BIT_V;
        }

        pReco[0][iX] = Clip( decYint, CHANNEL_TYPE_LUMA );
        pReco[1][iX] = Clip( decUint, CHANNEL_TYPE_CHROMA );
        pReco[2][iX] = Clip( decVint, CHANNEL_TYPE_CHROMA );
      }
      pReco[0] += uiRecoStride;
      pReco[1] += uiRecoStride;
      pReco[2] += uiRecoStride;
    }
  }
}

#endif
