/*
 * Copyright 2011-2019, European Organisation for the Exploitation of Meteorological Satellites (EUMETSAT)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "CDataField.h"

#include "CBuffer.h"

#include "CWBlock.h"
#include "CACCoder.h"
#include "CVLCCoder.h"
#include "CWTCoder.h"

namespace COMP 
{


void CWTCoder::CodeBufferBlock (const unsigned int i_Bs)
{
	COMP_TRYTHIS
	unsigned int bX, bY;
	
	// Write header
	m_Buf.write_marker (c_MarkerHEADER);
	m_Buf.real_write (m_Param.m_BitsPerPixel, 4);
	m_Buf.real_write (m_Image.GetW(), 16);
	m_Buf.real_write (m_Image.GetH(), 16);
	m_Buf.real_write (m_Param.m_nWTlevels - 3, 2);
	m_Buf.real_write (m_Param.m_PredMode - CWTParams::e_None, 2);
	m_Buf.real_write (i_Bs >> 5, 2);
	m_Buf.real_write (m_Param.m_RestartInterval, 16);
	m_Buf.real_write (m_Param.m_nLossyBitPlanes, 4);
	m_Buf.real_write (0, 2);	// Pad bits
	m_Buf.write_marker (c_MarkerDATA);

	// Computing the number of blocks along each dimension
	unsigned int nbW = m_Image.GetW() / i_Bs;
	const unsigned int rW = m_Image.GetW() % i_Bs;
	if (rW) nbW++;
	unsigned int nbH = m_Image.GetH() / i_Bs;
	const unsigned int rH = m_Image.GetH() % i_Bs;
	if (rH) nbH++;
	Assert (nbW > 0 && nbH > 0, Util::CParamException());
		
	// Start arithmetic coder
	CACCoder ac (m_Buf);
	ac.Start ();

	// VLC coder object creation
	CVLCCoder coder (ac);

	// Wavelet block object creation
	CWBlock wblk (i_Bs, i_Bs);

	// block loops
	unsigned int nbBlock = 0;
	unsigned int markerNum = 0;
	bool bACStopped = false;
	for (bY=0 ; bY<nbH ; bY++)
	{
		const unsigned int nH = bY == (nbH - 1) && rH ? rH : i_Bs;
		for (bX=0 ; bX<nbW ; bX++)
		{
			const unsigned int nW = bX == (nbW - 1) && rW ? rW : i_Bs;
			
			// Copy image block to CWBlock with line and column padding
			wblk.GetAndPad (m_Image, bX * i_Bs, bY * i_Bs, nW, nH);
			
			// Perform the forward m_nWTlevels iterations of 2D S(+P) transform
			switch (m_Param.m_PredMode)
			{
			case CWTParams::e_None:
				wblk.IterateSt (true, m_Param.m_nWTlevels);
				break;
			case CWTParams::e_PredA:
				wblk.IterateSptA (true, m_Param.m_nWTlevels);
				break;
			case CWTParams::e_PredB:
				wblk.IterateSptB (true, m_Param.m_nWTlevels);
				break;
			case CWTParams::e_PredC:
				wblk.IterateSptC (true, m_Param.m_nWTlevels);
				break;
			default:
				Assert (0, Util::CParamException())
			}
			
			// code the S+P coefs
			coder.Code (wblk, m_Param.m_nWTlevels, m_Param.m_nLossyBitPlanes);

			nbBlock++;
			if (m_Param.m_RestartInterval && nbBlock == m_Param.m_RestartInterval)
			{
				// Output a restart marker
				nbBlock = 0;
				ac.Stop ();
				m_Buf.write_marker (c_MarkerRESTART | (markerNum & 0xF));
				markerNum++;
				if (bX < (nbW - 1) || bY < (nbH - 1)) ac.Start ();
				else bACStopped = true;
				// Reset VLC decoder state (models probabilities)
				coder.Reset ();
			}
		}
	}

	// Stop the arithmetic coding
	if (! bACStopped) ac.Stop ();

	// Write the Footer marker
	m_Buf.write_marker (c_MarkerFOOTER);

	COMP_CATCHTHIS
}

void CWTCoder::CodeBufferFull (void)
{
	COMP_TRYTHIS
	// Write header
	m_Buf.write_marker (c_MarkerHEADER);
	m_Buf.real_write (m_Param.m_BitsPerPixel, 4);
	m_Buf.real_write (m_Image.GetW(), 16);
	m_Buf.real_write (m_Image.GetH(), 16);
	m_Buf.real_write (m_Param.m_nWTlevels - 3, 2);
	m_Buf.real_write (m_Param.m_PredMode - CWTParams::e_None, 2);
	m_Buf.real_write (3, 2); // Full block
	m_Buf.real_write (m_Param.m_RestartInterval, 16);
	m_Buf.real_write (m_Param.m_nLossyBitPlanes, 4);
	m_Buf.real_write (0, 2);	// Pad bits
	m_Buf.write_marker (c_MarkerDATA);

	// Start arithmetic coder
	CACCoder ac (m_Buf);
	ac.Start ();

	// Compute the Wavelet block size to allow m_nWTlevels S+P iterations
	int bW = m_Image.GetW();
	int bH = m_Image.GetH();
	const int iteS = 1U << m_Param.m_nWTlevels;
	bW = (bW + iteS - 1) & (- iteS);
	bH = (bH + iteS - 1) & (- iteS);
	
	// Wavelet block object creation
	CWBlock wblk (bW, bH);

	// Copy image to CWBlock with line and column padding
	wblk.GetAndPad (m_Image, 0, 0, m_Image.GetW(), m_Image.GetH());
	
	// Perform the forward m_nWTlevels iterations of 2D S(+P) transform
	switch (m_Param.m_PredMode)
	{
	case CWTParams::e_None:
		wblk.IterateSt (true, m_Param.m_nWTlevels);
		break;
	case CWTParams::e_PredA:
		wblk.IterateSptA (true, m_Param.m_nWTlevels);
		break;
	case CWTParams::e_PredB:
		wblk.IterateSptB (true, m_Param.m_nWTlevels);
		break;
	case CWTParams::e_PredC:
		wblk.IterateSptC (true, m_Param.m_nWTlevels);
		break;
	default:
		Assert (0, Util::CParamException())
	}

	// VLC coder object creation
	CVLCCoder coder (ac);

	// code the S+P coefs
	coder.Code (wblk, m_Param.m_nWTlevels, m_Param.m_nLossyBitPlanes);

	// Stop Arithmetic coding
	ac.Stop ();

	// Write the Footer marker
	m_Buf.write_marker (c_MarkerFOOTER);

	COMP_CATCHTHIS
}

void CWTCoder::CodeBuffer (void)
{
	COMP_TRYTHIS
	switch (m_Param.m_BlockMode)
	{
      case CWTParams::e_16x16Block:
		  Assert (m_Param.m_nWTlevels <= 4, Util::CParamException());
		  CodeBufferBlock (16U);
		  break;
      case CWTParams::e_32x32Block:
		  Assert (m_Param.m_nWTlevels <= 5, Util::CParamException());
		  CodeBufferBlock (32U);
		  break;
      case CWTParams::e_64x64Block:
		  // Assert not really needed... 
		  Assert (m_Param.m_nWTlevels <= 6, Util::CParamException());
		  CodeBufferBlock (64U);
		  break;
      case CWTParams::e_FullWidth:
		  CodeBufferFull ();
		  break;
      default:
		  Assert (0, Util::CParamException())
	}
	m_Buf.close ();
	COMP_CATCHTHIS
}


} // end namespace
