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

This code runs 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 

import pygame as pg
from numpy import cos, pi, sin, asarray

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

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

bgcolor = (255,255,255) #white
ecolor = (0,0,0) #electron color is black

# pygame draws pixel-art, not smoothed. Therefore I am drawing it
# bigger, then smoothly shrinking it down

img_height = 210
img_width = 900

final_height = 70
final_width = 300

# ~23 megapixel limit for wikipedia animated gifs
assert final_height * final_width * frames_in_anim < 22e6

#transmission line wire length and thickness, and y-coordinate of each wire
tl_length = img_width
tl_thickness = 27
tl_top_y = 66
tl_bot_y = img_height - tl_top_y - tl_thickness + 2

wavelength = 1.4 * tl_length

def rgb_from_V(V):
    """
    voltage V varies -1 to +1. Return a color as a function of V.
    Color is a 3-tuple red,green,blue, each 0 to 255.
    """
    return (200+55*V, 200-55*V, 200-55*V)

def tup_round(tup):
    """
    round each element of a tuple to nearest integer
    """
    return tuple(int(round(x)) for x in tup)

def make_wire_surf(phase_at_left):
    """
    make a pygame surface representing a colored wire. startphase is phase
    at left side of the wire.
    """
    imgarray = [[rgb_from_V(cos(phase_at_left + 2*pi*x/wavelength))
                 for y in range(tl_thickness)] for x in range(tl_length)]
    return pg.surfarray.make_surface(asarray(imgarray))

def e_path(param, phase_top_left):
    """
    as param goes 0 to 1, this returns {'pos': (x, y), 'phase':phi},
    where (x,y) is the coordinates of the corresponding point on the electron
    dot path, and phi is the phase for an electron at that point on the path.
    phase_top_left is phase of the left side of the top wire.
    """
    # d is a vertical offset between the electrons and the wires
    d = -21
    # pad is how far to extend the transmission line beyond the image borders
    # (since those electrons may enter the image a bit)
    pad = 36
    path_length = 2*(tl_length + 2*pad)
    howfar = param * path_length
    
    # move right across top transmission line
    if howfar <= path_length / 2:
        x = howfar - pad
        y = tl_top_y - d
        phase = phase_top_left + 2 * pi * x / wavelength
        return {'pos':(x,y), 'phase':phase}
    # then move left across the bottom transmission line
    x = path_length - howfar - pad
    y = tl_bot_y + tl_thickness + d
    phase = phase_top_left + 2 * pi * x / wavelength
    return {'pos':(x,y), 'phase':phase}

def main():
    #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_top_left = -2 * pi * frame / frames_in_anim
        
        #initialize surface
        surf = pg.Surface((img_width,img_height))
        surf.fill(bgcolor);
        
        #draw transmission line
        top_wire_surf = make_wire_surf(phase_top_left)
        bottom_wire_surf = make_wire_surf(phase_top_left + pi)
        surf.blit(top_wire_surf, (0, tl_top_y))
        surf.blit(bottom_wire_surf, (0, tl_bot_y))
        
        #draw electrons
        num_electrons = 70
        equilibrium_params = [x/(num_electrons-1) for x in range(num_electrons)]
        phases = [e_path(a, phase_top_left)['phase'] for a in equilibrium_params]
        now_params = [equilibrium_params[i] + sin(phases[i])/(0.45*num_electrons)
                           for i in range(num_electrons)]
        coords = [e_path(a, phase_top_left)['pos'] for a in now_params]
        for coord in coords:
            pg.draw.circle(surf, ecolor, tup_round(coord), 4, 0)
        
        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))
    # Use the "convert" command (part of ImageMagick) to build the animation
    command_list = ['convert', '-delay', frame_delay, '-loop', '0'] + filename_list + ['anim.gif']
    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)

main()