"""Learner that implements a greedy learning algorithm"""

import time

from pebl import network, result, evaluator
from pebl.util import *
from pebl.learner.base import *

class GreedyLearnerStatistics:
    def __init__(self):
        self.restarts = -1
        self.iterations = 0
        self.unimproved_iterations = 0
        self.best_score = 0
        self.start_time = time.time()

    @property
    def runtime(self):
        return time.time() - self.start_time

class GreedyLearner(Learner):
    #
    # Parameters
    #
    _params =  (
        config.IntParameter(
            'greedy.max_iterations',
            """Maximum number of iterations to run.""",
            default=1000
        ),
        config.IntParameter(
            'greedy.max_time',
            """Maximum learner runtime in seconds.""",
            default=0
        ),
        config.IntParameter(
            'greedy.max_unimproved_iterations',
            """Maximum number of iterations without score improvement before
            a restart.""", 
            default=500
        ),
        config.StringParameter(
            'greedy.seed',
            'Starting network for a greedy search.',
            default=''
        )
    )

    def __init__(self, data_=None, prior_=None, **options):
        """
        Create a learner that uses a greedy learning algorithm.

        The algorithm works as follows:

            1. start with a random network
            2. Make a small, local change and rescore network
            3. If new network scores better, accept it, otherwise reject.
            4. Steps 2-3 are repeated till the restarting_criteria is met, at
               which point we begin again with a new random network (step 1)
        
        Any config param for 'greedy' can be passed in via options.
        Use just the option part of the parameter name.

        For more information about greedy learning algorithms, consult:

            1. http://en.wikipedia.org/wiki/Greedy_algorithm
            2. D. Heckerman. A Tutorial on Learning with Bayesian Networks. 
               Microsoft Technical Report MSR-TR-95-06, 1995. p.35.
            
        """

        super(GreedyLearner, self).__init__(data_, prior_)
        self.options = options
        config.setparams(self, options)
        if not isinstance(self.seed, network.Network):
            self.seed = network.Network(self.data.variables, self.seed)
        
    def run(self):
        """Run the learner.

        Returns a LearnerResult instance. Also sets self.result to that
        instance.  
        
        """

        # max_time and max_iterations are mutually exclusive stopping critera
        if 'max_time' not in self.options:
            _stop = self._stop_after_iterations
        else:
            _stop = self._stop_after_time
            
        self.stats = GreedyLearnerStatistics()
        self.result = result.LearnerResult(self)
        self.evaluator = evaluator.fromconfig(self.data, self.seed, self.prior)
        self.evaluator.score_network(self.seed.copy())

        first = True
        self.result.start_run()
        while not _stop():
            self._run_without_restarts(_stop, self._restart, 
                                       randomize_net=(not first))
            first = False
        self.result.stop_run()

        return self.result

    def _run_without_restarts(self, _stop, _restart, randomize_net=True):
        self.stats.restarts += 1
        self.stats.unimproved_iterations = 0

        if randomize_net:
            self.evaluator.randomize_network()
         
        # set the default best score
        self.stats.best_score = self.evaluator.score_network()

        # continue learning until time to stop or restart
        while not (_restart() or _stop()):
            self.stats.iterations += 1

            try:
                curscore = self._alter_network_randomly_and_score()
            except CannotAlterNetworkException:
                return
            
            self.result.add_network(self.evaluator.network, curscore)

            if curscore <= self.stats.best_score:
                # score did not improve, undo network alteration
                self.stats.unimproved_iterations += 1
                self.evaluator.restore_network()
            else:
                self.stats.best_score = curscore
                self.stats.unimproved_iterations = 0

    #
    # Stopping and restarting criteria
    # 
    def _stop_after_time(self):
        return self.stats.runtime >= self.max_time

    def _stop_after_iterations(self):
        return self.stats.iterations >= self.max_iterations

    def _restart(self):
        return self.stats.unimproved_iterations >= self.max_unimproved_iterations

