text:
const USE_GPU     = haskey(ENV, "USE_GPU"    ) ? parse(Bool   , ENV["USE_GPU"    ]) : true
const do_viz      = haskey(ENV, "DO_VIZ"     ) ? parse(Bool   , ENV["DO_VIZ"     ]) : false
const do_save     = haskey(ENV, "DO_SAVE"    ) ? parse(Bool   , ENV["DO_SAVE"    ]) : false
const do_save_vis = haskey(ENV, "DO_SAVE_VIZ") ? parse(Bool   , ENV["DO_SAVE_VIZ"]) : false
const gpu_id      = haskey(ENV, "GPU_ID"     ) ? parse(Int    , ENV["GPU_ID"     ]) : 0
const nx          = haskey(ENV, "RESOL"      ) ? parse(Int    , ENV["RESOL"      ]) : 2048 - 1
const ny          = haskey(ENV, "RESOL"      ) ? parse(Int    , ENV["RESOL"      ]) : 2048 - 1
const vrpow       = haskey(ENV, "VISCR"      ) ? parse(Int    , ENV["VISCR"      ]) : 9
const nsub        = haskey(ENV, "NSUB"       ) ? parse(Int    , ENV["NSUB"       ]) : 1
const Re_mlt      = haskey(ENV, "REMULT"     ) ? parse(Float64, ENV["REMULT"     ]) : 1.0
const ρg0         = haskey(ENV, "RHOG0"      ) ? parse(Float64, ENV["RHOG0"      ]) : 0.0
const ρgi         = haskey(ENV, "RHOGI"      ) ? parse(Float64, ENV["RHOGI"      ]) : 0.0
const εbg         = haskey(ENV, "EBG"        ) ? parse(Float64, ENV["EBG"        ]) : 1.0
const simname     = haskey(ENV, "SIMNAME"    ) ?                ENV["SIMNAME"    ]  : ""
###
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 2)
else
    @init_parallel_stencil(Threads, Float64, 2)
end
using Plots, Printf, Statistics, LinearAlgebra, MAT, Random

@parallel function smooth!(A2::Data.Array, A::Data.Array, fact::Data.Number)
    @inn(A2) = @inn(A) + 1.0/4.1/fact*(@d2_xi(A) + @d2_yi(A))
    return
end

@parallel function compute_maxloc!(Musτ2::Data.Array, Musτ::Data.Array)
    @inn(Musτ2) = @maxloc(Musτ)
    return
end

@parallel function compute_iter_params!(dτ_Rho::Data.Array, Gdτ::Data.Array, Musτ::Data.Array, Vpdτ::Data.Number, Re::Data.Number, r::Data.Number, max_lxy::Data.Number)
    @all(dτ_Rho) = Vpdτ*max_lxy/Re/@all(Musτ)
    @all(Gdτ)    = Vpdτ^2/@all(dτ_Rho)/(r+2.0)
    return
end

@parallel function compute_P!(∇V::Data.Array, Pt::Data.Array, Vx::Data.Array, Vy::Data.Array, Gdτ::Data.Array, r::Data.Number, dx::Data.Number, dy::Data.Number)
    @all(∇V)  = @d_xa(Vx)/dx + @d_ya(Vy)/dy
    @all(Pt)  = @all(Pt) - r*@all(Gdτ)*@all(∇V)
    return
end

macro μ_veτ()    esc(:( 1.0/(1.0/@all(Gdτ) + 1.0/@all(Mus)) )) end
macro μ_veτ_av() esc(:( 1.0/(1.0/@av(Gdτ)  + 1.0/@av(Mus))  )) end
@parallel function compute_τ!(τxx::Data.Array, τyy::Data.Array, τxy::Data.Array, Vx::Data.Array, Vy::Data.Array, Mus::Data.Array, Gdτ::Data.Array, dx::Data.Number, dy::Data.Number)
    @all(τxx) = 2.0*@μ_veτ()*(@d_xa(Vx)/dx + @all(τxx)/@all(Gdτ)/2.0)
    @all(τyy) = 2.0*@μ_veτ()*(@d_ya(Vy)/dy + @all(τyy)/@all(Gdτ)/2.0)
    @all(τxy) = 2.0*@μ_veτ_av()*(0.5*(@d_yi(Vx)/dy + @d_xi(Vy)/dx) + @all(τxy)/@av(Gdτ)/2.0)
    return
end

@parallel function compute_dV!(Rx::Data.Array, Ry::Data.Array, dVx::Data.Array, dVy::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τxy::Data.Array, dτ_Rho::Data.Array, ρg::Data.Array, dx::Data.Number, dy::Data.Number)
    @all(Rx)  = @d_xi(τxx)/dx + @d_ya(τxy)/dy - @d_xi(Pt)/dx
    @all(Ry)  = @d_yi(τyy)/dy + @d_xa(τxy)/dx - @d_yi(Pt)/dy - @av_yi(ρg)
    @all(dVx) = @av_xi(dτ_Rho)*@all(Rx)
    @all(dVy) = @av_yi(dτ_Rho)*@all(Ry)
    return
end

@parallel function compute_V!(Vx::Data.Array, Vy::Data.Array, dVx::Data.Array, dVy::Data.Array)
    @inn(Vx) = @inn(Vx) + @all(dVx)
    @inn(Vy) = @inn(Vy) + @all(dVy)
    return
end

@parallel_indices (iy) function bc_x!(A::Data.Array)
    A[1  , iy] = A[2    , iy]
    A[end, iy] = A[end-1, iy]
    return
end

@parallel_indices (ix) function bc_y!(A::Data.Array)
    A[ix, 1  ] = A[ix, 2    ]
    A[ix, end] = A[ix, end-1]
    return
end

@parallel_indices (ix,iy) function create_inclusions!(Mus::Data.Array, ρg::Data.Array, Xi::Data.Array, Yi::Data.Array, r2::Data.Number, mus0::Data.Number, musi::Data.Number, ρg0::Data.Number, ρgi::Data.Number, dx::Data.Number, dy::Data.Number, lx::Data.Number, ly::Data.Number)
    if ix <= size(Mus,1) && iy <= size(Mus,2)
        x = (ix-1)*dx + 0.5*dx - 0.5*lx
        y = (iy-1)*dy + 0.5*dy - 0.5*ly
        isinside = false
        for idx = eachindex(Xi)
            rad2 = (x - Xi[idx])^2 + (y - Yi[idx])^2
            if rad2 < r2
                isinside = true
                break
            end
        end
        Mus[ix,iy] = isinside ? musi : mus0
        ρg[ix,iy]  = isinside ? ρgi  : ρg0
    end
    return
end

function generate_inclusions(lx,ly,nsub,ri)
    li    = min(3*nsub*ri,lx-2.5*ri)
    dx    = li/(nsub-1)
    dy    = dx*sqrt(3)/2
    if nsub == 1
        xs   = Float64[0]
        ys   = Float64[0]
        # random offset
        jitx = [0.25*lx*(2*rand()-1)]
        jity = [0.25*ly*(2*rand()-1)]
    else
        xs    = Float64[]
        ys    = Float64[]
        ox    = -li/2
        oy    = -li/2*sqrt(3)/2
        for j = 1:nsub
            for i = 1:(nsub-mod(j-1,2))
                push!(xs, ox + (i-1)*dx + dx/2*mod(j-1,2))
                push!(ys, oy + (j-1)*dy)
            end
        end
        # random offset
        jitx = 0.5*(dx-2*ri) .* (2 .* rand(length(xs)) .- 1.0)
        jity = 0.5*(dy-2*ri) .* (2 .* rand(length(xs)) .- 1.0)
        #exponenial falloff
        idx = xs.*jitx .> 0.0
        idy = ys.*jity .> 0.0
        jitx[idx] .*= exp.(-(xs[idx]/2).^2)
        jity[idy] .*= exp.(-(ys[idy]/2).^2)
    end
    # jitter
    xs .+= jitx
    ys .+= jity
    return xs,ys
end

@views function Stokes2D()
    # Set random seed for reproducibility
    Random.seed!(1855 + nsub)
    # Set CUDA device
    CUDA.device!(gpu_id)
    # Physics
    lx, ly    = 10.0, 10.0          # domain extends
    μs0       = 1.0                 # matrix viscosity
    μsi       = 10.0^(-vrpow)       # inclusion viscosity
    ri        = sqrt(lx*ly*0.005/π) # inclusion radius
    # Numerics
    iterMax   = 100max(nx,ny)       # maximum number of pseudo-transient iterations
    nout      = 200                 # error checking frequency
    ε_V       = 1e-5                # nonlinear absolute tolerence
    ε_∇V      = 1e-3                # nonlinear absolute tolerence
    CFL       = 0.9/sqrt(2)         # stability condition
    Re        = 5π*Re_mlt           # Reynolds number                     (numerical parameter #1)
    r         = 1.0                 # Bulk to shear elastic modulus ratio (numerical parameter #2)
    # Derived numerics
    dx, dy    = lx/nx, ly/ny # cell sizes
    max_lxy   = max(lx,ly)
    Vpdτ      = min(dx,dy)*CFL
    nsm       = 10
    # Array allocations
    Pt        = @zeros(nx  ,ny  )
    ∇V        = @zeros(nx  ,ny  )
    τxx       = @zeros(nx  ,ny  )
    τyy       = @zeros(nx  ,ny  )
    τxy       = @zeros(nx-1,ny-1)
    Rx        = @zeros(nx-1,ny-2)
    Ry        = @zeros(nx-2,ny-1)
    dVx       = @zeros(nx-1,ny-2)
    dVy       = @zeros(nx-2,ny-1)
    Mus       = @zeros(nx  ,ny  )
    Mus2      = @zeros(nx  ,ny  )
    Musτ      = @zeros(nx  ,ny  )
    Gdτ       = @zeros(nx  ,ny  )
    dτ_Rho    = @zeros(nx  ,ny  )
    ρg        = @zeros(nx  ,ny  )
    ρg2       = @zeros(nx  ,ny  )
    # Initial conditions
    Vx        =  zeros(nx+1,ny  )
    Vy        =  zeros(nx  ,ny+1)
    Vx        = Data.Array( -εbg.*[((ix-1)*dx -0.5*lx) for ix=1:size(Vx,1), iy=1:size(Vx,2)] )
    Vy        = Data.Array(  εbg.*[((iy-1)*dy -0.5*ly) for ix=1:size(Vy,1), iy=1:size(Vy,2)] )
    Xi,Yi     = Data.Array.(generate_inclusions(lx,ly,nsub,ri))
    @parallel create_inclusions!(Mus,ρg,Xi,Yi,ri^2,μs0,μsi,ρg0,ρgi,dx,dy,lx,ly)
    Mus2     .= Mus
    ρg2      .= ρg
    for _ = 1:nsm
        @parallel smooth!(Mus2, Mus, 1.0)
        Mus, Mus2 = Mus2, Mus
        @parallel smooth!(ρg2, ρg, 1.0)
        ρg, ρg2 = ρg2, ρg
    end
    Musτ     .= μs0
    # Time loop
    @parallel compute_iter_params!(dτ_Rho, Gdτ, Musτ, Vpdτ, Re, r, max_lxy)
    err_V=2*ε_V; err_∇V=2*ε_∇V; iter=0; err_evo1=[]; err_evo2=[]
    while !((err_V <= ε_V) && (err_∇V <= ε_∇V)) && (iter <= iterMax)
        if (iter==11)  global wtime0 = Base.time()  end
        @parallel compute_P!(∇V, Pt, Vx, Vy, Gdτ, r, dx, dy)
        @parallel compute_τ!(τxx, τyy, τxy, Vx, Vy, Mus, Gdτ, dx, dy)
        @parallel compute_dV!(Rx, Ry, dVx, dVy, Pt, τxx, τyy, τxy, dτ_Rho, ρg, dx, dy)
        @parallel compute_V!(Vx, Vy, dVx, dVy)
        @parallel (1:size(Vx,1)) bc_y!(Vx)
        @parallel (1:size(Vy,2)) bc_x!(Vy)
        iter += 1
        if iter % nout == 0
            Pmin, Pmax = minimum(Pt), maximum(Pt)
            Vmin, Vmax = minimum(Vy), maximum(Vy)
            norm_Rx    = norm(Rx)/(Pmax-Pmin)*lx/sqrt(length(Rx))
            norm_Ry    = norm(Ry)/(Pmax-Pmin)*lx/sqrt(length(Ry))
            norm_∇V    = norm(∇V)/(Vmax-Vmin)*lx/sqrt(length(∇V))
            err_V      = maximum([norm_Rx, norm_Ry])
            err_∇V     = norm_∇V
            push!(err_evo1, maximum([norm_Rx, norm_Ry])); push!(err_evo2,iter)
            @printf("Total steps = %d, err_V = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e], err_∇V = %1.3e \n", iter, err_V, norm_Rx, norm_Ry, err_∇V)
        end
    end
    # Performance
    wtime    = Base.time() - wtime0
    A_eff    = (3*2)/1e9*nx*ny*sizeof(Data.Number)  # Effective main memory access per iteration [GB] (Lower bound of required memory access: Te has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
    wtime_it = wtime/(iter-10)                      # Execution time per iteration [s]
    T_eff    = A_eff/wtime_it                       # Effective memory throughput [GB/s]
    @printf("Total steps = %d, time = %1.3e sec (@ T_eff = %1.2f GB/s) \n", iter, wtime, round(T_eff, sigdigits=2))
    # Visualisation
    if do_viz
        X, Y, Yv  = dx/2:dx:lx-dx/2, dy/2:dy:ly-dy/2, 0:dy:ly
        p1 = heatmap(X,  Y, Array(Pt)', aspect_ratio=1, xlims=(X[1],X[end]), ylims=(Y[1],Y[end]), c=:viridis, title="Pressure")
        p2 = heatmap(X, Yv, Array(Vy)', aspect_ratio=1, xlims=(X[1],X[end]), ylims=(Yv[1],Yv[end]), c=:viridis, title="Vy")
        p4 = heatmap(X[2:end-1], Yv[2:end-1], log10.(abs.(Array(Ry)')), aspect_ratio=1, xlims=(X[2],X[end-1]), ylims=(Yv[2],Yv[end-1]), c=:viridis, title="log10(Ry)")
        p5 = plot(err_evo2,err_evo1, legend=false, xlabel="# iterations", ylabel="log10(error)", linewidth=2, markershape=:circle, markersize=3, labels="max(error)", yaxis=:log10)
        display(plot(p1, p2, p4, p5))
    end
    if do_save_vis && Re_mlt ≈ 1.0 && vrpow == 9 # workaround to avoid saving too much data
        outdir = joinpath("..", "out_visu", simname, "VISCR_$(vrpow)_NSUB_$(nsub)_REMULT_$(Re_mlt)")
        !ispath(outdir) && mkpath(outdir)
        matwrite(joinpath(outdir,"Stokes2D.mat"), Dict("Pt_2D"=> Array(Pt), "Mus_2D"=> Array(Mus), "Txy_2D"=> Array(τxy), "Txx_2D"=> Array(τxx), "Tyy_2D"=> Array(τyy), "Vx_2D"=> Array(Vx), "Vy_2D"=> Array(Vy), "dx_2D"=> dx, "dy_2D"=> dy); compress = true)
    end
    if do_save
        !ispath("../output") && mkpath("../output")
        open("out_Stokes2D_$(simname)_param.txt","a") do io
            println(io, "$(vrpow) $(nsub) $(Re_mlt) $(iter) $(err_V) $(err_∇V)")
        end
    end
    return
end

Stokes2D()

ext: jl


text:
module TestGramian
using CovarianceFunctions
using Test
using LinearAlgebra
using ToeplitzMatrices
using CovarianceFunctions
using CovarianceFunctions: PeriodicInput, LazyMatrixSum
using BlockFactorizations

@testset "Gramian properties" begin
    n = 8
    x = randn(Float64, n)
    k = CovarianceFunctions.EQ()
    G = gramian(k, x)

    @test issymmetric(G)
    M = Matrix(G)
    @test M ≈ AbstractMatrix(G)
    @test G' ≈ M'
    @test transpose(G) ≈ transpose(M)
    @test M ≈ k.(x, x')
    @test issymmetric(G) && issymmetric(M)
    @test ishermitian(G) && ishermitian(M)
    @test isposdef(G) && isposdef(M) # if M is very low-rank, this can fail -> TODO: ispsd()

    G = gramian(k, x, x .+ randn(size(x)))
    M = Matrix(G)
    @test !issymmetric(G) && !issymmetric(M)
    @test !isposdef(G) && !isposdef(M)

    G = gramian(k, x, copy(x))
    @test issymmetric(G) # testing that value check also works
    @test isposdef(G)

    # testing standard Gramian
    G = gramian(x)
    @test G ≈ dot.(x, x')

    x = randn(Float64, n)
    k = CovarianceFunctions.EQ()
    G = gramian(k, x)
    # type promotion
    @test k(x[1], x[2]) isa typeof(G[1, 2]) # test promotion in inner constructor
    @test k(x[1], x[2]) isa eltype(G) # test promotion in inner constructor
    x = tuple.(x, x)
    G = gramian(k, x)
    @test k(x[1], x[2]) isa typeof(G[1,1])  # with tuples (or arrays)
    @test k(x[1], x[2]) isa eltype(G)

    # adding diagonal to Gramian is done lazily
    D = 1e-6I(n)
    @test D + G isa LazyMatrixSum
    @test G + D isa LazyMatrixSum

    # rectangular multiply
    x = randn(n)
    y = randn(2n)
    G = gramian(k, x, y)
    @test size(G) == (n, 2n)
    a = randn(2n)
    b = G*a
    @test length(b) == n
    @test b ≈ Matrix(G) * a

    y = randn(2n)
    G = gramian(k, x, y)
    @test size(G) == (n, 2n)
    p = 3
    A = randn(2n, p)
    B = G*A
    @test size(B) == (n, p)
    @test B ≈ Matrix(G) * A

    # testing scalar indexing
    G = gramian(k, x, y)
    for i in 1:n
        for j in 1:n
            @test G[i, j] ≈ k(x[i], y[j])
        end
    end

    # testing non-scalar indexing
    i, j = 2:n-1, 3:2n-4
    Gij = G[i, j]
    @test Gij isa AbstractMatrix
    @test Gij ≈ gramian(k, x[i], y[j])

    i, j = 2:n-1, 2
    Gij = G[i, j]
    @test Gij isa AbstractVector
    @test Gij ≈ vec(gramian(k, x[i], [y[j]]))

    Gij = G[j, i]
    @test Gij isa AbstractVector
    @test Gij ≈ vec(gramian(k, [x[j]], y[i]))
end

@testset "Gramian factorization" begin
    rtol = 1e-6
    k = CovarianceFunctions.EQ()
    n = 64
    x = randn(n)
    G = gramian(k, x)
    F = factorize(G)
    @test F isa CholeskyPivoted
    @test rank(F) < n ÷ 2 # this should be very low rank
    @test isapprox(Matrix(F), G, rtol = rtol)

    # pivoted cholesky
    G = gramian(k, x)
    F = cholesky(G, Val(true), check = false, tol = 1e-6)
    @test F isa CholeskyPivoted
    @test isapprox(Matrix(F), G, rtol = rtol)

    # regular cholesky
    k = CovarianceFunctions.Exponential() # does not yield low rank matrix
    G = gramian(k, x)
    F = cholesky(G)
    @test F isa Cholesky
    @test isapprox(Matrix(F), G, rtol = rtol)

    # testing gramian of matrix-valued anonymous kernel
    d = 3
    A = randn(d, d)
    k = (x, y)->A
    G = gramian(k, x)
    GA = Matrix(G)
    a = randn(d*n)
    @test size(G) == (d*n, d*n)
    @test G isa BlockFactorization
    @test G*a ≈ GA*a

    # adding diagonal by default maintains laziness
    @test I(d*n) + G isa LazyMatrixSum
    @test G + I(d*n) isa LazyMatrixSum

    # same thing for block diagonal matrices
    D = BlockFactorization(Diagonal([randn(d, d) for _ in 1:n]))
    @test D + G isa LazyMatrixSum
    @test G + D isa LazyMatrixSum
end

@testset "toeplitz structure" begin
    n = 32
    x = range(-1, 1, n)
    k = CovarianceFunctions.EQ()
    G = gramian(k, x)
    @test G isa SymmetricToeplitz
    @test size(G) == (n, n)
    @test G*x ≈ Matrix(G)*x

    @test Matrix(G) ≈ Matrix(Gramian(k, x))
    G = gramian(k, x, PeriodicInput()) # periodic boundary conditions
    @test G isa Circulant
    @test size(G) == (n, n)
    @test G*x ≈ Matrix(G)*x

    k = (x, y)->CovarianceFunctions.EQ()(x, y)
    G = gramian(k, x)
    @test G isa Gramian # if we can't infer the input trait of k, falls back to lazy representation
    @test Matrix(G) ≈ Matrix(Gramian(k, x))

    # testing that we can hook anonymous kernel into the system
    CovarianceFunctions.input_trait(::typeof(k)) = CovarianceFunctions.IsotropicInput()
    G = gramian(k, x)
    @test G isa SymmetricToeplitz
    @test Matrix(G) ≈ Matrix(Gramian(k, x))

    # testing non-symmetric Toeplitz systems
    y = x .+ randn() # shifted version of x with same step length
    G = gramian(k, x, y)
    @test G isa Toeplitz
    @test Matrix(G) ≈ Matrix(Gramian(k, x, y))

    y = range(-1, 1, n÷2) # if y has different step length, defaults to gramian
    G = gramian(k, x, y)
    @test G isa Gramian
    @test Matrix(G) ≈ Matrix(Gramian(k, x, y))

end # test solves and multiplications

end # TestGramian


text:
################################################################################
#
#  Local factor for dyadic primes
#
################################################################################

# Returns the local mass factor for dyadic primes
# The local scales must be be 0 or 1
# Otherwise the function returns 0

function _local_factor_dyadic(L::HermLat, p)
  @assert isdyadic(p)
  S = base_ring(L)
  E = nf(S)
  K = base_field(nf(S))
  P = prime_decomposition(S, p)[1][1]
  valscale = valuation(scale(L), P)
  @assert valscale >= 0
  val = div(valscale, 2)
  if !iszero(val )
    s = elem_in_nf(uniformizer(p))
    L = rescale(L, s^(-val))
  end

  G = _genus_symbol_kirschmer(L, p)::Vector{Tuple{Int, Int, Bool, Int, elem_type(K)}}
  # I should check if I can get this from the standard genus symbol

  if !(issubset([g[2] for g in G], [0, 1]))
    return zero(fmpq)
  end

  m0 = G[1][2] == 0 ? G[1][1] : 0
  m1 = G[end][2] == 1 ? G[end][1] : 0

  q = norm(p)
  m = m0 + m1
  @assert iseven(m1)
  @assert m >= 0
  @assert m0 >= 0
  @assert m1 >= 0
  k = div(m, 2)
  k0 = div(m0, 2)
  k1 = div(m1, 2)

  if isodd(m)
    return fmpq(1, 2) * _gauss(k, k1, q^2)
  end

  e = valuation(discriminant(S), p)
  f2 = div(e, 2)
  d = (-1)^((m - 1) * k) * det(rational_span(L))
  # b <=> L \isom M \perp H0^* \perp H1^* where M is unimodular.
  b = (m1 == 0) || (m0 != 0 && G[1][4] != G[2][4])
  lf = fmpq(1, 2) * _gauss(k, k0, q^2)
  l = div((b ? G[1][4] : G[end][4]), 2)

  if islocal_norm(E, d, p)
    h0 = m0 == 0 || G[1][4] == 2 * f2
    h1 = m1 == 0 || G[end][4] == 2 * div(e + 1, 2)
    # L_p = H(0)^k0 \perp H(1)^k1
    if h0 && h1
      t = iseven(e) ? k1 : k0
      return (q^t + 1) * lf
    end

    if !b
      if iseven(e)
        return q^(m * (f2 - l) - k1) * (q^m1 - 1) * lf
      else
        return q^(m * (f2 - l) - k0) * (q^m1 - 1) * lf
      end
    else
      if iseven(e)
        return q^(m * (f2 - 1 - l) + k1) * (q^m0 - 1) * lf
      else
        return q^(m * (f2 - 1 ) - k0)    * (q^m0 - 1) * lf
      end
    end
  else # non-hyperbolic case
    if isodd(e)
      if b && l == div(e - 1, 2)
        return (q^k0 - 1) * lf
      elseif b
        return q^(m * (f2 - l) - k0) * (q^m0 - 1) * lf
      else
        return q^(m * (f2 - l) + k0) * (q^m1 - 1) * lf
      end
    end

    if b
      return q^(m * (f2 - 1 - l) + k1) * (q^m0 - 1) * lf
    elseif l == div(e, 2) # e is even
      return q^(k1 - 1) * lf
    else
      return q^(m * (f2 - l) - k1) * (q^m1 - 1) * lf
    end
  end
  throw(error("Impossible"))
end

################################################################################
#
#  Local mass factor for maximal lattices
#
################################################################################

function _local_factor_maximal(L::HermLat, p)
  S = base_ring(L)
  E = nf(S)
  K = base_field(E)
  m = rank(L)
  lp = prime_decomposition(S, p)
  ram = length(lp) == 1 && lp[1][2] == 2
  if isodd(m) && ram
    return fmpq(1, 2)
  end
  G = gram_matrix_of_basis(L)
  disc = _discriminant(G)
  if !islocal_norm(E, K(disc), p)
    q = norm(p)
    if ram
      return fmpq(q^m - 1, 2*q + 2)
    else
      return fmpq(q^m - (-1)^m, q + 1)
    end
  end
  return fmpq(1)
end

################################################################################
#
#  Generic local mass factor
#
################################################################################

function _local_factor_generic(L::HermLat, p)
  K = fixed_field(L)
  a = _isdefinite(rational_span(L))
  def = !iszero(a)
  S = base_ring(L)
  lp = prime_decomposition(S, p)
  P = lp[1][1]
  val = valuation(norm(L) * S, P)
  if length(lp) == 1 && lp[1][2] == 2
    @assert val >= 0
    val = div(val, 2)
  end

  ss = uniformizer(p)^(-val)

  if def
    R = base_ring(base_ring(L))
    rlp = real_places(K)
    A::GrpAbFinGen, _exp, _log = infinite_primes_map(R, rlp, p)
    sa = ss * a
    t = (1 + _exp(A(Int[ sign(sa, rlp[j]) == 1 ? 0 : 1 for j in 1:length(rlp)]))::elem_type(R))
    @assert t - 1 in p
    @assert all(Bool[sign(t, rlp[i]) == sign(sa, rlp[i]) for i in 1:length(rlp)])
    ss = ss * t
  end
  L = rescale(L, ss)

  chain = typeof(L)[L]
  ok, LL = ismaximal_integral(L, p)
  while !ok
    push!(chain, LL)
    ok, LL = ismaximal_integral(LL, p)
  end
  f = _local_factor_maximal(L, p)
  for i in 1:(length(chain) - 1)
    M, E = maximal_sublattices(chain[i + 1], P, use_auto = false)
    lM = length(M)
    # should be use_auto = def)
    _f = 0
    for j in 1:lM
      if islocally_isometric(chain[i], M[j], p)
        _f += E[j]
      end
    end

    f = f * _f
    M, E = minimal_superlattices(chain[i], P, use_auto = false)
    # should be use_auto = def)
    lM = length(M)
    _f = 0
    for j in 1:lM
      if islocally_isometric(chain[i + 1], M[j], p)
        _f += E[j]
      end
    end
    f = divexact(f, _f)
  end
  return f
end

################################################################################
#
#  Local mass factor
#
################################################################################

function local_factor(L::HermLat, p)
  S = base_ring(L)
  E = nf(S)
  K = base_field(E)
  q = norm(p)
  lp = prime_decomposition(S, p)
  ram = length(lp) == 1 && lp[1][2] == 2
  if ram && isdyadic(p)
    lf = _local_factor_dyadic(L, p)
    # TODO: Use Cho's recipe if p unramified over Z.
    if iszero(lf)
      lf = _local_factor_generic(L, p)
      return lf
    end
    return lf
  end
  lp = prime_decomposition(S, p)
  split = length(lp) > 1 && !ram
  _, G, s = jordan_decomposition(L, p)
  if length(s) == 1 && !ram
    return fmpq(1)
  end

  m = rank(L)
  local f::fmpq
  if ram
    f = fmpq(group_order("Sp", 2 * div(m ,2), q))
  else
    f = fmpq(group_order(split ? "GL" : "U", m, q))
  end

  d = ram ? 1 : 2
  N = zero(fmpq)

  for i in 1:length(s)
    mi = ncols(G[i])
    ri = sum(Int[ncols(G[j]) for j in (i + 1):length(s)])
    if ram
      N = N - div(s[i], 2) * mi^2
      if isodd(s[i])
        N = N - (mi + 1) * div(mi, 2)
        f = divexact(f, group_order("Sp", mi, q))
      else
        det = K(_discriminant(G[i]))
        isnorm = islocal_norm(E, det, p)
        f = divexact(f, group_order(isnorm ? "O+" : "O-", mi, q))
      end
    else
      N = N - s[i] * mi^2
      f = divexact(f, group_order(split ? "GL" : "U", mi, q))
    end
    N = N - d * s[i] * mi * ri
    @assert mod(d * s[i] * mi * m, 2) == 0
    N = N + divexact(d * s[i] * mi * m, 2)
  end

  # Fix the difference coming from the discriminant
  if ram && iseven(m)
    N = N + div(m ,2)
  end

  return q^(Int(FlintZZ(N))) * f
end

function _standard_mass(L::HermLat, prec::Int = 10)
  S = base_ring(L)
  E = nf(S)
  K = base_field(E)
  n = absolute_degree(K)
  m = rank(L)
  stdmass = fmpq(2)^(1 - n * m)
  # K is totally real, so I can get the exact value "cheaply"
  for i in 2:2:m
    stdmass = stdmass * dedekind_zeta_exact(K, 1 - i)
  end

  wprec = 2 * prec

  RR = ArbField(wprec, cached = false)
  _stdmass = RR(stdmass)

  local relzeta::arb

  while true
    relzeta = prod(arb[_L_function(E, 1 - i, wprec) for i in 1:2:m])
    if radiuslttwopower(relzeta * _stdmass, -prec)
      break
    end
    wprec = wprec + 1
  end
  return _stdmass * relzeta
end

################################################################################
#
#  Mass of hermitian lattices
#
################################################################################

function mass(L::HermLat)
  @req isdefinite(L) "Lattice must be definite"
  m = rank(L)
  if m == 0
    return one(fmpq)
  end

  lm = local_mass(L)

  prec = 10
  mu = _minkowski_multiple(absolute_simple_field(nf(base_ring(L)))[1], rank(L))

  d = mu * numerator(lm)

  while true
    z = d * _standard_mass(L, prec)
    fl, b = unique_integer(z)
    if fl
      return abs(b)//d * lm
    end
    prec += 2
  end
end

function local_mass(L::HermLat)
  lf = fmpq(1)

  for p in bad_primes(L, discriminant = true)
    lf *= local_factor(L, p)
  end

  return lf
end

################################################################################
#
#  Misc
#
################################################################################

function _gauss(m, k, q)
  z = _gauss0(m, q)
  z = divexact(z, _gauss0(k, q))
  z = divexact(z, _gauss0(m - k, q))
  return z
end

function _gauss0(m, q)
  return fmpq(prod(fmpz[1 - q^i for i in 1:m]))
end


text:
module Junuran

import Base
import Libdl
export urgen_vnrou
export ursample

mutable struct UNUR_DISTR
end
mutable struct UNUR_PAR
end
mutable struct UNUR_GEN
end

struct UnuRanGenerator
  generator::Ptr{UNUR_GEN}
  type::String
  dim::Integer
  info::String
  unuran::Ptr{Nothing}
end

Base.show(io::IO, urg::UnuRanGenerator) = print(io, urg.info)
Base.show(io::IO, ::MIME"text/plain", urg::UnuRanGenerator) =
            print(io, "UNU.RAN generator\n", urg)

function urgen_vnrou(
  dim::Integer,
  pdf::Function,
  center::Union{Nothing,Vector{Float64}} = nothing,
  mode::Union{Nothing,Vector{Float64}} = nothing,
  lower::Union{Nothing,Vector{Float64}} = nothing,
  upper::Union{Nothing,Vector{Float64}} = nothing
)

  # Open the UNU.RAN library
  lib = Libdl.dlopen("/usr/local/lib/libunuran")

  # checks
  if dim < 2
    error("zzz")
  end

  if center !== nothing && length(center) != dim
    error("zzz")
  end

  if mode !== nothing && length(mode) != dim
    error("zzz")
  end

  if lower !== nothing && length(lower) != dim
    error("zzz")
  end

  if upper !== nothing && length(upper) != dim
    error("zzz")
  end

  # create continuous multivariate distribution
  distr = ccall(
    Libdl.dlsym(lib, :unur_distr_cvec_new),
    Ptr{UNUR_DISTR},
    (Cint, ),
    dim
  )

  # define the unnormalized pdf of the distribution
  pdf_j = function(xptr, distr)
    x = map(i -> unsafe_load(xptr, i), 1:dim)
    return pdf(x)::Cdouble
  end
  pdf_c = @cfunction(
    $pdf_j,
#    (xptr, distr) -> pdf(map(i -> unsafe_load(xptr, i), 1:dim))::Cdouble,
    Cdouble,
    (Ptr{Cdouble}, Ptr{UNUR_DISTR})
  )

  # set the pdf
  ccall(
    Libdl.dlsym(lib, :unur_distr_cvec_set_pdf),
    Cint,
    (Ptr{UNUR_DISTR}, Ptr{Cvoid}),
    distr, pdf_c
  )

  # set the rectangular domain
  if lower !== nothing && upper !== nothing
    ccall(
      Libdl.dlsym(lib, :unur_distr_cvec_set_domain_rect),
      Cint,
      (Ptr{UNUR_DISTR}, Ref{Cdouble}, Ref{Cdouble}),
      distr, lower, upper
    )
  end

  # set center
  if center !== nothing
    ccall(
      Libdl.dlsym(lib, :unur_distr_cvec_set_center),
      Cint,
      (Ptr{UNUR_DISTR}, Ref{Cdouble}),
      distr, center
    )
  end

  # set mode
  if mode !== nothing
    ccall(
      Libdl.dlsym(lib, :unur_distr_cvec_set_mode),
      Cint,
      (Ptr{UNUR_DISTR}, Ref{Cdouble}),
      distr, mode
    )
  end

  # create the parameters object
  par = ccall(
    Libdl.dlsym(lib, :unur_vnrou_new),
    Ptr{UNUR_PAR},
    (Ptr{UNUR_DISTR}, ),
    distr
  )

  # create the generator
  gen = ccall(
    Libdl.dlsym(lib, :unur_init),
    Ptr{UNUR_GEN},
    (Ptr{UNUR_PAR}, ),
    par
  )

  # TODO: check gen successful (i.e. not the null pointer)

  # destroy the distribution object
  ccall(
    Libdl.dlsym(lib, :unur_distr_free),
    Cvoid,
    (Ptr{UNUR_DISTR}, ),
    distr
  )

  # generator info
  info = ccall(
    Libdl.dlsym(lib, :unur_gen_info),
    Cstring,
    (Ptr{UNUR_GEN}, Cint),
    gen, 1
  )

  # output
  return UnuRanGenerator(
    gen,
    "cmv",
    dim,
    unsafe_string(info),
    lib
  )

end # urgen_vnrou

function ursample(
  urgen,
  n::Integer,
  seed::Union{Nothing,Integer} = nothing
)

  if seed !== nothing
    seed = convert(UInt32, abs(seed) + 1)
  end
  # set generator seed
  ccall(
    Libdl.dlsym(urgen.unuran, :unur_gen_seed),
    Cint,
    (Ptr{UNUR_GEN}, Culong),
    urgen.generator, seed === nothing ? rand(UInt32) + UInt32(1) : seed
  )

  out = Vector{Vector{Cdouble}}(undef, n)
  for i in 1:n
    out[i] = Vector{Cdouble}(undef, urgen.dim)
    ccall(
      Libdl.dlsym(urgen.unuran, :unur_sample_vec),
      Cint,
      (Ptr{UNUR_GEN}, Ref{Cdouble}),
      urgen.generator, out[i]
    )
  end
#=  ccall(
    Libdl.dlsym(urgen.unuran, :unur_reinit),
    Cint,
    (Ptr{UNUR_GEN}, ),
    urgen.generator
  )
=#
  return out
end # ursample

end # module


text:
using UniqueVectors: UniqueVector
using SparseArrays: spzeros, spdiagm
using LinearAlgebra: I, eigen, svd, diagm
using IterativeSolvers: lobpcg

export PSOModel
export lossless_modes_dense, lossy_modes_dense, lossless_modes_sparse

struct PSOModel{T,
        K<:AbstractMatrix{<:Number},
        G<:AbstractMatrix{<:Number},
        C<:AbstractMatrix{<:Number},
        P<:AbstractMatrix{<:Number}} <: AdmittanceModel{T}
    K::K
    G::G
    C::C
    P::P
    ports::UniqueVector{T}
    function PSOModel{T}(k, g, c, p, ports::UniqueVector{T}) where {T}
        l = size(k, 1)
        @assert all(isreal(m) for m in (k, g, c, p))
        @assert size(p) == (l, length(ports))
        @assert all(size(y) == (l, l) for y in (k, g, c))
        kk, gg, cc, pp = float.((k, g, c, p))
        K, G, C, P = typeof.((kk, gg, cc, pp))
        return new{T,K,G,C,P}(kk, gg, cc, pp, ports)
    end
end

"""
    PSOModel(K, G, C, P, ports::AbstractVector{T}) where {T}

Create a positive second order (PSO) model, representing the equations
`(K + G∂ₜ + C∂²ₜ)Φ = Px`, `y = Qᵀ∂ₜΦ`. (If the subscript `t` does not display in
the equations, install a font with greater Unicode coverage.)
"""
PSOModel(K, G, C, P, ports::AbstractVector{T}) where {T} =
    PSOModel{T}(K, G, C, P, UniqueVector(ports))

get_Y(pso::PSOModel) = (pso.K, pso.G, pso.C)

get_P(pso::PSOModel) = pso.P

get_ports(pso::PSOModel) = pso.ports

function partial_copy(pso::PSOModel;
        Y::Union{AbstractVector, Nothing} = nothing,
        P::Union{AbstractMatrix{<:Real}, Nothing} = nothing,
        ports::Union{AbstractVector, Nothing} = nothing)
    Y = isnothing(Y) ? get_Y(pso) : Y
    P = isnothing(P) ? get_P(pso) : P
    ports = isnothing(ports) ? get_ports(pso) : ports
    return PSOModel(Y..., P, UniqueVector(ports))
end

compatible(psos::AbstractVector{<:PSOModel}) = true

function canonical_gauge(pso::PSOModel)
    fullP = collect(pso.P)
    F = qr(fullP)
    (m, n) = size(fullP)
    transform = (F.Q * Matrix(I,m,m)) * [transpose(inv(F.R)) zeros(n, m-n);
        zeros(m-n, n) Matrix(I, m-n, m-n)]
    return apply_transform(pso, transform)
end

# NOTE we cannot write lossy_modes_sparse because KrylovKit.geneigsolve
# solves Av = λBv but requires that B is positive semidefinite and A is
# symmetric or Hermitian. For lossy modes, we can make B positive semidefinite
# and A nonsymmetric or A symmetric and B not positive semidefinite. If
# KrylovKit gains the ability to solve either type of eigenproblem we can
# write lossy_modes_sparse.

"""
    lossy_modes_dense(pso::PSOModel; min_freq::Real=0,
        max_freq::Union{Real, Nothing}=nothing, real_tol::Real=Inf, imag_tol::Real=Inf)

Use a dense eigenvalue solver to find the decaying modes of the PSOModel. Each mode consists
of an eigenvalue `λ` and a spatial distribution vector `v`. The frequency of the mode is
`imag(λ)/(2π)` and the decay rate is `-2*real(λ)/(2π)`. `min_freq` and `max_freq` give the
frequency range in which to find eigenvalues. Passivity is enforced using an inverse power
method solver, and `real_tol` and `imag_tol` are the tolerances used in that method to
determine convergence of the real and imaginary parts of the eigenvalue.
"""
function lossy_modes_dense(pso::PSOModel; min_freq::Real=0,
    max_freq::Union{Real, Nothing}=nothing, real_tol::Real=Inf, imag_tol::Real=Inf)
    K, G, C = get_Y(pso)
    Z = zero(K)
    M = collect([C Z; Z I]) \ collect([-G -K; I Z])
    M_scale = maximum(abs.(M))
    # for some reason eigen(Y \ X) seems to perform better than eigen(X, Y)
    values, vectors = eigen(M/M_scale)
    values *= M_scale
    freqs = imag.(values)/(2π)
    if max_freq == nothing
        max_freq = maximum(freqs)
    end
    inds = findall(max_freq .>= freqs .>= min_freq)
    values, vectors = values[inds], vectors[:, inds]
    # enforce passivity using inverse power method
    pred(x, y) = ((real(y) <= 0) && (abs(real(x-y)) < real_tol) &&
        (abs(imag(x-y)) < imag_tol))
    for i in 1:length(values)
        if real(values[i]) > 0
            λ, v = inv_power_eigen(M, v0=vectors[:, i], λ0=values[i], pred=pred)
            values[i] = λ
            vectors[:, i] = v
        end
    end
    # the second half of the eigenvector is the spatial part
    vectors = vectors[size(K, 1)+1:end, :]
    # normalize the columns
    return values, vectors./sqrt.(sum(abs2.(vectors), dims=1))
end

"""
    lossless_modes_dense(pso::PSOModel; min_freq::Real=0,
        max_freq::Union{Real, Nothing}=nothing)

Use a dense svd solver to find the modes of the PSOModel neglecting loss. Each mode consists
of an eigenvalue `λ` and an eigenvector `v`. The frequency of the mode is `imag(λ)/(2π)`
and the decay rate is `-2*real(λ)/(2π)`. `min_freq` and `max_freq` give the
frequency range in which to find eigenvalues.
"""
function lossless_modes_dense(pso::PSOModel; min_freq::Real = 0,
    max_freq::Union{Real, Nothing} = nothing)
    K, _, C = get_Y(pso)
    K = .5 * (transpose(K) + K)
    C = .5 * (transpose(C) + C)
    sqrt_C, sqrt_K = sqrt(collect(C)), sqrt(collect(K))
    M = sqrt_C \ sqrt_K
    scale = maximum(abs.(M))
    U, S, Vt = svd(M/scale)
    freqs = S * scale/(2π)
    vectors = real.(sqrt_C \ U)
    if max_freq == nothing
        max_freq = maximum(freqs)
    end
    inds = findall(max_freq .>= freqs .>= min_freq)
    freqs, vectors = freqs[inds], vectors[:, inds]
    inds = sortperm(freqs)
    return freqs[inds] * 2π * 1im, vectors[:, inds]
end

"""
    lossless_modes_sparse(pso::PSOModel; num_modes = 1, maxiter = 200)

Use a sparse generalized eigenvalue solver to find the `num_modes` lowest frequency modes
of the PSOModel neglecting loss. Each mode consists of an eigenvalue `λ` and an eigenvector
`v`. The frequency of the mode is `imag(λ)/(2π)` and the decay rate is `-2*real(λ)/(2π)`.
"""
function lossless_modes_sparse(pso::PSOModel; num_modes = 1, maxiter = 200)
    K, _, C = get_Y(pso)
    K = (K + transpose(K))/2
    C = (C + transpose(C))/2
    K_scale, C_scale = maximum(abs.(K)), maximum(abs.(C))
    result = lobpcg(K/K_scale, C/C_scale, false, num_modes, maxiter=maxiter)
    values = result.λ * (K_scale/C_scale)
    vectors = result.X
    return real.(sqrt.(values)) * 1im, vectors
end

"""
    port_matrix(c::Circuit, port_edges::Vector{<:Tuple})
Returns a port matrix suitable for constructing a `PSOModel` based on
`c::Circuit`, where the columns of the port matrix correspond to edges between
circuit vertices. The edges are given as tuples of vertex names in `port_edges`.
"""
function port_matrix(c::Circuit, port_edges::Vector{<:Tuple})
    coord_matrix = coordinate_matrix(c)
    if length(port_edges) > 0
        port_indices = [(findfirst(isequal(p[1]), c.vertices),
                         findfirst(isequal(p[2]), c.vertices)) for p in port_edges]
        return hcat([coord_matrix[:, i1] - coord_matrix[:, i2]
            for (i1, i2) in port_indices]...)
    else
        return spzeros(size(coord_matrix, 1), 0)
    end
end

"""
    circuit_to_pso_matrix(m::AbstractMatrix)
Convert one of the inverse inductance, conductance, or capacitance matrices
characterizing a `Circuit` into the form needed for a `PSOModel`.
"""
circuit_to_pso_matrix(m::AbstractMatrix) =
    (-m + spdiagm(0=>sum(m, dims=2)[:,1]))[2:end, 2:end]

"""
    PSOModel(circuit::Circuit, port_edges::Vector{<:Tuple},
        port_names::AbstractVector)

Create a PSOModel for a circuit with ports on given edges. The spanning tree
where the first vertex is chosen as ground and all other vertices are neighbors
of ground is used.
"""
function PSOModel(circuit::Circuit, port_edges::Vector{<:Tuple},
        port_names::AbstractVector)
    Y = [circuit_to_pso_matrix(m) for m in matrices(circuit)]
    P = port_matrix(circuit, port_edges)
    return PSOModel(Y..., P, port_names)
end

"""
    coordinate_transform(m::AbstractMatrix)
Construct a coordinate transformation matrix for a `PSOModel` from a matrix
given by e.g. `coordinate_matrix(c::Circuit{T}, tree::SpanningTree{T}) where T`.
"""
coordinate_transform(m::AbstractMatrix) = m[:,2:end] .- m[:,1]

"""
    coordinate_transform(coord_from::AbstractMatrix, coord_to::AbstractMatrix)
Construct a coordinate transformation matrix to go from a `PSOModel` constructed
with one spanning tree to a `PSOModel` constructed with another.
"""
function coordinate_transform(coord_from::AbstractMatrix, coord_to::AbstractMatrix)
    return coordinate_transform(coord_to) / coordinate_transform(coord_from)
end

"""
    PSOModel(circuit::Circuit, port_edges::Vector{<:Tuple},
        port_names::AbstractVector, tree::SpanningTree)

Create a PSOModel for a circuit with ports on given edges and using the given
spanning tree.
"""
function PSOModel(circuit::Circuit, port_edges::Vector{<:Tuple},
        port_names::AbstractVector, tree::SpanningTree)
    pso = PSOModel(circuit, port_edges, port_names)
    transform = transpose(coordinate_transform(coordinate_matrix(circuit, tree)))
    return apply_transform(pso, transform)
end


text:
#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################
# JuMP
# An algebraic modeling language for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################

using JuMP, GLPK, Test

"""
    example_sudoku(filepath)
A sudoku solver that uses a MIP to find solutions.
We have binary variables x[i, j, k] which, if = 1, say that cell (i, j) contains
the number k. The constraints are:
 1 - Each cell has one value only
 2 - Each row contains each number exactly once
 3 - Each column contains each number exactly once
 4 - Each 3x3 subgrid contains each number exactly once
We will take the initial grid as a CSV file at `filepath`, where 0s are blanks.
"""
function example_sudoku(filepath)
    initial_grid = zeros(Int, 9, 9)
    open(filepath, "r") do fp
        for row in 1:9
            line = readline(fp)
            initial_grid[row, :] .= parse.(Int, split(line, ","))
        end
    end

    model = Model(GLPK.Optimizer)

    @variable(model, x[1:9, 1:9, 1:9], Bin)

    @constraints(model, begin
        # Constraint 1 - Only one value appears in each cell
        cell[i in 1:9, j in 1:9], sum(x[i, j, :]) == 1
        # Constraint 2 - Each value appears in each row once only
        row[i in 1:9, k in 1:9], sum(x[i, :, k]) == 1
        # Constraint 3 - Each value appears in each column once only
        col[j in 1:9, k in 1:9], sum(x[:, j, k]) == 1
        # Constraint 4 - Each value appears in each 3x3 subgrid once only
        subgrid[i=1:3:7, j=1:3:7, val=1:9], sum(x[i:i + 2, j:j + 2, val]) == 1
    end)

    # Initial solution
    for row in 1:9, col in 1:9
        if initial_grid[row, col] != 0
            @constraint(model, x[row, col, initial_grid[row, col]] == 1)
        end
    end

    # Solve it
    JuMP.optimize!(model)

    term_status = JuMP.termination_status(model)
    primal_status = JuMP.primal_status(model)
    is_optimal = term_status == MOI.OPTIMAL

    # Check solution
    if is_optimal
        mip_solution = JuMP.value.(x)
        sol = zeros(Int, 9, 9)
        for row in 1:9, col in 1:9, val in 1:9
            if mip_solution[row, col, val] >= 0.9
                sol[row, col] = val
            end
        end
        return sol
    else
        error("The solver did not find an optimal solution.")
    end
end

function print_sudoku_solution(solution)
    println("Solution:")
    println("[-----------------------]")
    for row in 1:9
        print("[ ")
        for col in 1:9
            print(solution[row, col], " ")
            if col % 3 == 0 && col < 9
                print("| ")
            end
        end
        println("]")
        if row % 3 == 0
            println("[-----------------------]")
        end
    end
end

if length(ARGS) == 1
    sol = example_sudoku(ARGS[1])
    print_sudoku_solution(sol)
else
    solution = example_sudoku(joinpath(@__DIR__, "sudoku.csv"))
    @test solution == [
        3 1 7 9 5 8 2 6 4;
        4 6 9 3 2 7 8 1 5;
        8 2 5 1 6 4 7 9 3;
        2 7 1 6 4 5 3 8 9;
        9 3 8 7 1 2 5 4 6;
        5 4 6 8 3 9 1 7 2;
        1 8 4 2 9 3 6 5 7;
        6 9 2 5 7 1 4 3 8;
        7 5 3 4 8 6 9 2 1]
end


text:
@inline function metropolis_acc(kT::Real, dU::Real, sθa::Real, sθb::Real, ϵ::Real)
  return ϵ <= (exp(-dU / kT) * sθb / sθa);
end

function kawasaki_acc(kT::Real, dU::Real, ϵ::Real)
  boltz = exp(-dU / (2*kT));
  anti_boltz = exp(dU / (2*kT));
  return ( ϵ <= (boltz / (boltz + anti_boltz)) );
end

abstract type Acceptor end;

mutable struct Metropolis <: Acceptor
  logπ_prev::Float64;
  weight_function::WeightFunction;
end

logπ_chain(chain::EAPChain) = -chain.U / chain.kT + chain.Ω;

function logπ_chain(chain::EAPChain, log_wf::WeightFunction)
  return -chain.U / chain.kT + chain.Ω + log_wf(chain);
end

function Metropolis(chain::EAPChain, wf::WeightFunction)
  return Metropolis(logπ_chain(chain, wf), wf);
end

# alpha is a ratio of trial move probabilities (related to clustering)
function (metro::Metropolis)(chain::EAPChain, ϵ::Real; α::Real = 1.0)
  metro.weight_function(chain);
  logπ_curr = logπ_chain(chain, metro.weight_function) + log(α);
  if (logπ_curr >= metro.logπ_prev) || (ϵ < exp(logπ_curr - metro.logπ_prev))
    metro.logπ_prev = logπ_curr;
    return true;
  else
    return false;
  end
  return false;
end

max_stars_repo_name: grasingerm/polymer-stats

max_issues_repo_name: grasingerm/polymer-stats

max_forks_repo_name: grasingerm/polymer-stats


text:
module Javis

using Animations
import Cairo: CairoImageSurface, image
using FFMPEG
using Gtk
using GtkReactive
using Hungarian
using Images
import Interact
import Interact: @map, Widget, Widgets, @layout!, hbox, vbox # not exporting textbox & slider due to possible conflicts with Gtk & luxor
using LaTeXStrings
using LightXML
import Luxor
import Luxor: Point, @layer, translate, rotate, @imagematrix
using ProgressMeter
using Random
using Statistics
using VideoIO

const FRAMES_SYMBOL = [:same, :all]

abstract type AbstractAction end
abstract type AbstractObject end
abstract type AbstractTransition end

include("structs/Video.jl")
include("structs/Easing.jl")
include("structs/RFrames.jl")
include("structs/GFrames.jl")
include("structs/Frames.jl")
include("structs/Scale.jl")

"""
    Transformation

Defines a transformation which can be returned by an object to be accessible later.
This is further explained in the Javis tutorials.

# Fields
- `point::Point`: the translation part of the transformation
- `angle::Float64`: the angle component of the transformation (in radians)
- `scale::Tuple{Float64, Float64}`: the scaling component of the transformation
"""
struct Transformation
    point::Point
    angle::Float64
    scale::Scale
end

Transformation(p, a) = Transformation(p, a, 1.0)
Transformation(p, a, s::Float64) = Transformation(p, a, (s, s))
Transformation(p, a, s::Tuple{Float64,Float64}) = Transformation(p, a, Scale(s...))

include("structs/ObjectSetting.jl")
include("structs/Object.jl")
include("structs/Transitions.jl")
include("structs/Action.jl")
include("structs/LayerSetting.jl")
include("structs/LayerCache.jl")
include("structs/Layer.jl")

"""
    Line

A type to define a line by two points. Can be used i.e. in [`projection`](@ref)
We mean the mathematic definition of a continuous line and not a segment of a line.

# Fields
- `p1::Point`: start point
- `p2::Point`: second point to define the line
"""
struct Line
    p1::Point
    p2::Point
end

"""
    Base.:*(m::Array{Float64,2}, transformation::Transformation)

Convert the transformation to a matrix and multiplies m*trans_matrix.
Return a new Transformation
"""
function Base.:*(m::Array{Float64,2}, transformation::Transformation)
    p = transformation.point
    θ = transformation.angle
    s = transformation.scale
    trans_matrix = [
        s.x*cos(θ) -sin(θ) p.x
        sin(θ) s.y*cos(θ) p.y
        0 0 1
    ]
    res = m * trans_matrix
    return Transformation(Point(gettranslation(res)...), getrotation(res), getscale(res))
end

include("layers.jl")
include("util.jl")
include("scales.jl")
include("luxor_overrides.jl")
include("backgrounds.jl")
include("svg2luxor.jl")
include("morphs.jl")
include("action_animations.jl")
include("javis_viewer.jl")
include("latex.jl")
include("object_values.jl")

"""
    projection(p::Point, l::Line)

Return the projection of a point to a line.
"""
function projection(p::Point, l::Line)
    # move line to origin and describe it as a vector
    o = l.p1
    v = l.p2 - o
    # point also moved to origin
    x = p - o

    # scalar product <x,v>/<v,v>
    c = (x.x * v.x + x.y * v.y) / (v.x^2 + v.y^2)
    return c * v + o
end
"""
    centered_point(pos::Point, width::Int, height::Int)
Returns pre-centered points to be used in the place image functions
rather than using centered=true
https://github.com/JuliaGraphics/Luxor.jl/issues/155
# Returns
- `pt::Point`: the location of the center of a layer wrt global canvas
"""
function centered_point(pos::Point, width::Int, height::Int)
    Point(pos.x - width / 2, pos.y - height / 2)
end

"""
    preprocess_frames!(video::Video)

"""
function preprocess_frames!(video::Video)
    return preprocess_frames!([video.objects..., flatten(video.layers)...])
end

"""
    preprocess_frames!(objects::Vector{<:AbstractObject})

Computes the frames for each object(of both the main canvas and layers) and action based on the user defined frames that the
user can provide like [`RFrames`](@ref), [`GFrames`](@ref) and `:same`.

This function needs to be called before calling [`get_javis_frame`](@ref) as it computes
the actual frames for objects and actions.

# Returns
- `frames::Array{Int}` - list of all frames normally 1:...

# Warning
Shows a warning if some frames don't have a background.
"""
function preprocess_frames!(objects::Vector{<:AbstractObject})
    compute_frames!(objects)

    for (i, object) in enumerate(objects)
        compute_frames!(object.actions; parent = object, parent_counter = i)
    end

    # get all frames
    frames = Int[]
    for object in objects
        append!(frames, collect(get_frames(object)))
    end
    frames = unique(frames)
    if !(frames ⊆ CURRENT_VIDEO[1].background_frames)
        @warn(
            "Some of the frames don't have a background. In this case: $(setdiff(frames, CURRENT_VIDEO[1].background_frames)))"
        )
    end

    set_current_object(objects[1])
    return frames
end

"""
flatten(layers::Vector{AbstractObject})

Takes out all the objects from each layer and puts them into a single list.
This makes things easier for the [`preprocess_frames!`](@ref) method
# Returns
- `objects::Vector{AbstractObject}` - list of all objects in each layer
"""
function flatten(layers::Vector{AbstractObject})
    objects = AbstractObject[]
    for layer in layers
        flatten!(objects, layer)
    end
    return objects
end

function flatten!(objects::Array{AbstractObject}, l::Layer)
    for obj in l.layer_objects
        flatten!(objects, obj)
    end
end
# finally objects
flatten!(objects::Array{AbstractObject}, object::Object) = push!(objects, object)

"""
    render(
        video::Video;
        framerate=30,
        pathname="javis_GIBBERISH.gif",
        liveview=false,
        streamconfig::Union{StreamConfig, Nothing} = nothing,
        tempdirectory="",
        ffmpeg_loglevel="panic",
        rescale_factor=1.0,
        postprocess_frame=identity,
        postprocess_frames_flow=default_postprocess
    )

Renders all previously defined [`Object`](@ref) drawings to the user-defined `Video` as a gif or mp4.

# Arguments
- `video::Video`: The video which defines the dimensions of the output

# Keywords
- `framerate::Int`: The frame rate of the video
- `pathname::String`: The path for the rendered gif or mp4 (i.e `output.gif` or `output.mp4`)
    - **Default:** The animation is rendered as a gif with the `javis_` prefix and some gibberish afterwards
- `liveview::Bool`: Causes a live image viewer to appear to assist with animation development
- `streamconfig::Union{StreamConfig, Nothing}`: Contains livestream specific instructions, passed on to [`setup_stream`](@ref).
Streaming to Twitch or other platforms are not yet supported.
- `tempdirectory::String`: The folder where each frame is stored
    Defaults to a temporary directory when not set
- `ffmpeg_loglevel::String`:
    - Can be used if there are errors with ffmpeg. Defaults to panic:
    All other options are described here: https://ffmpeg.org/ffmpeg.html
- `rescale_factor::Float64` factor to which the frames should be rescaled for faster rendering
- `postprocess_frame::Function` function that is applied to the imagematrix of each frame after they have been computed
takes as argument `frame_image, frame, frames`, useful to apply a postprocessing e.g. blur to all or some
of the images. By default it is the identity and nothing happens.
- `postprocess_frames_flow::Function` function that is applied to the vector of the frames indices should return a new vector
where elements are a subset of the number of frames. Useful to reorder the frames, e.g. reverse the video with
`postprocess_frames_flow=reverse`. By default it is the identity and nothing happens.
"""
function render(
    video::Video;
    framerate = 30,
    pathname = "javis_$(randstring(7)).gif",
    liveview = false,
    streamconfig::Union{StreamConfig,Nothing} = nothing,
    tempdirectory = "",
    ffmpeg_loglevel = "panic",
    rescale_factor = 1.0,
    postprocess_frames_flow = identity,
    postprocess_frame = default_postprocess,
)
    layers = video.layers
    objects = video.objects
    frames = preprocess_frames!(video)

    if liveview
        if isdefined(Main, :IJulia) && Main.IJulia.inited
            return _jupyter_viewer(video, length(frames), objects, framerate)

        elseif isdefined(Main, :PlutoRunner)
            return _pluto_viewer(video, length(frames), objects)
        else
            _javis_viewer(video, length(frames), objects)
            return "Live Preview Started"
        end
    end

    path, ext = "", ""
    if !isempty(pathname)
        path, ext = splitext(pathname)
    end
    render_mp4 = ext == ".mp4"
    codec_props = (crf = 23, preset = "medium")
    video_io = nothing
    # if we render a gif and the user hasn't set a tempdirectory
    if !render_mp4 && isempty(tempdirectory)
        tempdirectory = mktempdir()
    end

    # Create a template useful in case of postprocessing
    frame_template = _convert_and_rescale_if_needed(
        get_javis_frame(video, objects, first(frames); layers = layers),
        rescale_factor,
    )

    # Check if postprocess_frames_flow is ill defined
    original_frames = copy(frames)
    frames = postprocess_frames_flow(frames)
    if !issubset(frames, original_frames)
        error(
            "postprocess_frames_flow should return a vector of frame indices contained in the original number of frames",
        )
    end

    # Memoization container for postprocessing
    frames_memory = Dict{Int,Matrix{RGB{N0f8}}}()

    # helps to check if the video is already being rendered in mp4 case
    started_encoding_video = false

    filecounter = 1
    @showprogress 1 "Rendering frames..." for frame in frames
        frame_image = _postprocess(
            video,
            objects,
            frame,
            layers = layers,
            frames_memory = frames_memory,
            postprocess_frame = postprocess_frame,
            frame_template = frame_template,
            filecounter = filecounter,
            frames = frames,
            rescale_factor = rescale_factor,
        )

        if !isempty(tempdirectory)
            Images.save("$(tempdirectory)/$(lpad(filecounter, 10, "0")).png", frame_image)
        end
        if render_mp4
            if !isa(frame_image, Matrix{RGB{N0f8}})
                frame_image = convert.(RGB{N0f8}, frame_image)
            end
            if !started_encoding_video
                started_encoding_video = true
                video_io = open_video_out(
                    pathname,
                    frame_image,
                    framerate = framerate,
                    encoder_options = codec_props,
                )
            end
            write(video_io, frame_image)
        end
        filecounter += 1
    end

    isempty(pathname) && return
    if ext == ".gif"
        # generate a colorpalette first so ffmpeg does not have to guess it
        ffmpeg_exe(`-loglevel $(ffmpeg_loglevel) -i $(tempdirectory)/%10d.png -vf
                    palettegen $(tempdirectory)/palette.png`)
        # then apply the palette to get better results
        ffmpeg_exe(
            `-loglevel $(ffmpeg_loglevel) -framerate $framerate -i $(tempdirectory)/%10d.png -i
               $(tempdirectory)/palette.png -lavfi paletteuse -y $pathname`,
        )
    elseif ext == ".mp4"
        close_video_out!(video_io)
    else
        @error "Currently, only gif and mp4 creation is supported. Not a $ext."
    end

    # check if livestream is used and livestream if that's the case
    _livestream(streamconfig, framerate, video.width, video.height, pathname)

    # clear all CURRENT_* constants to not accidentally use a previous video when creating a new one
    empty_CURRENT_constants()

    # even if liveview = false, show the rendered gif in the cell output
    if isdefined(Main, :IJulia) && Main.IJulia.inited
        display(MIME("text/html"), """<img src="$(pathname)">""")
    elseif isdefined(Main, :PlutoRunner)
        return PlutoViewer(pathname)
    end
    return pathname
end

"""
    render_objects(objects, video, frame; layer_frames=nothing)
Is called inside the [`get_javis_frame`](@ref) function and renders objects(both individual and ones belonging to a layer).

# Arguments
- `object::Object`: The object to be rendered
- `video::Video`: The video which defines the dimensions of the output
- `frame::Int`: The frame number to be rendered
- `layer_frames::UnitRange`: The frames of the layer to which the object belongs(`nothing` for independent objects)
"""
function render_objects(objects, video, frame; layer_frames = nothing)
    CURRENT_OBJECT[1] = objects[1]
    background_settings = ObjectSetting()
    origin()
    origin_matrix = cairotojuliamatrix(getmatrix())
    # this frame needs doing, see if each of the scenes defines it
    for object in objects
        # if object is not in global layer this sets the background_settings
        # from the parent background object
        update_object_settings!(object, background_settings)
        CURRENT_OBJECT[1] = object

        # checks if independent object is in the current frame
        # also checks the realtive frame of an object in a layer
        # if none is true then the object doesn't exist for that frame at all
        if !(
            (layer_frames == nothing && frame in get_frames(object)) ||
            layer_frames isa Frames &&
            (frame - first(layer_frames.frames) + 1) in get_frames(object)
        )
            continue
        end

        # check if the object should be part of the global layer (i.e Background)
        # or in its own layer (default)
        in_global_layer = get(object.opts, :in_global_layer, false)::Bool
        in_local_layer = get(object.opts, :in_local_layer, false)::Bool
        if !in_global_layer && !in_local_layer
            @layer begin
                draw_object(object, video, frame, origin_matrix, layer_frames)
            end
        else
            draw_object(object, video, frame, origin_matrix, layer_frames)
            # update origin_matrix as it's inside the global layer
            origin_matrix = cairotojuliamatrix(getmatrix())
        end

        # if object is in global layer this changes the background settings
        update_background_settings!(background_settings, object)
    end
end

"""
    get_layer_frame(video, layer, frame)

Is called inside [`get_javis_frame`](@ref) and does two things viz.
    - Creates a Luxor Drawing and renders the object of each layer
    - computes the actions applies on the layer and stores them

Returns the Drawing of the layer as an image matrix.
"""
function get_layer_frame(video, layer, frame)
    Drawing(layer.width, layer.height, :image)
    layer_frames = layer.frames
    render_objects(layer.layer_objects, video, frame, layer_frames = layer_frames)

    if frame in get_frames(layer)
        # call currently active actions and their transformations for each layer
        actions = layer.actions
        for action in actions
            get_frames(action) isa Nothing &&
                error("Frame range for the layer's action might be missing")
            rel_frame = frame - first(layer_frames.frames) + 1
            if rel_frame in get_frames(action)
                action.func(video, layer, action, rel_frame)
            elseif rel_frame > last(get_frames(action)) && action.keep
                # call the action on the last frame i.e. disappeared things stay disappeared
                action.func(video, layer, action, last(get_frames(action)))
            end
        end
    end
    img_layer = image_as_matrix()
    finish()
    return img_layer
end

"""
    apply_layer_settings(layer_settings, pos)

Applies the computed actions to the image matrix of the layer to it's image matrix
Actions supported:
- `anim_translate`:translates the entire layer to a specified position
- `setopacity`:changes the opacity of the entire layer
- `anim_rotate`:rotates the layer by a given angle
- `appear(:fade)`:fades in the layer
- `disappear(:fade)`:fades out the layer

It reads and applies the layer settings(computed by [`get_layer_frame`](@ref) function))
"""
function apply_layer_settings(layer_settings, pos)
    # final actions on the layer are applied here
    # currently scale and translate are support

    # translate origin to the center of the layer, apply the settings and
    # translate back to the previous position
    translate(pos)
    scale(layer_settings.scale)

    # routine for anim_rotate_around
    rot_around_settings = layer_settings.misc
    if get(rot_around_settings, :rotate_around, false)
        translate(get(rot_around_settings, :translate, pos))
        rotate(get(rot_around_settings, :angle, pos))
        translate(get(rot_around_settings, :translate_back, pos))
    end

    rotate(layer_settings.rotation_angle)
    translate(-pos)
end

"""
    place_layers(video, layers, frame)

Places the layers on an empty drawing
It does 4 things:
- creates an empty Drawing of the same size as video
- calls the [`apply_layer_settings`](@ref)
- places every layer's image matrix on the drawing
- Repeats the above two steps if the [`show_layer_frame`](@ref) is defined for that layer(and frame)
    But fetches image matrix, position and settings from the layer cache

Returns the Drawing containing all the layers as an image matrix.
"""
function place_layers(video, layers, frame)
    # create an empty drawing of size same as the main video
    Drawing(video.width, video.height, :image)
    origin()

    for layer in layers
        CURRENT_LAYER[1] = layer
        if frame in get_frames(layer)
            pt = centered_point(layer.position, layer.width, layer.height)
            @layer begin
                # any actions on the layer go in this block

                layer_settings = layer.current_setting

                apply_layer_settings(layer_settings, layer.position)
                placeimage(layer.image_matrix, pt, alpha = layer.current_setting.opacity)
            end
        end

        lc = layer.layer_cache
        if frame in lc.frames
            if lc.frame_counter < length(lc.layer_frames)
                lc.frame_counter += 1
            else
                lc.frame_counter = 1
            end
            pt = centered_point(lc.position[lc.frame_counter], layer.width, layer.height)
            @layer begin
                settings = lc.settings_cache[lc.frame_counter]
                apply_layer_settings(settings, lc.position[lc.frame_counter])
                placeimage(lc.matrix_cache[lc.frame_counter], pt, alpha = settings.opacity)
            end
        end
    end

    # matrix of a transparent drawing with all the layers
    img_layers = image_as_matrix()
    finish()
    return img_layers
end

"""
    get_javis_frame(video, objects, frame; layers = Layer[])

Is called inside the [`render`](@ref) function.
It is a 5-step process:
- for each layer fetch it's image matrix and store it into the layer's struct
- if the [`show_layer_frame`](@ref) method is defined for a layer, save the
position, image matrix and layer settings for that frame of the layer in a [`LayerCache`](@ref)
- place the layers on an empty drawing
- creates the main canvas and renders the independent objects
- places the drawing containing all the layers on the main drawing

Returens the final rendered frame
"""
function get_javis_frame(video, objects, frame; layers = Layer[])

    # check if any layers have been defined
    if !isempty(layers)
        starting_positions = Point[]
        # render each layer's objects and store the layer's Drawing as an image matrix
        for layer in layers
            push!(starting_positions, layer.position)
            if isempty(Javis.CURRENT_LAYER)
                push!(Javis.CURRENT_LAYER, layer)
            else
                Javis.CURRENT_LAYER[1] = layer
            end
            if frame in get_frames(layer)
                mat = get_layer_frame(video, layer, frame)
                layer.image_matrix = mat
            end

            # check if the layer's frame needs to be cached for vewing layer
            lc = layer.layer_cache
            if frame in lc.layer_frames
                push!(lc.position, deepcopy(layer.position))
                push!(lc.settings_cache, deepcopy(layer.current_setting))
                push!(lc.matrix_cache, mat)

            end
        end

        img_layers = place_layers(video, layers, frame)

        # resetting layer positions to originals
        for (position, layer) in zip(starting_positions, layers)
            layer.position = position
        end
    end

    empty!(CURRENT_LAYER)

    # now that the layers have been handled
    # finally render the independent objects on a separate/main drawing
    Drawing(video.width, video.height, :image)

    # sends over the independent objects for rendering
    render_objects(objects, video, frame)

    if !isempty(layers)
        # place the matrix containing all the layers over the global matrix
        placeimage(img_layers, Point(-video.width / 2, -video.height / 2))
    end
    img = image_as_matrix()
    finish()
    return img
end

"""
    draw_object(video, layer, frame, origin_matrix, layer_frames)
Is called inside the [`render`](@ref) and does everything handled for an `AbstractObject`.
It is a 4-step process:
- translate to the start position
- call the relevant actions
- call the object function
- save the result of the object if wanted inside `video.defs`
"""
function draw_object(object, video, frame, origin_matrix, layer_frames)
    # translate the object to it's starting position.
    # It's better to draw the object always at the origin and use `star_pos` to shift it
    translate(get_position(object.start_pos))

    # reset change keywords
    empty!(object.change_keywords)

    # first compute and perform the global transformations of this object
    # relative frame number for actions
    if layer_frames == nothing
        rel_frame = frame - first(get_frames(object)) + 1
    else
        # actions of objects in a layer
        # this is somewhat nested since object and action defined in a layer
        # both have their respective frame ranges that need to be calculated relatively
        rel_frame = frame - first(get_frames(object)) - first(layer_frames.frames) + 1
    end
    # call currently active actions and their transformations
    for action in object.actions
        if rel_frame in get_frames(action)
            action.func(video, object, action, rel_frame)
        elseif rel_frame > last(get_frames(action)) && action.keep
            # call the action on the last frame i.e. disappeared things stay disappeared
            action.func(video, object, action, last(get_frames(action)))
        end
    end

    # set the defaults for the frame like setline() and setopacity()
    # which can depend on the actions
    set_object_defaults!(object)

    # if the scale would be 0.0 `show_object` is set to false => don't show the object
    # (it wasn't actually scaled to 0 because it would break Cairo :D)
    cs = get_current_setting()
    !cs.show_object && return

    res = object.func(video, object, frame; collect(object.change_keywords)...)
    current_global_matrix = cairotojuliamatrix(getmatrix())
    # obtain current matrix without the initial matrix part
    current_matrix = inv(origin_matrix) * current_global_matrix

    # if a transformation let's save the global coordinates
    if res isa Point
        trans = current_matrix * Transformation(res, 0.0, 1.0)
        object.result[1] = trans
    elseif res isa Transformation
        trans = current_matrix * res
        object.result[1] = trans
    else # just save the result such that it can be used as one wishes
        object.result[1] = res
    end
end

"""
    set_object_defaults!(object)

Set the default object values
- line_width and calls `Luxor.setline`.
- opacity and calls `Luxor.opacity`.
- scale and calls `Luxor.scale`.
"""
function set_object_defaults!(object)
    cs = object.current_setting
    current_line_width = cs.line_width * cs.mul_line_width
    Luxor.setline(current_line_width)
    current_opacity = cs.opacity * cs.mul_opacity
    Luxor.setopacity(current_opacity)

    desired_scale = cs.desired_scale * cs.mul_scale
    scaleto(desired_scale)
end

const LUXOR_DONT_EXPORT = [
    :boundingbox,
    :Boxmaptile,
    :Sequence,
    :setline,
    :setopacity,
    :fontsize,
    :get_fontsize,
    :scale,
    :text,
    :background,
]

# Export each function from Luxor
for func in names(Luxor; imported = true)
    if !(func in LUXOR_DONT_EXPORT)
        eval(Meta.parse("import Luxor." * string(func)))
        eval(Expr(:export, func))
    end
end

# shorthands declarations
include("shorthands/JLine.jl")
include("shorthands/JCircle.jl")
include("shorthands/JRect.jl")
include("shorthands/JBox.jl")
include("shorthands/JEllipse.jl")
include("shorthands/JStar.jl")
include("shorthands/JPoly.jl")
include("shorthands/JShape.jl")

export render, latex
export Video, Object, Background, Action, RFrames, GFrames
export @JLayer, background
export Line, Transformation
export val, pos, ang, scl, get_value, get_position, get_angle, get_scale
export projection, morph_to
export appear, disappear, rotate_around, follow_path, change
export rev
export scaleto
export act!
export anim_translate, anim_rotate, anim_rotate_around, anim_scale

export @Frames, prev_start, prev_end, startof, endof
export JBox, JCircle, JEllipse, JLine, JPoly, JRect, JStar, @JShape

# custom override of luxor extensions
export setline, setopacity, fontsize, get_fontsize, scale, text
export setup_stream, cancel_stream

# scales
export scale_linear, @scale_layer

end


text:
include(Knet.dir("data", "mnist.jl"))
using ArgParse
using Dates

atype = gpu() >= 0 ? KnetArray{Float64} : Array{Float64}

function train(args="")
    s = ArgParseSettings()
    s.autofix_names = true
    s.description = "RAM implementation in Knet."

    @add_arg_table s begin
        ("--name"; help="experiment name"; default=nothing; required=true)
        ("--savedir"; help="checkpoint save directory"; default=SAVEDIR)
        ("--epochs"; help="# of epochs"; arg_type=Int; default=2000)
        ("--batchsize"; help="batch size"; arg_type=Int; default=256)
        ("--seed"; help="random seed"; arg_type=Int; default=-1)
        ("--atype"; help="array type"; default=string(atype))
        ("--lr"; help="learning rate"; arg_type=Float64; default=1e-4)
        ("--gclip"; help="gradient clip"; arg_type=Float64; default=0.1)
        ("--optim"; help="optimizer"; default="rmsprop")
        ("--loadfile"; help="load file"; default=nothing)
        ("--init"; help="init function"; default="randinit")
        ("--sigma"; help="std dev"; arg_type=Float64; default=.22)
        ("--patchsize"; help="patch size"; arg_type=Int; default=8)
        ("--num-patches"; help="# of patches"; arg_type=Int; default=1)
        ("--glimpse-scale"; help="glimpse scale"; arg_type=Float64; default=1.)
        ("--num-channels"; help="# of channels"; arg_type=Int; default=1)
        ("--num-glimpses"; help="# of glimpses"; arg_type=Int; default=6)
        ("--loc-hidden"; help="# of units in location net"; arg_type=Int;
         default=128)
        ("--glimpse-hidden"; help="# of units in glimpse net"; arg_type=Int;
         default=128)
        ("--hiddensize"; help="# of units in contoller rnn"; arg_type=Int;
         default=256)
        ("--rnntype"; help="type of controller (relu|lstm)"; default=":relu")
        ("--num-classes"; help="# of classes"; arg_type=Int; default=10)
    end

    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true)
    @info "Options parsed [$(now())]"
    println(o); flush(stdout)
    o[:atype] = eval(Meta.parse(o[:atype]))
    o[:seed] > 0 && Knet.seed!(o[:seed])
    optim = eval(Meta.parse(o[:optim]))
    o[:init] = eval(Meta.parse(o[:init]))
    o[:rnntype] = eval(Meta.parse(o[:rnntype]))
    o[:epochs] == -1 && return o

    xtrn, ytrn, xtst, ytst = mnist()

    ram = Net(
        o[:patchsize],
        o[:num_patches],
        o[:glimpse_scale],
        o[:num_channels],
        o[:loc_hidden],
        o[:glimpse_hidden],
        o[:sigma],
        o[:hiddensize],
        o[:num_classes],
        o[:num_glimpses]; rnnType=o[:rnntype], atype=o[:atype], init=o[:init])

    dtrn = minibatch(xtrn, ytrn, o[:batchsize]; xtype=o[:atype])
    dtst = minibatch(xtst, ytst, o[:batchsize]; xtype=o[:atype])

    bestmodel_path = abspath(joinpath(o[:savedir], o[:name]*"-best.jld2"))
    lastmodel_path = abspath(joinpath(o[:savedir], o[:name]*"-last.jld2"))
    history = []
    bestacc = 0.0
    o[:loadfile] != nothing && Knet.@load o[:loadfile] ram history
    loss(x, ygold) = sum(ram(x, ygold)[1:3])
    for epoch = 1:o[:epochs]
        progress!(optim(loss, dtrn; lr=o[:lr], gclip=o[:gclip]))
        trn_losses, trn_acc = validate(ram, dtrn)
        tst_losses, tst_acc = validate(ram, dtst)
        println(
            "epoch=$(1+length(history)) ",
            "trnloss=$(trn_losses), trnacc=$trn_acc, ",
            "tstloss=$(tst_losses), tstacc=$tst_acc")
        push!(history, ([trn_losses..., trn_acc, tst_losses..., tst_acc]))
        Knet.@save lastmodel_path ram history o

        if tst_acc > bestacc
            bestacc = tst_acc
            Knet.@save bestmodel_path ram history o
        end
    end
end

max_stars_repo_name: ilkerkesen/RAM

max_issues_repo_name: ilkerkesen/RAM

max_forks_repo_name: ilkerkesen/RAM


text:
using Py2ToPy3

filenames = ["ex.py"]
change(filenames)
#change(filenames, overwrite=true)


text:
# export
#     AbstractModel,
#     RigidBody,
#     dynamics,
#     discrete_dynamics,
#     jacobian!,
#     discrete_jacobian!,
# 	orientation,
# 	state_dim,
# 	control_dim,
# 	state_diff_size
#
# export
#     QuadratureRule,
# 	RK2,
#     RK3,
# 	RK4,
#     HermiteSimpson

"""
 	AbstractModel

Abstraction of a model of a dynamical system of the form ẋ = f(x,u), where x is the n-dimensional state vector
and u is the m-dimensional control vector.

Any inherited type must define the following interface:
ẋ = dynamics(model, x, u)
n,m = size(model)
"""
abstract type AbstractModel end

"""
	LieGroupModel <: AbstractModel

Abstraction of a dynamical system whose state contains at least one arbitrary rotation.
"""
abstract type LieGroupModel <: AbstractModel end

"""
	RigidBody{R<:Rotation} <: LieGroupModel

Abstraction of a dynamical system with free-body dynamics, with a 12 or 13-dimensional state
vector: `[p; q; v; ω]`
where `p` is the 3D position, `q` is the 3 or 4-dimension attitude representation, `v` is the
3D linear velocity, and `ω` is the 3D angular velocity.

# Interface
Any single-body system can leverage the `RigidBody` type by inheriting from it and defining the
following interface:
```julia
forces(::MyRigidBody, x, u)  # return the forces in the world frame
moments(::MyRigidBody, x, u) # return the moments in the body frame
inertia(::MyRigidBody, x, u) # return the 3x3 inertia matrix
mass(::MyRigidBody, x, u)  # return the mass as a real scalar
```

Instead of defining `forces` and `moments` you can define the higher-level `wrenches` function
	wrenches(model::MyRigidbody, z::AbstractKnotPoint)
	wrenches(model::MyRigidbody, x, u)

# Rotation Parameterization
A `RigidBody` model must specify the rotational representation being used. Any `Rotations.Rotation{3}`
can be used, but we suggest one of the following:
* `UnitQuaternion`
* `MRP`
* `RodriguesParam`
"""
abstract type RigidBody{R<:Rotation} <: LieGroupModel end

"Integration rule for approximating the continuous integrals for the equations of motion"
abstract type QuadratureRule end

"Specifier for continuous systems (i.e. no integration)"
abstract type Continuous <: QuadratureRule end

"Integration rules of the form x′ = f(x,u), where x′ is the next state"
abstract type Explicit <: QuadratureRule end

"Integration rules of the form x′ = f(x,u,x′,u′), where x′,u′ are the states and controls at the next time step."
abstract type Implicit <: QuadratureRule end

"Fourth-order Runge-Kutta method with zero-order-old on the controls"
abstract type RK4 <: Explicit end

"Second-order Runge-Kutta method with zero-order-old on the controls"
abstract type RK3 <: Explicit end

"Second-order Runge-Kutta method with zero-order-old on the controls (i.e. midpoint)"
abstract type RK2 <: Explicit end

abstract type Euler <: Explicit end

"Third-order Runge-Kutta method with first-order-hold on the controls"
abstract type HermiteSimpson <: Implicit end

"Default quadrature rule"
const DEFAULT_Q = RK3

abstract type DiffMethod end
struct ForwardAD <: DiffMethod end
struct FiniteDifference <: DiffMethod end
diffmethod(::AbstractModel) = ForwardAD()  # set default to ForwardDiff
gen_cache(model::AbstractModel) = gen_cache(diffmethod(model), model)
gen_cache(::ForwardAD, ::AbstractModel) = nothing
gen_cache(::FiniteDifference, model::AbstractModel) = FiniteDiff.JacobianCache(model)

gen_grad_cache(model::AbstractModel) = gen_grad_cache(diffmethod(model), model)
gen_grad_cache(::ForwardAD, model::AbstractModel) = nothing
gen_grad_cache(::FiniteDifference, model::AbstractModel) =
    FiniteDiff.GradientCache(model)

function FiniteDiff.JacobianCache(model::AbstractModel,
        fdtype::Union{Val{T1},Type{T1}} = Val(:forward),
        dtype::Type{T2} = Float64; colored::Bool=false,
        sparsity = detect_sparsity(DEFAULT_Q, model),
        kwargs...) where {T1,T2}
    n,m = size(model)
    if colored
        colorvec = matrix_colors(sparsity)
    else
        colorvec = 1:n+m
    end
    FiniteDiff.JacobianCache(zeros(T2,n+m), zeros(T2,n), fdtype, dtype;
        colorvec=colorvec, kwargs...)
end

function FiniteDiff.GradientCache(model::AbstractModel,
    fdtype = Val(:forward))
    FiniteDiff.GradientCache(zeros(state_dim(model)), zeros(sum(size(model))), fdtype)
end

"""
    detect_sparsity(Q, model, [samples])

Create a sparse matrix representing the nonzero elements of the discrete dynamics
Jacobian. Uses ForwardDiff to compute the Jacobian on `samples` randomly-sampled
states and controls.
"""
function detect_sparsity(::Type{Q}, model::AbstractModel, samples=10) where Q
    n,m = size(model)
    ∇f = spzeros(n,n+m)
    x,u = rand(model)
    z = StaticKnotPoint(x,u,0.1)
    for i = 1:samples  # try several inputs to get the sparsity pattern
        _discrete_jacobian!(ForwardAD(), Q, ∇f, model, z, nothing)
    end
    return ∇f .!= 0
end

#=
Convenient methods for creating state and control vectors directly from the model
=#
for method in [:rand, :zeros, :ones]
    @eval begin
        function Base.$(method)(model::AbstractModel)
            n,m = size(model)
            x = @SVector $(method)(n)
            u = @SVector $(method)(m)
            return x, u
        end
        function Base.$(method)(::Type{T}, model::AbstractModel) where T
            n,m = size(model)
            x = @SVector $(method)(T,n)
            u = @SVector $(method)(T,m)
            return x,u
        end
    end
end
function Base.fill(model::AbstractModel, val::Real)
    n,m = size(model)
    x = @SVector fill(val,n)
    u = @SVector fill(val,m)
    return x, u
end

@inline control_dim(::AbstractModel) = throw(ErrorException("control_dim not implemented"))
@inline state_dim(::AbstractModel) = throw(ErrorException("state_dim not implemented"))

"""Default size method for model (assumes model has fields n and m)"""
@inline Base.size(model::AbstractModel) = state_dim(model), control_dim(model)

############################################################################################
#                               CONTINUOUS TIME METHODS                                    #
############################################################################################
"""
	ẋ = dynamics(model, z::AbstractKnotPoint)
	ẋ = dynamics(model, x, u, [t=0])

Compute the continuous dynamics of a forced dynamical given the states `x`, controls `u` and
time `t` (optional).
"""
@inline dynamics(model::AbstractModel, z::AbstractKnotPoint) = dynamics(model, state(z), control(z), z.t)

# Default to not passing in t
@inline dynamics(model::AbstractModel, x, u, t) = dynamics(model, x, u)

"""
	jacobian!(∇f, model, z::AbstractKnotPoint, [cache])

Compute the `n × (n + m)` Jacobian `∇f` of the continuous-time dynamics.
Only accepts an `AbstractKnotPoint` as input in order to avoid potential allocations
associated with concatenation.

This method can use either ForwardDiff or FiniteDiff, based on the result of
`RobotDynamics.diffmethod(model)`. When using FiniteDiff, the cache should be
passed in for best performance. The cache can be generated using either of the
following:

    RobotDynamics.gen_cache(model)
    FiniteDiff.JacobianCache(model)

"""
function jacobian!(∇f::AbstractMatrix, model::AbstractModel, z::AbstractKnotPoint,
        cache=gen_cache(model))
    _jacobian!(diffmethod(model), ∇f, model, z, cache)
end

function _jacobian!(::ForwardAD, ∇f::AbstractMatrix,
        model::AbstractModel, z::AbstractKnotPoint, cache=gen_cache(model))
    ix, iu = z._x, z._u
	t = z.t
    f_aug(z) = dynamics(model, z[ix], z[iu], t)
    s = z.z
	ForwardDiff.jacobian!(get_data(∇f), f_aug, s)
end

function _jacobian!(::FiniteDifference, ∇f::AbstractMatrix, model::AbstractModel,
        z::AbstractKnotPoint, cache=FiniteDiff.JacobianCache(model))
    ix,iu,t = z._x, z._u, z.t
    f_aug!(ẋ,z) = copyto!(ẋ, dynamics(model, z[ix], z[iu]))
    cache.x1 .= z.z
    FiniteDiff.finite_difference_jacobian!(∇f, f_aug!, cache.x1, cache)
end

"""
    jvp!(grad, model, z, λ, [cache])

Compute the Jacobian-transpose vector product, `∇f'λ`. Can use either ForwardDiff or
FiniteDiff.
"""
function jvp!(grad, model::AbstractModel, z::AbstractKnotPoint, λ, cache=gen_grad_cache(model))
    _jvp!(diffmethod(model), grad, model, z, λ, cache)
end

function _jvp!(::ForwardAD, grad, model::AbstractModel, z::AbstractKnotPoint, λ, cache=gen_grad_cache(model))
    ix, iu = z._x, z._u
	t = z.t
    f_aug(z) = dot(dynamics(model, z[ix], z[iu], t), λ)
    s = z.z
	ForwardDiff.gradient!(grad, f_aug, s)
end

function _jvp!(::FiniteDifference, grad, model::AbstractModel, z::AbstractKnotPoint, λ, cache=gen_grad_cache(model))
    ix,iu,t = z._x, z._u, z.t
    f_aug!(z) = dot(dynamics(model, z[ix], z[iu]), λ)
    cache.c3 .= z.z
    FiniteDiff.finite_difference_gradient!(grad, f_aug!, cache.c3, cache)
end

"""
    ∇jacobian!(∇²f, model, z, b)

Evaluate the Jacobian of the Jacobian-transpose vector product: `∇f'b`,
for the continuous dynamics.
The output `∇²f` is of size `(n+m,n+m)`, and `b` is of size `(n,)`.
This is needed, for example, when computing the Hessian of the Lagrangian when
computing a full Newton step.
"""
function ∇jacobian!(∇f::AbstractMatrix, model::AbstractModel, z::AbstractKnotPoint, b::AbstractVector)
    ix,iu = z._x, z._u
    t = z.t
    f_aug(z) = dot(dynamics(model, z[ix], z[iu], t), b)
    ForwardDiff.hessian!(∇f, f_aug, z.z)
    return nothing
end
DynamicsJacobian(model::AbstractModel) = DynamicsJacobian(state_dim(model), control_dim(model))

############################################################################################
#                          EXPLICIT DISCRETE TIME METHODS                                  #
############################################################################################

# Set default integrator
@inline discrete_dynamics(model::AbstractModel, z::AbstractKnotPoint) =
    discrete_dynamics(DEFAULT_Q, model, z)

"""
    x′ = discrete_dynamics(model, model, z)  # uses $(DEFAULT_Q) as the default integration scheme
    x′ = discrete_dynamics(Q, model, x, u, t, dt)
    x′ = discrete_dynamics(Q, model, z::KnotPoint)

Compute the discretized dynamics of `model` using explicit integration scheme `Q<:QuadratureRule`.

The default integration scheme is stored in `TrajectoryOptimization.DEFAULT_Q`
"""
@inline discrete_dynamics(::Type{Q}, model::AbstractModel, z::AbstractKnotPoint) where Q<:Explicit =
    discrete_dynamics(Q, model, state(z), control(z), z.t, z.dt)

@inline discrete_dynamics(::Type{Q}, model::AbstractModel, x, u, t, dt) where Q =
    integrate(Q, model, x, u, t, dt)

"""
	propagate_dynamics(::Type{Q}, model, z_, z)

Evaluate the discrete dynamics of `model` using integration method `Q` at knot point `z`,
storing the result in the states of knot point `z_`.

Useful for propagating dynamics along a trajectory of knot points.
"""
function propagate_dynamics(::Type{Q}, model::AbstractModel, z_::AbstractKnotPoint, z::AbstractKnotPoint) where Q<:Explicit
    x_next = discrete_dynamics(Q, model, z)
    set_state!(z_, x_next)
end

"""
	discrete_jacobian!(Q, ∇f, model, z::AbstractKnotPoint)

Compute the `n × (n+m)` discrete dynamics Jacobian `∇f` of `model` using explicit
integration scheme `Q<:QuadratureRule`.

This method can use either ForwardDiff or FiniteDiff, based on the result of
`RobotDynamics.diffmethod(model)`. When using FiniteDiff, the cache should be
passed in for best performance. The cache can be generated using either of the
following:

    RobotDynamics.gen_cache(model)
    FiniteDiff.JacobianCache(model)

When using FiniteDiff, the coloring vector for sparse Jacobians can be calculated
and stored in the cache. To compute this automatically, use

    FiniteDiff.JacobianCache(model, colored=true, [sparsity=sparsity])

where `sparsity` is a sparse matrix with the non-zero entries of the discrete Jacobian.
If left out, it will be computed using `detect_sparsity(DEFAULT_Q, model)`.
"""
function discrete_jacobian!(::Type{Q}, ∇f, model::AbstractModel,
        z::AbstractKnotPoint{T,N,M}, cache=gen_cache(model)) where {T,N,M,Q<:Explicit}
    _discrete_jacobian!(diffmethod(model), Q, ∇f, model, z, cache)
end

function _discrete_jacobian!(::ForwardAD, ::Type{Q}, ∇f, model::AbstractModel,
		z::AbstractKnotPoint{T,N,M}, cache=gen_cache(model)) where {T,N,M,Q<:Explicit}
    ix,iu,dt = z._x, z._u, z.dt
    t = z.t
    fd_aug(s) = discrete_dynamics(Q, model, s[ix], s[iu], t, dt)
    ∇f .= ForwardDiff.jacobian(fd_aug, SVector{N+M}(z.z))
	return nothing
end

function _discrete_jacobian!(::FiniteDifference, ::Type{Q}, ∇f, model::AbstractModel,
        z::AbstractKnotPoint{T,N,M}, cache=gen_cache(model)) where {T,N,M,Q<:Explicit}
    if isnothing(cache)
        cache = FiniteDiff.JacobianCache(model)
    end
    ix,iu,t,dt = z._x, z._u, z.t, z.dt
    fd_aug!(ẋ,z) = copyto!(ẋ, discrete_dynamics(Q, model, z[ix], z[iu], t, dt))
    cache.x1 .= z.z
    FiniteDiff.finite_difference_jacobian!(∇f, fd_aug!, cache.x1, cache)
    return nothing
end

"""
    discrete_jvp!(Q, grad, model, z, λ, [cache])

Calculated the discrete Jacobian-vector product, `∇f'λ`. Can use either ForwardDiff
or FiniteDiff. The cache for FiniteDiff can be generated using either of the following

    RobotDynamics.gen_grad_cache(model)
    FiniteDiff.GradientCache(model)
"""
function discrete_jvp!(::Type{Q}, grad, model::AbstractModel,
        z::AbstractKnotPoint{T,N,M}, λ, cache=gen_grad_cache(model)) where {T,N,M,Q<:Explicit}
    _djvp!(Q, diffmethod(model), grad, model, z, λ, cache)
end

function _djvp!(::Type{Q}, ::ForwardAD, grad, model::AbstractModel,
        z::AbstractKnotPoint{T,N,M}, λ, cache=gen_grad_cache(model)) where {T,N,M,Q<:Explicit}
    ix,iu,dt = z._x, z._u, z.dt
    t = z.t
    fd_aug(s) = dot(discrete_dynamics(Q, model, s[ix], s[iu], t, dt), λ)
    ForwardDiff.gradient!(grad, fd_aug, z.z)
end

function _djvp!(::Type{Q}, ::FiniteDifference, grad, model::AbstractModel,
        z::AbstractKnotPoint{T,N,M}, λ, cache=gen_grad_cache(model)) where {T,N,M,Q<:Explicit}
    if isnothing(cache)
        cache = FiniteDiff.GradientCache(model)
    end
    ix,iu,t,dt = z._x, z._u, z.t, z.dt
    fd_aug(z) = dot(discrete_dynamics(Q, model, z[ix], z[iu], t, dt), λ)
    cache.c3 .= z.z
    FiniteDiff.finite_difference_gradient!(grad, fd_aug, cache.c3, cache)
end

"""
    ∇discrete_jacobian!(Q, ∇²f, model, z, b)

Evaluate the Jacobian of the Jacobian-transpose vector product: `∇f'b`,
for the discrete dynamics.
The output `∇²f` is of size `(n+m,n+m)`, and `b` is of size `(n,)`.
This is needed, for example, when computing the Hessian of the Lagrangian when
computing a full Newton step.
"""
function ∇discrete_jacobian!(::Type{Q}, ∇f::AbstractMatrix, model::AbstractModel,
        z::AbstractKnotPoint{<:Any,n,m}, b::AbstractVector) where {Q<:Explicit,n,m}
    ix,iu = z._x, z._u
    t,dt = z.t, z.dt
    fd_aug(z) = dot(discrete_dynamics(Q, model, z[ix], z[iu], t, dt), b)
    ForwardDiff.hessian!(∇f, fd_aug, z.z)
    # jac_z(z) = ForwardDiff.jacobian(fd_aug, z)
    # ForwardDiff.jacobian(z->jac_z(z), z.z)
    return nothing
end

############################################################################################
#                               STATE DIFFERENTIALS                                        #
############################################################################################

state_diff(model::AbstractModel, x, x0) = x - x0
@inline state_diff_jacobian(model::AbstractModel, x::SVector{N,T}) where {N,T} = I
@inline state_diff_size(model::AbstractModel) = size(model)[1]

@inline state_diff_jacobian!(G, model::AbstractModel, z::AbstractKnotPoint) =
	state_diff_jacobian!(G, model, state(z))

function state_diff_jacobian!(G, model::AbstractModel, x::StaticVector)
	for i in 1:length(x)
		G[i,i] = 1
	end
end

# function ∇²differential!(G, model::AbstractModel, x::StaticVector, dx::AbstractVector)
# 	G .= ∇²differential(model, x, dx)
# end


text:
using PyCall
@pyimport pygtsam
@pyimport numpy as np

noise=pygtsam.Diagonal[:Sigmas](np.array([0.2, 0.2, 0.1]))

x = pygtsam.Pose2(3,2,1)

f = pygtsam.PriorFactorPose2(0, x, noise)

fg = pygtsam.NonlinearFactorGraph()

fg[:add](f)
s = fg[:serialize]()
fg2 = pygtsam.deserializeGraph(s)
fg[:printf]("Factor graph before serialization")
println("-------------------")
fg2[:printf]("After serialization")

println()
v = pygtsam.Values()

v[:insert](0, x)

s = v[:serialize]()
v2 = pygtsam.deserializeValues(s)
println("Values before serialization")
v[:printf]()
println("-------------------")
println("Values after serialization")
v2[:printf]()


text:
@symbol_func function K_RP(cur_reactor::AbstractReactor)
  cur_epsilon = cur_reactor.epsilon

  cur_K = 5.6e-8

  cur_K *= cur_reactor.Z_eff

  cur_K /= cur_epsilon ^ 2

  cur_K /= cur_reactor.kappa_95

  cur_K /= ( 1 - 1.31 * sqrt(cur_epsilon) + 0.46 * cur_epsilon )

  cur_K
end


text:
# Construction of PV structure for University of Ottawa collaboration.
# See slides for schematic, slab structure is Si, Gap, InAsSbP, InAs (N cell),
# InAsSbP (P cell), InAs.
# All temperatures are in Kelvin, and all lengths are in microns.
# Absorption functions are current setup to calculate the number of photons
# absorbed as interband transitions.
# Alter code under the ***Optical models and absorption functions*** heading
# for alternative behavior.
####
# Number of division of the N-type InAs layer.
# Resolution of photon generation is thckInAs / divNCell.
# divNCell = 3
# Number of division of the P-type InAsSbP layer.
# divPCell = 3
# Temperature of the emitter.
# tEmit = 800.0
# Background temperature.
# tBck = 300.0

const e = 1.602*^(10,-19)
const eV = 1.60218*^(10,-19)
const hbar = 1.05457*^(10,-34) #m2kg/s
const hbar_eV = 6.5821*^(10,-16) #eV s
function uOttawaSlabs(tEmit::Float64, tBck::Float64, divNCell::Int, divPCell::Int)
	### Settings
	# Separation distances and thickness in uma
	distGap = 0.1
	# Total InAs slab.
	thckInAs = 0.5
	# Total InAsSbP slab.
	thckInAsSbP = 0.5
	# Quaternary thickness. thckQuat = 0.025
	# Protective layer thickness.
	thckProt = 0.015
	### Derived parameters
	## Layer numbers.
	resNCell = /(thckInAs, divNCell) #divNCell is the number
	resPCell = /(thckInAsSbP, divPCell)
	# Total number of layers.
	numLayers = 4 + divNCell + divPCell
	## Temperature list of the layers.
	tmpLst = fill(tBck, numLayers)
	# Set temperature of the emitter.
	tmpLst[1] = tEmit
	## Boundary locations.
	bdrLoc = Vector{Float64}(undef, numLayers - 1)
	bdrLoc[1] = 0.0
	bdrLoc[2] = distGap
	bdrLoc[3] = bdrLoc[2] + thckProt
	# Fill boundaries for N-type part of the cell.
	for ind = 1 : divNCell

		bdrLoc[3 + ind] = bdrLoc[2 + ind] + resNCell
	end
	# Fill boundaries for P-type part of the cell
	for ind = 1 : divPCell

		bdrLoc[3 + divNCell + ind] = bdrLoc[2 + divNCell + ind] + resPCell
	end
	## Optical models and absorption functions.
	# Silicon emitter.
	# Dopant densities, currently empty?  Francoeur set real part of p-doped Si to be same as real part of top layer. P7
	acpDpt = dptDsc(0.0, 0.0)
	dnrDpt = dptDsc(0.0, 0.0)
	# Calculate silicon model parameters, see siMod and ResponseModels.
	siModParamsE = prmMSi(tmpLst[1], dnrDpt, acpDpt)
	# Construct silicon response model.
	siRspE(enr) = siRsp(enr, siModParamsE)
	siAbsE(enr) = imag(siRspE(enr))
	# Gap response.
	gRsp(enr) = cstRsp(1.0 + im * ^(10.0, -7), enr)
	gAbs(enr) = 0.0
	# InAs N-type (PV Cell) doping concentration and temperature
	InAs_param =eps_InAs_struct(10.0^23,300.0)
	eps_nInAs(enr) = eps_InAsntype(enr,InAs_param)
	eps_nInAs_IBimag(enr) = imag(epsIB(enr*eV/hbar,InAs_param.N0,InAs_param.T,InAs_param.E0_T_InAs_value))
	# InAsSbP
	InAsSbPstructure = InAsSbP_struct(0.25,0.25)
	eps_InAsSbP(enr) = eps_InAsSbP_xy(enr,InAsSbPstructure)
	eps_InAsSbP_imag(enr) = eps_InAsSbP_imag_xy(enr,InAsSbPstructure)
	eps_gold_imag(enr) = imag(epsgold(enr))
	## Generate lists of optical responses and transfer factors.
	optRsp = []
	trfFacs = []  #absorption
	# Layers prior to N-type part of the cell.
	push!(optRsp, siRspE, gRsp,eps_InAsSbP ) #Si, air and protective InAsSbP
	push!(trfFacs, siAbsE, gAbs,eps_InAsSbP_imag)
	# N-type part of the PV cell. n-InAs
	for ind = 1 : divNCell

		push!(optRsp, eps_nInAs)
		push!(trfFacs, eps_nInAs_IBimag)
	end
	# P-type part of the PV cell. InAsSbP
	for ind = 1 : divPCell

		push!(optRsp, eps_InAsSbP)
		push!(trfFacs,eps_InAsSbP_imag)
	end
	# Optically thick backing, gold/p-InAs
	push!(optRsp, epsgold)
	push!(trfFacs,eps_gold_imag)
	## Set which layers transmission should be calculated for.
	lPairs = Array{Int64,2}(undef, 2, divNCell + divPCell)
	# N-type layer pairs.
	for ind = 1 : divNCell

		lPairs[1, ind] = 1
		lPairs[2, ind] = 3 + ind
	end
	# N-type layer pairs.
	for ind = 1:divPCell

		lPairs[1, ind + divNCell] = 1
		lPairs[2, ind + divNCell] = 3 + divNCell + ind
	end
	# Build layer description for heat transfer code.
	return (lyrDsc(bdrLoc, tmpLst, optRsp, trfFacs),lPairs)
end
precompile(uOttawaSlabs, (Float64, Float64, Int, Int))


text:
#solves Ax = b where A is a symmetric tridiagonal matrix
#a is the diagonal of A, c is the subdiagonal of A, b is the rhs,
#d is the diagonal of U, x is the solution
function serial_trid(a, b, c, d, x)
  n = size(a,1)
  d[1] = a[1]
  x[1] = b[1]
  for i = 2:n   # 3 reads and 2 writes
    l = c[i - 1] / d[i - 1]
    d[i] = a[i] - c[i - 1] * l
    x[i] = b[i] - l * x[i - 1]
  end

  x[n] = x[n] / d[n]
  for i = n - 1 : -1 : 1    # 3 reads and 1 write
    x[i] = (x[i] - c[i] * x[i + 1]) / d[i]
  end
end

#solves Ax = b in place, where A is a symmetric tridiagonal matrix
#a is the diagonal of A, overwritten with the diagonal of U
#c is the subdiagonal of A
#b is the rhs, overwritten with the solution
function serial_trid(a, b, c)
  n = size(a,1)
  for i = 2:n
    l = c[i - 1] / a[i - 1]
    a[i] -= c[i - 1] * l
    b[i] -= l * b[i - 1]
  end

  b[n] /= a[n]
  for i = n - 1 : -1 : 1
    b[i] = (b[i] - c[i] * b[i + 1]) / a[i]
  end
end

#compute A*x, where A is a tridiagonal matrix
#a is the diagonal of A
#b is the super diagonal of A
#c is the sub diagonal of A
#x is the column vector
function matrix_vector_product(a, b, c, x)
  n = size(x, 1)
  y = zeros(n)
  for i = 1 : n
    y[i] = a[i] * x[i]
  end
  for i = 2 : n
    y[i] += c[i - 1] * x[i - 1]
  end
  for i = 1 : n - 1
    y[i] += b[i] * x[i + 1]
  end
  y
end


text:
function save_noise!(W::NoiseWrapper)

end

function interpolate!(W::NoiseWrapper,u,p,t)
  W.source(u,p,t)
end

function interpolate!(out1,out2,W::NoiseWrapper,u,p,t)
  W.source(out1,out2,u,p,t)
end

function calculate_step!(W::NoiseWrapper,dt,u,p)
  if isinplace(W)
    W(W.dW,W.dZ,u,p,W.curt+dt)
    W.dW .-= W.curW
    if W.Z != nothing
      W.dZ .-= W.curZ
    end
  else
    new_W, new_Z = W(u,p,W.curt+dt)
    W.dW = new_W - W.curW
    if W.Z != nothing
      W.dZ = new_Z - W.curZ
    end
  end
  W.dt = dt
end

function accept_step!(W::NoiseWrapper,dt,u,p,setup_next=true)
  if isinplace(W)
    W.curW .+= W.dW
  else
    W.curW += W.dW
  end
  W.curt += W.dt
  push!(W.W,copy(W.curW))
  push!(W.t,copy(W.curt))
  if W.Z != nothing
    if isinplace(W)
      W.curZ .+= W.dZ
    else
      W.curZ += W.dZ
    end

    push!(W.Z,copy(W.curZ))
  end

  W.dt = dt #dtpropose
  if setup_next
    calculate_step!(W,dt,u,p)
  end
end

function reject_step!(W::NoiseWrapper,dtnew,u,p)
  calculate_step!(W::NoiseWrapper,dtnew,u,p)
end

function setup_next_step!(W::NoiseWrapper,u,p)
  calculate_step!(W,W.dt,u,p)
end


text:
function makeplot(datadir::String, field::String, df::DataFrame, sparsity::Int)
    if field != "time"
        p = plot(df, x = "nrange", y = field,
          Scale.y_continuous(minvalue=0.0, maxvalue=1.0))
    else
      p = plot(df, x = "nrange", y = field)
    end
    draw(PNG(joinpath(datadir,
                    "$(field)_sparsity$(sparsity).png"
            ), 3inch, 3inch), p)
    nothing
end


text:
using OrdinaryDiffEq, LinearAlgebra, Test

function dynamics!(dx,x,θ,t)
    zmin, zmax, M, g = θ
    z, dz, p = x

    dx[1] = dz
    dx[2] = -M*g -z + p
    dx[3] = p-w(t) # the last state (algebraic) p has to be equal to w

    if z <= zmin
        dx[2] = max(dx[2], 0.0)
    elseif z >= zmax
        dx[2] = min(dx[2], 0.0)
    end
end

w(t) = .1*(7.0+3.0sin(t))
zmin_cond(x,t,integrator) =  x[1]-zmin
zmax_cond(x,t,integrator) =  zmax-x[1]

function zmin_affect_neg!(integrator)
    integrator.u[1] = zmin
    integrator.u[2] = 0.0
end

function zmax_affect_neg!(integrator)
    integrator.u[1] = zmax
    integrator.u[2] = 0.0
end

cbs = CallbackSet(
    ContinuousCallback(zmin_cond,zmin_affect_neg!),
    ContinuousCallback(zmax_cond,zmax_affect_neg!)
)

tf = 20.0
tspan = (0.0,tf)

zmin = 1.0e-3
zmax = 2.0e-2
M = 0.05
g = 9.81
θ = zmin, zmax, M, g

x0 = [(zmin+zmax)/2,0.0,w(0.0)]

E = diagm([1.0,M,0.0])

f = ODEFunction(dynamics!,mass_matrix=E)
prob = ODEProblem(f,x0,tspan,θ)

sol1 = solve(prob,Rodas4(), callback = cbs, reltol=1e-6)
@test sol1(0.06692341688237893)[3] ≈ 0.72 atol=1e-2
sol1 = solve(prob,Rodas5(), callback = cbs, reltol=1e-6)
@test sol1(0.06692341688237893)[3] ≈ 0.72 atol=1e-2
sol1 = solve(prob,Rodas5P(), callback = cbs, reltol=1e-6)
@test sol1(0.06692341688237893)[3] ≈ 0.72 atol=1e-2

#=
sol1 = solve(prob,Rosenbrock23(),callback=cbs, reltol=1e-6)
@test sol1(1.0)[3] ≈ 0.95 atol=1e-2
=#