import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt
import seaborn as sns

import antipackage
import github.ellisonbg.misc.vizarray as va

"""
Explanation: Numpy Exercise 3
Imports
End of explanation
"""


def brownian(maxt, n):
    """Return one realization of a Brownian (Wiener) process with n steps and a max time of t."""
    t = np.linspace(0.0,maxt,n)
    h = t[1]-t[0]
    Z = np.random.normal(0.0,1.0,n-1)
    dW = np.sqrt(h)*Z
    W = np.zeros(n)
    W[1:] = dW.cumsum()
    return t, W

"""
Explanation: Geometric Brownian motion
Here is a function that produces standard Brownian motion using NumPy. This is also known as a Wiener Process.
End of explanation
"""



t, W =brownian(1.0,1000)

assert isinstance(t, np.ndarray)
assert isinstance(W, np.ndarray)
assert t.dtype==np.dtype(float)
assert W.dtype==np.dtype(float)
assert len(t)==len(W)==1000

"""
Explanation: Call the brownian function to simulate a Wiener process with 1000 steps and max time of 1.0. Save the results as two arrays t and W.
End of explanation
"""


plt.plot(t, W)

assert True # this is for grading

"""
Explanation: Visualize the process using plt.plot with t on the x-axis and W(t) on the y-axis. Label your x and y axes.
End of explanation
"""


dW=np.diff(W)

print dW.mean()

print dW.std()

assert len(dW)==len(W)-1
assert dW.dtype==np.dtype(float)

"""
Explanation: Use np.diff to compute the changes at each step of the motion, dW, and then compute the mean and standard deviation of those differences.
End of explanation
"""


import math
def geo_brownian(t, W, X0, mu, sigma):
    "Return X(t) for geometric brownian motion with drift mu, volatility sigma."""
    
    a = (mu-(sigma**2)/2)*t +(sigma)*W
    b = X0*(np.exp(a))
    return (t, b)

assert True # leave this for grading

"""
Explanation: Write a function that takes $W(t)$ and converts it to geometric Brownian motion using the equation:
$$
X(t) = X_0 e^{((\mu - \sigma^2/2)t + \sigma W(t))}
$$
Use Numpy ufuncs and no loops in your function.
End of explanation
"""


q, r = geo_brownian(t, W,1.0,0.5,0.3)

plt.plot(q, r)

assert True # leave this for grading

"""
Explanation: Use your function to simulate geometric brownian motion, $X(t)$ for $X_0=1.0$, $\mu=0.5$ and $\sigma=0.3$ with the Wiener process you computed above.
Visualize the process using plt.plot with t on the x-axis and X(t) on the y-axis. Label your x and y axes.
End of explanation
"""


