"""
Extra functions about molecule handling.
"""
import os
from ..molecule import HostMolecule


def visual_opt_process(dir, output):
    """
    Use .mol files in 'dir' to generate .xyz trajectory file for visualization.
    """
    if not output.endswith('.xyz'):
        output += '.xyz'
    info = []
    for file in os.listdir(dir):
        if file.endswith('.mol'):
            temp_mol = HostMolecule.init_from_mol_file(f'{dir}/{file}')

            info.append(f'{temp_mol.get_atom_number()}\n')
            info.append(f'generated by PyHostGuest\n')
            positions = temp_mol.get_positions()
            for i, atom in enumerate(temp_mol.get_atoms()):
                x = positions[i][0]
                y = positions[i][1]
                z = positions[i][2]
                # 将x、y、z坐标保留到小数点后8位（包含0）
                info.append(f'{atom.get_element():<2}                 {x:<.8f}   {y:<.8f}   {z:<.8f}\n')


    with open(output, 'w') as f:
        f.writelines(info)
        f.close()
    print(f'Optimization trajectory has been written to {output}.')
