/*!
 *****************************************************************************
 *
 * \file fmo.c
 *
 * \brief
 *    Support for Flexible Macroblock Ordering (FMO)
 *
 * \author
 *    Main contributors (see contributors.h for copyright, address and affiliation details)
 *    - Stephan Wenger      stewe@cs.tu-berlin.de
 *    - Karsten Suehring    suehring@hhi.de
 ******************************************************************************
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>


#include "global.h"
#include "elements.h"
#include "defines.h"
#include "header.h"
#include "fmo.h"

//#define PRINT_FMO_MAPS

int *MbToSliceGroupMap = NULL;
int *MapUnitToSliceGroupMap = NULL; 

static int NumberOfSliceGroups;    // the number of slice groups -1 (0 == scan order, 7 == maximum)

static void FmoGenerateType0MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );
static void FmoGenerateType1MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );
static void FmoGenerateType2MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );
static void FmoGenerateType3MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );
static void FmoGenerateType4MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );
static void FmoGenerateType5MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );
static void FmoGenerateType6MapUnitMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps, unsigned PicSizeInMapUnits );


/*!
 ************************************************************************
 * \brief
 *    Generates MapUnitToSliceGroupMap
 *    Has to be called every time a new Picture Parameter Set is used
 *
 * \param pps
 *    Picture Parameter set to be used for map generation
 * \param sps
 *    Sequence Parameter set to be used for map generation
 *
 ************************************************************************
 */
static int FmoGenerateMapUnitToSliceGroupMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps)
{
  unsigned int NumSliceGroupMapUnits;

  NumSliceGroupMapUnits = (sps->pic_height_in_map_units_minus1+1)* (sps->pic_width_in_mbs_minus1+1);

  if (pps->slice_group_map_type == 6)
  {
    if ((pps->num_slice_group_map_units_minus1+1) != NumSliceGroupMapUnits)
    {
      error ("wrong pps->num_slice_group_map_units_minus1 for used SPS and FMO type 6", 500);
    }
  }

  // allocate memory for MapUnitToSliceGroupMap
  if (MapUnitToSliceGroupMap)
    free (MapUnitToSliceGroupMap);
  if ((MapUnitToSliceGroupMap = malloc ((NumSliceGroupMapUnits) * sizeof (int))) == NULL)
  {
    printf ("cannot allocated %d bytes for MapUnitToSliceGroupMap, exit\n", (pps->num_slice_group_map_units_minus1+1) * sizeof (int));
    exit (-1);
  }

  if (pps->num_slice_groups_minus1 == 0)    // only one slice group
  {
    memset (MapUnitToSliceGroupMap, 0, NumSliceGroupMapUnits * sizeof (int));
    return 0;
  }

  switch (pps->slice_group_map_type)
  {
  case 0:
    FmoGenerateType0MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  case 1:
    FmoGenerateType1MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  case 2:
    FmoGenerateType2MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  case 3:
    FmoGenerateType3MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  case 4:
    FmoGenerateType4MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  case 5:
    FmoGenerateType5MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  case 6:
    FmoGenerateType6MapUnitMap (pps, sps, NumSliceGroupMapUnits);
    break;
  default:
    printf ("Illegal slice_group_map_type %d , exit \n", pps->slice_group_map_type);
    exit (-1);
  }
  return 0;
}


/*!
 ************************************************************************
 * \brief
 *    Generates MbToSliceGroupMap from MapUnitToSliceGroupMap
 *
 * \param pps
 *    Picture Parameter set to be used for map generation
 * \param sps
 *    Sequence Parameter set to be used for map generation
 *
 ************************************************************************
 */
static int FmoGenerateMbToSliceGroupMap (pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps)
{
  unsigned i;

  // allocate memory for MbToSliceGroupMap
  if (MbToSliceGroupMap)
    free (MbToSliceGroupMap);

  if ((MbToSliceGroupMap = malloc ((img->PicSizeInMbs) * sizeof (int))) == NULL)
  {
    printf ("cannot allocated %d bytes for MbToSliceGroupMap, exit\n", (img->PicSizeInMbs) * sizeof (int));
    exit (-1);
  }


  if ((sps->frame_mbs_only_flag)|| img->field_pic_flag)
  {
    for (i=0; i<img->PicSizeInMbs; i++)
    {
      MbToSliceGroupMap[i] = MapUnitToSliceGroupMap[i];
    }
  }
  else
    if (sps->mb_adaptive_frame_field_flag  &&  (!img->field_pic_flag))
    {
      for (i=0; i<img->PicSizeInMbs; i++)
      {
        MbToSliceGroupMap[i] = MapUnitToSliceGroupMap[i/2];
      }
    }
    else
    {
      for (i=0; i<img->PicSizeInMbs; i++)
      {
        MbToSliceGroupMap[i] = MapUnitToSliceGroupMap[(i/(2*img->PicWidthInMbs))*img->PicWidthInMbs+(i%img->PicWidthInMbs)];
      }
    }
  return 0;
}


/*!
 ************************************************************************
 * \brief
 *    FMO initialization: Generates MapUnitToSliceGroupMap and MbToSliceGroupMap.
 *
 * \param pps
 *    Picture Parameter set to be used for map generation
 * \param sps
 *    Sequence Parameter set to be used for map generation
 ************************************************************************
 */
int FmoInit(pic_parameter_set_rbsp_t* pps, seq_parameter_set_rbsp_t* sps)
{
#ifdef PRINT_FMO_MAPS
  unsigned i,j;
#endif

  FmoGenerateMapUnitToSliceGroupMap(pps, sps);
  FmoGenerateMbToSliceGroupMap(pps, sps);

  NumberOfSliceGroups = pps->num_slice_groups_minus1+1;

#ifdef PRINT_FMO_MAPS
  printf("\n");
  printf("FMO Map (Units):\n");

  for (j=0; j<img->PicHeightInMapUnits; j++)
  {
    for (i=0; i<img->PicWidthInMbs; i++)
    {
      printf("%c",48+MapUnitToSliceGroupMap[i+j*img->PicWidthInMbs]);
    }
    printf("\n");
  }
  printf("\n");
  printf("FMO Map (Mb):\n");

  for (j=0; j<img->PicHeightInMbs; j++)
  {
    for (i=0; i<img->PicWidthInMbs; i++)
    {
      printf("%c",48+MbToSliceGroupMap[i+j*img->PicWidthInMbs]);
    }
    printf("\n");
  }
  printf("\n");

#endif

  return 0;
}


/*!
 ************************************************************************
 * \brief
 *    Free memory allocated by FMO functions
 ************************************************************************
 */
int FmoFinit()
{
  if (MbToSliceGroupMap)
  {
    free (MbToSliceGroupMap);
    MbToSliceGroupMap = NULL;
  }
  if (MapUnitToSliceGroupMap)
  {
    free (MapUnitToSliceGroupMap);
    MapUnitToSliceGroupMap = NULL; 
  }
  return 0;
}


/*!
 ************************************************************************
 * \brief
 *    FmoGetNumberOfSliceGroup() 
 *
 * \par Input:
 *    None
 ************************************************************************
 */
int FmoGetNumberOfSliceGroup()
{
  return NumberOfSliceGroups;
}


/*!
 ************************************************************************
 * \brief
 *    FmoGetLastMBOfPicture() 
 *    returns the macroblock number of the last MB in a picture.  This
 *    mb happens to be the last macroblock of the picture if there is only
 *    one slice group
 *
 * \par Input:
 *    None
 ************************************************************************
 */
int FmoGetLastMBOfPicture()
{
  return FmoGetLastMBInSliceGroup (FmoGetNumberOfSliceGroup()-1);
}