import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button

plt.rcParams.update({'font.size': 14})  # global font size

# Samuelson's multiplier-accelerator model:
#   Y_{h+1} - alpha(1+beta) Y_h + alpha*beta Y_{h-1} = 1

N = 60                                   # number of time steps
alpha_0, beta_0 = 0.5, 1.0               # initial slider values (parameters)
Y0_0, Y1_0 = 1.0, 1.0                    # initial slider values (initial conditions)


def evolve(alpha, beta, Y0, Y1):
    Y = [Y0, Y1]
    for _ in range(N - 1):
        Y.append(alpha * (1 + beta) * Y[-1] - alpha * beta * Y[-2] + 1)
    return np.array(Y)


def regime(alpha, beta):
    """Classify behaviour from the roots of  l^2 - alpha(1+beta) l + alpha*beta = 0."""
    roots = np.roots([1, -alpha * (1 + beta), alpha * beta])
    rho = np.max(np.abs(roots))
    kind = "complex roots (oscillatory)" if np.any(np.iscomplex(roots)) else "real roots (monotone)"
    stab = "stable" if rho < 1 else "unstable"
    return f"{stab}, {kind}, max|root| = {rho:.3f}"


def thresholds(beta):
    """Thresholds for alpha (they depend on beta only)."""
    stab = 1 / beta if beta > 0 else np.inf
    osc = 4 * beta / (1 + beta) ** 2
    return (f"Thresholds for α (given β = {beta:.3f}):\n"
            f"  stability:    α < 1/β = {stab:.3f}\n"
            f"  oscillations: α < 4β/(1+β)² = {osc:.3f}")


fig, ax = plt.subplots(figsize=(10, 8))
fig.subplots_adjust(left=0.12, bottom=0.43)
info = fig.text(0.02, 0.01, thresholds(beta_0), fontsize=13, va='bottom',
                family='monospace')

h = np.arange(N + 1)
[line] = ax.plot(h, evolve(alpha_0, beta_0, Y0_0, Y1_0), '.-')
eq_line = ax.axhline(1 / (1 - alpha_0), color='gray', ls='--', lw=1,
                     label='equilibrium 1/(1-α)')
ax.set_xlabel('h')
ax.set_ylabel(r'$Y_h$')
ax.legend(loc='upper right')

# Sliders: two parameters and two initial conditions
alpha_ax = fig.add_axes([0.2, 0.33, 0.7, 0.03])
beta_ax = fig.add_axes([0.2, 0.28, 0.7, 0.03])
Y0_ax = fig.add_axes([0.2, 0.23, 0.7, 0.03])
Y1_ax = fig.add_axes([0.2, 0.18, 0.7, 0.03])
alpha_slider = Slider(alpha_ax, r'$\alpha$', 0.01, 0.99, valinit=alpha_0)
beta_slider = Slider(beta_ax, r'$\beta$', 0.0, 5.0, valinit=beta_0)
Y0_slider = Slider(Y0_ax, r'$Y_0$', 0.0, 5.0, valinit=Y0_0)
Y1_slider = Slider(Y1_ax, r'$Y_1$', 0.0, 5.0, valinit=Y1_0)
sliders = [alpha_slider, beta_slider, Y0_slider, Y1_slider]


def on_changed(val):
    a, b = alpha_slider.val, beta_slider.val
    y = evolve(a, b, Y0_slider.val, Y1_slider.val)
    line.set_ydata(y)
    eq_line.set_ydata([1 / (1 - a)] * 2)
    ax.set_title(regime(a, b), fontsize=14)
    info.set_text(thresholds(b))
    # rescale the y-axis, since unstable cases grow quickly
    lo, hi = y.min(), y.max()
    pad = 0.1 * (hi - lo + 1e-9)
    ax.set_ylim(lo - pad, hi + pad)
    fig.canvas.draw_idle()


for s in sliders:
    s.on_changed(on_changed)
on_changed(None)

# Reset button
reset_ax = fig.add_axes([0.8, 0.04, 0.1, 0.05])
reset_button = Button(reset_ax, 'Reset', hovercolor='0.975')
reset_button.on_clicked(lambda event: [s.reset() for s in sliders])

plt.show()