module Test
export generate_dot, generate_sum

using LinearAlgebra, Random
using ..EFT: two_prod

"""
    (x, y, d, c) = generate_dot(n, c)

Generate two Float64 vectors whose dot product is ill-conditioned.

Inputs:
  n -- vectors size
  c -- target condition number

Results:
  x, y -- vectors of size n
  d    -- accurate dot product, rounded to nearest
  c    -- actual condition number of the dot product
"""
function generate_dot(n, c::T) where {T}
    R = Rational{BigInt}

    # Initialization
    x = zeros(T, n)
    y = zeros(T, n)

    # First half of the vectors:
    #   random numbers within a large exponent range
    n2 = div(n, 2)
    b = log2(c)

    e = rand(n2) .* b/2
    e[1]  = b/2 + 1           # Make sure exponents b/2
    e[n2] = 0                 # and 0 actually occur
    for i in 1:n2
        x[i] = (2*rand(T)-1) * 2^(e[i])
        y[i] = (2*rand(T)-1) * 2^(e[i])
    end

    # Second half of the vectors such that
    #   (*) log2( dot (x[1:i], y[1:i]) ) decreases from b/2 to 0
    δe = -b/(2*(n-n2-1))
    e = b/2:δe:0
    for i in eachindex(e)
        # Random x[i]
        cx = (2*rand()-1) * 2^(e[i])
        x[i+n2] = cx

        # y[i] chosen according to (*)
        cy = (2*rand()-1) * 2^(e[i])
        y[i+n2] = (cy - T(dot(R.(x), R.(y)))) / cx
    end

    # Random permutation of x and y
    perm = randperm(n)
    X = x[perm]
    Y = y[perm]

    # Dot product, rounded to nearest
    d = T(dot(R.(X), R.(Y)))

    # Actual condition number
    c = 2 * dot(abs.(X), abs.(Y)) / abs(d)

    (X,Y,d,c)
end

"""
    (x, s, c) = generate_sum(n, c)

Generate a Float64 vectors whose sum is ill-conditioned.

Inputs:
  n -- vectors size
  c -- target condition number

Results:
  x -- vectors of size n
  s -- accurate sum, rounded to nearest
  c -- actual condition number of the sum
"""
function generate_sum(n, c::T) where {T}
    R = Rational{BigInt}

    (x, y, _, _) = generate_dot(n÷2, c)

    z = (two_prod.(x, y)
         |> Iterators.flatten
         |> collect)

    # Complete if necessary
    if length(z) < n
        push!(z, rand())
    end

    z = shuffle(z)

    # Sum, rounded to nearest
    s = T(sum(R.(z)))

    # Actual condition number
    c = sum(abs.(z)) / abs(s)

    (z, s, c)
end

end


## stieltjes

function stieltjes{L<:PeriodicLine,SS}(S::Space{SS,L},f,z::Number,s...)
    S2=setdomain(S,Circle())
    stieltjes(S2,f,mappoint(domain(S),Circle(),z),s...)+hilbert(S2,f,-1)*π
end

function hilbert{L<:PeriodicLine,SS}(S::Space{SS,L},f,z::Number)
    S2=setdomain(S,Circle())
    hilbert(S2,f,mappoint(domain(f),Circle(),z))-hilbert(S2,f,-1)
end

# we use the fact that C^± (z^k + z^(k-1)) = z^k + z^(k-1) and 0
# for k > 0 and
# C^± (z^k + z^(k-1)) = 0 and -z^k - z^(k-1)
# for k < 0, the formula H = im*C^+  +  im*C^-
# and C± 1 = ±1/2  (understood as a PV integral) so that H 1 = 0

bandinds{DD<:PeriodicLine}(H::ConcreteHilbert{LaurentDirichlet{DD}})=0,0
rangespace{DD<:PeriodicLine}(H::ConcreteHilbert{LaurentDirichlet{DD}})=domainspace(H)

function getindex{T}(H::ConcreteHilbert{LaurentDirichlet{PeriodicLine{false,T}}},k::Integer,j::Integer)
    if k==j && iseven(k)
        -T(im)
    elseif k==j && isodd(k) && k > 0
        T(im)
    else
        zero(T)
    end
end


function update_weights(data, e, closest_same, closest_other, weights, weights_mult, m, k, max_f_vals,
						min_f_vals, dm_vals_same, dm_vals_other, features_msk_same, features_msk_other)

	for t = 1:size(data, 2)

		# Penalty term
		penalty = ((abs.(e[t] .- closest_same[:, t])/(max_f_vals[t] .- min_f_vals[t] .+ eps(Float64))) -
				   dm_vals_same[:,t])[features_msk_same[:, t]]
		penalty = sum(penalty)

		# Reward term
		reward = ((weights_mult .* (abs.(e[t] .- closest_other[:, t])/(max_f_vals[t] .- min_f_vals[t] .+ eps(Float64)))) -
				  dm_vals_other[:,t])[features_msk_other[:, t]]
		reward = sum(reward)

		# Weights update
		weights[t] = weights[t] - penalty/(m*k) + reward/(m*k)
	end

	# Return updated weights.
	return vec(weights)
end


#=
An example implementaion of an Autoencoder inspired by
https://github.com/alrojo/tensorflow-tutorial/blob/master/lab5_AE/lab5_AE.ipynb
This example trains an undercomplete autoencoder with a bottleneck of three hidden neurons, the activation of which can be used for a low-dimensional visualization of the 768-dimensional MNIST data.

The plotting functionality require Plots.jl and a backend, e.g., PyPlot.jl or GR.jl
`Pkg.add("Plots")`
`Pkg.add("PyPlot")` or `Pkg.add("GR")`
=#
using MNIST, Plots, TensorFlow, ValueHistories

function showdigit(x; kwargs...)
    d = reshape(x,28,28)
    heatmap(d; yflip=true, color=:grays, kwargs...)
end

function printtime(t0)
    dt = time()-t0
    "$(Int(dt÷60)):$(round(Int,dt % 60))"
end

num_features = 28^2

include(Pkg.dir("TensorFlow","examples","mnist_loader.jl"))
include(Pkg.dir("TensorFlow","src","layers","fully_connected.jl"))
loader = DataLoader()
session = Session(Graph())

x = placeholder(Float32, shape=[-1, num_features], name="x_pl")#/255
variable_scope("ae", initializer=Distributions.Normal(0, .01)) do
    global W1 = get_variable("weights1", [num_features, 512], Float32)
    global W2 = get_variable("weights2", [512, 3], Float32)
    global W1o = get_variable("weights1o", [512, num_features], Float32)
    global W2o = get_variable("weights2o", [3, 512], Float32)
end

B1 = Variable(0.01ones(Float32,512), name="bias1")
B2 = Variable(0.01ones(Float32,3), name="bias2")
B1o = Variable(0.01ones(Float32,num_features), name="bias1o")
B2o = Variable(0.01ones(Float32,512), name="bias2o")

# These lines define the network structure
l_enc = x*W1 + B1 |> nn.relu
l_z = l_enc*W2 + B2
l_dec = l_z*W2o + B2o |> nn.relu
l_out = l_dec*W1o + B1o |> nn.sigmoid

# loss_CE = -reduce_sum(x.*log(l_out+1e-8) + (1-x).*log(1-l_out+1e-8))
loss_MSE = reduce_mean((x-l_out).^2)
train_step = train.minimize(train.AdamOptimizer(5e-4), loss_MSE)

run(session, global_variables_initializer())
saver = train.Saver()
checkpoint_path = mktempdir()

summary_writer = train.SummaryWriter(checkpoint_path)

testx, testy = load_test_set()
testx ./= 255
testy = sum(testy.*(0:9)',2)[:]

f = plot(legend=false)
history = MVHistory()
@time begin
    t0 = time()
    for epoch = 1:10_000
        batch,_ = next_batch(loader, 64)
        batch ./= 255 # Normalize input to [0,1]
        if epoch%500 == 0
            train_loss = run(session, loss_MSE, Dict(x=>batch))
            val_loss, center = run(session, [loss_MSE, l_z], Dict(x=>testx))
            push!(history,:loss_train, epoch, train_loss)
            push!(history,:loss_val, epoch, val_loss)
            plot(history, reuse=true)
            scatter3d(center[:,1],center[:,2],center[:,3], zcolor = testy, legend=false, title="Latent space", reuse=true)
            info("step $epoch, training loss $train_loss, time taken: $(printtime(t0))")
            train.save(saver, session, joinpath(checkpoint_path, "ae_mnist"), global_step=epoch)
        end
        run(session, train_step, Dict(x=>batch))
    end
end

test_loss, center, reconstruction = run(session, [loss_MSE, l_z, l_out], Dict(x=>testx))
info("test accuracy $test_loss")

# Plot som example reconstructions
offset = 0
plot(layout=(2,2*2))
for i = 1:4
    d = reshape(testx[i+offset,:],28,28)
    heatmap!(d, subplot=i,yflip=true, color=:grays, title="Input", colorbar = false)
    d = reshape(reconstruction[i+offset,:],28,28)
    heatmap!(d, subplot=i+4,yflip=true, color=:grays, title="Reconstruction", colorbar = false)
end


    import AlgebraicMultigrid: ruge_stuben, smoothed_aggregation,
    poisson, aspreconditioner

import IterativeSolvers: cg

function test_cycles()

    A = poisson((50,50))
    b = A * ones(size(A,1))

    reltol = 1e-8

    for method in [ruge_stuben, smoothed_aggregation]
        ml = method(A)

        for cycle in [AlgebraicMultigrid.V(),AlgebraicMultigrid.W(),AlgebraicMultigrid.F()]
            x,convhist = solve(ml, b, cycle; reltol = reltol, log = true)

            @debug "number of iterations for $cycle using $method: $(length(convhist))"
            @test norm(b - A*x) < reltol * norm(b)
        end

        for cycle in [AlgebraicMultigrid.V(),AlgebraicMultigrid.W(),AlgebraicMultigrid.F()]
            p = aspreconditioner(ml,cycle)

            x,convhist = cg(A, b, Pl = p; reltol = reltol, log = true)
            @debug "CG: number of iterations for $cycle using $method: $(convhist.iters)"
            @test norm(b - A*x) <= reltol * norm(b)
        end
    end
end


"""
    DifferentialRKHSMethod1d(epsilon, lambda)

Differential loss reproducing kernel Hilbert space (RKHS) method from [1], Section III.

### Parameters

- `epsilon`: variance of the Gaussian kernel, Eq. 17 in [1]
- `lambda`: regularization parameter, Eq. 20 in [1]

### Reference

[1] Radhakrishnan, A. & and Meyn, S. (2018). Feedback particle filter design using a differential-loss reproducing kernel Hilbert space. 2018 Annual American Control Conference (ACC). IEEE. https://doi.org/10.23919/ACC.2018.8431689
"""
struct DifferentialRKHSMethod <: GainEstimationMethod
    epsilon::Float64
    lambda::Float64
end

# Helper functions for DifferentialRKHSMethod1d
function DifferentialRKHS_MakeMatrices(vec::AbstractVector, eps::Number)
    # constructs kernel matrices for DifferentialRKHSMethod1d
    L = length(vec)
    eps1 = 2*eps
    eps2 = eps^2

    mat = zeros(Float64, L, L, 3)

    @inbounds @simd for i in 1:L
                mat[i,i,1] = 1
                mat[i,i,3] = 1/eps
                @simd for j in i+1:L
                    diff = vec[j] - vec[i]
                    sq = diff*diff
                    K = exp(-sq/eps1)
                    Kx = diff*K/eps
                    Kxy = (eps - sq) * K / eps2
                    mat[i,j,1] = mat[j,i,1] = K
                    mat[i,j,2] = Kx
                    mat[j,i,2] = -Kx
                    mat[i,j,3] = mat[j,i,3] = Kxy
                end
    end

    return view(mat, :, :, 1), view(mat, :, :, 2), view(mat, :, :, 3)
end

function DifferentialRKHS_M_and_b!(A::AbstractMatrix, B::AbstractMatrix, C::AbstractMatrix, v::AbstractVector, lamb::Number)
    # constructs matrix M and vector b for DifferentialRKHSMethod1d
    L = size(A,1)

    b = zeros(Float64, 2*L)
    LinearAlgebra.mul!(view(b,   1:L  ), A, v)
    LinearAlgebra.mul!(view(b, L+1:2*L), B, v)

    B2 = B*B
    BC = B*C
    C2 = C*C
    LinearAlgebra.lmul!(lamb, A)
    LinearAlgebra.lmul!(lamb, B)
    LinearAlgebra.lmul!(lamb, C)

    M = zeros(Float64, 2*L, 2*L)
    @simd for j in 1:L
        @simd for i in 1:L
            @inbounds M[i  ,j  ] =  A[i,j] - B2[i,j]
            @inbounds M[i+L,j  ] =  B[i,j] + BC[j,i]
            @inbounds M[i  ,j+L] = -B[i,j] - BC[i,j]
            @inbounds M[i+L,j+L] =  C[i,j] + C2[i,j]
        end
    end

    return M, b
end

function solve!(eq::PoissonEquation, method::DifferentialRKHSMethod)
    eps = method.epsilon
    N = no_of_particles(eq)
    lamb = N * method.lambda

    H = Htilde(eq)[1,:]

    # Gaussian kernel and partial derivative matrices
    K, Kx, Kxy = DifferentialRKHS_MakeMatrices(eq.positions[1,:], eps)

    # compute M and b
    M, b = DifferentialRKHS_M_and_b!(K, Kx, Kxy, H, lamb) # warning: this function multiplies K, Kx, and Kxy by lamb

    # solve linear system
    beta = M \ b

    # write potential and gain
    eq.potential[1,:] = ( K*view(beta, 1:N) - Kx*view(beta, N+1:2*N) ) / lamb
    broadcast!(-, eq.potential, eq.potential, Statistics.mean(eq.potential, dims=2))
    eq.gain[1,:,1] = ( Kx * view(beta, 1:N) + Kxy * view(beta, N+1:2*N) ) / lamb
end


function evaluate_fitness(state::Matrix{Float64},evo_config::POOLConfig)
    colors = unique(state)
    filter!(e->e!=0,colors)
    n_color = size(colors,1)

    matches = 0

    for j in 1:size(state,2)
        for i in 1:size(state,1)

            if state[i,j] == evo_config.flag[i,j]
                matches += 1
            end
        end
    end

    return exp(-(n_color - 3)^2)*matches

end

function evaluate_fitness(state::Matrix{Float64},evo_config::EMAPConfig)

    colors = unique(state)
    filter!(e->e!=0,colors)
    n_color = size(colors,1)

    if n_color in evo_config.families

        section_step = Int(floor(size(state,2)/n_color))

        choices = permutations(colors) |> collect

        results = []

        for layout in choices

            matches = 0

            start = 1

            for section in 1:n_color

                for j in (1 + (start - 1)*section_step):min(start*section_step,size(state,2))
                    for i in 1:size(state,1)

                        if state[i,j] == layout[section]
                            matches += 1
                        end
                    end
                end

                start += 1
            end

            push!(results,matches)
        end

        fingerprint = 1

        for i in colors
            if i == 1
                fingerprint += 2^3
            elseif i == 2
                fingerprint += 2^2
            elseif i == 3
                fingerprint += 2^1
            elseif i == 4
                fingerprint += 2^0
            else
                fingerprint = "error"
            end
        end

        return maximum(results),fingerprint
    else
        return 0,1
    end

end


@testset "343.integer-break.jl" begin
    for (i, j) in zip([2, 4, 5, 6, 7, 8, 9, 12, 46, 53],
            [1, 4, 6, 9, 12, 18, 27, 81, 19131876, 258280326])
        @test integer_break(i) == j
    end
end


# Copyright (c) 2021, Rolfe Power

"""
    DimensionSet

Struct representing fundamental quantities in an absolute dimension system.
"""
struct DimensionSet
    mass::Float64
    length::Float64
    time::Float64
end

"""
    mass(::DimensionSet)

Get fundamental mass quantity
"""
mass(ds::DimensionSet) = ds.mass

"""
    length(::DimensionSet)

Get fundamental length quantity
"""
Base.length(ds::DimensionSet) = ds.length

"""
    time(::DimensionSet)

Get fundamental time quantity
"""
Base.time(ds::DimensionSet) = ds.time

"""
    velocity(ds)

Return the velocity found as the ratio of length and time.
"""
velocity(ds) = length(ds) / time(ds)

"""
    acceleration(ds)

Return the acceleration found as the ratio of velocity and time.
"""
acceleration(ds) = velocity(ds) / time(ds)

"""
    area(ds)

Return the velocity found as the square of length.
"""
area(ds) = length(ds)^2

"""
    volume(ds)

Return the velocity found as the cube of length.
"""
volume(ds) = area(ds) * length(ds)

"""
    force(ds)

Return the force found as the product of acceleration and mass
"""
force(ds) = acceleration(ds) * mass(ds)

"""
    density(ds)

Return the acceleration found as the ratio of mass and volume.
"""
density(ds) = mass(ds) / volume(ds)

"""
    DimensionSet(value; mass=value, length=value, time=value)

Construct a [`DimensionSet`](@ref) specifying (optionally) `mass`, `length`, and `time`.

!!! note

    In addition to constructing the [`DimensionSet`](@ref), this function also validates
    the input arguments. Input dimensional quantities must be positive values!

## Examples
```jldoctest
julia> DimensionSet(2)
Pat.DimensionSet(2.0, 2.0, 2.0)

julia> DimensionSet(mass=2.0, length=3.0, time=4.0)
Pat.DimensionSet(2.0, 3.0, 4.0)

julia> DimensionSet(12.0, time=0.5)
Pat.DimensionSet(12.0, 12.0, 0.5)

julia> DimensionSet(-1.0)
ERROR: DomainError with -1.0:
Non-positive mass
[...]
```
"""
function DimensionSet(value=1.0; mass=value, length=value, time=value)
    mass   > 0 || throw(DomainError(mass, "Non-positive mass"))
    length > 0 || throw(DomainError(length, "Non-positive length"))
    time   > 0 || throw(DomainError(time, "Non-positive time"))
    DimensionSet(
        mass,
        length,
        time
    )
end


export TwoStageCost, two_stage_method

struct TwoStageCost{F,F2,D} <: Function
  cost_function::F
  cost_function2::F2
  estimated_solution::D
  estimated_derivative::D
end

(f::TwoStageCost)(p) = f.cost_function(p)
(f::TwoStageCost)(p,g) = f.cost_function2(p,g)

function decide_kernel(kernel)
    if kernel == :Epanechnikov
        return Epanechnikov_kernel
    elseif kernel == :Uniform
        return Uniform_kernel
    elseif kernel == :Triangular
        return Triangular_kernel
    elseif kernel == :Quartic
      return Quartic_Kernel
    elseif kernel == :Triweight
      return Triweight_Kernel
    elseif kernel == :Tricube
      return Tricube_Kernel
    elseif kernel == :Gaussian
      return Gaussian_Kernel
    elseif kernel == :Cosine
      return Cosine_Kernel
    elseif kernel == :Logistic
      return Logistic_Kernel
    elseif kernel == :Sigmoid
      return Sigmoid_Kernel
    else
      return Silverman_Kernel
    end
end

function construct_t1(t,tpoints)
    T1 = []
    for i in 1:length(tpoints)
        push!(T1,[1 tpoints[i]-t])
    end
    foldl(vcat,T1)
end
function construct_t2(t,tpoints)
    T2 = []
    for i in 1:length(tpoints)
        push!(T2,[1 tpoints[i]-t (tpoints[i]-t)^2])
    end
    foldl(vcat,T2)
end
function construct_w(t,tpoints,h,kernel_function)
    n = length(tpoints)
    W = zeros(n)
    for i in 1:n
        W[i] = kernel_function((tpoints[i]-t)/h)/h
    end
    Matrix(Diagonal(W))
end
function construct_estimated_solution_and_derivative!(estimated_solution,estimated_derivative,e1,e2,data,kernel_function,tpoints,h,n)
  for i in 1:n
      T1 = construct_t1(tpoints[i],tpoints)
      T2 = construct_t2(tpoints[i],tpoints)
      W = construct_w(tpoints[i],tpoints,h,kernel_function)
      estimated_solution[:,i] = e1'*inv(T1'*W*T1)*T1'*W*data'
      estimated_derivative[:,i] = e2'*inv(T2'*W*T2)T2'*W*data'
  end
end

function two_stage_method(prob::DiffEqBase.DEProblem,tpoints,data;kernel= :Epanechnikov,
                          loss_func = L2Loss,mpg_autodiff = false,
                          verbose = false,verbose_steps = 100,
                          autodiff_prototype = mpg_autodiff ? zeros(length(prob.p)) : nothing,
                          autodiff_chunk = mpg_autodiff ? ForwardDiff.Chunk(autodiff_prototype) : nothing)
    f = prob.f
    n = length(tpoints)
    h = (n^(-1/5))*(n^(-3/35))*((log(n))^(-1/16))
    estimated_solution = zeros(size(data)[1],n)
    estimated_derivative = zeros(size(data)[1],n)
    kernel_function = decide_kernel(kernel)
    e1 = [1;0]
    e2 = [0;1;0]
    construct_estimated_solution_and_derivative!(estimated_solution,estimated_derivative,e1,e2,data,kernel_function,tpoints,h,n)
    # Step - 2
    cost_function = function (p)
        du = similar(prob.u0, promote_type(eltype(prob.u0), eltype(p)))
        sol = Vector{typeof(du)}(undef,n)
        f = prob.f
        for i in 1:n
          est_sol = @view estimated_solution[:,i]
          f(du,est_sol,p,tpoints[i])
          sol[i] = copy(du)
        end
        sqrt(sum(abs2,vec(estimated_derivative)[i] - vec(VectorOfArray(sol))[i] for i in 1:length(vec(estimated_derivative))))
    end

    if mpg_autodiff
      gcfg = ForwardDiff.GradientConfig(cost_function, autodiff_prototype, autodiff_chunk)
      g! = (x, out) -> ForwardDiff.gradient!(out, cost_function, x, gcfg)
    else
      g! = (x, out) -> Calculus.finite_difference!(cost_function,x,out,:central)
    end
    if verbose
      count = 0 # keep track of # function evaluations
    end
    cost_function2 = function (p,grad)
      if length(grad)>0
        g!(p,grad)
      end
      loss_val = cost_function(p)
      if verbose
        count::Int += 1
        if mod(count,verbose_steps) == 0
          println("Iteration: $count")
          println("Current Cost: $loss_val")
          println("Parameters: $p")
        end
      end
      loss_val
    end

    return TwoStageCost(cost_function, cost_function2, estimated_solution, estimated_derivative)
end


# ---
# title: 376. Wiggle Subsequence
# id: problem376
# author: Indigo
# date: 2020-12-12
# difficulty: Medium
# categories: Dynamic Programming, Greedy
# link: <https://leetcode.com/problems/wiggle-subsequence/description/>
# hidden: true
# ---
#
# A sequence of numbers is called a **wiggle sequence** if the differences
# between successive numbers strictly alternate between positive and negative.
# The first difference (if one exists) may be either positive or negative. A
# sequence with fewer than two elements is trivially a wiggle sequence.
#
# For example, `[1,7,4,9,2,5]` is a wiggle sequence because the differences
# `(6,-3,5,-7,3)` are alternately positive and negative. In contrast,
# `[1,4,7,2,5]` and `[1,7,4,5,5]` are not wiggle sequences, the first because
# its first two differences are positive and the second because its last
# difference is zero.
#
# Given a sequence of integers, return the length of the longest subsequence
# that is a wiggle sequence. A subsequence is obtained by deleting some number
# of elements (eventually, also zero) from the original sequence, leaving the
# remaining elements in their original order.
#
# **Example 1:**
#
#
#
#     Input: [1,7,4,9,2,5]
#     Output: 6
#     Explanation: The entire sequence is a wiggle sequence.
#
# **Example 2:**
#
#
#
#     Input: [1,17,5,10,13,15,10,5,16,8]
#     Output: 7
#     Explanation: There are several subsequences that achieve this length. One is [1,17,10,13,10,16,8].
#
# **Example 3:**
#
#
#
#     Input: [1,2,3,4,5,6,7,8,9]
#     Output: 2
#
# **Follow up:**
# Can you do it in O( _n_ ) time?
#
#
## @lc code=start
using LeetCode
function wiggle_max_length(nums::Vector{Int})::Int
    len = length(nums)
    if len < 2
        return len
    end
    pre_diff = nums[2] - nums[1]
    res = (pre_diff == 0 ? 1 : 2)
    for i in 3:len
        diff = nums[i] - nums[i - 1]
        if diff > 0 && pre_diff <= 0 || diff < 0 && pre_diff > 0
            res += 1
            pre_diff = diff
        end
    end
    return res
end
## @lc code=end


const p2 = 2
function pow_array2(x::Vector{Float64})
s = 0.0
for y in x
s = s + y^p2
end
return s
end


using MultiScaleArrays
using OrdinaryDiffEq, DiffEqBase, Test, StochasticDiffEq, Statistics

struct Cell2{B} <: AbstractMultiScaleArrayLeaf{B}
    values::Vector{B}
end
struct Population2{T<:AbstractMultiScaleArray,B<:Number} <: AbstractMultiScaleArrayHead{B}
    nodes::Vector{T}
    values::Vector{B}
    end_idxs::Vector{Int}
end

cell1 = Cell2([1.0; 2.0; 3.0])
cell2 = Cell2([4.0; 5])
cell3 = Cell2([3.0; 2.0; 5.0])
cell4 = Cell2([4.0; 6])
pop = construct(Population2, deepcopy([cell1, cell2, cell3, cell4]))

function cell_ode4(dcell,cell,p,t)
  m = mean(cell)
  for i in eachindex(cell)
    dcell[i] = -m*cell[i]
  end
end

function f4(dembryo,embryo,p,t)
  for (cell, dcell) in LevelIter(1,embryo, dembryo)
    cell_ode4(dcell,cell,p,t)
  end
end

tstop = [0.5]

condition = function (u, t, integrator)
  t ∈ tstop
end

affect! = function (integrator)
  add_node!(integrator, integrator.u.nodes[1])
end

growing_cb = DiscreteCallback(condition, affect!)

prob = ODEProblem(f4, deepcopy(pop), (0.0, 1.0))

add_node!(pop,pop.nodes[1])

sol = solve(prob, Tsit5(), callback=growing_cb, tstops=tstop)

sol = solve(prob, Rosenbrock23(), callback=growing_cb,tstops=tstop)

@test length(sol[end]) == 13

affect_del! = function (integrator)
    remove_node!(integrator, 1)
end

shrinking_cb = DiscreteCallback(condition, affect_del!)

prob = ODEProblem(f4, deepcopy(pop), (0.0, 1.0))
sol = solve(prob, Tsit5(), callback=shrinking_cb, tstops=tstop)

prob = ODEProblem(f4, deepcopy(pop), (0.0, 1.0))
sol = solve(prob, Rosenbrock23(),callback=shrinking_cb,tstops=tstop)
@test length(sol[end]) == 10

println("Do the SDE Part")

function g4(du,u,p,t)
    for i in eachindex(u)
        du[i] = 0.1u[i]
    end
end

prob = SDEProblem(f4, g4, deepcopy(pop), (0.0, 1.0))

sol = solve(prob, SOSRI(), callback=growing_cb, tstops=tstop)

@test length(sol[end]) == 16

sol = solve(prob, SOSRI(), callback=shrinking_cb, tstops=tstop)

@test length(sol[end]) == 10


log_temperature(t) = 1 / log(t)

constant_temperature(t) = 1.0

function default_neighbor!(x::AbstractArray{T}, x_proposal::AbstractArray) where T
    @assert size(x) == size(x_proposal)
    for i in 1:length(x)
        @inbounds x_proposal[i] = x[i] + T(randn()) # workaround because all types might not have randn
    end
    return
end

struct SimulatedAnnealing{Tn, Ttemp} <: ZerothOrderOptimizer
    neighbor!::Tn
    temperature::Ttemp
    keep_best::Bool # not used!?
end

"""
# SimulatedAnnealing
## Constructor
```julia
SimulatedAnnealing(; neighbor = default_neighbor!,
                     temperature = log_temperature,
                     keep_best::Bool = true)
```

The constructor takes 3 keywords:
* `neighbor = a!(x_proposed, x_current)`, a mutating function of the current `x`,
and the proposed `x`
* `T = b(iteration)`, a function of the current iteration that returns a temperature
* `p = c(f_proposal, f_current, T)`, a function of the current temperature, current
function value and proposed function value that returns an acceptance probability

## Description
Simulated Annealing is a derivative free method for optimization. It is based on the
Metropolis-Hastings algorithm that was originally used to generate samples from a
thermodynamics system, and is often used to generate draws from a posterior when doing
Bayesian inference. As such, it is a probabilistic method for finding the minimum of a
function, often over a quite large domains. For the historical reasons given above, the
algorithm uses terms such as cooling, temperature, and acceptance probabilities.
"""
SimulatedAnnealing(;neighbor = default_neighbor!,
                    temperature = log_temperature,
                    keep_best::Bool = true) =
  SimulatedAnnealing(neighbor, temperature, keep_best)

Base.summary(::SimulatedAnnealing) = "Simulated Annealing"

mutable struct SimulatedAnnealingState{Tx,T} <: ZerothOrderState
    x::Tx
    iteration::Int
    x_current::Tx
    x_proposal::Tx
    f_x_current::T
    f_proposal::T
end
# We don't have an f_x_previous in SimulatedAnnealing, so we need to special case these
pick_best_x(f_increased, state::SimulatedAnnealingState) = state.x
pick_best_f(f_increased, state::SimulatedAnnealingState, d) = value(d)

function initial_state(method::SimulatedAnnealing, options, d, initial_x::AbstractArray{T}) where T

    value!!(d, initial_x)

    # Store the best state ever visited
    best_x = copy(initial_x)
    SimulatedAnnealingState(copy(best_x), 1, best_x, copy(initial_x), value(d), value(d))
end

function update_state!(nd, state::SimulatedAnnealingState{Tx, T}, method::SimulatedAnnealing) where {Tx, T}

    # Determine the temperature for current iteration
    t = method.temperature(state.iteration)

    # Randomly generate a neighbor of our current state
    method.neighbor!(state.x_current, state.x_proposal)

    # Evaluate the cost function at the proposed state
    state.f_proposal = value(nd, state.x_proposal)

    if state.f_proposal <= state.f_x_current
        # If proposal is superior, we always move to it
        copyto!(state.x_current, state.x_proposal)
        state.f_x_current = state.f_proposal

        # If the new state is the best state yet, keep a record of it
        if state.f_proposal < value(nd)
            nd.F = state.f_proposal
            copyto!(state.x, state.x_proposal)
        end
    else
        # If proposal is inferior, we move to it with probability p
        p = exp(-(state.f_proposal - state.f_x_current) / t)
        if rand() <= p
            copyto!(state.x_current, state.x_proposal)
            state.f_x_current = state.f_proposal
        end
    end

    state.iteration += 1
    false
end


using OrdinaryDiffEq, ForwardDiff, Zygote, Test
using DiffEqSensitivity
using Random

function dt!(du, u, p, t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end

n_par=3
Random.seed!(2)
u0=rand(2,n_par)
u0[:,1] = [1.0,1.0]
tspan = (0.0, 10.0)
p = [2.2, 1.0, 2.0, 0.4]
prob_ode = ODEProblem(dt!, u0[:,1], tspan)

function test_loss(p1,prob)

	function prob_func(prob, i, repeat)
		@show i
		remake(prob,u0=u0[:,i])
	end

	#define ensemble problem
	ensembleprob = EnsembleProblem(prob,prob_func = prob_func)

	u = Array(solve(ensembleprob, Tsit5(), EnsembleThreads(), trajectories=n_par,
	p=p1,
	sensealg = ForwardDiffSensitivity(),
	saveat = 0.1, dt=0.001))[:,end,:]
	loss=sum(u)
	return loss
end

test_loss(p,prob_ode)

@time gs = Zygote.gradient(p) do p
	test_loss(p,prob_ode)
end
@test gs[1] isa Vector

### https://github.com/SciML/DiffEqFlux.jl/issues/595

function fiip(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end

p = [1.5,1.0,3.0,1.0]; u0 = [1.0;1.0]
prob = ODEProblem(fiip,u0,(0.0,10.0),p)
sol = solve(prob,Tsit5())

function sum_of_solution(x)
    _prob = remake(prob,u0=x[1:2],p=x[3:end])
    sum(solve(_prob,Tsit5(),saveat=0.1))
end
Zygote.gradient(sum_of_solution,[u0;p])

# Testing ensemble problem. Works with ForwardDiff. Does not work with Zygote.
N = 3
eu0 = rand(N,2)
ep = rand(N,4)

ensemble_prob = EnsembleProblem(prob, prob_func=(prob, i, repeat)->remake(prob, u0=eu0[i,:], p=ep[i,:],saveat=0.1))
esol = solve(ensemble_prob, Tsit5(), trajectories=N)

cache = Ref{Any}()

function sum_of_e_solution(p)
    ensemble_prob = EnsembleProblem(prob, prob_func=(prob, i, repeat)->remake(prob, u0=eu0[i,:],p=p[i,:],saveat=0.1))
    sol = solve(ensemble_prob,Tsit5(),EnsembleSerial(),trajectories=N,abstol=1e-12,reltol=1e-12)
    z = Array(sol[1])
    cache[] = sol[1].t
    sum(z) # just test for the first solutions, gradients should be zero for others
end

sum_of_e_solution(ep)

x = ForwardDiff.gradient(sum_of_e_solution, ep)
y = Zygote.gradient(sum_of_e_solution, ep)[1] # Zygote second to test cache of forward pass
@test x ≈ y
@test cache[] == 0:0.1:10.0 # test prob.kwargs is forwarded


using Test

ns = [14, 1, 17, 0, 3, 20]

function saidwhen(ns)
    return Dict([(n, i) for (i, n) in enumerate(ns)])
end

function play(sofar, target)
    ns = copy(sofar)
    said = saidwhen(ns)
    turn = length(sofar)

    while turn < target
        last = ns[turn]
        if !haskey(said, last)
            append!(ns, 0)
            said[last] = turn
        else
            append!(ns, turn - said[last])
            said[last] = turn
        end
        turn += 1
    end
    return ns
end

@test last(play([0, 3, 6], 10)) == 0
@test last(play([1, 3, 2], 2020)) == 1
@test last(play([2, 1, 3], 2020)) == 10
@test last(play([1, 2, 3], 2020)) == 27
@test last(play([2, 3, 1], 2020)) == 78
@test last(play([3, 2, 1], 2020)) == 438
@test last(play([3, 1, 2], 2020)) == 1836

p1 = last(play(ns, 2020))
p2 = last(play(ns, 30000000))

@assert(p1 == 387)
@assert(p2 == 6428)

println("-----------------------------------------------------------------------")
println("rambunctious recitation -- part one :: $p1")
println("rambunctious recitation -- part two :: $p2")
println("-----------------------------------------------------------------------")


using GeometryTypes, Reactive, GLVisualize
using GLPlot

"""
Lorenz function
"""
function lorenz(t0, a, b, c, h)
    Point3f0(
        t0[1] + h * a * (t0[2] - t0[1]),
        t0[2] + h * (t0[1] * (b - t0[3]) - t0[2]),
        t0[3] + h * (t0[1] * t0[2] - c * t0[3]),
    )
end
# step through the `time`
function lorenz(array::Vector, a=5.0,b=2.0,c=6.0,d=0.01)
    t0 = Point3f0(0.1, 0, 0)
    for i=eachindex(array)
        t0 = lorenz(t0, a,b,c,d)
        array[i] = t0
    end
    array
end

# create sliders to interact with the parameters of the lorenz function
slidera = GLPlot.play_widget(linspace(0.1f0, 40f0, 30))
sliderb = GLPlot.play_widget(linspace(0.1f0, 40f0, 30))
sliderc = GLPlot.play_widget(linspace(0.1f0, 40f0, 30))
sliderd = GLPlot.play_widget(linspace(0.001f0, 0.5f0, 30))

# N time steps
n = 10_000
# foldp registers a callback (in this case lorenz), which updates the points
# whenever the slider changes.
points = foldp(lorenz, Array(Point3f0, n), slidera, sliderb, sliderc, sliderd)
# we set the boundingbox to nothing, since we don't need it and don't want to
# calculate it for every update (which is the default)
glplot(points, :lines, preferred_camera = :perspective)
# you can also visualize the points
glplot(
    (Circle(Point2f0(0), 0.01f0), points),
    preferred_camera = :perspective
)
# the above renders an antia aliased nice looking circle which could have outlines and glows.
# this makes it relatively slow. For optimal performance, one might need the command below
# which draws one (or n) pixel per point
#glplot(points, :speed, boundingbox=nothing, preferred_camera=:perspective)


include("./src/solver.jl")
using PyPlot

#=
solver.jl:
- Julia version: 0.6.2
- Author: jerrymei
- Date: 2018-04-27
=#

k(x, y)=0.6 * sin(2 * pi * x) + 0.2 * sin(4 * pi * y) + 0.1 * sin(10 * pi * x + 8 * pi * y)
N_mesh=512

x=linspace(0,1 - 1 / N_mesh,N_mesh)
y=linspace(0,1 - 1 / N_mesh,N_mesh)
sol=zeros(N_mesh,N_mesh)
for i=1:N_mesh
    for j=1:N_mesh
        sol[i,j]=k(x[i],y[j])
    end
end

f=get_f(sol,1.)

u,err=solver(f,1.,1e-6)
figure(1)
plot_surface(linspace(0,1 - 1 / N_mesh,N_mesh),linspace(0,1 - 1 / N_mesh,N_mesh),u)
println("Norm of the distance between Du and f ", vecnorm(get_f(u,1.) - f)/vecnorm(f))

u,err=multigrid_solver(f,1.,1e-6)
figure(1)
plot_surface(linspace(0,1 - 1 / N_mesh,N_mesh),linspace(0,1 - 1 / N_mesh,N_mesh),u)
println("Norm of the distance between Du and f ", vecnorm(get_f(u,1.) - f)/vecnorm(f))


using LinearAlgebra
using Test

using QD

function approxeq(::Type{T}, tolfac, xs, r0, r1) where {T}
    tol = max(abs.(big.(xs))..., abs(r0), abs(r1))
    iseq = abs(r0 - r1) <= tolfac * tol * big(eps(T))
    if !iseq
        @show T tolfac xs r0 r1 abs(r0 - r1)
    end
    iseq
end

include("test128.jl")
include("test256.jl")


using Flux
using Flux: Tracker, mse, throttle
using Base.Iterators: repeated

function norm_input(x, min, max)
     # res = []
     res = zeros(size(x))
     for i=1:length(res)
          res[i] = (x[i]-min)/(max-min)
          # push!(res,[(x[i]-min)/(max-min)])
     end
     return res
end

function extract_output(y, min, max)
     res = deepcopy(y)
     for i=1:length(y)
          res[i] = y[i]*((max-min)+min)
     end
     return res
end

# global parameters
input_dim = 1
hidden_dim = 2
output_dim = 2
α = 0.1                 # training rate
iterations = 5000

#### dataset ####
# input (control value)
raw_input = [1 1 1 1 1 1 2 2 2 2 2 1 1 1 1]
# raw_input = [1 1 1 1 1 1 2 2 2 2 2 1 1 1 1; 1 1 1 1 1 1 2 2 2 2 2 1 1 1 1]

# output (rotation speed)
raw_labels = [15 15 15 15 15 15 30 30 30 30 30 15 15 15 15]
# raw_labels = [15 15 15 15 15 15 30 30 30 30 30 15 15 15 15; 15 15 15 15 15 15 30 30 30 30 30 15 15 15 15]

x = norm_input(raw_input,0,10)
y = norm_input(raw_labels,0,50)

dataset = collect(repeated((x,y),iterations))
# dataset = (x,y)

#### model ####

# >>> without using Flux
# W1 = param(randn(hidden_dim, input_dim))
# b1 = param(randn(hidden_dim))
# layer1(x) = W1 * x .+ b1
#     # hidden neurons in between the two layers
# W2 = param(randn(output_dim, hidden_dim))
# b2 = param(randn(hidden_dim))
# layer2(x) = W2 * x .+ b2
# model(x) = layer2(layer1(x))

# using CuArrays (GPU)
# W, b, x, y = cu.((W, b, x, y))

# >>> with Flux
model = Chain(
    Dense(input_dim, hidden_dim),
    Dense(hidden_dim, output_dim))

#### loss function ####
loss(x, y) = mse(model(x), y)

# dataset = zip(x, y)
# dataset = repeated((x, y),100)
# opt = SGD(params(m), 0.1)          # Gradient descent with learning rate
# evalcb = () -> @show(loss(x, y))   # output the progress

#### training ####
function update!(ps, η = .1)
  for w in ps
    w.data .-= w.grad .* η
    w.grad .= 0
  end
end

for i = 1:length(dataset)
    xs = dataset[i][1]
    ys = dataset[i][2]
    back!(loss(xs, ys))             # automatic backpropagation
    update!(params(model), α)
    # update!((W1, b1, W2, b2), α)

    if i%1000 == 0  # do not print every step
        @show loss(x, y)
    end
end

# Flux.train!(loss, dataset, opt, cb = throttle(evalcb, 0.5))

# output final error
# @show(loss(x, y))

# testing the results
# @show m([1,0,0])    # should be close to 1
# @show m([0,1,0])    # should be close to 0


function rotate(n::Int, target::Char)
    ascii = Int(target)

    if 97 <= ascii <= 122
        # Lower case letters
        code = ascii - 97
        code_1 = (code + n) % 26

        return Char(code_1 + 97)

    elseif 65 <= ascii <= 90
        # Upper case letters
        code = ascii - 65
        code_1 = (code + n) % 26

        return Char(code_1 + 65)

    else
        return target
    end

end

function rotate(n::Int, target::String)
    return String([rotate(n, c) for c in target])
end

for i in 0:26
    macro_name = Symbol("R$(i)_str")
    @eval begin
        macro $macro_name(p)
            rotate($i, p)
        end
    end
end


include("model.jl")
using Flux, Flux.Tracker
import Random

Random.seed!(0);

# Generate Data
N = 1000
x = randn(N)
b0 = .5
b1 = 2.0
p = sigmoid.(b0 .+ b1 * x)
y = (p .> rand(N)) * 1.0

minibatch_size = 500
vp = VP()
metrics = Vector{Metrics}()
loss(y, x) = -elbo(y, x, vp, N, metrics) / N

params = Tracker.Params([getfield(vp, fn) for fn in fieldnames(typeof(vp))])
# grads = Tracker.gradient(() -> loss(y, x), params)

opt = ADAM(1e-1)
niters = 1000
Random.seed!(3);

@time for i in 1:niters
  idx = sample(1:N, minibatch_size)
  Flux.train!(loss, params, [(y[idx], x[idx])], opt)
  if i % (niters / 10) == 0
    println("Progress: $(i)/$(niters) | ELBO: $(-loss(y, x).data)")
  end
end

println("b0_true: $(b0) | b1_true: $(b1)")
println("b0_mean: $(vp.b0[1].data) | b1_mean: $(vp.b1[1].data)")
println("b0_sd: $(exp(vp.b0[2].data)) | b1_sd: $(exp(vp.b1[2].data))")


@testset "nri" begin

  #example om from NRI paper
  om1 = [0 0 0 0 0; 0 2 0 0 1;
         0 0 0 1 0; 0 0 3 0 0;
         0 1 0 0 0];
  om1 = sparse(om1);

  fullTPs, segTPs = NRI.compute_TPs(om1)

  @test segTPs[1] == 0
  @test segTPs[2] == 1
  @test segTPs[3] == 0
  @test segTPs[4] == 3
  @test segTPs[5] == 0

  fullFPs, segFPs = NRI.compute_FPs(om1)

  @test segFPs[1] == 0
  @test segFPs[2] == 1
  @test segFPs[3] == 0
  @test segFPs[4] == 0
  @test segFPs[5] == 1

  fullFNs, segFNs = NRI.compute_FNs(om1)

  @test segFNs[1] == 0
  @test segFNs[2] == 2
  @test segFNs[3] == 0
  @test segFNs[4] == 0
  @test segFNs[5] == 0

  fullNRI, segNRI, segNRIw = NRI.nri(om1)

  @test segNRI[2] == 0.4
  @test segNRI[4] == 1
  @test segNRI[5] == 0

end


using Pkg, Random
using Dates
using PyPlot, Logging
import Base.show

struct Node
    id::Int16
    x::Float32
    y::Float32
end

function show(io::IO, n::Node)
    print(string("Node:", n.id))
end

import Base.copy

mutable struct Firefly
    path::Array{Node}
    cost::Float32
end

function show(io::IO, n::Firefly)
    print(string("Firefly:", n.cost, " "))
end

copy(s::Firefly) = Firefly(s.path, s.cost)

function arc_distance(f1::Firefly, f2::Firefly)::Int
    """Calculates the arc distance between two firefly paths.
    Example:
        Firefly i -> 1 2 4 6 3 5 7
        Firefly j -> 1 2 4 3 5 7 6
        1-2, 2-4, 3-5, 5-7 arcs in firefly i are also presented
        in firefly j. But 4-6, 6-3 arcs are missing in
        firefly j. so the arc distance between firefly i
        and firefly j is 2 here.
    """
    missmatched_pair_counts = 0
    f1_pairs = zip(f1.path[1:end-1], f1.path[2:end])
    f2_pairs = zip(f2.path[1:end-1], f2.path[2:end])

    for f1_pair in f1_pairs
        if f1_pair ∉ f2_pairs missmatched_pair_counts += 1 end
    end

    return missmatched_pair_counts
end

function attractiveness(source_node::Firefly, destination_node::Firefly, λ::Float32)::Float32
    """Firefly’s attractiveness is proportional to the light
    intensity seen by adjacent fireflies, we can now define the
    attractiveness β of a firefly depending on distance between
    two fireflies.
    """
    β = source_node.cost
    r, _ = hamming_distance(source_node, destination_node)
    return β * exp(-λ ^ (r^2))
end

function create_distance_matrix(nodes::Array{Node})
    """Creates and returns the matrix of distances between given nodes.
    Uses euclidean_distance function to calculate distance.
    """
    node_num = length(nodes)
    result_matrix = zeros(node_num, node_num)

    for n1 in nodes
        for n2 in nodes
            result_matrix[n1.id,n2.id] = haversine_distance(n1, n2)
        end
    end
    return result_matrix
end

function generate_mutated_fireflies(firefly::Firefly, distance_matrix, r::Int, k::Int)::Array{Firefly}
    """Generates mutated fireflies using inversion_mutation function.
    Returns the new generated firefly array.
    """
    return [inversion_mutation(copy(firefly), distance_matrix, r) for _ in 1:k]
end

function hamming_distance(firefly_1::Firefly, firefly_2::Firefly)
    """Calculates the hamming distance between two fireflies.
    """
    dist = 0
    dist_info = []
    for (i, n) in enumerate(firefly_1.path)
        if n.id != firefly_2.path[i].id
            dist += 1
            push!(dist_info, true)
        else
            push!(dist_info, false)
        end
    end
    return dist, dist_info
end

function haversine_distance(n1::Node, n2::Node)::Float32
    #converting decimal degree to radians
    lat1=n1.y
    long1=n1.x
    lat2=n2.y
    long2=n2.x
    lat1,long1,lat2,long2=map(deg2rad,[lat1,long1,lat2,long2])

    # Using Haversine formula for distance calculation
    dlat=lat2-lat1
    dlong=long2-long1
    a=sin(dlat/2)^2 + cos(lat1) * cos(lat2) * sin(dlong/2)^2
    c=2 * asin(sqrt(a))
    r=6371 #radius of earth in km
    return c*r
end

function init_firefly_paths(fireflies::Array{Firefly}, distance_matrix)
    """Initializes random paths for the given firefly array.
    """
    println("inside init_firefly_paths")
    for f in fireflies
        shuffle!(f.path)
        f.cost = path_cost(f, distance_matrix)
    end
end

function inversion_mutation(firefly::Firefly, distance_matrix, r::Int)::Firefly
    """Reverses the random part of given firefly path with
    given difference.
    """
    length_of_mutation = rand(2:r);
    max_len = length(firefly.path)-length_of_mutation
    index1 = rand(1:max_len)
    index2 = index1 + length_of_mutation
    mutated_path = reverse(firefly.path, index1, index2)
    new_f = copy(firefly)
    new_f.path = mutated_path
    new_f.cost = path_cost(new_f, distance_matrix)
    return new_f
end

function light_intensity(source_node::Firefly, destination_node::Firefly, λ::Float32)::Float32
    """Calculates the light intensity for the source firefly.
    """
    I_0 = source_node.cost
    r, _ = hamming_distance(source_node, destination_node)
    return I_0 * ℯ ^ (-λ * (r^2))
end

function move_firefly(f1::Firefly, f2::Firefly, r::Int)
    """Moving f1 to f2 less than r swap operations.
    """
    number_of_swaps = Int(ceil(rand() * (r-2) ))
    d, d_info = hamming_distance(f1, f2)

    while number_of_swaps > 0
        d, d_info = hamming_distance(f1, f2)
        random_index = rand([i for i in 1:length(d_info) if d_info[i]])
        value_to_copy = f2.path[random_index]
        index_to_move = findfirst(x -> x==value_to_copy, f1.path)

        if number_of_swaps == 1 && f1.path[index_to_move] == f2.path[random_index] && f1.path[random_index] == f2.path[index_to_move]
            break
        end

        f1.path[random_index], f1.path[index_to_move] = f1.path[index_to_move], f1.path[random_index]
        # println("swapped! ", f1.path[random_index], " -- ", f1.path[index_to_move])
        if f1.path[index_to_move] == f2.path[index_to_move] number_of_swaps -= 1 end
        number_of_swaps -= 1
    end
end

function path_cost(firefly::Firefly, distance_matrix)::Float32
    """Calculates the total path cost from given distance matrix.
    """
    path = firefly.path
    path_len = length(path) - 1
    p_cost = distance_matrix[path[end].id, path[1].id]
    for i in range(1, length=path_len)
        p_cost +=  distance_matrix[path[i].id, path[i+1].id]
    end

    return p_cost
end

function read_nodes(cluster_info)::Array{Node}
    """Reading nodes from giving file path.
    """
    nodes = Node[]
    for coord in cluster_info
            id = parse(Int16, coord[1])
            x_coord = parse(Float32, coord[2])
            y_coord = parse(Float32, coord[3])
            push!(nodes, Node(id, x_coord, y_coord))
    end
    return nodes
end

function plot_firefly(best_x, best_y, x, y, x_cost, best_cost, file_name, tsp_file_name)

    plot(best_x, best_y, "-")
    plot(x, y, "ro", markersize=2.0)
    xlabel("X Dimension")
    ylabel("Y Dimension")
    title("Firefly Cost: $best_cost", fontsize=10)
    suptitle("File: $file_name", fontweight="bold")
    fig_name = "graphs/best_firefly_path_$file_name-cost_$best_cost.png"
    savefig(fig_name)
    #println("Saved figure:",fig_name)
    clf()
    plot(x_cost)
    grid(true)
    ylabel("Path Cost")
    xlabel("Iteration Number")
    title("Firefly Iteration Graphic", fontsize=12)
    fig_name = "graphs/iteration_$tsp_file_name-cost_$best_cost.png"
    savefig(fig_name)
    clf()
end

function main(cluster_data, name)
    final_route=[]

    for ncluster in 1:length(cluster_data)

        file_name=name[ncluster]
        nodes = read_nodes(cluster_data[ncluster])
        distance_matrix = create_distance_matrix(nodes)
        # Checking if logs direcitory is available!
        # If is not create one!
        if !isdir("logs") mkdir("logs"); end

        tsp_file_name = file_name
        io_log = open("logs/$tsp_file_name.txt", "w+")

        # Setting Logger!
        logger = SimpleLogger(io_log)
        global_logger(logger)

        # Setting hyperparameters!
        LIGHT_ABSORPTION_COEFF = 0.2
        ATTRACTION_COEFF = 1
        ITERATION_NUMBER = 100
        POPULATION_NUMBER = 150
        NUMBER_OF_MUTATION = 10

        # Main firefly algorithm
        bests = []
        pop = [Firefly(copy(nodes), -1.0) for _ in 1:POPULATION_NUMBER]
        init_firefly_paths(pop, distance_matrix)
        println(length(pop), " Sized population created!")
        for t in 1:ITERATION_NUMBER
            new_pop = []
            for f1 in pop
                for f2 in pop
                    if f1.cost < f2.cost
                        mutated_f1s = generate_mutated_fireflies(f2, distance_matrix, arc_distance(f1, f2), NUMBER_OF_MUTATION)
                        new_pop = [new_pop ; mutated_f1s]
                        push!(new_pop, f2)
                    end
                end
            end
            [push!(pop, n) for n in new_pop]
            sorted_pop = sort!(pop, by=p->p.cost)
            while length(pop) != POPULATION_NUMBER
                pop!(pop)
            end
            push!(bests, sorted_pop[1])
            current_best_cost = copy(bests[end].cost)
            now = string(Dates.now())
            println("Step:",t, " - Best ", bests[end])
            @info("Step:$t, Population Size:$POPULATION_NUMBER, Mutation Number:$NUMBER_OF_MUTATION, Date:$now",
                  current_best_cost)
        end

        close(io_log)

        # Checking if graphs direcitory is available for saving graphs!
        # If is not create one!
        if !isdir("graphs") mkdir("graphs"); end

        # Plotting and printing the best solutions.
        best = sort(pop, by=p->p.cost)[1]
        best_x = [n.x for n in best.path]
        best_y = [n.y for n in best.path]
        best_id= [n.id for n in best.path]
        best_cost = best.cost
        x = [n.x for n in nodes]
        y = [n.y for n in nodes]
        x_cost = [b.cost for b in bests]
        plot_firefly(best_x, best_y, x, y, x_cost, best_cost, file_name, tsp_file_name)
        z=[best_id, best_x, best_y]
        push!(final_route, z)
    end
    return final_route
end


include("../src/star.jl")
using DistributedArrays

scalar_times_vector(a, x, w) = x[:] = a * w

#compute x=u'vw
utvw_star(u, v, w, x) = Star.scan(scalar_times_vector, sum(Star.reduce(dot, u, v)), x, w)

function utvw_star_test(u, v, w, x)
  #compute x=u'vw
  t = @elapsed utvw_star(u, v, w, x)

  #code to test the results
  a = convert(Array{Float64, 1}, u)
  b = convert(Array{Float64, 1}, v)
  c = convert(Array{Float64, 1}, w)
  d = convert(Array{Float64, 1}, x)

  e = zeros(Float64, length(a))
  t1 = @elapsed begin
    e = dot(a, b) * c
  end

  println("parallel time ", t, " serial time ", t1)
  println("rel norm ", norm(d - e)/norm(e))
end

function utvw_no_star(u, v, w, x)
  remote_ref = [@spawnat p dot(localpart(u), localpart(v)) for p in workers()]
  partial_dot_product = [fetch(r) for r in remote_ref]
  s = sum(partial_dot_product)
  @sync [@spawnat p scalar_times_vector(s, localpart(x), localpart(w)) for p in workers()]
end


import LineSearches

let
   methods = [ LBFGS(), ConjugateGradient(),
               LBFGS(extrapolate=true, linesearch = LineSearches.interpbacktrack!) ]
   msgs = ["LBFGS Default Options: ",  "CG Default Options: ",
          "LBFGS + Backtracking + Extrapolation: " ]

   println("--------------------")
   println("Rosenbrock Example: ")
   println("--------------------")
   rosenbrock(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
   for (method, msg) in zip(methods, msgs)
      results = Optim.optimize(rosenbrock, zeros(2), method)
      println(msg, "g_calls = ", results.g_calls, ", f_calls = ", results.f_calls)
   end

   println("--------------------------------------")
   println("p-Laplacian Example (preconditioned): ")
   println("--------------------------------------")
   plap(U; n=length(U)) = (n-1) * sum( (0.1 + diff(U).^2).^2 ) - sum(U) / (n-1)
   plap1(U; n=length(U), dU = diff(U), dW = 4 * (0.1 + dU.^2) .* dU) =
                           (n-1) * ([0.0; dW] - [dW; 0.0]) - ones(U) / (n-1)
   precond(x::Vector) = precond(length(x))
   precond(n::Number) = spdiagm( ( -ones(n-1), 2*ones(n), -ones(n-1) ),
                                 (-1,0,1), n, n) * (n+1)
   df = DifferentiableFunction( X->plap([0;X;0]),
                                (X, G)->copy!(G, (plap1([0;X;0]))[2:end-1]) )
   GRTOL = 1e-6
   N = 100
   initial_x = zeros(N)
   P = precond(initial_x)
   methods = [ LBFGS(P=P), ConjugateGradient(P=P),
         LBFGS(extrapolate=true, linesearch = LineSearches.interpbacktrack!, P=P) ]

   for (method, msg) in zip(methods, msgs)
      results = Optim.optimize(df, copy(initial_x), method)
      println(msg, "g_calls = ", results.g_calls, ", f_calls = ", results.f_calls)
   end
end


import StatsBase.cov
# ------------------------------------------------------------------------------------------------ #
#          Run a pca on the (time series of a) feature matrix and plot the first two PC's          #
# ------------------------------------------------------------------------------------------------ #
@userplot LowDim2
@recipe function f(h::LowDim2; method=principalcomponents, classlabels=false, features=Catch22.featureDims(h.args[1]))
    # show(size(F))
    F = h.args[1]
    # if h.args > 1
    #     Fs = h.args[2:end]
    # end
    fs = Catch22.featureDims(F)
    seriestype := scatter
    M = project(standardise(Array(F)), method)
    F′ = embed(M, standardise(Array(F)), 1:2)

    if M isa MultivariateStats.PCA
        PV = principalvars(M)
        PV = PV[1:2]./sum(PV)
        wmax = mapslices(x -> sum(abs.(x)), featureweights(M, 1:2), dims=1)
        weights, fis = findmax(featureweights(M)[:, 1:2], dims=1)
        PV = [round(x*100, sigdigits=2) for x ∈ PV]
        w = [round(100*weights[i]/wmax[i], sigdigits=2) for i ∈ 1:length(weights)]
        #xguide --> "$(PV[1])% Variance ($(w[1])% $(fs[fis[1]]))"
        #yguide --> "$(PV[2])% Variance ($(w[2])% $(fs[fis[1]]))"
        xguide --> "PC1 ($(PV[1])%)"
        yguide --> "PC2 ($(PV[2])%)"
    end

    if classlabels isa Bool && classlabels
        classlabels = Catch22.timeseriesDims(F)
        uniquelabels = unique(classlabels)

        for c ∈ uniquelabels
            @series begin
                idxs = classlabels .== c
                label := string(c)
                (x, y) = (F′[1, idxs], F′[2, idxs])
            end
        end
    else
        @series begin
            (x, y) = (F′[1, :], F′[2, :])
        end
        # if length(h.args) > 1
        #     for F ∈ Fs
        #         @series begin
        #             F = embed(M, standardise(Array(F), mean(Array(F), dims=2), StatsBase.std(Array(F), dims=2)), 1:2)
        #             (x, y) = (F[1, :], F[2, :])
        #         end
        #     end
        # end
    end
end

@userplot LowDim3
@recipe function f(h::LowDim3; method=principalcomponents, classlabels=true, features=Catch22.featureDims(h.args[1]))
    # show(size(F))
    F = h.args[1]
    fs = Catch22.featureDims(F)
    seriestype := scatter
    M = project(standardise(Array(F)), method)
    F′ = embed(M, standardise(Array(F)), 1:3)

    markersize --> 2
    markerstrokewidth --> 0
    size --> (600, 400)
    if M isa MultivariateStats.PCA
        PV = principalvars(M)
        PV = PV[1:3]./sum(PV)
        wmax = mapslices(x -> sum(abs.(x)), featureweights(M)[:, 1:3], dims=1)
        weights, fis = findmax(featureweights(M)[:, 1:3], dims=1)
        PV = [round(x*100, sigdigits=2) for x ∈ PV]
        w = [round(100*weights[i]/wmax[i], sigdigits=2) for i ∈ 1:length(weights)]
        #xguide --> "$(PV[1])% Variance ($(w[1])% $(fs[fis[1]]))"
        #yguide --> "$(PV[2])% Variance ($(w[2])% $(fs[fis[1]]))"
        xguide --> "PC1 ($(PV[1])%)"
        yguide --> "PC2 ($(PV[2])%)"
        zguide --> "PC3 ($(PV[3])%)"
    end

    if classlabels isa Bool && classlabels
        classlabels = Catch22.timeseriesDims(F)
        uniquelabels = unique(classlabels)

        for c ∈ uniquelabels
            @series begin
                idxs = classlabels .== c
                label := string(c)
                (x, y) = (F′[1, idxs], F′[2, idxs], F′[3, idxs])
            end
        end
    else
        @series begin
            (x, y) = (F′[1, :], F′[2, :], F′[3, :])
        end
    end
end

# ---------------- Plot high, test and zero-dim datasets and in the PC2 same space --------------- #
@userplot LowDimBaselines
@recipe function f(h::LowDimBaselines; method=principalcomponents, ellipsealpha=0.2, doscatter=true, grouplabels=false)
    # show(size(F))
    F = h.args[1:3] # We will whiten according to the first input

    seriestype := scatter
    M = project(F[1], method)
    F′ = embed.((M,), F, (1:2,))

    # if M isa MultivariateStats.PCA
    #     PV = principalvars(M)
    #     PV = PV[1:2]./sum(PV)
    #     wmax = mapslices(x -> sum(abs.(x)), featureweights(M, 1:2), dims=1)
    #     weights, fis = findmax(featureweights(M)[:, 1:2], dims=1)
    #     PV = [round(x*100, sigdigits=2) for x ∈ PV]
    #     w = [round(100*weights[i]/wmax[i], sigdigits=2) for i ∈ 1:length(weights)]
    #     #xguide --> "$(PV[1])% Variance ($(w[1])% $(fs[fis[1]]))"
    #     #yguide --> "$(PV[2])% Variance ($(w[2])% $(fs[fis[1]]))"
    #     xguide --> "PC1 ($(PV[1])%)"
    #     yguide --> "PC2 ($(PV[2])%)"
    # end

    color_palette = [:cornflowerblue, :crimson, :forestgreen]
    color_palette --> color_palette
    for i = 1:3
        @series begin
            markerstrokewidth --> 0
            label := nothing
            seriescolor --> color_palette[i]
            markeralpha --> exp(-0.002*size(F′[i], 2))
            (x, y) = (F′[i][1, :], F′[i][2, :])
        end
        @series begin
            seriesalpha := ellipsealpha
            seriescolor := color_palette[i]
            seriestype := :shape
            if grouplabels isa Array
                label := grouplabels[i]
            else
                label := grouplabels
            end
            linewidth := 2
            linecolor := color_palette[i]
            linealpha := 1.0
            framestyle --> :box
            μ, S = StatsPlots._covellipse_args(Array.((mean(F′[i], dims=2)[:], cov((F′[i]), dims=2))); n_std=1.5)
            θ = range(0, 2π; length=1000)
            A = S * [cos.(θ)'; sin.(θ)']
            (μ[1] .+ A[1,:], μ[2] .+ A[2,:])
        end
    end
end

@userplot BaselineScatters
@recipe function f(h::BaselineScatters; ellipsealpha=0.2, doscatter=true, grouplabels=false)
    F = h.args[1:end]

    seriestype := :scatter

    color_palette = [:cornflowerblue, :crimson, :forestgreen]
    color_palette --> color_palette
    for i = 1:length(h.args)
        @series begin
            markerstrokewidth --> 0
            label := nothing
            seriescolor --> color_palette[i]
            markeralpha --> exp(-0.002*size(F[i], 2))
            (x, y) = (F[i][1, :], F[i][2, :])
        end
        @series begin
            seriesalpha := ellipsealpha
            seriescolor := color_palette[i]
            seriestype := :shape
            if grouplabels isa Array
                label := grouplabels[i]
            else
                label := grouplabels
            end
            linewidth := 2
            linecolor := color_palette[i]
            linealpha := 1.0
            framestyle --> :box
            μ, S = StatsPlots._covellipse_args(Array.((mean(F[i], dims=2)[:], cov((F[i]), dims=2))); n_std=1.5)
            θ = range(0, 2π; length=1000)
            A = S * [cos.(θ)'; sin.(θ)']
            (μ[1] .+ A[1,:], μ[2] .+ A[2,:])
        end
    end
end


"""
QQQQ
"""
function maph_init_test()
    # desired_absorbptions = [0.3,0.5,0.1,0.08,0.02]
    # desired_means = [2.0,3.1,2.3,4.5,0.2]
    # out = MAPHDist(20,desired_absorbptions,desired_means,[1.1,2.3,2.0,0.33,1.5])
    # # @show model_size(out)
    # # data = [rand(out) for _ in 1:10^5]
    # # pi_est = [count((x)->x.a == i, data)/length(data) for i in 1:5]
    # # μ_est = [mean(first.(filter((x)->x.a == i, data))) for i in 1:5]
    # (absorption_probs(out) ≈ desired_absorbptions') || return false

    Λ₄, λ45, λ54, Λ₅ = 20, 17, 7, 10
    μ41, μ42, μ43, μ51, μ52, μ53 = 1, 1, 1, 1, 1, 1
    T_example = [-Λ₄ λ45; λ54 -Λ₅]
    T0_example = [μ41 μ42 μ43; μ51 μ52 μ53]
    initial_dist = [0.5,0.5]
    maph = MAPHDist(initial_dist', 3*T_example, 3*T0_example)
    @show absorption_probs(maph)
    @show mean(maph)

    # mean(out) ≈ desired_means || return false
    # @show pi_est
    # @show μ_est
    # @show pi_est .* μ_est
    return true
end


using IcanProj
using NetCDF
using DataFrames
using CSV
using Statistics

function compute_statistics(var_coarse, var_ref, mask)

    n = size(var_coarse, 2)

    nse = fill(0.0, n)
    r = fill(0.0, n)
    bias = fill(0.0, n)
    rmse = fill(0.0, n)

    for i in 1:n

        ikeep = findall(mask[:, i])

        if length(ikeep) > 10

            sim = var_coarse[ikeep, i]

            ref = var_ref[ikeep, i]

            nse[i] = 1 - var(sim .- ref) / var(ref .- mean(ref))

            r[i] = cor(sim, ref)

            bias[i] = mean(sim .- ref)

            rmse[i] = sqrt(mean((sim .- ref).^2))

        else

            nse[i] = NaN

            r[i] = NaN

            bias[i] = NaN

            rmse[i] = NaN

        end

    end

    return nse, r, bias, rmse

end

function results_table(path, cfgs, variables, spaceres)

    # Create empty table

    file_tmp = joinpath(path, "results_1", "swe_$(spaceres).nc")

    idname = Symbol(ncgetatt(file_tmp, "id", "id"))

    ids = convert(Array{Int64}, ncread(file_tmp, "id"))

    df_res = DataFrame()

    df_res[idname] = ids

    for c in cfgs

        println("Running cfg $(c)")

        # Mask indicating presence of snow

        file_fine = joinpath(path, "results_$(c)", "swe_1km.nc")

        file_coarse = joinpath(path, "results_$(c)", "swe_$(spaceres).nc")

        df_links = link_results(file_fine, file_coarse)

        swe_coarse, swe_ref = unify_results(file_fine, file_coarse, df_links, "swe")

        mask = (swe_coarse .> 0) .| (swe_ref .> 0)

        # Compute statistics for all variables

        for v in variables

            file_fine = joinpath(path, "results_$(c)", "$(v)_1km.nc")

            file_coarse = joinpath(path, "results_$(c)", "$(v)_$(spaceres).nc")

            var_coarse, var_ref, ngrids = unify_results(file_fine, file_coarse, df_links, v)

            nse, r, bias, rmse = compute_statistics(var_coarse, var_ref, mask)

            nse_name = Symbol("nse_$(v)_cfg$(c)")

            r_name = Symbol("r_$(v)_cfg$(c)")

            bias_name = Symbol("bias_$(v)_cfg$(c)")

            rmse_name = Symbol("rmse_$(v)_cfg$(c)")

            df_res[nse_name] = nse

            df_res[r_name] = r

            df_res[bias_name] = bias

            df_res[rmse_name] = rmse

            df_res[:ngrids] = ngrids

        end

    end

    return df_res

end

path = "/data04/jmg/fsm_simulations/netcdf/fsmres_forest"

cfgs = 1:32

variables = ["gsurf", "hatmo", "latmo", "melt", "rnet", "rof", "snowdepth", "swe"]

for spaceres in ["5km", "10km", "25km", "50km"]

    println("Running spaceres $(spaceres)")

    df = results_table(path, cfgs, variables, spaceres)

    filesave = joinpath(dirname(pathof(IcanProj)), "..", "data", "table_errors_$(spaceres).txt")

    df |> CSV.write(filesave)

end


function jt06()
#
#  test polynomial suggested by Jenkins and Traub
#
    p = poly([.1;1.001;.998;1.00002;.99999])
    z = [.1;1.001;.998;1.00002;.99999];
    z = [z ones(5)]
    p, PolyZeros(z)
end


# This file contains algorithms needed for the integratiorn of
# transcendental functions from chapter 5 of the book
#
#    Manuel Bronstein, Symbolic Integration I: Transcendental Functions,
#    2nd ed., Springer 2005.
#

"""
    HermiteReduce(f, D) -> (g, h, r)

Hermite reduction - quadratic version.

Given a field `k`, a derivation `D` on `k(t)` and `f` in `k(t)`, return `g`, `h`, `r` in `k(t)`
such that `f=D(g)+h+r`, `h` is simple and `r` is reduced.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.4, p. 139.
"""
function HermiteReduce(f::FracElem{P}, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    iscompatible(f, D) || error("rational function f must be in the domain of derivation D")
    fp, fs, fn = CanonicalRepresentation(f, D)
    a = numerator(fn)
    d = denominator(fn)
    ds = Squarefree(d)
    g = zero(f)
    for i=2:length(ds)
        if degree(ds[i])>0
            v = ds[i]
            u = divexact(d, v^i)
            for j=i-1:-1:1
                b, c = gcdx(u*D(v), v, -1//j*a)
                g += b//v^j
                a = -j*c - u*D(b)
            end
            d = u*v
        end
    end
    q, r = divrem(a, d)
    g, r//d, q + fp + fs
end

"""
    PolynomialReduce(p, D) -> (q, r)

Polynomial reduction.

Given a field `k`, a derivation `D` on `k(t)` and `p` in `k[t]` where `t` is a nonlinear
monomial over `k`, return `q`, `r` in `k[t]`
such that `p=D(q)+r` and `degree(r)<degree(D(t))`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.4, p. 141.
"""
function PolynomialReduce(p::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    iscompatible(p, D) || error("polynomial p must be in the domain of derivation D")
    isnonlinear(D) || error("monomial of derivation D must be nonlinear")
    δ = degree(D)
    λ = leading_coefficient(D)
    t = gen(domain(D))
    if degree(p)<δ
        return (zero(p), p)
    end
    m = degree(p) - δ +1
    q0 = (leading_coefficient(p)//(m*λ))*t^m
    q, r = PolynomialReduce(p - D(q0), D)
    q0 + q, r
end

"""
    ResidueReduce(f, D) -> (ss, Ss, ρ)

Polynomial reduction.

Given a field `k`, a derivation `D` on `k(t)` and `f` in `k(t)`, return either `ρ=1`
and `ss=[s₁,...,sₘ]`, `sᵢ` in `k[z]`, `Ss=[S₁,...Sₘ]`, `Sᵢ` in `k[z,t]` where
`z` is a new indeterminate over `k` such that `f-D(g)` is in `k[t]` where
`g=∑ᵢ∑_{α:sᵢ(α)=0}α*log(Sᵢ(α,t))`, or `ρ=0` and such `ss` and `Ss` such that
`f+h` and `f+h-D(g)` do not have an elementary integral over `k(t)` for any
`h` in `k⟨t⟩`.

(Here, `k⟨t⟩` denotes the elements of `k(t)` which are reduced w.r.t. `D`.)

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.6, p. 153.
"""
function ResidueReduce(f::F, D::Derivation; symbol=:α) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(f, D) || error("rational function f must be in the domain of derivation D")
    issimple(f, D) || error("rational function f must be simple with respect to derivation D")
    d = denominator(f)
    if isone(d)
        return PolyElem[], PolyElem{PolyElem}[], 1
    end
    p, a = divrem(numerator(f), d)
    # For SubResultant with respect to t we have to construct the
    # polynomial ring k[z][t] with z, t in this order (!)
    kz, z = PolynomialRing(base_ring(d), symbol)
    kzt, t = PolynomialRing(kz, var(parent(d)))
    if degree(D(d))<=degree(d)
        r, Rs = SubResultant(d(t), a(t) - z*D(d)(t))
    else
        r, Rs = SubResultant(a(t) - z*D(d)(t), d(t))
    end
    κD = CoefficientLiftingDerivation(kz, BaseDerivation(D))
    ns, ss = SplitSquarefreeFactor(r, κD)
    ds = degree.(ss)
    ss = typeof(d)[ss[i] for i=1:length(ss) if ds[i]>0]
    Ss = typeof(t)[]
    ii = 1
    for i=1:length(ds)
        if ds[i]>0
            if i==degree(d)
                S = d(t)
            else
                m = findfirst(z->z==i, degree.(Rs))
                S =  Rs[m]
                As = Squarefree(leading_coefficient(S))
                for j=1:length(As)
                    S = divexact(S, gcd(As[j],ss[ii])^j)
                end
            end
            push!(Ss, S)
            ii+=1
        end
    end
    ρ = degree(prod(ns))<=0 ? 1 : 0
    ss, Ss, ρ
end

"""
    ConstantPart(ss, Ss, D) -> (αs, gs, ss1, Ss1)

Given the output `ss=[s₁,...,sₘ]`, `sᵢ` in `k[z]`, `Ss=[S₁,...Sₘ]`, `Sᵢ` in `k[z,t]`
of `ResidueReduce` and the derivation D on k(t) which was used by `ResidueReduce`
return  `αs=[α₁,...,αᵣ]` consisting of the roots of those `sᵢ` that have all its roots in `Const(k)` and the corresponding
`gs=[g₁,...,gᵣ]` where `gⱼ` in `k[t]` and `gⱼ(t) = Sᵢ(αⱼ,t)`. The remaining `sᵢ` that have at least one root not in `Const(k)`
and the corresponding `Sᵢ` are returned in `ss1` and `Ss1`.
"""
function ConstantPart(ss::Vector{P}, Ss::Vector{PP}, D::Derivation) where  {P<:PolyElem, PP<:PolyElem{P}}
    #TODO: better name for this function...
    length(ss)==length(Ss) || error("lengths must match")
    if isempty(ss)
        return Term[], 0,  ss, Ss
    end
    Ss1 = eltype(Ss)[]
    ss1 = eltype(ss)[]
    gs = Term[]
    Dg = 0
    D1 = CoefficientLiftingDerivation(parent(ss[1]), BaseDerivation(D))
    for i=1:length(ss)
        αs = constant_roots(ss[i], D1)
        if length(αs)==degree(ss[i]) # all roots found
            for α in αs
                g = map_coefficients(c->c(α), Ss[i])
                if !isconstant(g, D)
                    push!(gs, FunctionTerm(log, α, g))
                    Dg += α*D(g)//g
                end
            end
        else
            RT = LogToReal(SumOfLogTerms(ss[i], Ss[i]))
            αs = constant_roots(ss[i], D1, useQQBar=true)
            if all([isrational(real(α)) && isrational(imag(α)) for α in αs])
                for α in αs
                    u = rationalize(real(α))
                    v = rationalize(imag(α))
                    if iszero(v)
                        g = map_coefficients(c->c(u), Ss[i])
                        if !isconstant(g, D)
                            push!(gs, FunctionTerm(log, u, g))
                            Dg += u*D(g)//g
                        end
                    elseif v > 0
                        var = string(symbols(parent(Ss[i]))[1])
                        F = base_ring(ss[i])
                        if !iszero(u)
                            # Ignore log terms with contstant arguments. In some cases the denominator of the constant argument
                            # might be zero after substitutiong (u,v). So this avoids divison by zero in these cases.
                            if degree(RT.LT.arg)>0 || (!isconstant(numerator(constant_coefficient(RT.LT.arg))(u,v), BaseDerivation(D)) &&
                                                       !isconstant(denominator(constant_coefficient(RT.LT.arg))(u,v), BaseDerivation(D)))
                                # TODO: Think about avoiding division by zero like below for atan term.
                                # But maybe this case cannot occur anyway, at least I know of no example where it does.
                                g = polynomial(F, [numerator(c)(u, v)//denominator(c)(u, v) for c in coefficients(RT.LT.arg)], var)
                                push!(gs, FunctionTerm(log, RT.LT.coeff*u, g))
                                Dg += RT.LT.coeff*u*D(g)//g
                            end
                        end
                        for AT in RT.ATs
                            # Ignore atan terms with contstant arguments. In some cases the denominator of the constant argument
                            # might be zero after substitutiong (u,v). So this avoids divison by zero in these cases.
                            if degree(AT.arg)>0 || (!isconstant(numerator(constant_coefficient(AT.arg))(u,v), BaseDerivation(D)) &&
                                                    !isconstant(denominator(constant_coefficient(AT.arg))(u,v), BaseDerivation(D)))
                                if all([!iszero(denominator(c)(u, v)) for c in coefficients(AT.arg)])
                                    # Ignore atan term if substitution of (u,v) in argument would cause divion by zero.
                                    # This requires more thought, but it seems to work...
                                    g = polynomial(F, [numerator(c)(u, v)//denominator(c)(u, v) for c in coefficients(AT.arg)], var)
                                    push!(gs, FunctionTerm(atan, AT.coeff*v, g))
                                    Dg += AT.coeff*v*D(g)//(1 + g^2)
                                end
                            end
                        end
                    end
                end
            else
                push!(Ss1, Ss[i])
                push!(ss1, ss[i])
            end
        end
    end
    gs, Dg, ss1, Ss1
end

"""
    IntegratePrimitivePolynomial(p, D) -> (q, ρ)

Integration of polynomials in a primitive extension.

Given a field `k`, a derivation `D` on `k(t)` such that `t` is a primitive monomial over `k`
and `p` in `k[t]`, return either `ρ=1` and `q` in `k[t]` such that `p-D(q)` is in `k`, or
`ρ=0` and  `q` in `k[t]` such that `p-D(q)` does not have an elementary integral over
`k(t)`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.8, p. 158.
"""
function IntegratePrimitivePolynomial(p::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    iscompatible(p, D) || error("rational function p must be in the domain of derivation D")
    isprimitive(D) || error("monomial of derivation D must be primitive")
    if degree(p)<=0
        return zero(p), 1
    end
    a = leading_coefficient(p)
    b, c, β = LimitedIntegrate(a, leading_coefficient(D), BaseDerivation(D))
    if β<=0
        return zero(p), 0
    end
    m = degree(p)
    t = gen(parent(p))
    q0 = (c//(m+1))*t^(m+1) + b*t^m
    q, β = IntegratePrimitivePolynomial(p-D(q0), D)
    return q+q0, β
end

"""
    IntegrateHyperexponentialPolynomial(p, D) -> (q, ρ)

Integration of hyperexponential polynomials.

Given a field `k`, a derivation `D` on `k(t)` such that `t` is a hyperexponential monomial over `k`
and `p` in `k⟨t⟩=k[t,t⁻¹]`, return either `ρ=1` and `q` in `k[t,t⁻¹]` such that `p-D(q)` is in `k`, or
`ρ=0` and  `q` in `k[t,t⁻¹]` such that `p-D(q)` does not have an elementary integral over
`k(t)`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.9, p. 162.
"""
function IntegrateHyperexponentialPolynomial(p::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(p, D) || error("rational function p must be in the domain of derivation D")
    ishyperexponential(D) || error("monomial of derivation D must be hyperexponential")
    isreduced(p, D) || error("raional function p must be reduced.")
    q = zero(p)
    ρ = 1
    if iszero(p)
        # Note: with the conventions in Bronstein's book the for loop below would
        # run from i=+infinity to i=-infinity for p=0, i.e., not at all.
        # But it's cleaner to handle the case p=0 separately.
        return q, ρ
    end
    t = gen(base_ring(parent(p)))
    w = coeff(MonomialDerivative(D), 1) # Dt/t
    for i=valuation(p, t):-valuation_infinity(p)
        if i!=0
            a = coeff(numerator(p), i + degree(denominator(p)))
            v, ρ1 = RischDE(i*w, a, BaseDerivation(D))
            if ρ1==0
                ρ = 0
            else
                q += v*(t//1)^i
            end
        end
    end
    q, ρ
end

"""
    IntegrateHypertangentPolynomial(p, D) -> (q, c)

Integration of hypertangent polynomials.

Given a field `k` not conataining `√-1` a derivation `D` on `k(t)` such that `t` is a hypertangent monomial over `k`
and `p` in `k[t]`, return `q` in `k[t]` and `c` in `k` such that `p-D(q)-c*D(t²+1)/(t²+1)` is in `k` and
`p-D(q)` does not have an elementary integral over `k(t)` if `D(c)≠0`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.10, p. 167.
"""
function IntegrateHypertangentPolynomial(p::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    iscompatible(p, D) || error("rational function p must be in the domain of derivation D")
    ishypertangent(D) || error("monomial of derivation D must be hypertangent")
    t = gen(parent(p))
    q, r = PolynomialReduce(p, D)
    α = constant_coefficient(divexact(MonomialDerivative(D), t^2+1))
    c = coeff(r, 1)//(2*α)
    q, c
end

"""
    IntegrateHypertangentReduced(p, D) -> (q, ρ)

Integration of hypertangent reduced elements.

Given a field `k` not conataining `√-1` a derivation `D` on `k(t)` such that `t` is a hypertangent monomial over `k`
and `p` in `k⟨t⟩`, return either `ρ=1` and `q` in `k⟨t⟩` such that `p-D(q)` is in `k[t]`, or
`ρ=0` and  `q` in `k⟨t⟩` such that  `p-D(q)` does not have an elementary integral over `k(t).

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.10, p. 169.
"""
function IntegrateHypertangentReduced(p::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(p, D) || error("rational function p must be in the domain of derivation D")
    ishypertangent(D) || error("monomial of derivation D must be hypertangent")
    isreduced(p, D) || error("rational function p must be reduced.")
    Z = zero(parent(p))
    if iszero(p)
        return Z, 1
    end
    t = gen(base_ring(parent(p)))
    Q = t^2+1
    m = -(iszero(p) ? typemax(Int) : valuation(p, Q))
    if m<=0
        return Z, 1
    end
    h = numerator(Q^m*p)
    q, r = divrem(h, Q)
    a = coeff(r, 1)
    b = coeff(r, 0)
    η = constant_coefficient(divexact(MonomialDerivative(D), Q))
    c, d, ρ = CoupledDESystem(zero(parent(a)), 2*m*η, a, b, BaseDerivation(D))
    if ρ<=0
        return Z, 0
    end
    q0 = (c*t + d)//Q^m
    q, ρ = IntegrateHypertangentReduced(p - D(q0), D)
    q + q0, ρ
end

"""
    Integrate(f, D) -> (αs, lgs, g1, r, ρ)

Given a field `k` (with `√-1` not in `k` in the hypertangent case), a derivation
`D` on `k(t)`, and `f` in `k(t)`,
return either `ρ=1`, `αs=[α₁,...αᵣ]` with `αᵢ` in `Const(k)`, `lgs=[lg₁,...lgᵣ]`
with `lgᵢ` in `k[t]`,
and `g1` in `k(t)`  such that `g=g1+∑αᵢ*log(lgᵢ)` is elementary over `k(t)` and `r=f-D(g)` is in `k`, or
`ρ=0` and such `αs`, `lgs` and `g1` such that `r=f-D(g)` does not have an elementary integral over
`k(t)`.

This function corresponds to `IntegreatePrimitive`, `IntegrateHyperexponential` and
`IntegrateTangent` joined together, see [Bronstein's book](https://link.springer.com/book/10.1007/b138171),
Section 5.8 p. 160, Section 5.9 p. 163, and Section 5.10 p. 172, respectively.
"""
function Integrate(f:: F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    if isbasic(D)
        return IntegrateRationalFunction(f), zero(base_ring(parent(f))), 1
    end
    g1, h, r = HermiteReduce(f, D)
    ss, Ss, ρ = ResidueReduce(h, D)
    g2, Dg2, ss1, Ss1 = ConstantPart(ss, Ss, D)
    if !isempty(ss1)
        throw(AlgebraicNumbersInvolved())
        # From this point on all computations have to be performed in extension fields
        # over field of algebraic (instead of merely rational) numbers, i.e., replace
        # Nemo.QQ by Nemo.QQBar, see http://nemocas.github.io/Nemo.jl/latest/algebraic/
        # For the time being, we start all over again, now with QQBar enabled
        # and discard all results that have been obtained until now.
        # TODO: transform the so far obtained results to fields over QQBar and proceed.
    end
    if ρ<=0
        g = vcat(IdTerm(g1), g2)
        return g, f - D(g1) - Dg2, 0
    end
    p = h - Dg2 + r
    if isprimitive(D)
        @assert isone(denominator(p))
        q, ρ = IntegratePrimitivePolynomial(numerator(p), D)
        g = vcat(IdTerm(g1 + q), g2)
        f1 = f - D(g1 + q) - Dg2
    elseif ishyperexponential(D)
        q, ρ = IntegrateHyperexponentialPolynomial(p, D)
        g = vcat(IdTerm(g1 + q), g2)
        f1 = f - D(g1 + q) - Dg2
    elseif ishypertangent(D)
        q1, ρ = IntegrateHypertangentReduced(p, D)
        if ρ<=0
            g = vcat(IdTerm(g1 + q1), g2)
            f1 = f - D(g1 + q1) - Dg2
        else
            @assert isone(denominator(p - D(q1)))
            q2, c = IntegrateHypertangentPolynomial(numerator(p - D(q1)), D)
            if iszero(BaseDerivation(D)(c))
                t = gen(base_ring(parent(f)))
                η = constant_coefficient(divexact(MonomialDerivative(D), 1 + t^2))
                g = vcat(IdTerm(g1 + q1 + q2), g2, FunctionTerm(log, c, 1 + t^2))
                f1 = f - D(g1 + q1 + q2) - Dg2 - 2*c*η*t
            else
                g = vcat(IdTerm(g1 + q1 + q2), g2)
                f1 = f - D(g1 + q1 + q2) - Dg2
                ρ = 0
            end
        end
    else
        H = MonomialDerivative(D)
        throw(NotImplementedError("Integrate: monomial derivative = $H\n@ $(@__FILE__):$(@__LINE__)"))
    end
    if ρ<=0
        return g, f1, ρ
    end
    @assert isone(denominator(f1)) && degree(numerator(f1)) <= 0
    f1 = constant_coefficient(numerator(f1))
    h, f2,  ρ = Integrate(f1, BaseDerivation(D))
    @assert ρ<=0 || iszero(f2)
    vcat(h, g), f2, ρ
end

"""
    InFieldDerivative(f, D) -> (u, ρ)

Given a field `K`, a derivation `D` on `K` and `f` in `K`, return either `ρ=0`, in which case
the equation `D(u)=f` does not have a solution `u` in `K`, or `ρ=1` and a solution `u` in `K`
of that equation.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.12, p. 175.
"""
function InFieldDerivative(f::F, D::Derivation) where F<:FieldElement
    # base case f \in constant field, D must be the null derivation
    iscompatible(f, D) || error("field element f must be in the domain of derivation D")
    iszero(D) || error("base case only for null derivations")
    if iszero(f)
        return zero(f), 1
    else
        return zero(f), 0
    end
end

function InFieldDerivative(f::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    # See Bronstein's book, Section 5.12, p. 175
    iscompatible(f, D) || error("rational function f must be in the domain of derivation D")
    Z = zero(f)
    no_solution = (Z, 0)
    g, h, r = HermiteReduce(f, D)
    if !isone(denominator(h))
        # no solution, h returned by HermiteReduce not a polynomial
        return no_solution
    end
    p = h + r # reduced
    @assert isreduced(p, D)
    if isprimitive(D)
        # p reduced => polynomial in this case
        p = numerator(p)
        q, ρ = IntegratePrimitivePolynomial(p, D)
        if ρ<=0
            return no_solution
        end
        a0 = p-D(q)
        @assert degree(a0)<=0 # p-D(q) ∈ k
        a = constant_coefficient(a0)
        v, c, ρ = LimitedIntegrate(a, leading_coefficient(D), BaseDerivation(D))
        if ρ<=0
            return no_solution
        end
        t = gen(parent(p))
        return g + q + v + c*t, 1
    else
        if ishyperexponential(D)
            q, ρ = IntegrateHyperexponentialPolynomial(p, D)
            if ρ<=0
                return no_solution
            end
            a0 = p-D(q)
        elseif ishypertangent(D)
            q1, ρ = IntegrateHypertangentReduced(p, D)
            if ρ<=0
                return no_solution
            end
            p1 = p - D(q1)
            @assert isone(denominator(p1))
            q2, c = IntegrateHypertangentPolynomial(numerator(p1), D)
            if !iszero(c)
                # This is my (maybe too naive) interpretation of "prevent
                # introducing a new logarithm". See parenthetical comment
                # on top of p.176.
                return no_solution
            end
            q = q1 + q2
            a0 = p1 - D(q2)
        else
            H = MonomialDerivative(D)
            throw(NotImplementedError("InFieldDerivative: monomial deivative =$H\n@ $(@__FILE__):$(@__LINE__)"))
        end
        @assert isone(denominator(a0)) && degree(numerator(a0))<=0 # p-D(q) ∈ k
        a = constant_coefficient(numerator(a0))
        v, ρ = InFieldDerivative(a, BaseDerivation(D))
        if ρ<=0
            return no_solution
        end
        return g + q + v, 1
    end
end

"""
    InFieldLogarithmicDerivative(f, D) -> (u, ρ)

Given a field `K`, a derivation `D` on `K` and `f` in `K`, return either `ρ=0`, in which case
the equation `D(u)/u=f` does not have a nonzero  solution `u` in `K`, or `ρ=1` and a nonzero
solution `u` in `K` of that equation.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.12, p. 176.
"""
function InFieldLogarithmicDerivative(f::F, D::Derivation) where F<:FieldElement
    n, v, β = InFieldLogarithmicDerivativeOfRadical(f, D, expect_one=true)
    @assert β<=0 || n==1
    v, β
end

"""
    InFieldLogarithmicDerivativeOfRadical(f, D) -> (n, u, ρ)

Given a field `K`, a derivation `D` on `K` and `f` in `K`, return either `ρ=0`, in which case
the equation `D(u)/u=n*f` does not have a solution `u≠0` in `K` and integer `n≠0`, or `ρ=1` and a solution
`u` and `n` of that equation.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 5.12, p. 177.
"""
function InFieldLogarithmicDerivativeOfRadical(f::F, D::Derivation; expect_one::Bool=false) where F<:FieldElement
    # base case f \in constant field, D must be the null derivation
    iscompatible(f, D) || error("field element f must be in the domain of derivation D")
    iszero(D) || error("base case only for null derivations")
    if iszero(f)
        return 1, one(f), 1
    else
        return 0, zero(f), 0 # no solution (n must be !=0)
    end
end

function InFieldLogarithmicDerivativeOfRadical(f::F, D::Derivation; expect_one::Bool=false) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(f, D) || error("rational function f must be in the domain of derivation D")
    if iszero(f)
        return 1, one(f), 1
    end
    Z = zero(f)
    no_solution = (0, Z, 0)
    if !issimple(f, D)
        return no_solution
    end
    if isone(denominator(f))
        m = 1
        Dg = Z
        v = one(f)
    else
        Rs, Gs, ρ = ResidueReduce(f, D)
        @assert ρ>=1
        if length(Rs)==0
            m = 1
            Dg = Z
            v = one(f)
        else
            if !all([all(isrational.(coefficients(R))) for R in Rs])
                # no solution, Rothstein-Trager resultant has non-rational coefficients
                return no_solution
            end
            Rs1 = [map_coefficients(c->convert(fmpq, rationalize(c)), R) for R in Rs]
            As1 = [roots(R) for R in Rs1]
            if !(all([length(As1[i])==degree(Rs1[i]) for i=1:length(As1)]) &&
                all([all(isrational.(As1[i])) for i=1:length(As1)]) )
                # no solution, not all roots of the Rothstein-Trager resultant are rational
                return no_solution
            end
            As = [[rationalize_over_Int(a) for a in as] for as in As1]
            if expect_one
                if !all([all([isone(denominator(a)) for a in as]) for as in As])
                    # no solution, not all roots of the Rothstein-Trager resultant are integers"
                    return no_solution
                end
                m = 1
            else
                m = gcd([gcd(denominator.(as)) for as in As])
            end
            Dg = Z
            v = one(f)
            for i=1:length(As)
                for a in As[i]
                    lc = leading_coefficient(Gs[i])(a) # make monic (!!!)
                    g = map_coefficients(c->c(a)//lc, Gs[i])
                    Dg += a*D(g)//g
                    # Make g = g+Z an element of the field k(t),
                    # otherwise exponentiation with negative numbers would not work.
                    v *= (g+Z)^numerator(m*a)
                end
            end
        end
    end
    p = f - Dg
    if iszero(p)
        if expect_one && m!=1
            # no solution, m not == 1 as expected"
            return no_solution
        end
        return m, v, 1
    end
    if !(isone(denominator(p)))
        # no solution, p=f-Dg not a polynomial
        return no_solution
    end
    if !(degree(numerator(p))<max(1,degree(D)))
        return no_solution
    end
    H = MonomialDerivative(D)
    if isprimitive(D)
        p0 = constant_coefficient(numerator(p))
        n, u, ρ = InFieldLogarithmicDerivativeOfRadical(p0, BaseDerivation(D))
        if ρ<=0
            return no_solution
        end
        if expect_one && n!=1
            # no solution, n not == 1 as expected
            return no_solution
        end
        N = lcm(n, m)
        # Make u = u + Z an element of the field k(t),
        # otherwise exponentiation with negative numbers would not work.
        U = v^div(N, m)*(u+Z)^div(N, n)
        return N, U, 1
    elseif ishyperexponential(D)
        p0 = constant_coefficient(numerator(p))
        w = coeff(H, 1)
        n, e, u, ρ = ParametricLogarithmicDerivative(p0, w, BaseDerivation(D))
        if ρ<=0
            return no_solution
        end
        if expect_one && n!=1
            # no solution, m not == 1 as expected"
            return no_solution
        end
        N = lcm(n, m)
        t = gen(parent(numerator(f)))
        # Make u = u + Z, t = t + Z elements of the field k(t),
        # otherwise exponentiation with negative numbers would not work.
        U = v^div(N, m)*(u+Z)^div(N, n)*(t+Z)^div(e*N, n)
        return N, U, 1
    elseif ishypertangent(D)
        p0 = numerator(p)
        @assert degree(p0)<=1
        a = coeff(p0, 0)
        b = coeff(p0, 1)
        t = gen(parent(H))
        η = constant_coefficient(divexact(H, t^2+1))
        if !contains_I(parent(f))
            c = (1//2)*b//η
            if !isrational(c)
                return no_solution
            end
            n, u, ρ = InFieldLogarithmicDerivativeOfRadical(a, BaseDerivation(D))
            if ρ<=0
                return no_solution
            end
            e0 = n*rationalize_over_Int(c)
            if !isone(denominator(e0))
                return no_solution
            end
            e = numerator(e0)
            N = lcm(n, m)
            U = v^div(N, m)*(u+Z)^div(N, n)*(t^2+1+Z)^div(e*N, n)
            return N, U, 1
        elseif iszero(real(a))
            # Note: This case is not treated in 5.12 of Bronsteins's book, altough it seems
            # to be the one relevant for checking the condition of Theorem 5.10.1.
            # I did not prove that in the case of Theorem 5.12 real(a)=0 always holds true.
            ai = imag(a)
            c1 = (1//2)*(b+ai)//η
            c2 = (1//2)*(b-ai)//η
            if !isrational(c1) || !isrational(c2)
                return no_solution
            end
            # implicitely set u = 1 => D(u)//u = 0
            c1 = rationalize_over_Int(c1)
            c2 = rationalize_over_Int(c2)
            n = lcm(denominator(c1), denominator(c2))
            e1 = numerator(c1*n)
            e2 = numerator(c2*n)
            N = lcm(n, m)
            I = get_I(parent(Z))
            U = v^div(N, m)*(t-I)^div(e1*N, n)*(t+I)^div(e2*N, n)
            return N, U, 1
        else
            # TODO...
            throw(NotImplementedError("InFieldLogarithmicDerivativeOfRadical: hypertangent case with I in K and real(a)!=0\n@ $(@__FILE__):$(@__LINE__)"))
        end
    else
        throw(NotImplementedError("InFieldLogarithmicDerivativeOfRadical: monomial derivative $H\n@ $(@__FILE__):$(@__LINE__)"))
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    BisectPointPartition(normal, point)

A method for partitioning spatial data into two half spaces
defined by a `normal` direction and a reference `point`.
"""
struct BisectPointPartition{Dim,T} <: PartitionMethod
  normal::Vec{Dim,T}
  point::Point{Dim,T}

  function BisectPointPartition{Dim,T}(normal, point) where {Dim,T}
    new(normalize(normal), point)
  end
end

BisectPointPartition(normal::Vec{Dim,T}, point::Point{Dim,T}) where {Dim,T} =
  BisectPointPartition{Dim,T}(normal, point)

BisectPointPartition(normal::NTuple{Dim,T}, point::NTuple{Dim,T}) where {Dim,T} =
 BisectPointPartition(Vec(normal), Point(point))

function partition(object, method::BisectPointPartition)
  Dim = embeddim(object)
  T = coordtype(object)

  n = method.normal
  p = method.point

  x = MVector{Dim,T}(undef)

  left  = Vector{Int}()
  right = Vector{Int}()
  for location in 1:nelements(object)
    coordinates!(x, object, location)
    if (Point(x) - p) ⋅ n < zero(T)
      push!(left, location)
    else
      push!(right, location)
    end
  end

  Partition(object, [left,right])
end


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
# ==============================================================================
#
#   Functions related to the conversion from quaternion to DCM.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

export quat_to_dcm

"""
    quat_to_dcm(q::Quaternion)

Convert the quaternion `q` to a Direction Cosine Matrix (DCM).

# Examples

```jldoctest
julia> q = Quaternion(cosd(45/2), sind(45/2), 0, 0);

julia> quat_to_dcm(q)
3×3 StaticArrays.SMatrix{3, 3, Float64, 9} with indices SOneTo(3)×SOneTo(3):
 1.0   0.0       0.0
 0.0   0.707107  0.707107
 0.0  -0.707107  0.707107
```
"""
function quat_to_dcm(q::Quaternion)
    # Auxiliary variables.
    q0 = q.q0
    q1 = q.q1
    q2 = q.q2
    q3 = q.q3

    return DCM(
        q0^2+q1^2-q2^2-q3^2,   2(q1*q2+q0*q3)   ,   2(q1*q3-q0*q2),
          2(q1*q2-q0*q3)   , q0^2-q1^2+q2^2-q3^2,   2(q2*q3+q0*q1),
          2(q1*q3+q0*q2)   ,   2(q2*q3-q0*q1)   , q0^2-q1^2-q2^2+q3^2
    )'
end


# The purpose of this script is to compute the empirical mean squared error of
# various algorithms for line spetral estimation for a varying ambient dimension
# of the data while keeping the number of measurements fixed.
# The performance of classical algorithms for data from N measurements is then
# compared to the performance of matrix completion based methods with N sparse
# measurements randomly sampled from the N2 \geq N ambient dimensions.
# This allows to show how much sparse subsampling can increase the performance
# when the number of samples is kept fixed.
#
# Parameters:
# N: number of measurements
# N2range: vector of number of ambient dimensions
# K: number of sinusoids
# L: number of snapshots
# Niters: number of Monte Carlo iterations
# SNR_db: signal-to-noise-ratio [dB] of the AWGN model
# f_true: location of the frequencies of the sinusoids

include("MUSIC.jl")
include("rootMUSIC.jl")
include("ESPRIT.jl")
include("TLSesprit.jl")
include("MPM.jl")
include("ANM.jl")
include("EMAC.jl")
include("periodogram.jl")
include("findMinWrapDist.jl")
include("findMinWrapDistUnique.jl")
include("WrapDist.jl")

using Plots, Random

N2range = ceil.(Int,range(8,32,length=4))
N = 8
L = 1
SNR_db = 60
Niters = 1
f_true = [0.3 0.31]
K = length(f_true)
MSE = zeros(length(N2range),3)

A(f,N) = exp.(2*pi*1im*kron(f,(0:N-1)))
SNR = 10 .^(SNR_db./10)

for ii = 1:length(N2range)
    N2 = N2range[ii]
    for jj = 1:Niters

        noise = (randn(N,L)+1im*randn(N,L))/ (sqrt(2)*SNR)
        omega = sort(randperm(N2)[1:N])
        S = exp.(2*pi*rand(K,L)*1im)

        Y1 = A(f_true,N)*S + noise
        Y2 = A(f_true,N2)[omega,:]*S + noise

        t0 = time_ns()
        f_MUSIC = MPM(Y1,K)
        t1 = (time_ns()-t0)/1e9
        MinWrapDist = findMinWrapDistUnique(f_true,f_MUSIC)
        MSE[ii,1] += norm(MinWrapDist)^2

        f_EMAC = EMAC(Y2,N2,K,omega,1/SNR)
        MinWrapDist = findMinWrapDistUnique(f_true,f_EMAC)
        MSE[ii,2] += norm(MinWrapDist)^2

        f_ANM = ANM(Y2,N2,K,omega,1/SNR)
        MinWrapDist = findMinWrapDistUnique(f_true,f_ANM)
        MSE[ii,3] += norm(MinWrapDist)^2

    end
end

MSE /= Niters

plot(N2range, MSE[:,1],yaxis=:log,tickfontsize = 5,
label="Matrix Pencil Method",
#xlims = (0,1),
#xticks = round.(f_range.-f1; digits=2),
xlabel = "M",
ylabel = "MSE"
)
plot!(N2range,MSE[:,2],label="EMAC")
plot!(N2range,MSE[:,3],label="ANM")
#savefig("/home/thomas/atom-amd64/usr/share/atom/example/mse_N2.png")


"""
    slim_scoring_loss01_norm0(λ, X, y, C0)

SLIM (Supersparse Linear Integer Models) loss function that includes
the 0-1 loss and the norm-0 loss parts but not the norm-1 part (which is used
in the original paper). To include also the latter, use `slim_scoring_loss`.
"""
function slim_scoring_loss01_norm0(λ::Vector{N0},
    X::AbstractMatrix{N1}, y::AbstractVector{N2}, C0::Number) where {N0<:Number,N1<:Number,N2<:Number}
    ι = round.(Int, λ)
    loss01 = sum((y .* (X * ι)) .<= 0)/length(y)
    norm0 = sum(ι[2:end] .!= 0) # No loss for lambda[1] which is the intercept
    loss01 + C0*norm0
end

abstract type AbstractSLIMScoringSystem <: InterpretableModel end

"""
    BareSLIMScoringSystem

The "bare" SLIM Scoring System for classification represents only the
scoring system itself but does NOT now about the names of features
and the class to be predicted and thus cannot print itself nicely.
Use a SLIMScoringSystem to get a nicer interface and API.
"""
struct BareSLIMScoringSystem <: AbstractSLIMScoringSystem
    λ::Vector{Int}
    L::Int
    nonzeroindices::Vector{Int}
    trainaccuracy::Float64
    BareSLIMScoringSystem(λ::Vector{N}, L::Int, trainaccuracy::Float64 = -Inf) where {N<:Number} = begin
        @assert 0.0 <= trainaccuracy <= 1.0
        new(floor.(Int, λ), L, nonzeroindices(λ), trainaccuracy)
    end
end
function BareSLIMScoringSystem(λ::Vector{N}, L::Int,
    X::AbstractMatrix{N1}, y::AbstractVector{N2}) where {N<:Number,N1<:Number,N2<:Number}
    BareSLIMScoringSystem(λ, L, sign_accuracy(λ, X, y))
end
numrules(s::BareSLIMScoringSystem) = nnzeroasint(s.λ) - 1
trainaccuracy(s::BareSLIMScoringSystem) =
    (0.0 <= s.trainaccuracy <= 1.0) ? s.trainaccuracy : nothing

"""
    SLIMScoringSystem

A SLIM scoring system for classification which can present itself
and its predictions nicely.
"""
struct SLIMScoringSystem{S} <: AbstractSLIMScoringSystem
    s::BareSLIMScoringSystem
    colnames::Vector{S}
    classfeaturename::S
    classcolidx::Int
end
function SLIMScoringSystem(s::BareSLIMScoringSystem, df::DataFrame,
    classcolumnname::Union{Symbol,Int,S}) where {S<:AbstractString}
    cname, cidx = column_name_and_index(df, classcolumnname)
    ns = names(df)
    deleteat!(ns, cidx)
    SLIMScoringSystem(s, ns, cname, cidx)
end

function Base.show(io::IO, s::SLIMScoringSystem)
end

function column_name_and_index(df::DataFrame, idx::Int)
    @assert 1 <= idx <= size(df, 2)
    return names(df)[idx], idx
end

column_name_and_index(df::DataFrame, colname::Symbol) =
    column_name_and_index(df, string(colname))

function column_name_and_index(df::DataFrame, colname::AbstractString)
    ns = names(df)
    @assert in(colname, ns)
    colname, findfirst(n -> n == colname, ns)
end

function bboptimize_slim_scoring_system(X::AbstractMatrix{N1}, y::AbstractVector{N2},
    C0::Real;
    L::Int = 5,
    MaxTime = 60.0,
    PopSize = 50) where {N1<:Number,N2<:Number}

    C0 = float(C0)
    fitfn(λ::Vector{Float64}) = slim_scoring_loss01_norm0(λ, X, y, C0)

    P = size(X, 2)

    r = bboptimize(fitfn;
            SearchRange = bbo_floor_int_searchrange(L),
            NumDimensions = P,
            MaxTime = MaxTime,
            PopulationSize = PopSize)

    return BareSLIMScoringSystem(best_candidate(r), L, X, y)
end

function bboptimize_slim_scoring_system(df::DataFrame, classcolumn::Int, C0::Real; kws...) where {N1<:Number,N2<:Number}
    y = as_sign_labels(df[:, classcolumn])
    N, P = size(df)
    idxs = collect(1:P)
    deleteat!(idxs, classcolumn)
    X = hcat(ones(Float64, N), Matrix{Float64}(df[:, idxs]))
    bare_ss = bboptimize_slim_scoring_system(X, y, C0; kws...)
    return SLIMScoringSystem(bare_ss, df, classcolumn)
end