/*
 * scene.cc: screen sample generation and image reconstruction
 *
 * This file is part of Pyrit Ray Tracer.
 *
 * Copyright 2008  Radek Brich
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */

#include <math.h>

#include "common.h"
#include "scene.h"

void DefaultSampler::init()
{
	phase = 0;
	packetable = (subsample <= 1);
}

int DefaultSampler::initSampleSet()
{
	static const int gridsamples[] = {1,5,9,16};
	const int samples = gridsamples[oversample];
	const int &w = pixmap.getWidth(), &h = pixmap.getHeight();
	Float *&buffer = pixmap.getFloatData();

	if ( phase == 0 )
	{
		if (subsample > 1)
		{
			phase = 1;
			sx = -1;
			return (w/subsample+1)*(h/subsample+1);
		}
		else
		{
			phase = 2;
			sx = -1;
			return w*h*samples;
		}
	}
	if ( phase == 1 )
	{
		// finalize subsampling
		const Float subsample2 = 1.0f / (subsample*subsample);
		int num_samples = 0;
		Colour ic;
		phase = 2;
		sx = -1;
		for (int y = 0; y < h/subsample; y++)
			for (int x = 0; x < w/subsample; x++)
			{
				int x1 = x*subsample;
				int y1 = y*subsample;
				int x2 = (x+1)*subsample;
				int y2 = (y+1)*subsample;
				if (x2 > w-1)	x2 = w-1;
				if (y2 > h-1)	y2 = h-1;
				if (x1 == x2 || y1 == y2)
					continue;
				Float *p;
				p = buffer + 3*(y1*w + x1);
				Colour c1(*p, *(p+1), *(p+2));
				p = buffer + 3*(y1*w + x2);
				Colour c2(*p, *(p+1), *(p+2));
				p = buffer + 3*(y2*w + x1);
				Colour c3(*p, *(p+1), *(p+2));
				p = buffer + 3*(y2*w + x2);
				Colour c4(*p, *(p+1), *(p+2));
				Float m = (c1-c2).mag2();
				m = max(m, (c2-c3).mag2());
				m = max(m, (c3-c4).mag2());
				m = max(m, (c4-c1).mag2());
				if (m < 0.002)
				{
					// interpolate
					for (int i = 0; i < subsample; i++)
						for (int j = 0; j < subsample; j++)
						{
							ic = c1*(Float)((subsample-i)*(subsample-j)*subsample2)
								+ c2*(Float)((i)*(subsample-j)*subsample2)
								+ c3*(Float)((subsample-i)*(j)*subsample2)
								+ c4*(Float)((i)*(j)*subsample2);
							p = buffer + 3*((y1+j)*w + x1+i);
							*(p + 0) = ic.r;
							*(p + 1) = oversample ? -ic.g : ic.g;
							*(p + 2) = ic.b;
						}
				}
				else
				{
					// mark as to be computed
					num_samples += subsample * subsample;
					for (int i = 0; i < subsample; i++)
						for (int j = 0; j < subsample; j++)
							if (oversample || i != 0 || j != 0)
								*(buffer + 3*((y1+j)*w + x1+i)) = -1.;
				}
			}
		return num_samples;
	}
	if ( phase == 2 && oversample )
	{
		// finalize oversampling
		Float *buf;
		if (subsample > 1)
			for (buf = buffer; buf != buffer + w*h*3; buf += 3)
				if (*(buf+1) < 0)
				{
					// interpolated
					*(buf+1) = -*(buf+1);
				}
				else
				{
					*buf = *buf * (1.0f/samples);
					*(buf+1) = *(buf+1) * (1.0f/samples);
					*(buf+2) = *(buf+2) * (1.0f/samples);
				}
		else
			for (buf = buffer; buf != buffer + w*h*3; buf++)
				*buf = *buf * (1.0f/samples);
	}
	phase = -1;
	return 0;
}

bool DefaultSampler::nextSample(Sample* s)
{
	const int &w = pixmap.getWidth(), &h = pixmap.getHeight();
	Float *&buffer = pixmap.getFloatData();

	if (phase == 1)
	{
		// subsampling
		if (sx < 0)
		{
			// first sample
			s->x = -(Float)w/h/2.0f;
			s->y = -0.5f;
			sx = 0;
			sy = 0;
			osa_samp = 0;
		}
		else
		{
			if (sx == w-1)
			{
				if (sy == h-1)
					return false;
				sy += subsample;
				if (sy > h-1)
					sy = h-1;
				sx = 0;
			}
			else
			{
				sx += subsample;
				if (sx > w-1)
					sx = w-1;
			}

			s->x = (Float)sx/h - (Float)w/h/2.0f;
			s->y = (Float)sy/h - 0.5f;
		}
	}
	else if (phase == 2)
	{
		/* grid oversampling */
		static const int gridsamples[] = {1,4,9,16};
		static const Float osa4x[] = {-0.25f, +0.25f, +0.25f, -0.25f};
		static const Float osa4y[] = {-0.25f, -0.25f, +0.25f, +0.25f};
		static const Float osa9x[] = {-0.34f,  0.00f, +0.34f,
			-0.34f,  0.00f, +0.34f, -0.34f,  0.00f, +0.34f};
		static const Float osa9y[] = {-0.34f, -0.34f, -0.34f,
				0.00f,  0.00f,  0.00f, +0.34f, +0.34f, +0.34f};
		static const Float osa16x[] = {-0.375f, -0.125f, +0.125f, +0.375f,
			-0.375f, -0.125f, +0.125f, +0.375f, -0.375f, -0.125f, +0.125f, +0.375f,
			-0.375f, -0.125f, +0.125f, +0.375f};
		static const Float osa16y[] = {-0.375f, -0.375f, -0.375f, -0.375f,
			-0.125f, -0.125f, -0.125f, -0.125f, +0.125f, +0.125f, +0.125f, +0.125f,
			+0.375f, +0.375f, +0.375f, +0.375f};
		static const Float *osaSx[] = {NULL, osa4x, osa9x, osa16x};
		static const Float *osaSy[] = {NULL, osa4y, osa9y, osa16y};
		const int samples = gridsamples[oversample];
		const Float *osax = osaSx[oversample];
		const Float *osay = osaSy[oversample];

		if (sx < 0)
		{
			// first sample
			s->x = -(Float)w/h/2.0f;
			s->y = -0.5f;
			sx = 0;
			sy = 0;
			osa_samp = 0;
		}
		else
		{
			osa_samp++;

			if (oversample && oversample <= 3 && osa_samp < samples)
			{
				s->x = osax[osa_samp]/h + (Float)sx/h - (Float)w/h/2.0f;
				s->y = osay[osa_samp]/h + (Float)sy/h - 0.5f;
			}
			else
			{
				if (subsample > 1)
				{
					// find next not interpolated pixel
					do
					{
						sx++;
						if (sx >= w)
						{
							sx = 0;
							sy++;
						}
						if (sy >= h)
							return false;
					}
					while ( *(buffer + 3*(sy*w + sx)) >= 0. );
				}
				else if (!oversample && !(w&1) && !(h&1))
				{
					// generate good raster for packet tracing
					const int j = ((sy&1)<<1) + (sx&1);
					switch (j)
					{
						case 0:
						case 2:
							sx++;
							break;
						case 1:
							sx--;
							sy++;
							break;
						case 3:
							sx++;
							if (sx >= w)
							{
								sx = 0;
								sy++;
							}
							else
								sy--;
							if (sy >= h)
								return false;
							break;
					}
				}
				else
				{
					sx++;
					if (sx >= w)
					{
						sx = 0;
						sy++;
					}
					if (sy >= h)
						return false;
				}

				s->x = (Float)sx/h - (Float)w/h/2.0f;
				s->y = (Float)sy/h - 0.5f;
				osa_samp = 0;
			}
		}

		if (osa_samp == 0 && oversample && oversample <= 3)
		{
			s->x += osax[0]/h;
			s->y += osay[0]/h;
			Float *buf = buffer + 3*(sy*w + sx);
			*(buf++) = 0;
			*(buf++) = 0;
			*(buf++) = 0;
		}
	}

	s->sx = sx;
	s->sy = sy;
	s->osa_samp = osa_samp;

	return true;
}

void DefaultSampler::saveSample(const Sample &samp, const Colour &col)
{
	Float *buf = pixmap.getFloatData()
		+ 3*(samp.sy * pixmap.getWidth() + samp.sx);
	if (phase == 2 && oversample)
	{
		*(buf+0) += col.r;
		*(buf+1) += col.g;
		*(buf+2) += col.b;
	}
	else
	{
		*(buf++) = col.r;
		*(buf++) = col.g;
		*(buf++) = col.b;
	}
}