#----------------------------------------------------------------------
#  Script prova metodo SPH
#  Autore: Andrea Pavan
#  Data: 27/09/2022
#  Ref: https://www.youtube.com/watch?v=-0m05gzk8nk
#----------------------------------------------------------------------


import numpy as np;
import matplotlib.pyplot as plt;
from sklearn import neighbors;
from tqdm import tqdm;


#variabili
nMaxParticles = 125;
domainWidth = 40;
domainHeight = 80;
particleMass = 1;
isotropicExponent = 20;
baseDensity = 1;
smoothingLength = 5;
mu = 0.5;
wallDampingCoefficient = -0.9;
g = np.array([[0.0,-0.1]]);
dt = 0.01;
nTimesteps = 2500;
addedParticles = 50;
#plotSize = (4,6);
plotUpdate = 6;
pointSize = 2000;

domainXLim = np.array([smoothingLength, domainWidth-smoothingLength]);
domainYLim = np.array([smoothingLength, domainHeight-smoothingLength]);
krho = 315*particleMass/(64*np.pi*smoothingLength**9);
kpF = -45*particleMass/(5*np.pi*smoothingLength**6);
kvF = 45*mu*particleMass/(np.pi*smoothingLength**6);

#plt.style.use("dark_background");
#plt.figure(figsize=plotSize,dpi=160);
plt.figure();

#iterazioni
nParticles = 1;
x = np.zeros((nParticles,2));
v = np.zeros((nParticles,2));
F = np.zeros((nParticles,2));
for iter in tqdm(range(nTimesteps)):
    #aggiungo particelle
    if iter%addedParticles==0 and nParticles<=nMaxParticles:
        xnew = np.array([[10+np.random.rand(), domainYLim[1]],
                         [20+np.random.rand(), domainYLim[1]]]);
        vnew = np.array([[-3, -15],
                         [-3, +15]]);
        nParticles+=2;
        x = np.concatenate((x,xnew),axis=0);
        v = np.concatenate((v,vnew),axis=0);

    #calcolo distanze particelle vicine
    satIdx,satd = neighbors.KDTree(x).query_radius(x,smoothingLength,return_distance=True,sort_results=True);

    #calcolo densità
    rho = np.zeros(nParticles);
    for i in range(nParticles):
        jidx = 0;
        for j in satIdx[i]:
            rho[j] += krho*(smoothingLength**2 - satd[i][jidx]**2)**3;
            jidx+=1;
    satIdx = [np.delete(satItself,0) for satItself in satIdx];
    satd = [np.delete(xitself,0) for xitself in satd];
        
    #calcolo pressione
    p = isotropicExponent*(rho-baseDensity);
    F = np.zeros((nParticles,2));
    for i in range(nParticles):
        jidx = 0;
        for j in satIdx[i]:
            #forze pressione
            F[i] += -kpF*((x[j]-x[i])/satd[i][jidx])*((p[j]+p[i])/(2*rho[j]))*(smoothingLength-satd[i][jidx])**2;
            #forze viscose
            F[i] += kvF*((v[j]-v[i])/rho[j])*(smoothingLength-satd[i][jidx]);
            jidx+=1;
    #forza gravità
    F += g;
    
    #integrazione nel tempo
    v += dt*F/rho[:,np.newaxis];
    x += dt*v;
    
    #condizioni contorno parete
    outLeftParticles = x[:,0]<domainXLim[0];
    outRightParticles = x[:,0]>domainXLim[1];
    outBottomParticles = x[:,1]<domainYLim[0];
    outTopParticles = x[:,1]>domainYLim[1];
    v[outLeftParticles,0] *= wallDampingCoefficient;
    x[outLeftParticles,0] = domainXLim[0];
    v[outRightParticles,0] *= wallDampingCoefficient;
    x[outRightParticles,0] = domainXLim[1];
    v[outBottomParticles,1] *= wallDampingCoefficient;
    x[outBottomParticles,1] = domainYLim[0];
    v[outTopParticles,1] *= wallDampingCoefficient;
    x[outTopParticles,1] = domainYLim[1];
    
    #visualizzazione
    if iter%plotUpdate==0:
        #plt.scatter(x[:,0],x[:,1],s=pointSize,c=x[:,1],cmap="Wistia_r");
        #plt.xlim(domainXLim);
        #plt.ylim(domainYLim);
        #plt.xticks([],[]);
        #plt.yticks([],[]);
        #plt.tight_layout();
        #plt.draw();
        #plt.pause(0.0001);
        #plt.clf();
        plt.scatter(x[:,0],x[:,1]);
        plt.xlim(domainXLim);
        plt.ylim(domainYLim);
        plt.gca().set_aspect('equal',adjustable='box');
        plt.draw();
        plt.pause(0.001);
        plt.clf();
    