"""
Example: a univariate time series (local linear trend) with an increasing variance, going from 1 to 5.
"""

# Packages
using Plots

T = 1_000
y = zeros(T, 1)
F = [1, 0]
G = Matrix(I, 2, 2)
m = zeros(2, 1)
P = Matrix(I, 2, 2) * 1000
S = Matrix(I, 1, 1)

#s = StateSpaceModel(y, F, G, m, P, S, 0.95, 0.95)
s = LocalLevelTrend(y, m, P, S, 0.98, 0.98)

# Simulate
simulate!(s)

# Estimate
m = estimate(s)

plot([s.y m.predicted.μ m.filtered.μ], labels = ["Simulated" "Predicted" "Filtered"], legend = :topleft)

#plot([y m.predicted.μ m.filtered.μ], labels = ["True" "Predicted" "Filtered"], legend = :topleft)
#plot([fill(Σ, T) m.predicted.Σ[:, 1, 1] m.filtered.Σ[:, 1, 1]][2:end, :], labels = ["True" "Predicted" "Filtered"], legend = :topleft)


using OrdinaryDiffEq, Test, Statistics
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_2Dlinear
prob = deepcopy(prob_ode_2Dlinear)
prob2 = ODEProblem(prob.f,prob.u0,(1.0,0.0),1.01)

sol = solve(prob2,DP5(),dt=-1/4,tstops=[0.5])
@test sol.t == [1.0,.75,.5,0]

sol = solve(prob2,RK4(),dt=-1/4, adaptive=false)
@test sol.t == [1.00,.75,.5,.25,0]

sol = solve(prob2,RK4(),dt=-1/4,tstops=[0.5,0.33], adaptive=false)
@test sol.t ≈ [1.00,.75,.5,.33,0.08,0]


module Optimization

using MicroLogging
using KrylovMethods

using DeformableRegistration
using DeformableRegistration.Types

# include("../experimental/aspin.jl")

export checkStoppingCriteria, ArmijoLineSearch, optimizeGaussNewton, optimizeGaussNewtonAugmentedLagrangian

# function augmentedLagrangian(x, λ, μ, c::Function)
#     m = length(x)
#     F = - λ'*c(x)[1] + μ/2 * c(x)[1]'*c(x)[1]
#     dF = - c(x)[2]*λ   + μ * c(x)[2]'*c(x)[1]
#     d2F = spzeros(m,m)#pdiagm(-λ, 0)  +  μ * c(x)[2]'*c(x)[2] + μ * spdiagm(c(x)[1],0)
#     return [F, dF, d2F]
# end

function augmentedLagrangian(x, λ, μ, c::Function)
    m = length(x)
    cx, ∇c = c(x)
    F = - λ' * cx + μ/2 * cx' * cx
    ∇F = - ∇c * λ   + μ * ∇c' * cx
    ∇2F(y) = μ * ∇c' * ∇c * y
    return [F, ∇F, ∇2F]
end

"""
optimizeGaussNewtonAugmentedLagrangian(Jfunc, y, yReference, options; constraint)

Augmented lagrangian implementation following Nocelda & Wright, 2006, p. 515 (Framework 17.3)
"""
function optimizeGaussNewtonAugmentedLagrangian(Jfunc::Function,  # objective Function
                             y::Array{Float64,1}, yReference::Array{Float64,1}, options::regOptions;
                             constraint::Function = x -> [0, 0, 0], printFunction=x->[], gradientDescentOnly=false, displayCallback=x->x) #no constraint by default,

    c = constraint
    computeConstraint = ! (c == x -> [0, 0, 0])

    L(x, λ, μ; doDerivative=true, doHessian=true) = (c(y) != [0, 0, 0]) ? Jfunc(x, doDerivative=doDerivative, doHessian=doHessian) +
     augmentedLagrangian(x, λ, μ, c) : Jfunc(x, doDerivative=doDerivative, doHessian=doHessian)

    λ = spzeros(size(y)...)
    μ = 1

    JRef = L(yReference, λ, μ)[1]
    JOld = NaN
    oldNormCy = 99 * norm(c(y)[1], Inf)

    for iter = 1:options.maxIterGaussNewton
     if computeConstraint
         λ = λ - μ * c(y)[1]

         "Enforce reduction of infeasability"
         τReduction = 0.8
         if norm(c(y)[1], Inf) >  τReduction * oldNormCy;
             @debug @sprintf("... norm(c(y)[1], Inf) >  %1.1f * oldNormCy: %3.3e > %3.3e", τReduction, norm(c(y)[1], Inf), oldNormCy)
             @debug "... updating μ = 2 * μ = ", 2*μ
             μ = 2 * μ;
         end
         oldNormCy = norm(c(y)[1], Inf)

     end

     J, dJ, d2J = L(y, λ, μ, doDerivative=true, doHessian=true)

     @debug d2J(-dJ)
     if gradientDescentOnly
         dy = -dJ
         cgIterations = -1
     else
         dy,flag,resvec,cgIterations = KrylovMethods.cg(d2J,-dJ,maxIter=options.maxIterCG, tol=options.tolCG)[1:4] #2/3 of the time of this function is here
     end

        # check descent direction
        if( (dJ'*dy)[1] > 0)
         dy = -dy
         @warn "Changing sign of computed descent direction. This is a suspicious move."
        end

     # armijo line search (LS) method
     stepLength,LSiter,LSfailed = ArmijoLineSearch(x -> L(x, λ, μ), J, dJ, y, dy, tolLS=1e-4)

     if(LSfailed)
         @info "STOPPING after line search failed."
         break
     end

     # output
     # is this a IJulia notebook? -> delete previous iteration's output
     if isdefined(Main, :IJulia) && Main.IJulia.inited
        displayCallback(true)
        @info size(y)
     end
     D, α, S = printFunction(y)
     s = @sprintf("%03d | D %2.2e | S %2.2e | λ!=0: %d | |c(y)|= %2.2e | μ=%2.2e | LSiter: %2d | CGiter: %3d | J/Jref: %2.2e",
         iter, D[1], S[1], sum(λ.!=0), norm(c(y)[1], Inf), μ,  LSiter, cgIterations, J[1]/JRef[1],
         )
     @info s
     @debug "  λ=", λ
     @debug "  c=", c(y)[1]

     # update parameter y
     y = y + stepLength .* dy

     # stopping criteria
     if (iter>1)
         if (checkStoppingCriteria(J, JOld, JRef, dJ, y, stepLength*dy,
                 tolJ = options.stopping["tolJ"],    # tolerance: change of the objective function
                 tolY = options.stopping["tolY"],    # tolerance: change of the variables
                 tolG = options.stopping["tolG"],    # tolerance: change of the gradient
                 tolQ = options.stopping["tolQ"]))   # tolerance: change of quotient)
             break
         end
     end

     # updating JOld
     JOld = J[1]
     @info "\n"
    end

    return y

end

function optimizeGaussNewton(Jfunc::Function,  # objective Function
                             y::Array{Float64,1}, yInitial::Array{Float64,1}, options::regOptions; displayCallback=x->x)

    # Initilization of JRef: reference value of objective Function
    JRef = Jfunc(yInitial)[1]
    JOld = NaN

    for iter = 1:options.maxIterGaussNewton

        # get derivatives of objective function
        J, dJ, d2J = Jfunc(y,doDerivative=true,doHessian=true)

        # solve system of linear equations d2J*dy=-dJ
          # (dy: change of the variables / search direction)
        # number of parameters < 10 => use backslash operator
        # else (large systems) => use conjugate gradient (cg) method

        cgIterations=0

        if(length(y)<10 && !(typeof(d2J) <: Function))
            dy=d2J\-dJ
        else
            dy,flag,resvec,cgIterations = KrylovMethods.cg(d2J,-dJ,maxIter=options.maxIterCG, tol=options.tolCG)[1:4]
        end

	    # check descent direction
	    if( (dJ'*dy)[1] > 0)
            dy = -dy
            @warn "Changing sign of computed descent direction. This is a suspicious move."
	    end

        # armijo line search (LS) method
        stepLength,LSiter,LSfailed = ArmijoLineSearch(Jfunc,J,dJ,y,dy)
        if(LSfailed)
            break
        end

        # is this a IJulia notebook? -> hand over to user provided method to delete previous iteration's output
        if isdefined(Main, :IJulia) && Main.IJulia.inited
           displayCallback(true)
        end

        # output
        if(cgIterations==0)
          s = @sprintf("%3d: J %8.4e     LSiter: %2d    J/Jref: %1.2f \n",iter, J[1], LSiter, J[1]/JRef[1])
          @info s
        else
          s = @sprintf("%3d: J %8.4e     LSiter: %2d     CGiter: %3d     J/Jref: %1.2f \n",iter, J[1], LSiter,cgIterations, J[1]/JRef[1])
          @info s
        end

        # update parameter y
        y = y + stepLength.*dy

        # stopping criteria
        if (iter>1)
            if checkStoppingCriteria(J, JOld, JRef, dJ, y, stepLength*dy,
                    tolJ = options.stopping["tolJ"],    # tolerance: change of the objective function
                    tolY = options.stopping["tolY"],    # tolerance: change of the variables
                    tolG = options.stopping["tolG"],    # tolerance: change of the gradient
                    tolQ = options.stopping["tolQ"])    # tolerance: change of quotient)
                break
            end
        end

        # updating JOld
        JOld = J[1]

    end

    return y

end

function checkArmijoStepValid(Jfunc::Function, Jc::Float64, y::Array{Float64,1},
        stepLength::Float64, dy::Array{Float64,1}, tolLS::Float64,
        dJ::Array{Float64,1})
    @debug "Jfunc(y+stepLength*dy)[1] < (Jc + (stepLength.*tolLS.*(dJ'*dy))[1])" Jfunc(y+stepLength*dy)[1]  (Jc + (stepLength.*tolLS.*(dJ'*dy))[1])
    validStep = Jfunc(y+stepLength*dy)[1] < (Jc + (stepLength.*tolLS.*(dJ'*dy))[1])
    return validStep
end

function ArmijoLineSearch(Jfunc::Function,         # objective function
                          Jc::Float64,             # Jc = Jfunc(y)
                          dJ::Array{Float64,1},    # gradient of the objective funtion
                          y::Array{Float64,1},     # old variables
                          dy::Array{Float64,1};    # change of the variables / search direction
                          tolLS::Float64 = 1e-4)   # tolerance

    stepLength = 1.0; LSiter = 1; LSfailed = false;

    while( true )
        validStep = checkArmijoStepValid(Jfunc, Jc, y, stepLength, dy, tolLS, dJ)
        if validStep
            break
        end
        if(stepLength < 1e-4)
            LSfailed = true
            break
        end

        stepLength = 0.5 * stepLength
        LSiter = LSiter + 1

    end

    if(LSfailed)
        s = @sprintf("   Line search failed after %2d iterations.",LSiter)
        @info s
    end

    return stepLength,LSiter,LSfailed

end

# somehow slower with this signature
# function checkStoppingCriteria(J::Float64,JOld::Float64,JRef::Float64,    # value of the current objective function J(y+dy), the old J(y) and the reference J(yInital)
#                                dJ::Array{Float64,1},             # gradient of the objective funtion
#                                y::Array{Float64,1},              # old variables
#                                dy::Array{Float64,1};             # change of the variables / search direction
#                                tolJ::Float64 = 1e-3,    # tolerance: change of the objective function
#                                tolY::Float64 = 1e-2,    # tolerance: change of the variables
#                                tolG::Float64 = 1e-2,    # tolerance: change of the gradient
#                                tolQ::Float64 = 1e-4     # tolerance: change of quotient
#                                )

function checkStoppingCriteria(J,JOld,JRef,    # value of the current objective function J(y+dy), the old J(y) and the reference J(yInital)
                               dJ,             # gradient of the objective funtion
                               y,              # old variables
                               dy;             # change of the variables / search direction
                               tolJ = 1e-3,    # tolerance: change of the objective function
                               tolY = 1e-2,    # tolerance: change of the variables
                               tolG = 1e-2,    # tolerance: change of the gradient
                               tolQ = 1e-4     # tolerance: change of quotient
                               )

    # Initialize STOP-Array with false
    STOP = Array{Bool}(5); STOP[:] = false

    # Either STOP[1:3] or STOP[4]
    STOP[1] = abs(JOld-J) <= tolJ  * (1+abs(JOld))
    STOP[2] = norm(dy)    <= tolY  * (1+norm(y))
    STOP[3] = norm(dJ)    <= tolG  * (1+abs(JOld))
    STOP[4] = norm(dJ)    <= 1e6 * eps()
    STOP[5] = abs((JOld-J)/(JRef-J)) <= tolQ

    if(all(STOP[1:3]))
        s = @sprintf("STOPPING CRITERIA:\n")
        @info s
        s = @sprintf("   1. |JOld-J| = %3e <= %3e (tolJ * (1 + |JOld|)) \n",abs(JOld-J), tolJ*(1+abs(JOld)) )
        @info s
        s = @sprintf("   2. ||dy||   = %3e <= %3e (tolY * (1 + ||y||))\n",norm(dy),    tolY*(1+norm(y))  )
        @info s
        s = @sprintf("   3. ||dJ||   = %3e <= %3e (tolG * (1 + |JOld|)) \n",norm(dJ),    tolG*(1+abs(JOld)) )
        @info s
    end

    if(STOP[4])
        s = @sprintf("STOPPING CRITERION:\n")
        @info s
        s = @sprintf("   4. ||dJ|| = %3e <= %3e (1e6*eps)\n",norm(dJ), 1e6*eps())
        @info s
    end

    if(STOP[5])
        s = @sprintf("STOPPING CRITERION:\n")
        @info s
        s = @sprintf("   5. |(JOld-J)/(JRef-J)|= %3e <= %3e (tolQ)\n",abs((JOld-J)/(JRef-J)), tolQ)
        @info s
    end

    return all(STOP[1:3]) | STOP[4] | STOP[5]

end

end


# TODO: optimize

# Q(Θ) = Θ'AΘ + b'Θ + c
# A is stochastic average of .25 * X'X
# b is stochastic average of ∇f(Θ) - .25 * X'X * Θ
type LogRegMM{W <: Weight} <: OnlineStat{XYInput}
	β::VecF
    H::Matrix{Float64}  # Storage
    A::Matrix{Float64}
    b::VecF
    η::Float64
    weight::W
end
function LogRegMM(p::Integer, wt::Weight = LearningRate(); η::Float64 = 1.0)
    @assert 0.0 < η <= 1.0
    LogRegMM(zeros(p), zeros(p, p), zeros(p, p), zeros(p), η, wt)
end
function LogRegMM(x::AMat, y::AVec, wt::Weight = LearningRate())
    o = LogRegMM(size(x, 2), wt)
    fit!(o, x, y)
    o
end
function Base.show(io::IO, o::LogRegMM)
    printheader(io, name(o))
    print_item(io, "β", coef(o))
    print_item(io, "nobs", nobs(o))
end

function _fit!(o::LogRegMM, x::AVec, y::Real, γ::Float64)
    γ *= o.η
    BLAS.syrk!('U', 'N', γ, x / 2., 1.0 - γ, o.H)
    H = Symmetric(o.H)
    smooth!(o.A, H, γ)
    smooth!(o.b, (y - predict(o, x)) * x + H * o.β, γ)
    # try
    #     o.β = o.A \ o.b
    # end
    o.β = o.β + γ * inv(o.H) * (y - predict(o, x)) * x
end

value(o::LogRegMM) = coef(o)
coef(o::LogRegMM) = o.β
linpred(o::LogRegMM, x::AVec) = dot(x, o.β)
linpred(o::LogRegMM, x::AMat) = x * o.β
predict(o::LogRegMM, x) = predict(LogisticRegression(), linpred(o, x))
loss(o::LogRegMM, x, y) = loss(LogisticRegression(), y, linpred(o, x))


using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using IterativeSolvers
using LinearAlgebra

# The user can create his own orthogonalization function to use in IAR
function doubleGS_function!(VV, vv, h)
    h[:]=VV'*vv; vv[:]=vv-VV*h; g=VV'*vv; vv[:]=vv-VV*g;
    h[:] = h[:]+g[:]; β=norm(vv); vv[:]=vv/β; return β
end
# Then it is needed to create a type to access to this function
abstract type DoubleGS <: IterativeSolvers.OrthogonalizationMethod end
# And then introduce a function dispatch for this new type in order to use
# the defined orthogonalization function
import IterativeSolvers.orthogonalize_and_normalize!
function orthogonalize_and_normalize!(V,v,h,::Type{DoubleGS})
    doubleGS_function!(V, v, h) end

@testset "IAR" begin
    dep=nep_gallery("dep0");
    n=size(dep,1);

    @bench @testset "accuracy eigenpairs" begin
        (λ,Q)=iar(dep,σ=1.1,neigs=5,v=ones(n),
                  maxit=100,tol=eps()*100,errmeasure=ResidualErrmeasure(dep));
        verify_lambdas(5, dep, λ, Q, eps()*100)
    end

    @bench @testset "Solve by projection" begin
        (λ,Q)=iar(dep,σ=1.1,neigs=5,v=ones(n),
                  maxit=100,tol=eps()*100,errmeasure=ResidualErrmeasure(dep), proj_solve=true);
        verify_lambdas(5, dep, λ, Q, eps()*100)
    end

    @testset "Compute as many eigenpairs as possible (neigs=Inf)" begin
        (λ,Q)=iar(dep,σ=1.1,neigs=Inf,v=ones(n),
                  maxit=38,tol=eps()*100);
        verify_lambdas(6, dep, λ, Q, eps()*200)
    end

    @testset "orthogonalization" begin
        # NOW TEST DIFFERENT ORTHOGONALIZATION METHODS

        @bench @testset "DGKS" begin
            (λ,Q,V)=iar(dep,orthmethod=DGKS,σ=1.1,neigs=5,v=ones(n),maxit=100,tol=eps()*100)
            @test opnorm(V'*V - I) < 1e-6
        end

        @bench @testset "User provided doubleGS" begin
            (λ,Q,V)=iar(dep,orthmethod=DoubleGS,σ=1.1,neigs=5,v=ones(n),maxit=100,tol=eps()*100)
            @test opnorm(V'*V - I) < 1e-6
        end

        @bench @testset "ModifiedGramSchmidt" begin
            (λ,Q,V)=iar(dep,orthmethod=ModifiedGramSchmidt,σ=1.1,neigs=5,v=ones(n),maxit=100,tol=eps()*100)
            @test opnorm(V'*V - I) < 1e-6
        end

        @bench @testset "ClassicalGramSchmidt" begin
            (λ,Q,V)=iar(dep,orthmethod=ClassicalGramSchmidt,σ=1.1,neigs=5,v=ones(n),maxit=100,tol=eps()*100)
            @test opnorm(V'*V - I) < 1e-6
        end
    end

    @testset "Errors thrown" begin
        np=100;
        dep=nep_gallery("dep0",np);
        @test_throws NEPCore.NoConvergenceException (λ,Q)=iar(dep,σ=1.1,neigs=6,v=ones(np),
                  maxit=7,tol=eps()*100);
    end

end


"""
Abstract equation.
"""
abstract type Equation end

"""
    AdvectionEquation(domain, filter = IdentityFilter())

Filtered advection equation.
"""
Base.@kwdef struct AdvectionEquation{D<:Domain,F<:Filter} <: Equation
    domain::D
    filter::F = IdentityFilter()
end

"""
    DiffusionEquation(
        domain,
        filter = IdentityFilter(),
        f = nothing,
        g_a = nothing,
        g_b = nothing
    )

Filtered diffusion equation.
"""
Base.@kwdef struct DiffusionEquation{D<:Domain,F<:Filter} <: Equation
    domain::D
    filter::F = IdentityFilter()
    f::Union{Nothing,Function} = nothing
    g_a::Union{Nothing,Function} = nothing
    g_b::Union{Nothing,Function} = nothing
end

"""
    BurgersEquation(domain, filter = IdentityFilter())

Filtered Burgers equation.
"""
Base.@kwdef struct BurgersEquation{D<:Domain,F<:Filter} <: Equation
    domain::D
    filter::F = IdentityFilter()
end


using Statistics, SpecialFunctions

"""
    tnmom1c(c, a, b)

Computes ⟨(x - c)⟩ for the standard normal distribution truncated to the
interval [a, b]. This is more numerically stable than computing the mean
of x and subtracting c. Returns NaN if a > b.
"""
function tnmom1c(c::Real, a::Real, b::Real)
    if !(a ≤ b) || isnan(c)
        return oftype(middle(a, b), NaN)
    elseif a == b || isinf(c)
        return middle(a, b) - c
    elseif abs(a) > abs(b)
        return -tnmom1c(-c, -b, -a)
    elseif c ≤ 0
        #= At this point tnmean(a,b) ≥ 0.
        Therefore the subtraction is fine. =#
        return tnmean(a, b) - c
    elseif isinf(a) && isinf(b)
        return oftype(middle(a, b), -c)
    end

    @assert a < b
    @assert a ≤ 0 ≤ b || 0 ≤ a ≤ b
    @assert c > 0

    exΔm1 = expm1(middle(a, b) * (a - b))

    if a ≤ 0 ≤ b
        exa = √(2/π) * exp(-a^2 / 2)
        exb = √(2/π) * exp(-b^2 / 2)
        exa ≈ exb && return -c
        era = erf(a / √2)
        erb = erf(b / √2)
        era ≈ erb && return middle(a - c, b - c)

        m1 = (exa - exb) / (era - erb)
        c ≈ m1 && return zero(m1)

        fa = exa + c * era
        fb = exb + c * erb
        return (fa - fb) / (erb - era)
    elseif 0 < a < b
        era = erfcx(a / √2)
        erb = erfcx(b / √2)
        if xerfcx_asym_thresh(a / √2, 3)
            fa = (1 - c/a) * √(2/π) + √2 * c/a * xerfcx_pi(a / √2)
        else
            fa = √(2/π) - c * era
        end
        if xerfcx_asym_thresh(b / √2, 3)
            fb = (1 - c/b) * √(2/π) + √2 * c/b * xerfcx_pi(b / √2)
        else
            fb = √(2/π) - c * erb
        end
        fa = (1 - c/a) * √(2/π) + √2 * c/a * xerfcx_pi(a / √2)
        fb = (1 - c/b) * √(2/π) + √2 * c/b * xerfcx_pi(b / √2)
        exΔ = exp(middle(a, b) * (a - b))
        return (fa - fb * exΔ) / (era - erb * exΔ)
    end
end

"""
    tnmom1c(c, a, b, μ, σ)

Computes ⟨(x - c)⟩ for the normal distribution with mean μ and standard
deviation σ, truncated to the interval [a, b]. This is more numerically stable
than computing the mean of x and subtracting c.
"""
function tnmom1c(c, a, b, μ, σ)
    α = (a - μ) / σ
    β = (b - μ) / σ
    γ = (c - μ) / σ
    return σ * tnmom1c(γ, α, β)
end


const OrbitProjection = Equilibrium.Boundary

const in_orbit = in_boundary

function orbit_projection(M::AbstractEquilibrium, gcp; dx = 0.01, dy =0.01)
    hc = HamiltonianCoordinate(M,gcp)
    rlims, zlims = limits(M)

    x = range(rlims...,step=dx)
    y = range(zlims...,step=dy)

    mu = zeros(length(x), length(y))
    for ix in eachindex(x), iy in eachindex(y)
        B = norm(Bfield(M, x[ix], y[iy]))
        ir, iphi, iz = cylindrical_cocos_indices(cocos(M))
        g = poloidal_current(M, M(x[ix], y[iy])) # RBphi
        energy = 1e3*hc.energy*e0
        mu[ix,iy] = max(energy/B - (B/(2*hc.m))*((hc.p_phi - hc.q*e0*M(x[ix],y[iy]))/g)^2,0.0)
    end

    cl = Contour.contour(x,y,mu,hc.mu)
    ls = Contour.lines(cl)
    oproj = OrbitProjection([(gcp.r, gcp.z)])
    if length(ls) == 0
        return oproj
    end
    omin = Inf
    for (i,l) in enumerate(ls)
        r, z = Contour.coordinates(l)
        omin_tmp = minimum((r .- gcp.r).^2 .+ (z .- gcp.z).^2)
        if omin_tmp < omin
            omin = omin_tmp
            oproj = OrbitProjection(collect(zip(r,z)))
        end
    end

    return oproj
end


# Mel-cepstrum analysis
# re-coded from SPTK

!isdefined(Base, :FFTW) && using FFTW

function fill_al!{T<:AbstractFloat}(al::Vector{T}, α::AbstractFloat)
    al[1] = one(T)
    for i=2:length(al)
        @inbounds al[i] = -α*al[i-1]
    end
    al
end

function fill_toeplitz!{T}(A::AbstractMatrix{T}, t::AbstractVector{T})
    n = length(t)
    for j=1:n, i=1:n
        @inbounds A[i,j] = i-j+1 >= 1 ? t[i-j+1] : t[j-i+1]
    end
    A
end

function fill_hankel!{T}(A::AbstractMatrix{T}, h::AbstractVector{T})
    n = length(h)>>1 + 1
    for j=1:n, i=1:n
        @inbounds A[i,j] = h[i+j-1]
    end
    A
end

function fill_only_real_part!{T}(y::AbstractVector{Complex{T}},
                                 v::AbstractVector{T})
    for i=1:length(v)
        @inbounds y[i] = Complex(v[i], zero(T))
    end
    y
end

function update_hankel_elements!(he::AbstractVector, c::AbstractVector)
    for j=1:length(he)
        @inbounds he[j] = c[j]
    end
    for j=1:2:length(he)
        @inbounds he[j] -= c[1]
    end
    he
end

function update_toeplitz_elements!(te::AbstractVector, c::AbstractVector)
    for j=1:2:length(te)
        @inbounds te[j] += c[1]
    end
    te
end

function periodogram2mcep{T<:AbstractFloat}(periodogram::AbstractVector{T}, # modified periodogram
                                            order::Int=40,                  # order of mel-cepstrum
                                            α::AbstractFloat=0.41;          # all-pass constant
                                            miniter::Int=2,
                                            maxiter::Int=30,
                                            criteria::AbstractFloat=0.001,  # stopping criteria
                                            e::T=zero(T),                   # floor of periodogram
                                            verbose::Bool=false)
    logperiodogram = log.(periodogram + e)

    fftlen = (size(logperiodogram, 1)-1)*2
    const xh = fftlen>>1

    # create FFT workspace and plan
    y = Array{Complex{T},1}(xh+1)
    c = Array{T,1}(fftlen)
    # fplan = FFTW.Plan(c, y, 1, FFTW.ESTIMATE, FFTW.NO_TIMELIMIT)
    # bplan = FFTW.Plan(y, c, 1, FFTW.ESTIMATE, FFTW.NO_TIMELIMIT)
    fplan = plan_rfft(c)
    bplan = plan_brfft(y, fftlen)

    # Initial value of cepstrum
    fill_only_real_part!(y, logperiodogram)
    A_mul_B!(c, bplan, y)
    # FFTW.execute(T, bplan.plan)
    # FFTW.execute(bplan.plan, y, c)
    scale!(c, 1. / fftlen)
    # scale!(c, FFTW.normalization(c))
    c[1] /= 2.0
    c[xh+1] /= 2.0

    # Initial value of mel-cesptrum
    mc = freqt(view(c, 1:xh+1), order, α)
    czero = c[1]

    # Allocate memory for solving linear equation (Tm + Hm)d = b
    Tm = Array{T,2}(order+1, order+1)
    Hm = Array{T,2}(order+1, order+1)
    Tm_plus_Hm = Array{T,2}(order+1, order+1)
    he = Array{T,1}(2order+1) # elements of hankel matrix
    te = Array{T,1}(order+1)  # elements of toeplitz matrix
    b = Array{T,1}(order+1)   # right side of linear equation

    al = Array{T,1}(order+1)
    fill_al!(al, α)

    # Newton raphson roop
    ch = view(c, 1:xh+1)
    ch_copy = Array{T,1}(xh+1)
    c_frqtr = view(c, 1:2order+1)
    for i=1:maxiter
        fill!(c, zero(T))
        freqt!(ch, mc, -α)

        A_mul_B!(y, fplan, c)
        # FFTW.execute(fplan.plan, c, y)
        for i=1:length(y)
            @inbounds y[i] = Complex(periodogram[i] / exp(2real(y[i])), zero(T))
        end
        A_mul_B!(c, bplan, y)
        # FFTW.execute(bplan.plan, y, c)
        # scale!(c, FFTW.normalization(c))
        scale!(c, 1.0 / fftlen)

        copy!(ch_copy, ch)
        frqtr!(c_frqtr, ch_copy, α)

        # check convergence
        if i >= miniter
            err = abs((c[1]-czero)/c[1])
            verbose && println("czero nmse: $err")
            if err < criteria
                break
            end
            czero = c[1]
        end

        copy!(te, 1, c, 1, order+1)

        for j=1:order+1
            @inbounds b[j] = c[j] - al[j]
        end

        update_hankel_elements!(he, c)
        update_toeplitz_elements!(te, c)

        fill_hankel!(Hm, he)
        fill_toeplitz!(Tm, te)

        for i=1:length(Hm)
            @inbounds Tm_plus_Hm[i] = Hm[i] + Tm[i]
        end

        # Solve Ax = b
        # NOTE: both Tm_plus_Hm and b are overwritten
        A_ldiv_B!(lufact!(Tm_plus_Hm), b)

        # Add the solution vector to mel-cepstrum
        for i=1:length(b)
            @inbounds mc[i] += b[i]
        end
    end

    mc
end

function _mcep(x::AbstractVector,  # a *windowed* signal
               order=25,           # order of mel-cepstrum
               α=0.35;             # all-pass constant
               kargs...)
    periodgram = abs2.(rfft(x))
    periodogram2mcep(periodgram, order, α; kargs...)
end

function estimate(mgc::MelCepstrum, x::AbstractArray;
                  use_sptk::Bool=false,
                  kargs...)
    order = param_order(mgc)
    α = allpass_alpha(mgc)
    # Note that mcep in julia is more stable than SPTK.mcep
    mcepfunc::Function = use_sptk ? SPTK.mcep : _mcep
    data = mcepfunc(x, order, α; kargs...)
    SpectralParamState(mgc, data, true, true)
end

function mcep(x::AbstractArray, order=25, α=0.35; kargs...)
    estimate(MelGeneralizedCepstrum(order, α, 0.0), x; kargs...)
end


function parse_line(l)
    parse.(Int, collect(l))
end

const data = reduce(hcat, parse_line.(readlines()))

neighbours(pos) = [
    CartesianIndex(Tuple(pos) .- n)
    for n ∈ ((-1, 0), (1, 0), (0, -1), (0, 1))
]

# Part One
sum(CartesianIndices(data)) do i
    neib = [get(data, neib, Inf) for neib ∈ neighbours(i)]
    all(
        data[i] < n for n in neib
    ) ? data[i] + 1 : 0
end |> println

function dfs(pos, seen::Set)
    curr_n = 1
    push!(seen, pos)
    for neib in neighbours(pos)
        neib ∉ CartesianIndices(data) && continue
        if neib ∉ seen
            curr_n += dfs(neib, seen)
        end
    end
    curr_n
end

function solve_2()
    seen_pos = Set(findall(==(9), data))
    remain_pos = setdiff(CartesianIndices(data), seen_pos)
    basins = []
    while !isempty(remain_pos)
        next_pos = pop!(remain_pos)
        push!(basins, dfs(next_pos, seen_pos))
        setdiff!(remain_pos, seen_pos)
    end
    partialsort!(basins, 3; rev=true)
    prod(basins[1:3])
end

# Part Two
solve_2() |> println


#=
# 2D DG poisson
=#
if !@isdefined(Femshop)
    include("../Femshop.jl");
    using .Femshop
end
init_femshop("dgpoisson2d");

@useLog("dgpoisson2dlog")

@domain(2)
@solver(DG)
@functionSpace(LEGENDRE, 1)

n = 20;
@mesh(QUADMESH, n)

@variable(u)
@testSymbol(v)

@boundary(u, 1, DIRICHLET, 0)

@coefficient(f, "-2*pi*pi*sin(pi*x)*sin(pi*y)")
@coefficient(alpha, 2*n)
@weakForm(u, "-dot(grad(u), grad(v)) - f*v + surface(dot(ave(grad(u)), jump(v)) + dot(ave(grad(v)), jump(u)) - alpha*dot(jump(u), jump(v)))")

solve(u);

# # solution is stored in the variable's "values"
using Plots
pyplot();
display(plot(Femshop.grid_data.allnodes[1,:], Femshop.grid_data.allnodes[2,:], u.values[:], st=:surface))

@finalize()


function construct( e::ProjectionMatrix,
                   𝐊::AbstractArray,
                   𝐑::AbstractArray,
                    𝐭::AbstractArray )

    if size(𝐊) != (3,3) || size(𝐑) != (3,3)
        throw(ArgumentError("Expect 3 x 3 calibration and rotation matrices."))
    end
    if length(𝐭) != 3
        throw(ArgumentError("Expect length-3 translation vectors."))
    end
    # TODO: Reconcile this change in convention with the rest of the code.
    𝐏 = 𝐊*[𝐑' -𝐑'*𝐭]
    SMatrix{3,4,Float64,3*4}(𝐏)
end

function construct(e::ProjectionMatrix, 𝐅::AbstractArray)
    𝐞 = epipole(𝐅')
    𝐏₁ = Matrix{Float64}(I,3,4)
    𝐏₂ = [vec2antisym(𝐞) * 𝐅  𝐞]

    SMatrix{3,4,Float64,3*4}(𝐏₁), SMatrix{3,4,Float64,3*4}(𝐏₂)

end

function construct(e::ProjectionMatrix, 𝐄::AbstractArray, 𝒟::Tuple{AbstractArray, Vararg{AbstractArray}})
    𝐖 = SMatrix{3,3,Float64,3*3}([0 -1 0; 1 0 0; 0 0 1])
    𝐙 = SMatrix{3,3,Float64,3*3}([0 1 0; -1 0 0; 0 0 0])
    𝐔,𝐒,𝐕 = svd(𝐄)
    𝐭 = 𝐔[:,3]
    𝐏₁ = SMatrix{3,4,Float64,3*4}(1.0I)
    𝐏₂₁ = SMatrix{3,4,Float64,3*4}([𝐔*𝐖*𝐕'  𝐭])
    𝐏₂₂ = SMatrix{3,4,Float64,3*4}([𝐔*𝐖'*𝐕' 𝐭])
    𝐏₂₃ = SMatrix{3,4,Float64,3*4}([𝐔*𝐖*𝐕' -𝐭])
    𝐏₂₄ = SMatrix{3,4,Float64,3*4}([𝐔*𝐖'*𝐕' -𝐭])

    𝒳₁ = triangulate(DirectLinearTransform(), 𝐏₁, 𝐏₂₁, 𝒟)
    𝒳₂ = triangulate(DirectLinearTransform(), 𝐏₁, 𝐏₂₂, 𝒟)
    𝒳₃ = triangulate(DirectLinearTransform(), 𝐏₁, 𝐏₂₃, 𝒟)
    𝒳₄ = triangulate(DirectLinearTransform(), 𝐏₁, 𝐏₂₄, 𝒟)

    # Determine which projection matrix in the second view triangulated
    # the majority of points in front of the cameras.
    ℳ₁ = map(𝒳₁) do 𝐗
        𝐦 = 𝐏₂₁ * 𝐗
        𝐦[3] > 0
    end

    ℳ₂ = map(𝒳₂) do 𝐗
        𝐦 = 𝐏₂₂ * 𝐗
        𝐦[3] > 0
    end

    ℳ₃ = map(𝒳₃) do 𝐗
        𝐦 = 𝐏₂₃ * 𝐗
        𝐦[3] > 0
    end

    ℳ₄ = map(𝒳₄) do 𝐗
        𝐦 = 𝐏₂₄ * 𝐗
        𝐦[3] > 0
    end

    total, index = findmax((sum(ℳ₁), sum(ℳ₂), sum(ℳ₃), sum(ℳ₄)))

    if index == 1
        return 𝐏₁,  𝐏₂₁
    elseif index == 2
        return 𝐏₁,  𝐏₂₂
    elseif index == 3
        return 𝐏₁,  𝐏₂₃
    else
        return 𝐏₁,  𝐏₂₄
    end
end

function construct(e::ProjectionMatrices, ℋ::Tuple{AbstractArray, Vararg{AbstractArray}})
    construct(LatentVariables(HomographyMatrices()), ℋ)
end

function construct(lv::LatentVariables, ℋ::Tuple{AbstractArray, Vararg{AbstractArray}})
        N = length(ℋ)
        if N < 2
            throw(ArgumentError("Please supply at least two homography matrices."))
        end
        𝐗₁ = ℋ[1]
        𝛍 = zeros(N)
        𝐉 = Array{Float64}(undef,(3,(N-1)*6))
        i₁ = range(1, step = 6, length = N - 1)
        i₂ = range(6, step = 6, length = N - 1)
        for n = 2:N
            𝐗ₙ = ℋ[n]
            e₁, e₂ = find_nearest_eigenvalues(eigvals(Array(𝐗₁), Array(𝐗ₙ)))
            𝐘 = hcat(e₁ * 𝐗ₙ - 𝐗₁, e₂ * 𝐗ₙ - 𝐗₁)
            μ = (e₁ + e₂) / 2
            𝛍[n] = real(μ)
            𝐉[:,i₁[n-1]:i₂[n-1]] .= 𝐘
        end
        𝛈 = initialisation_procedure(𝐉, 𝛍, ℋ)
        𝐚  = @view 𝛈[1:9]
        𝐀 = reshape(𝐚, (3,3))
        𝐛 = @view 𝛈[10:12]
        𝐏₁ = SMatrix{3,4}(1.0I)
        𝐏₂ = SMatrix{3,4}(hcat(𝐀, 𝐛))
        𝐏₁, 𝐏₂
end

function initialisation_procedure(𝐉::AbstractArray, 𝛍::AbstractArray, ℋ::Tuple{AbstractArray, Vararg{AbstractArray}})
    N = length(ℋ)
    if N < 2
        throw(ArgumentError("Please supply at least two homography matrices."))
    end
    F = svd(𝐉)
    𝛈 = zeros(9 + 3 + (N*3) + N)
    𝐛 = real(F.U[:,1])
    𝐗₁ = ℋ[1]
    𝐀 = 𝐗₁
    𝐯₁ = SVector(0,0,0)
    wₙ = 1
    # pack 𝛈 = [𝐚,𝐛, 𝐯₁,...,𝐯ₙ, w₁, ..., wₙ]
    𝛈[1:9] .= vec(𝐀)
    𝛈[10:12] .= 𝐛
    for (n,i) in enumerate(range(13, step = 3, length = N))
        if n == 1
            𝛈[i:i+2] .= 𝐯₁
        else
            𝐗ₙ = ℋ[n]
            𝛈[i:i+2] .= 𝐯₁ +  (𝛍[n] * 𝐗ₙ - 𝐗₁)' * 𝐛 / (norm(𝐛)^2)
        end
    end
    𝛈[end-(N-1):end] .= wₙ
    𝛈
end

function find_nearest_eigenvalues(e::AbstractArray)
    dist = SVector(abs(e[1]-e[2]), abs(e[1]-e[3]), abs(e[2]-e[3]))
    minval, index = findmin(dist)
    if index == 3
        i₁ = 2
        i₂ = 3
    elseif index == 2
        i₁ = 1
        i₂ = 3
    else
        i₁ = 1
        i₂ = 2
    end
    e[i₁], e[i₂]
end

function unpack(e::LatentVariables, 𝛈::AbstractArray)
    N = div(length(𝛈) - 12,  4)
    𝐚  = @view 𝛈[1:9]
    𝐀 = reshape(𝐚, (3,3))
    𝐛 = @view 𝛈[10:12]
    𝐰 = @view 𝛈[end-(N-1):end]
    r = range(13, step = 3, length = N+1)
    @show first(r), last(r), N
    𝐯 = reshape(view(𝛈,first(r):last(r)-1), (3,N))
    𝒫 = Array{Array{Float64,2},1}(undef,(N,))
    for n = 1:N
        𝒫[n] = 𝐰[n]*𝐀 + 𝐛*𝐯[:,n]'
    end
    𝒫
end


using Statistics
using MATLAB
include("atn_module.jl")
# using .time_series

N = 10_000          # number of nodes
m = 30               # number of active links per active node
mu = 0.1            # probability I -> S
tmax = 1000          # maximum duration of the simulation
Ni = Int(0.1 * N)   # initial number of infects
gamma = -2.22       # exponent of the power law
NTrials = 5         # number of trials for each value of lambda
tolerance = 1e-15   # termination tolerance
TWindow = 1000      # time window for evaluating steady state
TMinConvergence = 500

# health state of agents S=0 I=1

# diffferent lambda
# lambda=0.1:0.1:1;
# IAvgLambda=zeros(length(lambda))
# beta=zeros(length(lambda))

lambda = 0.2

ITrend = atn_module.time_series(Ni, N, m, gamma, lambda, mu, tolerance, TWindow,
                                TMinConvergence, tmax);

mat"plot($ITrend)"


immutable PlaneWaveHH3DTD{T,P,F} <: TDFunctional{T}
    direction::P
    speed_of_light::T
    signature::F
    amplitude::T
end

#speedoflight(f::PlaneWaveHH3DTD) = f.speed_of_light

function planewave(direction, speedoflight::Number, signature, amplitude=one(speedoflight)::Number)
    PlaneWaveHH3DTD(direction, speedoflight, signature, amplitude)
end

*(a, f::PlaneWaveHH3DTD) = PlaneWaveHH3DTD(f.direction, f.speed_of_light, f.signature, a*f.amplitude)

function(f::PlaneWaveHH3DTD)(r,t)
    k = f.direction
    u = sum(k[i]*r[i] for i in 1:length(k))
    #u = dot(f.direction, r)
    a = f.amplitude
    c = f.speed_of_light
    h = f.signature
    a * h(c*t - u)
end


# # Quantum Circuit Born Machine

using Yao
using QuAlgorithmZoo

# ## DATA: Target Probability Distribution
# The gaussian probability disctribution in phase space of 2^6
"""
    gaussian_pdf(x, μ::Real, σ::Real)

gaussian probability density function.
"""
function gaussian_pdf(x, μ::Real, σ::Real)
    pl = @. 1 / sqrt(2pi * σ^2) * exp(-(x - μ)^2 / (2 * σ^2))
    pl / sum(pl)
end
nbit = 6
N = 1<<nbit
pg = gaussian_pdf(1:N, N/2-0.5, N/4);

# ## MODEL: Quantum Circuit and Loss
# Using a random differentiable circuit of depth 6 for training, the kernel function is universal RBF kernel
depth = 6
kernel = rbf_kernel(0:N-1, 0.25)
circuit = random_diff_circuit(nbit, depth, pair_ring(nbit)) |> autodiff(:QC);
dispatch!(circuit, :random)
qcbm = QCBM(circuit, kernel, pg);

# ## TRAINING: Adam Optimizer
# We probide the QCBMGo! iterative interface for training
niter = 100
optim = Adam(lr=0.1)
for info in QCBMGo!(qcbm, optim, niter)
    curr_loss = loss(qcbm, info["probs"])
    println("Step = ", info["step"], ", Loss = ", curr_loss)
end


function reflection_coefficient(ω::T, dwave::DiscretePlaneWaveMode{T},
    source::PlaneSource{T,2,1,Acoustic{T,2}}, material::Material{2,Halfspace{T,2}};
     # medium::Acoustic{T,2}, specie::Specie{T,2}; θin::T = 0.0,
     scheme::Symbol = :trapezoidal) where T <: AbstractFloat

    k = ω / source.medium.c
    M = dwave.basis_order
    σ = k .* integration_scheme(dwave.x; scheme=scheme) # multiple by k to be the same as the non-dimensional version

    θin = transmission_angle(source,material)

    R = T(2)*number_density(material.species) / (cos(θin)*k^2)*sum(
        im^T(m)*exp(im*k*dwave.x[j]*cos(θin) - im*θin*m)*dwave.amplitudes[j,m+M+1,1]*σ[j]
    for m=-M:M, j in eachindex(dwave.x))

    return R
end


# Linear Regression Tests

# OLS test:
OLSRegression = OLS(intercept=true, robust=false)
res = linregress(X, Y, OLSRegression; fast = true)

# OLS robust test:
OLSRegression = OLS(intercept=true, robust=HCEVariance())
res = solve(X, Y, OLSRegression)

# FGLS test
FGLSRegression = FGLS(intercept=true, robust=HCEVariance())
res = linregress(X, Y, FGLSRegression)

# IteratedFGLS test
IteratedFGLSRegression = IteratedFGLS(10, intercept=true, robust=HCEVariance())
res = linregress(X, Y, IteratedFGLSRegression)

true