"""Six reproducible animations. Run in Jupyter; requires numpy, matplotlib, IPython."""
import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML, display


def taylor_sin(x, degree, center=0.0):
    """Degree means polynomial degree, not number of nonzero sine terms."""
    x = np.asarray(x, dtype=np.float64)
    result = np.zeros_like(x)
    power = np.ones_like(x)
    derivatives = [np.sin(center), np.cos(center), -np.sin(center), -np.cos(center)]
    for k in range(degree + 1):
        result += derivatives[k % 4] * power
        power *= (x - center) / (k + 1)
    return result


def newton_states():
    states = []
    x = 3.0
    for k in range(6):
        next_x = x - (x*x - 2)/(2*x)
        states.append(dict(k=k, x=x, next=next_x, error=abs(x*x-2)))
        x = next_x
    return states


def objective(p):
    x, y = p
    return (1-x)**2 + 100*(y-x*x)**2


def gradient(p):
    x, y = p
    return np.array([-2*(1-x)-400*x*(y-x*x), 200*(y-x*x)])


def newton_opt_states():
    """Newton applied to f'(x)=0 for f(x)=x^4/4+x^2/2."""
    states = []
    x = 2.5
    for k in range(9):
        value, grad, hess = x**4/4 + x*x/2, x**3+x, 3*x*x+1
        next_x = x-grad/hess
        states.append(dict(k=k, x=x, next=next_x, value=value, grad=grad, hess=hess))
        if abs(grad) < 1e-12:
            break
        x = next_x
    return states


def bfgs_states():
    p = np.array([-1.2, 1.0], dtype=np.float64)
    H = np.eye(2)
    states = [dict(p=p.tolist(), value=float(objective(p)), alpha=0.0)]
    for _ in range(80):
        g = gradient(p)
        if np.linalg.norm(g) < 1e-7:
            break
        direction = -H @ g
        alpha = 1.0
        # Armijo backtracking prevents an oversized step up the steep valley wall.
        while objective(p + alpha*direction) > objective(p) + 1e-4*alpha*(g @ direction):
            alpha *= 0.5
            if alpha < 1e-14:
                raise RuntimeError('Line search failed')
        q = p + alpha*direction
        s, y = q-p, gradient(q)-g
        curvature = float(y @ s)
        if curvature > 1e-12:
            rho = 1/curvature
            V = np.eye(2) - rho*np.outer(s, y)
            H = V @ H @ V.T + rho*np.outer(s, s)
        p = q
        states.append(dict(p=p.tolist(), value=float(objective(p)), alpha=alpha))
    return states


def boost_states(histogram=False, rounds=16):
    """Educational squared-error GBDT: exact/depthwise versus binned/leafwise.

    These are small transparent implementations, not the xgboost/lightgbm libraries.
    Loss is 0.5*(prediction-target)**2, g=prediction-target, h=1.
    Both use lambda=1, gamma=0, eta=0.3, min_samples_leaf=3.
    """
    x = np.linspace(-3, 3, 72, dtype=np.float64)
    target = np.sin(1.5*x) + 0.35*x + 0.12*np.cos(11*x)
    prediction = np.full_like(x, target.mean())
    thresholds = (np.linspace(x.min(), x.max(), 13)[1:-1] if histogram
                  else (x[:-1]+x[1:])/2)
    states = [dict(pred=prediction.tolist(), tree=np.zeros_like(x).tolist(),
                   residual=(target-prediction).tolist(), mse=float(np.mean((target-prediction)**2)),
                   leaves=1, splits=[], gains=[])]
    for _ in range(rounds):
        residual = target-prediction
        g = prediction-target
        leaves = [dict(ids=np.arange(len(x)), depth=0)]
        splits, gains = [], []

        def best_split(leaf):
            ids = leaf['ids']
            best = None
            for threshold in thresholds:
                left = ids[x[ids] <= threshold]
                right = ids[x[ids] > threshold]
                if min(len(left), len(right)) < 3:
                    continue
                gain = 0.5*(g[left].sum()**2/(len(left)+1)
                            + g[right].sum()**2/(len(right)+1)
                            - g[ids].sum()**2/(len(ids)+1))
                if gain > 1e-12 and (best is None or gain > best[0]):
                    best = (float(gain), float(threshold), left, right)
            return best

        while True:
            candidates = [(i, best_split(leaf)) for i, leaf in enumerate(leaves)
                          if histogram or leaf['depth'] < 2]
            candidates = [(i, b) for i, b in candidates if b is not None]
            if not candidates or (histogram and len(leaves) >= 4):
                break
            # Depthwise chooses the shallowest available node; leafwise chooses gain.
            index, best = (max(candidates, key=lambda item: item[1][0]) if histogram
                           else min(candidates, key=lambda item: (leaves[item[0]]['depth'], item[0])))
            gain, threshold, left, right = best
            depth = leaves[index]['depth'] + 1
            leaves[index:index+1] = [dict(ids=left, depth=depth), dict(ids=right, depth=depth)]
            splits.append(threshold)
            gains.append(gain)
        tree = np.zeros_like(x)
        for leaf in leaves:
            ids = leaf['ids']
            tree[ids] = -g[ids].sum()/(len(ids)+1)
        prediction = prediction + 0.3*tree
        states.append(dict(pred=prediction.tolist(), tree=tree.tolist(), residual=residual.tolist(),
                           mse=float(np.mean((target-prediction)**2)), leaves=len(leaves),
                           splits=splits, gains=gains))
    return dict(x=x.tolist(), target=target.tolist(), states=states)


def make_animation(method):
    """Return a retained FuncAnimation; render with HTML(ani.to_jshtml())."""
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.4))
    fig.subplots_adjust(wspace=0.30, bottom=0.17, top=0.83)
    if method == 'newton':
        states = newton_states()
    elif method == 'newton_opt':
        states = newton_opt_states()
    elif method in ('taylor', 'maclaurin'):
        states = list(range(20))
    elif method == 'bfgs':
        states = bfgs_states()
    elif method in ('xgboost', 'lightgbm'):
        data = boost_states(method == 'lightgbm')
        states = data['states']
    else:
        raise ValueError('Unknown method')

    def animate(frame):
        a, b = axes
        for ax in axes:
            ax.clear()
            ax.grid(alpha=0.18)
        state = states[frame]
        if method == 'newton':
            grid = np.linspace(0, 3.3, 400)
            x, q = state['x'], state['next']
            a.plot(grid, grid**2-2, color='#176b64', label='f(x)=x^2-2')
            a.plot(grid, x*x-2+2*x*(grid-x), '--', color='#c45632', label='tangent')
            a.scatter([x, q], [x*x-2, 0], color='#c45632')
            a.axhline(0, color='gray', lw=0.7)
            a.set(xlim=(0, 3.3), ylim=(-3, 9), xlabel='x', ylabel='f(x)')
            b.semilogy(range(frame+1), [max(s['error'], 1e-16) for s in states[:frame+1]], 'o-')
            b.set(xlim=(0, 5), ylim=(1e-16, 10), xlabel='iteration k', ylabel='|f(x_k)|')
            title = f'Newton k={frame}, x={x:.8f}, next={q:.8f}'
        elif method == 'newton_opt':
            grid = np.linspace(-1, 2.8, 400)
            x, q = state['x'], state['next']
            model = state['value'] + state['grad']*(grid-x) + 0.5*state['hess']*(grid-x)**2
            a.plot(grid, grid**4/4+grid**2/2, color='#176b64', label='objective f(x)')
            a.plot(grid, model, '--', color='#c45632', label='quadratic model')
            a.scatter([x, q], [state['value'], q**4/4+q*q/2], color='#c45632', label='current / next on f')
            a.scatter([q], [state['value']-state['grad']**2/(2*state['hess'])], marker='x', color='#c45632', label='model minimum')
            a.set(xlim=(-1,2.8), ylim=(-2,21), xlabel='x', ylabel='objective / model')
            b.plot(grid, grid**3+grid, color='#176b64', label="f'(x)")
            b.plot(grid, state['grad']+state['hess']*(grid-x), '--', color='#c45632', label="tangent of f'")
            b.scatter([x,q], [state['grad'],0], color='#c45632')
            b.axhline(0, color='gray', lw=0.7)
            b.set(xlim=(-1,2.8), ylim=(-5,26), xlabel='x', ylabel='derivative')
            b.legend(fontsize=8)
            title = f"Newton optimization k={frame}, f={state['value']:.3g}, |gradient|={abs(state['grad']):.3g}"
        elif method in ('taylor', 'maclaurin'):
            center = 1.0 if method == 'taylor' else 0.0
            grid = np.linspace(-2*np.pi, 2*np.pi, 500, dtype=np.float64)
            approx = taylor_sin(grid, frame, center)
            a.plot(grid, np.sin(grid), '--', color='gray', label='sin(x)')
            a.plot(grid, approx, color='#c45632', label='polynomial')
            a.scatter([center], [np.sin(center)], color='#176b64')
            a.set(xlim=(-2*np.pi, 2*np.pi), ylim=(-2, 2), xlabel='x', ylabel='value (clipped to +/-2)')
            b.semilogy(grid, np.maximum(abs(approx-np.sin(grid)), 1e-16), color='#176b64')
            b.set(xlim=(-2*np.pi, 2*np.pi), ylim=(1e-16, 1e4), xlabel='x', ylabel='absolute error (floor 1e-16)')
            title = f'{method}: center={center}, degree={frame}'
        elif method == 'bfgs':
            xx, yy = np.meshgrid(np.linspace(-1.6, 1.6, 180), np.linspace(-0.6, 2, 160))
            a.contour(xx, yy, (1-xx)**2+100*(yy-xx*xx)**2, levels=[0.1,1,4,10,50,100,400], alpha=0.45)
            path = np.array([s['p'] for s in states[:frame+1]])
            a.plot(path[:, 0], path[:, 1], 'o-', ms=3, color='#c45632', label='BFGS path')
            a.scatter([1], [1], marker='*', s=110, color='#176b64')
            a.set(xlim=(-1.6, 1.6), ylim=(-0.6, 2), xlabel='x', ylabel='y')
            b.semilogy(range(frame+1), [max(s['value'], 1e-16) for s in states[:frame+1]])
            b.set(xlim=(0, len(states)-1), ylim=(1e-16, 100), xlabel='iteration', ylabel='objective')
            title = f'BFGS k={frame}, f={state["value"]:.3g}, alpha={state["alpha"]:.3g}'
        else:
            x, target = np.array(data['x']), np.array(data['target'])
            a.scatter(x, target, s=10, color='gray', label='training target')
            a.step(x, state['pred'], where='mid', color='#176b64', label='ensemble')
            a.set(xlim=(-3,3), ylim=(-2.3,2.3), xlabel='feature x', ylabel='prediction')
            b.scatter(x, state['residual'], s=10, color='gray', label='residual before this tree')
            b.step(x, state['tree'], where='mid', color='#c45632', label='new tree (before eta)')
            for split in state['splits']:
                b.axvline(split, color='#c45632', lw=0.7, alpha=0.3)
            b.set(xlim=(-3,3), ylim=(-2.3,2.3), xlabel='feature x', ylabel='residual / leaf weight')
            b.legend(fontsize=8)
            title = f'{method} teaching model: round={frame}, MSE={state["mse"]:.4f}, leaves={state["leaves"]}'
        a.legend(loc='upper left', fontsize=8)
        fig.suptitle(title)
        return tuple(axes)

    # Axes and labels change each frame, so blit=False is intentional.
    ani = FuncAnimation(fig, animate, frames=len(states), interval=650, blit=False, repeat=False)
    plt.close(fig)  # Suppress the extra static image in Jupyter.
    return ani


# Jupyter example:
# ani = make_animation('newton')
# HTML(ani.to_jshtml())
# Other names: newton_opt, taylor, maclaurin, bfgs, xgboost, lightgbm.