import numpy as np
import matplotlib.pyplot as plt
import IPython.display as ipd

def happy_bday(ts, note_fn):
    """
    Play the first 5 notes of happy birthday using
    some note generator specified functionally
    
    Parameters
    ----------
    ts: ndarray(N)
        Time samples for each notes
    note_fn:  function: (ts, p) => ndarray(N)
        A function that takes in ts and a note number
        p and returns an array with N samples of the note number
        p at the times in ts
    
    Returns
    -------
    ndarray(5N)
        The happy birthday samples with the particular note generator
    """
    y = note_fn(ts, 0)
    y = np.concatenate((y, note_fn(ts, 2)))
    y = np.concatenate((y, note_fn(ts, 0)))
    y = np.concatenate((y, note_fn(ts, 5)))
    y = np.concatenate((y, note_fn(ts, 4)))
    return y

sr = 44100
ts = np.arange(sr//2)/sr

pure_tone = lambda ts, p: np.array([]) ## TODO: This is a dummy method
square_wave = lambda ts, p: np.array([]) ## TODO: This is a dummy method

y = happy_bday(ts, pure_tone)
ipd.Audio(y, rate=sr)

y = happy_bday(ts, square_wave)
ipd.Audio(y, rate=sr)