//#####################################################################
// Copyright 2012.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include <Core/Nonlinear_Equations/NEWTONS_METHOD.h>
#include <Core/Krylov_Solvers/CONJUGATE_GRADIENT.h>
#include <Core/Krylov_Solvers/GMRES.h>
#include <Core/Krylov_Solvers/MINRES.h>
#include <Core/Log/DEBUG_SUBSTEPS.h>
#include <Core/Log/LOG.h>
#include <Core/Matrices/SYMMETRIC_MATRIX.h>
#include <Core/Nonlinear_Equations/LINE_SEARCH.h>
#include <Core/Read_Write/OCTAVE_OUTPUT.h>
#include <functional>
using namespace PhysBAM;
//#####################################################################
// Function Newtons_Method
//#####################################################################
template <class T> std::pair<bool,int> NEWTONS_METHOD<T>::
Newtons_Method(const NONLINEAR_FUNCTION<T(KRYLOV_VECTOR_BASE<T>&)>& F,KRYLOV_SYSTEM_BASE<T>& sys,KRYLOV_VECTOR_BASE<T>& x,ARRAY<KRYLOV_VECTOR_BASE<T>*>& av) const
{
    KRYLOV_SOLVER<T>::Ensure_Size(av,x,3+use_gradient_magnitude_objective);
    KRYLOV_VECTOR_BASE<T>& grad=*av.Pop_Value();
    KRYLOV_VECTOR_BASE<T>& dx=*av.Pop_Value();
    KRYLOV_VECTOR_BASE<T>& tm=*av.Pop_Value();
    KRYLOV_VECTOR_BASE<T>& eff_grad=use_gradient_magnitude_objective?*av.Pop_Value():grad;
    GMRES<T> gmres;
    MINRES<T> minres;
    CONJUGATE_GRADIENT<T> cg;
    cg.finish_before_indefiniteness=finish_before_indefiniteness;
    KRYLOV_SOLVER<T>* krylov=&minres;
    if(use_cg) krylov=&cg;
    if(use_gmres) krylov=&gmres;
    krylov->relative_tolerance=true;
    krylov->print_matrix=print_matrix;
    bool result=false;
    if(debug) Dump_Parameters();

    T last_E=FLT_MAX;
    F.Make_Feasible(x);
    int iterations_used;
    bool can_fallback=true;
    for(iterations_used=0;iterations_used<max_iterations;iterations_used++){
        T E=0;
        F.Compute(x,&sys,&grad,&E);
        T norm_grad=sqrt(sys.Inner_Product(grad,grad)),norm_eff_grad=norm_grad;
        if(use_gradient_magnitude_objective){
            E=norm_grad*norm_grad/2;
            sys.Multiply(grad,eff_grad,true);
            norm_eff_grad=sqrt(sys.Inner_Product(eff_grad,eff_grad));}
        T conv_norm=sys.Convergence_Norm(eff_grad);
        if(use_gradient_magnitude_objective)
            conv_norm=min(conv_norm,sys.Convergence_Norm(grad));

        if(debug){
            LOG::printf("GRAD STATS %.16g %.16g %.16g %.16g\n",E,(E-last_E),norm_grad,tolerance);
            PHYSBAM_DEBUG_WRITE_SUBSTEP("newton %d   %.16g %.16g %.16g %.16g",debug_substep_level,iterations_used,E,(E-last_E),norm_grad,tolerance);}

        if(conv_norm<tolerance && (iterations_used || !require_one_iteration || !conv_norm)){result=true;break;}
        dx*=0;
        tm.Copy(-1,grad);
        T local_krylov_tolerance=krylov_tolerance;
        if(!fixed_tolerance)
            local_krylov_tolerance=min((T).5,max(krylov_tolerance,adaptive_tolerance*(T)sqrt(max(norm_grad,tolerance))));
        if(debug) LOG::cout<<"local tolerance: "<<local_krylov_tolerance<<std::endl;
        if(!krylov->Solve(sys,dx,tm,av,local_krylov_tolerance,0,max_krylov_iterations) && fail_on_krylov_not_converged)
            break;
        krylov->print_matrix=nullptr;
        
        if(use_gradient_descent_failsafe)
            Make_Downhill_Direction(sys,dx,eff_grad,norm_eff_grad,tm,can_fallback);
        
        T a=Line_Search(F,sys,x,dx,grad,tm,eff_grad);
        if(a<=0) break;

        x.Copy(a,dx,x);
        F.Make_Feasible(x);
        last_E=E;}
    
    if(print_matrix)
    {
        krylov->print_matrix=print_matrix;
        krylov->Print_Matrix(sys,tm);
    }

    if(debug) LOG::printf("newton iterations: %P\n",iterations_used);
    
    if(use_gradient_magnitude_objective) av.Append(&eff_grad);
    av.Append(&tm);
    av.Append(&dx);
    av.Append(&grad);
    return {result,iterations_used};
}
//#####################################################################
// Function Make_Downhill_Direction
//#####################################################################
template <class T> void NEWTONS_METHOD<T>::
Make_Downhill_Direction(const KRYLOV_SYSTEM_BASE<T>& sys,
    KRYLOV_VECTOR_BASE<T>& dx,const KRYLOV_VECTOR_BASE<T>& grad,T norm_grad,
    KRYLOV_VECTOR_BASE<T>& tmp,bool& can_fallback) const
{
    T norm_dx=sqrt(sys.Inner_Product(dx,dx)),inner=sys.Inner_Product(dx,grad);
    if(inner>angle_tolerance*norm_dx*norm_grad)
    {
        dx*=-1;
        if(debug) LOG::printf("LOOK BACKWARDS inner: %P a_tol: %P norm_dx: %P norm_grad: %P\n",inner,angle_tolerance,norm_dx,norm_grad);
    }
    else if(inner>=-angle_tolerance*norm_dx*norm_grad)
    {
        if(inner<0 && alternate_gradient_fallback)
        {
            if(can_fallback) can_fallback=false;
            else
            {
                if(debug) LOG::printf("SKIP GRADIENT inner: %P a_tol: %P norm_dx: %P norm_grad: %P\n",inner,angle_tolerance,norm_dx,norm_grad);
                can_fallback=true;
                return;
            }
        }
        if(debug) LOG::printf("GRADIENT inner: %P a_tol: %P norm_dx: %P norm_grad: %P\n",inner,angle_tolerance,norm_dx,norm_grad);
        dx.Copy(-norm_dx/norm_grad,grad);
    }
}
namespace{
template<class T>
class GRADIENT_SQUARED_OBJECTIVE:public NONLINEAR_FUNCTION<T(KRYLOV_VECTOR_BASE<T>&)>
{
public:
    const NONLINEAR_FUNCTION<T(KRYLOV_VECTOR_BASE<T>&)>& F;
    KRYLOV_SYSTEM_BASE<T>& system;
    KRYLOV_VECTOR_BASE<T>& tmp;

    GRADIENT_SQUARED_OBJECTIVE(
        const NONLINEAR_FUNCTION<T(KRYLOV_VECTOR_BASE<T>&)>& F,
        KRYLOV_SYSTEM_BASE<T>& system,KRYLOV_VECTOR_BASE<T>& tmp)
        :F(F),system(system),tmp(tmp)
    {}

    ~GRADIENT_SQUARED_OBJECTIVE() override {}

    void Compute(const KRYLOV_VECTOR_BASE<T>& x,KRYLOV_SYSTEM_BASE<T>* h,
        KRYLOV_VECTOR_BASE<T>* g,T* e) const override
    {
        PHYSBAM_ASSERT(!h);
        F.Compute(x,g?&system:0,&tmp,0);
        if(e) *e=system.Inner_Product(tmp,tmp)/2;
        if(g) system.Multiply(tmp,*g,true);
    }
};
}
//#####################################################################
// Function Line_Search
//#####################################################################
template <class T> T NEWTONS_METHOD<T>::
Line_Search(const NONLINEAR_FUNCTION<T(KRYLOV_VECTOR_BASE<T>&)>& F,KRYLOV_SYSTEM_BASE<T>& sys,const KRYLOV_VECTOR_BASE<T>& x,const KRYLOV_VECTOR_BASE<T>& dx,KRYLOV_VECTOR_BASE<T>& tmp,KRYLOV_VECTOR_BASE<T>& tmp2,KRYLOV_VECTOR_BASE<T>& tmp3) const
{
    T a=1;
    GRADIENT_SQUARED_OBJECTIVE<T> GSO(F,sys,tmp3);
    const NONLINEAR_FUNCTION<T(KRYLOV_VECTOR_BASE<T>&)>& G=use_gradient_magnitude_objective?GSO:F;

    if(use_wolfe_search || use_backtracking || use_golden_section_search){
        auto f=[&](T t,T* dg)
            {
                tmp.Copy(t,dx,x);
                T g;
                G.Compute(tmp,0,dg?&tmp2:0,&g);
                if(dg) *dg=sys.Inner_Product(tmp2,dx);
                return g;
            };
        if(use_wolfe_search){
            if(!Line_Search_Wolfe_Conditions<T>(f,0,1,a,(T)1e-4,(T).9,objective_error,std::numeric_limits<T>::max()))
                a=0;}
        else if(use_backtracking){
            if(!Line_Search_Backtracking<T>(f,0,1,a,(T)1e-4,objective_error))
                a=0;}
        else if(use_golden_section_search){
            T phi=(1+sqrt(5))/2;
            a*=phi;
            while(f(phi*a,0)<=f(a,0)) a*=phi;
            T tau=(T).5*(sqrt((T)5)-1);
            T A=f(0,0),B=f(tau,0),C=f(1-tau,0),D=f(a,0),mx=max(A,B,C,D),mn=min(A,B,C,D);
            if(!Line_Search_Golden_Section<T>(f,0,a,a,max_golden_section_iterations,(T).25*tolerance))
                return 0;
            if(mx-mn<=1e-10*maxabs(mx,mn)){
                if(debug) LOG::printf("OVERRIDE\n");
                a=1;}}}

    if(debug) LOG::printf("ALPHA  %g\n",a);
    return a;
}
//#####################################################################
// Function Make_Vanilla_Newton
//#####################################################################
template <class T> void NEWTONS_METHOD<T>::
Make_Vanilla_Newton()
{
    use_golden_section_search=false;
    use_wolfe_search=false;
    use_backtracking=false;
    use_gradient_descent_failsafe=false;
    fixed_tolerance=true;
    finish_before_indefiniteness=false;
}
//#####################################################################
// Function Dump_Parameters
//#####################################################################
template <class T> void NEWTONS_METHOD<T>::
Dump_Parameters() const
{
    auto db=[](const char* s,bool f){if(f) LOG::cout<<" "<<s;};
    auto di=[](const char* s,int i){LOG::printf(" %s=%i",s,i);};
    auto df=[](const char* s,T t){LOG::printf(" %s=%.3g",s,t);};
    LOG::cout<<"flags:";
    db("bt",use_backtracking);
    db("cg",use_cg);
    db("db",debug);
    db("fi",finish_before_indefiniteness);
    db("fk",fail_on_krylov_not_converged);
    db("ft",fixed_tolerance);
    db("gd",use_gradient_descent_failsafe);
    db("gm",use_gmres);
    db("go",use_gradient_magnitude_objective);
    db("gs",use_golden_section_search);
    db("ri",require_one_iteration);
    db("ws",use_wolfe_search);
    di("gi",max_golden_section_iterations);
    di("ki",max_krylov_iterations);
    di("mi",max_iterations);
    df("at",angle_tolerance);
    df("kt",krylov_tolerance);
    df("tl",tolerance);
    df("oe",objective_error);
    df("at",adaptive_tolerance);
    LOG::cout<<std::endl;
}
namespace PhysBAM{
template struct NEWTONS_METHOD<float>;
template struct NEWTONS_METHOD<double>;
}
