from dolfin import *
from rbnics import *

"""
Explanation: TUTORIAL 13 - Elliptic Optimal Control
Keywords: optimal control, inf-sup condition, reduced basis
1. Introduction
This tutorial addresses a distributed optimal control problem for the Graetz conduction-convection equation on the domain $\Omega$ shown below:
<img src="data/mesh2.png" width="60%"/>
The problem is characterized by 3 parameters. The first parameter $\mu_0$ represents the Péclet number, which describes the heat transfer between the two domains. The second and third parameters, $\mu_1$ and $\mu_2$, control the parameter dependent observation function $y_d(\boldsymbol{\mu})$ such that:
$$ y_d(\boldsymbol{\mu})=
\begin{cases}
    \mu_1 \quad \text{in} \; \hat{\Omega}_1 \
    \mu_2 \quad \text{in} \; \hat{\Omega}_2
\end{cases}
$$
The ranges of the three parameters are the following: $$\mu_0 \in [3,20], \mu_1 \in [0.5,1.5], \mu_2 \in [1.5,2.5]$$
The parameter vector $\boldsymbol{\mu}$ is thus given by $$\boldsymbol{\mu}=(\mu_0,\mu_1,\mu_2)$$ on the parameter domain $$\mathbb{P}=[3,20] \times [0.5,1.5] \times [1.5,2.5].$$
In order to obtain a faster approximation of the optimal control problem, we pursue an optimize-then-discretize approach using the reduced basis method.
2. Parametrized Formulation
Let $y(\boldsymbol{\mu})$, the state function, be the temperature field in the domain $\Omega$ and $u(\boldsymbol{\mu})$, the control function, act as a heat source. The observation domain $\hat{\Omega}$ is defined as: $\hat{\Omega}=\hat{\Omega}_1 \cup \hat{\Omega}_2$.
Consider the following optimal control problem:
$$
\underset{y,u}{min} \; J(y,u;\boldsymbol{\mu}) = \frac{1}{2} \left\lVert y(\boldsymbol{\mu})-y_d(\boldsymbol{\mu})\right\rVert ^2_{L^2(\hat{\Omega})}, \
s.t. 
\begin{cases}
    -\frac{1}{\mu_0}\Delta y(\boldsymbol{\mu}) + x_2(1-x_2)\frac{\partial y(\boldsymbol{\mu})}{\partial x_1} = u(\boldsymbol{\mu}) \quad \text{in} \; \Omega, \
    \frac{1}{\mu_0} \nabla y(\boldsymbol{\mu}) \cdot \boldsymbol{n} = 0 \qquad \qquad \qquad \quad \enspace \; \text{on} \; \Gamma_N, \
    y(\boldsymbol{\mu})=1 \qquad \qquad \qquad \qquad \qquad \enspace  \text{on} \; \Gamma_{D1}, \
    y(\boldsymbol{\mu})=2 \qquad \qquad \qquad \qquad \qquad \enspace \text{on} \; \Gamma_{D2}
\end{cases}
$$
The corresponding weak formulation comes from solving for the gradient of the Lagrangian function as detailed in the previous tutorial. 
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
"""


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
        # Store the velocity expression
        self.vel = Expression("x[1] * (1 - x[1])", element=self.V.sub(0).ufl_element())
        # Customize linear solver parameters
        self._linear_solver_parameters.update({
            "linear_solver": "mumps"
        })

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

    # Return stability factor
    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):
        mu = self.mu
        if term in ("a", "a*"):
            theta_a0 = 1.0 / mu[0]
            theta_a1 = 1.0
            return (theta_a0, theta_a1)
        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 = mu[1]
            theta_g1 = mu[2]
            return (theta_g0, theta_g1)
        elif term == "h":
            theta_h0 = 0.24 * mu[1]**2 + 0.52 * mu[2]**2
            return (theta_h0,)
        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
            vel = self.vel
            a0 = inner(grad(y), grad(q)) * dx
            a1 = vel * y.dx(0) * q * dx
            return (a0, a1)
        elif term == "a*":
            z = self.z
            p = self.p
            vel = self.vel
            as0 = inner(grad(z), grad(p)) * dx
            as1 = - vel * p.dx(0) * z * dx
            return (as0, as1)
        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(1) + y * z * dx(2)
            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
            g0 = z * dx(1)
            g1 = z * dx(2)
            return (g0, g1)
        elif term == "h":
            h0 = 1.0
            return (h0,)
        elif term == "dirichlet_bc_y":
            bc0 = [DirichletBC(self.V.sub(0), Constant(i), self.boundaries, i) for i in (1, 2)]
            return (bc0,)
        elif term == "dirichlet_bc_p":
            bc0 = [DirichletBC(self.V.sub(2), Constant(0.0), self.boundaries, i) for i in (1, 2)]
            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
For this problem the affine decomposition is straightforward.
End of explanation
"""


mesh = Mesh("data/mesh2.xml")
subdomains = MeshFunction("size_t", mesh, "data/mesh2_physical_region.xml")
boundaries = MeshFunction("size_t", mesh, "data/mesh2_facet_region.xml")

"""
Explanation: 4. Main program
4.1. Read the mesh for this problem
The mesh was generated by the data/generate_mesh_2.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 = [(3.0, 20.0), (0.5, 1.5), (1.5, 2.5)]
problem.set_mu_range(mu_range)

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


reduced_basis_method = ReducedBasis(problem)
reduced_basis_method.set_Nmax(40)

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


lifting_mu = (3.0, 1.0, 2.0)
problem.set_mu(lifting_mu)
reduced_basis_method.initialize_training_set(100)
reduced_problem = reduced_basis_method.offline()

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


online_mu = (15.0, 0.6, 1.8)
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
"""


reduced_basis_method.initialize_testing_set(100)
reduced_basis_method.error_analysis()

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


reduced_basis_method.speedup_analysis()

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


