boltzmann

#!/home/vport/projects/scripts/.venv/bin/python

import os
import argparse
import glob

from pathlib import Path
from math import exp

from rich.console import Console
from rich.table import Table

K = 0.001987204259
XTBHESSEXT = ".sph"


class GibbsEnergyNotFound(Exception):
    pass


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Calculate the boltzmann distribution of a set of ORCA logfiles."
    )

    # Add temperature argument with a default value of 293 Kelvin
    parser.add_argument(
        "-t",
        "--temperature",
        type=float,
        default=293,
        help="Temperature in Kelvin (default: 293)",
    )

    # Add temperature argument with a default value of 293 Kelvin
    parser.add_argument(
        "-x",
        "--xtbhess",
        action="store_true",
        help="Use the single point hessian from xtb",
    )

    # Add temperature argument with a default value of 293 Kelvin
    parser.add_argument(
        "-l",
        "--latex",
        action="store_true",
        help="Print a latex formatted table",
    )

    parser.add_argument(
        "-r",
        "--remove",
        type=float,
        default=0,
        help="Remove all files above a certain threshhold",
    )

    parser.add_argument(
        "files", nargs="+", help="List of files (at least two are required)"
    )

    # Parse arguments
    args = parser.parse_args()

    # Check if at least two files are provided
    if len(args.files) < 2:
        parser.error("At least two files are required.")

    return args


def parse_gibbs_orca(file: str):
    energy = 0
    with open(file, "r") as f:
        lines = f.readlines()

    for line in lines:
        if "Final Gibbs free energy" in line:
            energy = float(line.split()[-2])
        # else:

    if energy == 0:
        return GibbsEnergyNotFound

    return energy


def parse_gibbs_orca_xtb(file: str):
    hessfile = Path(file).stem + XTBHESSEXT

    electronic_energy = 0
    grrho = 0

    with open(file, "r") as f:
        lines = f.readlines()

    for line in lines:
        if "FINAL SINGLE POINT ENERGY" in line:
            electronic_energy =  float(line.split()[-1])

    with open(hessfile, "r") as f:
        hesslines = f.readlines()

    for line in hesslines:
        if "G(RRHO) contrib" in line:
            grrho = float(line.replace(":", "").rstrip().split(" ")[-2])

    return electronic_energy + grrho

def generate_latex_table(populations, energy_dict, temperature):
    latex_table = []
    
    # LaTeX table header
    latex_table.append("\\begin{table}[h]")
    latex_table.append("\\centering")
    latex_table.append(f"\\caption{{Boltzmann population obtained at {temperature} K of the conformers for }}")
    latex_table.append("\\label{tab:conformer_population}")
    latex_table.append("\\begin{tabular}{c c c c}")
    latex_table.append("\\toprule")
    latex_table.append("Filename & Pop (\\%) & $\\Delta G$ (kcal/mol) & Rank \\\\")
    
    # Add rows to the LaTeX table
    count = 1
    for k, v in populations.items():
        filename = Path(k).stem.replace("_", " ")
        pop_percent = "{:.2f}\\%".format(v * 100)
        delta_g = "{:.3f}".format(energy_dict[k])
        rank = "{}".format(count)
        
        latex_table.append(f"{filename} & {pop_percent} & {delta_g} & {rank} \\\\")
        count += 1
    
    # LaTeX table footer
    latex_table.append("\\bottomrule")
    latex_table.append("\\end{tabular}")
    latex_table.append("\\end{table}")
    
    # Print the LaTeX table
    print("\n".join(latex_table))


def clean_files(energy_dict: dict[str, float], threshold):
    remove = [k for k, v in energy_dict.items() if v > threshold]
    remove = [Path(file).stem for file in remove]

    for rootname in remove:
        for filename in glob.glob(f"./{rootname}*"):
            os.remove(filename)


def gibbs_factor(energy, temperature):
    return exp(-energy / (K * temperature))


def main():
    args = parse_arguments()
    if args.xtbhess:
        energy_dict = {
            file: parse_gibbs_orca_xtb(file) * 627.503 for file in args.files
        }
    else:
        energy_dict = {file: parse_gibbs_orca(file) * 627.503 for file in args.files}

    energy_dict = {
        file: (energy - min(energy_dict.values()))
        for file, energy in energy_dict.items()
    }

    exponentials = {}

    for k, v in energy_dict.items():
        exponentials[k] = gibbs_factor(v, args.temperature)

    normalization_factor = sum(exponentials.values())
    populations = {k: (v / normalization_factor) for k, v in exponentials.items()}
    populations = {
        k: v
        for k, v in sorted(populations.items(), key=lambda item: item[1], reverse=True)
    }

    console = Console()

    table = Table()
    table.add_column("Filename")
    table.add_column("Pop (%)")
    table.add_column("deltaG")
    table.add_column("Rank")
    count = 1

    if args.latex:
        generate_latex_table(populations, energy_dict, args.temperature)

    else:
        for k, v in populations.items():
            table.add_row(
                f"{Path(k).stem}",
                "{:.2f}%".format(v * 100),
                "{:.3f} kcal/mol".format(energy_dict[k]),
                "{}".format(count),
            )
            count += 1

        console.print(table)

    if args.remove != 0:
        clean_files(energy_dict, args.remove)


# Example usage
if __name__ == "__main__":
    main()