#include "FXTimeSliceProcess.h"


double FXTimeSliceProcess::get_uncorrelated_grid_variance() const
{
	return m_uncorrelated_variance;
}

void FXTimeSliceProcess::set_uncorrelated_grid_variance(double variance)
{
	m_uncorrelated_variance = variance;
}

void FXTimeSliceProcess::fit_spot(const std::string& extrapolation)
{
	if (m_spot_fitted)
		return;

	gelib::d_vector log_spot(m_spot_values.size());

	// get the log spot values to fit
	if (m_log_fit)
		std::transform(m_spot_values.begin(), m_spot_values.end(), log_spot.begin(), [](double x) -> double {return std::log(x); });

	char c = gelib::PiecewisePoly1d::parse_extrapolation_string(extrapolation);

	// important here that the fit function is given a constant pointer to the grid or will be modified
	m_spot_fit->fit(get_grid()->data(), &log_spot[0], m_grid.size(), m_model->get_order(), c);

	m_spot_fitted = true;

}

void FXTimeSliceProcess::interpolate_spot(gelib::d_vector const& x, gelib::d_vector& f_x) const
{
	// interpolate log(f_x)
	TimeSliceProcess::interpolate_spot(x, f_x);

	// exponentiate interpolated values
	if (m_log_fit)
		std::transform(f_x.begin(), f_x.end(), f_x.begin(), [](double x) -> double {return std::exp(x); });
}

double FXTimeSliceProcess::interpolate_spot(double x) const
{
	return m_log_fit ? std::exp(TimeSliceProcess::interpolate_spot(x)) : TimeSliceProcess::interpolate_spot(x);
}

double FXTimeSliceProcess::mean_drift_at_T0() const
{
	double gridpoint_at_T0 = 0.0;
		return 
		std::log(m_foreign_currency_process->get_t1_discount_factor()) -
		std::log(m_domestic_currency_process->get_t1_discount_factor())
		+ m_forward_drift_function->evaluate(&gridpoint_at_T0);
}

gelib::d_vector& FXTimeSliceProcess::mean_drift(SliceProcessCorrelations3d const& fwd_increment_correlations, double x_arrow, double z_arrow, gelib::d_vector const& y_arrow, gelib::d_vector& mu_y) const
{
	if (!m_has_drift)
		throw std::runtime_error("Tried to calculate mean drift without setting it up first");

	double x = x_arrow;
	double z = z_arrow - fwd_increment_correlations.alpha() * x_arrow;
	double yarrow_minus_y = fwd_increment_correlations.beta() * x + fwd_increment_correlations.gamma() * z;

	double log_terms = 
		std::log(1.0 + m_domestic_currency_process->get_accrual_factor() * m_domestic_currency_process->interpolate_spot(x))
		- std::log(1.0 + m_foreign_currency_process->get_accrual_factor() * m_foreign_currency_process->interpolate_spot(z));

	if (std::isnan<double>(log_terms))
		throw std::runtime_error("Something wrong with the libor terms, if using general model, try using more gripoints");

	// load input to the drift function in to mu_y
	std::transform(y_arrow.begin(), y_arrow.end(), mu_y.begin(),
		[&](double yarr) -> double { return log_terms + yarr - yarrow_minus_y; });

	// load result of f(mu_y) into mu_y
	m_forward_drift_function->evaluate(mu_y, mu_y);

	return mu_y;
}


void FXTimeSliceProcess::set_up_drift(SliceProcessCorrelations3d const& fwd_increment_correlations, std::string const& extrapolation, gelib::d_vector const& fwd_y_grid, gelib::d_vector& fwd_drift_fitting_grid)
{
	// already called this function if true, so no need to do anything
	if (m_has_drift)
		return;

	// define function f(y) = y - 0.5 * (C_k - C_k-1)_22 on fwd_drift_fitting_grid
	std::transform(fwd_y_grid.begin(), fwd_y_grid.end(), fwd_drift_fitting_grid.begin(),
		[&fwd_increment_correlations](double y) -> double {return y - 0.5 * fwd_increment_correlations.get_increment_covariance()[4]; });

	char c = gelib::PiecewisePoly1d::parse_extrapolation_string(extrapolation);

	// important here that the fit function is given a constant pointer to the grid or will be modified
	m_forward_drift_function->fit(fwd_y_grid.data(), &fwd_drift_fitting_grid[0], fwd_y_grid.size(), m_model->get_order(), c);

	//fit the piecewise poly
	m_has_drift = true;
}

double FXTimeSliceProcess::get_spot_fx() const
{
	return m_spot_fx;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// After calling this constructor you should have a fully initialised FXTimeSliceProcess object
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
FXTimeSliceProcess::FXTimeSliceProcess(std::unique_ptr<TimeSliceProcess> base_process, std::unique_ptr<OptionPricer> pricer, double spot_fx, const InterestRateTimeSliceProcess* domestic_currency_process, const InterestRateTimeSliceProcess* foreign_currency_process, const MFMBase* model, bool log_fit)
	:
	TimeSliceProcess(std::move(base_process), std::move(pricer)),
	m_forward_drift_function(std::unique_ptr<gelib::PiecewisePoly1d>(new gelib::PiecewisePoly1d())),
	m_drift_fitting_grid(m_grid.size()),
	m_spot_fx(spot_fx),
	m_domestic_currency_process(domestic_currency_process),
	m_foreign_currency_process(foreign_currency_process),
	m_has_drift(false),
	m_log_fit(log_fit)
{
	// initialise spot grid to spot_fx at T0
	std::fill(m_spot_values.begin(), m_spot_values.end(), spot_fx);
}

FXTimeSliceProcess::FXTimeSliceProcess(std::unique_ptr<TimeSliceProcess> base_process, double spot_fx, const InterestRateTimeSliceProcess* domestic_currency_process, const InterestRateTimeSliceProcess* foreign_currency_process, const MFMBase* model, const std::string& ub_error, double beta, double mux_price, double muy_price, double Cxx_price, double Cyy_price, double Cxy_price)
	:
	FXTimeSliceProcess(
		std::move(base_process),// TODO check upper bound
		std::move(std::unique_ptr<HullWhiteOptionPricer>(new HullWhiteOptionPricer(ub_error, 0.0, MFMBase::plus_inf, beta,
			[&](double k) -> double { return k; },
			mux_price, muy_price, Cxx_price, Cyy_price, Cxy_price))),
		spot_fx, domestic_currency_process, foreign_currency_process, model, true)
{}

std::unique_ptr<FXTimeSliceProcess> FXTimeSliceProcess::construct(std::unique_ptr<TimeSliceProcess> base_process, double spot_fx, const InterestRateTimeSliceProcess* domestic_currency_process, const InterestRateTimeSliceProcess* foreign_currency_process, const MFMBase* model, const std::string& ub_error, double beta, double mux_price, double muy_price, double Cxx_price, double Cyy_price, double Cxy_price)
{
	return std::unique_ptr<FXTimeSliceProcess>(new FXTimeSliceProcess(std::move(base_process), spot_fx, domestic_currency_process, foreign_currency_process, model, ub_error, beta, mux_price, muy_price, Cxx_price, Cyy_price, Cxy_price));
}

bool FXTimeSliceProcess::log_fit() const
{
	return m_log_fit;
}

gelib::d_vector& FXTimeSliceProcess::get_drift_fitting_grid()
{
	return m_drift_fitting_grid;
}

std::vector<std::unique_ptr<FXTimeSliceProcess>> FXTimeSliceProcess::construct(
	double grid_spread, gelib::d_vector const& grid_centers, size_t number_of_gridpoints, const MFMBase* model,
	gelib::d_vector const& times, gelib::d_vector const& process_volatilities_squared, gelib::d_vector const& process_variances,
	double mean_reversion_param_domestic, double mean_reversion_param_foreign, double spot_fx,
	std::vector<std::unique_ptr<InterestRateTimeSliceProcess>> const& domestic_currency_processes, std::vector<std::unique_ptr<InterestRateTimeSliceProcess>> const& foreign_currency_processes, 
	gelib::d_vector const& mux_price, gelib::d_vector const& muy_price, gelib::d_vector const& Cxx_price, gelib::d_vector const& Cyy_price, gelib::d_vector const& Cxy_price)
{
	// TODO update with general correct formula
	double mean_reversion_param = mean_reversion_param_domestic;

	std::vector<std::unique_ptr<TimeSliceProcess>> processes = TimeSliceProcess::construct(grid_spread, number_of_gridpoints, model, mean_reversion_param, times, process_volatilities_squared, process_variances, grid_centers);

	std::vector<std::unique_ptr<FXTimeSliceProcess>> new_processes;

	new_processes.resize(processes.size());

	std::string ub_error = "Hull White digital FX in arrears price exceeded upper bound.";

	double beta = -1.0;

	new_processes[0] = FXTimeSliceProcess::construct(std::move(processes[0]), spot_fx, domestic_currency_processes[0].get(), foreign_currency_processes[0].get(), model, ub_error, beta, 0.0, 0.0, 0.0, 0.0, 0.0);

	for (size_t i = 1; i < processes.size(); ++i)
	{
		new_processes[i] = FXTimeSliceProcess::construct(std::move(processes[i]), spot_fx, domestic_currency_processes[i].get(), foreign_currency_processes[i].get(), model, ub_error, beta, mux_price[i - 1], muy_price[i - 1], Cxx_price[i - 1], Cyy_price[i - 1], Cxy_price[i - 1]);
	}

	return std::move(new_processes);
}
