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

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     Contexts.h

 *  \brief    Classes providing probability descriptions and contexts (header)

 */

#pragma once

#include "CommonDef.h"

#include "Slice.h"

#include <vector>

//! \ingroup CommonLib

//! \{

namespace vvenc {

static constexpr int     PROB_BITS   = 15;   // Nominal number of bits to represent probabilities

static constexpr int     PROB_BITS_0 = 10;   // Number of bits to represent 1st estimate

static constexpr int     PROB_BITS_1 = 14;   // Number of bits to represent 2nd estimate

static constexpr int     MASK_0      = ~(~0u << PROB_BITS_0) << (PROB_BITS - PROB_BITS_0);

static constexpr int     MASK_1      = ~(~0u << PROB_BITS_1) << (PROB_BITS - PROB_BITS_1);

static constexpr uint8_t DWS         = 8;   // 0x47 Default window sizes

struct BinFracBits

{

  uint32_t intBits[2];

};

class ProbModelTables

{

protected:

  static const BinFracBits m_binFracBits[256];

  static const uint8_t      m_RenormTable_32  [ 32];          // Std         MP   MPI

};

class BinProbModelBase : public ProbModelTables

{

public:

  static uint32_t estFracBitsEP ()                    { return  (       1 << SCALE_BITS ); }

  static uint32_t estFracBitsEP ( unsigned numBins )  { return  ( numBins << SCALE_BITS ); }

};

class BinProbModel : public BinProbModelBase

{

public:

  BinProbModel()

  {

    uint16_t half = 1 << (PROB_BITS - 1);

    m_state[0]    = half;

    m_state[1]    = half;

    m_rate        = DWS;

  }

public:

  void            init              ( int qp, int initId );

  void update(unsigned bin)

  {

    int rate0 = m_rate >> 4;

    int rate1 = m_rate & 15;

    m_state[0] -= ( m_state[0] >> rate0 ) & MASK_0;

    m_state[1] -= ( m_state[1] >> rate1 ) & MASK_1;

    m_state[0] += ( ( -static_cast< int >( bin ) & 0x7fffu ) >> rate0 ) & MASK_0;

    m_state[1] += ( ( -static_cast< int >( bin ) & 0x7fffu ) >> rate1 ) & MASK_1;

  }

  void setLog2WindowSize(uint8_t log2WindowSize)

  {

    int rate0 = 2 + ((log2WindowSize >> 2) & 3);

    int rate1 = 3 + rate0 + (log2WindowSize&  3);

    m_rate    = 16 * rate0 + rate1;

    CHECK(rate1 > 9, "Second window size is too large!");

  }

  void estFracBitsUpdate(unsigned bin, uint64_t &b)

  {

    b += estFracBits(bin);

    update(bin);

  }

  uint32_t        estFracBits(unsigned bin) const { return getFracBitsArray().intBits[bin]; }

  static uint32_t estFracBitsTrm(unsigned bin) { return (bin ? 0x3bfbb : 0x0010c); }

  const BinFracBits &getFracBitsArray() const { return m_binFracBits[state()]; }

public:

  uint8_t state() const { return (m_state[0] + m_state[1]) >> 8; }

  uint8_t mps() const { return state() >> 7; }

  uint8_t getLPS(unsigned range) const

  {

    uint16_t q = state();

    if (q & 0x80)

      q = q ^ 0xff;

    return ((q >> 2) * (range >> 5) >> 1) + 4;

  }

  static uint8_t  getRenormBitsLPS  ( unsigned LPS )                    { return    m_RenormTable_32  [LPS>>3]; }

  //static uint8_t  getRenormBitsRange( unsigned range )                  { return    1; }

  uint16_t getState() const { return m_state[0] + m_state[1]; }

  void     setState(uint16_t pState)

  {

    m_state[0] = (pState >> 1) & MASK_0;

    m_state[1] = (pState >> 1) & MASK_1;

  }

public:

  uint64_t estFracExcessBits(const BinProbModel &r) const

  {

    int n = 2 * state() + 1;

    return ((512 - n) * r.estFracBits(0) + n * r.estFracBits(1) + 256) >> 9;

  }

private:

  uint16_t m_state[2];

  uint8_t  m_rate;

};

class CtxSet

{

public:

  CtxSet( uint16_t offset, uint16_t size ) : Offset( offset ), Size( size ) {}

  CtxSet( const CtxSet& ctxSet ) : Offset( ctxSet.Offset ), Size( ctxSet.Size ) {}

  CtxSet( std::initializer_list<CtxSet> ctxSets );

public:

  uint16_t  operator()  ()  const

  {

    return Offset;

  }

  uint16_t  operator()  ( uint16_t inc )  const

  {

    CHECKD( inc >= Size, "Specified context increment (" << inc << ") exceed range of context set [0;" << Size - 1 << "]." );

    return Offset + inc;

  }

  bool operator== ( const CtxSet& ctxSet ) const

  {

    return ( Offset == ctxSet.Offset && Size == ctxSet.Size );

  }

  bool operator!= ( const CtxSet& ctxSet ) const

  {

    return ( Offset != ctxSet.Offset || Size != ctxSet.Size );

  }

  uint16_t size() const { return Size; }

public:

  uint16_t  Offset;

  uint16_t  Size;

};

class ContextSetCfg

{

public:

  // context sets: specify offset and size

  static const CtxSet   SplitFlag;

  static const CtxSet   SplitQtFlag;

  static const CtxSet   SplitHvFlag;

  static const CtxSet   Split12Flag;

  static const CtxSet   ModeConsFlag;

  static const CtxSet   SkipFlag;

  // keep these together to

  static const CtxSet   MergeFlag;

  static const CtxSet   RegularMergeFlag;

  static const CtxSet   MergeIdx;

  static const CtxSet   MmvdFlag;

  static const CtxSet   MmvdMergeIdx;

  static const CtxSet   MmvdStepMvpIdx;

  static const CtxSet   SubblockMergeFlag;

  static const CtxSet   AffMergeIdx;

  static const CtxSet   PredMode;

  static const CtxSet   CclmModeFlag;  

  static const CtxSet   CclmModeIdx;

  static const CtxSet   IntraChromaPredMode;

  // keep these for single cpy

  static const CtxSet   IntraLumaMpmFlag;

  static const CtxSet   IntraLumaPlanarFlag;

  static const CtxSet   MultiRefLineIdx;

  static const CtxSet   MipFlag;

  static const CtxSet   ISPMode;

  static const CtxSet   DeltaQP;

  static const CtxSet   InterDir;

  static const CtxSet   RefPic;

  static const CtxSet   AffineFlag;

  static const CtxSet   AffineType;

  static const CtxSet   Mvd;

  static const CtxSet   BDPCMMode;

  static const CtxSet   QtRootCbf;

  static const CtxSet   ACTFlag;

  static const CtxSet   QtCbf           [3];    // [ channel ]

  static const CtxSet   SigCoeffGroup   [2];    // [ ChannelType ]

  static const CtxSet   LastX           [2];    // [ ChannelType ]

  static const CtxSet   LastY           [2];    // [ ChannelType ]

  static const CtxSet   SigFlag         [6];    // [ ChannelType + State ]

  static const CtxSet   ParFlag         [2];    // [ ChannelType ]

  static const CtxSet   GtxFlag         [4];    // [ ChannelType + x ]

  static const CtxSet   TsSigCoeffGroup;

  static const CtxSet   TsSigFlag;

  static const CtxSet   TsParFlag;

  static const CtxSet   TsGtxFlag;

  static const CtxSet   TsLrg1Flag;

  static const CtxSet   TsResidualSign;

  static const CtxSet   MVPIdx;

  static const CtxSet   SaoMergeFlag;

  static const CtxSet   SaoTypeIdx;

  static const CtxSet   TransformSkipFlag;

  static const CtxSet   MTSIdx;

  static const CtxSet   LFNSTIdx;

  static const CtxSet   PLTFlag;

  static const CtxSet   SbtFlag;

  static const CtxSet   SbtQuadFlag;

  static const CtxSet   SbtHorFlag;

  static const CtxSet   SbtPosFlag;

  static const CtxSet   ChromaQpAdjFlag;

  static const CtxSet   ChromaQpAdjIdc;

  static const CtxSet   ImvFlag;

  static const CtxSet   BcwIdx;

  static const CtxSet   ctbAlfFlag;

  static const CtxSet   ctbAlfAlternative;

  static const CtxSet   AlfUseTemporalFilt;

  static const CtxSet   CcAlfFilterControlFlag;

  static const CtxSet   CiipFlag;

  static const CtxSet   SmvdFlag;

  static const CtxSet   IBCFlag;

  static const CtxSet   JointCbCrFlag;

  static const unsigned NumberOfContexts;

  // combined sets for less complex copying

  // NOTE: The contained CtxSet's should directly follow each other in the initalization list;

  //       otherwise, you will copy more elements than you want !!!

  static const CtxSet   Sao;

  static const CtxSet   Alf;

public:

  static const std::vector<uint8_t>&  getInitTable( unsigned initId );

private:

  static std::vector<std::vector<uint8_t> > sm_InitTables;

  static CtxSet addCtxSet( std::initializer_list<std::initializer_list<uint8_t> > initSet2d );

};

class CtxStore

{

public:

  CtxStore();

  CtxStore( bool dummy );

  CtxStore( const CtxStore& ctxStore );

public:

  void copyFrom   ( const CtxStore& src )                        { checkInit(); ::memcpy( m_Ctx,               src.m_Ctx,               sizeof( BinProbModel ) * ContextSetCfg::NumberOfContexts ); }

  void copyFrom   ( const CtxStore& src, const CtxSet& ctxSet )  { checkInit(); ::memcpy( m_Ctx+ctxSet.Offset, src.m_Ctx+ctxSet.Offset, sizeof( BinProbModel ) * ctxSet.Size ); }

  void init       ( int qp, int initId );

  void setWinSizes( const std::vector<uint8_t>&   log2WindowSizes );

  void loadPStates( const std::vector<uint16_t>&  probStates );

  void savePStates( std::vector<uint16_t>&        probStates )  const;

  const BinProbModel& operator[]      ( unsigned  ctxId  )  const { return m_Ctx[ctxId]; }

  BinProbModel&       operator[]      ( unsigned  ctxId  )        { return m_Ctx[ctxId]; }

  uint32_t            estFracBits     ( unsigned  bin,

                                        unsigned  ctxId  )  const { return m_Ctx[ctxId].estFracBits(bin); }

  const BinFracBits  &getFracBitsArray( unsigned  ctxId  )  const { return m_Ctx[ctxId].getFracBitsArray(); }

private:

  inline void checkInit() { if( m_Ctx ) return; m_CtxBuffer.resize( ContextSetCfg::NumberOfContexts ); m_Ctx = m_CtxBuffer.data(); }

private:

  std::vector<BinProbModel> m_CtxBuffer;

  BinProbModel*             m_Ctx;

};

typedef CtxStore FracBitsAccess;

class Ctx;

class SubCtx

{

  friend class Ctx;

public:

  SubCtx( const CtxSet& ctxSet, const Ctx& ctx ) : m_CtxSet( ctxSet          ), m_Ctx( ctx          ) {}

  SubCtx( const SubCtx& subCtx )                 : m_CtxSet( subCtx.m_CtxSet ), m_Ctx( subCtx.m_Ctx ) {}

  const SubCtx& operator= ( const SubCtx& ) = delete;

private:

  const CtxSet  m_CtxSet;

  const Ctx&    m_Ctx;

};

class Ctx : public ContextSetCfg

{

public:

  Ctx();

  Ctx( const BinProbModel*    dummy );

  Ctx( const Ctx&                 ctx   );

public:

  const Ctx& operator= ( const Ctx& ctx )

  {

    m_CtxStore  .copyFrom( ctx.m_CtxStore);

    return *this;

  }

  SubCtx operator= ( SubCtx&& subCtx )

  {

    m_CtxStore  .copyFrom( subCtx.m_Ctx.m_CtxStore,   subCtx.m_CtxSet );

    return std::move(subCtx);

  }

  void  init ( int qp, int initId )

  {

    m_CtxStore  .init( qp, initId );

  }

  void  loadPStates( const std::vector<uint16_t>& probStates )

  {

    m_CtxStore  .loadPStates( probStates );

  }

  void  savePStates( std::vector<uint16_t>& probStates ) const

  {

    m_CtxStore  .savePStates( probStates );

  }

  void  initCtxAndWinSize( unsigned ctxId, const Ctx& ctx, const uint8_t winSize )

  {

    m_CtxStore  [ctxId] = ctx.m_CtxStore      [ctxId];

    m_CtxStore  [ctxId] . setLog2WindowSize   (winSize);

  }

public:

  const Ctx&            getCtx       ()     const { return *this; }

  Ctx&                  getCtx       ()           { return *this; }

  explicit operator   const CtxStore&()     const { return m_CtxStore; }

  explicit operator         CtxStore&()           { return m_CtxStore; }

  const FracBitsAccess& getFracBitsAcess()  const { return m_CtxStore; }

private:

  CtxStore                      m_CtxStore;

};

typedef dynamic_cache<Ctx> CtxCache;

class TempCtx

{

  TempCtx( const TempCtx& ) = delete;

  const TempCtx& operator=( const TempCtx& ) = delete;

public:

  TempCtx ( CtxCache* cache )                     : m_ctx( *cache->get() ), m_cache( cache ) {}

  TempCtx ( CtxCache* cache, const Ctx& ctx    )  : m_ctx( *cache->get() ), m_cache( cache ) { m_ctx = ctx; }

  TempCtx ( CtxCache* cache, SubCtx&&   subCtx )  : m_ctx( *cache->get() ), m_cache( cache ) { m_ctx = std::forward<SubCtx>(subCtx); }

  ~TempCtx()                                      { m_cache->cache( &m_ctx ); }

  const Ctx& operator=( const Ctx& ctx )          { return ( m_ctx = ctx ); }

  SubCtx     operator=( SubCtx&&   subCtx )       { return m_ctx = std::forward<SubCtx>( subCtx ); }

  operator const Ctx& ()           const          { return m_ctx; }

  operator       Ctx& ()                          { return m_ctx; }

private:

  Ctx&      m_ctx;

  CtxCache* m_cache;

};

} // namespace vvenc

//! \}