text:
module Stacks

export Stack, StackUnderflowException, StackOverflowException
export peek, isfull

struct Stack{T}
  data::Vector{T}
  size::Int

  Stack{T}(data::Vector{T}, size::Int) where {T} = begin
    length(data) > size && throw(StackOverflowException("The length of the data exceeds the size of the stack"))
    size < 0 && throw(StackUnderflowException("The size of the stack can not be negative"))

    new(data, size)
  end
end

struct StackUnderflowException <: Exception
  msg
end

struct StackOverflowException <: Exception
  msg
end

Stack(data::Vector{T}, size::Int) where {T} = Stack{T}(data, size)
Stack(data::Vector{T}) where {T} = Stack{T}(data, length(data))

Stack(itr, size) where {T} = Stack{eltype(itr)}([itr...], size)
Stack(itr) where {T} = Stack{eltype(itr)}([itr...], length(itr))

# Stack(t::Type{T}, size::Int) where {T} = Stack{T}(Vector{T}(undef, size), size)

function Base.push!(s::Stack{T}, elems::Vararg{T}) where {T}
  length(s.data) + length(elems) > s.size && throw(StackOverflowException("Pushing $elems exceeds the size of the stack"))
  push!(s.data, elems...)

  s
end

function Base.pop!(s::Stack{T}) where {T}
  length(s.data) > 0 || throw(StackUnderflowException("Stack can not be empty"))

  pop!(s.data)
end

function peek(s::Stack{T}) where {T}
  length(s.data) > 0 || throw(StackUnderflowException("Stack can not be empty"))

  s.data[end]
end

function Base.isempty(s::Stack{T}) where {T}
  isempty(s.data)
end

function Base.length(s::Stack{T}) where {T}
  length(s.data)
end

function Base.size(s::Stack{T}) where {T}
  s.size
end

function isfull(s::Stack{T}) where {T}
  length(s) == size(s)
end

function Base.empty!(s::Stack{T}) where{T}
  empty!(s.data)
  s
end

function Base.iterate(s::Stack{T}, state::Union{Stack,Nothing} = s) where {T}
  (state === nothing || isempty(s)) && return nothing
  pop!(s), s
end

function Base.Array(s::Stack{T}) where {T}
  s.data |> reverse
end

function Base.eltype(s::Stack{T}) where {T}
  eltype(s.data)
end

end

max_stars_repo_path: Chapter 6/Stacks.jl

max_issues_repo_path: Chapter 6/Stacks.jl

max_forks_repo_path: Chapter 6/Stacks.jl


text:
push!(LOAD_PATH,"../src/")

using Documenter, PartonDensity
using Literate

# Generate examples notebooks

gen_content_dir = joinpath(@__DIR__, "src")

pdf_parametrisation_src = joinpath(@__DIR__, "..", "examples", "pdf_parametrisation.jl")
transfer_matrix_src = joinpath(@__DIR__, "..", "examples", "transfer_matrix.jl")
forward_model_src = joinpath(@__DIR__, "..", "examples", "forward_model.jl")
prior_bat_src = joinpath(@__DIR__, "..", "examples", "prior_bat_implementation.jl")

Literate.markdown(pdf_parametrisation_src, gen_content_dir, name="pdf_parametrisation")
Literate.markdown(transfer_matrix_src, gen_content_dir, name="transfer_matrix")
Literate.markdown(forward_model_src, gen_content_dir, name="forward_model")
Literate.markdown(prior_bat_src, gen_content_dir, name="prior_bat_implementation")

Introduction = "Introduction" => "index.md"

Examples = "Examples" => ["pdf_parametrisation.md",
                          "transfer_matrix.md", "forward_model.md",
                          "prior_bat_implementation.md"]

API = "API" => "api.md"

PAGES = [Introduction, Examples, API]

makedocs(modules=[PartonDensity], sitename="PartonDensity.jl", pages=PAGES)

deploydocs(
    devbranch = "main",
    repo = "github.com/cescalara/PartonDensity.jl.git",
)

ext: jl


text:
#!/usr/bin/env julia

using Eigenbroetler

function magic_phase(x::Number, y::Number)
    r = sqrt(x * x + y * y)
    return exp(-(r / 100) ^ 2) * exp(2 * im * x * (x ^ 2 - y ^ 2) / 200 ^ 2)
end

eb = Eigenbrot(magic_phase, 512, 512)
save("magic_phase.png", eb, ImageSetting(Phase, LinearScale))
save("magic_phase_dft.png", fft(eb), ImageSetting(Phase, LinearScale))


text:
using Distributions

Sigma = [1. 0.; 0. 1.];

mu1 = [1., -1.];
x1 = rand(MvNormal(mu1, Sigma), 200);

mu2 = [5.5, -4.5];
x2 = rand(MvNormal(mu2, Sigma), 200);

mu3 = [1., 4.];
x3 = rand(MvNormal(mu3, Sigma), 200);

mu4 = [6., 4.5];
x4 = rand(MvNormal(mu4, Sigma), 200);

mu5 = [9., 0.0];
x5 = rand(MvNormal(mu5, Sigma), 200);

x = [x1 x2 x3 x4 x5]

println(repr(x))

max_stars_repo_name: zxteloiv/julia-ex

max_issues_repo_name: zxteloiv/julia-ex

max_forks_repo_name: zxteloiv/julia-ex


text:
## check all votables and create a blacklist
##

using PyCall
using Distributions, Statistics , StatsBase

using Distances ,Random, Dates
using Glob

import PyPlot
sns= pyimport("seaborn")

## directory
rootdir = ENV["GAIA_ROOT"]
wdir    = "$rootdir/e2e_products"
votdir  = "$wdir/votable-TEST"
plotdir = "$wdir/plots-TEST"

push!(LOAD_PATH,"$rootdir/run/src")
using GaiaClustering

cd(wdir)

function plot_data(votname, df)

    s= size(df.data)
    npts= min(10000, s[2])
    idx= randperm(MersenneTwister(1234), s[2])[1:npts]

    PyPlot.plt.figure(figsize=(12.0,13.0))

    PyPlot.plt.subplot(2, 2, 1)
    PyPlot.plt.plot(df.data[1,idx],df.data[2,idx],"k.", markersize= 1)
    PyPlot.xlabel("l")
    PyPlot.ylabel("b")

    PyPlot.plt.subplot(2, 2, 2)
    PyPlot.plt.plot(df.raw[1,idx],df.raw[2,idx],"k.", markersize= 1)
    PyPlot.xlabel("α")
    PyPlot.ylabel("δ")

    PyPlot.plt.subplot(2, 2, 3)
    PyPlot.plt.plot(df.raw[1,idx],df.raw[5,idx],"k.", markersize= 1)
    PyPlot.xlabel("α")
    PyPlot.ylabel("parallax")

    PyPlot.plt.subplot(2, 2, 4)
    PyPlot.plt.plot(df.raw[6,idx],df.raw[7,idx],"k.", markersize= 1)
    PyPlot.xlabel("PM(α)")
    PyPlot.ylabel("PM(δ)")

    figname = plotdir*"/"*votname*".gaia.png"
    PyPlot.plt.savefig(figname)

    # PyPlot.plt.show()
end

function main(votlist)

    println("## Analyzing/plotting the raw data...\n")
    s=size(votlist)
    for i in 1:s[1]
        votname= votlist[i]
        data= read_votable(votdir*"/"*votname)
        df= filter_data(data)

        println("## Analyzing $votname")
        plot_data(votname,df)

    end
end

cd(votdir)
votlist= glob("*.vot")
cd(wdir)

main(votlist)

max_stars_repo_name: bosscha/gaia-shock

max_issues_repo_name: bosscha/gaia-shock

max_forks_repo_name: bosscha/gaia-shock


text:
@testset "UmData" begin

  #Test dataframe checker
  df1 = DataFrame(a=[1,2],b=[3,4])
  @test isequal(df1, Unmarked.check_data(df1,2))
  @test_throws ErrorException Unmarked.check_data(df1,3)
  df2 = DataFrame(a=[1,missing],b=[3,4])
  @test_throws ErrorException Unmarked.check_data(df2,2)
  df3 = DataFrame(_dummy=[1,1])
  @test isequal(Unmarked.check_data(nothing, 2), df3)

  #Test basic constructor
  y = [[1 0]; [0 1]; [1 1]]
  sc = DataFrame(a=[1,2,3], b=[3,4,5])
  oc = DataFrame(c=collect(1:6), d=collect(1:6))
  dat = UmData(y, sc, oc)
  @test isequal(dat.y, y)
  @test isequal(dat.site_covs, sc)
  @test isequal(dat.obs_covs, oc)
  sc2 = DataFrame(a=[1,2], b=[3,4])
  scno = DataFrame(_dummy=[1.0,1.0,1.0])
  @test_throws ErrorException UmData(y, sc2, oc)
  dat2 = UmData(y, nothing, oc)
  @test isequal(dat2.site_covs, scno)

  #Test missing values
  yna = [[missing 0]; [0 1]; [1 1]]
  dat2 = UmData(yna, sc, oc)
  @test isequal(dat2.y, yna)

  #Test outer constructor
  dat3 = UmData(y, obs_covs=oc, site_covs=sc)
  @test isequal(dat3.site_covs, sc)
  dat4 = UmData(y, obs_covs=oc)
  @test isequal(dat4.site_covs, DataFrame(_dummy=[1.0,1.0,1.0]))
  @test_throws ErrorException UmData(y, site_covs=oc, obs_covs=oc)
  dat5 = UmData(yna)
  @test isequal(dat5.y, yna)

  #Test subsetting
  ds1 = dat[1]
  @test isequal(size(ds1.y), (1,2))
  @test isequal(ds1.y[1,:], y[1,:])
  @test isequal(ds1.site_covs, sc[1:1,:])
  @test isequal(ds1.obs_covs, oc[1:2,:])

  ds2 = dat[2:3]
  ds3 = dat[[2,3]]
  @test isequal(ds2.y, ds3.y)
  @test isequal(ds2.y, y[2:3,:])
  @test isequal(ds2.site_covs, sc[2:3,:])
  @test isequal(ds2.obs_covs, oc[3:6,:])

end

max_stars_repo_name: kenkellner/julia-hme

max_issues_repo_name: kenkellner/julia-hme

max_forks_repo_name: kenkellner/julia-hme


text:
# by lazarusA # HIDE
#using GLMakie, Random # HIDE
using CairoMakie, Random
CairoMakie.activate!() #HIDE
let
    Random.seed!(13)
    n = 3000
    data = randn(n)
    a = fill(1, n)
    fig = Figure(resolution = (700, 800), font = "sans")
    ax1 = Axis(fig, xlabel = "variable", ylabel = "values",
                xticks = ([1],["normal Distribution"]) )
    ax2 = Axis(fig, xlabel = "values", ylabel = "variable",
        yticks = ([1],["normal Distribution"]), yticklabelrotation = pi/2)
    boxplot!(ax1, a, data; whiskerwidth = 1, width = 0.35,color = (:red, 0.45),
                whiskercolor = (:red, 1), mediancolor = :red, markersize = 8,
                strokecolor = :black, strokewidth = 1, label = "vertical")
    limits!(ax1, 0,2,-5,5)
    boxplot!(ax2,a,data;orientation=:horizontal, whiskerwidth = 1, width = 0.35,
            color = (:navy, 0.45),whiskercolor = (:navy, 1),mediancolor = :navy,
                markersize = 8, strokecolor = :black, strokewidth = 1,
                label = "horizontal")
    limits!(ax2, -5,5,0,2)
    axislegend(ax1, position = :rt)
    axislegend(ax2, position = :rt)
    fig[1,1] = ax1
    fig[2,1] = ax2
    fig
    save(joinpath(@__DIR__, "output", "BoxPlot.png"), current_figure(), px_per_unit = 2.0) # HIDE
end
using Pkg # HIDE
Pkg.status(["CairoMakie", "Random"]) # HIDE

max_stars_repo_name: MisterBiggs/BeautifulMakie

max_issues_repo_name: MisterBiggs/BeautifulMakie

max_forks_repo_name: MisterBiggs/BeautifulMakie


text:
include("filtering.jl");
##
odc = ODC(pokes)
@df filter(r -> r.PokeDuration < 4,odc) density(:PokeDuration,group = :Reward)
plot!(repeat([0.43],2), [0,8])
plot!(repeat([0.7],2), [0,8])
odc[!,:ClearDuration] = [r.Reward ? r.PokeDuration > 0.7 : r.PokeDuration < 0.4 for r in eachrow(odc)]
list = ["PreVehicle",
    "Altanserin",
    "SB242084",
    "Way_100135",
    "Methysergide",
    "Citalopram",
    "SB242084_opt",
    "Saline"]

filter!(r -> r.Treatment in list &&
    !ismissing(r.Pre_Interpoke) &&
    r.Pre_Interpoke < 0.5 &&
    r.ClearDuration &&
    r.Trial < 51 &&
    r.MouseID != "d5" &&
    r.MouseID != "pc7",# ["pc1","pc2","pc3","pc4","pc5","pc6","pc8","pc9","pc10"],
    odc)

odc[odc.Treatment .== "PreVehicle",:Treatment] .= "Control"
odc[odc.Treatment .== "Saline",:Treatment] = [o ? "Optogenetic" : "Control" for o in odc[odc.Treatment .== "Saline",:Stim]]
odc[odc.Treatment .== "SB242084_opt",:Phase] .=  "SB242084_opt"
odc[odc.Treatment .== "SB242084_opt",:Treatment] = [o ? "SB242084_opt" : "Control" for o in odc[odc.Treatment .== "SB242084_opt",:Stim]]

df1 = filter(r -> !r.Reward &&
    r.ExpDay > 3 &&
    !ismissing(r.PokeFromLastRew) &&
    0 < r.PokeFromLastRew &&
    !ismissing(r.ODC) &&
    r.Treatment in ["Control","Citalopram", "Optogenetic", "Altanserin","SB242084"] &&
    !in(r.Phase,["Methysergide", "Way_100135","SB242084_opt"]) ,odc)

categorical!(df1,[:MouseID,:Treatment,:Phase])

levels!(df1.Treatment, ["Control",
    "Citalopram",
    "Optogenetic",
    "Altanserin",
    "SB242084"])

transform!(groupby(df1,[:MouseID,:Phase,:Treatment]), :ODC => mean => :ODC_mean)
df1[!,:NODC] = df1.ODC ./ df1.ODC_mean
transform!(groupby(df1,[:MouseID,:Phase,:Treatment]), :NODC => binquantile => :QODC)
df1[!,:Leave] = [r == 0 for r in df1.PokeFromLeaving]
## Prepare 2 separate dataframes for pharma and opto
# pharma
################# from last reward
odc_fromlast = @formula(QODC ~ PokeFromLastRew * Treatment + (1|MouseID))
filtered = filter(r -> r.PokeFromLastRew < 11, df1)
cit = filter(r -> r.Phase == "Citalopram", filtered)
opto = filter(r -> r.Phase == "Optogenetic", filtered)
alt = filter(r -> r.Phase == "Altanserin", filtered)
sb = filter(r -> r.Phase == "SB242084", filtered)
cit_fromlast = fit!(LinearMixedModel(odc_fromlast,cit))
sb_fromlast = fit!(LinearMixedModel(odc_fromlast,sb))
opt_fromlast = fit!(LinearMixedModel(odc_fromlast,opto))
alt_fromlast = fit!(LinearMixedModel(odc_fromlast,alt))
ps = plot_QODC(filtered,:PokeFromLastRew)
savefig(joinpath(figs_loc,"ODC","ODC_from_last","QODC_lastrew50.pdf"))
########################## from leaving
odc_fromleave = @formula(QODC ~ PokeFromLeaving * Treatment + (1|MouseID))
filtered = filter(r -> r.PokeFromLeaving < 5, df1)
cit = filter(r -> r.Phase == "Citalopram", filtered)
opto = filter(r -> r.Phase == "Optogenetic", filtered)
alt = filter(r -> r.Phase == "Altanserin", filtered)
sb = filter(r -> r.Phase == "SB242084", filtered)
# cit_fromleave = fit!(LinearMixedModel(odc_fromleave,cit))
# sb_fromleave = fit!(LinearMixedModel(odc_fromleave,sb))
# opt_fromleave = fit!(LinearMixedModel(odc_fromleave,opto))
# alt_fromleave = fit!(LinearMixedModel(odc_fromleave,alt))

cit_fromleave = GeneralizedLinearMixedModel(odc_fromleave,cit,Poisson())
sb_fromleave = GeneralizedLinearMixedModel(odc_fromleave,sb,Poisson())
opt_fromleave = GeneralizedLinearMixedModel(odc_fromleave,opto,Poisson())
alt_fromleave = GeneralizedLinearMixedModel(odc_fromleave,alt,Poisson())
plot_QODC(filtered,:PokeFromLeaving; xflip = true)
savefig(joinpath(figs_loc,"ODC","ODC_from_leaving","QODC_leaving60.pdf"))
## test QODC exactly at leaving
filtered = filter(r -> r.PokeFromLeaving == 0, df1)
gd2 = groupby(filtered,[:MouseID,:Phase,:Treatment])
df2 = combine(gd2, :QODC => mean => :QODC)
gd3 = groupby(df2,[:Phase])
df3 = combine(gd3) do dd
    dd[:,:Treatment] = [t == "Control" ? t : "Manipulation" for t in dd.Treatment]
    unstack(dd,:Treatment, :QODC)
end
dropmissing!(df3)
gd4 = groupby(df3,:Phase)
df4 = combine(gd4) do dd
    wilcoxon(dd,:Control,:Manipulation)
end
df4[!,:Phase] = String.(df4.Phase)
Drug_colors!(df4)
plot_wilcoxon(df4)
savefig(joinpath(figs_loc,"ODC","ODC_from_leaving","WilcoxonQODC_leaving50trials.pdf"))
## heatmap

############################ pharma ######################################
basic_pharma = @formula(Leave ~ Poke_within_Trial + (1|MouseID));
odc_pharma = @formula(Leave ~ ODC * Poke_within_Trial + (1|MouseID));
full_pharma = @formula(Leave ~ ODC * Poke_within_Trial * Treatment  + (1|MouseID));
fBP = GeneralizedLinearMixedModel(basic_pharma, pharma,Bernoulli())
fOP = GeneralizedLinearMixedModel(odc_pharma, pharma,Bernoulli())
fFP = GeneralizedLinearMixedModel(full_pharma, pharma,Bernoulli())
AICcTest(fOP,fBP)
LoglikelihoodRatioTest(fBP,fOP)
AICcTest(fFP,fOP)
LoglikelihoodRatioTest(fOP,fFP)
############################ opto ######################################
basic_opto = @formula(Leave ~ Poke_within_Trial + (1|MouseID));
odc_opto = @formula(Leave ~ ODC * Poke_within_Trial + (1|MouseID));
full_opto = @formula(Leave ~ ODC * Poke_within_Trial * Treatment  + (1|MouseID));
fBO = GeneralizedLinearMixedModel(basic_opto, opto,Bernoulli())
fOO = GeneralizedLinearMixedModel(odc_opto, opto,Bernoulli())
fFO = GeneralizedLinearMixedModel(full_opto, opto,Bernoulli())
AICcTest(fOO,fBO)
LoglikelihoodRatioTest(fBO,fOO)
AICcTest(fFO,fOO)
LoglikelihoodRatioTest(fOO,fFO)
############################ pharma ######################################
basic_pharma = @formula(Leave ~ PokeFromLastRew + (1|MouseID));
odc_pharma = @formula(Leave ~ ODC * PokeFromLastRew + (1|MouseID));
full_pharma = @formula(Leave ~ ODC * PokeFromLastRew * Treatment  + (1|MouseID));
alt_pharma = @formula(Leave ~ PokeFromLastRew + ODC * Treatment  + (1|MouseID));
fBP = GeneralizedLinearMixedModel(basic_pharma, pharma,Bernoulli())
fOP = GeneralizedLinearMixedModel(odc_pharma, pharma,Bernoulli())
fFP = GeneralizedLinearMixedModel(full_pharma, pharma,Bernoulli())
f_alt = GeneralizedLinearMixedModel(alt_pharma,pharma,Bernoulli())
AICcTest(fOP,fBP)
LoglikelihoodRatioTest(fBP,fOP)
AICcTest(fFP,fOP)
LoglikelihoodRatioTest(fOP,fFP)
############################ opto ######################################
basic_opto = @formula(Leave ~ Poke_within_Trial + (1|MouseID));
odc_opto = @formula(Leave ~ ODC * Poke_within_Trial + (1|MouseID));
full_opto = @formula(Leave ~ ODC * Poke_within_Trial * Treatment  + (1|MouseID));
fBO = GeneralizedLinearMixedModel(basic_opto, opto,Bernoulli())
fOO = GeneralizedLinearMixedModel(odc_opto, opto,Bernoulli())
fFO = GeneralizedLinearMixedModel(full_opto, opto,Bernoulli())
AICcTest(fOO,fBO)
LoglikelihoodRatioTest(fBO,fOO)
AICcTest(fFO,fOO)
LoglikelihoodRatioTest(fOO,fFO)
##############################################################################
predict(fFP)
# afterlast
# afterlast + treatment
# afteralst + treatment + ODC
# plot histogram from data p leaving over binned duty cycle
# heatmap pleaving xaxis pokesafterlast yaxis binned duty cycle
##############################################################################
cit = filter(r -> r.Phase == "Citalopram", pharma)
treat = @formula(Leave ~ PokeFromLastRew + Treatment  + (1|MouseID));
with_odc = @formula(Leave ~ PokeFromLastRew + Treatment + ODC  + (1|MouseID));
duty = @formula(Leave ~  Treatment + ODC  + (1|MouseID));
cit_treat = GeneralizedLinearMixedModel(treat, cit,Bernoulli())
cit_odc = GeneralizedLinearMixedModel(with_odc, cit,Bernoulli())
cit_duty = GeneralizedLinearMixedModel(duty, cit,Bernoulli())
opt_treat = GeneralizedLinearMixedModel(treat, opto,Bernoulli())
opt_odc = GeneralizedLinearMixedModel(with_odc, opto,Bernoulli())
opt_duty = GeneralizedLinearMixedModel(duty, opto,Bernoulli())

##############################################################################
# ODC ~ afterlast * Treatment
# ODC ~ poke_toleave * treatment this shoud be opto
odc_after = @formula(ODC ~ PokeFromLastRew  * Treatment + (1|MouseID))
cit_after = fit!(LinearMixedModel(odc_after,cit))
opto_after = fit!(LinearMixedModel(odc_after,opto))
union(cit.Treatment)
cit2  = filter(r -> r.PokeFromLastRew < 6, cit)
cit2_after = fit!(LinearMixedModel(odc_after,cit2))

ll =-2(loglikelihood(fOP) - loglikelihood(fFP))

# check odc and pokeafterlastrew weight
open_html_table(df1[1:200,[:Poke,:Trial,:PokeFromLastRew,:PokeFromLeaving,:Leave]])
scatter(df1.Leave,df1.PokeFromLastRew)
m_basic = @formula(Leave ~ PokeFromLastRew + (1|MouseID));
f_basic = GeneralizedLinearMixedModel(m_basic, df1, Bernoulli());
m_odc = @formula(Leave ~ ODC * PokeFromLastRew + (1|MouseID));
f_odc = GeneralizedLinearMixedModel(m_basic, df1, Bernoulli());
AICcTest(f_odc, f_basic)
# full model
m_full = @formula(Leave ~ ODC * PokeFromLastRew * Treatment  + (1|MouseID));
f_full = GeneralizedLinearMixedModel(m_full, df1, Bernoulli());
# simplified model without significant simple factors
m_a1 = @formula(Leave ~ ODC + PokeFromLastRew +  Treatment  + (1|MouseID));
f_without_interaction = GeneralizedLinearMixedModel(m_a1, df1, Bernoulli());
m_a2 = @formula(Leave ~ ODC + PokeFromLastRew + ODC & PokeFromLastRew + ODC & Treatment & PokeFromLastRew  + (1|MouseID));
f_fromlast_interaction = GeneralizedLinearMixedModel(m_a2, df1, Bernoulli());
AICcTest(f_full,f_without_interaction)
AICcTest(f_full,f_fromlast_interaction)
## Plot regression
df2 = combine(groupby(df1,[:Treatment])) do dd
    summarize(dd,:PokeFromLastRew, :ODC)
end
Drug_colors!(df2)
@df filter(r -> r.Treatment in ["Control", "Citalopram"], df2) scatter(:Xaxis,:Mean, group = :Treatment, color = :color, markersize = 7)
@df filter(r -> r.Treatment in ["Control", "Optogenetic"], df2) scatter(:Xaxis,:Mean, group = :Treatment, color = :color, markersize = 7)
@df filter(r -> r.Treatment in ["Control",  "SB242084"], df2) scatter(:Xaxis,:Mean, group = :Treatment, color = :color, markersize = 7)
@df filter(r -> r.Treatment in ["Control", "Altanserin"], df2) scatter(:Xaxis,:Mean, group = :Treatment, color = :color, markersize = 7)

df2 = combine(groupby(df1,[:Phase,:Treatment])) do dd
    summarize(dd,:PokeFromLastRew, :ODC)
end
Drug_colors!(df2)
@df filter(r -> r.Phase == "Citalopram", df2) scatter(:Xaxis,:Mean, group = :Treatment, color = :color, markersize = 7)
@df filter(r -> r.Phase == "Optogenetic", df2) plot(:Xaxis,:Mean, group = :Treatment, linecolor = :color, markersize = 7)
@df filter(r -> r.Phase ==  "SB242084", df2) scatter(:Xaxis,:Mean, group = :Treatment, color = :color, markersize = 7)
@df filter(r -> r.Phase == "Altanserin", df2) plot(:Xaxis,:Mean, group = :Treatment, linecolor = :color, markersize = 7)
##Plot deltas
transform!(groupby(df1,[:MouseID,:Phase]), :ODC => mean)
df1[!,:ODC_norm] = df1.ODC ./ df1.ODC_mean
df2 = combine(groupby(df1,[:Treatment,:MouseID,:PokeFromLastRew])) do dd
    (ODC_norm = mean(dd.ODC_norm),)
end
open_html_table(df2)
df3 = unstack(df2,:Treatment,:ODC_norm)
for c in ["Citalopram","Optogenetic","Altanserin","SB242084"]
    df3[!,Symbol("Delta" * c)] = df3[:,Symbol(c)] - df3[:,:Control]
end
l = 18
plot_deltafromrew(df3,"Optogenetic"; limit = l)
plot_deltafromrew(df3,"Altanserin"; limit = l)
plot_deltafromrew(df3,"Citalopram"; limit = 11)
plot_deltafromrew(df3,"SB242084"; limit = 10)
# replot simple 2 lines
df4 = combine(groupby(df3,:MouseID), :DeltaOptogenetic => skipmean)
opto = filter(r -> !isnan(r.DeltaOptogenetic_skipmean),df4)
mean(opto.DeltaOptogenetic_skipmean)
## All controls joined
combined_controls = combine(groupby(df1,[:Treatment])) do dd
    summarize(dd,:PokeFromLastRew,:ODC_norm)
end
Drug_colors!(combined_controls)
l = 100
ongoing = filter(r -> r.Treatment in ["Control", "Altanserin"] && r.Xaxis < l, combined_controls)
@df  ongoing plot(:Xaxis,:Mean,
    ribbon = :SEM, group = :Treatment, color =:color)

ongoing = filter(r -> r.Treatment in ["Control", "Optogenetic"] && r.Xaxis < l, combined_controls)
@df ongoing plot(:Xaxis,:Mean,
    ribbon = :SEM, group = :Treatment, color =:color)

##
separate_controls = combine(groupby(df1,[:Phase,:Treatment])) do dd
    summarize(dd,:PokeFromLastRew,:ODC_norm)
end
Drug_colors!(separate_controls)
l = 12
ongoing = filter(r -> r.Phase == "Altanserin" && r.Xaxis < l, separate_controls)
@df ongoing plot(:Xaxis,:Mean, ribbon = :SEM, color = :color, group = :Treatment)
ongoing = filter(r -> r.Phase == "SB242084" && r.Xaxis < l, separate_controls)
@df ongoing plot(:Xaxis,:Mean, ribbon = :SEM, color = :color, group = :Treatment)
##
open_html_table(df3)
gd = groupby(df3,[:PokeFromLastRew])

df4 = combine(gd,:DeltaCitalopram => skipmean => :DeltaCitalopram,
    :DeltaOptogenetic => skipmean => :DeltaOptogenetic,
    :DeltaAltanserin=> skipmean => :DeltaAltanserin,
    :DeltaSB242084 => skipmean => :DeltaSB242084)
open_html_table(df4)

phase = "Optogenetic"
c_name = "Delta" * phase
c_symbol = Symbol(c_name)
c_symbol
df3.PokeFromLastRew
df3.DeltaOptogenetic
@df filter(r -> r.PokeFromLastRew < 15, df3) scatter(:PokeFromLastRew, cols(c_symbol), markercolor = :grey)
@df filter(r -> r.PokeFromLastRew < 15, df4) scatter!(:PokeFromLastRew, c_symbol, markersize = 8, markercolor = get(drug_colors,phase,:red))
Plots.abline!(0,0, legend = false)

@df filter(r -> r.PokeFromLastRew < 15, df3) scatter(:PokeFromLastRew, :DeltaAltanserin)
Plots.abline!(0,0)
##
c_baseline = f_full.β[1]
b_odc = f_full.β[2]
b_odcrew = f_full.β[8] ##ODC&PokeFromLast
b_sb = f_full.β[20]
b_alt = f_full.β[19]
b_opt = f_full.β[18]
b_cit = f_full.β[17]
r = range(minimum(df1.ODC),maximum(df1.ODC), length = 10)
plot(r, r .* b_odc)
## Plot logistic sinusoid
params = f_full.β
intercept = params[1]
O_D_C = params[2]
PFLR = params[3]
Citalopram = params[4]
Optogenetic = params[5]
Altanserin = params[6]
SB242084 = params[7]
ODC_PFLR = params[8]
ODC_Citalopram = params[9]
ODC_Optogenetic = params[10]
ODC_Altanserin = params[11]
ODC_SB242084 = params[12]
PFLR_Citalopram = params[13]
PFLR_Optogenetic = params[14]
PFLR_Altanserin = params[15]
PFLR_SB242084 = params[16]
ODC_PFLR_Citalopram = params[17]
ODC_PFLR_Optogenetic = params[18]
ODC_PFLR_Altanserin = params[19]
ODC_PFLR_SB242084 = params[20]

# by plugging in the mean as the value for odc,
# I'll be generating plots that show the relationship between PokesAfterLast and
# the outcome "for someone with an average ODC".
odc_val = minimum(df1.ODC)
afterlast_range = 1:10
# the reference group
a_logits = [intercept +
    O_D_C * odc_val +
    PFLR * af +
    Citalopram * 0 +
    Optogenetic * 0 +
    Altanserin * 0 +
    SB242084 * 0 +
    ODC_PFLR * odc_val * af +
    ODC_Citalopram * odc_val * 0 +
    ODC_Optogenetic * odc_val * 0 +
    ODC_Altanserin * odc_val * 0 +
    ODC_SB242084 * odc_val * 0 +
    PFLR_Citalopram * af * 0 +
    PFLR_Optogenetic * af * 0 +
    PFLR_Altanserin * af * 0 +
    PFLR_SB242084 * af * 0 +
    ODC_PFLR_Citalopram * odc_val * af +
    ODC_PFLR_Optogenetic * odc_val * af +
    ODC_PFLR_Altanserin * odc_val * af +
    ODC_PFLR_SB242084 * odc_val * af for af in afterlast_range]

a_probs = [exp(a) / (1 + exp(a)) for a in a_logits]
plot(1 .- a_probs)
bare = intercept + O_D_C * odc_val
exp(bare)/(1 + exp(bare))

d = countmap(df1[df1[:,:Leave],:PokeFromLastRew])
p = plot()
for (k,v) in d
    scatter!(k,v)
end
p


text:
using IcanProj
using JLD2
using PyPlot

"""
Plot analysis using Budyko kind of equations.
"""
function plot_budyko(opt, watershed)

    # Various equations to compute aet from prec and pet

    schreiber(pet_div_p) = 1 - exp(-pet_div_p)

    pike(pet_div_p) = 1 / sqrt(1 + (1/pet_div_p)^2)

    budyko(pet_div_p) = sqrt(pet_div_p * (1 - exp(-pet_div_p)) * tanh(1/pet_div_p))

    zhang(pet_div_p, w) = (1 + w*pet_div_p) / (1 + w*pet_div_p + 1/pet_div_p)

    # Load results

    file = joinpath(opt["eval_folder"], watershed, "res_1km.jld2")

    @load file res_vic

    # Get precipitation and potential evapotranspiration

    ivar = find(res_vic.var_names .== "prec")

    data_prec = res_vic.data_all[2:end,ivar[1],:]

    ivar = find(res_vic.var_names .== "pet_short")

    data_pet = res_vic.data_all[2:end,ivar[1],:]

    ivar = find(res_vic.var_names .== "evap")

    data_aet = res_vic.data_all[2:end,ivar[1],:]

    # Compute yearly average prec and pet for each gridcell

    mean_prec = 365 * mean(data_prec, 1)

    mean_pet = 365 * mean(data_pet, 1)

    mean_aet = 365 * mean(data_aet, 1)

    pet_div_p = mean_pet./mean_prec

    aet_div_p = mean_aet./mean_prec

    # Compute error in aet estimates by using the mean of the function or
    # the function of the mean

    zhang(x) = zhang(x, 0.3)

    aet_mean_func = []
    aet_func_mean = []

    for f in [schreiber, pike, budyko, zhang]

        push!(aet_mean_func, mean(f.(pet_div_p) .* mean_prec ))
        push!(aet_func_mean, f(mean(pet_div_p)) * mean(mean_prec))

    end

    # The error should be positive

    error_aet = aet_func_mean - aet_mean_func

    error_min = round(minimum(error_aet), 1)
    error_max = round(maximum(error_aet), 1)

    # Plot results

    ioff()

    fig = figure(figsize=(7, 5))

    ax = fig[:add_subplot](111)

    x_values = collect(0.001:0.001:2)

    plot(x_values, schreiber.(x_values), label = "schreiber")
    plot(x_values, pike.(x_values), label = "pike")
    plot(x_values, budyko.(x_values), label = "budyko")
    plot(x_values, zhang.(x_values, 0.3), label = "zhang (w = 0.3)")

    scatter(pet_div_p, aet_div_p)

    min_ratio = minimum(pet_div_p)
    max_ratio = maximum(pet_div_p)

    axvspan(min_ratio, max_ratio, alpha=0.5, color="gray", label = "Spread in observed\nconditions")

    error_str = "Min error = $(error_min) mm/year\nMax error = $(error_max) mm/year"

    ax[:text](0.6, 0.5, error_str, transform=ax[:transAxes])

    xlabel("PET / P")
    ylabel("ET / P")
    title(res_vic.name)
    legend()

    file_name = Pkg.dir("IcanProj", "plots", "budyko_$(watershed).png")

    rm(file_name, force=true)
    savefig(file_name, dpi = 300)

    close(fig)

end

# Get options

opt = get_options()

for watershed in opt["stat_sel"]

    plot_budyko(opt, watershed)

end


text:
using zmqtest
using Test

@testset "zmqtest.jl" begin
    # Write your tests here.
end

max_stars_repo_name: jwdevantier/julia-zmqtest

max_issues_repo_name: jwdevantier/julia-zmqtest

max_forks_repo_name: jwdevantier/julia-zmqtest


text:
#*******************************************************/
#* Copyright(c) 2018 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# This example demonstrates how to use Knitro to solve the following
# simple nonlinear optimization problem, while varying user
# options and bounds on variables or constraints.
# This model is test problem HS15 from the Hock & Schittkowski
# collection.
#
# min   100(x1 - x0^2)^2 +(1 - x0)^2
# s.t.  x0 x1 >= c0lb(initially 1)
#       x0 + x1^2 >= 0
#       x0 <= x0ub   (initially 0.5)
#
# We first solve the model with c0lb=1 and x0ub=0.5, and then
# re-solve for different values of the "bar_murule" user option.
# We then re-solve the model, while changing the value of the
# variable bound "x0ub".  Finally, we re-solve while changing the
# value of the constraint bound "c0lb".
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using KNITRO
using Printf

function example_restart(; verbose=true)
    #*------------------------------------------------------------------*
    #*     FUNCTION callbackEvalF                                       *
    #*------------------------------------------------------------------*
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.jl.
    # Only "obj" is set in the KNITRO.KN_eval_result structure.
    function callbackEvalF(kc, cb, evalRequest, evalResult, userParams)
        x = evalRequest.x

        # Evaluate nonlinear objective
        dTmp = x[2] - x[1] * x[1]
        evalResult.obj[1] = 100.0 * (dTmp * dTmp) + ((1.0 - x[1]) * (1.0 - x[1]))

        return 0
    end

    #*------------------------------------------------------------------*
    #*     FUNCTION callbackEvalG                                       *
    #*------------------------------------------------------------------*
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.h.
    # Only "objGrad" is set in the KNITRO.KN_eval_result structure.
    function callbackEvalG!(kc, cb, evalRequest, evalResult, userParams)
        x = evalRequest.x

        # Evaluate gradient of nonlinear objective
        dTmp = x[2] - x[1] * x[1]
        evalResult.objGrad[1] = (-400.0 * dTmp * x[1]) - (2.0 * (1.0 - x[1]))
        evalResult.objGrad[2] = 200.0 * dTmp

        return 0
    end

    #*------------------------------------------------------------------*
    #*     FUNCTION callbackEvalH                                       *
    #*------------------------------------------------------------------*
    # The signature of this function matches KNITRO.KN_eval_callback in knitro.h.
    # Only "hess" and "hessVec" are set in the KNITRO.KN_eval_result structure.
    function callbackEvalH!(kc, cb, evalRequest, evalResult, userParams)
        x = evalRequest.x
        # Scale objective component of hessian by sigma
        sigma = evalRequest.sigma

        # Evaluate the hessian of the nonlinear objective.
        # Note: Since the Hessian is symmetric, we only provide the
        #       nonzero elements in the upper triangle(plus diagonal).
        #       These are provided in row major ordering as specified
        #       by the setting KNITRO.KN_DENSE_ROWMAJOR in "KNITRO.KN_set_cb_hess()".
        # Note: The Hessian terms for the quadratic constraints
        #       will be added internally by Knitro to form
        #       the full Hessian of the Lagrangian.
        evalResult.hess[1] = sigma * ((-400.0 * x[2]) + (1200.0 * x[1] * x[1]) + 2.0) #(0,0)
        evalResult.hess[2] = sigma * (-400.0 * x[1]) #(0,1)
        evalResult.hess[3] = sigma * 200.0           #(1,1)

        return 0
    end

    #*------------------------------------------------------------------*
    #*     main                                                         *
    #*------------------------------------------------------------------*

    # Create a new Knitro solver instance.
    kc = KNITRO.KN_new()

    # Illustrate how to override default options by reading from
    # the knitro.opt file.
    options = joinpath(dirname(@__FILE__), "..", "examples", "knitro.opt")
    KNITRO.KN_load_param_file(kc, options)
    kn_outlev = verbose ? KNITRO.KN_OUTLEV_ITER : KNITRO.KN_OUTLEV_NONE
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_OUTLEV, kn_outlev)

    # Initialize knitro with the problem definition.

    # Add the 4 variables and set their bounds.
    # Note: any unset lower bounds are assumed to be
    # unbounded below and any unset upper bounds are
    # assumed to be unbounded above.
    KNITRO.KN_add_vars(kc, 2)
    KNITRO.KN_set_var_lobnds(kc, [-KNITRO.KN_INFINITY, -KNITRO.KN_INFINITY]) # not necessary since infinite
    KNITRO.KN_set_var_upbnds(kc, [0.5, KNITRO.KN_INFINITY])

    # Add the constraints and set their lower bounds
    KNITRO.KN_add_cons(kc, 2)
    KNITRO.KN_set_con_lobnds(kc, [1.0, 0.0])

    # Both constraints are quadratic so we can directly load all the
    # structure for these constraints.

    # First load quadratic structure x0*x1 for the first constraint
    KNITRO.KN_add_con_quadratic_struct(kc, 0, 0, 1, 1.0)

    # Load structure for the second constraint.  below we add the linear
    # structure and the quadratic structure separately, though it
    # is possible to add both together in one call to
    # "KNITRO.KN_add_con_quadratic_struct()" since this api function also
    # supports adding linear terms.

    # Add linear term x0 in the second constraint
    KNITRO.KN_add_con_linear_struct(kc, 1, 0, 1.0)

    # Add quadratic term x1^2 in the second constraint
    KNITRO.KN_add_con_quadratic_struct(kc, 1, 1, 1, 1.0)

    # Add a callback function "callbackEvalF" to evaluate the nonlinear
    #(non-quadratic) objective.  Note that the linear and
    # quadratic terms in the objective could be loaded separately
    # via "KNITRO.KN_add_obj_linear_struct()" / "KNITRO.KN_add_obj_quadratic_struct()".
    # However, for simplicity, we evaluate the whole objective
    # function through the callback.
    cb = KNITRO.KN_add_objective_callback(kc, callbackEvalF)

    # Also add a callback function "callbackEvalG" to evaluate the
    # objective gradient.  If not provided, Knitro will approximate
    # the gradient using finite-differencing.  However, we recommend
    # providing callbacks to evaluate the exact gradients whenever
    # possible as this can drastically improve the performance of Knitro.
    # We specify the objective gradient in "dense" form for simplicity.
    # However for models with many constraints, it is important to specify
    # the non-zero sparsity structure of the constraint gradients
    #(i.e. Jacobian matrix) for efficiency(this is true even when using
    # finite-difference gradients).
    KNITRO.KN_set_cb_grad(kc, cb, callbackEvalG!)

    # Add a callback function "callbackEvalH" to evaluate the Hessian
    #(i.e. second derivative matrix) of the objective.  If not specified,
    # Knitro will approximate the Hessian. However, providing a callback
    # for the exact Hessian(as well as the non-zero sparsity structure)
    # can greatly improve Knitro performance and is recommended if possible.
    # Since the Hessian is symmetric, only the upper triangle is provided.
    # Again for simplicity, we specify it in dense(row major) form.
    KNITRO.KN_set_cb_hess(kc, cb, KNITRO.KN_DENSE_ROWMAJOR, callbackEvalH!)

    # specify that the user is able to provide evaluations
    # of the hessian matrix without the objective component.
    # turned off by default but should be enabled if possible.
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_HESSIAN_NO_F, KNITRO.KN_HESSIAN_NO_F_ALLOW)

    # Set minimize or maximize(if not set, assumed minimize)
    KNITRO.KN_set_obj_goal(kc, KNITRO.KN_OBJGOAL_MINIMIZE)

    # Turn output off and use Interior/Direct algorithm
    KNITRO.KN_set_param(kc, "outlev", KNITRO.KN_OUTLEV_NONE)
    KNITRO.KN_set_param(kc, "algorithm", KNITRO.KN_ALG_BAR_DIRECT)

    # Solve the problem.
    #
    # Return status codes are defined in "knitro.h" and described
    # in the Knitro manual.

    # First solve for the 6 different values of user option "bar_murule".
    # This option handles how the barrier parameter is updated each
    # iteration in the barrier/interior-point solver.
    verbose && println("Changing a user option and re-solving...")
    for i in 1:6
        KNITRO.KN_set_param(kc, "bar_murule", i)
        # Reset original initial point
        KNITRO.KN_set_var_primal_init_values(kc, [-2.0, 1.0])
        nStatus = KNITRO.KN_solve(kc)
        if verbose
            if nStatus != 0
                println("  bar_murule=$i - Knitro failed to solve, status = $nStatus")
            else
                @printf("\n  bar_murule=%d - solved in %2d iters, %2d function evaluations, objective=%e",
                    i, KNITRO.KN_get_number_iters(kc), KNITRO.KN_get_number_FC_evals(kc), KNITRO.KN_get_obj_value(kc))
            end
        end
    end

    # Now solve for different values of the x0 upper bound.
    # Continually relax the upper bound until it is no longer
    # "active"(i.e. no longer restricting x0), at which point
    # there is no more significant change in the optimal solution.
    # Change to the active-set algorithm and do not reset the
    # initial point, so the re-solves are "warm-started".
    verbose && println("\nChanging a variable bound and re-solving...")
    KNITRO.KN_set_param(kc, "algorithm", KNITRO.KN_ALG_ACT_CG)
    i = 0

    for i = 1:20
        # Modify bound for next solve.
        tmpbound = 0.1*i
        KNITRO.KN_set_var_upbnds(kc, 0, tmpbound)

        nStatus = KNITRO.KN_solve(kc)
        nStatus, objSol, x, lambda_ = KNITRO.KN_get_solution(kc)
        if verbose
            if nStatus != 0
                @printf("\n  x0 upper bound=%e - Knitro failed to solve, status = %d", tmpbound, nStatus)
            else
                @printf("\n  x0 upper bound=%e - solved in %2d iters, x0=%e, objective=%e",
                        tmpbound, KNITRO.KN_get_number_iters(kc), x[1], objSol)
            end
        end

        if nStatus != 0 || x[1] < tmpbound - 1e-4
            break
        end
    end
    # Restore original value
    KNITRO.KN_set_var_upbnds(kc, 0, 0.5)

    # Now solve for different values of the c0 lower bound.
    # Continually relax the lower bound until it is no longer
    # "active"(i.e. no longer restricting c0), at which point
    # there is no more significant change in the optimal solution.
    verbose && println("\nChanging a constraint bound and re-solving...")
    i = 0
    for i = 1:20
        tmpbound = 1. - 0.1*i
        KNITRO.KN_set_con_lobnds(kc, 0, tmpbound)
        nStatus = KNITRO.KN_solve(kc)
        c0 = KNITRO.KN_get_con_values(kc, 0)
        if verbose
            if nStatus != 0
                @printf("\n  c0 lower bound=%e - Knitro failed to solve, status = %d", tmpbound, nStatus)
            else
                @printf("\n  c0 lower bound=%e - solved in %2d iters, c0=%e, objective=%e",
                    tmpbound, KNITRO.KN_get_number_iters(kc), c0, KNITRO.KN_get_obj_value(kc))
            end
        end
        if nStatus != 0 || c0 > tmpbound + 1e-4
            break
        end
    end

    # Delete the Knitro solver instance.
    KNITRO.KN_free(kc)
end

example_restart(; verbose=isdefined(Main, :KN_VERBOSE) ? KN_VERBOSE : true)


text:
# Script for SWMF Galileo flybys simulated magnetic field comparison.
#
# 1. Steady state
# 2. Time accurate
#
# Hongyang Zhou, hyzhou@umich.edu 02/03/2020

using Batsrus, CSV, Dates, PyPlot, Interpolations

struct BField{T<:AbstractFloat}
   t::Vector{DateTime}
   Bx::AbstractVector{T}
   By::AbstractVector{T}
   Bz::AbstractVector{T}
   Bmag::AbstractVector{T}
end

"""
	plotBSteady(flyby=8; file, dirobs, npict)

Plot quasi-steady state magnetic field comparison between simulation results from `file`
end observation from `dirobs`. Second norm is calculated to estimate the difference.
# Arguments
- `flyby::Int`: flyby index in [1,2,7,8,28,29].
"""
function plotBSteady(flyby=8; file="box_var_4_n00080000.out", dirobs=".", npict=1)

   # Read observation data
   fileobs = joinpath(dirobs, "Galileo_G"*string(flyby)*"_flyby_MAG.dat")
   t, df = getBObs(fileobs)
   BobsStrength = hypot.(df.Bx, df.By, df.Bz)

   # Read simulation data
   dirsim = dirname(file)
   if isempty(dirsim)
      data = readdata(file; npict)
   else
      data = readdata(basename(file); dir=dirsim, npict)
   end

   # Interpolate simulation data to observation data
   nx, nw = data.head.nx, data.head.nw

   x = data.x[:,1,1,1]
   y = data.x[1,:,1,2]
   z = data.x[1,1,:,3]

   w = data.w
   # Find the correct index of variables
   bx_ = findfirst(x->x=="bx", data.head.wnames)
   by_ = findfirst(x->x=="by", data.head.wnames)
   bz_ = findfirst(x->x=="bz", data.head.wnames)
   Bx = @view data.w[:,:,:,bx_]
   By = @view data.w[:,:,:,by_]
   Bz = @view data.w[:,:,:,bz_]

   # Gridded interpolation
   knots = (x, y, z)
   itpBx = interpolate(knots, Bx, Gridded(Linear()))
   BxSim = itpBx.(df.X,df.Y,df.Z)

   itpBy = interpolate(knots, By, Gridded(Linear()))
   BySim = itpBy.(df.X,df.Y,df.Z)

   itpBz = interpolate(knots, Bz, Gridded(Linear()))
   BzSim = itpBz.(df.X,df.Y,df.Z)

   BsimStrength = hypot.(BxSim, BySim, BzSim)

   # Visualization
   figure(figsize=(12,5))
   formatter = matplotlib.dates.DateFormatter("%H:%M")
   ax1 = subplot(411)
   plot(t, df.Bx, label=L"B_x")
   plot(t, BxSim)
   ax1.xaxis.set_major_formatter(formatter)
   legend()
   ax2 = subplot(412)
   plot(t, df.By, label=L"B_y")
   plot(t, BySim)
   ax2.xaxis.set_major_formatter(formatter)
   legend()
   ax3 = subplot(413)
   plot(t, df.Bz, label=L"B_z")
   plot(t, BzSim)
   ax3.xaxis.set_major_formatter(formatter)
   legend()
   ax4 = subplot(414)
   plot(t, BobsStrength, label=L"B")
   plot(t, BsimStrength)
   ax4.xaxis.set_major_formatter(formatter)
   legend()
   tight_layout()
   nothing
end

"""
    interp_simulation(file, dir, t, df, tStart, firstpict)

Return interpolated time-accurate B field stored in `file` from snapshot `firstpict` onto
timestamps `t` at satellite locations in `df`.
`tStart` is the starting DateTime for the synthetic satellite.

This is designed to handle any number of simulation snapshots. If the simulation time
interval is smaller than the observation's (which is usually the case), than the magnetic
field along the trajectory before time `tstart` is obtained from the first single snapshot
`firstpict`, and the field after the end time is obtained from the last single snapshot
ipict=npict. In this way we avoid the artificial jump due to periodic interpolation.
"""
function interp_simulation(file, dir, t, df, tStart, firstpict)

   data = readdata(file; dir)

   # Interpolate simulation data to observation data
   nx, nw = data.head.nx, data.head.nw
   npict = data.list.npictinfiles

   firstpict > npict && @error "firstpict out of range!"

   npict += 1 - firstpict

   kStart = findfirst(tStart .- t .< Second(0))

   tEnd  = tStart + Second(npict-1)
   kEnd = findfirst(tEnd .- t .< Second(0))

   nSim = kStart + (length(t) - kEnd + 1) + npict - 2
   BxSim = Vector{Float32}(undef, nSim)
   BySim = Vector{Float32}(undef, nSim)
   BzSim = Vector{Float32}(undef, nSim)
   tSim = Vector{DateTime}(undef, nSim)
   tSim[1:kStart] = t[1:kStart]

   # This is an approximation
   tSim[kStart:kStart+npict-1] .= timelinspace(tStart, tEnd, n=npict)
   tSim[kStart+npict-1:end] = t[kEnd:end]

   # Extract B field along the traj. before the starting time in one snapshot
   data = readdata(file; dir, npict=firstpict)

   # Interpolate simulation data to observation data
   x = data.x[:,1,1,1]
   y = data.x[1,:,1,2]
   z = data.x[1,1,:,3]

   bx_ = findfirst(x->x=="bx", data.head.wnames)
   by_ = findfirst(x->x=="by", data.head.wnames)
   bz_ = findfirst(x->x=="bz", data.head.wnames)
   Bx = @view data.w[:,:,:,bx_]
   By = @view data.w[:,:,:,by_]
   Bz = @view data.w[:,:,:,bz_]

   # Gridded interpolation
   knots = (x, y, z)
   itpBx = interpolate(knots, Bx, Gridded(Linear()))
   BxSim[1:kStart] = itpBx.(df.X[1:kStart],df.Y[1:kStart],df.Z[1:kStart])

   itpBy = interpolate(knots, By, Gridded(Linear()))
   BySim[1:kStart] = itpBy.(df.X[1:kStart],df.Y[1:kStart],df.Z[1:kStart])

   itpBz = interpolate(knots, Bz, Gridded(Linear()))
   BzSim[1:kStart] = itpBz.(df.X[1:kStart],df.Y[1:kStart],df.Z[1:kStart])

   # Extract the magnetic field along the trajectory from multiple snapshot
   for ipict = 1:npict-1
      println("ipict=$(round(ipict,digits=2))")
      data = readdata(file; dir, npict=firstpict+ipict)
      # Interpolate simulation data to observation data
      x = data.x[:,1,1,1]
      y = data.x[1,:,1,2]
      z = data.x[1,1,:,3]

      Bx = @view data.w[:,:,:,bx_]
      By = @view data.w[:,:,:,by_]
      Bz = @view data.w[:,:,:,bz_]

      k = argmin( abs.(tSim[kStart+ipict] .- t) )

      itpBx = interpolate(knots, Bx, Gridded(Linear()))
      BxSim[kStart+ipict] = itpBx(df.X[k], df.Y[k], df.Z[k])

      itpBy = interpolate(knots, By, Gridded(Linear()))
      BySim[kStart+ipict] = itpBy(df.X[k], df.Y[k], df.Z[k])

      itpBz = interpolate(knots, Bz, Gridded(Linear()))
      BzSim[kStart+ipict] = itpBz(df.X[k], df.Y[k], df.Z[k])
   end

   BxSim[kStart+npict:end] = itpBx.(df.X[kEnd+1:end], df.Y[kEnd+1:end], df.Z[kEnd+1:end])
   BySim[kStart+npict:end] = itpBy.(df.X[kEnd+1:end], df.Y[kEnd+1:end], df.Z[kEnd+1:end])
   BzSim[kStart+npict:end] = itpBz.(df.X[kEnd+1:end], df.Y[kEnd+1:end], df.Z[kEnd+1:end])

   BsimStrength = hypot.(BxSim, BySim, BzSim)

   Bsim = BField(tSim, BxSim, BySim, BzSim, BsimStrength)
end

function timelinspace(d1::DateTime, d2::DateTime; n=2)
   if n == 1
      return d1
   end
   Δ = d2 - d1
   T = typeof(Δ)
   δ = T(round(Int, Dates.value(Δ)/(n - 1)))
   d2 = d1 + δ*(n - 1)
   return d1:δ:d2
end

"Read observation data from `file` and return the timestamps and data."
function getBObs(file)

   df = CSV.File(file; header=2, delim=" ", ignorerepeated=true)
   t = DateTime.(df.yr, df.month, df.day, df.hr, df.min, floor.(Int, df.sec),
      floor.(Int, 1e3 .* (df.sec - floor.(df.sec))) )

   t, df
end

##
#plotBTimeAccurate("box_Hall_B_1200.outs";
#   ObsDir="/Users/hyzhou/Ganymede/GalileoData/galileomagdata",
#   SimDir="/Users/hyzhou/Ganymede/Hall_AMR3/GM",)

#=
# Plot time-accurate magnetic field comparisons.
flyby = 8
filename = ["box_Hall_B_1200.outs", "box_PIC_B_1200.outs"]
dir = "/Users/hyzhou/Documents/Computer/Julia/BATSRUS/VisAnaJulia"
firstpict, lastpict = 1, 1

t, df = getBObs()
BobsStrength = hypot.(df.Bx, df.By, df.Bz)

# Select the starting time for synthetic satellite
if flyby == 8
   #tStart = DateTime(1997,5,7,15,48,7)
   tStart = DateTime(1997,5,7,15,45,40)# PIC
   tStartHall = DateTime(1997,5,7,15,45,48)# Hall
elseif flyby == 28
   tStart = DateTime(2000,5,20,10,0,12)
end

Hall = interp_simulation(filename[1], dir, t,df, tStartHall, firstpict)

PIC = interp_simulation(filename[2], dir, t,df, tStart, firstpict)

fig, ax = plt.subplots(4,1,figsize=(10.0,5.0))
plt.rc("font", family="serif", size=14)

# Plot observation data
ax[1].plot(t, df.Bx, "k", label="Obs")
ax[2].plot(t, df.By, "k", label="Obs")
ax[3].plot(t, df.Bz, "k", label="Obs")
ax[4].plot(t, BobsStrength, "k", label="Obs")

# Plot simulation data
ax[1].plot(Hall.t, Hall.Bx, label="Hall MHD")
ax[2].plot(Hall.t, Hall.By, label="Hall MHD")
ax[3].plot(Hall.t, Hall.Bz, label="Hall MHD")
ax[4].plot(Hall.t, Hall.Bmag, label="Hall MHD")

ax[1].plot(PIC.t, PIC.Bx, label="MHD-EPIC")
ax[2].plot(PIC.t, PIC.By, label="MHD-EPIC")
ax[3].plot(PIC.t, PIC.Bz, label="MHD-EPIC")
ax[4].plot(PIC.t, PIC.Bmag, label="MHD-EPIC")

ax[1].set_ylabel("Bx [nT]")
ax[2].set_ylabel("By [nT]")
ax[3].set_ylabel("Bz [nT]")
ax[4].set_ylabel("B [nT]")
suptitle("Galileo G8 Flyby Magnetic Field")
for (i,a) in enumerate(ax)
   a.set_xlim(t[1500],t[end-2700])
   if i != 4
      a.axes.xaxis.set_ticklabels([])
   else
      formatter = matplotlib.dates.DateFormatter("%H:%M")
      a.xaxis.set_major_formatter(formatter)
   end
   a.tick_params(which="both", direction="in", top=true, right=true)
   a.minorticks_on()
   i == 1 &&
   a.legend(loc="lower left", bbox_to_anchor=(0.0, 0.86), ncol=3, frameon=false)
end

plt.subplots_adjust(hspace=0)
fig.align_ylabels(ax)
=#


text:
using Test
using Printf
using Distributed
using Random
using LinearAlgebra
import LinearMaps: LinearMap, _ismutating

print("Checking cores... ")
Sys.CPU_THREADS > 1 && addprocs(2)
println("done")
print("Loading KirchMig... ")
using KirchMig
println("done")

# Setup 2D
print("2D setup... ")
nR = nS = 64
dR = dS = 14
oR = oS = 0

ot = 1e-2
dt = 1e-4
nt = 901
t = ot:dt:ot+dt*(nt-1)

dx = dz = 100
x = 0:dx:oR+(nR-1)*dR
z = 0:dz:400
nx = length(x)
nz = length(z)
vel = 2700
println("done")

print("2D traveltimes... ")
src_x = oR:dR:oR+(nR-1)*dR
src_z = zeros(nR)
trav = eikonal_const_vel([src_z src_x], z, x, vel)
println("done")

try
   L = KirchMap(t, trav, trav[1:end-1, :, :])
catch err
    @test isa(err, DimensionMismatch)
end

# Dot test 2D
println("2D Kirchhoff...")
for pts in ["parallel", "threaded", "serial"]
    pts == "parallel" ? println("  Using ", nworkers(), " workers") : nothing
    pts == "threaded" ? println("  Using ", Threads.nthreads(),
                                " thread" * (Threads.nthreads() > 1 ? "s" : "") ) : nothing
    pts == "serial" ? println("  Using serial version") : nothing
    L = KirchMap(t, trav; parallel_threaded_serial=pts)
    Random.seed!(1234)
    u = rand(size(L, 2))
    v = rand(size(L, 1))
    print("    Forward:")
    @time v_hat = L*u;
    print("    Adjoint:")
    @time u_hat = L'*v;
    utu =  dot(u_hat, u)
    vtv =  dot(v_hat, v)
    println(@sprintf("    dot test: ⟨L'v, u⟩ = %.2f", utu))
    println(@sprintf("              ⟨L u, v⟩ = %.2f", vtv))
    @test abs(utu - vtv)/((utu+vtv)/2) <= 100*eps()

    issymmetric(L) == false
    isposdef(L) == false
    _ismutating(L.f)
    println("done")
    println("")
end

# Setup 3D
print("3D setup... ")

dy = 10
ny = 3
y = (0:ny-1)*dy
println("done")

print("3D traveltimes... ")
src_y = zeros(nR)
trav = eikonal_const_vel([src_x src_y src_z], x, y, z, vel)
println("done")

# Dot test 3D
println("3D Kirchhoff...")
for pts in ["parallel", "threaded", "serial"]
    pts == "parallel" ? println("  Using ", nworkers(), " workers") : nothing
    pts == "threaded" ? println("  Using ", Threads.nthreads(), " threads") : nothing
    pts == "serial" ? println("  Using serial version") : nothing
    L = KirchMap(t, trav; parallel_threaded_serial=pts)
    Random.seed!(1234)
    u = rand(size(L, 2))
    v = rand(size(L, 1))
    print("    Forward:")
    @time v_hat = L*u;
    print("    Adjoint:")
    @time u_hat = L'*v;
    utu =  dot(u_hat, u)
    vtv =  dot(v_hat, v)
    println(@sprintf("    dot test: ⟨L'v, u⟩ = %.2f", utu))
    println(@sprintf("              ⟨L u, v⟩ = %.2f", vtv))

    @test abs(utu - vtv)/((utu+vtv)/2) <= 100*eps()
    println("")
end

# cg
println("Conjugate gradients... ")
M, N = 10, 5
function fwd(x)
    y = zeros(eltype(x), M)
    y[1:N] = 2x .+ x[end:-1:1]
    return y
end
function adj(x)
    return 2x[1:N] .+ x[1:N][end:-1:1]
end
Id = LinearMap(fwd, adj, M, N)
u = rand(size(Id, 2))
v = rand(size(Id, 1))
v_hat = Id*u;
u_hat = Id'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨corr v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨conv u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()

b = rand(size(Id, 1))
x, hist_x, hist_r = cg(Id'Id, Id'b, log=true)
@test x ≈ inv(Matrix(Id'Id))*(Id'b)
println("done")
println("")

# Regularization
println("2D ConvMap... ")
ricker(t0, f) = @. (1 - 2pi^2 * f^2 * t0^2) * exp(-pi^2 * f^2 * t0^2)
rick_dt = ricker(t .- t[div(nt,3)], 15);
rick_dt[2:end-1] = @. (rick_dt[1:end-2] - rick_dt[3:end])/2(t[2] - t[1])

W = ConvMap(rick_dt, nS, nt);
Random.seed!(1234)
u = rand(size(W, 2))
v = rand(size(W, 1))
v_hat = W*u;
u_hat = W'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨corr v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨conv u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("3D ConvMap... ")
W = ConvMap(rick_dt, nR, nS, nt);
Random.seed!(1234)
u = rand(size(W, 2))
v = rand(size(W, 1))
v_hat = W*u;
u_hat = W'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨corr v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨conv u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("2D Laplacian... ")
Δ = KirchMig.LaplacianMap(nz, nx)

Random.seed!(1234)
u = rand(size(Δ, 2))
v = rand(size(Δ, 1))
v_hat = Δ*u;
u_hat = Δ'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨Δ'v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨Δ u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("3D Laplacian... ")
Δ = KirchMig.LaplacianMap(nz, nx, ny)

Random.seed!(1234)
u = rand(size(Δ, 2))
v = rand(size(Δ, 1))
v_hat = Δ*u;
u_hat = Δ'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨Δ'v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨Δ u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("3D DiffZ... ")
δz = KirchMig.DiffZMap(nz, nx, ny)

Random.seed!(1234)
u = rand(size(δz, 2))
v = rand(size(δz, 1))
v_hat = δz*u;
u_hat = δz'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨-δz v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨ δz u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("2D DiffX... ")
δx = KirchMig.DiffXMap(nz, nx)

Random.seed!(1234)
u = rand(size(δx, 2))
v = rand(size(δx, 1))
v_hat = δx*u;
u_hat = δx'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨-δx v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨ δx u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("2D GradDiv... ")
GD = KirchMig.GradDivMap(nz, nx)

Random.seed!(1234)
u = rand(size(GD, 2))
v = rand(size(GD, 1))
v_hat = GD*u;
u_hat = GD'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨-div v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨grad u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")

println("3D GradDiv... ")
GD = KirchMig.GradDivMap(nz, nx, ny)

Random.seed!(1234)
u = rand(size(GD, 2))
v = rand(size(GD, 1))
v_hat = GD*u;
u_hat = GD'v;
utu =  dot(u_hat, u)
vtv =  dot(v_hat, v)
println(@sprintf("  Dot test: ⟨-div v, u⟩ = %.5f", utu))
println(@sprintf("            ⟨grad u, v⟩ = %.5f", vtv))
@test abs(utu - vtv)/((utu+vtv)/2) <= 100eps()
println("")


text:
using Revise
	using DiffEqOperators, ForwardDiff
	using BifurcationKit, LinearAlgebra, Plots, SparseArrays, Parameters, Setfield
	const BK = BifurcationKit

norminf(x) = norm(x, Inf)
normbratu(x) = norm(x .* w) / sqrt(length(x))
##########################################################################################
# plotting function
plotsol!(x, nx = Nx, ny = Ny; kwargs...) = heatmap!(LinRange(0,1,nx), LinRange(0,1,ny), reshape(x, nx, ny)'; color = :viridis, xlabel = "x", ylabel = "y", kwargs...)
plotsol(x, nx = Nx, ny = Ny; kwargs...) = (plot();plotsol!(x, nx, ny; kwargs...))

function Laplacian2D(Nx, Ny, lx, ly, bc = :Neumann)
	hx = 2lx/Nx
	hy = 2ly/Ny
	D2x = CenteredDifference(2, 2, hx, Nx)
	D2y = CenteredDifference(2, 2, hy, Ny)

	if bc == :Neumann
		Qx = Neumann0BC(hx)
		Qy = Neumann0BC(hy)
	elseif bc == :Dirichlet
		Qx = Dirichlet0BC(typeof(hx))
		Qy = Dirichlet0BC(typeof(hy))
	end

	D2xsp = sparse(D2x * Qx)[1]
	D2ysp = sparse(D2y * Qy)[1]
	A = kron(sparse(I, Ny, Ny), D2xsp) + kron(D2ysp, sparse(I, Nx, Nx))
	return A, D2x
end

ϕ(u, λ)  = -10(u-λ*exp(u))
dϕ(u, λ) = -10(1-λ*exp(u))

function NL!(dest, u, p)
	@unpack λ = p
	dest .= ϕ.(u, λ)
	return dest
end

NL(u, p) = NL!(similar(u), u, p)

function Fmit!(f, u, p)
	mul!(f, p.Δ, u)
	f .= f .+ NL(u, p)
	return f
end

Fmit(u, p) = Fmit!(similar(u), u, p)

function dFmit(x, p, dx)
	f = similar(dx)
	mul!(f, p.Δ, dx)
	nl = d1NL(x, p, dx)
	f .= f .+ nl
end

function JFmit(x,p)
	J = p.Δ
	dg = dϕ.(x, p.λ)
	return J + spdiagm(0 => dg)
end

# computation of the derivatives
d1NL(x, p, dx) = ForwardDiff.derivative(t -> NL(x .+ t .* dx, p), 0.)

# compute 3-Jet
jet = BK.getJet(Fmit, JFmit)
####################################################################################################
Nx = 30
	Ny = 30
	lx = 0.5
	ly = 0.5

	Δ, = Laplacian2D(Nx, Ny, lx, ly)
	par_mit = (λ = .01, Δ = Δ)
	sol0 = 0*ones(Nx, Ny) |> vec
	const w = (lx .+ LinRange(-lx,lx,Nx)) * (LinRange(-ly,ly,Ny))' |> vec
	w .-= minimum(w)
####################################################################################################
eigls = EigArpack(20.5, :LM)
# eigls = EigKrylovKit(dim = 70)
# eigls = EigArpack()
	opt_newton = NewtonPar(tol = 1e-8, verbose = true, eigsolver = eigls, maxIter = 20)
	sol, = newton(Fmit, JFmit, sol0, par_mit, opt_newton, normN = norminf)

plotsol(sol)
####################################################################################################
function finSol(z, tau, step, br; k...)
	if length(br.specialpoint)>0
		if br.specialpoint[end].step == step
			BK._show(stdout, br.specialpoint[end], step)
		end
	end
	return true
end

function cb(x,f,J,res,it,itl,optN; kwargs...)
	_x = get(kwargs, :z0, nothing)
	fromNewton = get(kwargs, :fromNewton, false)
	if ~fromNewton
		return (norm(_x.u - x) < 20.5 && abs(_x.p - kwargs[:p]) < 0.05)
	end
	true
end

# optional parameters for continuation
kwargsC = (verbosity = 3,
	plot = true,
	recordFromSolution = (x, p) -> (x = normbratu(x), n2 = norm(x), n∞ = norminf(x)),
	plotSolution = (x, p; k...) -> plotsol!(x ; k...),
	callbackN = cb,
	finaliseSolution = finSol,
	normC = norminf
	)

opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.04, ds = 0.005, pMax = 3.5, pMin = 0.01, detectBifurcation = 3, nev = 50, plotEveryStep = 10, newtonOptions = (@set opt_newton.verbose = false), maxSteps = 251, precisionStability = 1e-6, nInversion = 6, maxBisectionSteps = 25)

	br, = @time continuation(
		Fmit, JFmit,
		sol0, par_mit, (@lens _.λ), opts_br; kwargsC...)

plot(br)
####################################################################################################
# automatic branch switching

br1, = continuation(jet..., br, 3; kwargsC...)

plot(br,br1,plotfold=false)

br2, = continuation(jet..., br1, 1; kwargsC...)

plot(br, br1, br2, plotfold=false)
####################################################################################################
# bifurcation diagram
function optionsCont(x,p,l; opt = opts_br)
	if l <= 1
		return opt
	elseif l==2
		return setproperties(opt ;detectBifurcation = 3,ds = 0.001, a = 0.75)
	else
		return setproperties(opt ;detectBifurcation = 3,ds = 0.00051, dsmax = 0.01)
	end
end

diagram = @time bifurcationdiagram(jet..., sol0, par_mit, (@lens _.λ), 3, optionsCont; kwargsC...,
	usedeflation = true,
	halfbranch = true,
	)

bifurcationdiagram!(jet..., getBranch(diagram, (14,)), (current = 3, maxlevel = 6), optionsCont;
	kwargsC..., usedeflation = true, halfbranch = true,)

code = ()
	plot(diagram; code = code,  plotfold = false, putspecialptlegend=false, markersize=2, vars = (:param, :x))
	# plot!(br)
	# xlims!(0.01, 0.4)
	title!("#branches = $(size(getBranch(diagram, code)))")
	# xlims!(0.01, 0.065, ylims=(2.5,6.5))

plot(getBranchesFromBP(diagram, 4); plotfold = false, legend = false, vars = (:param, :x))

BK.getBranch(diagram, (10,1,)) |> plot

####################################################################################################
# analyse 2d bifurcation point
bp2d = @time BK.computeNormalForm(jet..., br, 4,  verbose=true, nev=30)

BK.nf(bp2d)[2] |> println

using ProgressMeter
Nd = 100
	L = 3.9
	X = LinRange(-L,L, Nd)
	Y = LinRange(-L,L, Nd)
	P = LinRange(-0.001,0.001, Nd+1)

V1a = @showprogress [bp2d(Val(:reducedForm),[x1,y1], p1)[1] for p1 in P, x1 in X, y1 in Y]

Ind1 = findall( abs.(V1a) .<= 9e-4 * maximum(abs.(V1a)))

	V2a = @showprogress [bp2d(Val(:reducedForm),[X[ii[2]],Y[ii[3]]], P[ii[1]])[2] for ii in Ind1]

	Ind2 = findall( abs.(V2a) .<= 3e-3 * maximum(abs.(V2a)))
	@show length(Ind2)

resp = Float64[]
	resx = Vector{Float64}[]
	resnrm = Float64[]
	@showprogress for k in Ind2
		ii = Ind1[k]
		push!(resp, P[ii[1]])
		# push!(resx, max(X[ii[2]],Y[ii[3]]))
		push!(resnrm, sqrt(X[ii[2]]^2+Y[ii[3]]^2))
		push!(resx, [X[ii[2]], Y[ii[3]]])
	end

using LaTeXStrings

plot(
	scatter(1e4resp, map(x->x[1], resx), map(x->x[2], resx); label = "", markerstrokewidth=0, xlabel = L"10^4 \cdot \lambda", ylabel = L"x_1", zlabel = L"x_2", zcolor = resnrm, color = :viridis,colorbar=false),
	scatter(1e4resp, resnrm; label = "", markersize =2, markerstrokewidth=0, xlabel = L"10^4 \cdot \lambda", ylabel = L"\|x\|"))

####################################################################################################
bp2d = @time computeNormalForm(jet..., br, 2, nev = 30; issymmetric = true)

res, = BK.continuation(jet..., br, 2,
	setproperties(opts_br; detectBifurcation = 3, ds = 0.001, pMin = 0.01, maxSteps = 32 ) ;
	nev = 30, verbosity = 3,
	kwargsC...,
	)

plot(res..., br ;plotfold= false)

δp = 0.005
	deflationOp = DeflationOperator(2, 1.0, [zeros(2)])
		success = [0]
while sum(success) < 10
	outdef1, _, flag, _ = newton((x, p) -> bp2d(Val(:reducedForm), x, p[1]), rand(2), [δp], NewtonPar(maxIter = 50), deflationOp)
	@show flag
	flag && push!(deflationOp, outdef1)
	(flag==false) && push!(success, 1)
end
	println("--> found $(length(deflationOp)) solutions")

plotsol(bp(deflationOp[3], δp))
solbif, flag, _ = newton(Fmit, JFmit, bp2d.x0, bp2d(deflationOp[3], δp), (@set par_mit.λ = bp2d.p + δp), opts_br.newtonOptions)[1]

plotsol(solbif-0*bp2d(deflationOp[2], δp))

brnf1, = continuation(Fmit, JFmit, solbif, (@set par_mit.λ = bp2d.p + δp), (@lens _.λ), setproperties(opts_br; ds = 0.005);
	recordFromSolution = (x, p) -> norm(x),
	plotSolution = (x, p; kwargs...) -> plotsol!(x ; kwargs...),
	plot = true, verbosity = 3, normC = norminf)

branches2 = (br,br1,br2,brnf1)
push!(branches2, brnf1)
# plot([br,br1,br2])
# plot!(brnf1)

brnf2, = continuation(Fmit, JFmit, solbif, (@set par_mit.λ = bp2d.p + δp), (@lens _.λ), setproperties(opts_br; ds = -0.005);
	recordFromSolution = (x, p) -> norm(x),
	plotSolution = (x, p; kwargs...) -> plotsol!(x ; kwargs...),
	plot = true, verbosity = 3, normC = norminf)

# plot([br,br1,br2]);plot!(brnf1);plot!(brnf2)
plot(branches2...)
plot!(brnf2)
####################################################################################################
# find isolated branch, see Farrell et al.
deflationOp = DeflationOperator(2, 1.0, [out])
optdef = setproperties(opt_newton; tol = 1e-8, maxIter = 150)

# eigen-elements close to the second bifurcation point on the branch
# of homogenous solutions
vp, ve, _, _= eigls(JFmit(out, @set par_mit.λ = br.specialpoint[2].param), 5)

for ii=1:size(ve, 1)
		outdef1, _, flag, _ = @time newton(
			Fmit, JFmit,
			# initial guess for newton
			real.(br.specialpoint[2].x .+ 0.01 .* ve[ii] .* (1 .+ 0.01 .* rand(Nx*Ny))),
			(@set par_mit.λ = br.specialpoint[2].param + 0.005),
			optdef, deflationOp)
			flag && push!(deflationOp, outdef1)
	end
	length(deflationOp)

l = @layout grid(3,2)
	plot(layout = l)
	for ii=1:length(deflationOp)
		plotsol!(deflationOp[ii], title="$ii", subplot = ii, label = "", xlabel="$ii", colorbar=true)
	end
	title!("")

brdef1, _ = @time BK.continuation(
	Fmit, JFmit,
	deflationOp[3], (@set par_mit.λ = br.specialpoint[2].param + 0.005), (@lens _.λ),
	# bp2d([0.6,0.6], -0.01), br.specialpoint[2].param - 0.005,
	setproperties(opts_br;ds = 0.001, detectBifurcation = 0, dsmax = 0.01, maxSteps = 500);
	verbosity = 3, plot = true,
	recordFromSolution = (x, p) ->  normbratu(x),
	plotSolution = (x, p; kwargs...) -> plotsol!(x ; kwargs...), normC = norminf)

plot(br,br1,br2, brdef1,plotfold=false)

brdef2, _ = @time BK.continuation(
	Fmit, JFmit,
	deflationOp[5], (@set par_mit.λ = br.specialpoint[2].param + 0.005), (@lens _.λ),
	setproperties(opts_br;ds = -0.001, detectBifurcation = 0, dsmax = 0.02);
	verbosity = 3, plot = true,
	recordFromSolution = (x, p) ->  normbratu(x),
	plotSolution = (x, p; kwargs...) -> plotsol!(x ; kwargs...), normC = norminf)

plot(br,br1,br2, brdef1, brdef2,plotfold=false, putspecialptlegend = false)

plot(brdef1, brdef2,plotfold = false, putspecialptlegend = false)
####################################################################################################
# deflated continuation
brdef2, _ = @time BK.continuation(
	Fmit, JFmit, (@set par_mit.λ = 0.367), (@lens _.λ),
	ContinuationPar(opts_br; ds = -0.0001, maxSteps = 800000, plotEveryStep = 10, detectBifurcation = 0),
	DeflationOperator(2, 1., ([sol0]));
	plot=true, verbosity = 2,
	perturbSolution = (x,p,id) -> (x .+ 0.1 .* rand(length(x))),
	recordFromSolution = (x, p) ->  normbratu(x),
	plotSolution = (x, p; kwargs...) -> plotsol!(x ; kwargs...),
	normN = norminf)

plot(brdef2..., color=:red)


text:
#=

Code for transforming between Laplacian and SDDM solvers.

By Dan Spielman

=#

"""
    solveA = wrapInterface(solver::Function, A::AbstractMatrix; tol, maxits, maxtime, verbose, pcgIts=Int[],params...)
    solverConstructor = wrapInterface(A::AbstractMatrix; tol, maxits, maxtime, verbose, pcgIts=Int[],params...)

Returns a function that discards `tol`, `maxits`, `maxtime` and `verbose`,
sets `pcgIts` to 0 (because it might not be using pcg),
and passes whatever `params` are left to the solver.

# Examples

```julia
julia> a = randn(5,5);
julia> a = a * a';
julia> solvea = wrapInterface(X->cholesky(X,Val(true)), a, maxits=100, verbose=true);
julia> b = randn(5,1);
julia> norm(a*solvea(b, verbose=false)-b)
1.575705319704736e-14

julia> f = wrapInterface(X->cholesky(X,Val(true)))
julia> solvea = f(a, maxits=1000, maxtime = 1)
julia> norm(a*solvea(b, verbose=false, maxtime = 10)-b)
1.575705319704736e-14
```
"""
function wrapInterface(solver::Function, a::AbstractMatrix; tol=0, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[],params...)
    t1 = time()
    sol = solver(a)
    if verbose
        println("Solver build time: ", round((time() - t1),digits=3), " seconds.")
    end

    f = function(b; verbose=false, jnkargs...)
        if length(pcgIts) > 0
            pcgIts[1] = 1
        end

        t1 = time()
        if isa(sol,Factorization)
            x = sol \ b
        else
            x = sol(b; params...)
        end

        if verbose
            println("Solve time: ", round((time() - t1),digits=3), " seconds.")
        end

        return x
    end
    return f
end

function wrapInterface(solver::Function)
    f = function(a::AbstractArray; tol=0, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)
        return wrapInterface(solver, a; pcgIts=Int[], verbose=verbose, params...)
    end
    return f
end

function nullSolver(a;params...)
    return 0.0
end

"""
    solveSDDM = chol_sddm(sddm::AbstractMatrix; tol, maxits, maxtime, verbose, pcgIts=Int[])

This functions wraps cholfact so that it satsfies our interface.
It ignores all the keyword arguments.
"""
chol_sddm = wrapInterface(X->cholesky(X))

"""
    la = forceLap(a)

Create a Laplacian matrix from an adjacency matrix.
If the input looks like a Laplacian, throw a warning and convert it.
"""
function forceLap(a::AbstractArray)

    if minimum(a) < 0
        @warn "The input should be an adjacency matrix, whereas this one has negative entries."
        af = abs.(a)
        af = af - sparse(Diagonal(diag(af)))
    elseif sum(abs.(diag(a))) > 0
        @warn "The input should be an adjacency matrix, whereas this one has diagonal entries."
        af = a - sparse(Diagonal(diag(a)))
    else
        af = a
    end

    return sparse(Diagonal(vec(sum(af,dims=1)))) - af
end

"""
    f = lapWrapConnected(sddmSolver, a::AbstractMatrix; kwargs...)

Applies a `sddmSolver` to the Laplacian of the adjacency matrix `a` of a connected graph.
Passes on kwargs to the solver.
`sddmSolver` should be a solver that obeys the interface.
"""
function lapWrapConnected(solver, a::AbstractMatrix; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)
    la = forceLap(a)
    N = size(la)[1]

    ind = findmax(diag(la))[2]
    leave = [1:(ind-1);(ind+1):N]

    lasub = la[leave,leave]
    subSolver = solver(lasub; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts);

    tol_=tol
    maxits_=maxits
    maxtime_=maxtime
    verbose_=verbose
    pcgIts_=pcgIts

    f = function(b; tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_)

        bs = b[leave] .- mean(b)

        xs = subSolver(bs, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts)

        x = zeros(size(b))
        x[leave] = xs
        x = x .- mean(x)
        return x
    end

    return f
end

function lapWrapConnected(solver::Function)
    f(a::AbstractArray; tol=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...) = lapWrapConnected(solver, a; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts, params... )
    return f
end

"""Apply the ith solver on the ith component"""
function blockSolver(comps, solvers; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[])

    tol_=tol
    maxits_=maxits
    maxtime_=maxtime
    verbose_=verbose
    pcgIts_=pcgIts

    f = function(b; tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_)

        if length(pcgIts) > 0
            pcgIts[1] = 0
            pcgTmp = Int[0]
        else
            pcgTmp = Int[]
        end

        x = zeros(size(b))
        for i in 1:length(comps)
            ind = comps[i]
            bi = b[ind]
            x[ind] = (solvers[i])(bi;  tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgTmp)
            if length(pcgIts) > 0
                pcgIts[1] = max(pcgIts[1],pcgTmp[1])
            end

        end
        return x
    end

end

"""
    testZeroDiag(a)

Returns true if `a` has zero diagonal, false otherwise
"""
function testZeroDiag(a)
    n = size(a,1)
    for i in 1:n
        if a[i,i] != 0.0
            return false
        end
    end
    return true
end

"""
    f = lapWrapComponents(solver, a::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

Applies a Laplacian `solver` that satisfies our interface to each connected component of the graph with adjacency matrix `a`.
Passes kwargs on the solver.
"""
function lapWrapComponents(solver, a::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

    t1 = time()

    if !testZeroDiag(a)
        @warn "The matrix should not have any nonzero diagonal entries."
        a = a - sparse(Diagonal(diag(a)))
    end

    co = components(a)

    if maximum(co) == 1

        s = solver(a; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts, params... )
        if verbose
            println("Solver build time: ", round((time() - t1),digits=3), " seconds.")
        end

        # f(b; tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_) =
        return s

    else

        comps = vecToComps(co)

        solvers = []
        for i in 1:length(comps)
            ind = comps[i]

            asub = a[ind,ind]

            if (length(ind) == 1)
                subSolver = nullSolver

            elseif (length(ind) < 50)
                subSolver = lapWrapConnected(chol_sddm,asub)

            else

                subSolver = solver(asub; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts, params... );

            end
            push!(solvers, subSolver)
        end

        if verbose
            println("Solver build time: ", round((time() - t1),digits=3), " seconds.")
        end

        return blockSolver(comps,solvers; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts)

    end
end

function lapWrapComponents(solver::Function)
    f(a::AbstractArray; tol=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...) = lapWrapComponents(solver, a; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts, params... )
    return f
end

"""
    f = lapWrapSDDM(sddmSolver, A::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)
    f = lapWrapSDDM(sddmSolver)

Uses a `sddmSolver` to solve systems of linear equations in Laplacian matrices.
"""
function lapWrapSDDM(sddmSolver, a::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

    @assert(minimum(a) >= 0, "A must be nonnegative")
    @assert(sum(abs.(diag(a))) == 0, "A must have zero diagonal")

    f = Laplacians.lapWrapComponents(Laplacians.lapWrapConnected(sddmSolver))(a, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts)

    return f

end

function lapWrapSDDM(sddmSolver)

    return Laplacians.lapWrapComponents(Laplacians.lapWrapConnected(sddmSolver))

end

"""
    solver = chol_lap(A::AbstractArray)

Uses Cholesky Factorization to solve systems in Laplacians.
"""
chol_lap = lapWrapSDDM(chol_sddm)

"""
    f = sddmWrapLap(lapSolver, sddm::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

Uses a `lapSolver` to solve systems of linear equations in sddm matrices.
"""
function sddmWrapLap(lapSolver, sddm::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

    # Make a new adj matrix, a1, with an extra entry at the end.
    a, d = adj(sddm)
    a1 = extendMatrix(a,d)
    F = lapSolver(a1; tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts, params...)

    # make a function that solves the extended system, modulo the last entry
    tol_=tol
    maxits_=maxits
    maxtime_=maxtime
    verbose_=verbose
    pcgIts_=pcgIts

    f = function(b; tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_)
        xaug = F([b; -sum(b)], tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts)
        xaug = xaug .- xaug[end]
        return xaug[1:a.n]
    end

    return f

end

function sddmWrapLap(lapSolver)
    f = function(sddm::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)
        return sddmWrapLap(lapSolver, sddm;  tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts, params... )
    end
    return f
end

"""
    f = wrapCaptureRhs(sola::Function, rhss; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

Captures all the right-hand-sides that are passed to the solver `sola`.  It pushes them into an array called rhhs.
For example

```julia
julia> rhss = []
julia> a = wted_chimera(100)
julia> sola = approxchol_lap(a)
julia> wrappedSolver = wrapCaptureRhs(sola,rhss)
julia> b = randn(100)
julia> x = wrappedSolver(b,verbose=true)

PCG BLAS stopped after: 0.0 seconds and 11 iterations with relative error 3.160275810360986e-7.

julia> length(rhss[1])

100
```

"""
function wrapCaptureRhs(sola::Function, rhss; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

    tol_=tol
    maxits_=maxits
    maxtime_=maxtime
    verbose_=verbose
    pcgIts_=pcgIts

    f = function(b::AbstractArray; tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_)
        push!(rhss, b)
        x = sola(b, tol=tol, maxits=maxits, maxtime=maxtime, verbose=verbose, pcgIts=pcgIts)
    end
end

"""
    f = wrapCapture(solver::Function, mats, rhss; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)

This wraps a solver so that we can capture all the matrices that it solves and all the right-hand-sides.
Those are pushed into the arrays `mats` and `rhss`.
For example

```julia
julia> mats = []
julia> rhss = []
julia> solver = wrapCapture(approxchol_lap, mats, rhss)
julia> a = chimera(10)
julia> f = solver(a);
julia> size(mats[1])
(10,10)
julia> b = randn(10)
julia> x = f(b);
julia> rhss
1-element Array{Any,1}:
 [0.404962,-0.827718,0.704616,-0.403223,0.204891,-0.505589,0.907015,1.90266,-0.438115,0.0464351]
```
"""
function wrapCapture(solver::Function, mats, rhss)
    f = function(a::AbstractArray; tol::Real=1e-6, maxits=Inf, maxtime=Inf, verbose=false, pcgIts=Int[], params...)
        tol_=tol
        maxits_=maxits
        maxtime_=maxtime
        verbose_=verbose
        pcgIts_=pcgIts

        push!(mats,a)
        sol1 = solver(a, tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_, params...)
        return wrapCaptureRhs(sol1, rhss, tol=tol_, maxits=maxits_, maxtime=maxtime_, verbose=verbose_, pcgIts=pcgIts_)
    end
    return f
end


text:
_pymodelref(name) =
    "[`$name`](https://docs.python.org/3/reference/datamodel.html#object.$name)"

"""
    PyBase.__getattr__(self, ::Val{name::Symbol})
    PyBase.__getattr__(self, name::Symbol)

Python's attribute getter interface
[`__getattr__`](https://docs.python.org/3/reference/datamodel.html#object.__getattr__).

Roughly speaking, the default implementation is:

```julia
__getattr__(self, name::Symbol) = __getattr__(self, Val(name))
__getattr__(self, ::Val{name}) = getproperty(self, name)
```

To add a specific Python property to `self`, overload
`__getattr__(::MyType, ::Val{name})`.  Use `__getattr__(::MyType, ::Symbol)`
to fully control the attribute access in Python.

At the lowest level (which is invoked directly by Python), `PyBase`
defines `__getattr__(self, name::String) = __getattr__(self, Symbol(name))`
this can be used for advanced control when using predefined wrappers.
However, `__getattr__(::Shim{MyType}, ::String)` has to be overloaded
instead of `__getattr__(::MyType, ::String)`.
"""
PyBase.__getattr__

"""
    PyBase.__setattr__(self, ::Val{name::Symbol}, value)
    PyBase.__setattr__(self, name::Symbol, value)

Python's attribute setter interface
[`__setattr__`](https://docs.python.org/3/reference/datamodel.html#object.__setattr__).
Default implementation invokes `setproperty!`.
See [`PyBase.__getattr__`](@ref).
"""
PyBase.__setattr__

"""
    PyBase.delattr(self, ::Val{name::Symbol})
    PyBase.delattr(self, name::Symbol)

Python's attribute deletion interface
[`__delattr__`](https://docs.python.org/3/reference/datamodel.html#object.__delattr__).
Default implementation raises `AttributeError` in Python.
See [`PyBase.__getattr__`](@ref).
"""
PyBase.__delattr__

"""
Python's `__dir__`
"""
PyBase.__dir__

"""
    PyBase.convert_itemkey(self, key::Tuple)

Often [`__getitem__`](@ref), [`__setitem__`](@ref) and
[`__delitem__`](@ref) can be directly mapped to `getindex`,
`setindex!` and `delete!` respectively, provided that the
key/(multi-)index is mapped correctly.  This mapping is done by
`convert_itemkey`.  The default implementation does a lot of magics (=
run-time introspection) to convert Python/Numpy-like semantics to
Julia-like semantics.
"""
PyBase.convert_itemkey

"""
    __getitem__(self, key)

Python's
[`__getitem__`](https://docs.python.org/3/reference/datamodel.html#object.__getitem__).
Default implementation is:

```julia
getindex(self, convert_itemkey(self, key)...)
```
"""
PyBase.__getitem__

"""
    __setitem__(self, key, value)

Python's
[`__setitem__`](https://docs.python.org/3/reference/datamodel.html#object.__setitem__).
Default implementation is:

```julia
setindex!(self, value, convert_itemkey(self, key)...)
```

Note that the order of `key` and `value` for `__setitem__` is the opposite of
`setindex!`.
"""
PyBase.__setitem__

"""
    __delitem__(self, key)

Python's
[`__delitem__`](https://docs.python.org/3/reference/datamodel.html#object.__delitem__).
Default implementation is:

```julia
delete!(self, convert_itemkey(self, key)...)
```
"""
PyBase.__delitem__


text:
using YAActL, Printf

# define two functions for printing a message
function pr(msg)
    print(@sprintf("%s\n", msg))
    become(pr, "Next") # change behavior
end
pr(info, msg) = print(@sprintf("%s: %s\n", info, msg))

# a function for doing arithmetic
calc(op::F, x, y) where F<:Function = op(x, y)

# start an actor with the first behavior
myactor = Actor(pr)

cast!(myactor, "My first actor")     # send a message to it

cast!(myactor, "Something else")     # send again a message

become!(myactor, pr, "New behavior") # change the behavior to another one

cast!(myactor, "bla bla bla")        # and send again a message

become!(myactor, calc, +, 10);       # become a machine for adding to 10

call!(myactor, 5)                    # send a request to add 5

become!(myactor, ^);                 # become a exponentiation machine

call!(myactor, 123, 456)             # try it


text:
"""
    take_el_part(A, a, b, vibindices)

Take electric part specified by electric indices `a` and `b` from the A (type of Array).

"""
function take_el_part(A::Array, a, b, vibindices)
    a1 = vibindices[a][1]
    a2 = vibindices[a][end]
    b1 = vibindices[b][1]
    b2 = vibindices[b][end]

    return A[a1:a2, b1:b2]
end

"""
    MemoryKernel_1_traced(H_II_t, H_II_tau, W_bath, agg, FCProd, aggIndices, vibindices)

Calculate the first part of Memory Kernel with the definition

`` \\mathcal{M}_1(t, \\tau) = \\operatorname{tr}_B \\{ \\hat{H}_I^{(I)}(t) \\hat{H}_I^{(I)}(\\tau) W_\\text{bath} \\}``.

# Arguments
* `H_II_t`: Interaction Hamiltonian in interaction picutre at the time t, ``\\hat{H}_I^{(I)}(t)``.
* `H_II_tau`: Interaction Hamiltonian in interaction picutre at the time `tau`, ``\\hat{H}_I^{(I)}(\\tau)``.
* `W_bath`: Density matrix representing bath part of the density matrix, see [`get_rho_bath`](@ref).
* `agg`: Aggregate of molecules, see [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `vibindices`: Aggregate vibrational indices, see [`getVibIndices`](@ref).
"""
function MemoryKernel_1_traced(
    H_II_t::Array,
    H_II_tau::Array,
    W_bath::Array,
    agg,
    FCProd,
    aggIndices,
    vibindices
)
    aggIndLen = length(aggIndices)
    vibLen = length(vibindices[2])
    elLen = length(agg.molecules)
    elLen += 1
    MemoryKernel = zeros(ComplexF64, elLen, elLen, elLen, elLen)
    H_II_t_tau = H_II_t * H_II_tau

    for a = 1:elLen
        for c = 1:elLen
            H_II_t_tau_ac = take_el_part(H_II_t_tau, a, c, vibindices)
            for d = 1:elLen
                W_bath_cd = take_el_part(W_bath, c, d, vibindices)
                MK_big = H_II_t_tau_ac * W_bath_cd
                MemoryKernel[a, d, c, d] =
                    trace_bath_part(MK_big, a, d, agg, FCProd, aggIndices, vibindices)
            end
        end
    end
    return MemoryKernel
end

"""
    MemoryKernel_2_traced(H_II_t, H_II_tau, W_bath, agg, FCProd, aggIndices, vibindices)

Calculate the second part of Memory Kernel with the definition

`` \\mathcal{M}_1(t, \\tau) = \\operatorname{tr}_B \\{ \\hat{H}_I^{(I)}(t) W_\\text{bath} \\hat{H}_I^{(I)}(\\tau) \\}``.

# Arguments
* `H_II_t`: Interaction Hamiltonian in interaction picutre at the time t, ``\\hat{H}_I^{(I)}(t)``.
* `H_II_tau`: Interaction Hamiltonian in interaction picutre at the time `tau`, ``\\hat{H}_I^{(I)}(\\tau)``.
* `W_bath`: Density matrix representing bath part of the density matrix, see [`get_rho_bath`](@ref).
* `agg`: Aggregate of molecules, see [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `vibindices`: Aggregate vibrational indices, see [`getVibIndices`](@ref).
"""
function MemoryKernel_2_traced(
    H_II_t::Array,
    H_II_tau::Array,
    W_bath::Array,
    agg,
    FCProd,
    aggIndices,
    vibindices
)
    aggIndLen = length(aggIndices)
    vibLen = length(vibindices[2])
    elLen = length(agg.molecules)
    elLen += 1
    MemoryKernel = zeros(ComplexF64, elLen, elLen, elLen, elLen)

    for a = 1:elLen
        for b = 1:elLen
            for c = 1:elLen
                H_II_t_ac = take_el_part(H_II_t, a, c, vibindices)
                for d = 1:elLen
                    H_II_tau_db = take_el_part(H_II_tau, d, b, vibindices)
                    W_bath_cd = take_el_part(W_bath, c, d, vibindices)
                    MK_big = H_II_t_ac * W_bath_cd * H_II_tau_db
                    MemoryKernel[a, b, c, d] = trace_bath_part(
                        MK_big,
                        a,
                        b,
                        agg,
                        FCProd,
                        aggIndices,
                        vibindices,
                    )
                end
            end
        end
    end
    return MemoryKernel
end

"""
    MemoryKernel_3_traced(H_II_t, H_II_tau, W_bath, agg, FCProd, aggIndices, vibindices)

Calculate the third part of Memory Kernel with the definition

`` \\mathcal{M}_1(t, \\tau) = \\operatorname{tr}_B \\{ \\hat{H}_I^{(I)}(\\tau) W_\\text{bath} \\hat{H}_I^{(I)}(t) \\}``.

# Arguments
* `H_II_t`: Interaction Hamiltonian in interaction picutre at the time t, ``\\hat{H}_I^{(I)}(t)``.
* `H_II_tau`: Interaction Hamiltonian in interaction picutre at the time `tau`, ``\\hat{H}_I^{(I)}(\\tau)``.
* `W_bath`: Density matrix representing bath part of the density matrix, see [`get_rho_bath`](@ref).
* `agg`: Aggregate of molecules, see [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `vibindices`: Aggregate vibrational indices, see [`getVibIndices`](@ref).
"""
function MemoryKernel_3_traced(
    H_II_t::Array,
    H_II_tau::Array,
    W_bath::Array,
    agg,
    FCProd,
    aggIndices,
    vibindices
)
    aggIndLen = length(aggIndices)
    vibLen = length(vibindices[2])
    elLen = length(agg.molecules)
    elLen += 1
    MemoryKernel = zeros(ComplexF64, elLen, elLen, elLen, elLen)

    for a = 1:elLen
        for b = 1:elLen
            for c = 1:elLen
                H_II_tau_ac = take_el_part(H_II_tau, a, c, vibindices)
                for d = 1:elLen
                    W_bath_cd = take_el_part(W_bath, c, d, vibindices)
                    H_II_t_db = take_el_part(H_II_t, d, b, vibindices)
                    MK_big = H_II_tau_ac * W_bath_cd * H_II_t_db
                    MemoryKernel[a, b, c, d] = trace_bath_part(
                        MK_big,
                        a,
                        b,
                        agg,
                        FCProd,
                        aggIndices,
                        vibindices,
                    )
                end
            end
        end
    end
    return MemoryKernel
end

"""
    MemoryKernel_4_traced(H_II_t, H_II_tau, W_bath, agg, FCProd, aggIndices, vibindices)

Calculate the fourth part of Memory Kernel with the definition

`` \\mathcal{M}_1(t, \\tau) = \\operatorname{tr}_B \\{ W_\\text{bath} \\hat{H}_I^{(I)}(t) \\hat{H}_I^{(I)}(\\tau) \\}``.

# Arguments
* `H_II_t`: Interaction Hamiltonian in interaction picutre at the time t, ``\\hat{H}_I^{(I)}(t)``.
* `H_II_tau`: Interaction Hamiltonian in interaction picutre at the time `tau`, ``\\hat{H}_I^{(I)}(\\tau)``.
* `W_bath`: Density matrix representing bath part of the density matrix, see [`get_rho_bath`](@ref).
* `agg`: Aggregate of molecules, see [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `vibindices`: Aggregate vibrational indices, see [`getVibIndices`](@ref).
"""
function MemoryKernel_4_traced(
    H_II_t::Array,
    H_II_tau::Array,
    W_bath::Array,
    agg,
    FCProd,
    aggIndices,
    vibindices
)
    aggIndLen = length(aggIndices)
    vibLen = length(vibindices[2])
    elLen = length(agg.molecules)
    elLen += 1
    MemoryKernel = zeros(ComplexF64, elLen, elLen, elLen, elLen)
    H_II_tau_t = H_II_tau * H_II_t

    for a = 1:elLen
        for b = 1:elLen
            for d = 1:elLen
                W_bath_ad = take_el_part(W_bath, a, d, vibindices)
                H_II_tau_t_db = take_el_part(H_II_tau_t, d, b, vibindices)
                MK_big = W_bath_ad * H_II_tau_t_db
                MemoryKernel[a, b, a, d] =
                    trace_bath_part(MK_big, a, b, agg, FCProd, aggIndices, vibindices)
            end
        end
    end
    return MemoryKernel
end

"""
    MemoryKernel_traced(H_II_t, H_II_tau, W_bath, agg, FCProd, aggIndices, vibindices)

Calculate Memory Kernel with the definition

`` \\mathcal{M}(t, \\tau) = \\operatorname{tr}_B \\{ [ \\hat{H}_I^{(I)}(t), [ \\hat{H}_I^{(I)}(\\tau), W_\\text{bath} ]]\\}``.

# Arguments
* `H_II_t`: Interaction Hamiltonian in interaction picutre at the time t, ``\\hat{H}_I^{(I)}(t)``.
* `H_II_tau`: Interaction Hamiltonian in interaction picutre at the time `tau`, ``\\hat{H}_I^{(I)}(\\tau)``.
* `W_bath`: Density matrix representing bath part of the density matrix, see [`get_rho_bath`](@ref).
* `agg`: Aggregate of molecules, see [`Aggregate`](@ref).
* `aggIndices`: Aggregate indices, see [`getIndices`](@ref).
* `vibindices`: Aggregate vibrational indices, see [`getVibIndices`](@ref).
"""
function MemoryKernel_traced(
    H_II_t::Array,
    H_II_tau::Array,
    W_bath::Array,
    agg,
    FCProd,
    aggIndices,
    vibindices
)
    aggIndLen = length(aggIndices)
    vibLen = length(vibindices[2])
    elLen = length(agg.molecules)
    elLen += 1
    MemoryKernel = zeros(ComplexF64, elLen, elLen, elLen, elLen)
    H_II_tau_t = H_II_tau * H_II_t
    H_II_t_tau = H_II_t * H_II_tau

    for a = 1:elLen
        for c = 1:elLen
            H_II_t_tau_ac = take_el_part(H_II_t_tau, a, c, vibindices)
            H_II_tau_ac = take_el_part(H_II_tau, a, c, vibindices)
            H_II_t_ac = take_el_part(H_II_t, a, c, vibindices)
            for d = 1:elLen
                W_bath_cd = take_el_part(W_bath, c, d, vibindices)
                W_bath_cd = take_el_part(W_bath, c, d, vibindices)
                for b = 1:elLen
                    H_II_tau_db = take_el_part(H_II_tau, d, b, vibindices)

                    MK_big = -H_II_t_ac * W_bath_cd * H_II_tau_db

                    H_II_t_db = take_el_part(H_II_t, d, b, vibindices)
                    MK_big[:, :] -= H_II_tau_ac * W_bath_cd * H_II_t_db

                    if a == c
                        W_bath_ad = take_el_part(W_bath, a, d, vibindices)
                        H_II_tau_t_db = take_el_part(H_II_tau_t, d, b, vibindices)
                        MK_big[:, :] += W_bath_ad * H_II_tau_t_db
                    end

                    if d == b
                        MK_big[:, :] += H_II_t_tau_ac * W_bath_cd
                    end
                    MemoryKernel[a, b, c, d] = trace_bath_part(
                        MK_big,
                        a,
                        b,
                        agg,
                        FCProd,
                        aggIndices,
                        vibindices,
                    )
                end
            end
        end
    end
    return MemoryKernel
end