import itertools

import dolfin
import numpy as np
import pytest
import ufl_legacy as ufl
from pulse2 import kinematics
from pulse2.active_model import Passive
from pulse2.active_stress import ActiveStress


def W_fun(Ta: float, eta: float = 0.0) -> float:
    return 0.5 * Ta * ((4 - 1) + eta * ((12 - 3) - (4 - 1)))


@pytest.mark.parametrize(
    "eta, Ta",
    itertools.product(
        (0.0, 0.2, 0.5, 1.0),
        (0.0, 1.0, 2.0),
    ),
)
def test_transversely_active_stress(eta, Ta, mesh, u) -> None:
    f0 = dolfin.Constant((1.0, 0.0, 0.0))
    active_model = ActiveStress(f0, eta=eta)

    u.interpolate(dolfin.Expression(("x[0]", "x[1]", "x[2]"), degree=1))
    F = kinematics.DeformationGradient(u)
    W = active_model.strain_energy(F)

    assert active_model.Fe(F) is F
    active_model.activation.assign(Ta)

    assert np.isclose(
        dolfin.assemble(W * ufl.dx(domain=mesh)),
        W_fun(Ta=Ta, eta=eta),
    )

    assert np.isclose(
        dolfin.assemble(active_model.Ta * ufl.dx(domain=mesh)),
        Ta,
    )
    active_model.T_ref.assign(2.0)
    assert np.isclose(
        dolfin.assemble(active_model.Ta * ufl.dx(domain=mesh)),
        2 * Ta,
    )


def test_Passive(u) -> None:
    active_model = Passive()

    u.interpolate(dolfin.Expression(("x[0]", "x[1]", "x[2]"), degree=1))
    F = kinematics.DeformationGradient(u)

    assert active_model.Fe(F) is F
    W = active_model.strain_energy(F)
    assert np.isclose(
        dolfin.assemble(W * ufl.dx(domain=u.function_space().mesh())), 0.0
    )
