#pragma once
#include "../../general_model/TimeSliceCalibrator.h"
#include "../../generic/Algorithm.h"
#include "../ProcessCorrelations3d.h"
#include "../../1_dimensional/Grid1d.h"
#include "../../random_process/InterestRateTimeSliceProcess.h"
#include "../../random_process/FXTimeSliceProcess.h"
#include "../../single_currency_model/CalibratorSingleCurrency.h"
#include "../../general_model/TimeSlices.h"

//////////////////////////////////////////////////////////////////////////////////////
// General Arrow Debreu calibrator for a two currency model
//////////////////////////////////////////////////////////////////////////////////////
class TwoCurrencyArrowDebreuCalibrator : public SingleCurrencyCalibrator
{
	//////////////////////////////////////////////////////////////////////////////////////
	// If recycling the same fits, keep a pointer to the global arrow debreu calculator
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::BasisPiecewisePoly3d* m_arrow_debreu_calculator_ref;

	//////////////////////////////////////////////////////////////////////////////////////
	// 1 / (1 + alpha L(x)) at this timeslice
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_vector m_domestic_factors;

	//////////////////////////////////////////////////////////////////////////////////////
	// 1 / (1 + alpha L(x)) * arrow debreu prices at this timeslice
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_cube m_B0_C;

	//////////////////////////////////////////////////////////////////////////////////////
	// 3d arrow debreu prices
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_cube m_arrow_debreu;

	//////////////////////////////////////////////////////////////////////////////////////
	// Size 3 vector containing x,y and z grids
	//////////////////////////////////////////////////////////////////////////////////////
	std::vector<gelib::d_vector> m_larrow_grid_k;

	//////////////////////////////////////////////////////////////////////////////////////
	// means cubes for the arrow debreu integral
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_cube m_mu1;
	gelib::d_cube m_mu2;
	gelib::d_cube m_mu3;
	gelib::d_cube m_fx_drift_grid;

	//////////////////////////////////////////////////////////////////////////////////////
	// Working arrays for J0y calculation
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_vector m_fixedy_mu3_vector;
	gelib::d_vector m_fixedy_arrow_debreu_factors;

	//////////////////////////////////////////////////////////////////////////////////////
	// Working arrays for J0z calculation
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_vector m_calN_grid;
	gelib::d_vector m_calN_plusmu;
	gelib::d_vector m_calN_function;
	gelib::d_vector m_Hhat_function;
	gelib::d_vector m_Hhat_interpolation;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Driver function to calibrate at this time slice
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void calibrate(
		InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, InterestRateTimeSliceProcess* fgn_process, 
		SliceProcessCorrelations3d const& next_increment);

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Driver function to calculate next Arrow-Debreu and J0
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void calibrate_next(
		TwoCurrencyArrowDebreuCalibrator& next_calibrator,
		const InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* fgn_process,
		InterestRateTimeSliceProcess* next_dom_process, FXTimeSliceProcess* next_fx_process, InterestRateTimeSliceProcess* next_fgn_process, 
		SliceProcessCorrelations3d const& next_increment);

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// calibrate next which is only called for time slice zero
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void calibrate_next_at_T0(TimeSliceCalibrator& next_calibrator,
		std::vector<std::unique_ptr<TimeSliceProcess>>& current_processes, std::vector<std::unique_ptr<TimeSliceProcess>>& next_processes,
		SliceProcessCorrelations const& next_increment) override;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// calibrate next which is not called at T0
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void calibrate_next_T0plus(
		TwoCurrencyArrowDebreuCalibrator& next_calibrator,
		const InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* fgn_process,
		InterestRateTimeSliceProcess* next_dom_process, FXTimeSliceProcess* next_fx_process, InterestRateTimeSliceProcess* next_fgn_process,
		SliceProcessCorrelations3d const& next_increment);

protected:

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Driver functions to calculate next J0y and z - In this class the generalised version is used, but it may be overriden in 
	// derived classes, these are called in calibrate_next 
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	virtual void calculate_next_J0y(
		TwoCurrencyArrowDebreuCalibrator& next_calibrator,
		const InterestRateTimeSliceProcess* dom_process, const FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* fgn_process,
		InterestRateTimeSliceProcess* next_dom_process, FXTimeSliceProcess* next_fx_process, InterestRateTimeSliceProcess* next_fgn_process,
		SliceProcessCorrelations3d const& next_increment);

	virtual void calculate_next_J0z(
		TwoCurrencyArrowDebreuCalibrator& next_calibrator,
		const InterestRateTimeSliceProcess* dom_process, const FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* fgn_process,
		InterestRateTimeSliceProcess* next_dom_process, FXTimeSliceProcess* next_fx_process, InterestRateTimeSliceProcess* next_fgn_process,
		SliceProcessCorrelations3d const& next_increment);

public:

	// This flag is used to indicate to compiler whether this class calculates Arrow Debreu prices
	static const bool ARROW_DEBREU_CALIBRATOR_3D = true;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	//-------------------------------------------------------------------------------------------------------------------
	// Override methods to the general calibrator
	//-------------------------------------------------------------------------------------------------------------------
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

	virtual void initialise(const std::vector<size_t>& grid_sizes, const MFMBase* model) override;

	void calibrate_at_T0(std::vector<std::unique_ptr<TimeSliceProcess>>& processes_at_T0, SliceProcessCorrelations const& T0_to_T1_increment) override;

	void calibrate(std::vector<std::unique_ptr<TimeSliceProcess>>& processes, SliceProcessCorrelations const& next_increment) override;

	void calibrate_next(TimeSliceCalibrator& next_calibrator, std::vector<std::unique_ptr<TimeSliceProcess>>& current_processes, std::vector<std::unique_ptr<TimeSliceProcess>>& next_processes, SliceProcessCorrelations const& next_increment) override;

	std::vector<double>& get() override;

	virtual double roll_back(std::vector<double> const& spread) const override;

	// Override default behaviour here as we will have already calculated spot FX when this is called 
	void calculate_next_spot_values(std::vector<std::unique_ptr<TimeSliceProcess>>& next_processes) override;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	//-------------------------------------------------------------------------------------------------------------------
	// End of Overrides to the general calibrator
	//-------------------------------------------------------------------------------------------------------------------
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

	static std::unique_ptr<TimeSliceCalibrator> construct();

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// This can be used to set a common arrow debreu calculator across time slices if the grid sizes are kept constant
	// at each slice
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	template<class ArrowDebreuCalibrator3D>
	static std::unique_ptr<gelib::BasisPiecewisePoly3d> construct_arrow_debreu_calculator(TimeSlices<ArrowDebreuCalibrator3D>& time_slices, const std::array<size_t, 3>& ordering)// noexcept
	{
		size_t n = time_slices.size();
		size_t number_processes = 3;

		// First check we have an arrow debreu calculator and grid sizes are constant with respect to timeslice
		for (size_t i = 1; i < n - 1; ++i)
		{
			for (size_t j = 0; j < number_processes; ++j)
			{
				if (time_slices[i]->get_calibrator().get_s_l_arrow()[j] != time_slices[i + 1]->get_calibrator().get_s_l_arrow()[j])
					return std::unique_ptr<gelib::BasisPiecewisePoly3d>(nullptr);
			}
		}

		// construct arrow debreu calibrator to be used at all time slices by setting it up for a calculation at T1
		gelib::BasisPiecewisePoly3d* calculator =
			new gelib::BasisPiecewisePoly3d(
				// note T1 uses arrow grids from the calibrator at T2
				time_slices[2]->get_calibrator().get_s_l_arrow(),
				time_slices[1]->get_calibrator().get_B_0_i(),
				time_slices[1]->get_calibrator().get_mu3(),
				time_slices[1]->get_calibrator().get_mu2(),
				time_slices[1]->get_calibrator().get_mu1(),
				time_slices[2]->get_calibrator().get_arrow_debreu(), ordering);

		// Give each calibrator a reference to the calculator we just created.
		for (size_t i = 1; i < n - 1; ++i)
			time_slices[i]->get_calibrator().set_arrow_debreu_calculator(calculator);

		// wrap the calculator in a unique_ptr to prevent memory leak
		return std::move(std::unique_ptr<gelib::BasisPiecewisePoly3d>(calculator));
	}

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// store pointer to arrow debreu calculator across all time slices
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void set_arrow_debreu_calculator(gelib::BasisPiecewisePoly3d* arrow_debreu_calculator);

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// Getters
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	std::vector<gelib::d_vector>& get_s_l_arrow();

	gelib::d_cube const& get_B_0_i() const;

	gelib::d_cube const& get_mu1() const;

	gelib::d_cube const& get_mu2() const;

	gelib::d_cube const& get_mu3() const;

	gelib::d_cube const& get_arrow_debreu() const;

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// End of getters
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
};
