export sample_psdd

using StatsFuns
using LogicCircuits

"Samples an element from a Binomial distribution with p=0.5."
function sample_row(n::Int)::Int
    P = Float64[binompdf(n, 0.5, k) for k ∈ 1:n]
    c = 0.0
    r = rand()
    for (i, p) ∈ enumerate(P)
        c += p
        if r < c return i end
    end
    return n
end

"Samples a single combination uniformly following Robert Floyd's algorithm."
function sample_comb(n::Int, k::Int)::BitSet
    S = BitSet()
    for i ∈ n-k+1:n
        r = rand(1:i)
        if r ∉ S push!(S, r)
        else push!(S, i) end
    end
    return S
end

@inline function random_weights(n::Int)::Vector{Float64}
    W = rand(Float64, n)
    return log.(W /= sum(W))
end

@inline function get_lit(l::Int32, V::Vtree, L::Dict{Int32, StructProbLiteralNode})::StructProbLiteralNode
    if !haskey(L, l)
        node = StructProbLiteralNode(l, V)
        L[l] = node
        return node
    end
    return L[l]
end

"Options for SamplePSDD."
@enum SamplingOpts begin
    plain    = 0  # No compressions or merges.
    compress = 1  # Only compressions are allowed.
    merge    = 2  # Only merges are allowed.
    full     = 3  # Compressions and merges are allowed.
end

"""
Samples a partial partition.
"""
function sample_partition(ϕ::Bdd, Sc::BitSet, p::Real, k::Integer, ⊤_k::Integer, exact::Bool)::Dict{Bdd, Vector{Bdd}}
    X = intersect!(scope(ϕ), Sc)
    idem = isempty(X) || is_⊤(ϕ)
    O = shuffle!(idem ? collect(Sc) : X)
    E = Dict{Bdd, Vector{Bdd}}()
    if idem
        E[ϕ] = Vector{Bdd}()
        sample_idem_primes!(O, p, E[ϕ], ⊤_k)
    else sample_primes!(ϕ, O, E, k, exact) end
    return E
end

"""
Samples primes for the ⊤ case.
"""
function sample_idem_primes!(O::Vector{Int}, p::Real, P::Vector{Bdd}, k::Integer)
    e_count, Sc_len = 1, length(O)
    Q = Tuple{Int, Vector{Int}}[(1, Vector{Int}())]
    while !isempty(Q)
        i, V = popfirst!(Q)
        if e_count >= k || i > Sc_len
            push!(P, and(V))
            continue
        end
        x = O[i]
        c = rand() > p
        if i > 1 && c
            push!(Q, (i+1, V))
        else
            push!(Q, (i+1, push!(copy(V), x)))
            push!(Q, (i+1, push!(copy(V), -x)))
            e_count += 1
        end
    end
    nothing
end

"""
Samples primes for partial partition.
"""
function sample_primes!(ϕ::Bdd, O::Vector{Int}, E::Dict{Bdd, Vector{Bdd}}, k::Integer,
        exact::Bool)
    e_count, Sc_len = 1, length(O)
    Q = Tuple{Bdd, Int, Vector{Int}}[(ϕ, 1, Vector{Int}())]
    while !isempty(Q)
        ψ, i, V = popfirst!(Q)
        if (e_count >= k && !exact) || (i > Sc_len) || is_⊤(ψ)
            if !haskey(E, ψ) E[ψ] = Bdd[and(V)]
            else push!(E[ψ], and(V)) end
            continue
        end
        x = O[i]
        if x ∉ ψ push!(Q, (ψ, i+1, V))
        else
            α, β = ψ|x, ψ|-x
            if !is_⊥(α) push!(Q, (α, i+1, push!(copy(V), x))) end
            if !is_⊥(β) push!(Q, (β, i+1, push!(copy(V), -x))) end
            e_count += 1
        end
    end
    nothing
end

"""
Samples a PSDD from a BDD `ϕ` and vtree `V` with at most `k` elements in each disjunction node.
"""
@inline function sample_psdd(ϕ::Bdd, V::Vtree, k::Integer, D::DataFrame; opts::SamplingOpts = full,
        randomize_weights::Bool = false, pseudocount::Real = 1.0, fact_on_⊤::Bool = false,
        ⊤_k::Integer = k, p_mr::Real = 0.5, always_compress::Bool = false,
        always_merge::Bool = false, merge_branch::Real = 0.0, maxiter::Integer = 0)::StructProbCircuit
    memo = Dict{Tuple{Vtree, Bdd}, StructSumNode}()
    C = sample_psdd_r(ϕ, V, k, Dict{Int32, StructProbLiteralNode}(), randomize_weights, opts,
                      fact_on_⊤, ⊤_k, p_mr, always_compress, always_merge, memo, merge_branch > 0.0,
                      merge_branch, false, false)
    if maxiter > 0
        # Optionally grow the circuit by Strudel.
        loss(x) = heuristic_loss(x, D)
        C = struct_learn(C; primitives = [split_step], kwargs = Dict(split_step => (loss = loss,)),
                         maxiter, verbose = false)
    end
    !randomize_weights && estimate_parameters!(C, D; pseudocount)
    return C
end

function sample_psdd_r(ϕ::Bdd, V::Vtree, k::Integer, leaves::Dict{Int32, StructProbLiteralNode},
        randomize_weights::Bool, opts::SamplingOpts, fact_on_⊤::Bool, ⊤_k::Integer, p_mr::Real,
        always_compress::Bool, always_merge::Bool, repeats::Dict{Tuple{Vtree, Bdd}, StructSumNode},
        merge_this::Bool, merge_branch_pr::Float64, exact::Bool, anc_exact::Bool)::StructProbCircuit
    merge_branch = merge_branch_pr > 0.0
    if merge_branch
        r_p = (V, ϕ)
        if merge_this && (merge_branch_pr > rand()) && haskey(repeats, r_p) return repeats[r_p] end
    end
    if isleaf(V)
        v = convert(Int32, V.var)
        if v ∈ ϕ
            if is_lit(ϕ) return get_lit(to_lit(ϕ), V, leaves) end
            if is_⊤(ϕ|v) return get_lit(v, V, leaves) end
            if is_⊤(ϕ|-v) return get_lit(-v, V, leaves) end
            S = StructSumNode([get_lit(v, V, leaves), get_lit(-v, V, leaves)], V)
            if merge_branch repeats[r_p] = S end
            if randomize_weights S.log_probs = random_weights(2) end
            return S
        end
        S = StructSumNode([get_lit(v, V, leaves), get_lit(-v, V, leaves)], V)
        if merge_branch repeats[r_p] = S end
        if randomize_weights S.log_probs = random_weights(2) end
        return S
    elseif fact_on_⊤ && (is_⊤(ϕ) || isempty(intersect!(scopeset(ϕ), variables(V))))
        # When ϕ ≡ ⊤ and we want to simplify the circuit, fully factorize.
        left = sample_psdd_r(⊤, V.left, k, leaves, randomize_weights, opts, fact_on_⊤, ⊤_k, p_mr,
                             always_compress, always_merge, repeats, merge_branch, merge_branch_pr,
                             false, false)
        right = sample_psdd_r(⊤, V.right, k, leaves, randomize_weights, opts, fact_on_⊤, ⊤_k, p_mr,
                              always_compress, always_merge, repeats, merge_branch,
                              merge_branch_pr, false, false)
        S = StructSumNode([StructMulNode(left, right, V)], V)
        if merge_branch repeats[r_p] = S end
        return S
    elseif length(V.variables) == 2 # when |Sc(a,b)|=2, we have only two elements: (a,ϕ|a) and (¬a,ϕ|¬a).
        prime_var = convert(Int, V.left.var)
        C = Vector{StructProbCircuit}()
        # Left element.
        left_sub_ϕ = ϕ|prime_var
        if !is_⊥(left_sub_ϕ)
            left_prime = sample_psdd_r(bdd_var(prime_var), V.left, k, leaves,
                                       randomize_weights, opts, fact_on_⊤, ⊤_k, p_mr,
                                       always_compress, always_merge, repeats, merge_branch,
                                       merge_branch_pr, true, true)
            left_sub = sample_psdd_r(left_sub_ϕ, V.right, k, leaves, randomize_weights, opts,
                                     fact_on_⊤, ⊤_k, p_mr, always_compress, always_merge, repeats,
                                     merge_branch, merge_branch_pr, true, true)
            push!(C, StructMulNode(left_prime, left_sub, V))
        end
        right_sub_ϕ = ϕ|-prime_var
        if !is_⊥(right_sub_ϕ)
            right_prime = sample_psdd_r(bdd_var(-prime_var), V.left, k, leaves,
                                        randomize_weights, opts, fact_on_⊤, ⊤_k, p_mr,
                                        always_compress, always_merge, repeats, merge_branch,
                                        merge_branch_pr, true, true)
            right_sub = sample_psdd_r(right_sub_ϕ, V.right, k, leaves, randomize_weights, opts,
                                      fact_on_⊤, ⊤_k, p_mr, always_compress, always_merge, repeats,
                                      merge_branch, merge_branch_pr, true, true)
            push!(C, StructMulNode(right_prime, right_sub, V))
        end
        S = StructSumNode(C, V)
        if merge_branch repeats[r_p] = S end
        if randomize_weights S.log_probs = random_weights(length(C)) end
        return S
    end
    L, R = variables(V.left), variables(V.right)
    force_exact = (exact || anc_exact) && !is_⊤(ϕ)
    E = sample_partition(ϕ, L, p_mr, k, ⊤_k, force_exact)
    C = Vector{StructProbCircuit}()
    for (s, P) ∈ E
        # Single element with sub s.
        if length(P) == 1
            # Create (p, s) and move on.
            p = first(P)
            prime_node = sample_psdd_r(p, V.left, k, leaves, randomize_weights, opts, fact_on_⊤,
                                       ⊤_k, p_mr, always_compress, always_merge, repeats,
                                       merge_branch, merge_branch_pr, true, true)
            sub_node = sample_psdd_r(s, V.right, k, leaves, randomize_weights, opts, fact_on_⊤,
                                     ⊤_k, p_mr, always_compress, always_merge, repeats,
                                     merge_branch && isempty(intersect!(scopeset(s), R)),
                                     merge_branch_pr, false, anc_exact)
            push!(C, StructMulNode(prime_node, sub_node, V))
            continue
        end
        # Else, there are length(P) elements with sub s.
        if opts == plain # no compressions nor merges.
            for p ∈ P
                prime_node = sample_psdd_r(p, V.left, k, leaves, randomize_weights, opts, fact_on_⊤,
                                           ⊤_k, p_mr, always_compress, always_merge, repeats,
                                           merge_branch, merge_branch_pr, true, true)
                sub_node = sample_psdd_r(s, V.right, k, leaves, randomize_weights, opts, fact_on_⊤,
                                         ⊤_k, p_mr, always_compress, always_merge, repeats,
                                         merge_branch && isempty(intersect!(scopeset(s), R)),
                                         merge_branch_pr, false, anc_exact)
                push!(C, StructMulNode(prime_node, sub_node, V))
            end
            continue
        end
        n = length(P)
        do_compress = Int(opts) & Int(compress) > 0
        if do_compress
            if always_compress
                K = Tuple{Bdd, Vector{Int}}[(reduce(∨, P), Int[1])]
            else
                c = sample_row(n)
                if c == n # Compress everyone.
                    K = Tuple{Bdd, Vector{Int}}[(reduce(∨, P), Int[1])]
                elseif c == 1 # Compress no one.
                    K = Tuple{Bdd, Vector{Int}}[(P[i], Int[i]) for i ∈ 1:n]
                else # Compress c randomly selected elements.
                    K = Vector{Tuple{Bdd, Vector{Int}}}(undef, n-c+1)
                    comb = sample_comb(n, c)
                    j, ψ = 2, ⊥
                    @inbounds for i ∈ 1:n
                        if i ∈ comb
                            ψ = ψ ∨ P[i]
                        else
                            K[j] = (P[i], Int[j])
                            j += 1
                        end
                    end
                    K[1] = (ψ, Int[1])
                end
            end
        else
            K = Vector{Tuple{Bdd, Vector{Int}}}(undef, n)
            @inbounds for i ∈ 1:n K[i] = (P[i], Int[i]) end
        end
        m = length(K)
        if (Int(opts) & Int(merge)) > 0 # then merge.
            c = sample_row(m)
            if always_merge || c == m
                M = Vector{Int}(undef, m)
                @inbounds for i ∈ 2:m
                    M[i] = i
                    K[i] = (K[i][1], Int[])
                end
                M[1] = 1
                K[1] = (K[1][1], M)
            elseif c != 1 # if c == 1, no need to merge.
                comb = sample_comb(m, c)
                M = Vector{Int}(undef, c)
                j = 1
                for i ∈ 1:m
                    if i ∈ comb
                        M[j] = i
                        if j == 1 K[i] = (K[i][1], M)
                        else K[i] = (K[i][1], Int[]) end
                        j += 1
                    end
                end
            end
        end
        # Create subs.
        subs = Vector{StructProbCircuit}(undef, m)
        for (i, c) ∈ enumerate(K)
            M = c[2]
            if !isempty(M)
                sub_node = sample_psdd_r(s, V.right, k, leaves, randomize_weights, opts, fact_on_⊤,
                                         ⊤_k, p_mr, always_compress, always_merge, repeats,
                                         merge_branch && isempty(intersect!(scopeset(s), R)),
                                         merge_branch_pr, false, anc_exact)
                for j ∈ M subs[j] = sub_node end
            end
        end
        for (i, c) ∈ enumerate(K)
            α, M = c[1], c[2]
            prime_node = sample_psdd_r(α, V.left, k, leaves, randomize_weights, opts, fact_on_⊤,
                                       ⊤_k, p_mr, always_compress, always_merge, repeats,
                                       merge_branch, merge_branch_pr, true, true)
            prod_node = StructMulNode(prime_node, subs[i], V)
            push!(C, prod_node)
        end
    end
    S = StructSumNode(C, V)
    if merge_branch repeats[r_p] = S end
    if randomize_weights S.log_probs = random_weights(length(C)) end
    return S
end


##################################################################
# Source code taken from: https://github.com/jbytecode/MCCGA
# Credits to corresponding author.
##################################################################
include("utils.jl")

mutable struct MCCGA <: AbstractParameters
    N::Int # population size
    maxsamples::Int
    mutation::Float64
    probvector::Vector{Float64}
    use_local_search::Bool
end

"""
    MCCGA(;N, maxsamples)

### Parameters:

- `N` population size. Default is 100.
- `maxsamples` maximum number of samples. Default is 10000.

### Description

MCCGA method implements the Machine-coded genetic algorithms for real valued optimization problems.
The algorithm is based on the concept of a compact genetic algorithm but with a machine-coded representation
using IEEE-754 floating point encoding standard. In the first stage of the algorithm, `maxsamples` number
of samples are generated within the range of function domain. This process is required to obtain a vector
of probabilities for each single bit of the IEEE-754 representation. In classical CGAs, the initial vector
of probabilities is generated using the constant probability of 0.5 whereas in MCCGA, the probability of ith bit having
a value of 1 depends on the function domain. The second step performs a CGA search but with IEEE-754 bits again. Since
CGA does not use a population of solutions but a single vector of probabilities, the parameter `N` does not really mean
number of solutions. Instead, it means the amount of mutation in each iteration, e.g. 1/N.
In the last stage, a local search is performed for fine-tuning. In this implementation, `Optim` package is required.

### References

- Harik, G. R., Lobo, F. G., & Goldberg, D. E. (1999). The compact genetic algorithm. IEEE transactions on evolutionary computation, 3(4), 287-297.
- Satman, M. H. & Akadal, E. (2020). Machine Coded Compact Genetic Algorithms for Real Parameter Optimization Problems . Alphanumeric Journal , 8 (1) , 43-58 . DOI: 10.17093/alphanumeric.576919
- Mehmet Hakan Satman, Emre Akadal, Makine Kodlu Hibrit Kompakt Genetik Algoritmalar Optimizasyon Yöntemi, Patent, TR, 2022/01, 2018-GE-510239

### Example

```jldoctest

julia> import Optim

julia> f, bounds, solutions = Metaheuristics.TestProblems.rastrigin();

julia> result = optimize(f, bounds, MCCGA())

+=========== RESULT ==========+
  iteration: 42833
    minimum: 0
  minimizer: [-5.677669786379456e-17, 2.7942451898022582e-39, -3.60925916059986e-33, -6.609510861086017e-34, 2.998586759675011e-32, -1.8825832500775007e-38, -3.0729484147585938e-31, 1.7675578057632446e-38, 5.127944874215823e-16, 1.9623770480857484e-19]
    f calls: 86404
 total time: 3.2839 s
+============================+
```
"""
function MCCGA(;
        N = 100,
        maxsamples =10_000,
        mutation = 1 / N,
        use_local_search = true,
        information = Information(),
        options = Options()
    )

    parameters = MCCGA(N, maxsamples, mutation, [], use_local_search)

    Algorithm(
        parameters,
        information = information,
        options = options,
    )
end

function initialize!(
        status,
        parameters::MCCGA,
        problem,
        information,
        options,
        args...;
        kargs...
    )

    if options.iterations == 0
        options.iterations = 100_000_000
    end

    if options.f_calls_limit == 0
        options.f_calls_limit = options.iterations * parameters.N + 1
    end

    lower = problem.bounds[1,:]
    upper = problem.bounds[2,:]

    parameters.probvector = initialprobs(lower, upper, maxsamples = parameters.maxsamples)

    # sample a vector to create an initial State
    x = sample(parameters.probvector) |> floats
    initial_sol = create_solution(x, problem)
    return State(initial_sol, [initial_sol for i in 1:parameters.N])

end

function update_state!(
        status,
        parameters::MCCGA,
        problem,
        information,
        options,
        args...;
        kargs...
    )

    probvector = parameters.probvector
    chsize = length(probvector)

    # parents
    ch1 = sample(probvector)
    ch2 = sample(probvector)

    # evaluate cost function
    sol1 = create_solution(floats(ch1), problem)
    sol2 = create_solution(floats(ch2), problem)

    sol_winner = sol1
    winner = ch1
    loser  = ch2

    # check if ch2 (sol2) is better that ch1 (sol1)
    if is_better(sol2, sol1)
        sol_winner = sol2
        winner = ch2
        loser = ch1
    end

    # save in population the winner (informative only)
    # this is not used in the algorithm
    status.population[1 + status.iteration % parameters.N] = sol_winner

    # save best solution found so far
    is_better(sol_winner, status.best_sol) &&  (status.best_sol = sol_winner)

    for i = 1:chsize
        if winner[i] != loser[i]
            if winner[i] == 1
                probvector[i] += parameters.mutation
            else
                probvector[i] -= parameters.mutation
            end
            if probvector[i] > 1
                probvector[i] = 1
            elseif probvector[i] < 0
                probvector[i] = 0
            end
        end
    end
end

function final_stage!(
        status::State{xf_indiv}, # unconstrained case
        parameters::MCCGA,
        problem,
        information,
        options,
        args...;
        kargs...
    )

    if !parameters.use_local_search
        status.final_time = time()
        return
    end

    costfunction(x) = evaluate(x, problem)

    probvector = parameters.probvector
    sampledvector = sample(probvector)

    # initial solution for the local search
    initial_solution = floats(sampledvector)

    options.debug && @info "Running NelderMead..."
    local_result = Optim.optimize(costfunction, initial_solution, Optim.NelderMead())
    # display Nelder-Mead result
    options.debug && display(local_result)
    options.debug && @info "NelderMead done!"

    # save best solution found so far!
    sol = create_child(local_result.minimizer, local_result.minimum)

    status.final_time = time()
    status.f_calls = problem.f_calls

    if is_better(sol, status.best_sol)
        status.best_sol = sol
    end

    return
end

function final_stage!(
        status,
        parameters::MCCGA,
        problem,
        information,
        options,
        args...;
        kargs...
    )

    status.final_time = time()
    # nothing to do for constrained or multi-objective
    @warn "MCCGA has been designed for unconstrained problems (local search not performed)."
end

function stop_criteria!(status, parameters::MCCGA, problem, information, options)
    # check budget limitation
    if status.stop
        return
    end

    mutation = parameters.mutation
    status.stop = all(x -> (x <= mutation) || (x >= 1.0 - mutation), parameters.probvector)

    status.stop && options.debug && @info "MCCGA stopping criteria is met."
end


using LinearAlgebra

struct NullLattice <: Lattice end

mutable struct Branching
  k::Vector{Int}
  probs::Vector{Vector{Float64}}
  kMin::Int
  jMin::Int
  kMax::Int
  jMax::Int
end

function Branching()
  probs = Vector{Vector{Float64}}(undef, 3)
  probs[1] = zeros(0)
  probs[2] = zeros(0)
  probs[3] = zeros(0)
  return Branching(zeros(Int, 0), probs, typemax(Int), typemax(Int), typemin(Int), typemin(Int))
end

mutable struct TrinomialTree{S <: StochasticProcess} <: AbstractTree
  process::S
  timeGrid::TimeGrid
  dx::Vector{Float64}
  branchings::Vector{Branching}
  isPositive::Bool
end

function TrinomialTree(process::S, timeGrid::TimeGrid, isPositive::Bool = false) where {S <: StochasticProcess}
  x0 = process.x0
  dx = zeros(length(timeGrid.times))
  nTimeSteps = length(timeGrid.times) - 1
  jMin = 0
  jMax = 0
  branchings = Vector{Branching}(undef, nTimeSteps)

  for i = 1:nTimeSteps
    t = timeGrid.times[i]
    dt = timeGrid.dt[i]

    # Variance must be independent of x
    v2 = variance(process, t, 0.0, dt)
    v = sqrt(v2)
    dx[i+1] = v * sqrt(3.0)

    branching = Branching()

    @simd for j =jMin:jMax
      @inbounds x = x0 + j * dx[i]
      m = expectation(process, t, x, dt)
      @inbounds temp = round(Int, floor((m - x0) / dx[i+1] + 0.5))

      if isPositive
        @inbounds while (x0 + (temp - 1) * dx[i + 1] <= 0)
          temp += 1
        end
      end

      @inbounds e = m - (x0 + temp * dx[i + 1])
      e2 = e * e
      e3 = e * sqrt(3.0)

      p1 = (1.0 + e2 / v2 - e3 / v) / 6.0
      p2 = (2.0 - e2 / v2) / 3.0
      p3 = (1.0 + e2 / v2 + e3 / v) / 6.0

      add!(branching, temp, p1, p2, p3)
    end

    @inbounds branchings[i] = branching # check if we need copy

    jMin = branching.jMin
    jMax = branching.jMax
  end

  return TrinomialTree{S}(process, timeGrid, dx, branchings, isPositive)
end

function rebuild_tree!(tt::TrinomialTree, timeGrid::TimeGrid)
  x0 = tt.process.x0
  dx = zeros(length(timeGrid.times))
  nTimeSteps = length(timeGrid.times) - 1
  jMin = 0
  jMax = 0
  branchings = Vector{Branching}(undef, nTimeSteps)

  for i = 1:nTimeSteps
    t = timeGrid.times[i]
    dt = timeGrid.dt[i]

    # Variance must be independent of x
    v2 = variance(tt.process, t, 0.0, dt)
    v = sqrt(v2)
    dx[i+1] = v * sqrt(3.0)

    branching = Branching()

    @simd for j =jMin:jMax
      @inbounds x = x0 + j * dx[i]
      m = expectation(tt.process, t, x, dt)
      @inbounds temp = round(Int, floor((m - x0) / dx[i+1] + 0.5))

      if tt.isPositive
        @inbounds while (x0 + (temp - 1) * dx[i + 1] <= 0)
          temp += 1
        end
      end

      @inbounds e = m - (x0 + temp * dx[i + 1])
      e2 = e * e
      e3 = e * sqrt(3.0)

      p1 = (1.0 + e2 / v2 - e3 / v) / 6.0
      p2 = (2.0 - e2 / v2) / 3.0
      p3 = (1.0 + e2 / v2 + e3 / v) / 6.0

      add!(branching, temp, p1, p2, p3)
    end

    @inbounds branchings[i] = branching # check if we need copy

    jMin = branching.jMin
    jMax = branching.jMax
  end
  tt.dx = dx
  tt.timeGrid = timeGrid
  tt.branchings = branchings

  return tt
end

mutable struct TreeLattice1D{T <: TreeLattice} <: TreeLattice
  tg::TimeGrid
  impl::T
  statePrices::Vector{Vector{Float64}}
  n::Int
  statePricesLimit::Int
end

function TreeLattice1D(tg::TimeGrid, n::Int, impl::T) where {T <: TreeLattice}
  statePrices = Vector{Vector{Float64}}(undef, 1)
  statePrices[1] = ones(1)

  statePricesLimit = 1

  return TreeLattice1D{T}(tg, impl, statePrices, n, statePricesLimit)
end

function get_grid(lat::TreeLattice1D, t::Float64)
  i = return_index(lat.tg, t)
  grid = zeros(get_size(lat.impl, i))

  @simd for j in eachindex(grid)
    @inbounds grid[j] = get_underlying(lat.impl, i, j)
  end

  return grid
end

mutable struct TreeLattice2D{T <: TreeLattice} <: TreeLattice
  tg::TimeGrid
  impl::T
  statePrices::Vector{Vector{Float64}}
  n::Int
  statePricesLimit::Int
  tree1::TrinomialTree
  tree2::TrinomialTree
  m::Matrix{Float64}
  rho::Float64
end

function TreeLattice2D(tree1::TrinomialTree, tree2::TrinomialTree, correlation::Float64, impl::T) where {T <: TreeLattice}
  tg = tree1.timeGrid
  statePrices = Vector{Vector{Float64}}(undef, 1)
  statePrices[1] = ones(1)

  statePricesLimit = 1
  branch_num = branches(TrinomialTree)
  n = branch_num ^ 2
  m = zeros(branch_num, branch_num)
  rho = abs(correlation)

  if correlation < 0.0 && branch_num == 3
    m[1,1] = -1.0
    m[2,1] = -4.0
    m[3,1] =  5.0
    m[1,2] = -4.0
    m[2,2] =  8.0
    m[3,2] = -4.0
    m[1,3] =  5.0
    m[2,3] = -4.0
    m[3,3] = -1.0
  else
    m[1,1] =  5.0
    m[2,1] = -4.0
    m[3,1] = -1.0
    m[1,2] = -4.0
    m[2,2] =  8.0
    m[3,2] = -4.0
    m[1,3] = -1.0
    m[2,3] = -4.0
    m[3,3] =  5.0
  end

  return TreeLattice2D{T}(tg, impl, statePrices, n, statePricesLimit, tree1, tree2, m, rho)
end

get_size(tr::TreeLattice2D, i::Int) = get_size(tr.tree1, i) * get_size(tr.tree2, i)

function descendant(tr::TreeLattice2D, i::Int, idx::Int, branch::Int)
  modulo = get_size(tr.tree1, i)
  new_idx = idx - 1
  new_branch = branch - 1

  index1 = (new_idx % modulo) + 1
  index2 = (round(Int, floor(new_idx / modulo))) + 1

  branch1 = (new_branch % branches(TrinomialTree)) + 1
  branch2 = (round(Int, floor(new_branch / branches(TrinomialTree)))) + 1

  modulo = get_size(tr.tree1, i + 1)
  return ((descendant(tr.tree1, i, index1, branch1) - 1) + (descendant(tr.tree2, i, index2, branch2) - 1) * modulo) + 1
end

function probability(tr::TreeLattice2D, i::Int, idx::Int, branch::Int)
  modulo = get_size(tr.tree1, i)
  new_idx = idx - 1
  new_branch = branch - 1

  index1 = (new_idx % modulo) + 1
  index2 = (round(Int, floor(new_idx / modulo))) + 1

  branch1 = (new_branch % branches(TrinomialTree)) + 1
  branch2 = (round(Int, floor(new_branch / branches(TrinomialTree)))) + 1

  # println("$(index1) $(index2) $(branch1) $(branch2) $(idx) $(branch) $(modulo)")

  prob1 = probability(tr.tree1, i, index1, branch1)
  prob2 = probability(tr.tree2, i, index2, branch2)

  # println("$(prob1) $(prob2)")

  return prob1 * prob2 + tr.rho * tr.m[branch2,branch1] / 36.0 # this 36 could depend on the branches(TrinomialTree)
end

function get_state_prices!(t::TreeLattice, i::Int)
  if i > t.statePricesLimit
    compute_state_prices!(t, i)
  end

  return t.statePrices[i]
end

function compute_state_prices!(t::TreeLattice, until::Int)
  @inbounds @simd for i = t.statePricesLimit:until - 1
    push!(t.statePrices, zeros(get_size(t.impl, i + 1)))
    for j = 1:get_size(t.impl, i)
      disc = discount(t.impl, i, j)
      statePrice = t.statePrices[i][j]
      for l = 1:t.n
        t.statePrices[i + 1][descendant(t.impl, i, j, l)] += statePrice * disc * probability(t.impl, i, j, l)
      end
    end
  end

  t.statePricesLimit = until

  return t
end

function initialize!(lattice::TreeLattice, asset::DiscretizedAsset, t::Float64)
  #i = findfirst(lattice.tg.times .>= t)
  i = return_index(lattice.tg, t)
  set_time!(asset, t)
  reset!(asset, get_size(lattice.impl, i))
end

function rollback!(lattice::TreeLattice, asset::DiscretizedAsset, t::Float64)
  partial_rollback!(lattice, asset, t)
  adjust_values!(asset)

  return asset
end

function partial_rollback!(lattice::TreeLattice, asset::DiscretizedAsset, t::Float64)
  from = asset.common.time

  if QuantLib.Math.is_close(from, t)
    return
  end

  # iFrom = findfirst(lattice.tg.times .>= from)
  # iTo = findfirst(lattice.tg.times .>= t)
  iFrom = return_index(lattice.tg, from)
  iTo = return_index(lattice.tg, t)

  @simd for i = iFrom-1:-1:iTo
    newVals = zeros(get_size(lattice.impl, i))
    step_back!(lattice, i, asset.common.values, newVals)
    @inbounds asset.common.time = lattice.tg.times[i]
    asset.common.values = newVals
    # println("newVals before: ", newVals)
    if i != iTo
      adjust_values!(asset)
    end
    # println("newVals after: ", newVals)
  end

  return asset
end

function present_value(lattice::TreeLattice, asset::DiscretizedAsset)
  i = findfirst(lattice.tg.times .>= asset.common.time)
  return dot(asset.common.values, get_state_prices!(lattice, i))
end

function step_back!(lattice::TreeLattice, i::Int, vals::Vector{Float64}, newVals::Vector{Float64})
  # pragma omp parallel for
  for j = 1:get_size(lattice.impl, i)
    val = 0.0
    @simd for l = 1:lattice.n
      @inbounds val += probability(lattice.impl, i, j, l) * vals[descendant(lattice.impl, i, j, l)]
    end
    @inbounds val *= discount(lattice.impl, i, j)
    @inbounds newVals[j] = val
  end

  return newVals
end

get_size(b::Branching) = b.jMax - b.jMin + 1

function add!(branch::Branching, k::Int, p1::Float64, p2::Float64, p3::Float64)
  push!(branch.k, k)
  push!(branch.probs[1], p1)
  push!(branch.probs[2], p2)
  push!(branch.probs[3], p3)

  # maintain invariants
  branch.kMin = min(branch.kMin, k)
  branch.jMin = branch.kMin - 1
  branch.kMax = max(branch.kMax, k)
  branch.jMax = branch.kMax + 1

  return branch
end

descendant(b::Branching, idx::Int, branch::Int) = b.k[idx] - b.jMin - 1 + branch
probability(b::Branching, idx::Int, branch::Int) = b.probs[branch][idx]
branches(::Type{TrinomialTree}) = 3
get_size(t::TrinomialTree, i::Int) = i == 1 ? 1 : get_size(t.branchings[i-1])

function get_underlying(t::TrinomialTree, i::Int, idx::Int)
  if i == 1
    return t.process.x0
  else
    @inbounds return t.process.x0 + (t.branchings[i - 1].jMin + (idx - 1)) * t.dx[i]
  end
end

descendant(t::TrinomialTree, i::Int, idx::Int, branch::Int) = descendant(t.branchings[i], idx, branch)

probability(t::TrinomialTree, i::Int, j::Int, b::Int) = probability(t.branchings[i], j, b)


using BilevelJuMP
using Test, Gurobi, MathOptInterface, JuMP, Dualization
using MathOptFormat

const MOI  = MathOptInterface
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges
const MOIT = MathOptInterface.Test

# original problem
# min -4x -3y
# s.t.
#      2x + y <= 4
#       x +2y <= 4
#       x     >= 0
#           y >= 0
#
# sol x = y = 4/3
# obj = 28/3

# bilevel modifications
# min -4x -3y
# s.t.
# y = argmin_y y
#      2x + y <= 4
#       x +2y <= 4
#       x     >= 0
#           y >= 0
#
# sol: x = 2, y = 0
# obj_upper = -8
# obj_lower =  0

@testset "Simple LP" begin
    optimizer = Gurobi.Optimizer()
    bridged = MOIB.full_bridge_optimizer(optimizer, Float64)
    MOI.empty!(bridged)
    @test MOI.is_empty(bridged)

    # add 10 variables - only diagonal is relevant
    X = MOI.add_variables(bridged, 2)

    c1 = MOIU.normalize_and_add_constraint(bridged,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(2.0, X[1]),
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    c2 = MOIU.normalize_and_add_constraint(bridged,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1]),
            MOI.ScalarAffineTerm(2.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    b1 = MOIU.normalize_and_add_constraint(bridged,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1])
        ], 0.0), MOI.GreaterThan(0.0))

    b2 = MOIU.normalize_and_add_constraint(bridged,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.GreaterThan(0.0))

    MOI.set(bridged,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], [X[1], X[2]]), 0.0)
        )
    MOI.set(bridged, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(bridged)

    obj = MOI.get(bridged, MOI.ObjectiveValue())

    @test obj ≈ -9.33333 atol = 1e-2

    Xr = MOI.get(bridged, MOI.VariablePrimal(), X)

    @test Xr ≈ [1.3333, 1.3333] atol = 1e-2

end

@testset "Simple BLP" begin
    optimizer = Gurobi.Optimizer()
    bridged = MOIB.full_bridge_optimizer(optimizer, Float64)
    MOI.empty!(bridged)
    @test MOI.is_empty(bridged)

    X = MOI.add_variables(bridged, 2)

    MOI.set(bridged, MOI.VariableName(), X[1], "x_u")
    MOI.set(bridged, MOI.VariableName(), X[2], "y_u")

    MOI.set(bridged,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([-4.0, -3.0], [X[1], X[2]]), 0.0)
        )
    MOI.set(bridged, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    optimizer2 = Gurobi.Optimizer()
    bridged2 = MOIB.full_bridge_optimizer(optimizer2, Float64)
    MOI.empty!(bridged2)
    @test MOI.is_empty(bridged2)

    X2 = MOI.add_variables(bridged2, 2)
    MOI.set(bridged2, MOI.VariableName(), X2[1], "x_l")
    MOI.set(bridged2, MOI.VariableName(), X2[2], "y_l")

    c1 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(2.0, X2[1]),
            MOI.ScalarAffineTerm(1.0, X2[2])
        ], 0.0), MOI.LessThan(4.0))
    MOI.set(bridged2, MOI.ConstraintName(), c1, "l_c1")

    c2 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X2[1]),
            MOI.ScalarAffineTerm(2.0, X2[2])
        ], 0.0), MOI.LessThan(4.0))
    MOI.set(bridged2, MOI.ConstraintName(), c2, "l_c2")

    b1 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X2[1])
        ], 0.0), MOI.GreaterThan(0.0))
    MOI.set(bridged2, MOI.ConstraintName(), b1, "l_b1")

    b2 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X2[2])
        ], 0.0), MOI.GreaterThan(0.0))
    MOI.set(bridged2, MOI.ConstraintName(), b2, "l_b12")

    MOI.set(bridged2,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0], [X2[2]]), 0.0)
        )
    MOI.set(bridged2, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    links = Dict(X[1] => X2[1], X[2] => X2[2])

    parametric = [X2[1]]

    blp, _, _ = BilevelJuMP.build_bilivel(bridged, bridged2, links, parametric)

    # MOI.optimize!(bridged)

    # obj = MOI.get(bridged, MOI.ObjectiveValue())

    # @test obj ≈ -9.33333 atol = 1e-2

    # Xr = MOI.get(bridged, MOI.VariablePrimal(), X)

    # @test Xr ≈ [1.3333, 1.3333] atol = 1e-2

    # lp_model = MathOptFormat.MOF.Model()
    # MOI.copy_to(lp_model, blp)
    # MOI.write_to_file(lp_model, "my_model.LP")
    # @show pwd()

    optimizer = Gurobi.Optimizer()
    bridged3 = MOIB.full_bridge_optimizer(optimizer, Float64)
    MOI.empty!(bridged3)
    @test MOI.is_empty(bridged3)
    MOI.copy_to(bridged3, blp)

    MOI.optimize!(bridged3)

end

@testset "Simple BLP2" begin
    optimizer = Gurobi.Optimizer()
    bridged = MOIB.full_bridge_optimizer(optimizer, Float64)
    MOI.empty!(bridged)
    @test MOI.is_empty(bridged)

    x = MOI.add_variable(bridged)
    y = MOI.add_variable(bridged)

    MOI.set(bridged, MOI.VariableName(), x, "x_u")
    MOI.set(bridged, MOI.VariableName(), y, "y_u")

    MOI.set(bridged,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 3.0], [x, y]), 0.0)
        )
    MOI.set(bridged, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    optimizer2 = Gurobi.Optimizer()
    bridged2 = MOIB.full_bridge_optimizer(optimizer2, Float64)
    MOI.empty!(bridged2)
    @test MOI.is_empty(bridged2)

    x2 = MOI.add_variable(bridged2)
    y2 = MOI.add_variable(bridged2)

    MOI.set(bridged2, MOI.VariableName(), x2, "x_l")
    MOI.set(bridged2, MOI.VariableName(), y2, "y_l")

    c1 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2),
            MOI.ScalarAffineTerm(1.0, y2)
        ], 0.0), MOI.LessThan(8.0))
    MOI.set(bridged2, MOI.ConstraintName(), c1, "l_c1")

    c2 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2),
            MOI.ScalarAffineTerm(4.0, y2)
        ], 0.0), MOI.GreaterThan(8.0))
    MOI.set(bridged2, MOI.ConstraintName(), c2, "l_c2")

    c3 = MOIU.normalize_and_add_constraint(bridged2,
    MOI.ScalarAffineFunction([
        MOI.ScalarAffineTerm(1.0, x2),
        MOI.ScalarAffineTerm(2.0, y2)
    ], 0.0), MOI.LessThan(13.0))
    MOI.set(bridged2, MOI.ConstraintName(), c3, "l_c3")

    b1 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2)
        ], 0.0), MOI.GreaterThan(1.0))
    MOI.set(bridged2, MOI.ConstraintName(), b1, "l_b1")

    b2 = MOIU.normalize_and_add_constraint(bridged2,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2)
        ], 0.0), MOI.LessThan(6.0))
    MOI.set(bridged2, MOI.ConstraintName(), b2, "l_b12")

    MOI.set(bridged2,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([-1.0], [y2]), 0.0)
        )
    MOI.set(bridged2, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    links = Dict(x => x2, y => y2)

    parametric = [x2]

    blp, _, _ = BilevelJuMP.build_bilivel(bridged, bridged2, links, parametric)

    # MOI.optimize!(bridged)

    # obj = MOI.get(bridged, MOI.ObjectiveValue())

    # @test obj ≈ -9.33333 atol = 1e-2

    # Xr = MOI.get(bridged, MOI.VariablePrimal(), X)

    # @test Xr ≈ [1.3333, 1.3333] atol = 1e-2

    # lp_model = MathOptFormat.MOF.Model()
    # MOI.copy_to(lp_model, blp)
    # MOI.write_to_file(lp_model, "my_model.LP")
    # @show pwd()

    optimizer = Gurobi.Optimizer()
    bridged3 = MOIB.full_bridge_optimizer(optimizer, Float64)
    MOI.empty!(bridged3)
    @test MOI.is_empty(bridged3)
    MOI.copy_to(bridged3, blp)

    MOI.optimize!(bridged3)

end

@testset "Simple BLP JuMP" begin

    model = BilevelModel()

    @variable(UpperToLower(model), x)
    @variable(LowerToUpper(model), y)

    @objective(Upper(model), Min, -4x -3y)

    @objective(Lower(model), Min, y)

    @constraint(Lower(model), 2x+y <= 4)
    @constraint(Lower(model), x+2y <= 4)
    @constraint(Lower(model), x >= 0)
    @constraint(Lower(model), y >= 0)

    optimize!(model, Gurobi.Optimizer())

    primal_status(model)

    termination_status(model)

    @test objective_value(model) == -8

    @test value(x) ==  2
    @test value(y) ==  0

end

@testset "Simple BLP2 JuMP" begin

    model = BilevelModel()

    @variable(UpperToLower(model), x)
    @variable(LowerToUpper(model), y)

    @objective(Upper(model), Min, x + 3y)

    @objective(Lower(model), Min, -y)

    @constraint(Lower(model), x+y <= 8)
    @constraint(Lower(model), x+4y >= 8)
    @constraint(Lower(model), x+2y <= 13)
    @constraint(Lower(model), x >= 1)
    @constraint(Lower(model), x <= 6)
    @constraint(Lower(model), y >= 0)

    optimize!(model, Gurobi.Optimizer())

    primal_status(model)

    termination_status(model)

    @test objective_value(model) == 12

    @test value(x) == 6
    @test value(y) == 2

end

# alternative sintax
# @constraint(model, lower, [y] \in ArgMin([x]))


function logmvbeta(vals)
    r = 0.0
    for v in vals
        r += lgamma(v)
    end
    r -= lgamma(sum(vals))

    return r
end

function λ_objective(λ::Vector{Float64}, ∇λ::Vector{Float64},
        ν::Vector{Float64}, Ndivζ::Vector{Float64}, sumθ::Vector{Float64},
        μ::Vector{Float64}, invΣ::Matrix{Float64})

    diff = λ .- μ
    Eeη = exp.(λ .+ 0.5ν)

    if length(∇λ) > 0
        ∇λ .= -invΣ * diff .+ sumθ .- Ndivζ .* Eeη
    end

    return -0.5 * (diff' * invΣ * diff)[1] + sum(λ .* sumθ) - sum(Ndivζ .* Eeη)
end

function ν_objective(ν::Vector{Float64}, ∇ν::Vector{Float64},
        λ::Vector{Float64}, Ndivζ::Vector{Float64}, μ::Vector{Float64},
        invΣ::Matrix{Float64})

    Eeη = exp.(λ .+ 0.5ν)

    if length(∇ν) > 0
        ∇ν .= -0.5diag(invΣ) .- (Ndivζ / 2) .* Eeη .+ (1 ./ (2ν))
    end

    return -0.5 * tr(diagm(0 => ν) * invΣ) - sum(Ndivζ .* Eeη) + sum(log.(ν)) / 2
end

function α_objective(α::Vector{Float64}, ∇α::Vector{Float64},
        sum_Elnϕ::Float64, K::Int, V::Int)

    if length(∇α) > 0
        ∇α[1] = K * V * (digamma(V * α[1]) - digamma(α[1])) + sum_Elnϕ
    end

    return K * (lgamma(V * α[1]) - V * lgamma(α[1])) + α[1] * sum_Elnϕ
end

function check_convergence(metric::Vector{Vector{Float64}}; tol=1e-4)
    reldiff = maximum(abs.(metric[end - 1] .- metric[end]) ./ abs.(metric[end]))
    return reldiff < tol
end

function check_convergence(metric::Vector{Float64}; tol=1e-4)
    reldiff = maximum(abs(metric[end - 1] - metric[end]) / abs(metric[end]))
    return reldiff < tol
end


using Test
using IncrementalInference

##

@testset "test generic functions on distributions" begin
##

@test getDimension(Uniform(0,1)) == 1
@test getDimension(Normal()) == 1
@test getDimension(MvNormal([1;1;0.1])) == 3

p = manikde!(ContinuousScalar, [randn(1) for i in 1:100])
@test getDimension(p) == 1

##
end


using StatsBase
import DSP
import DSP: Periodograms
import LombScargle
export estimate_period

"""
    estimate_period(v::Vector, method, t=0:length(v)-1; kwargs...)
Estimate the period of the signal `v`, with accompanying time vector `t`,
using the given `method`.

If `t` is an AbstractArray, then it is iterated through to ensure that it's
evenly sampled (if necessary for the algorithm).  To avoid this, you can pass
any `AbstractRange`, like a `UnitRange` or a `LinRange`, which are defined to be
evenly sampled.

## Methods requiring evenly sampled data

These methods are faster, but some are error-prone.

* `:periodogram` or `:pg`: Use the fast Fourier transform to compute a
   periodogram (power-spectrum) of the given data.  Data must be evenly sampled.

* `:multitaper` or `mt`: The multitaper method reduces estimation bias by using
  multiple independent estimates from the same sample. Data tapers are then
  windowed and the power spectra are obtained.  Available keywords follow:
  `nw` is the time-bandwidth product, and `ntapers` is the number of tapers.
  If `window` is not specified, the signal is tapered with `ntapers` discrete
  prolate spheroidal sequences with time-bandwidth product `nw`.
  Each sequence is equally weighted; adaptive multitaper is not (yet) supported.
  If `window` is specified, each column is applied as a taper. The sum of
  periodograms is normalized by the total sum of squares of `window`.

* `:autocorrelation` or `:ac`: Use the autocorrelation function (AC). The value
  where the AC first comes back close to 1 is the period of the signal. The
  keyword `L = length(v)÷10` denotes the length of the AC (thus, given the
  default setting, this method will fail if there less than 10 periods in the
  signal). The keyword `ϵ = 0.2` (`\\epsilon`) means that `1-ϵ` counts as "1" for the AC.

## Methods not requiring evenly sampled data

These methods tend to be slow, but versatile and low-error.

* `:lombscargle` or `:ls`: Use the Lomb-Scargle algorithm to compute a
  periodogram.  The advantage of the Lomb-Scargle method is that it does not
  require an equally sampled dataset and performs well on undersampled datasets.
  Constraints have been set on the period, since Lomb-Scargle tends to have
  false peaks at very low frequencies.  That being said, it's a very flexible
  method.  It is extremely customizable, and the keyword arguments that can be
  passed to it are given [in the documentation](https://juliaastro.github.io/LombScargle.jl/stable/index.html#LombScargle.plan).

*  `:zerocrossing` or `:zc`: Find the zero crossings of the data, and use the
  average difference between zero crossings as the period.  This is a naïve
  implementation, with only linear interpolation; however, it's useful as a
  sanity check.  The keyword `line` controls where the "crossing point" is.
  It deffaults to `mean(v)`.

For more information on the periodogram methods, see the documentation of
DSP.jl and LombScargle.jl.
"""
function estimate_period(v, method, t = 0:length(v)-1; kwargs...)
    @assert length(v) == length(t)

    even_methods  = (:periodogram, :pg, :multitaper, :mt, :autocorrelation, :ac)
    other_methods = (:lombscargle, :ls, :zerocrossing, :zc)
    if method ∉ even_methods && method ∉ other_methods
        error("Unknown method (`$method`) given to `estimate_period`.")
    end

    period = if method ∈ even_methods
                isevenlysampled(t) ||
                error("Your time data was not evenly sampled,
                and the algorithm `$method` requires evenly sampled data.")

                if method == :periodogram || method == :pg
                    _periodogram_period(v, t; kwargs...)
                elseif method == :multitaper || method == :mt
                    _mt_period(v, t; kwargs...)
                elseif method == :autocorrelation || method == :ac
                    _ac_period(v, t; kwargs...)
                end
            else
                if method == :lombscargle || method == :ls
                    _ls_period(v, t; kwargs...)
                elseif method == :zerocrossing || method == :zc
                    _zc_period(v, t; kwargs...)
                end
            end

    return period
end

function isevenlysampled(t::AbstractVector)
    for i in 2:length(t)-1
        if !(t[nextind(t, i)] - t[i] ≈ t[i] - t[prevind(t, i)])
            return false
            break
        end
    end
    return true
end

# AbstractRanges are defined to be evenly sampled.
isevenlysampled(::AbstractRange) = true

################################################################################
#                           Autocorrelation Function                           #
################################################################################
"""
    _ac_period(v, t; ϵ = 0.2, L = length(v)÷10)

Use the autocorrelation function (AC). The value where the AC first
comes back close to 1 is the period of the signal. The keyword
`L = length(v)÷10` denotes the length of the AC (thus, given the default
setting, this method will fail if there less than 10 periods in the signal).
The keyword `ϵ = 0.2` means that `1-ϵ` counts as "1" for the AC.
"""
function _ac_period(v, t; ϵ = 0.2, L = length(v)÷10)
    err = "The autocorrelation did not become close to 1."
    ac = autocor(v, 0:L)
    j = 0
    local_maxima = findall(i -> ac[i-1] < ac[i] ≥ ac[i+1], 2:length(ac)-1)
    isempty(local_maxima) && error("AC did not have any local maxima")
    # progressively scan the local maxima and find the first within 1-ϵ
    for i in local_maxima
        if ac[i] > 1-ϵ
            j = i
            break
        end
    end
    j == 0 && error("No local maximum of the AC exceeded 1-ϵ")
    # since now it holds that ac[j] is a local maximum within 1-ϵ:
    period = t[j+1] - t[1]
    return period
end

################################################################################
#                                 Periodograms                                 #
################################################################################

########################################
#          Basic Periodogram           #
########################################

"""
    _periodogram_period(v, t; kwargs...)

Use the fast Fourier transform to compute a periodogram (power-spectrum) of the
given data.  Data must be evenly sampled.
"""
function _periodogram_period(v, t; kwargs...)

    p = Periodograms.periodogram(v; fs = length(t)/(t[end] - t[1]), kwargs...)

    return 1 / Periodograms.freq(p)[findmax(Periodograms.power(p))[2]]

end

########################################
#        Multitaper periodogram        #
########################################

"""
    _mt_period(v, t;
                    nw = 4, ntapers = DSP.ceil(2nw)-1,
                    window = DSP.dpss(length(s), nw, ntapers), kwargs...)

The multitaper method reduces estimation bias by obtaining multiple independent
estimates from the same sample. Data tapers are then windowed and the power
spectra are obtained.
`nw` is the time-bandwidth product, and `ntapers` is the number of tapers.
If `window` is not specified, the signal is tapered with `ntapers` discrete
prolate spheroidal sequences with time-bandwidth product `nw`.
Each sequence is equally weighted; adaptive multitaper is not (yet) supported.
If `window` is specified, each column is applied as a taper. The sum of
periodograms is normalized by the total sum of squares of `window`.
"""
function _mt_period(v, t;
                        nw = 4,
                        ntapers::Integer = ceil(Int, 2nw)-1,
                        window = DSP.dpss(length(v), nw, ntapers),
                        kwargs...
                    )

    Sys.WORD_SIZE == 32 &&
        error("multitaper method doesn't work on 32-bit systems. Sorry!")

    p = Periodograms.mt_pgram(v;
                                 fs = length(t)/(t[end] - t[1]),
                                 nw = nw,
                                 ntapers = ntapers,
                                 window = window,
                                 kwargs...
                             )

    return 1 / Periodograms.freq(p)[findmax(Periodograms.power(p))[2]]
end

########################################
#       Lomb-Scargle periodogram       #
########################################

"""
    _ls_period(v, t;
                minimum_period::Real = 2 * (t[end] - t[1]) / length(t),
                maximum_period::Real = (t[end] - t[1]) / 1.5,
                kwargs... # see the documentation of LombScargle.jl for these
        )

Uses the Lomb-Scargle algorithm to compute a periodogram.  The advantage of the Lomb-Scargle method is that it does not require an equally sampled dataset, and it performs well on undersampled datasets.
Constraints have been set on the period, since Lomb-Scargle tends to have false peaks at very low frequencies.  That being said, it's a very flexible method.  It is extremely customizable, and the keyword arguments that can be passed to it are given [in the documentation](https://juliaastro.github.io/LombScargle.jl/stable/index.html#LombScargle.plan).
"""
function _ls_period(v, t;
                    minimum_period::Real = 4 * (t[end] - t[1]) / length(t),
                    maximum_period::Real = (t[end] - t[1]) / 1.5,
                    kwargs...
            )

    plan = LombScargle.plan(t, v;
                            minimum_frequency = 1/maximum_period,
                            maximum_frequency = 1/minimum_period,
                            kwargs...
                        )

    p = LombScargle.lombscargle(plan)
    return LombScargle.findmaxperiod(p)[1]
end

################################################################################
#                                Zero crossings                                #
################################################################################
"""
    _zc_period(v, t; line = mean(v))

Find the zero crossings of the data, and use the average difference between
zero crossings as the period.  This is a naïve implementation, with no
interpolation; however, it's useful as a sanity check.

The keyword `line` controls where the "zero point" is.

The implementation of the function was inspired by [this gist](https://gist.github.com/endolith/255291),
and has been modified for performance and to support arbitrary time grids.
"""
function _zc_period(v, t; line = mean(v))
    # This line might be a little opaque, so I'll provide some more detail.
    # The macro @. applies the broadcast operator to all operations in the
    # expression given to it, except for those which have an $ in front of
    # them.  It also does some optimization to fuse the broadcasts into a
    # single operation.
    # So, @. a * b + c is more like fma.(a, b, c) than a .* b .+ c.
    # The `@views` macro provides a view into an array without actually copying it,
    # so this allocates less.
    # What this is doing is finding all rising line crossings, which means all
    # pairs of consecutive points for which the first is below the line and the
    # second is above it.
    inds = findall(@. ≥(line, $@view(v[2:end])) & <(line, $@view(v[1:end-1])))
    mean(diff(t[inds]))
end


# create dataframe of each group by hour
# calculate sum of probability at each hour

module ProbSumHour
    using DataFrames, CSV, Plots, FreqTables
    pyplot()

    function main(is_test=false)
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_600.txt")
        df_org = CSV.read(data_path, DataFrame,
                      header=["date", "time", "ir", "lidar"],
                      delim=' ')

        hours = [Int64(floor(e/10000)) for e in df_org.time]

        # add hour array to df
        df_h = DataFrame(hour=hours)
        df_new = hcat(df_org, df_h)

        # pivot
        freqs = freqtable(df_new, :lidar, :hour)
        probs = freqs / length(df_org.lidar)

        # sum
        probs_sum = [sum(probs[begin:end, i]) for i in Array(1:24)]
        plot(Array(0:23), probs_sum, label="Probability")
        println("Sum of probability = $(sum(probs))")

        if is_test == false
            save_path = joinpath(split(@__FILE__, "src")[1], "src/prob_stats/complex_dist/prob_sum_hour/prob_sum.png")
            savefig(save_path)
        end

        return sum(probs)
    end
end


function classify(number)
    factors = []
    if number == 1
        return "deficient"
    elseif number <= 0
        return throw(DomainError(number,"Classification is only possible for natural numbers."))
    end

    for i in 1:(number-1)
        if number%i == 0
            push!(factors, i)
        end
    end
    if sum(factors) > number
        return "abundant"
    elseif sum(factors) < number
        return "deficient"
    else
        return "perfect"
    end
end


@testset "313.super-ugly-number.jl" begin
    primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
              73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, 151,
              157, 163, 167, 173, 179, 181, 191, 193, 197, 199]
    @test nth_super_ugly_number(1, primes) == 1
    @test nth_super_ugly_number(22, primes) == 22
    @test nth_super_ugly_number(222222, primes) == 1922232
end


"""
	ENU(bounds::Bounds{LLA},
		lla_ref::LLA = center(bounds),
		datum::Ellipsoid = WGS84)

Convert LLA Bounds to ENU

there's not an unambiguous conversion, but for now,
returning the minimum bounds that contain all points contained
by the input bounds
"""
function ENU(bounds::Bounds{LLA}, lla_ref::LLA = center(bounds), datum::Ellipsoid = WGS84)

    max_x = max_y = -Inf
    min_x = min_y = Inf

    xs = [bounds.min_x, bounds.max_x]
    ys = [bounds.min_y, bounds.max_y]
    if bounds.min_y < 0.0 < bounds.max_y
        push!(ys, 0.0)
    end
    ref_x = getX(lla_ref)
    if bounds.min_x < ref_x < bounds.max_x ||
       (bounds.min_x > bounds.max_x && !(bounds.min_x >= ref_x >= bounds.max_x))
        push!(xs, ref_x)
    end

    for x_lla in xs, y_lla in ys
        pt = ENU(LLA(y_lla, x_lla), lla_ref, datum)
        x, y = getX(pt), getY(pt)

        min_x, max_x = min(x, min_x), max(x, max_x)
        min_y, max_y = min(y, min_y), max(y, max_y)
    end

    return Bounds{ENU}(min_y, max_y, min_x, max_x)
end

"""
	center(bounds::Bounds{ENU})

Get Center Point of ENU Bounds Region
"""
function center(bounds::Bounds{ENU})
    x_mid = (bounds.min_x + bounds.max_x) / 2
    y_mid = (bounds.min_y + bounds.max_y) / 2

    return ENU(x_mid, y_mid)
end

"""
	center(bounds::Bounds{LLA})

Get Center Point of LLA Bounds Region
"""
function center(bounds::Bounds{LLA})
    x_mid = (bounds.min_x + bounds.max_x) / 2
    y_mid = (bounds.min_y + bounds.max_y) / 2

    if bounds.min_x > bounds.max_x
        x_mid = x_mid > 0 ? x_mid - 180 : x_mid + 180
    end

    return LLA(y_mid, x_mid)
end

"""
    inbounds(loc::ENU, bounds::Bounds{ENU})

Check Whether a location `loc` is within bounds `bounds`
"""
function inbounds(loc::ENU, bounds::Bounds{ENU})
    x, y = getX(loc), getY(loc)
    bounds.min_x <= x <= bounds.max_x &&
    bounds.min_y <= y <= bounds.max_y
end

"""
    inbounds(loc::LLA, bounds::Bounds{LLA})

Check whether a location `loc` is within bounds `bounds`
"""
function inbounds(loc::LLA, bounds::Bounds{LLA})
    x, y = getX(loc), getY(loc)
    min_x, max_x = bounds.min_x, bounds.max_x
    (min_x > max_x ? !(max_x < x < min_x) : min_x <= x <= max_x) &&
    bounds.min_y <= y <= bounds.max_y
end

"""
	onbounds(loc::T, bounds::Bounds{T}) where T<:Union{LLA,ENU}

Check whether a location `loc` is onbounds `bounds`
Works only for points that have passed the inbounds test
"""
function onbounds(loc::T, bounds::Bounds{T}) where T<:Union{LLA,ENU}
    x, y = getX(loc), getY(loc)
    x == bounds.min_x || x == bounds.max_x ||
    y == bounds.min_y || y == bounds.max_y
end

"""
	boundary_point(p1::T, p2::T, bounds::Bounds{T}) where T<:Union{LLA,ENU}

Find the closest point within bounds
Works only for points where inbounds(p1) != inbounds(p2)
"""
function boundary_point(p1::T, p2::T, bounds::Bounds{T}) where T<:Union{LLA,ENU}
    x1, y1 = getX(p1), getY(p1)
    x2, y2 = getX(p2), getY(p2)

    x, y = Inf, Inf

    if bounds.min_x >  bounds.max_x && x1*x2 < 0

        if x1 < bounds.min_x && x2 < bounds.max_x || x2 < bounds.min_x && x1 < bounds.max_x
            x = bounds.min_x
            y = y1 + (y2 - y1) * (bounds.min_x - x1) / (x2 - x1)
        elseif x1 > bounds.max_x && x2 > bounds.min_x || x2 > bounds.max_x && x1 > bounds.min_x
            x = bounds.max_x
            y = y1 + (y2 - y1) * (bounds.max_x - x1) / (x2 - x1)
        end

        p3 = T(XY(x, y))
        inbounds(p3, bounds) && return p3
    end

    # Move x to x bound if segment crosses boundary
    if x1 < bounds.min_x < x2 || x1 > bounds.min_x > x2
        x = bounds.min_x
        y = y1 + (y2 - y1) * (bounds.min_x - x1) / (x2 - x1)
    elseif x1 < bounds.max_x < x2 || x1 > bounds.max_x > x2
        x = bounds.max_x
        y = y1 + (y2 - y1) * (bounds.max_x - x1) / (x2 - x1)
    end

    p3 = T(XY(x, y))
    inbounds(p3, bounds) && return p3

    # Move y to y bound if segment crosses boundary
    if y1 < bounds.min_y < y2 || y1 > bounds.min_y > y2
        x = x1 + (x2 - x1) * (bounds.min_y - y1) / (y2 - y1)
        y = bounds.min_y
    elseif y1 < bounds.max_y < y2 || y1 > bounds.max_y > y2
        x = x1 + (x2 - x1) * (bounds.max_y - y1) / (y2 - y1)
        y = bounds.max_y
    end

    p3 = T(XY(x, y))
    inbounds(p3, bounds) && return p3

    error("Failed to find boundary point.")
end


#---------------------------------------------------------------------#
#This function computes the norm of 2 vectors
#Written by F.X. Giraldo on April 19, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function compute_norm(qn,qe,Npoin,DFloat)
    norm2=norm(qn-qe,2)/norm(qe,2)
    return (norm2)
end


module RandomFloat

import Base: frexp, ldexp

export randfloat

# frexp([x,y]) == ([frexp(x)[1], frexp(y)[1]], [frexp(x)[2], frexp(y)[2]]])
# ldexp(frexp([x,y])...) == ldexp(frexp([x,y])[1], frexp([x,y])[2]) == [x,y]
@vectorize_1arg Real frexp
@vectorize_2arg Real ldexp

for (F,I) in ((:Float16, :Int16), (:Float32, :Int32), (:Float64, Int64))
  @eval begin
    function randfloat(r::FloatRange{$F}=sqrt(eps(one($F))):(one($F)/sqrt(eps(one($F)))), dims...)
      firstIsNeg = signbit(first(r))
      lastIsNeg  = signbit(last(r))
      if !firstIsNeg
        lo = reinterpret($I,first(r))
        hi = reinterpret($I,last(r))
        reinterpret($F, rand(lo:hi, dims...)
        #step = reinterpret($I, eps(eps(last(r)))/2)
        #reinterpret($F, rand(lo:step:hi, dims...))
      elseif lastIsNeg
        -randfloat((-last(r)):(-first(r)), dims...)
      else
        throw(ArgumentError("expecting range to be negative or positive, not mixed"))
      end
    end
  end
end

end # module


module GCNN

export RotGroupConv, RotGroupConvTranspose

using Flux
using Flux: conv, ∇conv_data, convfilter, glorot_uniform, @functor, DenseConvDims

"""
Add an extra dimension to the end of an array.
"""
expand_dims(x) = reshape(x, size(x)..., 1)

"""
A conveniently bundled group of rotation matrices.
"""
struct RotGroup{RT<:AbstractMatrix}
    Rs::Vector{RT}
end

@functor RotGroup

"""
Compute a sparse matrix that simultaneously does `nrots` interpolated rotations
of a m×n matrix. To compute the rotations, the image must be flattened, then
sparse matrix times dense vector multiply, then reshape the output vector into
a nrots×m×n array.
"""
function RotGroup(
        m::Integer, n::Integer, nrots::Integer, periodicity=2π, mask_disk::Bool=true)

    # TODO: until we test m != n case, I wouldn't trust it
    @assert m == n

    ci, cj = div(m, 2)+1, div(n, 2)+1
    c = min(ci, cj)

    Rs = Vector{Matrix{Float32}}(undef, nrots)

    for (r, θ) in enumerate(range(0, periodicity, length=nrots+1)[1:end-1])
        cosθ, sinθ = cos(θ), sin(θ)

        # COO sparse matrix values
        R = zeros(Float32, (m*n, m*n))

        for i in 1:m, j in 1:n
            if mask_disk && (i-ci)^2 + (j-cj)^2 > (c+0.5)^2
                continue
            end

            # flattened index
            l_input = (j-1)*m + i

            # compute fractional rotated indexes
            ir = cosθ * (i - ci) + sinθ * (j - cj) + ci
            jr = -sinθ * (i - ci) + cosθ * (j - cj) + cj

            # interpolate across 2x2 square
            ir1 = floor(Int, ir)
            ir2 = ir1 + 1
            jr1 = floor(Int, jr)
            jr2 = jr1 + 1

            wi = ir - ir1
            wj = jr - jr1

            for (irk, wik) in [(ir1, 1-wi), (ir2, wi)], (jrk, wjk) in [(jr1, 1-wj), (jr2, wj)]
                if 1 <= irk <= m && 1 <= jrk <= n
                    l_output = (jrk-1)*m + irk
                    R[l_output, l_input] = wik * wjk
                end
            end

            Rs[r] = R
        end
    end

    # TODO: It's probably better to use dense matrices since these will be
    # small. (E.g. to rotate a 5x5 rotation kernel, we only need a 25x25 matrix).

    return RotGroup(Rs)
end

"""
Apply each rotation matrix to an array w of shape [width, height, in_channels, out_channels]
produces a vector of rotated arrays each of the same shape.
"""
function (rg::RotGroup)(w::AbstractArray{T,4}) where {T}
    width, height, in_channels, out_channels = size(w)
    w_flat = reshape(w, (width*height, in_channels*out_channels))
    return [reshape(R*w_flat, (width, height, in_channels, out_channels)) for R in rg.Rs]
end

Base.length(rg::RotGroup) = length(rg.Rs)

"""
Group convolutional layer.
"""
struct RotGroupConv{F, WT<:AbstractArray, BT, RT<:AbstractMatrix, M, N}
    rg::RotGroup{RT}
    weight::WT
    bias::BT
    σ::F
    stride::NTuple{N,Int}
    pad::NTuple{M,Int}
    dilation::NTuple{N,Int}
end

@functor RotGroupConv

"""
    RotGroupConv(nrots, filter, in => out, σ = identity)

Construct a Group CNN (G-CNN) layer with `nrots` equally spaced rotations, and
a `(filter, filter)` convolution kernel, with `in` input channels and `out`
output channels, with an optional elementwise nonlinearity `σ`.
"""
function RotGroupConv(
        nrots::Integer, filter::Integer, ch::Pair{<:Integer,<:Integer}, σ=identity;
        init=glorot_uniform, stride::Int=1, pad=0, dilation=1, use_bias::Bool=true)

    weight = convfilter((filter, filter), ch; init)
    # bias = create_bias(weght, use_bias, size(weight, N))
    rg = RotGroup(filter, filter, nrots)

    return RotGroupConv(rg, weight, σ, stride, pad, dilation, use_bias)
end

function RotGroupConv(
        rg::RotGroup, weight::AbstractArray{T, N}, σ, stride, pad, dilation, use_bias) where {T, N}

    stride = Flux.expand(Val(N-2), stride)
    dilation = Flux.expand(Val(N-2), dilation)
    pad = Flux.calc_padding(Conv, pad, size(weight)[1:N-2], dilation, stride)
    bias = Flux.create_bias(weight, use_bias, size(weight, N))
    return RotGroupConv(rg, weight, bias, σ, stride, pad, dilation)
end

"""
Rotation group convolution applied to an array with shape [width, height, in_channels, batches] and
producing an array of shape [width, height, out_channels, batches, nrotations].
"""
function (lyr::RotGroupConv)(x::AbstractArray{T,4}) where {T}
    # Rotate the weights according to each rotation matrix, apply each matrix, then concatenate.
    bias = reshape(lyr.bias, ntuple(_ -> 1, length(lyr.stride))..., :, 1)
    cdims = DenseConvDims(x, lyr.weight; stride=lyr.stride, padding=lyr.pad, dilation=lyr.dilation)
    return cat([expand_dims(lyr.σ.(conv(x, Rw, cdims) .+ bias)) for Rw in lyr.rg(lyr.weight)]..., dims=5)
end

"""
Rotation group convolution applied to an array with shape `[width, height,
in_channels, batches, nrotations]` and producing an array of shape `[width,
height, out_channels, batches, nrotations]`.
"""
function (lyr::RotGroupConv)(x::AbstractArray{T,5}) where {T}
    # TODO:
    nrots = length(lyr.rg)
    @assert size(x, 5) == nrots

    bias = reshape(lyr.bias, ntuple(_ -> 1, length(lyr.stride))..., :, 1)
    cdims = DenseConvDims(size(x)[1:end-1], size(lyr.weight); stride=lyr.stride, padding=lyr.pad, dilation=lyr.dilation)

    Rws = lyr.rg(lyr.weight)
    Rxs = [view(x, :, :, :, :, i) for i in 1:nrots]

    # gradients don't work with zip
    #return cat([expand_dims(lyr.σ.(conv(Rx, Rw))) for (Rx, Rw) in zip(Rxs, Rws)]..., dims=5)

    return cat([expand_dims(lyr.σ.(conv(Rxs[k], Rws[k], cdims) .+ bias)) for k in 1:nrots]..., dims=5)
end

"""
Group transposed convolutional layer.
"""
struct RotGroupConvTranspose{F, WT<:AbstractArray, BT, RT<:AbstractMatrix, M, N}
    rg::RotGroup{RT}
    weight::WT
    bias::BT
    σ::F
    # TODO: Make this a tuple so everythig works.
    stride::NTuple{N,Int}
    pad::NTuple{M,Int}
    dilation::NTuple{N,Int}
end

@functor RotGroupConvTranspose

"""
    RotGroupConvTranspose(nrots, filter, in => out, σ = identity)

Construct a Group CNN (G-CNN) transposed convolution layer with `nrots` equally
spaced rotations, and a `(filter, filter)` convolution kernel, with `in` input
channels and `out` output channels, with an optional elementwise nonlinearity
`σ`.
"""
function RotGroupConvTranspose(
        nrots::Integer, filter::Integer, ch::Pair{<:Integer,<:Integer}, σ=identity;
        init=glorot_uniform, stride::Int=1, pad=0, dilation=1, use_bias::Bool=true)

    weight = convfilter((filter, filter), reverse(ch); init)

    # bias = create_bias(weght, use_bias, size(weight, N))
    rg = RotGroup(filter, filter, nrots)

    return RotGroupConvTranspose(rg, weight, σ, stride, pad, dilation, use_bias)
end

function RotGroupConvTranspose(
        rg::RotGroup, weight::AbstractArray{T, N}, σ, stride, pad, dilation, use_bias) where {T, N}
    stride = Flux.expand(Val(N-2), stride)
    dilation = Flux.expand(Val(N-2), dilation)
    pad = Flux.calc_padding(ConvTranspose, pad, size(weight)[1:N-2], dilation, stride)
    bias = Flux.create_bias(weight, use_bias, size(weight, N-1))
    return RotGroupConvTranspose(rg, weight, bias, σ, stride, pad, dilation)
end

function conv_transpose_dims(c::RotGroupConvTranspose, xsize::NTuple{N, Int}) where {N}
    # Calculate size of "input", from ∇conv_data()'s perspective...
    combined_pad = (c.pad[1:2:end] .+ c.pad[2:2:end])
    I = (xsize[1:end-2] .- 1).*c.stride .+ 1 .+ (size(c.weight)[1:end-2] .- 1).*c.dilation .- combined_pad
    C_in = size(c.weight)[end-1]
    batch_size = xsize[end]
    # Create DenseConvDims() that looks like the corresponding conv()
    return DenseConvDims((I..., C_in, batch_size), size(c.weight);
                        stride=c.stride,
                        padding=c.pad,
                        dilation=c.dilation,
    )
end

"""
Rotation group transposed convolution applied to an array with shape [width,
height, in_channels, batches] and producing an array of shape [width, height,
out_channels, batches, nrotations].
"""
function (lyr::RotGroupConvTranspose)(x::AbstractArray{T,4}) where {T}
    # Rotate the weights according to each rotation matrix, apply each matrix, then concatenate.
    cdims = conv_transpose_dims(lyr, size(x))
    bias = reshape(lyr.bias, map(_->1, lyr.stride)..., :, 1)

    # return cat([expand_dims(lyr.σ.(∇conv_data(x, Rw, cdims) .+ lyr.bias)) for Rw in lyr.rg(lyr.weight)]..., dims=5)
    return cat([expand_dims(lyr.σ.(∇conv_data(x, Rw, cdims) .+ bias)) for Rw in lyr.rg(lyr.weight)]..., dims=5)
end

"""
Rotation group transposed convolution applied to an array with shape `[width,
height, in_channels, batches, nrotations]` and producing an array of shape
`[width, height, out_channels, batches, nrotations]`.
"""
function (lyr::RotGroupConvTranspose)(x::AbstractArray{T,5}) where {T}
    nrots = length(lyr.rg)
    @assert size(x, 5) == nrots

    Rws = lyr.rg(lyr.weight)
    Rxs = [view(x, :, :, :, :, i) for i in 1:nrots]
    cdims = conv_transpose_dims(lyr, size(x)[1:end-1])
    bias = reshape(lyr.bias, map(_->1, lyr.stride)..., :, 1)

    # gradients don't work with zip
    #return cat([expand_dims(lyr.σ.(conv(Rx, Rw))) for (Rx, Rw) in zip(Rxs, Rws)]..., dims=5)

    return cat([expand_dims(lyr.σ.(∇conv_data(Rxs[k], Rws[k], cdims) .+ bias)) for k in 1:nrots]..., dims=5)
end

end # module


using ExactD, Test, LinearAlgebra

@testset "expectation values" begin
L = 5
N = 3
basis = make_LN_basis(L, N)
s = normalize(rand(length(basis)))
r = normalize(rand(length(basis)))

@testset "<s|n_p1*n_p2*...*n_pn|s>" begin
    p = 1
    @test expected_n(basis, s, p) ≈ s[1]^2+s[2]^2+s[3]^2+s[5]^2+s[6]^2+s[8]^2
    p = 4
    @test expected_n(basis, s, p) ≈ s[2]^2+s[3]^2+s[4]^2+s[8]^2+s[9]^2+s[10]^2

    p = [1, 3]
    @test expected_n(basis, s, p) ≈ s[1]^2+s[3]^2+s[6]^2
    p = [4, 2]
    @test expected_n(basis, s, p) ≈ s[2]^2+s[4]^2+s[9]^2

    p = [1, 3, 5]
    @test expected_n(basis, s, p) ≈ s[6]^2
    p = [4, 2, 3]
    @test expected_n(basis, s, p) ≈ s[4]^2

    p = [1, 2, 3, 5]
    @test expected_n(basis, s, p) ≈ 0. atol=1e-10
end

@testset "<r|n_p1*n_p2*...*n_pn|s>" begin
    p = 1
    @test expected_n(basis, r, s, p) ≈ s[1]*r[1]+s[2]*r[2]+s[3]*r[3]+s[5]*r[5]+s[6]*r[6]+s[8]*r[8]
    p = 4
    @test expected_n(basis, r, s, p) ≈ s[2]*r[2]+s[3]*r[3]+s[4]*r[4]+s[8]*r[8]+s[9]*r[9]+s[10]*r[10]

    p = [1, 3]
    @test expected_n(basis, r, s, p) ≈ s[1]*r[1]+s[3]*r[3]+s[6]*r[6]
    p = [4, 2]
    @test expected_n(basis, r, s, p) ≈ s[2]*r[2]+s[4]*r[4]+s[9]*r[9]

    p = [1, 3, 5]
    @test expected_n(basis, r, s, p) ≈ s[6]*r[6]
    p = [4, 2, 3]
    @test expected_n(basis, r, s, p) ≈ s[4]*r[4]

    p = [1, 2, 3, 5]
    @test expected_n(basis, r, s, p) ≈ 0. atol=1e-10
end

@testset "<s|b^+_p[1]*...*b^+_p[n]*b_q[1]*...*b_q[m]|s>" begin
    p = 1
    q = 1
    @test expected_pq(basis, s, p, q) ≈ s[1]^2+s[2]^2+s[3]^2+s[5]^2+s[6]^2+s[8]^2
    p = 4
    q = 4
    @test expected_pq(basis, s, p, q) ≈ s[2]^2+s[3]^2+s[4]^2+s[8]^2+s[9]^2+s[10]^2

    p = [1, 2]
    q = [1, 3]
    @test expected_pq(basis, s, p, q) ≈ s[2]*s[3] + s[5]*s[6]
    p = [4, 2]
    q = [1, 3]
    @test expected_pq(basis, s, p, q) ≈ s[6]*s[9]
end

@testset "<r|b^+_p[1]*...*b^+_p[n]*b_q[1]*...*b_q[m]|s>" begin
    p = 1
    q = 1
    @test expected_pq(basis, r, s, p, q) ≈ s[1]*r[1]+s[2]*r[2]+s[3]*r[3]+s[5]*r[5]+s[6]*r[6]+s[8]*r[8]
    p = 4
    q = 4
    @test expected_pq(basis, r, s, p, q) ≈ s[2]*r[2]+s[3]*r[3]+s[4]*r[4]+s[8]*r[8]+s[9]*r[9]+s[10]*r[10]

    p = [1, 2]
    q = [1, 3]
    @test expected_pq(basis, r, s, p, q) ≈ r[2]*s[3] + r[5]*s[6]
    p = [4, 2]
    q = [1, 3]
    @test expected_pq(basis, r, s, p, q) ≈ s[6]*r[9]
end
end


#=
MMFF94Parameters:
- Julia version:
- Author: Dan
- Date: 2021-07-30
=#
"""
Radii derived from [MMFF94 Validation Suite](http://www.ccl.net/cca/data/MMFF94/).
"""
const MMFF94radii = Float64[
     0.33, 0.0,
	 1.34, 0.90, 0.81, 0.77, 0.73, 0.72, 0.74, 0.0,
	 1.54, 1.30, 1.22, 1.15, 1.09, 1.03, 1.01, 0.0,
	 1.96, 1.74,
	 1.44, 1.36, 0.00, 0.00, 0.00,
	 0.00, 0.00, 0.00, 1.38, 1.31,
	 1.19, 1.20, 1.20, 1.16, 1.15, 0.0,
	 2.11, 1.92,
	 1.62, 1.48, 0.00, 0.00, 0.00,
	 0.00, 0.00, 0.00, 1.53, 1.48,
	 1.46, 1.40, 1.41, 1.35, 1.33, 0.0
]

"""
Stretch parameters derived from [MMFF94 Validation Suite](http://www.ccl.net/cca/data/MMFF94/).
"""
const MMFF94StrechParameters = Float64[
     2.20, 0.0,
	 0.97, 1.47, 2.01, 2.5, 3.07, 3.5, 4.10, 0.0,
	 1.01, 1.23, 1.47, 1.74, 2.06, 2.44, 2.83, 0.0,
	 0.91, 1.04,
	 1.3, 1.5, 1.6, 1.6, 1.5,
	 1.8, 1.8, 1.8, 1.9, 1.6,
	 1.82, 2.02, 2.20, 2.48, 2.74, 0.0,
	 0.89, 0.99,
	 1.3, 1.4, 1.6, 1.8, 1.9,
	 2.2, 2.2, 2.2, 1.9, 1.7,
	 1.49, 1.72, 1.82, 2.01, 2.21, 0.0
]


####    GridOperators

# Local finite-difference operator functions.

export advect
export grad, ∇_pos, ∇_neg

export ∇_z_flux, ∇_z_centered, ∇_z_dual, ∇_z_upwind
export Δ_z, Δ_z_dual

export advect_old
export UpwindAdvective,
       UpwindCollocated,
       OneSidedUp,
       OneSidedDn,
       CenteredUnstable

function advect(f, w, grid::Grid)
  @assert length(f)==3
  @assert length(w)==3
  if w[2] < 0
    return (f[3]-f[2])*grid.Δzi
  else
    return (f[2]-f[1])*grid.Δzi
  end
end

"""
    grad(f, grid::Grid{FT}) where FT

Computes local finite-difference gradient of `f`: `∇f`
"""
function grad(f, grid::Grid{FT}) where FT
  @assert length(f)==3 || length(f)==2
  if length(f)==2
    return (f[2]-f[1])*grid.Δzi
  else
    return (f[3]-f[1])*FT(0.5)*grid.Δzi
  end
end

"""
    ∇_pos(f, grid::Grid)

Computes a one-sided (up) local finite-difference gradient of `f`: `∇f`
"""
function ∇_pos(f, grid::Grid)
  @assert length(f)==3
  return (f[2]-f[1])*grid.Δzi
end

"""
    ∇_neg(f, grid::Grid)

Computes a one-sided (down) local finite-difference gradient of `f`: `∇f`
"""
function ∇_neg(f, grid::Grid)
  @assert length(f)==3
  return (f[3]-f[2])*grid.Δzi
end

"""
    ∇_z_flux(f, grid::Grid)

Computes the local derivative of field ``f``, which is
assumed to live on the dual grid: ``∇f``
"""
function ∇_z_flux(f, grid::Grid)
  @assert length(f)==2
  (f[2] - f[1])*grid.Δzi
end

"""
    ∇_z_centered(f, grid::Grid)

Computes the local derivative of field ``f``, which is
assumed to live on the dual grid: ``∇f``
"""
function ∇_z_centered(f, grid::Grid)
  @assert length(f)==3
  (f[3] - f[1])*grid.Δzi/2
end

"""
    ∇_z_dual(f, grid::Grid)

Computes the local derivative of field ``f``, which is
assumed to live on the dual grid: ``∇f``
"""
function ∇_z_dual(f, grid::Grid)
  @assert length(f)==3
  return [∇_z_flux(f[1:2], grid), ∇_z_flux(f[2:3], grid)]
end

"""
    ∇_z_upwind(ϕ, w, grid::Grid)

Local upwind gradient operator ``∇ϕ``. This
operator is stable but numerically diffusive.
"""
function ∇_z_upwind(ϕ, w, grid::Grid)
  @assert length(ϕ)==3
  @assert length(w)==3
  return w[2]>0 ? ∇_z_flux(ϕ[1:2], grid) : ∇_z_flux(ϕ[2:3], grid)
end

"""
    Δ_z(f, grid::Grid)

Computes the local Laplacian of field ``f``:
``∇ • (∇f)``
"""
function Δ_z(f, grid::Grid)
  @assert length(f)==3
  (f[3] - 2*f[2] + f[1])*grid.Δzi2
end

"""
    Δ_z(f, grid::Grid, K)

Computes the local Laplacian of field ``f``
with a variable coefficient:
``∇ • (K ∇f)``
"""
function Δ_z(f, grid::Grid, K)
  @assert length(f)==3
  K∇f_M = (K[1]+K[2])/2*(f[2]-f[1])*grid.Δzi
  K∇f_P = (K[2]+K[3])/2*(f[3]-f[2])*grid.Δzi
  return (K∇f_P - K∇f_M)*grid.Δzi
end

"""
    Δ_z_dual(f, grid::Grid, K)

Computes the local Laplacian of field ``f``
with a variable coefficient:
``∇ • (K ∇f)``
"""
function Δ_z_dual(f, grid::Grid, K)
  @assert length(f)==3
  @assert length(K)==2
  K∇f_M = K[1]*(f[2]-f[1])*grid.Δzi
  K∇f_P = K[2]*(f[3]-f[2])*grid.Δzi
  return (K∇f_P - K∇f_M)*grid.Δzi
end

abstract type AdvectForm end
struct UpwindAdvective <: AdvectForm end

abstract type ConservativeForm end
struct UpwindCollocated <: ConservativeForm end
struct OneSidedUp <: ConservativeForm end
struct OneSidedDn <: ConservativeForm end
struct CenteredUnstable <: ConservativeForm end

"""
    advect(ϕ, w, grid::Grid, ::AdvectForm, Δt)

Advection operator, which computes

 - ``w • ∇ϕ`` given `AdvectiveForm`
 - ``∇ • (wϕ)`` given `ConservativeForm`

which is dispatched by `AdvectForm`.
"""
function advect_old(ϕ, ϕ_dual, w, w_dual, grid::Grid, ::UpwindAdvective, Δt)
  @assert length(ϕ)==3
  @assert length(ϕ_dual)==2
  @assert length(w)==3
  @assert length(w_dual)==2
  if w[2] > 0
    return w[ 2 ]*(ϕ[2] - ϕ[1]) * grid.Δzi
  else
    return w[ 2 ]*(ϕ[3] - ϕ[2]) * grid.Δzi
  end
end
function advect_old(ϕ, ϕ_dual, w, w_dual, grid::Grid, ::UpwindCollocated, Δt)
  @assert length(ϕ)==3
  @assert length(ϕ_dual)==2
  @assert length(w)==3
  @assert length(w_dual)==2
  if sum(w)/length(w) > 0
    return (w[2]*ϕ[2] - w[1]*ϕ[1]) * grid.Δzi
  else
    return (w[3]*ϕ[3] - w[2]*ϕ[2]) * grid.Δzi
  end
end
function advect_old(ϕ, ϕ_dual, w, w_dual, grid::Grid, ::OneSidedUp, Δt)
  @assert length(ϕ)==3
  @assert length(ϕ_dual)==2
  @assert length(w)==3
  @assert length(w_dual)==2
  return (w[2]*ϕ[2] - w[1]*ϕ[1]) * grid.Δzi
end
function advect_old(ϕ, ϕ_dual, w, w_dual, grid::Grid, ::OneSidedDn, Δt)
  @assert length(ϕ)==3
  @assert length(ϕ_dual)==2
  @assert length(w)==3
  @assert length(w_dual)==2
  return (w[3]*ϕ[3] - w[2]*ϕ[2]) * grid.Δzi
end
function advect_old(ϕ, ϕ_dual, w, w_dual, grid::Grid, ::CenteredUnstable, Δt)
  @assert length(ϕ)==3
  @assert length(ϕ_dual)==2
  @assert length(w)==3
  @assert length(w_dual)==2
  return (w_dual[2]*ϕ_dual[2]-w_dual[1]*ϕ_dual[1])*grid.Δzi
end


using QuantumStatistics, Parameters, Printf
using Roots, Polylogarithms
using Plots
# using PackageCompiler
# --sysimage sys_plots.so
# create_sysimage(:Plots, sysimage_path="sys_plots.so", precompile_execution_file="precompile_plots.jl")
include("parameter.jl")
# gr()
# Plots.GRBackend()

@with_kw struct Para
    n::Int = 0 # external Matsubara frequency
    Qsize::Int = 17
    Tsize::Int = 16
    extQ::Vector{SVector{3,Float64}} = [@SVector [q, 0.0, 0.0] for q in LinRange(0.0, 32.0 * kF, Qsize)]
    extT::Vector{SVector{3,Float64}} = [@SVector [T, 0.0, 0.0] for T in LinRange(10, 610, Tsize)]
    fixQ::Float64 = 10* kF
end

function chemical_potential(beta)
    f(β,μ)=real(polylog(3/2,-exp(β*μ)))+4/3/π^0.5*(β)^(3/2)
    g(μ)=f(beta,μ)
    return find_zero(g, (-10000,1))
end

function chi0vsq(beta)
    @unpack n, Qsize, extQ = Para()

    μ=chemical_potential(beta)*EF
    println("μ=$μ, reduced β=$beta")

    pdata = Matrix{Float64}(undef, Qsize, 2)
    for (idx, q) in enumerate(extQ)
        q = q[1]
        p, err = TwoPoint.LindhardΩnFiniteTemperature(dim, q, n, μ, kF, β, me, spin)
        @printf("%10.6f  %10.6f ± %10.6f\n", q / kF, p, err)
        pdata[idx,1] = q/kF
        # pdata[idx,1] = (kF/q)^2
        pdata[idx,2] = -p
        # pdata[idx,1] = (kF/q)^4
        pdata[idx,2] = -p-2/3/π^2*EF^(3/2)/q^2
    end
    return pdata
end

function chi0vsT()
    @unpack n, Tsize, extT, fixQ = Para()
    pdata = Matrix{Float64}(undef, Tsize, 2)
    for (idx, T) in enumerate(extT)
        T = T[1]
        μ=chemical_potential(1.0/T)*EF
        println("μ=$μ, reduced T=$T, fixed extMom q=$fixQ")
        β_in = 1.0/T/kF^2

        p, err = TwoPoint.LindhardΩnFiniteTemperature(dim, fixQ, n, μ, kF, β_in, me, spin)
        @printf("%10.6f  %10.6f ± %10.6f\n", T, p, err)
        # pdata[idx,1] = T
        pdata[idx,1] = 1.0/T
        pdata[idx,2] = -p
    end
    return pdata, fixQ
end
# labels = [string(pdata[i,1]) for i in 1:Tsize]

# @unpack Tsize, extT = Para()
# for (idx, T) in enumerate(extT)
#     T = T[1]
#     pdata = chi0vsq(1.0/T)
#     plot!(pdata[:,1], pdata[:,2], xlabel="q/k_F", ylabel="χ_0-(1/q^2 term)", label="T=$T")
#     # plt = plot(pdata[:,1], pdata[:,2], xlim=(0, 0.1), ylim=(0,) xlabel="(k_F/q)^2", ylabel="χ_0", label="β=$beta", title = "ideal polarization")
#     # plt = plot(pdata[:,1], pdata[:,2], xlabel="(k_F/q)^4", ylabel="χ_0-(1/q^2 term)", label="β=$beta", title = "ideal polarization")
# end
# savefig("chi0vsq_T_v1.pdf")
# show(plt)

pdata, fixQ = chi0vsT()
qint = round(Int, fixQ/kF)
# plt = plot(pdata[:,1], pdata[:,2], xlim=(0,0.2), ylim=(0,0.0017), xlabel="√β", ylabel="χ_0(q=$qint k_F)", title = "ideal polarization", legend=:none)
plt = plot(pdata[:,1], pdata[:,2], xlim=(0,0.04), ylim=(0,0.0017), xlabel="β", ylabel="χ_0(q=$qint k_F)", title = "ideal polarization", legend=:none)
savefig("chi0vsβ_q$qint.pdf")
display(plt)


@testset "GaloisField" begin

  for p in [31, 11, 101]
    F = GF(p)
    G, mG = unit_group(F)
    #Test generator
    g = mG(G[1])
    g1 = g
    for i = 2:p-2
      g1 *= g
      @test !isone(g1)
    end
    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      elG = mG\x
      @test mG(elG) == x
    end

    G1, mG1 = unit_group(F, n_quo = 5)
    @test order(G1) == 5
    g = mG1(G1[1])
    g = g^(div(size(F)-1, 5))
    @test !isone(g)
    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      c = mG1\x
      el = mG1(c)
      if el != x
        dl = mG\(el*inv(x))
        @test iszero(mod(dl[1], 5))
      end
    end
  end
end

@testset "GaloisFmpzField" begin

  for p in [31, 11, 101]
    F = GF(fmpz(p))
    G, mG = unit_group(F)
    #Test generator
    g = mG(G[1])
    g1 = g
    for i = 2:p-2
      g1 *= g
      @test !isone(g1)
    end
    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      elG = mG\x
      @test mG(elG) == x
    end

    G1, mG1 = unit_group(F, n_quo = 5)
    @test order(G1) == 5

    g = mG1(G1[1])
    g = g^(div(size(F)-1, 5))
    @test !isone(g)
    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      c = mG1\x
      el = mG1(c)
      if el != x
        dl = mG\(el*inv(x))
        @test iszero(mod(dl[1], 5))
      end
    end
  end
end

@testset "FqNmodFiniteField" begin

  for p in [31, 11, 101]
    F = FiniteField(p, 2)[1]
    G, mG = unit_group(F)
    #Test generator
    g = mG(G[1])
    g1 = g
    for i = 2:p^2-2
      g1 *= g
      @test !isone(g1)
    end
    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      elG = mG\x
      @test mG(elG) == x
    end

    G1, mG1 = unit_group(F, n_quo = 5)
    @test order(G1) == 5
    g = mG1(G1[1])
    g = g^(div(size(F)-1, 5))
    @test !isone(g)

    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      c = mG1\x
      el = mG1(c)
      if el != x
        dl = mG\(el*inv(x))
        @test iszero(mod(dl[1], 5))
      end
    end
  end
end

@testset "FqFiniteField" begin

  for p in [31, 11, 101]
    F = FiniteField(fmpz(p), 2, "a")[1]
    G, mG = unit_group(F)
    #Test generator
    g = mG(G[1])
    g1 = g
    for i = 2:p^2-2
      g1 *= g
      @test !isone(g1)
    end
    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      elG = mG\x
      @test mG(elG) == x
    end

    G1, mG1 = unit_group(F, n_quo = 5)
    @test order(G1) == 5
    g = mG1(G1[1])
    g = g^(div(size(F)-1, 5))
    @test !isone(g)

    for i = 1:5
      x = rand(F)
      while iszero(x)
        x = rand(F)
      end
      c = mG1\x
      el = mG1(c)
      if el != x
        dl = mG\(el*inv(x))
        @test iszero(mod(dl[1], 5))
      end
    end
  end
end


###########################
#### Surface roughness ####
###########################

#' Roughness Reynolds Number
#'
#' calculates the Roughness Reynolds Number.
#'
#' - Tair      Air temperature (deg C)
#' - pressure  Atmospheric pressure (kPa)
#' - ustar     Friction velocity (m s-1)
#' - z0m       Roughness length (m)
#' - constants Kelvin - conversion degree Celsius to Kelvin
#'                  pressure0 - reference atmospheric pressure at sea level (Pa)
#'                  Tair0 - reference air temperature (K)
#'
#' # Details
 The Roughness Reynolds Number is calculated as in Massman 1999a:
#'
#'            ``Re = z0m * ustar / v``
#'
#'          where `v` is the kinematic viscosity (m2 s-1).
#'
#' # Value
 - Re -: Roughness Reynolds Number (-)
#'
#' #References
#' Massman, W_J., 1999a: A model study of kB H- 1 for vegetated surfaces using
#'            'localized near-field' Lagrangian theory. Journal of Hydrology 223, 27-43.
#'
#' ```@example; output = false
#' ```
#' Reynolds_Number(25,100,0.5,z0m=0.5)
#'
"""
"""
function Reynolds_Number(Tair,pressure,ustar,z0m,constants=BigleafConstants())

  v  = kinematic_viscosity(Tair,pressure,constants)
  Re = z0m*ustar/v

  return(Re)
end

#' Roughness Parameters
#'
#' A simple approximation of the two roughness parameters displacement height (d)
#'              and roughness length for momentum (z0m).
#'
#' - method    Method to use, one of `Val(:canopy_height),Val(:canopy_height_LAI),Val(:wind_profile)`
#'                  NOTE: if `method = Val(:canopy_height)`, only the following three arguments
#'                  are used. If `method = Val(:canopy_height_LAI)`, only `zh, LAI, cd`,
#'                  and `hs` are required.
#' - zh        Vegetation height (m)
#' - frac_d    Fraction of displacement height on canopy height (-)
#' - frac_{z0m}  Fraction of roughness length on canopy height (-)
#' - LAI       Leaf area index (-)
#' - zr        Instrument (reference) height (m)
#' - cd        Mean drag coefficient for individual leaves. Defaults to 0.2.
#'                  Only needed if `method = Val(:canopy_height_LAI)`.
#' - hs        roughness length of the soil surface (m). Only needed if `method = Val(:canopy_height_LAI)`
#'                  The following arguments are only needed if `method = Val(:wind_profile)`!
#' - data      DataFrame or matrix containing all required variables
#' - Tair      Air temperature (deg C)
#' - pressure  Atmospheric pressure (kPa)
#' - wind      Wind speed at height zr (m s-1)
#' - ustar     Friction velocity (m s-1)
#' - H         Sensible heat flux (W m-2)
#' - d         Zero-plane displacement height (m); optional
#' - z0m       Roughness length for momentum (m); optional
#' - stab_roughness   Should stability correction be considered? Default is `true`.
#' - stab_formulation Stability correction function used (If `stab_correction = true`).
#'                         Either `Val(:Dyer_1970)` or `Val(:Businger_1971)`.
#' - constants k - von-Karman constant (-)
#'                  Kelvin - conversion degree Celsius to Kelvin
#'                  cp - specific heat of air for constant pressure (J K-1 kg-1)
#'                  g - gravitational acceleration (m s-2)
#'                  se_median - conversion standard error (SE) of the mean to SE of the median
#'
#'
#' # Details
 The two main roughness parameters, the displacement height (d)
#'          and the roughness length for momentum (z0m) can be estimated from simple
#'          empirical relationships with canopy height (zh). If `method = Val(:canopy_height)`,
#'          the following formulas are used:
#'
#'            ``d = frac_d * zh``
#'
#'            ``z0m = frac_{z0m} * zh``
#'
#'          where frac_d defaults to 0.7 and frac_{z0m} to 0.1.
#'
#'          Alternatively, d and z0m can be estimated from both canopy height and LAI
#'          (If `method = Val(:canopy_height_LAI)`).
#'          Based on data from Shaw & Pereira 1982, Choudhury & Monteith 1988 proposed
#'          the following semi-empirical relations:
#'
#'            ``X = cd * LAI``
#'
#'            ``d = 1.1 * zh * ln(1 + X^(1/4))``
#'
#'            ``z0m = hs + 0.3 * zh * X^(1/2)   for 0 <= X <= 0.2``
#'
#'            ``z0m = hs * zh * (1 - d/zh)   for 0.2 < X``
#'
#'          If `method = Val(:wind_profile)`, z0m is estimated by solving
#'          the wind speed profile for z0m:
#'
#'            ``z0m = median((zr - d) * exp(-k*wind / ustar - psi_m)``
#'
#'          By default, d in this equation is fixed to 0.7*zh, but can be set to any
#'          other value. psi_m is 0 if `stab_roughness = false`.
#'
#' # Value
 a DataFrame with the following columns:
#'         - d: Zero-plane displacement height (m)
#'         - z0m: Roughness length for momentum (m)
#'         - z0m_se: Only if `method = wind_profile`: Standard Error of the median for z0m (m)
#'
#' #References
#' Choudhury, B. J., Monteith J_L., 1988: A four-layer model for the heat
#'             budget of homogeneous land surfaces. Q. J. R. Meteorol. Soc. 114, 373-398.
#'
#'             Shaw, R. H., Pereira, A., 1982: Aerodynamic roughness of a plant canopy:
#'             a numerical experiment. Agricultural Meteorology, 26, 51-65.
#'
#' #See also
#' [`wind_profile`](@ref)
#'
#' ```@example; output = false
#' ```
#' # estimate d and z0m from canopy height for a dense (LAI=5) and open (LAI=2) canopy
#' roughness_parameters(method=Val(:canopy_height_LAI),zh=25,LAI=5)
#' roughness_parameters(method=Val(:canopy_height_LAI),zh=25,LAI=2)
#'
#' # fix d to 0.7*zh and estimate z0m from the wind profile
#' df = DataFrame(Tair=c(25,25,25),pressure=100,wind=c(3,4,5),ustar=c(0.5,0.6,0.65),H=200)
#' roughness_parameters(method=Val(:wind_profile),zh=25,zr=40,frac_d=0.7,data=df)
#'
#' # assume d = 0.8*zh
#' roughness_parameters(method=Val(:wind_profile),zh=25,zr=40,frac_d=0.8,data=df)
#'
#' @importFrom stats median sd complete_cases
#' @export
function roughness_parameters(method=c(Val(:canopy_height),Val(:canopy_height_LAI),Val(:wind_profile)),zh,
                                 frac_d=0.7,frac_{z0m}=0.1,LAI,zr,cd=0.2,hs=0.01,data,Tair="Tair",pressure="pressure",
                                 wind="wind",ustar="ustar",H="H",d=nothing,z0m=nothing,
                                 stab_roughness=true,stab_formulation=c(Val(:Dyer_1970),Val(:Businger_1971)),
                                 constants=BigleafConstants())

  method           = match_arg(method)
  stab_formulation = match_arg(stab_formulation)

  if (method == Val(:canopy_height))

    d      = frac_d*zh
    z0m    = frac_{z0m}*zh
    z0m_se = missing

elseif (method == Val(:canopy_height_LAI))

    X = cd * LAI
    d = 1.1 * zh * log(1 + X^(1/4))

    if (X >= 0 & X <= 0.2)
      z0m = hs + 0.3 * X^(1/2)
else
      z0m = 0.3 * zh * (1 - d/zh)
end
    z0m_se = missing

elseif (method == Val(:wind_profile))

    check_input(data,Tair,pressure,wind,ustar,H)

    if (isnothing(d))

      d = frac_d * zh

end

    if (stab_roughness)

      zeta  = stability_parameter(data=data,Tair=Tair,pressure=pressure,ustar=ustar,H=H,
                                   zr=zr,d=d,constants=constants)
      psi_m = stability_correction(zeta,formulation=stab_formulation)[,"psi_m"]
      z0m_all = (zr - d) * exp(-constants.k*wind / ustar - psi_m)

else

      z0m_all = (zr - d) * exp(-constants.k*wind / ustar)

end

    z0m_all[z0m_all > zh] = missing

    z0m    = median(z0m_all,na_rm=true)
    z0m_se = constants.se_median] * (sd(z0m_all,na_rm=true) / sqrt(length(z0m_all[complete_cases(z0m_all))))

end

  return(DataFrame(d,z0m,z0m_se))
end

#' Wind Speed at Given Heights in the Surface Layer
#'
#' Wind speed at a given height above the canopy estimated from single-level
#'              measurements of wind speed.
#'
#' - data      DataFrame or matrix containing all required variables
#' - z         Heights above ground for which wind speed is calculated.
#'                  Needs to be of same length as `data` or of length 1
#' - Tair      Air temperature (deg C)
#' - pressure  Atmospheric pressure (kPa)
#' - ustar     Friction velocity (m s-1)
#' - H         Sensible heat flux (W m-2)
#' - wind      Wind speed at height zr (m s-1); only used if `stab_correction = true`
#' - zr        Instrument (reference) height (m)
#' - zh        Canopy height (m)
#' - d         Zero-plane displacement height (-)
#' - frac_d    Fraction of displacement height on canopy height (-);
#'                  only used if `d` is not available
#' - z0m       Roughness length (m), optional; only used if `stab_correction = false` (default=0.1)
#' - frac_{z0m}  Fraction of roughness length on canopy height (-), optional; only used if `z0m` is not provided.
#'                  Default is 0.1.
#' - estimate_z0m Should `z0m` be estimated from the logarithmic wind profile? If `true` (the default),
#'                     arguments `z0m` and `frac_{z0m}` are ignored.
#'                     See [`roughness_parameters`](@ref) for details.
#' - stab_correction Should stability correction be applied? Defaults to `true`
#' - stab_formulation Stability correction function used (If `stab_correction = true`).
#'                         Either `Val(:Dyer_1970)` or `Val(:Businger_1971)`.
#' - constants k - von-Karman constant (-)
#'                  Kelvin - conversion degree Celsius to Kelvin
#'                  cp - specific heat of air for constant pressure (J K-1 kg-1)
#'                  g - gravitational acceleration (m s-2)
#'
#' # Details
 The underlying assumption is the existence of a logarithmic wind profile
#'          above the height d + z0m (the height at which wind speed mathematically reaches zero
#'          according to the Monin-Obhukov similarity theory).
#'          In this case, the wind speed at a given height z is given by:
#'
#'            ``u(z) = (ustar/k) * (ln((z - d) / z0m) - \\psi_m``
#'
#'          The roughness parameters zero-plane displacement height (d) and roughness length (z0m)
#'          can be approximated from [`roughness_parameters`](@ref). ``\\psi_m`` is omitted
#'          if `stab_correction = false` (not recommended). If `estimate_z0m = true`,
#'          z0m is first estimated from the wind profile equation and then used in the equation
#'          above for the calculation of `u(z)` (see e.g. Newman & Klein 2014).
#'
#' # Note
#' Note that this equation is only valid for z >= d + z0m, and it is not
#'       meaningful to calculate values closely above d + z0m. All values in `heights`
#'       smaller than d + z0m will return 0.
#'
#' # Value
 A vector of wind speed at heights `z`.
#'
#' #References
#' Monteith, J_L., Unsworth, M_H., 2008: Principles of Environmental Physics.
#'             3rd edition. Academic Press, London.
#'
#'             Newman, J_F., Klein, P_M., 2014: The impacts of atmospheric stability on
#'             the accuracy of wind speed extrapolation methods. Resources 3, 81-105.
#'
#' #See also
#' [`roughness_parameters`](@ref)
#'
#' ```@example; output = false
#' ```
#' heights = seq(18,40,2)  # heights above ground for which to calculate wind speed
#' df = DataFrame(Tair=25,pressure=100,wind=c(3,4,5),ustar=c(0.5,0.6,0.65),H=c(200,230,250))
#' ws = DataFrame(matrix(missing,ncol=length(heights),nrow=nrow(df)))
#' colnames(ws) = paste0(heights,"m")
#' for (i in seq_along(heights))
#'   ws[,i] = wind_profile(df,z=heights[i],zr=40,zh=25,d=16)
#' }
#'
#' @export
function wind_profile(data,z,Tair="Tair",pressure="pressure",ustar="ustar",H="H",wind="wind",
                         zr,zh,d=nothing,frac_d=0.7,z0m=nothing,frac_{z0m}=nothing,estimate_z0m=true,
                         stab_correction=true,stab_formulation=c(Val(:Dyer_1970),Val(:Businger_1971)),
                         constants=BigleafConstants())

  stab_formulation = match_arg(stab_formulation)

  check_input(data,ustar)

  ## determine roughness parameters
  if (isnothing(d))
    if (isnothing(frac_d))
      stop("Either 'd' or 'frac_d' must be specified")
end
    d = frac_d * zh
end

  if (isnothing(z0m) & !estimate_z0m)
    if (isnothing(frac_{z0m}))
      stop("Either 'z0m' or 'frac_{z0m}' must be specified if 'estimate_z0m' = false")
end
    z0m = frac_{z0m} * zh
end

  if (estimate_z0m)

    if (!isnothing(z0m) | !isnothing(frac_{z0m}))
      cat("Note that arguments 'z0m' and 'frac_{z0m}' are ignored if 'estimate_z0m' = true. z0m is
           calculated from the logarithmic wind_profile equation.",fill=true)
end

    check_input(data,Tair,pressure,wind,ustar,H)

    z0m = roughness_parameters(method=Val(:wind_profile),zh=zh,zr=zr,d=d,data=data,
                                Tair=Tair,pressure=pressure,wind=wind,ustar=ustar,H=H,
                                stab_roughness=true,stab_formulation=stab_formulation,
                                constants=constants)[,"z0m"]
end

  if ( any(z < (d + z0m) & !ismissing(d + z0m)) )
    @warn"function is only valid for heights above d + z0m! Wind speed for heights below d + z0m will return 0!")
end

  ## calculate wind speeds at given heights z
  if (stab_correction)

    zeta  = stability_parameter(data=data,Tair=Tair,pressure=pressure,ustar=ustar,H=H,
                                 zr=z,d=d,constants=constants)
    psi_m = stability_correction(zeta,formulation=stab_formulation)[,"psi_m"]
    wind_heights = pmax(0,(ustar / constants.k) * (log(pmax(0,(z - d)) / z0m) - psi_m))

else

    wind_heights = pmax(0,(ustar / constants.k) * (log(pmax(0,(z - d)) / z0m)))

end

  return(wind_heights)
end