data_sample.jl

"""
$TYPEDEF

Data sample data structure.
Its main purpose is to store datasets generated by the benchmarks.
It has 3 main (optional) fields: features `x`, cost parameters `θ`, and solution `y`.
Additionally, it has a `context` field (solver kwargs, spread into the maximizer as
`maximizer(θ; sample.context...)`) and an `extra` field (non-solver data, never passed
to the maximizer).

The separation prevents silent breakage from accidentally passing non-solver data
(e.g. a scenario, reward, or step counter) as unexpected kwargs to the maximizer.

# Fields
$TYPEDFIELDS
"""
struct DataSample{
    CTX<:NamedTuple,
    EX<:NamedTuple,
    F<:Union{AbstractArray,Nothing},
    S<:Union{AbstractArray,Nothing},
    C<:Union{AbstractArray,Nothing},
}
    "input features (optional)"
    x::F
    "intermediate cost parameters (optional)"
    θ::C
    "output solution (optional)"
    y::S
    "context information as solver kwargs, e.g. instance, graph, etc."
    context::CTX
    "additional data, never passed to the maximizer, e.g. scenario, objective value, reward,
    step count, etc. Can be used for any purpose by the user, such as plotting utilities."
    extra::EX
end

"""
$TYPEDSIGNATURES

Constructor for `DataSample` with keyword arguments.

All keyword arguments beyond `x`, `θ`, `y`, and `extra` are collected into the `context`
field (solver kwargs). The `extra` keyword accepts a `NamedTuple` of non-solver data.

Fields in `context` and `extra` must be disjoint. An error is thrown if they overlap.
Both can be accessed directly via property forwarding.

# Examples
```julia
# Instance goes in context
d = DataSample(x=[1,2,3], θ=[4,5,6], y=[7,8,9], instance="my_instance")
d.instance  # "my_instance" (from context)

# Scenario goes in extra
d = DataSample(x=x, y=y, instance=inst, extra=(; scenario=ξ))
d.scenario  # ξ (from extra)

# State goes in context, reward in extra
d = DataSample(x=x, y=y, instance=state, extra=(; reward=-1.5))
d.instance  # state (from context)
d.reward    # -1.5 (from extra)
```
"""
function DataSample(; x=nothing, θ=nothing, y=nothing, extra=NamedTuple(), kwargs...)
    context = (; kwargs...)
    overlap = intersect(keys(context), keys(extra))
    if !isempty(overlap)
        error("Keys $(collect(overlap)) appear in both context and extra of DataSample")
    end
    return DataSample(x, θ, y, context, extra)
end

"""
$TYPEDSIGNATURES

Extended property access for `DataSample`.

Allows accessing `context` and `extra` fields directly as properties.
`context` is searched first; if the key is not found there, `extra` is searched.
"""
function Base.getproperty(d::DataSample, name::Symbol)
    if name in (:x, :θ, :y, :context, :extra)
        return getfield(d, name)
    else
        ctx = getfield(d, :context)
        if haskey(ctx, name)
            return getproperty(ctx, name)
        end
        return getproperty(getfield(d, :extra), name)
    end
end

"""
$TYPEDSIGNATURES

Return all property names of a `DataSample`, including both struct fields and forwarded
fields from `context` and `extra`.

This enables tab completion for all available properties.
"""
function Base.propertynames(d::DataSample, private::Bool=false)
    ctx_names = propertynames(getfield(d, :context), private)
    extra_names = propertynames(getfield(d, :extra), private)
    return (fieldnames(DataSample)..., ctx_names..., extra_names...)
end

"""
$TYPEDSIGNATURES

Display a `DataSample` with truncated array representations for readability.

Large arrays are automatically truncated with ellipsis (`...`), similar to standard Julia array printing.
"""
function Base.show(io::IO, d::DataSample)
    fields = String[]
    io_limited = IOContext(io, :limit => true, :compact => true)
    if !isnothing(d.x)
        x_str = sprint(show, d.x; context=io_limited)
        push!(fields, "x=$x_str")
    end
    if !isnothing(d.θ)
        θ_str = sprint(show, d.θ; context=io_limited)
        push!(fields, "θ_true=$θ_str")
    end
    if !isnothing(d.y)
        y_str = sprint(show, d.y; context=io_limited)
        push!(fields, "y_true=$y_str")
    end
    for (key, value) in pairs(d.context)
        value_str = sprint(show, value; context=io_limited)
        push!(fields, "$key=$value_str")
    end
    for (key, value) in pairs(d.extra)
        value_str = sprint(show, value; context=io_limited)
        push!(fields, "$key=$value_str")
    end
    return print(io, "DataSample(", join(fields, ", "), ")")
end

"""
$TYPEDSIGNATURES

Fit the given transform type (`ZScoreTransform` or `UnitRangeTransform`) on the dataset.
"""
function StatsBase.fit(transform_type, dataset::AbstractVector{<:DataSample}; kwargs...)
    x = hcat([d.x for d in dataset]...)
    return StatsBase.fit(transform_type, x; kwargs...)
end

"""
$TYPEDSIGNATURES

Transform the features in the dataset.
"""
function StatsBase.transform(t, dataset::AbstractVector{<:DataSample})
    return map(dataset) do d
        (; context, extra, x, θ, y) = d
        DataSample(StatsBase.transform(t, x), θ, y, context, extra)
    end
end

"""
$TYPEDSIGNATURES

Transform the features in the dataset, in place.
"""
function StatsBase.transform!(t, dataset::AbstractVector{<:DataSample})
    for d in dataset
        StatsBase.transform!(t, d.x)
    end
end

"""
$TYPEDSIGNATURES

Reconstruct the features in the dataset.
"""
function StatsBase.reconstruct(t, dataset::AbstractVector{<:DataSample})
    return map(dataset) do d
        (; context, extra, x, θ, y) = d
        DataSample(StatsBase.reconstruct(t, x), θ, y, context, extra)
    end
end

"""
$TYPEDSIGNATURES

Reconstruct the features in the dataset, in place.
"""
function StatsBase.reconstruct!(t, dataset::AbstractVector{<:DataSample})
    for d in dataset
        StatsBase.reconstruct!(t, d.x)
    end
end