diff --git a/waves.py b/waves.py
index 21d6fce..a3ccb53 100755
--- a/waves.py
+++ b/waves.py
@@ -2,9 +2,10 @@
 
 import cairo
 import math
-import datetime
+from datetime import datetime
 import calendar
 import random
+import colorsys
 from utils import random_color
 
 
@@ -25,53 +26,49 @@ DARK_BG = True
 # precision of the calculation
 PRECISION = 10000
 
-def days_color(date):
-    import colorsys
+def days_color(date: datetime) -> tuple[float, float, float]:
+    '''Return a color based on the progress through the year.'''
     # a day between 1 and 365 (inclusive)
     today = date.timetuple().tm_yday
-    year = date.year
-    days_in_year = 365 + calendar.isleap(year)
+    days_in_year = 365 + calendar.isleap(date.year)
     # between 0 and 1, how far through the year are we?
     progress = today/days_in_year
     return colorsys.hsv_to_rgb(progress, 1, 0.9)
 
-def days_amp(date, waves):
-    day = date.day
+def days_amp(date: datetime, waves) -> float:
+    '''Return the amplitude of waves, based on progress through the month.'''
     days_in_month = calendar.monthrange(date.year, date.month)[1]
     max_amp = 1/waves/2
-    return day/days_in_month * max_amp
+    return date.day/days_in_month * max_amp
 
-def days_count(date):
-    return date.month
-
-def create_wpotd(output):
-    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, WIDTH, HEIGHT)
+def create_wpotd(width: int, height: int):
+    #TODO: return a `file`-like object
+    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)
     ctx = cairo.Context(surface)
 
-    ctx.scale(WIDTH, HEIGHT)  # Normalizing the canvas
+    ctx.scale(width, height)  # Normalizing the canvas
     #ctx.set_antialias(cairo.Antialias.BEST)
 
-    step_size = 1/PRECISION
     lastpoints = [(x/PRECISION, 0) for x in range(PRECISION+1)]
 
-    #date = datetime.datetime.strptime('2021-06-30', '%Y-%m-%d')
-    date = datetime.datetime.today()
+    date = datetime.today()
     frequency = random.randint(10, 40)
     if DATE_BASED_COUNT:
-        waves = days_count(date)
+        # generate as many waves as the number of the current month
+        waves = date.month
     else:
         waves = 12
     if DATE_BASED_AMPLITUDE:
         amplitude = days_amp(date, waves)
     else:
         amplitude = 25
-    
     if MONOCHROME:
         if DATE_BASED_COLOR:
             r, g, b = days_color(date) #datetime.datetime.now())
         else:
             r, g, b = random_color()
-        alpha_step = 1/waves
+    alpha_step = 1/waves
+
     if DARK_BG:
         # make bg black
         ctx.rectangle(0, 0, 1, 1)
@@ -79,20 +76,24 @@ def create_wpotd(output):
         ctx.fill()
 
     wave_height = 1/waves
-    for num in range(waves+1):
+    step_size = 1/PRECISION
+
+    for wave_index in range(waves+1):
         if not MONOCHROME:
             r, g, b = random_color()
         points = []
         x = 0
         while x < 1:
-            y = amplitude * math.sin(frequency * x + (num * WAVE_OFFSET) ) 
-            points.append((x, ( y + (0.5+num)*wave_height)))
+            # step along, create points along the wave
+            y = amplitude * math.sin(frequency * x + (wave_index * WAVE_OFFSET))
+            points.append((x, (y + (0.5 + wave_index)*wave_height)))
             x += step_size
         #print(f'Draw {len(points)} points for curve {num}')
         if not MONOCHROME:
             ctx.set_source_rgb(r, g, b)
         else:
-            ctx.set_source_rgba(r, g, b, 1 - (alpha_step * num)) # make more transparent toward bottom
+            # make more transparent toward bottom
+            ctx.set_source_rgba(r, g, b, 1 - (alpha_step * wave_index)) 
         # draw waves
         ctx.move_to(*points[0])
         for p in points[1:]:
@@ -104,11 +105,12 @@ def create_wpotd(output):
         ctx.line_to(*points[0])
         ctx.fill()
         lastpoints = points
-    #return surface
-    surface.write_to_png(output)  # Output to PNG
+    return surface
+    #surface.write_to_png(output)  # Output to PNG
 
 def main():
-    create_wpotd('out/waves.png')
+    output = create_wpotd(1920, 1080)
+    output.write_to_png('out/waves.png')
 
 if __name__ == '__main__':
     main()