#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Jul 12 11:36:30 2022

@author: adomps
"""

#########################################################################
#### Marche aléatoire 1d depuis l'origine
########################################################################


##  L'unité de temps est de tau, durée de vol.
#  Un pas de pas correspond à la durée tau. 
#  L'unité de longueur est la libre parcours moyen.
#

# Attention
# rd.randint(a, b) : a et b inclus
# np.random.randint(a, n) : de a jusqu'à b EXCLU !

# Attention
# Pour t pair, toutes les particules se trouvent sur des abscisses
# pair, et symétriquement pour t impair.
# Cela introduit des fluctiations artificielles dans les populations,
# que l'on pourrait gommer en lissant les fréquences
# frequences[i] = 0.5 * frequences[i] + 0.25 *frequences[i-1]+ 0.25 * frequences[i:+1]
# mais je ne l'ai pas fait.


D = 1**2 / (2 * 1)



import numpy as np
import matplotlib.pyplot as plt
import random as rd
import time




# Une seule particule
# On fait N_pas, donc on a N_pas + 1 position en comptant le départ.
def marche_alea(N_pas, option) :
    ''' Pour une seule particule, renvoie le tableau des abscisses
    successives occupées après N_pas pas aléatoires à droite ou à gauche. 
    Deux version : la première naïve, la seconde bien plus rapide avec numpy.'''
    positions = np.zeros(N_pas+1, dtype = int)
    
    if option == 0 :
        for i in range(N_pas) :
            positions[i+1] = positions[i] + 2 * rd.randint(0, 1) - 1
            
    if option == 1 :
            deplacements = 2 * np.random.randint(0, 2, size = N_pas) - 1
            positions[1:] =  np.cumsum(deplacements)
            
    return positions

def trace_marche(positions, numfig) :
    plt.figure(numfig)
    plt.title('Marche aléatoire 1 d')
    plt.plot(t, positions)   
    plt.xlabel('t')
    plt.ylabel('abscisse X')     
    
    
N_pas = 1000
t = np.array([i for i in range(N_pas+1)])

t1 = time.time()
pos = marche_alea(N_pas, 0)
t2 = time.time()
print('durée', t2 - t1)
trace_marche(pos, 0)
    

t1 = time.time()
pos = marche_alea(N_pas, 1)
t2 = time.time()
print('durée', t2 - t1)
trace_marche(pos, 0)



# Mouvement d'un ensemble de particules. 
# On stocke dans un unique tableau les positions des
# différentes particules aux différents instants.
# Ce tableau positions est tel que positions[i, j]
# est l'abscisse à i instant de la  particule j. 
# La colonne j donne donc les abscisses successives de la particule j.
# La ligne i donne les abscisses de toutes les particules à l'instant i.

# Pour obtenir un temps de calcul raisonnable avec beaucoup de particules,
# il faut choisir la version rapide de marche_alea

#Version avec boucle
def marche_ensemble(N_pas, N_part) :
    ''' Renvoie les positions de toutes les particules à chaque pas de temps,
    sous forme de tableau positions[i, j]. i : instant, j : particule. '''
    positions = np.zeros((N_pas+1, N_part))
    for j in range(N_part) :
        positions[:, j] = marche_alea(N_pas, 1)
    return positions





# Version vectorisée avec tableaux de numpy.
# Au lieu d'utiliser la fonction marche_alea pour une seule particule,
# on tire directement un tableau de déplacements 2d.
# Légèrement plus rapide que la version précédente.
def marche_ensemble_bis(N_pas, N_part) :
    positions = np.zeros((N_pas+1, N_part)) # 4 * np.pi * D * instant
    deplacements = 2 * np.random.randint(0, 2, size = (N_pas, N_part)) - 1
    positions[1:] = np.cumsum(deplacements, axis = 0)
    return positions


# 
N_part = 20000
t1 = time.time()
positions = marche_ensemble(N_pas, N_part)
t2 = time.time()
print('durée', t2 - t1)


# extraction d'une particule au hasard
ind_part = np.random.randint(N_part)
pos = positions[:, ind_part]
trace_marche(pos, 0)

# x^2 moyen. On moyenne sur l'ensemble des particules
x2_moy = np.average(positions**2, axis = 1)

plt.figure(1)
plt.clf()
plt.plot(t[::20], x2_moy[::20], label = 'moyenne empirique', marker = '.', color = 'blue', linestyle = '')
plt.plot(2 * D * t, label  = '$<x^2> = 2Dt$', color = 'black' )
plt.xlabel('t')
plt.ylabel('$<x^2>$')
plt.legend()



# Répartition des particules à un instant donné

def distribution_particules(positions, instant, N_boites, x_min, x_max, numfig) :
    ''' Calcule et représente la distribution des particules à un instant donné.'''
    positions_instant = positions[instant, :]
    largeur_boites = (x_max - x_min) / N_boites
    
    plt.figure(numfig)
    plt.xlabel('x')
    plt.ylabel('population')    
    bords_boites = [x_min + largeur_boites * i for i in range(N_boites+1)]
    plt.hist(positions_instant, bins = bords_boites)  # tracé automatique de l'histogramme
    
    popu, bds_boites = np.histogram(positions_instant, bins = bords_boites, density = False)
    x_centres = (bds_boites[:-1] + bds_boites[1:])/2 # centres des boites

    return popu, x_centres

def trace_popu_theorique(instant, x_centres, largeur_boites) :
    if instant == 0 :
        return
    n_theo = N_part / np.sqrt(4 * np.pi * D * instant) * \
                np.exp(-x_centres**2 / (4 * D * instant)) * largeur_boites
    plt.plot(x_centres, n_theo, color = 'red', label = '$n(x,t) \Delta x$')


t_choisi = 501 # choisir une date

x_min = - 5 * np.sqrt(D*t_choisi)
x_max = + 5 * np.sqrt(D*t_choisi)
N_boites = 41 # prendre un nombre impair pour avoir une boite centrée en 0.
plt.figure(0)
plt.clf()
popu, x =  distribution_particules(positions, t_choisi, N_boites, x_min, x_max, 0)
trace_popu_theorique(t_choisi, x, (x_max - x_min) / N_boites)

##############################################################
### Juste pour moi, pour préparer mon cours. illustration de
# l'hitogramme au dessus du  << nuage >> de particules
################################################################
x_min = - 5 * np.sqrt(D*t_choisi)
x_max = + 5 * np.sqrt(D*t_choisi)
N_boites = 21 # prendre un nombre impair pour avoir une boite centrée en 0.
plt.figure(0)
plt.clf()
plt.subplot(2, 1, 1)
popu, x =  distribution_particules(positions, t_choisi, N_boites, x_min, x_max, 0)
plt.subplot(2, 1 , 2)
largeur_boites = (x_max - x_min) / N_boites
plt.yticks([])
bords_boites = [x_min + largeur_boites * i for i in range(N_boites+1)]
plt.scatter(positions[t_choisi], [rd.random() for i in range(N_part)]) # étalement fictif selon y
for xb in bords_boites :
    plt.vlines(xb, -0.1, 1.1, linestyle = 'dashed')        
plt.xlabel('x')



## visulation à divers instants pour voir l'étalement
positions = marche_ensemble(N_pas, N_part)
dates_choisies = [0, 200, 400, 600, 800, 1000]
N_boites = 21
N_dates = len(dates_choisies)
x_max = 4 * np.sqrt(D * N_pas)
x_min = - x_max
y_max = N_part /4
plt.figure(3)
plt.clf()
for i, t_choisi in enumerate(dates_choisies) :
    plt.subplot(N_dates, 1, i+1)
    plt.ylim(0, y_max)
    if i == 0 : plt.ylim(0, N_part)
    plt.text(0.9*x_min, y_max / 2, 't = ' + str(t_choisi))
    popu, x = distribution_particules(positions, t_choisi, N_boites, x_min, x_max, 3)  
    trace_popu_theorique(t_choisi, x, (x_max - x_min) / N_boites)    
    plt.xticks([-100, -75, -50, -25, 0, 25, 50, 75, 100], labels = '')     
    ax = plt.gca()
    ax.set_xlabel('')
    ax.yaxis.set_ticks_position('right')
    #plt.legend()
    #Xaxes.set_visible(False)  # efface tout !  
# On légende l'axe pour la figure du bas 
plt.xticks([-100, -75, -50, -25, 0, 25, 50, 75, 100], 
           ['-100', '-75', '-50','-25', '0', '25', '50', '75', '100']) 
ax.set_xlabel('x')