# code adapted from Markovic et al. (2021)
# https://github.com/dimarkov/aibandits/blob/master/inference_algorithms_comparison.ipynb

import numpy as np
from scipy.special import digamma, betaln
from scipy.special import betaincinv
import matplotlib.pyplot as plt
import pandas as pd
import math

# utility function
def softmax(a):
    a_max = max(a)
    x = np.exp(a-a_max)
    u = np.sum(x)
    return x/u

# define Agent 
class Agent():
    # what you should do at first
    def initialize_parameters(self, alpha, beta, mu, eta, omega, lam, a, b):
        self.alpha = alpha
        self.beta = beta
        self.mu = mu
        self.eta = eta
        self.omega = omega
        self.lam = lam
        self.alpha0 = alpha
        self.a = a
        self.b = b
        self.m = a/(a + b)

    # expected free energy is all thing that matters
    def G(self, alpha_t, beta_t, lam):
        alpha = math.exp(2 * lam)

        nu_t = alpha_t + beta_t
        mu_t = alpha_t / nu_t
    
        KL_a = - betaln(alpha_t, beta_t) + (alpha_t - alpha) * digamma(alpha_t)\
                 + (beta_t - 1) * digamma(beta_t) + (alpha + 1 - nu_t) * digamma(nu_t)
    
        H_a = - mu_t * digamma(alpha_t + 1) - (1-mu_t) * digamma(beta_t + 1) + digamma(nu_t + 1)
        return KL_a + H_a 
    

    # an agent takes action according to expected free energy 
    def calc_action_prob(self, G):
        return softmax(-G)


    # response probabilities according to softmax rule
    def take_action(self, prob):
        choice = np.random.multinomial(1, prob, size=1)
        c = np.where(choice[0] == 1)
        return int(list(c)[0])


    # after getting observations, parameters would be updated
    def update_parameters(self, o, c):
        # o : observation
        # c : choice
        alpha = self.alpha
        beta = self.beta
        mu = self.mu
        eta = self.eta
        omega = self.omega
        alpha0 = self.alpha0 
        a = self.a
        b = self.b
        m = self.m

        # these updates are according to Markovic et al. (2021)
        ll0 = mu[c] * o + (1 - mu[c]) * (1-o)
        ll1 = .5

        norm = ll1 * omega + ll0 * (1 - omega)
        et_post = ll1 * omega / norm
        om_post = m * (1 - et_post)
        self.alpha[c] = alpha[c] * ( 1 - et_post) + et_post * alpha0[c] + o
        self.beta[c] = beta[-1] * (1 - et_post) + et_post * beta[0] + (1 - o)
        self.mu = alpha/(alpha + beta)
        self.eta[c] = et_post
        self.omega = om_post
        self.a =  a + omega
        self.b =  b + 1 - eta[c] - omega
        self.m = a/(a + b)







#---------------------------------------------------------------------------------------#
class Environment_VI():
    # what you should do at first
    def initialize_enviroment(self, VI1, VI2, NUM_RFT, MINIMUM):
        RFT_TIME1 = np.random.exponential(VI1 - MINIMUM, NUM_RFT) + MINIMUM
        RFT_TIME2 = np.random.exponential(VI2 - MINIMUM, NUM_RFT) + MINIMUM
        
        self.RFT_TIME1 =  RFT_TIME1
        self.RFT_TIME2 =  RFT_TIME2
        self.iri1 = 0 # current inter-reinforcement interval
        self.iri2 = 0
        self.time = 0 # global_timer
        self.reward_set1 = 0
        self.reward_set2 = 0
        self.current_r1 = 0 # which RFTtime is used now
        self.current_r2 = 0 # which RFTtime is used now

    # update time
    def update_time(self, DT):
        self.iri1 = self.iri1 + DT
        self.iri2 = self.iri2 + DT
        self.time = self.time + DT

    # check reward preparation given schedule
    def reward_setting(self):
        if self.RFT_TIME1[self.current_r1] <= self.iri1:
            self.reward_set1 = 1
        else:
            self.reward_set1 = 0
            
        if self.RFT_TIME2[self.current_r2] <= self.iri2:
            self.reward_set2 = 1
        else:
            self.reward_set2 = 0
            
    # receive action and return observation o (reward or nothing)
    def return_observation(self, c):
        # c should be choce of agent generated by Agent.take_action()
        if c==0:
            if self.reward_set1 == 1:
                self.iri1 = 0
                self.current_r1 +=1
                o = 1
            else:
                o = 0

        if c==1:
            if self.reward_set2 == 1:
                self.iri2 = 0
                self.current_r2 +=1
                o = 1
            else:
                o = 0
        
        return o


#---------------------------------------------------------------------------------------#






'''

### quick test ###
# define agent
agent = Agent()
# initialize parameters
agent.initialize_parameters(alpha=np.array([1., 1.]), beta=np.array([1., 1.]), \
mu=np.array([.5, .5]), eta=np.array([.0, .0]), omega=.0, lam=.4, a=1., b=20.)
agent.alpha
agent.omega
# expected free energy
G = agent.G(agent.alpha, agent.beta, agent.lam)
# choice probability
prob = agent.calc_action_prob(G)
# perform choice
c = agent.take_action(prob)
# if o = 1 (rewarded), update parameter accordingly
o = 1
agent.update_parameters(o, c)

# environment
env = Environment_VI()
env.initialize_enviroment(10, 20, 10, 5)
np.mean(env.RFT_TIME1)
np.mean(env.RFT_TIME2)
env.RFT_TIME1[0]
env.update_time(25.)
env.iri1
env.reward_setting()
env.reward_set1
env.reward_set2
env.RFT_TIME1[env.current_r1]
env.return_observation(0)
'''