/**
 *    Copyright (C) 2016-2021 Grok Image Compression Inc.
 *
 *    This source code is free software: you can redistribute it and/or  modify
 *    it under the terms of the GNU Affero General Public License, version 3,
 *    as published by the Free Software Foundation.
 *
 *    This source code is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU Affero General Public License for more details.
 *
 *    You should have received a copy of the GNU Affero General Public License
 *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
#pragma once

#include "grk_includes.h"
#include <stdexcept>
#include <algorithm>

/*
 Various coordinate systems are used to describe regions in the tile component buffer.

 1) Canvas coordinates:  JPEG 2000 global image coordinates.

 2) Tile component coordinates: canvas coordinates with sub-sampling applied

 3) Band coordinates: coordinates relative to a specified sub-band's origin

 4) Buffer coordinates: coordinate system where all resolutions are translated
	to common origin (0,0). If each code block is translated relative to the origin of the
 resolution that **it belongs to**, the blocks are then all in buffer coordinate system

 Note: the name of any method or variable returning non canvas coordinates is appended
 with "REL", to signify relative coordinates.

 */

namespace grk
{
enum eSplitOrientation
{
	SPLIT_L,
	SPLIT_H,
	SPLIT_NUM_ORIENTATIONS
};

/**
 * Class: ResWindowBuffer
 *
 * Manage all buffers for a single windowed DWT resolution. This class
 * stores a buffer for the resolution, and also buffers for the 4 sub-bands
 * generated by DWT transform.
 *
 */
template<typename T>
struct ResWindowBuffer
{
	ResWindowBuffer(uint8_t numresolutions, uint8_t resno,
					grkBuffer2d<T, AllocatorAligned>* resWindowTopLevelREL,
					Resolution* tileCompAtRes, Resolution* tileCompAtLowerRes,
					grkRectU32 tileCompWindow, grkRectU32 tileCompWindowUnreduced,
					grkRectU32 tileCompUnreduced, uint32_t FILTER_WIDTH)
		: m_allocated(false), m_tileCompRes(tileCompAtRes), m_tileCompResLower(tileCompAtLowerRes),
		  m_resWindowBufferREL(new grkBuffer2d<T, AllocatorAligned>(tileCompWindow.width(),
																	tileCompWindow.height())),
		  m_resWindowBufferTopLevelREL(resWindowTopLevelREL), m_filterWidth(FILTER_WIDTH)
	{
		for(uint32_t i = 0; i < SPLIT_NUM_ORIENTATIONS; ++i)
			m_resWindowBufferSplitREL[i] = nullptr;
		// windowed decompression
		if(FILTER_WIDTH)
		{
			uint32_t numDecomps =
				(resno == 0) ? (uint32_t)(numresolutions - 1U) : (uint32_t)(numresolutions - resno);

			/*
			m_bandWindowPadded is used for determining which precincts and code blocks overlap
			the window of interest, in each respective resolution
			*/
			for(uint8_t orient = 0; orient < ((resno) > 0 ? BAND_NUM_ORIENTATIONS : 1); orient++)
			{
				m_bandWindowPadded.push_back(getBandWindow(numDecomps, orient,
														   tileCompWindowUnreduced,
														   tileCompUnreduced, 2 * FILTER_WIDTH));
			}

			if(m_tileCompResLower)
			{
				assert(resno > 0);
				for(uint8_t orient = 0; orient < BAND_NUM_ORIENTATIONS; orient++)
				{
					// todo: should only need padding equal to FILTER_WIDTH, not 2*FILTER_WIDTH
					auto bandWindow = getBandWindow(numDecomps, orient, tileCompWindowUnreduced,
													tileCompUnreduced, 2 * FILTER_WIDTH);
					auto bandFull = orient == BAND_ORIENT_LL ? *((grkRectU32*)m_tileCompResLower)
															 : m_tileCompRes->tileBand[orient - 1];
					auto bandWindowREL =
						bandWindow.pan(-(int64_t)bandFull.x0, -(int64_t)bandFull.y0);
					m_bandWindowBufferPaddedREL.push_back(
						new grkBuffer2d<T, AllocatorAligned>(&bandWindowREL));
				}
				auto winLow = m_bandWindowBufferPaddedREL[BAND_ORIENT_LL];
				auto winHigh = m_bandWindowBufferPaddedREL[BAND_ORIENT_HL];
				m_resWindowBufferREL->x0 =
					(std::min<uint32_t>)(2 * winLow->x0, 2 * winHigh->x0 + 1);
				m_resWindowBufferREL->x1 =
					(std::max<uint32_t>)(2 * winLow->x1, 2 * winHigh->x1 + 1);
				winLow = m_bandWindowBufferPaddedREL[BAND_ORIENT_LL];
				winHigh = m_bandWindowBufferPaddedREL[BAND_ORIENT_LH];
				m_resWindowBufferREL->y0 =
					(std::min<uint32_t>)(2 * winLow->y0, 2 * winHigh->y0 + 1);
				m_resWindowBufferREL->y1 =
					(std::max<uint32_t>)(2 * winLow->y1, 2 * winHigh->y1 + 1);

				// todo: shouldn't need to clip
				auto resBounds = grkRectU32(0, 0, m_tileCompRes->width(), m_tileCompRes->height());
				m_resWindowBufferREL->clip(&resBounds);

				// two windows formed by horizontal pass and used as input for vertical pass
				grkRectU32 splitResWindowREL[SPLIT_NUM_ORIENTATIONS];

				splitResWindowREL[SPLIT_L] = grkRectU32(
					m_resWindowBufferREL->x0, m_bandWindowBufferPaddedREL[BAND_ORIENT_LL]->y0,
					m_resWindowBufferREL->x1, m_bandWindowBufferPaddedREL[BAND_ORIENT_LL]->y1);

				m_resWindowBufferSplitREL[SPLIT_L] =
					new grkBuffer2d<T, AllocatorAligned>(&splitResWindowREL[SPLIT_L]);

				splitResWindowREL[SPLIT_H] = grkRectU32(
					m_resWindowBufferREL->x0,
					m_bandWindowBufferPaddedREL[BAND_ORIENT_LH]->y0 + m_tileCompResLower->height(),
					m_resWindowBufferREL->x1,
					m_bandWindowBufferPaddedREL[BAND_ORIENT_LH]->y1 + m_tileCompResLower->height());

				m_resWindowBufferSplitREL[SPLIT_H] =
					new grkBuffer2d<T, AllocatorAligned>(&splitResWindowREL[SPLIT_H]);
			}
			// compression or full tile decompression
		}
		else
		{
			// dummy LL band window
			m_bandWindowBufferPaddedREL.push_back(new grkBuffer2d<T, AllocatorAligned>(0, 0));
			assert(tileCompAtRes->numTileBandWindows == 3 || !tileCompAtLowerRes);
			if(m_tileCompResLower)
			{
				for(uint32_t i = 0; i < tileCompAtRes->numTileBandWindows; ++i)
				{
					auto b = tileCompAtRes->tileBand + i;
					m_bandWindowBufferPaddedREL.push_back(
						new grkBuffer2d<T, AllocatorAligned>(b->width(), b->height()));
				}
				// note: only dimensions of split resolution window buffer matter, not actual
				// coordinates
				for(uint32_t i = 0; i < SPLIT_NUM_ORIENTATIONS; ++i)
					m_resWindowBufferSplitREL[i] = new grkBuffer2d<T, AllocatorAligned>(
						tileCompWindow.width(), tileCompWindow.height() / 2);
			}
		}
	}
	~ResWindowBuffer()
	{
		delete m_resWindowBufferREL;
		for(auto& b : m_bandWindowBufferPaddedREL)
			delete b;
		for(uint32_t i = 0; i < SPLIT_NUM_ORIENTATIONS; ++i)
			delete m_resWindowBufferSplitREL[i];
	}
	bool alloc(bool clear)
	{
		if(m_allocated)
			return true;

		// if top level window is present, then all buffers attach to this window
		if(m_resWindowBufferTopLevelREL)
		{
			// ensure that top level window is allocated
			if(!m_resWindowBufferTopLevelREL->alloc2d(clear))
				return false;

			// don't allocate bandWindows for windowed decompression
			if(m_filterWidth)
				return true;

			// attach to top level window
			if(m_resWindowBufferREL != m_resWindowBufferTopLevelREL)
				m_resWindowBufferREL->attach(m_resWindowBufferTopLevelREL->getBuffer(),
											 m_resWindowBufferTopLevelREL->stride);

			// m_tileCompResLower is null for lowest resolution
			if(m_tileCompResLower)
			{
				for(uint8_t orientation = 0; orientation < m_bandWindowBufferPaddedREL.size();
					++orientation)
				{
					switch(orientation)
					{
						case BAND_ORIENT_HL:
							m_bandWindowBufferPaddedREL[orientation]->attach(
								m_resWindowBufferTopLevelREL->getBuffer() +
									m_tileCompResLower->width(),
								m_resWindowBufferTopLevelREL->stride);
							break;
						case BAND_ORIENT_LH:
							m_bandWindowBufferPaddedREL[orientation]->attach(
								m_resWindowBufferTopLevelREL->getBuffer() +
									m_tileCompResLower->height() *
										m_resWindowBufferTopLevelREL->stride,
								m_resWindowBufferTopLevelREL->stride);
							break;
						case BAND_ORIENT_HH:
							m_bandWindowBufferPaddedREL[orientation]->attach(
								m_resWindowBufferTopLevelREL->getBuffer() +
									m_tileCompResLower->width() +
									m_tileCompResLower->height() *
										m_resWindowBufferTopLevelREL->stride,
								m_resWindowBufferTopLevelREL->stride);
							break;
						default:
							break;
					}
				}
				m_resWindowBufferSplitREL[SPLIT_L]->attach(
					m_resWindowBufferTopLevelREL->getBuffer(),
					m_resWindowBufferTopLevelREL->stride);
				m_resWindowBufferSplitREL[SPLIT_H]->attach(
					m_resWindowBufferTopLevelREL->getBuffer() +
						m_tileCompResLower->height() * m_resWindowBufferTopLevelREL->stride,
					m_resWindowBufferTopLevelREL->stride);
			}
		}
		else
		{
			// resolution window is always allocated
			if(!m_resWindowBufferREL->alloc2d(clear))
				return false;

			// band windows are allocated if present
			for(auto& b : m_bandWindowBufferPaddedREL)
			{
				if(!b->alloc2d(clear))
					return false;
			}
			if(m_tileCompResLower)
			{
				m_resWindowBufferSplitREL[SPLIT_L]->attach(m_resWindowBufferREL->getBuffer(),
														   m_resWindowBufferREL->stride);
				m_resWindowBufferSplitREL[SPLIT_H]->attach(m_resWindowBufferREL->getBuffer() +
															   m_tileCompResLower->height() *
																   m_resWindowBufferREL->stride,
														   m_resWindowBufferREL->stride);
			}
		}
		m_allocated = true;

		return true;
	}

	/**
	 * Get band window (in tile component coordinates) for specified number
	 * of decompositions
	 *
	 * Note: if numDecomps is zero, then the band window (and there is only one)
	 * is equal to the unreduced tile component window
	 *
	 * See table F-1 in JPEG 2000 standard
	 *
	 */
	static grkRectU32 getBandWindow(uint32_t numDecomps, uint8_t orientation,
									grkRectU32 tileCompWindowUnreduced)
	{
		assert(orientation < BAND_NUM_ORIENTATIONS);
		if(numDecomps == 0)
			return tileCompWindowUnreduced;

		uint32_t tcx0 = tileCompWindowUnreduced.x0;
		uint32_t tcy0 = tileCompWindowUnreduced.y0;
		uint32_t tcx1 = tileCompWindowUnreduced.x1;
		uint32_t tcy1 = tileCompWindowUnreduced.y1;

		/* Map above tile-based coordinates to
		 * sub-band-based coordinates, i.e. origin is at tile origin  */
		/* See equation B-15 of the standard. */
		uint32_t bx0 = orientation & 1;
		uint32_t by0 = (uint32_t)(orientation >> 1U);

		uint32_t bx0Shift = (1U << (numDecomps - 1)) * bx0;
		uint32_t by0Shift = (1U << (numDecomps - 1)) * by0;

		return grkRectU32(
			(tcx0 <= bx0Shift) ? 0 : ceildivpow2<uint32_t>(tcx0 - bx0Shift, numDecomps),
			(tcy0 <= by0Shift) ? 0 : ceildivpow2<uint32_t>(tcy0 - by0Shift, numDecomps),
			(tcx1 <= bx0Shift) ? 0 : ceildivpow2<uint32_t>(tcx1 - bx0Shift, numDecomps),
			(tcy1 <= by0Shift) ? 0 : ceildivpow2<uint32_t>(tcy1 - by0Shift, numDecomps));
	}
	/**
	 * Get band window (in tile component coordinates) for specified number
	 * of decompositions (with padding)
	 *
	 * Note: if numDecomps is zero, then the band window (and there is only one)
	 * is equal to the unreduced tile component window (with padding)
	 */
	static grkRectU32 getBandWindow(uint32_t numDecomps, uint8_t orientation,
									grkRectU32 unreducedTileCompWindow,
									grkRectU32 unreducedTileComp, uint32_t padding)
	{
		assert(orientation < BAND_NUM_ORIENTATIONS);
		if(numDecomps == 0)
		{
			assert(orientation == 0);
			return unreducedTileCompWindow.grow(padding).intersection(&unreducedTileComp);
		}
		auto oneLessDecompWindow = unreducedTileCompWindow;
		auto oneLessDecompTile = unreducedTileComp;
		if(numDecomps > 1)
		{
			oneLessDecompWindow = getBandWindow(numDecomps - 1, 0, unreducedTileCompWindow);
			oneLessDecompTile = getBandWindow(numDecomps - 1, 0, unreducedTileComp);
		}

		return getBandWindow(
			1, orientation, oneLessDecompWindow.grow(2 * padding).intersection(&oneLessDecompTile));
	}

	bool m_allocated;

	Resolution* m_tileCompRes; // non-null will trigger creation of band window buffers
	Resolution* m_tileCompResLower; // null for lowest resolution

	std::vector<grkBuffer2d<T, AllocatorAligned>*> m_bandWindowBufferPaddedREL;
	std::vector<grkRectU32> m_bandWindowPadded;

	grkBuffer2d<T, AllocatorAligned>* m_resWindowBufferSplitREL[SPLIT_NUM_ORIENTATIONS];
	grkBuffer2d<T, AllocatorAligned>* m_resWindowBufferREL;
	grkBuffer2d<T, AllocatorAligned>* m_resWindowBufferTopLevelREL;

	uint32_t m_filterWidth;
};

template<typename T>
struct TileComponentWindowBuffer
{
	TileComponentWindowBuffer(bool isCompressor, bool lossless, bool wholeTileDecompress,
							  grkRectU32 tileCompUnreduced, grkRectU32 tileCompReduced,
							  grkRectU32 unreducedTileCompOrImageCompWindow,
							  Resolution* tileCompResolution, uint8_t numresolutions,
							  uint8_t reducedNumResolutions)
		: m_unreducedBounds(tileCompUnreduced), m_bounds(tileCompReduced),
		  m_numResolutions(numresolutions), m_compress(isCompressor),
		  m_wholeTileDecompress(wholeTileDecompress)
	{
		if(!m_compress)
		{
			// for decompress, we are passed the unreduced image component window
			auto unreducedImageCompWindow = unreducedTileCompOrImageCompWindow;
			m_bounds = unreducedImageCompWindow.rectceildivpow2(
				(uint32_t)(m_numResolutions - reducedNumResolutions));
			m_bounds = m_bounds.intersection(tileCompReduced);
			assert(m_bounds.is_valid());
			m_unreducedBounds = unreducedImageCompWindow.intersection(tileCompUnreduced);
			assert(m_unreducedBounds.is_valid());
		}
		// fill resolutions vector
		assert(reducedNumResolutions > 0);
		for(uint32_t resno = 0; resno < reducedNumResolutions; ++resno)
			m_resolution.push_back(tileCompResolution + resno);

		auto tileCompAtRes = tileCompResolution + reducedNumResolutions - 1;
		auto tileCompAtLowerRes =
			reducedNumResolutions > 1 ? tileCompResolution + reducedNumResolutions - 2 : nullptr;
		// create resolution buffers
		auto topLevel = new ResWindowBuffer<T>(
			numresolutions, (uint8_t)(reducedNumResolutions - 1U), nullptr, tileCompAtRes,
			tileCompAtLowerRes, m_bounds, m_unreducedBounds, tileCompUnreduced,
			wholeTileDecompress ? 0 : getFilterPad<uint32_t>(lossless));
		// setting top level prevents allocation of tileCompBandWindows buffers
		if(!useBandWindows())
			topLevel->m_resWindowBufferTopLevelREL = topLevel->m_resWindowBufferREL;

		for(uint8_t resno = 0; resno < reducedNumResolutions - 1; ++resno)
		{
			// resolution window ==  next resolution band window at orientation 0
			auto resDims = ResWindowBuffer<T>::getBandWindow((uint32_t)(numresolutions - 1 - resno),
															 0, m_unreducedBounds);
			m_resWindowBufferREL.push_back(new ResWindowBuffer<T>(
				numresolutions, resno, useBandWindows() ? nullptr : topLevel->m_resWindowBufferREL,
				tileCompResolution + resno, resno > 0 ? tileCompResolution + resno - 1 : nullptr,
				resDims, m_unreducedBounds, tileCompUnreduced,
				wholeTileDecompress ? 0 : getFilterPad<uint32_t>(lossless)));
		}
		m_resWindowBufferREL.push_back(topLevel);
	}
	~TileComponentWindowBuffer()
	{
		for(auto& b : m_resWindowBufferREL)
			delete b;
	}

	/**
	 * Transform code block offsets from canvas coordinates
	 * to either band coordinates (relative to sub band origin)
	 * or buffer coordinates (relative to associated resolution origin)
	 *
	 * @param resno resolution number
	 * @param orientation band orientation {LL,HL,LH,HH}
	 * @param offsetx x offset of code block in canvas coordinates
	 * @param offsety y offset of code block in canvas coordinates
	 *
	 */
	void toRelativeCoordinates(uint8_t resno, eBandOrientation orientation, uint32_t& offsetx,
							   uint32_t& offsety) const
	{
		assert(resno < m_resolution.size());

		auto res = m_resolution[resno];
		auto band = res->tileBand + getBandIndex(resno, orientation);

		uint32_t x = offsetx;
		uint32_t y = offsety;

		// get offset relative to band
		x -= band->x0;
		y -= band->y0;

		if(useBufferCoordinatesForCodeblock() && resno > 0)
		{
			auto resLower = m_resolution[resno - 1U];

			if(orientation & 1)
				x += resLower->width();
			if(orientation & 2)
				y += resLower->height();
		}
		offsetx = x;
		offsety = y;
	}
	/**
	 * Get code block destination window
	 *
	 * @param resno resolution number
	 * @param orientation band orientation {LL,HL,LH,HH}
	 *
	 */
	const grkBuffer2d<T, AllocatorAligned>*
		getCodeBlockDestWindowREL(uint8_t resno, eBandOrientation orientation) const
	{
		return (useBufferCoordinatesForCodeblock())
				   ? getResWindowBufferHighestREL()
				   : getBandWindowBufferPaddedREL(resno, orientation);
	}
	/**
	 * Get padded band window buffer
	 *
	 * @param resno resolution number
	 * @param orientation band orientation {0,1,2,3} for {LL,HL,LH,HH} band windows
	 *
	 * If resno is > 0, return LL,HL,LH or HH band window, otherwise return LL resolution window
	 *
	 */
	const grkBuffer2d<T, AllocatorAligned>*
		getBandWindowBufferPaddedREL(uint8_t resno, eBandOrientation orientation) const
	{
		assert(resno < m_resolution.size());
		assert(resno > 0 || orientation == BAND_ORIENT_LL);

		if(resno == 0 && (m_compress || m_wholeTileDecompress))
			return m_resWindowBufferREL[0]->m_resWindowBufferREL;

		return m_resWindowBufferREL[resno]->m_bandWindowBufferPaddedREL[orientation];
	}

	/**
	 * Get padded band window
	 *
	 * @param resno resolution number
	 * @param orientation band orientation {0,1,2,3} for {LL,HL,LH,HH} band windows
	 *
	 */
	const grkRectU32* getBandWindowPadded(uint8_t resno, eBandOrientation orientation) const
	{
		if(m_resWindowBufferREL[resno]->m_bandWindowPadded.empty())
			return nullptr;
		return &m_resWindowBufferREL[resno]->m_bandWindowPadded[orientation];
	}
	/*
	 * Get intermediate split window
	 *
	 * @param orientation 0 for upper split window, and 1 for lower split window
	 */
	const grkBuffer2d<T, AllocatorAligned>*
		getResWindowBufferSplitREL(uint8_t resno, eSplitOrientation orientation) const
	{
		assert(resno > 0 && resno < m_resolution.size());

		return m_resWindowBufferREL[resno]->m_resWindowBufferSplitREL[orientation];
	}
	/**
	 * Get resolution window
	 *
	 * @param resno resolution number
	 *
	 */
	const grkBuffer2d<T, AllocatorAligned>* getResWindowBufferREL(uint32_t resno) const
	{
		return m_resWindowBufferREL[resno]->m_resWindowBufferREL;
	}
	/**
	 * Get highest resolution window
	 *
	 *
	 */
	grkBuffer2d<T, AllocatorAligned>* getResWindowBufferHighestREL(void) const
	{
		return m_resWindowBufferREL.back()->m_resWindowBufferREL;
	}
	bool alloc()
	{
		for(auto& b : m_resWindowBufferREL)
		{
			if(!b->alloc(!m_compress))
				return false;
		}

		return true;
	}
	/**
	 * Get bounds of tile component (canvas coordinates)
	 * decompress: reduced canvas coordinates of window
	 * compress: unreduced canvas coordinates of entire tile
	 */
	grkRectU32 bounds() const
	{
		return m_bounds;
	}
	grkRectU32 unreducedBounds() const
	{
		return m_unreducedBounds;
	}
	uint64_t stridedArea(void) const
	{
		return getResWindowBufferHighestREL()->stride * m_bounds.height();
	}

	// set data to buf without owning it
	void attach(T* buffer, uint32_t stride)
	{
		getResWindowBufferHighestREL()->attach(buffer, stride);
	}
	// transfer data to buf, and cease owning it
	void transfer(T** buffer, uint32_t* stride)
	{
		getResWindowBufferHighestREL()->transfer(buffer, stride);
	}

  private:
	bool useBandWindows() const
	{
		return !m_wholeTileDecompress;
	}
	bool useBufferCoordinatesForCodeblock() const
	{
		return m_compress || !m_wholeTileDecompress;
	}
	uint8_t getBandIndex(uint8_t resno, eBandOrientation orientation) const
	{
		uint8_t index = 0;
		if(resno > 0)
		{
			index = (uint8_t)orientation;
			index--;
		}
		return index;
	}
	/******************************************************/
	// decompress: unreduced/reduced image component window
	// compress:  unreduced/reduced tile component
	grkRectU32 m_unreducedBounds;
	grkRectU32 m_bounds;
	/******************************************************/

	std::vector<Resolution*> m_resolution;
	// windowed bounds for windowed decompress, otherwise full bounds
	std::vector<ResWindowBuffer<T>*> m_resWindowBufferREL;

	// unreduced number of resolutions
	uint8_t m_numResolutions;

	bool m_compress;
	bool m_wholeTileDecompress;
};

} // namespace grk
