import pygame
import math
import random
import numpy as np

def field(x,y):
  #return (-y,x) # rotation 
  #return (-y-0.1*x,x-0.1*y) # spiral
  #return ((x*x*25+y*y*25-10)/10,y*5*(1-x*x*25)/10) # 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 ) # Seitenwind gleich 0
  #return (-0.5*x+y,-0.5*y+x*x ) # Exercise 8/2
  #return (y+3*np.sin(10.0*x)/10,-x) # many limit circles
  #return (-x*y+0.0, x*x+y*y-1 ) # Seitenwind gleich 0

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

# 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 = 500
initspeed = 0.05
reset = False

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

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

def intb(n):
  if n > bnd:
    ans = bnd
  elif n < -bnd:
    ans = -bnd
  elif 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 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 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 iterate():
  global plist,qlist
  plist = [ move(pt) for pt in plist ]
  qlist = [ movetr(trc) for trc in qlist ]

# -------- Print menue
print "p       mode point"
print "t       mode trace"
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 "SPC     clear"
print "q       quit"

# -------- Initialize paraneters 
scale,center,speed = initscale,(0,0),initspeed # parameters 
restart()

# -------- 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_c:
        centralize = True
      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 == "size":
    resize() 
  if mode == "speed":
    drive() # adjust speed
    iterate() # move points and traces
  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)

  #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):
        pygame.draw.line(screen,WHITE,toscreen(trc[i]),toscreen(trc[i+1]),3)
  for pt in plist:
    pygame.draw.circle(screen,RED,toscreen(pt),6)

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