module TestDPO
using Test
using Catlab.Graphs
using Catlab.Present
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra.FinSets
using Catlab, Catlab.Theories, Catlab.CategoricalAlgebra

# Wiring diagrams
#################

A, B, C, D = [:A], [:B], [:C], [:D];
f,g,h,i = Box(:f, A, B), Box(:g, B, C),Box(:h, C, D), Box(:i, [:D,:B], [:D,:B]);

GWD = WiringDiagram([:A,:A],[:D,:B])
fv, gv, hv, fv2 = add_box!(GWD, f), add_box!(GWD, g), add_box!(GWD, h), add_box!(GWD, f);
add_wires!(GWD, Pair[
  (input_id(GWD),1) => (fv,1),
  (fv,1) => (gv,1),
  (gv,1) => (hv,1),
  (hv,1) => (output_id(GWD),1),
  (input_id(GWD),2) => (fv2,1),
  (fv2,1) => (output_id(GWD),2)]);

RWD = WiringDiagram([], []);
fv, iv, gv = add_box!(RWD, Box(:f,[],B)), add_box!(RWD, i), add_box!(RWD, Box(:g,B,[]));
add_wires!(RWD, Pair[(iv,1) => (iv, 1),
                     (fv, 1) => (iv, 2),
                     (iv, 2)=>(gv,1)]);

LWD = WiringDiagram([], []);
fv, gv = add_box!(LWD, Box(:f,[],B)), add_box!(LWD, Box(:g,B,[]));
add_wires!(LWD, Pair[
  (fv,1) => (gv,1),
]);

IWD = WiringDiagram([], []);
fv, gv = add_box!(IWD, Box(:f,[],[])), add_box!(IWD, Box(:g,[],[]));

XWD = WiringDiagram([:A,:A],[:D,:B])
fv, gv, hv = add_box!(XWD, f), add_box!(XWD, g), add_box!(XWD, h)
fv2, iv = add_box!(XWD, f), add_box!(XWD, i);
add_wires!(XWD, Pair[
  (input_id(XWD),1) => (fv,1),
  (iv,1) => (iv, 1), (fv, 1) => (iv, 2), (iv, 2)=>(gv,1),
  (gv,1) => (hv,1),
  (hv,1) => (output_id(XWD),1),
  (input_id(XWD),2) => (fv2,1),
  (fv2,1) => (output_id(XWD),2)]);

L=ACSetTransformation(IWD.diagram, LWD.diagram, Box=[1,2]);
R=ACSetTransformation(IWD.diagram, RWD.diagram, Box=[1,3]);
m=ACSetTransformation(LWD.diagram, GWD.diagram, Box=[1,2],InPort=[2],OutPort=[1],Wire=[1]);
@test is_isomorphic(XWD.diagram, rewrite_match(L,R,m))

# Graphs with attributes
########################

@present TheoryDecGraph(FreeSchema) begin
  E::Ob
  V::Ob
  src::Hom(E,V)
  tgt::Hom(E,V)

  X::AttrType
  dec::Attr(E,X)
end

@present TheoryLabelledDecGraph <: TheoryDecGraph begin
  label::Attr(V,X)
end

@acset_type LabelledDecGraph(TheoryLabelledDecGraph, index=[:src,:tgt])

aI2 = @acset LabelledDecGraph{String} begin
  V = 2
  E = 0
  label = ["a","b"]
end

aarr = @acset LabelledDecGraph{String} begin
  V = 2
  E = 1
  src = [1]
  tgt = [2]
  dec = ["e1"]
  label = ["a","b"]
end

abiarr = @acset LabelledDecGraph{String} begin
  V = 2
  E = 2
  src = [1,2]
  tgt = [2,1]
  dec = ["e1","rev_e1"]
  label = ["a","b"]
end

aspan = @acset LabelledDecGraph{String} begin
  V = 3
  E = 2
  src = [1,1]
  tgt = [2,3]
  dec = ["e1","e2"]
  label = ["a","b","c"]
end

expected = @acset LabelledDecGraph{String} begin
  V = 3
  E = 3
  src = [1,1,2]
  tgt = [2,3,1]
  dec = ["e1","e2","rev_e1"]
  label = ["a","b","c"]
end

L = ACSetTransformation(aI2, aarr, V=[1,2]);
R = ACSetTransformation(aI2, abiarr, V=[1,2]);
m = ACSetTransformation(aarr, aspan, V=[2,1], E=[1]);  # sends 'a'->'b' and 'b'->'a'

@test_throws ErrorException("ACSet colimit does not exist: label attributes a != b") rewrite_match(L,R,m)

m = ACSetTransformation(aarr, aspan, V=[1,2], E=[1]);

@test is_isomorphic(expected, rewrite_match(L, R, m))

# Graphs
########

# Example graphs
I2 = Graph(2)
I3 = Graph(3)
#   e1   e2
# 1 <- 2 -> 3
span = Graph(3)
add_edges!(span, [2,2],[1,3])
# 1 -> 2
arr = path_graph(Graph, 2)
# 1 <-> 2
biarr = Graph(2)
add_edges!(biarr, [1,2],[2,1])
# 1 -> 2 -> 3 -> 1
tri = cycle_graph(Graph, 3)
# 4 <- 1 -> 2 and 2 <- 3 -> 4
dispan = Graph(4)
add_edges!(dispan, [1,1,3,3],[2,4,2,4])

#      e1
#    1 -> 2
# e2 |  / ^
#    vv   | e4
#    4 -> 3
#      e5
squarediag = Graph(4)
add_edges!(squarediag, [1,1,2,3,4],[2,4,4,2,3])

# Add a reverse arrow to span
span_w_arrow = Graph(3)
add_edges!(span_w_arrow,[1,1,2],[2,3,1])

L = CSetTransformation(I2, arr, V=[1,2])
R = CSetTransformation(I2, biarr, V=[1,2])
m = CSetTransformation(arr, span, V=[2,1], E=[1])
@test is_isomorphic(span_w_arrow, rewrite_match(L, R, m))

# Remove apex of a subspan (top left corner of squarediag, leaves the triangle behind)
L = CSetTransformation(I2, span, V=[1,3])
m = CSetTransformation(span, squarediag, V=[2,1,4], E=[1,2])
@test is_isomorphic(tri, rewrite_match(L,id(I2),m))

# Remove self-edge using a *non-monic* match morphism
two_loops = Graph(2)
add_edges!(two_loops,[1,2],[1,2]) # ↻1   2↺
one_loop = Graph(2)
add_edges!(one_loop,[2],[2]) # 1   2↺

L = CSetTransformation(I2, arr, V=[1,2])
m = CSetTransformation(arr, two_loops, V=[1, 1], E=[1])
@test is_isomorphic(one_loop, rewrite_match(L,id(I2),m))

# Simplest non-trivial, non-monic exmaple
@present TheoryFinSet(FreeSchema) begin
  X::Ob
end
@acset_type FinSetType(TheoryFinSet)

I, L, G = [FinSetType() for _ in 1:3]
add_parts!(I,:X,2)
add_parts!(L,:X,1)
add_parts!(G,:X,1)
l = CSetTransformation(I,L,X=[1,1])
m = CSetTransformation(L,G,X=[1])
@test can_pushout_complement(l,m)
ik, kg = pushout_complement(l,m)
# There are 3 functions `ik` that make this a valid P.C.
# codom=1 with [1,1], codom=2 with [1,2] or [2,1]
K = codom(ik)
@test nparts(K, :X) == 1 # algorithm currently picks the first option

# Non-discrete interface graph. Non-monic matching
arr_loop= Graph(2)
add_edges!(arr_loop,[1,2],[2,2])  # 1->2↺
arrarr = Graph(2)
add_edges!(arrarr, [1,1],[2,2]) #  1⇉2
arrarr_loop = Graph(2)
add_edges!(arrarr_loop,[1,1,2],[2,2,2]) # 1⇉2↺
arr_looploop = Graph(2)
add_edges!(arr_looploop, [1,2,2],[2,2,2]) # 1-> ↻2↺

L = CSetTransformation(arr, arr, V=[1,2],E=[1]) # identity
R = CSetTransformation(arr, arrarr, V=[1,2], E=[1])
m = CSetTransformation(arr, arr_loop, V=[2,2], E=[2]) # NOT MONIC
@test is_isomorphic(arr_looploop, rewrite_match(L,R,m))

# only one monic match
@test is_isomorphic(arrarr_loop, rewrite(L, R, arr_loop, monic=true))

# two possible morphisms L -> squarediag, but both violate dangling condition
L = CSetTransformation(arr, span, V=[1,2], E=[1]);
m = CSetTransformation(span, squarediag, V=[2,1,4], E=[1,2]);
@test (:src, 5, 4) in dangling_condition(ComposablePair(L,m))

# violate id condition because two orphans map to same point
L = CSetTransformation(I2, biarr, V=[1,2]); # delete both arrows
m = CSetTransformation(biarr, arr_loop, V=[2,2], E=[2,2]);
@test (1, 2) in id_condition(ComposablePair(L[:E],m[:E]))[2]
L = CSetTransformation(arr, biarr, V=[1,2], E=[1]); # delete one arrow
@test 1 in id_condition(ComposablePair(L[:E],m[:E]))[1]

span_triangle = Graph(3); # 2 <- 1 -> 3 (with edge 2->3)
add_edges!(span_triangle,[1,1,2], [2,3,3]);

L = CSetTransformation(arr, tri, V=[1,2], E=[1]);
m = CSetTransformation(tri, squarediag, V=[2,4,3], E=[3,5,4]);
@test is_isomorphic(span_triangle, rewrite_match(L,id(arr),m))

k, g = pushout_complement(L, m); # get PO complement to do further tests

# the graph interface is equal to the final graph b/c we only delete things
@test is_isomorphic(span_triangle, codom(k))

# Check pushout properties 1: apex is the original graph
@test is_isomorphic(squarediag, ob(pushout(L, k))) # recover original graph

# Check pushout properties 2: the diagram commutes
Lm = compose(L,m);
kg = compose(k,g);
for I_node in 1:2
  @test Lm[:V](I_node) == kg[:V](I_node)
end
@test Lm[:E](1) == kg[:E](1)

# Check pushout properties 3: for every pair of unmatched things between K and L, they are NOT equal
for K_node in 1:3
  @test m[:V](3) != g[:V](K_node)
end

for K_edge in 2:3
  @test m[:E](3) != g[:E](K_edge)
end

# Undirected bipartite graphs
#############################

# 1 --- 1
#    /
# 2 --- 2

z_ = UndirectedBipartiteGraph()
add_vertices₁!(z_, 2)
add_vertices₂!(z_, 2)
add_edges!(z_, [1,2,2], [1,1,2])

line = UndirectedBipartiteGraph()
add_vertices₁!(line, 1)
add_vertices₂!(line, 2)
add_edges!(line, [1], [1])

parallel = UndirectedBipartiteGraph()
add_vertices₁!(parallel, 2)
add_vertices₂!(parallel, 2)
add_edges!(parallel, [1,2], [1,2])

merge = UndirectedBipartiteGraph()
add_vertices₁!(merge, 2)
add_vertices₂!(merge, 2)
add_edges!(merge, [1,2], [1,1])

Lspan = UndirectedBipartiteGraph()
add_vertices₁!(Lspan, 1)
add_vertices₂!(Lspan, 2)
add_edges!(Lspan, [1,1],[1,2])

I = UndirectedBipartiteGraph()
add_vertices₁!(I, 1)
add_vertices₂!(I, 2)

L = CSetTransformation(I, Lspan, V₁=[1], V₂=[1,2])
R = CSetTransformation(I, line, V₁=[1], V₂=[1,2])
m1 = CSetTransformation(Lspan, z_, V₁=[1], V₂=[1,2], E=[1, 2])
m2 = CSetTransformation(Lspan, z_, V₁=[1], V₂=[2,1], E=[2, 1])

@test is_isomorphic(parallel, rewrite_match(L, R, m1))
@test is_isomorphic(merge, rewrite_match(L, R, m2))

end


using StateSpaceRoutines, JLD2, HDF5, Random
path = dirname(@__FILE__)

# Initialize arguments to function
@load "$path/reference/kalman_filter_args.jld2" y T R C Q Z D E z0 P0

file  = h5open("$path/reference/kalman_filter_out.h5", "r")
pred  = read(file, "pred")
vpred = read(file, "vpred")
filt  = read(file, "filt")
vfilt = read(file, "vfilt")
close(file)

# Run smoothers
states = Dict{Symbol, Matrix{Float64}}()
shocks = Dict{Symbol, Matrix{Float64}}()

states[:hamilton], shocks[:hamilton] = hamilton_smoother(y, T, R, C, Q, Z, D, E, z0, P0)
states[:koopman], shocks[:koopman] = koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0, pred, vpred)
states[:carter_kohn], shocks[:carter_kohn] = carter_kohn_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = false)
states[:durbin_koopman], shocks[:durbin_koopman] = durbin_koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = false)

# Check that last-period smoothed states equal last-period filtered states
@testset "Ensure last smoothed state = the last filtered state" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        @test filt[:, end] ≈ states[smoother][:, end]
    end
end

# Compare to expected output
exp_states, exp_shocks = h5open("$path/reference/smoothers_out.h5", "r") do file
    read(file, "exp_states"), read(file, "exp_shocks")
end

@testset "Check expected output for all smoothers" begin
    for smoother in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
        @test exp_states ≈ states[smoother]
        @test exp_shocks ≈ shocks[smoother]
    end
end

# Make sure that simulation smoothers run with `draw_states` on
carter_kohn_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = true)
durbin_koopman_smoother(y, T, R, C, Q, Z, D, E, z0, P0; draw_states = true)

# TODO: tests for regime-switching and make sure they match on states across smoothers too, including where the data is missing
# Furthermore, you should test that the implied observables are matching.
nreg = 7
TTTs = Vector{Matrix{Float64}}(undef, nreg)
RRRs = Vector{Matrix{Float64}}(undef, nreg)
CCCs = Vector{Vector{Float64}}(undef, nreg)
ZZs = Vector{Matrix{Float64}}(undef, nreg)
DDs = Vector{Vector{Float64}}(undef, nreg)
QQs = Vector{Matrix{Float64}}(undef, nreg)
EEs = Vector{Matrix{Float64}}(undef, nreg)
regime_inds0 = Vector{Vector{Int64}}(undef, nreg)
regime_inds = Vector{UnitRange{Int64}}(undef, nreg)

for i in 1:nreg
    TTTs[i] = h5read("reference/time_varying_system.h5", "T$(i)")
    RRRs[i] = h5read("reference/time_varying_system.h5", "R$(i)")
    CCCs[i] = h5read("reference/time_varying_system.h5", "C$(i)")
    ZZs[i] = h5read("reference/time_varying_system.h5", "Z$(i)")
    DDs[i] = h5read("reference/time_varying_system.h5", "D$(i)")
    QQs[i] = h5read("reference/time_varying_system.h5", "Q$(i)")
    EEs[i] = h5read("reference/time_varying_system.h5", "E$(i)")
    regime_inds0[i] = h5read("reference/time_varying_system.h5", "regime_inds$(i)")
    regime_inds[i] = regime_inds0[i][1]:regime_inds0[i][end]
end

y = h5read("reference/time_varying_system.h5", "data")

# Run smoothers
states = Dict{Symbol, Matrix{Float64}}()
shocks = Dict{Symbol, Matrix{Float64}}()

z0 = zeros(size(TTTs[1], 1))
P0 = StateSpaceRoutines.solve_discrete_lyapunov(TTTs[1], RRRs[1] * QQs[1] * RRRs[1]')
_, pred, vpred, ~ = kalman_filter(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0)

states[:hamilton], shocks[:hamilton] =
    hamilton_smoother(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0)
states[:koopman], shocks[:koopman] =
    koopman_smoother(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0, pred, vpred)
states[:carter_kohn], shocks[:carter_kohn] =
    carter_kohn_smoother(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0; draw_states = false)
states[:durbin_koopman], shocks[:durbin_koopman] =
    durbin_koopman_smoother(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0; draw_states = false)

Random.seed!(1793)
states[:carter_kohn_draw], shocks[:carter_kohn_draw] =
    carter_kohn_smoother(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0; draw_states = true)
states[:durbin_koopman_draw], shocks[:durbin_koopman_draw] =
    durbin_koopman_smoother(regime_inds, y, TTTs, RRRs, CCCs, QQs, ZZs, DDs, EEs, z0, P0; draw_states = true)

obs = Dict{Symbol, Matrix{Float64}}()
for k in [:hamilton, :koopman, :carter_kohn, :durbin_koopman, :carter_kohn_draw, :durbin_koopman_draw]
    obs[k] = zeros(size(y, 1), size(states[:hamilton], 2))
    for (i, inds) in enumerate(regime_inds)
        obs[k][:, inds] = ZZs[i] * states[k][:, inds] .+ DDs[i]
    end
end

@testset "Regime-switching in T, R, etc., with smoothers" begin
    @test states[:hamilton] ≈ states[:carter_kohn]
    @test states[:koopman] ≈ states[:durbin_koopman]
    @test shocks[:hamilton] ≈ shocks[:carter_kohn]
    @test shocks[:koopman] ≈ shocks[:durbin_koopman]
    @test maximum(abs.(states[:hamilton] - states[:koopman])) < 5e-3
    @test maximum(abs.(shocks[:hamilton] - shocks[:koopman])) < 1.5e-3
    for i in 1:size(y, 1)
        not_nan = findall(.!isnan.(y[i, :]))
        if !isempty(not_nan)
            for k in [:hamilton, :koopman, :carter_kohn, :durbin_koopman]
                @test obs[k][i, not_nan] ≈ y[i, not_nan] atol=6e-6
            end
            if i in [1, 3, 7, 8, 13]
                @test obs[:carter_kohn_draw][i, not_nan] ≈ y[i, not_nan] atol=1e-2
            elseif i in [4, 5]
                @test obs[:carter_kohn_draw][i, not_nan] ≈ y[i, not_nan] atol=5e-4
            else
                @test obs[:carter_kohn_draw][i, not_nan] ≈ y[i, not_nan] atol=5e-4
            end
            @test obs[:durbin_koopman_draw][i, not_nan] ≈ y[i, not_nan] atol=5e-4
        end
    end
end

nothing


module TimeSeriesA2P2

#Use this to print
#=
using Weave
codeName = "TimeSeriesA2P2"
weave(Pkg.dir("$(pwd())\\$(codeName).jl"),
  informat="script",
  out_path = "$(pwd())\\output\\$(codeName)_Appendix.html",
  doctype = "md2html")
=#

if pwd() ∉ LOAD_PATH
  push!(LOAD_PATH,pwd())
end

using DataFrames, Distributions, StatsBase, GZip, JLD,
  Gadfly, CTIO, CTReg, JuMP, NLopt, HCubature, StaticArrays, CTNumerical,
  ParallelAccelerator, Measures, Formatting

const DATA_PATH = pwd() * "\\data" #path where we store the data
const OUTPUT_PATH = pwd() * "\\output"
const YAHOO_DATE_FORMAT = "m/d/yyyy"
const SP500_NAME = "GSPC"
const DBV_NAME = "DBV"

const FOOTER_NAME = "footer.tex" #these are just headers for the tex tables
const HEADER_NAME = "header.tex"
const DECIMALS = 3
const PlotContainer = Union{Plot,Gadfly.Compose.Context,Vector{Plot}} #holds graphs
const OVERRIDE_STAR_STRINGS = ["\\ostar","\\dstar","\\tstar"]

#asignment specific constants
const FACTOR_NAME = "FMFactors"
const ASSET_NAME = "PORTSBM55"
const CUTOFF_DATE = 1963_07
const FACTOR_SYMS = [:Mkt_RF, :SMB, :HML]
const RFR_SYM = :RF
const MKT_SYM = :Mkt

BLAS.set_num_threads(1)

  n2s(x::T where T<:Real) = num2Str(x, DECIMALS, Ints=true, scaleHurdle=99.)

#pre-process the asset data
function preProcessMonthlyFA()::Void

  #read in the data (note the non-standard NA values)
  assetDF::DataFrame = readtable("$DATA_PATH\\$ASSET_NAME.csv", nastrings=["-999", "-99.99"])
  factorDF::DataFrame = readtable("$DATA_PATH\\$FACTOR_NAME.csv", nastrings=["-999", "-99.99"])

  factorDF[:,:Mkt] = factorDF[:,:Mkt_RF] .+ factorDF[:,:RF]

  #merge the dataframes
  DF::DataFrame = join(factorDF, assetDF, on = :DATE, kind=:inner)

  #write out a JLS binary file
  stream::IOStream = open("$DATA_PATH\\$(FACTOR_NAME)_$(ASSET_NAME).jls", "w")
  serialize(stream, DF)
  close(stream)
  return nothing
end

#gets the dataframe as a serial file
function getDF(pathStr::String, refreshData::Bool, preProcessFunc::Function)::DataFrame

  #re-process the data if needed or desired
  if !isfile(pathStr) || refreshData
    preProcessFunc()
  end

  #read the binary file
  stream::IOStream = open(pathStr)
  DF::DataFrame = deserialize(stream)
  close(stream)

  return DF
end

mutable struct PCA
  λ::Vector{Float64}
  ν::Vector{Vector{Float64}}
  svd::Base.LinAlg.SVD
  labels::Vector{Symbol}
  labelsInd::Dict
  normConstant::Float64
end

function PCA(Σ::Matrix{Float64},
    labels::Vector{Symbol} = ((i::Int)->Symbol("v$i")).(1:size(Σ,1)))::PCA
  #setup and factorize
  M::Int, N::Int = size(Σ)
  labelsInd::Dict = Dict(labels[i]=>i for i::Int ∈ 1:M)

  Σsvd::Base.LinAlg.SVD = svdfact(Σ)

  #recover the eigenvalues
  λ::Vector{Float64} = Σsvd[:S]
  normConstant::Float64 = sum(λ)
  λ ./= normConstant

  #recover the eigenvectors
  ν::Vector{Vector{Float64}} =
    ((i::Int)->(Σsvd[:U][:,i]./sum((abs2).(Σsvd[:U][:,i])))).(1:N)

  return PCA(λ, ν, Σsvd,labels,labelsInd,normConstant)
end

function PCA(DF::DataFrame, labels::Vector{Symbol})
  D::Matrix{Float64} = Matrix{Float64}(size(DF[:,labels]))
  D .= Matrix{Float64}(DF[:,labels])
  D .-= mean(D,1)
  return PCA(BLAS.gemm('T','N',1.0, D,D), labels) #(This is simply D'D)
end

#adds principal componenets to an existing dataframe
function PC(pca::PCA, DF::DataFrame;
    labels::Vector{Symbol}=pca.labels,
    dataLabels::Vector{Symbol}=labels,
    PCLabels::Vector{Symbol} = ((i::Int)->(Symbol("PC$(i)"))).(1:length(labels)),
    )::Tuple{DataFrame, Vector{Symbol}}

    D::Matrix{Float64} = DF[:,dataLabels]
    for i::Int ∈ 1:length(PCLabels)
      DF[:,PCLabels[i]] = BLAS.gemv('N', 1.0, D, pca.ν[pca.labelsInd[labels[i]]])
    end

    return (DF, PCLabels)
end

#Calculates the summary statistics
function summaryA2P2(FADF::DataFrame, PCLabels::Vector{Symbol}, cols::Vector{Symbol})
  #label the table
  colNames::Vector{Vector{String}} =
    [((s::Symbol)->replace(String(s),"_","")).(cols)]
  numCols = length(colNames[end])

  rowNames::Vector{String} = ((s::Symbol)->replace(String(s),"_","")).(PCLabels)

  numRows::Int = length(rowNames)
  content::Vector{Vector{String}} = [Vector{String}(numCols) for i::Int ∈ 1:numRows]

  #run the tests and record the results
  for r::Int ∈ 1:numRows, c::Int ∈ 1:numCols
    content[r][c] = n2s(cor(FADF[:,PCLabels[r]], FADF[cols[c]]))
  end

  #write the table
  summaryTable::String = texTable( "Correlations of Principal Components",
    """See Tex File""", #caption
    colNames, #colNames
    Vector{String}(),#contentRowNames
    Vector{Matrix{String}}(), #content
    rowNames, #descRowNames
    content, #descContent
    Vector{String}(), #notes
  )

  return summaryTable
end

#problem specific visualizations
function visualizeA2P2(pca::PCA, FADF::DataFrame, PCLabels::Vector{Symbol})::Void

  #convenience containers
  K::Int = length(pca.λ)
  plots::Vector{PlotContainer} = Vector{PlotContainer}()
  plotNames::Vector{String} = Vector{String}()

  Gadfly.push_theme(:default)

  #some basic data processing for the eigenvalues
  λLabels::Vector{Symbol} = ((i::Int)->(Symbol("λ$i"))).(1:K)
  λDataLabels::Vector{String} = (n2s).(pca.λ)
  λCum::Vector{Float64} = ((i::Int)->sum(pca.λ[1:i])).(1:K)
  λCumDataLabels::Vector{String} = (n2s).(λCum)

  push!(plotNames, "eigenvalues")
  push!(plots, plot(y=pca.λ, x=λLabels, label=λDataLabels, Geom.bar, Geom.label(position=:above),
    Guide.xlabel("λ"),Guide.ylabel("λᵢ / sum(λ)"),Coord.Cartesian(ymin=0., ymax=1.05),
    Guide.title("Eigenvalues λ")))

  push!(plotNames, "cumulative eig")
  push!(plots, plot(y=λCum, x=λLabels, label=λCumDataLabels,
    Geom.point, Geom.line, Geom.label(position=:above),
    Guide.xlabel("λ"),Guide.ylabel("Cumulative λᵢ / sum(λ)"),Coord.Cartesian(ymin=0., ymax=1.05),
    Guide.title("Cumulative Eigenvalues λ")))

  PC3DF::DataFrame = melt(DataFrame(labels=pca.labels,
    ν1=pca.ν[1], ν2=pca.ν[2], ν3=pca.ν[3]),:labels)
  push!(plotNames, "eigenvectors")
  push!(plots, plot(PC3DF, y=:value, x=:labels, color=:variable,
    Geom.bar(position=:dodge),# Geom.line, Geom.point,
    Guide.xlabel("λ"), Guide.ylabel("νᵢ/sum(ν.²)"),
    Guide.title("Eigenvectors ν")))

  push!(plotNames, "eigenvectors ν₁ vs ν₂")
  push!(plots,
    plot(x=pca.ν[1], y=pca.ν[2], label=(string).(pca.labels),
      Geom.point,Geom.label,
      Guide.xlabel("ν₁"), Guide.ylabel("ν₂"),
      Coord.Cartesian(ymin=-0.5, ymax=0.5, xmin=-0.5, xmax=0.5),
      Guide.title("Eigenvectors ν₁ vs ν₂")))
  push!(plotNames, "eigenvectors ν₂ vs ν₃")
  push!(plots,
    plot( x=pca.ν[2], y=pca.ν[3], label=(string).(pca.labels),
      Geom.point,Geom.label,# Geom.line, Geom.point,
      Guide.xlabel("ν₂"), Guide.ylabel("ν₃"),
      Coord.Cartesian(ymin=-0.5, ymax=0.5, xmin=-0.5, xmax=0.5),
      Guide.title("Eigenvectors ν₂ vs ν₃")))

  push!(plotNames, "eigenvectors ν₃ vs ν₁")
  push!(plots,
    plot(x=pca.ν[3], y=pca.ν[1], label=(string).(pca.labels),
      Geom.point,Geom.label,# Geom.line, Geom.point,
      Guide.xlabel("ν₃"), Guide.ylabel("ν₁"),
      Coord.Cartesian(ymin=-0.5, ymax=0.5, xmin=-0.5, xmax=0.5),
      Guide.title("Eigenvectors ν₃ vs ν₁")))

  #write the graphs
  for i::Int ∈ 1:length(plots)
    draw(SVG("$OUTPUT_PATH\\$(plotNames[i]).svg", 9inch, 7inch), plots[i])
  end

  return nothing

end

#problem specific code
function A2P2Script(refreshData::Bool = false)::Void
  FADF::DataFrame = getDF("$DATA_PATH\\$(FACTOR_NAME)_$(ASSET_NAME).jls",
    refreshData, preProcessMonthlyFA)
  FADF = FADF[FADF[:,:DATE].≥196307,:]

  #rename the columns to a common pattern
  rename!(FADF, [:SMALL_LoBM, :SMALL_HiBM, :BIG_LoBM, :BIG_HiBM],
    [:ME1_BM1, :ME1_BM5, :ME5_BM1, :ME5_BM5])

  #get portfolio handles
  FSyms::Vector{Symbol} = FACTOR_SYMS
  RFRSym::Symbol = RFR_SYM
  RSyms::Vector{Symbol} =
    names(FADF[:,setdiff(names(FADF), [:DATE; RFR_SYM; MKT_SYM; FACTOR_SYMS])])

  #do the eigenvalue decomposition
  Σpca::PCA = PCA(FADF, RSyms)

  #get the principal components
  (FADF, PCLabels::Vector{Symbol}) = PC(Σpca, FADF)
  visualizeA2P2(Σpca, FADF, PCLabels)

  summaryTable::String = summaryA2P2(FADF,PCLabels[1:5],[RFR_SYM; MKT_SYM; FACTOR_SYMS])

  writeTables2File([summaryTable],
    HEADER_NAME, FOOTER_NAME, path=OUTPUT_PATH,
    outName = "TimeSeries A2P2 corTable.tex")

  return nothing
end

#uncomment to run
@time begin
  A2P2Script()
end

end


export nonspatial, nonspatial!

function birth!(tumor::Vector{Cell}, mutations::Vector{Mutation}, parent, cur_id, cur_mutation, μ, t)

    pos = Float64[]

	new = Cell(cur_id, pos, copy(parent.mutations), parent.index, parent.b, t, SVector{0,Float64}())
	push!(tumor, new)

	m = 0
	if rand() < μ/2
		m += 1
		push!(mutations, Mutation(parent.index, isempty(parent.mutations) ? 0 : last(parent.mutations), t, length(tumor), SVector{0,Float64}()))
		push!(parent.mutations, cur_mutation+m)
	end
	if rand() < μ/2
		m += 1
		push!(mutations, Mutation(new.index, isempty(new.mutations) ? 0 : last(new.mutations), t, length(tumor), SVector{0,Float64}()))
		push!(new.mutations, cur_mutation+m)
	end
	return m
end

function nonspatial!( tumor::Vector{Cell}, mutations::Vector{Mutation}, until; b, d, μ, cur_id = 1, cur_mutation = 0, t = 0.0, seed=nothing, showprogress=true)

    isnothing(seed) || Random.seed!(seed)

    N = length(tumor)

    prog = showprogress ? ProgressUnknown("Progress: Tumor size ") : nothing
    while loop_condition(N,t, until)

        N==0 && error("Tumor died")
        row = rand(1:N)
        parent = tumor[row]

        p = rand()*(b+d) - b
        t += randexp()/((d+b)*N)

        if p < 0.
            cur_id += 1
            N += 1
            cur_mutation += birth!(tumor, mutations, parent, cur_id, cur_mutation, μ, t)
        elseif p < d
            N -= 1
            deleteat!(tumor, row)
        end
        showprogress && ProgressMeter.update!(prog, N )
    end
    return Dict{Symbol, Any}(:index => cur_id, :mutation => cur_mutation, :time => t)
end
nonspatial!( until; tumor::Vector{Cell}, mutations::Vector{Mutation}, simparams...) =  nonspatial!( tumor, mutations, until; simparams...)

function nonspatial( until; cur_mutation = 0, b, simparams...)
    p₀ = Float64[]
    tumor = [Cell(index=1, position=p₀, parent=0, mutations=collect(1:cur_mutation), b=b, t_birth=0.0, p_birth=SVector{0,Float64}(p₀))]

	mutations = [Mutation(; origin=1, ancestor=0, t_birth=0., N_birth=1, p_birth=SVector{0,Float64}(p₀)) for m=1:cur_mutation]

    output = nonspatial!(tumor, mutations, until; cur_mutation=cur_mutation, b=b, simparams...)
	output[:tumor] = tumor
	output[:mutations] = mutations

	return output
end


"""
    coef(MLM)

Extracts coefficients from Mlm object

# Arguments

- MLM = Mlm object

# Value

2d array of floats

"""
function coef(MLM::Mlm)

    return MLM.B
end

"""
    predict(MLM, newPredictors)

Calculates new predictions based on Mlm object

# Arguments

- MLM = Mlm object
- newPredictors = Predictors object. Defaults to the `data.predictors` field
  in the Mlm object used to fit the model.

# Value

Response object

"""
function predict(MLM::Mlm, newPredictors::Predictors=MLM.data.predictors)

  	# Include X and Z intercepts in new data if necessary
  	if MLM.data.predictors.hasXIntercept==true &&
       newPredictors.hasXIntercept==false
    	newPredictors.X = add_intercept(newPredictors.X)
    	newPredictors.hasXIntercept = true
    	println("Adding X intercept to newPredictors.")
  	end
  	if MLM.data.predictors.hasZIntercept==true &&
       newPredictors.hasZIntercept==false
    	newPredictors.Z = add_intercept(newPredictors.Z)
    	newPredictors.hasZIntercept = true
    	println("Adding Z intercept to newPredictors.")
  	end

  	# Remove X and Z intercepts in new data if necessary
  	if MLM.data.predictors.hasXIntercept==false &&
       newPredictors.hasXIntercept==true
    	newPredictors.X = remove_intercept(newPredictors.X)
    	newPredictors.hasXIntercept = false
    	println("Removing X intercept from newPredictors.")
  	end
  	if MLM.data.predictors.hasZIntercept==false &&
       newPredictors.hasZIntercept==true
    	newPredictors.Z = remove_intercept(newPredictors.Z)
    	newPredictors.hasZIntercept = false
    	println("Removing Z intercept from newPredictors.")
  	end

    # Calculate predictions
    return Response(calc_preds(newPredictors.X, newPredictors.Z, coef(MLM)))
end

"""
    fitted(MLM)

Calculate fitted values of an Mlm object

# Arguments

- MLM = Mlm object

# Value

Response object

"""
function fitted(MLM::Mlm)

    # Call the predict function with default newPredictors
    return predict(MLM)
end

"""
    resid(MLM, newData)

Calculates residuals of an Mlm object

# Arguments

- MLM = Mlm object
- newData = RawData object. Defaults to the `data` field in the Mlm object
  used to fit the model.

# Value

2d array of floats

"""
function resid(MLM::Mlm, newData::RawData=MLM.data)

    # Include X and Z intercepts in new data if necessary
    if MLM.data.predictors.hasXIntercept==true &&
       newData.predictors.hasXIntercept==false
        newData.predictors.X = add_intercept(newData.predictors.X)
        newData.predictors.hasXIntercept = true
        newData.p = newData.p + 1
        println("Adding X intercept to newData.")
    end
    if MLM.data.predictors.hasZIntercept==true &&
       newData.predictors.hasZIntercept==false
        newData.predictors.Z = add_intercept(newData.predictors.Z)
        newData.predictors.hasZIntercept = true
        newData.q = newData.q + 1
        println("Adding Z intercept to newData.")
    end

    # Remove X and Z intercepts in new data if necessary
    if MLM.data.predictors.hasXIntercept==false &&
       newData.predictors.hasXIntercept==true
        newData.predictors.X = remove_intercept(newData.predictors.X)
        newData.predictors.hasXIntercept = false
        newData.p = newData.p - 1
        println("Removing X intercept from newData.")
    end
    if MLM.data.predictors.hasZIntercept==false &&
       newData.predictors.hasZIntercept==true
        newData.predictors.Z = remove_intercept(newData.predictors.Z)
        newData.predictors.hasZIntercept = false
        newData.q = newData.q - 1
        println("Removing Z intercept from newData.")
    end

	# Calculate residuals
    return calc_resid(get_X(newData), get_Y(newData), get_Z(newData),
                      coef(MLM))
end


include("autodiff.jl")
include("utils.jl")

@info "Manifolds.jl Test settings:\n\n"*
    "Testing Float32:  $(TEST_FLOAT32)\n"*
    "Testing Double64: $(TEST_DOUBLE64)\n"*
    "Testing Static:   $(TEST_STATIC_SIZED)\n"

@testset "Ambiguities" begin
    # TODO: reduce the number of ambiguities
    @test length(Test.detect_ambiguities(ManifoldsBase)) <= 13
    @test length(Test.detect_ambiguities(Manifolds)) == 0
end

include("groups/group_utils.jl")
include("sized_abstract_array.jl")

# starting with tests of simple manifolds
include("circle.jl")
include("cholesky_space.jl")
include("euclidean.jl")
include("fixed_rank.jl")
include("generalized_grassmann.jl")
include("generalized_stiefel.jl")
include("grassmann.jl")
include("hyperbolic.jl")
include("rotations.jl")
include("skewsymmetric.jl")
include("sphere.jl")
include("stiefel.jl")
include("symmetric.jl")
include("symmetric_positive_definite.jl")

include("oblique.jl")

include("torus.jl")

#meta manifolds
include("product_manifold.jl")
include("power_manifold.jl")
include("vector_bundle.jl")
include("graph.jl")

include("metric.jl")
include("statistics.jl")

# Lie groups and actions
include("groups/groups_general.jl")
include("groups/array_manifold.jl")
include("groups/circle_group.jl")
include("groups/translation_group.jl")
include("groups/special_orthogonal.jl")
include("groups/product_group.jl")
include("groups/semidirect_product_group.jl")
include("groups/special_euclidean.jl")
include("groups/group_operation_action.jl")
include("groups/rotation_action.jl")
include("groups/translation_action.jl")
include("groups/metric.jl")


module TestUniform

using StatsBase
using Test
using TreeParzen
import TreeParzen: API, Trials

N = 10_000
COUNTMAX = 1300
COUNTMIN = 700

samples = [TreeParzen.Resolve.node(HP.Uniform(:u, 0.0, 10.0), TreeParzen.Trials.ValsDict()) for i in 1:N]

@test minimum(samples) > 0
@test maximum(samples) < 10
h = fit(StatsBase.Histogram, samples, nbins=10).weights
@test all(COUNTMIN .< h)
@test all(h .< COUNTMAX)

end
true


#!/usr/bin/env julia
#
# Benchmarking script for singular value problems

try
    Pkg.installed("Benchmarks")
catch
    Pkg.clone("https://github.com/johnmyleswhite/Benchmarks.jl")
end

try
    Pkg.installed("PROPACK")
catch
    Pkg.clone("https://github.com/andreasnoack/PROPACK.jl")
end

using Benchmarks
using IterativeSolvers
using JLD
using MAT
using PROPACK

BASEDIR = "florida"

if !isdir(BASEDIR)
    info("Setting up $BASEDIR collection")
    include("setup-florida.jl")
end

"""
Run benchmark function `runbenchmark` against all .mat files in all
subdirectories of BASEDIR
"""
function benchmarkall(BASEDIR)
    for group in readdir(BASEDIR)
        isdir(joinpath(BASEDIR, group)) || continue
        for matrix in readdir(joinpath(BASEDIR, group))
            filename = joinpath(BASEDIR, group, matrix)

            endswith(matrix, ".mat") || continue

        #If we already saved benchmarks, don't rerun
        benchmarkfilename = joinpath(BASEDIR, group, matrix[1:end-3]*"jld")
        isfile(benchmarkfilename) && continue

        #To debug this script, it's useful to run it on small matrices only
        #filesize(filename) < 100000 || continue

        runbenchmark(filename, benchmarkfilename)

        #debug: just try one
        #@goto stop
        end
    end
    @label stop
end

"""
Run several different computations for singular values
"""
function runbenchmark(filename, benchmarkfilename)
    mf = matread(filename)
    if !haskey(mf, "Problem")
        warn("Skipping unknown file $filename: No 'Problem' struct")
        return
    end

    pr = mf["Problem"]
    if !haskey(pr, "A")
        warn("Skipping unknown file $filename: 'Problem' struct has no matrix 'A'")
        return
    end
    A = pr["A"]
    #eltype(A) <: Real || warn("Skipping matrix with unsupported element type $(eltype(A))")
    info(filename*", matrix of type $(typeof(A)) and size $(size(A))")
    m, n = size(A)
    #Choose the same normalized unit vector to start with
    q = randn(n)
    eltype(A) <: Complex && (q += im*randn(n))
    rmul!(q, inv(norm(q)))
    qm = randn(m)
    eltype(A) <: Complex && (q += im*randn(m))
    rmul!(qm, inv(norm(qm)))

    #Number of singular values to request
    nv = min(m, n, 10)

    #Maximum number of iterations
    maxiter = max(m, n)

    #Tolerance, however the algorithm chooses to interpret this
    #Set convergence criterion to sqrt eps
    tol = √eps(real(one(eltype(A))))

    @info("svds (eigs on [0 A; A' 0])")
    #See #12890 :(
    #b_svds = @benchmark svds(A, v0=q, nsv=nv, tol=tol, maxiter=maxiter)
    b_svds = @benchmark svds(A, nsv=nv, tol=tol, maxiter=maxiter)

    @info("eigs on A'A or AA'")
    makeata(A) = (size(A, 1)≥ size(A,2) ? A*A' : A'A)
    B = makeata(A)
    b_ata = @benchmark makeata(A)
    b_eigs = if m≥n
        B = A'A #Size n x n
        b_eigs = @benchmark eigs(B, v0=q, nev=nv, tol=tol, maxiter=maxiter)
    else
        B = A*A' #Size m x m
        #eigs cannot use the same starting vector as everyone else, needs R^m not R^n
        b_eigs = @benchmark eigs(B, v0=qm, nev=nv, tol=tol, maxiter=maxiter)
    end

    @info("tsvd (PROPACK)")
    b_tsvd = try
        if m≥n #PROPACK implicitly assumes the input is tall
            @benchmark tsvdvals(A, initvec=qm, kmax=maxiter, k=nv, tolin=tol)
        else
            @benchmark tsvdvals(A', initvec=q, kmax=maxiter, k=nv, tolin=tol)
        end
    catch exc
        println("Exception: $exc")
    end

    @info("GKL with thick restart using Ritz values")
    b_tr = try
        @benchmark svdl(A, nv, v0=q, tol=tol, reltol=tol)
    catch exc
        println("Exception: $exc")
    end

    @info("GKL with thick restart using harmonic Ritz values")
        b_trh = try
            @benchmark svdl(A, nv, v0=q, tol=tol, reltol=tol, method=:harmonic)
        catch exc
            println("Exception: $exc")
        end

    data = Dict(
        "svds" => b_svds,
        "ata"  => b_ata,
        "eigs" => b_eigs,
        "tsvd" => b_tsvd,
        "tr"   => b_tr,
        "trh"  => b_trh
    )
    println("Timings:")
    for (k, v) in data
        println(k)
        display(v)
        println()
    end
    JLD.save(benchmarkfilename, data, compress=true)
end

benchmarkall(BASEDIR)


@testset "zygote_adjoints" begin

    rng = MersenneTwister(123456)
    x = rand(rng, 5)
    y = rand(rng, 5)
    r = rand(rng, 5)

    gzeucl = gradient(:Zygote, [x,y]) do xy
        evaluate(Euclidean(), xy[1], xy[2])
    end
    gzsqeucl = gradient(:Zygote, [x,y]) do xy
        evaluate(SqEuclidean(), xy[1], xy[2])
    end
    gzdotprod = gradient(:Zygote, [x,y]) do xy
        evaluate(KernelFunctions.DotProduct(), xy[1], xy[2])
    end
    gzdelta = gradient(:Zygote, [x,y]) do xy
        evaluate(KernelFunctions.Delta(), xy[1], xy[2])
    end
    gzsinus = gradient(:Zygote, [x,y]) do xy
        evaluate(KernelFunctions.Sinus(r), xy[1], xy[2])
    end

    gfeucl = gradient(:FiniteDiff, [x,y]) do xy
        evaluate(Euclidean(), xy[1], xy[2])
    end
    gfsqeucl = gradient(:FiniteDiff, [x,y]) do xy
        evaluate(SqEuclidean(), xy[1], xy[2])
    end
    gfdotprod = gradient(:FiniteDiff, [x,y]) do xy
        evaluate(KernelFunctions.DotProduct(), xy[1], xy[2])
    end
    gfdelta = gradient(:FiniteDiff, [x,y]) do xy
        evaluate(KernelFunctions.Delta(), xy[1], xy[2])
    end
    gfsinus = gradient(:FiniteDiff, [x,y]) do xy
        evaluate(KernelFunctions.Sinus(r), xy[1], xy[2])
    end

    @test all(gzeucl .≈ gfeucl)
    @test all(gzsqeucl .≈ gfsqeucl)
    @test all(gzdotprod .≈ gfdotprod)
    @test all(gzdelta .≈ gfdelta)
    @test all(gzsinus .≈ gfsinus)
end


using BenchmarkTools, DynamicalSystemsBase

const SUITE = BenchmarkGroup(["DynamicalSystemsBase"])

include("reconstructions.jl")
include("integrators.jl")


using Plots
x=-3:0.1:3
f(x)=x^2
y=f.(x)

gr()
plot(x,y,label="line")
scatter!(x,y,label="points")


export TracyWidom

struct TracyWidom <: ContinuousUnivariateDistribution
    beta::Integer
end

function cdf(d::TracyWidom, x::Real)
    cdf(RandomMatrices.TracyWidom, x, beta=d.beta)
end

_TWfs = Dict()

function _TWf(x, beta)
    if !(beta in keys(_TWfs))
        dom = Chebyshev(-5..5)
        D = Derivative(dom)
        _TWfs[beta] = D * Fun(x -> cdf(RandomMatrices.TracyWidom, x, beta=beta), dom)
    end
    _TWfs[beta](x)
end

function pdf(d::TracyWidom, x::Real)
    _TWf(x, d.beta)
end

_TWinvs = Dict()

function _TWinv(t, beta)
    (0 < t < 1) || throw(DomainError(t))

    if !(beta in keys(_TWinvs))
        # bounds outside which we use limiting representation
        u = 4
        l = -5

        # TODO: make a more principled grid using tail behaviour
        n = 1000
        y = range(l, u, length=n)
        x = cdf(RandomMatrices.TracyWidom, y, beta=1)

        # Transform limiting ends to maintain differentiability
        righta = pdf(TracyWidom(beta), u)/(beta * sqrt(u) * exp(-beta * u^(3/2)*2/3))
        rightb = x[end] - righta * (1 - exp(-beta * u^(3/2) * 2/3))

        lefta = pdf(TracyWidom(beta), l)/(beta * l^2/8 * exp(beta * l^3/24))
        leftb = x[1] - lefta * exp(-beta * abs(l)^3/24)

        itp = interpolate((x,), y, Gridded(Linear()))

        _TWinvs[beta] = [righta, rightb, x[end], lefta, leftb, x[1], itp]
    end

    # TODO: make the indexing less terrible
    if t > _TWinvs[beta][3]
        s = (t - _TWinvs[beta][2])/_TWinvs[beta][1]
        return (log(1/(1 - s)) * 3/(2 * beta))^(2/3)
    elseif t < _TWinvs[beta][6]
        s = (t - _TWinvs[beta][5])/_TWinvs[beta][4]
        return -(log(1/s) * 24/beta)^(1/3)
    else
        return _TWinvs[beta][end](t)
    end

end

function quantile(d::TracyWidom, q::Real)
    _TWinv(q, d.beta)
end


## Loading packages.
using MicroHH
using Tullio
using Statistics
using PyPlot
pygui(:qt5)

## Solve the Taylor Green vortex
function solve_taylorgreen(itot)
    # Loading settings.
    d_grid = Dict{String, Any}()
    d_grid["itot"] = itot
    d_grid["jtot"] = 1
    d_grid["ktot"] = itot ÷ 2

    d_grid["xsize"] = 1.
    d_grid["ysize"] = 1.
    d_grid["zsize"] = 0.5

    zsize = d_grid["zsize"]; ktot = d_grid["ktot"]
    dz = zsize / ktot
    d_grid["z"] = range(0.5*dz, step=dz, length=ktot) |> collect

    d_fields = Dict{String, Any}()
    d_fields["visc"] = (8 * pi^2 * 100)^(-1)
    d_fields["alpha"] = 0.

    d_boundary = Dict{String, Any}()
    d_boundary["mom_bot_type"] = "Neumann"
    d_boundary["mom_top_type"] = "Neumann"
    d_boundary["s_bot_type"] = "Neumann"
    d_boundary["s_top_type"] = "Neumann"

    d_timeloop = Dict{String, Any}()
    d_timeloop["start_time"] = 0.
    d_timeloop["end_time"] = 1.
    d_timeloop["save_time"] = 100.
    d_timeloop["check_time"] = 0.1
    d_timeloop["dt"] = 0.005 * (128 / itot)

    d_multidomain = Dict{String, Any}()
    d_multidomain["enable_nudge"] = false

    settings = [ Dict(
        "grid" => d_grid,
        "fields" => d_fields,
        "boundary" => d_boundary,
        "timeloop" => d_timeloop,
        "multidomain" => d_multidomain) ]

    # Initialize the model.
    n_domains = 1
    m = Model("taylorgreen", n_domains, settings, Float64)

    # Create the initials fields.
    f = m.fields[1]; g = m.grid[1]
    u = @view f.u[g.is:g.ie, g.js:g.je, g.ks:g.ke]
    w = @view f.w[g.is:g.ie, g.js:g.je, g.ks:g.keh]
    p = @view f.p[g.is:g.ie, g.js:g.je, g.ks:g.ke]
    x = @view g.x[g.is:g.ie]
    xh = @view g.xh[g.is:g.ie]
    z = @view g.z[g.ks:g.ke]
    zh = @view g.zh[g.ks:g.keh]
    @tullio u[i, j, k] = sin(2pi*xh[i])*cos(2pi*z[k])
    @tullio w[i, j, k] = -cos(2pi*x[i])*sin(2pi*zh[k])

    f_t = exp(-8. * pi^2 * d_fields["visc"] * d_timeloop["end_time"])
    u_ref = u[:, :, :] * f_t
    w_ref = w[:, :, :] * f_t
    p_ref = zeros(size(u_ref))
    @tullio p_ref[i, j, k] = (1/4 * (cos(4pi*x[i]) + cos(4pi*z[k])) - 1/4) * f_t^2

    # Run the model
    in_progress = prepare_model!(m)
    while in_progress
        in_progress = step_model!(m)
    end

    dx = 1. / itot
    u_err = sum(dx^2 * abs.(u .- u_ref))
    w_err = sum(dx^2 * abs.(w .- w_ref))
    p_err = sum(dx^2 * abs.(p .- p_ref))

    return u_err, w_err, p_err
end

## Solve the case at multiple resolutions.
u016, w016, p016 = solve_taylorgreen(16)
u032, w032, p032 = solve_taylorgreen(32)
u064, w064, p064 = solve_taylorgreen(64)
u128, w128, p128 = solve_taylorgreen(128)
u256, w256, p256 = solve_taylorgreen(256)

dxs = [ 1/16, 1/32, 1/64, 1/128, 1/256 ]
u_errs = [ u016, u032, u064, u128, u256 ]
w_errs = [ w016, w032, w064, w128, w256 ]
p_errs = [ p016, p032, p064, p128, p256 ]

## Plot the results
figure()
loglog(dxs, u_errs, "C0o-", label="u")
loglog(dxs, w_errs, "C1o-", label="w")
loglog(dxs, p_errs, "C2o-", label="p")
loglog(dxs, 0.6*dxs.^2, "k:", label="2nd")
legend(loc=0, frameon=false)
xlabel(L"$\Delta$ x")
ylabel("Error")
display(gcf())
show()


module Forecast

using CSV, Distributions, ColorSchemes, DataFrames, DataFramesMeta, Dates, CodecZlib,
    HypothesisTests, LinearAlgebra, Optim, Plots, PrettyTables, RecipesBase, Statistics, StatsBase

# types
export AR, CCF, STL

# methods
export acf, ar, arsim, boxcox, ccf, d, hma, iboxcox, loess, p, pacf, forecast,
       setnames!, sma, stl, transform, summarize

# datasets
export air, co2, london, quakes, seaborne

# types
include("AR.jl")
include("CCF.jl")
include("FORECAST.jl")
include("STL.jl")
include("SUMMARIZE.jl")

# source files
include("acf.jl")
include("ar.jl")
include("arsim.jl")
include("boxcox.jl")
include("ccf.jl")
include("d.jl")
include("datasets.jl")
include("hma.jl")
include("loess.jl")
include("p.jl")
include("pacf.jl")
include("forecast_ar.jl")
include("sma.jl")
include("stl.jl")
include("summarize.jl")
include("utils.jl")
include("utils_datetime.jl")

# plot recipes
##  type
include("plot_CCF.jl")
include("plot_DataFrame.jl")
include("plot_FORECAST.jl")
include("plot_STL.jl")

## user
include("splot.jl")

"""
Collection of methods for Time Series analysis

Featured Methods:

    acf:        Auto-correlation or auto-covariance of univariate series.
    ar:         Multivariate Autoregressive Model.
    arsim:      Simulated Multivariate Autoregressive Model.
    boxcox:     Box-Cox Transformations.
    ccf:        Cross-correlation or cros-covariance of two univariate series.
    d:          Lagged differences of a given order for Vector and Array.
    forecast:   Forecast values of fitted time series models.
    hma:        Henderson moving average filter.
    iboxcox:    Inverse Box-Cox Transformations.
    loess:      Locally estimated scatterplot smoothing.
    p:          Reverse lagged differences of a given order for types Vector and Array.
    pacf:       Partial Auto-correlation function.
    sma:        Simple moving average.
    splot:      Plot a seasonal plot for types Vector and TimeArray.
    stl:        Seasonal and Trend decomposition using loess.
    summarize:  Statistical summary.
"""
Forecast

end


if length(ARGS) >= 1
  nw = parse(ARGS[1])
  if nw >1
     addprocs(nw)
  end
else
  addprocs()
end

import Distributions; @everywhere using Distributions
import ABC; @everywhere using ABC

# Set Prior for Population Parameters
@everywhere theta_true = [0.0, 1.0]
@everywhere param_prior = Distributions.MvNormal(theta_true,ones(length(theta_true)))

# Code to generate simulated data given array of model parameters
@everywhere num_data_default = 100
@everywhere gen_data_normal(theta::Array, n::Integer = num_data_default) = rand(Distributions.Normal(theta[1],theta[2]),num_data_default)

# Function to adjust originally proposed model parameters, so that they will be valid
@everywhere normalize_theta2_pos(theta::Array) =  theta[2] = abs(theta[2])

# Function to test if the proposed model parameter are valid
@everywhere is_valid_theta2_pos(theta::Array) =  theta[2]>0.0 ? true : false

# Tell ABC what it needs to know for a simulation
@everywhere in_parallel = length(workers()) > 1 ? true : false
@everywhere abc_plan = ABC.abc_pmc_plan_type(gen_data_normal,ABC.calc_summary_stats_mean_var,ABC.calc_dist_max, param_prior; is_valid=is_valid_theta2_pos,num_max_attempt=10000, in_parallel=in_parallel);

# Generate "true/observed data" and summary statistics
@everywhere data_true = abc_plan.gen_data(theta_true)
@everywhere ss_true = abc_plan.calc_summary_stats(data_true)
#println("theta= ",theta_true," ss= ",ss_true, " d= ", 0.)

# Run ABC simulation
@time pop_out = run_abc(abc_plan,ss_true;verbose=true,in_parallel=in_parallel);

#=
# Optional plotting of results
using PyPlot

hist(pop_out.weights*length(pop_out.weights));
hist(pop_out.dist);

num_param = 2
num_grid_x = 100
num_grid_y = 100
limit = 1.0
x = collect(linspace(theta_true[1]-limit,theta_true[1]+limit,num_grid_x));
y = collect(linspace(theta_true[2]-limit,theta_true[2]+limit,num_grid_y));
z = zeros(Float64,(num_param,length(x),length(y)))
for i in 1:length(x), j in 1:length(y)
    z[1,i,j] = x[i]
    z[2,i,j] = y[j]
end
z = reshape(z,(num_param,length(x)*length(y)))
zz = [ ABC.pdf(ABC.GaussianMixtureModelCommonCovar(pop_out.theta,pop_out.weights,ABC.cov_weighted(pop_out.theta',pop_out.weights)),vec(z[:,i])) for i in 1:size(z,2) ]
zz = reshape(zz ,(length(x),length(y)));
levels = [exp(-0.5*i^2)/sqrt(2pi^num_param) for i in 0:5];
PyPlot.contour(x,y,zz',levels);
plot(pop_out.theta[1,:],pop_out.theta[2,:],".");
=#


using GalacticOptim, GalacticOptimJL, GalacticFlux, Test
using ForwardDiff, Zygote, ReverseDiff, FiniteDiff, Tracker
using ModelingToolkit
x0 = zeros(2)
rosenbrock(x, p=nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
l1 = rosenbrock(x0)

function g!(G, x)
    G[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
    G[2] = 200.0 * (x[2] - x[1]^2)
end

function h!(H, x)
    H[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
    H[1, 2] = -400.0 * x[1]
    H[2, 1] = -400.0 * x[1]
    H[2, 2] = 200.0
end

G1 = Array{Float64}(undef, 2)
G2 = Array{Float64}(undef, 2)
H1 = Array{Float64}(undef, 2, 2)
H2 = Array{Float64}(undef, 2, 2)

g!(G1, x0)
h!(H1, x0)

cons = (x, p) -> [x[1]^2 + x[2]^2]
optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoModelingToolkit(), cons=cons)
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoModelingToolkit(), nothing, 1)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2
@test optprob.cons(x0) == [0.0]
J = Array{Float64}(undef, 2)
optprob.cons_j(J, [5.0, 3.0])
@test J == [10.0, 6.0]
H3 = [Array{Float64}(undef, 2, 2)]
optprob.cons_h(H3, x0)
@test H3 == [[2.0 0.0; 0.0 2.0]]

function con2_c(x, p)
    [x[1]^2 + x[2]^2, x[2] * sin(x[1]) - x[1]]
end
optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoModelingToolkit(), cons=con2_c)
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoModelingToolkit(), nothing, 2)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2
@test optprob.cons(x0) == [0.0, 0.0]
J = Array{Float64}(undef, 2, 2)
optprob.cons_j(J, [5.0, 3.0])
@test all(isapprox(J, [10.0 6.0; -0.149013 -0.958924]; rtol=1e-3))
H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
optprob.cons_h(H3, x0)
@test H3 == [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoModelingToolkit(true, true), cons=con2_c)
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoModelingToolkit(true, true), nothing, 2)
using SparseArrays
sH = sparse([1, 1, 2, 2], [1, 2, 1, 2], zeros(4))
optprob.hess(sH, x0)
@test sH == H2
@test optprob.cons(x0) == [0.0, 0.0]
sJ = sparse([1, 1, 2, 2], [1, 2, 1, 2], zeros(4))
optprob.cons_j(sJ, [5.0, 3.0])
@test all(isapprox(sJ, [10.0 6.0; -0.149013 -0.958924]; rtol=1e-3))
sH3 = [sparse([1,2], [1, 2], zeros(2)), sparse([1, 1, 2], [1, 2, 1], zeros(3))]
optprob.cons_h(sH3, x0)
@test Array.(sH3) == [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff())
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoForwardDiff(), nothing)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, Optim.BFGS())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.Newton())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10 * sol.minimum < l1

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoZygote())
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoZygote(), nothing)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, Optim.BFGS())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.Newton())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10 * sol.minimum < l1

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoReverseDiff())
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoReverseDiff(), nothing)
optprob.grad(G2, x0)
@test G1 == G2
optprob.hess(H2, x0)
@test H1 == H2

prob = OptimizationProblem(optprob, x0)
sol = solve(prob, Optim.BFGS())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.Newton())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test 10 * sol.minimum < l1

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoTracker())
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoTracker(), nothing)
optprob.grad(G2, x0)
@test G1 == G2
@test_throws ErrorException optprob.hess(H2, x0)

prob = OptimizationProblem(optprob, x0)

sol = solve(prob, Optim.BFGS())
@test 10 * sol.minimum < l1

@test_throws ErrorException solve(prob, Newton())

optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoFiniteDiff())
optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoFiniteDiff(), nothing)
optprob.grad(G2, x0)
@test G1 ≈ G2 rtol = 1e-6
optprob.hess(H2, x0)
@test H1 ≈ H2 rtol = 1e-6

prob = OptimizationProblem(optprob, x0)
sol = solve(prob, Optim.BFGS())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.Newton())
@test 10 * sol.minimum < l1

sol = solve(prob, Optim.KrylovTrustRegion())
@test sol.minimum < l1 #the loss doesn't go below 5e-1 here

sol = solve(prob, Flux.ADAM(0.1), maxiters=1000)
@test 10 * sol.minimum < l1


using DelimitedFiles
using Printf
using Plots
using Statistics

"""
    imtz(t; clims=(-0.5,0.5), cmap=:pu_or, title="", anom=false, L=2*pi)

Plots a vertical rbsolve section at step `t`.

# Optional keywords
- `clims::Tuple{Real,Real}`: limits for the colormap
- `cmap::Symbol`: the colormap to use
- `title::String`: title of the plot
- `anom::Bool`: compute horizontal anomaly
- `L::Real`: X domain size
"""
function imtz(t; clims=(-0.5,0.5), cmap=:pu_or, title="", anom=false, L=2*pi)
    coord=readdlm("coord")
    (nz,nc)=size(coord)
    nz=convert(Int64,(nz-1)/2)+1
    coord=coord[1:nz,4]
    nz=nz-2
    filename=@sprintf("t%07d.zzz", t)
    tt=readdlm(filename)
    n=length(tt); n=n÷nz
    zz=coord[2:(nz+1)]
    xx=(0:(n-1))/n*L
    tt=reshape(tt, n, nz)
    if anom
       tt1=tt.-ones(n,1)*mean(tt,dims=1)
    else
       tt1=tt
    end
    plot(xx, zz.+0.5, tt1', st=:heatmap, c=cmap, clims=clims)
    xlabel!("x")
    ylabel!("z")
    title!(title)
end


# arules_testing.jl

using RCall
using DataFrames
using StatsBase                        # for sample() function

include("arules.jl")
include("utils.jl")

init_arules()                       # this is called automatically by `using AssociationRules`

# simulate transactions
groceries = ["milk", "bread", "eggs", "apples", "oranges", "beer"]
transacts = [sample(groceries, 4, replace = false) for x in 1:1000]
transacts_df = DataFrame(transactions(transacts))

# minimum support of 0.2, minimum confidence of 0.8
rules = apriori(transacts_df, 0.2, 0.8)

adult = dataset("adult")            # example dataset included in package
sanitize_input!(adult)              # remove excluded character patterns (e.g., " => ")

# minimum support of 0.1, minimum confidence of 0.8
rules = apriori(adult, 0.1, 0.8)


function test_sparsity(x, target_sparsity; atol = 0.05)
    number_of_zeros = count(ele -> (ele == 0.0), x)
    sparsity = number_of_zeros / sum(length, x)
    @test isapprox(target_sparsity, sparsity; atol)
end

@testset "MCDense GPU" begin
    dropout_rate = 0.35
    # Test MC Dense layer
    a = gpu(rand(Float32, 8, 32))
    layer = gpu(MCDense(8, 16, dropout_rate))
    output = layer(a)
    # Test if it's CuArray
    @test output isa CuArray
    @test isequal(size(output), (16, 32))
    test_sparsity(output, dropout_rate)
    # Test MC dense dropout toggle
    output = layer(a, dropout = false)
    test_sparsity(output, 0)
end

@testset "MCConv GPU" begin
    dropout_rate = 0.4
    # Test MC conv layer
    a = gpu(rand(Float32, 32, 32, 3, 32))
    layer = MCConv((5, 5), 3 => 6, dropout_rate) |> gpu
    output = layer(a)
    # Test if it's CuArray
    @test output isa CuArray
    # Test the output shape
    @test isequal(size(output), (28, 28, 6, 32))
    # Test the sparsity percentage in the array
    test_sparsity(output, dropout_rate)
    # Test MC conv dropout toggle
    output = layer(a, dropout = false)
    test_sparsity(output, 0)
end


#Demo 1: Simple sine wave plot
#-------------------------------------------------------------------------------
using InspectDR
using Colors
import NumericIO: UEXPONENT_SI

#==Input
===============================================================================#
#Constants
#-------------------------------------------------------------------------------
black = RGB24(0, 0, 0)
white = RGB24(1, 1, 1)
red = RGB24(1, 0, 0)
green = RGB24(0, 1, 0)
blue = RGB24(0, 0, 1)

#Input data
#-------------------------------------------------------------------------------
tcycle = 1e-9
npts = 2000 #Relatively large dataset
ncycles = 4
tmax = tcycle*ncycles
Δ = tmax/(npts-1)
t = collect(0:Δ:tmax)
y = sin(t*(ncycles*2pi/tmax))

t[end] = t[1] #x-values are no longer ordered

#Lower resolution t-vector:
t_lres = collect(0:(tmax/10):tmax)

#==Generate plot
===============================================================================#
plot = InspectDR.Plot2D(title="Sample Plot (λ)")
plot.layout = InspectDR.Layout(fontname="monospace", fontscale=1.5)
plot.layout.grid = grid(vmajor=true, vminor=true, hmajor=true)
plot.layout.legend.enabled = true
plot.layout.legend.width = 150
plot.layout.showtimestamp = true
plot.layout.xlabelformat.expdisplay = UEXPONENT_SI #Use SI notation on x-axis

style = :dashdot #solid/dashdot/...
wfrm = add(plot, t, y+1, id="sin(2πt)+1")
	wfrm.line = line(color=blue, width=1, style=style)
wfrm = add(plot, t, y-1, id="sin(2πt)-1")
	wfrm.line = line(color=red, width=5, style=style)
wfrm = add(plot, t_lres, -2+4*(t_lres./tmax), id="-2+4t/tmax")
	wfrm.line = line(color=blue, width=3, style=:dash)
	wfrm.glyph = glyph(shape=:*, size=10)
#= Supported shapes:
	:square, :diamond,
	:uarrow, :darrow, :larrow, :rarrow, #usually triangles
	:cross, :+, :diagcross, :x,
	:circle, :o, :star, :*,
=#

a = plot.annotation
a.xlabel = "Time (s)"
a.ylabel = "Signal Voltage (V)"

gplot = display(InspectDR.GtkDisplay(), plot)

InspectDR.write_png("export_plot.png", plot)
InspectDR.write_svg("export_plot.svg", plot)
InspectDR.write_eps("export_plot.eps", plot)
InspectDR.write_pdf("export_plot.pdf", plot)

:DONE


"""
    triangulate(points::AbstractArray{Float64, 2})

Triangulate a set of vertices in N dimensions. `points` is an array of vertices, where
each row of the array is a point.
"""
function triangulate(points::AbstractArray{Float64, 2})
    indices = Simplices.Delaunay.delaunayn(points)
    return indices
end

"""
    Triangulation(points::AbstractArray{Float64, 2}, impoints::AbstractArray{Float64, 2},
                simplex_inds::AbstractArray{Int, 2})
"""
@with_kw mutable struct Triangulation
    # The vertices of the triangulation
    points::AbstractArray{Float64, 2} = zeros(Float64, 0, 0)

    # The image vertices of the triangulation
    impoints::AbstractArray{Float64, 2} = zeros(Float64, 0, 0)

    # Array of indices referencing the vertices furnishing each simplex
    simplex_inds::AbstractArray{Int, 2} = zeros(Float64, 0, 0)

    # Some properties of the simplices furnishing the triangulation
    centroids::AbstractArray{Float64, 2} = zeros(Float64, 0, 0)
    radii::AbstractVector{Float64} = Float64[]
    centroids_im::AbstractArray{Float64, 2}  = zeros(Float64, 0, 0)
    radii_im::AbstractVector{Float64} = Float64[]
    orientations::AbstractVector{Float64} = Float64[]
    orientations_im::AbstractVector{Float64} = Float64[]
    volumes::AbstractVector{Float64} = Float64[]
    volumes_im::AbstractVector{Float64} = Float64[]
end

todict(t::Triangulation) = Dict([fn => getfield(t, fn) for fn = fieldnames(t)])
export todict

"""
    invariant_triangulation(embedding::Array{Float64, 2})

Create a `Triangulation` from an invariant `embedding`. Note: this function does not check
for invariance. This must have been done beforehand.
"""
function triang_from_embedding(e::SimplexSplitting.Embedding)
    points = e.embedding[1:end-1, :]
    simplex_inds = triangulate(points)
    impoints = e.embedding[2:end, :]
    c, r = SimplexSplitting.centroids_radii2(points, simplex_inds)
    cim, rim = SimplexSplitting.centroids_radii2(impoints, simplex_inds)
    vol = SimplexSplitting.simplex_volumes(points, simplex_inds)
    volim = SimplexSplitting.simplex_volumes(impoints, simplex_inds)
    orientations = SimplexSplitting.orientations(points, simplex_inds)
    orientations_im = SimplexSplitting.orientations(impoints, simplex_inds)

    Triangulation(
        points = points,
        impoints = impoints,
        simplex_inds = simplex_inds,
        centroids = c,
        radii = r,
        volumes = vol,
        centroids_im = cim,
        radii_im = rim,
        volumes_im = volim,
        orientations = orientations,
        orientations_im = orientations_im)
end

"""
    invariant_triangulation(embedding::Array{Float64, 2})

Create a `Triangulation` from an invariant `embedding`. Note: this function does not check
for invariance. This must have been done beforehand.
"""
function example_triangulation(n_simplices::Int)
    embedding = gaussian_embedding(n_simplices).embedding
    points = embedding[1:end-1, :]
    impoints = embedding[2:end, :]

    simplex_inds = triangulate(points)
    c, r = SimplexSplitting.centroids_radii2(points, simplex_inds)
    cim, rim = SimplexSplitting.centroids_radii2(impoints, simplex_inds)
    vol = SimplexSplitting.simplex_volumes(points, simplex_inds)
    volim = SimplexSplitting.simplex_volumes(impoints, simplex_inds)
    Triangulation(points, impoints, simplex_inds, c, r, vol, cim, rim, volim)
end

struct Simplex
    v::AbstractVector{AbstractVector{Float64}}
end

"""
Find the simplex with index i
"""
function find_simplex(t::Triangulation, i::Int)
    s = Vector{Vector{Float64}}(4)
    n_vertices = size(t.simplex_inds, 2)
    for j in 1:n_vertices
        s[j] = t.points[t.simplex_inds[i, j], :]
    end
    return Simplex(s)
end

"""
Find the image simplex with index i
"""
function find_imsimplex(t::Triangulation, i::Int)
    s = Vector{Vector{Float64}}(4)
    n_vertices = size(t.simplex_inds, 2)
    for j in 1:n_vertices
        s[j] = t.impoints[t.simplex_inds[i, j], :]
    end
    return Simplex(s)
end

function get_simplices(t::Triangulation)
    n_simplices = size(t.simplex_inds, 1)
    simplices = Vector{Simplex}(n_simplices)
    for i = 1:n_simplices
        simplices[i] = find_simplex(t, i)
    end
    simplices
end

function get_imagesimplices(t::Triangulation)
    n_simplices = size(t.simplex_inds, 1)
    simplices = Vector{Simplex}(n_simplices)
    for i = 1:n_simplices
        simplices[i] = find_imsimplex(t, i)
    end
    simplices
end

function newpoint!(pt::AbstractArray{Float64, 1}, s::Simplex, coeffs)
    pt .= 0.0
    for i in 1:length(s.v)
        pt .= pt .+ coeffs[i] * s.v[i]
    end
    pt
end

"""
Find the indices of the simplices in the original triangulation that potentially
intersect with the image simplex with index `image_i`.
"""
function maybeintersecting_simplices(t::Triangulation, image_i::Int)
    inds_potential_simplices = Int[]

    n_simplices = length(t.radii)

    @inbounds for i = 1:n_simplices
        dist_difference = (transpose(t.centroids_im[image_i] - t.centroids[i]) *
                            (t.centroids_im[image_i] - t.centroids[i]) - (t.radii_im[image_i] + t.radii[i])^2)[1]
        if dist_difference < 0
            push!(inds_potential_simplices, i)
        end
    end
    return inds_potential_simplices
end

"""
Find the indices of the image simplices in `t` that potentially intersect with
the original simplex with index `orig_i`.
"""
function maybeintersecting_imsimplices(t::Triangulation, orig_i::Int)
    inds_potential_simplices = Int[]

    n_simplices = length(t.radii)

    @inbounds for i = 1:n_simplices
        dist_difference = (transpose(t.centroids[orig_i] - t.centroids_im[i]) *
                            (t.centroids[orig_i] - t.centroids_im[i]) - (t.radii[orig_i] + t.radii_im[i])^2)[1]
        if dist_difference < 0
            push!(inds_potential_simplices, i)
        end
    end
    return inds_potential_simplices
end


@generated function Base.zero{U<:AbstractUnum}(::Type{U})
  c = unumConstants(U)
  :($(c.zero))
end

Base.one{U<:AbstractUnum}(::Type{U}) = convert(U, 1)

"exponent of the unum"
@generated function Base.exponent{U<:AbstractUnum}(u::U)
  c = unumConstants(U)
  :(reinterpret($(c.INT), u & $(c.emask)) >> $(c.fpos))
end

"significand (fraction) of an unum"
@generated function Base.significand{U<:AbstractUnum}(u::U)
  c = unumConstants(U)
  :(reinterpret($(c.INT), u & $(c.fmask)) >> $(c.ubitpos))
end

"Number of bits in the unum's exponent"
@generated function esize{U<:AbstractUnum}(u::U)
  c = unumConstants(U)
  :(1 + reinterpret($(c.INT), u & $(c.esizemask)) >> $(c.fsizepos))
end

"Number of bits in the unum's significand (fraction)"
@generated function fsize{U<:AbstractUnum}(u::U)
  c = unumConstants(U)
  :(1 + reinterpret($(c.INT), u & $(c.fsizemask)))
end

@generated function isnegative{U<:AbstractUnum}(u::U)
  c = unumConstants(U)
  :(reinterpret($(c.UINT), u & $(c.signbitmask)) != zero($(c.UINT)))
end

ispositive(u::AbstractUnum) = !isnegative(u)

@generated function isinexact{U<:AbstractUnum}(u::U)
  c = unumConstants(U)
  :(reinterpret($(c.UINT), u & $(c.ubitmask)) != zero($(c.UINT)))
end

isexact(u::AbstractUnum) = !isinexact(u)

@generated function typemin{U<:AbstractUnum}(::Type{U})
  c = unumConstants(U)
  :(c.neginf)
end

@generated function typemax{U<:AbstractUnum}(::Type{U})
  c = unumConstants(U)
  :(c.posinf)
end

@generated function maxubits{U<:AbstractUnum}(::Type{U})
  c = unumConstants(U)
  :(2 + $(c.esizesize) + $(c.fsizesize) + 2^$(c.esizesize) + 2^$(c.fsizesize))
end

# for func in [:isnan, :isnull, :isinf, :isfinite, :isposinf, :isneginf,
#               :isnegative, :iszero, :ispositive, :-, :isinteger, :issubnormal]

# -----------------------------------------------------------------------------

# @generated function +{U<:AbstractUnum}(u1::U, u2::U)

# end

# for op in [:+, :-, :*, :/, :(==), :<, :>, :<=, :>=, :isless, :isequal]


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
Random.seed!(233)

function acoustic_one_step(w,wold,phi,psi,sigma,tau,c,dt,hx,hy,nx,ny)
    acoustic_one_step_ = load_op_and_grad("../build/libADSeismic.so","acoustic_one_step_cpu", multiple=true)
    w,wold,phi,psi,sigma,tau,c,dt,hx,hy,nx,ny = convert_to_tensor(Any[w,wold,phi,psi,sigma,tau,c,dt,hx,hy,nx,ny], [Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64,Int64,Int64])
    acoustic_one_step_(w,wold,phi,psi,sigma,tau,c,dt,hx,hy,nx,ny)
end

# TODO: specify your input parameters
nx = 10
ny = 10
mask = ones(nx+2, ny+2)
mask = mask[:]

using Random; Random.seed!(1234)
hx = constant(0.1)
hy = constant(0.1)
dt = constant(0.1)
c = rand((nx+2)*(ny+2)) .* mask
phi = rand((nx+2)*(ny+2)) .* mask
psi = rand((nx+2)*(ny+2)) .* mask
wold = rand((nx+2)*(ny+2)) .* mask
w = rand((nx+2)*(ny+2)) .* mask
sigma = rand((nx+2)*(ny+2)) .* mask
tau = rand((nx+2)*(ny+2)) .* mask
u = acoustic_one_step(w,wold,phi,psi,sigma,tau,c,dt,hx,hy,nx,ny)
sess = Session(); init(sess)
@show run(sess, u)

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(x)
    return sum(sum(acoustic_one_step(w, wold,phi,psi,sigma,tau,c,dt,hx,hy,nx,ny).^2))
end
# w, psi

# psi, w

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand((nx+2)*(ny+2)) .* mask)
v_ = rand((nx+2)*(ny+2)) .* mask
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 0.01 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")
savefig("gradtest")


using ApproxFun, Base.Test
import Compat: view

## Try constructor variants

ff=(x,y)->exp(-10(x+.2)^2-20(y-.1)^2)*cos(x*y)
gg=x->exp(-10(x[1]+.2)^2-20(x[1]-.1)^2)*cos(x[1]*x[2])
f=Fun(ff,Interval()^2,10000)
@test_approx_eq f(0.,0.) ff(0.,0.)

f=Fun(gg,Interval()^2,10000)
@test_approx_eq f(0.,0.) ff(0.,0.)

f=Fun(ff,Interval()^2)
@test_approx_eq f(0.,0.) ff(0.,0.)
f=Fun(gg,Interval()^2)
@test_approx_eq f(0.,0.) ff(0.,0.)

f=Fun(ff)
@test_approx_eq f(0.,0.) ff(0.,0.)
f=Fun(gg)
@test_approx_eq f(0.,0.) ff(0.,0.)

# Fun +-* constant
f=Fun((x,y)->exp(x)*cos(y))

@test_approx_eq f(0.1,0.2)+2 (f+2)(0.1,0.2)
@test_approx_eq f(0.1,0.2)-2 (f-2)(0.1,0.2)
@test_approx_eq f(0.1,0.2)*2 (f*2)(0.1,0.2)

## ProductFun
u0   = ProductFun((x,y)->cos(x)+sin(y) +exp(-50x.^2-40(y-.1).^2)+.5exp(-30(x+.5).^2-40(y+.2).^2))

@test values(u0)-values(u0|>LowRankFun)|>norm < 1000eps()
@test chebyshevtransform(values(u0))-coefficients(u0)|>norm < 100eps()

##TODO: need to do adaptive to get better accuracy
@test sin(u0)(.1,.2)-sin(u0(.1,.2))|>abs < 10e-4

## LowRankFun

F = LowRankFun((x,y)->besselj0(10(y-x)),Chebyshev(),Chebyshev())

@test_approx_eq F(.123,.456) besselj0(10(.456-.123))

G = LowRankFun((x,y)->besselj0(10(y-x));method=:Cholesky)

@test_approx_eq G(.357,.246) besselj0(10(.246-.357))

F = LowRankFun((x,y)->hankelh1(0,10abs(y-x)),Chebyshev([1.0,2.0]),Chebyshev([1.0im,2.0im]))

@test_approx_eq F(1.5,1.5im) hankelh1(0,10abs(1.5im-1.5))

## Periodic
f=LowRankFun((x,y)->cos(x)*sin(y),PeriodicInterval(),PeriodicInterval())
@test_approx_eq f(.1,.2) cos(.1)*sin(.2)

f=LowRankFun((x,y)->cos(cos(x)+sin(y)),PeriodicInterval(),PeriodicInterval())
@test_approx_eq f(.1,.2) cos(cos(.1)+sin(.2))
@test norm(Float64[cos(cos(x)+sin(y)) for x=ApproxFun.vecpoints(f,1),y=ApproxFun.vecpoints(f,2)]-values(f))<10000eps()

f=ProductFun((x,y)->cos(cos(x)+sin(y)),PeriodicInterval()^2)
@test_approx_eq f(.1,.2) cos(cos(.1)+sin(.2))
x,y=points(f)
@test norm(Float64[cos(cos(x[k,j])+sin(y[k,j])) for k=1:size(f,1),j=1:size(f,2)]-values(f))<10000eps()

d=PeriodicInterval()^2
f=ProductFun((x,y)->exp(-10(sin(x/2)^2+sin(y/2)^2)),d)
@test (f.'-f|>coefficients|>norm)< 1000eps()

## Functional*Fun

d=Interval()
B=ldirichlet(d)
f=ProductFun((x,y)->cos(cos(x)*sin(y)),d^2)
@test norm(B*f-Fun(y->cos(cos(-1)*sin(y)),d))<20000eps()
@test norm(f*B-Fun(x->cos(cos(x)*sin(-1)),d))<20000eps()

## 1D in 2D

d=Interval((0.,0.),(1.,1.))
f=Fun(xy->exp(-xy[1]-2cos(xy[2])),d)
@test_approx_eq f(0.5,0.5) exp(-0.5-2cos(0.5))
@test_approx_eq f(FixedSizeArrays.Vec(0.5,0.5)) exp(-0.5-2cos(0.5))

f=Fun(xy->exp(-xy[1]-2cos(xy[2])),d,20)
@test_approx_eq f(0.5,0.5) exp(-0.5-2cos(0.5))

f=Fun((x,y)->exp(-x-2cos(y)),d)
@test_approx_eq f(0.5,0.5) exp(-0.5-2cos(0.5))

f=Fun((x,y)->exp(-x-2cos(y)),d,20)
@test_approx_eq f(0.5,0.5) exp(-0.5-2cos(0.5))

d=Circle((0.,0.),1.)
f=Fun(xy->exp(-xy[1]-2cos(xy[2])),Fourier(d),40)
@test_approx_eq f(cos(0.1),sin(0.1)) exp(-cos(0.1)-2cos(sin(0.1)))
@test_approx_eq f(FixedSizeArrays.Vec(cos(0.1),sin(0.1))) exp(-cos(0.1)-2cos(sin(0.1)))

f=Fun((x,y)->exp(-x-2cos(y)),Fourier(d),40)
@test_approx_eq f(cos(0.1),sin(0.1)) exp(-cos(0.1)-2cos(sin(0.1)))

f=Fun((x,y)->exp(-x-2cos(y)),Fourier(d))
@test_approx_eq f(cos(0.1),sin(0.1)) exp(-cos(0.1)-2cos(sin(0.1)))

println("    Calculus tests")

## Sum

ff=(x,y)->(x-y)^2*exp(-x^2/2.-y^2/2)
f=Fun(ff,[-4.,4.],[-4.,4.])

@test_approx_eq sum(f,1)(0.1) 2.5162377980828357
f=LowRankFun(f)
@test_approx_eq evaluate(f.A,0.1) map(f->f(0.1),f.A)

## Kron operator
Mx=Multiplication(Fun(cos),Chebyshev())
My=Multiplication(Fun(sin),Chebyshev())
K=Mx⊗My
@test_approx_eq ApproxFun.BandedBlockBandedMatrix(view(K,1:10,1:10)) K[1:10,1:10]
C=Conversion(Chebyshev()⊗Chebyshev(),Ultraspherical(1)⊗Ultraspherical(1))
@test_approx_eq C[1:100,1:100] Float64[C[k,j] for k=1:100,j=1:100]

# 2d derivative (issue #346)
let d = Chebyshev()^2
    f = Fun((x,y) -> sin(x) * cos(y), d)
    C=Conversion(Chebyshev()⊗Chebyshev(),Ultraspherical(1)⊗Ultraspherical(1))
    @test_approx_eq (C*f)(0.1,0.2) f(0.1,0.2)
    Dx = Derivative(d, [1,0])
    f = Fun((x,y) -> sin(x) * cos(y), d)
    fx = Fun((x,y) -> cos(x) * cos(y), d)
    @test (Dx*f)(0.2,0.3) ≈ fx(0.2,0.3)
    Dy = Derivative(d, [0,1])
    fy = Fun((x,y) -> -sin(x) * sin(y), d)
    @test (Dy*f)(0.2,0.3) ≈ fy(0.2,0.3)
    L=Dx+Dy
    @test_approx_eq (L*f)(0.2,0.3) (fx(0.2,0.3)+fy(0.2,0.3))

    B=ldirichlet(d[1])⊗ldirichlet(d[2])
    @test_approx_eq B*f f(-1.,-1.)

    B=Evaluation(d[1],0.1)⊗ldirichlet(d[2])
    @test_approx_eq B*f f(0.1,-1.)

    B=Evaluation(d[1],0.1)⊗Evaluation(d[2],0.3)
    @test_approx_eq B*f f(0.1,0.3)

    B=Evaluation(d,(0.1,0.3))
    @test_approx_eq B*f f(0.1,0.3)
end

let d = Space([0,1]) * Space([0,2])
    Dx = Derivative(d, [1,0])
    f = Fun((x,y) -> sin(x) * cos(y), d)
    fx = Fun((x,y) -> cos(x) * cos(y), d)
    @test (Dx*f)(0.2,0.3) ≈ fx(0.2,0.3)
    Dy = Derivative(d, [0,1])
    fy = Fun((x,y) -> -sin(x) * sin(y), d)
    @test (Dy*f)(0.2,0.3) ≈ fy(0.2,0.3)
    L=Dx+Dy
    @test_approx_eq (L*f)(0.2,0.3) (fx(0.2,0.3)+fy(0.2,0.3))

    B=ldirichlet(d[1])⊗ldirichlet(d[2])
    @test abs(B*f-f(0.,0.)) ≤ 10eps()

    B=Evaluation(d[1],0.1)⊗ldirichlet(d[2])
    @test_approx_eq B*f f(0.1,0.)

    B=Evaluation(d[1],0.1)⊗Evaluation(d[2],0.3)
    @test_approx_eq B*f f(0.1,0.3)

    B=Evaluation(d,(0.1,0.3))
    @test_approx_eq B*f f(0.1,0.3)
end


using Test
using LinearAlgebra
using jInv.Mesh;
using FactoredEikonalFastMarching;
using Statistics
include("sensitivityTest.jl");

@testset "Sensitivity test" begin
	runSensitivityTest2D();
	runSensitivityTest3D();
	include("../examples/runExperiments.jl")
end


@testset "Basic Monomial Functionalities" begin

    @testset "automatic sorting" begin
        x = semi_elementary_monomial((3,2,1),(1,0,0))
        @test x.sp_term == (0,1,2)
        x = semi_elementary_monomial((2,1,0),(1,0,0))
        @test x.sp_term == (0,1,2)
    end

    @testset "pull out common factors" begin
        x = semi_elementary_monomial((3,4,5),(0,0,0))
        @test x == semi_elementary_monomial((0,1,2),(0,0,3))
    end

    @testset "test <, 0 terms same" begin
        x = semi_elementary_monomial((3,5,5),(0,0,0))
        y = semi_elementary_monomial((3,4,6),(0,0,0))
        @test (x < y) == true
    end

    @testset "test <, 1 term same" begin
        x = semi_elementary_monomial((3,3,6),(0,0,0))
        y = semi_elementary_monomial((3,4,6),(0,0,0))
        @test (x < y) == true
    end

    @testset "test <, all sp terms same" begin
        x = semi_elementary_monomial((3,3,6),(0,0,0))
        y = semi_elementary_monomial((3,3,6),(1,0,0))
        @test (x < y) == true
    end

    @testset "test <, autosimplifying" begin
        x = semi_elementary_monomial((2,2,2),(1,0,0))
        y = semi_elementary_monomial((3,3,3),(0,0,0))
        @test (x < y) == false
    end

    @testset "test is_elementary" begin
        x = semi_elementary_monomial((3,3,3),(0,0,0))
        @test is_elementary(x) == true
        x = semi_elementary_monomial((1,0,0),(0,0,0))
        @test is_elementary(x) == true
        x = semi_elementary_monomial((0,0,1),(0,0,0))
        @test is_elementary(x) == true
    end

    @testset "test dim" begin
        x = semi_elementary_monomial((3,3,3),(0,0,0))
        @test dim(x) == 3
        x = semi_elementary_monomial((3,3,3,3),(0,0,0,0))
        @test dim(x) == 4
    end

    @testset "string conversion" begin
        x = semi_elementary_monomial((3,3,3),(0,0,0))
        @test to_string(x) == "e3^3 "
        x = semi_elementary_monomial((3,3,0),(0,0,0))
        @test to_string(x) == "S(0, 3, 3)"
        x = semi_elementary_monomial((3,3,1),(1,0,0))
        @test to_string(x) == "S(0, 2, 2)e1 e3 "
    end
end


using RCall
using Distributions
using StatsBase
using IterTools
using Gurobi
using JLD2

using EmpiricalBayes

nreps = 4
k_max = 10 #in total 25*16 replications
ks_par = 1:10#1:10

true_dists = [EmpiricalBayes.easy_dist, EmpiricalBayes.hard_dist]
xs = collect(-3.0:0.5:3.0)
minimax_σs = [0.2]
ms = [10_000]
targets =[ [PosteriorTarget(PosteriorMeanNumerator(x)) for x in xs],
           [PosteriorTarget(LFSRNumerator(x)) for x in xs]]
combs = []

for (true_dist, target, σ, m) in IterTools.product(true_dists, targets, minimax_σs, ms)
    push!(combs, Dict(:dist=>true_dist, :target=>target, :σ=>σ, :m=>m ))
end

comb_reps = []
for (l, k_par) in IterTools.product(1:length(combs), ks_par)
    push!(comb_reps, (l,k_par))
end

j = parse(Int64, ARGS[1])
comb_rep = comb_reps[j]
i = comb_rep[1]
k_rep = comb_rep[2]
comb = combs[i]

m = comb[:m]
σ = comb[:σ]
true_dist = comb[:dist]
target = comb[:target]

marginal_grid = collect(range(-6,stop=6,length=1001));
prior_grid = collect(range(-3,stop=3,length=121));
marginal_h = marginal_grid[2]-marginal_grid[1];

d_true = NormalConvolutionProblem(true_dist, marginal_grid)
ds = MixingNormalConvolutionProblem(Normal, σ, prior_grid, marginal_grid);

for k=1:k_max
    sim_array = []
    for l = 1:nreps
        @show l
        # get the sample
        Xs = rand(d_true,m)
        m_total = length(Xs)
        m_test= ceil(Int, m_total/2)
        m_train = m_total - m_test
        idx_test = sample(1:m_total, m_test, replace=false)
        idx_train = setdiff(1:m_total, idx_test)
        Xs_train = Xs[idx_train]
        Xs_test = Xs[idx_test]

        brad_fit =  fit(SmoothedBradDeconvolveR, Xs, 0.05; c0=0.01,prior_grid = collect(-3.6:0.05:3.6))
        f_nb = fit(BinnedMarginalDensityNeighborhood, Xs_train, ds; nboot=400)

        calibrator_res = []
        for t in target
            @show t
            CB_num = ComteButucea(t.num, m_train, marginal_grid)
            CB_denom = ComteButucea(t.denom, m_train, marginal_grid)
            ceb_res =  CEB_ci(Xs_train, Xs_test,
                        ds,
                        t,
                        CB_num, #
                        CB_denom; f_nb=f_nb)
            push!(calibrator_res, ceb_res)
        end
        push!(sim_array, (calibrator_res, brad_fit))
    end
    res = (i, comb, sim_array)
    @save "/scratch/users/ignat/sims/base_sim_Jan10/mysim_$(i)_$(k)_$(k_rep).jld" res
end


# phase_vocoder.jl

using Sound: phase_vocoder
using Test: @testset, @test, @inferred

@testset "vocoder" begin
    S = 8192
    N = S
    x = cos.(2π*400*(1:N)/S)
    y = @inferred phase_vocoder(x, S)
    @test length(y) == 2N
    y2 = @inferred phase_vocoder([x 2x], S)
    @test y2 ≈ [y 2y]
end


import Distributions
using SpecialFunctions: lgamma, lbeta, digamma

#########################
# abstract distribution #
#########################

abstract type Distribution{T} end

function random end
function logpdf end
get_return_type(::Distribution{T}) where {T} = T

export Distribution
export random
export logpdf
export logpdf_grad

#########
# dirac #
#########

struct Dirac <: Distribution{Float64} end

"""
    dirac(val::Real)

Deterministically return a `Float64` that is equal to the given value.
`logpdf(dirac, x, y)` returns `0` if `x == y` and `-Inf` otherwise.
"""
const dirac = Dirac()

logpdf(::Dirac, x::Real, y::Real) = (x == y ? 0. : -Inf)

random(::Dirac, y::Real) = Float64(y)

(::Dirac)(y) = random(Dirac(), y)

has_output_grad(::Dirac) = false
has_argument_grads(::Dirac) = (false,)
get_static_argument_types(::Dirac) = [Float64]

export dirac

#############
# bernoulli #
#############

struct Bernoulli <: Distribution{Bool} end

"""
    bernoulli(prob_true::Real)

Samples a `Bool` value which is true with given probability
"""
const bernoulli = Bernoulli()

function logpdf(::Bernoulli, x::Bool, prob::Real)
    x ? log(prob) : log(1. - prob)
end

function logpdf_grad(::Bernoulli, x::Bool, prob::Real)
    prob_grad = x ? 1. / prob : -1. / (1-prob)
    (nothing, prob_grad)
end

random(::Bernoulli, prob::Real) = rand() < prob

(::Bernoulli)(prob) = random(Bernoulli(), prob)

has_output_grad(::Bernoulli) = false
has_argument_grads(::Bernoulli) = (true,)
get_static_argument_types(::Bernoulli) = [Float64]

export bernoulli

##########
# normal #
##########

struct Normal <: Distribution{Float64} end

"""
    normal(mu::Real, std::Real)

Samples a `Float64` value from a normal distribution.
"""
const normal = Normal()

function logpdf(::Normal, x::Real, mu::Real, std::Real)
    var = std * std
    diff = x - mu
    -(diff * diff)/ (2.0 * var) - 0.5 * log(2.0 * pi * var)
end

function logpdf_grad(::Normal, x::Real, mu::Real, std::Real)
    precision = 1. / (std * std)
    diff = mu - x
    deriv_x = diff * precision
    deriv_mu = -deriv_x
    deriv_sigma = -0.5 * precision * (1. - precision * (diff * diff))
    (deriv_x, deriv_mu, deriv_sigma)
end

random(::Normal, mu::Real, std::Real) = mu + std * randn()

(::Normal)(mu, std) = random(Normal(), mu, std)

has_output_grad(::Normal) = true
has_argument_grads(::Normal) = (true, true)
get_static_argument_types(::Normal) = [Float64, Float64]

export normal

#########
# gamma #
#########

struct Gamma <: Distribution{Float64} end

"""
    gamma(shape::Real, scale::Real)

Sample a `Float64` from a gamma distribution.
"""
const gamma = Gamma()

function logpdf(::Gamma, x::Real, shape::Real, scale::Real)
    if x > 0.
        (shape - 1.0) * log(x) - (x / scale) - shape * log(scale) - lgamma(shape)
    else
        -Inf
    end
end

function logpdf_grad(::Gamma, x::Real, shape::Real, scale::Real)
    error("Not Implemented")
    (nothing, nothing, nothing)
end

function random(::Gamma, shape::Real, scale::Real)
    rand(Distributions.Gamma(shape, scale))
end

(::Gamma)(shape, scale) = random(Gamma(), shape, scale)

has_output_grad(::Gamma) = false
has_argument_grads(::Gamma) = (false, false)
get_static_argument_types(::Gamma) = [Float64, Float64]

export gamma

#################
# inverse gamma #
#################

struct InverseGamma <: Distribution{Float64} end

"""
    inv_gamma(shape::Real, scale::Real)

Sample a `Float64` from a inverse gamma distribution.
"""
const inv_gamma = InverseGamma()

function logpdf(::InverseGamma, x::Real, shape::Real, scale::Real)
    if x > 0.
        shape * log(scale) - (shape + 1) * log(x) - lgamma(shape) - (scale / x)
    else
        -Inf
    end
end

function logpdf_grad(::InverseGamma, x::Real, shape::Real, scale::Real)
    error("Not Implemented")
    (nothing, nothing, nothing)
end

function random(::InverseGamma, shape::Real, scale::Real)
    rand(Distributions.InverseGamma(shape, scale))
end

(::InverseGamma)(shape, scale) = random(InverseGamma(), shape, scale)

has_output_grad(::InverseGamma) = false
has_argument_grads(::InverseGamma) = (false, false)
get_static_argument_types(::InverseGamma) = [Float64, Float64]

export inv_gamma

########
# beta #
########

struct Beta <: Distribution{Float64} end

"""
    beta(alpha::Real, beta::Real)

Sample a `Float64` from a beta distribution.
"""
const beta = Beta()

function logpdf(::Beta, x::Real, alpha::Real, beta::Real)
    (alpha - 1) * log(x) + (beta - 1) * log1p(-x) - lbeta(alpha, beta)
end

function logpdf_grad(::Beta, x::Real, alpha::Real, beta::Real)
    if 0 <= x <= 1
        deriv_x = (alpha - 1) / x - (beta - 1) / (1 - x)
        deriv_alpha = log(x) - (digamma(alpha) - digamma(alpha + beta))
        deriv_beta = log1p(-x) - (digamma(beta) - digamma(alpha + beta))
    else
        error("x is outside of support: $x")
    end
    (deriv_x, deriv_alpha, deriv_beta)
end

function random(::Beta, alpha::Real, beta::Real)
    rand(Distributions.Beta(alpha, beta))
end

(::Beta)(alpha, beta) = random(Beta(), alpha, beta)

has_output_grad(::Beta) = true
has_argument_grads(::Beta) = (true, true)
get_static_argument_types(::Beta) = [Float64, Float64]

export beta

###############
# categorical #
###############

struct Categorical <: Distribution{Int} end

"""
    categorical(probs::AbstractArray{U, 1}) where {U <: Real}

Given a vector of probabilities `probs` where `sum(probs) = 1`, sample an `Int` `i` from the set {1, 2, .., `length(probs)`} with probability `probs[i]`.
"""
const categorical = Categorical()

function logpdf(::Categorical, x::Int, probs::AbstractArray{U,1}) where {U <: Real}
    log(probs[x])
end

function logpdf_grad(::Beta, x::Int, probs::AbstractArray{U,1})  where {U <: Real}
    grad = zeros(length(probs))
    grad[x] = 1.0
    (nothing, grad)
end

function random(::Categorical, probs::AbstractArray{U,1}) where {U <: Real}
    rand(Distributions.Categorical(probs))
end

(::Categorical)(probs) = random(Categorical(), probs)

has_output_grad(::Categorical) = false
has_argument_grads(::Categorical) = (true,)
get_static_argument_types(::Categorical) = [Vector{Float64}]

export categorical

####################
# uniform_discrete #
####################

struct UniformDiscrete <: Distribution{Int} end

"""
    uniform_discrete(low::Integer, high::Integer)

Sample an `Int` from the uniform distribution on the set {low, low + 1, ..., high-1, high}.
"""
const uniform_discrete = UniformDiscrete()

function logpdf(::UniformDiscrete, x::Int, low::Integer, high::Integer)
    d = Distributions.DiscreteUniform(low, high)
    Distributions.logpdf(d, x)
end

function logpdf_grad(::UniformDiscrete, x::Int, lower::Integer, high::Integer)
    (nothing, nothing, nothing)
end

function random(::UniformDiscrete, low::Integer, high::Integer)
    rand(Distributions.DiscreteUniform(low, high))
end

(::UniformDiscrete)(low, high) = random(UniformDiscrete(), low, high)

has_output_grad(::UniformDiscrete) = false
has_argument_grads(::UniformDiscrete) = (false, false)
get_static_argument_types(::UniformDiscrete) = [Int, Int]

export uniform_discrete

######################
# uniform_continuous #
######################

struct UniformContinuous <: Distribution{Float64} end

const uniform_continuous = UniformContinuous()

"""
    uniform(low::Real, high::Real)

Sample a `Float64` from the uniform distribution on the interval [low, high].
"""
const uniform = uniform_continuous

function logpdf(::UniformContinuous, x::Real, low::Real, high::Real)
    (x >= low && x <= high) ? -log(high-low) : -Inf
end

function logpdf_grad(::UniformContinuous, x::Real, low::Real, high::Real)
    inv_diff = 1. / (high-low)
    (0., inv_diff, -inv_diff)
end

function random(::UniformContinuous, low::Real, high::Real)
    rand() * (high - low) + low
end

(::UniformContinuous)(low, high) = random(UniformContinuous(), low, high)

has_output_grad(::UniformContinuous) = true
has_argument_grads(::UniformContinuous) = (true, true)
get_static_argument_types(::UniformContinuous) = [Float64, Float64]

export uniform_continuous, uniform

###########
# poisson #
###########

struct Poisson <: Distribution{Int} end

"""
    poisson(lambda::Real)

Sample an `Int` from the Poisson distribution with rate `lambda`.
"""
const poisson = Poisson()

function logpdf(::Poisson, x::Integer, lambda::Real)
    x * log(lambda) - lambda - lgamma(x+1)
end

function logpdf_grad(::Poisson, x::Integer, lambda::Real)
    error("Not implemented")
    (nothing, nothing)
end

function random(::Poisson, lambda::Real)
    rand(Distributions.Poisson(lambda))
end

(::Poisson)(lambda) = random(Poisson(), lambda)

has_output_grad(::Poisson) = false
has_argument_grads(::Poisson) = (false,)
get_static_argument_types(::Poisson) = [Float64]

export poisson


"""
Compute Rayleigh scattering cross-sections.

Most of the expressions are obtained from
Bodhaine et al. (1999) J. Atmosph. Ocean. Tech., 16 1854.

If invoked as julia Rayleigh.jl, computes the cross sections for 337 and 777 nm.
From the julia REPL you can compute cross sections for any wavelength as e.g.
```julia-repl
julia> include("rayleigh.jl")
julia> using Main.Ray
julia> rayleigh(337 * co.nano)
3.4371080894577986e-30
```
Alejandro Luque, IAA-CSIC, 2019
"""
module Ray

export AIR_STP, rayleigh, co

include("constants.jl")
const co = constants

# Depolarization function for different gases.  λ here is μm.
const F = Dict("N2" => λ -> 1.034 + 3.17e-4 / λ^2,
               "O2" => λ -> 1.096 + 1.385e-3 / λ^2 + 1.448e-4 / λ^4,
               "Ar" => λ -> 1.00,
               "CO2" => λ -> 1.15)

# These are densities from ICAO, extracted from Wikipedia.  Note that
# the expressions in Bodhaine assume 300 ppmv CO2 but we consider this
# close enough for our purposes.
const AIR_STP = Dict("N2" => 780840,
                     "O2" => 209476,
                     "Ar" => 9340,
                     "CO2" => 314)

# The density of air in constants.jl is defined at (273.15 K, 101.325 kPa)
# but the functions in Bodhaine et al. 1999 are valid for 288.15 K so we have to
# rescale
nair = co.nair * 273.15 / 288.15

"""
Compute the refractive index of air.

Takes the wavelength λ in m and returns (n - 1).
Uses the formula from Edlén (1953) referenced from Bodhaine et al. 1999.

"""
function edlen(λ)
    a, b, c = 6432.8, 2949810, 25540
    p, q = 146, 41

    λ_μm = λ / co.micro
    λ2 = 1 / λ_μm^2

    1.0 + 1e-8 * (a + b / (p - λ2) + c / (q - λ2))
end

"""
Compute the refractive index of air.

Takes the wavelength λ in m and returns (n - 1).
Uses the formula from Peck and Reeder (1953) referenced from Bodhaine et al. 1999.

"""
function peck(λ)
    a, b, c = 8060.51, 2480990, 17455.7
    p, q = 132.274, 39.32957

    λ_μm = λ / co.micro
    λ2 = 1 / λ_μm^2

    1.0 + 1e-8 * (a + b / (p - λ2) + c / (q - λ2))
end

"""
Compute the Rayleigh scattering cross-section.
"""
function σ(λ; composition=AIR_STP)
    n = peck(λ)
    P = sum(ppmv * F[gas](λ / co.micro) for (gas, ppmv) in composition) / 1e6

    N = (n^2 - 1)^2 / (n^2 + 2)^2
    24π^3 * N * P / (λ^4 * nair^2)
end

function main()
    println("# Rayleigh scattering cross-section for 337 nm")
    σ = rayleigh(337 * co.nano)
    println("σ: $σ")

    println()
    println("# Rayleigh scattering cross-section for 777 nm")
    σ = rayleigh(777 * co.nano)
    println("σ: $σ")
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end

# Depolarization functions
end


using Revise
using Test
using NNFEM
using PyCall
using PyPlot
using JLD2
using ADCME
using LinearAlgebra
np = pyimport("numpy")

"""
1  -  2  -  3  -  4

local node id: left to right
"""

ndofs = 2
nx = 4
nnodes, neles = (nx + 1), nx
ngp=2
T = 200
NT = 200
Δt = T/NT

function BoundaryCondition(tid::Int64)
# PARAMETER: distance between nodes 0-1 and 0-2
l = 1.0
nodes = zeros(Float64, nnodes, ndofs)
for i = 1:nx+1
    nodes[i, 1], nodes[i, 2] = (i - 1)*l, 0.0
end

EBC, g = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
#fix node 1
EBC[1, :] .= -1

function ggt(t)
    return zeros(sum(EBC.==-2)), zeros(sum(EBC.==-2))
end
gt = ggt

#pull in the x direction
FBC, fext = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)

# todo PARAMETER
FORCE_TYPE = "non-constant"

if FORCE_TYPE == "constant"
    FBC[nx+1, 1]  = -1
    fext[nx+1, 1] = 1.0e3
else
    FBC[nx+1, 1] = -2
end

#force load function
function ft(t)
    return  sin(pi*t/T) * (0.4*tid + 1.6)
end

npoints = nnodes
node_to_point = collect(1:npoints)
return nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point
end


__precompile__()

module SGSDist

    using Logging
    import DelimitedFiles: readdlm
    import Random: rand, randn, AbstractRNG, GLOBAL_RNG

    import DifferentialEquations: SDEProblem, solve, EulerHeun, EM, RKMil
    import Distributions: @distr_support, @check_args, ContinuousUnivariateDistribution, Normal, params, pdf, cdf, ccdf, logpdf, logcdf, logccdf, fit, mean, var, skewness, kurtosis
    import Interpolations: LinearInterpolation, Line
    import QuadGK: quadgk
    import SpecialFunctions: loggamma
    import StatsBase: var, skewness, kurtosis

    export SGS, params, mean, var, std, skewness, kurtosis, pdf, pdf_xmax, logN, logpdf, cdf, logcdf, ccdf, logccdf, fit, CAM1D, Hasselmann1D, RK4, Euler, fdt, rand, air, epsilon_condition, epsilon_constraint, ks_constraint, small_E_constraint, fit_mm, param_E, param_b, param_g, invitp

    include("constraints.jl")
    include("continuous.jl")
    include("fit.jl")
    include("markov.jl")
    include("sampledata.jl")

    # precompile hints
    # d = SGS()
    # rand(d, 1)

end