text:
#==
Linear
    Regression
==#
@doc """
      Linear Regression is a well-known linear function used for predicting
      continuous features with a mostly linear or semi-linear slope.\n
      --------------------\n
      ==PARAMETERS==\n
     [y] <- Fill with your trainY values. Should be an array of shape (0,1) or (1,0)\n
     [x] <- Fill in with your trainX values. Should be an array of shape (0,1) or (1,0)\n
     --------------------\n
     ==Functions==\n
     predict(xt) <- Returns a prediction from the model based on the xtrain value passed (xt)
       """
mutable struct LinearRegression{P} <: LinearModel
    A::Float64
    B::Float64
    predict::P
    regressors::Array{LinearModel}
    function LinearRegression(x::Array,y::Array)
        # a = ((∑y)(∑x^2)-(∑x)(∑xy)) / (n(∑x^2) - (∑x)^2)
        # b = (x(∑xy) - (∑x)(∑y)) / n(∑x^2) - (∑x)^2
        regressors = []
        if length(x) != length(y)
            throw(ArgumentError("The array shape does not match!"))
        end
        # Get our Summations:
        Σx = sum(x)
        Σy = sum(y)
        # dot x and y
        xy = x .* y
        # ∑dot x and y
        Σxy = sum(xy)
        # dotsquare x
        x2 = x .^ 2
        # ∑ dotsquare x
        Σx2 = sum(x2)
        # n = sample size
        n = length(x)
        # Calculate a
        a = (((Σy) * (Σx2)) - ((Σx * (Σxy)))) / ((n * (Σx2))-(Σx^2))
        # Calculate b
        b = ((n*(Σxy)) - (Σx * Σy)) / ((n * (Σx2)) - (Σx ^ 2))
        predict(xt::Array) = (xt = [i = a + (b * i) for i in xt])
        P = typeof(predict)
        return new{P}(a, b, predict, [])
    end
        function LinearRegression(x::DataFrame,y::Array)
            # a = ((∑y)(∑x^2)-(∑x)(∑xy)) / (n(∑x^2) - (∑x)^2)
            # b = (x(∑xy) - (∑x)(∑y)) / n(∑x^2) - (∑x)^2
            regressors = []
            count = 1
            [push!(regressors, LinearRegression(feature, y) for feature in x)]
            a = nothing
            b = nothing
            for m in regressors
                if a != nothing
                    a = mean(a, m.a)
                    b = mean(b, m.b)
                else
                    a = m.a
                    b = m.b
                end
            end
            predict(xt::DataFrame) = _compare_predCon(models, xt)
            P = typeof(predict)
            return new{P}(a, b, predict, regressors)
    end
end
#==
Linear
    Least
     Square
==#
"""

      Least Square regressors are ideal for predicting continous features.\n
      --------------------\n
      x = [7,6,5,6,5]\n
      y  = [3.4.5.6.3]\n
      xtrain = [7,5,4,5,3,5,7,8]\n
      Type = :LIN\n
      model = models.LeastSquare(x,y,Type)\n
      y_pred = models.predict(model,xtrain)\n
      -------------------\n
      HYPER PARAMETERS\n
      --------------------\n
      ==Functions==\n
      predict(xt) <- Returns a prediction from the model based on the xtrain value passed (xt)
       """
function LinearLeastSquare(x,y)
    if length(x) != length(y)
        throw(ArgumentError("The array shape does not match!"))
    end
        xy = x .* y
        sxy = sum(xy)
        n = length(x)
        x2 = x .^ 2
        sx2 = sum(x2)
        sx = sum(x)
        sy = sum(y)
        # Calculate the slope:
        a = ((n*sxy) - (sx * sy)) / ((n * sx2) - (sx)^2)
     # Calculate the y intercept
        b = (sy - (a*sx)) / n
    predict(xt) = (xt = [z = (a * x) + b for x in xt])
    (var)->(a;b;predict)
end
#==
Lasso
    Regression
==#
function lasso_cost(X, Y, B, alp)
     m = length(Y)
     cost = (sum(((X .* B) - Y).^2)/2) +(sum(broadcast(abs, B)))*alp/2
     return cost
end
function gradientDescent(X, Y, B, learningRate, numIterations, alp)
    m = length(Y)
    for iteration in 1:numIterations
        loss = (X .* B) - Y
        for i in 1:length(X)
            if(B[i]==0)
                B[i]=1
            end
        end
        gradient = ((X' * loss) .+ (alp .* (broadcast(abs, B) ./ B))) / m
        B = B - learningRate * gradient
    end
    return B
end
mutable struct LassoRegression{P}
    x::Array
    y::Array
    predict::P
    lambda::Float64
    i::Float64
    LearningRate::Float64
    n_iterations::Int64
    function LassoRegression(x, y, lambda = 0; lr = .01, n_iterations = 1200)
        i = 10 ^ 6
        predict(xt) = pred_lasso(x, y, xt, i, learningRate = lr, n_iterations = n_iterations)
        new{typeof(predict)}(x, y, predict, lambda, i, lr, n_iterations)
    end
    function pred_lasso(x, y, B, i; learningRate = .01, n_iterations = 1200)
        costm = 10 ^ 30
        newB = zeros(size(B))
          while i > 0.001
             B1 = gradientDescent(x, y, B, learningRate, 1200,i)
             cost = lasso_cost(x, y, B1, i)
            if (cost < costm)
                costm = cost
                newB = B1
                lambda = i
            end
         i=i/2
        end
        return(B .* newB)
    end
end


text:
using DrWatson
quickactivate(@__DIR__)
@show projectdir()
##
using CovidSurvey
using CSV
using DataFrames
using Dates
# using RCall
using Turing
using StatsPlots
using Plots.PlotMeasures
using Random
using Base.Threads
using StatsFuns
using ReverseDiff
using Memoization
Turing.setrdcache(true)
setadbackend(:reversediff)

## ============================================================================
nthreads()
name2model = Dict(
    "v3" => National.model_v3,
    "v2"        => National.model_v2,
    "deaths" => National.model_deaths,
    "cases" => National.model_cases,
    "hospit" => National.model_hospit,
)
#-----------------------------------------------------------------------------
# load data
ps = (
    warmup = 10,
    steps  = 100,
    seed   = nothing,
    observ = "2021-06-01",#"2021-01-13",#"2021-02-06",#"2021-03-25"
    cases  = "2020-06-01",
    model  = "cases",
    preds = "CF,CC,CR,CS",
    semipara=true,
)
!isnothing(ps.seed) && Random.seed!(ps.seed);
fname = savename("chains", ps, "")
@info ps

data = National.load_data(;
    observationsend  = ps.observ,
    predictors        = ps.preds |> CovidSurvey.parse_predictors,
    rwstep           = 1,
    epidemicstart    = 30,
    numimpute        = 6,
    deathmodel        = National.DeathInit(obs_stop="$( Date(ps.cases) + Day(60))"),
    casemodel         = National.CaseInit(obs_start=ps.cases),
    # hospitmodel       = National.HospitInit(obs_stop="$( Date(ps.cases) + Day(1))")
    link              = KLogistic(3.),
    invlink           = KLogit(3.),
    lockdown          = "2020-03-18",
    covariates_kwargs = Dict(
        :semiparametric => ps.semipara,
        :fname => projectdir("data","smoothed_contact_rates.csv"),
        :shift => -1,
        :startdate => "2020-11-10",
        :enddate => nothing))
turing_data = data.turing_data;
#----------------------------------------------------------------------------
# sample model
model = name2model[ps.model]
m = model(turing_data, false)
Turing.emptyrdcache()
m()
@time chain = sample(m, NUTS(ps.warmup, 0.95), ps.steps + ps.warmup; progress=true)

# @time chain = sample(m, NUTS(ps.warmup, 0.95), MCMCThreads(), ps.steps + ps.warmup, 3)
# using MCMCChains
chain = chain[ps.warmup+1:end,:,:]
# chain = chain[:,:,[2,4]]

@info "Saving at: $(projectdir("out", fname))"
# safesave(projectdir("out/tmp", fname*".jls"), chain)
## =====================================================================
diagnostics = gelmandiag(chain)
fname_diagnostics = projectdir("out/tmp", fname*"_GELMANDIAG.csv")
safesave(fname_diagnostics, diagnostics)
## ======================================================================
# data = ImperialUSAcases.Data(Dict(),turing_data,Dict("DK"=>dk.date),turing_data.deaths,"DK")
# m = model(turing_data..., false)
# chain = sample(m, Prior(), 2000)
m_pred = model(turing_data, true)
gq = Turing.generated_quantities(m_pred, chain)
generated_posterior = vectup2tupvec( reshape(gq, length(gq)) );
#---------------------------------------------------------------------------
# plot results
plotlyjs()
p = National.plot(data, generated_posterior);
savefig(p, "/home/and/tmp/figures/"*fname*".html")
# savefig(p, projectdir("figures/tmp", fname*".png") )
run(`firefox $("/home/and/tmp/figures/"*fname*".html")`, wait=false)
##
using Optim
init_vals = Float64.(vec(mean(Chains(chain, :parameters).value, dims=1)))
mle_estimate = optimize(m, MAP(), init_vals, LBFGS(), Optim.Options(iterations=1_000, allow_f_increases=true))
##
plotlyjs()
n = filter( x->!occursin(r"\[", x), String.(names(chain)))
p = plot(chain[n]);
savefig(p, "/home/and/tmp/figures/"*fname*"_CHAINS.html" )
run(`firefox  $("/home/and/tmp/figures/"*fname*"_CHAINS.html")`, wait=false)
##
using MCMCChains
plotlyjs()
n = filter( x->!occursin(r"\[", x), String.(names(chain)))
p = meanplot(chain[n]);
savefig(p, projectdir("/home/and/tmp/figures/", fname*"_MEANPLOT.html") )
run(`firefox $(projectdir("/home/and/tmp/figures/", fname*"_MEANPLOT.html"))`, wait=false)

##
p = let
    changepoint  = argmin( abs.(mean(Rt)[end-50:end] .- 1) )
    localmin  = argmin(mean(Rt)[end-50:end])
    tickdates = data.dates[1]:Month(2):data.dates[end]
    rt        = hcat(Rt...)
    t1        = Date(2020, 12, 24)
    t1_str    = t1 |> Symbol |> String
    t2        = data.dates[end-50:end][ localmin ]
    t2_str    = t2 |> Symbol |> String
    t3        = data.dates[end-50:end][ changepoint ]
    t3_str    = t3 |> Symbol |> String

    # p1 = plot()
    tickdates = data.dates[1]:Month(3):data.dates[end]
    p1 = plot(title=L"estimated $R_t$")
    rt = hcat(Rt...)
    plot_confidence_timeseries!(p1, data.dates, rt)
    vline!(p1, [t1], ls=:solid, lc=:black, label=t1_str)
    vline!(p1, [t2], ls=:dash, lc=:black, label=t2_str)
    vline!(p1, [t3], ls=:dot, lc=:black, label=t3_str)

    p2 = plot(title="Mobility: change in percent from baseline")
    y = dk.workplaces_percent_change_from_baseline |> rolling
    t = data.dates[4:end-3]
    plot!(p2, t, y, label="work") #data.turing_data.covariates[:,1]
    # plot!(p2, data.dates, -100*data.turing_data.covariates[:,1], label="non-residential") #data.turing_data.covariates[:,1]
    y = dk.residential_percent_change_from_baseline |> rolling
    plot!(p2, t, y, label="residential")

    p = plot(p1, p2, layout=(2,1), sharex=true, xticks=tickdates, legend = :outertopright)
    p
end
savefig(p, "FLA20a/Covid19/figures/reproduction_rate_vs_mobility.png")

rolling(xs) = [maximum(xs[i:i+6]) for i in 1:length(xs)-6]

data_nature = ImperialReport13.load_data(datadir("imperial-report13", "processed.rds"))
data_usa = ImperialUSAcases.load_data(datadir("imperial-usa-cases", "processed.rds"))
dk = begin
    dk = load(datadir("imperial-usa-cases", "DK.rds"))
    i = findfirst(>=(10), skipmissing(dk.deaths))
    startdate = dk.date[i] - Day(30)
    startidx = findfirst(==(startdate), dk.date)
    dk[startidx:end,:]
end
survey = CSV.File(projectdir("data", "survey", "predictors.csv"))|>DataFrame
survey.date = survey.date .|> parsetime

##
figname = savename("chains_seroprev_Nobs=$(length(turing_data.cases))", ps, "")
ImperialDKcases.plot_results(data, expected_daily_cases, expected_daily_deaths, predicted_daily_cases, Rt)
figpath = "/home/and/code/ImperialCollegeReplications/FLA20a/Covid19/figures/$(figname)"
savefig( figpath*".png" )
savefig( figpath*".html" )
run(`firefox /home/and/code/ImperialCollegeReplications/FLA20a/Covid19/figures/fig_Nobs=303_seed=1_steps=500_warmup=500.html`)# $(figpath*".html")`)

#----------------------------------------------------------------
# prior checks
# testing random walk
using StatsFuns
# invlink(x) = 10 / ( 1 + exp( -x ) )
# invlink(x) = exp( -x )

myinvlink(x) = 5 / ( 1 + exp( -x ) )

histogram(rand(Normal(.7, .3), 1000) .|> myinvlink)

@model function testrw(week_index, num_weeks, ::Type{TV} = Vector{Float64}) where {TV}
    # κ         ~ truncated(Normal(0,0.5), 0, Inf)
    # μ         ~ truncated(Normal(3.28, κ), 0, Inf)
    μ         ~ Normal(.7, .3)
    ρ1        ~ truncated(Normal(0.8,0.05), 0, 1)
    ρ2        ~ truncated(Normal(0.1,0.05), 0, 1) #-ρ1
    σw        ~ truncated(Normal(0,1), 0, Inf)
    #WARNING √x can fail during warm up as we dont guarantee x>0
    σw_star = σw * sqrt( max(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 1e-5) )#min(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 0) ) #sqrt() (1+ρ2)*(1-ρ1-ρ2)*(1+ρ1-ρ2)/(1-ρ2)
    weekly_effect = TV(undef, num_weeks)
    weekly_effect[1] ~ Normal(0,0.01)
    weekly_effect[2] ~ Normal(0,σw_star)
    for w in 3:num_weeks
        weekly_effect[w] ~ Normal(ρ1*weekly_effect[w-1]+ρ2*weekly_effect[w-2], σw_star)
    end
    # invlink.( weekly_effect[week_index] )
    # 3.28 * 2 * logistic.( - weekly_effect[week_index])
    myinvlink.( μ .+ weekly_effect[week_index])
end

num_weeks  = 52
week_index = vcat( [ones(Int64, 7)*i for i in 1:num_weeks]... )
num_days   = length(week_index)
m = testrw(week_index, num_weeks)
# m = testrw(num_days)
# plot(c)
c = sample(m, Prior(), 1000)
rt = Turing.generated_quantities(m, c)|>vec
y = convert(Matrix{Float64}, hcat(rt...))
p = plot_confidence_timeseries(y)
# plot!(rt[1:3])

## Test GP

using KernelFunctions
using PDMats

invlink(x) = 5 / ( 1 + exp( -x ) )
x = 1:365
k = KernelFunctions.transform(Matern32Kernel(), 1/7)
K = kernelpdmat(k, x[:,1:1], obsdim=1)
y = invlink.(rand(MvNormal(K)))

plot(x, y)
X   = (1:10:365)[:,1:1]
Y   = y[1:10:end] + randn(length(X))*0.01
scatter!(X,Y)

# a ~ LogNormal(-1, .5)
# l ~ LogNormal(5,.5)
# k = KernelFunctions.transform(a*Matern32Kernel(), 1/l)
# K = kernelmatrix(k, day;obsdim=1)
# ϵ ~ MvNormal(K)

@model function testgp(X,Y)
    σ ~ LogNormal(-2, 2)
    l ~ LogNormal(5,.5)
    k = KernelFunctions.transform(Matern32Kernel(), 1/l)
    K = kernelpdmat(k,X;obsdim=1)
    ŷ ~ MvNormal(K)
    y = logistic.(ŷ)
    Y ~ MvNormal(y, σ)
    y
end
m = testgp(X,Y)
c = sample(m, Prior(), 1000)
histogram(c["θ"])
Ŷs = Turing.generated_quantities(m,c)
plot(x, y)
plot!(X, Ŷs[2], lc=:black, α=0.5, lw=.5)

c = sample(m, NUTS(200, .65), 100)
scatter!(X,Y)
plot!(X, Ŷs[1:100], lc=:black, α=0.2, lw=.2)
## test auto correlated noise in Rt

histogram(invlink.(-rand(LogNormal(-1, .5), 1000)*1) * 2 )

@model function testgp(X)
    a ~ LogNormal(-1, .5)
    l ~ LogNormal(5,.5)
    k = KernelFunctions.transform(a*Matern32Kernel(), 1/l)
    K = kernelpdmat(k,X;obsdim=1)
    ϵ ~ MvNormal(K)
    logistic.(-ϵ)*2
end
m = testgp(x[:,1:1])
c = sample(m, Prior(), 1000)
Ŷs = Turing.generated_quantities(m,c)
plot(x, Ŷs[1:10], α=0.5)
## test NegBinomial observation model

@model function testobs()
    ϕ ~ truncated(Normal(0, 5), 0, Inf)
    x ~ NegativeBinomial2(10, ϕ)
end

using StatsPlots
n = 1000
ys = []
for x in 10 .^(1:2)
    ϕs = rand(truncated(Normal(0, 5), 0, Inf), n)
    ϕs = rand(LogNormal()) y = [rand(NegativeBinomial2(x, 5)) for ϕ in ϕs]
    push!(ys, y)
end
# plot(ylim=(0.1, 10^7), yaxis=:log)
boxplot(ys; outliers=false, leg=nothing)

histogram(rand(Gamma(2,1), 1000))

##
using StatsPlots
plot(turing_data.π)
function gamma(μ, cv)
    α = 1/cv^2
    θ = μ/α
    Gamma(α, θ)
end
d  = gamma(6.5, 0.62)
d1 = gamma(5.1, 0.86)
d2 = gamma(17.8, 0.45)

g  = turing_data.serial_intervals
plot(d)
plot!(g / sum(g))

g  = turing_data.π2
plot(d1)
plot!(g / sum(g))

g = turing_data.π
histogram(rand(d1, 10_000) .+ rand(d2, 10_000), normalize=:pdf)
plot!(g / sum(g))
#--------------------------------------------------------------
# Import Turing and Distributions.
# Import Turing and Distributions.
using Turing, Distributions, Bijectors
using Random

struct RandomWalk{T} <: AbstractMvNormal  where T
    n::Int
    ρ1::T#Float64rw
    ρ2::T
    σ::T
end

Distributions.length(d::RandomWalk) = d.n

function Base.rand(::AbstractRNG, d::RandomWalk{T} where T)
    #σw_star = d.σw * sqrt( max(1-d.ρ1^2-d.ρ2^2-2d.ρ1^2*d.ρ2/(1-d.ρ2), 1e-5) )#min(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 0) ) #sqrt() (1+ρ2)*(1-ρ1-ρ2)*(1+ρ1-ρ2)/(1-ρ2)
    rw = zeros(d.n)
    rw[1] = rand(Normal(0, .01))
    rw[2] = rand(Normal(0, d.σ))
    for n in 3:d.n
        rw[n] = rand( Normal( d.ρ1*rw[n-1]+d.ρ2*rw[n-2], d.σ ) )
    end
    rw
end

function Distributions.logpdf(d::RandomWalk, x::AbstractVector)
    l  = logpdf(Normal(0, .01), x[1])
    l += logpdf(Normal(0, d.σ), x[2])
	for i in 3:d.n
        l += logpdf( Normal( d.ρ1*x[i-1]+d.ρ2*x[i-2], d.σ ), x[i] )
    end
    l
end

rw = let
    ρ1 = .8
    ρ2 = .1
    σw = .2
    σ = σw * sqrt( max(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 1e-5) )
    RandomWalk(100, ρ1, ρ2, σ)
end
x = rand(rw) .+ randn(100)*0.01
logpdf(rw, x)

@model function testrw(X)
    ρ1        ~ truncated(Normal(0.8,0.05), 0, 1)
    ρ2        ~ truncated(Normal(0.1,0.05), 0, 1) #-ρ1
    σw        ~ truncated(Normal(0,.2), 0, Inf)
    σ = σw * sqrt( max(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 1e-5) )
    rw ~ RandomWalk(100, ρ1, ρ2, σ)
    X  ~ MvNormal(rw, 0.01)
    return rw
end

m = testrw(x)
Turing.setadbackend(:forwarddiff)
@time c =sample(m, NUTS(), 1000)

rws = Turing.generated_quantities(m, c)
let
    p = scatter(x)
    plot!(rws[1:10], lw=.5, lc=:black, α=.5)
end

## evalueate Rt prior
using StatsFuns

@model function testpred(covariates, num_weeks, week_index, ::Type{TV} = Vector{Float64}) where {TV}
    last_time_step = 300
    num_missing    = findfirst(x->!ismissing(x), covariates[:,1])-1

    ρ1        ~ truncated(Normal(0.8,0.05), 0, 1)
    ρ2        ~ truncated(Normal(0.1,0.05), 0, 1) #-ρ1
    σw        ~ truncated(Normal(0,.2), 0, Inf)
    σ = σw * sqrt( max(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 1e-5) )
    weekly_effect ~ RandomWalk(num_weeks, ρ1, ρ2, σ)

    α ~ filldist(Normal(0,0.5), 1)
    μ = 3.28
    Rt                     = TV(undef, last_time_step)
    Rt[1:num_missing]     .= μ * 2 * logistic.( - weekly_effect[week_index[1:num_missing]]);
    Rt[num_missing+1:end] .= μ * 2 * logistic.( - weekly_effect[week_index[num_missing+1:last_time_step]]
                                                - covariates[num_missing+1:last_time_step, :] * α)
    Rt
end

m = testpred(turing_data.covariates, turing_data.num_weeks, turing_data.week_index)
x = [m() for i in 1:100]
plot(x, lw=.5, lc=:black, α=.5, legend=false)

c = turing_data.covariates
nmis    = findfirst(x->!ismissing(x), c[:,1])-1
c = c[nmis+1:end,:]
n = size(c,1)
x = [2*logistic.(rand(truncated(Normal(0,0.5), 0, Inf), n) .* c) for i in 1:100]
plot(x, lw=.5, lc=:black, α=.5, legend=false)
plot!(c, c=:red)

#----------------------------------------------------------------------------
dk = load(datadir("imperial-usa-cases", "DK.rds"))
X = data.turing_data.covariates
data_nature = ImperialReport13.load_data(datadir("imperial-report13", "processed.rds"))
t = dk.date

function npi(x)
    x = diff(x)
    pushfirst!(x, 0)
    t = findfirst(==(1), x)
    data_nature.country_to_dates["Denmark"][t]
end

t_npi = let
    ts = []
    for i in 1:6
        t = data_nature.turing_data.covariates[1][:,i] |> npi
        push!(ts, t)
    end
    ts
end
##
# plot(X[:,[3,4]])
# histogram(rand(Beta(5, 35), 10000))
#
# logistic(-.1)
# X[:,3] |>plot
# X[:,2] .-= 1

@model function testrw(num_weeks, week_index, X, ::Type{TV} = Vector{Float64}) where {TV}
    num_missing = findfirst(x->!ismissing(x), X[:,end]) -1

    α          ~ Beta(35, 15)
    α2         ~ Beta(10, 25)
    α3         ~ Beta(10, 25)
    # α         ~ arraydist([Normal(-1.5, 2), Normal(2,2)])#, truncated(Normal(0,1), 0, Inf)])
    ρ1        ~ truncated(Normal(0.8,0.05), 0, 1)
    ρ2        ~ truncated(Normal(0.1,0.05), 0, 1) #-ρ1
    σw        ~ truncated(Normal(0,.2), 0, Inf)
    σw_star = σw * sqrt( max(1-ρ1^2-ρ2^2-2ρ1^2*ρ2/(1-ρ2), 1e-5) )
    weekly_effect = TV(undef, num_weeks)
    weekly_effect[1] ~ Normal(0,0.01)
    weekly_effect[2] ~ Normal(0,σw_star)
    for w in 3:num_weeks
        weekly_effect[w] ~ Normal(ρ1*weekly_effect[w-1]+ρ2*weekly_effect[w-2], σw_star)
    end
    lockdown =  1 .- α * X[:,1]
    cc = 1 .+ α2 * X[num_missing+1:end,2]
    cs = 1 .+ α3 * X[num_missing+1:end,3]
    β = @. 2 * logistic( weekly_effect[week_index[1:end]] ) * lockdown
    β[num_missing+1:end] .*= cs .* cc
    # dailyeffect = weekly_effect[week_index[1:end]] .+ X[1:end,[1,2]]*α[[1,2]]
	# dailyeffect[num_missing+1:end] .+= X[num_missing+1 : end, 3]*α[3]

    # Rt[1:num_missing]     .= 3.28  * 2 * logistic.( X[1:num_missing,1]*α[1] .+ weekly_effect[week_index[1:num_missing]]);
    # Rt[num_missing+1:end] .= 3.28  * 2 * logistic.( X[num_missing+1:end,:]*α .+ weekly_effect[week_index[num_missing+1:end]])
    3.28 * β
end

num_weeks = data.turing_data.num_weeks
week_index = data.turing_data.week_index
m = testrw(num_weeks, week_index, X)
# plot(X)
let
    t = data.dates
    x = [m() for i in 1:1000]
    p = plot(t, x[1:10], c=:black, α=.4, legend=false)
    plot_confidence_timeseries!(p, t, hcat(x...))
end
##
using StatsFuns
lockdown = Date("18/3/2020", dateformat"d/m/y")

2 * logistic.(rand(Normal(-2.2, .2), 1000)) |>histogram

plotlyjs()
plotnpi(x) = plot!(dk.date, x; ticks=:native, label=false)
p1 = plot(data.dates, data.cases, label="cases")
p2 = plot()
plot!(t, dk.workplace_closing, label="workplace")
plot!(t, dk.gatherings_restrictions, label="gatherings_restrictions")
plot!(t, dk.internal_movement_restrictions, label="internal_movement_restrictions")
plot!(t, dk.school_closing, label="school_closing")
plot!(t, dk.transport_closing, label="transport_closing")
plot!(t, dk.cancel_events, label="cancel_events")
plot!(t, dk.stay_home_restrictions, label="stay_home_restrictions")
p4 = vline(t_npi, labels=data_nature.covariate_names)
p3 = plot(t, dk.stringency_index, label="stringency")
plot(p1, p2, p4, p3, layout=(4,1), legend=:outertopright, sharex=true, link=:x )
#------------------------------------------------------------------
# student project
date = data.dates
serial_interval = National.padzeros(National.serialinterval(length(t)))
delay_distr = National.inf2hosp(length(t))
hospit = data.hospit
df = DataFrame((; date, hospit, serial_interval, delay_distr))
save("/home/and/code/StudentProject/covid19model.csv", df)

## ============================================================================
θ = turing_data
predict = false
TV = Vector{Float64}

@unpack num_observations, num_total_days, num_rt_steps, invlink, deathmodel = θ
# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_countries[m]`
num_time_steps   = predict ? num_total_days : num_observations
num_time_steps_2 = predict ? num_total_days : deathmodel.stop
############# 2.) time varying reproduction number

R0         = truncated(Normal(3., 1.), 1., 5.)  |>rand
R1         = truncated(Normal(.8, .1), .5, 1.1)  |>rand
σ_rt       = truncated(Normal(0.1, .05), 0, .25)  |>rand
latent_Rt  = CovidSurvey.RandomWalk(num_rt_steps, σ_rt, invlink(R1))  |>rand
# @code_warntype logpdf(CovidSurvey.RandomWalk(num_rt_steps, σ_rt, invlink(R1)), latent_Rt)

Rt = TV(undef, num_time_steps)
CovidSurvey.National.random_walks!(Rt, θ, predict, latent_Rt, R0)
# @code_warntype CovidSurvey.National.random_walks!(Rt, θ, predict, latent_Rt, R0)

############ 3.) infection dynamics
τ  = Exponential(1 / 0.03)  |> rand
T  = typeof(τ)
y  = truncated(Exponential(τ),T(0),T(1000))  |> rand

newly_infected       = TV(undef, num_time_steps)
cumulative_infected  = TV(undef, num_time_steps)
effective_Rt         = TV(undef, num_time_steps)

# infections!(newly_infected, cumulative_infected, θ, τ, y, Rt)
CovidSurvey.National.infections!(newly_infected, cumulative_infected, effective_Rt, θ, τ, y, Rt)
@code_warntype CovidSurvey.National.infections!(newly_infected, cumulative_infected, effective_Rt, θ, τ, y, Rt)

########### 4.) derive observables
μ_i2d = truncated(Normal(20., 4.), 10, 30)  |> rand
μ_i2c = truncated(Normal(4.5, 1.), 3, 7)  |> rand
iar   = Beta(1,10)  |> rand
ifr   = truncated(Normal(6/1000, 3/1000), 2/1000, 10/1000)  |> rand
# ihr   ~ truncated(Normal(1/100,1/100),.1/100,5/100)

### day of week effect
weekdayeffect_simplex = Dirichlet([7,7,7,7,7,3,1])  |> rand
weekdayeffect = TV(undef, 7)
weekdayeffect!(weekdayeffect, weekdayeffect_simplex)
@code_warntype weekdayeffect!(weekdayeffect, weekdayeffect_simplex)

### holiday effect
holidayeffect = Beta(1,1)  |> rand

########### 3.) observation model
expected_daily_deaths = TV(undef, num_time_steps_2)
expected_daily_cases  = TV(undef, num_time_steps)

ϕ_c = truncated(Normal(50, 10), 30, Inf)  |> rand
ϕ_d = truncated(Normal(30, 10), 5, Inf)  |> rand

cmodel = National.WeekdayHolidayObsModel(
    θ.casemodel, μ_i2d, iar, ϕ_c, expected_daily_cases, weekdayeffect, holidayeffect
)
dmodel = National.SimpleObsModel(θ.deathmodel, μ_i2d, ifr, ϕ_d, expected_daily_deaths)

National.expected!(expected_daily_deaths, dmodel, newly_infected)
@code_warntype National.expected!(expected_daily_deaths, dmodel, newly_infected)
National.expected!(expected_daily_cases, cmodel, newly_infected)
@code_warntype National.expected!(expected_daily_cases, cmodel, newly_infected)

ℓ  = zero(Float64)
ℓ += logpdf(cmodel, θ.cases)
ℓ += logpdf(dmodel, θ.deaths)
@code_warntype logpdf(cmodel, expected_daily_cases)
@code_warntype logpdf(dmodel, expected_daily_deaths)

ext: jl

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
"""
Quite crude way of checking that a number is a probability

The two steps are

1. The number should be of the `Float64` type
2. The number should belong to [0.0,1.0]
"""
function checkprob(p::T) where {T<:AbstractFloat}
    @assert 0.0 <= p <= 1.0
end

"""
Expected value of a single Bernoulli event

Simply f(p): p
"""
function i_esp(p::Float64)
    checkprob(p)
	return p
end

"""
Variance of a single Bernoulli event

f(p): p(1-p)
"""
function i_var(p::Float64)
    checkprob(p)
	return p*(1.0-p)
end

"""
Variance of a series of additive Bernoulli events

f(p): ∑(p(1-p))
"""
function a_var(p::Array{Float64})
	return reduce(+, map(i_var, p))
end

"""
Variance of a series of multiplicative Bernoulli events
"""
function m_var(p::Array{Float64})
	# LOLWUT
	return reduce(*, map((x) -> i_var(x) + i_esp(x)*i_esp(x), p)) - reduce(*, map((x) -> i_esp(x)*i_esp(x), p))
end


text:
using Random

using Dates, TimeZones
using MicroLogging
using ProgressMeter

using Mise

function run_corr()

    @info "Initiate Sequence.."
    flush(stdout); flush(stderr)
    df = read_jongro("/input/jongro_single.csv")
    #=
    first(df, 5)
    │ Row │ no     │ stationCode │ date                      │ lat      │ long      │ SO2     │ CO      │ O3      │ NO2     │ PM10    │ PM25    │ temp    │ u       │ v       │ p       │ prep     │ snow     │
    │     │ Int64  │ Int64       │ TimeZones.ZonedDateTime   │ Float64  │ Float64   │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64 │ Float64  │ Float64  │
    ├─────┼────────┼─────────────┼───────────────────────────┼──────────┼───────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┼─────────┤─────────┤──────────┤──────────┤
    │ 1   │ 1      │ 111123      │ 2015-01-01T01:00:00+09:00 │ 37.572   │ 127.005   │ 0.004   │ 0.2     │ 0.02    │ 0.009   │ 57.0    │ 0.0     │ -7.4    │ 4.41656 │ 1.60749 │ 1011.8  │ missing  │ missing  │
    │ 2   │ 26     │ 111123      │ 2015-01-01T02:00:00+09:00 │ 37.572   │ 127.005   │ 0.005   │ 0.2     │ 0.019   │ 0.008   │ 70.0    │ 3.0     │ -8.0    │ 4.22862 │ 1.53909 │ 1011.7  │ missing  │ missing  │
    │ 3   │ 51     │ 111123      │ 2015-01-01T03:00:00+09:00 │ 37.572   │ 127.005   │ 0.005   │ 0.2     │ 0.02    │ 0.006   │ 92.0    │ 5.0     │ -8.4    │ 3.57083 │ 1.29968 │ 1012.1  │ missing  │ missing  │
    │ 4   │ 76     │ 111123      │ 2015-01-01T04:00:00+09:00 │ 37.572   │ 127.005   │ 0.004   │ 0.2     │ 0.019   │ 0.005   │ 111.0   │ 2.0     │ -8.8    │ 4.60449 │ 1.6759  │ 1012.3  │ missing  │ missing  │
    │ 5   │ 101    │ 111123      │ 2015-01-01T05:00:00+09:00 │ 37.572   │ 127.005   │ 0.005   │ 0.2     │ 0.019   │ 0.006   │ 127.0   │ 5.0     │ -9.1    │ 5.35625 │ 1.94951 │ 1011.8  │ missing  │ missing  │
    =#
    features = [:SO2, :CO, :O3, :NO2, :PM10, :PM25, :temp, :u, :v, :pres, :humid, :prep, :snow]

    # For GPU
    default_FloatType::DataType = Float64

    norm_prefix = "norm_"
    norm_features = [Symbol(eval(norm_prefix * String(f))) for f in features]

    μσs = mean_and_std_cols(df, features)
    zscore!(df, features, norm_features)

    # convert Float types
    for fea in features
        df[!, fea] = default_FloatType.(df[!, fea])
    end

    for nfea in norm_features
        df[!, nfea] = default_FloatType.(df[!, nfea])
    end

    train_features = [:SO2, :CO, :O3, :NO2, :PM10, :PM25, :temp, :u, :v, :pres, :humid, :prep, :snow]
    norm_train_features = [Symbol(eval(norm_prefix * String(f))) for f in train_features]

    plot_totaldata(df, :PM25, "/mnt/")
    plot_totaldata(df, :PM10, "/mnt/")

    plot_pcorr(df, norm_features, features, "/mnt/")

    sample_size = 72
    output_size = 24
    epoch_size = 500
    batch_size = 32
    @info "training feature : " train_features
    @info "sizes (sample, output, epoch, batch) : ", sample_size, output_size, epoch_size, batch_size

    # slices for indicies
    # split into segment
    # sg_idxs = split_df(size(df, 1), sample_size)
    # split into window [[1,2,3,4],[2,3,4,5]...]
    train_sdate = ZonedDateTime(2012, 1, 1, 1, tz"Asia/Seoul")
    train_fdate = ZonedDateTime(2017, 12, 31, 23, tz"Asia/Seoul")
    test_sdate = ZonedDateTime(2018, 1, 1, 1, tz"Asia/Seoul")
    test_fdate = ZonedDateTime(2018, 12, 31, 23, tz"Asia/Seoul")

    # windowsed index (index in input)
    train_valid_wd_idxs = window_df(df, sample_size, output_size, train_sdate, train_fdate)
    test_wd_idxs = window_df(df, sample_size, output_size, test_sdate, test_fdate)

    # Befor to prepare (input, output) pairs, determine its size first
    train_valid_wd_size = length(train_valid_wd_idxs)
    test_wd_size = length(test_wd_idxs)
    train_size, valid_size = split_sizes2(length(train_valid_wd_idxs), batch_size)
    test_size = length(test_wd_idxs)

    # random permutation train_valid_wd_idxs itself for splitting train/valid set
    rng = MersenneTwister()
    Random.shuffle!(rng, train_valid_wd_idxs)
    # chunks for minibatch, only need train_chnk because of minibatch itself, i.e.
    train_chnk, valid_chnk = create_chunks(train_valid_wd_idxs, train_size, valid_size, batch_size)
    # pure indexes i.e. [1:2, 2:3, ...]
    train_idxs, valid_idxs = create_idxs(train_valid_wd_idxs, train_size, valid_size)

    # pure indexs
    test_idxs = create_idxs(test_wd_idxs, test_size)
    #test_idxs = collect((train_size + valid_size + 1):(train_size + valid_size + test_size))
    flush(stdout); flush(stderr)

    # simply collect dates, determine exact date for prediction (for 1h, 24h, and so on) later
    test_dates = collect(test_sdate + Hour(sample_size - 1):Hour(1):test_fdate - Hour(output_size))

    input_size = sample_size * length(train_features)

    days = 10
    # max
    lags = collect(0:(24 * days))
    output_dir = "/mnt/corr/"
    Base.Filesystem.mkpath(output_dir)
    for train_sym in [:norm_PM10, :norm_PM25]
        @info "    Compute corrs"
        p = Progress(length(norm_train_features), dt=1.0, barglyphs=BarGlyphs("[=> ]"), barlen=40, color=:yellow)
        for feature_sym in norm_train_features
            ProgressMeter.next!(p);
            corrs = corr_input(df, train_sym, feature_sym, lags)

            plot_corr_input(corrs,
                "Corr : $(string(train_sym)) & $(string(feature_sym))",
                output_dir, "corr_$(string(train_sym))_$(string(feature_sym))")
        end
    end
end

run_corr()


text:
@deprecate isNA(x) isna(x)
@deprecate anyNA(x) anyna(x)
@deprecate getAttrib getattrib
@deprecate setAttrib! setattrib!
@deprecate getNames getnames
@deprecate setNames! setnames!
@deprecate getClass getclass
@deprecate setClass! setclass!


text:
### A Pluto.jl notebook ###
# v0.17.2

using Markdown
using InteractiveUtils

# ╔═╡ 0469b379-82a6-4d53-9b52-262c04c7d6ba
begin
	let
		using Pkg
		Pkg.activate(mktempdir())
		Pkg.Registry.update()
		Pkg.add("PlutoUI")
		Pkg.add("Plots")
		Pkg.add("BenchmarkTools")
		Pkg.add(url="https://github.com/Dale-Black/DistanceTransforms.jl")
	end

	using PlutoUI
	using Plots
	using BenchmarkTools
	using DistanceTransforms
end

# ╔═╡ a77f43f6-4fb4-11ec-2a50-3181f56e24e0
md"""
## Import packages
"""

# ╔═╡ e3c3666c-7686-4aab-95ad-bc4f41b3319b
TableOfContents()

# ╔═╡ e99a10ab-1ab5-48d0-9342-9a724b750538
md"""
## Multi-threaded distance transforms
One of the advantages of a Julia-based distance transform library is how accessible something like multi-threading is for complex algorithms. The `SquaredEuclidean` distance transform is highly parallelizable and is set up to take advantage of multi-threaded hardware.
"""

# ╔═╡ 03a5c392-c134-42fa-b81c-3e3a3a7b11b4
threads = Threads.nthreads()

# ╔═╡ 3f7ba10d-d7f8-47cb-a879-1582f78e7828
array = [
	1 0 1 1 1
	0 0 1 0 1
	1 0 1 0 1
	1 1 0 1 1
]

# ╔═╡ 4a34664a-93e4-4835-ab9c-75dd10b9c483
bool_array = boolean_indicator(array)

# ╔═╡ 2d44897c-8430-4d35-8e26-4e2b03fba411
tfm = SquaredEuclidean(bool_array)

# ╔═╡ e85546b9-03f0-4a1c-abd4-cd7a86d3fe80
sq_euc_transform = transform!(bool_array, tfm, threads)

# ╔═╡ 4d35a592-5b22-44ef-8823-9025ce054643
md"""
## Timing
Let's compare the time difference between the regular `SquaredEuclidean` distance transform algorithm to the multi-threaded algorithm
"""

# ╔═╡ 390323c4-25a3-4048-b5e8-9ff99422c304
begin
	sedt_mean = []
	sedt_std = []

	sedtP_mean = []
	sedtP_std = []

	for n in 1:50:200

		# SEDT
		x2 = boolean_indicator(rand([0, 1], n, n))
		tfm2 = SquaredEuclidean(x2)
		sedt = @benchmark transform($x2, $tfm2)

		push!(sedt_mean, BenchmarkTools.mean(sedt).time)
		push!(sedt_std, BenchmarkTools.std(sedt).time)

		# SEDT threaded
		x3 = boolean_indicator(rand([0, 1], n, n))
		tfm3 = SquaredEuclidean(x3)
		nthreads = Threads.nthreads()
		sedtP = @benchmark transform!($x3, $tfm3, $nthreads)

		push!(sedtP_mean, BenchmarkTools.mean(sedtP).time)
		push!(sedtP_std, BenchmarkTools.std(sedtP).time)
	end
end

# ╔═╡ 9cbb098b-e5b4-45c1-8eeb-d8b3a91d565b
begin
	x = collect(1:length(sedt_mean))
	x_new = ones(Int32, size(x))
	for i in 2:length(x)
		x_new[i] = x[i] * (200^2)
	end
end

# ╔═╡ 6ccf4414-c76a-475d-8091-2aa03f268382
begin
	Plots.scatter(
		x_new,
		sedt_mean,
		label="squared euclidean DT",
		xlabel = "Array size (elements)",
		ylabel = "Time (ns)"
		)
	Plots.scatter!(x_new, sedtP_mean, label="squared euclidean DT threaded")
end

# ╔═╡ b88bea1c-73bd-40e5-8966-3e5af6c3a026
md"""
As you can see, for small arrays the multi-threading functionality doesn't offer any improvements in speed (in fact it limits it) but as the array size increases the speed improvements also increase by order(s) of magnitude.
"""

# ╔═╡ Cell order:
# ╟─a77f43f6-4fb4-11ec-2a50-3181f56e24e0
# ╠═0469b379-82a6-4d53-9b52-262c04c7d6ba
# ╠═e3c3666c-7686-4aab-95ad-bc4f41b3319b
# ╟─e99a10ab-1ab5-48d0-9342-9a724b750538
# ╠═03a5c392-c134-42fa-b81c-3e3a3a7b11b4
# ╠═3f7ba10d-d7f8-47cb-a879-1582f78e7828
# ╠═4a34664a-93e4-4835-ab9c-75dd10b9c483
# ╠═2d44897c-8430-4d35-8e26-4e2b03fba411
# ╠═e85546b9-03f0-4a1c-abd4-cd7a86d3fe80
# ╟─4d35a592-5b22-44ef-8823-9025ce054643
# ╠═390323c4-25a3-4048-b5e8-9ff99422c304
# ╠═9cbb098b-e5b4-45c1-8eeb-d8b3a91d565b
# ╠═6ccf4414-c76a-475d-8091-2aa03f268382
# ╟─b88bea1c-73bd-40e5-8966-3e5af6c3a026


text:
# SDE estimators

## Supporting functions ————————————————————————
"""
    logpdf_reg_refprior(a,T=1.0) -> q

Unnormalized log reference prior for normalized SDE(1) parameter `α`,
`α = a*T` regularly sampled at interval `T`

Implements eq:sde1-log-reg-samp-refprior in lyx doc
"""
logpdf_reg_refprior(α) = -1/2*log(exp(abs(2α))-1)

"Approximate uninformative prior on SDE parameters a"
sdeparprior(a,Tp) =
  sum(logpdf_reg_refprior(a/Tp^i) for (i,a) in enumerate(a))

uniform_rand(rng,x,w) = rand(rng,Uniform(x-w,x+w))

normparstd(x) = exp(abs(x))*√(1-exp(-2abs(x)))

"
$SIGNATURES

Model variation around process `proc` with variations in the
transformed domain of size `f`.
"
function modelvariation(rng,problem::ProblemLike,proc,f)
  𝓣 = transformation(problem,proc)
  θ = paramstuple(problem,proc)
  η = inverse(𝓣,θ)
  ηv = η.+f*rand(rng,length(η))
  θv = 𝓣(ηv)
  procv = SdeProcess(problem,θv)
  return procv
end

modelvariation(rng,problem::ProblemProb,proc,f) = modelvariation(rng,problem.like,proc,f)

modelvariation(rng,problem,proc,f,n) = map(1:n) do _
  modelvariation(rng,problem,proc,f)
end

"""
$SIGNATURES
Posterior mean
"""
function posterior_mean(sam::AbstractArray{SdeProcess{T,U}}) where {T,U}
  a =  mean(x->x.eq.a,sam)
  σϵ = mean(x->x.eq.σϵ,sam)
  eq = Sde(a,σϵ)
  procm = if V<:State
    z = mean(x->x.state.z,sam)
    SdeProcess(eq,State(z))
  else
    SdeProcess(eq)
  end
  return procm
end

## Types ————————————————————————
"""
Maximum Likelihood estimator
$FIELDS
"""
Base.@kwdef struct MlEstimator
  "Number of iterations"
  N::Int64 = 1000
  "Root truncation threshold"
  R_truncate::Float64 = 1e5
end

"""
Posterior Mean estimator
$FIELDS
"""
Base.@kwdef struct PmeanEstimator
  "Sampling random number generator"
  rng::MersenneTwister
  "Number of iterations"
  N::Int64 = 1000
  "Indication for number of warmup iterations"
  Nwarmup::Int64 = 1000
  "Number of chains"
  nchains::Int64 = 4
  "Size of model variation in the transformed (η) domain"
  fvar::Float64 = 0.1
end

## Supporting functions WhiteNoise —————————————————————
"
$SIGNATURES

`logpdf` of white noise model with analytical maximum likelihood estimate for `σϵ`.

Returns log probability `lp` and the maximum likelihood estimate `σϵ`.
"
function logpdf_ml_white_noise(data)
  y = getfield.(data,:y)
  σ2 = mean(y.^2)
  Nl = length(y)
  lp = -1/2*(Nl*log(σ2)+Nl+Nl*log(2π))
  σ = √σ2
  return lp,σ
end

"""
$SIGNATURES

Estimate white noise parameter (standard deviation) from `data`.
"""
function Distributions.estimate(::Type{WhiteNoise},data)
  _,σ = logpdf_ml_white_noise(data)
  return WhiteNoise(σ)
end

## DynamicHMC supporting functions Sde —————————————————————

"
$SIGNATURES

logpdf with analytical maximum likelihood estimate for `σϵ`.
`z` is the initial state. If `nothing`, use stationary state.

Returns log probability `lp` and the maximum likelihood estimate `σϵ`.
"
function logpdf_ml_σϵ(e::SdeEngine{T},state,data) where {T}
  TR = real(T)
  E = zero(TR); S = zero(TR)

  # First observation
  t = data[1].t; y = data[1].y
  μy,σ2y = obsdist(e,state)
  fixed_init_state = σ2y==0
  if fixed_init_state
    if !(isapprox(y,μy,rtol=1e-6)) error("Distribution inconsistent with observations: y[1] = $(y[1]) ≠ μy = $μy while σ2y==0") end
  else
    if !(σ2y>0) throw(DomainError(σ2y,"σ2y<0")) end
    S += log(σ2y)
    E += (y-μy)^2/σ2y
    state = measure(e,state,y,μy,σ2y)
  end

  N = length(data)
  for n in 2:N
    t = data[n].t; y = data[n].y
    t_prev = data[n-1].t
    state = predict(e,state,t-t_prev)

    μy,σ2y = obsdist(e,state)
    if !(σ2y>0) throw(DomainError(σ2y,"σ2y<0")) end
    S += log(σ2y)
    E += (y-μy)^2/σ2y
    state = measure(e,state,y,μy,σ2y)
  end

  Nl = fixed_init_state ? (N-1) : N
  σ2ϵ = E/Nl
  lp = -1/2*(Nl*log(σ2ϵ)+S+Nl+Nl*log(2π))
  σϵ = √σ2ϵ
  return lp,σϵ
end

## General parameter estimation ————————————————————————

"
$SIGNATURES

Estimate the SDE(1) model.
"
function estimate_sde1(problem::Problem,R=100)
  @unpack data,state_est = problem
  t = [x.t for x in data]; y = [x.y for x in data]
  σy = std(y,corrected=false,mean=0)
  Tav = mean(diff(t))
  _,yr = interpolate(t,y,Tav;method=:nearest)
  yb = yr[1:end-1]; σb = sqrt(mean(yb.^2))
  yf = yr[2:end]; σf = sqrt(mean(yf.^2))
  ah = mean(yb.*yf)/(σb*σf)
  a = ah≤0 ? -R : -log(ah)/Tav
  σϵ = √(2a)*σy
  est = SdeProcess(Sde([a],σϵ))
  return est
end

"""
$SIGNATURES
Maximum likelihood estimate given initial process estimate `procstart`,
defined in the state space corresponding the problem's `SdeEngine`.
"""
function ml_estimate_init(backend,tor,problem,procstart′)
  if procstart′∉problem error("Process\n$procstart′\nnot in range of problem:\n$problem") end
  ml_est_init_backend = if backend==:dynamichmc
    ml_estimate_init_optim
  elseif backend==:stan
    ml_estimate_init_stan
  end
  return ml_est_init_backend(tor,problem,procstart′)
end

@enum EstimationStage EXPAND TRUNCATE

"
Estimation information used in estimation trace

$FIELDS
"
Base.@kwdef mutable struct EstimateInfo
    "model order"
    p::Int
    "Estimation stage"
    stage::EstimationStage
    "Initial estimator"
    init_tor::Union{Nothing,InitialEstimate} = nothing
    "Initial estimate"
    proc_ini::Union{Nothing,TimeSeriesModel} = nothing
    "Likelihood initial model"
    L_ini::Float64 = -Inf
    "Optimization converged?"
    converged::Bool = false
    "Estimate"
    est::Union{Nothing,TimeSeriesModel} = nothing
    "Likelihood estimate"
    L::Float64 = -Inf
    "Fatal error during estimation"
    error = nothing
end

function Base.show(io::IO, ::MIME{Symbol("text/plain")},x::EstimateInfo)
  println(io,"EstimateInfo")
  for k in fieldnames(EstimateInfo)
    print(io,@sprintf("% 12s : %s\n",k,getfield(x,k)))
  end
end

"
$SIGNATURES

Estimation trace entry for initial estimate `proc_ini`. Modifies `Tc`.
"
function estimate_trace_init!(Tc::EstimateInfo,tor,problem,proc_ini)
  proc_ini′ = SdeProcess(sde(problem,proc_ini.eq))
  Tc.L_ini = logpdf(problem,proc_ini′)
  Tc.est,opt = ml_estimate_init(:dynamichmc,tor,problem,proc_ini′)
  Tc.converged = Optim.converged(opt)
  Tc.L = logpdf(problem,Tc.est)
  return nothing
end

"
$SIGNATURES

SDE estimate and, optionally, the estimation trace. Note that the returned estimate may be
of lower order than the requested `problem.p`, because a lower order model may have
a higher likelihood.
"
function Distributions.estimate(tor::MlEstimator,problem::ProblemLike;
      init_tors=default_init_tors(problem),ret_trace=false,throw_errors=false)

  data = problem.data

  trace = EstimateInfo[]

  # White noise model
  L,σ = logpdf_ml_white_noise(data)
  est_wn = WhiteNoise(σ)

  push!(trace,EstimateInfo(;p=0,stage=EXPAND,est=est_wn,L=L,converged=true))

  ml_order = TimeSeriesModel[est_wn] # Maximum likelihood per order

  for stage in (EXPAND,TRUNCATE)
    pcs = stage==EXPAND ? (1:problem.p) : (problem.p-1:-1:2)
    for pc in pcs
      prob = lowerorder(problem,pc)
      n_tors = length(init_tors)
      TA = Vector{Any}(missing,n_tors)
      # Data copies so that threads do not have to access shared memory ==> higher speed
      thprob = [deepcopy(prob) for _ in 1:nthreads()]
      @threads for i in 1:n_tors
        prob_c = thprob[threadid()]
        init_tor_c = init_tors[i]
        Tc = EstimateInfo(;p=pc,init_tor=init_tor_c,stage=stage)
        try
          proc_ini0 = initial_estimate(init_tor_c,prob_c,ml_order)
          if isnothing(proc_ini0) continue end
          Tc.proc_ini = bringin(problem,proc_ini0)
          estimate_trace_init!(Tc,tor,prob_c,Tc.proc_ini)
        catch e
          if throw_errors||(e isa MethodError) rethrow() else @debug("SDE($pc) $(typeof(init_tor_c)) $e"); Tc.error = e end
        end
        TA[i] = Tc
      end
      T = collect(skipmissing(TA))
      if isempty(T) continue end
      append!(trace,T)

      # Maximum likelihood for current order and stage
      i = argmax(getfield.(T,:L))
      est_c = T[i].est
      if !isnothing(est_c)
        push!(ml_order,est_c)
      end
    end
  end

  ## Overall maximum likelihood model. It can be of an order lower than `problem.p`
  i = argmax(getfield.(trace,:L))
  est = trace[i].est

  return ret_trace ? (est,trace) : est
end

function tracetable(data,trace,proc,t_kld)
  Href = -logpdf(proc,data)
  R = map(trace) do t
    D = if isnothing(t.est)
      missing
    else
      safe_kldiv(proc,t.est,t_kld)
    end
    u = (
      Hhat = -t.L-Href,
      D = D,
      init_tor_type = string(typeof(t.init_tor)),
      stage_str = string(t.stage)
    )
    tt = (;(x=>getfield(t,x) for x in fieldnames(typeof(t)))...)
    merge(tt,u)
  end |> DataFrame

  return R,Href
end

maxlike_est(trace) = trace[argmax(getfield.(trace,:L))]

"
$SIGNATURES

Selected `trace` estimates.
"
function trace_estimates(trace)
  R = []
  pmax = maximum(getfield.(trace,:p))

  # Overall MLE
  estinfo = maxlike_est(trace)
  push!(R,(method="MLE",pmethod=pmax,estinfo=estinfo))

  # MLE using only estimates from expansion phase
  estinfo = maxlike_est(filter(x->x.stage==EXPAND,trace))
  push!(R,(method="MLE EXPAND",pmethod=pmax,estinfo=estinfo))

  # MLE for order `pmax_c`
  for pmax_c in 0:pmax
    estinfo = maxlike_est(filter(x->x.p≤pmax_c,trace))
    push!(R,(method="MLE pmax",pmethod=pmax_c,estinfo=estinfo))
  end

  ## MLE from random initiation
  trace_random = filter(trace) do x
    x.init_tor isa AddComplexRoot && x.init_tor.rootmode==EstimateSDE.ALL_RANDOM
  end
  estinfo = maxlike_est(trace_random)
  push!(R,(method="MLE init ALL_RANDOM",pmethod=pmax,estinfo=estinfo))

  return R
end

"""
$SIGNATURES

Sample SDE posterior
"""
function posterior_samples(backend,tor,problem,procstarts)
  posterior_samples_backend = if backend==:dynamichmc
    posterior_samples_dynamichmc
  elseif backend==:stan
    posterior_samples_stan
  end
  return posterior_samples_backend(tor,problem,procstarts)
end

"
$SIGNATURES

Posterior mean estimate `procm` for SDE estimation `problem` based on HMC samples from starting
points concentrated around `procstart`. The posterior mean is taken over the
parameterization corresponding to the `problem`, so either SDE coefficients `a` or roots `r`.
"
function Distributions.estimate(tor::PmeanEstimator,problem::ProblemProb,procstart)
  procstart′ = SdeProcess(sde(problem.like,procstart.eq))
  procstarts = modelvariation(tor.rng,problem.like,procstart′,tor.fvar,tor.nchains)
  procs,R = posterior_samples(:dynamichmc,tor,problem,procstarts)
  procm = mean(procs,problem)
  return procm,procs,R
end

## DynamicHMC / Optim backend ————————————————————————

function ml_estimate_init_optim(tor,problem::ProblemLike,procstart)
  𝓣 = transformation(problem,procstart)
  fg! = function(F,G,η)
    f(η) = -(problem∘𝓣)(η)
    neg_lp,B = Zygote.pullback(f,η)
    if !isnothing(G)
      grad = isnothing(B) ? 0 : B(1)[1]
      G[:] .= isnothing(grad) ? 0 : grad
    end
    if !isnothing(F) return neg_lp end
    return nothing
  end
  θstart = paramstuple(problem,procstart)
  ηstart = inverse(𝓣,θstart)
  opt = optimize(Optim.only_fg!(fg!),ηstart,
    LBFGS(),Optim.Options(f_tol=1e-2,iterations=tor.N,store_trace=true))
  θ = 𝓣(Optim.minimizer(opt))
  est = SdeProcess(problem,θ)
  return est,opt
end

function posterior_samples_dynamichmc(tor,problem::ProblemProb,procstarts)
  R = Vector{Any}(undef,tor.nchains)

  warmup_stages = default_warmup_stages(;
    local_optimization=nothing,
    middle_steps=Int(round(tor.Nwarmup/7)),
    doubling_stages=3
  )

  # Create data copies so threads do not need to access shared memory
  # this should accelerate processing
  threadpars = map(1:nthreads()) do _
    prob = deepcopy(problem)
    𝓣 = transformation(prob,first(procstarts))
    P = TransformedLogDensity(𝓣,prob)
    ∇P = ADgradient(:Zygote,P)
    (problem=prob,𝓣=𝓣,∇P=∇P)
  end

  @threads for i in 1:tor.nchains
    tpar = threadpars[threadid()]
    procs = procstarts[i]
    θs = paramstuple(tpar.problem,procs)
    ηs = inverse(tpar.𝓣,θs)
    R[i] = DynamicHMC.mcmc_keep_warmup(tor.rng,tpar.∇P,tor.N,
      initialization=(q=ηs,),
      warmup_stages=default_warmup_stages(;local_optimization=nothing,M=Symmetric),
      reporter=LogProgressReport(i,100,600) ) # NoProgressReport
  end
  chains = vcat([r.inference.chain for r in R]...)
  𝓣 = transformation(problem,first(procstarts))
  procs = map(η->SdeProcess(problem,𝓣(η)),chains)

  return procs,R
end


text:
@testset "Type" begin
    @testset "SuperPixel" begin
        for T in (Lab{Float32}, RGB{Float32}, Gray{Float32}, Float32)
            img = rand(T, 4, 4)

            pos = (1:2, 3:4)
            sp1 = SuperPixel(img, pos)
            sp2 = SuperPixel(img[pos...], pos)

            pos = CartesianIndices(pos)
            sp3 = SuperPixel(img, pos)
            sp4 = SuperPixel(img[pos], pos)
            @test_nowarn SuperPixel(img, pos[:])

            @test eltype(sp1.color) <: Colorant # Number is promoted to Gray
            @test size(sp1.color) == size(sp1.position)
            @test sp1.color == img[sp1.position] # relative order is perserved
            @test sp2 == sp1
            @test sp3 == sp1
            @test sp4 == sp1

            @test color(sp1) == sp1.color
            @test position(sp1) == sp1.position

            @test isempty(sp1) == isempty(sp1.color)
            @test imsize(sp1) == size(pos)
            @test color_type(sp1) == (T <: Number ? Gray{T} : T)
        end
    end

    @testset "SuperPixel image" begin
        for T in (Lab{Float32}, RGB{Float32}, Gray{Float32}, Float32)
            img = rand(T, 4, 4)
            img_sp = [SuperPixel(img, (1:2, 1:2)),
                      SuperPixel(img, (1:2, 3:4)),
                      SuperPixel(img, (3:4, 1:2)),
                      SuperPixel(img, (3:4, 3:4))]

            @test imsize(img_sp) == size(img)
            @test color_type(img_sp) == (T <: Number ? Gray{T} : T)
        end
    end
end


text:
using FewShotAnomalyDetection
using Flux
using MLDataPattern
using FluxExtensions
using Adapt
using DataFrames
using CSV
using Statistics

import FewShotAnomalyDetection: loss, zparams
include("experimentalutils.jl")

function runExperiment(datasetName, train, test, inputDim, hiddenDim, latentDim, numLayers, nonlinearity, layerType, β, γ, batchSize = 100, numBatches = 10000)
    T = Float32
    svae = SVAEtwocaps(size(train[1], 1), hiddenDim, latentDim, numLayers, nonlinearity, layerType, :scalarsigma)
    FewShotAnomalyDetection.set_normal_μ_nonparam(svae, vcat(T(1), zeros(latentDim - 1)))
    learnRepresentation(data, labels) = wloss(svae, data, (x, y) -> FewShotAnomalyDetection.mmd_imq(x, y, γ))
    printing_learnRepresentation(data, labels) = printing_wloss(svae, data, (x, y) -> FewShotAnomalyDetection.mmd_imq(x, y, γ))
    opt = Flux.Optimise.ADAM(1e-4)
    cb = Flux.throttle(() -> println("$datasetName inner SVAE: $(printing_learnRepresentation(train[1], zero(train[2])))"), 10)
    println("Before training - $datasetName inner SVAE: $(printing_learnRepresentation(train[1], zero(train[2])))")
    Flux.train!(learnRepresentation, Flux.params(svae), RandomBatches((train[1], zero(train[2])), batchSize, numBatches), opt, cb = cb)

    pxv = vec(collect(.-pxexpectedz(svae, test[1])'))
    pzs = vec(collect(.-pz(svae, test[1])'))
    pxis = vec(collect(.-px(svae, Flux.Tracker.data(test[1]))'))
    # println(pxis)
    auc_pxv = computeauc(pxv, test[2] .- 1)
    auc_pz = computeauc(pzs, test[2] .- 1)
    auc_pxis = computeauc(pxis, test[2] .- 1)

    println("AUC P(X) Vita: $auc_pxv | AUC P(X) IS: $auc_pxis AUC | learnt prior P(Z): $auc_pz")

    zs = zparams(svae, test[1])[1]
    μz = mean(zs, dims = 2)
    FewShotAnomalyDetection.set_normal_μ(svae, μz)
    pzs = vec(collect(.-pz(svae, test[1])'))
    auc_pz = computeauc(pzs, test[2] .- 1)

    println("AUC mean prior P(Z): $auc_pz")

    return DataFrame(dataset = datasetName, idim = inputDim, hdim = hiddenDim, ldim = latentDim, layers = numLayers, β = β, γ = γ, auc_pxv = auc_pxv, auc_pxis = auc_pxis, auc_pz = auc_pz)
end

outputFolder = mainfolder * "experiments/kernelsizetest/"
mkpath(outputFolder)

datasets = ["breast-cancer-wisconsin", "cardiotocography"]
difficulties = ["easy", "easy"]
batchSize = 100
iterations = 10000

if length(ARGS) != 0
    datasets = [ARGS[1]]
    difficulties = ["easy"]
end

for i in 1:10
    for (dn, df) in zip(datasets, difficulties)
        println("Loading $dn : $df")
        train, test, clusterdness = loaddata(dn, df)

        println("$dn")

        evaluateOneConfig = p -> runExperiment(dn, train, test, size(train[1], 1), p..., batchSize, iterations)
        results = gridsearch(evaluateOneConfig, [32], [8], [3], ["relu"], ["Dense"], [1.], [0.001 0.01 0.1 1])

        CSV.write(outputFolder * "kernelsizetest-$dn-$df-$i.csv", vcat(results...))
    end
end


text:
"""
    brute_force_search(w::Vector{UInt8}, s::Vector{UInt8}, ix_start, ix_stop)
Searches for word `w` in template `s`. `w` and `s` are strings represented as UInt8 vectors for speed. `ix_start` and `ix_stop` limit the search to within these indices. Returns the number of exact matches `N_matches`. Also returns `pos`, the offset position of the word with respect to `ix_start` to find the match if the match is unique, or 0 otherswise (multiple matches or no match)
# Examples
```jldoctest
julia> w=b"julia"; s=b"I love julia language";

julia> brute_force_search(w,s,1,length(s))
(8, 1)

julia> brute_force_search(w,s,7,length(s))
(2, 1)

julia> brute_force_search(w,s,9,length(s))
(0, 0)

julia> brute_force_search(w,s,1,10)
(0, 0)

julia> brute_force_search(b"l",s,1,length(s))
(0, 3)
```
"""
function brute_force_search(w::Vector{UInt8}, s::Vector{UInt8}, ix_start::Int64, ix_stop::Int64) #pattern, template
    nw=length(w)
    s_cursor=ix_start
    s_cursor_max=ix_stop-nw+2
    j=s_cursor
    N_match=0
    pos=0
    while s_cursor<s_cursor_max
        for i=1:nw
            w[i]==s[j] ? j+=1 : break
        end
        if j==(s_cursor+nw)
            N_match+=1
            N_match==1 ? pos=s_cursor : pos=0
        end
        s_cursor+=1
        j=s_cursor
    end
    pos>0 ? pos-=(ix_start-1) : true
    return pos, N_match
end

function brute_force_search_wildcard(w::Vector{UInt8}, s::Vector{UInt8}, ix_start::Int64, ix_stop::Int64) #pattern, template, wildcard in word is 0x4e (N)
    nw=length(w)
    s_cursor=ix_start
    s_cursor_max=ix_stop-nw+2
    j=s_cursor
    N_match=0
    pos=0
    while s_cursor<s_cursor_max
        for i=1:nw
                ((w[i]==s[j]) || (w[i]==0x4e)) ? j+=1 : break
        end
        if j==(s_cursor+nw)
            N_match+=1
            N_match==1 ? pos=s_cursor : pos=0
        end
        s_cursor+=1
        j=s_cursor
    end
    pos>0 ? pos-=(ix_start-1) : true
    return pos, N_match
end

function brute_force_search_bestfuzzy(w::Vector{UInt8}, s::Vector{UInt8}, ix_start::Int64, ix_stop::Int64, maxerr::Int64=0)::NTuple{4,Int64} #pattern, template
    nw=length(w)
    s_cursor=ix_start
    s_cursor_max=ix_stop-nw+2
    j=s_cursor
    N_match=0
    N_fuzzy=0
    pos=0
    N_err=0
    while s_cursor<s_cursor_max
        n_err=0
        for i=1:nw
            if w[i]==s[j]
                j+=1
            else
                n_err+=1
                n_err>maxerr ? break : j+=1
            end
        end
        if j==(s_cursor+nw) && n_err<=maxerr
            if n_err==0
                N_err=0
                N_match+=1
                N_match==1 ? pos=s_cursor : pos=0
            else #n_err>0
                N_fuzzy+=1
                if N_match==0
                    if N_fuzzy==1 || N_err>n_err #if improve err then on
                        pos=s_cursor
                        N_err=n_err
                    else #N_fuzzy>1 and no improvement
                        pos=0
                    end
                end
            end
        end
        s_cursor+=1
        j=s_cursor
    end
    pos>0 ? pos-=(ix_start-1) : true
    return pos, N_err, N_match, N_fuzzy #pos is 0 unless N_match=1 | (N_match=0 & N_fuzzy=1) | (there is a single fuzzy match at the lowest err value). Return unique exact match pos (then n_err=0) or if none pos of unique best fuzzy match (then n_err>0). N_err is 0 if exact match, or the best error level found (even if there are multiple match at this error level)
end

function bestof2searches(s1::NTuple{4,Int64},s2::NTuple{4,Int64})::Tuple{UInt8,Int64,Int64,Int64,Int64}
    if s1[1]>0 && s2[1]>0
        if s1[2]<s2[2] #s1 wins
            bof=(0x02,s1[1],s1[2],s1[3]+s2[3],s1[4]+s2[4])
        elseif s1[2]>s2[2] #s2 wins
            bof=(0x03,s2[1],s2[2],s1[3]+s2[3],s1[4]+s2[4])
        else #tie, turns inot confusing out
            bof=(0x04,0,s1[2],s1[3]+s2[3],s1[4]+s2[4])
        end
    elseif s1[1]>0 #s1 has single best, s2 has either no match or confusing (s2[1]=0)
        if s2[2]<=s1[2]
            s2[3]+s2[4]>0 ? bof=(0x04,s2[1],s2[2],s1[3]+s2[3],s1[4]+s2[4]) : bof=(0x02,s1[1],s1[2],s1[3]+s2[3],s1[4]+s2[4]) # s2 has confusion and wins : s2 has no match so s1 wins
        else #s2[2]>s1[2] #s1 has no confusion and is better than s2 anyway, ezpz
            bof=(0x02,s1[1],s1[2],s1[3]+s2[3],s1[4]+s2[4])
        end
    elseif s2[1]>0 #s2 has single best, s1 has either no match or confusing (s1[1]=0)
        if s1[2]<=s2[2]
            s1[3]+s1[4]>0 ? bof=(0x04,s1[1],s1[2],s1[3]+s2[3],s1[4]+s2[4]) : bof=(0x03,s2[1],s2[2],s1[3]+s2[3],s1[4]+s2[4]) # s2 has confusion and wins : s2 has no match so s1 wins
        else #s2[2]>s1[2] #s1 has no confusion and is better than s2 anyway, ezpz
            bof=(0x03,s2[1],s2[2],s1[3]+s2[3],s1[4]+s2[4])
        end
    else #both s1 and s2 are either confusing or have no match
        if (s1[3]+s1[4]==0) && (s2[3]+s2[4]==0) #no match
            bof=(0x01,0,0,0,0)
        elseif (s1[3]+s1[4]>0) && (s2[3]+s2[4]>0) # both are confusing
            bof=(0x04,0,min(s1[2],s2[2]),s1[3]+s2[3],s1[4]+s2[4])
        else #only one of them is confusin
            (s1[3]+s1[4]==0) ? bof=(0x04,0,s2[2],s2[3],s2[4]) : bof=(0x04,0,s1[2],s1[3],s1[4])
        end
    end
    return bof
end

max_stars_repo_name: straightlab/chartools-dev

max_issues_repo_name: straightlab/chartools-dev

max_forks_repo_name: straightlab/chartools-dev


text:
module LibVulkan

import Libdl

paths = String[]
const libvulkan = Libdl.find_library(["libvulkan", "vulkan", "vulkan-1", "libvulkan.so.1"], paths)
@assert libvulkan != ""

using CEnum

const Ctm = Base.Libc.TmStruct
const Ctime_t = UInt
const Cclock_t = UInt

export Ctm, Ctime_t, Cclock_t

include(joinpath(@__DIR__, "..", "gen", "vk_common.jl"))
include(joinpath(@__DIR__, "..", "gen", "vk_api.jl"))

# export everything
foreach(names(@__MODULE__, all=true)) do s
   if startswith(string(s), "VK_") || startswith(string(s), "vk") || startswith(string(s), "VULKAN")
       @eval export $s
   end
end

end # module


text:
# Tests for the ImageTransformations.warp! function
# Checking the gradients manually by writing τ and ϕ
function ϕ(p, homo)
    q = inv(homo) * [p..., 1]
    q = q[1:2] ./ q[3]
    return q
end

function τ(X, q)
    y = interpolate(X, BSpline(Linear()))
    y = extrapolate(y, zero(eltype(X)))
    return y(q...)
end

function project(X, H) # our own warp
    Y = similar(X)
    for i in CartesianIndices(Y)
        v = τ(X, ϕ(i.I, H))
        Y[i] = v
    end
    return Y
end

function ChainRulesCore.rrule(::typeof(τ), X, q)
    y = τ(X, q)
    function τ_pb(Δy)
        Δy = RGB(Δy...)
        itp = extrapolate(interpolate(X, BSpline(Linear())), zero(eltype(X)))
        ∇X = NoTangent()
        ∇q = Interpolations.gradient(itp, q...)
        n = nfields(Δy) > 0 ? nfields(Δy) : 1
        return NoTangent(), ∇X, typeof(q)(Δy .⋅ ∇q .⋅ n)
    end
    return y, τ_pb
end

function ChainRulesCore.rrule_via_ad(::Zygote.ZygoteRuleConfig, ::typeof(ϕ), p, homo)
    Zygote.pullback(ϕ, p, homo)
end

function ChainRulesCore.rrule(::typeof(project), X, H)
    Y = project(X, H)
    function project_pb(Δy)
        ∇H = zeros(3,3)
        for p in CartesianIndices(Y)
            _, ∇τ = rrule(τ, X, Float32.(p.I))
            _, _, ∇τ = ∇τ(Δy[p])
            _, ∇ϕ = rrule_via_ad(Zygote.ZygoteRuleConfig(), ϕ, p.I, H)
            _, ∇h = ∇ϕ([∇τ...])
            ∇H += ∇h
        end
        return NoTangent(), NoTangent(), ∇H
    end
    return Y, project_pb
end

function image_mse(y, ŷ)
    l = map((x, y) -> (x - y), y, ŷ)
    l = mapreducec.(x->x^2, +, 0, l)
    l = sum(l)
    l
end

@testset "ImageTransformations.warp! gradient" begin
    h = DiffImages.Homography{Float64}()
    H = [1.0 0.0 0.0;0.0 1.0 0.0;0.0 0.0 1.0]

    _abs2(c) = mapreducec(v->v^2, +, 0, c)
    img = rand(RGB{Float32}, 10, 10)
    tgt = rand(RGB{Float32}, 10, 10)

    zy = Zygote.gradient(h) do trfm
        out = ImageTransformations.warp(img, trfm, axes(img), zero(eltype(img)))
        out = image_mse(out, tgt)
        out
    end

    man = Zygote.gradient(H) do trfm
        out = project(img, trfm)
        out = image_mse(out, tgt)
        out
    end

    @test_broken zy[1].H ≈ man[1]
    @test zy[1].H ≈ -man[1] # TODO: Remove this test and make the above @test_broken -> @test after resolving the sign thing
end

@testset "RotMatrix vs. Non-RotMatrix" begin
    img = rand(RGB{Float32}, 10, 10)
    tgt = rand(RGB{Float32}, 10, 10)

    f(t) = begin
        out = ImageTransformations.warp(img, t, axes(img), zero(eltype(img)))
        loss = image_mse(out, tgt)
        loss
    end

    tfm1 = recenter(RotMatrix(π/8), center(img)) # using RotMatrix
    tfm2 = AffineMap(SMatrix{2, 2, Float64, 4}([cos(π/8) -sin(π/8); sin(π/8) cos(π/8)]), tfm1.translation)

    zy1 = Zygote.gradient(f, tfm1)
    zy2 = Zygote.gradient(f, tfm2)

    # The gradients are the same, but if you use `RotMatrix`s, you get nested NamedTuples,
    # for example:
    # zy1[1] => (linear = (mat = (data = [37.49...
    # zy2[1] => (linear = [37.49...
    @test_broken zy1[1].linear === zy2[1].linear # TODO: Resolve this in a seperate issue
    @test zy1[1].linear.mat.data === zy2[1].linear
end

@testset "Homography struct tests" begin
    h = DiffImages.Homography{Float32}()
    # Checking type stability of (h::Homography{K})(x::AbstractVector{L})
    for i ∈ (2, 3)
        v = rand(Float32, i)
        @test @inferred(h(v)) == h(v)
    end
    # Checking type stability of (h::Homography{K})(x::SVector{Union{2, 3}, L})
    for i ∈ (2, 3)
        v = rand(SVector{i, Float32})
        @test @inferred(h(v)) == h(v)
    end
end


text:
"""
	load_segment(filename::String) -> LAR

Return segment from file.
"""
function load_segment(filename::String)::LAR
    io = open(filename, "r")
    LINES = readlines(io)
    close(io)

    b = [tryparse.(Float64,split(LINES[i], " ")) for i in 1:length(LINES)]
    dim = Int(length(b[1])/2)
    V = hcat(reshape.(b,dim,2)...)
    EV = [[i,i+1] for i in range(1, size(V,2), step=2)]
    return V,EV
end

"""
	load_points(filename::String) -> Points

Return points from file.
"""
function load_points(filename::String)::Points
    io = open(filename, "r")
    point = readlines(io)
    close(io)

    b = [tryparse.(Float64,split(point[i], " ")) for i in 1:length(point)]
    V = hcat(b...)
    return V
end

"""
	load_cells(filename::String) -> Cells

Return cells from file.
"""
function load_cells(filename::String)::Cells
    io = open(filename, "r")
    cells = readlines(io)
    close(io)

    b = [tryparse.(Int64,split(cells[i], " ")) for i in 1:length(cells)]
    return b
end

# function load_hyperplane(filename::String)::Hyperplane
#     io = open(filename, "r")
#     point = readlines(io)
#     close(io)
#
#     b = [tryparse.(Float64,split(point[i], " ")) for i in 1:length(point)]
#     V = hcat(b...)
#
#     hyperplane = Hyperplane(PointCloud(V[:,3:end]), V[:,1],V[:,2])
#     return hyperplane
# end

# """
# 	load_connected_components(filename::String) -> Cells
#
# Load edges.
# """
# function load_connected_components(filename::String)::Cells
# 	EV = Array{Int64,1}[]
# 	io = open(filename, "r")
#     string_conn_comps = readlines(io)
#     close(io)
#
# 	conn_comps = [tryparse.(Float64,split(string_conn_comps[i], " ")) for i in 1:length(string_conn_comps)]
# 	for comp in conn_comps
# 		for i in 1:(length(comp)-1)
# 			push!(EV, [comp[i],comp[i+1]])
# 		end
# 		push!(EV,[comp[end],comp[1]])
# 	end
# 	return EV
# end


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

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

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

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

	dfGen = inputs["dfGen"]

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

	STOR_SYMMETRIC = inputs["STOR_SYMMETRIC"]

	### Constraints ###

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

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

	return EP
end

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

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

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

	STOR_SYMMETRIC = inputs["STOR_SYMMETRIC"]

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

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

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

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

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

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

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

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

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

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

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

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

	return EP
end

max_stars_repo_name: Betristor/LTESOM

max_issues_repo_name: Betristor/LTESOM

max_forks_repo_name: Betristor/LTESOM


text:
export ElasticWave

"""
    struct ElasticWave{N} <: Application{N}
Defines the elastic wave equation in conservation form:
∂E/∂t + ∇⋅(-1/2(u⊗I+I⊗u)) = 0
∂(ρu)/∂t + ∇⋅(-C:E) = f
where ρ is the density, u is the velocity, E is the strain tensor, f is a
forcing function or source term, and I is the N × N identity tensor. Defining
λ,μ to be the Lame parameters, in isotropic media, stress tensor S=CE=2μE+λtr(E)I
"""
struct ElasticWave{N} <: Application{N}
    " Pressure wave velocity (actually speed)"
    vp
    " Shear wave velocity (actually speed) "
    vs
    " ρ = Density "
    ρ
    " λ = 1st Lame parameter "
    λ
    " μ = 2nd Lame parameter "
    μ
    ElasticWave{N}(vp,vs,ρ,λ,μ) where {N} = new{N}(vp,vs,ρ,λ,μ)
    function ElasticWave{N}(param_dict::Dict) where {N}
        vp = param_dict["vp"]
        vs = param_dict["vs"]
        ρ = param_dict["ρ"]
        λ = param_dict["λ"]
        μ = param_dict["μ"]
        new{N}(vp,vs,ρ,λ,μ)
    end
end
nstates(app::ElasticWave{N}) where {N} = Int64(N*(N+1)/2)+N
"""
    flux_c!(flux, app::ElasticWave{N}, u)
Computes the convective flux for the Elastic wave equation:
f_c([E,ρu]) = [-1/2(u⊗I+I⊗u), -CE] ∈ ℜ^(nstates × N)
"""
function flux_c!(flux, app::ElasticWave{N}, u) where {N}
    nstrains = N*(N+1)/2
    # Remember - we are using Voigt notation for strain tensor
    Ediag = @view u[1:N]
    Eupp  = @view u[N+1:nstrains]
    vel = u[nstrains+1:end]./app.ρ
    # Compute -1/2(vel⊗I + I⊗vel)
    flux .= 0.0
    for i = 1:N # diagonal of strain tensor
        flux[i,i] = -vel[i]
    end
    offset = 1
    for i = N:-1:2 # upper-triangular portion of strain tensor
        for j = i-1:-1:1
            flux[N+offset,i] = -1/2*vel[j]
            flux[N+offset,j] = -1/2*vel[i]
            offset += 1
        end
    end
    # Compute -CE = -2μE - λtr(E)I
    trEλ = app.λ*sum(Ediag)
    for i = 1:N
        flux[nstrains+i,i] = -2*app.μ*Ediag[i] - trEλ
    end
    offset = 1
    for i = N:-1:2
        for j = i-1:-1:1
            flux[nstrains+i,j] = -2*app.μ*Eupp[offset]
            flux[nstrains+j,i] = -2*app.μ*Eupp[offset]
            offset += 1
        end
    end
    flux
end
"""
    numerical_flux_c!(flux, app::ElasticWave, uK, uN, normal_k)
Given the state vector at this element (uK) and its neighbor (uN), compute the
Lax-Friedrichs flux function of this PDE dotted with the normal, and
store the results in flux.
"""
function numerical_flux_c!(flux, app::ElasticWave{N}, uK, uN, normal_k) where {N}
    @assert N == length(normal_k)
    flux_k = Array{Float64,2}(undef, nstates(app), N)
    flux_c!(flux_k, app, uK)
    flux_n = Array{Float64,2}(undef, nstates(app), N)
    flux_c!(flux_n, app, uN)
    # Max speed is pressure wave speed
    C = app.vp
    flux .= (flux_k+flux_n)*normal_k./2.0 - (C/2.0)*(uN-uK)
end
"""
    boundary_flux_c!(flux, app::ElasticWave, u, bc, normal_k)
Free surface boundary condition (v=0) and absorbing boundary condition (Z*v)
for elastic wave equation
"""
function boundary_flux_c!(flux, app::ElasticWave{N}, u, bc, normal_k) where {N}
    nstrains = N*(N+1)/2
    vel = u[nstrains+1:end]./app.ρ
    flux .= 0.0
    # Remember - we are using Voigt notation for strain tensor
    # Compute -1/2(vel⊗I + I⊗vel)*n for Strain fluxes
    for i = 1:N # diagonal of strain tensor
        flux[i] = -vel[i]*normal_k[i]
    end
    offset = 1
    for i = N:-1:2 # upper-triangular portion of strain tensor
        for j = i-1:-1:1
            flux[N+offset] += -1/2*vel[j]*normal_k[i]
            flux[N+offset] += -1/2*vel[i]*normal_k[j]
            offset += 1
        end
    end
    if bc == -1 # Free surface BC
        # flux[nstrains+1:end] = 0.0
    elseif bc == -2 # Absorbing/non-reflecting BC
        # flux[nstrains+1:end] = ρ (vp n⊗n + vs (I-n⊗n))*vel
        Z = app.ρ*(app.vp*normal_k*normal_k'+app.vs*(diagm(ones(N))-normal_k*normal_k'))
        flux[nstrains+1:end] .= Z*vel
    else
        error("FATAL: Boundary condition not a valid choice!")
    end
end
"""
    source!(s_val, app::ElasticWave, t=0.0, x=0.0)
Computes the source(t,x) of the elastic wave equation:
source(t,x) = [0, ρ*wave(t)*g(x)]
"""
function source!(s_val, app::ElasticWave{N}, t=0.0, x=0.0) where {N}
    nstrains = N*(N+1)/2
    pts_per_lambda_min = 4
    mesh_max_dx = 0.1 # TODO: Does this need to be an app param?
    max_freq = app.vs/(pts_per_lambda_min*mesh_max_dx)
    freq = 0.9*max_freq
    wave_amp = cos(2*π*freq*t)
    s_val[1:nstrains] .= 0.0
    # TODO: g(x) := 1 => source wave active at all locations
    s_val[nstrains+1:end] .= app.ρ*wave_amp*1
    s_val
end

max_stars_repo_name: NoseKnowsAll/DGToolkit

max_issues_repo_name: NoseKnowsAll/DGToolkit

max_forks_repo_name: NoseKnowsAll/DGToolkit


text:
struct Suppressed{T}
    item::T
end
Base.show(io::IO, x::Suppressed) = print(io, "<suppressed ", x.item, '>')

function print_var(io::IO, var::JuliaInterpreter.Variable)
    print("  | ")
    T = typeof(var.value)
    local val
    try
        val = repr(var.value)
        if length(val) > 150
            val = Suppressed("$(length(val)) bytes of output")
        end
    catch
        val = Suppressed("printing error")
    end
    println(io, var.name, "::", T, " = ", val)
end

print_locdesc(io::IO, frame::Frame) = println(io, locdesc(frame))

function print_locals(io::IO, frame::Frame)
    vars = JuliaInterpreter.locals(frame)
    for var in vars
        # Hide gensymmed variables
        startswith(string(var.name), "#") && continue
        print_var(io, var)
    end
end

function print_frame(io::IO, num::Integer, frame::Frame)
    print(io, "[$num] ")
    print_locdesc(io, frame)
    print_locals(io, frame)
end

function print_next_expr(io::IO, frame::Frame)
    maybe_quote(x) = (isa(x, Expr) || isa(x, Symbol)) ? QuoteNode(x) : x
    expr = pc_expr(frame)
    @assert expr !== nothing
    print(io, "About to run: ")
    isa(expr, Expr) && (expr = copy(expr))
    if isexpr(expr, :(=))
        expr = expr.args[2]
    end
    if isexpr(expr, :call) || isexpr(expr, :return)
        for i in 1:length(expr.args)
            val = try
                @lookup(frame, expr.args[i])
            catch err
                err isa UndefVarError || rethrow(err)
                expr.args[i]
            end
            expr.args[i] = maybe_quote(val)
        end
    end
    if isa(expr, Expr)
        for (i, arg) in enumerate(expr.args)
            try
                nbytes = length(repr(arg))
                if nbytes > max(40, div(200, length(expr.args)))
                    expr.args[i] = Suppressed("$nbytes bytes of output")
                end
            catch
                expr.args[i] = Suppressed("printing error")
            end
        end
    end
    print(io, expr)
    println(io)
end

function print_status(io::IO, frame::Frame)
    # Buffer to avoid flickering
    outbuf = IOContext(IOBuffer(), io)
    printstyled(outbuf, "In ", locdesc(frame), "\n"; color=:bold)
    loc = locinfo(frame)

    if loc !== nothing
        data = if isa(loc, BufferLocInfo)
                loc.data
            else
                read(loc.filepath, String)
            end
        print_sourcecode(outbuf, data,
            loc.line, loc.defline)
    else
        print_codeinfo(outbuf, frame)
    end
    print_next_expr(outbuf, frame)
    print(io, String(take!(outbuf.io)))
end

function print_codeinfo(io::IO, frame::Frame)
    buf = IOBuffer()
    src = frame.framecode.src
    show(buf, src)
    active_line = convert(Int, frame.pc[])

    code = filter!(split(String(take!(buf)), '\n')) do line
        !(line == "CodeInfo(" || line == ")" || isempty(line))
    end

    code .= replace.(code, Ref(r"\$\(QuoteNode\((.+?)\)\)" => s"\1"))

    for (lineno, line) in enumerate(code)
        (lineno < active_line - 3 || lineno > active_line + 2) && continue

        color = (lineno < active_line) ? :white : :normal
        if lineno == active_line
            printstyled(io, rpad(lineno, 4), bold = true, color = :yellow)
        else
            printstyled(io, rpad(lineno, 4), bold = true, color = color)
        end
        printstyled(io, line, color = color)
        println(io)
    end
    println(io)
end


text:
using GeometricalOptics
using Test

funcsList_zern=[zernike,zernike,zplane]
funcsList_sph=[spherical, spherical,zplane]
funcsList_asph1=[aspherical, aspherical, zplane]
funcsList_asph2=[aspherical, aspherical, zplane]

super_funcsList=[funcsList_zern,funcsList_sph,funcsList_asph1,funcsList_asph2]

coeffsList_zern=[[4.,1.8,0.,0.,0.,0.05],[4.,0.01,0.,0.,0.,-0.2],[8.]]
coeffsList_sph=[[1.,6.],[2.,-6.],[8.]]
coeffsList_asph1=[[1.,6.,-5.],[2.,-6.,-5.],[8.]]
coeffsList_asph2=[[5.,-6.,-1.],[1.,2.,-1.],[8.]]

super_coeffsList=[coeffsList_zern,coeffsList_sph,coeffsList_asph1,coeffsList_asph2]

nList_zern=[1.,1.5,1.]
nList_sph=[1.,1.5,1.]
nList_asph1=[1.,1.5,1.]
nList_asph2=[1.,-1.,1.]

super_nList=[nList_zern,nList_sph,nList_asph1,nList_asph2]

super_optSta=opticalstack.(super_coeffsList,super_funcsList,super_nList)

ind_u(x,y,z,coeffs)=coeffs[1]+coeffs[2]*(x^2+y^2)
dt=0.03
ind_coeffs=[[0.],[1.5, -1/12],[0.]]
funcsList=[zplane,zplane,zplane]
coeffsList=[[1.],[2.],[6]]
nList=[1.,ind_u,1.]
optSta_grin=opticalstack(coeffsList,funcsList,nList,ind_coeffs,dt)

append!(super_optSta,[optSta_grin])

volPh(x,y,z,coeffs)= coeffs[1]*(x^2+y^2)
diffList=[volPh,0.]
diffmList=[1,0]
diffCoeffsList=[[-1000.],[1.]]
funcsList=[zplane,zplane]
coeffsList=[[1.],[8.]]
nList=[1.,1.]
optSta_diff=opticalstack(coeffsList,funcsList,nList,diffList,diffCoeffsList,diffmList,0.00055)

append!(super_optSta,[optSta_diff])

rdom=0.01:0.2: 1.
thdom=0.:pi/10:2pi
r=rdom'.*ones(length(thdom))
th=ones(length(rdom))'.*thdom
bund1= bundle_as_array(r.*cos.(th),r.*sin.(th),0. .* ones(size(r)), 0. .* ones(size(r)),0.)

@testset "bundles" begin
    @test_nowarn bundle_as_array(r.*cos.(th),r.*sin.(th),0. .* ones(size(r)), 0. .* ones(size(r)),0.)
    @test_nowarn bundle([0.],(-1:0.4:1).+eps(),0.,0.,0.)
    @test_nowarn bundle([0.],(-1:0.4:1).+eps(),0.,0.14,0.)
    @test_nowarn bundle((-1:0.1:1).+eps(),(-1:0.1:1).+eps(),0.,0.,0.)
    @test_nowarn bundle((-1:0.1:1).+eps(),(-1:0.1:1).+eps(),0.,0.14,0.)
    @test_nowarn bundle(0.,0. .+eps(),(-0.1:0.01:0.1).+eps(),(-0.1:0.01:0.1).+eps(),0.)
    @test_nowarn bundle(0.,0.14,(-0.1:0.01:0.1).+eps(),(-0.1:0.01:0.1).+eps(),0.)
    traceObj=super_optSta[1](bund1;isbigtrace=true)
    @test_nowarn bigtrace_to_bundle(traceObj,4)
end

rdom=0.01:0.2: 1.
thdom=0.:pi/10:2pi
r=rdom'.*ones(length(thdom))
th=ones(length(rdom))'.*thdom
bund1= bundle_as_array(r.*cos.(th),r.*sin.(th),0. .* ones(size(r)), 0. .* ones(size(r)),0.)
bund2=bundle([0.],(-1:0.4:1).+eps(),0.,0.,0.)
bund3 = bundle([0.],(-1:0.4:1).+eps(),0.,0.14,0.)
bund4=bundle((-1:0.1:1).+eps(),(-1:0.1:1).+eps(),0.,0.,0.)
bund5=bundle((-1:0.1:1).+eps(),(-1:0.1:1).+eps(),0.,0.14,0.)
bund6=bundle(0.,0. .+eps(),(-0.1:0.01:0.1).+eps(),(-0.1:0.01:0.1).+eps(),0.)
bund7=bundle(0.,0.14,(-0.1:0.01:0.1).+eps(),(-0.1:0.01:0.1).+eps(),0.)
traceObj=super_optSta[1](bund1;isbigtrace=true)
bund8=bigtrace_to_bundle(traceObj,4)

super_bundle=[bund1,bund2,bund3,bund4,bund5,bund6,bund7,bund8]

@testset "trace,eval" begin
    for i in 1:length(super_optSta), j in 1:length(super_bundle)
        optSta=super_optSta[i]
        bund=super_bundle[j]
        @test_nowarn traceObj=optSta(bund)
        @test_nowarn splot(traceObj)
        @test_nowarn splot(traceObj;pos=2)
        @test_nowarn rac(traceObj,1,4)
        @test_nowarn rac(traceObj,2,4)
        @test_nowarn rac(traceObj,1,3)
        @test_nowarn rms_spot(traceObj)
        @test_nowarn rms_spot(traceObj;pos=3)
        @test_nowarn trace_extract_ray(traceObj,1,1)
        @test_nowarn trace_extract_terminus(traceObj,4)
        # println(i,j)
    end
end

@testset "bigtrace,eval" begin
    for i in 1:length(super_optSta), j in 1:length(super_bundle)
        optSta=super_optSta[i]
        bund=super_bundle[j]
        @test_nowarn traceObj=optSta(bund,isbigtrace=true)
        @test_nowarn splot(traceObj)
        @test_nowarn splot(traceObj;pos=2)
        @test_nowarn rac(traceObj,1,4)
        @test_nowarn rac(traceObj,2,4)
        @test_nowarn rac(traceObj,1,3)
        @test_nowarn rms_spot(traceObj)
        @test_nowarn rms_spot(traceObj;pos=3)
        @test_nowarn trace_extract_ray(traceObj,1,1)
        @test_nowarn trace_extract_terminus(traceObj,4)
        # println(i,j,"Big")
    end
end

@testset "plot sys" begin
    rendvec=["YZ","XZ","3Dcirc","3Dsq"]
    for i in 1:length(super_optSta), j in 1:length(super_bundle), k in 1:length(rendvec)
        optSta=super_optSta[i]
        bund=super_bundle[j]
        @test_nowarn plot_obj=optSta(bund;rend=rendvec[k])
        # println(i,j,k)
    end
end


text:
using ThreeBodyDecay
using Test

two_js = ThreeBodySpins(1,0,2; two_h0=3)

@testset "Three Body Masses structure" begin
    #
    @test two_js.two_h1 == two_js.two_λ1 == two_js.two_j1 == two_js[1] == 1
    @test two_js.two_h2 == two_js.two_λ2 == two_js.two_j2 == two_js[2] == 0
    @test two_js.two_h3 == two_js.two_λ3 == two_js.two_j3 == two_js[3] == 2
    @test two_js.two_h0 == two_js.two_λ0 == two_js.two_j0 == two_js[0] == two_js[4] == 3
    #
    @test_throws ErrorException ThreeBodySpins(0,1,1; two_js0=1)
    @test_throws ErrorException ThreeBodySpins(0,1,1)
    @test_throws ErrorException two_js.j1
    @test_throws ErrorException two_js[5]
end

@testset "operations and interate" begin
    @test iseven(sum(two_js))
    @test div(sum(two_js .|> x2), 2) == sum(two_js)
end

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


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

mutable struct IoK8sApiExtensionsV1beta1IngressStatus <: SwaggerModel
    loadBalancer::Any # spec type: Union{ Nothing, IoK8sApiCoreV1LoadBalancerStatus } # spec name: loadBalancer

    function IoK8sApiExtensionsV1beta1IngressStatus(;loadBalancer=nothing)
        o = new()
        validate_property(IoK8sApiExtensionsV1beta1IngressStatus, Symbol("loadBalancer"), loadBalancer)
        setfield!(o, Symbol("loadBalancer"), loadBalancer)
        o
    end
end # type IoK8sApiExtensionsV1beta1IngressStatus

const _property_map_IoK8sApiExtensionsV1beta1IngressStatus = Dict{Symbol,Symbol}(Symbol("loadBalancer")=>Symbol("loadBalancer"))
const _property_types_IoK8sApiExtensionsV1beta1IngressStatus = Dict{Symbol,String}(Symbol("loadBalancer")=>"IoK8sApiCoreV1LoadBalancerStatus")
Base.propertynames(::Type{ IoK8sApiExtensionsV1beta1IngressStatus }) = collect(keys(_property_map_IoK8sApiExtensionsV1beta1IngressStatus))
Swagger.property_type(::Type{ IoK8sApiExtensionsV1beta1IngressStatus }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiExtensionsV1beta1IngressStatus[name]))}
Swagger.field_name(::Type{ IoK8sApiExtensionsV1beta1IngressStatus }, property_name::Symbol) =  _property_map_IoK8sApiExtensionsV1beta1IngressStatus[property_name]

function check_required(o::IoK8sApiExtensionsV1beta1IngressStatus)
    true
end

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


text:
### Measurements.jl ---  Uncertainty propagation library
#
# Copyright (C) 2016, 2017 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: uncertainty, error propagation, physics
#
# This file is a part of Measurements.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# This file is the entry point of the package: it defines the new `Measurement'
# type, some functions to handle it within Julia and the new functions provided
# by the package and exposed to users.
#
### Code:

__precompile__()

module Measurements

# Calculus is used to calculate numerical derivatives in "@uncertain" macro.
using Calculus

# Functions provided by this package and exposed to users
export Measurement, measurement, ±

# Define the "Derivatives" type, used inside "Measurement" type.  This should be
# a lightweight and immutable dictionary.
include("derivatives-type.jl")

##### New Type: Measurement
# Definition.  The Measurement type is composed by the following fields:
#   * val: the nominal value of the measurement
#   * err: the uncertainty, assumed to be standard deviation
#   * tag: a (hopefully) unique identifier, it is used to identify a specific
#     measurement in the list of derivatives.  NOTE: independent measurements
#     have tag > 0.  For dependent quantities the tag is 0.
#   * der: the list of derivates.  It is a lightweight dictionary, whose keys
#     are the tuples (nominal value, uncertainty, tag) of all independent
#     variables from which the object has been derived, the corresponding value
#     is the partial derivative of the object with respect to that independent
#     variable.  This dictionary is useful to trace the contribution of each
#     measurement and propagate the uncertainty in the case of functions with
#     more than one argument (in order to deal with correlation between
#     arguments).
struct Measurement{T<:AbstractFloat} <: AbstractFloat
    val::T
    err::T
    tag::UInt64
    der::Derivatives{T}
end
function Measurement(val::V, err::E, tag::UInt64,
                     der::Derivatives{D}) where {V,E,D}
    T = promote_type(V, E, D)
    return Measurement(T(val), T(err), tag, Derivatives{T}(der))
end
Measurement{T}(x::Measurement{S}) where {T,S} = convert(Measurement{T}, x)
Measurement{T}(x::S) where {T,S} = measurement(x)
Measurement{T}(x::S) where {T,S<:Rational} = measurement(x)

# Functions to quickly create an empty Derivatives object.
@generated empty_der1(x::Measurement{T}) where {T<:AbstractFloat} = Derivatives{T}()
@generated empty_der2(x::T) where {T<:AbstractFloat} = Derivatives{x}()

const tag_counters = UInt64[1]
function __init__()
    nthr = Base.Threads.nthreads()
    resize!(tag_counters, nthr)[:] = range(UInt64(1), step=typemax(UInt64)÷nthr, length=nthr)
end

measurement(x::Measurement) = x
measurement(val::T) where {T<:AbstractFloat} = Measurement(val, zero(T), UInt64(0), empty_der2(val))
measurement(val::Real) = measurement(float(val))
function measurement(val::T, err::T) where {T<:AbstractFloat}
    newder = empty_der2(val)
    if iszero(err)
        Measurement{T}(val, err, UInt64(0), newder)
    else
        @inbounds tag = tag_counters[Base.Threads.threadid()] += 1
        return Measurement{T}(val, err, tag, Derivatives(newder, (val, err, tag)=>one(T)))
    end
end
measurement(val::Real, err::Real) = measurement(promote(float(val), float(err))...)
const ± = measurement

"""
    measurement(val::Real, [err::Real]) -> Measurement
    val ± err -> Measurement

Return a `Measurement` object with `val` as nominal value and `err` as
uncertainty.  `err` defaults to 0 if omitted.

The binary operator `±` is equivalent to `measurement`, so you can construct a
`Measurement` object by explicitely writing `123 ± 4`.
"""
measurement

# Type representation
Base.show(io::IO, measure::Measurement) =
    print(io, measure.val, get(io, :compact, false) ? "±" : " ± ", measure.err)
Base.show(io::IO, ::MIME"text/latex", measure::Measurement) =
    print(io, "\$", measure.val, " \\pm ", measure.err, "\$")
for mime in (MIME"text/x-tex", MIME"text/x-latex")
    @eval Base.show(io::IO, ::$mime, measure::Measurement) =
        print(io, measure.val, " \\pm ", measure.err)
end
# Representation of complex measurements.  Print something that is easy to
# understand and that can be meaningfully copy-pasted into the REPL, at least
# for standard numeric types.
function Base.show(io::IO, measure::Complex{<:Measurement})
    r, i = reim(measure)
    compact = get(io, :compact, false)
    print(io, "(", r, ")")
    if signbit(i) && !isnan(i)
        i = -i
        print(io, compact ? "-" : " - ")
    else
        print(io, compact ? "+" : " + ")
    end
    print(io, "(", i, ")im")
end
# This is adapted from base/show.jl for Complex type.
function Base.alignment(io::IO, measure::Measurement)
    m = match(r"^(.*[\±])(.*)$", sprint(show, measure, context=io, sizehint=0))
    m === nothing ? (length(sprint(show, measure, context=io, sizehint=0)), 0) :
        (length(m.captures[1]), length(m.captures[2]))
end

### Juno pretty printing
using Requires
@require Juno="e5e0dc1b-0480-54bc-9374-aad01c23163d" begin
    Juno.render(i::Juno.Inline, measure::Measurement) =
    Juno.render(i, Juno.Row(measure.val, Text(" ± "), measure.err))

    Juno.Row(measure.val, Text(" ± "), measure.err)

    function Juno.render(ji::Juno.Inline, cm::Complex{<:Measurement})
        r, i = reim(cm)
        if signbit(i) && !isnan(i)
            i = -i
            sss = " - "
        else
            sss = " + "
        end
        Juno.render(ji, Juno.Row("(", Juno.render(ji, r), ")", sss,
        "(", Juno.render(ji, i), ")im"))
    end
end

include("conversions.jl")
include("comparisons-tests.jl")
include("utils.jl")
include("math.jl")
include("parsing.jl")
include("plot-recipes.jl")

end # module


text:
module Portfolio

    using Pipe: @pipe
    using TimeSeries
    import StatsBase: std, mean

    include("types.jl")
    include("dsp.jl")
    include("performance-analytics.jl")
    include("operations.jl")
    include("signals.jl")
    include("strategies-single.jl")
    include("strategies-portfolio.jl")
    include("statistics.jl")
    include("utils.jl")
    include("rebalancing.jl")
    include("compute.jl")

    export
        # types
        Signal, Weights, Returns,
        PortfolioReturns, Brokerage,
        # dsp
        highpass, lowpass,
        # statistics
        semideviation,
        # utils
        equity_curve,
        # rebalancing
        rebalance_every_n_days, rebelance_once,
        # performance-analytics
        returns_annualized, std_annualized, semid_annualized, sharpe_annualized, sortino_annualized,
        # operations
        ensemble_parameters, merge_weights,
        # compute
        run_weights,
        # strategies-single
        strategy_positive_mom_binary, strategy_prive_above_sma_binary, strategy_vol_target_mean, strategy_vol_target, strategy_vol_threshold, strategy_semideviation_target_mean, strategy_semideviation_threshold,
        # strategies-portfolio
        strategy_static_weights, strategy_static_equal_weights

end


text:
n=rand(1:10)
H = Hankel([-n:n])
@test diag(full(H)) == [-n:2:n]


text:
push!(LOAD_PATH,abspath(@__DIR__,"..","..","src"))
using Documenter
using SmoothedSpectralAbscissa ; const SSA=SmoothedSpectralAbscissa
using Plots,NamedColors
using LinearAlgebra
using Random

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

makedocs(;
    modules=[SmoothedSpectralAbscissa],
    authors="Dylan Festa",
    repo="https://github.com/dylanfesta/SmoothedSpectralAbscissa.jl/blob/{commit}{path}#{line}",
    sitename="SmoothedSpectralAbscissa.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://dylanfesta.github.io/SmoothedSpectralAbscissa.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/dylanfesta/SmoothedSpectralAbscissa.jl",
    devbranch="main",
    devurl="dev",
    versions = ["stable" => "v^", "v#.#"]
)