import numpy as np
import matplotlib.pyplot as plt
from IPython.html import widgets
from IPython.display import display

lat = widgets.BoundedFloatTextWidget(value=-32, min=-90, max=90, description = 'Latitude:')
fdt = widgets.FloatTextWidget(value=0.05, description = 'Time step:')  # (To test convergence decrease the value.)

display(lat)
display(fdt)

# Coriolis.
twopi = 2*np.pi
f = 2 * twopi / (24 * 3600) * np.cos(np.pi / 180 * (lat.value))

# Inertial period.
TP = twopi / f
dt = fdt.value / f

# Sampling of the discrete solution every dtsol second.
dtsol = 1 * dt

# Intitial stuff.
t0 = 0
u0 = 0
v0 = 1

# Simulation time (note the time is used, not the number of steps).
tf = 10 * TP

def numcorio(f, dt, t0, tf, u0, v0, dtsol, alpha=0):
    """Function which discretises the inertial oscillation.
    If other equations need to be discretised, this can be done
    by adapting numcorio.py (see Numerical Exercises 2.7 and 2.8)."""
    # Initial condition.
    u, v, t = [], [], []
    u.append(u0)
    v.append(v0)
    t.append(t0)

    # Starting iterations.
    un = u0
    vn = v0
    time = t0

    # Parameters appearing in the discretisations.
    umafdt = (1 - alpha) * f * dt
    afdt = alpha * f * dt
    det1 = 1 / (1 + afdt ** 2)

    n = 0
    # Extract solution every dtsol/dt time steps.
    nsol = dtsol / dt
    isol = 0

    # Loop until final time tf is reached.
    while time < tf:
        # Right hand side of discretisation.
        ru = un + umafdt * vn
        rv = vn - umafdt * un
        # Solve linear system for the new velocities unp and vnp.
        unp = (ru + afdt * rv) * det1
        vnp = (rv - afdt * ru) * det1
        # Update time.
        time = time + dt
        # Call old value the values you just calculated.
        un = unp
        vn = vnp
        # Increase counter of time steps.
        n += 1
        # Save solution every nsol steps.
        if np.mod(n, nsol) == 0:
            isol += 1
            # print(isol)
            t.append(time)
            v.append(vn)
            u.append(un)

    return map(np.array, (u, v, t))

# Explicit.
ue, ve, te = numcorio(f, dt, t0, tf, u0, v0, dtsol, alpha=0)

# Implicit.
ui, vi, ti = numcorio(f, dt, t0, tf, u0, v0, dtsol, alpha=1)

# Trapezoidal.
ut, vt, tt = numcorio(f, dt, t0, tf, u0, v0, dtsol, alpha=0.5)

fig, ax = plt.subplots(figsize=(4, 4))
ax.plot(ue, ve, ':', label=r'Explicity ($\alpha$ = 0)')
ax.plot(ui, vi, '.', label=r'Implicity ($\alpha$ = 1)')
ax.plot(ut, vt, '-', label=r'Trapezoidal ($\alpha$ = 0.5)')
ax.set_xlabel('u')
ax.set_ylabel('v')
ax.legend()
_ = ax.axis([-2, 2, -2, 2])

HTML(html)