using Statistics
using DelimitedFiles
using Random
using Printf

struct Lattice
    nsites::Int64
    nup::Int64
    ndn::Int64
    ntot::Int64
    t::Float64
    U::Float64
end

struct QMC
    nblocks::Int64
    samples::Int64
    Δτ::Float64
    L::Int64
    γ::Float64
    Kmatrix::Array{Float64,2}
end

struct Opt  # Optimization
    stab_interval::Int64
    nburn_in::Int64
end

const FloatType = Union{Float64, Complex{Float64}}

include("./MatrixGenerator.jl")
include("./Calculations.jl")

function getinputs()
    lines = readlines("./InputFermiHubbard")

    # basic parameters
    β = parse(Float64, split(lines[1])[2])
    num_sites = parse(Int64, split(lines[2])[2])
    num_spinup = parse(Int64, split(lines[3])[2])
    num_spindn = parse(Int64, split(lines[4])[2])
    num_elec = num_spinup + num_spinup

    Δτ = parse(Float64, split(lines[5])[2])
    t = parse(Float64, split(lines[6])[2])
    U = parse(Float64, split(lines[7])[2])
    num_blocks = convert(Int64, parse(Float64, split(lines[8])[2]))
    num_samples = convert(Int64, parse(Float64, split(lines[9])[2]))

    # optimization parameters
    stab_interval = parse(Int64, split(lines[12])[2])
    num_burnins = convert(Int64, parse(Float64, split(lines[13])[2]))

    # constants for calculations
    L = convert(Int64, round(β / Δτ))
    γ = atanh(sqrt(tanh(Δτ * U / 4)))

    lattice = Lattice(num_sites, num_spinup, num_spindn, num_elec, t, U)

    K_matrix = kinetic_matrix(lattice)
    qmc = QMC(num_blocks, num_samples, Δτ, L, γ, exp(-K_matrix * Δτ / 2))

    opt = Opt(stab_interval, num_burnins)
    return lattice, qmc, opt
end

function mainbody()
    lattice, qmc, opt = getinputs()
    SubEk_array = Vector{Float64}()
    SubEp_array = Vector{Float64}()
    SubE_array = Vector{Float64}()
    SubEk_error_array = Vector{Float64}()
    SubEp_error_array = Vector{Float64}()
    SubE_error_array = Vector{Float64}()
    filename = string("output_Ns", lattice.nsites, "_U", lattice.U, "_beta", qmc.Δτ * qmc.L)
    outputfile = open(filename, "a+")
    @printf(outputfile, "       Canonical Ensemble AFQMC for Hubbard Model     \n\n")
    @printf(outputfile, "Number of Blocks: %d\n", qmc.nblocks)
    @printf(outputfile, "Length of Img Timestep: %.4f\n", qmc.Δτ)
    @printf(outputfile, "Stablization Interval: %d\n\n", opt.stab_interval)
    for i = 1 : qmc.nblocks
        SubE, SubE_error, SubEk, SubEk_error, SubEp, SubEp_error, sgn_mean, NiMax = MCIntegration(lattice, qmc, opt)
        SubE = SubE / sgn_mean
        SubEk = SubEk / sgn_mean
        SubEp = SubEp / sgn_mean

        @printf(outputfile, "Block # %d out of %d\n", i, qmc.nblocks)
        @printf(outputfile, "--------------------\n")
        @printf(outputfile, "Average Sign: %.2f\n", sgn_mean)
        @printf(outputfile, "Maximum Occupancy: %.2f\n", NiMax)
        @printf(outputfile, "Ek            Ep            Etot\n")
        @printf(outputfile, "%.5f     %.5f       %.5f\n", SubEk, SubEp, SubE)
        @printf(outputfile, "Error:\n")
        @printf(outputfile, "%.5f     %.5f       %.5f\n\n", SubEk_error, SubEp_error, SubE_error)

        push!(SubEk_array, SubEk)
        push!(SubEk_error_array, SubEk_error)
        push!(SubEp_array, SubEp)
        push!(SubEp_error_array, SubEp_error)
        push!(SubE_array, SubE)
        push!(SubE_error_array, SubE_error)
    end
    Ek_mean, Ek_error = mean(SubEk_array), mean(SubEk_error_array)
    Ep_mean, Ep_error = mean(SubEp_array), mean(SubEp_error_array)
    E_mean, E_error = mean(SubE_array), mean(SubE_error_array)

    @printf(outputfile, "--------------------\n")
    @printf(outputfile, "Results for this run:\n")
    @printf(outputfile, "Ek            Ep            Etot\n")
    @printf(outputfile, "%.5f     %.5f       %.5f\n", Ek_mean, Ep_mean, E_mean)
    @printf(outputfile, "Error:\n")
    @printf(outputfile, "%.5f     %.5f       %.5f\n\n", Ek_error, Ep_error, E_error)
    close(outputfile)
end

mainbody()


using QuantumInfo, RandomQuantum
using LinearAlgebra: svdvals, norm, ishermitian, eigvals
using StatsBase: var

using Test

trnorm(A) = sum(svdvals(A))

@testset "Randomized GinUE tests" begin
for _ in 1:10
    rows = rand(2:10)
    cols = rand(1:10)
    X = rand(GinibreEnsemble(rows,cols))
    @test var(vec(real(X))) > 0
    @test var(vec(imag(X))) > 0
end
end

@testset "Randomized FubiniStudyEnsemble tests" begin
for _ in 1:10
    d = rand(2:10)
    ψ = rand(FubiniStudyPureState(d))
    @test var(real(ψ)) > 0
    @test var(imag(ψ)) > 0
    @test isapprox(norm(ψ), 1.0)
end
end

@testset "Randomized FSOE tests" begin
for _ in 1:10
    d = rand(2:10)
    db = rand(2:10)
    ρ = rand(FubiniStudyMixedState(d,db))
    @test isapprox(trnorm(ρ), 1)
    @test ishermitian(ρ)
    @test ispossemidef(ρ)
    @test var(vec(real(ρ))) > 0
    @test var(vec(imag(ρ))) > 0
end
end

@testset "Randomized HSOE tests" begin
for _ in 1:10
    d = rand(2:10)
    ρ = rand(HilbertSchmidtMixedState(d))
    @test isapprox(trnorm(ρ), 1)
    @test ishermitian(ρ)
    @test ispossemidef(ρ)
    @test var(vec(real(ρ))) > 0
    @test var(vec(imag(ρ))) > 0
end
end

@testset "BuresMixedState" begin
for _ in 1:10
    d = rand(2:10)
    ρ = rand(BuresMixedState(d))
    @test isapprox(trnorm(ρ), 1)
    @test ishermitian(ρ)
    @test ispossemidef(ρ)
    @test var(vec(real(ρ))) > 0
    @test var(vec(imag(ρ))) > 0
end
end

@testset "ClosedHaarEnsemble" begin
for _ in 1:10
    d = rand(2:10)
    U = rand(ClosedHaarEnsemble(d))
    @test isapprox(U*U', QuantumInfo.eye(d), atol=1e-12)
    @test isapprox(U'*U, QuantumInfo.eye(d), atol=1e-12)
    @test isapprox(abs.(eigvals(U)),  ones(d), atol=1e-12)
end
end

@testset "OpenHaarEnsemble" begin
for _ in 1:10
    d = 2 # rand(2:10)
    db = rand(2:10)
    E = rand(OpenHaarEnsemble(d,db))
    @test istp(E,tol=1e-14)
    @test iscp(E,tol=1e-14)
    @test ischannel(E,tol=1e-14)
end
end

@testset "ClosedEvolution" begin
for _ in 1:10
    α = rand()
    d = rand(2:10)
    U = rand(RandomClosedEvolution(d,α))
    @test isapprox(U*U', QuantumInfo.eye(d), atol=1e-12)
    @test isapprox(U'*U, QuantumInfo.eye(d), atol=1e-12)
    @test isapprox(abs.(eigvals(U)), ones(d), atol=1e-12)
end
end

@testset "OpenEvolution" begin
for _ in 1:10
    α = rand()
    d = rand(2:10)
    db = rand(2:10)
    E = rand(RandomOpenEvolution(d,db,α))
    @test ischannel(E,tol=1e-12)
end
end


### Model
m = Model805()
TTT, RRR, CCC = solve(m)
meas = measurement(m, TTT, RRR, CCC)

nothing


"""
    log_sum_exp(z::Array)

A stable implementation f(x) = log ∘ sum ∘ exp (x).
Since exponentiation can lead to very large numbers.
"""
log_sum_exp(z) = log_sum_exp(z, maximum(z, dims = 1))
log_sum_exp(z, m) = log.(sum(exp.(z .- m), dims = 1)) .+ m


using HypothesisTests
using Distances
using DelayEmbeddings

function normalise_minmax(x, vmin, vmax)
    (x - vmin)/(vmax - vmin)
end

"""
    jdd(source, target; distance_metric = SqEuclidean(),
        B::Int = 10, D::Int = 2, τ::Int = 1) → Vector{Float64}

Compute the joint distance distribution [^1] from `source` to `target` using
the provided `distance_metric`, with `B` controlling the number of subintervals,
`D` the embedding dimension and `τ` the embedding lag.

## Example

```julia
using CausalityTools
x, y = rand(1000), rand(1000)

jdd(x, y)
```

## Keyword arguments

- **`distance_metric::Metric`**: An instance of a valid distance metric from `Distances.jl`.
    Defaults to `SqEuclidean()`.
- **`B`**: The number of equidistant subintervals to divide the interval `[0, 1]` into
    when comparing the normalised distances.
- **`D`**: Embedding dimension.
- **`τ`**: Embedding delay.

## References
[^1] Amigó, José M., and Yoshito Hirata. "Detecting directional couplings from multivariate flows by the joint distance distribution." Chaos: An Interdisciplinary Journal of Nonlinear Science 28.7 (2018): 075302.
"""
function jdd(source, target;
        distance_metric = Euclidean(),
        B::Int = 10,
        D::Int = 2,
        τ::Int = 1)

    js = ([1 for i = 1:D]...,)
    τs = (collect(0:-τ:-(D-1)*τ)...,)
    Ex = DelayEmbeddings.genembed(source, τs, js)
    Ey = DelayEmbeddings.genembed(target, τs, js)
    Mx = DelayEmbeddings.Matrix(Ex)
    My = DelayEmbeddings.Matrix(Ey)

    npts = length(Ex)
    Dx = pairwise(distance_metric, Mx, Mx, dims = 1)
    Dy = pairwise(distance_metric, My, My, dims = 1)

    # Normalise the distances to the interval [0, 1]
    Dx_min = minimum(Dx[Dx .> 0])
    Dy_min = minimum(Dy[Dy .> 0])
    Dx_max = maximum(Dx[Dx .> 0])
    Dy_max = maximum(Dy[Dy .> 0])

    Dx_norm = zeros(Float64, size(Dx))
    Dy_norm = zeros(Float64, size(Dy))

    for i in LinearIndices(Dx[Dx .> 0])
        Dx_norm[i] = normalise_minmax(Dx[i], Dx_min, Dx_max)
    end

    for i in LinearIndices(Dy[Dy .> 0])
        Dy_norm[i] = normalise_minmax(Dy[i], Dy_min, Dy_max)
    end

    mins_δ_yi_yj = fill(2.0, 2*B)

    for (k, b) in enumerate(1:2*B)
        bmin = (b-1)/(2*B)
        bmax = b/(2*B)

        # Find the indices of all pairs (i, j) in Dx whose distances fall inside the interval Ib
        #idxs_Dxs_in_Ib = findall(Dx[])

        # We don't need to store any indices or distances explicitly, but only
        # keep track of whether a smaller distance has has been detected.
        # The maximum possible distance after normalisation is 1.0, so this
        # value can only decrease as we update.
        min_δ_yi_yj = 1.0

        for i = 1:npts
            for j = (i+1):npts
                δ_xi_xj = Dx_norm[i, j]
                if bmin < δ_xi_xj <= bmax
                    δ_yi_yj = Dy_norm[i, j]
                    if δ_yi_yj < min_δ_yi_yj
                        min_δ_yi_yj = δ_yi_yj
                    end
                end
            end
        end
        mins_δ_yi_yj[k] = min_δ_yi_yj
    end

    Δjdd = [mins_δ_yi_yj[B + i] - mins_δ_yi_yj[i] for i in 1:B]

    return Δjdd
end

"""
    jdd(test::OneSampleTTest, source, target;
        distance_metric = SqEuclidean(), B::Int = 10, D::Int = 2, τ::Int = 1,
        μ0 = 0.0) → OneSampleTTest

Perform a one sample t-test to check that the joint distance distribution[^Amigo]
computed from `source` to `target` is biased towards positive values, using the null
hypothesis that the mean of the distribution is `μ0`.

The interpretation of the t-test is that if we can reject the null, then the
joint distance distribution is biased towards positive values, and then there
exists an underlying coupling from `source` to `target`.

## Example

```julia
using CausalityTools, HypothesisTests
x, y = rand(1000), rand(1000)

jdd(OneSampleTTest, x, y)
```

which gives

```julia
One sample t-test
-----------------
Population details:
    parameter of interest:   Mean
    value under h_0:         0.0
    point estimate:          0.06361857324022721
    95% confidence interval: (0.0185, 0.1087)

Test summary:
    outcome with 95% confidence: reject h_0
    two-sided p-value:           0.0082

Details:
    number of observations:   20
    t-statistic:              2.9517208721082873
    degrees of freedom:       19
    empirical standard error: 0.0215530451545668
```

The lower bound of the confidence interval for the mean of the joint
distance distribution is `0.0185` at confidence level `α = 0.05`. The
meaning that the test falsely detected causality from `x` to `y`
between these two random time series.

To get the confidence intervals at confidence level `α`, run

```confinf(jdd(OneSampleTTest, x, y), α)```.

If you just want the p-value at 95% confidence, run

```pvalue(jdd(OneSampleTTest, x, y), tail = :right)```

## Keyword arguments

- **`distance_metric::Metric`**: An instance of a valid distance metric from `Distances.jl`.
    Defaults to `SqEuclidean()`.
- **`B`**: The number of equidistant subintervals to divide the interval `[0, 1]` into
    when comparing the normalised distances.
- **`D`**: Embedding dimension.
- **`τ`**: Embedding delay.
- **`μ0`**: The hypothetical mean value of the joint distance distribution if there
    is no coupling between `x` and `y` (default is `μ0 = 0.0`).

## References
[^Amigo]: Amigó, José M., and Yoshito Hirata. "Detecting directional couplings from multivariate flows by the joint distance distribution." Chaos: An Interdisciplinary Journal of Nonlinear Science 28.7 (2018): 075302.
"""
function jdd(test::Type{OneSampleTTest}, source, target;
        distance_metric = SqEuclidean(),
        B::Int = 10,
        D::Int = 2, τ::Int = 1,
        μ0 = 0.0)

    Δjdd = jdd(source, target,
        distance_metric = distance_metric,
        B = B,
        D = D, τ = τ)

    OneSampleTTest(Δjdd, μ0)
end

export jdd


"""
    BaB(optimizer, ϵ::Float64)

BaB uses branch and bound to estimate the range of the output node.

# Problem requirement
1. Network: any depth, ReLU activation, single output
2. Input: hyperrectangle
3. Output: hyperrectangle (1d interval)

# Return
`CounterExampleResult` or `ReachabilityResult`

# Method
Branch and bound.
For branch, it uses iterative interval refinement.
For bound, it computes concrete bounds by sampling, approximated bound by optimization.
- `optimizer` default `GLPKSolverMIP()`
- `ϵ` is the desired accurancy for termination, default `0.1`.

# Property
Sound and complete.

# Reference
[R. Bunel, I. Turkaslan, P. H. Torr, P. Kohli, and M. P. Kumar,
"A Unified View of Piecewise Linear Neural Network Verification,"
*ArXiv Preprint ArXiv:1711.00455*, 2017.](https://arxiv.org/abs/1711.00455)
"""
@with_kw struct BaB
    optimizer::AbstractMathProgSolver = GLPKSolverMIP()
    ϵ::Float64                        = 0.1
end

function solve(solver::BaB, problem::Problem)
    (u_approx, u, x_u) = output_bound(solver, problem, :max)
    (l_approx, l, x_l) = output_bound(solver, problem, :min)
    bound = Hyperrectangle(low = [l], high = [u])
    reach = Hyperrectangle(low = [l_approx], high = [u_approx])
    return interpret_result(reach, bound, problem.output, x_l, x_u)
end

function interpret_result(reach, bound, output, x_l, x_u)
    if high(reach) < high(output) && low(reach) > low(output)
        return ReachabilityResult(:SAT, reach)
    end
    high(bound) > high(output)    && return CounterExampleResult(:UNSAT, x_u)
    low(bound) < low(output)      && return CounterExampleResult(:UNSAT, x_l)
    return ReachabilityResult(:Unknown, reach)
end

function output_bound(solver::BaB, problem::Problem, type::Symbol)
    nnet = problem.network
    global_concrete, x_star = concrete_bound(nnet, problem.input, type)
    global_approx = approx_bound(nnet, problem.input, solver.optimizer, type)
    doms = Tuple{Float64, Hyperrectangle}[(global_approx, problem.input)]
    index = ifelse(type == :max, 1, -1)
    while index * (global_approx - global_concrete) > solver.ϵ
        dom = pick_out(doms) # pick_out implements the search strategy
        subdoms = split_dom(dom[2]) # split implements the branching rule
        for i in 1:length(subdoms)
            dom_concrete, x = concrete_bound(nnet, subdoms[i], type) # Sample
            dom_approx = approx_bound(nnet, subdoms[i], solver.optimizer, type)
            if index * (dom_concrete - global_concrete) > 0
                (global_concrete, x_star) = (dom_concrete, x)
            end
            if index * (dom_approx - global_concrete) > 0
                add_domain!(doms, (dom_approx, subdoms[i]), type)
            end
        end
        global_approx = doms[1][1]
    end
    return (global_approx, global_concrete, x_star)
end

# Always pick the first dom
function pick_out(doms)
    return doms[1]
end

function add_domain!(doms::Vector{Tuple{Float64, Hyperrectangle}}, new::Tuple{Float64, Hyperrectangle}, type::Symbol)
    rank = length(doms) + 1
    index = ifelse(type == :max, 1, -1)
    for i in 1:length(doms)
        if index * (new[1] - doms[i][1]) >= 0
            rank = i
            break
        end
    end
    insert!(doms, rank, new)
end

# Always split the longest input dimention
function split_dom(dom::Hyperrectangle)
    max_value, index_to_split = findmax(dom.radius)
    return split_interval(dom, index_to_split)
end

# For simplicity
function concrete_bound(nnet::Network, subdom::Hyperrectangle, type::Symbol)
    points = [subdom.center, low(subdom), high(subdom)]
    values = Vector{Float64}(undef, 0)
    for p in points
        push!(values, sum(compute_output(nnet, p)))
    end
    value, index = ifelse(type == :min, findmin(values), findmax(values))
    return (value, points[index])
end

function approx_bound(nnet::Network, dom::Hyperrectangle, optimizer::AbstractMathProgSolver, type::Symbol)
    bounds = get_bounds(nnet, dom)
    model = Model(solver = optimizer)
    neurons = init_neurons(model, nnet)
    add_set_constraint!(model, dom, first(neurons))
    encode_Δ_lp!(model, nnet, bounds, neurons)
    index = ifelse(type == :max, 1, -1)
    o = sum(last(neurons))
    @objective(model, Max, index * o)
    status = solve(model, suppress_warnings = true)
    status == :Optimal && return getvalue(o)
    error("Could not find bound for dom: ", dom)
end


struct NormalProblem{T}
    func::Function
    domain::T
end

struct DifficultProblem{T}
    func::Function
    domain::T
end

prblm_1 = NormalProblem{IntervalBox}((x,y,z)->x^3+y^2+z, IntervalBox(2..3, 5..6, 4..8))
prblm_2 = NormalProblem{IntervalBox}((x,y,z)->((x^3+y^2)^4+z^2, IntervalBox(2..3, 5..6, 4..8))
prblm_3 = NormalProblem{IntervalBox}((x,y,z)->((x^3+y^2)^4+z^2, IntervalBox(2..3, 5..6, 4..8))


using .StatsBase

@inline @inbounds function skewedGaussianQuadratic(x::AbstractArray, p::Function, t::Real)
    (μ, η, κ) = p(t)
    x = μ*x[1] + μ*x[1]*exp(-κ*x[1])*(κ*x[1]-2.0)
    return SVector{1}(x)
end

@inline @inbounds function skewedGaussianQuadratic_σ(σ::AbstractArray, p::Function, t::Real)
    (μ, η, κ) = p(t)
    σ = η
end

"""Skewed Pitchfork Bifurcation"""
function skewedGaussianQuadratic(P::Process)
    seed(P.solver_rng)
    prob = SDEProblem(P.process, skewedGaussianQuadratic_σ, P.X0, (P.transient_t0, P.tmax), tuplef2ftuple(P.parameter_profile, P.parameter_profile_parameters))
    sol = dsolve(prob, P.alg; dt = P.dt, saveat=P.savedt, P.solver_opts...)
end

skewedGaussianQuadraticSim = Process(
    process = skewedGaussianQuadratic,
    X0 = [0.0],
    parameter_profile = (constantParameter, constantParameter, rampInterval),
    parameter_profile_parameters = ((-1.0,), (2.0,), (0.01, 0.05, 0.0, 1000.0)),
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.001,
    savedt = 0.001,
    tmax = 1000.0,
    alg = EM())
export skewedGaussianQuadraticSim

"""A system combining multiple classes of features. Combine Chen's system with gaussian noise and an exponential measurement function"""
function seeIfItSticks(P::Process)
    seed(P.solver_rng)
    prob = ODEProblem(chensSystem, P.X0, (P.transient_t0, P.tmax), tuplef2ftuple(P.parameter_profile[1:3], P.parameter_profile_parameters[1:3]), jac=chensSystem_J)
    sol = dsolve(prob, P.alg; dt = P.dt, saveat=P.savedt, P.solver_opts...)
    η = parameterseries(P, p=[4], transient=true)
    α = parameterseries(P, p=[5], transient=true)
    sol = sol .+ η.*randn(size(sol))
    sol = (sol.*exp.(α.*(sol.-mean(sol, dims=2))))'
end

seeIfItSticksSim = Process(
    process = seeIfItSticks,
    X0 = [-3.0, 2.0, 20.0],
    parameter_profile = (ramp, ramp, ramp, constant, constant),
    parameter_profile_parameters = ((42.0, 42.0, 0.0, 500.0), (11.0, 11.0, 0.0, 500.0), (28.0, 28.0, 0.0, 500.0), 2.0, 0.01), # (𝑎, 𝑏, 𝑐, η, α)
    transient_t0 = -100.0,
    t0 = 0.0,
    dt = 0.001,
    savedt = 0.025,
    tmax = 500.0,
    alg = RK4(),
    solver_opts = Dict(:adaptive => false))
export seeIfItSticksSim


using LinearAlgebra

const liblapack = Base.liblapack_name
const libblas = Base.libblas_name
import LinearAlgebra.BLAS.@blasfunc
import LinearAlgebra: BlasFloat, BlasInt

function ksvd!(A, U, S, VT, work, calc_work = false)
    m, n   = size(A)
    minmn  = min(m, n)
    lwork  = BlasInt(length(work))
    info   = Ref{BlasInt}()

    jobu = 'S'
    jobvt = 'S'

    # calculate work
    if calc_work
        lwork = BlasInt(-1)
    end

    ccall((@blasfunc(dgesvd_), liblapack), Cvoid,
                      (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt},
                       Ptr{Float64}, Ref{BlasInt}, Ptr{Float64}, Ptr{Float64},
                       Ref{BlasInt}, Ptr{Float64}, Ref{BlasInt}, Ptr{Float64},
                       Ref{BlasInt}, Ptr{BlasInt}),
                      jobu, jobvt, m, n, A, m, S, U, m, VT, n,
                      work, lwork, info)

    if calc_work
        lwork = BlasInt(real(work[1]))
        resize!(work, lwork)
    end
end

function kgemv!(tA, rank, alpha, A, x, beta, y)
    ccall((@blasfunc(dgemv_), libblas), Cvoid,
                (Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt}, Ref{Float64},
                 Ptr{Float64}, Ref{BlasInt}, Ptr{Float64}, Ref{BlasInt},
                 Ref{Float64}, Ptr{Float64}, Ref{BlasInt}),
                 tA, size(A,1), rank, alpha,
                 A, size(A,1), x, 1,
                 beta, y, 1)
end

# adapted from http://pi.math.cornell.edu/~web6140/TopTenAlgorithms/QRUpdate.html
# w is a work vector
function qrupdate!(Q, R, u, v, w)
    # Compute the QR factorization of Q*R + u*v':

    # Note that Q*R + u*v' = Q*(R + w*v') with w = Q'*u:
    mul!(w, Q', u)
    n = size(Q, 1)

    # Convert R+w*v' into upper-hessenberg form using n-1 Givens rotations:
    for k = n-1:-1:1
        c, s, r = LinearAlgebra.givensAlgorithm(w[k], w[k+1])
        w[k+1] = 0.; w[k] = r

        # Compute G*R[k:k+1,:] and Q[:,k:k+1]*G', where G = [c -s ; s c]
        for j = 1:n
            newrow = c*R[k,j] + s*R[k+1,j]
            R[k+1,j] = -s*R[k,j] + c*R[k+1,j]
            R[k,j] = newrow
            newcol = c*Q[j,k] + s*Q[j,k+1]
            Q[j,k+1] = -s*Q[j,k] + c*Q[j,k+1]
            Q[j,k] = newcol
        end
    end

    # R <- R + w*v' is now upper-hessenberg:
    R[1,:] .+= w[1].*v

    # Convert R from upper-hessenberg form to upper-triangular form using n-1 Givens rotations:
    for k = 1:n-1
        c, s, r = LinearAlgebra.givensAlgorithm(R[k,k], R[k+1,k])

        # Compute G*R[k:k+1,:] and Q[:,k:k+1]*G', where G = [c -s ; s c]
        for j = 1:n
            newrow = c*R[k,j] + s*R[k+1,j]
            R[k+1,j] = -s*R[k,j] + c*R[k+1,j]
            R[k,j] = newrow
            newcol = c*Q[j,k] + s*Q[j,k+1]
            Q[j,k+1] = -s*Q[j,k] + c*Q[j,k+1]
            Q[j,k] = newcol
        end
    end
end


abstract type AbstractWindShearModel end

struct PowerLawWindShear{TF} <: AbstractWindShearModel

    # model parameter
    shear_exponent::TF

end

#TODO add log shear

function adjust_for_wind_shear(locz, reference_velocity, reference_height, ground_height, model::PowerLawWindShear)

    # initialize adjusted wind speed to zero
    adjusted_wind_speed = 0.0
    shear_exp = model.shear_exponent

    # check that the point of interest is above ground level
    if locz >= ground_height
        # adjusted wind speed for wind shear if point is above ground
        adjusted_wind_speed = reference_velocity*((locz-ground_height)/(reference_height-ground_height))^shear_exp
    end

    return adjusted_wind_speed
end

# function adjust_for_wind_shear(loc, point_velocity_no_shear, reference_height, ground_height, shear_model::CubicWindShear)

#     # initialize adjusted wind speed to zero
#     adjusted_wind_speed = 0.0

#     # check that the point of interest is above ground level
#     if loc[3] >= ground_height
#         # adjusted wind speed for wind shear if point is above ground
#         adjusted_wind_speed = point_velocity_no_shear*((loc[3]-ground_height)/(reference_height-ground_height))^shear_exp
#     else
#         # if the point of interest is below ground, set the wind speed to 0.0
#         adjusted_wind_speed = 0.0
#     end

#     return adjusted_wind_speed
# end


## Metsa-Oy Forest Products Supply Chain
## Edward J. Xu <edxu96@outlook.com>
## Jan 24th, 2020

using JuMP, Gurobi

include("func.jl")
include("para.jl")
include("optim_1.jl")
include("optim_2.jl")
include("optim_3.jl")
include("result_2.jl")
include("result_3.jl")

function main()
  alpha_t, beta_t, a_i, b_i, e_i, c_i, r_saw, r_plywood, a_j, c_j,
    r_j, a_paper_t, b_paper_j, c_paper, r_paper, gamma_ik, delta_ik, gamma_jk,
    delta_jk, gamma_paper_k, delta_paper_k, p_fuel, sigma, omega_i, omega_j,
    omega_paper, nu_saw, nu_plywood, nu_j, nu_paper, o_saw, o_plywood, o_j,
    o_paper, i_t1, i_t2, j_t2, pi_n, rho_nm = get_para()

  optim_1(alpha_t, beta_t, a_i, b_i, e_i, c_i, r_saw, r_plywood, a_j, c_j, r_j,
    a_paper_t, b_paper_j, c_paper, r_paper, gamma_ik, delta_ik, gamma_jk,
    delta_jk, gamma_paper_k, delta_paper_k, p_fuel)

  optim_2(alpha_t, beta_t, a_i, b_i, e_i, c_i, r_saw, r_plywood, a_j, c_j,
    r_j, a_paper_t, b_paper_j, c_paper, r_paper, gamma_ik, delta_ik, gamma_jk,
    delta_jk, gamma_paper_k, delta_paper_k, p_fuel, sigma, omega_i, omega_j,
    omega_paper, nu_saw, nu_plywood, nu_j, nu_paper, o_saw, o_plywood, o_j,
    o_paper, i_t1, i_t2, j_t2)

  optim_3(alpha_t, beta_t, a_i, b_i, e_i, c_i, r_saw, r_plywood, a_j,
    c_j, r_j, a_paper_t, b_paper_j, c_paper, r_paper, gamma_ik, delta_ik,
    gamma_jk, delta_jk, gamma_paper_k, delta_paper_k, p_fuel, sigma, omega_i,
    omega_j, omega_paper, nu_saw, nu_plywood, nu_j, nu_paper, o_saw, o_plywood,
    o_j, o_paper, i_t1, i_t2, j_t2, pi_n, rho_nm)
end

main()


let
    using UnsteadyFlowSolvers

    # Set up a single vortex
    start_vc = 0.02 # Core radius
    start_vort = 1.
    re = 1e3
    ref_len = 1
    nu = start_vort / re
    vort = TwoDVVort(0., 0., start_vort, start_vc, 0., 0., 0.)
    field = TwoDVFlowField(nu)
    push!(field.tev, vort)

    # Solver settings
    dt = 0.015
    t0 = 0.
    tmax = 1.5
    t = t0
    nsteps = Int(ceil((tmax-t)/dt)) + 1

    # Comparison to the (Gaussian) Lamb-Oseen vortex
    rc = t->sqrt(4 * nu * (t-t0) + start_vc^2)

    for i = 1:nsteps
        t += dt
        UnsteadyFlowSolvers.wakeroll(field, dt)
        println("Step\t",i,":\tCore radius is \t",field.tev[1].vc, "\tLamb-Oseen radius is\t", rc(t))
    end
end


using LinearAlgebra, Arpack, IterativeSolvers
using Random, Distributions

struct Schur{
    T,
    M<:AbstractMatrix{T},
    N<:AbstractMatrix{T},
    O<:AbstractMatrix{T},
    P<:AbstractMatrix{T}
}
    W::M
    Ku::N
    Kux::O
    Kx::P
    ϵ::T
end

function Matrix(sc::Schur)
    return sc.Kx .- sc.W' * (sc.Kux .- sc.Ku * sc.W) .- sc.Kux' * sc.W + sc.ϵ * I
end
eltype(::Schur{T}) where T = T
size(sc::Schur, args...) = size(sc.Kx, args...)
function mul!(y, sc::Schur, x, α, β)
    p, n = size(sc.W)
    v1 = similar(x, p)
    v2 = similar(x, p)

    y .= y .* β .+ x .* (α * sc.ϵ)

    mul!(y, sc.Kx, x, α, 1)

    mul!(v1, sc.W, x)
    mul!(v2, sc.Ku, v1)

    mul!(y, sc.W', v2, α, 1)
    mul!(y, sc.Kux', v1, -α, 1)

    mul!(v1, sc.Kux, x)

    mul!(y, sc.W', v1, -α, 1)

    return y
end

function lanczos_arpack(A, k, v; maxiter, tol)
    T = eltype(A)
    n = size(A, 1)
    mulA! = (y, x) -> mul!(y, A, x)
    id = x -> x
    # in: (T, mulA!, mulB, solveSI, n, issym, iscmplx, bmat,
    #            nev, ncv, whichstr, tol, maxiter, mode, v0)
    # out: (resid, v, ldv, iparam, ipntr, workd, workl, lworkl, rwork, TOL)
    out = Arpack.aupd_wrapper(T, mulA!, id, id, n, true, false, "I",
                       1, k, "LM", tol, maxiter, 1, v)

    α = out[7][k+1:2*k-1]
    β = out[7][2:k-1]

    return out[2], α, β, out[1]
end

function _lanczos_logdet!(z, acc, sc, k; maxiter, tol, nsamples)
    for i in 1:nsamples
        rand!(Normal(), z)
        z .= sign.(z)
        Q, α, β, resid = lanczos_arpack(sc, k, z; maxiter=maxiter, tol=tol)
        T = SymTridiagonal(α, β)
        Λ = eigen(T)
        wts = Λ.vectors[1, :] .^ 2 .* norm(z) ^ 2
        acc += dot(wts, log.(Λ.values))
    end
    return acc / nsamples
end

function lanczos_logdet(sc, k; maxiter, tol, nsamples)
    z = Vector{eltype(sc)}(undef, size(sc, 1))
    return _lanczos_logdet!(z, 0, sc, k;
                            maxiter=maxiter, tol=tol, nsamples=nsamples)
end


using Test
using TwoPointFunctions

import TwoPointFunctions: ∂θ₁∂θ₂cosχ, ∂θ₂∂θ₁cosχ, ∂ϕ₁∂ϕ₂cosχ, ∂ϕ₂∂ϕ₁cosχ,
∂ϕ₂∂ϕ₂cosχ, ∂ϕ₁∂ϕ₁cosχ, ∂θ₁∂ϕ₁cosχ, ∂θ₂∂ϕ₂cosχ, ∂ϕ₁∂θ₁cosχ, ∂ϕ₂∂θ₂cosχ

isapproxzero(x) = isapprox(x, zero(x), atol=1e-14, rtol=1e-8)

@testset "cosχ" begin
	# cos(θ₁)cos(θ₂) + sin(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂)
    @test cosχ((0, 0), (0, 0)) == 1
    @test cosχ((0, 0), (0, pi/3)) == 1
    @test cosχ((0, 0), (pi, 0)) == -1
    @test cosχ((0, 0), (pi, pi/3)) == -1

    @test cosχ((0, 0), (0,0)) == 1
    @test cosχ((0, pi/3), (0,0)) == 1
    @test cosχ((pi, 0), (0,0)) == -1
    @test cosχ((pi, pi/3), (0,0)) == -1

    @test isapproxzero(cosχ((0, 0), (pi/2, 0)))
    @test isapproxzero(cosχ((0, 0), (pi/2, pi/3)))

    @test isapproxzero(cosχ((pi/2, 0), (0, 0)))
    @test isapproxzero(cosχ((pi/2, pi/3), (0, 0)))

    @test cosχ((0, 0), (pi/3, 0)) ≈ cos(pi/3)
    @test cosχ((0, 0), (pi/3, pi/3)) ≈ cos(pi/3)
end

@testset "first derivatives" begin
    @testset "∂θ₁cosχ" begin
		# -sin(θ₁)cos(θ₂) + cos(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂)
	    @test isapproxzero(∂θ₁cosχ((0, 0), (0, 0)))
	    @test isapproxzero(∂θ₁cosχ((0, 0), (0, pi/3)))
	    @test isapproxzero(∂θ₁cosχ((0, 0), (pi, 0)))
	    @test isapproxzero(∂θ₁cosχ((0, 0), (pi, pi/3)))
	    @test ∂θ₁cosχ((0, 0), (pi/2, 0)) ≈ 1
	    @test ∂θ₁cosχ((0, 0), (pi/2, pi/3)) ≈ 0.5
	    @test ∂θ₁cosχ((pi/2, 0), (0, 0)) ≈ -1
	    @test ∂θ₁cosχ((pi/2, pi/3), (0, 0)) ≈ -1
	end

	@testset "∂θ₂cosχ" begin
		# -cos(θ₁)sin(θ₂) + sin(θ₁)cos(θ₂)cos(ϕ₁ - ϕ₂)
	    @test isapproxzero(∂θ₂cosχ((0, 0), (0, 0)))
	    @test isapproxzero(∂θ₂cosχ((0, 0), (0, pi/3)))
	    @test isapproxzero(∂θ₂cosχ((0, 0), (pi, 0)))
	    @test isapproxzero(∂θ₂cosχ((0, 0), (pi, pi/3)))
	    @test ∂θ₂cosχ((0, 0), (pi/2, 0)) ≈ -1
	    @test ∂θ₂cosχ((0, 0), (pi/2, pi/3)) ≈ -1
	    @test ∂θ₂cosχ((pi/2, 0), (0, 0)) ≈ 1
	    @test ∂θ₂cosχ((pi/2, pi/3), (0, 0)) ≈ 0.5
	end

	@testset "∂ϕ₁cosχ" begin
		# -sin(θ₁)sin(θ₂)sin(ϕ₁ - ϕ₂)
	    @test isapproxzero(∂ϕ₁cosχ((0, 0), (0, 0)))
	    @test isapproxzero(∂ϕ₁cosχ((0, 0), (0, pi/3)))
	    @test isapproxzero(∂ϕ₁cosχ((0, 0), (pi, 0)))
	    @test isapproxzero(∂ϕ₁cosχ((0, 0), (pi, pi/3)))

	    @test isapproxzero(∂ϕ₁cosχ((0, 0), (0, 0)))
	    @test isapproxzero(∂ϕ₁cosχ((0, pi/3), (0, 0)))
	    @test isapproxzero(∂ϕ₁cosχ((pi, 0), (0, 0)))
	    @test isapproxzero(∂ϕ₁cosχ((pi, pi/3), (0, 0)))

	    @test isapproxzero(∂ϕ₁cosχ((pi/2, 0), (pi/2, 0)))
	    @test ∂ϕ₁cosχ((pi/2, 0), (pi/2, pi/6)) ≈ 0.5
	    @test ∂ϕ₁cosχ((pi/2, pi/6), (pi/2, 0)) ≈ -0.5
	    @test ∂ϕ₁cosχ((pi/2, pi/2), (pi/2, 0)) ≈ -1
	end

	@testset "∇ϕ₁cosχ" begin
	    # -sin(θ₂)sin(ϕ₁ - ϕ₂)
	    @test isapprox(∇ϕ₁cosχ((0, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₁cosχ((0, 0), (0, pi/3)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₁cosχ((0, 0), (pi, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₁cosχ((0, 0), (pi, pi/3)), 0, atol=1e-14)

	    @test isapprox(∇ϕ₁cosχ((0, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₁cosχ((0, pi/3), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₁cosχ((pi, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₁cosχ((pi, pi/3), (0, 0)), 0, atol=1e-14)

	    @test isapprox(∇ϕ₁cosχ((pi/2, 0), (pi/2, 0)), 0, atol=1e-14)
	    @test ∇ϕ₁cosχ((pi/2, 0), (pi/2, pi/6)) ≈ 0.5
	    @test ∇ϕ₁cosχ((pi/2, pi/6), (pi/2, 0)) ≈ -0.5
	    @test ∇ϕ₁cosχ((pi/2, pi/2), (pi/2, 0)) ≈ -1
	end

	@testset "∂ϕ₂cosχ" begin
		# sin(θ₁)sin(θ₂)sin(ϕ₁ - ϕ₂)
	    @test isapprox(∂ϕ₂cosχ((0, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∂ϕ₂cosχ((0, 0), (0, pi/3)), 0, atol=1e-14)
	    @test isapprox(∂ϕ₂cosχ((0, 0), (pi, 0)), 0, atol=1e-14)
	    @test isapprox(∂ϕ₂cosχ((0, 0), (pi, pi/3)), 0, atol=1e-14)

	    @test isapprox(∂ϕ₂cosχ((0, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∂ϕ₂cosχ((0, pi/3), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∂ϕ₂cosχ((pi, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∂ϕ₂cosχ((pi, pi/3), (0, 0)), 0, atol=1e-14)

	    @test isapprox(∂ϕ₂cosχ((pi/2, 0), (pi/2, 0)), 0, atol=1e-14)
	    @test ∂ϕ₂cosχ((pi/2, 0), (pi/2, pi/6)) ≈ -0.5
	    @test ∂ϕ₂cosχ((pi/2, pi/6), (pi/2, 0)) ≈ 0.5
	    @test ∂ϕ₂cosχ((pi/2, pi/2), (pi/2, 0)) ≈ 1
	end

	@testset "∇ϕ₂cosχ" begin
	    # sin(θ₁)sin(ϕ₁ - ϕ₂)
	    @test isapprox(∇ϕ₂cosχ((0, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₂cosχ((0, 0), (0, pi/3)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₂cosχ((0, 0), (pi, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₂cosχ((0, 0), (pi, pi/3)), 0, atol=1e-14)

	    @test isapprox(∇ϕ₂cosχ((0, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₂cosχ((0, pi/3), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₂cosχ((pi, 0), (0, 0)), 0, atol=1e-14)
	    @test isapprox(∇ϕ₂cosχ((pi, pi/3), (0, 0)), 0, atol=1e-14)

	    @test isapprox(∇ϕ₂cosχ((pi/2, 0), (pi/2, 0)), 0, atol=1e-14)
	    @test ∇ϕ₂cosχ((pi/2, 0), (pi/2, pi/6)) ≈ -0.5
	    @test ∇ϕ₂cosχ((pi/2, pi/6), (pi/2, 0)) ≈ 0.5
	    @test ∇ϕ₂cosχ((pi/2, pi/2), (pi/2, 0)) ≈ 1
	end
end

@testset "second derivatives" begin
    @testset "∂²θ₁cosχ" begin
        # -cos(θ₁)cos(θ₂) -sin(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂) == -cosχ
        @test isapproxzero(∂²θ₁cosχ((pi/2, 0), (0, 0)))
        @test ∂²θ₁cosχ((pi/2, 0), (pi/2, 0)) == -1
        @test ∂²θ₁cosχ((pi/2, 0), (pi/2, pi/3)) ≈ -0.5

        @test ∂²θ₁cosχ((0, 0), (0, 0)) ≈ -cosχ((0, 0), (0, 0))
        @test ∂²θ₁cosχ((0, 0), (pi, 0)) ≈ -cosχ((0, 0), (pi, 0))
        @test ∂²θ₁cosχ((0, 0), (pi/2, 0)) ≈ -cosχ((0, 0), (pi/2, 0))
        @test ∂²θ₁cosχ((0, pi/3), (pi/2, 0)) ≈ -cosχ((0, pi/3), (pi/2, 0))
    end

    @testset "∂²θ₂cosχ" begin
        # -cos(θ₁)cos(θ₂) - sin(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂) = -cosχ
        @test isapproxzero(∂²θ₂cosχ((pi/2, 0), (0, 0)))
        @test ∂²θ₂cosχ((pi/2, 0), (pi/2, 0)) == -1
        @test ∂²θ₂cosχ((pi/2, 0), (pi/2, pi/3)) ≈ -0.5

        @test ∂²θ₂cosχ((0, 0), (0, 0)) ≈ -cosχ((0, 0), (0, 0))
        @test ∂²θ₂cosχ((0, 0), (pi, 0)) ≈ -cosχ((0, 0), (pi, 0))
        @test ∂²θ₂cosχ((0, 0), (pi/2, 0)) ≈ -cosχ((0, 0), (pi/2, 0))
        @test ∂²θ₂cosχ((0, pi/3), (pi/2, 0)) ≈ -cosχ((0, pi/3), (pi/2, 0))
    end

    @testset "∂θ₁∂θ₂cosχ" begin
    	# sin(θ₁)sin(θ₂) + cos(θ₁)cos(θ₂)cos(ϕ₁ - ϕ₂)

    	@test ∂θ₁∂θ₂cosχ((0, 0), (0, 0)) ≈ 1
    	@test ∂θ₁∂θ₂cosχ((0, pi/3), (0, 0)) ≈ cos(pi/3)
    	@test ∂θ₁∂θ₂cosχ((0, pi/3), (pi, 0)) ≈ -cos(pi/3)

    	@test isapproxzero(∂θ₁∂θ₂cosχ((pi/2, pi/3), (0, 0)))
    	@test isapproxzero(∂θ₁∂θ₂cosχ((pi/2, pi/3), (0, pi/3)))
    	@test isapproxzero(∂θ₁∂θ₂cosχ((pi/2, pi/3), (pi, 0)))
    	@test isapproxzero(∂θ₁∂θ₂cosχ((pi/2, 0), (pi, 0)))
    	@test isapproxzero(∂θ₁∂θ₂cosχ((pi/2, 0), (pi, pi/3)))

    	@test ∂θ₁∂θ₂cosχ((pi/2, 0), (pi/2, 0)) ≈ 1
    	@test ∂θ₁∂θ₂cosχ((pi/2, 0), (pi/2, pi/3)) ≈ 1

        @testset "∂θ₂∂θ₁cosχ" begin
        	for θ₁ in LinRange(0,π,10), ϕ₁ in LinRange(0,2π,20),
            	θ₂ in LinRange(0,π,10), ϕ₂ in LinRange(0,2π,20)

            	@test isapprox(∂θ₂∂θ₁cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)) , ∂θ₁∂θ₂cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)), atol=1e-14, rtol=1e-8)
        	end
        end
    end

    @testset "∂²ϕ₁cosχ" begin
        # -sin(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂)

        @test isapproxzero(∂²ϕ₁cosχ((0, 0), (0, 0)))
        @test isapproxzero(∂²ϕ₁cosχ((0, 0), (pi,0)))
        @test isapproxzero(∂²ϕ₁cosχ((pi, π/4), (pi,0)))
        @test isapproxzero(∂²ϕ₁cosχ((0, pi/3), (0, 0)))
        @test isapproxzero(∂²ϕ₁cosχ((0, pi/3), (0, pi/2)))
        @test isapproxzero(∂²ϕ₁cosχ((pi/4, pi/2), (pi/4, 0)))

        @test ∂²ϕ₁cosχ((pi/4, 0), (pi/4, 0)) ≈ -sin(π/4)^2
        @test ∂²ϕ₁cosχ((pi/4, pi/6), (pi/4, pi/6)) ≈ -sin(π/4)^2

        @test ∂²ϕ₁cosχ((pi/4, pi/3), (pi/4, pi/6)) ≈ -sin(π/4)^2*cos(π/6)

        @testset "∂ϕ₁∂ϕ₁cosχ" begin
        	for θ₁ in LinRange(0,π,10), ϕ₁ in LinRange(0,2π,20),
            	θ₂ in LinRange(0,π,10), ϕ₂ in LinRange(0,2π,20)

            	@test isapprox(∂²ϕ₁cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)) , ∂ϕ₁∂ϕ₁cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)), atol=1e-14, rtol=1e-8)
        	end
        end
    end

    @testset "∂²ϕ₂cosχ" begin
        # -sin(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂)

        @test isapproxzero(∂²ϕ₂cosχ((0, 0), (0, 0)))
        @test isapproxzero(∂²ϕ₂cosχ((0, 0), (pi,0)))
        @test isapproxzero(∂²ϕ₂cosχ((pi, π/4), (pi,0)))
        @test isapproxzero(∂²ϕ₂cosχ((0, pi/3), (0, 0)))
        @test isapproxzero(∂²ϕ₂cosχ((0, pi/3), (0, pi/2)))
        @test isapproxzero(∂²ϕ₂cosχ((pi/4, pi/2), (pi/4, 0)))

        @test ∂²ϕ₂cosχ((pi/4, 0), (pi/4, 0)) ≈ -sin(π/4)^2
        @test ∂²ϕ₂cosχ((pi/4, pi/6), (pi/4, pi/6)) ≈ -sin(π/4)^2

        @test ∂²ϕ₂cosχ((pi/4, pi/3), (pi/4, pi/6)) ≈ -sin(π/4)^2*cos(π/6)

        @testset "∂ϕ₂∂ϕ₂cosχ" begin
        	for θ₁ in LinRange(0,π,10), ϕ₁ in LinRange(0,2π,20),
            	θ₂ in LinRange(0,π,10), ϕ₂ in LinRange(0,2π,20)

            	@test isapprox(∂²ϕ₂cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)) , ∂ϕ₂∂ϕ₂cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)), atol=1e-14, rtol=1e-8)
        	end
        end
    end

    @testset "∂ϕ₁∂ϕ₂cosχ" begin
        # sin(θ₁)sin(θ₂)cos(ϕ₁ - ϕ₂)

        @test isapproxzero(∂ϕ₁∂ϕ₂cosχ((0, 0), (0, 0)))
        @test isapproxzero(∂ϕ₁∂ϕ₂cosχ((0, 0), (pi,0)))
        @test isapproxzero(∂ϕ₁∂ϕ₂cosχ((pi, π/4), (pi,0)))
        @test isapproxzero(∂ϕ₁∂ϕ₂cosχ((0, pi/3), (0, 0)))
        @test isapproxzero(∂ϕ₁∂ϕ₂cosχ((0, pi/3), (0, pi/2)))
        @test isapproxzero(∂ϕ₁∂ϕ₂cosχ((pi/4, pi/2), (pi/4, 0)))

        @test ∂ϕ₁∂ϕ₂cosχ((pi/4, 0), (pi/4, 0)) ≈ sin(π/4)^2
        @test ∂ϕ₁∂ϕ₂cosχ((pi/4, pi/6), (pi/4, pi/6)) ≈ sin(π/4)^2

        @test ∂ϕ₁∂ϕ₂cosχ((pi/4, pi/3), (pi/4, pi/6)) ≈ sin(π/4)^2*cos(π/6)

        @testset "∂ϕ₂∂ϕ₁cosχ" begin
        	for θ₁ in LinRange(0,π,10), ϕ₁ in LinRange(0,2π,20),
            	θ₂ in LinRange(0,π,10), ϕ₂ in LinRange(0,2π,20)

            	@test isapprox(∂ϕ₂∂ϕ₁cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)) , ∂ϕ₁∂ϕ₂cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)), atol=1e-14, rtol=1e-8)
        	end
        end
    end

    @testset "∂θ₁∂ϕ₁cosχ" begin

        @testset "∂ϕ₁∂θ₁cosχ" begin
        	for θ₁ in LinRange(0,π,10), ϕ₁ in LinRange(0,2π,20),
            	θ₂ in LinRange(0,π,10), ϕ₂ in LinRange(0,2π,20)

            	@test isapprox(∂θ₁∂ϕ₁cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)) , ∂ϕ₁∂θ₁cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)), atol=1e-14, rtol=1e-8)
        	end
        end
    end

    @testset "∂θ₂∂ϕ₂cosχ" begin

        @testset "∂ϕ₂∂θ₂cosχ" begin
        	for θ₁ in LinRange(0,π,10), ϕ₁ in LinRange(0,2π,20),
            	θ₂ in LinRange(0,π,10), ϕ₂ in LinRange(0,2π,20)

            	@test isapprox(∂θ₂∂ϕ₂cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)) , ∂ϕ₂∂θ₂cosχ((θ₁, ϕ₁), (θ₂, ϕ₂)), atol=1e-14, rtol=1e-8)
        	end
        end
    end
end


using UniqueVectors: UniqueVector
using LinearAlgebra: issymmetric, diag, I
using SparseArrays: spzeros

export Circuit
export partial_copy, matrices
export get_inv_inductance, get_inductance
export get_conductance, get_resistance
export get_elastance, get_capacitance
export set_inv_inductance!, set_inductance!
export set_conductance!, set_resistance!
export set_capacitance!, set_elastance!

export SpanningTree, coordinate_matrix, unite_vertices

SparseMat{T} = SparseMatrixCSC{T, Int}

"""
    struct Circuit{T, K<:Number, G<:Number, C<:Number}
Model of a circuit, containing matrix representations of the inverse inductance,
conductance, and capacitance between any pair of named vertices. All vertices are
named.

## Fields
- `k`: inverse inductance matrix
- `g`: conductance matrix
- `c`: capacitance matrix
- `vertices`: vector of uniquely named vertices

## Note
The matrices are not in the conventional form, i.e. `c` is not a Maxwell capacitance
matrix. Rather, `c[i,j]` is only the direct capacitance between vertices `i` and `j`.
`iszero(c[i,i])` is demanded for all `i`.
"""
struct Circuit{T, K<:Number, G<:Number, C<:Number}
    k::SparseMat{K} # inverse inductances
    g::SparseMat{G} # conductances
    c::SparseMat{C} # capacitances
    vertices::UniqueVector{T}
    function Circuit{T}(k::SparseMat, g::SparseMat, c::SparseMat,
            vertices::UniqueVector{T}) where {T}
        for m in (k, g, c)
            @assert isreal(m) && issymmetric(m)
            @assert size(m) == (length(vertices), length(vertices))
            @assert all(iszero.(diag(m)))
            @assert all(m .>= zero(eltype(m)))
        end
        kk, gg, cc = float(k), float(g), float(c)
        K, G, C = eltype(kk), eltype(gg), eltype(cc)
        return new{T,K,G,C}(kk, gg, cc, vertices)
    end
end

"""
    Circuit(k::AbstractMatrix, g::AbstractMatrix, c::AbstractMatrix,
        vertices::AbstractVector{T}) where {T}
    Circuit(vertices::AbstractVector)

Construct a circuit with named vertices, defaulting to zero inverse inductance,
conductance, and capacitance between the vertices if the matrices are not given.
"""
function Circuit(k::AbstractMatrix, g::AbstractMatrix, c::AbstractMatrix,
        vertices::AbstractVector{T}) where {T}
    return Circuit{T}(sparse(float(k)), sparse(float(g)), sparse(float(c)),
        UniqueVector(vertices))
end

function Circuit(vertices::AbstractVector)
    z() = spzeros(length(vertices), length(vertices))
    return Circuit(z(), z(), z(), vertices)
end

"""
    partial_copy(circ::Circuit;
       k::Union{AbstractMatrix, Nothing} = nothing,
       g::Union{AbstractMatrix, Nothing} = nothing,
       c::Union{AbstractMatrix, Nothing} = nothing,
       vertices::Union{AbstractVector, Nothing} = nothing)

Create a new `Circuit` with the same fields except those given as keyword arguments.
"""
function partial_copy(circ::Circuit;
       k::Union{AbstractMatrix, Nothing} = nothing,
       g::Union{AbstractMatrix, Nothing} = nothing,
       c::Union{AbstractMatrix, Nothing} = nothing,
       vertices::Union{AbstractVector, Nothing} = nothing)
   k = isnothing(k) ? circ.k : k
   g = isnothing(g) ? circ.g : g
   c = isnothing(c) ? circ.c : c
   vertices = isnothing(vertices) ? circ.vertices : vertices
   return Circuit(k, g, c, vertices)
end

"""
    ==(circ1::Circuit, circ2::Circuit)
Test whether two circuits have the same circuit components and named vertices
(in order).
"""
# don't remove space between Base. and ==
function Base. ==(circ1::Circuit, circ2::Circuit)
    return all(getfield(circ1, name) == getfield(circ2, name)
        for name in fieldnames(Circuit))
end

"""
    isapprox(circ1::Circuit, circ2::Circuit)
Tests for approximate equality of the circuits' `k`, `g`, `c` matrices.
"""
function Base.isapprox(circ1::Circuit, circ2::Circuit)
    return all(getfield(circ1, name) ≈ getfield(circ2, name)
        for name in fieldnames(Circuit)[1:end-1])
end

matrices(c::Circuit) = (c.k, c.g, c.c)

"""
    get_matrix_element(c::Circuit{T}, matrix_name::Symbol, v0::T, v1::T) where {T}
Returns a matrix element from one of the circuit matrices corresponding to the
edge between vertices with names `v0` and `v1`. `matrix_name` must be in
`(:k, :g, :c)`.
"""
function get_matrix_element(c::Circuit{T}, matrix_name::Symbol, v0::T, v1::T) where {T}
    @assert matrix_name in (:k, :g, :c)
    i0 = findfirst(isequal(v0), c.vertices)
    i1 = findfirst(isequal(v1), c.vertices)
    return getfield(c, matrix_name)[i0, i1]
end

"""
    get_inv_inductance(c::Circuit{T,K}, v0, v1) where {T,K}
Returns the inverse of the inductance on the edge between between vertices `v0` and `v1`.
"""
get_inv_inductance(c::Circuit{T, K}, v0, v1) where {T, K} =
    get_matrix_element(c, :k, v0, v1)::K

"""
    get_inductance(c::Circuit, v0, v1)
Returns the inductance on the edge between between vertices `v0` and `v1`.
"""
get_inductance(c::Circuit, v0, v1) = inv(get_inv_inductance(c, v0, v1))

"""
    get_conductance(c::Circuit{T,K,G}, v0, v1) where {T,K,G}
Returns the conductance on the edge between between vertices `v0` and `v1`.
"""
get_conductance(c::Circuit{T, K, G}, v0, v1) where {T, K, G} =
    get_matrix_element(c, :g, v0, v1)::G

"""
    get_resistance(c::Circuit, v0, v1)
Returns the resistance on the edge between between vertices `v0` and `v1`.
"""
get_resistance(c::Circuit, v0, v1) = inv(get_conductance(c, v0, v1))

"""
    get_capacitance(c::Circuit{T,K,G,C}, v0, v1) where {T,K,G,C}
Returns the capacitance on the edge between between vertices `v0` and `v1`.
"""
get_capacitance(c::Circuit{T,K,G,C}, v0, v1) where {T,K,G,C} =
    get_matrix_element(c, :c, v0, v1)::C

"""
    get_elastance(c::Circuit, v0, v1)
Returns the elastance on the edge between between vertices `v0` and `v1`.
"""
get_elastance(c::Circuit, v0, v1) = inv(get_capacitance(c, v0, v1))

"""
    set_matrix_element!(c::Circuit, matrix_name::Symbol, v0, v1, value)
Sets a matrix element in one of the circuit matrices corresponding to the
edge between vertices with names `v0` and `v1`. `matrix_name` must be in
`(:k, :g, :c)`.
"""
function set_matrix_element!(c::Circuit, matrix_name::Symbol, v0, v1, value)
    @assert matrix_name in (:k, :g, :c)
    @assert value >= zero(value)
    @assert !isinf(value)
    i0 = findfirst(isequal(v0), c.vertices)
    i1 = findfirst(isequal(v1), c.vertices)
    @assert i0 != i1
    matrix = getfield(c, matrix_name)
    matrix[i0, i1] = value
    matrix[i1, i0] = value
    return c
end

"""
    set_inv_inductance!(c::Circuit, v0, v1)
Sets the inverse of the inductance on the edge between between vertices `v0` and `v1`.
"""
set_inv_inductance!(c::Circuit, v0, v1, value) =
    set_matrix_element!(c, :k, v0, v1, value)

"""
    set_inductance!(c::Circuit, v0, v1)
Sets the inductance on the edge between between vertices `v0` and `v1`.
"""
set_inductance!(c::Circuit, v0, v1, value) =
    set_inv_inductance!(c, v0, v1, inv(value))

"""
    set_conductance!(c::Circuit, v0, v1)
Sets the conductance on the edge between between vertices `v0` and `v1`.
"""
set_conductance!(c::Circuit, v0, v1, value) =
    set_matrix_element!(c, :g, v0, v1, value)

"""
    set_resistance!(c::Circuit, v0, v1)
Sets the resistance on the edge between between vertices `v0` and `v1`.
"""
set_resistance!(c::Circuit, v0, v1, value) =
    set_conductance!(c, v0, v1, inv(value))

"""
    set_capacitance!(c::Circuit, v0, v1)
Sets the capacitance on the edge between between vertices `v0` and `v1`.
"""
set_capacitance!(c::Circuit, v0, v1, value) =
    set_matrix_element!(c, :c, v0, v1, value)

"""
    set_elastance!(c::Circuit, v0, v1)
Sets the elastance on the edge between between vertices `v0` and `v1`.
"""
set_elastance!(c::Circuit, v0, v1, value) =
    set_capacitance!(c, v0, v1, inv(value))

#######################################################
# spanning trees
#######################################################

"""
    struct SpanningTree{T}
A spanning tree of a graph with vertices of type `T`.

## Fields
- `root`: the root of the spanning tree
- `edges`: a vector of tuples `(vertex1, vertex2)` defining tree edges.
"""
struct SpanningTree{T}
    root::T
    edges::Vector{Tuple{T, T}}
end

"""
    SpanningTree(vertices::AbstractVector{T}) where T
Construct a depth-1 spanning tree from a vector of vertices of a complete graph,
with the first vertex taken as the root of the spanning tree.
"""
function SpanningTree(vertices::AbstractVector{T}) where T
    @assert !isempty(vertices)
    root = first(vertices)
    edges = [(v, root) for v in vertices[eachindex(vertices)[2:end]]]
    return SpanningTree{T}(root, edges)
end

"""
    coordinate_matrix(c::Circuit{T}, tree::SpanningTree{T}) where T
Returns the `T` matrix described in section IIC of arXiv:1810.11510 for a
circuit `c` with a given spanning `tree`. The direction of the edges passed in
is ignored and enforced as being outwards from the root.
"""
function coordinate_matrix(c::Circuit{T}, tree::SpanningTree{T}) where T
    V = length(c.vertices)
    @assert tree.root in c.vertices
    @assert length(tree.edges) == V-1 # the correct number of edges for a spanning tree
    mapping = Dict(tree.root => spzeros(V-1)) # the function l in equation 2
    edge_set = Set(enumerate(tree.edges))
    while !isempty(edge_set)
        for e in edge_set
            (index, (v0, v1)) = e
            @assert !(v0 in keys(mapping) && v1 in keys(mapping)) # edges contains a loop
            if v0 in keys(mapping) || v1 in keys(mapping)
                delete!(edge_set, e)
                if v0 in keys(mapping) # enforces all edges are oriented away from the root
                    mapping[v1] = 1.0 * mapping[v0]
                    mapping[v1][index] += 1
                else
                    mapping[v0] = 1.0 * mapping[v1]
                    mapping[v0][index] += 1
                end
            end
        end
    end
    @assert length(mapping) == V # all vertices have been assigned a coordinate
    return hcat([mapping[v] for v in c.vertices]...)
end

"""
    coordinate_matrix(num_vertices::Int)
Returns the `T` matrix described in section IIC of arXiv:1810.11510 for a
circuit with `num_vertices`, where the spanning tree is depth-1 and rooted at
the first vertex of the circuit.
"""
coordinate_matrix(num_vertices::Int) = [spzeros(num_vertices - 1) I]

"""
    coordinate_matrix(num_vertices::Int)
Returns the `T` matrix described in section IIC of arXiv:1810.11510 for a
circuit `c`, where the spanning tree is depth-1 and rooted at the first vertex
of the circuit.
"""
coordinate_matrix(c::Circuit) = coordinate_matrix(length(c.vertices))

#######################################################
# cascade and unite
#######################################################

"""
    cascade(circs::AbstractVector{<:Circuit})
    cascade(circs::Circuit...)
Cascade the circuits into one larger block diagonal circuit.
"""
function cascade(circs::AbstractVector{<:Circuit})
    circuit_matrices = [cat(m..., dims=(1,2))
        for m in zip([matrices(circ) for circ in circs]...)]
    vertices = vcat([circ.vertices for circ in circs]...)
    return Circuit(circuit_matrices..., vertices)
end
cascade(circs::Circuit...) = cascade(collect(circs))

"""
    unite_vertices(circ::Circuit{T}, vertices::AbstractVector{T}) where T
    unite_vertices(circ::Circuit{T}, vertices::T...) where T
Unite the given vertices of a circuit.
"""
function unite_vertices(circ::Circuit{T}, vertices::AbstractVector{T}) where T
    length(vertices) <= 1 && return circ
    n = length(circ.vertices)
    vertex_inds = [findfirst(isequal(v), circ.vertices) for v in vertices]
    keep_inds = filter(!in(vertex_inds[2:end]), 1:n)
    function unite_indices(mat::AbstractMatrix)
        m = copy(mat)
        ind1 = vertex_inds[1]
        for ind2 in vertex_inds[2:end]
            m[ind1, :] += mat[ind2, :]
            m[:, ind1] += mat[:, ind2]
        end
        m = m[keep_inds, keep_inds]
        for i in 1:size(m, 1)
            m[i, i] = 0
        end
        return m
    end
    circuit_matrices = [unite_indices(mat) for mat in matrices(circ)]
    return Circuit(circuit_matrices..., circ.vertices[keep_inds])
end
unite_vertices(circ::Circuit{T}, vertices::T...) where T =
    unite_vertices(circ, collect(vertices))

"""
    connect(models::AbstractVector{<:Circuit})
    connect(models::Circuit...)
Cascade all models and unite ports with the same name. This results in a
combined circuit with common ports connected.
"""
function connect(circs::AbstractVector{<:Circuit})
    @assert !isempty(circs)
    if length(circs) == 1
        return first(circs)
    end
    # number the vertices so that the names are all distinct and then cascade
    vertex_number = 1
    function rename(circ)
        vertices = [(vertex_number + i - 1, vertex)
            for (i, vertex) in enumerate(circ.vertices)]
        vertex_number += length(vertices)
        return Circuit(matrices(circ)..., vertices)
    end
    circ = cascade(map(rename, circs))
    # merge all vertices with the same name
    original_vertices = vcat([c.vertices for c in circs]...)
    for vertex in unique(original_vertices)
        inds = findall([v[2] == vertex for v in circ.vertices])
        circ = unite_vertices(circ, circ.vertices[inds])
    end
    # remove numbering
    return Circuit(matrices(circ)..., [v[2] for v in circ.vertices])
end
connect(circs::Circuit...) = connect(collect(circs))


# This is test of the code for computing a transit model integrated over a time step.
using PyPlot

include("../src/integrate_transit.jl")

# Test it out:

include("../test/loglinspace.jl")

t1 = -1.2; t2 = 1.2; nt = 1000; dt = (t2-t1)/nt*5
t = linearspace(t1,t2,nt)
favg0 = zeros(nt)
favg1 = zeros(nt); depth1 = zeros(Int64,nt); diff1 = zeros(nt)
favg2 = zeros(nt); depth2 = zeros(Int64,nt); diff2 = zeros(nt)
favg3 = zeros(nt); depth3 = zeros(Int64,nt); diff3 = zeros(nt)
favg4 = zeros(nt); depth4 = zeros(Int64,nt); diff4 = zeros(nt)
# The following compares different tolerances and maxdepths:
r = 0.1; b0 = 0.5; u_n = [0.2,0.2,0.2,0.2,0.2]
trans = transit_init(r,b0,u_n,false)
param = [0.0,1.0,b0]
for i=1:nt
  trans.b = sqrt(param[3]^2+(param[2]*(t[i]-param[1]))^2)
  favg0[i] = transit_poly_d!(trans)
  favg1[i],depth1[i],diff1[i] = integrate_timestep_track(param,trans,t[i],dt,1e-5*r^2,4)
  favg2[i],depth2[i],diff2[i]  = integrate_timestep_track(param,trans,t[i],dt,1e-5*r^2,256)
  favg3[i],depth3[i],diff3[i]  = integrate_timestep_track(param,trans,t[i],dt,1e-10*r^2,4)
  favg4[i],depth4[i],diff4[i]  = integrate_timestep_track(param,trans,t[i],dt,1e-10*r^2,256)
end
favg1 /= dt
favg2 /= dt
favg3 /= dt
favg4 /= dt
# The depth arrays can be plotted to show how the maximum depth used varies across
# the light curve

# Now, carry out speed test:

@time for i=1:nt
        trans.b = sqrt(param[3]^2+(param[2]*(t[i]-param[1]))^2)
        favg0[i] = transit_poly_d!(trans)
      end
dtinv = inv(dt)
@time for i=1:nt
        favg1[i] = integrate_timestep(param,trans,t[i],dt,1e-5*r^2,8)*dtinv
      end

# Estimate the difference between time-integrated and point estimates of the flux
# by integrating the second order term in the Taylor-series expansion.

d2fdb2 = zeros(nt)
dfdb = zeros(nt)
# Initialize a BigFloat version of the transit structure:
trans_big = transit_init(big(r),big(b0),big.(u_n),true)
bigparam = [big(0.0),big(1.0),big(b0)]
dq = 1e-18
# Loop over time steps, and at each time step compute the second derivative
# with respect to b using finite difference of BigFloat for accuracy:
for i=1:nt
  trans_big.b = sqrt(bigparam[3]^2+(bigparam[2]*(t[i]-bigparam[1]))^2)
  trans_big.b -= dq
  transit_poly_d!(trans_big)
  dfdbm = trans_big.dfdrb[2]
  trans_big.b += 2dq
  transit_poly_d!(trans_big)
  dfdbp = trans_big.dfdrb[2]
  d2fdb2[i] = convert(Float64,(dfdbp-dfdbm)/(2*dq))
  dfdb[i] = convert(Float64,trans_big.dfdrb[2])
end

# Now plot the estimate of the difference:
plot(favg4-favg0)
b = sqrt.(param[3]^2+(param[2]*(t-param[1])).^2)
plot(dt^2*(d2fdb2.*(param[2]*(t-param[1])./b).^2+dfdb.*param[2]^2./b)/24.0)


function _trim(x::Array{T})::Tuple{Array{Int64},Array{Int64}} where {T<:AbstractFloat}
    ids = sortperm(x)
    distances = [x[ids[i+1]] - x[ids[i]] for i=1:length(x)-1]

    if length(x) < 3
        return Array{Int64, 1}(0), collect(1:length(x))
    end

    density = zeros(length(x))
    for i in eachindex(density)
        if i == 1
            density[i] = 0.25 / distances[1]
        elseif i == length(density)
            density[i] = 0.25 / distances[end]
        else
            density[i] = 1.0 / (distances[i-1] + distances[i])
        end
    end
    #normalize density
    density /= sum(density)

    #remove tails -> trim to 1σ
    cutoff = (1 - 0.6827) * 0.5

    prob = 0.0

    start = 1
    stop = length(x)

    for i in eachindex(x)
        prob += density[i]
        if prob >= cutoff
            start = max(1, i - 1)
            break
        end
    end

    j = length(x)
    prob = 0.0
    for i in eachindex(x)
        prob += density[j]
        if prob >= cutoff
            stop = min(length(x), j + 1)
            break
        end
        j -= 1
    end

    sort!(ids[[1:start..., stop:length(x)...]]), sort!(ids[[start+1:stop-1...]])
end

function trim(x::Array)
    deleteat!(x, _trim(x)[1])
    x
end

function trim(
    res::IntermediateResults{T},
    dotrimming::Bool)::Array{Int64, 1} where {T<:AbstractFloat}

    deleteids, remainingids = _trim(res.integrals)
    @debug "Trimming integration results: $(length(deleteids)) entries out of $(length(res.integrals)) deleted"

    if dotrimming
        deleteat!(res.integrals, deleteids)
        deleteat!(res.volumeID, deleteids)
        res.Y = res.Y[:, remainingids]
    end

    deleteids
end

function split_dataset(dataset::DataSet{T, I})::Tuple{DataSet{T, I}, DataSet{T, I}} where {T<:Real, I<:Integer}
    N = dataset.N
    n = floor(Int64, N / 2)

    ds1 = DataSet(dataset.data[:, 1:n], dataset.logprob[1:n], dataset.weights[1:n], dataset.nsubsets, dataset.subsetsize)
    ds2 = DataSet(dataset.data[:, n+1:N], dataset.logprob[n+1:N], dataset.weights[n+1:N], dataset.nsubsets, dataset.subsetsize)

    return ds1, ds2
end


struct Eigenstate{T<:Real,W<:AbstractWfn{T}}
    wfn::W
    E::T
end

function Eigenstate(ee::LinAlg.Eigen, n::Int, WfnType::Type{<:AbstractWfn}, args...)
    n in eachindex(ee[:values]) || error("Requested eigenstate of invalid n quantum number")
    wfn = WfnType(vec2mat(Complex, ee[:vectors][:,n]), args...)
    normalize!(wfn)
    Eigenstate(wfn, ee[:values][n])
end

Eigenstate(H::AbstractMatrix, n, WfnType, args...) =
    Eigenstate(eigfact(H), n, WfnType, args...)

function Eigenstate(H::AbstractLinearMap, n::Int, WfnType::Type{<:AbstractWfn}, args...)
    E, φ = eigs(H, which = :SR, nev = n)
    wfn = WfnType(vec2mat(Complex, φ[:,n]), args...)
    normalize!(wfn)
    Eigenstate(wfn, E[n])
end

export Eigenstate


typealias Constraint Vector{Int}

"""
    computescaleandshiftfromboundingbox(boundingBox::Vector{Float64})

Computes the scale and horizontal/vertical shift from a given bounding box.
This establishes an aspect-ratio-conserving transformation from the given
bounding box into a square that fits comfortably into the outer bounding triangle
of the DelaunayMeshes. The square (blue) fills ten per cent of the maximally
inscribed square (dashed).

![inscribed](inscribedsquare.svg)

Specifically, if ``(c_x, c_y)`` denote the center of the bounding box and
``\Delta = \max((x_\mathrm{max} - x_\mathrm{min}), (x_\mathrm{max} - x_\mathrm{min}))``
is the maximum extent of the bounding box, then the coordinate transformation to
the new primed coordinates is given by
``(x', y') = 0.1 (x - c_x, y - c_y)/\Delta + (c_x', c_y')``,
where ``(c_x', c_y') = (1.5, 1.25)`` denotes the center of the inscribed square.
"""
function computescaleandshiftfromboundingbox(boundingBox::Vector{Float64})
  local scale = max(boundingBox[2]-boundingBox[1], boundingBox[4]-boundingBox[3])
  local cx = 0.5*(boundingBox[2]+boundingBox[1])
  local cy = 0.5*(boundingBox[4]+boundingBox[3])
  return scale, cx, cy
end

"""
Defines a mesh. A mesh consists of a bounding box providing bounds to the
vertex coordinates, constraints consisting of closed polygons and the
corresponding Delaunay tesselation.
"""
type Mesh
   boundingBox::Nullable{Vector{Float64}}
   scale::Nullable{Float64}
   shiftX::Nullable{Float64}
   shiftY::Nullable{Float64}

   tesselation::DelaunayMeshes.DelaunayTesselation
   constraints::Vector{Constraint}
   faceLocation::Vector{Bool}
   voronoiVertices::Nullable{Array{Float64, 2}}

   Mesh() = new(
     Nullable{Vector{Float64}}(),
     Nullable{Vector{Float64}}(),
     Nullable{Vector{Float64}}(),
     Nullable{Vector{Float64}}(),
     DelaunayMeshes.DelaunayTesselation(),
     Vector{Constraint}(),
     [true, false],
     Nullable{Array{Float64, 2}}()
   )
end

"""
    getvoronoivertices(mesh::Mesh)

Computes the Voronoi vertices of the given DelaunayMeshes.

# Remarks
* The results are cached but need to be re-computed whenever a new vertex or
  a new constraint is pushed into the DelaunayMeshes.
"""
function getvoronoivertices(mesh::Mesh)
  if (!isnull(mesh.voronoiVertices))
    return mesh.voronoiVertices
  end

  mesh.voronoiVertices = fill(Inf, size(mesh.tesselation.faces, 1), 2)

  # go through all dual edges and compute barycenter of corresponding face
  for ei in 2:2:length(mesh.tesselation.edges)
    local curFace = mesh.tesselation.edges[ei].origin
    if get(mesh.voronoiVertices)[curFace, :] == [Inf, Inf]
      local curFaceEdge = rot(ei)
      local curFaceVertexIndices =
       lface_vertices(mesh.tesselation, curFaceEdge)
      local curFaceVertexCoordinates =
       map((i::Int) -> mesh.tesselation.vertices[i], curFaceVertexIndices)
      local baryCenter = sum(curFaceVertexCoordinates)/3.
      get(mesh.voronoiVertices)[curFace, :] = baryCenter
    end
  end

  return get(mesh.voronoiVertices)
end

function setboundingbox(mesh::Mesh, boundingBox::Vector{Float64})
  @assert isnull(mesh.boundingBox)
  mesh.boundingBox = boundingBox
  mesh.scale, mesh.shiftX, mesh.shiftY =
   computescaleandshiftfromboundingbox(get(mesh.boundingBox))
end

function scalepoints(mesh::Mesh, points::Array{Float64, 2})
  local scaledX = 0.05*(points[:, 1] - get(mesh.shiftX))/get(mesh.scale) + 1.5
  local scaledY = 0.05*(points[:, 2] - get(mesh.shiftY))/get(mesh.scale) + 1.25
  local scaledPoints = [scaledX scaledY]

  return scaledPoints
end

function unscalepoints(mesh::Mesh, points::Array{Float64, 2})
  local unScaledX = (points[:, 1]  - 1.5) * get(mesh.scale)/0.05 + get(mesh.shiftX)
  local unScaledY = (points[:, 2] - 1.25) * get(mesh.scale)/0.05 + get(mesh.shiftY)
  local unScaledPoints = [unScaledX unScaledY]

  return unScaledPoints
end

function push_scaled!(mesh::Mesh, scaledPoints::Array{Float64, 2})
  push!(mesh.tesselation, scaledPoints)

  # mark faces as interior
  local noMissingFaceLocations =
   length(mesh.tesselation.faces) - length(mesh.faceLocation)
  local truth = fill(true, noMissingFaceLocations)
  mesh.faceLocation = [mesh.faceLocation; truth ]

  # invalidate voronoi vertex cache
  mesh.voronoiVertices = Nullable{Array{Float64, 2}}()
end

"""
    push!(mesh::Mesh, points::Array{Float64, 2})

Inserts a number of vertices into the mesh.

# Remarks
* If the mesh was initialized without a bounding box, it is determined by
  the first set of vertices that are pushed into the triangualation. All
  consecutive push operations cannot insert vertices that exceed the bounding
  box limits
* New faces that are created by a push operation are marked as interior by
  default, i.e. it is forbidden to push vertices into an exterior region.
"""
function push!(mesh::Mesh, points::Array{Float64, 2})
  local minX = minimum(points[:,1])
  local maxX = maximum(points[:,1])
  local minY = minimum(points[:,2])
  local maxY = maximum(points[:,2])

  if isnull(mesh.boundingBox)
    setboundingbox(mesh, [minX, maxX, minY, maxY])
  else
    @assert (minX >= get(mesh.boundingBox)[1])
    @assert (maxX <= get(mesh.boundingBox)[2])
    @assert (minY >= get(mesh.boundingBox)[3])
    @assert (maxY <= get(mesh.boundingBox)[4])
  end

  local scaledPoints = scalepoints(mesh, points)
  push_scaled!(mesh, scaledPoints)
end

"""
    getfacesinsideregion(mesh::Mesh, ei::Int)

Finds all faces that are located inside the same region as the face identified
by `org(ei)`.

# Remarks
* The algorithm performs a breadth-first search over all adjacent triangles
  while omitting triangles that requires crossing a boundary.
"""
function getfacesinsideregion(mesh::Mesh, ei::Int)
  local foundFaces = Vector{Int}()
  local queuedEdges = [ei]
  while (length(queuedEdges) > 0)
    local currentEdge = shift!(queuedEdges)
    local currentFace = org(mesh.tesselation, currentEdge)

    # do nothing if it was already discovered
    if (currentFace ∉ foundFaces)
      push!(foundFaces, currentFace)

      # insert adjacent faces if they do not require boundary crossing
      local ec = currentEdge
      if (!isconstraint(mesh.tesselation, rot(ec)))
        push!(queuedEdges, sym(ec))
      end
      ec = onext(mesh.tesselation, ec)
      if (!isconstraint(mesh.tesselation, rot(ec)))
        push!(queuedEdges, sym(ec))
      end
      ec = onext(mesh.tesselation, ec)
      if (!isconstraint(mesh.tesselation, rot(ec)))
        push!(queuedEdges, sym(ec))
      end
    end
  end

  return foundFaces
end

"""
    addconstraint!(mesh::Mesh, vertices::Vector{Int})

Adds a constraint, given as a list of existing vertices, into the DelaunayMeshes.

# Remarks
* Constraints are closed polygons. A constraint edge from the last vertex to
  the first vertex is automatically inserted to ensure closure.
* All faces interior of the constrained region are marked as "exterior" to the
  DelaunayMeshes. A face is interior to the constrained region if it is located
  to the right of the given constraint polygon (i.e. the constraint vertices
  circle the interior region counter-clockwise).
* Constraint edges must not intersect. For performance reasons, this is not
  checked.
* Any vertex of the triangulation can intersect either zero or two constraint edges.
"""
function addconstraint!(mesh::Mesh, vertices::Vector{Int})
  # insert constraints between vertices and last one to close
  for i=1:length(vertices)-1
    insertconstraint!(mesh.tesselation, vertices[i], vertices[i+1])
  end
  local ci = insertconstraint!(mesh.tesselation, vertices[end], vertices[1])

  # find vertices to the right of the constraint
  local facesInside = getfacesinsideregion(mesh, rot(ci))
  mesh.faceLocation[facesInside] = false

  # invalidate voronoi vertex cache
  mesh.voronoiVertices = Nullable{Array{Float64, 2}}()
end

function istrianglevalid(mesh::Mesh, x::Tuple{Int, Int, Int})
  local rei = rot(x[1])
  local faceI = dest(mesh.tesselation, rei)
  return mesh.faceLocation[faceI]
end

function verticesfromtriangle(mesh::Mesh, x::Tuple{Int, Int, Int})
  local v1 = org(mesh.tesselation, x[1])
  local v2 = org(mesh.tesselation, x[2])
  local v3 = org(mesh.tesselation, x[3])

  return [v1 v2 v3]
end

"""
    convert(::Type{DiffEqPDEBase.SimpleFEMMesh}, mesh::Mesh)

Converts the given mesh to a format that can be used by the FEM-solvers.
"""
function convert(::Type{DiffEqPDEBase.SimpleFEMMesh}, mesh::Mesh)
  local allTriangles = DelaunayMeshes.getalltriangles(mesh.tesselation)
  local validTriangles = filter((x)->istrianglevalid(mesh,x), allTriangles)
  local triangles = map((x)->verticesfromtriangle(mesh,x), validTriangles)

  local redTriangles = reduce(vcat, triangles)
  local redVertices = reduce(hcat, mesh.tesselation.vertices)'

  return DiffEqPDEBase.SimpleFEMMesh(mesh.tesselation.vertices, redTriangles)
end


immutable Conv2dParameters
    x_max :: Float64
    y_max :: Float64
    filter :: Function
    filter_dx :: Function
    filter_dy :: Function
    pix_size_x :: Float64
    pix_size_y :: Float64
    guess_error_x :: Float64
    guess_error_y :: Float64
end
immutable Conv2d <: SuperRes
    x_max :: Float64
    y_max :: Float64
    filter :: Function
    filter_dx :: Function
    filter_dy :: Function
    eval_grid_x :: Array{Float64}
    eval_grid_y :: Array{Float64}
    guess_grid_x :: Array{Float64}
    guess_grid_y :: Array{Float64}
    n_x :: Int64
    n_y :: Int64
    function Conv2d(params :: Conv2dParameters)
        eval_grid_x_tmp = 0:params.pix_size_x:params.x_max
        eval_grid_y_tmp = 0:params.pix_size_y:params.y_max
        eval_grid_x = [x for x in eval_grid_x_tmp, y in eval_grid_y_tmp][:]
        eval_grid_y = [y for x in eval_grid_x_tmp, y in eval_grid_y_tmp][:]
        guess_grid_x_tmp = 0:params.guess_error_x:params.x_max
        guess_grid_y_tmp = 0:params.guess_error_y:params.y_max
        guess_grid_x = [x for x in guess_grid_x_tmp, y in guess_grid_y_tmp][:]
        guess_grid_y = [y for x in guess_grid_x_tmp, y in guess_grid_y_tmp][:]
        return new(params.x_max, params.y_max,
                   params.filter, params.filter_dx, params.filter_dy,
                   eval_grid_x, eval_grid_y,
                   guess_grid_x, guess_grid_y,
                   length(eval_grid_x_tmp), length(eval_grid_y_tmp))
    end
end
function psi(model :: Conv2d, theta :: Vector{Float64})
    # This function computes the direct problem for a single point theta
    return Float64[model.filter(model.eval_grid_x[i] - theta[1], model.eval_grid_y[i] - theta[2]) for i in 1:length(model.eval_grid_x)]
end

function dpsi(model :: Conv2d, theta :: Vector{Float64})
    # This function computes the gradient of psi and adds it to eval_array_grad field
    dx = Float64[- model.filter_dx(model.eval_grid_x[i] - theta[1], model.eval_grid_y[i] - theta[2]) for i in 1:length(model.eval_grid_x)]
    dy = Float64[- model.filter_dy(model.eval_grid_x[i] - theta[1], model.eval_grid_y[i] - theta[2]) for i in 1:length(model.eval_grid_x)]
    return [dx dy]
end

function phi(model :: Conv2d, parameters :: Matrix{Float64}, weights :: Vector{Float64})
    # Evaluates the direct problem
    return sum([weights[i] * psi(model, vec(parameters[:, i])) for i in 1:length(weights)])
end

function getStartingPoint(model :: Conv2d, v :: Vector{Float64})
    # This function gets an initial value for the location
    guess_array = map((x,y) -> dot(v, psi(model, [x;y])), model.guess_grid_x, model.guess_grid_y)
    ind = indmin(guess_array);
    res = [model.guess_grid_x[ind]; model.guess_grid_y[ind]]
    return res
end

parameterBounds(model :: Conv2d) = [0.0; 0.0], [model.x_max; model.y_max] # Bounds for the parameters
dim(model :: Conv2d) = 2

type DynamicConv2dParameters
    K :: Int64
    tau :: Float64
    v_max :: Float64
    num_v :: Int64
end

type DynamicConv2d<: SuperRes
    static :: Conv2d
    v_max :: Float64
    eval_grid_x :: Array{Float64}
    eval_grid_y :: Array{Float64}
    eval_grid_t :: Array{Float64}
    guess_grid_x :: Array{Float64}
    guess_grid_y :: Array{Float64}
    guess_grid_vx :: Array{Float64}
    guess_grid_vy :: Array{Float64}
    eval_array :: Vector{Float64}
    eval_array_grad :: Array{Float64}
    function DynamicConv2d(static :: Conv2d, params::DynamicConv2dParameters)
        times = linspace(-params.K*params.tau, params.K*params.tau, 2*params.K + 1)
        eval_grid_x = [x for x in static.eval_grid_x, t in times][:]
        eval_grid_y = [y for y in static.eval_grid_y, t in times][:]
        eval_grid_t = [t for x in static.eval_grid_y, t in times][:]
        velocities = linspace(-params.v_max, params.v_max, params.num_v)
        guess_grid_x = [x for x in static.guess_grid_x, vx in velocities, vy in velocities][:]
        guess_grid_y = [y for y in static.guess_grid_y, vx in velocities, vy in velocities][:]
        guess_grid_vx = [vx for y in static.guess_grid_y, vx in velocities, vy in velocities][:]
        guess_grid_vy = [vy for y in static.guess_grid_y, vx in velocities, vy in velocities][:]
        return new(static, params.v_max,
                   eval_grid_x, eval_grid_y, eval_grid_t,
                   guess_grid_x, guess_grid_y, guess_grid_vx, guess_grid_vy)
    end
end
function psi(model :: DynamicConv2d, theta :: Vector{Float64}, w = 1.0, acc=false)
    # This function computes the direct problem for a single point theta
    return Float64[model.static.filter(model.eval_grid_x[i] - theta[1] - theta[3]*model.eval_grid_t[i],
                                       model.eval_grid_y[i] - theta[2] - theta[4]*model.eval_grid_t[i]) for i in 1:length(model.eval_grid_x)]

end

function dpsi(model :: DynamicConv2d, theta :: Vector{Float64})
    # This function computes the gradient of psi and adds it to eval_array_grad field
    dx =  Float64[- model.static.filter_dx(model.eval_grid_x[i] - theta[1] - theta[3]*model.eval_grid_t[i],
                                           model.eval_grid_y[i] - theta[2] - theta[4]*model.eval_grid_t[i]) for i in 1:length(model.eval_grid_x)]
    dy =  Float64[- model.static.filter_dy(model.eval_grid_x[i] - theta[1] - theta[3]*model.eval_grid_t[i],
                                           model.eval_grid_y[i] - theta[2] - theta[4]*model.eval_grid_t[i]) for i in 1:length(model.eval_grid_x)]
    dvx =  Float64[-model.eval_grid_t[i] * model.static.filter_dx(model.eval_grid_x[i] - theta[1] - theta[3]*model.eval_grid_t[i],
                                                                  model.eval_grid_y[i] - theta[2] - theta[4]*model.eval_grid_t[i]) for i in 1:length(model.eval_grid_x)]
    dvy =  Float64[-model.eval_grid_t[i] * model.static.filter_dy(model.eval_grid_x[i] - theta[1] - theta[3]*model.eval_grid_t[i],
                                                                  model.eval_grid_y[i] - theta[2] - theta[4]*model.eval_grid_t[i]) for i in 1:length(model.eval_grid_x)]
    return [dx dy dvx dvy]
end

function phi(model :: DynamicConv2d, parameters :: Matrix{Float64}, weights :: Vector{Float64})
    # Evaluates the direct problem
    return sum([weights[i] * psi(model, vec(parameters[:, i])) for i in 1:length(weights)])
end

function getStartingPoint(model :: DynamicConv2d, v :: Vector{Float64})
    # This function gets an initial value for the location
    guess_array = map((x,y, vx, vy) -> dot(v, psi(model, [x;y;vx;vy])), model.guess_grid_x, model.guess_grid_y, model.guess_grid_vx, model.guess_grid_vy)
    ind = indmin(guess_array);
    return vec([model.guess_grid_x[ind]; model.guess_grid_y[ind]; model.guess_grid_vx[ind]; model.guess_grid_vy[ind]])
end

parameterBounds(model :: DynamicConv2d) = [0.0; 0.0; -model.v_max; -model.v_max], [model.static.x_max; model.static.y_max; model.v_max; model.v_max] # Bounds for the parameters


mutable struct DensityMatrix{T, d} <: QuantumVariable{T}
    head::Symbol
    id_hash::UInt64
    size::Tuple{Int, Int}
    value::Convex.ValueOrNothing
    vexity::Convex.Vexity
    function DensityMatrix{T}(d::Int) where {T <: Number}
        this = new{T, d}(:DensityMatrix, 0, (d,d), nothing, Convex.AffineVexity())
        this.id_hash = objectid(this)
        Convex.id_to_variables[this.id_hash] = this
        this
    end
    DensityMatrix(d::Int) = DensityMatrix{ComplexF64}(d)
end

Convex.constraints(ρ::DensityMatrix) = [ ρ ⪰ 0, tr(ρ) == 1 ]
Convex.sign(::DensityMatrix) = Convex.ComplexSign()
Convex.vartype(::DensityMatrix) = Convex.ContVar


using LinearAlgebra
using StaticArrays

abstract type AbstractUnitCell end
abstract type AbtractPeriodicCell <: AbstractUnitCell end
abstract type AbstractOrthorombicCell <: AbtractPeriodicCell end

struct NonPeriodic <: AbstractUnitCell end

struct VariableCell <: AbtractPeriodicCell
    cell::Vector{AbtractPeriodicCell}
end

struct CubicCell <: AbstractOrthorombicCell
    a::Float64
end

struct OrthorombicCell <: AbstractOrthorombicCell
    abc::SVector{3,Float64}
    OrthorombicCell(a::Real,b::Real,c::Real) = new([a,b,c])
    OrthorombicCell(a::Real) = new([a,a,a])
    function OrthorombicCell(a::AbstractVector)
        @assert length(a) == 3
        new(a)
    end
end

struct TriclinicCell <: AbtractPeriodicCell
    abc::SMatrix{3,3,Float64}
    function TriclinicCell(abc::AbstractMatrix)
        @assert size(abc) == (3,3)
        new(abc)
    end
    function TriclinicCell(a::AbstractVector)
        @assert length(a) == 3
        new(Diagonal(a))
    end
    function TriclinicCell(a::Real, b::Real, c::Real, α::Real, β::Real, γ::Real)
        tmp = zeros(3,3)
        tmp[1,1] = sind(γ)*a
        tmp[2,1] = cosd(γ)*a
        tmp[2,2] = b
        tmp[2,3] = cosd(α)*c
        tmp[1,3] = c*(cosd(β)-cosd(γ)*cosd(α))/sind(γ)
        tmp[3,3] = sqrt(c^2-tmp[1,3]^2-tmp[2,3]^2)
        new(tmp)
    end
end

parsecell(a::Number) = CubicCell(a)
function parsecell(a::Real, b::Real, c::Real)
    a == b == c && return CubicCell(a)
    return OrthorombicCell(a,b,c)
end

function parsecell(a::Real, b::Real, c::Real, α::Real, β::Real, γ::Real)
    if α == β == γ == 90
        a == b == c && return CubicCell(a)
        return OrthorombicCell(a,b,c)
    else
        return TriclinicCell(a,b,c,α,β,γ)
    end
end

Base.show(io::IO, c::CubicCell) = print(io, "CubicCell a=",c.a)
Base.show(io::IO, c::OrthorombicCell) = print(io, "OrthorombicCell")
Base.show(io::IO, c::TriclinicCell) = print(io, "TriclinicCell")
Base.show(io::IO, c::VariableCell) = print(io, "Variable cell of length ", length(c))

Base.length(c::VariableCell) = length(c.cell)
Base.getindex(c::VariableCell, i) = c.cell[i]
Base.lastindex(c::VariableCell) = length(c)
Base.firstindex(c::VariableCell) = 1

Base.iterate(c::VariableCell, state=1) = state > length(c) ? nothing : (c[state], state+1)

celldiag(c::AbtractPeriodicCell) = diag(cellmatrix(c))

cellmatrix(c::CubicCell) = Diagonal([c.a, c.a, c.a])
cellmatrix(c::OrthorombicCell) = Diagonal(c.abc)
cellmatrix(c::TriclinicCell) = c.abc
cellmatrix(c::VariableCell) = [ cellmatrix(x) for x in c]
cellvectorlengths(c::AbtractPeriodicCell) = [norm(x) for x in eachcol(cellmatrix(c))]
cellvectorlengths(c::AbstractOrthorombicCell) = (Diagonal ∘ cellmatrix)(c)
cellvectorlengths(c::VariableCell) = [ cellvectorlengths(x) for x in c ]

volume(c::CubicCell) = c.a^3
volume(c::OrthorombicCell) = prod(c.abc)
volume(c::AbtractPeriodicCell) = (abs ∘ det ∘ cellmatrix)(c)
volume(c::VariableCell) = [ volume(x) for x in c ]


using Gillespie

function F_dd(x,parms)
    (S1,S2,S3) = x
    (c1,c2,c3,c4) = parms
    [c1*S1,c2*S1*S1,c3*S2,c4*S2]
end

x0 = [10000,0,0]
nu = [[-1 0 0];[-2 1 0];[2 -1 0];[0 -1 1]]
parms = [1.0,0.002,0.5,0.04]
tf = 10.0
srand(1234)

result = ssa(x0,F_dd,nu,parms,tf)

data = ssa_data(result)


function solve()
    n, m = [parse(Int, x) for x in split(readline())]
    as = [parse(Int, x) for x in split(readline())]
    edges = Dict{Int, Set{Int}}()
    for i=1:n
        edges[i] = Set{Int}()
    end
    for _=1:m
        x, y = [parse(Int, x) for x in split(readline())]
        push!(edges[x], y)
    end

    dp = ones(Int, n) * 10^10
    for i=1:n
        for j in edges[i]
            dp[j] = min(dp[j], dp[i], as[i])
        end
    end
    maximum((as - dp)[2:end])
end

println(solve())


module TestUtils

using TowerFields.BLS12381Field: Q

function randfq()
    # number of UInt128 elements needed
    uint128bitsize = sizeof(UInt128) * 8
    n = Int(cld(log2(Q), uint128bitsize))
    bigrand = zero(BigInt)
    for r in rand(UInt128, n)
        bigrand = (bigrand << uint128bitsize) | r
    end
    return mod(bigrand, Q)
end

randfq2() = (randfq(), randfq())
randfq3() = (randfq(), randfq(), randfq())
randfq6() = (randfq2(), randfq2(), randfq2())
randfq12() = (randfq6(), randfq6())

end


# Statistical utilities.

# Ambegaokar, V., & Troyer, M. (2010). Estimating errors reliably in Monte
# Carlo simulations of the Ehrenfest model. American Journal of Physics, 78(2),
# 150-157.

"""
    stderr_bin(dt::Float64, xs::AbstractVector)

Compute the standard error and autocorrelation time of `xs` by binning.

Stops when there are too few bins left. Returns the largest error attained and
the corresponding autocorrelation time in units of `dt`.
"""
function stderr_bin(dt::Float64, xs::AbstractVector)
    stderr_full = std(xs) / sqrt(length(xs))
    stderr_max = stderr_full

    while length(xs) >= 60
        xs = [0.5 * (xs[i] + xs[i+1]) for i in (1+length(xs)%2):2:length(xs)]
        stderr_max = max(stderr_max, std(xs) / sqrt(length(xs)))
    end

    autocorrelation_steps = 0.5 * ((stderr_max / stderr_full)^2 - 1)

    stderr_max, autocorrelation_steps * dt
end

"""
    aggregate(dt::Float64, xs::AbstractVector)

Perform a statistical analysis of `xs` for time step `dt`.
"""
function aggregate(dt::Float64, xs::AbstractVector)
    mean(xs), stderr_bin(dt, xs)...
end


using StructJuMP

m = StructuredModel()

factories = 1:3
centers = 1:5
scenarios = ["lo", "mid", "hi"]

cap = [500, 450, 650]

demand = [150 160 170;
          100 120 135;
          250 270 300;
          300 325 350;
          600 700 800]

prob = [0.25, 0.5, 0.25]

transcost = [2.49 5.21 3.76 4.85 2.07;
             1.46 2.54 1.83 1.86 4.76;
             3.26 3.08 2.60 3.76 4.45]

prodcost = 14
price = 24
wastecost = 4

@variable(m, ship[factories,centers] >= 0)
@variable(m, product[factories] >= 0)
@variable(m, sales[centers] >= 0)
@variable(m, waste[centers] >= 0)
@variable(m, profit)

@variable(m, received[centers] >= 0)
@objective(m, Max, -sum(transcost[i,j]*ship[i,j] for i=factories, j=centers) + sum(prodcost*product[i] for i=factories))
for j in centers
    @constraint(m, received[j] == sum(ship[i,j] for i=factories))
end
for i in factories
    @constraint(m, product[i] == sum(ship[i,j] for j=centers))
end

for (s, elem) in enumerate(scenarios)
    bl = StructuredModel(parent=m, id=s)
    @variable(bl, 0 <= salesw[i=centers] <= demand[i,s])
    @variable(bl, wastew[centers] >= 0)
    @objective(bl, :Max, sum(price*prob[s]*salesw[j] for j=centers) - sum(wastecost*prob[s]*wastew[j] for j=centers))
    for j in centers
        @constraint(bl, received[j] == salesw[j]+wastew[j])
    end
end


import IntervalArithmetic

"""
    IntervalAbs(lo::Real, hi::Real)

Interval abstraction for real-valued numbers.
"""
struct IntervalAbs{T <: Real}
    interval::IntervalArithmetic.Interval{T}
end

IntervalAbs{T}(x::Integer) where {T <: Integer} =
    IntervalAbs{T}(T(x), T(x)+1)
IntervalAbs{T}(x::Union{AbstractFloat,Integer}) where {T <:AbstractFloat} =
    IntervalAbs{T}(T(x), T(x))
IntervalAbs{T}(a::T, b::T) where {T <: Real} =
    IntervalAbs{T}(IntervalArithmetic.Interval{T}(a, b))

Base.copy(x::IntervalAbs) = x

Base.convert(I::Type{IntervalAbs}, x::Real) = I(x)
Base.convert(I::Type{IntervalAbs{T}}, x::Integer) where {T <: Real} = I(x)

Base.show(io::IO, x::IntervalAbs) = Base.show(io, x.interval)

lub(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval ∪ b.interval)
glb(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval ∩ b.interval)

widen(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval ∪ b.interval)

Base.:(==)(a::IntervalAbs, b::IntervalAbs) = a.interval == b.interval

Base.:+(a::IntervalAbs) = a
Base.:-(a::IntervalAbs) = IntervalAbs(-a)

Base.:+(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval + b.interval)
Base.:-(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval - b.interval)
Base.:*(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval * b.interval)
Base.:/(a::IntervalAbs, b::IntervalAbs) = IntervalAbs(a.interval / b.interval)

Base.:+(a::IntervalAbs, b::Real) = IntervalAbs(a.interval + b)
Base.:-(a::IntervalAbs, b::Real) = IntervalAbs(a.interval - b)
Base.:*(a::IntervalAbs, b::Real) = IntervalAbs(a.interval * b)
Base.:/(a::IntervalAbs, b::Real) = IntervalAbs(a.interval / b)

Base.:+(a::Real, b::IntervalAbs) = IntervalAbs(a + b.interval)
Base.:-(a::Real, b::IntervalAbs) = IntervalAbs(a - b.interval)
Base.:*(a::Real, b::IntervalAbs) = IntervalAbs(a * b.interval)
Base.:/(a::Real, b::IntervalAbs) = IntervalAbs(a / b.interval)

equiv(a::IntervalAbs, b::IntervalAbs) =
    !isdisjoint(a.interval, b.interval)
nequiv(a::IntervalAbs, b::IntervalAbs) =
    isdisjoint(a.interval, b.interval)
Base.:<(a::IntervalAbs, b::IntervalAbs) =
    !IntervalArithmetic.precedes(b.interval, a.interval)
Base.:<=(a::IntervalAbs, b::IntervalAbs) =
    !IntervalArithmetic.strictprecedes(b.interval, a.interval)

equiv(a::IntervalAbs, b::Real) = b ⊆ a.interval
nequiv(a::IntervalAbs, b::Real) = !(a.interval.lo == a.interval.hi == b)
Base.:<(a::IntervalAbs, b::Real) = a.interval.lo < b
Base.:<=(a::IntervalAbs, b::Real) = a.interval.lo <= b

equiv(a::Real, b::IntervalAbs) = a ⊆ b.interval
nequiv(a::Real, b::IntervalAbs) = !(a == b.interval.lo == b.interval.hi)
Base.:<(a::Real, b::IntervalAbs) = a < b.interval.hi
Base.:<=(a::Real, b::IntervalAbs) = a <= b.interval.hi


using Combinatorics:combinations
using LinearAlgebra: dot, norm
using Statistics:var

"""
    ABOD(k = 5,
         metric = Euclidean(),
         algorithm = :kdtree,
         leafsize = 10,
         reorder = true,
         parallel = false,
         enhanced = false)

Determine outliers based on the angles to its nearest neighbors. This implements the `FastABOD` variant described in
the paper, that is, it uses the variance of angles to its nearest neighbors, not to the whole dataset, see [1].

*Notice:* The scores are inverted, to conform to our notion that higher scores describe higher outlierness.

Parameters
----------
$_k_param

$_knn_shared

    enhanced::Bool
When `enhanced=true`, it uses the enhanced ABOD (EABOD) adaptation proposed by [2].

Examples
--------
$(_score_unsupervised("ABOD"))

References
----------
[1] Kriegel, Hans-Peter; S hubert, Matthias; Zimek, Arthur (2008): Angle-based outlier detection in high-dimensional
data.

[2] Li, Xiaojie; Lv, Jian Cheng; Cheng, Dongdong (2015): Angle-Based Outlier Detection Algorithm with More Stable
Relationships.
"""
MMI.@mlj_model mutable struct ABOD <: UnsupervisedDetector
    k::Integer = 5::(_ > 0)
    metric::DI.Metric = DI.Euclidean()
    algorithm::Symbol = :kdtree::(_ in (:kdtree, :balltree))
    leafsize::Integer = 10::(_ ≥ 0)
    reorder::Bool = true
    parallel::Bool = false
    enhanced::Bool = false
end

struct ABODModel <: Model
    # We have to store the tree to efficiently retrieve the indices to the nearest neighbors. Additionally, we have to
    # store the raw training data `X` for later angle calculations.
    X::AbstractArray
    tree::NN.NNTree
end

function fit(detector::ABOD, X::Data)::Fit
    # use tree to calculate distances
    tree = buildTree(X, detector.metric, detector.algorithm, detector.leafsize, detector.reorder)
    idxs, _ = knn_others(tree, X, detector.k)
    scores = detector.enhanced ? _eabod(X, X, idxs, detector.k) : _abod(X, X, idxs, detector.k)
    Fit(ABODModel(X, tree), scores)
end

@score function score(detector::ABOD, model::Fit, X::Data)::Result
    # TODO: We could also paralellize the abod score calculation.
    if detector.parallel
        idxs, _ = knn_parallel(model.tree, X, detector.k)
        return detector.enhanced ? _eabod(X, model.X, idxs, detector.k) : _abod(X, model.X, idxs, detector.k)
    else
        idxs, _ = NN.knn(model.tree, X, detector.k)
        return detector.enhanced ? _eabod(X, model.X, idxs, detector.k) : _abod(X, model.X, idxs, detector.k)
    end
end

function _abod(X::AbstractArray, Xtrain::AbstractArray, idxs::AbstractVector, k::Int)::Scores
    # Calculate the ABOF for all instances in X.
    scores = Vector{Float64}(undef, length(idxs))
    for i in eachindex(idxs)
        @inbounds scores[i] = _abof(X[:, i], idxs[i], Xtrain, k)
    end
    scores
end

function _abof(p::AbstractVector, idxs::AbstractVector, X::AbstractArray, k::Int)::Real
    # Calculate the angle-based outlier factor (ABOF). The ABOF is the variance over the angles between the difference
    # vectors of a point `p` to all pairs of points in its nearest neighbors weighted by the distance of the points.

    # all two-neighbor combinations
    combs = combinations(idxs, 2)
    # we know that there are binomial(k, 2) results for all two-neighbor combinations and can thus pre-allocate
    result = Vector{Float64}(undef, binomial(k, 2))
    for (i, (idx1, idx2)) in enumerate(combs)
        neighbor1 = p .- X[:, idx1]
        neighbor2 = p .- X[:, idx2]
        @inbounds result[i] = dot(neighbor1, neighbor2) / (norm(neighbor1)^2 * norm(neighbor2)^2)
    end
    # NaN means that at least one norm was zero, we use -1, because higher scores should describe outlierness
    -1 * var(Iterators.filter(!isnan, result))
end

function _eabod(X::AbstractArray, Xtrain::AbstractArray, idxs::AbstractVector, k::Int)::Scores
    # Calculate the EABOF for all instances in X.
    scores = Vector{Float64}(undef, length(idxs))
    for i in eachindex(idxs)
        @inbounds scores[i] = _eabof(X[:, i], idxs[i], Xtrain, k)
    end
    scores
end

function _eabof(p::AbstractVector, idxs::AbstractVector, X::AbstractArray, k::Int)::Real
    # Calculate the enhanced angle-based outlier factor (EABOF).

    # all two-neighbor combinations
    combs = combinations(idxs, 2)
    # we know that there are binomial(k, 2) results for all two-neighbor combinations and can thus pre-allocate
    result = Vector{Float64}(undef, binomial(k, 2))
    for (i, (idx1, idx2)) in enumerate(combs)
        neighbor1 = p .- X[:, idx1]
        neighbor2 = p .- X[:, idx2]
        norm1 = norm(neighbor1)
        norm2 = norm(neighbor2)
        @inbounds result[i] = (1 / (norm1 + norm2)) * (dot(neighbor1, neighbor2) / (norm1^2 * norm2^2))
    end
    # NaN means that at least one norm was zero, we use -1, because higher scores should describe outlierness
    -1 * var(Iterators.filter(!isnan, result))
end


using Point2DTools
using SVG

function conv_svg(fname, points, upper, lower)
    cs = map(p->Circle(p.x, p.y, 2), points)
    hull = [Polyline(points[upper]; style=Style(;strokecolor="green")), Polyline(points[lower]; style=Style(;strokecolor="red"))]
    write("$fname.html", SVG.Svg(vcat(cs, hull)), 500, 500; inhtml = true)
end

function rand_svg()
    points(n=16) = 500 * rand(Point2D, n)
    ps = collect(points())
    conv_svg("rand", ps, convexperm(ps)...)
end

function rpoints_svg()
    rpoints = Point2D[[408, 22], [424, 65], [484, 416], [283, 161], [98, 358], [126, 288], [309, 56], [63, 410], [221, 8] ,[71, 378] ,[73, 138], [305, 105], [397, 475], [376, 212], [292, 304], [498, 303]])
    conv_svg("rpoints", rpoints, convexperm(rpoints)...)
end


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

# Compute X <- a X + b Y.
function matfun_axpby!(X,a,b,Y)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:m
        @inbounds for j=1:n
            if (b isa ValueOne)
                X[i,j]+=Y[i,j]
            else
                X[i,j]+=b*Y[i,j]
            end
        end
    end
end

@inline function exp_bbcs_m4(A)
    T=promote_type(eltype(A),ComplexF64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_bbcs_m4!(A_copy)
end

@inline function exp_bbcs_m4!(A)
    T=promote_type(eltype(A),ComplexF64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=6
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: B2 B3 Ba4 Bb4 B4 Ba5 Bb5 B5 T2k7
    # Computing B2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing B3 with operation: mult
    mul!(memslots3,memslots2,memslots1)
    # Computing Ba4 = x*A+x*B2+x*B3
    coeff1=0.0 + 0.13340427306445612im
    coeff2=0.020226020298183107 + 0.0im
    coeff3=-0.0 - 0.00674638241111651im
    memslots4 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots3
    # Computing Bb4 = x*A+x*B2+x*B3
    coeff1=0.0 + 0.13340427306445612im
    coeff2=0.020226020298183107 + 0.0im
    coeff3=-0.0 - 0.00674638241111651im
    memslots5 .= coeff1.*memslots1 .+ coeff2.*memslots2 .+ coeff3.*memslots3
    # Computing B4 with operation: mult
    mul!(memslots6,memslots4,memslots5)
    # Deallocating Ba4 in slot 4
    # Deallocating Bb4 in slot 5
    # Computing Ba5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=2.6958430691533257 + 0.0im
    coeff2=0.0 + 0.05272871327381115im
    coeff3=-0.09896214548845832 + 0.0im
    coeff4=0.0 + 0.007295441446830946im
    coeff5=1.0 + 0.0im
    memslots4 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots6
    mul!(memslots4,true,I*coeff1,true,true)
    # Computing Bb5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=2.6958430691533257 + 0.0im
    coeff2=0.0 - 1.3591092616886926im
    coeff3=-0.09896214548845832 + 0.0im
    coeff4=0.0 + 0.015964794632994668im
    coeff5=1.0 + 0.0im
    # Smart lincomb recycle B4
    memslots6 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots6
    mul!(memslots6,true,I*coeff1,true,true)
    # Computing B5 with operation: mult
    mul!(memslots5,memslots4,memslots6)
    # Deallocating Ba5 in slot 4
    # Deallocating Bb5 in slot 6
    # Computing T2k7 = x*I+x*A+x*B2+x*B3+x*B5
    coeff1=-6.267569853502023 + 0.0im
    coeff2=0.0 + 2.521796947120981im
    coeff3=0.05786296656487002 + 0.0im
    coeff4=0.0 - 0.0776668640807187im
    coeff5=1.0 + 0.0im
    # Smart lincomb recycle A
    memslots1 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots5
    mul!(memslots1,true,I*coeff1,true,true)
    # Deallocating B2 in slot 2
    # Deallocating B3 in slot 3
    # Deallocating B5 in slot 5
    return memslots1 # Returning T2k7
end


using PyPlot, Glob, Printf

include("ConvergenceTests/ConvergenceTests.jl")

defaultcolors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
removespine(side) = gca().spines[side].set_visible(false)
removespines(sides...) = [removespine(side) for side in sides]

u = ConvergenceTests.ForcedFlowFixedSlip.u
ξ = ConvergenceTests.ForcedFlowFixedSlip.ξ
ξ′ = ConvergenceTests.ForcedFlowFixedSlip.ξ′
f = ConvergenceTests.ForcedFlowFixedSlip.f
fₓ = ConvergenceTests.ForcedFlowFixedSlip.fₓ

filenameses = [
               glob("data/forced_fixed_slip_xy*Δt6.0e-06.jld2"),
               glob("data/forced_fixed_slip_xy*Δt6.0e-07.jld2"),
               glob("data/forced_fixed_slip_xy*Δt6.0e-08.jld2"),
              ]

Δt = [
      6.0e-06,
      6.0e-07,
      6.0e-08,
     ]

p̆(y) = - 1/4 * y^4 + 1/3 * y^3 + 143/144 * y^2 - 35/18 * y + 1/2 - 289/72 * cosh(y)/sinh(1)
p̃(y) =   1/4 * y^4 - 1/3 * y^3 - 575/144 * y^2 + 71/18 * y + 2   - 721/72 * cosh(y)/sinh(1) + 4 * cosh(y-1)/sinh(1)
p̂(y) = - 10/4 * y^3 - 15/4 * y + 90/16 * cosh(2y)/sinh(2)

p(x, y, t) = p̂(y) * cos(2*(x - ξ(t))) + p̆(y) * ξ′(t) * cos(x - ξ(t)) + p̃(y) * sin(x - ξ(t))

close("all")
fig, ax = subplots()

for (j, filenames) in enumerate(filenameses)

    errors = ConvergenceTests.compute_errors((x, y, z, t) -> u(x, y, t), filenames...)

    sizes = ConvergenceTests.extract_sizes(filenames...)

    Nx = map(sz -> sz[1], sizes)

    names = (L"(x, y)",)
    errors = (errors,)

    for i = 1:length(errors)

        @show error = errors[i]
        name = @sprintf("%s, \$ \\Delta t \$ = %.0e", names[i], Δt[j])

        L₁ = map(err -> err.L₁, error)
        L∞ = map(err -> err.L∞, error)

        @show size(L₁) size(Nx)

        loglog(Nx, L₁, color=defaultcolors[i + j - 1], alpha=0.6, mfc="None",
               linestyle="None", marker="o", label="\$L_1\$-norm, $name")

        loglog(Nx, L∞, color=defaultcolors[i + j - 1], alpha=0.6, mfc="None",
               linestyle="None", marker="^", label="\$L_\\infty\$-norm, $name")
    end

    if j == length(filenameses)
        L₁ = map(err -> err.L₁, errors[1])
        loglog(Nx, L₁[end] * (Nx[end] ./ Nx).^2, "k-", linewidth=1, alpha=0.6, label=L"\sim N_x^{-2}")
    end
end

filenames = filenameses[1]

errors = ConvergenceTests.compute_errors((x, y, z, t) -> u(x, y, t), filenames...)

sizes = ConvergenceTests.extract_sizes(filenames...)

Nx = map(sz -> sz[1], sizes)

legend()
xlabel(L"N_x")
ylabel("Norms of the absolute error, \$ | u_{\\mathrm{sim}} - u_{\\mathrm{exact}} | \$")
removespines("top", "right")
title("Convergence for forced fixed slip")
xticks(sort(Nx), ["\$ 2^{$(round(Int, log2(n)))} \$" for n in sort(Nx)])

savefig("figs/forced_fixed_slip_convergence.png", dpi=480)


using Test;

function run(fname)
    f = open(fname)
    lines = readlines(f)
    close(f)

    mask = ""
    mem_values = Dict()
    for line in lines
        parts = split(line, " = ")
        cmd, val = parts[1], parts[2]
        println("cmd = ", cmd)
        println("val = ", val)
        if cmd == "mask"
            # update the mask
            println("MASK VAL = ", val)
            mask = val
        else
            # update something in mem, applying mask
            mem_idx = parse(Int, cmd[5:length(cmd) - 1])
            num = parse(Int, parts[2])
            println("Mem_idx = ", mem_idx, " ... Num = ", num)

            # convert number to bit representation
            bits = bitstring(num)[29:64]  # want the last 36-bits
            # apply mask
            println("mask = ", mask)
            println("bits = ", bits)
            out = ""
            for (idx, m) in enumerate(mask)
                if m == 'X'
                    out *= bits[idx]
                else
                    out *= m
                end
            end
            println("new  = ", out)
            # store final number in memory
            new_val = parse(Int, out, base=2)
            println("newval = ", new_val)
            mem_values[mem_idx] = new_val
        end
    end
    total = 0
    for k in keys(mem_values)
        v = mem_values[k]
        total += v
    end

    println("[Part 1] Result of running ", fname, " was = ", total)
    return total
end

function run2(fname)
    # Bitmask modifies how we write to memory
    f = open(fname)
    lines = readlines(f)
    close(f)

    mask = ""
    mem_values = Dict()
    for line in lines
        parts = split(line, " = ")
        cmd, val = parts[1], parts[2]
        println("cmd = ", cmd)
        println("val = ", val)
        if cmd == "mask"
            # update the mask
            println("MASK VAL = ", val)
            mask = val
        else
            # update something in mem, applying mask
            mem_idx = parse(Int, cmd[5:length(cmd) - 1])
            num = parse(Int, parts[2])
            println("Mem_idx = ", mem_idx, " ... Num = ", num)

            # convert number to bit representation
            bits = bitstring(mem_idx)[29:64]  # want the last 36-bits
            # apply mask
            println("mask = ", mask)
            println("bits = ", bits)
            out = [""]

            # Generate memory address(es)
            for (idx, m) in enumerate(mask)
                if m == 'X'
                    # add both possible characters to each item
                    out1 = map((x) -> x * '0', out)
                    out2 = map((x) -> x * '1', out)
                    out = [out1..., out2...]
                elseif m == '1'
                    # overwrite with '1'
                    out = map((x) -> x * '1', out)
                elseif m == '0'
                    # don't override
                    out = map((x) -> x * bits[idx], out)
                end
            end
            # println("new  = ", out)

            # store number in memory
            for o in out
                addr = parse(Int, o, base=2)
                # println("Writing [", addr, "] => ", num)
                mem_values[addr] = num
            end
        end
    end

    total = 0
    # println(mem_values)
    for k in keys(mem_values)
        v = mem_values[k]
        total += v
    end

    println("[Part 2] Result of running ", fname, " was = ", total)
    return total
end

# @test run("14ex.txt") == 165
# run("14.txt")

@test run2("14ex.part2.txt") == 208
run2("14.txt")