import numpy as np
import fenics
fenics.set_log_level(fenics.LogLevel.ERROR)
import fenics_adjoint as fa
import ufl

from fenics_pymc3 import create_fem_theano_op
from fenics_pymc3 import to_numpy

# Create mesh for the unit square domain
n = 10
mesh = fa.UnitSquareMesh(n, n)

# Define discrete function spaces and functions
V = fenics.FunctionSpace(mesh, "CG", 1)
W = fenics.FunctionSpace(mesh, "DG", 0)

def solve_fenics(kappa0, kappa1):
    # This function inside should be traceable by fenics_adjoint
    f = fa.Expression(
        "10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2
    )

    u = fa.Function(V)
    bcs = [fa.DirichletBC(V, fa.Constant(0.0), "on_boundary")]

    inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
    v = fenics.TestFunction(V)
    F = inner(kappa0*grad(u), grad(v)) * dx - kappa1 * f * v * dx
    fa.solve(F == 0, u, bcs=bcs)
    return u

# Let's generate artificial data
true_kappa0 = fa.Constant(1.25)
true_kappa1 = fa.Constant(0.55)

true_solution = solve_fenics(true_kappa0, true_kappa1)
true_solution_numpy = to_numpy(true_solution)

# Perturb the solution with noise
noise_level = 0.05
noise = np.random.normal(scale=noise_level * np.linalg.norm(true_solution_numpy), size=true_solution_numpy.size)
noisy_solution = true_solution_numpy + noise

# Define FEniCS template representation of Theano/NumPy input
templates = (fa.Constant(0.0), fa.Constant(0.0))

# Now let's create Theano wrapper of `solve_fenics` function
theano_fem_solver = create_fem_theano_op(templates)(solve_fenics)

# `theano_fem_solver` can now be used inside PyMC3's model
import pymc3 as pm
import theano.tensor as tt

with pm.Model() as fit_poisson:
    sigma = pm.InverseGamma("sigma", alpha=3.0, beta=0.5)

    kappa0 = pm.TruncatedNormal(
        "kappa0", mu=1.0, sigma=0.5, lower=1e-5, upper=2.0, shape=(1,)
    )
    kappa1 = pm.TruncatedNormal(
        "kappa1", mu=0.7, sigma=0.5, lower=1e-5, upper=2.0, shape=(1,)
    )
    predicted_solution = pm.Deterministic("pred_sol", theano_fem_solver(kappa0, kappa1))

    d = pm.Normal("d", mu=predicted_solution, sd=sigma, observed=noisy_solution)

with fit_poisson:
    trace = pm.sample(500, chains=4, cores=4)

pm.summary(trace)
#                 mean     sd  hdi_3%  hdi_97%  ...  ess_sd  ess_bulk  ess_tail  r_hat
# sigma          0.015  0.001   0.013    0.017  ...   689.0     715.0     723.0   1.00
# kappa0[0]      1.247  0.377   0.586    1.926  ...   334.0     331.0     462.0   1.02
# kappa1[0]      0.586  0.179   0.267    0.900  ...   352.0     352.0     582.0   1.02