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

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

# ==================================================================================
# Cutting plane algorithm for cases with large memory length 
# ================================================================================== 


# ----------- McCormick Envelope Constraints for Bilinear Terms -----------
""" 
Adds McCormick envelope constraints. 
- 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 >= 0)      
    @constraint(m, -z + yu*x + xl*y - xl*yu >= 0)     
    @constraint(m, -z + yl*x + xu*y - xu*yl >= 0)     
    @constraint(m, z - yu*x - xu*y + xu*yu >= 0)     
end

 
# ----------- Attraction Value Calculation -----------
""" 
- 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)  
    # Construct historical assortment matrix h based on current state and ranking
    h = zeros(M+1, M)
    h[1, :] = s[i, :]  
    for m in 2:M+1 
        h[m, :] = h[m-1, :]  # Copy previous state
        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 

# -------- Convex lower estimation constraints using subgradient cuts ----------- 
""" 
Provides linear underestimation of the exponential function
- xhis_set: set of historical assortment points 
""" 
function exp_constr(model, M, N, T, xhis_set, u, beta, y, g, z)
    for xhis in xhis_set
        for i in 1:N 
            for t in 1:T 
                # Calculate attraction value at historical point xhis
                attrv = exp(u[i] + dot(beta[i, :], xhis))
                # Add subgradient constraint - provides convex lower bound
                @constraint(model, y[i,t] >= 
                    g[i,t] * attrv + 
                    sum((z[i,t,m] - g[i,t]*xhis[m]) * beta[i,m] * attrv for m in 1:M)
                )
            end 
        end 
    end 
end 



#  ------------------- Separation Oracle - Implementation of Algorithm 4 ---------------
"""
Identifies violated constraints and generates cutting planes 
- yv, gv, zv: current variable values 
- y, g, z: JuMP variables for constraint generation 
- delta: violation tolerance threshold
"""
function Separation_LargeM(M, N, T, yv, gv, zv, u, beta, betanew, y, g, z, s, delta)
    pool = []  # Pool to store newly generated constraints
    conc_cut = 0  # Counter for concave constraint cuts
    conv_cut = 0  # Counter for convex constraint cuts
    
    # ------- Concave part: Check for violated Lovász constraints -------
    for i in 1:N 
        for t in 1:T 
            # Find new permutation based on current solution values 
            sortv = [zv[i,t,m]*(1-s[i,m]) + (gv[i,t] - zv[i,t,m])*s[i,m] for m in 1:M]
            newrank = sortperm(sortv, rev=true)  # Sort in descending order

            # Define current historical state matrix h
            newh = zeros(M+1, M)
            newh[1, :] = s[i, :]
            for m in 2:M+1 
                newh[m, :] = newh[m-1, :]
                if s[i, newrank[m-1]] > 0 
                    newh[m, newrank[m-1]] = newh[m, newrank[m-1]] - 1
                else 
                    newh[m, newrank[m-1]] = newh[m, newrank[m-1]] + 1 
                end 
            end 

            # Calculate right-hand side value of the Lovász constraint
            rhsv = attraction(i, t, 0, M, u, betanew, newh) * 
                   (gv[i,t] - (1-s[i,newrank[1]])*zv[i,t,newrank[1]] - 
                    s[i,newrank[1]]*(gv[i,t]-zv[i,t,newrank[1]])) + 
                   sum(attraction(i, t, k, M, u, betanew, newh) *
                       ((1-s[i,newrank[k]])*zv[i,t,newrank[k]] + 
                        s[i,newrank[k]]*(gv[i,t]-zv[i,t,newrank[k]]) - 
                        (1-s[i,newrank[k+1]])*zv[i,t,newrank[k+1]] - 
                        s[i,newrank[k+1]]*(gv[i,t]-zv[i,t,newrank[k+1]]))   
                       for k in 1:M-1) +  
                   attraction(i, t, M, M, u, betanew, newh) * 
                   ((1-s[i,newrank[M]])*zv[i,t,newrank[M]] + 
                    s[i,newrank[M]]*(gv[i,t]-zv[i,t,newrank[M]]))
            
            # Check if constraint is violated beyond tolerance
            if yv[i,t] > (rhsv + delta)
                # Generate new cutting constraint
                cut = @build_constraint(y[i,t] <= 
                    attraction(i, t, 0, M, u, betanew, newh) * 
                    (g[i,t] - (1-s[i,newrank[1]])*z[i,t,newrank[1]] - 
                     s[i,newrank[1]]*(g[i,t]-z[i,t,newrank[1]])) + 
                    sum(attraction(i, t, k, M, u, betanew, newh) *
                        ((1-s[i,newrank[k]])*z[i,t,newrank[k]] + 
                         s[i,newrank[k]]*(g[i,t]-z[i,t,newrank[k]]) - 
                         (1-s[i,newrank[k+1]])*z[i,t,newrank[k+1]] - 
                         s[i,newrank[k+1]]*(g[i,t]-z[i,t,newrank[k+1]]))   
                        for k in 1:M-1) +  
                    attraction(i, t, M, M, u, betanew, newh) * 
                    ((1-s[i,newrank[M]])*z[i,t,newrank[M]] + 
                     s[i,newrank[M]]*(g[i,t]-z[i,t,newrank[M]]))
                ) 
                push!(pool, cut)
                conc_cut += 1 
            end 
        end 
    end 

    # ----------- Convex part: Check for violated subgradient constraints -----------
    # Calculate current historical assortments from the solution
    xhis = zeros(N, T, M)
    for i in 1:N, t in 1:T, m in 1:M 
        if t-m >= 1  # Ensure valid time index
            if gv[i,t] > 0  # Avoid division by zero
                xhis[i,t,m] = zv[i,t,m] / gv[i,t] 
            end  
        end 
    end 
    
    # Check subgradient constraints for each product and time period
    for i in 1:N , t in 1:T  
        # Calculate right-hand side value
        attrv = exp(u[i] + dot(beta[i, :], xhis[i,t,:]))
        rhsv = gv[i,t] * attrv + 
                sum((zv[i,t,m] - gv[i,t]*xhis[i,t,m]) * beta[i,m] * attrv for m in 1:M)  
    
        # Check if constraint is violated
        if yv[i,t] < (rhsv - delta)
            cut = @build_constraint(y[i,t] >= 
                g[i,t] * attrv + 
                sum((z[i,t,m] - g[i,t]*xhis[i,t,m]) * beta[i,m] * attrv for m in 1:M)
            ) 
            push!(pool, cut)
            conv_cut += 1 
        end  
    end 
    return pool, conc_cut, conv_cut 
end 

# ------------------- Main Algorithm Function - Large M Method with Callback ------------- 
"""
Implements the complete cutting plane algorithm for assortment optimization
- rank_set: set of initial rankings for Lovász constraints
- xhis_set: set of initial historical points for subgradient cuts 
"""
function LargeM_Callback(T, N, M, r, beta, u, v0, init, init_X, IP, time_bound, stop_gap, 
                        output, logname, s, betanew, rank_set, xhis_set, epsilon, delta, 
                        cardinal, showtime) 
    t0 = now()  # Record start time
     
    model = Model(Gurobi.Optimizer) 
    set_optimizer_attribute(model, "OutputFlag", output) 
 
    low_rho = 1 / (v0 + sum(exp(u[i] + dot(s[i,:], beta[i,:])) for i in 1:N))
    up_rho = 1 / v0

    # Define decision variables
    @variable(model, 0 <= x[1:N, 1:T] <= 1)      # start with continuous relaxation)
    @variable(model, low_rho <= rho[1:T] <= up_rho)   
    @variable(model, 0 <= y[1:N, 1:T])        
    @variable(model, 0 <= g[1:N, 1:T] <= up_rho) 
    @variable(model, 0 <= z[1:N, 1:T, 1:M] <= up_rho)  
 
    # Set objective function - maximize total revenue
    @objective(model, Max, sum(r[i] * y[i,t] for i in 1:N for t in 1:T))

    # Set initial solution (if provided)
    if init > 0
        for t in 1:T, i in 1:N
            set_start_value(x[i,t], init_X[i,t]) 
        end 
    end   
  
    # Add basic 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)

    # Add McCormick constraints for bilinear and trilinear term relaxations
    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)   
            for m in 1:M
                if t-m >= 1  # Ensure valid time index
                    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)   
                end 
            end 
        end 
    end 

    # Add convex envelope constraints (subgradient cuts) 
    exp_constr(model, M, N, T, xhis_set, u, beta, y, g, z)
  
    # Add concave envelope constraints (Lovász constraints) 
    for i in 1:N, t in 1:T 
        for item in rank_set
            lovasz_constr(model, i, t, M, item, u, betanew, y, g, z, s)
        end  
    end 
    
    # Solve continuous relaxation
    obj_old = 100 * T  # Initialize objective value
    optimize!(model)
    obj_new = objective_value(model)
     
    # Get current solution values
    xv = value.(x)
    yv = value.(y)
    zv = value.(z)
    gv = value.(g)

    t1 = now() 

    # Cutting plane phase for continuous relaxation
    Round = 1 
    total_lp_lovasz = 0   # Total Lovász cuts in LP phase
    total_lp_grad = 0     # Total subgradient cuts in LP phase
    
    # Iteratively add cutting planes until convergence
    println("Starting cutting plane iterations...")
    while (obj_old - obj_new) > epsilon
        obj_old = obj_new  
        # Call separation oracle to find violated constraints
        pool, conc, conv = Separation_LargeM(M, N, T, yv, gv, zv, u, beta, betanew, y, g, z, s, delta)
        # Update cutting plane counters
        total_lp_lovasz += conc 
        total_lp_grad += conv  
        # Add new cutting constraints
        if length(pool) > 0
            for k in 1:length(pool)
                add_constraint(model, pool[k]) 
            end
        else
            println("No violated constraints found, stopping cutting plane phase.")
            break
        end 
        
        # Re-solve with new constraints
        optimize!(model)
        obj_new = objective_value(model)
        
        # Update solution values
        yv = value.(y)
        zv = value.(z)
        gv = value.(g)
        
        println(" ******* Continuous Round ", Round, "  Added cuts: ", length(pool),
               "  Lovász cuts: ", conc, "  Subgradient cuts: ", conv, " *******")
        Round += 1
        
        # Safety check to prevent infinite loops
        if Round > 1000
            println("Maximum cutting plane iterations reached.")
            break
        end
    end
    
    t2 = now() 
    relax_add = (t2-t1).value / 1000  # Time spent in continuous relaxation phase

    # Lazy constraint callback function: called during branch-and-bound when integer solutions are found
    total_lazy_lovasz = 0  # Total lazy Lovász cuts
    total_lazy_grad = 0    # Total lazy subgradient cuts

    function lazy_cb(cb_data)   
        status = callback_node_status(cb_data, model)
        
        # When a feasible integer solution is found
        if status == MOI.CALLBACK_NODE_STATUS_INTEGER 
            # Get current solution values
            yv = callback_value.(Ref(cb_data), y)
            gv = callback_value.(Ref(cb_data), g)
            zv = callback_value.(Ref(cb_data), z)
 
            # Generate lazy constraints using separation oracle
            pool, conc, conv = Separation_LargeM(M, N, T, yv, gv, zv, u, beta, betanew, y, g, z, s, delta)
            
            # Add lazy constraints to the model
            for con in pool 
                MOI.submit(model, MOI.LazyConstraint(cb_data), con) 
            end  
            
            # Update lazy cut counters
            total_lazy_lovasz += conc
            total_lazy_grad += conv
        end 
        return 
    end

    # Solve integer programming problem if requested
    if IP == 1  
        println("Starting integer programming phase...") 
        set_binary.(x) 
        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 
        
        # Enable lazy constraints
        set_optimizer_attribute(model, "LazyConstraints", 1)  
        MOI.set(model, MOI.LazyConstraintCallback(), lazy_cb) 
        # Solve integer program
        optimize!(model)  
        t3 = now()
        gap = MOI.get(model, MOI.RelativeGap())  # Get final optimality gap
        ip_add = (t3-t2).value / 1000  # Time spent in integer programming phase
        alltime = (t3-t0).value / 1000  # Total solution time
  
        return (objective_value(model), alltime, gap, total_lp_lovasz, total_lp_grad, 
                total_lazy_lovasz, total_lazy_grad, relax_add, ip_add, value.(x))
    else 
        return (objective_value(model), total_lp_lovasz, total_lp_grad, relax_add)
    end 
end