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

import numpy as np
import time

from math import sqrt
from math import log
from math import ceil

from Pre_Functions import rev
from Pre_Functions import cands
from Pre_Functions import probs
from Pre_Functions import no_trun


# main LR-UCB policy
# pos: the position effect vector
# v: the preference vector
# r: the revenue vector
# v_max: \nu_{0}, an upper bound for all items' preferences
# T: horizon length
# simu_times: number of replications
# lam: lambda
# c_UCB: UCB constant ahead of the bonus
def LR_UCB(pos,v,r,v_max,T,simu_times,lam,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
    # ini: number of initialization epochs
    N = len(v)
    K = len(pos)
    ini = ceil(N/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)


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

        # ell: track the current epoch
        # t: track the current time period
        # ave: record data to compute the UCB and the estimate
        ell = 1
        t = 1
        ave = {}

        # initial exploration
        while t <= T and ell <= ini:
            expr_odr = [((ell - 1) * K + i - 1) % N + 1 for i in range(1, K + 1)]
            expr_rev = rev(expr_odr, pos, v, r)
            prob = probs(expr_odr, pos, v)

            # count purchase numbers of the items when applying the strategy of fix-assortment-until-no-purchase
            pur = np.zeros(K + 1, dtype = int)
            while pur[0] == 0 and t <= T:
                reg_reg[ii, t - 1] = opt_rev - expr_rev
                t += 1
                pur_new = np.random.multinomial(1, prob)
                pur += pur_new

            if pur[0] == 0:
                break

            # update ave
            for i, c in enumerate(expr_odr):
                if c not in ave:
                    ave[c] = [no_trun(pur[i+1], 0) * pos[i], lam + pos[i] * pos[i], 1]
                else:
                    ave[c][2] += 1
                    ave[c][0] += no_trun(pur[i+1], log(ave[c][2]) / log(1+1/v_max)) * pos[i]
                    ave[c][1] += pos[i] * pos[i]

            ell += 1

        # UCB Procedure
        while t <= T:
            # construct UCB
            v_UCB = [0] * N
            for i in ave:
                v_UCB[i - 1] = ave[i][0] / ave[i][1] + c_UCB * sqrt(log(1 + ell * log(ell) * log(ell))) / sqrt(ave[i][2])
            # print("True Preferences", v)
            # print("Preference_UCB", v_UCB)
            cur_odr, _ = cands(pos, v_UCB, r)
            # print("UCB Assortment and Positioning", cur_odr)
            cur_rev = rev(cur_odr, pos, v, r)
            prob = probs(cur_odr, pos, v)

            # count purchase numbers of the items when applying the strategy of fix-assortment-until-no-purchase
            pur = np.zeros(len(cur_odr) + 1, dtype=int)
            while pur[0] == 0 and t <= T:
                reg_reg[ii, t - 1] = opt_rev - cur_rev
                t += 1
                pur_new = np.random.multinomial(1, prob)
                pur += pur_new

            if pur[0] == 0:
                break

            # update ave
            for i, c in enumerate(cur_odr):
                if c not in ave:
                    ave[c] = [no_trun(pur[i+1], 0) * pos[i], lam + pos[i] * pos[i], 1]
                else:
                    ave[c][2] += 1
                    ave[c][0] += no_trun(pur[i+1], log(ave[c][2]) / log(1+1/v_max)) * pos[i]
                    ave[c][1] += pos[i] * pos[i]

            ell += 1

        print("True Preferences", v)
        print("Preference_UCB", v_UCB)
        print("Offering Numbers", ave)
        print(cur_odr)

        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