import numpy as np
from copy import deepcopy
from joblib import Parallel, delayed
import pandas as pd
import time

class BaseSimuation():

    def __repr__(self):
        return '{}'.format(self.__class__.__name__)

class SimpleSimulation(BaseSimuation):
    """Basic class for a sigle simulation for simple (non-contextual) environment and a particular policy
    Parameters:
        environment: Environment class instance
        policy : Policy class instance
        horizont: number of iterations
        seed: seed
        record_every: how often to record results from simulations
    """

    def __init__(self, environment, policy, horizont=10000, seed = None, record_every = 1):

        self.environment = environment
        self.policy = policy
        self.horizont = horizont
        self.seed = seed
        self.t = 0

        self.cumulative_reward = 0
        self.cumulative_regret = 0
        self.best_arm_pulls = 0
        self.best_arm = environment.get_optimal_arm()
        self.best_mean = environment.get_optimal_mean()
        self.record_every = record_every

        self.results = []

    def run_step(self):

        # pick action (arm) and observe reward
        arm_idx = self.policy.pick_action()
        arm_reward = self.environment.get_stochastic_reward(arm_idx)
        # policy(agent) observes reward
        self.policy.observe_reward(arm_idx, arm_reward)

        # accumulate data for regret
        self.cumulative_reward += arm_reward
        instant_regret = self.best_mean - self.environment.means[arm_idx]
        self.cumulative_regret += instant_regret
        self.best_arm_pulls += 1 if arm_idx == self.best_arm else 0
        if (self.t+1) % self.record_every == 0:
            self._data = {'t': self.t+1,
                     'instant_regret': instant_regret,
                     'cumulative_regret': self.cumulative_regret,
                     'best_arm_pulls': self.best_arm_pulls,
                     }
            self.results.append(self._data)

        self.t += 1

    def run(self):
        np.random.seed(self.seed)
        for t in range(self.horizont):
            self.run_step()

    def to_pandas(self):
        return pd.DataFrame(self.results)

def _run_simulation(simulation, sid = None, full_results = True):
    """ Helper function for Parralel computation of several simulations.
    Returns list of dictionaries of results for a particular simulation with id and policy
    type added
    """
    simulation.run()
    if full_results:
        result = simulation.results
    else:
        result = [simulation.results[-1]]
    # put simulation identifier
    for step_results in result:
        step_results.update({'id':sid,
                          'policyType': simulation.policy.__class__.__name__,
                          })
    return result

class ParallelSimulation(BaseSimuation):

    def __init__(self, simulation, n_simulations = 1, n_jobs = 1, full_results=True):
        """Runs a particular simulation n_simulations times

        """
        self.simulation = simulation
        self.n_simulations = n_simulations
        self.n_jobs = n_jobs
        self.ids = [f'id-{i+1}' for i in range(n_simulations)]
        self.full_results = full_results

    def run(self, verbose = 0):
        self.results = Parallel(n_jobs=self.n_jobs, verbose=verbose)(delayed(_run_simulation)(
                deepcopy(self.simulation), self.ids[i], self.full_results)
                for i in range(self.n_simulations))

    def to_pandas(self):
        if not hasattr(self, 'results'):
            raise Exception('Run the simulation to get results')
        else:
            return pd.concat([pd.DataFrame(res) for res in self.results])

    def get_results(self):
        return self.results

class ContextualSimulation():
    """Sigle simulation for contextual environment and a particular policy
    Parameters:
        environment: Environment class instance
        policy : Policy class instance
        horizont: number of iterations
        seed: seed
        record_every: how often to record results from simulations
    """

    def __init__(self, environment, policy, horizont=10000, seed = None, record_every = 1):

        self.environment = environment
        self.policy = policy
        self.horizont = horizont
        self.seed = seed
        self.t = 0

        self.cumulative_reward = 0
        self.cumulative_regret = 0
        self.best_arm_pulls = 0
#        self.best_arm = None
#        self.best_mean = None
        self.record_every = record_every

        self.results = []

    def run_step(self):
        # environment generates context and rewards (unknown to the agent/policy)
        context = self.environment.generate_context()
        # pick action (arm) and observe reward
        arm_idx = self.policy.pick_action(context)
        arm_reward = self.environment.get_stochastic_reward(arm_idx)
        # agent/agent observes reward
        self.policy.observe_reward(arm_idx, arm_reward)

        # accumulate data for regret
        self.cumulative_reward += arm_reward
        instant_regret = self.environment.current_optimal_mean - self.environment.current_means[arm_idx]
        self.cumulative_regret += instant_regret
        self.best_arm_pulls += 1 if arm_idx == self.environment.current_optimal_arm else 0

        if (self.t+1) % self.record_every == 0:
            self._data = {'t': self.t+1,
                         'instant_regret': instant_regret,
                         'cumulative_regret': self.cumulative_regret,
                         'best_arm_pulls': self.best_arm_pulls,
                         'pulled_arm' : arm_idx,
                         'reward': arm_reward,
                         'current_best_arm': self.environment.current_optimal_arm
                         }
            self.results.append(self._data)

        self.t += 1

    def run(self):
        np.random.seed(self.seed)
        for t in range(self.horizont):
            self.run_step()

    def to_pandas(self):
        return pd.DataFrame(self.results)