from meshgen.primitives import Quadrilateral, Pyramid, Hexahedron


def comsol_node_order(nodes, element_cls):
    """Change the node order to/from COMSOL style.

    It is exploited that COMSOL just swaps the positions of two nodes and not more
    complicated reorderings.

    Args:
        nodes (list): nodes of the element
        element_cls (object): type of element from which the nodes are

    Returns:
        list
    """
    if element_cls == Quadrilateral:
        return nodes[0], nodes[1], nodes[3], nodes[2]
    if element_cls == Pyramid:
        return nodes[0], nodes[1], nodes[3], nodes[2], nodes[4]
    if element_cls == Hexahedron:
        return (
            nodes[0],
            nodes[1],
            nodes[3],
            nodes[2],
            nodes[4],
            nodes[5],
            nodes[7],
            nodes[6],
        )
    return nodes
