# Python: příprava dat pro vizualizaci sférických harmonických funkcí
import numpy as np
import scipy

# Vytvoření polí s kartézskými souřadnicemi a daty
def createSphHarm(nth,nph):
  n=3; m=1                        # stupeň a řád sférické harmonické funkce
  nr=0                            # 0 pro jednu plochu
  thmin=0; thmax=np.pi            # meze souřadnice theta (poledník): theta_min a theta_max
  phmin=0; phmax=np.pi*2          # meze souřadnice phi (rovnoběžka): phi_min a phi_max
  x=np.empty((nr+1,nth+1,nph+1))  # pole kartézských souřadnic
  y=x.copy()
  z=x.copy()
  s=x.copy()                      # skalární data v uzlech sítě
  ir=nr                           # jedna plocha
  for ith in range(nth+1):        # v NumPy měníme indexy nejrychleji zprava
    th=thmin+ith*(thmax-thmin)/nth
    for iph in range(nph+1):
      ph=phmin+iph*(phmax-phmin)/nph
      s[ir,ith,iph]=scipy.special.sph_harm(m,n,ph,th).real
#     ss=1
      ss=abs(s[ir,ith,iph])       # faktor pro lokální deformaci kulové plochy
#     ss=s[ir,ith,iph]**2
      x[ir,ith,iph]=ss*np.sin(th)*np.cos(ph)
      y[ir,ith,iph]=ss*np.sin(th)*np.sin(ph)
      z[ir,ith,iph]=ss*np.cos(th)
  return (x,y,z,s)

# Vytvoření PNG-souboru v Matplotlibu (plocha v 3D)
def saveFig(x,y,z,s,file):
  import matplotlib.pyplot as plt
  import matplotlib.cm as cm
  import matplotlib.colors as colors
  with open(file,'w') as f:
    fig=plt.figure(layout='tight')
    ax=plt.axes(projection='3d')
    cmap=colors.ListedColormap(['yellow','yellow','yellow','black'])
    ax.plot_surface(x[0,:,:],y[0,:,:],z[0,:,:],cmap=cm.inferno)  # z-coordinate by predefined colormap
    # ax.plot_surface(x[0,:,:],y[0,:,:],z[0,:,:],cmap=cmap)        # z-coordinate by defined cmap
    # ax.plot_surface(x[0,:,:],y[0,:,:],z[0,:,:],facecolors=cmap(s[0,:,:])) # s-array by defined cmap
    ax.set_aspect('equal')
    ax.view_init(elev=10,azim=120)
    plt.savefig(file)
    plt.show()

# Vytvoření VTK-souboru pro ParaView (plocha v 3D)
# varianta STRUCTURED_GRID pro křivočarou síť/curvilinear structured grid
def saveVtkSG(x,y,z,s,file):
  (nr,nth,nph)=np.shape(x)   # počty uzlů
  fmt='8.4f'                 # formát výstupu
  with open(file,'w') as f:
    print('# vtk DataFile Version 3.0',file=f)
    print('vtk output',file=f)
    print('ASCII',file=f)
    print('DATASET STRUCTURED_GRID',file=f)
    print('DIMENSIONS',nr,nth,nph,file=f)
    print('POINTS',nr*nth*nph,'float',file=f)
    for iph in range(nph):   # pro ParaView měníme indexy nejrychleji zleva
      for ith in range(nth):
        for ir in range(nr):
          print(f'{x[ir,ith,iph]:{fmt}} {y[ir,ith,iph]:{fmt}} {z[ir,ith,iph]:{fmt}}',file=f)
    print('POINT_DATA',nr*nth*nph,file=f)
    print('SCALARS MyData float 1',file=f)
    print('LOOKUP_TABLE default',file=f)
    for iph in range(nph):
      for ith in range(nth):
        for ir in range(nr):
          print(f'{s[ir,ith,iph]:{fmt}}',file=f)

nth=90   # počet kroků sférické sítě ve směru theta (poledník)
nph=180  # počet kroků sférické sítě ve směru phi (rovnoběžka)
(x,y,z,s)=createSphHarm(nth,nph)
if True: soubor='Y:/Programovani24/10/sh.vtk'; saveVtkSG(x,y,z,s,soubor); print('Soubor',soubor,'uložen.')
if True: soubor='Y:/Programovani24/10/sh.png'; saveFig(x,y,z,s,soubor);   print('Soubor',soubor,'uložen.')