using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CSV
using DataFrames
using Dates
using Plots
using BSON
using StatsBase
using Underscores
using ColorSchemes
using PrettyTables
colors = ColorSchemes.tableau_10.colors

## ===========================================================================

function consistent!(raw; nonnegative = false, positive = false)
    df = @_ raw |>
        disallowmissing(__) |>
        sort(__, :date) |>
        continuous_time!(__)

    if nonnegative || positive
        dat = @_ df |>
            select(__, Not(:date)) |>
            Array(__)
        nonnegative && @assert all( dat .>= 0 )
        positive && @assert all( dat .> 0 )
    end
    df
end

function Base.filter!(conditions::Vector{Pair{Symbol, Function}}, df::AbstractDataFrame)
    for condition in conditions
        filter!(condition, df)
    end
    df
end
# Base.filter!(condition::Pair{Symbol, Function}, df::AbstractDataFrame) =
#     filter!(condition, df)
Base.filter!(nothing, df::AbstractDataFrame) = df

function read_survey(region; fnames_survey, shift = 0, kwargs...)
    @_ fnames_survey[region] |>
        CSV.read(__, DataFrame) |>
        consistent!(__, positive = true) |>
        transform(__, :date =>( x -> x + Day(shift) )=> :date)
end

function read_mobil(region; fnames_mobil, kwargs...)
    @_ fnames_mobil[region] |>
        CSV.read(__, DataFrame) |>
        consistent!(__)
end

# read_predictors(region; predictors = All(), predictors_kwargs...)
# read_mobil(region; predictors_kwargs...)

function read_predictors(region; predictors = All(), kwargs...)
    @_ outerjoin(
            read_survey(region; kwargs...),
            read_mobil(region; kwargs...);
            on=:date) |>
        select(__, :date, predictors) |>
        sort(__, :date)
end

function read_epidata(region; conditions = nothing, kwargs...)
    @_ projectdir("data/epidata/epidata_region=$(region).csv") |>
        CSV.read(__, DataFrame) |>
        consistent!(__, nonnegative = true) |>
        filter!(conditions, __)
end

function relativechange!(xs::AbstractVector; standardize=false, kwargs...)
    xs[1] == 0. && (@warn "predictor data stream starts with 0, so we assume it is already normalized."; return nothing)
    xs ./= abs(xs[1])
    xs .-= xs[1]
    standardize && (xs ./= std(xs))
    return nothing
end

function relativechange!(df::AbstractDataFrame; kwargs...)
    for n in names(df, Not(:date))
        relativechange!(df[!,n], kwargs...)
    end
    df
end

function continuous_time!(df)
    @assert all( df.date[1]:Day(1):df.date[end] .== df.date )
    df
end

function clean_predictors(pred, epid; conditions, kwargs...)
    df = @_ pred |>
        leftjoin(epid, __; on = :date) |>
        sort(__, :date) |>
        filter!(conditions, __) |>
        select(__, names(pred)) |>
        disallowmissing(__) |>
        relativechange!(__) |>
        continuous_time!(__)
end

function join_predictors!(epid, pred)
    @_ epid |>
        continuous_time!(__) |>
        disallowmissing!(__) |>
        leftjoin(__,
            continuous_time!(pred);
            on=:date) |>
        sort(__, :date) |>
        coalesce.(__, 0.) |>
        disallowmissing!(__)
end

function stepindex(n, j0 = 1; step = 1, kwargs...)
    index = Vector{Int64}(undef, n)
    j = j0
    for i in 1:n
        index[i] = j
        i % step == 0 && ( j += 1 )
    end
    index
end

function earlyphase!(df, startdate)
    j = findlast( <(startdate), df.date)
    !isnothing(j) && (df[1:j,:rw] .= 1)
end

function latephase!(df)
    i = findlast(x->!ismissing(x), df[!, :rw])
    df[i+1:end, :rw] .= df[i, :rw]
end

function last_step(df)
    i = findlast(x->!ismissing(x), df[!, :rw])
    isnothing(i) && return 1
    return df[i, :rw]
end

function step!(df, sdate, edate; kwargs...)
    i0 = last_step(df) + 1
    is = findfirst(==(sdate), df.date)
    ie = findfirst(==(edate), df.date)
    n  = ie - is + 1
    df[is:ie, :rw] = stepindex(n, i0; kwargs...)
end

function wait!(df, sdate, edate)
    i0 = last_step(df)
    is = findfirst(==(sdate), df.date)
    ie = findfirst(==(edate), df.date)
    df[is:ie, :rw] .= i0
end

function randomwalk!(data, pred; semiparam = true, startdate = nothing, enddate = nothing, kwargs...)
    rw = DataFrame(:date => data.date, :rw => Vector{Union{Missing, Int64}}(missing, size(data,1)))
    startdate = isnothing(startdate) ? rw.date[1] : startdate
    enddate   = isnothing(enddate) ? data.date[end] : enddate

    earlyphase!(rw, startdate)
    if semiparam
        s = pred.date[1]
        e = pred.date[end]
        step!(rw, startdate, s; kwargs...)
        wait!(rw, s + Day(1), e)
        step!(rw, e + Day(1), enddate; kwargs...)
    else
        step!(rw, startdate, enddate; kwargs...)
    end
    latephase!(rw)

    @_ leftjoin(data, rw; on=:date) |>
        sort(__, :date)# |>
        # disallowmissing!(__)
end

function load_data(region, epidata_kwargs, predictor_kwargs, rw_kwargs)
    pred = read_predictors(region ; predictor_kwargs... )
    epid = read_epidata(region; epidata_kwargs...)
    pred = clean_predictors(pred, epid; predictor_kwargs... )
    data = @_ join_predictors!(epid, pred) |>
                randomwalk!(__, pred; rw_kwargs... ) |>
                continuous_time!(__) |>
                disallowmissing!(__)
    data[!,:region] .= region
    data
end
## ===========================================================================
regions = ["France", "Hungary", "UK", "Sweden", "Denmark", "Germany", "Italy", "USA"]
params = (
      epidata_kwargs = Dict(
        :conditions => [:date => x -> x >= Date("2020-08-01")]
        )

    , predictors_kwargs = Dict(
        :fnames_survey => Dict(region => projectdir("data/contacts/int/smoothed_contact_rates_region=$(region).csv") for region in regions),
        :fnames_mobil  => Dict(region => projectdir("data/mobility/mobility_country=$(region).csv") for region in regions),
        :shift         => -1,
        :predictors    => [:strangers],#[:strangers, :google, :apple],
        :conditions    => [
            :cases_per_100000    => x -> x >
            20,
            :fraction_vaccinated => x -> x < 5/100
            ]
        )

    , rw_kwargs = Dict(
        :step      => 7,
        :startdate => nothing,
        :enddate   => nothing,
        :semiparam => true
        )
    )

function load_data(regions::AbstractVector, args...)
    data = load_data(regions[1], params...)
    for region in regions[2:end]
        @show region
        data = vcat(data,
            load_data(region, args...)
            )
    end
    data
end

load_data(regions, params...)


function Dfp(arg0::Dfp)
    return Dfp((Dfp,), arg0)
end

function abs(obj::Dfp)
    return jcall(obj, "abs", Dfp, ())
end

function acos(obj::Dfp)
    return jcall(obj, "acos", Dfp, ())
end

function acosh(obj::Dfp)
    return jcall(obj, "acosh", Dfp, ())
end

function add(obj::Dfp, arg0::Dfp)
    return jcall(obj, "add", Dfp, (Dfp,), arg0)
end

function add(obj::Dfp, arg0::jdouble)
    return jcall(obj, "add", Dfp, (jdouble,), arg0)
end

function asin(obj::Dfp)
    return jcall(obj, "asin", Dfp, ())
end

function asinh(obj::Dfp)
    return jcall(obj, "asinh", Dfp, ())
end

function atan(obj::Dfp)
    return jcall(obj, "atan", Dfp, ())
end

function atan2(obj::Dfp, arg0::Dfp)
    return jcall(obj, "atan2", Dfp, (Dfp,), arg0)
end

function atanh(obj::Dfp)
    return jcall(obj, "atanh", Dfp, ())
end

function cbrt(obj::Dfp)
    return jcall(obj, "cbrt", Dfp, ())
end

function ceil(obj::Dfp)
    return jcall(obj, "ceil", Dfp, ())
end

function classify(obj::Dfp)
    return jcall(obj, "classify", jint, ())
end

function copy_sign(obj::Dfp, arg0::Dfp)
    return jcall(obj, "copySign", Dfp, (Dfp,), arg0)
end

function copy_sign(obj::Dfp, arg0::jdouble)
    return jcall(obj, "copySign", Dfp, (jdouble,), arg0)
end

function copysign(::Type{Dfp}, arg0::Dfp, arg1::Dfp)
    return jcall(Dfp, "copysign", Dfp, (Dfp, Dfp), arg0, arg1)
end

function cos(obj::Dfp)
    return jcall(obj, "cos", Dfp, ())
end

function cosh(obj::Dfp)
    return jcall(obj, "cosh", Dfp, ())
end

function divide(obj::Dfp, arg0::Dfp)
    return jcall(obj, "divide", Dfp, (Dfp,), arg0)
end

function divide(obj::Dfp, arg0::jdouble)
    return jcall(obj, "divide", Dfp, (jdouble,), arg0)
end

function divide(obj::Dfp, arg0::jint)
    return jcall(obj, "divide", Dfp, (jint,), arg0)
end

function dotrap(obj::Dfp, arg0::jint, arg1::JString, arg2::Dfp, arg3::Dfp)
    return jcall(obj, "dotrap", Dfp, (jint, JString, Dfp, Dfp), arg0, arg1, arg2, arg3)
end

function equals(obj::Dfp, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function exp(obj::Dfp)
    return jcall(obj, "exp", Dfp, ())
end

function expm1(obj::Dfp)
    return jcall(obj, "expm1", Dfp, ())
end

function floor(obj::Dfp)
    return jcall(obj, "floor", Dfp, ())
end

function get_class(obj::Object)
    return jcall(obj, "getClass", Class, ())
end

function get_exponent(obj::Dfp)
    return jcall(obj, "getExponent", jint, ())
end

function get_field(obj::Dfp)
    return jcall(obj, "getField", DfpField, ())
end

function get_one(obj::Dfp)
    return jcall(obj, "getOne", Dfp, ())
end

function get_radix_digits(obj::Dfp)
    return jcall(obj, "getRadixDigits", jint, ())
end

function get_real(obj::Dfp)
    return jcall(obj, "getReal", jdouble, ())
end

function get_two(obj::Dfp)
    return jcall(obj, "getTwo", Dfp, ())
end

function get_zero(obj::Dfp)
    return jcall(obj, "getZero", Dfp, ())
end

function greater_than(obj::Dfp, arg0::Dfp)
    return jcall(obj, "greaterThan", jboolean, (Dfp,), arg0)
end

function hash_code(obj::Dfp)
    return jcall(obj, "hashCode", jint, ())
end

function hypot(obj::Dfp, arg0::Dfp)
    return jcall(obj, "hypot", Dfp, (Dfp,), arg0)
end

function int_log10(obj::Dfp)
    return jcall(obj, "intLog10", jint, ())
end

function int_value(obj::Dfp)
    return jcall(obj, "intValue", jint, ())
end

function is_infinite(obj::Dfp)
    return jcall(obj, "isInfinite", jboolean, ())
end

function is_na_n(obj::Dfp)
    return jcall(obj, "isNaN", jboolean, ())
end

function is_zero(obj::Dfp)
    return jcall(obj, "isZero", jboolean, ())
end

function less_than(obj::Dfp, arg0::Dfp)
    return jcall(obj, "lessThan", jboolean, (Dfp,), arg0)
end

function linear_combination(obj::Dfp, arg0::Dfp, arg1::Dfp, arg2::Dfp, arg3::Dfp)
    return jcall(obj, "linearCombination", Dfp, (Dfp, Dfp, Dfp, Dfp), arg0, arg1, arg2, arg3)
end

function linear_combination(obj::Dfp, arg0::Dfp, arg1::Dfp, arg2::Dfp, arg3::Dfp, arg4::Dfp, arg5::Dfp)
    return jcall(obj, "linearCombination", Dfp, (Dfp, Dfp, Dfp, Dfp, Dfp, Dfp), arg0, arg1, arg2, arg3, arg4, arg5)
end

function linear_combination(obj::Dfp, arg0::Dfp, arg1::Dfp, arg2::Dfp, arg3::Dfp, arg4::Dfp, arg5::Dfp, arg6::Dfp, arg7::Dfp)
    return jcall(obj, "linearCombination", Dfp, (Dfp, Dfp, Dfp, Dfp, Dfp, Dfp, Dfp, Dfp), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function linear_combination(obj::Dfp, arg0::Vector{Dfp}, arg1::Vector{Dfp})
    return jcall(obj, "linearCombination", Dfp, (Vector{Dfp}, Vector{Dfp}), arg0, arg1)
end

function linear_combination(obj::Dfp, arg0::Vector{jdouble}, arg1::Vector{Dfp})
    return jcall(obj, "linearCombination", Dfp, (Vector{jdouble}, Vector{Dfp}), arg0, arg1)
end

function linear_combination(obj::Dfp, arg0::jdouble, arg1::Dfp, arg2::jdouble, arg3::Dfp)
    return jcall(obj, "linearCombination", Dfp, (jdouble, Dfp, jdouble, Dfp), arg0, arg1, arg2, arg3)
end

function linear_combination(obj::Dfp, arg0::jdouble, arg1::Dfp, arg2::jdouble, arg3::Dfp, arg4::jdouble, arg5::Dfp)
    return jcall(obj, "linearCombination", Dfp, (jdouble, Dfp, jdouble, Dfp, jdouble, Dfp), arg0, arg1, arg2, arg3, arg4, arg5)
end

function linear_combination(obj::Dfp, arg0::jdouble, arg1::Dfp, arg2::jdouble, arg3::Dfp, arg4::jdouble, arg5::Dfp, arg6::jdouble, arg7::Dfp)
    return jcall(obj, "linearCombination", Dfp, (jdouble, Dfp, jdouble, Dfp, jdouble, Dfp, jdouble, Dfp), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function log(obj::Dfp)
    return jcall(obj, "log", Dfp, ())
end

function log10(obj::Dfp)
    return jcall(obj, "log10", Dfp, ())
end

function log10k(obj::Dfp)
    return jcall(obj, "log10K", jint, ())
end

function log1p(obj::Dfp)
    return jcall(obj, "log1p", Dfp, ())
end

function multiply(obj::Dfp, arg0::Dfp)
    return jcall(obj, "multiply", Dfp, (Dfp,), arg0)
end

function multiply(obj::Dfp, arg0::jdouble)
    return jcall(obj, "multiply", Dfp, (jdouble,), arg0)
end

function multiply(obj::Dfp, arg0::jint)
    return jcall(obj, "multiply", Dfp, (jint,), arg0)
end

function negate(obj::Dfp)
    return jcall(obj, "negate", Dfp, ())
end

function negative_or_null(obj::Dfp)
    return jcall(obj, "negativeOrNull", jboolean, ())
end

function new_instance(obj::Dfp)
    return jcall(obj, "newInstance", Dfp, ())
end

function new_instance(obj::Dfp, arg0::Dfp)
    return jcall(obj, "newInstance", Dfp, (Dfp,), arg0)
end

function new_instance(obj::Dfp, arg0::DfpField, arg1::DfpField_RoundingMode)
    return jcall(obj, "newInstance", Dfp, (DfpField, DfpField_RoundingMode), arg0, arg1)
end

function new_instance(obj::Dfp, arg0::JString)
    return jcall(obj, "newInstance", Dfp, (JString,), arg0)
end

function new_instance(obj::Dfp, arg0::jbyte)
    return jcall(obj, "newInstance", Dfp, (jbyte,), arg0)
end

function new_instance(obj::Dfp, arg0::jbyte, arg1::jbyte)
    return jcall(obj, "newInstance", Dfp, (jbyte, jbyte), arg0, arg1)
end

function new_instance(obj::Dfp, arg0::jdouble)
    return jcall(obj, "newInstance", Dfp, (jdouble,), arg0)
end

function new_instance(obj::Dfp, arg0::jint)
    return jcall(obj, "newInstance", Dfp, (jint,), arg0)
end

function new_instance(obj::Dfp, arg0::jlong)
    return jcall(obj, "newInstance", Dfp, (jlong,), arg0)
end

function next_after(obj::Dfp, arg0::Dfp)
    return jcall(obj, "nextAfter", Dfp, (Dfp,), arg0)
end

function notify(obj::Object)
    return jcall(obj, "notify", void, ())
end

function notify_all(obj::Object)
    return jcall(obj, "notifyAll", void, ())
end

function positive_or_null(obj::Dfp)
    return jcall(obj, "positiveOrNull", jboolean, ())
end

function pow(obj::Dfp, arg0::Dfp)
    return jcall(obj, "pow", Dfp, (Dfp,), arg0)
end

function pow(obj::Dfp, arg0::jdouble)
    return jcall(obj, "pow", Dfp, (jdouble,), arg0)
end

function pow(obj::Dfp, arg0::jint)
    return jcall(obj, "pow", Dfp, (jint,), arg0)
end

function power10(obj::Dfp, arg0::jint)
    return jcall(obj, "power10", Dfp, (jint,), arg0)
end

function power10k(obj::Dfp, arg0::jint)
    return jcall(obj, "power10K", Dfp, (jint,), arg0)
end

function reciprocal(obj::Dfp)
    return jcall(obj, "reciprocal", Dfp, ())
end

function remainder(obj::Dfp, arg0::Dfp)
    return jcall(obj, "remainder", Dfp, (Dfp,), arg0)
end

function remainder(obj::Dfp, arg0::jdouble)
    return jcall(obj, "remainder", Dfp, (jdouble,), arg0)
end

function rint(obj::Dfp)
    return jcall(obj, "rint", Dfp, ())
end

function root_n(obj::Dfp, arg0::jint)
    return jcall(obj, "rootN", Dfp, (jint,), arg0)
end

function round(obj::RealFieldElement)
    return jcall(obj, "round", jlong, ())
end

function scalb(obj::Dfp, arg0::jint)
    return jcall(obj, "scalb", Dfp, (jint,), arg0)
end

function signum(obj::Dfp)
    return jcall(obj, "signum", Dfp, ())
end

function sin(obj::Dfp)
    return jcall(obj, "sin", Dfp, ())
end

function sin_cos(obj::CalculusFieldElement)
    return jcall(obj, "sinCos", FieldSinCos, ())
end

function sinh(obj::Dfp)
    return jcall(obj, "sinh", Dfp, ())
end

function sqrt(obj::Dfp)
    return jcall(obj, "sqrt", Dfp, ())
end

function strictly_negative(obj::Dfp)
    return jcall(obj, "strictlyNegative", jboolean, ())
end

function strictly_positive(obj::Dfp)
    return jcall(obj, "strictlyPositive", jboolean, ())
end

function subtract(obj::Dfp, arg0::Dfp)
    return jcall(obj, "subtract", Dfp, (Dfp,), arg0)
end

function subtract(obj::Dfp, arg0::jdouble)
    return jcall(obj, "subtract", Dfp, (jdouble,), arg0)
end

function tan(obj::Dfp)
    return jcall(obj, "tan", Dfp, ())
end

function tanh(obj::Dfp)
    return jcall(obj, "tanh", Dfp, ())
end

function to_degrees(obj::Dfp)
    return jcall(obj, "toDegrees", Dfp, ())
end

function to_double(obj::Dfp)
    return jcall(obj, "toDouble", jdouble, ())
end

function to_radians(obj::Dfp)
    return jcall(obj, "toRadians", Dfp, ())
end

function to_split_double(obj::Dfp)
    return jcall(obj, "toSplitDouble", Vector{jdouble}, ())
end

function to_string(obj::Dfp)
    return jcall(obj, "toString", JString, ())
end

function unequal(obj::Dfp, arg0::Dfp)
    return jcall(obj, "unequal", jboolean, (Dfp,), arg0)
end

function wait(obj::Object)
    return jcall(obj, "wait", void, ())
end

function wait(obj::Object, arg0::jlong)
    return jcall(obj, "wait", void, (jlong,), arg0)
end

function wait(obj::Object, arg0::jlong, arg1::jint)
    return jcall(obj, "wait", void, (jlong, jint), arg0, arg1)
end


function _check_constraint_bounds(bounds::PSI.ConstraintBounds, valid_bounds::NamedTuple)
    @test bounds.coefficient.min == valid_bounds.coefficient.min
    @test bounds.coefficient.max == valid_bounds.coefficient.max
    @test bounds.rhs.min == valid_bounds.rhs.min
    @test bounds.rhs.max == valid_bounds.rhs.max
end

function _check_variable_bounds(bounds::PSI.VariableBounds, valid_bounds::NamedTuple)
    @test bounds.bounds.min == valid_bounds.min
    @test bounds.bounds.max == valid_bounds.max
end

@testset "Test Numerical Stability of Constraints" begin
    template = get_thermal_dispatch_template_network()
    c_sys5 = PSB.build_system(PSITestSystems, "c_sys5")
    valid_bounds =
        (coefficient = (min = 1.0, max = 1.0), rhs = (min = 0.4, max = 9.930296584))
    model = OperationsProblem(template, c_sys5; optimizer = GLPK_optimizer)
    @test build!(model; output_dir = mktempdir(cleanup = true)) == PSI.BuildStatus.BUILT

    bounds = PSI.get_constraint_numerical_bounds(model; verbose = false)
    _check_constraint_bounds(bounds, valid_bounds)

    model_bounds = PSI.get_constraint_numerical_bounds(model; verbose = true)
    valid_model_bounds = Dict(
        :CopperPlateBalance => (
            coefficient = (min = 1.0, max = 1.0),
            rhs = (min = 6.434489705000001, max = 9.930296584),
        ),
        :P_lb__ThermalStandard__RangeConstraint =>
            (coefficient = (min = 1.0, max = 1.0), rhs = (min = Inf, max = -Inf)),
        :P_ub__ThermalStandard__RangeConstraint =>
            (coefficient = (min = 1.0, max = 1.0), rhs = (min = 0.4, max = 6.0)),
    )
    for (constriant_key, constriant_bounds) in model_bounds
        _check_constraint_bounds(constriant_bounds, valid_model_bounds[constriant_key])
    end
end

@testset "Test Numerical Stability of Variables" begin
    template = get_template_basic_uc_simulation()
    c_sys5 = PSB.build_system(PSITestSystems, "c_sys5_uc")
    valid_bounds = (min = 0.0, max = 6.0)
    model = OperationsProblem(template, c_sys5; optimizer = GLPK_optimizer)
    @test build!(model; output_dir = mktempdir(cleanup = true)) == PSI.BuildStatus.BUILT

    bounds = PSI.get_variable_numerical_bounds(model; verbose = false)
    _check_variable_bounds(bounds, valid_bounds)

    model_bounds = PSI.get_variable_numerical_bounds(model; verbose = true)
    valid_model_bounds = Dict(
        :stop__ThermalStandard => (min = 0.0, max = 1.0),
        :start__ThermalStandard => (min = 0.0, max = 1.0),
        :P__ThermalStandard => (min = 0.4, max = 6.0),
        :On__ThermalStandard => (min = 0.0, max = 1.0),
    )
    for (variable_key, variable_bounds) in model_bounds
        _check_variable_bounds(variable_bounds, valid_model_bounds[variable_key])
    end
end


# Functions for time-domain modeling and inversion using devito
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January, 2017
# Updated, December 2020, Mathias Louboutin, mlouboutin3@gatech.edu

#############################################################################
# Containers
include("Types/ModelStructure.jl")    # model container
include("Types/GeometryStructure.jl") # source or receiver setup, recording time and sampling
include("Types/OptionsStructure.jl")

#############################################################################
# Abstract vectors
include("Types/abstract.jl")
include("Types/broadcasting.jl")
include("Types/judiWavefield.jl") # dense RHS (wavefield)
include("Types/judiWeights.jl")    # Extended source weight vector
include("Types/judiVector.jl")    # Julia data container
include("Types/judiComposites.jl")    # A composite type to work with hcat/vcat of judi types

# Utility types
const SourceType{T} = Union{Vector{T}, judiMultiSourceVector{T}, PhysicalParameter{T}}
const dmType{T} = Union{Vector{T}, PhysicalParameter{T}}

#############################################################################
# Utils
include("Utils/auxiliaryFunctions.jl")

#############################################################################
# Linear operators
include("LinearOperators/basics.jl")
include("LinearOperators/lazy.jl")
include("LinearOperators/operators.jl")

#############################################################################
# PDE solvers
include("Modeling/distributed.jl") # Modeling functions utilities
include("Modeling/python_interface.jl")  # forward/adjoint linear/nonlinear modeling
include("Modeling/time_modeling_serial.jl")  # forward/adjoint linear/nonlinear modeling
include("Modeling/misfit_fg.jl")  # FWI/LSRTM objective function value and gradient
include("Modeling/twri_objective.jl")  # TWRI objective function value and gradient
include("Modeling/propagation.jl")

#############################################################################
# Preconditioners and optimization
include("Utils/seismic_preconditioners.jl")

#############################################################################
if VERSION>v"1.2"
  function (F::judiPropagator)(m::Model)
    Fl = deepcopy(F)
    Fl.model.n = m.n
    Fl.model.d = m.d
    Fl.model.o = m.o
    for (k, v) in m.params
      Fl.model.params[k] = v
    end
    Fl
  end

  function (F::judiPropagator)(;kwargs...)
    Fl = deepcopy(F)
    for (k, v) in kwargs
      k in keys(Fl.model.params) && Fl.model.params[k] .= v
    end
    Fl
  end
end


"""
GenX: An Configurable Capacity Expansion Model
Copyright (C) 2021,  Massachusetts Institute of Technology
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
A complete copy of the GNU General Public License v2 (GPLv2) is available
in LICENSE.txt.  Users uncompressing this from an archive may not have
received this license file.  If not, see <http://www.gnu.org/licenses/>.
"""

@doc raw"""
	storage_symmetric(EP::Model, inputs::Dict, Reserves::Int)

Sets up variables and constraints specific to storage resources with symmetric charge and discharge capacities. See ```storage()``` in ```storage.jl``` for description of constraints.
"""
function storage_symmetric(EP::Model, inputs::Dict, Reserves::Int)
	# Set up additional variables, constraints, and expressions associated with storage resources with symmetric charge & discharge capacity
	# (e.g. most electrochemical batteries that use same components for charge & discharge)
	# STOR = 1 corresponds to storage with distinct power and energy capacity decisions but symmetric charge/discharge power ratings

	println("Storage Resources with Symmetric Charge/Discharge Capacity Module")

	dfGen = inputs["dfGen"]

	T = inputs["T"]     # Number of time steps (hours)

	STOR_SYMMETRIC = inputs["STOR_SYMMETRIC"]

	### Constraints ###

	# Storage discharge and charge power (and reserve contribution) related constraints for symmetric storage resources:
	if Reserves == 1
		EP = storage_symmetric_reserves(EP, inputs)
	else
		@constraints(EP, begin
			# Maximum charging rate must be less than symmetric power rating
			[y in STOR_SYMMETRIC, t in 1:T], EP[:vCHARGE][y,t] <= EP[:eTotalCap][y]

			# Max simultaneous charge and discharge cannot be greater than capacity
			[y in STOR_SYMMETRIC, t in 1:T], EP[:vP][y,t]+EP[:vCHARGE][y,t] <= EP[:eTotalCap][y]
		end)
	end

	return EP
end

@doc raw"""
	storage_symmetric_reserves(EP::Model, inputs::Dict)

Sets up variables and constraints specific to storage resources with symmetric charge and discharge capacities when reserves are modeled. See ```storage()``` in ```storage.jl``` for description of constraints.
"""
function storage_symmetric_reserves(EP::Model, inputs::Dict)

	dfGen = inputs["dfGen"]
	T = inputs["T"]

	STOR_SYMMETRIC = inputs["STOR_SYMMETRIC"]

	STOR_SYM_REG_RSV = intersect(STOR_SYMMETRIC, inputs["REG"], inputs["RSV"]) # Set of symmetric storage resources with both REG and RSV reserves

	STOR_SYM_REG = intersect(STOR_SYMMETRIC, inputs["REG"]) # Set of symmetric storage resources with REG reserves
	STOR_SYM_RSV = intersect(STOR_SYMMETRIC, inputs["RSV"]) # Set of symmetric storage resources with RSV reserves

	STOR_SYM_NO_RES = setdiff(STOR_SYMMETRIC, STOR_SYM_REG, STOR_SYM_RSV) # Set of symmetric storage resources with no reserves

	STOR_SYM_REG_ONLY = setdiff(STOR_SYM_REG, STOR_SYM_RSV) # Set of symmetric storage resources only with REG reserves
	STOR_SYM_RSV_ONLY = setdiff(STOR_SYM_RSV, STOR_SYM_REG) # Set of symmetric storage resources only with RSV reserves

	if !isempty(STOR_SYM_REG_RSV)
		# Storage units charging can charge faster to provide reserves down and charge slower to provide reserves up
		@constraints(EP, begin
			# Maximum charging rate plus contribution to regulation down must be less than symmetric power rating
			[y in STOR_SYM_REG_RSV, t in 1:T], EP[:vCHARGE][y,t]+EP[:vREG_charge][y,t] <= EP[:eTotalCap][y]

			# Max simultaneous charge and discharge rates cannot be greater than symmetric charge/discharge capacity
			[y in STOR_SYM_REG_RSV, t in 1:T], EP[:vP][y,t]+EP[:vREG_discharge][y,t]+EP[:vRSV_discharge][y,t]+EP[:vCHARGE][y,t]+EP[:vREG_charge][y,t] <= EP[:eTotalCap][y]
		end)
	end

	if !isempty(STOR_SYM_REG_ONLY)
		# Storage units charging can charge faster to provide reserves down and charge slower to provide reserves up
		@constraints(EP, begin
			# Maximum charging rate plus contribution to regulation down must be less than symmetric power rating
			[y in STOR_SYM_REG_ONLY, t in 1:T], EP[:vCHARGE][y,t]+EP[:vREG_charge][y,t] <= EP[:eTotalCap][y]

			# Max simultaneous charge and discharge rates cannot be greater than symmetric charge/discharge capacity
			[y in STOR_SYM_REG_ONLY, t in 1:T], EP[:vP][y,t]+EP[:vREG_discharge][y,t]+EP[:vCHARGE][y,t]+EP[:vREG_charge][y,t] <= EP[:eTotalCap][y]
		end)
	end

	if !isempty(STOR_SYM_RSV_ONLY)
		# Storage units charging can charge faster to provide reserves down and charge slower to provide reserves up
		@constraints(EP, begin
			# Maximum charging rate must be less than symmetric power rating
			[y in STOR_SYM_RSV_ONLY, t in 1:T], EP[:vCHARGE][y,t] <= EP[:eTotalCap][y]

			# Max simultaneous charge and discharge rates cannot be greater than symmetric charge/discharge capacity
			[y in STOR_SYM_RSV_ONLY, t in 1:T], EP[:vP][y,t]+EP[:vRSV_discharge][y,t]+EP[:vCHARGE][y,t] <= EP[:eTotalCap][y]
		end)
	end

	if !isempty(STOR_SYM_NO_RES)
		@constraints(EP, begin
			# Maximum charging rate must be less than symmetric power rating
			[y in STOR_SYM_NO_RES, t in 1:T], EP[:vCHARGE][y,t] <= EP[:eTotalCap][y]

			# Max simultaneous charge and discharge cannot be greater than capacity
			[y in STOR_SYM_NO_RES, t in 1:T], EP[:vP][y,t]+EP[:vCHARGE][y,t] <= EP[:eTotalCap][y]
		end)
	end

	return EP
end


module Simulation

export PlanktonSimulation
export update!, vel_copy!

using CUDA
using StructArrays
using LinearAlgebra

using PlanktonIndividuals.Architectures
using PlanktonIndividuals.Grids
using PlanktonIndividuals.Fields
using PlanktonIndividuals.Biogeochemistry
using PlanktonIndividuals.Parameters
using PlanktonIndividuals.Plankton
using PlanktonIndividuals.Diagnostics
using PlanktonIndividuals.Model
using PlanktonIndividuals.Output

using PlanktonIndividuals: AbstractMode, CarbonMode, QuotaMode, MacroMolecularMode

import Base: show

include("utils.jl")
include("simulations.jl")
include("update.jl")

end


@testset "SE3PDController" begin
    val = Valkyrie()
    mechanism = val.mechanism
    state = MechanismState(mechanism)
    rng = MersenneTwister(1)
    base = rand(rng, bodies(mechanism))
    body = rand(rng, bodies(mechanism))
    angular = let angle = π / 2, axis = SVector(1.0, 0.0, 0.0), y0 = one(Quat), yf = Quat(AngleAxis(angle, axis...))
        Interpolated(0.0, 1.0, y0, yf)
    end
    linear = Interpolated(0.0, 1.0, SVector(0.0, 1.0, 2.0), SVector(2.0, 3.0, 4.0))
    frame = default_frame(body)
    gains = SE3PDGains(FramePDGains(frame, PDGains(100.0, 20.0)), FramePDGains(frame, PDGains(1000.0, 200.0)))
    trajectory = SE3Trajectory(default_frame(body), default_frame(base), angular, linear)
    weight = Diagonal(vcat(zeros(3), fill(10.0, 3)))
    controller = SE3PDController(BodyID(base), BodyID(body), trajectory, weight, gains)
    controller(0.5, state)

    # FIXME: proper tests
end

@testset "parameterized contacts" begin
    # Construct a mechanism consisting of a single body which can
    # rotate about its origin
    world = RigidBody{Float64}("world")
    mechanism = Mechanism(world, gravity=SVector(0, 0, -9.81))
    frame = CartesianFrame3D("body")
    inertia = SpatialInertia(frame, SDiagonal(1., 1, 1), SVector(0., 0, 0), 10.0)
    body = RigidBody(inertia)
    joint = Joint("rx", Revolute(SVector(1., 0, 0)))
    attach!(mechanism, world, body, joint)

    contactmodel = SoftContactModel(hunt_crossley_hertz(k = 500e3), ViscoelasticCoulombModel(0.8, 20e3, 100.))
    add_contact_point!(body, Contact.ContactPoint(Point3D(default_frame(body), 0., 0, 0), contactmodel))

    # Make the contact normal a parameter, so that it updates its representation
    # in body frame to match a fixed orientation in world frame
    state = MechanismState(mechanism)
    model = Parametron.Model(defaultoptimizer())
    position = Point3D(default_frame(body), 0., 0, 0)
    μ = 1.0
    normal = let state = state, body = body
        Parameter(model) do
            transform(state, FreeVector3D(root_frame(state.mechanism), 0., 0, 1), default_frame(body))
        end
    end
    controller_contact = QPControl.ContactPoint{4}(position, normal, μ, state, model)
    controller_contact.maxnormalforce[] = 1e3
    @test QPControl.isenabled(controller_contact)

    # Constrain the contact force to be non-zero for testing
    @constraint(model, controller_contact.force_local.v == [0.0, 0.0, 1.0])

    for θ in range(-π, stop=π, length=10)
        set_configuration!(state, [θ])
        solve!(model)
        # Sanity check our constraint
        @test value.(Ref(model), controller_contact.force_local.v) ≈ [0.0, 0.0, 1.0]

        # No matter how we rotate the robot, the contact-aligned frame will still
        # be aligned to the contact normal, which is fixed in world frame. So the
        # linear component of the contact wrench in world frame will always be
        # along [0, 0, 1].
        @test value.(Ref(model), linear(controller_contact.wrench_world)) ≈ [0.0, 0.0, 1.0]
    end
end

@testset "fixed base joint space control, constrained = $constrained" for constrained in [true, false]
    Random.seed!(42)
    mechanism = rand_tree_mechanism(Float64, Prismatic{Float64}, Revolute{Float64}, Revolute{Float64})
    N = 4
    controller = MomentumBasedController{N}(mechanism, defaultoptimizer())
    tasks = Dict{Joint{Float64}, JointAccelerationTask}()
    for joint in tree_joints(mechanism)
        task = JointAccelerationTask(joint)
        tasks[joint] = task
        if constrained
            addtask!(controller, task)
        else
            weight = 1.0
            addtask!(controller, task, weight)
        end
        setdesired!(task, rand(num_velocities(joint)))
    end
    state = MechanismState(mechanism)
    rand!(state)
    τ = similar(velocity(state))
    controller(τ, 0.0, state)
    allocs = @allocated controller(τ, 0.0, state)
    @test allocs == 0

    result = DynamicsResult(mechanism)
    dynamics!(result, state, τ)
    for joint in tree_joints(mechanism)
        @test result.v̇[joint] ≈ tasks[joint].desired atol = 1e-10
    end
end

"""
Automatically load contact points from each body in the mechanism and add them
to the controller. If `parametric_contact_surface=true`, then the `normal` and
`μ` associated with each contact will be Parameters set to random, but fixed,
values.
"""
function set_up_valkyrie_contacts!(controller::MomentumBasedController; parametric_contact_surface=false)
    valmechanism = controller.state.mechanism
    for body in bodies(valmechanism)
        for point in RBD.contact_points(body)
            position = RBD.Contact.location(point)
            if parametric_contact_surface
                normal = let frame = position.frame, direction = normalize(randn(SVector{3}))
                    Parameter(controller.qpmodel) do
                        FreeVector3D(frame, direction)
                    end
                end
                μ = let μ = rand()
                    Parameter(() -> μ, controller.qpmodel)
                end
            else
                normal = FreeVector3D(position.frame, 0.0, 0.0, 1.0)
                μ = point.model.friction.μ
            end
            addcontact!(controller, body, position, normal, μ)
        end
    end
end

@testset "zero velocity free fall" begin
    Random.seed!(5354)
    val = Valkyrie()
    mechanism = val.mechanism
    floatingjoint = val.basejoint
    N = 4
    controller = MomentumBasedController{N}(mechanism, defaultoptimizer(), floatingjoint=floatingjoint)
    set_up_valkyrie_contacts!(controller)
    state = MechanismState(mechanism)
    τ = similar(velocity(state))

    zero!(state)
    rand_configuration!(state)
    for joint in tree_joints(mechanism)
        joint == floatingjoint && continue
        regularize!(controller, joint, 1.0)
    end
    controller(τ, 0., state)
    result = DynamicsResult(mechanism)
    accels = result.accelerations
    RBD.spatial_accelerations!(accels, state, controller.result.v̇)
    for joint in tree_joints(mechanism)
        if joint == floatingjoint
            baseaccel = relative_acceleration(accels, val.pelvis, root_body(mechanism))
            baseaccel = transform(state, baseaccel, frame_after(floatingjoint))
            angularaccel = FreeVector3D(baseaccel.frame, baseaccel.angular)
            @test isapprox(angularaccel, FreeVector3D(frame_after(floatingjoint), zeros(SVector{3})), atol = 1e-4)
            linearaccel = FreeVector3D(baseaccel.frame, baseaccel.linear)
            linearaccel = transform_to_root(state, linearaccel.frame) * linearaccel
            @test isapprox(linearaccel, mechanism.gravitational_acceleration; atol = 1e-4)
        else
            v̇joint = controller.result.v̇[velocity_range(state, joint)]
            @test isapprox(v̇joint, zeros(num_velocities(joint)); atol = 1e-4)
        end
    end
    allocs = @allocated controller(τ, 0., state)
    @test allocs == 0
end

const MAX_NORMAL_FORCE_FIXME = 1e9

@testset "achievable momentum rate" begin
    Random.seed!(533454)
    val = Valkyrie()
    mechanism = val.mechanism
    floatingjoint = val.basejoint
    state = MechanismState(mechanism)
    τ = similar(velocity(state))

    N = 4
    controller = MomentumBasedController{N}(mechanism, defaultoptimizer(), floatingjoint=floatingjoint)

    set_up_valkyrie_contacts!(controller; parametric_contact_surface=true)
    ḣtask = MomentumRateTask(mechanism, centroidal_frame(controller))
    addtask!(controller, ḣtask) #, 1.0)

    for joint in tree_joints(mechanism)
        regularize!(controller, joint, 1e-6)
    end

    for p in range(0., stop=1., length=5)
        rand!(state)
        com = center_of_mass(state)
        centroidal_to_world = Transform3D(centroidal_frame(controller), com.frame, com.v)
        world_to_centroidal = inv(centroidal_to_world)

        # set random active contacts and random achievable wrench
        fg = world_to_centroidal * (mass(mechanism) * mechanism.gravitational_acceleration)
        ḣdes = Wrench(zero(fg), fg)
        for body in keys(controller.contacts)
            for contact in controller.contacts[body]
                active = rand() < p
                if active
                    normal = Parametron.evalarg(contact.normal)
                    μ = Parametron.evalarg(contact.μ)
                    contact.weight[] = 1e-6
                    contact.maxnormalforce[] = MAX_NORMAL_FORCE_FIXME
                    fnormal = 50. * rand()
                    μreduced = sqrt(2) / 2 * μ # due to polyhedral inner approximation; assumes 4 basis vectors or more
                    ftangential = μreduced * fnormal * rand() * cross(normal, FreeVector3D(normal.frame, normalize(randn(SVector{3}))))
                    f = fnormal * normal + ftangential
                    @assert isapprox(ftangential ⋅ normal, 0., atol = 1e-12)
                    @assert norm(f - (normal ⋅ f) * normal) ≤ μreduced * (normal ⋅ f)
                    wrench = Wrench(contact.position, f)
                    ḣdes += transform(wrench, world_to_centroidal * transform_to_root(state, wrench.frame))
                else
                    disable!(contact)
                end
            end
        end
        setdesired!(ḣtask, ḣdes)

        controller(τ, 0., state)

        # Ensure that desired momentum rate is achieved.
        ḣ = Wrench(momentum_matrix(state), controller.result.v̇) + momentum_rate_bias(state)
        ḣ = transform(ḣ, world_to_centroidal)
        @test isapprox(ḣdes, ḣ; atol = 1e-3)

        allocs = @allocated controller(τ, 0., state)
        @test allocs == 0
    end
end

@enum TaskMode begin
    CONSTRAINT
    OBJECTIVE_TERM_SCALAR_WEIGHT
    OBJECTIVE_TERM_SCALAR_WEIGHT_PARAMETER
    OBJECTIVE_TERM_MATRIX_WEIGHT
    OBJECTIVE_TERM_MATRIX_WEIGHT_PARAMETER
end

@testset "spatial acceleration, mode = $mode" for mode in instances(TaskMode)
    Random.seed!(533)
    val = Valkyrie()
    mechanism = val.mechanism
    floatingjoint = val.basejoint
    state = MechanismState(mechanism)
    τ = similar(velocity(state))

    rand!(state)
    N = 4
    controller = MomentumBasedController{N}(mechanism, defaultoptimizer(), floatingjoint=floatingjoint)
    body = val.feet[left]
    base = val.palms[right]
    frame = default_frame(base)
    task = SpatialAccelerationTask(mechanism, path(mechanism, base, body), frame=frame)

    result = DynamicsResult(mechanism)
    accels = result.accelerations
    if mode == CONSTRAINT
        addtask!(controller, task)
        regularize!.(Ref(controller), tree_joints(mechanism), 1.0)
    elseif mode == OBJECTIVE_TERM_SCALAR_WEIGHT
        addtask!(controller, task, 1.0)
    elseif mode == OBJECTIVE_TERM_SCALAR_WEIGHT_PARAMETER
        weight = Parameter(controller.qpmodel, val=1.0)
        addtask!(controller, task, weight)
    elseif mode == OBJECTIVE_TERM_MATRIX_WEIGHT
        addtask!(controller, task, Matrix(I, 6, 6))
    elseif mode == OBJECTIVE_TERM_MATRIX_WEIGHT_PARAMETER
        weight = Parameter(controller.qpmodel, val=Matrix(I, 6, 6))
        addtask!(controller, task, weight)
    else
        error()
    end
    desiredaccel = rand(SpatialAcceleration{Float64}, default_frame(body), default_frame(base), frame)
    setdesired!(task, desiredaccel)
    controller(τ, 0., state)
    v̇ = controller.result.v̇
    spatial_accelerations!(accels, state, v̇)
    accel = relative_acceleration(accels, body, base)
    accel = transform(state, accel, frame)
    @test isapprox(accel, desiredaccel, atol = 1e-8)

    allocs = @allocated controller(τ, 0., state)
    @test allocs == 0
end


isempty(line::AbstractString) = !ismatch(r"[^\s\t\n\r]+", line)

iscomment(line::AbstractString) = ismatch(r"^#", line)

istitle(line::AbstractString) = ismatch(r"^\*\*\*(.*)MODEL (STATES|PARAMETERS|ODES)", line)

is_ode_name(name::AbstractString) = ismatch(r"^d/dt\((.*)\)$", name)

function get_title_type(line::AbstractString)
  if ismatch(r"^\*\*\*(.*)MODEL STATES(.*)", line)
    return "state"
  elseif ismatch(r"^\*\*\*(.*)MODEL PARAMETERS(.*)", line)
    return "parameter"
  elseif ismatch(r"^\*\*\*(.*)MODEL ODES(.*)", line)
    return "ode"
  else
    throw(ParseError("Unknown title \""*line*"\" type in model specification"))
  end
end

function parse_line(line::AbstractString, linetype::AbstractString, model::OdeModel)
  if linetype == "state"
    nTokens = parse_state_line(line, model.states)
  elseif linetype == "parameter"
    nTokens = parse_parameter_line(line, model.parameters)
  elseif linetype == "ode"
    nTokens = parse_ode_line(line, model.odes)
  else
    throw(ParseError("Unknown title \""*line*"\" type in model specification"))
  end

  return nTokens
end

function parse_state_line(line::AbstractString, dict::Dict{AbstractString, Float64})
  tokens = split(line, '=')
  nTokens = length(tokens)

  if nTokens != 2
    throw(ParseError("Line \""*line*"\" has wrong format"))
  end

  name, value = [strip(i) for i in tokens]

  haskey(dict, name) ? throw(KeyError("\""*name*"\" specified more than once")) : (dict[name] = Float64(parse(value)))

  return nTokens
end

function parse_parameter_line(line::AbstractString, dict::Dict{AbstractString, Float64})
  tokens = split(line, '=')
  nTokens = length(tokens)

  if nTokens == 1
    dict[strip(tokens)] = NaN
  elseif nTokens == 2
    name, value = [strip(i) for i in tokens]

    haskey(dict, name) ? throw(KeyError("\""*name*"\" specified more than once")) : (dict[name] = Float64(parse(value)))
  else
    throw(ParseError("Line \""*line*"\" has wrong format"))
  end

  return nTokens
end

function parse_ode_line(line::AbstractString, dict::Dict{AbstractString, NSE})
  tokens = split(line, '=')
  nTokens = length(tokens)

  if nTokens != 2
    throw(ParseError("Line \""*line*"\" has wrong format"))
  end

  name, value = [strip(i) for i in tokens]

  is_ode_name(name) ? (name = name[6:end-1]) : throw(KeyError("Time derivative \""*name*"\" has wrong format"))

  haskey(dict, name) ? throw(KeyError("\""*name*"\" specified more than once")) : (dict[name] = parse(value))

  return nTokens
end

function parse_model(file::AbstractString)
  odeModel = OdeModel(file, Dict{AbstractString, Float64}(), Dict{AbstractString, Float64}(), Dict{AbstractString, NSE}())

  s = open(file)

  while !eof(s)
    line = readline(s)
    if isempty(line) || iscomment(line)
      continue
    elseif istitle(line)
      title = get_title_type(line)
    else
      parse_line(line, title, odeModel)
    end
  end

  close(s)

  if sort(collect(keys(odeModel.states))) != sort(collect(keys(odeModel.odes)))
    throw(KeyError("States and their time derivatives in the ODE model are inconsistent"))
  end

  return odeModel
end


GridCardType = CardType{:grid}()
const GridCard = NastranCard{GridCardType}

type GRID <: GridCard
    id::Int64
    csys_id::Int64
    x::Float64
    y::Float64
    z::Float64
    ocsys_id::Int64
end
function convert(::Type{GRID},card)
    id = card[2]::Int64
    csys_id = card[3] == "" ? 0 : card[3]::Int64
    x = card[4] == "" ? 0.0 : card[4]::Float64
    y = card[5] == "" ? 0.0 : card[5]::Float64
    z = card[6] == "" ? 0.0 : card[6]::Float64
    ocsys_id = card[7] == "" ? 0 : card[7]::Int64
    GRID(id,csys_id,x,y,z,ocsys_id)
end

type SPOINT <: GridCard
    id::Int64
end
function convert(::Type{SPOINT},card)
    id = card[2]::Int64
    SPOINT(id)
end

type EPOINT <: GridCard
    id::Int64
end
function convert(::Type{EPOINT},card)
    id = card[2]::Int64
    EPOINT(id)
end


module RATTLE

export RattleIntegrator
include("./types.jl")
include("./main.jl")

end


"""
    initial_conditions(model)

Get initial conditions at all degrees of freedom.
"""
initial_conditions(model) =
    mapreduce((ρ, p) -> fill(ρ, size(p, 2)), vcat, model.ρ, model.mesh.points)


#sigmoid number conversions.

#let's put these functions here. We'll move them later.
import Base.convert

const IEEEFloat = Union{Float16, Float32, Float64}

function convert{N, ES, mode, I <: Signed}(T::Type{Sigmoid{N, ES, mode}}, int::I)
  #warn("conversion from integers not yet properly supported! $int")
  #throw(ErrorException("halp"))
  convert(T, convert(Float64, int))
end

@generated function convert{F <: IEEEFloat, N, ES, mode}(::Type{F}, x::Sigmoid{N, ES, mode})
  FInt  = Dict(Float16 => UInt16, Float32 => UInt32, Float64 => UInt64)[F]
  fbits = Dict(Float16 => 16    , Float32 => 32,     Float64 => 64)[F]
  ebits = Dict(Float16 => 5     , Float32 => 8,      Float64 => 11)[F]
  bias = (1 << (ebits - 1)) -1
  quote
    #check for two dropout values.
    (reinterpret((@UInt), x) == zero(@UInt)) && return zero(F)
    isfinite(x) || return F(Inf)

    #magically generate the subproperties of x.
    @breakdown x

    FP_sign = (($FInt)(x_sgn) << ($fbits - 1))

    #calculate the biased exponent for conversion to IEEE FP
    FP_biased_exponent::$FInt = ($bias + x_exp) << ($fbits - $ebits - 1)

    #convert to UInt64 and move the fraction even more (or back, if we started as a 64-bit int.)
    if $fbits < __BITS
      FP_fraction = ($FInt)(x_frc >> (($ebits + 1) + (__BITS - $fbits)))
    else
      FP_fraction = ($FInt)(x_frc) >> (($ebits + 1) - (($fbits) - __BITS))
    end

    #put the parts together and reassign as a Floating point.
    reinterpret(F, FP_sign | FP_biased_exponent | FP_fraction)
  end
end

doc"""
  fptrip(x)

  takes an IEEE floating point and returns the triplet (sign::Bool, exponent::Int, fraction::@UInt).
  this will then be used to generate the fractional value.

  NB: Doesn't work for Inf and zero values.
"""
@generated function fptrip{F <: IEEEFloat}(x::F)
  FInt  = Dict(Float16 => UInt16, Float32 => UInt32, Float64 => UInt64)[F]
  fbits = Dict(Float16 => 16    , Float32 => 32,     Float64 => 64)[F]
  ebits = Dict(Float16 => 5     , Float32 => 8,      Float64 => 11)[F]
  bias = (1 << (ebits - 1)) -1
  quote
    #first grab the sign bit.
    sign = (x < 0)
    #flip the sign.
    x = sign ? -x : x
    intform = reinterpret($FInt, x) & ((zero($FInt) - 1) >> 1)
    exponent = Int(intform >> ($fbits - $ebits - 1)) - $bias
    if $fbits < __BITS
      #convert to left-shifted form of the fraction, then expand, then shift more.
      fraction = ((@UInt)(intform << ($ebits + 1))) << (__BITS - $fbits)
    else
      #convert to the left-shifted but then go all the way over, then convert to @UInt
      fraction = (@UInt)(intform << ($ebits + 1) - (__BITS - $fbits))
    end

    (sign, exponent, fraction)
  end
end

function convert{N, ES, mode, F <: IEEEFloat}(::Type{Sigmoid{N, ES, mode}}, f::F)
  #handle the three corner cases of NaN, infinity and zero.
  isnan(f) && throw(NaNError(convert, [Sigmoid{N, ES, mode}, f]))
  isfinite(f) || return reinterpret(Sigmoid{N, ES, mode}, @signbit)
  (f == zero(F)) && return reinterpret(Sigmoid{N, ES, mode}, zero(@UInt))
  #retrieve the floating point triplet.
  __round(build_numeric(Sigmoid{N, ES, mode}, fptrip(f)...))
end

@generated function (::Type{Sigmoid{N, ES, mode}}){N, ES, mode, UI <: Unsigned}(i::UI)
  #conversion from unsigned integer will be interpreted as a desire to convert
  #logical representation; conversion from signed integer will be interpreted
  #as a desire to convert the semantic value.

  #first check to see that the number of bits in the sigmoid is less than the
  #number of bits in the unsigned integer.
  int_length = sizeof(i) * 8
  (N <= int_length) || throw(ArgumentError("insufficient bits in source integer to represent the sigmoid number."))
  quote
    #assume that the integer representation is right-aligned.  We're going to want
    #to move that to the LEFT.  First convert the value to the appropriate integer
    #type based on the global integer size, then shift right the appropriate number
    #of bits, then reinterpret it as a sigmoid number.
    reinterpret(Sigmoid{N, ES, mode}, (@UInt)(i) << (__BITS - N))
  end
end

function build_numeric{N, ES, mode}(::Type{Sigmoid{N, ES, mode}}, sign, exponent, fraction)
  if exponent < 0
    (regime, subexponent) = divrem(exponent, 1 << ES)
    if (subexponent != 0)
      regime -= 1
      subexponent += 1 << ES
    end

    fshift = -regime + 2 + ES

    body = (one(@UInt) << ES) | subexponent

    body <<= (__BITS + regime - 2 - ES)
  else
    (regime, subexponent) = divrem(exponent, 1 << ES)

    #set the prefix.  That's 2^regime - 1
    body = (((one(@UInt) << (regime + 1)) - 1) << (ES + 1)) + subexponent

    #shift the prefix over
    body <<= (__BITS - regime - 3 - ES)

    fshift = regime + 3 + ES
  end

  absval = body | (fraction >> fshift)

  reinterpret(Sigmoid{N, ES, mode}, sign ? -absval : absval)
end

#build_numeric is much simpler when ES = 0
function build_numeric{N, mode}(::Type{Sigmoid{N, 0, mode}}, sign, exponent, fraction)
  if exponent < 0
    fshift = -exponent + 2
    body = one(@UInt) << (__BITS + exponent - 2)
  else
    #set the prefix.  That's 2^exponent - 1
    prefix = (one(@UInt) << (exponent + 1)) - 1
    #shift the prefix over
    body = prefix << (__BITS - exponent - 2)
    fshift = exponent + 3
  end

  absval = body | (fraction >> fshift)

  reinterpret(Sigmoid{N, 0, mode}, sign ? -absval : absval)
end

#build_arithmetic is restricted to ES == 0 sigmoids.
function build_arithmetic{N, mode}(::Type{Sigmoid{N, 0, mode}}, sign, exponent, fraction)
  normal = (fraction & @signbit) != 0
  #check if it's denormal.  If it is, then we don't shift.
  fshift = normal * (exponent + 1) + 1
  #set the prefix.  That's 2^exponent - 1
  prefix = (one(@UInt) << (exponent + 1)) - 1

  body = normal * (prefix << (__BITS - exponent - 2))

  absval = body | ((fraction & ((@signbit) - 1)) >> fshift)

  reinterpret(Sigmoid{N, 0, mode}, sign ? -absval : absval)
end

Base.promote_rule{T<:Sigmoid}(::Type{T}, ::Type{Int64}) = T


using Test, MathJaxRenderer, VisualRegressionTests

# Test cases from Latexify.jl.

keywords = Dict(
    "small" => (
        x_zoom = 4,
        y_zoom = 4,
        background_color = :orange,
    ),
    "large" => (
        x_zoom = 2,
        y_zoom = 2,
        background_color = :lightblue,
    ),
)

@testset "MathJaxRenderer" begin
    @testset "Integration" begin
        tex = joinpath(@__DIR__, "data", "tex")
        for file in readdir(tex)
            name = first(splitext(file))
            file = joinpath(tex, file)
            math = read(file, Math)
            kws = keywords[name]
            @test !isempty(sprint(io -> show(io, MIME("image/svg+xml"), math; kws...)))
            @test !isempty(sprint(io -> show(io, MIME("image/png"), math; kws...)))
            @test !isempty(sprint(io -> show(io, MIME("application/pdf"), math; kws...)))
            @test !isempty(sprint(io -> show(io, MIME("application/postscript"), math; kws...)))
        end
    end
    @testset "Visuals" begin
        testfun = (filename) -> write(filename, Math("\\frac{1}{1 + x}"); x_zoom = 5, y_zoom = 5)
        @visualtest testfun joinpath(@__DIR__, "data/references/fraction.png") popup = false tol = 0.2
        testfun = (filename) -> write(filename, Math("\\begin{align}\n3.29e+07 =& 1.23e+00 =& P_{1} \\\\\n\\frac{x}{y} =& 1.00e+10 =& 1.29e+03\n\\end{align}"); x_zoom = 5, y_zoom = 5)
        @visualtest testfun joinpath(@__DIR__, "data/references/align.png") popup = false tol = 0.2
    end
end


function numb!(r)
    function change(side, num)
        return HorizonSide((Int(side) + num) % 4)
    end

    function go_around(side)
        num = 0
        while isborder(r, side)
            if !isborder(r, change(side, 1))
                move!(r, change(side, 1))
                num += 1
            else
                brea
            end
        end
        if isborder(r, side)
            for i in 1:num
                if !isborder(r, change(side, 3))
                    move!(r, change(side, 3))
                end
            end
            return false
        else
            move!(r, side)
            for i in 1:num
                if !isborder(r, change(side, 3))
                    move!(r, change(side, 3))
                end
            end
            return true
        end
    end

    function move_and_num(side, num)
        flag = false
        ok = true
        count = 0
        while true
            if (isborder(r, change(side, num)) && !flag)
                flag = true
                count += 1
            elseif !isborder(r, change(side, num))
                flag = false
            end
            if (isborder(r, side))
                ok = go_around(side)
            end
            if ok
                if !isborder(r, side)
                    move!(r, side)
                end
            else
                break
            end
        end
        return count
    end

    function num(side, move_side, ok)
        count = 0
        n = 1
        if ok
            n = 1
        else
            n = 3
        end
        while !isborder(r, side)
            count += move_and_num(move_side, n)
            move_side = change(move_side, 2)
            n = (n + 2) % 4
            move!(r, side)
        end
        return count
    end

    function solve()
        count = 0
        count += num(Nord, Ost, true)
        while !isborder(r, Sud)
            move!(r, Sud)
        end
        count += num(Ost, Nord, false)
        println(count)
    end

    solve()
end


# # Model construction and modification

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/@__NAME__.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/@__NAME__.ipynb)

# `COBREXA` can load models stored in `.mat`, `.json`, and `.xml` formats; and convert
# these into `StandardModel`s. However, it is also possible to construct models
# from scratch, and modify existing models. This will be demonstrated
# here.

# ## Model construction

# In `COBREXA`, model construction is primarily supported through `StandardModel`s.
# To begin, create an empty `StandardModel`.

using COBREXA

model = StandardModel("FirstModel") # assign model id = "FirstModel"

# Next, genes, metabolites and reactions need to be added to the model.

# ### Add genes to the model
gene_list = [Gene(string("g", num)) for num = 1:8]

#md # !!! warning "Warning: Don't accidentally overwrite the generic accessors"
#md #       It may be tempting to call a variable `genes`, `metabolites`, or
#md #       `reactions`. However, these names conflict with generic accessors
#md #       functions and will create problems downstream.

add_genes!(model, gene_list)

# ### Add metabolites to the model
metabolite_list = [Metabolite(string("m", num)) for num = 1:4]

metabolite_list[1].formula = "C6H12O6" # can edit metabolites, etc. directly

add_metabolites!(model, metabolite_list)

# ### Add reactions to the model

# There are two ways to create and add reactions to a model.
# These are using functions, or macros.

r_m1 = Reaction("EX_m1", Dict("m1" => -1.0), :bidirectional) # exchange reaction: m1 <-> (is the same as m1 ⟷ nothing)
r1 = Reaction("r1", Dict("m1" => -1.0, "m2" => 1.0), :forward)
r1.grr = [["g1", "g2"], ["g3"]] # add some gene reaction rules
r2 = Reaction("r2", Dict("m2" => -1.0, "m1" => 1.0), :backward)
r3 = Reaction("r3", Dict("m2" => -1.0, "m3" => 1.0), :bidirectional)

add_reactions!(model, [r1, r2, r3, r_m1]) # function approach

m1 = metabolite_list[1]
m2 = metabolite_list[2]
m3 = metabolite_list[3]
m4 = metabolite_list[4]

@add_reactions! model begin # macro approach
    "r4", m2 ⟶ m4, 0, 1000
    "r_m3", m3 ⟷ nothing, -1000, 1000
    "r_m4", m4 ⟶ nothing
    "r5", m4 ⟶ m2
end

model.reactions["r4"].grr = [["g5"], ["g6", "g7"], ["g8"]]

#md # !!! note "Note: Using reaction arrows"
#md #       `COBREXA` exports arrows that can be used to construct reactions.
#md #       Both the long and short arrows (`⟶ == →`) mean the same thing if they
#md #       point in the same direction.
#md #
#md #       These arrows are accessible by using the `LaTeX` completions built into
#md #       Julia. For example:
#md #       1. → is \rightarrow<tab>
#md #       2. ⟶ is \longrightarrow<tab>
#md #       3. ← is \leftarrow<tab>
#md #       4. ⟵ is \longleftarrow<tab>
#md #       5. ↔ is \leftrightarrow<tab>
#md #       6. ⟷ is \longleftrightarrow<tab>

# The constructed model can now be inspected.
model

# ## Modifying existing models

# It is also possible to modify a model by deleting certain genes.
# This is simply achieved by calling `remove_genes!`.

remove_genes!(model, ["g1", "g2"]; knockout_reactions = false)
model


using Abaco
using Test

interval=10
ts = nowts()

val1 = 100
val2 = 101
val3 = 102

expected_triggers = 3
actual_triggers = 0

#function onresult(ts, ne, name, value, inputs)
#    global actual_triggers
#    @debug "timestamp [$ts]: ne: [$ne] $name = $value"
#    if name === "offset_footprint"
#        @test value == x + val1 + val2 + val3
#    end
#    actual_triggers += 1
#end

function onresult(timestamp, ne, name, value::Abaco.PValue, inputs)
    global actual_triggers
    @debug "progressive timestamp [$ts]: ne: [$ne] $name = $value"
    if value.contribs == 1
        @test ne == "italy"
        @test timestamp == span(ts, interval)
        @test value.value == val1
    elseif value.contribs == 2
        @test value.value == val1 + val2
    elseif value.contribs == 3
        @test value.value == val1 + val2 + val3
    else
        # never here
        @test 1 == 0
    end
    actual_triggers += 1
end

abaco = abaco_init(onresult, interval=interval)

@debug "starting demo with time span: $interval, values timestamp: $ts"

setup_settings(abaco, "state", oncomplete=onresult)

formula(abaco, "total", "sum(region.city.footprint)", "state")

italy = node(abaco, "italy", "state")
veneto = node(abaco, italy, "veneto", "region")
trentino = node(abaco, italy, "trentino", "region")
belluno = node(abaco, veneto, "belluno", "city")
padova = node(abaco, veneto, "padova", "city")
rovereto = node(abaco, trentino, "rovereto", "city")
pergine = node(abaco, trentino, "pergine", "city")

#@info "TEST: veneto origins: $(abaco.origins[veneto.ne])"

#current_footprint = get_collected(abaco, italy.ne, "region.city.footprint")
#@debug "initial carbon footprints: $current_footprint"

ingest(abaco, ts, belluno.ne, Dict("footprint" => val1))

ingest(abaco, ts, padova.ne, Dict("footprint" => val2))
ingest(abaco, ts, rovereto.ne, Dict("footprint" => val3))

#current_footprint = get_collected(abaco, italy.ne, "region.city.footprint", ts)
#@debug "current carbon footprint: $current_footprint"

@test actual_triggers === expected_triggers


export PARAMETER_DEFAULT, PARAMETER_UNSTABLE_BUT_FAST, PARAMETER_STABLE_BUT_SLOW

using MathOptInterface
MOI = MathOptInterface

mutable struct SDOptimizer <: SDOI.AbstractSDOptimizer
    problem::SDPAProblem
    options::Dict{Symbol,Any}
    function SDOptimizer(; kwargs...)
        new(SDPAProblem(), Dict{Symbol, Any}(kwargs))
    end
end
Optimizer(; kws...) = SDOI.SDOIOptimizer(SDOptimizer(; kws...))

MOI.get(::SDOptimizer, ::MOI.SolverName) = "SDPA"

function MOI.empty!(optimizer::SDOptimizer)
    optimizer.problem = SDPAProblem()
end

function SDOI.init!(m::SDOptimizer, blkdims::Vector{Int}, nconstrs::Int)
    @assert nconstrs >= 0
    dummy = nconstrs == 0
    if dummy
        nconstrs = 1
        blkdims = [blkdims; -1]
    end
    m.problem = SDPAProblem()
    setParameterType(m.problem, PARAMETER_DEFAULT)
    setparameters!(m.problem, m.options)
    inputConstraintNumber(m.problem, nconstrs)
    inputBlockNumber(m.problem, length(blkdims))
    for (i, blkdim) in enumerate(blkdims)
        inputBlockSize(m.problem, i, blkdim)
        inputBlockType(m.problem, i, blkdim < 0 ? LP : SDP)
    end
    initializeUpperTriangleSpace(m.problem)
    if dummy
        SDOI.setconstraintconstant!(m, 1., 1)
        SDOI.setconstraintcoefficient!(m, 1., 1, length(blkdims), 1, 1)
    end
end

function SDOI.setconstraintconstant!(m::SDOptimizer, val, constr::Integer)
    @assert constr > 0
    #println("b[$constr] = $val")
    inputCVec(m.problem, constr, val)
end
function SDOI.setconstraintcoefficient!(m::SDOptimizer, coef, constr::Integer, blk::Integer, i::Integer, j::Integer)
    @assert constr > 0
    #println("A[$constr][$blk][$i, $j] = $coef")
    inputElement(m.problem, constr, blk, i, j, float(coef), false)
end
function SDOI.setobjectivecoefficient!(m::SDOptimizer, coef, blk::Integer, i::Integer, j::Integer)
    #println("C[$blk][$i, $j] = $coef")
    inputElement(m.problem, 0, blk, i, j, float(coef), false)
end

function MOI.optimize!(m::SDOptimizer)
    SDPA.initializeUpperTriangle(m.problem, false)
    SDPA.initializeSolve(m.problem)
    SDPA.solve(m.problem)
end

function MOI.get(m::SDOptimizer, ::MOI.TerminationStatus)
    status = getPhaseValue(m.problem)
    if status == noINFO
        return MOI.OPTIMIZE_NOT_CALLED
    elseif status == pFEAS
        return MOI.SLOW_PROGRESS
    elseif status == dFEAS
        return MOI.SLOW_PROGRESS
    elseif status == pdFEAS
        return MOI.OPTIMAL
    elseif status == pdINF
        return MOI.INFEASIBLE_OR_UNBOUNDED
    elseif status == pFEAS_dINF
        return MOI.DUAL_INFEASIBLE
    elseif status == pINF_dFEAS
        return MOI.INFEASIBLE
    elseif status == pdOPT
        return MOI.OPTIMAL
    elseif status == pUNBD
        return MOI.DUAL_INFEASIBLE
    elseif status == dUNBD
        return MOI.INFEASIBLE
    end
end

function MOI.get(m::SDOptimizer, ::MOI.PrimalStatus)
    status = getPhaseValue(m.problem)
    if status == noINFO
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS
        return MOI.FEASIBLE_POINT
    elseif status == dFEAS
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pdFEAS
        return MOI.FEASIBLE_POINT
    elseif status == pdINF
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS_dINF
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == pINF_dFEAS
        return MOI.INFEASIBLE_POINT
    elseif status == pdOPT
        return MOI.FEASIBLE_POINT
    elseif status == pUNBD
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == dUNBD
        return MOI.INFEASIBLE_POINT
    end
end

function MOI.get(m::SDOptimizer, ::MOI.DualStatus)
    status = getPhaseValue(m.problem)
    if status == noINFO
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == dFEAS
        return MOI.FEASIBLE_POINT
    elseif status == pdFEAS
        return MOI.FEASIBLE_POINT
    elseif status == pdINF
        return MOI.UNKNOWN_RESULT_STATUS
    elseif status == pFEAS_dINF
        return MOI.INFEASIBLE_POINT
    elseif status == pINF_dFEAS
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == pdOPT
        return MOI.FEASIBLE_POINT
    elseif status == pUNBD
        return MOI.INFEASIBLE_POINT
    elseif status == dUNBD
        return MOI.INFEASIBILITY_CERTIFICATE
    end
end

SDOI.getprimalobjectivevalue(m::SDOptimizer) = getPrimalObj(m.problem)
SDOI.getdualobjectivevalue(m::SDOptimizer) = getDualObj(m.problem)
SDOI.getX(m::SDOptimizer) = PrimalSolution(m.problem)
function SDOI.gety(m::SDOptimizer)
    unsafe_wrap(Array, getResultXVec(m.problem), getConstraintNumber(m.problem))
end
SDOI.getZ(m::SDOptimizer) = VarDualSolution(m.problem)


## Function so the user does not need to type it every time
# --------------------------------------> For now, resolution of simulation is every 1. Need to add modification in case user wants to go lower
function defModStruct()
    model_def = Dict();

    model_def["NameF"] = []; # Name for the model (scripts will be stored with this name) --> String
    model_def["nStat"] = []; # Nuber of states in the model --> Integer
    model_def["nPar"] = []; # Number of parameters of the mode --> Integer
    model_def["nInp"] = []; # Nuber of stimuly (inducer) in the model --> Integer
    model_def["stName"] = []; # Name for all the states --> Vector of strings
    model_def["parName"] = []; # Name for all the parameters --> Vector of strings
    model_def["inpName"] = []; # Name for all the inputs --> Vector of strings
    model_def["eqns"] = []; # Equations defining the model (left and right hand side) --> Vector of strings
    model_def["Y0eqs"] = []; # Steady state equations --> Vector of strings or empty vector if not used. If some element of the equation requires an experimental value for the calculation, please add exp at the beginning of the state(exemple: Cmrna -> expCmrna). Please do not name anything alp.
    model_def["Y0Sim"] = []; # If the ON simulation is required for the steady state --> true, false or a string saying Yes/yes, No/no. If empty the default would be No. Time-scale for the simulation is assumed in minutes. If this wants to be changed, what should be introduced here is a number for the time conversion (e.g. 1/60 if to convert to days or 60 if to convert to seconds)
    model_def["tols"] = []; # Relative and absolute tolerances (in that order) --> vector of 2 floats (can be left empty, where 1e-6 will be taken as default for both)
    model_def["solver"] = []; # IVP solver to solve the ODEs. If nothing specified, the default will be Tsit5(). For more info check https://diffeq.sciml.ai/v2.0/tutorials/ode_example.html

    return(model_def)
end

## Function with all the checks
# NOTE: some checks have been commented since they considered the option that the user would only introduce ODE and steady state equations and nothing else. Code has been changed to allow definition of variables, prints, ifs, fors, etc.

function checkStruct(model_def)

    # Check taht all the dictionary entries are correct
    entries = ["stName","inpName","eqns","nPar","NameF","nInp","parName","nStat","Y0eqs","Y0Sim","tols","solver"]
    if symdiff(entries,keys(model_def))!=[] && symdiff(entries,keys(model_def))!=["modelpath"]
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the entries of the dictionary, there is something wrong...")
        println(symdiff(entries,keys(model_def)))
        return
    end

    # Check one to be sure there is no empty entry
    if model_def["NameF"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce a name for the model!")
        return
    elseif model_def["nStat"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the number of states!")
        return
    elseif model_def["nPar"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the number of parameters!")
        return
    elseif model_def["nInp"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the number of stimuli!")
        return
    elseif model_def["stName"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the state names!")
        return
    elseif model_def["parName"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the parameter names!")
        return
    elseif model_def["inpName"] == [] && (model_def["nInp"] != 0 && model_def["nInp"] != [0])
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the stimuli names!")
        return
    elseif model_def["eqns"] == []
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, introduce the equations!")
        return
    end

    # Checks to see if the content of the fields is the expected (also consier if single value entries have been given as a number or a vector)
    if ((typeof(model_def["NameF"])!=Array{String,1}) && (typeof(model_def["NameF"])!=String))
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field NameF!")
        return
    elseif (typeof(model_def["nStat"][1])!=Int) && (typeof(model_def["nStat"])!=Int)
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field nStat!")
        return
    elseif (typeof(model_def["nPar"][1])!=Int) && (typeof(model_def["nPar"])!=Int)
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field nPar!")
        return
    elseif (typeof(model_def["nInp"][1])!=Int) && (typeof(model_def["nInp"])!=Int)
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field nInp!")
        return
    elseif (typeof(model_def["stName"])!=Array{String,1})
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field stName!")
        return
    elseif (typeof(model_def["parName"])!=Array{String,1})
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field parName!")
        return
    elseif (typeof(model_def["inpName"])!=Array{String,1}) && (typeof(model_def["inpName"])!=Array{Any,1})
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field inpName!")
        return
    elseif (typeof(model_def["eqns"])!=Array{String,1})
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field eqns!")
        return
    elseif (typeof(model_def["Y0eqs"])!=Array{Any,1}) && (typeof(model_def["Y0eqs"])!=Array{String,1})
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field Y0eqs!")
        return
    elseif ((typeof(model_def["solver"])!=Array{String,1}) && (typeof(model_def["solver"])!=String)) && ((model_def["solver"])!=[])
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field solver!")
        return
    elseif model_def["Y0Sim"]!=true && model_def["Y0Sim"]!="Yes" && model_def["Y0Sim"]!="yes" &&
    model_def["Y0Sim"]!= false && model_def["Y0Sim"] != "No" && model_def["Y0Sim"] != "no"&&
    model_def["Y0Sim"]!=[true] && model_def["Y0Sim"]!=["Yes"] && model_def["Y0Sim"]!=["yes"] &&
    model_def["Y0Sim"]!= [false] && model_def["Y0Sim"] != ["No"] && model_def["Y0Sim"] != ["no"]&&
    model_def["Y0Sim"]!= [] && model_def["Y0Sim"] != "" && (typeof(model_def["Y0Sim"])!=Array{Int,1}) &&
    (typeof(model_def["Y0Sim"])!=Array{Float64,1}) && (typeof(model_def["Y0Sim"])!=Array{Float32,1}) &&
    (typeof(model_def["Y0Sim"])!=Int) && (typeof(model_def["Y0Sim"])!=Float64) && (typeof(model_def["Y0Sim"])!=Float32)

        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field Y0Sim!")
        return
    elseif (typeof(model_def["tols"])!=Array{Any,1}) && (typeof(model_def["tols"])!=Array{Float64,1}) && (typeof(model_def["tols"])!=Array{Float32,1})
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field tols!")
        return
    elseif ((typeof(model_def["tols"])==Array{Float64,1}) || (typeof(model_def["tols"])==Array{Float32,1})) && (length(model_def["tols"]) != 2)
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the field tols!")
        return
    end

    if "alp" in model_def["parName"] || "alp" in model_def["inpName"] || "alp" in model_def["stName"]
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Sorry, but the name alp is reserved for something else in the model...")
        return
    end

    # Take single elements outside vector (if so) to generalise things after
    if typeof(model_def["NameF"]) == Array{String,1}
        model_def["NameF"] = model_def["NameF"][1];
    end

    if typeof(model_def["nStat"]) == Array{Int,1}
        model_def["nStat"] = model_def["nStat"][1];
    end

    if typeof(model_def["nPar"]) == Array{Int,1}
        model_def["nPar"] = model_def["nPar"][1];
    end

    if typeof(model_def["nInp"]) == Array{Int,1}
        model_def["nInp"] = model_def["nInp"][1];
    end

    if model_def["Y0Sim"]==true || model_def["Y0Sim"]=="Yes" || model_def["Y0Sim"]=="yes" ||
    model_def["Y0Sim"]==[true] || model_def["Y0Sim"]==["Yes"] || model_def["Y0Sim"]==["yes"]
        model_def["Y0Sim"]=true;
    elseif model_def["Y0Sim"]== false || model_def["Y0Sim"] == "No" || model_def["Y0Sim"] == "no"||
    model_def["Y0Sim"]== [false] || model_def["Y0Sim"] == ["No"] || model_def["Y0Sim"] == ["no"]||
    model_def["Y0Sim"]== [] || model_def["Y0Sim"] == ""
        model_def["Y0Sim"]=false;
    elseif typeof(model_def["Y0Sim"]) <: Array
        model_def["Y0Sim"]=model_def["Y0Sim"][1];
    end

    if model_def["tols"] == []
        model_def["tols"] = [1e-6, 1e-6];
    end

    if model_def["solver"] == []
        model_def["solver"] = "Tsit5";
    elseif typeof(model_def["solver"]) == Array{String,1}
        model_def["solver"] = model_def["solver"][1];
    end

    # Check that the user has introduced a correct solver
    if (model_def["solver"] != "AutoTsit5(Rosenbrock23())") && (model_def["solver"] != "AutoVern7(Rodas5())") &&
        (model_def["solver"] != "Tsit5") && (model_def["solver"] != "BS3") && (model_def["solver"] != "Vern7") &&
        (model_def["solver"] != "Rodas4") && (model_def["solver"] != "Rodas5") && (model_def["solver"] != "KenCarp4") &&
        (model_def["solver"] != "TRBDF2") && (model_def["solver"] != "RadauIIA") && (model_def["solver"] != "CVODE_BDF")
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Sorry, but you have introduced a wrong solver!")
        println("The options are AutoTsit5(Rosenbrock23()), AutoVern7(Rodas5()), Tsit5, BS3, Vern7, Rodas4, Rodas5, ")
        println("          KenCarp4, TRBDF2, RadauIIA, CVODE_BDF.")
        println("For more info, please check https://diffeq.sciml.ai/stable/tutorials/ode_example/#ode_example-1")
        return
    end

    # Check tha there are no spaces in determined fields
    if findfirst(" ", join(model_def["stName"])) != nothing
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, do NOT include spaces in the entries of the field stName!")
        return
    end

    if findfirst(" ", join(model_def["inpName"])) != nothing
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, do NOT include spaces in the entries of the field inpName!")
        return
    end

    if findfirst(" ", join(model_def["NameF"])) != nothing
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, do NOT include spaces in the entries of the field NameF!")
        return
    end

    if findfirst(" ", join(model_def["parName"])) != nothing
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, do NOT include spaces in the entries of the field parName!")
        return
    end

    # Check that contents make sense
    if length(model_def["stName"]) != model_def["nStat"]
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check stName and nStat, they do not make sense")
        return
    end
    if length(model_def["inpName"]) != model_def["nInp"]
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check inpName and nInp, they do not make sense")
        return
    end

    if length(model_def["parName"]) != model_def["nPar"]
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check parName and nPar, they do not make sense")
        return
    end

    #     if length(model_def["eqns"]) != model_def["nStat"]
    #         println("-------------------------- Process STOPPED!!! --------------------------")
    #         println("Please, chec eqns and nStat, the number of equations is not the same as the number of states")
    #         return
    #     end

    # ----------------------------------------> Might need to double check this part
    # Check if in the ODEs all the states are in the equations and with a d in front
    # First check that the equations contain an = sign and only 1

    # Check that in the ODEs there is only 1 = sign
    if (findfirst.("=", model_def["eqns"])!=findlast.("=", model_def["eqns"]))
        println("-------------------------- WARNING!!! --------------------------")
        println("Please, check the number of = signs in your OED equations!")
        println("Make sure it is something desired (ifs, fors, prints, etc.) and not that you put 2 equations in one same string or forgot a right-hand side. ")
    end
    if model_def["Y0eqs"] != [] # Avoind the check if no Steady state equations are given
        # Check that in the Steady State Equations there is only 1 equal sign
        if sum((findfirst.("=", model_def["Y0eqs"]).==findlast.("=", model_def["Y0eqs"]))) < model_def["nStat"]
            println("-------------------------- Process STOPPED!!! --------------------------")
            println("Please, check the number of = signs in your Steady State equations! You might have put 2 equations in one same string!")
            return
        end
        # Check that in the Steady state equations all entries have left and right hand side
        # if sum(occursin.("=", model_def["Y0eqs"])) != length(model_def["Y0eqs"])
        #     println("-------------------------- Process STOPPED!!! --------------------------")
        #     println("Please, check the equations! You might have forgoten a left hand side or add 2 equations in one same string in the Steady State Equation.")
        #     return
        # end
        # Check that in the Steady State equations all the states are given as equations
        ds2 = [model_def["Y0eqs"][i][1:findfirst.("=", model_def["Y0eqs"][i])[1]] for i in 1:length(model_def["Y0eqs"]) if !isnothing(findfirst.("=", model_def["Y0eqs"][i]))]
        if sum(occursin.(model_def["stName"], join(ds2))) !=  model_def["nStat"]
            println("-------------------------- Process STOPPED!!! --------------------------")
            println("Please, check the equations! You might have forgotten a Steady State Equation!.")
            return
        end
    end

    # Check that in the ODEs all entries have left and right hand side
    # if sum(occursin.("=", model_def["eqns"])) != length(model_def["eqns"])
    #     println("-------------------------- Process STOPPED!!! --------------------------")
    #     println("Please, check the equations! You might have forgoten a left hand side of the ODEs or add 2 equations in one same string.")
    #     return
    # end

    # Check that in the ODEs all the states are given as equations
    ds = [model_def["eqns"][i][1:findfirst.("=", model_def["eqns"][i])[1]] for i in 1:length(model_def["eqns"]) if !isnothing(findfirst.("=", model_def["eqns"][i]))];
    if sum(occursin.(string.("d", model_def["stName"]), join(ds))) != model_def["nStat"]
        println("-------------------------- Process STOPPED!!! --------------------------")
        println("Please, check the equations! You might have forgotten an ODE!.")
        println("Remember, ODEs should start with a d and no other equation can.")
        return
    end

    return(model_def)
end

## Function to generate all the necesary Julia scripts for the model simulations

function GenerateModel(model_def)

    model_def = checkStruct(model_def);

    # Packages needed for the simulations
    Head = string("
    using DifferentialEquations
    using OrdinaryDiffEq
    using DiffEqBase
    using Sundials
    using ODEInterfaceDiffEq
        ");

    # Generate ODEs function
    # String containing the equtions, ODEs and others
    tet = Array{Any,1}(undef, length(model_def["eqns"])); # Vector of strings that will contain the eauqtions for the function
    # odec = 1; # Counter for the ODEs to take into account equations that are not ODEs
    for i in 1:length(model_def["eqns"]) # Loop over number of equations
        if replace(model_def["eqns"][i], " "=>"")[1]=='d' # If the equation starts with a d (differential)
            for j in 1:model_def["nStat"]
                eqsg = findfirst("=", replace(model_def["eqns"][i], " "=>""))[1];
                if replace(model_def["eqns"][i], " "=>"")[1:eqsg-1] == string("d", model_def["stName"][j])
                    tet[i] = string("        du[", j, "] = ", model_def["eqns"][i],";\n");
                    break;
                else
                    tet[i] = string("        ", model_def["eqns"][i],";\n");
                end
            end
        else # Other equations
            tet[i] = string("            ", model_def["eqns"][i],";\n");
        end
    end;

    fun1 = string("
        # du -> Derivative
        # u -> State at time t
        # p-> paramter vector
        # t-> time as tuple (init, end)

        function ",join(model_def["NameF"]),"ODE!(du,u,p,t)
            ",join(model_def["stName"], ", ")," = u;

            ",join(model_def["parName"], ", "),", ",join(model_def["inpName"], ", ")," = p;

",join(tet),"

        end
        ");

    # Generate Steady State function
    y0eq = ""; # Vector that will contain the steady state equations (SSE)
    if model_def["Y0eqs"]==[] # If there is no equations, make the function return the same Y0 that the user gives
        y0eq = (string(join(model_def["stName"], ", ")," = I; \n        ",

                "       return(I) \n "));
    else
        tet = Array{Any,1}(undef, length(model_def["Y0eqs"])); # Vector of strings that will contain the eauqtions for the function
        odec = 1 # Counter for the SSEs to take into account equations that are not SSEs
        df = Array{Any,1}(undef, length(model_def["Y0eqs"]));
        for i in 1:length(model_def["Y0eqs"]) # Check for the equations that start with the name of the states (to tell apart from other equations)
            if !isnothing(findfirst.("=", model_def["Y0eqs"][i]))
                df[i] = model_def["Y0eqs"][i][1:findfirst.("=", model_def["Y0eqs"][i])[1]]
            else
                df[i] = "";
            end
        end;
        for i in 1:length(model_def["Y0eqs"]) # Loop over equations
            if sum(occursin.(model_def["stName"], df[i]))!=0 # If the equation looked begins with the name of a state
                tet[i] = string("        alp[", odec, "] = ", model_def["Y0eqs"][i],";\n");
                odec +=1;
            else
                tet[i] = string("        ", model_def["Y0eqs"][i],";\n");
            end
        end

        y0eq = (string(join([string("exp",model_def["stName"][k]) for k in 1:model_def["nStat"]], ", ")," = I; \n        ",

        join(model_def["parName"], ", "),", ",join(model_def["inpName"], ", ")," = p;\n \n \n ",

                "        alp = zeros(",length(model_def["stName"]),"); \n \n",

                join(tet),
                "        return(alp) \n "));

    end

    fun2 = string("
        # p -> parameter vector (plus inducer vector at the end)
        # I -> initial Y0 vector

        function ",join(model_def["NameF"]),"SteadyState(p,I)

            ",y0eq,"

        end
        ");

    # Generate Step Wise simulation function
    ONSim = "";
    if model_def["Y0Sim"]==false
        ONSim = "y0 = y_al;";
    elseif model_def["Y0Sim"]==true
        if model_def["solver"] == "CVODE_BDF"
            ONSim = string("
                prob = ODEProblem(",join(model_def["NameF"]),"ODE!,y_al,(0.0,Float64(24*60-1)),pSS);
                ssv = Sundials.solve(prob, CVODE_BDF(),reltol=",model_def["tols"][1],",abstol=",model_def["tols"][2],");
                y0 = ssv[:,end]; ");
        else
            ONSim = string("
                prob = ODEProblem(",join(model_def["NameF"]),"ODE!,y_al,(0.0,Float64(24*60-1)),pSS);
                ssv = DifferentialEquations.solve(prob, ",join(model_def["solver"]),"(),reltol=",model_def["tols"][1],",abstol=",model_def["tols"][2],");
                y0 = ssv[:,end]; ");
        end
    elseif typeof(model_def["Y0Sim"]) == Float64 || typeof(model_def["Y0Sim"]) == Float32 || typeof(model_def["Y0Sim"]) == Int
        if model_def["solver"] == "CVODE_BDF"
            ONSim = string("
                prob = ODEProblem(",join(model_def["NameF"]),"ODE!,y_al,(0.0,Float64(",model_def["Y0Sim"],"*24*60-1)),pSS);
                ssv = Sundials.solve(prob, CVODE_BDF(),reltol=",model_def["tols"][1],",abstol=",model_def["tols"][2],");
                y0 = ssv[:,end]; ");
        else
            ONSim = string("
                prob = ODEProblem(",join(model_def["NameF"]),"ODE!,y_al,(0.0,Float64(",model_def["Y0Sim"],"*24*60-1)),pSS);
                ssv = DifferentialEquations.solve(prob, ",join(model_def["solver"]),"(),reltol=",model_def["tols"][1],",abstol=",model_def["tols"][2],");
                y0 = ssv[:,end]; ");
        end
    end

    if model_def["solver"] == "CVOED_BDF"
        soso = "Sundials.";
    else
        soso = "DifferentialEquations.";
    end

    fun3 = string("
        # ts -> time vector going from t=0 to t=end every 1
        # p -> parameter vector
        # sp -> vector with switching times (t=0 and t=end needs to be included). Length of vector will be length(inducer steps)+1
        # inputs -> Vecor with the inducers (if more than 1 inducer in the system, the shape will be a1,b1,c1,..., a2,b2,c2,... aend,bend,cend,... where each leter represents a different inducer)
        # ivss -> Initial Y0 vector
        # pre -> Vector with the inducers in the ON. It can be empty

        function ",join(model_def["NameF"]),"_solvecoupledODE(ts, p, sp, inputs, ivss, pre=[])

            maxtime = length(ts);
            Nsp = length(sp);
            Nevents = length(sp)-1;
            Neq = ",join(model_def["nStat"]),";
            if pre != []
                pSS = vcat(p, pre);
            else
                pSS = p;
            end

            final = zeros(maxtime, Neq);

            y_al = ",join(model_def["NameF"]),"SteadyState(pSS,ivss) # Calculation of initial guesses for steady state

            ",join(ONSim),"
            initialV = y0;
            i = 1;

            for q in collect(1:Nevents)
                lts = length(ts[(sp[q]+1):sp[q+1]+1]);  # General way to define the number of elements in each event series
                Tevent = ts[(sp[q]+1):sp[q+1]+1];  # General way to extract the times of each event
                I = inputs[i:(i+(",model_def["nInp"]-1,"))];
                pSte = vcat(p, I);

                if q == 1
                    prob = ODEProblem(",join(model_def["NameF"]),"ODE!,initialV,(ts[(sp[q]+1)],ts[sp[q+1]+1]),pSte);
                    part1 = ",join(soso),"solve(prob, ",join(model_def["solver"]),"(),reltol=",model_def["tols"][1],",abstol=",model_def["tols"][2],",saveat=1);
                else
                    prob = ODEProblem(",join(model_def["NameF"]),"ODE!,initialV,(ts[(sp[q]+1)],ts[sp[q+1]+1]),pSte);
                    part1 = ",join(soso),"solve(prob, ",join(model_def["solver"]),"(),reltol=",model_def["tols"][1],",abstol=",model_def["tols"][2],",saveat=1);
                end

                initialV = part1[lts];
                i+=",model_def["nInp"],";

                for d in collect((sp[q]+1):(sp[q]+lts))
                    final[d,:] = part1[d-sp[q]];
                end
            end
            return(final)

        end
        ");

    # Function to simulate the system as many times as parameter draws are given
    fun4 = string("
        # ts -> time vector going from t=0 to t=end every 1
        # pD -> Matrix containing all the draws for the parameters (can be a single vector if nly 1 draw of parameters is considered)
        # sp -> vector with switching times (t=0 and t=end needs to be included). Length of vector will be length(inducer steps)+1
        # inputs -> Vecor with the inducers (if more than 1 inducer in the system, the shape will be a1,b1,c1,..., a2,b2,c2,... aend,bend,cend,... where each leter represents a different inducer)
        # ivss -> Initial Y0 vector
        # samps -> Vector of Time points that want to be extracted from the simulations (it considers the initial time point as 0)
        # pre -> Vector with the inducers in the ON. It can be empty

        function ",join(model_def["NameF"]),"_SolveAll(ts, pD, sp, inputs, ivss, samps, pre=[])

            if length(size(pD)) == 1
                pD = reshape(pD,size(pD)[1],1);
            end

            if size(pD)[2] != ",model_def["nPar"],"
                pD = pD';
            end

            if length(ivss)/",model_def["nStat"]," > 1
                if size(ivss)[2] != ",model_def["nStat"],"
                    ivss = ivss';
                end
            end

            AllSolTest = zeros(length(samps), ",model_def["nStat"],", length(pD[:,1]));

            for drawInd in collect(1:length(pD[:,1]))
                p = pD[drawInd,:];
                if length(ivss)/",model_def["nStat"]," > 1
                    ivss2 = ivss[drawInd,:];
                else
                    ivss2 = ivss;
                end
                temp = ",join(model_def["NameF"]),"_solvecoupledODE(ts, p, sp, inputs, ivss2, pre);
                AllSolTest[:,:,drawInd] = temp[convert.(Int,samps.+1),:];
            end

            return(AllSolTest);
        end
        ");

    finscr = string(Head, fun1, fun2, fun3, fun4);

    fidi = pwd(); #@__DIR__;

    if !isdir(string(fidi,"\\ModelsFunctions"))
        mkdir(string(fidi,"\\ModelsFunctions"));
    end

    open(string(fidi,"\\ModelsFunctions\\", model_def["NameF"], "_Model.jl"), "w") do io
       write(io, finscr);
    end;

    include(string(fidi,"\\ModelsFunctions\\", model_def["NameF"], "_Model.jl"))

    model_def["modelpath"] = string(fidi,"\\ModelsFunctions\\", model_def["NameF"], "_Model.jl");

    println("")
    println("----------------------------------------- MODEL GENERATION -----------------------------------------")
    println("The model has been generated in the directory: ")
    println(string("                 ", model_def["modelpath"]))
    println("--------------------------------------------------------------------------------------")
    println("")

    return(model_def)

end


# test constraints
# MPI Declarations
comm = MPI.COMM_WORLD
comm_world = MPI.MPI_COMM_WORLD
comm_self = MPI.COMM_SELF
my_rank = MPI.Comm_rank(comm)
comm_size = MPI.Comm_size(comm)

# opts = PdePumiInterface.get_defaults()
# # 2D mesh
# opts["order"] = 1
# opts["dimensions"] = 2
# opts["use_DG"] = true
# opts["operator_type"] = "SBPOmega"
# opts["dmg_name"] = "../src/mesh_files/2D_Airfoil.dmg"
# opts["smb_name"] = "../src/mesh_files/2D_Airfoil.smb"
# opts["numBC"] = 2
#
# # For 2DAirfoil
# opts["BC1"] = [8,11,14,17]
# opts["BC1_name"] = "FarField"
# opts["BC2"] = [5]
# opts["BC2_name"] = "Airfoil"
#
# opts["coloring_distance"] = 2 # 0 For CG Mesh 2 for DG Mesh
# opts["jac_type"] = 2
# opts["jac_method"] = 2
# opts["run_type"] = 5
#
# # Create PumiMesh and SBP objects
# sbp, mesh, pmesh, Tsol, Tres, Tmsh, mesh_time = PDESolver.createMeshAndOperator(opts, 1)

opts = Dict{ASCIIString, Any}()
opts["order"] = 1
opts["dimensions"] = 2
opts["use_DG"] = true
opts["Tsbp"] = Float64
opts["Tmsh"] = Complex128
opts["operator_type"] = "SBPOmega"
opts["dmg_name"] = "../src/mesh_files/2D_Airfoil.dmg"
opts["smb_name"] = "../src/mesh_files/2D_Airfoil.smb"
opts["numBC"] = 2

# For 2DAirfoil
opts["BC1"] = [8,11,14,17]
opts["BC1_name"] = "FarField"
opts["BC2"] = [5]
opts["BC2_name"] = "Airfoil"

opts["coloring_distance"] = 2 # 0 For CG Mesh 2 for DG Mesh
opts["jac_type"] = 2
opts["jac_method"] = 2
opts["run_type"] = 5

# Create PumiMesh and SBP objects
dofpernode = 1
ref_verts = [-1. 1 -1; -1 -1 1]
Tsbp = opts["Tsbp"]
Tmsh = opts["Tmsh"]
sbp = getTriSBPOmega(degree=opts["order"], Tsbp=opts["Tsbp"])
sbpface = TriFace{Tsbp}(opts["order"], sbp.cub, ref_verts.')
topo = 0
shape_type = 2
mesh = PumiMeshDG2(Tmsh, sbp, opts, sbpface, dofpernode=dofpernode,
                     shape_type=shape_type)

#mesh = PumiMeshDG2{Tmsh}(opts["dmg_name"], opts["smb_name"], opts["order"],
#                         sbp, opts, sbpface; dofpernode=dofpernode,
#                         coloring_distance=opts["coloring_distance"],
#                         shape_type=shape_type)

#geom_faces = opts["BC2"]
bc_nums = [2]
# Free Form deformation parameters
ndim = 3
order = [2,4,2]  # Order of B-splines in the 3 directions
nControlPts = [2,4,2]

# code here for creating rectilinear map
# Create Mapping object
map = PumiMapping{Tmsh}(ndim, order, nControlPts, mesh,
                            bc_nums=bc_nums)

# Create knot vector
calcKnot(map)

# Create Bounding box
offset = [0., 0., 0.5] # No offset in the X & Y direction
ffd_box = PumiBoundingBox{Tmsh}(map, mesh, sbp, offset)

# Control points
controlPoint(map, ffd_box)

# Populate map.xi
calcParametricMappingNonlinear(map, ffd_box, mesh, bc_nums)

@testset "--- Checking Linear-Plane Constraints ---" begin

  di = 2

  # test function for counting number of constraints
  numCnstr = numLinearPlaneConstraints(map, di)
  @test ( numCnstr )== 1*3*4

  # test function for counting variables in each constraint
  cntEq = zeros(Int, numCnstr)
  fncidx = 1
  fncidx = countVarsLinearPlaneConstraints!(map, di, cntEq, fncidx)
  @test ( fncidx )== numCnstr+1
  for k = 1:numCnstr
    @test ( cntEq[k] )== 4
  end

  # test function for filling sparse constraint Jacobian
  fncidx = 1
  ptr = 1
  iLfun = zeros(Int, sum(cntEq))
  jLvar = zeros(Int, sum(cntEq))
  LinG = zeros(Tmsh, sum(cntEq))
  Flow = rand(numCnstr) + 0im
  Fupp = rand(numCnstr) + 0im
  fncidx, ptr = setLinearPlaneConstraints!(map, di, iLfun, jLvar, LinG, Flow,
                                           Fupp, fncidx, ptr)
  @test ( fncidx )== numCnstr+1
  @test ( ptr )== sum(cntEq)+1
  @test isapprox( Flow, zeros(numCnstr), 1e-15)
  @test isapprox( Fupp, zeros(numCnstr), 1e-15)
  ptr = 1
  for k = 1:numCnstr
    for i = 1:cntEq[k]
      @test ( iLfun[ptr] )== k
      ptr += 1
    end
  end
  # !!!!!!!!!! should also include checks on jLvar and LinG, eventually
end

@testset "--- Checking Linear-Corner Constraints ---" begin

  di = 2

  # test function for counting number of constraints
  numCnstr = numLinearCornerConstraints(map, di)
  @test ( numCnstr )== 8

  # test function for counting variables in each constraint
  cntEq = zeros(Int, numCnstr)
  fncidx = 1
  fncidx = countVarsLinearCornerConstraints!(map, di, cntEq, fncidx)
  @test ( fncidx )== numCnstr+1
  for k = 1:numCnstr
    @test ( cntEq[k] )== 4
  end

  # test function for filling sparse constraint Jacobian
  fncidx = 1
  ptr = 1
  iLfun = zeros(Int, sum(cntEq))
  jLvar = zeros(Int, sum(cntEq))
  LinG = zeros(Tmsh, sum(cntEq))
  Flow = rand(numCnstr) + 0im
  Fupp = rand(numCnstr) + 0im
  fncidx, ptr = setLinearCornerConstraints!(map, di, iLfun, jLvar, LinG, Flow,
                                            Fupp, fncidx, ptr)
  @test ( fncidx )== numCnstr+1
  @test ( ptr )== sum(cntEq)+1
  @test isapprox( Flow, zeros(numCnstr), 1e-15)
  @test isapprox( Fupp, zeros(numCnstr), 1e-15)
  ptr = 1
  for k = 1:numCnstr
    for i = 1:cntEq[k]
      @test ( iLfun[ptr] )== k
      ptr += 1
    end
  end
  # !!!!!!!!!! should also include checks on jLvar and LinG, eventually
end

@testset "--- Checking Linear-Stretch Constraints ---" begin

  di = 2

  # test function for counting number of constraints
  numCnstr = numLinearStretchConstraints(map)
  @test ( numCnstr )== (2*4*2 - 1)

  # test function for counting variables in each constraint
  cntEq = zeros(Int, numCnstr)
  fncidx = 1
  fncidx = countVarsLinearStretchConstraints!(map, cntEq, fncidx)
  @test ( fncidx )== numCnstr+1
  for k = 1:numCnstr
    @test ( cntEq[k] )== 2
  end

  # test function for filling sparse constraint Jacobian
  fncidx = 1
  ptr = 1
  iLfun = zeros(Int, sum(cntEq))
  jLvar = zeros(Int, sum(cntEq))
  LinG = zeros(Tmsh, sum(cntEq))
  Flow = rand(numCnstr) + 0im
  Fupp = rand(numCnstr) + 0im
  fncidx, ptr = setLinearStretchConstraints!(map, di, iLfun, jLvar, LinG, Flow,
                                             Fupp, fncidx, ptr)

  @test ( fncidx )== numCnstr+1
  @test ( ptr )== sum(cntEq)+1
  @test isapprox( Flow, zeros(numCnstr), 1e-15)
  @test isapprox( Fupp, zeros(numCnstr), 1e-15)
  ptr = 1
  for k = 1:numCnstr
    for i = 1:cntEq[k]
      @test ( iLfun[ptr] )== k
      ptr += 1
    end
  end
  # !!!!!!!!!! should also include checks on jLvar and LinG, eventually
end

@testset "--- Checking Linear-Root Constraints ---" begin

  di = 2

  # test function for counting number of constraints
  numCnstr = numLinearRootConstraints(map)
  @test ( numCnstr )== 5

  # test function for counting variables in each constraint
  cntEq = zeros(Int, numCnstr)
  fncidx = 1
  fncidx = countVarsLinearRootConstraints!(map, cntEq, fncidx)
  @test ( fncidx )== numCnstr+1
  for k = 1:3
    @test ( cntEq[k] )== 1
  end
  @test ( cntEq[4] )== 2
  @test ( cntEq[5] )== 2

  # test function for filling sparse constraint Jacobian
  fncidx = 1
  ptr = 1
  iLfun = zeros(Int, sum(cntEq))
  jLvar = zeros(Int, sum(cntEq))
  LinG = zeros(Tmsh, sum(cntEq))
  Flow = rand(numCnstr) + 0im
  Fupp = rand(numCnstr) + 0im
  fncidx, ptr = setLinearRootConstraints!(map, di, iLfun, jLvar, LinG, Flow,
                                          Fupp, fncidx, ptr)
  @test ( fncidx )== numCnstr+1
  @test ( ptr )== sum(cntEq)+1
  ptr = 1
  for k = 1:numCnstr
    for i = 1:cntEq[k]
      @test ( iLfun[ptr] )== k
      ptr += 1
    end
  end
  for k = 1:3
    @test ( LinG[k] )== 1.0
  end
  @test ( LinG[4] )== 1.0
  @test ( LinG[5] )== -1.0
  @test ( LinG[6] )== 1.0
  @test ( LinG[7] )== -1.0
end