#!/usr/bin/env python
# coding: utf-8

# Testing 1-D random walk turbulence jittering for vertical particle distribution

# In[19]:


# imports modules and renames them short names
import numpy as np
import os
from matplotlib import pyplot as plt, animation
from datetime import datetime, timedelta
from dateutil.parser import parse
from IPython.display import HTML
from scipy.stats import skewnorm, gamma
import seaborn as sns

# makes the plots show up below the code cells
get_ipython().run_line_magic('matplotlib', 'inline')


# In[3]:


# Sets the default font size of matplotlib plots
plt.rcParams['font.size'] = 12


# Create a 1D random walk turbulence model to jitter particles in the vertical

# In[4]:


# Define parameters for the walk

dims = 1 # just 1 dimension, in my case it represents depth
step_n = 10000 # number of timesteps - when implementing this in parcels, divide your "duration" (days) variable by your timestep variable "dt" 
#(seconds) to get the total number of timesteps. It would look like this: (duration*24*60*60)/dt
step_set = [-1, 0, 1] # these are the random step choices, the particles can either go backwards (-1 increment, could be up in my case if depth
# variable is positive for increasing depth), not move at all (0), or go forwards (1 increment). Need to define increments somehow
origin = np.zeros((1,dims)) # this creates a new array, of x dimensions, filled with zeroes. The first 1 I think is basically the number of particles
# if I want to assign an initial depth distribution the array I would create would not be filled with zeroes, but would be all the particles
# and their starting depths


# In[8]:


# Simulate steps in 1D

step_shape = (step_n,dims)
steps = np.random.choice(a=step_set, size=step_shape)
path = np.concatenate([origin, steps]).cumsum(0)
start = path[:1]
stop = path[-1:]


# In[14]:


print()


# In[16]:


# Plot the path

fig = plt.figure(figsize=(8,4),dpi=200)
ax = fig.add_subplot(111)
#ax.scatter(np.arange(step_n+1), path, c='blue', alpha=0.5, s=0.05);
ax.plot(path, c='blue', alpha=0.5, lw=0.5, ls='-',);
ax.plot(0, start, c='red', marker='+')
ax.plot(step_n, stop, c='black', marker='o')

# The red + is the particles start, black circle is end. Blue line is its path along the 1 dimension over the timesteps


# In[12]:


fig, ax = plt.subplots(figsize=(4, 4))
l = ax.plot(path,  c='black', marker='o')


# Try to animate this

# In[ ]:


get_ipython().run_cell_magic('capture', '', "\nfig, ax = plt.subplots(figsize=(4, 4))\nl = ax.plot(path,  c='black', marker='o')\n\n# Init function\ndef init():\n    l.set_offsets(np.empty((0, 2)))\n    l.set_array(np.empty(0))\n    return l,\n\n# Animate function\ndef animate(hour):\n    l.set_offsets(np.vstack([data.lon[:, hour], data.z[:, hour]]).T)\n    l.set_array(data.z[:, hour])\n    return l,\n\n# Build animation\nanim = animation.FuncAnimation(fig, animate, init_func=init, frames=73, interval=100, blit=True)\n")


# In[ ]:


# Render animation
HTML(anim.to_html5_video())


# Create skewed Gaussian distribution modelled after my pteropod vertical distributions

# In[17]:


# create some random data from a skewnorm
data = skewnorm.rvs(5, loc=10, scale=70, size=1000).astype(np.int_)

# draw a histogram and kde of the given data
ax = sns.distplot(data, kde_kws={'label':'kde of given data'}, label='histogram')

# find parameters to fit a skewnorm to the data
params = skewnorm.fit(data, 5, loc=10, scale=70)

# draw the pdf of the fitted skewnorm
x = np.linspace(0, 255, 500)
ax.plot(x, skewnorm.pdf(x, *params), label='approximated skewnorm')
plt.legend()
plt.show()


# In[41]:


print(params)


# In[30]:


# Calculate the first four moments:
a = 2 # a is the input the scipy.stats.gamma takes as the shape parameter
#mean, var, skew, kurt = gamma.stats(a, moments='mvsk')


# In[33]:


#Display the probability density function (pdf):
fig, ax = plt.subplots(1, 1)
x = np.linspace(gamma.ppf(0, a), gamma.ppf(300, a), 1000)
ax.plot(x, gamma.pdf(x, a, loc = 10, scale = 100), 'r-', lw=5, alpha=0.6, label='gamma pdf')