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

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

# ==================================================================================
# Heuristic Algorithms for HAP 
# ================================================================================== 

# -------------- Compute revenue for a given assortment plan ------- 
""" 
Computes the total revenue for a given assortment plan and validates feasibility for memory length M ∈ {1,2}.
"""
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 

# -------------------- Single period MNL without constraints ------------- 
""" 
Solves the unconstrained single-period Multinomial Logit (MNL) assortment optimization.
"""
function MNL_unconstraint(v0, v, r, N) 
    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                              
    
    # ----------- Variables -----------
    @variable(model, low_rho <= rho <= up_rho)    
    @variable(model, 0 <= y[1:N])                
    
    # ----------- Constraints ----------- 
    @constraint(model, rho + sum(v[i] * y[i] for i in 1:N) == 1)  
    @constraint(model, [i in 1:N], y[i] <= rho)
    
    # ----------- Objective ----------- 
    @objective(model, Max, sum(r[i] * v[i] * y[i] for i in 1:N))  
 
    b = @timed optimize!(model)  
    x = value.(y) ./ value.(rho)  # Convert to assortment probabilities
    return objective_value(model), x, b.time
end 

# -------------------- Implements the Sequential-RO policy --------------
function SequentialRO(v0, u, beta, r, N, M, T) 
    X = zeros(N, T + M)  # Extended dimension for historical assortment computation
    Time = 0            # Total computation time tracker
    
    # ----------- Sequential Decision Making -----------
    for t in (M+1):(T+M)   
        
        # Compute current attraction value 
        v = [exp(u[i] + sum(beta[i, k] * X[i, t-k] for k in 1:M)) for i in 1:N]
        
        # Solve single-period unconstrained MNL optimization
        obj, xsolu, runtime = MNL_unconstraint(v0, v, r, N)
        
        # Store solution and update computation time
        X[:, t] = xsolu 
        Time += runtime
    end 
    
    # Round to binary decisions  
    X = round.(X; digits=0)  
    
    # Compute total revenue using the feasible solution
    rev = FeasibleVarValue(N, T, M, u, beta, v0, X[:, (M+1):(T+M)], r)
    
    return rev[1], Time, X[:, (M+1):(T+M)]
end 


# ------------------- Compute single period revenue --------------
function single_period_revenue(his_x, x, u, beta, v0, r, N) 
    u_change_m = his_x * transpose(beta)    # Matrix multiplication (N×M) × (M×N) = (N×N)
    u_change = LinearAlgebra.diag(u_change_m)  # Extract diagonal elements (N×1)
    v = exp.(u + u_change) 
    suma = v0 + sum(v[i] * x[i] for i in 1:N) 
    sumr = sum(r[i] * v[i] * x[i] for i in 1:N) / suma 
    return sumr 
end 

# ----------- Implements the Sequential Leave-One-Satiation-Product-Out policy ---------
function ro_oneout(N, M, T, r, u, beta, v0) 
    swt_x = zeros(N, T)  # Assortment plan  
    
    # ----------- Sequential Decision Making -----------
    for t in 1:T  
        his_x = zeros(N, M)  # Historical assortment matrix (N×M)
        if t > 1  
            for tt in 1:min(t-1, M)
                his_x[:, tt] = swt_x[:, t-tt]  # tt-th lag = t-tt period
            end 
        end  
        
        # ----------- Compute Current Utilities -----------
        u_change_m = his_x * transpose(beta)        
        u_change = LinearAlgebra.diag(u_change_m)   
        v = exp.(u + u_change)                      
        
        # ----------- Base Unconstrained Solution -----------
        resmnl = MNL_unconstraint(v0, v, r, N)
        xbase = resmnl[2]   
        
        # ----------- Identify and Evaluate Removals -----------
        r_base = 0          # Best revenue found so far
        cut = 0             # Index of product to remove (0 = no removal)
        
        # Identify satiation products (negative first-period memory effect)
        swt_item = Set{Int}()
        for i in 1:N 
            if (xbase[i] > 0) && (beta[i, 1] < 0)  # Product is offered AND has satiation effect
                push!(swt_item, i)
            end 
        end 
        
        # Evaluate removal of each satiation product
        for j in swt_item  
            cut_j_x = copy(xbase)
            cut_j_x[j] = 0  
            r_j = single_period_revenue(his_x, cut_j_x, u, beta, v0, r, N)  
            if r_j > r_base 
                r_base = r_j 
                cut = j 
            end 
        end 
        
        # ----------- Make Final Decision -----------
        if cut > 0  # Removal improves revenue
            cut_j_x = copy(xbase)
            cut_j_x[cut] = 0
            swt_x[:, t] = cut_j_x 
        else  # Keep base assortment
            swt_x[:, t] = xbase 
        end 
    end 
    
    # ----------- Compute Total Revenue -----------
    rev = FeasibleVarValue(N, T, M, u, beta, v0, swt_x, r) 
    return rev[1], swt_x 
end