from netCDF4 import Dataset
from yambopy import LetzElphElectronPhononDB,ConvertElectronPhononDB
import os
import shutil
import subprocess
"""
Calculate gauge-invariant electron-phonon matrix elements with LetzElPhC and convert them into Yambo format

- Usage:

>> yambopy l2y -ph phinp -b b1 b2 -par nq nk [--kernel kernel] [--lelphc lelphc] [--no_lelphc_dbs] [--no_gkkp]

- Input parameters:
	-ph                   : path to ph.x input file, e.g. dvscf/ph.in
	-b                    : initial and final band indices (counting from 1, both values included)
	-par [OPT]            : MPI pools for q and k (needs mpirun)
	--kernel [OPT]        : e-ph kernel type, default 'dfpt'
	--lelphc [OPT]        : path to lelphc executable, default 'lelphc', code will prompt
	--no_lelphc_dbs [OPT] : will remove LetzElPhC input and outputs
    --no_gkkp [OPT]       : do not generate the gkkp databases for Yambo/Lumen
 
- Prerequisites:

* ph.x phonon calculation must be complete, e.g. the phinp folder should contain:
    - ph.x input file
    - pw.x (scf) save directory
    - [prefix].dyn files
    - _ph* directories
* Yambo SAVE directory must be present. We run in the directory where the SAVE is.
* LetzElPhC must be installed
* mpirun must be linked for parallel runs
"""

def get_input(bands,kernel,pools,ph_path):
	"""
	Get LetzElPhC input
	"""
	input_file = []; app = input_file.append
	app("# LetzElPhC input for yambo generated by yambopy")
	app(f"nqpool      = {pools[0]}") # first pool is qpoints
	app(f"nkpool      = {pools[1]}") # second pool is kpoints
	app(f"start_bnd   = {bands[0]}")
	app(f"end_bnd     = {bands[1]}")
	app(f"save_dir    = ./SAVE")
	app(f"kernel      = {kernel}")
	app(f"ph_save_dir = {ph_path}/ph_save")
	app(f"convention = yambo")
	return "\n".join(input_file)

def checks(phinp,lelphc,bands,kernel,pools):

	## ph.x input file and directory
	if '/' in phinp:
		path_ph = phinp.rsplit('/',1)[0]
		inp_ph  = phinp.rsplit('/',1)[1]
	else:
		path_ph='./'
		inp_ph = phinp

	## check for SAVE
	if not os.path.isdir("SAVE"):
		raise FileNotFoundError("SAVE directory not found. Make sure to run in the directory where the SAVE is.")

	## check for lelphc executable
	is_lelphc = shutil.which(lelphc)
	if is_lelphc is None: lelphc=input("lelphc executable not found. Please provide absolute path here: \n")

	## check for mpirun
	if int(pools[0])*int(pools[1])>1:
		is_mpirun = shutil.which('mpirun')
		if is_mpirun is None: 
			print("[WARNING] mpirun not found, running in serial") 
			pools = [1,1]

	## check band indices
	try: assert(int(bands[0])<int(bands[1]))
	except: raise ValueError("[ERROR] band indices must be integers with b1<b2")

	## check kernel type
	if kernel not in ['dfpt','bare','dfpt_local','bare_local']:
		raise ValueError(f"[ERROR] Unrecognized kernel type {kernel}")

	## lelphc input file
	inp_lelphc = get_input(bands,kernel,pools,path_ph)
	inp_name = 'lelphc.in'

	return lelphc,path_ph,inp_ph,inp_lelphc,inp_name

def run_preprocessing(lelphc,path_ph,inp_ph):

	print(":: LetzElPhC pre-processing ::")
	pp_run_str = f"{lelphc} -pp --code=qe -F {inp_ph}"
	try:
		pp=subprocess.run(f'cd {path_ph} ; {pp_run_str} ; cd -',shell=True,check=True,capture_output=True,text=True)
	except subprocess.CalledProcessError as e:
		print("PP Error:", e)
		print("Description", e.stderr)

	if not os.path.isdir(f"{path_ph}/ph_save"):
		print("[ERROR] Something wrong with preprocessing step:")
		print(pp.stderr)
		exit()

def run_elph(lelphc,inp_lelphc,inp_name):

	def get_ntasks(inp):
		ntasks=1
		for line in inp.split('\n'):
			if 'pool' in line: ntasks *= int(line.split()[-1])
		return ntasks			

	print(":: LetzElPhC el-ph calculation ($> tail -f lelphc.log) ::")

	# write input file
	f = open(inp_name, "w")
	f.write(inp_lelphc)
	f.close()

	# manage parallel run
	ntasks = get_ntasks(inp_lelphc)
	if ntasks > 1 : elph_run_str = f"mpirun -np {ntasks} {lelphc} -F {inp_name}"
	elif ntasks==1: elph_run_str = f"{lelphc} -F {inp_name}"
	else: raise ValueError(f'Wrong MPI task selection: {ntasks} tasks?')

	# run lelphc
	try:
		with open('lelphc.log','w') as log_file:
			elph_run=subprocess.run(elph_run_str,shell=True,check=True,capture_output=False,text=True,stdout=log_file)
	except subprocess.CalledProcessError as e:
		print("Error:", e)
		print("Description:", e.stderr)

def letzelph_to_yambo():

	print(":: Load el-ph database ::")
	lelph_obj = LetzElphElectronPhononDB('ndb.elph',div_by_energies=False)

	print(":: Convert el-ph data to yambo format ::")
	l2y       = ConvertElectronPhononDB(lelph_obj,'lelphc','SAVE','SAVE')

def clean_lelphc(no_lelphc_dbs,inp_name,ph_path):
	if not no_lelphc_dbs:
		print(":: LetzElPhC databases ::")
		print("     - lelphc.in input written")
		print("     - lelphc.log output written")
		print("     - ph_save | ndb.elph | ndb.Dmats written")
	else:
		[os.remove(C) for C in [inp_name,'ndb.elph','ndb.Dmats', 'lelphc.log'] if os.path.isfile(C)]
		if os.path.isdir(f"{ph_path}/ph_save"): shutil.rmtree(f"{ph_path}/ph_save")

if __name__=="__main__":

	parser = argparse.ArgumentParser(description='Generate electron-phonon coupling databases via LetzElPhC')
	parser.add_argument('-ph','--ph_inp_path', type=str, help='<Required> Path to ph.x (dvscf) input file',required=True)
	parser.add_argument('-b','--bands',nargs='2',type=str,help="<Required> First and last band (counting from 1), e.g. 'b_i b_f'",required=True)
	parser.add_argument('-k','--kernel', type=str, default='dfpt',help="<Optional> Electron-phonon kernel type, e.g. 'dfpt', 'bare', ... (default 'dfpt')")
	parser.add_argument('-par','--pools',nargs='2',type=str,default=[1,1],help="<Optional> MPI tasks as 'nqpools nkpools' (default serial)")
	parser.add_argument('-lelphc','--lelphc',type=str,default='lelphc',help="<Optional> Path to lelphc executable (default assumed in Path, otherwise prompted)")
	parser.add_argument('-nl','--no_lelphc_dbs', action="store_true", help="Remove lelphc databases (False if not given)")
	parser.add_argument('-ng','--no_gkkp', action="store_true", help="Do not generate gkkp dbs (False if not given)")
	args = parser.parse_args()

	phinp     = args.ph_inp_path
	bands     = args.bands
	kernel    = args.kernel
	pools     = args.pools
	lelphc    = args.lelphc
	no_lelphc = args.no_lelphc_dbs
	no_gkkp   = args.no_gkkp

	# Check inputs
	lelphc,ph_path,inp_ph,inp_lelphc,inp_name = checks(phinp,lelphc,bands,kernel,pools)

	# run preprocessing
	run_preprocessing(lelphc,ph_path,inp_ph)

	# run el-ph calculation and rotation
	run_elph(lelphc,inp_lelphc,inp_name,pools)

	# load database and convert to yambo format
	if not no_gkkp: letzelph_to_yambo()

	# clean
	clean_lelphc(no_lelphc_dbs,inp_name,ph_path)	
