#!/usr/bin/env python
# This code extracts energy information from ENERGY_INFO.txt generated by dekode after running
# a set of strained unitcell and use them to calculated the deformation potential for VBM and CBM
# Assuming that the format (header) of ENERGY_INFO.txt is as followed:
# strain(%)      volume(A3)       total(eV)        core(eV)         VBM(eV)         CBM(eV)
# Please contact alireza@wustl.edu if you had any questions. Bug reports are very much appreciated.
# By: Alireza Faghaninia

import argparse
import numpy as np

def calc_def_potential(info_file):
	data =[]
	counter = 0
	with open(info_file, 'r') as energy_info:
#		enfo.seek(0)
#		next(enfo)
			for row in energy_info:
				counter += 1
				if counter > 1:
					data.append([float(i) for i in row.split()])
	
	data = np.mat(data)
	eqidx = np.where(data[:,0]==0)[0]
#	eqidx = abs(int(data[0,0])) + 1
	eqvol = data[eqidx,1]
	data[:,1] -= eqvol 
	data[:,4] -= data[:,3]
	data[:,5] -= data[:,3]
	
	x = np.array(np.ravel(data[:,1]))
	y = np.array(np.ravel(data[:,4]))
	z = np.polyfit(x, y, 6)
	VBM_deformation = abs(z[-2]*eqvol)
	
	x = np.array(np.ravel(data[:,1]))
	y = np.array(np.ravel(data[:,5]))
	z = np.polyfit(x, y, 6)
	CBM_deformation = abs(z[-2]*eqvol)
	
	return float(VBM_deformation), float(CBM_deformation)

if __name__ == "__main__":
### Check the input arguments
	parser = argparse.ArgumentParser()
	parser.add_argument("-f","--info_file", help="The name of the file containing energy information of the strained unitcells", required = False, default = 'ENERGY_INFO.txt')
	args = parser.parse_args()

	VBM_deformation, CBM_deformation = calc_def_potential(args.info_file)
	print(VBM_deformation)
	print(CBM_deformation)

