Source code for Muscat.Bridges.PlotlyBridge

from typing import Optional, Union, TYPE_CHECKING

import numpy as np

from Muscat.Simple import Mesh

if TYPE_CHECKING:
    from  plotly.graph_objects import Mesh3d



def GetPlotlySurfaceMesh(myMesh: Mesh, fieldToPlot : Optional[Union[str,np.ndarray]] = None):
    """Function to convert a Mesh and a Nodal fields (element fields to yet implemented) to
    a Plotly graph_object Mesh3d library. Internally the norm of the field is computed if
    the fieldToPlot is a vector field (like normals).

    (This functions need the plotly package installed)

    Parameters
    ----------
    myMesh : Mesh
        The mesh to be converted
    fieldToPlot : Optional[Union[str,np.ndarray]], optional
        point field or point field name to put the intensity field of
        the Mesh3d. by default None

    Returns
    -------
    plotly.graph_objects.Mesh3d
        the mesh converted to Plotly format

    Raises
    ------
    RuntimeError
        if field.shape[0]  != number of node of the mesh
    RuntimeError
        if the field name is not present in the mesh.nodesFields
    """

    import plotly.graph_objects as go

    mesh = myMesh.View()
    #for the moment the compute skin does permit to transport elemField data to the skin
    mesh.elemFields = {}

    from Muscat.Simple import ElementFilter, ED
    from Muscat.MeshTools.MeshModificationTools import ComputeSkin
    from Muscat.MeshTools.MeshInspectionTools import ExtractElementsByElementFilter

    ComputeSkin(mesh=mesh, inPlace=True)
    mesh2D = ExtractElementsByElementFilter(mesh,ElementFilter(dimensionality=2), copy=False)
    from Muscat.MeshTools.MeshTetrahedrization import Tetrahedrization

    mesh_of_tris = Tetrahedrization(mesh2D)

    if mesh_of_tris.nodes.shape[1]<3  :
        z = np.zeros_like(mesh_of_tris.nodes[:,0])
    else:
        z=mesh_of_tris.nodes[:,2]

    res = go.Mesh3d(
        x=mesh_of_tris.nodes[:,0],
        y=mesh_of_tris.nodes[:,1],
        z=z,
        i=mesh_of_tris.GetElementsOfType(ED.Triangle_3).connectivity[:,0],
        j=mesh_of_tris.GetElementsOfType(ED.Triangle_3).connectivity[:,1],
        k=mesh_of_tris.GetElementsOfType(ED.Triangle_3).connectivity[:,2],
        flatshading= True
    )
    if fieldToPlot is None:
        return res

    res.showscale=True

    if isinstance(fieldToPlot,str):
        if fieldToPlot in mesh.nodeFields:
            name = fieldToPlot
            fieldToPlot = mesh.nodeFields[fieldToPlot]
            intensitymode = 'vertex'
        elif fieldToPlot in mesh.elemFields:
            name = fieldToPlot
            fieldToPlot = mesh.elemFields[fieldToPlot]
            intensitymode = 'cell'
        else:
            raise RuntimeError("Error: fieldToPlot not in nodeFields neither elemFields")
    else:
        name = "data"
        if fieldToPlot.shape[0] == myMesh.GetNumberOfNodes():
            intensitymode = 'vertex'
        elif fieldToPlot.shape[0] == myMesh.GetNumberOfElements():
            intensitymode = 'cell'
        else:
            raise RuntimeError("Error: fieldToPlot must match the number of nodes or the number of elements")

    if len(fieldToPlot.shape) > 1:
        fieldToPlot = np.linalg.norm(fieldToPlot, axis= 1)
        name = f"norm({name})"

    res.intensity = fieldToPlot
    res.intensitymode = intensitymode
    res.colorbar=dict(title=dict(text=name))
    res.colorscale='Viridis'

    return res




def PlotMesh(mesh : Union[Mesh,'Mesh3d'], nodeFieldToPlot : Optional[Union[str,np.ndarray]] = None):
    """Plot a mesh using Plotly
    (This functions need the plotly package installed)

    Parameters
    ----------
    mesh : Union[Mesh,'Mesh3d']
        Mesh to plot (Muscat Mesh or plotly Mesh3D)
    fieldToPlot : Optional[Union[str,np.ndarray]], optional
        point field or point field name to put the intensity field of
        the Mesh3d. by default None

    Returns
    -------
    Plotly Figure
        the Plotly figure
    """

    import plotly.graph_objects as go

    if isinstance(mesh, Mesh):
        meshtoplot = GetPlotlySurfaceMesh(mesh, nodeFieldToPlot )
    else:
        meshtoplot = mesh

    fig = go.Figure(data=[meshtoplot])
    fig.show()

    return fig




def CheckIntegrity(GUI: bool = False):

    from Muscat.Helpers.CheckTools import SkipTest

    if SkipTest("PLOTLY_NO_FAIL"):  # pragma: no cover
        return "skip"
    try:
        import plotly
    except:  # pragma: no cover
        return "skip : plotly not installed"

    from Muscat.MeshTools.MeshTools import GetElementsCenters
    from Muscat.Simple import ReadMesh
    from Muscat.TestData import GetTestDataPath

    myMesh: Mesh = ReadMesh(GetTestDataPath() + "dent3D.msh")

    myMesh.nodeFields["xpos"] = myMesh.nodes[:,0]
    myMesh.nodeFields["pos"] = myMesh.nodes
    myMesh.elemFields["centers"] = GetElementsCenters(myMesh )

    plotlymesh = GetPlotlySurfaceMesh(myMesh)
    plotlymesh = GetPlotlySurfaceMesh(myMesh, "xpos")

    if GUI:
        PlotMesh(plotlymesh)
        PlotMesh(myMesh)

    return 'ok'


if __name__ == "__main__":
    print(CheckIntegrity(True))  # pragma: no cover