function mp_contignecy(data_ip, Total_sample, total_yr, br_no, base_weight_ip, weights)
    Cont_list = Any[]
    base_list = Any[] #(base_nw, cont_nw, cont_br)
    year_base = Any[]
    base_weight = base_weight_ip
    st_nw = 0
    end_nw = 0
    for i = 1:Total_sample*total_yr
         if i == 1
            st_nw = 2
            end_nw = 2
         else
             st_nw = end_nw + 2
             end_nw = end_nw + 2
        end
        push!(base_list, st_nw - 1)
        base_weight[st_nw - 1] = weights[i]
        for n = st_nw:end_nw
            push!(Cont_list, (st_nw - 1, n, br_no))
            data_ip["nw"]["$n"]["branchdc"]["$br_no"]["rateA"] = 0
        end
    end

    return data_ip, Cont_list, base_list, base_weight
end

function mp_contignecy_nocl(data_ip, Total_sample, br_no)
    Cont_list = Any[]
    base_list = Any[] #(base_nw, cont_nw, cont_br)
    st_nw = 0
    end_nw = 0
    for i = 1:Total_sample
         if i == 1
            st_nw = 2
            end_nw = 2
         else
             st_nw = end_nw + 2
            end_nw = end_nw + 2
        end
        push!(base_list, st_nw - 1)
        for n = st_nw:end_nw
            push!(Cont_list, (st_nw - 1, n, br_no))
            data_ip["nw"]["$n"]["branchdc"]["$br_no"]["rateA"] = 0
         end
    end
    return data_ip, Cont_list, base_list
end

function year_base_networks(Total_sample, total_yr, base_list)
    base_list_no = 1
    year_base =  Dict([(yr, []) for yr = 1: total_yr])
    for k = 1:total_yr
            for t = 1:Total_sample
                     push!(year_base[k], base_list[base_list_no])
                    base_list_no += 1
            end
    end
    return year_base
end

function mp_datainputs(data_ip,Total_sample, year, year_num,file)
    kk= 1
    for yr = 1:length(year)
        fname = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output.mat")
        vars = matread(fname)
        fname1 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output_MC.mat")
        vars1 = matread(fname1)
        fname2 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\wind_sample.mat")
        vars2 = matread(fname2)

        file_cl = string("cluster_", year_num[yr], ".jld2")
        cl = load(file_cl)
    for i = 1:Total_sample
        # sample = Int(round(8760*(1-rand())))
          sample = cl["scenarios"][i]
         for tt = 1:2
            data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = vars["Heq"][sample]
            if data_ip["nw"]["$kk"]["reserves"]["2"]["H"] < 0.5
                data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = 0.5
            end
            # data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2] = maximum(vars1["MC1"])*100*scenarios[1,i] for kmeans
            data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2] = vars1["MC1"][sample]*100
            rating = deepcopy(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
            data_ip["nw"]["$kk"]["gen"]["1"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
            data_ip["nw"]["$kk"]["gen"]["2"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
            if occursin("6bus", file)
                data_ip["nw"]["$kk"]["gen"]["4"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
            end
            data_ip["nw"]["$kk"]["gen"]["3"]["pmax"] = vars["Ptotal"][sample]/100
            data_ip["nw"]["$kk"]["load"]["1"]["pd"] = vars["Ptotal"][sample]/100
            # data_ip["nw"]["$kk"]["gen"]["3"]["pmax"] = maximum(vars["Ptotal"])*scenarios[2,i]/100 # for kmeans
            # data_ip["nw"]["$kk"]["load"]["1"]["pd"]  = maximum(vars["Ptotal"])*scenarios[2,i]/100 # for k means
            # what about load?
            kk +=1
        end
    end
    end
    return data_ip
end

    function mp_datainputs_nocl(data_ip, Total_sample, year, year_num, scenario, file)
        kk= 1
        for yr = 1:length(year)
            fname = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output.mat")
            vars = matread(fname)
            fname1 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output_MC.mat")
            vars1 = matread(fname1)
            fname2 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\wind_sample.mat")
            vars2 = matread(fname2)

        for i = 1:Total_sample
            # sample = Int(round(8760*(1-rand())))
            sample = scenario[i]
            for tt = 1:2
                data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = vars["Heq"][sample]
                if data_ip["nw"]["$kk"]["reserves"]["2"]["H"] < 0.5
                    data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = 0.5
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2] = vars1["MC1"][sample]*100
                rating = deepcopy(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                data_ip["nw"]["$kk"]["gen"]["1"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                data_ip["nw"]["$kk"]["gen"]["2"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                if occursin("6bus", file)
                    data_ip["nw"]["$kk"]["gen"]["4"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["pmax"] = vars["Ptotal"][sample]/100
                data_ip["nw"]["$kk"]["load"]["1"]["pd"] = vars["Ptotal"][sample]/100
                # what about load?
                kk +=1
            end
        end
        end
        return data_ip
    end

    function mp_datainputs_nocl_PLdim(data_ip, Total_sample, year, year_num, scenario, file)
        kk= 1
        for yr = 1:length(year)
            fname = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output.mat")
            vars = matread(fname)
            fname1 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output_MC.mat")
            vars1 = matread(fname1)
            fname2 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\wind_sample.mat")
            vars2 = matread(fname2)

        for i = 1:Total_sample
            # sample = Int(round(8760*(1-rand())))
            sample = scenario[i]
            for tt = 1:2
                data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = vars["Heq"][sample]
                if data_ip["nw"]["$kk"]["reserves"]["2"]["H"] < 0.5
                    data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = 0.5
                end
                data_ip["nw"]["$kk"]["gen"]["2"]["cost"][2] = vars1["MC1"][sample]*100
                # rating = deepcopy(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                # data_ip["nw"]["$kk"]["gen"]["1"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                if occursin("6bus", file)
                    data_ip["nw"]["$kk"]["gen"]["4"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                end
                data_ip["nw"]["$kk"]["gen"]["2"]["pmax"] = vars["Ptotal"][sample]/100
                data_ip["nw"]["$kk"]["load"]["1"]["pd"] = vars["Ptotal"][sample]/100
                # what about load?
                kk +=1
            end
        end
        end
        return data_ip
    end

    function mp_contignecy_ne(data_ip, Total_sample, base_weight_ip, weights)
        Cont_list = Any[]
        base_list = Any[] #(base_nw, cont_nw, cont_br)
        base_weight = base_weight_ip
        cont_no = 1
        for i = 1:Total_sample
             if i == 1
                global st_nw = 2
                global end_nw = length(data_ip["nw"]["1"]["branchdc_ne"]) + 1
             else
                st_nw = end_nw + 2
                end_nw = (st_nw - 1) + length(data_ip["nw"]["1"]["branchdc_ne"])
            end
            push!(base_list, st_nw - 1)
            base_weight[st_nw - 1] = weights[i]
            br_no = 1
            for n = st_nw:end_nw
                push!(Cont_list, (st_nw - 1, n, br_no))
                data_ip["nw"]["$n"]["branchdc_ne"]["$br_no"]["rateA"] = 0
                br_no += 1
                cont_no += 1
            end
        end
        return data_ip, Cont_list, base_list, base_weight
    end

    function mp_contignecy_ne_PL(data_ip, Total_sample, br_no, base_weight_ip, weights)
        Cont_list = Any[]
        base_list = Any[] #(base_nw, cont_nw, cont_br)
        base_weight = base_weight_ip
        cont_no = 1
        for i = 1:Total_sample
             if i == 1
                global st_nw = 2
                global end_nw = 2
             else
                st_nw = end_nw + 2
                end_nw = end_nw + 2
            end
            push!(base_list, st_nw - 1)
            base_weight[st_nw - 1] = weights[i]
            for n = st_nw:end_nw
                push!(Cont_list, (st_nw - 1, n, br_no))
                display(n)
                data_ip["nw"]["$n"]["branchdc_ne"]["$br_no"]["rateA"] = 0
                cont_no += 1
            end
        end
        return data_ip, Cont_list, base_list, base_weight
    end

    function mp_datainputs_ne(data_ip,Total_sample, year, year_num,file)
        kk= 1
        cost = Any[]
        for yr = 1:length(year)
            fname = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output.mat")
            vars = matread(fname)
            fname1 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output_MC.mat")
            vars1 = matread(fname1)
            fname2 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\wind_sample.mat")
            vars2 = matread(fname2)
            file_cl = string("cluster_", year_num[yr], ".jld2")
            cl = load(file_cl)
        for i = 1:Total_sample
            # sample = Int(round(8760*(1-rand())))
            sample = cl["scenarios"][i]
             for tt = 1:(length(data_ip["nw"]["1"]["branchdc_ne"])+1)
               data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = vars["Heq"][sample]
                if data_ip["nw"]["$kk"]["reserves"]["2"]["H"] < 0.5
                    data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = 0.5
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2] = vars1["MC1"][sample]*100
                rating = deepcopy(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                data_ip["nw"]["$kk"]["gen"]["1"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                data_ip["nw"]["$kk"]["gen"]["2"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                if occursin("6bus", file)
                    data_ip["nw"]["$kk"]["gen"]["4"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["pmax"] = vars["Ptotal"][sample]/100
                data_ip["nw"]["$kk"]["load"]["1"]["pd"] = vars["Ptotal"][sample]/100
                push!(cost,  data_ip["nw"]["$kk"]["gen"]["3"]["pmax"]*data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2])
                 kk +=1
            end
        end
        end
        return data_ip, cost
    end

    function mp_datainputs_ne_PL(data_ip,Total_sample, year, year_num,file)
        kk= 1
        cost = Any[]
        for yr = 1:length(year)
            fname = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output.mat")
            vars = matread(fname)
            fname1 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output_MC.mat")
            vars1 = matread(fname1)
            fname2 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\wind_sample.mat")
            vars2 = matread(fname2)
            file_cl = string("cluster_", year_num[yr], ".jld2")
            cl = load(file_cl)
        for i = 1:Total_sample
            # sample = Int(round(8760*(1-rand())))
            sample = cl["scenarios"][i]
            for tt = 1:2
                data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = vars["Heq"][sample]
                if data_ip["nw"]["$kk"]["reserves"]["2"]["H"] < 0.5
                    data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = 0.5
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2] = vars1["MC1"][sample]*100
                rating = deepcopy(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                data_ip["nw"]["$kk"]["gen"]["1"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                data_ip["nw"]["$kk"]["gen"]["2"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                if occursin("6bus", file)
                    data_ip["nw"]["$kk"]["gen"]["4"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["pmax"] = vars["Ptotal"][sample]/100
                data_ip["nw"]["$kk"]["load"]["1"]["pd"] = vars["Ptotal"][sample]/100
                # what about load?
                # display("Pgmax")
                display(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                push!(cost,  data_ip["nw"]["$kk"]["gen"]["1"]["pmax"]* data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2])
                 kk +=1
            end
        end
        end
        return data_ip, cost
    end

    function mp_datainputs_ne_r1(data_ip,Total_sample)
        kk= 1
        cost = Any[]
        for yr = 1:length(year)
            fname = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output.mat")
            vars = matread(fname)
            fname1 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\",year[yr],"\\output_MC.mat")
            vars1 = matread(fname1)
            fname2 = string("C:\\Users\\djaykuma\\OneDrive - Energyville\\Freq_TNEP_paper\\MATLAB\\NAT_Results\\wind_sample.mat")
            vars2 = matread(fname2)
            file_cl = string("cluster_", year_num[yr], ".jld2")
            cl = load(file_cl)
        for i = 1:Total_sample
            # sample = Int(round(8760*(1-rand())))
            sample = cl["scenarios"][i]
            # display("sample:$sample")
            for tt = 1:(length(data_sp["branchdc_ne"])+1)
               data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = vars["Heq"][sample]
                if data_ip["nw"]["$kk"]["reserves"]["2"]["H"] < 0.5
                    data_ip["nw"]["$kk"]["reserves"]["2"]["H"] = 0.5
                end
                data_ip["nw"]["$kk"]["gen"]["1"]["cost"][2] = vars1["MC1"][sample]*100
                data_ip["nw"]["$kk"]["gen"]["2"]["cost"][2] = vars1["MC1"][sample]*100
                rating = deepcopy(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                data_ip["nw"]["$kk"]["gen"]["1"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                data_ip["nw"]["$kk"]["gen"]["2"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                if occursin("6bus", file)
                    data_ip["nw"]["$kk"]["gen"]["4"]["pmax"] = vars2["winddata"][sample]/maximum(vars2["winddata"])*rating
                end
                data_ip["nw"]["$kk"]["gen"]["3"]["pmax"] = vars["Ptotal"][sample]/100
                data_ip["nw"]["$kk"]["load"]["1"]["pd"] = vars["Ptotal"][sample]/100
                # what about load?
                # display(data_ip["nw"]["$kk"]["gen"]["1"]["pmax"])
                push!(cost,  data_ip["nw"]["$kk"]["gen"]["3"]["pmax"]*data_ip["nw"]["$kk"]["gen"]["3"]["cost"][2])
                 kk +=1
            end
        end
        end
        return data_ip, cost
    end

    function mp_contignecy_nocl_sensitivity(data_ip, Total_sample, br_no)
        Cont_list = Any[]
        base_list = Any[] #(base_nw, cont_nw, cont_br)
        st_nw = 0
        end_nw = 0
        for i = 1:Total_sample
             if i == 1
                st_nw = 2
                end_nw = 2
             else
                 st_nw = end_nw + 2
                end_nw = end_nw + 2
            end
            push!(base_list, st_nw - 1)
            for n = st_nw:end_nw
                push!(Cont_list, (st_nw - 1, n, br_no))
                data_ip["nw"]["$n"]["branchdc"]["$br_no"]["rateA"] = 0
             end
        end
        return data_ip, Cont_list, base_list
    end


abstract type SDEIntegratorCache{DT,D,M} <: IntegratorCache{DT,D} end
abstract type PSDEIntegratorCache{DT,D,M} <: IntegratorCache{DT,D} end


############################################################
# judiWavefield ##############################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: June 2017

export judiWavefield, judiWavefieldException, judiDFTwavefield, muteWavefield, dump_wavefield, fft_wavefield

############################################################

mutable struct judiWavefield{vDT<:Number} <: joAbstractLinearOperator{vDT,vDT}
	name::String
	m::Integer
	n::Integer
	info::Info
    dt
	data
end

mutable struct judiWavefieldException <: Exception
	msg :: String
end

############################################################

## outer constructors

"""
judiWavefield
        name::String
        m::Integer
        n::Integer
        info::Info
        dt::Real
        data

Abstract vector for seismic wavefields.

Constructors
============

Construct wavefield vector from an info structure, a cell array of wavefields and the computational \\
time step dt:

    judiWavefield(info, nt, data)

"""
function judiWavefield(info,dt::Real,data::Union{Array, PyCall.PyObject, String};  vDT::DataType=Float32)

	# length of vector
	m = info.n * sum(info.nt)
	n = 1
	dataCell = [vDT.(data) for j=1:info.nsrc]

	return judiWavefield{vDT}("judiWavefield",m,n,info, Float32(dt), dataCell)
end

function judiWavefield(info, dt::Real, data::Union{Array{Any,1}, Array{Array{T, N}, 1}};vDT::DataType=Float32) where {T, N}
	# length of vector
	nsrc = length(data)
	nsrc != info.nsrc && throw("Different number of sources in info ($(info.nsrc)) and data array ($nsrc)")
	m = info.n * sum(info.nt)
	n = 1
	T != Float32 && (data = tof32.(data))
	return judiWavefield{vDT}("judiWavefield",m,n,info, Float32(dt), data)
end

####################################################################
## overloaded Base functions

# conj(jo)
conj(A::judiWavefield{vDT}) where vDT =
	judiWavefield{vDT}("conj("*A.name*")",A.m,A.n,A.info,A.dt,A.data)

# transpose(jo)
transpose(A::judiWavefield{vDT}) where vDT =
	judiWavefield{vDT}(""*A.name*".'",A.n,A.m,A.info,A.dt,A.data)

# adjoint(jo)
adjoint(A::judiWavefield{vDT}) where vDT = conj(transpose(A))

####################################################################

function vcat(ai::Vararg{judiWavefield{avDT}, N}) where {avDT, N}
	N == 1 && (return ai[1])
	N > 2 && (return vcat(ai[1], vcat(ai[2:end]...)))
	a, b = ai
	m = a.m + b.m
	n = 1
	nsrc = a.info.nsrc + b.info.nsrc
	data = Array{Array{avDT, 3}, 1}(undef, nsrc)
	nt = Array{Integer}(undef, nsrc)
	for j=1:a.info.nsrc
		data[j] = a.data[j]
		nt[j] = a.info.nt[j]
	end
	for j=a.info.nsrc+1:nsrc
		data[j] = b.data[j-a.info.nsrc]
		nt[j] = b.info.nt[j-a.info.nsrc]
	end
	info = Info(a.info.n, nsrc, nt)
	return judiWavefield{avDT}("judiWavefield",info.n * sum(info.nt), 1, info, a.dt ,data)
end

function push!(a::judiWavefield{T}, b::judiWavefield{T}) where T
	append!(a.data, b.data)
	a.m += b.m
	a.info.nsrc += b.info.nsrc
end

# DFT operator for wavefields, acts along time dimension
function fft_wavefield(x_in,mode)
	nsrc = x_in.info.nsrc
	nt = size(x_in.data[1], 1)
	if mode==1
		x = similar(x_in, Complex{Float32})
		for i=1:nsrc
			x.data[i] = fft(x_in.data[i], 1)/sqrt(nt)
		end
	elseif mode==-1
		x = similar(x_in, Float32)
		for i=1:nsrc
			x.data[i] = real(ifft(x_in.data[i], 1)) * sqrt(nt)
		end
	end
	return x
end

# Sampling mask to extract wavefields from full vector
subsample(u::judiWavefield,srcnum) = judiWavefield(u.info,u.data[srcnum];vDT=eltype(u))
similar(u::judiWavefield) = 0f0 * u
similar(u::judiWavefield, vDT::DataType) = judiWavefield{vDT}("judiWavefield",u.m, u.n, u.info, u.dt, Array{vDT}.(0f0.*u.data))

# norm
function norm(a::judiWavefield{avDT}, p::Real=2) where avDT
    if p == Inf
        return max([maximum(abs.(a.data[i])) for i=1:a.info.nsrc]...)
    end
    x = 0.f0
    for j=1:a.info.nsrc
        x += Float32(a.dt) * sum(abs.(vec(a.data[j])).^p)
    end
    return x^(1.f0/p)
end

# inner product
function dot(a::judiWavefield{avDT}, b::judiWavefield{bvDT}) where {avDT, bvDT}
	# Dot product for data containers
	size(a) == size(b) || throw(judiWavefieldException("dimension mismatch"))
	dotprod = 0f0
	for j=1:a.info.nsrc
		dotprod += Float32(a.dt) * dot(vec(a.data[j]),vec(b.data[j]))
	end
	return dotprod
end

# abs
function abs(a::judiWavefield{avDT}) where avDT
	b = deepcopy(a)
	for j=1:a.info.nsrc
		b.data[j] = abs.(a.data[j])
	end
	return b
end

function isapprox(x::judiWavefield, y::judiWavefield; rtol::Real=sqrt(eps()), atol::Real=0)
    x.info.nsrc == y.info.nsrc || throw(judiVectorException("Incompatible number of sources"))
    isapprox(x.data, y.data; rtol=rtol, atol=atol)
end

####################################################################################################

isfinite(x::judiWavefield) = all(all(isfinite.(x.data[i])) for i=1:length(x.data))


@with_kw mutable struct Agent
    get_initial_entity::Function # Returns an entity IC
    model::DriverModel # The driver model associated with this agent
    entity_dim::Int # The dimension of the entity
    disturbance_dim::Int # The disturbance dimension of this agent
    entity_to_vec::Function # A Function that converts the agent to a vector of length o_dim
    disturbance_to_vec::Union{Function, Nothing} = nothing # A Function that converts an agent action to a vector of length a_dim
    vec_to_entity::Union{Function, Nothing} = nothing # A Function that converts a vector of length o_dim to an entity
    vec_to_disturbance::Union{Function, Nothing} = nothing # A Function that converts a vector of length a_dim to an action
    disturbance_model = nothing  # Model of the disturbances (supports logpdf, rand and actions)
end

id(a::Agent) = a.get_initial_entity().id

# Construct a regular Blinker vehicle agent
function BlinkerVehicleAgent(get_veh::Function, model::TIDM;
    entity_dim = BLINKERVEHICLE_ENTITY_DIM,
    disturbance_dim=BLINKERVEHICLE_DISTURBANCE_DIM,
    entity_to_vec = BlinkerVehicle_to_vec,
    disturbance_to_vec = BlinkerVehicleControl_to_vec,
    vec_to_entity = vec_to_BlinkerVehicle,
    vec_to_disturbance = vec_to_BlinkerVehicleControl,
    disturbance_model = get_bv_actions())
    Agent(get_veh, model, entity_dim, disturbance_dim, entity_to_vec,
          disturbance_to_vec,  vec_to_entity, vec_to_disturbance, disturbance_model)
end

# Construct a regular adversarial pedestrian agent
function NoisyPedestrianAgent(get_ped::Function, model::AdversarialPedestrian;
    entity_dim = PEDESTRIAN_ENTITY_DIM,
    disturbance_dim = PEDESTRIAN_DISTURBANCE_DIM,
    entity_to_vec = NoisyPedestrian_to_vec,
    disturbance_to_vec = PedestrianControl_to_vec,
    vec_to_entity = vec_to_NoisyPedestrian_fn(DEFAULT_CROSSWALK_LANE),
    vec_to_disturbance = vec_to_PedestrianControl,
    disturbance_model = get_ped_actions())
    Agent(get_ped, model, entity_dim, disturbance_dim, entity_to_vec,
          disturbance_to_vec,  vec_to_entity, vec_to_disturbance, disturbance_model)
end

# Definition of the adversarial driving mdp
mutable struct AdversarialDrivingMDP <: MDP{Scene, Vector{Disturbance}}
    agents::Vector{Agent} # All the agents ordered by (adversaries..., sut, others...)
    vehid2ind::Dict{Int64, Int64} # Dictionary that maps vehid to index in agent list
    num_adversaries::Int64 # The number of adversaries
    roadway::Roadway # The roadway for the simulation
    dt::Float64 # Simulation timestep
    last_observation::Array{Float64} # Last observation of the vehicle state
    disturbance_model # Model used for disturbances. supports `logpdf` and `rand` and `actions`
    γ::Float64 # discount
    ast_reward::Bool # A function that gives action log prob.
    no_collision_penalty::Float64 # penalty for not getting a collision (for ast reward)
    scale_reward::Bool #whether or not to scale the AST reward
    end_of_road::Float64 # specify an early end of the road
end

# Constructor
function AdversarialDrivingMDP(sut::Agent, adversaries::Vector{Agent}, road::Roadway, dt::Float64;
                               other_agents::Vector{Agent} = Agent[],
                               γ = 1,
                               ast_reward = false,
                               no_collision_penalty = 1e3,
                               scale_reward = true,
                               end_of_road = Inf,)
    agents = [adversaries..., sut, other_agents...]
    d = Dict(id(agents[i]) => i for i=1:length(agents))
    Na = length(adversaries)
    o = Float64[] # Last observation

    m = combine_disturbance_models(adversaries)
    AdversarialDrivingMDP(agents, d, Na, road, dt, o, m, γ,
                         ast_reward, no_collision_penalty, scale_reward, end_of_road)
end

# Returns the intial state of the mdp simulator (as a Deteriministic distribution)
function POMDPs.initialstate(mdp::MDP{Scene, A}, rng::AbstractRNG = Random.GLOBAL_RNG) where A
     POMDPModelTools.Deterministic(Scene([a.get_initial_entity(rng) for a in agents(mdp)]))
 end

# The generative interface to the POMDP
function POMDPs.gen(mdp::MDP{Scene, A}, s::Scene, a::A, rng::Random.AbstractRNG = Random.GLOBAL_RNG) where A
    mdp.last_observation = convert_s(AbstractArray, s, mdp)
    sp = step_scene(mdp, s, a, rng)
    r = reward(mdp, s, a, sp)
    (sp=sp, r=r)
end

# Get the reward from the actions taken and the next state
function POMDPs.reward(mdp::AdversarialDrivingMDP, s::Scene, a::Vector{Disturbance}, sp::Scene)
    iscollision = length(sp) > 0 && ego_collides(sutid(mdp), sp)
    if mdp.ast_reward
        isterm = isterminal(mdp, sp)
        r = logpdf(mdp, s, a)
        r += iscollision * abs(mdp.no_collision_penalty)
        # r = (isterm && !iscollision)*(-abs(mdp.no_collision_penalty)) + logpdf(mdp, s, a)
        mdp.scale_reward && (r = r / abs(mdp.no_collision_penalty))
        return Float32(r)
    else
        return Float32(iscollision)
    end
end

# Discount factor for the POMDP (Set to 1 because of the finite horizon)
POMDPs.discount(mdp::MDP{Scene, A}) where A = mdp.γ

# The simulation is terminal if there is collision with the ego vehicle or if the maximum simulation time has been reached
POMDPs.isterminal(mdp::MDP{Scene, A}, s::Scene) where A = !(sutid(mdp) in s)|| any_collides(s)

# Define the set of actions, action index and probability
POMDPs.actions(mdp::AdversarialDrivingMDP) = get_actions(mdp.disturbance_model)
POMDPs.actionindex(mdp::AdversarialDrivingMDP, a::Vector{Disturbance}) = Int32(get_actionindex(mdp.disturbance_model, a))

# The default disturbance policy according to the disturbance models
#TODO: Switch the "Function policy" to something that can get logpdf?
default_policy(mdp::AdversarialDrivingMDP, rng::AbstractRNG = Random.GLOBAL_RNG) = FunctionPolicy((s) -> rand(rng, mdp.disturbance_model))

# POMDPs.actionindex(mdp::AdversarialDrivingMDP, a::Vector{Disturbance}) = findfirst(actions(mdp) .== a)
Distributions.logpdf(mdp::AdversarialDrivingMDP, a::Vector{Disturbance}) = logpdf(mdp.disturbance_model, a, mdp)
Distributions.logpdf(mdp::AdversarialDrivingMDP, s::Scene, a::Vector{Disturbance}) = logpdf(mdp, a)
Distributions.logpdf(mdp::AdversarialDrivingMDP, h::SimHistory) = sum([logpdf(mdp, s, a) for (s,a) in eachstep(h, (:s, :a))])
Distributions.logpdf(mdp::AdversarialDrivingMDP, as::Vector) = sum([logpdf(mdp, a) for a in as])

## Helper functions

# Step the scene forward by one timestep and return the next state
function step_scene(mdp::AdversarialDrivingMDP, s::Scene, actions::Vector{Disturbance}, rng::AbstractRNG = Random.GLOBAL_RNG)
    entities = []

    # Loop through the adversaries and apply the instantaneous aspects of their disturbance
    for (adversary, action) in zip(adversaries(mdp), actions)
        update_adversary!(adversary, action, s)
    end

    # Loop through the vehicles in the scene, apply action and add to next scene
    for (i, veh) in enumerate(s)
        m = model(mdp, veh.id)
        observe!(m, s, mdp.roadway, veh.id)
        a = rand(rng, m)
        bv = Entity(propagate(veh, a, mdp.roadway, mdp.dt), veh.def, veh.id)
        !end_of_road(bv, mdp.roadway, mdp.end_of_road) && push!(entities, bv)
    end
    isempty(entities) ? Scene(typeof(sut(mdp).get_initial_entity())) : Scene([entities...])
end

# Returns the list of agents in the mdp
agents(mdp::MDP{Scene,A}) where A = mdp.agents

# Returns the list of adversaries in the mdp
adversaries(mdp::MDP{Scene,A}) where A = view(mdp.agents, 1:mdp.num_adversaries)

# Returns the model associated with the vehid
model(mdp::MDP{Scene,A}, vehid::Int) where A = mdp.agents[mdp.vehid2ind[vehid]].model

# Returns the system under test
sut(mdp::MDP{Scene,A}) where A  = mdp.agents[mdp.num_adversaries + 1]

# Returns the sut id
sutid(mdp::MDP{Scene,A}) where A = id(sut(mdp))

function update_adversary!(adversary::Agent, action::Disturbance, s::Scene)
    index = findfirst(id(adversary), s)
    isnothing(index) && return nothing # If the adversary is not in the scene then don't update
    adversary.model.next_action = action # Set the adversaries next action
    veh = s[index] # Get the actual entity
    state_type = typeof(veh.state) # Find out the type of its state
    s[index] =  Entity(state_type(veh.state, noise = action.noise), veh.def, veh.id) # replace the entity in the scene
end

## SUT driving MDP
mutable struct DrivingMDP <: MDP{Scene, BlinkerVehicleControl}
    agents::Vector{Agent} # All the agents ordered by (adversaries..., sut, others...)
    vehid2ind::Dict{Int64, Int64} # Dictionary that maps vehid to index in agent list
    num_adversaries::Int64 # The number of adversaries
    roadway::Roadway # The roadway for the simulation
    dt::Float64 # Simulation timestep
    last_observation::Array{Float64} # Last observation of the vehicle state
    γ::Float64 # discount
    end_of_road::Float64 # Early stopping of road
    per_timestep_penalty::Float64
    v_des::Float64
end

# Constructor
function DrivingMDP(sut::Agent, adversaries::Vector{Agent}, road::Roadway, dt::Float64; γ = 1, end_of_road = Inf, per_timestep_penalty = 0, v_des = 25)
    agents = [adversaries..., sut]
    d = Dict(id(agents[i]) => i for i=1:length(agents))
    DrivingMDP(agents, d, length(adversaries), road, dt, Float64[], γ, end_of_road, per_timestep_penalty, v_des)
end

# Get the reward from the actions taken and the next state
function POMDPs.reward(mdp::DrivingMDP, s::Scene, a::BlinkerVehicleControl, sp::Scene)
    id = sutid(mdp)
    v = vel(get_by_id(s, id))

    r = -abs(mdp.per_timestep_penalty)
    # If the simulation ends but the SUT is not at the end of the road, big penalty
    if ego_collides(id, sp)
        r += -1
    elseif isterminal(mdp, sp)
        r += 1
    end
    r += -.1 * (v < 0)
    r += -0.001 * abs(v  - mdp.v_des)
    r
end

function step_scene(mdp::DrivingMDP, s::Scene, action::BlinkerVehicleControl, rng::AbstractRNG = Random.GLOBAL_RNG)
    entities = []
    sid = sutid(mdp)

    # Choose random actions for the adversaries
    for adversary in adversaries(mdp)
        adv_action = rand(rng, adversary.disturbance_model, adversary.vec_to_disturbance)
        update_adversary!(adversary, adv_action, s)
    end

    # Loop through the vehicles in the scene, apply action and add to next scene
    for (i, veh) in enumerate(s)
        if veh.id == sid # for the sut, use the prescribed action
            a = action
        else # For the other vehicles use their  model
            m = model(mdp, veh.id)
            observe!(m, s, mdp.roadway, veh.id)
            a = rand(rng, m)
        end
        bv = Entity(propagate(veh, a, mdp.roadway, mdp.dt), veh.def, veh.id)
        !end_of_road(bv, mdp.roadway, mdp.end_of_road) && push!(entities, bv)
    end
    isempty(entities) ? Scene(typeof(sut(mdp).get_initial_entity())) : Scene([entities...])
end

# Define the set of actions, action index and probability
POMDPs.actions(mdp::DrivingMDP) = [BlinkerVehicleControl(a = -4.), BlinkerVehicleControl(a= -2.), BlinkerVehicleControl(a = 0.), BlinkerVehicleControl(a = 1.5), BlinkerVehicleControl(a = 3.)]
POMDPs.actionindex(mdp::DrivingMDP, a::BlinkerVehicleControl) = findfirst([a] .== actions(mdp))


using OrderedCollections

abstract type PlottingRecipe end

Plots.plot(pr::PlottingRecipe, args...) = plot!(plot(), pr, args...)

# =============================================================================
# hospitalizations

struct ObservationsPlottingRecipe{Tsd, Ted, To, Te, Tl} <: PlottingRecipe
    startdate ::Tsd
    enddate   ::Ted
    observed  ::To
    expected  ::Te
    label     ::Tl
end

function expected(data, gp, label)
    @assert label in ["cases", "hospitalizations", "deaths"]
    dates  = data.dates
    if label == "cases"
        values = gp.expected_daily_cases
        return (; dates, values )
    elseif label == "hospitalizations"
        values = gp.expected_daily_hospit
        return (; dates, values )
    else
        values = gp.expected_daily_deaths
        return (; dates, values )
    end
end

function observed(data, label)
    @assert label in ["cases", "hospitalizations", "deaths"]
    dates  = data.hospit.date
    if label == "cases"
        values = data.cases.country
        return (; dates, values )
    elseif label == "hospitalizations"
        values = data.hospit.country
        return (; dates, values )
    else
        values = data.deaths.country
        return (; dates, values )
    end
end

function startdate(data::National.Data, label::String)
    @assert label in ["cases", "hospitalizations", "deaths"]
    s = if label == "cases"
        data.turing_data.casemodel.start
    elseif label == "hospitalizations"
        data.turing_data.hospitmodel.start
    else
        data.turing_data.deathmodel.start
    end
    return data.dates[s]
end

function enddate(data::National.Data, label::String)
    @assert label in ["cases", "hospitalizations", "deaths"]
    s = if label == "cases"
        data.turing_data.casemodel.stop
    elseif label == "hospitalizations"
        data.turing_data.hospitmodel.stop
    else
        data.turing_data.deathmodel.stop
    end
    return data.dates[s]
end

function ObservationsPlottingRecipe(data::National.Data, gp, label)
    e = expected(data, gp, label)
    o = observed(data, label)
    sd = startdate(data, label)
    ed = enddate(data, label)
    return ObservationsPlottingRecipe( sd, ed, o, e, label)
end

function Plots.plot!(p::Plots.Plot, r::ObservationsPlottingRecipe)
    o  = r.observed
    e  = r.expected
    ed = r.enddate
    sd = r.startdate

    plot!(p, o.dates, o.values, α=0.5, lc=:match, lab="observed $(r.label)", c=:midnightblue, lw=4, ylab="cases")
    vline!(p, [ed], lab="end observations", lw=2, lc=:black, hover="$ed")
    vline!(p, [sd], lab="start observations", lw=2, lc=:black, hover="$sd", ls=:dash)
    plot_confidence_timeseries!(p, e.dates, e.values; label = "expected $(r.label)") #Dict(hover=>strdates)
end

# ============================================================================
# seropos

struct SeroPlottingRecipe{E, O, L} <: National.PlottingRecipe
    expected  ::E
    observed  ::O
    label     ::L
end

function SeroPlottingRecipe(data::National.Data, gp, label)
    @unpack dates, mean, CIs, delay, population = data.turing_data.seromodel
    observed   = let
        mean *= 100
        error = [(m-l*100, u*100-m) for (m,(l,u)) in zip(mean,CIs)]
        (; dates, mean, error )
    end

    expected = let
        values = gp.expected_seropos  ./ population .* 100
        dates  = data.dates + Day(delay)
        (; dates, values)
    end
    SeroPlottingRecipe( expected, observed, label)
end

function Plots.plot!(p::Plots.Plot, r::SeroPlottingRecipe)
    e = r.expected
    o = r.observed

    plot_confidence_timeseries!(p, e.dates, e.values; r.label)
    y  = o.mean
    ye = o.error
    scatter!(p, o.dates, y, yerror=ye, hover=["$d: $val" for (d, val) in zip(o.dates, y)])
end
# ============================================================================
# reproduction number

struct RtPlottingRecipe{Tlo, Ted, Te, Tl} <: National.PlottingRecipe
    lockdown  ::Tlo
    enddates  ::Ted
    expected  ::Te
    label     ::Tl
end

function RtPlottingRecipe(data::National.Data, gp, label)
    dates      = data.dates
    values     = label == "reproduction number" ? gp.Rt : gp.effective_Rt
    expected   = (; dates, values )
    lockdown   = Date(data.lockdown)
    enddates   = Date(data.observations_end)
    RtPlottingRecipe( lockdown, enddates, expected, label)
end

function Plots.plot!(p::Plots.Plot, r::RtPlottingRecipe)
    e  = r.expected
    ed = r.enddates
    lo = r.lockdown

    vline!(p, [ed], lab="end observations", lw=2, lc=:black, hover="$ed")
    vline!(p, [lo], lab="lockdown", lw=2, lc=:black, hover="$lo", ls=:dash)
    plot_confidence_timeseries!(p, e.dates, e.values; r.label) #Dict(hover=>strdates)
end
# ============================================================================
# infection ascertainment rate

struct IARPlottingRecipe{Tso, Ted, Te} <: National.PlottingRecipe
    startdate ::Tso
    enddates  ::Ted
    expected  ::Te
end

function IARPlottingRecipe(data::National.Data, generated_posterior, args...)
    dates      = data.dates
    values     = generated_posterior.iar
    expected   = (; dates, values )
    startdate  = dates[ data.turing_data.iar_start_idx ]
    enddate    = Date(data.observations_end)
    IARPlottingRecipe( startdate, enddate, expected)
end

function Plots.plot!(p::Plots.Plot, r::IARPlottingRecipe)
    e  = r.expected
    ed = r.enddates
    sd = r.startdate

    vline!(p, [ed], lab="end observations", lw=2, lc=:black, hover="$ed")
    vline!(p, [sd], lab="start case observations", lw=2, lc=:black, hover="$sd", ls=:dash)
    plot_confidence_timeseries!(p, e.dates, e.values; label = "infections ascertainment rate") #Dict(hover=>strdates)
end
# ============================================================================
# region plot

struct OverviewPlottingRecipe{Tr,Tt} <: National.PlottingRecipe
    recipes::Tr
    titles ::Tt
end

posterior2recipe = OrderedDict(
    :expected_daily_cases  => National.ObservationsPlottingRecipe,
    :expected_daily_hospit => National.ObservationsPlottingRecipe,
    :expected_daily_deaths => National.ObservationsPlottingRecipe,
    :Rt                    => National.RtPlottingRecipe,
    :iar                   => National.IARPlottingRecipe,
    :effective_Rt          => National.RtPlottingRecipe,
    :expected_seropos      => National.SeroPlottingRecipe,
)

posterior2label = OrderedDict(
    :expected_daily_cases  => "cases",
    :expected_daily_hospit => "hospitalizations",
    :expected_daily_deaths => "deaths",
    :Rt                    => "reproduction number",
    :iar                   => "infections ascertainment rate",
    :effective_Rt          => "effective reproduction number",
    :expected_seropos      => "sero positive"
)

function OverviewPlottingRecipe(data::National.Data, generated_posterior)
    ks = keys(generated_posterior)
    recipes = Vector{National.PlottingRecipe}()
    labels  = Vector{String}()
    for (k,recipe) in posterior2recipe
        if k in ks
            label = posterior2label[k]
            r     = recipe(data, generated_posterior, label)
            push!(recipes, r)
            push!(labels, label)
        end
    end
    OverviewPlottingRecipe(recipes, labels)
end

function Plots.plot(r::OverviewPlottingRecipe)
    plots  = Vector{Plots.Plot}()
    nplots = length(r.recipes)
    for (recipe, title) in zip(r.recipes, r.titles)
        p = plot(; xaxis = true, legend = :outertopright, title)
        plot!(p, recipe)
        push!(plots, p)
    end
    plot(plots..., layout=(nplots,1), size=(1500, nplots*250), sharex=true, link=:x)
end

Plots.plot(data::National.Data, generated_posterior) =
    plot( OverviewPlottingRecipe(data, generated_posterior) )
# Plots.plot(data::National.Data, generated_posterior) = plot!(plot(), data, generated_posterior)

## ============================================================================
# out dated
function plot_results(data, expected_daily_cases, expected_daily_deaths, expected_daily_hospit, Rt, cumulative_cases, iar)
    l = @layout [a; b; c; d; e; f]

    xlims   = (data.dates[1], data.dates[end])
    # p1 = plot(; xaxis = false, xlims, legend = :outertopright)
    # plot_death_counts!(p1, data)
    p2 = plot(; xaxis = false, xlims, legend = :outertopright)
    plot_deaths_posterior!(p2, data, expected_daily_deaths)
    p3 = plot(; xaxis = false, xlims, legend = :outertopright)
    plot_cases_posterior!(p3, data, expected_daily_cases)
    p4 = plot(; xaxis = false, xlims, legend = :outertopright)
    plot_hospit_posterior!(p4, data, expected_daily_hospit)
    # plot_cumdeaths_posterior!(p5, data, edeaths)
    p5 = plot(; xaxis = true, xlims, legend = :outertopright)
    plot_cumcases_posterior!(p5, data, cumulative_cases)
    p6 = plot(; xaxis = false, xlims, legend = :outertopright)
    plot_Rt_posterior!(p6, data, Rt)
    p7 = plot(; xaxis = false, xlims, legend = :outertopright)
    plot_iar_posterior!(p7, data, iar)
    tickdates = data.dates[1]:Month(1):data.dates[end]
    p = plot(p3, p4, p2, p5, p6, p7, layout=l, size=(1500, 1000), sharex=true, link=:x, xticks = tickdates, legend=false)

    #dates = data.dates[1:20:end] .|> Symbol .|> String
    #Plots.xticks!(1:20:ntot, dates, xrotation=45)
    p
end

function plot_death_counts!(p, data)
    title!(p, "death counts")
    cases_start = data.num_observations - data.turing_data.cases_start
    tstart = data.dates[data.turing_data.epidemic_start]
    tNobs  = data.dates[data.num_observations]
    tcases = data.dates[cases_start]
    plot!(p, data.dates, data.deaths, lc=:match, c=:midnightblue, label=nothing, lw=4) #, label="Observed daily deaths"
    vline!(p, [tstart], label="epidemic start", linewidth=2, lc=:green, hover="$(tstart)")
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
    vline!(p, [tcases], label="start with case observations", linewidth=2, lc=:black, ls=:dash, hover="$(tcases)")
    # xlims!(p, data.dates[1], data.dates[end])
end

function plot_deaths_posterior!(p, data, expected_daily_deaths)
    edeaths = hcat(expected_daily_deaths...)
    tNobs  = data.dates[data.num_observations]
    title!(p, "posterior death counts")
    plot!(p, data.dates, data.deaths, alpha=0.5, lc=:match, c=:midnightblue, label=nothing, lw=4) #, label="Recorded daily deaths (observed)"
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
    plot_confidence_timeseries!(p, data.dates, edeaths; label = "Expected daily deaths") #Dict(hover=>strdates)
end

function plot_hospit_posterior!(p, data, expected_daily_hospit)
    ehospit  = hcat(expected_daily_hospit...)
    tNobs  = data.dates[data.num_observations]
    title!(p, "posterior hospitalization counts")
    plot!(p, data.dates, data.hospit, alpha=0.5, lc=:match, label=nothing, c=:midnightblue, lw=4) #, label="Recorded daily hospitalizations (observed)"
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
    plot_confidence_timeseries!(p, data.dates, ehospit; label = "Expected daily hospitalizations") #Dict(hover=>strdates)
end

function plot_Rt_posterior!(p, data, Rt)
    eRt    = hcat(Rt...)
    tNobs  = data.dates[data.num_observations]
    title!(p, "R_t")
    plot_confidence_timeseries!(p, data.dates, eRt; no_label = true)
    qs = [quantile(v, [0.025, 0.975]) for v in eachrow(eRt)]
    lq, hq = (eachrow(hcat(qs...))..., )
    ylims!(p, 0, maximum(hq) + 0.1)
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
end

function plot_iar_posterior!(p, data, iar)
    iar    = hcat(iar...)
    size(iar, 1) == 1 && (iar = repeat(iar, outer=( length(data.dates), 1)) )
    tNobs  = data.dates[data.num_observations]
    title!(p, "infection ascertainment rate")
    plot_confidence_timeseries!(p, data.dates, iar; no_label = true)
    qs = [quantile(v, [0.025, 0.975]) for v in eachrow(iar)]
    lq, hq = (eachrow(hcat(qs...))..., )
    ylims!(p, 0, maximum(hq) + 0.1)
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
end

function plot_cases_posterior!(p, data, expected_daily_cases)
    ecases  = hcat(expected_daily_cases...)
    cases_start = data.num_observations - data.turing_data.num_case_obs
    tcases = data.dates[cases_start]
    tNobs  = data.dates[data.num_observations]
    title!(p, "posterior case counts")
    plot!(p, data.dates, data.cases, alpha=0.5, lc=:match, label=nothing, c=:midnightblue, lw=4)
    vline!(p, [tcases], label="start with case observations", linewidth=2, lc=:black, ls=:dash, hover="$(tcases)")
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
    # plot!(p, [d1, d2, d2, d1], [low, low, up, up], fc=:orange, fα=0.4, leg=nothing, hover="$d1 - $d2", seriestype=:shape, lc=:orange)
    plot_confidence_timeseries!(p,data.dates, ecases; label = "observed daily cases")
    # plot_confidence_timeseries!(p,data.dates, ecases; label = "expected daily cases")
end

function plot_cumdeaths_posterior!(p, data, expected_daily_deaths)
    edeaths = hcat(expected_daily_deaths...)
    tNobs  = data.dates[data.num_observations]
    title!(p, "posterior cumulative death counts")
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
    plot_confidence_timeseries!(p, data.dates, cumsum(edeaths; dims=1); label = "Expected cumulative deaths")
    Plots.plot!(p, data.dates, cumsum(data.deaths; dims=1), label="Recorded cumulative deaths", color=:red)
end

function plot_cumcases_posterior!(p, data, cumulative_cases)
    ecumul  = hcat(cumulative_cases...)
    tNobs  = data.dates[data.num_observations]
    title!(p, "posterior cumulative cases counts")
    vline!(p, [tNobs], label="end of observations", linewidth=2, lc=:black, hover="$(tNobs)")
    plot_confidence_timeseries!(p, data.dates, ecumul ./ data.pop .* 100; label = "Expected cumulative cases")
    # plot_confidence_timeseries!(p, data.dates, cumsum(ecases, dims=1) ./ data.pop .* 100; label = "test")
    Plots.plot!(p, data.dates, cumsum(data.cases; dims=1) ./ data.pop .* 100, label="Recorded cumulative cases", color=:red)
    plot_seroprev!(p, data)

    # plot_seroprev!(p, Date("2020-06-09"), Date("2020-06-10"), 1.2, .7, 1.7)
    # plot_seroprev!(p, Date("2020")+Week(34), Date("2020")+Week(36), 2.2, 1.8, 2.6)
    # plot_seroprev!(p, Date("2020")+Week(38), Date("2020")+Week(51), 3.9, 3.3, 4.6)
    #TODO include seroprevalence study: 34 (1.2%) positive with 95% CI: 0,7%-1,7% in week 20-21
    #TODO include seroprevalence study: 2.2% positive with 95% CI: 1,8% - 2,6% in week 34-36
    #TODO include seroprevalence study: 3.9% positive with 95% CI: 3,3% - 4,6% in week 38-51
end

function plot_seroprev!(p, data)
    seroprev = data.seroprev
    le = [s.country-s.CI[1] for s in eachrow(seroprev)] * 100
    ue = [s.CI[2]-s.country for s in eachrow(seroprev)] * 100
    scatter!(p, seroprev.date, seroprev.country * 100, yerror=(le,ue), hover=["$d" for d in seroprev.date])
end

function plot_seroprev!(p, d1, d2, med, up, low)
    # rec = [(d1, low), (d1, low), (d2, up), (d2, up)]
    plot!(p, [d1, d2, d2, d1], [low, low, up, up], fc=:orange, fα=0.4, leg=nothing, hover="$d1 - $d2", seriestype=:shape, lc=:orange)
    plot!(p, [d1,d2], [med, med], lc=:orange, leg=nothing, lw=4, hover="$med")
end


import StatsBase:sample
export AbstractEnvironmentModel, AbstractSampleModel, AbstractDistributionModel,
       update!, getstates, getactions

"""
    AbstractEnvironmentModel

A supertype of environment models. An environment model is used to simulate the behavior of an environment.
"""
abstract type AbstractEnvironmentModel end

"""
    AbstractSampleModel <: AbstractEnvironmentModel

An `AbstractSampleModel` can be used to generate a random turn info to simulate an environment.

| Required Methods| Brief Description |
|:----------------|:------------------|
| `sample(model)` | Return a turn info (s, a, r, d, s′) |
| `update!(model, buffer, learner)` | Update the `model` given a `buffer` and a `learner` |
"""
abstract type AbstractSampleModel <: AbstractEnvironmentModel end

"""
    AbstractDistributionModel <: AbstractEnvironmentModel

An `AbstractDistributionModel` can return all the possible following up states, rewards, isdone info and corresponding probability.

| Required Methods| Brief Description |
|:----------------|:------------------|
| `sample(model, s, a)` | Return a vector of (s′, r, d, p) |
| `update!(model, buffer, learner)` | Update the `model` given a `buffer` and a `learner` |
| `getstates(model)` | Get the number of states of the model |
| `getactions(model)` | Get the number of actions of the model |
"""
abstract type AbstractDistributionModel <: AbstractEnvironmentModel end

function update! end
function sample end
function getstates end
function getactions end


module MathTeXEngine
# Adapted from matplotlib mathtext engine
# https://github.com/matplotlib/matplotlib/blob/master/lib/matplotlib/_mathtext.py

using AbstractTrees
using Automa
using FreeTypeAbstraction
using LaTeXStrings

import Automa.RegExp: @re_str
import DataStructures: Stack
import FreeTypeAbstraction:
    ascender, boundingbox, descender, get_extent, hadvance, inkheight, inkwidth,
    leftinkbound, rightinkbound, topinkbound, bottominkbound
import GeometryBasics: Point2f0
import REPL.REPLCompletions: latex_symbols

const re = Automa.RegExp

export TeXExpr, texparse
export TeXElement, TeXChar, VLine, HLine, generate_tex_elements

include("parser/texexpr.jl")
include("parser/command_data.jl")
include("parser/parser.jl")

include("engine/fonts.jl")
include("engine/texelements.jl")
include("engine/layout.jl")

end # module


module v2_tiny_voc

import ...flipdict, ...Settings

#20 classes on Voc dataset
# 2 dictionaries to access number<->class by O(1)
const namesdic = Dict("aeroplane"=>1,"bicycle"=>2,"bird"=>3, "boat"=>4,
            "bottle"=>5,"bus"=>6,"car"=>7,"cat"=>8,"chair"=>9,
            "cow"=>10,"diningtable"=>11,"dog"=>12,"horse"=>13,"motorbike"=>14,
            "person"=>15,"pottedplant"=>16,"sheep"=>17,"sofa"=>18,"train"=>19,"tvmonitor"=>20)

#Yolo V2-tiny pre-trained boxes width and height
const anchors = [(1.08,1.19),  (3.42,4.41),  (6.63,11.38),  (9.42,5.11),  (16.62,10.52)]

function load(;minibatch_size=1)
    return Settings(
        dataset_description = "VOC2007",
        source = "https://pjreddie.com/media/files/yolov2-tiny-voc.weights",
        weights_filepath = joinpath(@__DIR__,"v2_tiny_voc.weights"),
        image_shape = (416, 416),
        image_channels = 3,
        namesdic = namesdic,
        numsdic = flipdict(namesdic), #Flip the key value pair to help lookup by nums
        anchors = anchors,
        num_classes = length(namesdic),
        minibatch_size = minibatch_size,
        grid_x = 13,
        grid_y = 13,
        cell_bboxes = 5
    )
end

end #module


using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    include("../src/main_model.jl")

    m = getpage(scenario, use_permafrost, use_seaice)
    run(m)

    if scenario == "NDCs"
        while m[:Discontinuity,:occurdis_occurrencedummy] != [0.,0.,0.,0.,0.,0.,1.,1.,1.,1.]
            println(m[:Discontinuity,:occurdis_occurrencedummy])
            run(m)
        end
    else
        while m[:Discontinuity,:occurdis_occurrencedummy] != [0.,0.,0.,0.,0.,0.,0.,0.,0.,0.]
            println(m[:Discontinuity,:occurdis_occurrencedummy])
            run(m)
        end
    end

    if updatetestdata
        include("../src/utils/save_parameters.jl")
        savepagedata(m, :Population, :pop_population, "test/validationdata/$valdir/pop_population.csv")
        savepagedata(m, :GDP, :cons_percap_consumption, "test/validationdata/$valdir/cons_percap_consumption.csv")
        savepagedata(m, :GDP, :gdp, "test/validationdata/$valdir/gdp.csv")
        savepagedata(m, :Discontinuity, :rcons_per_cap_DiscRemainConsumption, "test/validationdata/$valdir/rcons_per_cap_DiscRemainConsumption.csv")
        savepagedata(m, :TotalAbatementCosts, :tct_per_cap_totalcostspercap, "test/validationdata/$valdir/tct_per_cap_totalcostspercap.csv")
        savepagedata(m, :AdaptiveCostsEconomic, :ac_adaptivecosts, "test/validationdata/ac_adaptationcosts_economic.csv")
        savepagedata(m, :AdaptiveCostsNonEconomic, :ac_adaptivecosts, "test/validationdata/ac_adaptationcosts_noneconomic.csv")
        savepagedata(m, :AdaptiveCostsSeaLevel, :ac_adaptivecosts, "test/validationdata/$valdir/ac_adaptationcosts_sealevelrise.csv")
        savepagedata(m, :TotalAdaptationCosts, :act_adaptationcosts_total, "test/validationdata/$valdir/act_adaptationcosts_tot.csv")
        savepagedata(m, :TotalAdaptationCosts, :act_percap_adaptationcosts, "test/validationdata/$valdir/act_percap_adaptationcosts.csv")
        savepagedata(m, :SLRDamages, :rcons_per_cap_SLRRemainConsumption, "test/validationdata/$valdir/rcons_per_cap_SLRRemainConsumption.csv")
        savepagedata(m, :SLRDamages, :cons_percap_aftercosts, "test/validationdata/$valdir/cons_percap_aftercosts.csv")
        savepagedata(m, :MarketDamagesBurke, :rcons_per_cap_MarketRemainConsumption, "test/validationdata/$valdir/rcons_per_cap_MarketRemainConsumption.csv")
        savepagedata(m, :MarketDamagesBurke, :rgdp_per_cap_MarketRemainGDP, "test/validationdata/$valdir/rgdp_per_cap_MarketRemainGDP.csv")
        savepagedata(m, :NonMarketDamages, :rcons_per_cap_NonMarketRemainConsumption, "test/validationdata/$valdir/rcons_per_cap_NonMarketRemainConsumption.csv")
        savepagedata(m, :NonMarketDamages, :rgdp_per_cap_NonMarketRemainGDP, "test/validationdata/$valdir/rgdp_per_cap_NonMarketRemainGDP.csv")
        savepagedata(m, :SLRDamages, :rcons_per_cap_SLRRemainConsumption, "test/validationdata/$valdir/rcons_per_cap_SLRRemainConsumption.csv")
        savepagedata(m, :SLRDamages, :rgdp_per_cap_SLRRemainGDP, "test/validationdata/$valdir/rgdp_per_cap_SLRRemainGDP.csv")
        savepagedata(m, :SLRDamages, :igdp_ImpactatActualGDPperCapSLR, "test/validationdata/$valdir/igdp_ImpactatActualGDPperCap_sea.csv")
        savepagedata(m, :SLRDamages, :isat_ImpactinclSaturationandAdaptationSLR, "test/validationdata/$valdir/isat_ImpactinclSaturationandAdaptation_SLRise.csv")
        savepagedata(m, :EquityWeighting, :wtct_percap_weightedcosts, "test/validationdata/$valdir/wtct_percap_weightedcosts.csv")
        savepagedata(m, :EquityWeighting, :pct_percap_partiallyweighted, "test/validationdata/$valdir/pct_percap_partiallyweighted.csv")
        savepagedata(m, :EquityWeighting, :dr_discountrate, "test/validationdata/$valdir/dr_discountrate.csv")
        savepagedata(m, :EquityWeighting, :dfc_consumptiondiscountrate, "test/validationdata/$valdir/dfc_consumptiondiscountrate.csv")
        savepagedata(m, :EquityWeighting, :pct_partiallyweighted, "test/validationdata/$valdir/pct_partiallyweighted.csv")
        savepagedata(m, :EquityWeighting, :pcdt_partiallyweighted_discounted, "test/validationdata/$valdir/pcdt_partiallyweighted_discounted.csv")
        savepagedata(m, :EquityWeighting, :wacdt_partiallyweighted_discounted, "test/validationdata/$valdir/wacdt_partiallyweighted_discounted.csv")
        savepagedata(m, :EquityWeighting, :aact_equityweightedadaptation_discountedaggregated, "test/validationdata/$valdir/aact_equityweightedadaptation_discountedaggregated.csv")
        savepagedata(m, :EquityWeighting, :wit_equityweightedimpact, "test/validationdata/$valdir/wit_equityweightedimpact.csv")
        savepagedata(m, :EquityWeighting, :addt_equityweightedimpact_discountedaggregated, "test/validationdata/$valdir/addt_equityweightedimpact_discountedaggregated.csv")
    end

    # climate component
    temp = m[:ClimateTemperature,:rt_g_globaltemperature]
    temp_compare = readpagedata(m, "test/validationdata/$valdir/rt_g_globaltemperature.csv")
    @test temp ≈ temp_compare rtol = 1e-3

    slr = m[:SeaLevelRise,:s_sealevel]
    slr_compare = readpagedata(m, "test/validationdata/$valdir/s_sealevel.csv")
    @test slr ≈ slr_compare rtol = 1e-2

    # Socio-Economics
    gdp = m[:GDP,:gdp]
    gdp_compare = readpagedata(m, "test/validationdata/$valdir/gdp.csv")
    @test gdp ≈ gdp_compare rtol = 1

    pop = m[:Population,:pop_population]
    pop_compare = readpagedata(m, "test/validationdata/$valdir/pop_population.csv")
    @test pop ≈ pop_compare rtol = 0.001

    # Abatement Costs
    abatement = m[:TotalAbatementCosts,:tct_totalcosts]
    abatement_compare = readpagedata(m, "test/validationdata/$valdir/tct_totalcosts.csv")
    @test abatement ≈ abatement_compare rtol = 1e-2

    # Adaptation Costs
    adaptation = m[:TotalAdaptationCosts,:act_adaptationcosts_total]
    adaptation_compare = readpagedata(m, "test/validationdata/$valdir/act_adaptationcosts_tot.csv")
    @test adaptation ≈ adaptation_compare rtol = 1e-2

    # Damages
    damages = m[:Discontinuity,:rcons_per_cap_DiscRemainConsumption]
    damages_compare = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_DiscRemainConsumption.csv")
    @test damages ≈ damages_compare rtol = 10
    # SLR damages
    slrdamages = m[:SLRDamages,:rcons_per_cap_SLRRemainConsumption]
    slrdamages_compare = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_SLRRemainConsumption.csv")
    @test slrdamages ≈ slrdamages_compare rtol = 0.1
    # Market damages
    mdamages = m[:MarketDamages,:rcons_per_cap_MarketRemainConsumption]
    mdamages_compare = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_MarketRemainConsumption.csv")
    @test mdamages ≈ mdamages_compare rtol = 1
    # NonMarket Damages
    nmdamages = m[:NonMarketDamages,:rcons_per_cap_NonMarketRemainConsumption]
    nmdamages_compare = readpagedata(m, "test/validationdata/$valdir/rcons_per_cap_NonMarketRemainConsumption.csv")
    @test nmdamages ≈ nmdamages_compare rtol = 1

    te = m[:EquityWeighting, :te_totaleffect]
    if scenario == "NDCs"
        te_compare = 1.0320923880568126e9
    elseif scenario == "2 degC Target"
        te_compare = 5.0496576104580307e8
    end
    @test te ≈ te_compare rtol = 1e4
end


export MACLearner

using Flux

"""
    MACLearner(;kwargs...)
   Keyword arguments
- `approximator`::[`ActorCritic`](@ref)
- `γ::Float32`, reward discount rate
-  `bootstrap::bool`, if false then Q function is approximated using monte carlo returns.
"""

Base.@kwdef mutable struct MACLearner{A<:ActorCritic} <: AbstractLearner
    approximator::A
    γ::Float32
    max_grad_norm::Union{Nothing,Float32} = nothing
    norm::Float32 = 0.0f0
    actor_loss::Float32 = 0.0f0
    critic_loss::Float32 = 0.0f0
    loss::Float32 = 0.0f0
    bootstrap::Bool = true
end

function (learner::MACLearner)(env::MultiThreadEnv)
    learner.approximator.actor(send_to_device(
        device(learner.approximator),
        get_state(env),
    )) |> send_to_host
end

function (learner::MACLearner)(env)
    s = get_state(env)
    s = Flux.unsqueeze(s, ndims(s) + 1)
    s = send_to_device(device(learner.approximator), s)
    learner.approximator.actor(s) |> vec |> send_to_host
end

function RLBase.update!(learner::MACLearner, t::AbstractTrajectory)
    isfull(t) || return

    states = t[:state]
    actions = t[:action]
    rewards = t[:reward]
    terminals = t[:terminal]

    AC = learner.approximator
    γ = learner.γ
    D = device(AC)

    states = send_to_device(D, states)
    states_flattened = flatten_batch(states) # (state_size..., n_thread * update_step)

    actions = flatten_batch(actions)
    actions = CartesianIndex.(actions, 1:length(actions))

    if learner.bootstrap
        next_state = select_last_frame(t[:next_state])
        next_state = send_to_device(D, next_state)
        next_state_values = AC.critic(next_state)

        gains = discount_rewards(
            rewards,
            γ;
            dims = 2,
            init = send_to_host(next_state_values),
            terminal = terminals,
        )
        gains = send_to_device(D, gains)
    else
        next_state_flattened = flatten_batch(t[:next_state])
        next_state_flattened = send_to_device(D, next_state_flattened)
        rewards_flattened = flatten_batch(rewards)
        rewards_flattened = send_to_device(D, rewards_flattened)
    end

    action_values = AC.critic(states_flattened)

    ps1 = Flux.params(AC.actor)
    gs1 = gradient(ps1) do
        logits = AC.actor(states_flattened)
        probs = softmax(logits)
        actor_loss = -mean(sum((probs .* Zygote.dropgrad(action_values)), dims = 1))
        loss = actor_loss
        ignore() do
            learner.actor_loss = actor_loss
        end
        loss
    end
    if !isnothing(learner.max_grad_norm)
        learner.norm = clip_by_global_norm!(gs1, ps1, learner.max_grad_norm)
    end
    update!(AC.actor, gs1)

    ps2 = Flux.params(AC.critic)
    gs2 = gradient(ps2) do
        if learner.bootstrap
            critic_loss = mean((vec(gains) .- vec(action_values[actions])) .^ 2)
        else
            next_state_values = AC.critic(next_state_flattened)
            target_action_values =
                vec(rewards_flattened) .+
                γ * vec(Zygote.dropgrad(sum(
                    next_state_values .* softmax(AC.actor(next_state_flattened)),
                    dims = 1,
                )))
            critic_loss =
                mean((vec(target_action_values) .- vec(action_values[actions])) .^ 2)
        end

        loss = critic_loss
        ignore() do
            learner.critic_loss = critic_loss
        end
        loss
    end
    if !isnothing(learner.max_grad_norm)
        learner.norm = clip_by_global_norm!(gs2, ps2, learner.max_grad_norm)
    end
    update!(AC.critic, gs2)
end

function (agent::Agent{<:QBasedPolicy{<:MACLearner},<:CircularCompactSARTSATrajectory})(
    ::Training{PreActStage},
    env,
)
    action = agent.policy(env)
    state = get_state(env)
    push!(agent.trajectory; state = state, action = action)
    update!(agent.policy, agent.trajectory)

    # the main difference is we'd like to flush the buffer after each update!
    if isfull(agent.trajectory)
        empty!(agent.trajectory)
        push!(agent.trajectory; state = state, action = action)
    end

    action
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using Logging

if haskey(ENV, "JULIAFEM_LOGLEVEL")
    ENV["JULIAFEM_LOGLEVEL"] == "DEBUG" && Logging.configure(level=DEBUG)
end

"""
This is JuliaFEM -- Finite Element Package
"""
module JuliaFEM

importall Base

module Testing

    if VERSION >= v"0.5-"
        using Base.Test
    else
        using BaseTestNext
    end

    export @test, @testset, @test_throws

end

include("io.jl")
export Xdmf, h5file, xmffile, has_child, get_child, new_dataitem

include("fields.jl")
export Field, DCTI, DVTI, DCTV, DVTV, CCTI, CVTI, CCTV, CVTV, Increment
include("types.jl")  # data types: Point, IntegrationPoint, ...
export AbstractPoint, Point, IntegrationPoint, IP, Node

### ELEMENTS ###
include("elements.jl") # common element routines
export Node, AbstractElement, Element, update!, get_connectivity, get_basis,
       get_dbasis, inside, get_local_coordinates
include("elements_lagrange.jl") # Continuous Galerkin (Lagrange) elements
export get_reference_coordinates, get_interpolation_polynomial
export Poi1,
       Seg2, Seg3,
       Tri3, Tri6, Tri7,
       Quad4, Quad8, Quad9,
       Tet4, Tet10,
       Wedge6,
       Hex8, Hex20, Hex27

include("elements_nurbs.jl")
export NSeg, NSurf, NSolid, is_nurbs

#include("hierarchical.jl") # P-elements

include("integrate.jl")  # default integration points for elements
export get_integration_points

include("sparse.jl")
export add!, SparseMatrixCOO, get_nonzero_rows, get_nonzero_columns

include("problems.jl") # common problem routines
export Problem, AbstractProblem, FieldProblem, BoundaryProblem,
       get_unknown_field_dimension, get_gdofs, Assembly,
       get_parent_field_name, get_elements

include("problems_elasticity.jl")
export Elasticity

include("problems_dirichlet.jl")
export Dirichlet

include("problems_heat.jl")
export Heat

export assemble!, postprocess!

function assemble!(problem::Problem, element::Element, time=0.0)
    assemble!(problem.assembly, problem, element, time)
end

### ASSEMBLY + SOLVE ###
include("assembly.jl")
include("solvers_utils.jl")
include("solvers.jl")
export AbstractSolver, Solver, Nonlinear, NonlinearSolver, Linear, LinearSolver,
       get_unknown_field_name, get_formulation_type, get_problems,
       get_field_problems, get_boundary_problems,
       get_field_assembly, get_boundary_assembly,
       initialize!, create_projection, eliminate_interior_dofs,
       is_field_problem, is_boundary_problem
include("solvers_modal.jl")
export Modal

include("optics.jl")
export find_intersection, calc_reflection, calc_normal

### Mortar methods ###
include("problems_mortar.jl")
include("problems_mortar_2d.jl")
include("problems_mortar_3d.jl")
include("problems_mortar_2d_autodiff.jl")
export calculate_normals, calculate_normals!, project_from_slave_to_master,
       project_from_master_to_slave, Mortar, get_slave_elements,
       get_polygon_clip

### Mortar methods, contact mechanics extension ###
include("problems_contact.jl")
include("problems_contact_2d.jl")
include("problems_contact_3d.jl")
include("problems_contact_2d_autodiff.jl")
export Contact

module Preprocess
include("preprocess.jl")
export create_elements, Mesh, add_node!, add_nodes!, add_element!,
       add_elements!, add_element_to_element_set!, add_node_to_node_set!,
       find_nearest_nodes, reorder_element_connectivity!
include("preprocess_abaqus_reader.jl")
export parse_abaqus, parse_section, parse_element_section,
       abaqus_read_mesh, abaqus_read_model
include("preprocess_aster_reader.jl")
export aster_create_elements, parse_aster_med_file, is_aster_mail_keyword,
       parse_aster_header, aster_parse_nodes, aster_renumber_nodes!,
       aster_renumber_elements!, aster_combine_meshes, aster_read_mesh,
       filter_by_element_set, filter_by_element_id, MEDFile, aster_read_data,
       aster_read_mesh_names, aster_read_node_sets, aster_read_nodes, RMEDFile
end

function get_mesh(mesh_name::AbstractString, args...; kwargs...)
    return get_mesh(Val{Symbol(mesh_name)}, args...; kwargs...)
end

function get_model(model_name::AbstractString, args...; kwargs...)
    return get_model(Val{Symbol(model_name)}, args...; kwargs...)
end

export get_mesh, get_model

module Postprocess

include("postprocess_utils.jl")
export calc_nodal_values!, get_nodal_vector, get_nodal_dict, copy_field!,
       calculate_area, calculate_center_of_mass,
       calculate_second_moment_of_mass, extract

include("postprocess_xdmf.jl")
export XDMF, xdmf_new_result!, xdmf_save_field!, xdmf_save!, DataFrame

end
export Postprocessor

module Abaqus
include("abaqus.jl")
export abaqus_read_model, abaqus_run_model, abaqus_open_results
end

end #


"""
Convenient structure for postprocessing. It is less performant than having `SArray`s but is easier to manipulate.
"""
struct SpacecraftTrajectory
    position::Array{Float64, 2}
    attitude::Array{Float64, 2}
    velocity::Array{Float64, 2}
    angular_velocity::Array{Float64, 2}
    thrust::Array{Float64, 2}
    angular_controls::Array{Float64, 2}
    time
    SpacecraftTrajectory(states, controls, time) = new(
        states[1:3, :],
        states[4:7, :],
        states[7:9, :],
        states[10:12, :],
        controls[1:1, :],
        controls[2:4, :],
        time)
end

function Base.write(model::Spacecraft, traj::SpacecraftTrajectory, filename)
    dict = Dict(
        "$name" => getproperty(traj, name) for name ∈ fieldnames(SpacecraftTrajectory)
    )
    open(filename, "w") do io
        write(io, json(dict))
    end
end

"""
Print some information about the solved problem.
"""
function postprocess(solver, x, u)
    println("Cost: ", cost(solver))
    println("Constraint violation: ", max_violation(solver))
    println("Iterations: ", iterations(solver))
    println("Final state: ", x[:, end])
    println("Maximum control: ", maximum(u))
end