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

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

# ==================== Optimization Model ====================================
# This module implements the projected cutting plane algorithm for solving (M+1)-cyclic policies. See EC.5.2 for more algorithm details. 
# ================================================================================


# ----------------- Bound-free formulation (12) ----------------
""" 
This implements formulation (12) in the paper with second-order cone constraints.
"""
function BoundFree(N, M, r, u, v0, beta, IP, stop_gap, time_bound, output, logname, s, cardinal)
    v = exp.(u)  # Exponential utilities
    K = M + 1    # Number of periods in cycle
 
    # 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

    # Initialize optimization model
    model = Model(Gurobi.Optimizer)

    # Decision Variables 
    @variable(model, 0 <= x[1:N, 1:K] <= 1)        
    @variable(model, low_rho <= rho[1:K] <= up_rho) 
    @variable(model, 0 <= g[1:N, 1:K])             # Auxiliary: rho * x  
    @variable(model, 0 <= Y[1:N, 1:N, 1:K])        # Cross-product terms Y[i,j,t]
    @variable(model, w[1:K])                       # Auxiliary: denominator terms

    # Objective Function 
    @objective(model, Max, sum(r[i] * v[i] * g[i, t] for i in 1:N for t in 1:K))

    # Constraints  
    @constraint(model, [i in 1:N], sum(x[i, t] for t in 1:K) <= 1) 
    @constraint(model, [t in 1:K], v0 * rho[t] + sum(v[i] * g[i, t] for i in 1:N) == 1) 
    @constraint(model, [t in 1:K], sum(x[i, t] for i in 1:N) <= cardinal)
    
    # Second-order cone formulation (replaces the nonlinear constraint rho[t] * w[t] = 1)
    @constraint(model, [t in 1:K], v0 + sum(v[i] * x[i, t] for i in 1:N) == w[t]) 
    @constraint(model, [t in 1:K], [rho[t] + w[t], 2, rho[t] - w[t]] in SecondOrderCone())

    # Constraints for g and Y
    @constraint(model, [t in 1:K, i in 1:N], g[i, t] <= rho[t]) 
    @constraint(model, [t in 1:K, i in 1:N], x[i, t] == g[i, t] + sum(v[j] * Y[i, j, t] for j in 1:N))
    @constraint(model, [t in 1:K, i in 1:N], Y[i, i, t] == g[i, t])
    for i in 1:N
        for j in 1:N 
            if i != j  
                @constraint(model, [t in 1:K], Y[i, j, t] <= g[i, t])
                @constraint(model, [t in 1:K], Y[i, j, t] <= g[j, t]) 
                @constraint(model, [t in 1:K], Y[i, j, t] >= g[i, t] + g[j, t] - rho[t])
            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), value.(x), solve_time(model), MOI.get(model, MOI.RelativeGap())   
    else 
        return objective_value(model), value.(x), solve_time(model), 0
    end
end 


# ------------------ Separation oracle for projected cutting plane algorithm (Algorithm 2) -------
""" 
Identifies violated inequalities for the convex hull of a single term.
Returns:
- cut: JuMP constraint representing the violated inequality, or 0 if no violation
"""
function separate1Term(x::VariableRef, y::Vector{VariableRef}, rho::VariableRef, v0, a, xv, yv, ρv, ind)
    sv = xv / ρv  # Normalized assortment value
    tv = yv / ρv  # Normalized probability values
    n = length(tv) # Number of products
    
    # CONVEX EVALUATION (Lower bound check) 
    A = []   
    for k = 1:n 
        if (tv[k] + tv[ind] > 1) && (k !== ind)  
            push!(A, k)        
        end
    end
    
    # Compute convex evaluation
    if length(A) == 0
        conv_eval = (v0 + a[ind]) * tv[ind]  
    else
        conv_eval = (v0 + a[ind]) * tv[ind] + sum(a[k] * (tv[ind] + tv[k] - 1) for k ∈ A)  
    end
    
    # CONCAVE EVALUATION (Upper bound check) 
    B = []   
    C = []   
    for k = 1:n 
        if tv[k] >= tv[ind]
            push!(B, k)  
        else
            push!(C, k) 
        end
    end
    
    # Compute concave evaluation
    if length(C) == 0
        conc_eval = v0 * tv[ind] + sum(a[k] * tv[ind] for k ∈ B) 
    else
        conc_eval = v0 * tv[ind] + sum(a[k] * tv[ind] for k ∈ B) + sum(a[k] * tv[k] for k ∈ C)
    end
    
    # CHECK FOR CONSTRAINT VIOLATIONS
    # If current point is above upper bound, add cutting plane
    if sv >= (conc_eval + 0.00001)  
        if length(C) == 0
            cut = @build_constraint(x <= v0 * y[ind] + sum(a[k] * y[ind] for k ∈ B))  
        else
            cut = @build_constraint(x <= v0 * y[ind] + sum(a[k] * y[ind] for k ∈ B) + sum(a[k] * y[k] for k ∈ C))  
        end
    # If current point is below lower bound, add supporting hyperplane
    elseif sv <= (conv_eval - 0.000001)
        if length(A) == 0
            cut = @build_constraint(x >= (v0 + a[ind]) * y[ind])
        else
            cut = @build_constraint(x >= (v0 + a[ind]) * y[ind] + sum(a[k] * (y[ind] + y[k] - rho) for k ∈ A))
        end
    else
        cut = 0  # No violation found
    end
    
    return cut
end

# -------------- Bound-free formulation with separation algorithm (Algorithm 3) ------------ 
function BoundFreewithSepa(N, M, r, u, v0, beta, IP, iter, stop_gap, time_bound, logname, s, cardinal)
    v = exp.(u)
    K = M + 1  # Cycle length 

    # Compute rho bounds
    low_rho = 1 / (v0 + sum(exp(u[i] + dot(s[i, :], beta[i, :])) for i in 1:N))
    up_rho = 1 / v0

    # Initialize model
    model = Model(Gurobi.Optimizer) 

    # Decision Variables 
    @variable(model, 0 <= x[1:N, 1:K] <= 1)         
    @variable(model, low_rho <= rho[1:K] <= up_rho) # Normalization factors
    @variable(model, 0 <= y[1:N, 1:K])             # McCormick variables: y ≈ x * rho
    @variable(model, w[1:K])                       # Auxiliary variables

    # Objective Functions 
    @objective(model, Max, sum(r[i] * v[i] * y[i, t] for i in 1:N for t in 1:K))

    # Constraints 
    @constraint(model, [i in 1:N], sum(x[i, t] for t in 1:K) <= 1)  # Non-overlapping
    @constraint(model, [t in 1:K], sum(x[i, t] for i in 1:N) <= cardinal) # Cardinality constraint
    @constraint(model, [t in 1:K], v0 * rho[t] + sum(v[i] * y[i, t] for i in 1:N) == 1)  # Normalization
    @constraint(model, [t in 1:K], v0 + sum(v[i] * x[i, t] for i in 1:N) == w[t]) # Auxiliary definition
    @constraint(model, [t in 1:K], [rho[t] + w[t], 2, rho[t] - w[t]] in SecondOrderCone()) # Conic constraint

    # McCormick relaxation for y[i,t] = x[i,t] * rho[t]
    for t in 1:K  
        for i in 1:N 
            McCormick(model, y[i, t], x[i, t], rho[t], 0, 1, low_rho, up_rho)
        end 
    end 
    
    # PHASE 1: Solve initial LP relaxation
    optimize!(model)
    run_time_1 = solve_time(model)
    xv = value.(x) 
    yv = value.(y) 
    rhov = value.(rho)
    
    # PHASE 2: Cutting plane generation
    t1 = now()
    Round = 1
    while Round <= iter
        pool = Any[]  # Store generated cuts
        
        # Generate cuts for each period and product
        for t = 1:K
            yvec = y[:, t]      # Probability variables for period t
            yvalue = yv[:, t]   # Current probability values for period t
            for i = 1:N
                if v[i] > 0.0000001  # Only consider products with positive utility
                    cut = separate1Term(x[i, t], yvec, rho[t], 1, v, xv[i, t], yvalue, rhov[t], i)
                    if cut !== 0
                        push!(pool, cut)
                    end    
                end
            end 
        end
        
        # Add generated cuts to model
        if length(pool) > 0
            for k in 1:length(pool)
                add_constraint(model, pool[k])
            end
        end 
        
        # Re-solve with new cuts
        optimize!(model)
        xv = value.(x) 
        yv = value.(y) 
        rhov = value.(rho)   
        Round = Round + 1
    end
    t2 = now()

    # PHASE 3: Solve integer program (if requested)
    if IP == 1
        set_binary.(x)
        set_time_limit_sec(model, time_bound)
        set_optimizer_attribute(model, "MIPGap", stop_gap)
        set_optimizer_attribute(model, "LogFile", logname)
        optimize!(model)
        gap = MOI.get(model, MOI.RelativeGap())
        run_time_2 = solve_time(model)
    else
        gap = 0 
        run_time_2 = 0
    end

    # COMPUTE TOTAL RUNTIME
    add_cut_time = (t2 - t1).value / 1000   
    if IP == 1 
        total = run_time_1 + run_time_2 + add_cut_time
    else 
        total = run_time_1 + add_cut_time 
    end 

    return objective_value(model), value.(x), solve_time(model), total, gap
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


# ============================ Cyclic policy without non-overlapping constraint ==============================
function Two_Cyclic(N,r,u,beta,v0,time_bound,IP,cardinal, stop_gap)
    v = exp.(u)
    d = exp.(u+beta) - exp.(u)
    K = 2 
    M = 1
    low_rho = 1/(v0 + sum( v[i] for i in 1:N))
    up_rho = 1/v0 
    beta_reshape = zeros(N,2)
    beta_reshape[:,1] = beta  
    model = Model(Gurobi.Optimizer)
    set_optimizer_attribute(model, "MIPGap", stop_gap)

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

    # ===============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:K )    )   
    # ========== constraint ===================
    for t in 1:K 
        @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 
            McCormick(model, y[i,t], x[i,t], rho[t], 0, 1, low_rho, up_rho )
            McCormick(model, w[i,t], delta[i], rho[t], 0, 1, low_rho, up_rho )  
            McCormick(model, delta[i], x[i,1], x[i,2], 0, 1, 0, 1 )  
        end 
    end 
 
    if time_bound >0
        set_time_limit_sec(model, time_bound)
    end 
    optimize!(model)

    # compute gap 
    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 


function Three_Cyclic(N,M,r,beta,u,v0,IP,cardinal,cone,stop_gap,time_bound)
    K = 3 
    # define an opt model 
    if cone>0
        model = Model(Mosek.Optimizer) 
    else
        model = Model(Gurobi.Optimizer)
        set_optimizer_attribute(model, "MIPGap", stop_gap)
    end 
    # define variables 
    low_rho = 1/(v0 + sum(exp(u[i]) for i in 1:N))
    up_rho = 1/v0 

    # assortment binary variable 
    if IP >0
        @variable(model, x[1:N,1:K],Bin)
    else
        @variable(model, 0<= x[1:N,1:K]<=1)
    end 
    @variable(model, low_rho<=rho[1:K]<=up_rho) # rho_t 
    @variable(model, 0<= y[1:N,1:K] ) # y_it = rho_t*x_it*alpha_it   
    @variable(model, 0<=g[1:N,1:K]<=up_rho ) #  g_it = rho_t * x_it \gamma 
    @variable(model, 0<=z[1:N,1:K,1:M]<=up_rho ) # z_itm = rho_t * x_it * x_it-m 

    # define constraints 
    @objective(model, Max, sum(r[i]*y[i,t] for i in 1:N for t in 1:K  )  )  
 
    # normalize and McCormick
    for t in 1:K
        @constraint(model, v0*rho[t] + sum(y[i,t] for i in 1:N) == 1 ) # normalize fractional
        @constraint(model, sum(x[i,t] for i in 1:N) <= cardinal) # cardinality 
        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, z[i,t,1], g[i,t], x[i,t-1], 0, up_rho, 0, 1 )
                McCormick(model, z[i,t,2], g[i,t], x[i,t-2], 0, up_rho, 0, 1 )
            elseif t == 2
                McCormick(model, z[i,t,1], g[i,t], x[i,t-1], 0, up_rho, 0, 1 )
                McCormick(model, z[i,t,2], g[i,t], x[i,K], 0, up_rho, 0, 1 )
            else
                McCormick(model, z[i,t,1], g[i,t], x[i,K], 0, up_rho, 0, 1 )
                McCormick(model, z[i,t,2], g[i,t], x[i,K-1], 0, up_rho, 0, 1 )
            end 
        end 
    end 

    # ----------convex envelope -- exponential cone 
    if cone == 1
        @variable(model, w[1:N, 1:K]) 
        for i in 1:N
            for t in 1:K                 
                @constraint(model, g[i,t]*u[i]+ sum(beta[i,m]*z[i,t,m] for m in 1:M) == w[i,t])
                @constraint(model, [w[i,t], g[i,t],y[i,t] ] in MOI.ExponentialCone())
            end 
        end 
    else
        Weak_env_M2(model,N,K,u,beta,y,g,z,1)
    end 
    # concave envelope 
    Weak_env_M2(model,N,K,u,beta,y,g,z,0)

    if time_bound >0
        set_time_limit_sec(model, time_bound)
    end 
    optimize!(model)
    
    # compute gap 
    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 
 

function cycle_policy(T,N,M,r,beta,u,v0,time_bound,IP,stop_gap,cardinal,cone)
    # use (M+1)-cyclic policy for general T 
    K = M+1 
    kapa = trunc(T/K)
    remain = Int(T-kapa*K) 
    if M==1 
        res_c = Two_Cyclic(N,r,u,beta,v0,time_bound,IP,cardinal,stop_gap)
    else
        res_c = Three_Cyclic(N,M,r,beta,u,v0,IP,cardinal,cone,stop_gap,time_bound)
    end  
    X = Matrix{Float64}(undef,N,T) 
    for k in 1:kapa  
        for ll in 1:K  
            X[:, Int(K*(k-1)+ll)] = res_c[2][:,ll]
        end 
    end 
    for ll in 1:remain 
        X[:,  Int(K*kapa + ll)] = res_c[2][:,ll]
    end      
    X = round.(X;digits=0)
    rev = FeasibleVarValue(N,T,M,u,beta,v0,X,r) 
    return rev[1] 
end 

function Weak_env_M2(model,N,T,u,beta,y,g,z,conv)
    if conv == 1
        for i in 1:N 
            for t in 1:T
                @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])) )
            end 
        end 
    else 
        for i in 1:N 
            for t in 1:T
                @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]) + z[i,t,2]*exp(beta[i,1] + beta[i,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]) + z[i,t,1]*exp(beta[i,1] + beta[i,2])  )
            end 
        end 
    end
end  
function FeasibleVarValue(N, T, M, u, beta, v0, x, r)  
    y = zeros(N, T)     
    rho = zeros(T)       
    g = zeros(N, T)      # Auxiliary variable g_it = ρ_t * x_it
    z = zeros(N, T, M)   # Auxiliary variable z_itm = g_it * x_{i,t-m}
    a = zeros(N, T)      # Attraction values
 
    for i in 1:N, t in 1:T
        if (t >= 3) && (M > 1)  # Two periods of memory available
            a[i, t] = exp(u[i] + beta[i, 1] * x[i, t-1] + beta[i, 2] * x[i, t-2])
        elseif t >= 2  # One period of memory available  
            a[i, t] = exp(u[i] + beta[i, 1] * x[i, t-1])
        else  # No memory (first period)
            a[i, t] = exp(u[i])
        end 
    end 
    
    for t in 1:T 
        rho[t] = 1 / (v0 + sum(x[i, t] * a[i, t] for i in 1:N))
        
        for i in 1:N 
            g[i, t] = rho[t] * x[i, t]                    
            y[i, t] = rho[t] * x[i, t] * a[i, t]      
             
            if t >= 3  # Both memory periods available
                for m in 1:M
                    z[i, t, m] = g[i, t] * x[i, t-m]
                end 
            elseif t >= 2  # Only first memory period available
                z[i, t, 1] = g[i, t] * x[i, t-1]
            else  # No memory periods available
                for k in 1:M
                    z[i, t, k] = 0
                end 
            end 
        end
    end 
    
    # Calculate total revenue across all periods
    total_r = sum(r[i] * y[i, t] for i in 1:N for t in 1:T)
    return total_r, x, rho, g, z, y 
end 