mutable struct JumpProblem{iip,P,A,C,J<:Union{Nothing,AbstractJumpAggregator},J2,J3,J4} <: DiffEqBase.AbstractJumpProblem{P,J}
  prob::P
  aggregator::A
  discrete_jump_aggregation::J
  jump_callback::C
  variable_jumps::J2
  regular_jump::J3
  massaction_jump::J4
end

DiffEqBase.isinplace(::JumpProblem{iip}) where {iip} = iip
JumpProblem(prob::JumpProblem) = prob

JumpProblem(prob,jumps::ConstantRateJump;kwargs...) = JumpProblem(prob,JumpSet(jumps);kwargs...)
JumpProblem(prob,jumps::VariableRateJump;kwargs...) = JumpProblem(prob,JumpSet(jumps);kwargs...)
JumpProblem(prob,jumps::RegularJump;kwargs...) = JumpProblem(prob,JumpSet(jumps);kwargs...)
JumpProblem(prob,jumps::MassActionJump;kwargs...) = JumpProblem(prob,JumpSet(jumps);kwargs...)
JumpProblem(prob,jumps::AbstractJump...;kwargs...) = JumpProblem(prob,JumpSet(jumps...);kwargs...)

JumpProblem(prob,aggregator::AbstractAggregatorAlgorithm,jumps::ConstantRateJump;kwargs...) = JumpProblem(prob,aggregator,JumpSet(jumps);kwargs...)
JumpProblem(prob,aggregator::AbstractAggregatorAlgorithm,jumps::VariableRateJump;kwargs...) = JumpProblem(prob,aggregator,JumpSet(jumps);kwargs...)
JumpProblem(prob,aggregator::AbstractAggregatorAlgorithm,jumps::RegularJump;kwargs...) = JumpProblem(prob,aggregator,JumpSet(jumps);kwargs...)
JumpProblem(prob,aggregator::AbstractAggregatorAlgorithm,jumps::MassActionJump;kwargs...) = JumpProblem(prob,aggregator,JumpSet(jumps);kwargs...)
JumpProblem(prob,aggregator::AbstractAggregatorAlgorithm,jumps::AbstractJump...;kwargs...) = JumpProblem(prob,aggregator,JumpSet(jumps...);kwargs...)
JumpProblem(prob,jumps::JumpSet;kwargs...) = JumpProblem(prob,NullAggregator(),jumps;kwargs...)

function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpSet;
                     save_positions = typeof(prob) <: DiffEqBase.AbstractDiscreteProblem ? (false,true) : (true,true),
                     rng = Xorshifts.Xoroshiro128Star(rand(UInt64)), kwargs...)

  ## Constant Rate Handling
  t,end_time,u = prob.tspan[1],prob.tspan[2],prob.u0
  if (typeof(jumps.constant_jumps) <: Tuple{}) && (jumps.massaction_jump === nothing)
    disc = nothing
    constant_jump_callback = CallbackSet()
  else
    disc = aggregate(aggregator,u,prob.p,t,end_time,jumps.constant_jumps,jumps.massaction_jump,save_positions,rng;kwargs...)
    constant_jump_callback = DiscreteCallback(disc)
  end

  iip = if prob isa DiscreteProblem && prob.f === DiffEqBase.DISCRETE_INPLACE_DEFAULT && jumps.regular_jump !== nothing
    # Just a default discrete problem f, so don't use it for iip
    DiffEqBase.isinplace(jumps.regular_jump)
  else
    DiffEqBase.isinplace(prob)
  end

  ## Variable Rate Handling
  if typeof(jumps.variable_jumps) <: Tuple{}
    new_prob = prob
    variable_jump_callback = CallbackSet()
  else
    new_prob = extend_problem(prob,jumps)
    variable_jump_callback = build_variable_callback(CallbackSet(),0,jumps.variable_jumps...)
  end
  callbacks = CallbackSet(constant_jump_callback,variable_jump_callback)
  JumpProblem{iip,typeof(new_prob),typeof(aggregator),typeof(callbacks),
              typeof(disc),typeof(jumps.variable_jumps),
              typeof(jumps.regular_jump),typeof(jumps.massaction_jump)}(
                        new_prob,aggregator,disc,
                        callbacks,
                        jumps.variable_jumps,
                        jumps.regular_jump, jumps.massaction_jump)
end

function extend_problem(prob::DiffEqBase.AbstractODEProblem,jumps)
  function jump_f(du::ExtendedJumpArray,u::ExtendedJumpArray,p,t)
    prob.f(du.u,u.u,p,t)
    update_jumps!(du,u,p,t,length(u.u),jumps.variable_jumps...)
  end
  u0 = ExtendedJumpArray(prob.u0,[-randexp() for i in 1:length(jumps.variable_jumps)])
  remake(prob,f=ODEFunction{true}(jump_f),u0=u0)
end

function extend_problem(prob::DiffEqBase.AbstractSDEProblem,jumps)
  function jump_f(du,u,p,t)
    prob.f(du.u,u.u,p,t)
    update_jumps!(du,u,p,t,length(u.u),jumps.variable_jumps...)
  end

  if prob.noise_rate_prototype === nothing
    jump_g = function (du,u,p,t)
      prob.g(du.u,u.u,p,t)
    end
  else
    jump_g = function (du,u,p,t)
      prob.g(du,u.u,p,t)
    end
  end

  u0 = ExtendedJumpArray(prob.u0,[-randexp() for i in 1:length(jumps.variable_jumps)])
  remake(prob,f=SDEFunction{true}(jump_f,jump_g),g=jump_g,u0=u0)
end

function extend_problem(prob::DiffEqBase.AbstractDDEProblem,jumps)
  jump_f = function (du,u,h,p,t)
    prob.f(du.u,u.u,h,p,t)
    update_jumps!(du,u,p,t,length(u.u),jumps.variable_jumps...)
  end
  u0 = ExtendedJumpArray(prob.u0,[-randexp() for i in 1:length(jumps.variable_jumps)])
  ramake(prob,f=DDEFunction{true}(jump_f),u0=u0)
end

# Not sure if the DAE one is correct: Should be a residual of sorts
function extend_problem(prob::DiffEqBase.AbstractDAEProblem,jumps)
  jump_f = function (out,du,u,p,t)
    prob.f(out.u,du.u,u.u,t)
    update_jumps!(du,u,t,length(u.u),jumps.variable_jumps...)
  end
  u0 = ExtendedJumpArray(prob.u0,[-randexp() for i in 1:length(jumps.variable_jumps)])
  remake(prob,f=DAEFunction{true}(jump_f),u0=u0)
end

function build_variable_callback(cb,idx,jump,jumps...)
  idx += 1
  condition = function (u,t,integrator)
    u.jump_u[idx]
  end
  affect! = function (integrator)
    jump.affect!(integrator)
    integrator.u.jump_u[idx] = -randexp()
  end
  new_cb = ContinuousCallback(condition,affect!;
                      idxs = jump.idxs,
                      rootfind = jump.rootfind,
                      interp_points = jump.interp_points,
                      save_positions = jump.save_positions,
                      abstol = jump.abstol,
                      reltol = jump.reltol)
  build_variable_callback(CallbackSet(cb,new_cb),idx,jumps...)
end

function build_variable_callback(cb,idx,jump)
  idx += 1
  condition = function (u,t,integrator)
    u.jump_u[idx]
  end
  affect! = function (integrator)
    jump.affect!(integrator)
    integrator.u.jump_u[idx] = -randexp()
  end
  new_cb = ContinuousCallback(condition,affect!;
                      idxs = jump.idxs,
                      rootfind = jump.rootfind,
                      interp_points = jump.interp_points,
                      save_positions = jump.save_positions,
                      abstol = jump.abstol,
                      reltol = jump.reltol)
  CallbackSet(cb,new_cb)
end

aggregator(jp::JumpProblem{P,A,C,J,J2}) where {P,A,C,J,J2} = A

@inline function extend_tstops!(tstops,jp::JumpProblem{P,A,C,J,J2}) where {P,A,C,J,J2}
  !(typeof(jp.jump_callback.discrete_callbacks) <: Tuple{}) && push!(tstops,jp.jump_callback.discrete_callbacks[1].condition.next_jump_time)
end

@inline function update_jumps!(du,u,p,t,idx,jump)
  idx += 1
  du[idx] = jump.rate(u,p,t)
end

@inline function update_jumps!(du,u,p,t,idx,jump,jumps...)
  idx += 1
  du[idx] = jump.rate(u,p,t)
  update_jumps!(du,u,p,t,idx,jumps...)
end


### Displays

Base.summary(prob::JumpProblem) = string(DiffEqBase.parameterless_type(prob)," with problem ",DiffEqBase.parameterless_type(prob.prob)," and aggregator ",typeof(prob.aggregator))
function Base.show(io::IO, A::JumpProblem)
  println(io,summary(A))
  println(io,"Number of constant rate jumps: ",A.discrete_jump_aggregation === nothing ? 0 : length(A.discrete_jump_aggregation.rates))
  println(io,"Number of variable rate jumps: ",length(A.variable_jumps))
  if A.regular_jump !== nothing
    println(io,"Have a regular jump")
  end
  if (A.massaction_jump !== nothing) && (get_num_majumps(A.massaction_jump) > 0)
    println(io,"Have a mass action jump")
  end
end

TreeViews.hastreeview(x::JumpProblem) = true
