export Lazy, eachdim′, eachcol′, eachrow′, eachslice′

### import many funtions
import Base.Broadcast: throwdm, AbstractArrayStyle, Unknown, combine_eltypes,
    preprocess, Broadcasted, DefaultArrayStyle, ischunkedbroadcast, chunkedcopyto!, bitcache_size,
    dumpbitcache, bitcache_chunks, materialize!, BroadcastStyle, combine_styles, instantiate,
    broadcastable, broadcast_unalias, Style, _broadcast_getindex, broadcasted
import Base: size, axes, setindex!, unalias, mightalias, unaliascopy, IndexStyle, parent,
    unsafe_setindex!, copyto!, IndexStyle, tail

# getsame
@noinline _getsamethrow(@nospecialize(f)) =
    throw(ArgumentError("Inputs with different $f"))
@inline getsame(f::F, x) where {F} = f(x)
@inline getsame(f::F, x, y, zs...) where {F} = begin
    fx = getsame(f, y, zs...)
    f(x) == fx || _getsamethrow(f)
    fx
end

const FilledBC = Broadcasted{<:AbstractArrayStyle{0}}
module LazyCollect
import Base: @propagate_inbounds, getindex
import Base.Broadcast: BroadcastStyle, broadcastable, extrude, newindex, broadcasted,
    instantiate, Broadcasted, DefaultArrayStyle, combine_styles
import Base.Iterators: ProductIterator

export Lazy

ndims(x::Tuple) = 1
size(x::Tuple) = (length(x),)
ndims(x) = Base.ndims(x)
size(x) = Base.size(x)
# Like reshape, but allow tuple inputs.
struct FakeDim{N,S,T,D} <: AbstractArray{T,N}
    data::D
    FakeDim{N,S}(data::T) where {N,S,T} = new{N,S,eltype(data),T}(data)
end
FakeDim{N,S}(data::FakeDim) where {N,S} = FakeDim{N,S}(data.data)
Base.axes(id::FakeDim{N,S}) where {N,S} = (ntuple(Returns(Base.OneTo(1)), Val(S - 1))..., axes(id.data)...)
Base.size(id::FakeDim{N,S}) where {N,S} = (ntuple(Returns(1), Val(S - 1))..., size(id.data)...)
extrude(x::FakeDim) = x
struct FakeIndex{N}
    inds::NTuple{N,Int}
    FakeIndex(inds::NTuple{N,Int}) where {N} = new{N}(inds)
end
@inline pick_inds(I::NTuple{N,Int}, ::Val{S}, ::Val{E}) where {N,S,E} = ntuple(i -> I[S + i - 1], Val(E - S + 1))
@inline newindex(::FakeDim{N,S}, I::CartesianIndex) where {N,S} = FakeIndex(pick_inds(I.I, Val(S), Val(N)))
@propagate_inbounds getindex(id::FakeDim, I::FakeIndex{N}) where {N} = id.data[I.inds...]
@propagate_inbounds getindex(id::FakeDim{N,S}, I::Vararg{Int, N}) where {N,S} = id.data[I[S:N]...]
@inline BroadcastStyle(::Type{ID}) where {ID<:FakeDim} =
    ID.parameters[4] <: Tuple ? DefaultArrayStyle{ID.parameters[1]}() :
                                BroadcastStyle(ID.parameters[4])

@inline fakedim(pre_ax::Tuple, data) = begin
    data′ = broadcastable(data)
    ax = axes(data)
    m, n = length(pre_ax), length(ax)
    data″ = m == 0 || n == 0 ? data′ : FakeDim{m + n, m + 1}(data′)
    data″, (pre_ax..., ax...)
end
@inline fakedim(pre_ax::Tuple, g::Base.Generator) = begin
    iter, ax = fakedim(pre_ax, g.iter)
    Broadcasted(g.f, (iter,), ax), ax
end
@inline fakedim(pre_ax::Tuple, i::Iterators.ProductIterator) = begin
    iter, ax = fakedims(pre_ax, i.iterators...)
    Broadcasted(tuple, iter, ax), ax
end
@inline fakedims(pre_ax::Tuple) = (), pre_ax
@inline fakedims(pre_ax::Tuple, x, y...) = begin 
    x′, pre_ax′ = fakedim(pre_ax, x) 
    y′, ax = fakedims(pre_ax′, y...)
    (x′, y′...), ax
end
# Lazy is used to wrap Generator to avoid collect before broadcast.
Lazy(x) = fakedim((), x)[1]

end
using .LazyCollect
## better each
@inline unsafe_view(A, I...) = Base.unsafe_view(A, to_indices(A, I)...)
eachcol′(A::AbstractVecOrMat) = Lazy(unsafe_view(A, :, i) for i in axes(A, 2))
eachrow′(A::AbstractVecOrMat) = Lazy(unsafe_view(A, i, :) for i in axes(A, 1))
function eachdim′(A::AbstractArray; dim::Val{D}) where {D}
    D <= ndims(A) || throw(DimensionMismatch("A doesn't have $dim dimensions"))
    axes_all = ntuple(d -> d == D ? Ref(:) : axes(A, d), ndims(A))
    Lazy(unsafe_view(A, i...) for i in Iterators.product(axes_all...))
end
function eachslice′(A::AbstractArray; dim::Val{D}) where {D}
    D <= ndims(A) || throw(DimensionMismatch("A doesn't have $dim dimensions"))
    inds_before = ntuple(Returns(:), D - 1)
    inds_after = ntuple(Returns(:), ndims(A) - D)
    Lazy(unsafe_view(A, inds_before..., i, inds_after...) for i in axes(A, D))
end

## faster Base.read!
function Base.read!(s::IO, a::StridedArray{T}) where {T}
    _style(a) = ndims(a) < 2 ? IndexLinear() : IndexStyle(a)
    if isbitstype(T) && _style(a) isa Union{IndexLinear, Base.IndexSCartesian2} && isone(stride(a,1))
        GC.@preserve a unsafe_read(s, pointer(a), sizeof(a))
        return a
    end
    Base.@invoke read!(s::IO, a::AbstractArray)
end

function Base.copyto_unaliased!(::IndexLinear, dest::AbstractArray, ::IndexLinear, src::AbstractArray)
    isempty(src) && return dest
    length(dest) < length(src) && throw(BoundsError(dest, LinearIndices(src)))
    Δi = firstindex(dest) - firstindex(src)
    for i in eachindex(IndexLinear(), src)
        @inbounds dest[i + Δi] = src[i]
    end
    return dest
end

splitloop(iter::CartesianIndices) = begin
    indices = iter.indices
    CartesianIndices(tail(indices)), first(indices)
end

function Base.copyto_unaliased!(::IndexLinear, dest::AbstractArray, ::IndexCartesian, src::AbstractArray)
    isempty(src) && return dest
    length(dest) < length(src) && throw(BoundsError(dest, LinearIndices(src)))
    iter, j = CartesianIndices(src), firstindex(dest) - 1
    inner_len = size(iter, 1)
    if inner_len >= 16
        # manually expand the inner loop similar to @simd
        outer, inner = splitloop(iter)
        @inbounds for II in outer
            for i in inner
                dest[j += 1] = src[i, II]
            end
        end
    else
        @inbounds for I in iter
            dest[j += 1] = src[I]
        end
    end
    return dest
end

function Base.copyto_unaliased!(::IndexCartesian, dest::AbstractArray, ::IndexLinear, src::AbstractArray)
    isempty(src) && return dest
    length(dest) < length(src) && throw(BoundsError(dest, LinearIndices(src)))
    iter, i = CartesianIndices(dest), firstindex(src) - 1
    inner_len = size(iter, 1)
    if inner_len >= 16
        # manually expand the inner loop similar to @simd
        outer, inner = splitloop(iter)
        if length(dest) == length(src)
            @inbounds for II in outer
                for j in inner
                    dest[j, II] = src[i += 1]
                end
            end
        else
            final = lastindex(src)
            @inbounds for II in outer
                if i + inner_len >= final
                    j = first(inner)
                    while i < final
                        dest[j, II] = src[i += 1]
                        j += step(inner)
                    end
                    break
                end
                for j in inner
                    dest[j, II] = src[i += 1]
                end
            end
        end
    elseif length(dest) == length(src)
        @inbounds for I in iter
            dest[I] = src[i += 1]
        end
    else
        # use zip based interator
        @inbounds for (I, J) in zip(eachindex(src), iter)
            dest[J] = src[I]
        end
    end
    return dest
end

function Base.copyto_unaliased!(::IndexCartesian, dest::AbstractArray, ::IndexCartesian, src::AbstractArray)
    isempty(src) && return dest
    length(dest) < length(src) && throw(BoundsError(dest, LinearIndices(src)))
    iterdest, itersrc = CartesianIndices(dest), CartesianIndices(src)
    if iterdest == itersrc
        iter = itersrc
        inner_len = size(iter, 1)
        if inner_len >= 16
            # manually expand the inner loop similar to @simd
            outer, inner = splitloop(iter)
            @inbounds for II in outer
                for i in inner
                    dest[i, II] = src[i, II]
                end
            end
        else
            @inbounds for I in iter
                dest[I] = src[I]
            end
        end
    else
        @inbounds for (J, I) in zip(iterdest, itersrc)
            dest[J] = src[I]
        end
    end
    return dest
end
