# Source code for macop.callbacks.UCBCheckpoint

    """UCB policy Checkpoint class implementation
    """

    # main imports
    import os
    import logging
    import numpy as np
    import pkgutil

    # module imports
    from .Callback import Callback
    from ..utils.color import macop_text, macop_line

    class UCBCheckpoint(Callback):
        """
        UCB checkpoint is used for loading previous UCB data and start again after loading checkpoint
        Need to be the same operators used during previous run

        Attributes:
            algo: {Algorithm} -- main algorithm instance reference
            every: {int} -- checkpoint frequency used (based on number of evaluations)
            filepath: {str} -- file path where checkpoints will be saved
        """
        def run(self):
            """
            Check if necessary to do backup based on `every` variable
            """
            # get current population
            currentEvaluation = self.algo.getGlobalEvaluation()

            # backup if necessary
            if currentEvaluation % self.every == 0:

                logging.info("UCB Checkpoint is done into " + self.filepath)

                with open(self.filepath, 'w') as f:

                    rewardsLine = ''

                    for i, r in enumerate(self.algo.policy.rewards):
                        rewardsLine += str(r)

                        if i != len(self.algo.policy.rewards) - 1:
                            rewardsLine += ';'

                    f.write(rewardsLine + '\n')

                    occurrencesLine = ''

                    for i, o in enumerate(self.algo.policy.occurrences):
                        occurrencesLine += str(o)

                        if i != len(self.algo.policy.occurrences) - 1:
                            occurrencesLine += ';'

                    f.write(occurrencesLine + '\n')

        def load(self):
            """
            Load backup lines as rewards and occurrences for UCB
            """
            if os.path.exists(self.filepath):

                logging.info('Load UCB data')
                with open(self.filepath) as f:

                    lines = f.readlines()
                    # read data for each line
                    rewardsLine = lines[0].replace('\n', '')
                    occurrencesLine = lines[1].replace('\n', '')

                    self.algo.policy.rewards = [
                        float(f) for f in rewardsLine.split(';')
                    ]
                    self.algo.policy.occurrences = [
                        float(f) for f in occurrencesLine.split(';')
                    ]

                print(
                    macop_text(
                        'Load of available UCB policy data from `{}`'.format(
                            self.filepath)))

            else:
                print(macop_text('No UCB data found, use default UCB policy'))
                logging.info("No UCB data found...")

            print(macop_line())