text:
using Distributions
using Lora

# using Gadfly

dpkeys = [:am, :assq, :bm, :bssq, :yssq, :a, :b, :y]
ndpkeys = length(dpkeys)
nudpkeys = ndpkeys-1

nrats, ndays = 30, 5
nvars = 9

X, = dataset("rats", "age")
X = vec(X)-mean(X)
Y, = dataset("rats", "weight")
v0 = Dict{Symbol, Any}(
  :am=>150.,
  :assq=>1.,
  :bm=>10.,
  :bssq=>1.,
  :yssq=>1.,
  :x=>X,
  :a=>fill(250., nrats),
  :b=>fill(6., nrats),
  :y=>vec(hcat([Y[i, :] for i in 1:30]...))
)

logtargets = Dict{Symbol, Function}(
  :am => p -> logpdf(Normal(0, 100), p),
  :assq => p -> logpdf(InverseGamma(0.01, 0.01), p),
  :bm => p -> logpdf(Normal(0, 100), p),
  :bssq => p -> logpdf(InverseGamma(0.01, 0.01), p),
  :yssq => p -> logpdf(InverseGamma(0.01, 0.01), p),
  :a => (p, am, assq) -> logpdf(MvNormal(am*ones(nrats), sqrt(assq)), p),
  :b => (p, bm, bssq) -> logpdf(MvNormal(bm*ones(nrats), sqrt(bssq)), p),
  :y => (p, yssq, a, b, x) -> logpdf(
    MvNormal(repeat(a, inner=[ndays])+repeat(b, inner=[ndays]).*repeat(x, outer=[nrats]), sqrt(yssq)),
      p
    )
)

function logtarget(am, assq, bm, bssq, yssq, a, b, y, x)
  logtargets[:am](am)+
  logtargets[:assq](assq)+
  logtargets[:bm](bm)+
  logtargets[:bssq](bssq)+
  logtargets[:yssq](yssq)+
  logtargets[:a](a, am, assq)+
  logtargets[:b](b, bm, bssq)+
  logtargets[:y](y, yssq, a, b, x)
end

nmcmc = 1000
nburnin = 500
npostburnin = nmcmc-nburnin

ns = Dict{Symbol, Int}(
  :am => 1,
  :assq => 1,
  :bm => 1,
  :bssq => 1,
  :yssq => 1,
  :a => 30,
  :b => 30
)

chains = Dict{Symbol, Array{Float64}}(
  :am => Array(Float64, npostburnin),
  :assq => Array(Float64, npostburnin),
  :bm => Array(Float64, npostburnin),
  :bssq => Array(Float64, npostburnin),
  :yssq => Array(Float64, npostburnin),
  :a => Array(Float64, npostburnin, 2),
  :b => Array(Float64, npostburnin, 2)
)

oldsample = deepcopy(v0)
delete!(oldsample, :x)

newsample = deepcopy(oldsample)

oldlogtarget = logtarget(Any[v0[k] for k in dpkeys]..., v0[:x])
newlogtarget = NaN

s = Dict{Symbol, Array{Float64}}(
  :am => 1.,
  :assq => 1.,
  :bm => 1.,
  :bssq => 1.,
  :yssq => 1.,
  :a => 1.,
  :b => 1.
)

type RobertsRosenthalMCTune
  σ::Real # Stepsize of MCMC iteration (for ex leapfrog in HMC or drift stepsize in MALA)
  myflag::Bool
  accepted::Int # Number of accepted MCMC samples during current tuning period
  proposed::Int # Number of proposed MCMC samples during current tuning period
  totproposed::Int # Total number of proposed MCMC samples during burnin
  rate::Real # Observed acceptance rate over current tuning period

  function RobertsRosenthalMCTune(σ::Real, myflag::Bool, accepted::Int, proposed::Int, totproposed::Int, rate::Real)
    @assert σ > 0 "Stepsize of MCMC iteration should be positive"
    @assert accepted >= 0 "Number of accepted MCMC samples should be non-negative"
    @assert proposed >= 0 "Number of proposed MCMC samples should be non-negative"
    @assert totproposed >= 0 "Total number of proposed MCMC samples should be non-negative"
    if !isnan(rate)
      @assert 0 < rate < 1 "Observed acceptance rate should be between 0 and 1"
    end
    new(σ, myflag, accepted, proposed, totproposed, rate)
  end
end

RobertsRosenthalMCTune(σ::Real=1., myflag::Bool=true, accepted::Int=0, proposed::Int=0, totproposed::Int=0) =
  RobertsRosenthalMCTune(σ, myflag, accepted, proposed, totproposed, NaN)

immutable RobertsRosenthalMCTuner
  targetrate::Real # Target acceptance rate
  period::Int # Tuning period over which acceptance rate is computed
  verbose::Bool # If verbose=false then the tuner is silent, else it is verbose

  function RobertsRosenthalMCTuner(targetrate::Real, period::Int, verbose::Bool)
    @assert 0 < targetrate < 1 "Target acceptance rate should be between 0 and 1"
    @assert period > 0 "Tuning period should be positive"
    new(targetrate, period, verbose)
  end
end

RobertsRosenthalMCTuner(
  targetrate::Real;
  period::Int=100,
  verbose::Bool=false
) =
  RobertsRosenthalMCTuner(targetrate, period, verbose)

function tune!(tune::RobertsRosenthalMCTune, tuner::RobertsRosenthalMCTuner)
  delta = min(0.01, (tune.totproposed/tuner.period)^-0.5)
  epsilon = tune.rate < tuner.targetrate ? -delta : delta
  tune.σ *= exp(epsilon)
end

function reset_burnin!(tune::RobertsRosenthalMCTune)
  tune.totproposed += tune.proposed
  (tune.accepted, tune.proposed, tune.rate) = (0, 0, NaN)
end

rate!(tune::RobertsRosenthalMCTune) = (tune.rate = tune.accepted/tune.proposed)

tune = Dict(zip(dpkeys, fill(RobertsRosenthalMCTune(), ndpkeys)))
delete!(tune, :y)

m = 1

for i in 1:nmcmc
  println(i)

  for j in 1:nudpkeys
    k = dpkeys[j]
    newsample[k] = oldsample[k]+tune[k].σ*randn(ns[k])

    ratio = logtarget(Any[v0[k] for k in dpkeys]..., v0[:x])

    if accept
      oldsample[k] = copy(newsample[k])
    end
  end

  if i < nburnin
  else
    chains[k][m] = oldsample[k]
    m += 1
  end
end


text:
path,_ = splitdir(@__FILE__)
using PyPlot
import JLD

#Packages I haven't developed
import Interpolations
# using PyPlot
import JLD
import CSV
import Dierckx

# Packages I played a very minor role
import PreComp
import Composites
import VTKtools

# Packages I played a minor role
import OptParams
import Atmosphere
import CCBlade #new CCBlade package w/o ND splined af data
import BeamFEA
import Snopt
import Gradients

# Packages I played a major role
import BrentMin
import MotorPower
import Akima
# import VLM
include("$path/../oldvlm.jl")
import BPM
import AirfoilPrep

printiter = 4600
#TODO verify this works with parallel

include("$path/../objective.jl")
include("$path/../propstructmisc.jl")
include("$path/../propwingmisc.jl")
include("$path/../setup.jl")
include("$path/../save_vtk.jl")

detailed_output_accel = []
detailed_output_climb = []
detailed_output_cruise = []
design_accel = []
design_climb = []
design_cruise = []
detailed_system = []

rc("figure", figsize=(60*1.8, 60))
rc("font", size=10.0)
rc("lines", linewidth=1.5)
rc("lines", markersize=3.0)
rc("legend", frameon=false)
# rc("figure", dpi=300)
rc("axes.spines", right=false, top=false)
rc("axes", linewidth=15)
# rc("figure.subplot", left=0.27, bottom=0.2, top=0.8, right=.95)
rc("figure.subplot", left=0.01, bottom=0.01, top=0.99, right=.99)
rc("axes", color_cycle=["348ABD", "A60628", "009E73", "7A68A6", "D55E00", "CC79A7"])
plot_cycle=["#348ABD", "#A60628", "#009E73", "#7A68A6", "#D55E00", "#CC79A7"]

ioff()

index = 70
failnum = 0
imagenum = 0
iterations = zeros(1)
objective1 = ones(1)*100
objectiveColor = ones(1)
indexcorrection = 0
background_transp = ones(1)
objcol = 0.0

while index<51521

    # try

close("all")

println("imagenum: $imagenum index: $index")
# success = true
# while success

if index == 2170
    index = 21070
    indexcorrection = 18900
end

filename = "$(path)/Journal_GIF/p8_cr100_db76correct_span$(index).jld"
summary = JLD.load(filename)
image = imread("$(path)/Journal_PNG/blade.$(lpad(imagenum,4,0)).png")
# success = false
# failnum = 0
# catch
#     failnum += 1
#     index += 1
#
# end

detailed_output_accel = summary["detailed_output_accel"]
detailed_output_climb = summary["detailed_output_climb"]
detailed_output_cruise = summary["detailed_output_cruise"]
design_accel = summary["design_accel"]
design_climb = summary["design_climb"]
design_cruise = summary["design_cruise"]
detailed_system = summary["detailed_system"]
obj = summary["objective"]

G = linspace(0,detailed_system.gama,100)-pi/2

xtr = detailed_system.dist_accel+detailed_system.radius*cos.(G)
ytr = detailed_system.radius*sin.(G)+detailed_system.radius

pathx = [0;detailed_system.dist_accel;xtr;detailed_system.dist_accel+detailed_system.dist_climb] # dist_climb includes detailed_system.xtr
plotx = [0,detailed_system.dist_accel,detailed_system.dist_accel+detailed_system.xtr,detailed_system.dist_accel+detailed_system.dist_climb]
pathy = [0;0;ytr; 15.24]

iterations = append!(iterations, iterations[end]+index-indexcorrection)

if pathx[end]<1.0

    objective1 = append!(objective1, objective1[end])

else

    objective1 = append!(objective1, pathx[end])

end

w_height = 0

imagenum += 1
index += 70

figname = "Journal_GIF$(lpad(imagenum,4,0))"
figure(figname)

# SETUP
axes()
axis("off")

#TOP TEXTS

# text(0.04, 0.95, "Props (#): 2 4    16 32     Noise (dBa): 60 65          Cruise (MPH): 50        150",
# verticalalignment="bottom", horizontalalignment="left",
# color=(0.0, 0.0, 0.0, 1.0), fontsize=180)
#
# text(0.212, 0.95, "8",
# verticalalignment="bottom", horizontalalignment="left",
# color=(0.0, 0.0, 0.0, 1.0), fontsize=180,
# bbox=Dict("facecolor"=>"none", "edgecolor"=>"black", "boxstyle"=>"round,pad=0.1","linewidth"=>"15"))
#
# text(0.57, 0.95, "70",
# verticalalignment="bottom", horizontalalignment="left",
# color=(0.0, 0.0, 0.0, 1.0), fontsize=180,
# bbox=Dict("facecolor"=>"none", "edgecolor"=>"black", "boxstyle"=>"round,pad=0.1","linewidth"=>"15"))
#
# text(0.843, 0.95, "100",
# verticalalignment="bottom", horizontalalignment="left",
# color=(0.0, 0.0, 0.0, 1.0), fontsize=180,
# bbox=Dict("facecolor"=>"none", "edgecolor"=>"black", "boxstyle"=>"round,pad=0.1","linewidth"=>"15"))
#

con_names = ["Weight",
"Noise",
"Range",
"Lift",
"Thrust",
"Tip Separation",
"Tip Speed",
"Stall",
"Cruise Speed",
"Composite Failure",
"Buckling",
"Battery Capacity"]

con_dic_names = [:totalmass,
:noise_accel,
:range,
:Lift_cruise,
:Drag_cruise,
:radiussize,
:machtip_accel,
:cl_max_accel,
:cruisespeed,
:p_stress_climb,
:p_buckling_climb,
:capacity_climb]

# for names in keys(detailed_system.optconvals)
#     println(names)
# end

constraint_y = collect(range(.82,-.07,length(con_names)))

objcol = 0

#CONSTRAINTS
for i = 1:length(con_names)
    try #catch if one of the constraints listed above is not in the list here
        background_transp = maximum(detailed_system.optconvals[con_dic_names[i]])*100
    catch
        background_transp = 0.0
    end

    # if background_transp< 0.0
    #     background_transp = 0.0
    # elseif background_transp > 1.0
    #     background_transp = 1.0
    # end

    if objcol < background_transp
        objcol = background_transp
    end

    coolwarm_constr_color = PyPlot.cm[:coolwarm](background_transp)

    # println(background_transp)
    #Background
    text(0.915, constraint_y[i], con_names[i],
    verticalalignment="bottom", horizontalalignment="center",
    color=Tuple(coolwarm_constr_color), fontsize=130,fontweight = 1000)
    # Actual text
    # text(0.915, constraint_y[i], con_names[i],
    # verticalalignment="bottom", horizontalalignment="center",
    # color=(0.0, 0.0, 0.0, 1.0), fontsize=140,fontweight = 540.0,)
end

objectiveColor = append!(objectiveColor, objcol)

#LABEL TEXTS
text(0.1, 0.91, "Props On Wing",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=160)

text(0.55, 0.90, "Takeoff Profile",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=160)

text(0.44, 0.468, "Optimization Progression",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=160)

text(0.915, 0.91, "Constraints",
verticalalignment="bottom", horizontalalignment="center",
color=(0.0, 0.0, 0.0, 1.0), fontsize=170)

# VTK Image
# Scaled by 1/2 when outputting in paraview
axes([-0.050, 0.04, 0.35, 0.68])
imshow(image)

text(210, 0, "Blade Design",
verticalalignment="bottom", horizontalalignment="center",
color=(0.0, 0.0, 0.0, 1.0), fontsize=160)

xoff = 40

text(285-xoff, 50, "Composite Failure",
verticalalignment="bottom", horizontalalignment="center",
color=(0.0, 0.0, 0.0, 1.0), fontsize=130)#,rotation = "vertical")

text(275-xoff, 321, "Buckling Failure",
verticalalignment="bottom", horizontalalignment="center",
color=(0.0, 0.0, 0.0, 1.0), fontsize=130)#,rotation = "vertical")

# TOP Bar Annotation
center = 178
offset = 95

text(205-xoff, center-offset, "Violated",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=95)

text(205-xoff, center, "Constrained",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=95)

text(205-xoff, center+offset, "Satisfied",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=95)

# BOTTOM Bar Annotation
center = 445
offset = 95

text(205-xoff, center-offset, "Violated",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=95)

text(205-xoff, center, "Constrained",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=95)

text(205-xoff, center+offset, "Satisfied",
verticalalignment="bottom", horizontalalignment="left",
color=(0.0, 0.0, 0.0, 1.0), fontsize=95)

axis("off")

# Wing and Props
axes([0.0, 0.56, 0.4, 0.6])
x_tips = [0.0,1/8,2/8,3/8,4/8,5/8,6/8,7/8,1.0]
x = ((x_tips[2:end]-x_tips[1:end-1])/2+x_tips[1:end-1])*11.4
y = ones(length(x))*w_height
areas = ones(length(x))*pi*design_accel.p_Rtip^2*85000
# areas = ones(length(x))*pi*.7125^2*85000 #tip to tip for tuning scaling factor
scatter(x,y,s=areas,alpha=0.5)
plot([0.1,11.3],ones(2)*w_height,linewidth=75.0)
axis("off")

# Takeoff
axes([0.47, 0.655, .34, 0.238])
plot(pathx[1:2],pathy[1:2],"-",linewidth=25.0,color = "#1a455e")
plot(pathx[2:102],pathy[2:102],"-",linewidth=25.0,color = "#28698f")
plot(pathx[102:103],pathy[102:103],"-",linewidth=25.0,color = "#348ABD")
tick_params(which="major",length=10*7,width=2*7,labelsize=16*7)
xlabel("Distance (m)",fontsize=140)
xticks(plotx, rotation = 35)
ylabel("Height (m)",fontsize=140)
yticks([0,15.24],["0","15.24\n (50 ft)"],rotation = 90)
ylim([0.0,20.0])
xlim([0.0,100.0])
axis("equal")

coolwarm_obj_color = PyPlot.cm[:coolwarm](objectiveColor)

# Objective
axes([0.38, 0.119, 0.40, 0.35])
for i = 1:length(iterations)-1
    semilogy(iterations[i:i+1]/1000/35, objective1[i:i+1],"-",c = Tuple(coolwarm_obj_color[i,:]),linewidth=25.0)
end
xlabel("Iterations (1k)",fontsize=140)
ylabel("Takeoff Distance (m)",fontsize=140)
yticks(rotation = 30)
# ylim([0,100])
# xlim([0,100])
tick_params(which="major",length=10*7,width=2*7,labelsize=16*7)
tick_params(which="minor",length=10*7,width=2*7,labelsize=16*7)

# legend(title = "Cruise Constraint (m/s)",loc="upper center", ncol = length(cr_factor), bbox_to_anchor=(0.5, 1.30))#loc="center left", bbox_to_anchor=(1, 0.5))
# legend(title = "Cruise Constraint (m/s)",loc="upper center", ncol = length(cr_factor), bbox_to_anchor=(0.5, -.20))
savefig("$path/Journal_GIF/Output/$figname.png")#,transparent = false)

    # catch
    #     index += 70
    #     failnum+=1
    #     println("$failnum")
    # end
end

# run(`cd Journal_GIF/Output ; pwd ; convert -delay 2 -loop 0 *.pdf Moore2018_3.gif`)

ext: jl


text:
@testset "FeatureTransforms" begin

    M1 = [0.0 1.0; 1.0 2.0; -0.5 0.0]
    M2 = [-2.0 4.0; 3.0 2.0; -1.0 -1.0]
    M3 = [0.0 1.0; -1.0 0.5; -0.5 0.0]
    M4 = [0.5 -1.0; -5.0 -2.0; 0.0 1.0]

    ds = KeyedDataset(
        flatten([
            :train => [
                :load => KeyedArray(M1; time=1:3, loc=[:x, :y]),
                :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
            ],
            :predict => [
                :load => KeyedArray(M3; time=1:3, loc=[:x, :y]),
                :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
            ]
        ])...
    )

    @testset "transform" begin
        @test is_transformable(ds)
    end

    @testset "OneToOne" begin
        T = FakeOneToOneTransform()

        @testset "default applies to all components" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 2); time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 2); time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using pattern" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(M1; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 2); time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(M3; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 2); time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, (:_, :price, :_))

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using symbol" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using dims" begin
            # replaces the first :loc column with ones(...)
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 1); time=1:3, loc=[:x]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 1); time=1:3, loc=[:x]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc; dims=2, inds=[1])

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end
    end

    @testset "OneToMany" begin
        T = FakeOneToManyTransform()

        @testset "default applies to all components" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using pattern" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(M1; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(M3; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, (:_, :price, :_))

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using symbol" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using dims" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :x]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :x]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc; dims=2, inds=[1])

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end
    end

    @testset "ManyToOne" begin
        T = FakeManyToOneTransform()

        @testset "default applies to all components" begin
            expected = KeyedDataset(
                # ideally we would drop the :time constraint when it gets reduced
                OrderedSet(Pattern[(:__, :time), (:__, :loc), (:__, :id)]),
                LittleDict(flatten([
                    :train => [
                        :load => KeyedArray(ones(2); loc=[:x, :y]),
                        :price => KeyedArray(ones(2); id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(2); loc=[:x, :y]),
                        :price => KeyedArray(ones(2); id=[:a, :b]),
                    ]
                ])...)
            )

            r = FeatureTransforms.apply(ds, T; dims=:time)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using pattern" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(M1; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(2); id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(M3; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(2); id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, (:_, :price, :_); dims=:time)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using symbol" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(2); loc=[:x, :y]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(2); loc=[:x, :y]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc; dims=:time)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end
    end

    # Note: There are no ManyToMany transforms implemented just yet
    @testset "ManyToMany" begin
        T = FakeManyToManyTransform()

        @testset "default applies to all components" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using pattern" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(M1; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(M3; time=1:3, loc=[:x, :y]),
                        :price => KeyedArray(ones(3, 4); time=1:3, id=[:a, :b, :a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, (:_, :price, :_))

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using symbol" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 4); time=1:3, loc=[:x, :y, :x, :y]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc)

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end

        @testset "using dims" begin
            expected = KeyedDataset(
                flatten([
                    :train => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :x]),
                        :price => KeyedArray(M2; time=1:3, id=[:a, :b]),
                    ],
                    :predict => [
                        :load => KeyedArray(ones(3, 2); time=1:3, loc=[:x, :x]),
                        :price => KeyedArray(M4; time=1:3, id=[:a, :b]),
                    ]
                ])...
            )

            r = FeatureTransforms.apply(ds, T, :loc; dims=2, inds=[1])

            @test r isa KeyedDataset
            @test isequal(r, expected)
            @test !isequal(ds, expected)
        end
    end

end


text:
export addnode!
################################################################
# Add node
################################################################
# Document
################################################################
# Any
"""
    addnode!(node, name, value, argAmlType)

To add nodes (single or a vector of nodes) as a child of given a node/document.
"""
function addnode!(aml::Document, name::String, value::T, argAmlType::Type{<:AbsDocOrNode}) where {T}

    if hasroot(aml)
        amlNode = root(aml)
        elm = addnode!(amlNode, name, value, argAmlType)
        return elm
    elseif hasfield(T, :aml)
        elm = setroot!(aml, value.aml)
        return elm
    else
        error("You cannot insert $(T) in the document directly. Define a @aml defined field for xml or html document struct")
    end

end

# Nothing
function addnode!(aml::Document, name::String, value::Nothing, argAmlType::Type{<:AbsDocOrNode})
# do nothing if value is nothing
end
################################################################
# Vector
"""
    addnode!(node, name, value, argAmlType)

Add a vector to a node/document
```
"""
function addnode!(aml::Document, name::String, value::Vector, argAmlType::Type{<:AbsDocOrNode})

    if hasroot(aml)
        amlNode = root(aml)
        elm = addnode!(amlNode, name, value, argAmlType)
        return elm
    else
        error("You cannot insert a vector in the document directly. Define a @aml defined field for xml or html document struct")
    end

end

################################################################
# Nodes
################################################################
# String
@transform function addnode!(aml::Node, name::String,value::AbstractString, argAmlType::Type{allsubtypes(AbsNormal)})
    if !isnothing(value) # do nothing if value is nothing
        elm = addelement!(aml, name, value)
        return elm
    end
end

function addnode!(aml::Node, name::String,value::AbstractString, argAmlType::Type{AbsAttribute})
    if !isnothing(value) # do nothing if value is nothing
        elm = link!(aml, AttributeNode(name, value))
        return elm
    end
end

# helper
function textlinkinfo(aml, indexstr)
    index = parse_textindex(indexstr)
    aml_elms = elements(aml)
    aml_elms_len = length(aml_elms)
    if length(aml_elms) == 0
        desired_node = aml
        linking_function! = link!
    elseif index < aml_elms_len
        desired_node = aml_elms[index]
        linking_function! = linkprev!
    else
        desired_node = aml_elms[end]
        linking_function! = linknext!
    end
    return desired_node, linking_function!
end

function addnode!(aml::Node, indexstr::String, value::AbstractString, argAmlType::Type{AbsText})
    desired_node, linking_function! = textlinkinfo(aml, indexstr)
    if !isnothing(value) # do nothing if value is nothing
        elm = linking_function!(desired_node, TextNode(value))
        return elm
    end
end

# Number (and also Bool <:Number)
@transform function addnode!(aml::Node, name::String, value::Number, argAmlType::Type{allsubtypes(AbsNode)})
    addnode!(aml, name, string(value), argAmlType)
end

# Other
function addnode!(aml::Node, name::String, value::T, argAmlType::Type{<:AbsNormal}) where {T}
    if hasfield(T, :aml)
        elm = link!(aml,value.aml)
    elseif Tables.istable(value)
        elm = link!(aml, amlTable(value))
    elseif hasmethod(aml, Tuple{T})
        elm = link!(aml,aml(value))
    else
        elm = addelement!(aml, name, string(value))
    end
    return elm
end

function addnode!(aml::Node, name::String, value::T, argAmlType::Type{AbsAttribute}) where {T}
    if hasfield(T, :aml)
        elm = link!(aml, AttributeNode(name, value.aml))
    elseif Tables.istable(value)
        elm = link!(aml, AttributeNode(name, amlTable(value)))
    elseif hasmethod(aml, Tuple{T})
        elm = link!(aml, AttributeNode(name, aml(value)))
    else
        elm = link!(aml, AttributeNode(name, string(value)))
    end
    return elm
end

function addnode!(aml::Node, indexstr::String, value::T, argAmlType::Type{AbsText}) where {T}
    desired_node, linking_function! = textlinkinfo(aml, indexstr)
    if hasfield(T, :aml)
        elm = linking_function!(desired_node, TextNode(value.aml))
    elseif Tables.istable(value)
        elm = linking_function!(desired_node, TextNode(amlTable(value)))
    elseif hasmethod(aml, Tuple{T})
        elm = linking_function!(desired_node, TextNode(aml(value)))
    else
        elm = linking_function!(desired_node, TextNode(string(value)))
    end
    return elm
end

# Nothing
@transform function addnode!(aml::Node, name::String, value::Nothing, argAmlType::Type{allsubtypes(AbsDocOrNode)})
    # do nothing
end
################################################################
# Vector

allsubtypes_butAbsText(t) = setdiff(allsubtypes(AbsDocOrNode), [AbsText])

@transform function addnode!(aml::Node, name::String, values::Vector, argAmlType::Type{allsubtypes_butAbsText(AbsDocOrNode)})
    elms = Vector{Union{Node, Nothing}}(undef, length(values))
    for (ielm, value) in enumerate(values)
        elms[ielm] = addnode!(aml, name, value, argAmlType)
    end
    return elms
end

function addnode!(aml::Node, indicesstr::String, values::Vector, argAmlType::Type{AbsText})
    indices = parse_textindices(indicesstr)
    if indices isa Colon
        indices = 1:length(elements(aml))
    end
    elms = Vector{Union{Node, Nothing}}(undef, length(indices))
    ielm = 1
    for (value, index) in zip(values, indices)
        elms[ielm] = addnode!(aml, string(index), value, argAmlType)
        ielm += 1
    end
    return elms
end

################################################################
# Dict

@transform function addnode!(aml::Node, name::String, values::AbstractDict, argAmlType::Type{allsubtypes(AbsDocOrNode)})
    # name is discarded now: actual names are stored in the Dict itself
    # elements are added directly
    # for AbsText, v_name is considered as the text index
    elms = Vector{Union{Node, Nothing}}(undef, length(values))
    ielm = 1
    for (v_name, v_value) in values
        elms[ielm] = addnode!(aml, v_name, v_value, argAmlType)
        ielm += 1
    end
end


text:
# 1.4.2 Deprecations

function _add_component_kwarg_deprecation(kwargs)
    if haskey(kwargs, :deserialization_in_progress)
        Base.depwarn(
            "Keyword deserialization_in_progress is deprecated, use allow_existing_time_series instead.",
            :add_component!,
        )
        kw = Dict(k => v for (k, v) in kwargs if k != :deserialization_in_progress)
        kw[:allow_existing_time_series] = kwargs[:deserialization_in_progress]
        return kw
    else
        return kwargs
    end
end


text:
#****************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#****************************************************************************

function periodicimages(pbc::Array{Int64,1}, a::Array{Float64,1}, b::Array{Float64,1}, c=nothing)
    # a = (a1, a2, a3), b = (b1,b2,b3), c=(c1,c2,c3) are three vectors defining the simulation box
    # r is the distance from the original parallelepiped to the bounding parallelepiped

    if c === nothing
        a = a[:];
        b = b[:];
        nd = length(a);
        px = zeros(nd,3^nd);

        px[:,1] = [0.0; 0.0];
        px[:,2] = a;
        px[:,3] = -a;
        px[:,4] = b;
        px[:,5] = -b;
        px[:,6] = a+b;
        px[:,7] = a-b;
        px[:,8] = -a+b;
        px[:,9] = -a-b;

        if (pbc[1]==0) && (pbc[2]==0)
            indx = 1;
        elseif (pbc[1]==1) && (pbc[2]==0)
            indx = [1 2 3];
        elseif (pbc[1]==0) && (pbc[2]==1)
            indx = [1 4 5];
        elseif (pbc[1]==1) && (pbc[2]==1)
            indx = [1 2 3 4 5 6 7 8 9];
        end
        if length(indx)==1
            ximages = px[:,indx];
        else
            ximages = px[:,indx[:]];
        end
    else
        a = a[:];
        b = b[:];
        c = c[:];
        nd = length(a);
        px = zeros(nd,3^nd);

        px[:,1] = [0.0; 0.0; 0.0];
        px[:,2] = a;
        px[:,3] = -a;
        px[:,4] = b;
        px[:,5] = -b;
        px[:,6] = a+b;
        px[:,7] = a-b;
        px[:,8] = -a+b;
        px[:,9] = -a-b;

        px[:,10] = c;
        px[:,11] = c+a;
        px[:,12] = c-a;
        px[:,13] = c+b;
        px[:,14] = c-b;
        px[:,15] = c+a+b;
        px[:,16] = c+a-b;
        px[:,17] = c-a+b;
        px[:,18] = c-a-b;

        px[:,19] = -c;
        px[:,20] = a-c;
        px[:,21] = -a-c;
        px[:,22] = b-c;
        px[:,23] = -b-c;
        px[:,24] = a+b-c;
        px[:,25] = a-b-c;
        px[:,26] = -a+b-c;
        px[:,27] = -a-b-c;

        if (pbc[1]==0) && (pbc[2]==0) && (pbc[3]==0)
            indx = 1;
        elseif (pbc[1]==1) && (pbc[2]==0) && (pbc[3]==0)
            indx = [1 2 3];
        elseif (pbc[1]==0) && (pbc[2]==1) && (pbc[3]==0)
            indx = [1 4 5];
        elseif (pbc[1]==0) && (pbc[2]==0) && (pbc[3]==1)
            indx = [1 10 19];
        elseif (pbc[1]==1) && (pbc[2]==1) && (pbc[3]==0)
            indx = reshape(collect(1:9), 1, 9);
        elseif (pbc[1]==1) && (pbc[2]==0) && (pbc[3]==1)
            indx = ones(1,9) + [0 1 2 9 10 11 18 19 20];
        elseif (pbc[1]==0) && (pbc[2]==1) && (pbc[3]==1)
            indx = ones(1,9) + [0 3 4 9 12 13 18 21 22];
        elseif (pbc[1]==1) && (pbc[2]==1) && (pbc[3]==1)
            indx = reshape(collect(1:27), 1, 27);
        end
        if length(indx)==1
            ximages = px[:,indx];
        else
            ximages = px[:,indx[:]];
        end

        return ximages;
    end

    end

max_stars_repo_name: cesmix-mit/MDP

max_issues_repo_name: cesmix-mit/MDP

max_forks_repo_name: cesmix-mit/MDP


text:
module CudaNNUtils

using CUDArt
using Quaternions
using QuaternionArgs
import Base: transpose, ctranspose

const ptxdict = Dict()
const mdlist = Array(CuModule, 0)
const block_default = 128

function init(devlist)
    global ptxdict
    global mdlist
    isempty(mdlist) || return
    for dev in devlist
        device(dev)
        mod_dir = joinpath(dirname(@__FILE__), "kernels")
        mod_path = joinpath(mod_dir, "kernels.ptx")
        md = CuModule(mod_path, false)
        ptxdict[(dev, "equiv", Complex{Float32}, Quaternion{Float32})] = CuFunction(md, "q2c_float")
        ptxdict[(dev, "equiv", Complex{Float64}, Quaternion{Float64})] = CuFunction(md, "q2c_double")
        ptxdict[(dev, "equiv", Quaternion{Float32}, Complex{Float32})] = CuFunction(md, "c2q_float")
        ptxdict[(dev, "equiv", Quaternion{Float64}, Complex{Float64})] = CuFunction(md, "c2q_double")
        ptxdict[(dev, "actfunc", Float32)] = CuFunction(md, "actfunc_real_float")
        ptxdict[(dev, "actfunc", Float64)] = CuFunction(md, "actfunc_real_double")
        ptxdict[(dev, "actfunc", Complex{Float32})] = CuFunction(md, "actfunc_comp_float")
        ptxdict[(dev, "actfunc", Complex{Float64})] = CuFunction(md, "actfunc_comp_double")
        ptxdict[(dev, "actfunc", Quaternion{Float32})] = CuFunction(md, "actfunc_quat_float")
        ptxdict[(dev, "actfunc", Quaternion{Float64})] = CuFunction(md, "actfunc_quat_double")
        ptxdict[(dev, "actfunc_multi", Complex{Float32})] = CuFunction(md, "actfunc_comp_multi_float")
        ptxdict[(dev, "actfunc_multi", Complex{Float64})] = CuFunction(md, "actfunc_comp_multi_double")
        ptxdict[(dev, "actfunc_multi", Quaternion{Float32})] = CuFunction(md, "actfunc_quat_multi_float")
        ptxdict[(dev, "actfunc_multi", Quaternion{Float64})] = CuFunction(md, "actfunc_quat_multi_double")
        ptxdict[(dev, "fill", Float32)] = CuFunction(md, "fill_real_float")
        ptxdict[(dev, "fill", Float64)] = CuFunction(md, "fill_real_double")
        ptxdict[(dev, "fill", Complex{Float32})] = CuFunction(md, "fill_comp_float")
        ptxdict[(dev, "fill", Complex{Float64})] = CuFunction(md, "fill_comp_double")
        ptxdict[(dev, "fill", Quaternion{Float32})] = CuFunction(md, "fill_quat_float")
        ptxdict[(dev, "fill", Quaternion{Float64})] = CuFunction(md, "fill_quat_double")
        ptxdict[(dev, "compare", Int32)] = CuFunction(md, "compare_real_int")
        ptxdict[(dev, "compare", Float32)] = CuFunction(md, "compare_real_float")
        ptxdict[(dev, "compare", Float64)] = CuFunction(md, "compare_real_double")
        ptxdict[(dev, "compare", Complex{Float32})] = CuFunction(md, "compare_comp_float")
        ptxdict[(dev, "compare", Complex{Float64})] = CuFunction(md, "compare_comp_double")
        ptxdict[(dev, "compare", Quaternion{Float32})] = CuFunction(md, "compare_quat_float")
        ptxdict[(dev, "compare", Quaternion{Float64})] = CuFunction(md, "compare_quat_double")
        ptxdict[(dev, "transpose_pre", Complex{Float32})] = CuFunction(md, "transpose_pre_float")
        ptxdict[(dev, "transpose_pre", Complex{Float64})] = CuFunction(md, "transpose_pre_double")
        ptxdict[(dev, "transpose", Quaternion{Float32})] = CuFunction(md, "transpose_float")
        ptxdict[(dev, "transpose", Quaternion{Float64})] = CuFunction(md, "transpose_double")
        ptxdict[(dev, "ctranspose", Quaternion{Float32})] = CuFunction(md, "ctranspose_float")
        ptxdict[(dev, "ctranspose", Quaternion{Float64})] = CuFunction(md, "ctranspose_double")
        push!(mdlist, md)
    end
end

function close()
    for md in mdlist
        unload(md)
    end
    empty!(mdlist)
    empty!(ptxdict)
end

function equiv!{T}(dst::CudaArray{Complex{T}}, src::CudaArray{Quaternion{T}})
    @assert(length(dst) == (4*length(src)))
    dev = device(dst)
    w = size(src,1)
    h = size(src,2)
    block = (16,16)
    grid = (ceil(Int,w/16),ceil(Int,h/16))
    func = ptxdict[(dev, "equiv", Complex{T}, Quaternion{T})]
    CUDArt.launch(func, grid, block, (dst, src, w, h))
end

function equiv!{T}(dst::CudaArray{Quaternion{T}}, src::CudaArray{Complex{T}})
    @assert((4*length(dst)) == length(src))
    dev = device(dst)
    w = size(dst,1)
    h = size(dst,2)
    block = (16,16)
    grid = (ceil(Int,w/16),ceil(Int,h/16))
    func = ptxdict[(dev, "equiv", Quaternion{T}, Complex{T})]
    CUDArt.launch(func, grid, block, (dst, src, w, h))
end

function equiv{T}(src::CudaArray{Quaternion{T}})
    w = size(src,1)
    h = size(src,2)
    dst = CudaArray(Complex{T}, w<<1, h<<1)
    equiv!(dst, src)
    return dst
end

function equiv{T}(src::CudaArray{Complex{T}})
    w = size(src,1)
    h = size(src,2)
    dst = CudaArray(Quaternion{T}, w>>1, h>>1)
    equiv!(dst, src)
    return dst
end

function actfunc!{T}(data::CudaArray{T}; block=block_default)
    dev = device(data)
    N = length(data)
    grid = ceil(Int,N/block)
    func = ptxdict[(dev, "actfunc", T)]
    CUDArt.launch(func, grid, block, (data,N))
end

function actfunc_multi!{T<:Complex}(data::CudaArray{T}, K::Integer; block=block_default)
    dev = device(data)
    N = length(data)
    grid = ceil(Int,N/block)
    func = ptxdict[(dev, "actfunc_multi", T)]
    CUDArt.launch(func, grid, block, (data,K,N))
end

function actfunc_multi!{T<:Quaternion}(data::CudaArray{T}, A::Integer, B::Integer, C::Integer; block=block_default)
    dev = device(data)
    N = length(data)
    grid = ceil(Int,N/block)
    func = ptxdict[(dev, "actfunc_multi", T)]
    CUDArt.launch(func, grid, block, (data,A,B,C,N))
end

function fill!{T}(data::CudaArray{T}, val::T; block=block_default)
    dev = device(data)
    N = length(data)
    grid = ceil(Int,N/block)
    func = ptxdict[(dev, "fill", T)]
    CUDArt.launch(func, grid, block, (data,val,N))
end

function fill!{T}(data::CudaArray{T}, val::T; block=block_default)
    dev = device(data)
    N = length(data)
    grid = ceil(Int,N/block)
    func = ptxdict[(dev, "fill", T)]
    CUDArt.launch(func, grid, block, (data,val,N))
end

function transpose_pre!{T}(dst::CudaArray{T}; block=block_default)
    dev = device(dst)
    w = size(dst,1)
    h = size(dst,2)
    block = (16,16)
    grid = (ceil(Int,w/16),ceil(Int,h/16))
    func = ptxdict[(dev, "transpose_pre", T)]
    CUDArt.launch(func, grid, block, (dst, w, h))
end

function transpose!{T<:Quaternion}(dst::CudaArray{T}, src::CudaArray{T}; block=block_default)
    dev = device(dst)
    w = size(dst,1)
    h = size(dst,2)
    block = (16,16)
    grid = (ceil(Int,w/16),ceil(Int,h/16))
    func = ptxdict[(dev, "transpose", T)]
    CUDArt.launch(func, grid, block, (dst, src, w, h))
end

function transpose{T<:Quaternion}(src::CudaArray{T}; block=block_default)
    w = size(src,1)
    h = size(src,2)
    dst = CudaArray(T, h,w)
    transpose!(dst, src)
    return dst
end

function ctranspose!{T<:Quaternion}(dst::CudaArray{T}, src::CudaArray{T}; block=block_default)
    dev = device(dst)
    w = size(dst,1)
    h = size(dst,2)
    block = (16,16)
    grid = (ceil(Int,w/16),ceil(Int,h/16))
    func = ptxdict[(dev, "ctranspose", T)]
    CUDArt.launch(func, grid, block, (dst, src, w, h))
end

function ctranspose{T<:Quaternion}(src::CudaArray{T}; block=block_default)
    w = size(src,1)
    h = size(src,2)
    dst = CudaArray(T, h,w)
    ctranspose!(dst, src)
    return dst
end
end # module


text:
# # Cells in Literate Julia
#
# The `{:cell}` syntax for [executable cells](#) is not limited to just
# markdown files. It will work the same across other source types as well.
# This page illustrates this using literate Julia.
#
# To use cells in a Julia file just add `{:cell}` at the end of a comment block
# before the code that you would like to be evaluated. Like so:
#
# {:cell}

f(x) = x + 1
f(1)

# The cells behave the same across all source types and so you will need to
# provide named cells to be able to share values between cells.
#
# {cell=example}

f(x::Integer) = x - 1
length(methods(f))

# As we can see above the function `f` only has 1 method defined since it does
# not know about the previous cell. Below we will define another cell but use
# the same name as before so that we can access it's definitions.
#
# {cell=example}

f(1)

# When the resulting value is all that you care about when writing a cell use
# can just discard the output stream and the cell itself using a little CSS
# by adding a `style="display: none"` to the cell. For example below we use
# the attributes `{cell=example output=false style="display: none"}` to hide
# all but the result. The cell content below is `methods(f)`:
#
# {cell=example output=false style="display: none"}

methods(f)


text:
import Base: show, showcompact

type Jagsmodel
  name::String
  ncommands::Int
  nchains::Int
  adapt::Int
  update::Int
  thin::Int
  monitor::Dict
  deviance::Bool
  dic::Bool
  popt::Bool
  model::String
  model_file::String
  data::Dict{String,Any}
  data_file::String
  init::Array{Dict{String,Any},1}
  command::Array{Base.AbstractCmd, 1}
  tmpdir::String
end

function Jagsmodel(;
  name::String="Noname",
  model::String="",
  model_file::String="",
  ncommands::Int=1,
  nchains::Int=4,
  adapt::Int=1000,
  update::Int=10000,
  thin::Int=10,
  monitor=Dict{String,Any}(),
  deviance::Bool=false,
  dic::Bool=false,
  popt::Bool=false,
  updatejagsfile::Bool=true,
  pdir::String=pwd())

  cd(pdir)

  tmpdir = Pkg.dir(pdir, "tmp")
  if !isdir(tmpdir)
    mkdir(tmpdir)
  end

  # Check if .bugs file needs to be updated.
  if length(model) > 0
    update_bugs_file(Pkg.dir(tmpdir, "$(name).bugs"), strip(model))
  elseif length(model) == 0 && length(model_file) >0 && isfile(model_file)
    cp(model_file, Pkg.dir(tmpdir, "$(name).bugs"))
  else
    println("No proper model defined.")
  end
  model_file = "$(name).bugs"

  # Remove old files created by previous runs
  for i in 1:ncommands
    isfile(Pkg.dir(tmpdir, "$(name)-cmd$(i)-index0.txt")) &&
      rm(Pkg.dir(tmpdir, "$(name)-cmd$(i)-index0.txt"));
    isfile(Pkg.dir(tmpdir, "$(name)-cmd$(i)-table0.txt")) &&
      rm(Pkg.dir(tmpdir, "$(name)-cmd$(i)-table0.txt"));
    for j in 0:nchains
       isfile(Pkg.dir(tmpdir, "$(name)-cmd$(i)-chain$(j).R")) &&
        rm(Pkg.dir(tmpdir, "$(name)-cmd$(i)-chain$(j).R"));
    end
  end

  # Create the command array which will be executed in parallel
  cmdarray = fill(``, ncommands)
  for i in 1:ncommands
    jfile = "$(name)-cmd$(i).jags"
    cmdarray[i] = @static is_windows() ? `cmd /c jags $(jfile)` : `jags $(jfile)`
  end

  # DIC needs at least 2 chains
  if (dic || popt) && nchains < 2
    nchains = 2
  end

  data = Dict{String, Any}()
  init = Dict{String, Any}[]

  if length(monitor) == 0
    println("No monitors defined!")
  end

  data_file = "$(name)-data.R"

  jm = Jagsmodel(name,
    ncommands, nchains,
    adapt, update, thin,
    monitor,
    deviance, dic, popt,
    model, model_file,
    data,
    data_file,
    init,
    cmdarray,
    tmpdir);

  if ncommands == 1
    updatejagsfile && update_jags_file(jm)
  else
    for i in 1:ncommands
      updatejagsfile && update_jags_file(jm, i)
    end
  end

  jm
end

#### Function to update the bugs, init and data files

function update_bugs_file(file::String, str::String)
  str2 = ""
  if isfile(file)
    str2 = open(readstring, file, "r")
    str != str2 && rm(file)
  end
  if str != str2
    println("\nFile $(file) will be updated.")
    strmout = open(file, "w")
    write(strmout, str)
    close(strmout)
  else
    println("\nFile $(file) not updated.")
  end
end

#### Function to update the jags file for ncommand == 1

function update_jags_file(model::Jagsmodel)
  jagsstr = "/*\n\tGenerated $(model.name).jags command file\n*/\n"
  if model.deviance || model.dic || model.popt
    jagsstr = jagsstr*"load dic\n"
  end
  jagsstr = jagsstr*"model in $(model.model_file)\n"
  jagsstr = jagsstr*"data in $(model.data_file)\n"
  jagsstr = jagsstr*"compile, nchains($(model.nchains))\n"
  for i in 1:model.nchains
    fname = "$(model.name)-inits$(i).R"
    jagsstr = jagsstr*"parameters in $(fname), chain($(i))\n"
  end
  jagsstr = jagsstr*"initialize\n"
  jagsstr = jagsstr*"update $(model.adapt)\n"
  if model.deviance
    jagsstr = jagsstr*"monitor deviance\n"
  end
  if model.dic
    jagsstr = jagsstr*"monitor pD\n"
    jagsstr = jagsstr*"monitor pD, type(mean)\n"
  end
  if model.popt
    jagsstr = jagsstr*"monitor popt, type(mean)\n"
  end
  for entry in model.monitor
    if entry[2]
      jagsstr = jagsstr*"monitor $(string(entry[1])), thin(1)\n"
    end
  end
  jagsstr = jagsstr*"update $(model.update)\n"
  jagsstr = jagsstr*"coda *, stem($(model.name)-cmd1-)\n"
  jagsstr = jagsstr*"exit\n"
  check_jags_file(Pkg.dir(model.tmpdir, "$(model.name)-cmd1.jags"), jagsstr)
end

#### Function to update the jags file for ncommand > 1

function update_jags_file(model::Jagsmodel, cmd::Int)
  m = max(model.nchains, model.ncommands)
  indx = filter(x -> x!=0, [%(i, (m+1)) for i in 1:3m])
  indx = indx[cmd:length(indx)]
  jagsstr = "/*\n\tGenerated $(model.name).jags command file\n*/\n"
  if model.deviance || model.dic || model.popt
    jagsstr = jagsstr*"load dic\n"
  end
  jagsstr = jagsstr*"model in $(model.model_file)\n"
  jagsstr = jagsstr*"data in $(model.data_file)\n"
  jagsstr = jagsstr*"compile, nchains($(model.nchains))\n"
  for i in 1:model.nchains
    fname = "$(model.name)-inits$(indx[i]).R"
    jagsstr = jagsstr*"parameters in $(fname), chain($(i))\n"
  end
  jagsstr = jagsstr*"initialize\n"
  jagsstr = jagsstr*"update $(model.adapt)\n"
  if model.deviance
    jagsstr = jagsstr*"monitor deviance\n"
  end
  if model.dic
    jagsstr = jagsstr*"monitor pD\n"
    jagsstr = jagsstr*"monitor pD, type(mean)\n"
  end
  if model.popt
    jagsstr = jagsstr*"monitor popt, type(mean)\n"
  end
  for entry in model.monitor
    if entry[2]
      jagsstr = jagsstr*"monitor $(string(entry[1])), thin(1)\n"
    end
  end
  jagsstr = jagsstr*"update $(model.update)\n"
  jagsstr = jagsstr*"coda *, stem($(model.name)-cmd$(cmd)-)\n"
  jagsstr = jagsstr*"exit\n"
  check_jags_file(Pkg.dir(model.tmpdir, "$(model.name)-cmd$(cmd).jags"), jagsstr)
end

function check_jags_file(file::String, str::String)
  str2 = ""
  if isfile(file)
    str2 = open(readstring, file, "r")
    str != str2 && rm(file)
  end
  if str != str2
    println("File $(file) will be updated.")
    strmout = open(file, "w")
    write(strmout, str)
    close(strmout)
  else
    println("File $(file) not updated.")
  end
end

function model_show(io::IO, m::Jagsmodel, compact::Bool=false)
  if compact==true
    println("Jagsmodel(", m.name, ", ",
      m.ncommands, ", ", m.nchains, ", ",
      m.adapt, ", ", m.update, ", ", m.thin, ", ",
      m.monitor, ", ",
      m.deviance, ", ", m.dic, ", ", m.popt, ", ",
      m.model_file, ", ", m.data_file, ", ", m.tmpdir, ")")
  else
    println("  name =                    \"$(m.name)\"")
    println("  ncommands =               $(m.ncommands)")
    println("  nchains =                 $(m.nchains)")
    println("  adapt =                   $(m.adapt)")
    println("  update =                  $(m.update)")
    println("  thin =                    $(m.thin)")
    println("  monitor =                 $(m.monitor)")
    println("  deviance =                $(m.deviance)")
    println("  dic =                     $(m.dic)")
    println("  popt =                    $(m.popt)")
    #println("  model =                   $(model)")
    println("  model_file =              $(m.model_file)")
    #println("  data =                    $(data)")
    println("  data_file =               $(m.data_file)")
    #println("  init =                    $(init)")
    println("  tmpdir =                  $(m.tmpdir)")
  end
end

show(io::IO, m::Jagsmodel) = model_show(io, m, false)
showcompact(io::IO, m::Jagsmodel) = model_show(io, m, true)


text:
using Test, Random
using StatsBase, Distributions, DependentBootstrap
using ForecastEval

@testset "Diebold-Mariano tests" begin
	Random.seed!(1234)
	x = randn(100)
	y = randn(100)
	z = randn(100)
	ld = (x - z).^2 - (y - z).^2
	Random.seed!(1234)
	bO = dm(ld, :boot)
	hO = dm(ld, :hac)
	#v0.7+
	@test (bO.rejH0, bO.bestinput) == (false, 1)
	@test isapprox(bO.pvalue, 0.52)
	@test isapprox(bO.teststat, -0.23500198562777389)
	@test (hO.rejH0, hO.bestinput) == (false, 1)
	@test isapprox(hO.pvalue, 0.5016201102228858)
	@test isapprox(hO.teststat, -0.6719428022372147)
	#------------------------------------------
	#v0.6
	#@test (bO.rejH0, bO.bestinput) == (false, 1)
	#@test isapprox(bO.pvalue, 0.42599999999999993)
	#@test isapprox(bO.teststat, -0.23500198562777386)
	#@test (hO.rejH0, hO.bestinput) == (false, 1)
	#@test isapprox(hO.pvalue, 0.5016201102228858)
	#@test isapprox(hO.teststat, -0.6719428022372147)
end

@testset "Reality Check tests" begin
	Random.seed!(1234)
	ld = randn(100, 20)
	rcOut = rc(ld)
	#v0.7+
	@test rcOut.rejH0 == false
	@test isapprox(rcOut.pvalue, 0.778)
	#------------------------------------------
	#v0.6
	# @test rcOut.rejH0 == false
	# @test isapprox(rcOut.pvalue, 0.762)
end

@testset "SPA tests" begin
	Random.seed!(1234)
	ld = randn(100, 5)
	spaOut = spa(ld)
	#v0.7+
	@test spaOut.rejH0 == false
	@test isapprox(sum(spaOut.mu_u), 0.026125507727482872)
	@test isapprox(sum(spaOut.mu_c), 0.026125507727482872)
	@test isapprox(sum(spaOut.mu_l), 0.21071619368258687)
	@test isapprox(spaOut.pvalue_u, 0.29)
	@test isapprox(spaOut.pvalue_c, 0.29)
	@test isapprox(spaOut.pvalue_l, 0.181)
	@test isapprox(spaOut.pvalueauto, 0.29)
	@test isapprox(spaOut.teststat, 1.5384635269840627)
	#------------------------------------------
	#v0.6
	# @test spaOut.rejH0 == false
	# @test isapprox(sum(spaOut.mu_u), 0.026125507727482858)
	# @test isapprox(sum(spaOut.mu_c), 0.026125507727482858)
	# @test isapprox(sum(spaOut.mu_l), 0.21071619368258687)
	# @test isapprox(spaOut.pvalue_u, 0.25)
	# @test isapprox(spaOut.pvalue_c, 0.25)
	# @test isapprox(spaOut.pvalue_l, 0.124)
	# @test isapprox(spaOut.pvalueauto, 0.25)
	# @test isapprox(spaOut.teststat, 1.5384635269840627)
end

@testset "MCS tests" begin
	Random.seed!(1234)
	l = randn(100, 8)
	mcsOut = mcs(l)
	#v0.7+
	@test mcsOut.inQF == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsOut.outQF == Int[]
	@test isapprox(sum(mcsOut.pvalueQF), 6.702999999999999)
	@test mcsOut.inMT == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsOut.outMT == Int[]
	@test isapprox(sum(mcsOut.pvalueMT), 6.754000000000001)
	#------------------------------------------
	#v0.6
	# @test mcsOut.inQF == Int[5,2,3,7,4,8,6,1]
	# @test mcsOut.outQF == Int[]
	# @test isapprox(sum(mcsOut.pvalueQF), 6.665)
	# @test mcsOut.inMT == Int[5,2,3,7,4,8,6,1]
	# @test mcsOut.outMT == Int[]
	# @test isapprox(sum(mcsOut.pvalueMT), 6.673)
	#------------------------------------------
	Random.seed!(1234)
	mcsLROut = mcs(l, MCSBootLowRAM(l))
	#v0.7+
	@test mcsLROut.inQF == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsLROut.outQF == Int[]
	@test isapprox(sum(mcsLROut.pvalueQF), 6.67)
	@test mcsLROut.inMT == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsLROut.outMT == Int[]
	@test isapprox(sum(mcsLROut.pvalueMT), 6.719)
	#------------------------------------------
	#v0.6
	# @test mcsLROut.inQF == Int[5,2,3,7,4,8,6,1]
	# @test mcsLROut.outQF == Int[]
	# @test isapprox(sum(mcsLROut.pvalueQF), 6.604)
	# @test mcsLROut.inMT == Int[5,2,3,7,4,8,6,1]
	# @test mcsLROut.outMT == Int[]
	# @test isapprox(sum(mcsLROut.pvalueMT), 6.563)
end


text:
using POMDPs
using ContinuousPOMDPTreeSearchExperiments
using POMDPToolbox
using POMCP
using POMCPOW
using ProgressMeter
using PmapProgressMeter
using ParticleFilters
using JLD
using CPUTime
using VDPTag

N = 1000

@everywhere begin
    using VDPTag
    using POMDPs
    using ContinuousPOMDPTreeSearchExperiments
    using POMDPToolbox
    using POMCP
    using POMCPOW
    using ParticleFilters
    using CPUTime

    pomdp = VDPTagPOMDP()
    p_rng = MersenneTwister(4)
    solvers = Dict{String, Union{Solver,Policy}}(

        "pomcpow" => begin
            rollout_policy = ToNextML(mdp(pomdp))
            solver = POMCPOWSolver(tree_queries=10_000,
                                   criterion=MaxUCB(40.0),
                                   final_criterion=MaxTries(),
                                   max_depth=20,
                                   k_action=8.0,
                                   alpha_action=1/20,
                                   k_observation=4.0,
                                   alpha_observation=1/20,
                                   estimate_value=FORollout(rollout_policy),
                                   check_repeat_act=false,
                                   check_repeat_obs=false,
                                   rng=MersenneTwister(13)
                                  )

        end,

        "bt_100" => begin
            rng3 = copy(p_rng)
            rollout_policy = ToNextML(mdp(pomdp))
            tree_queries = 1000
            node_updater = ObsAdaptiveParticleFilter(pomdp, LowVarianceResampler(100), 0.05, rng3)
            solver = POMCPDPWSolver(tree_queries=tree_queries,
                                    c=40.0,
                                    max_depth=20,
                                    k_action=8.0,
                                    alpha_action=1/20,
                                    k_observation=4.0,
                                    alpha_observation=1/20,
                                    estimate_value=FORollout(rollout_policy),
                                    node_belief_updater=node_updater,
                                    default_action=TagAction(false, 0.0),
                                    rng=rng3
                                   )
        end,

        "dpw" => begin
            rng3 = copy(p_rng)
            rollout_policy = ToNextML(mdp(pomdp))
            tree_queries = 10_000
            solver = POMCPDPWSolver(tree_queries=tree_queries,
                                    c=40.0,
                                    max_depth=20,
                                    k_action=8.0,
                                    alpha_action=1/20,
                                    k_observation=4.0,
                                    alpha_observation=1/20,
                                    estimate_value=FORollout(rollout_policy),
                                    default_action=TagAction(false, 0.0),
                                    rng=rng3
                                   )
        end,
    )
end

solver_keys = collect(keys(solvers))
rewards = Dict{String, AbstractVector{Float64}}()
# state_hists = Dict{String, AbstractVector{AbstractVector{state_type(pomdp)}}}()
times = Dict{String, AbstractVector{Float64}}()
steps = Dict{String, AbstractVector{Int}}()

for (j, sk) in enumerate(solver_keys)
    s_rewards = SharedArray(Float64, N)
    # s_hists = SharedArray(AbstractVector{state_type(pomdp)}, N)
    s_times = SharedArray(Float64, N)
    s_steps = SharedArray(Float64, N)
    prog = Progress(N, desc="$sk ($j of $(length(solver_keys)))...")
    pmap(prog, 1:N) do i
        sim_rng = MersenneTwister(i)
        up_rng = MersenneTwister(i+100_000)
        policy = solve(solvers[sk], pomdp)
        up = ObsAdaptiveParticleFilter(deepcopy(pomdp), LowVarianceResampler(10_000), 0.05, up_rng)
        sim = HistoryRecorder(max_steps=40, rng=sim_rng)
        s_times[i] = @CPUelapsed hist = simulate(sim, deepcopy(pomdp), policy, up)
        s_rewards[i] = discounted_reward(hist)
        s_steps[i] = n_steps(hist)
    end
    rewards[sk] = sdata(s_rewards)
    @show mean(rewards[sk])
    # state_hists[sk] = sdata(s_hists)
    times[sk] = sdata(s_times)
    steps[sk] = sdata(s_steps)
    @show sum(times[sk])/sum(steps[sk])
end

for k in solver_keys
    println("$k mean: $(mean(rewards[k])) sem: $(std(rewards[k])/sqrt(N))")
    println("$k time per step: $(sum(times[k])/sum(steps[k]))")
end

filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "data", "compare_$(Dates.format(now(), "E_d_u_HH_MM")).jld")
println("saving to $filename...")
@save(filename, solver_keys, rewards, times, steps)
println("done.")


text:
### FWF reader, due to @aplavin
### https://gist.github.com/aplavin/224a31ea457b6e0ef0f4c1a20bd28850
convert_val(::Type{String}, val::String) = val
convert_val(::Type{Symbol}, val::String) = Symbol(val)
convert_val(typ::Type{<:Integer}, val::String) = parse(typ, val)
convert_val(typ::Type{<:AbstractFloat}, val::String) = parse(typ, replace(val, "D" => "E"))  # tables output from Fortran often have floats as 1D+5 instead of 1E+5
# # usage:
# read_fwf(
#   "file.tab",
#   (
#     colname1=(1, 20, String),
#     colname2=(100, 120, Int),
#     # ...
#   ),
#   skiprows=[1, 2, 3, 7]
# )
function readfwf(io, colspecs; skiprows=[], missingstrings=[])
    cols = Dict(
        k => Vector{Union{typ, Missing}}(undef, 0)
        for (k, (from, to, typ)) in pairs(colspecs)
    )
    for (irow, line) in eachline(io) |> enumerate
        if irow ∈ skiprows continue end
        for (k, (from, to, typ)) in pairs(colspecs)
            s_val = from <= length(line) ? line[from:min(length(line), to)] : ""
            # @show irow line k typ s_val
            f_val = s_val in missingstrings ? missing : convert_val(typ, s_val)
            push!(cols[k], f_val)
        end
    end
    d = [k => identity.(cols[k]) for k in keys(colspecs)] # form Dict from data
    return DataFrame(d)
end

mpc_format_obscode = (Code=(1,3,String),
Long=(5,13,Float64),
cos=(14,21,Float64),
sin=(22,30,Float64),
Name=(31,80,String)
)

# lines in ObsCodes.txt corresponding to space telescopes
spaceobs = union([1,247,249,251,252,260,269], 1223:1232)

# MPC minor planet observatory code reader
readmpcobs(mpcfile::String=joinpath(dirname(pathof(NEOs)), "ObsCodes.txt")) = readfwf(mpcfile, mpc_format_obscode, skiprows=spaceobs)

mpcobscodes = readmpcobs()

# Functions get_eop_iau1980, get_eop_iau2000a were adapted from SatelliteToolbox.jl; MIT-licensed
# these functions avoid the use of @eval
# https://github.com/JuliaSpace/SatelliteToolbox.jl/blob/b95e7f54f85c26744c64270841c874631f5addf1/src/transformations/eop.jl#L87
function get_eop_iau1980(; force=false)
    eop_iau1980_rf = RemoteFiles.@RemoteFile(EOP_IAU1980_RF, "https://datacenter.iers.org/data/csv/finals.all.csv", file="EOP_IAU1980.TXT", updates=:weekly)
    RemoteFiles.download(EOP_IAU1980_RF, force=force)
    eop, ~ = readdlm(RemoteFiles.path(eop_iau1980_rf), ';'; header = true)
    SatelliteToolbox._parse_iers_eop_iau_1980(eop)
end
function get_eop_iau2000a(; force=false)
    eop_iau2000a_rf = RemoteFiles.@RemoteFile(EOP_IAU2000A_RF, "https://datacenter.iers.org/data/csv/finals2000A.all.csv", file="EOP_IAU2000A.TXT", updates=:weekly)
    RemoteFiles.download(EOP_IAU2000A_RF, force=force)
    eop, ~ = readdlm(RemoteFiles.path(eop_iau2000a_rf), ';'; header = true)
    SatelliteToolbox._parse_iers_eop_iau_2000A(eop)
end

eop_IAU1980 = get_eop_iau1980()
eop_IAU2000A = get_eop_iau2000a()

# et: ephemeris time (TDB seconds since J2000.0 epoch)
function observer_position(station_code::Union{Int, String}, et::T; eo::Bool=true,
        eop::Union{EOPData_IAU1980, EOPData_IAU2000A} = eop_IAU1980) where {T<:Number}
    # λ_deg: East longitude (deg)
    # u: distance from spin axis (km), u = r*cos(ϕ)
    # v: height above equatorial plane (km), v = r*sin(ϕ)
    st_code_str = lpad(string(station_code), 3, "0")
    # @show st_code_str filter(i->i.Code==st_code_str, mpcobscodes)
    obscode = filter(i->i.Code==st_code_str, mpcobscodes)
    if nrow(obscode) == 0
        @error "Unknown station code: $station_code"
    elseif nrow(obscode) > 1
        @warn "More than one observatory assigned to code: $station_code"
    end
    λ_deg = obscode.Long[1]
    u = obscode.cos[1]*RE
    v = obscode.sin[1]*RE
    # cartesian components of Earth-fixed position of observer
    λ_rad = deg2rad(λ_deg) # rad
    x_gc = u*cos(λ_rad) #km
    y_gc = u*sin(λ_rad) #km
    z_gc = v #km

    pos_geo = [x_gc, y_gc, z_gc] #km

    # mt2c, dmt2c, gast = t2c_rotation_iau_76_80(et, eo=eo)
    # r_c = mt2c*pos_geo
    # v_c = dmt2c*pos_geo

    utc_secs = et - tdb_utc(et)
    jd_utc = JD_J2000 + utc_secs/daysec
    sv_geo = SatelliteToolbox.satsv(jd_utc, pos_geo, zeros(3), zeros(3))
    # Transform state vector coordinates from geocentric, Earth-fixed frame to inertial (celestial) frame
    if eo
        sv_c = SatelliteToolbox.svECEFtoECI(sv_geo, Val(:ITRF), Val(:GCRF), jd_utc, eop)
    else
        sv_c = SatelliteToolbox.svECEFtoECI(sv_geo, Val(:ITRF), Val(:GCRF), jd_utc)
    end
    r_c = convert(Vector{eltype(sv_c.r)}, sv_c.r)
    v_c = convert(Vector{eltype(sv_c.v)}, sv_c.v)

    return r_c, v_c
end

# conversion of radians to arcseconds
rad2arcsec(x) = 3600rad2deg(x) # rad2deg(rad) -> deg; 3600deg -> arcsec

# conversion of arcseconds to radians
arcsec2rad(x) = deg2rad(x/3600) # arcsec/3600 -> deg; deg2rad(deg) -> rad

# conversion of milli-arcseconds to radians
mas2rad(x) = arcsec2rad(x/1000) # mas/1000 -> arcsec; arcsec2rad(arcsec) -> rad

# IAU 1976/1980 nutation-precession matrix
# tt: days since J2000.0 (TT)
# IAU 1980 nutation angles Δψ (nutation in longitude), Δϵ (nutation in obliquity), both in radians
# dde80, ddp80: nutation corrections wrt IAU 1976/1980 (in radians)
# eo: Earth orientation corrections from IERS
function nupr7680mat(tt, Δϵ_1980, ΔΨ_1980, dde80, ddp80)
    # IAU 1976 precession matrix, J2000.0 to date
    rp = Rz(-PlanetaryEphemeris.zeta(tt))*Ry(PlanetaryEphemeris.Theta(tt))*Rz(-PlanetaryEphemeris.Zeta(tt))
    # Add nutation corrections
    dpsi = ΔΨ_1980 + ddp80 #rad
    deps = Δϵ_1980 + dde80 #rad
    # IAU 1980 nutation matrix (ESAA 1992)
    # Eq. (5-152) from Moyer, 2003
    # ϵ0 (rad): mean obliquity obliquity of the ecliptic
    ϵ0 = PlanetaryEphemeris.ϵ̄(tt)
    # Δϵ (rad): nutation in obliquity
    Δϵ = deps
    # Δψ (rad): nutation in longitude
    Δψ = dpsi
    # ϵ (rad): true obliquity of date
    ϵ = ϵ0 + Δϵ
    rn = Rx(-ϵ)*Rz(-Δψ)*Rx(ϵ0)
    # Combine the matrices:  PN = N x P (with frame-bias B included)
    return rn*rp
end

# Polar motion matrix (TIRS->ITRS, IERS 1996)
# tt: days since J2000.0
function polarmotionmat(tt; eo::Bool=true)
    # Polar motion (arcsec->radians)
    if eo
        xp_arcsec, yp_arcsec = EarthOrientation.polarmotion(JD_J2000+tt)
        xp = deg2rad(xp_arcsec/3600)
        yp = deg2rad(yp_arcsec/3600)
    else
        xp = 0.0
        yp = 0.0
    end
    return Ry(-xp)*Rx(-yp)
end

# Terrestrial-to-celestial rotation matrix (including polar motion)
# Using 1976/1980 Earth orientation/rotation model
# et: ephemeris time (TDB seconds since J2000.0 epoch)
function t2c_rotation_iau_76_80(et::T; eo::Bool=true) where {T<:Number}
    # UTC (JD)
    utc_secs = et - tdb_utc(et)
    t_utc = JD_J2000 + utc_secs/daysec
    t_utc_00 = constant_term(constant_term(t_utc))
    # TT
    t0_tt = et + ttmtdb(et)
    tt = t0_tt/daysec

    #Nutation corrections wrt IAU 1976/1980
    # Output of `EarthOrientation.precession_nutation80` is in mas:
    # https://github.com/JuliaAstro/EarthOrientation.jl/blob/529f12425a6331b133f989443aeb3fbbafd8f324/src/EarthOrientation.jl#L413
    if eo
        ddp80_mas, dde80_mas = EarthOrientation.precession_nutation80(t_utc_00)
    else
        ddp80_mas, dde80_mas = 0.0, 0.0
    end
    # Convert mas -> radians
    ddp80 = mas2rad(ddp80_mas)
    dde80 = mas2rad(dde80_mas)

    et_days = et/daysec # TDB days since J2000.0
    # Mean obliquity (output in radians)
    epsa = PlanetaryEphemeris.ϵ̄(et_days) # rad
    # Lunar longitude of ascending node (measured from ecliptic)
    Ω_M = PlanetaryEphemeris.Ω(et_days) # rad
    # IAU 1980 nutation angles
    _, Δϵ_1980, ΔΨ_1980 = nutation_fk5(JD_J2000 + et_days)
    # Equation of the equinoxes `ee = GAST - GMST`, including nutation correction
    # Expression from Eq. (5-184) of Moyer (2003): Δθ = ∆ψ cosϵ̄ + 0''.00264 sinΩ + 0''.000063 sin2Ω
    # where the nutation in longitude includes corrections from IERS eop file
    ee = (ΔΨ_1980+ddp80)*cos(epsa) + arcsec2rad( 0.00264sin(Ω_M) + 0.000063sin(2Ω_M) )
    # ΔUT1 = UT1-UTC (seconds)
    if eo
        dut1 = EarthOrientation.getΔUT1(t_utc_00)
        # dut1 = eop_IAU1980.UT1_UTC(t_utc)
    else
        dut1 = 0.0
    end
    # UT1
    ut1 = utc_secs + dut1 # elapsed UT1 secs since J2000.0
    ut1_days = ut1/daysec # elapsed UT1 days since J2000.0
    # Greenwich apparent sidereal time (IAU 1982/1994)
    # Eq. (5-173) from Moyer (2003)
    gmst82 = J2000toGMST(ut1_days)
    gast = mod2pi( gmst82 + ee )
    β_dot = eo ? omega(EarthOrientation.getlod(t_utc_00)) : ω

    # For more details, see ESAA 2014, p. 295, Sec. 7.4.3.3, Eqs. 7.137-7.140
    Rz_minus_GAST = [
            cos(gast) -sin(gast) zero(gast);
            sin(gast) cos(gast) zero(gast);
            zero(gast) zero(gast) one(gast)
        ]
    dRz_minus_GAST = (β_dot)*[
            -sin(gast) -cos(gast) zero(gast);
            cos(gast) -sin(gast) zero(gast);
            zero(gast) zero(gast) zero(gast)
        ]

    # Polar motion matrix (TIRS->ITRS, IERS 1996)
    W = polarmotionmat(constant_term(constant_term(tt)), eo=eo)
    W_inv = convert(Matrix{Float64}, transpose(W))
    # _W_inv = SatelliteToolbox.rECEFtoECEF(DCM, ITRF(), PEF(), t_utc, eop_IAU1980)
    # W_inv = convert(Matrix{eltype(_W_inv)}, _W_inv)

    # IAU 76/80 nutation-precession matrix
    C = nupr7680mat(et_days, Δϵ_1980, ΔΨ_1980, dde80, ddp80)
    C_inv = transpose(C)
    # _C_inv = SatelliteToolbox.rECItoECI(DCM, TOD(), GCRF(), t_utc, eop_IAU1980)
    # C_inv = convert(Matrix{eltype(_C_inv)}, _C_inv)

    # Velocity transformation may be retrieved also from: SatelliteToolbox.svECEFtoECI
    mt2c = C_inv*Rz_minus_GAST*W_inv
    dmt2c = C_inv*dRz_minus_GAST*W_inv

    return mt2c, dmt2c, gast
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct IoK8sApiCustomMetricsV1beta1MetricValue <: SwaggerModel
    apiVersion::Any # spec type: Union{ Nothing, String } # spec name: apiVersion
    kind::Any # spec type: Union{ Nothing, String } # spec name: kind
    describedObject::Any # spec type: Union{ Nothing, IoK8sApiCoreV1ObjectReference } # spec name: describedObject
    metricName::Any # spec type: Union{ Nothing, String } # spec name: metricName
    timestamp::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgApisMetaV1Time } # spec name: timestamp
    windowSeconds::Any # spec type: Union{ Nothing, Int64 } # spec name: windowSeconds
    value::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgApiResourceQuantity } # spec name: value

    function IoK8sApiCustomMetricsV1beta1MetricValue(;apiVersion=nothing, kind=nothing, describedObject=nothing, metricName=nothing, timestamp=nothing, windowSeconds=nothing, value=nothing)
        o = new()
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("apiVersion"), apiVersion)
        setfield!(o, Symbol("apiVersion"), apiVersion)
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("kind"), kind)
        setfield!(o, Symbol("kind"), kind)
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("describedObject"), describedObject)
        setfield!(o, Symbol("describedObject"), describedObject)
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("metricName"), metricName)
        setfield!(o, Symbol("metricName"), metricName)
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("timestamp"), timestamp)
        setfield!(o, Symbol("timestamp"), timestamp)
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("windowSeconds"), windowSeconds)
        setfield!(o, Symbol("windowSeconds"), windowSeconds)
        validate_property(IoK8sApiCustomMetricsV1beta1MetricValue, Symbol("value"), value)
        setfield!(o, Symbol("value"), value)
        o
    end
end # type IoK8sApiCustomMetricsV1beta1MetricValue

const _property_map_IoK8sApiCustomMetricsV1beta1MetricValue = Dict{Symbol,Symbol}(Symbol("apiVersion")=>Symbol("apiVersion"), Symbol("kind")=>Symbol("kind"), Symbol("describedObject")=>Symbol("describedObject"), Symbol("metricName")=>Symbol("metricName"), Symbol("timestamp")=>Symbol("timestamp"), Symbol("windowSeconds")=>Symbol("windowSeconds"), Symbol("value")=>Symbol("value"))
const _property_types_IoK8sApiCustomMetricsV1beta1MetricValue = Dict{Symbol,String}(Symbol("apiVersion")=>"String", Symbol("kind")=>"String", Symbol("describedObject")=>"IoK8sApiCoreV1ObjectReference", Symbol("metricName")=>"String", Symbol("timestamp")=>"IoK8sApimachineryPkgApisMetaV1Time", Symbol("windowSeconds")=>"Int64", Symbol("value")=>"IoK8sApimachineryPkgApiResourceQuantity")
Base.propertynames(::Type{ IoK8sApiCustomMetricsV1beta1MetricValue }) = collect(keys(_property_map_IoK8sApiCustomMetricsV1beta1MetricValue))
Swagger.property_type(::Type{ IoK8sApiCustomMetricsV1beta1MetricValue }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCustomMetricsV1beta1MetricValue[name]))}
Swagger.field_name(::Type{ IoK8sApiCustomMetricsV1beta1MetricValue }, property_name::Symbol) =  _property_map_IoK8sApiCustomMetricsV1beta1MetricValue[property_name]

function check_required(o::IoK8sApiCustomMetricsV1beta1MetricValue)
    true
end

function validate_property(::Type{ IoK8sApiCustomMetricsV1beta1MetricValue }, name::Symbol, val)
end


text:
using Documenter, MathematicalSets

DocMeta.setdocmeta!(MathematicalSets, :DocTestSetup,
                    :(using MathematicalSets); recursive=true)

makedocs(;
    sitename="MathematicalSets.jl",
    modules=[MathematicalSets],
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        assets = ["assets/aligned.css"]),
    strict = true,
    pages = [
        "Home" => "index.md",
        "Library" => Any[
            "Types" => "lib/types.md",
            "Methods" => "lib/methods.md"],
        "About" => "about.md"
    ]
)

deploydocs(;
    repo = "github.com/JuliaReach/MathematicalSets.jl.git",
    push_preview = true
)


text:
abstract type AbstractInverseModel end
abstract type AbstractInverseSolution end

Base.@kwdef struct LinearLeastSquares{FT} <: AbstractInverseModel
    "Jacobian Matrix K"
    K::Matrix{FT}
    "Measurement vector y"
    y::Vector{FT}
end

Base.@kwdef struct WeightedLinearLeastSquares{FT,FT2<:AbstractArray{FT}}  <: AbstractInverseModel
    "Jacobian Matrix K"
    K::Matrix{FT}
    "Measurement covariance matrix Sₑ"
    Sₑ::FT2
    "Measurement vector y"
    y::Vector{FT}
end

Base.@kwdef struct OptimalEstimationProblem{FT,FT2<:AbstractArray{FT},FT3<:AbstractArray{FT}} <: AbstractInverseModel
    "Jacobian Matrix K"
    K::Matrix{FT}
    "Measurement covariance matrix Sₑ"
    Sₑ::FT2
    "Measurement vector δy (y-Kxₐ) in linear case, (y-F(xᵢ)+Kᵢ(xᵢ-xₐ)) in non-linear case"
    δy::Vector{FT}
    "Prior state covariance matrix Sₐ"
    Sₐ::FT3
end

Base.@kwdef struct OEsolution{FT} <: AbstractInverseSolution
    "Retrieved state vector x̂"
    x̂::Vector{FT}
    "Posterior Covariance Matrix Ŝ"
    Ŝ::Matrix{FT}
    "Averaging kernel Matrix A"
    A::Matrix{FT}
    "residual vector r"
    r::Vector{FT}
    "reduced χ²"
    χ²::FT

end

max_stars_repo_name: RadiativeTransfer/OptimalEstimation

max_issues_repo_name: RadiativeTransfer/OptimalEstimation

max_forks_repo_name: RadiativeTransfer/OptimalEstimation


text:
" Generate image overapproximations for each discrete state.
# Arguments
- `params::ExperimentParameters` - Experiment parameters structure
- `gp_info_dict::Dict` - Dictionary with GP set and RKHS info
"
function generate_region_images(params::ExperimentParameters, gp_info_dict::Dict)
	region_dict = create_region_data(params.domain, params.discretization_step)
	region_post_dict = Dict()

	# Minimum and maximum extents
	domain = params.domain
    σ_ubs = calculate_σ_global_bounds(domain, gp_info_dict)
    # TODO: Handle this flag better
	# σ_ubs = nothing

    Threads.@threads for i=1:region_dict.count-1
        region_post_dict[i] = bound_extent(region_dict[i], gp_info_dict, params.system_params.dependency_dims; known_component=params.system_params.known_dynamics_fcn, σ_ubs=σ_ubs)
    end

	region_data = Dict(:extents=>region_dict, :posts=>region_post_dict)

    return region_data
end

" Generate image overapproximations for each discrete state using local regressions.
# Arguments
- `params::ExperimentParameters` - Experiment parameters structure
- `x_train::AbstractArray{Any}` - Array of input data
- `y_train::AbstractArray{Any}` - Array of input data
"
function generate_region_images(params::ExperimentParameters, x_train::AbstractArray{Any}, y_train::AbstractArray{Any}; reuse_regions_flag=false)

    exp_dir = params.experiment_directory
    region_filename = "$exp_dir/regions.bson"

    # TODO: Remove this - reloading is handled elsewhere
    if reuse_regions_flag && isfile(region_filename)
        # TODO: Replace this fella!!!!
        region_data = BSON.load(region_filename)
    else
        region_dict = create_region_data(params.domain, params.discretization_step)
        num_regions = region_dict.count
        region_post_dict = Dict()
        region_gp_dict = Dict()

        # Minimum and maximum extents
        domain = params.domain
        min_ex = [domain[dim_key][1] for dim_key in keys(domain)]
        max_ex = [domain[dim_key][2] for dim_key in keys(domain)]

        # Store the predicted mean and covariance for each sampled point
        dim_keys = keys(domain)
        Threads.@threads for i=1:num_regions
            # Get subset of data here
            if i == num_regions
                extent = region_dict[i]
                lb = [extent[dim_key][1] for dim_key in dim_keys]
                ub = [extent[dim_key][2] for dim_key in dim_keys]
                # Assume 2D
                gp_set, gp_info_dict = generate_estimates(params, x_train, y_train, filename_appendix=i)
                save_gp_info(params, gp_set, gp_info_dict, save_global_gps=false, filename_appendix=i)
                region_post_dict[i] = bound_extent(extent, lb, ub, gp_info_dict, dim_keys, params.system_params.dependency_dims; known_part_flag=!isnothing(params.system_params.known_dynamics_fcn))
                region_gp_dict[i] = gp_info_dict
            else
                extent = region_dict[i]
                lb = [extent[dim_key][1] for dim_key in dim_keys]
                ub = [extent[dim_key][2] for dim_key in dim_keys]
                # Assume 2D
                center = [mean(lb), mean(ub)]

                kdtree = KDTree(x_train);

                num_neighbors = params.data_params.num_neighbors
                sub_idx, _ = knn(kdtree, center, num_neighbors, true)

                x_sub = @view x_train[:, sub_idx]
                y_sub = @view y_train[:, sub_idx]

                gp_set, gp_info_dict = generate_estimates(params, x_sub, y_sub, filename_appendix=i)
                # TODO: saving GPs might be a dumb idea to get a true sense of time.
                save_gp_info(params, gp_set, gp_info_dict, save_global_gps=false, filename_appendix=i)
                region_post_dict[i] = bound_extent(extent, lb, ub, gp_info_dict, dim_keys, params.system_params.dependency_dims; known_part_flag=!isnothing(params.system_params.known_dynamics_fcn))
                region_gp_dict[i] = gp_info_dict
            end
        end  # End threaded forloop

        region_data = Dict()
        region_data[:extents] = region_dict
        region_data[:posts] = region_post_dict
        region_data[:gps] = region_gp_dict
    end

    return region_data
end

function calculate_σ_global_bounds(domain, gp_info_dict)
    min_ex = [domain[dim_key][1] for dim_key in keys(domain)]
    max_ex = [domain[dim_key][2] for dim_key in keys(domain)]
    mx = zeros(2,1)
    mx[1] = max_ex[1]*1.1
    mx[2] = max_ex[2]*1.1
    # Store the predicted mean and covariance for each sampled point
    dim_keys = keys(domain)
    σ_ubs = Dict()

    # Calculate GP UBs
    for dim_key in dim_keys
        σ_ubs[dim_key] = 1.1*sqrt(predict_f(gp_info_dict[dim_key].gp, mx)[2][1])
    end

    return σ_ubs
end

" Generate overapproximations of posterior mean and covariance functions using one of several methods.
# Arguments
- `extent::Dict` - Discrete state extent
- `gp_info_dict::Dict` - Dictionary with GP set and RKHS info
- `data_deps::Dict` - Dictionary indicating input-output data dependencies
"
function bound_extent(extent::Dict, gp_info_dict::Dict, data_deps; known_component=nothing, σ_ubs=nothing, σ_approx_flag=false)
    # TODO: mod keyword handling and document
    dim_keys = keys(extent)
    lb = [extent[dim_key][1] for dim_key in dim_keys]
    ub = [extent[dim_key][2] for dim_key in dim_keys]
    post_extent = Dict()
    σ_bounds = []

    for dim_key in dim_keys
        lbf = lb[findall(.>(0), data_deps[dim_key][:])]
        ubf = ub[findall(.>(0), data_deps[dim_key][:])]
        x_lb, μ_L_lb, μ_L_ub = compute_μ_bounds_bnb(deepcopy(gp_info_dict[dim_key].gp), lbf, ubf)
        x_ub, μ_U_lb, μ_U_ub = compute_μ_bounds_bnb(deepcopy(gp_info_dict[dim_key].gp), lbf, ubf, max_flag=true)

        if σ_approx_flag
            _, σ_U_lb, σ_U_ub = compute_σ_ub_bounds_approx(gp_info_dict[dim_key].gp, lbf, ubf)
        elseif !isnothing(σ_ubs)
            _, σ_U_lb, σ_U_ub = compute_σ_ub_bounds_from_gp(gp_info_dict[dim_key].gp, lbf, ubf, ub=σ_ubs[dim_key])
        else
            _, σ_U_lb, σ_U_ub = compute_σ_ub_bounds(gp_info_dict[dim_key].gp, gp_info_dict[dim_key].Kinv, lbf, ubf)
        end

        if !(μ_L_lb <= μ_U_ub)
            @error "μ LB is greater than μ UB! ", μ_L_lb, μ_U_ub
            throw("aborting")
            @assert μ_L_lb <= μ_U_ub
        end
        if !(σ_U_lb <= σ_U_ub)
            @error "Sigma LB is greater than sigma UB! ", σ_U_lb, σ_U_ub
            throw("aborting")
            @assert σ_U_lb <= σ_U_ub
        end
        if !isnothing(known_component)
            # TODO: Assume an identity form for now. Requires its own bounding process.
            post_extent[dim_key] = [extent[dim_key][1] + μ_L_lb, extent[dim_key][2] + μ_U_ub]
        else
            post_extent[dim_key] = [μ_L_lb, μ_U_ub]
        end
        push!(σ_bounds, σ_U_ub)
    end
    return post_extent, σ_bounds
end


text:
# context data management tests

@testset "Basic operations" begin
    c = context()

    # property access
    @test c.generations == 1
    @test c.id > 0
    @test c.hex_id isa String
    @test length(c.hex_id) > 0
    @test c.path isa Vector{Symbol}
    @test length(c.path) == 0
    @test propertynames(c) ∩ (:id, :hex_id, :generations, :data, :path) |> length == 5

    # add entries to current context
    push!(c, "key1" => 1)
    push!(c, "key2" => "hey")
    @test length(c) == 2

    # Test direct data access
    @test c["key1"] == 1
    @test c["key2"] == "hey"

    # save current context to stack
    ContextTracking.save(c)
    @test c.generations == 2   # current + one saved context
    @test length(c) == 2

    # add more "color" to current context.  Assumes Dict.
    push!(c, "key3" => :cool)
    @test length(c) == 3
    @test sort(collect(keys(c.data))) == ["key1", "key2", "key3"]

    # restore the prior context
    ContextTracking.restore(c)
    @test c.generations == 1
    @test length(c) == 2
    @test sort(collect(keys(c.data))) == ["key1", "key2"]

    # test iterability
    @test length(collect(c)) == 2
    @test length(kv for kv in c) == 2

    # empty context data
    empty!(c)
    @test length(c) == 0

    # print
    let io = IOBuffer()
        show(io, c)
        @test length(String(take!(io))) > 0
    end
end

@testset "Custom container" begin
    # create custom context; use random id to avoid getting the default
    c = context(id = UInt(rand(UInt)), container = String[])
    push!(c, "hello")
    push!(c, "world")

    @test length(c.data) == 2
    @test c.data == [ "hello", "world"]
end

@testset "Error handling" begin
    c = context()
    @test_throws UndefVarError c.whatever
end


text:
"""
    distance(a, b, [datum = wgs84])

The Cartesian distance between points `a` and `b`. Uses `datum` to perform
transformations as necessary, and unlike other *Geodesy* functions, this defaults
to use the common WGS-84 datum for convenience.
"""
distance{T <: AbstractVector}(a::T, b::T) = vecnorm(a-b)

# Automatic transformations to ECEF
# Uses wgs84 datum as a default. In most cases, the datum choice will only make
# a small difference to the answer. Nevertheless, is this acceptable?
distance{T <: AbstractVector}(a::T, b::T, datum) = distance(a, b)
distance(a::LLA, b, datum = wgs84) = distance(ECEFfromLLA(datum)(a), b, datum)
distance(a::ECEF, b::LLA, datum = wgs84) = distance(a, ECEFfromLLA(datum)(b), datum)
distance(a::LatLon, b, datum = wgs84) = distance(ECEFfromLLA(datum)(LLA(a)), b, datum)
distance(a::ECEF, b::LatLon, datum = wgs84) = distance(a, ECEFfromLLA(datum)(LLA(b)), datum)
distance(a::UTMZ, b, datum = wgs84) = distance(ECEFfromUTMZ(datum)(a), b, datum)
distance(a::ECEF, b::UTMZ, datum = wgs84) = distance(a, ECEFfromUTMZ(datum)(b), datum)

"""
    distance(utm1, utm2, zone, isnorth, [datum = wgs84])
    distance(a, utm2, zone, isnorth, [datum = wgs84])
    distance(utm1, b, zone, isnorth, [datum = wgs84])

If one or both points are UTM, we need the zone (and particularly the hemisphere,
isnorth = true/false) to determine the Cartesian distance.
"""
distance(a::UTM, b::UTM, zone::Integer, isnorth::Bool, datum = wgs84) = distance(ECEFfromUTM(zone, isnorth, datum)(a), ECEFfromUTM(zone, isnorth, datum)(b), datum)
distance(a, b::UTM, zone::Integer, isnorth::Bool, datum = wgs84) = distance(a, ECEFfromUTM(zone, isnorth, datum)(b), datum)
distance(a::UTM, b, zone::Integer, isnorth::Bool, datum = wgs84) = distance(ECEFfromUTM(zone, isnorth, datum)(a), b, datum)

# Also add geodesic distances here


text:
### BatchBernoulli

struct BatchBernoulli{T<:AbstractArray{<:AbstractFloat}} <: DiscreteBatchDistribution
    p::T
end

Broadcast.broadcastable(b::BatchBernoulli) = Ref(b)

Bernoulli(p::AbstractArray) = BatchBernoulli(p)

function logpdf(b::BatchBernoulli, x)
    ϵ = eps(b.p)
    lp = x .* log.(b.p .+ ϵ) + (1 .- x) .* log.(1 + ϵ .- b.p)
    return lp
end

function rand(rng::AbstractRNG, b::BatchBernoulli, dims::Int...)
    u = randsimilar(rng, b.p, dims...)
    x = convert.(eltype(b.p), b.p .> u)
    return x
end

mean(b::BatchBernoulli) = b.p

function var(b::BatchBernoulli)
    p = mean(b)
    return p .* (1 .- p)
end

function mode(b::BatchBernoulli)
    T = eltype(b.p)
    return convert.(T, b.p .> T(0.5))
end

function kldiv(
    b1::BatchBernoulli{<:AbstractArray{T}},
    b2::BatchBernoulli{<:AbstractArray{T}}
) where {T<:AbstractFloat}
    ϵ = eps(T)
    kl = b1.p .* (log.(b1.p .+ ϵ) .- log.(b2.p .+ ϵ)) .+
         (1 .- b1.p) .* (log.(1 + ϵ .- b1.p) .- log.(1 + ϵ .- b2.p))
    return kl
end

### BatchBernoulliLogit

struct BatchBernoulliLogit{T<:AbstractArray{<:AbstractFloat}} <: DiscreteBatchDistribution
    logitp::T
end

Broadcast.broadcastable(bl::BatchBernoulliLogit) = Ref(bl)

BernoulliLogit(logitp) = BatchBernoulliLogit(logitp)

logpdf(bl::BatchBernoulliLogit, x) = x .* bl.logitp .- log.(1 .+ exp.(bl.logitp))

mean(bl::BatchBernoulliLogit) = StatsFuns.logistic.(bl.logitp)

function var(b::BatchBernoulliLogit)
    p = mean(b)
    return p .* (1 .- p)
end

mode(bl::BatchBernoulliLogit) = mode(Bernoulli(mean(bl)))


text:
struct ImageGenerator
  model
  img
  target_class::Int
end

function ImageGenerator(model, target_class::Int)
  target_class > 1000 && error("Invalid Target Class")
  img = param(rand(224, 224, 3, 1) * 255) |> gpu
  ImageGenerator(model.layers[1:end-1] |> gpu, img, target_class)
end

function (IG::ImageGenerator)(save_dir, img_name = "generated"; niters::Int = 150, lr = 6.0)
  isdir(save_dir) || error("Storage must be a directory")
  mask = zeros(1000, 1) |> gpu
  mask[IG.target_class, 1] = 1.0
  local mean_img = reshape([0.485, 0.456, 0.406], 1, 1, 3, 1) .* 255 |> gpu
  for i in 1:niters
    IG.img.data .= IG.img.data .- mean_img
    outputs = IG.model(IG.img)
    loss = outputs .* mask
    println("Loss after Iteration $i is $(sum(loss)) and Probability is $(softmax(outputs)[IG.target_class, 1])")
    Flux.back!(outputs, Flux.Tracker.data(loss))
    IG.img.data .= IG.img.data .+ lr * IG.img.grad/sqrt(mean(abs2.(IG.img.grad)))
    IG.img.grad .= 0.0
    save("$(save_dir)/$(img_name)_iteration_$i.png", generate_image(Flux.Tracker.data(IG.img)))
  end
end


text:
# functions for writing output

function writeCounts(mesh::PumiMeshDG2; fname="counts")
# write values needed for memory usage estimate
vals = Array{Int}(8)
vals[1] = mesh.numVert
vals[2] = mesh.numEdge
vals[3] = mesh.numEl
vals[4] = mesh.numBoundaryFaces
vals[5] = mesh.numNodesPerType[1]
vals[6] = mesh.numNodesPerType[2]
vals[7] = mesh.numNodesPerType[3]
vals[8] = mesh.numDofPerNode

#=
# estimate jacobian storage size
acc = 0
for i=1:mesh.numDof
  acc += mesh.sparsity_bnds[2, i] - mesh.sparsity_bnds[1, i]
end

size_nz = 64*acc
size_rowval = 64*acc
size_colptr = 64*mesh.numDof

vals[9] = size_nz + size_rowval + size_colptr
=#
# append mpi_rank, file extension
myrank = mesh.myrank
fname2 = string(fname, "_", myrank, ".txt")
writedlm(fname2, vals)

return nothing
end

function writeCounts(mesh::PumiMesh3D; fname="counts")
# write values needed for memory usage estimate
vals = Array{Int}(9)
vals[1] = mesh.numVert
vals[2] = mesh.numEdge
vals[3] = mesh.numFace
vals[4] = mesh.numEl
vals[5] = mesh.numBoundaryFaces
vals[6] = mesh.numNodesPerType[1]
vals[7] = mesh.numNodesPerType[2]
vals[8] = mesh.numNodesPerType[3]
vals[9] = mesh.numDofPerNode

#=
# estimate jacobian storage size
acc = 0
for i=1:mesh.numDof
  acc += mesh.sparsity_bnds[2, i] - mesh.sparsity_bnds[1, i]
end

size_nz = 64*acc
size_rowval = 64*acc
size_colptr = 64*mesh.numDof

vals[9] = size_nz + size_rowval + size_colptr
=#
# append mpi_rank, file extension
myrank = mesh.myrank
fname2 = string(fname, "_", myrank, ".txt")
writedlm(fname2, vals)

return nothing
end

function printBoundaryFaceNums(mesh::PumiMesh2D)

  n = mesh.numBC

  bndry = 1
  edges = Array{Ptr{Void}}(12)
  for i=1:n
    fname = string("boundary_edge_verts", i, ".dat")
    println("printing ", fname)
    f = open(fname, "a+")

    start_index = mesh.bndry_offsets[i]
    end_index = mesh.bndry_offsets[i+1] - 1
    num_edge = end_index - start_index + 1
    arr = Array{Ptr{Void}}(num_edge)

    for i=1:num_edge  # get the mesh edge pointers
      el = mesh.bndryfaces[bndry].element
      local_face = mesh.bndryfaces[bndry].face
      el_ptr = mesh.elements[el]
      apf.getDownward(mesh.m_ptr, el_ptr, mesh.dim-1, edges)
      arr[i] = edges[local_face]
      bndry += 1
    end

    apf.printFaceVertNumbers(arr, mesh.edge_Nptr, mesh.vert_Nptr; fstream=f)
    close(f)
  end

  return nothing
end  # end function

#=
function printdxidx(name::AbstractString, mat)
# print the 4d matrix formatted like dxidx

(p, p, m, n) = size(mat)

f = open(name, "a+")

for i=1:n  # loop over elements
  for j=1:m  # loop over nodes on an element
    println(f, "el ", i, " node ", j, " dxidx = ", mat[:, :, j, i])
  end
end

close(f)

return nothing
end
=#

function printcoords(name::AbstractString, coords)

  f = open(name, "a+")

  (d, numnodes, numel) = size(coords)

  for i=1:numel  # loop over elements
    for j=1:numnodes # loop over nodes
      print(f, i, " ", j)

      for k=1:d
	print(f, " ",  coords[k, j, i])
      end

      print(f, "\n")
    end
  end

  return nothing
end

function writeCounts(mesh::PumiMesh2; fname="counts_0.txt")
# write values needed for memory usage estimate
vals = Array{Int}(9)
vals[1] = mesh.numVert
vals[2] = mesh.numEdge
vals[3] = mesh.numEl
vals[4] = mesh.numBoundaryFaces
vals[5] = mesh.numNodesPerType[1]
vals[6] = mesh.numNodesPerType[2]
vals[7] = mesh.numNodesPerType[3]
vals[8] = mesh.numDofPerNode

# estimate jacobian storage size
acc = 0
for i=1:mesh.numDof
  acc += mesh.sparsity_bnds[2, i] - mesh.sparsity_bnds[1, i]
end

size_nz = 64*acc
size_rowval = 64*acc
size_colptr = 64*mesh.numDof

vals[9] = size_nz + size_rowval + size_colptr

writedlm(fname, vals)

return nothing
end

"""
  This function prints some overview statistics of the mesh.  Specifically,
  it prints the sum of the number of mesh:

   * vertices
   * edges
   * faces (3D only)
   * elements
   * degrees of freedom
   * nodes

   on all processes (for some entities, the value printed is greater than
   the number of *unique* entities because some entities are shared in
   parallel.  This function does not attempt to determine entity uniqueness

   **Inputs**

    * mesh: a mesh obj
    * f: an IO object to print to.  Only mesh.myrank == 0 prints to the IO.
         defaults to STDOUT
"""
function printStats(mesh::PumiMesh2D, f::IO=STDOUT)
  # this might overflow, but its only used for output
  counts = Int64[mesh.numVert, mesh.numEdge, mesh.numEl, mesh.numDof, mesh.numNodes]
  counts_recv = zeros(Int64, length(counts))
  MPI.Allreduce!(counts, counts_recv, MPI.SUM, mesh.comm)

  if mesh.myrank == 0
    println(f, "printin global mesh statistics:")
    println(f, "  numVert = ", counts_recv[1])
    println(f, "  numEdge = ", counts_recv[2])
    println(f, "  numEl = ", counts_recv[3])
    println(f, "  numDof = ", counts_recv[4])
    println(f, "  numNodes = ", counts_recv[5])
  end

  return nothing
end

# 3D method
function printStats(mesh::PumiMesh3D, f::IO=STDOUT)

  # this might overflow, but its only used for output
  counts = Int64[mesh.numVert, mesh.numEdge, mesh.numFace, mesh.numEl, mesh.numDof, mesh.numNodes]
  counts_recv = zeros(Int64, length(counts))
  MPI.Allreduce!(counts, counts_recv, MPI.SUM, mesh.comm)

  if mesh.myrank == 0
    println(f, "printin global mesh statistics:")
    println(f, "  numVert = ", counts_recv[1])
    println(f, "  numEdge = ", counts_recv[2])
    println(f, "  numEdge = ", counts_recv[3])
    println(f, "  numEl = ", counts_recv[4])
    println(f, "  numDof = ", counts_recv[5])
    println(f, "  numNodes = ", counts_recv[6])
  end

  return nothing
end


text:
# Test functions

function trait_test_cash()
    cash = Money("USD", 100.00)
    @show tradable(cash)
    @show marketprice(cash)
    return true
end

function trait_test_stock()
    aapl = Stock("AAPL", "Apple, Inc.")
    @show tradable(aapl)
    @show marketprice(aapl)
    return true
end

function trait_test_residence()
    try
        home = Residence("Los Angeles")
        @show tradable(home)    # returns false
        @show marketprice(home) # exception is raised
    catch ex
        println(ex)
    end
    return true
end

function trait_test_bond()
    try
        bill = TreasuryBill("123456789")
        @show tradable(bill)
        @show marketprice(bill) # exception is raised
    catch ex
        println(ex)
    end
    return true
end

#= REPL
julia> trait_test_cash();
tradable(cash) = true
marketprice(cash) = 100.0

julia> trait_test_stock();
tradable(aapl) = true
marketprice(aapl) = 235

julia> trait_test_residence();
tradable(home) = false
ErrorException("Price for illiquid asset Residence(\"Los Angeles\") is not available.")

julia> trait_test_bond();
tradable(bill) = true
ErrorException("Please implement pricing function for TreasuryBill")
=#

function trait_test_book()
    book = Book("Romeo and Juliet")
    @show tradable(book)
    @show marketprice(book)
    return true
end

#= REPL
julia> trait_test_book();
tradable(book) = true
marketprice(book) = 10.0
=#