/*
 * 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 "CBuffer.h"

#include "CWBlock.h"
#include "CACDecoder.h"
#include "CVLCDecoder.h"
#include "CWTDecoder.h"

namespace COMP 
{


void CWTDecoder::ZeroBlock (unsigned short i_from_line, 
						    unsigned short i_from_column,
						    unsigned short i_to_line, 
						    unsigned short i_to_column,
						    unsigned short i_Bs)
{
	COMP_TRYTHIS
	unsigned short bX = i_from_column;
	unsigned short bY = i_from_line;

	unsigned short nBW = (m_Image.GetW() + i_Bs - 1) / i_Bs;	// nb blocks per line
	unsigned short nBH = (m_Image.GetH() + i_Bs - 1) / i_Bs;	// nb blocks per column
	unsigned long  nB  = (unsigned long) nBH * nBW;				// total number of blocks

	CWBlock wblk (i_Bs, i_Bs);
	wblk.Zero();
	for(; bX<=i_to_column || bY<i_to_line ;)
	{
		if (bX * i_Bs >= m_Image.GetW())
		{
			bX = 0;
			bY ++;
		}
		if (bY > i_to_line) // go out of here 
			break;
//		m_Image.put_block (pixels_block, corner_column, corner_line);
		unsigned int nW = ((m_Image.GetW() - bX*i_Bs)< i_Bs) ? (m_Image.GetW() - bX*i_Bs) : i_Bs;  
		unsigned int nH = ((m_Image.GetH() - bY*i_Bs)< i_Bs) ? (m_Image.GetH() - bY*i_Bs) : i_Bs;
		wblk.Put (m_Image, bX * i_Bs, bY * i_Bs, nW, nH);

		bX ++;
	}
	COMP_CATCHTHIS
}
//-----------------------------------------------------------------------
short CWTDecoder::FindNextMarker (void)
{
	COMP_TRYTHIS
	// resynchronization with the next marker (or at least with the EOF)
	unsigned short code;

	m_Buf.byteAlign ();
	for (;;)
		if (m_Buf.read_marker (code))
		{
			if (code >= c_MarkerRESTART && code <= (c_MarkerRESTART + 15))
			return (code & 0x000F);
			if (code == c_MarkerFOOTER)
				return -2;		// must return a negative value. -2 is nice for debugging.
			m_Buf.seek (8);
		}
		else
		{
			if (m_Buf.reached_end ()) 
				return -1;
			else 
				m_Buf.seek (8);
		}
	COMP_CATCHTHIS
}
//-----------------------------------------------------------------------
bool CWTDecoder::PerformResync (unsigned int i_Bs,
								 unsigned int & t_markerNum,
								 unsigned int & t_nbBlock,
								 unsigned int & t_bX,
								 unsigned int & t_bY)
{
	COMP_TRYTHIS
	// Initialize things...
	//---Just to avoid passing these parameters...
	unsigned short nBW = (m_Image.GetW() + i_Bs - 1) / i_Bs;	// nb blocks per line
	unsigned short nBH = (m_Image.GetH() + i_Bs - 1) / i_Bs;	// nb blocks per column
	unsigned long  nB  = (unsigned long) nBH * nBW;				// total number of blocks

	unsigned int begin_bY	= 0;
	unsigned int new_bX		= 0;
	unsigned int new_bY		= 0;
	// Find the next restart marker
	short new_RSTm = FindNextMarker();

	// Found new marker
	begin_bY = (t_markerNum * m_Param.m_RestartInterval)/nBW;
	if (new_RSTm < 0) // no usable marker found
	{
		new_bX		= nBW;
		new_bY		= nBH-1;
		if (m_Param.m_RestartInterval > 0)
		{
			t_nbBlock	= nB % m_Param.m_RestartInterval;
			t_markerNum	= nB / m_Param.m_RestartInterval;
		}
		else
		{
			// not needed since already null
			//t_nbBlock	= 0;
			//t_markerNum	= 0;
		}
	}
	else
	{
		// a valid restart marker was found
		new_RSTm = new_RSTm - (t_markerNum % 16);
		// if new_RSTm == 0  --> assume we didn't miss 16 markers but that we found the marker we missed before
		t_markerNum += new_RSTm;
		unsigned long new_nB = (t_markerNum + 1) * m_Param.m_RestartInterval;
		if (new_nB > nB)
			new_nB = nB;
		t_nbBlock = m_Param.m_RestartInterval;
		new_bX = (new_nB-1) % nBW;
		new_bY = (new_nB-1) / nBW;
	}
	ZeroBlock( t_bY, t_bX, new_bY, new_bX, i_Bs);
	m_Qinfo.Negate (begin_bY * i_Bs, cmin (t_bY*i_Bs + i_Bs, m_Image.GetH())-1);
	if ( t_bY < new_bY)
		m_Qinfo.Zero( cmin (t_bY  * i_Bs + i_Bs, m_Image.GetH()),
		              cmin (new_bY* i_Bs + i_Bs, m_Image.GetH()));

	// modify the returned data
	t_bX = new_bX;
	t_bY = new_bY;

	m_LastGoodLine = cmin (new_bY* i_Bs + i_Bs, m_Image.GetH());

	if (new_RSTm < 0 )
		return false;	// do not eat the marker, it is the EOI marker
	else
		return true;
	COMP_CATCHTHIS
}
//-----------------------------------------------------------------------
bool CWTDecoder::DecodeBufferBlock (const unsigned int i_Bs)
{
	COMP_TRYTHIS
	unsigned int bX, bY;
	
	// Compute 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++;
	
	// Start the arithmetic decoder
	CACDecoder ad (m_Buf);
	ad.Start ();

	// VLC decoder object creation
	CVLCDecoder decoder (ad);

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

	// block loops
	unsigned int nbBlock	= 0;		// relative block index
	unsigned int markerNum	= 0;		// absolute block index
	bool goodline			= true;		// true if current line is ok from the beginning
	bool bACStopped = false;
	for (bY=0 ; bY<nbH ; bY++)
	{
		goodline = true;
		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;			

			// Decode the stream
			if (! decoder.Decode (wblk, m_Param.m_nWTlevels, m_Param.m_nLossyBitPlanes) 
				|| ad.IsMarkerReached())
			{
				// Decode failed and/or a marker was reached during decoding -> error
				if (goodline)
				{
					// Set quality info
					m_LastGoodLine = bY * i_Bs + nH;
					m_Qinfo.Set (bY * i_Bs, m_LastGoodLine - 1, (-1) * (int)bX * (int)i_Bs);
				}
				goodline = false;
				PerformResync(i_Bs, markerNum, nbBlock, bX, bY);
			}
			else
			{				
				switch (m_Param.m_PredMode)
				{
				case CWTParams::e_None:
					wblk.IterateSt (false, m_Param.m_nWTlevels);
					break;
				case CWTParams::e_PredA:
					wblk.IterateSptA (false, m_Param.m_nWTlevels);
					break;
				case CWTParams::e_PredB:
					wblk.IterateSptB (false, m_Param.m_nWTlevels);
					break;
				case CWTParams::e_PredC:
					wblk.IterateSptC (false, m_Param.m_nWTlevels);
					break;
				default:
					Assert (0, Util::CParamException())
				}
				// Copy CWBlock to image
				wblk.Put (m_Image, bX * i_Bs, bY * i_Bs, nW, nH);
				nbBlock++;
			}

			if (m_Param.m_RestartInterval && nbBlock == m_Param.m_RestartInterval)
			{
				unsigned short marker;

				// Check if the Restart marker is there
				nbBlock = 0;
				ad.Stop ();
				if (! m_Buf.read_marker (marker) || marker != (c_MarkerRESTART | (markerNum & 0xF)))
				{
					// Restart marker not found

					// Perform resynchronization
					if (goodline)
					{
						// Set quality info
						m_LastGoodLine = bY * i_Bs + nH;
						m_Qinfo.Set (bY * i_Bs, m_LastGoodLine - 1, (-1) * (int)(bX * i_Bs + nW));
					}
					goodline = false;
					bX++;	// the current block is complete and shouldn't be reset
					if (PerformResync(i_Bs, markerNum, nbBlock, bX, bY))
						m_Buf.seek (16);
					nbBlock = 0;
				}
				else 
				{
					// Restart marker found

					m_Buf.seek (16);
				}
				markerNum++;
				if (bX < (nbW - 1) || bY < (nbH - 1)) ad.Start ();
				else bACStopped = true;
				// Reset VLC decoder state (models probabilities)
				decoder.Reset ();
			}
		}
		if (goodline)
		{
			// Set quality info
			m_LastGoodLine = bY * i_Bs + nH;
			m_Qinfo.Set (bY * i_Bs, m_LastGoodLine - 1, m_Image.GetW());
		}
		goodline = true;
	}

	// Stop the arithmetic decoder
	if (! bACStopped) ad.Stop ();
	
	if (m_LastGoodLine < m_Image.GetH())
		DecodeBufferError (m_LastGoodLine, m_Image.GetH() );
	else
	{
		unsigned short marker = 0;
		// Check if Footer marker is there
		if (! m_Buf.read_marker (marker) || marker != c_MarkerFOOTER)
		{

			// FOOTER marker not found

			unsigned short nBW = (m_Image.GetW() + i_Bs - 1) / i_Bs;	// nb blocks per line
			unsigned short begin_bY = (markerNum * m_Param.m_RestartInterval)/nBW;
			
			m_Qinfo.Negate (begin_bY * i_Bs, 
							m_Image.GetH()-1);
		}
		m_Buf.seek (16);
	}

	return true;
	COMP_CATCHTHIS
}

bool CWTDecoder::DecodeBufferFull (void)
{
	COMP_TRYTHIS
	// Start the arithmetic decoder
	CACDecoder ad (m_Buf);
	ad.Start ();

	// Compute CWBlock size to allow m_Param.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);

	// VLC decoder object creation
	CVLCDecoder decoder (ad);

	// Decode stream
	if (! decoder.Decode (wblk, m_Param.m_nWTlevels, m_Param.m_nLossyBitPlanes) 
		|| ad.IsMarkerReached())
		// Decode failed and/or a marker was reached during decoding -> error
		return false;
	
	// Perform the inverse m_nWTlevels iterations of 2D S(+P) transform
	switch (m_Param.m_PredMode)
	{
	case CWTParams::e_None:
		wblk.IterateSt (false, m_Param.m_nWTlevels);
		break;
	case CWTParams::e_PredA:
		wblk.IterateSptA (false, m_Param.m_nWTlevels);
		break;
	case CWTParams::e_PredB:
		wblk.IterateSptB (false, m_Param.m_nWTlevels);
		break;
	case CWTParams::e_PredC:
		wblk.IterateSptC (false, m_Param.m_nWTlevels);
		break;
	default:
		Assert (0, Util::CParamException())
	}
	
	// Copy the Wavelet block coefs to image
	wblk.Put (m_Image, 0, 0, m_Image.GetW(), m_Image.GetH());

	// Stop the arithmetic decoder
	ad.Stop ();

	// Set quality info
	m_Qinfo.Set (0, m_Image.GetH() - 1, m_Image.GetW());
	m_LastGoodLine = m_Image.GetH();

	unsigned short marker = 0;
	// Check if Footer marker is there (there is no restart marker in full-image mode)
	if (! m_Buf.read_marker (marker) || marker != c_MarkerFOOTER)
	{

		// FOOTER marker not found

		m_Qinfo.Negate (0, m_Image.GetH()-1);
	}
	m_Buf.seek (16);

	return true;
	COMP_CATCHTHIS
}

void CWTDecoder::DecodeBuffer (void)
{
	COMP_TRYTHIS
	unsigned short marker;

	// reset the reading function
	m_Buf.real_rewind ();

	// Check if the Header marker is there
	if (! m_Buf.read_marker (marker) || marker != c_MarkerHEADER)
	{

		// HEADER marker not found

		DecodeBufferError (0, m_Image.GetH() );
		return;
	}
	m_Buf.real_seek (16);

	// Original image number of bits per pixel
	m_Param.m_BitsPerPixel = m_Buf.readN (4);
	m_Buf.real_seek (4);
	if (m_Param.m_BitsPerPixel == 0) m_Param.m_BitsPerPixel = 16;

	// Image width
	const unsigned int imgW = m_Buf.readN (16);
	m_Buf.real_seek (16);
	if (imgW != m_Image.GetW())
	{

		// Bad image width in header

		DecodeBufferError (0, m_Image.GetH() );
		return;
	}

	// Image height
	const unsigned int imgH = m_Buf.readN (16);
	m_Buf.real_seek (16);
	if (imgH != m_Image.GetH())
	{

		// Bad image height in header

		DecodeBufferError (0, m_Image.GetH() );
		return;
	}

	// Number of S+P transform iterations
	m_Param.m_nWTlevels = m_Buf.readN (2) + 3;
	m_Buf.real_seek (2);
	
	// S+P prediction type
	m_Param.m_PredMode = (CWTParams::EWTPredictionMode)(CWTParams::e_None + m_Buf.readN (2));
	m_Buf.real_seek (2);

	// Block coding type
	const unsigned int blockType = m_Buf.readN (2);
	m_Buf.real_seek (2);

	// Restart marker interval in number of blocks
	m_Param.m_RestartInterval = m_Buf.readN (16);
	m_Buf.real_seek (16);

	// Lossy coding parameter
	m_Param.m_nLossyBitPlanes = m_Buf.readN (4);
	m_Buf.real_seek (4);

	// Skip 2 pad bits to be sure next marker dont need align
	m_Buf.real_seek (2);

	// Check if the Data marker is there
	if (! m_Buf.read_marker (marker) || marker != c_MarkerDATA)
	{

		// DATA marker not found

		DecodeBufferError (0, m_Image.GetH() );
		return;
	}
	m_Buf.real_seek (16);

	m_Buf.resync ();

	m_LastGoodLine = 0;
	switch (blockType)
	{
	case 0:
		m_Param.m_BlockMode = CWTParams::e_16x16Block;
		if (! DecodeBufferBlock (16U))
			std::cerr << "Can not decode blocky (16x16) buffer" << std::endl;
		break;
	case 1:
		m_Param.m_BlockMode = CWTParams::e_32x32Block;
		if (! DecodeBufferBlock (32U))
			std::cerr << "Can not decode blocky (32x32) buffer" << std::endl;
		break;
	case 2:
		m_Param.m_BlockMode = CWTParams::e_64x64Block;
		if (! DecodeBufferBlock (64U))
			std::cerr << "Can not decode blocky (64x64) buffer" << std::endl;
		break;
	case 3:
		m_Param.m_BlockMode = CWTParams::e_FullWidth;
		if (! DecodeBufferFull ())
			std::cerr << "Can not decode full buffer" << std::endl;
		break;
	}

	COMP_CATCHTHIS
}


} // end namespace
