# -*- coding: utf-8 -*-
# # Multivariate Nonlinear Regression
#
# This example demonstrates how to use a product group to manipulate a multivariate regression model in an algebraically precise way. It builds on the example `monomial_regression.jl`.
#

# Let `m` be a model like
# ```julia
# f(x,y,z) = a*x + b*y
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i],z[i]) - target[i])^2
# ```
#
# So it is a bivariate least squares regression model.
# Define the set of transformations $T$ as the free monoid generated by
#
# ```julia
# x^i->x^i+1
# x^i->x^i-1
# y^i->y^i+1
# y^i->y^i-1
# ```
#
# The free monoid over this set of transformations is actually the product of cyclic groups $Z^2 = Z x Z$ with the
# group operation of $+:Z^2->Z^2$. so you get the set of models:
#
# ```julia
# f(x,y) = a*x^i + b*y^j
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i]) - target[i])^2
# ```
#
# for all combinations of integers $i,j$ which is a pretty cool structure. If you do the transformations with `i mod N` then you get $Z_N^2$ which is a finite abelian group. It doesn't get more computationally tractable than component-wise arithmetic mod N.

# We could then talk about the orbits of a given model under this group. For this group, there is only one orbit, but if you take subgroups like the group generated by $(1,1)$ aka $\langle 1,1\rangle$ you get multiple orbits. For example if you start with
#
# ```julia
# f(x,y) = a*x^1 + b*y^2
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i]) - target[i])^2
# ```
# the orbit is all of the models of the form:
# ```julia
# f(x,y) = a*x^1+i + b*y^2+i
# x,y = rand(Normal(0,1), (n,2))
# target[i] = f(x[i],y[i]) + rand(Normal(0,1), 1)
# min_{a,b} sum_i (f(x[i],y[i]) - target[i])^2
# ```
# where $i$ is in $Z_N$. You could reason about finding the best fitting model from a given orbit. and then comparing
# the orbits to see which one was contained the best fitting model.
#
# Here is an example implementation of this theory

# ## Example Model

#Our working example of a multivariate nonlinear regression model
expr = quote
    module Regression
    using Random
    using LsqFit
    using LinearAlgebra

    function f(X, β)
        a = β[1]
        b = β[2]
        x = X[:, 1]
        y = X[:, 2]
        return a.*x.^0 .+ b.*y.^0
    end

    function sample(g::Function, n)
        x = randn(Float64, n)
        target = g(x) .+ randn(Float64, n[1])./1600
        return x, target
    end

    function describe(fit)
        if !fit.converged
            error("Did not converge")
        end
        return (β = fit.param, r=norm(fit.resid,2))
    end
    #setup

    # Random.seed!(42)
    β = (1/2, 1/2)
    n = (1000,2)
    g(x) = β[1].*x[:,1].^2 + β[2].*x[:,2].^3
    X, target = sample(g, n)
    # x, y = X[:,1], X[:,2]
    # @show size(X), size(target)
    # loss(a) = sum((f.(a, x, y) .- target).^2)
    # @show loss.([-1,-1/2,-1/4, 0, 1/4,1/3,1/2,2/3, 1])

    a₀ = [1.5, 0.1]
    try
        ŷ₀ = f(X, a₀)
    catch
        error("Could not execute f on the initial data")
    end

    #solving
    fit = curve_fit(f, X, target, a₀; autodiff=:forwarddiff)
    result = describe(fit)
    end
end;

# ## Implementation Details
# The following code is the implementation details for representing the models as an `AbstractProblem` and representing the transformations as `Product{Tuple{Pow{Int}, Pow{Int}}}` and applying the transformations onto the models.

using SemanticModels
using SemanticModels.ModelTools
using SemanticModels.ModelTools.Transformations

import SemanticModels.ModelTools: model, AbstractModel
import SemanticModels.Parsers: findfunc, findassign
import Base: show

# +
""""    MultivariateLsq

A program that solves min_β || f(X,β) - y ||_2

Example:

`f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q`
"""
struct MultivariateLsq <: AbstractModel
    expr
    f
    coefficient
    p₀
end
# -

function show(io::IO, m::MultivariateLsq)
    write(io, "MultivariateLsq(\n  f=$(repr(m.f)),\n  coefficient=$(repr(m.coefficient)),\n  p₀=$(repr(m.p₀))\n)")
end

function model(::Type{MultivariateLsq}, expr::Expr)
    if expr.head == :block
        return model(MultivariateLsq, expr.args[2])
    end
    objective = :l2norm
    f = callsites(expr, :curve_fit)[end].args[3]
    coeff = callsites(expr, f)[1].args[end]
    p₀ = callsites(expr, :curve_fit)[end].args[end]
    return MultivariateLsq(expr, f, coeff, p₀)
end

"""    poly(m::MultivariateLsq)::Expr

find the part of the model that implements the polynomial model for regression.
"""
function poly(m::MultivariateLsq)
    func = findfunc(m.expr, m.f)[1]
    poly = func.args[2].args[end].args[1]
    return poly
end

"""    (t::Pow)(m::MultivariateLsq, i::Int)

apply the power transformation to the i-th coordinate of m.

Example:

If `m` is a program implementing `f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q`

a) and `t = Pow(2)` then `t(m, 1)` is the model implementing
`f(X, β) = β[1]*X[:,1]^p+2 + β[2]*X[:,2]^q`.

or

b) and `t = Pow(2)` then `t(m, 2)` is the model implementing
`f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q+2`.

See also [`(t::Pow)(m::MultivariateLsq, i::Int)`](@ref)
"""
function (t::Pow)(m::MultivariateLsq, i::Int)
    p = poly(m)
    comp = p.args[i+1]
    pow = callsites(comp, :(.^))
    pow[end].args[3] += t.inc
    return m
end

"""    (t::Product)(m::MultivariateLsq)

apply the product transformation to the model m.
For a MultivariateLsq model a product transformation works
by applying the transformation component-wise where each component represents a
column of the data.

Example:

If `m` is a program implementing `f(X, β) = β[1]*X[:,1]^p + β[2]*X[:,2]^q`
and `t = Pow(2,-1)` then `t(m)` is the model implementing
`f(X, β) = β[1]*X[:,1]^p+2 + β[2]*X[:,2]^q-1`.
"""
function (t::Product)(m::MultivariateLsq)
    m = m
    for (i, tᵢ) in enumerate(t.dims)
        m = tᵢ(m, i)
    end
    return m
end

# ## Model Augmentation with Group Actions
# the Group `Product{Tuple{Pow{Int},Pow{Int}}}` which is isomorphic to $(Z, +, 0)\times (Z,+,0)$ can act on our class of models
#
# Let's build an instance of the model object from the code snippet `expr`

m = model(MultivariateLsq, expr)
@show m
poly(m)

# We can get the poly from the model and show how the Pow group elements can act on it componentwise.

@show poly((one(Pow)^4)(deepcopy(m), 1))
@show poly((one(Pow)^4)(deepcopy(m), 2))

# Some *basis vectors* will come in handy for building elements of the product group.
# $T_1,T_2$ are generators for $Z\times Z$

@show T₁ = Product((one(Pow), zero(Pow)))
@show T₂ = Product((zero(Pow), one(Pow)))

# Some Checks that we get what is expected.

m′ = @show (T₁^2 ∘ T₂^3)(deepcopy(m))
p = poly(m)
p′ = poly(m′)
@assert p′.args[2].args[3].args[3] == p.args[2].args[3].args[3] + 2
@assert p′.args[3].args[3].args[3] == p.args[3].args[3].args[3]+3
p′

# ## Exploring the Orbits
#
# We use the Levenberg Marquardt algorithm for general least squares problems in data analysis. The goal is to fix a model class and find the best coefficients for that class. Our algebraic representation allows us to have a similar treatment of model class (in this case the exponents $i,j$ in our formula). We can sample from the orbits of the group when applied to the model and solve for the best coefficients in order to find the best model class.
#
# In this case we have reduced the code search space to a model search space to an algebraic formulation.

m = model(MultivariateLsq, expr)
results = []
for i in 1:3
    for j in 1:3
        m₀ = deepcopy(m)
        T = (T₁^i) ∘ (T₂^j)
        mᵢⱼ = T(m₀)
        p = poly(mᵢⱼ)
        M = eval(mᵢⱼ.expr)
        push!(results, (i=i, j=j, M.result..., p=p))
    end
end

# It turns out that this method recovers the *true* model order $2,3$
#
#
#
#

sort(results, by=x->x.r)

# ## Conclusions
#
# We have seen how abstract algebra can be applied to the category of models to build a systematic treatment of model augmentation. This proves that the model transformations can be arranged into a simple algebraic structure that can act on a model to build new models. The structure of the transformations are easier to analyze than the changes to the models themselves.


import ..Embeddings:
    AbstractEmbedding

"""
    minima_and_stepsizes(points, ϵ) -> (Vector{Float}, Vector{Float})

Find the minima along each axis of the embedding, and computes appropriate
`stepsizes` given `ϵ`, which provide instructions on how to grid the space.
Assumes each point is a column vector.

Specifically, the binning procedure is controlled by the type of `ϵ`:

1. `ϵ::Int` divides each axis into `ϵ` intervals of the same size.
2. `ϵ::Float` divides each axis into intervals of size `ϵ`.
3. `ϵ::Vector{Int}` divides the i-th axis into `ϵᵢ` intervals of the same size.
4. `ϵ::Vector{Float64}` divides the i-th axis into intervals of size `ϵᵢ`.
"""
function minima_and_stepsizes(points, ϵ)
    if size(points, 1) > size(points, 2)
        error("The dimension of the dataset exceeds the number of points.")
    end
    D = size(points, 1)
    n_pts = size(points, 2)

    axisminima = Vector{Float64}(undef, D)
    top = Vector{Float64}(undef, D)

    for i = 1:D
        axisminima[i] = minimum(points[i, :])
        top[i] = maximum(points[i, :])
    end
    axisminima = axisminima - (top - axisminima) / 100
    top = top + (top - axisminima) / 100

    stepsizes = Vector{Float64}(undef, D)
    if ϵ isa Float64
        stepsizes = [ϵ for i in 1:D]
    elseif ϵ isa Vector{Float64}
        stepsizes = ϵ
    elseif ϵ isa Int
        stepsizes = (top - axisminima) / ϵ
    elseif ϵ isa Vector{Int}
        stepsizes = (top - axisminima) ./ ϵ
    elseif ϵ isa Tuple{Vector{Tuple{Float64, Float64}}, Int}
        # We have predefined axis minima and axis maxima.
        n_bins = ϵ[2]
        dim = length(ϵ[1])
        stepsizes = zeros(Float64, dim)
        axisminima = zeros(Float64, dim)

        for i = 1:dim
            stepsizes[i] = (maximum(ϵ[1][i]) - minimum(ϵ[1][i]))/n_bins
            axisminima[i] = minimum(ϵ[1][i])
        end
    end

    axisminima, stepsizes
end

"""
    assign_integer_bin_label_to_eachpoint!(
        A::Array{Int, 2},
        embedding_points::Array{Float64, 2},
        axisminima::Vector{Float64},
        stepsizes::Vector{Float64},
        npts::Int) -> Array{Int, 2}

Assign integer bin labels to the points of an embedding, storing the
labels in the preallocated array `A`.
"""
function assign_integer_bin_label_to_eachpoint!(
            A::AbstractArray{Int, 2},
            pts::AbstractArray{T, 2},
            axisminima::Vector{Float64},
            δs::Vector{Float64},
            npts::Int) where T
    @inbounds for i = 1:npts
        A[:, i] .= floor.(Int, abs.(axisminima .- pts[:, i]) ./ δs)
    end
end

"""
    assign_bin_labels(points, ϵ)

Consider a rectangular grid specified by ϵ. Assign bin labels to the provided
`points` by checking which bins each point fall into. Each points is given a
label which is an integer encoding of the origin of the bin it falls into.

The following `ϵ` will work:

* `ϵ::Int` divide each axis into `ϵ` intervals of the same size.
* `ϵ::Float` divide each axis into intervals of size `ϵ`.
* `ϵ::Vector{Int}` divide the i-th axis into `ϵᵢ` intervals of the same size.
* `ϵ::Vector{Float64}` divide the i-th axis into intervals of size `ϵᵢ`.

The points are assumed to be provided as an array where each point is a column.
"""
function assign_bin_labels(points, ϵ)
    # Make sure that the array contains points as columns.
    if size(points, 1) > size(points, 2)
        error("The dimension of the dataset exceeds the number of points.")
    end

    # How many points are there in the embedding, and what is its dimension?
    D = size(points, 1)
    npts = size(points, 2)

    if npts <= D
        error("Number of pts is less than dim=$D. Is every column a point?")
    end
    # Find the minima of the embedding and generate step
    # sizes along each axis
    bottom, stepsizes = minima_and_stepsizes(points, ϵ)
    if typeof(stepsizes) <: Union{Float64, Int}
        stepsizes = [stepsizes]
    end

    # Each points of the embedding gets assigned to one bin.
    # The coordinates of each point of the original embedding are
    # assigned an integer number indicating which bin along the respective
    # dimension it falls into.
    visited_bins_inds = zeros(Int, D, npts)
    assign_integer_bin_label_to_eachpoint!(
        visited_bins_inds,
        points,
        bottom,
        stepsizes,
        npts)

    return visited_bins_inds
end

"""
    assign_bin_labels(E::AbstractEmbedding, ϵ) -> Array{Int, 2}

Find bins visited by the reconstructed orbit and assign
to them integer tuples uniquely identifying the bins.

Each point of the embedding gets assigned one integer tuple label,
stored as column vectors in the returned array. If some bin
is visited more than once, there will be repeated columns.

# Details
Given a bin size ϵ, whics is `Union{Int, Float}` for cubic bins,
and `Union{Vector{Int}, Vector{Float64}}` for rectangular bins,
assign a tuple of integers uniquely identifying that bin.

Take, for example, the pair `(pᵢ, (n₁, n₂, n₃)`. Having a
3-element tuple associated with pᵢ means that the reconstruction
is 3D, so we have three coordinate axes - x₁, x₂ and x₃ - to
consider.

This generic example should be read as:
for point `pᵢ`, one must jump n₁ steps along x₁, n₂ steps
along x₂ and n₃ steps along x₃ (in units of ϵ₁, ϵ₂ and ϵ₃)
from the minima of each axis to reach the coordinates of `pᵢ`.

Each visited bin is thus uniquely identified by a tuple of
integers. `visited_bins_indices` assigns one such tuple
to each of the points in the embedding. These are gathered
in an array where column vectors represent the tuples
associated to each point.
"""
function assign_bin_labels(E::AbstractEmbedding, ϵ)
    assign_bin_labels(E.points, ϵ)
end

"""
    assign_coordinate_labels_to_eachpoint!(
            A::Array{Float64, 2},
            pts::Array{Float64, 2},
            axisminima::Vector{Float64},
            δs::Vector{Float64},
            npts::Int)

Assign bin labels to the points of an embedding as coordinate
relative to the minima along each coordinate axis. Store the
coordinate labels in the preallocated array `bin_origins`.
"""
function assign_coordinate_labels_to_eachpoint!(
        bin_origins::AbstractArray{Float64, 2},
        visited_bin_inds::AbstractArray{Int, 2},
        pts::AbstractArray{Float64, 2},
        axisminima::Vector{Float64},
        δs::Vector{Float64},
        n_orbit_pts::Int)
    for i = 1:n_orbit_pts
        bin_origins[:, i] = axisminima .+ (δs .* visited_bin_inds[:, i])
    end
end

"""
    assign_coordinate_labels(points, ϵ) -> Array{Float64, 2}

Consider a rectangular grid specified by ϵ. Assume the bin labels to the provided
`points` by checking which bins each point fall into. Each points is given a
label which is the coordinates of the origin of the bin it falls into.

Together with `ϵ`, the coordinates of the origins of the bins provide
complete information about a coarse graining that covers all visited
states of the system (when `points` represent the orbit).

The following `ϵ` will work:

* `ϵ::Int` divide each axis into `ϵ` intervals of the same size.
* `ϵ::Float` divide each axis into intervals of size `ϵ`.
* `ϵ::Vector{Int}` divide the i-th axis into `ϵᵢ` intervals of the same size.
* `ϵ::Vector{Float64}` divide the i-th axis into intervals of size `ϵᵢ`.

The points are assumed to be provided as an array where each point is a column.
"""
function assign_coordinate_labels(points, ϵ)
    # Make sure that the array contains points as columns.
    if size(points, 1) > size(points, 2)
        error("The dimension of the dataset exceeds the number of points.")
    end

    # How many points are there in the embedding, and what is its dimension?
    D = size(points, 1)
    npts = size(points, 2)

    if npts <= D
        error("Number of pts is less than dim=$D. Is every column a point?")
    end

    # Find the minima of the embedding and generate step
    # sizes along each axis.
    axisminima, stepsizes = minima_and_stepsizes(points, ϵ)
    visited_bin_inds = assign_bin_labels(points, ϵ)
    bin_origins = zeros(Float64, D, npts)

    n_orbit_points = size(points, 2)
    assign_coordinate_labels_to_eachpoint!(
        bin_origins,
        assign_bin_labels(points, ϵ),
        points,
        axisminima,
        stepsizes,
        n_orbit_points
    )

    return bin_origins
end

"""
    assign_coordinate_labels(E::AbstractEmbedding, ϵ) -> Array{Float64, 2}

Consider a rectangular grid specified by ϵ. Assume the bin labels to the
provided embedding by checking which bins each point fall into. Each points is
given a label which is the coordinates of the origin of the bin it falls into.

Together with `ϵ`, the coordinates of the origins of the bins provide
complete information about a coarse graining that covers all visited
states in the reconstruction/embedding.

The following `ϵ` will work:

* `ϵ::Int` divide each axis into `ϵ` intervals of the same size.
* `ϵ::Float` divide each axis into intervals of size `ϵ`.
* `ϵ::Vector{Int}` divide the i-th axis into `ϵᵢ` intervals of the same size.
* `ϵ::Vector{Float64}` divide the i-th axis into intervals of size `ϵᵢ`.
"""
function assign_coordinate_labels(E::AbstractEmbedding, ϵ)
    assign_coordinate_labels(E.points, ϵ)
end

"""
    assign_coordinate_labels(visited_bin_inds, points, ϵ)  -> Array{Float64, 2}

Consider a rectangular grid specified by ϵ. Assume that, given `ϵ`,
integer bin  labels have been assigned to each point and are stored in
`visited_bin_inds`. This array contains one column of integer bin
labels each of the points in `points`.

Here, we return an array of the same size as `visited_bin_inds`. However,
instead of the labels being integers, they are converted into the
coordinates of the origin of the corresponding bins.

Together with `ϵ`, the coordinates of the origins of the bins provide
complete information about a coarse graining that covers all visited
states of the system (when `points` represent the orbit).

The following `ϵ` will work:

* `ϵ::Int` divide each axis into `ϵ` intervals of the same size.
* `ϵ::Float` divide each axis into intervals of size `ϵ`.
* `ϵ::Vector{Int}` divide the i-th axis into `ϵᵢ` intervals of the same size.
* `ϵ::Vector{Float64}` divide the i-th axis into intervals of size `ϵᵢ`.

The points are assumed to be provided as an array where each point is a column.
"""
function assign_coordinate_labels(visited_bin_inds, points, ϵ)
    # Make sure that the array contains points as columns.
    if size(points, 1) > size(points, 2)
        error("The dimension of the dataset exceeds the number of points.")
    end

    # How many points are there in the embedding, and what is its dimension?
    D = size(points, 1)
    npts = size(points, 2)

    if npts <= D
        error("Number of pts is less than dim=$D. Is every column a point?")
    end

    # Find the minima of the embedding and generate step
    # sizes along each axis.
    axisminima, stepsizes = minima_and_stepsizes(points, ϵ)

    bin_origins = zeros(Float64, size(visited_bin_inds))

    n_orbit_points = size(points, 2)
    assign_coordinate_labels_to_eachpoint!(
        bin_origins,
        visited_bin_inds,
        points,
        axisminima,
        stepsizes,
        n_orbit_points
    )

    return bin_origins
end

"""
    assign_coordinate_labels(visited_bin_inds,
        E::AbstractEmbedding, ϵ)  -> Array{Float64, 2}

Consider a rectangular grid specified by ϵ. Assume that, given `ϵ`,
integer bin  labels have been assigned to each point and are stored in
`visited_bin_inds`. This array contains one column of integer bin
labels each of the points of the embedding.

Here, we return an array of the same size as `visited_bin_inds`. However,
instead of the labels being integers, they are converted into the
coordinates of the origin of the corresponding bins.

Together with `ϵ`, the coordinates of the origins of the bins provide
complete information about a coarse graining that covers all visited
states of the reconstructed state space.

The following `ϵ` will work:

* `ϵ::Int` divide each axis into `ϵ` intervals of the same size.
* `ϵ::Float` divide each axis into intervals of size `ϵ`.
* `ϵ::Vector{Int}` divide the i-th axis into `ϵᵢ` intervals of the same size.
* `ϵ::Vector{Float64}` divide the i-th axis into intervals of size `ϵᵢ`.
"""
function assign_coordinate_labels(visited_bin_inds, E::AbstractEmbedding, ϵ)
    assign_coordinate_labels(visited_bin_inds, E.points, ϵ)
end

export
minima_and_stepsizes,
assign_integer_bin_label_to_eachpoint!,
assign_coordinate_labels_to_eachpoint!,
assign_bin_labels,
assign_coordinate_labels


# ------------------------------------------------------------------------------------------
# # 新手入门
#
# 主题：
# 1. 如何打印（print）
# 2. 如何给变量赋值
# 3. 如何注释（comment）
# 4. 基础数学语法
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## 如何打印（print）
#
# 在 Julia 中使用 `println()` 函数来打印输出（print）
# ------------------------------------------------------------------------------------------

println("I'm excited to learn Julia!")

# ------------------------------------------------------------------------------------------
# ## 如何给变量赋值
#
# 只需要变量名、变量值和一个等号！
#
# Julia 会自动判断类型。
# ------------------------------------------------------------------------------------------

my_answer = 42
typeof(my_answer)

my_pi = 3.14159
typeof(my_pi)

😺 = "开心的猫!"
typeof(😺)

# ------------------------------------------------------------------------------------------
# 要打出“开心的猫”，用 tab 键补全选择 emoji 的名字然后再按 tab 键即可
#
# > `\:smi` + `<tab>` ==> 上下键选定 + `<enter>` ==> `<tab>` + `<enter>` ==> 完成
# >
# > 即 `\:smile_cat:` + `<tab>`
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# 在给一个变量赋值后，可以给它重新赋任意类型的值而不会报错。
# ------------------------------------------------------------------------------------------

😺 = 1

typeof(😺)

# ------------------------------------------------------------------------------------------
# 注意：Julia 允许用超广泛的字符，就比如😺
#
# 这就允许我们可以把代码写成这样：
# ------------------------------------------------------------------------------------------

😀 = 0
😞 = -1

😺 + 😞 == 😀

# ------------------------------------------------------------------------------------------
# ## 如何注释（comment）
# ------------------------------------------------------------------------------------------

# 单行注释用井号键（pound/hash key）

#=

多行注释，
用“#= =#”组合

=#

# ------------------------------------------------------------------------------------------
# ## 基础数学语法
# ------------------------------------------------------------------------------------------

sum = 3 + 7  # 加

difference = 10 - 3  # 减

product = 20 * 5  # 乘

quotient = 100 / 10  # 除

power = 10 ^ 2  # 乘方

modulus = 101 % 2  # 取余

# ------------------------------------------------------------------------------------------
# ## 练习
#
# ### 1.1
# 查看函数 `convert` 的帮助文档。
# ------------------------------------------------------------------------------------------

# 在此作答

# ------------------------------------------------------------------------------------------
# ### 1.2
# 给变量 `days` 赋值`365`，将 `days` 转为浮点数（float）并将其赋值给变量 `days_float`
# ------------------------------------------------------------------------------------------

days = 0
days_float = days

@assert days == 365
@assert days_float == 365.0

# ------------------------------------------------------------------------------------------
# ### 1.3
# 看看下面两种用法的区别
#
# ```julia
# convert(Int64, "1")
# ```
# 和
#
# ```julia
# parse(Int64, "1")
# ```
# ------------------------------------------------------------------------------------------

convert(Int64, "1")

parse(Int64, "1")

# ------------------------------------------------------------------------------------------
# 请在完成练习运行 `@assert` 块以检测结果是否正确，若无报错即为正确。
# ------------------------------------------------------------------------------------------


"""

"""
function corey_gas_oil()

end # corey_gas_oil

"""
brooks_corey_gas_oil(labda=4.17, sor=0.2, pce_go=5000)

Original Brooks and Corey (1964)

Input parameters:

    + labda: sorting factor (pore-size distribution)
    + sor: residual oil saturation
    + pce_go: gas-oil pore entry capillary pressure
"""
function brooks_corey_gas_oil(labda=4.17, sor=0.2, pce_go=5000)
  kro(so)=((so-sor)/(1-sor)).^((2+3*labda)/labda)
  krg(so)=((1-so)/(1-sor)).^2.*(1-((so-sor)/(1-sor)).^((2+labda)/labda))
  pc_go(so)=pce_go*((so-sor)/(1-sor)).^(-1/labda)
  return kro, krg, pc_go
end # brooks_corey

"""

"""
function brooks_corey_general(kro_max=1.0, krg_max=1.0, ng=2.0, no=2.0)

end # brooks_corey


module DifferentialExpression
export
    de_mannwhitney,
    DESeq2

using GLM: glm, coeftable, LogLink, NegativeBinomial
using HypothesisTests:MannWhitneyUTest
using MultipleTesting
using Statistics: median
using StatsBase: geomean

# scRNASeq Differential Gene Expression with two-sample
# MannWhitneyUTest

function de_mannwhitney(X,Y)
    assert(size(X,1)==size(Y,1))
    n=size(X,1)
    p=zeros(n,1)
    for k in 1:n
        p[k]=pvalue(MannWhitneyUTest(X[k,:],Y[k,:]))
    end
    q=adjust(vec(p), BenjaminiHochberg())
    return p,q
end

add_constant(X) = hcat(ones(size(X, 1)), X)
loge_to_log2(β) = β |> exp |> log2
median_of_ratios_normalization(X) = normalize_geomean(X) |> scale_size_factors
normalize_geomean(X) = X ./ mapslices(geomean, X, dims=1)
scale_size_factors(X) = X .* mapslices(median, X, dims=2)

struct DESeq2_results
    β
    SE
    z
    p
end

function get_DESeq2_results(models)
    # Needs to be fixed to use Log2 in the future
    β  = [coeftable(model).cols[1][2] for model in models]
    SE = [coeftable(model).cols[2][2] for model in models]
    z  = [coeftable(model).cols[3][2] for model in models]
    p  = [coeftable(model).cols[4][2] for model in models]
    DESeq2_results(β, SE, z, p)
end

function DESeq2(X, y)
    X_norm = median_of_ratios_normalization(X)
    models = mapslices(X -> glm(add_constant(X), y, NegativeBinomial(), LogLink()), X_norm, dims=1)
    get_DESeq2_results(models)
end

"""
https://genomebiology.biomedcentral.com/articles/10.1186/s13059-014-0550-8#Sec2

Work in progress
----------------
Some low hanging fruit
1. Convert analysis to log2
2. Offer extra analysis/plots that DESeq2 provides

Challenging work
----------------
1. Ensuring this module behaves exactly like DESeq2
2. Going through the rest of the paper for details
"""

end


##############################################################################
#
# REFERENCES: Forbes et al. "Statistical Distributions"
#
##############################################################################

immutable BetaPrime <: ContinuousUnivariateDistribution
    alpha::Float64
    beta::Float64
    function BetaPrime(a::Real, b::Real)
        (a > zero(a) && b > zero(b)) || error("Alpha and beta must be positive")
        new(float64(a), float64(b))
    end
end

BetaPrime() = BetaPrime(1.0, 1.0)

cdf(d::BetaPrime, q::Real) = inc_beta(q / 1.0 + q, d.alpha, d.beta)

insupport(::BetaPrime, x::Real) = zero(x) < x
insupport(::Type{BetaPrime}, x::Real) = zero(x) < x

function mean(d::BetaPrime)
    d.beta > 1.0 || error("mean not defined when beta <= 1")
    d.alpha / (d.beta - 1.0)
end

mode(d::BetaPrime) = d.alpha > 1.0 ? (d.alpha - 1.0) / (d.beta + 1.0) : 0.0
modes(d::BetaPrime) = [mode(d)]

function pdf(d::BetaPrime, x::Real)
    α, β = d.alpha, d.beta
    (x^(α - 1.0) * (1.0 + x)^(-(α + β))) / beta(α, β)
end

rand(d::BetaPrime) = 1.0 / rand(Beta(d.alpha, d.beta))

function skewness(d::BetaPrime)
    α, β = d.alpha, d.beta
    β > 3.0 || error("skewness not defined when β <= 3")
    (2.0 * (2.0 * α + β - 1))/(β - 3.0) * sqrt((β - 2.0)/(α * (α + β - 1.0)))
end

function var(d::BetaPrime)
    α, β = d.alpha, d.beta
    β > 2.0 || error("var not defined when β <= 2")
    (α * (α + β - 1.0)) / ((β - 2.0) * (β - 1.0)^2)
end


@node MvNormalMeanCovariance Stochastic [ out, (μ, aliases = [ mean ]), (Σ, aliases = [ cov ]) ]

conjugate_type(::Type{ <: MvNormalMeanCovariance }, ::Type{ Val{ :out } }) = MvNormalMeanCovariance
conjugate_type(::Type{ <: MvNormalMeanCovariance }, ::Type{ Val{ :μ } })   = MvNormalMeanCovariance
conjugate_type(::Type{ <: MvNormalMeanCovariance }, ::Type{ Val{ :Σ } })   = InverseWishart

# default method for mean-field assumption
@average_energy MvNormalMeanCovariance (q_out::Any, q_μ::Any, q_Σ::Any) = begin
    # naive: (ndims(q_out) * log2π + mean(logdet, q_Σ) + tr(cholinv(mean(q_Σ))*(v_out + v_mean + (m_out - m_mean)*(m_out - m_mean)'))) / 2
    dim = ndims(q_out)

    m_mean, v_mean = mean_cov(q_μ)
    m_out, v_out   = mean_cov(q_out)
    inv_m_Σ        = mean(inv, q_Σ)

    result = zero(promote_type(eltype(q_out), eltype(q_μ), eltype(q_Σ)))
    result += mean(logdet, q_Σ)
    result += dim * log2π
    @inbounds for k1 in 1:dim, k2 in 1:dim   # optimize trace operation (indices can be interchanges because of symmetry)
        result += inv_m_Σ[k1,k2] * (v_out[k1,k2] + v_mean[k1,k2] + (m_out[k2]-m_mean[k2]) * (m_out[k1] - m_mean[k1]))
    end
    result /= 2

    return result
end

# default method for structured mean-field assumption
@average_energy MvNormalMeanCovariance (q_out_μ::Any, q_Σ::Any) = begin
    # naive: @views (d*log2π + mean(logdet, q_Σ) + tr(cholinv(mean(q_Σ))*( V[1:d,1:d] - V[1:d,d+1:end] - V[d+1:end,1:d] + V[d+1:end,d+1:end] + (m[1:d] - m[d+1:end])*(m[1:d] - m[d+1:end])' ))) / 2
    dim = div(ndims(q_out_μ), 2)

    m, V = mean_cov(q_out_μ)
    inv_m_Σ = mean(inv, q_Σ)

    result = zero(promote_type(eltype(q_out_μ), eltype(q_Σ)))
    result += mean(logdet, q_Σ)
    result += dim * log2π
    @inbounds for k1 in 1:dim, k2 in 1:dim   # optimize trace operation (indices can be interchanges because of symmetry)
        result += inv_m_Σ[k1,k2] * (V[k1,k2] + V[dim+k1,dim+k2] - V[dim+k1,k2] - V[k1,dim+k2] + (m[k1] - m[dim+k1]) * (m[k2] - m[dim+k2]))
    end
    result /= 2

    return result
end


@testset "CDS_MVp_mat2CDS_tests" begin
  TI=Int32
  TF=Float32;
  n1=30;
  n2=20;
  N=n1*n2;
  x=randn(TF,N);
  y=zeros(TF,N);

  #test structured matrix
  #comp_grid = compgrid( (25, 25),( n1,n2 ) );
  comp_grid = compgrid( (TF(25), TF(25)),( n1,n2 ) );
  (TV_OP, AtA_diag, dense, TD_n, banded)=get_TD_operator(comp_grid,"TV",TF);
  A=TV_OP'*TV_OP;
  A=convert(SparseMatrixCSC{TF,TI},A);

  Ax_native = A*x;

  (R,offset)=mat2CDS(A)
  fill!(y,0)
  Ax_CDS = CDS_MVp(size(R,1),size(R,2),R,offset,x,y)
  @test isapprox(Ax_native,Ax_CDS,rtol=10*eps())

  #test random matrix
  x=randn(1000);
  y=zeros(1000);
  A=sprandn(1000,1000,0.1)
  Ax_native = A*x;

  (R,offset)=mat2CDS(A)
  fill!(y,0)
  Ax_CDS = CDS_MVp(size(R,1),size(R,2),R,offset,x,y)
  @test isapprox(Ax_native,Ax_CDS,rtol=10*eps())

end


"""
    x = spectralfactor(a::Polynomial[,method=:roots])

Find a spectral factor `x` for a scalar univariate polynomial `a`.

For a scalar univariate polynomial `a` conjugate symmetric with respect to the imaginary axis, that is, `a=b̃b`, for some `b`, where `b̃(s)=b(-s)` (for complex polynomials the coefficients are complex-conjugated on top of substituting `-s` for `s`), find a polynomial `x` such that `x̃x=a` and `x` is Hurwitz stable (all roots in the left half plane).

Several methods are available: `:roots`, `:xxx`, `:xxx` (`xxx` is just a placeholder for methods to be added, the `:roots` methods (based on computing the roots) is just a baseline solution not indended for serious use.

# Examples

```julia
julia> b = Polynomial([1.0,2.0,3.0],:s)*Polynomial([1,-1],:s)
Polynomial(1.0 + 1.0*s + 1.0*s^2 - 3.0*s^3)

julia> a = cconj(b)*b
Polynomial(1.0 + 1.0*s^2 + 7.0*s^4 - 9.0*s^6)

julia> roots(a)
6-element Array{Complex{Float64},1}:
  -1.0000000000000004 + 0.0im
 -0.33333333333333276 - 0.4714045207910314im
 -0.33333333333333276 + 0.4714045207910314im
  0.33333333333333326 - 0.47140452079103173im
  0.33333333333333326 + 0.47140452079103173im
   1.0000000000000024 + 0.0im

julia> x = spectralfactor(a)
Polynomial(0.9999999999999987 + 2.9999999999999956*s + 4.999999999999998*s^2 + 3.0*s^3)

julia> roots(x)
3-element Array{Complex{Float64},1}:
  -1.0000000000000022 - 2.220446049250313e-16im
 -0.33333333333333287 + 0.47140452079103157im
 -0.33333333333333254 - 0.4714045207910313im
```
"""
function spectralfactor(a::Polynomial)
    an = a[degree(a)]
    r = roots(a)
    m = map(x->isless(real(x),0.0),r)       # Currently ignoring roots on imaginary axis. TO DO!
    s = r[m]                                # Picking just the roots with negative real parts.
    x = fromroots(s,var=a.var)*sqrt(abs(an))
    return x
end

"""
    x = spectralfactor(a::LaurentPolynomial[,method=:roots))

Find a spectral factor `x` for a scalar univariate Laurent polynomial `a`.

For a scalar univariate Laurent polynomial `a` conjugate symmetric with respect to the unit circle, that is, `a=b̃b`, for some `b`, where `b̃(z)=b(1/z)` (for complex polynomials the coefficients are complex-conjugated on top of substituting `1/z` for `z`), find a polynomial `x` such that `x̃x=a` and `x` is Schur stable (all roots in the unit disk).

Several methods are available: `:roots`, `:xxx`, `:xxx` (`xxx` is just a placeholder for methods to be added). The method based on computing the roots is just a baseline solution not indended for serious use.

# Examples

```julia
julia> b = convert(LaurentPolynomial,fromroots([3,-1/2],var=:z));
julia> a = b*dconj(b)
LaurentPolynomial(-1.5*z⁻² + 1.25*z⁻¹ + 9.5 + 1.25*z - 1.5*z²)

julia> roots(a)
4-element Array{Float64,1}:
 -2.000000000000001
 -0.5
  0.33333333333333354
  2.999999999999998

julia> x = spectralfactor(a)
LaurentPolynomial(-0.5000000000000003 + 0.4999999999999994*z + 3.0*z²)

julia> roots(x)
2-element Array{Float64,1}:
 -0.5
  0.33333333333333354
```
"""
function spectralfactor(a::LaurentPolynomial)
    an = a[degree(a)]
    r = roots(a)
    m = map(x->isless(abs(x),1.0),r)    # Currently ignoring roots on the unit circle. TO DO!
    s = r[m]                            # Picking just the roots inside the unit disk.
    x = convert(LaurentPolynomial,fromroots(s,var=a.var))
    xx = x*dconj(x)                     # Finding α could be more efficent, should there be a need.
    α = sqrt(a[0]/(xx[0]))
    x = α*x
    return x
end


@testset "Test DelayPECE method" begin
    function ϕ(x)
        if x == 0
            return 19.00001
        else
            return 19.0
        end
    end

    function f(t, y, ϕ)
        return 3.5*y*(1-ϕ/19)
    end

    h = 0.5
    α = 0.97
    τ = 0.8
    T = 1
    fddeprob = FDDEProblem(f, ϕ, α, τ, T)
    V, y = solve(fddeprob, h, DelayPECE())

    @test V≈[19.0, 19.0, 1.0]
    @test y≈[19.00001, 19.00001, 37.274176448220274]
end

@testset "Test DelayPI method" begin
    function ϕ(x)
        if x == 0
            return 19.00001
        else
            return 19.0
        end
    end
    function f(t, y, ϕ)
        return 3.5*y*(1-ϕ/19)
    end
    prob = FDDEProblem(f, ϕ, 0.97, 0.8, (0, 2))
    result = solve(prob, 0.5, DelayPI())
    @test result≈[19.00001, 19.00001, 19.00001, 18.99999190949352, 18.99997456359874]
end

@testset "Test DelayABM method" begin
    h=0.5; T=5; α=0.97; τ=2; q=0.5
    delayabmfun(t, ϕ, y) = 2*ϕ/(1+ϕ^9.65)-y
    prob = FDDEProblem(delayabmfun, q, α, τ, (0, T))
    x, y=solve(prob, h, DelayABM())

    @test isapprox(x, [1.078559863692747, 1.175963999045738, 1.1661317460354588, 1.128481756921719, 1.0016061526083417, 0.7724564325042358, 0.5974978685646778]; atol=1e-3)
    @test isapprox(y, [0.8889787467894421, 0.9404487875504524, 0.9667449499617093, 0.9803311436135411, 1.078559863692747, 1.175963999045738, 1.1661317460354588]; atol=1e-3)
end

@testset "Test DelayABM for FDDESystem" begin
    function EnzymeKinetics!(dy, y, ϕ, t)
        dy[1] = 10.5-y[1]/(1+0.0005*ϕ[4]^3)
        dy[2] = y[1]/(1+0.0005*ϕ[4]^3)-y[2]
        dy[3] = y[2]-y[3]
        dy[4] = y[3]-0.5*y[4]
    end
    q = [60, 10, 10, 20]; α = [0.95, 0.95, 0.95, 0.95]
    prob = FDDESystem(EnzymeKinetics!, q, α, 4, 6)
    sold, sol = solve(prob, 0.1, DelayABM())

    @test isapprox(sold, [ 56.3     11.3272   11.4284  22.4025
    56.3229  11.2293   11.4106  22.4194
    56.3468  11.141    11.3849  22.4332
    56.3731  11.0592   11.3533  22.4434
    56.4026  10.9824   11.3171  22.4495
    56.4358  10.9096   11.2773  22.4516
    56.4729  10.8398   11.2346  22.4493
    56.5144  10.7725   11.1896  22.4429
    56.5604  10.707    11.1429  22.4322
    56.6111  10.6431   11.0946  22.4175
    56.6668  10.5803   11.0452  22.3987
    56.7275  10.5186   10.9947  22.376
    56.7932  10.4577   10.9435  22.3495
    56.864   10.3975   10.8916  22.3194
    56.9398  10.3382   10.8392  22.2857
    57.0205  10.2795   10.7863  22.2488
    57.1059  10.2217   10.7332  22.2086
    57.1961  10.1647   10.6799  22.1653
    57.2908  10.1087   10.6265  22.1191
    57.3897  10.0537   10.5731  22.0702
    57.4928   9.99987  10.5198  22.0187]; atol=1e-3)

    @test isapprox(sol, [60.5329  10.8225  10.6355  20.6245
    60.3612  10.9555  10.6638  20.6616
    60.1978  11.0679  10.7015  20.6988
    60.0407  11.1632  10.7456  20.7379
    59.8889  11.2442  10.7935  20.7792
    59.7419  11.3129  10.8437  20.8231
    59.5991  11.3711  10.8946  20.8697
    59.4603  11.4202  10.9452  20.9187
    59.3251  11.4613  10.9947  20.9701
    59.1935  11.4954  11.0425  21.0235
    59.0652  11.5234  11.0883  21.0786
    58.9401  11.5461  11.1315  21.1351
    58.818   11.5642  11.1722  21.1927
    58.6988  11.5781  11.21    21.2511
    58.5824  11.5883  11.245   21.3098
    58.4686  11.5954  11.2772  21.3687
    58.3575  11.5997  11.3065  21.4274
    58.2489  11.6016  11.3331  21.4857
    58.1427  11.6013  11.357   21.5433
    58.0389  11.5992  11.3783  21.6001
    57.9373  11.5953  11.3971  21.6557]; atol=1e-3)
end

@testset "Test FDDE Matrix Form method" begin
    τ=3.1416; h=0.5; α=0.4; tspan = (0, 1)
    fx0(t) = [sin(t)*cos(t); sin(t)*cos(t); cos(t)^2-sin(t)^2; cos(t)^2-sin(t)^2]
    A=[0 0 1 0; 0 0 0 1; 0 -2 0 0; -2 0 0 0]
    B=[0 0 0 0; 0 0 0 0; -2 0 0 0; 0 -2 0 0]
    f=[0; 0; 0; 0]
    prob = FDDEMatrixProblem(α, τ, A, B, f, fx0, tspan)
    result=solve(prob, h, MatrixForm())

    @test isapprox(result, [0.139708   0.139708   0.96017    0.96017
    0.479462   0.479462   0.283662   0.283662
    0.378401   0.378401  -0.653644  -0.653644
   -0.07056   -0.07056   -0.989992  -0.989992
   -0.454649  -0.454649  -0.416147  -0.416147
   -0.420735  -0.420735   0.540302   0.540302
    0.0        0.0        1.0        1.0
    0.242142   0.242142   0.932342   0.932342
    0.401838   0.401838   0.60134    0.60134]; atol=1e-4)
end


using Flux
using Zygote: @adjoint
using Statistics

# CPU implementation

"""
  logadd(a, b)

Adds log-space `a` and `b` such that the result equals `log(exp(a)+exp(b))`
"""
function logadd(a, b)
  isinf(a) && return b
  isinf(b) && return a

  # always want the greater number on the left in the exponentiation;
  # the magnitude difference may end up making the number very positive
  # which will cause exp() to return Inf
  # E.g., a = -900, b = -800, will give exp(-800 - -900), which will be
  # Inf for Float32 values
  if a < b
    a, b = b, a
  end

  return a + log(1+exp(b-a))
end

"""
  logsum(a)

Sums the elements in `a` such that the result equals `log(sum(exp.(a)))`
"""
function logsum(a::AbstractArray)
  local s
  s = a[1]
  for item in a[2:end]
    s = logadd(s, item)
  end
  return s
end

"""
  F(A, blank)

Removes blanks and repetitions in the sequence `A`

This is the function `F` as defined in Graves (2012)
"""
function F(A, blank)
  prev = A[1]
  z = [prev]
  for curr in A[2:end]
    if curr != prev && curr != blank
      push!(z, curr)
    end
    prev = curr
  end
  return z
end

"""
  addBlanks(z)

Adds blanks to the start and end of `z`, and between item in `z`
"""
function addBlanks(z, blank)

  z′ = [blank]
  for label in z
    push!(z′, label)
    push!(z′, blank)
  end
  return z′
end

function ctc_(ŷ, y)

  typedZero = zero(ŷ[1])

  ŷ = logsoftmax(ŷ)
  blank = size(ŷ, 1)

  z = F(Base.argmax.([y[:,i] for i=1:size(y,2)]), blank)
  z′ = addBlanks(z, blank)
  T = size(ŷ, 2)
  U = length(z)
  U′ = length(z′)

  # Calculate α coefficients, from the upper-left, to the bottom-right
  α = fill(typedZero, T, U′)
  for t=1:T
    for u=1:U′
      if t == u == 1
        α[t,u] = ŷ[blank, t]
      elseif t == 1 && u == 2
        α[t,u] = ŷ[z′[2], t]
      elseif t == 1 && u > 2
        α[t,u] = log(typedZero)
      elseif u < U′ - 2(T - t) - 1
        α[t,u] = log(typedZero)
      else
        idx = u - 2
        idx += z′[u] == blank || (u > 2 && z′[u-2] == z′[u])
        idx = max(1, idx)

        α[t,u] = ŷ[z′[u], t] + logsum(α[t-1, idx:u])
      end
    end
  end

  # Calculate beta coefficients, from the bottom-right, to the upper-left
  β = fill(log(typedZero), T, U′)

  # Fill bottom-right corner so bounding errors can be avoided
  # by starting `u` at `U′-1`
  β[T,U′] = typedZero
  β[T,U′-1] = typedZero

  # start at T-1 so that β(T, u) = log(0) for all u < U′ - 1
  for t=(T-1):-1:1
    for u=U′:-1:1
      if u > 2t || u > U′ + 1
        continue
      end

      idx = u+2
      idx -= z′[u] == blank || (idx < U′ && z′[u+2] == z′[u])
      idx = min(idx, U′)

      v = [β[t+1,i] + ŷ[z′[i], t+1] for i=u:idx]
      β[t, u] = logsum(v)
    end
  end

  # Loss at each time t is taken as the sum of the product (sum in log space) of the
  # α and β coefficients for all the label classes at time t
  αβ = α + β
  losses = -1 .* mapslices(logsum, αβ, dims=2)

  accum = fill(log(typedZero), size(ŷ))
  grads = fill(log(typedZero), size(ŷ))

  for t=1:T
    for u=1:U′
      accum[z′[u], t] = logadd(accum[z′[u], t], α[t,u] + β[t,u])
    end
    for u=1:size(grads, 1)
      grads[u,t] = exp(ŷ[u, t]) - exp(accum[u, t] - -losses[t])
    end
  end

  return mean(losses), grads
end

"""
  ctc(ŷ, y)

Computes the connectionist temporal classification loss between `ŷ`
and `y`.

Both `ŷ` and `y` must be classes-by-time matrices, i.e., each row
represents a class and each column represents a time step.
Additionally, the `logsoftmax` function will be applied to `ŷ`, so
it must be the raw activation values from the neural network and
not, for example, the activations after being passed through a
`softmax` activation function.

Used for sequence to sequence classification problems such as
speech recognition and handwriting recognition where the exact
time-alignment of the output (e.g., letters) is not needed to
solve the problem. See [Graves et al. (2006)](https://www.cs.toronto.edu/~graves/icml_2006.pdf)
or [Graves (2012)](https://www.cs.toronto.edu/~graves/preprint.pdf#chapter.7)
for mathematical details.
"""
function ctc(ŷ::Array, y::Array)
  return ctc_(ŷ, y)[1]
end

@adjoint function ctc(ŷ::Array, y::Array)
  l,g = ctc_(ŷ, y)
  return l, Δ -> (Δ .* g, nothing)
end

# @adjoint function ctc_(ŷ, y)
#   ls, gs = ctc_(ŷ, y)
#   return ls, Δ -> (Δ .* gs, nothing)
# end


module SparseMatrixAssemblers

using Test
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.Fields
using Gridap.Algebra
using SparseArrays
using SparseMatricesCSR
using Gridap.FESpaces
using Gridap.CellData
using Gridap.Algebra

domain =(0,1,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)

reffe = ReferenceFE(lagrangian,Float64,1)
V = FESpace(model,reffe,dirichlet_tags=[1,2,3,4,6,5])
U = V

b(x) = x[2]

v = get_fe_basis(V)
u = get_trial_fe_basis(U)

degree = 2
trian = get_triangulation(model)
quad = CellQuadrature(trian,degree)

btrian = BoundaryTriangulation(model)
bquad = CellQuadrature(btrian,degree)

b0trian = Triangulation(model,Int[])
b0quad = CellQuadrature(b0trian,degree)

cellmat = integrate(∇(v)⊙∇(u),quad)
cellvec = integrate(v⊙b,quad)
cellmatvec = pair_arrays(cellmat,cellvec)
rows = get_cell_dof_ids(V,trian)
cols = get_cell_dof_ids(U,trian)
cellmat_c = attach_constraints_cols(U,cellmat,trian)
cellmat_rc = attach_constraints_rows(V,cellmat_c,trian)
cellvec_r = attach_constraints_rows(V,cellvec,trian)
cellmatvec_c = attach_constraints_cols(U,cellmatvec,trian)
cellmatvec_rc = attach_constraints_rows(V,cellmatvec_c,trian)

bcellmat = integrate(v*u,bquad)
bcellvec = integrate(v*3,bquad)
bcellmatvec = pair_arrays(bcellmat,bcellvec)
brows = get_cell_dof_ids(V,btrian)
bcols = get_cell_dof_ids(U,btrian)
bcellmat_c = attach_constraints_cols(U,bcellmat,btrian)
bcellmat_rc = attach_constraints_rows(V,bcellmat_c,btrian)
bcellvec_r = attach_constraints_rows(V,bcellvec,btrian)
bcellmatvec_c = attach_constraints_cols(U,bcellmatvec,btrian)
bcellmatvec_rc = attach_constraints_rows(V,bcellmatvec_c,btrian)

b0cellmat = integrate(v*u,b0quad)
b0cellvec = integrate(v*3,b0quad)
b0cellmatvec = pair_arrays(b0cellmat,b0cellvec)
b0rows = get_cell_dof_ids(V,b0trian)
b0cols = get_cell_dof_ids(U,b0trian)
@test length(b0rows) == 0
b0cellmat_c = attach_constraints_cols(U,b0cellmat,b0trian)
b0cellmat_rc = attach_constraints_rows(V,b0cellmat_c,b0trian)
b0cellvec_r = attach_constraints_rows(V,b0cellvec,b0trian)
b0cellmatvec_c = attach_constraints_cols(U,b0cellmatvec,b0trian)
b0cellmatvec_rc = attach_constraints_rows(V,b0cellmatvec_c,b0trian)

term_to_cellmat = [cellmat_rc, bcellmat_rc, b0cellmat_rc]
term_to_cellvec = [cellvec, bcellvec, b0cellvec]
term_to_rows = [rows, brows, b0rows]
term_to_cols = [cols, bcols, b0cols]
term_to_cellmatvec = [ cellmatvec, bcellmatvec, b0cellmatvec ]

mtypes = [
  SparseMatrixCSC{Float64,Int},
  SparseMatrixCSR{0,Float64,Int},
  SparseMatrixCSR{1,Float64,Int},
  SymSparseMatrixCSR{0,Float64,Int},
  SymSparseMatrixCSR{1,Float64,Int}]

for T in mtypes

  matvecdata = ( term_to_cellmatvec , term_to_rows, term_to_cols)
  matdata = (term_to_cellmat,term_to_rows,term_to_cols)
  vecdata = (term_to_cellvec,term_to_rows)
  data = (matvecdata,matdata,vecdata)

  assem = SparseMatrixAssembler(T,Vector{Float64},U,V)
  test_sparse_matrix_assembler(assem,matdata,vecdata,data)

  strategy = GenericAssemblyStrategy(row->row,col->col,row->true,col->true)

  assem2 = SparseMatrixAssembler(T,Vector{Float64},U,V,strategy)
  test_sparse_matrix_assembler(assem2,matdata,vecdata,data)

  matdata = ([cellmat],[rows],[cols])
  vecdata = ([cellvec],[rows])

  mat = assemble_matrix(assem,matdata)
  vec = assemble_vector(assem,vecdata)

  x = mat \ vec

  assemble_matrix!(mat,assem,matdata)
  assemble_vector!(vec,assem,vecdata)

  x2 = mat \ vec
  @test x ≈ x2

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333
  @test mat[1, 2]  ≈ -0.33333333333333
  @test mat[2, 2]  ≈ 2.666666666666666
  @test mat[3, 2]  ≈ -0.33333333333333
  @test mat[2, 3]  ≈ -0.33333333333333
  @test mat[3, 3]  ≈ 1.333333333333333

  data = (([cellmatvec],[rows],[cols]),([],[],[]),([],[]))
  mat, vec = allocate_matrix_and_vector(assem,data)
  assemble_matrix_and_vector!(mat,vec,assem,data)
  assemble_matrix_and_vector!(mat,vec,assem,data)

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333

  x3 = mat \ vec
  @test x ≈ x3

  mat, vec = assemble_matrix_and_vector(assem,data)

  x4 = mat \ vec
  @test x ≈ x4

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333

  mat, vec = assemble_matrix_and_vector(assem2,data)

  x4 = mat \ vec
  @test x ≈ x4

  @test vec ≈ [0.0625, 0.125, 0.0625]
  @test mat[1, 1]  ≈  1.333333333333333
  @test mat[2, 1]  ≈ -0.33333333333333

end

strian = SkeletonTriangulation(model)
squad = CellQuadrature(strian,degree)

scellmat = integrate(jump(v)*u.⁻,squad)
@test isa(scellmat[1],ArrayBlock)
@test scellmat[1][1,1] === nothing
@test scellmat[1][2,1] === nothing
@test scellmat[1][1,2] !== nothing
@test scellmat[1][2,2] !== nothing

scellvec = integrate(mean(v*3),squad)
@test isa(scellvec[1],ArrayBlock)
scellmatvec = pair_arrays(scellmat,scellvec)
srows = get_cell_dof_ids(V,strian)
scols = get_cell_dof_ids(U,strian)
zh = zero(V)
scellvals = get_cell_dof_values(zh,strian)
scellmatvec = attach_dirichlet(scellmatvec,scellvals)

matvecdata = ([scellmatvec],[srows],[scols])
matdata = ([scellmat],[srows],[scols])
vecdata = ([scellvec],[srows])
data = (matvecdata,matdata,vecdata)

assem = SparseMatrixAssembler(U,V)
test_sparse_matrix_assembler(assem,matdata,vecdata,data)

A = assemble_matrix(assem,matdata)
@test A == zeros(num_free_dofs(V),num_free_dofs(U))

assem = SparseMatrixAssembler(
  SparseMatrixBuilder(SparseMatrixCSC{Float64,Int},MinCPU()),U,V)
test_sparse_matrix_assembler(assem,matdata,vecdata,data)

assem = SparseMatrixAssembler(
  SparseMatrixBuilder(SparseMatrixCSC{Float64,Int},MinMemory()),U,V)
test_sparse_matrix_assembler(assem,matdata,vecdata,data)

end # module


using LinearAlgebra
using Distributions
using SpecialFunctions
using PDMats
using Printf

include("kernel.jl")
include("examples.jl")

"""
Define prediction/inference problem by supplying known parameters, including design matrices,
observation locations, prediction locations, nonlinear transform, prior marginals on theta and lambda,
and anticipated ranges for theta and lambda. Returns Bayesian predictive density function,
namely the integral of p(z_0| theta, lambda, z)*p(theta, lambda | z) over Theta x Lambda
(See Equations 8 and 12)

INPUTS:
X: k x p design matrix
X0: n x p design matrix
s: observed locations
s0: prediction locations
g: nonlinear transform parametrized by lambda. Usage: g(lambda, z)
gprime: derivative of g. Usage: gprime(z, lambda)
ptheta: prior marginal on theta; function handle
plambda: prior marginal on lambda; function handle
theta: structural parameter
lambda: transform parameter
z: observed random vector
z0: unobserved random vector

OUTPUTS:
probability density function z0 -> f(z_0|z)
cumulative density function z0 -> F(z0|z)
"""

function model(setting, g, gprime, pθ, pλ, rangeθ, rangeλ)
    s = setting.s; s0 = setting.s0; X = setting.X; X0 = setting.X0; z = setting.z;
    n = size(X, 1)
    p = size(X, 2)
    k = size(X0, 1)

    function func(θ)
        Eθ = K(s0, s0, θ, rbf)
        Σθ = K(s, s, θ, rbf)
        Bθ = K(s0, s, θ, rbf)
        choleskyΣθ = cholesky(Σθ) #precompute Cholesky decomposition of Sigma
        choleskyXΣX = cholesky(Hermitian(X'*(choleskyΣθ\X))) #precompute Cholesky decomposition of XSigma\X
        Dθ = Eθ - Bθ*(choleskyΣθ\Bθ')
        Hθ = X0 - Bθ*(choleskyΣθ\X)
        Cθ = Dθ + Hθ*(choleskyXΣX\Hθ')
        args = [Eθ, Σθ, Bθ, choleskyΣθ, choleskyXΣX, Dθ, Hθ, Cθ]
        return args
    end
    function density(λ, z0, args)
        Eθ, Σθ, Bθ, choleskyΣθ, choleskyXΣX, Dθ, Hθ, Cθ = args
        βhat = (X'*(choleskyΣθ\X))\(X'*(choleskyΣθ\g(z, λ)))
        qtilde = (expr = g(z, λ)-X*βhat; expr'*(choleskyΣθ\expr))
        m = Bθ*(cholesky(Σθ)\g(z, λ)) + Hθ*βhat
        J = λ -> abs(reduce(*, map(x -> gprime(x, λ), z)))
        post = det(cholesky(Σθ))^(-0.5)*det(choleskyXΣX)^(-0.5)*qtilde[1]^(-(n-p)/2)*J(λ)^(1-p/n)*pθ(θ)*pλ(λ)
        jac = abs(reduce(*, map(x -> gprime(x, λ), z0)))
        t = LocationScale(m[1], sqrt(qtilde[1]*Cθ[1]/(n-p)), TDist(n-p))
        return post*Distributions.pdf(t, g(z0, λ))*jac, post*Distributions.cdf(t, g(z0, λ))*jac
    end
    #simultaneously define PDF and CDF - (reuse integration nodes and weights)
    (z0 ->  int1D(θ -> (args = func(θ); int1D(λ -> density(λ, z0, args)[1], rangeλ)), rangeθ),
    z0 ->  int1D(θ -> (args = func(θ); int1D(λ -> density(λ, z0, args)[2], rangeλ)), rangeθ))
end

if false

end


"""
    abstract type AbstractMultipleOrthogonalBasis{P} <: AbstractPolynomialVectorBasis{P, Vector{P}} end

Polynomial basis such that ``\\langle p_i(x), p_j(x) \\rangle = 0`` if ``i \\neq j`` where
```math
\\langle p(x), q(x) \\rangle = \\int p(x)q(x) w(x) dx
```
where the weight is a product of weight functions
``w(x) = w_1(x_1)w_2(x_2) \\cdots w_n(x_n)`` in each variable.
The polynomial of the basis are product of univariate polynomials:
``p(x) = p_1(x_1)p_2(x_2) \\cdots p_n(x_n)``.
where the univariate polynomials of variable `x_i` form an univariate
orthogonal basis for the weight function `w_i(x_i)`.
Therefore, they satisfy the recurrence relation
```math
d_k p_k(x_i) = (a_k x_i + b_k) p_{k-1}(x_i) + c_k p_{k-2}(x_i)
```
where [`reccurence_first_coef`](@ref) gives `a_k`,
[`reccurence_second_coef`](@ref) gives `b_k`,
[`reccurence_third_coef`](@ref) gives `c_k` and
[`reccurence_deno_coef`](@ref) gives `d_k`.
"""
abstract type AbstractMultipleOrthogonalBasis{P} <: AbstractPolynomialVectorBasis{P, Vector{P}} end

"""
    reccurence_first_coef(B::Type{<:AbstractMultipleOrthogonalBasis}, degree::Integer)

Return `a_{degree}` in recurrence equation
```math
d_k p_k(x_i) = (a_k x_i + b_k) p_{k-1}(x_i) + c_k p_{k-2}(x_i)
```
"""
function reccurence_first_coef end

"""
    reccurence_second_coef(B::Type{<:AbstractMultipleOrthogonalBasis}, degree::Integer)

Return `b_{degree}` in recurrence equation
```math
d_k p_k(x_i) = (a_k x_i + b_k) p_{k-1}(x_i) + c_k p_{k-2}(x_i)
```
"""
function reccurence_second_coef end

"""
    reccurence_third_coef(B::Type{<:AbstractMultipleOrthogonalBasis}, degree::Integer)

Return `c_{degree}` in recurrence equation
```math
d_k p_k(x_i) = (a_k x_i + b_k) p_{k-1}(x_i) + c_k p_{k-2}(x_i)
```
"""
function reccurence_third_coef end

"""
    reccurence_deno_coef(B::Type{<:AbstractMultipleOrthogonalBasis}, degree::Integer)

Return `d_{degree}` in recurrence equation
```math
d_k p_k(x_i) = (a_k x_i + b_k) p_{k-1}(x_i) + c_k p_{k-2}(x_i)
```
"""
function reccurence_deno_coef end

"""
    univariate_orthogonal_basis(B::Type{<:AbstractMultipleOrthogonalBasis},
                                variable::MP.AbstractVariable, degree::Integer)

Return the vector of univariate polynomials of the basis `B` up to `degree`
with variable `variable`.
"""
function univariate_orthogonal_basis(
    B::Type{<:AbstractMultipleOrthogonalBasis}, variable::MP.AbstractVariable, degree::Integer)

    @assert degree >= 0
    if degree == 0
        return polynomial_type(B, typeof(variable))[one(variable)]
    elseif degree == 1
        return push!(univariate_orthogonal_basis(B, variable, 0),
                     degree_one_univariate_polynomial(B, variable))
    else
        previous = univariate_orthogonal_basis(B, variable, degree - 1)
        a = reccurence_first_coef(B, degree)
        b = reccurence_second_coef(B, degree)
        c = reccurence_third_coef(B, degree)
        d = reccurence_deno_coef(B, degree)
        next = MA.@rewrite((a * variable + b) * previous[degree] + c * previous[degree - 1])
        if !isone(d)
            next = next / d
        end
        push!(previous, next)
        return previous
    end
end

function _basis_from_monomials(B::Type{<:AbstractMultipleOrthogonalBasis}, variables, monos)
    univariate = [univariate_orthogonal_basis(
                      B, variable, maximum(mono -> degree(mono, variable), monos))
                  for variable in variables]
    return B([
        prod(i -> univariate[i][degree(mono, variables[i]) + 1],
             eachindex(variables))
        for mono in monos])
end

function maxdegree_basis(B::Type{<:AbstractMultipleOrthogonalBasis}, variables, maxdegree::Int)
    return _basis_from_monomials(B, variables, MP.monomials(variables, 0:maxdegree))
end

function basis_covering_monomials(B::Type{<:AbstractMultipleOrthogonalBasis}, monos::AbstractVector{<:AbstractMonomial})
    to_add = collect(monos)
    m = Set(monos)
    while !isempty(to_add)
        mono = pop!(to_add)
        for v in MP.variables(mono)
            step = even_odd_separated(B) ? 2 : 1
            vstep = v^step
            if MP.divides(vstep, mono)
                new_mono = MP.mapexponents(-, mono, vstep)
                if !(new_mono in m)
                    push!(m, new_mono)
                    push!(to_add, new_mono)
                end
            end
        end
    end
    return _basis_from_monomials(B, variables(monos), MP.monovec(collect(m)))
end


@testset " spline pp 1d " begin
    import GEMPIC: SplinePP, b_to_pp, uniform_bsplines_eval_basis, horner_1d

    ncells = 8

    domain = [0.0, 2π]
    delta_x = (domain[2] - domain[1]) / ncells

    b = rand(ncells)

    for d in 1:5
        xp = 4.8141437173987800

        spline_pp = SplinePP(d, ncells)

        pp = b_to_pp(spline_pp, ncells, b)

        xi = (xp - domain[1]) / delta_x
        index = floor(Int64, xi) + 1
        xi = xi - (index - 1)
        res1 = horner_1d(d, pp, xi, index)

        index = index - d

        val = uniform_bsplines_eval_basis(d, xi)

        res2 = 0.0

        for i in 1:(d + 1)
            index1d = (index + i - 2) % ncells + 1
            res2 += b[index1d] * val[i]
        end

        @test res1 ≈ res2

        @test spline_pp.degree ≈ d

        xp = rand()
        index = 1
        res1 = horner_1d(d, pp, xp, index)

        res2 = 0.0
        for i in 1:(d + 1)
            res2 += pp[i, 1] * xp^((d + 1) - i)
        end

        @test res1 ≈ res2
    end
end

@testset " spline pp 2d " begin
    using Random

    ncells = 50
    d = 3

    b = zeros(ncells * ncells)
    pp = zeros((d + 1) * (d + 1), ncells * ncells)

    dx = 2π / ncells
    dy = 2π / ncells

    rng = MersenneTwister(42)
    rand!(rng, b)

    spline1 = SplinePP(d, ncells)
    spline2 = SplinePP(d, ncells)

    GEMPIC.b_to_pp_2d!(pp, spline1, spline2, b)

    xp, yp = rand(rng, 2) .* 2π

    xi = xp / dx
    yi = yp / dy

    ind_x = floor(Int, xi) + 1
    ind_y = floor(Int, yi) + 1

    xi -= (ind_x - 1)
    yi -= (ind_y - 1)

    res1 = GEMPIC.horner_2d((d, d), pp, (xi, yi), (ind_x, ind_y), (ncells, ncells))

    ind_x -= d
    ind_y -= d

    val1 = uniform_bsplines_eval_basis(d, xi)
    val2 = uniform_bsplines_eval_basis(d, yi)

    res2 = 0.0
    for i in 1:(d + 1)
        idx1 = mod(ind_x + i - 2, ncells)
        for j in 1:(d + 1)
            idx2 = mod(ind_y + j - 2, ncells)
            index2d = idx1 + idx2 * ncells + 1
            res2 += b[index2d] * val1[i] * val2[j]
        end
    end

    @test res1 ≈ res2

    xp = rand(rng, 2)

    res1 = GEMPIC.horner_2d((d, d), pp, xp, [1, 1], [1, 1])

    res2 = 0.0
    for i in 1:(d + 1), j in 1:(d + 1)
        res2 += pp[i + (j - 1) * (d + 1), 1] * xp[1]^((d + 1) - i) * xp[2]^((d + 1) - j)
    end

    @test res1 ≈ res2
end


@testset "Angles" begin
  # 2D points
  @test ∠(P2(0,1),P2(0,0),P2(1,0)) ≈ T(-π/2)
  @test ∠(P2(1,0),P2(0,0),P2(0,1)) ≈ T(π/2)
  @test ∠(P2(-1,0),P2(0,0),P2(0,1)) ≈ T(-π/2)
  @test ∠(P2(0,1),P2(0,0),P2(-1,0)) ≈ T(π/2)
  @test ∠(P2(0,-1),P2(0,0),P2(1,0)) ≈ T(π/2)
  @test ∠(P2(1,0),P2(0,0),P2(0,-1)) ≈ T(-π/2)
  @test ∠(P2(0,-1),P2(0,0),P2(-1,0)) ≈ T(-π/2)
  @test ∠(P2(-1,0),P2(0,0),P2(0,-1)) ≈ T(π/2)

  # 3D points
  @test ∠(P3(1,0,0),P3(0,0,0),P3(0,1,0)) ≈ T(π/2)
  @test ∠(P3(1,0,0),P3(0,0,0),P3(0,0,1)) ≈ T(π/2)
  @test ∠(P3(0,1,0),P3(0,0,0),P3(1,0,0)) ≈ T(π/2)
  @test ∠(P3(0,1,0),P3(0,0,0),P3(0,0,1)) ≈ T(π/2)
  @test ∠(P3(0,0,1),P3(0,0,0),P3(1,0,0)) ≈ T(π/2)
  @test ∠(P3(0,0,1),P3(0,0,0),P3(0,1,0)) ≈ T(π/2)

  # 2D vectors
  @test ∠(V2(1, 0), V2(0, 1)) ≈ T(π / 2)
  @test ∠(V2(1, 0), V2(0, -1)) ≈ T(-π / 2)
  @test ∠(V2(1, 0), V2(-1, 0)) ≈ T(π)
  @test ∠(V2(0, 1), V2(-1, 0)) ≈ T(π / 2)
  @test ∠(V2(0, 1), V2(0, -1)) ≈ T(-π)
  @test ∠(V2(0, 1), V2(1, 1)) ≈ T(-π / 4)
  @test ∠(V2(0, -1), V2(1, 1)) ≈ T(π * 3 / 4)
  @test ∠(V2(-1, -1), V2(1, 1)) ≈ T(π)

  # 3D vectors
  @test ∠(V3(0, 0, 1), V3(1, 1, 0)) ≈ T(π / 2)
  @test ∠(V3(1, 0, 1), V3(1, 1, 0)) ≈ T(π / 3)
  @test ∠(V3(-1, -1, 0), V3(1, 1, 0)) ≈ T(π)
  @test ∠(V3(0, -1, -1), V3(0, 1, 1)) ≈ T(π)
  @test ∠(V3(0, -1, -1), V3(0, 1, 0)) ≈ T(π * 3 / 4)
  @test ∠(V3(0, 1, 1), V3(1, 1, 0)) ≈ T(π / 3)
end


# these are complex functions (acb_)
# the corresponding real function values are obtained with a zeroed imaginary part

# void acb_elliptic_k(acb_t res, const acb_t m, slong prec)

"""
    elliptic_k(modulus)

Computes the complete elliptic integral of the first kind  K(m)
"""
function elliptic_k(modulus::ArbComplex{P}) where {P}
    isone(modulus) && return ArbComplex{P}(Inf,0)
    result = ArbComplex{P}()
    ccall(@libarb(acb_elliptic_k), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Cint), result, modulus, P)
    return result
end

"""
    elliptic_k2(modulus)

Computes the complete elliptic integral of the first kind  K(m^2)
"""
elliptic_k2(modulus::ArbComplex{P}) where {P} = elliptic_k(modulus * modulus)

function elliptic_k(modulus::ArbReal{P}) where {P}
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_k(cmodulus)
    result   = real(cresult)
    return result
end

elliptic_k2(modulus::ArbReal{P}) where {P} = elliptic_k(modulus * modulus)

function elliptic_k(modulus::ArbFloat{P}) where {P}
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_k(cmodulus)
    result   = ArbFloat{P}(real(cresult))
    return result
end

elliptic_k2(modulus::ArbFloat{P}) where {P} = elliptic_k(modulus * modulus)

# the complete elliptic integral of the second kind  E(m)

function elliptic_e(modulus::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    ccall(@libarb(acb_elliptic_e), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Cint), result, modulus, P)
    return result
end

elliptic_e2(modulus::ArbComplex{P}) where {P} = elliptic_e(modulus * modulus)

function elliptic_e(modulus::ArbReal{P}) where {P}
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_e(cmodulus)
    result   = real(cresult)
    return result
end

elliptic_e2(modulus::ArbReal{P}) where {P} = elliptic_e(modulus * modulus)

function elliptic_e(modulus::ArbFloat{P}) where {P}
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_e(cmodulus)
    result   = ArbFloat{P}(real(cresult))
    return result
end

elliptic_e2(modulus::ArbFloat{P}) where {P} = elliptic_e(modulus * modulus)

# the complete elliptic integral of the third kind  Π(m)
# void acb_elliptic_pi(acb_t res, const acb_t n, const acb_t m, slong prec)

function elliptic_pi(nu::ArbComplex{P}, modulus::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    ccall(@libarb(acb_elliptic_pi), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, nu, modulus, P)
    return result
end

elliptic_pi2(nu::ArbComplex{P}, modulus::ArbComplex{P}) where {P} = elliptic_pi(nu, modulus * modulus)

function elliptic_pi(nu::ArbReal{P}, modulus::ArbReal{P}) where {P}
    cnu      = ArbComplex(nu, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_pi(cnu, cmodulus)
    result   = real(cresult)
    return result
end

elliptic_pi2(nu::ArbReal{P}, modulus::ArbReal{P}) where {P} = elliptic_pi(nu, modulus * modulus)

function elliptic_pi(nu::ArbFloat{P}, modulus::ArbFloat{P}) where {P}
    cnu      = ArbComplex(nu, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_pi(cnu, cmodulus)
    result   = ArbFloat{P}(real(cresult))
    return result
end

elliptic_pi2(nu::ArbFloat{P}, modulus::ArbFloat{P}) where {P} = elliptic_pi(nu, modulus * modulus)

# Legendre incomplete elliptic integral of the first kind F(phi, m)

function elliptic_f(phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_f), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Cint),
          result, phi, modulus, flag, P)
    return result
end

elliptic_f2(phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P} = elliptic_f(phi, modulus * modulus)

elliptic_k(phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P} = elliptic_f(phi, modulus)

elliptic_k2(phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P} = elliptic_f(phi, modulus * modulus)

function elliptic_f(phi::ArbReal{P}, modulus::ArbReal{P}) where {P}
    cphi     = ArbComplex(phi, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = ellipticf(cphi, cmodulus)
    result   = real(cresult)
    return result
end

elliptic_f2(phi::ArbReal{P}, modulus::ArbReal{P}) where {P} = elliptic_f(phi, modulus * modulus)

elliptic_k(phi::ArbReal{P}, modulus::ArbReal{P}) where {P} = elliptic_f(phi, modulus)

elliptic_k2(phi::ArbReal{P}, modulus::ArbReal{P}) where {P} = elliptic_f(phi, modulus * modulus)

function elliptic_f(phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P}
    cphi     = ArbComplex(phi, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_f(cphi, cmodulus)
    result   = ArbFloat{P}(real(cresult))
    return result
end

elliptic_f2(phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P} = elliptic_f(phi, modulus * modulus)

elliptic_k(phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P} = elliptic_f(phi, modulus)

elliptic_k2(phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P} = elliptic_f(phi, modulus * modulus)

# void acb_elliptic_e_inc(acb_t res, const acb_t phi, const acb_t m, int pi, slong prec)
# Legendre incomplete elliptic integral of the second kind E(phi, m)

function elliptic_e(phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_e_inc), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Cint),
          result, phi, modulus, flag, P)
    return result
end

elliptic_e2(phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P} = elliptice(phi, modulus * modulus)

function elliptic_e(phi::ArbReal{P}, modulus::ArbReal{P}) where {P}
    cphi     = ArbComplex(phi, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_e(cphi, cmodulus)
    result   = real(cresult)
    return result
end

elliptic_e2(phi::ArbReal{P}, modulus::ArbReal{P}) where {P} = elliptic_e_incomplete(phi, modulus * modulus)

function elliptic_e(phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P}
    cphi     = ArbComplex(phi, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_e(cphi, cmodulus)
    result   = ArbFloat{P}(real(cresult))
    return result
end

elliptic_e2(phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P} = elliptic_e(phi, modulus * modulus)

# void acb_elliptic_pi_inc(acb_t res, const acb_t n, const acb_t phi, const acb_t m, int pi, slong prec)
# the Legendre incomplete elliptic integral of the third kind, given by

function elliptic_pi(nu::ArbComplex{P}, phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_pi_inc), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Cint),
          result, nu, phi, modulus, flag, P)
    return result
end

elliptic_pi2(nu::ArbComplex{P}, phi::ArbComplex{P}, modulus::ArbComplex{P}) where {P} =
    elliptic_pi(nu, phi, modulus * modulus)

function elliptic_pi(nu::ArbReal{P}, phi::ArbReal{P}, modulus::ArbReal{P}) where {P}
    cnu      = ArbComplex(nu, 0)
    cphi     = ArbComplex(phi, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_pi(cnu, cphi, cmodulus)
    result   = real(cresult)
    return result
end

elliptic_pi2(nu::ArbReal{P}, phi::ArbReal{P}, modulus::ArbReal{P}) where {P} = elliptic_pi_incomplete(nu, phi, modulus * modulus)

function elliptic_pi(nu::ArbFloat{P}, phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P}
    cnu      = ArbComplex(nu, 0)
    cphi     = ArbComplex(phi, 0)
    cmodulus = ArbComplex(modulus, 0)
    cresult  = elliptic_pi(cnu, cphi, cmodulus)
    result   = ArbFloat{P}(real(cresult))
    return result
end

elliptic_pi2(nu::ArbFloat{P}, phi::ArbFloat{P}, modulus::ArbFloat{P}) where {P} = elliptic_pi_incomplete(nu, phi, modulus * modulus)

#=
# Carleson Symmetric Elliptic Integrals

first kind
void acb_elliptic_rf(acb_t res, const acb_t x, const acb_t y, const acb_t z, int flags, slong prec)
second kind
void acb_elliptic_rg(acb_t res, const acb_t x, const acb_t y, const acb_t z, int flags, slong prec)
third kind
void acb_elliptic_rj(acb_t res, const acb_t x, const acb_t y, const acb_t z, const acb_t p, int flags, slong prec)
=#

function elliptic_rf(x::ArbComplex{P}, y::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flags = 0
    ccall(@libarb(acb_elliptic_rf), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Cint),
          result, x, y, z, flags, P)
    return result
end

function elliptic_rg(x::ArbComplex{P}, y::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flags = 0
    ccall(@libarb(acb_elliptic_rg), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Cint),
          result, x, y, z, flags, P)
    return result
end

function elliptic_rj(x::ArbComplex{P}, y::ArbComplex{P}, z::ArbComplex{P}, p::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flags = 0
    ccall(@libarb(acb_elliptic_rf), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Cint),
          result, x, y, z, p, flags, P)
    return result
end

function elliptic_rf(x::ArbReal{P}, y::ArbReal{P}, z::ArbReal{P}) where {P}
    xx, yy, zz = ArbComplex{P}(x), ArbComplex{P}(y), ArbComplex{P}(z)
    result = ArbReal{P}(real(elliptic_rf(xx, yy, zz)))
    return result
end

function elliptic_rg(x::ArbReal{P}, y::ArbReal{P}, z::ArbReal{P}) where {P}
    xx, yy, zz = ArbComplex{P}(x), ArbComplex{P}(y), ArbComplex{P}(z)
    result = ArbReal{P}(real(elliptic_rg(xx, yy, zz)))
    return result
end

function elliptic_rj(x::ArbReal{P}, y::ArbReal{P}, z::ArbReal{P}, p::ArbReal{P}) where {P}
    xx, yy, zz, pp = ArbComplex{P}(x), ArbComplex{P}(y), ArbComplex{P}(z), ArbComplex{P}(p)
    result = ArbReal{P}(real(elliptic_rj(xx, yy, zz, pp)))
    return result
end

function elliptic_rf(x::ArbFloat{P}, y::ArbFloat{P}, z::ArbFloat{P}) where {P}
    xx, yy, zz = ArbComplex{P}(x), ArbComplex{P}(y), ArbComplex{P}(z)
    result = ArbFloat{P}(real(elliptic_rf(xx, yy, zz)))
    return result
end

function elliptic_rg(x::ArbFloat{P}, y::ArbFloat{P}, z::ArbFloat{P}) where {P}
    xx, yy, zz = ArbComplex{P}(x), ArbComplex{P}(y), ArbComplex{P}(z)
    result = ArbFloat{P}(real(elliptic_rg(xx, yy, zz)))
    return result
end

function elliptic_rj(x::ArbFloat{P}, y::ArbFloat{P}, z::ArbFloat{P}, p::ArbFloat{P}) where {P}
    xx, yy, zz, pp = ArbComplex{P}(x), ArbComplex{P}(y), ArbComplex{P}(z), ArbComplex{P}(p)
    result = ArbFloat{P}(real(elliptic_rj(xx, yy, zz, pp)))
    return result
end


abstract type Grid{D} end

struct Grid2d{M,N,T} <: Grid{2}
    x::Vector{T}
    y::Vector{T}
    x1::T
    x2::T
    y1::T
    y2::T
    normalisation::T
end

function Grid2d(M, N; x1=0, x2=2π, y1=0, y2=2π)
    x = [x1 + i*(x2-x1)/M for i in 0:M-1]
    y = [y1 + j*(y2-y1)/N for j in 0:N-1]
    normalisation = (x2 - x1) * (y2 - y1) / M^2 / N^2
    Grid2d{M,N,typeof(normalisation)}(x, y, x1, x2, y1, y2, normalisation)
end

# TODO Add constructor to read grid from HDF5.
# TODO Add function to write grid to HDF5.

function get_field(grid::Grid2d{M,N,T}) where {M,N,T}
    zeros(T, M, N)
end

@generated function frequencies(r::T, s::T, grid::Grid2d{M,N}) where {M, N, T <: Integer}
    m = circshift(collect(-div(M-1,2):div(M,2)), -div(M-1,2))
    n = collect(0:div(N,2))

    # println("Creating grid->frequency conversion table for (", M, ",", N, ", grid points:")
    # println("   m = ", m)
    # println("   n = ", n)

    quote
        return $m[r], $n[s]
    end
end


using FFTW
export stack_images, align_images, align_channel!, align_images_crop
export align_channel_crop!

@doc raw"""
    align_images(ref::Array, obj::Array, size::Tuple)

This function align two images, using the cross-correlation methods. It works for x-y translations but it fails if the images suffer of field-rotation.
"""
function align_images(ref::Array, obj::Array, size::Tuple)
    ref_fft = fftshift(fft(ref))
    obj_fft = fftshift(fft(obj))
    cx_ft = ref_fft .* conj.(obj_fft)
    cx = abs.(fftshift(ifft(cx_ft)))

    cx_c = Tuple.(indexin(maximum(cx), cx))
    im_c = Int.(round.(size./2))

    shift = im_c .- cx_c

    obj_shifted = zeros(size[1], size[2])

    if (shift[1] > 0) & (shift[2] > 0)

        dx = shift[1]
        dy = shift[2]

        obj_shifted[1:end-dx, 1:end-dy] .= obj[dx+1:end, dy+1:end]
        obj_shifted[end-dx:end, end-dy:end] .= 0.0

    end

    if (shift[1] < 0) & (shift[2] < 0)

        dx = abs(shift[1])
        dy = abs(shift[2])

        obj_shifted[dx+1:end, dy+1:end] = obj[1:end-dx, 1:end-dy]
        obj_shifted[1:dx, 1:dy] .= 0.0

    end

    if (shift[1] > 0) & (shift[2] < 0)

        dx = abs(shift[1])
        dy = abs(shift[2])

        obj_shifted[1:end-dx, dy+1:end] = obj[dx+1:end, 1:end-dy]
        obj_shifted[end-dx:end, 1:dy] .= 0.0
    end

    if (shift[1] < 0) & (shift[2] > 0)

        dx = abs(shift[1])
        dy = abs(shift[2])

        obj_shifted[dx+1:end, 1:end-dy] = obj[1:end-dx, dy+1:end]
        obj_shifted[1:dx, end-dy:end] .= 0.0
    end

    return obj_shifted
end

function align_channel!(channel::Array{Array{UInt16,2},1})

    ref = channel[1]

    for i in 2:length(channel)
            obj_aligned = align_images(ref, channel[i], size(ref))
            channel[i] = obj_aligned
    end

end

function align_channel_crop!(channel::Array{Array{UInt16,2},1}, x::Tuple, y::Tuple)

    ref = channel[1]
    for i in 2:length(channel)
        obj_aligned = align_images_crop(ref, channel[i], size(ref), x, y)
        channel[i] = obj_aligned
    end
end

function align_images_crop(ref::Array, obj::Array, size::Tuple, x::Tuple, y::Tuple)
    partial_r = ref[y[1]:y[2], x[1]:x[2]]
    partial_o = obj[y[1]:y[2], x[1]:x[2]]

    ref_fft = fftshift(fft(partial_r))
    obj_fft = fftshift(fft(partial_o))
    cx_ft = ref_fft .* conj.(obj_fft)
    cx = abs.(fftshift(ifft(cx_ft)))

    cx_c = Tuple.(indexin(maximum(cx), cx))
    partial_size = (y[2]-y[1], x[2]-x[1])
    im_c = Int.(round.(partial_size./2))

    shift = im_c .- cx_c

    obj_shifted = zeros(size[1], size[2])

    if (shift[1] > 0) & (shift[2] > 0)
        println("shift-1 ... ")
        dx = shift[1]
        dy = shift[2]

        obj_shifted[1:end-dx, 1:end-dy] .= obj[dx+1:end, dy+1:end]
        obj_shifted[end-dx:end, end-dy:end] .= 0.0

    end

    if (shift[1] < 0) & (shift[2] < 0)
        println("shift-2 ... ")
        dx = abs(shift[1])
        dy = abs(shift[2])

        obj_shifted[dx+1:end, dy+1:end] = obj[1:end-dx, 1:end-dy]
        obj_shifted[1:dx, 1:dy] .= 0.0

    end

    if (shift[1] > 0) & (shift[2] < 0)
        println("shift-3 ... ")
        dx = abs(shift[1])
        dy = abs(shift[2])

        obj_shifted[1:end-dx, dy+1:end] = obj[dx+1:end, 1:end-dy]
        obj_shifted[end-dx:end, 1:dy] .= 0.0
    end

    if (shift[1] < 0) & (shift[2] > 0)
        println("shift-2 ... ")
        dx = abs(shift[1])
        dy = abs(shift[2])

        obj_shifted[dx+1:end, 1:end-dy] = obj[1:end-dx, dy+1:end]
        obj_shifted[1:dx, end-dy:end] .= 0.0
    end

    return obj_shifted

end

@doc raw"""
    stack_images(img_aligned::Array)

Image staking
"""
function stack_images(img_aligned::Array)
    k = 0
    sum_img = zeros(size(img_aligned[1])[1], size(img_aligned[1])[2])
    for i in img_aligned[1:end]
        sum_img = sum_img .+ i
        k = k + 1
    end

    return sum_img ./ k

end


"""
    life_expectancy(table,age)
    life_expectancy(table,age,DeathDistribution)

Calcuate the remaining life expectancy. Assumes curtate life expectancy for tables if not Parametric or DeathDistribution given.

The life_expectancy of the last age defined in the table is set to be `0.0`, even if the table does not end with a rate of `1.0`.
"""
function life_expectancy(table,age)
    if age == lastindex(table)
        return 0.
    else
        sum(survival(table,age,age + dur) for dur in 1:lastindex(table) -age)
    end
end

function life_expectancy(table,age,dist)
    if age == lastindex(table)
        return 0.
    else
        QuadGK.quadgk(to -> survival(table,age,to+age,dist),0,lastindex(table)-age)[1]
    end

end

function life_expectancy(table::ParametricMortality,age)
    QuadGK.quadgk(to -> survival(table,age,to+age),0,Inf)[1]
end

life_expectancy(table::MortalityTable,args...) = throw(ArgumentError("The first argument should be a vector of rates instead of an entire table. E.g. `table.ulitmate` or `table.select[age]`."))


### ============================================================================
### Build quadratic/linear programming problem using JuMP and various solvers.
### ============================================================================

function MakeLassoOptimizationProblem(
        A::Matrix{Float64},
        y::Matrix{Float64},
        λ::Float64
    )
    """
        Phrase the quadratic programming problem for Lasso regularization
        problem.

    """
    m, n = size(A)
    @assert size(y, 1) == m && size(y, 2) == 1 "the label and"*
    " the data matrix doesn't have the maching dimension X:" *
    string(m, "×",  n) * string(";y: ", size(y))
    @assert λ ≤ 1 && λ ≥ 0 "Regularization λ should be between (0, 1)"

    Pb = Pm.ProgressUnknown("Building Optimization Model: ", spinner=true)

    model = Model(with_optimizer(COSMO.Optimizer); bridge_constraints = false)

    set_optimizer_attribute(model, MOI.Silent(), true)
    set_optimizer_attribute(model, "max_iter", 5000*10)

    @variable(model, x[1:n]); Pm.next!(Pb)
    setvalue.(x, A\y); Pm.next!(Pb)
    @variable(model, η[1:n]); Pm.next!(Pb)
    @constraint(model, -η .<= x); Pm.next!(Pb)
    @constraint(model, x .<= η); Pm.next!(Pb)
    @objective(model, Min, λ*sum(η) + sum((A*x - y).^2)); Pm.finish!(Pb)

    return model

end

function MakeQuantileLassoOptimizationProblem(
        A::Matrix{Float64},
        y::Matrix{Float64},
        λ::Float64=0.0,
        q::Float64=0.5,
    )
    """
        Phrase regularized, quantile regression problem as a linear programming problem.
        problem.

    """
    @assert size(A, 2)  == length(y) "The number of columns for matrix A"*
    " should equals to the number of elements in vector b. "
    @assert λ >= 0 "The L1 regularization parameter should be positive"
    @assert q <= 1 && q >= 0 "The quantile parameter is not percentile"*
    "It should be in between 0 and 1. "

end


#==
This example shows a simple linear bi-objective problem.
==#

using MultiJuMP, JuMP
using Clp

const mmodel = multi_model(Clp.Optimizer, linear = true)
const y = @variable(mmodel, 0 <= y <= 10.0)
const z = @variable(mmodel, 0 <= z <= 10.0)
@constraint(mmodel, y + z <= 15.0)

# objectives
const exp_obj1 = @expression(mmodel, -y +0.05 * z)
const exp_obj2 = @expression(mmodel, 0.05 * y - z)
const obj1 = SingleObjective(exp_obj1)
const obj2 = SingleObjective(exp_obj2)

# # setting objectives in the data
const multim = get_multidata(mmodel)
multim.objectives = [obj1, obj2]

optimize!(mmodel, method = WeightedSum())

using Plots: plot
plot(multim)


#= Code written by Valentín Martpinez Gama @valenmgama with assistance in
    development from Eduardo Garza, Mariela Romano, Jorge Serrano and Manek
    Raffiel.

    To participate in the Bastida Challenge 2020 - Traveling Salesman Problem

     Heuristic Genetic Algorithm"
=#

using Random, CSV, DataFrames
const file = "challenge2020_s1111.csv"

const pop_size = 10000
const elites = 100
const survivors = 4950
const mutts = 3500
const gens = 150

function generate_distance_matrix(file)
    #Importar datos de Excel
    nodes_import = CSV.read(file, header = true)

    #Definir que hay en cada columna del excel y cuantos nodos hay en total
    num_nodes = size(nodes_import,1)
    idcol = 1
    Xcol = 2
    Ycol = 3

    #Generar la Matriz de distancias a partir del numero de nodos
    distance_matrix = Array{Float64}(undef, (num_nodes, num_nodes))
    #Llenar la matriz de distancias entre todos los puntos usando Pitagoras
    for n in 1:num_nodes
        for s in 1:num_nodes
            distance_matrix[n,s] = sqrt((nodes_import[n,Ycol] - nodes_import[s,Ycol])^2 +
            (nodes_import[n,Xcol] - nodes_import[s,Xcol])^2)
        end
    end
    #println(pretty_table(nodes_import))
    #println(pretty_table(distance_matrix))
    return distance_matrix
end

const distance_matrix = generate_distance_matrix(file)
const num_nodes = size(distance_matrix, 1)

function fitness_function(path)
    return distance_matrix[path[end],path[1]] +
        sum(distance_matrix[path[i-1],path[i]] for i = 2:num_nodes)
end

struct Individual
    generation::Int
    genes::Vector{Int}
    fitness_score::Float64
end
Individual(gen, genes) = Individual(gen, genes, fitness_function(genes))

function init_poulation(pop_size)
    population = Vector{Individual}(undef,pop_size)
    for i in 1:pop_size
        population[i] = Individual(1,randperm(num_nodes))
    end
    return population
end

function new_generation!(population, mating_pool, gen)

    for p in 1:survivors
        parent1 = mating_pool[p]
        parent2 = mating_pool[rand(1:survivors)]

        for (c,I) in enumerate(cross_parents(parent1.genes, parent2.genes, gen))
            if c == 1
                population[elites+p] = I
            elseif c == 2
                population[elites+survivors+p] = I
            else
                print("Error")
            end
        end
    end

    for m in 1:(mutts÷2)
        mutt_1 = population[rand(10:pop_size)]
        mutt_2 = population[rand(10:pop_size)]
        mutate_type1!(mutt_1)
        mutate_type2!(mutt_2)
    end
end

function cross_parents(parent1_genes, parent2_genes, gen)

    cross_point = rand(1:num_nodes-1)

    ch1_genes = similar(parent1_genes)
    ch2_genes = similar(parent2_genes)
    ch1_genes[1:cross_point] = @view parent1_genes[1:cross_point]
    ch2_genes[1:cross_point] = @view parent2_genes[1:cross_point]
    parent1_left = @view parent1_genes[cross_point+1:end]
    parent2_left = @view parent2_genes[cross_point+1:end]

    ch1_pending = similar(parent1_left, Tuple{Int, Int})
    ch2_pending = similar(parent2_left, Tuple{Int, Int})

    for (index, value) in enumerate(parent1_left)
        ch1_pending[index] = (value, findfirst(isequal(value), parent2_genes)::Int64)
    end
    sort!(ch1_pending, by = last)
    for (index, (value, _)) in enumerate(ch1_pending)
        ch1_genes[cross_point+index] = value
    end

    for (i, v) in enumerate(parent2_left)
        ch2_pending[i] = (v, findfirst(isequal(v), parent1_genes)::Int64)
    end
    sort!(ch2_pending, by = last)
    for (i,(v, _)) in enumerate(ch2_pending)
        ch2_genes[cross_point+i] = v
    end

    return (Individual(gen,ch1_genes), Individual(gen,ch2_genes))
end

function mutate_type1!(mutt_1::Individual)
    mutt_gene = rand(1:num_nodes-1)
    i = findfirst(isequal(mutt_gene), mutt_1.genes)
    j = findfirst(isequal(mutt_gene+1), mutt_1.genes)
    mutt_1.genes[i] = mutt_gene+1
    mutt_1.genes[j] = mutt_gene
end

function mutate_type2!(mutt_2::Individual)
    mutt_gene = rand(1:num_nodes-2)
    i = findfirst(isequal(mutt_gene), mutt_2.genes)
    j = findfirst(isequal(mutt_gene+2), mutt_2.genes)
    mutt_2.genes[i] = mutt_gene+2
    mutt_2.genes[j] = mutt_gene
end

function GA()
    population = init_poulation(pop_size)
    mating_pool = Array{Individual}(undef, survivors)

    for g = 2:gens

        partialsort!(population, survivors, by = p -> p.fitness_score)
        mating_pool = @view population[1:survivors]

        new_generation!(population, mating_pool, g)

        if (population[1].fitness_score == population[elites].fitness_score) || ((g - population[1].generation) > 6)
            print("Convergencia en la GEN")
            println(g)
            println("TODOS ELITES IGUALES")
            break
        end
        if ((g - population[1].generation) > gens/5)
            print("Convergencia en la GEN")
            println(g)
            print("SIN CAMBIOS en x GENS x = ")
            println(gens/5)
            break
        end
        println(g)
    end
    println()
    sort!(population, by = p -> p.fitness_score)
    return population[1]
end

@time winner = GA()
print(winner)

function arcs_function(sequence)
    arcs = zeros(Int, num_nodes, num_nodes)
    i = sequence[1]
    for s in 2:length(sequence)
        j = sequence[s]
        arcs[i,j] = 1
        i = j
    end
    arcs[i,sequence[1]] = 1
    return arcs
end

arc_matrix = arcs_function(winner.genes)

df = DataFrame(arc_matrix)
CSV.write("tabla_arcos.csv", df)


nt = 1000
seed = 99
m = ABCDQR_simula(nt,seed)
y = m["y"]
u = m["u"]

############################### starting point
wi = collect(range(0.1, length=2, stop=3))
Ai,Bi,Ci,Di,Qi,Ri,Si,m1i,P1i = ABCDQR_em_starting(y,u,wi,1,true)

#######################
max_iter=10
tol=1e-6
txo=false
print("Testing ABCDQR_em: ")
A,B,C,D,Q,R,m1,P1,loglik,aic = ABCDQR_em(y,u,Ai,Bi,Ci,Di,Qi,Ri,m1i,P1i,max_iter,tol,txo)
if sum(A.^2) > 1e-6
    println("OK")
end


using Gadfly, DataFrames

#demonstration of BLE channel hopping noise avoidance
function gaussian(x, mu, sigmaSquared)
    sigma = sqrt(sigmaSquared)
    1 / (sigma * sqrt(2 * pi)) * exp(-1.0 * ((x - mu) ^ 2) / (2.0 * (sigmaSquared) ))
end

WiFiFreqs = [2412, 2437, 2462]
BLEFreqs = vcat([2404 : 2 : 2424], [2428 : 2 : 2478])
x = [2400 : 1 : 2500]

WiFiPlots = [2.0 * gaussian(ix, iy, 40.0) for ix=x, iy=WiFiFreqs]

BLEValues = [ ((findfirst(BLEFreqs, ix) > 0) ? 0.127 : 0/0) for ix=x]
BLEChannels = [findfirst(BLEFreqs, ix) for ix=x]
BLEChannelStrs = [(ix < 10)?"0$ix":"$ix" for ix=BLEChannels]
df = DataFrame(frequency=x, value=BLEValues, channel=BLEChannels)

plot(
    layer(x=x, y=WiFiPlots[:, 1], Geom.line, Theme(default_color=color("#95a5b5"), line_width=7px)),
    layer(x=x, y=WiFiPlots[:, 2], Geom.line, Theme(default_color=color("#95a5b5"), line_width=7px)),
    layer(x=x, y=WiFiPlots[:, 3], Geom.line, Theme(default_color=color("#95a5b5"), line_width=7px)),
    layer(df, x="frequency", y="value", color="channel", Geom.bar,
          Theme(default_color=color("#ff8585"), default_point_size=5px)),
     Scale.y_sqrt
)


@testset verbose=true "indexing" begin
    import SphericalFunctions: WignerHsize, WignerHrange, WignerHindex
    import SphericalFunctions: WignerDsize, WignerDrange, WignerDindex
    import SphericalFunctions: Ysize, Yrange, Yindex, deduce_limits
    import SphericalFunctions: theta_phi

    ell_max = 16
    ell_max_slow = ell_max ÷ 2

    @testset "WignerHrange" begin
        r(mp_max, ell_max) = hcat([
            [ell, mp, m] for ell in 0:ell_max
                for mp in -min(ell, mp_max):min(ell, mp_max)
                    for m in abs(mp):ell
                        ]...)'
        for ell_max in 0:ell_max
            a = WignerHrange(ell_max)  # Implicitly, mp_max=ell_max
            b = r(ell_max, ell_max)
            @test a == b
            for mp_max in 0:ell_max
                a = WignerHrange(mp_max, ell_max)
                b = r(mp_max, ell_max)
                @test a == b
            end
        end
    end

    @testset "WignerHsize" begin
        for ell_max in 0:ell_max
            @test WignerHsize(ell_max) == size(WignerHrange(ell_max, ell_max))[1]
            for mp_max in ell_max
                @test WignerHsize(mp_max, ell_max) == size(WignerHrange(mp_max, ell_max))[1]
            end
        end
    end

    @testset "WignerHindex" begin
        function fold_H_indices(ell, mp, m)
            if m < -mp
                if m < mp
                    return [ell, -mp, -m]
                else
                    return [ell, -m, -mp]
                end
            else
                if m < mp
                    return [ell, m, mp]
                else
                    return [ell, mp, m]
                end
            end
        end

        for ell_max_i in 0:ell_max_slow
            r = WignerHrange(ell_max_i)
            for ell in 0:ell_max_i
                for mp in -ell:ell
                    for m in -ell:ell
                        i = WignerHindex(ell, mp, m)
                        @test r[i, :] == fold_H_indices(ell, mp, m)
                    end
                end
            end
            for mp_max in 0:ell_max_i
                r = WignerHrange(mp_max, ell_max_i)
                for ell in 0:ell_max_i
                    for mp in -ell:ell
                        for m in -ell:ell
                            if abs(mp) > mp_max && abs(m) > mp_max
                                continue
                            end
                            i = WignerHindex(ell, mp, m, mp_max)
                            @test r[i, :] == fold_H_indices(ell, mp, m)
                        end
                    end
                end
            end
        end
    end

    @testset "WignerDrange" begin
        function r(ell_min, mp_max, ell_max)
            a = [
                [ℓ, mp, m]
                for ℓ in ell_min:ell_max
                for mp in -min(ℓ, mp_max):min(ℓ, mp_max)
                for m in -ℓ:ℓ
            ]
            collect(transpose(reshape(collect(Iterators.flatten(a)), (3, length(a)))))
        end

        for ell_max in 0:ell_max_slow÷2
            for ell_min in 0:ell_max
                a = WignerDrange(ell_min, ell_max)  # Implicitly, mp_max=ell_max
                b = r(ell_min, ell_max, ell_max)
                @test a == b
                for mp_max in 0:ell_max
                    a = WignerDrange(ell_min, mp_max, ell_max)
                    b = r(ell_min, mp_max, ell_max)
                    @test a == b
                end
            end
        end
    end

    @testset "WignerDsize" begin
        for ell_max in 0:ell_max
            for ell_min in 0:ell_max
                for mp_max in 0:ell_max
                    a = WignerDsize(ell_min, mp_max, ell_max)
                    b = size(WignerDrange(ell_min, mp_max, ell_max))[1]
                    @test a == b
                end
            end
        end
        for ell_max in 0:ell_max
            for ell_min in [0]
                for mp_max in 0:ell_max
                    a = WignerDsize(ell_min, mp_max, ell_max)
                    #a = WignerDsize_mpmax(ell_max, mp_max)
                    b = size(WignerDrange(ell_min, mp_max, ell_max))[1]
                    @test a == b
                end
            end
        end
        for ell_max in 0:ell_max
            for ell_min in 0:ell_max
                for mp_max in [ell_max]
                    a = WignerDsize(ell_min, mp_max, ell_max)
                    # a = WignerDsize_ellmin(ell_min, ell_max)
                    b = size(WignerDrange(ell_min, mp_max, ell_max))[1]
                    @test a == b
                end
            end
        end
        for ell_max in 0:ell_max
            for ell_min in [0]
                for mp_max in [ell_max]
                    a = WignerDsize(ell_min, mp_max, ell_max)
                    # a = WignerDsize(ell_max)
                    b = size(WignerDrange(ell_min, mp_max, ell_max))[1]
                    @test a == b
                end
            end
        end
    end

    @testset "WignerDindex" begin
        function r(ell_min, mp_max, ell_max)
            a = [
                [ℓ, mp, m]
                for ℓ in ell_min:ell_max
                for mp in -min(ℓ, mp_max):min(ℓ, mp_max)
                for m in -ℓ:ℓ
            ]
            collect(transpose(reshape(collect(Iterators.flatten(a)), (3, length(a)))))
        end

        for ellmax in 0:ell_max_slow
            r = WignerDrange(0, ellmax)
            for ell in 0:ellmax
                for mp in -ell:ell
                    for m in -ell:ell
                        i = WignerDindex(ell, mp, m)
                        @test r[i, :] == [ell, mp, m]
                    end
                end
            end
            for ell_min in 0:ellmax
                r = WignerDrange(ell_min, ellmax)
                for ell in ell_min:ellmax
                    for mp in -ell:ell
                        for m in -ell:ell
                            i = WignerDindex(ell, mp, m, ell_min)
                            @test r[i, :] == [ell, mp, m]
                        end
                    end
                end
                for mp_max in 0:ellmax
                    r = WignerDrange(ell_min, mp_max, ellmax)
                    for ell in ell_min:ellmax
                        for mp in -min(ell, mp_max):min(ell, mp_max)
                            for m in -ell:ell
                                i = WignerDindex(ell, mp, m, ell_min, mp_max)
                                @test r[i, :] == [ell, mp, m]
                            end
                        end
                    end
                end
            end
        end
    end

    @testset "Yrange" begin
        function r(ell_min, ell_max)
            a = [
                [ℓ, m]
                for ℓ in ell_min:ell_max
                for m in -ℓ:ℓ
            ]
            collect(transpose(reshape(collect(Iterators.flatten(a)), (2, length(a)))))
        end
        for ell_max in 0:ell_max
            for ell_min in 0:ell_max
                a = Yrange(ell_min, ell_max)
                b = r(ell_min, ell_max)
                @test a == b
            end
        end
    end

    @testset "Ysize" begin
        for ell_max in 0:ell_max
            for ell_min in 0:ell_max
                a = Ysize(ell_min, ell_max)
                b = size(Yrange(ell_min, ell_max))[1]
                @test a == b
            end
        end
    end

    @testset "deduce_limits" begin
        for ℓmax in 0:4096
            for ℓmin in 0:min(2, ℓmax)
                deduced = deduce_limits(Ysize(ℓmin, ℓmax), nothing)
                @test deduced == (ℓmin, ℓmax)
            end
        end
        for ℓmax in 0:ell_max
            for ℓmin in [0]
                deduced = deduce_limits(Ysize(ℓmin, ℓmax))
                @test deduced == (ℓmin, ℓmax)
            end
        end
        for ℓmax in 0:ell_max
            for ℓmin in 0:ℓmax
                deduced = deduce_limits(Ysize(ℓmin, ℓmax), ℓmin)
                @test deduced == (ℓmin, ℓmax)
            end
        end
    end

    @testset "Yindex" begin
        for ell_max in 0:ell_max
            for ell_min in [0]
                r = Yrange(ell_min, ell_max)
                for ell in ell_min:ell_max
                    for m in -ell:ell
                        i = Yindex(ell, m)
                        @test r[i, :] == [ell, m]
                    end
                end
            end
            for ell_min in 0:ell_max
                r = Yrange(ell_min, ell_max)
                for ell in ell_min:ell_max
                    for m in -ell:ell
                        i = Yindex(ell, m, ell_min)
                        @test r[i, :] == [ell, m]
                    end
                end
            end
        end
    end

end


###########################################################

# random sample of size n, from the qExponential(q,b) distribution
function rqexp(n,q,b)
   us = rand(Uniform(0,1),n)
   invals = (ones(n) - (ones(n) - us).^((q-1)/q))/(b*(q-1))
   return(invals)
end

###########################################################

# pdf of the qExponential(q,b) distribution
function dqexp(xs,q,b)
   ys = q*b*(ones(length(xs))-(q-1)*b*xs).^(1/(q-1))
   return(ys)
end

###########################################################
# Monte Carlo repetitions of the qExponential \beta-parameter and MSE calculation (Figure 3)
# Estimation methods: Shannon - Renyi (linear constraints) - Renyi (escort constraints)
# ss: sample size, nreps: number of Monte Carlo repetitions, q,b: parameters of the qExponential

function monteCarlo_betas(ss, nreps, q, b)
   qhat = zeros(nreps)
   bhat_sh = zeros(nreps)
   bhat_r1 = zeros(nreps)
   bhat_r2 = zeros(nreps)
   betas = zeros(nreps,3)
   for i in 1:nreps
      global data
      # generate qexp sample of size ss
      data = rqexp(ss,q,b)
      ##################################################################
      g(x) = 1/length(data)*sum(data ./ (1 .+ data ./ x)) * (1. + (length(data) / sum(log.(1 .+ data / x))) ) - x
      strue = 1/(b * (1-q))
      sigmahat = find_zero(g, rand(Uniform(strue-ee,strue+ee)))
      # calculate estimators
      thetahat = length(data)/sum(log.(1 .+ data / sigmahat))
      qhat[i] = thetahat / (thetahat + 1.)
      ##
      bhat_sh[i] = 1. / (sigmahat * (1. - qhat[i]))
      bhat_r1[i] = 1/((2q-1)*mean(data)) #linear constraint estimator
      bhat_r2[i] = 1/(q*mean(data)) #escort constraint estimator
   end
   betas[:,1] = bhat_sh
   betas[:,2] = bhat_r1
   betas[:,3] = bhat_r2
   return(betas)
end

###########################################################

# Monte Carlo repetitions of the qExponential \beta-parameter and MSE calculation (Figure 5)
# Estimation methods: Shannon - Renyi (linear constraints) - Renyi (escort constraints)
# ss: sample size, nreps: number of Monte Carlo repetitions, q,b: parameters of the qExponential
function monteCarloMSE(ss, nreps, q, b)
   # initialization
   qhat = zeros(nreps)
   bhat_sh = zeros(nreps)
   bhat_r1 = zeros(nreps)
   bhat_r2 = zeros(nreps)
   for i in 1:nreps
      global data
      # generate qexp sample of size ss
      data = rqexp(ss,q,b)
      ##################################################################
      g(x) = 1/length(data)*sum(data ./ (1 .+ data ./ x)) * (1. + (length(data) / sum(log.(1 .+ data / x))) ) - x
      strue = 1/(b * (1-q))
      sigmahat = find_zero(g, rand(Uniform(strue-ee,strue+ee)))
      # calculate estimators
      thetahat = length(data)/sum(log.(1 .+ data / sigmahat))
      qhat[i] = thetahat / (thetahat + 1.)
      ##
      bhat_sh[i] = 1. / (sigmahat * (1. - qhat[i])) # shannon estimator (coincides with MLE)
      bhat_r1[i] = 1/((2q-1)*mean(data)) # Renyi linear constraint estimator
      bhat_r2[i] = 1/(q*mean(data)) # Renyi escort constraint estimator
   end
   mser = zeros(4)
   # calculate MC MS errors
   mser[1] = var(qhat) + (mean(qhat)-q)^2
   mser[2] = var(bhat_sh) + (mean(bhat_sh)-b)^2
   mser[3] = var(bhat_r1) + (mean(bhat_r1)-b)^2
   mser[4] = var(bhat_r2) + (mean(bhat_r2)-b)^2
   return(mser)
end

###########################################################


using CairoMakie, Random, LaTeXStrings
Random.seed!(123)
m = 200
function spiral(; a = 1, n = 100, h = 0, k = 0)
    φ = LinRange(rand() + 1, 6π, n)
    h .+  rand(-1:2:1)*a*cos.(φ)./φ, k .+  rand(-1:2:1)*a*sin.(φ)./φ
end
x = raw"x=a\,\cos(\varphi)/\varphi,\quad "
y = raw"y=a\,\sin(\varphi)/\varphi,\quad "
curves = [spiral(; a = rand(), h = rand(-1:1)) for i in 1:m]

with_theme(theme_dark()) do
    fig = Figure(resolution = (600,400))
    ax = Axis(fig[1,1])
    series!(curves; color = to_colormap(:bone_1, m), linewidth = 1.5)
    text!(latexstring(x * y * "\\varphi>0"), position = (-1,-1))
    hidedecorations!(ax; grid = false)
    save("series.png", fig, px_per_unit = 2)
    display(fig)
end