#pragma once

#include <iostream>
#include "../generic/Algorithm.h"
#include "GaussianExpectationLib.h"
#include "TimeSliceCalibrator.h"
#include "../general_model/TimeSliceProcess.h"
#include "../general_model/ProcessCorrelations.h"
#include "../generic/ConnectedObject.h"
#include "MFMBase.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Stores model data relating to a given discrete time step Ti with an abitrary number of (possibly correlated)
// random processes
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <class CalibratorType>
class TimeSlice : public ConnectedObject<TimeSlice<CalibratorType>>
{

protected:
	TimeSlice<CalibratorType>(double time, const MFMBase* model, size_t number_of_random_processes)
		:
		ConnectedObject<TimeSlice<CalibratorType>>(),
		m_time(time),
		m_model(model),
		m_calibrator(std::unique_ptr<CalibratorType>()),
		m_correlations(std::unique_ptr<SliceProcessCorrelations>())
	{
		m_processes.resize(number_of_random_processes);
	}

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Random processes at this timeslice
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	std::vector<std::unique_ptr<TimeSliceProcess>> m_processes;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// model that this timeslice is attached to
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	const MFMBase* m_model;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// T_i
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	double m_time;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Contains calibration info at a time slice and everything needed to calibrate
	// at next time slice
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	std::unique_ptr<CalibratorType> m_calibrator;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Contains correlation data between processes at this timeslice
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	std::unique_ptr<SliceProcessCorrelations> m_correlations;

	static std::vector<std::unique_ptr<TimeSliceProcess>>& get_processes(TimeSlice* impl)
	{
		return impl->m_processes;
	}

public:

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Inside the time slice class we need to treat CalibratorType as TimeSliceCalibrator to use the general functions.
	// If the CalibratorType used when constructing the timeslice doesn't inherit from TimeSliceCalibrator,
	// then there will rightly be a compile error.
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	TimeSliceCalibrator& get_base_calibrator() const
	{
		return static_cast<TimeSliceCalibrator&>(*m_calibrator);
	}

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// call this before attaching calibrator
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void attach_correlations(std::unique_ptr<SliceProcessCorrelations>(*corr_constructor) (void), const std::vector<double>& prev_covariances, const std::vector<double>& covariances, const std::array<size_t, 3>& ordering)
	{
		m_correlations = corr_constructor();
		m_correlations->initialise(prev_covariances, covariances, ordering);
	}

	void attach_correlations(std::unique_ptr<SliceProcessCorrelations> corr)
	{
		m_correlations = std::move(corr);
	}

	virtual void attach_slice_calibrator(std::unique_ptr<CalibratorType>(*ad_constructor) (void))
	{
		m_calibrator = ad_constructor();

		std::vector<size_t> grid_sizes;

		grid_sizes.resize(m_processes.size());
		std::transform(m_processes.begin(), m_processes.end(), grid_sizes.begin(), [](std::unique_ptr < TimeSliceProcess>& p) -> size_t {return p->total_grid_points(); });

		get_base_calibrator().initialise(grid_sizes, m_model);

	}

	static std::unique_ptr<TimeSlice<CalibratorType>> construct(double time, const MFMBase* model, size_t number_of_random_processes)
	{
		std::unique_ptr<TimeSlice<CalibratorType>> Ti = std::unique_ptr<TimeSlice<CalibratorType>>(new TimeSlice<CalibratorType>(time, model, number_of_random_processes));
		return std::move(Ti);
	}

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Attach a random process at this time slice
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void add_random_process(size_t which_process, std::unique_ptr<TimeSliceProcess> process)
	{
		m_processes[which_process] = std::move(process);
	}

	virtual void calibrate_next(TimeSliceCalibrator& next_calibrator)
	{
		// Do any work required at slice k to calibrate time slice k + 1
		this->get_base_calibrator().calibrate(m_processes, this->next()->get_correlations());

		// calibrate time slice k + 1
		this->get_base_calibrator().calibrate_next(next_calibrator, this->m_processes, this->next()->m_processes, this->next()->get_correlations());

		// Now that k + 1 is calibrated, calculate spot values at k + 1
		this->get_base_calibrator().calculate_next_spot_values(this->next()->m_processes);
	}

	double get_option_in_arrears_price(size_t which_process, double strike) const
	{
		return m_processes[which_process]->get_option_in_arrears_price(strike);
	}

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

	///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Getters
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

	size_t total_grid_points(size_t which_process) const
	{
		return m_processes[which_process]->total_grid_points();
	}

	SliceProcessCorrelations const& get_correlations() const
	{
		return *m_correlations;
	}

	/////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Note here that we return whatever implementation of TimeSliceCalibrator was used when
	// constructing the TimeSlice, as it will likely contain functions specific to the model it was built
	// for that will need to be accessed.
	////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	CalibratorType& get_calibrator() const
	{
		return *m_calibrator;
	}

	TimeSliceProcess* get_process(size_t which_process) const
	{
		return m_processes[which_process].get();
	}

	double get_conditional_variance(size_t which_process) const
	{
		return m_processes[which_process]->get_conditional_variance();
	}

	double get_conditional_sigma(size_t which_process) const
	{
		return m_processes[which_process]->get_conditional_standard_deviation();
	}

	double get_process_variance(size_t which_process) const
	{
		return m_processes[which_process]->get_process_variance();
	}

	void get_grid(std::vector<double>& out, size_t which_process) const
	{
		out = *m_processes[which_process]->get_grid();
	}

	const gelib::d_vector& get_grid(size_t which_process) const
	{
		return *m_processes[which_process]->get_grid();
	}

	void get_conditional_means(std::vector<double>& out, size_t which_process) const
	{
		out = *m_processes[which_process]->get_conditional_means();
	}

	const gelib::d_vector& get_conditional_means(size_t which_process) const
	{
		return *m_processes[which_process]->get_conditional_means();
	}

	gelib::d_vector& get_J0_grid(size_t which_process) const
	{
		return *m_processes[which_process]->get_J0_grid();
	}

	void get_J0_grid(std::vector<double>& out, size_t which_process) const
	{
		out = *m_processes[which_process]->get_J0_grid();
	}

	void get_spot_grid(std::vector<double>& out, size_t which_process) const
	{
		out = *m_processes[which_process]->get_spot();
	}

	void get_arrow_debreu_grid(std::vector<double>& out) const
	{
		out = this->get_base_calibrator().get();
	}


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

};

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Special properties of first time slice T_0
//
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

template <class CalibratorType>
class TimeSliceZero : public TimeSlice<CalibratorType>
{
	TimeSliceZero(const MFMBase* model, size_t number_of_random_processes)
		:
		TimeSlice<CalibratorType>(0.0, model, number_of_random_processes)
	{}

public:

	static std::unique_ptr<TimeSlice<CalibratorType>> construct(const MFMBase* model, size_t number_of_random_processes)
	{
		return std::move(std::unique_ptr<TimeSliceZero<CalibratorType>>(new TimeSliceZero<CalibratorType>(model, number_of_random_processes)));
	}

	void calibrate_next(TimeSliceCalibrator& next_calibrator) override
	{
		// Do any special calibration that may be required at T0
		this->get_base_calibrator().calibrate_at_T0(this->m_processes, this->next()->get_correlations());

		// calibrate T1
		this->get_base_calibrator().calibrate_next_at_T0(next_calibrator, this->m_processes, this->get_processes(this->next()), this->next()->get_correlations());

		// Calculate spot values at T1
		this->get_base_calibrator().calculate_next_spot_values(this->get_processes(this->next()));
	}
};

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Special properties of last time slice T_n
//
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

template <class CalibratorType>
class TimeSliceLast : public TimeSlice<CalibratorType>
{
	TimeSliceLast(double time, const MFMBase* model, size_t number_of_random_processes)
		:
		TimeSlice<CalibratorType>(time, model, number_of_random_processes)
	{}

public:

	static std::unique_ptr<TimeSlice<CalibratorType>> construct(double time, const MFMBase* model, size_t number_of_random_processes)
	{
		std::unique_ptr<TimeSliceLast<CalibratorType>> Ti = std::unique_ptr<TimeSliceLast<CalibratorType>>(
			new TimeSliceLast<CalibratorType>(time, model, number_of_random_processes));
		return std::move(Ti);
	}

	// shouldn't do anything as this is last slice
	void calibrate_next(TimeSliceCalibrator& next_calibrator) override {}

};
