export JumpProcess, ReactionProcess, DiffusionProcess, LangevinProcess, JumpDiffusionProcess, ControlProcess,
       MarkovProcess, LagrangeMayer, Lagrange, Mayer, ExitProbability, TerminalSetProbability,
       inf_generator, extended_inf_generator

"""
    MarkovProcess

An abstract type to define stochastic processes that are supported by MarkovBounds.jl. 
"""
abstract type MarkovProcess end

"""
    JumpProcess(x, a, h, X = FullSpace();
                iv = PV{true}("t"), controls = [], poly_vars = Dict())

constructor returns `JumpProcess` object with fields
* `x` - vector containing the system state. Can be specified both as PolyVar type object
        from DynamicPolynomials.jl or Num type object generated by Catalyst.jl/ModelingToolkit.jl. 
        In the latter case, PolyVar variables with analogous names are generated automatically and
        a dictionary is assembled that keeps track of this transformation.
* `a` - vector containing the rate coefficients for each independent jump in the system
* `h` - vector of vectors containing the jump functions 
* `X` - basic semialgebraic set enclosing the state space of the process
* `iv` - independent variable of the process (usually time)
* `controls` - vector of control variables (empty if there are none)
* `poly_vars` - Dict mapping symbolic variables created by Catalyst.jl or ModelingToolkit.jl to generated PolyVar placeholders.
"""
mutable struct JumpProcess <: MarkovProcess
    x::Vector{<:PV} # state
    a::Vector{<:Polynomial} # propensities
    h::Vector{Vector{<:Polynomial}} # jumps
    X # state space enclosure
    iv::PV
    controls::Vector{<:PV}
    poly_vars::Dict # needed if process defined in terms of symbolics.jl variables
    function JumpProcess(x::Vector{<:PV}, a::Vector{<:APL}, h::Vector{<:Vector{<:APL}}, X = FullSpace();
                         iv = PV{true}("t"), controls = PV{true}[], poly_vars = Dict())
        return new(x, polynomial.(a), map(p -> polynomial.(p), h), X, iv, controls, poly_vars)
    end
end

JumpProcess(x::PV, a::APL, h::APL, X = FullSpace(); iv = PV{true}("t"), controls = PV{true}[]) =
            JumpProcess([x], [a], [[h]], X; iv=iv, controls = (controls isa Vector ? controls : [controls]))
JumpProcess(x::PV, a::Vector{<:APL}, h::Vector{<:APL}, X = FullSpace(); iv = PV{true}("t"), controls = PV{true}[]) =
            JumpProcess([x], a, [[hi] for hi in h], X; iv = iv, controls = (controls isa Vector ? controls : [controls]))

"""
    ReactionProcess(rn::ReactionSystem, params = Dict())

constructor returns `ReactionProcess` object generated from Catalyst.jl `ReactionSystem` object with given parameters.
The fields of ReactionProcess objects are
* `ReactionSystem`- underlying ReactionSystem object as generated by Catalyst.jl
* `JumpProcess` - JumpProcess object describing a `JumpProcess` that is equivalent to the stochastic reaction network described by ReactionSystem
* `species_to_index` - Dictionary mapping chemical species to its internally used variable species_to_index
* `species_to_state` - Dictionary mapping species to equivalent state in the JumpProcess
* `state_to_species` - Dictionary mapping state in `JumpProcess` to species in reaction network
"""
mutable struct ReactionProcess <: MarkovProcess
    ReactionSystem::ReactionSystem
    JumpProcess::JumpProcess
    species_to_index::Dict
    species_to_state::Dict
    state_to_species::Dict
    function ReactionProcess(rn::ReactionSystem, params = Dict())
        specs = species(rn)
        S = prodstoichmat(rn) - substoichmat(rn)
        n = length(specs)
        @polyvar(x[1:n])
        spec2idx = Dict(specs[i] => i for i in 1:n)
        spec2state = Dict(specs[i] => x[i] for i in 1:n)
        state2spec = Dict(x[i] => specs[i] for i in 1:n)
        props = reformat_reactions(reactions(rn), spec2idx, x, params)
        jumps = reformat_jumps(S, spec2idx, x)
        support = intersect([@set(x[i] >= 0) for i in 1:n]...)
        return new(rn, JumpProcess(x,props,jumps,support), spec2idx, spec2state, state2spec)
    end
end

"""
    DiffusionProcess(x, f, σ, X = FullSpace();
                     iv = PV{true}("t"), controls = [], poly_vars = Dict())

constructor returns `DiffusionProcess` object with fields
* `x` - vector containing the system state. Can be specified both as PolyVar type object
        from DynamicPolynomials.jl or Num type object generated by Catalyst.jl/ModelingToolkit.jl. 
        In the latter case, PolyVar variables with analogous names are generated automatically and
        a dictionary is assembled that keeps track of this transformation.
* `f` - vector of drift coefficients for each state
* `σ` - diffusion matrix of the process
* `X` - basic semialgebraic set enclosing the state space of the process
* `iv` - independent variable of the process (usually time)
* `controls` - vector of control variables (empty if there are none)
* `poly_vars` - Dict mapping symbolic variables created by Catalyst.jl or ModelingToolkit.jl to generated PolyVar placeholders.
"""
mutable struct DiffusionProcess <: MarkovProcess
    x::Vector{<:PV}  # state
    f::Vector{<:Polynomial} # drift
    σ::Matrix{<:Polynomial} # diffusion matrix
    X # support
    iv::PV
    controls::Vector{<:PV}
    poly_vars::Dict # needed if process defined in terms of symbolics.jl variables
    function DiffusionProcess(x::Vector{<:PV}, f::Vector{<:APL}, σ::Matrix{<:APL}, X = FullSpace();
                              iv = PV{true}("t"), controls = PV{true}[], poly_vars = Dict())
        return new(x, polynomial.(f), polynomial.(σ), X, iv, controls, poly_vars)
    end
end

DiffusionProcess(x::PV, f::APL, σ::APL, X = FullSpace(); iv::PV = PV{true}("t"), controls = PV{true}[]) =
                 DiffusionProcess([x], [f], reshape([σ],1,1), X; iv=iv, controls = (controls isa Vector ? controls : [controls]))
                
mutable struct LangevinProcess <: MarkovProcess
    ReactionSystem::ReactionSystem
    DiffusionProcess::DiffusionProcess
    species_to_index::Dict
    species_to_state::Dict
    state_to_species::Dict
    function LangevinProcess(rn::ReactionSystem, params = Dict())
        specs = species(rn)
        S = prodstoichmat(rn) - substoichmat(rn)
        n = length(specs)
        @polyvar(x[1:n])
        spec2idx = Dict(specs[i] => i for i in 1:n)
        spec2state = Dict(specs[i] => x[i] for i in 1:n)
        state2spec = Dict(x[i] => specs[i] for i in 1:n)
        props = polynomial.(reformat_reactions(reactions(rn), spec2idx, x, params))
        drift = S'*props
        diff = [sum(S[k,i]*S[k,j]*props[k] for k in eachindex(props)) for i in 1:n, j in 1:n]
        support = intersect([@set(x[i] >= 0) for i in 1:n]...)
        return new(rn, DiffusionProcess(x,drift,diff,support), spec2idx, spec2state, state2spec)
    end
end

"""
    JumpDiffusionProcess(x, a, h, f, σ, X = FullSpace()
                         iv = PV{true}("t"), controls = [], poly_vars = Dict())

constructor returns `JumpDiffusionProcess` object with fields
* `x` - vector containing the system state. Can be specified both as PolyVar type object
        from DynamicPolynomials.jl or Num type object generated by Catalyst.jl/ModelingToolkit.jl. 
        In the latter case, PolyVar variables with analogous names are generated automatically and
        a dictionary is assembled that keeps track of this transformation.
* `a` - vector containing the rate coefficients for each independent jump in the system
* `h` - vector of vectors containing the jump functions 
* `f` - vector of drift coefficients for each state
* `σ` - diffusion matrix of the process
* `X` - basic semialgebraic set enclosing the state space of the process
* `iv` - independent variable of the process (usually time)
* `controls` - vector of control variables (empty if there are none)
* `poly_vars` - Dict mapping symbolic variables created by Catalyst.jl or ModelingToolkit.jl to generated PolyVar placeholders.
"""
mutable struct JumpDiffusionProcess <: MarkovProcess
    x::Vector{<:PV} # state
    a::Vector{<:Polynomial} # propensities
    h::Vector{Vector{<:Polynomial}} # jumps
    f::Vector{<:Polynomial} # drift
    σ::Matrix{<:Polynomial} # diffusion matrix
    X # state space enclosure
    iv::PV
    controls::Vector{<:PV}
    poly_vars::Dict # needed if process defined in terms of symbolics.jl variables
    function JumpDiffusionProcess(x::Vector{<:PV}, a::Vector{<:APL}, h::Vector{<:Vector{<:APL}}, f::Vector{<:APL}, σ::Matrix{<:APL}, X = FullSpace();
                                  iv = PV{true}("t"), controls = PV{true}[], poly_vars = Dict())
        return new(x, polynomial.(a), map(p -> polynomial.(p), h), polynomial.(f), polynomial.(σ), X, iv, controls, poly_vars)
    end
end

JumpDiffusionProcess(x::PV, a::APL, h::APL, f::APL, σ::APL, X = Fullspace(); iv = PV{true}("t"), controls = PV{true}[]) =
                     JumpDiffusionProcess([x], [a], [[h]], [f], reshape(σ,1,1), X; iv=iv, controls=(controls isa Vector ? controls : [controls]))
JumpDiffusionProcess(x::PV, a::Vector{<:APL}, h::Vector{<:APL}, f::APL, σ::APL, X = Fullspace(); iv = PV{true}("t"), controls=PV{true}[]) =
                     JumpDiffusionProcess([x], a, [[hi] for hi in h], [f], reshape(σ,1,1), X; iv=iv, controls=(controls isa Vector ? controls : [controls]))

function JumpDiffusionProcess(JP::JumpProcess, DP::DiffusionProcess)
    @assert all(JP.x .== DP.x) "The jump and diffusion process must have the same state"
    return JumpDiffusionProcess(JP.x, JP.a, JP.h, DP.f, DP.σ, intersect(DP.X, JP.X), iv = JP.iv, controls = JP.controls)
end

"""
    ControlProcess(MP::MarkovProcess, T::Real, U, obj, PCs = [], TCs = [], dis_fac = 0)

constructor returns object of type `ControlProcess` with fields
* `MP` - MarkovProcess to be controlled
* `T` - length of the time horizon (can be Inf)
* `U` - closed basic semialgebraic set describing admissible control actions
* `PathChanceConstraints` - Array of chance path constraints
* `TerminalChanceConstraints` - Array of terminal chance constraints
* `discount_factor` - discount factor for accumulating stage cost
"""
mutable struct ControlProcess
    MP::MarkovProcess
    T::Real  # time horizon
    U # set of admissible controls
    Objective # objetive function
    PathChanceConstraints
    TerminalChanceConstraints
    discount_factor
    function ControlProcess(MP::MarkovProcess, T::Real, U, obj, PCs = [], TCs = [], dis_fac = 0)
        return new(MP, T, U, obj, PCs, TCs, dis_fac)
    end
end

"""
    ExitProbability

a type used to specify an objective function quantifiying the probability for a supported Markov process 
to exit a specified closed basic semialgebraic set on the time horizon specified in a `ControlProcess` object. 

Fields: 
* X - basic semialgebraic set whose exit probability is to be quantified
"""
mutable struct ExitProbability
    X::BasicSemialgebraicSet
end


"""
    TerminalSetProbability

a type used to specify an objective function quantifiying the probability for a supported Markov process 
to occupy a specified closed basic semialgebraic set at the end of the control horizon specified in a `ControlProcess` object. 

Fields: 
* X - basic semialgebraic set whose occupation probability is to be quantified at the end of the control horizon
"""
mutable struct TerminalSetProbability
    X::BasicSemialgebraicSet
end

@doc raw"""
    LagrangeMayer

a type used to specify an objective function consisting of a Lagrange and Mayer term, i.e., ``\mathbb{E}\left[ \int_{0}^T l(x(t), u(t)) \, dt + m(x(T))\right]``.
The Lagrange and Mayer term are specified in terms of polynomial functions ``l : X \times U \to \mathbb{R}`` and ``m : X \to \mathbb{R}``, respectively.


Fields: 
* `l` - `AbstractPolynomialLike` determining the Lagrange term
* `m` - `AbstractPolynomialLike` determining the Mayer term

If only a Lagrange or Mayer term is neede for specification of the objective function, one can make use of the functions
`Lagrange(l::AbstractPolynomialLike)` and `Mayer(m::AbstractPolynomialLike)` for more convenient model specification.
"""
mutable struct LagrangeMayer
    l::APL
    m::APL
end

Lagrange(l) = LagrangeMayer(l, polynomial(0)) # not elegant but works
Mayer(m) = LagrangeMayer(polynomial(0), m) # not elegant but works

mutable struct ChanceConstraint
    X::BasicSemialgebraicSet
    α::Real # confidence level
end


inf_generator(MP::JumpProcess, w::Polynomial) = sum(MP.a[i]*(subs(w, MP.x => MP.h[i]) - w) for i in 1:length(MP.a))
inf_generator(MP::ReactionProcess, w::Polynomial) = inf_generator(MP.JumpProcess,w)
inf_generator(MP::DiffusionProcess, w::Polynomial) = MP.f'*∂(w,MP.x) + 1/2*sum(∂²(w,MP.x,MP.x) .* MP.σ)
inf_generator(MP::JumpDiffusionProcess, w::Polynomial) = MP.f'*∂(w,MP.x) + 1/2*sum(∂²(w,MP.x,MP.x) .* MP.σ) + sum(MP.a[i]*(subs(w, MP.x => MP.h[i]) - w) for i in 1:length(MP.a))
extended_inf_generator(MP::MarkovProcess, w::Polynomial; scale = 1) = ∂(w,MP.iv) + scale*inf_generator(MP, w)

function inf_generator(MP::JumpProcess, w::Dict, idx::Tuple{Int, Int}, p::Partition)
    k, v = idx
    x_v = props(p.graph, v)[:cell].x
    gen = 0
    for i in 1:length(MP.a)
        if isempty(intersect(MP.a[i].x.vars, MP.controls)) 
            prop = MP.a[i](MP.x => x_v)
            skip = prop == 0 ? true : false
        else
            prop = subs(MP.a[i], MP.x => x_v)
            skip = false
        end
        if !skip
            x_u = [h(MP.x => x_v) for h in MP.h[i]]
            u = p.get_vertex(x_u)
            if props(p.graph, u)[:cell] isa Singleton
                gen += prop*(w[k,u] - w[k,v])
            else
                gen += prop*(subs(w[k,u], MP.x => x_u) - w[k,v])
            end
        end
    end
    return gen
end

function inf_generator(MP::JumpProcess, w::Dict, v::Int, p::Partition)
    x_v = props(p.graph, v)[:cell].x
    gen = 0
    for i in 1:length(MP.a)
        if isempty(intersect(MP.a[i].x.vars, MP.controls)) 
            prop = MP.a[i](MP.x => x_v)
            skip = prop == 0 ? true : false
        else
            prop = subs(MP.a[i], MP.x => x_v)
            skip = false
        end
        if !skip
            x_u = [h(MP.x => x_v) for h in MP.h[i]]
            u = p.get_vertex(x_u)
            if props(p.graph, u)[:cell] isa Singleton
                gen += prop*(w[u] - w[v])
            else
                gen += prop*(subs(w[u], MP.x => x_u) - w[v])
            end
        end
    end
    return gen
end

extended_inf_generator(MP::JumpProcess, w, v, p::Partition; scale = 1) = ∂(w[v],MP.iv) + scale*inf_generator(MP, w, v, p)

mutable struct Bound
    value::Real
    model::Model
    partition::Partition
    w
end
