# SPDX-FileCopyrightText: 2026 Yating Zhang
# SPDX-License-Identifier: MIT

import numpy as np 
import pandas as pd  
import matplotlib.pyplot as plt  
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
import seaborn as sns 
from matplotlib import gridspec
import collections

import os

import platform 
import subprocess
 

ROOT = os.path.dirname(os.path.dirname(__file__))   # project_root 
DATA_DIR = os.path.join(ROOT, "data")
RESULT_DIR = os.path.join(ROOT, "results", "case_study" )
LOG_DIR = os.path.join(ROOT, RESULT_DIR, "logs" )


# ===============================================
# Data Import and Initialization
# ===============================================

# Import revenue data and base utility data
rdf = pd.read_csv(os.path.join(DATA_DIR,"revenue.csv"))  
u_base = pd.read_csv(os.path.join(DATA_DIR,"utility.csv" )) 
assortment = np.load(os.path.join(DATA_DIR,"ast_matrix.npy") )

# Set benchmark days and initial beta parameters
bmk_days = [5,6,7,8,9] 
beta_init = [-0.49216979 ,-0.35776757 ,-0.28904397 ] 

# Data preprocessing
r = np.array(rdf)[:,0]  # Extract revenue data
base_utility = np.array(u_base)[:,3]  # Extract base utility data
bmk_menu = assortment[bmk_days,:].T  # Get benchmark menu and transpose
print(f"Benchmark menu shape: {bmk_menu.shape}")

# Set model parameters
(N,T) = bmk_menu.shape  # N: number of products, T: time periods
M = 3  # Model parameter

# ===============================================
# Import Different Menu Strategy Data
# ===============================================

menu_list = []

# Import MNL model menu
astdf = pd.read_csv( os.path.join(LOG_DIR,"case_menu_mnl.csv"), header=None)
menu0 = np.array(astdf)
menu_list.append(menu0)

# Import M=1,2,3 menu strategies
for m in [1,2,3]:
    astdf = pd.read_csv(os.path.join(LOG_DIR,"case_menu_m"+str(m)+".csv"), header=None)
    menu0 = np.array(astdf)
    menu_list.append(menu0)

# Import cyclic menu strategies
for m in [1,2,3]:
    astdf = pd.read_csv(os.path.join(LOG_DIR,"case_cycle_menu_"+str(m)+".csv"), header=None)
    menu0 = np.array(astdf)
    menu_list.append(menu0)

# ===============================================
# Parameter Setup and Data Preprocessing
# ===============================================

# Initialize beta parameter matrix
beta = np.zeros((N,M))
for i in range(N):
    beta[i,:] = beta_init 

# Data ordering (reverse order)
order_r = np.array([r[N-i-1] for i in range(N)])*15  # Price data (scaled by 15)
order_u = np.array([base_utility[N-i-1] for i in range(N)])  # Base utility data
order_a = np.exp(order_u)  # Base attraction values

# Define menu strategy names
menu_name_all  = ['MNL', 'M=1','M=2','M=3','Practice',"M=1 cyclic","M=2 cyclic","M=3 cyclic"]
menu_name_opt  = ['MNL', 'M=1','M=2','M=3', "M=1 cyclic","M=2 cyclic","M=3 cyclic"]

# ===============================================
# Analysis Function Definitions
# ===============================================

def HHI(XX):
    """
    Calculate Herfindahl-Hirschman Index (HHI) and number of unique items
    HHI measures market concentration, higher values indicate higher concentration
    
    Args:
        XX: Menu matrix
    
    Returns:
        [hhi, number of unique items]
    """
    XX = XX.T
    prods = []
    items = set()
    
    # Count all products offered across 5 days
    for t in range(5):
        ast = np.where(XX[t]>0)[0]  # Find products offered on day t
        prods.extend(ast)
        for i in ast:
            items.add(i)
    
    # Calculate product frequency
    frequency = collections.Counter(prods)
    dic_fre = dict(frequency)
    
    # Calculate HHI index
    hhi = 0
    for i in items:
        hhi += (dic_fre[i]/len(prods))**2
    
    return [hhi, len(items)]

def gene_rev_sale(xx, M, D, N):
    """
    Generate revenue and sales data
    Uses Multinomial Logit model to calculate expected revenue and sales
    
    Args:
        xx: Menu matrix
        M: Model parameter
        D: Number of days
        N: Number of products
    
    Returns:
        (revenue, sales)
    """
    xx = xx.T 
    extend_menu = np.zeros(shape=(N, D+M))
    
    # Build extended menu matrix
    for i in range(N):
        for t in range(D):
            extend_menu[i, t+M] = xx[t, i]
    
    sale = 0  # Total sales
    rev = 0   # Total revenue
    
    # Calculate daily revenue and sales
    for t in range(D):
        # Calculate attraction values
        a = np.zeros(N)
        for i in range(N):
            a[i] = np.exp(base_utility[i] + extend_menu[i, t:t+M]@beta[i,:].T)
        
        # Calculate total attraction
        totala = 1 + sum(a[i]*extend_menu[i,t+M] for i in range(N))
        
        # Calculate choice probabilities
        prob = [a[i]*extend_menu[i,t+M]/totala for i in range(N)]
        
        # Accumulate revenue and sales
        rev += prob@r.T 
        sale += prob@np.ones(N).T 
    
    return rev, sale

# ===============================================
# Menu Strategy Comparison Analysis
# ===============================================

# Calculate revenue and sales for each menu strategy
rev_list = []
sale_list = []

for k in range(len(menu_list)):
    rev, sale = gene_rev_sale(menu_list[k], M, T, N)
    rev_list.append(rev)
    sale_list.append(sale)

# Calculate benchmark menu revenue and sales
rev_bmk, sale_bmk = gene_rev_sale(bmk_menu, M, T, N) 

print("Strategy revenues:", rev_list)
print("Benchmark revenue:", rev_bmk)

# Calculate HHI index for each strategy
opt_hhi = [] 
for xx in menu_list:
    res = HHI(xx)
    opt_hhi.append(res[0])

bc_hhi = res[0]  # Benchmark HHI

# ===============================================
# Visualization Style Settings
# ===============================================

# Set color theme and styles
custom_palette = ['#2F7FC1', "#96C37D", "#F3D266", "#D8383A", "#A9B8C6", '#14517C', '#C497B2', '#F8AC8C']
sns.set_theme(style="whitegrid", palette=custom_palette)
sns.color_palette("husl", 9) 
sns.set_context(rc={"font.size":20, "axes.titlesize":22, "axes.labelsize":20, "axes.label":22})  

# ===============================================
# Product Information Analysis
# ===============================================

# Build product information dataframe
prod_info_np = []
mlist = ['MNL', 'M=1', 'M=2', 'M=3', 'Practice']

# Collect product info from first 4 strategies
for m in range(4):
    for i in np.where(menu_list[m].sum(axis=1)>0)[0]:
        prod_info_np.append([mlist[m], 15*r[i], base_utility[i], np.exp(base_utility)[i]]) 

# Collect product info from benchmark strategy
for i in np.where(bmk_menu.sum(axis=1)>0)[0]:
    prod_info_np.append(['Practice', 15*r[i], base_utility[i], np.exp(base_utility)[i]]) 

prod_info = pd.DataFrame(prod_info_np, columns=['Menu type', 'Price', 'Base utility', 'Base attraction value']) 

# ===============================================
# Offering Frequency Analysis
# ===============================================

# Analyze product offering frequency under different strategies
offer_times = [[] for m in range(5)]

# Calculate offering times for optimization strategies
for m in range(4):
    for k in range(1,6,1):
        offer_times[m].append(len(np.where(menu_list[m].sum(axis=1)<=k)[0]) - len(np.where(menu_list[m].sum(axis=1)<=0)[0]))

# Calculate offering times for benchmark strategy
for k in range(1,6,1):
    offer_times[4].append(len(np.where(bmk_menu.sum(axis=1)<=k)[0]) - len(np.where(bmk_menu.sum(axis=1)<=0)[0]))

print("Offering times:", offer_times)

# ========================================= 
# Visualization: Figure 3 
# =======================================


fig = plt.figure(figsize=(21,6)) 
gs = gridspec.GridSpec(1, 3)
 
 
ax3 = plt.subplot(gs[0])
ax1 = plt.subplot(gs[1])  
ax2 = plt.subplot(gs[2])  
 
markerlist = ['.','s','>','<','*','D','X', 'P','+'] 

plot1 = sns.boxplot(x=prod_info['Menu type'],
            y = prod_info['Price'], ax=ax1, width=0.3,gap= 0.1  )  
ax1.tick_params(axis='x', labelsize=20)
ax1.tick_params(axis='y', labelsize=20)
ax1.set_yticks(np.arange(4,14.1,2))
ax1.set_xlabel("")
ax1.set_ylabel("Price (RMB)")
ax1.set_title('(b) Price')


plot2 = sns.boxplot(x=prod_info['Menu type'],
            y = prod_info['Base attraction value'], ax=ax2, width=0.3,gap= 0.1  )  
ax2.tick_params(axis='x', labelsize=20)
ax2.tick_params(axis='y', labelsize=20)
ax2.set_yticks(np.arange(0,1.05,0.2))
ax2.set_title('(c) Base attraction value')
ax2.set_xlabel("")


line_list = ['solid','dashed','dotted','dashdot','dotted' ]
sslist = [200, 110,110,110,110,110,110]
markerlist = ['<','','','','s']
color_order = [2,1,0,3,4]
for m in [4,3,2,1,0]:
    ax3.plot(np.arange(1,6,1), offer_times[m] , label = mlist[m], linestyle=line_list[color_order[m]],marker = markerlist[m], linewidth= 3 )   

 
ax3.set_xlabel('Offering times in one week')
ax3.set_ylabel('Accumulated # of products')
ax3.set_title('(a) Offering times' )
ax3.tick_params(axis='x', labelsize=20)
ax3.tick_params(axis='y', labelsize=20)
ax3.set_ylim(-0.7,24)
ax3.set_xticks(np.arange(1,6,1))
ax3.set_yticks(np.arange(0,24.1,4))
ax3.set_xlim(0.8,8.3 )
ax3.text(5.07, offer_times[0][4]-0.3, str(offer_times[0][4]), fontsize=20, color='k')
ax3.text(5.03, offer_times[1][4]-0.3, str(offer_times[1][4]), fontsize=20, color='k')
ax3.text(5.03, offer_times[2][4]-0.3,  str(offer_times[2][4]), fontsize=20, color='k')
ax3.text(5.04, offer_times[3][4]-0.3,  str(offer_times[3][4]), fontsize=20, color='k')
ax3.text(5.05, offer_times[4][4]-0.3, str(offer_times[4][4]), fontsize=20, color='k')
ax3.legend(fontsize='20')
plt.tight_layout() 
 
fig.savefig(os.path.join(RESULT_DIR,'Figure3.png'),dpi = 500)
print('Successfully save figure 3.')


# ========================================= 
# Visualization: Figure 5 
# =======================================

plt.figure(figsize=(12,4.8))  

markerlist = ['.','s','>','<','H','X','P','D']
custom_palette = ['#2F7FC1', "#96C37D" , "#F3D266", "#D8383A","#14517C",'#A9B8C6','#C497B2','#F8AC8C']
menu_names = ['MNL','M=1','M=2','M=3','M=1 cyclic','M=2 cyclic','M=3 cyclic']
sslist = [150,200,200,200,150,200,200,200]
x = []
y = []
markers = [ ]
sizelist = []
colorlist = []
 
for k in [0,1,2,3]:
    x.append(opt_hhi[k] )
    y.append( round(rev_list[k]/rev_bmk,3) )
    markers.append(markerlist[k+1])
    sizelist.append(sslist[k])
    colorlist.append( custom_palette[k])
    plt.scatter( opt_hhi[k] , round(rev_list[k]/rev_bmk,3) , label = menu_names[k],marker= markerlist[k+1] ,s=sslist[k], color = custom_palette[k])


x.append(bc_hhi )
y.append( round(rev_bmk/ rev_bmk,3))
markers.append(markerlist[0])  
sizelist.append(220)
colorlist.append( 'k')



plt.scatter( bc_hhi, round(rev_bmk/ rev_bmk,3), label = 'Practice',marker= markerlist[0] ,s=220, color = 'k')

for k in [4,5,6]:
    plt.scatter( opt_hhi[k] ,round(rev_list[k]/rev_bmk,3) , label = menu_names[k],marker= markerlist[k+1] ,s=sslist[k], color = custom_palette[k])
    x.append(opt_hhi[k] )
    y.append( round(rev_list[k]/rev_bmk,3) )
    markers.append(markerlist[k+1])
    sizelist.append(sslist[k])
    colorlist.append( custom_palette[k])

adjust_x = [ 0.001,   0.001, 0, -0.005,       0.001,       0,  0.0013,   -0.006]
adjust_y =[0.005,    0.0025,  0.005, 0.006,   0.004,      -0.025,-0.025,  -0.035]
for i, label in enumerate(menu_name_all):
    plt.text(x[i]+adjust_x[i] , y[i]+adjust_y[i], label, fontsize=18,ha='left', va='bottom' )  



plt.xlabel('HHI' )
plt.ylabel('Normalized revenue\n(out-of-sample)',multialignment='center' ) 
 
plt.tick_params(axis='x', labelsize=20)
plt.tick_params(axis='y', labelsize=20)
plt.ylim(0.88,1.22)
plt.xlim(0.035,0.18) 

 
plt.tight_layout()
plt.savefig(os.path.join(RESULT_DIR,'Figure5.png'),dpi = 500) 
print('Successfully save figure 5.')

# ========================================= 
# Visualization: Figure 4
# =======================================

fig = plt.figure(figsize=(20,7)) 
gs = gridspec.GridSpec(2,2) 

ax4 = plt.subplot(gs[0]) 
ax5 = plt.subplot(gs[1])   
ax2 = plt.subplot(gs[2])   
ax3 = plt.subplot(gs[3])    
 
count_prod = 30 
start_n = 0
end_n = count_prod 
axlist = [ax4,ax5,ax2,ax3]
abc = ['(a) ', '(b) ',"(c) ","(d) "] 
act_menu = [0,1,2,3]
for kk in range(4): 
    axs = axlist[kk]
    axs.set_xlim(start_n,end_n )
    axs.set_ylim(5,0) 
    axs.set_yticks([])

    axs.xaxis.set_major_locator(MultipleLocator(1))
    axs.yaxis.set_major_locator(MultipleLocator(1))
    axs.xaxis.grid(True, which='major')
    axs.yaxis.grid(True, which='major')
    axs.set_ylabel('Weekday')
    axs.set_xlabel('Product index')
 
    for t in range(5):
        for i in range(N):
            if menu_list[act_menu[kk]][i,t] >0:
                axs.fill_between([N-1-i,N-i],t,t+1,facecolor= custom_palette[0]) 

    axs.tick_params(axis='x', labelsize=14)
    axs.tick_params(axis='y', labelsize=16)
    axs.set_title(abc[kk] + menu_name_all[kk]+" policy")
 
plt.tight_layout()
fig.savefig( os.path.join(RESULT_DIR,'Figure4_abcd.png'),dpi = 500)
print('Successfully save figure 4-abcd.')


fig = plt.figure(figsize=(19,4)) 
gs = gridspec.GridSpec(1, 2) 

ax0 = plt.subplot(gs[1])  
ax1 = plt.subplot(gs[0])  
 
sns.barplot( x=np.arange(1,31),y = order_a[0:30],ax= ax0, width=0.5 ) 
ax0.set_ylabel('Base attraction value',fontsize='20')
ax0.set_xlabel('Product index',fontsize='20')
ax0.tick_params(axis='x', labelsize=14)
ax0.tick_params(axis='y', labelsize=16)  

sns.barplot( x=np.arange(1,31),y =order_r[0:30],ax= ax1, width=0.5 ) 
ax1.set_ylabel('Price (RMB)',fontsize='20')
ax1.set_xlabel('Product index',fontsize='20')
ax1.tick_params(axis='x', labelsize=14)
ax1.tick_params(axis='y', labelsize=16) 

 
plt.tight_layout()
fig.savefig(os.path.join(RESULT_DIR,'Figure4_ef.png'),dpi = 500)
print('Successfully save figure 4-ef.')


# ========================================= 
# Visualization: Figure 1 
# ======================================= 

intro_fig = []  

intro_fig.append( ["Our Model", rev_list[3]  ,  len(np.where(menu_list[3].sum(axis=1)>0)[0]) ] ) 
intro_fig.append( [menu_name_all[0], rev_list[0] , len(np.where(menu_list[0].sum(axis=1)>0)[0])] )
intro_fig.append( ["Practice", 1, len(np.where(bmk_menu.sum(axis=1)>0)[0])] )

intro_fig = pd.DataFrame(intro_fig, columns=['Menu type','Revenue (out-of-sample)',"Assortment size"])


fig = plt.figure(figsize=(20,4)) 
gs = gridspec.GridSpec(1, 2)
width_ratios=[1,1] 
 
ax3 = plt.subplot(gs[0])   

sns.barplot(data=intro_fig, x = 'Menu type',y='Revenue (out-of-sample)',ax=ax3,  width=0.4,gap= 0.1  )
ax3.tick_params(axis='x', labelsize=18)
ax3.tick_params(axis='y', labelsize=18) 
ax3.set_xlabel("") 
for i in ax3.containers:
    ax3.bar_label(i,)
ax3.set_ylabel("Normalized revenue\n(out-of-sample)", multialignment='center')

 
ax4 = plt.subplot(gs[1])   

sns.barplot(data=intro_fig, x = 'Menu type',y='Assortment size',ax=ax4,  width=0.4,gap= 0.1  )
ax4.tick_params(axis='x', labelsize=18)
ax4.tick_params(axis='y', labelsize=18) 
for i in ax4.containers:
    ax4.bar_label(i,)
ax4.set_ylabel("Assortment size", multialignment='center') 

plt.tight_layout() 

fig.savefig(os.path.join(RESULT_DIR,'Figure1.png'),dpi = 500)
print('Successfully save figure 1.')
