"""
(C) Steven Byrnes, 2014. This code is released under the MIT license
http://opensource.org/licenses/MIT

This code should work in Python 2.7 or 3.3. It requires imagemagick to be
installed; that's how it assembles images into animated GIFs.
"""

from __future__ import division, print_function

import pygame as pg
from math import pi
from cmath import exp

import subprocess, os
directory_now = os.path.dirname(os.path.realpath(__file__))

frames_in_anim = 200
animation_loop_seconds = 5 #time in seconds for animation to loop one cycle

bgcolor = (255,255,255) #white

circle_radius = 12

img_height = 300
img_width = 1000

# pygame draws pixel-art, not smoothed. Therefore I am drawing it
# bigger, then smoothly shrinking it down
final_width = int(round(0.3 * img_width))
final_height = int(round(0.3 * img_height))

# Constants and function for calculating electron motion
def coef(n):
    """ f(t) = sum of Re[coef(n) * e^(int)] """
    return 0 if n % 2 == 0 else -1j / n

def exact(phase):
    """ exact square wave function """
    if phase == 0 or phase == pi:
        return 0
    if (phase % (2*pi)) < pi:
        return pi/4
    return -pi/4

center_x = 500
scale = 450 # scale multiplies all horizontal distances

def draw_arrow(surf, tail_xy, head_xy, width=2, color=(0,0,0)):
    """
    draw a horizontal arrow
    """
    # tail_xy and head_xy are 2-tuples. Unpack them first
    tail_x, tail_y = tail_xy
    head_x, head_y = head_xy
    assert head_y == tail_y
    h = 16 # arrowhead height
    b = 18 # arrowhead half-base
    if tail_x < head_x:
        # rightward arrow
        triangle = [(head_x, head_y),
                    (head_x - h, head_y - b),
                    (head_x - h, head_y + b)]
    else:
        # leftward arrow
        triangle = [(head_x, head_y),
                    (head_x + h, head_y - b),
                    (head_x + h, head_y + b)]
    pg.draw.line(surf, color, (tail_x, tail_y), (head_x, head_y), width)
    pg.draw.polygon(surf, color, triangle, 0)

def main():
    """ function for creating the animated GIF """
    # Make and save a drawing for each frame
    filename_list = [os.path.join(directory_now, 'temp' + str(n) + '.png')
                         for n in range(frames_in_anim)]

    
    for frame in range(frames_in_anim):
        phase = 2 * pi * frame / frames_in_anim
        
        # initialize surface
        surf = pg.Surface((img_width,img_height))
        surf.fill(bgcolor)
        
        # draw all the arrows
        x_now = center_x
        y_now = 30
        for n in range(1,13,2):
            new_term = (scale * coef(n) * exp(n * 1j * phase)).real
            draw_arrow(surf, (x_now, y_now), (x_now + new_term, y_now), width=8)
            x_now += new_term
            y_now += 30
        
        # draw the two circles
        pg.draw.circle(surf, (100,0,100),
                       (int(round(x_now)),int(round(y_now + 10))),
                       circle_radius, 0)
        pg.draw.circle(surf, (50,50,150),
                       (int(round(center_x + scale * exact(phase))),
                            int(round(y_now + 40))),
                       circle_radius, 0)

        # scale down then save the surface
        shrunk_surface = pg.transform.smoothscale(surf, (final_width, final_height))
        pg.image.save(shrunk_surface, filename_list[frame])

    seconds_per_frame = animation_loop_seconds / frames_in_anim
    frame_delay = str(int(seconds_per_frame * 100))
    command_list = ['convert', '-delay', frame_delay, '-loop', '0'] + filename_list + ['anim.gif']
    # Use the "convert" command (part of ImageMagick) to build the animation
    subprocess.call(command_list, cwd=directory_now)
    # Earlier, we saved an image file for each frame of the animation. Now
    # that the animation is assembled, we can (optionally) delete those files
    if True:
        for filename in filename_list:
            os.remove(filename)
    return

main()