# Skew-normal simulation example for c-posterior for mixtures
module Suite
using Distributions
using PyPlot
using HDF5, JLD
draw() = (pause(0.001); get_current_fig_manager()[:window][:raise_]())

# Settings
ns = [20,100,500,2000,10000]  # sample sizes n to use
alphas = [Inf,100]  # robustification params alpha to use
nreps = 5  # number of times to run the simulation
n_samples = 10^5  # number of MCMC iterations
nburn = int(n_samples/10)  # number of iterations to discard as burn-in
from_file = false  # load previous results from file

# Model parameters and sampler code
include("setup.jl")

nns = length(ns)
for (i_a,alpha) in enumerate(alphas)

    if from_file
        k_posteriors = load("k_posteriors-alpha=$alpha.jld","k_posteriors")
    else
        k_posteriors = zeros(m,nreps,nns)
        for (i_n,n) in enumerate(ns)
            for rep in 1:nreps
                srand(n+rep) # Reset RNG
                x = [(rand()<p1 ? skewrnd(1,l1,s1,a1) : skewrnd(1,l2,s2,a2))[1]::Float64 for j=1:n] # Sample data

                # Run sampler
                zeta = (1/n)/(1/n + 1/alpha)
                p,theta,k_r,v_r,art,arv,m_r,l_r = sampler(x,n_samples,m,c,sigma,zeta)

                # Compute posterior on k
                edges,counts = hist(k_r[nburn+1:end],[0:m])
                k_posteriors[:,rep,i_n] = counts/(n_samples-nburn)
            end
        end
        save("k_posteriors-alpha=$alpha.jld","k_posteriors",k_posteriors)
    end

    # Prior on k: p(k) \propto Binomial(k|m,lambda/m)I(k>0).
    pk = pdf(Binomial(m,lambda/m),1:m); pk = pk/sum(pk)

    # Plot
    figure(i_a,figsize=(8,3.2)); clf(); hold(true)
    subplots_adjust(bottom=0.2)
    colors = "bgyrm"
    shapes = "ds^v*"
    plot(1:m,pk,"ko-",label="prior",markersize=8,linewidth=2)
    for (i_n,n) in enumerate(ns)
        plot(1:m,vec(mean(k_posteriors[:,:,i_n],2)),"$(colors[i_n])$(shapes[i_n])-",label="n = $n",markersize=8,linewidth=2)
    end
    xlim(0,m)
    xticks(1:m)
    if alpha<Inf; title("\$\\mathrm{Coarsened\\,posterior}\$",fontsize=17)
    else; title("\$\\mathrm{Standard\\,posterior}\$",fontsize=17)
    end
    xlabel("\$k \\,\\,\\mathrm{(\\#\\,of\\,components)}\$",fontsize=17)
    ylabel("\$\\Pi(k|\\mathrm{data})\$",fontsize=15)
    legend(loc="upper right",numpoints=1)
    draw()
    savefig("skewsuite-alpha=$alpha.png",dpi=150)

end

end # module


# ------------------------ Air ------------------------
"""
    γ_AIR = 1.4

Adiabatic index or ratio of specific heats (dry air at 20º C).
"""
const γ_AIR = 1.4

"""
    R_AIR = 287.05287  (J/(Kg·K))

Specific gas constant for dry air.
"""
const R_AIR = 287.05287

# ------------------------ Earth ------------------------
"""
    gD = 9.80665  (m/s²)

Down component of gravity acceleration at Earth surface at 45º geodetic latitude.

1. Stevens, B. L., Lewis, F. L., & Johnson, E. N. (2015). Aircraft control and simulation:
   dynamics, controls design, and autonomous systems. John Wiley & Sons. Equation (page 33)
"""
const gD = 9.80665  # m/s²

# ------------------------ UNIT CONVERSION ------------------------
# ANGLES
const RAD2DEG = 57.29578  # Radians (rad) to degrees (deg)
const DEG2RAD = 1 / RAD2DEG
# LENGTH
const FT2M = 0.3048  # Feet (ft) to meter (m)
const M2FT = 1 / FT2M
# TEMPERATURE
const KEL2RANK = 1.8  # Kelvin to Rankine
const RANK2KEL = 1 / KEL2RANK
# MASS
const KG2LB = 2.20462262185  # Kilogram (Kg) to pound (lb)
const LB2KG = 1 / KG2LB

const SLUG2KG = 14.5939029372  # Slug to Kilogram (kg)
const KG2SLUG = 1 / SLUG2KG
# PRESSURE
const PSF2PA = LB2KG * gD / FT2M^2  # Pounds per square foot (PSF) to Pascal (Pa)
const PA2PSF = 1 / PSF2PA

const SLUGFT32KGM3 = SLUG2KG / FT2M^3  # slug/ft³ to Kg/m³
const KGM32SLUGFT3 = 1 / SLUGFT32KGM3


doc"""
    validate(g, parent_, ij, search_key)
validate `parent_` info for bfs on graph `g` with key `search_key` based on Graph500 result
validation(http://graph500.org/?page_id=12#sec-8)
Return `1` for no error
"""
function validate(
  g::Graph{T},
  parent_::Vector{T},
  ij::Array{T,2},
  search_key::T
  ) where T<:Integer

    out = 1 # Default: no error.

    if parent_[search_key] != search_key
      # root must be the parent of itself
      out = 0
      return out
    end

    N = nv(g)
    slice = findall(x->(x > zero(T)), parent_)

    #Compute levels and check for cycles.
    level_ = zeros(T, length(parent_))
    level_[slice] .= 1
    P = parent_[slice]
    mask = (P .!= search_key)
    k = 0
    while any(mask)
      level_[slice[mask]] .+= 1
      P = parent_[P]
      mask = (P .!= search_key)
      k = k + 1
      if k > N
        #There must be a cycle in the tree.
        out = -3
        return out
      end
    end

    #Check that there are no edges with only one end in the tree.
    #This also checks the component condition.
    lij = level_[ij]
    neither_in = zeros(Bool, size(lij)[2])
    not_neither_in_or_both_in = zeros(Bool, size(lij)[2])
    for i in 1:size(lij)[2]
      if ij[1,i] != ij[2,i]
        neither_in[i] = (lij[1,i] == 0 && lij[2,i] == 0)
        not_neither_in_or_both_in[i] = !(neither_in[i] || (lij[1,i] > 0 && lij[2,i] > 0))
      else
        neither_in[i] = true
        not_neither_in_or_both_in[i] = false
      end
    end
    if any(not_neither_in_or_both_in)
      out = -4
      return out
    end

    #Validate the distances/levels.
    respects_tree_level = trues(1, size(ij)[2])
    not_neither_in_or_respects_tree_level = zeros(Bool,size(ij)[2])
    level_difference = one(T);
    @simd for i in 1:size(ij)[2]
        # To prevent Unsigned integer overflows
        if(lij[1,i] > lij[2,i])
            respects_tree_level[i] = (lij[1,i] .- lij[2,i]) <= level_difference
        else
            respects_tree_level[i] = (lij[2,i] .- lij[1,i]) <= level_difference
        end

        not_neither_in_or_respects_tree_level[i] = !(neither_in[i] || respects_tree_level[i])
    end
    if any(not_neither_in_or_respects_tree_level)
      out = -5
      return out
    end

    return out
  end


using SLEEFInline
using Test
using Printf

using Base.Math: significand_bits

isnzero(x::T) where {T <: AbstractFloat} = signbit(x)
ispzero(x::T) where {T <: AbstractFloat} = !signbit(x)

function cmpdenorm(x::Tx, y::Ty) where {Tx <: AbstractFloat, Ty <: AbstractFloat}
    sizeof(Tx) < sizeof(Ty) ? y = Tx(y) : x = Ty(x) # cast larger type to smaller type
    (isnan(x) && isnan(y)) && return true
    (isnan(x) || isnan(y)) && return false
    (isinf(x) != isinf(y)) && return false
    (x == Tx(Inf)  && y == Ty(Inf))  && return true
    (x == Tx(-Inf) && y == Ty(-Inf)) && return true
    if y == 0
        (ispzero(x) && ispzero(y)) && return true
        (isnzero(x) && isnzero(y)) && return true
        return false
    end
    (!isnan(x) && !isnan(y) && !isinf(x) && !isinf(y)) && return sign(x) == sign(y)
    return false
end

# the following compares the ulp between x and y.
# First it promotes them to the larger of the two types x,y
const infh(::Type{Float64}) = 1e300
const infh(::Type{Float32}) = 1e37
function countulp(T, x::AbstractFloat, y::AbstractFloat)
    X, Y = promote(x, y)
    x, y = T(X), T(Y) # Cast to smaller type
    (isnan(x) && isnan(y)) && return 0
    (isnan(x) || isnan(y)) && return 10000
    if isinf(x)
        (sign(x) == sign(y) && abs(y) > infh(T)) && return 0 # relaxed infinity handling
        return 10001
    end
    (x ==  Inf && y ==  Inf) && return 0
    (x == -Inf && y == -Inf) && return 0
    if y == 0
        x == 0 && return 0
        return 10002
    end
    if isfinite(x) && isfinite(y)
        return T(abs(X - Y) / ulp(y))
    end
    return 10003
end

const DENORMAL_MIN(::Type{Float64}) = 2.0^-1074
const DENORMAL_MIN(::Type{Float32}) = 2f0^-149

function ulp(x::T) where {T<:AbstractFloat}
    x = abs(x)
    x == T(0.0) && return DENORMAL_MIN(T)
    val, e = frexp(x)
    return max(ldexp(T(1.0), e - significand_bits(T) - 1), DENORMAL_MIN(T))
end

countulp(x::T, y::T) where {T <: AbstractFloat} = countulp(T, x, y)

# get rid off annoying warnings from overwritten function
macro nowarn(expr)
    quote
        _stderr = stderr
        tmp = tempname()
        stream = open(tmp, "w")
        redirect_stderr(stream)
        result = $(esc(expr))
        redirect_stderr(_stderr)
        close(stream)
        result
    end
end

# overide domain checking that base adheres to
using Base.MPFR: ROUNDING_MODE
for f in (:sin, :cos, :tan, :asin, :acos, :atan, :asinh, :acosh, :atanh, :log, :log10, :log2, :log1p)
    @eval begin
        import Base.$f
        @nowarn function ($f)(x::BigFloat)
            z = BigFloat()
            ccall($(string(:mpfr_, f), :libmpfr), Int32, (Ref{BigFloat}, Ref{BigFloat}, Int32), z, x, ROUNDING_MODE[])
            return z
        end
    end
end

strip_module_name(f::Function) = last(split(string(f), '.')) # strip module name from function f

# test the accuracy of a function where fun_table is a Dict mapping the function you want
# to test to a reference function
# xx is an array of values (which may be tuples for multiple arugment functions)
# tol is the acceptable tolerance to test against
function test_acc(T, fun_table, xx, tol; debug = false, tol_debug = 5)
    @testset "accuracy $(strip_module_name(xfun))" for (xfun, fun) in fun_table
        rmax = 0.0
        rmean = 0.0
        xmax = map(zero, first(xx))
        for x in xx
            q = xfun(x...)
            c = fun(map(BigFloat, x)...)
            u = countulp(T, q, c)
            rmax = max(rmax, u)
            xmax = rmax == u ? x : xmax
            rmean += u
            if debug && u > tol_debug
                @printf("%s = %.20g\n%s  = %.20g\nx = %.20g\nulp = %g\n", strip_module_name(xfun), q, strip_module_name(fun), T(c), x, ulp(T(c)))
            end
        end
        rmean = rmean / length(xx)

        t = @test trunc(rmax, digits=1) <= tol

        fmtxloc = isa(xmax, Tuple) ? string('(', join((@sprintf("%.5f", x) for x in xmax), ", "), ')') : @sprintf("%.5f", xmax)
        println(rpad(strip_module_name(xfun), 18, " "), ": max ", @sprintf("%f", rmax),
            rpad(" at x = " * fmtxloc, 40, " "),
            ": mean ", @sprintf("%f", rmean))
    end
end

function runtests()
    @testset "SLEEF" begin
        include("exceptional.jl")
        include("accuracy.jl")
    end
end

runtests()


@doc raw"""
    SubGradientProblem <: Problem

A structure to store information about a subgradient based optimization problem

# Fields
* `manifold` – a [Manifold](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold)
* `cost` – the function $F$ to be minimized
* `subgradient` – a function returning a subgradient $\partial F$ of $F$

# Constructor

    SubGradientProblem(M, f, ∂f)

Generate the [`Problem`] for a subgradient problem, i.e. a function `f` on the
manifold `M` and a function `∂f` that returns an element from the subdifferential
at a point.
"""
mutable struct SubGradientProblem{mT <: Manifold} <: Problem
    M::mT
    cost::Function
    subgradient::Function
end
"""
    get_subgradient(p,x)

Evaluate the (sub)gradient of a [`SubGradientProblem`](@ref)` p` at the point `x`.
"""
function get_subgradient(p::P, x) where {P <: SubGradientProblem{M} where M <: Manifold}
    return p.subgradient(x)
end
"""
    SubGradientMethodOptions <: Options
stories option values for a [`subgradient_method`](@ref) solver

# Fields
* `retraction` – the retration to use within
* `stepsize` – a [`Stepsize`](@ref)
* `stop` – a [`StoppingCriterion`](@ref)
* `x` – (initial or current) value the algorithm is at
* `xOptimal` – optimal value
* `∂` the current element from the possivle subgradients at `x` that is used
"""
mutable struct SubGradientMethodOptions{P,T} <: Options where {P,T}
    retract!::Function
    stepsize::Stepsize
    stop::StoppingCriterion
    x::P
    xOptimal::P
    ∂::T
    function SubGradientMethodOptions{P}(
        M::TM,
        x::P,
        sC::StoppingCriterion,
        s::Stepsize,
        retr!::Function=exp!
        ) where {TM <: Manifold, P}
        o = new{P,typeof(zero_tangent_vector(M,x))}();
        o.x = x;
        o.xOptimal = x;
        o.stepsize = s; o.retract! = retr!;
        o.stop = sC;
        return o
    end
end
function SubGradientMethodOptions(
    M::Manifold,
    x::P,
    sC::StoppingCriterion,
    s::Stepsize,
    retr!::Function=exp!
) where {P}
    return SubGradientMethodOptions{P}(M, x, sC, s, retr!)
end


using Plots
using ImageMagick
using FileIO
using Images

using DataStructures

import Plots: default
default(show=false, reuse=true)

ENV["GKSwstype"]="nul"

scale = 2
fnt = font("Computer Modern", 10*scale)
default(titlefont=fnt, guidefont=fnt, tickfont=fnt, legendfont=fnt)

function convert_to_float(a::RGB)
    a = float(a)
    return (a.r + a.g + a.b)/3.0
end

function maze_to_canvas(filename)
    img = load(filename)
    return convert_to_float.(img)
end

function inbounds(canvas_size, loc)
    if minimum(Tuple(loc)) < 1
        return false
    elseif loc[2] > canvas_size[2]
        return false
    elseif loc[1] > canvas_size[1]
        return false
    else
        return true
    end
end

function create_circle!(canvas, loc::CartesianIndex, radius, threshold,
                        old_val)
    for i = 1:size(canvas)[1]
        for j = 1:size(canvas)[2]
            if ((radius-threshold)^2
                < (i-loc[1])^2+(j-loc[2])^2
                < (radius+threshold)^2)
                canvas[i,j] = old_val
            end
        end
    end
end

function color!(canvas, loc::CartesianIndex, old_val, new_val)
    # bounds check
    if !inbounds(size(canvas), loc)
        return
    end

    if (old_val == new_val)
        return
    elseif (canvas[loc] != old_val)
        return
    else
        canvas[loc] = new_val
    end
end

function find_neighbors(canvas, loc::CartesianIndex, old_val, new_val)

    neighbors =  []
    possible_neighbors = [loc + CartesianIndex(0, 1),
                          loc + CartesianIndex(1, 0),
                          loc + CartesianIndex(0, -1),
                          loc + CartesianIndex(-1, 0)]
    for possible_neighbor in possible_neighbors
        if inbounds(size(canvas), possible_neighbor) &&
           canvas[possible_neighbor] == old_val
            push!(neighbors, possible_neighbor)
        end
    end

    return neighbors
end

function stack_fill!(canvas, loc::CartesianIndex, old_val, new_val;
                     output=false)
    if new_val == old_val
        return
    end

    s = Stack{CartesianIndex}()
    #s = CartesianIndex[]
    push!(s, loc)

    id = 0

    while length(s) > 0
        temp_loc = pop!(s)
        if canvas[temp_loc] == old_val
            color!(canvas, temp_loc, old_val, new_val)
            if output
                plt = heatmap(canvas, color=:coolwarm, ratio=:equal,
                              size=(1000,1000))
                savefig("out"*lpad(string(id[1]),5,string(0))*".png")
                id += 1
                println("Outputting: ", id)
            end
            possible_neighbors = find_neighbors(canvas, temp_loc,
                                                old_val, new_val)
            for neighbor in possible_neighbors
                push!(s,neighbor)
            end
        end

    end
end

function queue_fill!(canvas, loc::CartesianIndex, old_val, new_val;
                     output=false)
    if new_val == old_val
        return
    end

    q = Queue{CartesianIndex}()
    enqueue!(q, loc)
    color!(canvas, loc, old_val, new_val)

    id = 0

    while length(q) > 0
        temp_loc = dequeue!(q)
        if output
            plt = heatmap(canvas, color=:coolwarm, ratio=:equal,
                          size=(500,500))
            savefig("out"*lpad(string(id[1]),5,string(0))*".png")
            id += 1
            println("Outputting: ", id)
        end

        possible_neighbors = find_neighbors(canvas, temp_loc, old_val, new_val)
        for neighbor in possible_neighbors
            color!(canvas, neighbor, old_val, new_val)
            enqueue!(q,neighbor)
        end

    end
end

function maze_fill!(canvas, loc::CartesianIndex, old_val, new_val;
                     output=false, end_loc = CartesianIndex(-1,-1))
    if new_val == old_val
        return
    end

    q = Queue{CartesianIndex}()
    enqueue!(q, loc)
    color!(canvas, loc, old_val, new_val)

    id = 0

    stop_solve = false
    while !stop_solve
        temp_loc = dequeue!(q)
        if output
            plt = heatmap(canvas, ratio=:equal,
                          size=(500,500), legend=:bottomright)
            scatter!(plt, Tuple(loc), label="Start")
            scatter!(plt, Tuple(end_loc), label="Finish")
            savefig("out"*lpad(string(id[1]),5,string(0))*".png")
            id += 1
            println("Outputting: ", id)
        end

        possible_neighbors = find_neighbors(canvas, temp_loc, old_val, new_val)
        for neighbor in possible_neighbors
            color!(canvas, neighbor, old_val, new_val)
            enqueue!(q,neighbor)
        end

        if (length(q) == 0 || temp_loc == end_loc)
            stop_solve = true
        end

    end
end

function recursive_fill!(canvas, loc::CartesianIndex, old_val, new_val;
                         output=false, id=id)

    if (old_val == new_val)
        return
    elseif (canvas[loc] != old_val)
        return
    else
        canvas[loc] = new_val
    end

    if output
        plt = heatmap(canvas, color=:coolwarm, ratio=:equal, size=(1000,1000))
        savefig("out"*lpad(string(id[1]),5,string(0))*".png")
        id[1] += 1
        println("Outputting: ", id[1])
    end

    possible_neighbors = find_neighbors(canvas, loc, old_val, new_val)
    for possible_neighbor in possible_neighbors
        recursive_fill!(canvas, possible_neighbor, old_val, new_val;
                        output=output,id=id)
    end
end

function init_canvas()
    canvas = zeros(30,30)
    loc = CartesianIndex(15,15)
    create_circle!(canvas, loc, 12, 2, 0.75)
    return canvas
end

function create_grid(dims)
    tuple_array = Array{Tuple{Int, Int}}(undef, dims[1]*dims[2])
    grid = Array{Float64}(undef, dims)
    for i = 1:dims[1]
        for j = 1:dims[2]
            tuple_array[(i-1)*dims[2]+j] = (i,j)
            grid[i,j] = i+j
        end
    end

    lines = Array{Tuple{Int, Int}}(undef, 10,2)
    for i = 1:5
        lines[i,:]=[tuple_array[(i-1)*5+1], tuple_array[i*5]]
        lines[i+5,:]=[tuple_array[i], tuple_array[20+i]]
    end

    println(lines)

    plt = heatmap(grid, color=:coolwarm, size=(600,500), aspect_ratio=:equal,
                 xlims=(0.5,5.5), ylims=(0.5,5.5))
    for i = 1:size(lines)[1]
        plot!(plt, lines[i,:], label="", color=:black, linewidth=3)
    end
    #plot!(plt, lines, label="")
    scatter!(plt, tuple_array, label = "", markersize=15)

    savefig("grid.png")

end

function main()

    create_grid((5,5))
#=
    canvas = init_canvas()
    loc2 = CartesianIndex(15,15)
    #recursive_fill!(canvas, loc2, 0.0, 0.5; output=true, id = [0])
    #queue_fill!(canvas, loc2, 0.0, 0.5; output=true)
    #stack_fill!(canvas, loc2, 0.0, 0.5; output=true)

    plt = heatmap(canvas, color=:coolwarm, size=(600,500), aspect_ratio=:equal,
                  xlims=(1,30), ylims=(1,30))
    savefig("simple_circle.png")

    # For maze
    maze = maze_to_canvas("maze2.png")
    maze_start = CartesianIndex(50,50)
    maze_end = CartesianIndex(96,96)
    maze_fill!(maze, maze_start, 1.0, 0.5; end_loc = maze_end, output=true)
    heatmap(maze)
=#

end

main()


#=
All DataElements should have a size and a getindex method
=#
abstract AbstractDataElement

lengths(x) = [length(el) for el in x.data]
function spansizes(data::Vector)
  lens = [length(el) for el in data]
  n = length(lens)
  res = fill(1, n)
  acc = 1
  if n > 0
    res[n] = 1
    for i=(n-1):-1:1
      acc *= lens[i+1]
      res[i]=acc
    end
  end
  return res
end
immutable CartesianDataElement <: AbstractDataElement
  data::Vector{Vector}
  isvertical::Bool
  spansizes::Vector{Int}

  function CartesianDataElement(data, isvertical)
    spsz = spansizes(data)
    return new(data, isvertical, spsz)
  end
end

isvertical(x::CartesianDataElement) = x.isvertical
function size(x::CartesianDataElement)
  if isvertical(x)
    n = length(x.data)
    m = prod(lengths(x))
  else
    n = prod(lengths(x))
    m = length(x.data)
  end
  return n, m
end

function getindex(x::CartesianDataElement, i::Int, j::Int)
#   println("i, j: $i, $j")
  spsz = x.spansizes
  lens = lengths(x)
  if isvertical(x)
    idx = mod(div(j-1, spsz[i]), lens[i]) + 1
    res = x.data[i][idx]
  else
    idx = mod(div(i-1, spsz[j]), lens[j]) + 1
    res = x.data[j][idx]
  end
  return res
end

immutable ExpandDataElement <: AbstractDataElement
  data
  n::Int
  m::Int
end

size(x::ExpandDataElement) = (x.n, x.m)
function size(x::ExpandDataElement, dim::Int)
  if dim==1
    return x.n
  elseif dim==2
    return x.m
  else
    error("dim must be 1|2")
  end
end
getindex(x::ExpandDataElement, i::Int, j::Int) = x.data


import BlockArrays: BlockSizes, BlockIndex, globalrange, nblocks, global2blockindex, blockindex2global

#=
[1,1  1,2] | [1,3  1,4  1,5]
--------------------------
[2,1  2,2] | [2,3  2,4  2,5]
[3,1  3,2] | [3,3  3,4  3,5]
----------------------------
[4,1  4,2] | [4,3  4,4  4,5]
[5,1  5,2] | [5,3  5,4  5,5]
[6,1  6,2] | [6,3  6,4  6,5]
=#

block_size = BlockSizes([1,2,3], [2, 3])

@testset "BlockSizes / BlockIndices" begin
@test nblocks(block_size) == (3,2)
@test nblocks(block_size, 1) == 3
@test nblocks(block_size, 2) == 2

@test @inferred(globalrange(block_size, (1,1))) == (1:1, 1:2)
@test @inferred(globalrange(block_size, (1,2))) == (1:1, 3:5)
@test @inferred(globalrange(block_size, (2,1))) == (2:3, 1:2)
@test @inferred(globalrange(block_size, (2,2))) == (2:3, 3:5)

# Test for allocations inside a function to avoid noise due to global
# variable references
wrapped_allocations = (bs, i) -> @allocated(globalrange(bs, i))
@test wrapped_allocations(block_size, (1, 1)) == 0

@test @inferred(global2blockindex(block_size, (3, 1))) == BlockIndex((2,1), (2,1))
@test @inferred(global2blockindex(block_size, (1, 4))) == BlockIndex((1,2), (1,2))
@test @inferred(global2blockindex(block_size, (4, 5))) == BlockIndex((3,2), (1,3))

wrapped_allocations = (bs, i) -> @allocated(global2blockindex(bs, i))
@test wrapped_allocations(block_size, (3, 1)) == 0

@test @inferred(blockindex2global(block_size, BlockIndex((2,1), (2,1)))) == (3, 1)
@test @inferred(blockindex2global(block_size, BlockIndex((1,2), (1,2)))) == (1, 4)
@test @inferred(blockindex2global(block_size, BlockIndex((3,2), (1,3)))) == (4, 5)

wrapped_allocations = (bs, i) -> @allocated(blockindex2global(bs, i))
@test wrapped_allocations(block_size, BlockIndex((2,1), (2,1))) == 0

@test block_size == BlockSizes(1:3, 2:3)

buf = IOBuffer()
print(buf, block_size)
@test String(take!(buf)) == "[1, 2, 3] × [2, 3]"

end


# -*- coding: utf-8 -*-
using Petri
using ModelingToolkit
import ModelingToolkit: Constant, Variable

mutable struct SIRState{T,F}
    S::T
    I::T
    R::T
    β::F
    γ::F
    μ::F
end

mutable struct ParamSIR{T, P}
    S::T
    I::T
    R::T
    params::P
end

mutable struct ParamSEIR{T, P}
    S::T
    E::T
    I::T
    R::T
    params::P
end

mutable struct SEIRState{T,F}
    S::T
    I::T
    R::T
    β::F
    γ::F
    μ::F
    E::T
    η::F
end

mutable struct SEIRDState{T,F}
    S::T
    I::T
    R::T
    β::F
    γ::F
    μ::F
    E::T
    η::F
    D::T
    ψ::F
end

macro grounding(ex)
    return ()
end

# +
function main()
    @grounding begin
        S => Noun(Susceptible, ontology=Snowmed)
        I => Noun(Infectious, ontology=ICD9)
        R => Noun(Recovered, ontology=ICD9)
        λ₁ => Verb(infection)
        λ₂ => Verb(recovery)
        λ₃ => Verb(loss_of_immunity)
    end
    @variables S, I, R, β, γ, μ
    N = +(S,I,R)
    ϕ = [(S > 0) * (I > 0),
         I > 0,
         R > 0]

    Δ = [(S~S-1, I~I+1),
        (I~I-1, R~R+1),
        (R~R-1, S~S+1)]

    Λ = [β*S*I/N,
        γ*I,
        μ*R]

    m = Petri.Model([S,I,R], Δ, Λ, ϕ)
    p = Petri.Problem(m, SIRState(100, 1, 0, 0.5, 0.15, 0.05), 150)

    @grounding begin
        E => Noun(Exposed, ontology=ICD9)
        λ₄ => Verb(exposure)
    end
    @variables E, η
    N = +(S,E,I,R)
    ϕ = [(S > 0) * (I > 0),
         E > 0]

    Δ = [(S~S-1, E~E+1),
         (E~E-1, I~I+1)]

    Λ = [β*S*I/N,
        η*E]
    m2 = Petri.Model([S,E], Δ, Λ, ϕ)
    f = Dict(S => S)
    m′ = deepcopy(m)
    Petri.rewrite!(m′, m2, f)
    m
    p2 = Petri.Problem(m′, SEIRState(100, 1, 0, 0.5, 0.15, 0.05, 0, 0.12), 150)

    @grounding begin
        D => Noun(Dead, ontology=ICD9)
        λ₅ => Verb(death)
    end
    @variables D, ψ
    ϕ = [I > 0]

    Δ = [(I~I-1, D~D+1)]

    Λ = [ψ*I]

    m3 = Petri.Model([D], Δ, Λ, ϕ)
    m′′ = deepcopy(m′)
    Petri.rewrite!(m′′, m3)
    p3 = Petri.Problem(m′′, SEIRDState(100, 1, 0, 0.5, 0.15, 0.05, 0, 0.12, 0, 0.1), 150)

    return p, p2, p3

end
# -

p, p2, p3 = main()

# +
@show "SIR"

Petri.solve(p)
@time Petri.solve(p)

mf = Petri.eval(Petri.funckit(p))
pf = Petri.Problem(mf, SIRState(100, 1, 0, 0.5, 0.15, 0.05), 150)
Petri.solve(pf)
@time Petri.solve(pf)

# +
@show "SEIR"

Petri.solve(p2)
@time Petri.solve(p2)

mf2 = Petri.eval(Petri.funckit(p2))
pf2 = Petri.Problem(mf2, SEIRState(100, 1, 0, 0.5, 0.15, 0.05, 0, 0.12), 150)
Petri.solve(pf2)
@time Petri.solve(pf2)

# +
@show "SEIRD"

Petri.solve(p3)
@time Petri.solve(p3)

mf3 = Petri.eval(Petri.funckit(p3))
pf3 = Petri.Problem(mf3, SEIRDState(100, 1, 0, 0.5, 0.15, 0.05, 0, 0.12, 0, 0.1), 150)
Petri.solve(pf3)
@time Petri.solve(pf3)


### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 6e487246-5d73-11eb-2f43-1f159557c68c
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using ReinforcementLearning
	using Flux
	using Statistics
	using Plots
	using LinearAlgebra:dot
end

# ╔═╡ 41796f22-5d73-11eb-0365-f35b80ba9bbb
md"""
# Chapter 13 Short Corridor
"""

# ╔═╡ 8fa18536-5d73-11eb-398d-b93cbd63aedc
begin
	Base.@kwdef mutable struct ShortCorridorEnv <: AbstractEnv
		position::Int = 1
	end

	RLBase.state_space(env::ShortCorridorEnv) = Base.OneTo(4)
	RLBase.action_space(env::ShortCorridorEnv) = Base.OneTo(2)

	function (env::ShortCorridorEnv)(a)
		if env.position == 1 && a == 2
			env.position += 1
		elseif env.position == 2
			env.position += a == 1 ? 1 : -1
		elseif env.position == 3
			env.position += a == 1 ? -1 : 1
		end
		nothing
	end

	function RLBase.reset!(env::ShortCorridorEnv)
		env.position = 1
		nothing
	end

	RLBase.state(env::ShortCorridorEnv) = env.position
	RLBase.is_terminated(env::ShortCorridorEnv) = env.position == 4
	RLBase.reward(env::ShortCorridorEnv) =  env.position == 4 ? 0.0 : -1.0
end

# ╔═╡ 9a269140-5d73-11eb-01a6-a78c2b583887
world = ShortCorridorEnv()

# ╔═╡ aab4885a-5d73-11eb-33c7-b97e45de4016
ns, na = length(state_space(world)), length(action_space(world))

# ╔═╡ ae7c57c4-5d73-11eb-1734-f531d5b4cc6e
function run_once(A)
    avg_rewards = []
    for ϵ in A
        p = TabularRandomPolicy(;table=Dict(s => [1-ϵ, ϵ] for s in 1:ns))
        env = ShortCorridorEnv()
        hook=TotalRewardPerEpisode()
        run(p, env, StopAfterEpisode(1000),hook)
        push!(avg_rewards, mean(hook.rewards[end-100:end]))
    end
    avg_rewards
end

# ╔═╡ 2159d8de-5d74-11eb-0f9b-9d4fbb340ad6
X = 0.05:0.05:0.95

# ╔═╡ c2a474fe-5d73-11eb-0073-2f4afc5a1ada
plot(X, mean([run_once(X) for _ in 1:10]), legend=nothing)

# ╔═╡ d39bd212-5d73-11eb-0baf-b54fd44c2337
md"""
## REINFORCE Policy

Based on descriptions in Chapter 13.1, we need to define a new customized approximator.
"""

# ╔═╡ 33c22aaa-5d74-11eb-2271-e1c21336c28a
begin
	Base.@kwdef struct LinearPreferenceApproximator{F,O} <: AbstractApproximator
		weight::Vector{Float64}
		feature_func::F
		actions::Int
		opt::O
	end

	function (A::LinearPreferenceApproximator)(s)
		h = [dot(A.weight, A.feature_func(s, a)) for a in 1:A.actions]
		softmax(h)
	end

	function RLBase.update!(A::LinearPreferenceApproximator, correction::Pair)
		(s, a), Δ = correction
		w, x = A.weight, A.feature_func
		w̄ = -Δ .* (x(s,a) .- sum(A(s) .* [x(s, b) for b in 1:A.actions]))
		Flux.Optimise.update!(A.opt, w, w̄)
	end
end

# ╔═╡ 4d00b372-5d74-11eb-0760-11feb1439567
begin
	Base.@kwdef struct ReinforcePolicy{A<:AbstractApproximator} <: AbstractPolicy
		approximator::A
		γ::Float64
	end

	(p::ReinforcePolicy)(env::AbstractEnv) = prob(p, state(env)) |> WeightedExplorer(;is_normalized=true)

	RLBase.prob(p::ReinforcePolicy, s) = p.approximator(s)

	function RLBase.update!(
		p::ReinforcePolicy,
		t::AbstractTrajectory,
		::AbstractEnv,
		::PostEpisodeStage
	)
		S, A, R = t[:state], t[:action], t[:reward]
		Q, γ = p.approximator, p.γ
		G = 0.

		for i in 1:length(R)
			s,a,r = S[end-i], A[end-i], R[end-i+1]
			G = r + γ*G

			update!(Q, (s, a) => G)
		end
	end

	function RLBase.update!(
		t::AbstractTrajectory,
		::ReinforcePolicy,
		::AbstractEnv,
		::PreEpisodeStage
	)
		empty!(t)
	end
end

# ╔═╡ 5f0926f8-5d74-11eb-114c-77cfe8c57a11
function run_once_RL(α)
    agent = Agent(
        policy=ReinforcePolicy(
            approximator=LinearPreferenceApproximator(
                weight=[-1.47, 1.47],  # init_weight
                feature_func=(s,a) -> a == 1 ? [0, 1] : [1, 0],
                actions=na,
                opt=Descent(α)
            ),
            γ=1.0
        ),
        trajectory=VectorSARTTrajectory()
    )

    env = ShortCorridorEnv()
    hook = TotalRewardPerEpisode()
    run(agent,env,StopAfterEpisode(1000;is_show_progress=false),hook)
    hook.rewards
end

# ╔═╡ 6310c5c4-5d74-11eb-331a-8339e03cd161
begin
	fig_13_1 = plot(legend=:bottomright)
	for x in [-13, -14]  # for -12, it seems not that easy to converge in short time
		plot!(fig_13_1, mean(run_once_RL(2. ^ x) for _ in 1:50), label="alpha = 2^$x")
	end
	fig_13_1
end

# ╔═╡ 9bf6ff86-5d74-11eb-2235-1bada5c181d4
md"""
Interested in how to reproduce figure 13.2? A PR is warmly welcomed! See you there!
"""

# ╔═╡ Cell order:
# ╟─41796f22-5d73-11eb-0365-f35b80ba9bbb
# ╠═6e487246-5d73-11eb-2f43-1f159557c68c
# ╠═8fa18536-5d73-11eb-398d-b93cbd63aedc
# ╠═9a269140-5d73-11eb-01a6-a78c2b583887
# ╠═aab4885a-5d73-11eb-33c7-b97e45de4016
# ╠═ae7c57c4-5d73-11eb-1734-f531d5b4cc6e
# ╠═2159d8de-5d74-11eb-0f9b-9d4fbb340ad6
# ╠═c2a474fe-5d73-11eb-0073-2f4afc5a1ada
# ╟─d39bd212-5d73-11eb-0baf-b54fd44c2337
# ╠═33c22aaa-5d74-11eb-2271-e1c21336c28a
# ╠═4d00b372-5d74-11eb-0760-11feb1439567
# ╠═5f0926f8-5d74-11eb-114c-77cfe8c57a11
# ╠═6310c5c4-5d74-11eb-331a-8339e03cd161
# ╟─9bf6ff86-5d74-11eb-2235-1bada5c181d4


#=
Created on Wednesday 08 July 2020
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Example of dessert distribution to illustrate the concept.
=#

using OTSIN, Plots, StatsPlots, Printf
using LaTeXStrings

a = [3, 3, 3, 2, 2, 2]
b = [5, 4, 3, 3]

@assert sum(a) ≈ sum(b)

desserts = [:merveilleux, :eclair, Symbol("chocolate mousse"), Symbol("carrot cake")]
persons = [:Bernard, :Jan, :Willem, :Laura, :Laure, :Margot]

M = [2 2 1 0;
     0 -2 -2 2;
     1 2 2 -1;
     -1 2 1 -2;
     -1 -1 -2 1;
     -1 0 0 2]

Q = optimaltransport(M, a, b, λ=10.0)

λs = [10^i for i in -2:0.05:1]

utilities = Float64[]
entropies = Float64[]
for λ in λs
        Q = optimaltransport(M, a, b, λ=λ)
        Q ./= sum(Q)  # normalize
        push!(utilities, utility(M, Q))
        push!(entropies, entropy(Q))
end

groupedbar(String.(persons), Q, labels=reshape(desserts, 1, :), bar_position = :stack,
        title="Neutrally-driven")

anim = @animate for λ in λs
  Q = optimaltransport(M, a, b, λ=λ)
  pQ = groupedbar(String.(persons), Q, labels=reshape(desserts, 1, :),
                bar_position = :stack,
                title=@sprintf("Optimal coupling\n lambda=%.2f", λ), ylabel="portions")
  put = plot(λs, utilities, label="utility", lw=2)
  plot!(λs, entropies, label="entropy", lw=2)
  xaxis!(:log)
  xlabel!(L"$\lambda$")
  vline!([λ], label="", lw=2, color=:red, ls=:dash)
  plot(pQ, put, size=(800, 500))
end

gif(anim, "figures/desserts.gif", fps = 5)


# Combinations iterator with tabu
using Combinatorics
import Base: iterate, eltype, length

struct CombinationsWithout{T}
     a::T
     t::Int
     w::Int
     l::Int
     L::Int
 end

 function Base.iterate(c::CombinationsWithout)
     s = zeros(Int, c.t)
     if 0 < c.t <= c.l
         s[1] = 1
         if c.w == 1
             s[1] = 2
         end
         i = 1
         while i < c.t
             i += 1
             s[i] = s[i-1] + 1
             if s[i] == c.w
                 s[i] += 1
             end
         end
     end
     Base.iterate(c, (s, Vector{eltype(c.a)}(undef, c.t), 1))
 end

 function Base.iterate(c::CombinationsWithout, state)
     s, comb, k =  state
     k > c.L && return nothing
     k += 1
     for i in 1:c.t
         comb[i] = c.a[s[i]]
     end
     for i in length(s):-1:1
         s[i] += 1
         if s[i] > c.l - (length(s)-i) - (s[i] <= c.w)
             continue
         end
         if s[i] == c.w
             s[i] += 1
         end
         for j = i+1:length(s)
             s[j] = s[j-1]+1
             if s[j] == c.w
                 s[j] += 1
             end
         end
         break
     end
     comb, (s, comb, k)
 end

 length(c::CombinationsWithout) = c.L
 eltype(::Type{CombinationsWithout{T}}) where {T} = Vector{eltype(T)}

 """
     combinations_without(a, n::Integer, w)

 Generate all combinations of `n` elements from an indexable object except those with index `w`. Note that the combinations
 are modified inplace.
 """
 function combinations_without(a, t::Integer, w)
     l = length(a)
     ins = 1 <= w <= l
     L = binomial(l - ins, t)
     if ins && length(a) == w # w is last index
         CombinationsWithout(a, t, 0, l - 1, L)
     elseif ins && l == t # t too big
         CombinationsWithout(a, t, w, l, 0)
     elseif ins
         CombinationsWithout(a, t, w, l, L)
     else
         CombinationsWithout(a, t, 0, l, L)
     end
 end


function jump_max_area_diameter(n)
    if mod(n,2)==0 && n>=6
        s = (2*sqrt(114)-7)/22;
        t = (3521*sqrt(114)-34010)/9196;
        c = (17328*(663157+3161*pi^2) - (1088031703 - 3918085*pi^2)*sqrt(114))/507398496;
        u = s*pi/n+t*pi/n^2-c*pi/n^3;
        v = (pi/2-u)/(n/2-1);
        d = asin(sin(u)+sin(u+3*v/2)/(2*cos(v/2)))-u-v;
        Al = sin(u)+sin(2*v)-sin(v+d)+(n/2-3)*(sin(v)-tan(v/2))+(cos(v-d)-cos(2*v)-1/2)*tan(v/2);
        Au = (n*sin(pi/n)-(n-1)*tan(pi/(2*n-2)))/2;
        a0 = zeros(Int(n/2));
        a0[1] = u;
        a0[2] = v+d;
        a0[3] = v-d;
        if n>=8
            a0[4:end] .= v;
        end
        model = Model(Ipopt.Optimizer)
        # Variables
        @variable(model, a[i = 1:Int(n/2)] >= 0, start = a0[i])
        # Area
        @NLobjective(model, Max, sin(a[1])+sum(sum((-1)^j*(sin(sum(a[k] for k in i-j:i+1))-
        sin(sum(a[k] for k in i-j:i))) for j in 0:i-2) for i = 2:Int(n/2-1)))
        # Optimal diameter graph
        @constraint(model, sum(a[k] for k in 1:Int(n/2)) == pi/2)
        @NLconstraint(model, sum((-1)^(j-1)*sin(sum(a[k] for k in 1:j)) for j in 1:Int(n/2-1)) == (-1)^(n/2)/2)
        @constraint(model, a[1] <= pi/6)
        @constraint(model, [i = 2:Int(n/2)], a[i] <= 2*a[1])
        # Bounds
        @NLconstraint(model, Al <= sin(a[1])+sum(sum((-1)^j*(sin(sum(a[k] for k in i-j:i+1))-
        sin(sum(a[k] for k in i-j:i))) for j in 0:i-2) for i = 2:Int(n/2-1)) <= Au)
        # Solving
        optimize!(model)
        A = objective_value(model)
        a = value.(a)
        ct = solve_time(model)
        return(A,a,ct)
    end
end


include("utils.jl")

@testset "Positive Numbers" begin
    M = PositiveNumbers()
    @testset "Positive Numbers Basics" begin
        @test repr(M) == "PositiveNumbers()"
        @test repr(PositiveVectors(2)) == "PositiveVectors(2)"
        @test repr(PositiveMatrices(2, 3)) == "PositiveMatrices(2, 3)"
        @test repr(PositiveArrays(2, 3, 4)) == "PositiveArrays(2, 3, 4)"
        @test representation_size(M) == ()
        @test manifold_dimension(M) == 1
        @test !is_manifold_point(M, -1.0)
        @test_throws DomainError is_manifold_point(M, -1.0, true)
        @test is_tangent_vector(M, 1.0, 0.0; check_base_point = false)
        @test flat(M, 1.0, FVector(TangentSpace, 1.0)) == FVector(CotangentSpace, 1.0)
        @test sharp(M, 1.0, FVector(CotangentSpace, 1.0)) == FVector(TangentSpace, 1.0)
        @test vector_transport_to(M, 1.0, 3.0, 2.0, ParallelTransport()) == 6.0
        @test retract(M, 1.0, 1.0) == exp(M, 1.0, 1.0)
        @test isinf(injectivity_radius(M))
        @test isinf(injectivity_radius(M, -2.0))
        @test isinf(injectivity_radius(M, -2.0, ExponentialRetraction()))
        @test isinf(injectivity_radius(M, ExponentialRetraction()))
        @test project(M, 1.5, 1.0) == 1.0
        @test zero_tangent_vector(M, 1.0) == 0.0
        X = similar([1.0])
        zero_tangent_vector!(M, X, 1.0)
        @test X == [0.0]
    end
    types = [Float64]
    TEST_FLOAT32 && push!(types, Float32)
    for T in types
        @testset "Type $T" begin
            pts = convert.(Ref(T), [1.0, 4.0, 2.0])
            test_manifold(
                M,
                pts,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_default_vector_transport = true,
                test_vee_hat = false,
                is_mutating = false,
            )
        end
    end
    @testset "Power of Positive Numbers" begin
        M2 = PositiveVectors(2)
        for T in types
            pts2 = [convert.(Ref(T), v) for v in [[1.0, 1.1], [3.0, 3.3], [2.0, 2.2]]]
            test_manifold(
                M2,
                pts2,
                test_forward_diff = false,
                test_reverse_diff = false,
                test_vector_spaces = false,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_default_vector_transport = true,
                test_vee_hat = false,
            )
        end
    end
end


# # Fitting SNAP using GaN data
#
# This example:
#    - Uses the atomic positions in the folder `example_GaN_data`
#    - Generates surrogate DFT data based on the GaN model presented in `10.1088/1361-648x/ab6cbe`
#    - Uses `snap.jl` and 80% of the GaN data to create the matrix A. The matrix is generated only with the energy block.
#    - Uses 80% of the GaN data to create the vector b. The reference energy (``E_{\rm ref}``) is assumed to be zero.
#    - Uses backslash to fit the parameters ``\boldsymbol{\beta}``, thus, solves ``\mathbf{A \cdot \boldsymbol{\beta}=y}``
#    - Uses 20% of the GaN data and the bispectrum components provided by `snap.jl`and the fitted parameters ``\boldsymbol{\beta}`` to validate the fitting.
#    - The error computed during the validation is < ~5%.
# TODO:
#    - Check if the GaN model is correct.
#    - Calculate reference energy (``E_{\rm ref}``).
#    - Check ill-conditioned system.

#
# ## Install and import dependencies
#
# First let's make sure we have all required packages installed

# ```julia
# using Pkg
# pkg"add LAMMPS"
# ```
#
# We start off by importing the necessary packages

using LAMMPS
using LinearAlgebra: norm, pinv
using Printf

# ##  Estimation of the SNAP coefficients
#
# The choice of the coefficients ``\boldsymbol{\beta}=(\beta_0^1, \tilde{\beta}^1, \dots, \beta_0^l, \tilde{\beta}^l)``
# is based on a system of linear equations which considers a large number of atomic
# configurations and ``l`` atom types. The matrix formulation for this system ``\mathbf{A \cdot \boldsymbol{\beta}=y}``
# is defined in the following equations (see 10.1016/j.jcp.2014.12.018):

# ```math
# \begin{equation*}
#    \mathbf{A}=
#    \begin{pmatrix}
#        \vdots &  &  & & & & & & \\
#        N_{s_1} & \sum_{i=1}^{N_{s_1}} B_1^i & \dots & \sum_{i=1}^{N_{s_1}} B_k^i  & \dots & N_{s_L} & \sum_{i=1}^{N_{s_L}} B_1^i & \dots & \sum_{i=1}^{N_{s_L}} B_k^i\\
#        \vdots &  &  & & & & & & \\
#        0 & -\sum_{i=1}^{N_{s_1}} \frac{\partial B_1^i}{\partial r_j^{\alpha}} & \dots & -\sum_{i=1}^{N_{s_1}} \frac{\partial B_k^i}{\partial r_j^{\alpha}} & \dots &  0 & -\sum_{i=1}^{N_{s_l}} \frac{\partial B_1^i}{\partial r_j^{\alpha}} & \dots & -\sum_{i=1}^{N_{s_l}} \frac{\partial B_k^i}{\partial r_j^{\alpha}} \\
#        \vdots &  &  & & & & & & \\
#        0 & - \sum_{j=1}^{N_{s_1}} r^j_{\alpha} \sum_{i=1}^{N_{s_1}} \frac{\partial B_1^i}{\partial r_j^{\beta}} & \dots & - \sum_{j=1}^{N_{s_1}} r^j_{\alpha} \sum_{i=1}^{N_{s_1}} \frac{\partial B_k^i}{\partial r_j^{\beta}} & \dots & 0 & - \sum_{j=1}^{N_{s_l}} r^j_{\alpha} \sum_{i=1}^{N_{s_l}} \frac{\partial B_1^i}{\partial r_j^{\beta}} & \dots & - \sum_{j=1}^{N_{s_l}} r^j_{\alpha} \sum_{i=1}^{N_{s_l}} \frac{\partial B_k^i}{\partial r_j^{\beta}}\\
#        \vdots &  &  & & & & & & \\
#    \end{pmatrix}
# \end{equation*}
# ```

# The indexes ``\alpha, \beta = 1,2,3`` depict the ``x``, ``y`` and ``z``
# spatial component, ``j`` is an individual atom, and ``s`` a particular configuration.
# All atoms in each configuration are considered. The number of atoms of type $m$ in the configuration
# ``s`` is ``N_{s_m}``.

# The RHS of the linear system is computed as:

# ```math
# \begin{equation*}
#  \mathbf{y}=  \begin{pmatrix}
#    \vdots \\
#   E^s_{\rm qm} -  E^s_{\rm ref} \\
#   \vdots \\\\
#   F^{s,j,\alpha}_{\rm qm} - F^{s,j,\alpha}_{\rm ref} \\
#   \vdots \\
#   W_{\rm qm}^{s,\alpha,\beta} - W_{\rm ref}^{s,\alpha,\beta} \\
#   \vdots \\
#    \end{pmatrix}
# \end{equation*}
# ```

# # Calculate ``A`` using `snap.jl`

# This example only fits the energy, thus, the first block of the matrix A.

include(joinpath(dirname(pathof(LAMMPS)), "..", "examples", "snap.jl"))
A

# # Calculate ``y``

# A molecular mechanics model for the interaction of gallium and nitride is used
# to generate surrogate DFT data (see 10.1088/1361-648x/ab6cbe).
# The total potential energy is computed for each configuration

# ```math
# E = \sum_{i \lt j;  r_{i,j} \leq rcut } E_{GaN}(r_{i,j}) \  \text{ where } \\
#
# E_{GaN}(r_{i,j}) =
# \left\{
#    \begin{array}{ll}
#        E_{C}(r_{i,j}) + E_{LJ}(r_{i,j}, \epsilon_{Ga,Ga}, \sigma_{Ga,Ga})  & \mbox{if both atoms are Ga} \\
#        E_{C}(r_{i,j}) + E_{LJ}(r_{i,j}, \epsilon_{N,N}, \sigma_{N,N})  & \mbox{if both atoms are N} \\
#        E_{C}(r_{i,j}) + E_{BM}(r_{i,j}, A_{Ga,N}, \rho_{Ga,N})  & \mbox{if one atom is Ga and the other is N}\\
#    \end{array}
# \right.
# ```

# Note: this is a work in progress, this model should be checked.

const N = N1 + N2
const M1 = M
const M2 = 61

const ε_Ga_Ga = 0.643
const σ_Ga_Ga = 2.390
const ε_N_N = 1.474
const σ_N_N = 1.981
const A_Ga_N = 608.54
const ρ_Ga_N = 0.435
const q_Ga = 3.0
const q_N = -3.0
const ε0 = 55.26349406 # e2⋅GeV−1⋅fm−1 ToDo: check this
const E_ref = zeros(M1)
E_LJ(r, ε = 1.0, σ = 1.0) = 4.0 * ε * ((σ / norm(r))^12 - (σ / norm(r))^6)
E_BM(r, A = 1.0, ρ = 1.0) = A * exp(-norm(r) / ρ)
E_C(r) = q_Ga * q_N / (4.0 * π * ε0 * norm(r))

function read_atomic_conf(m, N)
    rs = []
    open(joinpath(DATA, string(m), "DATA")) do f
        for i = 1:23
            readline(f)
        end
        for i = 1:N
            s = split(readline(f))
            r = [parse(Float64, s[3]),
                 parse(Float64, s[4]),
                 parse(Float64, s[5])]
            push!(rs, r)
        end
    end
    return rs
end

function calc_tot_energy(rcut, rs, N, ε_Ga_Ga, σ_Ga_Ga, ε_N_N, σ_N_N, A_Ga_N, ρ_Ga_N)
    E_tot_acc = 0.0
    for i = 1:N
        for j = i:N
            r_diff = rs[i] - rs[j]
            if norm(r_diff) <= rcut && norm(r_diff) > 0.0
                if i <= N1 && j <= N1
                    E_tot_acc += E_C(r_diff) + E_LJ(r_diff, ε_Ga_Ga, σ_Ga_Ga)
                elseif i > N1 && j > N1
                    E_tot_acc += E_C(r_diff) + E_LJ(r_diff, ε_N_N, σ_N_N)
                else
                    E_tot_acc += E_C(r_diff) + E_BM(r_diff, A_Ga_N, ρ_Ga_N)
                end
            end
        end
    end
    return E_tot_acc
end

# Vector ``y`` is calculated using the DFT data and ``E_{\rm ref}``

function calc_y(rcut, M1, N, ε_Ga_Ga, σ_Ga_Ga, ε_N_N, σ_N_N, A_Ga_N, ρ_Ga_N, E_ref)
    y = zeros(M1)
    for m = 1:M1
        rs = read_atomic_conf(m, N)
        y[m] = calc_tot_energy(rcut, rs, N, ε_Ga_Ga, σ_Ga_Ga,
                               ε_N_N, σ_N_N, A_Ga_N, ρ_Ga_N) - E_ref[m]
    end
    return y
end

y = calc_y(rcut, M1, N, ε_Ga_Ga, σ_Ga_Ga, ε_N_N, σ_N_N, A_Ga_N, ρ_Ga_N, E_ref)

# ## Calculate the optimal solution

# The optimal solution $\widehat{\mathbf{\beta}}$ is, thus:
# ```math
# \begin{equation*}
#     \widehat{\mathbf{\beta}} = \mathrm{argmin}_{\mathbf{\beta}} ||\mathbf{A \cdot \boldsymbol{\beta} - y}||^2 = \mathbf{A^{-1} \cdot y}
# \end{equation*}
# ```

β = A \ y

# ## Check results

# The local energy can be decomposed into separate contributions for each atom.
# SNAP energy can be written in terms of the bispectrum components of the atoms
# and a set of coefficients. ``K`` components of the bispectrum are considered
# so that ``\mathbf{B}^{i}=\{ B^i_1, \dots, B_K^i\}`` for each atom ``i``, whose
# SNAP energy is computed as follows:
# ```math
#    E^i_{\rm SNAP}(\mathbf{B}^i) = \beta_0^{\alpha_i} + \sum_{k=1}^K \beta_k^{\alpha_i} B_k^i =  \beta_0^{\alpha_i} + \boldsymbol{\tilde{\beta}}^{\alpha_i} \cdot \mathbf{B}^i
# ```
# where $\alpha_i$ depends on the atom type.

function calc_fitted_tot_energy(path, β, ncoeff, N1, N)
    ## Calculate y
    lmp = LMP(["-screen","none"])
    bs = run_snap(lmp, path, rcut, twojmax)

    E_tot_acc = 0.0
    for n in 1:N1
        E_atom_acc = β[1]
        for k in 2:ncoeff
            k2 = k - 1
            E_atom_acc += β[k] * bs[k2, n]
        end
        E_tot_acc += E_atom_acc
    end
    for n in N1+1:N
        E_atom_acc = β[ncoeff+1]
        for k in ncoeff+2:2*ncoeff
            k2 = k - ncoeff - 1
            E_atom_acc += β[k] * bs[k2, n]
        end
        E_tot_acc += E_atom_acc
    end
    command(lmp, "clear")
    return E_tot_acc
end

# Comparison between the potential energy calculated based on the DFT data, and
# the SNAP potential energy calculated based on the bispectrum components
# provided by ``snap.jl`` and the fitted coefficients ``\mathbf{\beta}``.

@printf("Potential Energy, Fitted Potential Energy, Error (%%)\n")
for m = M1+1:M2
    path = joinpath(DATA, string(m), "DATA")
    rs = read_atomic_conf(m, N)
    E_tot = calc_tot_energy(rcut, rs, N, ε_Ga_Ga, σ_Ga_Ga, ε_N_N, σ_N_N, A_Ga_N, ρ_Ga_N)
    E_tot_fit = calc_fitted_tot_energy(path, β, ncoeff, N1, N)
    @printf("%0.2f, %0.2f, %0.2f\n", E_tot, E_tot_fit,
            abs(E_tot - E_tot_fit) / E_tot * 100.)
end


@testset "$(rpad("Special Partitioned Additive Runge-Kutta Tableaus",80))" begin

    @test typeof(TableauSPARKGLRK(1))                      <: TableauSPARK
    @test typeof(TableauSPARKGLRK(2))                      <: TableauSPARK
    @test typeof(TableauSPARKLobattoIIIAIIIB(2))           <: TableauSPARK
    @test typeof(TableauSPARKLobattoIIIAIIIB(3))           <: TableauSPARK
    @test typeof(TableauSPARKGLRKLobattoIIIAIIIB(1))       <: TableauSPARK
    @test typeof(TableauSPARKGLRKLobattoIIIAIIIB(2))       <: TableauSPARK
    @test typeof(TableauSPARKLobABC(2))                    <: TableauSPARK
    @test typeof(TableauSPARKLobABC(3))                    <: TableauSPARK
    @test typeof(TableauSPARKLobABD(2))                    <: TableauSPARK
    @test typeof(TableauSPARKLobABD(3))                    <: TableauSPARK
    @test typeof(TableauSPARKGLVPRK(1))                    <: TableauSPARK
    @test typeof(TableauSPARKGLVPRK(2))                    <: TableauSPARK

    @test typeof(TableauGausspSymplectic(1))                <: TableauVPARK
    @test typeof(TableauGausspSymplectic(2))                <: TableauVPARK
    @test typeof(TableauLobattoIIIAIIIBpSymplectic(2))     <: TableauVPARK
    @test typeof(TableauLobattoIIIAIIIBpSymplectic(3))     <: TableauVPARK
    @test typeof(TableauLobattoIIIAIIIBpSymplectic(4))     <: TableauVPARK
    @test typeof(TableauLobattoIIIBIIIApSymplectic(2))     <: TableauVPARK
    @test typeof(TableauLobattoIIIBIIIApSymplectic(3))     <: TableauVPARK
    @test typeof(TableauLobattoIIIBIIIApSymplectic(4))     <: TableauVPARK

    @test typeof(TableauVSPARKGLRKpLobattoIIIAIIIB(1))                     <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpLobattoIIIAIIIB(2))                     <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpLobattoIIIBIIIA(1))                     <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpLobattoIIIBIIIA(2))                     <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedLobattoIIIAIIIB(1))             <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedLobattoIIIAIIIB(2))             <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedLobattoIIIBIIIA(1))             <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedLobattoIIIBIIIA(2))             <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpInternal(1))                            <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpInternal(2))                            <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedInternal(1))                    <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedInternal(2))                    <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpMidpoint(1))                            <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpMidpoint(2))                            <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedMidpoint(1))                    <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpModifiedMidpoint(2))                    <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpSymmetric(1))                           <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpSymmetric(2))                           <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpSymplectic(1))                          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKGLRKpSymplectic(2))                          <: TableauVSPARKprimary

    @test typeof(TableauVSPARKLobattoIIIAIIIBpLobattoIIIAIIIB(2))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpLobattoIIIAIIIB(3))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApLobattoIIIAIIIB(2))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApLobattoIIIAIIIB(3))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpLobattoIIIBIIIA(2))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpLobattoIIIBIIIA(3))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApLobattoIIIBIIIA(2))          <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApLobattoIIIBIIIA(3))          <: TableauVSPARKprimary

    @test typeof(TableauVSPARKLobattoIIIAIIIBpModifiedLobattoIIIAIIIB(2))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpModifiedLobattoIIIAIIIB(3))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpModifiedLobattoIIIBIIIA(2))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpModifiedLobattoIIIBIIIA(3))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApModifiedLobattoIIIAIIIB(2))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApModifiedLobattoIIIAIIIB(3))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApModifiedLobattoIIIBIIIA(2))  <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApModifiedLobattoIIIBIIIA(3))  <: TableauVSPARKprimary

    @test typeof(TableauVSPARKLobattoIIIAIIIBpMidpoint(2))                 <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpMidpoint(3))                 <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApMidpoint(2))                 <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApMidpoint(3))                 <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpModifiedMidpoint(2))         <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpModifiedMidpoint(3))         <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApModifiedMidpoint(2))         <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApModifiedMidpoint(3))         <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpSymmetric(2))                <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIAIIIBpSymmetric(3))                <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApSymmetric(2))                <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobattoIIIBIIIApSymmetric(3))                <: TableauVSPARKprimary

    @test typeof(TableauVSPARKLobABCCD(2))            <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobABCCE(2))            <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobABDE(2))             <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobABED(2))             <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobABD(2))              <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobABE(2))              <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobDE(2))               <: TableauVSPARKprimary
    @test typeof(TableauVSPARKLobED(2))               <: TableauVSPARKprimary

    @test typeof(TableauVSPARKLobattoIIIAB(2))        <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKLobattoIIIBA(2))        <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKLobattoIIICC̄(2))        <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKLobattoIIIC̄C(2))        <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKLobattoIIID(2))         <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKLobattoIIIE(2))         <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKGLRKLobattoIIIAB(2))    <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKGLRKLobattoIIIBA(2))    <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKGLRKLobattoIIICC̄(2))    <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKGLRKLobattoIIIC̄C(2))    <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKGLRKLobattoIIID(2))     <: TableauVSPARKsecondary
    @test typeof(TableauVSPARKGLRKLobattoIIIE(2))     <: TableauVSPARKsecondary

    @test typeof(TableauHPARKGLRK(1))                 <: TableauHPARK
    @test typeof(TableauHPARKGLRK(2))                 <: TableauHPARK
    @test typeof(TableauHPARKLobattoIIIAIIIB(2))      <: TableauHPARK
    @test typeof(TableauHPARKLobattoIIIAIIIB(3))      <: TableauHPARK
    @test typeof(TableauHPARKLobattoIIIBIIIA(2))      <: TableauHPARK
    @test typeof(TableauHPARKLobattoIIIBIIIA(3))      <: TableauHPARK

    @test TableauHPARKGLRK(1) == getTableauHPARK(:HPARKGLRK1, TableauGauss(1), TableauGauss(1))
    @test TableauHPARKGLRK(2) == getTableauHPARK(:HPARKGLRK2, TableauGauss(2), TableauGauss(2))

    @test typeof(TableauHSPARKGLRKpSymmetric(1))                       <: TableauHSPARKprimary
    @test typeof(TableauHSPARKGLRKpSymmetric(2))                       <: TableauHSPARKprimary
    @test typeof(TableauHSPARKLobattoIIIAIIIBpSymmetric(2))            <: TableauHSPARKprimary
    @test typeof(TableauHSPARKLobattoIIIAIIIBpSymmetric(3))            <: TableauHSPARKprimary
    @test typeof(TableauHSPARKLobattoIIIAIIIBpSymmetric(4))            <: TableauHSPARKprimary
    @test typeof(TableauHSPARKLobattoIIIBIIIApSymmetric(2))            <: TableauHSPARKprimary
    @test typeof(TableauHSPARKLobattoIIIBIIIApSymmetric(3))            <: TableauHSPARKprimary
    @test typeof(TableauHSPARKLobattoIIIBIIIApSymmetric(4))            <: TableauHSPARKprimary

    @test typeof(TableauHSPARKLobattoIIIAB(2))                         <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKLobattoIIIBA(2))                         <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKLobattoIIICC̄(2))                         <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKLobattoIIIC̄C(2))                         <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKLobattoIIID(2))                          <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKLobattoIIIE(2))                          <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKGLRKLobattoIIIAB(1))                     <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKGLRKLobattoIIIBA(1))                     <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKGLRKLobattoIIICC̄(1))                     <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKGLRKLobattoIIIC̄C(1))                     <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKGLRKLobattoIIID(1))                      <: TableauHSPARKsecondary
    @test typeof(TableauHSPARKGLRKLobattoIIIE(1))                      <: TableauHSPARKsecondary

    @test typeof(TableauSLRKLobattoIIIAB(2))          <: TableauSLRK
    @test typeof(TableauSLRKLobattoIIIBA(2))          <: TableauSLRK
    @test typeof(TableauSLRKLobattoIIICC̄(2))          <: TableauSLRK
    @test typeof(TableauSLRKLobattoIIIC̄C(2))          <: TableauSLRK
    @test typeof(TableauSLRKLobattoIIID(2))           <: TableauSLRK
    @test typeof(TableauSLRKLobattoIIIE(2))           <: TableauSLRK

end


using ITensorGaussianMPS
using LinearAlgebra
using ITensors

# Half filling
N = 50
Nf = N÷2

@show N, Nf

# Maximum MPS link dimension
_maxlinkdim = 100

@show _maxlinkdim

# DMRG cutoff
_cutoff = 1e-12

# Hopping
t = 1.0

# Electron-electron on-site interaction
U = 1.0

@show t, U

# Free fermion Hamiltonian
ampo = AutoMPO()
for n in 1:N-1
  ampo .+= -t, "Cdag", n, "C", n+1
  ampo .+= -t, "Cdag", n+1, "C", n
end

# Hopping Hamiltonian with N spinless fermions
h = hopping_hamiltonian(ampo)

# Get the Slater determinant
Φ = slater_determinant_matrix(h, Nf)

# Create an mps for the free fermion ground state
s = siteinds("Fermion", N; conserve_qns = true)
println("Making free fermion starting MPS")
@time ψ0 = slater_determinant_to_mps(s, Φ; eigval_cutoff = 1e-4,
                                           cutoff = _cutoff,
                                           maxdim = _maxlinkdim)
@show maxlinkdim(ψ0)

# Make an interacting Hamiltonian
for n in 1:N-1
  ampo .+= U, "N", n, "N", n+1
end
H = MPO(ampo, s)

# Random starting state
ψr = randomMPS(s, n -> n ≤ Nf ? "1" : "0")

println("\nRandom state starting energy")
@show flux(ψr)
@show inner(ψr, H, ψr)

println("\nFree fermion starting energy")
@show flux(ψ0)
@show inner(ψ0, H, ψ0)

println("\nRun dmrg with random starting state")
sweeps = Sweeps(20)
maxdim!(sweeps,10,20,40,_maxlinkdim)
cutoff!(sweeps,_cutoff)
@time dmrg(H, ψr, sweeps)

println("\nRun dmrg with free fermion starting state")
sweeps = Sweeps(4)
maxdim!(sweeps,_maxlinkdim)
cutoff!(sweeps,_cutoff)
@time dmrg(H, ψ0, sweeps)

nothing


@testset "Distances" begin
  p = P2(0,1)
  l = Line(P2(0,0), P2(1,0))
  @test evaluate(Euclidean(), p, l) == T(1)
  @test evaluate(Euclidean(), l, p) == T(1)

  p1, p2 = P2(1,0), P2(0,1)
  @test evaluate(Chebyshev(), p1, p2) == T(1)
end


mutable struct FiniteSignalType{T}
    s0::Vector{T} # default signal.
    a::Float64
    b::Float64

    s::Vector{T} # buffer.
    #a::Float64
    #b::Float64
    a_ind
    b_ind
end

function FiniteSignalType(s, a, b)
    return FiniteSignalType(copy(s), a, b, copy(s), 1, length(s))
end

function resetsignal!(S::FiniteSignalType{T}) where T

    S.a_ind = 1
    S.b_ind = length(S.s0)
    S.s[:] = S.s0

    return nothing
end

# positive is delay, or shift towards more positive.
function applyintegershift!(S::FiniteSignalType{T}, d::Int) where T
    #
    S.a_ind = S.a_ind-d
    S.b_ind = S.b_ind-d

    return nothing
end

function evalfinitesignal(s::Vector{T}, i::Int)::T where T
    if 1 <= i <= length(s)
        return s[i]
    end

    return zero(T)
end

function evalsignal(S::FiniteSignalType{T}, n::Int)::T where T

    #ind_t_Float = convertcompactdomain(t, S.a, S.b, S.a_ind, S.b_ind)
    #ind_t_Float = round(Int, ind_t_Float)
    #@assert 1 <= n <= length(S.s)

    return evalfinitesignal(S.s, S.a_ind + n-1)
end

#  https://lcondat.github.io/publis/condat_icip08_rotations.pdf
function getthiriancoeffs!(b::Vector{T}, τ::T) where T <: Real

    #resize!(b, N)
    N = length(b)
    for k = 1:N
        b[k] = (-1)^k * binomial(N,k) * prod( (τ-n)/(τ-n-k) for n = 0:N)
    end

    return nothing
end

# https://lcondat.github.io/publis/condat_icip08_rotations.pdf
# τ ∈ [0, 0.5]
# mutates s and b_buffer.
function applydelaythirian!(b_buffer::Vector{T}, s::Vector{T}, τ::T) where T <: Real

    getthiriancoeffs!(b_buffer, τ)

    N = length(b_buffer)
    M = length(s)

    for i = M-1:-1:0
        for k = 1:N

            #s[i+1] += b_buffer[k]*(s[i-k+1]-s[i+k+1])

            tmp = evalfinitesignal(s, i-k+1) - evalfinitesignal(s, i+k+1)

            s[i+1] += b_buffer[k]*tmp
        end
    end

    return nothing
end

function applyadvancethirian!(b_buffer::Vector{T}, s::Vector{T}, τ::T) where T <: Real

    getthiriancoeffs!(b_buffer, τ)

    N = length(b_buffer)
    M = length(s)

    #reverse!(s)
    for i = M-1:-1:0
        for k = 1:N

            # tmp = evalfinitesignal(s, i-k+1) - evalfinitesignal(s, i+k+1)
            # s[i+1] += b_buffer[k]*tmp

            # reverse: if we have the ORDERED lists:
            # U = {1, 2, ..., M}, V = {M, M-1, ... , 1}, then
            # j = M-i+1, where i is a given element from U, and j is the
            #   corresponding element in V.

            j = M-i #j = M-(i+1)+1
            j1 = j+k #j1 = M-(i-k+1)+1
            j2 = j-k #j2 = M-(i+k+1)+1
            tmp = evalfinitesignal(s, j1) - evalfinitesignal(s, j2)
            s[j] += b_buffer[k]*tmp
        end
    end
    #reverse!(s)

    return nothing
end

function applythirian!(b_buffer)

    if τ > 0.5
        return nothing # shift, then apply thirian.
    end

    return nothing # applythirian!()
end

function applyshift!(S::FiniteSignalType{T}, shift::T, b_thirian) where T

    d = round(Int, shift)
    τ = shift-d
    println(τ)

    applyintegershift!(S, d)
    if τ > 0
        applydelaythirian!(b_thirian, S.s, τ)
    else
        applyadvancethirian!(b_thirian, S.s, abs(τ))
    end

    return nothing
end

# other interesting articles about fractional delay: https://www.dsprelated.com/showarticle/1327.php


@testset "map" begin
    @inline foo(x, y) = exp(x) - sin(y)
    for T ∈ (Float32,Float64)
        @show T, @__LINE__
        for N ∈ [ 3, 371 ]
            a = rand(T, N); b = rand(T, N);
            c1 = map(foo, a, b);
            c2 = vmap(foo, a, b);
            @test c1 ≈ c2
            c2 = vmapnt(foo, a, b);
            @test c1 ≈ c2
            c2 = vmapntt(foo, a, b);
            @test c1 ≈ c2
            fill!(c2, NaN); @views vmapnt!(foo, c2[2:end], a[2:end], b[2:end]);
            @test @views c1[2:end] ≈ c2[2:end]
            fill!(c2, NaN); @views vmapntt!(foo, c2[2:end], a[2:end], b[2:end]);
            @test @views c1[2:end] ≈ c2[2:end]
        end

        c = rand(T,100); x = rand(T,10^4); y1 = similar(x); y2 = similar(x);
        map!(xᵢ -> clenshaw(xᵢ, c), y1, x)
        vmap!(xᵢ -> clenshaw(xᵢ, c), y2, x)
        @test y1 ≈ y2

    end
end


export calculate_sfc_fluxes_energy
export DryBulkFormulaWithRadiation, DryBulkFormula, LinearRelaxation
#surface_fluxes
"""
dummy for SurfaceFluxes.jl
"""
abstract type SurfaceFluxType end

struct LinearRelaxation <: SurfaceFluxType end
struct DryBulkFormula <: SurfaceFluxType end
struct DryBulkFormulaWithRadiation <: SurfaceFluxType end

calculate_sfc_fluxes_energy(forumation::LinearRelaxation, p, θ_sfc, θ1) = p.λ .* (θ_sfc .- θ1)

calculate_sfc_fluxes_energy(forumation::DryBulkFormula, p, θ_sfc, θ1, u_1, v_1, ρ_1) =
    p.Ch * p.C_p * ρ_1 * sqrt(u_1^2 + v_1^2) * (θ_sfc .- θ1)

function calculate_sfc_fluxes_energy(forumation::DryBulkFormulaWithRadiation, parameters, θ_sfc, θ1, u_1, v_1, ρ_1, t)

    p = parameters

    R_SW = (1 - p.α) * p.τ * p.F_sol * (1 .+ sin(t * 2π / p.τ_d))
    R_LW = p.ϵ * (p.σ * θ_sfc .^ 4 - p.σ * θ1 .^ 4)
    SH = p.Ch * p.C_p * ρ_1 * sqrt(u_1^2 + v_1^2) * (θ_sfc - θ1)

    F_tot = -(R_SW - R_LW - SH)
end


using Mimi, Distributions

include("MyModel.jl")
using .MyModel

m = construct_MyModel()

# Step 2: Define the Simulation

"""
@defsim macro is the first step in the process and returns a SimulationDef

allows users to define random variables as distributions, and associate model
parameters with the defined random variables
"""

sd = @defsim begin

    # Define random variables. The rv() is only required
    # when defining correlations or sharing an RV across parameters
    # otherwise, you can use the shortcut syntax
    # to assign a distribution to a parameter name
    rv(name1) = Normal(1, 0.2)
    rv(name2) = Uniform(0.75, 1.25)
    rv(name3) = LogNormal(20, 4)

    # Define the sampling strategy
    sampling(LHSData,corrlist = [(:name1, :name2, 0.7), (:name1, :name3, 0.5)])

    # assign RVs to model Parameters
    grosseconomy.share = Uniform(0.2, 0.8)

    # you can use the *= operator to replace the values in the parameter with the
    # product of the original value and the value of the RV for the current
    # trial (note that in both lines below, all indexed values will be mulitplied by the
    # same draw from the given random parameter (name2 or Uniform(0.8, 1.2))

    emissions.sigma[:,Region1] *= name2
    emissions.sigma[2020:5:2050, (Region2, Region3)] *= Uniform(0.8, 1.2)

    # For parameters that have a region dimension, you can assign an array of distributions,
    # keyed by region label, which must match the region labels in the model
    grosseconomy.depk = [Region1 => Uniform(0.7, .9),
                        Region2 => Uniform(0.8, 1.),
                        Region3 => Truncated(Normal(), 0, 1)]

    # Indicate which variables to save for each model run
    # The syntax is: component_name.variable_name
    save(grosseconomy.K, grosseconomy.YGROSS,
        emissions.E, emissions.E_Global)
end

# STEP 3: Run Simulation

# Run 100 trials and optionally save results to the indicated directories
si = run(sd, m, 100; trials_output_filename = "/tmp/trialdata.csv", results_output_dir = "/tmp/tutorial4")

# Explore the results saved in-memory by using getdataframe with the returned SimulationInstance.
# Values are saved from each trial for each variable or parameter specified by the call to "save()" at the end of the @defsim block.
K_results = getdataframe(si, :grosseconomy, :K)
E_results = getdataframe(si, :emissions, :E)

explore(si)


using BayesOpt
using Test
include("branin.jl")


############################################################################################
"""
$(TYPEDEF)

Holds information about updating tempering parameter, parameter for basic logistic function.

# Fields
$(TYPEDFIELDS)
"""
struct TemperingParameter{T<:AbstractFloat}
    L::T
    k::T
    x₀::T
    function TemperingParameter(L::T, k::T, x₀::T) where {T<:AbstractFloat}
        ArgCheck.@argcheck 0 < L
        ArgCheck.@argcheck 0 < k
        return new{T}(L, k, x₀)
    end
end
function TemperingParameter(::Type{T}, growthrate::F, iterations::Integer; normalizer::Float64 = 100.0
) where {F<:AbstractFloat, T<:AbstractFloat}
    return TemperingParameter(T(1.0), T(growthrate*normalizer/iterations), T(growthrate*iterations/normalizer))
end

############################################################################################
"""
$(SIGNATURES)
Initialize temperature.

# Examples
```julia
```

"""
function init(::Type{T}, adaption::UpdateFalse, val::F, parameter::TemperingParameter
) where {T<:AbstractFloat, F<:AbstractFloat}
    return T(val)
end
function init(::Type{T}, adaption::UpdateTrue, val::F, parameter::TemperingParameter
) where {T<:AbstractFloat, F<:AbstractFloat}
    return update(parameter, 1)
end

############################################################################################
"""
$(SIGNATURES)
Calculate new temperature based on current iteration.

# Examples
```julia
```

"""
function update(param::TemperingParameter{T}, index::Integer) where {T<:AbstractFloat}
    @unpack L, k, x₀ = param
    return L / (1 + exp(-k*(T(index) - x₀)))
end

############################################################################################
"""
$(SIGNATURES)
Helper function to determine temperature shape given param variables.

# Examples
```julia
```

"""
function checktemperature(param::TemperingParameter, iterations::Integer)
    return map(iter -> update(param, iter), Base.OneTo(iterations))
end

############################################################################################
"""
$(TYPEDEF)

Tuning container for iteration uptdates of temperature.

# Fields
$(TYPEDFIELDS)
"""
struct IterationTempering{A<:UpdateBool, T<:AbstractFloat} <: TemperingMethod
    "Checks if temperature will be updated after new iterations"
    adaption::A
    "Initial temperature for tempering."
    initial::ValueHolder{T}
    "Tuning parameter for temperature adjustment."
    parameter::TemperingParameter{T}
    function IterationTempering(adaption::A, val::ValueHolder{T}, parameter::TemperingParameter{T}
    ) where {A<:UpdateBool,T<:AbstractFloat}
        @argcheck 0.0 < val.current <= 1.0
        return new{A,T}(adaption, val, parameter)
    end
end

function IterationTempering(::Type{T}, adaption::B, val::F, idx::Integer
) where {T<:AbstractFloat, B<:UpdateBool, F<:AbstractFloat}
    # Assign temperature adaption
    parameter = TemperingParameter(T, val, idx)
    # Assign initial temperature
    val₀ = init(T, adaption, val, parameter)
    # Return tuning struct
    return IterationTempering(adaption, ValueHolder(val₀), parameter)
end

############################################################################################
function update!(tempering::IterationTempering, adaption::UpdateTrue, index::Integer)
    # Compute new temperature
    return update(tempering.parameter, index)
end
function update!(tempering::IterationTempering, adaption::UpdateFalse, index::Integer)
    return tempering.initial.current
end
function update!(tempering::IterationTempering, index::Integer)
    return update!(tempering, tempering.adaption, index)
end

############################################################################################
"""
$(SIGNATURES)
Return initial temperature of tuning container.

# Examples
```julia
```

"""
function initial(tempertune::IterationTempering)
    return tempertune.initial.current
end

############################################################################################
# Export
export
    IterationTempering,
    TemperingParameter,
    init,
    initial,
    update,
    update!


module MatrixPolynomials

using LinearAlgebra

export
    Pol, variable, rev,
    lc, lm, subs,
    zero, one, izero, polrank,
    coeffs, coeff, deg,
    companion, roots,
    toPencil, eig, colNullities, rowNullities, colIndices, rowIndices, infMultiplicities, eigNullities, eigMultiplicities,
    wilkinson,
    colBlock, blkdiag, rowBlock, infBlock, eigBlock

include("utils.jl")
include("polynomials/polynomials.jl")
include("polynomials/arithmetic.jl")
include("polynomials/specialPolynomials.jl")
include("polynomials/roots.jl")
include("matrices.jl")
include("eigenvalues/eigenvalues.jl")
include("eigenvalues/nullities.jl")
include("kronecker/blocks.jl")
end # module


"""
    gmps2mps(gmps, maxdim)

Construct a matrix product state with a maximum bond dimension
`max_dim` from a set of GaussianMPS gates `gmps`.

Since the unitary gates orthogonalize the correlation matrix, if they
are applied in revese order to the configuration state (which is the
state in the occupation basis) they will generate the state in the
original spatial basis. So, a classical MPS from the configurations is
made and then the gates are applied in reverse order.

"""
function gmps2mps(gmps::GaussianMPS,
                  maxdim::Int64;
                  symmetry::Type{<:AbstractCharge}=Trivial)

    ##TODO: why are these defined complex here?!
    if symmetry == Trivial
        mps = MPS{Float64}(gmps.lx, 2, gmps.initconf)
        _applygmpsgates!(mps, gmps, maxdim)

    elseif symmetry == U1
        mps = U1MPS(Float64, gmps.lx, 2, gmps.initconf)
        _applygmpsgates!(mps, gmps, maxdim)
    else
        throw(ArgumentError("Unkown symmetry : $symmetry"))
    end
    mps
end

function _applygmpsgates!(mps::MPState{Array{T,3}},
                          gmps::GaussianMPS,
                          maxdim::Int) where {T}
    twobodyugate = Matrix{T}(I, 4, 4)
    ## NOTE: the gates are applies in reverse order
    for n=length(gmps.xs):-1:1
        site = gmps.xs[n]
        θ = gmps.θs[n]
        twobodyugate[2:3, 2:3] = rotationmat(θ)

        ugatetensor = reshape(twobodyugate, 2, 2, 2, 2)

        # produce the counterclockwise convention starting from bottom left
        ##NOTE: this is needed just because I followed the convention,
        ##probably best to avoid id, but need to remain consistant
        operator = permutedims(ugatetensor, [1,2,4,3])

        ## TODO: choose a better order of site or site+1 and push_to!
        center_at!(mps, site)
        apply!(mps, operator, site, maxdim=maxdim, pushto=:R, svnormalize=true)
    end
    nothing
end

function _applygmpsgates!(mps::MPState{SymTensor{S,T,3}},
                          gmps::GaussianMPS,
                          maxdim::Int) where {S,T}

    Vd = VectorSpace{S}(0=>1, 1=>1)
    # u is the two-body U gate matrix
    u = eye(Float64, fuse(false, Vd,Vd))

    #   twobodyugate = Matrix{ComplexF64}(I, 4, 4)
    ## NOTE: the gates are applies in reverse order

    for n=length(gmps.xs):-1:1
        site = gmps.xs[n]
        θ = gmps.θs[n]
        u[Sector{S}(1,1)] =  rotationmat(θ)

        uten = splitleg(splitleg(u, 2, (dual(Vd), dual(Vd))), 1, (Vd,Vd))

        ## TODO: choose a better order of site or site+1 and push_to!
        center_at!(mps, site)
        apply!(mps, uten, site, maxdim=maxdim, pushto=:R, svnormalize=true)
    end
    nothing
end

function rotationmat(θ::Float64)
    R = [
        cos(θ) sin(θ);
        -sin(θ) cos(θ)
    ]
    R
end


export BoltzmannPolicy

"Boltzmann policy with softmax action noise."
struct BoltzmannPolicy{P <: PolicySolution, R <: AbstractRNG} <: PolicySolution
    policy::P
    temperature::Float64
    rng::R
end

BoltzmannPolicy(policy, temperature) =
    BoltzmannPolicy(policy, temperature, Random.GLOBAL_RNG)

"Convert vector of scores to probabiities."
function softmax(scores)
    ws = exp.(scores .- maximum(scores))
    z = sum(ws)
    return isnan(z) ? ones(length(scores)) ./ length(scores) : ws ./ z
end

get_action(sol::BoltzmannPolicy, state::State) =
    rand_action(sol, state)
best_action(sol::BoltzmannPolicy, state::State) =
    best_action(sol.policy, state)
get_value(sol::BoltzmannPolicy, state::State) =
    get_value(sol.policy, state)
get_value(sol::BoltzmannPolicy, state::State, action::Term) =
    get_value(sol.policy, state, action)
get_action_values(sol::BoltzmannPolicy, state::State) =
    get_action_values(sol.policy, state)

unzip_pairs(ps::AbstractDict) = unzip_pairs(collect(ps))
unzip_pairs(ps::AbstractArray{<:Pair}) = first.(ps), last.(ps)

function rand_action(sol::BoltzmannPolicy, state::State)
    actions, q_values = unzip_pairs(get_action_values(sol, state))
    if sol.temperature == 0
        probs = zeros(length(actions))
        probs[argmax(q_values)] = 1.0
    else
        probs = softmax(q_values ./ sol.temperature)
    end
    return sample(sol.rng, actions, Weights(probs, 1.0))
end


#=

  Lectures problem in Julia ConstraintSolver.js

  Biggs: Discrete Mathematics (2nd ed), page 187.
  """
  Suppose we wish to schedule six one-hour lectures, v1, v2, v3, v4, v5, v6.
  Among the the potential audience there are people who wish to hear both

   - v1 and v2
   - v1 and v4
   - v3 and v5
   - v2 and v6
   - v4 and v5
   - v5 and v6
   - v1 and v6

  How many hours are necessary in order that the lectures can be given
  without clashes?
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function lectures(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    n, _ = size(problem)
    @variable(model, 1 <= x[1:n] <= n, Int)
    @variable(model, 1 <= max_t <= n, Int)

    for (l1,l2) in eachrow(problem)
        @constraint(model, x[l1] != x[l2])
    end

    # symmetry breaking
    for i in 1:n
        @constraint(model, x[i] <= i)
    end
    @constraint(model, x .<= max_t)
    @objective(model,Min,max_t)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                max_t_val = convert.(Integer,JuMP.value.(max_t; result=sol))
                println("x:$x_val")
                println("max_t:$max_t_val")
                println()

            end
        end
    else
        println("status:$status")
    end

    return status
end

# The schedule requirements:
#   lecture a cannot be held at the same time as b
 lectures_problem = resize_matrix([[1, 2],
                                    [1, 4],
                                    [3, 5],
                                    [2, 6],
                                    [4, 5],
                                    [5, 6],
                                    [1, 6]])

@time lectures(lectures_problem)


using Alea, FactCheck

facts("GPC") do
  germ = GPCGerm(Normal(0,1), 4)
  #=
  @show rand(germ)
  @show pdf(germ, [0.3, 0.4, 0.5, 0.6])
  @show cdf(germ, [0.3, 0.4, 0.5, 0.6])
  @show cdf(germ, [0.0, 0.0, 0.0, 0.0])
  @show cdf(germ, [Inf, 0.0, 0.0, Inf])
  =#
  context("cdf") do
    @fact cdf(germ, zeros(4)) --> 0.5^4
    @fact cdf(germ, [0.0, 0.0, 0.0, Inf]) --> 0.5 ^ 3
    @fact cdf(germ, [0.0, -Inf, 0.0, Inf]) --> 0
  end
end


using Base.Test
using QuantumOptics

ωc = 1.2
ωa = 0.9
g = 1.0
Ncutoff = 10

T = Float64[0.,1.]

fockbasis = FockBasis(Ncutoff)
spinbasis = SpinBasis(1//2)
basis = tensor(spinbasis, fockbasis)

sx = sigmax(spinbasis)
sy = sigmay(spinbasis)
sz = sigmaz(spinbasis)
sp = sigmap(spinbasis)
sm = sigmam(spinbasis)

Ha = embed(basis, 1, 0.5*ωa*sz)
Hc = embed(basis, 2, ωc*number(fockbasis))
Hint = g/2*(sm ⊗ create(fockbasis) + sp ⊗ destroy(fockbasis))

H = Ha + Hc + Hint

Omega_n(n::Int) = sqrt((ωa-ωc)^2 + g^2*(n+1))
Alpha_n(n::Int) = atan2(g*sqrt(n+1),(ωa-ωc))

Em(n::Int) = ωc*(n+0.5) - 0.5*Omega_n(n)
Ep(n::Int) = ωc*(n+0.5) + 0.5*Omega_n(n)

psi_e(n::Int) = spinup(spinbasis) ⊗ fockstate(fockbasis, n)
psi_g(n::Int) = spindown(spinbasis) ⊗ fockstate(fockbasis, n+1)

psi_p(n::Int) = cos(Alpha_n(n)/2)*psi_e(n) + sin(Alpha_n(n)/2)*psi_g(n)
psi_m(n::Int) = -sin(Alpha_n(n)/2)*psi_e(n) + cos(Alpha_n(n)/2)*psi_g(n)
basisstates_p = Ket[psi_p(n) for n=0:Ncutoff-1]
basisstates_m = Ket[psi_m(n) for n=0:Ncutoff-1]

Emvec = Float64[Em(n) for n=0:Ncutoff-1]
Epvec = Float64[Ep(n) for n=-1:Ncutoff-1]

P = sortperm([Emvec; Epvec])
basisstates = Ket[basisstates_m; psi_g(-1); basisstates_p][P][1:(Ncutoff+1)]
E = [Emvec; Epvec][P][1:(Ncutoff+1)]

@test norm(operatorspectrum(full(H))[1:(Ncutoff+1)] - E) < 1e-12
@test norm(operatorspectrum_hermitian(full(H))[1:(Ncutoff+1)] - E) < 1e-12
@test norm(operatorspectrum(H)[1:(Ncutoff+1)] - E) < 1e-12
@test norm(operatorspectrum_hermitian(H)[1:(Ncutoff+1)] - E) < 1e-12

b = eigenstates(full(H))[1:(Ncutoff+1)]
for i=1:length(b)
    @test 1-abs(dagger(b[i])*basisstates[i])<1e-12
end

b = eigenstates_hermitian(full(H))[1:(Ncutoff+1)]
for i=1:length(b)
    @test 1-abs(dagger(b[i])*basisstates[i])<1e-12
end

b = eigenstates(H)[1:(Ncutoff+1)]
for i=1:length(b)
    @test 1-abs(dagger(b[i])*basisstates[i])<1e-12
end

b = eigenstates_hermitian(H)[1:(Ncutoff+1)]
for i=1:length(b)
    @test 1-abs(dagger(b[i])*basisstates[i])<1e-12
end