text:
module LinearElasticity

using ConstructionBase: constructorof
using Tensorial: SymmetricSecondOrderTensor, SymmetricFourthOrderTensor, Vec

export TensorStress,
    TensorStrain,
    StiffnessTensor,
    ComplianceTensor,
    EngineeringStress,
    EngineeringStrain,
    ComplianceMatrix,
    StiffnessMatrix

abstract type Stress{T,N} <: AbstractArray{T,N} end
abstract type Strain{T,N} <: AbstractArray{T,N} end
abstract type Stiffness{T,N} <: AbstractArray{T,N} end
abstract type Compliance{T,N} <: AbstractArray{T,N} end
struct TensorStress{T} <: Stress{T,2}
    data::SymmetricSecondOrderTensor{3,T,6}
end
TensorStress(m::AbstractMatrix) = TensorStress(SymmetricSecondOrderTensor{3}(m))
TensorStress(data...) = TensorStress(SymmetricSecondOrderTensor{3}(data...))
struct TensorStrain{T} <: Strain{T,2}
    data::SymmetricSecondOrderTensor{3,T,6}
end
TensorStrain(m::AbstractMatrix) = TensorStrain(SymmetricSecondOrderTensor{3}(m))
TensorStrain(data...) = TensorStrain(SymmetricSecondOrderTensor{3}(data...))
struct StiffnessTensor{T} <: Stiffness{T,4}
    data::SymmetricFourthOrderTensor{3,T}
end
struct ComplianceTensor{T} <: Compliance{T,4}
    data::SymmetricFourthOrderTensor{3,T}
end
struct EngineeringStress{T} <: Stress{T,1}
    data::Vec{6,T}
end
EngineeringStress(x::TensorStress) = convert(EngineeringStress{eltype(x)}, x)
EngineeringStress(v::AbstractVector) = EngineeringStress(Vec{6}(v))
EngineeringStress(data...) = EngineeringStress(Vec{6}(data...))
TensorStress(x::EngineeringStress) = convert(TensorStress{eltype(x)}, x)
struct EngineeringStrain{T} <: Strain{T,1}
    data::Vec{6,T}
end
EngineeringStrain(x::TensorStrain) = convert(EngineeringStrain{eltype(x)}, x)
EngineeringStrain(v::AbstractVector) = EngineeringStrain(Vec{6}(v))
EngineeringStrain(data...) = EngineeringStrain(Vec{6}(data...))
TensorStrain(x::EngineeringStrain) = convert(TensorStrain{eltype(x)}, x)
struct StiffnessMatrix{T} <: Stiffness{T,2}
    data::SymmetricSecondOrderTensor{6,T,21}
end
StiffnessMatrix(x::StiffnessTensor) = convert(StiffnessMatrix{eltype(x)}, x)
StiffnessMatrix(m::AbstractMatrix) = StiffnessMatrix(SymmetricSecondOrderTensor{6}(m))
StiffnessMatrix(data...) = StiffnessMatrix(SymmetricSecondOrderTensor{6}(data...))
StiffnessTensor(x::StiffnessMatrix) = convert(StiffnessTensor{eltype(x)}, x)
struct ComplianceMatrix{T} <: Compliance{T,2}
    data::SymmetricSecondOrderTensor{6,T,21}
end
ComplianceMatrix(x::ComplianceTensor) = convert(ComplianceMatrix{eltype(x)}, x)
ComplianceMatrix(m::AbstractMatrix) = ComplianceMatrix(SymmetricSecondOrderTensor{6}(m))
ComplianceMatrix(data...) = ComplianceMatrix(SymmetricSecondOrderTensor{6}(data...))
ComplianceTensor(x::ComplianceMatrix) = convert(ComplianceTensor{eltype(x)}, x)

Base.size(::Union{TensorStress,TensorStrain}) = (3, 3)
Base.size(::Union{StiffnessTensor,ComplianceTensor}) = (3, 3, 3, 3)
Base.size(::Union{EngineeringStress,EngineeringStrain}) = (6,)
Base.size(::Union{StiffnessMatrix,ComplianceMatrix}) = (6, 6)

Base.getindex(A::Union{Stress,Strain,Stiffness,Compliance}, i...) = getindex(A.data, i...)

Base.parent(A::Union{Stress,Strain,Stiffness,Compliance}) = A.data

Base.IndexStyle(::Type{<:Union{Stress,Strain,Stiffness,Compliance}}) = IndexLinear()

# See https://github.com/JuliaLang/julia/blob/cb9acf5/base/arraymath.jl#L19-L26
for op in (:*, :/)
    @eval Base.$op(A::Union{Stress,Strain,Stiffness,Compliance}, B::Number) =
        constructorof(typeof(A))(Base.broadcast_preserving_zero_d($op, A, B))
end
Base.:*(B::Number, A::Union{Stress,Strain,Stiffness,Compliance}) = A * B

Base.:-(A::Union{Stress,Strain,Stiffness,Compliance}) =
    constructorof(typeof(A))(Base.broadcast_preserving_zero_d(-, A))

for op in (:+, :-)
    for T in (
        :TensorStress,
        :TensorStrain,
        :EngineeringStress,
        :EngineeringStrain,
        :StiffnessMatrix,
        :ComplianceMatrix,
        :StiffnessTensor,
        :ComplianceTensor,
    )
        @eval Base.$op(A::$T, B::$T) = $T(Base.broadcast_preserving_zero_d($op, A, B))
    end
end

include("conversion.jl")
include("invariants.jl")
include("SymmetryCriteria.jl")
# include("StabilityCriteria.jl")
include("Isotropic.jl")
include("distort.jl")
include("energy.jl")
include("solve.jl")

end


text:
include("field_operations.jl")
include("spherical_harmonics.jl")
using Functors

struct Grid
    dx::AbstractFloat
    sz::Any
    origin::Any
    spherical_harmonics::Any
    r::Any
    dV::AbstractFloat
end
@functor Grid

"""
    Grid(dx::AbstractFloat, rmax::AbstractFloat; dims = 3, rank_max = 1)
"""
function Grid(dx::AbstractFloat, rmax::AbstractFloat; dims = 3, rank_max = 2)
    dV = dx^dims

    nrmax = round(Int, rmax / dx)
    origin = (nrmax + 1) * ones(dims)
    sz = fill(2nrmax + 1, dims)

    r, spherical_harmonics =
        make_spherical_harmonics(dx, sz, origin, rank_max; dims)
    Grid(dx, sz, origin, spherical_harmonics, r, dV)
end

"""
    Grid(
        dx::AbstractFloat,
        sz::Union{AbstractVector,Tuple};
        origin = nothing,
        rank_max = 1
        )
"""
function Grid(
    dx::AbstractFloat,
    sz::Union{AbstractVector,Tuple};
    origin = nothing,
    rank_max = 1,
)
    dims = length(sz)
    dV = dx^dims
    if origin === nothing
        origin = (sz .+ 1.0) / 2
    end
    r, spherical_harmonics =
        make_spherical_harmonics(dx, sz, origin, rank_max; dims)
    Grid(dx, sz, origin, spherical_harmonics, r, dV)
end

GRID_ALIASES = (
    x = x -> x.spherical_harmonics[1][1],
    y = x -> x.spherical_harmonics[1][2],
    z = x -> x.spherical_harmonics[1][3],
)

function Base.getproperty(grid::Grid, v::Symbol)
    v in GRID_ALIASES ? GRID_ALIASES[v](grid) : getfield(grid, v)
end

"""
    Field(; grid = nothing, rank = nothing, radfunc = nothing)
"""
function Field(; grid = nothing, rank = nothing, radfunc = nothing)
    if radfunc !== nothing
        r = radfunc.(grid.r)
        r = join([r])
        return rank == 0 ? r : field_prod(r, grid.spherical_harmonics[rank])
    end

    zeros(grid.sz..., 2rank + 1)
end

struct Props

    nranks::Any
    ranks::Any
    parities::Any
    slices::Any
end

function Props(ranks)
    Props(; ranks)
end
function Props(; nranks = nothing, ranks = nothing)
    if ranks === nothing
        ranks = vcat([fill(i - 1, n) for (i, n) in enumerate(nranks)]...)
    elseif nranks === nothing
        nranks = zeros(Int, maximum(ranks) + 1)
        for x in ranks
            nranks[x+1] += 1
        end
    end
    ends = Int.(cumsum(2 .* ranks .+ 1))
    slices = [(i == 1 ? 1 : ends[i-1] + 1):ends[i] for i in eachindex(ends)]

    # starts = [1, (1 .+ cumsum(nranks))...]
    # grouped_slices = [slices[starts[i]:starts[i+1]-1] for i = 1:length(nranks)]
    parities = (-1) .^ ranks
    Props(nranks, ranks, parities, slices)
end

"""
    (field::AbstractArray)(rvec, grid::Grid)

Makes field objects callable. Returns interpolated tensor value of the tensor field at position rvec.
"""
function Base.get(fields::AbstractArray, props::Props, i::Int)
    fields[:, :, :, props.slices[i]]
end

function Base.get(field::AbstractArray, grid::Grid, rvec::AbstractVector)
    sum([
        w * getindex.(eachslice(field, dims = 4), ix...)
        for (ix, w) in nearest(grid.dx, grid.origin, rvec)
    ])
end
function Base.get(
    fields::AbstractArray,
    props::Props,
    i::Int,
    rvec::AbstractVector,
)
    field = get(fields, props, i)
    get(field, props.grid, rvec)
end
function Base.get(fields::AbstractArray, props::Props, rvec::AbstractVector)
    get(fields, props, 1, rvec)
end

"""
    (field::AbstractArray)(i::Int)

Makes field objects callable. Returns ith component array (2d/3d) of tensor field.
"""
function (field::AbstractArray)(i::Int)
    selectdim(field, dim(field), i)
end

function nearest(dx, origin, rvec)
    ix = origin .+ rvec / dx
    ixfloor = floor.(Int, ix)
    er = ix - ixfloor
    [
        (ixfloor + [x, y, z], prod(ones(3) - abs.([x, y, z] - er)))
        for x = 0:1, y = 0:1, z = 0:1
    ]
end

"""
    put_point_source!(
        field::AbstractArray,
        grid::Grid,
        rvec::AbstractVector,
        val::AbstractVector)

"""
function put_point_source!(
    field::AbstractArray,
    grid::Grid,
    rvec::AbstractVector,
    val::AbstractVector,
)
    for (ix, w) in nearest(grid.dx, grid.origin, rvec)
        for (c, vi) in zip(1:size(field)[end], val)
            field[ix..., c] += w / grid.dV * vi
        end
    end
end

function put_point_source!(f, grid, rvecs::AbstractMatrix, vals::AbstractMatrix)
    for (val, rvec) in zip(eachcol(vals), eachcol(rvecs))
        put_point_source!(f, grid, rvec, val)
    end
end

"""
    field_rank(x::AbstractArray)
"""
field_rank(x::AbstractArray) = (size(x)[end] - 1) ÷ 2

ext: jl

max_stars_repo_name: aced-differentiate/equiv-docs

max_issues_repo_name: aced-differentiate/equiv-docs

max_forks_repo_name: aced-differentiate/equiv-docs


text:
import GitLab
using Base.Test
myauth = GitLab.authenticate(ENV["GITLAB_AUTH"]) # don't hardcode your access tokens!
println("Authentication successful")
options = Dict("private_token" => myauth.token)

myrepo = GitLab.repo_by_name("TestProject1"; headers=options)

issues, page_data = GitLab.issues(myrepo; params=options)
## @show issues

## get specific issues:
## state        string  no  Return all issues or just those that are opened or closed
## labels       string  no  Comma-separated list of label names, issues with any of the labels will be returned
## order_by     string  no  Return requests ordered by created_at or updated_at fields. Default is created_at
## sort         string  no  Return requests sorted in asc or desc order. Default is desc
options["labels"] = "MyLabel"
issues = GitLab.issues(myrepo; params=options)
## @show issues

issue = issues[1]

options["labels"] = "wronglabel"
issues, page_data = GitLab.issues(myrepo; params=options)
@test sizeof(issues) == 0
if haskey(options, "labels")
    delete!(options, "labels")
end

issue = GitLab.issue(myrepo, get(issue[1].id); params=options)
## @show issue

issue_data = Dict{AbstractString, Any}()
issue_data["title"] = "This is a test"
issue_data["description"] = "This is some description"
issue = GitLab.create_issue(myrepo; params=issue_data, headers=options)
## @show issue
@test get(issue.title) == "This is a test"

issue_data["title"] = "This is a test - edit"
issue_data["description"] = "This is some description - edit"
issue = GitLab.edit_issue(myrepo, get(issue.id); params=issue_data, headers=options)
## @show issue
@test get(issue.title) == "This is a test - edit"

println("Sleeping before deleting the issue !!")
sleep(10)

del_issue = GitLab.delete_issue(myrepo, get(issue.id); params=issue_data, headers=options)
## @show del_issue
@test get(del_issue.title) == "This is a test - edit"

println("Issue Tests Done !!!")


text:
# Function for loading sample data

"""
    get_data(dataset, poollength, datalength, horizon)

Function for importing one of the sample datasets in the repository. `dataset` can be one of
`:solar`, `:traffic`, `:exchange_rate` or `:electricity`. `poollength` gives the number of
timesteps to pool for the model, `datalength` determines the batch size of the output data,
and `horizon` determines the number of time steps that should be forecasted by the model.
"""
function get_data(dataset, poollength, datalength, horizon; normalise=true)
    admissible = [:solar, :traffic, :exchange_rate, :electricity]
    dataset in admissible || error("Sample data not implemented")

    datalength = datalength + poollength
    (dataset == :solar) && (BSON.@load "../data/solar_AL.bson" inp_raw)
    (dataset == :traffic) && (BSON.@load "../data/traffic.bson" inp_raw)
    (dataset == :exchange_rate) && (BSON.@load "../data/exchange_rate.bson" inp_raw)
    (dataset == :electricity) && (BSON.@load "../data/electricity.bson" inp_raw)

    (normalise == true) && (inp_raw = Flux.normalise(inp_raw, dims=1))
    out_ft = similar(inp_raw, size(inp_raw,2), poollength, 1, size(inp_raw,1))
    for i=0:poollength-1
        for j=poollength:min(datalength, size(inp_raw,1)-poollength)
            out_ft[:,i+1,1,j] = inp_raw[j-i,:]
        end
    end
    out_lb = circshift(inp_raw[1:datalength-poollength,1], -horizon)
    return out_ft[:,:,:,1:min(datalength, size(inp_raw,1)-poollength)-poollength], out_lb
end

max_stars_repo_name: cbdavis/FluxArchitectures

max_issues_repo_name: cbdavis/FluxArchitectures

max_forks_repo_name: cbdavis/FluxArchitectures


text:
#=

There are redundant exports in this file because it is also used to keep track of where different methods come from. Each source file inclusion should be followed by an export statement with an alphabetic list of all the exported functions with a method that comes from that file. Please keep this format.

=#

module Phylodendron2

using Random: randperm
using DelimitedFiles: readdlm
using PrettyTables

include("exceptions.jl")

include("species_directory.jl")
export
    addspecies!,
    delete!,
    getindex,
    in,
    length,
    rename!,
    SpeciesDirectory

include("species_data.jl")
export
    read_species_data,
    SpeciesDataMatrix,
    write_phylip

include("bipartition.jl")
export
    Bipartition

include("datablocks/types.jl")
export
    AbstractBranchDataBlock,
    AbstractNodeDataBlock,
    AbstractTreeDataBlock,
    BranchDataBlock,
    NodeDataBlock,
    TreeDataBlock

include("tree/types.jl")
export
    AbstractBranch,
    AbstractNode,
    AbstractTree,
    Branch,
    Node,
    NodeVector,
    Tree,
    TreeVector

include("clone.jl")
export
    clone

include("node_properties.jl")
export
    annotate!,
    annotation,
    deannotate!,
    label,
    label!,
    getspecies,
    setspecies!

include("find.jl")
# exports methods for findfirst and findall

include("peek_topology.jl")
export
    areneighbours,
    getbranch,
    istip,
    neighbours,
    n_neighbour

include("branch_properties.jl")
export
    brannotate!,
    brannotation,
    brdeannotate!,
    brlabel,
    brlabel!,
    brlength,
    brlength!

include("datablocks/manipulation.jl")
export
    delete_datablock!

include("tree_properties.jl")
export
    annotate!,
    annotation,
    deannotate!,
    isrooted,
    label,
    label!,
    origin,
    origin!,
    root!,
    setspecies!,
    unroot!

include("show.jl")

include("link.jl")
export
    graft!,
    link!,
    pluck!,
    unlink!

include("traversals.jl")
export
    postorder,
    postorder_vector,
    preorder,
    preorder_vector

include("ladderise.jl")
export
    ladderise!

include("subtree_funcs.jl")
export
    bipartitions,
    find_nonsplitting,
    getspecies,
    n_branch,
    n_node,
    n_species,
    n_tip,
    nodelabels,
    tiplabels,
    tips

include("topomanip.jl")
export
    pluck_nonsplitting!

include("node_paths.jl")
export
    brdist,
    nodepath,
    nodedist

include("bipartition_funcs.jl")
export
    are_compatible,
    are_conflicting,
    compute_bipartitions!,
    isinformative,
    istrivial,
    update_bipartition!

include("lexer.jl")
include("newick.jl")
export
    newick_string,
    parse_newick,
    read_newick,
    write_newick

include("randtree.jl")
export
    randtree

include("BrownianModels/BrownianModels.jl")

end # module Phylodendron2


text:
#!/usr/bin/env julia

function main()
    accum = 0
    fib_n  = 1
    fib_nn = 2
    fib_max = 4_000_000
    while fib_n <= fib_max
        if fib_n % 2 == 0
            accum += fib_n
        end
        tmp = fib_n + fib_nn
        fib_n = fib_nn
        fib_nn = tmp
    end
    @printf("The sum of even Fibonacci numbers up to %d is %d\n",
            fib_max, accum)
end

main()


text:
@testset "Equalized Odds Postprocessing" begin
	X, y, _ = @load_toydata;
	model = ConstantClassifier()
	wrappedModel = EqOddsWrapper(model; grp=:Sex)
	mach = machine(wrappedModel, X, y)
	fit!(mach)
	ŷ = predict(mach, X[6:10, :])
	@test all(ŷ .== 1)
	@test length(ŷ)==5
	# TODO:test the fitresult of machine
end

@testset "LinProgWrapper Postprocessing" begin
	X, y, _ = @load_toydata;
	model = ConstantClassifier()
	wrappedModel = LinProgWrapper(model, grp=:Sex, measure=true_positive_rate)
	mach = machine(wrappedModel, X, y)
	fit!(mach)
	ŷ = predict(mach, X[6:10, :])
	@test length(ŷ)==5
	# TODO:test the fitresult of machine
end


text:
#By convention from definetypes.jl, make title! update the first Title object in array
title!(ec::EChart; kwargs...) = text!(ec, 1; kwargs...)

#Push on to end of ec.title array
function text!(ec::EChart; kwargs...)

	temp = Title()

	for (k, v) in kwargs
	   setfield!(temp, k, v)
	end

	push!(ec.title, temp)

	return ec

end

#Modify a piece of text in place
function text!(ec::EChart, series::Int; kwargs...)

	for (k, v) in kwargs
	   setfield!(ec.title[series], k, v)
	end

	return ec

end


text:
using PortHamiltonianBenchmarkSystems, MAT

for n in [320, 980, 1805]
  E, J, R, B = poro_elasticity_model(n=n)
  filename = "PoroModelN$n"
  writeHDF5(filename*".h5"; E, J, R, B)
  writeMAT(filename*".mat"; E, J, R, B)
end


text:
#=------------------------------------------------------------------------------
  A storage file for old code that may be useful in the future.
------------------------------------------------------------------------------=#

#=-------------------------------Old Constructors-----------------------------=#
function SSSTensor(T::triangles_iterator)
   edges = Array{Tuple{Array{Int64,1},Float64}}(undef,0)
   for t in T
      push!(edges,([t.v1,t.v2,t.v3],1.0))
    end
    SSSTensor(edges)
  end

  function SSSTensor(T::triangles_iterator,n)
    edges = Array{Tuple{Array{Int64,1},Float64}}(undef,0)
    for t in T
      push!(edges,([t.v1,t.v2,t.v3],1.0))
    end
	SSSTensor(edges,n)
  end


text:
import LightGraphs: TreeBFSVisitorVector, bfs_tree!, tree
@testset "BFS" begin
    g5 = DiGraph(4)
    add_edge!(g5,1,2); add_edge!(g5,2,3); add_edge!(g5,1,3); add_edge!(g5,3,4)
    g6 = smallgraph(:house)

    for g in testdigraphs(g5)
      z = @inferred(bfs_tree(g, 1))
      visitor = LightGraphs.TreeBFSVisitorVector(zeros(eltype(g),nv(g)))
      LightGraphs.bfs_tree!(visitor, g, 1)
      t = visitor.tree
      @test t == [1,1,1,3]
      @test nv(z) == 4 && ne(z) == 3 && !has_edge(z, 2, 3)
    end
    for g in testgraphs(g6)
      @test @inferred(gdistances(g, 2)) == [1, 0, 2, 1, 2]
      @test @inferred(gdistances(g, [1,2])) == [0, 0, 1, 1, 2]
      @test @inferred(gdistances(g, [])) == [-1, -1, -1, -1, -1]
      @test @inferred(!is_bipartite(g))
      @test @inferred(!is_bipartite(g, 2))
    end

    gx = Graph(5)
    add_edge!(gx,1,2); add_edge!(gx,1,4)
    add_edge!(gx,2,3); add_edge!(gx,2,5)
    add_edge!(gx,3,4)

    for g in testgraphs(gx)
      @test @inferred(is_bipartite(g))
      @test @inferred(is_bipartite(g, 2))
    end

    # import LightGraphs: TreeBFSVisitorVector, bfs_tree!, tree

    function istree{T<:Integer}(parents::Vector{T}, maxdepth, n::T)
        flag = true
        for i in one(T):n
            s = i
            depth = 0
            while parents[s] > 0 && parents[s] != s
                s = parents[s]
                depth += 1
                if depth > maxdepth
                    return false
                end
            end
        end
        return flag
    end

    for g in testgraphs(g6)
      n = nv(g)
      visitor = @inferred(TreeBFSVisitorVector(n))
      @test length(visitor.tree) == n
      parents = visitor.tree
      @inferred(bfs_tree!(visitor, g, 1))

      @test istree(parents, n, n)
      t = tree(parents)
      @test is_directed(t)
      @test typeof(t) <: AbstractGraph
      @test ne(t) < nv(t)

    # test Dict{Int,Int}() colormap

      visitor = @inferred(TreeBFSVisitorVector(n))
      @test length(visitor.tree) == n
      parents = visitor.tree
      @inferred(bfs_tree!(visitor, g, 1, vertexcolormap = Dict{Int,Int}()))

      @test @inferred(istree(parents, n, n))
      t = tree(parents)
      @test is_directed(t)
      @test typeof(t) <: AbstractGraph
      @test ne(t) < nv(t)
    end

    g10 = CompleteGraph(10)
    for g in testgraphs(g10)
      @test @inferred(bipartite_map(g)) == Vector{eltype(g)}()
    end

    g10 = CompleteBipartiteGraph(10,10)
    for g in testgraphs(g10)
      T = eltype(g)
      @test @inferred(bipartite_map(g10)) == Vector{T}([ones(T, 10); 2*ones(T, 10)])

      h = blkdiag(g,g)
      @test @inferred(bipartite_map(h)) == Vector{T}([ones(T, 10); 2*ones(T, 10); ones(T, 10); 2*ones(T, 10)])
    end
end

max_stars_repo_name: AsileBcd/JohnsonCycles

max_issues_repo_name: AsileBcd/JohnsonCycles

max_forks_repo_name: AsileBcd/JohnsonCycles


text:
struct AliasTable <: Sampleable{Univariate,Discrete}
    accept::Vector{Float64}
    alias::Vector{Int}
end
ncategories(s::AliasTable) = length(s.alias)

function AliasTable(probs::AbstractVector)
    n = length(probs)
    n > 0 || throw(ArgumentError("The input probability vector is empty."))
    accp = Vector{Float64}(undef, n)
    alias = Vector{Int}(undef, n)
    StatsBase.make_alias_table!(probs, 1.0, accp, alias)
    AliasTable(accp, alias)
end

function rand(rng::AbstractRNG, s::AliasTable)
    i = rand(rng, 1:length(s.alias)) % Int
    u = rand(rng)
    @inbounds r = u < s.accept[i] ? i : s.alias[i]
    r
end

show(io::IO, s::AliasTable) = @printf(io, "AliasTable with %d entries", ncategories(s))


text:
module Counters2

# stdlib
using Dates

# for ordereddict
using DataStructures

# Measurement utilities
using MattDaemon
using CounterTools
using SystemSnoop
using DataCollection
using TimeSync

# The main thing
using AutoTM
using nGraph

#####
##### Global constants
#####

socket() = CounterTools.IndexZero(1)
port() = 2000
cpus() = 25:48
sampletime() = Dates.Millisecond(500)
database() = joinpath(dirname(@__DIR__), "autotm_counters.jls")
database_snoop() = joinpath(dirname(@__DIR__), "autotm_snooped_counters.jls")

const UncoreSelectRegister = CounterTools.UncoreSelectRegister
const CoreSelectRegister = CounterTools.CoreSelectRegister

function enable_timestamps()
    # Enable timestamping of ngraph kernels
    ENV["NGRAPH_KERNEL_TIMESTAMPS"] = true
end

#####
##### Interface AutoTM stuff with DataCollection
#####

# backend
DataCollection.lower(x::nGraph.Backend{T}) where {T} = string(T)

# model constructors
const ExperimentType = Union{
    AutoTM.Experiments.Resnet,
    AutoTM.Experiments.Inception_v4,
    AutoTM.Experiments.DenseNet,
    AutoTM.Experiments.Vgg
}
DataCollection.lower(x::ExperimentType) = AutoTM.Experiments.name(x)

# optimizers
DataCollection.lower(x::AutoTM.Optimizer.AbstractOptimizer) = AutoTM.Optimizer.name(x)

# cached
DataCollection.lower(x::T) where {T <: AutoTM.Profiler.AbstractKernelCache} = string(T)

#####
##### Counter Setup
#####

# Group IMC Counters into sets of four since we can only program four counters at a time.
const GROUPED_IMC_COUNTERS = [
    (
        UncoreSelectRegister(; event = 0xD3, umask = 0x01) => "Tag Hit",
        UncoreSelectRegister(; event = 0xD3, umask = 0x02) => "Tag Miss Clean",
        UncoreSelectRegister(; event = 0xD3, umask = 0x04) => "Tag Miss Dirty",
        UncoreSelectRegister(; event = 0x00, umask = 0x00) => "Uncore Clocks",
    ),
    (
        UncoreSelectRegister(; event = 0x04, umask = 0x03) => "DRAM Read",
        UncoreSelectRegister(; event = 0x04, umask = 0x0C) => "DRAM Write",
        UncoreSelectRegister(; event = 0xEA, umask = 0x02) => "PMM Read",
        UncoreSelectRegister(; event = 0xEA, umask = 0x04) => "PMM Write",
    )
]

function coreevents()
    return (
        # Just a dummy event for now
        events = (
            CoreSelectRegister(),
        ),
        fixed_events = (
            CounterTools.RetiredInstructions,
            CounterTools.ClockUnhaltedCore,
        ),
    )
end

#####
##### Run Parameters
#####

Base.@kwdef struct RunParameters <: DataCollection.AbstractParameters
    f::ExperimentType
    optimizer::AutoTM.Optimizer.AbstractOptimizer
    backend::nGraph.Backend
    cache::Any
    mode::String
end

#####
##### Run Loop
#####

maybeunwrap(x) = x
maybeunwrap(x::Tuple) = first(x)

function run(parameters::RunParameters)
    # Step 1 - create the desired AutoTM function and run it once to warm everything up.
    fex = AutoTM.Optimizer.factory(
        parameters.backend,
        parameters.f,
        parameters.optimizer;
        cache = parameters.cache
    ) |> maybeunwrap
    @time fex()

    sampletime_parameters = DataCollection.GenericParameters(; sampletime = sampletime())

    # Now, take IMC measurements
    for event_name_pairs in GROUPED_IMC_COUNTERS
        # Reset counters to generate a timeline after data collection.
        nGraph.reset_counters(fex.ex)

        # prepare imc events
        events = first.(event_name_pairs)

        # prepare core events
        core_nt = coreevents()
        core_events = core_nt.events
        core_fixed = core_nt.fixed_events

        # construct measurement to send to the monitor
        measurements = MattDaemon.@measurements (
            pretime = SystemSnoop.Timestamp(),
            imc = CounterTools.IMCMonitor(events, socket()),
            core = CounterTools.CoreMonitor(core_events, cpus(); fixed_events = core_fixed),
            posttime = SystemSnoop.Timestamp(),
        )

        payload = MattDaemon.ServerPayload(sampletime(), measurements)
        data, _, runtime = MattDaemon.run(fex, payload, port(); sleeptime = Second(1))
        @show runtime

        # Take counter differences and aggregate counters across all IMC Channels
        counter_values = data.imc
        deltas = CounterTools.aggregate.(diff(counter_values))
        df = DataCollection.load(database())

        for (i, name) in enumerate(last.(event_name_pairs))
            # Extract the values for this counter
            this_counter = getindex.(deltas, i)
            @show name

            dc_data = DataCollection.GenericData(; counter_values = this_counter)
            countername = DataCollection.GenericParameters(; counter_name = name)

            DataCollection.addrow!(
                df,
                dc_data,
                countername,
                parameters,
                sampletime_parameters;
                cols = :union,
            )
        end

        # Save the fixed events
        core_counter_values = data.core
        deltas = CounterTools.aggregate.(diff(counter_values))

        for (i, enum) in enumerate(core_fixed)
            this_counter = getindex.(deltas, i)
            dc_data = DataCollection.GenericData(; counter_values = this_counter)
            countername = DataCollection.GenericParameters(; counter_name = string(enum))

            DataCollection.addrow!(
                df,
                dc_data,
                countername,
                parameters,
                sampletime_parameters;
                cols = :union,
                force = true,
            )
        end

        DataCollection.save(df, database())
    end

    ### Collect heap data
    record = heap_record(fex, parameters.backend)
    df = DataCollection.load(database())
    DataCollection.addrow!(
        df,
        DataCollection.GenericData(; heap_record = record),

        # Need to add this to get around a limitation in DataCollection.jl
        DataCollection.GenericParameters(; record = "heap record"),
        parameters,
        sampletime_parameters;
        cols = :union,
    )

    DataCollection.save(df, database())
    return fex
end

function heap_record(fex, backend)
    function_data = AutoTM.Profiler.FunctionData(fex.ex.ngraph_function, backend)
    times = nGraph.get_performance(fex.ex)

    # Record metadata about each tensor that will be used for plot generation.
    tensor_records = map(collect(AutoTM.Profiler.tensors(function_data))) do tensor
        users = AutoTM.Profiler.users(tensor)
        return Dict(
             "name" => nGraph.name(tensor),
             "users" => nGraph.name.(users),
             "user_indices" => getproperty.(users, :index),
             "sizeof" => sizeof(tensor),
             "offset" => Int(AutoTM.Profiler.getoffset(tensor)),
        )
    end

    # Collect metadata on each node.
    node_records = map(AutoTM.Profiler.nodes(function_data)) do node
        outputs = nGraph.name.(AutoTM.Profiler.outputs(node))
        inputs = nGraph.name.(AutoTM.Profiler.inputs(node))
        time = get(times, nGraph.name(node), 0)
        delete!(times, nGraph.name(node))
        return Dict(
            "name" => nGraph.name(node),
            "inputs" => inputs,
            "outputs" => outputs,
            "time" => time,
        )
    end

    @assert isempty(times)

    # Combine together to the final record that will be saved.
    record = Dict(
        "tensors" => tensor_records,
        "nodes" => node_records,
    )

    return record
end

#####
##### Top Level Experiments
#####

function testrun()
    parameters = RunParameters(
        f = AutoTM.Experiments.test_vgg(),
        optimizer = AutoTM.Optimizer.Optimizer2LM(),
        cache = "nocache",
        backend = nGraph.Backend("CPU"),
        mode = "2LM"
    )

    run(parameters)
end

function experiment2lm(f)
    optimizer = AutoTM.Optimizer.Optimizer2LM()
    cache = "nocache"
    backend = nGraph.Backend("CPU")

    parameters = RunParameters(
        f = f,
        optimizer = optimizer,
        cache = cache,
        backend = backend,
        mode = "2LM",
    )

    return run(parameters)
end

function experiment1lm(f; limit = 185_000_000_000)
    optimizer = AutoTM.Optimizer.Synchronous(limit)
    cache = AutoTM.Profiler.CPUKernelCache(AutoTM.Experiments.CPU_CACHE)
    backend = nGraph.Backend("CPU")

    parameters = RunParameters(
        f = f,
        optimizer = optimizer,
        cache = cache,
        backend = backend,
        mode = "1LM",
    )

    return run(parameters)
end

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

# This run collects the timestamp of nGraph kernels to correlate kernel execution with
# counter values
function snooped_run(parameters::RunParameters)
    # Step 1 - create the desired AutoTM function and run it once to warm everything up.
    fex = AutoTM.Optimizer.factory(
        parameters.backend,
        parameters.f,
        parameters.optimizer;
        cache = parameters.cache
    ) |> maybeunwrap
    @time fex()

    sampletime_parameters = DataCollection.GenericParameters(; sampletime = Millisecond(100))

    # Now, take IMC measurements
    for event_name_pairs in GROUPED_IMC_COUNTERS
        # Reset counters to generate a timeline after data collection.
        nGraph.reset_counters(fex.ex)

        events = first.(event_name_pairs)
        measurements = MattDaemon.@measurements (
            pretime = SystemSnoop.Timestamp(),
            imc = CounterTools.IMCMonitor(events, socket()),
            posttime = SystemSnoop.Timestamp(),
        )

        payload = MattDaemon.ServerPayload(Millisecond(100), measurements)

        # Capture output stream to a buffer
        local data
        local runtime

        stdout_file = joinpath(dirname(@__DIR__), "stdout")
        open(stdout_file; write = true) do io
            redirect_stdout(io) do
                data, _, runtime = MattDaemon.run(fex, payload, port(); sleeptime = Second(1))
            end
        end

        # Now, associate time stamps with each kernel
        kernel_to_time = OrderedDict{String,Int}()
        open(stdout_file) do io
            # Each line has an entry
            #
            # "nodename timestamp"
            #
            # So it's pretty easy to parse
            for ln in eachline(io)
                sp = split(ln)
                nodename = String(sp[1])
                timestamp = parse(Int, sp[2])
                kernel_to_time[nodename] = timestamp
            end
        end

        # Open up the DataFrame where our data is being collected.
        df = DataCollection.load(database_snoop())

        #####
        ##### Node Time Processing
        #####

        # Adjust the c++ time to Julia time, then to a DateTime object
        jt = TimeSync.juliatime()
        ct = TimeSync.cxxtime()
        adjustment = jt - ct
        kernel_to_time = OrderedDict(
            k => unix2datetime((v + adjustment) / 1E6) for (k,v) in kernel_to_time
        )

        # Save the timings for each kernel
        node_params = DataCollection.GenericParameters(;
            counter_group = last.(event_name_pairs),
        )

        node_data = DataCollection.GenericData(;
            kernel_to_time = kernel_to_time
        )

        DataCollection.addrow!(
            df,
            node_data,
            node_params,
            parameters,
            sampletime_parameters;
            cols = :union
        )

        #####
        ##### Counter Processing
        #####

        # Take counter differences and aggregate counters across all IMC Channels
        counter_values = data.imc
        deltas = CounterTools.aggregate.(diff(counter_values))

        # Get timestamps in multiple formats for later review
        timestamps = data.pretime

        for (i, name) in enumerate(last.(event_name_pairs))
            # Extract the values for this counter
            this_counter = getindex.(deltas, i)
            @show name

            dc_data = DataCollection.GenericData(;
                counter_values = this_counter,
                timestamps = timestamps,
            )

            countername = DataCollection.GenericParameters(; counter_name = name)

            DataCollection.addrow!(
                df,
                dc_data,
                countername,
                parameters,
                sampletime_parameters;
                cols = :union,
            )

        end

        DataCollection.save(df, database_snoop())
    end

    ### Collect heap data
    record = heap_record(fex, parameters.backend)
    df = DataCollection.load(database_snoop())
    DataCollection.addrow!(
        df,
        DataCollection.GenericData(; heap_record = record),

        # Need to add this to get around a limitation in DataCollection.jl
        DataCollection.GenericParameters(; record = "heap record"),
        parameters,
        sampletime_parameters;
        cols = :union,
    )

    DataCollection.save(df, database_snoop())
    return fex
end

function experiment2lm_snooped(f)
    enable_timestamps()
    optimizer = AutoTM.Optimizer.Optimizer2LM()
    cache = "nocache"
    backend = nGraph.Backend("CPU")

    parameters = RunParameters(
        f = f,
        optimizer = optimizer,
        cache = cache,
        backend = backend,
        mode = "2LM",
    )

    return snooped_run(parameters)
end

end # module

max_stars_repo_name: darchr/AutoTM

max_issues_repo_name: darchr/AutoTM

max_forks_repo_name: darchr/AutoTM


text:
include("matrix.jl")

JETON_1 = "O"
JETON_2 = "X"

game = matrix::make_matrix(2, [3, 3], 0)
mutable struct P a; end
mutable struct p2D x; y; end
won = P(false)

function print_game()
    println()
    for i = 0 : 2
        for j = 0 : 2
            if game[i][j] == 0
                print(" ")
            elseif game[i][j] == 1
                print(JETON_1)
            elseif game[i][j] == 2
                print(JETON_2)
            else
                assert false
            end
            if j <= 1
                print("|")
            end
        end
        println()
        if i <= 1
            println("-----")
        end
    end
end

function jouer(numero :: Int64)
    indice = P(-1)
    dowhile
        println("Joueur ", numero)
        i = input_int() - 1
        indice.a = i
        if (i >= 0) && (i <= 8) && game[2 - div(i, 3)][i % 3] > 0
            println("Cannot play here")
            i = -1
        end
        (i < 0) || (i > 8)
    end
    i = indice.a
    game[2 - div(i, 3)][i % 3] = numero
    return p2D(2 - div(i, 3), i % 3)
end

function win(val :: Int64)
    println("Vainqueur : joueur ", val)
    won.a = true
end

function verifier(x :: Int64, y :: Int64)
    val = game[x][y]
    if (game[x][(y + 1) % 3] == val) && (game[x][(y + 2) % 3] == val)
        win(val)
    elseif (game[(x + 1) % 3][y] == val) && (game[(x + 2) % 3][y] == val)
        win(val)
    elseif (x == y) && (game[(x + 1) % 3][(y + 1) % 3] == val) && (game[(x + 2) % 3][(y + 2) % 3] == val)
        win(val)
    elseif ((x + y) % 3 == 2) && (game[(x - 1) % 3][(y + 1) % 3] == val) && (game[(x - 2) % 3][(y + 2) % 3] == val)
        win(val)
    end
end

function play()
    while !(won.a)
        joue = jouer(1)
        print_game()
        verifier(joue.x, joue.y)
        if !(won.a)
            joue = jouer(2)
            print_game()
            verifier(joue.x, joue.y)
        end
    end
end


text:
@testset "ReachNumber" begin
  heuristic(state, goal) = abs(goal - state)

  @testset "ReachNumberGoalDefaultHeuristic" begin
    start = 0
    goal = 10

    getneighbours(state) = [state - 1, state + 1]

    res = astar(getneighbours, start, goal)
    @test res.status == :success
    @test res.path == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test res.cost == 10
  end

  @testset "ReachNumberGoal" begin
    start = 0
    goal = 10

    getneighbours(state) = [state - 1, state + 1]

    res = astar(getneighbours, start, goal; heuristic)
    @test res.status == :success
    @test res.path == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test res.cost == 10
  end

  @testset "ReachNumberGoal Negative" begin
    start = 10
    goal = 0

    getneighbours(state) = [state - 1, state + 1]

    res = astar(getneighbours, start, goal; heuristic)
    @test res.status == :success
    @test res.path == [10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]
    @test res.cost == 10
  end

  @testset "Timeout" begin
    start = 1
    goal = 3

    getneighbours(state) = [state + 2, state + 1]

    res = astar(getneighbours, start, goal; heuristic, timeout = -1)# timeout immediately
    @test res.status == :timeout
    @test res.path == [1]
    @test res.cost == 0
  end

  @testset "Multiple paths" begin
    start = 1
    goal = 4

    getneighbours(state) = [state + 2, state + 1, state - 1]

    res = astar(getneighbours, start, goal; heuristic)
    @test res.status == :success
    @test res.path == [1, 3, 4]
    @test res.cost == 2
  end

  @testset "Cost argument" begin
    start = 0
    goal = 10

    getneighbours(state) = [state - 1, state + 1]
    cost(a, b) = a % 2 == 0 ? abs(b - a) : 2 * abs(b - a)

    res = astar(getneighbours, start, goal; heuristic, cost)
    @test res.status == :success
    @test res.path == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
    @test res.cost == 15
  end

  @testset "Multiple paths, a double step costs 3, exercise the new path" begin
    start = 1
    goal = 4

    getneighbours(state) = [state + 2, state + 1, state - 1]
    cost(a, b) = b - a == 2 ? 3 : 1

    res = astar(getneighbours, start, goal; heuristic, cost)
    @test res.status == :success
    @test res.path == [1, 2, 3, 4]
    @test res.cost == 3
  end
end


text:
@testset "utils" begin
    using KernelFunctions: vec_of_vecs, ColVecs, RowVecs
    rng, N, D = MersenneTwister(123456), 10, 4
    x = randn(rng, N)
    X = randn(rng, D, N)
    v = randn(rng, D)
    w = randn(rng, N)

    @testset "VecOfVecs" begin
        @test vec_of_vecs(X, obsdim = 2) == ColVecs(X)
        @test vec_of_vecs(X, obsdim = 1) == RowVecs(X)
    end
    # Test Matrix data sets.
    @testset "ColVecs" begin
        DX = ColVecs(X)
        @test DX == DX
        @test size(DX) == (N,)
        @test length(DX) == N
        @test getindex(DX, 5) isa AbstractVector
        @test getindex(DX, 5) == X[:, 5]
        @test getindex(DX, 1:2:6) isa ColVecs
        @test getindex(DX, 1:2:6) == ColVecs(X[:, 1:2:6])
        @test getindex(DX, :) == ColVecs(X)
        @test eachindex(DX) == 1:N
        @test first(DX) == X[:, 1]
        DX[2] = v
        @test DX[2] == v
        @test X[:, 2] == v

        Y = randn(rng, D, N + 1)
        DY = ColVecs(Y)
        @test KernelFunctions.pairwise(SqEuclidean(), DX) ≈ pairwise(SqEuclidean(), X; dims=2)
        @test KernelFunctions.pairwise(SqEuclidean(), DX, DY) ≈ pairwise(SqEuclidean(), X, Y; dims=2)
        K = zeros(N, N)
        KernelFunctions.pairwise!(K, SqEuclidean(), DX)
        @test K ≈ pairwise(SqEuclidean(), X; dims=2)
        K = zeros(N, N + 1)
        KernelFunctions.pairwise!(K, SqEuclidean(), DX, DY)
        @test K ≈ pairwise(SqEuclidean(), X, Y; dims=2)

        let
            @test Zygote.pullback(ColVecs, X)[1] == DX
            DX, back = Zygote.pullback(ColVecs, X)
            @test back((X=ones(size(X)),))[1] == ones(size(X))

            @test Zygote.pullback(DX->DX.X, DX)[1] == X
            X_, back = Zygote.pullback(DX->DX.X, DX)
            @test back(ones(size(X)))[1].X == ones(size(X))
        end
    end
    @testset "RowVecs" begin
        DX = RowVecs(X)
        @test DX == DX
        @test size(DX) == (D,)
        @test length(DX) == D
        @test getindex(DX, 2) isa AbstractVector
        @test getindex(DX, 2) == X[2, :]
        @test getindex(DX, 1:3) isa RowVecs
        @test getindex(DX, 1:3) == RowVecs(X[1:3, :])
        @test getindex(DX, :) == RowVecs(X)
        @test eachindex(DX) == 1:D
        @test first(DX) == X[1, :]
        DX[2] = w
        @test DX[2] == w
        @test X[2, :] == w

        Y = randn(rng, D + 1, N)
        DY = RowVecs(Y)
        @test KernelFunctions.pairwise(SqEuclidean(), DX) ≈ pairwise(SqEuclidean(), X; dims=1)
        @test KernelFunctions.pairwise(SqEuclidean(), DX, DY) ≈ pairwise(SqEuclidean(), X, Y; dims=1)
        K = zeros(D, D)
        KernelFunctions.pairwise!(K, SqEuclidean(), DX)
        @test K ≈ pairwise(SqEuclidean(), X; dims=1)
        K = zeros(D, D + 1)
        KernelFunctions.pairwise!(K, SqEuclidean(), DX, DY)
        @test K ≈ pairwise(SqEuclidean(), X, Y; dims=1)

        let
            @test Zygote.pullback(RowVecs, X)[1] == DX
            DX, back = Zygote.pullback(RowVecs, X)
            @test back((X=ones(size(X)),))[1] == ones(size(X))

            @test Zygote.pullback(DX->DX.X, DX)[1] == X
            X_, back = Zygote.pullback(DX->DX.X, DX)
            @test back(ones(size(X)))[1].X == ones(size(X))
        end
    end
    @testset "input checks" begin
        D = 3; D⁻ = 2
        N1 = 2; N2 = 3
        x = [rand(rng, D) for _ in 1:N1]
        x⁻ = [rand(rng, D⁻) for _ in 1:N1]
        y = [rand(rng, D) for _ in 1:N2]
        xx = [rand(rng, D, D) for _ in 1:N1]
        xx⁻ = [rand(rng, D, D⁻) for _ in 1:N1]
        yy = [rand(rng, D, D) for _ in 1:N2]

        @test KernelFunctions.dim("string") == 0
        @test KernelFunctions.dim(["string", "string2"]) == 0
        @test KernelFunctions.dim(rand(rng, 4)) == 1
        @test KernelFunctions.dim(x) == D

        @test_nowarn KernelFunctions.validate_inplace_dims(zeros(N1, N2), x, y)
        @test_throws DimensionMismatch KernelFunctions.validate_inplace_dims(zeros(N1, N1), x, y)
        @test_throws DimensionMismatch KernelFunctions.validate_inplace_dims(zeros(N1, N2), x⁻, y)
        @test_nowarn KernelFunctions.validate_inplace_dims(zeros(N1, N1), x)
        @test_nowarn KernelFunctions.validate_inplace_dims(zeros(N1), x)
        @test_throws DimensionMismatch KernelFunctions.validate_inplace_dims(zeros(N2), x)

        @test_nowarn KernelFunctions.validate_inputs(x, y)
        @test_throws DimensionMismatch KernelFunctions.validate_inputs(x⁻, y)

        @test_nowarn KernelFunctions.validate_inputs(xx, yy)
        @test_nowarn KernelFunctions.validate_inputs(xx⁻, yy)
    end
end


text:
""" Combinatorial maps and related structures, as C-sets.

In topological graph theory and graph drawing, an **embedded graph** is a
combinatorial structure representing a graph embedded in an (oriented) surface,
up to equivalence under (orientation-preserving) homeomorphism. This module
defines data structures for rotation systems, combinatorial maps, and other
combinatorial objects describing embedded graphs.
"""
module CombinatorialMaps
export σ, α, ϕ, trace_vertices, trace_edges, trace_faces,
  AbstractRotationGraph, RotationGraph, add_corolla!, pair_half_edges!,
  AbstractRotationSystem, RotationSystem,
  AbstractCombinatorialMap, CombinatorialMap

using Catlab, Catlab.CategoricalAlgebra.CSets, Catlab.Graphs
using Catlab.Graphs.BasicGraphs: TheoryHalfEdgeGraph
using Catlab.Permutations: cycles

# General properties
####################

""" Vertex permutation of combinatorial map or similar structure.
"""
σ(x::ACSet, args...) = subpart(x, args..., :σ)

""" Edge permutation of combinatorial map or similar structure.
"""
α(x::ACSet, args...) = subpart(x, args..., :α)

""" Face permutation of combinatorial map or similar structure.
"""
ϕ(x::ACSet, args...) = subpart(x, args..., :ϕ)

""" Trace vertices of combinatorial map or similar, returning a list of cycles.
"""
trace_vertices(x::ACSet) = cycles(σ(x))

""" Trace edges of combinatorial map or similar, return a listing of cycles.

Usually the cycles will be pairs of half edges but in a hypermap the cycles can
be arbitrary.
"""
trace_edges(x::ACSet) = cycles(α(x))

""" Trace faces of combinatorial map or similar, returning list of cycles.
"""
trace_faces(x::ACSet) = cycles(ϕ(x))

# Rotation graphs
#################

@present TheoryRotationGraph <: TheoryHalfEdgeGraph begin
  σ::Hom(H,H)

  compose(σ, vertex) == vertex
end

@abstract_acset_type AbstractRotationGraph <: AbstractHalfEdgeGraph
@acset_type RotationGraph(TheoryRotationGraph,
                          index=[:vertex]) <: AbstractRotationGraph

α(g::AbstractRotationGraph) = inv(g)
ϕ(g::AbstractRotationGraph) = sortperm(inv(g)[σ(g)]) # == (σ ⋅ inv)⁻¹

""" Add corolla to rotation graph, rotation system, or similar structure.

A *corolla* is a vertex together with its incident half-edges, the number of
which is its *valence*. The rotation on the half-edges is the consecutive one
induced by the half-edge part numbers.
"""
function add_corolla!(g::AbstractRotationGraph, valence::Int; kw...)
  v = add_vertex!(g; kw...)
  n = nparts(g, :H)
  add_parts!(g, :H, valence; vertex=v, σ=circshift((n+1):(n+valence), -1))
end

""" Pair together half-edges into edges.
"""
pair_half_edges!(g::AbstractRotationGraph, h, h′) =
  set_subpart!(g, [h; h′], :inv, [h′; h])

# Rotation systems
##################

@present TheoryRotationSystem(FreeSchema) begin
  H::Ob
  σ::Hom(H,H)
  α::Hom(H,H)

  compose(α, α) == id(H)
end

@abstract_acset_type AbstractRotationSystem
@acset_type RotationSystem(TheoryRotationSystem) <: AbstractRotationSystem

# ϕ == (σ⋅α)⁻¹ == α⁻¹ ⋅ σ⁻¹
ϕ(sys::AbstractRotationSystem) = sortperm(α(sys)[σ(sys)])

function add_corolla!(sys::AbstractRotationSystem, valence::Int)
  n = nparts(sys, :H)
  add_parts!(sys, :H, valence; σ=circshift((n+1):(n+valence), -1))
end

pair_half_edges!(sys::AbstractRotationSystem, h, h′) =
  set_subpart!(sys, [h; h′], :α, [h′; h])

# Combinatorial maps
####################

@present TheoryHypermap(FreeSchema) begin
  H::Ob
  σ::Hom(H,H)
  α::Hom(H,H)
  ϕ::Hom(H,H)

  compose(σ, α, ϕ) == id(H)
end

@present TheoryCombinatorialMap <: TheoryHypermap begin
  compose(α, α) == id(H)
end

@abstract_acset_type AbstractCombinatorialMap
@acset_type CombinatorialMap(TheoryCombinatorialMap) <: AbstractCombinatorialMap

# TODO: What kind of interface should we have for maps and hypermaps?

end


text:
#
#  - split off the linear transformation
#  - remove the abstract oneparticle basis crap and incorporate it into
#    the Product1pBasis -> renamed as OnepBasis
#  - spec simply specifies in what order the wrapped basis generates the
#    basis functions, basically don't allow doing it lazy anymore
#  - P_qa(r) is actually P_q(r, a) so the only issue with this is
#    if we want different basis length for different species, but even then
#    we could just fill in zeros
#  - Linear1pTransform: just a linear mapping
#

# MORE NOTES
#  - consider leaving FVAL in here for now, but move it out later...
#  - changed getval(fval, X) to evaluate(fval, X)
#  - can we get away with only proviging gradients w.r.t. the
#    symbols used in this component, or do we need the whole thing?

import NamedTupleTools
using NamedTupleTools: namedtuple

@doc raw"""
`struct B1pComponent:` Wraps a one-particle basis component.

`basis` is a structure that can compute a basis, i.e. a vector of scalars,
real or complex. each element of this vector is "specified" by a NamedTuple,
stored in `spec`. E.g., for an `Rn` type basis the spec is would be
```julia
   spec = [ (n = 1,), (n=2, ), ... ]
```
while for a Ylm type basis it would be
```julia
   spec = [ (l=0, m=0), (l=1, m=-1), (l=1, m=0), ... ]
```

There is no convenience constructor, the `B1pComponent` must be constructed
"by hand", for which several wrapper functions are written.
"""
struct B1pComponent{ISYMS, TT, TB, FVAL}
   basis::TB
   fval::FVAL
   spec::Vector{NamedTuple{ISYMS, TT}}
   degrees::Vector{Int}
   label::String
   meta::Dict{String,Any}
   # ------------ derived fields
   invspec::Dict{NamedTuple{ISYMS, TT}, Int}
   # todo - fields for temporary arrays ...
end

function B1pComponent(basis, fval, spec::AbstractVector{<: NamedTuple},
                      degrees::Vector{Int},
                      label::AbstractString,
                      meta = Dict{String, Any}())
   spec1 = collect(spec)
   invspec = Dict{eltype(spec1), Int}()
   for (i, b) in enumerate(spec)
      invspec[b] = i
   end
   return B1pComponent(basis, fval, spec, degrees, label, meta, invspec)
end

# -------------- management of the basis specification, in particular the
#                interaction with Product1pBasis

getlabel(basis::B1pComponent) = basis.label

Base.length(basis::B1pComponent) = length(basis.spec)

_idxsyms(basis::B1pComponent{ISYMS}) where {ISYMS} = ISYMS

get_spec(basis::B1pComponent) = copy(basis.spec)

get_spec(basis::B1pComponent, i::Integer) = basis.spec[i]

# TODO - I don't think this is needed anymore, but the sparsification will
#        have to be moved in here and something be done about that...
# function set_spec!(basis::XScal1pBasis, spec)
#    basis.spec = identity.(spec)
#    basis.coeffs = zeros(eltype(basis.coeffs), length(spec), length(basis.P))
#    return basis
# end

function isadmissible(b::NamedTuple{BSYMS}, basis::B1pComponent) where {BSYMS}
   ISYMS = _idxsyms(basis)
   # this is an assert since it should ALWAYS be true, if not there is a bug
   @assert all(sym in BSYMS for sym in ISYMS)
   # project to the ISYMS and check it is in the b1pcomponent
   return haskey(basis.invspec, b[ISYMS])
end

# TODO - LATER
# function sparsify!(basis::XScal1pBasis, spec)
#    # spec is the part of the basis we keep
#    inds = Vector{Int}(undef, length(spec))
#    for (ib, b) in enumerate(spec)
#       iold = findall(isequal(b), basis.spec)
#       @assert length(iold) == 1
#       inds[ib] = iold[1]
#    end
#    # keep the original order of the basis since we assume it is ordered
#    # by some sensible notion of degree.
#    p = sortperm(inds)
#    basis.spec = spec[p]
#    basis.coeffs = basis.coeffs[p, :]
#    return basis
# end

symbols(basis::B1pComponent) = [ _idxsyms(basis)... ]

function indexrange(basis::B1pComponent)
   ISYMS = _idxsyms(basis)
   minidx = Dict([sym => minimum(b[sym] for b in basis.spec) for sym in ISYMS]...)
   maxidx = Dict([sym => maximum(b[sym] for b in basis.spec) for sym in ISYMS]...)
   return Dict([sym => minidx[sym]:maxidx[sym] for sym in ISYMS]...)
end

# this is needed to generate the product 1p basis - it returns the
# index of the 1p basis function component used in the 1p basis function
# specified by b
function get_index(basis::B1pComponent, b::NamedTuple)
   ISYMS = _idxsyms(basis)
   b1 = b[ISYMS]
   if !haskey(basis.invspec, b1)
      error("B1pComponent ($(basis.label) : can't find $(b1) in spec")
   end
   return basis.invspec[b1]
end

# -------------- degree calculations

# this should be revisited, doesn't feel clean yet

function degree(b::NamedTuple, basis::B1pComponent)
   ISYMS = _idxsyms(basis)
   idx = basis.invspec[ b[ISYMS] ]
   return basis.degrees[idx]
end

function degree(b::NamedTuple, basis::B1pComponent, weight::Dict)
   ISYMS = _idxsyms(basis)
   return weight[ISYMS[1]] * degree(b, basis)
end

# --------------- IO operations

==(P1::B1pComponent, P2::B1pComponent) = (
      (P1.basis == P2.basis) && (P1.spec == P2.spec) &&
      (P1.label == P2.label) && (P1.fval == P2.fval) &&
      (P1.degrees == P2.degrees) )

function write_dict(basis::B1pComponent)
   ISYMS = _idxsyms(basis)
   return Dict("__id__" => "ACE_B1pComponent",
                 "syms" => [ string.(ISYMS) ...],
                "basis" => write_dict(basis.basis),
                 "fval" => write_dict(basis.fval),
                 "spec" => convert.(Dict, basis.spec),
              "degrees" => basis.degrees,
                "label" => basis.label )
end

function read_dict(::Val{:ACE_B1pComponent}, D::Dict)
   basis = read_dict(D["basis"])
   ISYMS = tuple(Symbol.(D["syms"])...)
   spec = NamedTuple{ISYMS}.(namedtuple.(D["spec"]))
   fval = read_dict(D["fval"])
   return B1pComponent(basis, fval, spec, Int.(D["degrees"]), D["label"])
end

function show(io::IO, basis::B1pComponent)
   vsyms = Transforms.get_symbols(basis.fval)
   strvsyms = filter(!isequal(':'), "$vsyms")
   if length(vsyms) == 1
      # this doesn't work with unicode!?!?!?
      # strvsyms = strvsyms[1:end-2] * ")"
      strvsyms = prod(split(strvsyms, ','))
   end
   print(io, "B1pComponent( $(basis.label)$(strvsyms)")
   for (key, rg) in indexrange(basis)
      print(io, ", $key = $(rg)")
   end
   print(io, ")")
end

# ------------------- preparation for evaluation and managing temporaries

valtype(basis::B1pComponent) = valtype(basis.basis)

valtype(basis::B1pComponent, X::AbstractState) =
      valtype(basis.basis, evaluate(basis.fval, X))

function gradtype(basis::B1pComponent, X::AbstractState)
   x = evaluate(basis.fval, X)
   # gradient type of the inner basis which knows nothing about states
   TDB = gradtype(basis.basis, x)
   # now we need to incorporate the grad type of fval itself
   # dx = evaluate_d(basis.fval, X)
   # the gradient will be a product of a TDB times a TDVAL
   return Transforms.grad_type_dP(TDB, basis.fval, X)
      # dstate_type(valtype(basis, X), X)
end

acquire_B!(basis::B1pComponent, args...) =
         Vector{valtype(basis, args...)}(undef, length(basis))

release_B!(basis::B1pComponent, args...) = nothing

acquire_dB!(basis::B1pComponent, X::AbstractState) =
            Vector{gradtype(basis, X)}(undef, length(basis))

release_dB!(basis::B1pComponent, args...) = nothing

# ------------------------ Evaluation code
#                          this is basically an interface for the inner basis

evaluate(basis::B1pComponent, X::AbstractState) =
         evaluate(basis.basis, evaluate(basis.fval, X))

# evaluate!(B, basis::B1pComponent, X::AbstractState) =
#       evaluate!(B, basis.basis, )

function evaluate_d(basis::B1pComponent, X::AbstractState)
   B, dB = evaluate_ed(basis, X)
   release!(B)
   return dB
end

function evaluate_ed(basis::B1pComponent, X::AbstractState)
   x = evaluate(basis.fval, X)
   B, dP = evaluate_ed(basis.basis, x)
   dB = rrule_evaluate(dP, basis.fval, X)
   release!(dP)
   return B, dB
end

# this one we probably only need for training so ...
#   - we can relax the efficiency a bit
#   - we actually never need the evaluate_dd but only the
#     associated backpropagation operation which may be a bit simpler
#     to implement.
# so we will likely remove this entirely.
# evaluate_dd(basis::B1pComponent, X::AbstractState)

# -------------- AD codes

#=   *** TODO
import ChainRules: rrule, ZeroTangent, NoTangent

function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState,
                         w::AbstractVector{<: Number})
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   a = _rrule_evaluate(basis.P, x, real.(w))
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

rrule(::typeof(evaluate), basis::Scal1pBasis, X::AbstractState) =
                  evaluate(basis, X),
                  w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w))

function _rrule_evaluate_d(basis::Scal1pBasis, X::AbstractState,
                           w::AbstractVector)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   w1 = [ _val(w, basis) for w in w ]
   a = _rrule_evaluate_d(basis.P, x, w1)
   TDX = ACE.dstate_type(a, X)
   return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
end

function rrule(::typeof(evaluate_d), basis::Scal1pBasis, X::AbstractState)
   @assert _varidx(basis) == 1
   x = _val(X, basis)
   dB_ = evaluate_d(basis.P, x)
   TDX = dstate_type(valtype(basis, X), X)
   dB = [ TDX( NamedTuple{(_varsym(basis),)}( (dx,) ) )  for dx in dB_ ]
   return dB,
          w -> (NoTangent(), NoTangent(), _rrule_evaluate_d(basis, X, w))
end

=#


text:
# p10.jl - polynomials and corresponding equipotential curves

N = 16; clf();
xx = -1.01:.005:1.01;
for i = 1:2
    i==1 && ( (s,x) = ("equispaced points", -1 + 2*(0:N)/N) );
    i==2 && ( (s,x) = ("Chebyshev points", cos.(pi*(0:N)/N)) );
    p = poly(x);

    # Plot p(x) over [-1,1]:
    xx = -1:.005:1; pp = p(xx);
    subplot(2,2,2*i-1);
    plot(x,0*x,"k.",markersize=6);
    plot(xx,pp);  grid(true);
    xticks(-1:.5:1);  title(s);

    # Plot equipotential curves:
    subplot(2,2,2*i)
    plot(real(x),imag(x),".",markersize=6);
    axis([-1.4,1.4,-1.12,1.12]);
    xx = -1.4:.02:1.4; yy = -1.12:.02:1.12;
    zz = xx' .+ 1im*yy;
    pp = p(zz); levels = 10.0.^(-4:0);
    contour(xx,yy,abs.(pp),levels,colors="k"); title(s);
end


text:
module InterleavedImages

using ImageMetadata, AxisArrays
const axes = Base.axes #for name conflict with AxisArrays

using CachedArrays #only for match_axisspacing

import Base: size, getindex, setindex!

export InterleavedImage

include("interleaved.jl")

end # module


text:
module GModel

using DocStringExtensions

using Random
using Distributions
using LinearAlgebra
using FFTW

export run_G
export run_G_ensemble
export lorenz_forward

"""
    GSettings{FT<:AbstractFloat, KT, D}

Structure to hold all information to run the forward model G

# Fields
$(DocStringExtensions.FIELDS)
"""
struct GSettings{FT <: AbstractFloat, KT, D}
    "model output"
    out::Array{FT, 1}
    "time period over which to run the model, e.g., `(0, 1)`"
    tspan::Tuple{FT, FT}
    "x domain"
    x::Array{FT, 1}
    "y domain"
    y::Array{FT, 1}
end

struct LSettings
    # Stationary or transient dynamics
    dynamics::Int32
    # G model statistics type
    stats_type::Int32
    # Integration time start
    t_start::Float64
    # Data collection length
    T::Float64
    # Integration length
    Ts::Float64
    # Duration of polynomial fit, number of samples of Ts to fit over
    Tfit::Float64
    # Initial perturbation
    Fp::Array{Float64}
    # Number of longitude steps
    N::Int32
    # Timestep
    dt::Float64
    # Simulation end time
    tend::Float64
    # For stats_type=2, number of frequencies to consider
    kmax::Int32
    # Should ω be learned using CES?
    ω_fixed::Bool
    # Truth ω
    ω_true::Float64
end

struct LParams
    # Mean forcing
    F::Float64
    # Forcing frequency for transient terms
    ω::Float64
    # Forcing amplitude for transient terms
    A::Float64
end

"""
    run_G_ensemble(params::Array{FT, 2},
                   settings::GSettings{FT},
                   update_params,
                   moment,
                   get_src;
                   rng_seed=42) where {FT<:AbstractFloat}

Run the forward model G for an array of parameters by iteratively
calling run_G for each of the N_ensemble parameter values.
Return g_ens, an array of size N_data x N_ensemble, where
g_ens[:,j] = G(params[:,j])

 - `params` - array of size N_parameters x N_ensemble containing the
              parameters for which G will be run
 - `settings` - a GSetttings struct

"""
function run_G_ensemble(params::Array{FT, 2}, settings::LSettings, rng_seed = 42) where {FT <: AbstractFloat}
    # Initilize ensemble
    N_ens = size(params, 2) # params is N_params x N_ens
    if settings.stats_type == 1
        nd = 2
    elseif settings.stats_type == 2
        nd = 1 + (2 * settings.kmax)
    elseif settings.stats_type == 3
        nd = 3
    elseif settings.stats_type == 4
        nd = Int64(2 * (settings.T / settings.Ts / settings.Tfit))
    elseif settings.stats_type == 5
        nd = Int64(settings.T / settings.Ts / settings.Tfit)
    end
    g_ens = zeros(nd, N_ens)
    # Lorenz parameters
    #Fp = rand(Normal(0.0, 0.01), N); # Initial perturbation
    F = params[1, :] # Forcing
    if settings.dynamics == 2
        A = params[2, :] # Transience amplitude
    else
        A = F .* 0.0 # Set to zero, it is not used
    end
    if settings.ω_fixed == false
        ω = params[3, :] # Transience frequency
    end

    Random.seed!(rng_seed)
    for i in 1:N_ens
        # run the model with the current parameters, i.e., map θ to G(θ)
        if settings.ω_fixed == false
            lorenz_params = GModel.LParams(F[i], ω[i], A[i])
        elseif settings.ω_fixed == true
            lorenz_params = GModel.LParams(F[i], settings.ω_true, A[i])
        end
        g_ens[:, i] = lorenz_forward(settings, lorenz_params)
    end

    return g_ens
end

# Forward pass of the Lorenz 96 model
# Inputs: settings: structure with stats_type, t_start, T
# F: scalar forcing Float64(1), Fp: initial perturbation Float64(N)
# N: number of longitude steps Int64(1), dt: time step Float64(1)
# tend: End of simulation Float64(1), nstep:
function lorenz_forward(settings::LSettings, params::LParams)
    # run the Lorenz simulation
    xn, t = lorenz_solve(settings, params)
    # Get statistics
    gt = stats(settings, xn, t)
    return gt
end

# Solve the Lorenz 96 system
# Inputs: F: scalar forcing Float64(1), Fp: initial perturbation Float64(N)
# N: number of longitude steps Int64(1), dt: time step Float64(1)
# tend: End of simulation Float64(1), nstep:
function lorenz_solve(settings::LSettings, params::LParams)
    # Initialize
    nstep = Int32(ceil(settings.tend / settings.dt))
    xn = zeros(Float64, settings.N, nstep)
    t = zeros(Float64, nstep)
    # Initial perturbation
    X = fill(Float64(0.0), settings.N)
    X = X + settings.Fp
    # March forward in time
    for j in 1:nstep
        t[j] = settings.dt * j
        if settings.dynamics == 1
            X = RK4(X, settings.dt, settings.N, params.F)
        elseif settings.dynamics == 2
            F_local = params.F + params.A * sin(params.ω * t[j])
            X = RK4(X, settings.dt, settings.N, F_local)
        end
        xn[:, j] = X
    end
    # Output
    return xn, t

end

# Lorenz 96 system
# f = dx/dt
# Inputs: x: state, N: longitude steps, F: forcing
function f(x, N, F)
    f = zeros(Float64, N)
    # Loop over N positions
    for i in 3:(N - 1)
        f[i] = -x[i - 2] * x[i - 1] + x[i - 1] * x[i + 1] - x[i] + F
    end
    # Periodic boundary conditions
    f[1] = -x[N - 1] * x[N] + x[N] * x[2] - x[1] + F
    f[2] = -x[N] * x[1] + x[1] * x[3] - x[2] + F
    f[N] = -x[N - 2] * x[N - 1] + x[N - 1] * x[1] - x[N] + F
    # Output
    return f
end

# RK4 solve
function RK4(xold, dt, N, F)
    # Predictor steps
    k1 = f(xold, N, F)
    k2 = f(xold + k1 * dt / 2.0, N, F)
    k3 = f(xold + k2 * dt / 2.0, N, F)
    k4 = f(xold + k3 * dt, N, F)
    # Step
    xnew = xold + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
    # Output
    return xnew
end

function regression(X, y)
    beta_tilde = [ones(length(X), 1) X] \ y
    v = beta_tilde[1]
    beta = beta_tilde[2]
    return beta, v
end

function spectra(signal, Ts, t)
    # Init
    Nt = length(t)
    if mod(Nt, 2) != 0
        t = t[1:(Nt - 1)]
        signal = signal[1:(Nt - 1)]
        Nt = Nt - 1
    end
    # Remove mean
    u = signal .- mean(signal, dims = 1)
    # FFT
    F = fft(u) |> fftshift
    freqs = fftfreq(length(t), 1.0 / Ts) |> fftshift
    #k = gen_k(Nt,Ts)  |> fftshift
    T = (Nt - 1) * Ts
    A = Ts^2 / (2 * pi * T)
    uhat = A * abs.(F .^ 2)
    # Output
    mid = Int32(Nt / 2) + 1
    f = freqs[mid:Nt]
    Fp = uhat[mid:Nt]
    return f, Fp
end

###############################
## Extract statistics
###############################
function stats(settings, xn, t)
    # Define averaging indices range
    indices = findall(x -> (x > settings.t_start) && (x < settings.t_start + settings.T), t)
    # Define statistics of interest
    if settings.stats_type == 1 # Mean
        # Average in time and over longitude
        gtm = mean(mean(xn[:, indices], dims = 2), dims = 1)
        # Variance
        gtc = mean(mean((xn[:, indices] .- gtm) .^ 2, dims = 2), dims = 1)
        # Combine statistics
        gt = vcat(gtm, gtc)
    elseif settings.stats_type == 2 #What about an integral under parts of the spectrum
        # Average in time and over longitude
        gtm = mean(mean(xn[:, indices], dims = 2), dims = 1)
        # Power spectra
        # Get size
        f, Fp = spectra(xn[1, indices]', settings.dt, t[indices])
        # Compute temporal spectrum for each longitude
        Fp = zeros(size(Fp, 1), settings.N)
        for i in 1:(settings.N)
            f, Fp[:, i] = spectra(xn[i, indices]', settings.dt, t[indices])
        end
        # Average spectra over periodic directions
        Fp = dropdims(mean(Fp, dims = 2), dims = 2)
        ys = partialsortperm(Fp, 1:(settings.kmax); rev = true)
        gt = vcat(gtm, Fp[ys]..., f[ys]...)
    elseif settings.stats_type == 3 # Structure function
        # Average in time and over longitude
        gtm = mean(mean(xn[:, indices], dims = 2), dims = 1)
        # Maximum
        mxval, mxind = findmax(xn[:, indices]; dims = 2)
        mxval_out = mean(mxval, dims = 1)
        # Power spectra
        # Get size
        f, Fp = spectra(xn[1, indices]', settings.dt, t[indices])
        # Compute temporal spectrum for each longitude
        Fp = zeros(size(Fp, 1), settings.N)
        for i in 1:(settings.N)
            f, Fp[:, i] = spectra(xn[i, indices]', settings.dt, t[indices])
        end
        # Average spectra over periodic directions
        Fp = dropdims(mean(Fp, dims = 2), dims = 2)
        ys = partialsortperm(Fp, 1; rev = true)
        # Period
        T = 1.0 / f[ys]
        mxval, r = findmin(abs.(t .- T); dims = 1)
        r = r[1]
        # Structure function
        xp = xn .- gtm
        st = zeros(settings.N)
        for i in 1:(settings.N)
            st[i] = mean(
                (
                    xp[1, indices[1]:(indices[length(indices)] - r)] .-
                    xp[1, (indices[1] + r):indices[length(indices)]]
                ) .^ 2,
            )
        end
        # Combine
        gt = vcat(gtm, mean(st)..., mxval_out...)
    elseif settings.stats_type == 4 # Variance sub-samples, linear fit over window
        T = settings.T
        # Calculate variance on subsamples
        # Variances over sliding windows
        Ts = settings.Ts
        N = Int64(T / Ts) # number of sliding windows
        nt = Int64(Ts / settings.dt) - 1 # length of sliding windows
        var_slide = zeros(N)
        t_slide = zeros(nt, N)
        for i in 1:N
            var_slide[i] = var(vcat(xn[:, indices[((i - 1) * nt + 1):(i * nt)]]...))
            t_slide[:, i] = t[((i - 1) * nt + 1):(i * nt)]
        end
        # Polynomial fit over a batch of sliding windows
        Tfit = Int64(settings.Tfit)
        NT = Int64(N / Tfit)
        a1 = zeros(NT)
        a2 = zeros(NT)
        t_start = zeros(NT)
        y = zeros(Tfit, NT)
        ty = zeros(Tfit, NT)
        ym = zeros(Tfit, NT)
        for i in 1:NT
            ind_low = (i - 1) * Tfit + 1
            ind_high = i * Tfit
            t_start[i] = t_slide[1, ind_low]
            ty[:, i] = t_slide[1, ind_low:ind_high] .- t_start[i]
            a1[i], a2[i] = regression(ty[:, i], var_slide[ind_low:ind_high])
            y[:, i] = a1[i] .* ty[:, i] .+ a2[i]
            ym[i] = mean(var_slide[ind_low:ind_high])
        end
        # Combine
        gt = vcat(a1..., a2...)
    elseif settings.stats_type == 5 # Variance sub-samples, mean over window
        T = settings.T
        # Calculate variance on subsamples
        # Variances over sliding windows
        Ts = settings.Ts
        N = Int64(T / Ts) # number of sliding windows
        nt = Int64(Ts / settings.dt) - 1 # length of sliding windows
        var_slide = zeros(N)
        t_slide = zeros(nt, N)
        for i in 1:N
            var_slide[i] = var(vcat(xn[:, indices[((i - 1) * nt + 1):(i * nt)]]...))
            t_slide[:, i] = t[((i - 1) * nt + 1):(i * nt)]
        end
        # Polynomial fit over a batch of sliding windows
        Tfit = Int64(settings.Tfit)
        NT = Int64(N / Tfit)
        a1 = zeros(NT)
        a2 = zeros(NT)
        t_start = zeros(NT)
        y = zeros(Tfit, NT)
        ty = zeros(Tfit, NT)
        ym = zeros(NT)
        for i in 1:NT
            ind_low = (i - 1) * Tfit + 1
            ind_high = i * Tfit
            t_start[i] = t_slide[1, ind_low]
            ty[:, i] = t_slide[1, ind_low:ind_high] .- t_start[i]
            a1[i], a2[i] = regression(ty[:, i], var_slide[ind_low:ind_high])
            y[:, i] = a1[i] .* ty[:, i] .+ a2[i]
            ym[i] = mean(var_slide[ind_low:ind_high])
        end
        # Combine
        gt = vcat(ym...)
    else
        ArgumentError("Setting " * string(settings.stats_type) * " not implemented.")
    end
    return gt
end

## Forward pass of the Lorenz 96 model
## Inputs: settings: structure with stats_type, t_start, T
## F: scalar forcing Float64(1), Fp: initial perturbation Float64(N)
## N: number of longitude steps Int64(1), dt: time step Float64(1)
## tend: End of simulation Float64(1), nstep:
#function lorenz_forward(settings::LSettings, F)
#	# run the Lorenz simulation
#	xn, t = lorenz_solve(settings, F)
#	# Define averaging indices range
#	indices = findall(x -> (x>settings.t_start)
#			  && (x<settings.t_start+settings.T), t)
#	# Define statistics of interest
#	if settings.stats_type == 1 # Mean
#		# Average in time and over longitude
#		gtm = mean(mean(xn[:,indices], dims=2), dims=1)
#		# Variance
#		gtc = mean(mean((xn[:,indices].-gtm).^2, dims=2), dims=1)
#		# Combine statistics
#		gt = vcat(gtm, gtc)
#	elseif settings.stats_type == 2 #What about an integral under parts of the spectrum
#		# Average in time and over longitude
#		gtm = mean(mean(xn[:,indices], dims=2), dims=1)
#		# Power spectra
#		# Get size
#		f, Fp = spectra(xn[1,indices]', settings.dt, t[indices]);
#		# Compute temporal spectrum for each longitude
#		Fp = zeros(size(Fp,1), settings.N)
#		for i in 1:settings.N
#			f, Fp[:,i] = spectra(xn[i,indices]', settings.dt, t[indices]);
#		end
#		# Average spectra over periodic directions
#		Fp = dropdims(mean(Fp, dims=2), dims=2)
#		ys = partialsortperm(Fp, 1:settings.kmax; rev=true)
#		gt = vcat(gtm, Fp[ys]..., f[ys]...)
#	elseif settings.stats_type == 3 # Structure function
#		# Average in time and over longitude
#		gtm = mean(mean(xn[:,indices], dims=2), dims=1)
#		# Maximum
#		mxval, mxind = findmax(xn[:,indices]; dims=2);
#		mxval_out = mean(mxval, dims=1);
#		# Power spectra
#		# Get size
#		f, Fp = spectra(xn[1,indices]', settings.dt, t[indices]);
#		# Compute temporal spectrum for each longitude
#		Fp = zeros(size(Fp,1), settings.N)
#		for i in 1:settings.N
#			f, Fp[:,i] = spectra(xn[i,indices]', settings.dt, t[indices]);
#		end
#		# Average spectra over periodic directions
#		Fp = dropdims(mean(Fp, dims=2), dims=2)
#		ys = partialsortperm(Fp, 1; rev=true)
#		# Period
#		T = 1. / f[ys]
#		mxval, r = findmin(abs.(t.-T); dims=1); r = r[1]
#                # Structure function
#		xp = xn .- gtm
#		st = zeros(settings.N)
#		for i in 1:settings.N
#			st[i] = mean( ( xp[1,indices[1]:indices[length(indices)]-r] .-
#				       xp[1,indices[1]+r:indices[length(indices)]] ).^2 );
#		end
#		# Combine
#		gt = vcat(gtm, mean(st)..., mxval_out...)
#	else
#		ArgumentError("Setting "*string(settings.stats_type)*" not implemented.")
#	end
#	return gt
#end
#

end # module


text:
using NPZ
include("./aq.jl")

data = NPZ.npzread("/amax/home/zhangjin/Anisotropic-Quantization/data/glove/glove-100-angular-normalized.npy")

M = 50
K = 16
train_data = data[1:10000,:]
aq_codebook,aq_codes = aq_train(M, K, train_data)
npzwrite("aq_codebook.npy", aq_codebook)

aq_codebook = codebook_trans2To3(aq_codebook,M)
aq_codes,~ = aq_encode(data,aq_codebook, M)

aq_codes .-= 1
npzwrite("aq_codes.npy", aq_codes)


text:
@testset "Quotient orders" begin
  G = small_group(8, 4)
  QG = QQ[G]
  ZG = Order(QG, basis(QG); cached = false)
  N, NtoG = subgroups(G, order = 2)[1]
  I = sum([QG(NtoG(n)) for n in N]) * ZG
  OO, h = Hecke.quotient_order(ZG, I)

  for i in 1:10
    for j in 1:10
      a = rand(QG, -10:10)
      b = rand(QG, -10:10)
      @test h(a * b) == h(a) * h(b)
      @test h(a) + h(b) == h(a) + h(b)

      a = rand(algebra(OO), -10:10)
      b = rand(algebra(OO), -10:10)
      @test h(h\(a * b)) == a * b
      @test h(h\(a + b)) == a + b
    end
  end

  @test OO == Order(algebra(OO), [h(elem_in_algebra(b)) for b in basis(ZG)])
end


text:
@testset "normal.jl" begin
    @testset "SKCE: basic example" begin
        skce = UnbiasedSKCE(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel()
        )

        # only two predictions, i.e., one term in the estimator
        normal1 = Normal(0, 1)
        normal2 = Normal(1, 2)
        @test @inferred(skce([normal1, normal1], [0, 0])) ≈ 1 - sqrt(2) + 1 / sqrt(3)
        @test @inferred(skce([normal1, normal2], [1, 0])) ≈
            exp(-sqrt(2)) *
              (exp(-1 / 2) - 1 / sqrt(2) - 1 / sqrt(5) + exp(-1 / 12) / sqrt(6))
        @test @inferred(skce([normal1, normal2], [0, 1])) ≈
            exp(-sqrt(2)) * (
            exp(-1 / 2) - exp(-1 / 4) / sqrt(2) - exp(-1 / 10) / sqrt(5) +
            exp(-1 / 12) / sqrt(6)
        )
    end

    @testset "SKCE: basic example (transformed)" begin
        skce = UnbiasedSKCE(
            ExponentialKernel(; metric=Wasserstein()) ⊗
            (SqExponentialKernel() ∘ ScaleTransform(0.5)),
        )

        # only two predictions, i.e., one term in the estimator
        normal1 = Normal(0, 1)
        normal2 = Normal(1, 2)
        @test @inferred(skce([normal1, normal1], [0, 0])) ≈
            1 - 2 / sqrt(1.25) + 1 / sqrt(1.5)
        @test @inferred(skce([normal1, normal2], [1, 0])) ≈
            exp(-sqrt(2)) *
              (exp(-1 / 8) - 1 / sqrt(1.25) - 1 / sqrt(2) + exp(-1 / 18) / sqrt(2.25))
        @test @inferred(skce([normal1, normal2], [0, 1])) ≈
            exp(-sqrt(2)) * (
            exp(-1 / 8) - exp(-1 / 10) / sqrt(1.25) - exp(-1 / 16) / sqrt(2) +
            exp(-1 / 18) / sqrt(2.25)
        )
    end

    @testset "SKCE: basic properties" begin
        skce = UnbiasedSKCE(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel()
        )

        estimates = map(1:10_000) do _
            predictions = map(Normal, randn(20), rand(20))
            targets = map(rand, predictions)
            return skce(predictions, targets)
        end

        @test any(x -> x > zero(x), estimates)
        @test any(x -> x < zero(x), estimates)
        @test mean(estimates) ≈ 0 atol = 1e-4
    end

    @testset "UCME: basic example" begin
        # one test location
        ucme = UCME(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel(),
            [Normal(0.5, 0.5)],
            [1],
        )

        # two predictions
        normal1 = Normal(0, 1)
        normal2 = Normal(1, 2)
        @test @inferred(ucme([normal1, normal2], [0, 0.5])) ≈
            (
            exp(-1 / sqrt(2)) * (exp(-1 / 2) - exp(-1 / 4) / sqrt(2)) +
            exp(-sqrt(5 / 2)) * (exp(-1 / 8) - 1 / sqrt(5))
        )^2 / 4

        # two test locations
        ucme = UCME(
            ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel(),
            [Normal(0.5, 0.5), Normal(-1, 1.5)],
            [1, -0.5],
        )
        @test @inferred(ucme([normal1, normal2], [0, 0.5])) ≈
            (
            (
                exp(-1 / sqrt(2)) * (exp(-1 / 2) - exp(-1 / 4) / sqrt(2)) +
                exp(-sqrt(5 / 2)) * (exp(-1 / 8) - 1 / sqrt(5))
            )^2 +
            (
                exp(-sqrt(5) / 2) * (exp(-1 / 8) - exp(-1 / 16) / sqrt(2)) +
                exp(-sqrt(17) / 2) * (exp(-1 / 2) - exp(-9 / 40) / sqrt(5))
            )^2
        ) / 8
    end

    @testset "kernels with input transformations" begin
        nsamples = 100
        ntestsamples = 5

        # create predictions and targets
        predictions = map(Normal, randn(nsamples), rand(nsamples))
        targets = randn(nsamples)

        # create random test locations
        testpredictions = map(Normal, randn(ntestsamples), rand(ntestsamples))
        testtargets = randn(ntestsamples)

        for γ in (1.0, rand())
            kernel1 =
                ExponentialKernel(; metric=Wasserstein()) ⊗
                (SqExponentialKernel() ∘ ScaleTransform(γ))
            kernel2 =
                ExponentialKernel(; metric=Wasserstein()) ⊗
                (SqExponentialKernel() ∘ ARDTransform([γ]))

            # check evaluation of the first two observations
            p1 = predictions[1]
            p2 = predictions[2]
            t1 = targets[1]
            t2 = targets[2]
            for f in (
                CalibrationErrors.unsafe_skce_eval_targets,
                CalibrationErrors.unsafe_ucme_eval_targets,
            )
                out1 = f(kernel1.kernels[2], p1, t1, p2, t2)
                out2 = f(kernel2.kernels[2], p1, t1, p2, t2)
                @test out2 ≈ out1
                if isone(γ)
                    @test f(SqExponentialKernel(), p1, t1, p2, t2) ≈ out1
                end
            end

            # check estimates
            for estimator in (UnbiasedSKCE, x -> UCME(x, testpredictions, testtargets))
                estimate1 = estimator(kernel1)(predictions, targets)
                estimate2 = estimator(kernel2)(predictions, targets)
                @test estimate2 ≈ estimate1
                if isone(γ)
                    @test estimator(
                        ExponentialKernel(; metric=Wasserstein()) ⊗ SqExponentialKernel()
                    )(
                        predictions, targets
                    ) ≈ estimate1
                end
            end
        end
    end
end