from mpi4py import MPI

import numpy as np
import ufl
from dolfinx.cpp.mesh import CellType
from dolfinx.fem import assemble_scalar, form
from dolfinx.mesh import create_unit_square, meshtags

msh = create_unit_square(MPI.COMM_WORLD, 1, 1, CellType.triangle)


# Setting different quadrature degrees and values for a and b, makes it fail
cells = np.array([0, 1])
tags = np.array([0, 1])
cell_tags = meshtags(msh, msh.topology.dim, cells, tags)
dx = ufl.dx(domain=msh, subdomain_data=cell_tags)
dx0 = dx(subdomain_id=(0, 1), degree=2)
dx1 = dx(subdomain_id=(0), degree=3)

a, b = 1.0, 2.0
ufl_form_0 = a * dx0
ufl_form_1 = b * dx1
ufl_form = ufl_form_0 + ufl_form_1

expected = 1.0 * a + 0.5 * b

computed_0 = assemble_scalar(form(ufl_form_0)) + assemble_scalar(form(ufl_form_1))
print(f"Two assemblies -- Expected: {expected}. Computed: {computed_0}")

computed_1 = assemble_scalar(form(ufl_form))
print(f"One assembly   -- Expected: {expected}. Computed: {computed_1}")