//By Stefano Jannuzzo: sjannuz@tiscali.it
//ONW adapted from Alan Jones public shader

#include <math.h>
#include <string.h>
#include <shader.h>

#define SHADER_NAME "sj_multi_mat"
#define SHADER_VERSION 1

//tiny helpers
#define sj_rgba_set(res, r, g, b, a) \
	(res)->r = (r),\
	(res)->g = (g),\
	(res)->b = (b),\
	(res)->a = (a)

#define sj_rgb_set(res, r, g, b, a) \
	(res)->r = (r),\
	(res)->g = (g),\
	(res)->b = (b)

#define sj_rgba_black(res) \
	(res)->r = (res)->g = (res)->b = (res)->a = 0.0f

#define sj_rgb_black(res) \
	(res)->r = (res)->g = (res)->b = 0.0f

#define sj_rgba_white(res) \
	(res)->r = (res)->g = (res)->b = (res)->a = 1.0f

#define sj_rgb_white(res) \
	(res)->r = (res)->g = (res)->b = 1.0f

#define sj_rgb_add(res, c0, c1) \
	(res)->r = (c0)->r + (c1)->r,\
	(res)->g = (c0)->g + (c1)->g,\
	(res)->b = (c0)->b + (c1)->b

#define sj_rgba_add(res, c0, c1) \
	sj_rgb_add(res, c0, c1),\
	(res)->a = (c0)->a + (c1)->a

#define sj_rgb_add_in_place(res, c0) \
	sj_rgb_add(res, res, c0)

#define sj_rgba_add_in_place(res, c0) \
	sj_rgba_add(res, res, c0)

#define sj_rgb_mul(res, c0, c1) \
	(res)->r = (c0)->r * (c1)->r,\
	(res)->g = (c0)->g * (c1)->g,\
	(res)->b = (c0)->b * (c1)->b

#define sj_rgba_mul(res, c0, c1) \
	sj_rgb_mul(res, c0, c1),\
	(res)->a = (c0)->a * (c1)->a

#define sj_rgb_mul_in_place(res, c0) \
	sj_rgb_mul(res, res, c0)

#define sj_rgba_mul_in_place(res, c0) \
	sj_rgba_mul(res, res, c0)

#define sj_rgb_smul(res, c0, s) \
	(res)->r = (c0)->r * (s),\
	(res)->g = (c0)->g * (s),\
	(res)->b = (c0)->b * (s)

#define sj_rgba_smul(res, c0, s) \
	sj_rgb_smul,\
	(res)->a = (c0)->a * (s)

#define sj_rgb_smul_in_place(res, s) \
	sj_rgb_smul(res, res, s)

#define sj_rgba_smul_in_place(res, s) \
	sj_rgba_smul(res, res, s)

#define sj_rgb_sdiv(res, c0, s) \
	(res)->r = (c0)->r / (s),\
	(res)->g = (c0)->g / (s),\
	(res)->b = (c0)->b / (s)

#define sj_rgba_sdiv(res, c0, s) \
	sj_rgb_sdiv,\
	(res)->a = (c0)->a / (s)

#define sj_rgb_sdiv_in_place(res, s) \
	sj_rgb_sdiv(res, res, s)

#define sj_rgba_sdiv_in_place(res, s) \
	sj_rgba_sdiv(res, res, s)


#define MAX_NB_LIGHTS 10 //must match the spdl in/out structs

typedef struct {
	miInteger	ll_i_lights; //lights list
	miInteger ll_n_lights;
	miTag			ll_lights[1];
	miInteger dModel[MAX_NB_LIGHTS];
	miInteger sModel[MAX_NB_LIGHTS];
	miColor   diffuseColor, specularColor;
	miScalar  onwIor, onwRoughness;
	miScalar  shinyness, roughness, blinnIor;
	miColor   ctIor;
} Parameters;


typedef struct {
	miColor dsSum, dSum, sSum; //dsSum = dSum + sSum
	miColor d[MAX_NB_LIGHTS], s[MAX_NB_LIGHTS]; //d and s for 10 lights
}	Output;


extern "C" {
DLLEXPORT void sj_multi_mat_init(miState*, Parameters*, miBoolean*);
DLLEXPORT void sj_multi_mat_exit(miState*, void*);
DLLEXPORT miBoolean sj_multi_mat(Output*, miState*, Parameters*);
DLLEXPORT int sj_multi_mat_version(void);
}


static void SetupONW(miState *state, miScalar ior, miScalar roughness,
							miScalar *reflAngle, miScalar *reflDot, miVector *reflAzi, 
							miScalar *C1A, miScalar *C1C, miScalar *C2A, miScalar *C3A);
static miScalar OrenNayarWolff(miState *state, miVector *incDir, miScalar ior,
													 miScalar reflAngle, miScalar reflDot, miVector *reflAzi, 
													 miScalar C1A, miScalar C1C, miScalar C2A, miScalar C3A);



DLLEXPORT int sj_multi_mat_version(void) 
{ return SHADER_VERSION; }

DLLEXPORT void sj_multi_mat_init(miState *state, Parameters *oParas, miBoolean *ii_req)
{
	if (!oParas)
	{
		*ii_req = miTRUE;
		mi_info("%s - V. %d", SHADER_NAME, SHADER_VERSION);
	}
}


DLLEXPORT void sj_multi_mat_exit(miState *state, void	*vparas) {}


enum D_MODES
{
	D_OFF=0,
	D_LAMBERT,
	D_ONW
};

enum S_MODES
{
	S_OFF=0,
	S_PHONG,
	S_BLINN,
	S_BLONG,
	S_COOKTORRANCE,
	S_WARD
};


DLLEXPORT miBoolean sj_multi_mat(Output *result, miState *state, Parameters *paras)
{
	memset(result, 0, sizeof(Output)); //all black

	miTag     *lights		= mi_eval_tag(&(paras->ll_lights));
	miInteger i_lights	= *mi_eval_integer(&(paras->ll_i_lights));
	miInteger	n_lights	= *mi_eval_integer(&(paras->ll_n_lights));

	if (n_lights < 1)
		return miFALSE;

	if (n_lights > MAX_NB_LIGHTS)
		n_lights = MAX_NB_LIGHTS;

	miColor   color, lColor, dColor, dResult, sColor;
	miVector  lDir;
	miInteger i, lSamples;
	miScalar  dot_nl, diffuse, specular;

	miInteger dModel[MAX_NB_LIGHTS];
	miInteger sModel[MAX_NB_LIGHTS];

	//setup

	bool doOnw = false, doDiffuse = false, doSpecular = false;
	for (i=0; i<n_lights; i++)
	{
		dModel[i] = *mi_eval_integer(&paras->dModel[i]);
		sModel[i] = *mi_eval_integer(&paras->sModel[i]);
		doDiffuse|= dModel[i] > D_OFF;
		doOnw|= dModel[i] == D_ONW;
		doSpecular|= sModel[i] > S_OFF;
	}

	if (!doDiffuse && !doSpecular)
		return miTRUE;

	miColor  diffuseColor, specularColor, ctIor;
	miScalar shinyness, roughness, blinnIor;
	miScalar onwIor, onwRoughness;
	miScalar reflAngle, reflDot, C1A, C1C, C2A, C3A;
	miVector reflAzi;

	if (doDiffuse)
		diffuseColor = *mi_eval_color(&paras->diffuseColor);
	if (doSpecular)
	{
		specularColor = *mi_eval_color(&paras->specularColor);
		shinyness = *mi_eval_scalar(&paras->shinyness);
		roughness = *mi_eval_scalar(&paras->roughness);
		blinnIor = *mi_eval_scalar(&paras->blinnIor);
		ctIor = *mi_eval_color(&paras->ctIor);
	}
	//onw stuff
	if (doOnw)
	{
		onwIor = *mi_eval_scalar(&paras->onwIor);
		onwRoughness = *mi_eval_scalar(&paras->onwRoughness);
		SetupONW(state, onwIor, onwRoughness,&reflAngle, &reflDot, &reflAzi, 
						 &C1A, &C1C, &C2A, &C3A);
	}

	for (i=0; i<n_lights; i++)
	{
		sj_rgb_black(&dResult);
		sj_rgb_black(&sColor);

		if (dModel[i] == D_OFF && sModel[i] == S_OFF)
			continue;

		lSamples=0;
		while(mi_sample_light(&lColor, &lDir, &dot_nl, state, lights[i + i_lights], &lSamples))
		{
			if (doDiffuse)
			{
				switch (dModel[i])
				{
					case D_LAMBERT:
						diffuse = dot_nl > 0.0f ? dot_nl : 0.0f;
						sj_rgb_smul(&dColor, &diffuseColor, diffuse);
						break;

					case D_ONW:
						diffuse = OrenNayarWolff(state, &lDir, onwIor, reflAngle, reflDot, &reflAzi, 
																		 C1A, C1C, C2A, C3A);
						sj_rgb_smul(&dColor, &diffuseColor, diffuse);
						break;

					default:
						break;
				}
				if (dModel[i] > D_OFF)
					sj_rgb_mul_in_place(&dColor, &lColor); //diffuse color for i-th light
			}

			if (doSpecular)
			{
				switch (sModel[i])
				{
					case S_PHONG:
						specular = mi_phong_specular(shinyness, state, &lDir);
						sj_rgb_smul(&sColor, &specularColor, specular);
						break;

					case S_BLINN:
						specular = mi_blinn_specular(&state->dir, &lDir, &state->normal, roughness, blinnIor);
						sj_rgb_smul(&sColor, &specularColor, specular);
						break;

					case S_BLONG:
						specular = mi_blong_specular(shinyness, state, &lDir);
						sj_rgb_smul(&sColor, &specularColor, specular);
						break;
					
					case S_COOKTORRANCE:
						if (mi_cooktorr_specular(&color, &state->dir, &lDir, &state->normal, roughness, &ctIor))
							sj_rgb_mul(&sColor, &specularColor, &color);
						break;

					case S_WARD:
						specular = mi_ward_glossy(&state->dir, &lDir, &state->normal, shinyness);
						sj_rgb_smul(&sColor, &specularColor, specular);
						break;

					default:
						break;
				}
				if (sModel[i] > S_OFF)
					sj_rgb_mul_in_place(&sColor, &lColor); //specular color for i-th light
			}
		} //while mi_sample_light

		if (lSamples)
		{
			if (doDiffuse)
			{
				sj_rgb_sdiv(&result->d[i], &dColor, (miScalar)lSamples);
				sj_rgb_add_in_place(&result->dSum, &result->d[i]);
				result->d[i].a = result->dSum.a = 1.0f; 
			}
			if (doSpecular)
			{
				sj_rgb_sdiv(&result->s[i], &sColor, (miScalar)lSamples);
				sj_rgb_add_in_place(&result->sSum, &result->s[i]);
				result->s[i].a = result->sSum.a = 1.0f;
			}
		}
	} //for n_lights

	sj_rgb_add(&result->dsSum, &result->dSum, &result->sSum);
	result->dsSum.a = 1.0f;

	return miTRUE;
}


static miScalar fresnelReflection(miScalar angle, miScalar ior) 
{
	miScalar transAngle = asinf(sinf(angle)/ ior);
	return miScalar(0.5f * (powf(sin(angle - transAngle), 2.0f) / powf(sinf(angle + transAngle), 2.0f)) * ( 1 + (pow(cos(angle + transAngle), 2) / pow(cos(angle - transAngle), 2) )));
}

#define C3BPi 0.4052847345693f

static void SetupONW(miState *state, miScalar ior, miScalar roughness,
							miScalar *reflAngle, miScalar *reflDot, miVector *reflAzi, 
							miScalar *C1A, miScalar *C1C, miScalar *C2A, miScalar *C3A)
{
	miScalar roughness2 = roughness * roughness;
	//calculate reflected angle
	miVector reflDir = state->dir;
	mi_vector_neg(&reflDir);
	*reflDot = mi_vector_dot(&state->normal, &reflDir);
	*reflAngle = acosf(*reflDot);

	//calculate reflected azimuth
	reflAzi->x = reflDir.x - (state->normal.x * *reflDot);
	reflAzi->y = reflDir.y - (state->normal.y * *reflDot);
	reflAzi->z = reflDir.z - (state->normal.z * *reflDot);
	mi_vector_normalize(reflAzi);

	*C1A = 1.0f - (0.5f * (roughness2 / (roughness2 + 0.33f)));
	*C1C = 1.0f - fresnelReflection(asinf(sinf(*reflAngle) / ior), 1.0f/ior);
	miScalar C23 = roughness2 / (roughness2 + 0.09f);
	*C2A = 0.45f * C23;
	*C3A = C23 / 8.0f;
}


static miScalar OrenNayarWolff(miState *state, miVector *incDir, miScalar ior,
													 miScalar reflAngle, miScalar reflDot, miVector *reflAzi, 
													 miScalar C1A, miScalar C1C, miScalar C2A, miScalar C3A)
{
	miScalar alpha, beta;
	miVector incAzi;
	miScalar aziDot;

	//calculate incidence angle
	miScalar incDot = mi_vector_dot(&state->normal, incDir);
	miScalar incAngle = acos(incDot);

	//C1B
	miScalar C1B = 1 - fresnelReflection(incAngle, ior);
	miScalar C1 = C1A * C1B * C1C;

	//pick alpha and beta
	if (incAngle > reflAngle)
	{
		alpha = incAngle;
		beta = reflAngle;
	}
	else
	{
		alpha = reflAngle;
		beta = incAngle;
	}

	//calculate indicidence azimuth
	incAzi.x = incDir->x - (state->normal.x * reflDot);
	incAzi.y = incDir->y - (state->normal.y * reflDot);
	incAzi.z = incDir->z - (state->normal.z * reflDot);
	mi_vector_normalize(&incAzi);

	//calculate azimuth dot product
	aziDot = mi_vector_dot(reflAzi, &incAzi);

	miScalar C2B = aziDot >= 0.0f ? sinf(alpha) : sinf(alpha) - powf((2.0f * beta)/(miScalar)M_PI, 3);
	miScalar C3B = powf(alpha * beta * C3BPi, 2);

	//get the illumination info
	miScalar C2 = C2A * C2B;
	miScalar C3 = C3A * C3B;
	miScalar luminance = incDot * (C1 + (aziDot * C2 * tan(beta)) + ((1.0f - fabs(aziDot)) * C3 * tan((alpha + beta) / 2.0f)));
	return luminance; 
}