using ModelingToolkitBase
using Symbolics: value
using DynamicQuantities
using Test

# Bounds
@variables u [bounds = (-1, 1)]
@test getbounds(u) == (-1, 1)
@test hasbounds(u)
@test ModelingToolkitBase.dump_variable_metadata(u).bounds == (-1, 1)

@variables y
@test !hasbounds(y)
@test !haskey(ModelingToolkitBase.dump_variable_metadata(y), :bounds)

@variables y[1:3]
@test !hasbounds(y)
@test getbounds(y)[1] == [-Inf, -Inf, -Inf]
@test getbounds(y)[2] == [Inf, Inf, Inf]
for i in eachindex(y)
    @test !hasbounds(y[i])
    b = getbounds(y[i])
    @test b[1] == -Inf && b[2] == Inf
end

@variables y[1:3] [bounds = (-ones(3), ones(3))]
@test hasbounds(y)
@test getbounds(y)[1] == -ones(3)
@test getbounds(y)[2] == ones(3)
for i in eachindex(y)
    @test hasbounds(y[i])
    b = getbounds(y[i])
    @test b[1] == -1.0 && b[2] == 1.0
end
@test getbounds(y[1:2])[1] == -ones(2)
@test getbounds(y[1:2])[2] == ones(2)

@variables y[1:2, 1:2] [bounds = (-ones(2, 2), [1.0 Inf; 2.0 3.0])]
@test hasbounds(y)
@test getbounds(y)[1] == [-1 -1; -1 -1]
@test getbounds(y)[2] == [1.0 Inf; 2.0 3.0]
for i in eachindex(y)
    @test hasbounds(y[i])
    b = getbounds(y[i])
    @test b[1] == -1 && b[2] == [1.0 Inf; 2.0 3.0][i]
end

@variables y[1:2] [bounds = (-Inf * ones(2), [1.0, Inf])]
@test hasbounds(y)
@test getbounds(y)[1] == [-Inf, -Inf]
@test getbounds(y)[2] == [1.0, Inf]
@test hasbounds(y[1])
@test getbounds(y[1]) == (-Inf, 1.0)
@test !hasbounds(y[2])
@test getbounds(y[2]) == (-Inf, Inf)

# Guess
@variables y [guess = 0]
@test getguess(y) == 0
@test hasguess(y) == true
@test ModelingToolkitBase.dump_variable_metadata(y).guess == 0

# Default
@variables y = 0
@test ModelingToolkitBase.getdefault(y) == 0
@test ModelingToolkitBase.hasdefault(y) == true
@test ModelingToolkitBase.dump_variable_metadata(y).default == 0

# Issue#2653
@variables y[1:3] [guess = ones(3)]
@test getguess(y) == ones(3)
@test hasguess(y) == true
@test ModelingToolkitBase.dump_variable_metadata(y).guess == ones(3)

for i in 1:3
    @test getguess(y[i]) == 1.0
    @test hasguess(y[i]) == true
    @test ModelingToolkitBase.dump_variable_metadata(y[i]).guess == 1.0
end

@variables y
@test hasguess(y) == false
@test !haskey(ModelingToolkitBase.dump_variable_metadata(y), :guess)

# Disturbance
@variables u [disturbance = true]
@test isdisturbance(u)
@test ModelingToolkitBase.dump_variable_metadata(u).disturbance

@variables y
@test !isdisturbance(y)
@test !haskey(ModelingToolkitBase.dump_variable_metadata(y), :disturbance)

# Tunable
@parameters u [tunable = true]
@test istunable(u)
@test ModelingToolkitBase.dump_variable_metadata(u).tunable

@parameters u2 [tunable = false]
@test !istunable(u2)
@test !ModelingToolkitBase.dump_variable_metadata(u2).tunable

@parameters y
@test istunable(y)
@test ModelingToolkitBase.dump_variable_metadata(y).tunable

# Distributions
struct FakeNormal end
d = FakeNormal()
@parameters u [dist = d]
@test hasdist(u)
@test getdist(u) == d
@test ModelingToolkitBase.dump_variable_metadata(u).dist == d

@parameters y
@test !hasdist(y)
@test !haskey(ModelingToolkitBase.dump_variable_metadata(y), :dist)

## System interface
@independent_variables t
Dₜ = Differential(t)
@variables x(t) = 0 [bounds = (-10, 10)] u(t) = 0 [input = true] y(t) = 0 [output = true]
@parameters T [bounds = (0, Inf)]
@parameters k [tunable = true, bounds = (0, Inf)]
@parameters k2 [tunable = false]
eqs = [
    Dₜ(x) ~ (-k2 * x + k * u) / T
    y ~ x
]
sys = System(eqs, t, name = :tunable_first_order)
unk_meta = ModelingToolkitBase.dump_unknowns(sys)
@test length(unk_meta) == 3
@test all(iszero, value(meta.default) for meta in unk_meta)
param_meta = ModelingToolkitBase.dump_parameters(sys)
@test length(param_meta) == 3
@test all(!haskey(meta, :default) for meta in param_meta)

p = tunable_parameters(sys)
sp = Set(p)
@test k ∈ sp
@test T ∈ sp
@test k2 ∉ sp
@test length(p) == 2

lb, ub = getbounds(p)
@test lb == [0, 0]
@test ub == [Inf, Inf]

b = getbounds(sys)
@test b[T] == (0, Inf)

b = getbounds(sys, unknowns(sys))
@test b[x] == (-10, 10)

p = tunable_parameters(sys, default = false)
sp = Set(p)
@test k ∈ sp
@test T ∉ sp
@test k2 ∉ sp
@test length(p) == 1

# Test that parameters appearing only in bounds are discovered
@independent_variables t2
Dₜ2 = Differential(t2)
@parameters Tₘ
@variables x2(t2) u2(t2) [bounds = (0, Tₘ)]
eqs2 = [Dₜ2(x2) ~ u2]
sys2 = System(eqs2, t2, name = :bounds_param_discovery)
@test Tₘ ∈ Set(parameters(sys2))

# Test that parameters in callable variable time arguments are discovered (x(tf) pattern)
@independent_variables t3
Dₜ3 = Differential(t3)
@parameters tf
@variables x3(..)
eqs3 = [Dₜ3(x3(t3)) ~ 1]
costs3 = [-x3(tf)]
sys3 = System(eqs3, t3; costs = costs3, name = :callable_param_discovery)
@test tf ∈ Set(parameters(sys3))

# Test that parameters in Integral domain bounds are discovered
@independent_variables t4
Dₜ4 = Differential(t4)
@parameters tf2
@variables x4(t4)
eqs4 = [Dₜ4(x4) ~ 1]
costs4 = [Symbolics.Integral(t4 in (0, tf2))(x4)]
sys4 = System(eqs4, t4; costs = costs4, name = :integral_param_discovery)
@test tf2 ∈ Set(parameters(sys4))

## Descriptions
@variables u [description = "This is my input"]
@test getdescription(u) == "This is my input"
@test hasdescription(u)
@test ModelingToolkitBase.dump_variable_metadata(u).desc == "This is my input"

@variables u
@test getdescription(u) == ""
@test !hasdescription(u)
@test !haskey(ModelingToolkitBase.dump_variable_metadata(u), :desc)

@independent_variables t
@variables u(t) [description = "A short description of u"]
@parameters p [description = "A description of p"]
@named sys = System([u ~ p], t)

@test_nowarn show(stdout, "text/plain", sys)

# Defaults, guesses overridden by system, parameter dependencies
@variables x(t) = 1.0 y(t) [guess = 1.0]
@parameters p = 2.0 q
@named sys = System(
    Equation[], t, [x, y], [p, q];
    bindings = [q => 2p], initial_conditions = Dict(x => 2.0, p => 3.0),
    guesses = Dict(y => 2.0)
)
sys = complete(sys)
unks_meta = ModelingToolkitBase.dump_unknowns(sys)
unks_meta = Dict([ModelingToolkitBase.getname(meta.var) => meta for meta in unks_meta])
@test value(unks_meta[:x].default) == 1.0
@test value(unks_meta[:x].initial_condition) == 2.0
@test value(unks_meta[:y].guess) == 2.0
params_meta = ModelingToolkitBase.dump_parameters(sys)
params_meta = Dict([ModelingToolkitBase.getname(meta.var) => meta for meta in params_meta])
@test value(params_meta[:p].default) == 2.0
@test value(params_meta[:p].initial_condition) == 3.0

# Connect
@variables x [connect = Flow]
@test hasconnect(x)
@test getconnect(x) == Flow
@test ModelingToolkitBase.dump_variable_metadata(x).connect == Flow
x = ModelingToolkitBase.setconnect(x, ModelingToolkitBase.Stream)
@test getconnect(x) == ModelingToolkitBase.Stream

struct BadConnect end
@test_throws Exception ModelingToolkitBase.setconnect(x, BadConnect)

# Unit
@variables x [unit = u"s"]
@test hasunit(x)
@test getunit(x) == u"s"
@test ModelingToolkitBase.dump_variable_metadata(x).unit == u"s"

# Misc data
@variables x [misc = [:good]]
@test hasmisc(x)
@test getmisc(x) == [:good]
x = ModelingToolkitBase.setmisc(x, "okay")
@test getmisc(x) == "okay"

# Variable Type
@variables x
@test ModelingToolkitBase.getvariabletype(x) == ModelingToolkitBase.VARIABLE
@test ModelingToolkitBase.dump_variable_metadata(x).variable_type == ModelingToolkitBase.VARIABLE
@test ModelingToolkitBase.dump_variable_metadata(x).variable_source == :variables
x = ModelingToolkitBase.toparam(x)
@test ModelingToolkitBase.getvariabletype(x) == ModelingToolkitBase.PARAMETER
@test ModelingToolkitBase.dump_variable_metadata(x).variable_source == :variables

@parameters y
@test ModelingToolkitBase.getvariabletype(y) == ModelingToolkitBase.PARAMETER

@brownians z
@test ModelingToolkitBase.getvariabletype(z) == ModelingToolkitBase.BROWNIAN

# Nominal
@variables x [nominal = 100.0]
@test getnominal(x) == 100.0
@test hasnominal(x)

@variables y
@test !hasnominal(y)
@test getnominal(y) == 1.0

y2 = setnominal(y, 50.0)
@test getnominal(y2) == 50.0
@test hasnominal(y2)

@variables z [nominal = 0.001]
@test getnominal(z) == 0.001

# Vector of variables
vals = getnominal([x, y, z])
@test vals == [100.0, 1.0, 0.001]

# Vector variables
@variables yv[1:3] [nominal = [100.0, 200.0, 300.0]]
@test hasnominal(yv)
@test getnominal(yv) == [100.0, 200.0, 300.0]
for i in 1:3
    @test hasnominal(yv[i])
end
@test getnominal(yv[1]) == 100.0
@test getnominal(yv[2]) == 200.0
@test getnominal(yv[3]) == 300.0

@variables wv[1:3]
@test !hasnominal(wv)
@test getnominal(wv[1]) == 1.0

# System-level accessor
@independent_variables t5
Dₜ5 = Differential(t5)
@variables x5(t5) [nominal = 100.0] y5(t5) z5(t5) [nominal = 0.001]
@parameters α5 = 1.5 β5 = 1.0
eqs5 = [Dₜ5(x5) ~ α5 * x5, Dₜ5(y5) ~ -β5 * y5, Dₜ5(z5) ~ x5 - z5]
sys5 = System(eqs5, t5, name = :nominal_test)
nv = getnominal(sys5, unknowns(sys5))
@test nv isa Dict
@test length(nv) == length(unknowns(sys5))
