/** \file     CNNFilter.cpp
    \brief    convolutional neural network-based filter class
*/
#include "CNNFilter.h"
#if CNN_FILTERING
#include "UnitTools.h"
#include "UnitPartitioner.h"
#include "CodingStructure.h"
#include "CommonLib/dtrace_codingstruct.h"
#include "CommonLib/dtrace_buffer.h"
#include <fstream>
#include <cmath>
#include <chrono>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <ctime>
#include <algorithm>
#include <numeric>
#include <sadl/model.h>
using namespace std;
  

//! \ingroup CommonLib
//! \{

CNNFilter::CNNFilter()
{
}
void CNNFilter::create( const int picWidth, const int picHeight, const ChromaFormat format, const int cnnlfNumParams)
{
  if (m_tempBuf.size() > 0)
    return;
  m_tempBuf.resize(cnnlfNumParams);
  for (int i = 0; i < cnnlfNumParams; i ++)
  {
    m_tempBuf[i].create(format, Area(0, 0, picWidth, picHeight));
  }

#if SCALE_NN_RESIDUE
  m_tempScaledBuf.resize(cnnlfNumParams);
  for (int i = 0; i < cnnlfNumParams; i ++)
  {
    m_tempScaledBuf[i].create(format, Area(0, 0, picWidth, picHeight));
  }
#endif
}
void CNNFilter::init(std::string interLuma, std::string interChroma, std::string intraLuma, std::string intraChroma)
{
  m_interLuma = interLuma;
  m_interChroma = interChroma;
  m_intraLuma = intraLuma;
  m_intraChroma = intraChroma;
}
void CNNFilter::destroy()
{
  for (int i = 0; i < m_tempBuf.size(); i++)
  {
    m_tempBuf[i].destroy();
  }

#if SCALE_NN_RESIDUE
  for (int i = 0; i < m_tempScaledBuf.size(); i++)
  {
    m_tempScaledBuf[i].destroy();
  }
#endif
}

template<typename T>
struct ModelData {
  sadl::Model<T> model;
  vector<sadl::Tensor<T>> inputs;
  int hor,ver;
  bool luma,inter;
};

template<typename T>
static std::vector<ModelData<T>> initSpace() {
  std::vector<ModelData<T>> v;
  v.reserve(5);
  return v;
}

#if NN_FIXED_POINT_IMPLEMENTATION
static std::vector<ModelData<int16_t>> models=initSpace<int16_t>();
#else
static std::vector<ModelData<float>> models=initSpace<float>();
#endif

template<typename T>
static ModelData<T> &getModel(int ver, int hor, bool luma, bool inter, const string modelName) {
  ModelData<T> *ptr = nullptr;
  for(auto &m: models)
  {
    if (m.luma == luma && m.inter == inter)
    {
      ptr = &m;
      break;
    }
  }
  if (ptr == nullptr)
  {
    if (models.size() == models.capacity())
    {
      std::cout << "[ERROR] increase cache" << std::endl;
      exit(-1);
    }
    models.resize(models.size()+1);
    ptr = &models.back();
    ModelData<T> &m = *ptr;
    if  (luma && inter)
    {
      cout << "[INFO] luma inter model loading" << endl;
    }
    else if (luma && !inter)
    {
      cout << "[INFO] luma intra model loading" << endl;
    }
    else if (!luma && inter)
    {
      cout << "[INFO] chroma inter model loading" << endl;
    }
    else if (!luma && !inter)
    {
      cout << "[INFO] chroma intra model loading" << endl;
    }
    ifstream file(modelName, ios::binary);
    m.model.load(file);
    m.luma = luma;
    m.inter = inter;
    m.ver = 0;
    m.hor = 0;
  }
  ModelData<T> &m = *ptr;
  if (m.ver != ver || m.hor != hor)
  {
    m.inputs = m.model.getInputsTemplate();
    if (luma)
    {
      for(auto &t: m.inputs)
      {
        sadl::Dimensions dims(std::initializer_list<int>({1, ver, hor, 1}));
        t.resize(dims);
      }
    }
    else if (inter) // inter chroma
    {
      int inputId = 0;
      for(auto &t: m.inputs)
      {
        if (t.dims()[3] == 1 && inputId != 4) // luma
        {
          sadl::Dimensions dims(std::initializer_list<int>({1, ver, hor, 1}));
          t.resize(dims);
        }
        else if (t.dims()[3] == 1 && inputId == 4) // qp
        {
          sadl::Dimensions dims(std::initializer_list<int>({1, ver/2, hor/2, 1}));
          t.resize(dims);
        }
        else
        {
          sadl::Dimensions dims(std::initializer_list<int>({1, ver/2, hor/2, 2}));
          t.resize(dims);
        }
        inputId ++;
      }
    }
    else // intra chroma
    {
      int inputId = 0;
      for(auto &t: m.inputs)
      {
        if (t.dims()[3] == 1 && inputId != 5 ) // luma
        {
          sadl::Dimensions dims(std::initializer_list<int>({1, ver, hor, 1}));
          t.resize(dims);
        }
        else if (t.dims()[3] == 1 && inputId == 5) // QP
        {
          sadl::Dimensions dims(std::initializer_list<int>({1, ver/2, hor/2, 1}));
          t.resize(dims);
        }
        else
        {
          sadl::Dimensions dims(std::initializer_list<int>({1, ver/2, hor/2, 2}));
          t.resize(dims);
        }
        inputId ++;
      }
    }
    if  (luma&&inter)        cout << "[INFO] luma inter model initilization" << endl;
    else if (luma&&!inter)   cout << "[INFO] luma intra model initilization" << endl;
    else if (!luma&&inter)   cout << "[INFO] chroma inter model initilization" << endl;
    else if (!luma&&!inter)  cout << "[INFO] chroma intra model initilization" << endl;
    if (!m.model.init(m.inputs))
    {
      cerr << "[ERROR] issue during initialization" << endl;
      exit(-1);
    }
    m.ver = ver;
    m.hor = hor;
  }
  return m;
}

template<typename T>
void prepareInputsLuma (Picture* pic, UnitArea inferArea, vector<sadl::Tensor<T>> &inputs, int qp, bool inter)
{
  double inputScale = 1024;
#if NN_FIXED_POINT_IMPLEMENTATION
  int shiftInput = NN_INPUT_PRECISION;
#else
  int shiftInput = 0;
#endif
#if FUSE_NN_AND_LF
  PelBuf bufRec = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Y);
#else
  PelBuf bufRec = pic->getRecoBuf(inferArea).get(COMPONENT_Y);
#endif
  PelBuf bufPred = pic->getPredBufCustom(inferArea).get(COMPONENT_Y);
  PelBuf bufPartition = pic->getPartitionBuf(inferArea).get(COMPONENT_Y);
  PelBuf bufBsMap = pic->getBsMapBuf(inferArea).get(COMPONENT_Y);
  
  sadl::Tensor<T>* inputRec, *inputPred, *inputPartition, *inputBs, *inputQp;
  if (inter)
  {
    inputRec = &inputs[0];
    inputPred = &inputs[1];
    inputBs = &inputs[2];
    inputQp = &inputs[3];
  }
  else
  {
    inputRec = &inputs[0];
    inputPred = &inputs[1];
    inputPartition = &inputs[2];
    inputBs = &inputs[3];
    inputQp = &inputs[4];
  }
  
  int hor = inferArea.lwidth();
  int ver = inferArea.lheight();
  
  for (int yy = 0; yy < ver; yy++)
  {
    for (int xx = 0; xx < hor; xx++)
    {
      (*inputRec)(0, yy, xx, 0) = bufRec.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputPred)(0, yy, xx, 0) = bufPred.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputBs)(0, yy, xx, 0) = bufBsMap.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputQp)(0, yy, xx, 0) = qp / 64.0 * (1 << shiftInput);
      if (!inter)
      {
        (*inputPartition)(0, yy, xx, 0) = bufPartition.at(xx, yy) / inputScale * (1 << shiftInput);
      }
    }
  }
}

template<typename T>
void prepareInputsChroma (Picture* pic, UnitArea inferArea, vector<sadl::Tensor<T>> &inputs, int qp, bool inter)
{
  double inputScale = 1024;
#if NN_FIXED_POINT_IMPLEMENTATION
  int shiftInput = NN_INPUT_PRECISION;
#else
  int shiftInput = 0;
#endif
#if FUSE_NN_AND_LF
  PelBuf bufRecY = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Y);
  PelBuf bufRecCb = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufRecCr = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Cr);
#else
  PelBuf bufRecY = pic->getRecoBuf(inferArea).get(COMPONENT_Y);
  PelBuf bufRecCb = pic->getRecoBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufRecCr = pic->getRecoBuf(inferArea).get(COMPONENT_Cr);
#endif
  
  PelBuf bufPredCb = pic->getPredBufCustom(inferArea).get(COMPONENT_Cb);
  PelBuf bufPredCr = pic->getPredBufCustom(inferArea).get(COMPONENT_Cr);
  
  PelBuf bufPartitionCb = pic->getPartitionBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufPartitionCr = pic->getPartitionBuf(inferArea).get(COMPONENT_Cr);
  
  PelBuf bufBsMapCb = pic->getBsMapBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufBsMapCr = pic->getBsMapBuf(inferArea).get(COMPONENT_Cr);
  
  sadl::Tensor<T>* inputRecCrossComponent, *inputRec, *inputPred, *inputPartition, *inputBs, *inputQp;
  
  if (inter)
  {
    inputRecCrossComponent = &inputs[0];
    inputRec = &inputs[1];
    inputPred = &inputs[2];
    inputBs = &inputs[3];
    inputQp = &inputs[4];
  }
  else
  {
    inputRecCrossComponent = &inputs[0];
    inputRec = &inputs[1];
    inputPred = &inputs[2];
    inputPartition = &inputs[3];
    inputBs = &inputs[4];
    inputQp = &inputs[5];
  }
  
  int hor = inferArea.lwidth();
  int ver = inferArea.lheight();
  
  int horC = inferArea.lwidth() >> 1;
  int verC = inferArea.lheight() >> 1;
  
  for (int yy = 0; yy < ver; yy++)
  {
    for (int xx = 0; xx < hor; xx++)
    {
      (*inputRecCrossComponent)(0, yy, xx, 0) = bufRecY.at(xx, yy) / inputScale * (1 << shiftInput);
    }
  }
  
  for (int yy = 0; yy < verC; yy++)
  {
    for (int xx = 0; xx < horC; xx++)
    {

      (*inputRec)(0, yy, xx, 0) = bufRecCb.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputRec)(0, yy, xx, 1) = bufRecCr.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputPred)(0, yy, xx, 0) = bufPredCb.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputPred)(0, yy, xx, 1) = bufPredCr.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputBs)(0, yy, xx, 0) = bufBsMapCb.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputBs)(0, yy, xx, 1) = bufBsMapCr.at(xx, yy) / inputScale * (1 << shiftInput);
      (*inputQp)(0, yy, xx, 0) = qp / 64.0 * (1 << shiftInput);
      if (!inter)
      {
        (*inputPartition)(0, yy, xx, 0) = bufPartitionCb.at(xx, yy) / inputScale * (1 << shiftInput);
        (*inputPartition)(0, yy, xx, 1) = bufPartitionCr.at(xx, yy) / inputScale * (1 << shiftInput);
      }
    }
  }
}

template<typename T>
void extractOutputsLuma (Picture* pic, sadl::Model<T> &m, PelStorage& tempBuf, PelStorage& tempScaledBuf, UnitArea inferArea, int extLeft, int extRight, int extTop, int extBottom, bool inter)
{
#if NN_FIXED_POINT_IMPLEMENTATION
  int log2InputScale = 10;
  int log2OutputScale = 10;
  int shiftInput = NN_OUPUTPUT_PRECISION - log2InputScale;
  int shiftOutput = NN_OUPUTPUT_PRECISION - log2OutputScale;
  int offset    = (1 << shiftOutput) / 2;
#else
  double inputScale = 1024;
  double outputScale = 1024;
#endif
  auto output = m.result(0);
  PelBuf bufDst = tempBuf.getBuf(inferArea).get(COMPONENT_Y);

#if SCALE_NN_RESIDUE
  PelBuf bufScaledDst = tempScaledBuf.getBuf(inferArea).get(COMPONENT_Y);
#endif
  
  int hor = inferArea.lwidth();
  int ver = inferArea.lheight();
  
#if FUSE_NN_AND_LF
  PelBuf bufRec = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Y);
#else
  PelBuf bufRec = pic->getRecoBuf(inferArea).get(COMPONENT_Y);
#endif
  
  for (int c = 0; c < 4; c++) // output includes 4 sub images
  {
    for (int y = 0; y < ver >> 1; y++)
    {
      for (int x = 0; x < hor >> 1; x++)
      {
        int yy = (y << 1) + c / 2;
        int xx = (x << 1) + c % 2;
        if (xx < extLeft || yy < extTop || xx >= hor - extRight || yy >= ver - extBottom)
        {
          continue;
        }
#if NN_FIXED_POINT_IMPLEMENTATION
        int out = ( output(0, y, x, c) + (bufRec.at(xx, yy) << shiftInput) + offset) >> shiftOutput;
#else
        int out = ( output(0, y, x, c) + bufRec.at(xx, yy) / inputScale ) * outputScale + 0.5;
#endif
        bufDst.at(xx, yy) = Pel(Clip3<int>( 0, 1023, out));

  #if SCALE_NN_RESIDUE
        bufScaledDst.at(xx, yy) = Pel(Clip3<int>(0, 1023 << NN_RESIDUE_ADDITIONAL_SHIFT, out * (1 << NN_RESIDUE_ADDITIONAL_SHIFT) ) );
  #endif
      }
    }
  }
}

template<typename T>
void extractOutputsChroma (Picture* pic, sadl::Model<T> &m, PelStorage& tempBuf, PelStorage& tempScaledBuf, UnitArea inferArea, int extLeft, int extRight, int extTop, int extBottom, bool inter)
{
#if NN_FIXED_POINT_IMPLEMENTATION
int log2InputScale = 10;
  int log2OutputScale = 10;
  int shiftInput = NN_OUPUTPUT_PRECISION - log2InputScale;
  int shiftOutput = NN_OUPUTPUT_PRECISION - log2OutputScale;
  int offset    = (1 << shiftOutput) / 2;
#else
  double inputScale = 1024;
  double outputScale = 1024;
#endif
  auto output = m.result(0);
  PelBuf bufDstCb = tempBuf.getBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufDstCr = tempBuf.getBuf(inferArea).get(COMPONENT_Cr);

#if SCALE_NN_RESIDUE
  PelBuf bufScaledDstCb = tempScaledBuf.getBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufScaledDstCr = tempScaledBuf.getBuf(inferArea).get(COMPONENT_Cr);
#endif
  
  int hor = inferArea.lwidth() >> 1;
  int ver = inferArea.lheight() >> 1;
  
#if FUSE_NN_AND_LF
  PelBuf bufRecCb = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufRecCr = pic->getUnfilteredRecBuf(inferArea).get(COMPONENT_Cr);
#else
  PelBuf bufRecCb = pic->getRecoBuf(inferArea).get(COMPONENT_Cb);
  PelBuf bufRecCr = pic->getRecoBuf(inferArea).get(COMPONENT_Cr);
#endif

  for (int c = 0; c < 4; c++) // output includes 4 sub images
  {
    for (int y = 0; y < ver >> 1; y++)
    {
      for (int x = 0; x < hor >> 1; x++)
      {
        int yy = (y << 1) + c / 2;
        int xx = (x << 1) + c % 2;
        if (xx < extLeft || yy < extTop || xx >= hor - extRight || yy >= ver - extBottom)
        {
          continue;
        }
        
#if NN_FIXED_POINT_IMPLEMENTATION
        int outCb = ( output(0, y, x, c) + (bufRecCb.at(xx, yy) << shiftInput) + offset) >> shiftOutput;
        int outCr = ( output(0, y, x, c+4) + (bufRecCr.at(xx, yy) << shiftInput) + offset) >> shiftOutput;
#else
        int outCb = ( output(0, y, x, c) + bufRecCb.at(xx, yy) / inputScale ) * outputScale + 0.5;
        int outCr = ( output(0, y, x, c+4) + bufRecCr.at(xx, yy) / inputScale ) * outputScale + 0.5;
#endif
        bufDstCb.at(xx, yy) = Pel(Clip3<int>( 0, 1023, outCb ) );
        bufDstCr.at(xx, yy) = Pel(Clip3<int>( 0, 1023, outCr ) );

  #if SCALE_NN_RESIDUE
        bufScaledDstCb.at(xx, yy) = Pel(Clip3<int>(0, 1023 << NN_RESIDUE_ADDITIONAL_SHIFT, outCb * (1 << NN_RESIDUE_ADDITIONAL_SHIFT ) ) );
        bufScaledDstCr.at(xx, yy) = Pel(Clip3<int>(0, 1023 << NN_RESIDUE_ADDITIONAL_SHIFT, outCr * (1 << NN_RESIDUE_ADDITIONAL_SHIFT ) ) );
  #endif
      }
    }
  }
}

template<typename T>
void CNNFilter::cnnFilterLumaBlock(Picture* pic, UnitArea inferArea, int extLeft, int extRight, int extTop, int extBottom, int paramIdx, bool inter)
{

  //at::init_num_threads(); // use all available threads
  
  const int border_to_skip = 0;
  if (border_to_skip>0) sadl::Tensor<float>::skip_border = true;
  
  // get model
  ModelData<T> &m = getModel<T>(inferArea.lheight(), inferArea.lwidth(), true, inter, inter ? m_interLuma : m_intraLuma);
  sadl::Model<T> &model = m.model;
  
  // get inputs
  vector<sadl::Tensor<T>> &inputs = m.inputs;

  int seqQp = pic->slices[0]->getPPS()->getPicInitQPMinus26() + 26;
  int sliceQp = pic->slices[0]->getSliceQp();
  int delta = inter ? paramIdx * 5 : paramIdx * 2;
  if ( pic->slices[0]->getTLayer() >= 4 && paramIdx >= 2 )
  {
    delta = 5 - delta;
  }
  int qp = inter ? seqQp - delta : sliceQp - delta;
  
  prepareInputsLuma<T>(pic, inferArea, inputs, qp, inter);
  
  // inference
  chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
  if (!model.apply(inputs))
  {
    cerr << "[ERROR] issue during luma model inference" << endl;
    exit(-1);
  }
  chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
  chrono::duration<double> dt = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
  cout << "[INFO] luma model takes "<< dt.count() * 1000. << " ms" <<endl;

  if (border_to_skip) cout<<"[INFO] discard border size="<<model.result().border_skip<<endl;
  
  // get outputs
  extractOutputsLuma(pic, model, m_tempBuf[paramIdx], m_tempScaledBuf[paramIdx], inferArea, extLeft, extRight, extTop, extBottom, inter);

}

template<typename T>
void CNNFilter::cnnFilterChromaBlock(Picture* pic, UnitArea inferArea, int extLeft, int extRight, int extTop, int extBottom, int paramIdx, bool inter)
{

  //at::init_num_threads();
  
  const int border_to_skip = 0;
  if (border_to_skip>0) sadl::Tensor<float>::skip_border = true;
  
  // get model
  ModelData<T> &m = getModel<T>(inferArea.lheight(), inferArea.lwidth(), false, inter, inter ? m_interChroma : m_intraChroma);
  sadl::Model<T> &model = m.model;
  
  // get inputs
  vector<sadl::Tensor<T>> &inputs = m.inputs;
  
  int seqQp = pic->slices[0]->getPPS()->getPicInitQPMinus26() + 26;
  int sliceQp = pic->slices[0]->getSliceQp();
  int delta = inter ? paramIdx * 5 : paramIdx * 2;
  if ( pic->slices[0]->getTLayer() >= 4 && paramIdx >= 2 )
  {
    delta = 5 - delta;
  }
  int qp = inter ? seqQp - delta : sliceQp - delta;
  
  prepareInputsChroma<T>(pic, inferArea, inputs, qp, inter);
  
  // inference
  chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
  if (!model.apply(inputs))
  {
    cerr << "[ERROR] issue during chroma model inference" << endl;
    exit(-1);
  }
  chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
  chrono::duration<double> dt = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
  cout << "[INFO] chroma model takes " << dt.count() * 1000. << " ms"<<endl;

  if (border_to_skip) cout << "[INFO] discard border size=" << model.result().border_skip << endl;
  
  // get outputs
  extractOutputsChroma(pic, model, m_tempBuf[paramIdx], m_tempScaledBuf[paramIdx], inferArea, extLeft, extRight, extTop, extBottom, inter);

}

void CNNFilter::cnnFilter(Picture* pic)
{
  CodingStructure&  cs            = *pic->cs;
  Slice* pcSlice                  = cs.slice;
  const PreCalcValues& pcv        = *cs.pcv;
  int extension = cs.sps->getCnnlfInferSizeExtension();
  for (int chal = 0; chal < MAX_NUM_CHANNEL_TYPE; chal ++)
  {
    const ChannelType chType = ChannelType( chal );
    CnnlfInferGranularity cnnlfInferGranularity = pcSlice->getCnnlfInferGranularity(chType);
    std::vector<uint8_t>* cnnlfParamIdx = pic->getCnnlfParamIdx(cnnlfInferGranularity);
    int blockSize = cs.sps->getCnnlfInferSize (cnnlfInferGranularity);
    for( int blockIdx = 0; blockIdx < pcv.sizeInCnnlfInferSize[cnnlfInferGranularity]; blockIdx++ )
    {
      int xPosInBlocks        = blockIdx % pcv.widthInCnnlfInferSize[cnnlfInferGranularity];
      int yPosInBlocks        = blockIdx / pcv.widthInCnnlfInferSize[cnnlfInferGranularity];
      int xPos = xPosInBlocks * blockSize;
      int yPos = yPosInBlocks * blockSize;
      int width = (xPos + blockSize > pcv.lumaWidth) ? (pcv.lumaWidth - xPos) : blockSize;
      int height = (yPos + blockSize > pcv.lumaHeight) ? (pcv.lumaHeight - yPos) : blockSize;
      
      const UnitArea block( cs.area.chromaFormat, Area( xPos, yPos, width, height ) );
      
      int extLeft = xPos > 0 ? extension : 0;
      int extRight = (xPos + width + extension > pcv.lumaWidth) ? (pcv.lumaWidth - xPos - width) : extension;
      int extTop = yPos > 0 ? extension : 0;
      int extBottom = (yPos + height + extension > pcv.lumaHeight) ? (pcv.lumaHeight - yPos - height) : extension;
      
      int extXPos = xPos - extLeft;
      int extYPos = yPos - extTop;
      int extWidth = width + extLeft + extRight;
      int extHeight =  height + extTop + extBottom;
      const UnitArea extBlock( cs.area.chromaFormat, Area( extXPos, extYPos, extWidth, extHeight ) );
      if (cnnlfParamIdx[chal][blockIdx] && chal == 0)
      {
#if NN_FIXED_POINT_IMPLEMENTATION
        cnnFilterLumaBlock<int16_t>(pic, extBlock, extLeft, extRight, extTop, extBottom, cnnlfParamIdx[chal][blockIdx]-1, pcSlice->getSliceType() != I_SLICE);
#else
        cnnFilterLumaBlock<float>(pic, extBlock, extLeft, extRight, extTop, extBottom, cnnlfParamIdx[chal][blockIdx]-1, pcSlice->getSliceType() != I_SLICE);
#endif
      }
      if (cnnlfParamIdx[chal][blockIdx] && chal > 0)
      {
#if NN_FIXED_POINT_IMPLEMENTATION
        cnnFilterChromaBlock<int16_t>(pic, extBlock, extLeft >> 1, extRight >> 1, extTop >> 1, extBottom >> 1, cnnlfParamIdx[chal][blockIdx]-1, pcSlice->getSliceType() != I_SLICE);
#else
        cnnFilterChromaBlock<float>(pic, extBlock, extLeft >> 1, extRight >> 1, extTop >> 1, extBottom >> 1, cnnlfParamIdx[chal][blockIdx]-1, pcSlice->getSliceType() != I_SLICE);
#endif
      }

#if SCALE_NN_RESIDUE
      if (cnnlfParamIdx[chal][blockIdx] && pcSlice->getNnScaleFlag(cnnlfParamIdx[chal][blockIdx] - 1, chType))
      {
        if (chType == CHANNEL_TYPE_LUMA)
        {
          scaleResidualBlock(pic, block, cnnlfParamIdx[chal][blockIdx] - 1, COMPONENT_Y);
        }
        else
        {
          scaleResidualBlock(pic, block, cnnlfParamIdx[chal][blockIdx] - 1, COMPONENT_Cb);
          scaleResidualBlock(pic, block, cnnlfParamIdx[chal][blockIdx] - 1, COMPONENT_Cr);
        }
      }
#endif
    }
  }
  for (int comp = 0; comp < MAX_NUM_COMPONENT; comp ++)
  {
    const ComponentID compID = ComponentID( comp );
    int chal = toChannelType(compID);
    CnnlfInferGranularity cnnlfInferGranularity = pcSlice->getCnnlfInferGranularity(toChannelType(compID));
    std::vector<uint8_t>* cnnlfParamIdx = pic->getCnnlfParamIdx(cnnlfInferGranularity);
    int blockSize = cs.sps->getCnnlfInferSize (cnnlfInferGranularity);
    
    for( int blockIdx = 0; blockIdx < pcv.sizeInCnnlfInferSize[cnnlfInferGranularity]; blockIdx++ )
    {
      int xPosInBlocks        = blockIdx % pcv.widthInCnnlfInferSize[cnnlfInferGranularity];
      int yPosInBlocks        = blockIdx / pcv.widthInCnnlfInferSize[cnnlfInferGranularity];
      int xPos = xPosInBlocks * blockSize;
      int yPos = yPosInBlocks * blockSize;
      int width = (xPos + blockSize > pcv.lumaWidth) ? (pcv.lumaWidth - xPos) : blockSize;
      int height = (yPos + blockSize > pcv.lumaHeight) ? (pcv.lumaHeight - yPos) : blockSize;
      
      const UnitArea block( cs.area.chromaFormat, Area( xPos, yPos, width, height ) );
      if (cnnlfParamIdx[chal][blockIdx])
      {
#if SCALE_NN_RESIDUE
        if (pcSlice->getNnScaleFlag(cnnlfParamIdx[chal][blockIdx] - 1, toChannelType(compID)))
          pic->getRecoBuf(block).get(compID).copyFrom(m_tempScaledBuf[cnnlfParamIdx[chal][blockIdx] - 1].getBuf(block).get(compID));
        else
#endif
        pic->getRecoBuf(block).get(compID).copyFrom(m_tempBuf[cnnlfParamIdx[chal][blockIdx] - 1].getBuf(block).get(compID));
      }
    }
  }
}

#if SCALE_NN_RESIDUE
void CNNFilter::scaleResidualBlock(Picture *pic, UnitArea inferArea, int paramIdx, ComponentID compID)
{

  CodingStructure &cs         = *pic->cs;
  Slice *          pcSlice    = cs.slice;

  const int scale  = pcSlice->getNnScale(paramIdx, toChannelType(compID));

  const int shift  = NN_RESIDUE_SCALE_SHIFT + NN_RESIDUE_ADDITIONAL_SHIFT;

  const int offset    = (1 << shift) / 2;

  PelBuf bufRec    = pic->getRecoBuf(inferArea).get(compID);

  int    strideRec  = bufRec.stride;
  Pel *  pRec      = bufRec.buf;

  int blockSizeHor = inferArea.lwidth();
  int blockSizeVer = inferArea.lheight();

  if (compID)
  {
    blockSizeHor      = blockSizeHor / 2;
    blockSizeVer      = blockSizeVer / 2;
  }

  int                idxDst        = 0;
  int                idxRec     = 0;
  int                blockSize  = blockSizeVer * blockSizeHor;
  
  PelBuf bufScaledDst = m_tempScaledBuf[paramIdx].getBuf(inferArea).get(compID);
  int    strideDst    = bufScaledDst.stride;
  Pel *  pScaledDst   = bufScaledDst.buf;

  for (int pixelIdx = 0, yy = 0, xx = 0; pixelIdx < blockSize; pixelIdx++)
  {
    xx = pixelIdx % blockSizeHor;
    yy = pixelIdx / blockSizeHor;

    idxDst       = yy * strideDst + xx;
    idxRec       = yy * strideRec + xx;
    pScaledDst[idxDst] = Clip3(0, 1023, pRec[idxRec] + (((pScaledDst[idxDst] - (pRec[idxRec] << NN_RESIDUE_ADDITIONAL_SHIFT)) * scale + offset) >> shift));
  }
}
#endif
#endif

//! \}
