from manimlib import *

import numpy as np
from scipy.special import sph_harm

# configure style
config.background_color = '#202020'
# config.tex_template.add_to_preamble(
#     r"\usepackage[p,osf]{scholax}"
#     r"\usepackage{amsmath}"
#     r"\usepackage[scaled=1.075,ncf,vvarbb]{newtxmath}"
# )

class FourierOnS1(ThreeDScene):
    @staticmethod
    def harmonic_cos(m, n, theta, phi):
        # get only the cos() part
        r = np.real(sph_harm(m, n, theta, phi)) * 5 # Y_n^m 

        # coordinates transformation
        x = r * np.cos(theta) * np.sin(phi)
        y = r * np.sin(phi) * np.sin(theta)
        z = r * np.cos(phi)

        return [x, y, z]

    @staticmethod
    def harmonic_surf(m, n):
        surf = Surface(uv_func=lambda u, v: FourierOnS1.harmonic_cos(m, n, u, v),
                       u_range=[0, 2 * PI], v_range=[-PI, PI], 
                       color=GREY, resolution=(101, 101), shadow=0.1)

        surf.mesh = SurfaceMesh(surf)
        surf.mesh.set_stroke(BLACK, 1, opacity=0.7)

        return surf

    def construct(self):
        axes = ThreeDAxes(x_range=[-2, 2], y_range=[-2, 2], z_range=[-3, 3])
        self.add(axes)

        frame = self.camera.frame
        frame.set_euler_angles(theta=-30*DEGREES, phi=70*DEGREES)
        frame.add_updater(lambda m, dt: m.increment_theta(-0.1 * dt))

        last_surf = None
        last_label = None
        for n in range(0, 5):
            for m in range(0, n):
                surf = FourierOnS1.harmonic_surf(m, n)
                label = Tex(r"Y_{%d}^{%d}(\theta, \phi)" % (n, m))

                if last_surf is None:
                    self.play(ShowCreation(surf))
                    self.play(ShowCreation(surf.mesh, lag_ratio=0.01))
                else:
                    self.play(FadeOut(last_surf), FadeOut(last_surf.mesh))
                    self.play(ShowCreation(surf))
                    self.play(ShowCreation(surf.mesh, lag_ratio=0.01))

                if last_label is not None:
                    self.remove(last_label)

                label.to_corner(UL)

                last_surf = surf 
                last_label = label

                self.wait(3)

        self.wait(5)