using Base: RefValue
using Base: ismutabletype

# Interfaces

accum() = nothing
accum(x) = x

accum(x, y) =
  x === nothing ? y :
  y === nothing ? x :
  x + y

accum(x, y, zs...) = accum(accum(x, y), zs...)

accum(x::Tuple, ys::Tuple...) = map(accum, x, ys...)
accum(x::AbstractArray, ys::AbstractArray...) = Base.broadcast_preserving_zero_d(accum, x, ys...)
accum(::Tuple{}, ::NamedTuple{}) = ()
accum(::NamedTuple{}, ::Tuple{}) = ()

@generated function accum(x::NamedTuple, y::NamedTuple)
  # assumes that y has no keys apart from those also in x
  fieldnames(y) ⊆ fieldnames(x) || throw(ArgumentError("$y keys must be a subset of $x keys"))

  grad(field) = field in fieldnames(y) ? :(y.$field) : :nothing
  Expr(:tuple, [:($f=accum(x.$f, $(grad(f)))) for f in fieldnames(x)]...)
end

function accum(x::RefValue, y::RefValue)
  @assert x === y
  return x
end
function accum(x::NamedTuple, ref::RefValue)
    # We do not actually do any accumulation here, because the ref will already have been mutated.
  fieldnames(typeof(ref[])) ⊆ fieldnames(typeof(x)) || throw(ArgumentError("$(ref[]) keys from Ref must be a subset of $x keys"))
  ref
end
function accum(ref::RefValue, x::NamedTuple)
    # We do not actually do any accumulation here, because the ref will already have been mutated.
  fieldnames(typeof(x)) ⊆ fieldnames(typeof(ref[])) || throw(ArgumentError("$x keys must be a subset of $(ref[]) keys from Ref"))
  ref
end

accum(x::NamedTuple, y::ChainRulesCore.Tangent) = accum(x, wrap_chainrules_output(y))
accum(x::ChainRulesCore.Tangent, y::NamedTuple) = accum(wrap_chainrules_output(x), y)

accum(x::Nothing, y::AbstractThunk) = y
accum(x::AbstractThunk, y::Nothing) = x

accum(x, y::AbstractThunk) = accum(x, unthunk(y))
accum(x::AbstractThunk, y) = accum(unthunk(x), y)
accum(x::AbstractThunk, y::AbstractThunk) = accum(unthunk(x), unthunk(y))

# Core functions
@_adjoint_keepthunks deepcopy(x) = deepcopy(x), ȳ -> (ȳ,)

@_adjoint_keepthunks (::Type{V})(x...) where V<:Val = V(x...), _ -> nothing

@_adjoint_keepthunks ifelse(cond::Bool, t, f) =
  ifelse(cond, t, f),
  Δ -> cond ? (nothing, Δ, zero(Δ)) : (nothing, zero(Δ), Δ)

@_adjoint_keepthunks Base.typeassert(x, T) = Base.typeassert(x, T), Δ -> (Δ, nothing)

accum_param(::Context{false}, _, Δ) = Δ
@generated function accum_param(cx::Context, x, Δ)
  isbitstype(x) && return :(Δ)
  quote
    if haskey(cache(cx), x)
      cache(cx)[x] = accum(cache(cx)[x],Δ)
      return
    else
      return Δ
    end
  end
end

function accum_global(cx::Context, ref, x̄)
  (x̄ === nothing || isconst(ref.mod, ref.name)) && return
  gs = cache(cx)
  gs[ref] = accum(get(gs, ref, nothing), x̄)
  return
end

unwrap(x) = x

@_adjoint_keepthunks unwrap(x) = unwrap(x), x̄ -> (accum_param(__context__, x, x̄),)

unwrap(ref, x) = x

@_adjoint_keepthunks unwrap(ref, x) = unwrap(x), function (x̄)
  accum_global(__context__, ref, x̄)
  (accum_param(__context__, x, x̄),)
end

function global_set(ref, val)
    setglobal!(ref.mod, ref.name, val)
end

@_adjoint_keepthunks! function global_set(ref, x)
  global_set(ref, x), function (x̄)
    gs = cache(__context__)
    x̄ = accum(get(gs, ref, nothing), x̄)
    gs[ref] = nothing
    return (nothing, x̄)
  end
end

# Tuples

using Base: tail

@_adjoint_keepthunks tuple(xs...) = xs, identity

@_adjoint_keepthunks function literal_getindex(xs::NTuple{N,Any}, ::Val{i}) where {N,i}
  val = xs[i]
  function back(Δ)
    accum_param(__context__, val, Δ) === nothing && return
    return ntuple(j -> i == j ? Δ : nothing, Val(N)), nothing
  end
  val, back
end

@_adjoint_keepthunks function getindex(xs::NTuple{N,Any}, i::Integer) where N
  val = xs[i]
  function back(Δ)
    accum_param(__context__, val, Δ) === nothing && return
    return ntuple(j -> i == j ? Δ : nothing, Val(N)), nothing
  end
  return val, back
end

@_adjoint_keepthunks getindex(xs::NTuple{N,Any}, r::AbstractUnitRange) where N =
  (xs[r], Δ -> (ntuple(j -> j in r ? Δ[findfirst(isequal(j), r)] : nothing, Val(N)), nothing))

@_adjoint_keepthunks function getindex(xs::NTuple{N,Any}, r::AbstractVector) where N
  val = xs[r]
  function back(Δ)
    dxs = ntuple(Val(length(xs))) do x
      total = nothing
      for r_i in eachindex(r)
        r[r_i] === x || continue
        total = accum(total, Δ[r_i])
      end
      return total
    end
    return (dxs, nothing)
  end
  val, back
end

function _pullback(cx::AContext, ::typeof(literal_indexed_iterate), xs::Tuple, ::Val{i}) where i
  y, b = _pullback(cx, literal_getindex, xs, Val(i))
  back(::Nothing) = nothing
  back(ȳ) = b(ȳ[1])
  (y, i+1), back
end

function _pullback(cx::AContext, ::typeof(literal_indexed_iterate), xs::Tuple, ::Val{i}, st) where i
  y, b = _pullback(cx, literal_getindex, xs, Val(i))
  back(::Nothing) = nothing
  back(ȳ) = (b(ȳ[1])..., nothing)
  (y, i+1), back
end

# Needed for iteration lowering
@_adjoint_keepthunks Core.getfield(xs::NTuple{N,Any}, i::Int) where N =
  (xs[i], Δ -> (ntuple(j -> i == j ? Δ : nothing, Val(N)), nothing))

@_adjoint_keepthunks Core.getfield(xs::NamedTuple{K,<:NTuple{N,Any}}, i::Int) where {K,N} =
  (xs[i], Δ -> (NamedTuple{K}(ntuple(j -> i == j ? Δ : nothing, Val(N))), nothing))

@_adjoint_keepthunks function Base.first(xs::Tuple)
  drest = map(_->nothing, tail(xs))
  first(xs), Δ -> ((Δ, drest...),)
end

@_adjoint_keepthunks Base.tail(xs::Tuple) = tail(xs), x̄s -> ((nothing, x̄s...),)

_empty(x) = length(x)
_empty(x::Union{Tuple,NamedTuple}) = map(_->nothing, x)

_unapply(t::Integer, xs) = xs[1:t], xs[t+1:end]
_unapply(t, xs) = first(xs), tail(xs)
_unapply(t::Tuple{}, xs) = (), xs

function _unapply(t::Tuple, xs)
  t1, xs1 = _unapply(first(t), xs)
  t2, xs2 = _unapply(tail(t), xs1)
  (t1, t2...), xs2
end

function _unapply(t::NamedTuple{K}, xs) where K
  t, rst = _unapply(Tuple(t), xs)
  NamedTuple{K}(t), rst
end

unapply(t, xs) = _unapply(t, xs)[1]

@_adjoint_keepthunks! function Core._apply(f, args...)
  y, back = Core._apply(_pullback, (__context__, f), args...)
  st = map(_empty, args)
  y, function (Δ)
    Δ = back(Δ)
    Δ === nothing ? nothing :
      (first(Δ), unapply(st, Base.tail(Δ))...)
  end
end

@_adjoint_keepthunks! function Core._apply_iterate(::typeof(iterate), f, args...)
  y, back = Core._apply(_pullback, (__context__, f), args...)
  st = map(_empty, args)
  y, function (Δ)
    Δ = back(Δ)
    Δ === nothing ? nothing :
      (nothing, first(Δ), unapply(st, Base.tail(Δ))...)
  end
end

# Structs

deref!(x) = x

function deref!(x::Ref)
  d = x[]
  x[] = nothing
  return d
end

@generated nt_nothing(x) = Expr(:tuple, [:($f=nothing) for f in fieldnames(x)]...)

@generated pair(::Val{k}, v, _=nothing) where k = :($k = v,)
@generated pair(::Val{k}, v, ::NamedTuple{keys}) where {k,keys} = k isa Int ? :($(getfield(keys, k)) = v,) : :($k = v,)

@_adjoint_keepthunks function literal_getfield(x, ::Val{f}) where f
  val = getfield(x, f)
  function back(Δ)
    accum_param(__context__, val, Δ) === nothing && return
    if isimmutable(x)
      dx = (; nt_nothing(x)..., pair(Val(f), Δ, x)...)
      (_project(x, dx), nothing)
    else
      dx = grad_mut(__context__, x)
      dx[] = (; dx[]..., pair(Val(f), accum(getfield(dx[], f), Δ))...)
      return (dx,nothing)
    end
  end
  unwrap(val), back
end

_pullback(cx::AContext, ::typeof(getfield), x, field_name::Symbol) =
  _pullback(cx, literal_getfield, x, Val(field_name))

# Field access by a non-literal integer index (e.g. `ColorTypes._comp(Val{N}, c)`
# lowers to `getfield(c, N)` with `N` a runtime `Int`). Convert the index to the
# field *name* and reuse the `Symbol` path, so the tangent is built with a valid
# field name rather than the bare integer. For a `Tuple`, `fieldname` returns the
# integer itself, which routes on to `literal_getindex` as usual. (Literal-integer
# `getfield` is already handled at IR level by `instrument_getfield!`.)
_pullback(cx::AContext, ::typeof(getfield), x, i::Int) =
  _pullback(cx, literal_getfield, x, Val(fieldname(typeof(x), i)))

function _pullback(cx::AContext, ::typeof(literal_getproperty), x::NamedTuple,
                   ::Val{property_name}) where {property_name}
  return _pullback(cx, literal_getfield, x, Val(property_name))
end
function _pullback(cx::AContext, ::typeof(literal_getindex), x::NamedTuple,
                   ::Val{key}) where {key}
  return _pullback(cx, literal_getfield, x, Val(key))
end

function _pullback(cx::AContext, ::typeof(literal_getproperty), x::Tuple,
                   ::Val{index}) where {index}
  return _pullback(cx, literal_getindex, x, Val(index))
end
function _pullback(cx::AContext, ::typeof(literal_getfield), x::Tuple,
                   ::Val{index}) where {index}
  return _pullback(cx, literal_getindex, x, Val(index))
end

# Reading a binding out of a `Module` is not differentiable: a module global is a
# constant from AD's point of view, never a differentiable input. Without these
# methods the source transform recurses through `getproperty(::Module, ::Symbol)`
# into the `getglobal` builtin, which has no IR and no adjoint — historically
# surfacing as a cryptic `UndefVarError: j` (issues #194, #252, #467, #619). The
# `accum_param` call keeps the legacy implicit-`Params` mode working when the
# binding's *value* (e.g. a qualified global array) is a tracked parameter.
function _pullback(cx::AContext, ::typeof(literal_getproperty), m::Module,
                   ::Val{property_name}) where {property_name}
  val = getproperty(m, property_name)
  module_getproperty_pullback(Δ) = (accum_param(cx, val, Δ); nothing)
  return val, module_getproperty_pullback
end
function _pullback(cx::AContext, ::typeof(literal_getfield), m::Module,
                   ::Val{field_name}) where {field_name}
  val = getfield(m, field_name)
  module_getfield_pullback(Δ) = (accum_param(cx, val, Δ); nothing)
  return val, module_getfield_pullback
end

grad_mut(x) = Ref{Any}(nt_nothing(x))

grad_mut(cx::Context, x) = _get!(() -> grad_mut(x), cache(cx), x)

# needed for reverse-over-reverse pending rrule for Base.get!
function _get!(default::Base.Callable, ch, x)
  if haskey(ch, x)
    ch[x]
  else
    ch[x] = default()
  end
end

@_adjoint_keepthunks! function setfield!(x, f, val)
  y = setfield!(x, f, val)
  g = grad_mut(__context__, x)
  y, function (_)
    Δ = getfield(g[], f)
    g[] = (;g[]...,pair(Val(f),nothing)...)
    (nothing, nothing, Δ)
  end
end

struct Jnew{T,G,splat}
  g::G
end

Jnew{T}(g) where T = Jnew{T,typeof(g)}(g)

@_adjoint_keepthunks! function __new__(T, args...)
  x = __new__(T, args...)
  g = !ismutabletype(T) || fieldcount(T) == 0 ? nothing : grad_mut(__context__, x)
  x, Jnew{T,typeof(g),false}(g)
end

@_adjoint_keepthunks! function __splatnew__(T, args)
  x = __splatnew__(T, args)
  g = !ismutabletype(T) || fieldcount(T) == 0 ? nothing : grad_mut(__context__, x)
  x, Jnew{T,typeof(g),true}(g)
end

# TODO captured mutables + multiple calls to `back`
@generated function (back::Jnew{T,G,false})(Δ::Union{NamedTuple,Nothing,RefValue}) where {T,G}
  !ismutabletype(T) && Δ == Nothing && return :nothing
  Δ = G == Nothing ? :Δ :
      Δ <: RefValue ? :(back.g[]) :
      :(accum(back.g[], Δ))
  quote
    x̄ = $Δ
    $(G == Nothing || :(back.g[] = nt_nothing($Δ)))
    (nothing, $(map(f -> :(x̄.$f), fieldnames(T))...))
  end
end

@generated function (back::Jnew{T,G,true})(Δ::Union{NamedTuple,Nothing,RefValue}) where {T,G}
  !ismutabletype(T) && Δ == Nothing && return :nothing
  Δ = G == Nothing ? :Δ : :(back.g)
  quote
    x̄ = $Δ
    $(G == Nothing || :($Δ = nt_nothing($Δ)))
    (nothing, ($(map(f -> :(x̄.$f), fieldnames(T))...),))
  end
end

(back::Jnew{T})(Δ) where T = error("Need an adjoint for constructor $T. Gradient is of type $(typeof(Δ))")
