"""This module defines different MultiStageScheme classes which can be passed to a RKSolver"""

# Copyright (C) 2013 Johan Hake
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN 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 Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
#
# First added:  2013-02-22
# Last changed: 2013-02-22

import numpy as np

# Import SWIG-generated extension module (DOLFIN C++)
import dolfin.cpp as cpp

# Import ufl
import ufl

# Import classes from dolfin python layer
from dolfin.functions.constant import Constant
from dolfin.fem.formmanipulations import derivative
from dolfin.fem.form import Form

# FIXME: Add support for algebraic parts (at least for implicit)
# FIXME: Add support for implicit/explicit split ala IMEX schemes
def _butcher_scheme_generator(a, b, c, solution, rhs_form):
    """
    Generates a list of forms and solutions for a given Butcher tableau 

    *Arguments*
        a (2 dimensional numpy array)
            The a matrix of the Butcher tableau.
        b (1-2 dimensional numpy array)
            The b vector of the Butcher tableau. If b is 2 dimensional the
            scheme includes an error estimator and can be used in adaptive
            solvers.
        c (1 dimensional numpy array)
            The c vector the Butcher tableau.
        solution (_Function_)
            The prognastic variable
        rhs_form (ufl.Form)
            A UFL form representing the rhs for a time differentiated equation 
    """
    if not (isinstance(a, np.ndarray) and (len(a) == 1 or \
            (len(a.shape)==2 and a.shape[0] == a.shape[1]))):
        raise TypeError("Expected an m x m numpy array as the first argument")
    if not (isinstance(b, np.ndarray) and len(b.shape) in [1,2]):
        raise TypeError("Expected a 1 or 2 dimensional numpy array as the second argument")
    if not (isinstance(c, np.ndarray) and len(c.shape) == 1):
        raise TypeError("Expected a 1 dimensional numpy array as the third argument")

    # Make sure a is a "matrix"
    if len(a) == 1:
        a.shape = (1, 1)

    # Get size of system
    size = a.shape[0]

    # If b is a matrix we expect it to have two rows 
    if len(b.shape) == 2:
        if not (b.shape[0] == 2 and b.shape[1] == size):
            raise ValueError("Expected a 2 row matrix with the same number "\
                             "of collumns as the first dimension of the a matrix.")
    elif len(b) != size:
        raise ValueError("Expected the length of the b vector to have the "\
                         "same size as the first dimension of the a matrix.")
        
    if len(c) != size:
        raise ValueError("Expected the length of the c vector to have the "\
                         "same size as the first dimension of the a matrix.")

    # Check if tableau is fully implicit
    for i in range(size):
        for j in range(i):
            if a[j, i] != 0:
                raise ValueError("Does not support fully implicit Butcher tableau.")

    if not isinstance(rhs_form, ufl.Form):
        raise TypeError("Expected a ufl.Form as the 5th argument.")
        
    # Check if form contains a cell or point integral
    if "cell" in rhs_form.integral_groups():
        DX = ufl.dx
    elif "point" in rhs_form.integral_groups():
        DX = ufl.dP
    else:
        raise ValueError("Expected either a cell or point integral in the form.")
    
    # Get test function
    arguments, coefficients = ufl.algorithms.extract_arguments_and_coefficients(rhs_form)
    if len(arguments) != 1:
        raise ValueError("Expected the form to have rank 1")
    v = arguments[0]

    # Create time step
    dt = Constant(0.1)

    # rhs forms
    dolfin_stage_forms = []
    ufl_stage_forms = []
    
    # Stage solutions
    k = [solution.copy(deepcopy=True) for i in range(size)]

    # Create the stage forms
    y_ = solution
    for i, ki in enumerate(k):

        # Check wether the stage is explicit
        explicit = a[i,i] == 0

        # Evaluation arguments for the ith stage
        evalargs = y_ + dt * sum([float(a[i,j]) * k[j] \
                                  for j in range(i+1)], ufl.zero(*y_.shape()))
        
        stage_form = ufl.replace(rhs_form, {y_:evalargs})

        if explicit:
            stage_forms = [stage_form]
        else:
            # Create a F=0 form and differentiate it
            stage_form -= ufl.inner(ki, v)*DX
            stage_forms = [stage_form, derivative(stage_form, ki)]
        ufl_stage_forms.append(stage_forms)

        dolfin_stage_forms.append([Form(form) for form in stage_forms])
        
    # Only one last stage
    if len(b.shape) == 1:
        last_stage = cpp.FunctionAXPY([(float(bi), ki) for bi, ki in zip(b, k)])
    else:
        # FIXME: Add support for addaptivity in RKSolver and MultiStageScheme
        last_stage = [cpp.FunctionAXPY([(float(bi), ki) for bi, ki in zip(b[0,:], k)]),
                      cpp.FunctionAXPY([(float(bi), ki) for bi, ki in zip(b[1,:], k)])]

    # Create the Function holding the solution at end of time step
    #k.append(solution.copy())

    # Generate human form of MultiStageScheme
    human_form = []
    for i in range(size):
        kterm = " + ".join("%sh*k_%s" % ("" if a[i,j] == 1.0 else \
                                         "%s*"% a[i,j], j) \
                           for j in range(size) if a[i,j] != 0)
        if c[i] in [0.0, 1.0]:
            cih = " + h" if c[i] == 1.0 else ""
        else:
            cih = " + %s*h" % c[i]
            
        if len(kterm) == 0:
            human_form.append("k_%(i)s = f(t_n%(cih)s, y_n)" % {"i": i, "cih": cih})
        else:
            human_form.append("k_%(i)s = f(t_n%(cih)s, y_n + %(kterm)s)" % \
                          {"i": i, "cih": cih, "kterm": kterm})

    parentheses = "(%s)" if np.sum(b>0) > 1 else "%s"
    human_form.append("y_{n+1} = y_n + h*" + parentheses % (" + ".join(\
        "%sk_%s" % ("" if b[i] == 1.0 else "%s*" % b[i], i) \
        for i in range(size) if b[i] > 0)))

    human_form = "\n".join(human_form)
    
    return ufl_stage_forms, dolfin_stage_forms, last_stage, k, dt, human_form

class MultiStageScheme(cpp.MultiStageScheme):
    """
    Base class for all MultiStageSchemes
    """
    def __init__(self, rhs_form, solution, t, bcs, a, b, c, order):
        bcs = bcs or []
        t = t or Constant(0.0)
        ufl_stage_forms, dolfin_stage_forms, last_stage, k, dt, human_form = \
                         _butcher_scheme_generator(a, b, c, solution, rhs_form)

        # Store data
        self._rhs_form = rhs_form
        self._ufl_stage_forms = ufl_stage_forms
        self._dolfin_stage_forms = dolfin_stage_forms
        self._t = t
        self._dt = dt
        self._last_stage = last_stage
        self._solution = solution
        self._k = k
        self.a = a
        self.b = b
        self.c = c

        cpp.MultiStageScheme.__init__(self, dolfin_stage_forms, last_stage, k, \
                                      solution, t, dt, c, order,
                                      self.__class__.__name__,
                                      human_form, bcs)
        
    def rhs_form(self):
        "Return the original rhs form"
        return self._rhs_form
        
    def ufl_stage_forms(self):
        "Return the ufl stage forms"
        return self._ufl_stage_forms
        
    def dolfin_stage_forms(self):
        "Return the dolfin stage forms"
        return self._dolfin_stage_forms

    def t(self):
        "Return the Constant used to describe time in the MultiStageScheme"
        return self._t
        
    def dt(self):
        "Return the Constant used to describe time in the MultiStageScheme"
        return self._dt

    def solution(self):
        "Return the solution Function"
        return self._solution

    def last_stage(self):
        "Return the AXPYFunction object describing the last stage"
        return self._last_stage

    def stage_solutions(self):
        "Return the stage solutions"
        return self._stage_solutions
        
class ERK1(MultiStageScheme):
    """
    Explicit first order Scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):
        a = np.array([0.])
        b = np.array([1.])
        c = np.array([0.])
        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 1)

class BDF1(MultiStageScheme):
    """
    Implicit first order scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):
        a = np.array([1.])
        b = np.array([1.])
        c = np.array([1.])
        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 1)

class ExplicitMidPoint(MultiStageScheme):
    """
    Explicit 2nd order scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):

        a = np.array([[0, 0],[0.5, 0.0]])
        b = np.array([0., 1])
        c = np.array([0, 0.5])
        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 2)

class CN2(MultiStageScheme):
    """
    Semi-implicit 2nd order scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):
        a = np.array([[0, 0],[0.5, 0.5]])
        b = np.array([0.5, 0.5])
        c = np.array([0, 1.0])

        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 2)

class ERK4(MultiStageScheme):
    """
    Explicit 4th order scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):
        a = np.array([[0, 0, 0, 0],
                      [0.5, 0, 0, 0],
                      [0, 0.5, 0, 0],
                      [0, 0, 1, 0]])
        b = np.array([1./6, 1./3, 1./3, 1./6])
        c = np.array([0, 0.5, 0.5, 1])
        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 4)

class ESDIRK3(MultiStageScheme):
    """
    Explicit implicit 3rd order scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):
        a = np.array([[ 0.        ,  0.        ,  0.        ,  0.        ,  0.        ],
                      [ 0.43586652,  0.43586652,  0.        ,  0.        ,  0.        ],
                      [ 0.14073777, -0.10836555,  0.43586652,  0.        ,  0.        ],
                      [ 0.1023994 , -0.37687845,  0.83861253,  0.43586652,  0.        ],
                      [ 0.1570249 ,  0.11733044,  0.61667803, -0.32689989,  0.43586652]])
        b = a[-1,:].copy()
        c = a.sum(1)
        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 3)

class ESDIRK4(MultiStageScheme):
    """
    Explicit implicit 4rd order scheme
    """
    def __init__(self, rhs_form, solution, t=None, bcs=None):
        a = np.array([[0,                   0,                   0,                   0,                   0                   ],
                      [0.435866521500000,   0.435866521500000,   0,                   0,                   0                   ],
                      [0.140737774731968,  -0.108365551378832,   0.435866521500000,   0,                   0                   ],
                      [0.102399400616089,  -0.376878452267324,   0.838612530151233,   0.435866521500000,   0                   ],
                      [0.157024897860995,   0.117330441357768,   0.616678030391680,  -0.326899891110444,   0.435866521500000   ]])

        b = a[-1,:].copy()
        c = a.sum(1)
        MultiStageScheme.__init__(self, rhs_form, solution, t, bcs, a, b, c, 4)

# Aliases
CrankNicolson = CN2
ExplicitEuler = ERK1
ForwardEuler = ERK1
ImplicitEuler = BDF1
BackwardEuler = BDF1
ERK = ERK1
RK4 = ERK4

__all__ = [name for name, attr in globals().items() \
           if isinstance(attr, type) and issubclass(attr, MultiStageScheme)]

__all__.append("MultiStageScheme")
