"""
(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 = 60
animation_loop_seconds = 2 #time in seconds for animation to loop one cycle

bgcolor = (255,255,255) #white
potential_curve_color = (0,0,0) #black
ecolor = (100,0,100) #electron is purple

linear_color = (0, 0, 150)
shg_color = (0, 150, 0)
const_color = (150, 0, 0)

eradius = 20

img_height = 500
img_width = 1000

top_arrow_y = 350
middle_arrow_y = 380
bottom_arrow_y = 410
arrow_width = 8

# Limits of the potential curve
xmin = 100
xmax = 900
ymin = 40
ymax = 300

# 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))

def potential_curve(x):
    """
    My potential curve y as a function of x
    """
    xscaled = (x-xmin) / (xmax - xmin)
    if xscaled < 0.2:
        yscaled = (0.2 - xscaled)**2 / (0.2**2)
    else:
        yscaled = (xscaled - 0.2)**2 / (0.8**2)
    # flip it, because higher y-coordinate is lower in pygame drawing
    yscaled = 1 - yscaled
    return ymin + (ymax - ymin) * yscaled

curve_bottom_x = 0.79 * xmin + 0.21 * xmax
curve_bottom_y = potential_curve(curve_bottom_x)

def electron_curve(x):
    """
    The path that the electron center travels along
    """
    # xscaled = (x-xmin) / (xmax - xmin)
    y = min(potential_curve(x), potential_curve(x+eradius), potential_curve(x-eradius))
    return y - eradius

# Constants and function for calculating electron motion
linear_coef = 0.3
shg_coef = 0.07
displacement = 0.32

def e_x(phase):
    """
    x-position of electron as a function of phase (from 0 to 2pi)
    """
    xscaled = (linear_coef * exp(1j * phase) + shg_coef * exp(2j * phase)
               + displacement).real
    return xmin + xscaled * (xmax - xmin)

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)]

    # Put the potential curve in the form of a list of points, to be drawn below
    xs = range(xmin, xmax + 1,1)
    ys = [potential_curve(x) for x in xs]
    potential_curve_path = zip(xs, ys)
    
    for frame in range(frames_in_anim):
        phase = 2 * pi * frame / frames_in_anim
        electron_x = e_x(phase)
        electron_y = electron_curve(electron_x)
        
        # initialize surface
        surf = pg.Surface((img_width,img_height))
        surf.fill(bgcolor)
        
        # draw potential curve
        pg.draw.lines(surf, potential_curve_color, False,
                      potential_curve_path, 10)
        
        # draw vertical line to first arrow
        pg.draw.line(surf, (0,0,0), (curve_bottom_x,curve_bottom_y),
                     (curve_bottom_x, top_arrow_y), 3)
        
        # draw three arrows
        linear_term = (linear_coef * exp(1j * phase)).real * (xmax - xmin)
        shg_term = (shg_coef * exp(2j * phase)).real * (xmax - xmin)
        
        draw_arrow(surf,
                   (curve_bottom_x, top_arrow_y),
                   (curve_bottom_x + linear_term, top_arrow_y),
                   width=arrow_width, color=linear_color)
        draw_arrow(surf,
                   (curve_bottom_x + linear_term, middle_arrow_y),
                   (curve_bottom_x + linear_term + shg_term, middle_arrow_y),
                   width=arrow_width, color=shg_color)
        draw_arrow(surf,
                   (curve_bottom_x + linear_term + shg_term, bottom_arrow_y),
                   (electron_x, bottom_arrow_y),
                   width=arrow_width, color=const_color)
        
        # draw electron
        pg.draw.circle(surf, ecolor,
                       ((int(round(electron_x)), int(round(electron_y)))),
                       eradius, 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))
    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()