"""
Calculates the negative log of the likelihood of our data given the paramters, using the Kalman filter. It uses the
predicted_observation_distributions from the kalman_filter function. The entries of this array in the second and
third columns represent the probability of the future observation of mRNA and Protein respectively, given our current knowledge.

Parameters
----------

protein_at_observations : numpy array.
    Observed protein. The dimension is m x n x 2, where m is the number of data sets, n is the
    number of observation time points. For each data set, the first column is the time,
    and the second column is the observed protein copy number at that time.

model_parameters : numpy array.
    An array containing the model parameters in the following order:
    repression_threshold, hill_coefficient, mRNA_degradation_rate,
    protein_degradation_rate, basal_transcription_rate, translation_rate,
    transcription_delay.

measurement_variance : float.
    The variance in our measurement. This is given by Sigma_e in Calderazzo et. al. (2018).

Returns
-------

log_likelihood : float.
    The log of the likelihood of the data.
"""
function calculate_log_likelihood_at_parameter_point(protein_at_observations,model_parameters,measurement_variance)
    if any(model_parameters .< 0)
        return -Inf
    end

    log_likelihood = 0
    _, _, _, _, predicted_observation_distributions, _, _ = kalman_filter(protein_at_observations,
                                                                          model_parameters,
                                                                          measurement_variance)
    observations = protein_at_observations[:,2]
    mean = predicted_observation_distributions[:,2]
    sd = sqrt.(predicted_observation_distributions[:,3])

    for observation_index in 1:length(observations)
        log_likelihood += logpdf(Normal(mean[observation_index],sd[observation_index]),observations[observation_index])
    end #for
    return log_likelihood
end # function

"""
Calculates the log of the likelihood, and the derivative of the negative log likelihood wrt each parameter, of our data given the
paramters, using the Kalman filter. It uses the predicted_observation_distributions, predicted_observation_mean_derivatives, and
predicted_observation_variance_derivatives from the kalman_filter function. It returns the log likelihood as in the
calculate_log_likelihood_at_parameter_point function, and also returns an array of the derivative wrt each parameter.

Parameters
----------

protein_at_observations : numpy array.
    Observed protein. The dimension is n x 2, where n is the number of observation time points.
    The first column is the time, and the second column is the observed protein copy number at
    that time.

model_parameters : numpy array.
    An array containing the moderowl parameters in the following order:
    repression_threshold, hill_coefficient, mRNA_degradation_rate,
    protein_degradation_rate, basal_transcription_rate, translation_rate,
    transcription_delay.

measurement_variance : float.
    The variance in our measurement. This is given by Sigma_e in Calderazzo et. al. (2018).

Returns
-------

log_likelihood : float.
    The log of the likelihood of the data.

log_likelihood_derivative : numpy array.
    The derivative of the log likelihood of the data, wrt each model parameter
"""
function calculate_log_likelihood_and_derivative_at_parameter_point(protein_at_observations,model_parameters,measurement_variance)
    number_of_parameters = size(model_parameters,1)

    _, _, _, _, predicted_observation_distributions, predicted_observation_mean_derivatives, predicted_observation_variance_derivatives = kalman_filter(protein_at_observations,
                                                                                model_parameters,
                                                                                measurement_variance)
    # calculate log likelihood as before
    log_likelihood = 0
    number_of_observations = size(protein_at_observations,1)
    observations = protein_at_observations[:,2]
    mean = predicted_observation_distributions[:,2]
    sd = sqrt.(predicted_observation_distributions[:,3])

    for observation_index in 1:length(observations)
        log_likelihood += logpdf(Normal(mean[observation_index],sd[observation_index]),observations[observation_index])
    end #for
    # now for the computation of the derivative of the negative log likelihood. An expression of this can be found
    # at equation (28) in Mbalawata, Särkkä, Haario (2013)
    observation_transform = [0.0 1.0]
    helper_inverse = 1.0./predicted_observation_distributions[:,3]

    log_likelihood_derivative = zeros(number_of_parameters)

    for parameter_index in 1:number_of_parameters
        for time_index in 1:number_of_observations
            log_likelihood_derivative[parameter_index] -= 0.5*((helper_inverse[time_index]*tr(observation_transform*predicted_observation_variance_derivatives[time_index,parameter_index,:,:]*
                                                               transpose(observation_transform))) -
                                                               (dot(helper_inverse[time_index]*transpose(observation_transform*predicted_observation_mean_derivatives[time_index,parameter_index,:,:]),
                                                               (observations[time_index] - mean[time_index]))) -
                                                               (helper_inverse[time_index]^2*(observations[time_index] - mean[time_index])^2*
                                                               dot(observation_transform,predicted_observation_variance_derivatives[time_index,parameter_index,:,:]*
                                                               transpose(observation_transform))) -
                                                               (dot(helper_inverse[time_index]*(observations[time_index] - mean[time_index]),
                                                               observation_transform*predicted_observation_mean_derivatives[time_index,parameter_index,:,:])))
         end #for
     end #for
    return log_likelihood, log_likelihood_derivative
end #function

"""
todo
"""
function log_likelihood_and_derivative_with_prior_and_transformation(protein_at_observations,proposed_position,mean_protein,measurement_variance)
    number_of_parameters = size(proposed_position,1)

    # reparameterise
    reparameterised_proposed_position = copy(proposed_position)
    reparameterised_proposed_position[[3,4,5,6]] .= exp.(reparameterised_proposed_position[[3,4,5,6]])

    # priors
    if any([pdf(Uniform(50,2*mean_protein),reparameterised_proposed_position[1]) == 0,
            pdf(Uniform(2,6),reparameterised_proposed_position[2]) == 0,
            pdf(Uniform(log(2)/150,log(2)/10),reparameterised_proposed_position[3]) == 0,
            pdf(Uniform(log(2)/150,log(2)/10),reparameterised_proposed_position[4]) == 0,
            pdf(Uniform(0.01,120),reparameterised_proposed_position[5]) == 0,
            pdf(Uniform(0.01,40),reparameterised_proposed_position[6]) == 0,
            pdf(Uniform(1,40),reparameterised_proposed_position[7]) == 0])
        return -Inf, zeros(number_of_parameters)
    end

    # quick and dirty clean up (make data start at '0')
    protein_at_observations[:,1] .-= protein_at_observations[1,1]
    log_likelihood, log_likelihood_derivative = calculate_log_likelihood_and_derivative_at_parameter_point(protein_at_observations,
                                                                                reparameterised_proposed_position,
                                                                                mean_protein,
                                                                                measurement_variance)
    for i in [3,4,5,6]
        log_likelihood_derivative[i] = reparameterised_proposed_position[i]*log_likelihood_derivative[i]
    end
    return log_likelihood, log_likelihood_derivative
end


using CVChannel
using Test

"""
This script shows the Siddhu and dephrasure channels tensored
admit multiplicative communication value.
"""

print("\n Here we verify the multiplicativity of the Siddhu and dephrasure")
print(" channel when ran in parallel. Note as a special")
print(" case this considers the erasure channel.")
@testset "Verifying Multiplicativity of Siddhu Channel with Dephrasure Channel" begin
    sid_scan_range =[0:0.1:0.5;]
    dephrasure_scan_range = [0:0.1:1;]
    sid_ctr, p_ctr, q_ctr = 1,1,1
    max_val = 0
    is_mult = true
    for s_id in sid_scan_range
        println("---Now scanning for s =", s_id, ".-----")
        for q_id in dephrasure_scan_range
            println("Now scanning for q =", q_id,".")
            target_val = 2*(2-q_id)
            for p_id in dephrasure_scan_range
                dephrasurepq(X) = dephrasureChannel(X,p_id,q_id)
                sidchan(X) = siddhuChannel(X,s_id)

                dephr_chan = Choi(dephrasurepq,2,3)
                sid_chan= Choi(sidchan,3,3)

                par_cv, = pptCV(parChoi(sid_chan, dephr_chan), :dual)

                non_mult = par_cv - target_val
                !isapprox(non_mult,0,atol=4e-6) ? is_mult = false : nothing
                max_val < non_mult ? max_val = non_mult : nothing

                p_ctr += 1
            end
            q_ctr += 1
            p_ctr = 1
        end
        sid_ctr += 1
        q_ctr = 1
    end
    @test is_mult
end


@testset "Optimizer" begin
    @testset "Global" begin
        f(x) =
            (x[2] - 5.1 / (4 * π^2) * x[1]^2 + 5 / π * x[1] - 6)^2 +
            10 * (1 - 1 / (8π)) * cos(x[1]) +
            10 +
            5 * x[1]
        optix = [-3.689285272296118, 13.629987729088747]
        optif = -16.64402157084319
        lb = [-5.0, 0.0]
        ub = [10.0, 15.0]
        for o in (CRSOptimizer, ISRESOptimizer)
            opt = o(maxeval = 10000, population = 100)
            fx, x, ret = minimize(opt, f, :Auto, lb, ub, (lb + ub) / 2)
            @test isapprox(fx, optif, rtol = 1e-3)
            @test isapprox(x, optix, rtol = 1e-3)
        end

    end
    @testset "Local" begin
        f(x) = sum(i -> i * x[i]^2, length(x))
        lb = fill(-10.0, 5)
        ub = fill(10.0, 5)
        optif = 0.0
        optix = zeros(5)
        for o in (LBFGSOptimizer, BOBYQAOptimizer)
            opt = o(maxeval = 1000)
            fx, x, ret = minimize(opt, f, :Auto, lb, ub, (lb + ub) / 2)
            @test isapprox(fx, optif, rtol = 1e-3)
            @test isapprox(x, optix, rtol = 1e-3)
        end
    end
end


using CompScienceMeshes

m = CompScienceMeshes.tetmeshsphere(1.0, 0.35)
@show numcells(m)

br = barycentric_refinement(m)
@show numcells(br)

@assert numcells(br) == 24*numcells(m)

V = 155
sm = submesh(tet -> V in tet, br.mesh)
@show numcells(sm)

bndry = boundary(sm)
import PlotlyJS
PlotlyJS.plot(patch(bndry))

E = 100
Edges = skeleton(m,1)
sm = submesh(br.mesh) do tet
    cells(Edges)[E][1] in tet && return true
    cells(Edges)[E][2] in tet && return true
    return false
end

PlotlyJS.plot(patch(boundary(sm)))


# This is an introductory julia script to help get new users quickly familiar with the basics
# It is recommended that you use JuliaPro, as it makes interacting with julia much easier

# If you are using JuliaPro, you can execute a line or block of code with Shift+Enter. Otherwise you can just use Julia in the console and copy-paste code to execute it
x = 0.0:0.01:1.0
# variable names can be unicode characters. Special characters like greek letters can be entered by typing them as in LaTeX, then hitting tab.
# for example \alpha -> α
μ = 4
# Functions are first class objects. You can assign them to variables and pass them to functions. In addition you can create anonymous functions using lambda expressions.
T = x -> μ*x*(1-x)

# Calling a function with a '.' at the end of its name applies it to every element of its input, usually an array. This is similar to Matlab
y = T.(x)

# Julia has many community packages that can be installed from the command line.
# Packages are loaded with the following syntax
using Plots
# The package Plots adds high level plotting functions
plot(x,y)
# When we want to plot a function, we can also pass it as an argument, along with the values we want to evaluate
plot(T,x, xlabel="x", ylabel="T(x)")

# Finally, variables all have a type, for instance Int64 or Float64. Julia uses type to optimize  code, especially functions
typeof(T(1))
typeof(T(0.5))
typeof(T(1//2))
# There are abstract types that encompass many concrete types, for example Real is the abstract type for real numbers, which encompasses Int64, Float64, etc. When declaring functions arguments with an abstract type, that function will work on any instance of that type.
function tent(x::Real, a::Real)
    if x <= 0.5
        return a*x
    else
        return a*(1-x)
    end
end

tent(1,3)
tent(0.5,3)
tent(1,3.1)


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

export PolCalFunc

struct PolCalFunc{N,T<:Number} <:Function
    params::NTuple{N,T}

    PolCalFunc(params::T...) where T = new{length(params),T}(params)
end

function (f::PolCalFunc{N,T})(x::U) where {N,T,U}
    R = promote_type(T, U)
    y = zero(R)
    xn = one(U)
    @inbounds for p in f.params
        y = R(fma(p, xn, y))
        xn *= x
    end
    y
end

function Base.convert(::Type{PolCalFunc{N,T}}, p::PropDict) where {N,T}
    funcstr = p.func
    if !(funcstr in ("pol1", "[0]+[1]*x", "pol2"))
        error("Unsupported function \"$funcstr\" for $(PolCalFunc{N,T})")
    end

    coeffs = zeros(T, N)
    for (i,v) in p.params
        coeffs[i + 1] = v
    end
    PolCalFunc(coeffs...)
end

function Base.convert(::Type{Dict{Int, PolCalFunc{N,T}}}, p::PropDict) where {N,T}
    calfuncs = Dict{Int, PolCalFunc{N,T}}()
    for (ch, caldict) in p
        calfuncs[ch] = convert(PolCalFunc{2,Float64}, caldict)
    end
    calfuncs
end

const CalFuncDict = Dict{Int,PolCalFunc{2,Float64}}
export CalFuncDict


@testset "Schreier Vectors and Stabilizers" begin
    Random.seed!(1);
    SIZE=30
    G = PermutationGroups.AbstractAlgebra.SymmetricGroup(SIZE);
    S = [rand(G) for _ in 1:3];

    pt = 1
    tr = Orbit1(Transversal, S, pt)

    import PermutationGroups.Word

    for OrbT in OrbitTypes
        schr = Schreier(OrbT, S, pt, ^)
        @test schr[pt] == one(G)

        g,h,k = S

        @test representative(S, schr.orb, pt) == one(G)
        @test representative(S, schr.orb, pt^g) == g
        @test representative(S, schr.orb, (pt^g)^h) == g*h
        @test representative(S, schr.orb, pt^(g*h)) == g*h

        @test Word(inv.(S), schr.orb, pt) == Word(Int[])
        @test Word(inv.(S), schr.orb, pt^g) == Word(Int[1])
        @test Word(inv.(S), schr.orb, (pt^g)^h) == Word(Int[2,1])
        @test Word(inv.(S), schr.orb, pt^(g*h)) == Word(Int[2,1])

        z = g*h*k
        δ = pt^z
        @test Word(inv.(S), schr.orb, δ, ^) == Word([3,2,1])

        @test schr[δ] == z
        @test all([pt^schr[o] == o for o in schr.orb])
        @test all(schr[o] == tr[o] for o in tr)
    end
end

@testset "Schreier-Sims unit tests" begin
    S = [perm"(1,2,4,3)(5)", perm"(2,5,4)"]
    @test PermutationGroups.initial_bsgs(S)[1] == [1,2]
    @test PermutationGroups.initial_bsgs(S)[2] == [S, [S[2]]]

    a,b = [perm"(1,2,4,3)(5)", perm"(1,2,5,4)"]
    B, S = PermutationGroups.initial_bsgs([a,b])
    @test B == [1]
    @test S[1] == [a,b]

    trs = [Schreier(gens, pt) for (gens, pt) in zip(S, B)]
    @test collect(trs[1]) == [1,2,4,5,3]
    depth = 1
    Δ = trs[depth]

    # first point on the orbit is β = 1:
    u = Δ[1]
    # the first Schreier generator is trivial
    @test isone(u*a*getinv(Δ, 1^a))
    # but the second is not
    @test u*b*getinv(Δ, 1^b) == perm"(2,5)(3,4)"
    c = perm"(2,5)(3,4)"
    @test sift(c, B, trs) == (c, 2)

    # so we extend the basis
    push!(B, firstmoved(c))
    # add c to generators at depth 2
    push!(S, [c])
    # and compute the orbit of the added point at depth 2
    push!(trs, Schreier(S[depth+1], 2))
    depth = 2
    @test sift(c, B, trs) == (perm"(5)", 3)

    u = Δ[2]
    @test u == a
    # Schreier generators are trivial
    @test u*a == Δ[2^a]
    @test u*b == Δ[2^b]
    # so we move to the next point on the orbit
    u = Δ[4] # β = 4
    @test u*a == Δ[4^a]
    @test u*b*getinv(Δ, 4^b) == a^2*b
    d = a^2*b
    @test sift(d, B, trs) == (d, 2)

    push!(S[2], d)
    trs[2] = Schreier(S[2], B[2])
    @test collect(trs[2]) == [2,5,3,4]

    @test sift(a^2*b, B, trs) == (perm"(5)", 3)

    @test prod(length, trs) == 20

    B_new, S_new, trs_new = schreier_sims([a,b])
    @test B_new == B
    @test S_new == S
    @test trs_new == trs

end


using StanSample, Random

Random.seed!(1233)

ProjDir = @__DIR__
cd(ProjDir)

include(joinpath(@__DIR__, "LBA_functions.jl"))

LBA_stan = "
functions{

     real lba_pdf(real t, real b, real A, real v, real s){
          //PDF of the LBA model

          real b_A_tv_ts;
          real b_tv_ts;
          real term_1;
          real term_2;
          real term_3;
          real term_4;
          real pdf;

          b_A_tv_ts = (b - A - t*v)/(t*s);
          b_tv_ts = (b - t*v)/(t*s);
          term_1 = v*Phi(b_A_tv_ts);
          term_2 = s*exp(normal_lpdf(b_A_tv_ts|0,1));
          term_3 = v*Phi(b_tv_ts);
          term_4 = s*exp(normal_lpdf(b_tv_ts|0,1));
          pdf = (1/A)*(-term_1 + term_2 + term_3 - term_4);

          return pdf;
     }

     real lba_cdf(real t, real b, real A, real v, real s){
          //CDF of the LBA model

          real b_A_tv;
          real b_tv;
          real ts;
          real term_1;
          real term_2;
          real term_3;
          real term_4;
          real cdf;

          b_A_tv = b - A - t*v;
          b_tv = b - t*v;
          ts = t*s;
          term_1 = b_A_tv/A * Phi(b_A_tv/ts);
          term_2 = b_tv/A   * Phi(b_tv/ts);
          term_3 = ts/A     * exp(normal_lpdf(b_A_tv/ts|0,1));
          term_4 = ts/A     * exp(normal_lpdf(b_tv/ts|0,1));
          cdf = 1 + term_1 - term_2 + term_3 - term_4;

          return cdf;

     }

     real lba_lpdf(matrix RT, real k, real A, vector v, real s, real tau){

          real t;
          real b;
          real cdf;
          real pdf;
          vector[rows(RT)] prob;
          real out;
          real prob_neg;

          b = A + k;
          for (i in 1:rows(RT)){
               t = RT[i,1] - tau;
               if(t > 0){
                    cdf = 1;

                    for(j in 1:num_elements(v)){
                         if(RT[i,2] == j){
                              pdf = lba_pdf(t, b, A, v[j], s);
                         }else{
                              cdf = (1-lba_cdf(t, b, A, v[j], s)) * cdf;
                         }
                    }
                    prob_neg = 1;
                    for(j in 1:num_elements(v)){
                         prob_neg = Phi(-v[j]/s) * prob_neg;
                    }
                    prob[i] = pdf*cdf;
                    prob[i] = prob[i]/(1-prob_neg);
                    if(prob[i] < 1e-10){
                         prob[i] = 1e-10;
                    }

               }else{
                    prob[i] = 1e-10;
               }
          }
          out = sum(log(prob));
          return out;
     }
}

data{
     int N;
     int Nc;
     vector[N] rt;
     vector[N] choice;
}

parameters {
     real<lower=0> k;
     real<lower=0> A;
     real<lower=0> tau;
     vector<lower=0>[Nc] v;
}

model {
     real s;
     matrix[N,2] RT;
     s=1;
     RT[:,1] = rt;
     RT[:,2] = choice;
     k ~ normal(.5,1)T[0,];
     A ~ normal(.5,1)T[0,];
     tau ~ normal(.5,.5)T[0,];
     for(n in 1:Nc){
          v[n] ~ normal(2,1)T[0,];
     }
     RT ~ lba(k,A,v,s,tau);
}
";

# First 2 runs are using the standard 2.19.1 version of cmdstan

# This run tests passing a data file name as data in the stan_sample() call

stanmodel = SampleModel("LBA", LBA_stan;
  method = StanSample.Sample(adapt = StanSample.Adapt(delta = 0.95)));

N = 200
v = [1.0, 1.5]
Nc = length(v)
data = simulateLBA(;Nd=N,v=v,A=.8,k=.2,tau=.4)

@time (sample_file, log_file) = stan_sample(stanmodel; data=data, n_chains=4)

if !(sample_file == Nothing)
  # Use StanSamples to read a chain in NamedTupla format
  nt = read_samples(stanmodel.sm; chain = 3)

  # Convert to an MCMCChains.Chains object
  chns = read_samples(stanmodel)

  # Describe the MCMCChains using MCMCChains statistics
  display(describe(chns))

  # Show the same output in DataFrame format
  #sdf = StanSample.read_summary(stanmodel)
  #display(sdf)
end


# SoftMax Arm selection algorithm

"""
    SoftMax Arm Selection
"""
struct SoftMax <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    τ::Float64              # temparature parameter
    play_count::Vector{Int64}    # Number of times each arm is played
    avgValue::Vector{Float64}   # Average value of each arm

    function SoftMax( noOfArms::Integer, τ::AbstractFloat )
        new( noOfArms,
             0,
             0,
             τ,
             zeros(Int64,noOfArms),
             zeros(Float64,noOfArms)
        )
    end
end

function get_arm_index( agent::SoftMax )
    pVec    = exp.(agent.avgValue/agent.τ) / sum(exp.(agent.avgValue/agent.τ))
    agent.lastPlayedArm = rand( Categorical(pVec) )
    return agent.lastPlayedArm
end

function update_reward!( agent::SoftMax, r::Real )
    agent.avgValue[agent.lastPlayedArm] = (agent.play_count[agent.lastPlayedArm]*agent.avgValue[agent.lastPlayedArm] + r) / (agent.play_count[agent.lastPlayedArm]+1)
    agent.play_count[agent.lastPlayedArm] += 1
    agent.noOfSteps += 1

    nothing
end

function reset!( agent::SoftMax )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.play_count    = zeros( Int64, agent.noOfArms )
    agent.avgValue      = zeros( Float64, agent.noOfArms )

    nothing
end

function info_str( agent::SoftMax, latex::Bool )
    if latex
        return @sprintf( "SoftMax (\$\\tau = %4.3f\$)", agent.τ )
    else
        return @sprintf( "SoftMax (τ = %4.3f)", agent.τ )
    end
end


using CairoMakie
n = 50
t = LinRange(0, 5*π, n)
x(t; a = 0.0, b = -2.4) = (a + b*t) * cos(t)
y(t; a = 0.0, b = -2.4) = (a + b*t) * sin(t)
with_theme(theme_dark()) do
    fig = Figure(resolution = (600,400))
    ax = Axis(fig[1,1]; aspect = DataAspect(), xlabel = "x", ylabel = "y")
    for (k,i) in enumerate(-4:-2)
        scatterlines!(x.(t; b = i), y.(t; b = i); color = 1:n,
      	colormap = (:Spectral_11, 0.5), markercolor = 1:n,
      	markersize = LinRange(3,30/k,n), markercolormap = :Spectral)
    end
    save("scatterlines_archimedes_spiral.png", fig, px_per_unit = 2)
    display(fig)
end


mutable struct SRC0 <: SRCType
end
function (obj::SRC0)(
              xy::AbstractVector,
              f::AbstractArray)
  f[:] = 0.0
  return nothing
end

mutable struct SRC1 <: SRCType
end
function (obj::SRC1,
              xy::AbstractVector,
              f::AbstractArray)
  f[:] = 1.0
  return nothing
end

#
# a scalar diffusion coefficient, 10
#
mutable struct SrcExpTrigPoly0thDiffn <: SRCType
end
function (obj::SrcExpTrigPoly0thDiffn)(
              xy::AbstractVector,
              src::AbstractArray)
  k = 2.0
  ss = sin(2*k*pi*xy[1])*sin(2*k*pi*xy[2])
  cs = cos(2*k*pi*xy[1])*sin(2*k*pi*xy[2])
  sc = sin(2*k*pi*xy[1])*cos(2*k*pi*xy[2])
  ex = exp(xy[1] + xy[2])
  src[:] = (2. - 8*k*k*pi*pi) * ss + 4*k*pi*(cs + sc)
  src[:] *= -10.*ex
  return nothing
end
#
# a scalar diffusion coefficient, 10
#
mutable struct SrcTrigPoly0thDiffn <: SRCType
end
function (obj::SrcTrigPoly0thDiffn)(
              xy::AbstractVector,
              src::AbstractArray)
  n = 2.0
  factor = 8.0*n*n*pi*pi
  ss = sin(2*n*pi*xy[1])*sin(2*n*pi*xy[2])
  src[:] = 0.0
  for i = 1:length(src)
    src[i] = factor*ss * 1.0e1
  end
  return nothing
end
#
# a 2nd order diffusion coefficient tensor
#
mutable struct SrcTrigPoly2ndDiffn <: SRCType
end
function (obj::SrcTrigPoly2ndDiffn)(
              xy::AbstractVector,
              src::AbstractArray)
  n = 2.0
  sc = sin(2*n*pi*xy[1])*cos(2*n*pi*xy[2])
  cs = cos(2*n*pi*xy[1])*sin(2*n*pi*xy[2])
  ss = sin(2*n*pi*xy[1])*sin(2*n*pi*xy[2])
  cc = cos(2*n*pi*xy[1])*cos(2*n*pi*xy[2])
  src[:] = 0.0
  for i = 1:length(src)
    src[i] += 6*n*pi*xy[1]*cs
    src[i] += 6*n*pi*xy[2]*sc
    src[i] += 4*n*n*pi*pi*(-2.0 - xy[1]*xy[1] - xy[2]*xy[2])*ss
    src[i] += 8*n*n*pi*pi*xy[1]*xy[2]*cc
    src[i] = -src[i]
  end
  return nothing
end

#
# a 6th order diffusion coefficient tensor
#
mutable struct SrcTrigPoly6thDiffn <: SRCType
end
function (obj::SrcTrigPoly6thDiffn)(
              xy::AbstractVector,
              src::AbstractArray)
  n = 2.0
  sc = sin(2*n*pi*xy[1])*cos(2*n*pi*xy[2])
  cs = cos(2*n*pi*xy[1])*sin(2*n*pi*xy[2])
  ss = sin(2*n*pi*xy[1])*sin(2*n*pi*xy[2])
  cc = cos(2*n*pi*xy[1])*cos(2*n*pi*xy[2])
  src[:] = 0.0
  for i = 1:length(src)
    src[i] += 2*n*pi*(6*xy[1]^5 + 3*xy[1]^3*xy[2]^2)*cs
    src[i] += 2*n*pi*(6*xy[2]^5 + 3*xy[1]^2*xy[2]^3)*sc
    src[i] += 4*n*n*pi*pi*(-2.0 - xy[1]^6 - xy[2]^6)*ss
    src[i] += 8*n*n*pi*pi*xy[1]^3*xy[2]^3*cc
    src[i] = -src[i]
  end
  return nothing
end

global const SRCDict = Dict{String, SRCType}(
 "SRC0" => SRC0(),
 "SRC1" => SRC1(),
 "SrcTrigPoly0thDiffn" => SrcTrigPoly0thDiffn(),
 "SrcExpTrigPoly0thDiffn" => SrcExpTrigPoly0thDiffn(),
 "SrcTrigPoly2ndDiffn" => SrcTrigPoly2ndDiffn(),
 "SrcTrigPoly6thDiffn" => SrcTrigPoly6thDiffn(),
)

function getSrcFuntors(mesh::AbstractMesh,
                       sbp::AbstractOperator,
                       eqn::EllipticData,
                       opts)
  eqn.src_func = SRCDict[opts["SRC_name"]]
  return nothing
end

function calcSource(mesh::AbstractMesh{Tmsh},
                    sbp::AbstractOperator,
                    eqn::EllipticData{Tsol, Tres, Tdim},
                    srcFunc::SRCType,
                    src::AbstractArray{Tres, 3}) where {Tmsh, Tsol, Tres, Tdim}

  numElems = mesh.numEl
  numNodesPerElement = mesh.numNodesPerElement
  # xy = Array{Tmsh}(Tdim)
  src_tmp = Array{eltype(src)}(size(src, 1))
  for elem = 1:numElems
    for n = 1:numNodesPerElement
      xy = sview(mesh.coords, :, n, elem)

      srcFunc(xy, sview(src, :, n, elem))
    end
  end
end


using LOWDER

# -----------------------------------------------
# Problem definition
# -----------------------------------------------

function f(x)

    return - 0.5 * x[1] + 10 * x[2]

end

function g(x)

    return x[1] ^ 2.0 + x[2] ^ 2.0

end

fmin_list = [f, g]
x = [1.0, 1.0]
a = [0.0, 0.0]
b = [5.0, 5.0]

# -----------------------------------------------
# LOWDER parameters
# -----------------------------------------------

δinit = 1.0
Δinit = 1.2
n_points = 3
verbosity = 3
path_to_file = "./data_files/simple_runtest_01.dat"

# -----------------------------------------------
# LOWDER call
# -----------------------------------------------

sol = LOWDER.lowder(fmin_list, x, a, b; δ = δinit, Δ = Δinit, m = n_points, verbose = verbosity, filename = path_to_file)


using JointPosteriors

struct BinaryClassification{T} <: parameters{T,1}
  x::Vector{T}
  p::ProbabilityVector{3,T}
end

function log_density(Θ::BinaryClassification, data::Data)

  log_π = data.αm_m1 * log(Θ.p[3]) + data.βm_m1 * log(1 - Θ.p[3]) + data.αp_m1*log(Θ.p[2]) + data.βp_m1 * log(1 - Θ.p[2]) + data.ατ_m1 * log(Θ.p[1]) + data.βτ_m1 * log(1 - Θ.p[1])

  for i ∈ eachindex(data.X)
    log_π += data.freq[i] * log( Θ.p[1] * (1 - Θ.p[3])^data.X[i] * Θ.p[3]^(data.NmX[i]) + (1-Θ.p[1]) * Θ.p[2]^data.X[i] * (1-Θ.p[2])^data.NmX[i] )
  end

  log_π

end

struct BinaryClassificationData <: Data
  X::Array{Int64,1}
  freq::Array{Int64,1}
  NmX::Array{Int64, 1}
  N::Int64
  αm_m1::Float64
  βm_m1::Float64
  αp_m1::Float64
  βp_m1::Float64
  ατ_m1::Float64
  βτ_m1::Float64

end

function BinaryClassificationData(x::Array{Int64, 1}, Freq::Array{Int64,1}; αm::Real = 1, βm::Real = 1, αp::Real = 1, βp::Real = 1, ατ::Real = 1, βτ::Real = 1)
  N = x[end]
  non_zero_freq = sum( Freq .!= 0 )
  X = zeros( Int64, non_zero_freq )
  freq = zeros( Int64, non_zero_freq )
  j = 1
  for i ∈ 1:non_zero_freq
    while Freq[j] == 0
      j += 1
    end
    X[i] = x[j]
    freq[i] = Freq[j]
    j += 1
  end
  BinaryClassificationData(X, freq, N .- X, N, αm - 1, βm - 1, αp - 1, βp - 1, ατ - 1, βτ - 1)
end

X = [i for i ∈ 0:9];
freq = [10, 2, 2, 1, 2, 0, 0, 3, 2, 16];
data = BinaryClassificationData(X, freq, βm = 2, βp = 2);

bc_model = Model(BinaryClassification)

jp = JointPosterior(bc_model, data);

τ(Θ::BinaryClassification) = Θ.p[1]
θ_plus(Θ::BinaryClassification) = Θ.p[2]
θ_minus(Θ::BinaryClassification) = Θ.p[3]

marginal_τ = marginal(jp, τ)

marginal_θ_plus = marginal(jp, θ_plus)
marginal_θ_minus = marginal(jp, θ_minus)

f(x) = x.p[1]
d = Distributions.Normal
q = SparseQuadratureGrids.GenzKeister

values_positive = Array{Float64}(length(jp.weights_positive));
values_negative = Array{Float64}(length(jp.weights_negative));
for i ∈ eachindex(values_positive)
  values_positive[i] = f(JointPosteriors.transform!(jp.M.Θ, jp.M.Grid.nodes_positive[:,i], jp.Θ_hat, jp.L, q))
end
for i ∈ eachindex(values_negative)
  values_negative[i] = f(JointPosteriors.transform!(jp.M.Θ, jp.M.Grid.nodes_negative[:,i], jp.Θ_hat, jp.L, q))
end
μ = dot(jp.weights_positive, values_positive) + dot(jp.weights_negative, values_negative)
σ = sqrt(dot(jp.weights_positive, values_positive .^ 2) + dot(jp.weights_negative, values_negative .^ 2) - μ^2)

poly = JointPosteriors.polynomial_interpolation(jp.weights_positive, jp.weights_negative, values_positive, values_negative, d)

cdf_error(zeros(6), poly)

function prec(n::Int, ρ::Real)
  p = zeros(n,n)
  np1 = n+1
  np1h2 = (np1/2)^2
  β2 = 1 / ( np1h2 - n )
  β1 = -np1*β2
  β0 = 2 + np1h2*β2
  p[1] = p[end] = 3
  det_im1 = 3
  det_im2 = 1
  a0 = 3
  for i ∈ 2:n
    p[i,i] = a = (β0 + β1*i + β2*i^2)
    b = a*a0*ρ/4
    p[i,i-1] = p[i-1,i] = -sqrt(b)
    det_i = a*det_im1-b*det_im2
    det_im2 = det_im1
    det_im1 = det_i
    a0 = a
  end
  p, det_im1
end
function cet(δ::Vector, n::Int, ρ::Real)
  np1 = n+1
  np1h2 = (np1/2)^2
  β2 = 1 / ( np1h2 - n )
  β1 = -np1*β2
  β0 = 2 + np1h2*β2
  det_im1 = 3.0
  det_im2 = 1.0
  out = 3.0 * δ[1]^2
  a0 = 3.0
  for i ∈ 2:n
    a = (β0 + β1*i + β2*i^2)
    b = a*a0*ρ/4
    out += δ[i]^2*a - 2δ[i-1]*δ[i]*√b

    det_i = a*det_im1-b*det_im2
    det_im2 = det_im1
    det_im1 = det_i
    a0 = a
  end
  out, det_im1
end

itp = JointPosteriors.InterpolateIntegral(jp.weights_positive, jp.weights_negative, values_positive, values_negative, d)

using JointPosteriors
struct HiWorld{T} <: parameters{T,1}
  x::Vector{T}
  β::RealVector{3,T}
  σ::PositiveVector{1,T}
end

struct HiWorldData <: Data
  X::Array{Float64,2}
  y::Vector{Float64}
end

function log_density(Θ::HiWorld, data::Data)
  lpdf_normal(Θ.β, 0, 10) + lpdf_normal(Θ.σ[1], 0, 1) + lpdf_normal(data.y, data.X * Θ.β, Θ.σ[1])
end

# True parameter values
sigma = 1;
beta = [1, 1, 2.5];

# Size of dataset
size = 100;

# Predictor variable
X = hcat(ones(size), randn(size,2));
X[:,3] .*= 0.2;

# Simulate outcome variable
y = X*beta .+ randn(size) .* sigma;

HW_data = HiWorldData(X, y);
HW_model = Model(HiWorld);
hw_jp = JointPosterior(HW_model, HW_data);

marginal(hw_jp, x -> x.β[1])
marginal(hw_jp, x -> x.β[2])
marginal(hw_jp, x -> x.β[3])
marginal(hw_jp, x -> x.σ[1])

const hw_stan = "data {
  int N;
  vector[N] y;
  matrix[N, 2] x;
}
parameters {
  real alpha;
  vector[2] beta;
  real<lower=0> sigma;
}
model {
  alpha ~ normal(0, 10);
  beta ~ normal(0, 10);
  sigma ~ normal(0, 1);
  y ~ normal(alpha + x * beta, sigma);
}"
using Stan, Mamba
set_cmdstan_home!("/mnt/ssd/cmdstan-2.14.0")
#set_cmdstan_home!("/home/christel/Downloads/cmdstan-2.14.0")

hw_stan_data = Dict("N" => length(y), "x" => X[:,2:3], "y" => y)
hw_stan_model = Stanmodel(Sample(), name = "HelloWorld", model = hw_stan, monitors = ["alpha", "beta.1", "beta.2", "sigma"]);
hw_stan_res = stan(hw_stan_model, [hw_stan_data])
describe(hw_stan_res[2])

function run_hw()
  hw_jp = JointPosterior(HW_model, HW_data);
  marginal(hw_jp, x -> x.β[1])
  marginal(hw_jp, x -> x.β[2])
  marginal(hw_jp, x -> x.β[3])
  marginal(hw_jp, x -> x.σ[1])
end
using BenchmarkTools
@benchmark run_hw()

using JointPosteriors

include("/mnt/ssd/MSA/Univariate/dataGen.jl")

tv = true_values_uni(15, √99, √0.6, √0.3, √.1);
y, yp, yo = gen_data(tv, 40, 12, 12);

d = LogDensities.TF_RE_ANOVA_Data(y, yp, yo);
anova = Model(LogDensities.TF_RE_ANOVA);

jp = JointPosterior(anova, d);

function rGT(Θ::LogDensities.TF_RE_ANOVA)
  σg2 = sum(Θ.σ[2:end])
  sqrt(σg2 / (σg2 + Θ.σ[1]))
end

marginal_rGT = marginal(jp, rGT)

function run_anova()
  jp = JointPosterior(anova, d)
  marginal_rGT = marginal(jp, rGT)
end
73, 67, 75, 81) seconds, 4.9 minutes total
Sampling took (70, 103, 69, 72)
const tfanova = "
data {
  int Np1;
  int p;
  int o;
  int POp1;
  matrix[p, o] cat_count;
  matrix[p, o] y_bars;
  real ns2h;
  real cauchy_spread;
}
transformed data {
  real n;
  real po;
  matrix[p, o] invRootCounts;
  invRootCounts = inv_sqrt(cat_count);
  n = sum(cat_count);
  po = p * o;
}
parameters {
  real mu_G;
  vector[p] theta_p;
  vector[o] theta_o;
  matrix[p,o] theta_po;
  real<lower = 0> sigma_p;
  real<lower = 0> sigma_o;
  real<lower = 0> sigma_po;
  real<lower = 0> sigma_r;
}
model {
  matrix[p,o] quad;
  for (i in 1:p){
    theta_po[i, ] ~ normal(0, sigma_po);
    for (j in 1:o){
      y_bars[i,j] ~ normal( mu_G + theta_p[i] + theta_o[j] + theta_po[i, j], invRootCounts[i,j] * sigma_r);
    }
  }
  mu_G ~ normal(0, 50);
  theta_p ~ normal(0, sigma_p);
  theta_o ~ normal(0, sigma_o);
  sigma_o ~ cauchy(0, cauchy_spread);
  target += ns2h / sigma_r^2 - (n - po) * log(sigma_r) ;
}
generated quantities {
  real sigma_g2;
  real sigma_g;
  real sigma_t;
  real rGT;
  sigma_g2 = sigma_o^2 + sigma_po^2 + sigma_r^2;
  sigma_g = sqrt(sigma_g2);
  sigma_t = sqrt(sigma_g2 + sigma_p^2);
  rGT = sigma_g / sigma_t;
}
"

using Stan, Mamba#, Gadfly
set_cmdstan_home!("/mnt/ssd/cmdstan-2.14.0")
anova_stan = Stanmodel(Sample(), name = "ANOVA", model = tfanova, monitors = ["sigma_g", "sigma_t", "rGT"]);

function StanDataANOVA(data::TF_RE_ANOVA_Data)
  Dict( "Np1" => data.N + 1, "p" => data.P, "Pp1" => data.P + 1, "o" => data.O, "POp1" => data.PO + 1, "cat_count" => fill(data.R, data.P, data.O), "y_bars" => data.δ .+ data.μ_hat, "ns2h" =>  - data.s2 / 2.0)
end

res = stan(anova_stan, [StanDataANOVA(data)])
describe(res[2])


### convolution.jl
#
# Copyright (C) 2021 Ido Kessler.
#
# Keywords: exact integer convolution
#
# This file is a part of ExactConvolution.jl.
#
# License is MIT.
#
### Commentary:
#
# This file contains implementation of convolutions using bigints multiplication.
#
### Code:

export exact_conv

include("conversion.jl")

"""
    exact_conv([T,] a, b)

Convolve arrays of integers `a` and `b`, and return the resulting array.

```math
exact\\_conv(a,b)[k] = \\sum_{i+j=k} a[i] \\cdot b[j]
```

# Notes:
- The implementation uses GMP's BigInt multiplication as the underling DSP engine (FFT for big enough inputs).
- T must be sufficient to support the result.
- Using negative values in `a` and `b` is supported but discouraged for runtime reasons.

# Examples
```jldoctest
julia> exact_conv(Int32, [1,2], [4,5,6])
4-element Vector{Int32}:
  4
 13
 16
 12
```
"""
function exact_conv end

exact_conv(a, b) = exact_conv(Int64, a, b)

function exact_conv(
    ::Type{T},
    a::AbstractArray{T1},
    b::AbstractArray{T2},
) where {T<:Integer,T1<:Integer,T2<:Integer}
    function make_positive(::Type{T}, v) where {T}
        min_v = min(minimum(v), 0) |> T
        return min_v, (min_v == 0 ? v : (v .- min_v))
    end

    min_a, a_pos = make_positive(T, a)
    min_b, b_pos = make_positive(T, b)

    res = positive_exact_conv(T, a_pos, b_pos)

    # handle the added -min_a * (b[...] - min_b)
    handle_negatives!(T, res, -min_a, b_pos)
    # handle the added -min_b * a[...]
    handle_negatives!(T, res, -min_b, a)

    res
end

# Only accept positive values
function positive_exact_conv(::Type{T}, a, b)::Vector{T} where {T<:Integer}
    nbits(x) = ndigits(x; base = 2)
    max_bits = nbits ∘ maximum

    lg2_max_mult = max_bits(a) + max_bits(b)
    lg2_max_len = [a, b] .|> length |> minimum |> nbits
    nb = lg2_max_mult + lg2_max_len
    nb = 5 * ceil(Int, nb // 5)

    big_a = array2bigint(a, nb)
    big_b = array2bigint(b, nb)
    ab = big_a * big_b
    n = length(a) + length(b) - 1
    bigint2array(T, ab, n, nb)
end

function handle_negatives!(::Type{T}, res, v1_added, v2_pos)::Nothing where {T}
    v1_added == 0 && return
    # min_v1 * v2_[i:j] was added
    n, n2 = length(res), length(v2_pos)

    v1_added = T(v1_added)
    extra = T(0)
    # v2[i] only participate in the conv from index i
    for (i, v) ∈ zip(1:n2, v2_pos)
        extra += v * v1_added
        res[i] -= extra
    end
    res[n2+1:end] .-= extra

    removed_extra = T(0)
    # v2[end-i] only participate in the conv upto index end-i
    for (i, v) ∈ zip(n-n2+1:n, v2_pos)
        res[i] += removed_extra
        removed_extra += v * v1_added
    end
    return
end


function levenberg_marquardt(f::Function, g::Function, x0::Vector{Float64}; tolX=1e-8, tolG=1e-12,
                             maxIter=100, lambda=100.0, show_trace=false, quiet=false)
	# finds argmin sum(f(x).^2) using the Levenberg-Marquardt algorithm
	#          x
	# The function f should take an input vector of length n and return an output vector of length m
	# The function g is the Jacobian of f, and should be an m x n matrix
	# x0 is an initial guess for the solution
	# fargs is a tuple of additional arguments to pass to f
	# available options:
	#   tolX - search tolerance in x
	#   tolG - search tolerance in gradient
	#   maxIter - maximum number of iterations
	#   lambda - (inverse of) initial trust region radius
	#   show_trace - print a status summary on each iteration if true
	# returns: x, J
	#   x - least squares solution for x
	#   J - estimate of the Jacobian of f at x

	const MAX_LAMBDA = 1e16 # minimum trust region radius
	const MIN_LAMBDA = 1e-16 # maximum trust region radius
	const MIN_STEP_QUALITY = 1e-3
	const MIN_DIAGONAL = 1e-6 # lower bound on values of diagonal matrix used to regularize the trust region step
	converged = false
	iterCt = 0
	x = x0
	delta_x = copy(x0)
	f_calls = 0
	g_calls = 0
	fcur = f(x)
	f_calls += 1
	residual = norm(fcur)^2

	while ~converged && iterCt < maxIter
		J = g(x)
		g_calls += 1
		# we want to solve:
		#    argmin 0.5*||J(x)*delta_x + f(x)||^2 + lambda*||diagm(J'*J)*delta_x||^2
		# Solving for the minimum gives:
		#    (J'*J + lambda*DtD) * delta_x == -J^T * f(x), where DtD = diagm(sum(J.^2,1))
		# Where we have used the equivalence: diagm(J'*J) = diagm(sum(J.^2, 1))
		# It is additionally useful to bound the elements of DtD below to help
		# prevent "parameter evaporation".
		DtD = diagm(Float64[max(x, MIN_DIAGONAL) for x in sum(J.^2,1)])
		delta_x = ( J'*J + sqrt(lambda)*DtD ) \ -J'*fcur

		# if the linear assumption is valid, our new residual should be:
		predicted_residual = norm(J*delta_x + fcur)^2
		# check for numerical problems in solving for delta_x by ensuring that the predicted residual is smaller
		# than the current residual
		if predicted_residual > residual + 2max(eps(predicted_residual),eps(residual))
			warn("""Problem solving for delta_x: predicted residual increase.
                             $predicted_residual (predicted_residual) >
                             $residual (residual) + $(eps(predicted_residual)) (eps)""")
		end
		# try the step and compute its quality
		trial_f = f(x + delta_x)
		f_calls += 1
		trial_residual = norm(trial_f)^2
		# step quality = residual change / predicted residual change
		rho = (trial_residual - residual) / (predicted_residual - residual)
		if rho > MIN_STEP_QUALITY
			x += delta_x
			fcur = trial_f
			residual = trial_residual
			# increase trust region radius
			lambda = max(0.1*lambda, MIN_LAMBDA)
		else
			# decrease trust region radius
			lambda = min(10*lambda, MAX_LAMBDA)
		end
		iterCt += 1
		# check convergence criteria:
		# 1. Small gradient: norm(J^T * fcur, Inf) < tolG
		# 2. Small step size: norm(delta_x) < tolX
		converged = (norm(J' * fcur, Inf) < tolG) || (norm(delta_x) < tolX*(tolX + norm(x)))
	end
	x
end

function nlsfitworker(model::Function, y::Matrix{Float64},
                          x::Vector{Float64}, p0::Vector{Float64})
  # performs Levenberg-Marquardt fitting
  nparams = length(p0)
  n = size(y,2)
  params = zeros(nparams, n)
  resids = similar(y)
  for j in 1:n
    f(p) = model(x, p) - y[:,j]
    g = Calculus.jacobian(f, :forward)
    results = levenberg_marquardt(f, g, p0; tolX=1e-3, tolG=1e-4,
                                        lambda=10.0)
    resids[:,j] = f(results) / sqrt(norm(y[:,j]))
    params[:,j] = results
  end
  (params, resids)
end

function nlsfit(f::Function, y::Matrix{Float64}, idxs::Vector{Int},
                x::Vector{Float64}, p0::Vector{Float64})
  tic()
  nt, n = size(y)
  nw = nworkers()
  nidxs = length(idxs)
  params = zeros(length(p0), n)
  resids = zeros(nt, n)
  dof = nt - length(p0)
  if nw > 1
    reflist = Any[]
    idxperworker = round(Int, nidxs / nw)
    workeridxs = Any[idxs[(w-1)*idxperworker+1:min(w*idxperworker,nidxs)] for w in 1:nw]
    @dprint "fitting $nt x $idxperworker points on each of $nw workers"
    workerids = workers()
    for w in 1:nw
      r = @spawnat workerids[w] nlsfitworker(f, y[:,workeridxs[w]], x, p0)
      push!(reflist, r)
    end
    for w in 1:nw
      params[:,workeridxs[w]], resids[:,workeridxs[w]] = fetch(reflist[w])
    end
  else   # run on one CPU core
    @dprint "running $nt x $nidxs points on one CPU core"
    params[:,idxs], resids[:,idxs] = nlsfitworker(f, y[:,idxs], x, p0)
  end
  t = toq()
  vps = nidxs / t
  @dprint @sprintf "processed %d voxels in %.1f s (%.1f vox/s)\n" nidxs t vps
  (params, resids, dof)
end


seed!(12345)

function testSortedUniforms(low, high, h, N)
  function cdf(x::Float64)
    return x
  end
  samples = sampleSortedUniforms(N)
  ## check that they are sorted first
  for j in 2:N
    if samples[j] < samples[j-1]
      return 0.0
    end
  end
  i::Int64 = 0
  function sampler()
    i = i + 1
    return samples[i]
  end
  return testGOFContinuous(cdf, sampler, low, high, h, N)
end

@test testSortedUniforms(0.01, 0.99, 0.01, 2^21) > 0.01
@test testSortedUniforms(0.02, 0.98, 0.02, 2^10) > 0.01


using OrdinaryDiffEq, DiffEqDevTools, Test, LinearAlgebra, Random
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
using OrdinaryDiffEq: maxeig!
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_linear, prob_ode_2Dlinear
probArr = Vector{ODEProblem}(undef, 2)
probArr[1] = prob_ode_linear
probArr[2] = prob_ode_2Dlinear

Random.seed!(123)
@testset "Power Iteration of Runge-Kutta-Chebyshev Tests" begin
  for i in 1:10, iip in [true, false]
    A = randn(20,20)
    test_f(u,p,t) = A*u
    test_f(du,u,p,t) = mul!(du, A, u)
    prob = ODEProblem{iip}(test_f, randn(20), (0,1.))
    integrator = init(prob, ROCK2())
    eigm = maximum(abs.(eigvals(A)))
    maxeig!(integrator, integrator.cache)
    eigest = integrator.eigen_est
    @test eigest ≈ eigm rtol=0.1eigm
  end
end

@testset "Runge-Kutta-Chebyshev Convergence Tests" begin
  dts = 1 .//2 .^(8:-1:4)
  testTol = 0.1
  for prob in probArr
    sim2 = test_convergence(dts,prob,ROCK2())
    @test abs(sim2.𝒪est[:l∞]-2) < testTol
  end
end


# Reference: Temporal Distance Metrics for Social Network Analysis,
#  John Tang, Mirco Musolesi, Cecilia Mascolo and Vito Latora,
# 2009.

# find the shortest temporal path via depth first search
function _DFS_shortest_temporal_path(g::AbstractEvolvingGraph,
                                     v1::Tuple,  # v1 and v2 are active nodes (v,t)
                                     v2::Tuple,
                                     path = TemporalPath(),
                                     shortest =  Union{};
                                     verbose::Bool = false)
    path = deepcopy(path)
    path.walks = [path.walks; v1]

    if verbose
        println("current:", path)
    end

    if v1 == v2
        return path
    end
    for node in forward_neighbors(g, v1)
        if !(has_node(path, node)) # avoid cycles
            if node[2] <= v2[2] # avoid searching nodes at timestamps > v2[2]
                if (shortest ==  Union{}) || (spatial_length(path) < spatial_length(shortest))
                    newPath = _DFS_shortest_temporal_path(g, node, v2, path, shortest, verbose = verbose)
                    if newPath !=  Union{}
                        shortest = newPath
                    end
                end
            end
        end
    end
    return shortest
end

"""
        shortest_temporal_path(g, v1, t1, v2, t2 [, verbose = false])

Find the shortest temporal path from node `v1` at timestamp `t1`
to node `v2` at timestamp `t2` on the evolving graph `g`.
If `verbose = true`, prints the current path at each search step.
"""
shortest_temporal_path(g::AbstractEvolvingGraph, v1, t1, v2, t2; verbose::Bool = false) =    _DFS_shortest_temporal_path(g, (v1, t1), (v2, t2), verbose = verbose)
shortest_temporal_path{V,  T}(g::EvolvingGraph{V, T},
                                                        v1::V, t1::T, v2::V, t2::T;
                                                        verbose::Bool = false) =
   _DFS_shortest_temporal_path(g, (v1, t1), (v2, t2), verbose = verbose)

function shortest_temporal_path{V, T}(g::EvolvingGraph{V, T}, v1, t1, v2, t2;
                                                                        verbose::Bool = false)
    v1 = find_node(g, v1)
    v2 = find_node(g, v2)
    shortest_temporal_path(g, v1, T(t1), v2, T(t2), verbose = verbose)
end

"""
        shortest_temporal_distance(g, v1, t1, v2, t2)

Find the shortest temporal distance from node `v1` at timestamp `t1` to
node `v2` at timestamp `t2` on the evolving graph `g`.
"""
function shortest_temporal_distance(g::AbstractEvolvingGraph, v1, t1, v2, t2)
    if shortest_temporal_path(g, v1, t1, v2, t2) == Union{}
        return Inf
    else
        return length(shortest_temporal_path(g, v1, t1, v2, t2)) - 1
    end
end

"""
        temporal_efficiency(g, v1, t1, v2, t2)

Find the temporal efficiency from node `v1` at timestamp `t1` to node `v2` at timestamp
`t2` on the evolving graph `g`.
"""
temporal_efficiency(g::AbstractEvolvingGraph, v1, t1, v2, t2) =
         one(Float64)/shortest_temporal_distance(g, v1, t1, v2, t2)

"""
        global_temporal_efficiency(g, t1, t2)

Return the global temporal efficiency of the evolving graph `g` between
timestamp `t1` and `t2`. The global temporal efficiency is a measure of
how well information flow between two given timestamps.
"""
function global_temporal_efficiency(g::AbstractEvolvingGraph, t1, t2)
    result = zero(Float64)
    for j in nodes(g)
        for i in nodes(g)
            result += temporal_efficiency(g, i,t1, j,t2)
        end
    end
    result
end


##
using Plots
using GraphRecipes
using SparseArrays
using DelimitedFiles
using LinearAlgebra
##
ei,ej = copy.(eachcol(Int.(readdlm("wiki-simple.edges"))))
xy = readdlm("wiki-simple.xy")'
pages = readlines("wiki-simple.nodes")

## A visual simulation of the random walk with fixed restart
gr(size=(300,300),dpi=300)
using Random
Random.seed!(1)
function sprrw_on_graph(ei, ej, xy, alpha, start, nstep, nframes)
  n = size(xy,2)
  A = sparse(ei,ej,1,n,n)
  B = A.*A'
  U = A - B # we need A' so that the ei, ej plots

  p = graphplot(findnz(triu(B,1))[1:2]..., x = xy[1,:], y = xy[2,:],
    markercolor=:lightblue, markerstrokecolor=:white, curves=false,
    markersize=7, markerstrokewidth=1,
    linecolor=1, linealpha=0.8, linewidth=0.7, arrow=:both, shorten=0.85,
    axis_buffer=0.5, background=nothing)

  graphplot!(findnz(U)[1:2]..., x = xy[1,:], y = xy[2,:],
    markercolor=nothing, markerstrokecolor=:white, curves=false, curvature_scalar = -0.4,
    markersize=0, markerstrokewidth=0,
     linecolor=2, linealpha=0.8, linewidth=0.7, arrow = :tail, shorten=0.85,
    axis_buffer=0.5, background=:black)

  X = start
  # random walk step
  #p1 = scatter!([xy[1,X]],[xy[2,X]],markersize=4, color=1)
  p2 = scatter!([xy[1,X]],[xy[2,X]],markersize=10, color=:orange)

  xlims!(xlims().*1.2)
  ylims!((-0.5,7.5))
  for i=1:10
    annotate!(xy[1,i],xy[2,i],text("$i",9,:white))
  end

  Xn = X
  jump = false
  anim = @animate for i=0:(nframes*nstep)
    stepi = mod(i,nframes) # this is the frame number in the step
    if stepi == 0
      X = Xn
      # with prob alpha, follow a rand. neig
      if rand() <= alpha && length(ej[ei .== X]) > 0
        Xn = rand(ej[ei .== X])
        jump = false
      else # with prob. (1-alpha)
        Xn = start # go back to start
        jump = true
        println("Jump")
      end
    end
    # interpolate between the two positions
    if jump
      pos = stepi/(nframes-2)
      if pos >= 1
        pos = 1
      else
        pos += 0.25*randn()
      end
      pos = min(max(pos,0),1) # cap in 0,1
    else
      pos = min(stepi/(nframes-2),1)
    end
    p[5] = [(1-pos)*xy[1,X] + pos*xy[1,Xn]],[(1-pos)*xy[2,X] + pos*xy[2,Xn]]
  end
  return anim
end
anim = sprrw_on_graph(ei,ej,xy,0.85,1,100,20)
gif(anim, "seeded-pr-rw-on-wiki-simple.gif", fps=40)


function parametric_sample(
    N::Int64,
    sample_fraction::Float64,
    bins::Vector{Int64},
    dist::Distributions.DiscreteNonParametric
    )
    W = dist.p
    K = size(W, 1)
    n_sample = rand(
        Distributions.DiscreteNonParametric(collect(1:K), W),
        Int(floor(sample_fraction * N)),
    )
    #check if each n_sample exists in bins, in case sample another one
    for n = 1:size(n_sample, 1)
        ns = copy(n_sample[n])
        while size(findall(bins .== ns), 1) == 0
            # if n_sample[n] > (K / 2)
            #     ns -= 1
            # else
            #     ns += 1
            # end
            ns = rand(Distributions.DiscreteNonParametric(collect(1:K), W),1)
        end
        n_sample[n] = copy(ns[1])
    end
    return n_sample::Vector{Int64}
end

function non_parametric_sample(
    N::Int64,
    sample_fraction::Float64
)
    return Distributions.sample(
        collect(1:N),
        Int(ceil(sample_fraction * N)),
        replace = true
    )
end


using PyPlot, Random

import Statistics.mean
import Random: randn

Random.seed!(1234)

stochastic = true

μ = 0.2
σ = 1/5
dt = 0.01
nsteps= 2001
T = 0:dt:(nsteps-1)*dt

# Theoretical results
if stochastic
  nens = 1000
  ΔW = sqrt(σ)*randn(nsteps, nens)/sqrt(dt)
  if μ == 0
       E_theory = 0.5*σ*T
    dEdt_theory = 0.5*σ*ones(size(T))
  else
       E_theory = @. σ/4μ * (1 - exp(-2μ*T))
    dEdt_theory = @. σ/2  * exp(-2μ*T)
  end
else
  nens = 1
  ΔW = sqrt(σ)*ones(nsteps, nens)
  if μ == 0
       E_theory = @. 0.5*σ*T^2
    dEdt_theory = σ*T
  else
       E_theory = @. σ/(2*μ^2) * (1 - exp(-μ*T))^2
    dEdt_theory = @. σ/μ * exp(-μ*T) * (1 - exp(-μ*T))
  end

end

# Numerical calculation
X = zeros(size(ΔW))
E_ito = zeros(size(ΔW))
E_str = zeros(size(ΔW))
E_numerical = zeros(size(ΔW))

for it = 1:nsteps-1 # time step the equation

    @views @.   X[it+1, :] = X[it, :] + (-μ*X[it, :] + ΔW[it, :])*dt

    @views @.   E_ito[it+1, :] = E_ito[it, :] + (-2*μ*E_ito[it, :]
                                    .+ σ/2)*dt + X[it, :]*ΔW[it, :]*dt

    Ebar = E_str[it, :] + (-2*μ*E_str[it, :])*dt + X[it, :].*ΔW[it, :]*dt
    @views @.   E_str[it+1, :] = E_str[it, :] + (-2*μ*(0.5*(E_str[it, :] +
                        Ebar)))*dt + (0.5*(X[it, :]+X[it+1, :]))*ΔW[it, :]*dt
end

# Energy
@views @. E_numerical = 0.5*X^2

# compute dE/dt numerically
dEdt_ito = mean((E_ito[2:nsteps, :] - E_ito[1:nsteps-1, :])/dt, dims=2)
dEdt_str = mean((E_str[2:nsteps, :] - E_str[1:nsteps-1, :])/dt, dims=2)

# compute the work and dissipation
work_ito = mean(ΔW[1:nsteps-1, :].* X[1:nsteps-1, :], dims=2) .+ σ/2
work_str = mean(ΔW[1:nsteps-1, :].*(X[1:nsteps-1, :]+X[2:nsteps, :])/2, dims=2)
diss_ito = 2*μ*(mean(E_ito[1:nsteps-1, :], dims=2))
diss_str = 2*μ*(mean(E_str[1:nsteps-1, :], dims=2))

# Make plots: compare E(t) evolution Ito, Stratonovich, direct 0.5*x^2

fig, axs = subplots(nrows=1, figsize=(8, 4))

sca(axs); cla()
plot(μ*T, E_numerical[:, 1],
                linewidth=3, label="\$\\frac{1}{2}x_t^2\$ ")
plot(μ*T, E_ito[:, 1],
                linestyle="--", linewidth=2, label="\$E_t\$ (Ito) ")
plot(μ*T, E_str[:, 1],
                linestyle="-.", linewidth=1, label="\$E_t\$ (Stratonovich) ")
xlabel(L"$\mu t$")
ylabel(L"$E$")
title(L"comparison of $E(t)$ for single realization")
legend()

savefig("energy_comparison.eps", dpi=240)
savefig("energy_comparison.png", dpi=240)

# Make plots: energy budgets for a realization of the Ito integration
fig, axs = subplots(nrows=2, figsize=(8, 6), sharex=true)

sca(axs[1]); cla()
axs[1][:plot](μ*T, E_theory,
                linewidth=3, label="theoretical \$\\langle E\\rangle\$")
axs[1][:plot](μ*T, mean(E_ito, dims=2), linestyle="--",
                linewidth=2, label="\$\\langle E\\rangle\$ from $nens ensemble member(s)")
xlabel(L"$\mu t$")
ylabel(L"$E$")
if stochastic; title(L"Ito: $\mathrm{d}E_t = (-2\mu E_t + \frac{\sigma}{2})\mathrm{d}t + \sqrt{\sigma}X_t\mathrm{d}W$")
else;          title(L"Ito: $\mathrm{d}X/\mathrm{d}t = -\mu X + \sqrt{\sigma}$")
end
legend()

sca(axs[2]); cla()
plot(μ*T[1:nsteps-1], dEdt_ito[1:nsteps-1, 1],
                linestyle="--", linewidth=2,
                label="numerical \$d\\langle E\\rangle/dt\$")
plot(μ*T[1:nsteps-1], work_ito[1:nsteps-1, 1] - diss_ito[1:nsteps-1, 1],
                linestyle="-.", linewidth=1,
                label="\$\\langle\$work \$-\$ dissipation\$\\rangle\$")
plot(μ*T[1:nsteps-1], dEdt_theory[1:nsteps-1],
                linewidth=3, label="theoretical \$d\\langle E\\rangle/dt\$")
xlabel(L"$\mu t$")
ylabel(L"$dE/dt$")

legend()

savefig("energy_budgets_Ito.eps", dpi=240)
savefig("energy_budgets_Ito.png", dpi=240)

# Make plots: energy budgets for a realization of the Stratonovich integration
fig, axs = subplots(nrows=2, figsize=(8, 6), sharex=true)

sca(axs[1]); cla()
plot(μ*T, E_theory,
                linewidth=3, label="theoretical \$\\langle E\\rangle\$")
plot(μ*T, mean(E_str, dims=2),
                linestyle="--", linewidth=2,
                    label="\$\\langle E\\rangle\$ from $nens ensemble member(s)")
xlabel(L"$\mu t$")
ylabel(L"$E$")
if stochastic; title(L"Stratonovich: $\mathrm{d}E_t = -2\mu E_t\mathrm{d}t + \sqrt{\sigma}x_t\circ\mathrm{d}W$")
else;          title(L"Stratonovich: $\mathrm{d}E/\mathrm{d}t = -2\mu E + \sqrt{\sigma} \dot{x}$")
end
legend()

sca(axs[2]); cla()
plot(μ*T[1:nsteps-1], dEdt_str[1:nsteps-1],
                linestyle="--", linewidth=2,
                label="numerical \$d\\langle E\\rangle/dt\$")
plot(μ*T[1:nsteps-1], work_str[1:nsteps-1] - diss_str[1:nsteps-1],
                linestyle="-.", linewidth=1,
                label="\$\\langle\$work \$-\$ dissipation\$\\rangle\$")
plot(μ*T[1:nsteps-1], dEdt_theory[1:nsteps-1],
                linewidth=3, label="theoretical \$d\\langle E\\rangle/dt\$")
xlabel(L"$\mu t$")
ylabel(L"$dE/dt$")

legend()

savefig("energy_budgets_Stratonovich.eps", dpi=240)
savefig("energy_budgets_Stratonovich.png", dpi=240)


## Seeded PageRank animation
# We are going to make an animation where we move from vertex to vertex
# following the nearest vertex in a seeded PageRank simulation

using MatrixNetworks
using Plots

function seeded_pagerank_animation(A,xy,seed,α,nsteps)
  curseed = seed
  started = Set{Int}() # this is the list of points we've started.
  @assert(nsteps <= size(A,1))
  @assert(is_connected((A)))
  anim = @animate for i=1:nsteps
    push!(started, curseed)
    f2 = seeded_pagerank(A, α, curseed)
    p = sortperm(f2) # this gets things to draw in the right z-order
    scatter(xy[p,1],xy[p,2], marker_z=-abs.(log10.(f2[p])).^0.5,colorbar=false,alpha=0.5,
      color=:magma,
      size=(1200,1200),
      markerstrokecolor=nothing,markerstrokewidth=0,label="",framestyle=:none,background=:black)
    curseed = -1
    for v in sortperm(f2,rev=true)
      if v in started
        continue
      else
        curseed = v
        break
      end
    end
  end
  return anim
end
A = load_matrix_network_all("geoclusters-example").A
xy = load_matrix_network_all("geoclusters-example").xy
anim = seeded_pagerank_animation(A,xy,1,0.5,100)
gif(anim,"geoclusters-animation.gif",fps=10)

## Now use the actual artistsim graph
names = readlines("../2-unit-1-demos/artistsim.names")
data = readdlm("../2-unit-1-demos/artistsim.smat")
xy = readdlm("artistsim.xy")
A = sparse(Int.(data[2:end,1]).+1,
           Int.(data[2:end,2]).+1,
           Int.(data[2:end,3]),
           Int(data[1,1]),Int(data[1,1]))
Asym = max.(A,A') # make the graph undirected
Acc,ccfilt = largest_component(Asym)
xycc = xy[ccfilt,:]

##
anim = seeded_pagerank_animation(Acc,xycc,200,0.5,200)
gif(anim,"artistsim-seeded-pagerank-animation.gif",fps=10)


using MathProgBaseMosek
using MathProgBase
using Base.Test

#const MPB = MathProgBase
include(joinpath(Pkg.dir("MathProgBase","test","linprog.jl")))
include(joinpath(Pkg.dir("MathProgBase","test","mixintprog.jl")))
include(joinpath(Pkg.dir("MathProgBase","test","quadprog.jl")))
include(joinpath(Pkg.dir("MathProgBase","test","nlp.jl")))
include(joinpath(Pkg.dir("MathProgBase","test","conicinterface.jl")))
include(joinpath(Pkg.dir("MathProgBase","test","linproginterface.jl")))

function test_mathprogbase(solver)
    @testset "linprog" begin
        linprogtest(solver)
    end

    @testset "mip" begin
        mixintprogtest(solver)
    end

    @testset "qp" begin
        #quadprogtest(solver)
        #qpdualtest(solver)
        #socptest(solver)
    end

    @testset "nlp" begin
        #nlptest(solver)
        #nlptest_nohessian(solver)
        #convexnlptest(solver)
        #rosenbrocktest(solver)
    end

    @testset "conic" begin
        coniclineartest(solver, duals=true)
        # Test conic fallback for LPs
        #coniclineartest(solver, duals=true)
    end

    @testset "linproginterface" begin
        linprogsolvertest(solver)
        linprogsolvertestextra(solver)
        # Test LP fallback for conics
        linprogsolvertest(solver)
    end

end

test_mathprogbase(MosekSolver(QUIET = true))

if (try Pkg.installed("Convex") != nothing catch false end)
    include("testconvex.jl")
end


@defcomp damages begin
    regions = Index()

    DAMFRAC = Variable(index=[time, regions]) # Damages as % of GDP

    TATM = Parameter(index=[time]) # Increase temperature of atmosphere (degrees C from 1900)
    SLRDAMAGES = Parameter(index=[time, regions])
    a1 = Parameter(index=[regions]) # Damage intercept
    a2 = Parameter(index=[regions]) # Damage quadratic term
    a3 = Parameter(index=[regions]) # Damage exponent

    function run_timestep(p, v, d, t)

        #Define function for DAMFRAC
        for r in d.regions
            v.DAMFRAC[t,r] = (((p.a1[r] * p.TATM[t]) + (p.a2[r] * p.TATM[t]^p.a3[r])) / 100) + (p.SLRDAMAGES[t,r] / 100)
        end
    end
end


include("TwistedHessianEC.jl")
using .TwistedHessianEC

function partI(m::W, n:: W, P::S, Q::S, Pother::S, Qother::S, l::U, e::V, E::T) where {T <: EllipticCurve, S <: EllipticCurveElem, U <: Integer, V <: Integer, W <: Integer}
    K = P*m + Q*n
    Ei = E
    Pi = Pother
    Qi = Qother
    for i in V(1):e
        Ki = (l^(e-i))*K
        phi = Isogeny(Ei, Ki)
        K = phi(K)
        Pi = phi(Pi)
        Qi = phi(Qi)
        Ei = phi.codomain
    end
    return Ei, Pi, Qi
end

function partII(m::W, n:: W, P::S, Q::S, l::U, e::V, E::T) where {T <: EllipticCurve, S <: EllipticCurveElem, U <: Integer, V <: Integer, W <: Integer}
    K = P*m + Q*n
    Ei = E
    for i in V(1):e
        Ki = (l^(e-i))*K
        phi = Isogeny(Ei, Ki)
        K = phi(K)
        Ei = phi.codomain
    end
    return jInvariant(Ei)
end

# Example:
#Public params
la = big(2)
lb = big(3)
ea = 63
eb = 41
OGp = la^ea*lb^eb*11-1
F = Fpsquare(OGp)
E = TwiHesEC(F, one(F), F([3, 2458499195747164226075539892235780232099]))

XP = F([big(264972986297324027107883370923360079519), big(483679027370140600263719808627289108176)])
YP = F([big(2550992266319579452440794756588926041584), big(2214394007774501716861032109362868970653)])
Z = one(F)
P = E(XP, YP, Z)

XQ = F([big(1287714346072412371952478415842888698876), big(427368407186067938537508979187813002488)])
YQ = F([big(528075014001520129686916245492679497447), big(857779714716146821676003610752395199871)])
Q = E(XQ, YQ, Z)

Pa = normalized(P*(big(3)^41*11))
Qa = normalized(Q*(big(3)^41*11))

Pb = normalized(P*(big(2)^63*11))
Qb = normalized(Q*(big(2)^63*11))

# Alice and Bob chooses secret params
# They cannot both be divisible prime by "their" prime...
function genSecret(la, lb)
    ma = rand(big(1):OGp-1)
    na = rand(big(1):OGp-1)
    while ma%la == 0 && na%la == 0
        ma = rand(big(1):OGp-1)
        na = rand(big(1):OGp-1)
    end
    mb = rand(big(1):OGp-1)
    nb = rand(big(1):OGp-1)
    while mb%lb == 0 && nb%lb == 0
        mb = rand(big(1):OGp-1)
        nb = rand(big(1):OGp-1)
    end
    return ma, na, mb, nb
end

ma, na, mb, nb = genSecret(la, lb)

# Alice do her part
E_A, P_B, Q_B = partI(ma, na, Pa, Qa, Pb, Qb, la, ea, E)
# Bob does his part
E_B, P_A, Q_A = partI(mb, nb, Pb, Qb, Pa, Qa, lb, eb, E)

#Both derive the shared key, hopefully it is equal
key_a = partII(ma, na, P_A, Q_A, la, ea, E_B)
key_b = partII(mb, nb, P_B, Q_B, lb, eb, E_A)
println("Alice key: $key_a")
println("Bob's key: $key_b")

# Del


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -3.412981244406108 <= q <= 4.209108534415106)

                     add_NL_constraint(m, :(softplus(-0.4372399910599052 + -0.43792093299500046*$(x[1]) + -0.746120027144912*$(x[2]) + 0.8963217772222452*$(x[3])) + softplus(0.8353036360644039 + 0.6875424158053405*$(x[1]) + 0.5184113900400429*$(x[2]) + -0.5247283462030659*$(x[3])) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


# we use warmstart to hint at a solution that has better objective value in pipe
# dimensioning than the dual bound of junction location?!

using MathOptInterface
using JuMP
using SCIP

using PipeLayout
using PipeLayout.JuncLoc
using PipeLayout.PipeDim
import PipeLayout: ploss_coeff_nice

const MOI = MathOptInterface

# instance data
pipelength = 100.0
width = height = pipelength / sqrt(2)
terminals = [Node(0,0), Node(width,height), Node(2width,0), Node(3width,height)]
demand = [400, -600, 600, -400]
bounds = fill(Bounds(40, 80), 4)
diams = map(Diameter, [1, 2], [1, 2])
inst = Instance(terminals, demand, bounds, diams, ploss_coeff_nice)

_scip = JuMP.optimizer_with_attributes(SCIP.Optimizer, "limits/gap" => 0.001)

function pipedim()
    println("### Pipe Dimensioning LP model, path topology ###")

    # topology for path
    path = Topology(terminals, [Arc(2,1), Arc(2,3), Arc(4,3)])

    # instead of calling `optimize` we do the steps manually
    solver = PipeDim.LP(_scip)
    model, π, l = PipeDim.make_model(inst, path, solver.lpsolver)
    JuMP.optimize!(model)

    println(model)
    @show JuMP.value.(l)
    @show JuMP.value.(π)
end

function juncloc()
    println("### Junction Location NLP model, FST with fixed locations ###")

    # topology data for FST
    junctions = [Node(0,-1), Node(0,-2)]
    nodes = vcat(terminals, junctions)
    fstw = Topology(nodes, [Arc(5,6), Arc(5,1), Arc(2,5), Arc(3,6), Arc(6,4)])

    # instead of calling `optimize` we do the steps manually
    solver = JuncLoc.NLP(_scip)
    model, x, y, L, l, π = JuncLoc.make_nlp(inst, fstw, solver)

    # fix values for the Steiner node positions
    # 1st steiner is at 2nd term
    @constraint(model, x[5] == inst.nodes[2].x)
    @constraint(model, y[5] == inst.nodes[2].y)
    # 2nd steiner is at 3rd term
    @constraint(model, x[6] == inst.nodes[3].x)
    @constraint(model, y[6] == inst.nodes[3].y)

    JuMP.optimize!(model)

    println(model)
    @show JuMP.value.(L)
    @show JuMP.value.(l)
    @show JuMP.value.(π)
end

pipedim()
println()
juncloc()

# demands = [400, -600, 600, -400]
#
# LP:                flow   flow^2
# 1: 1 --> 2  (a)    400     40000
# 2: 2 --> 3  (b)    200    160000
# 3: 3 --> 4  (c)    400     40000
#
#      2   4
#     / \ /
#    1   3
#
#
# NLP:                flow   flow squared
# 1: 5 --> 6 (b)      200       40000
# 2: 5 --> 1 (a)      400      160000
# 3: 2 --> 5 (-)      600      360000
# 4: 3 --> 6 (-)     -600      360000
# 5: 6 --> 4 (c)     -400      160000
#
#       2     _4
#     _ 5 -- 6
#    1       3
#


# for Q dynamics set α = r-λ*(exp(μ+0.5*δ^2)-1)
@sde_model Merton begin
  dS = S*α*dt + S*σ*dW + S*(ξ-1.0)*dN
  J ~ LogNormal(μ, δ)
end α σ λ μ δ

function sample(model::Merton, scheme::Exact, t0, x0)
  @unpack α, σ, λ, μ, δ = model
  N = rand(Poisson(model.λ*scheme.Δt))
  rand(LogNormal(log(x0)+(α-0.5*σ^2)*scheme.Δt+N*μ, sqrt(N*δ^2+scheme.Δt*σ^2)))
end


using LinearAlgebra
using Plots
using Images
using ImageView

n = 1000
X = zeros(n, n)
X[250:750, 250:750] .= 1
U, S, V = svd(X)
imshow(X)
p1 = plot(S, yaxis = :log)

nAngles = 12 # sweep through 12 angles, from 0:4:44

for j = 2:nAngles
    Y = imrotate(X, -(j-1)*4*pi/180, axes(X)) # rotate by (j-1)*4
    Y[isnan.(Y)] .= 0
    U2, S2, V2 = svd(Y)
    plot!(S2)
end


using Gridap

u_2d(x) = x[1]+x[2]
u_3d(x) = x[1]+x[2]+x[3]
domain_2d = (0,1,0,1)
domain_3d = (0,1,0,1,0,1)
cells_2d = (4,4)
cells_3d = (4,4,4)

for cells in (cells_2d,cells_3d)
  domain = length(cells) == 3 ? domain_3d : domain_2d
  u = length(cells) == 3 ? u_3d : u_2d
  model = CartesianDiscreteModel(domain,cells)
  order = 1
  reffe = ReferenceFE(lagrangian,Float64,order)
  V = TestFESpace(model,reffe,dirichlet_tags="boundary")
  U = TrialFESpace(u,V)
  Ω = Triangulation(model)
  dΩ = Measure(Ω,2*order)
  a(u,v) = ∫( ∇(u)⋅∇(v) )dΩ
  l(v) = 0
  op = AffineFEOperator(a,l,U,V)
  uh = solve(op)
  sum(∫( abs2(u-uh) )dΩ)
end


using Turing, Distributions
using Test
using Random

Random.seed!(125)

@model gdemo() = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  1.5 ~ Normal(m, sqrt(s))
  2.0 ~ Normal(m, sqrt(s))
  return s, m
end

N = 50
s1 = IPMCMC(10, N, 4, 2)
s2 = IPMCMC(10, N, 4)

c1 = sample(gdemo(), s1)
c2 = sample(gdemo(), s2)

# Very loose bound, only for testing constructor.
for c in [c1, c2]
  @test mean(c[:s]) ≈ 49/24 atol=1.0
  @test mean(c[:m]) ≈ 7/6 atol=1.0
end


using Gurobi
using Mosek
using JuMP
using Random
function convexRegress(Y, X, gamma, cut_fraction=1.0, epsilon = 1e-5)
    n = size(X,1); p = size(X,2);
    m = Model();
    set_optimizer(m, Gurobi.Optimizer);
    @variable(m, theta[1:n])
    @variable(m, ksi[1:n, 1:p])
    @objective(m, Min, 0.5*sum((Y-theta).^2) + 1/(2*gamma)*sum([sum(ksi[i,:].^2) for i=1:n-1]))
    # Creating constraints for ordered thetas
    for i=1:size(ksi,1)-1
        @constraint(m, theta[i] + sum(ksi[i,:].*(X[i+1,:] - X[i,:])) <= theta[i+1])
        @constraint(m, theta[i+1] + sum(ksi[i+1,:].*(X[i,:] - X[i+1,:])) <= theta[i])
    end
    @constraint(m, theta[n] + sum(ksi[n,:].*(X[1,:] - X[n,:])) <= theta[1])
    @constraint(m, theta[1] + sum(ksi[1,:].*(X[n,:] - X[1,:])) <= theta[n])
    # Adding constraint generator, since callbacks don't work
    violation = epsilon;
    cut_count = 1;
    while violation >= epsilon && cut_count <= size(X,1)^2/2*cut_fraction
        optimize!(m)
        ksis = getvalue.(ksi); thetas = getvalue.(theta);
        for i=1:n
            violation = 0;
            cut_index = 0;
            for k=1:n
                if (thetas[i] - thetas[k] + sum(ksis[i,:].*(X[k,:] - X[i,:]))) >= violation && (i != k)
                    violation = thetas[i] - thetas[k] + sum(ksis[i,:].*(X[k,:] - X[i,:]));
                    cut_index = k;
                end
            end
            if violation > epsilon
                @constraint(m, thetas[i] + sum(ksis[i,:].*(X[cut_index,:] - X[i,:])) <= thetas[cut_index])
                cut_count += 1;
#                 println("Cut ", cut_count, " added: ", [i,cut_index])
            end
        end
    end
    println("Cuts added: ", cut_count)
    return getvalue.(theta), getvalue.(ksi)
end

# function convexRegressDual(Y, X, gamma, k)
#     n = size(X,1); p = size(X,2);
#     m = Model(Gurobi.Optimizer)
#     @variable(m, z[1:p], Bin)
#     @variable(m, nu)
#     @variable(m, mu[1:n, 1:n])
#     # Objective
#     @objective(m, Max, nu)
#     # Subgradient and objective f'ns
#     function gz(Y, X, mu, z)
#         n = size(X,1); p = size(X,2);
#         return -0.5*sum([(y[i] + sum(mu[:,i]) - sum(mu[i,:])).^2 for i=1:n]) -
#                 0.5*gamma*sum(sum())
#     end
#     function dgz(Y, X, mu, z)
#         n = size(X,1); p = size(X,2);
#         dg = zeros(p);
#         for i = 1:p
#             dg[i] = -0.5/gamma*sum(sum(mu[i,j]*(X[i,p] - X[j,p]) for j=1:n)^2 for i=1:n)
#         end
#         return dg
#     end

#     # Initialize cuts
#     z0 = zeros(p)
#     z0[1:k] .= 1;

#     return getvalue.(theta), getvalue.(ksi)
# end


@doc raw"""
    velocity_curve(p, c::Curve, backend)

Velocity curve for a given curve `p` of type `c` calculated using differentiation
defined by `backend`.
If $c$ is a function such that $c\colon [0, 1] \to M$, then
`velocity_curve(c)(t)` is a vector tangent to `c` at `t`.
"""
function velocity_curve(M::AbstractManifold, p, backend::AbstractRiemannianDiffBackend)
    return t -> Manifolds.differential(M, p, t, backend)
end

function velocity_curve(M::AbstractManifold, p, backend::ProjectedDifferenceBackend)
    return t -> project(M.M, p(t), (embed(p(t + backend.dt)) - embed(p(t))) / backend.dt)
end

@doc raw"""
    curve_length(p, c::FunctionCurveSpace, a = 0.0, b = 1.0, dtype = Val(:continuous))

Returns length of the curve `c` between parameters `a` and `b`. By default
`a = 0.0` and `b = 1.0`. Calculates velocity using method `dtype`.
"""
function curve_length(
    M::AbstractManifold,
    p,
    a=0.0,
    b=1.0,
    backend=Manifolds.rdifferential_backend(),
)
    v = velocity_curve(M, p, backend)
    #TODO: add quadrature tolerances to curve_length?
    #TODO: add velocity calculation options
    i_val, error = quadgk(t -> norm(M, p(t), v(t)), a, b)
    return i_val
end

"""
    q_function(M::AbstractManifold, p, X)

Rescales given tangent vector `X` at point `p` to square root of its original length.
"""
function q_function(M::AbstractManifold, p, X)
    norm_X = norm(M, p, X)
    return if norm_X ≈ 0
        return zero_vector(M, p)
    else
        return (1 / sqrt(norm_X)) * X
    end
end

"""
    srvf(M, f, backend)

Calculate the Square Root Velocity Function of a given curve `f`.
Parameter `backend` describes type of differentiaion used. Supported values:
* `Val(:continuous)` -- symbolic differentiation resulting in a continuous
curve,
* `Val(:forward_diff)` -- forward finite differentiaion of a discretized curve
`c` based on its grid.
"""
function srvf(M::AbstractManifold, f, backend::AbstractRiemannianDiffBackend)
    vel = velocity_curve(M, f, backend)
    return t -> q_function(M, f(t), vel(t))
end

"""
    transport_srvf(M::AbstractManifold, c_p, c_X, p)

Transports all tangent vectors `c_X` of curve `c_p` to a given
point `p` along shortest geodesics.
`M` is the respective manifold of curves.
"""
function transport_srvf(M::AbstractManifold, c_p, c_X, p)
    c_out = allocate_result(M, transport_srvf, c_X, c_p)
    transport_srvf!(M, c_out, c_p, c_X, p)
    return c_out
end

"""
    tsrvf(M::AbstractManifold, c, p, backend::AbstractRiemannianDiffBackend)

Calculate SRVF of curve `c` from manifold `M` using method `backend` and transport to `p`.
"""
function tsrvf(M::AbstractManifold, c, p, backend::AbstractRiemannianDiffBackend)
    c_X = srvf(M, c, backend)
    return transport_srvf(M, c, c_X, p)
end

function reverse_srvf(M::AbstractManifold, c_X, initial_point)
    c_out = allocate_result(M, reverse_srvf, c_X)
    reverse_srvf!(M, c_out, c_X, initial_point)
    return c_out
end

"""
    srsf(M::CurveWarpingSpace, p)

Square Root Slope Function of curve warping `p`.
"""
function srsf(M::CurveWarpingSpace, p)
    grid = M.knots
    coeffs = map(p, grid)
    grads = [
        (coeffs[i + 1] - coeffs[i]) / (grid[i + 1] - grid[i]) for i in 1:(length(grid) - 1)
    ]
    push!(grads, grads[end])
    vals = map(sqrt, grads)
    return project(CurveWarpingSRSFSpace(grid), vals)
end

"""
    reverse_srsf(c)

Reverse Square Root Slope Function of a given SRSF of a curve warping.
"""
function reverse_srsf(M::CurveWarpingSpace, p)
    grid = M.knots
    steps = diff(grid)
    ys = accumulate(+, [zero(eltype(p)); steps .* p[1:(end - 1)] .^ 2])
    ys ./= ys[end]
    return make_warping(M, ys)
end
function reverse_srsf(M::CurveWarpingSpace, p::Identity{WarpingCompositionOperation})
    return reverse_srsf(M, identity_element(CurveWarpingGroup(M), p))
end


"""
    maxwellian(u, ρ, U, λ)
    maxwellian(u, prim)
    maxwellian(u, v, ρ, U, V, λ)
    maxwellian(u, v, prim)
    maxwellian(u, v, w, ρ, U, V, W, λ)
    maxwellian(u, v, w, prim)

Maxwellian in discrete form

* @args: particle velocity quadrature points
* @args: density, velocity and inverse of temperature
* @return: Maxwellian distribution function

"""
maxwellian(u, ρ, U, λ) = @. ρ * sqrt(λ / π) * exp(-λ * (u - U)^2) # 1V

maxwellian(u::Union{Real,AbstractVector{T}}, prim) where {T<:Real} =
    maxwellian(u, prim[1], prim[2], prim[end]) # in case of input with length 4/5

#--- 2V ---#
maxwellian(u, v, ρ, U, V, λ) = @. ρ * (λ / π) * exp(-λ * ((u - U)^2 + (v - V)^2))

maxwellian(u, v, prim::AbstractVector{T}) where {T<:Real} =
    maxwellian(u, v, prim[1], prim[2], prim[3], prim[end]) # in case of input with length 5

#--- 3V ---#
maxwellian(u, v, w, ρ, U, V, W, λ) =
    @. ρ * sqrt((λ / π)^3) * exp(-λ * ((u - U)^2 + (v - V)^2 + (w - W)^2))

maxwellian(u, v, w, prim::AbstractVector{T}) where {T<:Real} =
    maxwellian(u, v, w, prim[1], prim[2], prim[3], prim[4], prim[5])

"""
    * Maxwellian
    maxwellian!(M, u, ρ, U, λ)
    maxwellian!(M, u, ρ, U, V, λ)
    maxwellian!(M, u, ρ, U, V, W, λ)
    maxwellian!(M, u, prim)
    maxwellian!(M, u, v, prim)
    maxwellian!(M, u, v, w, prim)

    * Rykov
    maxwellian!(Ht, Bt, Rt, Hr, Br, Rr, u, prim, K, Kr)
    maxwellian!(Ht, Bt, Rt, Hr, Br, Rr, u, v, prim, K, Kr)

In-place Maxwellian

* @args: particle velocity quadrature points
* @args: density, velocity and inverse of temperature
* @return: Maxwellian distribution function

"""
function maxwellian!(
    M::AbstractVector{T1},
    u::AbstractVector{T2},
    ρ,
    U,
    λ,
) where {T1<:AbstractFloat,T2<:AbstractFloat}
    @. M = ρ * sqrt(λ / π) * exp(-λ * (u - U)^2) # 1V
    return nothing
end

maxwellian!(
    M::AbstractVector{T1},
    u::AbstractVector{T2},
    prim::AbstractVector{T3},
) where {T1<:AbstractFloat,T2<:AbstractFloat,T3<:Real} =
    maxwellian!(M, u, prim[1], prim[2], prim[end])

# Rykov
function maxwellian!(
    Ht::T1,
    Bt::T1,
    Rt::T1,
    Hr::T1,
    Br::T1,
    Rr::T1,
    u::T2,
    prim::T3,
    K,
    Kr,
) where {
    T1<:AbstractVector{<:AbstractFloat},
    T2<:AbstractVector{<:AbstractFloat},
    T3<:AbstractVector{<:Real},
}
    @. Ht = prim[1] * sqrt(prim[4] / π) * exp(-prim[4] * (u - prim[2])^2)
    @. Bt = Ht * K / (2.0 * prim[4])
    @. Rt = Ht * Kr / (2.0 * prim[5])

    @. Hr = prim[1] * sqrt(prim[3] / π) * exp(-prim[3] * (u - prim[2])^2)
    @. Br = Hr * K / (2.0 * prim[3])
    @. Rr = Hr * Kr / (2.0 * prim[3])

    return nothing
end

#--- 2V ---#
function maxwellian!(
    M::T1,
    u::T2,
    v::T2,
    ρ,
    U,
    V,
    λ,
) where {T1<:AbstractArray{<:AbstractFloat,2},T2<:AbstractArray{<:AbstractFloat}}
    @. M = ρ * (λ / π) * exp(-λ * ((u - U)^2 + (v - V)^2))
    return nothing
end

maxwellian!(
    M::T1,
    u::T2,
    v::T2,
    prim::T3,
) where {
    T1<:AbstractArray{<:AbstractFloat,2},
    T2<:AbstractArray{<:AbstractFloat},
    T3<:AbstractArray{<:Real,1},
} = maxwellian!(M, u, v, prim[1], prim[2], prim[3], prim[end])

# Rykov
function maxwellian!(
    Ht::T1,
    Bt::T1,
    Rt::T1,
    Hr::T1,
    Br::T1,
    Rr::T1,
    u::T2,
    v::T2,
    prim::T3,
    K,
    Kr,
) where {
    T1<:AbstractArray{<:AbstractFloat,2},
    T2<:AbstractArray{<:AbstractFloat},
    T3<:AbstractArray{<:Real,1},
}
    @. Ht = prim[1] * (prim[5] / π) * exp(-prim[5] * ((u - prim[2])^2 + (v - prim[3])^2))
    @. Bt = Ht * K / (2.0 * prim[5])
    @. Rt = Ht * Kr / (2.0 * prim[6])

    @. Hr = prim[1] * (prim[4] / π) * exp(-prim[4] * ((u - prim[2])^2 + (v - prim[3])^2))
    @. Br = Hr * K / (2.0 * prim[4])
    @. Rr = Hr * Kr / (2.0 * prim[4])

    return nothing
end

#--- 3V ---#
function maxwellian!(
    M::T1,
    u::T2,
    v::T2,
    w::T2,
    ρ,
    U,
    V,
    W,
    λ,
) where {T1<:AbstractArray{<:AbstractFloat,3},T2<:AbstractArray{<:AbstractFloat}}
    @. M = ρ * sqrt((λ / π)^3) * exp(-λ * ((u - U)^2 + (v - V)^2 + (w - W)^2))
    return nothing
end

maxwellian!(
    M::T1,
    u::T2,
    v::T2,
    w::T2,
    prim::T3,
) where {
    T1<:AbstractArray{<:AbstractFloat,3},
    T2<:AbstractArray{<:AbstractFloat},
    T3<:AbstractArray{<:Real,1},
} = maxwellian!(M, u, v, w, prim[1], prim[2], prim[3], prim[4], prim[5])

"""
    mixture_maxwellian(u, prim)
    mixture_maxwellian(u, v, prim)
    mixture_maxwellian(u, v, w, prim)

Multi-component Maxwellian in discrete form
"""
function mixture_maxwellian(
    u::AbstractMatrix{T1},
    prim::AbstractMatrix{T2},
) where {T1<:Real,T2<:Real}

    mixM = similar(u)
    for j in axes(mixM, 2)
        mixM[:, j] .= maxwellian(u[:, j], prim[:, j])
    end

    return mixM

end

function mixture_maxwellian(
    u::X,
    v::X,
    prim::AbstractMatrix{Y},
) where {X<:AbstractArray{<:AbstractFloat,3},Y<:Real}

    mixM = similar(u)
    for k in axes(mixM, 3)
        mixM[:, :, k] .= maxwellian(u[:, :, k], v[:, :, k], prim[:, k])
    end

    return mixM

end

function mixture_maxwellian(
    u::X,
    v::X,
    w::X,
    prim::AbstractMatrix{Y},
) where {X<:AbstractArray{<:AbstractFloat,4},Y<:Real}

    mixM = similar(u)
    for l in axes(mixM, 4)
        mixM[:, :, :, l] .=
            maxwellian(u[:, :, :, l], v[:, :, :, l], w[:, :, :, l], prim[:, l])
    end

    return mixM

end

"""
    mixture_maxwellian!(M, u, prim)
    mixture_maxwellian!(M, u, v, prim)
    mixture_maxwellian!(M, u, v, w, prim)

In-place multi-component Maxwellian

"""
function mixture_maxwellian!(
    M::AbstractMatrix{T1},
    u::T2,
    prim::AbstractMatrix{T3},
) where {T1<:Real,T2<:AbstractArray{<:Real,2},T3<:Real}

    for j in axes(M, 2)
        _M = @view M[:, j]
        maxwellian!(_M, u[:, j], prim[:, j])
    end

    return nothing

end

function mixture_maxwellian!(
    M::AbstractArray{T1,3},
    u::T2,
    v::T2,
    prim::AbstractMatrix{T3},
) where {T1<:Real,T2<:AbstractArray{<:Real,3},T3<:Real}

    for k in axes(M, 3)
        _M = @view M[:, :, k]
        maxwellian!(_M, u[:, :, k], v[:, :, k], prim[:, k])
    end

    return nothing

end

function mixture_maxwellian!(
    M::AbstractArray{T1,4},
    u::T2,
    v::T2,
    w::T2,
    prim::AbstractMatrix{T3},
) where {T1<:Real,T2<:AbstractArray{<:Real,4},T3<:Real}

    for l in axes(M, 4)
        _M = @view M[:, :, :, l]
        maxwellian!(_M, u[:, :, :, l], v[:, :, :, l], w[:, :, :, l], prim[:, l])
    end

    return nothing

end


using DifferentialEquations
using Plots
using Random: Random
########### Defines the model parameters for Figure 1 ###########
n_parasites = 100; # number of parasites
K = 100; # carrying capacity
c = 4.0; # host contact rate
ux = 0.2; # mortality rate of uninfected hosts
# mortality rate of infected hosts (always >= ux)
Random.seed!(1479);
ui = rand(200:1000, n_parasites) / 1000;
βy = 3 * (ui .- ux) ./ (ui .- ux .+ 1) # horizontal transmission rate
bx = 1.0 # birth rate of uninfected individuals
ei = fill(bx .- bi, n_parasites) # number of infected offspring from an infected host
# set up the model parameters
parameters = (bx=bx, βy=βy, ei=ei, c=c, K=K, ux=ux, bi=bi, ui=ui);
# define the numbers of individuals in each category
Y = zeros(Float64, n_parasites); # number of parasites of each strain
Y[1] = 1.0; # start with one initial parasite
X0 = 80; # number of initial uninfected hosts
new_U = vcat(X0, Y); # merge both together into one array
# define the limits of the simulation
windowstart = 0.0;
windowsize = 1000.0;
windowend = windowstart + windowsize;
N = zeros(Float64, (n_parasites + 1, (n_parasites - 1) * Int(windowsize) + 1));
figure_directory = "Figure1" # output folder for subplots of Figure 1


#
# inventory control problem
#
using DiscreteEvents, Distributions, Random

mutable struct Station
    q::Float64           # fuel amount
    t::Vector{Float64}   # time vector
    qt::Vector{Float64}  # fuel/time vector
    cs::Int              # customers served
    cl::Int              # customers lost
    qs::Float64          # fuel sold
    ql::Float64          # lost sales
end

function customer(c::Clock, s::Station, X::Distribution)
    function fuel(s::Station, x::Float64)
        s.q -= x             # take x from tank
        push!(s.t, c.time)   # record time, amount, customer, sale
        push!(s.qt, s.q)
        s.cs += 1
        s.qs += x
    end

    x = rand(X)              # calculate demand
    if s.q ≥ x               # demand can be met
        fuel(s, x)
    elseif s.q ≥ a           # only partially in stock
        s.ql += x - s.q      # count the loss
        fuel(s, s.q)         # give em all we have
    else
        s.cl += 1            # count the lost customer
        s.ql += x            # count the lost demand
    end
end

function replenish(c::Clock, s::Station, Q::Float64)
    if s.q < a
        push!(s.t, c.time)
        push!(s.qt, s.q)
        s.q += Q
        push!(s.t, c.time)
        push!(s.qt, s.q)
    end
end

Random.seed!(123)
const λ = 0.5      # ~ every two minutes a customer
const ρ = 1/180    # ~ every 3 hours a replenishment truck
const μ = 30       # ~ mean demand per customer
const σ = 10       #   standard deviation
const a = 5        #   minimum amount
const Q = 6000.0   # replenishment amount
const M₁ = Exponential(1/λ)  # customer arrival time distribution
const M₂ = Exponential(1/ρ)  # replenishment time distribution
const X = truncated(Normal(μ, σ), a, Inf)  # demand distribution

clock = Clock()
s = Station(Q, Float64[0.0], Float64[Q], 0, 0, 0.0, 0.0)
@event replenish(clock, s, Q) every M₂
@event customer(clock, s, X) every M₁
println(@run! clock 5000)

@show fuel_sold = s.qs;
@show loss_rate = s.ql/s.qs;
@show served_customers = s.cs;
@show lost_customers = s.cl;

# using Plots
# plot(s.t, s.qt, title="Filling Station", xlabel="time [min]", ylabel="inventory [L]", legend=false)
# savefig("invctrl.png")


# Parameters
#	N::Int
#	- the problem size
# Return Values
#	YOR::Array{Array{Array{SparseMatrixCSC, 1}, 1}, 1} (Young's Orthogonal Representations)
#	- YOR[n][p][k] is Young's Orthogonal Representation for the Adjacent Transposition (K, K + 1) for the pth Partition of n
#	PT::Array{Array{Array{Int, 1}, 1}, 1} (Partition Tree)
#	- output1 from partition_tree()
function yor(N::Int)
	P, WI = partitions(N)
	PT = partition_tree(N, P, WI)
	YS = ys_symbols(N, P, PT)
	YOR = Array(Array{Array{SparseMatrixCSC, 1}, 1}, N)
	for n = 1:N
		YSn = YS[n]
		Pn = P[n]
		Pn_L = length(Pn)
		YORn = Array(Array{SparseMatrixCSC, 1}, Pn_L)
		for p = 1:Pn_L
			YSnp = YSn[p]
			Pnp = Pn[p]
			R = length(Pnp)
			YORn[p] = yor_p(n, R, YSnp)
		end
		YOR[n] = YORn
	end
	return YOR, PT
end

# Parameters:
#	N::Int
#	- the size of P
#	R::Int
#	- the length of P
#	P::Array{Int, 1}
#	- P is a Partition of N
# Return Values:
#	YORp::Array{SparseMatrixCSC, 1}
#	- YORp[k] is the Young's Orthogonal Representation of P for the Adjacent Transposistion (k, k + 1)
function yor_p(N::Int, R::Int, YSymbols::Array{Array{Int, 1}, 1})
	YORp = Array(SparseMatrixCSC, N - 1)
	L = length(YSymbols)
	DA, IA = ys_information(N, R, YSymbols)
	for k = 1:(N - 1)
		YORp[k] = yor_pk(DA, IA, L, k)
	end
	return YORp
end

#Parameters
#	DA::Array{Int, 2}
#	- output1 from ys_information()
#	IA::Array{Int, 2}
#	- outpu2 from ys_information()
#	L::Int
#	- the number of Yamanouchi Symbols for the current Partition
#	K::Int
#	- represents the Adjacent Transposition (K, K + 1)
#Return Values
#	YORpk::SparseMatrixCSC
#	- Young's Orthongal Representation of the Adjacent Transposition (K, K + 1) for the Partition p
function yor_pk(DA::Array{Int, 2}, IA::Array{Int, 2}, L::Int, k::Int)
	n = L
	for i = 1:L
		if IA[i,k] != 0
			n += 1
		end
	end
	colptr = Array(Int, L + 1)
	colptr[1] = 1
	rowval = Array(Int, n)
	nzval = Array(Float64, n)
	cp = 1
	for i = 1:L
		D = 1/DA[i, k]
		I = IA[i,k]
		if I == 0
			colptr[i + 1] = colptr[i] + 1
			rowval[cp] = i
			nzval[cp] = D
			cp += 1
		else
			colptr[i + 1] = colptr[i] + 2
			if I > i
				rowval[cp] = i
				nzval[cp] = D
				cp += 1
				rowval[cp] = I
				nzval[cp] = sqrt(1 - D * D)
				cp += 1
			else
				rowval[cp] = I
				nzval[cp] = sqrt(1 - D * D)
				cp += 1
				rowval[cp] = i
				nzval[cp] = D
				cp += 1
			end
		end
	end
	YORpk = SparseMatrixCSC(L, L, colptr, rowval, nzval)
	return YORpk
end


using Nullables

#############
# 2D scenes #
#############

struct Point
    x::Float64
    y::Float64
end

Point(arr::Vector{Float64}) = Point(arr[1], arr[2])

array(point::Point) = [point.x, point.y]

function dist(a::Point, b::Point)
    dx = a.x - b.x
    dy = a.y - b.y
    sqrt(dx * dx + dy * dy)
end

Base.:+(a::Point, b::Point) = Point(a.x + b.x, a.y + b.y)
Base.:*(a::Real, b::Point) = Point(b.x * a, b.y * a)

abstract type Obstacle end

"""
A 2D canvas with a set of obstacles
"""
mutable struct Scene
    xmin::Float64
    xmax::Float64
    ymin::Float64
    ymax::Float64
    obstacles::Array{Obstacle,1}
end

function Scene(xmin::Real, xmax::Real, ymin::Real, ymax::Real)
    Scene(xmin, xmax, ymin, ymax, Obstacle[])
end

function add!(scene::Scene, object::Obstacle)
    push!(scene.obstacles, object)
end

function ccw(a::Point, b::Point, c::Point)
    (b.x - a.x) * (c.y - a.y) - (c.x - a.x) * (b.y - a.y)
end

function line_intersects_line(a1::Point, a2::Point, b1::Point, b2::Point)
        # http://algs4.cs.princeton.edu/91primitives/
        if ccw(a1, a2, b1) * ccw(a1, a2, b2) > 0
            return false
        end
        if ccw(b1, b2, a1) * ccw(b1, b2, a2) > 0
            return false
        end
        return true
end

function line_of_site(scene::Scene, a::Point, b::Point)
    for obstacle in scene.obstacles
        if intersects_path(obstacle, a, b)
            return false
        end
    end
    return true
end

"""
An arbitrary polygon
"""
struct Polygon <: Obstacle
    vertices::Array{Point,1}
end

function intersects_path(poly::Polygon, path_start::Point, path_end::Point)
    n = length(poly.vertices)
    for start_vertex_idx=1:n
        end_vertex_idx = start_vertex_idx % n + 1 # loop over to 1
        v1 = poly.vertices[start_vertex_idx]
        v2 = poly.vertices[end_vertex_idx]
        if line_intersects_line(v1, v2, path_start, path_end)
            return true
        end
    end
    return false
end

"""
A line with thickness
"""
struct Wall <: Obstacle
    start::Point
    orientation::Int # x is 1, y is 2
    length::Float64
    thickness::Float64
    height::Float64
    poly::Polygon
    function Wall(start::Point, orientation::Int, length::Float64,
                  thickness::Float64, height::Float64)
        if orientation != 1 && orientation != 2
            error("orientation must be either 1 (x) or 2 (y)")
        end
        vertices = Vector{Point}(undef, 4)
        vertices[1] = start
        dx = orientation == 1 ? length : thickness
        dy = orientation == 2 ? length : thickness
        vertices[2] = Point(start.x + dx, start.y)
        vertices[3] = Point(start.x + dx, start.y + dy)
        vertices[4] = Point(start.x, start.y + dy)
        poly = Polygon(vertices)
        new(start, orientation, length, thickness, height, poly)
    end
end

function intersects_path(wall::Wall, path_start::Point, path_end::Point)
    intersects_path(wall.poly, path_start, path_end)
end

"""
A square
"""
struct Tree <: Obstacle
    center::Point
    size::Float64
    poly::Polygon
    function Tree(center::Point; size::Float64=10.)
        vertices = Vector{Point}(undef, 4)
        vertices[1] = Point(center.x - size/2, center.y - size/2)
        vertices[2] = Point(center.x + size/2, center.y - size/2)
        vertices[3] = Point(center.x + size/2, center.y + size/2)
        vertices[4] = Point(center.x - size/2, center.y + size/2)
        poly = Polygon(vertices)
        new(center, size, poly)
    end
end

function intersects_path(tree::Tree, path_start::Point, path_end::Point)
    intersects_path(tree.poly, path_start, path_end)
end

##########
# render #
##########

using PyPlot

const patches = PyPlot.matplotlib[:patches]

function render(polygon::Polygon, ax)
    num_sides = length(polygon.vertices)
    vertices = Matrix{Float64}(undef, num_sides, 2)
    for (i, pt) in enumerate(polygon.vertices)
        vertices[i,:] = [pt.x, pt.y]
    end
    poly = patches[:Polygon](vertices, true, facecolor="black")
    ax[:add_patch](poly)
end

render(wall::Wall, ax) = render(wall.poly, ax)
render(tree::Tree, ax) = render(tree.poly, ax)

function render(scene::Scene, ax)
    ax[:set_xlim]((scene.xmin, scene.xmax))
    ax[:set_ylim]((scene.ymin, scene.ymax))
    for obstacle in scene.obstacles
        render(obstacle, ax)
    end
end


###
### **acb_elliptic.h** -- elliptic integrals and functions of complex variables
###

### Complete elliptic integrals
arbcall"void acb_elliptic_k(acb_t res, const acb_t m, slong prec)"
arbcall"void acb_elliptic_k_jet(acb_ptr res, const acb_t m, slong len, slong prec)"
arbcall"void _acb_elliptic_k_series(acb_ptr res, acb_srcptr m, slong mlen, slong len, slong prec)"
arbcall"void acb_elliptic_k_series(acb_poly_t res, const acb_poly_t m, slong len, slong prec)"
arbcall"void acb_elliptic_e(acb_t res, const acb_t m, slong prec)"
arbcall"void acb_elliptic_pi(acb_t res, const acb_t n, const acb_t m, slong prec)"

### Legendre incomplete elliptic integrals
arbcall"void acb_elliptic_f(acb_t res, const acb_t phi, const acb_t m, int pi, slong prec)"
arbcall"void acb_elliptic_e_inc(acb_t res, const acb_t phi, const acb_t m, int pi, slong prec)"
arbcall"void acb_elliptic_pi_inc(acb_t res, const acb_t n, const acb_t phi, const acb_t m, int pi, slong prec)"

### Carlson symmetric elliptic integrals
arbcall"void acb_elliptic_rf(acb_t res, const acb_t x, const acb_t y, const acb_t z, int flags, slong prec)"
arbcall"void acb_elliptic_rg(acb_t res, const acb_t x, const acb_t y, const acb_t z, int flags, slong prec)"
arbcall"void acb_elliptic_rj(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)"
arbcall"void acb_elliptic_rj_carlson(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)"
arbcall"void acb_elliptic_rj_integration(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)"
arbcall"void acb_elliptic_rc1(acb_t res, const acb_t x, slong prec)"

### Weierstrass elliptic functions
arbcall"void acb_elliptic_p(acb_t res, const acb_t z, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_p_jet(acb_ptr res, const acb_t z, const acb_t tau, slong len, slong prec)"
arbcall"void _acb_elliptic_p_series(acb_ptr res, acb_srcptr z, slong zlen, const acb_t tau, slong len, slong prec)"
arbcall"void acb_elliptic_p_series(acb_poly_t res, const acb_poly_t z, const acb_t tau, slong len, slong prec)"
arbcall"void acb_elliptic_invariants(acb_t g2, acb_t g3, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_roots(acb_t e1, acb_t e2, acb_t e3, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_inv_p(acb_t res, const acb_t z, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_zeta(acb_t res, const acb_t z, const acb_t tau, slong prec)"
arbcall"void acb_elliptic_sigma(acb_t res, const acb_t z, const acb_t tau, slong prec)"


function second_moment(q::Normal)
    return abs2(mean(q)) + var(q)
end

function second_moment(q::Normal, y::Real)
    return abs2(mean(q) - y) + var(q)
end

# This is not exactly an approximation but corresponds anyway
# to the expectation of the function σ(f)
function approx_expected_logistic(μ::Real, c::Real)
    lower, upper = LogExpFunctions._logistic_bounds(μ)
    return μ < lower ? zero(μ) : (μ > upper ? one(μ) : exp(μ / 2) * sech(c / 2) / 2)
end

# Same thing as above for softmax
function approx_expected_logisticsoftmax(
    μ::AbstractVector, c::AbstractVector, θ::AbstractVector
)
    σs = θ[1:(end - 1)] .* approx_expected_logistic.(μ, c)
    return σs / (θ[end] * logistic(0) + sum(σs))
end

transpose_nested(x::ArrayOfSimilarArrays) = nestedview(flatview(x)')