"""
(C) Steven Byrnes, 2013. 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 = 30
animation_loop_seconds = 2 #time in seconds for animation to loop one cycle

bgcolor = (255,255,255) #white
linecolor = (0,0,0) #outline of resistor is black
ecolor = (0,0,0) #electron color is black

img_height = 100
img_width = 300

#transmission line wire length and thickness, and y-coordinate of each wire
tl_length = img_width * 6//7
tl_thickness = 5
tl_top_y = img_height*4//9
tl_bot_y = img_height*5//9 - tl_thickness//2 #second term is to keep it symmetric

wavelength = 0.6 * tl_length

resistor_length = img_height//2
resistor_width = resistor_length//3

resistor_center = (img_width - resistor_width*3//2, img_height//2)

top_lead_path = [(tl_length, tl_top_y + tl_thickness-1),
                 (tl_length, img_height//9),
                 (resistor_center[0], img_height//9),
                 resistor_center]

bot_lead_path = [(x,img_height-y+1) for (x,y) in top_lead_path]

lead_thickness = 2

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 (100+100*V, 100 - 100*V, 100-100*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 make_resistor_surf(phase_at_top):
    """
    make a pygame surface representing the resistor. topphase is phase at top
    """
    imgarray = [[rgb_from_V(cos(phase_at_top) * (1 - 2*y/resistor_length))
                     for y in range(resistor_length)]
                    for x in range(resistor_width)]
    surf = pg.surfarray.make_surface(asarray(imgarray))
    pg.draw.rect(surf,linecolor,surf.get_rect(),1) #1-pixel black outline
    return surf

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 = 3 #pixels between electron path and corresponding wires
    path_length = ( 2*(tl_length - d) #transmission lines
                  + 2*(img_height//3) #left vertical leads
                  + 2*(resistor_center[0] - tl_length + 2*d + lead_thickness)
                  + 2*(resistor_length//2 - img_height//9) #right vertical leads
                  + resistor_length) #through resistor
    howfar = param * path_length
    
    #move right across top transmission line
    if howfar < tl_length - d:
        x = howfar
        y = tl_top_y - d
        phase = phase_top_left + 2 * pi * x / wavelength
        return {'pos':(x,y), 'phase':phase}
    howfar -= (tl_length - d)
    
    #move up lead
    if howfar < img_height//3:
        x = tl_length - d
        y = tl_top_y - d - howfar
        phase = phase_top_left + 2 * pi * tl_length / wavelength
        return {'pos':(x,y), 'phase':phase}
    howfar -= img_height//3
    
    #move right to above resistor
    if howfar < (resistor_center[0]- tl_length) + 2*d + lead_thickness:
        x = tl_length - d + howfar
        y = img_height//9 - d
        phase = phase_top_left + 2 * pi * tl_length / wavelength
        return {'pos':(x,y), 'phase':phase}
    howfar -= (resistor_center[0] - tl_length) + 2*d + lead_thickness
    
    #move down to top of resistor
    if howfar < (resistor_length//2 - img_height//9):
        x = resistor_center[0] + d + lead_thickness
        y = img_height//9 - d + howfar
        phase = phase_top_left + 2 * pi * tl_length / wavelength
        return {'pos':(x,y), 'phase':phase}
    howfar -= (resistor_length//2 - img_height//9)
    
    #move down resistor
    if howfar < resistor_length:
        x = resistor_center[0] + resistor_width//2 + d
        y = resistor_center[1] - resistor_length//2 + howfar
        phase = phase_top_left + 2 * pi * tl_length / wavelength
        return {'pos':(x,y), 'phase':phase}
    howfar -= resistor_length
    
    #beyond here use the mirror symmetry
    flipdata = e_path(1-param, phase_top_left)
    flipdata['pos'] = (flipdata['pos'][0], img_height - flipdata['pos'][1] + 2)
    return flipdata

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
        phase_top_right = phase_top_left + 2 * pi * tl_length / wavelength
        
        #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 lead wires
        color = rgb_from_V(cos(phase_top_right))
        pg.draw.lines(surf,color,False,top_lead_path,lead_thickness)
        color = rgb_from_V(cos(phase_top_right + pi))
        pg.draw.lines(surf,color,False,bot_lead_path,lead_thickness)
        
        #draw resistor
        resistor_surf = make_resistor_surf(phase_top_right)
        surf.blit(resistor_surf, (resistor_center[0] - resistor_width//2,
                                  resistor_center[1] - resistor_length//2))
        
        #draw electrons
        num_electrons = 100
        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])/(1.3*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), 2, 0)
        
        pg.image.save(surf, 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)

main()