%matplotlib inline
from matplotlib import pyplot as plt
import numpy as np

from IPython.html.widgets import interact, interactive, fixed
from IPython.display import display

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


def soliton(x, t, c, a):
    """Return phi(x, t) for a soliton wave with constants c and a."""
    return (0.5)*c*((1/np.cosh(((np.sqrt(c)*0.5)*(x-c*t-a)))**2))

assert np.allclose(soliton(np.array([0]),0.0,1.0,0.0), np.array([0.5]))

"""
Explanation: Using interact for animation with data
A soliton is a constant velocity wave that maintains its shape as it propagates. They arise from non-linear wave equations, such has the Korteweg–de Vries equation, which has the following analytical solution:
$$
\phi(x,t) = \frac{1}{2} c \mathrm{sech}^2 \left[ \frac{\sqrt{c}}{2} \left(x - ct - a \right) \right]
$$
The constant c is the velocity and the constant a is the initial location of the soliton.
Define soliton(x, t, c, a) function that computes the value of the soliton wave for the given arguments. Your function should work when the postion x or t are NumPy arrays, in which case it should return a NumPy array itself.
End of explanation
"""


tmin = 0.0
tmax = 10.0
tpoints = 100
t = np.linspace(tmin, tmax, tpoints)

xmin = 0.0
xmax = 10.0
xpoints = 200
x = np.linspace(xmin, xmax, xpoints)

c = 1.0
a = 0.0

"""
Explanation: To create an animation of a soliton propagating in time, we are going to precompute the soliton data and store it in a 2d array. To set this up, we create the following variables and arrays:
End of explanation
"""


phi=np.zeros((xpoints,tpoints)) #worked with Hunter Thomas
for i in x:
    for j in t:
        phi[i,j]=soliton(x[i],t[j],c,a)

assert phi.shape==(xpoints, tpoints)
assert phi.ndim==2
assert phi.dtype==np.dtype(float)
assert phi[0,0]==soliton(x[0],t[0],c,a)

"""
Explanation: Compute a 2d NumPy array called phi:

It should have a dtype of float.
It should have a shape of (xpoints, tpoints).
phi[i,j] should contain the value $\phi(x[i],t[j])$.
End of explanation
"""


def plot_soliton_data(i=0):
    """Plot the soliton data at t[i] versus x."""
    plt.plot(soliton(x,t[i],c,a))
    plt.xlabel('Time')
    plt.ylabel('Phi')
    plt.title('Solition wave vs. Time')
    plt.tick_params(axis='x', top='off', direction='out')
    plt.tick_params(axis='y', right='off', direction='out')

plot_soliton_data(0)

assert True # leave this for grading the plot_soliton_data function

"""
Explanation: Write a plot_soliton_data(i) function that plots the soliton wave $\phi(x, t[i])$. Customize your plot to make it effective and beautiful.
End of explanation
"""


interact(plot_soliton_data, i=(0.0,50.0,0.1));

assert True # leave this for grading the interact with plot_soliton_data cell

"""
Explanation: Use interact to animate the plot_soliton_data function versus time.
End of explanation
"""


