text:
using Polynomials
# using PyPlot

include("mala.jl")
include("zigzagtools.jl")

function maximalTravelTime(xi, theta, xi_stars, current_region)

  (dimension, n_regions)  = size(xi_stars)

  # crossingtimes = Vector{Float64}(n_regions)
  # for i=1:n_regions
  #   crossingtimes[i]
  #
  crossingtimes = sum((0.5*(xi_stars[:,current_region] .+ xi_stars) .- xi).* (xi_stars[:,current_region] .- xi_stars), 1)./(sum(theta .* (xi_stars[:,current_region] .- xi_stars), 1))
  crossingtimes[(crossingtimes .< 0) | (crossingtimes .==-Inf)] = Inf
  crossingtimes[current_region] = Inf
  return findmin(crossingtimes)
end

function currentRegion(xi, xi_stars)
  (dimension, n_regions)  = size(xi_stars)
  distances = sum((xi .- xi_stars).^2,1)
  return findmin(distances)[2]

end

function determineExtrema(dataX::Vector{Float64}, dataY::Vector{Int64}, rho::Real)

  val = rho + 2 * sum((1./(1 + exp(-dataX/2)) - dataY).*dataX)
  if (val <=0)
    bimodal = true
    n_regions = 3
    xi_stars = Array{Float64,2}(2, n_regions)
    # and so xi_1^* >= 1/2
    f = z -> rho*(z+1/2) + sum((1./(1+exp(-(z+1/2)*dataX))-dataY).*dataX)
    xi1star = findRootMonotoneFunction(f)+1/2
    f = z-> rho/2 + sum((1./(1+exp(-(z^2+1/2)*dataX))-dataY).*dataX)
    xi2hat = findRootMonotoneFunction(f)
    println("Bimodal, xi_1^* = ", xi1star, ", \\hat xi_2 = ", xi2hat)
    xi_stars[:,1] = [0.5, xi2hat]
    xi_stars[:,2] = [0.5,-xi2hat]
    xi_stars[:,3] = [xi1star, 0]
  else
    bimodal = false
    n_regions = 1
    xi_stars = Array{Float64,2}(2,n_regions)
    # and so xi_1^* < 1/2, have to look in reverse direction:
    f = z -> -rho*(1/2-z) - sum((1./(1+exp(-(1/2-z)*dataX))-dataY).*dataX)
    xi1star = 1/2-findRootMonotoneFunction(f)
    xi_stars[:,1] = [xi1star,0]
    println("Unimodal, xi_1^* = ", xi1star)
  end
  return xi_stars
end

function manyReferencePoints(xi_hat, n_referencepoints)
  println("Number of referencepoints: ", n_referencepoints)
  Xi_reference = Array{Float64,2}(2,n_referencepoints)
  xi2_max = 2 # hardcoded for the moment
  delta_xi2 = 2*xi2_max/(n_referencepoints-1)
  println("Delta xi_2 = ", delta_xi2)
  Xi_reference[2,:] = -xi2_max + delta_xi2 * Vector(0:n_referencepoints-1)
  Xi_reference[1,:] = xi_hat[1] .- Xi_reference[2,:].^2
  return Xi_reference
end

function cvNonidentifiableLogistic(dataX::Vector{Float64}, dataY::Vector{Int64}, n_proposedswitches::Int64, xi_init::Vector{Float64} = [0.0;0.0], rho::Real = 1)

  dim = 2;
  array_size = 1000000
  T = zeros(array_size);
  Xi = zeros(dim, array_size);
  Xi[:,1] = xi_init;
  Theta_current = ones(dim);
  Xi_current = Xi[:,1];
  n = length(dataX);
  a = zeros(2,3);
  extrema = determineExtrema(dataX, dataY, rho)
  xi_hat = extrema[:,end]; # MLE
  xi_stars = extrema

  # Gareth's idea
  # alpha = 0.5
  # xi_stars = manyReferencePoints(xi_hat, Integer(ceil(n^alpha)))

  count_proposed_switches = 0
  count_accepted_switches = 0
  current_region = currentRegion(Xi_current+.00000000001*Theta_current,xi_stars) # slight bit of symmetry breaking helps
  (T_max, new_region) = maximalTravelTime(Xi_current, Theta_current, xi_stars, current_region)
  Xi_star = xi_stars[:,current_region]
  x_max = maximum(abs(dataX))
  a[1,3]= n*x_max^2/2;
  a[2,3]= n*x_max^2*abs(Xi_star[2]);

  while count_proposed_switches < n_proposedswitches
    switched = false;
    T_travelled = 0;
    while(~switched)
      relpos = Xi_current - Xi_star
      bigfactor1 = n * x_max^2 * (abs(relpos[1])+2*max(abs(Xi_current[2]),abs(Xi_star[2]))*abs(relpos[2]))
      bigfactor2 = n * x_max^2 * (1 + 2 * abs(relpos[2]) + 2 * max(abs(Xi_current[2]),abs(Xi_star[2])))
      a[1,1]=Theta_current[1]*(relpos[1]) + bigfactor1/4;
      a[1,2]=rho+1/4*bigfactor2;
      a[2,1]=Theta_current[2]*rho*relpos[2] +2*n*x_max*abs(relpos[2]) + abs(Xi_star[2])*bigfactor1/2;
      a[2,2]=rho+2*n*x_max + abs(Xi_star[2])/2*bigfactor2;

      m = Vector{Poly}(dim); # vector of polynomials
      U = -log(rand(dim));
      T_proposed = zeros(dim);
      for j=1:dim
        if all(a[j,:] .<= 0)
          T_proposed[j] = Inf
        else
          m[j] = Poly(a[j,:]);
          # println("j = ", j)
          T_proposed[j] = findRootPolyMaxZero(m[j], U[j]);
        end
        # println("T[", j, "] = ", T_proposed[j])
      end
      (T_min,index_to_switch) = findmin(T_proposed);
      if T_min > T_max # we are entering a new region
        println("Entering region around ", xi_stars[:,new_region])
        T_travelled += T_max
        Xi_current +=  T_max * Theta_current
        current_region = new_region
        (T_max, new_region) = maximalTravelTime(Xi_current +0.00000000001*Theta_current, Theta_current, xi_stars, current_region)
        Xi_star = xi_stars[:,current_region]
      else # we can check if we are actually switching
        count_proposed_switches += 1
        T_travelled += T_min;
        T_max -= T_min # we have a smaller distance to travel to new region
        Xi_current += T_min * Theta_current;
        m_current = m[index_to_switch](T_min);
        J = rand(1:n) # random observation to use
        h_current = 1/(1 + exp(-(Xi_current[1]+Xi_current[2]^2)*dataX[J]));
        h_star = 1/(1 + exp(-(Xi_star[1]+Xi_star[2]^2)*dataX[J]))
        rel_pos = Xi_current - Xi_star
        if index_to_switch == 1
          lambda_current = pp(Theta_current[1]*(rho*rel_pos[1] + n*dataX[J]*(h_current -h_star)))
        else
          lambda_current = pp(Theta_current[2]*(rho*rel_pos[2] + 2*n*dataX[J]*((h_current - dataY[J])*Xi_current[2] - (h_star - dataY[J])*Xi_star[2])))
        end
        V = rand();
        if (lambda_current > m_current + 0.00001)
          println("CV: bound: ", m_current, ", actual value: ", lambda_current)
          println("CV: intermediate bound: ", Theta_current[2]*rho*rel_pos[2] + 2*n*abs(dataX[J])*(abs(rel_pos[2])+abs(Xi_star[2])*abs(dataX[J])*(abs(rel_pos[1])/4 + max(abs(Xi_current[2]),abs(Xi_star[2]))*abs(rel_pos[2]))))
          error("CV: Bound exceeded in iteration ", i, " while switching component ", index_to_switch)
        end
#      println("CV: Xi_current = ", Xi_current)
        if (V <= lambda_current/m_current)
          Theta_current[index_to_switch] = -Theta_current[index_to_switch];
          switched = true;
#        println("switched component ", index_to_switch)
          if count_accepted_switches+2 > array_size
            T_extra = Vector{Float64}(array_size)
            Xi_extra = Array{Float64}(2,array_size)
            T_new = [T; T_extra]
            Xi_new = [Xi Xi_extra]
            T = T_new
            Xi = Xi_new
            array_size *= 2
          end
          T[count_accepted_switches+2] = T[count_accepted_switches + 1] + T_travelled;
          Xi[:,count_accepted_switches+2] = Xi_current;
          count_accepted_switches += 1;
          (T_max,new_region) = maximalTravelTime(Xi_current, Theta_current,xi_stars,current_region)
        end
      end # of if to check whether we switched regions
    end # of while loop until a switch
  end # of for loop over iterations
  println("CV: Fraction of accepted switches: ", count_accepted_switches/count_proposed_switches)
  return (T[1:count_accepted_switches+1],Xi[:,1:count_accepted_switches+1])
end

function zzNonidentifiableLogistic(dataX::Vector{Float64}, dataY::Vector{Int64}, n_proposedswitches::Int64, xi_init::Vector{Float64} = [0.0;0.0], rho::Real = 1)

  dim = 2;
  T = zeros(n_proposedswitches+2);
  Xi = zeros(dim, n_proposedswitches+2);
  Xi[:,1] = xi_init;
  Theta_current = ones(dim);
  Xi_current = Xi[:,1];
  n = length(dataX);
  a = zeros(dim, 2); # coefficients of upper bound polynomial for 0 to 1st order
  a[1,2] = rho
  a[2,2] = rho + 2 * sum(abs(dataX))
  count_proposed_switches = 0
  count_accepted_switches = 0
  m = Vector{Poly}(dim); # vector of polynomials

  while count_proposed_switches < n_proposedswitches
    switched = false;
    T_travelled = 0;
    while(~switched)
      a[1,1] = Theta_current[1]*rho*Xi_current[1] + sum(abs(dataX));
      a[2,1] = Theta_current[2]*rho*Xi_current[2] + 2*abs(Xi_current[2])*sum(abs(dataX));
      # a[2,1] = n *(Theta_current[2] *(rho/n - y_bar) * Xi_current[2] + 2 * Xi_current[1]^2 + 2 * Xi_current[2]^2 + 8 * abs(Xi_current[2])^3);
      U = -log(rand(dim));
      T_proposed = zeros(dim);
      count_proposed_switches += 1
      for j=1:dim
        if all(a[j,:] .<= 0)
          T_proposed[j] = Inf
        else
          m[j] = Poly(a[j,:]);
          # println("j = ", j)
          T_proposed[j] = findRootPolyMaxZero(m[j], U[j]);
        end
        # println("T[", j, "] = ", T_proposed[j])
      end
      (T_min,index_to_switch) = findmin(T_proposed);
      T_travelled += T_min;
      Xi_current += T_min * Theta_current;
      m_current = m[index_to_switch](T_min);
      h_vec = 1 + exp(-(Xi_current[1] + Xi_current[2]^2)*dataX)
      big_sum = sum((1./h_vec  - dataY).*dataX)
      if (index_to_switch == 1)
        lambda_current = pp(Theta_current[1]*(rho *Xi_current[1] + big_sum))
      else
        lambda_current = pp(Theta_current[2]*(rho * Xi_current[2] + 2 * Xi_current[2] * big_sum));
      end
      V = rand();
      if (lambda_current > m_current + 0.001)
        println("bound: ", m_current, ", actual value: ", lambda_current)
        error("Bound exceeded in proposed switch ", count_proposed_switches, " while switching component ", index_to_switch)

      end
      if (V <= lambda_current/m_current)
        Theta_current[index_to_switch] = -Theta_current[index_to_switch];
        switched = true;
#        println("switched component ", index_to_switch)
        T[count_accepted_switches+2] = T[count_accepted_switches+1] + T_travelled;
        Xi[:,count_accepted_switches+2] = Xi_current;
        count_accepted_switches +=1;
      end
    end # of while loop until a switch
  end # of for loop over iterations
  println("ZZ: Fraction of accepted switches: ",count_accepted_switches/count_proposed_switches)
  return (T[1:count_accepted_switches+1],Xi[:,1:count_accepted_switches+1])
end # of function zzNonidentifiableGaussian

function consensusNonidentifiableLogistic(dataX::Vector{Float64}, dataY::Vector{Int64}, n_batches::Int64, stepsize::Float64, n_samples::Int64, xi_init::Vector{Float64} = [0.0;0.0], rho::Real = 1)

  n = length(dataX)
  batch_size = div(n,n_batches)
  Xbatches = reshape(dataX, n_batches, batch_size)
  Ybatches = reshape(dataY, n_batches, batch_size)
  # batch_means = mean(batches, 2)
  rho_batch = rho / n_batches
  # samples = zeros(2, n_samples, n_batches)
  U_fun = Vector{Function}(n_batches)
  grad_U = Vector{Function}(n_batches)
  for i=1:n_batches
    U_fun[i] = xi -> rho_batch/2*sumabs2(xi) + sum(log(1 + exp(g(xi).*Xbatches[i,:]))- g(xi) * Xbatches[i,:] .* Ybatches[i,:])
    grad_U[i] = xi -> [rho_batch * xi[1] + sum((1./(1+exp(-g(xi).*Xbatches[i,:]))-Ybatches[i,:]) .* Xbatches[i,:]);
                       rho_batch * xi[2] + 2 * xi[2] * sum((1./(1+exp(-g(xi).*Xbatches[i,:]))-Ybatches[i,:]).* Xbatches[i,:])]
  end

  return consensusMALA(U_fun, grad_U, stepsize, 2, n_samples, xi_init)
end

function SGLDNonidentifiableLogistic(dataX::Vector{Float64}, dataY::Vector{Int64},n_batches::Int64,stepsize::Float64, n_samples::Int64, xi_init::Vector{Float64} = [0.0;0.0], rho::Real = 1)

  n = length(dataX)
  batch_size = div(n,n_batches)
  Xbatches = reshape(dataX, n_batches, batch_size)
  Ybatches = reshape(dataY, n_batches, batch_size)
  rho_batch = rho / n_batches
  U_fun = Vector{Function}(n_batches)
  grad_U = Vector{Function}(n_batches)
  for i=1:n_batches
    # U_fun[i] = xi -> rho_batch/2*sumabs2(xi) + sum(log(1 + exp(g(xi).*Xbatches[i,:]))- g(xi) * Xbatches[i,:] .* Ybatches[i,:])
    grad_U[i] = xi -> [rho_batch * xi[1] + sum((1./(1+exp(-g(xi).*Xbatches[i,:]))-Ybatches[i,:]) .* Xbatches[i,:]);
                       rho_batch * xi[2] + 2 * xi[2] * sum((1./(1+exp(-g(xi).*Xbatches[i,:]))-Ybatches[i,:]).* Xbatches[i,:])]
  end
  I = rand(1:n_batches, n_samples)
  Xi = zeros(2, n_samples + 1)
  Xi_current = xi_init
  Xi[:,1] = Xi_current
  Z = randn(n_samples)

  for i=1:n_samples
    Xi_current += -stepsize * n/batch_size * grad_U[I[i]](Xi_current) + sqrt(2 * stepsize) * Z[i]
    Xi[:,i+1] = Xi_current
  end

  return Xi

end

max_stars_repo_name: jbierkens/zigzag-experiments

max_issues_repo_name: jbierkens/zigzag-experiments

max_forks_repo_name: jbierkens/zigzag-experiments


text:
"""
 Tests for https://arxiv.org/abs/1303.5103 with test problems as in the paper.
 DRI1, RI1, RI3, RI5, RI6, RDI1WM, RDI2WM, RDI3WM, RDI4WM
"""

import Statistics # for mean values of trajectories
import LinearAlgebra # for the normn
using StochasticDiffEq
using Test
using Random
using DiffEqDevTools
#using DiffEqGPU

function prob_func(prob, i, repeat)
    remake(prob,seed=seeds[i])
end

"""
 Test non-commutative noise SDEs (iip)
"""

@info "Non-commutative noise"

u₀ = [1.0,1.0]
function f2!(du,u,p,t)
  du[1] = -273//512*u[1]
  du[2] = -1//160*u[1]-(-785//512+sqrt(2)/8)*u[2]
end
function g2!(du,u,p,t)
  du[1,1] = 1//4*u[1]
  du[1,2] = 1//16*u[1]
  du[2,1] = (1-2*sqrt(2))/4*u[1]
  du[2,2] = 1//10*u[1]+1//16*u[2]
end
dts = 1 .//2 .^(3:-1:0)
tspan = (0.0,3.0)

h2(z) = z^2 # but apply it only to u[1]

prob = SDEProblem(f2!,g2!,u₀,tspan,noise_rate_prototype=zeros(2,2))
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h2(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(1e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,DRI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("DRI1:", sim.𝒪est[:weak_final])

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
sim = test_convergence(dts,ensemble_prob,RI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order 3.43
println("RI1:", sim.𝒪est[:weak_final])

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
sim = test_convergence(dts,ensemble_prob,RI3(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order 2.46
println("RI3:", sim.𝒪est[:weak_final])

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
sim = test_convergence(dts,ensemble_prob,RI5(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order 2.57
println("RI5:", sim.𝒪est[:weak_final])

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)
sim = test_convergence(dts,ensemble_prob,RI6(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI6:", sim.𝒪est[:weak_final])

numtraj = Int(1e3)
seed = 10
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI1WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI1WM:", sim.𝒪est[:weak_final])

numtraj = Int(1e5)
seed = 10
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI2WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order 2.517769274990593
println("RDI2WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI3WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI3WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI4WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=exp(-3.0)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI4WM:", sim.𝒪est[:weak_final])

ext: jl


text:
module NES

include("emulation.jl")

export reset!, loadgame, step!, stepframe!, stepframes!, frame, screen, setbuttons1!, setbuttons2!, setaudiosamplesperframe!

end # module


text:
""" TransformedIndex """
struct TransformedIndex{T,I,F} <: AbstractArray{T,1}
    index::I
    tform::F

    function TransformedIndex(index::AbstractRange{Int}, tform::CoordinateTransform{<:Any,NO}) where {NO}
        @assert !isempty(index)
        if NO === 1
            return new{Int,typeof(index),typeof(tform)}(index, tform)
        else
            return new{CartesianIndex{NO},typeof(index),typeof(tform)}(index, tform)
        end
    end
end

Base.length(x::TransformedIndex) = length(getfield(x, :index))

ArrayInterface.known_size(::Type{<:TransformedIndex{I}}) where {I} = known_size(I)

Base.IteratorSize(::Type{<:TransformedIndex}) = Base.HasLength()

@propagate_inbounds function Base.getindex(x::TransformedIndex, i::Vararg{Any})
    getfield(x, :tform)(getfield(x, :index)[i...])
end
@inline function Base.nextind(x::TransformedIndex{I,F}, state::Int) where {I,F}
    if known_step(I) === missing
        return state + step(getfield(x, :index))
    else
        return state + known_step(I)
    end
end
@inline function Base.isdone(x::TransformedIndex{I,F}, state::Int) where {I,F}
    if known_last(I) === missing
        return state === last(getfield(x, :index))
    else
        return state === known_last(I)
    end
end

@inline function Base.iterate(x::TransformedIndex{I}) where {I}
    if known_first(I) === missing
        state = first(getfield(x, :index))
    else
        state = known_first(I)
    end
    return getfield(x, :tform)(state), state
end

@inline function Base.iterate(x::TransformedIndex, state::Int)
    if isdone(x, state)
        return nothing
    else
        newstate = nextind(x, state)
        return getfield(x, :tform)(newstate), newstate
    end
end

""" GetIndex """
struct GetIndex{inbounds,X,NI,NO} <: CoordinateTransform{NI,NO}
    x::X

    function GetIndex{inbounds}(a::A) where {inbounds,A}
        new{inbounds::Bool,A,ndims(A),dynamic(ndims_index(A))}(A)
    end
end

(t::GetIndex{true})(i) = @inbounds(getfield(t, :x)[i])
function (t::GetIndex{false})(i)
    data = getfield(t, :x)
    @boundscheck checkbounds(data, i)
    return @inbounds(data[i])
end

""" TransformedIndices """
struct TransformedIndices{T,N,I,F} <: AbstractArray{T,N}
    indices::I
    tform::F
end

""" reindex """
reindex(x::AbstractRange) = x
@inline function reindex(x::AbstractArray{T,N}) where {T,N}
    if N === 1
        return TransformedIndex(ArrayInterface.indices(x), GetIndex{true}(x))
    else
        return x
    end
end
reindex(x::Integer) = Int(x)
reindex(x::StaticInt) = x


text:
include("Quadratures.jl")
@reexport using Gridap.Quadratures

include("DuffyQuadratures.jl")
@reexport using Gridap.DuffyQuadratures

include("QuadratureFactories.jl")
@reexport using Gridap.QuadratureFactories

include("CellQuadratures.jl")
@reexport using Gridap.CellQuadratures

include("Triangulations.jl")
@reexport using Gridap.Triangulations

include("CellIntegration.jl")
@reexport using Gridap.CellIntegration

include("BoundaryDescriptors.jl")
@reexport using Gridap.BoundaryDescriptors

include("BoundaryTriangulations.jl")
@reexport using Gridap.BoundaryTriangulations

include("SkeletonTriangulations.jl")
@reexport using Gridap.SkeletonTriangulations

include("NormalVectors.jl")
@reexport using Gridap.NormalVectors

include("BoundaryCellFields.jl")
@reexport using Gridap.BoundaryCellFields

include("SkeletonCellFields.jl")
@reexport using Gridap.SkeletonCellFields


text:
# Finds the radical of the given matrix Algebra by
# computing the center of the Algebra using
# the characteristic polynomial of the Gramian matrix
function find_radical(Algebra::Subspacik)
    As = basis(Algebra)
    F = base_ring(Algebra)
    n = dim(Algebra)

    traces = Dict{Tuple{Int, Int}, elem_type(F)}(
        (i, j) => tr(As[i] * As[j])
        for i in 1 : n
            for j in i : n
    )

    nnz_coords = [
        (i, j, traces[min(i, j), max(i, j)])
        for i in 1 : n
            for j in 1 : n
                if ! iszero(traces[min(i, j), max(i, j)])
    ]

    if isempty(nnz_coords)
        return As
    end

    ZZ = Nemo.GF(2^31 - 1)
    A = from_COO(n, n, nnz_coords, F)

    A = modular_reduction(A, ZZ)

    # wiedemannchik.jl
    PolySpace, _ = ZZ["x"]
    MSpace = MatrixSpace(ZZ, n, n)
    char_poly = minimal_polynomial_wiedemann(
        A,
        subspace_minpoly=__randomized_wiedemann_minpoly
    )
    # char_poly = charpoly(PolySpace, MSpace(to_dense(A)))

    radical_basis = []

    # hm??
    # Gleb: why "hm??"
    if (!iszero(coeff(char_poly, 0)))
        return radical_basis
    end

    while iszero(coeff(char_poly, 0))
        char_poly = shift_right_x(char_poly)
    end

    Image = evaluate_2(char_poly, A)

    Image = rational_reconstruction(Image)
    # Gleb: technically, there should be some way to check and rerun with larger prime
    # however, if this is not the winner-radical-algorithm, not needed really

    transpose!(Image)

    basis_cols = linear_span!(collect(values(Image.rows)))

    summator(x, y) = reduce(x, y, 1)
    for (piv, vect) in basis_cols.echelon_form
        vectors = []
        for j in 1 : size(Image, 2)
            if !iszero(vect[j])
                push!(vectors, scale(As[j], vect[j]))
            end
        end
        push!(radical_basis, reduce(summator, vectors))
    end

    @info "computed the radical of dimension $(length(radical_basis))"

    return radical_basis
end

# Finds the radical of the given matrix Algebra by
# computing the radical of the Algebra *directly*

"""
    find_radical_sup(Algebra)

Finds a basis of the radical of given matrix Algebra.
"""
function find_radical_sup(Algebra::Subspacik)
    As = basis(Algebra)
    F = base_ring(Algebra)
    n = dim(Algebra)

    if n == 0
        return []
    end

    A = gram_matrix(As)
    if isempty(A.nnz_rows)
        return As
    end

    pivs = zeros(Int64, n)
    V = Subspacik(F)
    for (piv, vec) in A.rows
        eat_sparsik!(V, vec)
    end

    orig_pivs = collect(keys(V.echelon_form))

    if length(orig_pivs) == n
        return []
    end

    radical_basis = []

    for i in orig_pivs
        pivs[i] = 1
    end

    for i in 1:n
       if pivs[i] == 0
           y = Sparsik(n, F, Int[], Dict{Int, elem_type(F)}())
           reduce!(y, i, one(F))
            for j in orig_pivs
                reduce!(y, j, -V.echelon_form[j][i])
            end
            push!(radical_basis, y)
        end
    end

    rad = []
    summator(x, y) = reduce(x, y, 1)
    for vect in radical_basis
        vectors = []
        for j in 1:length(As)
            if !iszero(vect[j])
                push!(vectors, scale(As[j], vect[j]))
            end
        end
        push!(rad, reduce(summator, vectors))
    end

    @info "computed the radical of dimension $(length(rad))"

    return rad
end

# Finds the radical of the given matrix Algebra by
# computing the center of the Algebra *directly*
# with the means of Nemo.jl functions
function find_radical_2(Algebra::Subspacik)
    As = basis(Algebra)
    F = base_ring(Algebra)
    n = dim(Algebra)

    # Gleb: to a function (see above)
    traces = Dict{Tuple{Int, Int}, elem_type(F)}(
        (i, j) => tr(As[i] * As[j])
        for i in 1 : n
            for j in i : n
    )

    nnz_coords = [
        (i, j, traces[min(i, j), max(i, j)])
        for i in 1 : n
            for j in 1 : n
                if ! iszero(traces[min(i, j), max(i, j)])
    ]

    if isempty(nnz_coords)
        return As
    end

    ZZ = Nemo.GF(2^31 - 1)
    A = zeros(ZZ, n, n)
    for (i, j, x) in nnz_coords
        A[i, j] = modular_reduction(x, ZZ)
    end

    MSpace = MatrixSpace(ZZ, n, n)
    sz, vectors = kernel(MSpace(A))

    sz == 0 && return As

    basis_cols = linear_span!([
                rational_reconstruction(
                    from_dense([vectors[:, i]...], ZZ)
                )
                for i in 1:sz
    ])

    radical_basis = []

    summator(x, y) = reduce(x, y, 1)
    for (piv, vect) in basis_cols.echelon_form
        vectors = []
        for j in 1:length(As)
            if !iszero(vect[j])
                push!(vectors, scale(As[j], vect[j]))
            end
        end
        push!(radical_basis, reduce(summator, vectors))
    end

    @info "computed the radical of dimension $(length(radical_basis))"

    return radical_basis
end

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

# returns an invariant subspace of the given Algebra
# in case the latter is semisimple
function invariant_subspace_semisimple(Algebra::Subspacik)
    es = basis(Algebra)
    n = size(first(es), 1)
    F = base_ring(Algebra)

    if n^2 == dim(Algebra)
        return []
    end

    # we want to do dense LA here
    MSpace = MatrixSpace(F, n, n)
    PSpace, x = F["x"]

    i, iters = 0, 100
    while i < iters
        M = random_element(Algebra, count = dim(Algebra))
        reconstruct!(M)

        chpoly = charpoly(PSpace, MSpace(to_dense(M)))

        factors = AbstractAlgebra.factor(chpoly).fac
        # factors is of type Fac
        # Fac.fac of form
        #   factor → degree  ,  (x-a)(x-b)..

        # mb we can use other factors to output more subspaces
        # Gleb: yes, this is very true, you can actually simply loop over the factors of the charpoly
        f = first(keys(factors))
        # f = irreducible^k
        f = f^factors[f]

        # computing the kernel of f(M)
        factored = evaluate(f, M)
        V = last(kernel(MSpace(to_dense(factored))))

        # convert to sparse repr
        V = [
            from_dense(v, Nemo.QQ)
            for v in [[V[:, j]...]
                for j in 1:size(V, 2)]
        ]

        # if is proper and invaiant
        if 0 < length(V) < n && check_invariance!(es, deepcopy(V))
            return V
        end

        i += 1
    end

    @warn "invariant subspaces exist but are not defined over Q, not implemented, sorry"
    # should we return hmm..
    return []
end

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

# finds the general kernel of the given matrices
function general_kernel(matrices::AbstractArray)
    stacked = vcat(matrices...)
    Space = MatrixSpace(Nemo.QQ, size(stacked)...)
    return Array(last(kernel(Space(stacked))))
end

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


text:
using ContinuousPOMDPTreeSearchExperiments
using ParticleFilters
using ARDESPOT
using BasicPOMCP
using POMCPOW
using POMDPs
using DiscreteValueIteration
using QMDP
using MCTS
using VDPTag2
using POMDPToolbox
using DataFrames

file_contents = readstring(@__FILE__())

pomdp = VDPTagPOMDP(mdp=VDPTagMDP(barriers=CardinalBarriers(0.2, 1.8)))
dpomdp = AODiscreteVDPTagPOMDP(pomdp, 30, 0.5)

@show max_time = 1.0
@show max_depth = 10
@show RO = RandomSolver
# @show RO = ToNextMLSolver

solvers = Dict{String, Union{Solver,Policy}}(
    "to_next" => ToNextML(mdp(pomdp)),
    "manage_uncertainty" => ManageUncertainty(pomdp, 0.01),

    "pomcpow" => begin
        rng = MersenneTwister(13)
        # ro = ToNextMLSolver(rng)::RO
        ro = RandomSolver(rng)::RO
        solver = POMCPOWSolver(tree_queries=10_000_000,
                               criterion=MaxUCB(100.0),
                               final_criterion=MaxQ(),
                               max_depth=max_depth,
                               max_time=max_time,
                               k_action=25.0,
                               alpha_action=1/20,
                               k_observation=6.0,
                               alpha_observation=1/100,
                               estimate_value=FORollout(ro),
                               next_action=RootToNextMLFirst(rng),
                               check_repeat_obs=false,
                               check_repeat_act=false,
                               default_action=ReportWhenUsed(TagAction(false, 0.0)),
                               rng=rng
                              )
    end,

    "pft" => begin
        rng = MersenneTwister(13)
        m = 10
        node_updater = ObsAdaptiveParticleFilter(deepcopy(pomdp),
                                           LowVarianceResampler(m),
                                           0.05, rng)
        # ro = ToNextMLSolver(rng)::RO
        ro = RandomSolver(rng)::RO
        ev = SampleRollout(solve(ro, pomdp), rng)
        solver = DPWSolver(n_iterations=typemax(Int),
                           exploration_constant=90.0,
                           depth=max_depth,
                           max_time=max_time,
                           k_action = 20.0,
                           alpha_action = 1/20,
                           k_state = 6.0,
                           alpha_state = 1/55,
                           check_repeat_state=false,
                           check_repeat_action=false,
                           estimate_value=ev,
                           next_action=RootToNextMLFirst(rng),
                           default_action=ReportWhenUsed(TagAction(false, 0.0)),
                           rng=rng
                          )
        belief_mdp = GenerativeBeliefMDP(deepcopy(pomdp), node_updater)
        solve(solver, belief_mdp)
    end,

    "pomcpdpw" => begin
        rng = MersenneTwister(13)
        # ro = ToNextMLSolver(rng)
        ro = RandomSolver(rng)::RO
        sol = PDPWSolver(max_depth=max_depth,
                    max_time=max_time,
                    c=65.0,
                    k_action=20.0,
                    alpha_action=1/10.0,
                    k_observation=4.0,
                    alpha_observation=1/100.0,
                    enable_action_pw=true,
                    check_repeat_obs=false,
                    check_repeat_act=false,
                    tree_queries=typemax(Int),
                    # default_action=ReportWhenUsed(1),
                    estimate_value=FORollout(ro),
                    next_action=RootToNextMLFirst(rng),
                    rng=rng
                   )
    end,

    "d_despot" => begin
        rng = MersenneTwister(13)
        # ro = ToNextMLSolver(rng)
        ro = ToNextMLSolver(rng)
        b = IndependentBounds(DefaultPolicyLB(ro), VDPUpper())
        sol = DESPOTSolver(lambda=0.01,
                     K=100,
                     D=max_depth,
                     max_trials=1_000_000,
                     T_max=max_time,
                     bounds=b,
                     random_source=MemorizingSource(500, 10, rng, min_reserve=8),
                     rng=rng)
        planner = solve(sol, dpomdp)
        translate_policy(planner, dpomdp, pomdp, dpomdp)
    end,

    "d_pomcp" => begin
        rng = MersenneTwister(13)
        ro = ToNextMLSolver(rng)
        sol = POMCPSolver(max_depth=max_depth,
                    max_time=max_time,
                    c=65.0,
                    tree_queries=typemax(Int),
                    default_action=ReportWhenUsed(1),
                    estimate_value=FORollout(ro),
                    rng=rng
                   )
        planner = solve(sol, dpomdp)
        translate_policy(planner, dpomdp, pomdp, dpomdp)
    end
)

@show N=1000

# alldata = DataFrame()
# for (k, solver) in solvers
# test = ["pomcpow", "pft"]
test = ["manage_uncertainty", "pomcpow", "pft"]
for (k, solver) in [(s, solvers[s]) for s in test]
    @show k
    if isa(solver, Solver)
        planner = solve(solver, pomdp)
    else
        planner = solver
    end
    sims = []
    for i in 1:N
        srand(planner, i+50_000)
        filter = SIRParticleFilter(deepcopy(pomdp), 100_000, rng=MersenneTwister(i+90_000))

        md = Dict(:solver=>k, :i=>i)
        sim = Sim(deepcopy(pomdp),
                  planner,
                  filter,
                  rng=MersenneTwister(i+70_000),
                  max_steps=100,
                  metadata=md
                 )

        push!(sims, sim)
    end

    data = run_parallel(sims)
    # data = run(sims)

    rs = data[:reward]
    println(@sprintf("reward: %6.3f ± %6.3f", mean(rs), std(rs)/sqrt(length(rs))))
end

#=
datestring = Dates.format(now(), "E_d_u_HH_MM")
copyname = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "icaps_2018", "data", "subhunt_table_$(datestring).jl")
write(copyname, file_contents)
filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "icaps_2018", "data", "subhunt_$(datestring).csv")
println("saving to $filename...")
writetable(filename, alldata)
println("done.")
=#


text:
# Uses new formulation from limbdark paper.
using Limbdark
import Limbdark: Transit_Struct, s2

#include("s2_stable.jl")
#include("../../../src/s2.jl")
#include("../../../src/define_constants.jl")

if VERSION >= v"0.7"
  using DelimitedFiles
end

using PyPlot
nb = 1001; nr = 1001
b=range(0.0,stop=2,length=nr)
r=range(0.0,stop=2,length=nr)

# We compute the `occultquad` flux in
# a separate Python call
fgrid = readdlm("s2_MA2002.txt")

fgrid_old = zeros(Float64,nr,nb)
fgrid_big = zeros(Float64,nr,nb)
resid = zeros(Float64,nr,nb)
for ib=1:nb
  for ir=1:nr
    fgrid_big[ir,ib]=convert(Float64,s2(big(r[ir]),big(b[ib])))
    resid[ir,ib] = abs.(fgrid[ir,ib]-fgrid_big[ir,ib])
    if resid[ir,ib] < 1e-15
      resid[ir,ib] = 1e-15
    end
  end
end

clf()
figure(figsize=(5, 4))
logresid = log10.(resid)
img=imshow(logresid,interpolation="nearest",origin="lower",
           extent=[0,2,0,2],cmap="plasma",vmin=-15, vmax=-3)
plot([0,1],[1,2],"k-")
plot([1,2],[0,1],"k-")
plot([0,0.38],[0,0.38],"k-")
plot([0.5,2],[0.5,2],"k-")
plot([0,1],[1,0],"k-")
annotate(L"$r = 1 + b$", xy=(0.5, 1.5), xytext=(0, 0), xycoords="data",
         textcoords="offset points", ha="center", va="bottom",
         fontsize=12, color="w", rotation=45)

annotate(L"$r = b$", xy=(1, 1), xytext=(0, 0), xycoords="data",
         textcoords="offset points", ha="center", va="bottom",
         fontsize=12, color="k", rotation=45)

annotate(L"$r = b - 1$", xy=(1.5, 0.5), xytext=(0, 0), xycoords="data",
         textcoords="offset points", ha="center", va="bottom",
         fontsize=12, color="k", rotation=45)

annotate(L"$r = 1 - b$", xy=(0.5, 0.5), xytext=(-9, -9), xycoords="data",
         textcoords="offset points", ha="center", va="center",
         fontsize=12, color="k", rotation=-45)

xlabel("Impact parameter")
ylabel("Radius ratio")
ticks = [-15, -12, -9, -6, -3]
cbar = colorbar(img, ticks=ticks)
cbar[:ax][:set_yticklabels]([L"$10^{-15}$",
                             L"$10^{-12}$",
                             L"$10^{-9}$",
                             L"$10^{-6}$",
                             L"$10^{-3}$"])

savefig("s2_residuals_MA2002.pdf")

max_issues_repo_name: rodluger/limbdark


text:
using Documenter, HFP

makedocs(modules = [HFP], sitename = "HFP.jl")

deploydocs(repo = "github.com/tdunning/HFP.jl.git")


text:
using NonlinearEigenproblems: nep_gallery, compute_Mder
using FEASTSolver

function info(Λ, X, residuals, c, r)
    inside(x) = in_contour(x, c, r)
    in_eig = Λ[inside.(Λ)]
    in_res = residuals[inside.(Λ)]
    print("number inside : ")
    println(size(Λ[inside.(Λ)])[1])
    if sum(inside.(Λ)) > 0
        in_res_conv = in_res[in_res .<= 1e-5]
        in_eig_conv = in_eig[in_res .<= 1e-5]
        print("number inside converged : ")
        println(size(in_eig_conv)[1])
        print("max res inside: ")
        println(maximum(residuals[inside.(Λ)]))
        if size(in_res_conv, 1) > 0
            print("max res inside non spurious: ")
            println(maximum(in_res_conv))
        end
    end
    # display(in_eig_conv)
end

nep = nep_gallery("nlevp_native_loaded_string")
T(x) = compute_Mder(nep, x)

n = size(nep, 1)
C = complex(800.0, 0.0)
R = 790

e, v, res = nlfeast_moments!(T, rand(ComplexF64,n,14), 2^4, 10, c=C, r=R, debug=true, moments=3, ϵ=10e-16, store=true)

info(e,v,res,C,R)


text:
## abstract types ##

"""
    AbstractHyperGraph

Abstract type that is the supertype of all hypergraph types.
"""
abstract type AbstractHyperGraph end

"""
    AbstractHyperEdge

Abstract type that is the supertype of all hyperedge types.
"""
abstract type AbstractHyperEdge end

abstract type AbstractIncidenceSet end

## concrete types ##

"""
    HyperEdge{T}

A type that represents unoriented, unweighted hyperedges. There are no restrictions to the structure; self-edges are allowed, and the hyperedge may be empty.

# Fields
- `V`: a set of vertices
"""
mutable struct HyperEdge{T} <: AbstractHyperEdge
    V::AbstractVector{T}
end
HyperEdge(v::T) where {T} = HyperEdge{T}([v])
HyperEdge{T}() where {T} = HyperEdge{T}([])
HyperEdge() = HyperEdge{Any}([])
Base.:(==)(e1::HyperEdge, e2::HyperEdge) = isequal(vertices(e1), vertices(e2))
Base.eltype(::HyperEdge{T}) where {T} = T

"""
    HyperGraph{T}

A type that represents unoriented, unweighted hypergraphs.

# Fields
- `V`: a set of vertices
- `E`: a set of hyperedges
"""
mutable struct HyperGraph{T} <: AbstractHyperGraph
    V::AbstractVector{T}
    E::AbstractVector{HyperEdge{T}}
end
HyperGraph(v::T, es::AbstractVector{HyperEdge{T}}) where {T} = HyperGraph{T}([v], es)
HyperGraph(vs::AbstractVector{T}, e::HyperEdge{T}) where {T} = HyperGraph{T}(vs, [e])
HyperGraph(v::T, e::HyperEdge{T}) where {T} = HyperGraph{T}([v], [e])
HyperGraph(es::AbstractVector{HyperEdge{T}}) where {T} = HyperGraph(union(vertices(es)...), es)
HyperGraph(e::HyperEdge{T}) where {T} = HyperGraph(union(vertices(e)), e)
HyperGraph{T}() where {T} = HyperGraph{T}([], [])
HyperGraph() = HyperGraph{Any}([], [])
(::Type{HyperGraph{T}})(e::HyperEdge{T}) where {T} = HyperGraph(e)
(::Type{HyperGraph{T}})(es::AbstractVector{HyperEdge{T}}) where {T} = HyperGraph(es)
Base.:(==)(x1::HyperGraph, x2::HyperGraph) = isequal(vertices(x1), vertices(x2)) && isequal(hyperedges(x1), hyperedges(x2))

Base.eltype(::HyperGraph{T}) where {T} = T

"""
    SpeciesSet{T}

A type that represents an incidence structure, that is a set of objects and their associated multiplicities, in the context of chemical reactions (where objects are chemical species and multiplicities are stoichiometries).

# Fields
- `objs`: a set of chemical species
- `mults`: a set of stoichiometries

# Notes
- multiplicities are restricted to integers
- if no stoichiometry is given, defaults to 1 for convenience
"""
mutable struct SpeciesSet{T} <: AbstractIncidenceSet
    objs::AbstractVector{T}
    mults::AbstractVector{Int}
    function SpeciesSet{T}(objs::AbstractVector{T}, mults::AbstractVector{Int}) where {T}
        (length(objs) == length(mults)) ? new{T}(objs, mults) : error("species and stoichiometry vectors must have the same length")
    end
end
SpeciesSet(species::AbstractVector{T}, stoich::AbstractVector{Int}) where {T} = SpeciesSet{T}(species, stoich)
SpeciesSet(species::T, stoich::Int) where {T} = SpeciesSet{T}([species], [stoich])
SpeciesSet(species::AbstractVector{T}) where {T} = SpeciesSet{T}(species, ones(Int, length(species)))
SpeciesSet(species::T) where {T} = SpeciesSet([species])
SpeciesSet{T}() where {T} = SpeciesSet{T}(T[], Int[])
SpeciesSet() = SpeciesSet{Any}()
species(S::SpeciesSet) = objects(S)
stoich(S::SpeciesSet) = multiplicities(S)
Base.:(==)(s1::SpeciesSet, s2::SpeciesSet) = isequal(species(s1), species(s2)) && isequal(stoich(s1), stoich(s2))
Base.eltype(::SpeciesSet{T}) where {T} = T
Base.length(S::SpeciesSet) = length(species(S))

"""
    ChemicalHyperEdge{T}

A type that represents a chemical hyperedge, that is an oriented, weighted hyperedge, in the context of chemical reactions (where the weight is the reaction rate).

# Fields
- `src`: a set of input species (negatively oriented)
- `tgt`: a set of output species (positively oriented)
- `weight`: reaction rate

# Notes
- if no rate is given, defaults to 1 for convenience
"""
mutable struct ChemicalHyperEdge{T} <: AbstractHyperEdge
    src::SpeciesSet{T}
    tgt::SpeciesSet{T}
    weight
end
ChemicalHyperEdge(inputs::SpeciesSet{T}, outputs::SpeciesSet{T}) where {T} = ChemicalHyperEdge{T}(inputs, outputs, one(Int))
ChemicalHyperEdge(inputs::AbstractVector{T}, outputs::AbstractVector{T}, rate) where {T} = ChemicalHyperEdge(SpeciesSet(inputs), SpeciesSet(outputs), rate)
ChemicalHyperEdge(inputs::AbstractVector{T}, outputs::AbstractVector{T}) where {T} = ChemicalHyperEdge(SpeciesSet(inputs), SpeciesSet(outputs))
ChemicalHyperEdge(inputs::AbstractVector{T}, outputs::SpeciesSet{T}) where {T} = ChemicalHyperEdge(SpeciesSet(inputs), outputs)
ChemicalHyperEdge(inputs::SpeciesSet{T}, outputs::AbstractVector{T}) where {T} = ChemicalHyperEdge(inputs, SpeciesSet(outputs))
ChemicalHyperEdge{T}() where {T} = ChemicalHyperEdge{T}(SpeciesSet{T}(), SpeciesSet{T}(), one(Int))
ChemicalHyperEdge() = ChemicalHyperEdge{Any}()
(::Type{ChemicalHyperEdge{T}})(inputs, outputs) where {T} = ChemicalHyperEdge(inputs, outputs)
inputs(e::ChemicalHyperEdge) = src(e)
outputs(e::ChemicalHyperEdge) = tgt(e)
rate(e::ChemicalHyperEdge) = weight(e)
src_stoich(e::ChemicalHyperEdge) = src_multiplicities(e)
tgt_stoich(e::ChemicalHyperEdge) = tgt_multiplicities(e)
inputs_stoich(e::ChemicalHyperEdge) = src_stoich(e)
outputs_stoich(e::ChemicalHyperEdge) = tgt_stoich(e)
Base.:(==)(e1::ChemicalHyperEdge, e2::ChemicalHyperEdge) = isequal(inputs(e1), inputs(e2)) && isequal(outputs(e1), outputs(e2)) && isequal(rate(e1), rate(e2))
Base.eltype(::ChemicalHyperEdge{T}) where {T} = T

"""
    ChemicalHyperGraph{T}

A type that represents oriented, weighted hypergraphs, in the context of chemical reactions.

# Fields
- `V`: a set of vertices
- `E`: a set of chemical hyperedges
"""
mutable struct ChemicalHyperGraph{T} <: AbstractHyperGraph
    V::AbstractVector{T}
    E::AbstractVector{ChemicalHyperEdge{T}}
end
ChemicalHyperGraph(vs::AbstractVector{T}, e::ChemicalHyperEdge{T}) where {T} = ChemicalHyperGraph{T}(vs, [e])
ChemicalHyperGraph(e::ChemicalHyperEdge{T}) where {T} = ChemicalHyperGraph{T}(union(vertices(e)), [e])
ChemicalHyperGraph(es::AbstractVector{ChemicalHyperEdge{T}}) where {T} = ChemicalHyperGraph(union(vertices(es)...), es)
(::Type{ChemicalHyperGraph{T}})(e::ChemicalHyperEdge{T}) where {T} = ChemicalHyperGraph(e)
(::Type{ChemicalHyperGraph{T}})(es::AbstractVector{ChemicalHyperEdge{T}}) where {T} = ChemicalHyperGraph(es)
ChemicalHyperGraph{T}() where {T} = ChemicalHyperGraph{T}([], [])
ChemicalHyperGraph() = ChemicalHyperGraph{Any}([], [])
Base.:(==)(x1::ChemicalHyperGraph, x2::ChemicalHyperGraph) = isequal(vertices(x1), vertices(x2)) && isequal(hyperedges(x1), hyperedges(x2))
Base.eltype(::ChemicalHyperGraph{T}) where {T} = T

## traits ##

# build unions of abstract types for easier dispatch
const AbstractStructs = Union{AbstractHyperGraph, AbstractHyperEdge}

# fundamental traits
@traitdef IsOriented{X <: AbstractStructs}
@traitdef IsWeighted{X <: AbstractStructs}

# functions and default behaviour: orientation
@traitimpl IsOriented{T} <- isoriented(T)
isoriented(::Type{T}) where {T} = false
isoriented(X::T) where {T} = isoriented(T)

# functions and default behaviour: weight
@traitimpl IsWeighted{T} <- isweighted(T)
isweighted(::Type{T}) where {T} = false
isweighted(X::T) where {T} = isweighted(T)

# build unions of types for easier dispatch
const OrientedStructs = Union{ChemicalHyperGraph, ChemicalHyperEdge}
const WeightedStructs = Union{ChemicalHyperGraph, ChemicalHyperEdge}

# implementing traits on types
@traitimpl IsOriented{OrientedStructs}
isoriented(::Type{T}) where {T<:OrientedStructs} = true
@traitimpl IsWeighted{WeightedStructs}
isweighted(::Type{T}) where {T<:WeightedStructs} = true

## core functions ##

# objects and multiplicities of incidence structures
objects(i::T) where {T<:AbstractIncidenceSet} = i.objs
multiplicities(i::T) where {T<:AbstractIncidenceSet} = i.mults

# vertices of hyperedges, of hypergraphs
@traitfn vertices(e::T::(!IsOriented)) where {T<:AbstractHyperEdge} = e.V
@traitfn vertices(e::T::IsOriented) where {T<:AbstractHyperEdge} = vcat(src(e), tgt(e))
vertices(es::AbstractVector{T}) where {T<:AbstractHyperEdge} = vertices.(es)
vertices(x::T) where {T<:AbstractHyperGraph} = x.V

# hyperedges of hypergraphs
hyperedges(x::T) where {T<:AbstractHyperGraph} = x.E

# source and target of oriented hyperedges
@traitfn src(e::T::IsOriented) where {T<:AbstractHyperEdge} = objects(e.src)
@traitfn tgt(e::T::IsOriented) where {T<:AbstractHyperEdge} = objects(e.tgt)
src(es::AbstractVector{T}) where {T<:AbstractHyperEdge} = [src(e) for e in es]
tgt(es::AbstractVector{T}) where {T<:AbstractHyperEdge} = [tgt(e) for e in es]

# multiplicities of oriented hyperedges
@traitfn src_multiplicities(e::T::IsOriented) where {T<:AbstractHyperEdge} = multiplicities(e.src)
@traitfn tgt_multiplicities(e::T::IsOriented) where {T<:AbstractHyperEdge} = multiplicities(e.tgt)

# weight of weighted hyperedges
@traitfn weight(e::T::IsWeighted) where {T<:AbstractHyperEdge } = e.weight
weights(es::AbstractVector{T}) where {T<:AbstractHyperEdge} = weight.(es)
weights(x::T) where {T<:AbstractHyperGraph} = weights(hyperedges(x))


text:
using ITensors
struct KitaevFermionsModel <: AbstractElectronHamiltonian
    t::Real
	Δ::Real
	μ::Real
end

function Hamiltonian(model::KitaevFermionsModel)
    ampo = AutoMPO()

    ampo += -model.t ,"Cdagup",1,"Cup",2
    ampo += -model.t ,"Cdagup",2,"Cup",1
    ampo += -model.t ,"Cdagdn",1,"Cdn",2
    ampo += -model.t ,"Cdagdn",2,"Cdn",1

    ampo += model.Δ,"Cdagup",1,"Cdagdn",2
    ampo += conj(model.Δ),"Cdn",2,"Cup",1
    ampo += model.Δ,"Cdagup",2,"Cdagdn",1
    ampo += conj(model.Δ),"Cdn",1,"Cup",2

    ampo += 0.5*model.μ,"Cdagup",1,"Cup",1
    ampo += 0.5*model.μ,"Cdagup",2,"Cup",2
    ampo += 0.5*model.μ,"Cdagdn",1,"Cdn",1
    ampo += 0.5*model.μ,"Cdagdn",2,"Cdn",2

    return Hamiltonian(ampo,"Electron",2)
end

struct tJModel <: AbstractElectronHamiltonian
    t::Real
    J::Real
end

function Hamiltonian(model::tJModel)
    ampo = AutoMPO()

    ampo += -model.t ,"Cdagup",1,"Cup",2
    ampo += -model.t ,"Cdagup",2,"Cup",1
    ampo += -model.t ,"Cdagdn",1,"Cdn",2
    ampo += -model.t ,"Cdagdn",2,"Cdn",1

    ampo += 1.0*model.J ,"Sz",1,"Sz",2
    ampo += 0.5*model.J ,"S+",1,"S-",2
    ampo += 0.5*model.J ,"S-",1,"S+",2

    ampo += -0.25*model.J, "Ntot",1,"Ntot",2

    return Hamiltonian(ampo,"Electron",2)
end


text:
"Convolves and resamples the input spectrum with a fixed kernel"
function conv_spectra(m::FixedKernelInstrument, ν, spectrum)
    s = imfilter(spectrum, m.kernel)
    interp_cubic = CubicSplineInterpolation(ν, s)
    return interp_cubic(m.ν_out)
end;

"Convolves and resamples the input spectrum with a variable kernel (per spectral position)"
function conv_spectra(m::VariableKernelInstrument, ν, spectrum; stride=1)
    FT = eltype(m.ν_out)
    FT2 = eltype(spectrum)
    # Define grid where to perform convolution:

    # Padding at both sides required:
    start = argmin(abs.(ν.-m.ν_out[1]))-1
    stop  = argmin(abs.(ν.-m.ν_out[end]))+1
    # @show start, stop
    off = ceil(Int, size(m.kernel, 1) / 2)
    @assert off ≤ start "Start range of model grid too close to output grid (needs buffer) $(off), $(start)"
    @assert off ≤ length(ν)-stop "Stop range of model grid too close to output grid (needs buffer) $(off), $(length(ν)-stop)"
    ind = start:stride:stop

    # knots where convolution will be applied to
    knots = view(ν, ind)
    te = LinearInterpolation(m.ν_out, FT.(m.ind_out); extrapolation_bc=Interpolations.Flat())
    spec_out = zeros(FT2, length(knots));
    for i in eachindex(knots)
        # Simple first, nearest neighbor ILS
        ind_fraction = round(Int, te(knots[i]));
        kernel = view(m.kernel, :, ind_fraction)
        for j in eachindex(kernel)
            spec_out[i] += kernel[j] * spectrum[ind[i] + j]
        end
    end
    # Change this later to only perform conv around output grid!
    fin = LinearInterpolation(ν[ind], spec_out; extrapolation_bc=Interpolations.Flat())
    return fin(m.ν_out)
end;


text:
export CuFunction

mutable struct CuFunction
    ptr::Ptr{Cvoid}
    mod::CuModule # avoid CuModule gc-ed
end

function CuFunction(mod::CuModule, name::String)
    ref = Ref{Ptr{Cvoid}}()
    @apicall :cuModuleGetFunction (Ptr{Ptr{Cvoid}},Ptr{Cvoid},Cstring) ref mod name
    CuFunction(ref[], mod)
end

function CuFunction(ptx::String)
    mod = CuModule(ptx)
    fnames = String[]
    for line in split(ptx,'\n')
        m = match(r".visible .entry (.+)\(", line) # find function name
        m == nothing && continue
        push!(fnames, String(m[1]))
    end
    length(fnames) > 1 && throw("Multiple functions are found.")
    CuFunction(mod, fnames[1])
end

Base.unsafe_convert(::Type{Ptr{Cvoid}}, f::CuFunction) = f.ptr


text:
using DocStringExtensions
"""
Computes the fluxes of all components in the model per timestep (all HRUs combined + slow component)

$(SIGNATURES)

The function returns the fluxes leaving the model (evaporation, discharge). It also returns the amount of water stored in each model component.
"""
function allHRU(bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters,
                Slowstorage::Float64, slow_parameters::Slow_Paramters)
    #this function runs thy model for different HRUs
    #bare rock HRU
    bare_outflow::Outflows, bare_storage::Storages, Bare_precipitation::Float64, Bare_Storages::Float64 = hillslopeHRU(bare_input, bare_storage, bare_parameters)
    # forest HRU
    forest_outflow::Outflows, forest_storage::Storages, Forest_precipitation::Float64, Forest_Storages::Float64= hillslopeHRU(forest_input, forest_storage, forest_parameters)
    # Grassland HRU
    grass_outflow::Outflows, grass_storage::Storages, Grass_precipitation::Float64, Grass_Storages::Float64= hillslopeHRU(grass_input, grass_storage, grass_parameters)
    # riparian HRU
    rip_outflow::Outflows, rip_storage::Storages, Rip_precipitation::Float64, Rip_Storages::Float64= riparianHRU(rip_input, rip_storage, rip_parameters)
    # total flow into groundwater is the weighted sum of the HRUs according to areal extent (this was already done in hillslopeHRU)
    Total_GWflow = bare_outflow.GWflow + forest_outflow.GWflow  + grass_outflow.GWflow
    # Groundwater storage
    Riparian_Discharge::Float64, Slow_Discharge::Float64, Slowstorage_New::Float64 = slowstorage(Total_GWflow, Slowstorage, rip_input.Area_HRU, slow_parameters.Ks, slow_parameters.Ratio_Riparian)
    #return all storage values, all evaporation values, Fast_Discharge and Slow_Discharge
    # calculate total discharge of the timestep using weighted sum of each HRU
    Total_Discharge::Float64 = bare_outflow.Fast_Discharge  + forest_outflow.Fast_Discharge + grass_outflow.Fast_Discharge  + rip_outflow.Fast_Discharge  + Slow_Discharge
    Total_Soil_Evaporation::Float64 = bare_outflow.Soil_Evaporation + forest_outflow.Soil_Evaporation + grass_outflow.Soil_Evaporation  + rip_outflow.Soil_Evaporation
    Total_Interception_Evaporation::Float64 = bare_outflow.Interception_Evaporation  + forest_outflow.Interception_Evaporation + grass_outflow.Interception_Evaporation + rip_outflow.Interception_Evaporation
    @assert Riparian_Discharge >= 0
    @assert Total_Discharge >= 0
    @assert Total_Interception_Evaporation >= 0
    @assert Total_Soil_Evaporation >= 0
    @assert Slowstorage_New >= 0
    Total_Flows = Total_Discharge + Total_Soil_Evaporation + Total_Interception_Evaporation + Riparian_Discharge
    Total_Storages = Bare_Storages * bare_input.Area_HRU + Forest_Storages * forest_input.Area_HRU + Grass_Storages * grass_input.Area_HRU + Rip_Storages * rip_input.Area_HRU + Slowstorage_New - Slowstorage
    Precipitation = Bare_precipitation * bare_input.Area_HRU + Forest_precipitation * forest_input.Area_HRU + Grass_precipitation * grass_input.Area_HRU + Rip_precipitation * rip_input.Area_HRU
    #@assert -0.00000001 <= Precipitation + rip_input.Riparian_Discharge - (Total_Flows + Total_Storages) <= 0.00000001
    Waterbalance = Precipitation + rip_input.Riparian_Discharge - (Total_Flows + Total_Storages)
    return Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage_New::Float64, Waterbalance::Float64, Precipitation::Float64
end

"""
Computes the fluxes of all components in the model per timestep (all HRUs combined + slow component)

$(SIGNATURES)

The function returns the fluxes leaving the model (evaporation, discharge). It also returns the amount of water stored in each model component.
"""
function allHRU_future(bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters,
                Slowstorage::Float64, slow_parameters::Slow_Paramters)
    #this function runs thy model for different HRUs
    #bare rock HRU
    bare_outflow::Outflows, bare_storage::Storages, Bare_precipitation::Float64, Bare_Storages::Float64, Bare_Melt::Float64 = hillslopeHRU_future(bare_input, bare_storage, bare_parameters)
    # forest HRU
    forest_outflow::Outflows, forest_storage::Storages, Forest_precipitation::Float64, Forest_Storages::Float64, Forest_Melt::Float64 = hillslopeHRU_future(forest_input, forest_storage, forest_parameters)
    # Grassland HRU
    grass_outflow::Outflows, grass_storage::Storages, Grass_precipitation::Float64, Grass_Storages::Float64, Grass_Melt::Float64 = hillslopeHRU_future(grass_input, grass_storage, grass_parameters)
    # riparian HRU
    rip_outflow::Outflows, rip_storage::Storages, Rip_precipitation::Float64, Rip_Storages::Float64, Rip_Melt::Float64 = riparianHRU_future(rip_input, rip_storage, rip_parameters)
    # total flow into groundwater is the weighted sum of the HRUs according to areal extent (this was already done in hillslopeHRU)
    Total_GWflow = bare_outflow.GWflow + forest_outflow.GWflow  + grass_outflow.GWflow
    # Groundwater storage
    Riparian_Discharge::Float64, Slow_Discharge::Float64, Slowstorage_New::Float64 = slowstorage(Total_GWflow, Slowstorage, rip_input.Area_HRU, slow_parameters.Ks, slow_parameters.Ratio_Riparian)
    #return all storage values, all evaporation values, Fast_Discharge and Slow_Discharge
    # calculate total discharge of the timestep using weighted sum of each HRU
    Total_Discharge::Float64 = bare_outflow.Fast_Discharge  + forest_outflow.Fast_Discharge + grass_outflow.Fast_Discharge  + rip_outflow.Fast_Discharge  + Slow_Discharge
    Total_Soil_Evaporation::Float64 = bare_outflow.Soil_Evaporation + forest_outflow.Soil_Evaporation + grass_outflow.Soil_Evaporation  + rip_outflow.Soil_Evaporation
    Total_Interception_Evaporation::Float64 = bare_outflow.Interception_Evaporation  + forest_outflow.Interception_Evaporation + grass_outflow.Interception_Evaporation + rip_outflow.Interception_Evaporation
    @assert Riparian_Discharge >= 0
    @assert Total_Discharge >= 0
    @assert Total_Interception_Evaporation >= 0
    @assert Total_Soil_Evaporation >= 0
    @assert Slowstorage_New >= 0
    Total_Flows = Total_Discharge + Total_Soil_Evaporation + Total_Interception_Evaporation + Riparian_Discharge
    Total_Melt = Bare_Melt + Forest_Melt + Grass_Melt + Rip_Melt
    Total_Storages = Bare_Storages * bare_input.Area_HRU + Forest_Storages * forest_input.Area_HRU + Grass_Storages * grass_input.Area_HRU + Rip_Storages * rip_input.Area_HRU + Slowstorage_New - Slowstorage
    Precipitation = Bare_precipitation * bare_input.Area_HRU + Forest_precipitation * forest_input.Area_HRU + Grass_precipitation * grass_input.Area_HRU + Rip_precipitation * rip_input.Area_HRU
    #@assert -0.00000001 <= Precipitation + rip_input.Riparian_Discharge - (Total_Flows + Total_Storages) <= 0.00000001
    Waterbalance = Precipitation + rip_input.Riparian_Discharge - (Total_Flows + Total_Storages)
    return Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage_New::Float64, Waterbalance::Float64, Precipitation::Float64, Total_Melt::Float64
end

"""
Runs the semi-distributed bucket model for a given area with one precipitation and temperature input.

$(SIGNATURES)
"""
function runmodel_alloutput(Area, Evaporation_Mean::Array{Float64,1}, Precipitation::Array{Float64,2}, Temp::Array{Float64,2},
                bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters, slow_parameters::Slow_Paramters, Total_Elevationbands, Elevation_Percentage)
    # the function takes as input the parameters of each HRU, the inital storage values of each HRU, the inital value of the slow storage
    # KS, ratio riparian, all inputs

    # define the maximum time
    tmax::Int128 = length(Precipitation[:,1])

    # make arrays for each Model Component
    Int_Evaporation::Array{Float64,1} = zeros(tmax) #interception evaporation
    Soil_Evaporation::Array{Float64,1} = zeros(tmax) #soil evaporation
    Discharge::Array{Float64,1} = zeros(tmax)

    # store the initial storage values, Does not take into account GW storage!!
    # Assumption_ GW storage is 0 at start
    #TO DO: average values over area
    # as form now only works if storage input is 0 at start
    Initial_Storage_bare::Float64 = bare_storage.Fast + sum(bare_storage.Interception) + sum(bare_storage.Snow) + bare_storage.Soil
    Initial_Storage_forest::Float64 = forest_storage.Fast + sum(forest_storage.Interception) + sum(forest_storage.Snow) + forest_storage.Soil
    Initial_Storage_grass::Float64 = grass_storage.Fast + sum(grass_storage.Interception) + sum(grass_storage.Snow) + grass_storage.Soil
    Initial_Storage_rip::Float64 = rip_storage.Fast + sum(rip_storage.Interception) + sum(rip_storage.Snow) + rip_storage.Soil
    Initial_Storage::Float64 = Initial_Storage_bare + Initial_Storage_forest + Initial_Storage_grass + Initial_Storage_rip + Slowstorage

    #OPTIONAL: store all storage states
    Interceptionstorage::Array{Float64,2} = zeros(tmax, 4) #storage interception
    Snowstorage::Array{Float64,1} = zeros(tmax)
    Soilstorage::Array{Float64,2} = zeros(tmax, 4) #stroage unsaturated zone
    Faststorage::Array{Float64,2} = zeros(tmax, 4) #storage fast
    GWstorage::Array{Float64,1} = zeros(tmax) #storage GW
    WBtotal::Array{Float64,1} = zeros(tmax)
    Snow_Extend::Array{Float64,2} = zeros(tmax, Total_Elevationbands)
    Precipitation_Total::Array{Float64,1} = zeros(tmax)
    Snow_Elevations::Array{Float64,2} = zeros(tmax, Total_Elevationbands)
    Bare_Snow::Array{Float64,2} = zeros(tmax, bare_input.Nr_Elevationbands)

    for t in 1:tmax
        # at each timestep new temp, precipitation and Epot values have to be delivered
        # areas don't change
        # riparian discharge from former timestep has to be used
        # gives the current precipitation, evaporation and temperature
        Evaporation_Mean_Current = Evaporation_Mean[t]
        Precipitation_Current = Precipitation[t, :]
        Temperature_Current = Temp[t, :]

        bare_input::HRU_Input = input_timestep(bare_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        forest_input::HRU_Input = input_timestep(forest_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        grass_input::HRU_Input = input_timestep(grass_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        rip_input::HRU_Input = input_timestep(rip_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)

        Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64, WB::Float64, Total_Prec::Float64 = allHRU(bare_input, forest_input, grass_input, rip_input,
                                                                                bare_storage, forest_storage, grass_storage, rip_storage,
                                                                                bare_parameters, forest_parameters, grass_parameters, rip_parameters,
                                                                                Slowstorage, slow_parameters)
        # give new riparian discharge as input for next timestep
        rip_input.Riparian_Discharge = Riparian_Discharge
        # save the fluxes of the current timestep
        Discharge[t]::Float64 = Total_Discharge/1000 * Area / (3600 * 24) # mm convert it to meter and than * area / seconds in one day
        Int_Evaporation[t]::Float64 = Total_Interception_Evaporation
        Soil_Evaporation[t]::Float64 = Total_Soil_Evaporation
        #OPTIONAL: store all storage states at each timestep
        #get the total value stored as mean value of elevations and areal extent of HRU
        Bare_Interceptionstorage::Float64, Bare_Snowstorage::Float64 = Storage_Total(bare_storage, bare_input)
        Forest_Interceptionstorage::Float64, Forest_Snowstorage::Float64 = Storage_Total(forest_storage, forest_input)
        Grass_Interceptionstorage::Float64, Grass_Snowstorage::Float64 = Storage_Total(grass_storage, grass_input)
        Rip_Interceptionstorage::Float64, Rip_Snowstorage::Float64 = Storage_Total(rip_storage, rip_input)

        # OPTIONAL: store the total amount of snow in each elevation (considering whole catchment)
        Snow_Elevations[t, :]::Array{Float64,1}, Snow_Extend[t,:]::Array{Float64,1} = snowperelevation(bare_input, forest_input, grass_input, rip_input, bare_storage, forest_storage, grass_storage, rip_storage, Total_Elevationbands)
        #Bare_Snow[t,:]::Array{Float64,1} = bare_storage.Snow
        Snow_Extend[t,:] = Snow_Extend[t,:] ./ Elevation_Percentage
        Snow_Elevations[t,:] = Snow_Elevations[t,:] ./Elevation_Percentage

        Interceptionstorage[t, :]::Array{Float64,1} = [Bare_Interceptionstorage, Forest_Interceptionstorage, Grass_Interceptionstorage, Rip_Interceptionstorage]
        Snowstorage[t] = Bare_Snowstorage * bare_input.Area_HRU + Forest_Snowstorage * forest_input.Area_HRU + Grass_Snowstorage * grass_input.Area_HRU + Rip_Snowstorage * rip_input.Area_HRU
        Soilstorage[t, :]::Array{Float64,1} = [bare_storage.Soil, forest_storage.Soil, grass_storage.Soil, rip_storage.Soil]
        Faststorage[t, :]::Array{Float64,1} = [bare_storage.Fast, forest_storage.Fast, grass_storage.Fast, rip_storage.Fast]
        GWstorage[t]::Float64 = Slowstorage
        WBtotal[t]::Float64 = WB
        Precipitation_Total[t]::Float64 = Total_Prec
    end

    # Check Water Balance
    # calculate the water balance at each timestep and sum it at the end for getting waterbalance over all timesteps
    Waterbalance = sum(WBtotal)::Float64
    @assert Waterbalance <= 10^(-10)
    return Discharge::Array{Float64,1}, Snow_Extend::Array{Float64,2}, GWstorage::Array{Float64,1}, Snowstorage::Array{Float64,1}, Snow_Elevations, Soilstorage::Array{Float64,2}, Faststorage::Array{Float64,2} #, Interceptionstorage::Array{Float64,2},
end

function runmodel_alloutput_glacier(Area, Evaporation_Mean::Array{Float64,1}, Glacier::Array{Float64,2}, Precipitation::Array{Float64,2}, Temp::Array{Float64,2},
                bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters, slow_parameters::Slow_Paramters, Total_Elevationbands, Elevation_Percentage)
    # the function takes as input the parameters of each HRU, the inital storage values of each HRU, the inital value of the slow storage
    # KS, ratio riparian, all inputs

    # define the maximum time
    tmax::Int128 = length(Precipitation[:,1])

    # make arrays for each Model Component
    Int_Evaporation::Array{Float64,1} = zeros(tmax) #interception evaporation
    Soil_Evaporation::Array{Float64,1} = zeros(tmax) #soil evaporation
    Discharge::Array{Float64,1} = zeros(tmax)

    # store the initial storage values, Does not take into account GW storage!!
    # Assumption_ GW storage is 0 at start
    #TO DO: average values over area
    # as form now only works if storage input is 0 at start
    Initial_Storage_bare::Float64 = bare_storage.Fast + sum(bare_storage.Interception) + sum(bare_storage.Snow) + bare_storage.Soil
    Initial_Storage_forest::Float64 = forest_storage.Fast + sum(forest_storage.Interception) + sum(forest_storage.Snow) + forest_storage.Soil
    Initial_Storage_grass::Float64 = grass_storage.Fast + sum(grass_storage.Interception) + sum(grass_storage.Snow) + grass_storage.Soil
    Initial_Storage_rip::Float64 = rip_storage.Fast + sum(rip_storage.Interception) + sum(rip_storage.Snow) + rip_storage.Soil
    Initial_Storage::Float64 = Initial_Storage_bare + Initial_Storage_forest + Initial_Storage_grass + Initial_Storage_rip + Slowstorage

    #OPTIONAL: store all storage states
    Interceptionstorage::Array{Float64,2} = zeros(tmax, 4) #storage interception
    Snowstorage::Array{Float64,1} = zeros(tmax)
    Soilstorage::Array{Float64,2} = zeros(tmax, 4) #stroage unsaturated zone
    Faststorage::Array{Float64,2} = zeros(tmax, 4) #storage fast
    GWstorage::Array{Float64,1} = zeros(tmax) #storage GW
    WBtotal::Array{Float64,1} = zeros(tmax)
    Snow_Extend::Array{Float64,2} = zeros(tmax, Total_Elevationbands)
    Precipitation_Total::Array{Float64,1} = zeros(tmax)
    Snow_Elevations::Array{Float64,2} = zeros(tmax, Total_Elevationbands)
    Bare_Snow::Array{Float64,2} = zeros(tmax, bare_input.Nr_Elevationbands)

    for t in 1:tmax
        # at each timestep new temp, precipitation and Epot values have to be delivered
        # areas don't change
        # riparian discharge from former timestep has to be used
        # gives the current precipitation, evaporation and temperature
        Evaporation_Mean_Current = Evaporation_Mean[t]
        Precipitation_Current = Precipitation[t, :]
        Temperature_Current = Temp[t, :]
        Glacier_Current = Glacier[:,t]

        bare_input::HRU_Input = input_timestep_glacier(bare_input, Evaporation_Mean_Current, Glacier_Current, Precipitation_Current, Temperature_Current)
        forest_input::HRU_Input = input_timestep(forest_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        grass_input::HRU_Input = input_timestep(grass_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        rip_input::HRU_Input = input_timestep(rip_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)

        Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64, WB::Float64, Total_Prec::Float64 = allHRU(bare_input, forest_input, grass_input, rip_input,
                                                                                bare_storage, forest_storage, grass_storage, rip_storage,
                                                                                bare_parameters, forest_parameters, grass_parameters, rip_parameters,
                                                                                Slowstorage, slow_parameters)
        # give new riparian discharge as input for next timestep
        rip_input.Riparian_Discharge = Riparian_Discharge
        # save the fluxes of the current timestep
        Discharge[t]::Float64 = Total_Discharge/1000 * Area / (3600 * 24) # mm convert it to meter and than * area / seconds in one day
        Int_Evaporation[t]::Float64 = Total_Interception_Evaporation
        Soil_Evaporation[t]::Float64 = Total_Soil_Evaporation
        #OPTIONAL: store all storage states at each timestep
        #get the total value stored as mean value of elevations and areal extent of HRU
        Bare_Interceptionstorage::Float64, Bare_Snowstorage::Float64 = Storage_Total(bare_storage, bare_input)
        Forest_Interceptionstorage::Float64, Forest_Snowstorage::Float64 = Storage_Total(forest_storage, forest_input)
        Grass_Interceptionstorage::Float64, Grass_Snowstorage::Float64 = Storage_Total(grass_storage, grass_input)
        Rip_Interceptionstorage::Float64, Rip_Snowstorage::Float64 = Storage_Total(rip_storage, rip_input)

        # OPTIONAL: store the total amount of snow in each elevation (considering whole catchment)
        Snow_Elevations[t, :]::Array{Float64,1}, Snow_Extend[t,:]::Array{Float64,1} = snowperelevation(bare_input, forest_input, grass_input, rip_input, bare_storage, forest_storage, grass_storage, rip_storage, Total_Elevationbands)
        #Bare_Snow[t,:]::Array{Float64,1} = bare_storage.Snow
        Snow_Extend[t,:] = Snow_Extend[t,:] ./ Elevation_Percentage
        Snow_Elevations[t,:] = Snow_Elevations[t,:] ./Elevation_Percentage

        Interceptionstorage[t, :]::Array{Float64,1} = [Bare_Interceptionstorage, Forest_Interceptionstorage, Grass_Interceptionstorage, Rip_Interceptionstorage]
        Snowstorage[t] = Bare_Snowstorage * bare_input.Area_HRU + Forest_Snowstorage * forest_input.Area_HRU + Grass_Snowstorage * grass_input.Area_HRU + Rip_Snowstorage * rip_input.Area_HRU
        Soilstorage[t, :]::Array{Float64,1} = [bare_storage.Soil, forest_storage.Soil, grass_storage.Soil, rip_storage.Soil]
        Faststorage[t, :]::Array{Float64,1} = [bare_storage.Fast, forest_storage.Fast, grass_storage.Fast, rip_storage.Fast]
        GWstorage[t]::Float64 = Slowstorage
        WBtotal[t]::Float64 = WB
        Precipitation_Total[t]::Float64 = Total_Prec
    end

    # Check Water Balance
    # calculate the water balance at each timestep and sum it at the end for getting waterbalance over all timesteps
    Waterbalance = sum(WBtotal)::Float64
    @assert Waterbalance <= 10^(-10)
    return Discharge::Array{Float64,1}, Snow_Extend::Array{Float64,2}, GWstorage::Array{Float64,1}, Snowstorage::Array{Float64,1}, Snow_Elevations, Soilstorage::Array{Float64,2}, Faststorage::Array{Float64,2} #, Interceptionstorage::Array{Float64,2},
end

function Storage_Total(Storage::Storages, Input::HRU_Input)
    #print([bare_storage.Interception[1] * bare_input.Area_Elevations[1], forest_storage.Interception[1] * forest_input.Area_Elevations[1] , grass_storage.Interception[1] * grass_input.Area_Elevations[1], rip_storage.Interception[1] * rip_input.Area_Elevations[1]])
    # this function gives the total storage stored in reservoirs that are elevation distributed
    Total_Interception_Storage = 0
    Total_Snow_Storage = 0
    for i in 1 : Input.Nr_Elevationbands
        Interception_Storage = Storage.Interception[i] * Input.Area_Elevations[i]
        Snow_Storage = Storage.Snow[i] * Input.Area_Elevations[i]
        Former_Total_Interception_Storage = Total_Interception_Storage
        Former_Total_Snow_Storage = Total_Snow_Storage

        Total_Interception_Storage += Interception_Storage
        Total_Snow_Storage += Snow_Storage

        @assert Total_Interception_Storage >= Interception_Storage
        @assert Total_Snow_Storage >= Snow_Storage
        @assert Total_Interception_Storage >= Former_Total_Interception_Storage
        @assert Total_Snow_Storage >= Former_Total_Snow_Storage
    end
    return Total_Interception_Storage::Float64, Total_Snow_Storage::Float64
end

function snowperelevation(Bare::HRU_Input, Forest::HRU_Input, Grass::HRU_Input, Rip::HRU_Input, Bare_Storage::Storages, Forest_Storage::Storages, Grass_Storage::Storages, Rip_Storage::Storages, Total_Elevationbands)
    Snowstorage = zeros(Total_Elevationbands)
    Snow_Cover = zeros(Total_Elevationbands)
    bare_count = 1
    forest_count = 1
    grass_count = 1
    rip_count = 1
    for i in 1: Total_Elevationbands
        if bare_count <= length(Bare.Elevation_Count) && Bare.Elevation_Count[bare_count] == i
            Snowstorage[i]+= Bare_Storage.Snow[bare_count] * Bare.Area_Elevations[bare_count] * Bare.Area_HRU
            Snow_Cover[i]+= Bare_Storage.Snow_Cover[bare_count] * Bare.Area_Elevations[bare_count] * Bare.Area_HRU
            bare_count+= 1
        end
        if forest_count <= length(Forest.Elevation_Count) && Forest.Elevation_Count[forest_count] == i
            Snowstorage[i]+= Forest_Storage.Snow[forest_count] * Forest.Area_Elevations[forest_count] * Forest.Area_HRU
            Snow_Cover[i]+= Forest_Storage.Snow_Cover[forest_count] * Forest.Area_Elevations[forest_count] * Forest.Area_HRU
            forest_count+= 1
        end
        if grass_count <= length(Grass.Elevation_Count) && Grass.Elevation_Count[grass_count] == i
            Snowstorage[i]+= Grass_Storage.Snow[grass_count] * Grass.Area_Elevations[grass_count] * Grass.Area_HRU
            Snow_Cover[i]+= Grass_Storage.Snow_Cover[grass_count] * Grass.Area_Elevations[grass_count] * Grass.Area_HRU
            grass_count += 1
        end
        if rip_count <= length(Rip.Elevation_Count) && Rip.Elevation_Count[rip_count] == i
            Snowstorage[i]+= Rip_Storage.Snow[rip_count] * Rip.Area_Elevations[rip_count] * Rip.Area_HRU
            Snow_Cover[i]+= Rip_Storage.Snow_Cover[rip_count] * Rip.Area_Elevations[rip_count] * Rip.Area_HRU
            rip_count += 1
        end
    end
    return Snowstorage::Array{Float64,1}, Snow_Cover::Array{Float64,1}
end

function input_timestep(Input::HRU_Input, Evaporation_Mean::Float64, Precipitation::Array{Float64,1}, Temperature::Array{Float64,1})
    #Input.Potential_Evaporation::Array{Float64,1} = Evaporation
    Input.Potential_Evaporation_Mean::Float64 = Evaporation_Mean
    # get the precipitation data of the necessary elevations
    Precipitation_HRU = Float64[]
    Temperature_HRU = Float64[]
    for i in Input.Elevation_Count
        push!(Precipitation_HRU, Precipitation[i])
        push!(Temperature_HRU, Temperature[i])
    end
    Input.Precipitation::Array{Float64,1} = Precipitation_HRU
    Input.Temp_Elevation::Array{Float64,1} = Temperature_HRU
    return Input::HRU_Input
end

function input_timestep_glacier(Input::HRU_Input, Evaporation_Mean::Float64, Glacier_Area::Array{Float64,1}, Precipitation::Array{Float64,1}, Temperature::Array{Float64,1})
    #Input.Potential_Evaporation::Array{Float64,1} = Evaporation
    Input.Potential_Evaporation_Mean::Float64 = Evaporation_Mean
    # get the precipitation data of the necessary elevations
    Precipitation_HRU = Float64[]
    Temperature_HRU = Float64[]
    Glacier_HRU = Float64[]
    for i in Input.Elevation_Count
        push!(Precipitation_HRU, Precipitation[i])
        push!(Temperature_HRU, Temperature[i])
        push!(Glacier_HRU, Glacier_Area[i])
    end
    Input.Precipitation::Array{Float64,1} = Precipitation_HRU
    Input.Temp_Elevation::Array{Float64,1} = Temperature_HRU
    Input.Area_Glacier::Array{Float64,1} = Glacier_HRU
    return Input::HRU_Input
end

max_stars_repo_name: sarah-hanus/hbv-mountain

max_issues_repo_name: sarah-hanus/hbv-mountain

max_forks_repo_name: sarah-hanus/hbv-mountain


text:
using Metis
using SparseArrays
using LightGraphs: add_edge!
using SimpleWeightedGraphs
using MetaGraphs
using ..Utils

export learn_vtree

function learn_vtree(data::DataFrame; α=0.0, alg=:bottomup)
    if alg==:topdown
        PlainVtree(num_features(data), :topdown; f=metis_top_down(data;α))
    elseif alg==:bottomup
        PlainVtree(num_features(data), :bottomup; f=blossom_bottom_up(data;α))
    elseif alg==:clt
        clt = learn_chow_liu_tree(data)
        learn_vtree_from_clt(clt, vtree_mode="balanced")
    else
        error("Vtree learner $(alg) not supported.")
    end
end

#############
# Metis top down method
#############

# Add edge weights to Metis.jl
using Metis: idx_t, ishermitian

struct WeightedGraph
    nvtxs::idx_t
    xadj::Vector{idx_t}
    adjncy::Vector{idx_t}
    adjwgt::Vector{idx_t} # edge weights
    WeightedGraph(nvtxs, xadj, adjncy, adjwgt) = new(nvtxs, xadj, adjncy, adjwgt)
end

function my_graph(G::SparseMatrixCSC; check_hermitian=true)
    if check_hermitian
        ishermitian(G) || throw(ArgumentError("matrix must be Hermitian"))
    end
    N = size(G, 1)
    xadj = Vector{idx_t}(undef, N+1)
    xadj[1] = 1
    adjncy = Vector{idx_t}(undef, nnz(G))
    adjncy_i = 0
    adjwgt = Vector{idx_t}(undef, nnz(G))
    @inbounds for j in 1:N
        n_rows = 0
        for k in G.colptr[j] : (G.colptr[j+1] - 1)
            i = G.rowval[k]
            if i != j # don't include diagonal elements
                n_rows += 1
                adjncy_i += 1
                adjncy[adjncy_i] = i
                adjwgt[adjncy_i] = G[i, j]
            end
        end
        xadj[j+1] = xadj[j] + n_rows
    end
    resize!(adjncy, adjncy_i)
    resize!(adjwgt, adjncy_i)
    return WeightedGraph(idx_t(N), xadj, adjncy, adjwgt)
end

function my_partition(G::WeightedGraph, nparts::Integer)
    part = Vector{Metis.idx_t}(undef, G.nvtxs)
    edgecut = fill(idx_t(0), 1)
    # if alg === :RECURSIVE
        Metis.METIS_PartGraphRecursive(G.nvtxs, idx_t(1), G.xadj, G.adjncy, C_NULL, C_NULL, G.adjwgt,
                                 idx_t(nparts), C_NULL, C_NULL, Metis.options, edgecut, part)
    # elseif alg === :KWAY
    #     Metis.METIS_PartGraphKway(G.nvtxs, idx_t(1), G.xadj, G.adjncy, C_NULL, C_NULL, G.adjwgt,
    #                         idx_t(nparts), C_NULL, C_NULL, Metis.options, edgecut, part)
    # else
    #     throw(ArgumentError("unknown algorithm $(repr(alg))"))
    # end
    return part
end

my_partition(G, nparts) = my_partition(my_graph(G), nparts)

"Metis top down method"
function metis_top_down(data::DataFrame;α)
    δINT = 999999
    MIN_INT = 1
    MAX_INT = δINT + MIN_INT

    weight=ones(Float64, num_examples(data))
    (_, mi) = mutual_information(data, weight; α=α)
    vars = Var.(collect(1:num_features(data)))
    info = to_long_mi(mi, MIN_INT, MAX_INT)

    function f(leafs::Vector{PlainVtreeLeafNode})::Tuple{Vector{PlainVtreeLeafNode}, Vector{PlainVtreeLeafNode}}
        var2leaf = Dict([(variable(x),x) for x in leafs])
        vertices = sort(variable.(leafs))
        sub_context = info[vertices, vertices]
        len = length(vertices)
        for i in 1 : len
            sub_context[i, i] = 0
        end
        g = convert(SparseMatrixCSC, sub_context)
        partition = my_partition(my_graph(g), 2)

        subsets = (Vector{PlainVtreeLeafNode}(), Vector{PlainVtreeLeafNode}())
        for (index, p) in enumerate(partition)
            push!(subsets[p], var2leaf[vertices[index]])
        end

        return subsets
    end
    return f
end


text:
# AUTO GENERATED FILE - DO NOT EDIT

export daq_precisioninput

"""
    daq_precisioninput(;kwargs...)

A PrecisionInput component.
A numeric input component that converts
an input value to the desired precision.
Keyword arguments:
- `id` (String; optional): The ID used to identify this compnent in Dash callbacks
- `className` (String; optional): Class to apply to the root component element.
- `disabled` (Bool; optional): If true, numeric input cannot be changed.
- `label` (optional): Description to be displayed alongside the scientific notation. To control styling,
pass an object with label and style properties.. label has the following type: String | lists containing elements 'style', 'label'.
Those elements have the following types:
  - `style` (Dict; optional)
  - `label` (String; optional)
- `labelPosition` (a value equal to: 'top', 'bottom'; optional): Where the numeric input label is positioned.
- `max` (Real; optional): The maximum value of the numeric input
- `min` (Real; optional): The minimum value of the numeric input
- `persisted_props` (Array of a value equal to: 'value's; optional): Properties whose user interactions will persist after refreshing the
component or the page. Since only `value` is allowed this prop can
normally be ignored.
- `persistence` (Bool | String | Real; optional): Used to allow user interactions in this component to be persisted when
the component - or the page - is refreshed. If `persisted` is truthy and
hasn't changed from its previous value, a `value` that the user has
changed while using the app will keep that change, as long as
the new `value` also matches what was given originally.
Used in conjunction with `persistence_type`.
- `persistence_type` (a value equal to: 'local', 'session', 'memory'; optional): Where persisted user changes will be stored:
memory: only kept in memory, reset on page refresh.
local: window.localStorage, data is kept after the browser quit.
session: window.sessionStorage, data is cleared once the browser quit.
- `precision` (Real; optional): Number of significant figures
- `size` (Real; optional): The size (length) of the numeric input in pixels
- `style` (Dict; optional): Style to apply to the root component element.
- `theme` (Dict; optional): Theme configuration to be set by a ThemeProvider
- `value` (Real; optional): The value of numeric input
"""
function daq_precisioninput(; kwargs...)
        available_props = Symbol[:id, :className, :disabled, :label, :labelPosition, :max, :min, :persisted_props, :persistence, :persistence_type, :precision, :size, :style, :theme, :value]
        wild_props = Symbol[]
        return Component("daq_precisioninput", "PrecisionInput", "dash_daq", available_props, wild_props; kwargs...)
end


text:
"""
    write_vcf(outfile, x1, x2, [chr], [pos], [ref], [alt], [sampleID], [SNPids])

Writes haplotypes in `x1` and `x2` into VCF file, tracking phase information.

# Inputs
- `outfile`: Output file name (ends in `.vcf` or `.vcf.gz`)
- `x1`: `n × p` matrix of the 1st haplotype for each sample. Each row is a haplotype,
    and each entry is 0 or 1
- `x2`: `n × p` matrix of the 2nd haplotype for each sample. `x = x1 + x2`.
    Each entry is 0 or 1

# Optional Inputs
- `chr`: Vector of String. `chr[i]` is chromosome number of SNP `i`.
- `pos`: Vector of Int. `pos[i]` is position of SNP `i`.
- `ref`: Vector of String. `ref[i]` is reference allele of SNP `i`.
- `alt`: Vector of String. `alt[i]` is alternate allele of SNP `i`.
- `sampleID`: Vector of String. `sampleID[i]` is reference allele of SNP `i`.
- `SNPids`: Vector of String. `SNPids[i]` is reference allele of SNP `i`.

# Reference
For multithreaded write, see
https://github.com/OpenMendel/MendelImpute.jl/blob/master/src/impute.jl#L23
"""
function write_vcf(
    outfile::AbstractString,
    x1::AbstractMatrix,
    x2::AbstractMatrix;
    chr::AbstractVector{String} = ["1"     for i in 1:size(x1, 2)],
    pos::AbstractVector{Int}    = [100i    for i in 1:size(x1, 2)],
    ref::AbstractVector{String} = ["A"     for i in 1:size(x1, 2)],
    alt::AbstractVector{String} = ["T"     for i in 1:size(x1, 2)],
    sampleID::AbstractVector{String} = [string(i) for i in 1:size(x1, 1)],
    SNPids::AbstractVector{String} = ["snp$i" for i in 1:size(x1, 2)],
    )
    n, p = size(x1)
    (n, p) == size(x2) || error("x1 and x2 have different dimensions!")
    io = openvcf(outfile, "w")
    pb = PipeBuffer()

    # minimum meta info
    print(pb, "##fileformat=VCFv4.2\n")
    print(pb, "##source=VCFTools.jl\n")
    print(pb, "##FORMAT=<ID=GT,Number=1,Type=String,Description=\"Genotype\">\n")

    # header line (includes sample ID)
    print(pb, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
    for id in sampleID
        print(pb, "\t", id)
    end
    print(pb, "\n")
    bytesavailable(pb) > 1048576 && write(io, take!(pb))

    # loop over genotypes
    pmeter = Progress(p, 1, "Writing VCF...")
    @inbounds for i in 1:p
        # write meta info (chrom/pos/snpid/ref/alt/imputation-quality)
        print(pb, chr[i], "\t", string(pos[i]), "\t", SNPids[i], "\t",
            ref[i], "\t", alt[i], "\t.\tPASS\t.\tGT")
        # print ith record
        write_snp!(pb, x1, x2, i)
        bytesavailable(pb) > 1048576 && write(io, take!(pb))
        next!(pmeter)
    end
    write(io, take!(pb))
    close(io); close(pb) # close io and buffer
end

"""
Helper function for saving the `i`th record (SNP), tracking phase information.
Here `X = X1 + X2`.
"""
function write_snp!(pb::IOBuffer, X1::AbstractMatrix, X2::AbstractMatrix, i::Int)
    x1 = @view(X1[:, i])
    x2 = @view(X2[:, i])
    n = length(x1)
    @assert n == length(x2)
    @inbounds for j in 1:n
        if x1[j] == x2[j] == 0
            print(pb, "\t0|0")
        elseif x1[j] == 0 && x2[j] == 1
            print(pb, "\t0|1")
        elseif x1[j] == 1 && x2[j] == 0
            print(pb, "\t1|0")
        elseif x1[j] == 1 && x2[j] == 1
            print(pb, "\t1|1")
        else
            error("phased genotypes can only be 0|0, 0|1, 1|0 or 1|1 but
                got $(x1[j])|$(x2[j])")
        end
    end
    print(pb, "\n")
    return nothing
end

"""
    write_vcf(outfile, x, [chr], [pos], [ref], [alt], [sampleID], [SNPids])

Writes genotypes in `x` into VCF file, not tracking phase information.

# Inputs
- `outfile`: Output file name (ends in `.vcf` or `.vcf.gz`)
- `x`: `n × p` genotype matrix. Each row is a sample, and each entry is 0, 1, or 2.

# Optional Inputs
- `chr`: Vector of String. `chr[i]` is chromosome number of SNP `i`.
- `pos`: Vector of Int. `pos[i]` is position of SNP `i`.
- `ref`: Vector of String. `ref[i]` is reference allele of SNP `i`.
- `alt`: Vector of String. `alt[i]` is alternate allele of SNP `i`.
- `sampleID`: Vector of String. `sampleID[i]` is reference allele of SNP `i`.
- `SNPids`: Vector of String. `SNPids[i]` is reference allele of SNP `i`.

# Reference
For multithreaded write, see
https://github.com/OpenMendel/MendelImpute.jl/blob/master/src/impute.jl#L23
"""
function write_vcf(
    outfile::AbstractString,
    x::AbstractMatrix;
    chr::AbstractVector{String} = ["1"     for i in 1:size(x, 2)],
    pos::AbstractVector{Int}    = [100i    for i in 1:size(x, 2)],
    ref::AbstractVector{String} = ["A"     for i in 1:size(x, 2)],
    alt::AbstractVector{String} = ["T"     for i in 1:size(x, 2)],
    sampleID::AbstractVector{String} = [string(i) for i in 1:size(x, 1)],
    SNPids::AbstractVector{String} = ["snp$i" for i in 1:size(x, 2)],
    )
    n, p = size(x)
    io = openvcf(outfile, "w")
    pb = PipeBuffer()

    # minimum meta info
    print(pb, "##fileformat=VCFv4.2\n")
    print(pb, "##source=VCFTools.jl\n")
    print(pb, "##FORMAT=<ID=GT,Number=1,Type=String,Description=\"Genotype\">\n")

    # header line (includes sample ID)
    print(pb, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
    for id in sampleID
        print(pb, "\t", id)
    end
    print(pb, "\n")
    bytesavailable(pb) > 1048576 && write(io, take!(pb))

    # loop over genotypes
    pmeter = Progress(p, 1, "Writing VCF...")
    @inbounds for i in 1:p
        # write meta info (chrom/pos/snpid/ref/alt/imputation-quality)
        print(pb, chr[i], "\t", string(pos[i]), "\t", SNPids[i], "\t",
            ref[i], "\t", alt[i], "\t.\tPASS\t.\tGT")
        # print ith record
        write_snp!(pb, x, i)
        bytesavailable(pb) > 1048576 && write(io, take!(pb))
        next!(pmeter)
    end
    write(io, take!(pb))
    close(io); close(pb) # close io and buffer
end

"""
Helper function for saving the `i`th record (SNP), not tracking phase information.
Here genotypes will be written as
0 => 0/0
1 => 1/0
2 => 1/1
"""
function write_snp!(pb::IOBuffer, X::AbstractMatrix, i::Int)
    x = @view(X[:, i])
    n = length(x)
    @inbounds for j in 1:n
        if x[j] == 0
            print(pb, "\t0/0")
        elseif x[j] == 1
            print(pb, "\t1/0")
        elseif x[j] == 2
            print(pb, "\t1/1")
        else
            error("unphased genotypes can only be 0, 1, or 2 but got $(x[j])")
        end
    end
    print(pb, "\n")
    return nothing
end


text:
export Agent

"""
    Agent(;kwargs...)

One of the most commonly used [`AbstractAgent`](@ref).

Generally speaking, it does nothing but update the trajectory and policy appropriately in different stages.

# Keywords & Fields

- `policy`::[`AbstractPolicy`](@ref): the policy to use
- `trajectory`::[`AbstractTrajectory`](@ref): used to store transitions between an agent and an environment
- `role=DEFAULT_PLAYER`: used to distinguish different agents
"""
Base.@kwdef mutable struct Agent{P<:AbstractPolicy,T<:AbstractTrajectory,R} <: AbstractAgent
    policy::P
    trajectory::T
    role::R = DEFAULT_PLAYER
end

RLBase.get_role(agent::Agent) = agent.role

#####
# EpisodicCompactSARTSATrajectory
#####
function (agent::Agent{<:AbstractPolicy,<:EpisodicCompactSARTSATrajectory})(
    ::PreEpisodeStage,
    obs,
)
    empty!(agent.trajectory)
    nothing
end

function (agent::Agent{<:AbstractPolicy,<:EpisodicCompactSARTSATrajectory})(
    ::PreActStage,
    obs,
)
    action = agent.policy(obs)
    push!(agent.trajectory; state = get_state(obs), action = action)
    update!(agent.policy, agent.trajectory)
    action
end

function (agent::Agent{<:AbstractPolicy,<:EpisodicCompactSARTSATrajectory})(
    ::PostActStage,
    obs,
)
    push!(agent.trajectory; reward = get_reward(obs), terminal = get_terminal(obs))
    nothing
end

function (agent::Agent{<:AbstractPolicy,<:EpisodicCompactSARTSATrajectory})(
    ::PostEpisodeStage,
    obs,
)
    action = agent.policy(obs)
    push!(agent.trajectory; state = get_state(obs), action = action)
    update!(agent.policy, agent.trajectory)
    action
end

#####
# Union{CircularCompactSARTSATrajectory, CircularCompactPSARTSATrajectory}
#####

function (
    agent::Agent{
        <:AbstractPolicy,
        <:Union{CircularCompactSARTSATrajectory,CircularCompactPSARTSATrajectory},
    }
)(
    ::PreEpisodeStage,
    obs,
)
    if length(agent.trajectory) > 0
        pop!(agent.trajectory, :state, :action)
    end
    nothing
end

function (
    agent::Agent{
        <:AbstractPolicy,
        <:Union{CircularCompactSARTSATrajectory,CircularCompactPSARTSATrajectory},
    }
)(
    ::PreActStage,
    obs,
)
    action = agent.policy(obs)
    push!(agent.trajectory; state = get_state(obs), action = action)
    update!(agent.policy, agent.trajectory)
    action
end

function (
    agent::Agent{
        <:AbstractPolicy,
        <:Union{CircularCompactSARTSATrajectory,CircularCompactPSARTSATrajectory},
    }
)(
    ::PostActStage,
    obs,
)
    push!(agent.trajectory; reward = get_reward(obs), terminal = get_terminal(obs))
    nothing
end

function (
    agent::Agent{
        <:AbstractPolicy,
        <:Union{CircularCompactSARTSATrajectory,CircularCompactPSARTSATrajectory},
    }
)(
    ::PostEpisodeStage,
    obs,
)
    action = agent.policy(obs)
    push!(agent.trajectory; state = get_state(obs), action = action)
    update!(agent.policy, agent.trajectory)
    action
end

#####
# VectorialCompactSARTSATrajectory
#####

function (agent::Agent{<:AbstractPolicy,<:VectorialCompactSARTSATrajectory})(
    ::PreEpisodeStage,
    obs,
)
    if length(agent.trajectory) > 0
        pop!(agent.trajectory, :state, :action)
    end
    nothing
end

function (agent::Agent{<:AbstractPolicy,<:VectorialCompactSARTSATrajectory})(
    ::PreActStage,
    obs,
)
    action = agent.policy(obs)
    push!(agent.trajectory; state = get_state(obs), action = action)
    update!(agent.policy, agent.trajectory)
    action
end

function (agent::Agent{<:AbstractPolicy,<:VectorialCompactSARTSATrajectory})(
    ::PostActStage,
    obs,
)
    push!(agent.trajectory; reward = get_reward(obs), terminal = get_terminal(obs))
    nothing
end

function (agent::Agent{<:AbstractPolicy,<:VectorialCompactSARTSATrajectory})(
    ::PostEpisodeStage,
    obs,
)
    action = agent.policy(obs)
    push!(agent.trajectory; state = get_state(obs), action = action)
    update!(agent.policy, agent.trajectory)
    action
end


text:
@with_kw struct System{T}
    foo::Int = 6
    L::Array{Complex{T},1}
    Et::Array{T,1}
    fiber::Fiber{T, Symbol}
    pulse::Pulse{T}
    g::Grid{T}
    gases::Array{Gas,1}
    χ3::Float64
    gw::Array{T,1}
    gt::Array{T,1}
    wb::Array{T,1}
    plasma::Bool
    thg::Bool
    planrfft::FFTW.rFFTWPlan{Float64,-1,false,1}
    planirfft::AbstractFFTs.ScaledPlan{Complex{Float64},FFTW.rFFTWPlan{Complex{Float64},1,false,1},Float64}
    Hw::Array{Complex{Float64},1}
    fr::Float64
    ret::Array{T,1} = zeros(Float64, g.npts)#
    Pe0::Array{T,1} = zeros(Float64, g.npts)
    Pe::Array{T,1}= zeros(Float64, g.npts)
    PeTemp::Array{T,1}= zeros(Float64, g.npts)
    Na::Array{T,1}= zeros(Float64, g.npts)
    Ne::Array{T,1}= zeros(Float64, g.npts)
    WE::Array{T,1}= zeros(Float64, g.npts)
    temp1::Array{T,1} = zeros(Float64, g.npts)
    temp2::Array{T,1} = zeros(Float64, g.npts)
    temp3::Array{T,1} = zeros(Float64, g.npts)
    temp4::Array{T,1} = zeros(Float64, g.npts)
    temp5::Array{T,1} = zeros(Float64, g.npts)
    Pe_l::Array{T,1} = zeros(Float64, g.npts)
    Pe_f::Array{T,1} = zeros(Float64, g.npts)
    PNL::Array{Complex{T},1} = zeros(Complex{Float64}, Int(g.npts/2+1))
    dEt::Array{T,1} = zeros(Float64, g.npts)
    Ew_0::Array{Complex{T},1} = zeros(Complex{Float64}, Int(g.npts/2+1))
    Ew2::Array{Complex{T},1} = zeros(Complex{Float64}, Int(g.npts/2+1))
    Rt::Array{T,1} = zeros(Float64, g.npts)
    preRt::Array{T,1} = zeros(Float64, g.npts)
    pltype::Symbol
end

max_stars_repo_name: SabbahMohammed/Dev

max_issues_repo_name: SabbahMohammed/Dev

max_forks_repo_name: SabbahMohammed/Dev