from dolfin import *
from rbnics import *

"""
Explanation: TUTORIAL 04 - Graetz problem 1
Keywords: successive constraints method
1. Introduction
This Tutorial addresses geometrical parametrization and the successive constraints method (SCM). In particular, we will solve the Graetz problem, which deals with forced heat convection in a channel $\Omega_o(\mu_0)$ divided into two parts $\Omega_o^1$ and $\Omega_o^2(\mu_0)$, as in the following picture:
<img src="data/graetz_1.png" width="70%"/>
Boundaries $\Gamma_{o, 1} \cup \Gamma_{o, 5} \cup \Gamma_{o, 6}$ are kept at low temperature (say, zero), while boundaries $\Gamma_{o, 2}(\mu_0) \cup \Gamma_{o, 4}(\mu_0)$ are kept at high temperature (say, one). The convection is characterized by the velocity $\boldsymbol{\beta} = (x_1(1-x_1), 0)$, being $\boldsymbol{x}o = (x{o, 0}, x_1)$ the coordinate vector on the parametrized domain $\Omega_o(\mu_0)$.
The problem is characterized by two parameters. The first parameter $\mu_0$ controls the shape of deformable subdomain $\Omega_2(\mu_0)$. The heat transfer between the domains can be taken into account by means of the Péclet number, which will be labeled as the parameter $\mu_1$. The ranges of the two parameters are the following:
$$\mu_0 \in [0.1,10.0] \quad \text{and} \quad \mu_1 \in [0.01,10.0].$$
The parameter vector $\boldsymbol{\mu}$ is thus given by 
$$
\boldsymbol{\mu} = (\mu_0, \mu_1)
$$
on the parameter domain
$$
\mathbb{P}=[0.1,10.0]\times[0.01,10.0].
$$
In order to obtain a faster (yet, provably accurate) approximation of the problem, and avoiding any remeshing, we pursue a model reduction by means of a certified reduced basis reduced order method from a fixed reference domain.
The successive constraints method will be used to evaluate the stability factors.
2. Parametrized formulation
Let $u_o(\boldsymbol{\mu})$ be the temperature in the domain $\Omega_o(\mu_0)$.
We will directly provide a weak formulation for this problem
<center>for a given parameter $\boldsymbol{\mu}\in\mathbb{P}$, find $u_o(\boldsymbol{\mu})\in\mathbb{V}_o(\boldsymbol{\mu})$ such that</center>
$$a_o\left(u_o(\boldsymbol{\mu}),v_o;\boldsymbol{\mu}\right)=f_o(v_o;\boldsymbol{\mu})\quad \forall v_o\in\mathbb{V}_o(\boldsymbol{\mu})$$
where

the function space $\mathbb{V}o(\boldsymbol{\mu})$ is defined as
$$
\mathbb{V}_o(\mu_0) = \left{ v \in H^1(\Omega_o(\mu_0)): v|{\Gamma_{o,1} \cup \Gamma_{o,5} \cup \Gamma_{o,6}} = 0, v|{\Gamma{o,2}(\mu_0) \cup \Gamma_{o,2}(\mu_0)} = 1 \right}
$$
Note that, as in the previous tutorial, the function space is parameter dependent due to the shape variation. 
the parametrized bilinear form $a_o(\cdot, \cdot; \boldsymbol{\mu}): \mathbb{V}o(\boldsymbol{\mu}) \times \mathbb{V}_o(\boldsymbol{\mu}) \to \mathbb{R}$ is defined by
$$a_o(u_o,v_o;\boldsymbol{\mu}) = \mu_1\int{\Omega_o(\mu_0)} \nabla u_o \cdot \nabla v_o \ d\boldsymbol{x} + \int_{\Omega_o(\mu_0)} x_1(1-x_1) \partial_{x} u_o\ v_o \ d\boldsymbol{x},$$
the parametrized linear form $f_o(\cdot; \boldsymbol{\mu}): \mathbb{V}_o(\boldsymbol{\mu}) \to \mathbb{R}$ is defined by
$$f_o(v_o;\boldsymbol{\mu}) = 0.$$

The successive constraints method will be used to compute the stability factor of the bilinear form $a_o(\cdot, \cdot; \boldsymbol{\mu})$.
End of explanation
"""


@SCM()
@PullBackFormsToReferenceDomain()
@ShapeParametrization(
    ("x[0]", "x[1]"),  # subdomain 1
    ("mu[0]*(x[0] - 1) + 1", "x[1]"),  # subdomain 2
)
class Graetz(EllipticCoerciveProblem):

    # Default initialization of members
    @generate_function_space_for_stability_factor
    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.ds = Measure("ds")(subdomain_data=boundaries)
        # Store the velocity expression
        self.vel = Expression("x[1]*(1-x[1])", element=self.V.ufl_element())
        # Customize eigen solver parameters
        self._eigen_solver_parameters.update({
            "bounding_box_minimum": {
                "problem_type": "gen_hermitian", "spectral_transform": "shift-and-invert",
                "spectral_shift": 1.e-5, "linear_solver": "mumps"
            },
            "bounding_box_maximum": {
                "problem_type": "gen_hermitian", "spectral_transform": "shift-and-invert",
                "spectral_shift": 1.e5, "linear_solver": "mumps"
            },
            "stability_factor": {
                "problem_type": "gen_hermitian", "spectral_transform": "shift-and-invert",
                "spectral_shift": 1.e-5, "linear_solver": "mumps"
            }
        })

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

    # Return theta multiplicative terms of the affine expansion of the problem.
    @compute_theta_for_stability_factor
    def compute_theta(self, term):
        mu = self.mu
        if term == "a":
            theta_a0 = mu[1]
            theta_a1 = 1.0
            return (theta_a0, theta_a1)
        elif term == "f":
            theta_f0 = 1.0
            return (theta_f0,)
        elif term == "dirichlet_bc":
            theta_bc0 = 1.0
            return (theta_bc0,)
        else:
            raise ValueError("Invalid term for compute_theta().")

    # Return forms resulting from the discretization of the affine expansion of the problem operators.
    @assemble_operator_for_stability_factor
    def assemble_operator(self, term):
        v = self.v
        dx = self.dx
        if term == "a":
            u = self.u
            vel = self.vel
            a0 = inner(grad(u), grad(v)) * dx
            a1 = vel * u.dx(0) * v * dx
            return (a0, a1)
        elif term == "f":
            f0 = Constant(0.0) * v * dx
            return (f0,)
        elif term == "dirichlet_bc":
            bc0 = [DirichletBC(self.V, Constant(0.0), self.boundaries, 1),
                   DirichletBC(self.V, Constant(1.0), self.boundaries, 2),
                   DirichletBC(self.V, Constant(1.0), self.boundaries, 4),
                   DirichletBC(self.V, Constant(0.0), self.boundaries, 5),
                   DirichletBC(self.V, Constant(0.0), self.boundaries, 6)]
            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
In order to obtain an affine decomposition, we proceed as in the previous tutorial and recast the problem on a fixed, parameter independent, reference domain $\Omega$. As reference domain which choose the one characterized by $\mu_0 = 1$ which we generate through the generate_mesh notebook provided in the data folder.
As in the previous tutorial, we pull back the problem to the reference domain $\Omega$.
End of explanation
"""


mesh = Mesh("data/graetz_1.xml")
subdomains = MeshFunction("size_t", mesh, "data/graetz_physical_region_1.xml")
boundaries = MeshFunction("size_t", mesh, "data/graetz_facet_region_1.xml")

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


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

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


problem = Graetz(V, subdomains=subdomains, boundaries=boundaries)
mu_range = [(0.1, 10.0), (0.01, 10.0)]
problem.set_mu_range(mu_range)

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


reduction_method = ReducedBasis(problem)
reduction_method.set_Nmax(30, SCM=20)
reduction_method.set_tolerance(1e-5, SCM=1e-3)

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


lifting_mu = (1.0, 1.0)
problem.set_mu(lifting_mu)
reduction_method.initialize_training_set(200, SCM=250)
reduced_problem = reduction_method.offline()

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


online_mu = (10.0, 0.01)
reduced_problem.set_mu(online_mu)
reduced_solution = reduced_problem.solve()
plot(reduced_solution, reduced_problem=reduced_problem)

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


reduction_method.initialize_testing_set(100, SCM=100)
reduction_method.error_analysis(filename="error_analysis")

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


reduction_method.speedup_analysis(filename="speedup_analysis")

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


