"""
Stage coordinates are held in a `Dict{Symbol,Float64}`.  The nominal units
are mm and radians.  The default labels are `:X`, `:Y`, `:Z`, `:R`, `:T`, and `:B`.

The `StageMapping` abstract class defines the relationship between stage coordinates
and image (or hyperimage) pixels through the functions `image2stage(...)` and
`stage2image(...)`.

Since every instrument is different, you will likely have to implement your own
`StageMapping` struct with its own functions.  `DefaultStageMapping` is provided
as an example of a simple mapping.
"""

abstract type StageMapping end

"""
The `DefaultStageMapping` assumes a matching Cartesian coordinate system for both the stage and image.  It only assumes
`:X`, `:Y` stage motion.
"""
struct DefaultStageMapping <: StageMapping end

"""
    image2stage(::Type{DefaultStageMapping}, stage_coord::Dict{Symbol,Any}, img_coord::Dict{Symbol,Any}, theta::Float64)

Given the stage coordinate at the center of the image `stage_coord`, the pixel coordinate of interest (in the same units as
`stage_coord`) and the image rotation `theta`, compute the stage coordinate that would bring the pixel into the center of the
image.
"""
function image2stage(::Type{DefaultStageMapping}, stage_coord::Dict, img_coord::Dict, theta::Float64)
    rot = [ cos(theta) -sin(theta) ; sin(theta) cos(theta) ]
    xy = [ stage_coord[:X], stage_coord[:Y] ] + rot * [ img_coord[:X], img_coord[:Y] ]
    return Dict(:X=>xy[1], :Y=>xy[2])
end

"""
    stage2image(::Type{DefaultStageMapping}, stage_coord::Dict{Symbol,Any}, centered_coord::Dict{Symbol,Any}, theta::Float64)

`stage2image(...)` is the inverse function of `image2stage(...)`.  Given the stage coordinate of the center of the image
`stage_coord` and the stage coordinate that would center the pixel of interest `centered_coord`, compute the pixel coordinate
corresponding to the `centered_coordinate` when the stage is at `stage_coord`.
"""
function stage2image(::Type{DefaultStageMapping}, stage_coord::Dict, centered_coord::Dict, theta::Float64)
    roti = [ cos(theta) sin(theta) ; -sin(theta) cos(theta) ]
    xy = roti * ([ centered_coord[:X], centered_coord[:Y] ] - [ stage_coord[:X], stage_coord[:Y] ])
    return Dict{Symbol,Float64}(:X=>xy[1], :Y=>xy[2])
end

"""
    compute_tilt(v1::Vector, v2::Vector, v3::Vector)

Compute the tilt and angle of tilt of the sample where `v1`, `v2` and `v3` are three points
forming a triangle with each focused on the surface of the sample.  Assumes a right-handed
coordinate system which may/may not match your stage.
"""
function compute_tilt(v1::Vector, v2::Vector, v3::Vector)
    vz = cross(v2-v1, v3-v1)
    n = vz / norm(vz)
    return rad2deg.( (acos(n[3]), atan(n[2],n[1])) )
end


@everywhere include("load_network.jl")

@everywhere using Graphs

@everywhere function create_graph(start_node, end_node)
    @assert Base.length(start_node)==Base.length(end_node)

    no_node = max(maximum(start_node), maximum(end_node))
    no_arc = Base.length(start_node)

    graph = simple_inclist(no_node)
    for i=1:no_arc
        add_edge!(graph, start_node[i], end_node[i])
    end
    return graph
end

@everywhere function get_vector(state, origin, destination, link_dic)
    current = destination
    parent = -1
    x = zeros(Int, maximum(link_dic))

    while parent != origin
        parent = state.parents[current]

        link_idx = link_dic[parent,current]

        if link_idx != 0
            x[link_idx] = 1
        end

        current = parent
    end

    return x
end

ta_data = load_ta_network("Sioux Falls")

# unpacking data from ta_data
network_name = ta_data.network_name

number_of_zones = ta_data.number_of_zones
number_of_nodes = ta_data.number_of_nodes
first_thru_node = ta_data.first_thru_node
number_of_links = ta_data.number_of_links

start_node = ta_data.start_node
end_node = ta_data.end_node
capacity = ta_data.capacity
link_length = ta_data.link_length

free_flow_time = ta_data.free_flow_time
B = ta_data.B
power = ta_data.power
speed_limit = ta_data.speed_limit
toll = ta_data.toll
link_type = ta_data.link_type
number_of_zones = ta_data.number_of_zones
total_od_flow = ta_data.total_od_flow
travel_demand = ta_data.travel_demand
od_pairs = ta_data.od_pairs

toll_factor = ta_data.toll_factor
distance_factor = ta_data.distance_factor

best_objective = ta_data.best_objective

# preparing a graph
graph = create_graph(start_node, end_node)
link_dic = sparse(start_node, end_node, 1:number_of_links)

@everywhere function BPR(x)
    # travel_time = free_flow_time .* ( 1.0 + B .* (x./capacity).^power )
    # generalized_cost = travel_time + toll_factor *toll + distance_factor * link_length
    # return generalized_cost

    bpr = similar(x)
    for i=1:length(bpr)
        bpr[i] = free_flow_time[i] * ( 1.0 + B[i] * (x[i]/capacity[i])^power[i] )
        bpr[i] += toll_factor *toll[i] + distance_factor * link_length[i]
    end
    return bpr
end

@everywhere function objective(x)
    # value = free_flow_time .* ( x + B.* ( x.^(power+1)) ./ (capacity.^power) ./ (power+1))
    # return sum(value)

    sum = 0
    for i=1:length(x)
        sum += free_flow_time[i] * ( x[i] + B[i]* ( x[i]^(power[i]+1)) / (capacity[i]^power[i]) / (power[i]+1))
        sum += toll_factor *toll[i] + distance_factor * link_length[i]
    end
    return sum
end

@everywhere function all_or_nothing_single(travel_time)
    state = []
    path = []
    x = zeros(size(start_node))

    for r=1:size(travel_demand)[1]
        # for each origin node r, find shortest paths to all destination nodes
        state = dijkstra_shortest_paths(graph, travel_time, r)

        for s=1:size(travel_demand)[2]
            # for each destination node s, find the shortest-path vector
            # load travel demand
            x = x + travel_demand[r,s] * get_vector(state, r, s, link_dic)
        end
    end

    return x
end

# # parallel computing version
@everywhere function all_or_nothing_parallel(travel_time)
    state = []
    path = []
    vv = zeros(size(start_node))
    x = zeros(size(start_node))

    x = x + @parallel (+) for r=1:size(travel_demand)[1]
        # for each origin node r, find shortest paths to all destination nodes
        # if there is any travel demand starting from node r.
        vv = zeros(size(start_node))

        if sum(travel_demand, 2)[r] > 0.0
            state = dijkstra_shortest_paths(graph, travel_time, r)

            for s=1:size(travel_demand)[2]
                # for each destination node s, find the shortest-path vector
                # v = get_vector(state, r, s, start_node, end_node)

                if travel_demand[r,s] > 0.0
                    # load travel demand
                    vv = vv + travel_demand[r,s] * get_vector(state, r, s, link_dic)
                end
            end

        end

        vv
    end

    return x
end

@everywhere function all_or_nothing(travel_time)
    if nprocs() > 1 # if multiple CPU processes are available
        all_or_nothing_parallel(travel_time)
    else
        all_or_nothing_single(travel_time)
        # when nprocs()==1, using @parallel just adds unnecessary setup time. I guess.
    end
end

# Finding a starting feasible solution
travel_time = BPR(zeros(number_of_links))
xl = all_or_nothing(travel_time)

max_iter_no = 1e10
l = 1
average_excess_cost = 1
tol = 1e-7

while l < max_iter_no
    # Finding yl
    travel_time = BPR(xl)
    yl = all_or_nothing(travel_time)
    xl = xl + (yl - xl)/ (l + 1)
    l = l + 1

    # Average Excess Cost
    average_excess_cost = ( dot(xl, travel_time) - dot(yl, travel_time) ) / dot(xl, travel_time)

    if average_excess_cost < tol
        break
    end
end

print(xl)


### GumbelSoftmax

abstract type AbstractGumbelSoftmax{T} <: ContinuousMultivariateDistribution end

const τ0 = 0.2  # default value for τ

"""
    GumbelSoftmax(p, τ)

The Gumbel-Softmax distribution with mean `p` and temperature `τ`.

Ref: https://arxiv.org/abs/1611.01144
"""
struct GumbelSoftmax{T<:AbstractVecOrMat{<:AbstractFloat}} <: AbstractGumbelSoftmax{T}
    p::T
    τ
    function GumbelSoftmax(p::T, τ::F) where {F<:AbstractFloat, T<:AbstractVecOrMat{F}}
        return new{T}(p, τ)
    end
end

GumbelSoftmax(p; τ=eltype(p)(τ0)) = GumbelSoftmax(p, τ)

mean(gs::AbstractGumbelSoftmax) = gs.p

function u2gumbel(u)
    ϵ = eps(u)
    return -log.(-log.(u .+ ϵ) .+ ϵ)
end

function g2softmax(g, p, τ)
    logit = g .+ log.(p .+ eps(g))
    return NNlib.softmax(logit ./ τ; dims=1)
end

function _rand_gs(rng::AbstractRNG, p, τ, n::Int=1)
    u = randsimilar(rng, p, n)
    g = u2gumbel(u)
    return g2softmax(g, p, τ)
end

rand(
    rng::AbstractRNG,
    gs::AbstractGumbelSoftmax;
    τ=gs.τ
) = _rand_gs(rng, gs.p, τ)
rand(gs::AbstractGumbelSoftmax; τ=gs.τ) = rand(GLOBAL_RNG, gs; τ=gs.τ)

rand(
    rng::AbstractRNG,
    gs::AbstractGumbelSoftmax{<:AbstractVector},
    n::Int;
    τ=gs.τ
) = _rand_gs(rng, gs.p, τ, n)
rand(gs::AbstractGumbelSoftmax{<:AbstractVector}, n::Int; τ=gs.τ) = rand(GLOBAL_RNG, gs, n; τ=gs.τ)

function _logpdf(gs::AbstractGumbelSoftmax, x; τ=gs.τ)
    k = size(x, 1)
    ϵ = eps(gs.p)
    term1 = StatsFuns.lgamma(k) + (k - 1) * log(τ)
    xpowτ = x .^ τ
    term2 = k * log.(sum(gs.p ./ xpowτ; dims=1) .+ ϵ)
    xpowτp1 = xpowτ .* x
    term3 = sum(log.(gs.p ./ xpowτp1 .+ ϵ); dims=1)
    return term1 .- term2 + term3
end

logpdf(gs::AbstractGumbelSoftmax{<:AbstractArray{T1,2}}, x::AbstractArray{T2,2}; τ=gs.τ) where {T1,T2,N} = _logpdf(gs, x; τ=τ)

"""
    GumbelSoftmax2D(p1, τ)

2-dimensional GumbelSoftmax where `p1` is the probability of the first dimension is 1.
"""
function GumbelSoftmax2D(p1::AbstractVector; τ=eltype(p1)(τ0))
    p1transpose = transpose(p1)
    p = [p1transpose; 1 .- p1transpose]
    return GumbelSoftmax(p, τ)
end
function GumbelSoftmax2D(p1::T; τ=T(τ0)) where {T<:AbstractFloat}
    return GumbelSoftmax([p1, 1 - p1], τ)
end

### GumbelBernoulli

abstract type AbstractGumbelBernoulli{T} <: ContinuousMultivariateDistribution end

"""
The Gumbel-Bernoulli distribution with mean `p` and temperature `τ`.

Ref: https://arxiv.org/abs/1611.01144
"""
struct GumbelBernoulli{T<:AbstractArray{<:AbstractFloat}} <: AbstractGumbelBernoulli{T}
    p::T
    τ
    function GumbelBernoulli(p::T, τ::F) where {F<:AbstractFloat, T<:AbstractArray{F}}
        return new{T}(p, τ)
    end
end

GumbelBernoulli(p; τ=eltype(p)(τ0)) = GumbelBernoulli(p, τ)

"""
    logrand(rng, gb)

Generate sample in log space from a Gumbel-Bernoulli distribution.

Ref: https://arxiv.org/abs/1611.00712
"""
function logrand(rng::AbstractRNG, gb::GumbelBernoulli; τ=gb.τ)
    # TODO: re-implement this `rand` using the same procedure for `GumbelBernoulliLogit`
    u0 = randsimilar(rng, gb.p); g0 = u2gumbel(u0)
    u1 = randsimilar(rng, gb.p); g1 = u2gumbel(u1)

    ϵ = eps(gb.p)

    logit0 = (g0 + log.(1 + ϵ .- gb.p)) / τ
    logit1 = (g1 + log.(gb.p .+ ϵ)) / τ

    logit_max = max.(logit0, logit1)
    logit0_minus_max = logit0 - logit_max
    logit1_minus_max = logit1 - logit_max

    logx = logit1_minus_max - log.(exp.(logit0_minus_max) + exp.(logit1_minus_max))
    return logx
end
logrand(gb::GumbelBernoulli; τ=gb.τ) = logrand(GLOBAL_RNG, gb; τ=τ)

rand(rng::AbstractRNG, gb::GumbelBernoulli; τ=gb.τ) = exp.(logrand(rng, gb; τ=τ))
rand(gb::GumbelBernoulli; τ=gb.τ) = rand(GLOBAL_RNG, gb; τ=gb.τ)

function _logpdf(gb::GumbelBernoulli, x; τ=gb.τ)
    ϵ = eps(gb.p)
    α = gb.p ./ (1 + ϵ .- gb.p)
    xstabe = x .+ ϵ
    omxstabe = 1 .- x .+ ϵ
    return log(τ) .+ log.(α) .+ (1 - τ) * (log.(xstabe) + log.(omxstabe)) - 2 * (log.(α .* xstabe.^(-τ) + omxstabe.^(-τ) .+ ϵ))
end

logpdf(gb::GumbelBernoulli, x; τ=gb.τ) = _logpdf(gb, x; τ=τ)
logpdf(gb::GumbelBernoulli, x::AbstractMatrix{<:AbstractFloat}; τ=gb.τ) = _logpdf(gb, x; τ=τ)

"""
    logpdfCoV(bgb, x)

Compute the log density of `GumbelBernoulli(x; p)` by using
that of `GumbelBernoulli(logitx; logitp)` by the change of variable (CoV) rule.
"""
function logpdfCoV(gb::GumbelBernoulli, x; τ=gb.τ)
    ϵ = eps(gb.p)
    om2ϵ = 1 - 2 * ϵ
    logitp = StatsFuns.logit.(gb.p * om2ϵ .+ ϵ)
    logitx = StatsFuns.logit.(x * om2ϵ .+ ϵ)
    lp = logpdflogit(GumbelBernoulliLogit(logitp; τ=τ), logitx)
    return lp - log.(x .* (1 .- x) .+ ϵ)
end

mean(gb::GumbelBernoulli) = gb.p

"""
    GumbelBernoulliLogit(p, τ)

The Gumbel-Softmax distribution with mean `logitp` in logit space and temperature `τ`.

Ref: https://arxiv.org/abs/1611.01144
"""
struct GumbelBernoulliLogit{T} <: AbstractGumbelBernoulli{T}
    logitp::T
    τ
    function GumbelBernoulliLogit(logitp::T, τ::F) where {F<:AbstractFloat, T<:AbstractArray{F}}
        return new{T}(logitp, τ)
    end
end

GumbelBernoulliLogit(p; τ=eltype(p)(τ0)) = GumbelBernoulliLogit(p, τ)

"""
    logitrand(rng, gbl)

Generate sample in logit space from a Gumbel-Bernoulli distribution
with parameters in logit space.

Ref: https://arxiv.org/abs/1611.00712
"""
function logitrand(rng::AbstractRNG, gbl::GumbelBernoulliLogit; τ=gbl.τ)
    ϵ = eps(gbl.logitp)
    u = randsimilar(rng, gbl.logitp)
    logit = log.(u .+ ϵ) - log.(1 + ϵ .- u)
    logα = gbl.logitp
    logitx = (logα + logit) ./ τ
    return logitx
end
logitrand(gbl::GumbelBernoulliLogit; τ=gbl.τ) = logitrand(GLOBAL_RNG, gbl; τ=τ)

rand(
    rng::AbstractRNG,
    gbl::GumbelBernoulliLogit;
    τ=gbl.τ
) = StatsFuns.logistic.(logitrand(rng, gbl; τ=τ))

"""
    logpdflogit(gbl, logitx)

Compute the log density of `GumbelBernoulli(logitx; logitp)`.

Ref: https://arxiv.org/abs/1611.00712
"""
function logpdflogit(gbl::GumbelBernoulliLogit, logitx; τ=gbl.τ)
    exp_term = gbl.logitp .- τ * logitx
    lp = exp_term - 2 * StatsFuns.softplus.(exp_term) .+ log(τ)
    return lp
end

mean(gbl::GumbelBernoulliLogit) = StatsFuns.logistic.(gbl.logitp)


using NNlib
import NNlib: softmax, ∇softmax, logsoftmax, ∇logsoftmax, conv, depthwiseconv, ∇conv_data, ∇depthwiseconv_data, maxpool, meanpool, σ, relu, batched_mul, batched_adjoint

drelu(x, Δ) = ifelse(x > 0, Δ, zero(x))

@adjoint function Base.Broadcast.broadcasted(::typeof(relu), x::Numeric)
  relu.(x), Δ -> (nothing, drelu.(x, Δ))
end

function dselu(x)
  λ = oftype(x/1, 1.0507009873554804934193349852946)
  α = oftype(x/1, 1.6732632423543772848170429916717)
  λ * ifelse(x > 0, 1, α * exp(x))
end

@adjoint function Base.Broadcast.broadcasted(::typeof(selu), x::Numeric)
  selu.(x), Δ -> (nothing, dselu.(x) .* Δ)
end

@adjoint function σ(x::Real)
    y = σ(x)
    return y, Δ -> (Δ * y * (1 - y),)
end

@adjoint softmax(xs; dims=1) = softmax(xs, dims=dims), Δ -> (∇softmax(Δ, xs, dims=dims),)

@adjoint logsoftmax(xs; dims=1) = logsoftmax(xs, dims=dims), Δ -> (∇logsoftmax(Δ, xs, dims=dims),)

@adjoint NNlib.DenseConvDims(args...; kwargs...) = NNlib.DenseConvDims(args...; kwargs...), _ -> nothing
@adjoint NNlib.DepthwiseConvDims(args...; kwargs...) = NNlib.DepthwiseConvDims(args...; kwargs...), _ -> nothing
@adjoint NNlib.PoolDims(args...; kwargs...) = NNlib.PoolDims(args...; kwargs...), _ -> nothing

colmajor(x) = colmajor(MemoryLayout(typeof(x)), x)
colmajor(_, x) = convert(Array, x)
colmajor(::AbstractColumnMajor, x) = x

@adjoint conv(x, w, cdims; kw...) =
  conv(x, w, cdims; kw...),
    Δ -> begin
       Δ = colmajor(Δ)
       return (
           NNlib.∇conv_data(Δ, w, cdims; kw...),
           NNlib.∇conv_filter(x, Δ, cdims; kw...),
           nothing,
       )
   end

@adjoint ∇conv_data(x, w, cdims; kw...) =
  ∇conv_data(x, w, cdims; kw...),
    Δ -> begin
       Δ = colmajor(Δ)
       return (
           NNlib.conv(Δ, w, cdims; kw...),
           NNlib.∇conv_filter(Δ, x, cdims; kw...),
           nothing,
       )
   end

@adjoint depthwiseconv(x, w, cdims; kw...) =
  depthwiseconv(x, w, cdims; kw...),
    Δ -> begin
       Δ = colmajor(Δ)
       return (
           NNlib.∇depthwiseconv_data(Δ, w, cdims; kw...),
           NNlib.∇depthwiseconv_filter(x, Δ, cdims; kw...),
           nothing,
       )
   end

@adjoint ∇depthwiseconv_data(x, w, cdims; kw...) =
  ∇depthwiseconv_data(x, w, cdims; kw...),
    Δ -> begin
       Δ = colmajor(Δ)
       return (
           NNlib.depthwiseconv(Δ, w, cdims; kw...),
           NNlib.∇depthwiseconv_filter(Δ, x, cdims; kw...),
           nothing,
       )
   end

@adjoint function maxpool(x, pdims::NNlib.PoolDims; kw...)
  y = maxpool(x, pdims; kw...)
  y, Δ -> (NNlib.∇maxpool(Δ, y, x, pdims; kw...), nothing)
end

@adjoint function meanpool(x, pdims::NNlib.PoolDims; kw...)
  y = meanpool(x, pdims; kw...)
  y, Δ -> (NNlib.∇meanpool(Δ, y, x, pdims; kw...), nothing)
end

@adjoint function batched_mul(A, B)
    C = batched_mul(A, B)
    C, Δ -> (batched_mul(Δ, batched_adjoint(B)), batched_mul(batched_adjoint(A), Δ))
end


export z, t, abbott

include("end_effects.jl")

"""
    Rz(R, λc=250)

Determine roughness of profile-component `component`. Filterwavelength in accordance with [DIN EN ISO 3274:1998](DIN_EN_ISO_3274_1998.pdf).
"""
function z(profile::SurfaceProfile, component::Symbol, eliminate_end_effect=zeropad)
    c = eliminate_end_effect(profile, component)
    bc = profile.λc * profile.resolution
    maximum(-[-(extrema(c[n+1:n+bc])...) for n ∈ 1:bc:(length(c)-bc)])
end

"""
    Wt(W, λc=250)

Determine waviness of the waviness-profile W. Filterwavelength in accordance with [DIN EN ISO 3274:1998](DIN_EN_ISO_3274_1998.pdf).
"""
function t(profile::SurfaceProfile, component::Symbol, eliminate_end_effect=zeropad)
    c = eliminate_end_effect(profile, component)
    maximum(c) - minimum(c)
end

"""
    abbott(surf)

Calculate an [Abbott-Firestone-Curve](https://en.wikipedia.org/wiki/Abbott-Firestone_curve) for the line scan given in `surf`.
"""
abbott(surf) = [quantile(surf, i) for i ∈ 1:-0.01:0]


"""
Power nonlinearity, with energy C ∫ρ^α where ρ is the total density
"""
struct PowerNonlinearity
    C::Real
    α::Real
end
(P::PowerNonlinearity)(basis) = TermPowerNonlinearity(basis, P.C, P.α)

struct TermPowerNonlinearity <: Term
    basis::PlaneWaveBasis
    C::Real
    α::Real
end

function ene_ops(term::TermPowerNonlinearity, ψ, occ; ρ, kwargs...)
    basis = term.basis
    dVol = basis.model.unit_cell_volume / prod(basis.fft_size)

    E = term.C * sum(ρ.real .^ term.α) * dVol
    potential = @. term.C * term.α * ρ.real^(term.α-1)

    ops = [RealSpaceMultiplication(basis, kpoint, potential) for kpoint in basis.kpoints]
    (E=E, ops=ops)
end

_pnl_kernel(C, α, ρ) = @. C * α * (α-1) * ρ^(α-2)

function compute_kernel(term::TermPowerNonlinearity, dρspin=nothing;
                        ρ::RealFourierArray, ρspin=nothing, kwargs...)
    @assert term.basis.model.spin_polarization in (:none, :spinless, :collinear)
    K = Diagonal(vec(_pnl_kernel(term.C, term.α, ρ.real)))

    # PNL kernel is independent of spin, so to apply it to (ρtot, ρspin)^T
    # and obtain the same contribution to Vα and Vβ the operator has the block structure
    #     ( K 0 )
    #     ( K 0 )
    n_spin = term.basis.model.n_spin_components
    n_spin == 1 ? K : [K 0I; K 0I]
end

function apply_kernel(term::TermPowerNonlinearity, dρ, dρspin=nothing;
                      ρ::RealFourierArray, ρspin=nothing, kwargs...)
    @assert term.basis.model.spin_polarization in (:none, :spinless, :collinear)
    kernel = from_real(term.basis, _pnl_kernel(term.C, term.α, ρ.real) .* dρ.real)
    n_spin = term.basis.model.n_spin_components
    n_spin == 1 ? (kernel, ) : (kernel, kernel)  # PNL term does not care about spin
end


push!(LOAD_PATH, string(pwd(),"/../../code/src"))
using ModifiedStokesSolver
using Base.Test

include("ModStokesFMM2D.jl")

alpha = 20.0

## define source and target points

ns = 500
nt = 600

srand(1)
stokeslet_str = 1 - 2*rand(2,ns)
src = 1 - 2*rand(2,ns)
fvec = 1 - 2*rand(2,ns)
nvec = 1 - 2*rand(2,ns)
targ = 1 - 2*rand(2,nt)

# Compute using the kernels from ModifiedStokesSolver
function mss_stokeslet_direct(src, targ, stokeslet_str, alpha)
    ns = size(src, 2)
    nt = size(targ, 2)
    u = zeros(2, nt)
    for i=1:ns
        for j=1:nt
            f = stokeslet_str[:, i]
            r = src[:,i]-targ[:,j]
            u[:,j] += stokeslet(r, f, alpha)
        end
    end
    return u
end
function mss_stresslet_direct(src, targ, fvec, nvec, alpha)
    ns = size(src, 2)
    nt = size(targ, 2)
    u = zeros(2, nt)
    for i=1:ns
        for j=1:nt
            r = src[:,i]-targ[:,j]
            u[:,j] += stresslet(r, fvec[:,i], nvec[:,i], alpha)
        end
    end
    return u
end

## Stokeslet
u1 = mss_stokeslet_direct(src, targ, stokeslet_str, alpha)
u2 = fmm_stokeslet_direct(src, targ, stokeslet_str, alpha)
u3 = fmm_stokeslet_targ(src, targ, stokeslet_str, alpha)

## Stresslet
d1 = mss_stresslet_direct(src, targ, fvec, nvec, alpha)
d2 = fmm_stresslet_direct(src, targ, fvec, nvec, alpha)
d3 = fmm_stresslet_targ(src, targ, fvec, nvec, alpha)

## Print + test
println(" * Stokeslet")
Gerrdirect = norm(u1[:]-u2[:]) / norm(u1[:], Inf)
Gerrfmm = norm(u1[:]-u3[:]) / norm(u1[:], Inf)
@show Gerrdirect
@show Gerrfmm

println(" * Stresslet")
Terrdirect = norm(d1[:]-d2[:]) / norm(d1[:], Inf)
Terrfmm = norm(d1[:]-d3[:]) / norm(d1[:], Inf)
@show Terrdirect
@show Terrfmm


# This julia file builds on the Initial.julia
# It is based on the "How to efficiently select an arbitrary Clifford group element"
# Robert Koenig and John Smolin
# arXivv:1406.2170v1

# Copyright Robin Harper 2015-2020

# The premise here is that in order to generate a random Clifford, we can be sure of
# proper (Haar) randomness if either we generate all the cliffords for a certain quibit size
# and select one randomly OR if we can have a 1-1 mapping between the integers and the Cliffords
# and then we can just randomly select an integer. The latter is the one adopted in Koenig's paper.
# getNumberOfCliffords(n) returns the number of cliffords with n qubits.

# This ties into the Aaronsen/Gottesman paper becauase the generated cliffords are generated
# in such a way it is possible to create the tableau that would result IF the start myVector
# |000000> had been acted on by the chosen clifford unitary.
# We can then (decompose) the tableau using the algorithms specified in the Aaronson/Gottesman paper
# (see Initial.jl) to work out what combination of one-qubit (phase and hadmard) and two quibit
# (cnot) gates would create the unitary in the first place.

# The initial part of this file is a port of the python code in the Koenig/Smolin paper
# followed by an implementation of the decomposition algorithm in the Aaronson/Gottesman paper.
# There is also code to:
# -  Draw a decomposed circuit
# -  Display the kets of a |00..0> vector transformed by the circuit
# -  Form a 2^n 2^n complex matrix representing the "raw circuit" for use outside this formalism
# -  There is brute force method for finding the minimum set of gates to construct a given unitary
#    (even with only 3qubits this takes a tediously long time to run)
# -  There is the start of rationalise method, just now it eliminates chains of 4 phase gates or 2 hadamard gates.
#     Maybe the way to go here is to eliminate recognised patterns. [REF] uses a reverse tree building approach
#     which is interesting.
#
# Use is made of the ability of Julia to manipulate its own commands (a sort of hybrid functional/imperative)
# approach. With decomposition we are using arrays stored with the Tableau
# The first "commands" is a text represenetation of the commands needed to rebuild a decomposed end state.
# the second "executeCommands" contains the actual Julia instructions to rebuild the state.
# They are only populated when the appropriate decompositions have been made.
# This is better documented in the appropriate methods.

# The generation of the Clifford requires an element from the symplectic group
#include("Initial.jl")
#using Main.CHP

# Scratch variable in the global spce of Symplectic to be used with eval.
_t = Tableau(3)

# This just shows how big the groups get, Julia overflow is going to be a problem

function getNumberOfCliffords(n::BigInt)::BigInt
		return 2^(n^2+2*n)*prod([4^x-1 for x =1:n])
end

"""
    getNumberOfCliffords(n::Integer)::BigInt

returns the number of Cliffords for n qubits.
"""
function getNumberOfCliffords(n::Integer)::BigInt
		return getNumberOfCliffords(BigInt(n))
end

"""
    getNumberOfSymplecticCliffords(n::Integer)::BigInt

returns the number of Cliffords for n qubits, modulo the signs on each row
of the tableau. Note

getNumberOfCliffords(n) == getNumberOfSymplecticCliffords(n)*getNumerOfBitStringsCliffords(n)
"""
function getNumberOfSymplecticCliffords(n::BigInt)::BigInt
	return 2^(n^2)*prod([4^x-1 for x =1:n])
end

function getNumberOfSymplecticCliffords(n::Integer)::BigInt
	return getNumberOfSymplecticCliffords(BigInt(n))
end

"""
    getNumberOfBitStringsCliffords(n)

returns the number of bitstrings representing a +/- on each row for the Tableau.
Note:

getNumberOfCliffords(n) == getNumberOfSymplecticCliffords(n)*getNumerOfBitStringsCliffords(n)
"""
function getNumberOfBitStringsCliffords(n)
	return 2^(2*n)
end

"""
    cliffordToTableau(qubits,clifford,signs)

Converts an integer into a Clifford. You need to specify the number of qubits in the
tableau, the clifford number and the bit signs corresponding to the rows of the Tableau.
The commands used to create the Clifford can be found in Tableau.commands and Tableau.executeCommands.?
It is fine to specify numbers that are too large (they wrap), but the helper functions
```Julia
getNumberOfSymplecticCliffords(n::BigInt)::BigInt
getNumberOfBitStringsCliffords(n)
```
Will show how many cliffords there are for a certain number of qubits (Warning this is a very large number for n > 2)
"""
function cliffordToTableau(qubits,clifford,signs)
    t = Tableau(qubits)
	t.state = stabiliseSymp(symplectic(clifford,qubits),signs)
	decomposeState(t)
	return t
end

"""
    tableauToClifford(t::Tableau)

Returns the 'symplectic' number that corresponds to the clifford embedded in the Tableau.
Note that this is modulo the bit signs. I.e. the tableaus corresponding to e.g.
cliffordToTableau(1,3,3) and cliffordToTableau(1,3,4) will both return 3.
"""
function tableauToClifford(t::Tableau)
	symplecticinverse(decomposeStateToSymplectic(t.state))
end

"""
    decomposeState(tableau::Tableau,rationalise=true)

Decomposes the state stabilised by the Tableau using the Aaronsen/Gottesman method

That is it will decompose the "arbitrary" state into a series of CHP (cnot, hadamard and phase) gates.
This is unlikely to be the most concise decomposition. Unless rationalise is set to
false there will be some naive trimming of gates, basically self eliminating gates (e.g. 4 phase gates)
are removed. The commands (and the execCommands) in the tableau are wiped and recreated. Following the
decomposition the state is re-created using the new commands.
"""
function decomposeState(tableau::Tableau,rationalise=true)
	global _t
	tableau.commands = []
	tableau.executeCommands = []
	#addCommand(tableau,"output(state)",Expr(:call,:show,:_t))
    ss1=tableau.state
	while getState(ss1) < 11
		nextStep(tableau)
	end
	j=tableau.qubits
	addCommand(tableau,"initialise($j)",Expr(:call,:initialise,:_t))
	reverseCommands(tableau)
	removeRedundancy(tableau)
	_t = tableau
	tracking = _t.trackCommands
	_t.trackCommands = false
	for i in tableau.executeCommands
		eval(i)
	end
	_t.trackCommands = tracking

end

function drawCircuit(t::Tableau)
	commands = t.commands
    qasmCommands=[]
    labels=[]
	currentDir = pwd()
	for i = 1:size(commands,1)
		m = match(r"initialise\((.*)\)",commands[i])
		if (m!=nothing)
			for idx=1:t.qubits
    			#push!(qasmCommands,"qubit q$i")
                push!(labels,"q$idx")
			end
    	else
    		m=match(r"hadamard\((.*)\)",commands[i])
    		if (m!=nothing)
    			push!(qasmCommands,("H","q$(m.captures[1])"))
    		else
    			m=match(r"phase\((.*)\)",commands[i])
    			if (m!=nothing)
                    push!(qasmCommands,("P","q$(m.captures[1])"))
    			else
    				m=match(r"cnot\((.*),(.*)\)",commands[i])
    				if (m!=nothing)
    					push!(qasmCommands,("CNOT","q$(m.captures[2])","q$(m.captures[1])"))
    				end
    			end
    		end
    	end
    end
    return (qasmCommands,labels)
end

function qiskitCircuit(t::Tableau,name = "Circuit")
	commands = t.commands
    qasmCommands=["#Pass in the circuit and qubit register","def create$name(circs,qreg):"]

    labels=[]
	currentDir = pwd()
	for i = 1:size(commands,1)
		m = match(r"initialise\((.*)\)",commands[i])
		if (m!=nothing)
			for idx=1:Meta.parse(m.captures[1])
    			#push!(qasmCommands,"qubit q$i")
                push!(labels,"q$idx")
			end
    	else
    		m=match(r"hadamard\((.*)\)",commands[i])
    		if (m!=nothing)
    			push!(qasmCommands,"    circs.h(qreg[$(m.captures[1])])")
    		else
    			m=match(r"phase\((.*)\)",commands[i])
    			if (m!=nothing)
                    push!(qasmCommands,"    circs.s(qreg[$(m.captures[1])])")
    			else
    				m=match(r"cnot\((.*),(.*)\)",commands[i])
    				if (m!=nothing)
    					push!(qasmCommands,"    circs.cx(qreg[$(m.captures[1])],qreg[$(m.captures[2])])")
    				end
    			end
    		end
    	end
    end
	return join(qasmCommands,"\n")
end

###################### Functions used to create Clifford tabelau.

#takes square arrays and places them in a larger array, as follows
#   m1 0
#   0 m2
function directsum(m1,m2)
  n1=size(m1)[1]
  n2=size(m2)[1]
  out = zeros(n1+n2,n1+n2)
  out[1:n1,1:n1]=m1
  out[(n1+1):end,(n1+1):end]=m2
  return out
end

# Helper function used in decomposing the clifford
# builds up an array of the commands used to call the hadamard/phase/cnot.
# We are using two globals here for the sake of simplicity.
# execute commands is the Julia code required to call the relevant functions.

function addCommand(t::Tableau,comToAdd,comm)
	t.commands = append!(t.commands,[comToAdd])
	t.executeCommands=append!(t.executeCommands,[comm]);
end

# Returns the symplectic inner product of two vectors.

function inner(v,w)
	t=0
	for i in 1:(size(v)[1]>>1)
		 t+=v[2*i-1]*w[2*i]
		 t+=w[2*i-1]*v[2*i]
    end
    return t%2
 end

 function transvection(k,v)
 	if (k==0) return mod.(v,2)
 	end
 	return mod.((v+inner(k,v)*k),2)
 end

 function int2bits(i,n)
 	out = zeros(n)
 	for j in 1:n
 		out[j]=i&1
 		i>>=1
 	end
 	return out
 end

 function findtransvection(x,y)
 	out = zeros(UInt8,2,size(x)[1])
 	if (x==y)
 		return out
 	end
 	if inner(x,y) ==1
 		out[1,:]=mod.((x+y),2)
 		return out
 	end
 	z=zeros(size(x)[1])
 	for i in 1:(size(x)[1]>>1)
 		ii=2*i
 		if ((x[ii-1] + x[ii]) != 0) && ((y[ii-1]+y[ii])!=0)
 			# found the pair
 			z[ii-1]=(x[ii-1]+y[ii-1])%2
 			z[ii]=(x[ii]+y[ii])%2
 			if (z[ii-1]+z[ii])==0 #the same
 				z[ii]=1
 				if (x[ii-1] != x[ii])
 					z[ii-1]=1
 				end
 			end
 			out[1,:]=mod.((x+z),2)
 			out[2,:]=mod.((y+z),2)
 			return out
 		end
 	end
 	# ok we didn't find a pair
 	# need places where x is 00 and y does not
 	# and vice versa.
 	for i in 1:(size(x)[1]>>1)
 		ii=2*i
 		if ((x[ii-1]+x[ii]) !=0) && ((y[ii-1]+y[ii]) == 0)
 			if x[ii-1]==x[ii]
 				z[ii]=1
 			else
 				z[ii]=x[ii-1]
 				z[ii-1]=x[ii]
 			end
 			break
 		end
 	end
 	for i in 1:(size(x)[1]>>1)
 		ii=2*i
 		if ((x[ii-1]+x[ii]) ==0) && ((y[ii-1]+y[ii]) !=0 )
 			if y[ii-1]==y[ii]
 				z[ii]=1
 			else
 				z[ii]=y[ii-1]
 				z[ii-1]=y[ii]
 			end
 			break
 		end
 	end
 	# (x+z)%2 doesn't work on arrays, so mod function instead.
 	out[1,:]=mod.((x+z),2)
 	out[2,:]=mod.((y+z),2)
 	return out
 end

# The steps are as set out on page 5 of the paper.
# We use bigInts because we can (and if n was greater than about 64 we would need to).
function symplectic(i,n)
	nn=2*n
	s=BigInt(2)^nn-1
	k=(i%s)+1
	i=div(i,s)
	# step 2
	# There is probably a better way to do this, but it works.
	f1=[x == '0' ? UInt8(0) : UInt8(1) for x  in reverse(string(k,base=2,pad=nn))]
	#step 3
	e1=zeros(UInt8,nn)
	#define the first basis
	e1[1]=1
	T=findtransvection(e1,f1)
	# step 4
	divisor = BigInt(2)^(nn-1)
	value = i%divisor
	bits = [x == '0' ? UInt8(0) : UInt8(1) for x  in reverse(string(value,base=2,pad=nn-1))]

	#step 5
	# Note that in Julia we need to expressly copy
	eprime=copy(e1)
	for j in  3:nn
		eprime[j]=bits[j-1]
	end
	h0=transvection(T[1:1,:]',eprime)
	h0=transvection(T[2:2,:]',h0)

	#step 6
	if bits[1]==1
		f1*=0
	end

	#step 7
	id2=[1 0;0 1]

	if n!=1
		g=directsum(id2,symplectic(i>>(nn-1),n-1))
	else
		g=id2
	end
	for j in 1:nn
		g[j,:]=transvection(T[1:1,:]',g[j:j,:]')
		g[j,:]=transvection(T[2:2,:]',g[j:j,:]')
		g[j,:]=transvection(h0,g[j:j,:]')
		g[j,:]=transvection(f1,g[j:j,:]')
	end
	return g
end

function numberofcosets(n)
	x = 2^(2*n-1)*(2^(2*n)-1)
	return x
end

function bits2int(b,nn)
	output = BigInt(0)
	tmp = BigInt(1)
	for j in 1:nn
		if b[j]==1
			output = output + tmp
		end
		tmp = tmp*2
	end
	return output
end

function symplecticinverse(gn)
	n = round(Int,size(gn)[1]/2)
	nn=2*n

	# step 1

	v = gn[1,:]
	w = gn[2,:]

	# step 2

	e1 = zeros(nn)
	e1[1] = 1
	T = findtransvection(v,e1)

	# step 3
	tw = copy(w)
	tw = transvection(T[1,:],tw)
	tw = transvection(T[2,:],tw)
	b = tw[1]
	h0 = zeros(nn)
	h0[1]=1
	h0[2]=0
	for j in 3:nn
		h0[j]=tw[j]
	end

	# step 4

	bb = zeros(nn-1)
	bb[1]=b
	for j in 3:nn
		bb[j-1] = tw[j]
	end

	zv = bits2int(v,nn)-1
	zw = bits2int(bb,nn-1)

	cvw = zw*(2^(2*n)-1)+zv

	if n == 1
		return cvw
	end

	gprime = copy(gn)
	if b == 0
		for j in 1:nn
			gprime[j,:] = transvection(T[2,:],transvection(T[1,:],gn[j,:]))
			gprime[j,:] = transvection(h0,gprime[j,:])
			gprime[j,:] = transvection(e1,gprime[j,:])
		end
	else
		for j in 1:nn
			gprime[j,:]=transvection(T[2,:],transvection(T[1,:],gn[j,:]))
			gprime[j,:]=transvection(h0,gprime[j,:])

		end
	end
	gnew = gprime[3:nn,3:nn]
	gnidx = symplecticinverse(gnew)*numberofcosets(n)+cvw
	return gnidx
end

#Splits up the symplectic into the alpha,beta,gamma and delta arrays that
#specify the action of the clifford unitary on the X and Z Paulis respectively

function parseSymplectic(symp)
	#note that here we have (so far) ignored that for a,b,c,d we need to multiply an r and s
	#then I am guessing we just use rj as bit string on rhs. ie r1 to rn
	# so with one bit, we have and additional 2*2 * 6
	# then with two bits its 2^2 * 2^2 = 16
	s2 = round(Int,(size(symp)[1]/2))
	a = zeros(s2,s2)
	b = zeros(s2,s2)
	c = zeros(s2,s2)
	d = zeros(s2,s2)
	for i=1:s2
		for j = 1:s2
			a[i,j]=symp[(i-1)*2+1,(j-1)*2+1]
			b[i,j]=symp[(i-1)*2+1,j*2]
			c[i,j]=symp[i*2,(j-1)*2+1]
			d[i,j]=symp[i*2,j*2]
		end
	end
	return (a,b,c,d)
end

# Another way of doing this (incomplete - just an aide memoir)
# reshape(collect(Iterators.flatten(collect(zip(a,b)))),3,6)

function decomposeStateToαβγδ(state)
    n = size(state)[1]
    boundary = round(Int,n/2)
    return (state[1:boundary,1:boundary],state[1:boundary,boundary+1:end],
        state[boundary+1:end,1:boundary],state[boundary+1:end,boundary+1:end])
end

function decomposeStateToSymplectic(state)
    n = round(Int,size(state)[1]/2)
    finalD = n*2
    sympy=zeros(finalD,finalD)
    (α,β,γ,δ)= decomposeStateToαβγδ(state)
    for i=1:n
        for j = 1:n
            sympy[(i-1)*2+1,(j-1)*2+1]= α[i,j]
            sympy[(i-1)*2+1,j*2]=β[i,j]
            sympy[i*2,(j-1)*2+1]= γ[i,j]
            sympy[i*2,j*2]= δ[i,j]
        end
    end
    return sympy
end

# Takes the symplectic, parses it and uses it to create an Aaronson/Gottesman tableau
#The bits specify the 'bit' pattern that controls the sign
#For example in the 1 qubit there are only 6 unique 'symplectic' patterns
#but each of the two rows in the "tableau" can have bit signs of 0 or 1, leading
# to 6*4 different cliffords.
#this is what the bits is.

function stabiliseSymp(symp,bits)
	(a,b,c,d)=parseSymplectic(symp)

	n=size(a)[1]
	top = hcat(a,b,[ (bits >> (x-1)) %2 for x=1:n])
	for co=1:n
		bits = div(bits,2)
	end
  	bottom = hcat(c,d,[ (bits >> (x-1)) %2 for x=1:n])
 	state = convert(Array{Int32,2},vcat(top,bottom))
 end

# Takes tableau state and works out what state we are in the the 11 step deomposition
# where we are taking an arbitrary tableau state and "decomposing" it with basic gates
# back to the initial |00...00> state.
function getState(state)
	as = state[:,1:end]
	# calculate the states, backwards, so we get the "most refined" first
	#define some useful comparison matrices
	# The tableau is split up as follows:
	# A | B
	# ______
	# C | D
	#
	n=div(size(state,1),2) # half the dimension of this 2n x (2n+1) matrix
	state_i = Matrix{Int32}(I,n,n)
	state_z = zeros(Int32,n,n)
	state_r = zeros(Int32,n,1)
	if (as==vcat(hcat(state_i,state_z,state_r),hcat(state_z,state_i,state_r)))
		return 11
	end
	#split up into four matrices
	A=as[1:n,1:n]
	B=as[1:n,n+1:2*n]
	C=as[n+1:2*n,1:n]
	D=as[n+1:2*n,n+1:2*n]
	R=as[1:end,end]
	R2=as[n+1:end,end:end]
	if (B==state_z && C==state_z && R2==state_r)
		return 10
	end
	if (C==state_z && A==B  && R2==state_r)
		return 9
	end
	if (A==state_i && C==state_z && D==state_i && R2==state_r)
		return 7
	end
	if (C==state_i && D==state_z && as[n+1:2*n,end:end] == state_r)
		return 6
	end
	if (D==state_z )
		return 5
	end
	if (C==D)
		return 4
	end
	if (C==state_i)
		return 2
	end
	if (rank(C) == n)
		return 1
	end
	return 0
end

# The next series of vectors each implement one of the manipulations specified in Aaronson/Gottesman
# Section VI Canonical form.

function getFullRank(tableau::Tableau)
	n=tableau.qubits # half the dimension of this 2n x (2n+1) matrix
	svec = tableau.state
	for ii=1:n
		for i=1:n
			r1 = rank(svec[n+1:2*n,1:n])
			if r1 == n
				return
			end
			hadamard(tableau,i,false)
			r2 = rank(svec[n+1:2*n,1:n])
			if (r2 == n )
				#println("hadamard(",i,")")
				#addCommand(tableau,"hadamard($i)",Expr(:call,:hadamard,:_t,i))
				return
			end
			if (r2 > r1)
				#println("hadamard(",i,")",Expr(:call,:hadamard,:svec,i))
				#addCommand(tableau,"hadamard($i)",Expr(:call,:hadamard,:_t,i))
				continue
			else
				hadamard(tableau,i,false)
			end
		end
		hadamard(tableau,ii,false)
	end
end

function getfirstOne(myVector)
	#println(myVector)
	for i=1:length(myVector)
		if myVector[i]==1
			#println("return ",i)
			return i
		end
	end
	#print("******* THIS MAY BE AN ERROR BUT GET FIRST ONE COULDNT FIND A 1********\n")
	#print("$myVector\n")
	return 0
end

function makeCtheI(tableau::Tableau,offset=-1)
	n=tableau.qubits # half the dimension of this 2n x (2n+1) matrix
	svec = tableau.state
	sanity = 100;
	if offset==-1 offset=n # no offset means we assume its the "C" we are making Int32
	end
	allDone = false;
	while !allDone
		for i=1:n
			#print("Checking for $i\n")
			allDone = identifytheRow(tableau,offset,i)
			if allDone == false
				break
			end
		end
		sanity = sanity -1;
		if (sanity < 1)
			print("Stuck in a loop making C, I\n")
			exit()
		end
	end
	# now we need double phases to make r_n+1..r_2n to be zero
end

function zapPhase(tableau::Tableau,offset=-1)
	n=tableau.qubits # half the dimension of this 2n x (2n+1) matrix
	svec = tableau.state
	if offset==-1 offset=n # no offset means we assume its the "C" we are making Int32
	end
	for i=1:n
		if (svec[offset+i,end] ==1)
			phase(tableau,i,false)
			#addCommand(t,"phase($i)",Expr(:call,:phase,:_t,i))
			phase(tableau,i,false)
			#addCommand(t,"phase($i)",Expr(:call,:phase,:_t,i))
		end
	end
end

# What I mean is take the passed in row, and use cnots to make it an identity element.
function identifytheRow(tableau::Tableau,offset,i)
	n=tableau.qubits # half the dimension of this 2n x (2n+1) state matrix
	#Make sure the diagonal element is 1, we can only use after the element to cnot with
	#println("Checking Diagonal ",i,",",i)
	svec = tableau.state
	if svec[offset+i,i] ==0 # need to get this diagonal equal to 1
		#println("It wasn't one in the diagonal we have offset ",offset, " i ", i)
		t = getfirstOne(svec[offset+i,i+1:n])
		#println(svec[offset+i,i:n])
		#println("Got a t back of ",t)
		if t==0
			#println("Couldnt make this the identity - redoing rank\n")
			getFullRank(tableau)
			return false
		end
		t=t+i
		#println("cnot(",t,",",i,")")
		cnot(tableau,t,i,false)
		#addCommand(tableau,"cnot($t,$i)",Expr(:call,:cnot,:_t,t,i))
		#println(svec[offset+i,1:n])
	end

	for j=1:i-1
		#println("Check ",i,",",j)
		if svec[offset+i,j]!= 0
			#println("Wasnt zero")
			#addCommand(tableau,"cnot($i,$j)",Expr(:call,:cnot,:_t,i,j))
			cnot(tableau,i,j,false)
			#println(svec[offset+i,1:n])
		end
	end
	#println("Up to diagonal cleared")
	#println(svec[offset+1,1:n])
	# so the diagonal element is now 1, make the rest zero
	for j=i+1:n
		#println("Checking to make zero",i,",",j)
		if svec[offset+i,j]==1
			#println("It was one")
			#println("cnot(",i,",",j,")")
			cnot(tableau,i,j,false)
			#addCommand(tableau,"cnot($i,$j)",Expr(:call,:cnot,:_t,i,j))
			#println(svec)
		end
	end
	return true
end

function diagonaliseD(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		if svec[n+i,n+i]== 0
			phase(tableau,i,false)
			#println("phase(",i,")");
			#addCommand("phase($i)",Expr(:call,:phase,:_t,i))
		end
	end
end

function diagonaliseB(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		if svec[i,n+i]== 0
			phase(tableau,i,false)
			#println("phase(",i,")");
			#addCommand(tableau,"phase($i)",Expr(:call,:phase,:_t,i))
		end
	end
end

function cmdm(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		for j=1:n
			if svec[n+i,j] == svec[n+i,n+j]
				continue
			end
			if svec[n+i,j] == 0
				cnot(tableau,j,i,false)
				#println("cnot(",j,",",i,")")
				#addCommand("cnot($j,$i)",Expr(:call,:cnot,:_t,j,i))
			else
				phase(tableau,j,false)
				#println("phase(",j,")")
				#addCommand("phase($j)",Expr(:call,:phase,:_t,i))
			end
		end
	end
end
function ambm(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		for j=1:n
			if svec[i,j] == svec[i,n+j]
				continue
			end
			if svec[i,j] == 0
				cnot(tableau,j,i,false)
				#println("cnot(",j,",",i,")")
				#addCommand(tableau,"cnot($j,$i)",Expr(:call,:cnot,:_t,j,i))
			else
				phase(tableau,j,false)
				#println("phase(",j,")")
				#addCommand(tableau,"phase($j)",Expr(:call,:phase,:_t,j))
			end
		end
	end
end

function zapD(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		for j=1:n
			if svec[n+i,n+j]==1
				phase(tableau,j,false)
				#println("phase(",j,")")
				#addCommand("phase($j)",Expr(:call,:phase,:_t,j))
			end
		end
	end
end

function zapB(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		for j=1:n
			if svec[i,n+j]==1
				phase(tableau,j,false)
				#println("phase(",j,")")
				#addCommand(tableau,"phase($j)",Expr(:call,:phase,:_t,j))
			end
		end
	end
end

function hadamardHard(tableau::Tableau)
	n=tableau.qubits
	svec = tableau.state
	for i=1:n
		hadamard(tableau,i,false)
		#println("hadamard(",i,")")
		#addCommand(tableau,"hadamard($i)",Expr(:call,:hadamard,:_t,i))
	end
end

function makeItR0(tableau::Tableau,offset=-1)
	n=tableau.qubits
	svec = tableau.state
	if offset<0
		offset=n
	end
	for i=1:n
		if svec[offset+i,2*n+1]== 1
			phase(tableau,i,false)
			phase(tableau,i,false)
			#println("phase(",i,")")
			#println("phase(",i,")")
			#addCommand(tableau,"phase($i)",Expr(:call,:phase,:_t,i))
			#addCommand(tableau,"phase($i)",Expr(:call,:phase,:_t,i))
		end
	end
end

function nextStep(tableau::Tableau)
	n=tableau.qubits # half the dimension of this 2n x (2n+1) matrix
	svec = tableau.state

	currentState = getState(svec)
	#print(currentState)
	if (currentState == 0)
		getFullRank(tableau)
	elseif currentState == 1
		makeCtheI(tableau)
	elseif currentState < 4
		diagonaliseD(tableau)
		cmdm(tableau)
	elseif currentState == 4
		zapD(tableau)
		makeItR0(tableau)
	elseif currentState == 5
		makeCtheI(tableau)
		zapPhase(tableau)
	elseif currentState == 6
		hadamardHard(tableau)
	elseif currentState < 9
		diagonaliseB(tableau)
		ambm(tableau)
	elseif currentState == 9
		zapB(tableau)
		makeItR0(tableau,0)
	elseif currentState == 10
		makeCtheI(tableau,0) # 0 = actually its A
		zapPhase(tableau,0)
	elseif currentState == 11
		println("Done")
	end
	return
end

# maxGates is the number of gates we are allowed (an integer offset to an array)
# gates is the "stack" <- an array of gates we are applying.
# increment index 1, if its greater then maxGate, increment the one above it.
function incrementGate!(gates,maxGate,offset=1)
	if (size(gates,1) < offset )
		push!(gates,1)
		return
	end
	currentValue = gates[offset]
	currentValue = currentValue +1
	if (currentValue > maxGate)
		gates[offset]=1
		incrementGate!(gates,maxGate,offset+1)
	else
		gates[offset]=currentValue
	end
end

function bruteForceBreadthFirst(clifford)
	global svec
	count = 0
	n=div(size(clifford,1),2) # half the dimension of this 2n x (2n+1) matrix
	if (n > 4)
		println("This might take some time!, there are ",n*(n+1) + n + n, " different gates , we might need ", n*n/log(n), " of them so I *might* have to check ", n*n/log(n)^(n*(n+1) + n + n), "combos !")
		println("Ctrl-C could be your friend")
	end
	gates = (Expr)[]
	svec = []
	# If we have a n quibit system then there are a total of
	# n*(n+1) possible cnots, n phase gates and n hadamards we can apply
	for i = 1:n
	  for j = 1:n
	  	if i!=j
	  		push!(gates,Expr(:call,:cnot,:svec,i,j,false))
	  	end
	  end
	  push!(gates,draw)
	  push!(gates,Expr(:call,:phase,:svec,i,false))
	end
	# so we now have an array of all possible gates we can apply to the system
	# we want to search through the gates, we "should" need at most n^2 of them.
	# The process will be to have a stack of gates to apply, check them out
	# if we match getState(11) i.e. they reverted the clifford to the initial state
	# we have found the gates needed
	# otherwise we increment our "bottom" gate, and if we have gone though them, recursively increment the
	# gate above it (incrementGate! does this for us)
	state = 0
	maxGates = size(gates,1);
	gatesToApply=(Int32)[]
	if getState(clifford) == 11
		println("Very funny you don't need to do anything")
		return
	end
	currentAt = 1
	while (state!=11 && size(gatesToApply,1) < n*n*n*n*n*n*n*n*n*n)
		svec = copy(clifford)
		incrementGate!(gatesToApply,maxGates)
		count += 1
		if (size(gatesToApply,1) > currentAt)
			print("Trying: ")
				for i=1:size(gatesToApply,1)
				print(gatesToApply[i]," ")
			end
			println("")
			currentAt+=1
		end
		for index=1:size(gatesToApply,1)
			#println("Going to apply ",gates[gatesToApply[index]])
			eval(gates[gatesToApply[index]])
		end
		state = getState(svec)
	end

	if (state == 11)
		println("Found it after (only) $count permuations")
		println("Gates applied")
		reverse!(gatesToApply)
		for i = 1:size(gatesToApply,1)
			println(gates[gatesToApply[i]])
		end
	else
		println("Didn't find it!")
	end
	return gatesToApply
end

#coding for current bits
# 1 = phase, 2 = hadamard
# 2 + x  = cnot control (with target at x)
# 2 + n + x = cnot target (with control at x)

function checkGate(tableau::Tableau,index,currentBits,n,toDelete)
	commands = tableau.commands
	executeCommands = tableau.executeCommands

	checking = commands[index]
	m = match(r"Tableau\((.*)\)",checking)
	if (m!=nothing)
    	return
    end
    m=match(r"hadamard\((.*)\)",checking)
    if (m!=nothing)
    	bit = parse(Int,(m.captures[1]))
    	if size(currentBits[bit],1) > 0
    		if currentBits[bit][end][1] == 2
    			# two hadamards make a nothing
    			(been,gone) = pop!(currentBits[bit])
    			push!(toDelete,gone)
    			push!(toDelete,index)
			else
				push!(currentBits[bit],(2,index))
			end
		else
				push!(currentBits[bit],(2,index))
		end
		return
	end
	m=match(r"phase\((.*)\)",checking)
    if (m!=nothing) # Its a phase, we need 4 of these
    	bit = parse(Int,(m.captures[1]))
    	if size(currentBits[bit],1) > 2
    		if currentBits[bit][end][1] == 1 && currentBits[bit][end-1][1] == 1 && currentBits[bit][end-2][1] == 1
    			# four phases make a nothing.
    			(b,g) = pop!(currentBits[bit])
    			push!(toDelete,g)
    			(b,g) = pop!(currentBits[bit])
    			push!(toDelete,g)
    			(b,g) = pop!(currentBits[bit])
    			push!(toDelete,g)
    			push!(toDelete,index)
			else
				push!(currentBits[bit],(1,index))
			end
		else
				push!(currentBits[bit],(1,index))
		end
		return
	end
	m=match(r"cnot\((.*),(.*)\)",checking)
    if (m!=nothing)
    		cbit = parse(Int,(m.captures[1]))
    		tbit = parse(Int,(m.captures[2]))
    		# so just now we are going to catch two cnots in a row
    		# needs to match both bits.
    		cbitNo = 2+tbit
    		tbitNo = 2+n+cbit
    		if (size(currentBits[cbit],1) > 0 && size(currentBits[tbit],1) > 0 && currentBits[cbit][end][1] == cbitNo &&
    			currentBits[tbit][end][1] == tbitNo)
    			(b,g)=pop!(currentBits[cbit])
    			push!(toDelete,g)
    			(b,g)=pop!(currentBits[tbit])
    			# dont need to push the deletion of this gate twice.
    			push!(toDelete,index)
    		else
    			push!(currentBits[cbit],(cbitNo,index))
    			push!(currentBits[tbit],(tbitNo,index))
    		end
    end
end

function removeRedundancy(tableau::Tableau)
	n = tableau.qubits
	executeCommands = tableau.executeCommands
	commands = tableau.commands
	toDelete = (Int32)[]
	bitsOf = (Array{Tuple{Int32,Int32}})[]
	for i=1:n
    	push!(bitsOf,[])
	end
	for i=1:size(commands,1)
		checkGate(tableau,i,bitsOf,n,toDelete)
	end
	deleteat!(commands,sort!(toDelete))
	deleteat!(executeCommands,toDelete)
end

##I am going to use the symplectic method to generate all the cliffords
# - Noting that we have 24, therefore there are 6 'symplectics' and 2*2=4 different bit patterns
# - Decompose, internally generates the stabilised regime and then works out what pattern of Hadamards and Phase gates are necessary to generate this, these are stored in "commands". [Here we are in the one Qubit regime - so there are no CNOTS]
# - I use a simple loop to execute the commands using a phase gate and hadamard gate. If we were to repeat this > 1 qubit, we would need to blow up the phase gates etc (for example, if we were in 3 qubit regime and it was phase(2) i.e. phase on the second qubit, we would multiply by $I\otimes P \otimes I$
# - Better still we could define our gates (here sphase, and shadamard) to have errors and get an erroneous clifford

"""
    generateRawCliffords(;phaseGate=[1 0;0 im],hadmardGate=1/sqrt(2)*[1 1;1 -1])::Array{Complex{Float64},2}[]

returns the single qubit cliffords, generated using the passed in  phase and hadamard gates (perfect gates default)
"""
function generateRawCliffords(;phaseGate=[1 0;0 im],hadamardGate=1/sqrt(2)*[1 1;1 -1])
	straightClifford = Array{Complex{Float64},2}[]
	state=[1 0;0 1]
	for i=1:6
    	for j=1:4
        	tab = cliffordToTableau(1,i,j)
        	for t in tab.commands
           		m = match(r"initialise\((.*)\)",t)
            	if (m!=nothing)
                	state=[1 0;0 1]
            	else
                	m=match(r"phase\((.*)\)",t)
                	if (m!=nothing)
           	        	state=state*phaseGate
            	    else
                    	m=match(r"hadamard\((.*)\)",t)
                    	if (m!=nothing)
                        	state=state*hadamardGate
                    	end
                	end
            	end
        	end
        	push!(straightClifford,round.(state,digits=10))
    	end
    end
    return straightClifford
end

"""
	drawCircuit(commands)
	Takes the command string and returns their qasm equivalent.
	Can then be passed into e.g. qiskit or something similar to draw them.
"""
function drawCircuit(commands)
	qasmCommands=[]
    labels=[]
	currentDir = pwd()
	for i = 1:size(commands,1)
		m = match(r"setup\((.*)\)",commands[i])
		if (m!=nothing)
			for idx=1:Meta.parse(m.captures[1])
    			#push!(qasmCommands,"qubit q$i")
                push!(labels,"q$idx")
			end
    	else
    		m=match(r"hadamard\((.*)\)",commands[i])
    		if (m!=nothing)
    			push!(qasmCommands,("H","q$(m.captures[1])"))
    		else
    			m=match(r"phase\((.*)\)",commands[i])
    			if (m!=nothing)
                    push!(qasmCommands,("P","q$(m.captures[1])"))
    			else
    				m=match(r"cnot\((.*),(.*)\)",commands[i])
    				if (m!=nothing)
    					push!(qasmCommands,("CNOT","q$(m.captures[2])","q$(m.captures[1])"))
    				end
    			end
    		end
    	end
    end
    return (qasmCommands,labels)
end

"""
	makeFromCommand(command)

	Takes the commands in a Tableau and uses them to build a matrix representing the Tableau (in the computational basis)
	Uses the original initialise command to determine the size of the system.
	Note: for many qubits this gets quite big (scales as \$2^n\$)
"""
function makeFromCommand(command)
	up=[1;0]
	down=[0;1]
	control=[1 0; 0 0]
	target= [0 0;0 1]
	id=[1 0;0 1]
	x=[0 1;1 0]
	maxBits=1
    currentState = [1 0 0 0;0 1 0 0;0 0 1 0;0 0 0 1]
    si = [1 0;0 1]
    phaseG = [1 0;0 im]
    sphase1=kron([1 0;0 im],si)
    sphase2=kron(si,[1 0;0 im])
    shadamard=1/sqrt(2)*[1 1;1 -1]
    cnot12 = [1 0 0 0;0 1 0 0;0 0 0 1;0 0 1 0]
    cnot21 = [1 0 0 0;0 0 0 1;0 0 1 0;0 1 0 0]
    for t in command
        m = match(r"initialise\((.*)\)",t)
        if (m!=nothing)
        	bit =parse(Int,(m.captures[1]))
        	maxBits=bit
            state = id;
            for i=2:bit
            	state = kron(state,id);
            end
            currentState = state
        else
           m=match(r"phase\((.*)\)",t)
           if (m!=nothing)
                bit = parse(Int,(m.captures[1]))
                tphase = phaseG
                for i=1:bit-1
                	tphase = kron(id,tphase)
				end
				for i=bit+1:maxBits
					tphase = kron(tphase,id)
				end
				currentState = currentState*tphase
            else
                m=match(r"hadamard\((.*)\)",t)
                if (m!=nothing)
                    bit = parse(Int,(m.captures[1]))
                    thad = shadamard
                    for i=1:bit-1
                    	thad=kron(id,thad)
                    end
                    for i=bit+1:maxBits
                    	thad=kron(thad,id)
                    end
                    currentState=currentState*thad
                else
                    m=match(r"cnot\((.*),(.*)\)",t)
                    if (m!=nothing)
                        cbit = parse(Int,(m.captures[1]))
                        tbit = parse(Int,(m.captures[2]))

                        # just doing this each way, its easier for my mind.
                        if cbit < tbit
                        	startC=cbit
                        	endC=tbit
                        	l1 = id
                        	l2  = x
                        	for i=cbit+1:tbit-1
                        		l1=kron(id,l1)
                        		l2=kron(id,l2)
                        	end
                        	cbitbit = kron(control,l1)+kron(target,l2)
                        else
                        	startC=tbit
                        	endC=cbit
                        	l1 = id
                        	l2  = x
                        	for i=tbit+1:cbit-1
                        		l1=kron(l1,id)
                        		l2=kron(l2,id)
                        	end
                        	cbitbit = kron(l1,control)+kron(l2,target)
                        end
                        for i = 1:startC-1
                        	cbitbit=kron(id,cbitbit)
                        end
                        for i=endC+1:maxBits
                        	cbitbit=kron(cbitbit,id)
                        end
                        currentState = currentState*cbitbit
                    end
                end
            end
        end
    end
    return  currentState
end

"""
	makeFromCommand(t::Tableau)

	Uses the commands stored in a Tableau and uses them to build a matrix representing the Tableau (in the computational basis)
	Uses the original initialise command to determine the size of the system.
	Note: for many qubits this gets quite big (scales as \$2^n\$)
"""
function makeFromCommand(t::Tableau)
	return makeFromCommand(t.commands)
end

"""
	makeInverseFromCommand(t::Tableau)

	Uses the commands stored in a Tableau and uses them to build a matrix representing the inverse of the Tableau (in the computational basis)
	Uses the original initialise command to determine the size of the system.
	This might be useful for instance if you have a Tableau represnting Pauli measurements and you wanted the Clifford that maps them to the computational basis.
	Note: for many qubits this gets quite big (scales as \$2^n\$)
"""
function makeInverseFromCommand(t::Tableau)
	return makeFromCommand(reverse(t.commands[2:end]))
end

function reverseCommands(t::Tableau)
	len = length(t.commands)
	@assert length(t.executeCommands)==len
	newCommands = []
	newExecute  = []
	for i = len:-1:1
		m=match(r"phase\((.*)\)",t.commands[i])
		if (m!=nothing)
			push!(newCommands,t.commands[i])
			push!(newCommands,t.commands[i])
			push!(newCommands,t.commands[i])
			push!(newExecute,t.executeCommands[i])
			push!(newExecute,t.executeCommands[i])
			push!(newExecute,t.executeCommands[i])
        else
        	push!(newCommands,t.commands[i])
			push!(newExecute,t.executeCommands[i])
        end
    end
	t.commands = newCommands
	t.executeCommands = newExecute
    return
end


using FractionalDim
using LightXML

function read_xml_mesh(mesh::AbstractString)::Vector{Segment{3}}
    xroot = root(parse_file(mesh))
    mesh = first(child_elements(xroot))
    @assert attribute(mesh, "celltype") == "interval"

    dim = parse(attribute(mesh, "dim"))

    nodes, cells, _ = child_elements(mesh)

    nnodes = parse(attribute(nodes, "size"))
    X = zeros((dim, nnodes))

    get_coords = if dim == 2
        (e, col) -> begin
            X[dim*col + 1] = parse(attribute(e, "x"))
            X[dim*col + 2] = parse(attribute(e, "y"))
        end
    else
        (e, col) -> begin
            X[dim*col + 1] = parse(attribute(e, "x"))
            X[dim*col + 2] = parse(attribute(e, "y"))
            X[dim*col + 3] = parse(attribute(e, "z"))
        end
    end

    col = 0
    for node in child_elements(nodes)
        get_coords(node, col)
        col += 1
    end

    segments = Vector{Segment{dim}}()
    for cell in child_elements(cells)
        v0 = X[:, parse(attribute(cell, "v0")) + 1]
        v1 = X[:, parse(attribute(cell, "v1")) + 1]

        push!(segments, Segment{dim}(Point{dim}(v0), Point{dim}(v1)))
    end
    segments
end

# Read
maybe_fractal = read_xml_mesh("vasc_mesh.xml");

# And now estimate the dimension
counter = BoxCounter(maybe_fractal)
sizes, counts = Vector{Float64}(), Vector{Int}()

state = start(counter, 16)

rate, rate_lsq = Inf, Inf
while !done(counter, state)
    state = next(counter, state)
    size, count = last(state)
    push!(sizes, size)
    push!(counts, count)

    if length(sizes) > 1
        rate = -log(counts[end]/counts[end-1])/log(sizes[end]/sizes[end-1])
    end

    println((size, count, rate))
    writedlm("vasc_mesh_FD.txt", [sizes counts])
end


@testset "sobolindices" begin
    x = RandomVariable.(Uniform(-π, π), [:x1, :x2, :x3])

    ishigami = Model(
        df -> sin.(df.x1) + 7 .* sin.(df.x2) .^ 2 + 0.1 .* df.x3 .^ 4 .* sin.(df.x1), :f
    )

    n = 1000

    firstorder_analytical = [0.3138, 0.4424, 0.00]
    totaleffect_analytical = [0.5574, 0.4424, 0.2436]

    @testset "Monte Carlo" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, MonteCarlo(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end

    @testset "Sobol" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, SobolSampling(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end

    @testset "Halton" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, MonteCarlo(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end

    @testset "Latin Hypercube" begin
        Random.seed!(8128)

        si = sobolindices([ishigami], x, :f, MonteCarlo(n))

        Random.seed!()

        @test all(isapprox(si.FirstOrder, firstorder_analytical; rtol=0.1))
        @test all(isapprox(si.TotalEffect, totaleffect_analytical; rtol=0.1))
    end
end


@doc raw"""
    hankel_image(v, channel_size, kernel_size)

Creates a Hankel array from `v` with channel and kernel dimensions
`channel_size` and `kernel_size`, respectively. The returned array `V` satisfies:

```math
V_{c,j,k,n} = V_{c,j+k-1,n}
```

where `c,j,k,n` are multi-indices, with `c` traversing channels,
`j` traversing the kernel, `k` the convolution output, and `n` batches.
"""
function hankel_image(
    img::AbstractArray, channel_sz::NTuple{C,Int}, kernel_sz::NTuple{N,Int},
    stride::NTuple{N,Int}, pad::NTuple{N2,Int}, dilation::NTuple{N,Int}
) where {C,N,N2}
    input_sz = size(img)[(C + 1):(C + N)]
    batch_sz = size(img)[(C + N + 1):end]
    @assert size(img) == (channel_sz..., input_sz..., batch_sz...)
    out_sz = output_size(kernel_sz, input_sz; stride, pad, dilation)
    pad_lo = CartesianIndex(splitpad(pad).lo)
    X = zeros(eltype(img), channel_sz..., kernel_sz..., out_sz..., batch_sz...)
    for k in CartesianIndices(out_sz), j in CartesianIndices(kernel_sz)
        i = out2in(j, k; dilation, stride) - pad_lo
        all(1 .≤ Tuple(i) .≤ input_sz) || continue
        for b in CartesianIndices(batch_sz), c in CartesianIndices(channel_sz)
            X[c,j,k,b] = img[c,i,b]
        end
    end
    return X
end

function hankel_image(
    img::AbstractArray, channel_sz::NTuple{C,Int}, kernel_sz::NTuple{N,Int};
    stride = 1, pad = 0, dilation = 1
) where {C,N}
    return hankel_image(
        img, channel_sz, kernel_sz,
        expand_tuple(Val(N), stride),
        parsepad(kernel_sz, dilation, pad),
        expand_tuple(Val(N), dilation)
    )
end

function hankel_image(img::AbstractArray, C::Int, kernel_size::TupleN{Int}; kwargs...)
    channel_size = size(img)[1:C]
    return hankel_image(img, channel_size, kernel_size; kwargs...)
end

@doc raw"""
    hankel_weight(w, channel_size, input_size)

Creates a Hankel array from `v` with channel and kernel dimensions
`channel_size` and `kernel_size`, respectively. The returned array `V` satisfies:

```math
V_{c,j,k,n} = V_{c,j+k-1,n}
```

where `c,j,k,n` are multi-indices, with `c` traversing channels,
`j` traversing the kernel, `k` the convolution output, and `n` batches.
"""
function hankel_weight(
    w::AbstractArray, channel_sz::NTuple{C,Int}, in_sz::NTuple{N,Int};
    stride = 1, pad = 0, dilation = 1
) where {C,N}
    kernel_sz = size(w)[(C + 1):(C + N)]
    return hankel_weight(
        w, channel_sz, in_sz,
        expand_tuple(Val(N), stride),
        parsepad(kernel_sz, dilation, pad),
        expand_tuple(Val(N), dilation)
    )
end

function hankel_weight(
    w::AbstractArray, channel_sz::NTuple{C,Int}, in_sz::NTuple{N,Int},
    stride::NTuple{N,Int}, pad::NTuple{N2,Int}, dilation::NTuple{N,Int}
) where {C,N,N2}
    kernel_sz = size(w)[(C + 1):(C + N)]
    hidden_sz = size(w)[(C + N + 1):end]
    @assert size(w) == (channel_sz..., kernel_sz..., hidden_sz...)
    out_sz = output_size(kernel_sz, in_sz; stride, pad, dilation)
    pad_lo = CartesianIndex(splitpad(pad).lo)
    A = zeros(eltype(w), channel_sz..., in_sz..., hidden_sz..., out_sz...)
    for k in CartesianIndices(out_sz), j in CartesianIndices(kernel_sz)
        i = out2in(j, k; dilation, stride) - pad_lo
        all(1 .≤ Tuple(i) .≤ in_sz) || continue
        for μ in CartesianIndices(hidden_sz), c in CartesianIndices(channel_sz)
            A[c,i,μ,k] = w[c,j,μ]
        end
    end
    return A
end

function hankel_weight(w::AbstractArray, C::Int, in_sz::TupleN{Int}; kwargs...)
    channel_size = size(w)[1:C]
    return hankel_weight(w, channel_size, in_sz; kwargs...)
end

"""
    hankel(convrbm, input_size)

Returns a dense `RBM` equivalent to `convrbm` with replicated weights.
"""
function hankel(rbm::ConvRBM, in_size::NTuple{N,Int}) where {N}
    @assert kernel_ndims(rbm) == N
    @assert rbm.groups == 1 # groups > 1 not supported
    w = hankel_weight(rbm.w, channel_size(rbm), in_size; rbm.stride, rbm.pad, rbm.dilation)
    out_size = output_size(rbm, in_size)
    vis = replicate(visible(rbm), in_size...)
    hid = replicate(hidden(rbm), out_size...)
    return RBM(vis, hid, w)
end

"""
    replicate(layer, n...)

Returns a new layer of size `(size(layer)..., n...)` by repeating the original `layer`
along the new dimensions.
"""
replicate(l::AbstractLayer, n::Int...) = repeat(l, map(Returns(1), size(l))..., n...)


export  ApproximateValueIteration,
        solve,
        NearestNeighborValueFunction,
        fit!,
        LocallyWeightedValueFunction,
        MultilinearValueFunction,
        SimplexValueFunction,
        LinearRegressionValueFunction

"""
    struct ApproximateValueIteration
        Uθ # initial parameterized value function that supports fit!
        S # set of discrete states for performing backups
        k_max # maximum number of iterations
    end
"""
struct ApproximateValueIteration
    Uθ # initial parameterized value function that supports fit!
    S # set of discrete states for performing backups
    k_max # maximum number of iterations
end

"""
    function solve(M::ApproximateValueIteration, 𝒫::MDP)
"""
function solve(M::ApproximateValueIteration, 𝒫::MDP)
    Uθ, S, k_max = M.Uθ, M.S, M.k_max
    for k in 1:k_max
        U = [backup(𝒫, Uθ, s) for s in S]
        fit!(Uθ, S, U)
    end

    return ValueFunctionPolicy(𝒫, Uθ)
end

"""
    mutable struct NearestNeighborValueFunction
        k # number of neighbors
        d # distance function d(s, s′)
        S # set of discrete states
        θ # vector of values at states in S
    end
"""
mutable struct NearestNeighborValueFunction
    k # number of neighbors
    d # distance function d(s, s′)
    S # set of discrete states
    θ # vector of values at states in S
end

"""
    function (Uθ::NearestNeighborValueFunction)(s)

Approximate value function at s as average of value function over k nearest states (whose value function value is already known).
"""
function (Uθ::NearestNeighborValueFunction)(s)
    # distances from s to each state s′ in state subset S
    dists = [Uθ.d(s,s′) for s′ in Uθ.S]
    # indicies of shortest 1:k distances (ind[1] == index of closest neighbor)
    ind = sortperm(dists)[1:Uθ.k]
    # average value of k nearest neighbors
    return mean(Uθ.θ[i] for i in ind)
end

"""
    function fit!(Uθ::NearestNeighborValueFunction, S, U)
"""
function fit!(Uθ::NearestNeighborValueFunction, S, U)
    Uθ.θ = U
    return Uθ
end

"""
    mutable struct LocallyWeightedValueFunction
        k # kernel function k(s, s′)
        S # set of discrete states
        θ # vector of values at states in S
    end
"""
mutable struct LocallyWeightedValueFunction
    k # kernel function k(s, s′)
    S # set of discrete states
    θ # vector of values at states in S
end

"""
    function (Uθ::LocallyWeightedValueFunction)(s)
"""
function (Uθ::LocallyWeightedValueFunction)(s)
    w = normalize([Uθ.k(s,s′) for s′ in Uθ.S], 1)
    return Uθ.θ ⋅ w
end

"""
    function fit!(Uθ::LocallyWeightedValueFunction, S, U)
"""
function fit!(Uθ::LocallyWeightedValueFunction, S, U)
    Uθ.θ = U
    return Uθ
end

# ------------------ MultilinearValueFunction ---------------------

# Also implemented in Interpolations.jl

"""
    mutable struct MultilinearValueFunction
        o # position of lower-left corner
        δ # vector of widths
        θ # vector of values at states in S
    end
"""
mutable struct MultilinearValueFunction
    o # position of lower-left corner
    δ # vector of widths
    θ # vector of values at states in S
end

"""
    function (Uθ::MultilinearValueFunction)(s)

Use multilinear interpolation to estimate the value of state vector `s` for known state values `θ` over a grid defined by lower-left vertex `o` and vector of widths `δ`.  Verticies of the grid can all be written as `o + δ.*i` for some non-negative integer vector `i`.
"""
function (Uθ::MultilinearValueFunction)(s)
    o, δ, θ = Uθ.o, Uθ.δ, Uθ.θ
    Δ = (s - o)./δ
    # Multidimensional index of lower-left cell
    i = min.(floor.(Int, Δ) .+ 1, size(θ) .- 1)
    vertex_index = similar(i)
    d = length(s)
    u = 0.0
    for vertex in 0:2^d-1
        weight = 1.0
        for j in 1:d
            # Check whether jth bit is set
            if vertex & (1 << (j-1)) > 0
                vertex_index[j] = i[j] + 1
                weight *= Δ[j] - i[j] + 1
            else
                weight *= i[j] - Δ[j]
            end
        end
        u += θ[vertex_index...]*weight
    end

    return u
end

"""
    function fit!(Uθ::MultilinearValueFunction, S, U)
"""
function fit!(Uθ::MultilinearValueFunction, S, U)
    Uθ.θ = U
    return Uθ
end

# --------------- SimplexValueFunction ------------------
#
# Also implemented in GridInterpolations.jl

"""
    mutable struct SimplexValueFunction
        o # position of lower-left corner
        δ # vector of widths
        θ # vector of values at states in S
    end
"""
mutable struct SimplexValueFunction
    o # position of lower-left corner
    δ # vector of widths
    θ # vector of values at states in S
end

"""
    function (Uθ::SimplexValueFunction)(s)

Estimate value of state vector `s` for known values `θ` over a grid defined by a lower-left vertex `o` and vector of widths `δ`.  Vertices of the grid can all be written `o + δ.*i` for some non-negative integral vector `i`.
"""
function (Uθ::SimplexValueFunction)(s)
    Δ = (s - Uθ.o)./Uθ.δ
    # Multidimensional index of upper-right cell
    i = min.(floor.(Int, Δ) .+ 1, size(Uθ.θ) .- 1) .+ 1
    u = 0.0
    s′ = (s - (Uθ.o + Uθ.δ.*(i.-2))) ./ Uθ.δ
    p = sortperm(s′) # increasing order
    w_tot = 0.0
    for j in p
        w = s′[j] - w_tot
        u += w*Uθ.θ[i...]
        i[j] -= 1
        w_tot += w
    end

    u += (1 - w_tot)*Uθ.θ[i...]
    return u
end

"""
    function fit!(Uθ::SimplexValueFunction, S, U)
"""
function fit!(Uθ::SimplexValueFunction, S, U)
    Uθ.θ = U
    return Uθ
end

# -------------- Linear Regression Value Function ---------------

"""
    mutable struct LinearRegressionValueFunction
        β # basis vector function
        θ # vector of paramters
    end
"""
mutable struct LinearRegressionValueFunction
    β # basis vector function
    θ # vector of paramters
end

"""
    function (Uθ::LinearRegressionValueFunction)(s)
"""
function (Uθ::LinearRegressionValueFunction)(s)
    return Uθ.β(s) ⋅ Uθ.θ
end

"""
    function fit!(Uθ::LinearRegressionValueFunction, S, U)
"""
function fit!(Uθ::LinearRegressionValueFunction, S, U)
    X = hcat([Uθ.β(s) for s in S]...)'
    Uθ.θ = pinv(X)*U
    return Uθ
end


include(joinpath(@__DIR__, "pedagogy_graph.jl"))
using Random

dims = [5, 5] # if > (14, 14) the display becomes all black, don't know why !
periodic = false
H, J = 0.0, 0.01

ising = PRS.Ising(dims, J, H; periodic=periodic)

using Random
rng = Random.MersenneTwister(123)
@time sample = PRS.generate_sample_gibbs_perfect(rng, ising)

path = joinpath(@__DIR__, "output", "ising", "ising_grid_prs.pdf")
p = plot(ising, sample, dims, path; nodelabel=nothing)


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Polytope{K,Dim,T}

We say that a geometry is a K-polytope when it is a collection of "flat" sides
that constitue a `K`-dimensional subspace. They are called polygon and polyhedron
respectively for 2D (`K=2`) and 3D (`K=3`) subspaces, embedded in a `Dim`-dimensional
space. The parameter `K` is also known as the rank or parametric dimension of the
polytope: https://en.wikipedia.org/wiki/Abstract_polytope.

The term polytope expresses a particular combinatorial structure. A polyhedron,
for example, can be decomposed into faces. Each face can then be decomposed into
edges, and edges into vertices. Some conventions act as a mapping between vertices
and higher dimensional features (edges, faces, cells...), removing the need to
store all features.

Additionally, the following property must hold in order for a geometry to be considered
a polytope: the boundary of a (K+1)-polytope is a collection of K-polytopes, which may
have (K-1)-polytopes in common. See https://en.wikipedia.org/wiki/Polytope.

### Notes

- Type aliases are `Polygon`, `Polyhedron`.
"""
abstract type Polytope{K,Dim,T} <: Geometry{Dim,T} end

(::Type{PL})(vertices::Vararg{P}) where {PL<:Polytope,P<:Point} = PL(SVector(vertices))
(::Type{PL})(vertices::AbstractVector{TP}) where {PL<:Polytope,TP<:Tuple} = PL(Point.(vertices))
(::Type{PL})(vertices::Vararg{TP}) where {PL<:Polytope,TP<:Tuple} = PL(collect(vertices))

"""
    paramdim(polytope)

Return the parametric dimension or rank of the polytope.
"""
paramdim(::Type{<:Polytope{K}}) where {K} = K

"""
    vertices(polytope)

Return the vertices of a `polytope`.
"""
vertices(p::Polytope) = p.vertices

"""
    nvertices(polytope)

Return the number of vertices in the `polytope`.
"""
nvertices(p::Polytope) = length(vertices(p))

"""
    p1 == p2

Tells whether or not polytopes `p1` and `p2` are equal.
"""
==(p1::Polytope, p2::Polytope) = vertices(p1) == vertices(p2)

"""
    centroid(polytope)

Return the centroid of the `polytope`.
"""
centroid(p::Polytope) = Point(sum(coordinates, vertices(p)) / length(vertices(p)))

"""
    unique(polytope)

Return a new `polytope` without duplicate vertices.
"""
Base.unique(p::Polytope) = unique!(deepcopy(p))

"""
    unique!(polytope)

Remove duplicate vertices in `polytope`.
"""
function Base.unique!(::Polytope) end

function Base.show(io::IO, p::Polytope)
  kind = prettyname(p)
  vert = join(vertices(p), ", ")
  print(io, "$kind($vert)")
end

function Base.show(io::IO, ::MIME"text/plain", p::Polytope{K,Dim,T}) where {K,Dim,T}
  kind = prettyname(p)
  lines = ["  └─$v" for v in vertices(p)]
  println(io, "$kind{$Dim,$T}")
  print(io, join(lines, "\n"))
end

prettyname(p::Polytope) = prettyname(typeof(p))
function prettyname(PL::Type{<:Polytope})
  n = string(PL)
  i = findfirst('{', n)
  isnothing(i) ? n : n[1:i-1]
end

# --------
# POLYGON
# --------

"""
    Polygon{Dim,T}

A polygon is a 2-polytope, i.e. a polytope with parametric dimension 2.
"""
const Polygon = Polytope{2}

"""
    area(polygon)

Return the area of the `polygon`.
"""
area(p::Polygon) = measure(p)

"""
    chains(polygon)

Return the outer and inner chains of the polygon.
"""
function chains(::Polygon) end

"""
    hasholes(polygon)

Tells whether or not the `polygon` contains holes.
"""
function hasholes(::Polygon) end

"""
    issimple(polygon)

Tells whether or not the `polygon` is simple.
See [https://en.wikipedia.org/wiki/Simple_polygon]
(https://en.wikipedia.org/wiki/Simple_polygon).
"""
issimple(p::Polygon) = issimple(typeof(p))

"""
    windingnumber(point, polygon)

Winding number of `point` with respect to the `polygon`.
"""
function windingnumber(::Point, ::Polygon) end

"""
    orientation(polygon)

Returns the orientation of the rings of the `polygon`
as either counter-clockwise (CCW) or clockwise (CW).
"""
orientation(p::Polygon) = orientation(p, WindingOrientation())

function orientation(p::Polygon, algo)
  o = [orientation(c, algo) for c in chains(p)]
  hasholes(p) ? o : first(o)
end

"""
    boundary(polygon)

Returns the boundary of the `polygon`.
"""
boundary(p::Polygon) = hasholes(p) ? Multi(chains(p)) : first(chains(p))

"""
    isconvex(polygon)

Tells whether or not the `polygon` is convex.
"""
isconvex(p::Polygon{Dim,T}) where {Dim,T} =
  issimple(p) && all(≤(T(π)), innerangles(boundary(p)))

"""
    bridge(polygon; width=0)

Transform `polygon` with holes into a single outer chain
via bridges of given `width` as described in Held 1998.
Return the outer chain and a vector with pairs of indices
for duplicate vertices. These indices can be used to undo
the bridges.

## References

* Held. 1998. [FIST: Fast Industrial-Strength Triangulation of Polygons]
  (https://link.springer.com/article/10.1007/s00453-001-0028-4)
"""
function bridge(p::Polygon{Dim,T}; width=zero(T)) where {Dim,T}
  # polygons without holes are trivial
  if !hasholes(p)
    outerchain = first(chains(p))
    duplicates = Tuple{Int,Int}[]
    return outerchain, duplicates
  end

  # retrieve chains as vectors of coordinates
  pchains = [coordinates.(vertices(open(c))) for c in chains(p)]

  # sort vertices lexicographically
  coords  = [coord for pchain in pchains for coord in pchain]
  indices = sortperm(sortperm(coords))

  # each chain has its own set of indices
  pinds = Vector{Int}[]; offset = 0
  for nvertex in length.(pchains)
    push!(pinds, indices[offset+1:offset+nvertex])
    offset += nvertex
  end

  # sort chains based on leftmost vertex
  leftmost = argmin.(pinds)
  minimums = getindex.(pinds, leftmost)
  reorder  = sortperm(minimums)
  leftmost = leftmost[reorder]
  minimums = minimums[reorder]
  pchains  = pchains[reorder]
  pinds    = pinds[reorder]

  # initialize outer boundary
  outer = first(pchains)
  oinds = first(pinds)

  # merge holes into outer boundary
  for i in 2:length(pchains)
    inner = pchains[i]
    iinds = pinds[i]
    l = leftmost[i]
    m = minimums[i]

    # find closest vertex in boundary
    dmin, jmin = Inf, 0
    for j in findall(oinds .≤ m)
      d = sum(abs, outer[j] - inner[l])
      if d < dmin
        dmin, jmin = d, j
      end
    end

    # create a bridge of given width δ
    # from line segment A--B. The point
    # A is split into A′ and A′′ and the
    # point B is split into B′ and B′′
    A = outer[jmin]
    B = inner[l]
    δ = width
    v = B - A
    u = Vec(-v[2], v[1])
    n = u / norm(u)
    A′  = A + δ/2 * n
    A′′ = A - δ/2 * n
    B′  = B + δ/2 * n
    B′′ = B - δ/2 * n

    # insert hole at closest vertex
    outer = [
      outer[begin:jmin-1]; [A′, B′];
      circshift(inner, -l+1)[2:end];
      [B′′, A′′]; outer[jmin+1:end]
    ]
    oinds = [
      oinds[begin:jmin];
      circshift(iinds, -l+1);
      [iinds[l]];
      oinds[jmin:end]
    ]
  end

  # find duplicate vertices
  duplicates = Tuple{Int,Int}[]
  occurred = Dict{Int,Int}()
  for (i, ind) in enumerate(oinds)
    if haskey(occurred, ind)
      push!(duplicates, (occurred[ind], i))
    else
      occurred[ind] = i
    end
  end

  # close outer boundary
  push!(outer, first(outer))

  outerchain = Chain(Point.(outer))

  outerchain, duplicates
end

# -----------
# POLYHEDRON
# -----------

"""
    Polyhedron{Dim,T}

A polyhedron is a 3-polytope, i.e. a polytope with parametric dimension 3.
"""
const Polyhedron = Polytope{3}

"""
   volume(polyhedron)

Return the volume of the `polyhedron`.
"""
volume(p::Polyhedron) = measure(p)

# ----------------
# IMPLEMENTATIONS
# ----------------

include("polytopes/segment.jl")
include("polytopes/ngon.jl")
include("polytopes/chain.jl")
include("polytopes/polyarea.jl")
include("polytopes/pyramid.jl")
include("polytopes/tetrahedron.jl")
include("polytopes/hexahedron.jl")


# ---
# title: 350. Intersection of Two Arrays II
# id: problem350
# author: zhwang
# date: 2022-01-19
# difficulty: Easy
# categories: Hash Table, Two Pointers, Binary Search, Sort
# link: <https://leetcode.com/problems/intersection-of-two-arrays-ii/description/>
# hidden: true
# ---
#
# Given two arrays, write a function to compute their intersection.
#
# **Example 1:**
#
#
#
#     Input: nums1 = [1,2,2,1], nums2 = [2,2]
#     Output: [2,2]
#
#
# **Example 2:**
#
#
#
#     Input: nums1 = [4,9,5], nums2 = [9,4,9,8,4]
#     Output: [4,9]
#
# **Note:**
#
#   * Each element in the result should appear as many times as it shows in both arrays.
#   * The result can be in any order.
#
# **Follow up:**
#
#   * What if the given array is already sorted? How would you optimize your algorithm?
#   * What if _nums1_ 's size is small compared to _nums2_ 's size? Which algorithm is better?
#   * What if elements of _nums2_ are stored on disk, and the memory is limited such that you cannot load all elements into the memory at once?
#
#
## @lc code=start
using LeetCode

function intersection_of_arrays(nums1::Vector{Int}, nums2::Vector{Int})::Vector{Int}
    nums1, nums2 = sort(nums1), sort(nums2)
    n1, n2 = length(nums1), length(nums2)
    i1, i2, res = 1, 1, Int[]
    while i1 <= n1 && i2 <= n2
        if nums1[i1] > nums2[i2]
            i2 += 1
        elseif nums1[i1] < nums2[i2]
            i1 += 1
        else
            push!(res, nums1[i1])
            i1 += 1
            i2 += 1
        end
    end
    return res
end

## @lc code=end


module EvalMetrics

import Base: show, precision
import DocStringExtensions: SIGNATURES
import Statistics: quantile
import StatsBase: RealVector
using RecipesBase
using Reexport

include("encodings/Encodings.jl")
@reexport using .Encodings

include("utilities.jl")
include("confusion_matrix.jl")

const CMVector{T<:Real} = AbstractVector{ConfusionMatrix{T}}

include("metrics.jl")
include("thresholds.jl")
include("curves.jl")

export
    # confusion matrix
    ConfusionMatrix,

    # performance metrics
    true_positive,
    true_negative,
    false_positive,
    false_negative,
    true_positive_rate,
        sensitivity,
        recall,
        hit_rate,
    true_negative_rate,
        specificity,
        selectivity,
    false_positive_rate,
        fall_out,
        type_I_error,
    false_negative_rate,
        miss_rate,
        type_II_error,
    precision,
        positive_predictive_value,
    negative_predictive_value,
    false_discovery_rate,
    false_omission_rate,
    threat_score,
        critical_success_index,
    accuracy,
    balanced_accuracy,
    error_rate,
    balanced_error_rate,
    f1_score,
    fβ_score,
    matthews_correlation_coefficient,
        mcc,
    quant,
    topquant,
    positive_likelihood_ratio,
    negative_likelihood_ratio,
    diagnostic_odds_ratio,
    prevalence,

    # threshold functions
    thresholds,
    threshold_at_tpr,
    threshold_at_tnr,
    threshold_at_fpr,
    threshold_at_fnr,
    threshold_at_k,

    #curves
    auc_trapezoidal,
    curve, auc,
    PRCurve, prcurve, au_prcurve, prplot,
    ROCCurve, roccurve, au_roccurve, rocplot,

    # utilities
    binary_eval_report,
    mergesorted

binary_eval_report(targets::AbstractVector, scores::RealVector, fpr = 0.05) =
    binary_eval_report(current_encoding(), targets, scores, fpr)

function binary_eval_report(enc::TwoClassEncoding, targets::AbstractVector, scores::RealVector, fpr = 0.05)
    t = threshold_at_fpr(enc, targets, scores, fpr)
    c = ConfusionMatrix(enc, targets, scores, t)

    return Dict(
        "accuracy@fpr$(fpr)" => accuracy(c),
        "au_prcurve" => au_prcurve(enc, targets, scores),
        "au_roccurve" => au_roccurve(enc, targets, scores),
        "precision@fpr$(fpr)" => precision(c),
        "prevalence" => prevalence(c),
        "recall@fpr$(fpr)" => recall(c),
        "samples" => length(targets),
        "true negative rate@fpr$(fpr)" => true_negative_rate(c)
    )
end

end


function determine_qdim(equation::Union{SDE,PSDE,SPSDE})
    nd = equation.d
    ns = equation.ns
    ni = equation.n

    if nd==ns==ni==1
        NQ = 1
    elseif ns==ni==1
        NQ = 2
    elseif ns==1 || ni==1
        NQ = 3
    else
        NQ = 4
    end

    return NQ
end


function hamiltonian!(ds, s, J::AbstractMatrix, h::AbstractVector, anisotropy, t)
    for i in 1:size(s, 1) # iterate over all spins i
        js_x, js_y, js_z = @uviews J s begin
            js_x = 0.5 * anisotropy[1] * get_mean_field(J, s, i, 1) + h[1]
            js_y = 0.5 * anisotropy[2] * get_mean_field(J, s, i, 2) + h[2]
            js_z = 0.5 * anisotropy[3] * get_mean_field(J, s, i, 3) + h[3]
            js_x, js_y, js_z
        end

        ds[i, 1] = s[i, 2] * js_z - s[i, 3] * js_y
        ds[i, 2] = s[i, 3] * js_x - s[i, 1] * js_z
        ds[i, 3] = s[i, 1] * js_y - s[i, 2] * js_x
    end
    nothing
end

function hamiltonian!(ds, s, J::AbstractMatrix, h::AbstractMatrix, anisotropy, t)
    for i in 1:size(s, 1) # iterate over all spins i
        js_x, js_y, js_z = @uviews J s begin
            js_x = anisotropy[1] * get_mean_field(J, s, i, 1) + h[i, 1]
            js_y = anisotropy[2] * get_mean_field(J, s, i, 2) + h[i, 2]
            js_z = anisotropy[3] * get_mean_field(J, s, i, 3) + h[i, 3]
            js_x, js_y, js_z
        end

        ds[i, 1] = s[i, 2] * js_z - s[i, 3] * js_y
        ds[i, 2] = s[i, 3] * js_x - s[i, 1] * js_z
        ds[i, 3] = s[i, 1] * js_y - s[i, 2] * js_x
    end
    nothing
end

get_mean_field(J_i::AbstractVector{T}, s::AbstractVector{T}) where T = J_i ⋅ s

function get_mean_field(J::AbstractMatrix{T}, s::AbstractMatrix{T}, i, α) where T
    get_mean_field(view(J, :, i), view(s, :, α))
end

function get_mean_field_unsafe(J::AbstractMatrix{T}, s::AbstractMatrix{T}, i, α) where T
    s_α = uview(s, :, α)
    J_i = uview(J, :, i)
    get_mean_field(J_i, s_α)
end


using MonteCarloX
using Random
using Distributions
using StatsBase
using Printf
using ProgressMeter
using LinearAlgebra
using DelimitedFiles
using StaticArrays

"""
Reproduce results similar to Fig.2 of Lloyed-Smith et al., Nature (2005) on the
the probability of exctinction (here after 3 weeks for mu=1/8) and the
occurrance of the first outbreak (here first day with >50 cases per trajectory).
"""
function reproduce_lloyed_smith_nature_2005_fig2(path)
    list_ratio_lambda_mu = collect(0.5:0.1:3)
    list_k = [10, 1, 1e-1, 1e-2]
    list_a = [3, 5, 10]
    N = Int(1e5)
    I0 = 1
    R0 = 0
    S0 = N - I0
    epsilon=0
    range_seed = 1:Int(1e6);

    # motivated by COVID-19 parameters [cf. Dehning et al, Science (2020)]
    mu = 1/8. # avg. infectious time of 5+3 days
    time_meas = 50 # days of simulation
    dT = 1

    for ratio_lambda_mu in list_ratio_lambda_mu
        lambda_0 = ratio_lambda_mu*mu
        list_P = []
        list_P_name = []
        for k in list_k
            P_Gamma = Gamma(k, lambda_0/k)
            push!(list_P, P_Gamma)
            push!(list_P_name, @sprintf("Gamma_k%.2e",k))
        end
        for a in list_a
            lambda_1 = mu/a
            lambda_2 = mu*a
            prior_1 = (lambda_0 - lambda_2)/(lambda_1 - lambda_2)
            prior_2 = 1 - prior_1
            if lambda_2 >= lambda_0
                P_Bimodal = MixtureModel(Normal[Normal(lambda_1, 0), Normal(lambda_2, 0)], [prior_1, prior_2])
                push!(list_P, P_Bimodal)
                push!(list_P_name, @sprintf("Bimodal_a%.2e",a))
            end
        end

        # simulations
        for (P,P_name) in zip(list_P, list_P_name)
            println("R=", ratio_lambda_mu, " P=", P_name)

            list_time = collect(range(0, time_meas, step=dT))
            trajectory_S = zeros(Int,length(list_time))
            trajectory_I = zeros(Int,length(list_time))
            trajectory_R = zeros(Int,length(list_time))

            # observables
            probability_exctinction_per_day = zeros(length(list_time))
            probability_first_day_more_than_50 = zeros(length(list_time))

            #generate and analyse trajectories
            @showprogress 1 for seed in range_seed
                rng = MersenneTwister(seed);
                system = SIR{typeof(P)}(rng, P, epsilon, mu, S0, I0, R0)
                trajectory!(rng, list_time, trajectory_S, trajectory_I, trajectory_R, system)

                first_index_more_than_50 = findfirst(x->x>50, trajectory_I)
                if !(typeof(first_index_more_than_50) == Nothing)
                    probability_first_day_more_than_50[first_index_more_than_50] += 1
                end

                probability_exctinction_per_day[trajectory_I .== 0] .+= 1
            end

            probability_exctinction_per_day ./= length(range_seed)
            probability_first_day_more_than_50 ./= length(range_seed)

            #write out
            filename =  @sprintf("%s/reproduce_lloyed_smith_nature_2005_fig2_%s_R%.2e_mu%.2e_N%.2e_I0%.2e.dat",
                                 path, P_name, ratio_lambda_mu, mu, N, I0)

            open(filename; write=true) do f
                write(f, "#P_ext(t) - probability of trajectory to be extinct at this day\n")
                write(f, "#P_50(t)  - probability that trajectory has more than 50 infections for the first time on this day\n")
                write(f, "#t\t P_ext(t)\t N_50\n")
                writedlm(f, [list_time  probability_exctinction_per_day probability_first_day_more_than_50])
            end
        end
    end
end

function example_dist_sir_hetero_rates(path;
                                       ratio_lambda_mu=1.0)
    # define default values (some specific initial conditions, I0=100)
    N = Int(1e5)
    list_I0 = Int[1,10,100]
    seeds = collect(Int, 1:1e6);

    # motivated by COVID-19 parameters [cf. Dehning et al, Science (2020)]
    mu = 1/8. # avg. infectious time of 5+3 days
    lambda_0 = ratio_lambda_mu*mu # focus on the "critical point" of sustained activity
    sigma = mu/10
    time_meas = 21 # 1 week observation of effect of initial perturbation
    list_k = [1,1e-1,1e-2]
    list_alpha = [1.5,2,3] # mean only defined for alpha > 1

    # Different distributions with same mean speading rate lambda_0
    list_P = []
    list_P_name = []
    # Compared to [Lloyd-Smith et al., Science (2020)] delta-distribution for
    # differential-equation model corresponds already to an exponential
    # distribution of R_0 i.e. a geometric offspring distribution (k=1)
    P_Delta = Normal(lambda_0, 0.0)
    push!(list_P, P_Delta)
    push!(list_P_name, "Delta")
    P_Gaussian = Normal(lambda_0, sigma)
    push!(list_P, P_Gaussian)
    push!(list_P_name, "Gaussian")
    # We cannot fully reproduce the negative-binomial offspring distribution by
    # a Gamma-Poisson mixture because we have no generation based model here.
    # Still, we should approximate the same behavior by choosing a Gamma
    # distribution of spreading rates with the same mean
    # https://en.wikipedia.org/wiki/Gamma_distribution
    #   mean(Gamma) = k*theta = lambda_0 i.e. theta = lambda_0/k
    for k in list_k
        P_Gamma = Gamma(k, lambda_0/k)
        push!(list_P, P_Gamma)
        push!(list_P_name, @sprintf("Gamma_k%.2e",k))
    end
    # Pareto distribution for scale-free spreading rates (motivated from
    # scale-free-ness of social networks)
    # P_Pareto(x_m, alpha) = GeneralizedPareto(x_m, x_m/alpha, 1/alpha)
    # for Pareto to give mean=lambda_0 : x_m(lambda_0,alpha) = (alpha-1)lambda_0/alpha
    for alpha in list_alpha
        x_m = (alpha-1)*lambda_0/alpha
        P_Pareto = GeneralizedPareto(x_m, x_m/alpha, 1/alpha)
        push!(list_P, P_Pareto)
        push!(list_P_name, @sprintf("Pareto_alpha%.2e",alpha))
    end
    # mean(lambda) = lambda_1*prior_1 + lambda_2*prior_2 = lambda_0
    # prior_1 + prior_2 = 1
    lambda_1 =  0.0
    lambda_2 = 10.0*mu
    prior_1 = (lambda_0 - lambda_2)/(lambda_1 - lambda_2)
    prior_2 = 1 - prior_1
    P_Bimodal = MixtureModel(Normal[Normal(lambda_1, 0), Normal(lambda_2, 0)], [prior_1, prior_2])
    push!(list_P, P_Bimodal)
    push!(list_P_name, @sprintf("Bimodal_left%.2e_right%.2e",lambda_1, lambda_2))
    # simulations
    for I0 in list_I0
        R0 = 0
        S0 = N-I0
        for (P,P_name) in zip(list_P, list_P_name)
            println(N, " ", I0, " ", P_name)
            list_time, list_dist = distribution(P, S0, I0, R0, time_meas, seeds, mu=mu);
            filename_base =  @sprintf("%s/distribution_%s_lambda%.2e_mu%.2e_N%.2e_I0%.2e",
                                 path, P_name, lambda_0, mu, N, I0)
            write_distributions(list_time, list_dist, filename_base)
        end
    end
end

function write_distributions(list_time, list_dist, filename_base)
    for (time,dist) in zip(list_time, list_dist)
        filename = @sprintf("%s_T%.2e.dat",filename_base,time)
        open(filename; write=true) do f
            write(f, "#I\t P(I)\n")
            writedlm(f, zip(collect(dist.edges[1]),dist.weights))
        end
    end
end

function sir_hetero_rates(P_lambda::D, S0, I0, R0, time_meas, seed;
                          mu=1.0, epsilon=0.0, dT=1) where D
    rng = MersenneTwister(seed);
    system = SIR{D}(rng, P_lambda, epsilon, mu, S0, I0, R0)

    list_T = collect(range(0, time_meas, step=dT))
    list_S = zeros(length(list_T));
    list_I = zeros(length(list_T));
    list_R = zeros(length(list_T));
    trajectory!(rng, list_T, list_S, list_I, list_R, system)

    return list_T, list_S, list_I, list_R, system
end

# memory-wise this is not the most efficient solution
# Alternative is to have one trajectory readout and directly add to
# distribution (which requires some a priori size though)
function distribution(P_lambda::D, S0, I0, R0, time_meas, seeds;
                      mu=1.0, epsilon=0, dT=1) where D

    list_T = collect(range(0, time_meas, step=dT))
    trajectories_S = [zeros(Int,length(list_T)) for i=1:length(seeds)]
    trajectories_I = [zeros(Int,length(list_T)) for i=1:length(seeds)]
    trajectories_R = [zeros(Int,length(list_T)) for i=1:length(seeds)]
    @showprogress 1 for (i,seed) in enumerate(seeds)
        rng = MersenneTwister(seed);
        system = SIR{D}(rng, P_lambda, epsilon, mu, S0, I0, R0)
        trajectory!(rng, list_T, trajectories_S[i], trajectories_I[i], trajectories_R[i], system)
    end
    #index [time,trajectory] -> a way to make this [trajectory,time] to use later trajectories[:,t]
    trajectories_S = hcat(trajectories_S...);
    trajectories_I = hcat(trajectories_I...);
    trajectories_R = hcat(trajectories_R...);

    # then fit distributions (for now only I)
    # TODO: generalize (function for this) to all cases
    list_dist = []
    for t=1:size(trajectories_I,1)
        min_I = 0
        max_I = maximum(trajectories_I[t,:])+2
        step = 1
        hist = fit(Histogram, trajectories_I[t,:], min_I:step:max_I)
        dist = normalize!(float(hist))
        push!(list_dist, dist)
    end

    return list_T, list_dist
end

"""
# generate a stochastic trajectory
SIR dynamics is goverened by differential equation
``
    \\frac{dI}{dt} = \\sum_{i=1}^{I}\\lambda_i\\frac{S}{N} - \\mu I
``

# Arguments
* distribution for the heterogeneous rates P(lambda)
"""
function trajectory!(rng, list_T, list_S, list_I, list_R, system)
    rates = current_rates(system)
    pass_update!(rates, index) = update!(rates, index, system, rng)

    list_S[1] = system.measure_S
    list_I[1] = system.measure_I
    list_R[1] = system.measure_R
    time_simulation = Float64(list_T[1])
    for i in 2:length(list_T)
        if time_simulation < list_T[i]
            dT = list_T[i] - time_simulation
            dT_sim = advance!(KineticMonteCarlo(), rng, rates, pass_update!, dT)
            time_simulation += dT_sim
        end
        list_S[i] = system.measure_S
        list_I[i] = system.measure_I
        list_R[i] = system.measure_R
    end
end

mutable struct SIR{D}
    epsilon::Float64
    mu::Float64
    P_lambda::D
    current_lambda::Vector{Float64}
    sum_current_lambda::Float64
    update_current_lambda::Int
    N::Int
    S::Int
    I::Int
    R::Int
    measure_S::Int
    measure_I::Int
    measure_R::Int

    function SIR{D}(rng::AbstractRNG, P_lambda::D, epsilon, mu, S0::Int, I0::Int, R0::Int) where D
        N = S0 + I0 + R0
        current_lambda = rand(rng, P_lambda, I0)
        sum_current_lambda = sum(current_lambda)
        new(epsilon, mu, P_lambda, current_lambda, sum_current_lambda, 0, N, S0, I0, R0, S0, I0, R0)
    end
end

"""
update rates and system according to the last event that happend:
recovery:  1
infection: 2
"""
function update!(rates::AbstractVector, index::Int, system::SIR, rng::AbstractRNG)
    system.measure_S = system.S
    system.measure_I = system.I
    system.measure_R = system.R
    if index == 1 # recovery
        random_I = rand(rng, 1:system.I)
        delete_from_system!(system, random_I)
        system.I -= 1
        system.R += 1
    elseif index == 2 # infection
        add_to_system!(system, rand(rng, system.P_lambda))
        system.S -= 1
        system.I += 1
        @assert length(system.current_lambda) == system.I
    elseif index == 0 # absorbing state of zero infected
        system.measure_S = system.S
        system.measure_I = system.I
        system.measure_R = system.R
    else
        throw(UndefVarError(:index))
    end

    rates .= current_rates(system)
end

function delete_from_system!(system, index)
    update_sum!(system, -system.current_lambda[index])
    deleteat!(system.current_lambda, index)
end

function add_to_system!(system, lambda)
    push!(system.current_lambda, lambda)
    update_sum!(system, lambda)
end

function update_sum!(system, change)
    system.sum_current_lambda += change
    system.update_current_lambda += 1
    if system.update_current_lambda > 1e5
        system.sum_current_lambda = sum(system.current_lambda)
        system.update_current_lambda = 0
    end
end

"""
evaluate current rates of SIR system
"""
function current_rates(system::SIR)
    rate_recovery = system.mu * system.I
    rate_infection = system.sum_current_lambda* system.S/system.N + system.epsilon
    return MVector{2,Float64}(rate_recovery, rate_infection)
end