import numpy as np
import matplotlib.pyplot as plt

# Define the functions
def rho(x):
  return np.exp(-x)

def cdf(x):
  return 1 - np.exp(-x)

xmax = 1.5
# Generate x values
x = np.linspace(0, xmax, 500)
x2 = np.linspace(np.log(2), 2 * np.log(2), 500)

# Create the plot
fig, axs = plt.subplots(2, 1, figsize=(8, 8))

# Plot the PDF

ax = axs[0]
axs[0].plot(x, rho(x))
axs[0].set_ylabel('Probability Density')
axs[0].legend()

ax.fill_between(x, 0, rho(x), color='lightblue')

ax.set_xlim(0, xmax)
ax.set_ylim(0, 1)
ax.set_xticks([])
ax.set_yticks([])

# Plot the CDF
ax = axs[1]
axs[1].plot(x, cdf(x))
axs[1].set_ylabel('Cumulative Probability')
axs[1].legend()

axs[1].plot([0, xmax], [1, 1], 'k--', linewidth=0.8) # Vertical line
axs[1].plot([np.log(2), np.log(2)], [0, 0.5], 'k--', linewidth=0.8) # Vertical line
axs[1].plot([0, np.log(2)], [0.5, 0.5], 'k--', linewidth=0.8) # Horizontal line
axs[1].plot([0, 2 * np.log(2)], [0, 1], 'r-')
axs[1].plot(x2, np.exp(x2) / 4, 'g-', label='Reflected across median')
ax.scatter(np.log(2), 0.5, marker='o', color='k', label='Median')

ax.fill_between(x, cdf(x), 1, color='lightgray', alpha=0.4)
xs = np.linspace(np.log(2), 2 * np.log(2), 100)
ax.fill_between(xs, np.exp(xs) / 4, xs/(2*np.log(2)), color='lightblue')
xs = np.linspace(0, np.log(2), 100)
ax.fill_between(xs, xs/(2*np.log(2)), cdf(xs), color='lightblue')


axs[1].legend()
ax.set_xlim(0, xmax)
ax.set_ylim(0, 1)
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()

plt.savefig("mean_greater_than_median.svg")
plt.show()