/*!
 *****************************************************************************
 *
 * \file fmo.c
 *
 * \brief
 *    Support for Flexible Macroblock Ordering for different Slice Group Modes: MBAmap handling
 *
 * \date
 *    16 June, 2002  Modified April 25, 2004
 *
 * \author
 *    Stephan Wenger   stewe@cs.tu-berlin.de
 *    Dong Wang (modify)	Dong.Wang@bristol.ac.uk
 * 
 *****************************************************************************/

/*!
 ****************************************************************************
 *   Notes by Dong Wang (April 25 2004)
 *
 *  Source codes are modified to support 7 slice group types (fmo modes).
 *  The functions for generating map are very similar to that in decoder, but have 
 *  a little difference. 
 *
 *  The MB map is calculated at the beginning of coding of each picture (frame or field).
 *
 *  'slice_group_change_cycle' in structure 'ImageParameters' is the syntax in the slice 
 *  header. It's set to be 1 before the initialization of FMO in function code_a_picture().
 *  It can be changed every time if needed.
 *
 **************************************************************************** 
 */

/*!
 *****************************************************************************
 *  How does a MBAmap look like?
 *
 *  An MBAmap is a one-diemnsional array of ints.  Each int 
 *  represents an MB in scan order.  A zero or positive value represents
 *  a slice group ID.  Negative values are reserved for future extensions.
 *  The numbering range for the SliceGroupIDs is 0..7 as per JVT-C167.
 *
 *  This module contains a static variable MBAmap.  This is the MBAmap of the
 *  picture currently coded.  It can be accessed only through the access
 *  functions.
 *****************************************************************************
*/

//#define PRINT_FMO_MAPS  1


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

#include "global.h"

#include "fmo.h"


static int FirstMBInSlice[MAXSLICEGROUPIDS];

int *MBAmap = NULL;   
int *MapUnitToSliceGroupMap = NULL; 
unsigned PicSizeInMapUnits;


static void FmoGenerateType0MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static void FmoGenerateType1MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static void FmoGenerateType2MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static void FmoGenerateType3MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static void FmoGenerateType4MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static void FmoGenerateType5MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static void FmoGenerateType6MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);


static int FmoGenerateMapUnitToSliceGroupMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps);
static int FmoGenerateMBAmap (ImageParameters * img, seq_parameter_set_rbsp_t* sps);


/*!
 ************************************************************************
 * \brief
 *    Generates MapUnitToSliceGroupMap
 *
 * \param img
 *    Image Parameter to be used for map generation
 * \param pps
 *    Picture Parameter set to be used for map generation
 *
 ************************************************************************
 */
static int FmoGenerateMapUnitToSliceGroupMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps)
{
  PicSizeInMapUnits = img->PicHeightInMapUnits * img->PicWidthInMbs;
  
  
  if (pps->slice_group_map_type == 6)
  {
    if ((pps->pic_size_in_map_units_minus1+1) != PicSizeInMapUnits)
    {
      error ("wrong pps->pic_size_in_map_units_minus1 for used SPS and FMO type 6", 500);
    }
  }
  
  // allocate memory for MapUnitToSliceGroupMap
  if (MapUnitToSliceGroupMap)
    free (MapUnitToSliceGroupMap);
  
  if ((MapUnitToSliceGroupMap = malloc ((PicSizeInMapUnits) * sizeof (int))) == NULL)
  {
    printf ("cannot allocated %d bytes for MapUnitToSliceGroupMap, exit\n", PicSizeInMapUnits * sizeof (int));
    exit (-1);
  }
  
  if (pps->num_slice_groups_minus1 == 0)    // only one slice group
  {
    memset (MapUnitToSliceGroupMap, 0,  PicSizeInMapUnits * sizeof (int));
    return 0;
  }
  
  switch (pps->slice_group_map_type)
  {
  case 0:
    FmoGenerateType0MapUnitMap (img, pps);
    break;
  case 1:
    FmoGenerateType1MapUnitMap (img, pps);
    break;
  case 2:
    FmoGenerateType2MapUnitMap (img, pps);
    break;
  case 3:
    FmoGenerateType3MapUnitMap (img, pps);
    break;
  case 4:
    FmoGenerateType4MapUnitMap (img, pps);
    break;
  case 5:
    FmoGenerateType5MapUnitMap (img, pps);
    break;
  case 6:
    FmoGenerateType6MapUnitMap (img, pps);
    break;
  default:
    printf ("Illegal slice_group_map_type %d , exit \n", pps->slice_group_map_type);
    exit (-1);
  }
  return 0;
}


/*!
 ************************************************************************
 * \brief
 *    Generates MBAmap from MapUnitToSliceGroupMap
 *
 * \param img
 *    Image Parameter to be used for map generation
 * \param sps
 *    Sequence Parameter set to be used for map generation
 *
 ************************************************************************
 */
static int FmoGenerateMBAmap (ImageParameters * img, seq_parameter_set_rbsp_t* sps)
{
  unsigned i;
  
  // allocate memory for MBAmap
  if (MBAmap)
    free (MBAmap);
  
  
  if ((MBAmap = malloc ((img->PicSizeInMbs) * sizeof (int))) == NULL)
  {
    printf ("cannot allocated %d bytes for MBAmap, exit\n", (img->PicSizeInMbs) * sizeof (int));
    exit (-1);
  }
  
  if ((sps->frame_mbs_only_flag) || img->field_picture)
  {
    for (i=0; i<img->PicSizeInMbs; i++)
    {
      MBAmap[i] = MapUnitToSliceGroupMap[i];
    }
  }
  else
    if (sps->mb_adaptive_frame_field_flag  &&  (! img->field_picture))
    {
      for (i=0; i<img->PicSizeInMbs; i++)
      {
        MBAmap[i] = MapUnitToSliceGroupMap[i/2];
      }
    }
    else
    {
      for (i=0; i<img->PicSizeInMbs; i++)
      {
        MBAmap[i] = MapUnitToSliceGroupMap[(i/(2*img->PicWidthInMbs))*img->PicWidthInMbs+(i%img->PicWidthInMbs)];
      }
    }
    return 0;
}


/*!
 ************************************************************************
 * \brief
 *    FMO initialization: Generates MapUnitToSliceGroupMap and MBAmap.
 *
 * \param img
 *    Image Parameter to be used for map generation
 * \param pps
 *    Picture Parameter set to be used for map generation
 * \param sps
 *    Sequence Parameter set to be used for map generation
 ************************************************************************
 */
int FmoInit(ImageParameters * img, pic_parameter_set_rbsp_t * pps, seq_parameter_set_rbsp_t * sps)
{
  
#ifdef PRINT_FMO_MAPS
  unsigned i,j;
  int bottom;
#endif
  
  int k;
  for (k=0;k<MAXSLICEGROUPIDS;k++)
    FirstMBInSlice[k] = -1;
  
  
  
  FmoGenerateMapUnitToSliceGroupMap(img, pps);
  FmoGenerateMBAmap(img, sps);
  
#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("%d ",MapUnitToSliceGroupMap[i+j*img->PicWidthInMbs]);
    }
    printf("\n");
  }
  printf("\n");
  
  if(sps->mb_adaptive_frame_field_flag==0)
  {
    printf("FMO Map (Mb):\n");	
    for (j=0; j<(img->PicSizeInMbs/img->PicWidthInMbs); j++)
    {
      for (i=0; i<img->PicWidthInMbs; i++)
      {
        printf("%d ",MBAmap[i+j*img->PicWidthInMbs]);
      }
      printf("\n");
    }
    printf("\n");
  }
  else 
  {
    printf("FMO Map (Mb in scan order for MBAFF):\n");	
    for (j=0; j<(img->PicSizeInMbs/img->PicWidthInMbs); j++)
    {
      for (i=0; i<img->PicWidthInMbs; i++)
      {
        bottom=(j%2);
        printf("%d ",MBAmap[(j-bottom)*img->PicWidthInMbs+i*2+bottom]);
      }
      printf("\n");
      
    }
    printf("\n");
    
  }
  
#endif
  
  return 0;
}


/*!
 ************************************************************************
 * \brief
 *    Free memory if allocated by FMO functions
 ************************************************************************
 */
void FmoUninit()
{
  if (MBAmap)
  {
    free (MBAmap);
    MBAmap = NULL;
  }
  if (MapUnitToSliceGroupMap)
  {
    free (MapUnitToSliceGroupMap);
    MapUnitToSliceGroupMap = NULL; 
  }
  
}


/*!
 ************************************************************************
 * \brief
 *    Generate interleaved slice group map type MapUnit map (type 0)
 *
 * \param img
 *    Image Parameter to be used for map generation
 * \param pps
 *    Picture Parameter set to be used for map generation
 ************************************************************************
 */
static void FmoGenerateType0MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps )
{
  unsigned iGroup, j;
  unsigned i = 0;
  do
  {
    for( iGroup = 0; 
    (iGroup <= pps->num_slice_groups_minus1) && (i < PicSizeInMapUnits); 
    i += pps->run_length_minus1[iGroup++] + 1)
    {
      for( j = 0; j <= pps->run_length_minus1[ iGroup ] && i + j < PicSizeInMapUnits; j++ )
        MapUnitToSliceGroupMap[i+j] = iGroup;
    }
  }
  while( i < PicSizeInMapUnits );
}

  
/*!
 ************************************************************************
 * \brief
 *    Generate dispersed slice group map type MapUnit map (type 1)
 *
 * \param img
 *    Image Parameter to be used for map generation
 * \param pps
 *    Picture Parameter set to be used for map generation
 ************************************************************************
 */
static void FmoGenerateType1MapUnitMap (ImageParameters * img, pic_parameter_set_rbsp_t * pps )
{
  unsigned i;
  for( i = 0; i < PicSizeInMapUnits; i++ )
  {
    MapUnitToSliceGroupMap[i] = ((i%img->PicWidthInMbs)+(((i/img->PicWidthInMbs)*(pps->num_slice_groups_minus1+1))/2))
      %(pps->num_slice_groups_minus1+1);
  }
}

/*!
 ************************************************************************
 * \brief
 *    Generate foreground with left-over slice group map type MapUnit map (type 2)
 *
 * \param img
 *    Image Parameter to be used for map generation