module FastBroadcast

using Base: Fix1
using Base.Broadcast: Broadcasted, materialize, materialize!

using ArrayInterface: indices_do_not_alias, flatten_tuples
using LinearAlgebra: Adjoint, Transpose

export @.., Threaded, Serial

"""
    Threaded()

Threading dispatch type for `@..`. When passed as `thread = Threaded()`,
enables Polyester-based multithreaded broadcasting. Using a type instead of
`true` allows compile-time dispatch and prevents invalidations when
Polyester is loaded (the `Serial` code path never references threaded methods).

# Example
```julia
using FastBroadcast, Polyester

destination = zeros(3)
source = [1, 2, 3]
@.. thread = Threaded() destination = source + 1
```
"""
struct Threaded end

"""
    Serial()

Threading dispatch type for `@..`. When passed as `thread = Serial()`,
uses serial broadcasting (the default). Using a type instead of `false`
allows compile-time dispatch.

# Example
```julia
using FastBroadcast

destination = zeros(3)
source = [1, 2, 3]
@.. thread = Serial() destination = source + 1
```
"""
struct Serial end

@inline function to_tup(::Val{M}, i::CartesianIndex{N}) where {M, N}
    return if M < N
        ntuple(Fix1(getindex, Tuple(i)), Val(M))
    elseif M == N
        Tuple(i)
    else # M > N
        (Tuple(i)..., ntuple(one, Val(M - N))...)
    end
end

@inline _rall(x::Tuple{}) = true
@inline _rall(x::Tuple{Vararg{Bool}}) = first(x) && _rall(Base.tail(x))

@inline _rany(x::Tuple{}) = false
@inline _rany(x::Tuple{Vararg{Bool}}) = (first(x)) || _rany(Base.tail(x))

@inline _rmap(f, ::Tuple{}) = ()
@inline _rmap(f, x::Tuple{X}) where {X} = (@inline(f(only(x))),)
@inline _rmap(f, x::Tuple{X, Y, Vararg}) where {X, Y} = (
    @inline(f(first(x))), _rmap(f, Base.tail(x))...,
)

@inline _rmap(f, ::Tuple{}, ::Tuple{}) = ()
@inline _rmap(f, a::Tuple{A}, x::Tuple{X}) where {A, X} = (@inline(f(only(a), only(x))),)
@inline _rmap(f, a::Tuple{A, B, Vararg}, x::Tuple{X, Y, Vararg}) where {A, B, X, Y} = (
    @inline(f(first(a), first(x))), _rmap(f, Base.tail(a), Base.tail(x))...,
)
@inline _rmap(
    f,
    a::Tuple{A, B, Vararg},
    m::Tuple{M, N, Vararg},
    x::Tuple{X, Y, Vararg}
) where {A, B, M, N, X, Y} = (
    @inline(f(first(a), first(m), first(x))),
    _rmap(f, Base.tail(a), Base.tail(m), Base.tail(x))...,
)

struct SingularAxis
end
function axismatch(ax1, ax2)
    ax1 isa SingularAxis && return true
    ax2 isa SingularAxis && return true
    return length(ax1) == length(ax2)
end

const VecAdjTrans = Union{Transpose{T, M}, Adjoint{T, M}} where {T, M <: AbstractVector}
semi_static_axes(A) = Base.axes(A)
semi_static_axes(A::VecAdjTrans) = (SingularAxis(), only(semi_static_axes(parent(A))))
semi_static_axes(A::Tuple{T}) where {T} = (SingularAxis(),)
function semi_static_axes(S::SubArray)
    parent_axes = semi_static_axes(parent(S))
    return flatten_tuples(_rmap(singularize_singular_slice, parent_axes, S.indices))
end
function singularize_singular_slice(parent_ax, ax)
    return ax isa Integer ? () : ax isa Base.Slice ? parent_ax : ax
end

check_no_dyn_broadcast(::Union{Number, Base.RefValue}, ::Tuple{Vararg{Any, N}}) where {N} = true
@inline function check_no_dyn_broadcast(B, ax::Tuple{Vararg{Any, N}}) where {N}
    bx = semi_static_axes(B)
    MN = min(length(ax), length(bx))
    subax = ntuple(Fix1(Base.getindex, ax), Val(MN))
    subbx = ntuple(Fix1(Base.getindex, bx), Val(MN))
    return _rall(_rmap(axismatch, subax, subbx))
end

@inline function check_no_dyn_broadcast(bc::Broadcasted, ax::Tuple{Vararg{Any, N}}) where {N}
    return _rall(_rmap(check_no_dyn_broadcast, bc.args, ntuple(Returns(ax), Val(length(bc.args)))))
end

broadcast_islinear(::Union{Number, Base.RefValue}, ::Tuple{Vararg{Any, N}}) where {N} = false
function broadcast_islinear(::Tuple{Vararg{Any, M}}, ::Tuple{Vararg{Any, N}}) where {M, N}
    return M != 1
end

function broadcast_islinear(B, ax::Tuple{Vararg{Any, N}}) where {N}
    bx = semi_static_axes(B)
    return if (IndexStyle(typeof(B)) === IndexLinear()) && length(bx) == N
        _rany(_rmap(isequal(SingularAxis()), bx))
    else
        true
    end
end

function broadcast_islinear(bc::Broadcasted, ax::Tuple{Vararg{Any, N}}) where {N}
    return _rany(_rmap(broadcast_islinear, bc.args, ntuple(Returns(ax), Val(length(bc.args)))))
end

@inline _fastindex(b::Number, _) = b
@inline _fastindex(b::Tuple{X}, _) where {X} = only(b)
@inline _fastindex(b::Base.RefValue, _) = b[]
@inline _fastindex(b::Tuple{X, Y, Vararg}, i) where {X, Y} = @inbounds b[i]
@inline _fastindex(b::Broadcasted, i) = b.f(_rmap(_fastindex, b.args, ntuple(Returns(i), length(b.args)))...)
@inline function _fastindex(A, i)
    i isa Int && return @inbounds A[i]
    axs = semi_static_axes(A)
    inds = _rmap(axs, to_tup(Val(length(axs)), i)) do ax, j
        ax isa SingularAxis ? 1 : j
    end
    return @inbounds A[inds...]
end

@inline function __fast_materialize!(dst, ::Val{true}, bc::Broadcasted, ::Val{false})
    @simd ivdep for i in eachindex(dst)
        @inbounds dst[i] = _fastindex(bc, i)
    end
    return dst
end
@inline function __fast_materialize!(dst, ::Val{true}, bc::Broadcasted, ::Val{true})
    @simd ivdep for i in CartesianIndices(dst)
        @inbounds dst[i] = _fastindex(bc, i)
    end
    return dst
end
@inline function __fast_materialize!(dst, ::Val{false}, bc::Broadcasted, ::Val{false})
    for i in eachindex(dst)
        @inbounds dst[i] = _fastindex(bc, i)
    end
    return dst
end
@inline function __fast_materialize!(dst, ::Val{false}, bc::Broadcasted, ::Val{true})
    for i in CartesianIndices(dst)
        @inbounds dst[i] = _fastindex(bc, i)
    end
    return dst
end

@inline function _fast_materialize!(dst, ::Val{NOALIAS}, bc::Broadcasted) where {NOALIAS}
    dstax = semi_static_axes(dst)
    _no_dyn_broadcast = check_no_dyn_broadcast(bc, semi_static_axes(dst))
    @boundscheck _no_dyn_broadcast || throw(
        DimensionMismatch("Some axes are not equal, or feature a dynamic broadcast!"),
    )
    _islinear = broadcast_islinear(bc, dstax)
    __fast_materialize!(dst, Val(NOALIAS), bc, Val(_islinear))
    return dst
end

# these are needed for `RecursiveArrayTools`'s fast broadcast extension
use_fast_broadcast(_) = false
use_fast_broadcast(::Type{<:Base.Broadcast.DefaultArrayStyle}) = true
use_fast_broadcast(::Type{<:Base.Broadcast.DefaultArrayStyle{0}}) = false

# Legacy 4-arg fallback (old API compat)
fast_materialize!(_, _, dst, x) = dst .= x

# Serial dispatch — the primary path. Never references fast_materialize_threaded!,
# so loading Polyester cannot invalidate any code compiled via Serial.
@inline fast_materialize!(::Serial, dst, x::Number) = fill!(dst, x)
@inline function fast_materialize!(::Serial, dst, x::AbstractArray)
    Base.Broadcast.check_broadcast_shape(size(dst), size(x))
    return copyto!(dst, x)
end
@inline fast_materialize!(::Serial, dst, x) = dst .= x

Base.@propagate_inbounds function fast_materialize!(
        ::Serial, dst::A, bc::Broadcasted{S}
    ) where {S, A}
    return if S === Base.Broadcast.DefaultArrayStyle{0}
        fill!(dst, bc[1])
    elseif S <: Base.Broadcast.DefaultArrayStyle
        _fast_materialize!(dst, Val(indices_do_not_alias(A)), bc)
    else
        materialize!(dst, bc)
    end
end

Base.@propagate_inbounds function fast_materialize(
        ::Serial, bc::Broadcasted{S}
    ) where {S}
    if S === Base.Broadcast.DefaultArrayStyle{0}
        return only(bc)
    elseif S <: Base.Broadcast.DefaultArrayStyle
        fast_materialize!(
            Serial(),
            similar(bc, Base.Broadcast.combine_eltypes(bc.f, bc.args)),
            bc
        )
    else
        materialize(bc)
    end
end
fast_materialize(::Serial, @nospecialize(x)) = x

# Threaded dispatch — delegates to extension-defined methods for Broadcasted,
# falls back to serial for non-broadcast assignments (threading overhead not worth it)
@inline fast_materialize!(::Threaded, dst, bc::Broadcasted) = fast_materialize_threaded!(dst, bc)
@inline fast_materialize!(::Threaded, dst, x) = dst .= x
@inline fast_materialize(::Threaded, bc::Broadcasted) = fast_materialize_threaded(bc)
@inline fast_materialize(::Threaded, x) = x

# Bool backward compat: runtime branch
fast_materialize!(t::Bool, dst, bc) = t ? fast_materialize!(Threaded(), dst, bc) : fast_materialize!(Serial(), dst, bc)
fast_materialize(t::Bool, bc) = t ? fast_materialize(Threaded(), bc) : fast_materialize(Serial(), bc)


@inline _view(A::AbstractArray{<:Any, N}, r, ::Val{N}) where {N} = view(
    A, ntuple(_ -> :, N - 1)..., r
)
@inline _view(A::AbstractArray, r, ::Val) = A
@inline _view(x, r, ::Val) = x
@inline __view(t::Tuple{T}, r, ::Val{N}) where {T, N} = (_view(first(t), r, Val(N)),)
@inline __view(t::Tuple{T, Vararg}, r, ::Val{N}) where {T, N} = (
    _view(first(t), r, Val(N)), __view(Base.tail(t), r, Val(N))...,
)
@inline function _view(
        bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{N}, Nothing},
        r,
        ::Val{N}
    ) where {N}
    return Base.Broadcast.Broadcasted(bc.f, __view(bc.args, r, Val(N)))
end
@inline _view(bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle}, _, ::Val{N}) where {N} = bc
@inline _view(t::Tuple{Vararg{AbstractRange, N}}, r, ::Val{N}) where {N} = (Base.front(t)..., r)

"""
    fast_materialize_threaded(bc)

Threaded version of `fast_materialize`. Requires the Polyester.jl package to be loaded
(which activates the FastBroadcastPolyesterExt extension).
"""
function fast_materialize_threaded end

"""
    fast_materialize_threaded!(dst, bc)

Threaded in-place version of `fast_materialize!`. Requires the Polyester.jl package
to be loaded (which activates the FastBroadcastPolyesterExt extension).
"""
function fast_materialize_threaded! end

function __init__()
    return Base.Experimental.register_error_hint(MethodError) do io, exc, _argtypes, _kwargs
        if exc.f === fast_materialize_threaded || exc.f === fast_materialize_threaded!
            if Base.get_extension(@__MODULE__, :FastBroadcastPolyesterExt) === nothing
                print(
                    io,
                    "\n\n`@.. thread=true` (and `FastBroadcast.fast_materialize_threaded`/",
                    "`fast_materialize_threaded!`) requires the Polyester.jl package to be loaded. ",
                    "Run `using Polyester` (or `import Polyester`) to load the ",
                    "FastBroadcastPolyesterExt extension."
                )
            end
        end
    end
end


_dim0(_) = false
_dim0(::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}}) = true
@inline function _broadcasted(f::F, args...) where {F}
    bc = Base.Broadcast.broadcasted(f, args...)
    return _dim0(bc) ? only(bc) : bc
end
@inline _broadcasted(::Colon, arg0, arg1) = arg0:arg1
@inline _broadcasted(::typeof(Base.maybeview), args...) = begin
    Base.maybeview(args...)
end

function _view!(ex::Expr)
    r = Expr(:ref)
    r.args = ex.args
    ex.head = :macrocall
    # `maybeview` doesn't return a `Broadcasted` object
    ex.args = Any[Symbol("@views"), Base.LineNumberNode(@__LINE__, @__FILE__), r]
    return nothing
end

@inline function literal_pow(x, ::Val{N}) where {N}
    if N == 0
        return one(x)
    elseif N == 1
        return x
    elseif N < 0
        return inv(literal_pow(x, Val(-N)))
    end
    p = N
    tz = trailing_zeros(p) + 1
    p >>= tz
    for _ in 1:(tz - 1)
        @fastmath x *= x
    end
    y = x
    @fastmath while p != 0
        tz = trailing_zeros(p) + 1
        p >>= tz
        for _ in 1:tz
            x *= x
        end
        y *= x
    end
    return y
end


function _fb_macro!(ex::Expr, threadarg, broadcastarg)
    ops = (:(+), :(-), :(*), :(/), :(\), :(÷), :(&), :(|), :(⊻), :(>>), :(>>>), :(<<), :(^))
    if Meta.isexpr(ex, :(.)) && length(ex.args) == 2
        args = ex.args[2]
        args isa QuoteNode && return
        resize!(ex.args, 1)
        ex.head = :call
        append!(ex.args, args.args)
    elseif Meta.isexpr(ex, :macrocall, 3) && ex.args[1] === Symbol("@view")
        ex3 = ex.args[3]
        ex.head = ex3.head
        ex.args = ex3.args
    end
    skip = 0
    if Meta.isexpr(ex, :call)
        if Meta.isexpr(ex.args[1], :($))
            ex.args[1] = only(ex.args[1].args)
            return
        elseif ex.args[1] == :(:)
            return
        end
        pushfirst!(ex.args, _broadcasted)
        fun = ex.args[2]
        ind = Base.findfirst(
            ==(fun),
            (
                Symbol(".+"),
                Symbol(".-"),
                Symbol(".*"),
                Symbol("./"),
                Symbol(".\\"),
                Symbol(".÷"),
                Symbol(".&"),
                Symbol(".|"),
                Symbol(".⊻"),
                Symbol(".>>"),
                Symbol(".>>>"),
                Symbol(".<<"),
                Symbol(".^"),
            )
        )
        if ind !== nothing
            ex.args[2] = fun = ops[ind]
        end
        if fun === :(^) && length(ex.args) == 4 && (p = ex.args[4]; p isa Int)
            ex.args[2] = literal_pow
            ex.args[4] = Val(p)
        end
    elseif Meta.isexpr(ex, :(=))
        ex.head = :call
        if broadcastarg
            pushfirst!(ex.args, materialize!)
        elseif threadarg === true
            pushfirst!(ex.args, fast_materialize!)
            insert!(ex.args, 2, Threaded())
        elseif threadarg === false
            pushfirst!(ex.args, fast_materialize!)
            insert!(ex.args, 2, Serial())
        else
            # Runtime threading variable: dispatch via fast_materialize!(threadval, dst, bc)
            pushfirst!(ex.args, fast_materialize!)
            insert!(ex.args, 2, threadarg)
        end
        a4 = ex.args[end]
        if Meta.isexpr(a4, :ref)
            _view!(a4)
            # Also wrap a `:ref` LHS in `@views` so fancy/range indexed assignment
            # writes to the underlying array. Otherwise `temp[diff] = data[diff]`
            # passes a `getindex` copy as dst and silently drops the write.
            a3 = ex.args[end - 1]
            if Meta.isexpr(a3, :ref)
                _view!(a3)
            end
            skip = length(ex.args)
        end
    elseif Meta.isexpr(ex, :($), 1)
        if (exarg = only(ex.args); exarg isa Expr)
            ex.head = exarg.head
            ex.args = exarg.args
            return
        end
    elseif Meta.isexpr(ex, :ref)
        return _view!(ex)
    elseif Meta.isexpr(ex, :let)
        skip = 1
    else
        ind = Base.findfirst(
            ==(ex.head),
            (
                :(+=),
                :(-=),
                :(*=),
                :(/=),
                :(\=),
                :(÷=),
                :(&=),
                :(|=),
                :(⊻=),
                :(>>=),
                :(>>>=),
                :(<<=),
                :(^=),
            )
        )
        if ind !== nothing
            op = ops[ind]
            # x op= f(args...) -> x = op(x, f(args...))
            @assert length(ex.args) == 2
            ex.args[2] = Expr(:call, op, ex.args[1], ex.args[2])
            ex.head = :(=)
            return _fb_macro!(ex, threadarg, broadcastarg)
        end
    end
    for i in (1 + skip):length(ex.args)
        x = ex.args[i]
        x isa Expr && _fb_macro!(x, threadarg, broadcastarg)
    end
    return
end

function fb_macro!(ex::Expr, threadarg, broadcastarg)
    iscall = Meta.isexpr(ex, :call)
    _fb_macro!(ex, threadarg, broadcastarg)
    if iscall
        ex = Expr(:call, ex)
        if broadcastarg
            pushfirst!(ex.args, materialize)
        elseif threadarg === true
            pushfirst!(ex.args, fast_materialize)
            insert!(ex.args, 2, Threaded())
        elseif threadarg === false
            pushfirst!(ex.args, fast_materialize)
            insert!(ex.args, 2, Serial())
        else
            # Runtime threading variable: dispatch via fast_materialize(threadval, bc)
            pushfirst!(ex.args, fast_materialize)
            insert!(ex.args, 2, threadarg)
        end
    end
    return esc(ex)
end

"""
    @.. [thread = false] [broadcast = false] expr

`@..` turns `expr` into a broadcast-like expression, similar to `@.`.
It additionally provides two optional keyword arguments:

- thread: Defaults to `false`. Controls multithreading via Polyester.jl.
    Accepts `Serial()`, `Threaded()`, `false`, or `true`. Using the type-based
    `Serial()`/`Threaded()` is preferred as it enables compile-time dispatch and
    avoids invalidations when Polyester is loaded. Requires `using Polyester`
    for threaded execution.
    This only works if broadcast=false.
- broadcast: Defaults to `false`. If `true`, it will broadcast axes with dynamic
    runtime sizes of `1` to larger sizes, if `false` only sizes known to be `1`
    at compile time will be supported, i.e. axes such that
    `ArrayInterface.known_length(typeof(semi_static_axes(x,i))) == 1` will
    be broadcast. Note that this differs from base broadcasting, in that
    base broadcasting only supports `broadcast=true`.

# Example
```julia
using FastBroadcast

destination = zeros(3)
source = [1, 2, 3]
@.. destination = 2source + 1
```
"""
macro (..)(ex)
    ex isa Expr || return esc(ex)
    return fb_macro!(ex, false, false)
end

function __process_kwarg(kwarg)
    return kwarg.args[2]
end
function _validate_kwarg(kwarg)
    @assert Meta.isexpr(kwarg, :(=), 2)
    argname = kwarg.args[1]
    @assert (argname === :thread) || (argname === :broadcast)
    return argname === :thread
end
function _process_kwarg(kwarg, threadarg = false, broadcastarg = false)
    if _validate_kwarg(kwarg)
        return __process_kwarg(kwarg), broadcastarg
    else
        return threadarg, __process_kwarg(kwarg)
    end
end

macro (..)(kwarg, ex)
    threadarg, broadcastarg = _process_kwarg(kwarg)
    return fb_macro!(ex, threadarg, broadcastarg)
end
macro (..)(kwarg0, kwarg1, ex)
    threadarg, broadcastarg = _process_kwarg(kwarg0)
    threadarg, broadcastarg = _process_kwarg(kwarg1, threadarg, broadcastarg)
    return fb_macro!(ex, threadarg, broadcastarg)
end


let # we could check `hasfield(Method, :recursion_relation)`, but I'd rather see an error if things change
    dont_limit = Returns(true)
    for f in (
            check_no_dyn_broadcast,
            broadcast_islinear,
            semi_static_axes,
            __view,
            _view,
            _rall,
            _rany,
            _rmap,
            _fastindex,
        )
        for m in methods(f)
            m.recursion_relation = dont_limit
        end
    end
end

using PrecompileTools: @compile_workload, @setup_workload

@setup_workload begin
    @compile_workload begin
        source = [1.0, 2.0, 3.0]
        destination = similar(source)
        @.. destination = source + 1
        @.. thread = Serial() destination = 2source + 1
        result = @.. source + 1
        result
    end
end

end
