#ifndef GELIB_IMPL_SPLINETREE_H_
#define GELIB_IMPL_SPLINETREE_H_

#include "../n_dimensional/IndexSet.h"
#include "../1_dimensional/Spline1d.h"
#include "../generic_classes/Structures.h"
#include "Gaussian3d.h"
#include "d_cube.h"

namespace gelib_impl
{

	class Spline3d
	{

		///////////////////////////////////////////////////////////////////////////////
		// gradient of linear transform for underlying grid points, 
		// one for each dimension
		///////////////////////////////////////////////////////////////////////////////
		std::vector<double> m_linear_transform_a;

		///////////////////////////////////////////////////////////////////////////////
		// intercept of linear transform for underlying grid points,
		// one for each dimension
		///////////////////////////////////////////////////////////////////////////////
		std::vector<double> m_linear_transform_b;

		///////////////////////////////////////////////////////////////////////////////
		// grids used internally, centered at zero, one for 
		// each dimension
		///////////////////////////////////////////////////////////////////////////////
		std::vector<d_vector> m_transformed_grids;

		///////////////////////////////////////////////////////////////////////////////
		// Wrapper for the indices of each spline to output to some form 
		// of n-dimensional storage.
		///////////////////////////////////////////////////////////////////////////////
		IndexSet<3> m_indices;

		///////////////////////////////////////////////////////////////////////////////
		// The splines indexed by the indexed set
		///////////////////////////////////////////////////////////////////////////////
		std::vector<std::vector<std::unique_ptr<Spline1d>>> m_splines;

	public:

		///////////////////////////////////////////////////////////////////////////////
		// construct the indices in one dimension only
		///////////////////////////////////////////////////////////////////////////////
		Spline3d(
			size_t varied_dimension,
			d_vector const& grid,
			d_cube const& function_data,
			d_cube& output_data);

		Spline3d(
			size_t varied_dimension,
			d_vector const& grid,
			d_cube const& function_data,
			d_cube& output_data, std::array<size_t, 3> ordering);


		///////////////////////////////////////////////////////////////////////////////
		// construct the indices in all dimensions
		///////////////////////////////////////////////////////////////////////////////
		Spline3d(
			std::vector<d_vector> const & grid_data,
			std::vector<const d_cube*>& function_data,
			std::vector<d_cube*>& output_data,
			std::array<size_t, 3> ordering);

		///////////////////////////////////////////////////////////////////////////////
		// linear transform arbitrary grid vectors to center at 0
		///////////////////////////////////////////////////////////////////////////////
		void transform_grid_vectors(
			std::vector<d_vector> const& grid_data,
			std::vector<d_vector>& transformed_grid_data);

		void reverse_transform_expectation(size_t which_grid, double* transform_a, double* transform_b, double* sigma, double* l, double* h) const;

		///////////////////////////////////////////////////////////////////////////////
		// fit splines to the chosen input cube passed to the impl_Spline_3d constructor
		///////////////////////////////////////////////////////////////////////////////
		void fit_input(d_cube const& input_data, size_t dimension, size_t order, char boundaries, double domain_tol, size_t which_fit);
		void fit_input(d_cube const& input_data, size_t dimension, size_t order, char boundaries, double domain_tol, size_t which_fit, size_t start, size_t end);

		///////////////////////////////////////////////////////////////////////////////
		// construct three dimensions and sort indices
		///////////////////////////////////////////////////////////////////////////////	
		void construct_indices(
			std::vector<d_vector>& grid_data,
			std::vector<const IndexedData<3>*>& function_data,
			std::vector<IndexedData<3>*>& output_data,
			std::array<size_t, 3> ordering);


		///////////////////////////////////////////////////////////////////////////////
		// get the size of index array with a dimension fixed
		///////////////////////////////////////////////////////////////////////////////
		const size_t num_indices(size_t which) const;

		///////////////////////////////////////////////////////////////////////////////
		// get the indices of the i'th (varied) indices with one dimension fixed
		///////////////////////////////////////////////////////////////////////////////
		const std::array<int, 3>& get_varied_indices(size_t which, size_t i) const;

		///////////////////////////////////////////////////////////////////////////////
		// get the i'th indexed spline with a dimension fixed
		///////////////////////////////////////////////////////////////////////////////
		const std::unique_ptr<Spline1d>& get_spline(size_t which, size_t i) const;

		///////////////////////////////////////////////////////////////////////////////
		// evaluate the 3d spline at data points provided by input_vector over a chosen dimension,
		// outputting to F_X
		///////////////////////////////////////////////////////////////////////////////
		void evaluate(d_vector& input_vector, size_t which, size_t dimension, d_cube& F_X);

		///////////////////////////////////////////////////////////////////////////////
		// For each pair of fixed points, transform input_vector (indexed by fixed points),
		// then evaluate
		///////////////////////////////////////////////////////////////////////////////
		void evaluate(d_vector& input_vector, const std::function<void(size_t, size_t, d_vector const&, d_vector&)>& transform, size_t which, size_t dimension, d_cube& F_X);

		///////////////////////////////////////////////////////////////////////////////
		// Calculate (single dimension) Gaussian expectations of a 3d spline tree, 
		// outputting to E_X
		///////////////////////////////////////////////////////////////////////////////
		void gaussian_expectation(Gaussian3d& gaussian, double l, double h, size_t dimension, d_cube& E_X, size_t which_fit) const;
		void gaussian_expectation(Gaussian3d& gaussian, double l, double h,
			size_t dimension, d_cube& E_X, size_t which_fit,
			size_t start, size_t end) const;

		///////////////////////////////////////////////////////////////////////////////
		// Calculate (single dimension) Gaussian expectations of the maximum over two 3d splines.
		// Assumes that st1 and st2 have identical grid points
		// and indexing
		///////////////////////////////////////////////////////////////////////////////
		static void gaussian_expectation_max(Spline3d const& st1, Spline3d const& st2, Gaussian3d& gaussian, double l, double h, size_t dimension, d_cube& E_X, size_t which_fit);
		static void gaussian_expectation_max(
			// expectation over max(st1, st2)
			Spline3d const& st1, Spline3d const& st2,
			// contains all info about r.v distribution
			Gaussian3d& gaussian,
			// limits of expectation
			double l, double h,
			size_t dimension, d_cube& E_X, size_t which_fit,
			size_t start, size_t end);

		static void fit_and_expectation(
			size_t which_fit, size_t dimension, double l, double h,
			std::array<const d_vector*, 3>& means, d_cube const& input, d_cube& output, Spline3d& st,
			Gaussian3d& gaussian, size_t order, char extrapolation, double domain_tol);

		static void fit_and_expectation_max(size_t which_fit, size_t dimension,
			double l, double h,
			std::array<const d_vector*, 3>& means,
			d_cube const& input1, d_cube const& input2, d_cube& output,
			Spline3d& st1, Spline3d& st2, Gaussian3d& gaussian, size_t order, char extrapolation, double domain_tol);

		///////////////////////////////////////////////////////////////////////////////
		// Allocate data in three cubes for three integrations
		///////////////////////////////////////////////////////////////////////////////
		static void prepare_expectation_3d(const std::array<d_vector*, 3>& a_hats_orig, std::vector<d_vector> const & grid_data, size_t P_1, size_t P_2, size_t P_3, d_cube& E_X_1, d_cube& E_X_2, d_cube& E_X_3);

		///////////////////////////////////////////////////////////////////////////////
		// Calculate Gaussian expectations of a 3d splines,
		// and covariance matrix C
		// outputting to E_X
		///////////////////////////////////////////////////////////////////////////////
		static void gaussian_expectation_correlated(
			d_vector & a_hat_1, d_vector & a_hat_2, d_vector & a_hat_3,
			d_matrix& C,
			// lower and higher extrapolation bounds in each dimension
			double l1, double h1, double l2, double h2, double l3, double h3,
			// grid points function defined on
			std::vector<d_vector> const & grid_data,

			d_cube const& function_data, d_cube& E_X,
			// spline fitting data
			size_t order, char boundaries, double domain_tol,
			size_t P_0, size_t P_1, size_t P_2
		);

		///////////////////////////////////////////////////////////////////////////////
		// Calculate Gaussian expectations of the max of two 3d splines,
		// E(max(st1, st2) with means in each dimension ahat_1, ahat_2, ahat_3
		// and covariance matrix C
		// outputting to E_X
		///////////////////////////////////////////////////////////////////////////////
		static void gaussian_expectation_correlated_max(
			d_vector& a_hat_1, d_vector& a_hat_2, d_vector& a_hat_3,
			d_matrix& C,
			// lower and higher extrapolation bounds in each dimension
			double l1, double h1, double l2, double h2, double l3, double h3,
			// grid points function defined on
			std::vector<d_vector> const & grid_data,
			// original function data
			d_cube& A1, d_cube& A2, d_cube& E_X,
			// spline fitting data
			size_t order, char boundaries, double domain_tol,
			size_t P_0, size_t P_1, size_t P_2, size_t num_threads = 1);

	};

}

#endif