from .forms import FormContext, MixedFormContext, FieldPair, ElementVector, vector_load_form
from .context_types import VolumeContext, SurfaceContext
from .space import (
    FESpaceBase,
    FESpace,
    FESpacePytree,
    make_space,
    make_space_pytree,
    make_tet_space,
    make_tet_space_pytree,
    make_tet10_space,
    make_tet10_space_pytree,
    make_hex_space,
    make_hex_space_pytree,
    make_hex20_space,
    make_hex20_space_pytree,
    make_hex27_space,
    make_hex27_space_pytree,
)
from .mixed_space import MixedFESpace, MixedProblem, MixedDirichletBC, MixedBlockSystem
from .data import MeshData, BasisData, SpaceData
from .basis import (
    make_hex_basis,
    make_hex_basis_pytree,
    make_hex20_basis,
    make_hex20_basis_pytree,
    make_hex27_basis,
    make_hex27_basis_pytree,
    make_tet_basis,
    make_tet_basis_pytree,
    make_tet10_basis,
    make_tet10_basis_pytree,
)
from .assembly import (
    AssemblyPolicy,
    assemble_bilinear_form,
    assemble_linear_form,
    assemble_functional,
    assemble_residual,
    assemble_jacobian,
    assemble_residual_scatter,
    assemble_jacobian_scatter,
    assemble_residual_elementwise,
    assemble_jacobian_elementwise,
    make_element_residual_kernel,
    make_element_bilinear_kernel,
    make_element_linear_kernel,
    make_element_jacobian_kernel,
    make_element_kernel,
    element_residual,
    element_jacobian,
    make_sparsity_pattern,
    chunk_pad_stats,
    BatchedAssembler,
    assemble_jacobian_values,
    assemble_mass_matrix,
    scalar_body_force_form,
    make_scalar_body_force_form,
    constant_body_force_form,
)
from .interp import eval_shape_functions_hex8, interpolate_field_at_element_points
from .weakform import (
    Expr,
    FieldRef,
    ParamRef,
    trial_ref,
    test_ref,
    unknown_ref,
    zero_ref,
    Params,
    LinearForm,
    BilinearForm,
    ResidualForm,
    kernel,
    compile_surface_linear,
    compile_surface_bilinear,
    compile_mixed_surface_residual,
    MixedWeakForm,
    make_mixed_residuals,
    compile_bilinear,
    compile_linear,
    compile_residual,
    compile_mixed_residual,
)
from .basis import (
    HexTriLinearBasis,
    HexTriLinearBasisPytree,
    HexSerendipityBasis20,
    HexSerendipityBasis20Pytree,
    HexTriQuadraticBasis27,
    HexTriQuadraticBasis27Pytree,
    TetLinearBasis,
    TetLinearBasisPytree,
    TetQuadraticBasis10,
    TetQuadraticBasis10Pytree,
)
import importlib

from .solver import spdirect_solve_cpu, spdirect_solve_gpu, spdirect_solve_jax, coo_to_csr

_MESH_EXPORTS = {
    "BaseMeshPytree",
    "SurfaceWithFacetMap",
    "bbox_predicate",
    "plane_predicate",
    "axis_plane_predicate",
    "slab_predicate",
    "HexMesh",
    "HexMeshPytree",
    "StructuredHexBox",
    "SurfaceMesh",
    "SurfaceMeshPytree",
    "SurfaceWithElemConn",
    "surface_with_elem_conn",
    "SurfaceSupermesh",
    "build_surface_supermesh",
    "tag_axis_minmax_facets",
    "TetMesh",
    "TetMeshPytree",
    "StructuredTetBox",
    "StructuredTetTensorBox",
    "load_gmsh_mesh",
    "load_gmsh_hex_mesh",
    "load_gmsh_tet_mesh",
    "make_surface_from_facets",
    "ContactCouplingMatrix",
    "assemble_contact_coupling_matrices",
    "assemble_contact_onesided_floor",
    "assemble_contact_interface_residual",
    "assemble_contact_interface_jacobian",
    "map_surface_facets_to_tet_elements",
    "map_surface_facets_to_hex_elements",
    "tri_area",
    "tri_quadrature",
    "facet_triangles",
    "facet_shape_values",
    "volume_shape_values_at_points",
    "quad_shape_and_local",
    "quad9_shape_values",
    "hex27_gradN",
    "ContactSurfaceSpace",
    "ContactSide",
    "ContactOperators",
    "OneSidedContact",
    "OneSidedContactSurfaceSpace",
    "OneToManyContactSurfaceSpace",
    "assemble_contact_operators",
    "assemble_contact_interface_residual",
    "assemble_contact_interface_jacobian",
    "assemble_contact_coupling_matrices",
    "assemble_contact_kkt",
    "solve_contact_kkt",
    "facet_gap_values",
    "active_contact_facets",
}

_SOLVER_EXPORTS = {
    "SparsityPattern",
    "FluxSparseMatrix",
    "FluxBlockMatrix",
    "coalesce_coo",
    "concat_flux",
    "block_diag_flux",
    "build_block_system",
    "split_block_matrix",
    "BlockSystem",
    "DirichletBC",
    "LinearSolver",
    "NonlinearSolver",
    "NewmarkResult",
    "newmark_solve_linear",
    "petsc_solve",
    "petsc_shell_solve",
    "petsc_is_available",
    "enforce_dirichlet_dense",
    "enforce_dirichlet_dense_jax",
    "enforce_dirichlet_fluxsparse",
    "enforce_dirichlet_sparse",
    "enforce_dirichlet_system",
    "split_dirichlet_matrix",
    "free_dofs",
    "restrict_flux_to_free",
    "condense_dirichlet_system",
    "condense_dirichlet_fluxsparse",
    "condense_dirichlet_fluxsparse_coo",
    "condense_dirichlet_dense",
    "expand_dirichlet_solution",
    "cg_solve",
    "cg_solve_jax",
    "build_cg_operator",
    "CGOperator",
    "make_block_jacobi_preconditioner",
    "NonlinearAnalysis",
    "NewtonLoopConfig",
    "LoadStepResult",
    "NewtonSolveRunner",
    "solve_nonlinear",
    "LinearAnalysis",
    "LinearSolveConfig",
    "LinearStepResult",
    "LinearSolveRunner",
    "newton_solve",
    "petsc_solve",
    "petsc_shell_solve",
    "petsc_is_available",
}

_SOLVER_BC_EXPORTS = {
    "SurfaceFormField",
    "SurfaceFormContext",
    "SurfaceBilinearContext",
    "vector_surface_load_form",
    "make_vector_surface_load_form",
    "assemble_surface_linear_form",
    "assemble_surface_bilinear_form",
}


def __getattr__(name: str):
    if name in _MESH_EXPORTS:
        module = importlib.import_module("..mesh", __name__)
    elif name in _SOLVER_EXPORTS:
        module = importlib.import_module("..solver", __name__)
    elif name in _SOLVER_BC_EXPORTS:
        module = importlib.import_module("..solver.bc", __name__)
    else:
        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

    value = getattr(module, name)
    globals()[name] = value
    return value

__all__ = [
    "FESpace",
    "FESpaceBase",
    "FESpacePytree",
    "MixedFESpace",
    "MixedProblem",
    "MixedDirichletBC",
    "MixedBlockSystem",
    "FormContext",
    "MixedFormContext",
    "VolumeContext",
    "SurfaceContext",
    "VolumeContext",
    "SurfaceContext",
    "FieldPair",
    "ElementVector",
    "vector_load_form",
    "make_space",
    "make_space_pytree",
    "make_hex_basis",
    "make_hex_basis_pytree",
    "make_hex_space",
    "make_hex_space_pytree",
    "make_hex20_basis",
    "make_hex20_basis_pytree",
    "make_hex20_space",
    "make_hex20_space_pytree",
    "make_hex27_basis",
    "make_hex27_basis_pytree",
    "make_hex27_space",
    "make_hex27_space_pytree",
    "make_tet_basis",
    "make_tet_basis_pytree",
    "make_tet_space",
    "make_tet_space_pytree",
    "make_tet10_basis",
    "make_tet10_basis_pytree",
    "make_tet10_space",
    "make_tet10_space_pytree",
    "assemble_bilinear_form",
    "assemble_linear_form",
    "AssemblyPolicy",
    "assemble_functional",
    "assemble_residual",
    "assemble_jacobian",
    "assemble_residual_scatter",
    "assemble_jacobian_scatter",
    "make_element_residual_kernel",
    "make_element_bilinear_kernel",
    "make_element_linear_kernel",
    "make_element_jacobian_kernel",
    "element_residual",
    "element_jacobian",
    "make_sparsity_pattern",
    "assemble_jacobian_values",
    "assemble_residual_elementwise",
    "assemble_jacobian_elementwise",
    "assemble_mass_matrix",
    "scalar_body_force_form",
    "make_scalar_body_force_form",
    "constant_body_force_form",
    "eval_shape_functions_hex8",
    "interpolate_field_at_element_points",
    "Expr",
    "FieldRef",
    "ParamRef",
    "trial_ref",
    "test_ref",
    "unknown_ref",
    "zero_ref",
    "Params",
    "LinearForm",
    "BilinearForm",
    "ResidualForm",
    "MixedWeakForm",
    "make_mixed_residuals",
    "kernel",
    "compile_bilinear",
    "compile_linear",
    "compile_residual",
    "compile_surface_linear",
    "compile_surface_bilinear",
    "compile_mixed_surface_residual",
    "compile_mixed_residual",
    "grad",
    "sym_grad",
    "outer",
    "dot",
    "sdot",
    "ddot",
    "inner",
    "action",
    "gaction",
    "normal",
    "ds",
    "dOmega",
    "I",
    "det",
    "inv",
    "transpose",
    "log",
    "transpose_last2",
    "matmul",
    "matmul_std",
    "einsum",
    "HexMesh",
    "HexMeshPytree",
    "StructuredHexBox",
    "SurfaceMesh",
    "SurfaceMeshPytree",
    "SurfaceWithElemConn",
    "surface_with_elem_conn",
    "SurfaceFormField",
    "SurfaceFormContext",
    "SurfaceBilinearContext",
    "vector_surface_load_form",
    "make_vector_surface_load_form",
    "assemble_surface_linear_form",
    "assemble_surface_bilinear_form",
    "load_gmsh_mesh",
    "load_gmsh_hex_mesh",
    "load_gmsh_tet_mesh",
    "make_surface_from_facets",
    "SurfaceSupermesh",
    "build_surface_supermesh",
    "ContactCouplingMatrix",
    "assemble_contact_coupling_matrices",
    "assemble_contact_onesided_floor",
    "assemble_contact_interface_residual",
    "assemble_contact_interface_jacobian",
    "map_surface_facets_to_tet_elements",
    "map_surface_facets_to_hex_elements",
    "tri_area",
    "tri_quadrature",
    "facet_triangles",
    "facet_shape_values",
    "volume_shape_values_at_points",
    "quad_shape_and_local",
    "quad9_shape_values",
    "hex27_gradN",
    "ContactSurfaceSpace",
    "ContactSide",
    "ContactOperators",
    "OneSidedContact",
    "OneSidedContactSurfaceSpace",
    "OneToManyContactSurfaceSpace",
    "assemble_contact_operators",
    "assemble_contact_interface_residual",
    "assemble_contact_interface_jacobian",
    "assemble_contact_coupling_matrices",
    "assemble_contact_kkt",
    "solve_contact_kkt",
    "facet_gap_values",
    "active_contact_facets",
    "TetMesh",
    "TetMeshPytree",
    "StructuredTetBox",
    "StructuredTetTensorBox",
    "tag_axis_minmax_facets",
    "BaseMeshPytree",
    "SurfaceWithFacetMap",
    "bbox_predicate",
    "plane_predicate",
    "axis_plane_predicate",
    "slab_predicate",
    "HexTriLinearBasis",
    "HexTriLinearBasisPytree",
    "HexSerendipityBasis20",
    "HexSerendipityBasis20Pytree",
    "HexTriQuadraticBasis27",
    "HexTriQuadraticBasis27Pytree",
    "TetLinearBasis",
    "TetLinearBasisPytree",
    "TetQuadraticBasis10",
    "TetQuadraticBasis10Pytree",
    "spdirect_solve_cpu",
    "spdirect_solve_gpu",
    "spdirect_solve_jax",
    "coo_to_csr",
    "SparsityPattern",
    "FluxSparseMatrix",
    "FluxBlockMatrix",
    "coalesce_coo",
    "concat_flux",
    "block_diag_flux",
    "LinearSolver",
    "NonlinearSolver",
    "NewmarkResult",
    "newmark_solve_linear",
    "enforce_dirichlet_dense",
    "enforce_dirichlet_dense_jax",
    "enforce_dirichlet_fluxsparse",
    "enforce_dirichlet_sparse",
    "enforce_dirichlet_system",
    "split_dirichlet_matrix",
    "free_dofs",
    "restrict_flux_to_free",
    "condense_dirichlet_system",
    "condense_dirichlet_fluxsparse",
    "condense_dirichlet_fluxsparse_coo",
    "condense_dirichlet_dense",
    "expand_dirichlet_solution",
    "cg_solve",
    "cg_solve_jax",
    "NonlinearAnalysis",
    "NewtonLoopConfig",
    "LoadStepResult",
    "NewtonSolveRunner",
    "solve_nonlinear",
    "LinearAnalysis",
    "LinearSolveConfig",
    "LinearStepResult",
    "LinearSolveRunner",
    "newton_solve",
    "jacobian_sparse",
    "eval_shape_functions_hex8",
    "interpolate_field_at_element_points",
    "MeshData",
    "BasisData",
    "SpaceData",
]
