import numpy as np
import matplotlib.pyplot as plt
from CoolProp.CoolProp import PropsSI
import matplotlib.cm as cm

# Ambient/reference conditions
T0 = 298.15  # K (25°C)
P0 = 101325  # Pa (1 bar)

# Pressures to consider (Pa)
pressures = [101325, 200000, 300000, 400000, 500000]  # 1 bar to 5 bar
pressure_labels = ['1 bar', '2 bar', '3 bar', '4 bar', '5 bar']

# Temperature range around saturation (°C)
T_min = 50
T_max = 200
n_points = 200

plt.figure(figsize=(5.6,5))

# Choose a colormap and pick N distinct colors
colors = cm.viridis(np.linspace(0, 1, len(pressures)))

for i, P in enumerate(pressures):
    T = np.linspace(T_min + 273.15, T_max + 273.15, n_points)
    exergies = []
    
    for temp in T:
        try:
            h = PropsSI('H','P',P,'T',temp,'Water')
            s = PropsSI('S','P',P,'T',temp,'Water')
            h0 = PropsSI('H','T',T0,'P',P0,'Water')
            s0 = PropsSI('S','T',T0,'P',P0,'Water')
            ex = (h - h0) - T0*(s - s0)
            exergies.append(ex / 1000)
        except:
            exergies.append(np.nan)
    
    plt.plot(T - 273.15, exergies, label=pressure_labels[i], color=colors[i], linewidth=2.2)

# Styling
plt.xlabel('$T$ / °C', fontsize=14)
plt.ylabel('$e$ / kJ/kg', fontsize=14)
plt.ylim(0, 800)
plt.xlim(60, 200)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.grid(False)
plt.legend(fontsize=14)
plt.tight_layout()
plt.show()