# #######################################
# Example 9.8 and 10.24 Grouped Multinomial Data of
# Robert, C. P., & Casella, G. (2010). Monte Carlo statistical methods (2. ed., softcover reprint of the hardcover 2. ed. 2004). New York, NY: Springer.
# #######################################

using Distributions
using Plots
using LaTeXStrings

# configure
a = [0.06, 0.14, 0.11, 0.09]
b = [0.17, 0.24, 0.19, 0.20]
x = [9, 15, 12, 7, 8]

# normal gibbs
function gibbs(M = 1000)
    Z = ones(M, 4)
    Params = ones(M, 2)
    Eparams = ones(M, 2)
    # prior
    alpha = [1/2, 1/2, 1/2]
    # initial
    mu = 0.1
    eta = 0.1
    for m in 1:M
        # sample zi
        zi = ones(4)
        for i = 1:2
            aimu = a[i] * mu
            zi[i] = rand(Binomial(x[i], aimu / (aimu + b[i])))
        end
        for i = 3:4
            aieta = a[i] * eta
            zi[i] = rand(Binomial(x[i], aieta / (aieta + b[i])))
        end
        # sample mu eta
        mu, eta, = rand(Dirichlet([ zi[1] + zi[2] + alpha[1], zi[3] + zi[4] + alpha[2], x[5] + alpha[3] ]))
        # store
        Z[m, :] = zi
        Params[m, :] = [mu, eta]
        #Eparams[m, :] = mean(Params[1:m, :], dims = 1)
    end
    Eparams[:, 1] = cumsum(Params[:, 1]) ./ Array(1:M)
    Eparams[:, 2] = cumsum(Params[:, 2]) ./ Array(1:M)
    return Z, Params, Eparams
end
Z, Params, Eparams = gibbs(10000)

plot(Eparams[:,1], label = L"\mu")
plot!(Eparams[:,2], label = L"\eta")
savefig("gibbs.png")

# Metropolization
function Metrogibbs(M = 1000)
    Z = ones(Int, M, 4)
    Params = ones(M, 2)
    Eparams = ones(M, 2)
    # prior
    alpha = [1/2, 1/2, 1/2]
    # initial
    mu = 0.1
    eta = 0.1
    for m in 1:M
        # sample zi
        zi = ones(Int, 4)
        for i = 1:4
            if i == 1 || i == 2
                aimu = a[i] * mu
            else
                aimu = a[i] * eta
            end
            p1 = aimu / (aimu + b[i])
            p2 = 1 - p1
            while true
                zi[i] = rand(Binomial(x[i], p1))
                if m == 1 || zi[i] != Z[m-1, i] # short-circuiting boolean
                    break
                end
            end
            # accept or not
            if m == 1
                rho = 1.1
            else
                num1 = 1 - binomial(Z[m-1, i], x[i]) * p1^Z[m-1,i] * p2^(x[i] - Z[m-1, i])
                num2 = 1 - binomial(zi[i], x[i]) * p1^zi[i] * p2^(x[i] - zi[i])
                rho = num1 / num2
            end
            if rand() > rho
                zi[i] = Z[m-1, i]
            end

        end
        # sample mu eta
        mu, eta, = rand(Dirichlet([ zi[1] + zi[2] + alpha[1], zi[3] + zi[4] + alpha[2], x[5] + alpha[3] ]))
        # store
        Z[m, :] = zi
        Params[m, :] = [mu, eta]
        #Eparams[m, :] = mean(Params[1:m, :], dims = 1)
    end
    Eparams[:, 1] = cumsum(Params[:, 1]) ./ Array(1:M)
    Eparams[:, 2] = cumsum(Params[:, 2]) ./ Array(1:M)
    return Z, Params, Eparams
end

Z, Params, Eparams = Metrogibbs(10000)
plot(Eparams[:,1], label = L"\mu")
plot!(Eparams[:,2], label = L"\eta")
savefig("metrogibbs.png")


"""
`sort_by_size(list)` returns a sorted copy of `list`
in which the elements are sorted by the `length` function
(from smallest to largest).
"""
function sort_by_size(list::Array)
    sort(list,lt=((x,y) -> length(x) < length(y)))
end

import Base.sqrt

function weave(c::Array{Int,1})::Array{Int,1}
    nc = length(c)
    if nc%2 == 0
        error("Attempting to weave $c, an even-length cycle, with itself")
    end
    half = Int((nc+1)/2)
    result = zeros(Int,nc)
    j = 1  # index into c
    jj = 1  # index into result
    while j <= nc
        result[j] = c[jj]
        j += 1
        jj = (jj+half) % nc
        if jj == 0
            jj = nc
        end
    end
    return result
end

function weave(c1::Array{Int,1}, c2::Array{Int,1})::Array{Int,1}
    n1 = length(c1)
    n2 = length(c2)
    @assert n1==n2 "Cycles $c1 and $c2 have different lengths"
    @assert n1%2==0 "Cycles must have even length to be weaved together"

    result = zeros(Int,2n1)
    for j=1:n1
        result[2j-1]=c1[j]
        result[2j]=c2[j]
    end
    return result
end
"""
`sqrt(p::Permutation)` returns a `Permutation` `q` such that
`q*q==p`. Note: There may be other square roots besides the one
returned. If `p` does not have a square root, an error is thrown.
"""
function sqrt(p::Permutation)::Permutation
    err_msg = "This permutation does not have a square root"
    n = length(p)
    clist = sort_by_size(cycles(p))
    nc = length(clist)
    result = zeros(Int,n)

    idx = 1
    while idx <= nc
        cyc = clist[idx]  # next cycle to process
        m = length(cyc)
        if m%2 == 1   # an odd cycle
            tmp = weave(cyc)
            for j=1:m-1
                result[tmp[j]] = tmp[j+1]
            end
            result[tmp[m]] = tmp[1]
            idx += 1
        else
            if idx == nc
                error(err_msg)
            end
            cyc2 = clist[idx+1]
            if length(cyc2) != length(cyc)
                error(err_msg)
            end
            tmp = weave(cyc,cyc2)
            for j=1:2m-1
                result[tmp[j]] = tmp[j+1]
            end
            result[tmp[2m]] = tmp[1]
            idx += 2
        end
    end
    return Permutation(result)
end

√(p::Permutation) = sqrt(p)


# ------------------------------------------------------------------------------------------
# # Creating a dynamical system
#
# Topics:
# * What *is* a dynamical system in the context of **DynamicalSystems.jl**?
# * How can we define a discrete or continuous system?
# * What are the differences in handling large and small systems?
#     * Handy Dandy definition table!
# * Comment on DifferentialEquations.jl
# * Docstrings
#
# ---
#
# # Dynamical System
#
# A "dynamical system" is a law that describes how some variables should evolve in time.
# This law is described by the "equations of motion" function $\vec{f}$.
#
# There are two types of dynamical systems (in our case):
#
# 1. One is called a map, where time is a discrete quantity (like "steps" or "generations").
# The equations of motion then look something like $\vec{u}_{n+1} = \vec{f}(\vec{u}_n, p,
# n)$ where $n$ is an integer and $p$ are the parameters of the system.
#
# 2. The other type is called an Ordinary Differential Equation (ODE), where time is a
# continuous quantity. Then, the equations of motion give the time derivatives of the
# variables: $\frac{d\vec{u}}{dt} = \vec{f}(\vec{u}, p, t)$.
#
# In both cases $\vec{u}$ is the state of the system, a vector of the variables that define
# the system. For example, $\vec{u}$ may be defined as $(x, y, z)$, or as $(\theta, I)$, or
# as $(V, I, R, C_1, C_2)$, etc.
#
# ---
#
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Creating a simple discrete system
#
#
# For all intents and purposes, a `DynamicalSystem` is a Julia `struct` that contains all
# the information of a dynamical system:
# 1. Equations of motion function $\vec{f}$
# 3. Initial state.
# 4. Parameter container (if the equations of motion are parameterized).
# 2. *Optional* : Jacobian of the equations of motion.
#
#
# For simplicity let's focus on creating a simple discrete chaotic system, the [Hénon
# map](https://en.wikipedia.org/wiki/H%C3%A9non_map).
#
# The constructor we need is:
#
# ```julia
# DiscreteDynamicalSystem(eom, state, p)
# ```
#
# In order to construct a Hénon map, let's discuss the input arguments we need to pass to
# this constructor!
#
# #### 1. Equations of motion
# The first argument of the above constructor --`eom` -- is a **function** representing the
# equations of motion.
#
# Therefore, the first step in creating a `DynamicalSystem` is to define the equations of
# motion. For the Hénon map they have the form:
#
# $$
# \begin{aligned}
# x_{n+1} &= 1 - ax^2_n+y_n \\
# y_{n+1} & = bx_n
# \end{aligned}
# $$
#
# Now we must make a function out of them. There is some strictness when writing this
# function in Julia. Specifically, it can only be one of two forms: out-of-place (oop) or
# in-place (iip). Here is how to define it:
#
# * **oop** : The `eom` function **must** be in the form `eom(x, p, t) -> SVector`
#   which means that given a state `x::SVector` and some parameter container
#   `p` it returns a
# [`SVector`](http://juliaarrays.github.io/StaticArrays.jl/stable/pages/api.html#SVector-1)
#   (from the [StaticArrays](https://github.com/JuliaArrays/StaticArrays.jl) module)
#   containing the next state.
# * **iip** : The `eom` function **must** be in the form `eom!(xnew, x, p, t)`
#   which means that given a state `x::Vector` and some parameter container `p`,
#   it writes in-place the new state in `xnew`.
#
# We advise to use **oop** for systems with dimension < 11, and use **iip** otherwise.
#
# *If you are familiar with [DifferentialEquations.jl](http://docs.juliadiffeq.org/latest/),
# then notice that the equations of motion are defined in an identical manner*
#
# Because the Hénon map is only 2-dimensional, we follow the advice of the documentation and
# take advantage of the [`StaticArrays`](https://github.com/JuliaArrays/StaticArrays.jl)
# module.
# ------------------------------------------------------------------------------------------

using DynamicalSystems
h_eom(x, p, t) = SVector{2}(1.0 - p[1]*x[1]^2 + x[2], p[2]*x[1])

# ------------------------------------------------------------------------------------------
# * Remember: *both* `p` (for parameters) and `t` (for time) must be included in the
# equations of motion function, irrespective of whether they are used or not!
#
# #### 2. State
#
# The second argument for the `DiscreteDynamicalSystem` constructor is a `state`, which
# represents the initial condition for the system:
# ------------------------------------------------------------------------------------------

state = zeros(2) # doesn't matter if we use Vector of SVector for the `state`

# ------------------------------------------------------------------------------------------
# #### 3. Parameters for our EOM
#
# Then, the last argument, `p`, is simply a parameter container.
# ------------------------------------------------------------------------------------------

p = [1.4, 0.3] # p = [a, b] from the equations of motion

# ------------------------------------------------------------------------------------------
# These 3 things are enough to make a `DynamicalSystem`:
# ------------------------------------------------------------------------------------------

henon = DiscreteDynamicalSystem(h_eom, state, p)

# ------------------------------------------------------------------------------------------
# ## Getting a trajectory
#
# Now that we have created the system, the first (and most basic) thing to do is to simply
# plot its time evolution and see what it looks like.
#
# The `trajectory` function is a convenient tool that evolves the system and returns the
# output at equally spaced time intervals. The call signature is simply:
#
# ```julia
# trajectory(ds::DynamicalSystem, T [, u]; kwargs...)
# ```
# which evolves a system for total time `T`, optionally starting from a different state `u`.
# ------------------------------------------------------------------------------------------

# trajectory from initial condition
tr = trajectory(henon, 100000)

# trajectory from a different starting point
tr2 = trajectory(henon, 100000, 0.01rand(2))

using PyPlot
figure(figsize=(6,4))
plot(tr[:, 1], tr[:, 2], lw = 0.0, marker = "o", ms = 0.1, alpha = 0.5);
plot(tr2[:, 1], tr2[:, 2], lw = 0.0, marker = "o", ms = 0.1, alpha = 0.5);
xlabel("x"); ylabel("y");

# ------------------------------------------------------------------------------------------
# ## Crash-course on `Dataset`
#
# `trajectory` returns an object that is a `Dataset`:
# ------------------------------------------------------------------------------------------

tr = trajectory(henon, 100000)
println(typeof(tr))
println(summary(tr))

# ------------------------------------------------------------------------------------------
# `Dataset` instances handle most data in **DynamicalSystems.jl**. `Dataset` is a wrapper of
# a `Vector` of `SVector`s (statically sized vectors).
# ------------------------------------------------------------------------------------------

typeof(tr.data) # this is where the `Dataset` has the data

# ------------------------------------------------------------------------------------------
# When accessed with one index, a `Dataset` behaves as a vector of vectors
# ------------------------------------------------------------------------------------------

tr[1] # variables at first time point

tr[1:5]

# ------------------------------------------------------------------------------------------
# When accessed with two indices, a `Dataset` behaves like a matrix with each column being
# the timeseries of each dynamic variable
# ------------------------------------------------------------------------------------------

tr[:, 1] # timeseries of first variable

tr[1:56, 2] # time points 1:56 of second variable

tr[12, 1] # value of first variable at 12th timepoint

tr[1:10, 1:2] # using two ranges returns a `Dataset`

# ------------------------------------------------------------------------------------------
# # Adding a Jacobian
#
# Let's see `henon` again:
# ------------------------------------------------------------------------------------------

henon

# ------------------------------------------------------------------------------------------
# * The last line, "`jacobian:  ForwardDiff`"
# says that the Jacobian function of the equations of motion was computed automatically
# using the module
# [`ForwardDiff`](http://www.juliadiff.org/ForwardDiff.jl/stable/user/api.html).
# * The Jacobian function is a crucial component of a dynamical system, and that is why if
# it is not given, it is computed automatically.
#
#
#
# Even though the automatic computation is very efficient, the best possible performance
# will come if you pass a "hard-coded" jacobian:
# ------------------------------------------------------------------------------------------

h_jacobian(x, p, t) = @SMatrix [-2*p[1]*x[1] 1.0; p[2] 0.0]

# ------------------------------------------------------------------------------------------
# * Notice that for **out-of-place** systems, the Jacobian must also have the same form as
# the equations of motion, namely to return an `SMatrix`.
#
# Now, we can pass this Jacobian function to the `DiscreteDynamicalSystem` constructor as a
# 4th argument:
# ------------------------------------------------------------------------------------------

henon_with_jac = DiscreteDynamicalSystem(h_eom, state, p, h_jacobian)

# ------------------------------------------------------------------------------------------
# To see the difference in performance, let's call this Jacobian function
# ------------------------------------------------------------------------------------------

using BenchmarkTools
a = rand(SVector{2})
@btime $(henon_with_jac.jacobian)($a, $henon.prob.p, 0);
@btime $(henon.jacobian)($a, $henon.prob.p, 0);

# ------------------------------------------------------------------------------------------
# You can see that even though `ForwardDiff` is truly performant, the hard-coded version is
# much faster.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Continuous System
# * The process of creating a continuous system is identical to that of a discrete system,
# except that the constructor `ContinuousDynamicalSystem` is used instead.
# * In this case the function `eom` returns the time derivatives and not a "next state".
#
#
# ---
#
#
#
# * We will take the opportunity to show the process of using in-place equations of motion
# for a continuous system, which is aimed to be used for large systems (dimensionality $\ge$
# 10).
#
# * In addition, the system we will use (Henon-Heiles) does not have any parameters.
# $$
# \begin{aligned}
# \dot{x} &= p_x \\
# \dot{y} &= p_y \\
# \dot{p}_x &= -x -2 xy \\
# \dot{p}_y &= -y - (x^2 - y^2)
# \end{aligned}
# $$
# ------------------------------------------------------------------------------------------

# Henon-heiles famous system
# in-place form of equations of motion
# du is the derivatives vector, u is the state vector
function hheom!(du, u, p, t)
    du[1] = u[3]
    du[2] = u[4]
    du[3] = -u[1] - 2u[1]*u[2]
    du[4] = -u[2] - (u[1]^2 - u[2]^2)
    return nothing
end

# pass `nothing` as the parameters, because the system doesn't have any
hh = ContinuousDynamicalSystem(hheom!, [0, -0.25, 0.42081, 0], nothing)

# ------------------------------------------------------------------------------------------
# Great, now we can get a trajectory of this system, by solving using
# DifferentialEquations.jl
# ------------------------------------------------------------------------------------------

tr = trajectory(hh, 100.0, dt = 0.05)

figure(figsize = (6,4))
plot(tr[:, 1], tr[:, 2]);
xlabel("\$q_1\$"); ylabel("\$q_2\$");

# ------------------------------------------------------------------------------------------
# # Handy Dandy definition table
#
# Depending on whether your system is small or large, you want to use out-of-place or in-
# place equations of motion. The Jacobian function (if you wish to provide it) must also be
# of the same form.
#
# Here is a handy table that summarizes what the definitions should look like:
#
# |          System Type         |    equations of motion    |            Jacobian
# |
# |:----------------------------:|:-------------------------:|:-----------------------------
# -:|
# | in-place (big systems)       | `eom!(du, u, p, t)`       | `jacobian!(J, u, p, t)`
# |
# | out-of-place (small systems) | `eom(u, p, t) -> SVector` | `jacobian(u, p, t) ->
# SMatrix` |
#
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Time-evolution of systems & DifferentialEquations.jl
# Discrete systems are evolved using internal algorithms. However, all time evolution of
# continuous systems is done through the
# [DifferentialEquations.jl](http://docs.juliadiffeq.org/latest/) library. In fact,
# `trajectory` for continuous systems simply wraps `solve` with some extra arguments.
#
# Keep in mind that by default all continuous systems are evolved using the solver `Vern9()`
# (9th order Verner solver) with tolerances `abstol = reltol = 1e-9`. This is especially
# important when one wants to compare benchmarks with different packages.
#
# It is almost certain that if you use **DynamicalSystems.jl** you want to use also
# DifferentialEquations.jl, due to the huge list of available features.
#
#
# ## When to use **DynamicalSystems.jl**?
# At this point in the tutorial you might be thinking:
# > How is DynamicalSystems.jl different from DifferentialEquations.jl? Seems the same to me
# so far...
#
# That's because we haven't seen any special features of **DynamicalSystems.jl** yet!
# `trajectory` is nothing more than a "convenient shortcut".
#
# The [contents page](https://juliadynamics.github.io/DynamicalSystems.jl/latest/#contents)
# of the documentation does a good job describing what is possible through
# **DynamicalSystems.jl**.
#
# In addition, for usage with DifferentialEquations.jl, we can create specialized
# integrators that evolve the system and the tangent space, or many states of the system in
# parallel (at *exactly* same times):
# * `tangent_integrator`
# * `parallel_integrator`
#
# These functions work for both continuous and discrete systems and also work regardless of
# whether the system is in-place (large), out-of-place (small) or auto-differentiated.
# Special attention has also been given to the performance of the integrators.
#
# Also, keep in mind that regardless of whether you use a `DynamicalSystem` to produce a
# timeseries or not, more than half of the **DynamicalSystems.jl** library is using
# numerical data as an input. Most of the time this numerical data is expected in the form
# of an `AbstractDataset` instance, which is what `trajectory` returns.
#
# ## Using DifferentialEquations.jl from a `DynamicalSystem`
#
# It is *very* likely that other features of DifferentialEquations.jl will be useful to
# someone using **DynamicalSystems.jl**. However, you can still use DifferentialEquations.jl
# *after* you have defined a continuous dynamical system, because the field `prob` gives an
# `ODEProblem`:
# ------------------------------------------------------------------------------------------

hh.prob

# ------------------------------------------------------------------------------------------
# * **please be careful when using this problem directly, because as you can see the
# `tspan[end]` field is `Inf`!!!**
#
#
# * One final comment: using Callbacks is *not* possible with **DynamicalSystems.jl**,
# because the the equations of motion function has to be assumed differentiable
# "everywhere".
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# # Docstrings
# ------------------------------------------------------------------------------------------

?DynamicalSystem

?trajectory

?Dataset

.nbexports/introductory-tutorials/broader-topics-and-ecosystem/introduction-to-dynamicalsystems.jl/1. Creating a dynamical system.jl


"""
    bernoulli(n, x)

 Calculates Bernoulli polynomials ``B_n(x)``
 e.g., see

 + [https://en.wikipedia.org/wiki/Bernoulli_polynomials](https://en.wikipedia.org/wiki/Bernoulli_polynomials)
 + [http://dlmf.nist.gov/24](http://dlmf.nist.gov/24)

## Arguments
* ``n`` `::Integer`: the index into the series, ``n=0,1,2,3,...``
* ``x`` `::Real`: the point at which to calculate the polynomial

## Examples
```jldoctest; setup = :(using Polylogarithms)
julia> bernoulli(6, 1.2)
0.008833523809524735
```
"""
function bernoulli(n::Integer, x::Real)
    if n<0
        throw(DomainError(n))
    end
    if n == 0
        return 1 # zeta formula doesn't hold for n=0, so return explicit value
    elseif n == 1 # get some easy cases out of the way quickly
        return x-0.5
    elseif n == 2
        return x^2 - x + 1.0/6.0
    elseif n == 3
        return x^3 - 1.5*x^2 + 0.5*x
    elseif n == 4
        return x^4 - 2.0*x^3 +     x^2 - 1/30.0
    elseif n == 5
         return x^5 - 2.5*x^4 +(5.0/3.0)*x^3 - x/6.0
    end

 # from the Hurwitz-zeta function using

 # ``ζ(-n,x) = -B_{n+1}(x)/(n+1), for Re(x)>0,``

 # which is faster than direct calculation of the polynomial except for n<=4, but for
 # small n (n<=5) we use the exact polynomials. For negative x, we use the recursive
 # formula to push it into a postive range.
     # return -n*SpecialFunctions.zeta(1-n, x)
    # see https://carma.newcastle.edu.au/resources/jon/Preprints/Papers/Published-InPress/Oscillatory%20(Tapas%20II)/Papers/coffey-zeta.pdf, p.341
    # however, I am not really happy with the results I am getting out of this
    # so exploit symmetries in B_n(x) to compute recursively for x<=0
    # if x >= 0
    #      return -n*SpecialFunctions.zeta(1-n, x)
    # else
    #     return bernoulli(n, x+1) - n*x^(n-1)
    # end

    # total = 0.0
    # for k = 0:n-1
    #     total -= binomial(n+1, k) *  bernoulli(k,x)
    # end
    # total = total / (n+1)
    # total += x^n
    # return total

    # alternative, a little better, but relies on bernoulli numbers, which I only have to 35
    # and still has some non-negligable errors, so I still would like a better piece in here
    n = Int64(n)
    if n <= 35
        total = 0.0
        for k = 0:n
            k = Int64(k)
            total += binomial(n, k) *  bernoulli(n-k) * x^k
            # this will overflow for large n and k, but that migth be avoidable if we were more clever?
        end
        return total
    else
        # this is faster but less accurate, and runs into some weird problems, e.g.,
        # SpecialFunctions.zeta(-6, -1.0) = NaN (probably should be 2)
        return -n*SpecialFunctions.zeta(1-n, x)
    end
end


module paqueteFourier

import Base: +, -, *, /

import Gadfly

# Notar que los métodos de +, -, *, / y diff se exportan solos
export trigPoly, innerProduct, nodes, vals, tpfplot

# El objeto básico de este paquete es un trigPoly, cuyo campo coefs contiene
# los coeficientes de la representación de un polinomio trigonométrico como
# combinación lineal de 1, exp(-im θ), exp(im θ), exp(-2im θ), exp(2im θ), …,
# en ese orden.
# Junto a la definición de este 'type' Julia automáticamente deja disponible un
# constructor, que es una función que recibe un vector de números complejos y
# devuelve un trigPoly con ese vector en el campo coefs.
# A continuación justificaré mi elección de ordenamiento para los coeficientes
# con la ayuda de una tabla.  Denoto al orden elegido (0, -1, 1, -2, 2, …) por
# A, al orden visto en clase (-N/2, -N/2+1, …, -1, 0, 1, …, N/2-2, N/2) por B y
# al orden por defecto de la FFT (0, 1, …, N/2-1, -N/2, -N/2+1, …, -1) por C (B
# y C son ligeramente distintos si N es impar, pero la tabla sigue valiendo).
#                                 ┌───┬───┬───┐
#                                 │ A │ B │ C │
# ┌───────────────────────────────┼───┼───┼───┤
# │                          Suma │ ✓ │ ✗ │ ✗ │
# │Interpretación de coeficientes │ ✓ │ ✗ │ ~ │
# │     Vector de números de onda │ ~ │ ✓ │ ✗ │
# │           Interpolación (FFT) │ ✗ │ ~ │ ✓ │
# ├───────────────────────────────┴───┴───┴───┤
# │   ✓ = Fácil, ~ = Moderado, ✗ = Difícil    │
# └───────────────────────────────────────────┘
"""
    trigPoly(x)

Polinomio trigonométrico que es combinación lineal de 1, exp(-im θ), exp(im θ), exp(-2im θ), … con los coeficientes indicados en `x`.
"""
type trigPoly
	coefs::Vector{Complex128}
end

"""
    innerProduct(f,g)

Producto interior L² entre los trigPoly `f` y `g`.
"""
function innerProduct(f::trigPoly, g::trigPoly)
	ac = 0.0 + 0.0im
	for k = 1:min(length(f.coefs),length(g.coefs))
		ac += f.coefs[k] * conj(g.coefs[k])
	end
	2π*ac
end

# Sin cargar ningún paquete Julia ya tiene una función llamada norm; la nuestra
# será un método más
"""
    norm(f::trigPoly)

Norma L² de `f`.
"""
Base.norm(f::trigPoly) = sqrt(real(innerProduct(f,f)))

"""
    nodes(n)

Arreglo de `n` nodos de Fourier.
"""
function nodes(n::Integer)
	2π*(0:n-1)/n
end

# Función que transforma un vector del orden C al orden A; Julia reemplazará T
# con el tipo de número que contenga el vector con el cual se alimente a esta
# función
function reordering{T}(v::Vector{T})
	n = length(v)
	out = Vector{T}(n) # Reserva de memoria
	floorhalfn = n÷2
	for k = 1:floorhalfn
		out[2*k] = v[n+1-k]
	end
	for k = 1:n-floorhalfn
		out[2*k-1] = v[k]
	end
	out
end

# Función que transforma un vector del orden A al orden C
function reverse_reordering{T}(v::Vector{T})
	n = length(v)
	floorhalfn = n÷2
	out = Vector{T}(n)
	for k = 1:floorhalfn
		out[n+1-k] = v[2*k]
	end
	for k = 1:n-floorhalfn
		out[k] = v[2*k-1]
	end
	out
end

# Interpolador que asume que los valores recibidos corresponden a la evaluación
# de una función en los nodos de Fourier
function interp{T}(v::Vector{T})
	modes = fft(v)/length(v)
	trigPoly(reordering(modes))
end

"""
    trigPoly(f, n)

Polinomio trigonométrico que interpola a la función `f` en `n` nodos de Fourier.
"""
function trigPoly(f::Function, n::Integer)
	vals = map(f, nodes(n)) # Vector con evaluaciones de f
	interp(vals)
end

"""
    vals(f)

Evaluación del polinomio trigonométrico `f` en tantos nodos de Fourier como coeficientes almacena `f`.
"""
function vals(f::trigPoly)
	ifft(reverse_reordering(f.coefs*length(f.coefs)))
end

"""
    vals(f, x)

Evaluación del polinomio trigonométrico `f` en el arreglo o escalar real `x`.
"""
function vals(f::trigPoly, x)
	coefsplus = f.coefs[3:2:end]
	coefsminus = f.coefs[2:2:end]
	# Podría escribir exp(im*x), pero si x es escalar Julia emite una
	# advertencia (no un error) que invita a usar la versión vectorizada de
	# exp; usando map x puede ser escalar o vectorial indistintamente
	expplusix = map(exp, im*x)
	expminusix = conj(expplusix)
	# Arreglo con ceros o escalar cero según lo que sea x
	outplus = zero(expplusix)
	outminus = zero(expplusix)
	for val in reverse(coefsplus)
		outplus = expplusix.*(outplus + val)
	end
	for val in reverse(coefsminus)
		outminus = expminusix.*(outminus + val)
	end
	outplus + outminus + f.coefs[1]
end

# Evaluador por defecto de un trigPoly; permite evaluar un trigPoly f en x con
# la sintaxis clásica f(x)
(f::trigPoly)(x) = vals(f, x)

# Suma
function +(f::trigPoly, g::trigPoly)
	fl = length(f.coefs)
	gl = length(g.coefs)
	retl = max(fl, gl)
	retcoefs = zeros(Complex128, retl)
	retcoefs[1:fl] = f.coefs
	retcoefs[1:gl] += g.coefs
	trigPoly(retcoefs)
end

# Cambio de signo
-(f::trigPoly) = trigPoly(-f.coefs)

# Resta
-(f::trigPoly, g::trigPoly) = f + (-g)

# Suma, resta, multiplicación y división por escalar
function +(f::trigPoly, a::Number)
	retcoefs = f.coefs
	retcoefs[1] += a
	trigPoly(retcoefs)
end
-(f::trigPoly, a::Number) = f + (-a)
*(f::trigPoly, a::Number) = trigPoly(f.coefs*a)
/(f::trigPoly, a::Number) = trigPoly(f.coefs/a)
*(a::Number, f::trigPoly) = f*a
+(a::Number, f::trigPoly) = f+a
-(a::Number, f::trigPoly) = (-f) + a

# Multiplicación
function *(f::trigPoly, g::trigPoly)
	ndof = length(f.coefs) + length(g.coefs)
	paddedf = trigPoly([f.coefs; zeros(Complex128,ndof-length(f.coefs))])
	paddedg = trigPoly([g.coefs; zeros(Complex128,ndof-length(g.coefs))])
	productvals = vals(paddedf) .* vals(paddedg)
	interp(productvals)
end

"""
    diff(f::trigPoly)

Derivada de `f`.
"""
function Base.diff(f::trigPoly)
	retcoefs = Vector{Complex128}(length(f.coefs))
	for k = 1:length(f.coefs)
		#             ↓     modo     ↓
		retcoefs[k] = (-1)^(k+1)*(k÷2) * im * f.coefs[k]
	end
	trigPoly(retcoefs)
end

"""
    diff(f::trigPoly, k::Int64)

`k`-ésima derivada de `f`.
"""
function Base.diff(f::trigPoly, k::Int64)
	@assert k>=0
	if k == 0
		return f
	elseif k == 1
		return diff(f)
	else
		return diff(f, k-1)
	end
end

"""
    tpfplot(f)
    tpfplot(f1, f2, ...)
    tpfplot(f1, f2, ..., labels=[label1, label2, ...])

Grafica a un polinomio trigonométrico (trigPoly) o función (Function) `f` o a la colección de polinomios trigonométricos o funciones `f1`, `f2`, ...; el argumento opcional `labels`, de aparecer, debe especificar nombres para las curvas ingresadas como un arreglo de cadenas de tipo `String`.

La figura obtenida se puede guardar mediante la sintaxis

    fh = tpfplot(...)
    Gadfly.draw(Gadfly.SVG("myplot.svg", 15Gadfly.cm, 9Gadfly.cm), fh)
"""
function tpfplot(vf::Vararg{Union{Function,trigPoly}}; labels=[])
	defaultNamesFlag = isempty(labels)
	layersReal = []
	layersImag = []
	layerLabels = []
	layerColors = []
	for i = 1:length(vf)
		c = ["blue" "green" "red" "deepskyblue" "magenta" "gold" "lightcoral" "black"][(i-1)%8+1]
		if defaultNamesFlag
			push!(labels, "Curve $(i)")
		end
		if typeof(vf[i]) <: trigPoly
			nsamples = max(150, 6*length(vf[i].coefs))
			x = nodes(nsamples)
			vfix = vals(trigPoly([vf[i].coefs;zeros(Complex128,nsamples-length(vf[i].coefs))]))
			# Agrego el punto de abscisa 2π
			x = [x;2π]
			vfix = [vfix;vfix[1]]
			# Ahora los puntos con abscisas en la colección de
			# nodos de Fourier que corresponde a vf[i]
			xn = nodes(length(vf[i].coefs))
			vfin = vals(vf[i])
			push!(layersReal, Gadfly.layer(x=x, y=real(vfix), Gadfly.Geom.line, Gadfly.Theme(default_color=c)))
			push!(layersReal, Gadfly.layer(x=xn, y=real(vfin), Gadfly.Geom.point, Gadfly.Theme(default_color=c)))
			push!(layersImag, Gadfly.layer(x=x, y=imag(vfix), Gadfly.Geom.line, Gadfly.Theme(default_color=c)))
			push!(layersImag, Gadfly.layer(x=xn, y=imag(vfin), Gadfly.Geom.point, Gadfly.Theme(default_color=c)))
			push!(layerColors, c); push!(layerLabels, labels[i])
		else
			x = linspace(0,2π,2000)
			vfix = map(vf[i], x)
			push!(layersReal, Gadfly.layer(x=x, y=real(vfix), Gadfly.Geom.line, Gadfly.Theme(default_color=c)))
			push!(layersImag, Gadfly.layer(x=x, y=imag(vfix), Gadfly.Geom.line, Gadfly.Theme(default_color=c)))
			push!(layerColors, c); push!(layerLabels, labels[i])
		end
	end
	labelSpecification = Gadfly.Guide.manual_color_key("", layerLabels, layerColors)
	Gadfly.hstack(
	       Gadfly.plot(layersReal..., Gadfly.Guide.XLabel("θ"), Gadfly.Guide.YLabel(""), Gadfly.Guide.Title("Real part"), Gadfly.Coord.cartesian(xmin=0, xmax=2π), Gadfly.Theme(key_position=:bottom), labelSpecification),
	       Gadfly.plot(layersImag..., Gadfly.Guide.XLabel("θ"), Gadfly.Guide.YLabel(""), Gadfly.Guide.Title("Imaginary part"), Gadfly.Coord.cartesian(xmin=0, xmax=2π), Gadfly.Theme(key_position=:bottom), labelSpecification)
	       )
end

end # module


using Roots

f(x) = x^3 - x - 1
@show fzeros(f)
println()

@show fzero(f, [-2,2])
println()

@show fzero(f, 1.0)
println()

f(x) = exp(x) - x^4
fp(x) = exp(x) - 4x^3
fpp(x) = exp(x) - 12x^2
newton(f, fp, 8)        # 8.613169456441398
newton(f, 8)
halley(f, fp, fpp, 8)
halley(f, 8)
secant_method(f, 8, 8.5)


module TestFunctions

"""
One dimensional Runge function defined over \$[-1, 1]\$
 (Taken from chapter 13 of "Approximation Theory and Approximation
 Practice" by Prof. Trefethen)
"""
1d_runge(x) = 1./(1 + 25x.*x)

"""
One dimensional Keane test function
    (Taken from "Engineering Design with Surrogate Modelling" by Prof. Keane
    et al.)
"""
1d_keane(x) = (6x - 2).^2.*sin.(12x-4)

"""
One dimensional Liu test function
 (A survey of adaptive sampling for global metamodeling (2017))
"""
1d_liu(x) = x.*x.*exp.(-4x).*(sin.(7x) + cos.(14x))

"""
Two dimensional Branin function from Schonolau's thesis.
"""
2d_branin(x) = (x[2] .- 5.1x[1].*x[1]./(4π*π) .+ 5x[1]/π .- 6).^2 + 10(1 - 1/(8π)).*cos.(x[1]) + 10

export 1d_runge, 1d_keane, 1d_liu,
       2d_branin

end

# vim:filetype=julia:foldmethod=syntax


function euler_step(f, t::Number, x, len)
    return len' .* f(t, x)
end

function mod_euler_step(f, t::Number, x, len)
    whole_step = len' .* f(t, x)
    return (whole_step + len' .* f(t .+ len, x .+ whole_step))/2
end

function heum_step(f, t::Number, x, len)
    whole_step = len' .* f(t, x)
    return (whole_step + 3*len' .* f(t .+ 2/3*len, x .+ 2/3*whole_step))/4
end

function RK4_step(f, t::Number, x, len)
    k1 = len' .* f(t, x)
    k2 = len' .* f(t .+ len/2, x .+ k1/2)
    k3 = len' .* f(t .+ len/2, x .+ k2/2)
    k4 = len' .* f(t .+ len, x .+ k3)
    return (k1 + 2*k2 + 2*k3 + k4)/6
end

macro define_order(f, ord::Integer)
    @eval function $f(s::Symbol)
        if s == :order
            return $ord
        end
        throw(ErrorException("Invalid symbol. Use :order"))
    end
end

@define_order(euler_step, 1)
@define_order(mod_euler_step, 2)
@define_order(heum_step, 2)
@define_order(RK4_step, 4)

function fixed_step_ode_solver(
    f,
    t0::Number,
    x0;
    method=RK4_step,
    step::Number = 1e-3,
    iter::Integer = 10000
)
    t = t0
    x = Array{typeof(x0 + step*f(t0, x0))}(undef, iter+1)
    x[1] = x0
    for j = 1:iter
        x[j+1] = x[j] + method(f, t, x[j], step)
        t += step
    end
    return x
end

function fixed_step_ode_solver(
    f,
    t0::Number,
    x0,
    stop_cond;
    method=RK4_step,
    step::Number = 1e-3,
    maxit::Integer = 10000
)
    t = t0
    x = Array{typeof(x0 + method(f, t, x0, step)), 1}()
    push!(x, x0)
    for j = 1:maxit
        push!(x, last(x) + method(f, t, last(x), step))
        t += step
        if stop_cond(t, x)
            break
        end
    end
    return x
end


module CoDa

export alpha, alphaTransform, Helmert
export ilr, clr

"""
    Helmert(n)

Calculate ``n×n`` Helmert matrix
"""
function Helmert(n)
    H = zeros(n,n)
  for i = 1:n, j = 1:n
      if i==1
        H[i,j] = 1.0/sqrt(n)
      elseif j < i
        H[i,j] = 1.0/sqrt(i*(i-1))
      elseif i==j
        H[i,j] = (1.0-i)/sqrt(i*(i-1))
      end
  end
    return -H
end

"""
    alphaTransform(X::Matrix, α::Float64)

Apply alpha-transform with Helmert product to ``X``
"""
function alphaTransform(X::Matrix{Float64}, α::Float64)
    D = size(X,2)
    H = Helmert(D)[2:end,:]
    Xᵃ = X.^α
    Xs = sum(Xᵃ, dims=2)
    U = (D*Xᵃ./Xs .- 1)./α
    return U*transpose(H)
end

"""
    alpha(X::Matrix, α::Float64)

Apply alpha-transform without Helmert product
"""
function alpha(X::Matrix{Float64}, α::Float64)
    D = size(X,2)
#    H = Helmert(D)[2:end,:]
    Xᵃ = X.^α
    Xs = sum(Xᵃ,dims=2)
    U = (D*Xᵃ./Xs .- 1)./α
    return U
end

"""
    ilr(X::Matrix)

Isometric logratio transform
"""
function ilr(X::Matrix{Float64})
    D = size(X,2)
    H = Helmert(D)[2:end,:]
    return log.(X) * transpose(H)
end

"""
    clr(X::Matrix)

Centered logratio transform
"""
function clr(X::Matrix{Float64})
    D = size(X,2)
    ratio = X ./ (prod(X,2).^(1/D))
    return log.(ratio)
end

end


# Linear Discriminant Analysis

#### Type to represent a linear discriminant functional

"""
A linear discriminant functional can be written as

```math
    f(\\mathbf{x}) = \\mathbf{w}^T \\mathbf{x} + b
```

Here, ``w`` is the coefficient vector, and ``b`` is the bias constant.
"""
struct LinearDiscriminant{T<:Real} <: RegressionModel
    w::Vector{T}
    b::T
end

#### function to solve linear discriminant
"""
    ldacov(C, μp, μn)

Performs LDA given a covariance matrix `C` and both mean vectors `μp` & `μn`.  Returns a linear discriminant functional of type [`LinearDiscriminant`](@ref).

*Parameters*
- `C`: The pooled covariane matrix (*i.e* ``(Cp + Cn)/2``)
- `μp`: The mean vector of the positive class.
- `μn`: The mean vector of the negative class.
"""
function ldacov(C::DenseMatrix{T},
                μp::DenseVector{T},
                μn::DenseVector{T}) where T<:Real

    w = cholesky(C) \ (μp - μn)
    ap = w ⋅ μp
    an = w ⋅ μn
    c = 2 / (ap - an)
    LinearDiscriminant(rmul!(w, c), 1 - c * ap)
end

"""
    ldacov(Cp, Cn, μp, μn)

Performs LDA given covariances and mean vectors. Returns a linear discriminant functional of type [`LinearDiscriminant`](@ref).

*Parameters*
- `Cp`: The covariance matrix of the positive class.
- `Cn`: The covariance matrix of the negative class.
- `μp`: The mean vector of the positive class.
- `μn`: The mean vector of the negative class.

**Note:** The coefficient vector is scaled such that ``w'μp + b = 1`` and ``w'μn + b = -1``.
"""
ldacov(Cp::DenseMatrix{T},
       Cn::DenseMatrix{T},
       μp::DenseVector{T},
       μn::DenseVector{T}) where T<:Real = ldacov(Cp + Cn, μp, μn)

"""
    evaluate(f, x::AbstractVector)

Evaluate the linear discriminant value, *i.e* ``w'x + b``, it returns a real value.
"""
evaluate(f::LinearDiscriminant, x::AbstractVector) = dot(f.w, x) + f.b

"""
    evaluate(f, X::AbstractMatrix)

Evaluate the linear discriminant value, *i.e* ``w'x + b``, for each sample in columns of `X`. The function returns a vector of length `size(X, 2)`.
"""
function evaluate(f::LinearDiscriminant, X::AbstractMatrix)
    R = transpose(X) * f.w
    if f.b != 0
        broadcast!(+, R, R, f.b)
    end
    return R
end

# RegressionModel interface

"""
    predict(f, x::AbstractVector)

Make prediction for the vector `x`. It returns `true` iff `evaluate(f, x)` is positive.
"""
predict(f::LinearDiscriminant, x::AbstractVector) = evaluate(f, x) > 0

"""
    predict(f, X::AbstractMatrix)

Make predictions for the matrix `X`.
"""
predict(f::LinearDiscriminant, X::AbstractMatrix) = Bool[y > 0 for y in evaluate(f, X)]

"""
    coef(f::LinearDiscriminant)

Return the coefficients of the linear discriminant model.
"""
coef(f::LinearDiscriminant) = (f.b, f.w)

"""
    coef(f::LinearDiscriminant)

Return the coefficients' names of the linear discriminant model.
"""
coefnames(f::LinearDiscriminant) = ["Bias", "Weights"]

"""
    dof(f::LinearDiscriminant)

Return the number of degrees of freedom in the linear discriminant model.
"""
dof(f::LinearDiscriminant) = length(f.w)+1

"""
    weights(f::LinearDiscriminant)

Return the linear discriminant model coefficient vector.
"""
weights(f::LinearDiscriminant) = f.w

"""
Get the length of the coefficient vector.
"""
length(f::LinearDiscriminant) = length(f.w)

"""
    fit(LinearDiscriminant, Xp, Xn; covestimator = SimpleCovariance())

Performs LDA given both positive and negative samples. The function accepts follwing parameters:

**Parameters**
- `Xp`: The sample matrix of the positive class.
- `Xn`: The sample matrix of the negative class.

**Keyword arguments:**
- `covestimator`: Custom covariance estimator for between-class covariance. The covariance matrix will be calculated as `cov(covestimator_between, #=data=#; dims=2, mean=zeros(#=...=#)`. Custom covariance estimators, available in other packages, may result in more robust discriminants for data with more features than observations.
"""
function fit(::Type{LinearDiscriminant}, Xp::DenseMatrix{T}, Xn::DenseMatrix{T};
             covestimator::CovarianceEstimator = SimpleCovariance()) where T<:Real
    μp = vec(mean(Xp, dims=2))
    μn = vec(mean(Xn, dims=2))
    Zp = Xp .- μp
    Zn = Xn .- μn
    Cp = calcscattermat(covestimator, Zp)
    Cn = calcscattermat(covestimator, Zn)
    ldacov(Cp, Cn, μp, μn)
end

#==============================================================================#

#### Multiclass LDA Stats

"""
Resulting statistics of the multi-class LDA evaluation.
"""
mutable struct MulticlassLDAStats{T<:Real, M<:AbstractMatrix{T}, N<:AbstractMatrix{T}}
    dim::Int              # sample dimensions
    nclasses::Int         # number of classes
    cweights::Vector{T}   # class weights
    tweight::T            # total sample weight
    mean::Vector{T}       # overall sample mean
    cmeans::Matrix{T}     # class-specific means
    Sw::M                 # within-class scatter matrix
    Sb::N                 # between-class scatter matrix
end

mean(S::MulticlassLDAStats) = S.mean
classweights(S::MulticlassLDAStats) = S.cweights
classmeans(S::MulticlassLDAStats) = S.cmeans

withclass_scatter(S::MulticlassLDAStats) = S.Sw
betweenclass_scatter(S::MulticlassLDAStats) = S.Sb

function MulticlassLDAStats(cweights::Vector{T},
                            mean::Vector{T},
                            cmeans::Matrix{T},
                            Sw::AbstractMatrix{T},
                            Sb::AbstractMatrix{T}) where T<:Real
    d, nc = size(cmeans)
    length(mean) == d || throw(DimensionMismatch("Incorrect length of mean"))
    length(cweights) == nc || throw(DimensionMismatch("Incorrect length of cweights"))
    tw = sum(cweights)
    size(Sw) == (d, d) || throw(DimensionMismatch("Incorrect size of Sw"))
    size(Sb) == (d, d) || throw(DimensionMismatch("Incorrect size of Sb"))
    MulticlassLDAStats(d, nc, cweights, tw, mean, cmeans, Sw, Sb)
end

function multiclass_lda_stats(nc::Int, X::AbstractMatrix{T}, y::AbstractVector{Int};
                              covestimator_within::CovarianceEstimator=SimpleCovariance(),
                              covestimator_between::CovarianceEstimator=SimpleCovariance()) where T<:Real
    # check sizes
    d = size(X, 1)
    n = size(X, 2)
    n ≥ nc || throw(ArgumentError("The number of samples is less than the number of classes"))
    length(y) == n || throw(DimensionMismatch("Inconsistent array sizes."))

    # compute class-specific weights and means
    cmeans, cweights, Z = center(X, y, nc)

    Sw = calcscattermat(covestimator_within, Z)

    # compute between-class scattering
    mean = cmeans * Vector{T}(cweights ./ n)
    U = rmul!(cmeans .- mean, Diagonal(sqrt.(cweights)))
    Sb = calcscattermat(covestimator_between, U)

    return MulticlassLDAStats(Vector{T}(cweights), mean, cmeans, Sw, Sb)
end

#### Multiclass LDA

"""
A multi-class linear discriminant model type has following fields:
- `proj` is the projection matrix
- `pmeans` is the projected means of all classes
- `stats` is an instance of [`MulticlassLDAStats`](@ref) type that captures all statistics computed to train the model (which we will discuss later).
"""
mutable struct MulticlassLDA{T<:Real} <: RegressionModel
    proj::Matrix{T}
    pmeans::Matrix{T}
    stats::MulticlassLDAStats{T}
end

"""
    size(M::MulticlassLDA)

Get the input (*i.e* the dimension of the observation space) and output (*i.e* the dimension of the transformed features) dimensions of the model `M`.
"""
size(M::MulticlassLDA) = size(M.proj)
"""
    length(M::MulticlassLDA)

Get the sample dimensions.
"""
length(M::MulticlassLDA) = M.stats.dim
"""
    projection(M::MulticlassLDA)

Get the projection matrix (of size *d x p*).
"""
projection(M::MulticlassLDA) = M.proj
"""
    mean(M::MulticlassLDA)

Get the overall sample mean vector (of length *d*).
"""
mean(M::MulticlassLDA) = mean(M.stats)

"""
    classmeans(M)

Get the matrix comprised of class-specific means as columns (of size ``(d, m)``).
"""
classmeans(M::MulticlassLDA) = classmeans(M.stats)
"""
    classweights(M)

Get the weights of individual classes (a vector of length ``m``). If the samples are not weighted,
the weight equals the number of samples of each class.
"""
classweights(M::MulticlassLDA) = classweights(M.stats)
"""
withinclass_scatter(M)

    Get the within-class scatter matrix (of size ``(d, d)``).
"""
withclass_scatter(M::MulticlassLDA) = withclass_scatter(M.stats)
"""
betweenclass_scatter(M)

    Get the between-class scatter matrix (of size ``(d, d)``).
"""
betweenclass_scatter(M::MulticlassLDA) = betweenclass_scatter(M.stats)

"""
    predict(M::MulticlassLDA, x)

Transform input sample(s) in `x` to the output space of MC-LDA model `M`. Here, `x` can be either a sample vector or a matrix comprised of samples in columns.
"""
predict(M::MulticlassLDA, x::AbstractVecOrMat{T}) where {T<:Real} = M.proj'x

"""
    fit(MulticlassLDA, nc, X, y; ...)

Perform multi-class LDA over a given data set `X` and collecttion of labels `y`.

This function returns the resultant multi-class LDA model as an instance of [`MulticlassLDA`](@ref).

*Parameters*

- `nc`:  the number of classes
- `X`:   the matrix of input samples, of size `(d, n)`. Each column in `X` is an observation.
- `y`:   the vector of class labels, of length `n`. Each element of `y` must be an integer between `1` and `nc`.

**Keyword arguments**

- `method`: The choice of methods:
    - `:gevd`: based on generalized eigenvalue decomposition (*default*).
    - `:whiten`: first derive a whitening transform from `Sw` and then solve the problem based on eigenvalue
    decomposition of the whiten `Sb`.
- `outdim`: The output dimension, i.e. dimension of the transformed space `min(d, nc-1)`
- `regcoef`: The regularization coefficient (*default:* `1.0e-6`). A positive value `regcoef * eigmax(Sw)`
    is added to the diagonal of `Sw` to improve numerical stability.
- `covestimator_between`: Custom covariance estimator for between-class covariance (*default:* `SimpleCovariance()`).
    The covariance matrix will be calculated as `cov(covestimator_between, #=data=#; dims=2, mean=zeros(#=...=#))`.
    Custom covariance estimators, available in other packages, may result in more robust discriminants for data
    with more features than observations.
- `covestimator_within`:  Custom covariance estimator for within-class covariance (*default:* `SimpleCovariance()`).
    The covariance matrix will be calculated as `cov(covestimator_within, #=data=#; dims=2, mean=zeros(nc))`.
    Custom covariance estimators, available in other packages, may result in more robust discriminants for data
    with more features than observations.

**Notes:**

The resultant projection matrix ``P`` satisfies:
```math
\\mathbf{P}^T (\\mathbf{S}_w + \\kappa \\mathbf{I}) \\mathbf{P} = \\mathbf{I}
```
Here, ``\\kappa`` equals `regcoef * eigmax(Sw)`. The columns of ``P`` are arranged in descending order of
the corresponding generalized eigenvalues.

Note that [`MulticlassLDA`](@ref) does not currently support the normalized version using ``\\mathbf{S}_w^*`` and
``\\mathbf{S}_b^*`` (see [`SubspaceLDA`](@ref)).
"""
function fit(::Type{MulticlassLDA}, nc::Int, X::DenseMatrix{T}, y::AbstractVector{Int};
             method::Symbol=:gevd,
             outdim::Int=min(size(X,1), nc-1),
             regcoef::T=T(1.0e-6),
             covestimator_within::CovarianceEstimator=SimpleCovariance(),
             covestimator_between::CovarianceEstimator=SimpleCovariance()) where T<:Real

    multiclass_lda(multiclass_lda_stats(nc, X, y;
                                        covestimator_within=covestimator_within,
                                        covestimator_between=covestimator_between);
                   method=method,
                   regcoef=regcoef,
                   outdim=outdim)
end

function multiclass_lda(S::MulticlassLDAStats{T};
                        method::Symbol=:gevd,
                        outdim::Int=min(S.dim, S.nclasses-1),
                        regcoef::T=T(1.0e-6)) where T<:Real

    P = mclda_solve(S.Sb, S.Sw, method, outdim, regcoef)
    MulticlassLDA(P, P'S.cmeans, S)
end

mclda_solve(Sb::AbstractMatrix{T}, Sw::AbstractMatrix{T}, method::Symbol, p::Int, regcoef::T) where T<:Real =
    mclda_solve!(copy(Sb), copy(Sw), method, p, regcoef)

function mclda_solve!(Sb::AbstractMatrix{T},
                      Sw::AbstractMatrix{T},
                      method::Symbol, p::Int, regcoef::T) where T<:Real

    p <= size(Sb, 1) || throw(ArgumentError("p cannot exceed sample dimension."))

    if method == :gevd
        regularize_symmat!(Sw, regcoef)
        E = eigen!(Symmetric(Sb), Symmetric(Sw))
        ord = sortperm(E.values; rev=true)
        P = E.vectors[:, ord[1:p]]

    elseif method == :whiten
        W = _lda_whitening!(Sw, regcoef)
        wSb = transpose(W) * (Sb * W)
        Eb = eigen!(Symmetric(wSb))
        ord = sortperm(Eb.values; rev=true)
        P = W * Eb.vectors[:, ord[1:p]]

    else
        throw(ArgumentError("Invalid method name $(method)"))
    end
    return P::Matrix{T}
end

function _lda_whitening!(C::AbstractMatrix{T}, regcoef::T) where T<:Real
    n = size(C,1)
    E = eigen!(Symmetric(C))
    v = E.values
    a = regcoef * maximum(v)
    for i = 1:n
        @inbounds v[i] = 1.0 / sqrt(v[i] + a)
    end
    return rmul!(E.vectors,  Diagonal(v))
end

#### SubspaceLDA

"""Subspace LDA model type has following fields:
- `projw`: the projection matrix of the subspace spanned by the between-class scatter
- `projLDA`: the projection matrix of the subspace spanned by the within-class scatter
- `λ`: the projection eigenvalues
- `cmeans`: the class centroids
- `cweights`: the class weights
"""
struct SubspaceLDA{T<:Real} <: RegressionModel
    projw::Matrix{T}
    projLDA::Matrix{T}
    λ::Vector{T}
    cmeans::Matrix{T}
    cweights::Vector{Int}
end

"""
    size(M)

Get the input (*i.e* the dimension of the observation space) and output (*i.e* the dimension of the subspace projection)
dimensions of the model `M`.
"""
size(M::SubspaceLDA) = (size(M.projw,1), size(M.projLDA, 2))
"""
    length(M)

Get dimension of the LDA model.
"""
length(M::SubspaceLDA) = size(M.projLDA, 2)
"""
    predict(M::SubspaceLDA, x)

Transform input sample(s) in `x` to the output space of LDA model `M`.
Here, `x` can be either a sample vector or a matrix comprised of samples in columns.
"""
predict(M::SubspaceLDA, x::AbstractVecOrMat{T}) where {T<:Real} = M.projLDA' * (M.projw' * x)
"""
    projection(M)

Get the projection matrix.
"""
projection(M::SubspaceLDA) = M.projw * M.projLDA

"""
    mean(M::SubspaceLDA)

Returns the mean vector of the subspace LDA model `M`.
"""
mean(M::SubspaceLDA) = vec(sum(M.cmeans * Diagonal(M.cweights / sum(M.cweights)), dims=2))

"""
    eigvals(M::SubspaceLDA)

Get the eigenvalues of the subspace LDA model `M`.
"""
eigvals(M::SubspaceLDA) = M.λ

classmeans(M::SubspaceLDA) = M.cmeans
classweights(M::SubspaceLDA) = M.cweights
fit(::Type{F}, X::AbstractMatrix{T}, nc::Int, label::AbstractVector{Int}) where {T<:Real, F<:SubspaceLDA} =
    fit(F, X, label, nc)

"""
    fit(SubspaceLDA, X, labels; normalize=true)

Fit an subspace projection of LDA model using the equivalent of ``\\mathbf{S}_w^*`` and ``\\mathbf{S}_b^*```.

Note: Subspace LDA also supports the normalized version of LDA via the `normalize` keyword.
"""
function fit(::Type{F}, X::AbstractMatrix{T},
             label::AbstractVector{Int},
             nc=maximum(label);
             normalize::Bool=false) where {T<:Real, F<:SubspaceLDA}
    d, n = size(X, 1), size(X, 2)
    n ≥ nc || throw(ArgumentError("The number of samples is less than the number of classes"))
    length(label) == n || throw(DimensionMismatch("Inconsistent array sizes."))
    # Compute centroids, class weights, and deviation from centroids
    # Note Sb = Hb*Hb', Sw = Hw*Hw'
    cmeans, cweights, Hw = center(X, label, nc)
    dmeans = cmeans .- (normalize ? mean(cmeans, dims=2) : cmeans * (cweights / T(n)))
    Hb = normalize ? dmeans : dmeans * Diagonal(convert(Vector{T}, sqrt.(cweights)))
    if normalize
        Hw /= T(sqrt(n))
    end
    # Project to the subspace spanned by the within-class scatter
    # (essentially, PCA before LDA)
    Uw, Σw, _ = svd(Hw, full=false)
    keep = Σw .> sqrt(eps(T)) * maximum(Σw)
    projw = Uw[:,keep]
    pHb = projw' * Hb
    pHw = projw' * Hw
    λ, G = lda_gsvd(pHb, pHw, cweights)
    SubspaceLDA(projw, G, λ, cmeans, cweights)
end

# Reference: Howland & Park (2006), "Generalizing discriminant analysis
# using the generalized singular value decomposition", IEEE
# Trans. Patt. Anal. & Mach. Int., 26: 995-1006.
function lda_gsvd(Hb::AbstractMatrix{T}, Hw::AbstractMatrix{T}, cweights::AbstractVector{Int}) where T<:Real
    nc = length(cweights)
    K = vcat(Hb', Hw')
    P, R, Q = svd(K, full=false)
    keep = R .> sqrt(eps(T))*maximum(R)
    R = R[keep]
    Pk = P[1:nc, keep]
    U, ΣA, W = svd(Pk)
    ncnz = sum(cweights .> 0)
    G = Q[:,keep]*(Diagonal(1 ./ R) * W[:,1:ncnz-1])
    # Normalize
    Gw = G' * Hw
    nrm = Gw * Gw'
    G = G ./ reshape(sqrt.(diag(nrm)), 1, ncnz-1)
    # Also get the eigenvalues
    Gw = G' * Hw
    Gb = G' * Hb
    λ = diag(Gb * Gb')./diag(Gw * Gw')
    λ, G
end

function center(X::AbstractMatrix{T}, label::AbstractVector{Int}, nc=maximum(label)) where T<:Real
    d, n = size(X,1), size(X,2)
    # Calculate the class weights and means
    cmeans = zeros(T, d, nc)
    cweights = zeros(Int, nc)
    for j = 1:n
        k = label[j]
        for i = 1:d
            cmeans[i,k] += X[i,j]
        end
        cweights[k] += 1
    end
    for j = 1:nc
        cw = cweights[j]
        cw == 0 && continue
        for i = 1:d
            cmeans[i,j] /= cw
        end
    end
    # Compute differences from the means
    dX = Matrix{T}(undef, d, n)
    for j = 1:n
        k = label[j]
        for i = 1:d
            dX[i,j] = X[i,j] - cmeans[i,k]
        end
    end
    cmeans, cweights, dX
end


# Density utilities
# Author: Gabrio Rizzuti, rizzuti.gabrio@gatech.edu
# Date: August 2020

export sample_gaussmixture, negLogDensity_gaussmixture

function sample_gaussmixture(nsamples::Int64, mu::Array{Float32, 2}, sigma2::Array{Float32, 1})
    nmodes = length(sigma2)
    X = Array{Float32, 4}(undef, (1, 1, 2, nsamples))
    for i = 1:nsamples
        imode = rand(1:nmodes)
        X[1, 1, :, i] = mu[imode, :]+sqrt(sigma2[imode])*randn(Float32, 2)
    end
    return X
end

function negLogDensity_gaussmixture(X::Array{Float32, 4}, mu::Array{Float32, 2}, sigma2::Array{Float32, 1}; ε::Float32 = 1f-10)
    nmodes = length(sigma2)
    p = zeros(Float32, 1, 1, 1, size(X, 4))
    g = zeros(Float32, size(X))
    for i = 1:nmodes
        p_i = exp.(-((X[:, :, 1:1, :].-mu[i, 1]).^2f0+(X[:, :, 2:2, :].-mu[i, 2]).^2f0)/(2f0*sigma2[i]))/(nmodes*2f0*pi*sigma2[i])
        p += p_i
        g += -p_i.*(X.-reshape(mu[i, :], 1, 1, 2, 1))/sigma2[i]
    end
    return -log.(p.+ε), -g./(p.+ε)
end


"""
    Dᶜ(N, Δ)

Returns a discrete 1D derivative operator for taking the derivative of a face-centered field with `N+1` grid points and `Δ` grid spacing and producing a cell-centered field with `N` grid points.
"""
function Dᶜ(N, Δ)
    D = zeros(N, N+1)
    for k in 1:N
        D[k, k]   = -1.0
        D[k, k+1] =  1.0
    end
    D = 1/Δ * D
    return D
end

"""
    Dᶠ(N, Δ)

Returns a discrete 1D derivative operator for taking the derivative of a cell-centered field with `N` grid points and `Δ` grid spacing and producing a face-centered field with `N+1` grid points.
"""
function Dᶠ(N, Δ)
    D = zeros(N+1, N)
    for k in 2:N
        D[k, k-1] = -1.0
        D[k, k]   =  1.0
    end
    D = 1/Δ * D
    return D
end

function cell_to_cell_derivative(D, data)
    face_data = D * data
    cell_data = 0.5 .* (@view(face_data[1:end-1]) + @view(face_data[2:end]))
    return cell_data
end


"""
Check if a solution x is in the core of G. Namely,

    - ∑ x < v(N)
    - ∑ₛ x < v(S) ∀ S

"""
function isincore(x::Vector{Float64}, G::Game)
    if sum(x) != G.v(G.N) return false end

    for S in subsets(G.N)
        if sum(x[S]) < G.v(S) return false end
    end

    return true

end