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

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

# ==================================================================================
# Dynamic Assortment Optimization Model: MIECP and ENV  
# ================================================================================== 
# Model Parameters:
# - N: number of products
# - T: number of time periods  
# - M: memory length (how many past periods affect current choice)
# - r: revenue vector (N×1)
# - beta: memory effect parameters (N×M)
# - u: base utility vector (N×1)
# - v0: no-purchase utility
#
# Decision Variables:
# - x[i,t]: binary, whether product i is offered in period t
# - rho[t]: market potential in period t
# - y[i,t]: purchase probability for product i in period t
# - g[i,t]: auxiliary variable = rho[t] * x[i,t]
# - z[i,t,m]: auxiliary variable = rho[t] * x[i,t] * x[i,t-m]
#
# Constraints:
# - cardinal: maximum number of products offered per period (default = N)
# - showtime: maximum periods each product can be offered (default = T)
#
# Optimization Parameters:
# - IP: 1 for integer programming, 0 for LP relaxation
# - time_bound: computation time limit in seconds
# - logname: log file name for solver output
# - cone: 1 for Conic (MIECP) formulation, 0 for Envelope (ENV) formulation
# - init: 1 if initial solution is provided
# - init_X: initial solution matrix (N×T)
# - s: indicator matrix (N×M) where s[i,m] = 1 if beta[i,m] > 0 
# ==================================================================================

# ----------- McCormick Envelope Constraints for Bilinear Terms -----------
"""  
- m: JuMP model
- z: bilinear variable (z = x*y)
- x, y: variables in the bilinear term
- xl, xu: lower and upper bounds for x
- yl, yu: lower and upper bounds for y
"""
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

# ----------- Convex Envelope Constraint for Purchase Probability -----------
""" 
Adds convex envelope constraints for purchase probability y[i,t].
Only applicable for memory length M ∈ {1,2} (see Remark 1 in paper).

- act_history: effective memory length plus one (min(M+1, t))
- i: product index
- t: time period index
"""
function ConvEnvelope(model, act_history, i, t, u, beta, y, g, z)  
    if act_history == 1  # No historical assortments (first period)
        @constraint(model, y[i,t] == g[i,t] * exp(u[i]))
    elseif act_history == 2  # One period of historical assortment
        @constraint(model, y[i,t] == g[i,t] * exp(u[i]) + 
                          z[i,t,1] * (exp(u[i] + beta[i,1]) - exp(u[i])))
    else # Two periods of historical assortment (M=2 case)
        if beta[i,1] * beta[i,2] > 0  # Same direction memory effects
            @constraint(model, y[i,t] * exp(-u[i]) >= g[i,t] + 
                              (exp(beta[i,1]) - 1) * z[i,t,1] + 
                              (exp(beta[i,2]) - 1) * z[i,t,2] )
            @constraint(model, y[i,t] * exp(-u[i]) >= 
                              (exp(beta[i,1] + beta[i,2]) - exp(beta[i,2])) * z[i,t,1] + 
                              (exp(beta[i,2] + beta[i,1]) - exp(beta[i,1])) * z[i,t,2] + 
                              g[i,t] * (exp(beta[i,1]) + exp(beta[i,2]) - exp(beta[i,1] + beta[i,2])))
        else  # Different direction memory effects
            @constraint(model, y[i,t] * exp(-u[i]) >= g[i,t] +
                              z[i,t,2] * (exp(beta[i,2]) - 1) +
                              z[i,t,1] * (exp(beta[i,1] + beta[i,2]) - exp(beta[i,2])) )
            @constraint(model, y[i,t] * exp(-u[i]) >= + g[i,t] +
                              z[i,t,1] * (exp(beta[i,1]) - 1) + 
                              z[i,t,2] * (exp(beta[i,1] + beta[i,2]) - exp(beta[i,1])) )
        end 
    end 
end  

# ----------- Attraction Value Calculation -----------
""" 
Calculates the attraction value for product i in period t.

- k: index of valid history memory (0,1,...,M)
- betanew: extended beta parameters (N × M × (M+1)) 
- h: history assortment matrix (M+1 × M) 
"""
function attraction(i, t, k, M, u, betanew, h)  
    return exp(u[i] + dot(betanew[i, :, min(M+1, t)], h[k+1, :]))
end 

# ----------- Concave Envelope Constraint (Lovász Extension) -----------
""" 
Adds concave envelope constraints using Lovász extension for product i in period t. 
- ranking: permutation of memory indices {1,...,M}   
"""
function lovasz_constr(model, i, t, M, ranking, u, betanew, y, g, z, s) 
    # Initialize history assortment matrix h
    h = zeros(M+1, M)
    h[1, :] = s[i, :]  # Starting point (initial configuration)
    
    # Build vertices of the hypercube for Lovász extension
    for m in 2:M+1 
        h[m, :] = h[m-1, :]  # Copy previous configuration
        # Switch the coordinate according to ranking
        if s[i, ranking[m-1]] > 0 
            h[m, ranking[m-1]] = h[m, ranking[m-1]] - 1   
        else 
            h[m, ranking[m-1]] = h[m, ranking[m-1]] + 1   
        end 
    end 
    
    # Add Lovász extension constraint (piecewise linear upper bound)
    @constraint(model, y[i,t] <=  
        attraction(i, t, 0, M, u, betanew, h) * 
        (g[i,t] - (1 - s[i, ranking[1]]) * z[i,t, ranking[1]] - 
         s[i, ranking[1]] * (g[i,t] - z[i,t, ranking[1]])) +  

        sum( attraction(i, t, k, M, u, betanew, h) * 
            ((1 - s[i, ranking[k]]) * z[i,t, ranking[k]] + 
             s[i, ranking[k]] * (g[i,t] - z[i,t, ranking[k]]) - 
             (1 - s[i, ranking[k+1]]) * z[i,t, ranking[k+1]] - 
             s[i, ranking[k+1]] * (g[i,t] - z[i,t, ranking[k+1]]))  for k in 1:M-1) + 
          
        attraction(i, t, M, M, u, betanew, h) * 
        ((1 - s[i, ranking[M]]) * z[i,t, ranking[M]] + 
         s[i, ranking[M]] * (g[i,t] - z[i,t, ranking[M]])))
end 

# ----------- Main Optimization Function -----------
""" 
- cone: 1 for MIECP formulation, 0 for ENV formulation
- return a tuple (objective_value, solve_time, optimality_gap, optimal_solution)
"""
function MIP(T, N, M, r, beta, u, v0, IP, time_bound, stop_gap, init, init_X, 
            output, logname, s, betanew, rank_set, cardinal, showtime, cone) 
    
    # ----------- Model Setup -----------
    if cone == 1
        model = Model(Mosek.Optimizer)
    else 
        model = Model(Gurobi.Optimizer)
        set_optimizer_attribute(model, "OutputFlag", output)
    end 
    
    # ----------- Bounds Calculation ----------- 
    low_rho = 1 / (v0 + sum(exp(u[i] + dot(s[i, :], beta[i, :])) for i in 1:N))
    up_rho = 1 / v0

    # ----------- Decision Variables -----------
    @variable(model, 0 <= x[1:N, 1:T] <= 1)           # Assortment decisions
    @variable(model, low_rho <= rho[1:T] <= up_rho)   # No-purchase probability 
    @variable(model, 0 <= y[1:N, 1:T])                # Purchase probabilities  
    @variable(model, 0 <= g[1:N, 1:T] <= up_rho)     # Auxiliary: rho * x
    @variable(model, 0 <= z[1:N, 1:T, 1:M] <= up_rho) # Auxiliary: rho * x * x_lag

    # ----------- Initial Solution -----------
    if init > 0
        for t in 1:T, i in 1:N
            set_start_value(x[i, t], init_X[i, t])
        end 
    end 

    # ----------- Objective Function ----------- 
    @objective(model, Max, sum(r[i] * y[i, t] for i in 1:N for t in 1:T))  

    # ----------- Constraints ----------- 
    @constraint(model, [t in 1:T], v0 * rho[t] + sum(y[i, t] 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 Bilinear Terms -----------
    for t in 1:T, i in 1:N
        # McCormick constraint: g[i,t] = rho[t] * x[i,t]
        McCormick(model, g[i, t], rho[t], x[i, t], low_rho, up_rho, 0, 1)  
        
        # McCormick constraints: z[i,t,m] = g[i,t] * x[i,t-m]
        for m in 1:M
            if t - m >= 1  # Historical period exists
                McCormick(model, z[i, t, m], g[i, t], x[i, t-m], 0, up_rho, 0, 1)
            else  # No historical period, set to zero
                @constraint(model, z[i, t, m] == 0)
            end 
        end 
    end 

    # ----------- Purchase Probability Constraints -----------
    if cone == 1   # MIECP Formulation: Use exponential cones
        @variable(model, w[1:N, 1:T])  
        @constraint(model, [i in 1:N, t in 1:T], 
                   g[i, t] * u[i] + sum(beta[i, m] * z[i, t, m] for m in 1:M) == w[i, t])
        @constraint(model, [i in 1:N, t in 1:T], 
                   [w[i, t], g[i, t], y[i, t]] in MOI.ExponentialCone())
    else 
        # ENV Formulation: Use convex envelope constraints
        for i in 1:N, t in 1:T
            ConvEnvelope(model, min(M+1, t), i, t, u, beta, y, g, z)
        end 
    end 
  
    # ----------- Concave Envelope Constraints ----------- 
    for i in 1:N, t in 1:T, ranking in rank_set
        lovasz_constr(model, i, t, M, ranking, u, betanew, y, g, z, s)
    end 

    # ----------- Solve Model -----------
    if IP == 1   
        set_binary.(x)  
        
        if cone == 1  
            set_optimizer_attribute(model, "MSK_DPAR_MIO_TOL_REL_GAP", stop_gap)
            if time_bound > 0 
                set_optimizer_attribute(model, "MSK_DPAR_OPTIMIZER_MAX_TIME", time_bound) 
            end  
            set_optimizer_attribute(model, MOI.RawOptimizerAttribute("MSK_IPAR_LOG"), 1)   

            println("\n","--- Starting Mosek Optimization via JuMP ---") 
            open(logname, "w") do file
                redirect_stdout(file) do
                    optimize!(model)
                end 
            end 
            println("\n", "--- Mosek Optimization Finished ---") 
 
        else  
            set_optimizer_attribute(model, "MIPGap", stop_gap)
            set_optimizer_attribute(model, "LogFile", logname) 
            if time_bound > 0
                set_time_limit_sec(model, time_bound)
            end
            optimize!(model)
        end 
        
        gap = MOI.get(model, MOI.RelativeGap())
        return objective_value(model), solve_time(model), gap, value.(x)
    else 
        optimize!(model)
        return objective_value(model), solve_time(model), 0, value.(x)
    end 
end

# ----------- A strong formulation for solving HAP with M=2 and satiation effects -----------
"""
    MIP_M2_Strong(T, N, M, r, beta, u, v0, time_bound, IP, cardinal, stop_gap, 
                  init, init_X, logname, s, output)

This model incorporates rotated second-order cone constraints for tighter relaxations. Only applicable for memory length M = 2 and negative history-dependent effects. This model can be used to compute IP benchmark when evaluating cyclic policies and heuristics and to solve the case study problem. 
"""
function MIP_M2_Satiation_Strong(T, N, M, r, beta, u, v0, time_bound, IP, cardinal, showtime, stop_gap,
                       init, init_X, logname, s, output)   
    model = Model(Gurobi.Optimizer)
    set_optimizer_attribute(model, "MIPGap", stop_gap)      
    set_optimizer_attribute(model, "OutputFlag", output)    

    # ----------- Bounds Calculation ----------- 
    low_rho = 1 / (v0 + sum(exp(u[i] + dot(s[i, :], beta[i, :])) for i in 1:N)) 
    up_rho = 1 / v0
   
    # ----------- Decision Variables -----------
    @variable(model, 0 <= x[1:N, 1:T] <= 1)          
    @variable(model, low_rho <= rho[1:T] <= up_rho)   
    @variable(model, 0 <= y[1:N, 1:T])                # Purchase probabilities (primal)
    @variable(model, 0 <= w[1:N, 1:T])                # Auxiliary purchase probabilities 
    @variable(model, 0 <= g[1:N, 1:T])                # Auxiliary: g_it = ρ_t * x_it
    @variable(model, 0 <= z[1:N, 1:T, 1:M])          # Auxiliary: z_itm = g_it * x_{i,t-m}
    @variable(model, 0 <= h[1:N, 1:T, 1:M])          # Auxiliary: h_itm = x_it * x_{i,t-m}

    # ----------- Objective Function ----------- 
    @objective(model, Max, sum(r[i] * y[i, t] for i in 1:N for t in 1:T))  

    # ----------- Initial Solution (Warm Start) -----------
    if init == 1 
        for t in 1:T, i in 1:N
            set_start_value(x[i, t], init_X[i, t]) 
        end 
    end 
  
    # ----------- Main Constraints -----------
    @constraint(model, [i in 1:N], sum(x[i, t] for t in 1:T) <= showtime)
    for t in 1:T 
        @constraint(model, v0 * rho[t] + sum(y[i, t] for i in 1:N) == 1) 
        @constraint(model, sum(x[i, t] for i in 1:N) <= cardinal)
        
        # Rotated Second-Order Cone Constraint for strengthening
        @constraint(model, [0.5 * rho[t]; v0 + sum(w[i, t] for i in 1:N); 1] 
                   in RotatedSecondOrderCone())

        for i in 1:N 
            McCormick(model, g[i, t], rho[t], x[i, t], low_rho, up_rho, 0, 1)   
            for m in 1:M
                if t - m >= 1   
                    # h_itm = x_it * x_{i,t-m}  
                    McCormick(model, h[i, t, m], x[i, t], x[i, t-m], 0, 1, 0, 1) 
                    # z_itm = g_it * x_{i,t-m}  
                    McCormick(model, z[i, t, m], g[i, t], x[i, t-m], 0, up_rho, 0, 1)  
                else  
                    @constraint(model, z[i, t, m] == 0)
                    @constraint(model, h[i, t, m] == 0)
                end  
            end 
        end 
    end 

    # ----------- Envelope Constraints -----------
    for i in 1:N, t in 1:T
        # Envelope constraints for primary variables (y, g, z)
        Envelope(model, i, t, u, beta, y, g, z)  
        # Envelope constraints for auxiliary variables (w, x, h)  
        Envelope(model, i, t, u, beta, w, x, h)
    end   

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

# ----------- Envelope constraints for the strong formulation -----------
function Envelope(model,i,t,u,beta,y,g,z)  
    if t==1 
        @constraint(model, y[i,t] == g[i,t]*exp(u[i]))
    elseif t==2
        @constraint(model, y[i,t]== g[i,t]*exp(u[i]) + z[i,t,1]*(exp(u[i]+beta[i,1]) - exp(u[i])))
    else  
        # convex envelope
        @constraint(model, y[i,t]*exp(-u[i])>= (exp(beta[i,1])-1)*z[i,t,1]
                                              +(exp(beta[i,2])-1)*z[i,t,2] + g[i,t] )
        @constraint(model, y[i,t]*exp(-u[i])>= (exp(beta[i,1]+beta[i,2])-exp(beta[i,2]))*z[i,t,1]
                                            +(exp(beta[i,2]+beta[i,1])-exp(beta[i,1]))*z[i,t,2]
                             + g[i,t]*(exp(beta[i,1]) + exp(beta[i,2]) - exp(beta[i,1] + beta[i,2])) )  
        
        # concave envelope 
        @constraint(model, y[i,t]*exp(-u[i]) <= (g[i,t] - z[i,t,1]) +
                                          (z[i,t,1] - z[i,t,2])*exp(beta[i,1]) +
                                           exp(beta[i,1] + beta[i,2])*z[i,t,2] ) 
        @constraint(model, y[i,t]*exp(-u[i]) <= (g[i,t] - z[i,t,2]) + 
                                           (z[i,t,2] - z[i,t,1])*exp(beta[i,2]) +
                                            exp(beta[i,1] + beta[i,2])*z[i,t,1] ) 
    end 
end  


#  ------------------- A strong formulation for M1 with second-order-cone constraints -----
function MIP_M1_Strong( T,N,M,r,beta,u,v0,IP, time_bound,stop_gap,init,init_X, output, logname, s,  cardinal, showtime  )   
    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)
    set_optimizer_attribute(model, "OutputFlag", output)
    set_optimizer_attribute(model, "MIPGap", stop_gap) 


    @variable(model, 0<= x[1:N,1:T]<=1)
    @variable(model, low_rho<=rho[1:T]<=up_rho) 
    @variable(model, y[1:N,1:T] ) 
    @variable(model, w[1:N,1:T] )    
    @variable(model, g[1:N,1:T]) 
    @variable(model, z[1:N,1:T])  

    # set initial solution 
    if init>0
        for t in 1:T
            for i in 1:N
                set_start_value(x[i,t], init_X[i,t])
            end 
        end 
    end 

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

    @constraint(model, [t in 1:T], v0*rho[t] + sum(y[i,t] for i in 1:N) == 1 ) # normalize fractional 
    @constraint(model, [t in 1:T], [0.5*rho[t]; v0+sum(w[i,t] for i in 1:N); 1] in RotatedSecondOrderCone())
    @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 ) 

    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>=2
                McCormick(model, z[i,t], rho[t], x[i,t-1], low_rho, up_rho, 0, 1 )  
            else
                @constraint(model, z[i,t] == 0  )
            end 
        end 
    end 

    # -------- envelope -- exponential cone 
   for i in 1:N 
        for t in 1:T
            if t>=2  
                # convex
                @constraint(model, y[i,t]*exp(-u[i])>= g[i,t]*exp(beta[i,1]) )
                @constraint(model, y[i,t]*exp(-u[i])>= g[i,t] + (exp(beta[i,1])-1)*z[i,t] )
                # concave 
                @constraint(model, y[i,t]*exp(-u[i])<= g[i,t] )
                @constraint(model, y[i,t]*exp(-u[i])<= g[i,t]*exp(beta[i,1]) + z[i,t]*(exp(beta[i,1])-1) + rho[t]*(1-exp(beta[i,1])))
            else 
                @constraint(model, y[i,t]*exp(-u[i]) == g[i,t])
            end 
        end 
    end 

    for i in 1:N 
        for t in 1:T
            if t>=2  
                # convex
                @constraint(model, w[i,t]*exp(-u[i])>= x[i,t]*exp(beta[i,1]) )
                @constraint(model, w[i,t]*exp(-u[i])>= x[i,t] + (exp(beta[i,1])-1)*x[i,t-1])
                # concave 
                @constraint(model, w[i,t]*exp(-u[i])<= x[i,t] )
                @constraint(model, w[i,t]*exp(-u[i])<= x[i,t]*exp(beta[i,1]) + x[i,t-1]*(exp(beta[i,1])-1) +  (1-exp(beta[i,1])))
            else 
                @constraint(model, w[i,t]*exp(-u[i]) == x[i,t])
            end 
        end 
    end 

    # optimization set up 
    if time_bound >0
        set_time_limit_sec(model, time_bound)
    end 
    set_optimizer_attribute(model, "LogFile", logname)

    if IP == 1 
        set_binary.(x)
    end 
    optimize!(model)

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



