#
#  Author:
#
#    Christian Hill
#
from wave_2d_pde import WaveEqn2D
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
#
print ( '' )
print ( 'raindrops():' )
print ( '  Wave equation in 2D with random raindrops.' )
#
# Raindrop probability (with each time tick) and intensity.
#
drop_probability = 0.01
max_intensity = 10
#
# Width of the Gaussian profile for each initial drop.
#
drop_width = 2
#
# Number of Gaussian widths to calculate to.
#
ndrop_widths = 3
#
# Size of the Gaussian template each drop is based on.
#
NDx = drop_width * ndrop_widths
NDy = drop_width * ndrop_widths
Dx = np.arange(0, NDx, 1, dtype=np.int32)
Dy = np.arange(0, NDy, 1, dtype=np.int32)
MDx, MDy = np.meshgrid(Dx, Dy)
#
# Create the 2D template of the initial drop.
#
cx = NDx // 2
cy = NDy // 2
gauss_template = np.exp ( -(((MDx-cx)/drop_width)**2 + ((MDy-cy)/drop_width)**2) )
#
#  Create "sim", the raindrop simulation.
#  nx and ny are the dimension of the domain;
#  c is the wave speed;
#  h is the space grid spacing;
#  dt is the time grid spacings;
#  If use_mur_abc is True, the Mur absorbing boundary
#  conditions will be used; if False, the Dirichlet
#  (reflecting) boundary conditions are used.
#
nx = 200
ny = 200
c = 0.2
h = 1
dt = 1
sim = WaveEqn2D ( nx, ny, c=c, h=h, dt=dt, use_mur_abc=True )

fig, ax = plt.subplots()
ax.axis ( "off" )
ax.set_title ( "Wave equation, random raindrops" )
img = ax.imshow ( sim.u[0], vmin=0, vmax=max_intensity, cmap='YlGnBu_r' )
#
#  Functions that control the animation.
#  Must be defined before calling FuncAnimation.
#
def update(i):
    """Advance the simulation by one tick."""
    # On random steps, pick a random (x,y) and start a raindrop there.
    #
    if np.random.random() < drop_probability:
        x = np.random.randint(NDx//2, nx-NDx//2-1)
        y = np.random.randint(NDy//2, ny-NDy//2-1)
        sim.u[0, y-NDy//2:y+NDy//2, x-NDx//2:x+NDx//2] = max_intensity * gauss_template

    sim.update()

def init():
    """
    Initialization, because we're blitting and need references to the
    animated objects.
    """
    return img,

def animate(i):
    """Draw frame i of the animation."""
    update ( i )
    img.set_data ( sim.u[0] )
    return img,

interval = 2 * sim.dt
nframes = 4000

ani = animation.FuncAnimation ( 
  fig, 
  animate, 
  frames = nframes,
  repeat = False,
  init_func = init, 
  interval = interval, 
  blit = True )

plt.show()