"""
This python file corresponds to the S-ETC policy
"""

import numpy as np
import random
import time
import collections

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
# kappa: a constant in the algorithm
# rsd: np.random.seed and random.seed
def S_ETC(pos,v,r,T,simu_times,kappa,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()

        numer = collections.defaultdict(int)
        denom = collections.defaultdict(int)
        offer_t = collections.defaultdict(int)
        t = 1
        pr = 1

        # initialize, offer each product as an assortment until no-purchase
        while t <= T and pr <= N:
            expr_odr = [pr]
            loc = random.sample(range(1,1+K),1)
            expr_rev = gen_rev(expr_odr,pos,v,r,loc)
            prob = gen_probs(expr_odr,pos,v,loc)
            no_pur = False

            while t <= T and not no_pur:
                reg_reg[ii,t-1] = opt_rev - expr_rev
                offer_t[pr] += 1
                pur = np.random.multinomial(1,prob,1)
                purch = pur[0].tolist()
                t += 1
                if purch[0] == 0:
                    numer[pr] += 1
                else:
                    denom[pr] += 1
                    no_pur = True
            pr += 1

        # print("Numerator",numer)
        # print("Denominator",denom)
        # print("Offer Times",offer_t)

        # exploration-exploitation
        while t <= T:
            v_e = [1] * N
            for i in range(1,N+1):
                v_e[i-1] = numer[i]/denom[i]
            # print("v_Estimate", v_e)
            est_opt_odr, est_opt_rev = cands(est_pos, v_e, r)

            # print("Estimated Optimal Assortment",est_opt_odr)
            # print("Estimated Optimal Revenue",est_opt_rev)

            potential_opt_set = []
            for i in range(1, N + 1):
                if r[i - 1] >= est_opt_rev and offer_t[i] <= kappa * log(t):
                    potential_opt_set.append(i)

            # print("Potentially Optimal Products",potential_opt_set)

            if potential_opt_set:
                max_item = min(K,len(potential_opt_set))
                # print("Max Item Number",max_item)
                assort = random.sample(potential_opt_set,max_item)
            else:
                assort = est_opt_odr

            # print("Assortment",assort)
            for j in assort:
                offer_t[j] += 1
            # print("Offer Times", offer_t)

            loc = random.sample(range(1, K + 1), len(assort))
            cur_rev = gen_rev(assort, pos, v, r, loc)
            reg_reg[ii, t - 1] = opt_rev - cur_rev
            t += 1

            prob = gen_probs(assort, pos, v, loc)
            pur = np.random.multinomial(1, prob, 1)
            purch = pur[0].tolist()
            if purch[0] == 0:
                for j in range(len(assort)):
                    numer[assort[j]] += purch[j + 1]
            else:
                for j in range(len(assort)):
                    denom[assort[j]] += 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