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
"""


# YOUR CODE HERE
# raise NotImplementedError()
plt.plot(t, W);
plt.xlabel('t')
plt.ylabel('W(t)');

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
"""


# YOUR CODE HERE
# raise NotImplementedError()
dW = np.diff(W)
print "mean;", dW.mean()
print "standard deviation;", 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
"""


def geo_brownian(t, W, X0, mu, sigma):
    "Return X(t) for geometric brownian motion with drift mu, volatility sigma."""
    # YOUR CODE HERE
#     raise NotImplementedError()
    a = (mu-0.5*sigma**2)*t + sigma*W
    b = sigma*W
    c = a + b
    e = np.exp(c)
    f = X0*e
    return f

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
"""


# YOUR CODE HERE
# raise NotImplementedError()
x = geo_brownian(t, W, 1.0, 0.5, 0.3)
plt.plot(t, x)
plt.xlabel('t')
plt.ylabel('x(t)');

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
"""


