# Manage and access multinetwork dimensions


## Dimension data structure

function _initialize_dim()
    Dict{Symbol,Any}(
        :pos    => NamedTuple(), # Position (order) of each dimension
        :prop   => Dict{Symbol,Dict{Int,Dict{String,Any}}}(), # Data relating to the elements of the dimensions
        :meta   => Dict{Symbol,Any}(), # Data relating to dimensions
        :offset => 0 # Offset of nw ids: the id of the first nw is offset+1
        # li (linear indices): generated by add_dimension!()
        # ci (cartesian indices): generated by add_dimension!()
    )
end

"""
    add_dimension!(data, name, properties; metadata)

Add dimension `name` to `data` specifying some `properties` of the dimension ids.

# Arguments
- `data::Dict{String,Any}`: a single-network data dictionary.
- `name::Symbol`: the name to use to refer to the dimension in the code.
- `properties::Dict{Int,Dict{String,Any}}`: properties associated to the dimension ids. The
  outer dictionary is indexed with the ids along the dimension (consecutive `Int`s starting
  from 1). The inner dictionaries, one for each id, store the properties.
- `metadata::Dict{String,Any} = Dict{String,Any}()`: optional metadata describing the
  dimension as a whole.

# Examples
```julia-repl
julia> add_dimension!(data, :scenario, Dict(s => Dict{String,Any}("probability"=>1/4) for s in 1:4))
```

# Extended help
The functions `dim_prop`, `dim_meta` and `dim_length` can be used to access properties,
metadata and length (cardinality) of a dimension. They apply to both data dictionaries
(`Dict{String,Any}`) and powermodels (`PowerModels.AbstractPowerModel`).
"""
function add_dimension!(data::Dict{String,Any}, name::Symbol, properties::Dict{Int,Dict{String,Any}}; metadata::Dict{String,Any}=Dict{String,Any}())
    dim = get!(data, "dim", _initialize_dim())

    if haskey(dim[:pos], name)
        Memento.error(_LOGGER, "A dimension named \"$name\" is already present in data.")
    end
    if Set(keys(properties)) != Set(1:length(properties))
        Memento.error(_LOGGER, "Keys of `properties` Dict must range from 1 to the number of $(name)s.")
    end

    dim[:pos] = (; dim[:pos]..., name => length(dim[:pos])+1)
    dim[:prop][name] = properties
    dim[:meta][name] = metadata
    dim[:li] = LinearIndices(Tuple(1:length(dim[:prop][nm]) for nm in keys(dim[:pos]))).+dim[:offset]
    dim[:ci] = CartesianIndices(dim[:li])

    return dim
end

"""
    add_dimension!(data, name, size; metadata)

Add dimension `name` to `data` specifying the dimension `size`.

# Examples
```julia-repl
julia> add_dimension!(data, :hour, 24)
```
"""
function add_dimension!(data::Dict{String,Any}, name::Symbol, size::Int; metadata::Dict{String,Any}=Dict{String,Any}())
    properties = Dict{Int,Dict{String,Any}}(i => Dict{String,Any}() for i in 1:size)
    add_dimension!(data, name, properties; metadata)
end

"""
    shift_ids!(dim::Dict{Symbol,Any}, offset)
    shift_ids!(data::Dict{String,Any}, offset)

Shift by `offset` the network ids in `dim` or `data`.

The `offset` argument is added to the existing offset.
Return a vector containing the shifted network ids.
`data` must be a single-network `data` dictionary.
"""
function shift_ids! end

function shift_ids!(dim::Dict{Symbol,Any}, offset::Int)
    dim[:offset] += offset
    dim[:li] .+= offset
    return vec(dim[:li])
end

function shift_ids!(data::Dict{String,Any}, offset::Int)
    if _IM.ismultinetwork(data)
        Memento.error(_LOGGER, "`shift_ids!` can only be applied to single-network data dictionaries.")
    end
    shift_ids!(dim(data), offset)
end

"""
    merge_dim!(dim1, dim2, dimension)

Merge `dim1` and `dim2` structures along `dimension`.

The ids of `dim2` must be contiguous to those of `dim1`.
"""
function merge_dim!(dim1::Dict{Symbol,Any}, dim2::Dict{Symbol,Any}, dimension::Symbol)
    dim = Dict{Symbol,Any}()

    if dim1[:pos] != dim2[:pos]
        Memento.error(_LOGGER, "The dimensions to be merged have different names and/or order:\nfirst:  $(dim1[:pos])\nsecond: $(dim1[:pos])")
    end
    dim[:pos] = dim1[:pos]

    if any(dim1[:prop][d] != dim2[:prop][d] for d in delete!(Set(keys(dim[:pos])), dimension))
        diff = join(", $d" for d in delete!(Set(keys(dim[:pos])), dimension) if dim1[:prop][d] != dim2[:prop][d])[2:end]
        Memento.error(_LOGGER, "Different properties found in the following dimension(s):$diff.")
    end
    dim[:prop] = dim1[:prop]
    offset = length(dim1[:prop][dimension])
    for (k,v) in dim2[:prop][dimension]
        dim[:prop][dimension][k+offset] = deepcopy(v)
    end

    if dim1[:meta] != dim2[:meta]
        diff = join(", $d" for d in keys(dim[:pos]) if dim1[:meta][d] != dim2[:meta][d])[2:end]
        Memento.error(_LOGGER, "Different metadata found in the following dimension(s):$diff.")
    end
    dim[:meta] = dim1[:meta]

    if dim2[:li][1] != dim1[:li][end]+1
        Memento.error(_LOGGER, "Multinetworks to be merged must have contiguous ids.")
    end
    dim[:offset] = min(dim1[:offset], dim2[:offset])
    dim[:li] = LinearIndices(Tuple(1:length(dim[:prop][nm]) for nm in keys(dim[:pos]))).+dim[:offset]
    dim[:ci] = CartesianIndices(dim[:li])

    return dim
end

"""
    slice, ids = slice_dim(dim::Dict{Symbol,Any}; kwargs...)

Slice `dim` structure keeping the networks that have the coordinates specified by `kwargs`.

`kwargs` must be of the form `name = <value>`, where `name` is the name of a dimension of
`dim` and `<value>` is an `Int` coordinate of that dimension.

Return `slice`, a sliced `dim` structure whose networks have ids starting at 1, and `ids`, a
vector containing the ids that the networks making up `slice` have in `dim`.
The coordinates of the dimensions at which `dim` is sliced are accessible with
`dim_meta(slice, <name>, "orig_id")`, where `<name>` is the name of one of those dimensions.

# Examples
```julia-repl
julia> slice_dim(dim; hour = 24)
julia> slice_dim(dim; hour = 24, scenario = 3)
```
"""
function slice_dim(dim::Dict{Symbol,Any}; kwargs...)
    slice = Dict{Symbol,Any}()
    slice[:pos] = dim[:pos]
    slice[:prop] = Dict{Symbol,Dict{Int,Dict{String,Any}}}()
    for d in keys(dim[:pos])
        if d ∈ keys(kwargs)
            slice[:prop][d] = Dict{Int,Dict{String,Any}}(1 => deepcopy(dim[:prop][d][kwargs[d]]))
        else
            slice[:prop][d] = deepcopy(dim[:prop][d])
        end
    end
    slice[:meta] = deepcopy(dim[:meta])
    for (d, i) in kwargs
        slice[:meta][d]["orig_id"] = i
    end
    slice[:offset] = 0
    slice[:li] = collect(LinearIndices(Tuple(1:length(slice[:prop][nm]) for nm in keys(slice[:pos]))))
    slice[:ci] = CartesianIndices(slice[:li])

    names = keys(dim[:pos])
    li = dim[:li]
    ids = vec(li[ntuple(i -> get(kwargs, names[i], axes(li,i)), ndims(li))...])

    return slice, ids
end

"""
    require_dim(data, dimensions...)

Verify that the specified `dimensions` are present in `data`; if not, raise an error.
"""
function require_dim(data::Dict{String,Any}, dimensions::Symbol...)
    if !haskey(data, "dim")
        Memento.error(_LOGGER, "Missing `dim` dict in `data`. Use `add_dimension!` to fix.")
    end
    for d in dimensions
        if !haskey(dim(data)[:prop], d)
            Memento.error(_LOGGER, "Missing dimension \"$d\" in `data`. Use `add_dimension!` to fix.")
        end
    end
end


## Access (subsets of) nw ids

"""
    nw_ids(pm::PowerModels.AbstractPowerModel; kwargs...)
    nw_ids(data::Dict{String,Any}; kwargs...)
    nw_ids(dim::Dict{Symbol,Any}; kwargs...)

Sorted vector containing nw ids of `pm`, `data`, or `dim`, optionally filtered by the coordinates of one or more dimensions.

`kwargs` must be of the form `name = <value>` or `name = <interval>` or `name = <subset>`,
where `name` is the name of a dimension of `pm`.

# Examples
```julia-repl
julia> nw_ids(pm)
julia> nw_ids(pm; hour = 24)
julia> nw_ids(pm; hour = 13:24)
julia> nw_ids(pm; hour = [6,12,18,24])
julia> nw_ids(pm; hour = 24, scenario = 3)
```
"""
function nw_ids end

function nw_ids(dim::Dict{Symbol,Any}; kwargs...)::Vector{Int}
    names = keys(dim[:pos])
    li = dim[:li]
    nws = li[ntuple(i -> get(kwargs, names[i], axes(li, i)), ndims(li))...]
    ndims(nws) >= 1 ? vec(nws) : [nws]
end

function nw_ids(data::Dict{String,Any}; kwargs...)::Vector{Int}
    haskey(data, "dim") ? nw_ids(dim(data); kwargs...) : [0]
end

function nw_ids(pm::_PM.AbstractPowerModel; kwargs...)::Vector{Int}
    haskey(pm.ref[:it][_PM.pm_it_sym], :dim) ? nw_ids(dim(pm); kwargs...) : [0]
end


## Compute nw ids given another nw id

"""
    similar_ids(pm::PowerModels.AbstractPowerModel, n::Int; kwargs...)
    similar_ids(data::Dict{String,Any}, n::Int; kwargs...)
    similar_ids(dim::Dict{Symbol,Any}, n::Int; kwargs...)

Sorted vector containing nw ids that have the same coordinates along dimensions as `n` except for dimensions passed in `kwargs`.

`kwargs` must be of the form `name = <value>` or `name = <interval>` or `name = <subset>`,
where `name` is the name of a dimension of `pm`, `data` or `dim`.

# Examples
```julia-repl
julia> similar_ids(pm, 3; hour = 24)
julia> similar_ids(pm, 3; hour = 13:24)
julia> similar_ids(pm, 3; hour = [6,12,18,24])
julia> similar_ids(pm, 3; hour = 24, scenario = 3)
```
"""
function similar_ids end

function similar_ids(dim::Dict{Symbol,Any}, n::Int; kwargs...)::Vector{Int}
    names = keys(dim[:pos])
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    nws = li[ntuple(i -> get(kwargs, names[i], ci_n[i]), ndims(li))...]
    ndims(nws) >= 1 ? vec(nws) : [nws]
end

similar_ids(data::Dict{String,Any}, n::Int; kwargs...) = similar_ids(dim(data), n; kwargs...)
similar_ids(pm::_PM.AbstractPowerModel, n::Int; kwargs...) = similar_ids(dim(pm), n; kwargs...)

"""
    similar_id(pm::PowerModels.AbstractPowerModel, n::Int; kwargs...)
    similar_id(data::Dict{String,Any}, n::Int; kwargs...)
    similar_id(dim::Dict{Symbol,Any}, n::Int; kwargs...)

Nw id that has the same coordinates along dimensions as `n` except for dimensions passed in `kwargs`.

`kwargs` must be of the form `name = <value>`, where `name` is the name of a dimension of
`pm`, `data` or `dim`.

# Examples
```julia-repl
julia> similar_id(pm, 3; hour = 24)
julia> similar_id(pm, 3; hour = 24, scenario = 3)
```
"""
function similar_id end

function similar_id(dim::Dict{Symbol,Any}, n::Int; kwargs...)::Int
    names = keys(dim[:pos])
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    li[ntuple(i -> get(kwargs, names[i], ci_n[i])::Int, ndims(li))...]
end

similar_id(data::Dict{String,Any}, n::Int; kwargs...) = similar_id(dim(data), n; kwargs...)
similar_id(pm::_PM.AbstractPowerModel, n::Int; kwargs...) = similar_id(dim(pm), n; kwargs...)

"""
    first_id(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol...)
    first_id(data::Dict{String,Any}, n::Int, dimension::Symbol...)
    first_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)

Return the first network in `pm`, `data` or `dim` along `dimension` while keeping the other dimensions fixed.
"""
function first_id end

function first_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)
    names = keys(dim[:pos])
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    li[ntuple(i -> names[i] in dimension ? 1 : ci_n[i], ndims(li))...]
end

first_id(data::Dict{String,Any}, n::Int, dimension::Symbol...) = first_id(dim(data), n, dimension...)
first_id(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol...) = first_id(dim(pm), n, dimension...)

"""
    last_id(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol...)
    last_id(data::Dict{String,Any}, n::Int, dimension::Symbol...)
    last_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)

Return the last network in `pm`, `data` or `dim` along `dimension` while keeping the other dimensions fixed.
"""
function last_id end

function last_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)
    names = keys(dim[:pos])
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    li[ntuple(i -> names[i] in dimension ? size(li,i) : ci_n[i], ndims(li))...]
end

last_id(data::Dict{String,Any}, n::Int, dimension::Symbol...) = last_id(dim(data), n, dimension...)
last_id(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol...) = last_id(dim(pm), n, dimension...)

"""
    prev_id(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol)
    prev_id(data::Dict{String,Any}, n::Int, dimension::Symbol)
    prev_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)

Return the previous network in `pm`, `data` or `dim` along `dimension` while keeping the other dimensions fixed.
"""
function prev_id end

function prev_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)::Int
    pos_d = dim[:pos][dimension]
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    li[ntuple(i -> i == pos_d ? ci_n[i]-1 : ci_n[i], ndims(li))...]
end

prev_id(data::Dict{String,Any}, n::Int, dimension::Symbol) = prev_id(dim(data), n, dimension)
prev_id(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol) = prev_id(dim(pm), n, dimension)

"""
    prev_ids(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol)
    prev_ids(data::Dict{String,Any}, n::Int, dimension::Symbol)
    prev_ids(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)

Return the previous networks in `pm`, `data` or `dim` along `dimension` while keeping the other dimensions fixed.
"""
function prev_ids end

function prev_ids(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)::Vector{Int}
    pos_d = dim[:pos][dimension]
    offset = dim[:offset]
    li = dim[:li]
    ci_n = Tuple(dim[:ci][n-offset])
    li[CartesianIndex(ci_n[1:pos_d-1]), CartesianIndices((1:ci_n[pos_d]-1,)), CartesianIndex(ci_n[pos_d+1:end])]
end

prev_ids(data::Dict{String,Any}, n::Int, dimension::Symbol) = prev_ids(dim(data), n, dimension)
prev_ids(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol) = prev_ids(dim(pm), n, dimension)

"""
    next_id(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol)
    next_id(data::Dict{String,Any}, n::Int, dimension::Symbol)
    next_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)

Return the next network in `pm`, `data` or `dim` along `dimension` while keeping the other dimensions fixed.
"""
function next_id end

function next_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)::Int
    pos_d = dim[:pos][dimension]
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    li[ntuple(i -> i == pos_d ? ci_n[i]+1 : ci_n[i], ndims(li))...]
end

next_id(data::Dict{String,Any}, n::Int, dimension::Symbol) = next_id(dim(data), n, dimension)
next_id(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol) = next_id(dim(pm), n, dimension)

"""
    next_ids(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol)
    next_ids(data::Dict{String,Any}, n::Int, dimension::Symbol)
    next_ids(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)

Return the next networks in `pm`, `data` or `dim` along `dimension` while keeping the other dimensions fixed.
"""
function next_ids end

function next_ids(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)::Vector{Int}
    pos_d = dim[:pos][dimension]
    offset = dim[:offset]
    li = dim[:li]
    ci_n = Tuple(dim[:ci][n-offset])
    li[CartesianIndex(ci_n[1:pos_d-1]), CartesianIndices((ci_n[pos_d]+1:size(li,pos_d),)), CartesianIndex(ci_n[pos_d+1:end])]
end

next_ids(data::Dict{String,Any}, n::Int, dimension::Symbol) = next_ids(dim(data), n, dimension)
next_ids(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol) = next_ids(dim(pm), n, dimension)


## Query properties of nw ids

"""
    coord(pm::_PM.AbstractPowerModel, n::Int, dimension::Symbol)
    coord(data::Dict{String,Any}, n::Int, dimension::Symbol)
    coord(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)

Return the coordinate along `dimension` of nw `n` of `pm`, `data` or `dim`.
"""
function coord end

function coord(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol)
    pos = dim[:pos]
    offset = dim[:offset]
    ci_n = dim[:ci][n-offset]
    ci_n[pos[dimension]]
end

coord(data::Dict{String,Any}, args...) = coord(dim(data), args...)
coord(pm::_PM.AbstractPowerModel, args...) = coord(dim(pm), args...)

"""
    is_first_id(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol...)
    is_first_id(data::Dict{String,Any}, n::Int, dimension::Symbol...)
    is_first_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)

Return whether the network `n` is the first along `dimension` in `pm`, `data` or `dim`.
"""
function is_first_id end

function is_first_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)
    pos = dim[:pos]
    offset = dim[:offset]
    ci_n = dim[:ci][n-offset]
    all(ci_n[pos[d]] == 1 for d in dimension)
end

is_first_id(data::Dict{String,Any}, args...) = is_first_id(dim(data), args...)
is_first_id(pm::_PM.AbstractPowerModel, args...) = is_first_id(dim(pm), args...)

"""
    is_last_id(pm::PowerModels.AbstractPowerModel, n::Int, dimension::Symbol...)
    is_last_id(data::Dict{String,Any}, n::Int, dimension::Symbol...)
    is_last_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)

Return whether the network `n` is the last along `dimension` in `pm`, `data` or `dim`.
"""
function is_last_id(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol...)
    pos = dim[:pos]
    offset = dim[:offset]
    li = dim[:li]
    ci_n = dim[:ci][n-offset]
    all(ci_n[pos[d]] == size(li,pos[d]) for d in dimension)
end

is_last_id(data::Dict{String,Any}, args...) = is_last_id(dim(data), args...)
is_last_id(pm::_PM.AbstractPowerModel, args...) = is_last_id(dim(pm), args...)


## Access data relating to dimensions

"""
    dim(data::Dict{String,Any})
    dim(pm::PowerModels.AbstractPowerModel)

Return `dim` data structure.
"""
function dim end
dim(data::Dict{String,Any}) = data["dim"]
dim(pm::_PM.AbstractPowerModel) = pm.ref[:it][_PM.pm_it_sym][:dim]

"""
    dim_prop(dim::Dict{Symbol,Any}[, dimension[, id[, key]]])
    dim_prop(data::Dict{String,Any}[, dimension[, id[, key]]])
    dim_prop(pm::PowerModels.AbstractPowerModel[, dimension[, id[, key]]])

Properties associated to the `id`s of `dimension`.

    dim_prop(dim::Dict{Symbol,Any}, n, dimension[, key])
    dim_prop(data::Dict{String,Any}, n, dimension[, key])
    dim_prop(pm::PowerModels.AbstractPowerModel, n, dimension[, key])

Properties associated to `dimension` of a network `n`.
"""
function dim_prop end
dim_prop(dim::Dict{Symbol,Any}) = dim[:prop]
dim_prop(dim::Dict{Symbol,Any}, dimension::Symbol) = dim[:prop][dimension]
dim_prop(dim::Dict{Symbol,Any}, dimension::Symbol, id::Int) = dim[:prop][dimension][id]
dim_prop(dim::Dict{Symbol,Any}, dimension::Symbol, id::Int, key::String) = dim[:prop][dimension][id][key]

dim_prop(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol) = dim[:prop][dimension][coord(dim,n,dimension)]
dim_prop(dim::Dict{Symbol,Any}, n::Int, dimension::Symbol, key::String) = dim[:prop][dimension][coord(dim,n,dimension)][key]

dim_prop(data::Dict{String,Any}, args...) = dim_prop(dim(data), args...)
dim_prop(pm::_PM.AbstractPowerModel, args...) = dim_prop(dim(pm), args...)

"""
    dim_meta(dim::Dict{Symbol,Any}[, dimension[, key]])
    dim_meta(data::Dict{String,Any}[, dimension[, key]])
    dim_meta(pm::PowerModels.AbstractPowerModel[, dimension[, key]])

Metadata associated to `dimension`.
"""
function dim_meta end
dim_meta(dim::Dict{Symbol,Any}) = dim[:meta]
dim_meta(dim::Dict{Symbol,Any}, dimension::Symbol) = dim[:meta][dimension]
dim_meta(dim::Dict{Symbol,Any}, dimension::Symbol, key::String) = dim[:meta][dimension][key]
dim_meta(data::Dict{String,Any}, args...) = dim_meta(dim(data), args...)
dim_meta(pm::_PM.AbstractPowerModel, args...) = dim_meta(dim(pm), args...)

"""
    dim_length(dim::Dict{Symbol,Any}[, dimension])
    dim_length(data::Dict{String,Any}[, dimension])
    dim_length(pm::PowerModels.AbstractPowerModel[, dimension])

Return the number of networks or, if `dimension` is specified, return its size.
"""
function dim_length end
dim_length(dim::Dict{Symbol,Any}) = length(dim[:li])
dim_length(dim::Dict{Symbol,Any}, dimension::Symbol) = size(dim[:li], dim[:pos][dimension])
dim_length(data::Dict{String,Any}, args...) = dim_length(dim(data), args...)
dim_length(pm::_PM.AbstractPowerModel, args...) = dim_length(dim(pm), args...)
