import numpy as np
import matplotlib.pyplot as plt
 
def makeLayer(y, height):
    neg=0.0
    pos=0.0
    if y > 0:
        if y - height >= 0:
            pos = height
            y -= pos
        else: 
            pos = y
    elif y < 0:
        if y + height <= 0:
            neg = height
            y += neg
        else: 
            neg = -y
    return pos, neg
 
def horizonPlot(x, y, nfolds=3, negcolor='CornflowerBlue', poscolor='DarkGreen', **figkwargs):
    fig, ax = plt.subplots(**figkwargs)
    alpha = 1.0/nfolds
    vlayer = np.vectorize(makeLayer)
    height = np.ceil(np.abs(y).max()/nfolds)
    while (y != 0).any():
 
        layer = vlayer(y,height)
        y -= layer[0];
        y += layer[1]
        ax.fill_between(x, 0, layer[0], color=poscolor, alpha=alpha)
        ax.fill_between(x, height-layer[1], height, color=negcolor, alpha=alpha)
        
        ax.set_yticks([])
        ax.set_xlim([x.min(), x.max()])
        ax.set_ylim([0, height])
        ax.xaxis.grid(False)
        ax.yaxis.grid(False)
        
    return fig
 

# data
x = np.linspace(0, np.pi*4, 137)
y = (2*np.random.normal(size=137) + x**2)
xx = np.hstack([-1*x[::-1], x])
yy = np.hstack([-1*y[::-1], y])

# figure properties
figkwargs = dict(figsize=(12,2), dpi=300)

fig, ax = plt.subplots(**figkwargs)
pos = np.ma.masked_less(yy, 0)
neg = np.ma.masked_greater(yy, 0)
ax.fill_between(xx, neg, y2=0, color='CornflowerBlue', alpha=0.5)
ax.fill_between(xx, pos, y2=0, color='DarkGreen', alpha=0.5)
ax.xaxis.grid(False)
ax.yaxis.grid(False)

ax.set_xlim([xx.min(), xx.max()])
ax.set_ylim([neg.min(), pos.max()])
ax.set_yticks([])


fig = horizonPlot(xx, yy, nfolds=3, **figkwargs)