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

"""
Use estimators as optimization input and optimize menus for memory length = 0,1,2,3.

"""

# =======================================================
# For quick replication, reduce the running time. 
time_bound = 60*60*2
# =======================================================

function enable_stay_awake()
    if Sys.iswindows()
        # Windows: Call Kernel32 API
        # ES_CONTINUOUS (0x80000000) | ES_SYSTEM_REQUIRED (0x00000001)
        # Tells the system that this thread requires the system to be in the working state.
        try
            ccall((:SetThreadExecutionState, "kernel32"), UInt32, (UInt32,), 0x80000001)
            println(">> Windows: System sleep prevention enabled.")
        catch
            println(">> Windows: Failed to set sleep prevention via API.")
        end
    elseif Sys.isapple()
        # macOS: Use caffeinate
        # -i: prevent idle sleep
        # -w: wait for the process with the given PID to exit
        run(`caffeinate -i -w $(getpid())`, wait=false)
        println(">> macOS: System sleep prevention enabled (caffeinate).")
    elseif Sys.islinux()
        # Linux: Use systemd-inhibit
        # It creates a lock that prevents sleep/idle.
        try
            run(`systemd-inhibit --why="Numerical Experiments" --who="Julia" --what=idle:sleep sleep infinity`, wait=false)
            println(">> Linux: System sleep prevention enabled (systemd-inhibit).")
        catch
            println(">> Linux: systemd-inhibit not found. System may still sleep.")
        end
    end
end

enable_stay_awake()

using JuMP,Gurobi,BenchmarkTools, LinearAlgebra  
import MathOptInterface as MOI  
using MosekTools  
using Mosek  
using CSV, DataFrames
 

# --- PATH SETUP ---
const SCRIPT_DIR = @__DIR__
const PROJECT_ROOT = dirname(SCRIPT_DIR) 

# ------------------

include(joinpath(PROJECT_ROOT, "src", "LargeM_CuttingPlane.jl")) 
include(joinpath(PROJECT_ROOT, "src", "Project_CuttingPlane.jl")) 
include(joinpath(PROJECT_ROOT, "src", "Conic_Env.jl")) 

# change a matrix to dataframe form 
function tupe2mat(a)
    n = length(a[1])
    m = ones(length(a), n)
    for i in 1:length(a)
         m[i,:] = [t for t in a[i]]
    end
    return m
end

# generate historical assortments for gradient constraint 
function gene_his_set(M)
    xhis_set = []
    push!(xhis_set, zeros(M))
    for k in 1:M 
        h = zeros(M)
        h[k] = 1
        push!(xhis_set, h)
    end 
    for k in 1:M-1 
        for kk in k+1:M 
            h = zeros(M)
            h[k] = 1
            h[kk] = 1
            push!(xhis_set, h)
        end 
    end 
    push!(xhis_set,ones(M))
end 

# generate initial rankings for lovasz extension 
using Random
function gene_rank(M,num_rank)
    items = collect(1:M)
    rank_set = Set()
    while length(rank_set)< num_rank 
        new_rank = shuffle(items)
        if (new_rank in rank_set)
            nothing 
        else
            push!(rank_set, new_rank)
            push!(rank_set, reverse(new_rank))
        end  
    end  
    return rank_set 
end 

# solve MNL
function MNL(v0,v,r,N, cardinal)
    model = Model(Gurobi.Optimizer)
    set_optimizer_attribute(model, "OutputFlag", 0)
    low_rho = 1/(v0 + sum(v[i] for i in 1:N))  
    up_rho = 1/v0 
    @variable(model, low_rho<= rho <=up_rho)  
    @variable(model, 0 <= y[1:N] )  
    @constraint(model, rho + sum( v[i]*y[i] for i in 1:N) == 1) 
    @constraint(model,[i in 1:N], y[i]<= rho )
    @constraint(model, sum(y[i] for i in 1:N)<= cardinal*rho) 
    @objective(model, Max, sum(r[i]*v[i]*y[i] for i in 1:N )  )  
    b = @timed optimize!(model)  
    x = value.(y) ./ value.(rho)
    return  objective_value(model), x , b.time
end 


# This ensures code doesn't crash if folder is missing
data_folder = joinpath(PROJECT_ROOT, "data")
result_folder = joinpath(PROJECT_ROOT, "results","case_study")
log_folder = joinpath(result_folder, "logs")

mkpath(log_folder)  

# ============ import optimization data ============
rdf = CSV.read(joinpath(data_folder,"revenue.csv"), DataFrame)
udf = CSV.read(joinpath(data_folder,"utility.csv"), DataFrame)
betadf = CSV.read(joinpath(data_folder,"beta.csv"), DataFrame)
beta_init = betadf[:,1:4]
r = round.(rdf[:,1], digits = 2)
u_init = udf[:,:]


# =============== optimization configurations ==============

v0 = 1 
stop_gap = 0.005
init = 0
init_X = Any[]
output = 1
IP = 1 

T = 5 
N = length(r)
cardinal = 6
showtime = T 

for M in [1,2,3] # for each memory length 
    beta_value = beta_init[M,1:M]
    beta = zeros(N,M)
    for i in 1:N 
        for m in 1:M
            beta[i,m] = beta_value[m]
        end 
    end  
    u = u_init[:,M+1]

    s = zeros(N,M)
    for i in 1:N 
        for m in 1:M 
            if beta[i,m]>0 
                s[i,m]=1
            end 
        end
    end 

    betanew = zeros(N,M,M+1)
    for t in 2:M+1
        for i in 1:N  
            if (t<=M)&(t>=2)
                for m in 1:t-1 
                    betanew[i,m,t] = beta[i,m]
                end 
            else 
                for m in 1:M 
                    betanew[i,m,t] =  beta[i,m] 
                end 
            end 
        end 
    end  

    logname = joinpath(log_folder,"case_solve_m"*string(M)*".txt")

    if M == 1
        res1 = MIP_M1_Strong( T,N,M,r,beta,u,v0,IP, time_bound ,stop_gap,init, init_X, output, logname, s,  cardinal, showtime  )  
        solution = res1[2]
    elseif M ==2 
        res1 = MIP_M2_Satiation_Strong(T, N, M, r, beta, u, v0, time_bound, IP, cardinal, showtime, stop_gap, init, init_X, logname, s, output)   
        solution = res1[2] 
    else
        rank_set = [ [1,2,3], [1,3,2], [2,1,3], [2,3,1],[3,2,1],[3,1,2] ]
        res1 = MIP(T, N, M, r, beta, u, v0, IP, time_bound, stop_gap, init, init_X, 
            output, logname, s, betanew, rank_set, cardinal, showtime, 0) 
        solution = res1[4]
    end 

    f1 = joinpath(log_folder,"case_menu_m"*string(M)*".csv")
    temp_data = DataFrame(solution, :auto) 
    CSV.write(f1,  temp_data, writeheader=false)


    # solve M+1 cyclic  
    lognamec = joinpath(log_folder,"case_solve_cycle_"*string(M)*".txt")
    resbf1 = BoundFreewithSepa(N,M,r,u,v0,beta, IP ,1, stop_gap, time_bound, lognamec, s, cardinal )
    # generate solution for T periods 
    cycle_solu = zeros(N,T)
    for t in 1:T 
        ind = t%(M+1)
        if ind == 0
            cycle_solu[:, t] = resbf1[2][:,M+1]
        else
            cycle_solu[:, t] = resbf1[2][:, ind]
        end 
    end  

    f2 = joinpath(log_folder,"case_cycle_menu_"*string(M)*".csv")
    temp_data = DataFrame(cycle_solu, :auto) 
    CSV.write(f2,  temp_data, writeheader=false)

end 
# ----------- solving constrained MNL ---------------
u = u_init[:,1]
v = exp.(u)
mnlres = MNL(v0,v,r,N, cardinal) 
mnlres[2]

mnlsolu = zeros(N,T)
for t in 1:T 
    mnlsolu[:,t] = mnlres[2]
end  
temp_data_mnl = DataFrame(mnlsolu, :auto) 
CSV.write(joinpath(log_folder,"case_menu_mnl.csv"),  temp_data_mnl, writeheader=false)

print("Finish optimizing menus")

exit(0)