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

import PISM

has_siple = False
try:
  import siple
  has_siple = True
except:
  pass

has_tao = False
try:
  PISM.__getattribute__('TaoInitializer')
  has_tao = True
except:
  pass

if not PISM.imported_from_sphinx:
  from petsc4py import PETSc

import math

from PISM.logging import logMessage

class SSAForwardRun(PISM.ssa.SSARun):
  """Subclass of :class:`PISM.ssa.SSAFromInputFile` where the underlying SSA implementation is an
  :cpp:class:`IP_SSATaucForwardProblem` or :cpp:class:`IP_SSAHardavForwardProblem`.
  It is responsible for putting together a :class:`PISM.model.ModelData` containing the auxilliary data
  needed for solving the SSA (:cpp:class:`IceModelVec`\s, :cpp:class:`EnthalpyConverter`, etc.) as well
  as the instance of :cpp:class:`IP_SSATaucForwardProblem` that will solve the SSA repeatedly in the course
  of solving an inverse problem.  This class is intended to be subclassed by test cases where the data
  is not provided from an input file.  See also :class:`SSAForwardRunFromInputFile`."""

  def __init__(self, design_var):
    PISM.ssa.SSARun.__init__(self)
    assert(design_var in ssa_forward_problems.keys())
    self.grid = PISM.Context().newgrid()
    self.config = self.grid.config
    self.design_var = design_var
    self.design_var_param = createDesignVariableParam(self.config, self.design_var)

  def designVariable(self):
    """:returns: String description of the design variable of the forward problem (e.g. 'tauc' or 'hardness')"""
    return self.design_var

  def designVariableParameterization(self):
    """:returns: Object that performs zeta->design variable transformation."""
    return self.design_var_param

  def _setFromOptions(self):
    """Initialize internal parameters based on command-line flags. Called from :meth:`PISM.ssa.SSARun.setup`."""
    for o in PISM.OptionsGroup(PISM.Context().com,"","SSAForwardRun"):
      self.is_regional = PISM.optionsIsSet("-regional")

  def _constructSSA(self):
    """Returns an instance of :cpp:class:`IP_SSATaucForwardProblem` rather than
       a basic :cpp:class:`SSAFEM` or :cpp:class:`SSAFD`. Called from :meth:`PISM.ssa.SSARun.setup`."""
    md = self.modeldata
    return createSSAForwardProblem(md.grid,md.enthalpyconverter,self.design_var_param,self.config,self.design_var)

  def _initSSA(self):
    """One-time initialization of the :cpp:class:`IP_SSATaucForwardProblem`. Called from :meth:`PISM.ssa.SSARun.setup`."""

    vecs = self.modeldata.vecs

    if vecs.has('vel_bc'):
      self.ssa.set_boundary_conditions(vecs.bc_mask,vecs.vel_bc)

    self.ssa.init(vecs.asPISMVars())

    # Cache the values of the coefficeints at quadrature points once here.
    # Subsequent solves will then not need to cache these values.
    self.ssa.cacheQuadPtValues();


class SSAForwardRunFromInputFile(SSAForwardRun):
  """Subclass of :class:`SSAForwardRun` where the vector data
  for the run is provided in an input :file:`.nc` file."""

  def __init__(self,input_filename,inv_data_filename,design_var):
    """
    :param input_filename:    :file:`.nc` file containing generic PISM model data.
    :param inv_data_filename: :file:`.nc` file containing data specific to inversion (e.g. observed SSA velocities).
    """
    SSAForwardRun.__init__(self,design_var)
    self.input_filename = input_filename
    self.inv_data_filename = inv_data_filename

  def _initGrid(self):
    """Initialize grid size and periodicity. Called from :meth:`PISM.ssa.SSARun.setup`."""
    # The implementation in PISM.ssa.SSAFromInputFile uses a non-periodic
    # grid only if the run is regional and "ssa_method=fem" in the config
    # file.  For inversions, we always use an FEM type method, so for
    # regional inversions, we always use a non-periodic grid.
    periodicity = PISM.XY_PERIODIC
    (pstring,pflag) = PISM.optionsListWasSet(self.grid.com,'-periodicity',"Grid periodicity",['x','y','xy', 'none'],'xy')
    if pflag:
      pdict = {'x':PISM.X_PERIODIC,   'y':PISM.Y_PERIODIC,
               'xy':PISM.XY_PERIODIC, 'none':PISM.NOT_PERIODIC }
      periodicity = pdict[pstring]
    else:
      periodicity = PISM.XY_PERIODIC
      if self.is_regional:
        periodicity=PISM.NOT_PERIODIC
    PISM.model.initGridFromFile(self.grid,self.input_filename,periodicity);

  def _initPhysics(self):
    """Override of :meth:`SSARun._initPhysics` that sets the physics based on command-line flags."""
    config = self.config

    enthalpyconverter = PISM.EnthalpyConverter(config)
    if PISM.getVerbosityLevel() >3:
      enthalpyconverter.viewConstants(PETSc.Viewer.STDOUT())

    if PISM.optionsIsSet("-ssa_glen"):
      config.set_string("ssa_flow_law", "isothermal_glen")
      config.scalar_from_option("ice_softness", "ice_softness")
    else:
      config.set_string("ssa_flow_law", "gpbld")

    self.modeldata.setPhysics(enthalpyconverter)

  def _initSSACoefficients(self):
    """Reads SSA coefficients from the input file. Called from :meth:`PISM.ssa.SSARun.setup`."""
    self._allocStdSSACoefficients()

    # Read PISM SSA related state variables
    #
    # Hmmm.  A lot of code duplication with SSAFromInputFile._initSSACoefficients.

    vecs = self.modeldata.vecs
    thickness = vecs.thickness; bed = vecs.bed; enthalpy = vecs.enthalpy
    mask = vecs.ice_mask; surface = vecs.surface

    # Read in the PISM state variables that are used directly in the SSA solver
    for v in [thickness, bed, enthalpy]:
      v.regrid(self.input_filename,True)

    # variables mask and surface are computed from the geometry previously read
    sea_level = 0 # FIXME setFromOption?
    gc = PISM.GeometryCalculator(sea_level, self.config)
    gc.compute(bed,thickness,mask,surface)


    grid = self.grid

    # Compute yield stress from PISM state variables
    # (basal melt rate, tillphi, and basal water height) if they are available

    if( PISM.util.fileHasVariable(self.input_filename,'bmelt') and
        PISM.util.fileHasVariable(self.input_filename,'bwat') and
        PISM.util.fileHasVariable(self.input_filename,'tillphi') ):

      bmr   = PISM.model.createBasalMeltRateVec(grid)
      tillphi = PISM.model.createTillPhiVec(grid)
      bwat = PISM.model.createBasalWaterVec(grid)
      for v in [bmr,tillphi,bwat]:
         v.regrid(self.input_filename,True)
         vecs.add(v)
      if self.is_regional:
        yieldstress = PISM.PISMRegionalDefaultYieldStress(self.modeldata.grid,self.modeldata.config)
      else:
        yieldstress = PISM.PISMMohrCoulombYieldStress(self.modeldata.grid,self.modeldata.config)
      yieldstress.init(vecs.asPISMVars())
      yieldstress.basal_material_yield_stress(vecs.tauc)
    elif PISM.util.fileHasVariable(self.input_filename,'tauc'):
      vecs.tauc.regrid(self.input_filename,True)

    if PISM.util.fileHasVariable(self.input_filename,'ssa_driving_stress_x'):
      vecs.add( PISM.model.createDrivingStressXVec(self.grid))
      vecs.ssa_driving_stress_x.regrid(self.input_filename,critical=True)

    if PISM.util.fileHasVariable(self.input_filename,'ssa_driving_stress_y'):
      vecs.add( PISM.model.createDrivingStressYVec(self.grid))
      vecs.ssa_driving_stress_y.regrid(self.input_filename,critical=True)

    if self.is_regional:
      vecs.add( PISM.model.createNoModelMaskVec(self.grid), 'no_model_mask' )
      vecs.no_model_mask.regrid(self.input_filename,True)
      vecs.add( vecs.surface, 'usurfstore')
      vecs.setPISMVarsName('usurfstore','usurfstore')

    if self.config.get_flag('ssa_dirichlet_bc'):
      vecs.add( PISM.model.create2dVelocityVec( self.grid, name='_ssa_bc', desc='SSA velocity boundary condition',intent='intent' ), "vel_ssa_bc" )
      has_u_ssa_bc = PISM.util.fileHasVariable(self.input_filename,'u_ssa_bc');
      has_v_ssa_bc = PISM.util.fileHasVariable(self.input_filename,'v_ssa_bc');
      if (not has_u_ssa_bc) or (not has_v_ssa_bc):
        PISM.verbPrintf(2,self.grid.com, "Input file '%s' missing Dirichlet boundary data u/v_ssa_bc; using zero default instead." % self.input_filename)
        vecs.vel_ssa_bc.set(0.)
      else:
        vecs.vel_ssa_bc.regrid(self.input_filename,True)

      if self.is_regional:
        vecs.add( vecs.no_model_mask, 'bc_mask')
      else:
        vecs.add( PISM.model.createBCMaskVec( self.grid ), 'bc_mask' )
        bc_mask_name = vecs.bc_mask.string_attr("name")
        if PISM.util.fileHasVariable(self.input_filename,bc_mask_name):
          vecs.bc_mask.regrid(self.input_filename,True)
        else:
          PISM.verbPrintf(2,self.grid.com,"Input file '%s' missing Dirichlet location mask '%s'.  Default to no Dirichlet locations." %(self.input_filename,bc_mask_name))
          vecs.bc_mask.set(0)
      # We call this variable 'bc_mask' in the python code, it is called
      # 'bcflag' when passed between pism components, and it has yet
      # another name when written out to a file.  Anyway, we flag its
      # export to PISMVars name here.
      vecs.setPISMVarsName('bc_mask','bcflag')

    if PISM.util.fileHasVariable(self.inv_data_filename, 'vel_misfit_weight'):
        vecs.add( PISM.model.createVelocityMisfitWeightVec(self.grid) )
        vecs.vel_misfit_weight.regrid(self.inv_data_filename,True)


class InvSSASolver:
  """Abstract base class for SSA inverse problem solvers."""

  def __init__(self,ssarun,method):
    """
    :param ssarun: The :class:`PISM.invert.ssa.SSAForwardRun` defining the forward problem.
    :param method: String describing the actual algorithm to use. Must be a key in :attr:`tao_types`."""

    self.ssarun = ssarun
    self.config = ssarun.config
    self.method = method

  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.
    """
    raise NotImplementedError()

  def addIterationListener(self,listener):
    """Add a listener to be called after each iteration.  See :ref:`Listeners`."""
    raise NotImplementedError()

  def addDesignUpdateListener(self,listener):
    """Add a listener to be called after each time the design variable is changed."""
    raise NotImplementedError()

  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 :math:`\zeta`.
    :param u_obs: :cpp:class:`IceModelVec2V` of observed surface velocities.
    :param zeta_inv: :cpp:class:`zeta_inv` starting value of :math:`\zeta` for minimization of the Tikhonov functional.
    :returns: A :cpp:class:`TerminationReason`.
    """
    raise NotImplementedError()

  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."""
    raise NotImplementedError()


def createInvSSASolver(ssarun,method=None):
  """Factory function returning an inverse solver appropriate for the config variable ``inv_ssa_method``.

  :param ssarun: an instance of :class:`SSAForwardRun:` or :class:`SSAForwardRunFromInputFile`.
  :param method: a string correpsonding to config variable ``inv_ssa_method`` describing the inversion method to be used.
  """
  if method is None:
    method = ssarun.config.get_string('inv_ssa_method')
  if method == 'tikhonov_gn':
    import ssa_gn
    return ssa_gn.InvSSASolver_Tikhonov(ssarun,method)
  elif method.startswith('tikhonov'):
    if not has_tao:
      raise RuntimeError("Inversion method '%s' requires the TAO library.\nInstall from http://www.mcs.anl.gov/tao and rebuild PISM with TAO support." % method)
    import ssa_tao
    return ssa_tao.InvSSASolver_Tikhonov(ssarun,method)
  if method == 'sd' or method == 'nlcg' or method == 'ign':
    if not has_siple:
      raise RuntimeError("Inversion method '%s' requires the siple python library.\nInstall from https://github.com/damaxwell/siple" % method)
    import ssa_siple
    return ssa_siple.InvSSASolver_Gradient(ssarun,method)
  raise Exception("Unknown inverse method '%s'; unable to construct solver.",method)


design_param_types = {"ident":PISM.IPDesignVariableParamIdent,
                      "square":PISM.IPDesignVariableParamSquare,
                      "exp":PISM.IPDesignVariableParamExp,
                      "trunc":PISM.IPDesignVariableParamTruncatedIdent }

def createDesignVariableParam( config, design_var_name, param_name=None ):
  """Factory function for creating subclasses of :cpp:class:`IPDesignVariableParameterization` based on command-line flags."""
  if param_name is None:
    param_name = config.get_string("inv_design_param")
  design_param = design_param_types[param_name]()
  design_param.set_scales(config,design_var_name)
  return design_param

ssa_forward_problems = {'tauc'   : PISM.IP_SSATaucForwardProblem,
                        'hardav' : PISM.IP_SSAHardavForwardProblem}

def createSSAForwardProblem(grid,ec,design_param,config,design_var):
  """Returns an instance of an SSA forward problem (e.g. :cpp:class:`IP_SSATaucForwardProblem`)
  suitable for the value of `design_var`"""
  ForwardProblem = ssa_forward_problems[design_var]
  if ForwardProblem is None:
    raise RuntimeError("Design variable %s is not yet supported.", design_var)
    
  return ForwardProblem(grid,ec,design_param,config)


def createGradientFunctionals(ssarun):
  """Returns a tuple ``(designFunctional,stateFunctional)`` of :cpp:class:`IP_IPFunctional`\s 
  for gradient-based inversions.  The specific functionals are constructed on the basis of 
  command-line parameters ``inv_state_func`` and ``inv_design_func``.  

  :param ssarun: The instance of :class:`PISM.ssa.SSARun` that encapsulates the forward problem,
                 typically a :class:`SSAForwardRunFromFile`.
  """

  vecs = ssarun.modeldata.vecs
  grid = ssarun.grid


  for o in PISM.OptionsGroup(grid.com,"","InvSSAFunctionals"):
      useGroundedIceOnly = PISM.optionsIsSet("-inv_ssa_grounded_ice_tauc","Computed norms for tau_c only on elements with all grounded ice.")

  misfit_type = grid.config.get_string("inv_state_func")
  if misfit_type != 'meansquare':
    inv_method = grid.config.get_string("inv_ssa_method")
    raise Exception("'-inv_state_func %s' is not supported with '-inv_method %s'.\nUse '-inv_state_func meansquare' instead" % (misfit_type,inv_method))

  design_functional = grid.config.get_string("inv_design_func")
  if design_functional != "sobolevH1":
    inv_method = grid.config.get_string("inv_ssa_method")
    raise Exception("'-inv_design_func %s' is not supported with '-inv_method %s'.\nUse '-inv_design_func sobolevH1' instead" % (design_functional,inv_method))

  designFunctional = createHilbertDesignFunctional(grid,vecs,useGroundedIceOnly)

  stateFunctional = createMeanSquareMisfitFunctional(grid,vecs)

  return (designFunctional,stateFunctional)

def createTikhonovFunctionals(ssarun):
  """Returns a tuple ``(designFunctional,stateFunctional)`` of :cpp:class:`IP_Functional`\s 
  for Tikhonov-based inversions.  The specific functionals are constructed on the basis of 
  command-line parameters ``inv_state_func`` and ``inv_design_func``.  
  
  :param ssarun: The instance of :class:`PISM.ssa.SSARun` that encapsulates the forward problem,
                 typically a :class:`SSATaucForwardRunFromFile`.
"""
  vecs = ssarun.modeldata.vecs
  grid = ssarun.grid

  for o in PISM.OptionsGroup(grid.com,"","SSATikhonovFunctionals"):
      useGroundedIceOnly = PISM.optionsIsSet("-inv_ssa_grounded_ice_tauc","Computed norms for tau_c only on elements with all grounded ice.")

  misfit_type = grid.config.get_string("inv_state_func")
  if misfit_type == "meansquare":  
    stateFunctional = createMeanSquareMisfitFunctional(grid,vecs)
  elif misfit_type == "log_ratio":
    vel_ssa_observed = vecs.vel_ssa_observed
    scale = grid.config.get("inv_log_ratio_scale")
    velocity_eps = grid.config.get("inv_ssa_velocity_eps", "m/year", "m/second")
    misfit_weight = None
    if vecs.has('vel_misfit_weight'):
      misfit_weight = vecs.vel_misfit_weight
    stateFunctional = PISM.IPLogRatioFunctional(grid, vel_ssa_observed, velocity_eps, misfit_weight)
    stateFunctional.normalize(scale);
  elif misfit_type == "log_relative":
    vel_ssa_observed = vecs.vel_ssa_observed
    velocity_scale = grid.config.get("inv_ssa_velocity_scale", "m/year", "m/second")
    velocity_eps = grid.config.get("inv_ssa_velocity_eps", "m/year", "m/second")
    misfit_weight = None
    if vecs.has('vel_misfit_weight'):
      misfit_weight = vecs.vel_misfit_weight
    stateFunctional = PISM.IPLogRelativeFunctional(grid, vel_ssa_observed, velocity_eps, misfit_weight)
    stateFunctional.normalize(velocity_scale);
  else:
    raise RuntimeError("Unknown inv_state_func '%s'; unable to construct solver.",misfit_type)

  design_functional = grid.config.get_string("inv_design_func")
  if design_functional == "sobolevH1":
    designFunctional = createHilbertDesignFunctional(grid,vecs,useGroundedIceOnly)
  elif design_functional == "tv":
    area = 4*grid.Lx*grid.Ly;
    velocity_scale = grid.config.get("inv_ssa_velocity_scale", "m/year", "m/second")
    length_scale = grid.config.get("inv_ssa_length_scale");
    lebesgue_exponent = grid.config.get("inv_ssa_tv_exponent");
    cTV = 1/area
    cTV *= (length_scale)**(lebesgue_exponent)
    
    zeta_fixed_mask = None
    if vecs.has('zeta_fixed_mask'):
      zeta_fixed_mask = vecs.zeta_fixed_mask

    for o in PISM.OptionsGroup(grid.com, title="Total Variation Functional"):
      strain_rate_eps = PISM.optionsReal("-inv_ssa_tv_eps",
        "regularization constant for 'total variation' functional", default=None)
    if strain_rate_eps is None:
      schoofLen = grid.config.get("Schoof_regularizing_length", "km", "m");
      strain_rate_eps = 1/schoofLen

    designFunctional = PISM.IPTotalVariationFunctional2S(grid,cTV,lebesgue_exponent,strain_rate_eps,zeta_fixed_mask)
  else:
    raise Exception("Unknown inv_design_func '%s'; unable to construct solver." % design_functional)

  return (designFunctional,stateFunctional)

def createMeanSquareMisfitFunctional(grid,vecs):
  """Creates a :cpp:class:`IPMeanSquareFunctional2V` suitable for use for a 
  state variable function for SSA inversions."""

  misfit_weight = None
  if vecs.has('vel_misfit_weight'):
    misfit_weight = vecs.vel_misfit_weight

  velocity_scale = grid.config.get("inv_ssa_velocity_scale", "m/year", "m/second")
  stateFunctional = PISM.IPMeanSquareFunctional2V(grid,misfit_weight);    
  stateFunctional.normalize(velocity_scale);
  return stateFunctional

def createHilbertDesignFunctional(grid,vecs,useGroundedIceOnly):
  """Creates a :cpp:class:`IP_H1NormFunctional2S` or a :cpp:class`IPGroundedIceH1NormFunctional2S` suitable
  for use for a design variable functional.
  
  :param grid: computation grid
  :param vecs: model vecs
  :param useGroundedIceOnly: flag, ``True`` if a :cpp:class`IPGroundedIceH1NormFunctional2S` should be created.
"""
  cL2 = grid.config.get("inv_design_cL2");
  cH1 = grid.config.get("inv_design_cH1");

  area = 4*grid.Lx*grid.Ly;
  length_scale = grid.config.get("inv_ssa_length_scale");
  cL2 /= area;
  cH1 /= area;
  cH1 *= (length_scale*length_scale);

  zeta_fixed_mask = None
  if vecs.has('zeta_fixed_mask'):
    zeta_fixed_mask = vecs.zeta_fixed_mask

  if useGroundedIceOnly:
    ice_mask =vecs.ice_mask
    designFunctional = PISM.IPGroundedIceH1NormFunctional2S(grid,cL2,cH1,ice_mask,zeta_fixed_mask)
  else:
    designFunctional = PISM.IP_H1NormFunctional2S(grid,cL2,cH1,zeta_fixed_mask)

  return designFunctional

def printIteration(invssa_solver,it,data):
  logMessage("----------------------------------------------------------\n");
  logMessage("Iteration %d\n" % it)    

def printTikhonovProgress(invssasolver,it,data):
    eta = data.tikhonov_penalty
    stateVal = data.JState
    designVal = data.JDesign
    grid = invssasolver.ssarun.grid
    sWeight = 1
    dWeight = 1./eta
    
    logMessage("design objective %.8g; weighted %.8g\n" % (designVal,designVal*dWeight)) 
    if data.has_key('grad_JTikhonov'):
      logMessage("gradient: design %.8g state %.8g sum %.8g\n" % (data.grad_JDesign.norm(PETSc.NormType.NORM_2)*dWeight,data.grad_JState.norm(PETSc.NormType.NORM_2)*sWeight,data.grad_JTikhonov.norm(PETSc.NormType.NORM_2)))
    else:
      logMessage("gradient: design %.8g state %.8g; constraints: %.8g\n" % (data.grad_JDesign.norm(PETSc.NormType.NORM_2)*dWeight,data.grad_JState.norm(PETSc.NormType.NORM_2)*sWeight,data.constraints.norm(PETSc.NormType.NORM_2)))
    logMessage("tikhonov functional: %.8g\n" % (stateVal*sWeight + designVal*dWeight) )

class RMSMisfitReporter:
  def __init__(self):
    self.J = None

  def __call__(self, invssa_solver, it, data):

    grid = invssa_solver.ssarun.grid

    if self.J is None:
      vecs = invssa_solver.ssarun.modeldata.vecs      
      self.J = createMeanSquareMisfitFunctional(grid, vecs)
    
    Jmisfit = self.J.valueAt(data.residual)
    rms_misfit = math.sqrt(Jmisfit) * grid.config.get("inv_ssa_velocity_scale")

    PISM.logging.logMessage("Diagnostic RMS Misfit: %0.8g (m/a)\n" % rms_misfit)


class MisfitLogger:
  def __init__(self):
    self.misfit_history = []
    self.misfit_type = None
    
  def __call__(self,invssa_solver,it,data):
    """
    :param inverse_solver: 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.
    """

    grid = invssa_solver.ssarun.grid

    if self.misfit_type is None:
        self.misfit_type = grid.config.get_string("inv_state_func")

    method = invssa_solver.method
    if method == 'ign' or method == 'sd' or method == 'nlcg':
      import PISM.invert.sipletools
      fp = invssa_solver.forward_problem
      r=PISM.invert.sipletools.PISMLocalVector(data.residual)
      Jmisfit = fp.rangeIP(r,r)
    elif data.has_key('JState'):
      Jmisfit = data.JState
    else:
      raise RuntimeError("Unable to report misfits for inversion method: %s" % method)

    if self.misfit_type == "meansquare":
        velScale_m_per_year = grid.config.get("inv_ssa_velocity_scale")
        velScale_m_per_s = grid.convert(velScale_m_per_year, "m/year", "m/second")

        rms_misfit = math.sqrt(Jmisfit) * velScale_m_per_year
        
        logMessage("Misfit: sqrt(J_misfit) = %.8g (m/a)\n" % rms_misfit);
        self.misfit_history.append(rms_misfit)
    else:
        logMessage("Misfit: J_misfit = %.8g (dimensionless)\n" % Jmisfit);
        self.misfit_history.append(Jmisfit)

  def write(self,output_filename):
    """Saves a history of misfits as :ncvar:`inv_ssa_misfit`

    :param output_filename: filename to save misfits to."""
    if PISM.Context().rank == 0:
      nc = PISM.netCDF.Dataset(output_filename, 'a')  # append
      nc.createDimension('inv_ssa_iter',len(self.misfit_history))
      nc_misfit = nc.createVariable('inv_ssa_misfit','f8',dimensions=('inv_ssa_iter'))
      if self.misfit_type == "meansquare":
        nc_misfit.setncattr('_units','m/a')
      nc_misfit[:] = self.misfit_history[:]
      nc.close()


class ZetaSaver:
  """Iteration listener used to save a copy of the current value
  of :math:`\zeta` (i.e. a parameterized design variable such as :math:`\tau_c` or hardness) 
  at each iteration during an inversion. The intent is to use a saved value to restart 
  an inversion if need be.
  """
  def __init__(self,output_filename):
    """:param output_filename: file to save iterations to."""
    self.output_filename = output_filename

  def __call__(self,inverse_solver,count,data):
    zeta = data.zeta
    # The solver doesn't care what the name of zeta is, and we
    # want it called 'zeta_inv' in the output file, so we rename it.
    zeta.rename('zeta_inv', 'last iteration of parameterized basal yeild stress computed by inversion','')
    zeta.write(self.output_filename)