module YAXTools
using Distributed
import ..YAXArrays: YAXdir
export freshworkermodule, passobj, @everywhereelsem, @loadOrGenerate, PickAxisArray, Window

struct PickAxisArray{T,N,AT<:AbstractArray,P,PERM}
    parent::AT
end

struct Window
    i::Int
    pre::Int
    after::Int
end


function PickAxisArray(parent, indmask, perm = nothing)
    f = findall(isequal(true), indmask)
    f2 = findall(isequal(Colon()), indmask)
    f3 = findall(i -> isa(i, Tuple{Int,Int}), indmask)
    o = sort([f; f2; f3])
    o = isempty(f2) ? o : replace(o, map(i -> i => Colon(), f2)...)
    o = isempty(f3) ? o : replace(o, map(i -> i => Window(i, indmask[i]...), f3)...)
    nsub = 0
    for i = 1:length(o)
        if o[i] isa Colon
            nsub += 1
        elseif o[i] isa Window
            o[i] = Window(o[i].i - nsub, o[i].pre, o[i].after)
        else
            o[i] = o[i] - nsub
        end
    end
    if perm !== nothing
        length(perm) != length(f2) + length(f3) && error("Not a valid permutation")
        perm = (perm...,)
    end
    PickAxisArray{eltype(parent),length(indmask),typeof(parent),(o...,),perm}(parent)
end
indmask(::PickAxisArray{<:Any,<:Any,<:Any,i}) where {i} = i
getind(i, j) = i[j]
getind(i, j::Colon) = j
function getind(i, j::Window)
    c = i[j.i]
    c-j.pre:c+j.after
end

permout(::PickAxisArray{<:Any,<:Any,<:Any,<:Any,P}, x) where {P} = permutedims(x, P)
permout(::PickAxisArray{<:Any,<:Any,<:Any,<:Any,nothing}, x) = x
function Base.view(p::PickAxisArray, i::Integer...)
    inew = map(j -> getind(i, j), indmask(p))
    r = permout(p, view(p.parent, inew...))
    r
end
function Base.getindex(p::PickAxisArray, i::Integer...)
    inew = map(j -> getind(i, j), indmask(p))
    permout(p, getindex(p.parent, inew...))
end
anycol(::Tuple{}) = false
anycol(t::Tuple) = anycol(first(t), Base.tail(t))
anycol(::Union{Colon,Window}, t::Tuple) = true
anycol(i, ::Tuple{}) = false
anycol(::Union{Colon,Window}, ::Tuple{}) = true
anycol(i, t::Tuple) = anycol(first(t), Base.tail(t))


ncol(t::Tuple) = ncol(first(t), Base.tail(t), 0)
ncol(::Union{Colon,Window}, t::Tuple, n) = ncol(first(t), Base.tail(t), n + 1)
ncol(i, ::Tuple{}, n) = n
ncol(::Union{Colon,Window}, ::Tuple{}, n) = n + 1
ncol(i, t::Tuple, n) = ncol(first(t), Base.tail(t), n)

function Base.eltype(p::PickAxisArray{T}) where {T}
    im = indmask(p)
    if anycol(im)
        Array{T,ncol(im)}
    else
        T
    end
end
Base.getindex(p::PickAxisArray, i::CartesianIndex) = p[i.I...]

"""
macro loadOrGenerate(x...,expression)
  
Takes a list of `variablename=>"Storage Name"` pairs. Checks if all datasets can be found
on disk and loads them. If not, the datasets will be regenerated by evaluating the given expression.
  
To force recalculation, call `YAXArrays.recalculate(true)` before evaluating the macro.
  
### Example
  
The following lines will check if cubes with the names "Filled" and "Normalized"
exist on disk, load them and assign the variable names `cube_filled` and `cube_norm`.
If the datasets to not exist on disk, they are generated and saved under the given names.
    
````julia
    @loadOrGenerate cube_filled=>"Filled" cube_norm=>"Normalized" begin
    cube_norm   = mapCube(normalize_TS,d)
end  
````
"""
macro loadOrGenerate(x...)
    code = x[end]
    x = x[1:end-1]
    x2 = map(x) do i
        isa(i, Symbol) ? (i, string(i)) :
        (i.head == :call && i.args[1] == :(=>)) ? (i.args[2], i.args[3]) :
        error("Wrong Argument type")
    end
    xnames = map(i -> i[2], x2)
    loadEx = map(x2) do i
        :($(i[1]) = loadcube($(i[2])))
    end
    loadEx = Expr(:block, loadEx...)
    saveEx = map(x2) do i
        :(savecube($(i[1]), $(i[2])))
    end
    saveEx = Expr(:block, saveEx...)
    rmEx = map(x2) do i
        :(rmcube($(i[2])))
    end
    rmEx = Expr(:block, rmEx...)
    esc(quote
        if !YAXArrays.recalculate() && all(i -> isdir(joinpath(YAXdir(), i)), $xnames)
            $loadEx
        else
            $rmEx
            $code
            $saveEx
        end
    end)
end

end
