#pragma once

#include "Timeslices.h"
#include "MFMBase.h"

//////////////////////////////////////////////////////////////////////////////////////
// General Markov Functional model, with an arbitrary number of random processes
// at each timeslice 
//
//
//////////////////////////////////////////////////////////////////////////////////////
template <class CalibratorType, class ... ProcessTypes>
class MFMGeneral : public MFMBase
{

protected:

	std::array<size_t, sizeof ... (ProcessTypes)> m_ordering;

	//////////////////////////////////////////////////////////////////////////////////////
	// Set to true when model is calibrated for valuing options
	//////////////////////////////////////////////////////////////////////////////////////
	bool m_bls_calibrated;

	//////////////////////////////////////////////////////////////////////////////////////
	// Array of timeslices, with first time always at 0
	//////////////////////////////////////////////////////////////////////////////////////
	size_t m_number_of_slices;
	TimeSlices<CalibratorType> m_time_slices;

	//////////////////////////////////////////////////////////////////////////////////////
	// largest possible grid-length in any time slice
	//////////////////////////////////////////////////////////////////////////////////////
	size_t m_gridsize_ub;

	size_t m_processes_per_slice;

	//////////////////////////////////////////////////////////////////////////////////////
	// Tell each time slice to construct a chosen calibrator using a the constructor
	// function passed as the argument
	//////////////////////////////////////////////////////////////////////////////////////
	void attach_slice_calibrator(std::unique_ptr<CalibratorType>(*ad_constructor) (void))
	{
		for (size_t i = 0; i < m_time_slices.size(); ++i)
			m_time_slices.get(i)->attach_slice_calibrator(ad_constructor);
	}

	//////////////////////////////////////////////////////////////////////////////////////
	// Set up the processes at each time slice once they are constructed.
	// In the input parameter, processes[i] signifies the i'th time slice,
	// while processes[i][j] is the j'th process at the slice
	//////////////////////////////////////////////////////////////////////////////////////
	void populate_processes(std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>>& processes)
	{
		assert(m_processes_per_slice == processes.size());
		for (size_t j = 0; j < processes.size(); ++j)
		{
			assert(m_number_of_slices == processes[j].size());
			for (size_t i = 0; i < m_number_of_slices; ++i)
				m_time_slices.get(i)->add_random_process(j, std::move(processes[j][i]));
		}
	}

	///////////////////////////////////////////////////////////////////////////////////////////////
	// Retrive array of DerviedProcess pointers from array of DerviedProcess unique_ptrs -
	// useful if model class requires references to the processes which are stored in timeslices
	///////////////////////////////////////////////////////////////////////////////////////////////
	template <class DerivedProcess>
	void switch_to_base_process(size_t process_number, std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>>& base_processes, std::vector<std::unique_ptr<DerivedProcess>>& derived_process)
	{

		base_processes[process_number].resize(derived_process.size());

		std::transform(derived_process.begin(), derived_process.end(), base_processes[process_number].begin(), [](std::unique_ptr<DerivedProcess>& p) -> std::unique_ptr<TimeSliceProcess> { return std::move(p); });

	}

	template <class DerivedProcess, class ... MoreDerivedProcesses>
	void switch_to_base_process(size_t process_number, std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>>& base_processes, 
		std::vector<std::unique_ptr<DerivedProcess>>& derived_process, std::vector<std::unique_ptr<MoreDerivedProcesses>>&... more_derived_process)
	{
		base_processes[process_number].resize(derived_process.size());

		switch_to_base_process<DerivedProcess>(process_number, base_processes, derived_process);

		switch_to_base_process(process_number + 1, base_processes, more_derived_process...);
	}


	template <const size_t N>
	void add_correlations(std::unique_ptr<SliceProcessCorrelations>(*corr_constructor)(void), const std::array<gelib::d_vector, N>& piecewise_variances, const std::array<gelib::d_vector, N>& piecewise_covariances, const std::array<size_t, N>& ordering)
	{ // TODO clean this up
		assert(piecewise_variances[0].size() == m_time_slices.size() - 1);
		for (size_t i = 0; i < N; ++i)
			assert(piecewise_variances[i].size() == piecewise_covariances[i].size());

		std::vector<double> prev_cov, current_cov;

		// T0
		m_time_slices.get(0)->attach_correlations(std::move(std::unique_ptr<SliceProcessCorrelationsTimeZero>(new SliceProcessCorrelationsTimeZero(N))));

		size_t k;

		// T1
		k = 0;
		for (size_t j = 0; j < N; ++j)
		{

			// elements before diagonal
			for (size_t l = 0; l < j; ++l)
			{
				prev_cov.push_back(0.0);
				current_cov.push_back(piecewise_covariances[k][0]);
				++k;
			}

			// jth diagonal
			prev_cov.push_back(0.0);
			current_cov.push_back(piecewise_variances[j][0]);

		}

		assert(prev_cov.size() == N * (N + 1) / 2);
		m_time_slices.get(1)->attach_correlations(corr_constructor, prev_cov, current_cov, ordering);
		prev_cov.clear();
		current_cov.clear();


		// Ti where i > 1

		prev_cov.reserve(N * (N + 1) / 2);
		current_cov.reserve(N * (N + 1) / 2);

		for (size_t i = 1; i < m_time_slices.size() - 1; ++i)
		{
			k = 0;
			for (size_t j = 0; j < N; ++j)
			{

				// elements before diagonal
				for (size_t l = 0; l < j; ++l)
				{
					prev_cov.push_back(piecewise_covariances[k][i - 1]);
					current_cov.push_back(piecewise_covariances[k][i]);
					++k;
				}

				// jth diagonal
				prev_cov.push_back(piecewise_variances[j][i - 1]);
				current_cov.push_back(piecewise_variances[j][i]);

			}

			assert(prev_cov.size() == N * (N + 1) / 2);
			m_time_slices.get(i + 1)->attach_correlations(corr_constructor, prev_cov, current_cov, ordering);
			prev_cov.clear();
			current_cov.clear();
		}
	}

	void initialise(size_t grid_size_ub, gelib::d_vector const& times, size_t number_of_processes, std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>>& processes)
	{
		m_gridsize_ub = grid_size_ub;
		m_processes_per_slice = number_of_processes;
		m_bls_calibrated = false;

		size_t n = times.size();
		m_number_of_slices = n + 1;

		this->m_time_slices.claim(0, TimeSliceZero<CalibratorType>::construct(this, number_of_processes));

		for (size_t i = 1; i < n; ++i)
			this->m_time_slices.claim(i, TimeSlice<CalibratorType>::construct(times[i - 1], this, number_of_processes));

		this->m_time_slices.claim(n, TimeSliceLast<CalibratorType>::construct(times[n - 1], this, number_of_processes));

		this->m_time_slices.connect();

		this->populate_processes(processes);

		// now that all the time slices are constructed, make sure we have the largest grid length
		for (size_t i = 0; i < m_number_of_slices; ++i)
			m_gridsize_ub = std::max<size_t>(m_gridsize_ub, m_time_slices[i]->total_grid_points(0));
	}

	MFMGeneral(size_t poly_order, size_t number_of_slices)
		:
		MFMBase(poly_order),
		m_time_slices(number_of_slices)
	{}

	void initialise(size_t grid_size_ub, gelib::d_vector const& times, size_t number_of_processes, std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>>& processes, std::unique_ptr<CalibratorType>(*ad_constructor)(void))
	{
		initialise(grid_size_ub, times, number_of_processes, processes);
		attach_slice_calibrator(ad_constructor);
	}
	// TODO think its safe to remove initialse
	template <const size_t N>
	void initialise(size_t grid_size_ub, gelib::d_vector const& times, size_t number_of_processes, std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>>& processes, std::unique_ptr<CalibratorType>(*ad_constructor) (void), std::unique_ptr<SliceProcessCorrelations>(*corr_constructor)(void),
		const std::array<gelib::d_vector, N>& variances, const std::array<gelib::d_vector, N>& covariances, const std::array<size_t, 3>& ordering)
	{
		initialise(grid_size_ub, times, number_of_processes, processes);
		add_correlations<N>(corr_constructor, variances, covariances, ordering);
		attach_slice_calibrator(ad_constructor);
	}

public:

	MFMGeneral(
		size_t poly_order, size_t grid_size_ub, gelib::d_vector const& times,
		std::unique_ptr<CalibratorType>(*ad_constructor) (void),
		std::vector<std::unique_ptr<ProcessTypes>>... constructed_processes) 
		:
		MFMGeneral(poly_order, times.size() + 1)
	{
		const size_t number_of_processes = sizeof...(ProcessTypes);
		// need base class pointers to populate time slices
		std::vector<std::vector<std::unique_ptr<TimeSliceProcess>>> b_processes;
		b_processes.resize(number_of_processes);
		MFMGeneral::switch_to_base_process(0, b_processes, constructed_processes...);
		this->initialise(grid_size_ub, times, number_of_processes, b_processes, ad_constructor);
	}

	MFMGeneral(
		size_t poly_order, size_t grid_size_ub, gelib::d_vector const& times,
		std::unique_ptr<CalibratorType>(*ad_constructor) (void),
		std::unique_ptr<SliceProcessCorrelations>(*corr_constructor)(void),
		const std::array<gelib::d_vector, sizeof...(ProcessTypes)>& variances,
		const std::array<gelib::d_vector, sizeof...(ProcessTypes)>& covariances,
		const std::array<size_t, sizeof...(ProcessTypes)>& ordering,
		std::vector<std::unique_ptr<ProcessTypes>>... constructed_processes)
		:
		MFMGeneral(poly_order, grid_size_ub, times, ad_constructor, std::move(constructed_processes)...)
	{
		std::copy(ordering.begin(), ordering.end(), m_ordering.begin());
		this->add_correlations<sizeof...(ProcessTypes)>(corr_constructor, variances, covariances, ordering);
		attach_slice_calibrator(ad_constructor);
	}

	//////////////////////////////////////////////////////////////////////////////////////
	// prepare model to value options
	//////////////////////////////////////////////////////////////////////////////////////
	virtual void calibrate()
	{

		for (size_t i = 0; i < m_time_slices.size() - 1; ++i)
			m_time_slices.get(i)->calibrate_next(m_time_slices.get(i + 1)->get_base_calibrator());

		m_bls_calibrated = true;
	}

	//////////////////////////////////////////////////////////////////////////////////////
	// Calculate the option in arrears function/function inverse at a particular time slice and process. 
	// These functions are specific to how they are implemented in the processes 
	//////////////////////////////////////////////////////////////////////////////////////
	double get_option_in_arrears_price(size_t slice, size_t which_process, double strike) const
	{
		return m_time_slices[slice]->get_option_in_arrears_price(which_process, strike);
	}

	double get_option_in_arrears_price_inverse(size_t slice, size_t which_process, double value) const
	{
		return m_time_slices[slice]->get_option_in_arrears_price_inverse(which_process, value);
	}


	//////////////////////////////////////////////////////////////////////////////////////
	// Extract data from underlying objects
	//////////////////////////////////////////////////////////////////////////////////////

	void get_gridded_data(void (TimeSlice<CalibratorType>::* data_func) (std::vector<double>&) const, gelib::d_matrix& out) const
	{
		// If grids aren't all the same size, the shorter grids are filled with zeros
		out.resize(m_gridsize_ub, m_number_of_slices);
		gelib::d_vector v_data(m_gridsize_ub);
		// iterate over columns (times slices)
		for (size_t i = 0; i < m_number_of_slices; ++i)
		{
			(m_time_slices[i]->*data_func)(v_data);
			out.insert_vector(v_data, -1, (int)i);
		}
	}

	void get_gridded_data(void (TimeSlice<CalibratorType>::* data_func) (std::vector<double>&, size_t) const, size_t which_process, gelib::d_matrix& out) const
	{
		// If grids aren't all the same size, the shorter grids are filled with zeros
		out.resize(m_gridsize_ub, m_number_of_slices);
		gelib::d_vector v_data(m_gridsize_ub);
		// iterate over columns (times slices)
		for (size_t i = 0; i < m_number_of_slices; ++i)
		{
			(m_time_slices[i]->*data_func)(v_data, which_process);
			out.insert_vector(v_data, -1, (int)i);
		}
	}

	void get_grids(gelib::d_matrix& out, size_t which_process) const
	{
		get_gridded_data(&TimeSlice<CalibratorType>::get_grid, which_process, out);
	}

	void get_conditional_means(gelib::d_matrix& out, size_t which_process) const
	{
		get_gridded_data(&TimeSlice<CalibratorType>::get_conditional_means, which_process, out);
	}

	void get_ad_grids(gelib::d_matrix& out) const
	{
		get_gridded_data(&TimeSlice<CalibratorType>::get_arrow_debreu_grid, out);
	}

	void get_J0_grids(gelib::d_matrix& out, size_t which_process) const
	{
		get_gridded_data(&TimeSlice<CalibratorType>::get_J0_grid, which_process, out);
	}

	void get_spot_grids(gelib::d_matrix& out, size_t which_process) const
	{
		get_gridded_data(&TimeSlice<CalibratorType>::get_spot_grid, which_process, out);
	}

	void get_variances(gelib::d_vector& out, size_t which_process) const
	{
		out.resize(m_number_of_slices);
		for (size_t i = 0; i < m_number_of_slices; ++i)
			out[i] = m_time_slices[i]->get_process_variance(which_process);
	}

	void get_slices(std::vector<const TimeSlice<CalibratorType>*>& slices) const
	{
		slices.resize(m_number_of_slices);
		for (size_t i = 0; i < m_number_of_slices; ++i)
			slices[i] = m_time_slices[i];
	}


	//////////////////////////////////////////////////////////////////////////////////////
	// Retrieve a pointer to chosen process at a particular time slice
	//////////////////////////////////////////////////////////////////////////////////////
	template <class ProcessType>
	ProcessType* get_process(size_t which_process, size_t slice) const
	{
		return static_cast<ProcessType*>(m_time_slices[slice]->get_process(which_process));
	}

	//////////////////////////////////////////////////////////////////////////////////////
	// End of getters
	//////////////////////////////////////////////////////////////////////////////////////

	static void constructor_check(size_t gridpoints_per_slice, double sd_spread, size_t poly_order, gelib::d_vector const& times)
	{
		std::string err = "";
		const size_t max_poly_order = gelib_impl::Polynomial1d::MAXIMUM_DEGREE;
		const size_t max_grid_points = 100000000;

		for (size_t i = 0; i < times.size() - 1; i++)
		{
			if (times[i + 1] - times[i] < 0.0)
			{
				err = "At i = " + std::to_string(i) + " T_{i+1} - T_{i} is negative.";
			}
		}

		if (times.back() > 1000)
			err = "Times should not exceed 1000";
		if (sd_spread <= 0)
			err = "Spread should be positive.";
		if (poly_order > max_poly_order)
			err = "Polynomial order is too large or a negative parameter was supplied.";
		if (times.size() < 2)
			err = "There should be at least 2 times";
		if (gridpoints_per_slice > max_grid_points)
			err = "Maximum grid points exceeded, was a negative parameter supplied?";

		gelib_impl::IO::throw_on_input_error(err);
	}

};