from dolfin import *
from rbnics import *

"""
Explanation: TUTORIAL 13 - Elliptic Optimal Control
Keywords: distributed optimal control, geometric parametrization, inf-sup condition, POD-Galerkin
1. Introduction
This tutorial addresses a distributed elliptic optimal control for the Laplace equation with geometrical parametrization. We consider an "original" domain $\Omega_o(\boldsymbol{\mu})$ divided into two parts $\Omega_o^1$ and $\Omega_o^2(\boldsymbol{\mu})$, as in the following picture:
<img src="data/mesh1.png" width="70%"/>
The problem is characterized by two parameters. The first parameter $\mu_0$ controls the shape of the deformable subdomain $\Omega_o^2$. The second parameter $\mu_1$ controls the parameter dependent observation function $y_d(\boldsymbol{\mu})$ such that 
$$y_d(\boldsymbol{\mu}) = 
\begin{cases} 
    1 \; \text{in} \; \Omega_o^1 \ 
    \mu_1 \; \text{in} \; \Omega_o^2 (\boldsymbol{\mu})
\end{cases} $$
The ranges of the two parameters are the following: $$\mu_0 \in [1, 3.5] \; \text{and} \; \mu_1 \in [0.5, 2.5]$$
Thus, the parameter vector $\boldsymbol{\mu}$ is given by $$\boldsymbol{\mu}=(\mu_0,\mu_1)$$ on the parameter domain $$\mathbb{P} = [1.0,3.5] \times [0.5,2.5].$$
In order to obtain a faster approximation of the optimal control problem, and avoid any remeshing, we pursue an optimize-then-discretize approach using the POD-Galerkin method from a fixed reference domain. 
2. Parametrized formulation
Let $y_0(\boldsymbol{\mu})$ be the temperature in the domain $\Omega_o(\boldsymbol{\mu}).$
Consider the following optimal control problem:
$$\underset{y_0,u_o}{min} \; J(y_o(\boldsymbol{\mu}), u_o(\boldsymbol{\mu});\boldsymbol{\mu}) = \frac{1}{2} \left\lVert y_o(\boldsymbol{\mu})-y_d(\boldsymbol{\mu})\right\rVert^2_{L^2(\Omega_o)} + \frac{\alpha}{2} \left\lVert u_o(\boldsymbol{\mu}) \right\rVert^2_{U_o)}, $$
such that
$$\begin{cases} 
    -\Delta y_o(\boldsymbol{\mu}) = u_o(\boldsymbol{\mu}) \quad \text{in} \; \Omega_o(\boldsymbol{\mu}), \ 
    y_o(\boldsymbol{\mu}) = g_D \quad \quad \; \quad \text{on} \; \Gamma_D^o(\boldsymbol{\mu})=\partial \Omega_o(\boldsymbol{\mu})
\end{cases}$$
where 
* $y_o$ and $u_o$ are the state and control functions defined on the original domain
* the Dirichlet boundary condition is given by $g_D=1$ 
Note that the functional spaces are parameter dependent due to the shape variation.
The corresponding weak formulation comes from the Lagrangian method to derive the solution to the optimal control problem:
Using the Lagrangian functional, the problem becomes: <center> find $(y_o,p_o,u_o) \in \mathbb{Y_o} \times \mathbb{Q_o} \times \mathbb{U_o}\; : \; \nabla L(y_o,p_o,u_o)[(z_o,q_o,v_o)]=0 \quad \forall (z_o,q_o,v_o) \in \mathbb{Y_o} \times \mathbb{Q_o} \times \mathbb{U_o}$ </center>
Which gives:
<center>
    $
    \begin{cases}
        L_{o,p} = f(q_o) + c(u_o,q_o) - a(y_o,q_o) \
        L_{o,y} = m(y_o,z_o) - g(y_d,z_o) - a^(z_o,p_o) \
        L_{o,u} = \alpha n(u_o,v_o) + c^(v_o,p_o)
    \end{cases}
    $
</center>
where
* the parametrized bilinear form $a(\cdot,\cdot; \boldsymbol{\mu}): \mathbb{Y_o} \times \mathbb{Q_o} \rightarrow \mathbb{R}$ is defined as $$a(y_o,q_o; \boldsymbol{\mu}) = \int_{\Omega_o} \nabla y_o \cdot \nabla q_o \ d \Omega$$
* the parametrized bilinear form $c(\cdot,\cdot; \boldsymbol{\mu}): \mathbb{U_o} \times \mathbb{Q_o} \rightarrow \mathbb{R}$ is defined as $$c(u,q; \boldsymbol{\mu}) = \int_{\Omega_o} u_o \cdot q_o \ d \Omega$$
* the parametrized linear form $f(\cdot; \boldsymbol{\mu}): \mathbb{Y_o} \rightarrow \mathbb{R}$ is defined as $$f(q_o; \boldsymbol{\mu}) = 0$$
* the parametrized bilinear form $m(\cdot,\cdot; \boldsymbol{\mu}): \mathbb{Z_o} \times \mathbb{Z_o} \rightarrow \mathbb{R}$ is defined as $$m(y_o,z_o; \boldsymbol{\mu}) = \int_{\Omega_o} y_o \cdot z_o \ d \Omega$$
* the parametrized bilinear form $n(\cdot,\cdot; \boldsymbol{\mu}): \mathbb{U_o} \times \mathbb{U_o} \rightarrow \mathbb{R}$ is defined as $$n(u_o,v_o; \boldsymbol{\mu}) = \int_{\Omega_o} u_o \cdot v_o \ d \Omega$$
* the parametrized bilinear form $g(\cdot,\cdot; \boldsymbol{\mu}): \mathbb{Z_o} \times \mathbb{Y_o} \rightarrow \mathbb{R}$ is defined as $$g(y_d,z_o; \boldsymbol{\mu}) = \mu_1 \int_{\Omega_o} y_d \cdot z_o \ d \Omega$$
and
* the functional space $\mathbb{Y}_o(\boldsymbol{\mu})$ is defined as $$\mathbb{Y}_o(\boldsymbol{\mu})=H_o^1(\Omega_o(\boldsymbol{\mu})$$
* the functional space $\mathbb{U}_o(\boldsymbol{\mu})$ is defined as $$\mathbb{U}_o(\boldsymbol{\mu})=L^2(\Omega_o(\boldsymbol{\mu}))$$
* the functional space $\mathbb{Q}_o=\mathbb{Y}_o$
* the functional space $\mathbb{Z}_o \supset \mathbb{Y}_o$
Since this problem is recast in the framework of saddle-point problems, the reduced basis problem must satisfy the inf-sup condition, thus an aggregated space for the state and adjoint variables is defined.
End of explanation
"""


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

    # Default initialization of members
    def __init__(self, V, **kwargs):
        # Call the standard initialization
        EllipticOptimalControlProblem.__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"]
        yup = TrialFunction(V)
        (self.y, self.u, self.p) = split(yup)
        zvq = TestFunction(V)
        (self.z, self.v, self.q) = split(zvq)
        self.dx = Measure("dx")(subdomain_data=subdomains)
        self.ds = Measure("ds")(subdomain_data=boundaries)
        # Regularization coefficient
        self.alpha = 0.01
        # Desired state
        self.y_d = Constant(1.0)
        # Customize linear solver parameters
        self._linear_solver_parameters.update({
            "linear_solver": "mumps"
        })

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

    # Return theta multiplicative terms of the affine expansion of the problem.
    def compute_theta(self, term):
        mu = self.mu
        if term in ("a", "a*"):
            theta_a0 = 1.0
            return (theta_a0,)
        elif term in ("c", "c*"):
            theta_c0 = 1.0
            return (theta_c0,)
        elif term == "m":
            theta_m0 = 1.0
            return (theta_m0,)
        elif term == "n":
            theta_n0 = self.alpha
            return (theta_n0,)
        elif term == "f":
            theta_f0 = 1.0
            return (theta_f0,)
        elif term == "g":
            theta_g0 = 1.0
            theta_g1 = mu[1]
            return (theta_g0, theta_g1)
        elif term == "h":
            theta_h0 = 1.0
            theta_h1 = mu[1]**2
            return (theta_h0, theta_h1)
        elif term == "dirichlet_bc_y":
            theta_bc0 = 1.
            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.
    def assemble_operator(self, term):
        dx = self.dx
        if term == "a":
            y = self.y
            q = self.q
            a0 = inner(grad(y), grad(q)) * dx
            return (a0,)
        elif term == "a*":
            z = self.z
            p = self.p
            as0 = inner(grad(z), grad(p)) * dx
            return (as0,)
        elif term == "c":
            u = self.u
            q = self.q
            c0 = u * q * dx
            return (c0,)
        elif term == "c*":
            v = self.v
            p = self.p
            cs0 = v * p * dx
            return (cs0,)
        elif term == "m":
            y = self.y
            z = self.z
            m0 = y * z * dx
            return (m0,)
        elif term == "n":
            u = self.u
            v = self.v
            n0 = u * v * dx
            return (n0,)
        elif term == "f":
            q = self.q
            f0 = Constant(0.0) * q * dx
            return (f0,)
        elif term == "g":
            z = self.z
            y_d = self.y_d
            g0 = y_d * z * dx(1)
            g1 = y_d * z * dx(2)
            return (g0, g1)
        elif term == "h":
            y_d = self.y_d
            h0 = y_d * y_d * dx(1, domain=mesh)
            h1 = y_d * y_d * dx(2, domain=mesh)
            return (h0, h1)
        elif term == "dirichlet_bc_y":
            bc0 = [DirichletBC(self.V.sub(0), Constant(1.0), self.boundaries, i) for i in range(1, 7)]
            return (bc0,)
        elif term == "dirichlet_bc_p":
            bc0 = [DirichletBC(self.V.sub(2), Constant(0.0), self.boundaries, i) for i in range(1, 7)]
            return (bc0,)
        elif term == "inner_product_y":
            y = self.y
            z = self.z
            x0 = inner(grad(y), grad(z)) * dx
            return (x0,)
        elif term == "inner_product_u":
            u = self.u
            v = self.v
            x0 = u * v * dx
            return (x0,)
        elif term == "inner_product_p":
            p = self.p
            q = self.q
            x0 = inner(grad(p), grad(q)) * dx
            return (x0,)
        else:
            raise ValueError("Invalid term for assemble_operator().")

"""
Explanation: 3. Affine Decomposition
In order to obtain an affine decomposition, we recast the problem on a fixed, parameter independent, reference domain $\Omega$. We choose the reference domain characterized by $\mu_0$=1 which we generate through the generate_mesh notebook provided in the data folder.
End of explanation
"""


mesh = Mesh("data/mesh1.xml")
subdomains = MeshFunction("size_t", mesh, "data/mesh1_physical_region.xml")
boundaries = MeshFunction("size_t", mesh, "data/mesh1_facet_region.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
"""


scalar_element = FiniteElement("Lagrange", mesh.ufl_cell(), 1)
element = MixedElement(scalar_element, scalar_element, scalar_element)
V = FunctionSpace(mesh, element, components=["y", "u", "p"])

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


problem = EllipticOptimalControl(V, subdomains=subdomains, boundaries=boundaries)
mu_range = [(1.0, 3.5), (0.5, 2.5)]
problem.set_mu_range(mu_range)

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


pod_galerkin_method = PODGalerkin(problem)
pod_galerkin_method.set_Nmax(10)

"""
Explanation: 4.4. Prepare reduction with a POD-Galerkin method
End of explanation
"""


lifting_mu = (1.0, 1.0)
problem.set_mu(lifting_mu)
pod_galerkin_method.initialize_training_set(100)
reduced_problem = pod_galerkin_method.offline()

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


online_mu = (3.0, 0.6)
reduced_problem.set_mu(online_mu)
reduced_solution = reduced_problem.solve()
print("Reduced output for mu =", online_mu, "is", reduced_problem.compute_output())

plot(reduced_solution, reduced_problem=reduced_problem, component="y")

plot(reduced_solution, reduced_problem=reduced_problem, component="u")

plot(reduced_solution, reduced_problem=reduced_problem, component="p")

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


pod_galerkin_method.initialize_testing_set(100)
pod_galerkin_method.error_analysis()

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


pod_galerkin_method.speedup_analysis()

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


