import numpy as np
import matplotlib.pyplot as plt

def S_(N, V, U):
"""
Just a simple function to calculate entropy as defined in week 2 problem 10
"""
return ((N*V*U)**(1/3))

Ntot = 50 #total amount - keep the same
Vtot = 13 #total amount - keep the same
Utot = 80 #total amount - keep the same
NA = 30 #Can change in order to alter NA and NB
VA = 9 #Can change in order to alter VA and VB
NB = Ntot-NA
VB = Vtot-VA
UAlist = np.arange(0, Utot, 0.1) #creates UA to calculate S(UA) and then graph

#This cell calculates S(UA) for each UA
SAlist = np.empty(len(UAlist))
for i in range(len(UAlist)):
SAlist[i] = S_(NA, VA, UAlist[i])+S_(NB, VB, Utot-UAlist[i])

plt.plot(UAlist, SAlist)
plt.xlabel(r"$U_A$ ( J )")
plt.ylabel(r"S ( J K$^{-1} )$")
plt.show()
UAmax = UAlist[np.where(SAlist == np.max(SAlist))]
print("The value of U_A that maximizes S(U_A) is ", UAmax)
print("The maximum value of S(U_A) is ", np.max(SAlist))