"""
This python file corresponds to the A-UCB policy
"""

import numpy as np
import random
import time

from math import sqrt
from math import log

from Pre_Functions import gen_rev
from Pre_Functions import cands
from Pre_Functions import gen_probs


# main A-UCB-V policy
# pos: the position effect vector
# v: the preference vector
# r: the revenue vector
# T: horizon length
# simu_times: number of replications
# c_UCB: UCB constant ahead of the bonus
# rsd: np.random.seed and random.seed
def A_UCB(pos,v,r,T,simu_times,c_UCB,rsd):

    # calculate the true optimal assortment and positioning (opt_odr), and its expected revenue (opt_rev)
    opt_odr,opt_rev = cands(pos,v,r)
    print("Optimal Assortment and Positioning",opt_odr)
    print("Maximum Revenue",opt_rev)


    # initialize
    # N: number of items
    # K: number of positions
    # est_pos: misspecified position effects
    N = len(v)
    K = len(pos)
    est_pos = [1]*K


    # initialize data to save
    # reg_reg: expected regret at each time period, for each replication
    # time_consume: time consumption for each replication
    reg_reg = np.zeros((simu_times,T))
    time_consume = [0]*simu_times


    # random seed
    np.random.seed(rsd)
    random.seed(rsd)


    # algorithm starts
    print("Start")
    for ii in range(simu_times):
        cur_time = time.time()

        v_e = [1] * N
        est_odr, _ = cands(est_pos, v_e, r)
        loc = random.sample(range(1, K + 1), len(est_odr))
        est_rev = gen_rev(est_odr, pos, v, r, loc)
        prob = gen_probs(est_odr, pos, v, loc)

        l = len(est_odr)
        pur_num = [0] * l

        t = 1
        ell = 1
        ave = {}
        while t <= T:
            reg_reg[ii, t - 1] = opt_rev - est_rev
            pur = np.random.multinomial(1, prob, 1)
            purch = pur[0].tolist()
            if purch[0] == 0:
                for i in range(l):
                    pur_num[i] += purch[i+1]
            else:
                for i,c in enumerate(est_odr):
                    if c not in ave:
                        ave[c] = [pur_num[i], 1]
                    else:
                        ave[c][0] += pur_num[i]
                        ave[c][1] += 1
                for i in ave:
                    v_e[i - 1] = ave[i][0]/ave[i][1] + c_UCB * (
                                sqrt(ave[i][0]/ave[i][1] * 48 * log(1 + sqrt(N) * ell) / ave[i][1]) + 48 * log(1 + sqrt(N) * ell) /
                                ave[i][1])
                # print("Preference_UCB",v_e)
                est_odr, _ = cands(est_pos, v_e, r)
                loc = random.sample(range(1, K + 1), len(est_odr))
                est_rev = gen_rev(est_odr, pos, v, r, loc)
                l = len(est_odr)
                pur_num = [0] * l
                # print("Estimated Assortment",est_odr)
                prob = gen_probs(est_odr, pos, v, loc)
                # print("Probabilities",prob)
                ell += 1
            t += 1

        elapsed_time = time.time() - cur_time
        time_consume[ii] = elapsed_time
        print("Rep", ii + 1)
        print("Time Consumption", elapsed_time)

    print("End")
    print(time_consume)

    reg_mean = np.mean(reg_reg,axis=0)
    print(np.sum(reg_reg,axis=1))

    return reg_mean