from susmost import load_lattice_task
from susmost import mc
import numpy as np
from mpi4py import MPI
lt = load_lattice_task('./')
lt.set_property('coverage', {'A':1.0, 'B':1.0} )
lt.set_property('A_cnt', {'A':1})
lt.set_property('B_cnt', {'B':1})

L = 36 # lattice size
temperatures = [100.,150.] # K
muB = 0.
lt.set_ads_energy('B', muB) # kJ/mol
kB = 0.008314 # kJ/(mol*K)

# Create a file to write the curent valuses of interest
f = open('results.dat','w')
f.write('T	mu	coverage	H	U	Qd	Cp''\n')
# Perform the simulation for a series of TEtB chemical potentials
for mu in np.arange(30., 0.-0.0001, -2.0 ):
	lt.set_ads_energy('A', mu)	# kJ/mol
	m = mc.make_metropolis(lt, L, temperatures, kB)
	mc.run(m, log_periods_cnt=10, log_period_steps = 1000*m.cells_count, relaxation_steps = 100000*m.cells_count, \
		traj_fns = ["T={}.xyz".format(T) for T in temperatures])
	mc.stat_digest(m)

	if MPI.COMM_WORLD.Get_rank() == 0:
		E_2 = 0
		covE = 0
		cov_2 = 0
		for T, param_log in zip(temperatures, m.full_params_log):
			stat = np.mean(param_log,axis=0)[[0, -2, -3]] # coverage, H, U
			for log_row in param_log:
				E_2 += log_row[-3]**2.0
				covE += log_row[0]*log_row[-3]
				cov_2 += log_row[0]**2.0
			E_2 /= len(param_log)
			covE /= len(param_log)
			cov_2 /= len(param_log)
			Qd = -(covE - stat[0]*stat[2])/(cov_2-stat[0]**2.0)
			Cp = (L^2)*(E_2 - stat[2]*stat[2])/(T*T*kB)
			f.write( ('{}\t'*7).format(T, mu, *stat, Qd, Cp) + '\n')
		f.flush()
f.close()