text:
using Oceananigans.Architectures: architecture
using Oceananigans.Grids: interior_parent_indices
using Statistics: norm, dot
using LinearAlgebra

mutable struct PreconditionedConjugateGradientSolver{A, G, L, T, F, M, P}
               architecture :: A
                       grid :: G
          linear_operation! :: L
                  tolerance :: T
         maximum_iterations :: Int
                  iteration :: Int
                       ρⁱ⁻¹ :: T
    linear_operator_product :: F
           search_direction :: F
                   residual :: F
              precondition! :: M
     preconditioner_product :: P
end

no_precondition!(args...) = nothing

initialize_precondition_product(precondition, template_field) = similar(template_field)
initialize_precondition_product(::Nothing, template_field) = nothing

maybe_precondition(::Nothing, ::Nothing, r, args...) = r

function maybe_precondition(precondition!, z, r, args...)
    precondition!(z, r, args...)
    return z
end

"""
    PreconditionedConjugateGradientSolver(linear_operation;
                                          template_field,
                                          maximum_iterations = size(template_field.grid),
                                          tolerance = 1e-13,
                                          precondition = nothing)

Returns a PreconditionedConjugateGradientSolver that solves the linear equation
``A x = b`` using a iterative conjugate gradient method with optional preconditioning.
The solver is used by calling

```
solve!(x, solver::PreconditionedConjugateGradientOperator, b, args...)
```

for `solver`, right-hand side `b`, solution `x`, and optional arguments `args...`.

Arguments
=========

* `template_field`: Dummy field that is the same type and size as `x` and `b`, which
                    is used to infer the `architecture`, `grid`, and to create work arrays
                    that are used internally by the solver.

* `linear_operation`: Function with signature `linear_operation!(p, y, args...)` that calculates
                     `A*y` and stores the result in `p` for a "candidate solution `y`. `args...`
                     are optional positional arguments passed from `solve!(x, solver, b, args...)`.

* `maximum_iterations`: Maximum number of iterations the solver may perform before exiting.

* `tolerance`: Tolerance for convergence of the algorithm. The algorithm quits when
               `norm(A * x - b) < tolerance`.

* `precondition`: Function with signature `preconditioner!(z, y, args...)` that calculates
                  `P * y` and stores the result in `z` for linear operator `P`.
                  Note that some precondition algorithms describe the step
                  "solve `M * x = b`" for precondition `M`"; in this context,
                  `P = M⁻¹`.

See [`solve!`](@ref) for more information about the preconditioned conjugate-gradient algorithm.
"""
function PreconditionedConjugateGradientSolver(linear_operation;
                                               template_field::AbstractField,
                                               maximum_iterations = prod(size(template_field)),
                                               tolerance = 1e-13, #sqrt(eps(eltype(template_field.grid))),
                                               precondition = nothing)

    arch = architecture(template_field)
    grid = template_field.grid

    # Create work arrays for solver
    linear_operator_product = similar(template_field) # A*xᵢ = qᵢ
    search_direction = similar(template_field) # pᵢ
            residual = similar(template_field) # rᵢ

    # Either nothing (no precondition) or P*xᵢ = zᵢ
    precondition_product = initialize_precondition_product(precondition, template_field)

    return PreconditionedConjugateGradientSolver(arch,
                                                 grid,
                                                 linear_operation,
                                                 tolerance,
                                                 maximum_iterations,
                                                 0,
                                                 0.0,
                                                 linear_operator_product,
                                                 search_direction,
                                                 residual,
                                                 precondition,
                                                 precondition_product)
end

"""
    solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)

Solve `A * x = b` using an iterative conjugate-gradient method, where `A * x` is
determined by `solver.linear_operation`

See figure 2.5 in

> The Preconditioned Conjugate Gradient Method in "Templates for the Solution of Linear Systems: Building Blocks for Iterative Methods" Barrett et. al, 2nd Edition.

Given:
  * Linear Preconditioner operator `M!(solution, x, other_args...)` that computes `M * x = solution`
  * A matrix operator `A` as a function `A()`;
  * A dot product function `norm()`;
  * A right-hand side `b`;
  * An initial guess `x`; and
  * Local vectors: `z`, `r`, `p`, `q`

This function executes the algorithm

```
β  = 0
r = b - A(x)
iteration  = 0

Loop:
     if iteration > maximum_iterations
        break
     end

     ρ = r ⋅ z

     z = M(r)
     β = ρⁱ⁻¹ / ρ
     p = z + β * p
     q = A(p)

     α = ρ / (p ⋅ q)
     x = x + α * p
     r = r - α * q

     if |r| < tolerance
        break
     end

     iteration += 1
     ρⁱ⁻¹ = ρ
```
"""
function solve!(x, solver::PreconditionedConjugateGradientSolver, b, args...)

    # Initialize
    solver.iteration = 0

    # q = A*x
    q = solver.linear_operator_product
    solver.linear_operation!(q, x, args...)

    # r = b - A*x
    solver.residual .= b .- q

    @debug "PreconditionedConjugateGradientSolver, |b|: $(norm(b))"
    @debug "PreconditionedConjugateGradientSolver, |A(x)|: $(norm(q))"

    while iterating(solver)
        iterate!(x, solver, b, args...)
    end

    fill_halo_regions!(x, solver.architecture) # blocking

    return x
end

function iterate!(x, solver, b, args...)
    r = solver.residual
    p = solver.search_direction
    q = solver.linear_operator_product

    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |r|: $(norm(r))"

    # Preconditioned:   z = P * r
    # Unpreconditioned: z = r
    z = maybe_precondition(solver.precondition!, solver.preconditioner_product, r, args...)
    ρ = dot(z, r)

    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), ρ: $ρ"
    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |z|: $(norm(z))"

    if solver.iteration == 0
        p .= z
    else
        β = ρ / solver.ρⁱ⁻¹
        p .= z .+ β .* p

        @debug "PreconditionedConjugateGradientSolver $(solver.iteration), β: $β"
    end

    # q = A * p
    solver.linear_operation!(q, p, args...)
    α = ρ / dot(p, q)

    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), |q|: $(norm(q))"
    @debug "PreconditionedConjugateGradientSolver $(solver.iteration), α: $α"

    x .+= α .* p
    r .-= α .* q

    solver.iteration += 1
    solver.ρⁱ⁻¹ = ρ

    return nothing
end

function iterating(solver)
    # End conditions
    solver.iteration >= solver.maximum_iterations && return false
    norm(solver.residual) <= solver.tolerance && return false
    return true
end

function Base.show(io::IO, solver::PreconditionedConjugateGradientSolver)
    print(io, "Oceananigans-compatible preconditioned conjugate gradient solver.\n")
    print(io, " Problem size = "  , size(solver.q), '\n')
    print(io, " Grid = "  , solver.grid)
    return nothing
end


text:
push!(LOAD_PATH,"/home/simone/Documents/uni/postdoc/projects/FeedbackParticleFilters.jl/")
using Documenter, FeedbackParticleFilters

makedocs(
    modules = [FeedbackParticleFilters],
    format = Documenter.HTML(),
    clean = false,              # do we clean build dir
    sitename = "FeedbackParticle Filters.jl",
    authors = "Simone Carlo Surace",
    pages = Any[ # Compat: `Any` for 0.4 compat
        "Home" => "index.md",
        "Getting started" => "getting_started.md",
        "Background" => Any[
             "filtering.md",
             "fpf.md",
        ],
        "Manual" => Any[
             "basic.md",
             "hidden.md",
             "observation.md",
             "algorithms.md",
             "gainest.md",
        ],
    ],
)

#deploydocs(
   # repo = "github.com/simsurace/FeedbackParticleFilters.jl.git",
  #  target = "build",
    # julia = "1.0",
    # osname = "linux",
    # make = nothing,
    # deps = nothing,
    # deps   = Deps.pip("mkdocs", "python-markdown-math"),
#)
deploydocs(
    repo = "github.com/simsurace/FeedbackParticleFilters.jl.git",
)

ext: jl

max_stars_repo_name: simsurace/FeedbackParticleFilters

max_issues_repo_name: simsurace/FeedbackParticleFilters

max_forks_repo_name: simsurace/FeedbackParticleFilters


text:
module MathOptInterfaceECOS

export ECOSOptimizer

using MathOptInterface
const MOI = MathOptInterface
const CI = MOI.ConstraintIndex
const VI = MOI.VariableIndex

const MOIU = MOI.Utilities

const SF = Union{MOI.SingleVariable, MOI.ScalarAffineFunction{Float64}, MOI.VectorOfVariables, MOI.VectorAffineFunction{Float64}}
const SS = Union{MOI.EqualTo{Float64}, MOI.GreaterThan{Float64}, MOI.LessThan{Float64}, MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives, MOI.SecondOrderCone, MOI.ExponentialCone}

using ECOS

include("types.jl")

mutable struct ECOSOptimizer <: MOI.AbstractOptimizer
    cone::Cone
    maxsense::Bool
    data::Union{Void, Data} # only non-Void between MOI.copy! and MOI.optimize!
    sol::Solution
    function ECOSOptimizer()
        new(Cone(), false, nothing, Solution())
    end
end

function MOI.isempty(instance::ECOSOptimizer)
    !instance.maxsense && instance.data === nothing
end

function MOI.empty!(instance::ECOSOptimizer)
    instance.maxsense = false
    instance.data = nothing # It should already be nothing except if an error is thrown inside copy!
end

MOI.canaddvariable(instance::ECOSOptimizer) = false
MOIU.needsallocateload(instance::ECOSOptimizer) = true
MOI.supports(::ECOSOptimizer, ::MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}) = true
MOI.supportsconstraint(::ECOSOptimizer, ::Type{<:SF}, ::Type{<:SS}) = true
MOI.copy!(dest::ECOSOptimizer, src::MOI.ModelLike; copynames=true) = MOIU.allocateload!(dest, src, copynames)

# Implements optimize! : translate data to ECOSData and call ECOS_solve
include("solve.jl")

# Implements getter for result value and statuses
include("attributes.jl")

end # module


text:
VERSION >= v"0.4.0" && __precompile__(true)

module Indicators
    export
        runmean, runsum, runvar, runsd, runcov, runcor, runmax, runmin, runmad, runquantile, runacf,
        wilder_sum, mode, diffn,
        sma, trima, wma, ema, mma, kama, mama, hma, swma, dema, tema, alma, zlema, vwma, vwap, hama,
        mlr_beta, mlr_slope, mlr_intercept, mlr, mlr_se, mlr_ub, mlr_lb, mlr_bands, mlr_rsq,
        aroon, donch, momentum, roc, macd, rsi, adx, psar, kst, wpr, cci, stoch, smi,
        bbands, tr, atr, keltner,
        crossover, crossunder,
        renko,
        maxima, minima, support, resistance,
        rsrange, hurst

    include("run.jl")
    include("ma.jl")
    include("reg.jl")
    include("mom.jl")
    include("vol.jl")
    include("trendy.jl")
    include("utils.jl")
    include("patterns.jl")
    include("chaos.jl")

    using Temporal
    include("temporal.jl")
end


text:
using PkgBenchmark

benchmarkpkg("AugmentedGaussianProcesses")


text:
using LightGraphs: LightGraphs, DiGraph, add_edge!, add_vertex!, dst, edges, src

struct Dependency <: Graph
    g::DiGraph
    N::Vector{VarNode}
    I::Dict{VarNode,Int}
    M::Dict{Symbol,VarInfo}
end

dependency(M::Dict{Symbol,VarInfo}) = Dependency(DiGraph(), VarNode[], Dict{VarNode,VarInfo}(), M)
dependency(V::Vector{VarInfo}) = begin
    M = Dict{Symbol,VarInfo}()
    for v in V
        for n in names(v)
            M[n] = v
        end
    end
    d = dependency(M)
    add!(d, V)
    d
end
dependency(::Type{S}) where {S<:System} = dependency(geninfos(S))
dependency(::S) where {S<:System} = dependency(S)

graph(d::Dependency) = d.g

node!(d::Dependency, n::VarNode) = begin
    if !haskey(d.I, n)
        add_vertex!(d.g)
        push!(d.N, n)
        d.I[n] = length(d.N)
        #@show "new vertex at $(d.I[n])"
    end
    n
end
node!(d::Dependency, v::VarInfo, t::VarStep) = node!(d, VarNode(v, t))
node!(d::Dependency, v::Symbol, t::VarStep) = node!(d, VarNode(d.M[v], t))
prenode!(d::Dependency, v) = node!(d, v, PreStep())
mainnode!(d::Dependency, v) = node!(d, v, MainStep())
postnode!(d::Dependency, v) = node!(d, v, PostStep())

firstnode(d::Dependency, v::VarInfo) = begin
    for t in (PreStep, MainStep, PostStep)
        n = VarNode(v, t())
        haskey(d.I, n) && return n
    end
    @error "no node found for $a"
end

extract(v::VarInfo; equation=true, tag=true, include=(), exclude=()) = begin
    pick(a) = let k, v; extractfuncargdep(@capture(a, k_=v_) ? v : a) end
    pack(A) = Iterators.flatten(filter!(!isnothing, pick.(A))) |> Tuple
    eq = equation ? pack(v.args) : ()
    #@show eq
    #HACK: exclude internal tags (i.e. _type)
    #HACK: filter included/excluded tags
    #TODO: share logic with filterconstructortags() in macro?
    tagfilter(t) = !startswith(String(t), "_") && (isempty(include) ? true : t ∈ include) && (isempty(exclude) ? true : t ∉ exclude)
    par = tag ? Tuple(Iterators.flatten(filter!(!isnothing, [extractfuncargdep(p[2]) for p in v.tags if tagfilter(p[1])]))) : ()
    #@show par
    Set([eq..., par...]) |> collect
end

link!(d::Dependency, v::VarInfo, n::VarNode; kwargs...) = begin
    A = extract(v; kwargs...)
    #HACK: skip missing refs to allow const variable patch syntax (i.e. @system S{x=1})
    V = [d.M[a] for a in A if haskey(d.M, a)]
    for v0 in V
        # self reference is handled within a single node
        if v0 == v
            continue
        elseif v0.state == :Bisect
            n0 = prenode!(d, v0)
            link!(d, n0, n)
        elseif isnothing(v0.state) && istag(v0, :context)
            n1 = mainnode!(d, v0)
            n2 = postnode!(d, v0)
            link!(d, n1, n)
            link!(d, n, n2)
        else
            n1 = mainnode!(d, v0)
            link!(d, n1, n)
        end
    end
end

link!(d::Dependency, a::VarNode, b::VarNode) = begin
    #@show "link: add edge $(a.info.name) ($(d.I[a])) => $(b.info.name) ($(d.I[b]))"
    add_edge!(d.g, d.I[a], d.I[b])
end

add!(d::Dependency, v::VarInfo) = begin
    #@show "add! $v"
    if v.state == :Accumulate || v.state == :Capture
        n1 = mainnode!(d, v)
        n2 = postnode!(d, v)
        link!(d, n1, n2)
        # needs `time` tags update, but equation args should be excluded due to cyclic dependency
        #HACK: support `when` tag while avoiding cyclic dependency
        #TODO: more elegant way to handle tags include/exclude
        link!(d, v, n1; equation=false, exclude=(:when,))
        link!(d, v, n2; equation=false, include=(:when,))
        link!(d, v, n2)
    elseif v.state == :Bisect
        n0 = prenode!(d, v)
        n1 = mainnode!(d, v)
        link!(d, n0, n1)
        # needs `lower/upper` tags
        link!(d, v, n0; equation=false)
        link!(d, v, n1)
        # needs access to context in Bisect constructor (otherwise convergence would fail)
        c = mainnode!(d, :context)
        link!(d, c, n0)
    elseif v.state == :Produce
        n0 = prenode!(d, v)
        n1 = mainnode!(d, v)
        n2 = postnode!(d, v)
        link!(d, n0, n1)
        link!(d, n1, n2)
        # no tag available for produce, but just in case we need one later
        link!(d, v, n0; equation=false)
        link!(d, v, n2)
        # make sure context get updated before updating subtree
        c = mainnode!(d, :context)
        link!(d, c, n1)
    elseif isnothing(v.state) && istag(v, :context)
        n0 = prenode!(d, v)
        n1 = mainnode!(d, v)
        n2 = postnode!(d, v)
        link!(d, n0, n1)
        link!(d, n1, n2)
        link!(d, v, n0)
        link!(d, v, n2)
    else
        n = mainnode!(d, v)
        link!(d, v, n)
    end
    if istag(v, :parameter)
        c = mainnode!(d, :config)
        n = firstnode(d, v)
        link!(d, c, n)
    end
end
add!(d::Dependency, V::Vector{VarInfo}) = begin
    for v in V
        add!(d, v)
    end
end

Base.sort(d::Dependency) = begin
    C = LightGraphs.simplecycles(d.g)
    !isempty(C) && error("no cyclic dependency allowed: $([[d.N[i].info.name for i in c] for c in C])")
    J = LightGraphs.topological_sort_by_dfs(d.g)
    [d.N[i] for i in J]
end

label(n::VarNode; alias=false) = begin
    v = n.info
    name = alias && !isnothing(v.alias) ? v.alias : v.name
    name = string(name)
    tag = string(n.step)
    tag * name
end
labels(d::Dependency; kw...) = label.(d.N; kw...)

edgestyle(d::Dependency, a::VarNode, b::VarNode) = ""
edgestyles(d::Dependency; kw...) = Dict(
    let a = src(e), b = dst(e)
        (a, b) => edgestyle(d, d.N[a], d.N[b])
    end for e in edges(d.g)
)

Base.show(io::IO, d::Dependency) = print(io, "Dependency")
Base.show(io::IO, ::MIME"text/plain", d::Dependency) = begin
    color = get(io, :color, false)
    VC = tokencolor(VarColor(); color)
    MC = tokencolor(MiscColor(); color)
    print(io, MC("["))
    print(io, join(VC.(label.(sort(d))), MC(" → ")))
    print(io, MC("]"))
end


text:
using DocumenterGitLabCI
using Test

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


text:
module HuberLR

export LinearRegression, fit!, predict

using Optim

type LinearRegression
    δ::Float64
    w::Vector{Float64}
    function LinearRegression(δ)
        if δ < 0.
            error("δ should be non-negative")
        end
        new(δ, Float64[])
    end
end

function Base.show(io::IO, lr::LinearRegression)
    print(io, "LinearRegression:\n")
    print(io, "    δ: $(lr.δ)\n")
    ws = ASCIIString[]
    for w in lr.w
        s = @sprintf "%.4f" w
        push!(ws, s)
    end
    print(io, "    w: $(join(ws, ", "))")
end

immutable Data
    x::Matrix{Float64}
    y::Vector{Float64}
    function Data(x, y)
        if size(x, 2) != length(y)
            error("the number of data points must match")
        end
        new(x, y)
    end
end

Base.size(data::Data) = length(data.y)
dim(data::Data) = size(data.x, 1)

function fit!(lr::LinearRegression, x::Matrix{Float64}, y::Vector{Float64}; method=:l_bfgs, show_trace=false)
    D, N = size(x)
    x = vcat(ones(N)', x)
    w₀ = randn(D + 1)
    data = Data(x, y)
    if method ∈ [:l_bfgs, :bfgs]
        opt = optimize(w -> loss(w, lr, data), (w, ∇) -> ∇loss!(w, ∇, lr, data), w₀, method=method, show_trace=show_trace)
    elseif method ∈ [:nelder_mead]
        opt = optimize(w -> loss(w, lr, data), w₀, method=method, show_trace=show_trace)
    else
        error("method $method is not supported")
    end
    lr.w = opt.minimum
    lr
end

function predict(lr::LinearRegression, x::Matrix{Float64})
    D, N = size(x)
    if isempty(lr.w)
        error("the model is not yet trained")
    elseif length(lr.w) != D + 1
        error("dimention mismatch")
    end
    ŷ = Array(Float64, N)
    @inbounds for n in 1:N
        s = lr.w[1]
        for j in 1:D
            s += lr.w[j+1] * x[j,n]
        end
        ŷ[n] = s
    end
    ŷ
end

# dot product along the n-th column
function dotat(w::Vector{Float64}, x::Matrix{Float64}, n::Int)
    s = 0.
    @inbounds for i in 1:length(w)
        s += w[i] * x[i,n]
    end
    s
end

# Huber loss function
function loss(w::Vector{Float64}, lr::LinearRegression, data::Data)
    N = size(data)
    δ = lr.δ
    ∑loss = 0.
    @inbounds for n in 1:N
        r = data.y[n] - dotat(w, data.x, n)
        ∑loss += ifelse(abs(r) < δ, r^2 / 2., abs(r) * δ - δ^2 / 2.)
    end
    ∑loss
end

# first derivative of Huber loss function
function ∇loss!(w::Vector{Float64}, ∇::Vector{Float64}, lr::LinearRegression, data::Data)
    N = size(data)
    D = dim(data)
    δ = lr.δ
    fill!(∇, 0.)
    @inbounds for n in 1:N
        r = data.y[n] - dotat(w, data.x, n)
        ∂loss = ifelse(abs(r) < δ, r, sign(r) * δ)
        for j in 1:D
            ∂r = -data.x[j,n]
            ∇[j] += ∂loss * ∂r
        end
    end
    ∇
end

end # module


text:
using Plots
using StaticArrays

mutable struct NeoOrb
    fper::Float64
    dtau::Float64
    dtaumax::Float64
    v0::Float64
    n_e::Int32
    n_d::Int32
    integmode::Int32
    relerr::Float64
    firstrun::Bool

    NeoOrb() = new()
  end

mutable struct FieldCan
    Ath::Float64
    Aph::Float64
    hth::Float64
    hph::Float64
    Bmod::Float64

    dAth::SVector{3, Float64}
    dAph::SVector{3, Float64}
    dhth::SVector{3, Float64}
    dhph::SVector{3, Float64}
    dBmod::SVector{3, Float64}

    # second derivatives: drdr, drdhth, drdph, dthdth, dthdph, dphdph
    d2Ath::SVector{6, Float64}
    d2Aph::SVector{6, Float64}
    d2hth::SVector{6, Float64}
    d2hph::SVector{6, Float64}
    d2Bmod::SVector{6, Float64}

    H::Float64
    pth::Float64
    vpar::Float64

    dvpar::SVector{4, Float64}
    dH::SVector{4, Float64}
    dpth::SVector{4, Float64}

    # order of second derivatives:
    # d2dr2, d2drdth, d2drph, d2dth2, d2dthdph, d2dph2,
    # d2dpphdr, d2dpphdth, d2dpphdph, d2dpph2
    d2vpar::SVector{10, Float64}
    d2H::SVector{10, Float64}
    d2pth::SVector{10, Float64}

    mu::Float64
    ro0::Float64

    FieldCan() = new()
end

norb = NeoOrb()
if !@isdefined(firstrun) | firstrun == true
    global firstrun, norb
    ccall((:__neo_orb_global_MOD_init_field, "libsimple"),
    Cvoid, (Ref{NeoOrb}, Ref{Int32}, Ref{Int32}, Ref{Int32}, Ref{Int32}),
    norb, 5, 5, 3, 0)
    firstrun = false
end

Z_charge = 2
m_mass = 4
c=2.9979e10
e_charge=4.8032e-10
e_mass=9.1094e-28
p_mass=1.6726e-24
ev=1.6022e-12

bmod_ref = 5.0e4
Z_charge = 2
m_mass = 4
E_kin = 3.5e6
trace_time = 1.0e-2

v0 = sqrt(2.0*E_kin*ev/(m_mass*p_mass))    # alpha particle velocity, cm/s
rlarm=v0*m_mass*p_mass*c/(Z_charge*e_charge*bmod_ref) # reference gyroradius
tau=trace_time*v0

npoiper2 = 64

p_rmajor = cglobal((:__new_vmec_stuff_mod_MOD_rmajor, "libsimple"), Float64)
rmajor = unsafe_load(p_rmajor)
println("R0 = ", rmajor)
rbig = rmajor*1.0e2
dtaumin = 2.0*pi*rbig/npoiper2
dtau = dtaumin
ccall((:__neo_orb_global_MOD_init_params, "libsimple"), Cvoid,
      (Ref{NeoOrb}, Ref{Int32}, Ref{Int32}, Ref{Float64}, Ref{Float64},
       Ref{Float64}, Ref{Float64}),
      norb, Z_charge, m_mass, E_kin, dtau, dtaumin, 1e-8)

s = [0.5]
th = [0.0]
ph = [0.314]
pabs = [1.0]
lam = [0.22]

z0 = [0.5, 0.0, 0.314, 1.0, 0.22]

rtol = 1e-12

kbuf = 0
kt = 0
k = 0

dt = dtaumin/sqrt(2e0) # factor 1/sqrt(2) due to velocity normalisation different from other modules

f = FieldCan()
ccall((:__field_can_mod_MOD_eval_field_can, "libsimple"), Cvoid,
        (Ref{FieldCan}, Ref{Float64}, Ref{Float64}, Ref{Float64}, Ref{Int32}),
        f, s, th, ph, 0)

mu = 0.5 * z0[4] * (1.0 - z0[5]^2)/f.Bmod*2.0 # mu by factor 2 from other modules

ro0_parmot = unsafe_load(
    cglobal((:__parmot_mod_MOD_ro0, "libsimple"), Float64)
)

f.ro0 = ro0_parmot/sqrt(2.0) # ro0 = mc/e*v0, different by sqrt(2) from other modules
f.vpar = z0[4]*z0[5]*sqrt(2.0) # vpar_bar = vpar/sqrt(T/m), different by sqrt(2) from other modules

# ! initialize canonical variables
z = zeros(4)
z[1:3] = z0[1:3]  # r, th, ph
z[4] = f.vpar*f.hph + f.Aph/f.ro0
ccall((:__field_can_mod_MOD_get_val, "libsimple"), Cvoid,
        (Ref{FieldCan}, Ref{Float64}),
        f, z[4])  # for pth

# call plag_coeff(nlag+1, 0, 1d0, 1d0*(/(k,k=-nlag,0)/), si%coef)

# ccall((:__field_can_mod_MOD_get_val, "lib/libsimple"), Cvoid,
#                    (Ref{FieldCan}, Ref{Float64}), f, pphi)

function timestep!(z)

end

max_stars_repo_name: itpplasma/SIMPLE

max_issues_repo_name: landreman/SIMPLE

max_forks_repo_name: landreman/SIMPLE


text:
# using Revise,
using Test, ForwardDiff, Parameters, Setfield, LinearAlgebra
# using Plots
using BifurcationKit, Test
const BK = BifurcationKit

norminf = x -> norm(x, Inf)
####################################################################################################
function Lor(u, p)
	@unpack α,β,γ,δ,G,F,T = p
	X,Y,Z,U = u
	[
		-Y^2 - Z^2 - α*X + α*F - γ*U^2,
		X*Y - β*X*Z - Y + G,
		β*X*Y + X*Z - Z,
		-δ*U + γ*U*X + T
	]
end
jet = BK.getJet(Lor;matrixfree=false)
parlor = (α = 1//4, β = 1, G = .25, δ = 1.04, γ = 0.987, F = 1.7620532879639, T = .0001265)

opts_br = ContinuationPar(pMin = -1.5, pMax = 3.0, ds = 0.002, dsmax = 0.15, nInversion = 6, detectBifurcation = 3, maxBisectionSteps = 25, nev = 4, maxSteps = 200, plotEveryStep = 30)

z0 =  [2.9787004394953343, -0.03868302503393752,  0.058232737694740085, -0.02105288273117459]

br, = @time continuation(jet[1], jet[2], z0, setproperties(parlor;T=0.04,F=3.), (@lens _.F), opts_br;
	recordFromSolution = (x, p) -> (X = x[1], Y = x[2], Z= x[3], U = x[4]),
	normC = norminf,
	tangentAlgo = BorderedPred(),
	bothside = true)
####################################################################################################
# this part is for testing the spectrum
recordFromSolutionLor(x, p) = (u = BK.getVec(x);(X = u[1], Y = u[2], Z = u[3], U = u[4]))
@set! opts_br.newtonOptions.verbose = false

sn_codim2_test, = continuation(jet[1:2]..., br, 4, (@lens _.T), ContinuationPar(opts_br, pMax = 3.2, pMin = -0.1, detectBifurcation = 1, dsmin=1e-5, ds = -0.001, dsmax = 0.005, nInversion = 8, saveSolEveryStep = 1, maxSteps = 60) ;
	normC = norminf,
	detectCodim2Bifurcation = 1,
	updateMinAugEveryStep = 1,
	startWithEigen = true,
	recordFromSolution = recordFromSolutionLor,
	d2F = jet[3], d3F = jet[4],
	bdlinsolver = MatrixBLS(),
	)

hp_codim2_test, = continuation(jet[1:2]..., br, 1, (@lens _.T), ContinuationPar(opts_br, ds = -0.001, dsmax = 0.02, dsmin = 1e-4, nInversion = 6, saveSolEveryStep = 1) ;
	normC = norminf,
	tangentAlgo = BorderedPred(),
	detectCodim2Bifurcation = 1,
	updateMinAugEveryStep = 1,
	startWithEigen = true,
	recordFromSolution = recordFromSolutionLor,
	d2F = jet[3], d3F = jet[4],
	bothside = true,
	bdlinsolver = MatrixBLS())
####################################################################################################
"""
This function test if the eigenvalues are well computed during a branch of Hopf/Fold.
This test works if the problem is not updated because we dont record prob.a and prob.b
"""
function testEV(br, verbose = false)
	verbose && println("\n\n\n"*"="^50)
	prob = br.functional
	eig = DefaultEig()
	lens1 = br.lens
	lens2 = prob.lens
	par0 = br.params
	ϵ = br.contparams.newtonOptions.tol

	for (ii, pt) in enumerate(br.branch)
		# we make sure the parameters are set right
		step = pt.step
		verbose && (println("="^50); @info step ii)
		x0 = BK.getVec(br.sol[ii].x, prob)
		p0 = BK.getP(br.sol[ii].x, prob)[1]
		if prob isa HopfProblemMinimallyAugmented
			ω0 = BK.getP(br.sol[ii].x, prob)[2]
		end
		p1 = br.sol[ii].p
		@test p1 == get(pt, lens1)
		@test p0 == get(pt, lens2)

		# we test the functional
		par1 = set(par0, lens1, p1)
		par1 = set(par1, lens2, p0)
		@test par1.T == pt.T && par1.F == pt.F
		resf = prob.F(x0, par1)
		@test norminf(resf) < ϵ
		if prob isa FoldProblemMinimallyAugmented
			res = prob(x0, p0, set(par0, lens1, p1))
		else
			res = prob(x0, p0, ω0, set(par0, lens1, p1))
		end
		@test resf == res[1]
		verbose && @show res
		@test norminf(res[1]) < 100ϵ
		@test norminf(res[2]) < 100ϵ
		par0 = merge(br.params, pt)

		# we test the eigenvalues
		vp = eigenvals(br, step)
		J = jet[2](x0, par1)
		vp2 = eig(J, 4)[1]
		verbose && display(hcat(vp, vp2))
		@test vp == vp2
	end
end

testEV(sn_codim2_test)
testEV(hp_codim2_test)
####################################################################################################
@set! opts_br.newtonOptions.verbose = false

sn_codim2, = continuation(jet[1:2]..., br, 4, (@lens _.T), ContinuationPar(opts_br, pMax = 3.2, pMin = -0.1, detectBifurcation = 1, dsmin=1e-5, ds = -0.001, dsmax = 0.015, nInversion = 10, saveSolEveryStep = 1, maxSteps = 30, maxBisectionSteps = 55) ;
	normC = norminf,
	detectCodim2Bifurcation = 2,
	updateMinAugEveryStep = 1,
	startWithEigen = true,
	recordFromSolution = recordFromSolutionLor,
	d2F = jet[3], d3F = jet[4],
	bdlinsolver = MatrixBLS())

@test sn_codim2.specialpoint[1].type == :bt
@test sn_codim2.specialpoint[2].type == :zh
@test sn_codim2.specialpoint[3].type == :zh
@test sn_codim2.specialpoint[4].type == :bt

btpt = computeNormalForm(jet..., sn_codim2, 1; nev = 4)
HC = BK.predictor(btpt, Val(:HopfCurve), 0.)
	HC.hopf(0.)

# plot(sn_codim2, vars=(:F, :T), branchlabel = "SN")
# 	_S = LinRange(0., 0.001, 100)
# 	plot!([HC.hopf(s)[1] for s in _S], [HC.hopf(s)[2] for s in _S], label = "Hpred")
# 	# plot!(hp_codim2_1, vars=(:F, :T), branchlabel = "Hopf1")
####################################################################################################
hp_codim2_1, = continuation(jet[1:2]..., br, 2, (@lens _.T), ContinuationPar(opts_br, ds = -0.001, dsmax = 0.02, dsmin = 1e-4, nInversion = 6, saveSolEveryStep = 1, detectBifurcation = 1) ;
	normC = norminf,
	tangentAlgo = BorderedPred(),
	detectCodim2Bifurcation = 2,
	updateMinAugEveryStep = 1,
	startWithEigen = true,
	recordFromSolution = recordFromSolutionLor,
	d2F = jet[3], d3F = jet[4],
	bothside = true,
	# linearAlgo = MatrixBLS(),
	bdlinsolver = MatrixBLS())

@test hp_codim2_1.specialpoint[1].type == :bt
@test hp_codim2_1.specialpoint[2].type == :gh
@test hp_codim2_1.specialpoint[3].type == :hh

# plot(hp_codim2_1, vars=(:F,:T))

computeNormalForm(jet..., hp_codim2_1, 1; nev = 4, verbose=true)
computeNormalForm(jet..., hp_codim2_1, 2; nev = 4, verbose=true)

# hp_codim2_2, = continuation(jet..., sn_codim2, 4, ContinuationPar(opts_br, ds = -0.001, dsmax = 0.02, dsmin = 1e-4, nInversion = 6, detectBifurcation = 0) ;
# 	normC = norminf,
# 	tangentAlgo = BorderedPred(),
# 	detectCodim2Bifurcation = 2,
# 	updateMinAugEveryStep = 1,
# 	startWithEigen = true,
# 	recordFromSolution = recordFromSolutionLor,
# 	# bothside = true,
# 	bdlinsolver = MatrixBLS())

# computeNormalForm(jet..., hp_codim2_2, 1; nev = 4, verbose=true)

# plot(hp_codim2_1, vars=(:F,:T),ylims=(0,0.06), xlims=(1,3))
# plot(hp_codim2_2, vars=(:F,:T),ylims=(-0.06,0.06), xlims=(1,3))
#
#
# plot(sn_codim2, vars=(:F, :T), branchlabel = "SN")
# 	plot!(hp_codim2_1, vars=(:F, :T), branchlabel = "Hopf1")
# 	plot!(hp_codim2_2, vars=(:F, :T), branchlabel = "Hopf2")
# 	ylims!(-0.06,0.09);xlims!(1,3.5)
#
# plot(sn_codim2, vars=(:X, :U), branchlabel = "SN")
# 	plot!(hp_codim2_1, vars=(:X, :U), branchlabel = "Hopf1")
# 	plot!(hp_codim2_2, vars=(:X, :U), branchlabel = "Hopf2")
# 	ylims!(-0.7,0.75);xlims!(0.95,1.5)