text:
function __init__()\
    Unitful.register(HydrationOfCement)
end

#using PeriodicTable, UnitfulMoles, Unitful
using CSV, DataFrames

file_location = joinpath("src", "cemdata18_std_thermo_properties.txt")
d = DataFrame(CSV.File(file_location))

function Avrami(u, p, t)
    k = p[1]
    n = p[2]
    if 0 <= u < 1
        return k/n * log(1-u) * (-log(1-u))^(1-n)
    end
end

function Jander(u, p, t)
    if 0 < u
        return (p * (1 - u)^(2/3)) / (1- (1 - u)^(1/3))
    end
end

function Power_law(α, k , n)
    if 0 <= α <= 1
        return k * (1 - α)^n
    end
end

function Parrot_Killoh(α, k, n)
#using DifferentialEquations
#p=...
#u0=...
#prob = ODEProblem(Avrami,u₀,tspan,p)
#sol = solve(prob)
end

struct Molecule
    name::String
    chemicalFormula::AbstractString
end

struct Anhydre{Molecule}
    name::String
    chemicalFormula::AbstractString
end

function read_molecule(fich)
    #TODO
end

"Defines a cool function. Returns some stuff"
function coolFunc(x::Int)
    return x+2
end

"""
Defines an even cooler function. ``LaTeX``.

### Returns
 * Markdown works in here
"""
function coolFunc2(y::Int)
    return y+3
end


text:
"""
    PerspectiveMap()

Construct a perspective transformation. The perspective transformation takes,
e.g., a point in 3D space and "projects" it onto a 2D virtual screen of an ideal
pinhole camera (at distance `1` away from the camera). The camera is oriented
towards the positive-Z axis (or in general, along the final dimension) and the
sign of the `x` and `y` components is preserved for objects in front of the
camera (objects behind the camera are also projected and therefore inverted - it
is up to the user to cull these as necessary).

This transformation is designed to be used in composition with other coordinate
transformations, defining e.g. the position and orientation of the camera. For
example:

    cam_transform = PerspectiveMap() ∘ inv(AffineMap(cam_rotation, cam_position))
    screen_points = map(cam_transform, points)

(see also `cameramap`)
"""
struct PerspectiveMap <: Transformation
end

function (::PerspectiveMap)(v::AbstractVector)
    scale =  1/v[end]
    return [v[i] * scale for i in 1:length(v)-1]
end

@inline function (::PerspectiveMap)(v::StaticVector)
    return pop(v) * inv(v[end])
end

Base.isapprox(::PerspectiveMap, ::PerspectiveMap; kwargs...) = true

"""
    cameramap()
    cameramap(scale)
    cameramap(scale, offset)

Create a transformation that takes points in real space (e.g. 3D) and projects
them through a perspective transformation onto the focal plane of an ideal
(pinhole) camera with the given properties.

The `scale` sets the scale of the screen. For a standard digital camera, this
would be `scale = focal_length / pixel_size`. Non-square pixels are supported
by providing a pair of scales in a tuple, `scale = (scale_x, scale_y)`. Positive
scales represent a camera looking in the +z axis with a virtual screen in front
of the camera (the x,y coordinates are not inverted compared to 3D space). Note
that points behind the camera (with negative z component) will be projected
(and inverted) onto the image coordinates and it is up to the user to cull
such points as necessary.

The `offset = (offset_x, offset_y)` is used to define the origin in the imaging
plane. For instance, you may wish to have the point (0,0) represent the top-left
corner of your imaging sensor. This measurement is in the units after applying
`scale` (e.g. pixels).

(see also `PerspectiveMap`)
"""
cameramap() = PerspectiveMap()
cameramap(scale::Number) =
    LinearMap(UniformScaling(scale)) ∘ PerspectiveMap()
cameramap(scale::Tuple{Number, Number}) =
    LinearMap(@SMatrix([scale[1] 0; 0 scale[2]])) ∘ PerspectiveMap()
cameramap(scale::Number, offset::Tuple{Number,Number}) =
    AffineMap(UniformScaling(scale), SVector(-offset[1], -offset[2])) ∘ PerspectiveMap()
cameramap(scale::Tuple{Number, Number}, offset::Tuple{Number,Number}) =
    AffineMap(@SMatrix([scale[1] 0; 0 scale[2]]), SVector(-offset[1], -offset[2])) ∘ PerspectiveMap()

ext: jl


text:
using Plots; gr()

function draw_neuron(x, y, r; c=:blue)

    θs = 0:0.1:2pi
    xs = x .+ r.*cos.(θs)
    ys = y .+ r.*sin.(θs)

    plot!(xs, ys, seriestype=:shape, c=c, alpha=0.5, aspect_ratio=1, leg=false)

end

#neuron_coords(x, N, spacing) = range(-(N - 1)/2 * spacing, spacing, N)

"""
Vertical position of neuron in layer i, position j, with a total of N neurons
"""
neuron_coords(j, N, y_spacing) = (-(N - 1)/2 + j) * y_spacing

function draw_neurons(x, N, spacing, r; c=:blue)

    ys = neuron_coords(x, N, spacing)

    draw_neuron.(x, ys, r; c=c)

end

function draw_layer(x, spacing, N1, N2, r)

    plot!(framestyle=:none, grid=:none)

    first_x = x
    second_x = x + 1

    first = neuron_coords(x + 1, N1, spacing)
    second = neuron_coords(x, N2, spacing)

    draw_neurons(x, N1, 1, r; c=:blue)
    draw_neurons(x+1, N2, 1, r; c=:red)

    for i in 1:N1
        for j in 1:N2

            vec = [second_x - first_x, second[j] - first[i]]
            normalize!(vec)

            start = [first_x, first[i]] + 1.2*r*vec
            finish = [second_x, second[j]] - 1.2*r*vec

            plot!([start[1], finish[1]], [start[2], finish[2]], c=:black, alpha=0.5)
        end
    end

end

#draw_layer(1, 1, 3, 4, 0.2)

function draw_link(x1, y1, x2, y2, r)
    vec = [x2 - x1, y2 - y1]
    normalize!(vec)

    start = [x1, y1] + 1.2 * r * vec
    finish = [x2, y2] - 1.2 * r * vec

    plot!([start[1], finish[1]], [start[2], finish[2]], c=:black, alpha=0.5)
end

"""
Takes a vector of neurons per layer
"""
function draw_network(neurons_per_layer)

    x_spacing = 1
    y_spacing = 1
    r = 0.2

    num_layers = length(neurons_per_layer)

    plot(framestyle=:none, grid=:none)

    # draw input links
    N1 = neurons_per_layer[1]

    for j in 1:N1
        y = neuron_coords(j, N1, y_spacing)
        draw_link(0, y, 1, y, r)
    end

    # draw neurons
    for layer in 1:length(neurons_per_layer)
        N = neurons_per_layer[layer]

        if layer == 1
            c = :green
        elseif layer == num_layers
            c = :red
        else
            c = :blue
        end

        for j in 1:N

            draw_neuron(layer, neuron_coords(j, N, y_spacing), r, c=c)

        end
    end

    # draw links
    for layer in 1:length(neurons_per_layer)-1
        N1 = neurons_per_layer[layer]
        N2 = neurons_per_layer[layer+1]

        for j1 in 1:N1
            for j2 in 1:N2

                draw_link(layer,  neuron_coords(j1, N1, y_spacing), layer+1,
                neuron_coords(j2, N2, y_spacing), r)
            end

        end
    end

    # draw output links
    N_last = neurons_per_layer[end]

    for j in 1:N_last
        y = neuron_coords(j, N_last, y_spacing)
        draw_link(num_layers, y, num_layers+1, y, r)
    end

    plot!()

end

draw_network([3, 2, 2])

max_stars_repo_name: grenkoca/JuliaTutorials

max_issues_repo_name: grenkoca/JuliaTutorials

max_forks_repo_name: grenkoca/JuliaTutorials


text:
"""
    Drop <: Imputor

Removes missing values from the `AbstractArray` or `DataFrame` provided.
"""
struct Drop <: Imputor end

"""
    impute!(imp::Drop, ctx::Context, data::AbstractVector)

Uses `filter!` to remove missing elements from the array.

# Arguments
* `imp::Drop`: this `Imputor` method
* `ctx::Context`: contextual information for missing data
* `data::AbstractVector`: the data to impute

# Returns
* `AbstractVector`: our data array with missing elements removed
"""
function impute!(imp::Drop, ctx::Context, data::AbstractVector)
    return filter!(x -> !ismissing(ctx, x), data)
end

"""
    impute!(imp::Drop, ctx::Context, data::AbstractMatrix)

Finds the missing rows in the matrix and uses a mask (Vector{Bool}) to return the
`data` with those rows removed. Unfortunately, the mask approach requires copying the matrix.

NOTES (or premature optimizations):
* We use `view`, but this will change the type of the `data` by returning a `SubArray`
* We might be able to do something clever by:
    1. reshaping the data to a vector
    2. running `deleteat!` for the appropriate indices and
    3. reshaping the data back to the desired shape.

# Arguments
* `imp::Drop`: this `Imputor` method
* `ctx::Context`: contextual information for missing data
* `data::AbstractMatrix`: the data to impute

# Returns
* `AbstractMatrix`: a new matrix with missing rows removed
"""
function impute!(imp::Drop, ctx::Context, data::AbstractMatrix)
    ctx.num = size(data, 1)
    mask = map(i -> !ismissing(ctx, data[i, :]), 1:size(data, 1))
    return data[mask, :]
end

"""
    impute!(imp::Drop, ctx::Context, data::DataFrame)

Finds the missing rows in the `DataFrame` and deletes them.

NOTE: this isn't quite as fast as `dropnull` in `DataFrames`s as we're using an arbitrary
`missing` function rather than using the precomputed `dt.isnull` vector of bools.

# Arguments
* `imp::Drop`: this `Imputor` method
* `ctx::Context`: contextual information for missing data
* `data::DataFrame`: the data to impute

# Returns
* `DataFrame`: our data with the missing rows removed.
"""
function impute!(imp::Drop, ctx::Context, data::DataFrame)
    ctx.num = size(data, 1)
    m = typeof(data).name.module
    m.deleterows!(data, findall(r -> ismissing(ctx, r), m.eachrow(data)))
    return data
end


text:
# This file is a part of RadiationDetectorSignals.jl, licensed under the MIT License (MIT).

# TODO: Replace by true custom types to avoid type piracy on NamedTuple
# (even though the named tuple types used here are quite specific).

const ArrayOfDims{N} = Array{T,N} where T
const AosAOfDims{M,N} = ArrayOfSimilarArrays{T,M,N} where T
const SArrayOfDims{N} = SArray{TP,T,N,L} where {TP,L,T}

Base.@pure recursive_ndims(::Type{<:AbstractArray{T,N}}) where {T,N} = N + recursive_ndims(T)
Base.@pure recursive_ndims(::Type{<:RealQuantity}) = 0


text:
@testset "Fitting" begin
  @testset "Block" begin
    # TODO
  end

  @testset "PoT" begin
    # TODO
  end
end


text:
# following variables need to be set before inclusion of this script
# CHILDNAME
# LATENT_D
# SIGMA_U
# SIGMA_V
# SIGMA_R
# LAMBDAREF
# MAXNEVENTS
# MAXT

using PDMP, JLD

params = Dict(
    "CHILDNAME"  => CHILDNAME,
    "LATENT_D"   => LATENT_D,
    "SIGMA_U"    => SIGMA_U,
    "SIGMA_V"    => SIGMA_V,
    "SIGMA_R"    => SIGMA_R,
    "LAMBDAREF"  => LAMBDAREF,
    "MAXNEVENTS" => MAXNEVENTS,
    "MAXT"       => MAXT
)

start = time()
scriptdir = dirname(@__FILE__)
println(scriptdir)
a = readdlm(scriptdir * "/data/ratings.csv", ',', Int)
R = a[:,1:3]

println("($(time()-start)s) -- read the data")

# srand(135)
latentD = LATENT_D     # dimension of latent space
sigmaU  = SIGMA_U
sigmaV  = SIGMA_V
sigmaR  = SIGMA_R

### there may be discrepancy with lines missing etc.
# -> use unique
nU = maximum(R[:,1])
nV = maximum(R[:,2])

# factors: create N factors for the users,
mvgU             = MvDiagonalGaussian(zeros(latentD), sigmaU)
gllU(x)          = gradloglik(mvgU, x)
nexteventU(x, v) = nextevent_bps(mvgU, x, v)
factorU(k)       = Factor(nexteventU, gllU, k)

allfactors = [factorU(k) for k in 1:nU]

# factors: create M factors for the movies,
mvgV             = MvDiagonalGaussian(zeros(latentD), sigmaV)
gllV(x)          = gradloglik(mvgV, x)
nexteventV(x, v) = nextevent_bps(mvgV, x, v)
factorV(k)       = Factor(nexteventV, gllV, nU+k)

allV = [factorV(k) for k in 1:nV]

push!(allfactors, allV...)

# factors: create nz(R) factors for the ratings
maskU(x) = x[1:latentD]
maskV(x) = x[latentD+1:end]

# each factor connected to its own var
structure = [[k] for k in 1:(nU+nV)]

# -----------------------------
# factors  variables
#  fU1       U1 = [1]
#  ...       ...
#  fUnU      UnU = [nU]
#  fV1       V1  = [nU+1]
#  ...       ...
#  fVnV      VnV = [1+nV]
# (for appropriate ij)
#  fRij      Ui,Vj = [i, nU+j]
# -----------------------------

for k in 1:size(R,1)
    i,j, rij = R[k,:]
    # the likelihood
    gij = PMFGaussian(rij, sigmaR, latentD)
    # the factor
    fij = Factor( (x,w)->nextevent_bps(gij,x,w),
                   x->gradloglik(gij, x),
                   nU+nV+k )
    push!(allfactors, fij)
    push!(structure, [i, j+nU])
end

fg = FactorGraph(structure, allfactors)

lambdaref  = LAMBDAREF
maxnevents = MAXNEVENTS
T          = MAXT

nvars      = nU+nV

x0 = [randn(latentD) for i in 1:nvars]
v0 = [randn(latentD) for i in 1:nvars]
v0 = map(v->v/norm(v), v0)

lsim = LocalSimulation(fg, x0, v0, T, maxnevents, lambdaref)

println("($(time()-start)s) -- created the graph + sim")

(all_evlist, details) = simulate(lsim)

println("($(time()-start)s) -- finished the simulation")

filepath = scriptdir * "/child_$(CHILDNAME).jld"
save(filepath,
        "evlist", all_evlist.evl,
        "details", details,
        "params", params)

println("($(time()-start)s) -- saved the results")

resourcegroup = "mortest42"
saspath  = scriptdir * "/secrets/azure_vm_pool_mortest42_sas_storage_container_data.txt"
storagescript = scriptdir * "/az-storage.py"
pushcommand = `python $storagescript $resourcegroup put -i $filepath --sas-path $saspath`

run(pushcommand)

println("($(time()-start)s) -- pushed the results to Azure")

max_stars_repo_name: alan-turing-institute/azure-batch-tools

max_issues_repo_name: alan-turing-institute/azure-batch-tools

max_forks_repo_name: alan-turing-institute/azure-batch-tools


text:
module VDPTag

using POMDPs
using StaticArrays
using Parameters
using Plots
using Distributions
using POMDPToolbox
using ParticleFilters
using AutoHashEquals

const Vec2 = SVector{2, Float64}

importall POMDPs
# import POMDPs: generate_s, generate_sr, generate_sor
# import POMDPs: reward, discount, isterminal
# import POMDPs: actions, initial_state_distribution
# import POMDPs: n_actions, n_states, n_observations
# import POMDPs: observation, pdf
# import POMDPs: action
import Base: rand, eltype, isnull, convert
import MCTS: next_action, n_children
import ParticleFilters: obs_weight

export
    TagState,
    TagAction,
    VDPTagMDP,
    VDPTagPOMDP,

    DiscreteVDPTagMDP,
    DiscreteVDPTagPOMDP,
    AODiscreteVDPTagPOMDP,
    ADiscreteVDPTagPOMDP,
    TranslatedPolicy,
    translate_policy,
    cproblem,

    convert_s,
    convert_a,
    convert_o,
    obs_weight,

    ToNextML,
    ToNextMLSolver,
    NextMLFirst,
    DiscretizedPolicy,
    ManageUncertainty,
    mdp

@auto_hash_equals immutable TagState
    agent::Vec2
    target::Vec2
end

@auto_hash_equals immutable TagAction
    look::Bool
    angle::Float64
end

@with_kw immutable VDPTagMDP <: MDP{TagState, Float64}
    mu::Float64         = 2.0
    dt::Float64         = 0.1
    step_size::Float64  = 0.5
    tag_radius::Float64 = 0.1
    tag_reward::Float64 = 100.0
    step_cost::Float64  = 1.0
    pos_std::Float64    = 0.05
    tag_terminate::Bool = true
    discount::Float64   = 0.95
end

@with_kw immutable VDPTagPOMDP <: POMDP{TagState, TagAction, Float64}
    mdp::VDPTagMDP          = VDPTagMDP()
    bearing_std::Float64    = deg2rad(2.0)
    bearing_cost::Float64   = 5.0
end

const VDPTagProblem = Union{VDPTagMDP,VDPTagPOMDP}
mdp(p::VDPTagMDP) = p
mdp(p::VDPTagPOMDP) = p.mdp

function next_ml_target(p::VDPTagMDP, pos::Vec2)
    steps = round(Int, p.step_size/p.dt)
    for i in 1:steps
        pos = rk4step(p, pos)
    end
    return pos
end

function generate_s(pp::VDPTagProblem, s::TagState, a::Float64, rng::AbstractRNG)
    p = mdp(pp)
    pos = next_ml_target(p, s.target)
    return TagState(s.agent+p.step_size*SVector(cos(a), sin(a)), pos+p.pos_std*SVector(randn(rng), randn(rng)))
end

function generate_sr(p::VDPTagProblem, s::TagState, a::Float64, rng::AbstractRNG)
    sp = generate_s(p, s, a, rng)
    return sp, reward(p, s, a, sp)
end

function reward(pp::VDPTagProblem, s::TagState, a::Float64, sp::TagState)
    p = mdp(pp)
    if norm(sp.agent-sp.target) < p.tag_radius
        return p.tag_reward
    else
        return -p.step_cost
    end
end

discount(pp::VDPTagProblem) = mdp(pp).discount
isterminal(pp::VDPTagProblem, s::TagState) = mdp(pp).tag_terminate && norm(s.agent-s.target) < mdp(pp).tag_radius

immutable AngleSpace end
rand(rng::AbstractRNG, ::AngleSpace) = 2*pi*rand(rng)
actions(::VDPTagMDP) = AngleSpace()

generate_s(p::VDPTagPOMDP, s::TagState, a::TagAction, rng::AbstractRNG) = generate_s(p, s, a.angle, rng)

immutable POVDPTagActionSpace end
rand(rng::AbstractRNG, ::POVDPTagActionSpace) = TagAction(rand(rng, Bool), 2*pi*rand(rng))
actions(::VDPTagPOMDP) = POVDPTagActionSpace()

function reward(p::VDPTagPOMDP, s::TagState, a::TagAction, sp::TagState)
    return reward(mdp(p), s, a.angle, sp) - a.look*p.bearing_cost
end

immutable NullableAngleNormal
    null::Bool
    mean::Float64
    std::Float64

    NullableAngleNormal() = new(true)
    NullableAngleNormal(mean::Float64, std::Float64) = new(false, mean, std)
end
isnull(n::NullableAngleNormal) = n.null
sampletype(::Type{NullableAngleNormal}) = Float64
function pdf(d::NullableAngleNormal, o::Float64)
    if isnull(d)
        return 1.0
    else
        dir_diff = abs(o-d.mean)
        while dir_diff > pi
            dir_diff -= 2*pi
        end
        dir_diff = abs(dir_diff)
        return exp(-dir_diff^2/(2*d.std^2))
    end
end
function rand(rng::AbstractRNG, d::NullableAngleNormal)
    if isnull(d)
        return 2*pi*rand(rng)
    else
        return d.mean + d.std*randn(rng)
    end
end

function observation(p::VDPTagPOMDP, a::TagAction, sp::TagState)
    if a.look
        diff = sp.target-sp.agent
        bearing = atan2(diff[2], diff[1])
        return NullableAngleNormal(bearing, p.bearing_std)
    else
        return NullableAngleNormal()
    end
end

include("rk4.jl")
include("initial.jl")
include("discretized.jl")
include("visualization.jl")
include("heuristics.jl")

end # module


text:
using Test
using NQCDynamics
using NQCDynamics.InitialConditions.QuantisedDiatomic
using LinearAlgebra: norm
using Unitful
using UnitfulAtomic

@testset "separate/combine_slab_and_molecule" begin

    @testset "3D, with slab" begin
        atom_indices = [2, 3]
        r0 = [
            1 2 3 4 5;
            6 7 8 9 10;
            11 12 13 14 15
        ]
        molecule, slab = QuantisedDiatomic.separate_slab_and_molecule(atom_indices, r0)
        @test molecule == [2 3; 7 8; 12 13]
        @test slab == [1 4 5; 6 9 10; 11 14 15]

        r = QuantisedDiatomic.combine_slab_and_molecule(atom_indices, molecule, slab)
        @test r == r0
    end

    @testset "3D, no slab" begin
        atom_indices = [1, 2]
        r0 = [1 2; 3 4; 5 6]
        molecule, slab = QuantisedDiatomic.separate_slab_and_molecule(atom_indices, r0)
        @test molecule == r0

        r = QuantisedDiatomic.combine_slab_and_molecule(atom_indices, molecule, slab)
        @test r == r0
    end

    @testset "1D, no slab" begin
        atom_indices = [1]
        r0 = [1;;]
        molecule, _ = QuantisedDiatomic.separate_slab_and_molecule(atom_indices, r0)
        @test molecule == r0
    end
end

@testset "calculate_force_constant" begin
    bond_lengths = 0.5:0.01:5.0
    model = NQCModels.DiatomicHarmonic(r₀=2.0)
    environment = QuantisedDiatomic.EvaluationEnvironment([1, 2], (3, 2), zeros(3,0), 0.0, [0, 0, 1.0])
    binding_curve = QuantisedDiatomic.calculate_binding_curve(bond_lengths, model, environment)

    @test QuantisedDiatomic.calculate_force_constant(binding_curve) ≈ 1 atol=0.1

    model = NQCModels.DiatomicHarmonic(r₀=3.0)
    binding_curve = QuantisedDiatomic.calculate_binding_curve(bond_lengths, model, environment)
    @test QuantisedDiatomic.calculate_force_constant(binding_curve) ≈ 1 atol=0.1
end

@testset "subtract_centre_of_mass!" begin
    atoms = Atoms([:O, :O])
    r = rand(3, 2)
    r = QuantisedDiatomic.subtract_centre_of_mass(r, atoms.masses)
    @test r[:,1] ≈ -r[:,2]
end

@testset "apply_random_rotation!" begin
    x = fill(5.0, 3, 2)
    y = fill(-4.0, 3, 2)
    @test all(x .== 5)
    @test all(y .== -4)
    QuantisedDiatomic.apply_random_rotation!(x, y)
    @test all(x .!= 5)
    @test all(y .!= -4)
    for col in eachcol(x)
        @test norm(col) ≈ norm([5, 5, 5])
    end
    for col in eachcol(y)
        @test norm(col) ≈ norm([-4, -4, -4])
    end
end

@testset "generate and check results" begin
    atoms = Atoms([:N, :O])
    model = DiatomicHarmonic(r₀=2)
    sim = Simulation(atoms, model)

    numbers = [2, 5]
    configs = generate_configurations(sim, numbers[1], numbers[2]; samples=10, translational_energy=1u"eV")
    for config in configs
        ν, J = quantise_diatomic(sim, config...)
        @test ν ≈ numbers[1] rtol=1e-1
        @test J ≈ numbers[2] rtol=1e-1
    end
end

@testset "position_above_surface!" begin
    r = zeros(3, 2)
    height = 10
    QuantisedDiatomic.position_above_surface!(r, [0, 0, height], InfiniteCell())
    @test r == [0 0; 0 0; height height]
end

@testset "velocity_from_energy" begin
    v = QuantisedDiatomic.velocity_from_energy([100, 100], 10u"eV")
    v2 = QuantisedDiatomic.velocity_from_energy([100, 100], austrip(10u"eV"))
    @test v ≈ v2
end

@testset "apply_translational_impulse!" begin
    atoms = Atoms([:N, :O])
    model = DiatomicHarmonic(r₀=2)
    sim = Simulation(atoms, model)

    masses = [100, 100]
    r = rand(3, 2) .+ [0 2; 0 2; 0 2]
    v = rand(3, 2) ./ 100
    v_before = copy(v)
    e = 10u"eV"
    direction = [0, 0.5, 0.5]
    QuantisedDiatomic.apply_translational_impulse!(v, masses, e, direction)
    @test v_before[1,:] == v[1,:]
    @test all(v_before[2:3,:] .< v[2:3,:])

    νi, Ji = quantise_diatomic(sim, v, r)
    QuantisedDiatomic.apply_translational_impulse!(v, sim.atoms.masses, 1, rand(3))
    νf, Jf = quantise_diatomic(sim, v, r)
    @test νi ≈ νf
    @test Ji ≈ Jf
end

@testset "find_total_energy" begin
    environment = QuantisedDiatomic.EvaluationEnvironment([1, 2], (3, 2), zeros(3,0), 0.0, [0, 0, 1.0])
    model = DiatomicHarmonic(r₀=2.0)
    bond_lengths = 0.5:0.01:5.0
    binding_curve = QuantisedDiatomic.calculate_binding_curve(bond_lengths, model, environment)
    μ = 1.0
    J = 2
    ν = 1
    V = QuantisedDiatomic.EffectivePotential(μ, J, binding_curve)
    E, bounds = QuantisedDiatomic.find_total_energy(V, ν)
end

@testset "find_integral_bounds" begin
    environment = QuantisedDiatomic.EvaluationEnvironment([1, 2], (3, 2), zeros(3,0), 0.0, [0, 0, 1.0])
    model = DiatomicHarmonic(r₀=2)
    bond_lengths = 0.5:0.01:5.0
    binding_curve = QuantisedDiatomic.calculate_binding_curve(bond_lengths, model, environment)
    total_energy = 1
    J = 0
    μ = 1.0
    equilibrium_bond_length = 2
    V = QuantisedDiatomic.EffectivePotential(μ, J, binding_curve)
    r₁, r₂ = QuantisedDiatomic.find_integral_bounds(total_energy, V)
    @test r₁ < equilibrium_bond_length
    @test r₂ > equilibrium_bond_length
end

@testset "calculate_binding_curve" begin
    model = Morse()
    bond_lengths = 0.5:0.01:5.0
    environment = QuantisedDiatomic.EvaluationEnvironment([1], (1, 1), zeros(1,0), 0.0, [0.0])
    binding_curve = QuantisedDiatomic.calculate_binding_curve(bond_lengths, model, environment)
    @test binding_curve.equilibrium_bond_length ≈ model.x₀
end

@testset "calculate_diatomic_energy" begin
    @testset "3D, no slab" begin
        model = DiatomicHarmonic(r₀=2)
        environment = QuantisedDiatomic.EvaluationEnvironment([1, 2], (3, 2), zeros(3,0), 0.0, [0, 0, 1.0])
        @test QuantisedDiatomic.calculate_diatomic_energy(2.0, model, environment) ≈ 0.0 atol=1e-3
        @test QuantisedDiatomic.calculate_diatomic_energy(3.0, model, environment) ≈ 0.5
    end

    @testset "1D, no slab" begin
        model = Harmonic(dofs=1)
        bond_length = 1.5
        environment = QuantisedDiatomic.EvaluationEnvironment([1], (1, 1), zeros(1,0), 0.0, [0.0])
        @test QuantisedDiatomic.calculate_diatomic_energy(bond_length, model, environment) ≈ bond_length^2 / 2
    end
end

@testset "assemble_evaluation_geometry" begin
    bond_length = 4.3

    @testset "3D, no slab" begin
        environment = QuantisedDiatomic.EvaluationEnvironment([1, 2], (3, 4), rand(3,2), 2.0, [0, 0, 1.0])
        r = QuantisedDiatomic.assemble_evaluation_geometry(bond_length, environment)
        @test norm(r[:,1] .- r[:,2]) == 4.3
    end

    @testset "1D, no slab" begin
        environment = QuantisedDiatomic.EvaluationEnvironment([1], (1, 1), zeros(1,0), 0.0, [0.0])
        r = QuantisedDiatomic.assemble_evaluation_geometry(bond_length, environment)
        @test r == [4.3;;]
    end
end

@testset "build_molecule" begin
    bond_length = 4.3

    @testset "3D, no slab" begin
        environment = QuantisedDiatomic.EvaluationEnvironment([1, 2], (3, 4), rand(3,2), 2.0, [0, 0, 1.0])
        r = QuantisedDiatomic.build_molecule(bond_length, environment)
        @test norm(r[:,1] .- r[:,2]) == 4.3
    end

    @testset "1D, no slab" begin
        environment = QuantisedDiatomic.EvaluationEnvironment([1], (1, 1), zeros(0,0), 0.0, [0.0])
        r = QuantisedDiatomic.build_molecule(bond_length, environment)
        @test r == [4.3;;]
    end
end

@testset "fit_binding_curve and plot_binding_curve" begin
    model = Morse()
    bond_lengths = 0.5:0.01:5.0
    binding_curve = potential.(model, bond_lengths)
    fit = QuantisedDiatomic.fit_binding_curve(bond_lengths, binding_curve)
    QuantisedDiatomic.plot_binding_curve(bond_lengths, binding_curve, fit)
end

@testset "generate_configurations 1D" begin
    sim = Simulation(Atoms(2000), Morse())
    r, v = QuantisedDiatomic.generate_1D_vibrations(Morse(), 1000.0, 0; samples=100)
end


text:
using LinearAlgebra
using SparseArrays
using Plots

"This routine solves the periodic advection equation using finite volumes."

m = 200 # number of points

# define spatial grid
xv = LinRange(-1,1,m+2)
Δx = xv[2]-xv[1]
x = xv[1:end-1] .+ Δx/2 # m+1 cell centers

a = 1 # advection speed
Δt = .9*Δx / a # timestep
T = 2.0 # final time

# initial condition and forcing
u0(x) = exp(-25*x^2)
u = u0.(x)
f(uL,uR) = uL

# indexing for periodic boundaries
index_left = [m+1; 1:m+1]
index_right = [1:m+1; 1]
fv_flux = zeros(m+2) # storage at interfaces

Nsteps = ceil(Int,T/Δt)
Δt = T / Nsteps
interval = 10
plot()
unorm = zeros(Nsteps)
@gif for k = 1:Nsteps
    for i = 1:m+2 # loop over cell interfaces
        left = index_left[i]
        right = index_left[i]
        fv_flux[i] = f(u[left],u[right])
    end
    for i = 1:m+1
        u[i] = u[i] - a*Δt/Δx * (fv_flux[i+1]-fv_flux[i])
    end

    if k % interval==0
        plot(x,u,linewidth=2,legend=false,title="Solution at time $(k*Δt)",ylims=(-1.5,1.5))
        println("on timestep $k out of $Nsteps.")
    end
end every interval
# plot(x,u,linewiΔth=2,legend=false,title="Solution at final time $(T)",ylims=(-1.5,1.5))


text:
import LinearAlgebra: I, mul!

struct MESH
    Δx::Float64
    Δy::Float64
    x::Array{Float64,1}
    y::Array{Float64,1}
    X::Array{Float64,2}
    Y::Array{Float64,2}
    xsize::Int64
    ysize::Int64
end

#Eulerian Mesh
function EulerianMesh(xsize=Int(3e5), ysize=Int(2e5), xnum=61, ynum=41)
    Δx = xsize/(xnum-1)
    Δy = ysize/(ynum-1)
    x = 0:Δx:xsize
    y = 0:Δy:ysize
    X = x'.*ones(length(y))
    Y = y.*ones(1,length(x))
    return MESH(Δx,Δy,x,y,X,Y,Int(xsize),Int(ysize))
end

mutable struct Particles
    x::Array{Float64,1}
    y::Array{Float64,1}
    n::Int64 # number of particles
end

#Lagrangian Particles
function LagrangianParticles(Δx,Δy,mxcell=2, mycell=2, xnum=61, ynum=41)
    mΔx = Δx/mxcell
    mΔy = Δy/mycell
    mxnum = (xnum-1)*mxcell
    mynum = (ynum-1)*mycell
    marknum = mxnum*mynum
    mx = Vector{Float64}(undef,marknum)
    my = similar(mx)
    @inbounds for i=1:mxnum, j=1:mynum
        m = i + (j-1)*mxnum
        mx[m] = mΔx*(i-0.5)
        my[m] = mΔy*(j-0.5)
    end
    return Particles(mx,my,marknum)
end

struct Velocity
    x::Array{Float64,2}
    y::Array{Float64,2}
end

# Eulerian velocity field
function EulerianVelocity(M)
    Vscale = 1
    xnum, ynum = size(M.X,2), size(M.X,1)
    Vx = Array{Float64,2}(undef,ynum,xnum)
    Vy = similar(Vx)
    kx, ky = 1, 1 # wavenumber
    vx0 = 1e-9/Vscale
    vy0 = vx0*kx/ky*M.ysize/M.xsize # Incompressibility
    cte₁ = kx*π/M.xsize
    cte₂ = ky*π/M.ysize
    @inbounds for i=1:xnum, j=1:ynum
        Vx[j,i]=-vx0*sin(cte₁*M.x[i])*cos(cte₂*M.y[j])
        Vy[j,i]= vy0*cos(cte₁*M.x[i])*sin(cte₂*M.y[j])
    end
    return Velocity(Vx, Vy)
end

struct Gradient
    ∂x∂x::Array{Float64,2}
    ∂x∂y::Array{Float64,2}
    ∂y∂y::Array{Float64,2}
    ∂y∂x::Array{Float64,2}
end

# Velocity gradient
function VelocityGradient(M)
    xnum, ynum = size(M.X,2), size(M.X,1)
    dVxdx = Array{Float64,2}(undef,ynum,xnum)
    dVxdy, dVydx, dVydy = similar(dVxdx), similar(dVxdx), similar(dVxdx)
    Vscale = 1
    vx0 = 1e-9/Vscale
    kx, ky = 1, 1 # wavenumber
    vy0 = vx0*kx/ky*M.ysize/M.xsize # Incompressibility
    cte₁ = kx*π/M.xsize
    cte₂ = ky*π/M.ysize
    @inbounds for i=1:xnum, j=1:ynum
        dVxdx[j,i] =-vx0*cos(cte₁*M.x[i])*cte₁*cos(cte₂*M.y[j]);
        dVxdy[j,i] = vx0*sin(cte₁*M.x[i])*sin(cte₂*M.y[j])*cte₂;
        dVydx[j,i] =-vy0*sin(cte₁*M.x[i])*cte₁*sin(cte₂*M.y[j]);
        dVydy[j,i] = vy0*cos(cte₁*M.x[i])*cos(cte₂*M.y[j])*cte₂;
    end
    return Gradient(dVxdx,dVxdy,dVydy,dVydx)
end

function interp∇U!(l,m,∇U,P,M)
    xnum, ynum = size(M.X,2), size(M.X,1)
    i = floor(Int,P.x[m]/M.Δx) + 1
    j = floor(Int,P.y[m]/M.Δy) + 1

    if i<1
        i = 1
        Δxx = 0.0
    elseif i > xnum-1
        i = xnum-1
        Δxx = 1.0
    else
        Δxx = (P.x[m]-M.x[i])/M.Δx
    end
    if j<1
        j = 1
        Δyy = 0.0
    elseif j > ynum-1
        j = ynum-1
        Δyy = 1.0
    else
        Δyy = (P.y[m]-M.y[j])/M.Δy
    end

    w1 = (1.0-Δxx)*(1.0-Δyy)
    w2 = (1.0-Δxx)*Δyy
    w3 = Δxx*(1.0-Δyy)
    w4 = Δxx*Δyy

    @inbounds l[1,1] = ∇U.∂x∂x[j,i]*w1 + ∇U.∂x∂x[j+1,i]*w2 + ∇U.∂x∂x[j,i+1]*w3 + ∇U.∂x∂x[j+1,i+1]*w4
    @inbounds l[1,2] = ∇U.∂x∂y[j,i]*w1 + ∇U.∂x∂y[j+1,i]*w2 + ∇U.∂x∂y[j,i+1]*w3 + ∇U.∂x∂y[j+1,i+1]*w4
    @inbounds l[2,1] = ∇U.∂y∂y[j,i]*w1 + ∇U.∂y∂y[j+1,i]*w2 + ∇U.∂y∂y[j,i+1]*w3 + ∇U.∂y∂y[j+1,i+1]*w4
    @inbounds l[2,2] = ∇U.∂y∂x[j,i]*w1 + ∇U.∂y∂x[j+1,i]*w2 + ∇U.∂y∂x[j,i+1]*w3 + ∇U.∂y∂x[j+1,i+1]*w4

end

function interpVel(mxA,myA,M,U)
    xnum, ynum = size(M.X,2), size(M.X,1)

    i = floor(Int, mxA/M.Δx)+1
    j = floor(Int, myA/M.Δy)+1

    if i<1
        i = 1
        Δxx = 0.0
    elseif i > xnum-1
        i = xnum-1
        Δxx = 1.0
    else
        Δxx = (mxA-M.x[i])/M.Δx
    end
    if j<1
        j = 1
        Δyy = 0.0
    elseif j > ynum-1
        j = ynum-1
        Δyy = 1.0
    else
        Δyy = (myA-M.y[j])/M.Δy
    end

    w1 = (1.0-Δxx)*(1.0-Δyy)
    w2 = (1.0-Δxx)*Δyy
    w3 = Δxx*(1.0-Δyy)
    w4 = Δxx*Δyy

    @inbounds Ux = U.x[j,i]*w1 + U.x[j+1,i]*w2 + U.x[j,i+1]*w3 + U.x[j+1,i+1]*w4
    @inbounds Uz = U.y[j,i]*w1 + U.y[j+1,i]*w2 + U.y[j,i+1]*w3 + U.y[j+1,i+1]*w4
    return Ux, Uz
end

function advectparticle!(P,M,U,Δt,m)
    mxA = P.x[m]
    myA = P.y[m]
    # 4th order RK
    mVxA, mVyA = interpVel(mxA,myA,M,U)
    mxB = mxA + 0.5*mVxA*Δt; myB = myA + 0.5*mVyA*Δt
    mVxB, mVyB = interpVel(mxB,myB,M,U)
    mxC = mxA + 0.5*mVxB*Δt; myC = myA + 0.5*mVyB*Δt
    mVxC, mVyC = interpVel(mxC,myC,M,U)
    mxD = mxA + mVxC*Δt; myD = myA + mVyC*Δt
    mVxD, mVyD = interpVel(mxD,myD,M,U)
    P.x[m] += (mVxA + 2*(mVxB + mVxC) + mVxD)/6*Δt
    P.y[m] += (mVyA + 2*(mVyB + mVyC) + mVyD)/6*Δt
end

mutable struct RK
    Δ1::Array{Float64,2}
    Δ2::Array{Float64,2}
    Δ3::Array{Float64,2}
    Δ4::Array{Float64,2}
end

function updateFSE(F0,l,Δt,RK4)
    Fᵢ = deepcopy(F0)
    # # 4th-order Runge-Kutta
    # RK4.Δ1 .= l*Fᵢ*Δt
    # Fᵢ = F0 + 0.5*RK4.Δ1
    # RK4.Δ2 .= l*Fᵢ*Δt
    # Fᵢ = F0 + 0.5*RK4.Δ2
    # RK4.Δ3 .= l*Fᵢ*Δt
    # Fᵢ = F0 + RK4.Δ3
    # RK4.Δ4 .= l*Fᵢ*Δt

    # 4th-order Runge-Kutta
    mul!(RK4.Δ1, l,Fᵢ*Δt)
    Fᵢ = F0 + 0.5*RK4.Δ1
    mul!(RK4.Δ2, l,Fᵢ*Δt)
    Fᵢ = F0 + 0.5*RK4.Δ2
    mul!(RK4.Δ3, l,Fᵢ*Δt)
    Fᵢ = F0 + RK4.Δ3
    mul!(RK4.Δ4, l,Fᵢ*Δt)

    return F0 += (RK4.Δ1 + 2*(RK4.Δ2 + RK4.Δ3) + RK4.Δ4)/6
end

max_stars_repo_name: albert-de-montserrat/ForwardFSE

max_issues_repo_name: albert-de-montserrat/ForwardFSE

max_forks_repo_name: albert-de-montserrat/ForwardFSE


text:
"""
    ModalTrav(A::AbstractMatrix)

    takes coefficients as provided by the Zernike polynomial layout of FastTransforms.jl and
    makes them accessible sorted such that in each block the m=0 entries are always in first place,
    followed by alternating sin and cos terms of increasing |m|.
"""
struct ModalTrav{T, AA<:AbstractMatrix{T}} <: AbstractBlockVector{T}
    matrix::AA
    function ModalTrav{T, AA}(matrix::AA) where {T,AA<:AbstractMatrix{T}}
        m,n = size(matrix)
        if isfinite(m)
            isfinite(n) && isodd(n) && m == n ÷ 4 + 1 || throw(ArgumentError("size must match"))
        end
        new{T,AA}(matrix)
    end
end

ModalTrav{T}(matrix::AbstractMatrix{T}) where T = ModalTrav{T,typeof(matrix)}(matrix)
ModalTrav(matrix::AbstractMatrix{T}) where T = ModalTrav{T}(matrix)

function ModalTrav{T}(::UndefInitializer, n::Int) where T
    N =  (isqrt(8n+1)-1) ÷ 2
    @assert sum(1:N) == n
    m = N ÷ 2 + 1
    n = 4(m-1) + 1
    ModalTrav(Matrix{T}(undef, m, n))
end

convert(::Type{ModalTrav{T,M}}, v::ModalTrav) where {T,M} = ModalTrav{T,M}(convert(M, v.matrix))

function convert(::Type{ModalTrav{T,M}}, v_in::AbstractVector) where {T,M}
    N =  (isqrt(8length(v_in)+1)-1) ÷ 2
    v = PseudoBlockVector(v_in, OneTo(N))
    m = N ÷ 2 + 1
    n = 4(m-1) + 1
    mat = zeros(T, m, n)
    for K in blockaxes(v,1)
        K̃ = Int(K)
        w = v[K]
        if isodd(K̃)
            mat[K̃÷2 + 1,1] = w[1]
            for j = 2:2:K̃-1
                mat[K̃÷2-j÷2+1,2(j-1)+2] = w[j]
                mat[K̃÷2-j÷2+1,2(j-1)+3] = w[j+1]
            end
        else
            for j = 1:2:K̃
                mat[K̃÷2-j÷2,2(j-1)+2] = w[j]
                mat[K̃÷2-j÷2,2(j-1)+3] = w[j+1]
            end
        end
    end
    ModalTrav{T,M}(mat)
end

ModalTrav{T,M}(v::AbstractVector) where {T,M} = convert(ModalTrav{T,M}, v)
ModalTrav{T}(v::AbstractVector) where T = ModalTrav{T,Matrix{T}}(v)
ModalTrav(v::AbstractVector{T}) where T = ModalTrav{T}(v)

copy(A::ModalTrav) = ModalTrav(copy(A.matrix))

_diviffinite(n) = div(n,2,RoundUp)
_diviffinite(n::InfiniteCardinal) = n

axes(A::ModalTrav) = (blockedrange(oneto(_diviffinite(size(A.matrix,2)))),)

getindex(A::ModalTrav, K::Block{1}) = _modaltravgetindex(A.matrix, K)

_modaltravgetindex(mat, K) = _modaltravgetindex(MemoryLayout(mat), mat, K)
function _modaltravgetindex(_, mat, K::Block{1})
    k = Int(K)
    m = k ÷ 2 + 1
    n = 4(m-1) + 1
    _modaltravgetindex(Matrix(mat[1:m, 1:n]), K)
end

function _modaltravgetindex(::AbstractStridedLayout, mat, K::Block{1})
    k = Int(K)
    k == 1 && return mat[1:1]
    k == 2 && return mat[1,2:3]
    st = stride(mat,2)
    if isodd(k)
        # nonnegative terms
        p = mat[range(k÷2+1, step=4*st-1, length=k÷2+1)]
        # negative terms
        n = mat[range(k÷2+3*st, step=4*st-1, length=k÷2)]
        interlace(p,n)
    else
        # negative terms
        n = mat[range(st+k÷2, step=4*st-1, length=k÷2)]
        # positive terms
        p = mat[range(2st+k÷2, step=4*st-1, length=k÷2)]
        interlace(n,p)
    end
end

getindex(A::ModalTrav, k::Int) = A[findblockindex(axes(A,1), k)]
function setindex!(A::ModalTrav, v, k::Int)
    Kk = findblockindex(axes(A,1), k)
    K,j = block(Kk),blockindex(Kk)
    K̃ = Int(K)
    mat = A.matrix
    if isodd(K̃)
        if j == 1
            mat[K̃÷2 + 1,1] = v
        elseif iseven(j)
            mat[K̃÷2-j÷2+1,2(j-1)+2] = v
        else
            mat[K̃÷2-(j-1)÷2+1,2(j-2)+3] = v
        end
    else
        if iseven(j)
            mat[K̃÷2-(j-1)÷2,2(j-2)+3] = v
        else
            mat[K̃÷2-j÷2,2(j-1)+2] = v
        end
    end
    A
end

similar(A::ModalTrav, ::Type{T}) where T = ModalTrav(similar(A.matrix, T))
function fill!(A::ModalTrav, x)
    fill!(A.matrix, x)
    A
end

struct ModalTravStyle <: AbstractBlockStyle{1} end

ModalTravStyle(::Val{1}) = ModalTravStyle()

BroadcastStyle(::Type{<:ModalTrav}) = ModalTravStyle()
BroadcastStyle(a::ModalTravStyle, b::DefaultArrayStyle{0}) = ModalTravStyle()
BroadcastStyle(a::DefaultArrayStyle{0}, b::ModalTravStyle) = ModalTravStyle()
BroadcastStyle(a::ModalTravStyle, b::DefaultArrayStyle{M}) where {M} = BroadcastStyle(BlockStyle{1}(), b)
BroadcastStyle(a::DefaultArrayStyle{M}, b::ModalTravStyle) where {M} = BroadcastStyle(a, BlockStyle{1}())

function similar(bc::Broadcasted{ModalTravStyle}, ::Type{T}) where T
    N = blocklength(axes(bc,1))
    n = 2N-1
    m = n ÷ 4 + 1
    ModalTrav(Matrix{T}(undef,m,n))
end

_modal2matrix(a::ModalTrav) = a.matrix
_modal2matrix(a::Broadcasted) = broadcasted(a.f, map(_modal2matrix, a.args)...)
_modal2matrix(a) = a

function copyto!(dest::ModalTrav, bc::Broadcasted{ModalTravStyle})
    broadcast!(bc.f, dest.matrix, map(_modal2matrix, bc.args)...)
    dest
end

function resize!(a::ModalTrav, N::Block{1})
    n = 2Int(N)-1
    m = n ÷ 4 + 1
    ModalTrav(a.matrix[1:m,1:n])
end

"""
    ModalInterlace(ops, (M,N), (l,u))

interlaces the entries of a vector of banded matrices
acting on the different Fourier modes. That is, a ModalInterlace
multiplying a DiagTrav is the same as the operators multiplying the matrix
that the DiagTrav wraps. We assume the same operator acts on the Sin and Cos.
"""
struct ModalInterlace{T, MMNN<:Tuple} <: AbstractBandedBlockBandedMatrix{T}
    ops
    MN::MMNN
    bandwidths::NTuple{2,Int}
end

ModalInterlace{T}(ops, MN::NTuple{2,Integer}, bandwidths::NTuple{2,Int}) where T = ModalInterlace{T,typeof(MN)}(ops, MN, bandwidths)
ModalInterlace(ops::AbstractVector{<:AbstractMatrix}, MN::NTuple{2,Integer}, bandwidths::NTuple{2,Int}) = ModalInterlace{eltype(eltype(ops))}(ops, MN, bandwidths)

axes(Z::ModalInterlace) = blockedrange.(oneto.(Z.MN))

blockbandwidths(R::ModalInterlace) = R.bandwidths
subblockbandwidths(::ModalInterlace) = (0,0)
copy(M::ModalInterlace) = M

function Base.view(R::ModalInterlace{T}, KJ::Block{2}) where T
    K,J = KJ.n
    dat = Matrix{T}(undef,1,J)
    l,u = blockbandwidths(R)
    if iseven(J-K) && -l ≤ J - K ≤ u
        sh = (J-K)÷2
        if isodd(K)
            k = K÷2+1
            dat[1,1] = R.ops[1][k,k+sh]
        end
        for m in range(2-iseven(K); step=2, length=J÷2-max(0,sh))
            k = K÷2-m÷2+isodd(K)
            dat[1,m] = dat[1,m+1] = R.ops[m+1][k,k+sh]
        end
    else
        fill!(dat, zero(T))
    end
    _BandedMatrix(dat, K, 0, 0)
end

getindex(R::ModalInterlace, k::Integer, j::Integer) = R[findblockindex.(axes(R),(k,j))...]

struct ModalInterlaceLayout <: AbstractBandedBlockBandedLayout end
struct LazyModalInterlaceLayout <: AbstractLazyBandedBlockBandedLayout end

MemoryLayout(::Type{<:ModalInterlace}) = ModalInterlaceLayout()
MemoryLayout(::Type{<:ModalInterlace{<:Any,NTuple{2,InfiniteCardinal{0}}}}) = LazyModalInterlaceLayout()
sublayout(::Union{ModalInterlaceLayout,LazyModalInterlaceLayout}, ::Type{<:NTuple{2,BlockSlice{<:BlockOneTo}}}) = ModalInterlaceLayout()

function sub_materialize(::ModalInterlaceLayout, V::AbstractMatrix{T}) where T
    kr,jr = parentindices(V)
    KR,JR = kr.block,jr.block
    M,N = Int(last(KR)), Int(last(JR))
    R = parent(V)
    ModalInterlace{T}([layout_getindex(R.ops[m],1:(M-m+2)÷2,1:(N-m+2)÷2) for m=1:min(N,M)], (M,N), R.bandwidths)
end

# act like lazy array
BroadcastStyle(::Type{<:ModalInterlace{<:Any,NTuple{2,InfiniteCardinal{0}}}}) = LazyArrayStyle{2}()

# TODO: overload muladd!
function *(A::ModalInterlace, b::ModalTrav)
    M = b.matrix
    ret = ModalTrav{promote_type(eltype(A), eltype(b))}(undef, size(A,1)).matrix
    mul!(view(ret,:,1), A.ops[1], M[:,1])
    for j = 1:size(ret,2)÷4
        mul!(@view(ret[1:end-j,4j-2]), A.ops[2j], @view(M[1:end-j,4j-2]))
        mul!(@view(ret[1:end-j,4j-1]), A.ops[2j], @view(M[1:end-j,4j-1]))
        mul!(@view(ret[1:end-j,4j]), A.ops[2j+1], @view(M[1:end-j,4j]))
        mul!(@view(ret[1:end-j,4j+1]), A.ops[2j+1], @view(M[1:end-j,4j+1]))
    end
    ModalTrav(ret)
end

function \(A::ModalInterlace, b::ModalTrav)
    M = b.matrix
    ret = similar(M, promote_type(eltype(A),eltype(b)))
    ldiv!(view(ret,:,1), A.ops[1], M[:,1])
    for j = 1:size(M,2)÷4
        ldiv!(@view(ret[1:end-j,4j-2]), A.ops[2j], @view(M[1:end-j,4j-2]))
        ldiv!(@view(ret[1:end-j,4j-1]), A.ops[2j], @view(M[1:end-j,4j-1]))
        ldiv!(@view(ret[1:end-j,4j]), A.ops[2j+1], @view(M[1:end-j,4j]))
        ldiv!(@view(ret[1:end-j,4j+1]), A.ops[2j+1], @view(M[1:end-j,4j+1]))
    end
    ModalTrav(ret)
end


text:
function expression_list()
    a = (1, 2)
    b = (7 * 8, 5 - 6)
    c = (1, 2)
    @assert(a == (1, 2))
    @assert(b == (56, -1))
    @assert(c == (1, 2))
    @assert((a, 4) == ((1, 2), 4))
    @assert(isa(a, Tuple))
end

function starred_item()
    a, b..., c = 0:5
    @assert(a == 0)
    @assert(b == [1, 2, 3, 4])
    @assert(c == 5)
    @assert(Dict("x" => 1, Dict("y" => 2)...) == Dict("x" => 1, "y" => 2))
    @assert([0:3..., 4] == [0, 1, 2, 3, 4])
    @assert(Set([0:3..., 4]) == Set([0, 1, 2, 3, 4]))
end

function starred_list()
    n = [1, 2, 3, 4]
    @assert([n..., 5, 6] == [1, 2, 3, 4, 5, 6])
    @assert(([1]..., [2]..., 3) == (1, 2, 3))
    @assert(([1, 2, 3]..., 4 * 1, 5) == (1, 2, 3, 4, 5))
end

if abspath(PROGRAM_FILE) == @__FILE__
    expression_list()
    starred_item()
    starred_list()
end


text:
"""
    linear interpolation between two vectors over horizon T
"""
function linear_interpolation(x0, xf, T)
    n = length(x0)
    X = [copy(Array(x0)) for t = 1:T]
    for t = 1:T
        for i = 1:n
            X[t][i] = (xf[i] - x0[i]) / (T - 1) * (t - 1) + x0[i]
        end
    end
    return X
end

function row_col!(row, col, r, c)
    for cc in c
        for rr in r
            push!(row, convert(Int, rr))
            push!(col, convert(Int, cc))
        end
    end
    return row, col
end

function row_col_cartesian!(row, col, r, c)
    for i = 1:length(r)
        push!(row, convert(Int, r[i]))
        push!(col, convert(Int, c[i]))
    end
    return row, col
end

"""
    random controls
"""
function random_controls(model, T, r = 1.0)
	[r * ones(model.m) for t = 1:T-1]
end

"""
    control bounds trajectories
"""
function control_bounds(model, T, l, u)
    l_traj = [l * ones(model.m) for t = 1:T-1]
    u_traj = [u * ones(model.m) for t = 1:T-1]

    return l_traj, u_traj
end

function control_bounds(model, T,
        l::AbstractVector{S} = -Inf * ones(model.m),
        u::AbstractVector{S} = Inf * ones(model.m)) where S

    l_traj = [copy(l) for t = 1:T-1]
    u_traj = [copy(u) for t = 1:T-1]

    return l_traj, u_traj
end

"""
    state bounds trajectories
"""
function state_bounds(model, T,
        l::AbstractVector{S} = -Inf * ones(model.n),
        u::AbstractVector{S} = Inf * ones(model.n);
        x1 = Inf * ones(model.n),
        xT = Inf * ones(model.n)) where S

    l_traj = [copy(l) for t = 1:T]
    u_traj = [copy(u) for t = 1:T]

    # initial and final conditions
    x1_mask = isfinite.(x1)
    xT_mask = isfinite.(xT)

    if length(x1_mask) > 0
        l_traj[1][x1_mask] = copy(x1[x1_mask])
        u_traj[1][x1_mask] = copy(x1[x1_mask])
    end

    if length(xT_mask) > 0
        l_traj[T][xT_mask] = copy(xT[xT_mask])
        u_traj[T][xT_mask] = copy(xT[xT_mask])
    end

    return l_traj, u_traj
end

function check_slack(Z, prob)
    model = prob.prob.model
    idx = prob.prob.idx
	T = prob.prob.T
    S̄ = [Z[idx.u[t]][model.idx_s][1] for t = 1:T-1]
    @show norm(S̄, Inf)
end

function state_to_configuration(X, nq = convert(Int, floor(length(X[1]) / 2)))
    T = length(X)
    [X[1][1:nq], [X[t][nq .+ (1:nq)] for t = 1:T]...]
end

function configuration_to_state(Q)
    T = length(Q)
    [[Q[t-1]; Q[t]] for t = 2:T]
end

function additive_noise_model(model::T) where T
	model_ft = typeof(model)([f == :d ? model.n : getfield(model,f)
	 	for f in fieldnames(typeof(model))]...)
	return model_ft
end

function get_time(u; idx = size(u[1], 1))
	T = length(u)
	h = [u[t][idx] for t = 1:T]
	tf = sum(h)
	t = range(0.0, stop = tf, length = T+1)

	return tf, t, h
end

# colors
goldenrod_color = RGBA(0.7686, 0.7647, 0.1333, 1.0)
red_color = RGBA(1.0, 0.0, 0.1529, 1.0)
darkslateblue_color = RGBA(0.1333, 0.2784, 0.5882, 1.0)


text:
using Literate, Documenter, ROSTransforms

# Literate.markdown("examples/SimpleExample.jl", "src/"; documenter=true)

Documenter.makedocs(
    sitename = "ROSTransforms.jl",
    source  = "src",
    build   = "build",
    clean   = true,
    doctest = true,
    # repo = "https://github.com/aabouman/ROSTransforms.jl",
    modules = [ROSTransforms],
    pages = ["Documentation" => "documentation.md"
             ]
)

deploydocs(;
    repo = "github.com/aabouman/ROSTransforms.jl",
    branch = "gh-pages",
    devurl = "docs",
)


text:
#Written By Johnathan Bizzano
module NeuralNet

        using ..Utils
        using ..Data
        using ..Functions

        include("Processor.jl")
        include("Initializer.jl")
        include("ActivationFun.jl")
        include("Layer/Layer.jl")
        include("Designer.jl")
        include("Model.jl")
end


text:
module AngleBetweenVectors

import Base: angle

import LinearAlgebra: norm

@inline unitize(p) = p ./ norm(p)

"""
    angle(point1::T, point2::T) where {T}

Accurately ascertains the undirected angle (0 <= radians < pi)
between two points given in N-dimensional Cartesian coordinates.

Prefer this to `acos` alternatives
- more reliably accurate
- more consistently stable

Suggested when any |coordinate| of either point may be outside 2^±20 or [1/1_000_000, 1_000_000].
Strongly recommended when any |coordinate| is outside 2^±24 or [1/16_000_000, 16_000_000].

If one of the points is at the origin, the result is zero.

You *must* define a tuple constructor `Tuple(x::YourPointType) = ...` if one does not already exist.
"""
function angle(point1::A, point2::A) where {N,T<:Real,NT<:NTuple{N,T}, V<:Vector{T}, A<:Union{NT,V}}
    unitpoint1 = unitize(point1)
    unitpoint2 = unitize(point2)

    y = unitpoint1 .- unitpoint2
    x = unitpoint1 .+ unitpoint2

    a = 2 * atan(norm(y) / norm(x))

    !(signbit(a) || signbit(T(pi) - a)) ? a : (signbit(a) ? zero(T) : T(pi))
end

@inline angle(point1::T, point2::T) where {T} = angle(Tuple(point1), Tuple(point2))

end # AngleBetweenVectors


text:
using Molly
using DelimitedFiles

data_dir = normpath(dirname(pathof(Molly)), "..", "data")
ff_dir = joinpath(data_dir, "force_fields")
openmm_dir = joinpath(data_dir, "openmm_6mrr")

ff = OpenMMForceField(joinpath.(ff_dir, ["ff99SBildn.xml", "tip3p_standard.xml", "his.xml"])...)
velocities = SVector{3}.(eachrow(readdlm(joinpath(openmm_dir, "velocities_300K.txt"))))u"nm * ps^-1"
s = System(joinpath(data_dir, "6mrr_equil.pdb"), ff; velocities=velocities)
simulator = VelocityVerlet(dt=0.0005u"ps", coupling=AndersenThermostat(300.0u"K", 1.0u"ps"))
n_steps = 500

parallel = false
simulate!(s, simulator, 5; parallel=parallel)
@time simulate!(s, simulator, n_steps; parallel=parallel)

parallel = true
simulate!(s, simulator, 5; parallel=parallel)
@time simulate!(s, simulator, n_steps; parallel=parallel)


text:
using CSV, DataFrames, Plots, StatsPlots, Missings
using Flux
using Flux: relu
using Flux: mse, crossentropy
using Flux.Optimise: update!
using Zygote
using MLDataUtils
using StatsBase
using Random;
Random.seed!(0);
using ColorSchemes
using Base.Iterators
using BSON: @save, @load

data = CSV.read("julia/cleaned.csv");

function df_to_data(dataframe)
    # Get the numeric data
    extracted = Matrix(dataframe[[:Lat, :Long, :PosEncoding, :HotDogEncoding]])
    # calculate cartesian coordinates from lat/long
    x = cos.(extracted[:, 1]) .* cos.(extracted[:, 2])
    y = cos.(extracted[:, 1]) .* sin.(extracted[:, 2])
    z = sin.(extracted[:, 1])

    X = hcat(x, y, z, extracted[:, 3], dataframe[:, :RelativeTime])
    dt = fit(UnitRangeTransform, X, dims = 1)
    norm_X = StatsBase.transform(dt, X)
    norm_X = transpose(norm_X)
    Y = convert(Array{Float32}, extracted[:, 4])

    return X, Y
end

X, Y = df_to_data(data);
X = transpose(X);

# Define a logistic regression model; output is sigmoid
embedder = Chain(
    Dense(5, 12, tanh),
    # Dense(16, 8, tanh),
    Dense(12, 2, tanh),
)
predictor = Dense(2, 1, σ)
# combine the models together
dumber_model = Chain(embedder, predictor)

# if isfile("julia/trained_model.bson") == false
calc_loss(x) = mse(dumber_model(x), transpose(Y[:, :]))

dumber_model(X)

# Setup training logistics
θ = Flux.params(dumber_model);
optimizer = NADAM(1e-3);

if isfile("julia/trained_model.bson")
    for i in range(1, stop = 50000)
        grads = gradient(() -> calc_loss(X), θ)
        for p in θ
            update!(optimizer, p, grads[p])
        end
        loss = calc_loss(X)
        if mod(i, 1000) == 0
            println("Current loss is $loss")
        end
    end
else
    @load "julia/trained_model.bson" θ
    Flux.loadparams!(dumber_model, θ)
end

# Evaluate feature importance with input gradients; throw
# noise at the model and see how each affects the final predicted
# sigmoid value
durr = rand(5, 1000)

gs = gradient(Flux.params(durr)) do
    sum(dumber_model(durr))
end

# normalize by gradient magnitude
inp_grads = abs.(gs[durr])
inp_grads = mean(inp_grads, dims=2)
inp_grads ./= maximum(inp_grads);
# Make a plot of the input gradients
gr(size = (600, 500))
chart = bar(
    inp_grads,
    color = "#b85231",
    leg = false,
    lw = 1.0,
    bar_width = 0.3,
    background_color = "#f9efde",
    titlefont = font(14, "Open Sans"),
)
title!("What decides what a sandwich?")
ylabel!("Importance", font = font(12, "Open Sans"))
xlabel!("Feature", font = font(12, "Open Sans"))
xticks!([1, 2, 3, 4, 5], ["X", "Y", "Z", "Career stage", "Registration time"])
savefig("julia/feature_importance.png")

# Visualizing the embeddings.

# The second to last layer is of 2D, with the idea
# that this corresponds to a 2D encoding of the "dataset"
# that we can use to evaluate a continuous decision boundary
embeddings = embedder(X)
pred_Y = dumber_model(X)

gr(size = (900, 500))

plot(
    scatter(
        embeddings[1, :],
        embeddings[2, :],
        zcolor = Y[:],
        leg = false,
        title = "Truth",
    ),
    scatter(
        embeddings[1, :],
        embeddings[2, :],
        zcolor = pred_Y[:],
        leg = false,
        title = "Predicted",
    ),
)
xlabel!("X")
ylabel!("Y")

groups = groupby(data, :PosEncoding)

# See how the postdocs look
sliced_X, sliced_Y = df_to_data(groups[3])

slice_embeddings = embedder(transpose(sliced_X))

gr(size=(600, 500))
plot(scatter(
    slice_embeddings[1, :],
    slice_embeddings[2, :],
    zcolor = sliced_Y[:],
    leg = false,
    title = "Truth",
    alpha = 1.0,
),
# scatter(embeddings[1,:], embeddings[2,:], zcolor=pred_Y[:], leg=false, title="Predicted")
)
scatter!(embeddings[1, :], embeddings[2, :], zcolor = Y[:], leg = false, alpha = 0.1)
xlabel!("X")
ylabel!("Y")

function generate_grid(start, stop, length)
    linear_array = range(start, stop = stop, length = length)
    grid = zeros(2, length * length)
    index = 1
    for a in linear_array
        for b in linear_array
            grid[1, index] = a
            grid[2, index] = b
            index = index + 1
        end
    end
    return grid
end

grid_array = generate_grid(-1.0, 1.0, 100)

herp = range(-1.0, stop = 1.0, length = 100);
continuous = predictor(grid_array);

embeddings = embedder(X)
pred_Y = dumber_model(X)

contourf(
    herp,
    herp,
    continuous;
    levels = collect(range(0, stop = 1, length = 10)),
    c = :Spectral_6,
    background_color = "#f9efde",
)
scatter!(
    embeddings[1, :],
    embeddings[2, :],
    zcolor = Y[:],
    lw = 1.0,
    markerstrokecolor = "white",
    leg = false,
    # c = :RdBu_10,
)
xlabel!("X")
ylabel!("Y")
title!("Wrong Decision Boundary")
savefig("julia/decision_boundary.png")

plots = [];
for (index, group) in enumerate(groupby(data, :PosEncoding))
    name = unique(group.Position)[1]
    temp_X, temp_Y = df_to_data(group)
    temp_embedding = embedder(transpose(temp_X))
    temp_plot = contourf(
        herp,
        herp,
        continuous;
        levels = collect(range(0, stop = 1, length = 10)),
        c = :Spectral_6,
        background_color = "#f9efde",
        alpha=0.8
    )
    scatter!(
        temp_embedding[1, :],
        temp_embedding[2, :],
        zcolor = temp_Y[:],
        lw = 1.0,
        markerstrokecolor = "white",
        leg = false
        # c = :RdBu_10,
    )
    scatter!(
        embeddings[1, :],
        embeddings[2, :],
        zcolor = Y[:],
        leg = false,
        alpha=0.1
    ),
    title!(name)
    push!(plots, temp_plot)
    savefig("julia/$name-contour.png")
end

plots[4]
# dump the weights to disk
# @save "julia/trained_model.bson" θ

sum(data.HotDogEncoding)


text:
register(DataDep(
    "pudl-raw-eia923",
    """
    Dataset: PUDL Raw EIA Form 923
    Website: https://doi.org/10.5281/zenodo.4127040
    Author: Gosnell, Christina, Selvans, Zane
    Date of Publication: October 24, 2020
    License:

    Raw US Energy Information Administration (EIA) Form 923 data, archived from https://www.eia.gov/electricity/data/eia923/

    EIA Form 923 collects detailed monthly and annual electric power data on electricity generation, fuel consumption, fossil fuel stocks, and receipts at the power plant and prime mover level.

    $PUDL_DESCRIPTION
    """,
    [
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/datapackage.json",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2001.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2002.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2003.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2004.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2005.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2006.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2007.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2008.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2009.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2010.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2011.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2012.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2013.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2014.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2015.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2016.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2017.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2018.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2019.zip",
        "https://zenodo.org/api/files/7a446075-263a-4188-9deb-180dab065609/eia923-2020.zip",
    ],
    "e41da21e10cb8bf22024410a26e104a077c6de14f1159c191dc2feb46b77ab80",
))


text:
using Clang.cindex
using Clang.wrap_c
using Compat

AF_INCLUDE = "/usr/include/af"

af_header = ["/usr/include/arrayfire.h"]

# Set up include paths
clang_includes = [".", AF_INCLUDE]

# Callback to test if a header should actually be wrapped (for exclusion)
function wrap_header(top_hdr, cursor_header)
    return startswith(dirname(cursor_header), AF_INCLUDE)
end

lib_file(hdr) = "af_lib"
output_file(hdr) = "../src/wrap.jl"

function wrap_cursor(name, cursor)
    if startswith(name, "AF") || startswith(name, "DEPRECATED") ||
        name == "bool" || isempty(name) || startswith(name, "SIZE_T") || endswith(name, ".h")
        return false
    end
    true
end

function return_val(typ, arg)
    if typ.args[2] == :af_array
        Expr(:call, :AFArray!, Expr(:ref, arg))
    else
        Expr(:ref, arg)
    end
end

function return_val(typ, arg, expr)
    if typ.args[2] == :af_array
        Expr(:call, expr, Expr(:ref, arg))
    else
        Expr(:ref, arg)
    end
end

const renames = Dict("sign" => "signbit", "product" => "prod", "init" => "afinit", "info" => "afinfo",
                     "copy_array" => "copy", "get_version" => "afversion",
                     "min" => "minimum", "max" => "maximum", "any_true" => "any", "all_true" => "all",
                     "select_scalar_l" => "select", "select_scalar_r" => "select", "replace_scalar" => "replace",
                     "is_sparse" => "issparse", "sparse_to_dense" => "full",
		     "cplx" => "complex", "conjg" => "conj", "diag_extract" => "diag",
                     "diag_create" => "diagm")

const ignore = Set(["example_function", "create_array", "retain_array", "get_data_ref_count", "info_string",
                    "device_info", "alloc_host", "free_host", "alloc_pinned", "free_pinned",
                    "get_type", "get_numdims", "join_many", "eval_multiple", "get_size_of", "cast",
                    "constant", "constant_complex", "constant_long", "constant_ulong", "transpose",
                    "release_array", "flat", "select", "get_last_error", "err_to_string", "moddims",
                    "randu", "randn", "set_seq_indexer", "release_indexers", "assign_gen", "join", "svd",
                    "sort", "sort_index", "mean", "mean_weighted", "var", "var_weighted",
                    "stdev", "median", "set_array_indexer", "set_seq_param_indexer", "clamp", "where",
                    "get_scalar", "cholesky", "cplx2", "grid", "fir", "iir", "eval"])

const recast = Dict(:Cint => :Integer, :UInt32 => :Integer, :Cdouble => :Real)

const booleans1 = Set(["iszero", "isinf", "isnan", "not"])
const booleans2 = Set(["lt", "gt", "le", "ge", "eq", "neq", "and", "or"])
const maths     = Set(["add", "sub", "mul", "div", "rem", "mod", "atan2", "root", "pow", "dot",
                       "minof", "maxof", "hypot", "cplx2", "matmul",
                       "bitshiftl", "bitshiftr", "bitxor", "bitand", "bitor"])
const floats    = Set(["signbit"])
const c2rs      = Set(["fft_c2r", "fft2_c2r", "fft3_c2r", "real", "imag"])
const complexes = Set(["fft_r2c", "fft2_r2c", "fft3_r2c", "complex"])
const unknowns = Set(["diag", "diagm"])

const exports = []

function rewrite(line::Expr)
    if line.head == :function
        hdr = line.args[1].args
        name = replace("$(hdr[1])", "af_", "", 1)
        name = get(renames, name, name)
        if in(name, ignore) || contains(name, "fft")
	    return []
	end
        hdr[1] = Symbol(name)
        args = hdr[2:end]
        body = line.args[2].args
        types = body[1].args[3].args
        vals = view(body[1].args, 4:length(body[1].args))

        num_input_arrays = 0
        num_output_arrays = 0

        for k in 1:length(args)
            if isa(args[k], Expr)
                if args[k].args[2] == :af_array
                    args[k].args[2] = :AFArray
                    vals[k] = Expr(:., vals[k], QuoteNode(:arr))
                    num_input_arrays += 1
                elseif args[k].args[2] == :af_dtype
                    args[k].args[2] = :Type
                    vals[k] = Expr(:call, :af_type, vals[k])
                elseif haskey(recast, args[k].args[2])
                    key = args[k].args[2]
                    val = recast[key]
                    args[k].args[2] = val
                    if vals[k] == :dim
                        vals[k] = Expr(:call, key, Expr(:call, :-, vals[k], 1))
                    else
                        vals[k] = Expr(:call, key, vals[k])
                    end
                end
            end
        end

        ret_type = body[1].args[2]
        if ret_type == :af_err
            body[1] = Expr(:call, :_error, body[1])
        end
        num_out = 0
        for k = 1:length(types)
            t = types[k]
            if isa(t, Expr) && t.args[1] == :Ptr && t.args[2] == :af_array
                num_output_arrays += 1
            end
            if isa(t, Expr) && t.args[1] == :Ptr && t.args[2] != :Void
                num_out = k
            else
                break
            end
        end
        if num_out > 0
            if name in floats
                hdr[1] = Expr(:curly, hdr[1], :T, :N)
                args[2].args[2] = Expr(:curly, :AFArray, :T, :N)
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, :Float32, :N)))
            elseif name in complexes
                hdr[1] = Expr(:curly, hdr[1], :T, :N)
                args[2].args[2] = Expr(:curly, :AFArray, :T, :N)
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, Expr(:curly, :Complex, :T), :N)))
            elseif name in c2rs
                hdr[1] = Expr(:curly, hdr[1], :T, :N)
                args[2].args[2] = Expr(:curly, :AFArray, Expr(:curly, :Complex, :T), :N)
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, :T, :N)))
            elseif name in booleans1
                hdr[1] = Expr(:curly, hdr[1], :T, :N)
                args[2].args[2] = Expr(:curly, :AFArray, :T, :N)
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, :Bool, :N)))
            elseif name in booleans2
                hdr[1] = Expr(:curly, hdr[1], :T1, :N1, :T2, :N2)
                args[2].args[2] = Expr(:curly, :AFArray, :T1, :N1)
                args[3].args[2] = Expr(:curly, :AFArray, :T2, :N2)
                n = Expr(:call, :batched, :N1, :N2)
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, :Bool, n)))
            elseif name in maths
                hdr[1] = Expr(:curly, hdr[1], :T1, :N1, :T2, :N2)
                args[2].args[2] = Expr(:curly, :AFArray, :T1, :N1)
                args[3].args[2] = Expr(:curly, :AFArray, :T2, :N2)
                t = Expr(:call, :typed, :T1, :T2)
                n = Expr(:call, :batched, :N1, :N2)
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, t, n)))
            elseif num_input_arrays == 1 && num_output_arrays == 1 && num_out == 1 &&
                !contains(name, "_sparse") && !contains(name, "_true") && !(name in unknowns)
                hdr[1] = Expr(:curly, hdr[1], :T, :N)
                for k = 1:length(args)
                    if isa(args[k], Expr) && args[k].args[2] == :AFArray
                        args[k].args[2] = Expr(:curly, :AFArray, :T, :N)
                    end
                end
                push!(body, return_val(types[1], args[1], Expr(:curly, :AFArray, :T, :N)))
            elseif num_out == 1
                if num_output_arrays > 0
                    @printf("No type inference for %s\n", name)
                end
                push!(body, return_val(types[1], args[1]))
            else
                if num_output_arrays > 0
                    @printf("No type inference for %s\n", name)
                end
                push!(body, Expr(:tuple, map(k->return_val(types[k], args[k]), 1:num_out)...))
            end
            for k in num_out:-1:1
                deleteat!(hdr, 2)
                t = Expr(:curly, :RefValue, types[k].args[2])
                c = types[k].args[2] == :Cstring ? Expr(:call, t) : Expr(:call, t, 0)
                unshift!(body, Expr(:(=), args[k], c))
            end
        end
        push!(exports, Symbol(name))
    elseif line.head == :const || line.head == :typealias
        if isa(line.args[1], Symbol)
            push!(exports, line.args[1])
        else
            push!(exports, line.args[1].args[1])
        end
        if line.head == :typealias
            line.head = :(=)
            line = Expr(:const, line)
        end
    end
    line
end

rewrite(line) = line

function rewriter(input)
    out = vcat(map(rewrite, input)...)
    exp = []
    te = []
    for ex in sort(exports)
        if length(string(Expr(:export, te...))) > 100
            push!(exp, Expr(:export, te...))
            empty!(te)
        end
        push!(te, ex)
    end
    if !isempty(te)
        push!(exp, Expr(:export, te...))
    end
    empty!(exports)
    vcat(exp, out)
end

const wc = wrap_c.init(;
                       headers             = af_header,
                       common_file         = "../src/common.jl",
                       clang_includes      = clang_includes,
                       header_wrapped      = wrap_header,
                       header_library      = lib_file,
                       header_outputfile   = output_file,
                       rewriter            = rewriter,
                       cursor_wrapped      = wrap_cursor)

run(wc)


text:
## OUTPUT FUNCTION
# Duty Cycle for 33250A
set_duty(obj::INSTR{:Agilent33250A},ch="1",func="square",duty=0.01) = write(obj,"SOURce$ch:FUNC:square:DCYC $duty")

## BURST
# Burst on
set_burst_stat(obj::T;ch="1",st="off") where {T<:F332x1} = write(obj,"BURS:STAT $st")
# Burst mode
set_burst_mode(obj::T;ch="1",md="triggered") where {T<:F332x1} = write(obj,"BURS:MODE $md")
# Burst polarity
set_burst_pol(obj::T;ch="1",pol="normal") where {T<:F332x1} = write(obj,"BURS:GATE:POL $pol")
# Burst ncycles
set_burst_ncyc(obj::T;ch="1",nc=1) where {T<:F332x1} = write(obj,"BURS:NCYC $nc")
# Burst period
set_burst_per(obj::T;ch="1",per=0.01) where {T<:F332x1} = write(obj,"BURS:INT:PER $per")
# Burst phase
set_burst_phase(obj::T;ch="1",unit="deg",phase=0) where {T<:F332x1} = write(obj,"UNIT:ANGL $unit\n;BURS:PHAS $phase")

## TRIGGER
# Output trigger source
set_out_trig_sour(obj::T;sour=1) where {T<:F332x1} = write(obj,"OUTP:TRIG:SOUR ch$sour")


text:
@test ip"127.0.0.1" == IPv4(127,0,0,1)
@test ip"192.0" == IPv4(192,0,0,0)
@test ip"192.0xFFF" == IPv4(192,0,15,255)
@test ip"192.0xFFFF" == IPv4(192,0,255,255)
@test ip"192.0xFFFFF" == IPv4(192,15,255,255)
@test ip"192.0xFFFFFF" == IPv4(192,255,255,255)
@test_throws Base.parseipv4("192.0xFFFFFFF")
@test ip"022.0.0.1" == IPv4(18,0,0,1)

@test_throws Base.parseipv4("192.0xFFFFFFFFF")
@test_throws Base.parseipv4("192.")

@test ip"::1" == IPv6(1)
@test ip"2605:2700:0:3::4713:93e3" == IPv6(parseint(Uint128,"260527000000000300000000471393e3",16))

@test ip"2001:db8:0:0:0:0:2:1" == ip"2001:db8::2:1" == ip"2001:db8::0:2:1"

@test ip"0:0:0:0:0:ffff:127.0.0.1" == IPv6(0xffff7f000001)

# RFC 5952 Compliance

@test repr(ip"2001:db8:0:0:0:0:2:1") == "ip\"2001:db8::2:1\""
@test repr(ip"2001:0db8::0001") == "ip\"2001:db8::1\""
@test repr(ip"2001:db8::1:1:1:1:1") == "ip\"2001:db8:0:1:1:1:1:1\""
@test repr(ip"2001:db8:0:0:1:0:0:1") == "ip\"2001:db8::1:0:0:1\""
@test repr(ip"2001:0:0:1:0:0:0:1") == "ip\"2001:0:0:1::1\""

c = Base.Condition()
@async begin
	s = listen(2134)
	Base.notify(c)
	sock = accept(s)
	write(sock,"Hello World\n")
	close(s)
	close(sock)
end
wait(c)
@test readall(connect(2134)) == "Hello World\n"

socketname = @windows ? "\\\\.\\pipe\\uv-test" : "testsocket"
@unix_only isfile(socketname) && Base.FS.unlink(socketname)
@async begin
	s = listen(socketname)
	Base.notify(c)
	sock = accept(s)
	write(sock,"Hello World\n")
	close(s)
	close(sock)
end
wait(c)
@test readall(connect(socketname)) == "Hello World\n"

try
    getaddrinfo(".invalid")
catch e
    @test typeof(e) == Base.UVError # E.g. not method error
end

try
    # This should be an invalid port
    connect("localhost",21452)
    @test false
catch e
    @test typeof(e) == Base.UVError
end

server = listen(2134)
@async @test_throws accept(server);
sleep(0.1)
close(server)

server = listen(2134)
@async connect("localhost",2134)
s1 = accept(server)
@test_throws accept(server,s1)
close(server)

try
    connect(".invalid",80)
catch e
    @test typeof(e) == Base.UVError
end

a = UdpSocket()
b = UdpSocket()
bind(a,ip"127.0.0.1",2134)
bind(b,ip"127.0.0.1",2135)

c = Condition()
@async begin
    @test bytestring(recv(a)) == "Hello World"
    notify(c)
end
send(b,ip"127.0.0.1",2134,"Hello World")
wait(c)

@test_throws bind(UdpSocket(),2134)

close(a)
close(b)

max_stars_repo_name: alyst/julia

max_issues_repo_name: alyst/julia

max_forks_repo_name: alyst/julia


text:
function _precompile_()

end


text:
#
# This file is part of the Simulate.jl Julia package, MIT license
#
# Paul Bayer, 2020
#
# This is a Julia package for discrete event simulation
#

"""
```
event!([clk::AbstractClock], ex::Action, t::Number;
       scope::Module=Main, cycle::Number=0.0, spawn=false)::Float64
event!([clk::AbstractClock], ex::Action, T::Timing, t::Number; kw...)
```
Schedule an event for a given simulation time.

# Arguments
- `clk::AbstractClock`: it not supplied, the event is scheduled to 𝐶,
- `ex::Action`: an expression or Fun or a tuple of them,
- `T::Timing`: a timing, one of `at`, `after` or `every`,
- `t::Real` or `t::Time`: simulation time, if t < clk.time set t = clk.time,

# Keyword arguments
- `scope::Module=Main`: scope for expressions to be evaluated in,
- `cycle::Float64=0.0`: repeat cycle time for an event,
- `spawn::Bool=false`: if true, spawn the event at other available threads,
- `cid::Int=clk.id`: if cid ≠ clk.id, assign the event to the parallel clock
    with id == cid. This overrides `spawn`,
- `sync::Bool=false`: if true, force a synchronization of all parallel clocks
    before executing the event.

# returns
Scheduled internal simulation time (unitless) for that event.
May return a time `> t` from repeated applications of `nextfloat(t)`
if there are events scheduled for `t`.

# Examples
```jldoctest
julia> using Simulate, Unitful

julia> import Unitful: s, minute, hr

julia> myfunc(a, b) = a+b
myfunc (generic function with 1 method)

julia> event!(𝐶, Fun(myfunc, 1, 2), 1) # a 1st event to 1
1.0
julia> event!(𝐶, Fun(myfunc, 2, 3), 1) #  a 2nd event to the same time
1.0000000000000002

julia> event!(𝐶, Fun(myfunc, 3, 4), 1s)
Warning: clock has no time unit, ignoring units
1.0000000000000004

julia> setUnit!(𝐶, s)
0.0 s

julia> event!(𝐶, Fun(myfunc, 4, 5), 1minute)
60.0

julia> event!(Fun(myfunc, 5, 6), after, 1hr)
3600.0
```
"""
function event!(clk::AbstractClock, ex::Action, t::Number;
                scope::Module=Main, cycle::Number=0.0,
                spawn::Bool=false, cid::Int=clk.id, sync::Bool=false) :: Float64
    (t isa Unitful.Time) && (t = tadjust(clk, t))
    (cycle isa Unitful.Time) && (cycle = tadjust(clk, cycle))
    (t < clk.time) && (t = clk.time)

    assign(clk, DiscreteEvent(ex, scope, t, cycle), spawn ? spawnid(clk) : 0)
end
function event!(clk::AbstractClock, ex::Action, T::Timing, t::Number;
                scope::Module=Main,
                spawn::Bool=false, cid::Int=0, sync::Bool=false) :: Float64
    (t isa Unitful.Time) && (t = tadjust(clk, t))
    if T == after
        event!(clk, ex, t+clk.time, scope=scope, spawn=spawn)
    elseif T == every
        event!(clk, ex, clk.time, scope=scope, cycle=t, spawn=spawn)
    else
        event!(clk, ex, t, scope=scope, spawn=spawn)
    end
end
event!(ex::Action, t::Number; kw...) = event!(𝐶, ex, t; kw...)
event!(ex::Action, T::Timing, t::Number; kw...) = event!(𝐶, ex, T, t; kw...)

"""
```
event!([clk::AbstractClock], ex::Action, cond::Action; scope::Module=Main)::Float64
```
Schedule a conditional event.

It is executed immediately if the conditions are met, else the condition is
checked at each clock tick Δt. A conditional event is triggered only once. After
that it is removed from the clock. If no sampling rate Δt is setup, a default
sampling rate is setup depending on the scale of the remaining simulation time
``Δt = scale(t_r)/100`` or ``0.01`` if ``t_r = 0``.

# Arguments
- `clk::AbstractClock`: if no clock is supplied, the event is scheduled to 𝐶,
- `ex::Union{SimExpr, Tuple{SimExpr}}`: an expression or Fun or a tuple of them,
- `cond::Union{SimExpr, Tuple{SimExpr}}`: a condition is an expression or Fun
    or a tuple of them. It is true only if all expressions or Funs
    therein return true,
- `scope::Module=Main`: scope for the expressions to be evaluated

# returns
current simulation time `tau(clk)`.

# Examples
```jldoctest
julia> using Simulate

julia> c = Clock()   # create a new clock
Clock thread 1 (+ 0 ac): state=Simulate.Undefined(), t=0.0 , Δt=0.0 , prc:0
  scheduled ev:0, cev:0, sampl:0

julia> event!(c, Fun((x)->println(tau(x), ": now I'm triggered"), c), Fun(>=, Fun(tau, c), 5))
0.0

# julia> c              # a conditional event turns sampling on  ⬇
Clock thread 1 (+ 0 ac): state=Simulate.Undefined(), t=0.0 , Δt=0.01 , prc:0
  scheduled ev:0, cev:1, sampl:0

julia> run!(c, 10)   # sampling is not exact, so it takes 501 sample steps to fire the event
5.009999999999938: now I'm triggered
"run! finished with 0 clock events, 501 sample steps, simulation time: 10.0"

julia> c   # after the event sampling is again switched off ⬇
Clock thread 1 (+ 0 ac): state=Simulate.Idle(), t=10.0 , Δt=0.0 , prc:0
  scheduled ev:0, cev:0, sampl:0
```
"""
function event!(clk::AbstractClock, ex::Action, cond::Action; scope::Module=Main, spawn=false)
    if busy(clk) && all(evExec(cond))   # all conditions met
        evExec(ex)                      # execute immediately
    else
        assign(clk, DiscreteCond(cond, ex, scope), spawn ? spawnid(clk) : 0)
    end
    return tau(clk)
end
event!( ex::Action, cond::Action; kw...) = event!(𝐶, ex, cond; kw...)

"""
    spawnid(clk::Clock) :: Int

Return a random number out of the thread ids of all available parallel clocks.
This is used for `spawn`ing tasks or events to them.

!!! note
    This function may be used for workload balancing between threads
    in the future.
"""
function spawnid(clk::Clock) :: Int
    if isempty(clk.ac)
        return 1
    else
        return rand(rng, (1, (i.id for i in clk.ac)...))
    end
end

"""
    scale(n::Number)::Float64

calculate the scale from a given number
"""
function scale(n::Number)::Float64
    if n > 0
        i = 1.0
        while !(10^i ≤ n < 10^(i+1))
            n < 10^i ? i -= 1 : i += 1
        end
        return 10^i
    else
        return 1
    end
end

# ---------------------------------------------------------------
# assign and register events and samples to a clock
# ---------------------------------------------------------------

"""
    assign(c::AbstractClock, ev::AbstractEvent, id::Int=c.id)

Assign an abstract event to a clock.

There are several ways to do it:
1. assign it directly to a clock or an active clock or
2. assign it via a clock via a clock to another one, if given a different id.
    In this case the event is sent over the channel to the target clock.
    - The master clock (id=0) can directly send to an active clock.
    - An active clock can send directly to master.
    - An active clock can only send via master to another active clock.
"""
function assign(c::AbstractClock, ev::AbstractEvent, id::Int=c.id)
    if id == c.id
        register!(c, ev)       # 1: register it yourself
    else
        register(c, ev, id)    # 2: register it to another clock
    end
end

"""
```
register!(c::Clock, ev::DiscreteEvent)
register!(c::Clock, cond::DiscreteCond)
register!(c::Clock, sp::Sample)
register!(ac::ActiveClock, ev::AbstractEvent)
```
Register a concrete event directly to a clock.
"""
function register!(c::Clock, ev::DiscreteEvent)
    t = ev.t
    while any(i->i==t, values(c.sc.events)) # in case an event at that time exists
        t = nextfloat(float(t))                  # increment scheduled time
    end
    return c.sc.events[ev] = t
end
function register!(c::Clock, cond::DiscreteCond)
    (c.Δt == 0) && (c.Δt = scale(c.end_time - c.time)/100)
    push!(c.sc.cevents, cond)
end
register!(c::Clock, sp::Sample) = push!(c.sc.samples, sp)
register!(ac::ActiveClock, ev::AbstractEvent) = register!(ac.clock, ev)

"""
```
register(c::Clock, ev::AbstractEvent, id::Int)
register(ac::ActiveClock, ev::AbstractEvent, id::Int)
```
Register an event to another clock via a channel.

# Arguments
- `c::Clock`: a master clock can forward events to active clocks,
- `ac::ActiveClock`: active clocks can forward events only through master, he then
    does the distribution for them,
- `ev::AbstractEvent`: the event to register,
- `id::Int`: the id of the clock it should get registered to.
"""
function register(c::Clock, ev::AbstractEvent, id::Int)
    if id > 0 && threadid() == 1       # only master can forward events
        put!(c.ac[id].forth, Register(ev))
    else
        register!(c, ev)               # otherwise take it yourself
    end
end
function register(ac::ActiveClock, ev::AbstractEvent, id::Int)
    if ac.id == id                     # if id is your id
        register!(ac.clock, ev)        # take it yourself
    elseif ac.master[].state == Busy() # if master handles its channels ↯↯↯↯↯ this has yet to be changed
        put!(ac.back, Forward(ev, id)) # put it over the channel
    else
        register(ac.master[], ev, id)  # otherwise call him directly
    end
end


text:
module SequGen

export Sequence
export sample

"""
    Sequence

Base type for synthetic sequences.
"""
abstract type Sequence end

"""
    x = sample(sequence, range)

Compute a sample in the given `range`.
"""
function sample end

include("composite/operation.jl")
include("composite/transformation.jl")
include("deterministic/constant.jl")
include("deterministic/trend.jl")
include("sthocastic/gaussian-noise.jl")
include("waves/sine-wave.jl")

end


text:
# ---
# title: 378. Kth Smallest Element in a Sorted Matrix
# id: problem378
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Binary Search, Heap
# link: <https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix/description/>
# hidden: true
# ---
#
# Given a _n_ x _n_ matrix where each of the rows and columns are sorted in
# ascending order, find the kth smallest element in the matrix.
#
# Note that it is the kth smallest element in the sorted order, not the kth
# distinct element.
#
# **Example:**
#
#
#
#     matrix = [
#        [ 1,  5,  9],
#        [10, 11, 13],
#        [12, 13, 15]
#     ],
#     k = 8,
#
#     return 13.
#
#
# **Note:**
# You may assume k is always valid, 1 ≤ k ≤ n2.
#
#
## @lc code=start
using LeetCode

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


text:
eye(A) = A^0
isnormal(A) = size(A,1) == size(A,2) && A'*A == A*A'
isunitary(A) = size(A,1) == size(A,2) && A'*A == eye(A)

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# handle grouping by DataFrame column
function Plots.extractGroupArgs(group::Symbol, df::AbstractDataFrame, args...)
    Plots.extractGroupArgs(collect(df[group]))
end

# if it's one symbol, set the guide and return the column
function handle_dfs(df::AbstractDataFrame, d::KW, letter, sym::Symbol)
    get!(d, Symbol(letter * "guide"), string(sym))
    processDFarg(df, sym)
end

# if it's one symbol, set the guide and return the column
function handle_dfs(df::AbstractDataFrame, d::KW, letter, expr::Expr)
    get!(d, Symbol(letter * "guide"), replace(string(expr), ":", ""))
    processDFarg(df, expr)
end

# if it's an array of symbols, set the labels and return a Vector{Any} of columns
function handle_dfs(df::AbstractDataFrame, d::KW, letter, syms::AbstractArray)
    get!(d, :label, reshape(Symbol.(syms), 1, length(syms)))
    processDFarg(df, syms)
end

# for anything else, no-op
function handle_dfs(df::AbstractDataFrame, d::KW, letter, anything)
    anything
end

hcat_extend(a::AbstractArray, b::Number) = hcat(a, fill(b, size(a,1)))
hcat_extend(a::Number, b::AbstractArray) = hcat(fill(a, size(b,1)), b)
hcat_extend(a, b) = hcat(a, b)

# allows the passing of expressions including DataFrame columns as symbols
processDFarg(df::AbstractDataFrame, sym::Symbol) = haskey(df, sym) ? collect(df[sym]) : sym
function processDFarg(df::AbstractDataFrame, syms::AbstractMatrix)
    ret = processDFcol(df, syms[:,1])
    for i in 2:size(syms, 2)
        tmp = processDFcol(df, syms[:,i])
        ret = hcat_extend(ret, tmp)
    end
    ret
end

processDFarg(df::AbstractDataFrame, expr::Expr) = eval(processDFsym(df, expr))
processDFarg(df::AbstractDataFrame, anything) = anything

processDFcol{T <: Union{Expr, Symbol}}(df, col::AbstractArray{T}) = processDFarg(df, col[1])
processDFcol(df, anything) = anything

# the processDFsym! functions work with expressions and pass results to processDFarg for final eval - this to allow recursion without eval
processDFsym(df::AbstractDataFrame, s::Symbol) = haskey(df,s) ? :(collect($(df)[$s])) : :($s)
processDFsym(df::AbstractDataFrame, s::QuoteNode) = haskey(df,s.value) ? :(collect($(df)[$s])) : :($s)
processDFsym(df::AbstractDataFrame, s) = :($s)

function processDFsym(df::AbstractDataFrame, expr::Expr)
    arg = map(x->processDFsym(df,x), expr.args)
    ret = copy(expr)
    ret.args = arg
    return ret
end

# if a DataFrame is the first arg, lets swap symbols out for columns
@recipe function f(df::AbstractDataFrame, args...)
    # if any of these attributes are symbols, swap out for the df column
    for k in (:fillrange, :line_z, :marker_z, :markersize, :ribbon, :weights, :xerror, :yerror, :hover, :bar_width, :fill_z, :fillalpha, :fillcolor, :linealpha, :linecolor, :linestyle, :linewidth, :markeralpha, :markercolor, :markershape, :markerstrokecolor, :series_annotations, :seriesalpha, :seriescolor)
        if haskey(d, k)
            d[k] = processDFarg(df, d[k])
        end
    end

    # return a list of new arguments
    tuple(Any[handle_dfs(df, d, (i==1 ? "x" : i==2 ? "y" : "z"), arg) for (i,arg) in enumerate(args)]...)
end


text:
#!/home/dss/git/julia/julia
using DCEMRI
results = fitdata(datafile="../out/dce_recon.mat",workers=4)

max_stars_repo_name: chixindebaoyu/qetsr

max_issues_repo_name: chixindebaoyu/qetsr

max_forks_repo_name: chixindebaoyu/qetsr


text:
module GalacticNonconvex

using Nonconvex, GalacticOptim, GalacticOptim.SciMLBase, ChainRulesCore

include("nonconvex_bayesian.jl")
include("nonconvex_pavito.jl")
include("nonconvex_juniper.jl")
include("nonconvex_ipopt.jl")
include("nonconvex_nlopt.jl")
include("nonconvex_mma.jl")
# include("nonconvex_multistart.jl")
include("nonconvex_percival.jl")
include("nonconvex_search.jl")

struct NonconvexADWrapper{F,P}
    f::F
    prob::P
end

(f::NonconvexADWrapper)(θ) = begin
    x = f.f.f(θ, f.prob.p)
    return first(x)
end

## Create ChainRule
function ChainRulesCore.rrule(f::NonconvexADWrapper, θ::AbstractVector)
    val = f.f.f(θ, f.prob.p)
    G = similar(θ)
    f.f.grad(G, θ)
    val, Δ -> (ChainRulesCore.NoTangent(), Δ * G)
end

function separate_kwargs(orig_kwargs)
    sub_options = :sub_options .∈ Ref(keys(orig_kwargs)) ? orig_kwargs[:sub_options] : nothing
    convergence_criteria = :convcriteria .∈ Ref(keys(orig_kwargs)) ? orig_kwargs[:convcriteria] : nothing
    opt_kwargs_names = keys(orig_kwargs)[[(keys(orig_kwargs) .∉ Ref([:sub_options, :convcriteria]))...]]
    opt_kwargs = (; zip(opt_kwargs_names, [orig_kwargs[j] for j in opt_kwargs_names])...)

    return opt_kwargs, sub_options, convergence_criteria
end

function __map_optimizer_args(prob::OptimizationProblem, opt::Nonconvex.NonconvexCore.AbstractOptimizer;
    callback=nothing,
    maxiters::Union{Number,Nothing}=nothing,
    maxtime::Union{Number,Nothing}=nothing,
    abstol::Union{Number,Nothing}=nothing,
    reltol::Union{Number,Nothing}=nothing,
    integer=nothing,
    kwargs...)

    opt_kwargs, sub_options, convergence_criteria = separate_kwargs(kwargs)
    opt_kwargs = convert_common_kwargs(opt, opt_kwargs, callback=callback, maxiters=maxiters, maxtime=maxtime, abstol=abstol, reltol=reltol)
    mapped_args = _create_options(opt, opt_kwargs=opt_kwargs, sub_options=sub_options, convergence_criteria=convergence_criteria)

    integer = isnothing(integer) ? fill(false, length(prob.u0)) : integer

    return mapped_args, integer
end

function SciMLBase.__solve(prob::OptimizationProblem, opt::Nonconvex.NonconvexCore.AbstractOptimizer;
    callback=nothing,
    maxiters::Union{Number,Nothing}=nothing,
    maxtime::Union{Number,Nothing}=nothing,
    abstol::Union{Number,Nothing}=nothing,
    reltol::Union{Number,Nothing}=nothing,
    progress=false,
    surrogate_objective::Union{Symbol,Nothing}=nothing,
    integer=nothing,
    kwargs...)

    # local x

    maxiters = GalacticOptim._check_and_convert_maxiters(maxiters)
    maxtime = GalacticOptim._check_and_convert_maxtime(maxtime)

    f = GalacticOptim.instantiate_function(prob.f, prob.u0, prob.f.adtype, prob.p)

    _loss = NonconvexADWrapper(f, prob)

    opt_f(θ) = _loss(θ)

    opt_args, integer = __map_optimizer_args(prob, opt, callback=callback, maxiters=maxiters, maxtime=maxtime, abstol=abstol, reltol=reltol, integer=integer; kwargs...)

    opt_set = Nonconvex.Model()
    if isnothing(surrogate_objective)
        Nonconvex.set_objective!(opt_set, opt_f)
    else
        Nonconvex.set_objective!(opt_set, opt_f, flags=[surrogate_objective])
    end

    Nonconvex.addvar!(opt_set, prob.lb, prob.ub, integer=integer)

    if !isnothing(prob.f.cons)
        @warn "Equality constraints are current not passed on by GalacticOptim"
        #add_ineq_constraint!(opt_set, f)
    end

    if !isnothing(prob.lcons)
        @warn "Inequality constraints are current not passed on by GalacticOptim"
        #add_ineq_constraint!(opt_set, f)
    end

    if !isnothing(prob.ucons)
        @warn "Inequality constraints are current not passed on by GalacticOptim"
        #add_ineq_constraint!(opt_set, f)
    end

    t0 = time()
    if check_optimizer_backend(opt)
        opt_res = Nonconvex.optimize(opt_set, opt; opt_args...)
    else
        opt_res = Nonconvex.optimize(opt_set, opt, prob.u0; opt_args...)
    end
    t1 = time()

    opt_ret = hasproperty(opt_res, :status) ? Symbol(string(opt_res.status)) : nothing

    SciMLBase.build_solution(prob, opt, opt_res.minimizer, opt_res.minimum; (isnothing(opt_ret) ? (; original=opt_res) : (; original=opt_res, retcode=opt_ret))...)
end

end


text:
using AtomBase
using Test

@testset "Atom.jl" begin
    atom1 = Atom()
    atom2 = Atom("H", [1.0, 0.0, 0.0])
    io = IOBuffer()
    show(io, atom2)
    @test atom2() == ("H", [1.0, 0.0, 0.0])
    @test !(atom1 == atom2)
    @test String(take!(io)) == "H\t[1.0, 0.0, 0.0]"
end

@testset "Atoms.jl Properties" begin
    property1 = Properties()
    property2 = Properties{Float64}("Force", 1.0)
    property3 = Properties("energy")
    energy1 = Energy()
    energy2 = Energy(1.0)
    energy3 = Energy(2)
    @test eltype(property1) == eltype(property2) == eltype(energy1) == eltype(energy2) == eltype(energy3) == Float64
    @test property2() == 1.0
    @test energy3() == 2.0
    energy1(4.0)
    @test energy1() == 4.0
    io = IOBuffer()
    show(io, property2)
    @test String(take!(io)) == "Force=1.0"
    show(io, energy3)
    @test String(take!(io)) == "energy=2.0"
end

@testset "Atoms.jl Atoms" begin
    atom = Atom("H", [1.0, 0.0, 0.0])
    property = Energy(1.0)
    atoms0 = Atoms()
    atoms1 = Atoms([atom, atom], 2, [property, property])
    atoms2 = Atoms([atom, atom, atom], [property,])
    atoms3 = Atoms([atom,])
    atoms4 = Atoms(atom, [property, ])
    atoms5 = Atoms(atom)
    atoms6 = Atoms([atom, atom], 2, property)
    atoms7 = Atoms([atom, atom], property)
    atoms8 = Atoms(atom, property)
    @test atoms2.number == 3
    addAtom!(atoms2, atom)
    @test atoms2.number == 4
    addProperty!(atoms3, property)
    @test atoms3.properties == [property,]
    @test iterate(atoms2) == (atoms2.atom[1], 2)
    @test iterate(atoms2, 2) == (atoms2.atom[2], 3)
    @test length(atoms2) == 4
    @test getindex(atoms2, 1) == atom
    @test firstindex(atoms2) == 1
    @test lastindex(atoms2) == 4
    @test getindex(atoms2, 1.0) == atom
    @test getindex(atoms2, [1, 2]) == [atom, atom]
    io = IOBuffer()
    show(io, atoms8)
    @test String(take!(io)) == "1\nenergy=1.0 \nH\t[1.0, 0.0, 0.0]\n"
end


text:
"""
    Double pendulum
"""
struct Impact end
struct Nominal end

struct DoublePendulum{T,X} <: Model{T}
    nq::Int
    nu::Int
    nw::Int
    nc::Int

    m1::T    # mass link 1
    J1::T    # inertia link 1
    l1::T    # length link 1
    lc1::T   # length to COM link 1

    m2::T    # mass link 2
    J2::T    # inertia link 2
    l2::T    # length link 2
    lc2::T   # length to COM link 2

    g::T     # gravity

    b1::T    # joint friction
    b2::T
end

lagrangian(model::DoublePendulum, q, q̇) = 0.0

function kinematics(model::DoublePendulum, x)
    [model.l1 * sin(x[1]) + model.l2 * sin(x[1] + x[2]),
     -1.0 * model.l1 * cos(x[1]) - model.l2 * cos(x[1] + x[2])]
end

function kinematics_elbow(model::DoublePendulum, x)
    [model.l1 * sin(x[1]),
     -1.0 * model.l1 * cos(x[1])]
end

function M_func(model::DoublePendulum, x)
    a = (model.J1 + model.J2 + model.m2 * model.l1 * model.l1
         + 2.0 * model.m2 * model.l1 * model.lc2 * cos(x[2]))

    b = model.J2 + model.m2 * model.l1 * model.lc2 * cos(x[2])

    c = model.J2

    [a b;
     b c]
end

function τ_func(model::DoublePendulum, x)
    a = (-1.0 * model.m1 * model.g * model.lc1 * sin(x[1])
         - model.m2 * model.g * (model.l1 * sin(x[1])
         + model.lc2 * sin(x[1] + x[2])))

    b = -1.0 * model.m2 * model.g * model.lc2 * sin(x[1] + x[2])

    [a, b]
end

function c_func(model::DoublePendulum, q, q̇)
    a = -2.0 * model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[2]
    b = -1.0 * model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[2]
    c = model.m2 * model.l1 * model.lc2 * sin(q[2]) * q̇[1]
    d = 0.0

    [a b;
     c d]
end

function B_func(model::DoublePendulum, x)
    [0.0; 1.0]
end

function C_func(model::DoublePendulum, q, q̇)
    c_func(model, q, q̇) * q̇ - τ_func(model, q)
end

function ϕ_func(model, q)
    [0.5 * π - q[2], q[2] + 0.5 * π]
end

function P_func(model, q)
    ϕ = ϕ_func(model, q)
    Symbolics.jacobian(ϕ, q)
end

function dynamics(model::DoublePendulum{T,Impact}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * u1[1]
        + transpose(P_func(model, q2)) * λ1
        - h[1] * 0.5 .* vm2) # damping
end

function dynamics(model::DoublePendulum{T,Nominal}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ B_func(model, qm2) * u1[1]
        - h[1] * 0.5 .* vm2) # damping
end

function residual(model::DoublePendulum{T,Impact}, z, θ, κ) where T
    nq = model.nq
    nu = model.nu
    nc = model.nc

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]
    λ1 = z[nq .+ (1:nc)]
    s1 = z[nq + nc .+ (1:nc)]

    [
     dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b),
        h, q0, q1, u1, zeros(model.nw), λ1, q2);
     s1 .- ϕ_func(model, q2);
     λ1 .* s1 .- κ;
    ]

end

function residual(model::DoublePendulum{T,Nominal}, z, θ, κ) where T
    nq = model.nq
    nu = model.nu
    nc = model.nc

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]
    return dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b),
        h, q0, q1, u1, zeros(model.nw), zeros(model.nc), q2)
end

acrobot_impact = DoublePendulum{Float64,Impact}(2, 1, 0, 2,
    1.0, 0.333, 1.0, 0.5, 1.0, 0.333, 1.0, 0.5, 9.81, 0.0, 0.0)

acrobot_nominal = DoublePendulum{Float64,Nominal}(2, 1, 0, 0,
    1.0, 0.333, 1.0, 0.5, 1.0, 0.333, 1.0, 0.5, 9.81, 0.0, 0.0)


text:
include("test_basic_inputs.jl")
include("test_nuisance_inputs.jl")


text:
"Integrate an implicit DAE with a specialised partitioned additive Runge-Kutta integrator."
function integrate_step!(int::AbstractIntegratorSPARK{DT,TT}, cache::IntegratorCacheSPARK{DT,TT}) where {DT,TT}
    # update nonlinear solver parameters from cache
    update_params!(int.params, cache)

    # compute initial guess
    initial_guess!(int, cache)

    # reset cache
    reset!(cache, timestep(int))

    # call nonlinear solver
    solve!(int.solver)

    # print solver status
    print_solver_status(int.solver.status, int.solver.params, cache.n)

    # check if solution contains NaNs or error bounds are violated
    check_solver_status(int.solver.status, int.solver.params, cache.n)

    # compute vector fields at internal stages
    compute_stages!(int.solver.x, cache, int.params)

    # compute final update
    update_solution!(cache.q, cache.qₑᵣᵣ, cache.Vi, int.params.t_q.b, timestep(int))
    update_solution!(cache.p, cache.pₑᵣᵣ, cache.Fi, int.params.t_p.b, timestep(int))

    # compute projection
    update_solution!(cache.q, cache.qₑᵣᵣ, cache.Up, int.params.t_q.β, timestep(int))
    update_solution!(cache.p, cache.pₑᵣᵣ, cache.Gp, int.params.t_p.β, timestep(int))
    # TODO # update_multiplier!(cache.λ, cache.Λp, int.params.t_λ.b)

    # copy solution to initial guess
    update!(int.iguess, cache.t, cache.q, cache.p, cache.v, cache.f)

    # take care of periodic solutions
    cut_periodic_solution!(cache, equation(int).periodicity)
end