#include "SliceValuatorTwoCurrency.h"

const std::string SliceValuatorTwoCurrency::EXTRAPOLATION = "linear";
const char SliceValuatorTwoCurrency::CHAR_EXTRAPOLATION = gelib::PiecewisePoly1d::parse_extrapolation_string(EXTRAPOLATION);
const size_t SliceValuatorTwoCurrency::STORE_VALUES_SIZE_CUTOFF = 30 * 30 * 30;

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//===================================================================================================================
// General Valuator
//===================================================================================================================
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

SliceValuatorTwoCurrency::SliceValuatorTwoCurrency(const MFMBase* model, SliceProcessCorrelations3d const& next_increment_covariance, const InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* for_process, double mean_shift)
	:
	m_y_mean_shift(mean_shift),
	m_model(model),
	m_dom_process(dom_process),
	m_fx_process(fx_process),
	m_for_process(for_process),
	m_ahat_1(dom_process->get_grid()->size()),
	m_ahat_2(fx_process->get_grid()->size()),
	m_ahat_3(for_process->get_grid()->size()),
	m_storing_values((dom_process->get_grid()->size() * fx_process->get_grid()->size() * for_process->get_grid()->size()) < SliceValuatorTwoCurrency::STORE_VALUES_SIZE_CUTOFF)
{
	m_Vk.fix_dimensions(dom_process->get_grid()->size(), fx_process->get_grid()->size(), for_process->get_grid()->size());
	m_Ck_Ckminus1 = next_increment_covariance.get_increment_covariance();
}

gelib::d_cube const& SliceValuatorTwoCurrency::roll_back_european_payoff()
{
	// Use non-max as we know we are passed the first roll back step
	gelib::d_cube& payoff = integrate_european_payoff();

	// get the payoff on the right grid
	interpolate_payoff(payoff);

	multiply_by_discount_factors();

	// call the roll back function from the previous time slice
	return prev()->roll_back_european_payoff();
}

gelib::d_cube const& SliceValuatorTwoCurrency::roll_back_european_payoff(gelib::d_cube& known_value)
{
	// Use max as we know we are on the first step
	gelib::d_cube& payoff = integrate_european_payoff_first(known_value);

	// get the payoff on the right grid
	interpolate_payoff(payoff);

	multiply_by_discount_factors();

	// call the roll back function from the previous time slice
	return prev()->roll_back_european_payoff();
}

double SliceValuatorTwoCurrency::roll_back_european_payoff(gelib::d_cube& known_value, TimeSliceCalibrator const& arrow_debreu_calibrator)
{

	// max integrate one timestep to get a smooth function
	gelib::d_cube& Vk_on_a2hat = integrate_european_payoff_first(known_value);

	gelib::d_cube const& Vk = interpolate_payoff(Vk_on_a2hat);

	return arrow_debreu_calibrator.roll_back(Vk);
}

void SliceValuatorTwoCurrency::integrate_european_payoff(gelib::d_vector const& grid1, gelib::d_vector const& grid2, gelib::d_vector const& grid3, gelib::d_cube& current_value, gelib::d_cube& out, gelib::d_matrix& covariance)
{
	std::string extrap = SliceValuatorTwoCurrencyAtTi::EXTRAPOLATION;

	// rollback the spread value
	gelib::correlated_gaussian_expectation_3d(
		out, grid1, grid2, grid3, current_value,
		m_ahat_1, m_ahat_2, m_ahat_3,
		covariance,
		MFMBase::minus_inf, MFMBase::plus_inf,
		MFMBase::minus_inf, MFMBase::plus_inf,
		MFMBase::minus_inf, MFMBase::plus_inf,
		m_model->get_order(), extrap, 2, 3, 1, MFMBase::DEBUG_ON);

	// clear up the memory for the input to the integrating function 
	if (!m_storing_values)
		std::vector<double>().swap(current_value);
}

void SliceValuatorTwoCurrency::integrate_european_payoff_first(gelib::d_vector const& grid1, gelib::d_vector const& grid2, gelib::d_vector const& grid3, gelib::d_cube& current_value, gelib::d_cube& out, gelib::d_matrix& covariance)
{
	std::string extrap = SliceValuatorTwoCurrencyAtTi::EXTRAPOLATION;

	// create cube filled with zeros
	gelib::d_cube zeros;
	zeros.resize(grid1.size(), grid2.size(), grid3.size());

	// rollback the spread value
	gelib::correlated_gaussian_expectation_max_3d(
		out, grid1, grid2, grid3, current_value, zeros,
		m_ahat_1, m_ahat_2, m_ahat_3,
		covariance,
		MFMBase::minus_inf, MFMBase::plus_inf,
		MFMBase::minus_inf, MFMBase::plus_inf,
		MFMBase::minus_inf, MFMBase::plus_inf,
		m_model->get_order(), extrap, 2, 3, 1, MFMBase::DEBUG_ON);

	// clear up the memory for the input to the integrating function
	if (!m_storing_values)
		std::vector<double>().swap(current_value);
}

double SliceValuatorTwoCurrency::value_spread_eur_payoff(gelib::d_cube& spread, TimeSliceCalibrator const& arrow_debreu_calibrator)
{
	return roll_back_european_payoff(spread, arrow_debreu_calibrator);
}

double SliceValuatorTwoCurrency::value_spread_eur_payoff(gelib::d_cube& spread)
{
	return roll_back_european_payoff(spread)[0];
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//===================================================================================================================
//
// Valuator at timeslice Ti where 0 < i < n 
//
//===================================================================================================================
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

gelib::d_cube const& SliceValuatorTwoCurrencyAtTi::interpolate_payoff(gelib::d_cube& payoff)
{
	// set up drift function here in case we don't have one
	m_fx_process->set_up_drift(m_next_increment_correlations, "linear", *m_next_fx_process->get_grid(), m_next_fx_process->get_drift_fitting_grid());

	const auto y_to_a2hat = [&](size_t z_index, size_t x_index, gelib::d_vector const& y_larrow, gelib::d_vector& ahat2)
	{
		// set ahat2 equal to the drift
		m_fx_process->mean_drift(m_next_increment_correlations, m_x_larrow_k[x_index], m_z_larrow_k[z_index], m_y_larrow_k, ahat2);

		double z = m_z_larrow_k[z_index] - m_next_increment_correlations.alpha() * m_x_larrow_k[x_index];
		double lhs = m_next_increment_correlations.beta() * m_x_larrow_k[x_index] + m_next_increment_correlations.gamma() * z;

		// calculate mu3(gi) at each gridpoint from the drift term
		std::transform(ahat2.begin(), ahat2.end(), ahat2.begin(), [&lhs](double drift) { return lhs + drift; });

	//	if (ahat2.front() < m_ahat_2.front())
	//		throw std::runtime_error("Extrapolation of " + std::to_string(m_ahat_2.front() - ahat2.front()) + " to the left of ahat2");

	//	if (ahat2.back() > m_ahat_2.back())
	//		throw std::runtime_error("Extrapolation of " + std::to_string(m_ahat_2.back() - ahat2.back()) + " to the right of ahat2");
	};


	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// UNCOMMENT BELOW FOR LOG-FIT
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	//std::transform(m_Vk_on_ahat.begin(), m_Vk_on_ahat.end(), m_Vk_on_ahat.begin(), [](double x) -> double {return std::log(x); });
	//if (std::any_of(m_Vk_on_ahat.begin(), m_Vk_on_ahat.end(), std::isnan<double>))
	//		throw std::runtime_error("Failed due to negative values in log-fit");

	// prepare cube to store slice value
	m_Vk.resize(m_Vk_on_ahat.rows(), m_Vk_on_ahat.cols(), m_Vk_on_ahat.layers());

	// prepare interpolant over y dimension
	gelib::PiecewisePoly3d interpolant3d(
			1,
			m_ahat_2,
			m_Vk_on_ahat,
			m_Vk, std::array<size_t, 3>({ 1, 2, 0 }));

	// fit piecewise polynomials in y dimension on ahat2
	interpolant3d.fit_input(m_Vk_on_ahat, 1, m_model->get_order(), CHAR_EXTRAPOLATION, GridMaker1d::INTERP_TOLERANCE_DEFAULT, 0);

	// clean up m_Vk_on_ahat, we don't need it anymore because we have the fits.
	std::vector<double>().swap(m_Vk_on_ahat);

	// evaluate vector transform(m_y_larrow_k) on each piecewise poly 
	interpolant3d.evaluate(m_y_larrow_k, y_to_a2hat, 0, 1, m_Vk);


	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// UNCOMMENT BELOW FOR LOG-FIT
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	//std::transform(m_Vk.begin(), m_Vk.end(), m_Vk.begin(), [](double x) -> double {return std::exp(x); });

	return m_Vk;
}

void SliceValuatorTwoCurrencyAtTi::multiply_by_discount_factors()
{
	const auto mult_func = [](double& a, double& b)
	{
		double prod = a * b;
		std::swap(prod, b);
	};

	// multiply every row by 1 / (1 + alpha L(x))
	for (size_t i = 0; i < m_Vk.cols(); ++i)
		for (size_t j = 0; j < m_Vk.layers(); ++j)
			m_Vk.vector_op(mult_func, m_expectation_to_value_x_factors, -1, (int)i, (int)j);
}

gelib::d_cube & SliceValuatorTwoCurrencyAtTi::integrate_european_payoff()
{
	SliceValuatorTwoCurrency::integrate_european_payoff(
		SliceValuatorTwoCurrency::get_grid_x_larrow(next()),
		SliceValuatorTwoCurrency::get_grid_y_larrow(next()),
		SliceValuatorTwoCurrency::get_grid_z_larrow(next()),
		SliceValuatorTwoCurrency::get_valuations(next()),
		m_Vk_on_ahat, m_s_rightarrow_k_increment_covariance);

	// note that m_Vk contains the final result of the integration and transform 
	return m_Vk;
}

gelib::d_cube & SliceValuatorTwoCurrencyAtTi::integrate_european_payoff_first(gelib::d_cube& known_value)
{
	SliceValuatorTwoCurrency::integrate_european_payoff_first(
		get_grid_x(next()),
		get_grid_y(next()),
		get_grid_z(next()),
		known_value,
		m_Vk_on_ahat,
		m_Ck_Ckminus1);

	// note that m_Vk contains the final result of the integration and transform 
	return m_Vk;
}

SliceValuatorTwoCurrencyAtTi::SliceValuatorTwoCurrencyAtTi(
	const MFMBase* model,
	SliceProcessCorrelations3d const& first_fwd_increment_covariance, SliceProcessCorrelations3d const& second_fwd_increment_covariance,
	const InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* for_process,
	FXTimeSliceProcess* next_fx_process)
	:
	SliceValuatorTwoCurrency(model, first_fwd_increment_covariance, dom_process, fx_process, for_process,
		0.5 * first_fwd_increment_covariance.get_increment_covariance()[4]),
		//0.5 * first_fwd_increment_covariance.get_Mk_Ck_Ckminus1_Mkt()[4]),
	m_next_increment_correlations(first_fwd_increment_covariance),
	m_next_fx_process(next_fx_process),
	// TODO this is already in the processcorrelations object
	m_s_rightarrow_k_increment_covariance((second_fwd_increment_covariance.get_M()* m_Ck_Ckminus1)* second_fwd_increment_covariance.get_Mt())
{

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// M_{k + 1} * C_{k} * M_{k + 1}^T
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	gelib::d_matrix ahat_var_mat = first_fwd_increment_covariance.get_Mk_Ckminus1_Mkt();

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// construct ahat grids - to integrate m_V_{k + 1} TODO note this is 20% larger
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	GridMaker1d::make_grid_zero_center(dom_process->get_grid()->size(), dom_process->get_grid_spread_factor(), ahat_var_mat[0], m_ahat_1);
	GridMaker1d::make_grid(fx_process->get_grid()->size(), fx_process->get_grid_spread_factor() * 1.2, next_fx_process->get_uncorrelated_grid_variance(), fx_process->get_uncorrelated_grid_variance(), m_ahat_2);
	GridMaker1d::make_grid_zero_center(for_process->get_grid()->size(), for_process->get_grid_spread_factor(), ahat_var_mat[8], m_ahat_3);

	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	// construct grids for m_Vk - spread value at time slice k
	//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
	m_x_larrow_k = m_ahat_1;
	GridMaker1d::make_grid(m_ahat_2.size(), fx_process->get_grid_spread_factor(), fx_process->get_grid_center(), fx_process->get_uncorrelated_grid_variance(), m_y_larrow_k);
	m_z_larrow_k = m_ahat_3;

	// set the dimensions of the cubes for the indexing but don't actually allocate the memory. 
	//m_Vk_on_ahat.fix_dimensions(m_ahat_1.size(), m_ahat_2.size(), m_ahat_3.size());

	m_expectation_to_value_x_factors.resize(dom_process->get_grid()->size());

	for (size_t i = 0; i < m_expectation_to_value_x_factors.size(); ++i)
		m_expectation_to_value_x_factors[i] = 1.0 / (1.0 + m_dom_process->get_accrual_factor() * m_dom_process->get_spot(i));

}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//===================================================================================================================
// Valuator at first timeslice T0
//===================================================================================================================
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

SliceValuatorTwoCurrencyAtTZero::SliceValuatorTwoCurrencyAtTZero(
	const MFMBase* model, SliceProcessCorrelations3d const& first_fwd_increment_covariance, SliceProcessCorrelations3d const& second_fwd_increment_covariance, const InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* for_process, double dom_t1_discount)
	:
	SliceValuatorTwoCurrency(
		model, first_fwd_increment_covariance, dom_process, fx_process, for_process, 
		// y mean at T1
		std::log(for_process->get_t1_discount_factor()) - std::log(dom_process->get_t1_discount_factor()) - 0.5 * first_fwd_increment_covariance.get_increment_covariance()[4]),
	m_t1_discount(dom_t1_discount)
{
	m_ahat_1[0] = dom_process->get_grid_center();
	m_ahat_2[0] = m_y_mean_shift;
	m_ahat_3[0] = for_process->get_grid_center();
	m_s_rightarrow_k_increment_covariance = (second_fwd_increment_covariance.get_M() * m_Ck_Ckminus1) * second_fwd_increment_covariance.get_Mt();// TODO move to initialiser list
}

SliceValuatorTwoCurrency* SliceValuatorTwoCurrencyAtTZero::prev()
{
	// no prev valuator from T0
	throw;
	return nullptr;
}

gelib::d_cube const& SliceValuatorTwoCurrencyAtTZero::roll_back_european_payoff()
{

	integrate_european_payoff();

	multiply_by_discount_factors();

	return m_Vk;
}

gelib::d_cube const& SliceValuatorTwoCurrencyAtTZero::interpolate_payoff(gelib::d_cube& payoff)
{
	// Don't need to interpolate after rolling back from T1
	return payoff;
}

void SliceValuatorTwoCurrencyAtTZero::multiply_by_discount_factors()
{
	assert(m_Vk.size() == 1);

	// multiply by domestic discount factor at T1
	m_Vk[0] *= m_t1_discount;
}

double SliceValuatorTwoCurrencyAtTZero::roll_back_european_payoff(gelib::d_cube& known_value, TimeSliceCalibrator const& arrow_debreu_calibrator)
{
	// don't need to use the arrow debreu prices when rolling back T1 -> T0
	return roll_back_european_payoff(known_value)[0];
}

gelib::d_cube const& SliceValuatorTwoCurrencyAtTZero::roll_back_european_payoff(gelib::d_cube& known_value)
{

	integrate_european_payoff_first(known_value);
	
	multiply_by_discount_factors();

	return m_Vk;
}

gelib::d_cube & SliceValuatorTwoCurrencyAtTZero::integrate_european_payoff()
{
	SliceValuatorTwoCurrency::integrate_european_payoff(
		SliceValuatorTwoCurrency::get_grid_x_larrow(next()),
		SliceValuatorTwoCurrency::get_grid_y_larrow(next()),
		SliceValuatorTwoCurrency::get_grid_z_larrow(next()),
		SliceValuatorTwoCurrency::get_valuations(next()),
		m_Vk, m_s_rightarrow_k_increment_covariance);

	return m_Vk;
}

gelib::d_cube & SliceValuatorTwoCurrencyAtTZero::integrate_european_payoff_first(gelib::d_cube& known_value)
{
	SliceValuatorTwoCurrency::integrate_european_payoff_first(
		SliceValuatorTwoCurrency::get_grid_x(next()),
		SliceValuatorTwoCurrency::get_grid_y(next()),
		SliceValuatorTwoCurrency::get_grid_z(next()),
		known_value,
		m_Vk,
		m_Ck_Ckminus1);

	return m_Vk;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
//===================================================================================================================
//
//  Valuator at second to Last timeslice Tn-1
//
//===================================================================================================================
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

SliceValuatorTwoCurrency* SliceValuatorTwoCurrencyAtTLast::next()
{
	throw std::runtime_error("Tried to access next valuator from the last valuator");
	return nullptr;
}


// This function should never be used as we never roll back from the next valuator at the last valuator.
gelib::d_cube& SliceValuatorTwoCurrencyAtTLast::integrate_european_payoff()
{
	throw std::runtime_error("Tried to use non-max integrate to roll back from last time slice");
	return m_Vk;
}

gelib::d_cube& SliceValuatorTwoCurrencyAtTLast::integrate_european_payoff_first(gelib::d_cube& known_value)
{
	SliceValuatorTwoCurrency::integrate_european_payoff_first(
		*m_dom_process_last->get_grid(),
		*m_fx_process_last->get_grid(),
		*m_for_process_last->get_grid(),
		known_value,
		m_Vk_on_ahat,
		m_Ck_Ckminus1);

	return m_Vk;
}


SliceValuatorTwoCurrencyAtTLast::SliceValuatorTwoCurrencyAtTLast(
	const MFMBase* model,
	SliceProcessCorrelations3d const& first_fwd_increment_covariance,
	const InterestRateTimeSliceProcess* dom_process, FXTimeSliceProcess* fx_process, const InterestRateTimeSliceProcess* for_process,
	const InterestRateTimeSliceProcess* dom_process_last, FXTimeSliceProcess* fx_process_last, const InterestRateTimeSliceProcess* for_process_last)
	:
	SliceValuatorTwoCurrencyAtTi(
		model,
		first_fwd_increment_covariance, SliceProcessCorrelations3d(3),
		dom_process, fx_process, for_process, fx_process_last),
	m_dom_process_last(dom_process_last),
	m_fx_process_last(fx_process_last),
	m_for_process_last(for_process_last)
{}
