#include "reflection.h"
#include "pbrt.h"
#include "material.h"
#include "reflection.h"

#include "reflection.h"
#include "paramset.h"
#include "texture.h"
#include "shape.h"
#include "mc.h"
#include "bessel.h"

#include <iostream>
using namespace std;

#define LAMBDA_MIN (400e-7f)
#define LAMBDA_MAX (800e-7f)

float Bessel( int order, float x )
{
  return besseljn(order, x);
}

float absf( float x )
{
  if( x < 0 ) 
    x *= -1;
  return x;
}

float sinc(float x)
{
  if( absf(x) <= 1e-5 )
    return 1;
  else
    return sin(M_PI*x) / (M_PI*x);
}

float tent( float A, float B )
{
   float v = 1 - absf( A/B );
   if( v < 0 )
     return 0;
   return v;
}

Vector wavelength2color( float lambda, float intensity )
{
   Vector rgb(0,0,0);
   rgb[0] = tent( (650e-7-lambda), 100e-7 ) * intensity / 2;
   rgb[1] = tent( (550e-7-lambda), 100e-7 ) * intensity / 2;
   rgb[2] = tent( (450e-7-lambda), 100e-7 ) * intensity / 2;
   return rgb;
}

class SinGratingTransmission : public BxDF {
public:
    SinGratingTransmission(const DifferentialGeometry& geom,
    		float m, float period, float scale);
    Spectrum f(const Vector &, const Vector &) const;

    Spectrum Sample_f(const Vector &wo, Vector *wi,
                      float u1, float u2, float *pdf) const;

    float Pdf(const Vector &wo, const Vector &wi) const;

private:
    // SpecularReflection Private Data
    float mM;
    float mPeriod;
    float mScale;

    DifferentialGeometry geom;
};

class SinGratingReflection : public BxDF {
public:
    // SpecularReflection Public Methods
    SinGratingReflection(const DifferentialGeometry& geom,
    		float m, float period, float scale );
    Spectrum f(const Vector &, const Vector &) const;

    Spectrum Sample_f(const Vector &wo, Vector *wi,
                      float u1, float u2, float *pdf) const;

    float Pdf(const Vector &wo, const Vector &wi) const;

private:
    // SpecularReflection Private Data
    float mM;
    float mPeriod;
    float mScale;

    DifferentialGeometry geom;
};


class SinGratingMaterial : public Material {
public:
    // MeasuredMaterial Public Methods
    SinGratingMaterial(float amp, float period, float scale,
    		const Reference<Texture<Spectrum> >& Kr,
    		const Reference<Texture<Spectrum> >& Kd,
    		const Reference<Texture<Spectrum> >& Kt,
    		bool reflection, bool transmission );
    BSDF *GetBSDF(const DifferentialGeometry &dgGeom,
                  const DifferentialGeometry &dgShading ) const;
private:
    // MeasuredMaterial Private Data
    float mAmp;
    float mPeriod;
    float mScale;
    Reference<Texture<Spectrum> > mKr;
    Reference<Texture<Spectrum> > mKd;
    Reference<Texture<Spectrum> > mKt;

    bool mReflection;
    bool mTransmission;
};


// theta = ]-Pi/2,Pi/2[, x = [0.0,1.0[
static Spectrum eval(float theta, float x, float mM, float p, float scale) {
	//x = (0.5f - x) * scale;
    float colors[3] = { 0.0f, 0.0f, 0.0f };

	static bool init = false;
	static float bessels[121][121];
	if( !init ) {
		for(int i = -60; i <= 60; i++)
			for(int j = -60; j <= 60; j++)
				bessels[i + 60][j + 60] = Bessel(i, mM / 2) * Bessel(j, mM / 2);
		init = true;
	}

	for(int n = -5; n <= 5; n++) {
		for(int m = -5; m <= 5; m++) {
			float lambda = 0.0f;
			const float r = bessels[n + 60][m + 60];

	        Vector rgb;
	        if( (m+n == 0) && (fabsf(theta) < 5e-4f) ) {
	        	rgb = Vector( r, r, r );
	        } else if( (m + n != 0) && (fabsf(theta) > 5e-4f) ) {
	          lambda = 2 * theta * p / (m+n);
	          if( lambda < LAMBDA_MIN || lambda > LAMBDA_MAX )
	        	  continue;

	          rgb = wavelength2color( lambda, r );
	        } else
	        	continue;

			float phase = 2 * M_PI * (m - n) * x / p;
			float cp = cosf(phase);
			colors[0] += rgb[0] * cp;
			colors[1] += rgb[1] * cp;
			colors[2] += rgb[2] * cp;
		}
	}

	Spectrum F(colors);
	return F;
}

static void dump_table(FILE *file, float scale, float m, float p,
						unsigned int theta_samples, unsigned int x_samples) {
	fprintf(file, "# scale %f amplitude %f period %f\n", scale, m, p);
	fprintf(file, "%u %u\n", theta_samples, x_samples);
	for(unsigned int i = 0; i < theta_samples; ++i) {
		const float theta = acosf( (float(i) / (theta_samples - 1)) );
		for(unsigned int j = 0; j < x_samples; ++j) {
			const float x = (float(j) / x_samples);
			Spectrum S = eval(theta, x, m, p, scale);
			S.Print(file);
		}
		fprintf(file, "\n");
	}
}

SinGratingTransmission::SinGratingTransmission(const DifferentialGeometry& g, float m, float period, float scale )
        : BxDF(BxDFType(BSDF_SPECULAR | BSDF_GLOSSY | BSDF_TRANSMISSION )),
           mM( m ), mPeriod( period ), mScale(scale), geom( g ) {
    }

Spectrum SinGratingTransmission::f(const Vector &wo, const Vector& wi) const {
	if( SameHemisphere(wo, wi) )
		return Spectrum(0.0f);

#ifdef SAME_Y
	if( (wi.y + wo.y) > 1e-2f )
		return Spectrum(0.0f);
#endif

	const float li = sqrtf(1.0f - (wi.y*wi.y));
	const float lr = sqrtf(1.0f - (wo.y*wo.y));
	const float cd = ((-wi.x * wo.x) + (-wi.z * wo.z)) / (li * lr);
	float theta = acosf( Clamp(cd, -1.0f, 1.0f) );

	return eval(theta, geom.u, mM, mPeriod, mScale);
}

Spectrum SinGratingTransmission::Sample_f(const Vector &wo, Vector *wi,
		float u1, float u2, float *pdf) const {

	// sample theta in [0,Pi/2]
	float theta = (powf(u1,5) * M_PI*0.5f);
	if( u2 < 0.5f ) theta = -theta;

	// Projected reflection direction
	const Vector wr( -wo.x, -wo.y, -wo.z );

	// Rotate wr_p around Y, theta rads
	const float costheta = cosf(theta);
	const float sintheta = sinf(theta);
	*wi = Vector(costheta * wr.x - sintheta * wr.z, wr.y,
				sintheta * wr.x + costheta * wr.z);

	*pdf = 1.0f / M_PI;

	if( SameHemisphere(wo, *wi) ) {
		*pdf = 0.0f;
		return Spectrum(0.0f);
	}

   Spectrum S = f(wo, *wi);
   return S;
}

float SinGratingTransmission::Pdf(const Vector &wo, const Vector &wi) const {
	if( SameHemisphere(wo, wi) )
		return 0.0f;

#ifdef SAME_Y
	if( (wi.y + wo.y) > 1e-2f )
		return 0.0f;
#endif

	return 1.0f / M_PI;
}

SinGratingReflection::SinGratingReflection(const DifferentialGeometry& g,
		float m, float period, float scale )
        : BxDF(BxDFType(BSDF_REFLECTION | BSDF_SPECULAR | BSDF_GLOSSY )),
           mM( m ), mPeriod( period ), mScale(scale),
           geom( g ) {
    }


Spectrum SinGratingReflection::f(const Vector &wo, const Vector& wi) const {
	if( !SameHemisphere(wo, wi) )
		return Spectrum(0.0f);

#ifdef SAME_Y
	if( (wi.y + wo.y) > 1e-2f )
		return Spectrum(0.0f);
#endif

	const float li = sqrtf(1.0f - (wi.y*wi.y));
	const float lr = sqrtf(1.0f - (wo.y*wo.y));
	const float cd = ((-wi.x * wo.x) + (wo.z * wi.z)) / (li * lr);
	float theta = acosf( Clamp(cd, -1.0f, 1.0f) );

	return eval(theta, geom.u, mM, mPeriod, mScale);
}

Spectrum SinGratingReflection::Sample_f(const Vector &wo, Vector *wi, float u1,
		float u2, float *pdf) const {

	// sample theta in [-Pi/2,Pi/2]
	float theta = (powf(u1,5) * M_PI*0.5f);
	if( u2 < 0.5f ) theta = -theta;

	// Projected reflection direction
	const Vector wr( -wo.x, -wo.y, wo.z );

	// Rotate wr_p around Y, theta rads
	const float costheta = cosf(theta);
	const float sintheta = sinf(theta);
	*wi = Vector(costheta * wr.x - sintheta * wr.z, wr.y,
				sintheta * wr.x + costheta * wr.z);

	*pdf = 1.0f / M_PI;

	if( !SameHemisphere(wo, *wi) ) {
		*pdf = 0.0f;
		return Spectrum(0.0f);
	}

	Spectrum S = f(wo, *wi);
   return S;
}

float SinGratingReflection::Pdf(const Vector &wo, const Vector &wi) const {
	if( !SameHemisphere(wo, wi) )
		return 0.0f;

#ifdef SAME_Y
	if( (wi.y + wo.y) > 1e-2f )
		return 0.0f;
#endif

	return 1.0f / M_PI;
}

SinGratingMaterial::SinGratingMaterial(float amp, float period, float scale,
		const Reference<Texture<Spectrum> >& Kr, const Reference<Texture<
				Spectrum> >& Kd, const Reference<Texture<Spectrum> >& Kt,
		bool reflection, bool transmission) :
	mAmp(amp), mPeriod(period), mScale(scale), mKr(Kr), mKd(Kd), mKt(Kt), mReflection(
			reflection), mTransmission(transmission) {
	if( scale <= 0 || scale > 1 )
		printf("scale should be between 0 < scale <= 1\n");
	if( amp <= 0 || amp > 1 )
		printf("amp should be between 0 < amp <= 1\n");
}


BSDF* SinGratingMaterial::GetBSDF(const DifferentialGeometry &dgGeom,
                  const DifferentialGeometry &dgShading ) const
{
    BSDF *bsdf = BSDF_ALLOC(BSDF)(dgGeom, dgGeom.nn);

	if( mTransmission )
		bsdf->Add( BSDF_ALLOC(SinGratingTransmission)(dgGeom, mAmp, mPeriod, mScale) );
    if( mReflection )
      bsdf->Add( BSDF_ALLOC(SinGratingReflection)(dgGeom, mAmp, mPeriod, mScale) );

    Spectrum R = mKr->Evaluate(dgShading).Clamp();

    if (!R.Black())
      bsdf->Add(BSDF_ALLOC(SpecularReflection)(R, BSDF_ALLOC(FresnelNoOp)()));

    R = mKd->Evaluate(dgShading).Clamp();
    if( !R.Black() )
      bsdf->Add(BSDF_ALLOC(Lambertian)(R));

    Spectrum T = mKt->Evaluate(dgShading).Clamp();
    if (!T.Black())
		bsdf->Add(BSDF_ALLOC(SpecularTransmission)(T, 1., 1.));

    return bsdf;
}

extern "C" DLLEXPORT Material* CreateMaterial( const Transform &xform,
        const TextureParams &mp )
{
   float mScale = mp.FindFloat("scale", 1.0f);
   float mAmp = mp.FindFloat("amplitude", 0.1f);
   float mPeriod = mp.FindFloat("period", 0.0001f);
   Reference<Texture<Spectrum> > Kr = mp.GetSpectrumTexture("Kr", Spectrum(.2f));
   Reference<Texture<Spectrum> > Kd = mp.GetSpectrumTexture("Kd", Spectrum(.5f));
   Reference<Texture<Spectrum> > Kt = mp.GetSpectrumTexture("Kt", Spectrum(.0f));
   bool ref = mp.FindBool("reflection", true);
   bool trans = mp.FindBool("transmission", true);

   if( mp.FindBool("generate-table", false) ) {
	   printf("Generating table\n");
	   FILE *file = fopen("singrating.dat", "w+");
	   unsigned int theta_samples = 4096;
	   unsigned int x_samples = 2048;
	   dump_table(file, mScale, mAmp, mPeriod, theta_samples, x_samples);
	   fclose(file);
   }

   return new SinGratingMaterial( mAmp, mPeriod, mScale, Kr, Kd, Kt, ref, trans );
}
