<commit_before>"""
Class which contains various methods to print the output of different
ways of using FLARE, such as training a GP from an AIMD run,
or running an MD simulation updated on-the-fly.
"""
import datetime
import logging
import time
import numpy as np

from logging import FileHandler, StreamHandler, Logger
from os.path import isfile
from shutil import move as movefile
from typing import Union

from flare.struc import Structure
from flare.utils.element_coder import Z_to_element


# Unit conversions.
eva_to_gpa = 160.21766208


class Output:
    """
    This is an I/O class that hosts the log files for OTF and Trajectories
    class. It is also used in get_neg_like_grad and get_neg_likelihood in
    gp_algebra to print intermediate results.

    It opens and print files with the basename prefix and different
    suffixes corresponding to different kinds of output data.

    :param basename: Base output file name, suffixes will be added
    :type basename: str, optional
    :param verbose: print level. The same as logging level. It can be
                    CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
    :type verbose: str, optional
    :param always_flush: Always write to file instantly
    :type always_flus: bool, optional
    """

    def __init__(self, basename: str = 'otf_run',
                 verbose: str = 'INFO',
                 print_as_xyz: bool = False,
                 always_flush: bool = False):
        """
        Construction. Open files.
        """
        self.basename = f"{basename}"
        self.print_as_xyz = print_as_xyz
        self.always_flush = always_flush

        filesuffix = {'log': '.out', 'hyps': '-hyps.dat'}
        if print_as_xyz:
            filesuffix['xyz'] = '.xyz'

        self.logger = []
        for filetype in filesuffix:
            self.open_new_log(filetype, filesuffix[filetype], verbose)

    def conclude_run(self):
        """
        destruction function that closes all files
        """

        logger = logging.getLogger(self.basename+'log')
        logger.info('-' * 20)
        logger.info('Run complete.')
        logging.shutdown()
        self.logger = []

    def open_new_log(self, filetype: str, suffix: str, verbose='info'):
        """
        Open files.  If files with the same
        name are exist, they are backed up with a suffix "-bak".

        :param filetype: the key name for logging
        :param suffix: the suffix of the file to be opened
        :param verbose: the verbose level for the logger
        """

        if filetype not in self.logger:
            set_logger(self.basename+filetype, stream=False,
                       fileout_name=self.basename+suffix,
                       verbose=verbose)
            self.logger += [filetype]

    def write_to_log(self, logstring: str, name: str = "log",
                     flush: bool = False):
        """
        Write any string to logfile

        :param logstring: the string to write
        :param name: the key name of the file to logger named 'log'
        :param flush: whether it should be flushed
        """
        logger = logging.getLogger(self.basename+name)
        logger.info(logstring)

        if flush or self.always_flush:
            logger.handlers[0].flush()

    def write_header(self, gp_str: str,
                     dt: float = None,
                     Nsteps: int = None, structure: Structure = None,
                     std_tolerance: Union[float, int] = None,
                     optional: dict = None):
        """
        TO DO: this should be replace by the string method of GP and OTF, GPFA

        Write header to the log function. Designed for Trajectory Trainer and
        OTF runs and can take flexible input for both.

        :param gp_str: string representation of the GP
        :param dt: timestep for OTF MD
        :param Nsteps: total number of steps for OTF MD
        :param structure: initial structure
        :param std_tolerance: tolarence for active learning
        :param optional: a dictionary of all the other parameters
        """

        f = logging.getLogger(self.basename+'log')
        f.info(f'{datetime.datetime.now()}')

        if isinstance(std_tolerance, tuple):
            std_string = 'Relative uncertainty tolerance: ' \
                         f'{std_tolerance[0]} times noise hyperparameter \n'
            std_string += 'Absolute uncertainty tolerance: ' \
                          f'{std_tolerance[1]} eV/A\n'
        elif std_tolerance < 0:
            std_string = \
                f'Uncertainty tolerance: {np.abs(std_tolerance)} eV/A\n'
        elif std_tolerance > 0:
            std_string = \
                f'Uncertainty tolerance: {np.abs(std_tolerance)} ' \
                'times noise hyperparameter \n'
        else:
            std_string = ''

        headerstring = '\n'
        headerstring += gp_str
        headerstring += '\n'
        headerstring += std_string
        if dt is not None:
            headerstring += f'Timestep (ps): {dt}\n'
        headerstring += f'Number of frames: {Nsteps}\n'
        if structure is not None:
            headerstring += f'Number of atoms: {structure.nat}\n'
            headerstring += \
                f'System species: {set(structure.species_labels)}\n'
            headerstring += 'Periodic cell (A): \n'
            headerstring += str(np.array(structure.cell))+'\n'

        if optional:
            for key, value in optional.items():
                headerstring += f"{key}: {value} \n"

        # report previous positions
        if structure is not None:
            headerstring += '\nPrevious positions (A):\n'
            for i in range(len(structure.positions)):
                headerstring += f'{structure.species_labels[i]:5}'
                for j in range(3):
                    headerstring += f'{structure.prev_positions[i][j]:10.4f}'
                headerstring += '\n'
        headerstring += '-' * 80

        f.info(headerstring)

        if self.always_flush:
            f.handlers[0].flush()


    def write_md_header(self, dt, curr_step, dft_step):
        string = ''
        # Mark if a frame had DFT forces with an asterisk
        if not dft_step:
            string += '-' * 80 + '\n'
            string += f"-Frame: {curr_step} "
        else:
            string += f"\n*-Frame: {curr_step} "

        string += f'\nSimulation Time: {(dt * curr_step):.3} ps \n'
        return string


    def write_md_config(
            self, dt, curr_step, structure, temperature, KE, start_time, dft_step,
            velocities):
        """ write md configuration in log file

        :param dt: timestemp of OTF MD
        :param curr_step: current timestep of OTF MD
        :param structure: atomic structure
        :param temperature: current temperature<commit_msg>report more decimals in simulation time<commit_after>"""
Class which contains various methods to print the output of different
ways of using FLARE, such as training a GP from an AIMD run,
or running an MD simulation updated on-the-fly.
"""
import datetime
import logging
import time
import numpy as np

from logging import FileHandler, StreamHandler, Logger
from os.path import isfile
from shutil import move as movefile
from typing import Union

from flare.struc import Structure
from flare.utils.element_coder import Z_to_element


# Unit conversions.
eva_to_gpa = 160.21766208


class Output:
    """
    This is an I/O class that hosts the log files for OTF and Trajectories
    class. It is also used in get_neg_like_grad and get_neg_likelihood in
    gp_algebra to print intermediate results.

    It opens and print files with the basename prefix and different
    suffixes corresponding to different kinds of output data.

    :param basename: Base output file name, suffixes will be added
    :type basename: str, optional
    :param verbose: print level. The same as logging level. It can be
                    CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET
    :type verbose: str, optional
    :param always_flush: Always write to file instantly
    :type always_flus: bool, optional
    """

    def __init__(self, basename: str = 'otf_run',
                 verbose: str = 'INFO',
                 print_as_xyz: bool = False,
                 always_flush: bool = False):
        """
        Construction. Open files.
        """
        self.basename = f"{basename}"
        self.print_as_xyz = print_as_xyz
        self.always_flush = always_flush

        filesuffix = {'log': '.out', 'hyps': '-hyps.dat'}
        if print_as_xyz:
            filesuffix['xyz'] = '.xyz'

        self.logger = []
        for filetype in filesuffix:
            self.open_new_log(filetype, filesuffix[filetype], verbose)

    def conclude_run(self):
        """
        destruction function that closes all files
        """

        logger = logging.getLogger(self.basename+'log')
        logger.info('-' * 20)
        logger.info('Run complete.')
        logging.shutdown()
        self.logger = []

    def open_new_log(self, filetype: str, suffix: str, verbose='info'):
        """
        Open files.  If files with the same
        name are exist, they are backed up with a suffix "-bak".

        :param filetype: the key name for logging
        :param suffix: the suffix of the file to be opened
        :param verbose: the verbose level for the logger
        """

        if filetype not in self.logger:
            set_logger(self.basename+filetype, stream=False,
                       fileout_name=self.basename+suffix,
                       verbose=verbose)
            self.logger += [filetype]

    def write_to_log(self, logstring: str, name: str = "log",
                     flush: bool = False):
        """
        Write any string to logfile

        :param logstring: the string to write
        :param name: the key name of the file to logger named 'log'
        :param flush: whether it should be flushed
        """
        logger = logging.getLogger(self.basename+name)
        logger.info(logstring)

        if flush or self.always_flush:
            logger.handlers[0].flush()

    def write_header(self, gp_str: str,
                     dt: float = None,
                     Nsteps: int = None, structure: Structure = None,
                     std_tolerance: Union[float, int] = None,
                     optional: dict = None):
        """
        TO DO: this should be replace by the string method of GP and OTF, GPFA

        Write header to the log function. Designed for Trajectory Trainer and
        OTF runs and can take flexible input for both.

        :param gp_str: string representation of the GP
        :param dt: timestep for OTF MD
        :param Nsteps: total number of steps for OTF MD
        :param structure: initial structure
        :param std_tolerance: tolarence for active learning
        :param optional: a dictionary of all the other parameters
        """

        f = logging.getLogger(self.basename+'log')
        f.info(f'{datetime.datetime.now()}')

        if isinstance(std_tolerance, tuple):
            std_string = 'Relative uncertainty tolerance: ' \
                         f'{std_tolerance[0]} times noise hyperparameter \n'
            std_string += 'Absolute uncertainty tolerance: ' \
                          f'{std_tolerance[1]} eV/A\n'
        elif std_tolerance < 0:
            std_string = \
                f'Uncertainty tolerance: {np.abs(std_tolerance)} eV/A\n'
        elif std_tolerance > 0:
            std_string = \
                f'Uncertainty tolerance: {np.abs(std_tolerance)} ' \
                'times noise hyperparameter \n'
        else:
            std_string = ''

        headerstring = '\n'
        headerstring += gp_str
        headerstring += '\n'
        headerstring += std_string
        if dt is not None:
            headerstring += f'Timestep (ps): {dt}\n'
        headerstring += f'Number of frames: {Nsteps}\n'
        if structure is not None:
            headerstring += f'Number of atoms: {structure.nat}\n'
            headerstring += \
                f'System species: {set(structure.species_labels)}\n'
            headerstring += 'Periodic cell (A): \n'
            headerstring += str(np.array(structure.cell))+'\n'

        if optional:
            for key, value in optional.items():
                headerstring += f"{key}: {value} \n"

        # report previous positions
        if structure is not None:
            headerstring += '\nPrevious positions (A):\n'
            for i in range(len(structure.positions)):
                headerstring += f'{structure.species_labels[i]:5}'
                for j in range(3):
                    headerstring += f'{structure.prev_positions[i][j]:10.4f}'
                headerstring += '\n'
        headerstring += '-' * 80

        f.info(headerstring)

        if self.always_flush:
            f.handlers[0].flush()


    def write_md_header(self, dt, curr_step, dft_step):
        string = ''
        # Mark if a frame had DFT forces with an asterisk
        if not dft_step:
            string += '-' * 80 + '\n'
            string += f"-Frame: {curr_step} "
        else:
            string += f"\n*-Frame: {curr_step} "

        string += f'\nSimulation Time: {(dt * curr_step):10.3f} ps \n'
        return string


    def write_md_config(
            self, dt, curr_step, structure, temperature, KE, start_time, dft_step,
            velocities):
        """ write md configuration in log file

        :param dt: timestemp of OTF MD
        :param curr_step: current timestep of OTF MD
        :param structure: atomic structure
        :param temperature: current temperature
{'fraction_non_alphanumeric': 0.09066355208409373, 'fraction_numerical': 0.004379881743192934, 'mean_word_length': 2.0269553689792312, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '569473', 'n_tokens_mistral': 3899, 'n_tokens_neox': 3682, 'n_words': 1276}