from dolfin import *
from rbnics import *

"""
Explanation: TUTORIAL 05 - Discrete Empirical Interpolation Method for non-affine elliptic problems
Keywords: discrete empirical interpolation method
1. Introduction
In this Tutorial, we consider steady heat conduction in a two-dimensional square domain $\Omega = (-1, 1)^2$.
The boundary $\partial\Omega$ is kept at a reference temperature (say, zero). The conductivity coefficient is fixed to 1, while the heat source is characterized by the following expression
$$
g(\boldsymbol{x}; \boldsymbol{\mu}) = \exp{ -2 (x_0-\mu_0)^2 - 2 (x_1 - \mu_1)^2} \quad \forall \boldsymbol{x} = (x_0, x_1) \in \Omega.
$$
The parameter vector $\boldsymbol{\mu}$, given by 
$$
\boldsymbol{\mu} = (\mu_0,\mu_1)
$$
affects the center of the Gaussian source $g(\boldsymbol{x}; \boldsymbol{\mu})$, which could be located at any point $\Omega$. Thus, the parameter domain is
$$
\mathbb{P}=[-1,1]^2.
$$
In order to obtain a faster evaluation (yet, provably accurate) of the problem we propose to use a certified reduced basis approximation for the problem. In order to preserve the affinity assumption (for the sake of performance) the discrete empirical interpolation method will be used on the forcing term $g(\boldsymbol{x}; \boldsymbol{\mu})$.
2. Parametrized formulation
Let $u(\boldsymbol{\mu})$ be the temperature in the domain $\Omega$.
We will directly provide a weak formulation for this problem
<center>for a given parameter $\boldsymbol{\mu}\in\mathbb{P}$, find $u(\boldsymbol{\mu})\in\mathbb{V}$ such that</center>
$$a\left(u(\boldsymbol{\mu}),v;\boldsymbol{\mu}\right)=f(v;\boldsymbol{\mu})\quad \forall v\in\mathbb{V}$$
where

the function space $\mathbb{V}$ is defined as
$$
\mathbb{V} = \left{ v \in H^1(\Omega(\mu_0)): v|_{\partial\Omega} = 0\right}
$$
Note that, as in the previous tutorial, the function space is parameter dependent due to the shape variation. 
the parametrized bilinear form $a(\cdot, \cdot; \boldsymbol{\mu}): \mathbb{V} \times \mathbb{V} \to \mathbb{R}$ is defined by
$$a(u,v;\boldsymbol{\mu}) = \int_{\Omega} \nabla u \cdot \nabla v \ d\boldsymbol{x}$$
the parametrized linear form $f(\cdot; \boldsymbol{\mu}): \mathbb{V} \to \mathbb{R}$ is defined by
$$f(v;\boldsymbol{\mu}) = \int_\Omega g(\boldsymbol{\mu}) v  \ d\boldsymbol{x}.$$
End of explanation
"""


@DEIM()
class Gaussian(EllipticCoerciveProblem):

    # Default initialization of members
    def __init__(self, V, **kwargs):
        # Call the standard initialization
        EllipticCoerciveProblem.__init__(self, V, **kwargs)
        # ... and also store FEniCS data structures for assembly
        assert "subdomains" in kwargs
        assert "boundaries" in kwargs
        self.subdomains, self.boundaries = kwargs["subdomains"], kwargs["boundaries"]
        self.u = TrialFunction(V)
        self.v = TestFunction(V)
        self.dx = Measure("dx")(subdomain_data=subdomains)
        self.f = ParametrizedExpression(
            self, "exp(- 2 * pow(x[0] - mu[0], 2) - 2 * pow(x[1] - mu[1], 2))", mu=(0., 0.),
            element=V.ufl_element())
        # note that we cannot use self.mu in the initialization of self.f, because self.mu has not been initialized yet

    # Return custom problem name
    def name(self):
        return "GaussianDEIM"

    # Return the alpha_lower bound.
    def get_stability_factor_lower_bound(self):
        return 1.

    # Return theta multiplicative terms of the affine expansion of the problem.
    def compute_theta(self, term):
        if term == "a":
            return (1.,)
        elif term == "f":
            return (1.,)
        else:
            raise ValueError("Invalid term for compute_theta().")

    # Return forms resulting from the discretization of the affine expansion of the problem operators.
    def assemble_operator(self, term):
        v = self.v
        dx = self.dx
        if term == "a":
            u = self.u
            a0 = inner(grad(u), grad(v)) * dx
            return (a0,)
        elif term == "f":
            f = self.f
            f0 = f * v * dx
            return (f0,)
        elif term == "dirichlet_bc":
            bc0 = [DirichletBC(self.V, Constant(0.0), self.boundaries, 1),
                   DirichletBC(self.V, Constant(0.0), self.boundaries, 2),
                   DirichletBC(self.V, Constant(0.0), self.boundaries, 3)]
            return (bc0,)
        elif term == "inner_product":
            u = self.u
            x0 = inner(grad(u), grad(v)) * dx
            return (x0,)
        else:
            raise ValueError("Invalid term for assemble_operator().")

"""
Explanation: 3. Affine decomposition
The parametrized bilinear form $a(\cdot, \cdot; \boldsymbol{\mu})$ is trivially affine.
The discrete empirical interpolation method will be used on the forcing term $g(\boldsymbol{x}; \boldsymbol{\mu})$ to obtain an efficient (approximately affine) expansion of $f(\cdot; \boldsymbol{\mu})$.
End of explanation
"""


mesh = Mesh("data/gaussian.xml")
subdomains = MeshFunction("size_t", mesh, "data/gaussian_physical_region.xml")
boundaries = MeshFunction("size_t", mesh, "data/gaussian_facet_region.xml")

"""
Explanation: 4. Main program
4.1. Read the mesh for this problem
The mesh was generated by the data/generate_mesh.ipynb notebook.
End of explanation
"""


V = FunctionSpace(mesh, "Lagrange", 1)

"""
Explanation: 4.2. Create Finite Element space (Lagrange P1)
End of explanation
"""


problem = Gaussian(V, subdomains=subdomains, boundaries=boundaries)
mu_range = [(-1.0, 1.0), (-1.0, 1.0)]
problem.set_mu_range(mu_range)

"""
Explanation: 4.3. Allocate an object of the Gaussian class
End of explanation
"""


reduction_method = ReducedBasis(problem)
reduction_method.set_Nmax(20, DEIM=21)
reduction_method.set_tolerance(1e-4, DEIM=1e-8)

"""
Explanation: 4.4. Prepare reduction with a reduced basis method
End of explanation
"""


reduction_method.initialize_training_set(50, DEIM=60)
reduced_problem = reduction_method.offline()

"""
Explanation: 4.5. Perform the offline phase
End of explanation
"""


online_mu = (0.3, -1.0)
reduced_problem.set_mu(online_mu)
reduced_solution = reduced_problem.solve()
plot(reduced_solution, reduced_problem=reduced_problem)

"""
Explanation: 4.6.1. Perform an online solve
End of explanation
"""


reduced_solution_11 = reduced_problem.solve(DEIM=11)
plot(reduced_solution_11, reduced_problem=reduced_problem)

"""
Explanation: 4.6.2. Perform an online solve with a lower number of DEIM terms
End of explanation
"""


reduced_solution_1 = reduced_problem.solve(DEIM=1)
plot(reduced_solution_1, reduced_problem=reduced_problem)

"""
Explanation: 4.6.3. Perform an online solve with an even lower number of DEIM terms
End of explanation
"""


reduction_method.initialize_testing_set(50, DEIM=60)
reduction_method.error_analysis(filename="error_analysis")

"""
Explanation: 4.7.1. Perform an error analysis
End of explanation
"""


reduction_method.error_analysis(
    with_respect_to=exact_problem, filename="error_analysis__with_respect_to_exact")

"""
Explanation: 4.7.2. Perform an error analysis with respect to the exact problem
End of explanation
"""


reduction_method.error_analysis(
    with_respect_to=exact_problem, DEIM=11, filename="error_analysis__with_respect_to_exact__DEIM_11")

"""
Explanation: 4.7.3. Perform an error analysis with respect to the exact problem, but employing a smaller number of DEIM terms
End of explanation
"""


