# Copyright (C) 2011, 2012, 2013, 2014, 2015 David Maxwell and Constantine Khroulev
#
# 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

"""Module containing classes managing SSA forward runs and
SSA verification test cases."""

import PISM
import math
from PISM import util, model

# Conversion from command-line arguments to classes of SSA solver.
SSAAlgorithms = {"fem": PISM.SSAFEM, "fd": PISM.SSAFD}

class SSARun(object):

    """Mediates solving PISM's SSA model from a minimal set of data, without the constrution of an :cpp:class:`iceModel`.
       It codifies the steps needed to put together the data for an SSA run; subclasses do the work of
       implementing the steps in :meth:`_setFromOptions`, :meth:`_initGrid`, etc.  Uses include:

         * Running SSA test cases.
         * Running the SSA in standalone mode (e.g. via :command:`ssaforward.py`)
         * The SSA inversion code.

       Usage:  After construction (of a subclass),

         1. Call :meth:`setup` to run through the various
               steps needed to set up an environment for solving the SSA.
         2. Solve the SSA with :meth:`solve`.
         3. Optionally write the the model vectors and solution to a file with :meth:`write`."""

    def __init__(self):
        """Do little constructor.  Real work is done by :meth:`setup` which should be called prior to :meth:`solve`."""
        self.grid = None         #: The computation grid; will be set by :meth:`_initGrid`
        self.config = None       #: Placeholder for config dictionary; set indirectly by :meth:`_constructModelData`

        #: Instance of :class:`PISM.model.ModelData` that stores all data needed for solving the SSA. Much of the work of
        #: the :class:`SSARun` is involved in setting up this object. Tasks include setting up :cpp:class:IceModelVec
        #: variables as well as model physics (e.g. :cpp:class:`EnthalpyConverter`).
        self.modeldata = None
        self.ssa = None          #: Subclass of :cpp:class:`SSA` that sovles the SSA.

    def setup(self):
        """Orchestrates the steps of setting up an environment for running the SSA.  The following methods
           are called in order, and should be impelmeneted by a subclass.

             1. :meth:`_setFromOptions` to set any parameters from command-line options
             2. :meth:`_initGrid` to determine the computation grid, to be stored as :attr:`grid`
             3. :meth:`_constructModelData` provide a :class:`ModelData` object (a default implementation is provided)
             4. :meth:`_initPhysics` to set the non-vec members of the :class:`ModelData`, e.g. the :cpp:class:`EnthalpyConverter`.
             5. :meth:`_constructSSA` to build the actual subclass of :cpp:class:`SSA` that will be used to solve the SSA
             6. :meth:`_initSSACoefficients` enter all of the vecs needed for solving the SSA into the :class:`ModelData`.
             7. :meth:`_initSSA` initialize the :cpp:class:`SSA` returned in step 5
             """
        self._setFromOptions()

        self._initGrid()
        if self.grid is None:
            raise RuntimeError("SSARun failed to provide a grid.")

        self.modeldata = self._constructModelData()
        if self.modeldata is None:
            raise RuntimeError("SSARun._constructModelData failed to provide a ModelData.")
        self.config = self.modeldata.config

        self._initPhysics()
        if self.modeldata.enthalpyconverter is None:
            raise RuntimeError("SSARun._initPhysics failed to initialize the physics of the underlying SSA solver.")

        self.ssa = self._constructSSA()
        if self.ssa is None:
            raise RuntimeError("SSARun._constructSSA failed to provide an SSA.")

        self._initSSACoefficients()
        # FIXME: is there a reasonable check to do here?

        self._initSSA()

    def solve(self):
        """Solve the SSA by calling the underlying PISM :cpp:class:`SSA`'s
        :cpp:member:`update` method. Returns the solution vector (owned by
        self.ssa, but you should not need to know about ownership).

        """
        vecs = self.modeldata.vecs

        # make sure vecs is locked!
        self.ssa.init()
        if vecs.has('vel_bc'):
            self.ssa.set_boundary_conditions(vecs.bc_mask, vecs.vel_bc)

        melange_back_pressure = PISM.IceModelVec2S()
        melange_back_pressure.create(self.grid, "melange_back_pressure", PISM.WITHOUT_GHOSTS)
        melange_back_pressure.set_attrs("diagnostic",
                                        "melange back pressure fraction", "1", "")

        PISM.verbPrintf(2, self.grid.com, "* Solving the SSA stress balance ...\n")

        fast = False

        self.ssa.update(fast, melange_back_pressure)

        return self.ssa.velocity()

    def write(self, filename):
        """Saves all of :attr:`modeldata`'s vecs (and the solution) to an
        output file."""
        grid = self.grid
        vecs = self.modeldata.vecs

        pio = PISM.PIO(grid.com, "netcdf3")
        pio.open(filename, PISM.PISM_READWRITE_MOVE)
        PISM.define_time(pio, grid.ctx().config().get_string("time_dimension_name"),
                         grid.ctx().config().get_string("calendar"),
                         grid.ctx().time().units_string(),
                         grid.ctx().unit_system())
        PISM.append_time(pio, grid.ctx().config().get_string("time_dimension_name"), 0.0)
        pio.close()

        # Save time & command line
        PISM.util.writeProvenance(filename)

        vecs.writeall(filename)

        vel_ssa = self.ssa.velocity()
        vel_ssa.write(filename)

        sys = self.grid.ctx().unit_system()

        velbar_mag = model.createCBarVec(self.grid)
        velbar_mag.set_to_magnitude(vel_ssa)
        velbar_mag.mask_by(vecs.thk, PISM.convert(sys, -0.01, "m/year", "m/second"))
        velbar_mag.write(filename)

    def _setFromOptions(self):
        """Optionally override to set any data from command line variables."""
        pass

    def _constructModelData(self):
        """Optionally override to return a custom :class:`PISM.model.ModelData` instance."""
        return model.ModelData(self.grid)

    def _initGrid(self):
        """Override to return the computation grid."""
        raise NotImplementedError()

    def _initPhysics(self):
        """Override to set the non-var parts  of :attr:`modeldata` (e.g. the basal yeild stress model and the enthalpy converter)"""
        raise NotImplementedError()

    def _allocStdSSACoefficients(self):
        """Helper method that allocates the standard :cpp:class:`IceModelVec` variables used to solve the SSA and stores them
        in :attr:`modeldata```.vecs``:

          * ``surface``
          * ``thickness``
          * ``bed``
          * ``tauc``
          * ``enthalpy``
          * ``mask``
          * ``age`` if -age is given

        Intended to be called from custom implementations of :meth:`_initSSACoefficients` if desired."""
        vecs = self.modeldata.vecs
        grid = self.grid

        vecs.add(model.createIceSurfaceVec(grid))
        vecs.add(model.createIceThicknessVec(grid))
        vecs.add(model.createBedrockElevationVec(grid))
        vecs.add(model.createYieldStressVec(grid), 'tauc')
        vecs.add(model.createEnthalpyVec(grid), 'enthalpy')
        vecs.add(model.createIceMaskVec(grid), 'mask')

        # The SIA model might need the "age" field
        if grid.ctx().config().get_boolean("do_age"):
            vecs.add(model.createAgeVec(grid), "age")

    def _allocateBCs(self, velname='_bc', maskname='bc_mask'):
        """Helper method that allocates standard Dirichlet data
            :cpp:class:`IceModelVec` variable and stores them in
            :attr:`modeldata` ``.vecs``:

          * ``vel_bc``
          * ``bc_mask``

        """
        vecs = self.modeldata.vecs
        vecs.add(model.create2dVelocityVec(self.grid,
                                           name=velname,
                                           desc='SSA velocity boundary condition',
                                           intent='intent'),
                 "vel_bc")
        vecs.add(model.createBCMaskVec(self.grid, name=maskname),
                 "bc_mask")

    def _initSSACoefficients(self):
        """Override to allocate and initialize all :cpp:class:`IceModelVec` variables in :attr:`modeldata` ``.vecs``
           needed for solving the SSA."""
        raise NotImplementedError()

    def _constructSSA(self):
        """Optionally override to return an instance of :cpp:class:`SSA` (e.g. :cpp:class:`SSAFD` or :cpp:class:`SSAFEM`)
           that will be used for solving the SSA."""
        md = self.modeldata
        return SSAAlgorithms[md.config.get_string("ssa_method")](md.grid, md.enthalpyconverter)

    def _initSSA(self):
        """Optionally perform any final initialization of :attr:`ssa`."""
        pass


class SSAExactTestCase(SSARun):

    """Base class for implmentation of specific SSA test cases.  Provides a mechanism for comparing
    computed and exact values.  Simply construct with a grid size and then call :meth:`run`"""

    def __init__(self, Mx, My):
        """Initialize with a grid of the specified size."""
        SSARun.__init__(self)
        self.Mx = Mx
        self.My = My

        # For convenience, provide a grid. It will get initialized later
        # on when _initGrid is called by our setup method.
        self.grid = None

    def run(self, output_file):
        """Main command intended to be called by whatever code executes the test case.
        Calls :meth:`setup`, :meth:`solve`, :meth:`report`, and :meth:`write`."""
        self.setup()
        self.solve()
        self.report()
        self.write(output_file)

    def report(self):
        """Compares computed and exact solution values and displays a summary report."""
        grid = self.grid

        ssa_stdout = self.ssa.stdout_report()
        PISM.verbPrintf(3, grid.com, ssa_stdout)

        maxvecerr = 0.0
        avvecerr = 0.0
        avuerr = 0.0
        avverr = 0.0
        maxuerr = 0.0
        maxverr = 0.0

        if (self.config.get_boolean("do_pseudo_plastic_till") and
                self.config.get_double("pseudo_plastic_q") != 1.0):
            PISM.verbPrintf(1, grid.com, "WARNING: numerical errors not valid for pseudo-plastic till\n")
        PISM.verbPrintf(1, grid.com, "NUMERICAL ERRORS in velocity relative to exact solution:\n")

        vel_ssa = self.ssa.velocity()

        vel_ssa.begin_access()

        exactvelmax = 0
        gexactvelmax = 0
        for (i, j) in self.grid.points():
            x = grid.x(i)
            y = grid.y(j)
            (uexact, vexact) = self.exactSolution(i, j, x, y)
            exactnormsq = math.sqrt(uexact * uexact + vexact * vexact)
            exactvelmax = max(exactnormsq, exactvelmax)
            solution = vel_ssa[i, j]
            uerr = abs(solution.u - uexact)
            verr = abs(solution.v - vexact)
            avuerr += uerr
            avverr += verr
            maxuerr = max(maxuerr, uerr)
            maxverr = max(maxverr, verr)
            vecerr = math.sqrt(uerr * uerr + verr * verr)
            maxvecerr = max(maxvecerr, vecerr)
            avvecerr = avvecerr + vecerr

        vel_ssa.end_access()

        N = grid.Mx() * grid.My()
        gexactvelmax = PISM.GlobalMax(grid.com, exactvelmax)
        gmaxuerr = PISM.GlobalMax(grid.com, maxuerr)
        gmaxverr = PISM.GlobalMax(grid.com, maxverr)
        gavuerr = PISM.GlobalSum(grid.com, avuerr) / N
        gavverr = PISM.GlobalSum(grid.com, avverr) / N
        gmaxvecerr = PISM.GlobalMax(grid.com, maxvecerr)
        gavvecerr = PISM.GlobalSum(grid.com, avvecerr) / N

        sys = grid.ctx().unit_system()

        m_year = PISM.UnitConverter(sys, "m / second", "m / year")

        if abs(gexactvelmax) > 0.0:
            relative_vel_error = (gavvecerr / gexactvelmax) * 100.0
        else:
            relative_vel_error = 0.0

        PISM.verbPrintf(1, grid.com, "velocity  :  maxvector   prcntavvec      maxu      maxv       avu       avv\n")
        PISM.verbPrintf(1, grid.com,
                        "           %11.4f%13.5f%10.4f%10.4f%10.4f%10.4f\n",
                        m_year(gmaxvecerr),
                        relative_vel_error,
                        m_year(gmaxuerr),
                        m_year(gmaxverr),
                        m_year(gavuerr),
                        m_year(gavverr))
        PISM.verbPrintf(1, grid.com, "NUM ERRORS DONE\n")

    def exactSolution(self, i, j, xi, xj):
        """Override to provide the exact value of the solution at grid index (``i``, ``j``) with
        coordinates (``xi``, ``xj``)."""
        raise NotImplementedError()

    def write(self, filename):
        """Override of :meth:`SSARun.write`.  Does all of the above, and saves a copy of the exact solution."""
        SSARun.write(self, filename)

        grid = self.grid
        exact = model.create2dVelocityVec(grid, name="_exact", desc="SSA exact solution", intent="diagnostic")
        exact.begin_access()
        for (i, j) in grid.points():
            exact[i, j] = self.exactSolution(i, j, grid.x(i), grid.y(j))
        exact.end_access()
        exact.write(filename)


class SSAFromInputFile(SSARun):

    """Class for running the SSA based on data provided in an input file."""

    def __init__(self, boot_file):
        SSARun.__init__(self)
        self.grid = None
        self.config = PISM.Context().config
        self.boot_file = boot_file
        self.phi_to_tauc = False
        self.is_regional = False

    def _setFromOptions(self):
        self.phi_to_tauc = PISM.OptionBool("-phi_to_tauc",
                                           "Recompute pseudo yield stresses from till friction angles.")
        self.is_regional = PISM.OptionBool("-regional", "enable 'regional' mode")

    def _initGrid(self):
        """Override of :meth:`SSARun._initGrid`. Sets periodicity based on
        ``-periodicity`` command-line option."""
        # FIXME: allow specification of Mx and My different from what's
        # in the boot_file.
        periodicity = PISM.XY_PERIODIC
        (pstring, pflag) = PISM.optionsListWasSet('-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:
            if self.is_regional and (self.config.get_string("ssa_method") == "fem"):
                periodicity = PISM.NOT_PERIODIC

        self.grid = PISM.IceGrid.FromFile(PISM.Context().ctx, self.boot_file, "enthalpy",
                                          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.OptionString("-ssa_glen", "SSA flow law Glen exponent").is_set():
            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 _allocExtraSSACoefficients(self):
        """Allocate storage for SSA coefficients."""
        vecs = self.modeldata.vecs
        if util.fileHasVariable(self.boot_file, 'ssa_driving_stress_x'):
            vecs.add(model.createDrivingStressXVec(self.grid))

        if util.fileHasVariable(self.boot_file, 'ssa_driving_stress_y'):
            vecs.add(model.createDrivingStressYVec(self.grid))

        no_model_mask = None
        # For a regional run we'll need no_model_mask, usurfstore, thkstore
        if self.is_regional:
            no_model_mask = model.createNoModelMaskVec(self.grid)
            vecs.add(no_model_mask, 'no_model_mask')
            vecs.add(model.createIceSurfaceStoreVec(self.grid))
            vecs.add(model.createIceThicknessStoreVec(self.grid))

        if self.config.get_boolean('ssa_dirichlet_bc'):
            vecs.add(model.create2dVelocityVec(self.grid, name='_ssa_bc',
                                               desc='SSA velocity boundary condition',
                                               intent='intent'),
                     "vel_ssa_bc")

            if self.is_regional:
                vecs.add(no_model_mask, 'bc_mask')
            else:
                vecs.add(model.createBCMaskVec(self.grid), 'bc_mask')

        if self.phi_to_tauc:
            vecs.add(PISM.model.createBasalMeltRateVec(self.grid))
            vecs.add(PISM.model.createTillPhiVec(self.grid))
            vecs.add(PISM.model.createBasalWaterVec(self.grid))

    def _initSSACoefficients(self):
        """Override of :meth:`SSARun._initSSACoefficients` that initializes variables from the
        contents of the input file."""
        # Build the standard thickness, bed, etc
        self._allocStdSSACoefficients()
        self._allocExtraSSACoefficients()

        vecs = self.modeldata.vecs

        thickness = vecs.land_ice_thickness
        bed = vecs.bedrock_altitude
        enthalpy = vecs.enthalpy
        mask = vecs.mask
        surface = vecs.surface_altitude

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

        # The SIA model might need the age field.
        if self.config.get_boolean("do_age"):
            vecs.age.regrid(self.boot_file, 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)

        if util.fileHasVariable(self.boot_file, 'ssa_driving_stress_x'):
            vecs.ssa_driving_stress_x.regrid(self.boot_file, critical=True)

        if util.fileHasVariable(self.boot_file, 'ssa_driving_stress_y'):
            vecs.ssa_driving_stress_y.regrid(self.boot_file, critical=True)

        # For a regional run we'll need no_model_mask, usurfstore, thkstore
        if self.is_regional:
            vecs.no_model_mask.regrid(self.boot_file, True)

            if util.fileHasVariable(self.boot_file, 'usurfstore'):
                vecs.usurfstore.regrid(self.boot_file, True)
            else:
                vecs.usurfstore.copy_from(vecs.surface_altitude)

            if util.fileHasVariable(self.boot_file, 'thkstore'):
                vecs.thkstore.regrid(self.boot_file, True)
            else:
                vecs.thkstore.copy_from(vecs.land_ice_thickness)

        # Compute yield stress from PISM state variables
        # (basal melt rate, tillphi, and basal water height)
        grid = self.grid

        if self.phi_to_tauc:
            for v in [vecs.bmr, vecs.tillphi, vecs.bwat]:
                v.regrid(self.boot_file, True)
                vecs.add(v)

            if self.is_regional:
                yieldstress = PISM.RegionalDefaultYieldStress(self.modeldata.grid)
            else:
                yieldstress = PISM.MohrCoulombYieldStress(self.modeldata.grid)

            # make sure vecs is locked!
            yieldstress.init()
            yieldstress.set_till_friction_angle(vecs.tillphi)
            yieldstress.update(0, 1)
            vecs.tauc.copy_from(yieldstress.basal_material_yield_stress())
        else:
            vecs.tauc.regrid(self.boot_file, True)

        if self.config.get_boolean('ssa_dirichlet_bc'):
            has_u_ssa_bc = util.fileHasVariable(self.boot_file, 'u_ssa_bc')
            has_v_ssa_bc = util.fileHasVariable(self.boot_file, 'v_ssa_bc')

            if (not has_u_ssa_bc) or (not has_v_ssa_bc):
                PISM.verbPrintf(2, grid.com,
                                "Input file '%s' missing Dirichlet boundary data u/v_ssa_bc;"
                                " using zero default instead." % self.boot_file)
                vecs.vel_ssa_bc.set(0.0)
            else:
                vecs.vel_ssa_bc.regrid(self.boot_file, True)

            if not self.is_regional:
                bc_mask_name = vecs.bc_mask.metadata().get_string("short_name")
                if util.fileHasVariable(self.boot_file, bc_mask_name):
                    vecs.bc_mask.regrid(self.boot_file, True)
                else:
                    PISM.verbPrintf(2, grid.com,
                                    "Input file '%s' missing Dirichlet location mask '%s'."
                                    "  Default to no Dirichlet locations." % (self.boot_file, bc_mask_name))
                    vecs.bc_mask.set(0)

    def _constructSSA(self):
        """Constructs an instance of :cpp:class:`SSA` for solving the SSA based on command-line flags ``-regional`` and ``-ssa_method``"""
        md = self.modeldata
        if self.is_regional and (md.config.get_string("ssa_method") == "fd"):
            algorithm = PISM.SSAFD_Regional
        else:
            algorithm = SSAAlgorithms[md.config.get_string("ssa_method")]
        return algorithm(md.grid, md.enthalpyconverter)
