text:
module Edges

using ..Paths
using ..Nodes
using ..Blocks
using ..Graphs

export edges, deparalog!

"""
    struct Position
        path  :: Path
        node  :: Tuple{Node{Block},Node{Block}}
        index :: Tuple{Int,Int} # positions on path
        locus :: Int # breakpoint on sequence
    end

Store a single position of an edge/breakpoint between homologous pancontigs in an individual genome.
`path` is the containing `Path` object.
`node` stores the junction of nodes that represent the position of the breakpoint.
`index` is the indices of `node` within `path`.
`locus` is the physical location on the genome of the breakpoint.
"""
struct Position
    path  :: Path
    node  :: Tuple{Node{Block},Node{Block}}
    index :: Tuple{Int,Int} # positions on path
    locus :: Int # breakpoint on sequence
end

"""
    isolates(positions::Array{Position,1})

Compute the array of `Position` values for each isolate.
"""
function isolates(positions::Array{Position,1})
    isolate = Dict{String,Array{Position,1}}()
    for pos in positions
        if pos.path.name ∈ keys(isolate)
            push!(isolate[pos.path.name],pos)
        else
            isolate[pos.path.name] = [ pos ]
        end
    end

    return isolate
end

"""
    next(x::Position, blk::Block)

Compute the next position from `x` that is connected through block `blk`.
Use to traverse the path of an individual genome through a pangraph.
"""
function next(x::Position, blk::Block)
    function index(i::Int)
        i ≤ 0 && return length(x.path)+i
        i > length(x.path) && return i - length(x.path)
        return i
    end

    function isright(num, of)
        if num == 1 && of == length(x.path)
            return true
        end
        if of == 1 && num == length(x.path)
            return false
        end

        return num > of
    end

    i, j = (blk == x.node[1].block) ? (x.index[1], x.index[2]) : (x.index[2], x.index[1])
    n, m, k, l, pos = if isright(j, i) #j > i
        ι = index(i-1)
        x.path.node[i], x.path.node[ι], i, ι, x.path.position[j]
    else
        ι = index(i+1)
        x.path.node[i], x.path.node[ι], i, ι, x.path.position[i]
    end

    if (n.block.uuid > m.block.uuid) || (n.block.uuid == m.block.uuid && pointer_from_objref(n) > pointer_from_objref(m))
        n, m = m, n
        k, l = l, k
    end

    function imax(a,b)
        L = length(x.path.node)
        if (a == L && b == 1) || (b == L && a == 1)
            return 1
        end
        return max(a,b)
    end

    pos = x.path.position[imax(k,l)]

    return Position(x.path, (n, m), (k, l), pos)
    #=
    i,s,d = if blk == x.node[1].block
        x.index[2],
        x.node[1].strand ? +1 : -1,
        index(minimum(x.index)+1)
    else
        x.index[1],
        x.node[2].strand ? +1 : -1,
        index(minimum(x.index)+1)
    end
    p = index(i + s)

    if x.locus != x.path.position[d]
        @infiltrate
    end

    n₁  = x.path.node[i]
    n₂  = x.path.node[p]
    pos = x.path.position[index(i+s+d)]

    if n₁.block.uuid > n₂.block.uuid
        i, p = p, i
        n₁, n₂ = n₂, n₁
    end

    return Position(
        x.path,
        (n₁, n₂),
        (i, p),
        pos
    )
    =#
end

Base.show(io::IO, x::Position) = Base.show(io,
    (
        path=x.path.name,
        blocks=(x.node[1].block, x.node[2].block),
        strand=(x.node[1].strand, x.node[2].strand),
        index=(x.index[1], x.index[2]),
        locus=x.locus
    )
)

"""
    mutable struct Edge
        block  :: Tuple{Block, Block}
        invert :: Bool # changes strand
        nodes  :: Array{Position}
    end

Store a unique edge within a pangraph.
An edge is undirected and is defined by the two juxtaposed blocks, as well as a relative orientation.
Contain all positions of all genomes that contain the edge.
"""
mutable struct Edge
    block  :: Tuple{Block, Block}
    invert :: Bool # changes strand
    nodes  :: Array{Position}
end

Breakpoints = Tuple{Array{Tuple{Int,Int},1},Array{Tuple{Int,Int},1}}

function equals(e::Edge, block1::Block, block2::Block, invert::Bool)
    return e.block[1] == block1 && e.block[2] == block2 && invert == e.invert
end

Edge(block₁::Block, block₂::Block, invert::Bool, pos::Position) = Edge((block₁, block₂), invert, Position[pos])

"""
    edges(G)

Compute all edges associated with pangraph `G`.
"""
function edges(G)
    edgeset = Edge[]

    function addedge!(side₁, side₂, path, locus)
        if (side₁.node.block.uuid > side₂.node.block.uuid) || (side₁.node.block.uuid == side₂.node.block.uuid && pointer_from_objref(side₁.node) > pointer_from_objref(side₂.node))
            side₁, side₂ = side₂, side₁
        end

        invert   = side₁.node.strand ⊻ side₂.node.strand
        position = Position(
            path,
            (side₁.node, side₂.node),
            (side₁.index,side₂.index),
            locus
        )
        for edge in edgeset
            if equals(edge, side₁.node.block, side₂.node.block, invert)
                push!(edge.nodes, position)
                return
            end
        end
        # if we are here, then this is a new edge
        push!(edgeset, Edge(side₁.node.block, side₂.node.block, invert, position))
    end

    for path in values(G.sequence)
        for (i,link) in enumerate(path.node[2:end])
            node = path.node[i]
            addedge!(
                (node=node,index=i),
                (node=link,index=i+1),
                path, path.position[i+1]
            )
        end

        if path.circular
            addedge!(
                (node=path.node[end],index=length(path.node)),
                (node=path.node[1],index=1),
                path, path.position[1]
            )
        end
    end

    return edgeset
end

"""
    deparalog!(G)

Split duplicated blocks that have non-intersecting, seperable paths that run in parallel and transitively connect associated genomes.
Use to simplify high copy number blocks found in all individuals in equivalent contexts within pangraph `G`.
"""
function deparalog!(G)
    edgeset = edges(G)

    # collect all edges for each block, segregated by position
    vertex = Dict(
        block => (
            edge = Set{Edge}(),
            link = Dict{Position,Edge}(),
        ) for block in values(G.block)
    )

    function duplicates(vertex, block)
        loci = [
            let
                if k.node[1].block == block
                    (k.path.name,k.index[1])
                elseif k.node[2].block == block
                    (k.path.name,k.index[2])
                else
                    error("bad indexing")
                end
            end for k in keys(vertex)
        ] |> unique
        names = unique(map(first, loci))
        return length(loci) > length(names)
    end

    function connect(edge, block)
        for locus in edge.nodes
            push!(vertex[block].edge, edge)
            vertex[block].link[locus] = edge
        end
    end

    for edge in edgeset
        block = edge.block
        connect(edge, block[1])
        connect(edge, block[2])
    end

    touched = Set{Block}()
    blocks  = collect(values(G.block))

    for block ∈ blocks
        block ∉ touched || continue
        vert = vertex[block]

        duplicates(vert.link,block) || continue

        split = true
        links = Dict{Tuple{Edge,Edge}, Set{Position}}()
        for edge in vert.edge
            dest = Dict{Edge,Array{Position,1}}()
            for pos in edge.nodes
                x = next(pos,block)
                x ∈ keys(vert.link) || @goto nosplit

                e = vert.link[x]
                if e ∈ keys(dest)
                    push!(dest[e], x)
                else
                    dest[e] = [x]
                end
            end

            if length(dest) > 1
                # XXX: last effort to split paralogs:
                # if the set positions contained by dest edges
                # is equal to the the set positions we flow to,
                # then we are still transitive
                linkpos = reduce(union, map((link)->Set(link.nodes), collect(keys(dest))))
                edgepos = reduce(union, map((v)->Set(v), collect(values(dest))))
                if linkpos == edgepos
                    for link in keys(dest)
                        if (link,edge) ∉ keys(links)
                            linkpos = Set(link.nodes)
                            links[edge,link] = Set(
                                pos for pos in edge.nodes if next(pos, block) ∈ linkpos
                            )
                        end
                    end
                else @label nosplit
                    split = false
                    break
                end
            else
                link = collect(keys(dest))[1]
                if (link,edge) ∉ keys(links)
                    links[(edge,link)] = Set(edge.nodes)
                end
            end
        end

        (split && length(links) > 1) || continue

        seennodes = Set{Node}()
        for positions in collect(values(links))[2:end]
            new = Block(block.sequence, block.gaps)
            for pos in positions
                i,j = if pos.node[1].block == block
                    (1,2)
                elseif pos.node[2].block == block
                    (2,1)
                else
                    error("bad edge")
                end
                push!(touched, pos.node[j].block)

                # rewire node to new block
                n = pos.node[i]
                push!(seennodes,n)
                new.mutate[n] = pop!(block.mutate, n)
                new.insert[n] = pop!(block.insert, n)
                new.delete[n] = pop!(block.delete, n)
                n.block = new
            end

            G.block[new.uuid] = new
        end
    end
end

end

max_stars_repo_name: neherlab/pangraph

max_issues_repo_name: neherlab/pangraph

max_forks_repo_name: neherlab/pangraph


text:
#function render(wall::Wall, ax; kwargs...)
    #plot([wall.a.x, wall.b.x], [wall.a.y, wall.b.y]; kwargs...)
#end
#
#function render(scene::Scene, ax; kwargs...)
    #for wall in scene.walls
        #render(wall, ax; kwargs...)
    #end
    #xspan = scene.bounds.xmax - scene.bounds.xmin
    #yspan = scene.bounds.ymax - scene.bounds.ymin
    #ax[:set_xlim]((scene.bounds.xmin - xspan * 0.05, scene.bounds.xmax + xspan * 0.05))
    #ax[:set_ylim]((scene.bounds.ymin - yspan * 0.05, scene.bounds.ymax + yspan * 0.05))
#end
function render_tree(tree::RRTTree, thickness)
    for node in 1:tree.num_nodes
        a = Point(tree.confs[1,node], tree.confs[2,node])
        parent = tree.parents[node]
        if parent == 0
            continue
        end
        b = Point(tree.confs[1,parent],tree.confs[2,parent])
        plot([a.x, b.x], [a.y, b.y], color="black", linewidth=thickness)
    end
end

ext: jl


text:
using Test, VIEstimators, Random, LinearAlgebra

# Helper functionality for testing, primarily checking custom adjoints.
include("test_util.jl")

@testset "VIEstimators" begin
    include("util.jl")
    include("gaussian.jl")
    # include("elbo_estimators.jl")
end


text:
#using HiOscSolver
using Test
include("testfftbig.jl")
include("testexpmatrices.jl")
include("testpolyexp.jl")
# include("testcoefexp_ab.jl")
# include("testpreparephi.jl")
# include("testtwoscales_pure_ab.jl")
include("testinterface.jl")
include("testcompfloat.jl")


text:
using Gurobi
using Test

# When adding new tests, be mindful about creating new environments. Either
# re-use an existing environment in the module, or name the test function
# `test_MULTI_ENV_xxx` to trap the specific Gurobi error indicating that an
# environment could not be created.
const GRB_ENV = Gurobi.Env(output_flag = 0)

@testset "MathOptInterface Tests" begin
    for file in readdir("MOI")
        include(joinpath("MOI", file))
    end
end


text:
using Flux, Flux.Data.CMUDict
using Flux: onehot, batchseq
using Base.Iterators: partition

tokenise(s, α) = [onehot(c, α) for c in s]

function getdata(args)
    dict = cmudict()
    alphabet = [:end, CMUDict.alphabet()...]
    args.Nin = length(alphabet)

    phones = [:start, :end, CMUDict.symbols()...]
    args.phones_len = length(phones)

    # Turn a word into a sequence of vectors
    tokenise("PHYLOGENY", alphabet)
    # Same for phoneme lists
    tokenise(dict["PHYLOGENY"], phones)

    words = sort(collect(keys(dict)), by = length)

    # Finally, create iterators for our inputs and outputs.
    batches(xs, p) = [batchseq(b, p) for b in partition(xs, 50)]

    Xs = batches([tokenise(word, alphabet) for word in words],
             onehot(:end, alphabet))

    Ys = batches([tokenise([dict[word]..., :end], phones) for word in words],
             onehot(:end, phones))

    Yo = batches([tokenise([:start, dict[word]...], phones) for word in words],
             onehot(:end, phones))

    data = collect(zip(Xs, Yo, Ys))
    return data, alphabet, phones
end

max_stars_repo_name: shawn-kern/model-zoo

max_issues_repo_name: shawn-kern/model-zoo

max_forks_repo_name: shawn-kern/model-zoo


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.0
#   kernelspec:
#     display_name: Julia 1.4.2
#     language: julia
#     name: julia-1.4
# ---

# + active=""
# # Harris Corner Detector
#
# - [実践 コンピュータビジョン](https://www.oreilly.co.jp/books/9784873116075/)
# - https://en.wikipedia.org/wiki/Harris_Corner_Detector
# - https://www.researchgate.net/publication/299711769_Revisiting_Harris_Corner_Detector_Algorithm_A_Gradual_Thresholding_Approach
# - [解説コンピュータビジョンのための 画像の特徴点の抽出](http://www.iim.cs.tut.ac.jp/~kanatani/papers/harris.pdf)
# -

using Images
using TestImages
using ImageShow
using ImageFiltering

img = testimage("c")

typeof(img), eltype(img)

# +
# img = testimage("jetplane")
# img = Gray.(img) # GrayA to Gray
# -

imfilter(img, 8 .* Kernel.sobel())

∂x, ∂y = 8 .* Kernel.sobel()

Ix = imfilter(img, ∂x)

Iy = imfilter(img, ∂y)

Iy .^ 2.

# ## Calculate Harris response

# +
#=
W = [Wxx Wxy
     Wxy Wyy]y
  = G * I
where G * I is the convolution between Gauss Operator G and I
I = [Ix^2  IxIy
    IyIx  Iy^2]
  = [Ix * [Ix Iy]
    Iy]
=#

σ = 3.

Wxx = imfilter(Ix .^ 2., Kernel.gaussian([σ,σ], [3,3]))
Wyy = imfilter(Iy .^ 2., Kernel.gaussian([σ,σ], [3,3]))
Wyx = Wxy = imfilter(Ix .* Iy, Kernel.gaussian([σ,σ], [3,3]))

detW = @. Wxx * Wyy - Wxy * Wyx
trW = @. Wxx + Wyy
κ = detW ./ trW
# -

thresh = 0.08
corner_thresh = maximum(κ) * thresh
iscorner = 1. .* (κ .> corner_thresh)
colorview(Gray, iscorner)

# # Apply NMS (Non maximum suppression)

# +
mindist = 10
indices_ = sortperm(κ[:],rev=true) # reshape Mat to Vector
indices = Int[]
for i in indices_
    if κ[i] > corner_thresh
        push!(indices, i)
    end
end

indices;
# -

loc = ones(Bool, img |> size)
H, W = img |> size
valids = zeros(img|>size)
for i in indices
    ix = min(div(i, H)+1,W)
    iy = min(rem(i, H)+1,H)
    if loc[iy, ix] == true
        valids[iy, ix] = 1
        for s in -mindist:mindist
            for t in -mindist:mindist
                if 1<= iy+s <= H && 1<= ix+t <=W
                    loc[iy+s, ix+t] = false
                end
            end
        end
    end
end

colorview(Gray, hcat(channelview(img), valids))


text:
function convert(::Type{Polynomial{T}}, p::Polynomial{S}) where {T,S}
    Polynomial{T}(p.exponents, convert(Vector{T}, p.coefficients), p.variables, p.homogenized)
end

function convert(::Type{Polynomial{T}}, p::MP.AbstractPolynomialLike) where {T}
    Polynomial{T}(p)
end

function convert(::Type{Vector{Polynomial{T}}}, ps::Vector{<:MP.AbstractPolynomialLike}) where {T}
    variables = sort!(union(Iterators.flatten(MP.variables.(ps))), rev=true)
    [Polynomial{T}(p, variables) for p in ps]
end

function promote_rule(::Type{Polynomial{T}}, x::Type{Polynomial{S}}) where {T,S}
    Polynomial{promote_type(S,T)}
end
function promote_rule(::Type{Polynomial{T}}, x::Type{S}) where {T,S}
    Polynomial{promote_type(S,T)}
end


text:
###################################################
## Tools to subset networks
###################################################

"""
subset the network: only keep the given nodes
"""
function subsetNetwork(n::Network, keepN::Vector{Int})
    nodesIndices = Dict{Int,Int}()
    roads = Dict{Tuple{Int,Int},Road}()
    nodes = Array{Node}(length(keepN))

    for (k,i) in enumerate(keepN)
        nodes[k] = n.nodes[i]
        nodesIndices[i] = k
    end

    g = DiGraph(length(nodes))
    for (s,d) in keys(n.roads)
        if haskey(nodesIndices,s) && haskey(nodesIndices,d)
            add_edge!(g,nodesIndices[s],nodesIndices[d])
            prevRoad = n.roads[s,d]
            roads[nodesIndices[s],nodesIndices[d]] = Road(nodesIndices[s], nodesIndices[d], prevRoad.distance, prevRoad.roadType)
        end
    end

    return updateProjection!(Network(g,nodes,roads))
end

"""
Remove nodes from network
"""
function removeNodes(n::Network, rem::Vector{Int})
    keep = ones(Bool, length(n.nodes))
    for i in rem
        keep[i] = false
    end
    keepN = collect(1:length(n.nodes))[keep]
    Tuple{Float64,Float64}[(n.x,n.y) for n in n.nodes]
    return subsetNetwork(n,keepN)
end

"""
Extract single nodes from network
"""
singleNodes(n::Network) = singleNodes(n.graph)
function singleNodes(g::DiGraph)
    singles = Int[]
    for i in vertices(g)
        if degree(g,i) == 0
            push!(singles, i)
        end
    end
    return singles
end

"""
Return strongly connected component starting from a node
"""
function stronglyConnected(n::Network, node::Int)
    indices  = singleNodes(bfs_tree(n.graph,node))
    indices2 = singleNodes(bfs_tree(LightGraphs.reverse(n.graph), node))
    return removeNodes(n,[indices;indices2])
end

"""
Returns nodes that are inside a Polygon
"""
function inPolygon{T<:AbstractFloat}(n::Network, poly::Vector{Tuple{T,T}})
    inPoly = Int[]
    for i in vertices(n.graph)
        if pointInsidePolygon(n.nodes[i].lon, n.nodes[i].lat, poly)
            push!(inPoly,i)
        end
    end
    return inPoly
end

"""
Only keep roads that are of some types
"""
function roadTypeSubset(n::Network, roadTypes::AbstractArray{Int})
    keep = Bool[(i in roadTypes) for i in 1:8]
    roads = Dict{Tuple{Int,Int},Road}()
    g = DiGraph(length(n.nodes))

    for (o,d) in keys(n.roads)
        if keep[n.roads[o,d].roadType]
            roads[o,d] = n.roads[o,d]
            add_edge!(g,o,d)
        end
    end
    n2 = Network(g,deepcopy(n.nodes), roads)
    n2.projcenter = n.projcenter
    return n2
end

"""
Only keep nodes that are intersections, suppose strong connectivity
"""
function intersections(n::Network)
    g = n.graph
    #################################
    # First pass: detect intersections nodes
    #################################
    nodes = Node[]
    index = Int[]
    revIndex = zeros(Int,nv(g))
    for i in vertices(n.graph)
        if length(all_neighbors(n.graph,i)) != 2 || indegree(n.graph,i) != outdegree(n.graph,i)
            push!(nodes,n.nodes[i])
            revIndex[i] = length(nodes)
            push!(index,i)
        end
    end

    #################################
    # second pass:  recreate graph
    #################################

    g2 = DiGraph(length(nodes))
    roads = Dict{Tuple{Int,Int},Road}()
    for i in vertices(g2)
        for j in out_neighbors(g,index[i])
            dist = n.roads[index[i],j].distance
            roadType = n.roads[index[i],j].roadType
            k = j
            k2 = index[i]
            while revIndex[k] == 0
                next = out_neighbors(g,k)[1]
                if next == k2
                    next = out_neighbors(g,k)[2]
                end
                dist += n.roads[k,next].distance
                roadType = max(roadType, n.roads[k,next].roadType)
                k2 = k
                k = next
            end
            if i != revIndex[k]
                if has_edge(g2,i,revIndex[k])
                    dist = min(dist, roads[i,revIndex[k]].distance)
                    roadType = max(roadType, roads[i,revIndex[k]].roadType)
                else
                    add_edge!(g2, i, revIndex[k])
                end
                roads[i,revIndex[k]] = Road(i,revIndex[k],dist,roadType)
            end

        end
    end
    n2 = Network(g2,nodes,roads)
    updateProjection!(n2)
    return n2
end

"""
    Take a list of RoadPosition, and returns the list of nodes that are needed to
    drive/walk from any position to any other position.
"""
function connectingNodes(n::Network, positions::Vector{RoadPosition})
    cartimes = maxSpeedTimes(n)
    dists = roadDistances(n)
    walktimes = max.(dists, dists')
    walkgraph = Graph(n.graph)
    allNodes = Set{Int}()
    for p in positions
        push!(allNodes, p.nodeO)
        push!(allNodes, p.nodeD)
    end
    allNodes = sort(collect(allNodes))

    seenNodes = Set{Int}()
    println(length(allNodes))
    for (i, origin) in enumerate(allNodes)
        print(i, "/", length(allNodes), "\r")
        union!(seenNodes, getConnectingNodes(n.graph, cartimes, origin, allNodes))
        union!(seenNodes, getConnectingNodes(walkgraph, cartimes, origin, allNodes))
    end
    return seenNodes
end

"""
    Take an origin node and a graph, with a set of destination nodes, and return all the
    nodes needed to take a shortest path.
"""
function getConnectingNodes(graph, times, origin, dests)
    dijkstraResult = dijkstra_shortest_paths(graph, origin, times)
    paths  = enumerate_paths(dijkstraResult, dests)
    nodes = Set{Int}()
    for p in paths
        union!(nodes, p)
    end
    return nodes
end

max_stars_repo_name: andrew-zhm/RoutingNetworks

max_issues_repo_name: andrew-zhm/RoutingNetworks

max_forks_repo_name: andrew-zhm/RoutingNetworks


text:
include("../../src/ttv.jl")
include("/Users/ericagol/Computer/Julia/regress.jl")

using PyPlot

# This routine takes derivative of transit times with respect
# to the initial orbital elements.
#n = 8
n = 3
n_body = n
t0 = 7257.93115525
#t0 =  0.0
#t0 =  randn()
#h  = 0.12
h  = 0.07
#tmax = 600.0
#tmax = 800.0
#tmax = 600.0
#tmax = 800.0
tmax = 1000.0

# Read in initial conditions:
elements = readdlm("elements.txt",',')
# Make masses of planets bigger
#elements[2,1] *= 10.0
#elements[3,1] *= 10.0

ntt = zeros(Int64,n)

# Make an array, tt,  to hold transit times:
# First, though, make sure it is large enough:
for i=2:n
  ntt[i] = ceil(Int64,tmax/elements[i,2])+3
end
dtdq0 = zeros(n,maximum(ntt),7,n)
tt  = zeros(n,maximum(ntt))
tt1 = zeros(n,maximum(ntt))
tt_save = zeros(5,n,maximum(ntt))
# Save a counter for the actual number of transit times of each planet:
count = zeros(Int64,n)
count1 = zeros(Int64,n)
# Call the ttv function:
rstar = 1e12
elements_big = convert(Array{BigFloat,2},elements)
hbig = big(h)
t0big = big(t0)
tmaxbig = big(tmax)
tt1big = big.(tt1)
rstarbig = big(rstar)
dqbig = ttv_elements!(n,t0big,hbig,tmaxbig,elements_big,tt1big,count,big(0.0),0,0,rstarbig)
tt_save[1,:,:]=convert(Array{Float64,2},tt1big)
#dq = ttv_elements!(n,t0,h,tmax,elements,tt1,count1,0.0,0,0,rstar)
#tt_save[1,:,:]=tt1
# Now call with 1/10 the timestep:
#dq = ttv_elements!(n,t0,h/10.,tmax,elements,tt1,count,0.0,0,0,rstar)

mask = zeros(Bool, size(dtdq0))
for jq=1:n_body
  for iq=1:7
    if iq == 7; ivary = 1; else; ivary = iq+1; end  # Shift mass variation to end
    for i=2:n
      for k=1:count[i]
        # Ignore inclination & longitude of nodes variations:
        if iq != 5 && iq != 6 && ~(jq == 1 && iq < 7) && ~(jq == i && iq == 7)
          mask[i,k,iq,jq] = true
        end
      end
    end
  end
end

# Now, compute derivatives (with respect to initial cartesian positions/masses):
dtdelements0 = zeros(n,maximum(ntt),7,n)
#dtdelements0 = ttv_elements!(n,t0,h,tmax,elements,tt,count,dtdq0,rstar)
dtdq2 = zeros(n,maximum(ntt),7,n)
dtdelements2 = zeros(n,maximum(ntt),7,n)
#dtdelements2 = ttv_elements!(n,t0,h/2.,tmax,elements,tt1,count,dtdq2,rstar)
dqbig = ttv_elements!(n,t0big,hbig/2,tmaxbig,elements_big,tt1big,count,big(0.0),0,0,rstarbig)
tt_save[2,:,:]=convert(Array{Float64,2},tt1big)
#dq = ttv_elements!(n,t0,h/2.,tmax,elements,tt1,count,0.0,0,0,rstar)
#tt_save[2,:,:]=tt1
dtdq4 = zeros(n,maximum(ntt),7,n)
dtdelements4 = zeros(n,maximum(ntt),7,n)
#dtdelements4 = ttv_elements!(n,t0,h/4.,tmax,elements,tt1,count,dtdq4,rstar)
dqbig = ttv_elements!(n,t0big,hbig/4,tmaxbig,elements_big,tt1big,count,big(0.0),0,0,rstarbig)
tt_save[3,:,:]=convert(Array{Float64,2},tt1big)
#dq = ttv_elements!(n,t0,h/4.,tmax,elements,tt1,count,0.0,0,0,rstar)
#tt_save[3,:,:]=tt1
dtdq8 = zeros(n,maximum(ntt),7,n)
dtdelements8 = zeros(n,maximum(ntt),7,n)
#dtdelements8 = ttv_elements!(n,t0,h/8.,tmax,elements,tt1,count,dtdq8,rstar)
#dq = ttv_elements!(n,t0,h/8.,tmax,elements,tt1,count,0.0,0,0,rstar)
#tt_save[4,:,:]=tt1
dqbig = ttv_elements!(n,t0big,hbig/8,tmaxbig,elements_big,tt1big,count,big(0.0),0,0,rstarbig)
tt_save[4,:,:]=convert(Array{Float64,2},tt1big)
dtdq16 = zeros(n,maximum(ntt),7,n)
dtdelements16 = zeros(n,maximum(ntt),7,n)
#dtdelements16 = ttv_elements!(n,t0,h/16.,tmax,elements,tt1,count,dtdq16,rstar)
#dq = ttv_elements!(n,t0,h/16.,tmax,elements,tt1,count,0.0,0,0,rstar)
#tt_save[5,:,:]=tt1
dqbig = ttv_elements!(n,t0big,hbig/16,tmaxbig,elements_big,tt1big,count,big(0.0),0,0,rstarbig)
tt_save[5,:,:]=convert(Array{Float64,2},tt1big)
#println("Maximum error on derivative: ",maximum(abs.(dtdelements0-dtdelements2)))
#println("Maximum error on derivative: ",maximum(abs.(dtdelements2-dtdelements4)))
#println("Maximum error on derivative: ",maximum(abs.(dtdelements4-dtdelements8)))
println("Maximum error on derivative: ",maximum(abs.(asinh.(dtdelements0[mask])-asinh.(dtdelements2[mask]))))
println("Maximum error on derivative: ",maximum(abs.(asinh.(dtdq0)-asinh.(dtdq2))))
println("Maximum error on derivative: ",maximum(abs.(asinh.(dtdelements2[mask])-asinh.(dtdelements4[mask]))))
println("Maximum error on derivative: ",maximum(abs.(asinh.(dtdq2)-asinh.(dtdq4))))
println("Maximum error on derivative: ",maximum(abs.(asinh.(dtdelements4[mask])-asinh.(dtdelements8[mask]))))
println("Maximum error on derivative: ",maximum(abs.(asinh.(dtdq4)-asinh.(dtdq8))))

# Make a plot of transit time errors versus stepsize:
ntrans = sum(count)
clf()
sigt = zeros(n-1,5)
tab = 0
h_list = [h,h/2.,h/4.,h/8.,h/8]
hlabel = ["h-h/16","h/2-h/16","h/4-h/16","h/8-h/16","h/16-big(h/16)"]
ch = ["black","red","green","blue","orange"]
for i=2:n
  for j=1:4
    tti1 = tt_save[j,i,1:count[i]]
    tti16 = tt_save[5,i,1:count[i]]
    diff = tti1-tti16
    sigt[i-1,j]=std(diff)
    if i == n
      plot(tti1,diff/h_list[j]^4,linestyle="dashed",c=ch[j])
    else
      plot(tti1,diff/h_list[j]^4,label=hlabel[j],c=ch[j])
    end
  end
end
xlabel("Transit time")
ylabel("Difference in transit time / h^4")
legend(loc="lower left")

read(STDIN,Char)

# Make a plot of timing errors versus stepsize:
ntrans = sum(count)
clf()
sigt = zeros(n-1,5)
tab = 0
h_list = [h,h/2.,h/4.,h/8.,h/8]
hlabel = ["h-h/16","h/2-h/16","h/4-h/16","h/8-h/16","h/16-big(h/16)"]
ch = ["black","red","green","blue","orange"]
for i=2:n
  fn = zeros(Float64,2,count[i])
  sig = ones(count[i])
  tti16 = tt_save[5,i,1:count[i]]
  for j=1:count[i]
    fn[1,j] = 1.0
    fn[2,j] = round(Int64,(tti16[j]-elements[i,3])/elements[i,2])
  end
  for j=1:4
    tti1 = tt_save[j,i,1:count[i]]
    coeff,cov = regress(fn,tti1-tti16,sig)
    diff = tti1-tti16-coeff[1]-coeff[2]*fn[2,:]
    sigt[i-1,j]=std(diff)
    if i == n
      plot(tti1,diff/h_list[j]^4,linestyle="dashed",c=ch[j])
    else
      plot(tti1,diff/h_list[j]^4,label=hlabel[j],c=ch[j])
    end
  end
end
legend(loc="lower left")

read(STDIN,Char)
clf()

# Make a plot of some TTVs:
loglog(h_list,sigt[1,:]*24.*3600.,".",markersize=15,label="Inner planet")
loglog(h_list,sigt[1,1]*24.*3600.*(h_list/h[1]).^4,label=L"$\propto h^4$")
loglog(h_list,sigt[2,:]*24.*3600.,".",markersize=15,label="Outer planet")
loglog(h_list,sigt[2,1]*24.*3600.*(h_list/h[1]).^4,label=L"$\propto h^4$")
legend(loc = "upper left")
ylabel("RMS timing error [sec]")
xlabel("Step size [day]")

PyPlot.savefig("timing_error_vs_h_big.pdf",bbox_inches="tight")

max_forks_repo_name: langfzac/nbg-papers


text:
include("../apm.jl")

# Select server
server = "http://byu.apmonitor.com"

# Set application name
app = "distill"

# Clear previous application
apm(server,app,"clear all")

# Load model file
apm_load(server,app,"distill.apm")

# Load time points for future predictions
csv_load(server,app,"horizon.csv")

# APM Variable Classification
# class = FV, MV, SV, CV
#   F or FV = Fixed value - parameter may change to a new value every cycle
#   M or MV = Manipulated variable - independent variable over time horizon
#   S or SV = State variable - model variable for viewing
#   C or CV = Controlled variable - model variable for control
FVs = ["feed","x_feed","alpha","atray","acond","areb"]
MVs = ["rr","fbot","sp_x[1]","sp_x[32]"]
SVs = ["x[2]","x[5]","x[10]","x[15]","x[20]","x[25]","x[30]","x[31]"]
CVs = ["x[1]","x[32]"]

# Set up variable classifications for data flow
for x in FVs
	apm_info(server,app,"FV",x)
end
for x in MVs
    apm_info(server,app,"MV",x)
end
for x in SVs
    apm_info(server,app,"SV",x)
end
for x in CVs
    apm_info(server,app,"CV",x)
end

# Options

# time units (1=sec,2=min,3=hrs,etc)
apm_option(server,app,"nlc.ctrl_units",2)
apm_option(server,app,"nlc.hist_units",3)

# set controlled variable error model type
apm_option(server,app,"nlc.cv_type",1)
apm_option(server,app,"nlc.ev_type",1)

# controller mode (1=simulate, 2=predict, 3=control)
apm_option(server,app,"nlc.reqctrlmode",1)

# read discretization from CSV file
apm_option(server,app,"nlc.csv_read",1)

# turn on historization to see past results
apm_option(server,app,"nlc.hist_hor",200)

# set web plot update frequency
apm_option(server,app,"nlc.web_plot_freq",10)

# Controlled variable (c)
apm_option(server,app,"x[1].sphi",0.95)
apm_option(server,app,"x[1].splo",0.94)
apm_option(server,app,"x[1].tau",20.0)
apm_option(server,app,"x[1].fstatus",0)

apm_option(server,app,"x[32].sphi",0.05)
apm_option(server,app,"x[32].splo",0.04)
apm_option(server,app,"x[32].tau",20.0)
apm_option(server,app,"x[32].fstatus",0)

# Manipulated variables (u)
apm_option(server,app,"rr.upper",10)
apm_option(server,app,"rr.dmax",0.2)
apm_option(server,app,"rr.lower",1)
apm_option(server,app,"rr.fstatus",0)

apm_option(server,app,"fbot.upper",1)
apm_option(server,app,"fbot.dmax",0.05)
apm_option(server,app,"fbot.lower",0)
apm_option(server,app,"fbot.fstatus",0)

# Measured Disturbances
apm_option(server,app,"feed.fstatus",1)
apm_option(server,app,"x_feed.fstatus",1)

# imode (1=ss, 2=mpu, 3=rto, 4=sim, 5=mhe, 6=nlc)
apm_option(server,app,"nlc.imode",1)
apm_option(server,app,"nlc.sensitivity",1)

# steady state solution
solver_output = apm(server,app,"solve")
print(solver_output)

# imode (1=ss, 2=mpu, 3=rto, 4=sim, 5=mhe, 6=nlc)
apm_option(server,app,"nlc.imode",6)
apm_option(server,app,"nlc.sensitivity",0)

# turn on controller
apm_option(server,app,"nlc.reqctrlmode",3)
# turn on overhead composition control
apm_option(server,app,"rr.status",1)
apm_option(server,app,"x[1].sphi",0.952)
apm_option(server,app,"x[1].splo",0.952)
apm_option(server,app,"x[1].status",1)
apm_meas(server,app,"sp_x[1]",0.952)
# turn on bottoms composition control
apm_option(server,app,"fbot.status",1)
apm_option(server,app,"x[32].sphi",0.019)
apm_option(server,app,"x[32].splo",0.019)
apm_option(server,app,"x[32].status",1)
apm_meas(server,app,"sp_x[32]",0.019)

# Run NLC on APM server
solver_output = apm(server,app,"solve")
println(solver_output)

# Open Web Viewer and Display Link
println("Web viewer")
url = apm_web(server,app)


text:
## Jacobi Constant
# =============================================================================
function JC(r, v, μ₃)
    # jacobiConstant(r,v,μ₃)
    # r is [x;y;z] in nondimensional units
    # v is [x';y';z'] in nondimensional units, with as many columns as you like
    # μ₃ is 3-body mu, or m2/(m1+m2)
    # outputs row vector of C-values
    #
    # Alex Hoffman
    # 10/31/2020
    # notation from AAE632 notes

    x, y, z = r[1:1, :], r[2:2, :], r[3:3, :]
    xd, yd, zd = v[1:1, :], v[2:2, :], v[3:3, :]
    d = sqrt.((x .+ μ₃) .^ 2 .+ y .^ 2 .+ z .^ 2)
    r = sqrt.((x .- 1 .+ μ₃) .^ 2 .+ y .^ 2 .+ z .^ 2)
    C =
        -(xd .^ 2 .+ yd .^ 2 .+ zd .^ 2) .+ x .^ 2 .+ y .^ 2 .+
        2(1 - μ₃) ./ d .+ 2μ₃ ./ r

    return C
end

function xtest(a, b)
    return a
end

## Directions Cosine Matrix
# =============================================================================
function DCM(angle, XYZ, isPassive)
    # rot = DCM(angle,XYZ,isPassive)
    # angle is in radians
    # XYZ is 'x' 'y' or 'z'
    # isPassive is a logical stating whether you want a passive rotation
    #   (which rotates the reference frame, not the vector)
    # output is the direction cosines matrix
    # you should use passive when creating a transformation matrix from euler angles
    #
    # Alex Hoffman
    # 10/30/2020

    if XYZ == 'x'
        rot = [1 0 0; 0 cos(angle) sin(angle); 0 -sin(angle) cos(angle)]
    elseif XYZ == 'y'
        rot = [cos(angle) 0 -sin(angle); 0 1 0; sin(angle) 0 cos(angle)]
    elseif XYZ == 'z'
        rot = [cos(angle) sin(angle) 0; -sin(angle) cos(angle) 0; 0 0 1]
    else
        error("Rotation input XYZ must be 'x','y', or 'z'.")
    end
    if ~isPassive
        return transpose(rot)
    else
        return rot
    end
end

## CR#BP Equations of Motion
# =============================================================================
function CR3BP_EOM(state,μ₃,xyzddot)
    # calculates the accelerations in 3BP
    #
    # Alex Hoffman
    # 10/31/2020
    x, y, z, xd, yd, zd = state
    r = state[1:3]
    if xyzddot == 'x'
        dd = 2yd + U_star(r,μ₃,'x')
    elseif xyddot == 'y'
        dd = -2xd + U_star(r,μ₃,'y')
    elseif xyzddot == 'z'
        dd = U_star(r,μ₃,'z')
    else
        error("xyzddot must be 'x', 'y', or 'z'.")
    end
    return dd
end

## Pseudopotential
# =============================================================================
function U_star(r,μ₃,deriv)
    # return U_star(r,mu3,deriv)
    # solves 3-Body U* pseudo-potential and its derivatives
    # X,Y,Z are the position coordinates
    # μ₃ is 3-body mu
    # deriv is 'U','i','ij', where i and j are either 'x','y','z'
    # this represents evaluation of the function and its partial derivatives
    x,y,z = r
    mu3 = μ₃
    a = length(deriv)
    if deriv[1] == 'U'
        U = 0.5*(x^2 + y^2) + (1-mu3)/sqrt((x+mu3)^2+y^2+z^2) + mu3/sqrt((x-1+mu3)^2+y^2+z^2)
        # turn U into symbolic with x,y,z to get derivatives
        return U
    elseif deriv[1] == 'x'
        if a == 1
            Ux = x + ((2*mu3 + 2*x)*(mu3 - 1))/(2*((mu3 + x)^2 + y^2 + z^2)^(3/2)) - (mu3*(2*mu3 + 2*x - 2))/(2*((mu3 + x - 1)^2 + y^2 + z^2)^(3/2))
            return Ux
        else #second partial
            if deriv[2] == 'x'
                Uxx = (mu3 - 1)/((mu3 + x)^2 + y^2 + z^2)^(3/2) - mu3/((mu3 + x - 1)^2 + y^2 + z^2)^(3/2) + (3*mu3*(2*mu3 + 2*x - 2)^2)/(4*((mu3 + x - 1)^2 + y^2 + z^2)^(5/2)) - (3*(2*mu3 + 2*x)^2*(mu3 - 1))/(4*((mu3 + x)^2 + y^2 + z^2)^(5/2)) + 1
                return Uxx
            elseif deriv[2] == 'y'
                Uxy = (3*mu3*y*(2*mu3 + 2*x - 2))/(2*((mu3 + x - 1)^2 + y^2 + z^2)^(5/2)) - (3*y*(2*mu3 + 2*x)*(mu3 - 1))/(2*((mu3 + x)^2 + y^2 + z^2)^(5/2))
                return Uxy
            elseif deriv[2] == 'z'
                Uxz = (3*mu3*z*(2*mu3 + 2*x - 2))/(2*((mu3 + x - 1)^2 + y^2 + z^2)^(5/2)) - (3*z*(2*mu3 + 2*x)*(mu3 - 1))/(2*((mu3 + x)^2 + y^2 + z^2)^(5/2))
                return Uxz
            else
                error("incorrect derivative expression")
            end
        end
    elseif deriv[1] == 'y'
        if a == 1
            Uy = y - (mu3*y)/((mu3 + x - 1)^2 + y^2 + z^2)^(3/2) + (y*(mu3 - 1))/((mu3 + x)^2 + y^2 + z^2)^(3/2)
            return Uy
        else #second partial
            if deriv[2] == 'x'
                Uyx = (3*mu3*y*(2*mu3 + 2*x - 2))/(2*((mu3 + x - 1)^2 + y^2 + z^2)^(5/2)) - (3*y*(2*mu3 + 2*x)*(mu3 - 1))/(2*((mu3 + x)^2 + y^2 + z^2)^(5/2))
                return Uyx
            elseif deriv[2] == 'y'
                Uyy = (mu3 - 1)/((mu3 + x)^2 + y^2 + z^2)^(3/2) - mu3/((mu3 + x - 1)^2 + y^2 + z^2)^(3/2) - (3*y^2*(mu3 - 1))/((mu3 + x)^2 + y^2 + z^2)^(5/2) + (3*mu3*y^2)/((mu3 + x - 1)^2 + y^2 + z^2)^(5/2) + 1
                return Uyy
            elseif deriv[2] == 'z'
                Uyz = (3*mu3*y*z)/((mu3 + x - 1)^2 + y^2 + z^2)^(5/2) - (3*y*z*(mu3 - 1))/((mu3 + x)^2 + y^2 + z^2)^(5/2)
                return Uyz
            else
                error("incorrect derivative expression")
            end
        end
    elseif deriv[1] == 'z'
        if a == 1
            Uz = (z*(mu3 - 1))/((mu3 + x)^2 + y^2 + z^2)^(3/2) - (mu3*z)/((mu3 + x - 1)^2 + y^2 + z^2)^(3/2)
            return Uz
        else #second partial
            if deriv[2] == 'x'
                Uzx = (3*mu3*z*(2*mu3 + 2*x - 2))/(2*((mu3 + x - 1)^2 + y^2 + z^2)^(5/2)) - (3*z*(2*mu3 + 2*x)*(mu3 - 1))/(2*((mu3 + x)^2 + y^2 + z^2)^(5/2))
                return Uzx
            elseif deriv[2] == 'y'
                Uzy = (3*mu3*y*z)/((mu3 + x - 1)^2 + y^2 + z^2)^(5/2) - (3*y*z*(mu3 - 1))/((mu3 + x)^2 + y^2 + z^2)^(5/2)
                return Uzy
            elseif deriv[2] == 'z'
                Uzz = (mu3 - 1)/((mu3 + x)^2 + y^2 + z^2)^(3/2) - mu3/((mu3 + x - 1)^2 + y^2 + z^2)^(3/2) - (3*z^2*(mu3 - 1))/((mu3 + x)^2 + y^2 + z^2)^(5/2) + (3*mu3*z^2)/((mu3 + x - 1)^2 + y^2 + z^2)^(5/2)
                return Uzz
            else
                error("incorrect derivative expression")
            end
        end
    else
        error("incorrect derivative expression")
    end
end


text:
module PointObjs
export PointObj

mutable struct PointObj
    class       :: UInt8
    x           :: Float32
    y           :: Float32
    z           :: Float32
    radius      :: Float32
    parent      :: Int32
end

function PointObj( p::Tuple )
    @assert length(p) == 6
    PointObj( UInt8(p[1]), Float32(p[2]), Float32(p[3]), Float32(p[4]), Float32(p[5]), Int32(p[6]) )
end

function PointObj( data::Vector{UInt8} )
    @assert length(data) == 21
    PointObj(   data[1],
             reinterpret(Float32, data[2:5])[1],
             reinterpret(Float32, data[6:9])[1],
             reinterpret(Float32, data[10:13])[1],
             reinterpret(Float32, data[14:17])[1],
             reinterpret(Int32, data[18:21])[1])
end

function get_parent( self::PointObj )
    self.parent
end

function serialize(self::PointObj)
    io = IOBuffer(read=false, write=true, maxsize=21)
    write(io, self.class)
    write(io, self.x)
    write(io, self.y)
    write(io, self.z)
    write(io, self.radius)
    write(io, self.parent)
    take!(io)
end

function Base.String(self::PointObj)
    "$(self.class) $(self.x) $(self.y) $(self.z) $(self.radius) $(self.parent)"
end

function Base.isequal(self::PointObj, other::PointObj)
    self.class == other.class &&
    self.x == other.x &&
    self.y == other.y &&
    self.z == other.z &&
    self.parent == other.parent
end

function Base.:(==)(self::PointObj, other::PointObj) isequal(self, other) end

function euclidean_distance(self::PointObj, other::PointObj)
    norm( [self.x - other.x, self.y - other.y, self.z - other.z] )
end

end # end of module


text:
@enum ResponseType RELATIVE ABSOLUTE

type Led <: Control
    info::ControlInfo
    min::Float64
    max::Float64
end

function Led(;min=0.0, max=1.0, kwargs...)
    Led(ControlInfo(; kwargs...), min, max)
end

typestring(::Led, landscape::Bool) = "led"

function addxmlinfo!(xml, led::Led)
    set_attribute(xml, "scalef", led.min)
    set_attribute(xml, "scalet", led.max)

    nothing
end

type Push <: Control
    info::ControlInfo
    onval::Float64
    offval::Float64
    local_off::Bool
end

function Push(;offval=0.0, onval=1.0, local_off=false, kwargs...)
    Push(ControlInfo(; kwargs...), onval, offval, local_off)
end

typestring(::Push, landscape::Bool) = "push"

function addxmlinfo!(xml, tog::Push)
    set_attribute(xml, "scalef", tog.offval)
    set_attribute(xml, "scalet", tog.onval)
    set_attribute(xml, "local_off", tog.local_off)

    nothing
end

type Toggle <: Control
    info::ControlInfo
    onval::Float64
    offval::Float64
    local_off::Bool
end

function Toggle(;offval=0.0, onval=1.0, local_off=false, kwargs...)
    Toggle(ControlInfo(; kwargs...), onval, offval, local_off)
end

typestring(::Toggle, landscape::Bool) = "toggle"

function addxmlinfo!(xml, tog::Toggle)
    set_attribute(xml, "scalef", tog.offval)
    set_attribute(xml, "scalet", tog.onval)
    set_attribute(xml, "local_off", tog.local_off)

    nothing
end

type Label <: Control
    info::ControlInfo
    text::ASCIIString
    textsize::Int
    orientation::Orientation
end

function Label(; text="", textsize=24, orientation=HORIZONTAL, kwargs...)
    Label(ControlInfo(; kwargs...), text, textsize, orientation)
end

function typestring(label::Label, landscape::Bool)
    if landscape
        label.orientation == VERTICAL ? "labelh" : "labelv"
    else
        label.orientation == VERTICAL ? "labelv" : "labelh"
    end
end

function addxmlinfo!(xml, label::Label)
    set_attribute(xml, "text", base64encode(label.text))
    set_attribute(xml, "size", label.textsize)

    nothing
end

type Fader <: Control
    info::ControlInfo
    min::Float64
    max::Float64
    orientation::Orientation
    responsetype::ResponseType
    inverted::Bool
    centered::Bool
end

function Fader(;min=0.0, max=1.0, orientation=HORIZONTAL, responsetype=ABSOLUTE,
          inverted=false, centered=false, kwargs...)
    Fader(ControlInfo(; kwargs...), min, max, orientation, responsetype, inverted, centered)
end

function typestring(fader::Fader, landscape::Bool)
    if landscape
        fader.orientation == VERTICAL ? "faderh" : "faderv"
    else
        fader.orientation == VERTICAL ? "faderv" : "faderh"
    end
end

function addxmlinfo!(xml, fader::Fader)
    set_attribute(xml, "response", lowercase(string(fader.responsetype)))
    set_attribute(xml, "scalef", fader.min)
    set_attribute(xml, "scalet", fader.max)
    set_attribute(xml, "inverted", fader.inverted)
    set_attribute(xml, "centered", fader.centered)

    nothing
end

type Encoder <: Control
    info::ControlInfo
    min::Float64
    max::Float64
end

function Encoder(;min=0.0, max=1.0, kwargs...)
    Encoder(ControlInfo(; kwargs...), min, max)
end

typestring(enc::Encoder, landscape::Bool) = "encoder"

function addxmlinfo!(xml, enc::Encoder)
    set_attribute(xml, "scalef", enc.min)
    set_attribute(xml, "scalet", enc.max)

    nothing
end


text:
################################################################################
#  Copyright 2020, Marta Vanin, Tom Van Acker                                  #
################################################################################
# PowerModelsDistributionStateEstimation.jl                                    #
# An extention package of PowerModels(Distribution).jl for Static Power System #
# State Estimation.                                                            #
################################################################################
"""
    objective_mc_se
"""
function objective_mc_se(pm::_PMD.AbstractUnbalancedPowerModel)
    return JuMP.@objective(pm.model, Min,
    sum(
        sum(
            sum(_PMD.var(pm, nw, :res, i)[idx] for idx in 1:length(_PMD.var(pm, nw, :res, i)) )
        for i in _PMD.ids(pm, nw, :meas))
    for (nw, nw_ref) in _PMD.nws(pm) )
    )
end


text:
using Flux: ADAM, Momentum

#=
Each self-play job is independent of all others; it takes the latest network
snapshot, produces a game and makes it available to the training job by
writing it to a shared replay buffer.
=#

using Base.Threads

struct SelfplayCallback
    period::Integer
    fn::Function
end

(x::SelfplayCallback)() = x.fn()

function getnextactions(
    nodes::Vector{Node},
    plans::AbstractArray;
    num_moves,
    num_steps,
    temp_fn,
    )
    if !isempty(plans)
        return getfield.(popfirst!.(plans), :idx)
    end
    select_action.(
        nodes;
        num_moves=num_moves,
        num_steps=num_steps,
        temp_fn=temp_fn
    )
end

#=
Core Monte Carlo Tree Search algorithm.
To decide on an action, we run N simulations, always starting at the root of
the search tree and traversing the tree according to the UCB formula until we
reach a leaf node.
=#

function select_action(
    node::Node;
    num_moves,
    num_steps,
    temp_fn,
    )
    softmax_sample(
        [(child.visit_count, action) for (action, child) in node.children],
        temp_fn(num_moves, num_steps)
    )
end

# Select the child with the highest UCB score.
function select_child!(
    node::Node,
    min_max_stats::MinMaxStats,
    config,
    )
    score_action_child = [
        (
            ucbscore(node, child, min_max_stats, config),
            action,
            child
        )
        for (action, child) in node.children
    ]
    scores = [x[1] for x in score_action_child]
    max_score = maximum(scores)
    selected = rand(score_action_child[scores .== max_score ])
    _, action, child = selected
    action, child
end

#=
The score for a node is based on its value, plus an exploration bonus based on
the prior.
=#

function ucbscore(
    parent::Node,
    child::Node,
    min_max_stats::MinMaxStats,
    config;
    action=nothing,
    )
    discount = config.discount
    pb_c_init = config.pb_c_init
    pb_c_base = config.pb_c_base
    pb_c = log((parent.visit_count + pb_c_base + 1) / pb_c_base) + pb_c_init
    pb_c *= sqrt(parent.visit_count) / (child.visit_count + 1)
    prior_score = pb_c * child.prior
    value_score = 0
    neg_nextstate_value = -1*(discount * normalize(min_max_stats, calc_value(child)))
    if child.visit_count > 0
        value_score = child.reward + neg_nextstate_value
    end
    ucb_score = prior_score + value_score
    if !isnothing(action)
        if parent.state.board == [0 2 0; 0 0 1; 0 1 0]
            @show action, value_score, child.value_sum, child.visit_count #child.prior, prior_score, ucb_score
        end
    end
    ucb_score
end

#=
At the start of each search, we add dirichlet noise to the prior of the root
to encourage the search to explore new actions.
=#
using Distributions

function addexplorationnoise(
    node::Node;
    root_dirichlet_alpha,
    root_exploration_fraction
    )
    actions = collect(keys(node.children))
    noise = rand(
        Dirichlet(
            length(actions),
            root_dirichlet_alpha
        )
    )
    for (a, n) in zip(actions, noise)
        node.children[a].prior = (
            node.children[a].prior
            * (1 - root_exploration_fraction)
            + (n * root_exploration_fraction)
        )
    end
end

function add_exploration_noise(
    config::MuZeroConfig,
    nodes::Vector{Node}
    )
    addexplorationnoise.(
        nodes;
        root_dirichlet_alpha=config.root_dirichlet_alpha,
        root_exploration_fraction=config.root_exploration_fraction
    )
end

function add_exploration_noise(config::MuZeroConfig, node::Node)
    addexplorationnoise(
        node;
        root_dirichlet_alpha=config.root_dirichlet_alpha,
        root_exploration_fraction=config.root_exploration_fraction
    )
end

# function play_me()
#     network = test_network
#     g = config.game_type()
#     to_play = PLAYER_ID
#     output = initial_inference(network, current_image(g)) |> SingleNetworkOutput
#     function get_next(output)
#         to_play = (to_play == PLAYER_ID) ? OPPONENT_ID : PLAYER_ID
#         if to_play == OPPONENT_ID
#             println(); print_game(g); println();
#             next = OPPONENT_ID
#             while !islegal(g, next)
#                 println(); print("Your move [1-9]: ")
#                 next = parse(Int, readline())
#             end
#             return next
#         end
#         greedy_action(output)
#     end
#     while !is_terminal(g)
#         next = get_next(output)
#         if !islegal(g, next)
#             println("bot made illegal move... ($next)")
#             return
#         end
#         apply!(g, next)
#         output = recurrent_inference(network, hidden_state(output), [next]) |> SingleNetworkOutput
#     end
#     println(); print_game(g); println();
#     println(); println("Done.")
#     g
# end

# function dynamics_passed(game_type, network)
#     g = game_type()
#     # @todo This needs attention (player 1 is can only be first for certain games)
#     player = rand(1:2) == PLAYER_ID ? PLAYER_ID : OPPONENT_ID
#     output = NetworkOutput(initial_inference(network, current_image(g)))
#     function get_next(output)
#         @assert false "Fix logic below for OPPONENT_ID"
#         player *= OPPONENT_ID
#         if player == OPPONENT_ID
#             return rand(legal_actions(g))
#         end
#         greedy_action(output)
#     end
#     while !is_terminal(g)
#         next = get_next(output)
#         if !islegal(g, next)
#             return false, g
#         end
#         apply!(g, next)
#         output = recurrent_inference(network, output.hidden_state, [next]) |> SingleNetworkOutput
#     end
#     true, g
# end

function check_dynamics(config::MuZeroConfig, network)
    passes = 0
    fails = 0
    failed_games = []
    for _ in 1:1000
        res, g = dynamics_passed(config.game_type, network)
        if res
            passes += 1
        else
            push!(failed_games, g)
            fails += 1
        end
    end
    passes, fails, failed_games
end

#=
At the end of a simulation, we propagate the evaluation all the way up the
tree to the root.
=#

function backpropagate!(
    search_path::Vector{Node},
    last_value::Float32,
    to_play::Integer,
    discount::Float32,
    min_max_stats::MinMaxStats
    )
    value = last_value
    #backprops = []
    for node in reverse(search_path)
        # @todo check on the logic below and make sure value_sum is correct
        node.value_sum += (gettoplay(node).idx == to_play) ? value : -value
        #node.info = "$(node.to_play == to_play ? "v" : "-v") -> $(node.value_sum)"
        node.visit_count += 1
        update!(min_max_stats, calc_value(node))
        value = Float32(-1)*node.reward + discount * value
        #push!(backprops, (state=key(node.state), value_sum=node.value_sum))
    end
    #backprops
    return
end

# MuZero training is split into two independent parts: Network training and
# self-play data generation.
# These two parts only communicate by transferring the latest network checkpoint
# from the training to the self-play, and the finished games from the self-play
# to the training.

using Base.Threads: @spawn, @threads

launch_job(f, args...) = f(args...)


text:
##############################################################################################
module TestSingleIterationSmootherExps
##############################################################################################
##############################################################################################
# imports and exports
using DataAssimilationBenchmarks.SmootherExps
using JLD2, Statistics

##############################################################################################
##############################################################################################
# run and analyze the IEnKS for state estimation with the Lorenz-96 model
# arguments are
# time_series, method, seed, nanl, lag, shift, mda, obs_un, obs_dim, γ, N_ens, infl = args

function run_sda_smoother_state_L96()
    try
        path = joinpath(@__DIR__, "../src/data/time_series/")
        L96_test_ts = "L96_time_series_seed_0000_dim_40_diff_0.000_F_08.0" *
                      "_tanl_0.05_nanl_05000_spin_1500_h_0.050.jld2"

        args = (path * L96_test_ts, "etks", 0, 3500, 10, 1, false,
                1.0, 40, 1.00, 21, 1.02)
        single_iteration_state(args)
        true
    catch
        false
    end
end

function analyze_sda_smoother_state_L96()
    try
        # test if the filter RMSE for standard simulation falls below adequate threshold
        path = joinpath(@__DIR__, "../src/data/etks-single-iteration/")
        data = load(path * "etks-single-iteration_L96_state_seed_0000_diff_0.000_sysD_40_" *
                    "obsD_40_obsU_1.00_gamma_001.0_nanl_03500_tanl_0.05_h_0.05_lag_010_" *
                    "shift_001_mda_false_nens_021_stateInfl_1.02.jld2")
        filt_rmse = data["filt_rmse"]
        post_rmse = data["post_rmse"]

        # note, we use a small burn-in to reach more regular cycles
        if mean(filt_rmse[501:end]) < 0.2 && mean(post_rmse[501:end]) < 0.10
            true
        else
            false
        end
    catch
        false
    end
end

function run_mda_smoother_state_L96()
    try
        path = joinpath(@__DIR__, "../src/data/time_series/")
        L96_test_ts = "L96_time_series_seed_0000_dim_40_diff_0.000_F_08.0" *
                      "_tanl_0.05_nanl_05000_spin_1500_h_0.050.jld2"

        args = (path * L96_test_ts, "etks", 0, 3500, 10, 1, true,
                1.0, 40, 1.00, 21, 1.02)
        single_iteration_state(args)
        true
    catch
        false
    end
end

function analyze_mda_smoother_state_L96()
    try
        # test if the filter RMSE for standard simulation falls below adequate threshold
        path = joinpath(@__DIR__, "../src/data/etks-single-iteration/")
        data = load(path * "etks-single-iteration_L96_state_seed_0000_diff_0.000_sysD_40_" *
                    "obsD_40_obsU_1.00_gamma_001.0_nanl_03500_tanl_0.05_h_0.05_lag_010_" *
                    "shift_001_mda_true_nens_021_stateInfl_1.02.jld2")
        filt_rmse = data["filt_rmse"]
        post_rmse = data["post_rmse"]

        # note, we use a small burn-in to reach more regular cycles
        if mean(filt_rmse[501:end]) < 0.2 && mean(post_rmse[501:end]) < 0.10
            true
        else
            false
        end
    catch
        false
    end
end

##############################################################################################
# run and analyze the IEnKS for joint state-parameter estimation with the Lorenz-96 model
# time_series, method, seed, nanl, lag, shift, mda, obs_un, obs_dim, γ,
# param_err, param_wlk, N_ens, state_infl, param_infl = args

function run_sda_smoother_param_L96()
    try
        path = joinpath(@__DIR__, "../src/data/time_series/")
        L96_test_ts = "L96_time_series_seed_0000_dim_40_diff_0.000_F_08.0" *
                      "_tanl_0.05_nanl_05000_spin_1500_h_0.050.jld2"

        args = (path * L96_test_ts, "etks", 0, 3500, 10, 1, false, 1.0, 40, 1.0,
                0.10, 0.0010, 21, 1.02, 1.0)
        single_iteration_param(args)
        true
    catch
        false
    end
end

function analyze_sda_smoother_param_L96()
    try
        # test if the filter RMSE for standard simulation falls below adequate threshold
        path = joinpath(@__DIR__, "../src/data/etks-single-iteration/")
        data = load(path * "etks-single-iteration_L96_param_seed_0000_diff_0.000_sysD_41" *
                    "_obsD_40_obsU_1.00_gamma_001.0_paramE_0.10_paramW_0.0010_nanl_03500" *
                    "_tanl_0.05_h_0.05_lag_010_shift_001_mda_false_nens_021_" *
                    "stateInfl_1.02_paramInfl_1.00.jld2")
        filt_rmse = data["filt_rmse"]
        post_rmse = data["post_rmse"]
        para_rmse = data["param_rmse"]

        # note, we use a small burn-in to reach more regular cycles
        if (mean(filt_rmse[501:end]) < 0.2) && (mean(post_rmse[501:end]) < 0.10) &&
            (mean(para_rmse[501:end]) < 0.01)
            true
        else
            false
        end
    catch
        false
    end
end

function run_mda_smoother_param_L96()
    try
        path = joinpath(@__DIR__, "../src/data/time_series/")
        L96_test_ts = "L96_time_series_seed_0000_dim_40_diff_0.000_F_08.0" *
                      "_tanl_0.05_nanl_05000_spin_1500_h_0.050.jld2"

        args = (path * L96_test_ts, "etks", 0, 3500, 10, 1, true, 1.0, 40, 1.0,
                0.10, 0.0010, 21, 1.02, 1.0)
        single_iteration_param(args)
        true
    catch
        false
    end
end

function analyze_mda_smoother_param_L96()
    try
        # test if the filter RMSE for standard simulation falls below adequate threshold
        path = joinpath(@__DIR__, "../src/data/etks-single-iteration/")
        data = load(path * "etks-single-iteration_L96_param_seed_0000_diff_0.000_sysD_41_" *
                    "obsD_40_obsU_1.00_gamma_001.0_paramE_0.10_paramW_0.0010_nanl_03500_" *
                    "tanl_0.05_h_0.05_lag_010_shift_001_mda_true_nens_021_stateInfl_1.02_" *
                    "paramInfl_1.00.jld2")
        filt_rmse = data["filt_rmse"]
        post_rmse = data["post_rmse"]
        para_rmse = data["param_rmse"]

        # note, we use a small burn-in to reach more regular cycles
        if (mean(filt_rmse[501:end]) < 0.2) && (mean(post_rmse[501:end]) < 0.10) &&
            (mean(para_rmse[501:end]) < 0.01)
            true
        else
            false
        end
    catch
        false
    end
end

##############################################################################################
# end module

end

max_stars_repo_name: cgrudz/DataAssimilationBenchmarks

max_issues_repo_name: cgrudz/DataAssimilationBenchmarks


text:
r = try
    parse(Int,ARGS[begin])
  catch
    error("First arg is required and must be an integer")
end

wh_ratio = 8//15 # a guess

h = 2r
w = h // wh_ratio
origin = (h//2, w//2)
x_origin, y_origin = origin

get_char(x::Number, y::Number) = begin
  d = √((x_origin - x)^2 + ((y_origin - y) * wh_ratio)^2)
  d = Int(floor(d))
  d = d + 1 # treat distance as slightly larger for better edge rendering
  if d < r
    'o'
  elseif d == r
    θorigin = atan((y_origin - y)//(x_origin - x))
    θedge = π/2 + θorigin # 0 to π
    if π/8 ≤ θedge < 3π/8
      '\\'
    elseif 3π/8 ≤ θedge ≤ 5π/8
      '-'
    elseif 5π/8 < θedge ≤ 7π/8
      '/'
    else
      '|'
    end
  else
    '.'
  end
end

for x = 1:h
  for y = 1:w
    print(get_char(x,y))
  end
  println()
end

max_stars_repo_name: spenserblack/ascii-circle

max_issues_repo_name: spenserblack/ascii-circle

max_forks_repo_name: spenserblack/ascii-circle


text:
### A Pluto.jl notebook ###
# v0.19.4

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 74bc6fb9-8b52-46bf-b9cc-4a63883c5196
using Pkg; Pkg.activate("/Users/jj/JuliaProjects/LaserLab/")

# ╔═╡ f9cb42a4-188e-4929-8c23-b84cfe9df792
begin
	using PlutoUI
	using Statistics
	using StatsBase
	using Plots
	using Unitful
	using Peaks
	using Glob
	using FFTW
	using DSP
	using DataFrames
	using CSV
	using LsqFit
end

# ╔═╡ d87e4feb-a0d7-47ad-b9e2-2c92f5e467aa
import Unitful:
    nm, μm, mm, cm, m, km,
    mg, g, kg,
    ps, ns, μs, ms, s, minute, hr, d, yr, Hz, kHz, MHz, GHz,
    eV,
    μJ, mJ, J,
	μW, mW, W,
    A, N, mol, mmol, V, L, M

# ╔═╡ 948d5fee-cc84-11ec-1981-e71e1585c8bb
PlutoUI.TableOfContents(title="PMT Laser Analysis", indent=true)

# ╔═╡ 1a05bb05-5231-4168-8eab-ef84660deb9e
function ingredients(path::String)
	# this is from the Julia source code (evalfile in base/loading.jl)
	# but with the modification that it returns the module instead of the last object
	name = Symbol(basename(path))
	m = Module(name)
	Core.eval(m,
        Expr(:toplevel,
             :(eval(x) = $(Expr(:core, :eval))($name, x)),
             :(include(x) = $(Expr(:top, :include))($name, x)),
             :(include(mapexpr::Function, x) = $(Expr(:top, :include))(mapexpr, $name, x)),
             :(include($path))))
	m
end

# ╔═╡ fca3e7c4-c160-43a1-ada1-78836dd2fe43
lfi = ingredients("../src/LaserLab.jl")

# ╔═╡ e56888f9-3747-4085-853f-6114956b72d3
md"""
# Analysis
"""

# ╔═╡ abc85752-30e4-48ac-87da-8b1050e05e43
begin
	tl = 1.0μs
	lf = uconvert(MHz, 1.0/tl)

md" - Parameters:
 - laser frequency (MHz) = $lf
 - Laser pulses every    = $tl
"
end

# ╔═╡ cc7ea918-066b-4485-80d2-64eba6af76a7
begin
	pmtdir  ="/Users/jj/JuliaProjects/LaserLab/labdata/PMT2"
	odir  = "/Users/jj/JuliaProjects/LaserLab/data/PMT"
	dfiles ="*.trc"
	dplots ="*.png"
	md"""
	CMOS dir = $pmtdir
	"""
	#close(io)
end

# ╔═╡ c907c77c-0291-4c80-bcba-ccbcd8e864cb
let
	readdir(pmtdir)
	dirs = lfi.LaserLab.getbolddirs(pmtdir)
	md""" Select experiment : $(@bind sexp Select(dirs))"""
end

# ╔═╡ a4fdfa7a-57bc-4d9f-ab90-bf99d3b48277
#sexp = "Bold_068_A1"

# ╔═╡ 740681e5-7dc2-4fa1-bfb7-1ee5557ab885
let
	path = joinpath(pmtdir,sexp)
	readdir(path)
	dirs = lfi.LaserLab.getbolddirs(path)
	md""" Select filter : $(@bind sflt Select(dirs))"""
end

# ╔═╡ d3009c78-6318-4b5c-adfe-fb5905414081
begin
	xpath = joinpath(pmtdir,sexp,sflt)
	xfiles = Glob.glob(dfiles, xpath)
	md""" Select a file to analyze (1 to $(length(xfiles))): $(@bind xnf NumberField(0:length(xfiles), default=1))"""
end

# ╔═╡ 513a0b87-409e-4e1e-b872-6c37fd84e80b
xfiles

# ╔═╡ db6868e3-9f96-49ef-b343-309a4b778591
begin
	zpath = joinpath(pmtdir,xfiles[xnf])
	zio = open(zpath, "r")
	WAVEDESC = lfi.LaserLab.wavedesc(zio)
	address = lfi.LaserLab.scope_rawdata(zio, WAVEDESC)
	wvfm = lfi.LaserLab.xydata(zio, address)
	tw = round(wvfm.time[end] - wvfm.time[1], sigdigits=2) * μs
	sr = lfi.LaserLab.sampling_rate(wvfm, tw)
	sp = lfi.LaserLab.sampling_period(wvfm, tw)
	stats = lfi.LaserLab.wstats(wvfm; nsigma=3.0)
	close(zio)
	md"""
	- Waveform length = $(tw) μs
	- sampling rate   = $(round(sr/GHz, sigdigits=2)) GHz
	- sampling period = $(round(sp/ns, sigdigits=4)) ns
	- baseline at     => $(round(stats.mean, sigdigits=3)) mV
	- baseline std    => $(round(stats.std, sigdigits=3)) mV
	"""
end

# ╔═╡ 880dfedd-2df4-4d43-bd66-13962f910055
typeof(wvfm.ampl)

# ╔═╡ b30c1537-465b-4938-97be-17219bef9402
md""" Select max filter frequency (in MHz): $(@bind ffmhz NumberField(0:1500, default=300))"""

# ╔═╡ 92e61e22-5666-43d9-a965-ef1d030a644b
begin
	ffhz = ffmhz*1e+6
    fwvfm = lfi.LaserLab.filter_signal_lp(wvfm, tw; filtertype= "Butterworth",
	                                      flhz=ffhz)
	fstats = lfi.LaserLab.wstats(fwvfm; nsigma=3.0)
end

# ╔═╡ afeb9110-2058-4f09-9d27-4872f813256e
md""" Click to plot waveforms: $(@bind zpwf CheckBox(default=true))"""

# ╔═╡ 00778f59-2f53-48db-82a0-2adc18fd1aab
if zpwf
	pwvfm = lfi.LaserLab.plot_waveform(wvfm, stats; sct=true, window=false, ws=1, we=50000, trace=false)
	pfft = lfi.LaserLab.plot_fft(wvfm,tw; sct=true, window=false, ws=1, we=50000)
	pfwvfm = lfi.LaserLab.plot_waveform(fwvfm, fstats; sct=true, window=false, ws=1, we=50000, trace=false)

	plot(size=(850,550), pwvfm, pfft, pfwvfm, layout=(1,3), titlefontsize=8)
end

# ╔═╡ cb696494-5030-4568-a65b-aad9cbc01cf1
wvfm

# ╔═╡ 99e8afb0-f38e-4b79-88d2-6e43828b6e2b
md""" Select prominence cut: $(@bind spcut NumberField(0:10, default=2))"""

# ╔═╡ 5d2cce97-ad00-46be-b086-8c4a35cdb008
begin
	promsel = Float64(spcut)  #prominence cut

	md"
	- Peak Prominence cut = $sprom8n5k1m
	"
	rfpeaks = lfi.LaserLab.select_filtered_peaks(fwvfm, fstats.thrp;
                                       promsel=promsel, wsel=0.0)

	if rfpeaks != Nothing
		wvflt, fpeaks = rfpeaks
		lfi.LaserLab.plot_filtered_peaks(wvflt, fpeaks)
	end
end

# ╔═╡ bd6faf80-4f68-44a0-9ad6-319a2b4f2316
md"## Optimizing prominence cut"

# ╔═╡ 551cd586-e702-41fb-989c-5139632e29b6
md""" Click to compute proms histos: $(@bind zprom CheckBox(default=false))"""

# ╔═╡ 68bce42f-1417-4a8c-b8e7-05168fe33b9f
if zprom
	md""" Select min prominence cut: $(@bind spcutm NumberField(0.5:10.0, default=0.5))"""
end

# ╔═╡ c135833d-e9ae-41de-84a3-29b7b2e023ba
if zprom
	md""" Select how many files to analyze: $(@bind nnf NumberField(1:length(xfiles), default=ceil(length(xfiles)/2)))"""
end

# ╔═╡ cb8c269e-67f7-4239-bd56-066fce1886d2
if zprom
	promin = Float64(spcutm)
	if promin < 0.5
		promin = 0.5
	end
	proms, widths = lfi.LaserLab.proms(xfiles, 1, nnf;
	                                   flhz=ffhz, promsel=promin, nsigma=3.0)
	hproms  = lfi.LaserLab.hist1d(proms, 20, 0.0, 10.0)
	hwidths = lfi.LaserLab.hist1d(widths, 20, 5.0, 15.0)

	phproms = lfi.LaserLab.hist1d(hproms, "Prominence (promsel=$promin)")
    phwidths =lfi.LaserLab.hist1d(hwidths, "Width")
	plot(phproms, phwidths, layout=(1,2), titlefontsize=8)
end

# ╔═╡ 05fe45dc-0745-4d21-8648-a1257625fc72
md"## Selection of events"

# ╔═╡ e6d7c71d-7eac-4f4e-af04-89d4d92d54f4
md""" Select  prominence cut: $(@bind zpcut NumberField(0.5:10.0, default=2.0))"""

# ╔═╡ 0e1b6a85-af4b-4493-92e2-33aa4e1374e9

md""" Select how many files to analyze: $(@bind nnf2 NumberField(1:length(xfiles), default=ceil(length(xfiles)/2)))"""

# ╔═╡ 67dac32c-11df-4843-9e52-e4845bc1a2d4
md""" Click to compute time histos: $(@bind ztimes CheckBox(default=false))"""

# ╔═╡ 01c601a6-1531-41df-b011-614c0b70d476
tl/1.0μs

# ╔═╡ 931aad69-ae0f-4168-8245-619590a046a2
md""" Select  prominence cut: $(@bind pcut NumberField(0.5:10.0, default=2.5))"""

# ╔═╡ dc3b9b68-4478-4f53-875d-3ecb6449edb8
coefw

# ╔═╡ 268f59b5-5b40-4ba0-a064-edb534bb3bab
md"""
# Functions
"""

# ╔═╡ 5c11564b-ba5c-49b4-b23e-6bef4fb242c1
function subtract_laser(tv::Vector{Float64}, ltus::Float64)
	tv - ltus * floor.(tv/ltus)
end

# ╔═╡ d557f2b0-7ccf-45c7-a93b-5bf499cfafad
if rfpeaks != Nothing
	subtract_laser(fpeaks.xs, tl/1.0μs)
end

# ╔═╡ 4c222697-e58c-4199-8114-38bc5a457cbb
function tfit(htime, pa0, i0, l0, ffun, mffun)

	tdata = htime.centers
	vdata = htime.weights

	fitb = curve_fit(mffun, tdata[i0:l0], vdata[i0:l0], pa0)

	fit = curve_fit(mffun, tdata[i0:l0], vdata[i0:l0], pa0)
	coef(fit), stderror(fit), ffun.(tdata, coef(fit)...)
end

# ╔═╡ 2e4d3f1b-bf12-4729-adcb-8ffb95902ec6
function tfit_bkg(htime; pa0=[100.0, 0.5], i0=10, l0=20)

	ffun(t, a,b) = a + b * t
	mffun(t, p) = p[1] .+ t .* p[2]

	tfit(htime, pa0, i0, l0, ffun, mffun)
end

# ╔═╡ 0d61f471-d272-4dfb-9cf9-45e08d7b1c82
function tfit_sgn(htime; cbkg, pa0=[100.0, 0.5], i0=10, l0=20)

	bkg(t) = cbkg[1] + cbkg[2] * t

	ffun(t, N1, λ) = N1 * exp(-t/λ)  +  bkg(t)
	mffun(t, p) = p[1] * exp.(-t/p[2]) +  bkg.(t)

	tfit(htime, pa0, i0, l0, ffun, mffun)
end

# ╔═╡ ae323f97-3bda-4933-9652-68224c939938
function tfit_sgn2(htime; cbkg, csgn, pa0=[100.0, 0.5], i0=1, l0=2)

	bkg(t) = cbkg[1] + cbkg[2] * t
	sgn(t) = csgn[1] * exp(-t/csgn[2])  +  bkg(t)
	ffun(t, N1, λ) = N1 * exp(-t/λ)  +  sgn(t)

	mffun(t, p) = p[1] * exp.(-t/p[2]) +  sgn.(t)

	tfit(htime, pa0, i0, l0, ffun, mffun)
end

# ╔═╡ cce4cc88-be0b-47a2-b37f-08cc0e40cf70
function tfit_2exp(htime; pa0=[1000.0, 10., 1000.0, 10.0], i0, l0)

	ffun(t, N1, λ1, N2, λ2) = N1 * exp(-t/λ1) + N2 * exp(-t/λ2)
	mffun(t, p) = p[1] * exp.(-t/p[2]) + p[3] * exp.(-t/p[4])

	tfit(htime, pa0, i0, l0, ffun, mffun)
end

# ╔═╡ 368cd59a-e9e2-436f-872b-92ef5bb8c1d8
function fit_peaks(htime; pa0=[100.0, 0.5], i0=1, il=-1)
	expo(t, N, λ) = N*exp(-t/λ)
	mexp(t, p) = p[1] * exp.(-t/p[2])
	tdata = htime.centers
	vdata = htime.weights
	if il < 0
		il = length(tdata)
	end
	fit = curve_fit(mexp, tdata[i0:il], vdata[i0:il], pa0)
	coef(fit), stderror(fit), expo.(tdata, coef(fit)...)
end

# ╔═╡ 7fa04ed2-113a-4166-9f29-eef98393ab6d
function create_dir!(dir)
	if isdir(dir) == false
		mkdir(dir)
	end
end

# ╔═╡ e44d66e8-bac5-4ec0-9c94-aa0dbbe05d06
function output_dirs!(odir, sexp)
	create_dir!(joinpath(odir, sexp))
	csvdir = joinpath(odir, sexp,"csv")
	pngdir = joinpath(odir, sexp, "png")
	create_dir!(csvdir)
	create_dir!(pngdir)
	csvdir, pngdir
end

# ╔═╡ ce8097c9-9608-49c8-aa5b-7338d090d743
begin
	#namex = split(sexp,"_")
	csvdir, pngdir = output_dirs!(odir, sexp)

	md"""
	- csv dir = $csvdir
	- png dir = $pngdir
	"""
end

# ╔═╡ 8436f3b4-43eb-40bd-afa4-ad3dd87310ea
pmtdf = lfi.LaserLab.load_df_from_csv(csvdir, string(sflt,".csv"), lfi.LaserLab.enG);

# ╔═╡ 3e2d1b5f-de3a-4149-a90f-cf56a2c5a244
pmtFpr = filter(row -> row[:pr] > pcut, pmtdf);

# ╔═╡ 12d461dd-1e80-4aab-90f4-8fd58344ee73
begin
	hpr = lfi.LaserLab.h1d(pmtFpr.pr, 20, 0.0, 10.0)
	phpr = lfi.LaserLab.plot_h1d(hpr, "prom")
	htpr = lfi.LaserLab.h1d(pmtFpr.ts, 20, 0.0, tl/1.0μs)
	phtpr = lfi.LaserLab.plot_h1d(htpr, "times")
	hws = lfi.LaserLab.h1d(pmtFpr.ws, 20, 5.0, 10.0)
	phws = lfi.LaserLab.plot_h1d(hws, "width")
	hvs = lfi.LaserLab.h1d(pmtFpr.vs, 20, 0.0, 20.0)
	phvs = lfi.LaserLab.plot_h1d(hvs, "voltage")
	Nothing
end

# ╔═╡ 39f8b749-3c41-4bff-9140-22c0be92e8dc
begin
	hprx = lfi.LaserLab.h1d(pmtdf.pr, 20, 0.0, 10.0)
	phprx = lfi.LaserLab.plot_h1d(hpr, "prom")
end

# ╔═╡ 8adbeffc-5ac4-43da-b668-9bc17aca8bfc
lfi.LaserLab.plot_h1d(htpr, "times"; i0=2)

# ╔═╡ a3843632-e13c-47df-a7c7-72b4d841dbad
plot(phpr, phtpr, phws, phvs, size=(750,750), layout=(2,2),titlefontsize=8)

# ╔═╡ adafbbc7-9535-4d04-96e3-27cb386770a1
coefx, stderx, tftx = tfit_bkg(htpr; pa0=[10.0, 0.5], i0=13, l0=20)

# ╔═╡ 1146ca40-9d0d-4779-8d62-1ae4f991ada6
coefy, stdery, tfty = tfit_sgn(htpr; cbkg=coefx, pa0=[1000.0, 0.5], i0=2, l0=20)

# ╔═╡ c6e5ed98-aecd-4bba-9911-e3ecbc57bfd2
coefy

# ╔═╡ e90932f1-9a1b-4a77-8750-1e8b1322b788
pngpath = joinpath(pngdir, string(sflt,".png"))

# ╔═╡ fc65ca7f-aac8-49e5-8c13-2c47fb6a4a83
function runsel(csvf::Vector{String}, fi::Integer, fe::Integer;
	            flhz::Float64, promsel::Float64, tlmus::Float64, nsigma::Float64)

	spks = lfi.LaserLab.select_peaks(csvf, fi, fe; flhz=flhz, promsel=promsel,
	                                 nsigma=nsigma)
	df = DataFrame()
	df.pr = reduce(vcat,[spks[i].proms  for i in 1:length(spks)])
	df.ws    = reduce(vcat,[spks[i].widths  for i in 1:length(spks)])
	df.ts   = reduce(vcat,[subtract_laser(spks[i].xs, tlmus) for i in 1:length(spks)])
	df.vs   = reduce(vcat,[spks[i].ys  for i in 1:length(spks)])
	df
end

# ╔═╡ 799cb553-5093-4efc-91c6-5dc4161e6961
if ztimes
	ppcut = Float64(zpcut)
	tdf = runsel(xfiles, 1, nnf2;flhz=ffhz, promsel=ppcut, tlmus=tl/1.0μs, nsigma=3.0)

	dfpath = joinpath(csvdir,string(sflt,".csv"))
	CSV.write(dfpath, tdf)
end

# ╔═╡ e3b0ed46-122e-44d3-9a9a-01a69e032421
md"""
## Histogram functions
"""

# ╔═╡ 70217e57-b672-4f17-9d4e-ab27e9a47057
function add_h1d(h1::Histo1d, h2::Histo1d)
	Histo1d(h1.edges, h1.weights + h2.weights, h1.centers)
end

# ╔═╡ ec2cd649-7293-426c-ad9f-7ace1bf53f45
md"""
## Plotting functions
"""

# ╔═╡ 11ee36e8-2eca-4577-93eb-207d1efa4cbb
function plot_fit(htime, coeff, tft; savefig=false, fn="")
	tdata =htime.centers
	vdata =htime.weights
	ps1 = scatter(tdata, vdata, yerr=sqrt.(abs.(vdata)), markersize=2,
				  color = :black,
	    		  label="data",
				  fmt = :png)
	pp = plot(ps1, tdata, tft, lw=2, label="μ = $(round(coeff[2]*1000, sigdigits=2)) ns", fmt = :png)
	xlabel!("t (μs)")
	ylabel!("frequency")
	if savefig
		png(pp, fn)
	end
	return pp
end

# ╔═╡ 2c8543f7-878d-4cdf-bc90-d1a0ef989dd3
function plot_fit2(htime, label, tft; savefig=false, fn="")
	tdata =htime.centers
	vdata =htime.weights
	ps1 = scatter(tdata, vdata, yerr=sqrt.(abs.(vdata)), markersize=2,
				  color = :black,
	    		  label="data",
				  fmt = :png)
	pp = plot(ps1, tdata, tft, lw=2, label=label, fmt = :png)
	xlabel!("t (μs)")
	ylabel!("frequency")
	if savefig
		png(pp, fn)
	end
	return pp
end

# ╔═╡ 3e9b0679-decd-40b0-81a7-94853101ac85
let
	s1 = " cte = $(round(coefx[1]))\n"
   	ptime = plot_fit2(htpr, s1, tftx; savefig=false,fn=pngpath)
	plot(ptime)
end

# ╔═╡ 37ba3d41-8816-470d-9536-8133bcd05cfd
let
	s1 = " cte = $(round(coefx[1]))\n"
	s2 = " μ1 = $(round(coefy[2]*1000)) ns\n"
	label = string(s1,s2)
    ptime2 = plot_fit2(htpr, label, tfty; savefig=true,
	                 fn=pngpath)
	plot(ptime2)
end

# ╔═╡ 76cdb001-b78c-456b-b285-743abc1bf007
begin
	coefz, stderz, tftz = tfit_sgn2(htpr; cbkg=coefx, csgn=coefy, pa0=[1000.0, 0.5], i0=1, l0=2)
	sf1 = " cte = $(round(coefx[1]))\n"
	sf2 = " μ1 = $(round(coefy[2]*1000)) ns \n"
	sf3 = " μ2 = $(round(coefz[2]*1000)) ns \n"
	label = string(sf1,sf2,sf3)
    ptime3 = plot_fit2(htpr, label, tftz; savefig=true,
	                 fn=pngpath)
	plot(ptime3)
end

# ╔═╡ Cell order:
# ╠═74bc6fb9-8b52-46bf-b9cc-4a63883c5196
# ╠═f9cb42a4-188e-4929-8c23-b84cfe9df792
# ╠═d87e4feb-a0d7-47ad-b9e2-2c92f5e467aa
# ╠═948d5fee-cc84-11ec-1981-e71e1585c8bb
# ╠═1a05bb05-5231-4168-8eab-ef84660deb9e
# ╠═fca3e7c4-c160-43a1-ada1-78836dd2fe43
# ╠═e56888f9-3747-4085-853f-6114956b72d3
# ╠═abc85752-30e4-48ac-87da-8b1050e05e43
# ╠═cc7ea918-066b-4485-80d2-64eba6af76a7
# ╠═c907c77c-0291-4c80-bcba-ccbcd8e864cb
# ╠═a4fdfa7a-57bc-4d9f-ab90-bf99d3b48277
# ╠═ce8097c9-9608-49c8-aa5b-7338d090d743
# ╠═740681e5-7dc2-4fa1-bfb7-1ee5557ab885
# ╠═d3009c78-6318-4b5c-adfe-fb5905414081
# ╠═513a0b87-409e-4e1e-b872-6c37fd84e80b
# ╠═db6868e3-9f96-49ef-b343-309a4b778591
# ╠═880dfedd-2df4-4d43-bd66-13962f910055
# ╠═b30c1537-465b-4938-97be-17219bef9402
# ╠═92e61e22-5666-43d9-a965-ef1d030a644b
# ╠═afeb9110-2058-4f09-9d27-4872f813256e
# ╠═00778f59-2f53-48db-82a0-2adc18fd1aab
# ╠═cb696494-5030-4568-a65b-aad9cbc01cf1
# ╠═99e8afb0-f38e-4b79-88d2-6e43828b6e2b
# ╠═5d2cce97-ad00-46be-b086-8c4a35cdb008
# ╠═d557f2b0-7ccf-45c7-a93b-5bf499cfafad
# ╠═bd6faf80-4f68-44a0-9ad6-319a2b4f2316
# ╠═551cd586-e702-41fb-989c-5139632e29b6
# ╠═68bce42f-1417-4a8c-b8e7-05168fe33b9f
# ╠═c135833d-e9ae-41de-84a3-29b7b2e023ba
# ╠═cb8c269e-67f7-4239-bd56-066fce1886d2
# ╠═05fe45dc-0745-4d21-8648-a1257625fc72
# ╠═e6d7c71d-7eac-4f4e-af04-89d4d92d54f4
# ╠═0e1b6a85-af4b-4493-92e2-33aa4e1374e9
# ╠═67dac32c-11df-4843-9e52-e4845bc1a2d4
# ╠═01c601a6-1531-41df-b011-614c0b70d476
# ╠═799cb553-5093-4efc-91c6-5dc4161e6961
# ╠═8436f3b4-43eb-40bd-afa4-ad3dd87310ea
# ╠═39f8b749-3c41-4bff-9140-22c0be92e8dc
# ╠═931aad69-ae0f-4168-8245-619590a046a2
# ╠═3e2d1b5f-de3a-4149-a90f-cf56a2c5a244
# ╠═12d461dd-1e80-4aab-90f4-8fd58344ee73
# ╠═8adbeffc-5ac4-43da-b668-9bc17aca8bfc
# ╠═a3843632-e13c-47df-a7c7-72b4d841dbad
# ╠═e90932f1-9a1b-4a77-8750-1e8b1322b788
# ╠═adafbbc7-9535-4d04-96e3-27cb386770a1
# ╠═3e9b0679-decd-40b0-81a7-94853101ac85
# ╠═1146ca40-9d0d-4779-8d62-1ae4f991ada6
# ╠═37ba3d41-8816-470d-9536-8133bcd05cfd
# ╠═76cdb001-b78c-456b-b285-743abc1bf007
# ╠═dc3b9b68-4478-4f53-875d-3ecb6449edb8
# ╠═c6e5ed98-aecd-4bba-9911-e3ecbc57bfd2
# ╠═268f59b5-5b40-4ba0-a064-edb534bb3bab
# ╠═5c11564b-ba5c-49b4-b23e-6bef4fb242c1
# ╠═2e4d3f1b-bf12-4729-adcb-8ffb95902ec6
# ╠═0d61f471-d272-4dfb-9cf9-45e08d7b1c82
# ╠═ae323f97-3bda-4933-9652-68224c939938
# ╠═cce4cc88-be0b-47a2-b37f-08cc0e40cf70
# ╠═4c222697-e58c-4199-8114-38bc5a457cbb
# ╠═368cd59a-e9e2-436f-872b-92ef5bb8c1d8
# ╠═7fa04ed2-113a-4166-9f29-eef98393ab6d
# ╠═e44d66e8-bac5-4ec0-9c94-aa0dbbe05d06
# ╠═fc65ca7f-aac8-49e5-8c13-2c47fb6a4a83
# ╠═e3b0ed46-122e-44d3-9a9a-01a69e032421
# ╠═70217e57-b672-4f17-9d4e-ab27e9a47057
# ╠═ec2cd649-7293-426c-ad9f-7ace1bf53f45
# ╠═11ee36e8-2eca-4577-93eb-207d1efa4cbb
# ╠═2c8543f7-878d-4cdf-bc90-d1a0ef989dd3

max_stars_repo_name: jjgomezcadenas/LaserLab

max_issues_repo_name: jjgomezcadenas/LaserLab

max_forks_repo_name: jjgomezcadenas/LaserLab


text:
module ChpSim

using Random

import Base.show, Base.convert

export ChpState, cnot!, hadamard!, phase!, measure!, reset_all_qubits!

"""
    ChpState(num_qubits[, bitpack=false])

The state of a quantum stabilizer circuit simulation.  I.e. the simulated state
of a quantum computer after applying only Clifford operations to qubits that all
start in the `|0⟩` state.

Supported operations: `cnot!`, `hadamard!`, `phase!`, `measure!`
"""
struct ChpState{MatrixT<:AbstractMatrix{Bool}, VectorT<:AbstractVector{Bool}}
    n::Int
    x::MatrixT
    z::MatrixT
    r::VectorT
    _tmp::VectorT
    _tmp_s::VectorT

    function ChpState(num_qubits, x, z, r)
        @assert size(x) == (2num_qubits, num_qubits)
        @assert size(z) == (2num_qubits, num_qubits)
        @assert size(r) == (2num_qubits,)
        tmp = similar(r)
        tmp_s = similar(r, 1)
        new{typeof(x), typeof(r)}(num_qubits, x, z, r, tmp, tmp_s)
    end
end

function ChpState(MatrixT::DataType, VectorT::DataType, num_qubits::Integer)
    x = similar(MatrixT, 2num_qubits, num_qubits)
    z = similar(MatrixT, 2num_qubits, num_qubits)
    r = similar(VectorT, 2num_qubits)
    reset_all_qubits!(ChpState(num_qubits, x, z, r))
end

function ChpState(num_qubits::Integer; bitpack::Bool=false)
    ChpState(bitpack ? BitMatrix : Matrix{Bool},
             bitpack ? BitVector : Vector{Bool},
             num_qubits)
end

"""
    reset_all_qubits!(state)

Resets all qubits to |0⟩.
"""
function reset_all_qubits!(state::ChpState)
    state.x .= 0
    state.x[CartesianIndex.(1:state.n, 1:state.n)] .= 1  # Primary diagonal
    state.z .= 0
    # Diagonal offset down by n
    state.z[CartesianIndex.(state.n+1:2state.n, 1:state.n)] .= 1
    state.r .= 0
    state
end

function Base.show(io::IO, self::ChpState)
    write(io, "$(typeof(self))($(self.n), $(self.x), $(self.z), $(self.r))")
end

function Base.show(io::IO, ::MIME"text/plain", self::ChpState)
    function _cell(row::Int, col::Int)::String
        k = self.x[row, col] + 2*self.z[row, col]
        [".", "X", "Z", "Y"][k+1]
    end
    function _row(row::Int)::String
        result = self.r[row] ? "-" : "+"
        result *= join(_cell(row, col) for col in 1:self.n)
    end
    z_obs = (_row(row) for row in 1:self.n)
    sep = '-' ^ (self.n + 1)
    x_obs = (_row(row) for row in self.n+1:2*self.n)
    write(io, join((z_obs..., sep, x_obs...), "\n"))
end

"""
    MeasureResult(value, determined)

The result of a qubit measurement returned by `measure!`.
"""
struct MeasureResult
    value::Bool
    determined::Bool

    function MeasureResult(value; determined)
        new(value, determined)
    end
end

function convert(::Type{T}, v::ChpSim.MeasureResult)::T where T<:Number
    return convert(T, v.value)
end

function Base.show(io::IO, self::MeasureResult)
    write(io, "$(typeof(self))($(Int(self.value)), "
              * "determined=$(self.determined))")
end

function Base.show(io::IO, ::MIME"text/plain", self::MeasureResult)
    write(io, "$(Int(self.value)) "
              * "($(self.determined ? "determined" : "random"))")
end

"""
    cnot!(state, control, target)

Perform a CNOT gate and update the state.
"""
function cnot!(state::ChpState, control::Int, target::Int)::Nothing
    @assert control != target
    state.r .⊻= ((@view state.x[:, control])
                 .& (@view state.z[:, target])
                 .& ((@view state.x[:, target])
                     .⊻ (@view state.z[:, control])
                     .⊻ true))
    state.x[:, target] .⊻= @view state.x[:, control]
    state.z[:, control] .⊻= @view state.z[:, target]
    nothing
end

"""
    hadamard!(state, qubit)

Perform a Hadamard gate and update the state.
"""
function hadamard!(state::ChpState, qubit::Int)::Nothing
    state.r .⊻= (@view state.x[:, qubit]) .& (@view state.z[:, qubit])
    state._tmp .= @view state.x[:, qubit]
    state.x[:, qubit] .= @view state.z[:, qubit]
    state.z[:, qubit] .= state._tmp
    nothing
end

"""
    phase!(state, qubit)

Perform a phase gate (S gate) and update the state.
"""
function phase!(state::ChpState, qubit::Int)::Nothing
    state.r .⊻= (@view state.x[:, qubit]) .& (@view state.z[:, qubit])
    state.z[:, qubit] .⊻= @view state.x[:, qubit]
    nothing
end

"""
    measure!(state, qubit; rng=GLOBAL_RNG, bias=0.5)

Perform a measurement in the Z basis and return the result.
If the qubit was in superposition, pick a random result and update the state.

# Arguments
- `qubit::Int`: The qubit index to measure.
- `rng::AbstractRNG`: The random number generator to use.
- `bias::Real`: The artificial probablility of measuring a 1 randomly.
"""
function measure!(state::ChpState, qubit::Int;
                  rng::AbstractRNG=Random.GLOBAL_RNG, bias::Real=0.5
                 )::MeasureResult
    for p in 1:state.n
        if state.x[state.n+p, qubit]
            return _measure_random(state, qubit, p, rng, bias)
        end
    end
    _measure_determined(state, qubit)
end

function _measure_determined(state::ChpState, qubit::Int)::MeasureResult
    state._tmp .= 0
    state._tmp_s .= 0
    tx = @view state._tmp[1:state.n]
    tz = @view state._tmp[state.n+1:2state.n]
    tr = @view state._tmp_s[]
    for i in 1:state.n
        if state.x[i, qubit]
            _row_mult(tx, tz, state._tmp_s,
                      (@view state.x[i + state.n, :]),
                      (@view state.z[i + state.n, :]),
                      (@view state.r[i + state.n]))
        end
    end
    MeasureResult(tr[], determined=true)
end

function _measure_random(state::ChpState, qubit::Int, p::Int, rng::AbstractRNG,
                         bias::Real)::MeasureResult
    @assert state.x[state.n + p, qubit]
    state.x[p, :] .= @view state.x[state.n + p, :]
    state.x[state.n + p, :] .= 0
    state.z[p, :] .= @view state.z[state.n + p, :]
    state.z[state.n + p, :] .= 0
    state.r[p] = state.r[state.n + p]
    state.r[state.n + p] = 0
    state.z[state.n + p, qubit] = 1
    state.r[state.n + p] = rand(rng) < bias

    px = @view state.x[p, :]
    pz = @view state.z[p, :]
    pr = @view state.r[p]
    for i in 1:2state.n
        if state.x[i, qubit] && i != p && i != state.n + p
            _row_mult((@view state.x[i, :]),
                      (@view state.z[i, :]),
                      (@view state.r[i]),
                      px, pz, pr)
        end
    end
    MeasureResult(state.r[state.n + p], determined=false)
end

function _row_mult(xrow1, zrow1, rrow1, xrow2, zrow2, rrow2)
    rrow1[] = _row_product_sign(xrow1, zrow1, rrow1, xrow2, zrow2, rrow2)
    xrow1 .⊻= xrow2
    zrow1 .⊻= zrow2
end

function _row_product_sign(xrow1, zrow1, rrow1, xrow2, zrow2, rrow2)
    pauli_phases = sum(
        _pauli_product_phase(xrow1[j], zrow1[j], xrow2[j], zrow2[j])
        for j in 1:length(xrow1)
    )
    @assert pauli_phases & 1 == 0
    p = Bool((pauli_phases >> 1) & 1)
    rrow1[] ⊻ rrow2[] ⊻ p
end

function _pauli_product_phase(x1::Bool, z1::Bool, x2::Bool, z2::Bool)::Int
    ixzy1 = x1 | z1 << 1
    ixzy2 = x2 | z2 << 1
    # If either pauli is I or both are the same
    (ixzy1 == 0 || ixzy2 == 0) && return 0
    # +1 if p1->p2 is increasing X->Y, Y->Z, or Z->X
    # -1 if decending
    # 0 if p1 == p2
    return mod(ixzy1 - ixzy2 + 1, 3) - 1
end

end


text:
#--------------------------------------------------------------------------------------------------
# Functions on [0, 1] tests
#--------------------------------------------------------------------------------------------------
using Test
using Hensel

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

@test Hensel.pIndex(0.5     , 2) == 1
@test Hensel.pIndex(0.4999  , 2) == 0
@test Hensel.pIndex(0.499999, 3) == 1
@test Hensel.pIndex(0.699999, 3) == 2

@test Hensel.hVector(0.5, 2, 3)   == [1, 0, 0]
@test Hensel.hVector(0.999, 2, 3) == [1, 1, 1]
@test Hensel.hVector(0.499, 2, 3) == [0, 1, 1]
@test Hensel.hVector(0.3, 2, 3)   == [0, 1, 0]

@test Hensel.iValue([1, 1, 0], 2) == 3
@test Hensel.iValue([1, 0, 1], 2) == 5
@test Hensel.iValue([1, 1, 1], 2) == 7

@test Hensel.hVector(0, 2, 3) == [0, 0, 0]
@test Hensel.hVector(1, 2, 3) == [1, 0, 0]
@test Hensel.hVector(3, 2, 3) == [1, 1, 0]
@test Hensel.hVector(5, 2, 3) == [1, 0, 1]
@test Hensel.hVector(7, 2, 3) == [1, 1, 1]

@test Hensel.hVector(0.0, 2, 3) == [0, 0, 0]
@test Hensel.hVector(0.9, 2, 3) == [1, 1, 1]
@test Hensel.hVector(1.0, 2, 3) == [1, 1, 1]
@test Hensel.hVector(0.0, 7, 3) == [0, 0, 0]
@test Hensel.hVector(0.005, 7, 3) == [0, 0, 1]
@test Hensel.hVector(0.995, 7, 3) == [6, 6, 5]
@test Hensel.hVector(1.0, 7, 3) == [6, 6, 6]

@test ((x) -> Hensel.hVector(x, 2, 5)).(0:9) == [
    [0, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [0, 1, 0, 0, 0],
    [1, 1, 0, 0, 0],
    [0, 0, 1, 0, 0],
    [1, 0, 1, 0, 0],
    [0, 1, 1, 0, 0],
    [1, 1, 1, 0, 0],
    [0, 0, 0, 1, 0],
    [1, 0, 0, 1, 0]
]
@test ((x) -> Hensel.hVector(x, 3, 5)).(0:9) == [
    [0, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [2, 0, 0, 0, 0],
    [0, 1, 0, 0, 0],
    [1, 1, 0, 0, 0],
    [2, 1, 0, 0, 0],
    [0, 2, 0, 0, 0],
    [1, 2, 0, 0, 0],
    [2, 2, 0, 0, 0],
    [0, 0, 1, 0, 0]
]
@test ((x) -> Hensel.hVector(x, 7, 5)).(0:9) == [
    [0, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [2, 0, 0, 0, 0],
    [3, 0, 0, 0, 0],
    [4, 0, 0, 0, 0],
    [5, 0, 0, 0, 0],
    [6, 0, 0, 0, 0],
    [0, 1, 0, 0, 0],
    [1, 1, 0, 0, 0],
    [2, 1, 0, 0, 0]
]

@test ((x) -> Hensel.hVector(x, 2, 5)).([n/10.0 for n in 0:10]) == [
    [0, 0, 0, 0, 0],
    [0, 0, 0, 1, 1],
    [0, 0, 1, 1, 0],
    [0, 1, 0, 0, 1],
    [0, 1, 1, 0, 0],
    [1, 0, 0, 0, 0],
    [1, 0, 0, 1, 1],
    [1, 0, 1, 1, 0],
    [1, 1, 0, 0, 1],
    [1, 1, 1, 0, 0],
    [1, 1, 1, 1, 1]
]
@test ((x) -> Hensel.hVector(x, 3, 5)).([n/10.0 for n in 0:10]) == [
    [0, 0, 0, 0, 0],
    [0, 0, 2, 2, 0],
    [0, 1, 2, 1, 0],
    [0, 2, 2, 0, 0],
    [1, 0, 1, 2, 1],
    [1, 1, 1, 1, 1],
    [1, 2, 1, 0, 1],
    [2, 0, 0, 2, 2],
    [2, 1, 0, 1, 2],
    [2, 2, 0, 0, 2],
    [2, 2, 2, 2, 2]
]
@test ((x) -> Hensel.hVector(x, 7, 5)).([n/10.0 for n in 0:10]) == [
    [0, 0, 0, 0, 0],
    [0, 4, 6, 2, 0],
    [1, 2, 5, 4, 1],
    [2, 0, 4, 6, 2],
    [2, 5, 4, 1, 2],
    [3, 3, 3, 3, 3],
    [4, 1, 2, 5, 4],
    [4, 6, 2, 0, 4],
    [5, 4, 1, 2, 5],
    [6, 2, 0, 4, 6],
    [6, 6, 6, 6, 6]
]

@test ((x) -> Hensel.rValue(Hensel.hVector(x, 2, 5), 2)).(0:9) == [
    0.015625,
    0.515625,
    0.265625,
    0.765625,
    0.140625,
    0.640625,
    0.390625,
    0.890625,
    0.078125,
    0.578125
]
@test ((x) -> Hensel.rValue(Hensel.hVector(x, 3, 5), 5)).(0:9) == [
    0.00016,
    0.20015999999999998,
    0.40015999999999996,
    0.04016,
    0.24016,
    0.44016,
    0.08016,
    0.28016,
    0.48016,
    0.00816,
]
@test ((x) -> Hensel.rValue(Hensel.hVector(x, 7, 5), 7)).(0:9) == [
    2.9749509133099306e-5,
    0.14288689236627597,
    0.2857440352234188,
    0.4286011780805617,
    0.5714583209377045,
    0.7143154637948473,
    0.8571726066519902,
    0.02043791277443922,
    0.16329505563158206,
    0.30615219848872494
]

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

function hTest(x, p, sz)
    @test Hensel.hVector(Hensel.rValue(Hensel.hVector(x, p, sz), p), p, sz) == Hensel.hVector(x, p, sz)
end

hTest(0.5, 2, 3)
hTest(0.5, 3, 4)
hTest(0.009, 3, 14)
hTest(0.333, 3, 14)
hTest(0.999, 3, 14)

dist(x, p, sz) = abs(Hensel.rValue(Hensel.hVector(x, p, sz), p) - x)

function dTest(x, p, sz)
    dist(x, p, sz) <= 1/p^(sz)
end

dTest(0.5, 2, 3)
dTest(0.5, 3, 4)
dTest(0.009, 3, 14)
dTest(0.333, 3, 14)
dTest(0.999, 3, 14)

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

max_stars_repo_name: ayarodionov/Hensel

max_issues_repo_name: ayarodionov/Hensel

max_forks_repo_name: ayarodionov/Hensel


text:
const config = Dict(:names => Symbol[])

function _var_expr(i, n)
    return RevisedAffineForm(0, [i == j for j in 1:n], 0)
end

macro rafvars(x...)
    return _rafvars(x...)
end

function _rafvars(x::Symbol)
    rhs = _var_expr(1, 1)
    config[:names] = [x]
    return :($(esc(x)) = $rhs)
end

function _rafvars(x::Expr)
    s = x.args[1]
    start = x.args[2].args[2]
    stop = x.args[2].args[3]
    n = stop - start + 1
    rhs = [_var_expr(i, n) for i in start:stop]
    varnames = [Symbol(s, _lower(i)) for i in start:stop]
    config[:names] = varnames
    return :($(esc(s)) = $rhs)
end

function _rafvars(x...)
    n = length(x)
    ex = quote end
    for (i, s) in enumerate(x)
        rhs = _var_expr(i, n)
        push!(ex.args, :($(esc(s)) = $rhs))
    end
    config[:names] = collect(x)
    return ex
end

macro affinize(block)
    if block.head == :tuple
        nvars = length(block.args)
        names = Vector{Symbol}(undef, nvars)
        ex = quote end
        for (i, var) in enumerate(block.args)
            name, rhs = _affinize(var, nvars, i)
            push!(ex.args, :($(esc(name)) = $rhs))
            names[i] = name
        end
        config[:names] = names
        return ex
    elseif block.head == :(=)
        name, rhs = _affinize(block, 1, 1)
        config[:names] = [name]
        return :($(esc(name)) = $rhs)
    else
        throw(ArgumentError("Invalid Statement"))
    end
end

function _affinize(ex::Expr, n::Int, i::Int)
    if ex.head == :(=)
        name = ex.args[1]
        p = eval(ex.args[2])
        c, r = midpoint_radius(p)
        coeffs = zeros(typeof(r), n)
        coeffs[i] = r
        return name, RevisedAffineForm(c, coeffs, 0)
    else
        throw(ArgumentError("Invalid Statement $ex"))
    end
end


text:
abstract type NamedRange{N,T,R} <: StaticRangeWrapper{T,R} end

Base.names(r::NamedRange{N,T,R}) where {N,T,R} = N
Base.names(r::Type{<:NamedRange{N,T,R}}) where {N,T,R} = N

NamedRange(n::NTuple{N,Symbol}, r::SRange) where N = NamedSRange{n}(r)
NamedRange(n::NTuple{N,Symbol}, r::MRange) where N = NamedMRange{n}(r)

struct NamedSRange{N,T,R} <: NamedRange{N,T,R}

    function NamedSRange{N,T,R}() where {N,T,R}
        isstatic(R) || error("NamedSRange only supports static ranges.")
        new{N,T,R}()
    end
end

NamedSRange{N}(r::SRange{T}) where {N,T} = NamedSRange{N,T,typeof(r)}()

"""
    NamedMRange

A mutable range that may be indexed by
"""
mutable struct NamedMRange{N,T,R} <: NamedRange{N,T,R}
    names::N
    parent::R

    NamedMRange{N}(r::AbstractRange{T}) where {N,T} = new{N,T,typeof(r)}(names, r)
end

Base.parent(r::NamedMRange) = getfield(r, :parent)

@inline getindex(r::NamedRange, i::Integer) = getindex(parent(r), i)

@inline getindex(r::NamedRange, s::Symbol) = getindex(r, SSymbol(s))

@inline @generated function getindex(r::NamedRange, ::SSymbol{S}) where S
    idx = findfirst(y->y==S::Symbol, names(r))
    :(parent(r)[$idx])
end

#@inline Base.setindex!(x::LArray,y,i...) = getfield(x,:__x)[i...] = y

@inline function getindex(r::NamedRange, inds::AbstractVector{<:Number})
    _namedrange(names(r)[inds], getindex(parent(r), inds))
end

_namedrange(n::NTuple{N,Symbol}, r::Vector) where N = NamedTuple{n}(r...)
_namedrange(n::NTuple{N,Symbol}, r::SRange) where N = NamedSRange{n}(r)
_namedrange(n::NTuple{N,Symbol}, r::AbstractRange) where N = NamedMRange{n}(r)

function getindex(x::NamedRange, s::AbstractArray{Symbol,1})
    NamedTuple{(s...,)}([getindex(x, si) for si in s]...)
end


text:
module MaxHelpingHandRainbowSpinnerColorTrigger

using ..Ahorn, Maple

@mapdef Trigger "MaxHelpingHand/RainbowSpinnerColorTrigger" RainbowSpinnerColorTrigger(x::Integer, y::Integer, width::Integer=Maple.defaultTriggerWidth, height::Integer=Maple.defaultTriggerHeight,
    colors::String="89E5AE,88E0E0,87A9DD,9887DB,D088E2", gradientSize::Number=280.0, loopColors::Bool=false, centerX::Number=0.0, centerY::Number=0.0, gradientSpeed::Number=50.0, persistent::Bool=false)

const placements = Ahorn.PlacementDict(
    "Rainbow Spinner Color (max480's Helping Hand)" => Ahorn.EntityPlacement(
        RainbowSpinnerColorTrigger,
        "rectangle"
    )
)

end

max_stars_repo_name: Kitty-Cats/MaxHelpingHand

max_issues_repo_name: Kitty-Cats/MaxHelpingHand

max_forks_repo_name: Kitty-Cats/MaxHelpingHand


text:
using CUDArt,CUBLAS,CUSPARSE,Knet,Base.Test
using Knet: gemm!, gpusync

# include("isapprox.jl")

csc2csr{T}(x::SparseMatrixCSC{T})=CudaSparseMatrixCSR{T}(CudaArray(convert(Vector{Cint},x.colptr)), CudaArray(convert(Vector{Cint},x.rowval)), CudaArray(x.nzval), (x.n,x.m), convert(Cint,length(x.nzval)), device())

function LinAlg.A_mul_B!{T}(Y::StridedMatrix{T}, X::StridedMatrix{T}, A::SparseMatrixCSC{T})
    mX, nX = size(X)
    nX == A.m || throw(DimensionMismatch())
    fill!(Y,0)
    rowval = A.rowval
    nzval = A.nzval
    @inbounds for multivec_row=1:mX, col = 1:A.n, k=A.colptr[col]:(A.colptr[col+1]-1)
        Y[multivec_row, col] += X[multivec_row, rowval[k]] * nzval[k]
    end
    Y
end

# Experiment 1:
# nh = 1000
# nd = 100000
# nb = 100

# Experiment 2:
nh = 1500
nd = 10000
nb = 20

info("1000x CPU forw: 0.2ms")
w = x = y = nothing; gc()
w = rand(Float32,  nh, nd)      # N
x = sponehot(Float32, nd, nb)      # N
rand!(x.rowval, 1:nd)
y = Array(Float32, nh, nb)      # N
@time for i=1:1000              # 1:0.5ms, 2:0.2ms
    A_mul_B!(y, w, x)
end
@test @show to_host(y)==to_host(w)[:,x.rowval]

# info("1000x GPU forw using KUsparse A_mul_B!(y,w,x) 2.5ms")
# w1 = x1 = y1 = nothing; gc()
# w1 = CudaArray(w)               # N
# x1 = convert(KUsparse{CudaArray,Float32}, x) # N
# y1 = fill!(CudaArray(y),0)                   # N
# @time for i=1:1000
#     A_mul_B!(y1, w1, x1)        # 1:2.5ms
#     gpusync()
# end
# @test @show to_host(y1)==y
# CUBLAS.geam!('T','T',1f0,w1,0f0,w1,w2)

# if false

info("10000x GPU forw using csrmm!(x',w',y') 0.056ms")
w2 = x2 = y2 = nothing; gc()
x2 = CudaSparseMatrixCSR(x')    # T
w2 = CudaArray(w')              # T
y2 = fill!(CudaArray(y'),0)     # T
@time for i=1:10000             # 1:0.08ms 2:0.056ms
    CUSPARSE.csrmm!('N',1f0,x2,w2,0f0,y2,'O')
    gpusync()
end
@test @show to_host(y2)==y'

info("10000x GPU forw using csrmm2!(x',w,y') 0.036ms")
w3 = x3 = y3 = nothing; gc()
x3 = x2                         # T
w3 = CudaArray(w)               # N
y3 = similar(y2)                # T
@time for i=1:10000             # 1:0.04ms, 2:0.036ms
    CUSPARSE.csrmm2!('N','T',1f0,x3,w3,0f0,y3,'O')
    gpusync()
end
@test @show to_host(y3)==y'

info("10000x GPU forw using csrmm2!(x',w,y')+transpose 0.046ms")
w4 = x4 = y4 = nothing; gc()
x4 = csc2csr(x)                 # T: direct copy possible
w4 = w3                         # N
y4 = fill!(CudaArray(y),0)      # N
z4 = similar(y2)                # T
@time for i=1:10000             # 1:0.05ms 2:0.046ms
    CUSPARSE.csrmm2!('N','T',1f0,x4,w4,0f0,z4,'O')
    CUBLAS.geam!('T','T',1f0,z4,0f0,z4,y4) # need temp variable :(
    gpusync()
end
@test @show to_host(y4)==y

info("10000x GPU forw using A_mul_B!(y,w,x) 0.074ms")
w4a = x4a = y4a = nothing; gc()
x4a = CudaSparseMatrixCSC(x)     # N
w4a = w3                         # N
y4a = CudaArray(similar(y)) # N
@time for i=1:10000             #
    A_mul_B!(y4a, w4a, x4a)
    gpusync()
end
@test @show convert(Array,y4a)==y

# end # if false

#### going back: dw = dy * x',

# we want sparse dw, so we need to use csrgemm.
# new design: we have x in csc, which is equiv to x' csr
# we'll have dw and dy also in csc and run A_mul_Bt(dw,dy,x)

# we already have x' sparse and csr.
# we convert dy from dense to csr.  dw needs lots of allocation which hopefully can be avoided.

# The true cost should include: any transposes, sparsifications,
# densifications, allocations, and the final update dw+=iw, w+=dw
# we always have x sparse, dy,w dense. we have a choice for iw and dw.
# if seqlen=t=20, we have 20 (iw=dy*x'; dw+=iw) followed by one w+=dw.
# D=10K, h=1K, b=10, t=10
# x:(D,b), nnz(x)=b, D>>b, avoid densification if possible.
# w,dw,iw:(h,D), nnz(iw):h*b, nnz(dw):h*b*t
# y,dy:(h,b)

#    ........x............................................
#    ..................x..................................
#    ..x..................................................
# ...
# ...
# ...
# ...
# ...

info("1x CPU back: dense x sparse (single iter) 41ms")
dw = dy = nothing; gc()
dy = rand!(similar(y))
@time dw = dy * x'              # 1:280ms 2:41.6ms

# if false

info("1000x CPU back: all sparse 0.78ms")
dw0 = dy0 = nothing; gc()
dy0 = sparse(dy)
@time for i=1:1000              # 1:3.26ms 2:0.78ms
    dw0 = dy0 * x'
end
@test @show full(dw0) == dw

# dw = dy * x'
info("1000x GPU back (all sparse) using dw=gemm(dy,x') 1.08ms")
dw5 = x5 = dy5 = nothing; gc()
x5 = csc2csr(x)                 # T
dy5 = sparse(CudaArray(dy)) # N
dw5 = nothing
@time for i=1:1000              # 1:8.55ms 2:1.08ms
    dw5 = CUSPARSE.gemm('N','N',dy5,x5,'O','O','O')
    gpusync()
end
@test @show full(to_host(dw5)) == dw

# end # if false

# dw = dy * x'
info("1000x GPU back (all sparse) using gemm!(dy,x',dw) 0.72ms")
dw6 = x6 = dy6 = nothing; gc()
x6 = csc2csr(x)                 # T
dy6 = sparse(CudaArray(dy))     # N
dw6 = CudaSparseMatrixCSR(spzeros(Float32,size(w)...)) # N
@time for i=1:1000              # 1:7.26ms 2:0.72ms
    gemm!('N','N',dy6,x6,dw6)
    gpusync()
end
@test @show full(to_host(dw6)) == dw

# dw = dy * x'
info("1000x GPU back (all sparse) using A_mul_Bt (including sparse(dy)) 1.28ms")
dw6a = x6a = dy6a = nothing; gc()
x6a = CudaSparseMatrixCSC(x)
dy6a = CudaArray(dy)     # N
dw6a = CudaSparseMatrixCSR(spzeros(Float32,size(w)...)) # N
@time for i=1:1000              # 2:1.28ms
    A_mul_Bt!(dw6a,dy6a,x6a)
    gpusync()
end
@test @show full(to_host(dw6a)) == dw

# dw = dy * x'
info("1000x GPU back (all sparse) using A_mul_Bt (including sparse(dy) and axpy) 2.21ms")
dw6b = x6b = dy6b = nothing; gc()
x6b = CudaSparseMatrixCSC(x)
dy6b = CudaArray(dy)     # N
dw6b = CudaSparseMatrixCSR(spzeros(Float32,size(w)...)) # N
iw6b = CudaSparseMatrixCSR(spzeros(Float32,size(w)...)) # N
@time for i=1:1000              # 2:2.21ms
    A_mul_Bt!(iw6b,dy6b,x6b)
    Base.axpy!(1,iw6b,dw6b)
    gpusync()
end
@test @show full(to_host(iw6b)) == dw
@show full(to_host(dw6b)) == 1000*full(to_host(iw6b))
@show isapprox(full(to_host(dw6b)), 1000*full(to_host(iw6b)))

# if false

# dw = dy * x'
info("1000x GPU back (all sparse) using dw'=gemm(x,dy') 2.39ms")
dw7a = x7a = dy7a = nothing; gc()
x7a = CudaSparseMatrixCSR(x)                 # N
dy7a = sparse(CudaArray(dy')) # T
dw7a = nothing                # T
@time for i=1:1000            # 1:3.39ms 2:2.39ms
    dw7a = CUSPARSE.gemm('N','N',x7a,dy7a,'O','O','O')
    gpusync()
end
@test @show full(to_host(dw7a)) == dw'

info("1000x GPU back (all sparse) using gemm!(x,dy',dw') 2.09ms")
dw7 = x7 = dy7 = nothing; gc()
x7 = CudaSparseMatrixCSR(x)                 # N
dy7 = sparse(CudaArray(dy'))                # T
dw7 = CudaSparseMatrixCSR(spzeros(Float32,size(w')...)) # T
@time for i=1:1000              # 1:2.13ms 2:2.09ms
    gemm!('N','N',x7,dy7,dw7)
    gpusync()
end
@test @show full(to_host(dw7)) == dw'

# dw'= x * dy'
info("1000x GPU back (sparse x dense)->dense using csrmm!(x,dy',dw') 3.40ms")
dw8 = x8 = dy8 = nothing; gc()
x8 = CudaSparseMatrixCSR(x)                 # N
dy8 = CudaArray(dy')                        # T
dw8 = fill!(CudaArray(w'),0)                # T
@time for i=1:1000              # 1:22.3ms 2:3.40ms
    CUSPARSE.csrmm!('N',1f0,x8,dy8,0f0,dw8,'O')
    gpusync()
end
@test @show to_host(dw8) == dw'

# dw = dy * x'
info("1000x GPU back (all dense) using gemm!(dy,x',dw) 1.59ms")
dw9 = x9 = dy9 = nothing; gc()
dy9 = CudaArray(dy)             # N
x9 = CudaArray(full(x'))        # T
dw9 = fill!(CudaArray(w),0)     # N
@time for i=1:1000              # 1:16.6s 2:1.59ms
    A_mul_B!(dw9,dy9,x9)
    gpusync()
end
@test @show to_host(dw9) == dw

# end # if false