using Symbolics: get_variables
"""
    inputs(sys)

Return all variables that mare marked as inputs. See also [`unbound_inputs`](@ref)
See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref)
"""
inputs(sys) = collect(get_inputs(sys))

"""
    outputs(sys)

Return all variables that mare marked as outputs. See also [`unbound_outputs`](@ref)
See also [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
"""
function outputs(sys)
    return collect(get_outputs(sys))
end

"""
    bound_inputs(sys)

Return inputs that are bound within the system, i.e., internal inputs
See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
"""
bound_inputs(sys) = filter(x -> is_bound(sys, x), inputs(sys))

"""
    unbound_inputs(sys)

Return inputs that are not bound within the system, i.e., external inputs
See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
"""
unbound_inputs(sys) = filter(x -> !is_bound(sys, x), inputs(sys))

"""
    bound_outputs(sys)

Return outputs that are bound within the system, i.e., internal outputs
See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
"""
bound_outputs(sys) = filter(x -> is_bound(sys, x), outputs(sys))

"""
    unbound_outputs(sys)

Return outputs that are not bound within the system, i.e., external outputs
See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
"""
unbound_outputs(sys) = filter(x -> !is_bound(sys, x), outputs(sys))

"""
    is_bound(sys, u)

Determine whether input/output variable `u` is "bound" within the system, i.e., if it's to be considered internal to `sys`.
A variable/signal is considered bound if it appears in an equation together with variables from other subsystems.
The typical usecase for this function is to determine whether the input to an IO component is connected to another component,
or if it remains an external input that the user has to supply before simulating the system.

See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
"""
function is_bound(sys, u, stack = [])
    #=
    For observed quantities, we check if a variable is connected to something that is bound to something further out.
    In the following scenario
    julia> observed(syss)
        2-element Vector{Equation}:
        sys₊y(tv) ~ sys₊x(tv)
        y(tv) ~ sys₊x(tv)
    sys₊y(t) is bound to the outer y(t) through the variable sys₊x(t) and should thus return is_bound(sys₊y(t)) = true.
    When asking is_bound(sys₊y(t)), we know that we are looking through observed equations and can thus ask
    if var is bound, if it is, then sys₊y(t) is also bound. This can lead to an infinite recursion, so we maintain a stack of variables we have previously asked about to be able to break cycles
    =#
    u ∈ Set(stack) && return false # Cycle detected
    eqs = equations(sys)
    eqs = filter(eq -> has_var(eq, u), eqs) # Only look at equations that contain u
    # isout = isoutput(u)
    for eq in eqs
        vars = [get_variables(eq.rhs); get_variables(eq.lhs)]
        for var in vars
            var === u && continue
            if !same_or_inner_namespace(u, var)
                return true
            end
        end
    end
    # Look through observed equations as well
    oeqs = observed(sys)
    oeqs = filter(eq -> has_var(eq, u), oeqs) # Only look at equations that contain u
    for eq in oeqs
        vars = [get_variables(eq.rhs); get_variables(eq.lhs)]
        for var in vars
            var === u && continue
            if !same_or_inner_namespace(u, var)
                return true
            end
            if is_bound(sys, var, [stack; u]) && !inner_namespace(u, var) # The variable we are comparing to can not come from an inner namespace, binding only counts outwards
                return true
            end
        end
    end
    false
end

"""
    same_or_inner_namespace(u, var)

Determine whether `var` is in the same namespace as `u`, or a namespace internal to the namespace of `u`.
Example: `sys.u ~ sys.inner.u` will bind `sys.inner.u`, but `sys.u` remains an unbound, external signal. The namespaced signal `sys.inner.u` lives in a namespace internal to `sys`.
"""
function same_or_inner_namespace(u, var)
    nu = get_namespace(u)
    nv = get_namespace(var)
    nu == nv ||           # namespaces are the same
        startswith(nv, nu) || # or nv starts with nu, i.e., nv is an inner namespace to nu
        occursin(NAMESPACE_SEPARATOR, string(getname(var))) &&
        !occursin(NAMESPACE_SEPARATOR, string(getname(u))) # or u is top level but var is internal
end

function inner_namespace(u, var)
    nu = get_namespace(u)
    nv = get_namespace(var)
    nu == nv && return false
    startswith(nv, nu) || # or nv starts with nu, i.e., nv is an inner namespace to nu
        occursin(NAMESPACE_SEPARATOR, string(getname(var))) &&
        !occursin(NAMESPACE_SEPARATOR, string(getname(u))) # or u is top level but var is internal
end

"""
    get_namespace(x)

Return the namespace of a variable as a string. If the variable is not namespaced, the string is empty.
"""
function get_namespace(x)
    sname = string(getname(x))
    parts = split(sname, NAMESPACE_SEPARATOR)
    if length(parts) == 1
        return ""
    end
    join(parts[1:(end - 1)], NAMESPACE_SEPARATOR)
end

"""
    has_var(eq, x)

Determine whether an equation or expression contains variable `x`.
"""
function has_var(eq::Equation, x)
    has_var(eq.rhs, x) || has_var(eq.lhs, x)
end

has_var(ex, x) = x ∈ Set(get_variables(ex))

# Build control function

"""
    (f_oop, f_ip), x_sym, p_sym, io_sys = generate_control_function(
            sys::System,
            inputs                   = unbound_inputs(sys),
            disturbance_inputs       = disturbances(sys);
            known_disturbance_inputs = nothing,
            implicit_dae             = false,
            simplify                 = false,
            split                    = true,
        )

For a system `sys` with inputs (as determined by [`unbound_inputs`](@ref) or user specified), generate functions with additional input argument `u`

The returned functions are the out-of-place (`f_oop`) and in-place (`f_ip`) forms:
```
f_oop : (x,u,p,t)      -> rhs         # basic form
f_oop : (x,u,p,t,w)    -> rhs         # with known_disturbance_inputs
f_ip  : (xout,x,u,p,t) -> nothing     # basic form
f_ip  : (xout,x,u,p,t,w) -> nothing   # with known_disturbance_inputs
```

The return values also include the chosen state-realization (the remaining unknowns) `x_sym` and parameters, in the order they appear as arguments to `f`.

# Disturbance Handling

- `disturbance_inputs`: Unknown disturbance inputs. The generated dynamics will preserve any state and dynamics associated with these disturbances, but the disturbance inputs themselves will not be included as function arguments. This is useful for state observers that estimate unmeasured disturbances.

- `known_disturbance_inputs`: Known disturbance inputs. The generated dynamics will preserve state and dynamics, and the disturbance inputs will be added as an additional input argument `w` to the generated function: `(x,u,p,t,w)->rhs`.

# Example

```julia
using ModelingToolkit: generate_control_function, varmap_to_vars, defaults
f, x_sym, ps = generate_control_function(sys, expression=Val{false}, simplify=false)
p = varmap_to_vars(defaults(sys), ps)
x = varmap_to_vars(defaults(sys), x_sym)
t = 0
f[1](x, inputs, p, t)
```
"""
function generate_control_function(sys::AbstractSystem, inputs = unbound_inputs(sys),
        disturbance_inputs = disturbances(sys);
        known_disturbance_inputs = nothing,
        disturbance_argument = false,
        implicit_dae = false,
        simplify = false,
        eval_expression = false,
        eval_module = @__MODULE__,
        split = true,
        kwargs...)
    isempty(inputs) && @warn("No unbound inputs were found in system.")

    # Handle backward compatibility for disturbance_argument
    if disturbance_argument
        Base.depwarn("The `disturbance_argument` keyword argument is deprecated. Use `known_disturbance_inputs` instead. " *
                     "For `disturbance_argument=true`, pass `known_disturbance_inputs=disturbance_inputs, disturbance_inputs=nothing`. " *
                     "For `disturbance_argument=false`, use `disturbance_inputs` as before.",
                     :generate_control_function)
        if known_disturbance_inputs !== nothing
            error("Cannot specify both `disturbance_argument=true` and `known_disturbance_inputs`")
        end
        known_disturbance_inputs = disturbance_inputs
        disturbance_inputs = nothing
    end

    # Collect all disturbance inputs for mtkcompile
    all_disturbances = vcat(
        disturbance_inputs === nothing ? [] : disturbance_inputs,
        known_disturbance_inputs === nothing ? [] : known_disturbance_inputs
    )

    if !isscheduled(sys)
        sys = mtkcompile(sys; inputs, disturbance_inputs=all_disturbances, split)
    end

    # Add all disturbances to inputs for the purposes of io processing
    if !isempty(all_disturbances)
        inputs = [inputs; all_disturbances]
    end

    dvs = unknowns(sys)
    ps = parameters(sys; initial_parameters = true)
    ps = setdiff(ps, inputs)

    # Remove unknown disturbances from inputs (we don't want them as actual inputs to the dynamics)
    if disturbance_inputs !== nothing
        inputs = setdiff(inputs, disturbance_inputs)
    end

    inputs = map(value, inputs)

    # Prepare disturbance arrays for substitution and function arguments
    unknown_disturbances = disturbance_inputs === nothing ? [] : unwrap.(disturbance_inputs)
    known_disturbances = known_disturbance_inputs === nothing ? [] : unwrap.(known_disturbance_inputs)

    eqs = [eq for eq in full_equations(sys)]

    # Set unknown disturbance inputs to zero (we just want to keep the disturbance state)
    if !isempty(unknown_disturbances)
        subs = Dict(unknown_disturbances .=> 0)
        eqs = [eq.lhs ~ substitute(eq.rhs, subs) for eq in eqs]
    end
    check_operator_variables(eqs, Differential)
    # substitute x(t) by just x
    rhss = implicit_dae ? [_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs] :
           [eq.rhs for eq in eqs]

    # TODO: add an optional check on the ordering of observed equations
    p = reorder_parameters(sys, ps)
    t = get_iv(sys)

    # Construct args with known disturbances if provided
    if !isempty(known_disturbances)
        args = (dvs, inputs, p..., t, known_disturbances)
    else
        args = (dvs, inputs, p..., t)
    end
    if implicit_dae
        ddvs = map(Differential(get_iv(sys)), dvs)
        args = (ddvs, args...)
    end
    f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
        p_end = length(p) + 2 + implicit_dae, kwargs...)
    f = eval_or_rgf.(f; eval_expression, eval_module)
    f = GeneratedFunctionWrapper{(
        3 + implicit_dae, length(args) - length(p) + 1, is_split(sys))}(f...)
    # Return parameters excluding both control inputs and all disturbances
    ps = setdiff(parameters(sys), inputs, all_disturbances)
    (; f = (f, f), dvs, ps, io_sys = sys)
end

"""
Turn input variables into parameters of the system.
"""
function inputs_to_parameters!(state::TransformationState, inputsyms)
    check_bound = inputsyms === nothing
    @unpack structure, fullvars, sys = state
    @unpack var_to_diff, graph, solvable_graph = structure
    @assert solvable_graph === nothing

    inputs = BitSet()
    var_reidx = zeros(Int, length(fullvars))
    ninputs = 0
    nvar = 0
    new_parameters = []
    input_to_parameters = Dict()
    new_fullvars = []
    for (i, v) in enumerate(fullvars)
        if isinput(v) && !(check_bound && is_bound(sys, v))
            if var_to_diff[i] !== nothing
                error("Input $(fullvars[i]) is differentiated!")
            end
            push!(inputs, i)
            ninputs += 1
            var_reidx[i] = -1
            p = toparam(v)
            push!(new_parameters, p)
            input_to_parameters[v] = p
        else
            nvar += 1
            var_reidx[i] = nvar
            push!(new_fullvars, v)
        end
    end
    if ninputs == 0
        @set! sys.inputs = OrderedSet{BasicSymbolic}()
        @set! sys.outputs = OrderedSet{BasicSymbolic}(filter(isoutput, fullvars))
        state.sys = sys
        return state
    end

    nvars = ndsts(graph) - ninputs
    new_graph = BipartiteGraph(nsrcs(graph), nvars, Val(false))

    for ie in 1:nsrcs(graph)
        for iv in 𝑠neighbors(graph, ie)
            iv = var_reidx[iv]
            iv > 0 || continue
            add_edge!(new_graph, ie, iv)
        end
    end

    new_var_to_diff = DiffGraph(nvars, true)
    for (i, v) in enumerate(var_to_diff)
        new_i = var_reidx[i]
        (new_i < 1 || v === nothing) && continue
        new_v = var_reidx[v]
        @assert new_v > 0
        new_var_to_diff[new_i] = new_v
    end
    @set! structure.var_to_diff = complete(new_var_to_diff)
    @set! structure.graph = complete(new_graph)

    @set! sys.eqs = isempty(input_to_parameters) ? equations(sys) :
                    fast_substitute(equations(sys), input_to_parameters)
    @set! sys.unknowns = setdiff(unknowns(sys), keys(input_to_parameters))
    ps = parameters(sys)

    @set! sys.ps = [ps; new_parameters]
    @set! sys.inputs = OrderedSet{BasicSymbolic}(new_parameters)
    @set! sys.outputs = OrderedSet{BasicSymbolic}(filter(isoutput, fullvars))
    @set! state.sys = sys
    @set! state.fullvars = Vector{BasicSymbolic}(new_fullvars)
    @set! state.structure = structure
    return state
end

"""
    DisturbanceModel{M}

The structure represents a model of a disturbance, along with the input variable that is affected by the disturbance. See [`add_input_disturbance`](@ref) for additional details and an example.

# Fields:

  - `input`: The variable affected by the disturbance.
  - `model::M`: A model of the disturbance. This is typically a `System`, but type that implements [`ModelingToolkit.get_disturbance_system`](@ref)`(dist::DisturbanceModel) -> ::System` is supported.
"""
struct DisturbanceModel{M}
    input::Any
    model::M
    name::Symbol
end
DisturbanceModel(input, model; name) = DisturbanceModel(input, model, name)

# Point of overloading for libraries, e.g., to be able to support disturbance models from ControlSystemsBase
function get_disturbance_system(dist::DisturbanceModel{System})
    dist.model
end

"""
    (f_oop, f_ip), augmented_sys, dvs, p = add_input_disturbance(sys, dist::DisturbanceModel, inputs = Any[])

Add a model of an unmeasured disturbance to `sys`. The disturbance model is an instance of [`DisturbanceModel`](@ref).

The generated dynamics functions `(f_oop, f_ip)` will preserve any state and dynamics associated with disturbance inputs, but the disturbance inputs themselves will not be included as inputs to the generated function. The use case for this is to generate dynamics for state observers that estimate the influence of unmeasured disturbances, and thus require state variables for the disturbance model, but without disturbance inputs since the disturbances are not available for measurement.

`dvs` will be the states of the simplified augmented system, consisting of the states of `sys` as well as the states of the disturbance model.

For MIMO systems, all inputs to the system has to be specified in the argument `inputs`

# Example

The example below builds a double-mass model and adds an integrating disturbance to the input

```julia
using ModelingToolkit
using ModelingToolkitStandardLibrary
using ModelingToolkitStandardLibrary.Mechanical.Rotational
using ModelingToolkitStandardLibrary.Blocks
t = ModelingToolkitStandardLibrary.Blocks.t

# Parameters
m1 = 1
m2 = 1
k = 1000 # Spring stiffness
c = 10   # Damping coefficient

@named inertia1 = Inertia(; J = m1)
@named inertia2 = Inertia(; J = m2)
@named spring = Spring(; c = k)
@named damper = Damper(; d = c)
@named torque = Torque(; use_support = false)

eqs = [connect(torque.flange, inertia1.flange_a)
       connect(inertia1.flange_b, spring.flange_a, damper.flange_a)
       connect(inertia2.flange_a, spring.flange_b, damper.flange_b)]
model = System(eqs, t; systems = [torque, inertia1, inertia2, spring, damper],
                  name = :model)
model = complete(model)
model_outputs = [model.inertia1.w, model.inertia2.w, model.inertia1.phi, model.inertia2.phi]

# Disturbance model
@named dmodel = Blocks.StateSpace([0.0], [1.0], [1.0], [0.0]) # An integrating disturbance
@named dist = ModelingToolkit.DisturbanceModel(model.torque.tau.u, dmodel)
(f_oop, f_ip), augmented_sys, dvs, p = ModelingToolkit.add_input_disturbance(model, dist)
```

`f_oop` will have an extra state corresponding to the integrator in the disturbance model. This state will not be affected by any input, but will affect the dynamics from where it enters, in this case it will affect additively from `model.torque.tau.u`.
"""
function add_input_disturbance(sys, dist::DisturbanceModel, inputs = Any[]; kwargs...)
    t = get_iv(sys)
    @variables d(t)=0 [disturbance = true]
    @variables u(t)=0 [input = true] # New system input
    dsys = get_disturbance_system(dist)

    if isempty(inputs)
        all_inputs = [u]
    else
        i = findfirst(isequal(dist.input), inputs)
        if i === nothing
            throw(ArgumentError("Input $(dist.input) indicated in the disturbance model was not found among inputs specified to add_input_disturbance"))
        end
        all_inputs = convert(Vector{Any}, copy(inputs))
        all_inputs[i] = u # The input where the disturbance acts is no longer an input, the new input is u
    end

    eqs = [dsys.input.u[1] ~ d
           dist.input ~ u + dsys.output.u[1]]
    augmented_sys = System(eqs, t, systems = [dsys], name = gensym(:outer))
    augmented_sys = extend(augmented_sys, sys)
    ssys = mtkcompile(augmented_sys, inputs = all_inputs, disturbance_inputs = [d])

    f, dvs, p, io_sys = generate_control_function(ssys, all_inputs,
        [d]; kwargs...)
    f, augmented_sys, dvs, p, io_sys
end
