#pragma once

#include <iostream>
#include <functional>
#include "GaussianExpectationLib.h"
#include "../general_model/MFMGeneral.h"
#include "../random_process/InterestRateTimeSliceProcess.h"
#include "../random_process/FXTimeSliceProcess.h"
#include "SliceValuatorTwoCurrency.h"
#include "arrow_debreu/CalibratorTwoCurrency.h"
#include "../general_model/Valuator.h"
#include "SpreadFactory3d.h"
#include <string>

//////////////////////////////////////////////////////////////////////////////////////
// A general 3 factor Markov functional model with two currencies and an FX rate as processes,
// this class implements the common functionality between 
// a Hull-White model "HullWhiteMFMTwoCurrency" 
// or the more general "ADMFMTwoCurrency"
//////////////////////////////////////////////////////////////////////////////////////
template <class CalibratorType>
class MFMTwoCurrency : public MFMGeneral<CalibratorType, InterestRateTimeSliceProcess, FXTimeSliceProcess, InterestRateTimeSliceProcess>
{
private:
	void set_up_valuator()
	{
		size_t i = 0;

		// Used for valuing spreads at T1
		m_valuator.claim(i,
			new SliceValuatorTwoCurrencyAtTZero(this,
				static_cast<SliceProcessCorrelations3d const&>(this->m_time_slices[i + 1]->get_correlations()),
				static_cast<SliceProcessCorrelations3d const&>(this->m_time_slices[i + 2]->get_correlations()),
				get_domestic_process(i), get_fx_process(i), get_foreign_process(i),
				m_t1_discount_factor_domestic));

		++i;

		// Used for valuing spreads at T2 to Tn-1
		for (; i < m_valuator.size() - 1; ++i)
			m_valuator.claim(i,
				new SliceValuatorTwoCurrencyAtTi(this,
					static_cast<SliceProcessCorrelations3d const&>(this->m_time_slices[i + 1]->get_correlations()),
					static_cast<SliceProcessCorrelations3d const&>(this->m_time_slices[i + 2]->get_correlations()),
					get_domestic_process(i), get_fx_process(i), get_foreign_process(i), get_fx_process(i + 1)
				));

		// Used for valuing spreads at Tn
		m_valuator.claim(i,
			new SliceValuatorTwoCurrencyAtTLast(this,
				static_cast<SliceProcessCorrelations3d const&>(this->m_time_slices[i + 1]->get_correlations()),
				get_domestic_process(i), get_fx_process(i), get_foreign_process(i),
				get_domestic_process(i + 1), get_fx_process(i + 1), get_foreign_process(i + 1)
			));

		// Set up the valuator so it can communicate with valuators at adjacent time slices
		m_valuator.connect();
	}

protected:

	static std::array<gelib::d_vector, 3> variances_to_squared_volatilites(double mean_reversion_param, gelib::d_vector const& times, std::array<gelib::d_vector, 3> const& process_variances)
	{
		std::array<gelib::d_vector, 3> process_volatilities_squared;
		for (size_t i = 0; i < process_volatilities_squared.size(); i++)
			TimeSliceProcess::variances_to_squared_volatilites(mean_reversion_param, times, process_variances[i], process_volatilities_squared[i]);
		return std::move(process_volatilities_squared);
	}

	Valuator<SliceValuatorTwoCurrency> m_valuator;

	double m_mean_reversion_param_domestic;
	double m_mean_reversion_param_foreign;

	double m_t1_discount_factor_domestic;
	double m_t1_discount_factor_foreign;

	//////////////////////////////////////////////////////////////////////////////////////
	// multiple of standard deviations to make the grids - all grids will use the same currently
	//////////////////////////////////////////////////////////////////////////////////////
	double m_spread_factor;

	//////////////////////////////////////////////////////////////////////////////////////
	// input info about y process
	//////////////////////////////////////////////////////////////////////////////////////
	gelib::d_vector m_y_larrow_mean;
	gelib::d_vector m_y_larrow_variance; // TODO try and make the constructor indepenendent of these

	//////////////////////////////////////////////////////////////////////////////////////
	// This object is synced with compatible arrow debreu calibrators - if grid sizes are constant 
	// across all time slices. The calculation will be far more efficient as only one 
	// set of fits is needed.
	//////////////////////////////////////////////////////////////////////////////////////
	std::unique_ptr<gelib::BasisPiecewisePoly3d> m_arrow_debreu_calculator;

	InterestRateTimeSliceProcess* get_domestic_process(size_t slice) {
		return this->get_process<InterestRateTimeSliceProcess>(0, slice);
	}

	FXTimeSliceProcess* get_fx_process(size_t slice) {
		return this->get_process<FXTimeSliceProcess>(1, slice);
	}

	InterestRateTimeSliceProcess* get_foreign_process(size_t slice) {
		return this->get_process<InterestRateTimeSliceProcess>(2, slice);
	}

	//////////////////////////////////////////////////////////////////////////////////////
	// Decide whether to value spread with arrow debreu or otherwise depending on what
	// calibrator we have
	//////////////////////////////////////////////////////////////////////////////////////
	virtual double value_spread(size_t expiry_time_index, gelib::d_cube& spread)
	{
		// use the arrow debreu method
		if constexpr (CalibratorType::ARROW_DEBREU_CALIBRATOR_3D)
			return this->m_valuator.get(expiry_time_index - 1)->value_spread_eur_payoff(spread, this->m_time_slices.get(expiry_time_index - 1)->get_calibrator());
		else
			// roll back each increment
			return this->m_valuator.get(expiry_time_index - 1)->value_spread_eur_payoff(spread);
	}


public:

	//////////////////////////////////////////////////////////////////////////////////////
	// Constructor for a two currency MFM
	//////////////////////////////////////////////////////////////////////////////////////
	MFMTwoCurrency(
		size_t number_of_slices, size_t poly_order,
		const std::array<double, 2>& mean_reversion_params,
		const std::array<double, 2>& discount_factors_t1,
		gelib::d_vector const& y_larrow_mean, gelib::d_vector const& y_larrow_variance,
		size_t gridpoints_per_slice, double sd_spread, 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, 3>& variances,
		const std::array<gelib::d_vector, 3>& covariances,
		std::vector<std::unique_ptr<InterestRateTimeSliceProcess>> dom_processes,
		std::vector<std::unique_ptr<InterestRateTimeSliceProcess>> for_processes,
		std::vector<std::unique_ptr<FXTimeSliceProcess>> fx_processes)
		:
		// Call MFMGeneral constructor specialised with the three processes
		MFMGeneral<CalibratorType, InterestRateTimeSliceProcess, FXTimeSliceProcess, InterestRateTimeSliceProcess>(poly_order, gridpoints_per_slice, times,
			// provide a constructor for the calibrator used by this model 
			ad_constructor, corr_constructor,
			variances, covariances, std::array<size_t, 3>({ 1, 2, 0 }), // TODO don't hard-code ordering here
			std::move(dom_processes), std::move(fx_processes), std::move(for_processes)),
		m_valuator(number_of_slices - 1),
		m_t1_discount_factor_domestic(discount_factors_t1[0]),
		m_t1_discount_factor_foreign(discount_factors_t1[1]),
		m_mean_reversion_param_domestic(mean_reversion_params[0]),
		m_mean_reversion_param_foreign(mean_reversion_params[1]),
		m_y_larrow_mean(y_larrow_mean),
		m_y_larrow_variance(y_larrow_variance),
		m_spread_factor(sd_spread)
	{}
	
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// This function can be called to set a common arrow debreu calculator across time slices if the grid sizes are kept constant
	// at each slice, note that it's compilation success depends on the CalibratorType of this model being a TwoCurrencyArrowDebreuCalibrator
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	void set_up_arrow_debreu_calculator()
	{
		m_arrow_debreu_calculator = std::move(TwoCurrencyArrowDebreuCalibrator::construct_arrow_debreu_calculator(this->m_time_slices, this->m_ordering));
	}
	
	//////////////////////////////////////////////////////////////////////////////////////
	// override to default calibrator
	//////////////////////////////////////////////////////////////////////////////////////
	virtual void calibrate() override
	{
		// need to set variances for before calibrating TODO probably remove once fx variances are calculated within code rather than input 
		for (size_t i = 1; i < this->m_number_of_slices; ++i)
			get_fx_process(i)->set_uncorrelated_grid_variance(m_y_larrow_variance[i - 1]);

		// Try and set up efficient arrow debreu calculator
		if constexpr (CalibratorType::ARROW_DEBREU_CALIBRATOR_3D)
			set_up_arrow_debreu_calculator();

		MFMGeneral<CalibratorType, InterestRateTimeSliceProcess, FXTimeSliceProcess, InterestRateTimeSliceProcess>::calibrate();

		// Now that we have libors/FX, we can construct everything needed for product valuation
		set_up_valuator();

	}

	double value_payoff_of_1(size_t expiry_time_index)
	{
		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		gelib::d_cube spread;
		spread.resize(get_domestic_process(expiry_time_index)->total_grid_points(),
			get_fx_process(expiry_time_index)->total_grid_points(),
			get_foreign_process(expiry_time_index)->total_grid_points());

		std::fill(spread.begin(), spread.end(), 1.0);

		// The way the spread is valued depends on the derived class implementation of value_spread
		return value_spread(expiry_time_index, spread);
	}

	double value_dom_fx_spread(size_t expiry_time_index, double A, double B)
	{
		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		gelib::d_cube spread;
		SpreadFactory3d::calculate_dom_fx_spread(spread, A, B,
			get_domestic_process(expiry_time_index), get_fx_process(expiry_time_index), get_foreign_process(expiry_time_index));

		// The way the spread is valued depends on the derived class implementation of value_spread
		return value_spread(expiry_time_index, spread);
	}

	double value_fgn_fx_spread(size_t expiry_time_index, double A, double B)
	{
		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		gelib::d_cube spread;
		SpreadFactory3d::calculate_fgn_fx_spread(spread, A, B,
			get_domestic_process(expiry_time_index), get_fx_process(expiry_time_index), get_foreign_process(expiry_time_index));

		// The way the spread is valued depends on the derived class implementation of value_spread
		return value_spread(expiry_time_index, spread);
	}

	double value_dom_fgn_spread(size_t expiry_time_index, double A, double B)
	{
		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		gelib::d_cube spread;
		SpreadFactory3d::calculate_dom_fgn_spread(spread, A, B,
			get_domestic_process(expiry_time_index), get_fx_process(expiry_time_index), get_foreign_process(expiry_time_index));

		return value_spread(expiry_time_index, spread);
	}

	// TODO, move the MFM general get_arrow_debreu one to the 1d model
	void get_arrow_debreu(gelib::d_matrix& out, size_t which_slice)
	{
		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model trying to retrieve arrow debreu");

		if (!CalibratorType::ARROW_DEBREU_CALIBRATOR_3D)
			throw std::runtime_error("Calibrator used by this model does not create Arrow Debreu prices");

		out.resize(
			get_domestic_process(which_slice)->total_grid_points() * get_foreign_process(which_slice)->total_grid_points(),
			get_fx_process(which_slice)->total_grid_points());

		std::copy(this->m_time_slices.get(which_slice)->get_calibrator().get().begin(), this->m_time_slices.get(which_slice)->get_calibrator().get().end(), out.begin());

	}

	void get_domfx_xy_slice(gelib::d_matrix& out, double A, double B, size_t expiry_time_index, size_t which_value)
	{

		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		if (!m_valuator.get(which_value)->storing_values())
			throw std::runtime_error("The grids being used are too large. So the values at each time slice are not stored.");

		gelib::d_cube spread;
		SpreadFactory3d::calculate_dom_fx_spread(spread, A, B,
			get_domestic_process(expiry_time_index), get_fx_process(expiry_time_index), get_foreign_process(expiry_time_index));

		if (which_value == expiry_time_index)
		{
			spread.extract_matrix_layer(out, 0);
			return;
		}

		value_dom_fx_spread(expiry_time_index, A, B);

		// extract first layer to output
		SliceValuatorTwoCurrency::get_valuations(m_valuator.get(which_value)).extract_matrix_layer(out, 0);
	}

	void get_value_payoff_of_one(gelib::d_matrix& out, size_t expiry_time_index, size_t which_value)
	{

		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		if (!m_valuator.get(which_value)->storing_values())
			throw std::runtime_error("The grids being used are too large. So the values at each time slice are not stored.");

		gelib::d_cube spread;
		spread.resize(get_domestic_process(expiry_time_index)->total_grid_points(),
			get_fx_process(expiry_time_index)->total_grid_points(),
			get_foreign_process(expiry_time_index)->total_grid_points());

		std::fill(spread.begin(), spread.end(), 1.0);

		if (which_value == expiry_time_index)
		{
			out = spread.flatten();
			return;
		}

		value_payoff_of_1(expiry_time_index);

		// extract first layer to output
		out = SliceValuatorTwoCurrency::get_valuations(m_valuator.get(which_value)).flatten();
	}

	double value_single_arrow_debreu(size_t expiry_time_index, size_t x_index, size_t y_index, size_t z_index)
	{

		if (!this->m_bls_calibrated)
			throw std::runtime_error("Calibrate the model before spread valuation");

		gelib::d_cube spread;
		spread.resize(get_domestic_process(expiry_time_index)->total_grid_points(),
			get_fx_process(expiry_time_index)->total_grid_points(),
			get_foreign_process(expiry_time_index)->total_grid_points());

		std::fill(spread.begin(), spread.end(), 0.0);

		spread[spread.index(x_index, y_index, z_index)] = 1.0;

		return value_spread(expiry_time_index, spread);
	}

};