function gen_upper_mul_quote(T, N, N2)
    q, qa = initial_quote(N2)
    for i ∈ 1:N2
        push!(qa, :($(Symbol(:B_, i)) = B.data[$i]))
    end
    for i ∈ 1:N
        lti = small_triangle(i)
        for j ∈ 1:i
            C_i = Symbol(:C_, lti + j)
            B_i = Symbol(:B_, lti + j)
            A_i = Symbol(:A_, small_triangle(j) + j)
            push!(qa, :($C_i = $A_i * $B_i) )
            for k ∈ j+1:i
                A_i = Symbol(:A_, small_triangle(k) + j)
                B_i = Symbol(:B_, lti + k)
                push!(qa, :($C_i += $A_i * $B_i) )
            end
        end
    end
    push!(q.args, :($T( ( @ntuple $N2 C ) )))
    q
end

@generated function Base.:*(A::UpperTriangularMatrix{T,N,N2}, B::UpperTriangularMatrix{T,N,N2}) where {T,N,N2}
    gen_upper_mul_quote(UpperTriangularMatrix{T,N,N2}, N, N2)
end

function gen_lower_mul_quote(T, N, N2)
    q, qa = initial_quote(N2)
    for i ∈ 1:N2
        push!(qa, :($(Symbol(:B_, i)) = B.data[$i]))
    end
    for i ∈ 1:N
        lti = small_triangle(i)
        for j ∈ i:N
            ltj = small_triangle(j)
            C_i = Symbol(:C_, ltj + i)
            B_i = Symbol(:B_, lti + i)
            A_i = Symbol(:A_, ltj + i)
            push!(qa, :($C_i = $A_i * $B_i) )
            for k ∈ i+1:j
                A_i = Symbol(:A_, ltj + k)
                B_i = Symbol(:B_, small_triangle(k) + i)
                push!(qa, :($C_i += $A_i * $B_i) )
            end
        end
    end
    push!(q.args, :($T( ( @ntuple $N2 C ) )))
    q
end

@generated function Base.:*(A::LowerTriangularMatrix{T,N,N2}, B::LowerTriangularMatrix{T,N,N2}) where {T,N,N2}
    gen_lower_mul_quote(LowerTriangularMatrix{T,N,N2}, N, N2)
end

function gen_upper_xxt_quote!(qa,T,N,L, A = :A, B = :B, extract = id_symbol)
    for i ∈ 1:N
        lti = small_triangle(i)
        for j ∈ 1:i
            B_i = extract(B, lti + j)
            A1_i = extract(A, lti + j)
            A2_i = extract(A, lti + i)
            push!(qa, :($B_i = $A1_i * $A2_i) )
            for k ∈ i+1:N
                ltk = small_triangle(k)
                A1_i = extract(A, ltk + j)
                A2_i = extract(A, ltk + i)
                push!(qa, :($B_i += $A1_i * $A2_i) )
            end
        end
    end
end

function gen_upper_xtx_quote!(qa,T,N,L)
    for i ∈ 1:N
        lti = small_triangle(i)
        for j ∈ 1:i
            ltj = small_triangle(j)
            B_i = Symbol(:B_, lti + j)
            A1_i = Symbol(:A_, ltj + 1)
            A2_i = Symbol(:A_, lti + 1)
            push!(qa, :($B_i = $A1_i * $A2_i) )
            for k ∈ 2:j
                A1_i = Symbol(:A_, ltj + k)
                A2_i = Symbol(:A_, lti + k)
                push!(qa, :($B_i += $A1_i * $A2_i) )
            end
        end
    end
end

function gen_xxt_quote!(qa,T,M,N,L)
    N2 = big_triangle(M) #number of elements in output
    q, qa = initial_quote(M*N)
    for i ∈ 1:M
        lti = small_triangle(i)
        for j ∈ 1:i
            B_i = Symbol(:B_, lti + j)
            A1_i = Symbol(:A_, j)
            A2_i = Symbol(:A_, i)
            push!(qa, :($B_i = $A1_i * $A2_i) )
            for k ∈ 1:N-1
                Nk = M*k
                A1_i = Symbol(:A_, Nk + j)
                A2_i = Symbol(:A_, Nk + i)
                push!(qa, :($B_i += $A1_i * $A2_i) )
            end
        end
    end
end
function gen_xtx_quote!(qa,T,M,N,L)
    for i ∈ 1:N
        offseti = (i-1)*M
        lti = small_triangle(i)
        for j ∈ 1:i
            offsetj = (j-1)*M
            B_i = Symbol(:B_, lti + j)
            A1_i = Symbol(:A_, offseti + 1)
            A2_i = Symbol(:A_, offsetj + 1)
            push!(qa, :($B_i = $A1_i * $A2_i) )
            for k ∈ 2:M
                A1_i = Symbol(:A_, offseti + k)
                A2_i = Symbol(:A_, offsetj + k)
                push!(qa, :($B_i += $A1_i * $A2_i) )
            end
        end
    end
end

@generated function xxt(A::UpperTriangularMatrix{T,N,L}) where {T,N,L}
    q, qa = initial_quote(L)
    gen_upper_xxt_quote!(qa,T,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$N,$L}( ( @ntuple $L B ) )))
    q
end
@generated function xtx(A::UpperTriangularMatrix{T,N,L}) where {T,N,L}
    q, qa = initial_quote(L)
    gen_upper_xtx_quote!(qa,T,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$N,$L}( ( @ntuple $L B ) )))
    q
end
@generated function xxt(A::Union{SMatrix{M,N,T},StaticRecursiveMatrix{T,M,N}}) where {T,M,N}
    q, qa = initial_quote(M*N)
    L = big_triangle(M)
    gen_xxt_quote!(qa,T,M,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$M,$L}( ( @ntuple $L B ) )))
    q
end
@generated function xtx(A::Union{SMatrix{M,N,T},StaticRecursiveMatrix{T,M,N}}) where {T,M,N}
    q, qa = initial_quote(M*N)
    L = big_triangle(N)
    gen_xtx_quote!(qa,T,M,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$N,$L}( ( @ntuple $L B ) )))
    q
end

xxt(A::LowerTriangularMatrix{T,N,N2}) where {T,N,N2} = xtx(UpperTriangularMatrix{T,N,N2}(A.data))
xtx(A::LowerTriangularMatrix{T,N,N2}) where {T,N,N2} = xxt(UpperTriangularMatrix{T,N,N2}(A.data))

@generated function xxt!(A::UpperTriangularMatrix{T,N,L}) where {T,N,L}
    q, qa = initial_quote(L)
    gen_upper_xxt_quote!(qa,T,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$N,$L}( ( @ntuple $L B ) )))
    q
end
@generated function xtx!(A::UpperTriangularMatrix{T,N,L}) where {T,N,L}
    q, qa = initial_quote(L)
    gen_upper_xtx_quote!(qa,T,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$N,$L}( ( @ntuple $L B ) )))
    q
end
@generated function xxt!(A::Union{MMatrix{M,N,T},RecursiveMatrix{T,M,N}}) where {T,M,N}
    q, qa = initial_quote(M*N)
    L = big_triangle(M)
    gen_xxt_quote!(qa,T,M,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$M,$L}( ( @ntuple $L B ) )))
    q
end
@generated function xtx!(A::Union{MMatrix{M,N,T},RecursiveMatrix{T,M,N}}) where {T,M,N}
    q, qa = initial_quote(M*N)
    L = big_triangle(N)
    gen_xtx_quote!(qa,T,M,N,L)
    push!(q.args, :(SymmetricMatrix{$T,$N,$L}( ( @ntuple $L B ) )))
    q
end

function chol2inv!()

end

# output is an N x N2 SMatrix
function AU_quote!(qa, N, N2)
    for i ∈ 1:N2
        lmi = (i-1)*N
        lti = small_triangle(i)
        for j ∈ 1:N
            C_i = Symbol(:C_, j + lmi)
            A_i = Symbol(:A_, j)
            U_i = Symbol(:U_, lti + 1)
            push!(qa, :($C_i = $A_i * $U_i) )
            for k ∈ 2:i
                A_i = Symbol(:A_, j + (k-1)*N)
                U_i = Symbol(:U_, lti + k)
                push!(qa, :($C_i += $A_i * $U_i) )
            end
        end
    end
end
function UA_quote!(qa, N, N2)
    for i ∈ 1:N2
        lmi = (i-1)*N
        # lti = small_triangle(i)
        for j ∈ 1:N
            # ltj = (j-1)*N
            C_i = Symbol(:C_, j + lmi)
            A_i = Symbol(:A_, j + lmi)
            U_i = Symbol(:U_, big_triangle(j)) # equivalent to j + small_triangle(j)
            push!(qa, :($C_i = $A_i * $U_i) )
            for k ∈ j+1:N
                A_i = Symbol(:A_, k + lmi)
                U_i = Symbol(:U_, j + small_triangle(k))
                push!(qa, :($C_i += $A_i * $U_i) )
            end
        end
    end
end

# Shoot, L is row-major storate.
# Must fix, or write row major extractor? #Writing for SMatrices anyway
function LU_quote!(qa, N)
    for i ∈ 1:N
        lmi = (i-1)*N
        lti = small_triangle(i)
        for j ∈ 1:N
            ltj = small_triangle(j)
            C_i = Symbol(:C_, j + lmi)
            L_i = Symbol(:L_, 1 + ltj)
            U_i = Symbol(:U_, 1 + lti)
            push!(qa, :($C_i = $L_i * $U_i) )
            for k ∈ 2:min(i,j)
                L_i = Symbol(:L_, k + ltj)
                U_i = Symbol(:U_, k + lti)
                push!(qa, :($C_i += $L_i * $U_i) )
            end
        end
    end
end
function UL_quote!(qa, N)
    for i ∈ 1:N
        lmi = (i-1)*N
        lti = small_triangle(i)
        for j ∈ 1:N
            k_start = max(i,j)
            ltk = small_triangle(k_start)
            C_i = Symbol(:C_, j + lmi)
            U_i = Symbol(:U_, ltk + j)
            L_i = Symbol(:L_, ltk + i)
            push!(qa, :($C_i = $U_i * $L_i) )
            for k ∈ k_start+1:N
                ltk = small_triangle(k)
                U_i = Symbol(:U_, ltk + j)
                L_i = Symbol(:L_, ltk + i)
                push!(qa, :($C_i += $U_i * $L_i) )
            end
        end
    end
end
@generated function Base.:*(U::UpperTriangularMatrix{T,N,N2}, L::LowerTriangularMatrix{T,N,N2}) where {T,N,N2}
    # println("Okay...")
    q, qa = initial_quote(N2, :U)
    extract_linear!(qa, N2, :L)
    UL_quote!(qa, N)
    push!(q.args, :( SMatrix{$N,$N}( @ntuple $(N*N) C ) ))
    q
end

function AS_quote!(qa, N, N2)
    for i ∈ 1:N
        lmi = (i-1)*N2
        lti = small_triangle(i)
        for j ∈ 1:N2
            C_i = Symbol(:C_, j + lmi)
            A_i = Symbol(:A_, j)
            S_i = Symbol(:S_, 1 + lti)
            push!(qa, :($C_i = $A_i * $S_i) )
            for k ∈ 2:i #we're in upper triangle of S
                A_i = Symbol(:A_, j + (k-1)*N2)
                S_i = Symbol(:S_, k + lti)
                push!(qa, :($C_i += $A_i * $S_i) )
            end
            for k ∈ 1+i:N
                #we're in mirrored lower triangle
                #thus, k is column, i is row
                A_i = Symbol(:A_, j + (k-1)*N2)
                S_i = Symbol(:S_, i + small_triangle(k))
                push!(qa, :($C_i += $A_i * $S_i) )
            end
        end
    end
end

function gen_AU_quote(N3, N, N2, T)
    Ntotal = N3*N
    q, qa = initial_quote(Ntotal, :A)
    extract_linear!(qa, N2, :U)
    AU_quote!(qa, N3, N)
    push!(q.args, :(StaticArrays.SArray{Tuple{$N3,$N},$T,2,$Ntotal}( ( @ntuple $Ntotal C ))))
    q
end
@generated function Base.:*(A::SMatrix{N3,N,S}, U::UpperTriangularMatrix{T,N,N2}) where {S,T,N,N2,N3}
    gen_AU_quote(N3, N, N2, promote_type(S,T))
end
function gen_LA_quote(N3, N, N2, T)
    Ntotal = N3*N
    q, qa = initial_quote(Ntotal, :L)
    extract_transpose!(qa, N3, N, :A)
    AU_quote!(qa, N3, N)
    push!(q.args, output_transpose(N,N3,T,:C))
    q
end
@generated function Base.:*(L::LowerTriangularMatrix{T,N,N2}, A::SMatrix{N3,N,S}) where {S,T,N,N2,N3}
    gen_LA_quote(N3, N, N2, promote_type(S,T))
end

function gen_UA_quote(N3::Int, N::Int, N2::Int, ::Type{T}) where T
    Ntotal = N3*N
    q, qa = initial_quote(Ntotal, :A)
    extract_linear!(qa, N2, :U)
    UA_quote!(qa, N, N3)
    push!(q.args, :(StaticArrays.SArray{Tuple{$N,$N3},$T,2,$Ntotal}( ( @ntuple $Ntotal C ))))
    q
end
@generated function Base.:*(U::UpperTriangularMatrix{T,N,L}, A::SMatrix{N,L2,S}) where {S,T,N,L,L2}
    gen_UA_quote(L2, N, L, promote_type(S,T))
end

function gen_AL_quote(N3, N, N2, T)
    Ntotal = N3*N
    q, qa = initial_quote(N2, :L)
    extract_transpose!(qa, N3, N, :A)
    UA_quote!(qa, N, N3)
    push!(q.args, output_transpose(N3,N,T,:C))
    q
end
@generated function Base.:*(A::SMatrix{N3,N,S}, L::LowerTriangularMatrix{T,N,N2}) where {S,T,N,N2,N3}
    gen_AL_quote(N3, N, N2, promote_type(S,T))
end

function rank_downdate!()

end

# @generated function Base.:*(A::SMatrix{N3,N}, B::LowerTriangularMatrix{T,N,N2}) where {T,N,N2,N3}
#     gen_lower_mul_quote(LowerTriangularMatrix{T,N,N2}, N, N2)
# end
# @generated function Base.:*(A::LowerTriangularMatrix{T,N,N2}, B::SMatrix{N3,N}) where {T,N,N2,N3}
#     gen_lower_mul_quote(LowerTriangularMatrix{T,N,N2}, N, N2)
# end

function US_quote!()

end
function SU_quote!()

end
function USU_quote!()

end
function LSL_quote!()

end


#=
"A Graph-Free Approach to Data-Flow Analysis" implementation – constant folding propagation example

xref: https://www.semanticscholar.org/paper/A-Graph-Free-Approach-to-Data-Flow-Analysis-Mohnen/5ad8cb6b477793ffb5ec29dde89df6b82dbb6dba?p2df
NOTE: the code is originally adapted from https://github.com/JeffBezanson/dataflow.jl/blob/master/dataflow.jl
=#

# implementation
# ==============

# Part 1: program P
# -----------------

abstract type Exp end

struct Sym <: Exp
    name::Symbol
end

struct Num <: Exp
    val::Int
end

struct Call <: Exp
    head::Sym
    args::Vector{Exp}
end

abstract type Instr end

struct Assign <: Instr
    lhs::Sym
    rhs::Exp
end

struct Goto <: Instr
    label::Int
end

struct GotoIf <: Instr
    label::Int
    cond::Exp
end

const Program = Vector{Instr}

# Part 2: lattice L
# -----------------

# partial order, meet, join, top, bottom, and their identities

import Base: ≤, ==, <, show

abstract type LatticeElement end

struct Const <: LatticeElement
    val::Int
end

struct TopElement <: LatticeElement end
struct BotElement <: LatticeElement end

const ⊤ = TopElement()
const ⊥ = BotElement()

show(io::IO, ::TopElement) = print(io, '⊤')
show(io::IO, ::BotElement) = print(io, '⊥')

≤(x::LatticeElement, y::LatticeElement) = x≡y
≤(::BotElement,      ::TopElement)      = true
≤(::BotElement,      ::LatticeElement)  = true
≤(::LatticeElement,  ::TopElement)      = true

# NOTE: == and < are defined such that future LatticeElements only need to implement ≤
==(x::LatticeElement, y::LatticeElement) = x≤y && y≤x
<(x::LatticeElement,  y::LatticeElement) = x≤y && !(y≤x)

# join
⊔(x::LatticeElement, y::LatticeElement) = x≤y ? y : y≤x ? x : ⊤

# meet
⊓(x::LatticeElement, y::LatticeElement) = x≤y ? x : y≤x ? y : ⊥

# abstract state

# NOTE: the paper (https://api.semanticscholar.org/CorpusID:28519618) uses U+1D56E MATHEMATICAL BOLD FRAKTUR CAPITAL C for this
const AbstractState = Dict{Symbol,LatticeElement}

# extend lattices of abstract values to lattices of mappings of variables to abstract values;
# ⊓ and ⊔ operate pair-wise, and from there we can just rely on the Base implementation for
# dictionary equality comparison

⊔(X::AbstractState, Y::AbstractState) = AbstractState( v => X[v] ⊔ Y[v] for v in keys(X) )
⊓(X::AbstractState, Y::AbstractState) = AbstractState( v => X[v] ⊓ Y[v] for v in keys(X) )

<(X::AbstractState, Y::AbstractState) = X⊓Y==X && X≠Y

# Part 3: abstract semantics ![.!]
# --------------------------------

abstract_eval(x::Num, s::AbstractState) = Const(x.val)

abstract_eval(x::Sym, s::AbstractState) = get(s, x.name, ⊥)

function abstract_eval(x::Call, s::AbstractState)
    f = getfield(@__MODULE__, x.head.name)

    argvals = Int[]
    for arg in x.args
        arg = abstract_eval(arg, s)
        arg === ⊥ && return ⊥ # bail out if any of call arguments is non-constant
        push!(argvals, unwrap_val(arg))
    end

    return Const(f(argvals...))
end

# unwrap our lattice representation into actual Julia value
unwrap_val(x::Num)   = x.val
unwrap_val(x::Const) = x.val

# Part 4: initial state a₀
# ------------------------

a₀ = AbstractState(:x => ⊤, :y => ⊤, :z => ⊤, :r => ⊤)

# algorithm
# ---------

# original, wrong version
function max_fixed_point(prog::Program, a₀::AbstractState, eval)
    n = length(prog)
    init = AbstractState( v => ⊤ for v in keys(a₀) )
    s = [ a₀; [ init for i = 2:n ] ]
    W = BitSet(0:n-1)

    while !isempty(W)
        pc = first(W)
        while pc ≠ n
            delete!(W, pc)
            I = prog[pc+1]
            new = s[pc+1]
            if isa(I, Assign)
                # for an assignment, outgoing value is different from incoming
                new = copy(new)
                new[I.lhs.name] = eval(I.rhs, new)
            end

            if isa(I, Goto)
                pc´ = I.label
            else
                pc´ = pc+1
                if isa(I, GotoIf)
                    l = I.label
                    if new < s[l+1]
                        push!(W, l)
                        s[l+1] = new
                    end
                end
            end
            if pc´≤n-1 && new < s[pc´+1]
                s[pc´+1] = new
                pc = pc´
            else
                pc = n
            end
        end
    end

    return s
end

# NOTE: in this problem, we make sure that states will always move to _lower_ position in lattice, so
# - initialize states with `⊤`
# - we use `⊓` (meet) operator to update states,
# - and the condition we use to check whether or not the statement makes a change is `new ≠ prev`
function max_fixed_point(prog::Program, a₀::AbstractState, eval)
    n = length(prog)
    init = AbstractState( v => ⊤ for v in keys(a₀) )
    s = [ a₀; [ init for i = 2:n ] ]
    W = BitSet(0:n-1)

    while !isempty(W)
        pc = first(W)
        while pc ≠ n
            delete!(W, pc)
            I = prog[pc+1]
            new = s[pc+1]
            if isa(I, Assign)
                # for an assignment, outgoing value is different from incoming
                new = copy(new)
                new[I.lhs.name] = eval(I.rhs, new)
            end

            if isa(I, Goto)
                pc´ = I.label
            else
                pc´ = pc+1
                if isa(I, GotoIf)
                    l = I.label
                    if new ≠ s[l+1]
                        push!(W, l)
                        s[l+1] = new ⊓ s[l+1]
                    end
                    # # alternatively,
                    # new = new ⊓ s[l]
                    # if new < s[l]
                    #     push!(W, l)
                    #     s[l] = new
                    # end
                end
            end
            if pc´≤n-1 && new ≠ s[pc´+1]
                s[pc´+1] = new ⊓ s[pc´+1]
                pc = pc´
            else
                pc = n
            end
            # # alternatively,
            # pc´≤n && (new = new ⊓ s[pc´])
            # if pc´≤n && new < s[pc´]
            #     s[pc´] = new
            #     pc = pc´
            # else
            #     pc = n+1
            # end
        end
    end

    return s
end

# example problem

prog0 = [Assign(Sym(:x), Num(1)),                              # I₀
         Assign(Sym(:y), Num(2)),                              # I₁
         Assign(Sym(:z), Num(3)),                              # I₂
         Goto(8),                                              # I₃
         Assign(Sym(:r), Call(Sym(:(+)), [Sym(:y), Sym(:z)])), # I₄
         GotoIf(7, Call(Sym(:(≤)), [Sym(:x), Sym(:z)])),       # I₅
         Assign(Sym(:r), Call(Sym(:(+)), [Sym(:z), Sym(:y)])), # I₆
         Assign(Sym(:x), Call(Sym(:(+)), [Sym(:x), Num(1)])),  # I₇
         GotoIf(4, Call(Sym(:(<)), [Sym(:x), Num(10)])),       # I₈
         ]::Program

using MacroTools

macro prog(blk)
    Instr[Instr(x) for x in filter(!islnn, blk.args)]::Program
end

function Instr(x)
    if @capture(x, lhs_ = rhs_)               # => Assign
        Assign(Instr(lhs), Instr(rhs))
    elseif @capture(x, @goto label_)          # => Goto
        Goto(label)
    elseif @capture(x, cond_ && @goto label_) # => GotoIf
        GotoIf(label, Instr(cond))
    elseif @capture(x, f_(args__))            # => Call
        Call(Instr(f), Instr.(args))
    elseif isa(x, Symbol)                     # => Sym
        Sym(x)
    elseif isa(x, Int)                        # => Num
        Num(x)
    else
        error("invalid expression: $(x)")
    end
end

islnn(@nospecialize(_)) = false
islnn(::LineNumberNode) = true

prog0 = @prog begin
    x = 1             # I₀
    y = 2             # I₁
    z = 3             # I₂
    @goto 8           # I₃
    r = y + z         # I₄
    x ≤ z && @goto 7  # I₅
    r = z + y         # I₆
    x = x + 1         # I₇
    x < 10 && @goto 4 # I₀
end

max_fixed_point(prog0, a₀, abstract_eval) # The solution contains the `:r => Const(5)`, which is not found in the program

# comparison with Julia's inference
# =================================

# generate valid Julia code from the "`Instr` syntax"
macro prog′(blk)
    prog′ = Expr(:block)
    bns = [gensym(Symbol(:instruction, i-1)) for i in 1:length(blk.args)] # pre-generate labels for all instructions

    for (i,x) in enumerate(filter(!islnn, blk.args))
        x = MacroTools.postwalk(x) do x
            return if @capture(x, @goto label_)
                Expr(:symbolicgoto, bns[label+1]) # fix `@goto i` into valid goto syntax
            else
                x
            end
        end

        push!(prog′.args, Expr(:block, Expr(:symboliclabel, bns[i]), x)) # label this statement
    end

    return prog′
end

@macroexpand @prog′ begin
    x = 1             # I₀
    y = 2             # I₁
    z = 3             # I₂
    @goto 8           # I₃
    r = y + z         # I₄
    x ≤ z && @goto 7  # I₅
    r = z + y         # I₆
    x = x + 1         # I₇
    x < 10 && @goto 4 # I₈
end

code_typed(; optimize = false) do
    @prog′ begin
        x = 1             # I₀
        y = 2             # I₁
        z = 3             # I₂
        @goto 8           # I₃
        r = y + z         # I₄
        x ≤ z && @goto 7  # I₅
        r = z + y         # I₆
        x = x + 1         # I₇
        x < 10 && @goto 4 # I₈

        x, y, z, r # to check the result of abstract interpretation
    end
end |> first

@prog′ begin
    x = 1             # I₀
    y = 2             # I₁
    z = 3             # I₂
    @goto 8           # I₃
    r = y + z         # I₄
    x ≤ z && @goto 7  # I₅
    r = z + y         # I₆
    x = x + 1         # I₇
    x < 10 && @goto 4 # I₈

    x, y, z, r # to check the result of actual execution
end


# %% Standard map island shrinking
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot
grid = 0.0:0.5:2π
sm = Systems.standardmap()
ks = [0.1, 1.0, 4.0]

fig, axs = subplots(1,3; sharey = true)

for i in 1:3
    set_parameter!(sm, 1, ks[i])
    ax = axs[i]
    for θ in grid, p in grid
        u0 = SVector(θ, p)
        tr = trajectory(sm, 1000, u0)
        c = DynamicalSystems.lyapunovspectrum(sm, 4000, 1; u0 = u0)[1] > 0.01 ? "C0" : "C2"
        ax.plot(columns(tr)...; ls = "None", marker = "o", ms = 0.2, c = c)
    end
    ax.set_xlim(0, 2π)
    ax.set_ylim(0, 2π)
    ax.text(4.5, 5.5, "\$k=$(ks[i])\$"; bbox = bbox)
    ax.set_xticks(0:2:6)
end

axs[1].set_ylabel("\$p\$", labelpad = -5)
axs[2].set_xlabel("\$\\theta\$", labelpad = -15)
fig.tight_layout(pad=0.3)
# fig.subplots_adjust(left = 0.05, bottom = 0.14, right = 0.95, top = 0.97, wspace = 0.06)
wsave(plotsdir("8","kam_sm"), fig)


using OrdinaryDiffEq: phi, phi, phiv, phiv_timestep, expv, expv_timestep, arnoldi, getH, getV
using OrdinaryDiffEq: WOperator, set_gamma!, calc_W!
using OrdinaryDiffEq, LinearAlgebra, SparseArrays, Random, Test, DiffEqOperators

@testset "Exponential Utilities" begin
  # Scalar phi
  K = 4
  z = 0.1
  P = fill(0., K+1); P[1] = exp(z)
  for i = 1:K
    P[i+1] = (P[i] - 1/factorial(i-1))/z
  end
  @test phi(z, K) ≈ P

  # Matrix phi
  A = [0.1 0.2; 0.3 0.4]
  P = Vector{Matrix{Float64}}(undef, K+1); P[1] = exp(A)
  for i = 1:K
    P[i+1] = (P[i] - 1/factorial(i-1)*I) / A
  end
  @test phi(A, K) ≈ P

  # Krylov
  n = 20; m = 5
  Random.seed!(0)
  A = randn(n, n)
  t = 1e-2
  b = randn(n)
  @test exp(t * A) * b ≈ expv(t, A, b; m=m)
  P = phi(t * A, K)
  W = fill(0., n, K+1)
  for i = 1:K+1
    W[:,i] = P[i] * b
  end
  Ks = arnoldi(A, b; m=m)
  W_approx = phiv(t, Ks, K)
  @test W ≈ W_approx

  # Happy-breakdown in Krylov
  v = normalize(randn(n))
  A = v * v' # A is Idempotent
  Ks = arnoldi(A, b)
  @test Ks.m == 2

  # Arnoldi vs Lanczos
  A = Hermitian(randn(n, n))
  Aperm = A + 1e-10 * randn(n, n) # no longer Hermitian
  w = expv(t, A, b; m=m)
  wperm = expv(t, Aperm, b; m=m)
  @test w ≈ wperm

  # Internal time-stepping for Krylov (with adaptation)
  n = 100
  K = 4
  t = 5.0
  tol = 1e-7
  A = spdiagm(-1=>ones(n-1), 0=>-2*ones(n), 1=>ones(n-1))
  B = randn(n, K+1)
  Phi_half = phi(t/2 * A, K)
  Phi = phi(t * A, K)
  uhalf_exact = sum((t/2)^i * Phi_half[i+1] * B[:,i+1] for i = 0:K)
  u_exact = sum(t^i * Phi[i+1] * B[:,i+1] for i = 0:K)
  U = phiv_timestep([t/2, t], A, B; adaptive=true, tol=tol)
  @test norm(U[:,1] - uhalf_exact) / norm(uhalf_exact) < tol
  @test norm(U[:,2] - u_exact) / norm(u_exact) < tol
  # p = 0 special case (expv_timestep)
  u_exact = Phi[1] * B[:, 1]
  u = expv_timestep(t, A, B[:, 1]; adaptive=true, tol=tol)
  @test norm(u - u_exact) / norm(u_exact) < tol
end

@testset "Derivative Utilities" begin
  @testset "WOperator" begin
    Random.seed!(0); y = zeros(2); b = rand(2)
    mm = I; _J = rand(2,2)
    W = WOperator(mm, 1.0, DiffEqArrayOperator(_J))
    set_gamma!(W, 2.0)
    _W = mm - 2.0 * _J
    @test convert(AbstractMatrix,W) ≈ _W
    @test W * b ≈ _W * b
    mul!(y, W, b); @test y ≈ _W * b
  end

  @testset "calc_W!" begin
    A = [-1.0 0.0; 0.0 -0.5]; mm = [2.0 0.0; 0.0 1.0]
    u0 = [1.0, 1.0]; tmp = zeros(2)
    tspan = (0.0,1.0); dt = 0.01; dtgamma = 0.5dt
    concrete_W = mm - dtgamma * A
    concrete_Wt = mm/dtgamma - A

    # Out-of-place
    fun = ODEFunction((u,p,t) -> A*u;
                      mass_matrix=mm,
                      jac=(u,p,t) -> A)
    integrator = init(ODEProblem(fun,u0,tspan), ImplicitEuler(); adaptive=false, dt=dt)
    W = calc_W!(integrator, integrator.cache, dtgamma, false)
    @test convert(AbstractMatrix, W) ≈ concrete_W
    @test W \ u0 ≈ concrete_W \ u0

    # In-place
    fun = ODEFunction((du,u,p,t) -> mul!(du,A,u);
                      mass_matrix=mm,
                      jac_prototype=DiffEqArrayOperator(A))
    integrator = init(ODEProblem(fun,u0,tspan), ImplicitEuler(); adaptive=false, dt=dt)
    calc_W!(integrator, integrator.cache, dtgamma, false)
    @test convert(AbstractMatrix, integrator.cache.W) ≈ concrete_W
    ldiv!(tmp, lu!(integrator.cache.W), u0); @test tmp ≈ concrete_W \ u0
  end

  @testset "Implicit solver with lazy W" begin
    A = sparse([-1.0 0.0; 0.0 -0.5])
    mm = sparse([2.0 0.0; 0.0 1.0])
    u0 = [1.0, 1.0]; tspan = (0.0,1.0)

    _f = (u,p,t) -> t*(A*u); _f_ip = (du,u,p,t) -> lmul!(t,mul!(du,A,u))
    fun1 = ODEFunction(_f; mass_matrix=mm)
    fun2 = ODEFunction(_f; mass_matrix=mm, jac=(u,p,t) -> t*A)
    fun1_ip = ODEFunction(_f_ip; mass_matrix=mm)
    fun2_ip = ODEFunction(_f_ip; mass_matrix=mm,
    jac_prototype=DiffEqArrayOperator(similar(A); update_func=(J,u,p,t) -> (J .= t .* A; J)))

    for Alg in [ImplicitEuler, Rosenbrock23]
      println(Alg)
      sol1 = solve(ODEProblem(fun1,u0,tspan), Alg(); adaptive=false, dt=0.01)
      sol2 = solve(ODEProblem(fun2,u0,tspan), Alg(); adaptive=false, dt=0.01)
      @test sol1(1.0) ≈ sol2(1.0)

      sol1_ip = solve(ODEProblem(fun1_ip,u0,tspan), Alg(); adaptive=false, dt=0.01)
      sol2_ip = solve(ODEProblem(fun2_ip,u0,tspan), Alg(linsolve=LinSolveFactorize(lu)); adaptive=false, dt=0.01)
      @test sol1_ip(1.0) ≈ sol2_ip(1.0)
    end
  end
end


using Gadfly
using Distributions

f(u) = exp(-u^2/2)/√(2pi)

const curve_plot =
    plot(f, -5, +5,
        Theme(line_width=2Gadfly.mm,
              major_label_font_size=40Gadfly.px,
              minor_label_font_size=25Gadfly.px)
    )

const u1=Uniform(-5.0,+5.0)
const u2=Uniform(0.0, 0.5)

function simulate_pt()
    x=rand(u1)
    y=rand(u2)
    y<f(x)
end

function simulate(num, running, current_approx)
     push!(running, true)
     hits=0
     for (i in 1:num)
        hits = hits + simulate_pt()
        if (i % 10000) == 0
            push!(current_approx, (i, hits/i*(0.5*10) ))
        end
    end
    estimate = hits/num*(0.5*10)
    push!(running, false)
end

function run_simulate(N, running, current_approx)
    println("Starting Simulations for N=$N")
    @async simulate(N, running, current_approx)
end

function main(window)
    push!(window.assets, "tex")
    push!(window.assets, "widgets")

    N = Signal(10000)
    btn = Signal(Any, leftbutton)

    current_approx = Signal((0, 1.0)) # current approximation
    buffered_approx = foldp((Any[0],Any[1.0]), current_approx) do prev, current
        idx, val = current
        push!(prev[1], idx), push!(prev[2], val)
    end

    running = Signal(false)
    result = map(sampleon(btn, N); typ=Any, init=nothing) do n
        run_simulate(n, running, current_approx)
    end

    vbox(
        h1("Interactive Simulations"),
        hbox(
            vbox(
                 md"""We want to estimate the following integral using a
                 *Monte Carlo* Simulation""",
                 tex("\\int_{-5}^{5} \\frac{exp(-u^2/2)}{\\sqrt{2pi}}du" , block=true),
            ),
            curve_plot
         ) |> packacross(center),
         map(buffered_approx, running) do approx, r
             xs, ys = approx
             vbox(
                 hbox("Number of runs", slider(10^3:10^3:5*10^6) >>> N, hskip(2em), hbox(button("Start", raised=true, disabled=r) >>> btn)) |> packacross(center),
                 hbox("Current value: ", hskip(1em), @sprintf("%2.4f", ys[end]) |> emph, hskip(1em)," at iteration :", hskip(1em), string(xs[end])),
                 plot(x=xs, y=ys, Geom.line) |> Escher.pad(2em)
             )
         end
    ) |> Escher.pad(2em)
end


# by lazarusA # HIDE
# using GLMakie # HIDE
using CairoMakie, ColorSchemes
let
    xs = 0:0.01:1
    p = -10:1:10
    p= filter(x->x != 0,collect(p))
    ys = zeros(length(xs), length(p))
    for (indx,i) in enumerate(p)
        if i <= -1
            ys[:,indx]= xs .^(1/abs(i))
        elseif i>= 1
            ys[:,indx]= xs .^i
        end
    end
    cbarPal = :thermal
    cmap = get(colorschemes[cbarPal], LinRange(0,1,length(p)))

    fig = Figure(resolution = (700, 450), font =:sans, fontsize = 18)
    ax = Axis(fig, aspect = 1, xlabel = "x", ylabel = "xᵖ")
    lins = [lines!(xs, ys[:,v], color = cmap[v]) for v in 1:length(p)]
    leg = Legend(fig, lins, string.(p), "p", nbanks = 2, labelsize = 12,
     valign = :center)
    fig[1, 1] = ax
    fig[1, 2] = leg
    fig
    #save("MulticoloredLinesLegOut.png", fig, px_per_unit = 2)
    save(joinpath(@__DIR__, "output", "MulticoloredLinesLegOut.png"), fig, px_per_unit = 2) # HIDE
end
using Pkg # HIDE
Pkg.status(["CairoMakie", "ColorSchemes"]) # HIDE


if !(VERSION.major==0&&VERSION.minor<=6)
	using Test
else
	using Base.Test
end
using FinancialToolbox

#Test Parameters
testToll=1e-14;

print_colored("Starting Standard Test\n",:green)
spot=10;K=10;r=0.02;T=2.0;sigma=0.2;d=0.01;
assert_(value,toll)=@test abs(value)<toll
#EuropeanCall Option
PriceCall=blsprice(spot,K,r,T,sigma,d);
PriceCallBlack=blkprice(spot,K,r,T,sigma);
DeltaCall=blsdelta(spot,K,r,T,sigma,d);
ThetaCall=blstheta(spot,K,r,T,sigma,d);
RhoCall=blsrho(spot,K,r,T,sigma,d);
VannaCall=blsvanna(spot,K,r,T,sigma,d);
PsiCall=blspsi(spot,K,r,T,sigma,d);
LambdaCall=blslambda(spot,K,r,T,sigma,d);
SigmaCall=blsimpv(spot, K, r, T, PriceCall, d);
SigmaCallBlack=blkimpv(spot, K, r, T, PriceCallBlack);

#EuropeanPut Option
PricePut=blsprice(spot,K,r,T,sigma,d,false);
PricePutBlack=blkprice(spot,K,r,T,sigma,false);
DeltaPut=blsdelta(spot,K,r,T,sigma,d,false);
ThetaPut=blstheta(spot,K,r,T,sigma,d,false);
VannaPut=blsvanna(spot,K,r,T,sigma,d,false);
PsiPut=blspsi(spot,K,r,T,sigma,d,false);
LambdaPut=blslambda(spot,K,r,T,sigma,d,false);
RhoPut=blsrho(spot,K,r,T,sigma,d,false);
SigmaPut=blsimpv(spot,K,r,T,PricePut,d,false);
SigmaPutBlack=blkimpv(spot,K,r,T,PricePutBlack,false);

#Equals for both Options
Gamma=blsgamma(spot,K,r,T,sigma,d);
Vega=blsvega(spot,K,r,T,sigma,d);

print_colored("---  European Call: Price and Sensitivities\n",:yellow)
#Standard Test European Call Option
print_colored("-----Testing Price\n",:blue);
assert_(PriceCall-1.191201316999582,testToll)
print_colored("-----Testing Black Price\n",:blue);
assert_(PriceCallBlack-1.080531820066428,testToll)
print_colored("-----Testing Delta\n",:blue);
assert_(DeltaCall-0.572434050810368,testToll)
print_colored("-----Testing Theta\n",:blue);
assert_(ThetaCall+0.303776337550247,testToll)
print_colored("-----Testing Rho\n",:blue);
assert_(RhoCall-9.066278382208203,testToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(LambdaCall-4.805518955034612,testToll)
print_colored("-----Testing Implied Volatility\n",:blue);
assert_(SigmaCall-0.2,testToll)
print_colored("-----Testing Implied Volatility Black\n",:blue);
assert_(SigmaCallBlack-0.2,testToll)
#Low Xtol for blsimpv
tol_low=1e-16;
tol_high=1e-3;
SigmaLowXTol=blsimpv(spot, K, r, T, PriceCall, d,true,tol_low,tol_high);
print_colored("-----Testing Implied Volatility Low X tol, High Y tol\n",:blue);
@test((abs(blsprice(spot,K,r,T,SigmaLowXTol,d,true)-PriceCall)<tol_high)||abs(SigmaLowXTol-sigma)<tol_low)
#Low Ytol for blsimpv
SigmaLowYTol=blsimpv(spot, K, r, T, PriceCall, d,true,tol_high,tol_low);
print_colored("-----Testing Implied Volatility Low Y tol, High X tol\n",:blue);
@test((abs(blsprice(spot,K,r,T,SigmaLowYTol,d,true)-PriceCall)<tol_low)||abs(SigmaLowYTol-sigma)<tol_high)

print_colored("---  European Put: Price and Sensitivities\n",:yellow)
#Standard Test European Put Option
print_colored("-----Testing Price\n",:blue);
assert_(PricePut-0.997108975455260,testToll)
print_colored("-----Testing Price Black\n",:blue);
assert_(PricePutBlack-1.080531820066428,testToll)
print_colored("-----Testing Delta\n",:blue);
assert_(DeltaPut+0.407764622496387,testToll)
print_colored("-----Testing Theta\n",:blue);
assert_(ThetaPut+0.209638317050458,testToll)
print_colored("-----Testing Rho\n",:blue);
assert_(RhoPut+10.149510400838260,testToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(LambdaPut+4.089468980160465,testToll)
print_colored("-----Testing Implied Volatility\n",:blue);
assert_(SigmaPut-0.2,testToll)
print_colored("-----Testing Implied Volatility Black\n",:blue);
assert_(SigmaPutBlack-0.2,testToll)

#Standard Test for Common Sensitivities
print_colored("-----Testing Gamma\n",:blue);
assert_(Gamma-0.135178479404601,testToll)
print_colored("-----Testing Vega\n",:blue);
assert_(Vega-5.407139176184034,testToll)
print_colored("Standard Test Passed\n",:green)
println("")

#TEST OF INPUT VALIDATION
print_colored("Starting Input Validation Test Real\n",:magenta)
print_colored("----Testing Negative Spot Price\n",:cyan)
@test_throws(ErrorException, blsprice(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blkprice(-spot,K,r,T,sigma))
@test_throws(ErrorException, blsdelta(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blstheta(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blsrho(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blsvega(-spot,K,r,T,sigma,d));
@test_throws(ErrorException, blspsi(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blslambda(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(-spot,K,r,T,sigma,d))
@test_throws(ErrorException, blsimpv(-spot,K,r,T,PriceCall,d))
@test_throws(ErrorException, blkimpv(-spot,K,r,T,PriceCall))

print_colored("----Testing Negative Strike Price\n",:cyan)
@test_throws(ErrorException, blsprice(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blkprice(spot,-K,r,T,sigma))
@test_throws(ErrorException, blsdelta(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blstheta(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blsrho(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blsvega(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blspsi(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blslambda(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(spot,-K,r,T,sigma,d))
@test_throws(ErrorException, blsimpv(spot,-K,r,T,PriceCall,d))
@test_throws(ErrorException, blkimpv(spot,-K,r,T,PriceCall))

print_colored("----Testing Negative Time to Maturity\n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blkprice(spot,K,r,-T,sigma))
@test_throws(ErrorException, blsdelta(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blsgamma(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blstheta(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blsrho(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blsvega(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blspsi(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blslambda(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blsvanna(spot,K,r,-T,sigma,d))
@test_throws(ErrorException, blsimpv(spot,K,r,-T,PriceCall,d))
@test_throws(ErrorException, blkimpv(spot,K,r,-T,PriceCall))

print_colored("----Testing Negative Volatility\n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blkprice(spot,K,r,T,-sigma))
@test_throws(ErrorException, blsdelta(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blsgamma(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blstheta(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blsrho(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blspsi(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blslambda(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blsvanna(spot,K,r,T,-sigma,d))
@test_throws(ErrorException, blsvega(spot,K,r,T,-sigma,d))
print_colored("----Testing Negative Option Price\n",:cyan)
@test_throws(ErrorException, blsimpv(spot,K,r,T,-PriceCall,d))
@test_throws(ErrorException, blkimpv(spot,K,r,T,-PriceCall))

print_colored("----Testing Negative Tollerance\n",:cyan)
@test_throws(ErrorException, blsimpv(spot,K,r,T,PriceCall,d,true,-1e-12,1e-12))
@test_throws(ErrorException, blkimpv(spot,K,r,T,PriceCall,true,-1e-12,1e-12))
@test_throws(ErrorException, blsimpv(spot,K,r,T,PriceCall,d,true,1e-12,-1e-12))
@test_throws(ErrorException, blkimpv(spot,K,r,T,PriceCall,true,1e-12,-1e-12))

#Too low tollerance
@test_throws(ErrorException, blsimpv(spot,K,r,T,PriceCall,d,true,0.0,0.0))
@test_throws(ErrorException, blkimpv(spot,K,r,T,PriceCall,true,0.0,0.0))

print_colored("Real Input Validation Test Passed\n",:magenta)

#End of the Test


# fbp_ramp.jl

export fbp_ramp
export ramp_flat
export ramp_arc

using MIRT: SinoFanFlat, SinoFanArc

"""
    h, nn = fbp_ramp(how, n, ds, dsd)
'ramp-like' filters for parallel-beam and fan-beam FBP reconstruction.
This sampled band-limited approach avoids the aliasing that would be
caused by sampling the ramp directly in the frequency domain.

in
- `how::Symbol`
  `:arc` (3rd generation CT) or `:flat` (for parallel too)
- `n::Int` : # of samples (must be even)
- `ds::RealU` : sample spacing (in distance units, e.g., cm)
- `dsd::RealU` : source-to-detector distance, for `:arc` case

out
- `h::AbstractMatrix{<:Number}` : samples of band-limited ramp filter

"""
fbp_ramp(::SinoPar, n::Int, ds::RealU, dsd::RealU) = ramp_flat(n, ds)
fbp_ramp(::SinoFanFlat, n::Int, ds::RealU, dsd::RealU) = ramp_flat(n, ds)
fbp_ramp(::SinoFanArc, n::Int, ds::RealU, dsd::RealU) = ramp_arc(n, ds, dsd)

function ramp_arc(n::Int, ds::RealU, dsd::RealU)
    isodd(n) && throw("n must be even")

    if n/2 * ds / dsd > 0.9 * π/2
        println("angle is $(rad2deg(n/2 * ds / dsd)) degrees: too large")
        throw("physically impossible arc geometry")
    end

    nn = -(n÷2):(n÷2-1)
    h = zeros(size(nn))
    h[nn .== 0] .= 1 / (4 * abs2(ds))
    odd = isodd.(nn)
    h[odd] = @.( -1 / abs2((pi * dsd * sin(nn[odd] * ds / dsd))))

    return h, nn
end

function ramp_flat(n::Int, ds::RealU)

    isodd(n) && throw("n must be even")

    nn = -(n÷2):(n÷2-1)
    h = zeros(size(nn))
    h[n÷2+1] = 1 / 4
    odd = isodd.(nn)
    h[odd] .= -1 ./ abs2.(π * nn[odd])
    h ./= abs2(ds)

    return h, nn
end


### A symmetric tridiagonal matrix that supports appending onto the bottom
### right corner of the matrix and keeping track of eigenvalues of the matrix.
### * Support dynamic LU decomposition in real time, without pivoting.
### * Fast characteristic polynomial computations for shifted system of this
###   matrix.

mutable struct DynamicSymTridiagonal{T<:AbstractFloat}
    # Parameters groupd 1
    alphas::Vector{T}       # Diaognal
    betas::Vector{T}        # Lower & upper Diaognal
    L::Vector{T}            # The lower diagonal of the unit-bidiagonal matrix L
    U::Vector{T}            # Lower diagonal of the upper bi-diagonal matrix U
    k::Int64                # Size of the matrix.

    # Parameters group 2
    last_update::Int64          # last iteration where eigenvalues are updated for this matrix.
    thetas::Vector{T}           # Ritz values.
    converged::Vector{Bool}     # Indicates convergence for matching index ritz value from last iteration.
    converged_this_step::Vector{T}
    velocity::Vector{T}
    velocity_tol::Float64

    function DynamicSymTridiagonal{T}(
        alpha::T;
        velocity_tol::S=1e-10
    ) where {T<:AbstractFloat, S<:AbstractFloat}
        this = new{T}()
        this.alphas = Vector{T}()
        push!(this.alphas, alpha)
        this.betas = Vector{T}()
        this.L = Vector{T}()
        this.U = Vector{T}()
        push!(this.U, alpha)
        this.k = 1
        this.last_update = 1
        this.thetas = Vector{T}(); push!(this.thetas, alpha)
        this.converged = Vector{T}(); push!(this.converged, false)
        this.velocity = Vector{T}();push!(this.velocity, Inf)
        this.velocity_tol = velocity_tol # <- Default value.
        this.converged_this_step = Vector{T}()

    return this end

    function DynamicSymTridiagonal(alpha::Number)
        return DynamicSymTridiagonal{Float64}(convert(Float64, alpha))
    end
end

"""
    Opverting the opertor (alpha,beta) to append new element to the matrix.
    append an alpha and a beta (diagonal and sub & super diagonal) elemnt to
    the current sym diagonal matrix.
"""
function (this::DynamicSymTridiagonal{T})(alpha::T, beta::T) where {T<:AbstractFloat}
    push!(this.alphas, alpha)
    push!(this.betas, beta)
    push!(this.U, alpha - beta^2/this.U[end])
    push!(this.L, beta/this.U[end - 1])
    this.k += 1
return this end

"""
    Evaluate for an instance of the symmetric tridiagonal symmetric matrix for
    a given shift value of x, which evalutes: det(A - xI) for the instance matrix.

    * It will return Nan if the polynomial exploded.
"""
function CharPolyShifted(this::DynamicSymTridiagonal{T}, x::Number=0.0) where {T<:AbstractFloat}
    pPrevious = 1
    pNow = this.alphas[1] - x
    for j in 2:this.k
        pNew = (this.alphas[j] - x)*pNow - (this.betas[j - 1]^2)*pPrevious
        pPrevious = pNow
        pNow = pNew
    end
return pNow end

"""
    Evalute the derivetive of the characteristic polynoial for a shifted quantity x.
    det(T - xI)
"""
function CharacteristicPolyDerivative(
    this::DynamicSymTridiagonal{T},
    x::Number
) where {T<:AbstractFloat}
    error("I haven't implement it yet. ")
return end

"""
    Locate the eigenvalue given only the left bound on the eigenvalue.
    * Eigenvalue CANNOT be on the boundary that passed in, they might cause
    sign problem violating the bisection pre-conditions.
"""
function EigenValueLocate(
    this::DynamicSymTridiagonal{T},
    left_bound::T,
    right_bound::T
) where {T<:AbstractFloat}
    @assert !(isnan(left_bound) || isnan(right_bound)) "Any of the bound for "*
    "the eigenvalue for dynamic tridiagonal symmetric matrix cannot be nan. "
    @assert !(isinf(left_bound) && isinf(right_bound)) "Both left and "*
    "right_bound cannot be inf. "
    @assert left_bound < right_bound "The left boundary should be less than the"*
    "right boundary, this is by the definition of interval. "

    # perform search when one of the bound is infinity.
    P(x) = CharPolyShifted(this, x)
    if isinf(left_bound) || isinf(right_bound)
        stepSize = isinf(left_bound) ? -1e-8 : 1e-8
        if stepSize > 0
            while sign(P(left_bound + stepSize)) == sign(P(left_bound))
                    stepSize += 2*stepSize
            end
        else
            while sign(P(right_bound + stepSize)) == sign(P(right_bound))
                stepSize += 2*stepSize
            end
        end
        if isinf(left_bound)
            left_bound = right_bound + stepSize
        else
            right_bound = left_bound + stepSize
        end
    end

    # test preconditions for bisection
    pL = P(left_bound); pR = P(right_bound)
    @assert sign(pL) != sign(pR) "Bisection search error;"*
    "the left and right bondary turns out to have the same sign, please check the "*
    "caller routine which is responsible. the left and right bounds are"*
    ": $(left_bound), $(right_bound), function values is: $(pL), $(pR); "*
    "please consider adjusing the velocity tolerance."
    @assert !isnan(pL) && !isnan(pR) "Polynomial exploded with Nan on one of "*
    "the boundaries OR both boundaries after expontneial probing "*
    "(or without the probing). Please refine search interval. "

    # perform bisection search to locate a root.
    midX = (left_bound + right_bound)/2
    midP = P(midX)
    itrLimit = 100
    while midP != 0 &&
        itrLimit > 0 &&
        right_bound != midX &&
        left_bound != midX

        if sign(midP) == sign(P(left_bound))
            left_bound = midX
        else
            right_bound = midX
        end
        itrLimit -= 1
        midX = (left_bound + right_bound)/2
        if right_bound - left_bound < this.velocity_tol
            break
        end
        midP = P(midX)
        # println("Bisection: [$left_bound, $right_bound]")
    end

return midX end

"""
    Update the eigenvalues based on all the eigenvalues
    from the previous sub-matrix.
        * If current eigenvalues are not from the previous iterations, this
        will give an error, it can only compute eigenvalues for this iteration
        based on the eigenvalues from the previous iterations.
"""
function EigenvaluesUpdate(this::DynamicSymTridiagonal{T}) where {T <: AbstractFloat}
    if this.k == 1  # current matrix is 1 x 1.
        EstablishEigenSystem(this)
        this.last_update = this.k
        return
    end

    if this.k == this.last_update  # eigen system already updated.
        return
    elseif this.k - this.last_update == 1
        # update the eigensystem using the interlace properties and search routine.
        if all(this.converged)
            # method will re-establish eigen system.
            EstablishEigenSystem(this)
            return
        end
        v = this.thetas
        c = this.converged
        k = length(v) + 1
        δ = 0
        push!(v, convert(T, Inf64))
        pushfirst!(v, convert(T, -Inf64))
        newThetas = Array{T}(undef, k)
        newConverged = Array{Bool}(undef, k)
        Δ = Array{T}(undef, k)
        empty!(this.converged_this_step)
        for II in 1: k
            if II < min(ceil(k/2), floor(k/2) + 1)
                if c[II]
                    newThetas[II] = v[II + 1]
                    newConverged[II] = true
                else
                    newThetas[II] = EigenValueLocate(this, v[II] + δ, v[II + 1] - δ)
                    newConverged[II] = abs(newThetas[II] - v[II + 1]) <= this.velocity_tol
                    if II != 1 && II != k
                        newConverged[II] = newConverged[II] && c[II - 1]
                    end
                    if newConverged[II]
                        push!(this.converged_this_step, newThetas[II])
                    end
                end
            elseif II > max(ceil(k/2), floor(k/2) + 1)
                if c[II - 1]
                    newThetas[II] = v[II]
                    newConverged[II] = true
                else
                    newThetas[II] = EigenValueLocate(this, v[II] + δ, v[II + 1] - δ)
                    newConverged[II] = abs(newThetas[II] - v[II]) <= this.velocity_tol
                    if II != 1 && II != k
                        newConverged[II] = newConverged[II] && c[II]
                    end
                    if newConverged[II]
                        push!(this.converged_this_step, newThetas[II])
                    end
                end
            else
                newThetas[II] = EigenValueLocate(this, v[II] + δ, v[II + 1] - δ)
                newConverged[II] = false
            end

        end
        # Velocity update
        for II in 1:k
            if II < min(ceil(k/2), floor(k/2) + 1)
                Δ[II] = newThetas[II] - v[II + 1]
            elseif II > max(ceil(k/2), floor(k/2) + 1)
                Δ[II] = newThetas[II] - v[II]
            else
                Δ[II] = Inf
            end
        end

        this.thetas = newThetas
        this.converged = newConverged
        this.velocity = Δ
        this.last_update = this.k
        return
    else
        # Maybeshould consdier re-establishing the system.
        # error("Must updated after each iterations, or else we lost track of the eigenvalues. ")
        EstablishEigenSystem(this)
    end

return end

"""
    Function establish the eigensystem of the current Tridiagonal Matrix using Lapack
    Routine, and erase all running parameters about the eigensystem.

    * Compute all the eigenvalues of the tridiagonal system and mark all as
    convernged.
"""
function EstablishEigenSystem(this::DynamicSymTridiagonal{T}) where {T<:AbstractFloat}
    TMatrix = GetT(this)
    eigenvalues = eigvals(TMatrix)
    sort!(eigenvalues)
    this.thetas = eigenValues
    this.betas = fill(true, length(eigenValues))
return end

"""
    Get the L matrix for this instance.
"""
function GetL(this::DynamicSymTridiagonal)
    return Bidiagonal(fill(1, this.k), convert(Vector{Float64}, this.L), :L)
end

"""
    Get the U matrix for this instance.
"""
function GetU(this::DynamicSymTridiagonal)
    return Bidiagonal(convert(Vector{Float64}, this.U), this.betas, :U)
end

"""
    Get the Tridiagonal Matrix
"""
function GetT(this::DynamicSymTridiagonal)
    return SymTridiagonal(this.alphas, this.betas)
end

function CountConverged(this::DynamicSymTridiagonal)
    return convert(Vector{Int64}, this.converged) |> sum
end

function GetConverged(this::DynamicSymTridiagonal) return this.thetas[this.converged] end

function ChangeVelocityTol!(
    this::DynamicSymTridiagonal,
    velocity_tol::T
) where {T<:AbstractFloat}
    this.velocity_tol = velocity_tol
end


using Test, StaticArrays, DelayEmbeddings

println("\nTesting reconstruct...")

@testset "reconstruct" begin

    data = Dataset(rand(10001,3))
    s = data[:, 1]; N = length(s)

    @testset "standard" begin

    	@testset "D = $(D), τ = $(τ)" for D in [1,2], τ in [2,3]

    		R = reconstruct(s, D, τ)

    		@test R[(1+τ):end, 1] == R[1:end-τ, 2]
    		@test size(R) == (length(s) - τ*D, D+1)
            @test embed(s, D+1, τ) == R
    	end
    end

    @testset "weighted" begin

    	@testset "D = $(D), τ = $(τ)" for D in [1,2], τ in [2,3]

            w = 0.1
    		R1 = reconstruct(s, D, τ)
    		R2 = reconstruct(s, D, τ, w)

            for γ in 0:D
                @test  (w^γ) * R1[1, γ+1][1] == R2[1, γ+1]
                @test  (w^γ) * R1[5, γ+1][1] == R2[5, γ+1]
            end
    	end
    end

    @testset "multi-time" begin

        D = 2
        τ1 = [2, 4]
        τ2 = [4, 8]

        R0 = reconstruct(s, D, 2)
        R1 = reconstruct(s, D, τ1)
        R2 = reconstruct(s, D, τ2)

        @test R1 == R0

        R2y = R2[:, 2]
        @test R2y == R0[5:end, 1]
        @test R2[:, 1] == R0[1:end-4, 1]
        @test size(R2) == (N-maximum(τ2), 3)

        @test_throws ArgumentError reconstruct(s, 4, τ1)

    end

    @testset "multidim " begin
        @testset "D = $(D), B = $(B)" for  D in [2,3], B in [2,3]

            τ = 3
            si = Matrix(data[:,1:B])
            sizedsi = Size(N,B)(si)
            R = reconstruct(sizedsi, D, τ)
            tr = Dataset(si)
            R2 = reconstruct(tr, D, τ)

            @test R == R2

            for dim in 1:B
                @test R[(1+τ):end, dim] == R[1:end-τ, dim+B]
                @test R2[(1+τ):end, dim] == R[1:end-τ, dim+B]
            end
            @test size(R) == (size(s,1) - τ*D, (D+1)*B)
        end
    end

    @testset "multidim multi-time" begin

        taus = [2 3; 4 6; 6 8]
        data2 = data[:, 1:2]
        data3 = Size(N, 2)(Matrix(data2))
        R1 = reconstruct(data2, 3, taus)
        R2 = reconstruct(data3, 3, taus)

        @test R1 == R2
        @test R1[:, 1] == data2[1:end-8, 1]
        @test R1[:, 2] == data2[1:end-8, 2]
        @test R1[:, 3] == data2[3:end-6, 1]

        # test error throws:
        taus = [0 0 0; 2 3 0; 4 6 0; 6 8 0]
        @test_throws ArgumentError reconstruct(data2, 5, taus)
        @test_throws ArgumentError reconstruct(data2, 4, taus)

    end
end


using Plots, LaTeXStrings, FourierGPE
gr(titlefontsize=12,size=(500,300),colorbar=false)

# Units:
# Units of ξ for length, 1/μ for time
# related by μ ≡ ħ²/mξ² = mc^2 where
# c = ħ/mξ is the speed of sound of
# the uniform system.

## set simulation parameters
L=(16.,16.,16.)
N=(64,64,64)
sim = Sim(L,N)
@unpack_Sim sim

γ = 0.05
tf = 4/γ
Nt = 200
t = LinRange(0.,tf,Nt)

## useful state functions
ψ0(x,y,z,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,z,0.0)/μ,0.0)+im*0.0)
healing(x,y,z,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,z,μ,g)))

## make initial state
x,y,z = X
ψi = ψ0.(x,y',reshape(z,1,1,length(z)),μ,g)
ψi = randn(N)+im*randn(N)
ϕi = kspace(ψi,sim)

@pack_Sim! sim

## Evolve in k space
sol = runsim(sim)

## plot
# function showpsi(x,y,ψ)
#     p1 = heatmap(x,y,abs2.(ψ),aspectratio=1)
#     xlabel!(L"x/\xi");ylabel!(L"y/\xi")
#     title!(L"|\psi|^2")
#     p2 = heatmap(x,y,angle.(ψ),aspectratio=1)
#     xlabel!(L"x/\xi");ylabel!(L"y/\xi")
#     title!(L"\textrm{phase}(\psi)")
#     p = plot(p1,p2,size=(600,300))
#     return p
# end

# pull out the ground state
ϕg = sol[end]
ψg = xspace(ϕg,sim)
showpsi(x,y,ψg[:,:,1])

# animate a slice
# dipole decay
anim = @animate for i=1:Nt
    showpsi(x,y,xspace(sol[i],sim)[:,:,1])
end

gif(anim,"./examples/3dquench.gif",fps=30)

## animate isosurface in Makie
using Makie, AbstractPlotting

function dense(i)
    ψm = xspace(sol[i],sim)
    density = abs2.(ψm)
    pmax = maximum(density)
    return density/pmax
end

function densityfilm(Nt,saveto="examples/3dquenchisotest.gif")
scene = Scene()
tindex = Node(1)

scene = volume(lift(i -> dense(i), tindex), algorithm = :iso,show_axis=false)

record(scene, saveto, 1:Nt-10) do i
    tindex[] = i
end
end

densityfilm(Nt)


@doc raw"""
    SkewHermitianMatrices{n,𝔽} <: AbstractEmbeddedManifold{𝔽,TransparentIsometricEmbedding}

The [`AbstractManifold`](@ref) $ \operatorname{SkewHerm}(n)$ consisting of the real- or
complex-valued skew-hermitian matrices of size ``n × n``, i.e. the set

````math
\operatorname{SkewHerm}(n) = \bigl\{p  ∈ 𝔽^{n × n}\ \big|\ p^{\mathrm{H}} = -p \bigr\},
````
where $\cdot^{\mathrm{H}}$ denotes the Hermitian, i.e. complex conjugate transpose,
and the field $𝔽 ∈ \{ ℝ, ℂ, ℍ\}$.

Though it is slightly redundant, usually the matrices are stored as ``n × n`` arrays.

Note that in this representation, the real-valued part of the diagonal must be zero,
which is also reflected in the
[`manifold_dimension`](@ref manifold_dimension(::SkewHermitianMatrices{N,𝔽}) where {N,𝔽}).

# Constructor

    SkewHermitianMatrices(n::Int, field::AbstractNumbers=ℝ)

Generate the manifold of ``n × n`` skew-hermitian matrices.
"""
struct SkewHermitianMatrices{n,𝔽} <:
       AbstractEmbeddedManifold{𝔽,TransparentIsometricEmbedding} end

function SkewHermitianMatrices(n::Int, field::AbstractNumbers=ℝ)
    return SkewHermitianMatrices{n,field}()
end

@doc raw"""
    SkewSymmetricMatrices{n}

Generate the manifold of ``n × n`` real skew-symmetric matrices.
This is equivalent to [`SkewHermitianMatrices(n, ℝ)`](@ref).

# Constructor

    SkewSymmetricMatrices(n::Int)
"""
const SkewSymmetricMatrices{n} = SkewHermitianMatrices{n,ℝ}

SkewSymmetricMatrices(n::Int) = SkewSymmetricMatrices{n}()
@deprecate SkewSymmetricMatrices(n::Int, 𝔽) SkewHermitianMatrices(n, 𝔽)

function allocation_promotion_function(
    ::SkewHermitianMatrices{<:Any,ℂ},
    ::typeof(get_vector),
    args::Tuple,
)
    return complex
end

@doc raw"""
    check_point(M::SkewHermitianMatrices{n,𝔽}, p; kwargs...)

Check whether `p` is a valid manifold point on the [`SkewHermitianMatrices`](@ref) `M`, i.e.
whether `p` is a skew-hermitian matrix of size `(n,n)` with values from the corresponding
[`AbstractNumbers`](@ref) `𝔽`.

The tolerance for the skew-symmetry of `p` can be set using `kwargs...`.
"""
function check_point(M::SkewHermitianMatrices{n,𝔽}, p; kwargs...) where {n,𝔽}
    mpv = check_point(decorated_manifold(M), p; kwargs...)
    mpv === nothing || return mpv
    if !isapprox(p, -p'; kwargs...)
        return DomainError(
            norm(p + p'),
            "The point $(p) does not lie on $M, since it is not skew-hermitian.",
        )
    end
    return nothing
end

"""
    check_vector(M::SkewHermitianMatrices{n}, p, X; kwargs... )

Check whether `X` is a tangent vector to manifold point `p` on the
[`SkewHermitianMatrices`](@ref) `M`, i.e. `X` must be a skew-hermitian matrix of size `(n,n)`
and its values have to be from the correct [`AbstractNumbers`](@ref).
The tolerance for the skew-symmetry of `p` and `X` can be set using `kwargs...`.
"""
function check_vector(M::SkewHermitianMatrices, p, X; kwargs...)
    return check_point(M, X; kwargs...)  # manifold is its own tangent space
end

decorated_manifold(M::SkewHermitianMatrices{N,𝔽}) where {N,𝔽} = Euclidean(N, N; field=𝔽)

function get_basis(M::SkewHermitianMatrices, p, B::DiagonalizingOrthonormalBasis)
    Ξ = get_basis(M, p, DefaultOrthonormalBasis()).data
    κ = zeros(real(eltype(p)), manifold_dimension(M))
    return CachedBasis(B, κ, Ξ)
end

function get_coordinates!(
    M::SkewSymmetricMatrices{N},
    Y,
    p,
    X,
    ::DefaultOrthonormalBasis{ℝ,TangentSpaceType},
) where {N}
    dim = manifold_dimension(M)
    @assert size(Y) == (dim,)
    @assert size(X) == (N, N)
    @assert dim == div(N * (N - 1), 2)
    k = 1
    for i in 1:N, j in (i + 1):N
        @inbounds Y[k] = X[i, j] * sqrt(2)
        k += 1
    end
    return Y
end
function get_coordinates!(
    M::SkewHermitianMatrices{N,ℂ},
    Y,
    p,
    X,
    ::DefaultOrthonormalBasis{ℂ,TangentSpaceType},
) where {N}
    dim = manifold_dimension(M)
    @assert size(Y) == (dim,)
    @assert size(X) == (N, N)
    @assert dim == N^2
    k = 1
    for i in 1:N, j in i:N
        if i == j # real-part zero on the diagonal -> just one basis vector per diag entry
            @inbounds Y[k] = imag(X[i, j])
            k += 1
        else
            @inbounds Y[k] = real(X[i, j]) * sqrt(2)
            @inbounds Y[k + 1] = imag(X[i, j]) * sqrt(2)
            k += 2
        end
    end
    return Y
end

function get_vector!(
    M::SkewSymmetricMatrices{N},
    Y,
    p,
    X,
    ::DefaultOrthonormalBasis{ℝ,TangentSpaceType},
) where {N}
    dim = manifold_dimension(M)
    @assert size(X) == (dim,)
    @assert size(Y) == (N, N)
    k = 1
    for i in 1:N
        Y[i, i] = convert(eltype(p), 0.0)
    end
    for i in 1:N, j in (i + 1):N
        @inbounds Y[i, j] = X[k] / sqrt(2)
        @inbounds Y[j, i] = -X[k] / sqrt(2)
        k += 1
    end
    return Y
end
function get_vector!(
    M::SkewHermitianMatrices{N,ℂ},
    Y,
    p,
    X,
    ::DefaultOrthonormalBasis{ℂ,TangentSpaceType},
) where {N}
    dim = manifold_dimension(M)
    @assert size(X) == (dim,)
    @assert size(Y) == (N, N)
    k = 1
    for i in 1:N, j in i:N
        if i == j # real zero on the diag
            @inbounds Y[i, j] = X[k] * 1im
            k += 1
        else
            @inbounds Y[i, j] = (X[k] + X[k + 1] * 1im) / sqrt(2)
            k += 2
            @inbounds Y[j, i] = -conj(Y[i, j])
        end
    end
    return Y
end

@doc raw"""
    manifold_dimension(M::SkewHermitianMatrices{n,𝔽})

Return the dimension of the [`SkewHermitianMatrices`](@ref) matrix `M` over the number
system `𝔽`, i.e.

````math
\dim \mathrm{SkewHerm}(n,ℝ) = \frac{n(n+1)}{2} \dim_ℝ 𝔽 - n,
````

where ``\dim_ℝ 𝔽`` is the [`real_dimension`](@ref) of ``𝔽``. The first term corresponds to
only the upper triangular elements of the matrix being unique, and the second term
corresponds to the constraint that the real part of the diagonal be zero.
"""
function manifold_dimension(::SkewHermitianMatrices{N,𝔽}) where {N,𝔽}
    return div(N * (N + 1), 2) * real_dimension(𝔽) - N
end

@doc raw"""
    project(M::SkewHermitianMatrices, p)

Projects `p` from the embedding onto the [`SkewHermitianMatrices`](@ref) `M`, i.e.

````math
\operatorname{proj}_{\operatorname{SkewHerm}(n)}(p) = \frac{1}{2} \bigl( p - p^{\mathrm{H}} \bigr),
````

where $\cdot^{\mathrm{H}}$ denotes the Hermitian, i.e. complex conjugate transposed.
"""
project(::SkewHermitianMatrices, ::Any)

function project!(M::SkewHermitianMatrices, q, p)
    q .= (p .- p') ./ 2
    return q
end

@doc raw"""
    project(M::SkewHermitianMatrices, p, X)

Project the matrix `X` onto the tangent space at `p` on the [`SkewHermitianMatrices`](@ref) `M`,

````math
\operatorname{proj}_p(X) = \frac{1}{2} \bigl( X - X^{\mathrm{H}} \bigr),
````

where $\cdot^{\mathrm{H}}$ denotes the Hermitian, i.e. complex conjugate transposed.
"""
project(::SkewHermitianMatrices, ::Any, ::Any)

project!(M::SkewHermitianMatrices, Y, p, X) = project!(M, Y, X)

function Base.show(io::IO, ::SkewHermitianMatrices{n,F}) where {n,F}
    return print(io, "SkewHermitianMatrices($(n), $(F))")
end
function Base.show(io::IO, ::SkewSymmetricMatrices{n}) where {n}
    return print(io, "SkewSymmetricMatrices($(n))")
end


using Test
using StaticArrays, IntervalSets, LinearAlgebra

import ClimaCore: slab, Domains, Meshes, Topologies, Spaces
import ClimaCore.Geometry: Geometry

@testset "1d domain space" begin
    FT = Float64
    domain = Domains.RectangleDomain(
        Geometry.XPoint{FT}(-3)..Geometry.XPoint{FT}(5),
        Geometry.YPoint{FT}(0)..Geometry.YPoint{FT}(0),
        x1periodic = true,
        x2periodic = false,
        x2boundary = (:south, :north),
    )
    mesh = Meshes.EquispacedRectangleMesh(domain, 1, 1)
    grid_topology = Topologies.GridTopology(mesh)

    quad = Spaces.Quadratures.GLL{4}()
    points, weights = Spaces.Quadratures.quadrature_points(FT, quad)

    space = Spaces.SpectralElementSpace1D(grid_topology, quad)

    coord_data = Spaces.coordinates_data(space)
    @test eltype(coord_data) == Geometry.XPoint{Float64}

    array = parent(Spaces.coordinates_data(space))
    @test size(array) == (4, 1, 1)
    coord_slab = slab(Spaces.coordinates_data(space), 1)
    @test coord_slab[1] == Geometry.XPoint{FT}(-3)
    @test coord_slab[4] == Geometry.XPoint{FT}(5)

    local_geometry_slab = slab(space.local_geometry, 1)
    dss_weights_slab = slab(space.dss_weights, 1)

    for i in 1:4
        @test Geometry.components(local_geometry_slab[i].∂x∂ξ) ≈
              @SMatrix [8 / 2]
        @test Geometry.components(local_geometry_slab[i].∂ξ∂x) ≈
              @SMatrix [2 / 8]
        @test local_geometry_slab[i].J ≈ (8 / 2)
        @test local_geometry_slab[i].WJ ≈ (8 / 2) * weights[i]
        if i in (1, 4)
            @test dss_weights_slab[i] ≈ 1 / 2
        else
            @test dss_weights_slab[i] ≈ 1
        end
    end

end

@testset "1×1 domain space" begin
    FT = Float32
    domain = Domains.RectangleDomain(
        Geometry.XPoint{FT}(-3)..Geometry.XPoint{FT}(5),
        Geometry.YPoint{FT}(-2)..Geometry.YPoint{FT}(8),
        x1periodic = true,
        x2periodic = false,
        x2boundary = (:south, :north),
    )
    mesh = Meshes.EquispacedRectangleMesh(domain, 1, 1)
    grid_topology = Topologies.GridTopology(mesh)

    quad = Spaces.Quadratures.GLL{4}()
    points, weights = Spaces.Quadratures.quadrature_points(FT, quad)

    space = Spaces.SpectralElementSpace2D(grid_topology, quad)

    array = parent(Spaces.coordinates_data(space))
    @test size(array) == (4, 4, 2, 1)
    coord_slab = slab(Spaces.coordinates_data(space), 1)
    @test coord_slab[1, 1] ≈ Geometry.XYPoint{FT}(-3.0, -2.0)
    @test coord_slab[4, 1] ≈ Geometry.XYPoint{FT}(5.0, -2.0)
    @test coord_slab[1, 4] ≈ Geometry.XYPoint{FT}(-3.0, 8.0)
    @test coord_slab[4, 4] ≈ Geometry.XYPoint{FT}(5.0, 8.0)

    local_geometry_slab = slab(space.local_geometry, 1)
    dss_weights_slab = slab(space.dss_weights, 1)

    for i in 1:4, j in 1:4
        @test Geometry.components(local_geometry_slab[i, j].∂x∂ξ) ≈
              @SMatrix [8/2 0; 0 10/2]
        @test Geometry.components(local_geometry_slab[i, j].∂ξ∂x) ≈
              @SMatrix [2/8 0; 0 2/10]
        @test local_geometry_slab[i, j].J ≈ (10 / 2) * (8 / 2)
        @test local_geometry_slab[i, j].WJ ≈
              (10 / 2) * (8 / 2) * weights[i] * weights[j]
        if i in (1, 4)
            @test dss_weights_slab[i, j] ≈ 1 / 2
        else
            @test dss_weights_slab[i, j] ≈ 1
        end
    end

    @test length(space.boundary_surface_geometries) == 2
    @test keys(space.boundary_surface_geometries) == (:south, :north)
    @test sum(parent(space.boundary_surface_geometries.north.sWJ)) ≈ 8
    @test parent(space.boundary_surface_geometries.north.normal)[1, :, 1] ≈
          [0.0, 1.0]
end


# An implementation of a FISTA-like method, where the smooth part of the objective function can be strongly convex.

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalOperators: Zero
using LinearAlgebra
using Printf

"""
    SFISTAIteration(; <keyword-arguments>)

Instantiate the FISTA-like method in [3] for solving strongly-convex composite optimization problems of the form

    minimize f(x) + h(x),

where h is proper closed convex and f is a continuously differentiable function that is μ-strongly convex and whose gradient is
Lf-Lipschitz continuous.

The scheme is based on Nesterov's accelerated gradient method [1, Eq. (4.9)] and Beck's method for the convex case [2]. Its full
definition is given in [3, Algorithm 2.2.2.], and some analyses of this method are given in [3, 4, 5]. Another perspective is that
it is a special instance of [4, Algorithm 1] in which μh=0.

# Arguments
- `y0`: initial point; must be in the domain of h.
- `f=Zero()`: smooth objective term.
- `h=Zero()`: proximable objective term.
- `μf=0` : strong convexity constant of f (see above).
- `Lf` : Lipschitz constant of ∇f (see above).
- `adaptive=false` : enables the use of adaptive stepsize selection.

# References
1. Nesterov, Y. (2013). Gradient methods for minimizing composite functions. Mathematical Programming, 140(1), 125-161.
2. Beck, A., & Teboulle, M. (2009). A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1), 183-202.
3. Kong, W. (2021). Accelerated Inexact First-Order Methods for Solving Nonconvex Composite Optimization Problems. arXiv preprint arXiv:2104.09685.
4. Kong, W., Melo, J. G., & Monteiro, R. D. (2021). FISTA and Extensions - Review and New Insights. arXiv preprint arXiv:2107.01267.
5. Florea, M. I. (2018). Constructing Accelerated Algorithms for Large-scale Optimization-Framework, Algorithms, and Applications.
"""
Base.@kwdef struct SFISTAIteration{R,C<:Union{R,Complex{R}},Tx<:AbstractArray{C},Tf,Th}
    y0::Tx
    f::Tf = Zero()
    h::Th = Zero()
    Lf::R
    μf::R = real(eltype(Lf))(0.0)
end

Base.IteratorSize(::Type{<:SFISTAIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct SFISTAState{R, Tx}
    λ::R                                # stepsize (mutable if iter.adaptive == true).
    yPrev::Tx                           # previous main iterate.
    y:: Tx = zero(yPrev)                # main iterate.
    xPrev::Tx = copy(yPrev)             # previous auxiliary iterate.
    x::Tx  = zero(yPrev)                # auxiliary iterate (see [3]).
    xt::Tx = zero(yPrev)                # prox center used to generate main iterate y.
    τ::R = real(eltype(yPrev))(1.0)     # helper variable (see [3]).
    a::R = real(eltype(yPrev))(0.0)     # helper variable (see [3]).
    APrev::R = real(eltype(yPrev))(1.0) # previous A (helper variable).
    A::R = real(eltype(yPrev))(0.0)     # helper variable (see [3]).
    gradf_xt::Tx = zero(yPrev)          # array containing ∇f(xt).
end

function Base.iterate(
    iter::SFISTAIteration,
    state::SFISTAState = SFISTAState(λ=1/iter.Lf, yPrev=copy(iter.y0))
)
    # Set up helper variables.
    state.τ = state.λ * (1 + iter.μf * state.APrev)
    state.a = (state.τ + sqrt(state.τ ^ 2 + 4 * state.τ * state.APrev)) / 2
    state.A = state.APrev + state.a
    state.xt .= (state.APrev / state.A) .* state.yPrev + (state.a / state.A) .* state.xPrev
    gradient!(state.gradf_xt, iter.f, state.xt)
    λ2 = state.λ / (1 + state.λ * iter.μf)
    # FISTA acceleration steps.
    prox!(state.y, iter.h, state.xt - λ2 * state.gradf_xt, λ2)
    state.x .= state.xPrev .+ (state.a / (1 + state.A * iter.μf)) .* ((state.y .- state.xt) ./ state.λ .+ iter.μf .* (state.y .- state.xPrev))
    # Update state variables.
    state.yPrev .= state.y
    state.xPrev .= state.x
    state.APrev = state.A
    return state, state
end

## Solver.

struct SFISTA{R, K}
    maxit::Int
    tol::R
    termination_type::String
    verbose::Bool
    freq::Int
    kwargs::K
end

# Different stopping conditions (sc). Returns the current residual value and whether or not a stopping condition holds.
function check_sc(state::SFISTAState, iter::SFISTAIteration, tol, termination_type)
    if termination_type == "AIPP"
        # AIPP-style termination [4]. The main inclusion is: r ∈ ∂_η(f + h)(y).
        r = (iter.y0 - state.x) / state.A
        η = (norm(iter.y0 - state.y) ^ 2 - norm(state.x - state.y) ^ 2) / (2 * state.A)
        res = (norm(r) ^ 2 + max(η, 0.0)) / max(norm(iter.y0 - state.y + r) ^ 2, 1e-16)
    else
        # Classic (approximate) first-order stationary point [4]. The main inclusion is: r ∈ ∇f(y) + ∂h(y).
        λ2 = state.λ / (1 + state.λ * iter.μf)
        gradf_y, = gradient(iter.f, state.y)
        r = gradf_y - state.gradf_xt + (state.xt - state.y) / λ2
        res = norm(r)
    end
    return res, (res <= tol || res ≈ tol)
end

# Functor ('function-like object') for the above type.
function (solver::SFISTA)(y0; kwargs...)
    raw_iter = SFISTAIteration(; y0=y0, solver.kwargs..., kwargs...)
    stop(state::SFISTAState) = check_sc(state, raw_iter, solver.tol, solver.termination_type)[2]
    disp((it, state)) = @printf("%5d | %.3e\n", it, check_sc(state, raw_iter, solver.tol, solver.termination_type)[1])
    iter = take(halt(raw_iter, stop), solver.maxit)
    iter = enumerate(iter)
    if solver.verbose
        iter = tee(sample(iter, solver.freq), disp)
    end
    num_iters, state_final = loop(iter)
    return state_final.y, num_iters
end

SFISTA(; maxit=10_000, tol=1e-6, termination_type="", verbose=false, freq=(maxit < Inf ? Int(maxit/100) : 100), kwargs...) =
    SFISTA(maxit, tol, termination_type, verbose, freq, kwargs)


@testset "LSQR" begin

	@testset "small test" begin
		A   = [1 1; 0 -1; 1 0.]
		r1  = A[:,1]   # compatible
		r2  = randn(3) # incompatible
		r3  = [1;1;-1.] # in the nullpace of A'

		x1 = lsqr(A,r1,out=2,atol=1e-8,doBidiag=true,storeInterm=true)
		@test norm(x1[1][:,end]-[1;0])<1e-10
		@test x1[2]==1
		x2 = lsqr(A,r2,out=1,atol=1e-8)
		@test x2[2]==2
		x3 = lsqr(A,r3,out=2)
		@test all(x3[1].==0)
	end

	@testset "sparse random" begin
		A   = sprandn(100,10,.2)
		rhs = randn(100)
		xgt = full(A)\rhs
		Af = (flag,x,a=0.0,v=0.0) -> ((flag=='F') ? A*x +a*v : A'*x+a*v)

		x1  = lsqr(LinearOperator(A),rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10,out=1)
		x2  = lsqr(A,rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)
		x3  = lsqr(Af,rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)

		@test norm(xgt-x1[1])/norm(xgt) < 1e-6
		@test norm(xgt-x2[1])/norm(xgt) < 1e-6
		@test norm(xgt-x3[1])/norm(xgt) < 1e-6
		@test norm(x1[1]-x2[1])/norm(x1[1]) < 1e-12
		@test norm(x2[1]-x3[1])/norm(x3[1]) < 1e-12
	end
	@testset "hilbert matrix" begin
		A   = matrixdepot("hilb",3,3)
		rhs = [1;0;0.]
		xgt = A\rhs

		x1  = lsqr(LinearOperator(A),rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)
		x2  = lsqr(A,rhs,atol=1e-10,btol=1e-10,condlim=1e5,maxIter=10)

		@test norm(xgt-x1[1])/norm(xgt) < 1e-6
		@test norm(xgt-x2[1])/norm(xgt) < 1e-6
		@test norm(x1[1]-x2[1])/norm(x1[1]) < 1e-12
		conA = norm(vec(A))*norm(vec(pinv(A)))
		@test norm(conA-x1[3][end,end])/conA < 1e-12
		@test norm(conA-x2[3][end,end])/conA < 1e-12
	end

end


# # Model Forecasting

# ## Set the State-Space model
using Plots
using AnalogDataAssimilation
using DifferentialEquations

σ = 10.0
ρ = 28.0
β = 8.0 / 3

dt_integration = 0.01
dt_states = 1
dt_obs = 8
parameters = [σ, ρ, β]
var_obs = [1]
nb_loop_train = 100
nb_loop_test = 10
sigma2_catalog = 0.0
sigma2_obs = 2.0

ssm = StateSpaceModel(
    lorenz63,
    dt_integration,
    dt_states,
    dt_obs,
    parameters,
    var_obs,
    nb_loop_train,
    nb_loop_test,
    sigma2_catalog,
    sigma2_obs,
)

# - compute u0 to be in the attractor space

u0 = [8.0; 0.0; 30.0]
tspan = (0.0, 5.0)
prob = ODEProblem(ssm.model, u0, tspan, parameters)
u0 = last(solve(prob, save_everystep = false))

# ## Generate data

xt, yo, catalog = generate_data(ssm, u0);

plot(xt.t, xt[1])
scatter!(yo.t, yo[1]; markersize = 2)

# ## Data assimilation with model forecasting

np = 100
DA = DataAssimilation(ssm, xt)
@time x̂ = forecast(DA, yo, AnPF(np), progress = false);
println(RMSE(xt, x̂))

# ## Plot the times series

plot(xt.t, x̂[1])
scatter!(xt.t, xt[1]; markersize = 2)
plot!(xt.t, x̂[2])
scatter!(xt.t, xt[2]; markersize = 2)
plot!(xt.t, x̂[3])
scatter!(xt.t, xt[3]; markersize = 2)

# ## Plot the phase-space plot

p = plot3d(
    1,
    xlim = (-25, 25),
    ylim = (-25, 25),
    zlim = (0, 50),
    title = "Lorenz 63",
    marker = 2,
)
for x in eachrow(vcat(x̂.u'...))
    push!(p, x...)
end
p


using QuadraticTools
using Test

@testset "QuadraticTools.jl" begin
    # test the validity checker
    @test_throws DomainError QuadraticTools._validitycheck(0)
    @test_nowarn QuadraticTools._validitycheck(0.1)
    @test_nowarn QuadraticTools._validitycheck(0 + 5im)

    # test the delta calculation
    @test calcdelta(1, 1, -1) == 5
    @test calcdelta(1, 1, 1) == -3
    @test calcdelta(2, -7, -5) == 89

    # test conversions from the expanded form
    @test tofactored(Expanded(1, 5, 6)) == Factored(1, -3.0, -2.0)
    @test tovertex(Expanded(1, 5, 6)) == Vertex(1, -2.5, -0.25)
    @test calcvalue(Expanded(1, 5, 6), 10) == 156

    # test conversions from the factored form
    @test toexpanded(Factored(1, -3.0, -2.0)) == Expanded(1, 5.0, 6.0)
    @test tovertex(Factored(1, -3.0, -2.0)) == Vertex(1, -2.5, -0.25)
    @test calcvalue(Factored(1, -3.0, -2.0), 10) == 156.0
    @test tofactored(Vertex(-1, -2.5, -0.25)) == Factored(-1, -2.5 - 0.5im, -2.5 + 0.5im)

    # test conversions from the vertex form
    @test toexpanded(Vertex(1, -2.5, -0.25)) == Expanded(1, 5.0, 6.0)
    @test tofactored(Vertex(1, -2.5, -0.25)) == Factored(1, -3.0, -2.0)
    @test calcvalue(Vertex(1, -2.5, -0.25), 10) == 156.0
end


# This example is used to generate plots in the paper, "Reflection from a multi-species material and its transmitted effective wavenumber." Proc. R. Soc. (2018): 20170864.

# Everything related to ploting has been commented so that this example can run with requiring the Plots package. That way, this example can be checked everytime the package is modified.

using EffectiveWaves

# using LaTeXStrings
# using Plots
# height=500
 # pyplot(linewidth=3, size=(2.6*height,height), border=false)

 # Plots.scalefontsizes(1.7)

## choose material

  # concrete
  medium = LimeStone
  inclusion1 = AirDry
  inclusion2 = Brick

  ωfactor = 50.0;
  ωs = 20*ωfactor.*LinRange(real(medium.c/10000),real(medium.c),400) # k from 0 to 1

  volfrac = 0.16
  r1 = 0.1/ωfactor; vol1 = 0.06
  r2 = 1.0/ωfactor; vol2 = volfrac - vol1

  sp1 = Specie(ρ=inclusion1.ρ ,r=r1, c=inclusion1.c, volfrac = vol1)
  sp2 = Specie(ρ=inclusion2.ρ ,r=r2, c=inclusion2.c, volfrac = volfrac-vol1)

  # True wavenumber
  kTs  = wavenumber_low_volfrac(ωs, medium, [sp1,sp2]; tol=0.5e-5);
  # Approximate challis
  kTCs = wavenumber_challis(ωs, medium, [sp1,sp2]; hankel_order=10);

  kTs_arr = [kTs,kTCs];
  speed_arr = [ ωs./real(ks) for ks in kTs_arr];
  push!(speed_arr, 0 .*ωs .+ real(medium.c))
  atten_arr = imag.(kTs_arr)

  styles = [:solid :dashdot :dot]
  # labs = [L"Effective $k_{*}$" L"Approximate $k_{*C}$" "Lime Stone"]
  ys_arr = speed_arr;
  xs = r1.*(ωs./real(medium.c));
  m =5;
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  # p1 = plot(xs, ys_arr, xlabel=L"a_S k", ylabel="sound speed (m/s)", labels=labs, line = styles)

  ys_arr = atten_arr;
  # labs = [L"Effective $k_{*}$" L"Approximate $k_{*C}$"]
  # p2 = plot(xs, ys_arr, labels=labs, xlabel=L"a_S k", ylabel="attenuation (1/m)", line=styles)

  # plot(p1,p2)
  # try mkdir("media") end
  # savefig("media/compare_concrete_large-w.png")
  # savefig("media/compare_concrete_large-w.pdf")
  # gui()

# Plots.scalefontsizes(1/1.7)


module Mixtures

using Roots, SpecialFunctions, StatsFuns
using Random, Distributions, Statistics

export fitmixture, EM

function MLE(x; γ=missing)
    if ismissing(γ)
        γ = ones(size(x))
    end
    N = sum(γ)
    μ = sum((γ.*x))/N
    σ = log(μ) - sum(γ.*log.(x))/N

    # initial guess
    α = ((3 - σ) + sqrt((σ-3)^2 + 24σ)) / (12σ)
    if isinf(α)
        α = 10
    end

    try
        F(α) = log(α) - digamma(α) - σ
        α = find_zero(F, α, Order1())
    catch
        println("ERROR:")
        println("bounds=($(minimum(x)), $(maximum(x)))")
        println("N=$(N)")
        println("σ=$(σ)")
        println("μ=$(μ)")
        println(((3 - σ) + sqrt((σ-3)^2 + 24σ)) / (12σ))
    end
    β = α/μ

    return α, β
end

function weights(γ)
    return sum(γ, dims=1)'/sum(γ)
end

logρ(x,α,β) = (α*log(β) - loggamma(α)) .+ (α-1).*log.(x) .- β.*x
ρ(x,α,β)    = (x != 0) ? exp(logρ(x,α,β)) : β^α / gamma(α)

function posterior(x, Π, α, β; dropout=false)
    γ = Array{Float64, 2}(undef, length(x), length(α))

    for k = 1:length(α)
        γ[:,k] = Π[k]*ρ.(x, α[k], β[k])
    end

    if dropout
        χ = minimum(x[x.>0])
        p = vec([gammacdf(α[k], 1/β[k], χ) for k in 1:length(α)])
        γ[findall(x.==0),:] .= p'
    end

    if any(isinf.(γ))
        γ[isinf.(γ)] .= 1
    end

    γ = γ ./ sum(γ, dims=2)

    return γ
end

loglikelihood(x, Π, α, β) = sum(log.(sum(Π' .* hcat([ρ.(x,α[k],β[k]) for k in 1:length(α)]...), dims=2)))

function sample(Π, α, β, N; k=missing)
    if ismissing(k)
        MM = MixtureModel(Gamma, [(α[k], 1/β[k]) for k in 1:length(α)], vec(Π))
    else
        MM = Gamma(α[k], 1/β[k])
    end
    return rand(MM, N)
end

function kmeans(x, k::Integer; verbose = false, maxᵢ::Integer = 1000)
    q  = [quantile(x, q) for q in range(0, 1, length=k+1)]
    μ  = [mean(x[(q[i-1] .<= x) .& (x .<= q[i])]) for i in 2:length(q)]
    S₀ = zeros(length(x), 1)
    for n in 1:maxᵢ
        S = getindex.(argmin(abs.(x .- μ'), dims=2), 2)
        if !any(x->x!=0, S-S₀)
            break
        end
        if verbose
            println("iteration $(n): μ = $(μ)")
        end
        μ  = [mean(x[findall(S.==i)]) for i in 1:k]
        S₀ = S
    end

    σ = [std(x[findall(S₀.==i)]) for i in 1:k]

    return μ, σ, S₀
end

function EM(x, K::Integer; ϵ::Real = 1e-5, maxᵢ::Integer = 10000, verbose::Bool = false)
    μ, σ, S₀ = kmeans(x, K)

    β = μ./σ.^2
    α = μ.*β
    Π = [sum(S₀.==k)/length(S₀) for k in 1:K]

    δ  = Inf
    L₀ = loglikelihood(x, Π, α, β)
    i  = 0
    while abs(δ) > ϵ && i < maxᵢ
        # expectation + maximization
        γ = posterior(x, Π, α, β)

        Π = weights(γ)
        Θ = [MLE(x, γ=γ[:,k]) for k in 1:K]
        α = [θ[1] for θ in Θ]
        β = [θ[2] for θ in Θ]

        L = loglikelihood(x, Π, α, β)
        if verbose
            println("iteration $(i). loglikelihood = $(L)")
        end
        δ  = L - L₀
        L₀ = L
        i += 1
    end

    ι = sortperm(α./β)

    return L₀, (Π[ι], α[ι], β[ι])
end

function cluster(x, Π, α, β; dropout=false)
    println("computing posterior")
    γ = posterior(x, Π, α, β, dropout=dropout)
    println("getting index")
    χ = getindex.(argmax(γ, dims=2), 2)
    println("returning")
    return [x[findall(χ .== k)] for k in 1:length(α)]
end

# XXX: right now just returns the posteriors for component with largest mean
function fitmixture(data)
    _, params = EM(data.+1e-6, 2)
    ρ = posterior(data, params...)
    return ρ[:,2]
end

end


# # Housing data

# In this example, we create a linear regression model that predicts housing data.
# It replicates the housing data example from the [Knet.jl readme](https://github.com/denizyuret/Knet.jl).
# Although we could have reused more of Flux (see the MNIST example), the library's abstractions are very
# lightweight and don't force you into any particular strategy.

# A linear model can be created as a neural network with a single layer.
# The number of inputs is the same as the features that the data has.
# Each input is connected to a single output with no activation function.
# Then, the output of the model is a linear function that predicts unseen data.

# ![singleneuron](img/singleneuron.svg)

# Source: [Dive into Deep Learning](http://d2l.ai/chapter_linear-networks/linear-regression.html#from-linear-regression-to-deep-networks)

# To run this example, we need the following packages:

using Flux
using Flux: gradient
using Flux.Optimise: update!
using DelimitedFiles, Statistics
using Parameters: @with_kw

# We set default values for the learning rate (for the training routine) and the percentage of
# the data that we use when testing the model:

@with_kw mutable struct Hyperparams
    ## Learning rate
    lr::Float64 = 0.1
    ## Train Test split ratio, define percentage of data to be used as Test data
    split_ratio::Float64 = 0.1
end

# ## Data

# We create the function `get_processed_data` to load the housing data, normalize it,
# and finally split it into train and test datasets:

function get_processed_data(args)
    isfile("housing.data") ||
        download("https://raw.githubusercontent.com/MikeInnes/notebooks/master/housing.data",
            "housing.data")

    rawdata = readdlm("housing.data")'

    ## The last feature is our target -- the price of the house.
    split_ratio = args.split_ratio ## For the train test split

    x = rawdata[1:13,:]
    y = rawdata[14:14,:]

    ## Normalise the data
    x = (x .- mean(x, dims = 2)) ./ std(x, dims = 2)

    ## Split into train and test sets
    split_index = floor(Int,size(x,2)*split_ratio)
    x_train = x[:,1:split_index]
    y_train = y[:,1:split_index]
    x_test = x[:,split_index+1:size(x,2)]
    y_test = y[:,split_index+1:size(x,2)]

    train_data = (x_train, y_train)
    test_data = (x_test, y_test)

    return train_data,test_data
end

# This function performs the following tasks:

# 1. Downloads the housing data. The original size of the data is 505 rows and 14 columns.
# 2. Loads the data as a 14x505 matrix. This is the shape that Flux expects.
# 3. Splits the data into features and a target. Notice that the 14th row corresponds to the target for each example.
# 4. Normalizes the data. For more information on normalizing data, see [How to Use StandardScaler and MinMaxScaler Transforms in Python](https://machinelearningmastery.com/standardscaler-and-minmaxscaler-transforms-in-python/).
# 5. Splits the data into train and test datasets.

# ## Model
# We use a struct to define the model’s parameters.
# It contains an array for holding the weights *W* and a vector for the bias term *b*:

mutable struct model
    W::AbstractArray
    b::AbstractVector
end

# Also, we create the function `predict` to compute the model’s output:

predict(x, m) = m.W*x .+ m.b

# Notice that the function `predict` takes as an argument the model struct we defined above.

# ## Loss function

# The most commonly used loss function for Linear Regression is Mean Squared Error (MSE).
# We define the MSE function as:

meansquarederror(ŷ, y) = sum((ŷ .- y).^2)/size(y, 2)

# **Note:** An implementation of the MSE function is also available in
# [Flux](https://fluxml.ai/Flux.jl/stable/models/losses/#Flux.Losses.mse).

# ## Train function
# Finally, we define the `train` function so that the model learns the best parameters (*W* and *b*):

function train(; kws...)
    ## Initialize the Hyperparamters
    args = Hyperparams(; kws...)

    ## Load the data
    (x_train,y_train),(x_test,y_test) = get_processed_data(args)

    ## The model
    m = model((randn(1,13)),[0.])

    loss(x, y) = meansquarederror(predict(x, m), y)

    ## Training
    η = args.lr
    θ = params(m.W, m.b)

    for i = 1:500
        g = gradient(() -> loss(x_train, y_train), θ)
        for x in θ
            update!(x, g[x]*η)
        end
        if i%100==0
            @show loss(x_train, y_train)
        end
    end

    ## Predict the RMSE on the test set
    err = meansquarederror(predict(x_test, m),y_test)
    println(err)
end

# The function above initializes the model’s parameters *W* and *b* randomly.
# Then, it sets the learning rate η and θ as a
# [params object](https://fluxml.ai/Flux.jl/stable/training/training/#Flux.params)
# that points to  W and b. Also, it sets a
# [custom training loop](https://fluxml.ai/Flux.jl/stable/training/training/#Custom-Training-loops)
# which is the [Gradient descent algorithm](https://en.wikipedia.org/wiki/Gradient_descent).
# Finally, it computes the MSE for the test set.

# ## Run the example
# We call the `train` function to run the Housing data example:

cd(@__DIR__)
train()


using TSVD: tsvd

"""
    PCA(ncomps; options...)
    PCA(;ncomps=nothing, options...)

Use principal components analysis (PCA) to form a low-rank orthonormal basis of the input. Uses deterministic singular-value decomposition (SVD) to decompose data.

If `ncomps` is `nothing`, the basis will not be truncated (i.e. `ncomps` is equal to the number of frames). `ncomps` can be set to `:noise` or `:pratio` to automatically choose the number of components using the residual frame noise or principal ratio, respectively. For more information, see the extended help.

# References
1. [Soummer, Pueyo, and Larkin (2012)](https://ui.adsabs.harvard.edu/abs/2012ApJ...755L..28S) "Detection and Characterization of Exoplanets and Disks Using Projections on Karhunen-Loève Eigenimages"

# Extended help

## Optimizing `ncomps`

There are a few ways to optimize `ncomps` using the input data. Additional options for the optimization are listed below
1. `ncomps=:noise` - residual noise optimization
2. `ncomps=:pratio` - principal ratio optimization

### Residual noise optimization

This technique progressively increases `ncomps` at each step measuring the pixel-to-pixel noise (standard deviation) in the residual data. Iteration will stop when the noise is not improving beyond a threshold. This is suited for data with similar statistical characteristics, such as an annulus more so than a full-frame cube.
* `collapse=false` - if true, the temporal median of the residual data will be used for measuring the noise.
* `noise_error=1e-3` - the threshold for the minimal noise improvement looking back 2 iterations

### Principal ratio optimization

This technique chooses the number of components required to explain some ratio of the total variance in the data. This is known as the *principal ratio* or the *explained variance ratio*. The explained variance is measured by transforming the singular values of the SVD decomposition (`Λ = @. S^2 / (n - 1)`).
* `pratio=0.9` - the target principal ratio (between 0 and 1)
"""
struct PCA <: ADIAlgorithm
    ncomps
    opts
end
PCA(ncomps; options...) = PCA(ncomps, options)
PCA(; ncomps=nothing, options...) = PCA(ncomps, options)

function fit(alg::PCA, data::AbstractMatrix; ref=data, kwargs...)
    # get number of components (using dispatch for symbolic args)
    k = get_ncomps(alg.ncomps, ref; alg.opts...)
    # fit SVD to get principal subspace of reference
    decomp = svd(collect(ref))
    # Get the principal components (principal subspace) and weights
    P = decomp.Vt[1:k, :]
    weights = data * P'
    return LinearDesign(P, weights)
end

# get ncomps using given value or num frames, whichever is smaller
get_ncomps(n::Int, data; kwargs...) = min(n, size(data, 1))
get_ncomps(::Nothing, data; kwargs...) = size(data, 1)

# get ncomps using automatic methods
function get_ncomps(s::Symbol, data; kwargs...)
    if s === :noise
        noise_decay_ncomps(data; kwargs...)
    elseif s === :pratio
        pratio_ncomps(data; kwargs...)
    else
        error("Invalid `ncomps`. Did you mean :noise or :pratio?")
    end
end

function noise_decay_ncomps(data; collapse=false, noise_error=1e-3)
    if collapse
        μ = mean(data; dims=1)
        σ2 = var(data; dims=1, mean=μ)
        X = @. (data - μ) / σ2
    else
        X = data .- mean(data, dims=1)
    end
    P = svd(X).Vt
    tmpr = similar(data)
    τ1 = τ2 = 0
    @progress name="Optimizing ncomps using residual noise" for ncomp in axes(data, 1)
        Pv = @view P[1:ncomp, :]
        tmpr .= X * (I - Pv'Pv)
        # calculate noise (standard deviation) optionally collapsing
        noise = collapse ? std(median(tmpr, dims=1)) : std(tmpr)
        # test if we've reached the noise decay tolerance
        if ncomp > firstindex(data, 1) + 2
            px_noise_decay = τ2 - noise
            @debug noise_decay=px_noise_decay noise=noise
            if px_noise_decay < noise_error
                @info "noise threshold reached with $ncomp components"
                return ncomp
            end
        end
        # update recursion variables
        τ2, τ1 = τ1, noise
    end
    return lastindex(data, 1)
end

function pratio_ncomps(data; pratio=0.9)
    @debug "Choosing ncomps required to explain $(pratio*100)% of data's temporal variance"
    X = data .- mean(data, dims=1)
    Λ = svd!(X).S
    n = length(Λ)
    exp_var = @. Λ^2 / (n - 1)
    ratio_cumsum = cumsum(exp_var ./ sum(exp_var))
    k = last(searchsorted(ratio_cumsum, pratio))
    @info "$(pratio*100)% of variance explained with $k components"
    return k
end

"""
    TPCA(;ncomps=nothing, options...)
    TPCA(ncomps; options...)

Perform principal components analysis (PCA) using truncated SVD (TSVD; provided by TSVD.jl) instead of deterministic SVD. This is often faster than [`PCA`](@ref) but is non-deterministic, so the results may be different. See the [`PCA`](@ref) docstring for more information about the options.

# See Also
* [`PCA`](@ref), [`TSVD.tsvd`](https://github.com/JuliaLinearAlgebra/TSVD.jl)
"""
struct TPCA <: ADIAlgorithm
    ncomps
    opts
end
TPCA(ncomps; options...) = TPCA(ncomps, options)
TPCA(; ncomps=nothing, options...) = TPCA(ncomps, options)

function fit(alg::TPCA, data::AbstractMatrix; ref=data, kwargs...)
    # get number of components (using dispatch for symbolic args)
    # TODO using automatic methods is not valid for TPCA using different SVD methods
    k = get_ncomps(alg.ncomps, ref; alg.opts...)
    # fit SVD to get principal subspace of reference
    U, Σ, V = tsvd(collect(ref), k)
    # Get the principal components (principal subspace) and weights
    P = V'
    weights = data * P'
    return LinearDesign(P, weights)
end


"""
    gevfitpwm(...)

Estimate the GEV parameters with the probability weighted moments.

# Implementation

Estimation with the probability weighted moments, as described by [Hosking *et al*. (1985)](https://www.tandfonline.com/doi/abs/10.1080/00401706.1985.10488049),
is only possible in the stationary case.

See also [`gevfitpwm`](@ref) for the other methods, [`gevfit`](@ref), [`gevfitbayes`](@ref) and [`BlockMaxima`](@ref).

# Reference

Hosking, J. R. M., Wallis, J. R. and Wood, E. F. (1985). Estimation of the generalized extreme-value
    distribution by the method of probability-weighted moments. *Technometrics*, 27:251-261.
"""
function gevfitpwm end

"""
    gevfitpwm(y::Vector{<:Real})

Estimate the GEV parameters with the probability weighted moments.

"""
function gevfitpwm(y::Vector{<:Real})::pwmEVA

    model = BlockMaxima(Variable("y", y))

    fittedmodel = gevfitpwm(model)

    return fittedmodel

end

"""
    gevfitpwm(df::DataFrame, datacol::Symbol)

Estimate the GEV parameters with the probability weighted moments.

Block maxima data are in the column `datacol` of the dataframe `df`.

"""
function gevfitpwm(df::DataFrame, datacol::Symbol)::pwmEVA

    model = BlockMaxima(Variable(string(datacol), df[:, datacol]))

    fittedmodel = gevfitpwm(model)

    return fittedmodel

end

"""
    gevfitpwm(model::BlockMaxima)

Estimate the GEV parameters with the probability weighted moments.

"""
function gevfitpwm(model::BlockMaxima)::pwmEVA

    model = validatestationarity(model)

    y = model.data.value

    # Computing the estimates of the probability weighted moments M_{1,q,0} for q ∈ {0,1,2}.
    b₀ = pwm(y,1,0,0)
    b₁ = pwm(y,1,1,0)
    b₂ = pwm(y,1,2,0)

    # GEV parameters estimations. Expressions retrieved from Hosking et al. (1985).
    c = (2b₁ - b₀)/(3b₂ - b₀) - log(2)/log(3)
    k = 7.859c + 2.9554c^2
    σ̂ = k *( 2b₁-b₀ ) /(1-2^(-k))/SpecialFunctions.gamma(1+k)
    μ̂ = b₀ - σ̂/k*( 1-SpecialFunctions.gamma(1+k) )

    ξ̂ = -k
    ϕ̂ = log(σ̂)

    θ̂ = [μ̂; ϕ̂; ξ̂]

    fm = pwmEVA{BlockMaxima, GeneralizedExtremeValue}(model, θ̂)

    return fm

end


function run_naive(num_sims::Integer) :: Float64

    # `x` and `y` are vectors
    x = rand(Float64, num_sims)
    y = rand(Float64, num_sims)

    # dot-syntax: broadcast
    r = sqrt.( x.^2 + y.^2 )

    area = mapreduce(
                     el -> el <= 1.0 ? 1.0 : 0.0, # map
                     +, # reduce
                     r) / num_sims
    return 4*area
end


using LatticeDiracOperators
using Gaugefields
using Test
using LinearAlgebra

@testset "LatticeDiracOperators.jl" begin
    @testset "Wilson HMC" begin
        println("Wilson HMC")
        include("wilsonhmc.jl")
        @test true
    end

    @testset "Staggered HMC" begin
        println("Staggered HMC")
        include("hmc.jl")
        @test true
    end

    @testset "Basic operations" begin
        include("basic.jl")
    # Write your tests here.
        @test true
    end
end


# ---
# title: 1296. Divide Array in Sets of K Consecutive Numbers
# id: problem1296
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Greedy
# link: <https://leetcode.com/problems/divide-array-in-sets-of-k-consecutive-numbers/description/>
# hidden: true
# ---
#
# Given an array of integers `nums` and a positive integer `k`, find whether
# it's possible to divide this array into sets of `k` consecutive numbers
# Return `True` if its possible ** ** otherwise return `False`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3,3,4,4,5,6], k = 4
#     Output: true
#     Explanation: Array can be divided into [1,2,3,4] and [3,4,5,6].
#
#
# **Example 2:**
#
#
#
#     Input: nums = [3,2,1,2,3,4,3,4,5,9,10,11], k = 3
#     Output: true
#     Explanation: Array can be divided into [1,2,3] , [2,3,4] , [3,4,5] and [9,10,11].
#
#
# **Example 3:**
#
#
#
#     Input: nums = [3,3,2,2,1,1], k = 3
#     Output: true
#
#
# **Example 4:**
#
#
#
#     Input: nums = [1,2,3,4], k = 3
#     Output: false
#     Explanation: Each array should be divided in subarrays of size 3.
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 10^5`
#   * `1 <= nums[i] <= 10^9`
#   * `1 <= k <= nums.length`
#
# **Note:** This question is the same as 846:
# <https://leetcode.comhttps://leetcode.com/problems/hand-of-straights/>
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end