#ifndef _NIS_SOLVER_BASE_H_
#define _NIS_SOLVER_BASE_H_

//#define OPEN_MP_SUPPORT

#define PARDISO_SUPPORT

#ifdef PARDISO_SUPPORT
#include <Eigen/PardisoSupport>
#endif

#include "..\headers\global_headers.h"
#include "..\headers\math_headers.h"
#include "constraint.h"
#include "..\timer_wrapper.h"

#include <queue>
#include <Eigen\SparseLU>

typedef enum
{
	OPTIMIZATION_LBS,
	OPTIMIZATION_LOCAL_GLOBAL,
	OPTIMIZATION_PENALTY,
	OPTIMIZATION_PENALTY_LOCAL_GLOBAL_BLEND,
	OPTIMIZATION_LOCALLY_INJECTIVE_MAPPING,
	OPTIMIZATION_MOSEK_WITH_ITERATIONS,
	OPTIMIZATION_MOSEK_WITHOUT_ITERATIONS,
	OPTIMIZATION_BARRIER_METHOD,
	OPTIMIZATION_NONLINEAR_RESCALING,
	OPTIMIZATION_NEWTON_PRIMAL_DUAL,
	OPTIMIZATION_PENALTY_ON_DETF,
	OPTIMIZATION_SPECTRAL_SDP,
	OPTIMIZATION_PENNON,
	OPTIMIZATION_LOCALIZED_SDP_SCHUR_COMPLEMENT,
	OPTIMIZATION_AUXILIARY_SOLVE,
	OPTIMIZATION_ALTERNATIVE_PROJECTION,
	OPTIMIZATION_ONE_TANGENT_PLANE,
	OPTIMIZATION_EXTERIOR_POINTS,
	OPTIMIZATION_TAGENT_SPACE,
	OPTIMIZATION_ITERATIVE_QP,
	OPTIMIZATION_DISCONNECTED_MESH,
	OPTIMIZATION_COROT_NEWTON,
	OPTIMIZATION_COROT_LOCAL_GLOBAL_STYLE,
	OPTIMIZATION_COROT_CONSTANT_HESSIAN_1,
	OPTIMIZATION_COROT_CONSTANT_HESSIAN_2,
	OPTIMIZATION_NEOHOOKEAN_GRADIENT_DESCENT,
	OPTIMIZATION_NEWHOOKEAN_CONSTANT_HESSIAN,
	OPTIMIZATION_NONINVERSION_CONSTANT_HESSIAN,
	OPTIMIZATION_TOTAL_NUM
}OptimizationMethod;

typedef enum
{
	QP_STATE_INIT,
	QP_STATE_GENERATING_S_PAIR,
	QP_STATE_PREPARED_TO_SOLVE,
	QP_STATE_FEASIBLE,
	QP_STATE_INFEASIBLE,
	QP_STATE_CONVERGED,
	QP_STATE_TOTAL_STATE
}QPState;

typedef enum
{
	BARRIER_TYPE_LOG,
	BARRIER_TYPE_LOG_PLUS_LINEAR,
	BARRIER_TYPE_SPLINE,
	BARRIER_TYPE_QUADRATIC,
	BARRIER_TYPE_MODIFIED_LOG,
	BARRIER_TYPE_NR_TRANSFORMATION,
	BARRIER_TYPE_LSL,
	BARRIER_TYPE_TOTAL_NUM
} BarrierType;

typedef enum
{
	PENALTY_TYPE_QUADRATIC,
	PENALTY_TYPE_CUBIC,
	PENALTY_TYPE_SPLINE,
	PENALTY_TYPE_TOTAL_NUM
} PenaltyType;

class NISSolverBase
{
public:
	NISSolverBase();
	virtual ~NISSolverBase();

	// TO Ming:
	void InitParametersInPhysbam();

	void Reset();
	void ResetParameterToOptimizationMethod();
	void ResetFactorizationFlag();

	// write your own Init function in inherited class;
	// Init(arg_list): initialize "m_constraints", allocate space for "m_handle_vertices", "m_per_vertex_rotation_2d" and "m_per_vertex_rotation_3d"
	// write your own ComputeHandleInfo function in inherited class;
	// ComputeHandleInfo(arg_list): calculate exrotation and attachement infomation
	
	// preparation
	void UpdateHandleInfo();
	void ResetInitialRescalingParameter(ScalarType initial_rescaling_parameter = -1);
	void SetInitialLambda(const VectorX& x, VectorX& lambda);
	void SetNoninversionCriteria();
	void SetRotationFieldType();
	void UpdateRotation(const VectorX& x);
	// update
	bool UpdateOneIteration(VectorX& x, VectorX& lambda, ScalarType& progress, OptimizationMethod method = OPTIMIZATION_TOTAL_NUM); // return true when converged
	bool UpdateMatlabMosekYalmip(VectorX& x, bool has_max_it = false);
	void UpdateMatlabMosekEpigraph(VectorX& x);

	// alternating projection
	bool CorrectFeasibilityProjectionLocal(VectorX& f);
	void CorrectFeasibilityProjectionGlobal(VectorX& f, VectorX& x);
	inline void SetPrefactorizationFlag() { m_need_prefactorization_flag = true; }

protected:
	// solve related subroutines from high to low;
	// highest level, optimization method
	bool updateGradientDescent(VectorX& x);
	bool updateLocalGlobal(VectorX& x);
	bool updatePenaltyMethod(VectorX& x);
	bool updateSpectralSDP(VectorX& x);
	bool updatePennon(VectorX& x);
	bool updatePenaltyMethodLocalGlobalBlend(VectorX& x);
	bool updatePenaltyMethodOnDetF(VectorX& x);
	bool updateLIM(VectorX& x);
	bool updateBarrierMethod(VectorX& x, VectorX& lambda);
	bool updatePrimalDual(VectorX& x, VectorX& lambda);
	// update corot model
	bool updateCorotNewton(VectorX& x);
	bool updateCorotLocalGlobal(VectorX& x);
	bool updateCorotConstantHessian1(VectorX& x);
	bool updateCorotConstantHessian2(VectorX& x);
	bool updateAlternatingProjection(VectorX& x);
	bool updateOneTangentPlane(VectorX& x);

	// prefactorization
	void prefactorizeGTG();

	// update neohookean model
	bool updateNeoHookeanGradientDescent(VectorX& x);
	bool updateNeoHookeanConstantHessian(VectorX& x);

	bool updateNoninversionConstantHessian(VectorX& x);


	// middle level, evaluate energy/gradient/hessian, linesearch
	ScalarType evaluateEnergy(const VectorX& x);
	void evaluateGradient(const VectorX& x, VectorX& gradient);
	void evaluateHessian(const VectorX& x, EigenSparseMatrix& hessian);
	ScalarType lineSearch(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir, bool throwable = true);

	// local global
	void localGlobalPrefactorization(const VectorX& x);
	unsigned int localGlobalChooseMultiplier(ScalarType multiplier);
	void localGlobalLocalStep(const VectorX& x);
	ScalarType evaluateEnergyLocalGlobal(const VectorX& x);
	void evaluateGradientLocalGlobal(const VectorX& x, VectorX& gradient);
	void evaluateHessianLocalGlobal(const VectorX& x, EigenSparseMatrix& hessian);

	// penalty method
	ScalarType evaluateEnergyPenaltyMethod(const VectorX& x, ScalarType (*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void evaluateGradientPenaltyMethod(const VectorX& x, VectorX& gradient, ScalarType (*d_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void evaluateHessianPenaltyMethod(const VectorX& x, EigenSparseMatrix& hessian, ScalarType (*d_psi)(ScalarType, ScalarType), ScalarType (*d2_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	ScalarType lineSearchPenaltyMethod(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir, ScalarType (*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	bool allElementsHasZeroPenaltyEnergy(const VectorX& x, ScalarType (*psi)(ScalarType, ScalarType));
	// spectral method
	ScalarType evaluateEnergySpectralSDP(const VectorX& x, ScalarType (*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void evaluateGradientSpectralSDP(const VectorX& x, VectorX& gradient, ScalarType (*d_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void evaluateHessianSpectralSDP(const VectorX& x, EigenSparseMatrix& hessian, ScalarType (*d_psi)(ScalarType, ScalarType), ScalarType (*d2_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	ScalarType lineSearchSpectralSDP(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir, ScalarType (*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void initSpectralSDPLagrangian();
	void updateSpectralSDPLagrangian(const VectorX& x_old, const VectorX& x_new, ScalarType (*d_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	//pennon
	ScalarType evaluateEnergyPennon(const VectorX& x, ScalarType(*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void evaluateGradientPennon(const VectorX& x, VectorX& gradient, ScalarType(*d_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void evaluateHessianPennon(const VectorX& x, EigenSparseMatrix& hessian, ScalarType(*d_psi)(ScalarType, ScalarType), ScalarType(*d2_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	ScalarType lineSearchPennon(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir, ScalarType(*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	void initPennonLagrangian();
	void updatePennonLagrangian(const VectorX& x_old, const VectorX& x_new, ScalarType(*d_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);


	// penalty method parallelized
	void evaluateEnergyGradientHessianPenaltyMethod(const VectorX& x, ScalarType (*psi)(ScalarType, ScalarType), ScalarType (*d_psi)(ScalarType, ScalarType), ScalarType (*d2_psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);
	ScalarType retrieveEnergyPenaltyMethod();
	void retrieveGradientPenaltyMethod(VectorX& gradient, unsigned size);
	void retrieveHessianPenaltyMethod(EigenSparseMatrix& hessian, unsigned size);
	ScalarType lineSearchPenaltyMethod(const ScalarType currentObjectiveValue, const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir, ScalarType (*psi)(ScalarType, ScalarType), ScalarType multiplier = 1.0);

	ScalarType evaluateEnergy(const VectorX& x, ScalarType (*psi)(ScalarType), bool throwable = true);
	void evaluateGradient(const VectorX& x, VectorX& gradient, ScalarType (*d_psi)(ScalarType));
	void evaluateHessian(const VectorX& x, EigenSparseMatrix& hessian, ScalarType (*d_psi)(ScalarType), ScalarType (*d2_psi)(ScalarType));

	void evaluateGradient(const VectorX& x, const VectorX& lambda, VectorX& gradient, ScalarType (*d_psi)(ScalarType));
	void evaluateHessian(const VectorX& x, const VectorX& lambda, EigenSparseMatrix& hessian, ScalarType (*d_psi)(ScalarType), ScalarType (*d2_psi)(ScalarType));
	//ScalarType lineSearch(const VectorX& x, const VectorX& lambda, const VectorX& current_residual, const VectorX& primal_dual_descent_dir, const unsigned int x_size, const unsigned int lambda_size);

	ScalarType evaluateEnergyNoConstraint(const VectorX& x);
	void evaluateGradientNoConstraint(const VectorX& x, VectorX& gradient);
	void evaluateHessianNoConstraint(const VectorX& x, EigenSparseMatrix& hessian);
	ScalarType lineSearchNoConstraint(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir);

	void evaluatePrimalDualResidual(const VectorX& x, const VectorX& lambda, VectorX& res);
	void evaluatePrimalDualKKT(const VectorX& x, const VectorX& lambda, EigenSparseMatrix& KKT);
	ScalarType lineSearchPrimalDual(const VectorX& primal_dual_x_lambda, const VectorX& current_residual, const VectorX& primal_dual_descent_dir, const unsigned int x_size, const unsigned int lambda_size);

	ScalarType evaluateEnergyCorot(const VectorX& x);
	void evaluateGradientCorot(const VectorX& x, VectorX& gradient);
	void evaluateHessianCorot(const VectorX& x, EigenSparseMatrix& hessian);
	ScalarType lineSearchCorot(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir);
	void evaluateHessianLocalGlobalStyle(const VectorX& x, EigenSparseMatrix& hessian);
	void evaluateHessianConst1(const VectorX& x, EigenSparseMatrix& hessian);
	void evaluateHessianConst2(const VectorX& x, EigenSparseMatrix& hessian);

	ScalarType evaluateEnergyNeoHookean(const VectorX& x);
	void evaluateGradientNeoHookean(const VectorX& x, VectorX& gradient);
	ScalarType lineSearchNeoHookean(const VectorX& x, const VectorX& gradient_dir, const VectorX& descent_dir);

	void setBarrierFunction();
	void setRescalingParameter(const VectorX& x, const VectorX& lambda);

	ScalarType limChooseAlpha(const VectorX& x);

	// feasibility correction
	void evaluateGMatrix(EigenSparseMatrix& G);
	void evaluateTMatrix(EigenSparseMatrix& T);
	void evaluateAMatrix(EigenSparseMatrix& A);

	// matlab solve related
	void evaluateGMatrix(std::vector<SparseMatrixTriplet>& g_triplets, unsigned int& dim1, unsigned int& dim2);

	// low level, specific solvers
	void factorizeDirectSolverLDLT(const EigenSparseMatrix& A, Eigen::SimplicialLDLT<EigenSparseMatrix, Eigen::Upper>& ldltSolver, char* warning_msg = ""); // factorize matrix A using LDLT decomposition
	void factorizeSparseLU(const EigenSparseMatrix& A, Eigen::SparseLU<EigenSparseMatrix>& luSolver, char* warning_msg = "");
	void factorizeDirectSolverLLT(const EigenSparseMatrix& A, Eigen::SimplicialLLT<EigenSparseMatrix, Eigen::Upper>& lltSolver, char* warning_msg = ""); // factorize matrix A using LLT decomposition
	void factorizeDirectSolverLLTWithoutSymbolicFactorization(const EigenSparseMatrix& A, Eigen::SimplicialLLT<EigenSparseMatrix, Eigen::Upper>& lltSolver, char* warning_msg = "");
#ifdef PARDISO_SUPPORT
	void factorizePardisoLDLT(const EigenSparseMatrix& A, Eigen::PardisoLDLT<EigenSparseMatrix, Eigen::Upper>& ldltSolver, char* warning_msg = ""); // factorize matrix A using LDLT decomposition
	void factorizePardisoLDLTWithoutSymbolicFactorization(const EigenSparseMatrix& A, Eigen::PardisoLDLT<EigenSparseMatrix, Eigen::Upper>& ldltSolver, char* warning_msg = ""); // factorize matrix A using LDLT decomposition
	void factorizePardisoLLT(const EigenSparseMatrix& A, Eigen::PardisoLLT<EigenSparseMatrix, Eigen::Upper>& lltSolver, char* warning_msg = ""); // factorize matrix A using LLT decomposition
	void factorizePardisoLLTWithoutSymbolicFactorization(const EigenSparseMatrix& A, Eigen::PardisoLLT<EigenSparseMatrix, Eigen::Upper>& lltSolver, char* warning_msg = "");
#endif

public:
	// stiffness
	ScalarType m_stiffness_attachment; // spring stiffness
	ScalarType m_mu; // lame's second coefficient
	ScalarType m_lambda; // lame's first coefficient
	ScalarType m_alpha; // stiffness for w00 
	ScalarType m_beta; // stiffness for det(w(0:1, 0:1)) in 3d, or detW in 2d
	ScalarType m_gamma; // stiffness for detW in 3d, not used in 2d
	bool m_auto_multiplier;
	ScalarType m_multiplier_mu;
	ScalarType m_multiplier;

	// state control
	OptimizationMethod m_optimization_method;
	bool m_enable_non_inversion; // enable/disable non-inversion
	bool m_enable_modified_barrier; //TOBE removed
	EnergyType m_energy_type; // switch ex-rotation or co-rotation
	BarrierType m_barrier_type; // barrier function type
	NoninversionCriteria m_noninversion_criteria; // swtich between W being psd or det(F) > 0
	PenaltyType m_penalty_type; // penalty function type
	ScalarType (*m_psi)(ScalarType);
	ScalarType (*m_d_psi)(ScalarType);
	ScalarType (*m_d2_psi)(ScalarType);

	// hard coded collision
	bool m_collision_enable;
	ScalarType m_collision_y_plane;
	ScalarType m_collision_stiffness;
	std::vector<unsigned int> m_surface_indices;

	// Local Global
	unsigned int m_inner_loop_iterations;
	bool m_local_global_factorization_flag;
	Eigen::SimplicialLLT<EigenSparseMatrix, Eigen::Upper> m_local_global_llt_solvers[16];
	ScalarType m_local_global_multipliers[16];

	// penalty method
#ifdef PARDISO_SUPPORT
	Eigen::PardisoLLT<EigenSparseMatrix, Eigen::Upper> m_penalty_hessian_llt_solver;
#else
	Eigen::SimplicialLLT<EigenSparseMatrix, Eigen::Upper> m_penalty_hessian_llt_solver;
#endif
	bool m_penalty_hessian_symbolic_factorization_flag;

	// alternating projection
	VectorX m_f;
	VectorX m_last_x;
	// prefactorization
	bool m_need_prefactorization_flag;
	EigenSparseMatrix m_G;
	Eigen::SimplicialLDLT<EigenSparseMatrix, Eigen::Upper> m_GTG_solver;
	EigenSparseMatrix m_T;
	EigenSparseMatrix m_A;
	EigenSparseMatrix m_M;
	Eigen::SimplicialLDLT<EigenSparseMatrix, Eigen::Upper, Eigen::NaturalOrdering<IndexType>> m_M_solver;

	// Spectral SDP
	bool m_spectral_sdp_is_first_iteration;
	VectorX m_spectral_sdp_last_x;

	// LIM
	bool m_enable_substepping;
	ScalarType m_regularization_mu;
	ScalarType m_lim_alpha_update_last_alpha;

	// linesearch related
	bool m_enable_line_search;
	ScalarType m_ls_alpha;
	ScalarType m_ls_beta;
	ScalarType m_ls_stepsize;

	// number of tri/tet constraints
	unsigned int m_constraint_num;
	unsigned int m_attachment_constraint_num;
	// number of tri/tet constraints
	unsigned int m_f_size;
	// number of vertices
	unsigned int m_vertices_num;
	// number of tri/tet constraints
	unsigned int m_element_num;

	// number of iterations
	unsigned int m_iteration_num;

	// timer
	TimerWrapper m_main_timer;

	// depricated varibales, keep them for antweakbar
	bool m_ex_rot_barrier; // depricated
	ScalarType m_slackness;//depricated TOBE removed
	bool m_low_rank_update; // depricated
	int m_pyramid_side_num; // depricated

public:
	
	// key component: constraints
	std::vector<Constraint*> m_constraints;

	// static component
	bool m_mode; 
	// dynamics components: update from outside
	std::vector<VectorX> m_handle_vertices; // vertices in each handle
	std::vector<ScalarType> m_per_element_rotation_2d; // 2d rotations
	std::vector<Eigen::Quaternion<ScalarType, Eigen::DontAlign>> m_per_element_rotation_3d; // 3d rotations
};


#endif