/**
 * @file gust.cpp
 *
 * Computes wind gusts
 */

#include <boost/lexical_cast.hpp>

#include "forecast_time.h"
#include "gust.h"
#include "level.h"
#include "logger.h"
#include "numerical_functions.h"
#include "plugin_factory.h"
#include <boost/thread.hpp>

#include "hitool.h"
#include "neons.h"
#include "radon.h"
#include <NFmiLocation.h>
#include <NFmiMetTime.h>

using namespace std;
using namespace himan;
using namespace himan::plugin;

struct deltaT
{
	vector<double> deltaT_100;
	vector<double> deltaT_200;
	vector<double> deltaT_300;
	vector<double> deltaT_400;
	vector<double> deltaT_500;
	vector<double> deltaT_600;
	vector<double> deltaT_700;

	deltaT() {}
};

struct deltaTot
{
	vector<double> deltaTot_100;
	vector<double> deltaTot_200;
	vector<double> deltaTot_300;
	vector<double> deltaTot_400;
	vector<double> deltaTot_500;
	vector<double> deltaTot_600;
	vector<double> deltaTot_700;

	deltaTot() {}
};

struct intT
{
	vector<double> intT_100;
	vector<double> intT_200;
	vector<double> intT_300;
	vector<double> intT_400;
	vector<double> intT_500;
	vector<double> intT_600;
	vector<double> intT_700;

	intT() {}
};

struct intTot
{
	vector<double> intTot_100;
	vector<double> intTot_200;
	vector<double> intTot_300;
	vector<double> intTot_400;
	vector<double> intTot_500;
	vector<double> intTot_600;
	vector<double> intTot_700;

	intTot() {}
};

void DeltaT(shared_ptr<const plugin_configuration> conf, info_t T_lowestLevel, const forecast_time& ftime,
            const forecast_type& ftype, size_t gridSize, deltaT& dT);
void DeltaTot(deltaTot& dTot, info_t T_lowestLevel, size_t gridSize);
void IntT(intT& iT, const deltaT& dT, size_t gridSize);
void IntTot(intTot& iTot, const deltaTot& dTot, size_t gridSize);
void LowAndMiddleClouds(vector<double>& lowAndMiddleClouds, info_t lowClouds, info_t middleClouds, info_t highClouds,
                        info_t totalClouds);

gust::gust() { itsLogger = logger("gust"); }
void gust::Process(std::shared_ptr<const plugin_configuration> conf)
{
	Init(conf);

	param theRequestedParam("FFG2-MS", 417, 0, 2, 22);

	theRequestedParam.Unit(kMs);

	SetParams({theRequestedParam});

	Start();
}

/*
 * Calculate()

 * This function does the actual calculation.
 */

void gust::Calculate(shared_ptr<info> myTargetInfo, unsigned short threadIndex)
{
	NFmiMetTime theTime(boost::lexical_cast<short>(myTargetInfo->Time().ValidDateTime().String("%Y")),
	                    boost::lexical_cast<short>(myTargetInfo->Time().ValidDateTime().String("%m")),
	                    boost::lexical_cast<short>(myTargetInfo->Time().ValidDateTime().String("%d")),
	                    boost::lexical_cast<short>(myTargetInfo->Time().ValidDateTime().String("%H")),
	                    boost::lexical_cast<short>(myTargetInfo->Time().ValidDateTime().String("%M")));

	auto myThreadedLogger = logger("gust_pluginThread #" + boost::lexical_cast<string>(threadIndex));

	/*
	 * Required source parameters
	 *
	 */
	const size_t gridSize = myTargetInfo->Grid()->Size();

	const param BLHParam("MIXHGT-M");                                     // boundary layer height
	const param WSParam("FF-MS");                                         // wind speed
	const param GustParam("FFG-MS");                                      // wind gust
	const param TParam("T-K");                                            // temperature
	const param TopoParam("Z-M2S2");                                      // geopotential height
	const params LowCloudParam = {param("NL-PRCNT"), param("NL-0TO1")};   // low cloud cover
	const params MidCloudParam = {param("NM-PRCNT"), param("NM-0TO1")};   // middle cloud cover
	const params HighCloudParam = {param("NH-PRCNT"), param("NH-0TO1")};  // high cloud cover
	const params TotalCloudParam = {param("N-PRCNT"), param("N-0TO1")};   // total cloud cover

	level H0, H10;

	producer prod = itsConfiguration->SourceProducer(0);

	HPDatabaseType dbtype = itsConfiguration->DatabaseType();

	long lowestHybridLevelNumber = kHPMissingInt;

	if (dbtype == kNeons || dbtype == kNeonsAndRadon)
	{
		auto n = GET_PLUGIN(neons);

		lowestHybridLevelNumber = boost::lexical_cast<long>(n->ProducerMetaData(prod.Id(), "last hybrid level number"));
	}

	if ((dbtype == kRadon || dbtype == kNeonsAndRadon) && lowestHybridLevelNumber == kHPMissingInt)
	{
		auto r = GET_PLUGIN(radon);

		lowestHybridLevelNumber =
		    boost::lexical_cast<long>(r->RadonDB().GetProducerMetaData(prod.Id(), "last hybrid level number"));
	}

	assert(lowestHybridLevelNumber != kHPMissingInt);

	level lowestHybridLevel(kHybrid, lowestHybridLevelNumber);

	if (myTargetInfo->Producer().Id() == 240 || myTargetInfo->Producer().Id() == 243)
	{
		H0 = level(kGround, 0);
		H10 = H0;
	}
	else
	{
		H0 = level(kHeight, 0);
		H10 = level(kHeight, 10);
	}

	info_t GustInfo, T_LowestLevelInfo, BLHInfo, TopoInfo, LCloudInfo, MCloudInfo, HCloudInfo, TCloudInfo;

	// Current time and level

	forecast_time forecastTime = myTargetInfo->Time();
	level forecastLevel = myTargetInfo->Level();
	forecast_type forecastType = myTargetInfo->ForecastType();

	GustInfo = Fetch(forecastTime, H10, GustParam, forecastType, false);
	T_LowestLevelInfo = Fetch(forecastTime, lowestHybridLevel, TParam, forecastType, false);
	BLHInfo = Fetch(forecastTime, H0, BLHParam, forecastType, false);
	TopoInfo = Fetch(forecastTime, H0, TopoParam, forecastType, false);
	LCloudInfo = Fetch(forecastTime, H0, LowCloudParam, forecastType, false);
	MCloudInfo = Fetch(forecastTime, H0, MidCloudParam, forecastType, false);
	HCloudInfo = Fetch(forecastTime, H0, HighCloudParam, forecastType, false);
	TCloudInfo = Fetch(forecastTime, H0, TotalCloudParam, forecastType, false);

	if (!GustInfo || !T_LowestLevelInfo || !LCloudInfo || !MCloudInfo || !HCloudInfo || !TCloudInfo)
	{
		itsLogger.Error("Unable to find all source data");
		return;
	}

	deltaT dT;
	vector<double> lowAndMiddleClouds(gridSize, kFloatMissing);

	boost::thread t(&DeltaT, itsConfiguration, T_LowestLevelInfo, forecastTime, forecastType, gridSize, boost::ref(dT));
	boost::thread t2(&LowAndMiddleClouds, boost::ref(lowAndMiddleClouds), LCloudInfo, MCloudInfo, HCloudInfo,
	                 TCloudInfo);

	// calc boundary layer height
	vector<double> z_boundaryl(gridSize, kFloatMissing);
	vector<double> z_one_third_boundaryl(gridSize, kFloatMissing);
	vector<double> z_two_third_boundaryl(gridSize, kFloatMissing);
	vector<double> z_zero(gridSize, 0);
	for (size_t i = 0; i < gridSize; ++i)
	{
		if (BLHInfo->Data()[i] == kFloatMissing)
		{
			continue;
		}

		z_boundaryl[i] = BLHInfo->Data()[i];

		if (BLHInfo->Data()[i] >= 200)
		{
			z_boundaryl[i] = 0.5 * BLHInfo->Data()[i] + 100;

			if (BLHInfo->Data()[i] > 1200)
			{
				z_boundaryl[i] = 700;
			}
		}

		z_one_third_boundaryl[i] = z_boundaryl[i] / 3;
		z_two_third_boundaryl[i] = 2 * z_one_third_boundaryl[i];
	}

	// maybe need adjusting
	auto h = GET_PLUGIN(hitool);

	h->Configuration(itsConfiguration);
	h->Time(forecastTime);
	h->ForecastType(forecastType);

	vector<double> maxEstimate;
	try
	{
		maxEstimate = h->VerticalAverage(WSParam, z_two_third_boundaryl, z_boundaryl);
	}

	catch (const HPExceptionType& e)
	{
		if (e != kFileDataNotFound)
		{
			throw runtime_error("Caught exception " + boost::lexical_cast<string>(e));
		}
		else
		{
			myThreadedLogger.Error("hitool was unable to find data");
			t.join();
			t2.join();
			return;
		}
	}

	vector<double> BLtop_ws;
	try
	{
		BLtop_ws = h->VerticalAverage(WSParam, z_one_third_boundaryl, z_boundaryl);
	}

	catch (const HPExceptionType& e)
	{
		if (e != kFileDataNotFound)
		{
			throw runtime_error("Caught exception " + boost::lexical_cast<string>(e));
		}
		else
		{
			myThreadedLogger.Error("hitool was unable to find data");
			t.join();
			t2.join();
			return;
		}
	}

	vector<double> meanWind;
	try
	{
		meanWind = h->VerticalAverage(WSParam, z_zero, z_boundaryl);
	}

	catch (const HPExceptionType& e)
	{
		if (e != kFileDataNotFound)
		{
			throw runtime_error("Caught exception " + boost::lexical_cast<string>(e));
		}
		else
		{
			myThreadedLogger.Error("hitool was unable to find data");
			t.join();
			t2.join();
			return;
		}
	}

	vector<double> t_diff;
	try
	{
		t_diff = h->VerticalMaximum(TParam, 0, 200);
		for (size_t i = 0; i < gridSize; ++i)
		{
			t_diff[i] = t_diff[i] - T_LowestLevelInfo->Data()[i];
		}
	}

	catch (const HPExceptionType& e)
	{
		if (e != kFileDataNotFound)
		{
			throw runtime_error("Caught exception " + boost::lexical_cast<string>(e));
		}
		else
		{
			myThreadedLogger.Error("hitool was unable to find data");
			t.join();
			t2.join();
			return;
		}
	}

	vector<double> BLbottom_ws;
	try
	{
		BLbottom_ws = h->VerticalAverage(WSParam, 0, 200);
	}

	catch (const HPExceptionType& e)
	{
		if (e != kFileDataNotFound)
		{
			throw runtime_error("Caught exception " + boost::lexical_cast<string>(e));
		}
		else
		{
			myThreadedLogger.Error("hitool was unable to find data");
			t.join();
			t2.join();
			return;
		}
	}

	t.join();
	t2.join();

	deltaTot dTot;
	intT iT;
	intTot iTot;

	DeltaTot(dTot, T_LowestLevelInfo, gridSize);
	IntT(iT, dT, gridSize);
	IntTot(iTot, dTot, gridSize);

	string deviceType = "CPU";

	LOCKSTEP(myTargetInfo, GustInfo, T_LowestLevelInfo, TopoInfo, BLHInfo, LCloudInfo, MCloudInfo)
	{
		size_t i = myTargetInfo->LocationIndex();

		double topo = TopoInfo->Value();
		double esto = kFloatMissing;
		double esto_tot = kFloatMissing;
		double gust = kFloatMissing;
		double turb_lisa = 0;
		double turb_kerroin = 0;
		double pilvikerroin = 0;
		double cloudCover = (LCloudInfo->Value() + MCloudInfo->Value()) * 100;
		topo *= himan::constants::kIg;

		NFmiLocation theLocation(myTargetInfo->LatLon().X(), myTargetInfo->LatLon().Y());
		double elevationAngle = theLocation.ElevationAngle(theTime);

		/* Calculations go here */

		if (z_boundaryl[i] >= 200)
		{
			esto = iT.intT_100[i] + iT.intT_200[i];
			esto_tot = iTot.intTot_100[i] + iTot.intTot_200[i];
		}

		if (z_boundaryl[i] >= 300)
		{
			esto = iT.intT_100[i] + iT.intT_200[i] + iT.intT_300[i];
			esto_tot = iTot.intTot_100[i] + iTot.intTot_200[i] + iTot.intTot_300[i];
		}

		if (z_boundaryl[i] >= 400)
		{
			esto = iT.intT_100[i] + iT.intT_200[i] + iT.intT_300[i] + iT.intT_400[i];
			esto_tot = iTot.intTot_100[i] + iTot.intTot_200[i] + iTot.intTot_300[i] + iTot.intTot_400[i];
		}

		if (z_boundaryl[i] >= 500)
		{
			esto = iT.intT_100[i] + iT.intT_200[i] + iT.intT_300[i] + iT.intT_400[i] + iT.intT_500[i];
			esto_tot =
			    iTot.intTot_100[i] + iTot.intTot_200[i] + iTot.intTot_300[i] + iTot.intTot_400[i] + iTot.intTot_500[i];
		}

		if (z_boundaryl[i] >= 600)
		{
			esto = iT.intT_100[i] + iT.intT_200[i] + iT.intT_300[i] + iT.intT_400[i] + iT.intT_500[i] + iT.intT_600[i];
			esto_tot = iTot.intTot_100[i] + iTot.intTot_200[i] + iTot.intTot_300[i] + iTot.intTot_400[i] +
			           iTot.intTot_500[i] + iTot.intTot_600[i];
		}

		if (z_boundaryl[i] == 700)
		{
			esto = iT.intT_100[i] + iT.intT_200[i] + iT.intT_300[i] + iT.intT_400[i] + iT.intT_500[i] + iT.intT_600[i] +
			       iT.intT_700[i];
			esto_tot = iTot.intTot_100[i] + iTot.intTot_200[i] + iTot.intTot_300[i] + iTot.intTot_400[i] +
			           iTot.intTot_500[i] + iTot.intTot_600[i] + iTot.intTot_700[i];
		}

		if (z_boundaryl[i] >= 200 && esto >= 0 && esto <= esto_tot)
		{
			gust = ((BLbottom_ws[i] - BLtop_ws[i]) / esto_tot) * esto + BLtop_ws[i];
		}

		if (z_boundaryl[i] >= 200 && esto <= 0 && esto >= -200)
		{
			gust = ((BLtop_ws[i] - maxEstimate[i]) / 200) * esto + BLtop_ws[i];
		}

		if (z_boundaryl[i] >= 200 && esto < -200)
		{
			gust = maxEstimate[i];
		}

		if (z_boundaryl[i] >= 200 && esto > esto_tot)
		{
			gust = BLbottom_ws[i];
		}

		if (z_boundaryl[i] < 200)
		{
			gust = meanWind[i];
		}

		if (gust == kFloatMissing || gust < 1)
		{
			gust = 1;
		}

		if (topo > 15 && topo < 400 && t_diff[i] > 0 && t_diff[i] <= 4 && BLbottom_ws[i] < 7 && elevationAngle < 15)
		{
			gust = ((1 - gust) / 4) * t_diff[i] + gust;
		}

		if (topo > 15 && topo < 400 && t_diff[i] > 4 && BLbottom_ws[i] < 7)
		{
			gust = 1;
		}

		if (cloudCover >= 30 && cloudCover <= 70)
		{
			pilvikerroin = -0.025 * cloudCover + 1.75;
		}

		if (cloudCover < 30)
		{
			pilvikerroin = 1;
		}

		if (elevationAngle >= 20 && topo > 10 && z_boundaryl[i] >= 200 && esto < 0)
		{
			turb_lisa = 0.133333333 * elevationAngle - 2.666666667;

			if (elevationAngle > 50)
			{
				turb_lisa = 4;
			}

			if (gust > 4 && gust <= 14)
			{
				turb_kerroin = -0.1 * gust + 1.4;
			}

			if (gust <= 4)
			{
				turb_kerroin = 1;
			}
		}

		gust = gust + turb_lisa * turb_kerroin * pilvikerroin;

		if (topo > 400 || T_LowestLevelInfo->Value() == kFloatMissing || BLHInfo->Value() == kFloatMissing)
		{
			gust = GustInfo->Value() * 0.95;
		}

		myTargetInfo->Value(gust);
	}

	himan::matrix<double> filter_kernel(3, 3, 1, kFloatMissing);
	filter_kernel.Fill(1.0 / 9.0);
	himan::matrix<double> gust_filtered = numerical_functions::Filter2D(myTargetInfo->Data(), filter_kernel);

	myTargetInfo->Grid()->Data(gust_filtered);

	myThreadedLogger.Info("[" + deviceType + "] Missing values: " +
						  boost::lexical_cast<string>(myTargetInfo->Data().MissingCount()) + "/" +
						  boost::lexical_cast<string>(myTargetInfo->Data().Size()));
}

void DeltaT(shared_ptr<const plugin_configuration> conf, info_t T_lowestLevel, const forecast_time& ftime,
            const forecast_type& ftype, size_t gridSize, deltaT& dT)
{
	auto h = GET_PLUGIN(hitool);

	h->Configuration(conf);
	h->Time(ftime);
	h->ForecastType(ftype);

	const param TParam("T-K");

	try
	{
		// Potential temperature differences
		dT.deltaT_100 = h->VerticalValue(TParam, 100);
		dT.deltaT_200 = h->VerticalValue(TParam, 200);
		dT.deltaT_300 = h->VerticalValue(TParam, 300);
		dT.deltaT_400 = h->VerticalValue(TParam, 400);
		dT.deltaT_500 = h->VerticalValue(TParam, 500);
		dT.deltaT_600 = h->VerticalValue(TParam, 600);
		dT.deltaT_700 = h->VerticalValue(TParam, 700);

		for (size_t i = 0; i < gridSize; ++i)
		{
			if (T_lowestLevel->Data()[i] == himan::kFloatMissing)
			{
				dT.deltaT_100[i] = himan::kFloatMissing;
				dT.deltaT_200[i] = himan::kFloatMissing;
				dT.deltaT_300[i] = himan::kFloatMissing;
				dT.deltaT_400[i] = himan::kFloatMissing;
				dT.deltaT_500[i] = himan::kFloatMissing;
				dT.deltaT_600[i] = himan::kFloatMissing;
				dT.deltaT_700[i] = himan::kFloatMissing;
				continue;
			}
			if (dT.deltaT_100[i] != himan::kFloatMissing)
				dT.deltaT_100[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (100.0 / 1000.0));
			if (dT.deltaT_200[i] != himan::kFloatMissing)
				dT.deltaT_200[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (200.0 / 1000.0));
			if (dT.deltaT_300[i] != himan::kFloatMissing)
				dT.deltaT_300[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (300.0 / 1000.0));
			if (dT.deltaT_400[i] != himan::kFloatMissing)
				dT.deltaT_400[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (400.0 / 1000.0));
			if (dT.deltaT_500[i] != himan::kFloatMissing)
				dT.deltaT_500[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (500.0 / 1000.0));
			if (dT.deltaT_600[i] != himan::kFloatMissing)
				dT.deltaT_600[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (600.0 / 1000.0));
			if (dT.deltaT_700[i] != himan::kFloatMissing)
				dT.deltaT_700[i] -= T_lowestLevel->Data()[i] + 9.8 * (0.010 - (700.0 / 1000.0));
		}
	}
	catch (const HPExceptionType& e)
	{
		if (e != kFileDataNotFound)
		{
			throw runtime_error("DeltaT() caught exception " + boost::lexical_cast<string>(e));
		}
	}
}

void DeltaTot(deltaTot& dTot, info_t T_lowestLevel, size_t gridSize)
{
	dTot.deltaTot_100 = vector<double>(gridSize, 0);
	dTot.deltaTot_200 = vector<double>(gridSize, 0);
	dTot.deltaTot_300 = vector<double>(gridSize, 0);
	dTot.deltaTot_400 = vector<double>(gridSize, 0);
	dTot.deltaTot_500 = vector<double>(gridSize, 0);
	dTot.deltaTot_600 = vector<double>(gridSize, 0);
	dTot.deltaTot_700 = vector<double>(gridSize, 0);

	for (size_t i = 0; i < gridSize; ++i)
	{
		dTot.deltaTot_100[i] =
		    (T_lowestLevel->Data()[i] + 6 * (0.010 - (100.0 / 1000.0))) -
		    (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (100.0 / 1000.0)));  // => -3.8*(0.010-(100/1000)) ???
		dTot.deltaTot_200[i] =
		    (T_lowestLevel->Data()[i] + 6 * (0.010 - (200.0 / 1000.0))) -
		    (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (200.0 / 1000.0)));  // => -3.8*(0.010-(200/1000))
		dTot.deltaTot_300[i] = (T_lowestLevel->Data()[i] + 6 * (0.010 - (300.0 / 1000.0))) -
		                       (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (300.0 / 1000.0)));  // ... etc.
		dTot.deltaTot_400[i] = (T_lowestLevel->Data()[i] + 6 * (0.010 - (400.0 / 1000.0))) -
		                       (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (400.0 / 1000.0)));
		dTot.deltaTot_500[i] = (T_lowestLevel->Data()[i] + 6 * (0.010 - (500.0 / 1000.0))) -
		                       (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (500.0 / 1000.0)));
		dTot.deltaTot_600[i] = (T_lowestLevel->Data()[i] + 6 * (0.010 - (600.0 / 1000.0))) -
		                       (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (600.0 / 1000.0)));
		dTot.deltaTot_700[i] = (T_lowestLevel->Data()[i] + 6 * (0.010 - (700.0 / 1000.0))) -
		                       (T_lowestLevel->Data()[i] + 9.8 * (0.010 - (700.0 / 1000.0)));
	}
}

void IntT(intT& iT, const deltaT& dT, size_t gridSize)
{
	iT.intT_100 = vector<double>(gridSize, himan::kFloatMissing);
	iT.intT_200 = vector<double>(gridSize, himan::kFloatMissing);
	iT.intT_300 = vector<double>(gridSize, himan::kFloatMissing);
	iT.intT_400 = vector<double>(gridSize, himan::kFloatMissing);
	iT.intT_500 = vector<double>(gridSize, himan::kFloatMissing);
	iT.intT_600 = vector<double>(gridSize, himan::kFloatMissing);
	iT.intT_700 = vector<double>(gridSize, himan::kFloatMissing);

	for (size_t i = 0; i < gridSize; ++i)
	{
		if (dT.deltaT_100[i] != himan::kFloatMissing)
			iT.intT_100[i] = 0.5 * dT.deltaT_100[i] * 100;  // why 0.5 * here? Would make sense in case of
		                                                    // (dT.deltaT_100[i] + dT.deltaT_0[i]) if this is
		                                                    // integration using trapezoidal rule
		if (dT.deltaT_100[i] != himan::kFloatMissing && dT.deltaT_200[i] != himan::kFloatMissing)
			iT.intT_200[i] = 0.5 * (dT.deltaT_200[i] + dT.deltaT_100[i]) * 100;
		if (dT.deltaT_200[i] != himan::kFloatMissing && dT.deltaT_300[i] != himan::kFloatMissing)
			iT.intT_300[i] = 0.5 * (dT.deltaT_300[i] + dT.deltaT_200[i]) * 100;
		if (dT.deltaT_300[i] != himan::kFloatMissing && dT.deltaT_400[i] != himan::kFloatMissing)
			iT.intT_400[i] = 0.5 * (dT.deltaT_400[i] + dT.deltaT_300[i]) * 100;
		if (dT.deltaT_400[i] != himan::kFloatMissing && dT.deltaT_500[i] != himan::kFloatMissing)
			iT.intT_500[i] = 0.5 * (dT.deltaT_500[i] + dT.deltaT_400[i]) * 100;
		if (dT.deltaT_500[i] != himan::kFloatMissing && dT.deltaT_600[i] != himan::kFloatMissing)
			iT.intT_600[i] = 0.5 * (dT.deltaT_600[i] + dT.deltaT_500[i]) * 100;
		if (dT.deltaT_600[i] != himan::kFloatMissing && dT.deltaT_700[i] != himan::kFloatMissing)
			iT.intT_700[i] = 0.5 * (dT.deltaT_700[i] + dT.deltaT_600[i]) * 100;
	}
}

void IntTot(intTot& iTot, const deltaTot& dTot, size_t gridSize)
{
	iTot.intTot_100 = vector<double>(gridSize, 0);
	iTot.intTot_200 = vector<double>(gridSize, 0);
	iTot.intTot_300 = vector<double>(gridSize, 0);
	iTot.intTot_400 = vector<double>(gridSize, 0);
	iTot.intTot_500 = vector<double>(gridSize, 0);
	iTot.intTot_600 = vector<double>(gridSize, 0);
	iTot.intTot_700 = vector<double>(gridSize, 0);

	for (size_t i = 0; i < gridSize; ++i)
	{
		iTot.intTot_100[i] = 0.5 * dTot.deltaTot_100[i] * 100;  // 0.5 * ???
		iTot.intTot_200[i] = 0.5 * (dTot.deltaTot_200[i] + dTot.deltaTot_100[i]) * 100;
		iTot.intTot_300[i] = 0.5 * (dTot.deltaTot_300[i] + dTot.deltaTot_200[i]) * 100;
		iTot.intTot_400[i] = 0.5 * (dTot.deltaTot_400[i] + dTot.deltaTot_300[i]) * 100;
		iTot.intTot_500[i] = 0.5 * (dTot.deltaTot_500[i] + dTot.deltaTot_400[i]) * 100;
		iTot.intTot_600[i] = 0.5 * (dTot.deltaTot_600[i] + dTot.deltaTot_500[i]) * 100;
		iTot.intTot_700[i] = 0.5 * (dTot.deltaTot_700[i] + dTot.deltaTot_600[i]) * 100;
	}
}

void LowAndMiddleClouds(vector<double>& lowAndMiddleClouds, info_t lowClouds, info_t middleClouds, info_t highClouds,
                        info_t totalClouds)
{
	for (size_t i = 0; i < lowAndMiddleClouds.size(); ++i)
	{
		if (highClouds->Data()[i] == himan::kFloatMissing || highClouds->Data()[i] == 0.0)
		{
			if (totalClouds->Data()[i] == himan::kFloatMissing)
				lowAndMiddleClouds[i] = himan::kFloatMissing;
			else
				lowAndMiddleClouds[i] = totalClouds->Data()[i] * 100;
		}
		else
		{
			if (lowClouds->Data()[i] != himan::kFloatMissing)
			{
				if (middleClouds->Data()[i] != himan::kFloatMissing)
					lowAndMiddleClouds[i] = max(lowClouds->Data()[i], middleClouds->Data()[i]) * 100;
				else
					lowAndMiddleClouds[i] = lowClouds->Data()[i] * 100;
			}
			else
			{
				if (middleClouds->Data()[i] == himan::kFloatMissing)
					lowAndMiddleClouds[i] = himan::kFloatMissing;
				else
					lowAndMiddleClouds[i] = middleClouds->Data()[i] * 100;
			}
		}
	}
}


/**
 *    Copyright (C) 2015 MongoDB, Inc.
 *
 *    This program 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 program 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/>.
 *
 *    As a special exception, the copyright holders give permission to link the
 *    code of portions of this program with the OpenSSL library under certain
 *    conditions as described in each individual source file and distribute
 *    linked combinations including the program with the OpenSSL library. You
 *    must comply with the GNU Affero General Public License in all respects for
 *    all of the code used other than as permitted herein. If you modify file(s)
 *    with this exception, you may extend this exception to your version of the
 *    file(s), but you are not obligated to do so. If you do not wish to do so,
 *    delete this exception statement from your version. If you delete this
 *    exception statement from all source files in the program, then also delete
 *    it in the license file.
 */

#define MONGO_LOG_DEFAULT_COMPONENT ::mongo::logger::LogComponent::kQuery

#include "mongo/platform/basic.h"

#include "mongo/db/pipeline/document_source_sample_from_random_cursor.h"

#include <boost/math/distributions/beta.hpp>

#include "mongo/db/client.h"
#include "mongo/db/pipeline/document.h"
#include "mongo/db/pipeline/expression.h"
#include "mongo/db/pipeline/expression_context.h"
#include "mongo/db/pipeline/value.h"
#include "mongo/util/log.h"

namespace mongo {
using boost::intrusive_ptr;

DocumentSourceSampleFromRandomCursor::DocumentSourceSampleFromRandomCursor(
    const intrusive_ptr<ExpressionContext>& pExpCtx,
    long long size,
    std::string idField,
    long long nDocsInCollection)
    : DocumentSource(pExpCtx),
      _size(size),
      _idField(std::move(idField)),
      _seenDocs(pExpCtx->getValueComparator().makeUnorderedValueSet()),
      _nDocsInColl(nDocsInCollection) {}

const char* DocumentSourceSampleFromRandomCursor::getSourceName() const {
    return "$sampleFromRandomCursor";
}

namespace {
/**
 * Select a random value drawn according to the distribution Beta(alpha=1, beta=N). The kth smallest
 * value of a sample of size N from a Uniform(0, 1) distribution has a Beta(k, N + 1 - k)
 * distribution, so the return value represents the smallest value from such a sample. This is also
 * the expected distance between the values drawn from a uniform distribution, which is how it is
 * being used here.
 */
double smallestFromSampleOfUniform(PseudoRandom* prng, size_t N) {
    boost::math::beta_distribution<double> betaDist(1.0, static_cast<double>(N));
    double p = prng->nextCanonicalDouble();
    return boost::math::quantile(betaDist, p);
}
}  // namespace

DocumentSource::GetNextResult DocumentSourceSampleFromRandomCursor::getNext() {
    pExpCtx->checkForInterrupt();

    if (_seenDocs.size() >= static_cast<size_t>(_size))
        return GetNextResult::makeEOF();

    auto nextResult = getNextNonDuplicateDocument();
    if (!nextResult.isAdvanced()) {
        return nextResult;
    }

    // Assign it a random value to enable merging by random value, attempting to avoid bias in that
    // process.
    auto& prng = pExpCtx->opCtx->getClient()->getPrng();
    _randMetaFieldVal -= smallestFromSampleOfUniform(&prng, _nDocsInColl);

    MutableDocument md(nextResult.releaseDocument());
    md.setRandMetaField(_randMetaFieldVal);
    if (pExpCtx->needsMerge) {
        // This stage will be merged by sorting results according to this random metadata field, but
        // the merging logic expects to sort by the sort key metadata.
        md.setSortKeyMetaField(BSON("" << _randMetaFieldVal));
    }
    return md.freeze();
}

DocumentSource::GetNextResult DocumentSourceSampleFromRandomCursor::getNextNonDuplicateDocument() {
    // We may get duplicate documents back from the random cursor, and should not return duplicate
    // documents, so keep trying until we get a new one.
    const int kMaxAttempts = 100;
    for (int i = 0; i < kMaxAttempts; ++i) {
        auto nextInput = pSource->getNext();
        switch (nextInput.getStatus()) {
            case GetNextResult::ReturnStatus::kAdvanced: {
                auto idField = nextInput.getDocument()[_idField];
                uassert(28793,
                        str::stream()
                            << "The optimized $sample stage requires all documents have a "
                            << _idField
                            << " field in order to de-duplicate results, but encountered a "
                               "document without a "
                            << _idField
                            << " field: "
                            << nextInput.getDocument().toString(),
                        !idField.missing());

                if (_seenDocs.insert(std::move(idField)).second) {
                    return nextInput;
                }
                LOG(1) << "$sample encountered duplicate document: "
                       << nextInput.getDocument().toString();
                break;  // Try again with the next document.
            }
            case GetNextResult::ReturnStatus::kPauseExecution: {
                MONGO_UNREACHABLE;  // Our input should be a random cursor, which should never
                                    // result in kPauseExecution.
            }
            case GetNextResult::ReturnStatus::kEOF: {
                return nextInput;
            }
        }
    }
    uasserted(28799,
              str::stream() << "$sample stage could not find a non-duplicate document after "
                            << kMaxAttempts
                            << " while using a random cursor. This is likely a "
                               "sporadic failure, please try again.");
}

Value DocumentSourceSampleFromRandomCursor::serialize(
    boost::optional<ExplainOptions::Verbosity> explain) const {
    return Value(DOC(getSourceName() << DOC("size" << _size)));
}

DocumentSource::GetDepsReturn DocumentSourceSampleFromRandomCursor::getDependencies(
    DepsTracker* deps) const {
    deps->fields.insert(_idField);
    return SEE_NEXT;
}

intrusive_ptr<DocumentSourceSampleFromRandomCursor> DocumentSourceSampleFromRandomCursor::create(
    const intrusive_ptr<ExpressionContext>& expCtx,
    long long size,
    std::string idField,
    long long nDocsInCollection) {
    intrusive_ptr<DocumentSourceSampleFromRandomCursor> source(
        new DocumentSourceSampleFromRandomCursor(expCtx, size, idField, nDocsInCollection));
    return source;
}
}  // mongo


//
// ********************************************************************
// * License and Disclaimer                                           *
// *                                                                  *
// * The  Geant4 software  is  copyright of the Copyright Holders  of *
// * the Geant4 Collaboration.  It is provided  under  the terms  and *
// * conditions of the Geant4 Software License,  included in the file *
// * LICENSE and available at  http://cern.ch/geant4/license .  These *
// * include a list of copyright holders.                             *
// *                                                                  *
// * Neither the authors of this software system, nor their employing *
// * institutes,nor the agencies providing financial support for this *
// * work  make  any representation or  warranty, express or implied, *
// * regarding  this  software system or assume any liability for its *
// * use.  Please see the license in the file  LICENSE  and URL above *
// * for the full disclaimer and the limitation of liability.         *
// *                                                                  *
// * This  code  implementation is the result of  the  scientific and *
// * technical work of the GEANT4 collaboration.                      *
// * By using,  copying,  modifying or  distributing the software (or *
// * any work based  on the software)  you  agree  to acknowledge its *
// * use  in  resulting  scientific  publications,  and indicate your *
// * acceptance of all terms of the Geant4 Software license.          *
// ********************************************************************
//
// $Id: PreampMedipix.cc $
//
/// \file src/PreampMedipix.cc
/// \brief Implementation of the PreampMedipix class

#include "PreampMedipix.hh"
#include "G4RunManager.hh"
#include "G4Run.hh"

#include <stdio.h>
#include <iostream>
#include <fstream>

#include <boost/concept_check.hpp>

#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/ini_parser.hpp>

using namespace std;

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
PreampMedipix::PreampMedipix(boost::property_tree::ptree pt)
{

    G4RunManager *fRM = G4RunManager::GetRunManager();
    myDet = (DetectorConstructionBase *)(fRM->GetUserDetectorConstruction());
    useCSM = myDet->GetCsmMode();

    G4int detectorType = myDet->GetDetectorType();
    // set gain modes for Medipix3RX and preamp
    if(detectorType == 0){
        if(useCSM == true){
            elecSigma = pt.get<G4double>("chip_medipix3rx.csm_enc");
        }else {
            elecSigma = pt.get<G4double>("chip_medipix3rx.spm_enc");
        }
        Cf = myDet->GetDetectorGain();
    } else if(detectorType == 1) { //Timepix
        elecSigma = pt.get<G4double>("chip_timepix.enc");
        Cf = 8e-15; // 8fF same as Medipix2MXR
    } else if(detectorType == 2) { // Dosepix
        elecSigma = pt.get<G4double>("chip_dosepix.enc");
        Cf = 8e-15;
    } else if (detectorType == 3) { // Timepix3
        elecSigma = pt.get<G4double>("chip_timepix3.enc");
        Cf = 3e-15; // 3fF
    }

    // get material and load material properties from ini file
    G4String material =  myDet->GetSensorMaterial()->GetName();
    G4String sensor;
    if (material == "G4_Si")
        sensor = "sensor_silicon";
    else if (material == "G4_CADMIUM_TELLURIDE")
        sensor = "sensor_cdte";

    pulsePrecision =    pt.get<G4double>("computation.pulsePrecision") * ns;  //must be ns!
    maxPulseTime =      pt.get<G4double>("computation.maxPulseTime") * ns;

    nPulseArrayElements = (G4int) maxPulseTime / pulsePrecision;
    ampResponseTime =   pt.get<G4double>("computation.ampResponseTime") * ns;
    nAmpResponseElements = (G4int)(ampResponseTime / pulsePrecision);

    nElectronHolePairs = pt.get<G4double>(sensor + ".nElectronHolePairs");
    minEnergy = 0.*keV;

    DEBUG = pt.get<G4bool>("testing.debug");
    writePeakToFile = pt.get<G4bool>("testing.writePeakToFile");

    preampType = pt.get<G4int>("computation.preampType");

    Ikrum = pt.get<G4double>("chip.Ikrum") * 1e-9;

    detector = MpxDetector::GetInstance();
    thresholdkeV = detector->GetTpxThreshold();
    thresholdCharge = thresholdkeV * 1000 / nElectronHolePairs ;

    //set the transfer functions
    SetTransferFunctions();

    //THL dispersion
    nPixel = myDet->GetNbPixels();
    TpxThlDisp = new G4double[nPixel * nPixel]();

    SetThresholdDispersion();
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::SetThresholdDispersion()
{
    //load .ini file with configuration data
    boost::property_tree::ptree pt;
    boost::property_tree::ini_parser::read_ini("DetectorConfig.ini", pt);

    //sensor properties
    G4String material =  myDet->GetSensorMaterial()->GetName();
    G4String sensor;
    if (material == "G4_Si")
        sensor = "sensor_silicon";
    else if (material == "G4_CADMIUM_TELLURIDE")
    {
        sensor = "sensor_cdte";
    }

    G4int detectorType = myDet->GetDetectorType();

    G4double disp = 0;

    if (detectorType == 0){ //Medipix3RX
        if(useCSM == true){
            disp = pt.get<G4double>("chip_medipix3rx.csm_threshold_dispersion");
        } else {
            disp = pt.get<G4double>("chip_medipix3rx.spm_threshold_dispersion");
        }

    }else if (detectorType == 1){ //Timepix1
        disp = pt.get<G4double>("chip_timepix.threshold_dispersion");

    }else if(detectorType == 2){ //Dosepix
        disp = pt.get<G4double>("chip_dosepix.threshold_dispersion");

    }else if (detectorType == 3) { // Timepix3
        disp = pt.get<G4double>("chip_timepix3.threshold_dispersion");
    }

    for (G4int i = 0; i < nPixel * nPixel; i++) {
        //G4cout << "DEBUG: Setting Threshold Dispersion " << disp << " i: " << i << " nPixel: " << nPixel*nPixel << G4endl;
        TpxThlDisp[i] = CLHEP::RandGauss::shoot(0, disp); //sigma of thl dispersion
    }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::SetTransferFunctions(){

    //Preamp transfer function
    G4double Cd = 25e-15;  //Detector capacitance

    G4double Gmfb   = Ikrum/2.*20;  //Transconductance in the feedback loop (approximation from Weak Inversion)
    G4double Rf     = 2./Gmfb;       //feedback resistor

    G4double gm0=30e-6;  //transconductance parameter
    //G4double R0=20e6;    //Output resistance amplifier
    G4double C0=25e-15;  //Output capacitance amplifier
    G4double Ct=Cd*Cf+Cd*C0+Cf*C0;

    G4double rcpr=Ct/(gm0*Cf);
    G4double rcpf=Cf*Rf;

    G4double w1 = 1./rcpr;
    G4double w2 = 1./rcpf;

    //create transfer function:
    nImpResEl = nAmpResponseElements;
    impRes = new G4double[nImpResEl];
    for (G4int i = 0; i < nImpResEl; i++) {
        impRes[i] = exp(-i* pulsePrecision * 1e-9 *w2 - exp(-i* pulsePrecision * 1e-9* w1));
        //G4cout << "Debug1: " << i << " value: " << impRes[i] << G4endl;
    }
    //Shaper transfer function
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::GetPixelResponse(map<pair<G4int, G4int>, G4double *> *inducedPixelContent, MpxDigitCollection* digitCollection, G4int event)
{
    thresholdkeV = detector->GetTpxThreshold();
    thresholdCharge = thresholdkeV * 1000 / nElectronHolePairs ;

    if(preampType == 0){
        chargeIntegrationPreamp(inducedPixelContent, event, digitCollection);
    } else if(preampType == 1){
        convolutionPreamp(inducedPixelContent, event, digitCollection);
    }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::chargeIntegrationPreamp(map<pair<G4int, G4int>, G4double *> *inducedPixelContent, G4int event, MpxDigitCollection *digitCollection)
{
    map<pair<G4int, G4int>, G4double> pixelsContent;

    map<pair<G4int, G4int>, G4double * >::iterator iPCItr;
    iPCItr = inducedPixelContent->begin();
    for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {

        G4double maxCharge = 0;
        G4double *tempArray = (*iPCItr).second;
        G4double *ampResponse = new G4double[nAmpResponseElements]();

        //preamp response
        if(preampType == 0){
            for (G4int i = 0; i < nPulseArrayElements; i++) {
                maxCharge += tempArray[i];
            }
        }

        //convert maximum in preamp response back to energy //FIXME here TOT
        G4double energy = ElectronicsNoise(maxCharge) * nElectronHolePairs * eV;

        //safe new pixel energy, not yet Digit(SPM and CSM)
        if (energy > minEnergy) {
            std::pair<G4int, G4int> newPixel;
            newPixel.first  = (*iPCItr).first.first;
            newPixel.second = (*iPCItr).first.second;
            pixelsContent[newPixel] = energy;
        }
        //write preamp response to file
        if (writePeakToFile == true) {

            G4int k = (*iPCItr).first.first * 1000 + (*iPCItr).first.second;
            char filename[64];
            std::sprintf(filename, "preampSingle%d", k);

            writeToFile(filename, ampResponse, nAmpResponseElements);
        }

    }

    //print the energy in every pixel to the console
    if (DEBUG) {
        std::map<std::pair<G4int, G4int>, G4double>::iterator testItr =  pixelsContent.begin();
        G4cout << pixelsContent.size() << " number of entries in pixel Content" << G4endl;

        for (; testItr != pixelsContent.end() ; testItr++) {
            G4cout << (*testItr).first.first << " " << (*testItr).first.second << " " << (*testItr).second / keV << G4endl;
        }
        G4cout << G4endl;
    }

    //create digits. one per pixel
    std::map<std::pair<G4int, G4int>, G4double>::iterator pCItr =  pixelsContent.begin();
    std::map<std::pair<G4int, G4int>, G4double>::iterator tempPCItr =  pixelsContent.begin();

    //charge summing mode
    if (useCSM == true) {

        map<pair<G4int, G4int>, G4double> csmSumContent;
        std::pair<G4int, G4int> tempPixelCSM;
        std::pair<G4int, G4int> iteratePixel;

        //calculate the 2x2 sums and store in csmSumContent, equivalent of the sum/pixel in the chip
        for (; pCItr != pixelsContent.end() ; pCItr++) {

            G4double sumEnergy = 0;

            //get tempPixel
            tempPixelCSM = pCItr->first;
            iteratePixel = tempPixelCSM;

            //sumEnergy of pixel
            if (pixelsContent[tempPixelCSM] > 0) sumEnergy = pixelsContent[tempPixelCSM];

            iteratePixel.first += 1;
            tempPCItr = pixelsContent.find(iteratePixel);
            if (tempPCItr != pixelsContent.end()) {
                if (tempPCItr->second > 0) sumEnergy += tempPCItr->second;
            }

            iteratePixel.second -= 1;
            tempPCItr = pixelsContent.find(iteratePixel);
            if (tempPCItr != pixelsContent.end()) {
                if (tempPCItr->second > 0) sumEnergy += tempPCItr->second;
            }

            iteratePixel.first -= 1;
            tempPCItr = pixelsContent.find(iteratePixel);
            if (tempPCItr != pixelsContent.end()) {
                if (tempPCItr->second > 0) sumEnergy += tempPCItr->second;
            }

            //store the sum with the pixel ID
            csmSumContent[tempPixelCSM] = sumEnergy;
        }

        //compare the sums
        std::map<std::pair<G4int, G4int>, G4double>::iterator csmItr =  csmSumContent.begin();
        for (; csmItr != csmSumContent.end() ; csmItr++) {

            //get tempPixel
            tempPixelCSM = csmItr->first;
            iteratePixel = tempPixelCSM;

            //neighbors
            G4double cmpEnergy1 = 0;
            G4double cmpEnergy2 = 0;
            G4double cmpEnergy3 = 0;
            G4double cmpEnergy4 = 0;

            //diagonal
            G4double cmpEnergy5 = 0;
            G4double cmpEnergy6 = 0;
            G4double cmpEnergy7 = 0;
            G4double cmpEnergy8 = 0;

            //get sumEnergy of pixel
            G4double pixelEnergy = csmSumContent[tempPixelCSM];

            iteratePixel = tempPixelCSM;
            iteratePixel.first += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy1 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy2 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.second += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy3 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.second -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy4 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first -= 1;
            iteratePixel.second -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy5 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first += 1;
            iteratePixel.second -= 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy6 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first -= 1;
            iteratePixel.second += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy7 += tempPCItr->second;
            }

            iteratePixel = tempPixelCSM;
            iteratePixel.first += 1;
            iteratePixel.second += 1;
            tempPCItr = csmSumContent.find(iteratePixel);
            if (tempPCItr != csmSumContent.end()) {
                cmpEnergy8 += tempPCItr->second;
            }

            //create digit
            if (pixelEnergy > cmpEnergy1 && pixelEnergy > cmpEnergy2 && pixelEnergy > cmpEnergy3 && pixelEnergy > cmpEnergy4 && pixelEnergy > cmpEnergy5 && pixelEnergy > cmpEnergy6 && pixelEnergy > cmpEnergy7 && pixelEnergy > cmpEnergy8) {
                Digit *digit = new Digit;
                digit->SetColumn(tempPixelCSM.first);
                digit->SetLine(tempPixelCSM.second);
                digit->SetEnergy(pixelEnergy / keV);
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }

    } else {  //no CSM
        for (; pCItr != pixelsContent.end() ; pCItr++) {
            //convert all events into digits
            G4double energy = (*pCItr).second;
            if (energy > minEnergy) {
                Digit *digit = new Digit;
                digit->SetColumn((*pCItr).first.first);
                digit->SetLine((*pCItr).first.second);
                digit->SetEnergy((energy / keV));
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }
    }

    //free memory
    iPCItr = inducedPixelContent->begin();
    for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
        G4double *tempArray = (*iPCItr).second;
        delete[] tempArray;
    }
    inducedPixelContent->clear();
    delete inducedPixelContent;
    pixelsContent.clear();

}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::convolutionPreamp(map<pair<G4int, G4int>, G4double *> *inducedPixelContent, G4int event, MpxDigitCollection *digitCollection)
{
    //SPM mode
    if(useCSM == false){

        map<pair<G4int, G4int>, G4double * >::iterator iPCItr;
        iPCItr = inducedPixelContent->begin();

        //iterate over induced currents in pixels
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {

            G4double *ampResponse = new G4double[nAmpResponseElements]();
            G4double *inducedChargeArray = (*iPCItr).second;

            //Preamp
            convolve(inducedChargeArray, ampResponse, nPulseArrayElements, impRes, nImpResEl);

            //noise
            for(G4int i=0; i<nAmpResponseElements; i++){
                ampResponse[i] = ElectronicsNoise(ampResponse[i]);
            }

            //getEnergy
            boost::tuple<G4double,G4double,G4double> totPeakToa;
            G4int pixelIndex = ((*iPCItr).first.first) * nPixel + ((*iPCItr).first.second);
            G4double thlDisp = TpxThlDisp[pixelIndex];

            totPeakToa = getToTandPeak(ampResponse, thresholdCharge, thlDisp);
            delete[] ampResponse;

            //convert to digits from peak value //FIXME here ToT
            if (boost::get<1>(totPeakToa) > minEnergy) {
                Digit *digit = new Digit;
                digit->SetColumn((*iPCItr).first.first);
                digit->SetLine((*iPCItr).first.second);
                digit->SetEnergy((boost::get<1>(totPeakToa) / keV));
                digit->SetToT(boost::get<0>(totPeakToa));
                digit->SetToA(boost::get<2>(totPeakToa));
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }

        //free memory
        iPCItr = inducedPixelContent->begin();
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
            G4double *tempArray = (*iPCItr).second;
            delete[] tempArray;
        }
        inducedPixelContent->clear();
        delete inducedPixelContent;
    }

    //CSM mode
    if(useCSM == true){

        map<pair<G4int, G4int>, G4double * >::iterator iPCItr;
        iPCItr = inducedPixelContent->begin();

        map<pair<G4int, G4int>, G4double * >* preampResponse    = new map<pair<G4int, G4int>, G4double * >;
        map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >* sumPreampResponse = new map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >;
        pair<G4int,G4int> tempPixel;
        G4double* emptyDummy = new G4double[nAmpResponseElements]();

        //preampresponse of every pixel
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
            G4double* ampResponse = new G4double[nAmpResponseElements]();
            G4double* inducedChargeArray = (*iPCItr).second;

            //Preamp
            convolve(inducedChargeArray, ampResponse, nPulseArrayElements, impRes, nImpResEl);

            //write single preamp response to file
            if (writePeakToFile == true) {

                G4int k = (*iPCItr).first.first * 1000 + (*iPCItr).first.second;
                char filename[64];
                std::sprintf(filename, "preampSingle%d", k);

                writeToFile(filename, ampResponse, nAmpResponseElements);
            }

            //store single preampResponse
            tempPixel.first = (*iPCItr).first.first;
            tempPixel.second = (*iPCItr).first.second;
            preampResponse->operator [](tempPixel) = ampResponse;
        }

        //sum the responses and store
        map<pair<G4int, G4int>, G4double * >::iterator pRItr;
        map<pair<G4int, G4int>, G4double * >::iterator tempPRItr;
        pRItr = preampResponse->begin();

        for(; pRItr != preampResponse->end(); pRItr++){

            pair<G4int,G4int>mainPixel = pRItr->first;
            pair<G4int,G4int>iteratePixel = mainPixel;

            //get the preamp arrays
            G4double *tempPreampReponse;
            //sum the responses
            G4double* csmSumPreamp = new G4double[nAmpResponseElements]();

            G4int sumArrLow = 0;
            G4int sumArrHigh = 1;

            for(G4int i=sumArrLow; i<=sumArrHigh; i++){
                for(G4int j=sumArrLow; j<=sumArrHigh; j++){
                    //set pixel
                    iteratePixel.first = mainPixel.first + i;
                    iteratePixel.second = mainPixel.second + j;

                    //find response pulse
                    tempPRItr = preampResponse->find(iteratePixel);
                    if(tempPRItr != preampResponse->end()){
                        tempPreampReponse = tempPRItr->second;

                        //add to sum
                        for(G4int k=0; k < nAmpResponseElements; k++){
                            csmSumPreamp[k] += tempPreampReponse[k];
                        }
                    }
                }
            }

            //add noise
            for(G4int i=0; i < nAmpResponseElements; i++){
                csmSumPreamp[i] = ElectronicsNoise(csmSumPreamp[i]);
            }

            //write CSM preamp response to file
            if (writePeakToFile == true) {

                G4int k = (*pRItr).first.first * 1000 + (*pRItr).first.second;
                char filename[64];
                std::sprintf(filename, "preampCSM%d", k);

                writeToFile(filename, csmSumPreamp, nAmpResponseElements);
            }

            //get energy
            boost::tuple<G4double,G4double,G4double> totPeakToa;
            G4int pixelIndex = ((*pRItr).first.first) * nPixel + ((*pRItr).first.second);
            G4double thlDisp = TpxThlDisp[pixelIndex];

            totPeakToa = getToTandPeak(csmSumPreamp, thresholdCharge, thlDisp);

            sumPreampResponse->operator [](mainPixel) = totPeakToa;

            //delete the array
            delete[] csmSumPreamp;
        }

        //compare the sums
        map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >::iterator sPRItr = sumPreampResponse->begin();
        map<pair<G4int, G4int>, boost::tuple<G4double,G4double,G4double> >::iterator tempSPRItr;
        for (; sPRItr != sumPreampResponse->end() ; sPRItr++) {

            //get tempPixel
            pair<G4int,G4int> tempPixelCSM = sPRItr->first;
            pair<G4int,G4int> iteratePixel = tempPixelCSM;

            //first value: ToT, second: peak amplitude
            boost::tuple<G4double,G4double,G4double> cmpEnergy(0,0,0);

            //get sumEnergy of pixel
            boost::tuple<G4double,G4double,G4double> pixelEnergy = sumPreampResponse->operator [](tempPixelCSM);
            G4bool isMax = true;

            //sets the borders for the comparison
            G4int compArrLow    = -1;
            G4int compArrHigh   = 1;

            //compare neighbouring pixels with main pixel
            for(G4int i=compArrLow; i<=compArrHigh; i++){
                for(G4int j=compArrLow; j<=compArrHigh; j++){
                    if(!(i == 0 && j == 0)){
                        iteratePixel = tempPixelCSM;
                        iteratePixel.first += i;
                        iteratePixel.second += j;
                        tempSPRItr = sumPreampResponse->find(iteratePixel);
                        if (tempSPRItr != sumPreampResponse->end()) {
                            cmpEnergy = tempSPRItr->second;
                            if(boost::get<1>(cmpEnergy) > boost::get<1>(pixelEnergy)) isMax = false;
                        }
                    }
                }
            }

            //create digit
            if (isMax == true) {
                Digit *digit = new Digit;
                digit->SetColumn(tempPixelCSM.first);
                digit->SetLine(tempPixelCSM.second);
                digit->SetEnergy((boost::get<1>(pixelEnergy) / keV));
                digit->SetToT(boost::get<0>(pixelEnergy));
                digit->SetToA(boost::get<2>(pixelEnergy));
                digit->SetEvent(event);
                digitCollection->insert(digit);
            }
        }

        //free memory
        //induced pulses
        iPCItr = inducedPixelContent->begin();
        for (; iPCItr != inducedPixelContent->end() ; iPCItr++) {
            G4double *tempArray = (*iPCItr).second;
            delete[] tempArray;
        }
        inducedPixelContent->clear();
        delete inducedPixelContent;

        //preamp response
        pRItr = preampResponse->begin();

        for(; pRItr != preampResponse->end(); pRItr++){
            G4double* tempArray = (*pRItr).second;
            delete[] tempArray;
        }
        preampResponse->clear();
        delete preampResponse;

        //preamp sum
        sumPreampResponse->clear();
        delete sumPreampResponse;

        delete[] emptyDummy;
    }
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
G4double PreampMedipix::ElectronicsNoise(G4double charge)
{
    G4double sigma = elecSigma;
    if(useCSM == true && preampType != 2) sigma = sigma * 2;
    return CLHEP::RandGauss::shoot(charge, sigma);
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
G4double PreampMedipix::convolve(G4double* in, G4double* out, G4int inLength, G4double* kernel, G4int kernel_length)
{
    G4double maxCharge = 0;

    for(G4int i=0; i<(kernel_length); i++){

        out[i] = 0.0;
        for(G4int k=0; k<kernel_length; k++){
            if((i-k) >=0 && (i-k) < inLength){
                out[i] += in[i-k] * kernel[k];
            }
        }
        if(out[i] > maxCharge) maxCharge = out[i];
    }

    return maxCharge;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
boost::tuple<G4double,G4double,G4double> PreampMedipix::getToTandPeak(G4double* preamp, G4double threshold, G4double thlDisp){
    G4int risingEdge = 0;

    boost::tuple<G4double,G4double,G4double> returnValues(0,0,0);
    // 0 = ToT, 1 = peakCharge, 2 = ToA

    for(G4int i=0; i<nAmpResponseElements; i++){

        if(risingEdge == 0 && preamp[i] >= (threshold+thlDisp)){
            risingEdge = i;
            // Set timeOfArrival
            if(boost::get<2>(returnValues) == 0) {
                boost::get<2>(returnValues) = i * pulsePrecision;
            }
        }

        if(risingEdge != 0 && preamp[i] <= (threshold+thlDisp)){
            // Set timeOverThreshold
            if((i - risingEdge) * pulsePrecision > boost::get<0>(returnValues)) {
                boost::get<0>(returnValues) = (i - risingEdge) * pulsePrecision;
            }
            risingEdge = 0;
        }

        // set the peakCharge
        if(boost::get<1>(returnValues) < (preamp[i]+thlDisp)) boost::get<1>(returnValues) = (preamp[i]+thlDisp);
    }

    boost::get<1>(returnValues) = boost::get<1>(returnValues) * nElectronHolePairs * eV;

    return returnValues;
}

//....oooOO0OOooo........oooOO0OOooo........oooOO0OOooo........oooOO0OOooo......
void PreampMedipix::writeToFile(G4String filename, G4double* data, G4int nElements){
    ofstream myfile;

    myfile.open(filename);

    for (G4int it = 0; it < nElements; it++) {
        myfile << data[it] << " ";
    }
    myfile.close();
}


// export POPLAR_ENGINE_OPTIONS='{"target.workerStackSizeInBytes":1024}'

#include <iostream>
#include <vector>
#include <boost/timer/timer.hpp>

//~ poplar
#include <poplar/DeviceManager.hpp>
#include <poplar/IPUModel.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Program.hpp>

//~ poplin
#include <poplin/codelets.hpp>

//~ popops
#include <popops/ElementWise.hpp>
#include <popops/ExprOp.hpp>
#include <popops/Pad.hpp>
#include <popops/codelets.hpp>

//~ poputil
#include <poputil/TileMapping.hpp>

//~ KF header
#include "KalmanFilter.h"

/* contents
~01 KalmanFilter::connectToIPU
~02 KalmanFilter::skipSwitch
~03 KalmanFilter::smoothingState
~04 KalmanFilter::appendTo
~05 KalmanFilter::inverse
~06 KalmanFilter::packHits
~07 KalmanFilter::iterate
~08 KalmanFilter::product
~09 KalmanFilter::scaledAdd
~10 KalmanFilter::filter
~11 KalmanFilter::propagateState
~12 KalmanFilter::jacobian
~13 KalmanFilter::projectEKF
~14 KalmanFilter::project
~15 KalmanFilter::packIterationTensors
~16 KalmanFilter::calcResidual
~17 KalmanFilter::calcChiSq
~18 KalmanFilter::chiSqTest
~19 KalmanFilter::smooth
*/

using namespace poplar;
using namespace poplar::program;
using namespace popops;

Device KalmanFilter::connectToIPU()
{

    DeviceManager manager = DeviceManager::createDeviceManager();
    Device dev;

    bool success = false;

    //~ looping over each IPU
    for (auto &d : manager.getDevices(poplar::TargetType::IPU, 1))
    {
        dev = std::move(d);
        std::cerr << "Trying to attach to IPU " << dev.getId();
        if ((success = dev.attach()))
        {
            std::cerr << " - attached" << std::endl;
            break;
        }
        else
        {
            std::cerr << std::endl;
        }
    }

    if (!success)
    {
        std::cerr << "Error attaching to device" << std::endl;
        exit(1);
    }

    return dev;

} //~ end KalmanFilter::connectToIPU

std::tuple<ComputeSet, Tensor, Tensor, Tensor>
KalmanFilter::skipSwitch(Graph &graph,
                         const Tensor &inputPOld,
                         const Tensor &inputCOld,
                         const Tensor &inputPNew,
                         const Tensor &inputCNew,
                         const Tensor &chiSq,
                         uint tile)
{

    ComputeSet computeSet = graph.addComputeSet("SkipSwitch" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "SkipSwitch");

    graph.setCycleEstimate(vtx, 100);

    Tensor outP =
        graph.addVariable(FLOAT, {2, 1}, "switchOutP" + std::to_string(tile));
    Tensor outC =
        graph.addVariable(FLOAT, {2, 2}, "switchOutC" + std::to_string(tile));
    Tensor outD =
        graph.addVariable(FLOAT, {1, 1}, "switchOutD" + std::to_string(tile));

    graph.connect(vtx["inputPOld"], inputPOld);
    graph.connect(vtx["inputCOld"], inputCOld);
    graph.connect(vtx["inputPNew"], inputPNew);
    graph.connect(vtx["inputCNew"], inputCNew);
    graph.connect(vtx["chiSq"], chiSq);
    graph.connect(vtx["outP"], outP);
    graph.connect(vtx["outC"], outC);
    graph.connect(vtx["outD"], outD);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(outP, tile);
    graph.setTileMapping(outC, tile);
    graph.setTileMapping(outD, tile);

    return std::make_tuple(computeSet, outP, outC, outD);

} //~ end KalmanFilter::skipSwitch

std::tuple<ComputeSet, Tensor>
KalmanFilter::smoothingState(Graph &graph,
                             const Tensor &states,
                             const Tensor &itr,
                             uint tile)
{
    ComputeSet computeSet =
        graph.addComputeSet("smoothingState" + std::to_string(tile));

    VertexRef vtx = graph.addVertex(computeSet, "StateForSmoothing");

    graph.setCycleEstimate(vtx, 100);

    Tensor tmp =
        states.reshape({states.shape()[0], states.shape()[1] * states.shape()[2]});

    Tensor state =
        graph.addVariable(FLOAT,
                          {states.shape()[1], states.shape()[2]},
                          "state" + std::to_string(tile));

    graph.connect(vtx["state"], state.reshape({state.shape()[0] * state.shape()[1]}));
    graph.connect(vtx["states"], tmp);
    graph.connect(vtx["itr"], itr);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(tmp, tile);
    graph.setTileMapping(state, tile);

    return std::make_pair(computeSet,
                          state.reshape({states.shape()[1],
                          states.shape()[2]}));

} //~ end KalmanFilter::smoothingState

std::tuple<ComputeSet, Tensor>
KalmanFilter::appendTo(Graph &graph,
                       const Tensor &state,
                       const Tensor &itr,
                       Tensor &states,
                       uint tile)
{
    ComputeSet computeSet = graph.addComputeSet("append" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "AppendState");
    graph.setCycleEstimate(vtx, 100);

    Tensor tmp = states.reshape(
            {states.shape()[0], states.shape()[1] * states.shape()[2]});

    Tensor newStates =
            graph.addVariable(FLOAT, tmp.shape(), "newStates" + std::to_string(tile));

    graph.connect(vtx["state"], state.reshape({state.shape()[0] * state.shape()[1]}));
    graph.connect(vtx["states"], tmp);
    graph.connect(vtx["newStates"], newStates);
    graph.connect(vtx["itr"], itr);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(tmp, tile);
    graph.setTileMapping(newStates, tile);

    return std::make_pair(computeSet,
                          newStates.reshape({states.shape()[0],
                          states.shape()[1],
                          states.shape()[2]}));

} //~ end KalmanFilter::appendTo

std::tuple<ComputeSet, Tensor>
KalmanFilter::inverse(Graph &graph,
                      const Tensor &in,
                      uint tile,
                      uint dim)
{

    ComputeSet computeSet = graph.addComputeSet("inverse" + std::to_string(tile));

    VertexRef vtx =
        graph.addVertex(computeSet, "MatrixInverse" + std::to_string(dim));

    graph.setCycleEstimate(vtx, 100);

    Tensor out =
            graph.addVariable(FLOAT, {dim, dim}, "invOut" + std::to_string(tile));

    graph.connect(vtx["in"], in);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::inverse

std::tuple<ComputeSet, Tensor>
KalmanFilter::packHits(Graph &graph,
                       const Tensor &inputHits,
                       const Tensor &itr,
                       uint tile)
{

    ComputeSet computeSet =
            graph.addComputeSet("packHits4" + std::to_string(tile));

    VertexRef vtx = graph.addVertex(computeSet, "PackHits4");

    graph.setCycleEstimate(vtx, 10);

    // Same shape as ps
    Tensor out =
            graph.addVariable(FLOAT, {4, 1}, "hitsThisItr" + std::to_string(tile));

    graph.connect(vtx["inputHits"], inputHits);
    graph.connect(vtx["itr"], itr);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::packHits

std::tuple<ComputeSet, Tensor>
KalmanFilter::iterate(Graph &graph,
                      const Tensor &iterator,
                      uint tile)
{

    ComputeSet computeSet = graph.addComputeSet("iterate" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "Iterate");
    graph.setCycleEstimate(vtx, 10);

    Tensor out = graph.addVariable(INT, {1}, "itrOut" + std::to_string(tile));

    graph.connect(vtx["in"], iterator);
    graph.connect(vtx["out"], out);
    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::iterate

std::tuple<ComputeSet, Tensor>
KalmanFilter::product(Graph &graph,
                      const Tensor &first,
                      const Tensor &second,
                      uint tile)
{

    ComputeSet computeSet = graph.addComputeSet("product" + std::to_string(tile));
    VertexRef vtx = graph.addVertex(computeSet, "MatrixProduct");
    graph.setCycleEstimate(vtx, 100);

    uint outRows = first.shape()[0];
    uint outCols = second.shape()[1];

    Tensor out = graph.addVariable(FLOAT, {outRows, outCols},
                                   "prodOut" + std::to_string(tile));

    graph.connect(vtx["first"], first);
    graph.connect(vtx["second"], second);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::product

std::tuple<ComputeSet, Tensor>
KalmanFilter::scaledAdd(Graph &graph,
                        const Tensor &first,
                        const Tensor &second,
                        uint tile,
                        float s,
                        std::string name)
{

    ComputeSet computeSet =
        graph.addComputeSet("scaledAdd" + name + std::to_string(tile));

    VertexRef vtx = graph.addVertex(computeSet, "ScaledAdd");
    graph.setCycleEstimate(vtx, 100);

    uint outRows = first.shape()[0];
    uint outCols = first.shape()[1];

    Tensor out = graph.addVariable(FLOAT, {outRows, outCols},
                                   "addOut" + name + std::to_string(tile));

    graph.connect(vtx["first"], first);
    graph.connect(vtx["second"], second);
    graph.connect(vtx["s"], s);
    graph.connect(vtx["out"], out);

    graph.setTileMapping(vtx, tile);
    graph.setTileMapping(out, tile);

    return std::make_pair(computeSet, out);

} //~ end KalmanFilter::scaledAdd

std::tuple<Sequence, std::vector<Tensor>, std::vector<Tensor>>
KalmanFilter::filter(Graph &graph,
                     const std::vector<Tensor> &H,
                     const std::vector<Tensor> &G,
                     const std::vector<Tensor> &hits,
                     const std::vector<Tensor> &p,
                     const std::vector<Tensor> &C)
{

    Sequence op;

    std::vector<Tensor> outp(p.size());
    std::vector<Tensor> outc(p.size());

    for (uint i = 0; i < p.size(); i++)
    {

        outp[i] = graph.addVariable(FLOAT, {4, 1});
        graph.setTileMapping(outp[i], i);

        outc[i] = graph.addVariable(FLOAT, {4, 4});
        graph.setTileMapping(outc[i], i);

        auto [computeHG, HG] = product(graph, H[i], G[i], i);
        op.add(Execute(computeHG));

        // inv_C_proj + HG @ H
        auto [computeHGH, HGH] = product(graph, HG, H[i], i);
        op.add(Execute(computeHGH));

        auto [computeInverseC, invOutC] = inverse(graph, C[i], i, 4);
        op.add(Execute(computeInverseC));

        auto [computeAddCHGH, CHGH_HGH] = scaledAdd(graph, invOutC, HGH, i);
        op.add(Execute(computeAddCHGH));

        auto [computeInverseCHGH, invOutCHGH] = inverse(graph, CHGH_HGH, i, 4);
        op.add(Execute(computeInverseCHGH));

        auto [computeHGhits, HGhits] = product(graph, HG, hits[i], i);
        op.add(Execute(computeHGhits));

        auto [computePFilt, Cp_filt] = product(graph, invOutC, p[i], i);
        op.add(Execute(computePFilt));

        auto [computeAddCpHits, Cp_filt_hits] = scaledAdd(graph, Cp_filt, HGhits, i);
        op.add(Execute(computeAddCpHits));

        auto [computeCHGHPFilt, CHGHPFilt] = product(graph, invOutCHGH, Cp_filt_hits, i);
        op.add(Execute(computeCHGHPFilt));

        op.add(Copy(CHGHPFilt, outp[i]));
        op.add(Copy(invOutCHGH, outc[i]));

    }

    return std::make_tuple(op, outp, outc);

} //~ end KalmanFilter::filter

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::propagateState(Graph &graph,
                             std::vector<Tensor> &ps,
                             std::vector<Tensor> &d)
{

    Sequence op;

    std::vector<Tensor> outP(ps.size());

    for (uint i = 0; i < ps.size(); i++)
    {

        outP[i] = graph.addVariable(FLOAT, {2, 1});
        graph.setTileMapping(outP[i], i);

        // Need to allocate this earlier, and not during every function invocation!
        Tensor updMFlat = graph.addConstant<float>(FLOAT, {4, 1}, {0., 1., 0., 0.});
        Tensor updM = graph.addVariable(FLOAT, {2, 2});

        graph.setTileMapping(updMFlat, i);
        graph.setTileMapping(updM, i);

        updM = updMFlat.reshape({2, 2}); // Can't seem to set nD constant

        graph.setTileMapping(updM, i);

        // xTan = tan([x0, x1])
        Tensor pTan = popops::map(graph, popops::expr::UnaryOpType::TAN, ps[i], op);
        graph.setTileMapping(pTan, i);

        // x += d * [[0, 1], [0, 0]] * xTan -> x0 += d * tan(x1), x1 += 0

        auto [computeProdUpdTan, updTan] = product(graph, updM, pTan, i);
        op.add(Execute(computeProdUpdTan));

        auto [computeProdUpdD, updD] = product(graph, updTan, d[i], i);
        op.add(Execute(computeProdUpdD));

        auto [computeAddP, addP] = scaledAdd(graph, ps[i], updD, i);
        op.add(Execute(computeAddP));

        op.add(Copy(addP, outP[i]));

    }

    return std::make_tuple(op, outP);

} //~ end KalmanFilter::propagateState

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::jacobian(Graph &graph,
                       const std::vector<Tensor> &ps,
                       const std::vector<Tensor> &d)
{

    Sequence op;

    std::vector<Tensor> jacs(ps.size());

    for (uint i = 0; i < ps.size(); i++)
    {

        Tensor updMFlat = graph.addConstant<float>(FLOAT, {4, 1}, {0., 1., 0., 0.});
        Tensor updM = graph.addVariable(FLOAT, {2, 2});

        jacs[i] = graph.addVariable(FLOAT, {2, 2});

        graph.setTileMapping(updMFlat, i);
        graph.setTileMapping(updM, i);
        graph.setTileMapping(jacs[i], i);

        updM = updMFlat.reshape({2, 2});

        Tensor xCos = popops::map(graph, popops::expr::UnaryOpType::COS, ps[i], op);
        Tensor xCosSq =
            popops::map(graph, popops::expr::UnaryOpType::SQUARE, xCos, op);
        Tensor xCosSqI =
            popops::map(graph, popops::expr::UnaryOpType::INVERSE, xCosSq, op);

        Tensor xCosM = popops::pad(graph, xCosSqI,
                                   1,  // pad lower
                                   0,  // pad upper,
                                   1,  // pad dim
                                   0); // pad value

        graph.setTileMapping(xCos, i);
        graph.setTileMapping(xCosM, i);
        graph.setTileMapping(xCosSq, i);
        graph.setTileMapping(xCosSqI, i);

        auto [computeJac, jac] = product(graph, updM, xCosM, i);
        op.add(Execute(computeJac));

        Tensor updM2Flat =
            graph.addConstant<float>(FLOAT, {4, 1}, {1., 0., 0., 1.});

        Tensor updM2 = graph.addVariable(FLOAT, {2, 2});

        graph.setTileMapping(updM2Flat, i);
        graph.setTileMapping(updM2, i);

        updM2 = updM2Flat.reshape({2, 2});

        auto [computeAddJacUpd, jac_upd] = scaledAdd(graph, jac, updM2, i);
        op.add(Execute(computeAddJacUpd));

        op.add(Copy(jac_upd, jacs[i]));

    }

    return std::make_tuple(op, jacs);

} //~ end KalmanFilter::jacobian

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::projectEKF(Graph &graph,
                         const std::vector<Tensor> &jacs,
                         const std::vector<Tensor> &qs,
                         std::vector<Tensor> &covs)
{

    Sequence op;

    std::vector<Tensor> outC(covs.size());

    for (uint i = 0; i < jacs.size(); i++)
    {

        outC[i] = graph.addVariable(FLOAT, {2, 2});
        graph.setTileMapping(outC[i], i);

        auto [computeCovT1, covT1] = product(graph, jacs[i], covs[i], i);
        op.add(Execute(computeCovT1));

        Tensor jacsT = jacs[i].transpose();
        graph.setTileMapping(jacsT, i);

        auto [computeCovT2, covT2] = product(graph, covT1, jacsT, i);
        op.add(Execute(computeCovT2));

        auto [computeAddCovQ, covQ] = scaledAdd(graph, covT2, qs[i], i);
        op.add(Execute(computeAddCovQ));

        op.add(Copy(covQ, outC[i]));

    }

    return std::make_tuple(op, outC);

} //~ end KalmanFilter::projectEKF

std::tuple<Sequence, std::vector<Tensor>, std::vector<Tensor>>
KalmanFilter::project(Graph &graph,
                      const std::vector<Tensor> &ps,
                      const std::vector<Tensor> &C,
                      const std::vector<Tensor> &F,
                      const std::vector<Tensor> &Q)
{

    Sequence op;

    std::vector<Tensor> outP(ps.size());
    std::vector<Tensor> outC(ps.size());

    for (uint i = 0; i < outP.size(); i++)
    {

        auto [computePproj, p_proj] = product(graph, F[i], ps[i], i);
        op.add(Execute(computePproj));

        auto [computeFC, fc] = product(graph, F[i], C[i], i);
        op.add(Execute(computeFC));

        // Check this works wrt tile mapping
        auto [computeFCF, fcf] = product(graph, fc, F[i].transpose(), i);
        op.add(Execute(computeFCF));

        auto [computeCproj, c_proj] = scaledAdd(graph, fcf, Q[i], i);
        op.add(Execute(computeCproj));

        outC[i] = graph.addVariable(FLOAT, {4, 4});
        outP[i] = graph.addVariable(FLOAT, {4, 1});

        graph.setTileMapping(outP[i], i);
        graph.setTileMapping(outC[i], i);

        op.add(Copy(c_proj, outC[i]));
        op.add(Copy(p_proj, outP[i]));

    }

    return std::make_tuple(op, outP, outC);

} //~ end KalmanFilter::project

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::packIterationTensors(Graph &graph,
                                   const std::vector<Tensor> &loop,
                                   const std::vector<Tensor> &scatterInto,
                                   const std::vector<Tensor> &inputs)
{

    Sequence op;
    std::vector<Tensor> hits_packed(inputs.size());

    for (uint i = 0; i < inputs.size(); i++)
    {

        std::string iStr = std::to_string(i);

        hits_packed[i] = graph.addVariable(FLOAT, {4, 1}, "hits_packed" + iStr);
        graph.setTileMapping(hits_packed[i], i);

        auto [computeH, h] = packHits(graph, inputs[i], loop[i], i);
        op.add(Execute(computeH));
        op.add(Copy(h, hits_packed[i]));

    }

    return std::make_tuple(op, hits_packed);

} //~ end KalmanFilter::packIterationTensors

// Residual at step t
std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::calcResidual(Graph &graph,
                           const std::vector<Tensor> &hits,
                           const std::vector<Tensor> &p_filt,
                           const std::vector<Tensor> &H)
{

    Sequence op;
    std::vector<Tensor> res_out(p_filt.size());

    for (uint i = 0; i < res_out.size(); i++)
    {

        std::string iStr = std::to_string(i);

        res_out[i] = graph.addVariable(FLOAT, {4, 1}, "res" + iStr);
        graph.setTileMapping(res_out[i], i);

        auto [computeHp, Hp] = product(graph, H[i], p_filt[i], i);
        op.add(Execute(computeHp));

        // Scaled add with prefactor of -1
        auto [computeRes, res] = scaledAdd(graph, hits[i], Hp, i, -1.0, "res");
        op.add(Execute(computeRes));
        op.add(Copy(res, res_out[i]));

    }

    return std::make_tuple(op, res_out);

} //~ end KalmanFilter::calcResidual

// ChiSq for a plane of hits
std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::calcChiSq(Graph &graph,
                        const std::vector<Tensor> &res_filt,
                        const std::vector<Tensor> &G,
                        const std::vector<Tensor> &C_proj,
                        const std::vector<Tensor> &p_proj,
                        const std::vector<Tensor> &p_filt)
{

    Sequence op;
    std::vector<Tensor> chiSq(p_filt.size());

    for (uint i = 0; i < chiSq.size(); i++)
    {

        std::string iStr = std::to_string(i);

        Tensor rT = res_filt[i].transpose();
        graph.setTileMapping(rT, i);

        auto [computeRTG, rTG] = product(graph, rT, G[i], i);
        op.add(Execute(computeRTG));

        auto [computeResTerm, resTerm] = product(graph, rTG, res_filt[i], i);
        op.add(Execute(computeResTerm));

        auto [computeStateDiff, stateDiff] =
            scaledAdd(graph, p_filt[i], p_proj[i], i, -1.0);
        op.add(Execute(computeStateDiff));

        auto [computeInverseC, C_inv] = inverse(graph, C_proj[i], i, 4);
        op.add(Execute(computeInverseC));

        Tensor diffT = stateDiff.transpose();
        graph.setTileMapping(diffT, i);

        auto [computeStateTC, stateTC] = product(graph, diffT, C_inv, i);
        op.add(Execute(computeStateTC));

        auto [computeStateTerm, stateTerm] =
            product(graph, stateTC, stateDiff, i);
        op.add(Execute(computeStateTerm));

        auto [computeChiSqTot, chiSqTot] =
            scaledAdd(graph, resTerm, stateTerm, i, 1.0);
        op.add(Execute(computeChiSqTot));

        chiSq[i] = graph.addVariable(FLOAT, {1}, "chiSq" + iStr);
        graph.setTileMapping(chiSq[i], i);

        op.add(Copy(chiSqTot, chiSq[i]));

    }

    return std::make_tuple(op, chiSq);

} //~ end KalmanFilter::calcChiSq

std::tuple<Sequence, std::vector<Tensor>>
KalmanFilter::chiSqTest(Graph &graph,
                        const std::vector<Tensor> &chiSq,
                        const std::vector<Tensor> &threshold)
{

    Sequence op;
    std::vector<Tensor> chiSqPred(chiSq.size());

    for (uint i = 0; i < chiSqPred.size(); i++)
    {

        std::string iStr = std::to_string(i);

        chiSqPred[i] = popops::map(graph,
                                   popops::expr::BinaryOpType::GREATER_THAN,
                                   chiSq[i], threshold[i], op);

        graph.setTileMapping(chiSqPred[i], i);

    }

    return std::make_tuple(op, chiSqPred);

} //~ end KalmanFilter::chiSqTest

std::tuple<Sequence, std::vector<Tensor>, std::vector<Tensor>>
KalmanFilter::smooth(Graph &graph,
                     const std::vector<Tensor> &p_smooth_prev,
                     const std::vector<Tensor> &C_smooth_prev,
                     const std::vector<Tensor> &p_filt,
                     const std::vector<Tensor> &C_filt,
                     const std::vector<Tensor> &p_proj,
                     const std::vector<Tensor> &C_proj,
                     const std::vector<Tensor> &F)
{

    Sequence op;

    std::vector<Tensor> p_smoothOut(p_filt.size());
    std::vector<Tensor> C_smoothOut(p_filt.size());

    for (uint i = 0; i < p_filt.size(); i++)
    {

        std::string iStr = std::to_string(i);

        // Compute A

        auto [computeInverseC_proj, invOutC_proj] = inverse(graph, C_proj[i], i, 4);
        op.add(Execute(computeInverseC_proj));

        auto [computeFC, fc] = product(graph, F[i].transpose(), invOutC_proj, i);
        op.add(Execute(computeFC));

        auto [computeA, a] = product(graph, C_filt[i], fc, i);
        op.add(Execute(computeA));

        auto [computePDiff, pDiff] =
            scaledAdd(graph, p_smooth_prev[i], p_proj[i], i, -1);
        op.add(Execute(computePDiff));

        auto [computeCDiff, cDiff] =
            scaledAdd(graph, C_smooth_prev[i], C_proj[i], i, -1);
        op.add(Execute(computeCDiff));

        auto [computeAPDiff, aPDiff] = product(graph, a, pDiff, i);
        op.add(Execute(computeAPDiff));

        auto [computeP_smooth, p_smooth] = scaledAdd(graph, p_filt[i], aPDiff, i);
        op.add(Execute(computeP_smooth));

        auto [computeCdiffA, cDiffA] = product(graph, cDiff, a.transpose(), i);
        op.add(Execute(computeCdiffA));

        auto [computeACdiffA, aCDiffA] = product(graph, a, cDiffA, i);
        op.add(Execute(computeACdiffA));

        auto [computeC_smooth, C_smooth] = scaledAdd(graph, C_filt[i], aCDiffA, i);
        op.add(Execute(computeC_smooth));

        p_smoothOut[i] = graph.addVariable(FLOAT, {4, 1}, "p_smooth" + iStr);
        C_smoothOut[i] = graph.addVariable(FLOAT, {4, 4}, "C_smooth" + iStr);

        graph.setTileMapping(p_smoothOut[i], i);
        graph.setTileMapping(C_smoothOut[i], i);

        op.add(Copy(p_smooth, p_smoothOut[i]));
        op.add(Copy(C_smooth, C_smoothOut[i]));

    }

    return std::make_tuple(op, p_smoothOut, C_smoothOut);

} //~ end KalmanFilter::smooth


/*
 Copyright (C) 2020 Quaternion Risk Management Ltd
 All rights reserved.

 This file is part of ORE, a free-software/open-source library
 for transparent pricing and risk analysis - http://opensourcerisk.org

 ORE is free software: you can redistribute it and/or modify it
 under the terms of the Modified BSD License.  You should have received a
 copy of the license along with this program.
 The license is also available online at <http://opensourcerisk.org>

 This program is distributed on the basis that it will form a useful
 contribution to risk analytics and model standardisation, but WITHOUT
 ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 FITNESS FOR A PARTICULAR PURPOSE. See the license for more details.
*/

#include <orea/aggregation/dimcalculator.hpp>
#include <ored/utilities/log.hpp>
#include <ored/utilities/vectorutils.hpp>
#include <ql/errors.hpp>
#include <ql/time/calendars/weekendsonly.hpp>
#include <ql/version.hpp>

#include <ql/math/distributions/normaldistribution.hpp>
#include <ql/math/generallinearleastsquares.hpp>
#include <ql/math/kernelfunctions.hpp>
#include <ql/methods/montecarlo/lsmbasissystem.hpp>
#include <ql/time/daycounters/actualactual.hpp>

#include <qle/math/nadarayawatson.hpp>
#include <qle/math/stabilisedglls.hpp>

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/error_of_mean.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/stats.hpp>

using namespace std;
using namespace QuantLib;

using namespace boost::accumulators;

namespace ore {
namespace analytics {

DynamicInitialMarginCalculator::DynamicInitialMarginCalculator(
    const boost::shared_ptr<Portfolio>& portfolio, const boost::shared_ptr<NPVCube>& cube,
    const boost::shared_ptr<CubeInterpretation>& cubeInterpretation,
    const boost::shared_ptr<AggregationScenarioData>& scenarioData, Real quantile, Size horizonCalendarDays,
    const std::map<std::string, Real>& currentIM)
    : portfolio_(portfolio), cube_(cube), cubeInterpretation_(cubeInterpretation), scenarioData_(scenarioData),
      quantile_(quantile), horizonCalendarDays_(horizonCalendarDays), currentIM_(currentIM) {

    QL_REQUIRE(portfolio_, "portfolio is null");
    QL_REQUIRE(cube_, "cube is null");
    QL_REQUIRE(cubeInterpretation_, "cube interpretation is null");
    QL_REQUIRE(scenarioData_, "aggregation scenario data is null");

    boost::shared_ptr<RegularCubeInterpretation> regCubeInt =
        boost::dynamic_pointer_cast<RegularCubeInterpretation>(cubeInterpretation_);
    cubeIsRegular_ = (regCubeInt != NULL);
    datesLoopSize_ = cubeIsRegular_ ? cube_->dates().size() - 1 : cube_->dates().size();

    Size dates = cube_->dates().size();
    Size samples = cube_->samples();

    if (!cubeInterpretation_->hasMporFlows(cube_)) {
        WLOG("cube holds no mpor flows, will assume no flows in the dim calculation");
    }

    // initialise aggregate NPV and Flow by date and scenario
    set<string> nettingSets;
    for (Size i = 0; i < portfolio_->size(); ++i) {
        string tradeId = portfolio_->trades()[i]->id();
        string nettingSetId = portfolio_->trades()[i]->envelope().nettingSetId();
        if (nettingSets.find(nettingSetId) == nettingSets.end()) {
            nettingSets.insert(nettingSetId);
            nettingSetNPV_[nettingSetId] = vector<vector<Real>>(dates, vector<Real>(samples, 0.0));
            nettingSetCloseOutNPV_[nettingSetId] = vector<vector<Real>>(dates, vector<Real>(samples, 0.0));
            nettingSetFLOW_[nettingSetId] = vector<vector<Real>>(dates, vector<Real>(samples, 0.0));
            nettingSetDeltaNPV_[nettingSetId] = vector<vector<Real>>(dates, vector<Real>(samples, 0.0));
            nettingSetDIM_[nettingSetId] = vector<vector<Real>>(dates, vector<Real>(samples, 0.0));
            nettingSetExpectedDIM_[nettingSetId] = vector<Real>(dates, 0.0);
        }
        for (Size j = 0; j < datesLoopSize_; ++j) {
            for (Size k = 0; k < samples; ++k) {
                Real defaultNpv = cubeInterpretation_->getDefaultNpv(cube_, i, j, k);
                Real closeOutNpv = cubeInterpretation_->getCloseOutNpv(cube_, i, j, k);
                Real mporFlow =
                    cubeInterpretation_->hasMporFlows(cube_) ? cubeInterpretation_->getMporFlows(cube_, i, j, k) : 0.0;
                nettingSetNPV_[nettingSetId][j][k] += defaultNpv;
                nettingSetCloseOutNPV_[nettingSetId][j][k] += closeOutNpv;
                nettingSetFLOW_[nettingSetId][j][k] += mporFlow;
            }
        }
    }

    for (auto n : nettingSets)
        nettingSetIds_.push_back(n);

    dimCube_ = boost::make_shared<SinglePrecisionInMemoryCube>(cube_->asof(), nettingSetIds_, cube_->dates(),
                                                               cube_->samples());
}

const vector<vector<Real>>& DynamicInitialMarginCalculator::dynamicIM(const std::string& nettingSet) {
    if (nettingSetDIM_.find(nettingSet) != nettingSetDIM_.end())
        return nettingSetDIM_[nettingSet];
    else
        QL_FAIL("netting set " << nettingSet << " not found in DIM results");
}

const vector<Real>& DynamicInitialMarginCalculator::expectedIM(const std::string& nettingSet) {
    if (nettingSetExpectedDIM_.find(nettingSet) != nettingSetExpectedDIM_.end())
        return nettingSetExpectedDIM_[nettingSet];
    else
        QL_FAIL("netting set " << nettingSet << " not found in expected DIM results");
}

const vector<vector<Real>>& DynamicInitialMarginCalculator::cashFlow(const std::string& nettingSet) {
    if (nettingSetFLOW_.find(nettingSet) != nettingSetFLOW_.end())
        return nettingSetFLOW_[nettingSet];
    else
        QL_FAIL("netting set " << nettingSet << " not found in DIM results");
}

void DynamicInitialMarginCalculator::exportDimEvolution(ore::data::Report& dimEvolutionReport) {

    Size samples = dimCube_->samples();
    Size stopDatesLoop = datesLoopSize_;
    Date asof = cube_->asof();

    dimEvolutionReport.addColumn("TimeStep", Size())
        .addColumn("Date", Date())
        .addColumn("DaysInPeriod", Size())
        .addColumn("AverageDIM", Real(), 6)
        .addColumn("AverageFLOW", Real(), 6)
        .addColumn("NettingSet", string())
        .addColumn("Time", Real(), 6);

    for (auto nettingSet : dimCube_->ids()) {

        LOG("Export DIM evolution for netting set " << nettingSet);
        for (Size i = 0; i < stopDatesLoop; ++i) {
            Real expectedFlow = 0.0;
            for (Size j = 0; j < samples; ++j) {
                expectedFlow += nettingSetFLOW_[nettingSet][i][j] / samples;
            }

            Date defaultDate = dimCube_->dates()[i];
	    Time t = ActualActual().yearFraction(asof, defaultDate);
            Size days = cubeInterpretation_->getMporCalendarDays(dimCube_, i);
            dimEvolutionReport.next()
                .add(i)
                .add(defaultDate)
                .add(days)
                .add(nettingSetExpectedDIM_[nettingSet][i])
                .add(expectedFlow)
                .add(nettingSet)
	        .add(t);
        }
    }
    dimEvolutionReport.end();
    LOG("Exporting expected DIM through time done");
}

} // namespace analytics
} // namespace ore