############################## Parameters ##############################
# The name of the lower-boiling point component                        #
low_bp = 'methanol'
# The name of the higher-boiling point component                       #
high_bp = 'water'
########################################################################
# Total pressure, measured in Pa                                       #
P = 101325
assert P > 0, 'Total pressure must be greater than 0 (Pa).'            #
########################################################################
# The reflux ratio                                                     #
R = 1.5
assert R > 0, 'The reflux ratio must be greater than 0.'               #
########################################################################
# The parameter q                                                      #
#   Superheated vapor feed  : q < 0                                    #
#   Saturated vapor feed    : q = 0                                    #
#   Partially vaporized feed: 0 < q < 1                                #
#   Saturated liquid feed   : q = 1                                    #
#   Subcooled liquid feed   : q > 1                                    #
q = 1
########################################################################
# The distillate composition                                           #
# (measured in mole fraction of the lower-boiling feed component
# in the liquid phase; the same shall apply hereinafter)               #
x_D = .95
assert 0 < x_D < 1, \
'The distillate composition must be greater than 0 and less than 1.'   #
########################################################################
# The bottoms composition                                              #
x_W = .05
assert 0 < x_W < 1, \
'The bottoms composition must be greater than 0 and less than 1.'      #
########################################################################
# The feed composition                                                 #
z_f = .40
assert 0 < z_f < 1, \
'The feed composition must be greater than 0 and less than 1.'         #
########################################################################

y_D = x_D
y_W = x_W

antoine_constants = {
	# Antoine equation:
	#   log10(p) = A - B / (T - C)
	#   where p is the vapor pressure in kilopascal
	#   and T is temperature in kelvin
	# Data:
	#   Substance: (A, B, C)
	#   (ref: http://memoirs.lib-e.yamaguchi-u.ac.jp/621/01.pdf)
	## Alkane
	'2-methylbutane'         : (5.93330, 1029.602, 38.856),
	'pentane'                : (5.99028, 1071.187, 40.384),
	'2-methylpentane'        : (5.99479, 1152.210, 44.579),
	'3-methylpentane'        : (5.99139, 1162.069, 44.870),
	'hexane'                 : (6.01098, 1176.102, 48.251),
	'heptane'                : (6.02701, 1267.592, 56.354),
	'2,3-dimethylpentane'    : (5.98293, 1240.404, 51.056),
	'octane'                 : (6.04394, 1351.938, 64.030),
	'2,2,4-trimethylpentane' : (5.92751, 1252.340, 53.060),
	## cyclic compound
	'cyclohexane'            : (6.00569, 1223.273, 48.061),
	'benzene'                : (6.01905, 1204.637, 53.081),
	'toluene'                : (6.08436, 1347.620, 53.363),
	## ether
	'diethyl ether'          : (6.04920, 1061.391, 45.090),
	'methul t-butyl ether'   : (6.070343, 1158.912, 43.200),
	'ethyl t-butyl ether'    : (6.073724, 1206.874, 49.190),
	't-amyl methyl ether'    : (6.067822, 1256.258, 50.100),
	'diisopropyl ether'      : (5.97081, 1137.408, 54.634),
	'dibutyl ether'          : (5.92274, 1298.256, 82.006),
	## ketone
	'acetone'                : (6.25017, 1214.208, 43.148),
	'methyl ethyl ketone'    : (6.18397, 1258.940, 51.425),
	'diethyl ketone'         : (6.14570, 1307.941, 59.182),
	'methyl propyl ketone'   : (6.13931, 1309.629, 58.585),
	'methyl isopropyl ketone': (6.09024, 1265.595, 57.631),
	'methyl isobutyl ketone' : (5.81291, 1176.833, 80.225),
	## alcohol
	'methanol'               : (7.24693, 1605.615, 31.317),
	'ethanol'                : (7.24222, 1595.811, 46.702),
	'1-propanol'             : (6.87065, 1438.587, 74.598),
	'2-propanol'             : (6.86634, 1360.183, 75.557),
	'1-butanol'              : (6.54068, 1335.028, 96.496),
	'2-butanol'              : (6.35079, 1169.924, 103.413),
	'water'                  : (7.06252, 1650.270, 46.804),
}

import matplotlib.pyplot as plt
import numpy as np
import sys

try:
	from scipy import optimize
	def inv(func, y, x0=0, *args, **kwargs):
		return optimize.fsolve(lambda x: func(x) - y, x0)[0]
except ImportError:
	def inv(func, y, x0=0, x1=None, eps=2*sys.float_info.epsilon):
		assert not x0 == x1, 'x1 cannot be equal to x0.'
		f = lambda x: func(x) - y
		if not x1:
			x1 = x0 + 1
		elif x0 > x1:
			x0, x1 = x1, x0

		# scant method
		while abs(f(x1)) > eps:
			s = (x0 * f(x1) - x1 * f(x0)) / (func(x1) - func(x0))
			x0, x1 = x1, s
			if abs(x0 - x1) < eps:
				break
		return x1


def Antoine(func):
	def wrapped(*args, **kwargs):
		res = func(*args, **kwargs)
		return (10 ** res) * 1000
	return wrapped


@Antoine
def vp(consts, T):
	return consts[0] - consts[1] / (T - consts[2])


def vp_1(T):
	return vp(antoine_constants[low_bp], T)


def vp_2(T):
	return vp(antoine_constants[high_bp], T)


def temperature(x):
	def total_P(T):
		# Raoult's law
		p_2 = vp_2(T) * (1 - x)
		p_1 = vp_1(T) * x

		# Dalton's law
		return p_2 + p_1
	return inv(total_P, y=P, x0=373.15, eps=1e-5)


@np.vectorize
def vle(x):
	T = temperature(x)
	return vp_1(T) * x / P


def draw(z_f, x_W, x_D, R, q, P):
	plt.close()
	t = np.linspace(0, 1, 200)

	func_r = lambda x: R / (R + 1) * (x - x_D) + y_D  # rectifying section
	if q == 1:
		Q_x = z_f
		Q_y = func_r(z_f)
		min_s = (y_D - vle(z_f)) / (x_D - z_f)
	else:
		func_q = lambda x: -q / (1 - q) * x + z_f / (1 - q)
		func_diff = lambda x: func_r(x) - func_q(x)
		f0 = func_diff(0)
		Q_x = -f0 / (func_diff(1) - f0)
		Q_y = func_q(Q_x)
		p_x = inv(lambda x: vle(x) - func_q(x), 0, z_f-.1, x1=z_f+.1, eps=1e-5)
		p_y = func_q(p_x)
		min_s = (y_D - p_y) / (x_D - p_x)
	func_s = lambda x: (y_W - Q_y) / (x_W - Q_x) * (x - x_W) + y_W  # stripping section
	func = lambda x: func_r(x) if x >= Q_x else func_s(x)
	min_R = min_s / (1 - min_s)

	plt.plot([0, 1], [0, 1], 'k')  # diagonal line
	plt.plot(t, vle(t), 'b')  # VLE line

	plt.plot([x_D, x_D], [0, y_D], 'k--')  # distillate composition line
	plt.plot([x_W, x_W], [0, y_W], 'k--')  # bottoms composition line
	plt.plot([z_f, z_f], [0, z_f], 'k--')  # feed composition line

	label_opt = {
		'verticalalignment': 'top',
		'horizontalalignment': 'center'
	}
	plt.text(x_D, -.0105, f'{x_D:.2f}', **label_opt)
	plt.text(x_W, -.0105, f'{x_W:.2f}', **label_opt)

	plt.plot([Q_x, x_D], [Q_y, y_D], 'g')  # rectifying operating line
	plt.plot([x_W, Q_x], [y_W, Q_y], 'g')  # stripping section operating line
	plt.plot([z_f, Q_x], [z_f, Q_y], 'b')  # q-line

	if R > min_R:
		y1 = y_D
		x1 = x_D
		cnt = 1
		while (x1 - x_W) > 0.01:
			x2 = inv(vle, y=y1, x0=0, x1=x1, eps=1e-5)
			y2 = max(func(x2), x2)
			T = round(temperature(x2), 1)

			if y1 >= Q_y:
				if q == 0:
					if y2 <= Q_y <= y1:
						feed = cnt, T
				elif q == 1:
					if x2 <= Q_x <= x1:
						feed = cnt, T
				else:
					F_x = inv(func_q, y=y1, x0=x2, x1=x1, eps=1e-5)
					if (y1 < vle(F_x)) or (y2 <= func_q(x2) <= y1):
						feed = cnt, T

			plt.text((x1 - x2) / 2 + x2, y1 + .005,
				f'{T:.1f} K',
				horizontalalignment='center'
			)
			yield f'{cnt}: x = {round(x2, 2):.2f}, y = {round(y1, 2):.2f} at {T:.1f} K'

			plt.plot([x1, x2], [y1, y1], 'c')
			plt.plot([x2, x2], [y1, y2], 'c')
			x1, y1 = x2, y2
			cnt += 1
		yield f'Feed tray: Step #{feed[0]} ({feed[1]:.1f} K)'
	else:
		yield f'R <= {round(min_R, 2):.2f}'

	plt.xlim(xmin=0.0, xmax=1.0)
	plt.ylim(ymin=0.0, ymax=1.05)
	plt.title(f'{low_bp}-{high_bp} (q = {q}, R = {R}, P = {round(P, 1):.1f} Pa)')
	plt.xlabel('x')
	plt.ylabel('y')
	plt.show()


def validate_compound():
	global low_bp, high_bp
	if low_bp == high_bp:
		print('Same compound')
		return False
	for compound in [low_bp, high_bp]:
		if compound not in antoine_constants:
			print(f'No data: {compound}')
			return False

	bp_opt = [P, 273.15, 373.15, 1e-5]
	bp1 = inv(vp_1, *bp_opt)
	bp2 = inv(vp_2, *bp_opt)
	if bp1 > bp2:
		print(f'b.p.: {low_bp} ({bp1:.1f} K) > {high_bp} ({bp2:.1f} K)')
		print('      They will be swapped.')
		low_bp, high_bp = high_bp, low_bp

	return True


def main():
	if not validate_compound():
		return

	for R in [1.5, 3.0, 0.5]:
		print(f'R = {R}')
		for data in draw(z_f, x_W, x_D, R, q, P):
			print(data)

if __name__ == '__main__':
	main()