rmsd

#!python

from rich.progress import Progress
from rich.console import Console
from rich.table import Table
from rich import box

from pathlib import Path
from InquirerPy import inquirer

import os
import glob
import subprocess
import argparse
import json
import subprocess
import operator


PYTHON_PATH = "/home/vport/projects/scripts/.scripts_environment/bin/python"
SCRIPT_PATH = "/home/vport/projects/scripts/calculate_rmsd"
CHEMCRAFT_PATH = "/home/vport/.wine/drive_c/chemcraft/Chemcraft.exe"


CACHE_FILE = ".rmsd"


class Rmsd:
    def __init__(self, file1, file2, rmsd):
        self.file1 = file1
        self.file2 = file2
        self.rmsd = rmsd

    def toJSON(self):
        return json.dumps(self.__dict__)

    def __eq__(self, other):
        return self.file1 == other.file1 and self.file2 == other.file2

    def calculate_rmsd(self):
        if self.rmsd is None:
            command = f"{PYTHON_PATH} {SCRIPT_PATH} {self.file1} {self.file2} --reorder"
            result = subprocess.run(command, shell=True, capture_output=True, text=True)
            output = result.stdout.strip()

            if output:
                try:
                    # Extrair o valor de RMSD do output
                    rmsd = float(output.split()[-1])
                except ValueError:
                    print(f"Erro ao converter RMSD para float: {output}")
                    rmsd = None
            else:
                print("Erro: Saída do cálculo do RMSD está vazia.")
                rmsd = None

            self.rmsd = rmsd

    def view(self):
        file1 = ["wine", CHEMCRAFT_PATH, self.file1]
        file2 = ["wine", CHEMCRAFT_PATH, self.file2]
        subprocess.Popen(file1, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        subprocess.Popen(file2, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def handle_args():
    parser = argparse.ArgumentParser(
        description="Calculate RMSD between files matching a pattern."
    )

    parser.add_argument(
        "-r",
        "--remove",
        action="store_true",
        help="Remove every file that has the root of the same xyz file that was filtered.",
    )

    parser.add_argument(
        "-t",
        "--tolerance",
        type=float,
        help="Tolerance for RMSD comparison.",
    )

    parser.add_argument(
        "-v",
        "--view",
        action="store_true",
        help="Open files in chemcraft for inspection.",
    )

    parser.add_argument(
        "files",
        nargs="+",
        help="List of xyz files (at least two are required), format is mandatory to be .xyz",
    )

    return parser.parse_args()


def load_cache():
    rmsds = []
    try:
        with open(CACHE_FILE) as cache_data:
            for blob in json.load(cache_data):
                rmsds.append(Rmsd(**blob))
    except FileNotFoundError:
        print("No cache file found.")
    return rmsds


def clean_file(file, rmsds):
    for blob in rmsds:
        if blob.file1 == file or blob.file2 == file:
            rmsds.remove(blob)

    # clean all files that contain the rootname of the xyz file
    rootname = Path(file).stem

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


def generate_rmsd_list(cache, files):
    for idx1 in range(0, len(files)):
        file_1 = files[idx1]
        for idx2 in range(idx1 + 1, len(files)):
            file_2 = files[idx2]

            obj = Rmsd(file_1, file_2, None)

            if obj in cache:
                pass
            else:
                cache.append(obj)

    return cache


def gen_rich_table(rmsds, tolerance):
    table = Table(box=box.MINIMAL)
    table.add_column("File 1", justify="center")
    table.add_column("File 2", justify="center")
    table.add_column("RMSD", justify="center")

    for obj in rmsds:
        if tolerance is not None:
            if obj.rmsd < tolerance:
                table.add_row(obj.file1, obj.file2, "{:.3f}".format(obj.rmsd))
        else:
            table.add_row(obj.file1, obj.file2, "{:.3f}".format(obj.rmsd))

    return table


def ask_view_pair(rmsds, tolerance):
    rmsd_choices = []
    for obj in rmsds:
        if tolerance is not None:
            if obj.rmsd < tolerance:
                rmsd_choices.append(obj)
        else:
            rmsd_choices = rmsds

    string_choices = [
        "file1: {}, file2: {}, rmsd: {}".format(i.file1, i.file2, i.rmsd)
        for i in rmsd_choices
    ]

    choice = inquirer.fuzzy(
        message="Select a file pair",
        choices=string_choices,
        default=None,
    ).execute()

    return rmsd_choices[string_choices.index(choice)]


if __name__ == "__main__":
    args = handle_args()
    cache = load_cache()

    rmsds = generate_rmsd_list(cache, args.files)

    total_tasks = len(rmsds)

    filtered = []

    try:
        with Progress() as progress:
            task = progress.add_task("", total=total_tasks)
            for obj in rmsds:
                obj.calculate_rmsd()
                progress.update(task, advance=1)
                if args.tolerance is not None:
                    if obj.rmsd < args.tolerance:
                        filtered.append(obj.file2)

        if args.remove:
            for file in filtered:
                clean_file(file, rmsds)

        with open(CACHE_FILE, "w") as c:
            json.dump(rmsds, c, default=lambda o: o.__dict__, indent=4)

    except KeyboardInterrupt:
        with open(CACHE_FILE, "w") as c:
            json.dump(rmsds, c, default=lambda o: o.__dict__, indent=4)

    rmsds.sort(key=operator.attrgetter('rmsd'), reverse = True)
    table = gen_rich_table(rmsds, args.tolerance)
    console = Console()
    console.print(table)

    if args.view:
        another = True
        while another:
            pair = ask_view_pair(rmsds, args.tolerance)
            pair.view()

            another = inquirer.confirm(
                message="Open another pair?", default=False
            ).execute()