# Copyright (C) 2012 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

import PISM
from PISM.logging import logError, logMessage
from PISM.util import Bunch

if not PISM.imported_from_sphinx:
  from petsc4py import PETSc

import sys
import math

import siple
import PISM.invert.sipletools
from PISM.invert.ssa import InvSSASolver
from PISM.invert.sipletools import PISMLocalVector
from siple.gradient.forward import NonlinearForwardProblem
from siple.gradient.nonlinear import InvertNLCG, InvertIGN

siple.reporting.clear_loggers()
siple.reporting.add_logger(PISM.invert.sipletools.pism_logger)
siple.reporting.set_pause_callback(PISM.invert.sipletools.pism_pause)

WIDE_STENCIL = 2

class InvSSASolver_Gradient(InvSSASolver):
  """Inverse SSA solver based on `siple` iterative gradient methods."""

  def __init__(self,ssarun,method):
    """
    :param ssarun: The :class:`PISM.ssa.SSARun` defining the forward problem. See :class:`PISM.inv_ssa.SSAForwardRunFromInputFile`.
    :param method: String describing the actual ``siple`` algorithm to use. One of ``sd``, ``nlcg`` or ``ign``."""
    InvSSASolver.__init__(self,ssarun,method)

    for o in PISM.OptionsGroup(ssarun.grid.com,"","InvSSASolver_Gradient"):
      self.target_misfit = PISM.optionsReal("-inv_target_misfit","m/year; desired root misfit for inversions",default=None)
      
      monitor_adjoint = PISM.optionsFlag("-inv_monitor_adjoint","Track accuracy of the adjoint during computation",default=False)
      morozov_scale_factor  = PISM.optionsFlag("-inv_morozov_scale","Scale factor (>=1) for Morozov discrepancy principle",default=1.0)
      ls_verbose = PISM.optionsFlag("-inv_ls_verbose","Turn on a verbose linesearch.",default=False)
      ign_theta  = PISM.optionsReal("-ign_theta","theta parameter for IGN algorithm",default=0.5)
      max_it     = PISM.optionsInt("-inv_max_it","maximum iteration count",default=1000)

    if self.target_misfit is None:
      raise RuntimeError("Missing required option -inv_target_misfit")

    velocity_scale = ssarun.grid.config.get("inv_ssa_velocity_scale", "m/year", "m/year")
    self.target_misfit /=  velocity_scale
  
    self.forward_problem = SSAForwardProblem(ssarun)

    # Determine the inversion algorithm, and set up its arguments.
    if self.method == "ign":
      Solver = InvertSSAIGN
    else:
      Solver = InvertSSANLCG

    params=Solver.defaultParameters()
    params.ITER_MAX = max_it
    if self.method == "sd":
      params.steepest_descent = True
      params.ITER_MAX=10000
    elif self.method =="ign":
      params.linearsolver.ITER_MAX=10000
      params.linearsolver.verbose = True
    if ls_verbose:
      params.linesearch.verbose = True
    params.verbose   = True
    params.thetaMax = ign_theta
    params.mu = morozov_scale_factor

    params.deriv_eps = 0.

    self.siple_solver=Solver(self.forward_problem,params=params)

    if monitor_adjoint:
      self.addIterationListener(MonitorAdjoint())
      if self.method=='ign':
        self.addLinearIterationListener(MonitorAdjointLin())
    

  def solveForward(self,zeta,out=None):
    """Given a parameterized design variable value :math:`\zeta`, solve the SSA.
    See :cpp:class:`IP_TaucParam` for a discussion of parameterizations.
    
    :param zeta: :cpp:class:`IceModelVec` containing :math:`\zeta`.
    :param out: optional :cpp:class:`IceModelVec` for storage of the computation result.
    :returns: An :cpp:class:`IceModelVec` contianing the computation result.
    """
    if out is None:
      out = self.forward_problem.F(PISMLocalVector(zeta))
    else:
      out = self.forward_problem.F(PISMLocalVector(zeta),out=PISMLocalVector(out))
    return out.core()

  def addIterationListener(self,listener):
    """Add a listener to be called after each iteration.  See FIXME."""
    self.siple_solver.addIterationListener(SipleIterationListenerAdaptor(self,listener))

  def addDesignUpdateListener(self,listener):
    self.siple_solver.addXUpdateListener(SipleDesignUpdateListenerAdaptor(self,listener))

  def addLinearIterationListener(self,listener):
    self.siple_solver.addLinearIterationListener(SipleLinearIterationListenerAdaptor(self,listener))

  def solveInverse(self,zeta0,u_obs,zeta_inv):
    """Executes the inversion algorithm.
    
    :param zeta0: The best `a-priori` guess for the value of the parameterized design variable value :math:`\zeta`.
                  Ignored by siple gradient algorithms in deference to `zeta_inv`.
    :param u_obs: :cpp:class:`IceModelVec2V` of observed surface velocities.
    :param zeta_inv: :cpp:class:`zeta_inv` starting value of :math:`\zeta` for iterative minimization.
    :returns: A :cpp:class:`TerminationReason`.
    """
    try:
      vecs = self.ssarun.modeldata.vecs
      if vecs.has('zeta_fixed_mask'):
        self.ssarun.ssa.set_tauc_fixed_locations(vecs.zeta_fixed_mask)

      (self.zeta_i,self.u_i) = self.siple_solver.solve(zeta_inv,u_obs,self.target_misfit)
    except Exception:
      import traceback
      exc_type, exc_value, exc_traceback = sys.exc_info()
      description = ""
      for l in traceback.format_exception(exc_type, exc_value,exc_traceback):
        description += l
      # It would be nice to make siple so that if the inverse solve fails
      # you can still keep the most recent iteration. 
      self.u_i = None
      self.zeta_i = None
      return PISM.GenericTerminationReason(-1,description)
    return PISM.GenericTerminationReason(1,"Morozov Discrepancy Met")

  def inverseSolution(self):
    """Returns a tuple ``(zeta,u)`` of :cpp:class:`IceModelVec`\s corresponding to the values
    of the design and state variables at the end of inversion."""
    return (self.zeta_i,self.u_i)


class SSAForwardProblem(NonlinearForwardProblem):
  """Subclass of a :class:`siple.NonlinearForwardProblem` defining the forward problem for ``siple``.
  The heavy lifting is forwarded to a :cpp:class:`IP_SSATaucForwardProblem` or 
  :cpp:class:`IP_SSAHardnessForwardProblem` living inside a :class:`PISM.invert.ssa.SSATaucForwardRun`."""

  def __init__(self,ssarun):
    """
    :param ssarun:  

            A :class:`PISM.invert.ssa.SSAForwardRun`
            or :class:`PISM.invert.ssa.SSAForwardRunFromInputFile` defining
            the forward problem."""

    self.ssarun = ssarun
    self.ssa = self.ssarun.ssa
    self.grid = ssarun.grid

    self.tmpV = PISM.IceModelVec2V();
    self.tmpV.create(self.grid, "work vector (2V)", PISM.WITHOUT_GHOSTS, WIDE_STENCIL);

    self.tmpS = PISM.IceModelVec2S();
    self.tmpS.create(self.grid, "work vector (2S)", PISM.WITHOUT_GHOSTS, WIDE_STENCIL);

    self.tmpS2 = PISM.IceModelVec2S();
    self.tmpS2.create(self.grid, "work vector (2S)", PISM.WITHOUT_GHOSTS, WIDE_STENCIL);

    ksp_rtol = 1e-12;
    self.ksp = PETSc.KSP();
    self.ksp.create(self.grid.com)
    self.ksp.setTolerances(ksp_rtol,PETSc.DEFAULT,PETSc.DEFAULT,PETSc.DEFAULT);
    self.ksp.getPC().setType('bjacobi')
    self.ksp.setFromOptions()

    (self.designFunctional,self.stateFunctional) = PISM.invert.ssa.createGradientFunctionals(ssarun)
    
    self.designForm = None


  def F(self, x,out=None,guess=None):
    """
    Returns the value of the forward problem at the design variable  ``x``.

    Nonlinear problems often make use of an initial guess; this can be provided in ``guess``.

    Storage in ``out``, if given, is used for the return value.
    """
    if out is None:
      out = self.rangeVector()
    reason = self.ssa.linearize_at(x.core())
    if reason.failed():
      raise PISM.AlgorithmFailureException(reason.description())
    out.core().copy_from(self.ssa.solution())
    return out

  def T(self,d,out=None):
    """
    Returns the value of the linearization :math:`T` of the forward problem at the design variable :math:`x` 
    specified previously in :meth:`linearizeAt`, in the direction `d`.

    Storage in `out`, if given, is used for the return value.
    """
    if out is None:
      out = self.rangeVector()
    self.ssa.apply_linearization(d.core(),out.core())
    return out

  def TStar(self,r,out=None):
    """
    Let :math:`T` be the linearization of the forward problem :math:`F` at the design variable `x` (as specified previously in :meth:`linearizeAt`).  
    Its adjoint is :math:`T^*`.  This method returns the value of :math:`T^*` in the direction `r`.

    Storage in `out`, if given, is used for the return value.
    """
    if out is None:
      out = self.domainVector()

    if self.designForm is None:
      da2 = self.grid.get_dm(1, self.grid.max_stencil_width)
      self.designForm = da2.getMatrix("baij")
      self.designFunctional.assemble_form(self.designForm)

    # First step
    self.stateFunctional.gradientAt(r.core(),self.tmpV)
    self.tmpV.scale(0.5)
    self.ssa.apply_linearization_transpose(self.tmpV,self.tmpS)

    # Second step
    self.ksp.setOperators(self.designForm,self.designForm,PETSc.Mat.Structure.SAME_NONZERO_PATTERN)
    self.ksp.solve(self.tmpS.get_vec(),self.tmpS2.get_vec())
    reason = self.ksp.getConvergedReason()
    if reason < 0:
      raise RuntimeError('TStarB linear solve failed to converge (KSP reason %s)\n\n' % reason)
    else:
      PISM.logging.logPrattle("TStarB converged (KSP reason %s)\n" % reason )
    
    out.core().copy_from(self.tmpS2)
    return out

  def linearizeAt(self,x,guess=None):
    """
    Instructs the class that subsequent calls to T and TStar will be conducted for the given value of x.

    Nonlinear problems often make use of an initial guess; this can be provided in 'guess'.
    """
    reason = self.ssa.linearize_at(x.core())
    if reason.failed():
      raise Exception(reason.description())

  def evalFandLinearize(self,x,out=None,guess=None):
    """
    Computes the value of F(x) and locks in a linearization.  Sometimes there are efficiencies that
    can be acheived this way.

    Default implementation simply calls F, then linearizeAt.
    """
    if out is None:
      out = self.rangeVector()
    self.linearizeAt(x)
    out.core().copy_from(self.ssa.solution())
    return out

  def rangeIP(self,a,b):
    """
    Computes the inner product of two vectors in the range (i.e. state) space.
    """
    return self.stateFunctional.dot(a.core(),b.core())

  def domainIP(self,a,b):
    """
    Computes the inner product of two vectors in the domain (i.e. design) space.
    """
    return self.designFunctional.dot(a.core(),b.core())

  def rangeVector(self):
    """Constructs a brand new vector from the range (i.e. state) vector space"""
    v = PISM.IceModelVec2V()
    v.create(self.grid,"",True,WIDE_STENCIL)

    # Add appropriate meta data.
    intent = "?inverse?" # FIXME
    desc = "SSA velocity computed by inversion"
    v.set_attrs(intent, "%s%s" %("X-component of the ",desc), "m s-1", "", 0);
    v.set_attrs(intent, "%s%s" %("Y-component of the ",desc), "m s-1", "", 1);
    v.set_glaciological_units("m year-1");
    v.write_in_glaciological_units = True
    huge_vel = self.grid.convert(1e6, "m/year", "m/second")
    attrs = [ ("valid_min", -huge_vel), ("valid_max", huge_vel), ("_FillValue", 2*huge_vel) ]
    for a in attrs: 
      for component in range(2):
        v.metadata(component).set_double(a[0],a[1])

    return PISMLocalVector(v)

  def domainVector(self):
    """Constructs a brand new vector from the domain (i.e. design) vector space"""
    v = PISM.IceModelVec2S()
    v.create(self.grid,"",True,WIDE_STENCIL)
    return PISMLocalVector(v)


class InvertSSANLCG(InvertNLCG):
  """Subclass of :class:`siple.gradient.nonlinear.InvertNLCG` for inversion of SSA velocities
  from :math:`\\tau_c` or hardness."""

  @staticmethod
  def defaultParameters():
    params = InvertNLCG.defaultParameters()
    return params

  def __init__(self,forward_problem,params=None):
    """:param forward_problem: A :class:`SSATaucForwardProblem` that defines the forward problem.
       :param params: A :class:`siple.params.Parameters` containing run-time parameters.
    """
    InvertNLCG.__init__(self,params)
    self.forward_problem = forward_problem

  def forwardProblem(self):
    """:returns: the associated :class:`SSATaucForwardProblem` provided at construction"""
    return self.forward_problem

  def stopConditionMet(self,count,x,Fx,y,r):
    """
    Determines if minimization should be halted (based, e.g. on a Morozov discrepancy principle)

    :param count: current iteration count
    :param x:     point in domain of potential minimizer.
    :param Fx:    value of nonlinear function at `x`
    :param r:     current residual, i.e. :math:`y-F(x)`
    :returns: boolean, `True` if termination condition is met
    """

    misfit = math.sqrt(abs(self.forward_problem.rangeIP(r,r)));

    if( misfit < self.misfit_goal ):
      siple.reporting.msg('Stop condition met')
      return True
    return False

  def initialize(self,x,y,deltaLInf):
    """
    Hook called at the start of :meth:`solve`.  This gives the class a chance to massage the input.

    The remaining arguments are passed directly from solve, and can be used for determining the
    final stopping criterion.

    Returns vectors corresponding to the initial value of `x` and the desired value of `y` where :math:`y=F(x)`.    
    """
    xv = PISMLocalVector(x)
    yv = PISMLocalVector(y)

    self.misfit_goal = self.params.mu*deltaLInf

    return (xv,yv)

  def finalize(self,x,y):
    """
    Hook called at the end of :meth:`solve`.  Gives the chance to massage the return values.
    """
    zeta = x.core()
    u = y.core()
    return (zeta,u)

class InvertSSAIGN(InvertIGN):
  """Subclass of :class:`siple.gradient.nonlinear.InvertIGN` for inversion of SSA velocities
  from :math:`\\tau_c` or hardness."""

  @staticmethod
  def defaultParameters():
    params = InvertIGN.defaultParameters()
    return params

  def __init__(self,forward_problem,params=None):
    """:param forward_problem: A :class:`SSATaucForwardProblem` that defines the forward problem.
       :param params: A :class:`siple.params.Parameters` containing run-time parameters.
    """
    InvertIGN.__init__(self,params)
    self.forward_problem = forward_problem

  def forwardProblem(self):
    """:returns: the associated :class:`SSATaucForwardProblem` provided at construction"""
    return self.forward_problem

  def temper_d(self, x,d,y,r):
    """Method called during iterations in :meth:`solve` to ensure we don't make to large a step.
    :param x: current design parameter
    :param d: step in design space
    :param y: desired value of state paramter
    :param r: residual of desired and current state parameter values (:math:`y-F(x)`)
    
    Changes `d` as a side-effect to temper step size.
    """
    dnorm = d.norm('linf');  xnorm = x.norm('linf')
    if dnorm > 2*xnorm:
      siple.reporting.msg('wild change predicted by linear step. scaling')
      d.scale(2*xnorm/dnorm)

  def initialize(self,x,y,target_misfit):
    """
    Hook called at the start of :meth:`solve`.  This gives the class a chance to massage the input.

    The remaining arguments are passed directly from solve, and can be used for determining the
    final stopping criterion.

    Returns vectors corresponding to the initial value of `x` and the desired value of `y` where :math:`y=F(x)`.    
    """
    xv = PISMLocalVector(x)
    yv = PISMLocalVector(y)

    return (xv,yv,target_misfit)

  def finalize(self,x,y):
    """
    Hook called at the end of :meth:`solve`.  Gives the chance to massage the return values.
    """

    zeta = x.core()
    u = y.core()

    return (zeta,u)


class SipleIterationListenerAdaptor:
  """Adaptor for passing listening events from `siple`-based solvers to a python object."""
  def __init__(self,owner,listener):
    """:param owner: The :class:`InvSSATaucSolver_Tikhonov` that constructed us
       :param listener: The python-based listener.
     """
    self.owner = owner
    self.listener = listener
  def __call__(self,siplesolver,it,x,Fx,y,d,r,*args):    
    """Callback from `siple`.  Gathers together the long list of arguments
    into a dictionary and passes it along in a standard form to the python listener."""   
    data = Bunch(zeta=x.core(),u=Fx.core(),zeta_step=d.core(),residual=r.core(),target_misfit=self.owner.target_misfit)

    if self.owner.method == 'ign':
      data.update(T_zeta_step=args[0].core())
    else:
      data.update(TStar_residual=args[0].core())
    try:
      self.listener(self.owner,it,data)
    except Exception as e:
      logError("\nWARNING: Exception occured during an inverse solver listener callback:\n%s\n\n" % str(e))

class SipleLinearIterationListenerAdaptor:
  """Adaptor for passing listening events the linear steps of `siple`-based `ign` solvers to a python object."""
  def __init__(self,owner,listener):
    """:param owner: The :class:`InvSSATaucSolver_Tikhonov` that constructed us
       :param listener: The python-based listener.
     """
    self.owner = owner
    self.listener = listener

  def __call__(self,siplesolver,it,x, y, d, r, Td, TStarR):    
    """Callback from `siple`.  Gathers together the long list of arguments
    into a dictionary and passes it along in a standard form to the python listener."""   

    data = Bunch(x=x.core(),y=y.core(),r=r.core(),d=d.core(),Td=Td.core(),TSTarR=TStarR.core())
    try:
      self.listener(self.owner,it,data)
    except Exception as e:
      logError("\nWARNING: Exception occured during an inverse solver listener callback:\n%s\n\n" % str(e))

class SipleDesignUpdateListenerAdaptor:
  """Adaptor for design variable update events of `siple`-based solvers to a python object."""
  def __init__(self,owner,listener):
    """:param owner: The :class:`InvSSATaucSolver_Tikhonov` that constructed us
       :param listener: The python-based listener.
     """
    self.owner = owner
    self.listener = listener

  def __call__(self,siplesolver,it,zeta,u,u_obs,r):    
    """Callback from `siple`.  Gathers together the long list of arguments
    into a dictionary and passes it along in a standard form to the python listener."""   
    data = Bunch(zeta=zeta.core(),u=u.core(),r=r.core(),u_obs=u_obs.core())
    try:
      self.listener(self.owner,it,data)
    except Exception as e:
      logError("\nWARNING: Exception occured during an inverse solver DesignUpdate listener callback:\n%s\n\n" % str(e))

class MonitorAdjoint:
  """Iteration listener that can be used to verify the correctness of the implementation of an adjoint.
  For adjoint-based interative inverse methods, a residual ``r`` is known in state space and a step direction ``d`` is 
  known in design space.  A linearized forward problem :math:`T` maps from design space to state space, and its adjoint :math:`T^*`
  goes in the opposite direction.  The inner products :math:`\left<Td,r\\right>_{\\rm State}` 
  and :math:`\left<d,T^*r\\right>_{\\rm Design}` should always be the same; this is a good diagnostic to determine
  of an adjoint has been coded correctly. The listener prints a comparison of the values of the two inner products
  at each iteration.  
  """

  def __init__(self):
    self.Td = None
    self.TStarR = None
    self.didWarning = False

  def __call__(self,inverse_solver,count,data):
    """
    :param inverse_sovler: the solver (e.g. :class:`~InvSolver_Tikhonov`) we are listening to.
    :param count: the iteration number.
    :param data: dictionary of data related to the iteration.
    """
    import PISM.invert.sipletools
    method = inverse_solver.method
    if method != 'sd' and method !='nlcg' and method != 'ign':
      if not self.didWarning:
        PISM.verbPrintf(1,PISM.Context().com,'\nWarning: unable to monitor adjoint for inverse method: %s\nOption -inv_monitor_adjoint ignored\n' % method)
      self.didWarning = True
      return
    fp = inverse_solver.forward_problem
    d = PISM.invert.sipletools.PISMLocalVector(data.d)
    r = PISM.invert.sipletools.PISMLocalVector(data.r)
    self.Td=fp.T(d,self.Td)
    self.TStarR=fp.TStar(r,out=self.TStarR)
    ip1 = fp.domainIP(d,self.TStarR)
    ip2 = fp.rangeIP(self.Td,r)
    logMessage("adjoint test: <Td,r>=%g <d,T^*r>=%g (percent error %g)",ip1,ip2,(abs(ip1-ip2))/max(abs(ip1),abs(ip2)))


class MonitorAdjointLin:
  def __init__(self):
    self.Td = None
    self.TStarR = None

  def __call__(self,inverse_solver,count,data):
    """
    :param inverse_sovler: the solver (e.g. :class:`~InvSSASolver_Tikhonov`) we are listening to.
    :param count: the iteration number.
    :param data: dictionary of data related to the iteration.
    """
    import PISM.invert.sipletools
    fp = inverse_solver.forward_problem
    r = PISM.invert.sipletools.PISMLocalVector(data.r)
    d = PISM.invert.sipletools.PISMLocalVector(data.d)
    self.Td=fp.T(d,self.Td)
    self.TStarR=fp.TStar(r,out=self.TStarR)
    ip1 = fp.domainIP(d,self.TStarR)
    ip2 = fp.rangeIP(self.Td,r)
    logMessage("adjoint test: <Td,r>=%g <d,T^*r>=%g (percent error %g)",ip1,ip2,(abs(ip1-ip2))/max(abs(ip1),abs(ip2)))