import argparse
import rdkit.Chem as Chem
import os
import shutil
import subprocess
import sys

env = {}
env.update(os.environ)

def parse_arguments():
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('smiles', type=str, help="SMILES string, or file containing SMILES string (.smi extension)")
    parser.add_argument('--nconfs', type=int, default=10, help="Number of conformers generated by obabel. Default is 10")
    parser.add_argument('--ff', type=str, default="uff", help="Forcefield used by obabel for conformer generation and minimization. Default is UFF")
    parser.add_argument('--chrg', type=int, help="Charge used for crest screening. Specifying this overrides charge auto-detection")
    parser.add_argument('--theory', type=str, default="gfnff", help="Level of theory used in crest screening (gfnff or gfn2). Default is GFN-FF")
    parser.add_argument('--implicit', type=str, help="Implicit solvation model for crest screening (gbsa or alpb).")
    parser.add_argument('--solvent', type=str, help="Solvent used for crest screening (see crest docs for list of solvents).")
    parser.add_argument('--rescue', default=True, action='store_true', help="Enable rescue protocol when conformer generation fails. Enabled by default")
    parser.add_argument('--no-rescue', dest='rescue', action='store_false', help="Disable rescue protocol when conformer generation fails.")

    args = parser.parse_args()
    if args.nconfs <= 0:
        raise parser.error('Must have at least one conformer!')
    return args

def monitor_errors(proc):
    error_status = False
    while proc.poll() is None:
        line = proc.stdout.readline()
        print(line.decode('utf-8').strip('\n'))
        if b"CorrectStereoAtoms" in line:
            print("Stereochemistry conversion error.")
            sys.exit(1)
        if b"3D coordinate generation failed" in line:
            error_status = True
        if b"Initial conformer count: 0" in line:
            error_status = True
        if b"Initial geometry optimization failed" in line:
            error_status = True
    return error_status

def smiles_to_mol(smiles, ff):
    if ".smi" in smiles:
        input = ["-ismi", smiles]
    else:
        input = ["-:"+smiles]
    args = ["obabel"] + input + ["--gen3d", "--ff", ff, "-O", "step1.mol"]
    proc = subprocess.Popen(args, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    coord_gen_failed = monitor_errors(proc)
    proc.wait()
    if coord_gen_failed:
        print("Failed to generate 3D coordinates. Canonicalizing and retrying")
        smi_can = subprocess.Popen(["obabel"] + input + ["-ocan"], shell=False, stdout=subprocess.PIPE)
        can_sdf = subprocess.Popen(["obabel", "-ismi", "-h", "--gen2d", "-osdf"], shell=False, stdin=smi_can.stdout, stdout=subprocess.PIPE)
        sdf_mol = subprocess.Popen(["obabel", "-isdf", "--gen3d", "--ff", ff, "-omol", "-O", "step1.mol"], shell=False, stdin=can_sdf.stdout, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        canonical_gen_failed = monitor_errors(sdf_mol)
        sdf_mol.wait()
        if canonical_gen_failed:
            print("Canonical SMILES rescue failed.")
            sys.exit(2)

    # Check if step1.mol is generated
    if not(os.path.exists("step1.mol")) or os.path.getsize("step1.mol") == 0: 
        print("SMILES is invalid.")
        sys.exit(3)

    # obabel can generate poorly optimized in --gen3d
    # This ensures structure is good enough for conformer generation
    args = ["obabel", "step1.mol", "--minimize", "--ff", ff, "-O", "step2.mol"]
    proc = subprocess.Popen(args, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    proc.wait()

def obabel_conformers(nconfs, ff, rescue, retried):
    args = ["obabel", "step2.mol", "--conformer", "--nconf", nconfs, "--ff", ff, "--writeconformers", "-O", "conformers_obabel.xyz"]
    obabel_conf_proc = subprocess.Popen(args, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    conf_gen_failed = monitor_errors(obabel_conf_proc)
    obabel_conf_proc.wait()
    if conf_gen_failed: 
        print("Conformers were not properly generated!")
        if retried == True:
            print("Rescue failed. xtb structure will be sole conformer.") 
            shutil.copyfile('rescue/xtbopt.xyz', 'conformers_obabel.xyz')
            return 
        elif rescue == True:
            print("Attempting rescue protocol.")
            rescue_protocol()
            obabel_conformers(nconfs, ff, rescue, retried=True)
        else:
            print("Rescue protocol disabled. Quitting.")
            sys.exit(4)

def rescue_protocol():
    os.mkdir('rescue')
    # First, convert to xyz to get rid of any broken bonding
    xyz_conversion = subprocess.Popen(["obabel", "step2.mol", "-O", "rescue/step2b.xyz"], shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    xyz_conversion.wait()
    if os.path.isfile('.CHRG'):
        shutil.copyfile('.CHRG', 'rescue/.CHRG')
    if os.path.isfile('.UHF'):
        shutil.copyfile('.UHF', 'rescue/.UHF')
    
    # Then, optimize with xtb, which is not a force field and thus not bonding sensitive.
    os.chdir('rescue')
    xtb_opt_proc = subprocess.Popen(["xtb", "step2b.xyz", "--opt"], shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env)
    monitor_errors(xtb_opt_proc)
    xtb_opt_proc.wait()
    if not(os.path.exists(".xtboptok")):
        print("xtb optimization in rescue protocol failed.")
        sys.exit(5)
    os.chdir('..')

    # Convert back to mol, hoping that the bonding is now properly detected
    mol_conversion = subprocess.Popen(["obabel", "rescue/xtbopt.xyz", "-O", "step2.mol"], shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    mol_conversion.wait()

def write_charge(chrg):
    # If charge is not specified, automatically detect with rdkit
    # Less hacky ;)
    if not(chrg):
        chrg = (Chem.GetFormalCharge(Chem.MolFromMolFile("step1.mol")))
    # .CHRG file is autodetected by crest for optimization
    with open('.CHRG', 'w') as chrg_file:
        chrg_file.write(str(chrg))

def launch_crest(nconfs, ff, theory, implicit, solvent, rescue, retried=False):
    crest_args = ["crest", "--"+args.theory]
    if args.implicit is not None and args.solvent is not None:
        crest_args += ["--"+args.implicit, args.solvent]
    crest_args += ["--screen", "conformers_obabel.xyz"]
    crest_proc = subprocess.Popen(crest_args, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, env=env)
    crest_geom_error = monitor_errors(crest_proc)
    crest_proc.wait()
    if crest_geom_error:
        print("obabel generated conformer is bad.")
        if retried == True:
            print("Rescue protocol could not save crest conformer generation.")
            sys.exit(6)
        elif rescue == True:
            print("Attempting rescue protocol to regenerate conformers.")
            rescue_protocol()
            obabel_conformers(nconfs, ff, rescue, retried=True)
            launch_crest(nconfs, ff, theory, implicit, solvent, rescue, retried=True)
        else:
            print("Rescue protocol disabled. Quitting")
            sys.exit(7)

if __name__ == "__main__":
    args = parse_arguments()
    smiles_to_mol(args.smiles, args.ff)
    write_charge(args.chrg)
    obabel_conformers(str(args.nconfs), args.ff, args.rescue, retried=False)
    launch_crest(str(args.nconfs), args.ff, args.theory, args.implicit, args.solvent, args.rescue, retried=False)
