#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jan 20 23:08:15 2022

@author: maximebouchereau
"""

import numpy as np
import statistics as stat
import matplotlib.pyplot as plt
from matplotlib import animation
import random as rd
import sys

def D(X,i,j):
    """
    Computation of the euclidian distance between two individuals i and j, X is ths collection of the positions
    of all the individuals

    Parameters
    ----------
    X : TYPE: array of shape (2,I) where I is the number of individuals
        DESCRIPTION: i-th column of X is the position of individual i
    i : TYPE: integer between 0 and I-1
        DESCRIPTION: First individual
    j : TYPE: integer between O and I-1
        DESCRIPTION: Second individual

    Returns a float: the euclidian norm of the array X[:,i]-X[:,j] of shape (2,1)
    -------
    None.

    """
    
    return np.linalg.norm(X[:,i]-X[:,j])


def dist(i,delta,X):
    """
    Searches individuals which are at distance of individual i lower than delta

    Parameters
    ----------
    i : TYPE: integer between 0 and I-1
        DESCRIPTION: i-th column of X is the position of individual i
    delta : TYPE: float (positive)
        DESCRIPTION: Maximal distance of individual i where others individuals are searched
    X : TYPE: array of shape (2,I) where I is the nuber of individuals
        DESCRIPTION: i-th colum of X is the position of individual i

    Returns a list of integers: the individuals which are at a distance of i less than delta
    -------
    None.

    """
    
    I = X.shape[1] # Number of individuals
    L = [delta - D(X,i,j) for j in range(0,I)] # The indices i where L[i] > 0 correspond to the indivuduals searched
    
    return list(np.where(np.array(L) > 0)[0]) 

def Traj(I,v,T,N,delta):
    """
    Compute the trajectory of individuals which follow the Vicsek model

    Parameters
    ----------
    I : TYPE: integer bigger than 2
        DESCRIPTION: Number of individuals
    v : TYPE: float (positive)
        DESCRIPTION: Average speed of moving of individuals
    T : TYPE: float (positive)
        DESCRIPTION: Duration of the simulation
    N : TYPE: integer
        DESCRIPTION: Number of time steps for the simulation
    delta : TYPE: float (positive)
        DESCRIPTION: Maximal distance where two individual will align

    Returns X,Theta: arrays of shape (2,I,N) and (1,I,N)
    X[:,i,n] is the position of individual i at step n and Theta[:,i,n] is the sense of direction of individual i at step n
    
    -------
    None.

    """
    
    dt = T/N # Time step of the simulation
    
    X = np.zeros((N+1,2,I))
    Theta = np.zeros((N+1,1,I))
    
    X[0,:,:] = np.random.uniform(0,1,(2,I)) # Random initial position of the individuals
    Theta[0,:,:] = np.random.uniform(0,2*np.pi,(1,I)) # Random initial sense of direction of the individuals
    
    print("Computation of the trajectories...")
    for n in range(N):
        sys.stdout.write("\r%d   "% int(100*(n+1)/N)+"%")
        sys.stdout.flush()
        for i in range(I):
            Theta[n+1,0,i] = stat.mean([Theta[n,0,j] for j in dist(i,delta,X[n,:,:])])
        X[n+1,:,:] = X[n,:,:] + v*dt*np.concatenate((np.cos(Theta[n,:,:]),np.sin(Theta[n,:,:])),axis=0).reshape(2,I)
        X[n+1,:,:] = X[n+1,:,:] - np.floor(X[n+1,:,:])
    
    return X, Theta
    

def Simul(I,v,T,N,delta):
    """
    Print the simulation of the Vicsek model

    Parameters
    ----------
    I : TYPE: integer bigger than 2
        DESCRIPTION: Number of individuals
    v : TYPE: float (positive)
        DESCRIPTION: Average speed of moving of individuals
    T : TYPE: float (positive)
        DESCRIPTION: Duration of the simulation
    N : TYPE: integer
        DESCRIPTION: Number of time steps for the simulation
    delta : TYPE: float (positive)
        DESCRIPTION: Maximal distance where two individual will align

    Returns a video of the simulation
    -------
    None.

    """
    
    X,Theta = Traj(I,v,T,N,delta)
    
    deltat = 15000/N # Duration between two frames (ms)

    #plt.title("Evolution of the individuals")
    #plt.xlabel("x")
    #plt.ylabel("y")
    
    Fx = X[0,0,:].T
    Fy = X[0,1,:].T
    U = 0.5*np.cos(Theta[0,:,:]).T
    V = 0.5*np.sin(Theta[0,:,:]).T
    
    Pos = np.concatenate((Fx,Fy),axis=0)
    
    fig, ax = plt.subplots(1,1)
    Q = ax.quiver(Fx, Fy, U, V, pivot='mid', color='k', units='inches')
    
    
    def animate(n,Q,Pos):
        U = 0.5*np.cos(Theta[n,:,:]).T
        V = 0.5*np.sin(Theta[n,:,:]).T
        Pos = X[n,:,:].T
        
        Q.set_offsets(Pos)
        Q.set_UVC(U,V)

    return Q
    
    #anim = animation.FuncAnimation(fig, animate, frames=N, fargs=(Q,Pos), blit=True, interval=deltat, repeat=True)
    anim = animation.FuncAnimation(fig, animate, fargs=(Q,Pos),interval=deltat, blit=False)
    fig.tight_layout()
    plt.show()
    #%matplotlib qt
    
    pass
    
    
N_simul = 200

X,Theta = Traj(I=200,v=0.02,T=200,N=N_simul,delta=0.03)
    
deltat = 15000/N_simul # Duration between two frames (ms)

#plt.title("Evolution of the individuals")
#plt.xlabel("x")
#plt.ylabel("y")

Fx = X[0,0,:].T
Fy = X[0,1,:].T
U = 0.05*np.cos(Theta[0,:,:]).T
V = 0.05*np.sin(Theta[0,:,:]).T
Pos = np.concatenate((Fx,Fy),axis=0)

fig, ax = plt.subplots(1,1)
Q = ax.quiver(Fx, Fy, U, V, pivot='mid', color='k', units='inches')


def animate(n,Q,Pos):
    U = 0.05*np.cos(Theta[n,:,:]).T
    V = 0.05*np.sin(Theta[n,:,:]).T
    Pos = X[n,:,:].T
    
    Q.set_offsets(Pos)
    Q.set_UVC(U,V)

    return Q

#anim = animation.FuncAnimation(fig, animate, frames=N, fargs=(Q,Pos), blit=True, interval=deltat, repeat=True)
anim = animation.FuncAnimation(fig, animate, fargs=(Q,Pos),interval=deltat, blit=False, frames=N_simul)
anim.save("sample.gif", writer="pillow")
fig.tight_layout()
plt.show()