from sys import stdout
import time
from .algorithms import *
from .arms import *
from .core import *

import numpy as np
from pprint import pprint


class Simulation:

    """
    Class to combine the Bandit - Algorithm interaction.

    Args:
        bandit_inst: first argument is an instance of a Bandit (StoBandit, LinBandit, AdvBandit)
        algorithm_inst: second argument is an instance of an Algorithm

    Methods:
        ()
        reset()

    TODO:
    * what else to track? total reward, calculate regret, ...

    """

    def __init__(self, bandit_inst, algorithm_inst):
        """
        Initializes the Simulation class, and passes certain attributes so that
        Bandit objects and Algorithm objects have acces to certain information.
        """
        self.alg = algorithm_inst
        self.bandit = bandit_inst
        self.alg.bandit = bandit_inst
        self.total_regret = 0
        self.iterations = 0
        self.arm_info = [each_arm.info() for each_arm in self.bandit.arms]


    def reset(self):
        """
        Resets the Simulation class the Bandit subclass.
        """
        self.iterations = 0
        self.total_regret = 0
        self.bandit.reset()


    def _print_info(self):
        """
        Prints out information about the Simulation.

        The 'try' statement attempts to print out the .__name__ attribute of
        functions, rather than the object itself; if it fails, it falls back to
        just printing the object itself
        """
        print("\nAverage Regret: {0}".format(
            self.total_regret / self.iterations))

        for i in self.alg.var_dict:
            try:
                print("{0} : {1}".format(i, self.alg.var_dict[i].__name__))
            except AttributeError:
                print("{0} : {1}".format(i, self.alg.var_dict[i]))

        print("{0} Arms: {1}".format(self.bandit.n_arms, self.arm_info))
        print("Horizon: {0}".format(self.horizon))
        print ("Total runtime: {0} seconds".format(self.runtime))
        print ("Average runtime per iteration: {0} seconds".format(
            self.runtime / self.iterations))
        print("\n" + "-"*50 + "\n")
        return None


    def runStep(self, num, horizon):
        """
        Runs the Simulation for one step.

        This method is called during animations.
        """
        self.bandit.horizon = np.full(self.bandit.n_arms, horizon, dtype = int)
        self.horizon = horizon

        for m in range(num):
            self.bandit.pullArm(self.alg.giveArm())
        return None


    def runView(self, horizon):
        """
        Runs a single iteration to a horizon of n and prints out information.

        runView is especially useful for debugging bandit algorithms or seeing
        everything in significant detail as the bandit runs.
        """
        self.bandit.horizon = np.full(self.bandit.n_arms, horizon, dtype = int)
        self.horizon = horizon

        start_time = time.clock()
        for m in range(self.bandit.n_arms):
            self.bandit.pullArm(m)
            self.bandit.fullInfo()
        for m in range(horizon - self.bandit.n_arms):
            self.bandit.pullArm(self.alg.giveArm())
            self.bandit.fullInfo()
        stop_time = time.clock()

        self.total_regret += self.bandit.giveRegret()
        self.iterations += 1
        self.bandit.reset()
        self.runtime = stop_time - start_time
        self._print_info()
        return None


    def runStandard(self, horizon, cycles):
        """
        A standard experiment-running method.

        Runs the bandit for a certain number of cycles each to a certain
        horizon, and prints out the result on screen.

        It uses _print_info to print information about the bandit itself upon
        termination.
        """
        self.bandit.horizon = np.full(self.bandit.n_arms, horizon, dtype = int)
        self.horizon = horizon

        print("-"*50 + "\n")
        start_time = time.clock()

        for i in range(cycles):
            for i in range(self.bandit.n_arms):
                self.bandit.pullArm(i)

            for l in range(horizon - self.bandit.n_arms):
                self.bandit.pullArm(self.alg.giveArm())

            self.total_regret += self.bandit.giveRegret()
            self.iterations += 1
            self.bandit.reset()

            stdout.write(
                "\r----------  "
                "{0:d} out of {1} cycles"
                "----------".format(self.iterations, cycles))
            stdout.flush()

        stop_time = time.clock()
        self.runtime = stop_time - start_time
        self._print_info()
        return None


    def runHist(self, horizon, cycles, out_file):
        """
        Produces Histogram - like data.

        Every cycle, prints the bandit regret it out in the file, followed by a
        newline. These data files are used by many other methods containing
        hist in their name.
        """
        self.bandit.horizon = np.full(self.bandit.n_arms, horizon, dtype = int)
        self.horizon = horizon

        for i in range(cycles):
            for j in range(self.bandit.n_arms):
                self.bandit.pullArm(j)

            for l in range(horizon - self.bandit.n_arms):
                self.bandit.pullArm(self.alg.giveArm())

            self.total_regret += self.bandit.giveRegret()
            self.iterations += 1

            with open(out_file, "a") as outfile:
                outfile.write("{0}\n".format(self.bandit.giveRegret()))
            self.bandit.reset()
        return None


    def runVar(self, horizon, cycles, out_file):
        """
        Produces Variable - like data.

        Runs the simulation for a certain number of cycles, then appends that
        data to the text file
        """
        self.bandit.horizon = np.full(self.bandit.n_arms, horizon, dtype = int)
        self.horizon = horizon

        for i in range(cycles):
            for i in range(self.bandit.n_arms):
                self.bandit.pullArm(i)

            for l in range(horizon - self.bandit.n_arms):
                self.bandit.pullArm(self.alg.giveArm())

            self.total_regret += self.bandit.giveRegret()
            self.iterations += 1

            self.bandit.reset()

        with open(out_file, "a") as outfile:
            outfile.write("{0}\n".format(self.total_regret/self.iterations))
        return None

    def selfCheck(self):
        """
        Checks the simulation for errors
        """
        errors = ""
        return errors

class ObjectDict:
    ArmDict = {
        'Bernoulli' : BernoulliArm,
        'Normal' : NormalArm,
        'Linear' : LinearArm
    }
    AlgDict = {
        'var_dict' : {},
        'random' : random,
        'greedy' : greedy,
        'greedy_ep' : greedy_ep,
        'UCB' : UCB,
        'UCB_KL' : UCB_KL,
        'TS_Beta' : TS_Beta,
        'TS_Gauss' : TS_Gauss,
        'Bayes_Gauss' : Bayes_Gauss,
        'UCB_Lin' : UCB_Lin,
        'TS_Lin' : TS_Lin
    }
    IndexDict = {
        'B1' : B1,
        'B2' : B2,
        'B3' : B3,
        'B4' : B4,
        'B5' : B5,
        'B6' : B6,
        'B7' : B7,
        'B8' : B8,
        'B9' : B9
    }


def ReMapSim(sim_dict):
    """
    Takes a simulation dictionary and turns string names into objects. It uses
    the ObjectDict class in order to create references to the actual objects.
    """
    for alg_key in list(sim_dict['Algorithm'].keys()):
        if alg_key == 'incr':
            sim_dict['Algorithm'][alg_key] = \
                ObjectDict.IndexDict[sim_dict['Algorithm'][alg_key]]
        elif alg_key == 'algtype':
            sim_dict['Algorithm'][alg_key] = \
                ObjectDict.AlgDict[sim_dict['Algorithm'][alg_key]]

    sim_dict['Algorithm'] = Algorithm(**sim_dict['Algorithm'])
    sim_dict['arm_object_list'] = [ObjectDict.ArmDict[arm[0]](arm[1])
        for arm in sim_dict['Bandit']['ArmList']]

    # different parsing for Linear and non-Linear bandits
    if sim_dict['Bandit']['ArmList'][0][0] in ('Linear'):
        sim_dict['vector_mean'] = sim_dict['Bandit']['MeanVector']
        sim_dict['Bandit'] = LinBandit(
            sim_dict['arm_object_list'],
            sim_dict['vector_mean'],
            normalized=sim_dict['Normalized'])
        del[sim_dict['vector_mean']]
    elif sim_dict['Bandit']['ArmList'][0][0] in ('Bernoulli', 'Normal'):
        sim_dict['Bandit'] = StoBandit(sim_dict['arm_object_list'])

    sim_dict['Simulation'] = Simulation(
        sim_dict['Bandit'],
        sim_dict['Algorithm'])

    # clean up residual keys
    del[sim_dict['Bandit']]
    del[sim_dict['Algorithm']]
    del[sim_dict['arm_object_list']]

    return None
