"""
This python file corresponds to EI-TLR's exploration phase and estimation errors for preferences and position effects
"""

import numpy as np
import time

from Pre_Functions import dist
from Pre_Functions import exist_zero
from Pre_Functions import min_optimize
from Pre_Functions import probs


# partial EI-TLR policy, contains only exploration phase and joint estimation of preferences and position effects
# pos: the position effect vector
# v: the preference vector
# exp_epoch: number of minimum exploration epochs
# simu_times: number of replications
# thr: threshold in applying Alternate Minimization Algorithm
# 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_Explore(pos,v,exp_epoch,simu_times,thr,end_cri,rsd):

    # N: number of items
    # K: number of positions
    N = len(v)
    K = len(pos)


    # initialize data to save
    # 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
    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)


    # random seed
    np.random.seed(rsd)


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

        # ell: track the current epoch
        # t: track the current time period
        ell = 1
        t = 1

        # 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 ell <= 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)
            prob = probs(expr_odr, pos, v)

            pur = np.zeros(K + 1, dtype=int)
            while pur[0] == 0:
                t += 1
                pur_new = np.random.multinomial(1,prob)
                pur += pur_new

            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

        real_exp_epoch[ii] = ell - 1
        # 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, _ = min_optimize(pq, thr)
        pos_est = np.append(1, beta)
        pre_est = alpha

        # 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))

        # print("Offering Numbers", pq_n)
        # print("True Preferences", v)
        # print("Preferences Estimate", pre_est)
        # print("True Position Effects", pos)
        # print("Position Effects Estimate", pos_est)
        # print("Rep", ii + 1)

    print("End")
    elapsed_time = time.time() - cur_time
    print("Time Consumption", elapsed_time)

    pre_est_err_mean = np.mean(pre_est_err)
    pos_est_err_mean = np.mean(pos_est_err)

    return pre_est_err_mean, pos_est_err_mean