from dolfin import *
from rbnics import *

"""
Explanation: TUTORIAL 06 - Unsteady Thermal block problem
Keywords: certified reduced basis method, scalar problem
1. Introduction
In this Tutorial, we consider unsteady heat conduction in a two-dimensional domain $\Omega$.
<img src="data/thermal_block.png" />
We define two subdomains $\Omega_1$ and $\Omega_2$, such that
1. $\Omega_1$ is a disk centered at the origin of radius $r_0=0.5$, and
2. $\Omega_2=\Omega/\ \overline{\Omega_1}$. 
The conductivity $\kappa$ is assumed to be constant on $\Omega_1$ and $\Omega_2$, i.e.
$$
\kappa|{\Omega_1}=\kappa_0 \quad \textrm{and} \quad \kappa|{\Omega_2}=1.
$$
For this problem, we consider $P=2$ parameters:
1. the first one is related to the conductivity in $\Omega_1$, i.e. $\mu_0\equiv\kappa_0$ (note that parameters numbering is zero-based);
2. the second parameter $\mu_1$ takes into account the constant heat flux over $\Gamma_{base}$.
The parameter vector $\boldsymbol{\mu}$ is thus given by 
$$
\boldsymbol{\mu} = (\mu_0,\mu_1)
$$
on the parameter domain
$$
\mathbb{P}=[0.1,10]\times[-1,1].
$$
In this problem we model the heat transfer process due to the heat flux over the bottom boundary $\Gamma_{base}$ and the following conditions on the remaining boundaries:
* the left and right boundaries $\Gamma_{side}$ are insulated,
* the top boundary $\Gamma_{top}$ is kept at a reference temperature (say, zero),
with the aim of measuring the average temperature on $\Gamma_{base}$.
In order to obtain a faster evaluation (yet, provably accurate) of the output of interest we propose to use a certified reduced basis approximation for the problem.
2. Parametrized formulation
Let $u(t;\boldsymbol{\mu})$ be the temperature in the domain $\Omega\times[0,t_f]$.
The strong formulation of the parametrized problem is given by:
<center>for a given parameter $\boldsymbol{\mu}\in\mathbb{P}$, for $t\in[0,t_f]$, find $u(t;\boldsymbol{\mu})$ such that</center>
$$
\begin{cases}
    \partial_tu(t;\boldsymbol{\mu})- \text{div} (\kappa(\mu_0)\nabla u(t;\boldsymbol{\mu})) = 0 & \text{in } \Omega\times[0,t_f],\
    u(t=0;\boldsymbol{\mu}) = 0 & \text{in } \Omega, \ 
    u(t;\boldsymbol{\mu}) = 0 & \text{on } \Gamma_{top}\times[0,t_f],\
    \kappa(\mu_0)\nabla u(t;\boldsymbol{\mu})\cdot \mathbf{n} = 0 & \text{on } \Gamma_{side}\times[0,t_f],\
    \kappa(\mu_0)\nabla u(t;\boldsymbol{\mu})\cdot \mathbf{n} = \mu_1 & \text{on } \Gamma_{base}\times[0,t_f].
\end{cases}
$$
<br>
where 
* $\mathbf{n}$ denotes the outer normal to the boundaries $\Gamma_{side}$ and $\Gamma_{base}$,
* the conductivity $\kappa(\mu_0)$ is defined as follows:
$$
\kappa(\mu_0) =
\begin{cases}
    \mu_0 & \text{in } \Omega_1,\
    1 & \text{in } \Omega_2,\
\end{cases}
$$
The corresponding weak formulation reads:
<center>for a given parameter $\boldsymbol{\mu}\in\mathbb{P}$, for $t\in[0,t_f]$, find $u(t;\boldsymbol{\mu})\in\mathbb{V}$ such that</center>
$$m\left(\partial_tu(t;\boldsymbol{\mu}),v;\boldsymbol{\mu}\right) + a\left(u(t;\boldsymbol{\mu}),v;\boldsymbol{\mu}\right)=f(v;\boldsymbol{\mu})\quad \forall v\in\mathbb{V},\quad \forall t\in[0,t_f]$$
where

the function space $\mathbb{V}$ is defined as
$$
\mathbb{V} = {v\in H^1(\Omega) : v|{\Gamma{top}}=0}
$$
the parametrized bilinear form $m(\cdot, \cdot; \boldsymbol{\mu}): \mathbb{V} \times \mathbb{V} \to \mathbb{R}$ is defined by
$$m(u, v;\boldsymbol{\mu})=\int_{\Omega} \partial_tu(t)v \ d\boldsymbol{x},$$
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} \kappa(\mu_0)\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})= \mu_1\int_{\Gamma_{base}}v \ ds,$$

The (compliant) output of interest $s(t;\boldsymbol{\mu})$ is given by
$$s(t;\boldsymbol{\mu}) = \mu_1\int_{\Gamma_{base}} u(t;\boldsymbol{\mu})$$
is computed for each $\boldsymbol{\mu}$.
End of explanation
"""


class UnsteadyThermalBlock(ParabolicCoerciveProblem):

    # Default initialization of members
    def __init__(self, V, **kwargs):
        # Call the standard initialization
        ParabolicCoerciveProblem.__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=self.subdomains)
        self.ds = Measure("ds")(subdomain_data=self.boundaries)

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

    # Return the alpha_lower bound.
    def get_stability_factor_lower_bound(self):
        return min(self.compute_theta("a"))

    # Return theta multiplicative terms of the affine expansion of the problem.
    def compute_theta(self, term):
        mu = self.mu
        if term == "m":
            theta_m0 = 1.
            return (theta_m0, )
        elif term == "a":
            theta_a0 = mu[0]
            theta_a1 = 1.
            return (theta_a0, theta_a1)
        elif term == "f":
            theta_f0 = mu[1]
            return (theta_f0,)
        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 == "m":
            u = self.u
            m0 = u * v * dx
            return (m0, )
        elif term == "a":
            u = self.u
            a0 = inner(grad(u), grad(v)) * dx(1)
            a1 = inner(grad(u), grad(v)) * dx(2)
            return (a0, a1)
        elif term == "f":
            ds = self.ds
            f0 = v * ds(1)
            return (f0,)
        elif term == "dirichlet_bc":
            bc0 = [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,)
        elif term == "projection_inner_product":
            u = self.u
            x0 = u * v * dx
            return (x0,)
        else:
            raise ValueError("Invalid term for assemble_operator().")

"""
Explanation: 3. Affine decomposition
For this problem the affine decomposition is straightforward:
$$m(u,v;\boldsymbol{\mu})=\underbrace{1}{\Theta^{m}_0(\boldsymbol{\mu})}\underbrace{\int{\Omega}uv \ d\boldsymbol{x}}{m_0(u,v)},$$
$$a(u,v;\boldsymbol{\mu})=\underbrace{\mu_0}{\Theta^{a}0(\boldsymbol{\mu})}\underbrace{\int{\Omega_1}\nabla u \cdot \nabla v \ d\boldsymbol{x}}{a_0(u,v)} \ + \  \underbrace{1}{\Theta^{a}1(\boldsymbol{\mu})}\underbrace{\int{\Omega_2}\nabla u \cdot \nabla v \ d\boldsymbol{x}}{a_1(u,v)},$$
$$f(v; \boldsymbol{\mu}) = \underbrace{\mu_1}{\Theta^{f}0(\boldsymbol{\mu})} \underbrace{\int{\Gamma_{base}}v \ ds}{f_0(v)}.$$
We will implement the numerical discretization of the problem in the class
class UnsteadyThermalBlock(ParabolicCoerciveProblem):
by specifying the coefficients $\Theta^{m}(\boldsymbol{\mu})$, $\Theta^{a}_(\boldsymbol{\mu})$ and $\Theta^{f}(\boldsymbol{\mu})$ in the method
def compute_theta(self, term):
and the bilinear forms $m_(u, v)$, $a(u, v)$ and linear forms $f_(v)$ in
def assemble_operator(self, term):
End of explanation
"""


mesh = Mesh("data/thermal_block.xml")
subdomains = MeshFunction("size_t", mesh, "data/thermal_block_physical_region.xml")
boundaries = MeshFunction("size_t", mesh, "data/thermal_block_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, two components)
End of explanation
"""


problem = UnsteadyThermalBlock(V, subdomains=subdomains, boundaries=boundaries)
mu_range = [(0.1, 10.0), (-1.0, 1.0)]
problem.set_mu_range(mu_range)
problem.set_time_step_size(0.05)
problem.set_final_time(3)

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


reduction_method = ReducedBasis(problem)
reduction_method.set_Nmax(20, POD_Greedy=4)
reduction_method.set_tolerance(1e-5, POD_Greedy=1e-2)

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


reduction_method.initialize_training_set(100)
reduced_problem = reduction_method.offline()

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


online_mu = (8.0, -1.0)
reduced_problem.set_mu(online_mu)
reduced_solution = reduced_problem.solve()
plot(reduced_solution, reduced_problem=reduced_problem, every=5, interval=500)

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


reduction_method.initialize_testing_set(10)
reduction_method.error_analysis()

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


reduction_method.initialize_testing_set(10)
reduction_method.speedup_analysis()

"""
Explanation: 4.8. Perform a speedup analysis
End of explanation
"""


