#ifndef ENCODERSTATE_H_
#define ENCODERSTATE_H_
/*****************************************************************************
 * This file is part of uvg266 VVC encoder.
 *
 * Copyright (c) 2021, Tampere University, ITU/ISO/IEC, project contributors
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted 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 Tampere University or ITU/ISO/IEC nor the names of its
 *   contributors may be used to endorse or promote products derived from
 *   this software without specific prior written permission.
 * 
 * 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
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 * INCLUDING NEGLIGENCE OR OTHERWISE ARISING IN ANY WAY OUT OF THE USE OF THIS
 ****************************************************************************/

/**
 * \ingroup Control
 * \file
 * Top level of the encoder implementation.
 */

#include "bitstream.h"
#include "cabac.h"
#include "cu.h"
#include "encoder.h"
#include "global.h" // IWYU pragma: keep
#include "image.h"
#include "imagelist.h"
#include "uvg266.h"
#include "tables.h"
#include "threadqueue.h"
#include "videoframe.h"

struct uvg_rc_data;

typedef enum {
  ENCODER_STATE_TYPE_INVALID = 'i',
  ENCODER_STATE_TYPE_MAIN = 'M',
  ENCODER_STATE_TYPE_SLICE = 'S',
  ENCODER_STATE_TYPE_TILE = 'T',
  ENCODER_STATE_TYPE_WAVEFRONT_ROW = 'W',
} encoder_state_type;


typedef struct lcu_stats_t {
  //! \brief Number of bits that were spent
  uint32_t bits;

  uint32_t pixels;

  //! \brief Weight of the LCU for rate control
  double weight;

  double original_weight;

  //! \brief Lambda value which was used for this LCU
  double lambda;

  double adjust_lambda;

  //! \brief Rate control alpha parameter
  double rc_alpha;

  //! \brief Rate control beta parameter
  double rc_beta;
  double distortion;
  int i_cost;

  int8_t qp;
  int8_t adjust_qp;
  uint8_t skipped;
} lcu_stats_t;


typedef struct encoder_state_config_frame_t {
  /**
   * \brief Frame-level lambda.
   *
   * Use state->lambda or state->lambda_sqrt for cost computations.
   *
   * \see encoder_state_t::lambda
   * \see encoder_state_t::lambda_sqrt
   */
  double lambda;

  int32_t num;       /*!< \brief Frame number */
  int32_t poc;       /*!< \brief Picture order count */
  int8_t gop_offset; /*!< \brief Offset in the gop structure */
  int32_t irap_poc;  /*!< \brief POC of the associated IRAP picture */

  /**
   * \brief Frame-level quantization parameter
   *
   * \see encoder_state_t::qp
   */
  int8_t QP;
  //! \brief quantization factor
  double QP_factor;

  //! Current pictures available for references
  image_list_t *ref;
  int8_t ref_list;

  //! L0 and L1 reference index list
  uint8_t ref_LX[2][16];
  //! L0 reference index list size
  uint8_t ref_LX_size[2];

  bool is_irap;
  uint8_t pictype;
  enum uvg_slice_type slicetype;

  //! Total number of bits written.
  uint64_t total_bits_coded;

  //! Number of bits written in the current GOP.
  uint64_t cur_gop_bits_coded;

  //! Number of bits written in the current frame.
  uint64_t cur_frame_bits_coded;

  //! Number of bits targeted for the current GOP.
  double cur_gop_target_bits;

  //! Number of bits targeted for the current picture.
  double cur_pic_target_bits;

  // Parameters used in rate control
  double rc_alpha;
  double rc_beta;

  /**
   * \brief Indicates that this encoder state is ready for encoding the
   * next frame i.e. uvg_encoder_prepare has been called.
   */
  bool prepared;

  /**
   * \brief Indicates that the previous frame has been encoded and the
   * encoded data written and the encoding the next frame has not been
   * started yet.
   */
  bool done;

  /**
   * \brief Information about the coded LCUs.
   *
   * Used for rate control.
   */
  lcu_stats_t *lcu_stats;

  pthread_mutex_t rc_lock;

  struct uvg_rc_data *new_ratecontrol;

  struct encoder_state_t const *previous_layer_state;

  /**
  * \brief Calculated adaptive QP offset for each LCU.
  */
  double *aq_offsets;

  int8_t max_qp_delta_depth;

  /**
   * \brief Whether next NAL is the first NAL in the access unit.
   */
  bool first_nal;
  double icost;
  double remaining_weight;
  double i_bits_left;

  double *c_para;
  double *k_para;

  bool jccr_sign; 

} encoder_state_config_frame_t;

typedef struct encoder_state_config_tile_t {
  //Current sub-frame
  videoframe_t *frame;
  
  int32_t id;

  //Tile: offset in LCU for current encoder_state in global coordinates
  int32_t lcu_offset_x;
  int32_t lcu_offset_y;

  //Tile: offset in pixels
  int32_t offset_x;
  int32_t offset_y;

  //Position of the first element in tile scan in global coordinates
  int32_t lcu_offset_in_ts;
  
  // This is a buffer for the non-loopfiltered bottom pixels of every LCU-row
  // in the tile. They are packed such that each LCU-row index maps to the
  // y-coordinate.
  yuv_t *hor_buf_search;
  // This is a buffer for the non-loopfiltered rightmost pixels of every
  // LCU-column. They are packed such that each LCU-column index maps to the
  // x-coordinate.
  yuv_t *ver_buf_search;

  // This is a buffer for the deblocked bottom pixels of every LCU in the
  // tile. They are packed such that each LCU-row index maps to the
  // y-coordinate.
  yuv_t *hor_buf_before_sao;

  // This is a buffer for the deblocked right pixels of every LCU in the
  // tile. They are packed such that each LCU-column index maps to the
  // x-coordinate.
  yuv_t *ver_buf_before_sao;

  //Jobs for each individual LCU of a wavefront row.
  threadqueue_job_t **wf_jobs;
  threadqueue_job_t **wf_recon_jobs;

} encoder_state_config_tile_t;

typedef struct encoder_state_config_alf_t {
  //ALF adaptation parameter set
  struct alf_aps* apss; //[ALF_CTB_MAX_NUM_APS];
  struct cc_alf_filter_param* cc_filter_param;
  int tile_group_num_aps;
  int8_t* tile_group_luma_aps_id;
  int tile_group_chroma_aps_id;
  bool tile_group_cc_alf_cb_enabled_flag;
  bool tile_group_cc_alf_cr_enabled_flag;
  int tile_group_cc_alf_cb_aps_id;
  int tile_group_cc_alf_cr_aps_id;
  //struct param_set_map *param_set_map; //mahdollisesti define during run
  uint32_t num_of_param_sets;
  bool tile_group_alf_enabled_flag[3/*MAX_NUM_COMPONENT*/];
} encoder_state_config_alf_t;

typedef struct encoder_state_config_slice_t {
  int32_t id;
  
  //Global coordinates
  int32_t start_in_ts;
  int32_t end_in_ts;
  
  //Global coordinates
  int32_t start_in_rs;
  int32_t end_in_rs;

  encoder_state_config_alf_t *alf;
  //unsigned num_hor_virtual_boundaries;
  //unsigned num_ver_virtual_boundaries;
  //unsigned virtual_boundaries_pos_x[3];
  //unsigned virtual_boundaries_pos_y[3];
} encoder_state_config_slice_t;

typedef struct encoder_state_config_wfrow_t {
  //Row in tile coordinates of the wavefront
  int32_t lcu_offset_y;
} encoder_state_config_wfrow_t;

typedef struct lcu_order_element {
  //This it used for leaf of the encoding tree. All is relative to the tile.
  int id;
  int index;
  struct encoder_state_t *encoder_state;
  vector2d_t position;
  vector2d_t position_px; //Top-left
  vector2d_t size;
  int first_column;
  int first_row;
  int last_column;
  int last_row;
  lcu_coeff_t *coeff;
  
  struct lcu_order_element *above;
  struct lcu_order_element *below;
  struct lcu_order_element *left;
  struct lcu_order_element *right;
} lcu_order_element_t;

typedef struct encoder_state_t {
  const encoder_control_t *encoder_control;
  encoder_state_type type;

  //List of children, the last item of this list is a pseudo-encoder with encoder_control = NULL
  //Use for (i = 0; encoder_state->children[i].encoder_control; ++i) {
  struct encoder_state_t *children;
  struct encoder_state_t *parent;
  
  //Pointer to the encoder_state of the previous frame
  struct encoder_state_t *previous_encoder_state;
    
  encoder_state_config_frame_t  *frame;
  encoder_state_config_tile_t   *tile;
  encoder_state_config_slice_t  *slice;
  encoder_state_config_wfrow_t  *wfrow;
  
  int is_leaf; //A leaf encoder state is one which should encode LCUs...
  lcu_order_element_t *lcu_order;
  uint32_t lcu_order_count;
  
  bitstream_t stream;
  cabac_data_t cabac;
  cabac_data_t search_cabac;

  uint32_t stats_bitstream_length; //Bitstream length written in bytes

  //! \brief Lambda for SSE
  double lambda;
  //! \brief Lambda for SAD and SATD
  double lambda_sqrt;
  //! \brief Quantization parameter for the current LCU
  int8_t qp;

  double c_lambda;
  double chroma_weights[4];

  /**
   * \brief Whether a QP delta value must be coded for the current LCU.
   */
  bool must_code_qp_delta;

  /**
   * \brief QP value of the last CU in the last coded quantization group.
   *
   * A quantization group is a square of width
   * (LCU_WIDTH >> encoder_control->max_qp_delta_depth). All CUs of in the
   * same quantization group share the QP predictor value, but may have
   * different QP values.
   *
   * Set to the frame QP at the beginning of a wavefront row or a tile and
   * updated when the last CU of a quantization group is coded.
   */
  int8_t last_qp;

  //Jobs to wait for
  threadqueue_job_t * tqj_recon_done; //Reconstruction is done
  threadqueue_job_t * tqj_bitstream_written; //Bitstream is written
  threadqueue_job_t*  tqj_alf_process; //ALF processed for the slice

  //Constraint structure  
  void * constraint;

  // Since lfnst needs the collocated luma intra mode for
  // dual tree if the chroma mode is cclm mode and getting all of
  // the information that would be necessary to get the collocated
  // luma mode in the lfnst functions, instead store the current
  // collocated luma mode in the state.
  int8_t collocated_luma_mode;

  quant_block quant_blocks[3]; // luma, ISP, chroma
  rate_estimator_t rate_estimator[4]; // luma, cb, cr, isp
} encoder_state_t;

void uvg_encode_one_frame(encoder_state_t * const state, uvg_picture* frame);

void uvg_encoder_prepare(encoder_state_t *state);


int uvg_encoder_state_match_children_of_previous_frame(encoder_state_t * const state);

coeff_scan_order_t uvg_get_scan_order(int8_t cu_type, int intra_mode, int depth);

void uvg_encoder_create_ref_lists(const encoder_state_t *const state);

lcu_stats_t* uvg_get_lcu_stats(encoder_state_t *state, int lcu_x, int lcu_y);


int uvg_get_cu_ref_qp(const encoder_state_t *state, int x, int y, int last_qp);

/**
 * Whether the parameter sets should be written with the current frame.
 */
static INLINE bool encoder_state_must_write_vps(const encoder_state_t *state)
{
  const int32_t frame = state->frame->num;
  const int32_t vps_period = state->encoder_control->cfg.vps_period;

  return (vps_period >  0 && frame % vps_period == 0) ||
         (vps_period >= 0 && frame == 0);
}


/**
 * \brief Returns true if the CU is the last CU in its containing
 * quantization group.
 *
 * \param state   encoder state
 * \param x       x-coordinate of the left edge of the CU
 * \param y       y-cooradinate of the top edge of the CU
 * \param depth   depth in the CU tree
 * \return true, if it's the last CU in its QG, otherwise false
 */
static INLINE bool is_last_cu_in_qg(const encoder_state_t *state, const cu_loc_t* const cu_loc)
{
  if (state->frame->max_qp_delta_depth < 0) return false;
  
  const int qg_width = LCU_WIDTH >> state->frame->max_qp_delta_depth;
  const int right  = cu_loc->x + cu_loc->width;
  const int bottom = cu_loc->y + cu_loc->height;
  return (right % qg_width == 0 || right >= state->tile->frame->width) &&
         (bottom % qg_width == 0 || bottom >= state->tile->frame->height);
}

static const uint8_t g_group_idx[64] = {
  0, 1, 2, 3, 4, 4, 5, 5,
  6, 6, 6, 6, 7, 7, 7, 7,
  8, 8, 8, 8, 8, 8, 8, 8,
  9, 9, 9, 9, 9, 9, 9, 9,
 10,10,10,10,10,10,10,10,
 10,10,10,10,10,10,10,10,
 11,11,11,11,11,11,11,11,
 11,11,11,11,11,11,11,11
};

/*
static const uint8_t g_group_idx[128] = {
  0, 1, 2, 3, 4, 4, 5, 5, 6, 6,
  6, 6, 7, 7, 7, 7, 8, 8, 8, 8,
  8, 8, 8, 8, 9, 9, 9, 9, 9, 9,
  9, 9,10,10,10,10,10,10,10,10,
 10,10,10,10,10,10,10,10,11,11,
 11,11,11,11,11,11,11,11,11,11,
 11,11,11,11,12,12,12,12,12,12,
 12,12,12,12,12,12,12,12,12,12,
 12,12,12,12,12,12,12,12,12,12,
 12,12,12,12,12,12,13,13,13,13,
 13,13,13,13,13,13,13,13,13,13,
 13,13,13,13,13,13,13,13,13,13,
 13,13,13,13,13,13,13,13 };
 */

static const uint8_t g_min_in_group[14] = {
  0,1,2,3,4,6,8,12,16,24,32,48,64,96 };


#define C1FLAG_NUMBER 8 // maximum number of largerThan1 flag coded in one chunk
#define C2FLAG_NUMBER 1 // maximum number of largerThan2 flag coded in one chunk

//Get the data for vertical buffer position at the left of LCU identified by the position in pixel
#define OFFSET_VER_BUF(position_x, position_y, cur_pic, i) ((position_y) + i + ((position_x)/LCU_WIDTH - 1) * (cur_pic)->height)
#define OFFSET_VER_BUF_C(position_x, position_y, cur_pic, i) ((position_y/2) + i + ((position_x)/LCU_WIDTH - 1) * (cur_pic)->height / 2)

//Get the data for horizontal buffer position at the top of LCU identified by the position in pixel
#define OFFSET_HOR_BUF(position_x, position_y, cur_pic, i) ((position_x) + i + ((position_y)/LCU_WIDTH - 1) * (cur_pic)->width)
#define OFFSET_HOR_BUF_C(position_x, position_y, cur_pic, i) ((position_x/2) + i + ((position_y)/LCU_WIDTH - 1) * (cur_pic)->width / 2)
  
/** @} */

#endif //ENCODERSTATE_H_