"""
boids.py 

Implementacja boidów Craiga Reynolda

Autor: Mahesh Venkitachalam
"""

import sys, argparse
import math
import numpy as np
import matplotlib.pyplot as plt 
import matplotlib.animation as animation
from scipy.spatial.distance import squareform, pdist, cdist
from numpy.linalg import norm

width, height = 640, 480

class Boids:
    """klasa reprezentująca symulację algorytmu stada"""
    def __init__(self, N):
        """inicjowanie symulacji algorytmu stada"""
        # inicjowanie pozycji i prędkości
        self.pos = [width/2.0, height/2.0] + 10*np.random.rand(2*N).reshape(N, 2)
        # znormalizowane prędkości losowej
        angles = 2*math.pi*np.random.rand(N)
        self.vel = np.array(list(zip(np.sin(angles), np.cos(angles))))
        self.N = N
        # minimalny dystans podejścia
        self.minDist = 25.0
        # maksymalna wielkość prędkości obliczanych według "reguł"
        self.maxRuleVel = 0.03
        # maksymalna wielkość finalnej prędkości
        self.maxVel = 2.0

    def tick(self, frameNum, pts, beak):
        """aktualizacja symulacji o jeden krok czasowy"""
        # uzyskanie odległości między parami punktów
        self.distMatrix = squareform(pdist(self.pos))
        # zastosowanie reguł
        self.vel += self.applyRules()
        self.limit(self.vel, self.maxVel)
        self.pos += self.vel
        self.applyBC()
        # aktualizowanie danych
        pts.set_data(self.pos.reshape(2*self.N)[::2], 
                     self.pos.reshape(2*self.N)[1::2])
        vec = self.pos + 10*self.vel/self.maxVel
        beak.set_data(vec.reshape(2*self.N)[::2], 
                      vec.reshape(2*self.N)[1::2])

    def limitVec(self, vec, maxVal):
        """ograniczanie długości wektora 2D"""
        mag = norm(vec)
        if mag > maxVal:
            vec[0], vec[1] = vec[0]*maxVal/mag, vec[1]*maxVal/mag
    
    def limit(self, X, maxVal):
        """ograniczanie długości wektorów 2D w tablicy X do wartości maksymalnej"""
        for vec in X:
            self.limitVec(vec, maxVal)
            
    def applyBC(self):
        """zastosowanie warunków brzegowych"""
        deltaR = 2.0
        for coord in self.pos:
            if coord[0] > width + deltaR:
                coord[0] = - deltaR
            if coord[0] < - deltaR:
                coord[0] = width + deltaR    
            if coord[1] > height + deltaR:
                coord[1] = - deltaR
            if coord[1] < - deltaR:
                coord[1] = height + deltaR
    
    def applyRules(self):
        # zastosowanie reguły #1 — rozdzielność
        D = self.distMatrix < 25.0
        vel = self.pos*D.sum(axis=1).reshape(self.N, 1) - D.dot(self.pos)
        self.limit(vel, self.maxRuleVel)

        # inny próg odległości
        D = self.distMatrix < 50.0

        # zastosowanie reguły #2 — wyrównanie
        vel2 = D.dot(self.vel)
        self.limit(vel2, self.maxRuleVel)
        vel += vel2;

        # zastosowanie reguły #1 — spójność
        vel3 = D.dot(self.pos) - self.pos
        self.limit(vel3, self.maxRuleVel)
        vel += vel3

        return vel

    def buttonPress(self, event):
        """procedura obsługi biblioteki matplotlib dla wciskania przycisków"""
        # kliknięcie lewym przyciskiem myszki — dodanie boida
        if event.button is 1:
            self.pos = np.concatenate((self.pos, 
                                       np.array([[event.xdata, event.ydata]])), 
                                      axis=0)
            # losowa prędkość
            angles = 2*math.pi*np.random.rand(1)
            v = np.array(list(zip(np.sin(angles), np.cos(angles))))
            self.vel = np.concatenate((self.vel, v), axis=0)
            self.N += 1 
        # kliknięcie prawym przyciskiem myszki — rozproszenie
        elif event.button is 3:
            # dodanie prędkości rozproszenia 
            self.vel += 0.1*(self.pos - np.array([[event.xdata, event.ydata]]))
        
def tick(frameNum, pts, beak, boids):
    # wyświetlanie frameNum
    """aktualizacja funkcji do animacji"""
    boids.tick(frameNum, pts, beak)
    return pts, beak

# funkcja main()
def main():
  # użycie w razie potrzeby sys.argv
  print('uruchamianie boidów...')

  parser = argparse.ArgumentParser(description="Implementacja algorytmu stada Craiga Reynoldsa...")
  # dodanie argumentów
  parser.add_argument('--num-boids', dest='N', required=False)
  args = parser.parse_args()

  # liczba boidów
  N = 100
  if args.N:
      N = int(args.N)

  # tworzenie boidów
  boids = Boids(N)

  # konfigurowanie rysunku
  fig = plt.figure()
  ax = plt.axes(xlim=(0, width), ylim=(0, height))

  pts, = ax.plot([], [], markersize=10, 
                  c='k', marker='o', ls='None')
  beak, = ax.plot([], [], markersize=4, 
                  c='r', marker='o', ls='None')
  anim = animation.FuncAnimation(fig, tick, fargs=(pts, beak, boids), 
                                 interval=50)

  # dodanie procedury obsługi zdarzenia wciśnięcia przycisku
  cid = fig.canvas.mpl_connect('button_press_event', boids.buttonPress)

  plt.show()

# wywołanie main
if __name__ == '__main__':
  main()
