# # Dihedral symmetry of the Robinson form

# Based on version from SumOfSquares by Benoît Legat
# **Adapted from**: Example 5.4 of [GP04]
#
# [GP04] Gatermann, Karin and Parrilo, Pablo A.
# *Symmetry groups, semidefinite programs, and sums of squares*.
# Journal of Pure and Applied Algebra 192.1-3 (2004): 95-128.

using LinearAlgebra
using SparseArrays

using SymbolicWedderburn
using SymbolicWedderburn.StarAlgebras
using DynamicPolynomials
using SumOfSquares
using SCS

OPTIMIZER = optimizer_with_attributes(
    SCS.Optimizer,
    "acceleration_lookback" => 10,
    "max_iters" => 3_000,
    "alpha" => 1.2,
    "eps" => 1e-6,
    "linear_solver" => SCS.DirectSolver,
)

@polyvar x y
const robinson_form =
    x^6 + y^6 - x^4 * y^2 - y^4 * x^2 - x^4 - y^4 - x^2 - y^2 + 3x^2 * y^2 + 1

# The Robinson form is invariant under the following action of the Dihedral group on monomials:
# The action of each element of the groups is to map the variables `x, y` to:
#
# | id | rotation | reflection |
# |----|----------|------------|
# | 0  | x, y     | y, x       |
# | 1  | -y, x    | -x, y      |
# | 2  | -x, -y   | -y, -x     |
# | 3  | y, -x    | x, -y      |

# implementation of the dihedral group:
include(joinpath(@__DIR__, "dihedral.jl"))
# definitions of general actions on polynomials:
include(joinpath(@__DIR__, "action_polynomials.jl"))

struct DihedralAction <: OnMonomials end

SymbolicWedderburn.coeff_type(::DihedralAction) = Float64
function SymbolicWedderburn.action(
    ::DihedralAction,
    el::DihedralElement,
    mono::AbstractMonomial,
)
    x, y = variables(mono)
    if iseven(el.reflection + el.id)
        var_x, var_y = x, y
    else
        var_x, var_y = y, x
    end
    sign_x = 1 <= el.id <= 2 ? -1 : 1
    sign_y = 2 <= el.id ? -1 : 1
    return mono([x, y] => [sign_x * var_x, sign_y * var_y])
end

G = DihedralGroup(4)
for g in G
    @assert SymbolicWedderburn.action(DihedralAction(), g, robinson_form) ==
            robinson_form
end

# A simple Gramm-matrix like formulation of SOS problem:
# * a single large pdd constraint and
# * one linear constraint for each monomial in the monomial basis of f.
no_symmetry = let f = robinson_form
    stats = @timed let m = Model(OPTIMIZER)
        @variable m t
        @objective m Max t
        @constraint m f - t in SOSCone()
        m
    end
    m = stats.value
    optimize!(m)
    @info (m,) termination_status(m) objective_value(m) creation_time =
        (stats.time) solve_time(m)

    @assert isapprox(value(m[:t]), -3825 / 4096, rtol = 1e-2)

    (
        status = termination_status(m),
        objective = objective_value(m),
        creation_t = stats.time,
        solve_t = solve_time(m),
    )
end

# Let's decompose basis into semi-simple summands and prepare the approximately
# reduced optimization problem. Since this is an isomorphism we reduce the
# number of variables from N² to n₁² + … + nₖ², where Σᵢ nᵢ = N. That is instead
# of one psd constraint of size N we have k constraints of sizes n₁,… nₖ.
# The number of linear constraints remains the same as before.
semisimple_dec = let f = robinson_form, G = DihedralGroup(4)
    basis = monomials([x, y], 0:(DynamicPolynomials.maxdegree(f)÷2))

    stats = @timed begin
        sa_basis, symmetry_adaptation_time, = @timed let
            ssimple_basis =
                SymbolicWedderburn.symmetry_adapted_basis(
                    Rational{Int},
                    G,
                    DihedralAction(),
                    basis,
                    semisimple = true,
                )
            sp = sparse.(ssimple_basis)
            droptol!.(sp, 1e-12)
        end

        m = SOSModel(OPTIMIZER)

        let basis = basis, R = sa_basis
            @variable m t
            @objective m Max t

            soses = @variable m [r in R] SOSPoly(
                FixedPolynomialBasis(r * basis),
            )
            @constraint m f - t == sum(soses)
        end
        m
    end
    m = stats.value
    optimize!(m)

    @assert isapprox(value(m[:t]), -3825 / 4096, rtol = 1e-2)

    @info (m,) termination_status(m) objective_value(m) solve_time(m) creation_time =
        (stats.time) symmetry_adaptation_time
    (
        status = termination_status(m),
        objective = objective_value(m),
        creation_t = stats.time,
        solve_t = solve_time(m),
    )
end

# Now given that f must be constant on orbits of monomials under the action of G
# let's exploit that to reduce the number of cnstraints (one for each orbit).
# Here we keep the single large psd constraint of size N×N.
include(joinpath(@__DIR__, "util.jl"))

orbit_dec = let f = robinson_form, G = DihedralGroup(4), T = Float64
    basis_monoms = monomials([x, y], 0:DynamicPolynomials.maxdegree(f))

    invariant_vs, symmetry_adaptation_time, = @timed let G = G
        tblG = SymbolicWedderburn.Characters.CharacterTable(Rational{Int}, G)
        invariant_vectors(
            tblG,
            DihedralAction(),
            StarAlgebras.Basis{UInt32}(basis_monoms),
        )
    end

    M =
        let basis_full = StarAlgebras.Basis{UInt32}(basis_monoms),
            basis_half =
                monomials([x, y], 0:(DynamicPolynomials.maxdegree(f)÷2))

            [basis_full[x*y] for x in basis_half, y in basis_half]
        end

    stats = @timed let m = JuMP.Model(OPTIMIZER)
        JuMP.@variable m t
        JuMP.@objective m Max t
        P = JuMP.@variable(m, [1:size(M, 1), 1:size(M, 1)], PSD)

        # preallocating
        M_orb = similar(M, T)

        C = DynamicPolynomials.coefficients(f - t, basis_monoms)

        for iv in invariant_vs
            c = dot(C, iv)

            # invariant_constraint! is defined locally in util.jl
            # it essentially amounts to averaging constraint matrices defined by M
            # with weights provided by iv and storing the result in (dense) matrix M_orb
            M_orb = invariant_constraint!(M_orb, M, iv)

            #=
            if !iszero(M_orb)
                monoms = basis_monoms[unique!(M[map(!iszero, M_orb)])]
                @info "Coefficient of the $monoms: $([DynamicPolynomials.coefficient(f-t, m) for m in monoms])" dot_C_v=c
                @show sparse(M_orb)
            else
                monoms = basis_monoms[unique!(M[SparseArrays.nonzeroinds(iv)])]
                @info "Coefficient of the $monoms: $([DynamicPolynomials.coefficient(f-t, m) for m in monoms])" dot_C_v=c
                @info "M_orb is zero!"
            end
            =#

            JuMP.@constraint m dot(M_orb, P) == c
        end
        m
    end

    m = stats.value
    optimize!(m)

    @assert isapprox(value(m[:t]), -3825 / 4096, rtol = 1e-2)

    @info (m,) termination_status(m) objective_value(m) solve_time(m) creation_time =
        (stats.time) symmetry_adaptation_time
    (
        status = termination_status(m),
        objective = objective_value(m),
        creation_t = stats.time,
        solve_t = solve_time(m),
    )
end

wedderburn_dec = let f = robinson_form, G = DihedralGroup(4), T = Float64
    basis_half = monomials([x, y], 0:(DynamicPolynomials.maxdegree(f)÷2))

    wedderburn, symmetry_adaptation_time, = @timed WedderburnDecomposition(
        T,
        G,
        DihedralAction(),
        monomials([x, y], 0:DynamicPolynomials.maxdegree(f)),
        basis_half,
    )

    M = [
        SymbolicWedderburn.basis(wedderburn)[x*y] for x in basis_half,
        y in basis_half
    ]

    stats = @timed let m = JuMP.Model(OPTIMIZER)
        JuMP.@variable m t
        JuMP.@objective m Max t
        psds = [
            JuMP.@variable(m, [1:d, 1:d] in PSDCone()) for
            d in size.(direct_summands(wedderburn), 1)
        ]

        # preallocating
        Mπs = zeros.(T, size.(psds))
        M_orb = similar(M, T)
        tmps = SymbolicWedderburn._tmps(wedderburn)

        C = DynamicPolynomials.coefficients(
            f - t,
            SymbolicWedderburn.basis(wedderburn),
        )

        for iv in invariant_vectors(wedderburn)
            c = dot(C, iv)

            # invariant_constraint! is defined locally in util.jl
            # it essentially amounts to averaging constraint matrices defined by M
            # with weights provided by iv and storing the result in (dense) matrix M_orb
            M_orb = invariant_constraint!(M_orb, M, iv)

            #=
            if !iszero(M_orb)
                monoms = SymbolicWedderburn.basis(wedderburn)[unique!(M[map(!iszero, M_orb)])]
                @info "Coefficient of the $monoms: $([DynamicPolynomials.coefficient(f-t, m) for m in monoms])" dot_C_v=c
                @show sparse(M_orb)
            else
                monoms = SymbolicWedderburn.basis(wedderburn)[unique!(M[SparseArrays.nonzeroinds(iv)])]
                @info "Coefficient of the $monoms: $([DynamicPolynomials.coefficient(f-t, m) for m in monoms])" dot_C_v=c
                @info "M_orb is zero!"
            end
            =#

            Mπs = SymbolicWedderburn.diagonalize!(Mπs, M_orb, wedderburn, tmps)

            JuMP.@constraint m sum(
                dot(Mπ, Pπ) for (Mπ, Pπ) in zip(Mπs, psds) if !iszero(Mπ)
            ) == c
        end
        m
    end
    m = stats.value
    optimize!(m)

    @info value(m[:t])
    @assert isapprox(value(m[:t]), -3825 / 4096, rtol = 1e-2)

    @info (m,) termination_status(m) objective_value(m) solve_time(m) creation_time =
        (stats.time)
    (
        status = termination_status(m),
        objective = objective_value(m),
        creation_t = stats.time,
        solve_t = solve_time(m),
    )
end


#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    nrbcirc(radius::F=1.0,center::Vector{F}=[0.0;0.0;0.0],
            sang::F=0.0,eang::F=2.0*π)::NURBS1D where {F<:AbstractFloat}

Constructs NURBS data structure for a circular arc in the x-y plane. If no
arguments are supplied a unit circle with center `[0.0;0.0]` is constructed.

Angles are defined as positive in the anti-clockwise direction.

# Arguments:
- `radius`: Radius of the circle, default `1.0`
- `center`: Center of the circle, default `[0.0;0.0;0.0]`
- `sang`: Start angle, default `0` degrees
- `eang`: End angle `360` degrees

# Output:
- `crv`: NURBS curve for a circular arc.

# Examples:
```julia
julia> crv = nrbcirc()
julia> crv = nrbcirc(radius)
julia> crv = nrbcirc(radius,center)
julia> crv = nrbcirc(radius,center,sang,eang)
```
"""
function nrbcirc(radius::F=1.0,center::Vector{F}=[0.0;0.0;0.0],
                 sang::F=0.0,eang::F=2.0*π)::NURBS1D where {F<:AbstractFloat}
# TODO: maybe revisit this function and split it up
# TODO: optional arguments in julia are always at the end and providing [] as
#       argument input is a bad habbit i changed the argument behaviour for
#       center > examine if that did change the behaviour in this toolbox

# sweep angle of arc
sweep = eang - sang;
if sweep < 0
  sweep = 2 * π + sweep;
end # if

if abs(sweep) <= π/2
  narcs = 1;
  knots = vec([0.0 0.0 0.0 1.0 1.0 1.0]);
elseif abs(sweep) <= π
  narcs = 2;
  knots = vec([0.0 0.0 0.0 1/2 1/2 1.0 1.0 1.0]);
elseif abs(sweep) <= π*3/2
  narcs = 3;
  knots = vec([0.0 0.0 0.0 1/3 1/3 2/3 2/3 1.0 1.0 1.0]);
else
  narcs = 4;
  knots = vec([0.0 0.0 0.0 1/4 1/4 1/2 1/2 3/4 3/4 1.0 1.0 1.0]);
end # if

# arc segment sweep angle/2
dsweep = sweep / (2 * narcs);

# determine middle control point and weight
wm = cos(dsweep);
x  = radius * wm;
y  = radius * sin(dsweep);
xm = x + y * tan(dsweep);

# arc segment control points
ctrlpt = [ x wm*xm x;  # w*x - coordinate
          -y   0   y;  # w*y - coordinate
           0   0   0;  # w*z - coordinate
           1   wm  1]; # w   - coordinate

# build up complete arc from rotated segments
coefs = zeros(F,(4,2*narcs+1)); # nurbs control points of arc
xx = vecrotz(sang + dsweep);
coefs[:,1:3] = xx * ctrlpt;     # rotate to start angle
xx = vecrotz(2 * dsweep);
for n in 2:narcs
  m = (0:1) .+ (2*n);
  coefs[:,m] = xx * coefs[:,m .- 2];
end # if

# vectrans arc if necessary
if !(all(i->i==0,center))
  xx = vectrans(center);
  coefs = xx * coefs;
end # if

return nrbmak(coefs,knots);
end # nrbcirc


# # Testing the stability of a CKSVAR model

# Here, we test the stability of the CKSVAR model estimated for Japan in:
# - Ikeda, Li, Mavroeidis and Zanetti (2021), "[Testing the effectiveness of monetary policy in Japan and the United States](https://arxiv.org/abs/2012.15158)." Working paper.
# This paper is hereafter referred to as ILMZ21.
##
# In ILMZ21, the effectiveness of unconventional monetary policy in Japan is investigated via estimation of a CKSVAR(2) model in inflation, output gap and the Call Rate for Japan, over the period 1985q3-2019q1, with shadow call rate as the latent variable.
#
# Converting this model into TAR form, we determine upper bounds on the state-constrained joint spectral radius (SCJSR)
##
using DataFrames, DataFramesMeta, CSV, HTTP
using ThresholdStability

extract_est(parname, df) = @subset(df, in([parname]).(:Parameter)).:Estimate[1]

JP = CSV.read(HTTP.get("https://raw.githubusercontent.com/samwycherley/ThresholdStability.jl/master/examples/src/estimates/estimates_JP.csv").body, DataFrame, copycols=true)  # Japan data from CSV
##
##
# This loads the Japan estimates as a dataframe and defines a function to extract estimates. Next, we build the matrices:
##
##
JPβtilde = zeros(2)
for i in 1:2
    JPβtilde[i] = extract_est("tildebeta_$i", JP)
end

JPCstar = zeros(3, 2)
for i in 1:3
    for j in 1:2
        JPCstar[i, j] = extract_est("Eq.$i lSR_$j", JP)
    end
end

Cbars = []
for j in 1:2
    Cbar_j = zeros(3, 3)
    for i in 1:3
        Cbar_j[i, 1] = extract_est("Eq.$i INFL_$j", JP)
        Cbar_j[i, 2] = extract_est("Eq.$i YGAP_BOJ_$j", JP)
        Cbar_j[i, 3] = extract_est("Eq.$i SR_$j", JP)
    end
    push!(Cbars, Cbar_j)
end
JPCbar = hcat(Cbars[1], Cbars[2])

JPC = copy(JPCbar); k = size(JPC, 1)  # convert C̄ to C
for i in 1:2
    JPC[:, i*k] -= JPCstar[:, i]
end

Σ_JP, X_JP = CKSVAR_to_TAR(JPC, JPCstar, JPβtilde, 2)
##
# Write this as a state-dependent switched linear system:
##
##
G_JP = automaton_constructor(Σ_JP)
s_JP = discreteswitchedsystem(Σ_JP, G_JP, X_JP)
##
@show sosbound_γ(s_JP, 2)
##
# We see that the model estimated for Japan is stable


function polyMap(Θ,z)
  g = 0.
  for i in 1:1:length(Θ)
    g += Θ[i]⋅z^(i-1)
  end
  return g
end

function cfy(yi, yp, yn, theta)
  return ( (yi - polyMap(theta, yp)) ^ 2 + (yn - polyMap(theta, yi)) ^ 2)
end

function noiseMix3(p1::Float64, lam1::Float64, lam2::Float64)
  if rand() < p1
    z = rand(Normal(0, sqrt(1 / lam1)))
  else
    z = rand(Normal(0, sqrt(1 / lam2)))
  end
  return z
end

function genData(npar, n, Θ, x₀, p1, lam1, lam2, seed)
  srand(seed)
  x = zeros(n)

  if npar .== 0
    x[1] = polyMap(Θ,x₀) + rand(Normal(0, sqrt(1 / lam1)))
    for i in 2:1:n
      x[i] = polyMap(Θ,x[i-1]) + rand(Normal(0, sqrt(1 / lam1)))
    end
  elseif npar .== 1
    x[1] = polyMap(Θ,x₀) + noiseMix1()
    for i in 2:1:n
      x[i] = polyMap(Θ,x[i-1]) + noiseMix1()
    end
  elseif npar .== 2
    x[1] = polyMap(Θ,x₀) + noiseMix2(4)
    for i in 2:1:n
      x[i] = polyMap(Θ,x[i-1]) + noiseMix2(4)
    end
  else
    x[1] = polyMap(Θ,x₀) + noiseMix3(p1, lam1, lam2)
    for i in 2:1:n
      x[i] = polyMap(Θ,x[i-1]) + noiseMix3(p1, lam1, lam2)
    end
  end

  return x
end

function unifmixrnd(x)
  local U::Float64
  n = length(x)
  nc = 0
  for i in 1:2:n
    nc += x[i + 1] - x[i]
  end

  U = 0
  u = rand()
  prob = 0
  for i in 1:2:n
    prob += (x[i + 1] - x[i]) / nc
    if u .< prob
      U = x[i] + (x[i + 1] - x[i]) * rand()
      break
    end
  end
  return(U)
end

function rangeIntersection(A::Array{Float64}, B::Array{Float64})
#=
Purpose: Range/interval intersection

 A and B two ranges of closed intervals written
 as vectors [lowerbound1 upperbound1 lowerbound2 upperbound2]
 or as matrix [lowerbound1, lowerbound2, lowerboundn;
               upperbound1, upperbound2, upperboundn]
 A and B have to be sorted in ascending order

 out is the mathematical intersection A n B

 EXAMPLE USAGE:
   >> out=rangeIntersection([1 3 5 9],[2 9])
   	out =  [2 3 5 9]
   >> out=rangeIntersection([40 44 55 58], [42 49 50 52])
   	out =  [42 44]
=#

# Allocate, as we don't know yet the size, we assume the largest case
  out1 = zeros(length(B)+(length(A)-2))
  k = 1

  while isempty(A) .== 0 && isempty(B) .== 0
  # make sure that first is ahead second
    if A[1] .> B[1]
      temp = copy(B)
      B = copy(A)
      A = copy(temp)
    end

    if A[2] .< B[1]
      A = copy(A[3:end])
      continue
    elseif A[2] .== B[1]
      out1[k] = B[1]
      out1[k + 1] = B[1]
      k = k + 2

      A = A[3:end]
      continue
    else
      if A[2] .== B[2]
        out1[k] = B[1]
        out1[k+1] = B[2]
        k = k + 2

        A = A[3:end]
        B = B[3:end]

      elseif A[2] .< B[2]
        out1[k] = B[1]
        out1[k+1] = A[2]
        k = k + 2

        A = A[3:end]
      else
        out1[k] = B[1]
        out1[k+1] = B[2]
        k = k + 2

        B = B[3:end]
      end
     end
    end

  # Remove the tails
  out = copy(out1[1:k-1])
  return(out)
end

function etgeornd(p::Float64, k::Vector{Int64})
local z::Vector{Int64}
  z = floor(log.(rand(length(k))) ./ log.(1 - p)) + k
  return(z)
end


abstract type AbstractFeatureNormalize <: AbstractFeatureConstructor end

"""
    MinMaxNormalize{T} <: AbstractFeatureNormalize

Preprocessor which normalizes from a known set of min and max.
"""
struct MinMaxNormalize{T} <: AbstractFeatureNormalize
    # Main Arguments
    min::Array{T, 1}
    max::Array{T, 1}
    min_p_max::Array{T, 1}
    MinMaxNormalize(min::Array{T, 1}, max::Array{T, 1}) where {T} =
        new{T}(min, max, max .- min)
end

create_features(fc::MinMaxNormalize, s) = (s .- fc.min)./fc.min_p_max
feature_size(fc::MinMaxNormalize) = length(fc.min)

"""
    MeanStdNormalize{T} <: AbstractFeatureNormalize

Normalize the vector to have 0 mean and 1 standard deviation
"""
struct MeanStdNormalize{T} <: AbstractFeatureNormalize
    mean::Array{T, 1}
    std::Array{T, 1}
end

create_features(fc::MeanStdNormalize, s) = (s .- fc.mean)./fc.std
feature_size(fc::MeanStdNormalize) = length(fc.mean)


"""
Subroutine dprsrch

This subroutine uses a projected search to compute a step
that satisfies a sufficient decrease condition for the quadratic

    q(s) = 0.5*s'*A*s + g'*s,

where A is a symmetric matrix in compressed column storage,
and g is a vector. Given the parameter alpha, the step is

    s[alpha] = P[x + alpha*w] - x,

where w is the search direction and P the projection onto the
n-dimensional interval [xl,xu]. The final step s = s[alpha]
satisfies the sufficient decrease condition

    q(s) <= mu_0*(g'*s),

where mu_0 is a constant in (0,1).

The search direction w must be a descent direction for the
quadratic q at x such that the quadratic is decreasing
in the ray x + alpha*w for 0 <= alpha <= 1.

MINPACK-2 Project. March 1999.
Argonne National Laboratory.
Chih-Jen Lin and Jorge J. More'.
"""
function dprsrch(n,x,xl,xu,A,g,w,wa1,wa2)
    one = 1.0
    p5 = 0.5

    # Constant that defines sufficient decrease.
    mu0 = 0.01

    # Interpolation factor.

    interpf = 0.5

    # Set the initial alpha = 1 because the quadratic function is
    # decreasing in the ray x + alpha*w for 0 <= alpha <= 1.

    alpha = one
    nsteps = 0

    # Find the smallest break-point on the ray x + alpha*w.
    nbrpt,brptmin,brptmax = dbreakpt(n,x,xl,xu,w)

    search = true
    while (search && alpha > brptmin)

        # Calculate P[x + alpha*w] - x and check the sufficient
        # decrease condition.

        nsteps = nsteps + 1
        dgpstep(n,x,xl,xu,alpha,w,wa1)
        dssyax(n,A,wa1,wa2)
        gts = ddot(n,g,1,wa1,1)
        q = p5*ddot(n,wa1,1,wa2,1) + gts
        if q <= mu0*gts
            search = false
        else

            # This is a crude interpolation procedure that
            # will be replaced in future versions of the code.

            alpha = interpf*alpha
        end
    end

    # Force at least one more constraint to be added to the active
    # set if alpha < brptmin and the full step is not successful.
    # There is sufficient decrease because the quadratic function
    # is decreasing in the ray x + alpha*w for 0 <= alpha <= 1.

    if (alpha < one && alpha < brptmin)
        alpha = brptmin
    end

    # Compute the final iterate and step.

    dgpstep(n,x,xl,xu,alpha,w,wa1)
    daxpy(n,alpha,w,1,x,1)
    dmid(n,x,xl,xu)
    dcopy(n,wa1,1,w,1)

    return
end


# find hyperparameter values using gird search

include(pwd()*"/dlmodel.jl")

using Distributions
using ProgressMeter

N = 1000

dist = Uniform(0, 0.1)

rand(dist, 3)

hyperparameters = zeros(N,3)
model_fit_regularization = zeros(N,3+3+1+1)
model_fit_train = zeros(N,2000)
model_fit_val = zeros(N,2000)

# find hyperparameters using random search
@time begin
@showprogress for i in 1:N

    Random.seed!(1234) # set random numbers

    lambda = rand(dist,3)

    hyperparameters[i,:] = lambda

    w = Any[ xavier(n_hidden_1,n_input), zeros(n_hidden_1,1),
             xavier(n_hidden_2,n_hidden_1), zeros(n_hidden_2,1),
             xavier(n_out,n_hidden_2), zeros(n_out,1)]

    optim = optimizers(w, Adam)

    dropout_percentage = 0
    loss_train_vec, loss_val_vec = train(2000, dropout_percentage, lambda, w, optim)

    model_fit_train[i,:] = loss_train_vec
    model_fit_val[i,:] = loss_val_vec

    # calc predictions
    y_pred_train_proc = predict(w, x_train, 0)
    y_pred_val_proc = predict(w, x_val, 0)

    y_pred_train = class_nll(y_pred_train_proc,[1,2])
    y_pred_val = class_nll(y_pred_val_proc,[1,2])

    # loss
    model_fit_regularization[i,end-1] = loss(w,x_train, y_train, 0, lambda)
    model_fit_regularization[i,end] = loss(w,x_val, y_val, 0, lambda)

    # classification results
    model_fit_regularization[i,1:3] = classification_results(y_pred_train, y_train, [1,2])
    model_fit_regularization[i,4:6] = classification_results(y_pred_val, y_val, [1,2])

end
end

# add run w.o regularization

w = Any[ xavier(n_hidden_1,n_input), zeros(n_hidden_1,1),
         xavier(n_hidden_2,n_hidden_1), zeros(n_hidden_2,1),
         xavier(n_out,n_hidden_2), zeros(n_out,1)]

optim = optimizers(w, Adam)

# training
dropout_percentage = 0.
lambda = [0.0,0.0,0.0]
@time loss_train_vec, loss_val_vec = train(2000, dropout_percentage, lambda, w, optim)

# calc predictions
y_pred_train_proc = predict(w, x_train, 0)
y_pred_val_proc = predict(w, x_val, 0)

y_pred_train = class_nll(y_pred_train_proc,[1,2])
y_pred_val = class_nll(y_pred_val_proc,[1,2])

# loss
training_loss_no_reg = loss(w,x_train, y_train, 0, lambda)
val_loss_no_reg = loss(w,x_val, y_val, 0, lambda)

# classification results
training_classification_res_no_reg = classification_results(y_pred_train, y_train, [1,2])
val_classification_res_no_reg = classification_results(y_pred_val, y_val, [1,2])

# find best hyperparameter settings
(~, idx_best_training_loss) = findmin(model_fit_regularization[:,end-1])
(~, idx_best_val_loss) = findmin(model_fit_regularization[:,end])

(~, idx_best_training_tp) = findmax(model_fit_regularization[:,2])
(~, idx_best_val_tp) = findmax(model_fit_regularization[:,5])

(~, idx_best_training_tn) = findmax(model_fit_regularization[:,3])
(~, idx_best_val_tn) = findmax(model_fit_regularization[:,6])

(~, idx_best_training_acc) = findmax(model_fit_regularization[:,1])
(~, idx_best_val_acc) = findmax(model_fit_regularization[:,4])

@printf "Best lambda_1 %.10f\n" lambda_1[1]
@printf "Best lambda_2 %.10f\n" lambda_2[4]
@printf "Best lambda_3 %.10f\n" lambda_3[1]


using Postprocessing
using Test
using LinearAlgebra, Random

@testset "total variation denoising" begin
  println(devnull, reference(total_variation_denoising))

  for T in (Float32, Float64)
    x = range(zero(T), 2*one(T), length=1000)
    y_true = sinpi.(x)
    y_noisy = y_true .+ 0.1 .* randn(T, length(x))
    y_denoised = total_variation_denoising(y_noisy, 1)
    @test norm(y_denoised - y_true) < norm(y_noisy - y_true)
    @test_throws ArgumentError total_variation_denoising(y_noisy, -1)
    @test_throws DimensionMismatch total_variation_denoising!(similar(y_noisy, T, length(y_noisy)-1), y_noisy, 1)
  end
end

@testset "group sparse total variation denoising" begin
  println(devnull, reference(group_sparse_total_variation_denoising))

  for T in (Float32, Float64)
    x = range(zero(T), 2*one(T), length=1000)
    y_true = sinpi.(x)
    y_noisy = y_true .+ 0.1 .* randn(T, length(x))
    y_denoised = group_sparse_total_variation_denoising(y_noisy, 1, group_size=3)
    @test norm(y_denoised - y_true) < norm(y_noisy - y_true)
    @test total_variation(y_denoised) < total_variation(y_noisy)
    @test_throws ArgumentError group_sparse_total_variation_denoising(y_noisy, -1, group_size=3)
    @test_throws DimensionMismatch group_sparse_total_variation_denoising!(similar(y_noisy, T, length(y_noisy)-1), y_noisy, 1)
  end
end


struct BundleMethod <: UnconstrainedNonSmoothSolver
  solver_factory # Must support QPs with linear constraints.
  µ::Float64 # Used in the proximal step. Considered to be constant for now. TODO: schedules
  K::Float64
  n_iter::Int
end

function solve(p::UnconstrainedNonSmoothProblem, params::BundleMethod, x0::Vector{Float64}; info_callback::Union{Nothing, Function}=nothing)
  @assert p.sense == Minimise # TODO: current limitation.
  @assert length(x0) == p.dimension

  # Build the bundle model.
  m = JuMP.Model(params.solver_factory)
  f = @JuMP.variable(m)
  bx = @JuMP.variable(m, [1:p.dimension])

  # Create the first cutting plane with the starting iterate (i.e. the program will never be unbounded).
  @JuMP.constraint(m, f >= p.f(x0) + dot(p.g(x0), bx - x0))

  # Start iterating.
  x = copy(x0)
  for k in 1:params.n_iter
    t0 = time_ns()

    # Get the new test point (proximal step).
    @JuMP.objective(m, Min, f + params.µ * sum((bx[i] - x[i])^2 for i in 1:3))
    JuMP.optimize!(m)
    @assert JuMP.termination_status(m) == MOI.OPTIMAL

    y = value.(bx)
    @JuMP.constraint(m, f >= p.f(y) + dot(p.g(y), bx - y))

    # Is this point optimum?
    @JuMP.objective(m, Min, f)
    c = @JuMP.constraint(m, bx .== y)
    JuMP.optimize!(m)
    @assert JuMP.termination_status(m) == MOI.OPTIMAL

    f_bundle_y = JuMP.value(f)
    JuMP.delete(m, c)

    f = p.f(x)
    v = f_bundle_y - f
    if v >= 0
      break
    end

    # Perform a step:
    if p.f(y) <= f + params.K * v
      # Descent step.
      x = y
    else
      # Short step.
      nothing
    end

    if info_callback !== nothing
      info_callback(k, f, x, g, f_best, x_best, time_ns() - t0)
    end
  end

  return x
end


using PSDMatrices
using Test

using PSDMatrices: X_A_Xt
using LinearAlgebra

@testset "PSDMatrices.jl" begin

    d = 5
    L = LowerTriangular(randn(d,d))
    LLT = L*L'
    A = PSDMatrix(L)
    B = PSDMatrix(LowerTriangular(rand(d,d)))

    @testset "Basics" begin
        @test A == LLT
        @test all(A .== LLT)
        @test size(A) == size(LLT)
        @test A[:] == LLT[:]

        @test PSDMatrix(Matrix(A)) ≈ A

        @test copy(A) isa PSDMatrix
        copy!(B, A)
        @test B == A

        @test Matrix(A) == LLT
        @test typeof(Matrix(A)) === typeof(LLT)
    end

    @testset "LinearAlgebra" begin
        @test diag(A) == diag(LLT)
        @test det(A) ≈ det(LLT)
        @test logdet(A) ≈ logdet(LLT)

        @test inv(A) ≈ inv(LLT)
        x = rand(d)
        @test A\x ≈ LLT\x

        @test cholesky(A) == cholesky(LLT)
    end

    @testset "X_A_Xt" begin
        X = rand(d,d)
        @test X_A_Xt(A, X) isa PSDMatrix
        @test X_A_Xt(A, X) ≈ X*A*X'

        D = Diagonal(rand(d))
        @test X_A_Xt(A, D) isa PSDMatrix
        @test X_A_Xt(A, D) ≈ D*A*D'
    end

    @testset "Arithmetic" begin
        L2 = LowerTriangular(randn(d,d))
        B = PSDMatrix(L2)
        LLT2 = L2*L2'
        @test A+B isa PSDMatrix
        @test A+B ≈ LLT+LLT2

        c = rand()
        @test c*A isa PSDMatrix
        @test 0*A isa PSDMatrix
        @test !(-1*A isa PSDMatrix)
        @test c*A == A*c
    end

    @testset "BigFloats" begin
        L = LowerTriangular(big.(randn(d,d)))
        @test PSDMatrix(L) isa PSDMatrix
    end
end


using SemiclassicalOrthogonalPolynomials, ClassicalOrthogonalPolynomials, Test
import ClassicalOrthogonalPolynomials: orthogonalityweight, Weighted, associated

@testset "Two Band" begin
    @testset "TwoBandWeight" begin
        w = TwoBandWeight(0.6,1,2,3)
        @test w[0.7] == w[-0.7] == 0.7^6 * (0.7^2-0.6^2)^2 * (1-0.7^2)
        @test_throws BoundsError w[0.4]
        @test_throws BoundsError w[-0.4]
        @test_throws BoundsError w[1.1]
        @test copy(w) == w
        @test_broken sum(w)
    end
    @testset "Chebyshev case" begin
        ρ = 0.5
        T = TwoBandJacobi(ρ, -1/2, -1/2, 1/2)
        @test T[0.6,1:5] ≈ [T[0.6,k] for k=1:5]

        U = TwoBandJacobi(ρ, 1/2, 1/2, -1/2)
        @test U[0.6,1:5] ≈ [U[0.6,k] for k=1:5]

        @test copy(T) == T
        @test U ≠ T
        @test orthogonalityweight(T) == TwoBandWeight(ρ, -1/2, -1/2, 1/2)

        # bug
        @test !issymmetric(jacobimatrix(T)[1:10,1:10])
    end

    @testset "Hilbert" begin
        ρ = 0.5
        w = TwoBandWeight(ρ, -1/2, -1/2, 1/2)
        x = axes(w,1)
        H = inv.(x .- x')
        @test iszero(H*w)
        @test sum(w) == π/2

        T = TwoBandJacobi(ρ, -1/2, -1/2, 1/2)
        Q = associated(T)
        @test (Q[0.6,1:100]' * (Q[:,Base.OneTo(100)] \ exp.(x))) ≈ exp(0.6)
        @test (Q[-0.6,1:100]' * (Q[:,Base.OneTo(100)] \ exp.(x))) ≈ exp(-0.6)

        @test (Q * (Q \ exp.(x)))[0.6] ≈ exp(0.6)
        @test_broken Q \ (H * Weighted(T)) # need to deal with Hcat

        @test_broken H*TwoBandWeight(ρ, 1/2, 1/2, -1/2)
    end
end


###########
# Library #
###########
# using MatrixLM
# using Distributions, Random, Statistics, LinearAlgebra, StatsBase
# using Random
using MatrixLMnet
using Helium
using Test

####################################################
# TEST Lasso vs Elastic Net (𝛼=1) - Simulated Data #
####################################################

#=
Description:
-----------

Model: 𝐘 = 𝐗 𝛃 𝐙' + 𝜎𝜖, with 𝜖∼𝑁(0,1)

Simulate data set consisting  of 20/20/200 observations and 8 predictors.
We let 𝛽₁ = (3, 1.5, 0, 0, 2, 0, 0, 0), 𝛽₂ = (0, 1.5, 0, 3.5, 2, 0, 0 , 2) where
𝛃 = [𝛽₁, 𝛽₂] and 𝜎 = 3.
The pairwise correlation between 𝑋ᵢ and 𝑋ⱼ was set to be 𝑐𝑜𝑟(𝑖,𝑗)=(0.5)^|𝑖−𝑗|.
Here, the Z matrix is an identity matrix.
=#

# Data testing directory name
dataDir = realpath(joinpath(@__DIR__,"data"))

# Get predictors
X = Helium.readhe(joinpath(dataDir, "Xmat.he"))

# Get response
Y = Helium.readhe(joinpath(dataDir, "Ymat.he"))

# Get Z matrix
Z = Helium.readhe(joinpath(dataDir, "Zmat.he"))

# Build raw data object from MatrixLM.jl
dat = RawData(Response(Y), Predictors(X, Z));

# Hyper parameters
λ = [10.0, 5.0, 3.0]
α = [1.0]

rng = MatrixLMnet.Random.MersenneTwister(2021)

############################################
# TEST 1 Lasso vs Elastic Net (𝛼=1) - ista #
############################################

# Elastic net penalized regression
est_ista_1 = MatrixLMnet.mlmnet(dat, λ, α, method = "ista", hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_ista_1 = est_ista_1.B[:, :, 3, 1];

# Elastic net penalized regression
est_ista_2 = MatrixLMnet.mlmnet(dat, λ, method = "ista",  hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_ista_2 = est_ista_2.B[:, :, 3, 1];

# Lasso penalized regression - ista
B_ista= Helium.readhe(joinpath(dataDir, "B_ista.he"))

println("Lasso vs Elastic Net when α=1 test 1 - ista: ", @test (B_Net_ista_1 ≈ B_ista) && (B_Net_ista_2 ≈ B_ista))

#############################################
# TEST 2 Lasso vs Elastic Net (𝛼=1) - fista #
#############################################

# Elastic net penalized regression
est_fista_1 = MatrixLMnet.mlmnet(dat, λ, α, method = "fista", hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_fista_1 = est_fista_1.B[:, :, 3, 1];

# Elastic net penalized regression
est_fista_2 = MatrixLMnet.mlmnet(dat, λ, method = "fista",  hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_fista_2 = est_fista_2.B[:, :, 3, 1];

# Lasso penalized regression - fista
B_fista= Helium.readhe(joinpath(dataDir, "B_fista.he"))

println("Lasso vs Elastic Net when α=1 test 2 - fista: ", @test (B_Net_fista_1 ≈ B_fista) && (B_Net_fista_2 ≈ B_fista))

##########################################################
# TEST 3 Lasso vs Elastic Net (𝛼=1) - fista backtracking #
##########################################################

# Elastic net penalized regression
est_fistabt_1 = MatrixLMnet.mlmnet(dat, λ, α, method = "fista_bt", hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_fistabt_1 = est_fistabt_1.B[:, :, 3, 1];

# Elastic net penalized regression
est_fistabt_2 = MatrixLMnet.mlmnet(dat, λ, method = "fista_bt",  hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_fistabt_2 = est_fistabt_2.B[:, :, 3, 1];

# Lasso penalized regression - fista-bt
B_fistabt = Helium.readhe(joinpath(dataDir, "B_fistabt.he"))

println("Lasso vs Elastic Net when α=1 test 3 - fista-bt: ", @test (B_Net_fistabt_1 ≈ B_fistabt) && (B_Net_fistabt_2 ≈ B_fistabt))

############################################
# TEST 4 Lasso vs Elastic Net (𝛼=1) - admm #
############################################

# Elastic net penalized regression
est_admm_1 = MatrixLMnet.mlmnet(dat, λ, α, method = "admm", hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_admm_1 = est_admm_1.B[:, :, 3, 1];

# Elastic net penalized regression
est_admm_2 = MatrixLMnet.mlmnet(dat, λ, method = "admm",  hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_admm_2 = est_admm_2.B[:, :, 3, 1];

# Lasso penalized regression - admm
B_admm = Helium.readhe(joinpath(dataDir, "B_admm.he"))

println("Lasso vs Elastic Net when α=1 test 4 - admm: ", @test (B_Net_admm_1 ≈ B_admm) && (B_Net_admm_2 ≈ B_admm))

##########################################
# TEST 5 Lasso vs Elastic Net (𝛼=1) - cd #
##########################################

# Elastic net penalized regression
MatrixLMnet.Random.seed!(rng)
est_cd_1 = MatrixLMnet.mlmnet(dat, λ, α, method = "cd", hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_cd_1 = est_cd_1.B[:, :, 3, 1];

# Elastic net penalized regression
MatrixLMnet.Random.seed!(rng)
est_cd_2 = MatrixLMnet.mlmnet(dat, λ, method = "cd",  hasZIntercept = false, hasXIntercept = false, isVerbose = false);
B_Net_cd_2 = est_cd_2.B[:, :, 3, 1];

# Lasso penalized regression - cd
B_cd = Helium.readhe(joinpath(dataDir, "B_cd.he"))

println("Lasso vs Elastic Net when α=1 test 5 - cd: ", @test ≈(B_Net_cd_1,  B_cd; atol=1.2e-4) && ≈(B_Net_cd_2, B_cd;  atol=1.2e-4))

println("Tests mlmnet finished!")


##
function get_pdf_samples(samples, trans_samples, log_posterior, prop_dist, lb, ub, names)
    n = size(samples, 2)
    n == size(trans_samples,2) || throw(DimensionMismatch("samples dimensions don't match"))

    g = zeros(n)
    p = zeros(n)

    for i in 1:n
        g_sample = @views(trans_samples[:,i])
        p_sample = @views(samples[:,i])
        logdetJ = sum(logjacobian_contribution.(transform.(p_sample, lb, ub), lb, ub))
        g[i] = logpdf(prop_dist, g_sample)
        p[i] = isnothing(names) ? log_posterior(p_sample) + logdetJ : log_posterior(NamedTuple{Tuple(names)}(p_sample)) + logdetJ
    end

    return p, g, p-g
end

## Iterative algo
function iterative_algorithm(l₁, l₂, n₁, n₂; tol, maxiter)

    lstar = median(l₁)
    r = 0.0
    s₁ = n₁ / (n₁ + n₂)
    s₂ = n₂ / (n₁ + n₂)

    logml = -Inf
    ϵ = 0.0
    for i = 1:maxiter
        numterm = 0.0
        denomterm = 0.0
        for (l₁ⱼ, l₂ⱼ) in zip(l₁, l₂)
            el1 = clamp(exp(l₁ⱼ - lstar), 0.0, 1e100) # Avoiding numerical errors
            el2 = clamp(exp(l₂ⱼ - lstar), 1e-100, 1e100)
            numterm += el2 / (s₁ * el2 + s₂ * r)
            denomterm += 1 / (s₁ * el1 + s₂ * r)

        end
        rnew = (1/n₂) * numterm / (denomterm / n₁)

        logmlnew = log(rnew) + lstar
        ϵ = abs((logmlnew - logml) / logmlnew)
        if ϵ < tol
            return logmlnew, i
        end
        logml = logmlnew
        r = rnew
    end
    @warn "Maximum number of iterations ($maxiter) reached before convergence under the tolerance level $tol"
    return logml, maxiter
end


### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9bc349e8-efb0-11ea-1392-cd735828af33
# packaged loads, etc
begin
	using DelimitedFiles
	using NIfTI
	using Images
	using StatsBase
	using Plots
	using PlutoUI
	using Printf
	using LinearAlgebra
end

# ╔═╡ df5232e0-04e2-11eb-0c26-f19b9d9c7b14
begin
	using GLM
	using DataFrames
end

# ╔═╡ 8029070e-efb0-11ea-1d0a-7f402f2e9989

# ╔═╡ 4df55a88-f98a-11ea-1c9f-cd52a4b5062f
md"""
# How do we analyze data here?

## General idea

What happens if I multiply each point in a timeseries by a number (or "weight", "multiplier", ...) ?
"""

# ╔═╡ c14bf31e-f98c-11ea-33d1-f5de8f2ad0c1
md"""
**weight:** $(@bind weight Slider(range(0.25, 3.0, step=0.25), default=0.75, show_value=true))
"""

# ╔═╡ 2b2f6be4-f07f-11ea-1459-775749ca949f
md"""## Load in that CSV file """

# ╔═╡ 368e2640-f8e0-11ea-1ddf-b37069ba7588
d, h = readdlm("data/stimulus_timing.csv", ',', Float64, header=true);

# ╔═╡ cac0339a-f8e4-11ea-0876-f91dc7db1454
md"""## Stimulus timing $\rightarrow$ fMRI response timing?

"""

# ╔═╡ 1058934e-f8e1-11ea-1896-35d75c4b6a8e
begin
    md"""
*Make a slider that controls the plot*

**row:** $(@bind whichRow Slider(1:2:size(d,1), default=1, show_value=true))
"""

end

# ╔═╡ c6cab78c-f836-11ea-2b5b-23f0006f8e93
begin
	plot(view(d, 1:whichRow,1), view(d,1:whichRow,2), xlabel="Time (volumes)",
		ylabel="fMRI response", label="$(h[2])", xlims=(0, 160), ylims=(-5, 5), lw=2)
	plot!(view(d, 1:whichRow,1), view(d,1:whichRow,3),
		label="$(h[3])", xlims=(0, 160), ylims=(-5, 5), lw=3, color=:red)

end

# ╔═╡ da4ac85c-f8e4-11ea-192e-03d5b6862d5a
md"""## Scaling = multiplying"""

# ╔═╡ 1c3b6c56-f8e2-11ea-3969-f59871025c5e
begin
    md"""
How can we *scale* the idealised prediction to fit the data better?

**multiplier:** $(@bind multiplier Slider(-10:0.5:10, default=1, show_value=true))
"""
end

# ╔═╡ ddf6018c-f921-11ea-0d74-c16f4f0b7426
md"""here I am plotting the data (blue line) and the prediction for what an "ideal" face-selective response would look like... and we can scale it"""

# ╔═╡ 39d465a6-f8e2-11ea-2d71-75cbe586428c
begin
	plot(d[:,1], d[:,5],
		xlim=(0, 160), ylim=(-5, 5),
		xlabel="Time (volumes)",
	ylabel="fMRI response (%signal change)", label="data")
	plot!(d[:,1], d[:,3] * multiplier, lw=2, label="model * $(@sprintf "%.1f" multiplier )")

end

# ╔═╡ a1b031a4-fc08-11ea-2314-13e2e79c2ca3
md"""### Now: more than one component to fit"""

# ╔═╡ e6309486-fc08-11ea-3ade-c93a24b6866b
begin
	plot(d[:,1], d[:,3],
		xlim=(0, 160), ylim=(-1, 1),
		xlabel="Time (volumes)",
		ylabel="fMRI response (%signal change)",
		label="response (FACES)",
		color=:red, lw=2, layout=(2, 1), subplot=1)
	plot!(d[:,1], d[:,4],
		color=:blue, lw=2,
		label="response (OBJECTS)", subplot=2,xlim=(0, 160), ylim=(-1, 1))

end

# ╔═╡ 9c3b0258-fc08-11ea-1922-fd1bec66b4ff
begin
    md"""
How can we *scale* the idealised prediction to fit the data better?

**multiplier (faces):** $(@bind mf Slider(-10:0.5:10, default=1, show_value=true))

**multiplier (objects):** $(@bind mo Slider(-10:0.5:10, default=1, show_value=true))

"""
end

# ╔═╡ 975860d2-fc08-11ea-2def-ad1c45826911
begin
	plot(d[:,1], d[:,5],
		xlim=(0, 160), ylim=(-5, 5),
		xlabel="Time (volumes)",
	ylabel="fMRI response (%signal change)", label="data")
	plot!(d[:,1], d[:,[3,4]] * [mf; mo], lw=2, label="model * $(@sprintf "%.1f" multiplier )")

end

# ╔═╡ e7fb832e-04e0-11eb-388d-f7520b672a6e
md"""

## Calculating the weights for the best fit

*Linear regression* can be used to find/calculate the weights that will give the **best** fit of the mixture of the two curves.

**Best** here means the settings that will result in the *least squares fit*, the fit that makes the *sum of squared errors* between each measurement in the timeseries ($\text{data}_i$) and the corresponding model description ($\text{model}_i$) smallest.

$\sum_i (\text{data}_i - \text{model}_i)^2$

"""

# ╔═╡ 2b070720-04e2-11eb-3101-cbc40b17682b
# X, the design matrix
designMatrix = d[:,3:4]

# ╔═╡ 4b968440-04e2-11eb-35ce-5f050dcd4b6b
y_1 = d[:,5] #

# ╔═╡ 6464f806-04e2-11eb-37cc-7363fa76c80e
β_1 = designMatrix \ y_1

# ╔═╡ 111aea9a-04e5-11eb-066d-4d245b04dc5e
md"""
##

turn data into a `dataframe` and use `GLM` module to get a linear model fit

"""

# ╔═╡ 1e370abc-04e3-11eb-064a-21fbb9fafc7c
theData = DataFrame(d[:, 3:6], [:face_response, :object_response, :Voxel1, :Voxel2 ])

# ╔═╡ 1d60545a-04e4-11eb-1720-b5f04c5da239
ols1 = lm(@formula(Voxel1 ~ face_response + object_response), theData)

# ╔═╡ 86d10676-04e4-11eb-2121-2f0c8945cc1a
coefnames(ols1)

# ╔═╡ 979732dc-04e4-11eb-3803-7302dc91e777
coef(ols1)

# ╔═╡ a33a18f2-04e4-11eb-0587-6117c036c641
GLM.confint(ols1)

# ╔═╡ c60d9136-1d5f-11eb-0839-d760afb98985

# ╔═╡ c47f840a-1d5f-11eb-1e88-d3835fa021e3

# ╔═╡ aa954778-1d5f-11eb-00bf-034690d40e94
md"""## Plot using the `calculated` weights"""

# ╔═╡ 5d60ed36-1d5f-11eb-3e90-9315dd9bc07a
begin
	plot(d[:,1], d[:,5],
		xlim=(0, 160), ylim=(-5, 5),
		xlabel="Time (volumes)",
	ylabel="fMRI response (%signal change)", label="data")
	plot!(d[:,1], d[:,[3,4]] * coef(ols1)[2:3], lw=2, label="model * [$(@sprintf "%.2f; %.2f" coef(ols1)[2] coef(ols1)[3] )]")

end

# ╔═╡ cc5dcb6e-1d5f-11eb-0be7-57d4795af804

# ╔═╡ cb65baa0-1d5f-11eb-346f-bdb3a85d0fe7

# ╔═╡ ca5b1768-1d5f-11eb-0459-13e4d2de80c7

# ╔═╡ c3899100-04e4-11eb-37aa-394d3f60bd53
import Pkg; Pkg.add("StatsPlots")

# ╔═╡ d5c0b236-f98b-11ea-0444-6fe340642707
md"""

## Some settings and helper functions

"""

# ╔═╡ 79d08b08-fa8b-11ea-0b4b-01ceef03660a
# calculate sum of squared errors for different betas
begin

	betaVals = -10:0.5:10 ;# range(-10, 10, length=nPoints)
	e = zeros(size(betaVals))
	X = Vector(d[:,3])
	theY = Vector(d[:,5])
	for (i, theB) in enumerate(betaVals)
		evec = X * theB .- theY
		e[i] = transpose(evec) * evec
	end
	# e
	md"""calculate squared errors"""
end

# ╔═╡ 58c63ca8-fa8b-11ea-296d-0b7f613e8976
begin
	# plot error curve and current choice of M
    function plotError(betas, es, currentM, subplotNum)
	plot!(betas, es, lw=2, color=:black,
		label="error", subplot=subplotNum,)
	idx = betas .== currentM
	if es[idx]... ≈ minimum(es)
		theColor = :green
		theLabel = "min SSE"
	else
		theColor = :red
		theLabel = "SSE"

	end
	plot!(betas[idx], es[idx],
		marker=:circle,
		ms=10,
		color=theColor,
		subplot=subplotNum,
		xlabel="β value",
		ylabel="SSE",
		label=theLabel)
    end
	md"""a function that plots error for a given choice of m `plotError` """
end

# ╔═╡ 96844d74-fa87-11ea-0966-89bbbc8a4353
begin
    function plotWithBeta(xData, yData, yModel, m)
	l = @layout [ grid(2, 1) a{0.30w}
					 ]
	plot(xData, yData,
		xlim=(0, 160), ylim=(-5, 5),
		ylabel="fMRI response (%change)",
		label="data",
		layout=l,
	subplot=1)
	plot!(xData, yModel * m, lw=2,
		subplot=1,
		label="model * $(@sprintf "%.1f" m )")
	sticks!(xData, yData .- yModel * m,
			color=:gray,
			label="data-model",
				xlabel="Time (volumes)",
		subplot=2,
	xlim=(0, 160), ylim=(-5, 5))

	# betaVals from other cell
	# ditto for e
	plotError(betaVals, e, m, 3)
    end
	md"""a function that plots data + model with a
	particular choice of m `plotWithBeta` """
end

# ╔═╡ bfdd7d56-f98a-11ea-3064-a9033ba994a5
# plotly()
gr()

# ╔═╡ eda9e2dc-f98b-11ea-3eca-539ac938107c
nPoints = 20; # number of points in plot above (not used)

# ╔═╡ 778800f8-f98a-11ea-1b87-1108f500ad5f
begin
	t = 0:1:(nPoints - 1);
	y = weight .* t;

	plot(t,y, lw=2, color=:red,
		aspect_ratio=1,
		xlims=(0, 15), ylims=(0, 15),
		label="fit curve")
	plot!(t,1.25.*t, seriestype=:scatter, marker=:square, ms=5, color=:black, label=:"data points")

end

# ╔═╡ f6e4754c-f98b-11ea-2a42-034d82f52ca1
addNoise = 0.5 .* randn(nPoints);

# ╔═╡ Cell order:
# ╟─9bc349e8-efb0-11ea-1392-cd735828af33
# ╟─8029070e-efb0-11ea-1d0a-7f402f2e9989
# ╟─4df55a88-f98a-11ea-1c9f-cd52a4b5062f
# ╟─778800f8-f98a-11ea-1b87-1108f500ad5f
# ╟─c14bf31e-f98c-11ea-33d1-f5de8f2ad0c1
# ╟─2b2f6be4-f07f-11ea-1459-775749ca949f
# ╠═368e2640-f8e0-11ea-1ddf-b37069ba7588
# ╟─cac0339a-f8e4-11ea-0876-f91dc7db1454
# ╟─1058934e-f8e1-11ea-1896-35d75c4b6a8e
# ╟─c6cab78c-f836-11ea-2b5b-23f0006f8e93
# ╟─da4ac85c-f8e4-11ea-192e-03d5b6862d5a
# ╟─1c3b6c56-f8e2-11ea-3969-f59871025c5e
# ╟─ddf6018c-f921-11ea-0d74-c16f4f0b7426
# ╟─39d465a6-f8e2-11ea-2d71-75cbe586428c
# ╟─a1b031a4-fc08-11ea-2314-13e2e79c2ca3
# ╟─e6309486-fc08-11ea-3ade-c93a24b6866b
# ╟─9c3b0258-fc08-11ea-1922-fd1bec66b4ff
# ╠═975860d2-fc08-11ea-2def-ad1c45826911
# ╟─e7fb832e-04e0-11eb-388d-f7520b672a6e
# ╠═2b070720-04e2-11eb-3101-cbc40b17682b
# ╟─4b968440-04e2-11eb-35ce-5f050dcd4b6b
# ╟─6464f806-04e2-11eb-37cc-7363fa76c80e
# ╟─df5232e0-04e2-11eb-0c26-f19b9d9c7b14
# ╟─111aea9a-04e5-11eb-066d-4d245b04dc5e
# ╟─1e370abc-04e3-11eb-064a-21fbb9fafc7c
# ╠═1d60545a-04e4-11eb-1720-b5f04c5da239
# ╠═86d10676-04e4-11eb-2121-2f0c8945cc1a
# ╠═979732dc-04e4-11eb-3803-7302dc91e777
# ╟─a33a18f2-04e4-11eb-0587-6117c036c641
# ╟─c60d9136-1d5f-11eb-0839-d760afb98985
# ╟─c47f840a-1d5f-11eb-1e88-d3835fa021e3
# ╟─aa954778-1d5f-11eb-00bf-034690d40e94
# ╟─5d60ed36-1d5f-11eb-3e90-9315dd9bc07a
# ╟─cc5dcb6e-1d5f-11eb-0be7-57d4795af804
# ╟─cb65baa0-1d5f-11eb-346f-bdb3a85d0fe7
# ╟─ca5b1768-1d5f-11eb-0459-13e4d2de80c7
# ╟─c3899100-04e4-11eb-37aa-394d3f60bd53
# ╟─d5c0b236-f98b-11ea-0444-6fe340642707
# ╟─79d08b08-fa8b-11ea-0b4b-01ceef03660a
# ╟─58c63ca8-fa8b-11ea-296d-0b7f613e8976
# ╠═96844d74-fa87-11ea-0966-89bbbc8a4353
# ╟─bfdd7d56-f98a-11ea-3064-a9033ba994a5
# ╠═eda9e2dc-f98b-11ea-3eca-539ac938107c
# ╠═f6e4754c-f98b-11ea-2a42-034d82f52ca1


module Utils
    using JSON, ..Metrics, ..Mutation, ..Recombination, ..FunctionsCollection
    function get_algo_params()
        file_path = pwd() * "/algo_inputs.json"
        parameters = JSON.parse(read(file_path, String))
        parameters["error_function"] = select_error_function(parameters["error_function"])
        parameters["recombination_type"] = select_recombination_type(parameters["recombination_type"])
        parameters["mutation_type"] = select_mutation_type(parameters["mutation_type"])
        parameters["evaluate_function"] = select_evaluate_function(parameters["evaluate_function"])
        parameters["output_file"] = select_output_file(parameters["output_file"])
        return parameters
    end

    function select_output_file(output_file)
        letters = Set(output_file)
        if length(letters) == 1 && ' ' in letters
            return nothing
        elseif output_file == ""
            return nothing
        else
            return output_file
        end
    end

    function select_evaluate_function(evaluate_function)
        if evaluate_function[end] == '1'
            return FunctionsCollection.function1
        elseif evaluate_function[end] == '2'
            return FunctionsCollection.function2
        elseif evaluate_function[end] == '3'
            return FunctionsCollection.function3
        elseif evaluate_function[end] == '4'
            return FunctionsCollection.function4
        elseif evaluate_function[end] == '5'
            return FunctionsCollection.function5
        else
            throw(ArgumentError("Function to evaluate requested not available. Options are:
                                function1, function2, function3, function4, function5"))
        end
    end

    # Get error function based on possibilities
    function select_error_function(error_function)
        if error_function == "mean_sqr_error"
            return Metrics.mean_sqr_error
        elseif error_function == "mean_absolute_error"
            return Metrics.mean_absolute_error
        elseif error_function == "weighted_mean_abs_error"
            return Metrics.weighted_mean_abs_error
        else
            throw(ArgumentError("Error function requested not available. Options are: mean_sqr_error,
                                mean_absolute_error, weighted_mean_abs_error"))
        end
    end

    # Get recombination based on possibilities
    function select_recombination_type(recombination_type)

        if recombination_type == "intermediary_recombination"
            Recombination.intermediary_recombination
        elseif recombination_type == "discrete_recombination"
            Recombination.discrete_recombination
        else
            throw(ArgumentError("Recombination type requested not available. Options are:
                                intermediary_recombination, discrete_recombination"))
        end
    end

    # Get mutation type based on possibilities
    function select_mutation_type(mutation_type)
        if mutation_type == "uncorr_mutation_onestepsize"
            return Mutation.uncorr_mutation_onestepsize
        elseif mutation_type == "uncorr_mutation_n_stepsize"
            return Mutation.uncorr_mutation_n_stepsize
        else
            throw(ArgumentError("Mutation type requested not available. Options are:
                                uncorr_mutation_onestepsize, uncorr_mutation_n_stepsize"))
        end

    end

    # Linspace like utils
    function generate_linespace(min, max, n_values)
        step = (max-min)/n_values
        return collect(min:step:max)
    end

end


export maxima,
       minima

for (funcname, comp) in ((:maxima, :<),
                         (:minima, :>))

    @eval begin
        function ($funcname)(x::AbstractVector,
                        w::Integer=1,
                        includebounds::Bool=false)
            w > 0 || throw(ArgumentError("window cannot be negative"))
            xlen = length(x)
            idxs = Int[]

            # There can't be more than one peak every `w` elements
            maxN = includebounds ? cld(xlen,2*w) : fld(xlen,2*w)
            N = 0

            ww = [ collect(w:-1:1); collect(-w:-1) ]

            sizehint!(idxs,maxN)

            i = 1

            if includebounds
                @inbounds while i <= w
                    peak = true
                    for j in max(i-w, 1):(i+w)
                        if ($comp)(x[i], x[j])
                            peak &= false
                            break # No reason to continue checking if the rest of the elements qualify
                        end
                    end

                    if peak
                        push!(idxs,i)
                        N += 1
                        i += w # There can't be another peak for at least `w` more elements
                    end
                    i += 1
                end
            end

            i = 1+w
            maxI = xlen-w

            @inbounds while i <= maxI
                peak = true
                for j in ww # For all elements within the window
                    if ($comp)(x[i], x[i+j])
                        peak &= false
                        break # No reason to continue checking if the rest of the elements qualify
                    end
                end

                if peak
                    push!(idxs,i)
                    N += 1
                    i += w # There can't be another peak for at least `w` more elements
                end
                i += 1
            end

            if includebounds
                @inbounds while i <= xlen
                    peak = true
                    for j in (i-w):min((i+w),xlen)
                        if ($comp)(x[i], x[j])
                            peak &= false
                            break # No reason to continue checking if the rest of the elements qualify
                        end
                    end

                    if peak
                        push!(idxs,i)
                        N += 1
                        i += w # There can't be another peak for at least `w` more elements
                    end
                    i += 1
                end
            end

            resize!(idxs,N)
            sizehint!(idxs,N)

            return idxs
        end
    end
end

@doc """
    maxima(x[, w=1, includebounds=false])

Find the local maxima of `x`.

`w` sets the minimum allowed distance between maxima. If `includebounds` is `true`, maxima
are allowed to be less than `w` away from the bounds of `x`, otherwise, maxima may not be
located any closer to the ends of `x` than `w`.
"""
maxima

@doc """
    minima(x[, w=1, includebounds=false])

Find the local minima of `x`.

`w` sets the minimum allowed distance between minima. If `includebounds` is `true`, minima
are allowed to be less than `w` away from the bounds of `x`, otherwise, minima may not be
located any closer to the ends of `x` than `w`.
"""
minima


using Clustering, DataFrames

n = 1_000

srand(1)

data = randn(n, 2)

centers = [0.0 0.0;
           5.0 5.0;
           10.0 0.0;
           15.0 -5.0;]

assignments = zeros(Int64, n)

for i = 1:n
    assignments[i] = rand(1:4)
    data[i, :] += centers[assignments[i], :]
end

results = dp_means(data, 50.0)

df = DataFrame()
df["Label"] = assignments
df["Cluster"] = results.assignments
by(df, ["Label", "Cluster"], nrow)

# Should really look at active clusters, not created clusters


import Base.vcat, Base.hcat, Base.hvcat, Base.start
export vcat, hcat, hvcat

function vcat(args::AffineExpr...)
  num_cols = args[1].size[2]
  num_rows = 0
  # Calculate numbers of rows after vertical concatenation
  for arg in args
    if arg.size[2] != num_cols
      error("All arguments must have same number of columns")
    end
    num_rows += arg.size[1]
  end

  vars_to_coeffs_map = Dict{UInt64, Constant}()
  vec_sz = num_rows * num_cols
  constant = Constant(spzeros(vec_sz, 1))
  index_start = 0

  for i = 1:length(args)
    for (v, c) in args[i].vars_to_coeffs_map
      if !haskey(vars_to_coeffs_map, v)
        vars_to_coeffs_map[v] = Constant(spzeros(vec_sz, c.size[2]))
      end
      for j in 1 : num_cols
        index = index_start + (j - 1) * num_rows
        vars_to_coeffs_map[v].value[index + 1 : index + args[i].size[1], :] =
            c.value[(j - 1) * args[i].size[1] + 1 : j * args[i].size[1], :]
      end
    end

    for j in 1 : num_cols
      index = index_start + (j - 1) * num_rows
      constant.value[index + 1 : index + args[i].size[1]] =
          args[i].constant.value[(j - 1) * args[i].size[1] + 1 : j * args[i].size[1]]
    end

    index_start += args[i].size[1]
  end

  this = AffineExpr(:vcat, tuple(args...), vars_to_coeffs_map, constant, (num_rows, num_cols))
  this.evaluate = ()->Base.vcat([arg.evaluate() for arg in args]...)
  return this
end

function hcat(args::AffineExpr...)
  num_rows = args[1].size[1]
  num_cols = 0

  for arg in args
    if arg.size[1] != num_rows
      error("All arguments must have same number of rows")
    end
    num_cols += arg.size[2]
  end

  vars_to_coeffs_map = Dict{UInt64, Constant}()
  vec_sz = num_rows * num_cols
  constant = Constant(spzeros(vec_sz, 1))
  index = 0

  for i = 1:length(args)
    arg_vec_sz = args[i].size[1] * args[i].size[2]

    for (v, c) in args[i].vars_to_coeffs_map
      if !haskey(vars_to_coeffs_map, v)
        vars_to_coeffs_map[v] = Constant(spzeros(vec_sz, c.size[2]))
      end
      vars_to_coeffs_map[v].value[index + 1 : index + arg_vec_sz, :] = c.value
    end
    constant.value[index + 1 : index + arg_vec_sz] = args[i].constant.value

    index += arg_vec_sz
  end

  this = AffineExpr(:hcat, tuple(args...), vars_to_coeffs_map, constant, (num_rows, num_cols))
  this.evaluate = ()->Base.hcat([arg.evaluate() for arg in args]...)
  return this
end

function hvcat(rows::Tuple{Vararg{Int64}}, args::AffineExpr...)
  row_exprs = AffineExpr[]
  index = 0
  for row_size in rows
    push!(row_exprs, hcat(args[index + 1 : index + row_size]...))
    index += row_size
  end
  return vcat(row_exprs...)
end


using Makie
using MultipleViewGeometry, Test, Random
using MultipleViewGeometry.ModuleCostFunction
using MultipleViewGeometry.ModuleTypes
using MultipleViewGeometry.ModuleConstraints
using MultipleViewGeometry.ModuleConstruct
using MultipleViewGeometry.ModuleDraw
using MultipleViewGeometry.ModuleMove
using LinearAlgebra
using StaticArrays
using GeometryTypes
using Test

𝒳 = [Point3D(x,y,rand(50:100)) for x = -100:5:100 for y = -100:5:100]
𝒳 = 𝒳[1:50:end]
X = reshape(reinterpret(Float64,𝒳),(3,length(𝒳)))

# Specify the coordinate systems of the world, the camera frame and the picture
# plane.
world_basis = (Vec(1.0, 0.0, 0.0), Vec(0.0, 1.0, 0.0), Vec(0.0, 0.0, 1.0))
camera_basis = (Point(0.0, 0.0, 0.0), Vec(-1.0, 0.0, 0.0), Vec(0.0, -1.0, 0.0), Vec(0.0, 0.0, 1.0))
picture_basis = (Point(0.0, 0.0), Vec(-1.0, 0.0), Vec(0.0, -1.0))

# The focal length for both cameras is one.
f = 1
image_width = 640 / 10
image_height = 480 / 10

camera₁ = Pinhole(image_width, image_height, f, camera_basis..., picture_basis...)
camera₂ = Pinhole(image_width, image_height, f, camera_basis..., picture_basis...)

# Rotate and translate camera one.
𝐑₁ = Matrix{Float64}(I,3,3)
𝐭₁ = [-50.0, -2.0, 0.0]
relocate!(camera₁, 𝐑₁, 𝐭₁)

# Rotate and translate camera two.
𝐑₂ = Matrix{Float64}(I,3,3)
𝐭₂ = [50.0, 2.0, 0.0]
relocate!(camera₂, 𝐑₂, 𝐭₂)

scale = 20.0f0
x = Vec3f0(0); baselen = 0.2f0 * scale ; dirlen = 1f0 * scale
# create an array of differently colored boxes in the direction of the 3 axes
rectangles = [
    (HyperRectangle(Vec3f0(x), Vec3f0(dirlen, baselen, baselen)), RGBAf0(1,0,0,1)),
    (HyperRectangle(Vec3f0(x), Vec3f0(baselen, dirlen, baselen)), RGBAf0(0,1,0,1)),
    (HyperRectangle(Vec3f0(x), Vec3f0(baselen, baselen, dirlen)), RGBAf0(0,0,1,1))
]
meshes = map(GLNormalMesh, rectangles)

scene = mesh(merge(meshes))
scatter!(scene, X[1,:],X[2,:], X[3,:], markersize = 3, color = :red)
draw!(camera₁, scene)
draw!(camera₂, scene)

𝐑₁′, 𝐭₁′ = ascertain_pose(camera₁, world_basis... )
𝐊₁′ = obtain_intrinsics(camera₁, CartesianSystem())
𝐑₂′, 𝐭₂′ = ascertain_pose(camera₂, world_basis... )
𝐊₂′ = obtain_intrinsics(camera₂, CartesianSystem())

# Camera projection matrices.
𝐏₁ = construct(ProjectionMatrix(),𝐊₁′,𝐑₁′,𝐭₁′)
𝐏₂ = construct(ProjectionMatrix(),𝐊₂′,𝐑₂′,𝐭₂′)

# Set of corresponding points.
ℳ = project(camera₁,𝐏₁,𝒳)
ℳʹ = project(camera₂,𝐏₂,𝒳)

# Estimate of the fundamental matrix and the true fundamental matrix.
𝐅 = estimate(FundamentalMatrix(), DirectLinearTransform(), (ℳ, ℳʹ))

#𝐅ₜ = construct(FundamentalMatrix(), 𝐏₁, 𝐏₂)
𝐅ₜ = construct(FundamentalMatrix(), 𝐊₁′,𝐑₁′,-𝐭₁′,𝐊₂′,𝐑₂′,-𝐭₂′)

# Ensure the estimated and true matrix have the same scale and sign.
𝐅 = 𝐅 / norm(𝐅)
𝐅 = 𝐅 / sign(𝐅[3,1])
𝐅ₜ = 𝐅ₜ / norm(𝐅ₜ)
𝐅ₜ = 𝐅ₜ / sign(𝐅ₜ[3,1])

@test 𝐅 ≈ 𝐅ₜ

# Check that the fundamental matrix satisfies the corresponding point equation.
npts = length(ℳ)
residual = zeros(Float64,npts,1)
for correspondence in zip(1:length(ℳ),ℳ, ℳʹ)
    i, m , mʹ = correspondence
    𝐦  = hom(m)
    𝐦ʹ = hom(mʹ)
    residual[i] = (𝐦ʹ'*𝐅*𝐦)
end

@test isapprox(sum(residual), 0.0; atol = 1e-7)

# Test the Fundamental Numerical Scheme on the Fundamental matrix problem.
Λ₁ =  [SMatrix{3,3}(Matrix(Diagonal([1.0,1.0,0.0]))) for i = 1:length(ℳ)]
Λ₂ =  [SMatrix{3,3}(Matrix(Diagonal([1.0,1.0,0.0]))) for i = 1:length(ℳ)]
𝐅₀ = estimate(FundamentalMatrix(),DirectLinearTransform(),  (ℳ, ℳʹ))
𝐅 = estimate(FundamentalMatrix(),
              FundamentalNumericalScheme(vec(𝐅₀), 5, 1e-10),
               (Λ₁,Λ₂), (ℳ, ℳʹ))

# Ensure the estimated and true matrix have the same scale and sign.
𝐅 = 𝐅 / norm(𝐅)
𝐅 = 𝐅 / sign(𝐅[3,1])

@test 𝐅 ≈ 𝐅ₜ

# Test the Bundle Adjustment estimator on the Fundamental matrix problem.
𝐅, lsqFit = estimate(FundamentalMatrix(),
                      BundleAdjustment(vec(𝐅₀), 5, 1e-10),
                        (ℳ, ℳʹ))
𝐅 = 𝐅 / norm(𝐅)
𝐅 = 𝐅 / sign(𝐅[3,1])
@test 𝐅 ≈ 𝐅ₜ


"""
Canonical polyadic N-mode tensor decomposition (CANDECOMP/PARAFAC).
"""
immutable CANDECOMP{T<:Number, N} <: TensorDecomposition{T, N}
  factors::NTuple{N, Matrix{T}} # rank X i-th tensor mode factor matrices
  lambdas::Vector{T}
  props::Dict{Symbol, Any}      # extra properties

  CANDECOMP{T, N}(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) where {T<:Number, N} =
    new(factors, lambdas, Dict{Symbol, Any}())

  (::Type{CANDECOMP}){T, N}(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T}) =
    CANDECOMP{T, N}(factors, lambdas)
end

"""
Returns CANDECOMP rank.
"""
Base.rank(decomp::CANDECOMP) = length(decomp.lambdas)

"""
Re-composes the tensor from CANDECOMP decomposition.
"""
@generated function compose!{T,N}(dest::Array{T,N}, factors::NTuple{N, Matrix{T}}, lambdas::Vector{T})
  quote
    @nloops $N i dest begin
        elm = zero(T)
        # FIXME traversal of factors is not very efficient, would be better to have them transposed
        for j in 1:length(lambdas)
            elm += lambdas[j] * (*(@ntuple($N, k -> factors[k][i_k, j])...))
        end
        @nref($N, dest, i) = elm
    end
    dest
  end
end

@generated function compose{T,N}(factors::NTuple{N, Matrix{T}}, lambdas::Vector{T})
  quote
    compose!(Array{T}(@ntuple $N i -> size(factors[i], 1)), factors, lambdas)
  end
end

compose(decomp::CANDECOMP) = compose(decomp.factors, decomp.lambdas)

compose!{T,N}(dest::Array{T,N}, decomp::CANDECOMP{T,N}) = compose!(dest, decomp.factors, decomp.lambdas)

"""
Calculates canonical polyadic tensor decomposition (CANDECOMP/PARAFAC).

Returns:
  `CANDECOMP` object
"""
function candecomp{T,N}(tnsr::StridedArray{T,N},
                   r::Integer,
                   initial_guess::NTuple{N, Matrix{T}};
                   method::Symbol=:ALS,
                   tol::Float64=1e-5,
                   maxiter::Integer=100,
                   compute_error::Bool=false,
                   verbose::Bool=true)

    # _check_tensor(tnsr, r) # decompositions with some of the sizes equal to the original tensor should be allowed
    verbose && info("initializing factor matrices...")
    all([(size(tnsr, i), r) == size(initial_guess[i]) for i in 1:N]) || throw(ArgumentError("dimension of initial guess does not match input tensor."))
    factors = collect(Matrix{T}, initial_guess)
    verbose && info("applying candecomp $method method...")
    res = _candecomp(Val{method}, tnsr, r, factors, tol, maxiter, verbose)
    if compute_error
      _set_rel_residue(res, tnsr)
    end
    return res
end

"Stub for non-implemented CANDECOMP algorithms."
_candecomp{T,N,S<:Symbol}(
  method::Val{S},
  tnsr::StridedArray{T,N},
  r::Integer,
  factors::Vector{Matrix{Float64}},
  tol::Float64,
  maxiter::Integer,
  verbose::Bool) = throw(ArgumentError("Unknown CANDECOMP method: $method"))

"""
Computes CANDECOMP by ALS (Alternating Least Squares) method.
"""
function _candecomp{T,N}(
    method::Type{Val{:ALS}},
    tnsr::StridedArray{T,N},
    r::Integer,
    factors::Vector{Matrix{Float64}},
    tol::Float64,
    maxiter::Integer,
    verbose::Bool)

    gram = [F'F for F in factors]
    tnsr_norm = vecnorm(tnsr)
    tnsr_flat = _col_unfold(tnsr, N)
    tnsr_size = size(tnsr)
    niters = 0
    converged = false
    resid = tnsr_norm
    lbds = Matrix{Float64}(1, r)
    V = Matrix{Float64}(div(length(tnsr), minimum(tnsr_size)), r)
    while !converged && niters < maxiter
        VB = 0
        @inbounds for i in 1:N
            idx = [N:-1:i + 1; i - 1:-1:1]
            VB = prod(k -> tnsr_size[k], idx)
            V[1:VB, :] = reduce(khatrirao, factors[idx])
            factors[i] = _row_unfold(tnsr, i) * V[1:VB, :] / reduce(((x, y) -> broadcast(*, x, y)), gram[idx])
            sum!(lbds, abs.(factors[i]))
            factors[i] ./= lbds
            At_mul_B!(gram[i], factors[i], factors[i])
        end
        resid_old = resid
        resid = vecnorm(V[1:VB, :] * broadcast(*, factors[N], lbds)' - tnsr_flat)
        converged = abs.(resid - resid_old) < tol * resid_old
        niters += 1
    end

    verbose && _iter_status(converged, niters, maxiter)

    return CANDECOMP((factors...), squeeze(lbds, 1))
end

"""
Computes CANDECOMP by SGSD (Simultaneous Generalized Schur Decomposition) method.
"""
function _candecomp{T,N}(
    method::Type{Val{:SGSD}},
    tnsr::StridedArray{T,N},
    r::Int,
    factors::Vector{Matrix{Float64}},
    tol::Float64,
    maxiter::Integer,
    verbose::Bool)

    ndims(tnsr) == 3 || throw(ArgumentError("This algorithm only applies to 3-mode tensors."))
    (n1, n2, n3) = size(tnsr)
    IB = [min(n1 - 1, r), (n2 == r) ? 2 : 1]
    Q = qr(factors[1], thin=false)[1]
    Z = qr(flipdim(factors[2], 2), thin=false)[1]
    q = Array{Float64}(n1, n1)
    z = Array{Float64}(n2, n2)

    R = zeros(size(tnsr))
    @tensor R[4,5,3] = tnsr[1,2,3] * Q[1,4] * Z[2,5]

    res = vecnorm(tnsr)
    converged = false
    niters = 0
    @inbounds while !converged && niters < maxiter
        q = eye(n1)
        z = eye(n2)

        for i in 1:IB[1]
            q[:, i:n1] *= svd(q[:, i:n1]' * view(R, :, n2 - r + i, :))[1]
        end

        @tensor R[4,2,3] = R[1,2,3] * q[1,4]

        for i in r:-1:IB[2]
            z[:, 1:n2 - r + i] *= flipdim(svd(view(R, i, :, :)' * z[:, 1:n2 - r + i])[3], 2)
        end
        Q *= q
        Z *= z

        @tensor R[4,5,3] = tnsr[1,2,3] * Q[1,4] * Z[2,5]

        res_old = res
        res = vecnorm(tril(squeeze(sum(R .^ 2, 3), 3), n2 - r - 1))
        converged = abs.(res - res_old) < tol
        niters += 1
    end

    verbose && _iter_status(converged, niters, maxiter)

    R = R[1:r, n2-r+1:n2, :]
    M = cat(3, eye(r, r), eye(r, r))
    @inbounds for i in (r - 1):-1:1, j = (i + 1):r
        if i + 1 < j
            d = (i + 1):(j - 1)
            println(d)

            M[i, j, :] = hcat(r[j, j, :][:], r[i, i, :][:]) \ (r[i, j, :][:] - mapslices(R3 -> sum(M[i, d, 1] * broadcast(*, diag(R3)[d], M[d, j, 2])), R, [1, 2])[:])
        end
    end

    factors[1] = Q[:, 1:r] * M[:, :, 1]
    factors[2] = Z[:, n2 - r + 1:n2] * M[:, :, 2]'
    factors[3] = _row_unfold(tnsr, 3) * khatrirao(factors[2], factors[1]) / broadcast(*, factors[2]'factors[2], factors[1]'factors[1])

    lbds = ones(1, r)
    for i in 1:3
        lbd = mapslices(vecnorm, factors[i], 1)
        factors[i] ./= lbd
        lbds = broadcast(*, lbds, lbd)
    end

    return CANDECOMP((factors...), squeeze(lbds, 1))
end

"""
Codes of implemented CANDECOMP methods.
"""
const CANDECOMP_methods = Symbol[t.parameters[1] for t in filter(t -> isa(t, DataType),
                                                                 [Base.unwrap_unionall(m.sig).parameters[2].parameters[1] for m in methods(_candecomp)])] |> Set{Symbol}


## utility functions ##

function stepup!{T<:AbstractFloat}(sortedPValues::Vector{T}, multiplier::Function, n::Integer = length(sortedPValues))
    sortedPValues[n] *= multiplier(0, n)
    for i in 1:(n-1)
        sortedPValues[n-i] = min(sortedPValues[n-i+1], sortedPValues[n-i] * multiplier(i, n))
    end
    return sortedPValues
end

# multiplier stepdown
function stepdown!{T<:AbstractFloat}(sortedPValues::Vector{T}, multiplier::Function, n::Integer = length(sortedPValues))
  stepfun(p::T, i::Int, n::Int) = p*multiplier(i,n)
  general_stepdown!(sortedPValues, stepfun, n)
  return sortedPValues
end

function general_stepdown!{T<:AbstractFloat}(sortedPValues::Vector{T}, stepfun::Function, n::Integer = length(sortedPValues))
    sortedPValues[1] = stepfun(sortedPValues[1], 1, n)
    for i in 2:n
        sortedPValues[i] = max(sortedPValues[i-1], stepfun(sortedPValues[i],i, n))
    end
    return sortedPValues
end

function reorder{T<:Number}(values::Vector{T})
    newOrder = sortperm(values)
    oldOrder = sortperm(newOrder)
    return newOrder, oldOrder
end

function validPValues{T<:AbstractFloat}(x::Vector{T})
    if !isin(x)
        throw(DomainError())
    end
end

function isin(x::Real, lower::Real = 0., upper::Real = 1.)
    x >= lower && x <= upper
end

function isin{T<:Real}(x::Vector{T}, lower::Real = 0., upper::Real = 1.)
    ex = extrema(x)
    ex[1] >= lower && ex[2] <= upper
end

function isotonicregression(y::Array{Float64,1},w::Array{Float64,1})
  #todo: ignore zero weights
  y=copy(y)
  w=copy(w)
  m = length(y)
  cnts = ones(Int64,m)
  i = 2
  # ... not most efficient way but could be fun to (ab)use iterator protocol
  while (!done(y,i))
    if y[i]<y[i-1]
      y[i-1]=(w[i]*y[i]+w[i-1]*y[i-1])/(w[i]+w[i-1])
      w[i-1]=w[i]+w[i-1]
      cnts[i-1] += cnts[i]
      deleteat!(y,i)
      deleteat!(w,i)
      deleteat!(cnts,i)
      i = max(i-2,1)
    end
    i += 1
  end
  yisotonic = vcat([y[idx]*ones(Float64,cnt) for (idx,cnt) in enumerate(cnts)]...)
end

function isotonicregression(y::Array{Float64,1})
  isotonicregression(y, ones(Float64, length(y)))
end

function isotonic_regression(y::Vector{Float64}, weights::Vector{Float64})
    n = length(y)
    if n <= 1
        return y
    end
    if n != length(weights)
        throw(DimensionMismatch("Lengths of values and weights mismatch"))
    end
    @inbounds begin
        n -= 1
        while true
            i = 1
            is_pooled = false
            while i <= n
                k = i
                while k <= n && y[k] >= y[k+1]
                    k += 1
                end
                if y[i] != y[k]
                    numerator = 0.0
                    denominator = 0.0
                    for j in i:k
                        numerator += y[j] * weights[j]
                        denominator += weights[j]
                    end
                    m = numerator / denominator
                    for j in i:k
                        y[j] = m
                    end
                    is_pooled = true
                end
                i = k + 1
            end
            if !is_pooled
               break
            end
        end
    end
    return y
end

function isotonic_regression(y::Vector{Float64})
    isotonic_regression(y, ones(y))
end

function grenander(pv::Vector{Float64})
    pv_sorted = sort(pv)
    ## ecdf that handles duplicated values
    pv_sorted_unique, counts = rle(pv_sorted)
    ecdf_value = cumsum(counts)
    ecdf_value = ecdf_value ./ ecdf_value[end]

    Δx = diff(pv_sorted_unique)
    Δy = diff(ecdf_value)

    f = Δy ./ Δx
    f = -isotonic_regression(-f, Δx)
    F  = ecdf_value[1] + vcat(0, cumsum(f .* Δx))
    f = push!(f, f[end])

    return pv_sorted_unique, f, F
end


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

"""
    DensityRatioWeighter(tdata, [vars]; [options])

Density ratio weights based on empirical distribution of
variables in target data `tdata`. Default to all variables.

## Optional parameters

* `estimator` - Density ratio estimator (default to `LSIF()`)
* `optlib`    - Optimization library (default to `default_optlib(estimator)`)

### Notes

Estimators from `DensityRatioEstimation.jl` are supported.
"""
struct DensityRatioWeighter <: AbstractWeighter
  tdata
  vars
  dre
  optlib
end

function DensityRatioWeighter(tdata, vars=nothing; estimator=LSIF(),
                              optlib=default_optlib(estimator))
  validvars = collect(name.(variables(tdata)))
  wvars = isnothing(vars) ? validvars : vars
  @assert wvars ⊆ validvars "invalid variables ($wvars) for spatial data"
  DensityRatioWeighter(tdata, wvars, estimator, optlib)
end

function weight(sdata, weighter::DensityRatioWeighter)
  # retrieve method parameters
  tdata  = weighter.tdata
  vars   = weighter.vars
  dre    = weighter.dre
  optlib = weighter.optlib

  @assert vars ⊆ name.(variables(sdata)) "invalid variables ($vars) for spatial data"

  # numerator and denominator samples
  Ωnu = view(tdata, vars)
  Ωde = view(sdata, vars)
  xnu = collect(Tables.rows(values(Ωnu)))
  xde = collect(Tables.rows(values(Ωde)))

  # perform denstiy ratio estimation
  ratios = densratio(xnu, xde, dre, optlib=optlib)

  SpatialWeights(domain(sdata), ratios)
end


# Include packages
using Interpolations
using Printf: @printf, @sprintf
using Polynomials: polyfit, polyval
using DSP: welch_pgram
using DelimitedFiles: readdlm
using Dates

####################
# function
####################
function loadcsvsample()
    # define the filepath & filename
    fdir = "./data"
    fname = "dat_climate.csv"
    dataorg = readdlm(joinpath(fdir,fname), ',', skipstart=1)
    return dataorg
end
####################

####################
## main
####################

# read csv data
dataorg = loadcsvsample()
torg = DateTime.(dataorg[:,1], "yyyy/mm/dd HH:MM")
tsecorg = convert.(Float64, Dates.value.((torg-torg[1])/1000)) # ms to s
V = convert.(Float64, dataorg[:,2])
dataorg = nothing

# Interpolate
t = collect(torg[1]:Day(1):torg[end])
tsec = convert.(Float64, Dates.value.((t-t[1])/1000)) # ms to s
nt = size(t,1)
itp = interpolate((tsecorg, ), V, Gridded(Linear()))
#itp = interpolate((tsecorg, ), V, Gridded(Constant()))
Vint = itp(tsec)

# regression
lin_p = polyfit(tsec, Vint, 1)

# power spectral
# (注)360*2は結果を合わせにいった値のため根拠なし
pdg = welch_pgram(Vint, 360*2, onesided=true; fs=1.0)
days = 1.0./convert.(Float64, pdg.freq)
PSD = pdg.power
maxval, Tc = findmax(PSD[2:end])

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

####################
## plot
####################

# directory where figures are printed
figdir="./fig"
if !isdir(figdir); mkdir(figdir); end

using Plots
pyplot()
#gr()
#plotlyjs()

# ticklabel format for Figure 1
xt = collect(DateTime(2012):Year(1):DateTime(2017))
xtl = Dates.format.(xt, DateFormat("Y"))
xt = Dates.value.(xt)

########
## Figure 1
#  trend
# constant interval
plot(t, Vint, line=(:solid, 1), lab="Uniformed", size=(800,600),
     tickfont=12, legend=:topleft, legendfont=14,
     xticks=(xt, xtl), xlims=(xt[1]-(0.5*365*86400*1e3), xt[end]),
     xlabel="Year", guidefont=12,
     ylims=(62,72),
     )
# plot!(torg, V, line=(:dash, 1), lab="Raw")

# regression
plot!(t, polyval(lin_p, tsec), lab="Regression 1")

# save
savefig(joinpath(figdir,"uniformed_data.png"))
########

########
## Figure 2
#  periods
plot(days[2:end], PSD[2:end], xscale=:log10, tickfont=12, size=(800, 600),
     lab="Power", legend=:topleft, legendfont=14,
     xlabel="Days", guidefont=14,
     xlims=(1e+00, 1e3),
     )
# natral frequency
scatter!([days[Tc+1]], [PSD[Tc+1]], ms=10.,
         lab="Period: "*@sprintf("%.0f",days[Tc+1])*"days",
         )
         #ylims=(0., 550.),
# save
savefig(joinpath(figdir,"trend.png"))
########

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


"""
Set of functions for generating signals.
"""

function random_walk(process_noise,
                    measurement_noise,
                    mean_state0,
                    precision_state0;
                    time_horizon=100)
    "Generate data according to a random walk"

    # Dimensionality
    d = size(process_noise, 1)

    # Preallocate
    y = zeros(d, time_horizon)
    x = zeros(d, time_horizon+1)

    # Initialize state
    x[:, 1] = sqrt(inv(precision_state0))*randn(d) + mean_state0

    for t = 1:time_horizon

        # Evolve state
        x[:, t+1] = sqrt(process_noise)*randn(d) + x[:, t]

        # Observe
        y[:, t] = sqrt(measurement_noise)*randn(d) + x[:, t+1]

    end
    return y, x
end

function LGDS(A, C, Q, R, x0; T=100)
    """
    Generate data according to a linear Gaussian dynamical system.

    Note: assumes univariate observations
    """

    # Dimensionality
    Dx = size(A,1)
    Dy = size(C,1)

    # Check for Dy = 1
    if Dy != 1; error("Assumes univariate observations"); end

    # Preallocate
    x = zeros(Dx,T)
    y = zeros(T,)

    # Initialize state
    x_kmin1 = x0

    for k = 1:T

        # Evolve state
        x[:,k] = A*x_kmin1 .+ cholesky(Q).U*randn(Dx,)

        # Observe
        y[k,:] = C*x[:,k] .+ sqrt(R)*randn(Dy,)[1]

        # Update previous state
        x_kmin1 = x[:,k]

    end
    return y, x
end

function NLGDS(transition_function,
               emission_function,
               process_noise,
               measurement_noise,
               state0_params;
               time_horizon=100)
    "Generate data according to a nonlinear Gaussian dynamical system"

    # Dimensionality
    d = size(process_noise, 1)

    # Preallocate
    x = zeros(d, time_horizon+1)
    y = zeros(d, time_horizon)

    # Prior parameters
    mean0, var0 = state0_params

    # Initialize state
    x[:, 1] = sqrt(var0)*randn(d) .+ mean0

    for t = 1:time_horizon

        # Evolve state
        x[:, t+1] = sqrt(process_noise)*randn(d) .+ transition_function(x[:,t])

        # Observe
        y[:, t] = sqrt(measurement_noise)*randn(d) .+ emission_function(x[:,t+1])

    end
    return y, x
end


using Printf  # for Julia version 1.0+

function isnarcissist(n, b=10)
    -1 < n || return false
    d = digits(n, base=b)
    m = length(d)
    n == mapreduce((x)->x^m, +, d)
end

function findnarcissist(verbose=false)
    goal = 25
    ncnt = 0
    verbose && println("Finding the first ", goal, " Narcissistic numbers:")
    for i in 0:typemax(1)
        isnarcissist(i) || continue
        ncnt += 1
        verbose && println(@sprintf "    %2d %7d" ncnt i)
        ncnt < goal || break
    end
end

findnarcissist()
@time findnarcissist(true)


using jInv.Utils
using Base.Test
using jInv.ForwardShare
using jInv.Mesh
using OED

n = 100
nrhs = 2
# A = abs(sprandn(N,n,.2))
A = sdiag(1./linspace(1,n,n).^4)
A = A[1:2:end,:]
N = size(A,1)
x = rand(n,nrhs); x[4:7,:] = 1;
D = A*x  + 1e-1* randn(N,nrhs)/1e3
OEDparamA = OEDTypeAParam(A,D,zeros(0,0))

# setup OED type B param
r = 20;
l = 8;
Mimg  = getRegularMesh([-1 1 -1 1.],[10,10])
Mgrid = getRegularMesh([-1.2 2.4 -1.6 1.4],[20,20])
OEDparamB = getOEDTypeBParam(r,l,x,Mimg,Mgrid,noiseLevel=0.)

OEDs = (OEDparamA, OEDparamB)

println("--- solve inequality constrained problem ---")
bounds = [.3 .9]
L = .1*speye(n)
Ce = ones(1,n)
ce = Ce*x
Ci = [speye(n);-speye(n)]
ci = [bounds[1]*ones(n,nrhs); -bounds[2]*ones(n,nrhs)]
pForBox    = getIneqconstrainedOEDParam(OEDparamA, L,Ce,ce,Ci,ci,maxIterIPQP=40,outIPQP=-1)
pForNonNeg = getIneqconstrainedOEDParam(OEDparamA, L,[],[],speye(n),0*ones(n,nrhs),maxIterIPQP=40,outIPQP=-1)

pFors = (pForBox, pForNonNeg)
for j=1:length(OEDs)
    for k=1:length(pFors)
		pFors[k].Design = OEDs[j]
		display(pFors[k])

		np  =  getNumberOfDesignParameters(OEDs[j])
		p0      = rand(np)
		xr, = getData(p0,pFors[k])
		@test norm(pFors[k].Ce*xr-pFors[k].ce) < 1e-12
		@test all(pFors[k].Ci*xr.>= pFors[k].ci)
		println("rel.error: $(norm(xr-x)/norm(x))")

		println("--- test derivative of LS solution w.r.t. design ---")
		checkDerivative(rand(np),pFors[k])

		println("--- adjoint test ---")
		(nw,nv) = getSensMatSize(pFors[k])
		v = randn(nv)
		w = randn(nw)
		t1 = dot(w,getSensMatVec(v,p0,pFors[k]))
		t2 = dot(v,getSensTMatVec(w,p0,pFors[k]))
		@test abs(t1-t2)/abs(t1) < 1e-10
	end
end


using jInv.Mesh
using Test

# regular mesh
domain = [1.2 2.4 0.0 1.1 0.0 2.3]
n      = [3 4 5]
M3D      = getRegularMesh(domain,n)
display(M3D)
M2D      = getRegularMesh(domain[1:4],n[1:2])
display(M2D)

# tensor mesh
h1 = rand(4)
h2 = rand(5)
h3 = rand(6)
MT = getTensorMesh3D(h1,h2,h3)
display(MT)


# interface to C++ library "figtree" from
#   Vlad I. Morariu, Balaji Vasan Srinivasan, Vikas C. Raykar,
#   Ramani Duraiswami, and Larry S. Davis. Automatic online tuning for
#   fast Gaussian summation. Advances in Neural Information Processing
#   Systems (NIPS), 2008.
# and website that hosts code is
#   http://www.umiacs.umd.edu/~morariu/figtree/

using LinearAlgebra, Plots
gr(size=(900,300))

#include("../../src/part/fgt.jl")  # loads in functions to be tests

function fgtTest()
    d = 2                           # dimension of data
    M = 1_000                      # number of targets
    N = 2
    #N = 1_000                      # number of sources
    h = 0.8                         # bandwidth
    ε = 1e-2                        # max tolerance
    x = [-0.5,0.5]
    #x = rand(N*d)                   # source vectors, reshape'd to 1D array
    y = rand(M*d)                   # target vectors, reshape'd to 1D array
    W = 1                           # number of sources being evaluated
    q=[0.5,1.0]
    #q = rand(N*W)                   # weights for sources
    v1 = rand(M)                     # output values array
    v2 = rand(M)                    # output 2
    #v3 = rand(M)                    # output 3

    temp = norm(v1-v2)

    fgt!(v1,d,M,N,h,ε,x,y,q,W)
    dgt!(v1,d,M,N,h,ε,x,y,q,W)
    #mydgt!(v3,d,M,N,h,ε,x,y,q,W)

    println("fast (C++) transform: ")
    @time fgt!(v1,d,M,N,h,ε,x,y,q,W)
    println("slow (C++) transform: ")
    @time dgt!(v2,d,M,N,h,ε,x,y,q,W)
    #println("slow (Julia) transform: ")
    #@time mydgt!(v3,d,M,N,h,ε,x,y,q,W)
    println()

    println("(check that two algorithms give same result)")
    println("initial norm: ",temp)
    println("later norm:   ",norm(v1-v2))
end

function fgt1Dtest()
    d = 1                           # dimension of data
    M = 10_000                      # number of targets
    #N = 2
    N = 10_000                      # number of sources
    h = 0.2                         # bandwidth
    ε = 1e-2                        # max tolerance
    #x = [-0.5,0.5]
    x = rand(N*d)                   # source vectors, reshape'd to 1D array
    y = collect(range(-2,stop=2,length=M))                   # target vectors, reshape'd to 1D array
    W = 1                           # number of sources being evaluated
    #q=[0.5,1.0]
    q = rand(N*W)                   # weights for sources
    v1 = rand(M)                     # output values array
    v2 = rand(M)                    # output 2
    v3 = rand(M)                    # output 3

    @time fgt!(v1,d,M,N,h,ε,x,y,q,W)
    @time dgt!(v2,d,M,N,h,ε,x,y,q,W)
    @time mydgt!(v3,d,M,N,h,ε,x,y,q,W)

    p1 = plot(y,v1,label="fgt (c++)")
    p2 = plot(y,v2,label="dgt (c++)")
    p3 = plot(y,v3,label="mydgt (julia)")

    l = @layout [a a a]
    p = plot(p1,p2,p3,layout = l)
    gui(p)
end


# TODO: better handling for nonlinear functions than an enumeration. Based on Convex.jl?
@enum NonlinearFunction begin
    SquareRoot
    Square
end

"""
    struct NonlinearCombinatorialInstance <: CombinatorialInstance

A nonlinear combinatorial-optimisation problem. Such a problem is built on top
of a linear problem, `combinatorial_structure`, providing all the constants
required for the problem (e.g., a graph); rewards from this object are ignored.
The optimisation sense is given by this object (i.e., minimise or maximise).

The nonlinearity is only contained in the nonlinear objective function,
supposed to have the following expression:

``a^T x + f(b^T x)``

where ``a`` is encoded as `linear_coefficients` and ``b`` as
`nonlinear_coefficients`. The function `f` is indicated by
`nonlinear_function`, currently a value of the enumeration `NonlinearFunction`.

`ε` is a parameter often used for nonlinear instances indicating the (additive)
precision at which the problem should be solved.
"""
struct NonlinearCombinatorialInstance <: CombinatorialInstance
    combinatorial_structure::CombinatorialInstance
    linear_coefficients::Union{Vector{Float64}, Dict{<:Any, Float64}}
    nonlinear_coefficients::Union{Vector{Float64}, Dict{<:Any, Float64}}
    nonlinear_function::NonlinearFunction
    # TODO: move these parameters to the specific algorithms (they are not really defining the problem to solve).
    ε::Float64 # Only for approximations.
    linear_algo::CombinatorialAlgorithm # Only for approximations.
    all_budgets_at_once::Bool # Only for approximation.
    formulation::CombinatorialLinearFormulation # Only for exact approach. Could be merged with linear_algo, but with a loss of precision in type (potentially useful for users).
end

abstract type NonlinearCombinatorialAlgorithm <: CombinatorialAlgorithm end


export TransformedSequence

struct TransformedSequence <: Sequence
  seq
  f
end

function sample(seq::TransformedSequence, t::AbstractVector)
  seq.f.(sample(seq.seq, t))
end

for f in (:exp, :log, :sin, :cos)
  @eval begin
    function Base.$f(s::Sequence)
      TransformedSequence(s, $f)
    end
  end
end


# 60x59 med rectangular (uses QR Lanczos)
using Pseudospectra, Test

@testset "Medium Rectangular" begin
    G = Pseudospectra.grcar(60)
    # drop a column to test rectangular case
    A = G[:,1:end-1]
    # for a rectangular mtx, must provide axes
    opts = Dict{Symbol,Any}(:ax => [-1.0,3.0,-3.0,3.0],:npts=>60)

    ps_data = new_matrix(A,opts)
    driver!(ps_data,opts,gs)
    @test iscomputed(ps_data)
end


#-----------------------------------------------------------------------
#
#   OUTLIER DETECTION FUNCTIONS
#
#-----------------------------------------------------------------------

#-----------------------------------------------------------------------
"""
    combX( X::RealMtx )

The `combX` function constructs and returns a vector of ``4×4`` matrices
consisting of all choices of 4 column vectors from the ``4×n_{\\rm e}`` matrix
`X`.

"""
function combX( X::RealMtx )
    tpfl=typeof(X[1,1])
    np = size(X)[2]

    if np==4
        return [X]
    elseif np>4
        Z = [zeros(tpfl,4) for _ in 1:np]   # Create containers

        for i=1:np
            Z[i] = copy(X[:,i])
        end

        W = collect(combinations(Z,4))      # Combinations
        nc = length(W)
        Y = [zeros(tpfl,4,4) for _ in 1:nc]

        for i=1:nc
            Y[i][:,1] = W[i][1]
            Y[i][:,2] = W[i][2]
            Y[i][:,3] = W[i][3]
            Y[i][:,4] = W[i][4]
        end

        return Y
    else
        return [zeros(tpfl,4,4)]
    end
end     #---------------------------------------------------------------

#-----------------------------------------------------------------------
"""
    odetc( XC::Vector , ξ::Real=1e1 )

The `odetc` function takes a vector of vectors as input and returns a
vector of the vectors that differ by less than a threshold of `ξ`.

"""
function odetc( XC::Vector , ξ::Real=1e1 )
    tpfl    = typeof(XC[1][1])
    ns      = length(XC)

    Xc  = zeros(tpfl,4,ns)
    ΔX  = zeros(tpfl,4,ns)
    Xa  = zeros(tpfl,4)
    V   = zeros(tpfl,ns)
    BV  = zeros(Bool,ns)
    σ   = zeros(tpfl,ns)

    for i=1:ns
        Xc[:,i] = XC[i][:]
    end

    for i=1:4
        Xa[i] = median(Xc[i,:])
    end

    for i=1:ns
        for j=1:4
            ΔX[j,i] = (Xc[j,i] - Xa[j])
        end
    end

    for i=1:ns
        V[i] = dot( ΔX[:,i] , ΔX[:,i] )
    end

    IV1 = partialsortperm( V , 1:ns )
    nr = 0

    if ns > 1
        for i=1:ns
            σ[i]    = sqrt(sum(V[IV1][1:i])/i)
            if i==1
                BV[i]   =  true
                nr = 1
            elseif i>1
                if σ[i] < ξ
                    BV[i]   =  true
                    nr += 1
                end
            end
        end
    end
    if nr > 1
        IV2 = partialsortperm( BV , 1:ns , rev=true )
        return (XC[IV1[IV2]][1:nr],IV1[IV2])
    else
        return (XC[IV1][1:1],nr*ones(Int,ns))
    end
end     #---------------------------------------------------------------


Base.@kwdef struct LinearAnalysis{I}
    npoints::Int=200
    dropcollinear::Bool=false
    interval::I=automatic
    level::Float64=0.95
end

function add_intercept_column(x::AbstractVector{T}) where {T}
    mat = similar(x, float(T), (length(x), 2))
    fill!(view(mat, :, 1), 1)
    copyto!(view(mat, :, 2), x)
    return mat
end

# TODO: add multidimensional version
function (l::LinearAnalysis)(input::ProcessedLayer)
    output = map(input) do p, n
        x, y = p
        weights = get(n, :weights, similar(x, 0))
        default_interval = length(weights) > 0 ? nothing : :confidence
        interval = l.interval === automatic ? default_interval : l.interval
        # FIXME: handle collinear case gracefully
        lin_model = GLM.lm(add_intercept_column(x), y; wts=weights, l.dropcollinear)
        x̂ = range(extrema(x)..., length=l.npoints)
        pred = GLM.predict(lin_model, add_intercept_column(x̂); interval, l.level)
        return if !isnothing(interval)
            # TODO: fix this in GLM https://github.com/JuliaStats/GLM.jl/pull/467
            ŷ, lower, upper = map(vec, pred) # GLM prediction returns matrices
            (x̂, ŷ), (; lower, upper)
        else
            ŷ = pred
            (x̂, ŷ), (;)
        end
    end
    default_plottype = isempty(output.named) ? Lines : LinesFill
    plottype = Makie.plottype(output.plottype, default_plottype)
    return ProcessedLayer(output; plottype)
end

"""
    linear(; interval=automatic, level=0.95, dropcollinear=false, npoints=200)

Compute a linear fit of `y ~ 1 + x`. An optional named mapping `weights` determines the weights.
Use `interval` to specify what type of interval the shaded band should represent,
for a given coverage `level` (the default `0.95` equates `alpha = 0.05`).
Valid values of `interval` are `:confidence`, to delimit the uncertainty of the predicted
relationship, and `:prediction`, to delimit estimated bounds for new data points.
Use `interval = nothing` to only compute the line fit, without any uncertainty estimate.
By default, this analysis errors on singular (collinear) data. To avoid that,
it is possible to set `dropcollinear=true`.
`npoints` is the number of points used by Makie to draw the shaded band.
"""
linear(; options...) = transformation(LinearAnalysis(; options...))


function gc2gc!(dst_ceps::AbstractVector, dst_γ,
                src_ceps::AbstractVector, src_γ)
    fill!(dst_ceps, zero(eltype(dst_ceps)))
    dst_order = length(dst_ceps)-1
    m1 = length(src_ceps)-1
    dst_ceps[1] = src_ceps[1]

    for m=2:dst_order+1
        ss1, ss2 = 0.0, 0.0
        min = (m1 < m-1) ? m1 : m - 2

        for k=2:min+1
            @inbounds cc = src_ceps[k] * dst_ceps[m-k+1]
            ss2 += (k-1) * cc
            ss1 += (m-k) * cc
        end

        if m <= m1+1
            @inbounds dst_ceps[m] = src_ceps[m] + (dst_γ*ss2 - src_γ*ss1)/(m-1)
        else
            @inbounds dst_ceps[m] = (dst_γ*ss2 - src_γ*ss1)/(m-1)
        end
    end

    dst_ceps
end

function gc2gc(src_ceps::AbstractVector, src_γ=0.0,
               dst_order=length(src_ceps)-1, dst_γ=0.0)
    gc2gc!(Array(eltype(src_ceps), dst_order+1), dst_γ, src_ceps, src_γ)
end

function gc2gc{T<:MelGeneralizedCepstrum}(state::SpectralParamState{T},
                                          dst_order=param_order(paramdef(state)),
                                          dst_γ=glog_gamma(paramdef(state)))
    assert_not_ready_to_filt(state)

    def = paramdef(state)
    order = param_order(def)
    src_α = allpass_alpha(def)
    src_γ = glog_gamma(def)
    data = rawdata(state)
    normalized = (dst_γ == zero(dst_γ))
    newdef = MelGeneralizedCepstrum(order, src_α, dst_γ)
    SpectralParamState(newdef, gc2gc(data, src_γ, dst_order, dst_γ),
                       has_loggain(state), normalized)
end


export TR_ls_step_computation
"""
Compute the direction used in the one dimensionnal trust region scheme.
Return either the Newton direction or the boundary of the trust region if the
Newton direction cannot be used. Inputs:
  - h, value of h(t).
  - g, value of h'(t).
  - dN, value of the Newton direction.
  - Δn, lower bound of the trust region.
  - Δp, upper bound of the trust region.
"""
function TR_ls_step_computation(h :: Float64, g :: Float64, dN :: Float64,
                                Δn :: Float64, Δp ::Float64)

  if h > 0.0
    if g > 0.0
      d = max(Δn, dN)
    else
      d = min(dN, Δp)
    end
  else
    if g > 0.0
      d = Δn
    else
      d = Δp
    end
  end

  return d
end


module TensorProductQuadraturesTests

using Gridap.ReferenceFEs

using Test

degrees = (2,2)
quad = Quadrature(QUAD,tensor_product,degrees)
test_quadrature(quad)
@test sum(get_weights(quad)) ≈ 1

D = 2
degree = 3
quad = Quadrature(QUAD,tensor_product,degree)
test_quadrature(quad)
@test sum(get_weights(quad)) ≈ 1

@test isa(get_name(quad),String)

end # module


# PREAMBLE

# PKG_SETUP

# ## Setup

using DirectTrajectoryOptimization
using LinearAlgebra
using Plots

# ## horizon
T = 51

# ## car
nx = 3
nu = 2
nw = 0

function car(x, u, w)
    [u[1] * cos(x[3]); u[1] * sin(x[3]); u[2]]
end

function midpoint_implicit(y, x, u, w)
    h = 0.1 # timestep
    y - (x + h * car(0.5 * (x + y), u, w))
end

# ## model
dt = Dynamics(midpoint_implicit, nx, nx, nu, nw=nw)
dyn = [dt for t = 1:T-1]
model = DynamicsModel(dyn)

# ## initialization
x1 = [0.0; 0.0; 0.0]
xT = [1.0; 1.0; 0.0]

# ## objective
ot = (x, u, w) -> 0.0 * dot(x - xT, x - xT) + 1.0 * dot(u, u)
oT = (x, u, w) -> 0.0 * dot(x - xT, x - xT)
ct = Cost(ot, nx, nu, nw, [t for t = 1:T-1])
cT = Cost(oT, nx, 0, nw, [T])
obj = [ct, cT]

# ## constraints
ul = -0.5 * ones(nu)
uu = 0.5 * ones(nu)
bnd1 = Bound(nx, nu, [1], xl=x1, xu=x1, ul=ul, uu=uu)
bndt = Bound(nx, nu, [t for t = 2:T-1], ul=ul, uu=uu)
bndT = Bound(nx, 0, [T], xl=xT, xu=xT)

p_obs = [0.5; 0.5]
r_obs = 0.1
function obs(x, u, w)
    e = x[1:2] - p_obs
    return [dot(e, e) - r_obs^2.0]
end
cont = StageConstraint(obs, nx, nu, nw, [t for t = 1:T-1], :inequality)
conT = StageConstraint(obs, nx, 0, nw, [T], :inequality)
cons = ConstraintSet([bnd1, bndt, bndT], [cont, conT])

# ## problem
trajopt = TrajectoryOptimizationProblem(obj, model, cons)
s = Solver(trajopt)

# ## initialize
x_interpolation = linear_interpolation(x1, xT, T)
u_guess = [0.001 * randn(nu) for t = 1:T-1]
z0 = zeros(s.p.num_var)
for (t, idx) in enumerate(s.p.trajopt.model.idx.x)
    z0[idx] = x_interpolation[t]
end
for (t, idx) in enumerate(s.p.trajopt.model.idx.u)
    z0[idx] = u_guess[t]
end
initialize!(s, z0)

# ## solve
@time solve!(s)

# ## solution
@show trajopt.x[1]
@show trajopt.x[T]

# ## state
plot(hcat(trajopt.x...)[1, :], hcat(trajopt.x...)[2, :], label = "", color = :orange, width=2.0)
pts = Plots.partialcircle(0.0, 2.0 * π, 100, r_obs)
cx, cy = Plots.unzip(pts)
plot!(Shape(cx .+ p_obs[1], cy .+ p_obs[2]), color = :black, label = "", linecolor = :black)

# ## control
plot(hcat(trajopt.u[1:end-1]..., trajopt.u[end-1])', linetype = :steppost)