text:
module MLJTuning

## METHOD EXPORT

# defined in tuned_models.jl:
export TunedModel

# defined in strategies/:
export Explicit, Grid, RandomSearch, LatinHypercube

# defined in selection_heuristics/:
export NaiveSelection

# defined in learning_curves.jl:
export learning_curve!, learning_curve

## METHOD IMPORT

import MLJBase
using MLJBase
import MLJBase: Bounded, Unbounded, DoublyUnbounded,
    LeftUnbounded, RightUnbounded, _process_accel_settings, chunks
using RecipesBase
using Distributed
import Distributions
import ComputationalResources: CPU1, CPUProcesses,
    CPUThreads, AbstractResource
using Random
using ProgressMeter
using LatinHypercubeSampling

## CONSTANTS

const DEFAULT_N = 10 # for when `default_n` is not implemented

## INCLUDE FILES

include("utilities.jl")
include("tuning_strategy_interface.jl")
include("selection_heuristics.jl")
include("strategies/explicit.jl")
include("strategies/grid.jl")
include("strategies/random_search.jl")
include("strategies/latin_hypercube.jl")
include("tuned_models.jl")
include("range_methods.jl")
include("plotrecipes.jl")
include("learning_curves.jl")

end


text:
# -*- coding: utf-8 -*-
# # Agent Based Model Augmentation
#
# SemanticModels supports model augmentation, which is the derivation of new models from old models with different (usually more advanced) capabilities. The approach is to define a type to represent a class of models and then a set of transformations that can act on that type to change the capabilities of the model.
#
# We can apply our model augmentation framework to models that are not defined as an analytical mathematical expression.
# A widely used class of models for complex systems are *agent based* in that they have an explicit representation of the agents with states and functions to represent their behavior and interactions. This notebook examines how to apply model transformations to augment agent based simulations.

# We are going to use the simulation in `examples/agentbased.jl` as a baseline simulation and add capabilities to the simulation with SemanticModels transformations. The simulation in question is an implementation of a basic SIRS model on a static population. We will make two augmentations.
#
# 1. Add *un estado de los muertos* or *a state for the dead*, transforming the model from SIRS to SIRD
# 2. Add *vital dynamics* a represented by net population growth
#
# These changes to the model could easily be made by changing the source code to add the features. However, this notebook shows how those changes could be scripted by a scientist. As we all know, once you can automate a scientific task by introducing a new technology, you free the mind of the scientist for more productive thoughts.
#
# In this case we are automating the implementation of model changes to free the scientist to think about *what augmentations to the model should I make?* instead of *how do I implement these augmentations?*

using SemanticModels
using SemanticModels.ExprModels.Parsers
using SemanticModels.ExprModels.ExpStateModels
import Base: push!

samples = 7
nsteps = 10
finalcounts = Any[]

println("Running Agent Based Simulation Augmentation Demo")
println("================================================")
println("demo parameters:\n\tsamples=$samples\n\tnsteps=$nsteps")

# ## Baseline SIRS model
#
# Here is the baseline model, which is read in from a text file. You could instead of using `parsefile` use a `quote/end` block to code up the baseline model in this script.
#
#
#
# Agents progress from the susceptible state to infected and then recovered, and get become susceptible again after recovery. See the file `../examples/agentbased.jl` for a full description of this model

expr = parsefile("agentbased.jl")
m = model(ExpStateModel, expr)

println("\nRunning basic model")
AgentModels = eval(m.expr)
for i in 1:samples
    newsam, counts = AgentModels.main(nsteps)
    push!(finalcounts, (model=:basic, counts=counts))
end

m

# ## Adding the Dead State
#
#
#
# We are going to add an additional state to the model to represent the infectious disease fatalities. The user must specify what that concept means in terms of the name for the new state and the behavior of that state. `D` is a terminal state for a finite automata.

# +
println("\nThe system states are $(m.states.args)")
println("\nAdding un estado de los muertos")

put!(m, ExpStateTransition(:D, :((x...)->:D)))

println("\nThe system states are $(m.states.args)")
# once you are dead, you are dead forever
println("\nThere is no resurrection in this model")
println("\nInfected individuals recover or die in one step")

# add a transition rule for infected -> recovered, dead, or infected
m[:I] = :((x...)->begin
        roll = mod(rand(Int),3)
        if roll == 1
            return :R
        elseif roll == 2
            return :D
        else
            return :I
        end
    end
)
@show m[:I]
# -

println("\nRunning SIRD model")
AgentModels = eval(m.expr)
for i in 1:samples
    newsam, counts = AgentModels.main(nsteps)
    push!(finalcounts, (model=:sird, counts=counts))
end

# Some utilities for manipulating functions at a higher level than expressions.

# +

struct Func end

function push!(::Func, func::Expr, ex::Expr)
    push!(bodyblock(func), ex)
end
# -

# ## Population Growth
#
# Another change we can make to our model is the introduction of population growth. Our model for population is that on each timestep, one new suceptible person will be added to the list of agents. We use the `tick!` function as an anchor point for this transformation.
#
#

println("\nAdding population growth to this model")
stepr = filter(x->isa(x,Expr), findfunc(m.expr, :tick!))[1]
@show stepr
push!(Func(), stepr, :(push!(sm.agents, :S)))
println("------------------------")
@show stepr;

println("\nRunning growth model")
AgentModels = eval(m.expr)
for i in 1:samples
    newsam, counts = AgentModels.main(nsteps)
    push!(finalcounts, (model=:growth, counts=counts))
end

# ## Presentation of results
#
# We have accumulated all of our simulation runs into the list `finalcounts` we process those simulation runs into summary tables describing the results of those simulations. This table can be used to make decisions and drive further inquiry.

# +
println("\nModel\t Counts")
println("-----\t ------")
for result in finalcounts
    println("$(result.model)\t$(result.counts)")
end

function groupagg(x::Vector{Tuple{S,T}}) where {S,T}
    c = Dict{S, Tuple{Int, T}}()
    # c2 = Dict{S, T}()
    for r in x
        g = first(r)
        c[g] = get(c, g,(0, 0.0)) .+ (1, last(r))
    end
    return c
end

mean_healthy_frac = [(r.model,
                  map(last, filter(x->(x.first == :R || x.first == :S), r.counts))[1] / sum(map(last, r.counts))[1])
                 for r in finalcounts] |> groupagg

num_unhealthy = [(r.model,
                  map(last,
                      sum(map(last, filter(x->(x.first != :R && x.first != :S),
                             r.counts)))))
                 for r in finalcounts] |> groupagg

println("\nModel\t Count \t Num Unhealthy \t Mean Healthy %")
println("-----\t ------\t --------------\t  --------------")
for (g, v) in mean_healthy_frac
    μ = last(v)/first(v)
    μ′ = round(μ*100, sigdigits=5)
    x = round(last(num_unhealthy[g]) / first(num_unhealthy[g]), sigdigits=5)
    println("$g\t   $(first(v))\t  $(rpad(x, 6))\t   $(μ′)")
end

ext: jl


text:
module TestStatsModels
using DataFrames
using Base.Test

# Tests for statsmodel.jl

# A dummy RegressionModel type
immutable DummyMod <: RegressionModel
    beta::Vector{Float64}
    x::Matrix
    y::Vector
end

## dumb fit method: just copy the x and y input over
StatsBase.fit(::Type{DummyMod}, x::Matrix, y::Vector) =
    DummyMod(collect(1:size(x, 2)), x, y)
StatsBase.model_response(mod::DummyMod) = mod.y
## dumb coeftable: just prints the "beta" values
StatsBase.coeftable(mod::DummyMod) =
    CoefTable(reshape(mod.beta, (size(mod.beta,1), 1)),
              ["'beta' value"],
              ["" for n in 1:size(mod.x,2)],
              0)

## Test fitting
d = DataFrame()
d[:y] = [1:4;]
d[:x1] = [5:8;]
d[:x2] = [9:12;]
d[:x3] = [13:16;]
d[:x4] = [17:20;]

f = y ~ x1 * x2
m = fit(DummyMod, f, d)
@test model_response(m) == d[:y]

## test prediction method
## vanilla
StatsBase.predict(mod::DummyMod) = mod.x * mod.beta
@test predict(m) == [ ones(size(d,1)) d[:x1] d[:x2] d[:x1].*d[:x2] ] * collect(1:4)

## new data from matrix
StatsBase.predict(mod::DummyMod, newX::Matrix) = newX * mod.beta
mm = ModelMatrix(ModelFrame(f, d))
@test predict(m, mm.m) == mm.m * collect(1:4)

## new data from DataFrame (via ModelMatrix)
@test predict(m, d) == predict(m, mm.m)

d2 = deepcopy(d)
d2[3, :x1] = NA
@test length(predict(m, d2)) == 4

## test copying of names from Terms to CoefTable
ct = coeftable(m)
@test ct.rownms == ["(Intercept)", "x1", "x2", "x1 & x2"]

## show with coeftable defined
@show m

## with categorical variables
d[:x1p] = PooledDataArray(d[:x1])
f2 = y ~ x1p
m2 = fit(DummyMod, f2, d)

@test coeftable(m2).rownms == ["(Intercept)", "x1p - 6", "x1p - 7", "x1p - 8"]

## predict w/ new data missing levels
## FAILS: mismatch between number of model matrix columns
## @test predict(m2, d[2:4, :]) == predict(m2)[2:4]

## Another dummy model type to test fall-through show method
immutable DummyModTwo <: RegressionModel
    msg::String
end

StatsBase.fit(::Type{DummyModTwo}, ::Matrix, ::Vector) = DummyModTwo("hello!")
Base.show(io::IO, m::DummyModTwo) = println(io, m.msg)

m2 = fit(DummyModTwo, f, d)
@show m2

end


text:
# ALL TESTS THAT CONCERN MODFICIATION

# include
using LatticePhysics
using Base.Test

# begin the Modification testblock
path_testset = @testset "Paths (in momentum space)" begin

################################################################################
#
#   METHODS FOR CONSTRUCTION OF PATHS INSIDE THE BZ OF A UNITCELL
#
################################################################################

################################################################################
#
#   1) TYPE PATH
#       - type definition
#       - printInfo function
#       - path String function (e.g. X-G-M-X)
#
################################################################################

# begin the testset
@testset "Type definition and information functions" begin

    # Constructors - how to construct a path
    @testset "Constructors" begin

        # construct some points
        points = Array{Float64,1}[
            [0.0, 0.0],
            [1.0, 1.0],
            [-pi, 0.0],
            [ pi,  pi]
        ]
        # construct some point names
        point_names = String[
            "A",
            "B",
            "A'",
            "something"
        ]
        # construct segment resolution
        segments = Int64[
            100,
            200,
            1200
        ]

        # Test the three possible constructors
        @test typeof(Path(points, point_names, segments)) == Path
        @test typeof(Path(points, point_names)) == Path
        @test typeof(Path()) == Path

    # end the testset here
    end;

    # Information functions printInfo and String
    @testset "Information functions" begin

        # construct some points
        points = Array{Float64,1}[
            [0.0, 0.0],
            [1.0, 1.0],
            [-pi, 0.0],
            [ pi,  pi]
        ]
        # construct some point names
        point_names = String[
            "A",
            "B",
            "A'",
            "something"
        ]
        # construct segment resolution
        segments = Int64[
            100,
            200,
            1200
        ]
        # construct the path
        path = Path(points, point_names, segments)

        # Test the path string
        @test typeof(getPathString(path)) == String
        # Test the printInfo function
        @test printInfo(path) == nothing
        @test printInfo(path, detailed=true) == nothing

    # end the testset here
    end;

# end the testset here
end;

################################################################################
#
#   2) CONSTRUCTION FUNCTIONS PATH
#       - function to add points
#       - TODO function to remove points
#       - function to scale the total resolution
#       - function to set the total resolution
#
################################################################################

# begin the testset
@testset "Construction functions for paths" begin

    # How to add points to a path
    @testset "Adding points" begin

        # construct an empty path
        path = Path()

        # Test three point addings
        @test addPointToPath!(path, [0.0, 0.0], "Gamma") == nothing
        @test addPointToPath!(path, [ pi,  pi], "K") == nothing
        @test addPointToPath!(path, [-pi, 1.0], "A", 200) == nothing

    # end the testset here
    end;

    # Modifying the segment resolution
    @testset "Changing resolution" begin

        # construct some points
        points = Array{Float64,1}[
            [0.0, 0.0],
            [1.0, 1.0],
            [-pi, 0.0],
            [ pi,  pi]
        ]
        # construct some point names
        point_names = String[
            "A",
            "B",
            "A'",
            "something"
        ]
        # construct the path (resolution up to hear: 300)
        path = Path(points, point_names)

        # Test to scale the total resolution by 2
        @test scaleResolution!(path, 2.0) == nothing
        @test sum(path.segment_resolution) == 600

        # Test to scale the total resolution by 0.25
        @test scaleResolution!(path, 0.25) == nothing
        @test sum(path.segment_resolution) == 150

        # construct the path again (resolution now: 300)
        path = Path(points, point_names)

        # Test to set the total resolution to 900
        @test setTotalResolution!(path, 900) == nothing
        @test sum(path.segment_resolution) == 900

        # Test to set the total resolution to 450
        @test setTotalResolution!(path, 450) == nothing
        @test sum(path.segment_resolution) == 450

    # end the testset here
    end;

# end the testset here
end;

################################################################################
#
#   3) DEFAULT PATHS
#       - getDefaultPathTriangular
#       - getDefaultPathSquare
#       - getDefaultPathSquareOctagon
#       - getDefaultPathFCC
#
################################################################################

# begin the testset
@testset "Default Paths" begin

    @testset "Triangular" begin
        # just get the path and test it
        @test typeof(getDefaultPathTriangular()) == Path
        # just get the path with a fixed resolution and test it
        @test typeof(getDefaultPathTriangular(resolution=1000)) == Path
    end;

    @testset "Square" begin
        # just get the path and test it
        @test typeof(getDefaultPathSquare()) == Path
        # get different versions of the path and test them
        @test typeof(getDefaultPathSquare(1)) == Path
        @test typeof(getDefaultPathSquare(2)) == Path
        # just get the path with a fixed resolution and test it
        @test typeof(getDefaultPathSquare(resolution=1000)) == Path
    end;

    @testset "Square Octagon" begin
        # just get the path and test it
        @test typeof(getDefaultPathSquareOctagon()) == Path
        # just get the path with a fixed resolution and test it
        @test typeof(getDefaultPathSquareOctagon(resolution=1000)) == Path
    end;

    @testset "FCC" begin
        # just get the path and test it
        @test typeof(getDefaultPathFCC()) == Path
        # just get the path with a fixed resolution and test it
        @test typeof(getDefaultPathFCC(resolution=1000)) == Path
    end;

# end the testset here
end;

# end the path test block here
end;


text:
using Pkg
Pkg.activate(pwd())

# # Logistic regression
# ## Loading and preparing data

using StatsPlots
using RDatasets

iris = dataset("datasets", "iris")

# ### Exercise:
# Create the `iris_reduced` dataframe in the following way:
# - Label "setosa" will be deleted.
# - Label "versicolor" will be the negative class.
# - Label "virginica" will be the positive class.
# - Add the `intercept` column with ones as entries.
# For the features, consider only petal length and petal width.
#
# **Hint**: Use the `Query` package or do it manually via the `!insertcols` function.
#
# ---
# ### Solution:

# ---

X = Matrix(iris_reduced[:, 1:3])
y = iris_reduced.label

# ### Exercise:
# Since X has two features (columns), it is simple to visualize. Use scatter plot to show
# the data. Use different colours for different classes. Try to produce a nice graph by
# including names of classes and axis labels (petal length and petal width).
#
# ---
# ### Solution:

# ---
#
# ## Training the classifier
# ### Exercise:
# Write a function log_reg which takes as an input the dataset, the labels and the initial
# point. It should use Newton's method to find the optimal weights ``w``. Print the results
# when started from zero.
#
# It would be possible to use the code optim(f, g, x, s::Step) from the previous lecture
# and define only the step function s for the Newton's method. However, sometimes it may be
# better to write simple functions separately instead of using more complex machinery.
#
# ---
# ### Solution:

# ---
#
# ## Analyzing the solution

using LinearAlgebra

separ(x::Real, w) = (-w[3]-w[1]*x)/w[2]

xlims = extrema(iris_reduced.PetalLength) .+ [-0.1, 0.1]
ylims = extrema(iris_reduced.PetalWidth) .+ [-0.1, 0.1]

@df iris_reduced scatter(
    :PetalLength,
    :PetalWidth;
    group = :Species,
    xlabel = "Petal length",
    ylabel = "Petal width",
    legend = :topleft,
    xlims,
    ylims,
)

plot!(xlims, x -> separ(x,w); label = "Separation", line = (:black,3))

#+

y_hat = σ.(X*w)
grad = X'*(y_hat.-y) / size(X,1)
norm(grad)

# ### Exercise:
# Compute how many samples were correctly and incorrectly classified.
#
# ---
# ### Solution:


text:
using Sesam
using ModelingToolkit, DifferentialEquations

@named s = sesam3()
@named pl = plant_const()

@named sp = plant_sesam_system(s,pl)
states(sp)

i_inspect_equations = () -> begin
    tmp_s = Dict([e.lhs for e in equations(s)] .=> equations(s))
    tmp_pl = Dict([e.lhs for e in equations(pl)] .=> equations(pl))
    [k ∈ keys(tmp_s) for k in keys(tmp_pl)]
    tmp = merge(tmp_s, tmp_pl)
    tmp3 = sort(string.(states(s)))
    tmp2 = sort(replace.(string.(keys(tmp)), "Differential(t)(" => ""))
    tmp4 = [(s,eq) for (s,eq) in zip(tmp3[1:length(tmp3)], tmp2)]
    tmp4[(1*20) .+ (1:20)]
    tmp4[(4*20) .+ (1:20)]
    tmp4[100:end]
end
p = pC = Dict(
    s.ϵ_tvr => 0.45,   # carbon use efficiency of microbial tvr (part by predators
        #which respire and corresponding amount of N must be mineralized)
    s.κ_E => 0.8,     ##<< amount of recycling enzyme turnover by biomass (
        # added to assimilable, i.e. uptake instead of R)
    s.a_E => 0.001*365,   ##<< C biomass allocated to enzymes 1/day /microbial biomass
    s.m => 0.005*365,    ##<< maintenance respiration rate   1/day /microbial biomass,
    s.τ => 1/60*365,  ##<< biomass turnover rate (12 days)
    s.k_L => 5.0,       ##<< 1/(x years)   # formerly 1 year
    s.k_R => 1/(20.0),        ##<< 1/(x years) # to demonstrate changes on short time scale
    s.k_mN_L => 0.05 * 60, # enzyme half-saturation constant, in magnitude of enzymes *
        # /yr enzyme turnover 60 times a year
    s.k_mN_R => 0.05 * 60, # enzyme half-saturation constant, in magnitude of enzymes *
        # /yr enzyme turnover 60 times a year
    s.ϵ => 0.5,      ##<< carbon use efficiency for growth respiration
    #i_L => t -> 1 - exp(-t),  # litter input
    s.ρ_CBtvr => 0.0,    # no carbon resorption
    pl.i_L0 => 400.0,         # g/m2 input per year (half NPP)
    pl.β_Ni0 => 25,
    pl.i_IN0 => 0,   ##<< input of mineral N,
    pl.β_Pi0 => 25*20, ## leaf litter N:P ~20(massratio Kang10)
    pl.i_IP0 => 0.65,   ##<< input of mineral P, weathering: Table3 mixed sedimentary rocks 0.65g/m2/yr Hartmann14 10.1016/j.chemgeo.2013.10.025
    pl.s_EP0 => 0.5, # plant 1/20 of typical total microbial enzyme synthesis flux
)
pN = Dict(
    s.i_BN => 0.4, ##<< potential immobilization flux rate
    s.β_NEnz => 3.1,     # Sterner02: Protein (Fig. 2.2.), high N investment (low P) need
    s.β_NB => 11.0,
    s.ρ_NBtvr => 0.0,    # no nitrogen resorption
    #s.l_N => 0.96,       #0.00262647*365     ##<< leaching rate of mineralN lN IN
    s.l_N => 0.0,
    s.ν_N =>  0.9,     # microbial N use efficiency accounting for apparent
)
pP = Dict(
    s.k_LP => pC[s.k_L], # TODO 1/x years, assume same rate as depolymerizing enzyme
    s.k_RP => pC[s.k_R], # TODO
    s.i_BP => pN[s.i_BN],       ##<< potential immobilization flux rate
    s.β_PEnz => 50.0,     # TODO Sterner02: Protein (Fig. 2.2.), high N investment (low P) need
    s.β_PB => 40.0, # Sterner02: low P in microbial cell walls, more in genetic machinary and energy compounds
    s.ρ_PBtvr => 0.0,    # no phosphorus resorption
    s.l_P => 0.01,      # almost no leaching
    s.ν_P =>  0.9,     # microbial P use efficiency accounting for apparent
    # mineralization, pertains only to SOM depolimerization, biomineralization is
    # all mineralized
    # s.k_mN_Pl => 0.05 * 60, # enzyme half-saturation constant, in magnitude of enzymes *
        # /yr enzyme turnover 60 times a year
)
p = p0 = merge(pC, pN, pP)

u0 = u0C = Dict(
    s.B => 17,
    s.L => 100,
    s.R => 1100,
    #s.cumresp => 0.0,
    s.α_R => 0.1,
    s.α_RP => 0.1,
    s.α_LP => 0.1,
)
#u0C[s.α_L] = 1.0 - u0C[s.α_R] - u0C[s.α_LP] - u0C[s.α_RP]
u0N = Dict(
    s.I_N => 0.04, ##<< inorganic pool gN/m2
    s.L_N => u0[s.L]/p[pl.β_Ni0],
    s.R_N => u0[s.R]/calculate_β_NR_sesam3(p,s) #p[s.β_NB],
    )
u0P = Dict(
    s.I_P => 0.04, ##<< TODO inorganic pool gN/m2
    s.L_P => u0[s.L]/p[pl.β_Pi0],
    s.R_P => u0[s.R]/calculate_β_PR_sesam3(p,s) #p[s.β_NB],
    )
u00 = u0 = merge(u0C, u0N, u0P)
#u0[s.R]/u0[s.R_N] # smaller p[s.β_NB]

tspan = (0.0,200.0)

#prob = ODEProblem(sp,[t for t in u0], tspan, [t for t in p])
#prob = ODEProblem(sp, remove_units(u0), tspan, remove_units(p))
prob = ODEProblem(sp, u0, tspan, p)
#prob = ODEProblem(sp,u0, tspan, p, jac=true)
#sol = sol_sesam3 = solve(prob);
#sol = sol_sesam3 = solve(prob, Tsit5());
sol = sol_sesam3 = solve(prob, Rodas4());
#sol = sol_sesam3 = solve(prob, Tsit5(), callback=PositiveDomain(prob.u0));

i_plot = () -> begin
    #import Plots
    ts = tspan
    ts = (2,2.2)
    Plots.plot(sol)
   Plots.plot(sol, vars=[s.R])
   Plots.plot(sol, vars=[s.calculate_β_PR_sesam3])
   Plots.plot(sol, vars=[s.lim_C, s.lim_N, s.lim_P], tspan=ts)
   Plots.plot(sol, vars=[s.α_L, s.α_R, s.α_LP, s.α_RP])
   Plots.plot(sol, vars=[s.α_L, s.α_R, s.α_LP, s.α_RP], tspan=ts)
   Plots.plot(sol, vars=[s.revenue_L, s.revenue_R, s.revenue_LP, s.revenue_RP], tspan=ts)
   Plots.plot(sol, vars=[s.dec_RP, s.dec_RPPlant], tspan=ts)
   Plots.plot(sol, vars=[s.I_P])
   Plots.plot(sol, vars=[s.u_PlantP, s.u_immPPot])
   Plots.plot(sol, vars=[s.p_uPmic])

   Plots.plot(sol, vars=[p[s.a_E] * s.B])
   Plots.plot(sol, vars=[s.α_LP, s.α_RP])
   Plots.plot(sol, vars=[s.α_LP, s.α_RP], tspan=(68,92))
   Plots.plot(sol, vars=[s.α_LPT, s.α_RPT], tspan=(80,92))

    Plots.plot(sol, vars=[s.β_PR], tspan=ts)
    Plots.plot(sol, vars=[s.tvr_B/p[s.β_PB], s.tvr_B/p[s.β_PBtvr], s.resorp_P, s.tvr_B/p[s.β_PBtvr] + s.resorp_P], tspan=ts)

    Plots.plot(sol, vars=[s.β_PBtvr, p[s.β_PB] * (1-p[s.ρ_CBtvr])/(1-p[s.ρ_PBtvr])], tspan=ts)
    calculate_β_PR_sesam3(p,s), sol[s.β_PR,end]

end

@testset "no P recycled" begin
    #include("test/test_sesam3_sol.jl")
    include("test_sesam3_sol.jl")
end;

@testset "P recycled" begin
    p2 = copy(p)
    p2[s.ρ_PBtvr] = 0.1  # 10% N recycled during turnover
    prob2 = ODEProblem(sp, u0, tspan, p2)
    sol = sol2 = solve(prob2, Rodas4());
    #include("test/test_sesam3_sol.jl")
    include("test_sesam3_sol.jl")
end;


text:
module LazyDicts

import Base: show, getindex, setindex!, keys, values, keytype, valtype, map

export LazyDict

struct LazyDict{F, K, V, D <: AbstractDict{K, V}}
    f :: F
    d :: D
end

LazyDict(f, kv) = LazyDict(f, Dict(kv))
LazyDict(f, kv::Pair...) = LazyDict(f, Dict(kv...))
LazyDict(f, K::Type, V::Type) = LazyDict(f, Dict{K, V}())
LazyDict(f) = LazyDict(f, Any, Any)

function getindex(d::LazyDict, key)
    get!(()->(d.f)(key), d.d, key)
end

function getindex(d::LazyDict, key_first, key_rest...)
    get!(()->(d.f)(key_first, key_rest...), d.d, (key_first, key_rest...))
end

function setindex!(d::LazyDict, value, key...)
    setindex!(d.d, value, key...)
end

keytype(::Type{LazyDict{F, K, V, D}}) where {F, K, V, D} = K
valtype(::Type{LazyDict{F, K, V, D}}) where {F, K, V, D} = V

keys(d::LazyDict) = keys(d.d)
values(d::LazyDict) = values(d.d)

function map(g, d::LazyDict)
    LazyDict(g ∘ d.f, Dict(k => g(d[k]) for k in keys(d)))
end

(d::LazyDict)(k...) = d[k...]

end # module

#############
### tests ###
#############

# plus = LazyDict(+, Tuple{Int, Int}, Int)
#
# plus[1, 2] == 3
# plus(1, 2) == 3
#
# keytype(typeof(plus)) == Tuple{Int, Int}
# valtype(typeof(plus)) == Int
#
# fibs = LazyDict(0 => BigInt(0), 1 => BigInt(1)) do x
#     fibs[x-1] + fibs[x-2]
# end
#
# first_fibs = [0,1,1,2,3,5,8,13,21,34,55]
# map(fibs, 0:10) == first_fibs
#
#
# doubled_fibs = map(x->2x, fibs)
# doubled_fibs.(0:10) == 2 * first_fibs
#
# @time fibs(10000)
# @time fibs(10000)


text:
using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter

#----------------------------------------------------------------------------
# step 1: create a computational domain
# fixed on the bottom and push from right

m = 40
n = 20
h = 1/n
NT = 50
Δt = 0.01/NT

elements = []
# prop = Dict("name"=> "PlaneStrain", "rho"=> 2000, "E"=> 1e10, "nu"=> 0.35)
prop = Dict("name"=> "ViscoelasticityMaxwell", "rho"=> 2000, "E"=> 1e10, "nu"=> 0.35, "eta"=>1e5)

coords = zeros((m+1)*(n+1), 2)
for j = 1:n
    for i = 1:m
        idx = (m+1)*(j-1)+i
        elnodes = [idx; idx+1; idx+1+m+1; idx+m+1]
        ngp = 2
        nodes = [
        (i-1)*h (j-1)*h
        i*h (j-1)*h
        i*h j*h
        (i-1)*h j*h
        ]
        coords[elnodes, :] = nodes
        push!(elements, SmallStrainContinuum(nodes, elnodes, prop, ngp))
    end
end

# fixed on the bottom, push on the right
EBC = zeros(Int64, (m+1)*(n+1), 2)
FBC = zeros(Int64, (m+1)*(n+1), 2)
g = zeros((m+1)*(n+1), 2)
f = zeros((m+1)*(n+1), 2)

for i = 1:m+1
    EBC[i+n*(m+1),:] .= -1
end

Edge_Traction_Data = Array{Int64}[]
for i = 1:n
  elem = elements[(i-1)*m + m]
  for k = 1:4
    if elem.coords[k,1]>2-0.001 && elem.coords[k+1>4 ? 1 : k+1,1]>2-0.001
      push!(Edge_Traction_Data, [(i-1)*m + m, k, 1])
    end
  end
end
Edge_Traction_Data = hcat(Edge_Traction_Data...)'|>Array

ndims = 2
domain = Domain(coords, elements, ndims, EBC, g, FBC, f, Edge_Traction_Data)
globaldata = example_global_data(domain)

# visualize_boundary(domain)
#----------------------------------------------------------------------------
# step 2: Simulation

Dstate = zeros(domain.neqs)
state = zeros(domain.neqs)
velo = zeros(domain.neqs)
acce = zeros(domain.neqs)
EBC_func = nothing
FBC_func = nothing
Body_func = nothing
function Edge_func(x, y, t, idx)
  return [-1e6*ones(length(x)) zeros(length(x))]
end

globaldata = GlobalData(state, Dstate, velo, acce, domain.neqs, EBC_func, FBC_func,Body_func, Edge_func)
assembleMassMatrix!(globaldata, domain)

@showprogress for i = 1:NT
    global globaldata, domain = GeneralizedAlphaSolverStep(globaldata, domain, Δt)
end

#----------------------------------------------------------------------------
# step 3: visualize
d_ = hcat(domain.history["state"]...)'|>Array
# visualize_von_mises_stress_on_scoped_body(d_, domain, scale_factor=10.0)
p = visualize_total_deformation_on_scoped_body(d_, domain, scale_factor=10.0)
saveanim(p, "julia.gif")
close("all")

#----------------------------------------------------------------------------
# step 4: verification with AdFem

max_stars_repo_name: ADCMEMarket/ADCMEImages

max_issues_repo_name: ADCMEMarket/ADCMEImages

max_forks_repo_name: ADCMEMarket/ADCMEImages


text:
using StaggeredGridCalculus
using Test
# using Statistics: mean
using AbbreviatedTypes

Base.isapprox(a::Tuple, b::Tuple; kws...) = all(p -> isapprox(p...; kws...), zip(a,b))

# @testset "StaggeredGridCalculus" begin

include("enumtype.jl")
include("util.jl")
include("grid.jl")
# include("interval.jl")
# include("gridgen.jl")
include("partial.jl")
include("curl.jl")
include("divergence.jl")
include("gradient.jl")
include("mean.jl")
include("permutation.jl")
include("pml.jl")

# end  # @testset "StaggeredGridCalculus"


text:
using LinearAlgebra
using SparseArrays
using Laplacians
using Random

function lmdmin_val(L,n)
	I=spzeros(n,n);
	for i=1:n
		I[i,i]=n;
	end
	for i=1:n
		for j=1:n
			L[i,j]=I[i,j]-L[i,j];
		end
	end
	u=randn(n);
	u=u./norm(u);
	y=L*u;
	y=y./norm(y);
	while norm(y-u)>1e-2
		u=y;
		y=L*u;
		y=y./norm(y);
	end
	return n-u'*L*y;
end

function lmdedge(L,s,k)
	for i=1:k
		L[Int(s[i]),Int(s[i])]+=1;
	end
	f=eigen(L);
	return f.values[1];
end

function lmd_LG(L,s,n)
	S=union(1:n);
	setdiff!(S,s);
	LG=L[S,S];
	f=eigvals(LG);
	k=1;
	while f[k]<1e-8
		k+=1;
	end
	return f[k];
end

function lap(G :: Graph)
    F = zeros(G.n, G.n);
    for i = 1 : G.m
        F[G.u[i], G.v[i]] -= 1
        F[G.v[i], G.u[i]] -= 1
        F[G.u[i], G.u[i]] += 1
        F[G.v[i], G.v[i]] += 1
    end
    return F
end

function lapsp(G :: Graph)
    F = spzeros(G.n, G.n)
    for i = 1 : G.m
        F[G.u[i], G.v[i]] -= 1
        F[G.v[i], G.u[i]] -= 1
        F[G.u[i], G.u[i]] += 1
        F[G.v[i], G.v[i]] += 1
    end
    return F
end

function getB(G :: Graph, mm)
    B :: SparseMatrixCSC{Float64} = spzeros(mm,G.n)
	jishu=0;
    for i = 1 : G.m
    #    if (G.u[i]>G.n0+G.n1 && G.v[i]>G.n0+G.n1)
			jishu+=1;
            B[jishu, G.u[i]] = 1
            B[jishu, G.v[i]] = -1
    #    end
    end
    return B
end

function getD(G :: Graph)
    D = spzeros(G.n-G.n0-G.n1,G.n-G.n0-G.n1);
    for i = 1 : G.m
        if (G.u[i]<=G.n0+G.n1) && (G.v[i]>G.n0+G.n1)
            D[G.v[i]-G.n0-G.n1] += 1;
        end
        if (G.v[i]<=G.n0+G.n1) && (G.u[i]>G.n0+G.n1)
            D[G.u[i]-G.n0-G.n1] += 1;
        end
    end
    return D;
end

function calc(G :: Graph ,i ,Linv, sel)
    e=spzeros(G.n-G.n0-G.n1,1)
    e[i,1]=1
    vect = Linv*e-Linv*e*(e'*Linv*sel);
    ans = sum(vect)/(1+Linv[i,i]);
    return ans
end

function findmaxexa(G :: Graph, Linv ,sel,eta)
    bestv = 0;
    tmp = 0
    maxinc = 0
    for i = G.n0+G.n1+1 : G.n
        if sel[i-(G.n0+G.n1)]< G.n1
            tmp = calc(G, i-G.n0-G.n1 , Linv ,sel)
            if tmp > maxinc && (Linv*sel)[i-G.n0-G.n1]<eta
                maxinc = tmp
                bestv = i
            end
			#print(tmp," ")
        end
    end
    return bestv
end

function exa(G :: Graph,k,eta,fout)
    L = lap(G);
    sel = zeros(G.n - G.n0- G.n1);
    for i = G.n0+G.n1+1 : G.n
        sel[i-(G.n0+G.n1)] =  sum(L[i,1:G.n0]) - sum(L[i,1:G.n0+G.n1]);
    end
	T1=time();
    Linv = inv(L[G.n0+G.n1+1:G.n,G.n0+G.n1+1:G.n]);
	firans = sum(Linv * sel);
	j=0;
	T2=time();
	#println(" ");
    for i = 1 : k
        j = findmaxexa(G, Linv ,sel,eta);
		#print(j-G.n0-G.n1," ")
        sel[j-G.n0-G.n1] += 1;
        e=spzeros(G.n-G.n0-G.n1,1);
        e[j-G.n0-G.n1,1]=1;
        Linv = Linv - Linv*e*e'*Linv/(1+Linv[j-G.n0-G.n1,j-G.n0-G.n1]);
    end
	#println("");
	T3=time();
	println(fout,"exact basetime=",T2-T1,",",k," times=",T3-T2);
    ansv = Linv * sel;
    return sum(ansv);
end

function gre(G :: Graph, k ,ep,eta,fout)
	finans= 0 ;
    d0=ep/(3*sqrt(G.n));
	#d0=eps/(3);
    #d=4*eps/(3*n^2);
	d=4*ep/(3*G.n^2)/(G.n^2);
	JLfac=1.0;
    #t= round(Int64,12*12*24*log2(G.n)/eps^2);
	t= round(Int64,JLfac*log(G.n)/(ep^2))+1;
	L = lapsp(G);
	nn = G.n-G.n0-G.n1;
    LL = L[G.n0+G.n1+1:G.n,G.n0+G.n1+1:G.n];
    mm = 0;
    for i = 1: G.m
        if ((G.u[i]>G.n0+G.n1) && (G.v[i]>G.n0+G.n1))
            mm += 1;
        end
    end
	B = getB(G,mm);
    D = getD(G);
	Dadd = spzeros(G.n-G.n0-G.n1,G.n-G.n0-G.n1);
    #sqrt.(D);
    rng = MersenneTwister(Int(round(time() * 10000.0)));
    #f = approxchol_sddm(LL, tol=d0)

    sel = zeros(nn);
    for i = G.n0+G.n1+1 : G.n
        sel[i-(G.n0+G.n1)] =  sum(L[i,1:G.n0]) - sum(L[i,1:G.n0+G.n1]);
    end
    o1 = ones(nn);
	maxx = 0;
	z1 = zeros(nn, 1)
	z2 = zeros(nn, 1)
	f = approxchol_sddm(LL, tol=1e-8);
	firans =sum(f(sel));
	#r2 = randn(rng, nn, 1);
	#r1 = randn(rng, mm, 1);

	T1=time();

    for i = 1 : k

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

####f = Laplacians.approxchol_sddm(LL+Dadd, tol=d0);
f = Laplacians.approxchol_sddm(LL, tol=d0,verbose=false);
        h1=f(o1);
        h2=f(sel);

#这里就是用求解器的步骤，f = approxchol_sddm(L, tol=d0);L是拉普拉斯矩阵，tol=d0是误差，选取1e-6完全够了，
#h1=f(o1)就是返回L的伪逆*o1向量
#####################

		Xe = zeros(nn);
		Ye = zeros(nn);
        for j = 1 : t
            r2 = randn(rng, nn, 1);
			r1 = randn(rng, mm, 1);
			xx = B'*r1;
			yy = sqrt.(D)*r2;
            f = approxchol_sddm(LL, tol=d);
			z1[:, 1] = f(xx[:, 1]);
			z2[:, 1] = f(yy[:, 1]);
			for p = 1 : nn
				Xe[p] += z1[p,1]^2;
				Ye[p] += z2[p,1]^2;
			end
		end
		add = zeros(G.n-G.n0-G.n1);
		for j = 1 : G.n-G.n0-G.n1
			add[j] += h1[j]*(1-h2[j])/(1+Xe[j]/t+Ye[j]/t);
		end
		maxx=0;
		bestv=0;
		for j = 1 : nn
			if (sel[j]<G.n1 && add[j]>maxx && h2[j]<eta)
				maxx = add[j]
				bestv = j
			end
		end
		#println(" ");
		#for j =1 : G.n-G.n0-G.n1
		#	print(add[j]," ")
		#end
		sel[bestv]+=1;
		#LL[bestv,bestv]+=1;
		D[bestv,bestv]+=1;
		Dadd[bestv,bestv]+=1;
		finans += maxx;
		#print(bestv," ")
	end
	T2=time();
	#println(" ");
	f = approxchol_sddm(LL+Dadd, tol=1e-8);
	finans =sum(f(sel));
	println(fout,"greed ",k," times=",T2-T1," eps=",ep);
	return finans;
end

function appxInvTrace(parL::SparseMatrixCSC{Float64}, parB::SparseMatrixCSC{Float64})

  f = approxchol_sddm(parL,tol=1e-6, verbose=false)
  m = size(parB, 1)
  n = size(parL, 1)
  k = round(Int, log(n)) # number of dims for JL
  print(n)
  rst = 0

  for i = 1:k
  #  V[:,i] .= f(UR[:,i])
    r = rand([1,-1],m,1)
    #ur = U'*r
    v = r'*parB
    print(v)
    v = f(v')
    #rst += sum(sqrt(v.*v))
  end

  return rst/k;

end


text:
# [ include("../src/"*s) for s in readdir("../src") ]
using PolyChaos
# using LinearAlgebra
# import FFTW
# import SpecialFunctions
# numerically compute recurrence coefficients for (almost) Gaussian density w(t)
N = 10
w(t) = exp(-t^2)
lb, ub = -Inf, Inf
@time α, β = rm_compute(w,lb,ub;Nquad=200,Npoly=N)
# analytical solution
α_ana = zeros(N)
β_ana = [ √π; [0.5*k for k=1:N-1] ]
# compare
display("Deviation for α: $(α-α_ana)")
display("Deviation for β: $(β-β_ana)")
## do the same for Chebyshev polynomials #4
v(t) = sqrt(1-t)/sqrt(1+t)
lb, ub = -1+1e-8, 1-1e-8
@time α, β = rm_compute(v,lb,ub;Nquad=2000,Npoly=N)
# analytical solution
α_ana = [-0.5; zeros(N-1) ]
β_ana = [ π; [0.25 for k=1:N-1] ]
# compare
display("Deviation for α: $(α-α_ana)")
display("Deviation for β: $(β-β_ana)")


text:
"""
Strip comments from a string

    striplinecomment{T<:String,U<:String}(a::T, cchars::U="#;")

# Arguments
- `a::String`: the string from which the comments has to be stripped
- `cchars::String`: the characters that defines comments
From https://rosettacode.org/wiki/Strip_comments_from_a_string#Julia
- `whitespace::Bool`: remove leading whitespaces.

```julia
strip_comments("test1")
strip_comments("test2 # with a comment")
strip_comments("# just a comment")
""
```
"""
function strip_comments(a::String, cchars::String = "#;"; whitespace = true)
    if whitespace
        b = strip(a)
    else
        b = rstrip(a)
    end

    0 < length(cchars) || return b
    for c in cchars
        r = Regex(@sprintf "\\%c.*" c)
        b = replace(b, r => "")
    end
    rstrip(b)
end


text:
export NormalMeanPrecision

import StatsFuns: log2π, invsqrt2π

struct NormalMeanPrecision{T <: Real} <: ContinuousUnivariateDistribution
    μ :: T
    w :: T
end

NormalMeanPrecision(μ::Real, w::Real)       = NormalMeanPrecision(promote(μ, w)...)
NormalMeanPrecision(μ::Integer, w::Integer) = NormalMeanPrecision(float(μ), float(w))
NormalMeanPrecision(μ::Real)                = NormalMeanPrecision(μ, one(μ))
NormalMeanPrecision()                       = NormalMeanPrecision(0.0, 1.0)

Distributions.@distr_support NormalMeanPrecision -Inf Inf

Distributions.support(dist::NormalMeanPrecision) = Distributions.RealInterval(minimum(dist), maximum(dist))

weightedmean(dist::NormalMeanPrecision) = precision(dist) * mean(dist)

Distributions.mean(dist::NormalMeanPrecision)    = dist.μ
Distributions.median(dist::NormalMeanPrecision)  = mean(dist)
Distributions.mode(dist::NormalMeanPrecision)    = mean(dist)
Distributions.var(dist::NormalMeanPrecision)     = inv(dist.w)
Distributions.std(dist::NormalMeanPrecision)     = sqrt(var(dist))
Distributions.cov(dist::NormalMeanPrecision)     = var(dist)
Distributions.invcov(dist::NormalMeanPrecision)  = dist.w
Distributions.entropy(dist::NormalMeanPrecision) = (1 + log2π - log(precision(dist))) / 2

Distributions.pdf(dist::NormalMeanPrecision, x::Real)    = (invsqrt2π * exp(-abs2(x - mean(dist)) * precision(dist) / 2)) * sqrt(precision(dist))
Distributions.logpdf(dist::NormalMeanPrecision, x::Real) = -(log2π - log(precision(dist)) + abs2(x - mean(dist)) * precision(dist)) / 2

Base.precision(dist::NormalMeanPrecision)     = invcov(dist)
Base.eltype(::NormalMeanPrecision{T}) where T = T

Base.convert(::Type{ NormalMeanPrecision }, μ::Real, w::Real) = NormalMeanPrecision(μ, w)
Base.convert(::Type{ NormalMeanPrecision{T} }, μ::Real, w::Real) where { T <: Real } = NormalMeanPrecision(convert(T, μ), convert(T, w))

vague(::Type{ <: NormalMeanPrecision }) = NormalMeanPrecision(0.0, tiny)

prod_analytical_rule(::Type{ <: NormalMeanPrecision }, ::Type{ <: NormalMeanPrecision }) = ProdAnalyticalRuleAvailable()

function Base.prod(::ProdPreserveType, left::NormalMeanPrecision, right::NormalMeanPrecision)
    w = precision(left) + precision(right)
    μ = (mean(left) * precision(left) + mean(right) * precision(right)) / w
    return NormalMeanPrecision(μ, w)
end

function Base.prod(::ProdAnalytical, left::NormalMeanPrecision, right::NormalMeanPrecision)
    w = precision(left) + precision(right)
    xi = mean(left)*precision(left) + mean(right)*precision(right)
    return NormalWeightedMeanPrecision(xi, w)
end


text:
using QuantEcon
using Statistics
using LinearAlgebra

include(string(pwd(),"/src/gensys.jl"))

function log_like_dsge(par,data;kalman_tol = 1e-10)
    #include(string(pwd(),"/matrices_carlos_giba.jl"))

    #par = [0.25 0.989 1 # alfa, beta, phi
    #       0.6 -0.2 -0.3 # A0
    #       0.6 0.1 -0.3 0.1 0.5 -0.2 0.1 0.8 0.5 # A1
    #       0.3 0.5 0.6 # shocks_star
    #       0.68 1.27 0.58 0.07 0.82 0.13 0.92 0.02 0.79 0.60 0.12 0.88 5.57
    #      10.34 2.08 1.10 0.54 0.30 1.99 0.70 0.31 0.79] # estimated parameters
           #29 dec: Calibrated the EUA VAR for the data we have

    # calibrated parameters
    alfa    = par[1]
    bet     = par[2]
    phi     = par[3]
    a_0piy  = par[4]
    a_0iy   = par[5]
    a_0ipi  = par[6]
    a_1yy   = par[7]
    a_1ypi  = par[8]
    a_1yi   = par[9]
    a_1piy  = par[10]
    a_1pipi = par[11]
    a_1pii  = par[12]
    a_1iy   = par[13]
    a_1ipi  = par[14]
    a_1ii   = par[15]
    sigma_y_star  = par[16]
    sigma_pi_star = par[17]
    sigma_i_star  = par[18]

    # estimated parameters
    h           = par[19]
    sig         = par[20]
    eta         = par[21]
    delta_D     = par[22]
    theta_D     = par[23]
    delta_I     = par[24]
    theta_I     = par[25]
    chi         = par[26]
    rho_a       = par[27]
    rho_gamma   = par[28]
    rho_cp      = par[29]
    rho_phi     = par[30]
    sigma_a     = par[31]
    sigma_gamma = par[32]
    sigma_cp    = par[33]
    sigma_phi   = par[34]
    rho_1       = par[35]
    rho_2       = par[36]
    lambda_pi   = par[37]
    lambda_y    = par[38]
    lambda_s    = par[39]
    sigma_i     = par[40]

    G0 = zeros(24,24)
    G1 = zeros(24,24)
    Psi = zeros(24,8)
    Pi = zeros(24,6)

    # Equation 1 (Euler)

    G0[1,1] = -1
    G0[1,2] = -(1-h)/sig
    G0[1,3] = (1-h)/sig
    G0[1,7] = 1+h
    G0[1,8] = (1-h)/sig

    G1[1,3] = (1+h)/sig
    G1[1,7] = h

    Pi[1,1] = -1
    Pi[1,2] = -(1-h)/sig
    Pi[1,3] = (1-h)/sig

    # Equation 2 (goods market equilibrium)

    G0[2,9]  = 1
    G0[2,7]  = -(1-alfa)
    G0[2,10] = -alfa*eta*(2-alfa)
    G0[2,11] = -alfa*eta
    G0[2,17] = -alfa

    # Equation 3 (tot)

    G0[3,10] = 1
    G0[3,13] = -1
    G0[3,12] = 1

    G1[3,10] = 1

    # Equation 4 (relationship q and tot)

    G0[4,10] = -(1-alfa)
    G0[4,14] = 1
    G0[4,11] = -1

    # Equation 5 (relationship q, s, pi)

    G0[5,14] = 1
    G0[5,15] = -1
    G0[5,18] = -1

    G1[5,2]  = -1
    G1[5,14] = 1
    G1[5,15] = -1

    # Equation 6 (phillips curve - domestic)

    G0[6,5]  = -bet
    G0[6,12] = 1+bet*delta_D
    G0[6,24] = -bet*(1-theta_D)*(1-theta_D*bet)/theta_D

    G1[6,12] = delta_D

    Pi[6,4] = -bet

    # Equation 7 (marginal cost)

    G0[7,24] = 1
    G0[7,10] = -alfa
    G0[7,9]  = -phi
    G0[7,7]  = -sig/(1-h)
    G0[7,20] = 1+phi

    G1[7,7] = -(sig*h)/(1-h)

    # Equation 8  (phillips curve - importing)

    G0[8,6]  = -bet
    G0[8,13] = 1+bet*delta_I
    G0[8,11] = -(1-theta_I)*(1-theta_I*bet)/theta_I
    G0[8,22] = -1

    G1[8,13] = delta_I

    Pi[8,5]  = -bet

    # Equation 9 (relationship pi and tot)

    G0[9,12] = -1
    G0[9,10] = -alfa

    G1[9,10] = -alfa
    G1[9,2]  = -1

    # Equation 10 (bugdet constraint)

    G0[10,21] = 1
    G0[10,10] = alfa
    G0[10,11] = alfa
    G0[10,9]  = -1
    G0[10,7]  = 1

    G1[10,21] = 1/bet

    # Equation 11 (UIP)

    G0[11,4]  = -1
    G0[11,15] = 1
    G0[11,8]  = 1
    G0[11,19] = -1
    G0[11,21] = chi
    G0[11,23] = -1

    Pi[11,6] = -1

    # Equation 12 (Taylor RMI)

    G0[12,8]  = 1
    G0[12,16] = -rho_1
    G0[12,9]  = -(1-rho_1-rho_2)*lambda_y
    G0[12,15] = -(1-rho_1-rho_2)*lambda_s

    G1[12,16] = rho_2
    G1[12,2]  = (1-rho_1-rho_2)*lambda_pi
    G1[12,15] = -(1-rho_1-rho_2)*lambda_s

    Psi[12,5] = sigma_i

    ###################
    ## Foreign Block ##
    ###################

    # System (Equations 13, 14 and 15)

    G0[13,17] = 1
    G0[14,17] = a_0piy
    G0[14,18] = 1
    G0[15,17] = a_0iy
    G0[15,18] = a_0ipi
    G0[15,19] = 1

    G1[13,17] = a_1yy
    G1[13,18] = a_1ypi
    G1[13,19] = a_1yi
    G1[14,17] = a_1piy
    G1[14,18] = a_1pipi
    G1[14,19] = a_1pii
    G1[15,17] = a_1iy
    G1[15,18] = a_1ipi
    G1[15,19] = a_1ii

    Psi[13,6] = sigma_y_star
    Psi[14,7] = sigma_pi_star
    Psi[15,8] = sigma_i_star

    #################
    ## Shocks block #
    #################

    # Equation 16 (a)

    G0[16,20] = 1

    G1[16,20] = rho_a

    Psi[16,1] = sigma_a

    # Equation 17 (gamma)

    G0[17,3] = 1

    G1[17,3] = rho_gamma

    Psi[17,2] = sigma_gamma

    # Equation 18 (epsilon_cp)

    G0[18,22] = 1

    G1[18,22] = rho_cp

    Psi[18,3] = sigma_cp

    # Equation 19 (phi)

    G0[19,23] = 1

    G1[19,23] = rho_phi

    Psi[19,4] = sigma_phi

    ####################
    ## Identity block ##
    ####################

    # Equation 20 (c)

    G0[20,7] = 1
    G1[20,1] = 1

    # Equation 21 (s)

    G0[21,15] = 1
    G1[21,4]  = 1

    # Equation 22 (pi_D)

    G0[22,12] = 1
    G1[22,5]  = 1

    # Equation 23 (pi_I)

    G0[23,13] = 1
    G1[23,6]  = 1

    # Equation 24 (i)

    G0[24,16] = 1
    G1[24,8]  = 1

    #nobs = 100
    nobs = size(data,1) # number of observations in data

    p = size(G1,1) # number of endogenous vars

    #pobs = 8
    pobs = size(data,2)

    nshocks = size(Psi,2) # number of shocks

    nexpec = size(Pi,2) # number of expectation errors

    sol = gensys(G0,G1,Psi,Pi; verbose = false)
    if sum(sol.eu) != 2
        return -99999999999
    end

    #Sig = zeros(p,p)
    #Sig[4,4] = par[6]

    G = zeros(pobs,p)
    G[1,9]  = 1 # y
    G[2,2]  = 1 # pi(t+1)
    G[3,8]  = 1 # i
    G[4,10] = 1 # tot
    G[5,15] = 1 # s
    G[6,17] = 1 # y_star
    G[7,18] = 1 # pi_star
    G[8,19] = 1 # i_star

    A = sol.Theta1
    R = Diagonal(repeat([1e-8],pobs)::AbstractVector)
    #Sigma = zeros(p,p)
    #Sigma[20,20] = par[31]^2 # sigma_a^2
    #Sigma[3,3]   = par[32]^2 # sigma_gamma^2
    #$Sigma[22,22] = par[33]^2 # sigma_cp^2
    # Sigma[23,23] = par[34]^2 # sigma_phi^2
    # Sigma[8,8]   = par[40]^2 # sigma_i^2
    # Sigma[17,17] = par[16]^2 # sigma_y_star^2
    # Sigma[18,18] = par[17]^2 # sigma_pi_star^2
    # Sigma[19,19] = par[18]^2 # sigma_i_star^2
    #Q = sol.Theta2*Sigma*sol.Theta2'
    Q = sol.Theta2*sol.Theta2'

    kalman_res = Kalman(A,G,Q,R) #create a Kalman filter instance

    y_mean = mean(data,dims=1)
    y_var = var(data,dims=1)
    y_var = reshape(y_var,size(y_var,2)) # transposição de forma a evitar problemas
    #y_var = repeat([1],p)

    x_hat = zeros(p) # initial mean of the state
    x_var = I(p) # variance initial of state
    # x_var = x_var*x_var'
    set_state!(kalman_res,x_hat,x_var)

    #fit = zeros(nobs,pobs) # antes de pobs (8 agora) estava 4

    llh = zeros(nobs)

    burnin_kalman = 10

    for j in 1:(nobs-1)
        med = kalman_res.cur_x_hat
        #fit[j,:] = med
        varian = kalman_res.cur_sigma
        #println(det(varian))
        eta = data[j+1,:] - kalman_res.G*med #mean loglike
        P = kalman_res.G*varian*kalman_res.G' + kalman_res.R
        teste_cond = 1/cond(P)
        if teste_cond < kalman_tol
            llh[j] = -500
        #elseif det(P) <= 0
        #    llh[j] = -500
        else
            llh[j] = -(p*log(2*pi) + logdet(P) .+ eta'*inv(P)*eta)/2
            QuantEcon.update!(kalman_res,data[j+1,:]) #updating the kalman estimates
        end
        #println(kalman_res.cur_sigma)
        #println(det(varian))
        #println(P)
    end
    llh = llh[burnin_kalman+1:length(llh)]
    return sum(llh)
end

max_stars_repo_name: danmrc/DSGE-HMC

max_issues_repo_name: danmrc/DSGE-HMC

max_forks_repo_name: danmrc/DSGE-HMC


text:
parent_type(::NfRelOrdElem{T, U}) where {T, U} = NfRelOrd{T, fractional_ideal_type(order_type(parent_type(T))), U}
parent_type(::Type{NfRelOrdElem{T, U}}) where {T, U} = NfRelOrd{T, fractional_ideal_type(order_type(parent_type(T))), U}

################################################################################
#
#  Deepcopy
#
################################################################################

function Base.deepcopy_internal(x::NfRelOrdElem, dict::IdDict)
  z = parent(x)()
  z.elem_in_nf = Base.deepcopy_internal(x.elem_in_nf, dict)
  if x.has_coord
    z.has_coord = true
    z.coordinates = Base.deepcopy_internal(x.coordinates, dict)
  end
  return z
end

################################################################################
#
#  Parent object overloading
#
################################################################################
#=
@doc Markdown.doc"""
      (O::NfRelOrd)(a::NumFieldElem, check::Bool = true) -> NfRelOrdElem

Given an element $a$ of the ambient number field of $\mathcal O$, this
function coerces the element into $\mathcal O$. If `check` is `true`
it will be checked that $a$ is contained in $\mathcal O$.
"""
=#
function (O::NfRelOrd{S, T, U})(a::U, check::Bool = true) where {S, T, U}
  if check
    x, y = _check_elem_in_order(a, O)
    !x && error("Number field element not in the order.")
    return NfRelOrdElem{S, U}(O, deepcopy(a), y)
  else
    return NfRelOrdElem{S, U}(O, deepcopy(a))
  end
end

#=
@doc Markdown.doc"""
      (O::NfRelOrd)(a::NfRelOrdElem, check::Bool = true) -> NfRelOrdElem

Given an element $a$ of some order in the ambient number field of
$\mathcal O$, this function coerces the element into $\mathcal O$.
If `check` is `true` it will be checked that $a$ is contained in
$\mathcal O$.
"""
=#
(O::NfRelOrd{S, T, U})(a::NfRelOrdElem{S, U}, check::Bool = true) where {S, T, U} = O(nf(O)(a.elem_in_nf), check)

function (O::NfRelOrd)(a::Vector{T}, check::Bool = true) where T
  t = nf(O)()
  basis = basis_nf(O, copy = false)
  for i = 1:degree(O)
    t += a[i]*basis[i]
  end
  s = O(t, check)
  s.coordinates = a
  s.has_coord = true
  return s
end

(O::NfRelOrd)(a::NfOrdElem, check::Bool = true) = O(nf(O)(a.elem_in_nf), check)

(O::NfRelOrd)(a::IntegerUnion) = O(nf(O)(a))

#=
@doc Markdown.doc"""
      (O::NfRelOrd)() -> NfRelOrdElem

Constructs a new element of $\mathcal O$ which is set to $0$.
"""
=#
(O::NfRelOrd{T, S, U})() where {T, S, U} = NfRelOrdElem{T, U}(O)

################################################################################
#
#  "Assure" functions for fields
#
################################################################################

function assure_has_coord(a::NfRelOrdElem)
  if a.has_coord
    return nothing
  else
    x, y = _check_elem_in_order(a.elem_in_nf, parent(a))
    !x && error("Not a valid order element.")
    a.coordinates = y
    a.has_coord = true
    return nothing
  end
end

################################################################################
#
#  Coordinates
#
################################################################################

@doc Markdown.doc"""
      coordinates(a::NfRelOrdElem{T}) -> Vector{T}

Returns the coefficient vector of $a$.
"""
function coordinates(a::NfRelOrdElem; copy::Bool = true)
  assure_has_coord(a)
  if copy
    return deepcopy(a.coordinates)
  else
    return a.coordinates
  end
end

################################################################################
#
#  String I/O
#
################################################################################

function show(io::IO, a::NfRelOrdElem)
  print(io, a.elem_in_nf)
end

################################################################################
#
#  Relative Norm and Trace
#
################################################################################

norm(a::NfRelOrdElem, k::Union{NumField, FlintRationalField }) = base_ring(parent(a))(norm(a.elem_in_nf, k))
tr(a::NfRelOrdElem, k::Union{NumField, FlintRationalField }) = base_ring(parent(a))(tr(a.elem_in_nf, k))

################################################################################
#
#  Conversion
#
################################################################################

(K::NfRel)(a::NfRelOrdElem) = elem_in_nf(a)

(K::NfRelNS)(a::NfRelOrdElem) = elem_in_nf(a)

################################################################################
#
#  Representation matrix
#
################################################################################

function representation_matrix(a::NfRelOrdElem)
  O = parent(a)
  A = representation_matrix(elem_in_nf(a))
  A = basis_matrix(O, copy = false)*A*basis_mat_inv(O, copy = false)
  return A
end


text:
module TestExamples1DAdvection

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "1d")

@testset "Linear scalar advection" begin
  @testset "elixir_advection_basic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_basic.jl"),
      l2   = [6.0388296447998465e-6],
      linf = [3.217887726258972e-5])
  end

  @testset "elixir_advection_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_amr.jl"),
      l2   = [0.3540206249507417],
      linf = [0.9999896603382347])
  end

  @testset "elixir_advection_amr_nonperiodic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_amr_nonperiodic.jl"),
      l2   = [4.283508859843524e-6],
      linf = [3.235356127918171e-5])
  end
end

end # module


text:
function connectedGreen(para, hop::Vector{BareHoppingId}, subdiagram = false; name = Symbol("Gc$(length(hop)*2)"), resetuid = false, even = true)
    # @assert para.diagType == GreenNDiag
    # @assert length(extT) == length(orbital) == length(site)
    # println("wip")
    # println(hop)
    N = length(hop)

    resetuid && uidreset()

    Gc = []

    # for paired Green's function, odd number of legs always leads to zero
    # if even && (length(site) % 2 == 1)
    #     return nothing
    # end

    Gfull = fullGreen(para, hop, true; resetuid = false, even = even)
    if isnothing(Gfull)
        return nothing
    end
    push!(Gc, Gfull)

    for (lind, rind) in partitions(collect(1:N), 2)
        #this partition will not generate fermionic sign because the hopping term is always a bosonic operator
        # if even && (length(lind) % 2 == 1)
        #     continue
        # end
        subGc = connectedGreen(para, hop[lind], true; resetuid = false, even = even)
        subGn = fullGreen(para, hop[rind], true; resetuid = false, even = even)
        if isnothing(subGn) || isnothing(subGc)
            continue
        end
        push!(Gc, Diagram(GenericId(para), Prod(), [subGc, subGn], factor = -1.0)) #additional minus sign because Gc(s) = Gn(s) - \sum_o Gc(o)Gn(s-o)
    end

    extT, orbital, site, creation = [], [], [], []
    for h in hop
        append!(extT, h.extT)
        append!(site, h.site)
        append!(creation, [true, false])
        append!(orbital, h.orbital)
    end
    return Diagram(ConnectedGreenNId(para, orbital = orbital, t = extT, r = site, creation = creation), Sum(), Gc, name = name)
end

# function connectedGreen(para, site::AbstractVector, orbital::AbstractVector, extT::AbstractVector = collect(1:length(orbital)), subdiagram = false; name = Symbol("Gc$(length(site))"), resetuid = false, even = true)
#     # @assert para.diagType == GreenNDiag
#     @assert length(extT) == length(orbital) == length(site)
#     N = length(extT)

#     resetuid && uidreset()

#     Gc = []

#     # for paired Green's function, odd number of legs always leads to zero
#     if even && (length(site) % 2 == 1)
#         return nothing
#     end

#     Gfull = fullGreen(para, site, orbital, extT, true; resetuid = false, even = even)
#     if isnothing(Gfull)
#         return nothing
#     end
#     push!(Gc, Gfull)

#     for (lind, rind) in partitions(collect(1:N), 2)
#         if even && (length(lind) % 2 == 1)
#             continue
#         end
#         lR, lo, lt = site[lind], orbital[lind], extT[lind]
#         rR, ro, rt = site[rind], orbital[rind], extT[rind]
#         subGc = connectedGreen(para, lR, lo, lt, true; resetuid = false, even = even)
#         subGn = fullGreen(para, rR, ro, rt, true; resetuid = false, even = even)
#         if isnothing(subGn) || isnothing(subGc)
#             continue
#         end
#         p = parity(vcat(lind, rind))
#         push!(Gc, Diagram(GenericId(para), Prod(), [subGc, subGn], factor = -p)) #additional minus sign because Gc(s) = Gn(s) - \sum_o Gc(o)Gn(s-o)
#     end

#     return Diagram(ConnectedGreenNId(para, orbital, extT, site), Sum(), Gc, name = name)
# end


text:
using Oceananigans.Grids
using Oceananigans.Utils: prettytime, hours, day, days, years

using Statistics
using JLD2
using Printf
using CairoMakie

output_prefix = "near_global_lat_lon_1440_600_1_fine_surface"
filepath = output_prefix * ".jld2"
file = jldopen(filepath)

Nx = file["grid/underlying_grid/Nx"]
Ny = file["grid/underlying_grid/Ny"]
Lλ = file["grid/underlying_grid/Lx"]
Lφ = file["grid/underlying_grid/Ly"]
Lz = file["grid/underlying_grid/Lz"]

grid = LatitudeLongitudeGrid(size = (Nx, Ny, 1),
                             longitude = (-180, 180),
                             latitude = (-Lφ/2, Lφ/2),
                             z = (-Lz, 0))

x, y, z = nodes((Center, Center, Center), grid)

bottom = Float32.(file["grid/immersed_boundary/bottom_height"][4:end-3,3:end-3,1])
bottom[ bottom .>0 ] .=  NaN
bottom[ bottom .<0 ] .= 0.0

iter = Observable(0)
iters = parse.(Int, keys(file["timeseries/t"]))
ζ′ = @lift file["timeseries/ζ/" * string($iter)][:, :, 1]
title = @lift(@sprintf("Surface Vorticity in Hydrostatic Model at time = %s", prettytime(file["timeseries/t/" * string($iter)])))
fig = Figure(resolution = (2000, 1000))
ax = Axis(fig[1,1], xlabel = "longitude", ylabel = "latitude", title=title)
heatmap_plot = heatmap!(ax, ζ′, colormap=:balance, colorrange=(-1e-6, 1e-6), nan_color=:black)
Colorbar(fig[1,2], heatmap_plot , width=25)
display(fig)

record(fig, output_prefix * ".mp4", iters[2:end], framerate=6) do i
    @info "Plotting iteration $i of $(iters[end])..."
    iter[] = i
end


text:
struct BinomialVanillaEngine{P <: AbstractBlackScholesProcess, T <: BinomialTreeType} <: AbstractVanillaEngine
  process::P
  timeSteps::Int
  treeClass::Type{T}
end

function _calculate!(pe::BinomialVanillaEngine, opt::VanillaOption)
  rfdc = pe.process.riskFreeRate.dc
  divdc = pe.process.dividendYield.dc
  voldc = pe.process.blackVolatility.dc
  volcal = pe.process.blackVolatility.calendar

  s0 = state_variable(pe.process).value
  s0 > 0.0 || error("negative or null underlying given")

  maturityDate = opt.exercise.dates[end]
  v = black_vol(pe.process.blackVolatility, maturityDate, s0)
  r = zero_rate(pe.process.riskFreeRate, maturityDate, rfdc, ContinuousCompounding(), NoFrequency()).rate
  q = zero_rate(pe.process.dividendYield, maturityDate, divdc, ContinuousCompounding(), NoFrequency()).rate

  refDate = reference_date(pe.process.riskFreeRate)

  # binomial trees with constant coefficient
  flatRiskFree = FlatForwardTermStructure(refDate, r, rfdc)
  flatDividends = FlatForwardTermStructure(refDate, q, divdc)
  flatVol = BlackConstantVol(refDate, volcal, v, voldc)

  payoff = opt.payoff

  maturity = year_fraction(rfdc, refDate, maturityDate)

  bs = GeneralizedBlackScholesProcess(state_variable(pe.process), flatRiskFree, flatDividends, flatVol)

  grid = TimeGrid(maturity, pe.timeSteps)

  tree = pe.treeClass(bs, maturity, pe.timeSteps, payoff.strike)
  lattice = BlackScholesLattice(tree, r, maturity, pe.timeSteps)

  option = DiscretizedVanillaOption(opt, pe.process, grid, lattice.treeLattice)

  initialize!(option, lattice.treeLattice, maturity)
  rollback!(option, grid[3])

  va2 = copy(option.common.values)
  length(va2) == 3 || error("Expect 3 nodes in grid at second step")

  p2d, p2m, p2u = va2

  s2u = get_underlying(lattice, 3, 3) # up price
  s2m = get_underlying(lattice, 3, 2) # middle price
  s2d = get_underlying(lattice, 3, 1) # down (low) price

  # calculate gamma by taking the first derivative of the two deltas
  delta2u = (p2u - p2m) / (s2u - s2m)
  delta2d = (p2m - p2d) / (s2m - s2d)
  gamma = (delta2u - delta2d) / ((s2u - s2d) / 2)

  # rollback to the second-last step, get option vals at this point
  rollback!(option, grid[2])
  va = copy(option.common.values)
  length(va) == 2 || error("Expect 2 nodes in grid at first step")

  p1d, p1u = va

  s1u = get_underlying(lattice, 2, 2) # up (high) price
  s1d = get_underlying(lattice, 2, 1) # down (low) price

  delta = (p1u - p1d) / (s1u - s1d)

  # finally, rollback to t0
  rollback!(option, 0.0)

  p0 = present_value(option)

  opt.results.value = p0
  opt.results.delta = delta
  opt.results.gamma = gamma
  opt.results.theta = black_scholes_theta(pe.process, opt.results.value, opt.results.delta, opt.results.gamma)

  return pe, opt
end


text:
function matrix_chain(n)

  reps = parse(Int64, ENV["LAMP_REPS"])
  m = Int(floor(n / 5))

  A = randn(n,n)
  B = randn(n,n)
  C = randn(m,n)

  Benchmarker.add_data(csv, "mc_l_r_orig", Benchmarker.measure(reps, left_to_right_chain_original, A, B, C))
  Benchmarker.add_data(csv, "mc_l_r_guid", Benchmarker.measure(reps, left_to_right_chain_guided, A, B, C))

  @test isapprox(left_to_right_chain_original(A, B, C)[2], left_to_right_chain_guided(A, B, C)[2])

  C = randn(n,m)
  Benchmarker.add_data(csv, "mc_r_l_orig", Benchmarker.measure(reps, right_to_left_chain_original, A, B, C))
  Benchmarker.add_data(csv, "mc_r_l_guid", Benchmarker.measure(reps, right_to_left_chain_guided, A, B, C))

  @test isapprox(right_to_left_chain_original(A, B, C)[2], right_to_left_chain_guided(A, B, C)[2])

  A = randn(n,n)
  B = randn(n,m)
  C = randn(m,n)
  D = randn(n,n)

  Benchmarker.add_data(csv, "mc_mixed_orig", Benchmarker.measure(reps, mixed_chain_original, A, B, C, D))
  Benchmarker.add_data(csv, "mc_mixed_guid", Benchmarker.measure(reps, mixed_chain_guided, A, B, C, D))

  @test isapprox(mixed_chain_original(A, B, C, D)[2], mixed_chain_guided(A, B, C, D)[2])
end

function left_to_right_chain_original(A::DenseArray{Float64, 2}, B::DenseArray{Float64, 2}, C::DenseArray{Float64, 2})
  result = zeros(size(C,1), size(B,2))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = C * A * B

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function left_to_right_chain_guided(A::DenseArray{Float64, 2}, B::DenseArray{Float64, 2}, C::DenseArray{Float64, 2})
  result = zeros(size(C,1), size(B,2))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = (C * A) * B

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function right_to_left_chain_original(A::DenseArray{Float64, 2}, B::DenseArray{Float64, 2}, C::DenseArray{Float64, 2})
  result = zeros(size(A,1), size(C,2))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = A * B * C

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function right_to_left_chain_guided(A::DenseArray{Float64, 2}, B::DenseArray{Float64, 2}, C::DenseArray{Float64, 2})
  result = zeros(size(A,1), size(C,2))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = A * (B * C)

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function mixed_chain_original(A::DenseArray{Float64, 2}, B::DenseArray{Float64, 2}, C::DenseArray{Float64, 2}, D::DenseArray{Float64, 2})
  result = zeros(size(A,1), size(D,2))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = A * B * C * D

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end

function mixed_chain_guided(A::DenseArray{Float64, 2}, B::DenseArray{Float64, 2}, C::DenseArray{Float64, 2}, D::DenseArray{Float64, 2})
  result = zeros(size(A,1), size(D,2))
  start::Float64 = 0.0
  finish::Float64 = 0.0
  Benchmarker.cachescrub()
  Benchmarker.gcscrub()
  start = time_ns()

  result = (A * B) * (C * D)

  finish = time_ns()
  return ((finish-start)*1e-9, result)
end


text:
"""
    WSymbol(name::Union{String, Symbol})
    W"..."

A Wolfram language symbol. The `W""` string macro can be used as a short form.

A `WSymbol` object is callable, which can be used to construct [`WExpr`](@ref)s (but doesn't evaluate it):
```julia
julia> W"Sin"
W"Sin"

julia> W"Sin"(1.0)
W"Sin"(1.0)

julia> weval(W"Sin"(1.0))
0.8414709848078965
```

Keyword arguments are passed as options:
```julia
julia> ex = W`Sqrt[x^2]`
W"Sqrt"(W"Power"(W"x", 2))

julia> assume = W`x<0`
W"Less"(W"x", 0)

julia> weval(W"Simplify"(ex, Assumptions=assume))
W"Times"(-1, W"x")
```
"""
struct WSymbol
    name::String
end
WSymbol(sym::Symbol) = WSymbol(String(sym))
Base.print(io::IO, s::WSymbol) = print(io, s.name)
Base.show(io::IO, s::WSymbol) = print(io, 'W', '"', s.name, '"')
Base.:(==)(a::WSymbol, b::WSymbol) = a.name == b.name

macro W_str(str)
    WSymbol(str)
end

"""
    WReal(str::String)

A Wolfram arbitrary-precision real number.
"""
struct WReal
    value::String
end
Base.show(io::IO, x::WReal) = print(io, x.value)
Base.:(==)(a::WReal, b::WReal) = a.value == b.value

"""
    WInteger(str::String)

A Wolfram arbitrary-precision integer.
"""
struct WInteger
    value::String
end
Base.show(io::IO, x::WInteger) = print(io, x.value)
Base.:(==)(a::WInteger, b::WInteger) = a.value == b.value

"""
    WExpr(head, args)

A Wolfram language expression. Like [`WSymbol`](@ref) it is callable to construct more complicated expressions.

```julia
julia> W`Function[x,x+1]`
W"Function"(W"x", W"Plus"(W"x", 1))

julia> W`Function[x,x+1]`(2)
W"Function"(W"x", W"Plus"(W"x", 1))(2)

julia> weval(W`Function[x,x+1]`(2))
3
```
"""
struct WExpr
    head
    args
end
function Base.:(==)(a::WExpr, b::WExpr)
    a.head == b.head || return false
    length(a.args) == length(b.args) || return false
    for (aa,bb) in zip(a.args, b.args)
        aa == bb || return false
    end
    return true
end

function Base.show(io::IO, w::WExpr)
    show(io, w.head)
    print(io, '(')
    isfirst = true
    for arg in w.args
        if !isfirst
            print(io, ", ")
        else
            isfirst = false
        end
        show(io, arg)
    end
    print(io, ')')
end

function (w::WSymbol)(args...; kwargs...)
    if !isempty(kwargs)
        args = [args..., [W"Rule"(WSymbol(k), v) for (k,v) in kwargs]...]
    end
    WExpr(w, args)
end
function (w::WExpr)(args...; kwargs...)
    if !isempty(kwargs)
        args = [args..., [W"Rule"(WSymbol(k), v) for (k,v) in kwargs]...]
    end
    WExpr(w, args)
end


text:
using Each
using Test

@testset "Each.jl" begin
    @test sum.(each(Tuple{Vararg{Number}},((1,1),(1,2),(1,3)))) == (2,3,4)
    @test sum.(each(Tuple{Vararg{Number}},(1,1))) == (2,)
    @test_throws ErrorException each(String,((1,2),(3,4)))
    @test eltype(each(String,"something")) == String
end


text:
####### IterData
@with_kw mutable struct IterData{ XT, YT, DT } <: AbstractIterData{XT,YT,DT}
    x :: XT = MIN_PRECISION[]
    fx :: YT = MIN_PRECISION[]
    Δ :: DT = zero(MIN_PRECISION)

    x_index :: Int = -1
    num_iterations :: Int = 0
    num_model_improvements :: Int = 0
    it_stat :: ITER_TYPE = SUCCESSFULL
end

# getters
get_x( id :: IterData ) = id.x
get_fx( id :: IterData ) = id.fx
get_delta( id :: IterData ) = id.Δ

get_x_index( id:: IterData ) = id.x_index

it_stat( id :: IterData ) = id.it_stat

# setters
function _set_x!( id :: IterData, x :: Vec ) :: Nothing
    id.x = x
    return nothing
end

function set_fx!( id :: IterData, y :: Vec ) :: Nothing
    id.fx = y
    return nothing
end

function set_delta!( id :: IterData, Δ :: NumOrVec ) :: Nothing
    id.Δ = Δ
    return nothing
end

function set_x_index!( id :: IterData, N :: Int ) :: Nothing
    id.x_index = N;
    nothing
end

get_num_iterations( id :: IterData ) :: Int = id.num_iterations
get_num_model_improvements( id :: IterData ) :: Int = id.num_model_improvements

function set_num_iterations!( id :: IterData, N :: Int = 0 ) :: Nothing
    id.num_iterations = N
    return nothing
end

function set_num_model_improvements!( id :: IterData, N :: Int = 0 ) :: Nothing
    id.num_model_improvements = N
    return nothing
end

function it_stat!( id :: IterData, t :: ITER_TYPE ) :: Nothing
    id.it_stat = t
    return nothing
end

function init_iter_data( ::Type{<:IterData}, x :: Vec, fx :: Vec, Δ :: NumOrVec )
    return IterData(; x, fx, Δ )
end

####### IterSaveable

# Note: I don't save x,fx, so i set some arbritrary types for XT and YT
# Also, I store all additional meta data as Float64, input gets automatically converted.
struct IterSaveable{
        DT <: NumOrVecF, C <: ContainerSaveable } <: AbstractIterSaveable{Vec64,Vec64,DT,C}

    Δ :: DT

    x_index :: Int
    num_iterations :: Int
    num_model_improvements
    it_stat :: ITER_TYPE

    # additional information for stamping
    x_trial_index :: Int
    ρ :: Float64
    stepsize :: Float64
    ω :: Float64

    model_meta :: C
end

function get_saveable( :: Type{<:IterSaveable}, id :: AbstractIterData;
    x_trial_index, ρ, stepsize, ω, sc = nothing, kwargs... )
    return IterSaveable(
        get_delta(id),
        get_x_index(id),
        get_num_iterations(id),
        get_num_model_improvements(id),
        it_stat(id),
        x_trial_index, ρ,stepsize, ω,
        get_saveable(sc)
    )
end


text:
using SeisIO, SeisIO.Quake, SeisIO.SeisHDF
import Printf

# US FDSN example: 5 stations, 2 networks, all channels, last 600 seconds
println(stdout, "Some real data acquisition examples...\n\n")

println(stdout, "preliminary: change directory")
path = Base.source_dir()
println(stdout, "cd(", path, ")")
cd(path)

printstyled(stdout, "FDSN get_data\n", color=:green, bold=true)
CHA = "CC.PALM, UW.HOOD, CC.TIMB, CC.HIYU, UW.TDH"
s = -600
t = u2d(time())
printstyled(stdout, "Command: ", color=:green)
println(stdout, "S_fdsn = get_data(\"FDSN\", \"", CHA, "\", src=\"IRIS\", s=", s, ", t=", t, ")")
S_fdsn = get_data("FDSN", CHA, src="IRIS", s=s, t=t)
printstyled(stdout, "Results: ", color=:green)
println(stdout, "S_fdsn")
show(S_fdsn)

# IRIS example: 6 channels, 30 minutes
printstyled(stdout, "\n\nIRIS get_data\n", color=:green, bold=true)
STA = ["CC.TIMB..BHE", "CC.TIMB..BHN", "CC.TIMB..BHZ", "UW.HOOD..HHE", "UW.HOOD..HHN", "UW.HOOD..HHZ"]
st = -3600
en = -1800
printstyled(stdout, "Command: ", color=:green)
println(stdout, "S_iris = get_data(\"IRIS\", ", STA, ", s=", st, ", t=", en, ")")
S_iris = get_data("IRIS", STA, s=st, t=en)
printstyled(stdout, "Results: ", color=:green)
println(stdout, "S_iris")
show(S_iris)

# The Tohoku-Oki great earthquake, from IRIS FDSN, recorded by boreholes in WA (USA)
printstyled(stdout, "\n\nFDSNevt\n", color=:green, bold=true)
printstyled(stdout, "Command: ", color=:green)
println(stdout, "S_evt = FDSNevt(\"201103110547\", \"PB.B004..EH?,PB.B004..BS?,PB.B001..BS?,PB.B001..EH?\")")
S_evt = FDSNevt("201103110547", "PB.B004..EH?,PB.B004..BS?,PB.B001..BS?,PB.B001..EH?")
printstyled(stdout, "Results: ", color=:green)
println(stdout, "S_evt")
show(S_evt)

# SeisComp3 SeedLink session, IRIS server, TIME mode
printstyled(stdout, "\n\nSeedLink\n", color=:green, bold=true)
sta = "UW.GRUT,UW.H1K,UW.MDW"
s1 = -120
t1 = 120

printstyled(stdout, "Commands: ", color=:green)

println(stdout, "S_sl = seedlink(\"TIME\", \"", sta, "\", s=", s1, ", t=", t1, ")")
S_sl = seedlink("TIME", sta, s=s1, t=t1)

println(stdout, "          seedlink!(S_sl, \"DATA\", \"SampleFiles/SL_long_test.conf\")")
seedlink!(S_sl, "DATA", "SampleFiles/SL_long_test.conf")

println(stdout, "          sleep(30)")
sleep(30)

println(stdout, "          for conn in S_sl.c; close(conn); end")
for conn in S_sl.c; close(conn); end

printstyled(stdout, "Results: ", color=:green)
println(stdout, "S_sl")
show(S_sl)

printstyled(stdout, "\n\n\nNote: ", color=:white, bold=true)
println(stdout, "ALL data from these examples can be written to file.")
printstyled(stdout, "wseis(\"fname.seis\", S)", color=7)
println(stdout, "       write S to low-level SeisIO native format in file fname.seis")
printstyled(stdout, "writesac(S)", color=7)
println(stdout, "                  write S to SAC files with auto-generated names")
printstyled(stdout, "write_hdf5(\"fname.h5\", S)", color=7)
println(stdout, "    write S to ASDF (HDF5) file \"fname.h5\"")


text:
# Test _all_function_variables
@testset "_all_function_variables" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 0 <= par <= 1)
    @infinite_variable(m, inf(par))
    @point_variable(m, inf(0), pt)
    @hold_variable(m, hold)
    meas = MeasureRef(m, 1)
    m.meas_to_name[1] = "meas"
    # test for variable reference
    @testset "Variable" begin
        @test InfiniteOpt._all_function_variables(par) == [par]
        @test InfiniteOpt._all_function_variables(inf) == [inf]
        @test InfiniteOpt._all_function_variables(meas) == [meas]
    end
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = meas + 2par + hold
        aff2 = zero(GenericAffExpr{Float64, GeneralVariableRef})
        # test expressions
        @test InfiniteOpt._all_function_variables(aff1) == [meas, par, hold]
        @test InfiniteOpt._all_function_variables(aff2) == GeneralVariableRef[]
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = pt^2 + inf * pt - meas + 2par + hold
        quad2 = pt^2 + inf * pt
        quad3 = zero(GenericQuadExpr{Float64, GeneralVariableRef})
        # test expressions
        @test InfiniteOpt._all_function_variables(quad1) == [meas, par, hold, pt, inf]
        @test InfiniteOpt._all_function_variables(quad2) == [pt, inf]
        @test InfiniteOpt._all_function_variables(quad3) == GeneralVariableRef[]
    end
    # test backup
    @testset "Fallback" begin
        @variable(Model(), x)
        @test_throws ErrorException InfiniteOpt._all_function_variables(x)
    end
end

# Test comparisons
@testset "Comparisons" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 0 <= par <= 1)
    @infinite_parameter(m, 0 <= par2 <= 1)
    @infinite_variable(m, inf(par))
    @infinite_variable(m, inf2(par, par2))
    @point_variable(m, inf(0), pt)
    @hold_variable(m, hold)
    red = InfiniteOpt.ReducedInfiniteVariableRef(m, 42)
    # test AffExpr comparison
    @testset "Base.:(==) AffExpr" begin
        @test par + par2 + inf - 2 == par + (par2 + inf) - 2
        @test 0.25red - inf == 0.25red - inf
        @test red + 3 - inf != par + red
        @test red + 3 - inf != red + 3 - hold
    end
    # test QuadExpr comparison
    @testset "Base.:(==) QuadExpr" begin
        @test par * inf + par2 + inf - 2 == par * inf + (par2 + inf) - 2
        @test red * inf - inf == red * inf - inf
        @test red * inf + 3 - inf != par * inf2 + red
        @test par * par2 + inf * inf2 == par * par2 + inf * inf2
        @test par * par2 + inf * inf2 != par * par2
        @test par * par2 + 2 * inf * inf2 !=  par * par2 + inf * inf2
        @test par * inf + par2 + inf - 2 != par * inf + (par2 + inf) - 3
    end
end

# Test _all_parameter_refs
@testset "_all_parameter_refs" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 0 <= par <= 1)
    @infinite_parameter(m, 0 <= par2 <= 1)
    @infinite_variable(m, inf(par))
    @infinite_variable(m, inf2(par, par2))
    @point_variable(m, inf(0), pt)
    @hold_variable(m, hold)
    red = InfiniteOpt.ReducedInfiniteVariableRef(m, 42)
    m.reduced_info[42] = ReducedInfiniteInfo(inf2, Dict(1 => 0.5))
    # test for finite variable reference
    @testset "FiniteVariable" begin
        @test InfiniteOpt._all_parameter_refs(pt) == ()
        @test InfiniteOpt._all_parameter_refs(hold) == ()
    end
    # test for infinite variable reference
    @testset "InfiniteVariable" begin
        @test InfiniteOpt._all_parameter_refs(inf) == (par, )
    end
    # test for parameter reference
    @testset "Parameter" begin
        @test InfiniteOpt._all_parameter_refs(par) == (par, )
    end
    # test for reduced infinite variable reference
    @testset "ReducedInfinite" begin
        @test InfiniteOpt._all_parameter_refs(red) == (par2, )
    end
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = inf + inf2 + pt - 3
        aff2 = pt + hold - 2
        # test expressions
        @test InfiniteOpt._all_parameter_refs(aff1) == (par, par2)
        @test InfiniteOpt._all_parameter_refs(aff2) == ()
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = inf * inf2 + inf + inf2 + pt - 3 - par
        quad2 = pt * pt + pt + hold - 2
        # test expressions
        @test InfiniteOpt._all_parameter_refs(quad1) == (par, par2)
        @test InfiniteOpt._all_parameter_refs(quad2) == ()
    end
end

# Test _remove_variable
@testset "_remove_variable" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 0 <= par <= 1)
    @infinite_variable(m, inf(par))
    @point_variable(m, inf(0), pt)
    @hold_variable(m, hold)
    meas = MeasureRef(m, 1)
    m.meas_to_name[1] = "meas"
    # test for GenericAffExpr
    @testset "AffExpr" begin
        # make expressions
        aff1 = meas + 2par + hold
        aff2 = zero(GenericAffExpr{Float64, GeneralVariableRef})
        # test expressions
        @test isa(InfiniteOpt._remove_variable(aff1, hold), Nothing)
        @test !haskey(aff1.terms, hold)
        @test isa(InfiniteOpt._remove_variable(aff2, inf), Nothing)
    end
    # test for GenericQuadExpr
    @testset "QuadExpr" begin
        # make expressions
        quad1 = pt^2 + inf * pt - meas + 2par + hold
        quad2 = zero(GenericQuadExpr{Float64, GeneralVariableRef})
        # test expressions
        quad = copy(quad1)
        @test isa(InfiniteOpt._remove_variable(quad, inf), Nothing)
        @test !haskey(quad.aff.terms, inf)
        @test !haskey(quad.terms, UnorderedPair{GeneralVariableRef}(inf, pt))
        quad = copy(quad1)
        @test isa(InfiniteOpt._remove_variable(quad, pt), Nothing)
        @test !haskey(quad.terms, UnorderedPair{GeneralVariableRef}(inf, pt))
        @test !haskey(quad.terms, UnorderedPair{GeneralVariableRef}(pt, pt))
        @test isa(InfiniteOpt._remove_variable(quad2, inf), Nothing)
    end
end

# Test destructive_add!
# @testset "JuMP.destructive_add!" begin
#     # initialize model and references
#     m = InfiniteModel()
#     @infinite_parameter(m, 0 <= par <= 1)
#     @infinite_parameter(m, 0 <= par2 <= 1)
#     @infinite_variable(m, inf(par))
#     @infinite_variable(m, inf2(par, par2))
#     @hold_variable(m, hold)
#     # test variable, constant, variable
#     @testset "Var, Constant, Var" begin
#         @test destructive_add!(inf, 2., inf2) == inf + 2inf2
#         @test destructive_add!(inf, 2., hold) == inf + 2hold
#     end
# end

# Test _set_variable_coefficient!
@testset "_set_variable_coefficient!" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, 0 <= par <= 1)
    @infinite_parameter(m, 0 <= par2 <= 1)
    @infinite_variable(m, x(par))
    @infinite_variable(m, y(par, par2))
    @hold_variable(m, z)
    # test with GeneralVariableRef
    @testset "GeneralVariableRef" begin
        @test InfiniteOpt._set_variable_coefficient!(x, x, 2) == 2 * x
        @test InfiniteOpt._set_variable_coefficient!(z, x, 2) == z + 2x
    end
    # test with GenericAffExpr
    @testset "AffExpr" begin
        @test InfiniteOpt._set_variable_coefficient!(x + z, x, 2) == 2x + z
        @test InfiniteOpt._set_variable_coefficient!(x + z, y, 2) == x + z + 2y
    end
    # test with GenericQuadExpr
    @testset "QuadExpr" begin
        @test InfiniteOpt._set_variable_coefficient!(y ^2 + x, x, 2) == y^2 + 2x
        @test InfiniteOpt._set_variable_coefficient!(y^2 + x, y, 2) == y^2 + x + 2y
    end
    # test fallabck
    @testset "Fallback" begin
        @test_throws ErrorException InfiniteOpt._set_variable_coefficient!(2, x, 2)
    end
end


text:
using Garch
using Documenter

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

makedocs(;
    modules=[Garch],
    authors="banachtech <balaji@banach.tech> and contributors",
    repo="https://github.com/banachtech/Garch.jl/blob/{commit}{path}#{line}",
    sitename="Garch.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://banachtech.github.io/Garch.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/banachtech/Garch.jl",
    devbranch="main",
)


text:
# ECC math utils implementation in Julia
# Copyright (C): Jaka Smrekar (vinctux) <vinctux@outlook.com>, 2016
#
# PyECC (https://github.com/amintos/PyECC) as Pythonic reference

importall Base
using Krypto

immutable CurveFP
    p::Integer
    a::Integer
    b::Integer
    function CurveFP(p, a, b)
        assert(4 * a^3 + 27 * b^2 % p != 0)
        new(p, a, b)
    end
end

# Slow addition
function +(C::CurveFP, p1, p2)
    x1, y1 = p1
    x2, y2 = p2
    if (x1 - x2) % n != 0
        s = ((y1 - y2) * ~(x1 - x2, n)) % n
        x = (s * s - x1 - x2) % n
        return (x, n - (y1 + s * (x - x1)) % n)
    elseif (y1 + y2) % n != 0
        s = ((3 * x1 * x1 - p) * ~(2 * y1, n)) % n
        x = (s * s - 2 * x1) % n
        return (x, n - (y1 + s * (x - x1)) % n)
    end
end

# Fast addition
function ++(C::CurveFP, jp1, jp2)
    x1, y1, z1, z1s, z1c = jp1
    x2, y2, z2, z2s, z2c = jp2
    s1 = (y1 * z2c) % C.p
    s2 = (y2 * z1c) % C.p
    u1 = (x1 * z2s) % C.p
    u2 = (x2 * z1s) % C.p
    if (u1 - u2) % n
        h = (u2 - u1) % C.p
        r = (s2 - s1) % C.p
        hs = (h * h) % C.p
        hc = (hs * h) % C.p
        x3 = (-hc - 2 * u1 * hs + r * r) % C.p
        y3 = (-s1 * hc + r * (u1 * hs - x3)) % C.p
        z3 = (z1 * z2 * h) % C.p
        z3s = (z3 * z3) % C.p
        z3c = (z3s * z3) % C.p
        return (x3, y3, z3, z3s, z3c)
    else
        if (s1 + s2) % n == 0 return doublef(C, jp1)
        else error("Something went wrong.") end
    end
end

function doublef(C::CurveFP, jp)
    x1, y1, z1, z1p2, z1p3 = jp
    y1p2 = (y1 * y1) % n
    a = (4 * x1 * y1p2) % n
    b = (3 * x1 * x1 - p * z1p3 * z1) % n
    x3 = (b * b - 2 * a) % n
    y3 = (b * (a - x3) - 8 * y1p2 * y1p2) % n
    z3 = (2 * y1 * z1) % n
    z3p2 = (z3 * z3) % n
    return x3, y3, z3, z3p2, (z3p2 * z3) % n
end

# Scalar multiplication p1 * c = p1 + p1 + ... + p1 (c times) in O(log(n))
function *(C::CurveFP, p1, c)
    r = 0
    while c > 0
        if c & 1 != 0 res = C + res + p1 end
        c >>= 1
        p1 = C + p1 + p1
    end
    return res
end

# Optimized scalar multiplication
function mulf(C::CurveFP, jp1, c)
    res = 0
    jp0 = ~(jp1, n)
    for s in signbin(c)
        res = doublef(C, res)
        if s != 0 res = s > 0 ? C ++ res ++ jp1 : C ++ res ++ jp0 end
    end
    return res
end

# Hyper-optimized scalar multiplication
function muladdf(C::CurveFP, jp1, c1, jp2, c2)
    s1 = signbin(c1)
    s2 = signbin(c2)
    diff = len(s2) - len(s1)
    if diff > 0 s1 = vcat([0 for i in 1:(diff)], s1)
    elseif diff < 0 s2 = vcat([0 for i in 1:(-diff)], s2) end
    jp1p2 = C ++ jp1 ++ jp2
    jp1n2 = C ++ jp1 ++ ~(jp2, n)
    precomp = ((0, jp2, ~(jp2, n)),
               (jp1, jp1p2, jp1n2),
               (~(jp1, n),  ~(jp1n2, n), ~(jp1p2, n)))
    res = 0
    for (i, j) in zip(s1, s2)
        res = doublef(C, res)
        if i != 0 || j != 0 res = C ++ res ++ precomp[i][j] end
    end
    return res
end

# Hyper-optimized scalar multiplication for xy-coordinates
muladdp(C::CurveFP, p1, c1, p2, c2) = proj2norm(muladdf(C, norm2proj(p1), c1, norm2proj(p2), c2), n)

# Multiply point p by c using fast multiplication
mulp(C::CurveFP, p1, c) = proj2norm(mulf(C, norm2proj(p1), c), n)

# Find curve parameter q mod n having point (x, y) and parameter p
curve_q(x, y, p, n) = ((x * x - p) * x - y * y) % n

# Test, whether the given point is on the curve (p, q, n)
in(p, C::CurveFP) = p != 0 ? (p[1]^3 - p * p[1] - q) % n == y^2 % n : true

# Transform point p given as (x, y) to projective coordinates
norm2proj(p) = p != 0 ? (p[0], p[1], 1, 1, 1) : 0

# Transform a point from projective coordinates to (x, y) mod n
proj2norm(jp, n) = jp != 0 ? ((jp[0] * imod(jp[3], n)) % n, (jp[1] * imod(jp[4], n)) % n) : 0

# Compute the inverse point to p in any coordinate system
~(p, n) = p != 0 ? (p[0], (n - p[1]) % n) + p[2:end] : 0


text:
import Celeste: SEP
using Base.Test
using FITSIO

@testset "sep" begin

    # use test image in SEP source directory
    sep_testdata_dir = joinpath(dirname(@__FILE__), "..", "deps", "src",
                                "sep-1.0.1", "data")

    data = read(FITS(joinpath(sep_testdata_dir, "image.fits"))[1])
    back_sextractor = read(FITS(joinpath(sep_testdata_dir, "back.fits"))[1])
    rms_sextractor = read(FITS(joinpath(sep_testdata_dir, "rms.fits"))[1])

    # test background
    bkg = SEP.Background(data)
    @test collect(bkg) ≈ back_sextractor

    for T in (Float32, Float64)
        A = zeros(T, size(data))
        back = collect(T, bkg)
        @test back ≈ back_sextractor
    end

    # test that broadcast subtraction methods work and return same result
    A = zeros(Float32, size(data))
    A .-= bkg
    B = zeros(Float32, size(data)) .- bkg
    C = zeros(Float32, size(data)) - bkg
    @test A == B == C

    # test source extraction
    data .-= bkg
    catalog = SEP.extract(data, 3.0; noise=SEP.rms(bkg))
    @test length(catalog.x) == 39
end


text:
using .Gadfly
using Colors

function Gadfly.plot(
    zep::Zeppelin,
    rows;
    autoklms = true,
    klms = [],
    xmin = 0.0,
    xmax = missing,
    norm = :None,
    yscale = 1.05,
    ytransform = identity,
    style = NeXLSpectrum.NeXLSpectrumStyle,
    palette = NeXLCore.NeXLPalette,
)
    plot(
        Spectrum[filter(s->!ismissing(s), map(r -> spectrum(zep, r, false),rows))...],
        autoklms = autoklms,
        klms = klms,
        xmin = xmin,
        xmax = xmax,
        norm = norm,
        yscale = yscale,
        ytransform = ytransform,
        style = style,
        palette = palette,
    )
end

function Gadfly.plot(ai::AlignIntermediary, df, df2, idx=1, xs=:x, ys=:y)
    df2p = align(df2, ai, idx, xs, ys)
    xex, yex, x2ex, y2ex = extrema(df[:,xs]), extrema(df[:,ys]), extrema(df2p[:,xs]), extrema(df2p[:,ys])
    xmin, xmax = max(xex[1], x2ex[1]), min(xex[2], x2ex[2])
    ymin, ymax = max(yex[1], y2ex[1]), min(yex[2], y2ex[2])
    extra = 0.1*max(xmax-xmin, ymax-ymin)
    plot(
        layer(x=xs, y=ys, color=[ colorant"red" ], df2p, Geom.point, Theme(alphas=[0.4], highlight_width=0pt, point_size=2pt)),
        layer(x=xs, y=ys, color=[ colorant"blue" ], df, Geom.point, Theme(alphas=[1.0], highlight_width=0pt, point_size=2pt)),
        Guide.xlabel("X"), Guide.ylabel("Y"),
        Coord.cartesian(xmin=xmin-extra, xmax=xmax+extra, ymin=ymin-extra, ymax=ymax+extra))
end