/*
 * Polygon_RayTrace
 * Copylight (C) 2013 mocchi
 * mocchi_2003@yahoo.co.jp
 * License: Boost ver.1
 */

#include <cstdio>
#include <cmath>
#include <algorithm>

#include "ximage.h"
#include "ximapng.h"

#include "opennurbs.h"

#include "pasync.h"

#include "rapidjson/rapidjson.h"
#include "rapidjson/document.h"
#include "rapidjson/reader.h"
#include "rapidjson/prettywriter.h"
#include "rj_wrapper.h"

#include "PhisicalProperties.h"

#include "MonteCarlo.h"

namespace rj = rapidjson;

#define RAY_IOTA_PROGRESS 0.005
#define TRACE_TERMINAL_LENGTH 10.0

#define THREADS 3
#define THREAD_BLOCK 1000

/// 参考文献: http://homepage2.nifty.com/yees/RayTrace/RayTraceVersion001b.pdf
struct RefractParams{
	double cost1, cos2t1;
	double sin2t1, sin2t2, cos2t2, cost2;
	ON_3dVector incident, nrm; // incident, nrm は必ず単位ベクトル
	double n1, n2, alpha; // alpha = n1 / n2
	int flg; // 1 : n1とn2が同じ 2 : 全反射 0 : それ以外

	// nrm, incident は単位ベクトルであること
	void Reset(const ON_3dVector &nrm_, const ON_3dVector &incident_, double n1_, double n2_){
		n1 = n1_, n2 = n2_;
		nrm = nrm_, incident = incident_;
		if (n1 == n2){
			flg = 1; return;
		}
		alpha = n1 / n2;
		cost1 = ON_DotProduct(nrm_, incident_);
		cos2t1 = cost1 * cost1;
		if (cos2t1 > 1.0) cos2t1 = 1.0;
		sin2t1 = 1.0 - cos2t1;
		sin2t2 = sin2t1 * alpha;
		cos2t2 = 1.0 - sin2t2;
		if (cos2t2 < 0){
			flg = 2; return;
		}
		cost2 = std::sqrt(cos2t2);
		flg = 0;
	}
	RefractParams(){
		n1 = n2 = 0; flg = 0;
		incident.Zero(), nrm.Zero();
	}
	// nrm, incident は単位ベクトルであること
	RefractParams(const ON_3dVector &nrm, const ON_3dVector &incident, double n1_, double n2_){
		Reset(nrm, incident, n1_, n2_);
	}
	bool IsSameIndex() const { return (flg == 1); }
	bool IsTotalReflection() const { return (flg == 2); }
};

bool CalcRefractDir(const RefractParams &rp, ON_3dVector &emit){
	if (rp.IsSameIndex()){
		emit = rp.incident;
		return true;
	}else if (rp.IsTotalReflection()){
		return false;
	}
	double gamma = rp.cost2 - rp.alpha * rp.cost1;
	if (rp.cost1 < 0) gamma = -gamma;
	emit = rp.incident * rp.alpha + rp.nrm * gamma;
	emit.Unitize();
	return true;
}

// 参考文献: http://qiita.com/edo_m18/items/b145f2f5d2d05f0f29c9
bool CalcReflectDir(const ON_3dVector &incident, const ON_3dVector &nrm, ON_3dVector &emit){
	double nrm_part = ON_DotProduct(incident, nrm);
	emit = incident - nrm * (nrm_part * 2.0);
	emit.Unitize();
	return true;
}

struct FresnelParam{
	double tp, rp, ts, rs; // 振幅
};

bool CalcFresnel(const RefractParams &rp, const ON_3dVector &emit, FresnelParam &fp){
	if (rp.IsSameIndex()){
		fp.tp = fp.ts = 1.0;
		fp.rp = fp.rs = 0.0;
		return true;
	}else if (rp.IsTotalReflection()){
		fp.tp = fp.ts = 0.0;
		fp.rp = fp.rs = 1.0;
		return true;
	}
	double n1_cost1 = rp.n1 * rp.cost1, n2_cost1 = rp.n2 * rp.cost1;
	double n1_cost2 = rp.n1 * rp.cost2, n2_cost2 = rp.n2 * rp.cost2;
	fp.tp = n1_cost1 * 2.0        / (n2_cost1 + n1_cost2);
	fp.rp = (n2_cost1 - n1_cost2) / (n2_cost1 + n1_cost2);
	fp.ts = n1_cost1 * 2.0        / (n1_cost1 + n2_cost2);
	fp.rs = (n1_cost1 - n2_cost2) / (n1_cost1 + n2_cost2);
	return true;
}

float get_ieee754(ON__UINT8 p[4]){
	return *reinterpret_cast<float *>(p);
}

// 参考文献: http://ja.wikipedia.org/wiki/Standard_Triangulated_Language
bool ConvertFromBinarySTL(ON_SimpleArray<ON__UINT8> &stl, ON_Mesh &mesh){
	mesh.Destroy();
	if (stl.Count() < 84) return false;
	ON__UINT32 num_facets =
		static_cast<ON__UINT32>(stl[80]) + static_cast<ON__UINT32>(stl[81]) * 256 +
		static_cast<ON__UINT32>(stl[82]) * 65536 + static_cast<ON__UINT32>(stl[83]) * 16777216;

	if (stl.Count() < 84 + static_cast<int>(num_facets) * 50) return false;

	mesh.DestroyDoublePrecisionVertices();
	mesh.SetSinglePrecisionVerticesAsValid();
	ON__UINT8 *p = stl.First() + 84;
	for (ON__UINT32 i = 0; i < num_facets; ++i, p += 50){
		int vcnt = mesh.m_V.Count();
		mesh.m_FN.AppendNew().Set(get_ieee754(p   ), get_ieee754(p+ 4), get_ieee754(p+ 8));
		mesh.m_FN.Last()->Unitize();
		mesh.m_V.AppendNew().Set(get_ieee754(p+12), get_ieee754(p+16), get_ieee754(p+20));
		mesh.m_V.AppendNew().Set(get_ieee754(p+24), get_ieee754(p+28), get_ieee754(p+32));
		mesh.m_V.AppendNew().Set(get_ieee754(p+36), get_ieee754(p+40), get_ieee754(p+44));
		mesh.SetTriangle(mesh.m_F.Count(), vcnt, vcnt+1, vcnt+2);
	}
	mesh.CombineIdenticalVertices();
	mesh.Compact();
	return true;
}

bool ReadFile(const char *filename, ON_SimpleArray<ON__UINT8> &data){
	FILE *fp = std::fopen(filename, "rb");
	if (!fp) return false;
	std::fseek(fp, 0, SEEK_END);
	data.SetCapacity(static_cast<int>(std::ftell(fp)));
	data.SetCount(data.Capacity());
	std::rewind(fp);
	std::fread(data, 1, data.Count(), fp);
	std::fclose(fp);
	return true;
}

struct MeshRayIntersection{
	ON_Mesh *mesh;
	ON_RTree *rtree;
	ON_3dPoint model_center;
	double rough_radius;

	// 計算結果
	ON_RTreeCapsule cc;
	int face_idx;
	void Initialize(ON_Mesh *mesh_, ON_RTree *rtree_){
		mesh = mesh_;
		rtree = rtree_;

		ON_BoundingBox tbb;
		mesh_->GetTightBoundingBox(tbb);
		rough_radius = tbb.Diagonal().Length() * 0.5;
		model_center = tbb.Center();
	}
	bool RayIntersection(const ON_3dRay &ray){
		face_idx = -1;
		for (int i = 0; i < 3; ++i) cc.m_point[0][i] = ray.m_P[i];
		double len_max = rough_radius + ray.m_P.DistanceTo(model_center);
		ON_3dPoint pt = ray.m_P + ray.m_V * len_max;
		for (int i = 0; i < 3; ++i) cc.m_point[1][i] = pt[i];
		cc.m_radius = 0;
		cc.m_domain[0] = 0, cc.m_domain[1] = 1;
		rtree->Search(&cc, IntersectCallback, this);
		return (face_idx >= 0);
	}
	static bool IntersectCallback(void* a_context, ON__INT_PTR a_id){
		MeshRayIntersection *this_ = static_cast<MeshRayIntersection *>(a_context);
		ON_Mesh &m = *this_->mesh;
		ON_MeshFace &f = m.m_F[a_id];
		ON_RTreeCapsule &cc = this_->cc;
		ON_Plane pln(
			ON_3dPoint(m.m_V[f.vi[0]]), ON_3dPoint(m.m_V[f.vi[1]]), ON_3dPoint(m.m_V[f.vi[2]]));
		double t;
		ON_Line lin(cc.m_point[0], cc.m_point[1]);
		ON_Intersect(lin, pln.plane_equation, &t);
		if (t < 0 || t > 1.0) return true;
		ON_3dPoint pt = lin.PointAt(t);
		for (int i = 0; i < 3; ++i) cc.m_point[1][i] = pt[i];
		this_->face_idx = a_id;
		return true;
	}
};

#if 0
struct LightSource{
	// -500, -470, 470 〜 -500, -530, 530 の範囲から(1,-1,0)の方向に照射する平行光
	ON_3dVector dir;
	LightSource(){
		dir.Set(1, 1, -1);
		dir.Unitize();
	}
	int Count() const{ return 80000; }
	void Get(int idx, ON_3dRay &ray, double &flux, double &wavelength) const{
		ray.m_P.x = -500;
		ray.m_P.y = -470 - CalcHaltonSequence(10 + idx, 2) * 60;
		ray.m_P.z = 470 + CalcHaltonSequence(10 + idx, 3) * 60;
		ray.m_V = dir;
		flux = 1.0;
		wavelength = 555;
	}
};

bool GetWavelengthDependentParams(const char *matinfo, double wavelength, double &refl_rate, double &refr_index, double &transmittance){
	
	return true;
}

bool CalcBSDF(const char *bsdf_info, const ON_3dRay &emitray, bool transmitted, int &seed, ON_3dVector &scattered_dir, double &flux){
	scattered_dir = emitray.m_V;
	return true;
}
#endif

struct CommonInfo{
	LightSources *light_src;
	Materials *materials;
	int cnt_10;
	double ref_index;
	pt::mutex mtx_read, mtx_write;

	ON_SimpleArray<int> block_seed;
	void CreateBlockSeed(int num_raies, int blocksize){
		div_t bpr = std::div(num_raies, blocksize);
		if (bpr.rem) bpr.quot++, bpr.rem = 0;
		block_seed.SetCapacity(bpr.quot);
		for (int i = 0; i < bpr.quot; ++i){
			double sd = CalcHaltonSequence(i + 11, 7);
			block_seed.Append(static_cast<int>(sd * num_raies) + 9);
		}
	}

	// read
	int ray_cursor;
	ON_Mesh *cshape;
	ON_RTree *rtree;
	ON_SimpleArray<unsigned int> *shapeidx2fidx;
	ON_SimpleArray<int> *shape2matidx;

	// write
	ON_ClassArray<ON_ClassArray<ON_3dRay> > raies_last;
	ON_ClassArray<ON_SimpleArray<double> > flux_last;
	ON_ClassArray<ON_Polyline> traces;
};

void RayTrace(const ON_3dRay &ray_, double flux, double wavelength, MeshRayIntersection &mri, RefractParams &rp, CommonInfo *ci, bool enable_fresnel, int &block_seed, ON_3dRay &ray_o, double &flux_o, ON_Polyline *trace){
	ON_3dRay ray = ray_;
//	ON_Polyline &trace = traces[k];
	if (trace) trace->Append(ray.m_P);
	bool is_inside = false;
	ON_Mesh &cshape = *mri.mesh;
	while (mri.RayIntersection(ray)){
		// 面番号から形状番号を取得
		int shape_idx = static_cast<int>(
			std::upper_bound(ci->shapeidx2fidx->First(),
				ci->shapeidx2fidx->Last()+1, static_cast<unsigned int>(mri.face_idx)) - ci->shapeidx2fidx->First()
		) - 1;

		double refl_rate, refr_index, transmittance, absorpt_coef;
		ci->materials->GetParams((*ci->shape2matidx)[shape_idx], wavelength, refl_rate, refr_index, transmittance, absorpt_coef);

		// 反射、透過を両方設定してある場合は半々の確率で反射/屈折を選んで、代わりに反射率/透過率を倍にする。
		if (refl_rate > 0 && transmittance > 0){
			double r = CalcHaltonSequence(block_seed, 7);
			if (r >= 0.5) transmittance *= 2.0, refl_rate = 0;
			else transmittance = 0, refl_rate *= 2.0;
		}

		// 材質内部の時は吸収係数を適用
		if (is_inside && absorpt_coef > 0){
			double dist = ray.m_P.DistanceTo(ON_3dPoint(mri.cc.m_point[1]));
			flux *= std::exp(-dist * absorpt_coef);
		}

		bool transmitted = false;
		if (transmittance > 0){
			// 屈折
			flux *= transmittance;
			rp.Reset(cshape.m_FN[mri.face_idx], ray.m_V, is_inside ? refr_index : 1.0, is_inside ? 1.0 : refr_index);
			if (rp.IsTotalReflection()){
				// 全反射
				CalcReflectDir(rp.incident, rp.nrm, ray.m_V);
			}else{
				CalcRefractDir(rp, ray.m_V);
				if (enable_fresnel){
					FresnelParam fp;
					CalcFresnel(rp, ray.m_V, fp);
					double rr = (fp.rp * fp.rp + fp.rs * fp.rs) * 0.5;
					double hs = CalcHaltonSequence(block_seed++, 5);
					if (rr > hs){
						// フレネル反射
						CalcReflectDir(rp.incident, rp.nrm, ray.m_V);
					}else{
						is_inside ^= true;
					}
				}else{
					is_inside ^= true;
				}
				transmitted = true;
			}
		}else if (refl_rate > 0){
			// 反射
			flux *= refl_rate;
			double n12 = is_inside ? refr_index : 1.0;
			rp.Reset(cshape.m_FN[mri.face_idx], ray.m_V, n12, n12);
			CalcReflectDir(rp.incident, rp.nrm, ray.m_V);
		}else{
			// 吸収
			flux = 0;
			ray.m_P = ON_3dPoint(mri.cc.m_point[1]);
			if (trace) trace->Append(mri.cc.m_point[1]);
			break;
		}
		ci->materials->CalcBSDF(ray.m_V, transmitted, block_seed, ray.m_V);
		if (trace) trace->Append(mri.cc.m_point[1]);
		ray.m_P = ON_3dPoint(mri.cc.m_point[1]) + ray.m_V * RAY_IOTA_PROGRESS;
	}
	if (trace) trace->Append(ray.m_P + ray.m_V * TRACE_TERMINAL_LENGTH);
	ray_o = ray;
	flux_o = flux;
}

#include "RenderingTest.cpp"

int main(int argc, char *argv[]){
	if (argc == 1){
		return 0;
	}
	rj::Document args;
	{
		FILE *fp = std::fopen(argv[1], "r");
		std::fseek(fp, 0, SEEK_END);
		ON_SimpleArray<char> settingfile(std::ftell(fp));
		std::rewind(fp);
		settingfile.SetCount(settingfile.Capacity());
		std::fread(settingfile, 1, settingfile.Count(), fp);
		std::fclose(fp);
		args.Parse<0>(settingfile);
	}

	ON_ClassArray<ON_Mesh> shapes;
	ON_SimpleArray<int> shape2matidx;

	// 全てのメッシュの合成
	ON_Mesh cshape;

	// 合成後の面番号から合成前の形状番号を取得するために使用する。
	ON_SimpleArray<unsigned int> shapeidx2fidx;

	std::fprintf(stderr, "Reading shapes.\n");
	rj::Value &shapesv = args["shapes"];
	if (shapesv.IsArray()){
		// とりあえず、すべてBinarySTLとして扱う
		for (rj::SizeType k = 0; k < shapesv.Size(); ++k){
			ON_Mesh &shape = shapes.AppendNew();
			const char *filename = rj_getstring(shapesv[k]["filename"]);
			if (filename[0] == '\0') continue;
			std::fprintf(stderr, "  %s\n", filename);

			ON_SimpleArray<ON__UINT8> data;
			if (!ReadFile(filename, data)) continue;
			ConvertFromBinarySTL(data, shape);
		}
	}

	std::fprintf(stderr, "Constructing tree.\n");
	shapeidx2fidx.Append(0U);
	for (int k = 0; k < shapes.Count(); ++k){
		cshape.Append(shapes[k]);
		shapeidx2fidx.Append(*shapeidx2fidx.Last() + shapes[k].FaceCount());
	}
	ON_RTree rtree_shape;
	rtree_shape.CreateMeshFaceTree(&cshape);

	// 材質の定義
	Materials mats(args["materials"], shapesv, shape2matidx);
	// まずはすべてアクリルで。
	double ref_index = 1.1;

	std::fprintf(stderr, "Defining lightsource.\n");

	// 光源データ
	LightSources light_src(args["lightsources"]);

	std::fprintf(stderr, "Raytracing.\n");

	// 光線追跡
	CommonInfo ci;
	ci.light_src = &light_src;
	ci.materials = &mats;
	ci.cnt_10 = light_src.RayCount() / 10;
	ci.ref_index = ref_index;
	ci.ray_cursor = 0;
	ci.cshape = &cshape;
	ci.rtree = &rtree_shape;
	ci.shapeidx2fidx = &shapeidx2fidx;
	ci.shape2matidx = &shape2matidx;
	ci.CreateBlockSeed(light_src.RayCount(), THREAD_BLOCK);

	{
		ON_ClassArray<ON_ClassArray<ON_3dRay> > &raies_last = ci.raies_last;
		raies_last.SetCapacity(light_src.LSCount());
		raies_last.SetCount(light_src.LSCount());
		ON_ClassArray<ON_SimpleArray<double> > &flux_last = ci.flux_last;
		flux_last.SetCapacity(light_src.LSCount());
		flux_last.SetCount(light_src.LSCount());
		for (int i = 0; i < light_src.LSCount(); ++i){
			raies_last[i].SetCapacity(light_src.RayCount(i));
			raies_last[i].SetCount(raies_last[i].Capacity());
			flux_last[i].SetCapacity(light_src.RayCount(i));
			flux_last[i].SetCount(flux_last[i].Capacity());
		}
	}

	struct ThreadCalculator : public pt::thread{
		MeshRayIntersection mri;
		CommonInfo *ci;
		ThreadCalculator() : pt::thread(false){}
		void execute(){
			ci->mtx_read.enter();
			mri.Initialize(ci->cshape, ci->rtree);

			std::fprintf(stderr, " thread (%d)\n", get_id());
			std::fprintf(stderr, "  models center:%f, %f, %f\n",
				mri.model_center.x, mri.model_center.y, mri.model_center.z);
			std::fprintf(stderr, "  models rough radius:%f\n", mri.rough_radius);
			ci->mtx_read.leave();
			ON_ClassArray<ON_3dRay> raies;
			ON_SimpleArray<int> lsidx;
			ON_SimpleArray<int64_t> rayidx;
			ON_SimpleArray<double> flux, wavelength;
			for(;;){
				ci->mtx_read.enter();
				int cnt = ci->light_src->RayCount() - ci->ray_cursor;
				int ray_cursor = ci->ray_cursor;
				int block_idx = ray_cursor / THREAD_BLOCK;
				if (cnt > 0){
					if (cnt > THREAD_BLOCK) cnt = THREAD_BLOCK;
					if (ci->ray_cursor / ci->cnt_10 < (ci->ray_cursor + cnt) / ci->cnt_10){
						std::fprintf(stderr, "... %d\n", ((ci->ray_cursor + cnt) / ci->cnt_10) * ci->cnt_10);
						std::fflush(stderr);
					}
					ci->ray_cursor += cnt;
				}
				ci->mtx_read.leave();
				if (cnt <= 0) break;

				int &block_seed = ci->block_seed[block_idx];

				raies.SetCapacity(cnt), raies.SetCount(0);
				lsidx.SetCapacity(cnt), lsidx.SetCount(0);
				rayidx.SetCapacity(cnt), rayidx.SetCount(0);
				traces.Destroy();
				for (int i = 0, ir = ray_cursor; i < cnt; ++i, ++ir){
					ci->light_src->Get(ir, block_seed, lsidx.AppendNew(), rayidx.AppendNew(), raies.AppendNew(),
						flux.AppendNew(), wavelength.AppendNew());
				}
				Calc(cnt, raies, flux, wavelength, block_seed);

				ci->mtx_write.enter();
				for (int i = 0; i < cnt; ++i){
					ci->raies_last[lsidx[i]][rayidx[i]] = raies_last[i];
					ci->flux_last[lsidx[i]][rayidx[i]] = flux_last[i];
				}
				ci->traces.Append(traces.Count(), traces);
				ci->mtx_write.leave();
			}
		}

		ON_ClassArray<ON_3dRay> raies_last;
		ON_SimpleArray<double> flux_last;
		ON_ClassArray<ON_Polyline> traces;
		void Calc(int cnt, ON_3dRay *raies_start, double *flux, double *wavelength, int block_seed){
			RefractParams rp;
			raies_last.SetCapacity(cnt), raies_last.SetCount(cnt);
			flux_last.SetCapacity(cnt), flux_last.SetCount(cnt);
			//			trace.SetCapacity(cnt), trace.SetCount(cnt);
			for (int k = 0; k < cnt; ++k){
				RayTrace(raies_start[k], flux[k], wavelength[k], mri, rp, ci, true, block_seed, raies_last[k], flux_last[k], (k % 50) ? 0 : &traces.AppendNew());
			}
		}
	};
	ON_ClassArray<ThreadCalculator> tcs(THREADS);

	for (int i = 0; i < THREADS; ++i){
		ThreadCalculator &tc = tcs.AppendNew();
		tc.ci = &ci;
	}
	for (int i = 0; i < tcs.Count(); ++i){
		tcs[i].start();
	}

	for (int i = 0; i < tcs.Count(); ++i){
		tcs[i].waitfor();
	}

	// Rayの書き出し
	FILE *fp = std::fopen("raies.txt", "w");
	for (int j = 0; j < ci.raies_last.Count(); ++j){
		std::fprintf(fp, "lsidx:%d\n", j);
		ON_ClassArray<ON_3dRay> &ray_last = ci.raies_last[j];
		ON_SimpleArray<double> f_last = ci.flux_last[j];
		for (int i = 0; i < ray_last.Count(); ++i){
			ON_3dRay &r = ray_last[i];
			std::fprintf(fp, " %d\t%f\t%f\t%f\t%f\t%f\t%f\t%f\n",
				i, r.m_P.x, r.m_P.y, r.m_P.z, r.m_V.x, r.m_V.y, r.m_V.z, f_last[i]); 
		}
	}
	std::fclose(fp);

	std::fprintf(stderr, "output models and trace\n");
	ONX_Model model;
	{
		ONX_Model_Object &obj = model.m_object_table.AppendNew();
		obj.m_object = cshape.Duplicate();
		obj.m_bDeleteObject = true;
	}
	for (int i = 0; i < ci.traces.Count(); ++i){
		ONX_Model_Object &obj = model.m_object_table.AppendNew();
		obj.m_object = ON_PolylineCurve(ci.traces[i]).Duplicate();
		obj.m_bDeleteObject = true;
	}
	model.Write("result.3dm", 4);

//	RenderingTest(tcs, ci);
	return 0;
}
