// Copyright (C) 2012, 2013, 2014  David Maxwell
//
// This file is part of PISM.
//
// PISM is free software; you can redistribute it and/or modify it under the
// terms of the GNU General Public License as published by the Free Software
// Foundation; either version 3 of the License, or (at your option) any later
// version.
//
// PISM is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
// details.
//
// You should have received a copy of the GNU General Public License
// along with PISM; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

#ifndef IP_SSATAUCFORWARDPROBLEM_HH_4AEVR4Z
#define IP_SSATAUCFORWARDPROBLEM_HH_4AEVR4Z


#include "SSAFEM.hh"
#include "IPDesignVariableParameterization.hh"

//! Implements the forward problem of the map taking \f$\tau_c\f$ to the corresponding solution of the %SSA.
/*! The class SSAFEM solves the %SSA, and the solution depends on a large number of parameters.  Considering
all of these to be fixed except for \f$\tau_c\f$, we obtain a map \f$F_{\rm SSA}\f$ taking
\f$\tau_c\f$ to the corresponding solution \f$u_{\rm SSA}\f$ of the %SSA.  This is a forward problem which
we would like to invert; given \f$u_{\rm SSA}\f$ determine \f$\tau_c\f$  such that 
\f$F_{\rm SSA}(\tau_c) = u_{\rm SSA}\f$.  The forward problem actually implemented by 
IP_SSATaucForwardProblem is a mild variation \f$F_{\rm SSA}\f$.

First, given the constraint \f$\tau_c\ge 0\f$ it can be helpful to parameterize \f$\tau_c\f$ by some other 
parameter \f$\zeta\f$,
\f[
\tau_c = g(\zeta),
\f]
where the function \f$g\f$ is non-negative.  The function \f$g\f$ is specified by an instance
of IPDesignVariableParameterization.

Second, there may be locations where the value of \f$\tau_c\f$ (and hence \f$\zeta\f$) 
is known a-priori, and should not be adjusted.  Let \f$\pi\f$ be the map that replaces
the values of \f$\zeta\f$ with known values at these locations.

IP_SSATaucForwardProblem implements the forward problem
\f[
F(\zeta) = F_{\rm SSA}(g(\pi(\zeta))).
\f]

Algorithms to solve the inverse problem make use of variations on the linearization
of \f$F\f$, which are explained below.

The solution of the %SSA in SSAFEM is implemented using SNES to solve
a nonlinear residual function of the form
\f[
\mathcal{R}_{SSA}(u;\tau_c,\text{other parameters})= \vec 0,
\f]
usually using Newton's method. Let
\f[
\mathcal{R}(u,\zeta) = \mathcal{R}_{SSA}(u;g(\pi(\zeta)),\text{other parameters}).
\f]

The derivative of \f$\mathcal{R}\f$ with respect to \f$ u\f$ is called the state Jacobian,
\f$J_{\rm State}\f$. Specifically, if \f$u=[U_j]\f$ then 
\f[
(J_{\rm State})_{ij} = \frac{d \mathcal{R}_i}{dU_j}.
\f]
This is exactly the same Jacobian that is computed by SSAFEM for solving the %SSA via SNES. The
matrix for the design Jacobian can be obtained with \ref assemble_jacobian_state.

The derivative of \f$\mathcal{R}\f$ with respect to \f$ \zeta\f$ is called the design Jacobian,
\f$J_{\rm Design}\f$. Specifically, if \f$\zeta=[Z_k]\f$ then 
\f[
(J_{\rm Design})_{ik} = \frac{d \mathcal{R}_i}{dZ_k}.
\f]
The map \f$J_{\rm Design}\f$ can be applied to a vector \f$d\zeta\f$ with
apply_jacobian_design.  For inverse methods using adjoints, one also
needs to be able to apply the transpose of \f$J_{\rm Design}\f$,
which is done using \ref apply_jacobian_design_transpose.

The forward problem map \f$F\f$ solves the implicit equation
\f[
\mathcal{R}(F(\zeta),\zeta) = 0.
\f]
Its linearisation \f$DF\f$ then satisfies
\f[
J_{\rm State}\; DF\; d\zeta + J_{\rm Design}\; d\zeta = 0
\f]
for any perturbation \f$d\zeta\f$.  Hence
\f[
DF = -J_{\rm State}^{-1}\circ J_{\rm Design}.
\f]
This derivative is sometimes called the reduced gradient in the literature.  To
apply \f$DF\f$ to a perturbation \f$d\zeta\f$, use \ref apply_linearization.  Adjoint
methods require the transpose of this map; to apply \f$DF^t\f$ to \f$du\f$ use
\ref apply_linearization_transpose.
*/
class IP_SSATaucForwardProblem : public SSAFEM
{
public:

  typedef IceModelVec2S DesignVec; ///< The function space for the design variable, i.e. \f$\tau_c\f$.
  typedef IceModelVec2V StateVec;  ///< The function space for the state variable, \f$u_{\rm SSA}\f$.

  //! Constructs from the same objects as SSAFEM, plus a specification of how \f$\tau_c\f$ is parameterized.
  IP_SSATaucForwardProblem(IceGrid &g, EnthalpyConverter &e,
                           IPDesignVariableParameterization &tp,
                           const PISMConfig &c);

  virtual ~IP_SSATaucForwardProblem();

  //! Selects nodes where \f$\tau_c\f$ (more specifically \f$\zeta\f$) should not be adjusted.
  /*! The paramter \a locations should be set to 1 at each node where \f$\tau_c\f$
      is fixed. The forward map then effectively treats the design space as the subspace
      of nodes where \a locations is 0. Tangent vectors to this subspace, as would be
      generated by, e.g., \f$J_{\rm Design}^t\f$ are represented as vectors in the full
      space with entries set to zero in the fixed locations.  These can safely be added
      to preexisting values of \f$\zeta\f$ without changing the entries of \f$\zeta\f$ at the 
      fixed locations.  Inversion can be done by setting an initial value of \f$\zeta\f$
      having the desired values in the fixed locations, and using set_tauc_fixed_locations()
      to indicate the nodes that should not be changed.
  */
  virtual PetscErrorCode set_tauc_fixed_locations(IceModelVec2Int &locations)
  { 
    m_fixed_tauc_locations = &locations;
    return 0;
  }

  //! Returns the last solution of the %SSA as computed by \ref linearize_at.
  virtual IceModelVec2V &solution() {
    return m_velocity;
  }

  //! Exposes the \f$\tau_c\f$ parameterization being used.
  virtual IPDesignVariableParameterization & tauc_param() {
    return m_tauc_param;
  }

  virtual PetscErrorCode set_design( IceModelVec2S &zeta);

  virtual PetscErrorCode linearize_at( IceModelVec2S &zeta, TerminationReason::Ptr &reason);

  virtual PetscErrorCode assemble_residual(IceModelVec2V &u, IceModelVec2V &R);
  virtual PetscErrorCode assemble_residual(IceModelVec2V &u, Vec R);

  virtual PetscErrorCode assemble_jacobian_state(IceModelVec2V &u, Mat J);

  virtual PetscErrorCode apply_jacobian_design(IceModelVec2V &u,IceModelVec2S &dzeta,IceModelVec2V &du);
  virtual PetscErrorCode apply_jacobian_design(IceModelVec2V &u,IceModelVec2S &dzeta, Vec du);
  virtual PetscErrorCode apply_jacobian_design(IceModelVec2V &u,IceModelVec2S &dzeta, PISMVector2 **du_a);

  virtual PetscErrorCode apply_jacobian_design_transpose(IceModelVec2V &u,IceModelVec2V &du,IceModelVec2S &dzeta);
  virtual PetscErrorCode apply_jacobian_design_transpose(IceModelVec2V &u,IceModelVec2V &du,Vec dzeta);
  virtual PetscErrorCode apply_jacobian_design_transpose(IceModelVec2V &u,IceModelVec2V &du,double **dzeta);

  virtual PetscErrorCode apply_linearization(IceModelVec2S &dzeta, IceModelVec2V &du);
  virtual PetscErrorCode apply_linearization_transpose(IceModelVec2V &du, IceModelVec2S &dzeta);

  //! Exposes the DMDA of the underlying grid for the benefit of TAO.
  virtual PetscErrorCode get_da(DM *da) {
    *da = SSADA;
    return 0;
  }
  
protected:

  PetscErrorCode construct();
  PetscErrorCode destruct();

  IceGrid &m_grid;                           ///< Cache of the comptuation grid.

  IceModelVec2S   *m_zeta;                   ///< Current value of zeta, provided from caller.
  IceModelVec2S   m_dzeta_local;             ///< Storage for d_zeta with ghosts, if needed when an argument d_zeta is ghost-less.

  IceModelVec2Int *m_fixed_tauc_locations;   ///< Locations where \f$\tau_c\f$ should not be adjusted.

  IPDesignVariableParameterization &m_tauc_param;     ///< The function taking \f$\zeta\f$ to \f$\tau_c\f$.

  IceModelVec2V  m_du_global;                ///< Temporary storage when state vectors need to be used without ghosts.
  IceModelVec2V  m_du_local;                 ///< Temporary storage when state vectors need to be used with ghosts.

  FEElementMap m_element_index;
  FEQuadrature m_quadrature;
  FEDOFMap     m_dofmap;

  KSP  m_ksp;                                ///< KSP used in \ref apply_linearization and \ref apply_linearization_transpose  
  Mat  m_J_state;                            ///< Mat used in \ref apply_linearization and \ref apply_linearization_transpose

  SNESConvergedReason m_reason;

  bool m_rebuild_J_state;                    ///< Flag indicating that the state jacobian matrix needs rebuilding.
};


#endif /* end of include guard: IP_SSATAUCFORWARDPROBLEM_HH_4AEVR4Z */