import numpy as np
import matplotlib.pyplot as plt

def psi_flipped(x):
    return np.exp(-x**2)

def weird_cost(y_grid, x_center):
    diff = np.abs(y_grid - x_center)
    s = 2 * (np.cos(3 * x_center) + 1)
    c = s * diff**2 + (4 - s) * diff**4
    return c

xlims = (-1.5, 1.5)
ylims = (0.1, 1)

xs = np.linspace(xlims[0], xlims[1], 1000)
ys = psi_flipped(xs)

randxs = np.linspace(xlims[0], xlims[1], 100)
randys = psi_flipped(randxs)

y_fill_top = 7.5 

fig, ax = plt.subplots(figsize=(9, 6), dpi=150)

for x0, y0 in zip(randxs, randys):
    cost_ys = y0 + weird_cost(xs, x0)
    
    ax.plot(xs, cost_ys, color='forestgreen', linewidth=0.8, alpha=0.6)
    
    ax.fill_between(xs, cost_ys, y_fill_top, color='forestgreen', alpha=0.02, linewidth=0)

ax.plot(xs, ys, color='#0088CC', linewidth=1.5, zorder=5)

ax.scatter(randxs, randys, color='#D6604D', edgecolors='black', s=40, zorder=6, linewidth=0.8)

ax.set_xlim(xlims)
ax.set_ylim(ylims)

ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)

ax.grid(True, which='both', linestyle='-', linewidth=0.5, color='gray', alpha=0.15)

ax.tick_params(direction='in', length=4)

plt.tight_layout()
plt.savefig('c-transformation of a gaussian curve.svg')

plt.show()