text:
module Observations

using Oceananigans
using Oceananigans: short_show, fields
using Oceananigans.Grids: AbstractGrid
using Oceananigans.Fields
using Oceananigans.Utils: SpecifiedTimes
using JLD2

import Oceananigans.Fields: set!

abstract type AbstractObservation end

include("normalization.jl")

"""
    OneDimensionalTimeSeries{F, G, T, P, M} <: AbstractObservation

A time series of horizontally-averaged observational or LES data
gridded as Oceananigans fields.
"""
struct OneDimensionalTimeSeries{F, G, T, P, M, N} <: AbstractObservation
    field_time_serieses :: F
                   grid :: G
                  times :: T
                   path :: P
               metadata :: M
          normalization :: N
end

obs_str(ts::OneDimensionalTimeSeries) = "OneDimensionalTimeSeries of $(keys(ts.field_time_serieses)) on $(short_show(ts.grid))"

tupleit(t) = try
    Tuple(t)
catch
    tuple(t)
end

const not_metadata_names = ("serialized", "timeseries")

read_group(group::JLD2.Group) = NamedTuple(Symbol(subgroup) => read_group(group[subgroup]) for subgroup in keys(group))
read_group(group) = group

get_grid(field_time_serieses) = field_time_serieses[1].grid

function OneDimensionalTimeSeries(path; field_names, normalize=IdentityNormalization, times=nothing)
    field_names = tupleit(field_names)
    field_time_serieses = NamedTuple(name => FieldTimeSeries(path, string(name); times) for name in field_names)
    grid = get_grid(field_time_serieses)
    times = first(field_time_serieses).times

    # validate_data(fields, grid, times) # might be a good idea to validate the data...
    file = jldopen(path)
    metadata = NamedTuple(Symbol(group) => read_group(file[group]) for group in filter(n -> n ∉ not_metadata_names, keys(file)))
    close(file)

    normalization = Dict(name => normalize(field_time_serieses[name]) for name in keys(field_time_serieses))

    return OneDimensionalTimeSeries(field_time_serieses, grid, times, path, metadata, normalization)
end

observation_times(observation) = observation.times
observation_times(obs::Vector) = observation_times(first(obs))

#####
##### set! for simulation models and observations
#####

default_initial_condition(ts, name) = 0

function set!(model, ts::OneDimensionalTimeSeries, index=1)
    # Set initial condition
    for name in keys(fields(model))

        model_field = fields(model)[name]

        if name in keys(ts.field_time_serieses)
            ts_field = ts.field_time_serieses[name][index]
            set!(model_field, ts_field)
        else
            set!(model_field, 0) #default_initial_condition(ts, Val(name)))
        end
    end

    return nothing
end

"""
    column_ensemble_interior(observations::Vector{<:OneDimensionalTimeSeries}, field_name, time_indices::Vector, N_ens)

Returns an `N_cases × N_ens` array of the interior of a field `field_name` defined on a
`OneDimensionalEnsembleGrid` of size `N_cases × N_ens × Nz`, given a list of `OneDimensionalTimeSeries` objects
containing the `N_cases` single-column fields at the corresponding time index in `time_indices`.
"""
function column_ensemble_interior(observations::Vector{<:OneDimensionalTimeSeries}, field_name, time_indices::Vector, ensemble_size)
    batch = @. get_interior(observations, field_name, time_indices)
    batch = cat(batch..., dims = 2) # (n_batch, n_z)
    return cat([batch for i = 1:ensemble_size]..., dims = 1) # (ensemble_size, n_batch, n_z)
end

struct FieldTimeSeriesCollector{G, D, F, T}
    grid :: G
    field_time_serieses :: D
    collected_fields :: F
    times :: T
end

"""
    FieldTimeSeriesCollector(fields)

Returns a `FieldTimeSeriesCollector` for `fields` of `simulation`.

`fields` is a `NamedTuple` of `AbstractField`s that are to be collected.
"""
function FieldTimeSeriesCollector(collected_fields, times; architecture=CPU())

    grid = collected_fields[1].grid
    field_time_serieses = Dict{Symbol, Any}()

    for name in keys(collected_fields)
        field = collected_fields[name]
        LX, LY, LZ = location(field)
        field_time_series = FieldTimeSeries{LX, LY, LZ}(architecture, field.grid, times)
        field_time_serieses[name] = field_time_series
    end

    # Convert to NamedTuple
    field_time_serieses = NamedTuple(name => field_time_serieses[name] for name in keys(collected_fields))

    return FieldTimeSeriesCollector(grid, field_time_serieses, collected_fields, times)
end

function (collector::FieldTimeSeriesCollector)(simulation)
    for field in collector.collected_fields
        compute!(field)
    end

    current_time = simulation.model.clock.time
    time_index = findfirst(t -> t >= current_time, collector.times)

    for name in keys(collector.collected_fields)
        field_time_series = collector.field_time_serieses[name]
        set!(field_time_series[time_index], collector.collected_fields[name])
    end

    return nothing
end

function initialize_simulation!(simulation, ts, time_series_collector, time_index=1)
    set!(simulation.model, ts, time_index)

    times = observation_times(ts)

    initial_time = ts.times[time_index]
    simulation.model.clock.time = initial_time
    simulation.model.clock.iteration = 0

    # Zero out time series data
    for time_series in time_series_collector.field_time_serieses
        time_series.data .= 0
    end

    simulation.callbacks[:data_collector] = Callback(time_series_collector, SpecifiedTimes(times...))

    simulation.stop_time = times[end]

    return nothing
end

summarize_metadata(::Nothing) = ""
summarize_metadata(metadata) = keys(metadata)

Base.show(io::IO, ts::OneDimensionalTimeSeries) =
    print(io, "OneDimensionalTimeSeries with fields $(propertynames(ts.field_time_serieses))", '\n',
              "├── times: $(ts.times)", '\n',
              "├── grid: $(short_show(ts.grid))", '\n',
              "├── path: \"$(ts.path)\"", '\n',
              "├── metadata: ", summarize_metadata(ts.metadata), '\n',
              "└── normalization: $(short_show(ts.normalization))")

end # module


text:
@doc raw"""
    LinearScalarAdvectionEquation3D

The linear scalar advection equation
```math
\partial_t u + a_1 \partial_1 u + a_2 \partial_2 u + a_3 \partial_3 u = 0
```
in three space dimensions with constant velocity `a`.
"""
struct LinearScalarAdvectionEquation3D{RealT<:Real} <: AbstractLinearScalarAdvectionEquation{3, 1}
  advectionvelocity::SVector{3, RealT}
end

function LinearScalarAdvectionEquation3D(a::NTuple{3,<:Real})
  LinearScalarAdvectionEquation3D(SVector(a))
end

function LinearScalarAdvectionEquation3D(a1::Real, a2::Real, a3::Real)
  LinearScalarAdvectionEquation3D(SVector(a1, a2, a3))
end

varnames(::typeof(cons2cons), ::LinearScalarAdvectionEquation3D) = ("scalar", )
varnames(::typeof(cons2prim), ::LinearScalarAdvectionEquation3D) = ("scalar", )

# Set initial conditions at physical location `x` for time `t`
"""
    initial_condition_constant(x, t, equations::LinearScalarAdvectionEquation1D)

A constant initial condition to test free-stream preservation.
"""
function initial_condition_constant(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  return SVector(2.0)
end

"""
    initial_condition_convergence_test(x, t, equations::LinearScalarAdvectionEquation1D)

A smooth initial condition used for convergence tests.
"""
function initial_condition_convergence_test(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  c = 1.0
  A = 0.5
  L = 2
  f = 1/L
  omega = 2 * pi * f
  scalar = c + A * sin(omega * sum(x_trans))
  return SVector(scalar)
end

"""
    initial_condition_gauss(x, t, equations::LinearScalarAdvectionEquation1D)

A Gaussian pulse.
"""
function initial_condition_gauss(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  scalar = exp(-(x_trans[1]^2 + x_trans[2]^2 + x_trans[3]^2))
  return SVector(scalar)
end

"""
    initial_condition_sin(x, t, equations::LinearScalarAdvectionEquation1D)

A sine wave in the conserved variable.
"""
function initial_condition_sin(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  scalar = sin(2 * pi * x_trans[1]) * sin(2 * pi * x_trans[2]) * sin(2 * pi * x_trans[3])
  return SVector(scalar)
end

"""
    initial_condition_linear_z(x, t, equations::LinearScalarAdvectionEquation1D)

A linear function of `x[3]` used together with
[`boundary_condition_linear_z`](@ref).
"""
function initial_condition_linear_z(x, t, equation::LinearScalarAdvectionEquation3D)
  # Store translated coordinate for easy use of exact solution
  x_trans = x - equation.advectionvelocity * t

  return SVector(x_trans[3])
end

"""
    boundary_condition_linear_z(u_inner, orientation, direction, x, t,
                                surface_flux_function,
                                equation::LinearScalarAdvectionEquation1D)

Boundary conditions for
[`boundary_condition_linear_z`](@ref).
"""
function boundary_condition_linear_z(u_inner, orientation, direction, x, t,
                                     surface_flux_function,
                                     equation::LinearScalarAdvectionEquation3D)
  u_boundary = initial_condition_linear_z(x, t, equation)

  # Calculate boundary flux
  if direction in (2, 4, 6) # u_inner is "left" of boundary, u_boundary is "right" of boundary
    flux = surface_flux_function(u_inner, u_boundary, orientation, equation)
  else # u_boundary is "left" of boundary, u_inner is "right" of boundary
    flux = surface_flux_function(u_boundary, u_inner, orientation, equation)
  end

  return flux
end

# Pre-defined source terms should be implemented as
# function source_terms_WHATEVER(u, x, t, equation::LinearScalarAdvectionEquation3D)

# Calculate 1D flux in for a single point
@inline function flux(u, orientation::Integer, equation::LinearScalarAdvectionEquation3D)
  a = equation.advectionvelocity[orientation]
  return a * u
end

# Calculate maximum wave speed for local Lax-Friedrichs-type dissipation
@inline function max_abs_speed_naive(u_ll, u_rr, orientation::Integer, equation::LinearScalarAdvectionEquation3D)
  λ_max = abs(equation.advectionvelocity[orientation])
end

# Calculate 1D flux for a single point in direction of a normal vector
@inline function flux(u, normal_vector::AbstractVector, equation::LinearScalarAdvectionEquation3D)
  a = dot(equation.advectionvelocity, normal_vector) # velocity in normal direction
  return a * u
end

# Calculate maximum wave speed in direction of a normal vector for local Lax-Friedrichs-type dissipation
@inline function max_abs_speed_naive(u_ll, u_rr, normal_vector::AbstractVector, equation::LinearScalarAdvectionEquation3D)
  a = dot(equation.advectionvelocity, normal_vector) # velocity in normal direction
  return abs(a)
end

@inline have_constant_speed(::LinearScalarAdvectionEquation3D) = Val(true)

@inline function max_abs_speeds(equation::LinearScalarAdvectionEquation3D)
  return abs.(equation.advectionvelocity)
end

# Convert conservative variables to primitive
@inline cons2prim(u, equation::LinearScalarAdvectionEquation3D) = u

# Convert conservative variables to entropy variables
@inline cons2entropy(u, equation::LinearScalarAdvectionEquation3D) = u

# Calculate entropy for a conservative state `cons`
@inline entropy(u::Real, ::LinearScalarAdvectionEquation3D) = 0.5 * u^2
@inline entropy(u, equation::LinearScalarAdvectionEquation3D) = entropy(u[1], equation)

# Calculate total energy for a conservative state `cons`
@inline energy_total(u::Real, ::LinearScalarAdvectionEquation3D) = 0.5 * u^2
@inline energy_total(u, equation::LinearScalarAdvectionEquation3D) = energy_total(u[1], equation)

ext: jl


text:
module TweetNaCl

import BinDeps

type SignatureFailure <: Exception end

#@BinDeps.load_dependencies
if( isfile(joinpath(dirname(@__FILE__), "..", "deps", "deps.jl")) )
    include("../deps/deps.jl")
else
    error("TweetNaCl not properly built.  Please run Pkg.build(\"TweetNaCl\").")
end

nacl_sign_public_key_length = 0;
nacl_sign_secret_key_length = 0;
nacl_sign_max_extra_length  = 0;

function __init__()
    global nacl_sign_public_key_length,
           nacl_sign_secret_key_length,
           nacl_sign_max_extra_length

    retval = ccall( (:init_randombytes, tweetnacl), Cint, () )
    if retval != 0
        error("Failed to initialise cryptographic PRNG.")
    end

    nacl_sign_public_key_length = ccall(
        (:nacl_get_sign_pk_length, tweetnacl), Cint, () )

    nacl_sign_secret_key_length = ccall(
        (:nacl_get_sign_sk_length, tweetnacl), Cint, () )

    nacl_sign_max_extra_length = ccall(
        (:nacl_get_sign_sig_extra_length, tweetnacl), Cint, () )

end

function randombytes!(str::Array{UInt8})
    ccall( (:randombytes, TweetNaCl.tweetnacl), Void, (Ptr{UInt8}, Culonglong), str, length(str))
end

function randombytes(count::Int64)
    str = Array(UInt8, count)
    randombytes!(str)
    return str
end

function crypto_sign_keypair()
    public_key = Array(UInt8, nacl_sign_public_key_length)
    secret_key = Array(UInt8, nacl_sign_secret_key_length)

    ccall( (:crypto_sign_ed25519_tweet_keypair, TweetNaCl.tweetnacl), Cint,
           (Ptr{UInt8}, Ptr{UInt8}), public_key, secret_key)

    return public_key, secret_key
end

function crypto_sign(message::AbstractString, secret_key)

    if(length(secret_key) != TweetNaCl.nacl_sign_secret_key_length)
        error("crypto_sign: Secret key is of incorrect length.")
    end

    signature = zeros(UInt8, length(message)+nacl_sign_max_extra_length)
    signature_length = zeros(UInt64, 1)

    ccall( (:crypto_sign_ed25519_tweet, TweetNaCl.tweetnacl), Cint,
           (Ptr{UInt8}, Ptr{Culonglong}, Ptr{UInt8}, Culonglong, Ptr{UInt8}),
           signature, signature_length, message, length(message), secret_key)

    return signature[1:signature_length[1]]
end

function crypto_sign_open(signature::Array{UInt8}, public_key)

    if(length(public_key) != TweetNaCl.nacl_sign_public_key_length)
        error("crypto_sign_open: Public key is of incorrect length.")
    end

    message = zeros(UInt8, length(signature))
    message_length = zeros(UInt64, 1)

    result = ccall( (:crypto_sign_ed25519_tweet_open, TweetNaCl.tweetnacl), Cint,
           (Ptr{UInt8}, Ptr{Culonglong}, Ptr{UInt8}, Culonglong, Ptr{UInt8}),
           message, message_length, signature, length(signature), public_key)

    if result != 0
        throw(SignatureFailure())
    end

    return message[1:message_length[1]]

end

end # module

max_stars_repo_name: LachlanGunn/TweetNaCl

max_issues_repo_name: LachlanGunn/TweetNaCl

max_forks_repo_name: LachlanGunn/TweetNaCl


text:
function run(s)
    # Your code here
    lines = eachline(IOBuffer(s))
    rules = Dict{SubString, SubString}()
    for line in lines
        if isempty(line)
            break
        end
        (idx, rule) = split(line, ": ", limit=2)
        rules[idx] = rule
    end
    memory = Dict{SubString,String}()
    function resolve_rule(rule::SubString)::String
        # memoize here
        get!(memory, rule) do
            if '"' in rule
                return rule[2:2]
            elseif '|' in rule
                left, right = split(rule, " | ")
                return "(?:$(resolve_rule(left))|$(resolve_rule(right)))"
            else
                parts = split(rule, " ")
                return join((resolve_rule(rules[i]) for i in parts), "")
            end
        end
    end

    regex_str = resolve_rule(rules["0"])
    regex = Regex("^$(regex_str)\$")
    return sum(occursin.(regex, lines))
end

#########################################

function main()
    run(ARGS[1])
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time * 1000)")
    println(res)
end

main()


text:
# Generate sequences of the Poisson events and plot histogram to check predictive distribution

include("./metadata.jl")
using Metadata

function run_sim()
    n_sim = 500
    n = 100
    w = 30
    seqs = cumsum(rand(1 : w, n))
    for i in 1 : n_sim
        new_dat = cumsum(rand(1 : w, n))
        seqs = vcat(seqs, new_dat)
    end
    println("sim_poisson_hist.jl > Save generated data ino ", outdir*"sim_poisson_hist.txt")
    writedlm(outdir*"sim_poisson_hist.txt", seqs)
end
run_sim()

max_stars_repo_name: av-maslov/Recurrency

max_issues_repo_name: av-maslov/Recurrency

max_forks_repo_name: av-maslov/Recurrency


text:
module Test_bouncingBall

#=
ContinuousCallback: affect! is not always called at the crossed condition

https://github.com/JuliaDiffEq/DifferentialEquations.jl/issues/551

Event restart: Stepsize not adapted for Tsit5 #558

https://github.com/JuliaDiffEq/DifferentialEquations.jl/issues/558

Fixed in the model by adding
    auto_dt_reset!(integrator)
    set_proposed_dt!(integrator, integrator.dt)
=#

using DifferentialEquations, Sundials, PyPlot

function f(du,u,p,t)
    du[1] = u[2]
    du[2] = -9.81
end

function condition(u,t,integrator) # Event when event_f(u,t) == 0
    z = u[1] + 1e-12
    # println("condition called at time = ", t, ", z = ", z)
    return z
end

function affect_neg(integrator)
    println("Event (affect_neg) at time = ", integrator.t, ", h = ", integrator.u[1])
    integrator.u[2] = -0.8*integrator.u[2]
    #set_proposed_dt!(integrator, 1e-5)
    #println("    after set_proposed_dt!: integrator.dt = ", integrator.dt)
    println("   before reset: integrator.dt = ", integrator.dt)
    auto_dt_reset!(integrator)
    println("   after reset : integrator.dt = ", integrator.dt)
    set_proposed_dt!(integrator, integrator.dt)
end

cb = ContinuousCallback(condition,nothing,affect_neg! = affect_neg)

u0 = [1.0,0.0]
tspan = (0.0,3.0)
prob = ODEProblem(f,u0,tspan)
sol = solve(prob,Tsit5(),saveat=0.01,callback=cb)
#sol = solve(prob,CVODE_BDF(),saveat=0.01,callback=cb)
#=
println("propertynames(prob) = ", propertynames(prob))
println("propertynames(sol) = ", propertynames(sol))
println("propertynames(sol.prob) = ", propertynames(sol.prob))
println("typeof(sol.prob.problem_type) = ", typeof(sol.prob.problem_type))
println("typeof(sol.alg) = ", typeof(sol.alg))
=#

figure(1)
clf()
plot(sol.t, getindex.(sol.u, 1), label="h")

end


text:
module CImGuiFrontEnd

export DoGui

# Utils
export TreeNode,
    WithColumns,
    WithStyleColors,
    @SameLine

using Reexport
using ArgCheck

@reexport using CImGui
# @reexport using CImGui.CSyntax
# @reexport using CImGui.CSyntax.CStatic
using CImGui.GLFWBackend
using CImGui.OpenGLBackend
using CImGui.GLFWBackend.GLFW
using CImGui.OpenGLBackend.ModernGL

g_glfw_window = nothing

const default_clear_color = Cfloat[0.45, 0.55, 0.60, 1.00]

function DoGui(func ; width=1280, height=720,
               name="Demo",
               clear_color=default_clear_color, style=:light,
               vsync=true,
               ini_filename=C_NULL)
    # OpenGL 3.0 + GLSL 130
    glsl_version = 130
    GLFW.WindowHint(GLFW.CONTEXT_VERSION_MAJOR, 3)
    GLFW.WindowHint(GLFW.CONTEXT_VERSION_MINOR, 0)
    # GLFW.WindowHint(GLFW.OPENGL_PROFILE, GLFW.OPENGL_CORE_PROFILE) # 3.2+ only
    # GLFW.WindowHint(GLFW.OPENGL_FORWARD_COMPAT, GL_TRUE) # 3.0+ only

    # setup GLFW error callback
    error_callback(err::GLFW.GLFWError) = @error "GLFW ERROR: code $(err.code) msg: $(err.description)"
    GLFW.SetErrorCallback(error_callback)

    # create window
    window = GLFW.CreateWindow(width, height, name)
    @assert window != C_NULL
    GLFW.MakeContextCurrent(window)

    vsync && GLFW.SwapInterval(1)

    # setup Dear ImGui context
    ctx = CImGui.CreateContext()

    # setup Dear ImGui style
    if style == :dark
        CImGui.StyleColorsDark()
    elseif style == :classic
        CImGui.StyleColorsClassic()
    elseif style == :light
        CImGui.StyleColorsLight()
    else
        error("Unknown style $style")
    end

    # setup Platform/Renderer bindings
    ImGui_ImplGlfw_InitForOpenGL(window, true)
    ImGui_ImplOpenGL3_Init(glsl_version)

    io = CImGui.GetIO()
    io.IniFilename = ini_filename

    try
        while !GLFW.WindowShouldClose(window)
            GLFW.PollEvents()
            # start the Dear ImGui frame
            ImGui_ImplOpenGL3_NewFrame()
            ImGui_ImplGlfw_NewFrame()
            CImGui.NewFrame()

            global g_glfw_window = window
            func()

            # rendering
            CImGui.Render()

            GLFW.MakeContextCurrent(window)
            display_w, display_h = GLFW.GetFramebufferSize(window)
            glViewport(0, 0, display_w, display_h)
            glClearColor(clear_color...)
            glClear(GL_COLOR_BUFFER_BIT)

            ImGui_ImplOpenGL3_RenderDrawData(CImGui.GetDrawData())

            GLFW.MakeContextCurrent(window)
            GLFW.SwapBuffers(window)

            yield()
        end
    finally
        # cleanup
        ImGui_ImplOpenGL3_Shutdown()
        ImGui_ImplGlfw_Shutdown()
        CImGui.DestroyContext(ctx)

        GLFW.DestroyWindow(window)
    end
end

##############################
# * Utils
#----------------------------
FrameSize(window=g_glfw_window) = GLFW.GetFramebufferSize(window)

######################################
# * Context wrappers
#------------------------------------

"""
    TreeNode(func, name)
    TreeNode(func, name, always_after)

Intended to be used as `TreeNode("label") do ... end`

The second form allows for extra commands to be given in `always_after` which will
occur immediately after the call to `CImGui.TreeNode`. This is necessary when
using other features such as overlays or drag and drop.
"""
function TreeNode(func, name, always_after=()->nothing)
    ret = CImGui.TreeNode(name)
    always_after()
    if ret
        func()
        CImGui.TreePop()
    end
end

function WithColumns(func, num, headings=nothing, widths=nothing ; separator=true, end_separator=separator)
    CImGui.Columns(num)

    if headings != nothing
        widths == nothing && (widths = fill(nothing, num))

        @argcheck length(headings) == num
        for (heading,width) in zip(headings,widths)
            CImGui.Text(heading)

            cur_width = CImGui.GetColumnWidth()
            width == :min && (width = CImGui.GetItemRectSize().x)
            width != nothing && width < cur_width && CImGui.SetColumnWidth(-1, width)

            CImGui.NextColumn()
        end
    end
    separator && CImGui.Separator()

    func()

    CImGui.Columns(1)
    end_separator && CImGui.Separator()
end

function WithStyleColors(func, args::Pair...)
    for (name,col) in args
        CImGui.PushStyleColor(name, col)
    end

    func()

    CImGui.PopStyleColor(length(args))
end

##########################################
# * Other useful tools
#----------------------------------------

"""
    @SameLine block

Places `CImGui.SameLine()` after every statement, except for the last.
"""
macro SameLine(ex)
    @argcheck ex.head == :block

    old_args = ex.args

    ex.args = Expr[]

    last = true

    for item in reverse(old_args)
        if !isa(item, LineNumberNode)
            !last && pushfirst!(ex.args, :(CImGui.SameLine()))
            last = false
        end
        pushfirst!(ex.args, item)
    end

    ex.args = esc.(ex.args)
    ex
end

end


text:
# TODO: Improve docstring https://docs.julialang.org/en/v1/manual/documentation/
"stage_states(mdp::MDP, epoch::Int64)"
function stage_states end

# TODO: Improve docstring
"stage_stateindex(mdp::MDP, s, epoch::Int64)::Integer"
function stage_stateindex end

# TODO: Docstring
abstract type HorizonLength end

"If HorizonLength(m::Union{MDP,POMDP}) == FiniteHorizon(), horizon(m) should be implemented and return an integer"
struct FiniteHorizon <: HorizonLength end
"HorizonLength(m::Union{MDP,POMDP}) == InfiniteHorizon() indicates that horizon(m) should not be called."
struct InfiniteHorizon <: HorizonLength end

HorizonLength(m::Union{MDP,POMDP}) = HorizonLength(typeof(m))
HorizonLength(::Type{<:Union{MDP,POMDP}}) = InfiniteHorizon()

# TODO Improve docstring
"""
    horizon(m::Union{MDP,POMDP})::Integer

Return the number of *steps* that will be taken in the (PO)MDP.

A simulation of a (PO)MDP with `horizon(m) == d` should contain *d+1* states and *d* actions and rewards.
"""
function horizon end

"""
    epoch(s::State)::Integer

Return the epoch of given state
"""
function epoch end


text:
using RoME
using RoMEPlotting

fg = initfg()

addVariable!(fg, :l2, Point2)
addFactor!(fg, [:l2;], Prior(MvNormal([0.0;0],Matrix(Diagonal([0.01;0.01])))), graphinit=false)

addVariable!(fg, :x13, Pose2)
p2br = Pose2Point2BearingRange(Normal(0,0.03),Normal(20.0,0.5))
addFactor!(fg, [:x13;:l2], p2br, graphinit=false)

tree = batchSolve!(fg, treeinit=true, drawpdf=true)

drawPosesLandms(fg)

plotPose(fg, :x13)

#


text:
export kx_tc_sp_coo, kx_tc_sp, kx_ct_sp, zx_tc_sp_coo, zx_tc_sp, zx_ct_sp, ε⁻¹_sp_coo
export ε⁻¹_sp, nng⁻¹_sp, M̂_sp, M̂ₖ_sp, 𝓕_dense, 𝓕⁻¹_dense

using SparseArrays
using Zygote: Buffer

# both rrules below copied from
# https://discourse.julialang.org/t/how-to-efficiently-differentiate-backslash-operator-for-sparse-matrix/43778/18

# function ChainRulesCore.rrule(::typeof(SparseArrays.SparseMatrixCSC),
#     m::Integer, n::Integer, pp::Vector, ii::Vector, Av::Vector)
#     A = SparseMatrixCSC(m,n,pp,ii,Av)
#     function SparseMatrixCSC_pullback(dA)
#         # Pick out the entries in `dA` corresponding to nonzeros in `A`
#         dAv = Vector{eltype(dA)}(undef, length(Av))
#         for j = 1:n, p = pp[j]:pp[j+1]-1
#             dAv[p] = dA[ii[p],j]
#         end
#         return (NO_FIELDS, DoesNotExist(), DoesNotExist(), DoesNotExist(), DoesNotExist(), dAv)
#     end
#
#     return A, SparseMatrixCSC_pullback
# end

function ChainRulesCore.rrule(::typeof(SparseArrays.SparseMatrixCSC),
    m::Integer, n::Integer, pp::Vector, ii::Vector, Av::Vector)
    A = SparseMatrixCSC(m,n,pp,ii,Av)
    function SparseMatrixCSC_pullback(dA::AbstractMatrix)
        dAv = Vector{eltype(dA)}(undef, length(Av))
        for j = 1:n, p = pp[j]:pp[j+1]-1
            dAv[p] = dA[ii[p],j]
        end
        return (NO_FIELDS, DoesNotExist(), DoesNotExist(), DoesNotExist(), DoesNotExist(), dAv)
    end
    function SparseMatrixCSC_pullback(dA::SparseMatrixCSC)
        @assert getproperty.(Ref(A), (:m,:n,:colptr,:rowval)) == getproperty.(Ref(dA), (:m,:n,:colptr,:rowval))
        return (NO_FIELDS, DoesNotExist(), DoesNotExist(), DoesNotExist(), DoesNotExist(), dA.nzval)
    end

    return A, SparseMatrixCSC_pullback
end

function kx_tc_sp_coo(mag::AbstractArray{T,2},mn) where T<:Real
    Nx, Ny = size(mag)
    I = Int32[]
    J = Int32[]
    V = Float64[]
    # kxt2c_matrix_buf = Zygote.bufferfrom(zeros(Complex{T},(3*Nx*Ny),(2*Nx*Ny)))
    for ix=1:Nx,iy=1:Ny,a=1:3 #,b=1:2
        q =  Nx * (iy - 1) + ix

        # reference from kcross_t2c!:
        #       ds.d[1,i,j,k] = ( ds.H[1,i,j,k] * ds.mn[1,2,i,j,k] - ds.H[2,i,j,k] * ds.mn[1,1,i,j,k] ) * -ds.kpg_mag[i,j,k]
        #       ds.d[2,i,j,k] = ( ds.H[1,i,j,k] * ds.mn[2,2,i,j,k] - ds.H[2,i,j,k] * ds.mn[2,1,i,j,k] ) * -ds.kpg_mag[i,j,k]
        #       ds.d[3,i,j,k] = ( ds.H[1,i,j,k] * ds.mn[3,2,i,j,k] - ds.H[2,i,j,k] * ds.mn[3,1,i,j,k] ) * -ds.kpg_mag[i,j,k]
        # which implements d⃗ = k×ₜ₂c ⋅ H⃗
        # Here we want to explicitly define the matrix k×ₜ₂c
        # the general indexing scheme:
        # kxt2c_matrix_buf[ (3*q-2)+a-1 ,(2*q-1) + (b-1) ] <==> mn[a,b,ix,iy,iz], mag[ix,iy,iz]
        # b = 1  ( m⃗ )
        # kxt2c_matrix_buf[(3*q-2)+a-1,(2*q-1)] = mn[a,2,ix,iy] * mag[ix,iy]
		push!(I,(3*q-2)+a-1)
		push!(J,2*q-1)
		push!(V,mn[2,a,ix,iy] * mag[ix,iy])
	        # b = 2  ( n⃗ )
	        # kxt2c_matrix_buf[(3*q-2)+a-1,(2*q-1)+1] = mn[a,1,ix,iy] * -mag[ix,iy]
		push!(I,(3*q-2)+a-1)
		push!(J,(2*q-1)+1)
		push!(V,mn[1,a,ix,iy] * -mag[ix,iy])
    end
    # return copy(kxt2c_matrix_buf)
    return sparse(I,J,V)
end

function kx_tc_sp(mag::AbstractArray{T,2},mn::AbstractArray{T,4}) where T<:Real
    Nx, Ny = size(mag)
    NN = length(mag)
    V = Buffer(mag,0)
    for iy=1:Ny, ix=1:Nx
    	push!(V,(mn[2,:,ix,iy] * mag[ix,iy])...)
    	push!(V,(mn[1,:,ix,iy] * -mag[ix,iy])...)
    end
    SparseMatrixCSC(
    	3*NN,	# m
    	2*NN,	# n
    	collect(Int32,1:3:6*NN+3), # colptr
    	convert(Vector{Int32},repeat(0:NN-1,inner=6).*3 .+ repeat([1,2,3,1,2,3],NN)), # rowval
    	copy(V),	# nzval
    )
end

function kx_tc_sp(mag::AbstractArray{T,2},m⃗::AbstractArray{SVector{3,T},2},n⃗::AbstractArray{SVector{3,T},2}) where T<:Real
    Nx, Ny = size(mag)
    NN = length(mag)
    V = Buffer(mag,0)
    for iy=1:Ny, ix=1:Nx
    	push!(V,(n⃗[ix,iy] * mag[ix,iy])...)
    	push!(V,(m⃗[ix,iy] * -mag[ix,iy])...)
    end
    SparseMatrixCSC(
    	3*NN,	# m
    	2*NN,	# n
    	collect(Int32,1:3:6*NN+3), # colptr
    	convert(Vector{Int32},repeat(0:NN-1,inner=6).*3 .+ repeat([1,2,3,1,2,3],NN)), # rowval
    	copy(V),	# nzval
    )
end

function kx_tc_sp(k,grid::Grid{ND,T}) where {ND,T<:Real}
    mag,m⃗,n⃗ = mag_m_n(k,g⃗(grid))
    Ns = size(grid)
    NN = N(grid)
    V = Buffer(mag,0)
    for iy=1:Ns[2], ix=1:Ns[1]
    	push!(V,(n⃗[ix,iy] * mag[ix,iy])...)
    	push!(V,(m⃗[ix,iy] * -mag[ix,iy])...)
    end
    SparseMatrixCSC(
    	3*NN,	# m
    	2*NN,	# n
    	collect(Int32,1:3:6*NN+3), # colptr
    	convert(Vector{Int32},repeat(0:NN-1,inner=6).*3 .+ repeat([1,2,3,1,2,3],NN)), # rowval
    	copy(V),	# nzval
    )
end

kx_ct_sp(k,grid::Grid) = -kx_tc_sp(k,grid::Grid)'

function zx_tc_sp_coo(mag::AbstractArray{T,2},mn) where T<:Real
    Nx, Ny = size(mn)[3:4]
    NN = Nx * Ny
    I = Int32[]
    J = Int32[]
    V = Float64[]
    for ix=1:Nx,iy=1:Ny
        q =  Nx * (iy - 1) + ix
		# reference from zcross_t2c!:
		#          ds.d[1,i,j,k] = -Hin[1,i,j,k] * ds.mn[2,1,i,j,k] - Hin[2,i,j,k] * ds.mn[2,2,i,j,k]
		#          ds.d[2,i,j,k] =  Hin[1,i,j,k] * ds.mn[1,1,i,j,k] + Hin[2,i,j,k] * ds.mn[1,2,i,j,k]
		#          ds.d[3,i,j,k] = 0
		# which implements d⃗ = z×ₜ₂c ⋅ H⃗
		# Here we want to explicitly define the matrix z×ₜ₂c
		# the general indexing scheme:
		# zxt2c_matrix_buf[ (3*q-2)+a-1 ,(2*q-1) + (b-1) ] <==> mn[a,b,ix,iy,iz]
		# a = 1  ( x̂ ), b = 1  ( m⃗ )
		# zxt2c_matrix_buf[(3*q-2),(2*q-1)] = -mn[2,1,ix,iy,iz]
		push!(I,3*q-2)
		push!(J,2*q-1)
		push!(V,-mn[1,2,ix,iy])
		# a = 1  ( x̂ ), b = 2  ( n⃗ )
		# zxt2c_matrix_buf[(3*q-2),2*q] = -mn[2,2,ix,iy,iz]
		push!(I,3*q-2)
		push!(J,2*q)
		push!(V,-mn[2,2,ix,iy])
		# a = 2  ( ŷ ), b = 1  ( m⃗ )
		# zxt2c_matrix_buf[(3*q-2)+1,(2*q-1)] = mn[1,1,ix,iy,iz]
		push!(I,3*q-1)
		push!(J,2*q-1)
		push!(V,mn[1,1,ix,iy])
		# a = 2  ( ŷ ), b = 2  ( n⃗ )
		# zxt2c_matrix_buf[(3*q-2)+1,2*q] = mn[1,2,ix,iy,iz]
		push!(I,3*q-1)
		push!(J,2*q)
		push!(V,mn[2,1,ix,iy])
    end
    return sparse(I,J,V,3*NN,2*NN)
end

function zx_tc_sp(mag::AbstractArray{T,2},m⃗::AbstractArray{SVector{3,T},2},n⃗::AbstractArray{SVector{3,T},2}) where T<:Real
    Nx, Ny = size(mag)
    NN = length(mag)
    V = Buffer(mag,0)
    for iy=1:Ny, ix=1:Nx
    	push!(V, -m⃗[ix,iy][2] )
		push!(V, m⃗[ix,iy][1] )
		push!(V, -n⃗[ix,iy][2] )
		push!(V, n⃗[ix,iy][1] )
    end
    SparseMatrixCSC(
    	3*NN,	# m
    	2*NN,	# n
		collect(Int32,1:2:4*NN+2), # colptr
		convert(Vector{Int32},repeat(0:NN-1,inner=4).*3 .+ repeat([1,2,1,2],NN)), # rowval
    	copy(V),	# nzval
    )
end

function zx_tc_sp(k,grid::Grid{ND,T}) where {ND,T<:Real}
    mag,m⃗,n⃗ = mag_m_n(k,g⃗(grid))
    Ns = size(grid)
    NN = N(grid)
    V = Buffer(mag,0)
	for iy=1:Ns[2], ix=1:Ns[1]
    	push!(V, -m⃗[ix,iy][2] )
		push!(V, m⃗[ix,iy][1] )
		push!(V, -n⃗[ix,iy][2] )
		push!(V, n⃗[ix,iy][1] )
    end
    SparseMatrixCSC(
    	3*NN,	# m
    	2*NN,	# n
		collect(Int32,1:2:4*NN+2), # colptr
		convert(Vector{Int32},repeat(0:NN-1,inner=4).*3 .+ repeat([1,2,1,2],NN)), # rowval
    	copy(V),	# nzval
    )
end

zx_ct_sp(k,grid::Grid) = zx_tc_sp(k,grid::Grid)'

function ε⁻¹_sp_coo(ω,geom::Geometry,grid::Grid) #(ω,geom::Vector{<:Shape},grid::Grid{ND,T})
	ei = εₛ⁻¹(ω,geom,grid) #wl doesnt affect anything for constant indices  #make_εₛ⁻¹(shapes,grid)
	Ns = size(grid)
    NN = N(grid)
    I = Int32[]
    J = Int32[]
    V = Float64[]
	for i=1:Ns[1],j=1:Ns[2],a=1:3,b=1:3
		q = (Ny * (j-1) + i)
		eiv = ei[i,j][a,b]
		# if !iszero(eiv)
			# ei_matrix_buf[(3*q-2)+a-1,(3*q-2)+b-1] = ei_field[a,b,i,j,1]
		push!(I,(3*q-2)+a-1)
		push!(J,(3*q-2)+b-1)
		push!(V,eiv)
		# end
	end
	return sparse(I,J,V,3*NN,3*NN)
end

function ε⁻¹_sp(ω,geom::Geometry,grid::Grid) #(ω,geom::Vector{<:Shape},grid::Grid{ND,T})
	ei = εₛ⁻¹(ω,geom,grid) #wl doesnt affect anything for constant indices  #make_εₛ⁻¹(shapes,grid)
	Ns = size(grid)
    NN = N(grid)
    V = Buffer([3.2],0)
	for iy=1:Ns[2], ix=1:Ns[1] #,a=1:3,b=1:3
		push!(V,ei[ix,iy]...) #[a,b])
	end
	SparseMatrixCSC(
    	3*NN,	# m
    	3*NN,	# n
    	collect(Int32,1:3:9*NN+3), # colptr
    	convert(Vector{Int32},repeat(0:NN-1,inner=9).*3 .+ repeat([1,2,3,1,2,3,1,2,3],NN)), # rowval
    	copy(V),	# nzval
    )
end

function nng⁻¹_sp(ω,geom::Geometry,grid::Grid) #(ω,geom::Vector{<:Shape},grid::Grid{ND,T})
	nnginv = nngₛ⁻¹(ω,geom,grid)
	Ns = size(grid)
    NN = N(grid)
    V = Buffer([3.2],0)
	for iy=1:Ns[2], ix=1:Ns[1] #,a=1:3,b=1:3
		push!(V,nnginv[ix,iy]...) #[a,b])
	end
	SparseMatrixCSC(
    	3*NN,	# m
    	3*NN,	# n
    	collect(Int32,1:3:9*NN+3), # colptr
    	convert(Vector{Int32},repeat(0:NN-1,inner=9).*3 .+ repeat([1,2,3,1,2,3,1,2,3],NN)), # rowval
    	copy(V),	# nzval
    )
end

function nng⁻¹_sp(nnginv,grid::Grid) #(ω,geom::Vector{<:Shape},grid::Grid{ND,T})
	Ns = size(grid)
    NN = N(grid)
    V = Buffer([3.2],0)
	for iy=1:Ns[2], ix=1:Ns[1] #,a=1:3,b=1:3
		push!(V,nnginv[ix,iy]...) #[a,b])
	end
	SparseMatrixCSC(
    	3*NN,	# m
    	3*NN,	# n
    	collect(Int32,1:3:9*NN+3), # colptr
    	convert(Vector{Int32},repeat(0:NN-1,inner=9).*3 .+ repeat([1,2,3,1,2,3,1,2,3],NN)), # rowval
    	copy(V),	# nzval
    )
end

function M̂_sp(ω,k,geom,grid::Grid{2})
	Ns = size(grid)
	Ninv = 1. / N(grid)
	kxtcsp = kx_tc_sp(k,grid)
	eisp = ε⁻¹_sp(ω,geom,grid)
	𝓕 = Zygote.ignore() do
        LinearMap{ComplexF64}(d::AbstractVector{ComplexF64} -> vec(fft(reshape(d,(3,Ns...)),(2:3)))::AbstractVector{ComplexF64},*(3,Ns...),ishermitian=true,ismutating=false)
    end
    𝓕⁻¹ = Zygote.ignore() do
        LinearMap{ComplexF64}(d::AbstractVector{ComplexF64} -> vec(bfft(reshape(d,(3,Ns...)),(2:3)))::AbstractVector{ComplexF64},*(3,Ns...),ishermitian=true,ismutating=false)
    end
	Ninv * kxtcsp' * 𝓕⁻¹ * eisp * 𝓕 * kxtcsp
end

function M̂ₖ_sp(ω,k,geom::Geometry,grid::Grid{2})
	Ns = size(grid)
	Ninv = 1. / N(grid)
	kxtcsp = kx_tc_sp(k,grid)
	zxtcsp = zx_tc_sp(k,grid)
	eisp = nng⁻¹_sp(ω,geom,grid) #ε⁻¹_sp(ω,geom,grid)
	𝓕 = Zygote.ignore() do
        LinearMap{ComplexF64}(d::AbstractVector{ComplexF64} -> vec(fft(reshape(d,(3,Ns...)),(2:3)))::AbstractVector{ComplexF64},*(3,Ns...),ishermitian=true,ismutating=false)
    end
    𝓕⁻¹ = Zygote.ignore() do
        LinearMap{ComplexF64}(d::AbstractVector{ComplexF64} -> vec(bfft(reshape(d,(3,Ns...)),(2:3)))::AbstractVector{ComplexF64},*(3,Ns...),ishermitian=true,ismutating=false)
    end
	Ninv * kxtcsp' * 𝓕⁻¹ * eisp * 𝓕 * zxtcsp
end

function M̂ₖ_sp(k,nnginv,grid::Grid{2})
	Ns = size(grid)
	Ninv = 1. / N(grid)
	kxtcsp = kx_tc_sp(k,grid)
	zxtcsp = zx_tc_sp(k,grid)
	nnginvsp = nng⁻¹_sp(nnginv,grid)
	𝓕 = Zygote.ignore() do
        LinearMap{ComplexF64}(d::AbstractVector{ComplexF64} -> vec(fft(reshape(d,(3,Ns...)),(2:3)))::AbstractVector{ComplexF64},*(3,Ns...),ishermitian=true,ismutating=false)
    end
    𝓕⁻¹ = Zygote.ignore() do
        LinearMap{ComplexF64}(d::AbstractVector{ComplexF64} -> vec(bfft(reshape(d,(3,Ns...)),(2:3)))::AbstractVector{ComplexF64},*(3,Ns...),ishermitian=true,ismutating=false)
    end
	Ninv * kxtcsp' * 𝓕⁻¹ * nnginvsp * 𝓕 * zxtcsp
end

𝓕_dense(grid::Grid{3}) = Matrix(LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(fft(reshape(d,(3,grid.Nx,grid.Ny,grid.Nz)),(2:4)))::AbstractArray{ComplexF64,1},*(3,grid.Nx,grid.Ny,grid.Nz),ishermitian=false,ismutating=false))
𝓕⁻¹_dense(grid::Grid{3}) = Matrix(LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(bfft(reshape(d,(3,grid.Nx,grid.Ny,grid.Nz)),(2:4)))::AbstractArray{ComplexF64,1},*(3,grid.Nx,grid.Ny,grid.Nz),ishermitian=false,ismutating=false))

𝓕_dense(grid::Grid{2}) = Matrix(LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(fft(reshape(d,(3,grid.Nx,grid.Ny)),(2:3)))::AbstractArray{ComplexF64,1},*(3,grid.Nx,grid.Ny),ishermitian=false,ismutating=false))
𝓕⁻¹_dense(grid::Grid{2}) = Matrix(LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(bfft(reshape(d,(3,grid.Nx,grid.Ny)),(2:3)))::AbstractArray{ComplexF64,1},*(3,grid.Nx,grid.Ny),ishermitian=false,ismutating=false))

# ## set discretization parameters and generate explicit dense matrices
# Δx          =   6.                    # μm
# Δy          =   4.                    # μm
# Δz          =   1.
# Nx          =   16
# Ny          =   16
# Nz          =   1
# kz          =   p0[1] #1.45
# # ω           =   1 / λ
# p = p0 #[kz,w,t_core,θ,n_core,n_subs,edge_gap] #,Δx,Δy,Δz,Nx,Ny,Nz]
# eid = ei_dot_rwg(p;Δx,Δy,Δz,Nx,Ny,Nz)
# g = MaxwellGrid(Δx,Δy,Δz,Nx,Ny,Nz)
# ds = MaxwellData(p[1],g)
# ei = make_εₛ⁻¹(ridge_wg(p[2],p[3],p[4],p[7],p[5],p[6],Δx,Δy),g)
# # eii = similar(ei); [ (eii[a,b,i,j,k] = inv(ei[:,:,i,j,k])[a,b]) for a=1:3,b=1:3,i=1:Nx,j=1:Ny,k=1:Nz ] # eii = epsilon tensor field (eii for epsilon_inverse_inverse, yea it's dumb)
# Mop = M̂!(ei,ds)
# Mop2 = M̂(ei,ds)
# Mₖop = M̂ₖ(ei,ds.mn,ds.kpg_mag,ds.𝓕,ds.𝓕⁻¹)
# M = Matrix(Mop)
# dMdk = Matrix(Mₖop)
# mag,mn = calc_kpg(p[1],OptiMode.make_MG(Δx, Δy, Δz, Nx, Ny, Nz).g⃗)
# eid = ei_dot_rwg(p0)
#
# make_M(eid,mag,mn) ≈ M
# make_Mₖ(eid,mag,mn) ≈ -dMdk
# make_Mₖ_eidot(p,eid) ≈ -dMdk
# make_Mₖ(p0) ≈ -dMdk
#
# 𝓕 = plan_fft(randn(ComplexF64, (3,Nx,Ny,Nz)),(2:4))
# 𝓕⁻¹ = plan_ifft(randn(ComplexF64, (3,Nx,Ny,Nz)),(2:4))
#
# ##
# # M̂(ε⁻¹,mn,kpg_mag,𝓕,𝓕⁻¹) = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> M(H,ε⁻¹,mn,kpg_mag,𝓕,𝓕⁻¹)::AbstractArray{ComplexF64,1},*(2,size(ε⁻¹)[end-2:end]...),ishermitian=true,ismutating=false)
# # function M(H,ε⁻¹,mn,kpg_mag,𝓕::FFTW.cFFTWPlan,𝓕⁻¹)
# #     kx_c2t( 𝓕⁻¹ * ε⁻¹_dot( 𝓕 * kx_t2c(H,mn,kpg_mag), ε⁻¹), mn,kpg_mag)
# # end
# kxt2c_op = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> vec( kx_t2c( reshape(H,(2,ds.Nx,ds.Ny,ds.Nz)), ds.mn, ds.kpg_mag ) )::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),*(2,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# kxt2c = Matrix(kxt2c_op)
# F_op = LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(ds.𝓕*reshape(d,(3,ds.Nx,ds.Ny,ds.Nz)))::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# # F_op = LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(fft(reshape(d,(3,ds.Nx,ds.Ny,ds.Nz)),(2:4)))::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# F = Matrix(F_op)
# einv_op = LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec( ε⁻¹_dot( reshape(d,(3,ds.Nx,ds.Ny,ds.Nz)), ei ) )::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# einv = Matrix(einv_op)
# Finv_op = LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(ds.𝓕⁻¹*reshape(d,(3,ds.Nx,ds.Ny,ds.Nz)))::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# # Finv_op = LinearMap{ComplexF64}(d::AbstractArray{ComplexF64,1} -> vec(bfft(reshape(d,(3,ds.Nx,ds.Ny,ds.Nz)),(2:4)))::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# Finv = Matrix(Finv_op)
# kxc2t_op = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> vec( kx_c2t( reshape(H,(3,ds.Nx,ds.Ny,ds.Nz)), ds.mn, ds.kpg_mag ) )::AbstractArray{ComplexF64,1},*(2,ds.Nx,ds.Ny,ds.Nz),*(3,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# kxc2t = Matrix(kxc2t_op)
# zxt2c_op = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> vec( zx_t2c( reshape(H,(2,ds.Nx,ds.Ny,ds.Nz)), ds.mn ) )::AbstractArray{ComplexF64,1},*(3,ds.Nx,ds.Ny,ds.Nz),*(2,ds.Nx,ds.Ny,ds.Nz),ishermitian=false,ismutating=false)
# zxt2c = Matrix(zxt2c_op)
#
# @assert -kxc2t * Finv * einv * F * kxt2c ≈ M
# @assert kxc2t * Finv * einv * F * zxt2c ≈ dMdk # wrong sign?
# @assert make_M(p;Δx,Δy,Δz,Nx,Ny,Nz) ≈ M
# @assert make_M_eidot(p,eid;Δx,Δy,Δz,Nx,Ny,Nz) ≈ M
# @assert ei_dot_rwg(p;Δx,Δy,Δz,Nx,Ny,Nz) ≈ einv
# # if Finv is ifft
# @assert F' ≈  Finv * ( size(F)[1]/3 )
# @assert Finv' * ( size(F)[1]/3 ) ≈  F
# # # if Finv is bfft
# # @assert F' ≈ Finv
# # @assert Finv' ≈  F
# @assert kxc2t' ≈ -kxt2c
# @assert kxt2c' ≈ -kxc2t
#
# # ix = 8
# # iy = 4
# # q = Nx * (iy - 1) + ix
# # 3q-2:3q+3 # 3q-2:3q-2+6-1
# # 2q-1:2q+2 # 2q-1:2q-1+4-1
# #
# # real(kxt2c[3q-2:3q+3,2q-1:2q+2])
# @assert kxt2c_matrix(p0) ≈ kxt2c
# @assert kxt2c_matrix(mag,mn) ≈ kxt2c
# @assert zxt2c_matrix(mn) ≈ zxt2c


text:
function continuous!(f, r, diam)
    dr = r[2] - r[1]
    m = 0

    for (i, j) in enumerate(r)
        if j < diam
            m = i
        end
    end

    fpd = -274.0 * f[m + 1] + 600.0 * f[m + 2] - 600.0 * f[m + 3]
    fpd += 400.0 * f[m + 4] - 150.0 * f[m + 5] + 24.0 * f[m + 6]
    fpd /= 120.0 * dr

    fppd = 225.0 * f[m + 1] - 770.0 * f[m + 2] + 1070.0 * f[m + 3]
    fppd += -780.0 * f[m + 4] + 305.0 * f[m + 5] - 50.0 * f[m + 6]
    fppd /= 60.0 * dr^2

    fd = fpd * (diam - r[m + 1]) + f[m + 1]

    fpi = 274.0 * f[m] - 600.0 * f[m - 1] + 600.0 * f[m - 2]
    fpi += -400.0 * f[m - 3] + 150.0 * f[m - 4] - 24.0 * f[m - 5]
    fpi /= 120.0 * dr

    fppi = 225.0 * f[m] - 770.0 * f[m - 1] + 1070.0 * f[m - 2]
    fppi += -780.0 * f[m - 3] + 305.0 * f[m - 4] - 50.0 * f[m - 5]
    fppi /= 60.0 * dr^2

    fi = fpi * (diam - r[m]) + f[m]
    delf = fd - fi
    delfp = fpd - fpi
    delfpp = fppd - fppi
    a0 = delf - diam * delfp + diam^2 * delfpp / 2.0
    a1 = delfp - diam * delfpp
    a2 = delfpp / 2.0

    for i in 1:m
        f[i] += a0 + a1 * r[i] + a2 * r[i]^2
    end

    return a0, a1, a2
end

function remove_continuous!(f, q, diam, params)
    (a0, a1, a2) = params

    for i in eachindex(f)
        if q[i] > 0.0
            part1 = a0 + 2.0 * diam * a1 + 3.0 * diam^2 * a2
            part1 -= 6.0 * a2 / q[i]^2
            part1 *= sin(q[i] * diam) / q[i]^3

            part2 = -diam * a0 + 2.0 * a1 / q[i]^2 - diam^2 * a1
            part2 += 6.0 * diam * a2 / q[i]^2
            part2 = (part2 - diam^3 * a2) * cos(q[i] * diam) / q[i]^2

            part3 = -2.0 * a1 / q[i]^4

            f[i] -= 4.0 * π * (part1 + part2 + part3)
        else
            f[i] = 0.0
        end
    end

    return nothing
end

@inline function extrapolation(g, f, rho, drho)
    a = (f .- g) ./ drho
    b = @. f - a * rho
    result = @. a * (rho + drho) + b

    return result
end

function inner_product(f, df)
    (integral, _) = romberg(df, f)

    return integral
end

function pairwise_inner_product(x::AbstractMatrix, dx::AbstractRange)
    n = size(x, 1)
    result_matrix = zeros(eltype(x), (n, n))

    # First, deal with the diagonal
    for i in 1:n
        temp_prod = x[i, :] .* x[i, :]
        result_matrix[i, i] = inner_product(temp_prod, dx)
    end

    for i in 1:n
        for j in (i + 1):n
            temp_prod = x[i, :] .* x[j, :]
            result_matrix[i, j] = inner_product(temp_prod, dx)
            result_matrix[j, i] = result_matrix[i, j] # Fill in the symmetric part
        end
    end

    return result_matrix
end


text:
include("settings.jl")
include("SGSolver.jl")
include("DLRSolver.jl")
include("plotting.jl")

using NPZ;

close("all");

s = Settings(1000,30,25);
#s = Settings(600,10,10);
N = s.N;
s.NCons = 0;
s.iCons = 0;

# trun on optimized used of stabilizing terms
optimized = false;

r = [2; 3; 4; 5; 6; 7; 8; 9; 10; 11; 12; 13; 14; 15; 16; 17; 18; 19; 20; 21; 22; 23; 24; 25]
errorExpBack = zeros(length(r));
errorVarBack = zeros(length(r));
errorExpFor = zeros(length(r));
errorVarFor = zeros(length(r));
errorExpSG = zeros(length(r));
errorVarSG = zeros(length(r));

s.useStabilizingTermsS = true;
s.useStabilizingTermsL = true;

s.NCons = 2; # number of conserved basis functions
s.iCons = [1 1; 1 2]; # indices of conserved basis funtions

solver = DLRSolver(s);
PhiQuad = Array(solver.basis.PhiQuad');
PhiQuadCons = Array(solver.basis.PhiQuadCons');
PhiQuadW = Array(solver.basis.PhiQuadW');
fXi = 0.25;
timingsPS = zeros(length(r))
timingsUI = zeros(length(r))
timingsSG = zeros(length(r))

for k = 1:length(r)
    s.r = r[k];
    if optimized
        s.useStabilizingTermsS = false;
        s.useStabilizingTermsL = true;
    end

    solver = DLRSolver(s);
    X,S,W,uCons = FactorIC(solver);

    timingsPS[k] = @elapsed tEnd, X,S,W,uCons = SolveSplitProjectorSplittingIntegrator(solver,X,S,W,uCons);

    # compute moments
    uQ = X*S*W'*PhiQuad + uCons*PhiQuadCons;
    uDLRCons = uQ*Array(solver.basis.PhiQuadWFull')*fXi;
    uDLRCons = Array(uDLRCons')

    plotSolution = Plotting(s,solver.basis,solver.q,tEnd);
    errorExpBack[k], errorVarBack[k] = L2ErrorExpVar(plotSolution,uDLRCons)
    println("rank ",r[k]," : runtime ",timingsPS[k]," ",errorExpBack[k]," ",errorVarBack[k])
end
println("-> DLR projector-splitting integrator DONE.")

s.useStabilizingTermsS = true;
s.useStabilizingTermsL = true;

for k = 1:length(r)
    s.r = r[k];
    if optimized
        s.useStabilizingTermsS = true;
        s.useStabilizingTermsL = false;
    end

    solver = DLRSolver(s);
    X,S,W,uCons = FactorIC(solver);

    timingsUI[k] = @elapsed tEnd, X,S,W,uCons = SolveSplitUnconventionalIntegrator(solver,X,S,W,uCons);
    # compute moments
    uQ = X*S*W'*PhiQuad + uCons*PhiQuadCons;
    uDLRCons = uQ*Array(solver.basis.PhiQuadWFull')*fXi;
    uDLRCons = Array(uDLRCons')

    plotSolution = Plotting(s,solver.basis,solver.q,tEnd);
    errorExpFor[k], errorVarFor[k] = L2ErrorExpVar(plotSolution,uDLRCons)
    println("rank ",r[k]," : runtime ",timingsUI[k]," ",errorExpFor[k]," ",errorVarFor[k])
end
println("-> DLR unconventional integrator DONE.")

s.NCons = 0;
s.iCons = 0;

for k = 1:length(r)
    s.N = r[k];
    if s.N > 5
        #break;
    end
    solver = Solver(s);
    timingsSG[k] = @elapsed tEnd, u = Solve(solver);
    plotSolution = Plotting(s,solver.basis,solver.q,tEnd);
    errorExpSG[k], errorVarSG[k] = L2ErrorExpVar(plotSolution,u)
    println("moments ",r[k]^2," : runtime ",timingsSG[k]," ",errorExpSG[k]," ",errorVarSG[k])
end
println("-> SG DONE.")

print("Computing SG fine...")
s.N = N;
solver = Solver(s);
timeSGFull = @elapsed tEnd, u = Solve(solver);
plotSolution = Plotting(s,solver.basis,solver.q,tEnd);
errorExpSGFull, errorVarSGFull = L2ErrorExpVar(plotSolution,u)
println(" DONE.")

###### plot results ######

fig = figure("Figure6a",figsize=(10, 8), dpi=100)
ax = gca()
ax.plot(r.^2,errorExpSG, "k--o", linewidth=2, label="SG", alpha=1.0)
ax.plot(r,errorExpBack, "g-.<", linewidth=2, label="DLRA, projector-splitting", alpha=1.0)
ax.plot(r,errorExpFor, "m:>", linewidth=2, label="DLRA, unconventional", alpha=1.0)
ax.plot(r,errorExpSGFull*ones(size(r)), "r-", linewidth=2, alpha=0.5)
ylabel("rel. error expectation", fontsize=20)
xlabel("rank/moments", fontsize=20)
ax.set_xlim([r[1],r[end]])
ax.legend(loc="upper right", fontsize=20)
ax.tick_params("both",labelsize=20)
fig.canvas.draw() # Update the figure

fig = figure("Figure6b",figsize=(10, 8), dpi=100)
ax = gca()
ax.plot(r.^2,errorVarSG, "k--o", linewidth=2, label="SG", alpha=1.0)
ax.plot(r,errorVarBack, "g-.<", linewidth=2, label="DLRA, projector-splitting", alpha=1.0)
ax.plot(r,errorVarFor, "m:>", linewidth=2, label="DLRA, unconventional", alpha=1.0)
ax.plot(r,errorVarSGFull*ones(size(r)), "b-", linewidth=2, alpha=0.5)
ylabel("rel. error variance", fontsize=20)
xlabel("rank/moments", fontsize=20)
ax.set_xlim([r[1],r[end]])
ax.legend(loc="upper right", fontsize=20)
ax.tick_params("both",labelsize=20)
fig.canvas.draw() # Update the figure

#### log plot

fig = figure("Figure6aLog",figsize=(10, 8), dpi=100)
ax = gca()
ax.plot(r.^2,errorExpSG, "k--o", linewidth=2, label="SG", alpha=1.0)
ax.plot(r,errorExpBack, "g-.<", linewidth=2, label="DLRA, projector-splitting", alpha=1.0)
ax.plot(r,errorExpFor, "m:>", linewidth=2, label="DLRA, unconventional", alpha=1.0)
ax.plot(r,errorExpSGFull*ones(size(r)), "b-", linewidth=2, alpha=0.5)
ylabel("rel. error expectation", fontsize=20)
xlabel("rank/moments", fontsize=20)
yscale("log")
ax.set_xlim([r[1],r[end]])
ax.legend(loc="upper right", fontsize=20)
ax.tick_params("both",labelsize=20)
fig.canvas.draw() # Update the figure

fig = figure("Figure6bLog",figsize=(10, 8), dpi=100)
ax = gca()
ax.plot(r.^2,errorVarSG, "k--o", linewidth=2, label="SG", alpha=1.0)
ax.plot(r,errorVarBack, "g-.<", linewidth=2, label="DLRA, projector-splitting", alpha=1.0)
ax.plot(r,errorVarFor, "m:>", linewidth=2, label="DLRA, unconventional", alpha=1.0)
ax.plot(r,errorVarSGFull*ones(size(r)), "r-", linewidth=2, alpha=0.5)
ylabel("rel. error variance", fontsize=20)
xlabel("rank/moments", fontsize=20)
yscale("log")
ax.set_xlim([r[1],r[end]])
ax.legend(loc="upper right", fontsize=20)
ax.tick_params("both",labelsize=20)
fig.canvas.draw() # Update the figure

endSG = 17;

# plot timings
fig = figure("Figure7a",figsize=(10, 8), dpi=100)
ax = gca()
ax.plot(timingsSG[1:endSG],errorExpSG[1:endSG], "k:o", linewidth=2, label="SG", alpha=1.0, ms = 10)
ax.plot(timingsPS,errorExpBack, "g:<", linewidth=2, label="DLRA, projector-splitting", alpha=1.0, ms = 10)
ax.plot(timingsUI,errorExpFor, "m:>", linewidth=2, label="DLRA, unconventional", alpha=1.0, ms = 10)
#ax.plot(r,errorExpSGFull*ones(size(r)), "r-", linewidth=2, alpha=0.5)
ylabel("rel. error expectation", fontsize=20)
xlabel("runtime [sec]", fontsize=20)
yscale("log")
xscale("log")
grid("on", which="both")
#ax.set_xlim([r[1],r[end]])
ax.legend(loc="upper right", fontsize=20)
ax.tick_params("both",labelsize=20)
fig.canvas.draw() # Update the figure

# plot timings
fig = figure("Figure7b",figsize=(10, 8), dpi=100)
ax = gca()
ax.plot(timingsSG[1:endSG],errorVarSG[1:endSG], "k:o", linewidth=2, label="SG", alpha=1.0, ms = 10)
ax.plot(timingsPS,errorVarBack, "g:<", linewidth=2, label="DLRA, projector-splitting", alpha=1.0, ms = 10)
ax.plot(timingsUI,errorVarFor, "m:>", linewidth=2, label="DLRA, unconventional", alpha=1.0, ms = 10)
#ax.plot(r,errorExpSGFull*ones(size(r)), "r-", linewidth=2, alpha=0.5)
ylabel("rel. error variance", fontsize=20)
xlabel("runtime [sec]", fontsize=20)
yscale("log")
xscale("log")
grid("on", which="both")
#ax.set_xlim([r[1],r[end]])
ax.legend(loc="upper right", fontsize=20)
ax.tick_params("both",labelsize=20)
fig.canvas.draw() # Update the figure

npzwrite("results/ranksN$(N)Nx$(s.Nx).jld", r)

npzwrite("results/timingsSGN$(N)Nx$(s.Nx).jld", timingsSG)
npzwrite("results/timingsPSN$(N)Nx$(s.Nx).jld", timingsPS)
npzwrite("results/timingsUIN$(N)Nx$(s.Nx).jld", timingsUI)

npzwrite("results/errorExpSGN$(N)Nx$(s.Nx).jld", errorExpSG)
npzwrite("results/errorExpPSN$(N)Nx$(s.Nx).jld", errorExpBack)
npzwrite("results/errorExpUIN$(N)Nx$(s.Nx).jld", errorExpFor)

npzwrite("results/errorVarSGN$(N)Nx$(s.Nx).jld", errorVarSG)
npzwrite("results/errorVarPSN$(N)Nx$(s.Nx).jld", errorVarBack)
npzwrite("results/errorVarUIN$(N)Nx$(s.Nx).jld", errorVarFor)


text:
# ---
# title: Legend merging
# cover: assets/legend_merging.png
# description: Multiple scales for the same variable.
# author: "[Pietro Vertechi](https://github.com/piever)"
# ---

using AlgebraOfGraphics, CairoMakie
set_aog_theme!() #src

N = 40

x = [1:N; 1:N]
y = [cumsum(randn(N)); cumsum(randn(N))]
grp = [fill("a", N); fill("b", N)]

df = (; x, y, grp)

layers = visual(Lines) + visual(Scatter) * mapping(marker = :grp)
plt = data(df) * layers * mapping(:x, :y, color = :grp)

fg = draw(plt)

# save cover image #src
mkpath("assets") #src
save("assets/legend_merging.png", fg) #src


text:
using DifferentialEquations
using LinearAlgebra
using Plots
using ForwardDiff
using ProgressMeter

# 2D Ball Example
function physics(u, p, t)
    # x, dx = u[1, 3]
    # y, dy = u[2, 4]

    # g = gravity, b = drag
    g, b = p[1], p[2]

    dx = u[2]
    ddx = b * u[1]^2
    dy = u[4]
    ddy = -g - b * u[3]^2

    [dx, ddx, dy, ddy]
end

# Callback
condition_cont(u,t,integrator) = u[3]
# condition(u,t,integrator) = u[3]<=0
affect!(integrator) = terminate!(integrator)
# cb = DiscreteCallback(condition,affect!)
cb = ContinuousCallback(condition_cont,nothing,affect!)

#Lets wrap this into a "predict" method
function predict(vel)
    # Simply roll this forward in time and we get our solution
    tspan = (0.0, 10.0)
    tspan_dual = convert.(eltype(vel),tspan)
    p = [9.8, 0.0]

    start_state = [0.0, vel[1], 10.0, vel[2]]
    prob = ODEProblem(physics, start_state, tspan_dual, p)
    sol = solve(prob, callback=cb)

    sol
end

## Now let's try to optimize the system to hit a particular point
# want point (1.0, 0)
function loss(vel)
    simulation = predict(vel)
    final_pos = simulation[[1, 3], end]
    l = (final_pos[1] - 1.0)^2 + final_pos[2]^2 # x -> 1.0, y -> 0
    return l
end

# Plots loss surface over velocities in x and y
function plot_loss_surf()
    v_x = 0.0:0.01:3
    v_y = 0.0:0.01:2

    loss_mat = zeros(length(v_x), length(v_y))
    for x in 1:length(v_x)
        for y in 1:length(v_y)
            loss_mat[x, y] = loss([v_x[x], v_y[y]])
        end
    end

    # sol_f(x, y) = loss_mat[x, y]
    # plot(1:length(v_x), 1:length(v_y), sol_f, st=:surf, color=:viridis)
    plot(loss_mat, st=:heatmap, color=:viridis)

end

# Plots loss over just x velocities
function plot_loss_x()

    v_x = 0.05:0.01:3

    loss_mat = zeros(length(v_x))
    for x in 1:length(v_x)
        loss_mat[x] = loss([v_x[x], 0.0])
    end

    plot(v_x, loss_mat)

end

# Optimize the xy velocities needed
function optimize(vel_opt, η)
    @showprogress for idx in 1:1000
        grads = ForwardDiff.gradient(s -> loss(s), vel_opt) # Magic
        vel_opt .-= η*grads

        if loss(vel_opt) < 0.00001
            break
        end
    end

    vel_opt
end

function optimize_and_save(vel_opt, η)
    solns = []
    for idx in 1:1000
        if idx % 10 == 1
            append!(solns, [deepcopy(vel_opt)])
            # @show idx
            # @show vel_opt
            # @show solns
        end

        grads = ForwardDiff.gradient(s -> loss(s), vel_opt) # Magic
        vel_opt .-= η*grads

        if loss(vel_opt) < 0.00001
            break
        end

    end
    display(solns)
    solns
end

# Example run of optimize
function run_test()
    p = [1.0, 0.2]
    init_sol = predict(p);
    v = optimize(p, 0.2);
    final_sol = predict(v);

    display(plot(init_sol, vars=(1, 3), label="Initial")); display(plot!(final_sol, vars=(1,3), label="Optimized to land at x=1.0, y=0.0"))
end

function run_and_save()
    p = [2.0, 0.4]
    vs = optimize_and_save(p, 0.03);
    display(plot(title="Ballistic Optimizer: Goal final position (1.0, 0.0)",
                 xlabel="x distance [m]",
                 ylabel="y distance [m]"))
    i = 0
    for v in vs
        display(
            plot!(predict(v),
                vars=(1,3),
                label="Iteration " * string(i),
                color=:Blues,
                colorbar=:none,
                legend=:bottomleft,
                line_z=i+1)
            )
        i += 10
    end

    savefig("plots/ballistic_example.png")
end


text:
import StringUtils

function char_to_digit(c::Char)
    if isdigit(c)
        return Int(c) - Int('0')
    else
        throw(ArgumentError(@sprintf("%c is not a digit", c)))
    end
end

function main()
    if length(ARGS) > 1
        filename = ARGS[1]
        n = parse(Int64, ARGS[2])
        a = readlines(filename)
        s = StringUtils.conc(a)
        max_product = -1
        for i in 1:(length(s)-n+1)
            product = reduce(*,
                             map(char_to_digit,
                                 collect(s[i:(i+n-1)])))
            if product > max_product
                max_product = product
            end
        end
        @printf("The maximum %d length product has value %d\n", n, max_product)
    else
        error("this program needs a single numeric argument")
    end
end

main()


text:
using LinearAlgebra

@testset "U1Tensor" begin
    u1chs, u1ds, io = U1Charges(-2:2), fill(2,5), InOut(-1,1)
    @testset "Initialization" begin
        a = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io)

        @test symmetry(a) == U1()
        @test chargetype(a) == chargetype(U1())
        @test chargestype(a) == chargestype(U1())
        @test eltype(a) == Float64
        @test isempty(tensor(a))

        #Initialize with zeros
        b0 = initwithzero!(deepcopy(a))
        @test isinvariant(b0)
        @tensor bs = b0[1,1]

        @test charge(b0) == zero(U1Charge)
        @test bs == zero(Float64)

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        bd = todense(br)
        @tensor begin
            sc = br[1,1]
            scd = bd[1,1]
        end

        @test sc isa Float64
        @test scd isa Float64
        @test sc ≈ scd
    end

    @testset "TensorOperations" begin
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test all(iszero(v) for v in values(fr0))

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test all(iszero(v) for v in values(fr2))

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
        @tensor ref = nfr[1,2,3] * conj(nfr)[1,2,3]
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DASTensor{Complex{Float64},3}(U1(), ntuple(x -> u1chs, 3),
                 ntuple(x -> u1ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with charge
        ar = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io, U1Charge(1))
        initwithrand!(ar)
        @test charge(ar) == U1Charge(1)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ## with truncation
        ar = DASTensor{Float64,2}(U1(), (U1Charges(0:0), U1Charges(0:0)), ([200], [200]), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DASTensor{Float64,2}(U1(), (u1chs, u1chs), (u1ds, u1ds), io)
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        fr = DASTensor{Float64,3}(U1(), (u1chs, u1chs, u1chs), (u1ds, u1ds, u1ds), InOut(1,1,-1))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar = DASTensor{Float64,2}(
            U1(),
            (U1Charges(-1:1), U1Charges(-1:1)),
            ([10,10,10], [10,10,10]),
            InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test all(v ≈ I for v in values(id))

        ar = DASTensor{Float64,2}(
            U1(),
            (U1Charges(-1:1), U1Charges(-1:1)),
            (fill(10,3),fill(4,3)),
            InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test all(v ≈ I for v in values(id))
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DASTensor{Complex{Float64},6}(U1(), ntuple(x -> u1chs, 6),
                 ntuple(x -> u1ds, 6), InOut(-1,-1,-1,1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DASTensor{Complex{Float64},4}(U1(), ntuple(x -> u1chs, 4),
                 ntuple(x -> u1ds, 4), InOut(-1,-1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end

@testset "Z2" begin
    z2chs, z2ds, io = Z2Charges(), fill(2,2), InOut(-1,1)
    @testset "Initialization" begin
        a = DASTensor{Float64,2}(Z2(), (z2chs, z2chs), (z2ds, z2ds), io)
        @test symmetry(a) == Z2()
        @test chargetype(a) == chargetype(Z2())
        @test chargestype(a) == chargestype(Z2())
        @test eltype(a) == Float64
        @test isempty(tensor(a))

        #Initialize with zeros
        b0 = initwithzero!(deepcopy(a))
        @tensor bs = b0[1,1]
        @test isinvariant(b0)

        @test charge(b0) == zero(Z2Charge)
        @test bs == zero(Float64)

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        bd = todense(br)
        @tensor begin
            sc = br[1,1]
            scd = bd[1,1]
        end

        @test sc isa Float64
        @test scd isa Float64
        @test sc ≈ scd
    end

    @testset "TensorOperations" begin
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test all(iszero(v) for v in values(fr0))

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test all(iszero(v) for v in values(fr2))

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DASTensor{Float64,2}(Z2(), (z2chs, z2chs), (z2ds, z2ds), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DASTensor{Complex{Float64},3}(Z2(), ntuple(x -> z2chs, 3),
                 ntuple(x -> z2ds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with charge
        ar = DASTensor{Float64,2}(Z2(), (z2chs, z2chs), (z2ds, z2ds), io, Z2Charge(1))
        initwithrand!(ar)
        @test charge(ar) == Z2Charge(1)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ## with truncation
        ar = DASTensor{Float64,2}(Z2(), (Z2Charges(), Z2Charges()), ([200,200], [200,200]), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DASTensor{ComplexF64,2}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),2),
            ntuple(i -> [10,10], 2),
            InOut(ntuple(i -> -1, 2)...))
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        fr = DASTensor{ComplexF64,3}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),3),
            ntuple(i -> [10,10], 3),
            InOut(1,1,-1))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar = DASTensor{ComplexF64,2}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),2),
            ([4,4], [4,4]),
            InOut(ntuple(i -> -1, 2)...))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test all(v ≈ I for v in values(id))

        ar = DASTensor{ComplexF64,2}(ZN{2}(),
            ntuple(i -> ZNCharges{2}(),2),
            ([10,10], [4,4]),
            InOut(ntuple(i -> -1, 2)...))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test all(v ≈ I for v in values(id))
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DASTensor{Complex{Float64},6}(Z2(), ntuple(x -> z2chs, 6),
                 ntuple(x -> z2ds, 6), InOut(-1,-1,-1,1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DASTensor{Complex{Float64},4}(Z2(), ntuple(x -> z2chs, 4),
                 ntuple(x -> z2ds, 4), InOut(-1,-1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end

@testset "NDAS" begin
    ndaschs, ndasds, io = NDASCharges(Z2Charges(), U1Charges(-2:2)), fill(2,10), InOut(-1,1)
    @testset "Initialization" begin
        a = DASTensor{Float64,2}(NDAS(Z2(),U1()), (ndaschs, ndaschs), (ndasds, ndasds), io)
        @test symmetry(a) == NDAS(Z2(),U1())
        @test chargetype(a) == chargetype(a)
        @test chargestype(a) == chargestype(a)
        @test eltype(a) == Float64
        @test isempty(tensor(a))

        #Initialize with zeros
        b0 = initwithzero!(deepcopy(a))
        @test isinvariant(b0)
        @tensor bs = b0[1,1]

        @test charge(b0) == zero(chargetype(a))
        @test bs == zero(Float64)

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        bd = todense(br)
        @tensor begin
            sc = br[1,1]
            scd = bd[1,1]
        end

        @test sc isa Float64
        @test scd isa Float64
        @test sc ≈ scd
    end

    @testset "TensorOperations" begin
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test all(iszero(v) for v in values(fr0))

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test all(iszero(v) for v in values(fr2))

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()), (ndaschs, ndaschs), (ndasds, ndasds), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with charge
        ar = DASTensor{Float64,2}(
            NDAS(Z2(),U1()),
            (ndaschs, ndaschs),
            (ndasds, ndasds),
            io,
            NDASCharge(Z2Charge(1), U1Charge(1)))
        initwithrand!(ar)
        @test charge(ar) == NDASCharge(Z2Charge(1), U1Charge(1))
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ## with truncation
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [200,200]), io)
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(sizes(sa,1) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [200,200]), io)
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        f = DASTensor{Complex{Float64},3}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 3),
                 ntuple(x -> ndasds, 3), InOut(-1,-1,1))
        fr = initwithrand!(deepcopy(f))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [200,200]), InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test all(v ≈ I for v in values(id))

        ar = DASTensor{Float64,2}(NDAS(Z2(),U1()),
            (NDASCharges(Z2Charges(), U1Charges(0:0)),
             NDASCharges(Z2Charges(), U1Charges(0:0))),
             ([200,200], [20,20]), InOut(1,-1))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test all(v ≈ I for v in values(id))
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DASTensor{Complex{Float64},6}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 6),
                 ntuple(x -> ndasds, 6), InOut(-1,-1,-1,1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DASTensor{Complex{Float64},4}(NDAS(Z2(),U1()), ntuple(x -> ndaschs, 4),
                 ntuple(x -> ndasds, 4), InOut(-1,-1,1,1))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end

@testset "DTensor" begin
    @testset "Initialization" begin
        #Initialize with zeros
        a = DTensor{Complex{Float64}}((8,8))
        b0 = initwithzero!(deepcopy(a))
        @tensor bs = b0[1,1]

        @test bs == zero(Complex{Float64})

        #Initialize with rand
        br = initwithrand!(deepcopy(a))
        @tensor sc = br[1,1]
        @test sc isa Complex{Float64}
    end

    @testset "TensorOperations" begin
        f = DTensor{Complex{Float64},3}((4,4,4))
        fr = initwithrand!(deepcopy(f))
        #permutation
        perms = (p for p in Iterators.product(ntuple(i -> 1:3, 3)...)
            if sort(collect(p)) == 1:3)
        @test all(toarray(tensorcopy(fr, p)) ≈ tensorcopy(toarray(fr), p)
            for p in perms)
        ##inner product direct
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        ##inner product direct with permutation
        @test ref ≈ @tensor fr'[2,1,3] * fr[2,1,3]
        ##outer product then trace
        @test ref ≈ @tensor begin
            fr2[1,2,3,4,5,6] := fr[1,2,3] * fr'[4,5,6]
            fr2[1,2,3,1,2,3]
        end
        ##partial contraction then partial trace
        @test ref ≈ @tensor begin
            fr3[1,2,3,4] := fr[1,2,-1] * fr'[3,4,-1]
            fr3[1,2,1,2]
        end
    end

    @testset "KrylovKit" begin
        #compare operations from krylovkit
        f = DTensor{Complex{Float64},3}((8,8,8))
        fr = initwithrand!(deepcopy(f))

        ##fill!
        fr0 = fill!(deepcopy(fr), 0)
        @test iszero(fr0.array)

        ##mul! -> w = a .* v
        fr2 = deepcopy(fr)
        mul!(fr2, fr, 2)
        @test fr2 ≈ fr * 2
        mul!(fr2, fr, 1/2)
        @test fr2 ≈ fr/2
        mul!(fr2, fr, 1im)
        @test fr2 ≈ fr * 1im

        ##rmul! w = a .* w
        fr2 = deepcopy(fr)
        rmul!(fr2, 2)
        @test fr2 ≈ fr * 2
        rmul!(fr2, 0)
        @test iszero(fr2.array)

        ##axpy! a,v,w -> w = a * v + w
        fr2 = deepcopy(fr)
        axpy!(1,fr, fr2)
        @test fr2 ≈ 2 * fr
        axpy!(0,fr, fr2)
        @test fr2 ≈ 2 * fr

        ##axpby! a,v,b,w -> w = a * v + b * w
        fr2 = deepcopy(fr)
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ 2 * fr
        initwithrand!(fr2)
        axpby!(1,fr,0,fr2)
        @test fr2 ≈ fr
        initwithrand!(fr2)
        @tensor frsum[1,2,3] := fr[1,2,3] + fr2[1,2,3]
        axpby!(1,fr,1,fr2)
        @test fr2 ≈ frsum

        ##dot, norm
        f = DTensor{Complex{Float64},3}((8,8,8))
        fr = initwithrand!(deepcopy(f))
        @tensor ref = fr[1,2,3] * fr'[1,2,3]
        @test ref ≈ dot(fr,fr)
        @test ref ≈ norm(fr)^2
        nfr = fr / norm(fr)
        @test norm(nfr) ≈ 1
    end

    @testset "TensorFactorization" begin
        ##w/o splitting and fusing
        ar = DTensor{Float64,2}((20,20))
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar)
        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        @test ar2 ≈ ar
        @test sum(diag(sa)) ≈ @tensor sa[1,1]
        @tensor sa2[1,2] := ua'[-1,1] * ar[-1,-2] * vda'[2,-2]
        @test sa2 ≈ sa

        ##with splitting and fusing
        f = DTensor{Complex{Float64},3}((5,5,5))
        fr = initwithrand!(deepcopy(f))
        uf, sf, vdf = tensorsvd(fr,((1,2),3))
        @tensor fr2[1,2,3] := uf[1,2,-1] * sf[-1,-2] * vdf[-2,3]
        @test fr2 ≈ fr

        ## with truncation
        ar = DTensor{Float64,2}((200,200))
        initwithrand!(ar)
        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxerror(0.001))
        @test all(size(sa) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ua, sa, vda = tensorsvd(ar, svdtrunc = svdtrunc_maxcumerror(0.001))
        @test all(size(sa) .<= 200)

        @tensor ar2[1,2] := ua[1,-1] * sa[-1,-2] * vda[-2,2]
        ar2 /= norm(ar2)
        ar /= norm(ar)
        @test dot(ar2,ar) > 1-0.001

        ##tensorqr
        ar = DTensor{Float64,2}((200,200))
        initwithrand!(ar)
        q,r = tensorqr(ar,((1,),(2,)))
        @tensor id[1,2] := q[1,-1] * q'[2,-1]
        @test toarray(id) ≈ I
        @tensor id[1,2] := q'[1,-1] * q[2,-1]
        @test toarray(id) ≈ I
        @test ar ≈ @tensor ar2[1,2] := q[1,-1] * r[-1,2]

        fr = DTensor{Complex{Float64},3}((5,5,5))
        fr = initwithrand!(fr)
        q,r = tensorqr(fr, ((1,2),(3,)))
        @tensor id[1,2] := q'[-1,-2,1] * q[-1,-2,2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := q[1,2,-1] * r[-1,3]

        r,q = tensorrq(fr, ((1,),(2,3)))
        @tensor id[1,2] := q'[1,-1,-2] * q[2,-1,-2]
        @test toarray(id) ≈ I
        @test fr ≈ @tensor fr2[1,2,3] := r[1,-1] * q[-1,2,3]
    end

    @testset "pinv" begin
        ar =  DTensor{ComplexF64,2}((200,200))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ar[1,-1] * ari[-1,2]
        @test toarray(id) ≈ I

        ar =  DTensor{ComplexF64,2}((200,20))
        initwithrand!(ar)
        ari = pinv(ar)
        @tensor id[1,2] := ari[1,-1] * ar[-1,2]
        @test toarray(id) ≈ I
    end

    @testset "Reshaping" begin
        ##without permutation
        f = DTensor{Complex{Float64},6}(ntuple(i -> 2, 6))
        fr = initwithrand!(deepcopy(f))
        frf, rs2 = fuselegs(fr, ((1,2),(3,4),(5,6)))
        frff, rs3 = fuselegs(frf, ((1,2,3),))
        @test abs(norm(frf)  - norm(fr)) < 1e-10
        @test abs(norm(frff) - norm(fr)) < 1e-10

        frffs = initwithzero!(deepcopy(frf))
        frffss = initwithzero!(deepcopy(fr))
        splitlegs!(frffs,  frff,  ((1,1,1),(1,1,2),(1,1,3)), rs3...)
        splitlegs!(frffss, frffs, ((1,1,1),(1,1,2),(2,2,1),(2,2,2),(3,3,1),(3,3,3)), rs2...)
        @test frffss ≈ fr

        ##with permutation
        f = DTensor{Complex{Float64},4}(ntuple(i -> 2, 4))
        fr = initwithrand!(deepcopy(f))
        frf, rs = fuselegs(fr, ((1,3),(2,4)))
        @test abs(norm(frf)  - norm(fr)) < 1e-10

        frfs = initwithzero!(deepcopy(fr))
        splitlegs!(frfs,  frf,  ((1,1,1),(2,2,1),(1,1,2),(2,2,2)), rs...)
        @test frfs ≈ fr
    end
end


text:
function mass(u, q, w, v, x, t, mu, eta)
    m = 1.0;
    return m;
end
function flux(u, q, w, v, x, t, mu, eta)
    f = mu[1]*q;
    return f;
end
function source(u, q, w, v, x, t, mu, eta)
    s = 0.0;
    return s;
end
function ubou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    ub = mu[2];
    return ub;
end
function fbou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    f = flux(u, q, w, v, x, t, mu, eta);
    fb = f[1]*n[1] + f[2]*n[2] + tau[1]*(u[1]-0.0);
    return fb;
end
function initu(x, mu, eta)
    x1 = x[1];
    x2 = x[2];
    u0 = sin(pi*x1)*sin(pi*x2);
    return u0;
end

max_stars_repo_name: glwagner/Exasim

max_issues_repo_name: glwagner/Exasim

max_forks_repo_name: glwagner/Exasim


text:
# This file includes node dynamics

abstract type AbstractNode end

"""
    $TYPEDEF

Lorenz dyanmics

# Fields

    $TYPEDFIELDS
"""
@kwdef mutable struct Lorenz <: AbstractNode
    σ::Float64 = 10.0
    β::Float64 = 8 / 3
    ρ::Float64 = 35.0
    x::Vector{Float64} = rand(3)
    t::Float64 = 0.0
end

function (node::Lorenz)(dx, x, u, t)
    dx[1] = node.σ * (x[2] - x[1])
    dx[2] = x[1] * (node.ρ - x[3]) - x[2]
    dx[3] = x[1] * x[2] - node.β * x[3]
end

dimension(::Type{Lorenz}) = 3

"""
    $TYPEDEF

Rossler dyanmics

# Fields

    $TYPEDFIELDS
"""
@kwdef mutable struct Rossler <: AbstractNode
    a::Float64 = 0.2
    b::Float64 = 0.2
    c::Float64 = 5.7
    x::Vector{Float64} = rand(3)
    t::Float64 = 0.0
end

function (node::Rossler)(dx, x, u, t)
    dx[1] = -x[2] - x[1]
    dx[2] = x[1] + node.a * x[2]
    dx[3] = node.b + x[3] * (x[1] - node.c)
end

dimension(::Type{Rossler}) = 3

# Piecewise Chua diode
@kwdef struct Diode
    a::Float64 = -8 / 7
    b::Float64 = -5 / 7
end

function (diode::Diode)(x)
    if x > 1
        diode.b * x + diode.a - diode.b
    elseif abs(x) ≤ 1
        diode.a * x
    else
        diode.b * x - diode.a + diode.b
    end
end

@kwdef mutable struct Chua <: AbstractNode
    diode::Diode = Diode()
    α::Float64 = 9.0
    β::Float64 = 100 / 7
    x::Vector{Float64} = 1e-3 * rand(3)
    t::Float64 = 0.00
end

function (ds::Chua)(dx::AbstractVector, x::AbstractVector, u, t)
    dx[1] = ds.α * (x[2] - x[1] - ds.diode(x[1]))
    dx[2] = x[1] - x[2] + x[3]
    dx[3] = -ds.β * x[2]
end


text:
module REPL

using SHA, Logger, Genie.Configuration, Genie, Genie.Generator, Tester, Toolbox, App, Util
SEARCHLIGHT_ON && eval(:(using SearchLight, SearchLight.Generator, Migration))

if is_dev()
  @eval using Revise
end

const JULIA_PATH = joinpath(JULIA_HOME, "julia")

"""
    secret_token() :: String

Generates a random secret token to be used for configuring the SECRET_TOKEN const.
"""
function secret_token() :: String
  sha256("$(randn()) $(Dates.now())") |> bytes2hex
end

"""
    function new_app(path = "."; db_support = false, skip_dependencies = false) :: Void

Creates a new Genie app at the indicated path.
"""
function new_app(path = "."; db_support = false, skip_dependencies = false, autostart = true) :: Void
  cp(joinpath(Pkg.dir("Genie"), "files", "new_app"), abspath(path))

  chmod(joinpath(path, "bin", "server"), 0o700)
  chmod(joinpath(path, "bin", "repl"), 0o700)

  open(joinpath(path, "config", "secrets.jl"), "w") do f
    write(f, """const SECRET_TOKEN = "$(secret_token())" """)
  end

  Logger.log("Done! New app created at $(abspath(path))", :info)

  skip_dependencies || add_dependencies(db_support = db_support)

  is_windows() && setup_windows_bin_files(path)

  autostart || return nothing

  Logger.log("Starting your brand new Genie app - hang tight!", :info)
  run_repl_app(path)

  nothing
end

"""
    run_repl_app() :: Void

Runs a new Genie REPL app within the current thread.
"""
function run_repl_app(path = ".") :: Void
  cd(abspath(path))

  run(`$JULIA_PATH -L $(abspath("genie.jl")) --color=yes --depwarn=no -q`)

  nothing
end

"""

"""
function setup_windows_bin_files(path = ".")
  open(joinpath(path, "bin", "repl.bat"), "w") do f
    write(f, "$JULIA_PATH -L ../genie.jl --color=yes --depwarn=no -q -- %*")
  end

  open(joinpath(path, "bin", "server.bat"), "w") do f
    write(f, "$JULIA_PATH -L ../genie.jl --color=yes --depwarn=no -q -- s %*")
  end
end

"""
    add_dependencies(; db_support = false) :: Void

Attempts to install Genie's dependencies - packages that are not part of METADATA.
"""
function add_dependencies(; db_support = false) :: Void
  Logger.log("Looking for dependencies", :info)

  Logger.log("Checking for Flax rendering engine support", :info)
  Util.package_added("Flax") || Pkg.clone("https://github.com/essenciary/Flax.jl")

  if db_support
    Logger.log("Checking for DBI support", :info)
    Util.package_added("DBI") || Pkg.clone("https://github.com/JuliaDB/DBI.jl")

    Logger.log("Checking for PostgreSQL support", :info)
    Util.package_added("PostgreSQL") || Pkg.clone("https://github.com/JuliaDB/PostgreSQL.jl")

    Logger.log("Checking for SearchLight ORM support", :info)
    Util.package_added("SearchLight") || Pkg.clone("https://github.com/essenciary/SearchLight.jl")
  end

  Logger.log("Finished adding dependencies", :info)

  nothing
end

"""
    new_model(model_name::String) :: Void

Creates a new `model` file.
"""
function new_model(model_name::String) :: Void
  SearchLight.Generator.new_model(Dict{String,Any}("model:new" => model_name))
  App.load_resources()

  nothing
end

"""
    new_controller(controller_name::String) :: Void

Creates a new `controller` file.
"""
function new_controller(controller_name::String) :: Void
  Genie.Generator.new_controller(Dict{String,Any}("controller:new" => controller_name))
  App.load_resources()

  nothing
end

"""
    new_channel(channel_name::String) :: Void

Creates a new `channel` file.
"""
function new_channel(channel_name::String) :: Void
  Genie.Generator.new_channel(Dict{String,Any}("channel:new" => channel_name))
  App.load_resources()

  nothing
end

"""
    new_resource(resource_name::String) :: Void

Creates all the files associated with a new resource.
"""
function new_resource(resource_name::String) :: Void
  Genie.Generator.new_resource(Dict{String,Any}("resource:new" => resource_name))
  App.load_resources()

  nothing
end

"""
    new_migration(migration_name::String) :: Void

Creates a new migration file.
"""
function new_migration(migration_name::String) :: Void
  SearchLight.Generator.new_migration(Dict{String,Any}("migration:new" => migration_name))
end

"""
"""
function new_table_migration(migration_name::String) :: Void
  SearchLight.Generator.new_table_migration(Dict{String,Any}("migration:new" => migration_name))
end

"""
    new_task(task_name::String) :: Void

Creates a new `Task` file.
"""
function new_task(task_name::String) :: Void
  endswith(task_name, "Task") || (task_name = task_name * "Task")
  Toolbox.new(Dict{String,Any}("task:new" => task_name), App.config)
end

"""
    write_secrets_file() :: Void

Generates a valid secrets.jl file with a random SECRET_TOKEN.
"""
function write_secrets_file() :: Void
  open(joinpath(Genie.CONFIG_PATH, "secrets.jl"), "w") do f
    write(f, """const SECRET_TOKEN = "$(secret_token())" """)
  end

  Logger.log("Generated secrets.jl file in $(Genie.CONFIG_PATH)", :info)

  nothing
end

"""
    reload_app() :: Void
"""
function reload_app() :: Void
  Logger.log("Attempting to reload the Genie's core modules. If you get unexpected errors or things don't work as expected, simply exit this Julia session and start a new one to fully reload Genie.", :warn)

  is_dev() && Revise.revise(App)

  App.load_configurations()
  App.load_initializers()

  Logger.log("The app was reloaded.", :info)

  nothing
end

end