"""
This python file corresponds to the EI-TLR policy
"""

import numpy as np
import time

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

from Pre_Functions import dist
from Pre_Functions import exist_zero
from Pre_Functions import min_optimize
from Pre_Functions import rev
from Pre_Functions import rev_new
from Pre_Functions import cands
from Pre_Functions import cands_new
from Pre_Functions import probs
from Pre_Functions import probs_new
from Pre_Functions import soft_trun


# main EI-TLR 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
# thr: threshold in applying Alternate Minimization Algorithm
# exp_c: constant affecting number of exploration epochs
# end_cri: binary ending criteria for the exploration phase
# default value of end_cri should be 1 as used for simulation, relaxed as 0 for real application
# rsd: np.random.seed
def EI_TLR(pos,v,r,v_max,T,simu_times,lam,c_UCB,thr,exp_c,end_cri,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, fused into exploration epochs
    # exp_epoch: number of minimal exploration epochs
    N = len(v)
    K = len(pos)
    ini = ceil(N/K)
    exp_epoch = ceil(exp_c*T**(2/3))
    print("Exploration Epochs",exp_epoch)


    # initialize data to save
    # reg: expected regret at each time period, for each replication
    # pos_est_err: estimation errors of the position effects from the exploration phase data
    # pre_est_err: estimation errors of the preferences from the exploration phase data
    # real_exp_epoch: real exploration epochs
    # real_exp_time: real exploration time periods
    # time_consume: time consumption for each replication
    # itn_simu: iteration times within the Alternate Minimization algorithm
    reg = np.zeros((simu_times,T))
    reg_exp2 = np.zeros((simu_times, 2))
    pos_est_err = np.zeros(simu_times)
    pre_est_err = np.zeros(simu_times)
    real_exp_epoch = np.zeros(simu_times, dtype = int)
    real_exp_time = np.zeros(simu_times, dtype = int)
    itn_simu = np.zeros(simu_times, dtype = int)
    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
        ell = 1
        t = 1

        # record data from the exploration phase, for future use in the UCB phase
        rec = {}

        # pq: mean-purchase matrix
        # pq_t: matrix of total number of purchases
        # pq_n: number-of-epoch matrix
        # pq_e: either pq_t or pq_n, depending on the binary ending criteria end_cri
        pq = np.zeros((N, K))
        pq_t = np.zeros((N, K), dtype = int)
        pq_n = np.zeros((N, K), dtype = int)
        pq_e = np.zeros((N, K), dtype = int)

        # exploration phase
        while t <= T and (ell <= max(ini,exp_epoch) or exist_zero(pq_e) == 0):
            expr_odr = [((ell - 1) + i - 1) % N + 1 for i in range(1, K + 1)]
            # print("Assortment and Positioning",expr_odr)
            expr_rev = rev(expr_odr, pos, v, r)
            prob = probs(expr_odr, pos, v)

            pur = np.zeros(K + 1, dtype = int)
            while pur[0] == 0 and t <= T:
                reg[ii, t - 1] += opt_rev - expr_rev
                reg_exp2[ii, 0] += opt_rev - expr_rev
                t += 1
                pur_new = np.random.multinomial(1, prob)
                pur += pur_new

            if pur[0] == 0:
                break

            rec[ell] = (expr_odr, pur)

            for i,c in enumerate(expr_odr):
                pq_t[c - 1][i] += pur[i + 1]
                pq_n[c - 1][i] += 1

            if end_cri == 1:
                pq_e = pq_t
            else:
                pq_e = pq_n

            ell += 1

        exp_ell = ell - 1
        real_exp_epoch[ii] = exp_ell
        print("Real exploration epochs", exp_ell)
        print("Expected exploration epochs", exp_epoch)
        real_exp_time[ii] = t - 1
        print("Real exploration time periods", t - 1)

        # calculate the pq matrix
        for i in range(N):
            for k in range(K):
                pq[i][k] = pq_t[i][k]/pq_n[i][k]

        # conduct Alternate Minimization Algorithm to estimate the preferences and position effects
        alpha, beta, itn = min_optimize(pq, thr)
        pos_est = np.append(1, beta)
        pre_est = alpha
        print("Iteration times", itn)
        itn_simu[ii] = itn

        # calculate estimation errors of preferences and position effects
        pre_est_err[ii] = dist(pre_est, np.array(v))
        pos_est_err[ii] = dist(pos_est, np.array(pos))

        # incorporate the exploration phase data as the foundation for generalized TLR-UCB procedure
        pre_est = pre_est.tolist()
        pos_est = pos_est.tolist()
        pos_max = max(pos_est)
        pos_min = min(pos_est)

        # ave: record data to compute the UCB and the estimate
        ave = {}
        for j in range(exp_ell):
            expr_odr, pur = rec[j + 1]
            for i, c in enumerate(expr_odr):
                if c not in ave:
                    ave[c] = [soft_trun(pur[i + 1], 0) * pos_est[i], lam + pos_est[i] * pos_est[i], 1]
                else:
                    ave[c][2] += 1
                    ave[c][0] += soft_trun(pur[i + 1], log(ave[c][2]) / log(1 + 1 / v_max)) * pos_est[i]
                    ave[c][1] += pos_est[i] * pos_est[i]

        # generalized TLR-UCB procedure
        while t <= T:
            print(t)
            v_UCB = [0] * N
            for i in ave:
                v_UCB[i - 1] = ave[i][0] / ave[i][1] + c_UCB * (2 * (log(T) / log(1 + 1 / v_max)) * sqrt(
                    2 * log(N * T) + log(1 + ave[i][2] * pos_max ** 2 / lam)) + (v_max + 1) * pi / sqrt(6) + sqrt(
                    lam) * v_max) / (pos_min * sqrt(ave[i][2]))
            cur_odr = cands_new(pos_est, v_UCB, r)
            # print("UCB Assortment and Positioning", cur_odr)
            cur_rev = rev_new(cur_odr, pos, v, r)
            prob = probs_new(cur_odr, pos, v)

            pur = np.zeros(len(cur_odr) + 1, dtype = int)
            while pur[0] == 0 and t <= T:
                reg[ii, t - 1] = opt_rev - cur_rev
                reg_exp2[ii, 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, (u, w) in enumerate(cur_odr.items()):
                if u not in ave:
                    ave[u] = [soft_trun(pur[i + 1], 0) * pos_est[w], lam + pos_est[w] * pos_est[w], 1]
                else:
                    ave[u][2] += 1
                    ave[u][0] += soft_trun(pur[i + 1], log(ave[u][2]) / log(1 + 1 / v_max)) * pos_est[w]
                    ave[u][1] += pos_est[w] * pos_est[w]

            ell += 1

        print("Offering Numbers", pq_n)
        print("True Preferences", v)
        print("Preferences Estimates", pre_est)
        print("True Position Effects", pos)
        print("Position Effects Estimate", pos_est)

        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,axis=0)
    print(np.sum(reg,axis=1))

    reg_exp2_mean = np.mean(reg_exp2, axis=0)
    print(reg_exp2_mean)

    return reg_mean, reg_exp2_mean