import pygame
import math
import random
import numpy as np


# Define some colors
BLACK    = (   0,   0,   0)
WHITE    = ( 255, 255, 255)
RED      = ( 255,   0,   0)
GREEN    = (   0, 255,   0)
BLUE     = (   0,   0, 255)
YELLOW   = ( 255, 255,   0)

pygame.init()

Lambda = 0.5

# Set the width and height of the screen [width, height]
size = (1000, 1000)
screen = pygame.display.set_mode(size)
bnd = 183370316
lifetime = 60*5 # 5 seconds

pygame.display.set_caption("Dynamisches System")

# Loop until the user clicks the close button.
done = False
centralize = False
freeze = False

# Used to manage how fast the screen updates
clock = pygame.time.Clock()

# define default values
initscale = 400
initspeed = 0.005
reset = False

def restart():
  global plist,qlist,curve
  global act,mode
  #plist,qlist = [ranpt()],[[ranpt()]] # list of points, list of traces
  plist,qlist,curve = [],[],[] # list of points, list of traces
  act,mode = "","" # interactive action

background_image = pygame.image.load("bluetrees.jpg").convert()

def intb(n):
  if n > bnd or n<-bnd or np.isnan(n):
    ans = bnd
  else:
    ans = int(n)
  return (ans)

#scale,center,speed = initscale,(0,0),initspeed # parameters 

def toscreen(pt):
  # convert point to pixel coordinates
  xco = (pt[0]-center[0])*scale+size[0]/2
  yco = (pt[1]-center[1])*scale+size[1]/2
  return ( intb(xco),size[1]-intb(yco) )  
  
def fromscreen(ipt):
  # grab screen point and return correct point in plane
  refl = size[1]-ipt[1]
  xco = (ipt[0]-size[0]/2.0)/scale+center[0] # make sure to convert to float
  yco = (refl-size[1]/2.0)/scale+center[1]
  return ( xco,yco )

def ranpt():
  ranscreen = (random.randint(0,size[0]),random.randint(0,size[1]))
  return (fromscreen(ranscreen))

def newpoint():
  # add/remove points
  global plist
  if act == "add":
    plist = plist + [fromscreen(pygame.mouse.get_pos())]
  if act == "remove" and len(plist) > 0:
    plist = plist[:-1]
  if act == "init":
    plist = []
  if act == "ranrev":
    plist = plist + [ranpt() for i in range(50)]

def newtrace():
  # add/remove traces
  global qlist
  if act == "add":
    qlist = qlist + [[fromscreen(pygame.mouse.get_pos())]]
  if act == "remove":
    qlist = qlist[:-1]
  if act == "init":
    qlist = []
  if act == "ranrev":
    qlist = qlist + [[ranpt()] for i in range(25)]

def newcurve():
  global curve
  if mode == "pen":
    curve = curve + [fromscreen(pygame.mouse.get_pos())]

def resize():
  global scale
  if act == "add":
    scale = scale*1.1
  if act == "remove":
    scale = scale/1.1
  if act == "init":
    scale = initscale

def drive():
  # adjust speed
  global speed
  if act == "add":
    speed = speed*1.5
  if act == "remove":
    speed = speed/1.5
  if act == "init":
    speed = initspeed
  if act == "ranrev":
    speed = -speed

def field(x,y):
  #return (-y,x) # rotation 
  #return (y+2*x*y+2*y*y*y,-x-y*y) # transformed rotation
  #return (-y-0.1*x,x-0.1*y) # spiral
  #return (np.sin(x)+np.sin(y),np.sin(x)+2*np.sin(y))
  #return ((x*x*25+y*y*25-10)/10,y*5*(1-x*x*25)/10) # heteroclinic limit orbits
  #return (x*x+y*y-10,y*(1-x*x)) # heteroclinic limit orbits
  #return (y+2*x*y-2*y*y*y,x-y*y) # saddle with nonlinear stable/unstable curves
  #return (x*y, x*x+y*y-0.5 ) # exercise 8/1
  #return (-0.5*x+y,-0.5*y+x*x ) # Exercise 8/3
  #return (y+Lambda*np.sin(10.0*x),-x) # many limit circles
  #return (-x*y+0.0, x*x+y*y-0.5 ) # heteroclinic
  #return (-x*y+par, x*x+y*y-0.5 ) # almost heteroclinic 
  #return ( x*(0.5-x*x-y*y)-y*x*x,y*(0.5-x*x-y*y)+x*x*x ) # does not work
  #return ( x*(0.5-x*x-y*y)-y,y*(0.5-x*x-y*y)+x ) # exercise 9/2
  #return (-y-x/2+y*y/2-2*x*y-2*y*y*y,x-y/2+y*y) # exercise 9/1
  # rad = np.sqrt(x*x+y*y)
  #return ( x*(1-rad*rad)-y*(1+x/rad), y*(1-rad*rad)+x*(1+x/rad) )
  #return (-16*x+16.0*y,-2*(x+y)*(x+y-8.0) )
  # c1 = 0.3
  # c2 = 1
  #return (-x*(c2*y-c1),c2*x*x) 
  #return(0,x-y) # structurally stable
  #return(0,y**2-x) # 1d saddle-node
  gg = (y-x*x+1)*(y+x*x-1)*x
  gx = -5*x*x*x*x+6*x*x+y*y-1
  gy = 2*x*y
  return (-gy-gg*gx,gx-gg*gy)

def move(pt):
  # Eulers polygon method
  #vect = field(pt[0],pt[1])
  #ptnew = ( pt[0]+speed*vect[0],pt[1]+speed*vect[1] )
  # two step method
  vect = field(pt[0],pt[1])
  pthalf = ( pt[0]+speed*vect[0]/2,pt[1]+speed*vect[1]/2 )
  vect = field(pthalf[0],pthalf[1])
  ptnew = ( pt[0]+speed*vect[0],pt[1]+speed*vect[1] )
  return (ptnew)

def movetr(trc):
  # extend a list of points by moving the last
  if len(trc) == lifetime:
    newtrc = trc[1:] + [move(trc[-1])]
  else:
    newtrc = trc + [move(trc[-1])]
  return (newtrc)

def movecurve(crv):
  return ([ move(pt) for pt in crv ])

def iterate():
  global plist,qlist,curve
  plist = [ move(pt) for pt in plist ]
  qlist = [ movetr(trc) for trc in qlist ]
  curve = movecurve(curve)

# -------- Print menue
print "p       mode point"
print "t       mode trace"
print "z       mode curve"
print "s       mode size"
print "g       go (mode speed)"
print "+       increase/add"
print "-       decrease/remove"
print "r       random points/reverse speed"
print "0       set to start value"
print "c       center"
print "a       adapt lambda"
print "SPC     clear"
print "q       quit"

# -------- Initialize paraneters 
scale,center,speed = initscale,(0,0),initspeed # parameters 
qlist = []
plist = [ranpt() for i in range(500)]
curve = []
act,mode = "","speed"
par = 0.05

# -------- Main Program Loop -----------
while not done:
  # --- Main event loop
  for event in pygame.event.get(): # User did something
    if event.type == pygame.KEYDOWN:
      if event.key == pygame.K_p:
        mode = "point"
      if event.key == pygame.K_t:
        mode = "trace"
      if event.key == pygame.K_s:
        mode = "size"
      if event.key == pygame.K_g:
        mode = "speed"
      if event.key == pygame.K_PLUS:
        act = "add"
      if event.key == pygame.K_MINUS:
        act = "remove"
      if event.key == pygame.K_0:
        act = "init"
      if event.key == pygame.K_r:
        act = "ranrev"
      if event.key == pygame.K_z:
        mode = "pen"
      if event.key == pygame.K_c:
        centralize = True
      if event.key == pygame.K_a:
        mode = "param"
      if event.key == pygame.K_SPACE:
        reset = True
      if event.key == pygame.K_q:
        done = True

  # --- Game logic should go here

  if mode == "point":
    newpoint()
  if mode == "trace":
    newtrace() 
  if mode == "pen":
    newcurve() 
  if mode == "size":
    resize() 
  if mode == "speed":
    drive() # adjust speed
    iterate() # move points and traces
  if mode == "param":
    # adjust parameter
    if act == "add":
      par = par+0.01
      print "lambda=", par
    if act == "remove":
      par = par-0.01
      print "lambda=", par
    if act == "init":
      par = 0.0
      print "lambda=", par
    if act == "ranrev":
      par = 1
  act = "" # action only takes place once, mode does not change

  if centralize:
    centralize = False
    center = fromscreen(pygame.mouse.get_pos())

  if reset:
    reset = False
    restart()

  # --- Drawing code should go here

  # First, clear the screen to white. Don't put other drawing commands
  # above this, or they will be erased with this command.
  screen.fill(BLACK)

  # dont draw too long lines (want curves!)

  MAXLENGTH = 50

  # screen.blit(background_image, [0, 0])
  for trc in qlist:
    n = len(trc)
    if n == 1:
      pygame.draw.circle(screen,WHITE,toscreen(trc[0]),3)
    if n > 1:
      for i in range(n-1):
        pt1 = toscreen(trc[i])
        pt2 = toscreen(trc[i+1])
        if abs(pt1[0]-pt2[0])+abs(pt1[1]-pt2[1]) < MAXLENGTH:
          pygame.draw.line(screen,WHITE,pt1,pt2,3)
  for pt in plist:
    drawpt = toscreen(pt)
    if not(drawpt[0] == bnd) and not(drawpt[1] == bnd): # remove hanmging points
      pygame.draw.circle(screen,RED,toscreen(pt),6)
  # draw curve
  n = len(curve)
  if n > 1:
    for i in range(n-1):      
      pt1 = toscreen(curve[i])
      pt2 = toscreen(curve[i+1])
      if abs(pt1[0]-pt2[0])+abs(pt1[1]-pt2[1]) < MAXLENGTH:
        pygame.draw.line(screen,YELLOW,pt1,pt2,3)

  # --- Go ahead and update the screen with what we've drawn.
  pygame.display.flip()

  # --- Limit to 60 frames per second
  clock.tick(60)

# Close the window and quit.
pygame.quit()