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

"""
Multinomial Logit (MNL) Choice Model with Memory Effects and Cross-Validation 

This script implements a complete MNL choice model estimation with:
- Choice parameter estimation
- Bootstrap significance testing
- Cross-validation for model selection 
- Results export for further analysis

The model considers both fixed product features and historical memory effects
to predict customer choices in assortment optimization problems.

Note: To quickly reproduce the case study, reduce `n_bootstrap` to a smaller value.
"""

# =======================================================
# For quick replication, reduce the number of bootstrap times. 
n_bootstrap = 100
# =======================================================

import platform
import os
import subprocess

def enable_stay_awake():
    system = platform.system()
    
    if system == "Windows":
        # Windows: Use ctypes to call SetThreadExecutionState
        # 0x80000001: ES_CONTINUOUS | ES_SYSTEM_REQUIRED
        import ctypes
        try:
            ctypes.windll.kernel32.SetThreadExecutionState(0x80000001)
            print(">> Windows: System sleep prevention enabled.")
        except Exception as e:
            print(f">> Windows: Failed to enable sleep prevention: {e}")

    elif system == "Darwin":  # macOS
        # macOS: Start caffeinate linked to the current Process ID (PID)
        try:
            subprocess.Popen(["caffeinate", "-i", "-w", str(os.getpid())])
            print(">> macOS: System sleep prevention enabled (caffeinate).")
        except Exception as e:
            print(f">> macOS: Failed to start caffeinate: {e}")

    elif system == "Linux":
        # Linux: Try systemd-inhibit
        try:
            subprocess.Popen(["systemd-inhibit", "--why='Numerical Experiments'", "sleep", "infinity"])
            print(">> Linux: System sleep prevention enabled (systemd-inhibit).")
        except FileNotFoundError:
            print(">> Linux: systemd-inhibit not found. System may still sleep.")

enable_stay_awake()

# ============================================================================
# IMPORT LIBRARIES AND DEPENDENCIES
# ============================================================================
import numpy as np 
import pandas as pd 
from scipy.optimize import minimize
from scipy import stats 
from scipy.stats import norm  
from datetime import datetime, time
from itertools import combinations
# from scipy.stats import chi2
import random 
import time
import pickle 

 

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

 

# ============================================================================
# CONFIGURATION AND GLOBAL PARAMETERS
# ============================================================================

# Define product feature columns
fea_col = ['price', 'fish', 'chicken', 'duck', 'beef', 'bean_root', 'eggplant',
           'pepper', 'mushroom', 'leaf', 'marinate', 'flavor0', 'flavor2','type1','type2']

# Model structure parameters
fixfea = len(fea_col)           # Number of fixed product features  
Mub = 4                         # Maximum memory length 
L = fixfea + Mub                # Total number of features 
print("Number of features =", L)
 
# Define column names for estimation results
res_cols = ['M', 'K']                      # Memory length and CV fold
res_cols.extend(fea_col)                   # Add feature column names
res_cols.extend(['beta1', 'beta2', 'beta3', 'beta4'])  # Add memory coefficient names

# Create feature name list for final results 
feaname = fea_col.copy()
feaname.extend(['beta1', 'beta2', 'beta3', 'beta4'])
# ============================================================================
# DATA LOADING AND PREPARATION
# ============================================================================

print("\n" + "="*60)
print("STAGE 1: DATA LOADING AND PREPARATION")
print("="*60)

# Load simulation data files
product_features = np.load( os.path.join(DATA_DIR,"product_features.npy")  )       # Product feature matrix (N x features)
ast_matrix = np.load(  os.path.join(DATA_DIR, "ast_matrix.npy")  )                 # Assortment availability matrix (T x N)
historical_offerings = np.load(  os.path.join(DATA_DIR,"historical_offerings.npy")) # Historical offering data (T x N x memory)
with open( os.path.join(DATA_DIR,'trans_data.pkl'), 'rb') as file:                 # Transaction data 
    trans_data = pickle.load(file)

# Extract dimensions
(T, N) = ast_matrix.shape
print("Number of products =", N)
print("Number of days =", T)
print("\n" + "="*60)
print("STAGE 2: OBTAIN CUSTOMER CHOICE MATRIX")
print("="*60)

# Simulation parameters 
buyvec = {}                     # Product features with historical information for each customer
astvec = {}                     # Assortment vector on different days
ast = {}                        # Set of available products per day
choices = {}                    # Customer choice records
histvec = {}                    # Historical vectors for each day (without customer-specific variations) 
nt = np.zeros(T)                # Number of customers for each day 
buy_cus=[[] for t in range(T)]  # Index of customers having a purchase
nobuy_cus=[[] for t in range(T)]# Index of customers buying nothing 
for t in range(T):  
    # Initialize product feature matrix for day t
    buyvec_init = np.zeros((N, L))
    
    # Fill product features: fixed features + historical offerings
    for i in range(N):
        buyvec_init[i, :fixfea] = product_features[i]          # Fixed product features
        buyvec_init[i, fixfea:] = historical_offerings[t, i]   # Historical memory features
        
    # complete historical vector
    histvec[t] = buyvec_init 
    astvec[t] = ast_matrix[t].copy()
    ast[t] = np.where(astvec[t] > 0)[0]   

    # individual customer choices
    for k in range(len(trans_data[t])) :  
        buyvec_entry = buyvec_init.copy() 
        for i in range(N):
            for m in range(Mub):
                buyvec_entry[i, fixfea+m ] = historical_offerings[t, i][m]*trans_data[t][k][1+m]    
        buyvec[(t, k)] = buyvec_entry 
        choices[(t, k)] = int(trans_data[t][k][0])    
        if trans_data[t][k][0]<N:
            buy_cus[t].append(k)
        else:
            nobuy_cus[t].append(k)   
    nt[t] = len(trans_data[t]) 
nt = np.array(nt)

 

print(f"✓ Customer choice matrix completed")
# ============================================================================
# IMPORTANT FUNCTIONS 
# ============================================================================
 

# ============================================================================
# CROSS-VALIDATION DATA SPLITTING
# ============================================================================

def split_days_kfold(T, k): 
    """
    Split T days into k folds for cross-validation
    
    Args:
        T: Total number of days
        k: Number of folds
    
    Returns:
        list: List of day groups for each fold
    """
    days = list(range(T))

    # Calculate base size and remainder for balanced splits
    base_size = T // k
    remainder = T % k

    groups = [[] for i in range(k)]
    start_idx = 0

    tset = set(range(T))
    # Create k groups with approximately equal sizes
    for i in range(k): 
        groups[i] = random.sample(list(tset), base_size)
        tset = tset - set(groups[i])
    for t in tset:
        rind = np.random.randint(0,k)
        groups[rind].append(t) 
    
    return groups

# ============================================================================
# MAXIMUM LIKELIHOOD ESTIMATION
# ============================================================================

def MLE_ESTIMATION_EM(N, num_para, train_item, choices, buyvec, astvec, ast, nt, input_beta, input_lamb, lamb_bound, opt_iter, regularizer, runem, em_round ): 
    """
    Maximum Likelihood Estimation for MNL model (use EM algorithm for censor data). EM algorithm is based on Section 9.4.1.2 in the book `The Theory and Practice of Revenue Management'. 

    Args:
        N (int): Number of products.
        num_para (int): Number of parameters to estimate.
        train_item (list): List of training day indices.
        choices (dict): Dictionary of customer choices.
        buyvec (dict): Dictionary of product feature vectors per customer.
        astvec (dict): Dictionary of assortment vectors per day.
        ast (dict): Available product indices per day.
        nt (list): List of number of customers per day.
        input_beta (float): Initial parameter value of parameters for optimization.
        input_lamb (float): Initial parameter value of arrival rate for optimization.
        lamb_bound (tuple of float): lower and upper bound of arrival rate.
        opt_iter (int): Maximum number of optimization iterations.
        regularizer: Weight of regularization term. 
        runem: whether use EM algorithm.
        em_round: Round limitation for EM algorithm 
    Returns:
        np.array: Estimated parameters and arrival rate. 
    """ 

    def MLE_OBJ(parameters): 
        part1 = 0
        part2 = 0
        
        for day in train_item:  
            for c in range(int(nt[day])):    
                utility = buyvec[(day, c)][0:N, 0:num_para] @ parameters
                attract = np.exp(utility)
                
                if choices[(day, c)] < N :  
                    part1 += utility[ int(choices[(day, c)]) ] 
                    part2 += np.log(astvec[day] @ attract + 1)  
                else: 
                    part2 += np.log(astvec[day] @ attract + 1)  
        
        return -part1 + part2 + regularizer* parameters@parameters 

    def MLE_GRADIENT(parameters):  
        gradvec = np.zeros(num_para)  
        
        for d_idx in range(num_para): # Renamed 'd' to 'd_idx' to avoid conflict with 'day'
            part1_grad = 0 
            part2_grad = 0 
            
            for day in train_item:  
                for c in range(int(nt[day])):   
                    utility = buyvec[(day, c)][0:N, 0:num_para] @ parameters
                    attract = np.exp(utility)
                    denominator = astvec[day] @ attract + 1
                    
                    if choices[(day, c)] < N:
                        part1_grad += buyvec[(day, c)][int(choices[(day, c)]), d_idx]  
                        numerator = sum(attract[j] * buyvec[(day, c)][j, d_idx] for j in ast[day])
                        part2_grad += numerator / denominator
                    else:
                        numerator = sum(attract[j] * buyvec[(day, c)][j, d_idx] for j in ast[day])
                        part2_grad += numerator / denominator
            
            gradvec[d_idx] = -part1_grad + part2_grad + 2*parameters[d_idx]*regularizer
        
        return gradvec
 
        
    def MLE_OBJ_EM(parameters): 
        part1 = 0
        part2 = 0
        
        for day in train_item:  
            for c in range(int(nt[day])):   
                utility = buyvec[(day, c)][0:N, 0:num_para] @ parameters
                attract = np.exp(utility)
                
                if choices[(day, c)] < N :
                    part1 += utility[ int(choices[(day, c)]) ] + np.log(lamb_star)
                    part2 += np.log(astvec[day] @ attract + 1)  
                else: 
                    part1 += at[day][c]*np.log(lamb_star) + (1-at[day][c])*np.log(1-lamb_star)
                    part2 += at[day][c]*np.log(astvec[day] @ attract + 1)  
        
        return -part1 + part2 + regularizer* parameters@parameters 

    def MLE_GRADIENT_EM(parameters):  
        gradvec = np.zeros(num_para)  
        
        for d_idx in range(num_para): # Renamed 'd' to 'd_idx' to avoid conflict with 'day'
            part1_grad = 0 
            part2_grad = 0 
            
            for day in train_item:  
                for c in range(int(nt[day])):  
                    utility = buyvec[(day, c)][0:N, 0:num_para] @ parameters
                    attract = np.exp(utility)
                    denominator = astvec[day] @ attract + 1
                    
                    if choices[(day, c)] < N :
                        part1_grad += buyvec[(day, c)][int(choices[(day, c)]), d_idx]  
                        numerator = sum(attract[j] * buyvec[(day, c)][j, d_idx] for j in ast[day])
                        part2_grad += numerator / denominator
                    else:
                        numerator = at[day][c] * sum(attract[j] * buyvec[(day, c)][j, d_idx] for j in ast[day])
                        part2_grad += numerator / denominator
            
            gradvec[d_idx] = -part1_grad + part2_grad + 2*parameters[d_idx]*regularizer
        
        return gradvec
    if runem ==0:
        lamb_star = 1 
        init_parameters = np.ones(num_para) * input_beta   
        res = minimize(MLE_OBJ, init_parameters, method="BFGS", jac=MLE_GRADIENT,
                        options={"gtol": 1e-6, "maxiter": opt_iter}) 
        para_star = res.x
    else:
        at = [ np.ones(int(nt[day])) for day in range(T)] 
        em_stage = 0 
        flag = 1 
        while flag:
            # initialize 
            if em_stage == 0:
                init_parameters = np.ones(num_para) * input_beta   
                init_lambda = input_lamb 
            else:
                init_parameters = para_star.copy()   
                init_lambda = lamb_star  
            for day in train_item: 
                for k in range(int(nt[day])): # for each no-purchase record 
                    if choices[(day,k)] >=N:
                        utility = buyvec[(day,k)][0:N, 0:num_para]@init_parameters
                        attrac = np.exp(utility)
                        suma = attrac@astvec[day] + 1 
                        no_buy = 1/suma 
                        at[day][k] = init_lambda*no_buy /(init_lambda*no_buy + 1 - init_lambda)

            # estimate parameters 
            # -- estimate lambda 
            total_at = 0.0
            total_n  = 0.0
            for day in train_item:
                total_at += np.sum(at[day])
                total_n  += nt[day]
            lamb_star = np.clip(total_at / total_n, lamb_bound[0], lamb_bound[1])


            # -- estimate beta  
            res = minimize(MLE_OBJ_EM, init_parameters, method="BFGS", jac=MLE_GRADIENT_EM,
                        options={"gtol": 1e-6, "maxiter": opt_iter})  
            para_star = res.x 

            # convergence test
            if (np.abs(lamb_star - init_lambda)<= 1e-3) or ((para_star - init_parameters)@(para_star - init_parameters) <= 1e-5):
                flag = 0 
            else:
                em_stage +=1 
                if em_stage > em_round:
                    flag = 0
    
    return para_star, lamb_star 
 
# ============================================================================
# MODEL EVALUATION FUNCTIONS
# ============================================================================

def COMPUTE_LOSS(N, num_para, test_item, choices, astvec, ast, nt, beta_star, lamb_star, histvec): 
    """
    Compute prediction loss on test data using chi-squared metric
    
    Args:
        N: Number of products
        num_para: Number of parameters
        test_item: List of test day indices
        choices: Customer choice data
        astvec: Assortment availability vectors
        ast: Available product sets
        nt: Number of customers per day
        beta_star: Estimated parameters
        buyvec: Product features 
    
    Returns:
        tuple: (mean_loss, total_loss)
    """
    # Calculate actual choice probabilities from observed data
    act_prob = []
    for day in test_item:
        b_count = np.zeros(N + 1)   # Count choices for each product + no-purchase
        for c in range(int(nt[day])): 
            bought = choices[(day, c)]
            b_count[bought] += 1 
        b_count[N] = b_count[N] - nt[day]*(1-lamb_star)  
        b_ratio = b_count / sum(b_count) # Convert counts to probabilities
        act_prob.append(b_ratio)

    # Calculate predicted choice probabilities using estimated model
    pred_prob = []
    for day in test_item: 
        utility = histvec[day][0:N, 0:num_para] @ beta_star
        attract = np.exp(utility)
        suma = 1 + astvec[day] @ attract   # Denominator for MNL probabilities
        
        ast_pred = np.zeros(N)
        for i in ast[day]:
            ast_pred[i] = attract[i] / suma   # MNL probability for each available product
        pred_prob.append(ast_pred)

    # Calculate chi-squared loss between actual and predicted probabilities
    num_cases = len(test_item) 
    loss = np.zeros(num_cases) 

    for t in range(num_cases): 
        loss_value = sum((act_prob[t][i] - pred_prob[t][i])**2 / pred_prob[t][i] 
                        for i in ast[test_item[t]])
        loss[t] = loss_value  
    
    return loss.mean(), loss.sum()

def GENE_LIKELIHOOD(N, num_para, test_item, choices, astvec, nt, beta_star):
    """
    Generate likelihood on test data
    
    Args:
        N: Number of products
        num_para: Number of parameters
        test_item: Test day indices
        choices: Customer choices
        astvec: Assortment vectors
        nt: Number of customers per day
        beta_star: Estimated parameters
    
    Returns:
        float: Negative log-likelihood on test data
    """
    part1 = 0 
    part2 = 0 
    
    for day in test_item:  
        for c in range(int(nt[day])):   
            utility = buyvec[(day, c)][0:N, 0:num_para] @ beta_star
            attract = np.exp(utility)
            
            # Add utility of chosen product
            if choices[(day, c)] < N:
                part1 += utility[choices[(day, c)]]
            
            # Add normalization term
            part2 += np.log(astvec[day] @ attract + 1)  
    
    return -part1 + part2
 
# ============================================================================
# FULL SAMPLE ESTIMATION
# ============================================================================

print("\n" + "="*60)
print("STAGE 3: FULL SAMPLE ESTIMATION")
print("="*60)

estimation = []
opt_iter = 100  
init_b = -1      # Initial parameter value
regularizer = 0 
input_lamb = 0.954 
lamb_bound = [0.95,0.99999] # no censor data 
runem = 1  
em_round = 100
 
print("Estimating parameters on full dataset for all memory lengths...")

# Estimate on full sample for each memory length
t1 = time.time()
for M in [0,1,2,3,4]:
    num_para = fixfea + M 
    print(f"  M = {M}: Estimating {num_para} parameters on {T} days")
    
    # Run MLE estimation on all days 
    beta_star_all, lamb_star = MLE_ESTIMATION_EM(N, num_para, list(range(T)), choices, buyvec, astvec, ast, nt, init_b, input_lamb, lamb_bound, opt_iter, regularizer, runem, em_round ) 
    print(" ========lamb_star:" , lamb_star)
    # Extend to full parameter vector 
    beta_star_all = list(beta_star_all)
    beta_star_all.extend([0 for j in range(Mub - M)])
    beta_star_all = np.array(beta_star_all)
    estimation.append(beta_star_all)
    
    print(f" ✓ Completed: {beta_star_all }...")

t2 = time.time() 
print("✓ Full sample estimation done. Running time: " , t2-t1, " seconds")


# ============================================================================
# OPTIMIZATION DATA PREPARATION
# ============================================================================

print("\n" + "="*60)
print("STAGE 4: OUTPUT FOR OPTIMIZATION")
print("="*60)

# Prepare data for optimization algorithms
r = product_features[:, 0]      # Revenue vector (price column)

# Create utility matrix for different memory lengths
u = np.zeros(shape=(N, Mub + 1))
beta = np.zeros(shape=(Mub, Mub))

for M in [0, 1, 2, 3, 4]:
    est_alpha = estimation[M][0:fixfea]      # Fixed feature coefficients
    u[:, M] = product_features @ est_alpha   # Utility for memory length M
    
    if M > 0:
        for m in range(M):
            beta[M-1, m] = estimation[M][fixfea + m]  # Memory coefficients


# Export optimization data 
pd.DataFrame(r).to_csv( os.path.join(DATA_DIR,'revenue.csv'), index=False)
pd.DataFrame(u).to_csv( os.path.join(DATA_DIR, 'utility.csv'), index=False)
pd.DataFrame(beta).to_csv(  os.path.join(DATA_DIR,'beta.csv'), index=False)
print("✓ Optimization data files exported")


# ============================================================================
# BOOTSTRAP SIGNIFICANCE TESTING
# ============================================================================

print("\n" + "="*60) 
print("STAGE 5: COMPUTE SIGNIFICANCE")
print("="*60)

# Create list of available choice sets for bootstrap sampling
astlist = []
for t in range(T):    
    astlist.append(set(ast[t]))

RESULT = []                      # Storage for final results with significance tests
n_choice_sets = len(astlist)     # Number of choice sets (days) 

print(f"Running bootstrap significance tests with {n_bootstrap} samples for each memory length...")

# Run bootstrap tests for each memory length
for M in [0, 1, 2, 3, 4]:
    num_para = fixfea + M 
    beta_star = estimation[M][0:num_para]    # Full sample estimates
    successful_boots = 0 
    bootstrap_params = []
    
    print(f"\n--- Bootstrap Testing for M = {M} ---")
    
    # Generate bootstrap samples
    for i in range(n_bootstrap):  
        # Sample choice sets (days) with replacement
        boot_choice_set = np.random.choice(list(range(n_choice_sets)), 
                                         size=n_choice_sets, 
                                         replace=True)

        # Fit model on bootstrap sample 
        boot_params, lamb_value = MLE_ESTIMATION_EM(N, num_para, boot_choice_set, choices, buyvec, astvec, ast, nt, init_b, lamb_star, lamb_bound, opt_iter, regularizer, 0, 0  )
        
        bootstrap_params.append(boot_params)  
        successful_boots += 1  

    # Compute statistical significance tests
    bootstrap_params = np.array(bootstrap_params)

    # Calculate standard errors from bootstrap distribution
    se_bootstrap = np.std(bootstrap_params, axis=0)

    # Calculate t-statistics and p-values
    t_stats = beta_star / se_bootstrap
    p_values = list(2 * (1 - stats.t.cdf(np.abs(t_stats), df=successful_boots - 1)))

    # Store results (extend with zeros for unused parameters)
    beta_star_full = list(beta_star) + [0 for j in range(Mub - M)] + [M]
    p_values_full = list(p_values) + [0 for j in range(Mub - M)] + [M]
    se_bootstrap_full = list(se_bootstrap) + [0 for j in range(Mub - M)] + [M]
 
    
    RESULT.append(beta_star)    # Estimates 

    RESULT.append(p_values_full)     # P-values  
    RESULT.append(se_bootstrap_full) # Standard errors   

print(f"\n✓ Bootstrap significance testing completed")

# Add memory length identifier to feature names
feaname.append("M")

# Create final results DataFrame with significance tests
result_pv = pd.DataFrame(RESULT, columns=feaname)

# Create key column for result identification (estimator, pvalue, std repeated 5 times)
key_list = ["estimator", "pvalue", "std"] * 5
result_pv['name'] = key_list

# Export final results with significance tests 
result_simple = result_pv[['name','price', 'beta1', 'beta2', 'beta3', 'beta4']]
namelist = [ ]
for M in range(0, Mub + 1): 
    for qq in range(3):
        namelist.append("M"+str(M)) 
result_simple.insert(0, "Memory length", namelist )

result_simple.to_excel( os.path.join(RESULT_DIR, "Table2_estimation.xlsx"), index=False)
print("✓ Final estimation results with significance tests exported")


# ============================================================================
# CROSS-VALIDATION ANALYSIS
# ============================================================================

print("\n" + "="*60)
print("STAGE 6: CROSS-VALIDATION ANALYSIS")
print("="*60)

K = 5                                    # Number of cross-validation folds
day_groups = split_days_kfold(T, K)      # Split days into folds
est_res = []                             # Storage for estimation results
pred_res = []                            # Storage for prediction results
opt_iter = 100                           # Maximum optimization iterations

print(f"Running {K}-fold cross-validation for memory lengths M = 0,1,2,3,4")

# Loop over different memory lengths
for M in [0, 1, 2, 3, 4]:
    num_para = fixfea + M                # Total parameters = fixed features + memory features
    print(f"\n--- Testing Memory Length M = {M} ({num_para} parameters) ---")
    
    # Loop over cross-validation folds
    for k in range(K):
        print(f"  Fold {k+1}/{K}: ", end="")
        
        # Split data into training and testing sets
        test_item = day_groups[k]
        train_item = list(set(range(T)) - set(test_item))
        
        print(f"Train on {len(train_item)} days, test on {len(test_item)} days")
        
        # Estimate parameters on training data  
        beta_star, lamb_star  = MLE_ESTIMATION_EM(N, num_para, train_item, choices, buyvec, astvec, ast, nt, init_b, lamb_star, lamb_bound, opt_iter, regularizer, 0, 0 )
        print(f"    ✓ Estimation completed: {beta_star[:min(3, len(beta_star))]}...")

        # Evaluate performance on test data
        loss = COMPUTE_LOSS(N, num_para, test_item, choices, astvec, ast, nt, beta_star, lamb_star, histvec)
        likeli = GENE_LIKELIHOOD(N, num_para, test_item, choices, astvec, nt, beta_star)
        
        print(f"    ✓ Test performance - Loss: {loss[0]:.4f}, Likelihood: {likeli:.2f}")
        
        # Store estimation results (pad with zeros for unused memory coefficients)
        est = [M, k]
        est.extend(beta_star)
        est.extend([0 for j in range(Mub - M)])  # Pad unused coefficients with 0
        est_res.append(est)
        
        # Store prediction results
        pred = [M, k]
        pred.extend(loss)
        pred.append(likeli)
        pred_res.append(pred)

print(f"\n✓ Cross-validation completed: {len(est_res)} total estimations")

# ============================================================================
# RESULTS EXPORT AND IMPROVEMENT ANALYSIS
# ============================================================================

print("\n" + "="*60)
print("STAGE 7: RESULTS EXPORT AND IMPROVEMENT ANALYSIS")
print("="*60)


# Create result DataFrames
res_df = pd.DataFrame(est_res, columns=res_cols)
pred_df = pd.DataFrame(pred_res, columns=['M', 'K', 'mean_loss', 'total_loss', 'likeli'])
total_sample = sum(nt)
def calculate_mean_likeli(x, y):
    test_sample = sum(nt[int(day)] for day in day_groups[int(x)])
    train_sample = total_sample - test_sample 

    return y/train_sample 
pred_df['mean_likeli'] = pred_df.apply(lambda row: calculate_mean_likeli(row['K'], row['likeli']), axis=1)

# Export cross-validation results 
# pred_df.to_excel('results/case_estimation/cv_loss.xlsx', index=False)
print("✓ Cross-validation done")

# Calculate improvement metrics compared to baseline (M=0)
consider_index = [2, 5]  # Indices for mean_loss and likelihood
baseline = np.array(pred_df[pred_df['M'] == 0])[:, consider_index].sum(axis=0) / K
baseline = np.insert(baseline,1,0 )
baseline = np.insert(baseline,3,0 )

improve = [baseline]


# Calculate improvements for each memory length
for M in [1, 2, 3, 4]: 
    mres = np.array(pred_df[pred_df['M'] == M])[:, consider_index].sum(axis=0) / K
    lift1 =  (-100)*(mres[0] - baseline[0]) / baseline[0]
    lift2 =  (-100)*(mres[1] - baseline[2]) / baseline[2]
    specific_m_res = [mres[0], lift1, mres[1], lift2]
    improve.append(specific_m_res) 

# Export improvement analysis
liftdf = pd.DataFrame(improve, columns = ['chi_square' , 'csq_lift (100%)', 'mean_log_likelihood', "lkh_lift (100%)"]) 
namelist = []
for M in range(0, Mub + 1): 
    namelist.append("M"+str(M)) 
 
liftdf.insert(0, "Memory length", namelist)
liftdf['mean_log_likelihood'] = (-1)*liftdf["mean_log_likelihood"]
liftdf.to_excel( os.path.join(RESULT_DIR, "Table2_improvement.xlsx"), index=False)
print("✓ Model improvement metrics exported")

# ============================================================================
# ANALYSIS COMPLETION SUMMARY
# ============================================================================

print("\n" + "="*80)
print("ANALYSIS COMPLETED SUCCESSFULLY")
print("="*80) 

print(f"\nGenerated Output Files:")
files_generated = [ 
    "Table2_estimation.xlsx - Bootstrap significance test results", 
    "revenue.csv - Product revenue vector for optimization",
    "utility.csv - Estimated utility matrix for optimization", 
    "beta.csv - Memory effect parameters for optimization",
    "Table2_improvement.xlsx - Model improvement metrics", 
]

for i, file_desc in enumerate(files_generated, 1):
    print(f"  {i}. {file_desc}")

print(f"\n✓ All analysis stages completed successfully!")
print(f"Analysis finished at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
 