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

using JuMP, Gurobi, BenchmarkTools, LinearAlgebra  
import MathOptInterface as MOI   

# ==================== Optimization Model ====================================
# This module implements the multilinear extension based formulation for HAP. 
# ================================================================================


# ------------- Generate utility values and interaction effects for the multilinear extension ------
""" 
Returns:
- v: exponential of base utilities exp(u)
- d: interaction effects matrix (N × (2^M-1))
# Note: only for memory length M ≤ 2 
"""
function Gene_V(u, beta, M, N) 
    v = exp.(u)  # Base exponential utilities
    d = Matrix{Float64}(undef, N, 2^M-1)  # Initialize interaction effects matrix
    if M>2
        print("Invalid memory length: M is no more than 2.")
    else 
        if M == 1   
            d = exp.(u + beta) - v
        end 
        if M == 2  
            d[:, 1] = exp.(u + beta[:, 1]) - v                   
            d[:, 2] = exp.(u + beta[:, 2]) - v                   
            d[:, 3] = exp.(u + beta[:, 1] + beta[:, 2]) - v - d[:, 1] - d[:, 2]   
        end 
    end 
    return v, d 
end 

# ----------------- Solve multilinear extension based formulation -------------- 
function Multilinear(T, N, M, r, u, beta, v0, IP, stop_gap, time_bound, output, logname, s, cardinal, showtime)
    # Generate multilinear extension parameters 
    v, d = Gene_V(u, beta, M, N) 
    K = 2^M - 1  # Number of interaction terms
    
    # Compute bounds for rho 
    low_rho = 1 / (v0 + sum(exp(u[i] + dot(s[i, :], beta[i, :])) for i in 1:N)) 
    up_rho = 1 / v0 

    model = Model(Gurobi.Optimizer)
    # Decision Variables 
    @variable(model, 0 <= x[1:N, 1:T] <= 1)           # Assortment decisions 
    @variable(model, low_rho <= rho[1:T] <= up_rho)   # No purchase probabilities 
    @variable(model, g[1:N, 1:T])                     # Auxiliary: rho[t] * x[i,t]
    @variable(model, delta[1:N, 1:T, 1:K])            # Interaction terms between periods
    @variable(model, p[1:N, 1:T])                     # Auxiliary: x[i,t-1] * x[i,t-2]
    @variable(model, y[1:N, 1:T, 1:K])                # Auxiliary: rho[t] * delta[i,t,k]
 
    # Objective Function 
    @objective(model, Max, 
        sum(r[i] * v[i] * g[i, t] + sum(r[i] * d[i, k] * y[i, t, k] for k in 1:K) 
            for i in 1:N for t in 1:T)
    )
  
    # Constraints  
    @constraint(model, [t in 1:T], 
        v0 * rho[t] + sum(v[i] * g[i, t] + sum(d[i, k] * y[i, t, k] for k in 1:K) for i in 1:N) == 1
    )
     
    @constraint(model, [t in 1:T], sum(x[i, t] for i in 1:N) <= cardinal) 
    @constraint(model, [i in 1:N], sum(x[i, t] for t in 1:T) <= showtime)
    
    # McCormick constraints 
    for t in 1:T
        for i in 1:N 
            # Linearize y[i,t,k] = rho[t] * delta[i,t,k]
            for k in 1:K
                McCormick(model, y[i, t, k], rho[t], delta[i, t, k], low_rho, up_rho, 0, 1)
            end 
            
            # Linearize g[i,t] = rho[t] * x[i,t]
            McCormick(model, g[i, t], rho[t], x[i, t], low_rho, up_rho, 0, 1)

            # Linearize p[i,t] = x[i,t-1] * x[i,t-2] (for second-order interactions)
            if (t >= 3) && (M > 1)
                McCormick(model, p[i, t], x[i, t-1], x[i, t-2], 0, 1, 0, 1)
            else 
                @constraint(model, p[i, t] == 0)  # Set to zero if not applicable
            end 

            # Define interaction terms 
            if M > 1   
                McCormick(model, delta[i, t, 3], x[i, t], p[i, t], 0, 1, 0, 1)
                if t >= 3  
                    for k in 1:2  # Two periods of history 
                        McCormick(model, delta[i, t, k], x[i, t], x[i, t-k], 0, 1, 0, 1)
                    end
                elseif t == 2  # Only one period of history
                    McCormick(model, delta[i, t, 1], x[i, t], x[i, t-1], 0, 1, 0, 1) 
                    @constraint(model, delta[i, t, 2] == 0)  # No t-2 available
                else  # No history available
                    @constraint(model, delta[i, t, 1] == 0) 
                    @constraint(model, delta[i, t, 2] == 0) 
                end 
            else  # M == 1: Only single interaction term
                if t >= 2 
                    McCormick(model, delta[i, t, 1], x[i, t], x[i, t-1], 0, 1, 0, 1) 
                else
                    @constraint(model, delta[i, t, 1] == 0)  
                end 
            end 
        end
    end       

    # Solver Configuration
    if IP == 1  
        set_binary.(x)  
        set_optimizer_attribute(model, "MIPGap", stop_gap)
        set_optimizer_attribute(model, "OutputFlag", output)
        set_optimizer_attribute(model, "LogFile", logname) 
        if time_bound > 0
            set_time_limit_sec(model, time_bound)
        end 
    end  
     
    optimize!(model)
 
    if IP == 1  
        return objective_value(model), solve_time(model), MOI.get(model, MOI.RelativeGap()), value.(x)
    else 
        return objective_value(model), solve_time(model), 0, value.(x)
    end 
end


# ----------- McCormick Envelope Constraints for Bilinear Terms -----------
function McCormick(m, z, x, y, xl, xu, yl, yu) 
    @constraint(m, z >= yl*x + xl*y - xl*yl)    # Lower bound constraint 1
    @constraint(m, z <= yu*x + xl*y - xl*yu)    # Upper bound constraint 1
    @constraint(m, z <= yl*x + xu*y - xu*yl)    # Upper bound constraint 2
    @constraint(m, z >= yu*x + xu*y - xu*yu)    # Lower bound constraint 2
end 
 
# --------------- Naive formulation -----------------------------------
function NaiveM2(T,N,r,u,v0,beta,time_bound, IP,cardinal, stop_gap,output,logname, model)
    M = 2 
    v,d = Gene_V(u,beta,M,N) 
    K = 2^M-1

    low_rho = 1/(v0 + sum(v[i] for i in 1:N)) 
    up_rho = 1/v0 
    if IP >0
        @variable(model, x[1:N,1:T],Bin) 
        set_optimizer_attribute(model, "MIPGap", stop_gap)
    else
        @variable(model, 0<=x[1:N,1:T]<=1) 
    end 
    @variable(model, low_rho<=rho[1:T]<=up_rho) # normalize fractional in each period 
    @variable(model, g[1:N,1:T]) # rho*xit 
    @variable(model, delta[1:N,1:T,1:K]) # x*x[i,t-1] x[i,t-2] x[i,t-1]*x[i,t-2]
    @variable(model, p[1:N,1:T]) # x[i,t-1]*x[i,t-2]
    @variable(model, y[1:N,1:T,1:K]) # rho*delta   
 
   
    # normalize 
    for t in 1:T
        @constraint(model, v0*rho[t]+sum(v[i]*g[i,t] + sum(d[i,k]*y[i,t,k] for k in 1:K) for i in 1:N ) == 1 ) 
        @constraint(model, sum(x[i,t] for i in 1:N)<=cardinal)
    end 
    
    # add McCormick constraints 
    for t in 1:T
        for i in 1:N 
            McCormick(model, g[i,t], rho[t], x[i,t],low_rho,up_rho,0,1)
            if t>=3
                McCormick(model, p[i,t],x[i,t-1],x[i,t-2],0,1,0,1)
            else 
                @constraint(model, p[i,t] == 0 )
            end 

            for k in 1:K
                McCormick(model, y[i,t,k], rho[t],delta[i,t,k], low_rho,up_rho, 0,1)
            end 

            McCormick(model, delta[i,t,3],x[i,t], p[i,t],0,1,0,1) # 最后一个维度由p控制了不用分类讨论
            # 前面两个维度需要分类讨论
            if t>=3
                for k in 1:2 
                    McCormick(model,delta[i,t,k],x[i,t],x[i,t-k],0,1,0,1 )
                end 
            elseif t==2
                McCormick(model, delta[i,t,1],x[i,t],x[i,t-1],0,1,0,1)
                @constraint(model, delta[i,t,2] == 0)
            else
                @constraint(model, delta[i,t,1] == 0)
                @constraint(model, delta[i,t,2] == 0)
            end 
        end
    end         
    # set objective function 
    @objective(model, Max, sum(r[i]*v[i]*g[i,t] + sum(r[i]*d[i,k]*y[i,t,k] for k in 1:K) 
    for i in 1:N for t in 1:T)  )
   
    if time_bound >0
        set_time_limit_sec(model, time_bound)
    end 
    if logname != 0 
        set_optimizer_attribute(model, "LogFile", logname)
    end
    set_optimizer_attribute(model, "OutputFlag", output)
    # optimize the model 
    optimize!(model)
    if IP ==1
        gap = MOI.get(model, MOI.RelativeGap())
        return objective_value(model), solve_time(model), gap 
    else 
        return objective_value(model), solve_time(model), 0 
    end 
end 
 
function NaiveM1(T,N,r,u,beta,v0,time_bound,IP,cardinal,stop_gap,output,logname, model)
    M =1 
    v = exp.(u)
    d = exp.(u+beta) - exp.(u)
    low_rho = 1/(v0 + sum( v[i] for i in 1:N))
    up_rho = 1/v0  
    set_optimizer_attribute(model, "MIPGap", stop_gap)

    # define variables 
    if IP >0
        @variable(model, x[1:N,1:T],Bin)
    else
        @variable(model, 0<= x[1:N,1:T]<=1)
    end 
    @variable(model, low_rho<=rho[1:T]<=up_rho) # rho_t 
    @variable(model, 0<= y[1:N,1:T] ) # y_it = rho_t*x_it 
    @variable(model, 0<= w[1:N,1:T] ) #  w_it = rho_t * delta_i
    @variable(model, 0<= delta[1:N,1:T] ) # delta_it  =  x_it * x_it-1 

    # ===============set objective function======== 
 
    @objective(model, Max,  sum(r[i]*v[i]*y[i,t] + r[i]*d[i]*w[i,t]  for i in 1:N for t in 1:T )    )  

    # ========== constraint ===================
    for t in 1:T
        @constraint(model, v0*rho[t] +sum(v[i]*y[i,t] + d[i]*w[i,t] for i in 1:N) == 1) # linearize fractional 
        @constraint(model, sum(x[i,t] for i in 1:N) <= cardinal) # cardinality 
        
        for i in 1:N 
            if t>=2
                McCormick(model, delta[i,t], x[i,t],x[i,t-1],0,1,0,1)
            else 
                @constraint(model, delta[i,t] == 0)
            end 
            McCormick(model, y[i,t], rho[t],x[i,t],low_rho,up_rho,0,1)
            McCormick(model, w[i,t], rho[t],delta[i,t],low_rho,up_rho,0,1)
        end         
    end 
 
    if time_bound >0
        set_time_limit_sec(model, time_bound)
    end 
    if logname != 0 
        set_optimizer_attribute(model, "LogFile", logname)
    end
    set_optimizer_attribute(model, "OutputFlag", output)
    optimize!(model)

    # compute gap 
    if IP ==1
        gap = MOI.get(model, MOI.RelativeGap())
        return objective_value(model),solve_time(model), gap ,value.(rho)
    else 
        return objective_value(model), solve_time(model), 0 
    end
end 