text: # pylint: disable=not-callable, no-member, invalid-name, missing-docstring
  import math
  import os
  import subprocess
  import argparse
  import shutil
  import tqdm
  
  import plotly.graph_objs as go
  import torch
  
  from e3nn import o3, rsh
  
  
  def rsh_surface(l, m, scale, tr, rot):
      n = 50
      a = torch.linspace(0, 2 * math.pi, 2 * n)
      b = torch.linspace(0, math.pi, n)
      a, b = torch.meshgrid(a, b)
  
      f = rsh.spherical_harmonics_alpha_beta([l], a, b)
      f = torch.einsum('ij,...j->...i', o3.irr_repr(l, *rot), f)
      f = f[..., l + m]
  
      x, y, z = o3.angles_to_xyz(a, b)
  
      r = f.abs()
      x = scale * r * x + tr[0]
      y = scale * r * y + tr[1]
      z = scale * r * z + tr[2]
  
      max_value = 0.5
  
      return go.Surface(
          x=x.numpy(),
          y=y.numpy(),
          z=z.numpy(),
          surfacecolor=f.numpy(),
          showscale=False,
          cmin=-max_value,
          cmax=max_value,
          colorscale=[[0, 'rgb(0,50,255)'], [0.5, 'rgb(200,200,200)'], [1, 'rgb(255,50,0)']],
      )
  
  
  def main(lmax, resolution, steps):
      scale = 0.5 * math.sqrt(4 * math.pi) / math.sqrt(2 * lmax + 1)
  
      axis = dict(
          showbackground=False,
          showticklabels=False,
          showgrid=False,
          zeroline=False,
          title='',
          nticks=3,
          range=[-lmax / 2 - 0.5, lmax / 2 + 0.5]
      )
  
      layout = dict(
          width=resolution,
          height=resolution,
          scene=dict(
              xaxis=axis,
              yaxis=axis,
              zaxis=axis,
              aspectmode='manual',
              aspectratio=dict(x=1, y=1, z=1),
              camera=dict(
                  up=dict(x=0, y=0, z=1),
                  center=dict(x=0, y=0, z=0),
                  eye=dict(x=0, y=-1.3, z=0),
                  projection=dict(type='perspective'),
              ),
          ),
          paper_bgcolor='rgba(255,255,255,255)',
          plot_bgcolor='rgba(0,0,0,0)',
          margin=dict(l=0, r=0, t=0, b=0)
      )
  
      shutil.rmtree('sh')
      os.makedirs('sh')
  
      for i in tqdm.tqdm(range(steps)):
          rot = 2 * math.pi * i / steps
          a, b, c = 0, math.pi / 4, 0
          abc = o3.compose(-c, -b, -a, *o3.compose(0, 0, rot, a, b, c))
  
          surfaces = [
              rsh_surface(l, m, scale, [l + (m if m < 0 else 0) - lmax / 2, 0, lmax / 2 - l + (m if m > 0 else 0)], abc)
              for l in range(lmax + 1)
              for m in range(-l, l + 1)
          ]
  
          fig = go.Figure(surfaces, layout=layout)
          fig.write_image('sh/{:03d}.png'.format(i))
  
      subprocess.check_output(["convert", "-delay", "3", "-loop", "0", "sh/*.png", "output.gif"])
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
  
      parser.add_argument("--lmax", type=int, default=2)
      parser.add_argument("--resolution", type=int, default=500)
      parser.add_argument("--steps", type=int, default=30)
  
      args = parser.parse_args()
  
      main(args.lmax, args.resolution, args.steps)
  