import torch


def get_node_features(omega, theta, phi, sc=False):
    """
    Extract node features
    
    Parameters:
    -----------
    omega : torch.Tensor, (L, L) or (2, L, L)
        omega angles or sines and cosines

    theta : torch.Tensor, (L, L) or (2, L, L)
        theta angles or sines and cosines

    phi : torch.Tensor, (L, L) or (2, L, L)
        phi angles or sines and cosines

    Returns:
    --------
    * : torch.Tensor, (L, 10)
        {sin, cos}×(ωi, φi, φ_ri, ψi, ψ_ri). 
    """

    if (omega.sum() == 0.0) and (theta.sum() == 0.0) and (phi.sum() == 0.0):
        return torch.zeros(omega.shape[0], 10)

    def get_features(omega, theta, phi):
        # omega is symmetric, n1 is omega angle relative to prior
        device = omega.device
        n1 = torch.cat((torch.tensor([0.], device=device), torch.diagonal(omega, offset=1)))

        # theta is asymmetric, n2 and n3 relative to prior
        n2 = torch.cat((torch.diagonal(theta, offset=1), torch.tensor([0.], device=device)))
        n3 = torch.cat((torch.tensor([0.], device=device), torch.diagonal(theta, offset=-1)))

        # phi is asymmetric n4 and n5 relative to prior
        n4 = torch.cat((torch.diagonal(phi, offset=1), torch.tensor([0.], device=device)))
        n5 = torch.cat((torch.tensor([0.], device=device), torch.diagonal(phi, offset=-1)))

        ns = torch.stack([n1, n2, n3, n4, n5], dim=1)
        return ns

    if not sc:
        ns = get_features(omega, theta, phi)
        s = torch.sin(ns)
        c = torch.cos(ns)
    else:
        s = get_features(omega[0], theta[0], phi[0])
        c = get_features(omega[1], theta[1], phi[1])
    return torch.cat([s, c], dim=1)