text:
module ASTInterpreter2

using DebuggerFramework
using DebuggerFramework: FileLocInfo, BufferLocInfo, Suppressed
using Base.Meta
using REPL.LineEdit
using REPL
import Base: +, convert, isless
using Core: CodeInfo, SSAValue, SlotNumber, TypeMapEntry, SimpleVector, LineInfoNode, GotoNode, Slot,
            GeneratedFunctionStub, MethodInstance, NewvarNode, TypeName
using Markdown

export @enter, @make_stack, @interpret, Compiled, JuliaStackFrame

"""
`Compiled` is a trait indicating that any `:call` expressions should be evaluated
using Julia's normal compiled-code evaluation. The alternative is to pass `stack=JuliaStackFrame[]`,
which will cause all calls to be evaluated via the interpreter.
"""
struct Compiled end

"""
    JuliaProgramCounter(next_stmt::Int)

A wrapper specifying the index of the next statement in the lowered code to be executed.
"""
struct JuliaProgramCounter
    next_stmt::Int
end
+(x::JuliaProgramCounter, y::Integer) = JuliaProgramCounter(x.next_stmt+y)
convert(::Type{Int}, pc::JuliaProgramCounter) = pc.next_stmt
isless(x::JuliaProgramCounter, y::Integer) = isless(x.next_stmt, y)

Base.show(io::IO, pc::JuliaProgramCounter) = print(io, "JuliaProgramCounter(", pc.next_stmt, ')')

# A type used transiently in renumbering CodeInfo SSAValues (to distinguish a new SSAValue from an old one)
struct NewSSAValue
    id::Int
end

"""
`JuliaFrameCode` holds static information about a method or toplevel code.
One `JuliaFrameCode` can be shared by many `JuliaFrameState` calling frames.

Important fields:
- `scope`: the `Method` or `Module` in which this frame is to be evaluated
- `code`: the `CodeInfo` object storing (optimized) lowered code
- `methodtables`: a vector, each entry potentially stores a "local method table" for the corresponding
  `:call` expression in `code` (undefined entries correspond to statements that do not
  contain `:call` expressions)
- `used`: a `BitSet` storing the list of SSAValues that get referenced by later statements.
"""
struct JuliaFrameCode
    scope::Union{Method,Module}
    code::CodeInfo
    methodtables::Vector{TypeMapEntry} # line-by-line method tables for generic-function :call Exprs
    used::BitSet
    wrapper::Bool
    generator::Bool
    # Display options
    fullpath::Bool
end

function JuliaFrameCode(frame::JuliaFrameCode; wrapper = frame.wrapper, generator=frame.generator, fullpath=frame.fullpath)
    JuliaFrameCode(frame.scope, frame.code, frame.methodtables, frame.used,
                   wrapper, generator, fullpath)
end

function JuliaFrameCode(scope, code::CodeInfo; wrapper=false, generator=false, fullpath=true)
    code = optimize!(copy_codeinfo(code), moduleof(scope))
    used = find_used(code)
    methodtables = Vector{TypeMapEntry}(undef, length(code.code))
    return JuliaFrameCode(scope, code, methodtables, used, wrapper, generator, fullpath)
end

"""
`JuliaStackFrame` represents the current execution state in a particular call frame.

Important fields:
- `code`: the [`JuliaFrameCode`](@ref) for this frame
- `locals`: a vector containing the input arguments and named local variables for this frame.
  The indexing corresponds to the names in `frame.code.code.slotnames`.
- `ssavalues`: a vector containing the
  [Static Single Assignment](https://en.wikipedia.org/wiki/Static_single_assignment_form)
  values produced at the current state of execution
- `sparams`: the static type parameters, e.g., for `f(x::Vector{T}) where T` this would store
  the value of `T` given the particular input `x`.
- `pc`: the [`JuliaProgramCounter`](@ref) that typically represents the current position
  during execution. However, note that some internal functions instead maintain the `pc`
  as a local variable, and only update the frame's `pc` when pushing a frame on the stack.
"""
struct JuliaStackFrame
    code::JuliaFrameCode
    locals::Vector{Union{Nothing,Some{Any}}}
    ssavalues::Vector{Any}
    sparams::Vector{Any}
    exception_frames::Vector{Int}
    last_exception::Base.RefValue{Any}
    pc::Base.RefValue{JuliaProgramCounter}
    # A vector from names to the slotnumber of that name
    # for which a reference was last encountered.
    last_reference::Dict{Symbol,Int}
    callargs::Vector{Any}  # a temporary for processing arguments of :call exprs
end

function JuliaStackFrame(framecode::JuliaFrameCode, frame::JuliaStackFrame, pc::JuliaProgramCounter; kwargs...)
    pcref = frame.pc
    pcref[] = pc
    if !isempty(kwargs)
        framecode = JuliaFrameCode(framecode; kwargs...)
    end
    JuliaStackFrame(framecode, frame.locals,
                    frame.ssavalues, frame.sparams,
                    frame.exception_frames, frame.last_exception,
                    pcref, frame.last_reference, frame.callargs)
end

JuliaStackFrame(frame::JuliaStackFrame, pc::JuliaProgramCounter; kwargs...) =
    JuliaStackFrame(frame.code, frame, pc; kwargs...)

"""
`framedict[method]` returns the `JuliaFrameCode` for `method`. For `@generated` methods,
see [`genframedict`](@ref).
"""
const framedict = Dict{Method,JuliaFrameCode}()                # essentially a method table for lowered code
"""
`genframedict[(method,argtypes)]` returns the `JuliaFrameCode` for a `@generated` method `method`,
for the particular argument types `argtypes`.

The framecodes stored in `genframedict` are for the code returned by the generator
(i.e, what will run when you call the method on particular argument types);
for the generator itself, its framecode would be stored in [`framedict`](@ref).
"""
const genframedict = Dict{Tuple{Method,Type},JuliaFrameCode}() # the same for @generated functions

const junk = JuliaStackFrame[]      # to allow re-use of allocated memory (this is otherwise a bottleneck)

const empty_svec = Core.svec()

include("localmethtable.jl")
include("interpret.jl")
include("builtins.jl")

function moduleof(x)
    if isa(x, JuliaStackFrame)
        x = x.code.scope
    end
    return _moduleof(x)
end
_moduleof(scope::Method) = scope.module
_moduleof(scope::Module) = scope
Base.nameof(frame) = isa(frame.code.scope, Method) ? frame.code.scope.name : nameof(frame.code.scope)

is_loc_meta(expr, kind) = isexpr(expr, :meta) && length(expr.args) >= 1 && expr.args[1] === kind

function DebuggerFramework.locdesc(frame::JuliaStackFrame, specslottypes = false)
    sprint() do io
        if frame.code.scope isa Method
            meth = frame.code.scope
            argnames = frame.code.code.slotnames[2:meth.nargs]
            spectypes = Any[Any for i=1:length(argnames)]
            print(io, meth.name,'(')
            first = true
            for (argname, argT) in zip(argnames, spectypes)
                first || print(io, ", ")
                first = false
                print(io, argname)
                !(argT === Any) && print(io, "::", argT)
            end
            print(io, ") at ",
                frame.code.fullpath ? meth.file :
                basename(String(meth.file)),
                ":",meth.line)
        else
            println("not yet implemented")
        end
    end
end

function DebuggerFramework.print_locals(io::IO, frame::JuliaStackFrame)
    for i = 1:length(frame.locals)
        if !isa(frame.locals[i], Nothing)
            # #self# is only interesting if it has values inside of it. We already know
            # which function we're in otherwise.
            val = something(frame.locals[i])
            if frame.code.code.slotnames[i] == Symbol("#self#") && (isa(val, Type) || sizeof(val) == 0)
                continue
            end
            DebuggerFramework.print_var(io, frame.code.code.slotnames[i], frame.locals[i], nothing)
        end
    end
    if frame.code.scope isa Method
        for i = 1:length(frame.sparams)
            DebuggerFramework.print_var(io, frame.code.scope.sparam_syms[i], frame.sparams[i], nothing)
        end
    end
end

const SEARCH_PATH = []
function __init__()
    append!(SEARCH_PATH,[joinpath(Sys.BINDIR,"../share/julia/base/"),
            joinpath(Sys.BINDIR,"../include/")])
    return nothing
end

function loc_for_fname(file, line, defline)
    if startswith(string(file),"REPL[")
        hist_idx = parse(Int,string(file)[6:end-1])
        isdefined(Base, :active_repl) || return nothing, ""
        hp = Base.active_repl.interface.modes[1].hist
        return BufferLocInfo(hp.history[hp.start_idx+hist_idx], line, 0, defline)
    else
        for path in SEARCH_PATH
            fullpath = joinpath(path,string(file))
            if isfile(fullpath)
                return FileLocInfo(fullpath, line, 0, defline)
            end
        end
    end
    return nothing
end

function DebuggerFramework.locinfo(frame::JuliaStackFrame)
    if frame.code.scope isa Method
        meth = frame.code.scope
        loc_for_fname(meth.file, location(frame), meth.line)
    else
        println("not yet implemented")
    end
end

function DebuggerFramework.eval_code(state, frame::JuliaStackFrame, command)
    expr = Base.parse_input_line(command)
    if isexpr(expr, :toplevel)
        expr = expr.args[end]
    end
    local_vars = Any[]
    local_vals = Any[]
    for i = 1:length(frame.locals)
        if !isa(frame.locals[i], Nothing)
            push!(local_vars, frame.code.code.slotnames[i])
            push!(local_vals, QuoteNode(something(frame.locals[i])))
        end
    end
    ismeth = frame.code.scope isa Method
    for i = 1:length(frame.sparams)
        ismeth && push!(local_vars, frame.code.scope.sparam_syms[i])
        push!(local_vals, QuoteNode(frame.sparams[i]))
    end
    res = gensym()
    eval_expr = Expr(:let,
        Expr(:block, map(x->Expr(:(=), x...), zip(local_vars, local_vals))...),
        Expr(:block,
            Expr(:(=), res, expr),
            Expr(:tuple, res, Expr(:tuple, local_vars...))
        ))
    eval_res, res = Core.eval(moduleof(frame), eval_expr)
    j = 1
    for i = 1:length(frame.locals)
        if !isa(frame.locals[i], Nothing)
            frame.locals[i] = Some{Any}(res[j])
            j += 1
        end
    end
    for i = 1:length(frame.sparams)
        frame.sparams[i] = res[j]
        j += 1
    end
    eval_res
end

function maybe_quote(x)
    (isa(x, Expr) || isa(x, Symbol)) ? QuoteNode(x) : x
end

function DebuggerFramework.print_next_state(io::IO, state, frame::JuliaStackFrame)
    print(io, "About to run: ")
    expr = plain(pc_expr(frame, frame.pc[]))
    isa(expr, Expr) && (expr = copy(expr))
    if isexpr(expr, :(=))
        expr = expr.args[2]
    end
    if isexpr(expr, :call) || isexpr(expr, :return)
        expr.args = map(var->maybe_quote(@lookup(frame, var)), expr.args)
    end
    if isa(expr, Expr)
        for (i, arg) in enumerate(expr.args)
            try
                nbytes = length(repr(arg))
                if nbytes > max(40, div(200, length(expr.args)))
                    expr.args[i] = Suppressed("$nbytes bytes of output")
                end
            catch
                expr.args[i] = Suppressed("printing error")
            end
        end
    end
    print(io, expr)
    println(io)
end

const all_commands = ("q", "s", "si", "finish", "bt", "loc", "ind", "shadow",
    "up", "down", "nc", "n", "se")

function DebuggerFramework.language_specific_prompt(state, frame::JuliaStackFrame)
    if haskey(state.language_modes, :julia)
        return state.language_modes[:julia]
    end
    julia_prompt = LineEdit.Prompt(DebuggerFramework.promptname(state.level, "julia");
        # Copy colors from the prompt object
        prompt_prefix = state.repl.prompt_color,
        prompt_suffix = (state.repl.envcolors ? Base.input_color : state.repl.input_color),
        complete = REPL.REPLCompletionProvider(),
        on_enter = REPL.return_callback)
    # 0.7 compat
    if isdefined(state.main_mode, :repl)
        julia_prompt.repl = state.main_mode.repl
    end
    julia_prompt.hist = state.main_mode.hist
    julia_prompt.hist.mode_mapping[:julia] = julia_prompt

    julia_prompt.on_done = (s,buf,ok)->begin
        if !ok
            LineEdit.transition(s, :abort)
            return false
        end
        xbuf = copy(buf)
        command = String(take!(buf))
        ok, result = DebuggerFramework.eval_code(state, command)
        REPL.print_response(state.repl, ok ? result : result[1], ok ? nothing : result[2], true, true)
        println(state.repl.t)

        if !ok
            # Convenience hack. We'll see if this is more useful or annoying
            for c in all_commands
                !startswith(command, c) && continue
                LineEdit.transition(s, state.main_mode)
                LineEdit.state(s, state.main_mode).input_buffer = xbuf
                break
            end
        end
        LineEdit.reset_state(s)
    end
    julia_prompt.keymap_dict = LineEdit.keymap([REPL.mode_keymap(state.main_mode);state.standard_keymap])
    state.language_modes[:julia] = julia_prompt
    return julia_prompt
end

function DebuggerFramework.debug(meth::Method, args...)
    stack = [enter_call(meth, args...)]
    DebuggerFramework.RunDebugger(stack)
end

function to_function(x)
    if isa(x, GlobalRef)
        getfield(x.mod, x.name)
    else
        x
    end
end

_Typeof(x) = isa(x,Type) ? Type{x} : typeof(x)

function is_function_def(ex)
    (isexpr(ex, :(=)) && isexpr(ex.args[1], :call)) ||
    isexpr(ex,:function)
end

function is_generated(meth)
    isdefined(meth, :generator)
end

function namedtuple(kwargs)
    names, types, vals = Symbol[], [], []
    for pr in kwargs
        if isa(pr, Expr)
            push!(names, pr.args[1])
            val = pr.args[2]
            push!(types, typeof(val))
            push!(vals, val)
        elseif isa(pr, Pair)
            push!(names, pr.first)
            val = pr.second
            push!(types, typeof(val))
            push!(vals, val)
        else
            error("unhandled entry type ", typeof(pr))
        end
    end
    return NamedTuple{(names...,), Tuple{types...}}(vals)
end

"""
    frun, allargs = prepare_args(fcall, fargs, kwargs)

Prepare the complete argument sequence for a call to `fcall`. `fargs = [fcall, args...]` is a list
containing both `fcall` (the `#self#` slot in lowered code) and the positional
arguments supplied to `fcall`. `kwargs` is a list of keyword arguments, supplied either as
list of expressions `:(kwname=kwval)` or pairs `:kwname=>kwval`.

For non-keyword methods, `frun === fcall`, but for methods with keywords `frun` will be the
keyword-sorter function for `fcall`.

# Example

```jldoctest; setup=(using ASTInterpreter2; empty!(ASTInterpreter2.junk))
julia> mymethod(x) = 1
mymethod (generic function with 1 method)

julia> mymethod(x, y; verbose=false) = nothing
mymethod (generic function with 2 methods)

julia> ASTInterpreter2.prepare_args(mymethod, [mymethod, 15], ())
(mymethod, Any[mymethod, 15])

julia> ASTInterpreter2.prepare_args(mymethod, [mymethod, 1, 2], [:verbose=>true])
(getfield( Symbol("#kw##mymethod"))(), Any[#kw##mymethod(), Any[:verbose, true], mymethod, 1, 2])
```
"""
function prepare_args(f, allargs, kwargs)
    if !isempty(kwargs)
        of = f
        f = Core.kwfunc(f)
        allargs = [f,namedtuple(kwargs),allargs...]
    elseif f === Core._apply
        f = to_function(allargs[2])
        allargs = Base.append_any((allargs[2],), allargs[3:end]...)
    end
    return f, allargs
end

"""
    framecode, frameargs, lenv, argtypes = prepare_call(f, allargs; enter_generated=false)

Prepare all the information needed to execute lowered code for `f` given arguments `allargs`.
`f` and `allargs` are the outputs of [`prepare_args`](@ref).
For `@generated` methods, set `enter_generated=true` if you want to extract the lowered code
of the generator itself.

On return `framecode` is the [`JuliaFrameCode`](@ref) of the method.
`frameargs` contains the actual arguments needed for executing this frame (for generators,
this will be the types of `allargs`);
`lenv` is the "environment", i.e., the static parameters for `f` given `allargs`.
`argtypes` is the `Tuple`-type for this specific call (equivalent to the signature of the `MethodInstance`).

# Example

```jldoctest; setup=(using ASTInterpreter2; empty!(ASTInterpreter2.junk))
julia> mymethod(x::Vector{T}) where T = 1
mymethod (generic function with 1 method)

julia> framecode, frameargs, lenv, argtypes = ASTInterpreter2.prepare_call(mymethod, [mymethod, [1.0,2.0]]);

julia> framecode
ASTInterpreter2.JuliaFrameCode(mymethod(x::Array{T,1}) where T in Main at none:1, CodeInfo(
1 ─     return 1
), Core.TypeMapEntry[#undef], BitSet([]), false, false, true)

julia> frameargs
2-element Array{Any,1}:
 mymethod
 [1.0, 2.0]

julia> lenv
svec(Float64)

julia> argtypes
Tuple{typeof(mymethod),Array{Float64,1}}
```
"""
function prepare_call(f, allargs; enter_generated = false)
    args = allargs[2:end]
    argtypes = Tuple{map(_Typeof,args)...}
    method = try
        which(f, argtypes)
    catch err
        @show typeof(f)
        println(f)
        println(argtypes)
        rethrow(err)
    end
    argtypes = Tuple{_Typeof(f), argtypes.parameters...}
    args = allargs
    sig = method.sig
    isa(method, TypeMapEntry) && (method = method.func)
    # Get static parameters
    (ti, lenv::SimpleVector) = ccall(:jl_type_intersection_with_env, Any, (Any, Any),
                        argtypes, sig)::SimpleVector
    enter_generated &= is_generated(method)
    if is_generated(method) && !enter_generated
        framecode = get(genframedict, (method, argtypes), nothing)
    else
        framecode = get(framedict, method, nothing)
    end
    if framecode === nothing
        if is_generated(method) && !enter_generated
            # If we're stepping into a staged function, we need to use
            # the specialization, rather than stepping through the
            # unspecialized method.
            code = Core.Compiler.get_staged(Core.Compiler.code_for_method(method, argtypes, lenv, typemax(UInt), false))
            generator = false
        else
            if is_generated(method)
                args = Any[_Typeof(a) for a in args]
                code = get_source(method.generator)
                generator = true
            else
                code = get_source(method)
                generator = false
            end
        end
        framecode = JuliaFrameCode(method, code; generator=generator)
        if is_generated(method) && !enter_generated
            genframedict[(method, argtypes)] = framecode
        else
            framedict[method] = framecode
        end
    end
    return framecode, args, lenv, argtypes
end

function prepare_thunk(mod::Module, thunk::Expr)
    if isexpr(thunk, :thunk)
        framecode = JuliaFrameCode(mod, thunk.args[1])
    elseif isexpr(thunk, :error)
        error("lowering returned an error, ", thunk)
    elseif isa(thunk, Expr)
        error("expected thunk or module, got ", thunk.head)
    else
        error("expected expression, got ", typeof(thunk))
    end
    return prepare_locals(framecode, [])
end

function prepare_toplevel(mod::Module, expr::Expr)
    if expr.head == :using || expr.head == :import || expr.head == :export
        Core.eval(mod, expr)
        return nothing
    else
        thunk = Meta.lower(mod, expr)
    end
    return prepare_thunk(mod, thunk)
end

"""
    framecode, frameargs, lenv, argtypes = determine_method_for_expr(expr; enter_generated = false)

Prepare all the information needed to execute a particular `:call` expression `expr`.
For example, try `ASTInterpreter2.determine_method_for_expr(:(sum([1,2])))`.
See [`ASTInterpreter2.prepare_call`](@ref) for information about the outputs.
"""
function determine_method_for_expr(expr; enter_generated = false)
    f = to_function(expr.args[1])
    allargs = expr.args
    # Extract keyword args
    local kwargs = Expr(:parameters)
    if length(allargs) > 1 && isexpr(allargs[2], :parameters)
        kwargs = splice!(allargs, 2)
    end
    f, allargs = prepare_args(f, allargs, kwargs.args)
    # Can happen for thunks created by generated functions
    if !isa(f, Core.Builtin) && !isa(f, Core.IntrinsicFunction)
        return prepare_call(f, allargs; enter_generated=enter_generated)
    end
    nothing
end

function get_source(meth)
    if isa(meth.source, Array{UInt8,1})
        return ccall(:jl_uncompress_ast, Any, (Any, Any), meth, meth.source)
    else
        return meth.source
    end
end

function get_source(g::GeneratedFunctionStub)
    b = g(g.argnames...)
    b isa CodeInfo && return b
    return eval(b)
end

function copy_codeinfo(code::CodeInfo)
    newcode = ccall(:jl_new_struct_uninit, Any, (Any,), CodeInfo)::CodeInfo
    for (i, name) in enumerate(fieldnames(CodeInfo))
        if isdefined(code, name)
            val = getfield(code, name)
            ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), newcode, i-1, val===nothing ? val : copy(val))
        end
    end
    return newcode
end

const calllike = Set([:call, :struct_type])

function extract_inner_call!(stmt, idx, once::Bool=false)
    isa(stmt, Expr) || return nothing
    once |= stmt.head ∈ calllike
    for (i, a) in enumerate(stmt.args)
        isa(a, Expr) || continue
        ret = extract_inner_call!(a, idx, once) # doing this first extracts innermost calls
        ret !== nothing && return ret
        iscalllike = a.head ∈ calllike
        if once && iscalllike
            stmt.args[i] = NewSSAValue(idx)
            return a
        end
    end
    return nothing
end

function replace_ssa!(stmt, ssalookup)
    isa(stmt, Expr) || return nothing
    for (i, a) in enumerate(stmt.args)
        if isa(a, SSAValue)
            stmt.args[i] = SSAValue(ssalookup[a.id])
        elseif isa(a, NewSSAValue)
            stmt.args[i] = SSAValue(a.id)
        else
            replace_ssa!(a, ssalookup)
        end
    end
    return nothing
end

function renumber_ssa!(stmts::Vector{Any}, ssalookup)
    for (i, stmt) in enumerate(stmts)
        if isa(stmt, GotoNode)
            new_code[i] = GotoNode(ssalookup[stmt.label])
        elseif isa(stmt, SSAValue)
            new_code[i] = SSAValue(ssalookup[stmt.id])
        elseif isa(stmt, NewSSAValue)
            new_code[i] = SSAValue(stmt.id)
        elseif isa(stmt, Expr)
            replace_ssa!(stmt, ssalookup)
            if stmt.head == :gotoifnot && isa(stmt.args[2], Int)
                stmt.args[2] = ssalookup[stmt.args[2]]
            end
        end
    end
    return stmts
end

function lookup_global_refs!(ex::Expr)
    for (i, a) in enumerate(ex.args)
        if isa(a, GlobalRef)
            r = getfield(a.mod, a.name)
            ex.args[i] = QuoteNode(r)
        elseif isa(a, Expr)
            lookup_global_refs!(a)
        end
    end
end

"""
    optimize!(code::CodeInfo, mod::Module)

Perform minor optimizations on the lowered AST in `code` to reduce execution time
of the interpreter.
Currently it looks up `GlobalRef`s (for which it needs `mod` to know the scope in
which this will run) and ensures that no statement includes nested `:call` expressions
(splitting them out into multiple SSA-form statements if needed).
"""
function optimize!(code::CodeInfo, mod::Module)
    code.inferred && error("optimization of inferred code not implemented")
    # TODO: because of builtins.jl, for CodeInfos like
    #   %1 = Core.apply_type
    #   %2 = (%1)(args...)
    # it would be best to *not* resolve the GlobalRef at %1

    ## Replace GlobalRefs with QuoteNodes
    for (i, stmt) in enumerate(code.code)
        if isa(stmt, GlobalRef)
            code.code[i] = QuoteNode(getfield(stmt.mod, stmt.name))
        elseif isa(stmt, Expr)
            if stmt.head == :call && isa(stmt.args[1], GlobalRef)
                # Special handling of cglobal, which requires constants for its arguments
                r = stmt.args[1]::GlobalRef
                f = getfield(r.mod, r.name)
                if f === Base.cglobal
                    code.code[i] = QuoteNode(Core.eval(mod, stmt))
                else
                    lookup_global_refs!(stmt)
                end
            else
                lookup_global_refs!(stmt)
            end
        end
    end

    # ## Un-nest :call expressions (so that there will be only one :call per line)
    # # This will allow us to re-use args-buffers rather than having to allocate new ones each time.
    # old_code, old_codelocs = code.code, code.codelocs
    # code.code = new_code = eltype(old_code)[]
    # code.codelocs = new_codelocs = Int32[]
    # ssainc = fill(1, length(old_code))
    # for (i, stmt) in enumerate(old_code)
    #     loc = old_codelocs[i]
    #     inner = extract_inner_call!(stmt, length(new_code)+1)
    #     while inner !== nothing
    #         push!(new_code, inner)
    #         push!(new_codelocs, loc)
    #         ssainc[i] += 1
    #         inner = extract_inner_call!(stmt, length(new_code)+1)
    #     end
    #     push!(new_code, stmt)
    #     push!(new_codelocs, loc)
    # end
    # # Fix all the SSAValues and GotoNodes
    # ssalookup = cumsum(ssainc)
    # renumber_ssa!(new_code, ssalookup)
    # code.ssavaluetypes = length(new_code)
    return code
end

plain(stmt) = stmt

function prepare_locals(framecode, argvals::Vector{Any})
    if isa(framecode.scope, Method)
        meth, code = framecode.scope::Method, framecode.code
        ssavt = code.ssavaluetypes
        ng = isa(ssavt, Int) ? ssavt : length(ssavt::Vector{Any})
        nargs = length(argvals)
        if !isempty(junk)
            oldframe = pop!(junk)
            locals, ssavalues, sparams = oldframe.locals, oldframe.ssavalues, oldframe.sparams
            exception_frames, last_reference = oldframe.exception_frames, oldframe.last_reference
            callargs = oldframe.callargs
            last_exception, pc = oldframe.last_exception, oldframe.pc
            # for check_isdefined to work properly, we need locals and sparams to start out unassigned
            resize!(resize!(locals, 0), length(code.slotflags))
            resize!(ssavalues, ng)
            resize!(resize!(sparams, 0), length(meth.sparam_syms))
            empty!(exception_frames)
            empty!(last_reference)
            last_exception[] = nothing
            pc[] = JuliaProgramCounter(1)
        else
            locals = Vector{Union{Nothing,Some{Any}}}(undef, length(code.slotflags))
            ssavalues = Vector{Any}(undef, ng)
            sparams = Vector{Any}(undef, length(meth.sparam_syms))
            exception_frames = Int[]
            last_reference = Dict{Symbol,Int}()
            callargs = Any[]
            last_exception = Ref{Any}(nothing)
            pc = Ref(JuliaProgramCounter(1))
        end
        for i = 1:meth.nargs
            if meth.isva && i == meth.nargs
                locals[i] = nargs < i ? Some{Any}(()) : (let i=i; Some{Any}(ntuple(k->argvals[i+k-1], nargs-i+1)); end)
                break
            end
            locals[i] = nargs >= i ? Some{Any}(argvals[i]) : Some{Any}(())
        end
        # add local variables initially undefined
        for i = (meth.nargs+1):length(code.slotnames)
            locals[i] = nothing
        end
    else
        code = framecode.code
        locals = Vector{Union{Nothing,Some{Any}}}(undef, length(code.slotflags))
        ssavalues = Vector{Any}(undef, length(code.code))
        sparams = Any[]
        exception_frames = Int[]
        last_reference = Dict{Symbol,Int}()
        callargs = Any[]
        last_exception = Ref{Any}(nothing)
        pc = Ref(JuliaProgramCounter(1))
    end
    JuliaStackFrame(framecode, locals, ssavalues, sparams, exception_frames, last_exception,
                    pc, last_reference, callargs)
end

"""
    frame = build_frame(framecode::JuliaFrameCode, frameargs, lenv)

Construct a new `JuliaStackFrame` for `framecode`, given lowered-code arguments `frameargs` and
static parameters `lenv`. See [`ASTInterpreter2.prepare_call`](@ref) for information about how to prepare the inputs.
"""
function build_frame(framecode, args, lenv)
    frame = prepare_locals(framecode, args)
    # Add static parameters to environment
    for i = 1:length(lenv)
        T = lenv[i]
        isa(T, TypeVar) && continue  # only fill concrete types
        frame.sparams[i] = T
    end
    return frame
end

"""
    frame = enter_call_expr(expr; enter_generated=false)

Build a `JuliaStackFrame` ready to execute the expression `expr`. Set `enter_generated=true`
if you want to execute the generator of a `@generated` function, rather than the code that
would be created by the generator.

# Example

```jldoctest; setup=(using ASTInterpreter2; empty!(ASTInterpreter2.junk))
julia> mymethod(x) = x+1
mymethod (generic function with 1 method)

julia> ASTInterpreter2.enter_call_expr(:(\$mymethod(1)))
JuliaStackFrame(ASTInterpreter2.JuliaFrameCode(mymethod(x) in Main at none:1, CodeInfo(
1 ─ %1 = (\$(QuoteNode(+)))(x, 1)
└──      return %1
), Core.TypeMapEntry[#undef, #undef], BitSet([1]), false, false, true), Union{Nothing, Some{Any}}[Some(mymethod), Some(1)], Any[#undef, #undef], Any[], Int64[], Base.RefValue{Any}(nothing), Base.RefValue{ASTInterpreter2.JuliaProgramCounter}(JuliaProgramCounter(1)), Dict{Symbol,Int64}(), Any[])

julia> mymethod(x::Vector{T}) where T = 1
mymethod (generic function with 2 methods)

julia> a = [1.0, 2.0]
2-element Array{Float64,1}:
 1.0
 2.0

julia> ASTInterpreter2.enter_call_expr(:(\$mymethod(\$a)))
JuliaStackFrame(ASTInterpreter2.JuliaFrameCode(mymethod(x::Array{T,1}) where T in Main at none:1, CodeInfo(
1 ─     return 1
), Core.TypeMapEntry[#undef], BitSet([]), false, false, true), Union{Nothing, Some{Any}}[Some(mymethod), Some([1.0, 2.0])], Any[#undef], Any[Float64], Int64[], Base.RefValue{Any}(nothing), Base.RefValue{ASTInterpreter2.JuliaProgramCounter}(JuliaProgramCounter(1)), Dict{Symbol,Int64}(), Any[])
```

See [`enter_call`](@ref) for a similar approach not based on expressions.
"""
function enter_call_expr(expr; enter_generated = false)
    r = determine_method_for_expr(expr; enter_generated = enter_generated)
    if r !== nothing
        return build_frame(r[1:end-1]...)
    end
    nothing
end

"""
    frame = enter_call(f, args...; kwargs...)

Build a `JuliaStackFrame` ready to execute `f` with the specified positional and keyword arguments.

# Example

```jldoctest; setup=(using ASTInterpreter2; empty!(ASTInterpreter2.junk))
julia> mymethod(x) = x+1
mymethod (generic function with 1 method)

julia> ASTInterpreter2.enter_call(mymethod, 1)
JuliaStackFrame(ASTInterpreter2.JuliaFrameCode(mymethod(x) in Main at none:1, CodeInfo(
1 ─ %1 = ($(QuoteNode(+)))(x, 1)
└──      return %1
), Core.TypeMapEntry[#undef, #undef], BitSet([1]), false, false, true), Union{Nothing, Some{Any}}[Some(mymethod), Some(1)], Any[#undef, #undef], Any[], Int64[], Base.RefValue{Any}(nothing), Base.RefValue{ASTInterpreter2.JuliaProgramCounter}(JuliaProgramCounter(1)), Dict{Symbol,Int64}(), Any[])

julia> mymethod(x::Vector{T}) where T = 1
mymethod (generic function with 2 methods)

julia> ASTInterpreter2.enter_call(mymethod, [1.0, 2.0])
JuliaStackFrame(ASTInterpreter2.JuliaFrameCode(mymethod(x::Array{T,1}) where T in Main at none:1, CodeInfo(
1 ─     return 1
), Core.TypeMapEntry[#undef], BitSet([]), false, false, true), Union{Nothing, Some{Any}}[Some(mymethod), Some([1.0, 2.0])], Any[#undef], Any[Float64], Int64[], Base.RefValue{Any}(nothing), Base.RefValue{ASTInterpreter2.JuliaProgramCounter}(JuliaProgramCounter(1)), Dict{Symbol,Int64}(), Any[])
```

For a `@generated` function you can use `enter_call((f, true), args...; kwargs...)`
to execute the generator of a `@generated` function, rather than the code that
would be created by the generator.

See [`enter_call_expr`](@ref) for a similar approach based on expressions.
"""
function enter_call(@nospecialize(finfo), @nospecialize(args...); kwargs...)
    if isa(finfo, Tuple)
        f = finfo[1]
        enter_generated = finfo[2]::Bool
    else
        f = finfo
        enter_generated = false
    end
    f, allargs = prepare_args(f, Any[f, args...], kwargs)
    # Can happen for thunks created by generated functions
    if isa(f, Core.Builtin) || isa(f, Core.IntrinsicFunction)
        error(f, " is a builtin or intrinsic")
    end
    r = prepare_call(f, allargs; enter_generated=enter_generated)
    if r !== nothing
        return build_frame(r[1:end-1]...)
    end
    return nothing
end

function maybe_step_through_wrapper!(stack)
    length(stack[1].code.code.code) < 2 && return stack
    last = plain(stack[1].code.code.code[end-1])
    isexpr(last, :(=)) && (last = last.args[2])
    stack1 = stack[1]
    is_kw = stack1.code.scope isa Method && startswith(String(Base.unwrap_unionall(stack1.code.scope.sig).parameters[1].name.name), "#kw")
    if is_kw || isexpr(last, :call) && any(x->x==SlotNumber(1), last.args)
        # If the last expr calls #self# or passes it to an implementation method,
        # this is a wrapper function that we might want to step through
        frame = stack1
        pc = frame.pc[]
        while pc != JuliaProgramCounter(length(frame.code.code.code)-1)
            pc = next_call!(Compiled(), frame, pc)
        end
        stack[1] = JuliaStackFrame(JuliaFrameCode(frame.code; wrapper=true), frame, pc)
        newcall = Expr(:call, map(x->@lookup(frame, x), last.args)...)
        pushfirst!(stack, enter_call_expr(newcall))
        return maybe_step_through_wrapper!(stack)
    end
    stack
end

lower(mod, arg) = false ? expand(arg) : Meta.lower(mod, arg)

cant_be_lowered(head) = head == :module || head == :toplevel

function checkfor_head(f, arg)
    isa(arg, Expr) || return false
    if f(arg.head)
        return true
    end
    for a in arg.args
        if isa(a, Expr)
            checkfor_head(f, a) && return true
        end
    end
    return false
end

"""
    interpret!(stack, mod::Module, expr::Expr)

Interpret `expr` at top-level in `mod`. This should be used for handling expressions
that define structures, methods, or otherwise look like code that defines modules.
"""
function interpret!(stack, mod::Module, expr::Expr)
    if expr.head != :toplevel
        expr = Expr(:toplevel, expr)
    end
    local ret
    for arg in expr.args
        isa(arg, Expr) || continue
        if checkfor_head(cant_be_lowered, arg)
            head = arg.head
            if head == :module
                # Create the module
                innermod = Core.eval(mod, Expr(:module, arg.args[1], arg.args[2], quote end))
                # Interpret the module body at toplevel inside the new module
                modbody = Expr(:toplevel)
                modbody.args = arg.args[3].args
                ret = interpret!(stack, innermod, modbody)
            elseif head == :toplevel
                ret = interpret!(stack, mod, arg)
            elseif head == :block
                for a in arg.args
                    if isa(a, Expr)
                        ret = interpret!(stack, mod, a)
                    end
                end
            else
                error("fixme unhandled ", arg)
            end
        elseif isa(arg, Expr)
            frame = ASTInterpreter2.prepare_toplevel(mod, arg)
            frame === nothing && continue
            ret = ASTInterpreter2.finish_and_return!(stack, frame, true)
        end
    end

    return ret
end

# This is a version of gen_call_with_extracted_types, except that is passes back the call expression
# for further processing.
function extract_args(__module__, ex0)
    if isa(ex0, Expr)
        kws = collect(filter(x->isexpr(x,:kw),ex0.args))
        if !isempty(kws)
            names = []
            values = Tuple(map(x-> begin
                push!(names,x.args[1])
                x.args[2]
            end,kws))
            names = Tuple(names)
            return Expr(:tuple,:(Core.kwfunc($(ex0.args[1]))),
                Expr(:call, NamedTuple{names,typeof(values)}, values),
                map(x->isexpr(x, :parameters) ? QuoteNode(x) : x,
                filter(x->!isexpr(x, :kw),ex0.args))...)
        else
            return Expr(:tuple,
                map(x->isexpr(x,:parameters) ? QuoteNode(x) : x, ex0.args)...)
        end
    end
    if isa(ex0, Expr) && ex0.head == :macrocall # Make @edit @time 1+2 edit the macro by using the types of the *expressions*
        return error("Macros are not supported in @enter")
    end
    ex = lower(__module__, ex0)
    exret = Expr(:none)
    if !isa(ex, Expr)
        return error("expression is not a function call or symbol")
    elseif ex.head == :call
        return Expr(:tuple,
                map(x->isexpr(x,:parameters) ? QuoteNode(x) : x, ex.args)...)
    elseif ex.head == :body
        a1 = ex.args[1]
        if isa(a1, Expr) && a1.head == :call
            a11 = a1.args[1]
            if a11 == :setindex!
                return Expr(:tuple,
                map(x->isexpr(x,:parameters) ? QuoteNode(x) : x, arg.args)...)
            end
        end
    end
    return error("expression is not a function call, "
               * "or is too complex for @enter to analyze; "
               * "break it down to simpler parts if possible")
end

function _make_stack(mod, arg)
    args = try
        extract_args(mod, arg)
    catch e
        return :(throw($e))
    end
    quote
        theargs = $(esc(args))
        stack = [ASTInterpreter2.enter_call_expr(Expr(:call,theargs...))]
        ASTInterpreter2.maybe_step_through_wrapper!(stack)
        stack[1] = ASTInterpreter2.JuliaStackFrame(stack[1], ASTInterpreter2.maybe_next_call!(Compiled(), stack[1]))
        stack
    end
end

macro make_stack(arg)
    _make_stack(__module__, arg)
end

macro enter(arg)
    quote
        let stack = $(_make_stack(__module__,arg))
            DebuggerFramework.RunDebugger(stack)
        end
    end
end

"""
    @interpret f(args; kwargs...)

Evaluate `f` on the specified arguments using the interpreter.

# Example

```jldoctest; setup=(using ASTInterpreter2; empty!(ASTInterpreter2.junk))
julia> a = [1, 7]
2-element Array{Int64,1}:
 1
 7

julia> sum(a)
8

julia> @interpret sum(a)
8
```
"""
macro interpret(arg)
    args = try
        extract_args(__module__, arg)
    catch e
        return :(throw($e))
    end
    quote
        theargs = $(esc(args))
        stack = JuliaStackFrame[]
        frame = ASTInterpreter2.enter_call_expr(Expr(:call,theargs...))
        empty!(framedict)  # start fresh each time; kind of like bumping the world age at the REPL prompt
        empty!(genframedict)
        finish_and_return!(stack, frame)
    end
end

include("commands.jl")

end # module


text:
########## Parameters for scripts ##########

###### bspline-MLE ######

#### cross-scale only ####

# sfct = 1 # factor of scale
# vm = 1  # vanishing moments

cscl = Int(Npd*sfct)  # central scale
wscl = 10  # half scale window width, e.g. <=10
minscl, maxscl = max(1, cscl-wscl), cscl+wscl
# maxscl = Npd - 0*Npd÷4
# # minscl = maxscl ÷ 2
# minscl = maxscl - Npd÷2

wfct = 40  # factor of moving window

parms = (
    estimator = "bspline_MLE_scale",
    method = :optim,  # method of optimizer

    # Parameters for B-Spline wavelet transform
    scalefactor = sfct,
    sclrng = 2*(minscl:maxscl),  # index range of inertial scales
    vm = 1,  # vanishing moment of the wavelet
    fmode = :causal,  # mode of filtration
    ρ = 0,  # threshold for truncation of small eigen values in the covariance matrix
    partial = true,  # true for whittle estimator. The case partial=false cannot be handled correctly for now.

    # Parameters for moving window
    windowfactor = wfct,
    wsize = cscl*wfct,  # size of moving window
    pov = Npd÷1,  # step of moving window
    boundary = :hard,  # boundary condition
)

base_estimator = x -> FracFin.bspline_MLE_estim(x, parms.sclrng, parms.vm; ρ=parms.ρ, partial=parms.partial)

ext: jl

max_stars_repo_name: yanncalec/FracFin

max_issues_repo_name: yanncalec/FracFin

max_forks_repo_name: yanncalec/FracFin


text:
__precompile__(true)
module SymEngineExtras

using SymEngine
using Base.Iterators
using ResumableFunctions
using Multisets
using Combinatorics

SymbolicType = SymEngine.SymbolicType
BasicType = SymEngine.BasicType

# package code goes here
include("iterator-utils.jl")
include("patternmatch.jl")
include("simplify.jl")

end # module


text:
using Rimu
using Rimu.Hamiltonians: hopnextneighbour

const BoseOccupiedModeMap{N} = OccupiedModeMap{N,BitStringAddresses.BoseFSIndex}
const FermiOccupiedModeMap{N} = OccupiedModeMap{N,BitStringAddresses.FermiFSIndex}

"""
    momentum_transfer_excitation(add, chosen, map; fold=true) -> nadd, α, p, q, -k
    momentum_transfer_excitation(add_a, add_b, chosen, map_a, map_b; fold=true)

Apply the momentum transfer operator to Fock address (or pair of addresses) `add` (or
`add_a`, `add_b`):

```math
a^†_{p + k} a^†_{q - k} a_q a_p |\\mathtt{add}⟩
```

The `fold` argument controls whether the terms `p + k` and `q - k` are done modulo M. If
not, zero is returned when either of those terms is less than 1 or larger than M.
It is expected that `map == OccupiedModeMap(add)`.

Return the new address(es), the value, modes `p` and `q`, and the momentum change `-k`.

See [`excitation`](@ref), [`OccupiedModeMap`](@ref).
"""
@inline function momentum_transfer_excitation(
    add::BoseFS, chosen, map::OccupiedModeMap; fold=true
)
    M = num_modes(add)
    singlies = length(map)
    doublies = count(i -> i.occnum ≥ 2, map)

    double = chosen - singlies * (singlies - 1) * (M - 2)

    if double > 0
        # Both moves from the same mode.
        double, mom_change = fldmod1(double, M - 1)
        idx = first(map) # placeholder
        for i in map
            double -= i.occnum ≥ 2
            if double == 0
                idx = i
                break
            end
        end
        src_indices = (idx, idx)
    else
        # Moves from different modes.
        pair, mom_change = fldmod1(chosen, M - 2)
        fst, snd = fldmod1(pair, singlies - 1) # where the holes are to be made
        if snd < fst # put them in ascending order
            f_hole = snd
            s_hole = fst
        else
            f_hole = fst
            s_hole = snd + 1 # as we are counting through all singlies
        end
        src_indices = getindex.(Ref(map), (f_hole, s_hole))
        f_mode, s_mode = src_indices[1].mode, src_indices[2].mode
        if mom_change ≥ s_mode - f_mode
            mom_change += 1 # to avoid putting particles back into the holes
        end
    end
    # For higher dimensions, replace mod1 here with some geometry.
    src_modes = getproperty.(src_indices, :mode)
    dst_modes = src_modes .+ (mom_change, -mom_change)
    if fold
        dst_modes = mod1.(dst_modes, M)
    elseif !all(1 .≤ dst_modes .≤ M)
        return add, 0.0, src_modes..., -mom_change
    end
    dst_indices = find_mode(add, dst_modes)
    return excitation(add, dst_indices, src_indices)..., src_modes..., -mom_change
end

function momentum_transfer_excitation(add::FermiFS, chosen::Integer, map; fold=true)
    return add, 0.0, 0, 0, 0
end

@inline function momentum_transfer_excitation(
    add_a, add_b, chosen, map_a, map_b; fold=true
)
    M = num_modes(add_a)
    M == num_modes(add_b) || throw(ArgumentError("Addresses must have the same number of modes"))

    src_a, remainder = fldmod1(chosen, (M - 1) * length(map_b))
    dst_a, src_b = fldmod1(remainder, length(map_b))

    src_a_index = map_a[src_a]
    src_b_index = map_b[src_b]
    src_a_mode = src_a_index.mode

    if dst_a ≥ src_a_mode
        dst_a += 1 # to skip the src_a
    end
    mom_change = dst_a - src_a_mode # change in momentun
    dst_a_index = find_mode(add_a, dst_a)
    src_b_mode = src_b_index.mode
    dst_b = src_b_mode - mom_change

    # Additional info returned with result:
    params = (src_a_mode, src_b_mode, -mom_change)

    if fold
        dst_a = mod1(dst_a, M)
        dst_b = mod1(dst_b, M) # enforce periodic boundary condition
    elseif !(0 < dst_a ≤ M) || !(0 < dst_b ≤ M)
        return add_a, add_b, 0.0, params...
    end

    dst_b_index = find_mode(add_b, dst_b)

    new_add_a, val_a = excitation(add_a, (dst_a_index,), (src_a_index,))
    new_add_b, val_b = excitation(add_b, (dst_b_index,), (src_b_index,))

    return new_add_a, new_add_b, val_a * val_b, params...
end

"""
    momentum_transfer_diagonal(map)

The diagonal part of [`momentum_transfer_excitation`](@ref).
"""
function momentum_transfer_diagonal(map::BoseOccupiedModeMap)
    onproduct = 0
    for i in 1:length(map)
        occ_i = map[i].occnum
        onproduct += occ_i * (occ_i - 1)
        for j in 1:i-1
            occ_j = map[j].occnum
            onproduct += 4 * occ_i * occ_j
        end
    end
    return float(onproduct)
end
function momentum_transfer_diagonal(
    map_a::FermiOccupiedModeMap, map_b::FermiOccupiedModeMap
)
    onproduct = 0
    n1 = length(map_a)
    n2 = length(map_b)

    return float(2 * n1 * n2)
end

"""
    transcorrelated_three_body_excitation(add↑, add↓, i, map↑, map↓)
    -> nadd↑, nadd↓, value, k, l

Apply the following operator to two addresses:

```math
a^†_{p+k,↑} a^†_{q+l,↑} a^†_{s-k-l,↓} a_{s,↓} a_{q,↑} a_{p,↑} |\\mathtt{add↑}⟩⊗|\\mathtt{add↓}⟩
```

The index `i` enumerates the possible non-zero terms and determines ``p, q, s, k, l``.
It is expected that `map↑, map↓ == OccupiedModeMap(add↑), OccupiedModeMap(add↓)`.

Return new addresses, prefactor `value`, `k`, and `l`. Note: If either `k` or `l` are zero,
or the excitation is diagonal, the function returns `value == 0`.

See [`transcorrelated_diagonal`](@ref), [`Transcorrelated1D`](@ref).
"""
function transcorrelated_three_body_excitation(add_a, add_b, i, map_a, map_b)
    N1 = length(map_a)
    N2 = length(map_b)
    M = num_modes(add_a)
    M == num_modes(add_b) || throw(ArgumentError("Addresses must have the same number of modes"))

    if add_a isa FermiFS # TODO: this is better done in the same way as momentum transfer
        p, q, s, p_k, q_l = Tuple(CartesianIndices((N1, N1 - 1, N2, M, M))[i])
        if q ≥ p
            q += 1
        end
    else
        p, q, s, p_k, q_l = Tuple(CartesianIndices((N1, N1, N2, M, M))[i])
    end

    p_index = map_a[p]
    q_index = map_a[q]
    s_index = map_b[s]

    k = p_index.mode - p_k
    l = q_l - q_index.mode
    s_kl = s_index.mode + k - l

    if k == 0 || l == 0
        # Zero because Q_kl == 0
        return add_a, add_b, 0.0, k,l
    elseif p_index.mode == q_l && q_index.mode == p_k
        # Diagonal
        return add_a, add_b, 0.0, k,l
    elseif s_kl > M || s_kl < 1
        # Out of bounds
        return add_a, add_b, 0.0, k,l
    end
    p_k_index, q_l_index = find_mode(add_a, (p_k, q_l))
    s_kl_index = find_mode(add_a, s_kl)
    new_add_a, val1 = excitation(add_a, (p_k_index, q_l_index), (q_index, p_index))
    new_add_b, val2 = excitation(add_b, (s_kl_index,), (s_index,))

    return new_add_a, new_add_b, val1 * val2, k,l
end

"""
    momentum_external_potential_excitation(ep, add, i, map::OccupiedModeMap) -> nadd, α

The momentum space version of an external potential. `ep` may be a discrete Fourier
transform of a real-space potential.

```math
\\sum_{k,q} \\mathtt{ep}_k a^†_{q+k} a_{q} |\\mathtt{add}⟩
```

Return the new address `nadd` and value `α` of the `i`th of the
`(num_modes(add)-1) * num_occupied_modes(add)` terms in the sum, excluding the diagonal
term `∝ |add⟩`. It is expected that `map == OccupiedModeMap(add)`.

See [`momentum_external_potential_diagonal`](@ref), [`OccupiedModeMap`](@ref),
[`num_occupied_modes`](@ref), [`num_modes`](@ref).
"""
function momentum_external_potential_excitation(ep, add, i, map::OccupiedModeMap)
    M = num_modes(add)
    p, q = fldmod1(i, M - 1) # i == (p-1)*(M-1) + q
    p_index = map[p] # p-th occupied mode in add
    q += q ≥ p_index.mode # leave out diagonal matrix element
    q_index = find_mode(add, q) # q-th mode in add (not counting p)
    k = p_index.mode - q # change in momentum
    factor = ep[mod(k, M) + 1]
    new_add, value = excitation(add, (q_index,), (p_index,)) # a_q^† a_p |add⟩
    return new_add, value * factor
end

"""
    momentum_external_potential_diagonal(ep, add, map::OccupiedModeMap)

The diagonal part of [`momentum_external_potential_excitation`](@ref).
"""
function momentum_external_potential_diagonal(ep, add, map::OccupiedModeMap)
    onproduct = sum(map) do index
        index.occnum
    end
    return onproduct * ep[1]
end


text:
using Test
using CompScienceMeshes
using LinearAlgebra

verts = [
    point(0,0,0),
    point(1,0,0),
    point(0,1,0),
    point(1,1,0),]

nodes = [
    (1,),
    (2,),
    (3,),
    (4,)]

edges = [
    (1,2,),
    (2,3,),
    (3,1,),
    (2,4),
    (4,3),]

faces = [
    (1,2,3,),
    (-2,4,5)]

T = eltype(verts[1])
P = eltype(verts)
m = CompScienceMeshes.SComplex2D{T,P}(verts,nodes,edges,faces)

@test numcells(m) == 2

ch = chart(m, first(cells(m)))
@test ch.vertices[1] == verts[3]
@test ch.vertices[2] == verts[1]
@test ch.vertices[3] == verts[2]
@test dot(normal(center(ch)),point(0,0,1)) > 0

x1 = iterate(cells(m))
@test x1 != nothing
x2 = iterate(cells(m), x1[2])
@test x2 != nothing
c2 = x2[1]

ch2 = chart(m, c2)
@test ch2.vertices[1] == verts[4]
@test ch2.vertices[2] == verts[3]
@test ch2.vertices[3] == verts[2]
@test dot(normal(center(ch2)),point(0,0,1)) > 0

m0 = skeleton(m,0)
@test dimension(m0) == 0

m1 = skeleton(m,1)
@test dimension(m1) == 1

m2 = skeleton(m,2)
@test dimension(m2) == 2

D = connectivity(m1,m2)
@test sum(D,dims=1) == [1 0 1 1 1]

# b = CompScienceMeshes.SComplex1D{T,P}(verts, nodes, edges[[1,3,4,5]])
b = boundary(m2)
@test numcells(b) == 4
b0 = skeleton(b,0)
@test numcells(b0) == 4

@test inclosure_gpredicate(b).(verts) == [true, true, true, true]

mm2 = weld(m,m,seam=b)
mm1 = skeleton(mm2,1)
mm0 = skeleton(mm2,0)
@test numcells(mm2) == 2*numcells(m)
@test numcells(mm1) == 6
@test numcells(mm0) == 4

r_old = meshrectangle(1.0, 1.0, 0.5, 3)
r_new = CompScienceMeshes.SComplex2D(r_old)

@test numcells(r_new) == 8
@test numcells(skeleton(r_new,1)) == 16
@test numcells(skeleton(r_new,0)) == 9

b_new = boundary(r_new)
@test dimension(b_new) == 1
@test numcells(b_new) == 8

# r_new1 = skeleton(r_new,1)
# p = interior_tpredicate(r_new)
# for e in cells(r_new1)
#     println(cartesian(center(chart(r_new1,e))))
#     println(p(e))
# end

rr = weld(r_new, r_new, seam=boundary(r_new))
rr1 = skeleton(rr,1)
rr0 = skeleton(rr,0)
@test numcells(rr) == 16
@test numcells(rr1) == 24
@test numcells(rr0) == 10

fm = -m
C = collect(cells(fm))
ch1 = chart(fm,C[1])
ch2 = chart(fm,C[2])

@test dot(normal(center(ch1)), point(0,0,1)) < 0
@test dot(normal(center(ch2)), point(0,0,1)) < 0

dl = weld(m,-m,seam=boundary(m))
@test numcells(dl) == 4
@test numcells(skeleton(dl,1)) == 6
@test numcells(skeleton(dl,0)) == 4


text:
using Photometry
using Documenter

setup = quote
    using Photometry
    using Random
    Random.seed!(123456)
end
DocMeta.setdocmeta!(Photometry, :DocTestSetup, setup; recursive = true)

makedocs(
    modules = [Photometry],
    authors = "Miles Lucas <mdlucas@hawaii.edu>",
    repo = "https://github.com/JuliaAstro/Photometry.jl/blob/{commit}{path}#L{line}",
    sitename = "Photometry.jl",
    format = Documenter.HTML(;
        prettyurls = get(ENV, "CI", "false") == "true",
        canonical = "https://juliaastro.github.io/Photometry.jl",
    ),
    pages = [
        "Home" => "index.md",
        "Background Estimation" => [
            "Getting Started" => "background/index.md",
            "Background Estimators" => "background/estimators.md",
            "Background Interpolators" => "background/interpolators.md"
        ],
        "Source Detection" => [
            "Getting Started" => "detection/index.md",
            "Algorithms" => "detection/algs.md",
        ],
        "Aperture Photometry" => [
            "Getting Started" => "apertures/index.md",
            "Apertures" => "apertures/apertures.md",
            "Examples" => "apertures/examples.md"
        ]
    ])

deploydocs(
    repo = "github.com/JuliaAstro/Photometry.jl",
    push_preview = true)


text:
#######################################################################
#
#######################################################################

#--------------------------------------------------
#
#--------------------------------------------------

module JuliaPkgsList

#--------------------------------------------------
# Exports
#--------------------------------------------------

export genTopPkgsList, getSortedPkgsInfo, getSortedPkgs
export PKGS_INFO_URL, loadPkgsData

#--------------------------------------------------
# Imports
#--------------------------------------------------

using JSON

#--------------------------------------------------
# Data and parameters
#--------------------------------------------------

# URL with JSON data on registered Julia packages
# from official Julia Computing web page (JuliaHub)
const PKGS_INFO_URL = "https://juliahub.com/app/packages/info"

# Path to the folder with data, relative to this file
const DATA_DIR = abspath(joinpath(@__DIR__, "..", "data"))

mkDataPath(fileName :: String) = joinpath(DATA_DIR, fileName)

# Default number of top packages
const TOP_PKGS_NUM = 100

# Default path to the output file with the list of packages
const PKGS_LIST_FILE = mkDataPath("top-pkgs-list.txt")
# Default path to the JSON file with packages info
const PKGS_INFO_FILE = mkDataPath("julia-pkgs-info.json")
# Default file with git-repos of excluded packages
const EXCLUDED_REPOS_FILE = mkDataPath("excluded.txt")

# If `true`, functions print some status information
VERBOSE = true

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# Code
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

#--------------------------------------------------
# Library
#--------------------------------------------------

include("utils.jl")
include("lib.jl")

#--------------------------------------------------
# Main entry function
#--------------------------------------------------

"""
    genTopPkgsList(...)

Entry function for generating the list of most starred Julia packages
using information from Julia Hub.

# Arguments
- `topPkgsNum::Int`: the number of most starred packages (natural number; if 0, means all packages)
- `reloadPkgsInfo::Bool`: if `true`, reloads JSON with packages info
- `pkgsListOutputFile::String`: ...
"""
genTopPkgsList(
    topPkgsNum          :: Int      = TOP_PKGS_NUM,
    reloadPkgsInfo      :: Bool     = false
;
    pkgsListOutputFile  :: String   = PKGS_LIST_FILE,
    addPkgNum           :: Bool     = true,
    showName            :: Bool     = false,
    includeVersion      :: Bool     = false,
    pkgsInfoFile        :: String   = PKGS_INFO_FILE,
    excludedReposFile   :: String   = EXCLUDED_REPOS_FILE,
    useStdOut           :: Bool     = false,
    useExcluded         :: Bool     = true
) = begin
    topPkgsNum >= 0 || exitErrWithMsg("The number of packages cannot be negative")
    @status("Packages list generation STARTED")
    # obtain package info
    if reloadPkgsInfo || !isfile(pkgsInfoFile)
        downloadPkgsInfo(pkgsInfoFile)
    end
    # load excluded packages
    excluded = useExcluded ? loadExcludedRepos(excludedReposFile) : String[]
    # get pkgs data
    pkgs = loadPkgsInfo(pkgsInfoFile)
    sortedPkgs =
        try
            getSortedPkgs(pkgs, showName;
                excludedRepos=excluded, includeVersion=includeVersion)
        catch err
            exitErrWithMsg("Unable to process packages data", err)
        end
    # outputting information
    if topPkgsNum > length(sortedPkgs)
        topPkgsNum = length(sortedPkgs)
        @status("The number of packages is too big -- set to the max $(topPkgsNum)")
    end
    outputPkgsInfo(sortedPkgs, topPkgsNum, useStdOut, pkgsListOutputFile, addPkgNum)
    @status("Packages list generation COMPLETED")
end

#--------------------------------------------------
# Aux functions
#--------------------------------------------------

downloadPkgsInfo(pkgsInfoFile :: String) = begin
    @status("Downloading packages info into $(pkgsInfoFile)...")
    try
        download(PKGS_INFO_URL, pkgsInfoFile)
    catch err
        exitErrWithMsg(
            "Unable to download Julia packages information from Julia Hub",
            err)
    end
end

loadExcludedRepos(excludedReposFile :: String) :: Vector{String} = begin
    @status("Loading the list of excluded repositories from $(excludedReposFile)...")
    text =
        try
            read(excludedReposFile, String)
        catch err
            exitErrWithMsg(
                "Unable to load the list of excluded repositories",
                err)
        end
    repos = filter(!isempty, splitlines(text))
    Vector{String}(repos)
end

loadPkgsInfo(pkgsInfoFile :: String) :: Vector = begin
    @status("Loading packages data from $(pkgsInfoFile)...")
    try
        loadPkgsData(pkgsInfoFile)
    catch err
        exitErrWithMsg(
            "Unable to parse packages information",
            err)
    end
end

outputPkgsInfo(
    sortedPkgs :: Vector, topPkgsNum :: Int,
    useStdOut :: Bool, pkgsListOutputFile :: String, addPkgNum :: Bool
) = begin
    resultPkgs = topPkgsNum == 0 ? sortedPkgs : sortedPkgs[1:topPkgsNum]
    output = join(map(string, resultPkgs), '\n')
    if useStdOut
        println(output)
    else
        if addPkgNum && topPkgsNum > 0
            (dirPath, fileName) = splitdir(pkgsListOutputFile)
            pkgsListOutputFile = joinpath(dirPath, "$(topPkgsNum)-$(fileName)")
        end
        @status("Writing output to $(pkgsListOutputFile)...")
        try
            open(pkgsListOutputFile, "w") do io
                write(io, output)
            end
        catch err
            exitErrWithMsg(
                "Unable to write to the output file (possibly, the folder doesn't exist)",
                err)
        end
    end
end

end


text:
include("newton.jl")
include("krylov.jl")


text:
# Copyright (c) 2020 California Institute of Technology (“Caltech”). U.S.
# Government sponsorship acknowledged.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#   • Redistributions of source code must retain the above copyright notice,
#     this list of conditions and the following disclaimer.
#   • Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions and the following disclaimer in the documentation
#     and/or other materials provided with the distribution.
#   • Neither the name of Caltech nor its operating division, the Jet
#   Propulsion Laboratory, nor the names of its contributors may be used to
#   endorse or promote products derived from this software without specific
#   prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

@doc raw"""
    WindowChains

This module exposes several ways to calculate a chain of window functions. All
require a cache to be created that speeds up subsequent calls. The type of
cache determines which method is used.

To create the cache, simply call one of
```
julia> cache = SFB.WindowChains.WindowChainsCacheWignerChain(win, wmodes, amodes)
julia> cache = SFB.WindowChains.WindowChainsCacheFullWmix(win, wmodes, amodes)
julia> cache = SFB.WindowChains.WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
julia> cache = SFB.WindowChains.WindowChainsCacheSeparableWmix(win, wmodes, amodes)
julia> cache = SFB.WindowChains.WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
julia> cache = SFB.WindowChainsCache(win, wmodes, amodes)
```
The last one will automatically select the typically fastest algorithm.

Then to calculate an element of a chain of windows, use `window_chain(ell, n1,
n2, cache)`. For example, for a chain of four windows,
```
julia> ell = [0, 1, 2, 3]
julia> n1 = [1, 2, 3, 4]
julia> n2 = [5, 6, 7, 8]
julia> Wk = SFB.window_chain(ell, n1, n2, cache)
```
"""
module WindowChains

export WindowChainsCache
export window_chain

using ..Windows
using WignerSymbols
using WignerFamilies
using Healpix
using ..HealpixHelpers
#using ..HealPy
using ..NDIterators
using ..SeparableArrays
using ..WignerChains
using ..Modes

######### WindowChainsCache ##################################################

abstract type WindowChainsCache{T<:Real} end

struct WindowChainsCacheWignerChain{T} <: WindowChainsCache{T}
    I_LM_ln_ln::Array{Complex{T},5}
    LMcache::Array{Array{Int,1},1}
end

struct WindowChainsCacheFullWmix{T} <: WindowChainsCache{T}
    wmix::Array{Complex{T},2}
    wmix_negm::Array{Complex{T},2}
    amodes::AnlmModes
end

struct WindowChainsCacheFullWmixOntheflyWmix{T} <: WindowChainsCache{T}
    I_LM_ln_ln::Array{Complex{T},5}
    LMcache::Array{Array{Int,1},1}
end

struct WindowChainsCacheSeparableWmix{T} <: WindowChainsCache{T}
    Wlmlm::Array{Complex{T},2}
    Wlmlm_negm::Array{Complex{T},2}
    LMcache::Array{Array{Int,1},1}
    Ilnln::Array{T,4}
    amodes::AnlmModes
end

struct WindowChainsCacheSeparableWmixOntheflyWlmlm{T} <: WindowChainsCache{T}
    Wlm::Array{Complex{T},1}
    LMcache::Array{Array{Int,1},1}
    Ilnln::Array{T,4}
    amodes::AnlmModes
end

# Using wigner chain
function WindowChainsCacheWignerChain(win, wmodes, amodes)
    lmax = amodes.lmax
    nmax = amodes.nmax
    nmax_l = amodes.nmax_l
    check_nsamp(amodes, wmodes)

    # Wr_lm
    LMAX = 2 * amodes.lmax
    Wr_lm = calc_Wr_lm(win, LMAX, amodes.nside)
    alm = Alm(LMAX, LMAX)
    LMcache = [almIndex(alm, L, 0:L) for L=0:LMAX]
    lmsize = numberOfAlms(LMAX)
    #@show LMcache
    #@show typeof(LMcache)

    # gnlr
    r, Δr = window_r(wmodes)
    gnl = amodes.basisfunctions
    gnlr = fill(NaN, length(r), size(gnl.knl)...)
    for l=0:amodes.lmax, n=1:amodes.nmax_l[l+1]
        @. gnlr[:,n,l+1] = gnl(n,l,r)
    end

    # I_LM_ln_ln
    I_LM_ln_ln = fill(im*NaN, lmsize, lmax+1, nmax, lmax+1, nmax)
    for l2=0:lmax, n2=1:nmax_l[l2+1], l1=0:lmax, n1=1:nmax_l[l1+1], iLM=1:lmsize
        #@show l2, n2, l1, n1, iLM, lmsize
        !isnan(I_LM_ln_ln[iLM,l1+1,n1,l2+1,n2]) && continue
        I = Δr * sum(@. r^2 * gnlr[:,n1,l1+1] * gnlr[:,n2,l2+1] * Wr_lm[:,iLM])
        I_LM_ln_ln[iLM,l1+1,n1,l2+1,n2] = I
        I_LM_ln_ln[iLM,l2+1,n2,l1+1,n1] = I
    end

    return WindowChainsCacheWignerChain(I_LM_ln_ln, LMcache)
end

# full wmix
function WindowChainsCacheFullWmix(win, wmodes, amodes)
    wmix = calc_wmix(win, wmodes, amodes)
    wmix′ = calc_wmix(win, wmodes, amodes, neg_m=true)
    wmix′ = collect(wmix′')  # put negative m into the second index
    return WindowChainsCacheFullWmix(wmix, wmix′, amodes)
end

# full wmix with on-the-fly wmix
function WindowChainsCacheFullWmixOntheflyWmix(win, wmodes, amodes)
    cache = WindowChainsCacheWignerChain(win, wmodes, amodes)
    return WindowChainsCacheFullWmixOntheflyWmix(cache.I_LM_ln_ln, cache.LMcache)
end

# separable wmix
function WindowChainsCacheSeparableWmix(win, wmodes, amodes)
    Wlmlms = calc_Wlmlm(win.mask, amodes.lmax, amodes.nside)
    Ilnln = calc_Ilnln(win.phi, wmodes, amodes)
    return WindowChainsCacheSeparableWmix(Wlmlms..., Ilnln, amodes)
end

# separable wmix with on-the-fly Wlmlm
function WindowChainsCacheSeparableWmixOntheflyWlmlm(win, wmodes, amodes)
    Wlm, LMcache = calc_Wlm(win.mask, amodes.lmax, amodes.nside)
    Ilnln = calc_Ilnln(win.phi, wmodes, amodes)
    return WindowChainsCacheSeparableWmixOntheflyWlmlm(Wlm, LMcache, Ilnln, amodes)
end

# choose best
WindowChainsCache(win, wmodes, amodes) = WindowChainsCacheFullWmix(win, wmodes, amodes)
WindowChainsCache(win::SeparableArray, wmodes, amodes) = WindowChainsCacheSeparableWmix(win, wmodes, amodes)

############ window_chains() #################################################

# wigner chains
function window_chain(ell, n1, n2, cache::WindowChainsCacheWignerChain{T}) where {T}
    I_LM_l_l = NeqLView(cache.I_LM_ln_ln, ell, n2, n1)
    return window_chain_wigner_chain(ell, I_LM_l_l, cache.LMcache)
end

# full wmix
function window_chain(ell, n1, n2, cache::WindowChainsCacheFullWmix{T}) where {T}
    Wk = T(0)
    m = NDIterator(-ell, ell)
    while advance(m)
        nl2 = getidx(cache.amodes, n2[end], ell[end], 0)
        nl1 = getidx(cache.amodes, n1[1], ell[1], 0)
        w = get_wmix(cache.wmix, cache.wmix_negm, nl2, m[end], nl1, m[1])
        for i=2:length(ell)
            nl2 = getidx(cache.amodes, n2[i-1], ell[i-1], 0)
            nl1 = getidx(cache.amodes, n1[i], ell[i], 0)
            w *= get_wmix(cache.wmix, cache.wmix_negm, nl2, m[i-1], nl1, m[i])
        end
        Wk += real(w)
        #@show m,w
    end
    return Wk
end

# full wmix, on-the-fly
function window_chain(ell, n1, n2, cache::WindowChainsCacheFullWmixOntheflyWmix{T}) where {T}
    I_LM_l_l = NeqLView(cache.I_LM_ln_ln, ell, n2, n1)
    return window_chain_onthefly_wmix(ell, I_LM_l_l, cache.LMcache)
end

# separable wmix
function window_chain(ell, n1, n2, cache::WindowChainsCacheSeparableWmix{T}) where {T}
    Wk = T(0)
    m = NDIterator(-ell, ell)
    while advance(m)
        w = get_wlmlm(cache, ell[end], m[end], ell[1], m[1])
        for i=2:length(ell)
            w *= get_wlmlm(cache, ell[i-1], m[i-1], ell[i], m[i])
        end
        #@show m,w
        Wk += real(w)
    end
    Ik = calc_kprod(cache.Ilnln, ell .+ 1, n1, n2)
    return Ik * Wk
end

# separable wmix with on-the-fly wmix
function window_chain(ell, n1, n2, cache::WindowChainsCacheSeparableWmixOntheflyWlmlm{T}) where {T}
    Wk = window_chain_onthefly_wmix(ell, cache.Wlm, cache.LMcache)
    Ik = calc_kprod(cache.Ilnln, ell .+ 1, n1, n2)
    return Ik * Wk
end

############ specialty window_chain() ############

# calculate full window chain matrix
function window_chain(k, win, wmodes::ConfigurationSpaceModes, cmodes::ClnnModes; cache=nothing)
    if isnothing(cache)
        cache = WindowChainsCache(win, wmodes, cmodes.amodes)
    end
    lnnsize = getlnnsize(cmodes)
    Wk = fill(NaN, fill(lnnsize,k)...)
    ell = Array{Int}(undef, k)
    nn1 = Array{Int}(undef, k)
    nn2 = Array{Int}(undef, k)
    lnn = NDIterator(1, lnnsize; N=k)
    while advance(lnn)
        for i=1:k
            l, n1, n2 = getlnn(cmodes, lnn[i])
            ell[i] = l
            nn1[i] = n1
            nn2[i] = n2
        end
        Wk[lnn...] = window_chain(ell, nn1, nn2, cache)
    end
    return Wk
end

# symmetry for some of the n1,n2
@doc raw"""
    window_chain(ell, n1, n2, cache, symmetries)

This version adds up several window chains taking into account the
`symmetries`. `symmetries` is an array of pairs of numbers specifying the
symmetries to consider. Each pair specifies the `ℓnn′` index and the type of
symmetry. For example, when `k≥3`, then `symmetries = [1=>0, 2=>1, 3=>2]`
would specify that no symmetries are taken into account for the first `lnn′`
combination, the second symmetry will flip `n` and `n′` and add the result only
when `n ≠ n′`, and the third will add the result regardless whether the `n`
equal or not.
"""
function window_chain(ell, n1, n2, cache, symmetries)
    if isempty(symmetries)
        #@show ell n1 n2
        return window_chain(ell, n1, n2, cache)
    end

    Wk = window_chain(ell, n1, n2, cache, symmetries[2:end])

    i, sym = symmetries[1]

    if sym >= 1
        if n1[i] == n2[i]
            if sym == 2
                Wk *= 2
            end
        else
            n1[i], n2[i] = n2[i], n1[i]
            Wk += window_chain(ell, n1, n2, cache, symmetries[2:end])
            n1[i], n2[i] = n2[i], n1[i]
        end
    end

    return Wk
end

############# helpers for WindowChainsCache ##################################

function calc_Wr_lm(win, LMAX, Wnside)
    nr = size(win,1)
    Wr_lm = fill(NaN*im, nr, getlmsize(LMAX))
    for i=1:nr
        W = udgrade(win[i,:], Wnside)
        Wr_lm[i,:] .= mymap2alm(W, lmax=LMAX).alm
    end
    return Wr_lm
end

function calc_Wr_lm(win::SeparableArray, LMAX, Wnside)  # specialize
    mask = udgrade(win.mask, Wnside)
    wlm = mymap2alm(mask, lmax=LMAX).alm
    return SeparableArray(win.phi, wlm, name1=:phi, name2=:wlm)
end

function calc_Wlm(mask, lmax, nside)
    LMAX = 2 * lmax
    mask = udgrade(mask, nside)
    Wlm = mymap2alm(mask, lmax=LMAX)
    LMcache = [almIndex(Wlm, L, 0:L) for L=0:LMAX]
    return Wlm.alm, LMcache
end

function calc_Wlmlm(mask, lmax, nside)
    wlm, LMcache = calc_Wlm(mask, lmax, nside)
    LMAX = 2 * lmax
    lmsize = numberOfAlms(LMAX)
    wlmlm = fill(NaN*im, lmsize, lmsize)
    wlmlm_negm = fill(NaN*im, lmsize, lmsize)
    for l=0:lmax, m=0:l, L=0:lmax, M=0:L
        i = LMcache[l+1][abs(m)+1]
        j = LMcache[L+1][abs(M)+1]
        wlmlm[i,j] = window_wmix(l, m, L, M, wlm, LMcache)
        wlmlm_negm[i,j] = window_wmix(l, m, L, -M, wlm, LMcache)
    end
    return wlmlm, wlmlm_negm, LMcache
end

function get_wlmlm(cache::WindowChainsCacheSeparableWmix, l::Int, m::Int, L::Int, M::Int)
    i = cache.LMcache[l+1][abs(m)+1]
    j = cache.LMcache[L+1][abs(M)+1]
    if m >= 0
        if M >= 0
            return cache.Wlmlm[i,j]
        else
            return cache.Wlmlm_negm[i,j]
        end
    else # m < 0
        w = if M >= 0
            cache.Wlmlm_negm[i,j]
        else
            cache.Wlmlm[i,j]
        end
        return conj(isodd(m+M) ? -w : w)
    end
end

function calc_Ilnln(phi, wmodes, amodes)
    check_nsamp(amodes, wmodes)

    # gnlr
    r, Δr = window_r(wmodes)
    gnl = amodes.basisfunctions
    gnlr = fill(NaN, length(r), size(gnl.knl)...)
    for l=0:amodes.lmax, n=1:amodes.nmax_l[l+1]
        @. gnlr[:,n,l+1] = gnl(n,l,r)
    end

    # Ilnln
    lmax = amodes.lmax
    nmax = amodes.nmax
    nmax_l = amodes.nmax_l
    Ilnln = fill(NaN, lmax+1, nmax, lmax+1, nmax)
    for l2=0:lmax, n2=1:nmax_l[l2+1], l1=0:lmax, n1=1:nmax_l[l1+1]
        #@show l2, n2, l1, n1
        !isnan(Ilnln[l1+1,n1,l2+1,n2]) && continue
        I = Δr * sum(@. r^2 * gnlr[:,n1,l1+1] * gnlr[:,n2,l2+1] * phi)
        Ilnln[l1+1,n1,l2+1,n2] = I
        Ilnln[l2+1,n2,l1+1,n1] = I
    end

    return Ilnln
end

########### helpers for window_chain() #######################################

# This struct encodes how to access I_LM_lᵢnᵢ_lⱼnⱼ such that we can access it as I[LM,i,j].
struct NeqLView{T}
    I_LM_ln_ln::Array{T,5}
    ell::Array{Int,1}
    n1::Array{Int,1}
    n2::Array{Int,1}
end
NeqLView(I, ell, n1, n2) = NeqLView(I, convert(Array{Int}, ell), convert(Array{Int}, n1), convert(Array{Int}, n2))

Base.getindex(v::NeqLView, LM::Int, i::Int, j::Int) = v.I_LM_ln_ln[LM, v.ell[i]+1, v.n1[i], v.ell[j]+1, v.n2[j]]
Base.getindex(v::NeqLView, ::Colon, i::Int, j::Int) = v.I_LM_ln_ln[:, v.ell[i]+1, v.n1[i], v.ell[j]+1, v.n2[j]]

function window_chain_wigner_chain(ell, I_LM_l_l, LMcache)
    # This implementation is actually quite slow. Instead, use 'window_chain()' below.
    T = Float64
    k = length(ell)

    # only iterate over triangles
    Lmin = fill(0, k)
    Lmax = fill(0, k)
    Lmin[1] = abs(ell[1] - ell[end])
    Lmax[1] = ell[1] + ell[end]
    for i=2:k
        Lmin[i] = abs(ell[i] - ell[i-1])
        Lmax[i] = ell[i] + ell[i-1]
    end
    #@show Lmin Lmax

    wk = T(0)
    L = NDIterator(Lmin, Lmax)
    while advance(L)
        twoLplus1_w3j000 = √(2*L[1] + 1) * wigner3j(T, ell[end], ell[1], L[1], 0, 0, 0)
        for i=2:k
            twoLplus1_w3j000 *= √(2*L[i] + 1) * wigner3j(T, ell[i-1], ell[i], L[i], 0, 0, 0)
        end

        # sum over M[i]
        M = NDIterator(-L, L)
        while advance(M)
            w3jk = wigner3j_chain(ell, L, M)
	    #@show L,M,w3jk
            (w3jk == 0) && continue

            LM = LMcache[L[1]+1][abs(M[1])+1]
            I = I_LM_l_l[LM, k, 1]
            (M[1] < 0) && (I = conj(I))
            #@show L[1],M[1],I
            Iprod = I
            for i=2:k
                LM = LMcache[L[i]+1][abs(M[i])+1]
                I = I_LM_l_l[LM, i-1, i]
                (M[i] < 0) && (I = conj(I))
                Iprod *= I
                #@show L[i],M[i],I
            end

            dwk = twoLplus1_w3j000 * real(Iprod) * w3jk  # imag must sum out
            wk += dwk
	    #@show L,M,dwk,Iprod
        end
    end
    wk *= (4*π)^(-k/2) * prod(@. 2 * ell + 1)
    return wk
end

function calc_kprod(arr, ell, n1, n2)
    Ik = arr[ell[end],n2[end],ell[1],n1[1]]
    for i=2:length(ell)
        Ik *= arr[ell[i-1],n2[i-1],ell[i],n1[i]]
    end
    return Ik
end

########### helpers for on-the-fly wmix #######################################

# Faster implementation using a direct summation.
function window_chain_onthefly_wmix(ell, I_LM_l_l::NeqLView, LMcache)
    # Note: can probably speed this up a bit by calculating the all Wlmlm for
    # given ell before summing over the m.
    T = Float64
    k = length(ell)
    wk = T(0)
    m = NDIterator(-ell, ell)
    while advance(m)
        w = window_wmix(ell[end], m[end], ell[1], m[1], I_LM_l_l[:,k,1], LMcache)
        for i=2:k
            w *= window_wmix(ell[i-1], m[i-1], ell[i], m[i], I_LM_l_l[:,i-1,i], LMcache)
        end
        wk += real(w)
        #@show m,w
    end
    return wk
end

# Faster implementation using a direct summation.
function window_chain_onthefly_wmix(ell, Wlm, LMcache)
    # Note: can probably speed this up a bit by calculating the all Wlmlm for
    # given ell before summing over the m.
    T = Float64
    k = length(ell)
    wk = T(0)
    m = NDIterator(-ell, ell)
    while advance(m)
        w = window_wmix(ell[end], m[end], ell[1], m[1], Wlm, LMcache)
        for i=2:k
            w *= window_wmix(ell[i-1], m[i-1], ell[i], m[i], Wlm, LMcache)
        end
        wk += real(w)
    end
    return wk
end

# calculate window function given lm,LM.
function window_wmix_wignerfamilies(l, m, L, M, Wlm, LMcache)
    T = Float64
    wigs000 = wigner3j_f(l, L, 0, 0)
    #@show l,L,-m,M
    wigs = wigner3j_f(l, L, -m, M)
    #@show eachindex(wigs000) eachindex(wigs)
    #@show length(wigs000) length(wigs)
    #@assert length(wigs000) >= length(wigs)
    #@show wigs000 wigs
    M1 = m - M
    w = T(0)*im
    for L1 in eachindex(wigs)
        L1M1 = LMcache[L1+1][abs(M1)+1]
        Iterm = Wlm[L1M1]
        if M1 < 0
            Iterm = (-1)^M1 * conj(Iterm)
        end
        w += √T(2*L1 + 1) * wigs000[L1] * wigs[L1] * Iterm
    end
    return (-1)^m * √((2*l + 1)*(2*L + 1) / (4*T(π))) * w
end
function window_wmix_wignersymbols(l, m, L, M, Wlm, LMcache)
    T = Float64
    M1 = m - M
    w = T(0)*im
    for L1 in max(abs(L-l),abs(M1)):(L+l)
        L1M1 = LMcache[L1+1][abs(M1)+1]
        Iterm = Wlm[L1M1]
        if M1 < 0
            Iterm = (-1)^M1 * conj(Iterm)
        end
        wig000 = Windows.wigner3j000(l, L, L1)
        w3j = wigner3j(T, l, L, L1, -m, M)
        w += √T(2*L1 + 1) * wig000 * w3j * Iterm
    end
    return (-1)^m * √((2*l + 1)*(2*L + 1) / (4*T(π))) * w
end
window_wmix = window_wmix_wignerfamilies
#window_wmix = window_wmix_wignersymbols

function get_wmix(w, w′, nl, m, NL, M)
    (m >= 0 && M >= 0) && return w[nl+m, NL+M]
    (m >= 0 && M < 0)  && return w′[nl+m, NL-M]
    (m < 0  && M >= 0) && return (-1)^(m+M) * conj(w′[nl-m, NL+M])
    return (-1)^(m-M) * conj(w[nl-m, NL-M])
end

end

# vim: set sw=4 et sts=4 :


text:
export spglcore

"""
# Info
    Use:
        spglcore{ETxg<:Number, Txg<:AbstractVector{ETxg}, Tidx<:BitArray}(init::spgInit{ETxg, Txg, Tidx})
# Input
    init:
        - An instance of the composite type spgInit which contains all initalized variables
            created in the preceding call to spgl1.\n

# Outputs
    init:
        - An instance of the composite type spgInit that contains all updates to variables
            initialized in the preceding call to spgl1\n
    rNorm:
        - Value of the objective\n

# Author
    Keegan Lensink
        Seismic Laboratory for Imaging and Modeling
        The University of British Columbia
        keeganlensink@gmail.com

    This code is an adaptation of Michael P. Friedlander, Ewout van den Berg,
    and Aleksandr Aravkin's MATLAB program SPGL1.
"""
function spglcore(init::spgInit{TA, ETb, ETx, ETg, ETr, Tidx}) where
            {TA<:Union{joAbstractLinearOperator, AbstractArray, Function},
             ETb<:Number, ETx<:Number, ETg<:Number, ETr<:Number, Tidx<:BitArray}

    #DEVNOTE# Create spgInit type to hold all these initialized paramaters

    (init.options.verbosity > 1) && println("script has entered spglcore\n
            Begin Main Loop.")

    # Pull out options type for easier use
    options = init.options
    params = init.params

    # Initialise rErr
    rErr = zero(ETr)

    #Main Loop
    while true

        # ================================================================================
        # Test Exit Conditions
        # ================================================================================
        gNorm = zero(ETg)

        if (options.proxy)
            #@code_warntype options.dual_norm(init.g2, options.weights, init.params)
            # prev prod: 1829
            gNorm = options.dual_norm(init.g2, options.weights, init.params)
        else
            gNorm = options.dual_norm(init.g,  options.weights, init.params)
        end

        # rNorm and f are the same thing
        rNorm = copy(init.f)
        #println("rNorm = $rNorm")

        tmp_proj, tmp_itn = project(init.x - init.g,
                                        init.tau, init.timeProject, options, params)
        Err = norm(init.x - tmp_proj)
        rErr = Err/max(1,init.f)

        aError1 = rNorm - init.sigma
        aError2 = rNorm^2 - init.sigma^2
        rError1 = abs(aError1) / max(1,rNorm)
        rError2 = abs(aError2) / max(1,init.f)

        #= Count number of consecutive iterations with identical support

        nnzOld = deepcopy(init.nnzIdx)

        #DEVNOTE# -Performance: Expensive function call
        if ~(options.proxy)
            nnzX,nnzG,init.nnzIdx,nnzDiff = activevars(init.x, init.g, init.nnzIdx, options, params)
        else
            nnzX,nnzG,init.nnzIdx,nnzDiff = activevars(init.x, init.g2, init.nnzIdx, options, params)
        end

        if (nnzDiff == -1)
            init.nnzIter = 0
        end

        options.verbosity > 1 && println("fin CompConditions")

        if isempty(init.x)
            println("DEVNOTE x is empty in spglcore, check spgl1 90-99")
        else
            init.nnzIter += 1
            (init.nnzIter >= options.activeSetIt) && (init.exit_status.triggered = 9)
            tmp_nnzX1 = min(0.1,10*options.optTol)::Float64
            tmp_nnzX2::BitArray{1} = (abs.(init.x) .>= tmp_nnzX1)
            nnzX = sum(tmp_nnzX2)
        end
        =#

        # SingleTau, check if optimal
        # Second condition guards against large tau
        if init.singleTau
            if ((rErr .<= options.optTol) | (rNorm .< options.optTol*init.bNorm))
                init.exit_status.triggered = 4
            end
        else
            # Test for LS solution found
            if gNorm <= options.lsTol
                init.exit_status.triggered = 3
            end

            if (rErr <= max(options.optTol, rError2)) | (rError1 <= options.optTol)

                # Problem is nearly optimal for current tau
                # Check optimailty of current root
                test1 = (rNorm <= options.bpTol*init.bNorm)::Bool
                test3 = (rError1 <= options.optTol)::Bool
                test4 = (rNorm <= init.sigma)::Bool

                test1 && (init.exit_status.triggered = 7)
                test3 && (init.exit_status.triggered = 1)
                test4 && (init.exit_status.triggered = 2)

            end

            testRelChange1 = (abs(init.f - init.fOld) <= options.decTol * init.f)::Bool
            testRelChange2 = (abs(init.f - init.fOld) <= 1e-1*init.f*(abs(rNorm - init.sigma)))

            (false) && println("""
            decTol: $(options.decTol)
            f: $(init.f)
            fOld: $(init.fOld)
            $(abs(init.f - init.fOld)) <= $(options.decTol * init.f)
            testRelChange1: $testRelChange1
            testRelChange2: $testRelChange2
            """)

            init.testUpdateTau = ((testRelChange1 & (rNorm >  2*init.sigma)) |
                            (testRelChange2 & (rNorm <= 2*init.sigma))) &
                            isnull(init.exit_status.triggered) &
                            ~init.testUpdateTau

            if init.testUpdateTau

                (options.verbosity > 1) && println("WARNING: Update Tau")
                if (options.quitPareto & (init.iter >= options.minPareto))
                    init.exit_status.triggered = 10
                end

                tauOld = copy(init.tau)
                init.tau = max(zero(typeof(init.tau)), init.tau + (aError1)/ gNorm)::typeof(tauOld)
                init.nNewton += one(typeof(init.nNewton))

                init.printTau = (abs(tauOld - init.tau) >= (1e-6 * init.tau))

                if init.tau < tauOld
                    (options.verbosity > 1) && println("WARNING: Tau Decreasing")
                    init.x, tmp_itn = project(init.x, init.tau, init.timeProject, options, params)

                end

            end

        end

        if (isnull(init.exit_status.triggered) & ((init.iter+1) >= options.iterations))
            init.exit_status.triggered = 5
        end

        (options.verbosity > 1) && println("fin CheckConverge")

        # ===============================================================================
        # Print log, update history, and check exit conditions
        # ===============================================================================

        if (options.verbosity > 0) | init.singleTau | init.printTau |
                                 (init.iter == 0) | ~isnull(init.exit_status.triggered)

            if init.singleTau

                s = @sprintf "%5i   %13.7e  %13.7e  %9.2e   %6.1f" init.iter rNorm rErr gNorm log10(init.stepG) #nnzX nnzG
                (options.verbosity > 0) && println(s)

                if init.subspace
                    #println("$itnLSQR")
                end
            else

                if init.printTau | init.subspace
                    s = @sprintf "%5i  %13.7e  %13.7e  %9.2e  %9.3e  %6.1f  %13.7e" init.iter rNorm rErr rError1 gNorm log10(init.stepG) #=nnzX nnzG=# init.tau
                else
                    s = @sprintf "%5i  %13.7e  %13.7e  %9.2e  %9.3e  %6.1f" init.iter rNorm rErr rError1 gNorm log10(init.stepG) #nnzX nnzG
                end
                    (options.verbosity > 0) && println(s)
            end

        end

        init.printTau = false
        init.subspace = false

        #Update History
        if isempty(init.x)
            init.xNorm1[init.iter+1] = 0.
        else
            init.xNorm1[init.iter+1] = options.primal_norm(init.x, options.weights, params)
        end
        init.rNorm2[init.iter+1] = copy(rNorm)
        init.lambda[init.iter+1] = copy(gNorm)

        # Act on exit conditions
        ~(isnull(init.exit_status.triggered)) && break
        # ================================================================================
        # Begin Iterations
        # ================================================================================
        init.iter += 1
        xOld = copy(init.x)
        init.fOld = copy(init.f)
        rOld = copy(init.r)

        try
            # ================================================================================
            # Projected gradient step and line search
            # ================================================================================

            (options.verbosity > 1) && println("begin LineSearch")
            (options.verbosity > 1) && println("Tau: $(init.tau)")

            init.f, init.x, init.r, nLine, init.stepG, lnErr, localProdA = spglinecurvy(init.A,
                                                                    init.x,
                                                                    init.gStep*init.g,
                                                                    maximum(init.lastFv),
                                                                    init.funForward,
                                                                    options.funPenalty,
                                                                    init.b,
                                                                    project,
                                                                    init.timeProject,
                                                                    init.tau,
                                                                    options,
                                                                    params)

            (options.verbosity > 1) && println("fin LineSearch")
            init.nLineTot += nLine
            init.nProdA += localProdA

            if lnErr == -1
                println("WARNING: Line Search Error not set in call to spglinecurvy")
            end

            # If an error was triggered in the line search
            if (lnErr !== 0)

                (options.verbosity > 1) && println("Line Error: $lnErr \n begin FeasLineSearch")

                # Projected backtrack failed. Retry with feasible dir'n line search
                init.x = copy(xOld)

                # In-place scaling of gradient and updating of x
                if ~isempty(xOld)

                    dx = project(xOld - init.gStep.*init.g, init.tau, init.timeProject,
                                                            options, params)[1] - xOld
                else
                    throw(error("Empty X")) # This should never be invoked
                end

                gtd = dot(init.g,dx)

                if options.linear

                    init.f, step, init.r, nLine, lnErr, localProdA = spgline(init.A,
                                                                init.f,
                                                                dx,
                                                                gtd,
                                                                rOld,
                                                                maximum(init.lastFv),
                                                                init.funForward,
                                                                options.funPenalty,
                                                                params,
                                                                init.b,
                                                                options.feasSrchIt,
                                                                options.linear,
                                                                options,
                                                                init.timeProject)
                else

                    init.f, step, init.r, nLine, lnErr, localProdA = spgline(init.A,
                                                                init.f,
                                                                dx,
                                                                gtd,
                                                                init.x,
                                                                maximum(init.lastFv),
                                                                init.funForward,
                                                                options.funPenalty,
                                                                params,
                                                                init.b,
                                                                options.feasSrchIt,
                                                                options.linear,
                                                                options,
                                                                init.timeProject)
                end

                init.nProdA += localProdA

                (options.verbosity > 1) && println("Fin FeasLineSearch")

                if isempty(xOld)
                    init.x = step*dx
                else
                    init.x = xOld + step*dx
                end

                init.x, tmp_itn = project(init.x, init.tau, init.timeProject, options, params)

            end

            # Failed again, Revert to Previous iterates and damp max BB step
            if (lnErr !== 0)
                options.stepMax /= 10
                println("WARNING Line Search Failed. Line Error: $(lnErr)") #DEVNOTE# Include more info
            end

            doSubspaceMin = false
            if options.subspaceMin
                throw(error("subspaceMin is in development"))
            end

            primNorm_x = options.primal_norm(init.x,options.weights,params)
            targetNorm = init.tau + options.optTol

            if options.ignorePErr

                if primNorm_x > targetNorm
                    println("WARNING: Primal norm of projected x is larger than expected")
                end

            end

            (options.verbosity > 1) && println("fin UpdateX")

            gOld = copy(init.g)

            init.f, init.g, init.g2 = options.funCompositeR(init.A, init.x, init.r, init.funForward,
            options.funPenalty, init.nProdAt, params)

            # xOld plays the role of s
            xOlds = init.x - xOld

            y = init.g - gOld
            sts = dot(xOlds,xOlds)
            sty = dot(xOlds,y)

            #DEVNOTE# Double check that it is okay to only compare the real part of sty
            if real(sty) <= 0
                init.gStep = options.stepMax
            else
                init.gStep = min(options.stepMax,max(options.stepMin, real(sts/sty)))
            end

            (options.verbosity > 1) && println("fin CompScaling")

        catch exc

            #DEVNOTE# Add in MaxMatVec error catch
            println("""
            ===============================CAUGHT_ERROR=====================================
            """)
            println("ERROR: ", exc)

            display(stacktrace(catch_backtrace()))

            println("""
            ================================================================================
            """)
            break
        end

        # ================================================================================
        # Update Function History
        # ================================================================================

        # Dont update if superoptimal
        if init.singleTau | (init.f > init.sigma)
            init.lastFv[mod(init.iter, options.nPrevVals)+1] = init.f
            if init.fBest > init.f
                init.fBest = copy(init.f)
                init.xBest = copy(init.x)
            end
        end

    end #Main Loop

    # Restore best solution(only if solving single problem)
    if (init.singleTau) & (init.f >init.fBest)
        if options.restore
            rNorm = init.fBest
            (options.verbosity > 1) && println("Restoring best iterate to objective: $(rNorm)")
            init.x = init.xBest
            init.r = init.b - init.funForward(init.A, init.x, Array{ETx,1}(), params)
            init.f, init. g, init.g2 = options.funCompositeR(init.A, init.x, init.r, init.funForward,
                                            options.funPenalty, init.nProdAt, params)
            if options.proxy
                gNorm = options.dual_norm(init.g2, options.weights)
            else
                gNorm = options.dual_norm(init.g, options.weights)
            end
            rNorm = init.f
        else
            (options.verbosity > 0) && println("""
                Note: Solution not actually optimal. Best objective value is $(init.fBest)
                """)
        end
    end

    return init, rNorm, gNorm, rErr
end

"""
Use: nnzX,nnzG,nnzIdx,nnzDiff = activevars(x,g,nnzIdx,options, params)

Find the current active set.
nnzX    is the number of nonzero x.
nnzG    is the number of elements in nnzIdx.
nnzIdx  is a vector of primal/dual indicators.
nnzDiff is the no. of elements that changed in the support.
"""
function activevars(x::Txg,g::Txg, nnzIdx::Ti, options::spgOptions,params::Dict{String,Any}) where
            {Ti<:BitArray{1}, ETxg<:Number, Txg<:AbstractVector{ETxg}}

    xTol = min(.1,10*options.optTol)
    gTol = min(.1,10*options.optTol)

    #DEVNOTE# -Performance: Expensive Line
    gNorm = options.dual_norm(g,options.weights, params)

    if isnull(nnzIdx)
        nnzOld_inner = Ti()
    else
        nnzOld_inner = copy(nnzIdx)
    end

    #Reduced costs for postive and negative parts of x
    z1 = gNorm .+ g
    z2 = gNorm .- g

    #Primal/dual based indicators
    xPos = BitArray{1}()
    xNeg = BitArray{1}()

    if(~options.proxy)
        # Unwrap temp vars for performance
        a = x .> xTol
        b = z1 .< gTol
        xPos = a .& b

        c = x .< -xTol
        d = z2 .< gTol
        xNeg = c .& d

        nnzIdx = xPos .| xNeg
    end

    nnzX = sum(abs.(x) .>= xTol)::Int64
    nnzG = sum(nnzIdx)::Int64

    if isempty(nnzOld_inner)
        #DEVNOTE# Use -1 instead of Inf? Inf creates a type stability
                # Need to document what this means though
        nnzDiff = -one(Int64)
    else
        e = nnzIdx .!== nnzOld_inner
        nnzDiff = sum(e)::Int64
    end

    return nnzX, nnzG, nnzIdx, nnzDiff
end

mtypeof(c...) = [typeof(a) for a in c]


text:
module Schwarzschild

include("SpectralSolver.jl")
include("NewtonSolver.jl")
using .NewtonSolver
using .SpectralSolver
using NLsolve
using CSV

function LoadLowSchwarzschildQnms()
    A = CSV.read(joinpath(@__DIR__, "SwData.csv"),NamedTuple)
    ωlist = Dict{NamedTuple,Complex{Float64}}()
    for i=1:length(A[:l])
        merge!(ωlist,Dict((l = A[:l][i], s = A[:s][i], n = A[:n][i]) => A[:ωreal][i]+im*A[:ωimag][i]))
    end
    ωlist
end

SchwMODESLOWWℓ = LoadLowSchwarzschildQnms();

function High_l_Guess(l,n,s)
    #Taken Directly from Eq: 17-22 of https://arxiv.org/pdf/0908.0329.pdf
    L = l+0.5
    N = n+0.5
    β = 1-s^2

    ω₋₁ = 1
    ω₀ᵢ = 1
    ω₁ = β/3 - (5/36)*N^2 - (115/432)
    ω₂ᵢ = β/9 + (235/3888)*N^2 - (1415/15552)
    ω₃ = -(1/27)*β^2 + β*(204*N^2 +211)/3888 + (854160*N^4 - 1664760*N^2 - 776939)/(40310784)
    ω₄ᵢ = (1/27)*β^2 + (1100*N^2 - 2719)/(46656) + (11273136*N^4 - 52753800*N^2 + 66480535)/(2902376448)

    ω = (ω₋₁*L + ω₁*(L^(-1)) + ω₃*(L^(-3))) - im*N*(ω₀ᵢ + ω₂ᵢ*(L^(-2)) - ω₄ᵢ*(L^(-4)))
    ω = (1/sqrt(27))*ω
end

function Low_l_Guess(l,n,s)
    SchwMODESLOWWℓ[(l=l,s=s,n=n)]
end

function High_n_Guess(l,n,s)
    print("Used higher n guess")
    # From H.-J. Blome, B. Mashhoon, Phys. Lett. A 110, 231 (1984)
    return (0.5 + l - im*(0.5 + n))/(3.0*sqrt(3))
end

function SchwarzschildGuess(l,n,s)
    if l >= 3
        return High_l_Guess(l,n,s)
    else
        return Low_l_Guess(l,n,s)
    end
end

function SchwarzschildModes(l,n,s)
    #if l < 3
    #   return SchwarzschildGuess(l,n,s),l*(l+1) - s*(s+1)
    #end
    ωguess = SchwarzschildGuess(l,n,s)
    #println(ωguess)
    Almsch = l*(l+1) - s*(s+1)
    P = ModeParameters((l=l,s=s,m=0,n=n,a=0.0,ω = ωguess,Alm = Almsch, Nmax = 10000, lmax = 10))
    #Newton!(P,0.0000001)

    function CFSch(P::ModeParameters)
        n = P.n; N = P.Nmax; ω = P.ω; Alm = P.Alm
        αₙsch(n) = (1+n)*(1+n-s-4*im*ω)
        βₙsch(n) = -1 - Alm - 2*n^2 - s + n*(-2 + im*16*ω) + im*6*ω + 24*ω^2 + ω*(im*2 + 8*ω)
        γₙsch(n) = (n-4*im*ω)*(n+s-4*im*ω)
        ContinuedFraction(αₙsch,βₙsch,γₙsch,n,N)
    end

    function f!(F,x)
        #ωᵢ = P.ω
        P.ω = x[1]+im*x[2]
        val = CF(P)
        #val = CFSch(P)
        #P.ω = ωᵢ
        F[1] = real(val)
        F[2] = imag(val)
    end

    NmaxTrials = [100,300,500,1000,5_000,10_000,50_000,100_000]
    ω₀ = P.ω; Alm₀ = P.Alm
    ωnew = P.ω + 1.0;
    iterations = 1;
    while (abs(ωnew-ω₀) > 1e-10) & (iterations < 8)
        P.Nmax = NmaxTrials[iterations]
        ω₀ = P.ω;
        sol = nlsolve(f!, [real(ωguess),imag(ωguess)],ftol = 1e-14, method = :newton)
        ωnews = sol.zero
        P.ω = ωnews[1] + im*ωnews[2]
        ωnew = P.ω;
        iterations += 1
        #println("ω diff = ",abs(ωnew-ω₀))
    end
    #println("Number of Iterations = ", iterations)
    P.ω, P.Alm
end

export SchwarzschildModes

#SchwarzschildModes(2,6,0) |> print

#0.30105345461250616 - 0.4782769832230698im
end

max_stars_repo_name: Potatoasad/AttemptAtQNM

max_issues_repo_name: Potatoasad/AttemptAtQNM

max_forks_repo_name: Potatoasad/AttemptAtQNM


text:
#####
##### Constant diffusivity
#####

"""
    ConstantIsotropicDiffusivity{FT, K}

Parameters for constant isotropic diffusivity models.
"""
struct ConstantIsotropicDiffusivity{FT, K} <: IsotropicViscosity{FT}
    ν :: FT
    κ :: K
    function ConstantIsotropicDiffusivity{FT}(ν, κ) where FT
        return new{FT, typeof(κ)}(ν, convert_diffusivity(FT, κ))
    end
end

"""
    ConstantIsotropicDiffusivity([FT=Float64;] ν, κ)

Returns parameters for a constant isotropic diffusivity model with constant viscosity `ν`
and constant thermal diffusivities `κ` for each tracer field in `tracers`
`ν` and the fields of `κ` may represent molecular diffusivities in cases that all flow
features are explicitly resovled, or turbulent eddy diffusivities that model the effect of
unresolved, subgrid-scale turbulence.
`κ` may be a `NamedTuple` with fields corresponding
to each tracer, or a single number to be a applied to all tracers.

By default, a molecular viscosity of `ν = 1.05×10⁻⁶` m² s⁻¹ and a molecular thermal
diffusivity of `κ = 1.46×10⁻⁷` m² s⁻¹ is used for each tracer. These molecular values are
the approximate viscosity and thermal diffusivity for seawater at 20°C and 35 psu,
according to Sharqawy et al., "Thermophysical properties of seawater: A review of existing
correlations and data" (2010).
"""
ConstantIsotropicDiffusivity(FT=Float64; ν=ν₀, κ=κ₀) = ConstantIsotropicDiffusivity{FT}(ν, κ)

function with_tracers(tracers, closure::ConstantIsotropicDiffusivity{FT}) where FT
    κ = tracer_diffusivities(tracers, closure.κ)
    return ConstantIsotropicDiffusivity{FT}(closure.ν, κ)
end

calculate_diffusivities!(K, arch, grid, closure::ConstantIsotropicDiffusivity, args...) = nothing

@inline function ∇_κ_∇c(i, j, k, grid, closure::ConstantIsotropicDiffusivity,
                        c, ::Val{tracer_index}, args...) where tracer_index

    @inbounds κ = closure.κ[tracer_index]

    return ∂ⱼκᵢⱼ∂ᵢc(i, j, k, grid, κ, c)
end

@inline ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure::ConstantIsotropicDiffusivity, U, args...) =
    ∂ⱼνᵢⱼ∂ᵢu(i, j, k, grid, closure.ν, U.u)

@inline ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure::ConstantIsotropicDiffusivity, U, args...) =
    ∂ⱼνᵢⱼ∂ᵢv(i, j, k, grid, closure.ν, U.v)

@inline ∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure::ConstantIsotropicDiffusivity, U, args...) =
    ∂ⱼνᵢⱼ∂ᵢw(i, j, k, grid, closure.ν, U.w)