"""
(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.
"""

# Use Python 3 style division: a/b is real division, a//b is integer division
from __future__ import division 

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

import pygame as pg
from numpy import pi, asarray, real, exp

frames_in_anim = 30
animation_loop_seconds = 2 #time in seconds for animation to loop one cycle

bgcolor = (255,255,255) #white
split_line_color = (0,0,0) #line down the middle is black
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 = 240
img_width = 900

final_height = 80
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 = int(img_width * .9)
tl_thickness = 27
tl_open_top_y = int(img_height*.1)
tl_open_bot_y = tl_open_top_y + 42
tl_short_top_y = int(img_height*.62)
tl_short_bot_y = tl_short_top_y + 42

wavelength = 0.6 * 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(f_phase_at_right, r_phase_at_right):
    """
    make a pygame surface representing a colored wire. f_phase and r_phase
    are the phases of the forward and reverse waves respectively.
    """
    def V(x):
        z = tl_length-x-1
        return 0.5*real(exp(1j*(f_phase_at_right + 2*pi*z/wavelength))
                      + exp(1j*(r_phase_at_right - 2*pi*z/wavelength)))
    imgarray = [[rgb_from_V( V(x) )
                 for y in range(tl_thickness)] for x in range(tl_length)]
    return pg.surfarray.make_surface(asarray(imgarray))

def e_path(param, f_phase_top_right, r_phase_top_right, which):
    """
    as param goes 0 to 1, this returns a dictionary: 'pos' is (x,y), the
    coordinates of the corresponding point on the electron
    dot path; 'f_phase' and 'r_phase' are the phases for the forward and
    reflected waves for an electron at that point on
    the path. top_right means right side of the top wire. which is either
    'open' or 'short' for which transmission line we're talking about.
    """
    d = -18 #pixels between electron path and corresponding wires
    
    #### Open transmission line ####
    
    if which == 'open':
        path_length = 2 * tl_length
        howfar = param * path_length
        
        #go right along top transmission line
        if howfar < tl_length:
            x = howfar
            y = tl_open_top_y - d
            f_phase = f_phase_top_right + 2 * pi * (tl_length-x) / wavelength
            r_phase = r_phase_top_right - 2 * pi * (tl_length-x) / wavelength
            return {'pos':(x,y), 'f_phase':f_phase, 'r_phase':r_phase}
        
        #go left along bottom transmission line
        x = 2*tl_length - howfar
        y = tl_open_bot_y + tl_thickness + d
        f_phase = f_phase_top_right + 2 * pi * (tl_length-x) / wavelength
        r_phase = r_phase_top_right - 2 * pi * (tl_length-x) / wavelength
        return {'pos':(x,y), 'f_phase':f_phase, 'r_phase':r_phase}
    
    #### Short transmission line ####
    
    path_length = (2 * tl_length + 3 * tl_thickness + 4*d +
                   + (tl_short_bot_y - tl_short_top_y))
    howfar = param * path_length
    
    #at the beginning, go right along top wire
    if howfar < tl_length:
        x = howfar
        y = tl_short_top_y - d
        f_phase = f_phase_top_right + 2 * pi * (tl_length-x) / wavelength
        r_phase = r_phase_top_right - 2 * pi * (tl_length-x) / wavelength
        return {'pos':(x,y), 'f_phase':f_phase, 'r_phase':r_phase}
    
    #at the end, go left along bottom wire
    if (path_length - howfar) < tl_length:
        x = path_length - howfar
        y = tl_short_bot_y + tl_thickness + d
        f_phase = f_phase_top_right + 2 * pi * (tl_length-x) / wavelength
        r_phase = r_phase_top_right - 2 * pi * (tl_length-x) / wavelength
        return {'pos':(x,y), 'f_phase':f_phase, 'r_phase':r_phase}
    
    #in the middle...
    f_phase = f_phase_top_right
    r_phase = r_phase_top_right
    
    #top part of short...
    if tl_length < howfar < tl_length + tl_thickness + d:
        x = howfar
        y = tl_short_top_y - d
    #bottom part of short...
    elif tl_length < (path_length - howfar) < tl_length + tl_thickness + d:
        x = path_length - howfar
        y = tl_short_bot_y + tl_thickness + d
    #vertical part of short...
    else:
        x = tl_length + tl_thickness + d
        y = (tl_short_top_y - d) + (howfar - (tl_length + tl_thickness + d))
    return {'pos':(x,y), 'f_phase':f_phase, 'r_phase':r_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):
        f_phase_open_top_right = -2 * pi * frame / frames_in_anim + pi/2
        r_phase_open_top_right = f_phase_open_top_right
        
        f_phase_short_top_right = -2 * pi * frame / frames_in_anim
        r_phase_short_top_right = f_phase_short_top_right + pi
                
        #initialize surface
        surf = pg.Surface((img_width,img_height))
        surf.fill(bgcolor);
        
        #draw transmission line
        open_top_wire_surf = make_wire_surf(f_phase_open_top_right,
                                            r_phase_open_top_right)
        surf.blit(open_top_wire_surf, (0, tl_open_top_y))
        
        open_bot_wire_surf = make_wire_surf(f_phase_open_top_right + pi,
                                            r_phase_open_top_right + pi)
        surf.blit(open_bot_wire_surf, (0, tl_open_bot_y))
        
        short_top_wire_surf = make_wire_surf(f_phase_short_top_right,
                                            r_phase_short_top_right)
        surf.blit(short_top_wire_surf, (0, tl_short_top_y))
        
        short_bot_wire_surf = make_wire_surf(f_phase_short_top_right + pi,
                                            r_phase_short_top_right + pi)
        surf.blit(short_bot_wire_surf, (0, tl_short_bot_y))
        
        #draw short wire
        color = rgb_from_V(0)
        pg.draw.line(surf,color,
                     (tl_length + tl_thickness//2,tl_short_top_y),
                     (tl_length + tl_thickness//2,tl_short_bot_y+tl_thickness-1),
                    tl_thickness)
        
        #draw line down the middle
        pg.draw.line(surf,split_line_color, (0,img_height//2),
                     (img_width,img_height//2), 12)
        
        #draw electrons
        num_electrons = 60
        equilibrium_params = [x/(num_electrons-1) for x in range(num_electrons)]
        for eq_a in equilibrium_params:
            for which in ['open', 'short']:
                f_phase_top_right = (f_phase_open_top_right if which == 'open'
                                        else f_phase_short_top_right)
                r_phase_top_right = (r_phase_open_top_right if which == 'open'
                                        else r_phase_short_top_right)
                temp = e_path(eq_a, f_phase_top_right,
                               r_phase_top_right, which)
                f_phase = temp['f_phase']
                r_phase = temp['r_phase']
                #displacement is always pi/2 out of phase with current. But
                #compared to voltage, it's +pi/2 for forward and -pi/2 for
                #reverse, because voltage reflection is negative of current
                #reflection.
                displacement = 0.5*real(exp(1j*(f_phase+pi/2))
                                  + exp(1j*(r_phase-pi/2)))
                now_a = eq_a + displacement/(.8*num_electrons)
                now_pos = e_path(now_a, f_phase_top_right,
                                            r_phase_top_right, which)['pos']
                pg.draw.circle(surf, ecolor, tup_round(now_pos), 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()