"""
This python file prepares functions needed to define main policies
"""

import numpy as np
from math import sqrt

# function that calculates the Euclidean distance between two vectors
def dist(a,b):
    return sqrt(np.inner(a-b,a-b))


# function that judges whether there exist zero entries in a matrix
# used in the ending criteria for the exploration phase in EI-TLR
def exist_zero(pq):
    N = len(pq[:, 0])
    K = len(pq[0, :])
    for i in range(N):
        for k in range(K):
            if pq[i][k] == 0:
                return 0
    return 1


# function that calculates the value of the loss function
def L(alpha,beta,pq):
    N = len(alpha)
    K = len(beta) + 1
    s = 0
    for i in range(N):
        s += (pq[i,0]-alpha[i])**2
        for k in range(1,K):
            s += (pq[i,k]-alpha[i]*beta[k-1])**2
    return s


# function that calculates one marginal optimization procedure in the Alternate Minimization Algorithm
def H_1(alpha,pq):
    K = len(pq[0,:])
    beta = np.zeros(K-1)
    denominator = np.inner(alpha,alpha)
    for k in range(1,K):
        beta[k-1] = np.inner(alpha,pq[:,k])/denominator
    return beta


# function that calculates one marginal optimization procedure in the Alternate Minimization Algorithm
def H_2(beta,pq):
    N = len(pq[:,0])
    alpha = np.zeros(N)
    denominator = 1 + np.inner(beta,beta)
    for i in range(N):
        alpha[i] = (pq[i,0] + np.inner(beta,pq[i,1:]))/denominator
    return alpha


# function that conducts the Alternate Minimization Algorithm
# pq: the mean-purchase matrix Q in the manuscript
# threshold: accuracy level to end the alternate minimization procedure
def min_optimize(pq,threshold):
    N = len(pq[:, 0])
    alpha = np.zeros(N)
    alpha_new = np.ones(N)
    n = 0
    while np.inner(alpha_new - alpha, alpha_new - alpha) > threshold**2:
        n += 1
        alpha = alpha_new
        beta = H_1(alpha,pq)
        alpha_new = H_2(beta,pq)
    alpha = alpha_new
    beta = H_1(alpha,pq)
    return alpha,beta,n


# function that maps an assortment and positioning (without gap, encoded as a list) to its expected revenue
# ass: assortment and positioning; example: ass = [2,1] refers to the assortment {1,2}, where item 2 is placed on
# position 1, and item 1 is placed on position 2
# pos: the position effect vector
# v: the preference vector
# r: the revenue vector
def rev(ass,pos,v,r):
    s = 0
    deno = 1
    for i,c in enumerate(ass):
        s += r[c-1]*pos[i]*v[c-1]
        deno += pos[i]*v[c-1]
    return s/deno


# a more general function that maps an assortment and positioning to its expected revenue
# ass: assortment
# loc: corresponding positions for the items in the assortment, i.e., ass
# pos: the position effect vector
# v: the preference vector
# r: the revenue vector
def gen_rev(ass,pos,v,r,loc):
    s = 0
    deno = 1
    for i in range(len(ass)):
        s += r[ass[i]-1]*pos[loc[i]-1]*v[ass[i]-1]
        deno += pos[loc[i]-1]*v[ass[i]-1]
    return s/deno


# function that maps an assortment and positioning (encoded as a dictionary instead of a list) to its expected revenue
# ass: assortment and positioning; example: ass = {2:0,1:1} refers to the assortment {1,2}, where item 2 is placed on
# position 1(=0+1), and item 1 is placed on position 2(=1+1)
# pos: the position effect vector
# v: the preference vector
# r: the revenue vector
def rev_new(ass,pos,v,r):
    s = 0
    deno = 1
    for i, (u, w) in enumerate(ass.items()):
        s += r[u-1]*pos[w]*v[u-1]
        deno += pos[w]*v[u-1]
    return s/deno


# function that solves the static optimization problem, i.e., outputs the assortment and positioning
# with maximal expected revenue given position effect vector, preference vector and revenue vector
# pos: the position effect vector, restrained to be sorted in decreasing order
# v: the preference vector
# r: the revenue vector
def cands(pos,v,r):
    K = len(pos)
    N = len(v)

    v_r = [0] * N
    for i in range(N):
        v_r[i] = v[i] * r[i]
    v_s = list(zip(v, v_r, range(1, N + 1)))
    v_s.sort(reverse=True)
    odr = [z for x, y, z in v_s]

    interset = []
    for i in range(N):
        interset.append([r[i],0,i+1])
        for j in range(i+1,N):
            if v[i] != v[j]:
                ipt = (v[i]*r[i]-v[j]*r[j])/(v[i]-v[j])
                interset.append([ipt,i+1,j+1])
    interset.sort()

    cand_odr = [odr[:K]]
    odr_idx = [0]*N
    for i,c in enumerate(odr):
        odr_idx[c-1] = i

    negnum = 0
    cur = 0
    while negnum < N:
        if interset[cur][1] == 0:
            negnum += 1
        else:
            i_l,j_l = interset[cur][1],interset[cur][2]
            odr[odr_idx[i_l-1]],odr[odr_idx[j_l-1]] = j_l,i_l
            odr_idx[i_l-1],odr_idx[j_l-1] = odr_idx[j_l-1],odr_idx[i_l-1]
        cand_odr.append(odr[:min(K, N - negnum)])
        cur += 1
    cand_odr.pop()

    res = []
    maxm = -float("inf")
    maxm_idx = -1
    for i in range(len(cand_odr)):
        ass = cand_odr[i]
        cur_res = rev(ass,pos,v,r)
        res.append(cur_res)
        if cur_res > maxm:
            maxm = cur_res
            maxm_idx = i

    opt_odr = cand_odr[maxm_idx]
    opt_rev = rev(opt_odr,pos,v,r)
    return opt_odr,opt_rev


# a second function that solves the static optimization problem,
# but more tailored to estimated position effect,
# and outputs the optimal assortment and positioning as a dictionary
# pos: the (estimated) position effect vector, thus not necessarily in a decreasing order
# v: the preference vector
# r: the revenue vector
def cands_new(pos,v,r):
    K = len(pos)
    N = len(v)

    v_r = [0] * N
    for i in range(N):
        v_r[i] = v[i] * r[i]
    v_s = list(zip(v, v_r, range(1, N + 1)))
    v_s.sort(reverse=True)
    odr = [z for x, y, z in v_s]

    interset = []
    for i in range(N):
        interset.append([r[i],0,i+1])
        for j in range(i+1,N):
            if v[i] != v[j]:
                ipt = (v[i]*r[i]-v[j]*r[j])/(v[i]-v[j])
                interset.append([ipt,i+1,j+1])
    interset.sort()

    cand_odr = [odr[:K]]
    odr_idx = [0]*N
    for i,c in enumerate(odr):
        odr_idx[c-1] = i

    negnum = 0
    cur = 0
    while negnum < N:
        if interset[cur][1] == 0:
            negnum += 1
        else:
            i_l,j_l = interset[cur][1],interset[cur][2]
            odr[odr_idx[i_l-1]],odr[odr_idx[j_l-1]] = j_l,i_l
            odr_idx[i_l-1],odr_idx[j_l-1] = odr_idx[j_l-1],odr_idx[i_l-1]
        cand_odr.append(odr[:min(K, N - negnum)])
        cur += 1
    cand_odr.pop()

    res = []
    maxm = -float("inf")
    maxm_idx = -1
    for i in range(len(cand_odr)):
        ass = cand_odr[i]
        cur_res = rev(ass,pos,v,r)
        res.append(cur_res)
        if cur_res > maxm:
            maxm = cur_res
            maxm_idx = i

    opt_odr = cand_odr[maxm_idx]
    new_odr = {}

    dc_1 = {}
    for i, y in enumerate(pos):
        dc_1[y] = i

    pos_sort = sorted(pos, reverse=True)

    for i, y in enumerate(pos_sort):
        if i < len(opt_odr):
            new_odr[opt_odr[i]] = dc_1[y]

    return new_odr


# function that maps an assortment and positioning to the customer choice probabilities
# ass: the assortment and positioning (without gap, encoded as a list)
# pos: the position effect vector
# v: the preference vector
def probs(ass,pos,v):
    l = len(ass)
    prob = [0] * (l+1)
    deno = 1
    for i, c in enumerate(ass):
        deno += pos[i] * v[c - 1]
    prob[0] = 1/deno
    for i in range(1,l+1):
        prob[i] = pos[i-1] * v[ass[i-1]-1]/deno
    return prob


# a more general function that maps an assortment and positioning to the customer choice probabilities
# ass: assortment
# loc: corresponding positions for the items in the assortment, i.e., ass
# pos: the position effect vector
# v: the preference vector
def gen_probs(ass,pos,v,loc):
    l = len(ass)
    prob = [0] * (l+1)
    deno = 1
    for i in range(l):
        deno += pos[loc[i]-1] * v[ass[i]-1]
    prob[0] = 1/deno
    for i in range(1,l+1):
        prob[i] = pos[loc[i-1]-1] * v[ass[i-1]-1]/deno
    return prob


# a second function that maps an assortment and positioning to the customer choice probabilities
# ass: the assortment and positioning (encoded as a dictionary instead of a list)
# pos: the position effect vector
# v: the preference vector
def probs_new(ass,pos,v):
    l = len(ass.items())
    prob = [0] * (l+1)
    deno = 1
    for i, (u, w) in enumerate(ass.items()):
        deno += pos[w] * v[u - 1]
    prob[0] = 1/deno
    for i, (u, w) in enumerate(ass.items()):
        prob[i+1] = pos[w] * v[u - 1]/deno
    return prob


# to apply truncation in the algorithm
def soft_trun(x,a):
    if abs(x) > a:
        return np.sign(x)*a
    else:
        return x


# to apply no truncation in the algorithm
def no_trun(x,a):
    return x


# implement simple linear regression
def lin_reg(x,y):
    x_bar = np.mean(x)
    y_bar = np.mean(y)
    s_xy = 0
    s_xx = 0
    for i in range(len(x)):
        s_xy += (x[i]-x_bar)*(y[i]-y_bar)
        s_xx += (x[i]-x_bar)*(x[i]-x_bar)
    beta = s_xy/s_xx
    alpha = y_bar - beta*x_bar
    return alpha, beta


# accumulate period-wise regrets
def acu(reg):
    l = len(reg)
    for i in range(1,l):
        reg[i] += reg[i-1]
    return reg


# shrink the preference vector towards its minimum, to vary the suboptimality gap
# c: contraction parameter
def contract_v(v,c):
    v_m = min(v)
    N = len(v)
    v_new = [0]*N
    for i in range(N):
        v_new[i] = v_m + c*(v[i]-v_m)
    return v_new