import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Computer Modern Roman']
plt.rcParams['mathtext.fontset'] = 'cm'

# Setup
dpi = 300
width_px, height_px = 1276, 877
fig, axs = plt.subplots(1, 3, figsize=(width_px/dpi, height_px/dpi), dpi=dpi)
x_range = np.linspace(-2, 2, 1600)
param_vals = [1/3, 3/2, -1/2, -4/3, 0]
colors = ["#387FC7", "#F5C443", "#40904E", "#AD3049", "#000000"]

# Plot quadratic functions with varying parameters
params = [
    ('c', lambda x, p: x**2 + p),      # y = x^2 + c
    ('b', lambda x, p: x**2 + p * x),  # y = x^2 + bx
    ('a', lambda x, p: p * x**2)       # y = ax^2
]
for i, (label, func) in enumerate(params):
    for val, color in zip(param_vals, colors):
        axs[i].plot(x_range, func(x_range, val), color=color, lw=1.5, zorder=i+1)
    if label == 'b':
        # Plot the vertices parabolic curve when varying `b`
        axs[i].plot(x_range, -x_range**2, color='grey', lw=1.5, linestyle=':', zorder=0)
    axs[i].set_xlabel(rf"$\boldsymbol{{{label}}}$", fontsize=14)

# Format axes
for ax in axs:
    ax.set_xlim(-2, 2)
    ax.set_xticks(np.arange(-2, 3))
    ax.set_ylim(-2, 2)
    ax.set_yticks(np.arange(-2, 3))
    ax.axhline(0, color='grey', linewidth=0.5, zorder=0)
    ax.axvline(0, color='grey', linewidth=0.5, zorder=0)
    for spine in ax.spines.values():
        spine.set_edgecolor('grey')
        spine.set_linewidth(0.5)
    ax.tick_params(axis='both', colors='grey', labelsize=4, width=0.5)

# Legend
labels = [
    r"$-\frac{4}{3}$,", 
    r"$-\frac{1}{2}$, ",
    r"$0$,", 
    r"$+\frac{1}{3}$,", 
    r"$+\frac{3}{2}$"
]
annot_colors = [colors[3], colors[2], colors[4], colors[0], colors[1]]
handles = [Line2D([], [], color=c, lw=0) for c in colors]
leg = fig.legend(
    handles, labels,
    loc=("upper center"), bbox_to_anchor=(0.5, 0.99),
    frameon=False, ncol=len(labels), handlelength=0.0, handletextpad=0.1,
    columnspacing=0.2, fontsize=18,
 )
for txt, c in zip(leg.get_texts(), annot_colors):
    txt.set_color(c)

# Show
plt.tight_layout(rect=(0, 0, 1, 0.80), w_pad=0.5)
plt.show()
fig.savefig("Quadratic_equation_coefficients_standard_ranges.png", dpi=dpi)