# SPDX-License-Identifier: CECILL-2.1

"""
FastDataset(supervisions, recordings, partition)

Constructor for dataset represented as JSONL files (a.k.a. manifests).
"""
struct SpeechDataset <: MLUtils.AbstractDataContainer
    idxs::Vector{AbstractString}
    superivions::Dict{AbstractString, Supervision}
    recordings::Dict{AbstractString, Recording}
    partition::Symbol
end

"""
dataset(manifestroot, subset)

Load `SpeechDataset` from manifest files stored in `manifestroot`.
Partition is specified by `subset`, e.g. `:train`, `:test`.

Each item of the dataset is a nested tuple `((samples, sampling_rate), Supervision.data)`.

See also [`Supervision`](@ref).

# Examples
```julia-repl
julia> ds = dataset("./manifests", :train)
SpeechDataset(
    ...
)
julia> ds[1]
(
    (samples=[...], sampling_rate=16_000),
    Dict(
        "text" => "Supervision text here"
    )
)
```
"""
function dataset(manifestroot::AbstractString, subset)
    supervisions = load(Supervision, manifestroot, subset)
    recordings = load(Recording, manifestroot, subset)
    idxs = collect(keys(supervisions))
    SpeechDataset(idxs, supervisions, recordings, Symbol(subset))
end

function Base.getindex(d::SpeechDataset, key::AbstractString)
    sup = d.superivions[key]
    rec = d.recordings[sup.recording_id]
    samples, sr = load(rec, sup)
    (samples=samples, sampling_rate=sr), sup.data
end
Base.getindex(d::SpeechDataset, idx::Integer) = getindex(d, d.idxs[idx])
# Fix1 -> partial funcion with fixed 1st argument
Base.getindex(d::SpeechDataset, idxs::AbstractVector) = map(Base.Fix1(getindex, d), idxs)

Base.length(d::SpeechDataset) = length(d.idxs)