"""
An evaluator specialised to the travelling salesman problem.
"""
type TSPEvaluator <: Evaluator
  stage::AbstractString
  cities::Int
  threads::Int
  distance::Array{Float, 2}

  TSPEvaluator(st::AbstractString, n::Int, t::Int, d::Array{Float, 2}) =
    new(st, n, t, d)
end

"""
Contains a definition of a TSP evaluator.
"""
type TSPEvaluatorDefinition <: EvaluatorDefinition
  stage::AbstractString
  file::AbstractString
  threads::Int

  TSPEvaluatorDefinition() = new("", "", 1)
  TSPEvaluatorDefinition(s::AbstractString, f::AbstractString, t::Int) =
    new(s, f, t)
end

"""
Composes a TSP evaluator from its definition.
"""
function compose!(ev::TSPEvaluatorDefinition, p::Population)
  # Find the name of the stage that is used for evaluation if none is provided.
  if ev.stage == ""
    ev.stage = p.demes[1].species.genotype.label
  end

  # Load the contents of the file and convert into a list of co-ordinates.
  coords = open(ev.file, "r") do f
    c = Vector{Float}[]
    for ln in eachline(f)
      !isempty(ln) && push!(c, Float[float(x) for x in split(ln)])
    end
    c
  end
  num = length(coords)

  # Generate a distance matrix from the set of co-ordinates.
  matrix = Array(Float, num, num)
  for i = 1:num
    for j = 1:num
      matrix[i, j] = sqrt(sum((coords[i] - coords[j]) .^ 2))
    end
  end

  # Build and return the evaluator.
  return TSPEvaluator(ev.stage, num, ev.threads, matrix)
end

"""
Evaluates the fitness of a TSP tour for a pre-determined geometric problem.

**Positional Arguments**:

* `file::AbstractString`, the path to the file containing the co-ordinates of
  each of the cities for the geometric TSP instance being solved, relative to
  the current working directory.

**Keyword Arguments**:

* `threads::Int`, the number of threads that the evaluation process should be
  split across. Defaults to `1`.
"""
tsp(file::AbstractString; stage::AbstractString = "", threads = 1) =
  TSPEvaluatorDefinition(stage, file, max(1, threads))

"""
Computes the length of a provided tour and transforms it into a fitness value
according to a provided fitness scheme.
"""
function evaluate!(e::TSPEvaluator, sch::FitnessScheme, tour::Vector{Int})
  length = zero(Float)
  for i in 1:(e.cities - 1)
    length += e.distance[tour[i], tour[i+1]]
  end
  length += e.distance[tour[end], tour[1]]
  assign(sch, length)
end


using EconPDEs, Distributions, LookupTables

mutable struct BansalYaronModel
    # consumption process parameters
    μbar::Float64
    vbar::Float64
    κμ::Float64
    νμ::Float64
    κv::Float64
    νv::Float64

    # utility parameters
    ρ::Float64
    γ::Float64
    ψ::Float64

    # lookup
    μn::Int64
    vn::Int64
    τ::LookupTable2D
end

function BansalYaronModel(;μbar = 0.018, vbar = 0.00073, κμ = 0.252, νμ = 0.528, κv = 0.156, νv = 0.00354, ρ = 0.024, γ = 7.5, ψ = 1.5, μn=30, vn=30, τ=LookupTable2D(30, 30))
    BansalYaronModel(μbar, vbar, κμ, νμ, κv, νv, ρ, γ, ψ, μn, vn, τ)
end

function initialize_stategrid(m::BansalYaronModel)
    μbar = m.μbar ; vbar = m.vbar ; κμ = m.κμ ; νμ = m.νμ ; κv = m.κv ; νv = m.νv ; ρ = m.ρ ; γ = m.γ ; ψ = m.ψ
    μn = m.μn ; vn = m.vn

    σ = sqrt(νμ^2 * vbar / (2 * κμ))
    μmin = quantile(Normal(μbar, σ), 0.025)
    μmax = quantile(Normal(μbar, σ), 0.975)
    μs = range(μmin, stop = μmax, length = μn)

    α = 2 * κv * vbar / νv^2
    β = νv^2 / (2 * κv)
    vmin = quantile(Gamma(α, β), 0.025)
    vmax = quantile(Gamma(α, β), 0.975)
    vs = range(vmin, stop = vmax, length = vn)
    OrderedDict(:μ => μs, :v => vs)
end

function initialize_y(m::BansalYaronModel, stategrid::OrderedDict)
    OrderedDict(:p => ones(length(stategrid[:μ]), length(stategrid[:v])))
end

function update_τ(m::BansalYaronModel, stategrid::OrderedDict, τmat::Array{Float64, 2})
    m.τ = LookupTable2D(stategrid[:μ], stategrid[:v], τmat)
    m.τ.xvar = stategrid[:μ]
    m.τ.yvar = stategrid[:v]
end

function (m::BansalYaronModel)(state::NamedTuple, y::NamedTuple)
    μbar = m.μbar ; vbar = m.vbar ; κμ = m.κμ ; νμ = m.νμ ; κv = m.κv ; νv = m.νv ; ρ = m.ρ ; γ = m.γ ; ψ = m.ψ ; τ = m.τ
    μ, v = state.μ, state.v
    p, pμ, pv, pμμ, pμv, pvv = y.p, y.pμ, y.pv, y.pμμ, y.pμv, y.pvv

    τvar = lookup(τ, μ, v)
    τnew = μ^2 - v * μ

    # drift and volatility of c, μ, σ, p
    μc = μ
    σc = sqrt(v)
    μμ = κμ * (μbar - μ)
    σμ = νμ * sqrt(v)
    μv = κv * (vbar - v)
    σv = νv * sqrt(v)
    σp_Zμ = pμ / p * σμ
    σp_Zv = pv / p * σv
    σp2 = σp_Zμ^2 + σp_Zv^2
    μp = pμ / p * μμ + pv / p * μv + 0.5 * pμμ / p * σμ^2 + 0.5 * pvv / p * σv^2

    # Market price of risk κ
    κ_Zc = γ * σc
    κ_Zμ = - (1 - γ * ψ) / (ψ - 1) * σp_Zμ
    κ_Zv = - (1 - γ * ψ) / (ψ - 1) * σp_Zv
    κ2 = κ_Zc^2 + κ_Zμ^2 + κ_Zv^2

    # Risk free rate r
    r = ρ + μc / ψ - (1 + 1 / ψ) / 2 * γ * σc^2 - (γ * ψ - 1) / (2 * (ψ - 1)) * σp2

    # Market Pricing
    # pt = p * (1 / p + μc + μp - r - κ_Zc * σc - κ_Zμ * σp_Zμ - κ_Zv * σp_Zv)
    pt = p * (1 / p - ρ + (1 - 1 / ψ) * (μc - 0.5 * γ * σc^2) + μp + 0.5 * (1 / ψ - γ) / (1 - 1 / ψ) * σp2)

    # return (pt,), (μμ, μv), (p = p, r = r, κ_Zc = κ_Zc, κ_Zμ = κ_Zμ, κ_Zv = κ_Zv, σμ = σμ, σμ_Zv = 0.0, σv_Zμ = 0.0, σv = σv, μ = μ, v = v, σμ2 = σμ^2, σv2 = σv^2, σμv = 0.0, μμ = μμ, μv = μv, σp2 = σp2, σp_Zμ = σp_Zμ, σp_Zv = σp_Zv)
    return (pt,), (μμ, μv), (p = p, μμ = μμ, σμ = σμ, μv = μv, σv = σv, τnew = τnew, μ = μ, v = v)
end

# Bansal Yaron (2004)
m = BansalYaronModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
update_τ(m, stategrid, zeros(m.μn, m.vn))
y, result, distance = pdesolve(m, stategrid, y0)

update_τ(m, stategrid, result[:τnew])


ranges, my_ticket, all_tickets = split(read("day16/input.txt", String), "\n\n")
# ranges, my_ticket, all_tickets = split(read("day16/small.txt", String), "\n\n")

ranges = split(ranges, "\n")
my_ticket = parse.(Int, split(split(my_ticket, "\n")[2], ","))
all_tickets = split(all_tickets, "\n")[2:end]

all_ranges = []

for r in ranges
    s1, e1, s2, e2 = parse.(Int, match(r"(\d+)-(\d+) or (\d+)-(\d+)", r).captures)
    push!(all_ranges, Set([s1:e1..., s2:e2...]))
end

all_valid_nums = union(all_ranges...)

filter!(r -> all(n -> (n in all_valid_nums), parse.(Int, split(r, ","))), all_tickets)

all_tickets = parse.(Int, vcat(permutedims.(split.(all_tickets, ","))...))

matches_col(c, valid_set) = all(n -> n in valid_set, all_tickets[:, c])

matches_per_rule = [Set(filter(c -> matches_col(c, all_ranges[i]), 1:size(all_tickets, 2))) for i in 1:length(all_ranges)]

println(matches_per_rule)

function range_solver(start_col, taken_cols)
    if start_col == length(all_ranges)+1
        return []
    end
    for c in 1:size(all_tickets, 2)
        if c in taken_cols
            continue
        end
        if c in matches_per_rule[start_col]
            result = range_solver(start_col+1, Set([taken_cols..., c]))
            if result !== nothing
                return vcat([c], result)
            end
        end
    end
end

cols = range_solver(1, Set())

println(prod(i -> my_ticket[i], cols[1:6]))

# not proud of this guy, but it does the job


type KMeansOutput
  assignments::Vector{Int64}
  centers::Matrix{Float64}
  iterations::Int64
  rss::Float64
  converged::Bool
end

KMeansOutput() = KMeansOutput(Int64[], Float64[], 0, 0.0, false)

function show(io::IO, results::KMeansOutput)
  println(io, "Cluster Assignments:")
  println(io, results.assignments)
  println(io)
  println(io, "Cluster Centers:")
  println(io, results.centers)
  println(io)
  print(io, "Number of Iterations: ")
  println(io, results.iterations)
  println(io)
  print(io, "Final RSS: ")
  println(io, results.rss)
  println(io)
  print(io, "Algorithm Converged: ")
  println(io, results.converged)
  println(io)
end

type DPMeansOutput
  assignments::Vector{Int64}
  centers::Matrix{Float64}
  k::Int64
  iterations::Int64
  rss::Float64
  converged::Bool
end

DPMeansOutput() = DPMeansOutput(Int64[], Float64[], 0, 0, 0.0, false)

function show(io::IO, results::DPMeansOutput)
  println(io, "Cluster Assignments:")
  println(io, results.assignments)
  println(io)
  println(io, "Cluster Centers:")
  println(io, results.centers)
  println(io)
  println(io, "Cluster Created:")
  println(io, results.k)
  println(io)
  print(io, "Number of Iterations: ")
  println(io, results.iterations)
  println(io)
  print(io, "Final RSS: ")
  println(io, results.rss)
  println(io)
  print(io, "Algorithm Converged: ")
  println(io, results.converged)
end

## Mostly following R's hclust class
type Hclust{T<:Real}
    merge::Matrix{Int}
    height::Vector{T}
    order::Vector{Int}
    labels::Vector
    method::Symbol
end


"""
heq_lu_test()

Does the H-equation module do what it's supposed to?
"""
function heq_lu_test()
    n = 32
    c = 0.5
    FS = ones(n)
    x0 = ones(n)
    FPS = ones(n, n)
    FPSS = ones(Float32, n, n)
    hdata = heqinit(x0, c)
    nsoloutfd = nsolheq(x0, FS, FPS, hdata)
    nsoloutbos = nsol(heqbos!, x0, FS, FPS; pdata = c, sham = 1)
    dbos = norm(nsoloutbos.solution - nsoloutfd.solution)
    bosok = dbos < 1.e-7
    if ~bosok
        println("Bosma and DeRooij test fails in H-equation")
    end
    nsoloutsp = nsolheq(x0, FS, FPSS, hdata; diff = :exact)
    nsoloutdp = nsolheq(x0, FS, FPS, hdata; diff = :exact)
    dsp = norm(nsoloutsp.history - nsoloutfd.history)
    ddp = norm(nsoloutdp.history - nsoloutfd.history)
    dsolsp = norm(nsoloutsp.solution - nsoloutfd.solution)
    dsoldp = norm(nsoloutsp.solution - nsoloutdp.solution)
    spok = (dsp < 1.e-7) && (dsolsp < 1.e-9) && (ddp < 1.e-7) && (dsoldp < 1.e-9)
    if ~spok
        println("Mixed precision test fails in H-equation")
    end
    #
    # change c and use old solution as initial iterate
    #
    h5 = nsoloutfd.solution
    setc!(hdata, 0.7)
    nsoloutfd7 = nsolheq(h5, FS, FPS, hdata)
    contok = (nsoloutfd7.history[4] < 1.e-12)
    if ~contok
        println("Update c test fails in H-equation")
    end
    heqok = spok && bosok && contok
    if ~heqok
        println("H-equation Chapter 2 test fails")
    end
    return heqok
end


# This file is part of the TaylorSeries.jl Julia package, MIT license
#
# Luis Benet & David P. Sanders
# UNAM
#
# Handles Taylor series of arbitrary but finite order

"""
    TaylorSeries

A Julia package for Taylor expansions in one or more independent variables.

The basic constructors are [`Taylor1`](@ref) and [`TaylorN`](@ref);
see also [`HomogeneousPolynomial`](@ref).

"""
module TaylorSeries

using InteractiveUtils: subtypes
using SparseArrays: SparseMatrixCSC
using Markdown
import LinearAlgebra: norm, mul!
if VERSION ≥ v"1.0.0"
    # export gradient
else
    import LinearAlgebra: gradient
end

import Base: ==, +, -, *, /, ^

import Base: iterate, size, eachindex, firstindex, lastindex,
    eltype, length, getindex, setindex!

import Base: zero, one, zeros, ones, isinf, isnan, iszero,
    convert, promote_rule, promote, show,
    real, imag, conj, adjoint,
    rem, mod, mod2pi, abs, abs2,
    sqrt, exp, log, sin, cos, tan,
    asin, acos, atan, sinh, cosh, tanh,
    power_by_squaring,
    rtoldefault, isfinite, isapprox, rad2deg, deg2rad

export Taylor1, TaylorN, HomogeneousPolynomial, AbstractSeries

export getcoeff, derivative, integrate, differentiate,
    evaluate, evaluate!, inverse,
    show_params_TaylorN, show_monomials, displayBigO, use_show_default,
    get_order, get_numvars,
    set_variables, get_variables,
    get_variable_names, get_variable_symbols,
    # jacobian, hessian, jacobian!, hessian!,
    ∇, taylor_expand, update!, constant_term

include("parameters.jl")
include("hash_tables.jl")
include("constructors.jl")
include("conversion.jl")
include("auxiliary.jl")
include("arithmetic.jl")
include("power.jl")
include("functions.jl")
include("other_functions.jl")
include("evaluate.jl")
include("calculus.jl")
include("dictmutfunct.jl")
include("printing.jl")

end # module


module PlanarConvexHullsTest

using PlanarConvexHulls
using StaticArrays
using Test
using Random
using LinearAlgebra
using Statistics

using PlanarConvexHulls: vertex_order

const Point{T} = SVector{2, T}

@testset "is_ordered_and_strongly_convex" begin
    @test is_ordered_and_strongly_convex([Point(1, 2)], CCW)
    @test is_ordered_and_strongly_convex([Point(1, 2), Point(3, 4)], CCW)
    @test is_ordered_and_strongly_convex([Point(1, 2), Point(3, 4), Point(2, 4)], CCW)
    @test !is_ordered_and_strongly_convex([Point(1, 2), Point(3, 4), Point(2, 3)], CCW) # on a line

    v1 = Point(0, 0)
    v2 = Point(1, 0)
    v3 = Point(1, 1)
    v4 = Point(0, 1)
    vertices = [v1, v2, v3, v4]
    for i in 0 : 4
        shifted = circshift(vertices, i)
        @test is_ordered_and_strongly_convex(shifted, CCW)
        @test !is_ordered_and_strongly_convex(reverse(shifted), CCW)
        @test is_ordered_and_strongly_convex(reverse(shifted), CW)
    end

    for i in eachindex(vertices)
        for j in eachindex(vertices)
            if i != j
                vertices′ = copy(vertices)
                vertices′[i] = vertices[j]
                vertices′[j] = vertices[i]
                @test !is_ordered_and_strongly_convex(vertices′, CCW)
            end
        end
    end
end

@testset "area" begin
    @test area(ConvexHull{CCW}(SVector((SVector(1, 1),)))) == 0

    @test area(ConvexHull{CCW}(SVector(SVector(1, 1), SVector(2, 3)))) == 0

    triangle = ConvexHull{CCW}(SVector(Point(1, 1), Point(2, 1), Point(3, 3)))
    @test area(triangle) == 1.0

    square = ConvexHull{CCW}(SVector(Point(1, 1), Point(4, 1), Point(4, 3), Point(1, 3)))
    @test area(square) == 3 * 2
end

@testset "in" begin
    @testset "point" begin
        rng = MersenneTwister(1)
        p = SVector(1, 1)
        C = ConvexHull{CCW}(SVector((p,)))
        @test p ∈ C
        @test Float64.(p) ∈ C
        for i in 1 : 10
            @test p + SVector(randn(rng), randn(rng)) ∉ C
        end
    end

    @testset "line segment" begin
        p1 = SVector(1, 1)
        p2 = SVector(3, 5)
        linesegment = ConvexHull{CCW}([p1, p2])
        @test p1 ∈ linesegment
        @test p2 ∈ linesegment
        @test div.(p1 + p2, 2) ∈ linesegment
        @test p1 + 2 * (p2 - p1) ∉ linesegment
    end

    @testset "triangle" begin
        rng = MersenneTwister(1)
        triangle = ConvexHull{CCW}(SVector(Point(1, 1), Point(2, 1), Point(3, 3)))
        for p in vertices(triangle)
            @test p ∈ triangle
        end
        for i = 1 : 100_000
            weights = normalize(rand(rng, SVector{3}), 1)
            p = reduce(+, vertices(triangle) .* weights)
            @test p ∈ triangle
        end
    end

    @testset "rectangle" begin
        rng = MersenneTwister(1)
        width = 4
        height = 3
        origin = Point(2, 4)
        rectangle = ConvexHull{CCW}(map(x -> x + origin, SVector(Point(0, 0), Point(width, 0), SVector(width, height), SVector(0, height))))
        for p in vertices(rectangle)
            @test p ∈ rectangle
        end
        for i = 1 : 100_000
            p = origin + SVector(width * rand(rng), height * rand(rng))
            @test p ∈ rectangle
        end
        for i = 1 : 10
            p = origin + SVector(width * rand(rng), height * rand(rng))
            @test setindex(p, origin[1] + width + rand(rng), 1) ∉ rectangle
            @test setindex(p, origin[1] - rand(rng), 1) ∉ rectangle
            @test setindex(p, origin[2] + height + rand(rng), 2) ∉ rectangle
            @test setindex(p, origin[2] - rand(rng), 2) ∉ rectangle
        end
    end
end

@testset "centroid" begin
    @testset "point" begin
        p = Point(1, 2)
        @test centroid(ConvexHull{CCW}([p])) === Float64.(p)
    end

    @testset "line segment" begin
        p1 = SVector(1, 1)
        p2 = SVector(3, 5)
        linesegment = ConvexHull{CCW}([p1, p2])
        @test centroid(linesegment) == Point(2.0, 3.0)
    end

    @testset "triangle" begin
        triangle = ConvexHull{CCW}(SVector(Point(1, 1), Point(2, 1), Point(3, 3)))
        @test centroid(triangle) ≈ mean(vertices(triangle)) atol=1e-15
    end
end

function convex_hull_alg_test(hull_alg!)
    @testset "random $order" for order in [CCW, CW]
        hull = ConvexHull{order, Float64}()
        rng = MersenneTwister(2)
        for n = 1 : 10
            sizehint!(hull, n) # just to get code coverage
            for _ = 1 : 10_000
                points = [rand(rng, Point{Float64}) for i = 1 : n]
                hull_alg!(hull, points)
                @test is_ordered_and_strongly_convex(vertices(hull), order)
            end
        end
    end

    @testset "collinear input $order" for order in [CCW, CW]
        hull = ConvexHull{order, Float64}()
        points = [Point(0., 0.), Point(0., 1.), Point(0., 2.), Point(1., 0.), Point(1., 1.), Point(1., 2.)]
        for i = 1 : 10
            shuffle!(points)
            hull_alg!(hull, points)
            @test is_ordered_and_strongly_convex(vertices(hull), order)
            @test isempty(symdiff(vertices(hull), [Point(0., 0.), Point(1., 0.), Point(1., 2.), Point(0., 2.)]))
        end
    end

    @testset "numerical problem $order" for order in [CCW, CW]
        hull = ConvexHull{order, Float64}()
        points = Point{Float64}[
            [-0.17559141363793285, 0.1777893970092983],
            [-0.17574071124767324, 0.04921003844438379],
            [0.08480367424885946, 0.17745906024313485],
            [0.08465437663911904, 0.04887970167822034],
            [-0.13059226250542266, 0.17773231116261107],
            [-0.13074156011516305, 0.04915295259769656],
            [-0.1753504285942606, -0.1772960226863406],
            [-0.1756647405301049, -0.04875022362480264],
            [0.08503864201675085, -0.17672638211747324],
            [0.08472433008090652, -0.04818058305593528],
            [-0.13035231731355584, -0.17719758249586814],
            [-0.13066662924940015, -0.04865178343433017]]
        hull_alg!(hull, points)
        @test is_ordered_and_strongly_convex(vertices(hull), order)
    end
end

@testset "jarvis_march!" begin
    convex_hull_alg_test(jarvis_march!)
end

@testset "closest_point $order" for order in [CCW, CW]
    hull = ConvexHull{order, Float64}()
    rng = MersenneTwister(3)
    for n = 1 : 10
        for _ = 1 : 100
            points = [rand(rng, Point{Float64}) for i = 1 : n]
            jarvis_march!(hull, points)

            for point in vertices(hull)
                closest = closest_point(point, hull)
                @test closest == point
            end

            for _ = 1 : 100
                point = rand(rng, Point{Float64})
                closest = closest_point(point, hull)
                if point ∈ hull
                    @test closest == point
                else
                    closest_to_closest = closest_point(closest, hull)
                    @test closest_to_closest ≈ closest atol=1e-14
                end
            end
        end
    end

    # visualize = true
    # if visualize
    #     flush(stdout)
    #     hull = ConvexHull{CCW, Float64}()
    #     rng = MersenneTwister(3)
    #     n = 5
    #     points = [rand(rng, Point{Float64}) for i = 1 : n]
    #     jarvis_march!(hull, points)
    #     projections = map(1 : 1_000_000) do _
    #         closest_point(rand(rng, Point{Float64}), hull)
    #     end
    #     plt = scatterplot(getindex.(projections, 1), getindex.(projections, 2))
    #     scatterplot!(plt, getindex.(vertices(hull), 1), getindex.(vertices(hull), 2), color = :red)
    #     display(plt)
    # end
end

function make_oversized(A::AbstractMatrix, b::AbstractVector)
    n = length(b)
    n_oversized = n + 1
    similar(A, (n_oversized, 2)), similar(b, n_oversized)
end

function make_oversized(A::SMatrix, b::SVector{n}) where n
    n_oversized = n + 1
    similar(A, Size(n_oversized, 2)), similar(b, Size(n_oversized))
end

function hreptest(hull::H, rng) where {H<:ConvexHull}
    A, b = hrep(hull)
    A_oversized, b_oversized = make_oversized(A, b)
    hrep!(A_oversized, b_oversized, hull)
    for _ = 1 : 100
        testpoint = rand(rng, Point{Float64})
        @test (testpoint ∈ hull) == all(A * testpoint .<= b) == all(A_oversized * testpoint .<= b_oversized)
    end
    if Length(vertices(hull)) !== Length{StaticArrays.Dynamic()}()
        @test(@allocated(hrep(hull)) == 0)
    end
end

@testset "hrep $order" for order in [CCW, CW]
    dynamichull = ConvexHull{order, Float64}()
    rng = MersenneTwister(4)
    for n = 2 : 10
        for _ = 1 : 100
            points = [rand(rng, Point{Float64}) for i = 1 : n]
            jarvis_march!(dynamichull, points)
            hreptest(dynamichull, rng)

            h = num_vertices(dynamichull)
            statichull = ConvexHull{order}(SVector{h}(vertices(dynamichull)))
            hreptest(statichull, rng)
        end
    end
end

struct CustomPoint{T} <: FieldVector{2, T}
    x::T
    y::T
end

@testset "Constructors" begin
    width = 4
    height = 3
    origin = Point(2, 4)
    verts = map(x -> x + origin, [Point(0, 0), Point(width, 0), SVector(width, height), SVector(0, height)])
    rectangle = ConvexHull{CW}(reverse(verts))

    srectangle = SConvexHull{4, Float64}(rectangle)
    drectangle = DConvexHull{Float64}(rectangle)
    @test vertices(srectangle) isa SVector{4}
    @test vertices(drectangle) isa Vector
    for hull in [srectangle, drectangle]
        @test vertex_order(hull) == CCW
        @test eltype(hull) == Float64
        @test is_ordered_and_strongly_convex(vertices(hull), vertex_order(hull))
    end

    let hull = ConvexHull{CCW, Float64}(verts)
        @test vertices(hull) == verts
        @test eltype(hull) == Float64
    end
    let hull = ConvexHull{CCW, Float64, CustomPoint{Float64}}(verts)
        @test vertices(hull) == verts
        @test eltype(hull) == Float64
        @test typeof(vertices(hull)) == Vector{CustomPoint{Float64}}
    end
    let hull = ConvexHull{CCW, Float64, Point{Float64}, SVector{4, Point{Float64}}}(verts)
        @test vertices(hull) == verts
        @test eltype(hull) == Float64
        @test typeof(vertices(hull)) == SVector{4, SVector{2, Float64}}
    end

    @test isempty(vertices(DConvexHull{Float64}()))
end

@testset "benchmarks" begin
    include(joinpath("..", "perf", "runbenchmarks.jl"))
end

end # module


## Assumptions
## http://docs.sympy.org/0.7.2/modules/assumptions/index.html

## TODO: Deprecate Qeven, etc. in favor of Q.even

"""
refine: http://docs.sympy.org/dev/modules/assumptions/refine.html
"""
refine(ex, assumpts...) = sympy_meth(:refine, ex, assumpts...)
export refine

"""

ask. Returns true, false or nothing

ask(Qinteger(x*y), And(Qinteger(x), Qinteger(y)))
## really slow isprime:
filter(x -> ask(Qprime(x)), [1:1000])

"""
ask(x::Sym, args...) = sympy_meth(:ask, x, args...)
export ask

## should we support & and | for (sym,sym) pairs? Not sure
## dependso on what happens to | (x==0) and ex == x usage for ==
## for now, we can combine terms logically with And, Or, Not...
## these are in logic module

## simple methods (x, args) -> y (y coercion happens via PyCall)
logic_sympy_methods = (
                     :And, :Or, :Not,
                     :Xor, :Nand, :Nor, :Implies,
                     :Equivalent,
                     :satisfiable
                     )

## This is a bit messed up, as we use Qeven in place of Q.even, ...
## This is now deprecated. See below
Q_predicates = (:even, :odd, :prime, :nonzero,
                :complex, :extended_real, :imaginary, :infinitesimal,
                :integer, :irrational, :real,
                :positive, :negative,
                :bounded, :commutative)
Q_nms = [@compat(Symbol("Q" * string(i))) for i in Q_predicates]

for (fn, meth) in zip(Q_nms, Q_predicates)
    nm = string(meth)
    @eval begin
        @doc """
`$($nm)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($nm)
""" ->
        ($fn)(x::SymbolicObject) = PyCall.pyeval("f(x)", Sym, f=SymPy.sympy[:Q][($nm)], x=SymPy.project(x))::Sym
    end
    eval(Expr(:export, fn))
end

@deprecate Qeven(x) Q.even(x)
@deprecate Qodd(x) Q.odd(x)
@deprecate Qprime(x) Q.prime(x)
@deprecate Qnonzero(x) Q.nonzero(x)
@deprecate Qcomplex(x) Q.complex(x)
@deprecate Qextended_real(x) Q.extended_real(x)
@deprecate Qimaginary(x) Q.imaginary(x)
@deprecate Qinfinitesimal(x) Q.infinitesimal(x)
@deprecate Qinteger(x) Q.integer(x)
@deprecate Qirrational(x) Q.irrational(x)
@deprecate Qreal(x) Q.real(x)
@deprecate Qpositive(x) Q.positive(x)
@deprecate Qnegative(x) Q.negative(x)
@deprecate Qbounded(x) Q.bounded(x)
@deprecate Qcommutative(x) Q.commutative(x)

## We make a module Q to hold the assumptions
## this follows this page http://docs.sympy.org/0.7.5/_modules/sympy/assumptions/ask.html
module Q
import SymPy
import PyCall

##http://docs.sympy.org/dev/_modules/sympy/assumptions/ask.html#ask
Q_predicates = (:antihermitian,
                :bounded, :finite, # bounded deprecated
                :commutative,
#                :complex,
                :composite,
                :even,
                :extended_real,
                :hermitian,
                :imaginary,
                :infinitesimal,
                :infinity, :infinite, # :infinity deprecated
                :integer,
                :irrational,
                :rational,
                :algebraic,
                :transcendental,
                :negative,
                :nonzero, :zero,
                :positive,
                :prime,
                :real,
                :odd,
                :is_true,
                :nonpositive,
                :nonnegative,
                :symmetric,
                :invertible,
                :singular,
                :orthogonal,
                :unitary,
                :normal,
                :positive_definite,
                :upper_triangular,
                :lower_triangular,
                :diagonal,
                :triangular,
                :unit_triangular,
                :fullrank,
                :square,
                :real_elements,
                :complex_elements,
                :integer_elements)

for meth in Q_predicates
        nm = string(meth)
        @eval begin
            @doc """
`$($nm)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($nm)
""" ->
            ($meth)(x) = PyCall.pyeval("f(x)", SymPy.Sym, f=SymPy.sympy[:Q][($nm)], x=SymPy.project(x))::SymPy.Sym
        end
    end
end

export Q


# Example 5.4 of
# A. Ahmadi, R. Jungers, P. Parrilo and M. Roozbehani,
# "Joint spectral radius and path-complete graph Lyapunov functions."
# SIAM J. CONTROL OPTIM 52(1), 687-717, 2014.
# The JSR is 3.917384715148

@testset "[AJPR14] Example 5.4" begin
    s = discreteswitchedsystem([[-1 -1; -4 0], [3 3; -2 1]])
    @test sprint(show, s) == "Hybrid System with automaton OneStateAutomaton(2)"
    qub = 4.31959610746
    @test quicklb(s) ≈ 3
    @test quickub(s) ≈ qub
    qb = quickb(s)
    @test typeof(qb) == NTuple{2,Float64}
    @test qb[1] ≈ 3
    @test qb[2] ≈ qub
    # pnorm=Inf: [17.2409,7.42377,5.65853,4.92192,4.52597,4.28099,4.11275,3.99139,3.8991,3.827,3.76891,3.72121,3.6813,3.64744,3.61834,3.59307,3.57092,3.55135,3.53393,3.51832,3.50426,3.49152,3.47994]
    # pnorm=1:   [15.1245,7.35044,5.47402,4.84722,4.45847,4.23242,4.0709,3.95644,3.8687,3.80015,3.74485,3.69944,3.66141,3.62914,3.6014,3.5773,3.55616,3.53748,3.52085,3.50595,3.49253,3.48037,3.4693,3.45919]
    # pnorm=2:   [14.1535,6.81871,5.32651,4.7096,4.37067,4.15723,4.01117,3.90488,3.82407,3.7606,3.70945,3.66735,3.6321,3.60215,3.57639,3.554,3.53437,3.517,3.50154,3.48768,3.47519]

    @test pradius(s, 2) ≈ 3.234535151244 rtol=1e-9
    @test pradius(s, 2, BruteForce(), pnorm=2) ≈ 3.632097274822649

    pρlb = [3.23453515151, 3.42756015606, 3.54374928704]
    pρub = [4.57432347886, 4.07607892469, 3.97772408342]
    @testset "p-radius with $algo" for algo in (VeroneseLift(), KroneckerLift())
        for d in 1:3
            lb, ub = pradiusb(s, 2d, algo)
            @test isapprox(lb, pρlb[d])
            @test isapprox(ub, pρub[d])
        end
    end

    tmp = periodicswitching(s, [1, 2])
    tmp = periodicswitching(tmp.s, tmp.period, round(tmp.growthrate, digits=4)) # Avoid difference of floating point rounding
    @test sprint(show, tmp) == "PSW(3.9174, [1, 2])" #"Periodic switching of growth rate 3.9174 and modes: [1, 2] for the transitions [1, 2]"
    smp = periodicswitching(s, [1, 2])
    #@test smp != tmp
    @test smp != periodicswitching(s, [2, 1, 2, 1])
    @test smp != periodicswitching(s, [1, 2, 1, 2])
    @test smp != periodicswitching(s, [1, 1])
    #@test smp != periodicswitching(discreteswitchedsystem([[-1 -1; -4 0],[3 3; -2 1]]), [1, 2])
    @testset "JSR with $factory" for factory in sdp_factories
        sosdata(s).lb = 0
        tol = ismosek(factory) ? 1e-5 : 5e-4
        lb, ub = soslyapb(s, 1, factory=factory, tol=tol)
        @test log(lb) ≈ log(2.814640557) rtol=tol
        @test log(ub) ≈ log(3.980502849) rtol=tol
        lb, ub = soslyapb(s, 2, factory=factory, tol=tol)
        @test log(lb) ≈ log(3.299750624) rtol=tol
        @test log(ub) ≈ log(3.924086919) rtol=tol
        @test_throws ArgumentError sosbuildsequence(s, 1, v_0 = 2)
        for d in (1, 2)
            for p_0 in (:Primal, :Random)
                seq = sosbuildsequence(s, d, p_0=p_0)
                psw = findsmp(seq)
                @test psw !== nothing
                @test psw == smp
            end
        end
    end
    if isempty(sdp_factories)
        @test getub(s) ≈ pρub[end]
    else
        @test log(getub(s)) ≈ log(3.924086919) rtol=5e-4
        @test getsmp(s) == smp
        @test getlb(s) ≈ 3.917384715148 rtol=1e-12
    end
end


@testset "exceptions" begin

    @test_broken @interval("[+infinity]") === emptyinterval()

    @test interval(+Inf,-Inf) === emptyinterval()

    @test_broken interval(nai()) === emptyinterval()

    @test_broken @interval("[1.0000000000000001, 1.0000000000000002]") === Interval(1.0, 0x1.0000000000001p+0)

end


# Performance tuning

v = [[1, 2, 3], [1, 2, 3], [1, 2, 3], [2, 3, 5], [1, 3, 4], [1, 2, 5], [2, 3, 4], [1, 4, 5], [3, 4, 5]]
v = [[1, 2], [1, 3], [2, 4], [1, 2, 3], [1, 2, 4], [1, 3, 4], [1, 2, 3, 4], [1, 2, 3, 5], [2, 3, 4, 6]]
# v = [rand([1, 2, 3, 4, 5], 10) for x = 1:1000];
# @code_warntype freq_itemsets(v, 0.5)

# testing freq_itemsets()
v = [[1, 2, 3], [1, 2, 3], [1, 2, 3],  [1, 2, 5], [1, 3, 4], [1, 4, 5], [2, 3, 4], [2, 3, 4], [2, 3, 5], [3, 4, 5]]
freq_itemsets(v, 0.2)

# testing gen_rules()
rules = Vector{Rule}(0)
freq = [1, 2]
consq = [Int[], [1], [2], [3], [4], [5]]
trans = [[1, 2], [1, 3], [2, 4], [1, 2, 3], [1, 2, 4], [1, 3, 4], [1, 2, 3, 4], [1, 2, 3, 5], [2, 3, 4, 6]]
ap_genrules!(freq, consq, trans, 0.20, 0.01, rules)
rules

rules = Vector{Rule}(0)
freq = [1, 2, 3, 4]
consq = [[1], [2], [3], [4]]
trans = [[1], [2], [1, 2], [1, 3], [2, 4], [1, 2, 3], [1, 2, 4], [1, 3, 4], [1, 2, 3, 4], [1, 2, 3, 5], [2, 3, 4, 6]]
ap_genrules!(freq, consq, trans, 0.2, 0.01, rules)
rules

rule1 = Vector{Rule}(0)
@time gen_onerules!([1], [[1], [2], [3]], rule1, [[1, 2], [2, 3], [1, 3]], 0.2, 0.01)

v = [rand([1, 2, 3, 4, 5, 6], 10) for x = 1:100];
rules = apriori(v, 0.01, 0.08)
show_rulestats(rules)

# Example using UCI Adult dataset

using DataFrames

d = readtable("./data/adult.csv", nastrings = ["?"])
complete_cases!(d)

dmat = convert(Array, d)
transactions = make_transactions(dmat)
@time rules = apriori(transactions[1:1000], 0.60, 0.8);

display(rules)


#include("example_nnet.jl")
using LinearAlgebra
function ComputeValid(predict_func::Function, W, xvalid, yvalid, valid)
    yhat = predict_func(W,[ones(size(xvalid,1)) xvalid],nHidden)
    push!(valid,sum((yhat-yvalid).^2)/length(yvalid))

end

function BS_Select(;validation_array,f_array, gradcalc::Function,xtrain,ytrain,W,nHidden, alist, Blist)
    #Get most recent step and batch sizes
    a=alist[end]
    B=Blist[end]
    n,d=size(xtrain)

    #Get Gradient Angle
    i=rand(1:n,B)
    j=rand(1:n,B)
    _,gi=gradcalc(W, xtrain[i,:],ytrain[i],nHidden)
    _,gj=gradcalc(W,xtrain[j,:],ytrain[j], nHidden)
    costheta=dot(gi,gj)/(norm(gi)*norm(gj))
    #theta=acos(costheta)

    #alpha decreases
    if length(validation_array) >= 4 && validation_array[end-2]>validation_array[end]
        a_adjust = (1 - abs(validation_array[end-2] - validation_array[end])/validation_array[end-2])^(1/4)
        a = a*a_adjust
    end
    push!(alist,a)
    push!(Blist,B)
    print(a," ")

    return a,B
end

function SGDBabysitter(gradcalc::Function, nn_predict::Function, maxIter, nHidden,
                        nParams, xtrain, ytrain, xvalid, yvalid )

    n = size(xtrain,1)
    valid=[]
    flist=[]
    alist=[]
    Blist=[]
    #placehold for initialization
    B=100
    a = 1e-4
    push!(Blist,B)
    push!(alist,a)
    W= randn(nParams,1)

    for t in 1:maxIter

    	i = rand(1:n,B)
    	f,g = gradcalc(W, xtrain[i,:], ytrain[i], nHidden)
    	W = W - a*g
        #Every few iterations, compute validation and append f,g,validation

          if (mod(t-1,round(maxIter/100)) == 0)
              # print("Training iteration = $(t-1) ")
              push!(flist,f)
              ComputeValid(nn_predict, W, xvalid, yvalid, valid)
          end
        if (mod(t-1,round(maxIter/50)) == 0)
            a,B=BS_Select(validation_array=valid,f_array=flist, gradcalc=gradcalc,xtrain=xtrain,ytrain=ytrain,W=W,nHidden=nHidden, alist=alist, Blist=Blist)

          end
    end

    display(plot(1:length(valid), valid))

    return W, valid
end


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

"""
    LCV(ranges, nfolds=10)

Likelihood cross-validation with parameter ranges
`ranges` and number of folds `nfolds`.

## References

* Sugiyama et al. 2008. Direct importance estimation for
  covariate shift adaptation.

### Author

* Júlio Hoffimann (julio.hoffimann@gmail.com)
"""
struct LCV <: EstimatorFitter
  ranges::NamedTuple
  nfolds::Int
end

LCV(ranges::NamedTuple) = LCV(ranges, 10)

function _fit(::Type{<:KLIEP}, x_nu, x_de,
              fitter::EstimatorFitter,
              optlib::Type{<:OptimizationLibrary})
  # retrieve parameters
  ranges = fitter.ranges
  nfolds = fitter.nfolds
  npts   = length(x_nu)

  @assert nfolds ≤ npts "number of folds must be smaller than number of numerator samples"

  # partition numerator samples into folds
  folds = collect(Iterators.partition(1:npts, npts ÷ nfolds))

  # initialize loss and optimal estimator
  Ĵmax, dre⭐ = -Inf, KLIEP()

  # loop over hyperparameters
  for σ in ranges.σ, b in ranges.b
    # density ratio estimator
    dre = KLIEP(σ=σ, b=b)

    # estimate loss with cross-validation
    Ĵₖ = map(1:nfolds) do k
      # training and hold-out samples
      train = [ind for i in vcat(1:k-1, k+1:nfolds) for ind in folds[i]]
      hold  = folds[k]

      # perform estimation with training samples
      r = densratiofunc(x_nu[train], x_de, dre, optlib=optlib)

      # evaluate loss with hold-out samples
      mean(log(r(x_nu[j])) for j in hold)
    end

    # mean loss for hyperparameters
    Ĵ = mean(Ĵₖ)

    # update and continue
    if Ĵ > Ĵmax
      Ĵmax = Ĵ
      dre⭐ = dre
    end
  end

  # optimal estimator
  dre⭐
end


export Camera, get_primary_rays

# ------ #
# Camera #
# ------ #

"""
    FixedCameraParams

The Parameters of the [`Camera`](@ref) which should not be updated in inverse rendering
are wrapped into this object.

### Fields:

* `vup`    - Stores the World UP Vector (In most cases this is the Y-Axis Vec3(0.0, 1.0, 0.0))
* `width`  - Width of the image to be rendered
* `height` - Height of the image to be rendered
"""
struct FixedCameraParams{T} <: FixedParams
    vup::Vec3{T}
    width::Int
    height::Int
end

Base.show(io::IO, fcp::FixedCameraParams) =
    print(io, "    Fixed Parameters:\n        World UP - ", fcp.vup,
          "\n        Screen Dimensions - ", fcp.height, " × ", fcp.width)

Base.zero(fcp::FixedCameraParams) = FixedCameraParams(zero(fcp.vup), zero(fcp.width),
                                                      zero(fcp.height))

"""
    Camera

The Perspective View Camera Model

### Fields:

* `lookfrom`    - Position of the Camera in 3D World Space
* `lookat`      - Point in the 3D World Space where the Camera is Pointing
* `vfov`        - Field of View of the Camera
* `focus`       - The focal length of the Camera
* `fixedparams` - An instance of [`FixedCameraParams`](@ref)

### Available Constructors:

* `Camera(;lookfrom = Vec3(0.0f0), lookat = Vec3(0.0f0),
           vup = Vec3(0.0f0, 1.0f0, 0.0f0), vfov = 90.0f0,
           focus = 1.0f0, width = 128, height = 128)`
* `Camera(lookfrom::Vec3{T}, lookat::Vec3{T}, vup::Vec3{T}, vfov::R,
          focus::R, width::Int, height::Int) where {T<:AbstractArray, R<:Real}`
"""
struct Camera{T}
    lookfrom::Vec3{T}
    lookat::Vec3{T}
    vfov::T
    focus::T
    fixedparams::FixedCameraParams{T}
end

Base.show(io::IO, cam::Camera) =
    print(io, "CAMERA Configuration:\n    Lookfrom - ", cam.lookfrom,
          "\n    Lookat - ", cam.lookat, "\n    Field of View - ", cam.vfov[],
          "\n    Focus - ", cam.focus[], "\n", cam.fixedparams)

@diffops Camera

Camera(;lookfrom = Vec3(0.0f0), lookat = Vec3(0.0f0),
        vup = Vec3(0.0f0, 1.0f0, 0.0f0), vfov = 90.0f0,
        focus = 1.0f0, width = 128, height = 128) =
    Camera(lookfrom, lookat, vup, vfov, focus, width, height)

function Camera(lookfrom::Vec3{T}, lookat::Vec3{T}, vup::Vec3{T}, vfov::R,
                focus::R, width::Int, height::Int) where {T<:AbstractArray, R<:Real}
    fixedparams = FixedCameraParams(vup, width, height)
    return Camera(lookfrom, lookat, [vfov], [focus], fixedparams)
end

"""
    get_primary_rays(c::Camera)

Takes the configuration of the camera and returns the origin and the direction
of the primary rays.
"""
function get_primary_rays(c::Camera)
    width = c.fixedparams.width
    height = c.fixedparams.height
    vup = c.fixedparams.vup
    vfov = c.vfov[]
    focus = c.focus[]

    aspect_ratio = width / height
    half_height = tan(deg2rad(vfov / 2))
    half_width = typeof(half_height)(aspect_ratio * half_height)

    origin = c.lookfrom
    w = normalize(c.lookfrom - c.lookat)
    u = normalize(cross(vup, w))
    v = normalize(cross(w, u))

    # Lower Left Corner
    llc = origin - half_width * focus * u - half_height * focus * v - w
    hori = 2 * half_width * focus * u
    vert = 2 * half_height * focus * v

    s = repeat((collect(0:(width - 1)) .+ 0.5f0) ./ width, outer = height)
    t = repeat((collect((height - 1):-1:0) .+ 0.5f0) ./ height, inner = width)

    direction = normalize(llc + s * hori + t * vert - origin)

    return (origin, direction)
end

"""
    get_transformation_matrix(c::Camera)

Returns the `camera_to_world` transformation matrix. This is used for transforming
the coordinates of a Point in 3D Camera Space to 3D World Space using the
[`camera2world`](@ref) function.
"""
function get_transformation_matrix(c::Camera{T}) where {T}
    forward = normalize(c.lookfrom - c.lookat)
    right = normalize(cross(c.fixedparams.vup, forward))
    up = normalize(cross(forward, right))

    return [     right.x[]      right.y[]      right.z[] zero(eltype(T));
                    up.x[]         up.y[]         up.z[] zero(eltype(T));
               forward.x[]    forward.y[]    forward.z[] zero(eltype(T));
            c.lookfrom.x[] c.lookfrom.y[] c.lookfrom.z[]  one(eltype(T))]
end

"""
    compute_screen_coordinates(c::Camera, film_aperture::Tuple,
                               inch_to_mm::Real = 25.4)

Computes the coordinates of the 4 corners of the screen and returns
`top`, `right`, `bottom`, and `left`.
"""
function compute_screen_coordinates(c::Camera, film_aperture::Tuple,
                                    inch_to_mm::Real = 25.4)
    width = c.fixedparams.width
    height = c.fixedparams.height
    vfov = c.vfov[]
    focus = c.focus[]

    film_aspect_ratio = film_aperture[1] / film_aperture[2]
    device_aspect_ratio = width / height

    top = ((film_aperture[2] * inch_to_mm / 2) / focus)
    right = ((film_aperture[1] * inch_to_mm / 2) / focus)

    xscale = 1
    yscale = 1

    if film_aspect_ratio > device_aspect_ratio
        xscale = device_aspect_ratio / film_aspect_ratio
    else
        yscale = film_aspect_ratio / device_aspect_ratio
    end

    right *= xscale
    top *= yscale

    bottom = -top
    left = -right

    return top, right, bottom, left
end


using DrWatson
@quickactivate "PECUZAL_Julia"

using DelayEmbeddings
using DynamicalSystemsBase
using Test
using DelimitedFiles
using Random
using PyPlot
pygui(true)

## We here focus on the dependence of the L-statistic on its parameter Tw and
# exemplary look at the Lorenz63 and the Lorenz96 system, as well as on random
# uniformly distributed noise.

## Generate Data:
# Lorenz63

# The time series have been computed in the script `/test/timeseries/produce_timeseries.jl`
s = readdlm("./test/timeseries/lorenz_pecora_uni_x.csv")
s1 = vec(s[1:5000]) # input timeseries = x component of lorenz
w1 = estimate_delay(vec(s), "mi_min")    # 1st minimum of mutual inf

# Lorenz96
idx = 487
N = 8 # number of oscillators
Fs = 3.5:0.002:5 # parameter spectrum
F = Fs[idx]
dt = 0.1 # sampling time
total = 5000  # time series length

# Parameters analysis:
ε = 0.05  # recurrence threshold
dmax = 10   # maximum dimension for traditional tde
lmin = 2   # minimum line length for RQA
trials = 80 # trials for MCDTS
taus = 0:100 # possible delays

#pick one time series
t_idx = 2

# init Lorenz96
u0 = [0.590; 0.766; 0.566; 0.460; 0.794; 0.854; 0.200; 0.298]
lo96 = Systems.lorenz96(N, u0; F = F)

#set_parameter!(lo96, 1, F)
data = trajectory(lo96, total*dt; dt = dt, Ttr = 2500 * dt)
s2 = data[1:5000,t_idx]
w2 = estimate_delay(vec(s2), "mi_min")    # 1st minimum of mutual inf

# Random numbers
Random.seed!(1234)
s3 = rand(1000)
w3 = estimate_delay(vec(s3), "mi_min")    # 1st minimum of mutual inf

# normalize data
D1 = regularize(Dataset(s1))
D2 = regularize(Dataset(s2))
D3 = regularize(Dataset(s3))

# generate embedding candidate
Y_t1 = DelayEmbeddings.hcat_lagged_values(D1, vec(Matrix(D1)), w1)
Y_t2 = DelayEmbeddings.hcat_lagged_values(D2, vec(Matrix(D2)), w2)
Y_t3 = DelayEmbeddings.hcat_lagged_values(D3, vec(Matrix(D3)), w3)

Tw = 100    # maximum Tw
L11 = zeros(Tw)
L12 = zeros(Tw)
L21 = zeros(Tw)
L22 = zeros(Tw)
L31 = zeros(Tw)
L32 = zeros(Tw)
for T = 1:Tw
    L11[T] = uzal_cost(D1; w = w1, samplesize = 1, Tw = T)
    L12[T] = uzal_cost(Y_t1; w = w1, samplesize = 1, Tw = T)
    L21[T] = uzal_cost(D2; w = w2, samplesize = 1, Tw = T)
    L22[T] = uzal_cost(Y_t2; w = w2, samplesize = 1, Tw = T)
    L31[T] = uzal_cost(D3; w = w3, samplesize = 1, Tw = T)
    L32[T] = uzal_cost(Y_t3; w = w3, samplesize = 1, Tw = T)
end

# Plotting

panelnames = ["A" "B" "C"]

fig = figure(figsize=[15,5])

lwg = 1.5         # linewidth of the graph
lwa = 2         # linewidth of the axis
fsp = 18        # Fontsize of panelnames
fsa = 12        # Fontsize of the axis
fsl = 10        # Fontsize of the legendentries
fst = 20        # Fontsize of title
axislabelsize = 12 # axislabelsize
ticklabelsize = 10  # labelsize of ticks
ms = 20         # markersize of maxima
ms_style = "*"  # markerstyle

subplots_adjust(
#left = 0.07,
#right = 0.96,
#bottom = 0.05,
#top = 0.96,
wspace = 0.35,
#hspace = 0.55)
)

rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
font0 = Dict(
        "font.size" => fsa,
        "font.weight" => "bold",
        "axes.labelweight" => "bold",
        "axes.titleweight" => "bold",
        "axes.labelsize" => axislabelsize,
        "axes.linewidth" => lwa,
        "xtick.labelsize" => ticklabelsize,
        "ytick.labelsize" => ticklabelsize,
        "legend.fontsize" => fsl)
        #"text.usetex" => true)
merge!(rcParams, font0)

ax = fig.add_subplot(1, 3, 1)
p1 = plot(1:Tw, L11, label="time series", linewidth=lwg)
color1 = p1[1].get_color()
scatter(1:Tw, L11, ms, label="")
p2 = plot(1:Tw, L12, label="2-dim embedding", linewidth=lwg)
color2 = p2[1].get_color()
scatter(1:Tw, L12, ms, label="")
ylims = ax.get_ylim()
ax.fill_between(1:Tw, L11, ylims[1], where = L11 .>= L12, facecolor = "gray", alpha = 0.5)
legend()
ax.text(-0.2, 1.1, panelnames[1], transform=ax.transAxes,
     fontsize=fsp, fontweight="bold", va="top")
title("Lorenz 63")
xlabel(L"T_M")
ylabel("L")
grid()

ax = fig.add_subplot(1, 3, 2)
p1 = plot(1:Tw, L21, label="time series", linewidth=lwg)
color1 = p1[1].get_color()
scatter(1:Tw, L21, ms, label="")
p2 = plot(1:Tw, L22, label="2-dim embedding", linewidth=lwg)
color2 = p2[1].get_color()
scatter(1:Tw, L22, ms, label="")
ylims = ax.get_ylim()
ax.fill_between(1:Tw, L21, ylims[1], where = L21 .>= L22, facecolor = "gray", alpha = 0.5)
#legend()
ax.text(-0.2, 1.1, panelnames[2], transform=ax.transAxes,
     fontsize=fsp, fontweight="bold", va="top")
title("Lorenz 96")
xlabel(L"T_M")
grid()

ax = fig.add_subplot(1, 3, 3)
p1 = plot(1:Tw, L31, label="time series", linewidth=lwg)
color1 = p1[1].get_color()
scatter(1:Tw, L31, ms, label="")
p2 = plot(1:Tw, L32, label="2-dim embedding", linewidth=lwg)
color2 = p2[1].get_color()
scatter(1:Tw, L32, ms, label="")
ylims = ax.get_ylim()
ax.fill_between(1:Tw, L31, ylims[1], where=L31 .>= L32, facecolor = "gray", alpha = 0.5)
#legend()
ax.text(-0.2, 1.1, panelnames[3], transform=ax.transAxes,
     fontsize=fsp, fontweight="bold", va="top")
title("Random numbers")
xlabel(L"T_M")
grid()


using Test
using SafeTestsets

@safetestset "Quaternion Tests" begin
    include("quaternion_test.jl")
end

include("diff/diff_test.jl")

@safetestset "Factorization Test" begin
    include("factorization_test.jl")
end

include("initialization/initialization_test.jl")

@safetestset "UI Test" begin
    include("ui_test.jl")
end

@safetestset "URDF Tests" begin
    include("urdf_test.jl")
end

@safetestset "Optionals Tests" begin
    include("optionals_test.jl")
end

@safetestset "Dynamics Tests" begin
    include("dynamics/dynamics_test.jl")
end

@safetestset "Example Tests" begin
    include("example_test.jl")
end


using jInv.Mesh
using jInv.Utils
using jInv.InverseSolve

using EikonalInv
using MAT

#############################################################################################################
plotting = false;
if plotting
	using jInvVis
	using  PyPlot
	close("all");
end
#############################################################################################################
dataDir = pwd();
resultsDir = pwd();

###############################################################################################################
# this filename can be used to save the model during the iterations. If set to "", no images or files will be saved.
# resultsFilenamePrefix = "";
resultsFilenamePrefix = string(resultsDir,"/shortTravelTimeInvSEG");
################################################################################################################

modelDir = pwd();
include("../drivers/prepareTravelTimeDataFiles.jl");
include("../drivers/readModelAndGenerateMeshMref.jl");
include("../drivers/setupTravelTimeTomography.jl");

########################################################################################################
######################################## for SEG 256X512 ###############################################
#######################################################################################################
dim     = 2;
pad     = 0;
jump    = 5;
newSize = [256,128];
offset  = ceil(Int64,(newSize[1]*(6/13.5)));;
(m,Minv,mref,boundsHigh,boundsLow) = readModelAndGenerateMeshMref(modelDir,"SEGmodel2Dsalt.dat",dim,pad,[0.0,13.5,0.0,4.2],newSize,1.752,2.9);
dataFilenamePrefix = string(dataDir,"/shortDATA_SEG",tuple((Minv.n+1)...));

limits = (minimum(m),maximum(m));

if plotting
	figure()
	plotModel(m,true,Minv,pad,limits);
	figure()
	plotModel(mref,true,Minv,pad,limits);
end

prepareTravelTimeDataFiles(m,Minv,mref,boundsHigh,boundsLow,dataFilenamePrefix,pad,jump,offset);

(Q,P,pMisRFs,SourcesSubInd,contDiv,Iact,sback,mref,boundsHigh,boundsLow,resultsFilename,Dobs) = setupTravelTimeTomography(m,dataFilenamePrefix, resultsFilenamePrefix);

########################################################################################################
##### Set up Inversion #################################################################################
########################################################################################################
############### Inversion for the velocity:
# modfun = velocityToSlowSquared;
################################################################################################################
############### USE A BOUND MODEL For velocity #################################################################
################################################################################################################
# maxStep=0.1*maximum(boundsHigh);
# a = minimum(boundsLow);
# b = maximum(boundsHigh);
# function modfun(m)
	# bm,dbm = getBoundModel(m,a,b);
	# bs,dbs = velocityToSlowSquared(bm);
	# dsdm = dbm*dbs;
	# return bs,dsdm;
# end
# mref = getBoundModelInv(mref,a,b);
# boundsHigh = boundsHigh*10000000.0;
# boundsLow = -boundsLow*100000000.0;
################################################################################################################
################ Inversion for the squared slowness: ###########################################################
################################################################################################################

mref 		= velocityToSlowSquared(mref)[1];
t    		= copy(boundsLow*0.9);
boundsLow 	= velocityToSlowSquared(boundsHigh*1.1)[1];
boundsHigh 	= velocityToSlowSquared(t)[1]; t = 0;
###################################################
a = minimum(boundsLow);
b = maximum(boundsHigh);
maxStep=0.1*maximum(boundsHigh);

############### Standard bounded optimization:
modfun 		= identityMod;

############### USE A BOUND MODEL For slow squared #############################################################

# modfun(m) = getBoundModel(m,a,b);
# mref = getBoundModelInv(mref,a,b);
# boundsHigh = boundsHigh*10000000.0;
# boundsLow = -boundsLow*100000000.0

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

cgit  = 8;
maxit = 10;
alpha = 5e-1;
pcgTol = 1e-1;

HesPrec=getSSORCGRegularizationPreconditioner(1.0,1e-5,100)

################################################# GIT VERSION OF JINV #################################################

regparams = [1.0,1.0,1.0,1e-6];
#### Use smoothness regularization
regfun(m, mref, M) = wdiffusionRegNodal(m, mref, M, Iact=Iact, C = regparams);
#### Use TV regularization
# regfun(m, mref, M) = wTVRegNodal(m, mref, M, Iact=Iact, C = regparams);

function dump(mc,Dc,iter,pInv,pMis)
	if resultsFilename!=""
		ntup = tuple((pInv.MInv.n+1)...);
		fullMc = slowSquaredToVelocity(reshape(Iact*pInv.modelfun(mc)[1] + sback,ntup))[1];
		Temp = splitext(resultsFilename);
		Temp = string(Temp[1],"_GN",iter,Temp[2]);
		writedlm(Temp,convert(Array{Float16},fullMc));
		if plotting
			close(888);
			figure(888);
			# plotModel(fullMc,false,[],0,limits,splitdir(Temp)[2]); # no information on axes
			plotModel(fullMc,true,Minv,pad,limits,splitdir(Temp)[2]); # with domain information on axes
		end
	end
end

pInv = getInverseParam(Minv,modfun,regfun,alpha,mref[:],boundsLow,boundsHigh,
                     maxStep=maxStep,pcgMaxIter=cgit,pcgTol=pcgTol,
					 minUpdate=1e-3, maxIter = maxit,HesPrec=HesPrec);

D0RFs, = computeMisfit(pInv.modelfun(mref[:])[1],pMisRFs,false);

#### Projected Gauss Newton
mc,Dc,flag,his = projGNCG(copy(mref[:]),pInv,pMisRFs,indCredit = [],dumpResults = dump);
#### Barrier Gauss Newton
# mc,Dc,flag,his = barrierGNCG(copy(mref[:]),pInv,pMisRFs,indCredit = [],dumpResults = dump);

if plotting
	figure()
	Fvals = his.F;
	Jvals = his.Jc;
	numiter = length(Fvals);
	semilogy(0:numiter-1,Fvals,"-*");
	xlabel("Gauss Newton iterations")
	ylabel("Misfit value")
end

D0 = Array(Array{Float64,2},length(pMisRFs))
for k = 1:length(pMisRFs)
	D0[k] = fetch(D0RFs[k]);
end
D0 = arrangeRemoteCallDataIntoLocalData(D0); D0RFs = 0;

Dpred = Array(Array{Float64,2},length(pMisRFs))
for k = 1:length(pMisRFs)
	Dpred[k] = fetch(Dc[k]);
end
Dpred = arrangeRemoteCallDataIntoLocalData(Dpred); Dc = 0;

if plotting
	figure()
	imshow(Dpred');
	colorbar();
	xlabel("Sources")
	ylabel("Receivers")
	figure()
	imshow(D0');
	colorbar();
	xlabel("Sources")
	ylabel("Receivers")
	figure()
	imshow(abs(Dpred - Dobs)');
	colorbar();
	xlabel("Sources")
	ylabel("Receivers")
	figure()
	imshow(abs(D0 - Dobs)');
	colorbar();
	xlabel("Sources")
	ylabel("Receivers")
end

if resultsFilename!=""
	Temp = splitext(resultsFilename);
	writedlm(string(Temp[1],"_predictedData",Temp[2]),convert(Array{Float16},Dpred));
	writedlm(string(Temp[1],"_recoveredModel",Temp[2]),convert(Array{Float16},reshape(Iact*modfun(mc)[1] + sback,tuple((pInv.MInv.n+1)...))));
end

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


# This function calculates the largest eigenvalue
# (absolute value wise) by power iteration.

function maxeig!(integrator, cache::OrdinaryDiffEqConstantCache)
  isfirst = integrator.iter == 1 || integrator.u_modified
  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
  maxiter = 50
  safe = 1.2
  # Initial guess for eigenvector `z`
  if isfirst
    fz = fsalfirst
    z = f(fz, p, t)
  else
    z = cache.zprev
  end
  # Perturbation
  u_norm = integrator.opts.internalnorm(uprev)
  z_norm = integrator.opts.internalnorm(z)
  pert   = eps(u_norm)
  sqrt_pert = sqrt(pert)
  is_u_zero = u_norm == zero(u_norm)
  is_z_zero = z_norm == zero(z_norm)
  # Normalize `z` such that z-u lie in a circle
  if ( !is_u_zero && !is_z_zero )
    dz_u = u_norm * sqrt_pert
    quot = dz_u/z_norm
    z = uprev + quot*z
  elseif !is_u_zero
    dz_u = u_norm * sqrt_pert
    z = uprev + uprev*dz_u
  elseif !is_z_zero
    dz_u = pert
    quot = dz_u/z_norm
    z *= quot
  else
    dz_u = pert
    z = dz_u
  end # endif
  # Start power iteration
  integrator.eigen_est = 0
  for iter in 1:maxiter
    fz = f(z, p, t)
    tmp = fz - fsalfirst
    Δ  = integrator.opts.internalnorm(tmp)
    eig_prev = integrator.eigen_est
    integrator.eigen_est = Δ/dz_u * safe
    # Convergence
    if iter >= 2 && abs(eig_prev - integrator.eigen_est) < integrator.eigen_est*0.05
      # Store the eigenvector
      cache.zprev = z
      return true
    end
    # Next `z`
    if Δ != zero(Δ)
      quot = dz_u/Δ
      z = uprev + quot*tmp
    else
      # An arbitrary change on `z`
      cache.z = shuffle(cache.z)
    end
  end
  return false
end

function maxeig!(integrator, cache::OrdinaryDiffEqMutableCache)
  isfirst = integrator.iter == 1 || integrator.u_modified
  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
  fz, z, atmp = cache.k, cache.tmp, cache.atmp
  ccache = cache.constantcache
  maxiter = 50
  safe = 1.2
  # Initial guess for eigenvector `z`
  if isfirst
    @. fz = u
    f(z, fz, p, t)
  else
    @. z = ccache.zprev
  end
  # Perturbation
  u_norm = integrator.opts.internalnorm(uprev)
  z_norm = integrator.opts.internalnorm(z)
  pert   = eps(u_norm)
  sqrt_pert = sqrt(pert)
  is_u_zero = u_norm == zero(u_norm)
  is_z_zero = z_norm == zero(z_norm)
  # Normalize `z` such that z-u lie in a circle
  if ( !is_u_zero && !is_z_zero )
    dz_u = u_norm * sqrt_pert
    quot = dz_u/z_norm
    @. z = uprev + quot*z
  elseif !is_u_zero
    dz_u = u_norm * sqrt_pert
    @. z = uprev + uprev*dz_u
  elseif !is_z_zero
    dz_u = pert
    quot = dz_u/z_norm
    @. z *= quot
  else
    dz_u = pert
    @. z = dz_u
  end # endif
  # Start power iteration
  integrator.eigen_est = 0
  for iter in 1:maxiter
    f(fz, z, p, t)
    @. atmp = fz - fsalfirst
    Δ  = integrator.opts.internalnorm(atmp)
    eig_prev = integrator.eigen_est
    integrator.eigen_est = Δ/dz_u * safe
    # Convergence
    if iter >= 2 && abs(eig_prev - integrator.eigen_est) < integrator.eigen_est*0.05
      # Store the eigenvector
      @. ccache.zprev = z
      return true
    end
    # Next `z`
    if Δ != zero(Δ)
      quot = dz_u/Δ
      @. z = uprev + quot*atmp
    else
      # An arbitrary change on `z`
      nind = length(uprev)
      ind = 1 + iter % nind
      z[ind] = uprev[ind] - (z[ind] - uprev[ind])
    end
  end
  return false
end

"""
    choosedeg!(cache) -> nothing

Calculate `ms[mdeg]` (the degree of the Chebyshev polynomial)
and `cache.recind` (the index of recurrence parameters for that
degree), where `recf[recind:(recind+ms[mdeg]-2)]` are the `μ,κ` pairs
for the `mdeg` degree method.
"""
function choosedeg!(cache::T) where T
  isconst = T <: OrdinaryDiffEqConstantCache
  isconst || ( cache = cache.constantcache )
  @unpack ms, fp1, fp2, recf, zprev = cache
  recind = 0
  @inbounds for i in 1:46
    recind += ms[i]
    if ms[i] > cache.mdeg
      cache.mdeg = i
      cache.recind = recind
      break
    end # end if
  end # end for
  return nothing
end


using Zygote, NNlib, Test, Random, LinearAlgebra, Statistics, FillArrays, FFTW
using Zygote: gradient
using NNlib: conv, ∇conv_data, depthwiseconv
using Base.Broadcast: broadcast_shape

function ngradient(f, xs::AbstractArray...)
  grads = zero.(xs)
  for (x, Δ) in zip(xs, grads), i in 1:length(x)
    δ = sqrt(eps())
    tmp = x[i]
    x[i] = tmp - δ/2
    y1 = f(xs...)
    x[i] = tmp + δ/2
    y2 = f(xs...)
    x[i] = tmp
    Δ[i] = (y2-y1)/δ
  end
  return grads
end

gradcheck(f, xs...) =
  all(isapprox.(ngradient(f, xs...),
                gradient(f, xs...), rtol = 1e-5, atol = 1e-5))

gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)
gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)

Random.seed!(0)

@test gradient(//, 2, 3) === (1//3, -2//9)

@test gradtest((x, W, b) -> σ.(W*x .+ b), 5, (2,5), 2)
@test gradtest((x, W, b) -> σ.(W*x .+ b), (5,3), (2,5), 2)
@test gradtest((x, W, b) -> logσ.(W*x .+ b), 5, (2,5), 2)
@test gradtest((x, W, b) -> logσ.(W*x .+ b), (5,3), (2,5), 2)

@test gradtest((w, x) -> w'*x, randn(10, 2), randn(10))
@test gradtest((w, x) -> transpose(w)*x, randn(5,5), randn(5,5))

@test gradtest(x -> sum(x, dims = (2, 3)), (3,4,5))
@test gradtest(x -> sum(abs2, x), randn(4, 3, 2))
@test gradtest(x -> sum(abs2, x; dims=1), randn(4, 3, 2))
@test gradtest(x -> prod(x, dims = (2, 3)), (3,4,5))
@test gradtest(x -> prod(x), (3,4,5))

@test gradtest(x -> softmax(x).*(1:3), 3)
@test gradtest(x -> softmax(x).*(1:3), (3,5))
@test gradtest(x -> logsoftmax(x).*(1:3), 3)
@test gradtest(x -> logsoftmax(x).*(1:3), (3,5))

@test gradtest(x -> x', rand(5))

@test gradtest(det, (4, 4))
@test gradtest(logdet, map(x -> x*x', (rand(4, 4),))[1])
@test gradtest(x -> logabsdet(x)[1], (4, 4))

@test gradtest(x -> view(x,:,2,:), (3,4,5))
@test gradtest(x -> view(x,1:2,3:4), (3,4))

@testset "conv" begin
  for spatial_rank in (1, 2, 3)
    x = rand(repeat([10], spatial_rank)..., 3, 2)
    w = rand(repeat([3], spatial_rank)..., 3, 3)
    cdims = DenseConvDims(x, w)
    @test gradtest((x, w) -> conv(x, w, cdims), x, w)
    y = conv(x, w, cdims)
    @test gradtest((y, w) -> ∇conv_data(y, w, cdims), y, w)
    dcdims = DepthwiseConvDims(x, w)
    @test gradtest((x, w) -> depthwiseconv(x, w, dcdims), x, w)
  end
end

@testset "pooling" begin
  for spatial_rank in (1, 2)
    x = rand(repeat([10], spatial_rank)..., 3, 2)
    pdims = PoolDims(x, 2)
    @test gradtest(x -> maxpool(x, pdims), x)
    @test gradtest(x -> meanpool(x, pdims), x)
  end
end

@test gradtest(x -> permutedims(x), rand(2))
@test gradtest(x -> permutedims(x), rand(2,3))
@test gradtest(x -> permutedims(x, [3,1,2]), rand(4,5,6))
@test gradtest(x -> PermutedDimsArray(x, (3,1,2)), rand(4,5,6))
let
  y, back = Zygote.pullback(permutedims, randn(3))
  @test first(back(randn(1, 3))) isa Vector
end

@test gradtest(x -> repeat(x; inner=2), rand(5))
@test gradtest(x -> repeat(x; inner=2, outer=3), rand(5))
@test gradtest(x -> repeat(x; inner=(2,2,1), outer=(1,1,3)), rand(5,4,3))

@test gradtest(tr, rand(4, 4))

@testset "fill" begin
  rng, N, M, P = MersenneTwister(123456), 11, 6, 5
  @test gradtest(x->fill(first(x), N), randn(rng, 1))
  @test gradtest(x->fill(first(x), N, M), randn(rng, 1))
  @test gradtest(x->fill(first(x), N, M, P), randn(rng, 1))
end

@testset "dot" begin
  rng = MersenneTwister(123456)
  @test gradtest((x, y)->dot(x[1], y[1]), [randn(rng)], [randn(rng)])
  @test gradtest(dot, randn(rng, 10), randn(rng, 10))
  @test gradtest(dot, randn(rng, 10, 3), randn(rng, 10, 3))
end

@test gradtest(kron, rand(5), rand(3))
@test gradtest(kron, rand(5), rand(3), rand(8))
@test gradtest(kron, rand(5,1), rand(3,1))
@test gradtest(kron, rand(5,1), rand(3,1), rand(8,1))
@test gradtest(kron, rand(5,2), rand(3,2), rand(8,2))

@testset "map" begin
  @test gradtest(xs -> sum(map(x -> x^2, xs)), rand(2,3))
  @test gradtest((xss...) -> sum(map((xs...) -> sqrt(sum(xs.^2)), xss...)), [rand(5) for _ in 1:6]...)
  function foo(y)
    bar = (x) -> x*y
    sum(map(bar, 1:5))
  end
  @test gradtest(foo, 3)
  @test gradient(v -> sum([x for x in v]), [1.1,2.2,3.3]) == ([1, 1, 1],)
end

@testset "mean" begin
  @test gradtest(mean, rand(2, 3))

  @test gradtest(x -> mean(x, dims=1), rand(2, 3))
  @test gradtest(x -> mean(x, dims=2), rand(2, 3))
  @test gradtest(x -> mean(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> mean(x, dims=[1, 2]), rand(2, 3, 4))
end

@testset "var" begin
  @test gradtest(var, rand(2, 3))
  @test gradtest(x -> var(x, dims=1), rand(2, 3))
  @test gradtest(x -> var(x, dims=2), rand(2, 3))
  @test gradtest(x -> var(x, dims=3), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2]), rand(2, 3, 4))

  @test gradtest(x -> var(x, corrected=false), rand(2, 3))
  @test gradtest(x -> var(x, dims=1, corrected=false), rand(2, 3))
  @test gradtest(x -> var(x, dims=2, corrected=false), rand(2, 3))
  @test gradtest(x -> var(x, dims=3, corrected=false), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2], corrected=false), rand(2, 3, 4))

  @test gradtest(x -> var(x, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> var(x, dims=1, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> var(x, dims=2, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> var(x, dims=3, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2], mean=mean(x, dims=[1, 2])), rand(2, 3, 4))

  @test gradtest(x -> var(x, corrected=false, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> var(x, dims=1, corrected=false, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> var(x, dims=2, corrected=false, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> var(x, dims=3, corrected=false, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> var(x, dims=[1, 2], corrected=false, mean=mean(x, dims=[1, 2])), rand(2, 3, 4))
end

@testset "std" begin
  @test gradtest(std, rand(2, 3))
  @test gradtest(x -> std(x, dims=1), rand(2, 3))
  @test gradtest(x -> std(x, dims=2), rand(2, 3))
  @test gradtest(x -> std(x, dims=3), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2]), rand(2, 3, 4))

  @test gradtest(x -> std(x, corrected=false), rand(2, 3))
  @test gradtest(x -> std(x, dims=1, corrected=false), rand(2, 3))
  @test gradtest(x -> std(x, dims=2, corrected=false), rand(2, 3))
  @test gradtest(x -> std(x, dims=3, corrected=false), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2], corrected=false), rand(2, 3, 4))

  @test gradtest(x -> std(x, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> std(x, dims=1, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> std(x, dims=2, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> std(x, dims=3, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2], mean=mean(x, dims=[1, 2])), rand(2, 3, 4))

  @test gradtest(x -> std(x, corrected=false, mean=mean(x)), rand(2, 3))
  @test gradtest(x -> std(x, dims=1, corrected=false, mean=mean(x, dims=1)), rand(2, 3))
  @test gradtest(x -> std(x, dims=2, corrected=false, mean=mean(x, dims=2)), rand(2, 3))
  @test gradtest(x -> std(x, dims=3, corrected=false, mean=mean(x, dims=3)), rand(2, 3, 4))
  @test gradtest(x -> std(x, dims=[1, 2], corrected=false, mean=mean(x, dims=[1, 2])), rand(2, 3, 4))
end

@testset "maximum" begin
  @test gradtest(maximum, rand(2, 3))

  @test gradtest(x -> maximum(x, dims=1), rand(2, 3))
  @test gradtest(x -> maximum(x, dims=2), rand(2, 3))
  @test gradtest(x -> maximum(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> maximum(x, dims=[1, 2]), rand(2, 3, 4))

  @test gradient(x -> 1 / maximum(x), [1., 2, 3])[1] == [0, 0, -1/9]
end

@testset "minimum" begin
  @test gradtest(minimum, rand(2, 3))

  @test gradtest(x -> minimum(x, dims=1), rand(2, 3))
  @test gradtest(x -> minimum(x, dims=2), rand(2, 3))
  @test gradtest(x -> minimum(x, dims=3), rand(2, 3, 4))

  @test gradtest(x -> minimum(x, dims=[1, 2]), rand(2, 3, 4))
end

@testset "dropdims" begin
  @test gradtest(x -> dropdims(x, dims = 3), rand(2, 2, 1, 2))
  @test gradtest(x -> dropdims(x, dims = (2, 3)), rand(2, 1, 1, 3))
  @test gradtest(x -> dropdims(x, dims = (1, 2, 3)), rand(1, 1, 1, 3))
end

@testset "$f(::AbstractArray)" for f in (real, conj, imag)
  rng, N = MersenneTwister(123456), 3
  Ts = (Float64, ComplexF64)
  @testset "$f(::Array{$IT})" for IT in Ts
    A = randn(IT, N, N)
    y, back = Zygote.pullback(f, A)
    y2, back2 = Zygote.pullback(x->f.(x), A)
    @test y == y2
    @testset "back(::Array{$BT})" for BT in Ts
      ȳ = randn(BT, N, N)
      @test back(ȳ)[1] == back2(ȳ)[1]
    end
  end
end

@testset "(p)inv" begin
  rng, P, Q = MersenneTwister(123456), 13, 11
  A, B, C = randn(rng, P, Q), randn(rng, P, P), randn(Q, P)
  @test gradtest(pinv, A)
  @test gradtest(inv, B)
  @test gradtest(pinv, C)
end

@testset "multiplication" begin
  @testset "matrix-matrix" begin
    rng, M, P, Q = MersenneTwister(123456), 13, 7, 11
    @test gradtest(*, randn(rng, M, P), randn(rng, P, Q))
    @test gradtest(*, randn(rng, M, P), randn(rng, P))
    @test gradtest(*, randn(rng, M, 1), randn(rng, 1, Q))
    @test gradtest(*, randn(rng, M), randn(rng, 1, Q))
    @test gradtest(*, randn(rng, 10)', randn(rng, 10))

    let
      y, back = Zygote.pullback(*, randn(rng, M, P), randn(rng, P))
      @test last(back(randn(rng, M))) isa Vector
    end
    let
      y, back = Zygote.pullback(*, randn(rng, M), randn(rng, 1, P))
      @test first(back(randn(rng, M, P))) isa Vector
    end
  end
end

@testset "backsolve" begin
  rng, M, P, Q = MersenneTwister(123456), 13, 10, 9
  X, Y, y = randn(rng, P, P), randn(rng, P, Q), randn(rng, P)
  A, B = randn(rng, P, M), randn(P, Q)
  D = collect(Diagonal(randn(rng, P)))
  L = collect(LowerTriangular(randn(rng, P, P)))
  L[diagind(L)] .= 1 .+ 0.01 .* randn(rng, P)
  U = collect(UpperTriangular(randn(rng, P, P)))
  U[diagind(U)] .= 1 .+ 0.01 .* randn(rng, P)

  # \ (Dense square)
  @test gradtest(\, X, Y)
  @test gradtest(\, X, y)

  # \ (Dense rectangular)
  @test gradtest(\, A, Y)
  @test gradtest(\, A, y)
  @test gradtest(\, B, Y)
  @test gradtest(\, B, y)

  # \ (Diagonal)
  @test gradtest(\, D, Y)
  @test gradtest(\, D, y)
  @test gradtest((D, Y)-> Diagonal(D) \ Y, D, Y)
  @test gradtest((D, Y)-> Diagonal(D) \ Y, D, y)

  # \ (LowerTriangular)
  @test gradtest(\, L, Y)
  @test gradtest(\, L, y)
  @test gradtest((L, Y) -> LowerTriangular(L) \ Y, L, Y)
  @test gradtest((L, Y) -> LowerTriangular(L) \ Y, L, y)

  # \ (UpperTriangular)
  @test gradtest(\, U, Y)
  @test gradtest(\, U, y)
  @test gradtest((U, Y) -> UpperTriangular(U) \ Y, U, Y)
  @test gradtest((U, Y) -> UpperTriangular(U) \ Y, U, y)

  # /
  @test gradtest(/, Y', X)
  @test gradtest((y, X)->y' / X, y, X)

  # / (rectangular)
  @test gradtest(/, Y', A')
  @test gradtest((y, A)->y' / A', y, A)
  @test gradtest(/, Y', B')
  @test gradtest((y, A)->y' / A', y, B)

  # / (Diagonal)
  @test gradtest((D, Y) -> Y' / D, D, Y)
  @test gradtest((D, Y) -> Y' / D, D, y)
  @test gradtest((D, Y)-> Y' / Diagonal(D), D, Y)
  @test gradtest((D, Y)-> Y' / Diagonal(D), D, y)

  # / (LowerTriangular)
  @test gradtest((L, Y) -> Y' / L, L, Y)
  @test gradtest((L, Y) -> Y' / L, L, y)
  @test gradtest((L, Y) -> Y' / LowerTriangular(L), L, Y)
  @test gradtest((L, Y) -> Y' / LowerTriangular(L), L, y)

  # / (UpperTriangular)
  @test gradtest((U, Y) -> Y' / U, U, Y)
  @test gradtest((U, Y) -> Y' / U, U, y)
  @test gradtest((U, Y) -> Y' / UpperTriangular(U), U, Y)
  @test gradtest((U, Y) -> Y' / UpperTriangular(U), U, y)

  @testset "Cholesky" begin

    # Check that the forwards pass computes the correct thing.
    @test Zygote.pullback(X->cholesky(X * X' + I) \ Y, X)[1] == cholesky(X * X' + I) \ Y
    @test gradtest(X->cholesky(X * X' + I) \ Y, X)
    @test gradtest(Y->cholesky(X * X' + I) \ Y, Y)
    @test gradtest(X->cholesky(X * X' + I) \ y, X)
    @test gradtest(y->cholesky(X * X' + I) \ y, y)
  end
end

@testset "Symmetric" begin
  @testset "real" begin
    rng, P = MersenneTwister(123456), 7
    A = randn(rng, P, P)
    @testset "uplo=$uplo" for uplo in (:U, :L)
      @test gradtest(x->Symmetric(x, uplo), A)
      y, back = Zygote.pullback(Symmetric, A, uplo)
      @test y isa Symmetric

      @testset "back(::Diagonal)" begin
        D̄ = Diagonal(randn(rng, P))
        @test back(Diagonal(D̄))[1] isa Diagonal
        @test back(Diagonal(D̄))[1] ≈ back(Matrix(D̄))[1]
      end

      @testset "back(::$TTri)" for TTri in (LowerTriangular,UpperTriangular)
        D̄ = TTri(randn(rng, P, P))
        @test back(D̄)[1] isa Matrix
        @test back(D̄)[2] === nothing
        @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      end
    end
  end

  @testset "complex" begin
    rng, P = MersenneTwister(123456), 7
    Re = randn(rng, P, P)
    Im = randn(rng, P, P)
    A = complex.(Re, Im)

    @testset "gradcheck dense" begin
      for uplo in (:U, :L)
        @test gradcheck(Re,Im) do a, b
          c = Symmetric(complex.(a, b), uplo)
          d = exp.(c)
          sum(real.(d) + imag.(d))
        end
      end
    end

    @testset "uplo=$uplo" for uplo in (:U, :L)
      y, back = Zygote.pullback(Symmetric, A, uplo)
      @test y isa Symmetric

      @testset "back(::Diagonal)" begin
        D̄ = Diagonal(complex.(randn(rng, P), randn(rng, P)))
        @test back(Diagonal(D̄))[1] isa Diagonal
        @test back(Diagonal(D̄))[1] ≈ back(Matrix(D̄))[1]
      end

      @testset "back(::$TTri)" for TTri in (LowerTriangular,UpperTriangular)
        D̄ = TTri(complex.(randn(rng, P, P), randn(rng, P, P)))
        @test back(D̄)[1] isa Matrix
        @test back(D̄)[2] === nothing
        @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      end
    end
  end
end

@testset "Hermitian" begin
  rng, P = MersenneTwister(123456), 7
  Re = randn(rng, P, P)
  Im = randn(rng, P, P)
  A = complex.(Re, Im)

  @testset "gradcheck dense" begin
    for uplo in (:U, :L)
      @test gradcheck(Re,Im) do a, b
        c = Hermitian(complex.(a, b), uplo)
        d = exp.(c)
        sum(real.(d) + imag.(d))
      end
    end
  end

  @testset "uplo=$uplo" for uplo in (:U, :L)
    y, back = Zygote.pullback(Hermitian, A, uplo)
    _, back_sym = Zygote.pullback(Symmetric, A, uplo)
    @test y isa Hermitian

    @testset "back" begin
      D̄ = randn(rng, P, P)
      @test back(D̄)[1] ≈ back_sym(D̄)[1]
    end

    @testset "back(::Diagonal)" begin
      D̄ = Diagonal(complex.(randn(rng, P), randn(rng, P)))
      @test back(D̄)[1] isa Diagonal
      @test back(D̄)[2] === nothing
      @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      @test back(real(D̄))[1] ≈ back_sym(real(D̄))[1]
    end

    @testset "back(::$TTri)" for TTri in (LowerTriangular,UpperTriangular)
      D̄ = TTri(complex.(randn(rng, P, P), randn(rng, P, P)))
      @test back(D̄)[1] isa Matrix
      @test back(D̄)[2] === nothing
      @test back(D̄)[1] ≈ back(Matrix(D̄))[1]
      @test back(real(D̄))[1] ≈ back_sym(real(D̄))[1]
    end
  end
end

@testset "diag" begin
  rng, P = MersenneTwister(123456), 10
  A = randn(rng, P, P)
  @test gradtest(diag, A)
end

@testset "Diagonal" begin
  rng, P = MersenneTwister(123456), 10
  d = randn(rng, P)
  @test gradtest(Diagonal, d)
  y, back = Zygote.pullback(Diagonal, d)
  D̄ = randn(rng, P, P)
  @test back(D̄) == back(Diagonal(D̄))
  @test back(D̄) == back((diag=diag(D̄),))
end

@testset "dense + UniformScaling" begin
  rng = MersenneTwister(123456)
  A, λ = randn(rng, 10, 10), randn(rng)
  @test gradtest(A->A + 5I, A)
  @test gradtest(λ->A + λ[1] * I, [λ])
end

@testset "cholesky" begin
  @testset "cholesky - dense" begin
    rng, N = MersenneTwister(123456), 5
    A = randn(rng, N, N)
    @test cholesky(A' * A + I) == first(Zygote.pullback(A->cholesky(A' * A + I), A))
    @test gradtest(A->cholesky(A' * A + I).U, A)
    @test gradtest(A->logdet(cholesky(A' * A + I)), A)
    @test gradtest(B->cholesky(Symmetric(B)).U, A * A' + I)
    @test gradtest(B->logdet(cholesky(Symmetric(B))), A * A' + I)
  end
  @testset "cholesky - scalar" begin
    rng = MersenneTwister(123456)
    y, back = Zygote.pullback(cholesky, 5.0 * ones(1, 1))
    y′, back′ = Zygote.pullback(cholesky, 5.0)
    C̄ = randn(rng, 1, 1)
    @test back′((factors=C̄,))[1] isa Real
    @test back′((factors=C̄,))[1] ≈ back((factors=C̄,))[1][1, 1]
  end
  @testset "cholesky - Diagonal" begin
    rng, N = MersenneTwister(123456), 3
    D = Diagonal(exp.(randn(rng, N)))
    Dmat = Matrix(D)
    y, back = Zygote.pullback(cholesky, Dmat)
    y′, back′ = Zygote.pullback(cholesky, D)
    C̄ = (factors=randn(rng, N, N),)
    @test back′(C̄)[1] isa Diagonal
    @test diag(back′(C̄)[1]) ≈ diag(back(C̄)[1])
  end
end

@testset "lyap" begin
  rng, N = MersenneTwister(6865943), 5
  for i = 1:5
    A = randn(rng, N, N)
    C = randn(rng, N, N)
    @test gradtest(lyap, A, C)
  end
  @test gradcheck(x->lyap(x[1],x[2]),[3.1,4.6])
end

@testset "matrix exponential" begin
  @testset "real dense" begin
    rng, N = MersenneTwister(6865931), 8
    for i = 1:5
      A = randn(rng, N, N)
      @test gradtest(exp, A)
    end
  end

  @testset "complex dense" begin
    rng, N = MersenneTwister(6865931), 8
    for i = 1:5
      A = randn(rng, N, N)
      B = randn(rng, N, N)
      @test gradcheck(A, B) do a,b
        c = complex.(a, b)
        d = exp(c)
        return sum(real.(d) + 2 .* imag.(d))
      end
    end
  end
end

@testset "eigen(::RealHermSymComplexHerm)" begin
  @testset "eigen(::Symmetric{<:Real})" begin
    rng, N = MersenneTwister(123), 7
    A = Symmetric(randn(rng, N, N))
    @test gradtest(collect(A)) do (x)
      d, Q = eigen(Symmetric(x))
      return Q * Diagonal(exp.(d)) * transpose(Q)
    end
    y = Zygote.pullback(eigen, A)[1]
    y2 = eigen(A)
    @test y.values ≈ y2.values
    @test y.vectors ≈ y2.vectors
    @testset "low rank" begin
      U = eigvecs(A)
      A2 = Symmetric(U * Diagonal([randn(rng), zeros(N-1)...]) * U')
      @test_broken gradtest(collect(A2)) do (x)
        d, Q = eigen(Symmetric(x))
        return Q * Diagonal(exp.(d)) * transpose(Q)
      end
    end
  end

  @testset "eigen(::Hermitian{<:Real})" begin
    rng, N = MersenneTwister(456), 7
    A = Hermitian(randn(rng, N, N))
    @test gradtest(collect(A)) do (x)
      d, Q = eigen(Hermitian(x))
      return Q * Diagonal(exp.(d)) * transpose(Q)
    end
    y = Zygote.pullback(eigen, A)[1]
    y2 = eigen(A)
    @test y.values ≈ y2.values
    @test y.vectors ≈ y2.vectors
    @testset "low rank" begin
      U = eigvecs(A)
      A2 = Hermitian(U * Diagonal([randn(rng), zeros(N-1)...]) * U')
      @test_broken gradtest(collect(A2)) do (x)
        d, Q = eigen(Hermitian(x))
        return Q * Diagonal(exp.(d)) * transpose(Q)
      end
    end
  end

  @testset "eigen(::Hermitian{<:Complex})" begin
    rng, N = MersenneTwister(789), 7
    A = Hermitian(randn(rng, ComplexF64, N, N))
    @test gradtest(reim(collect(A))...) do a,b
      d, U = eigen(Hermitian(complex.(a, b)))
      X = U * Diagonal(exp.(d)) * U'
      return real.(X) .+ imag.(X)
    end
    y = Zygote.pullback(eigen, A)[1]
    y2 = eigen(A)
    @test y.values ≈ y2.values
    @test y.vectors ≈ y2.vectors
    @testset "low rank" begin
      U = eigvecs(A)
      A2 = Hermitian(U * Diagonal([randn(rng), zeros(N-1)...]) * U')
      @test_broken gradtest(reim(collect(A2))...) do a,b
        d, U = eigen(Hermitian(complex.(a, b)))
        X = U * Diagonal(exp.(d)) * U'
        return real.(X) .+ imag.(X)
      end
    end
  end
end

@testset "eigvals(::RealHermSymComplexHerm)" begin
  @testset "eigvals(::Symmetric{<:Real})" begin
    rng, N = MersenneTwister(123), 7
    A = Symmetric(randn(rng, N, N))
    @test gradtest(x->eigvals(Symmetric(x)), collect(A))
    @test Zygote.pullback(eigvals, A)[1] ≈ eigvals(A)
  end

  @testset "eigvals(::Hermitian{<:Real})" begin
    rng, N = MersenneTwister(456), 7
    A = Hermitian(randn(rng, N, N))
    @test gradtest(x->eigvals(Hermitian(x)), collect(A))
    @test Zygote.pullback(eigvals, A)[1] ≈ eigvals(A)
  end

  @testset "eigvals(::Hermitian{<:Complex})" begin
    rng, N = MersenneTwister(789), 7
    A, B = randn(rng, N, N), randn(rng, N, N)
    @test gradtest(A, B) do a,b
      c = Hermitian(complex.(a, b))
      return eigvals(c)
    end
    @test Zygote.pullback(eigvals, Hermitian(A))[1] ≈ eigvals(Hermitian(A))
  end
end

using Distances

Zygote.refresh()

@testset "distances" begin
  rng, P, Q, D = MersenneTwister(123456), 10, 9, 8

   # Check sqeuclidean.
  let
    x, y = randn(rng, D), randn(rng, D)
    @test gradtest(x->sqeuclidean(x, y), x)
    @test gradtest(y->sqeuclidean(x, y), y)
  end

   # Check binary colwise.
  let
    X, Y = randn(rng, D, P), randn(rng, D, P)
    @test gradtest(X->colwise(SqEuclidean(), X, Y), X)
    @test gradtest(Y->colwise(SqEuclidean(), X, Y), Y)
  end

   # Check binary pairwise.
  let
    X, Y = randn(rng, D, P), randn(rng, D, Q)
    @test gradtest(X->pairwise(SqEuclidean(), X, Y; dims=2), X)
    @test gradtest(Y->pairwise(SqEuclidean(), X, Y; dims=2), Y)
  end
  let
    Xt, Yt = randn(rng, P, D), randn(rng, Q, D)
    @test gradtest(Xt->pairwise(SqEuclidean(), Xt, Yt; dims=1), Xt)
    @test gradtest(Yt->pairwise(SqEuclidean(), Xt, Yt; dims=1), Yt)
  end

   # Check unary pairwise.
  @test gradtest(X->pairwise(SqEuclidean(), X; dims=2), randn(rng, D, P))
  @test gradtest(Xt->pairwise(SqEuclidean(), Xt; dims=1), randn(rng, P, D))

  @testset "colwise(::Euclidean, X, Y; dims=2)" begin
    rng, D, P = MersenneTwister(123456), 2, 3
    X, Y, D̄ = randn(rng, D, P), randn(rng, D, P), randn(rng, P)
    gradtest((X, Y)->colwise(Euclidean(), X, Y), X, Y)
  end
  @testset "pairwise(::Euclidean, X, Y; dims=2)" begin
    rng, D, P, Q = MersenneTwister(123456), 2, 3, 5
    X, Y, D̄ = randn(rng, D, P), randn(rng, D, Q), randn(rng, P, Q)
    gradtest((X, Y)->pairwise(Euclidean(), X, Y; dims=2), X, Y)
  end
  @testset "pairwise(::Euclidean, X; dims=2)" begin
    rng, D, P = MersenneTwister(123456), 2, 3
    X, D̄ = randn(rng, D, P), randn(rng, P, P)
    gradtest(X->pairwise(Euclidean(), X; dims=2), X)
  end
end

function cat_test(f, A::Union{AbstractVector, AbstractMatrix}...)
  @test gradtest(f, A...)
  Z, back = Zygote.pullback(f, A...)
  Ā = back(randn(size(Z)))
  @test all(map((a, ā)->ā isa typeof(a), A, Ā))
end

@testset "vcat" begin
  # Scalar
  @test gradient((x,y) -> sum(vcat(x,y)), 1,2) == (1,1)
  @test gradient((x,y) -> sum([x;y]), 1,2) == (1,1)

  # Vector-only.
  cat_test(vcat, randn(1))
  cat_test(vcat, randn(10))

  cat_test(vcat, randn(1), randn(1))
  cat_test(vcat, randn(5), randn(1))
  cat_test(vcat, randn(1), randn(6))
  cat_test(vcat, randn(5), randn(6))

  # Matrix-only.
  for c in [1, 2], r1 in [1, 2], r2 in [1, 2]
    cat_test(vcat, randn(r1, c))
    cat_test(vcat, randn(r1, c), randn(r2, c))
  end

  # Matrix-Vector / Vector-Matrix / Vector-Matrix-Vector.
  for r in [1, 2]
    cat_test(vcat, randn(r, 1), randn(r))
    cat_test(vcat, randn(r), randn(r, 1))
    cat_test(vcat, randn(r), randn(r, 1), randn(r))
  end
end

@testset "hcat" begin
  # Scalar
  @test gradient((x,y) -> sum(hcat(x,y)), 1,2) == (1,1)
  @test gradient((x,y) -> sum([x y]), 1,2) == (1,1)
  @test gradient((a,b,c,d) -> sum(sqrt, [a b;c d]), 1,1,1,4) == (0.5, 0.5, 0.5, 0.25)

  # Vector-only.
  for r in [1, 2]
    cat_test(hcat, randn(r))
    cat_test(hcat, randn(r), randn(r))
  end

  # Matrix-only.
  for r in [1, 2], c1 in [1, 2], c2 in [1, 2]
    cat_test(hcat, randn(r, c1), randn(r, c2))
  end

  # Matrix-Vector / Vector-Matrix / Vector-Matrix-Vector.
  for r in [1, 2], c in [1, 2]
    cat_test(hcat, randn(r, c), randn(r))
    cat_test(hcat, randn(r), randn(r, c))
    cat_test(hcat, randn(r), randn(r, c), randn(r))
  end
end

@testset "hvcat" begin
  @test gradient(xs -> hvcat((2,2),xs...)[1,1], [1,2,3,4])[1] == (1,0,0,0)
  @test gradient(xs -> hvcat((2,2),xs...)[2,1], [1,2,3,4])[1] == (0,0,1,0)
  @test gradient(xs -> hvcat((2,2),xs...)[1,2], [1,2,3,4])[1] == (0,1,0,0)
  @test gradient(xs -> hvcat((2,2),xs...)[2,2], [1,2,3,4])[1] == (0,0,0,1)
end

@testset "cat(..., dims = $dim)" for dim in 1:5
  catdim = (x...) -> cat(x..., dims = dim)
  @test gradtest(catdim, rand(5), rand(5))
  @test gradtest(catdim, rand(2,5), rand(2,5), rand(2,5))
  @test gradtest(catdim, rand(2,5,3), rand(2,5,3), rand(2,5,3))
end

@testset "one(s) and zero(s)" begin
  @test Zygote.gradient(x->sum(ones(size(x))), randn(5))[1] isa Nothing
  @test Zygote.gradient(x->sum(one(x)), randn(3, 3))[1] isa Nothing
  @test Zygote.gradient(x->sum(zeros(size(x))), randn(7))[1] isa Nothing
  @test Zygote.gradient(x->sum(zero(x)), randn(3))[1] isa Nothing
end

@testset "fma and muladd" begin
    @test gradcheck(x -> fma(x[1], x[2], x[3]), [2.0, 3.0, 5.0])
    @test gradcheck(x -> muladd(x[1], x[2], x[3]), [2.0, 3.0, 5.0])
end

import StatsFuns

Zygote.refresh()

@testset "xlogx" begin
  @test gradcheck(x->2.5 * StatsFuns.xlogx(x[1]), [1.0])
  @test gradcheck(x->2.5 * StatsFuns.xlogx(x[1]), [2.45])
end

@testset "logistic" begin
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [-5.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [-1.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [-eps()])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [0.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [eps()])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [1.0])
  @test gradcheck(x->3.0 * StatsFuns.logistic(x[1]), [5.0])
end

@testset "logit" begin
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.1])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.3])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.5])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.7])
  @test gradcheck(x->5.0 * StatsFuns.logit(x[1]), [0.9])
end

function test_log1pexp(T, xs)
  y = T(4.3)
  for x in xs
    @test gradcheck(x->y * StatsFuns.log1pexp(x[1]), [x])
  end
end

@testset "log1pexp" begin
  @testset "Float64" begin
    @testset "x ∈ (-∞, 18.0)" begin
      test_log1pexp(Float64, [-1000.0, -50.0, -25.0, -10.0, 0.0, 10.0, 18.0 - eps()])
    end
    @testset "x ∈ [18.0, 33.3)" begin
      test_log1pexp(Float64, [18.0, 18.0 + eps(), 33.3 - eps()])
    end
    @testset "x ∈ [33.3, ∞)" begin
      test_log1pexp(Float64, [33.3, 33.3 + eps(), 100.0])
    end
  end
end

@testset "log1psq" begin
  rng = MersenneTwister(123456)
  @testset "Float64" begin
    for x in [-10.0, -5.0, -1.0, -eps(), 0.0, eps(), 1.0, 5.0, 10.0]
      @test gradcheck(x->5.1 * StatsFuns.log1psq(x[1]), [x])
    end
  end
end

@testset "logsumexp" begin
  rng = MersenneTwister(123456)
  @testset "Float64" begin
    @test gradtest(StatsFuns.logsumexp, randn(rng, 1))
    @test gradtest(StatsFuns.logsumexp, randn(rng, 1, 1))
    @test gradtest(StatsFuns.logsumexp, randn(rng, 3))
    @test gradtest(StatsFuns.logsumexp, randn(rng, 3, 4, 5))
  end
end

@testset "* sizing" begin
  @test size(Zygote.gradient((x, y)->sum(x * y), randn(1, 1), randn(1, 10))[1]) == (1, 1)
  @test size(Zygote.gradient((x, y)->sum(x * y), randn(1, 1), randn(1, 10))[2]) == (1, 10)
end

@testset "broadcast" begin
  @test gradient(x -> sum(sin.(x)), Diagonal(randn(3)))[1][2] == 1
end

using Zygote: Buffer

@testset "Buffer" begin
  @test gradient([1, 2, 3]) do x
    b = Buffer(x)
    b[:] = x
    return sum(copy(b))
  end == ([1,1,1],)

  function vstack(xs)
    buf = Buffer(xs, length(xs), 5)
    for i = 1:5
      buf[:, i] = xs
    end
    return copy(buf)
  end

  @test gradient(x -> sum(vstack(x)), [1, 2, 3]) == ([5, 5, 5],)

  buf = Buffer([1, 2, 3])
  buf[1] = 1
  copy(buf)
  @test_throws ErrorException buf[1] = 1
  @test eltype(buf) === Int
  @test length(buf) === 3
  @test ndims(buf) === 1
  @test size(buf) === (3, )
  @test size(buf, 2) === 1
  @test axes(buf) == (1:3, )
  @test axes(buf, 2) == 1:1
  @test eachindex(buf) == 1:3
  @test stride(buf, 2) === 3
  @test strides(buf) === (1, )

  @test gradient([1, 2, 3]) do xs
    b = Zygote.Buffer(xs)
    b .= xs .* 2
    return sum(copy(b))
  end == ([2,2,2],)
end

@testset "FillArrays" begin
  @test gradcheck(x->sum(Fill(x[], (2, 2))), [0.1])
  @test first(Zygote.gradient(sz->sum(Ones(sz)), 6)) === nothing
  @test first(Zygote.gradient(sz->sum(Zeros(sz)), 6)) === nothing
end

@testset "AbstractArray Addition / Subtraction / Negation" begin
  rng, M, N, P = MersenneTwister(123567), 3, 7, 11
  A, B = randn(rng, M, N, P), randn(rng, M, N, P)
  @test gradtest(+, A, B)
  @test gradtest(-, A, B)
  @test gradtest(-, A)
end

@testset "FFTW" begin
  x=[-0.353213 -0.789656 -0.270151; -0.95719 -1.27933 0.223982]
  # gradient of ifft(rfft) must be 1
  @test gradient((x)->real(ifft(fft(x))[1]),x)[1][1] == 1.0+0.0im
  @test gradient((x)->real(fft(ifft(x))[1]),x)[1][1] == 1.0+0.0im

  # check ffts for individual dimensions
  @test gradient((x)->sum(abs.(FFTW.fft(x))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
  @test gradient((x)->abs(sum((FFTW.fft(x)))),x)[1] ≈ gradient((x)->abs(sum(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
  @test gradient((x, dims)->sum(abs.(FFTW.fft(x,dims))),x,(1,2))[1] ≈ gradient((x)->sum(abs.(FFTW.fft(x))),x)[1]
  @test gradient((x)->sum(abs.(FFTW.fft(x,(1,2)))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.fft(FFTW.fft(x,1),2))),x)[1]
  @test gradient((x, dims)->sum(abs.(FFTW.ifft(x,dims))),x,(1,2))[1] ≈ gradient((x)->sum(abs.(FFTW.ifft(x))),x)[1]
  @test gradient((x)->sum(abs.(FFTW.ifft(x,(1,2)))),x)[1] ≈ gradient((x)->sum(abs.(FFTW.ifft(FFTW.ifft(x,1),2))),x)[1]

  @test gradcheck(x->sum(abs.(FFTW.fft(x))), x)
  @test gradcheck(x->sum(abs.(FFTW.ifft(x))), x)
  @test gradcheck(x->sum(abs.(FFTW.fft(x, 1))), x)
  @test gradcheck(x->sum(abs.(FFTW.ifft(x, 1))), x)

end

@testset "FillArrays" begin
  rng, M, N = MersenneTwister(123456), 7, 11
  x, y = randn(rng), randn(rng)
  @test Zygote.gradient(x->sum(Fill(x, N)), x)[1] == N
  @test Zygote.gradient(x->sum(Fill(x, N, 3, 4)), x)[1] == N * 3 * 4
  @test Zygote.gradient((x, y)->sum(Fill(x, N)), x, y) == (N, nothing)

  let
    out, back = Zygote.pullback(sum, Fill(x, N))
    @test back(nothing) isa Nothing
  end

  z = randn(rng, N)
  @test gradtest(x->Fill(first(x), N), [x])
  let
    out, back = Zygote.pullback(x->Fill(x, N), x)
    @test out == Fill(x, N)
    @test first(back(Fill(y, N))) ≈ y * N
  end

  # Test unary broadcasting gradients.
  out, back = Zygote.pullback(x->exp.(x), Fill(x, N))
  @test out isa Fill
  @test out == Fill(exp(x), N)
  @test back(Ones(N))[1] isa Fill
  @test back(Ones(N))[1] == Ones(N) .* exp(x)
  @test back(ones(N))[1] isa Vector
  @test back(ones(N))[1] == ones(N) .* exp(x)
  @test gradtest(x->exp.(Fill(3 * first(x), N)), [x])

  @testset "broadcast + and *" begin
    for sx in [(M, N), (M, 1), (1, N), (1, 1)]
      for sy in [(M, N), (M, 1), (1, N), (1, 1)]
        z = randn(rng, broadcast_shape(sx, sy))

        # Addition
        @test gradtest((x, y)->Fill(first(x), sx...) .+ Fill(first(y), sy...), [x], [y])
        @test gradtest(x->Fill(first(x), sx...) .+ Ones(sy...), [x])
        @test gradtest(x->Fill(first(x), sx...) .+ Zeros(sy...), [x])
        @test gradtest(y->Ones(sx...) .+ Fill(first(y), sy...), [y])
        @test gradtest(y->Zeros(sx...) .+ Fill(first(y), sy...), [y])

        # Multiplication
        @test gradtest((x, y)->Fill(first(x), sx...) .* Fill(first(y), sy...), [x], [y])
        @test gradtest(x->Fill(first(x), sx...) .* Ones(sy...), [x])
        @test gradtest(x->Fill(first(x), sx...) .* Zeros(sy...), [x])
        @test gradtest(y->Ones(sx...) .* Fill(first(y), sy...), [y])
        @test gradtest(y->Zeros(sx...) .* Fill(first(y), sy...), [y])
      end
    end
  end
end

@testset "@nograd" begin
  @test gradient(x -> findfirst(ismissing, x), [1, missing]) == (nothing,)
  @test gradient(x -> findlast(ismissing, x), [1, missing]) == (nothing,)
  @test gradient(x -> findall(ismissing, x)[1], [1, missing]) == (nothing,)
end


#==============================================================================
    function to return all neighbors in arbitrary dimensions, from an index
        getneighbors is 2D version
        getnbs is N dimensional version

    tested:     YES
    author:     discourse.julialang.org/u/Crown421
 =============================================================================#

using   BenchmarkTools

function getneighbours(M, cartIdx, radius = 1)
    return [M[i,j] for i in cartIdx[1]-radius:cartIdx[1]+radius, j in cartIdx[2]-radius:cartIdx[2]+radius
           if !((i == cartIdx[1]) & (j == cartIdx[2]))&& (1<= i <= size(M, 1)) && (1<= j <= size(M, 2))]
end

function getnbs(M, startIdx::CartesianIndex; radius = 1)
    idx = CartesianIndices(ntuple(i->(-radius:radius), ndims(M))) .+ startIdx
    idx = idx[:]
    deleteat!(idx, ceil(Int, length(idx)/2))
    filter!(id -> checkbounds(Bool, M, id) , idx)
    return M[idx]
end

M = LinearIndices((300,300))

# M = Int.(t)   # t not defined. not sure intention here, but works without it

testIdx = CartesianIndex((30,30))

nabes = getneighbours(M, testIdx)

display(@benchmark getneighbours(M, testIdx))
display(@benchmark getnbs(M, testIdx))


########################################################################
# Expected latent info -  for 1-D latent and 1PL pars                  #
########################################################################
"""
```julia
_latent_information(latent::Latent1D, parameters::Parameters1PL)
```

# Description

It computes the information (second derivative of the likelihood) with respect to the 1-dimensional latent variable under the 1PL model.
It follows the parametrization ``a(θ - b)``.

# Arguments

- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.
- **`parameters::Parameters1PL`** : Required. A 1-parameter logistic parameters object.

# Output

A `Float64` scalar.
"""
function _latent_information(latent::Latent1D, parameters::Parameters1PL)
    p = _probability(latent, parameters)
    return p * (1 - p)::Float64
end

########################################################################
# Expected latent info -  for 1-D latent and 2PL pars                  #
########################################################################
"""
```julia
_latent_information(latent::Latent1D, parameters::Parameters2PL)
```

# Description

It computes the information (-second derivative of the likelihood) with respect to the 1-dimensional latent variable under the 2PL model.
It follows the parametrization ``a(θ - b)``.

# Arguments

- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.
- **`parameters::Parameters2PL`** : Required. A 2-parameter logistic parameters object.

# Output

A `Float64` scalar.
"""
function _latent_information(latent::Latent1D, parameters::Parameters2PL)
    p = _probability(latent, parameters)
    return Float64(p * (1 - p) * parameters.a^2)::Float64
end

########################################################################
# Expected latent info -  for 1-D latent and 3PL pars                  #
########################################################################
"""
```julia
_latent_information(latent::Latent1D, parameters::Parameters3PL)
```

# Description

It computes the information (-second derivative of the likelihood) with respect to the 1-dimensional latent variable under the 3PL model.
It follows the parametrization ``a(θ - b)``.

# Arguments

- **`latent::Latent1D`** : Required. A 1-dimensional `Latent1D` latent variable.
- **`parameters::Parameters3PL`** : Required. A 3-parameter logistic parameters object.

# Output

A `Float64` scalar.
"""
function _latent_information(latent::Latent1D, parameters::Parameters3PL)
    p = _probability(latent, parameters)
    return (p - parameters.c)^2 * (1 - p) * parameters.a^2 / (1 - parameters.c)^2 / p::Float64
end


export brainscan_poincaresos

"""
    brainscan_poincaresos(A::Dataset, j::Int; kwargs...)
    brainscan_poincaresos(As::Vector{Dataset}, j::Int; kwargs...)
Launch an interactive application for scanning a Poincare surface of section of `A`
like a "brain scan", where the plane that defines the section can be arbitrarily
moved around via a slider. Return `figure, ax3D, ax2D`.

The input dataset must be 3 dimensional, and here the crossing plane is always
chosen to be when the `j`-th variable of the dataset crosses a predefined value.
The slider automatically gets all possible values the `j`-th variable can obtain.

If given multiple datasets, the keyword `colors` attributes a color to each one, e.g.
`colors = [JULIADYNAMICS_COLORS[mod1(i, 6)] for i in 1:length(As)]`.

The keywords `linekw, scatterkw` are named tuples or dictionaries that are propagated to
as keyword arguments to the line
and scatter plot respectively, while the keyword `direction = -1` is propagated
to the function `DyamicalSystems.poincaresos`.
"""
function brainscan_poincaresos(A::DynamicalSystems.AbstractDataset, j::Int; kwargs...)
    brainscan_poincaresos([A], j; kwargs...)
end

function brainscan_poincaresos(
    As::Vector{<:DynamicalSystems.AbstractDataset}, j::Int;
    linekw = (), scatterkw = (), direction = -1,
    colors = [CYCLIC_COLORS[i] for i in 1:length(As)]
)

for A in As; @assert size(A, 2) == 3; end
@assert j ∈ 1:3
mi, ma = total_minmaxima(As)
ms = 25maximum(abs(ma[i] - mi[i]) for i in 1:3)

otheridxs = DynamicalSystems.SVector(setdiff(1:3, j)...)

figure = Figure(resolution = (2000, 800))
display(figure)
ax = figure[1, 1] = Axis3(figure)
axp = figure[1, 2] = Axis(figure)
sll = labelslider!(
    figure, "$(('x':'z')[j]) =", range(mi[j], ma[j]; length = 100);
    sliderkw = Dict(:startvalue => (ma[j]+mi[j])/2)
)
figure[2, :] = sll.layout
y = sll.slider.value

for i in 1:length(As)
    A = As[i]
    # plot 3D trajectory
    lines!(ax, A.data; color = colors[i], transparency = false, linekw...)

    # Poincare sos
    psos = lift(y) do y
        DynamicalSystems.poincaresos(A, (j, y); direction, warning = false)
    end
    psos2d = lift(p -> p[:, otheridxs].data, psos)
    psos3d = lift(p -> p.data, psos)

    Makie.scatter!(axp, psos2d; color = colors[i], scatterkw...)
    Makie.scatter!(ax, psos3d; color = colors[i], markersize = ms, scatterkw...)
end

xlims!(axp, mi[otheridxs[1]], ma[otheridxs[1]])
ylims!(axp, mi[otheridxs[2]], ma[otheridxs[2]])

# plot transparent plane
ss = [mi...]
ws = [(ma - mi)...]
ws[j] = 0

p = lift(y) do y
    ss[j] = y
    o = Makie.Point3f0(ss...)
    w = Makie.Point3f0(ws...)
    Makie.FRect3D(o, w)
end

a = Makie.RGBAf0(0,0,0,0)
c = Makie.RGBAf0(0.2, 0.2, 0.25, 1.0)
img = Makie.ImagePattern([c a; a c]);
Makie.mesh!(ax, p; color = img);

return figure, ax, axp
end

function total_minmaxima(As::Vector{<:DynamicalSystems.AbstractDataset})
    mi, ma = DynamicalSystems.minmaxima(As[1])
    mi = Vector(mi); ma = Vector(ma)
    for j in 2:length(As)
        mi2, ma2 = DynamicalSystems.minmaxima(As[j])
        for i in 1:size(As[j], 2)
            mi[i] = min(mi2[i], mi[i])
            ma[i] = max(ma2[i], ma[i])
        end
    end
    return mi, ma
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
using Plots
pyplot()

function halftorus(n=25, a=5, b=10; lims=(-11, 11), size=(500, 400))
    u = range(-π/2, π/2, length=n+1)
    v = range(0, 2π, length=2n+1)
    x = @. (b + a * cos(u')) * cos(v)
    y = @. (b + a * cos(u')) * sin(v)
    z = @.      a * sin(u')  * one(v)
    surface(x, y, z; lims, colorbar=false, size)
end

halftorus()

# %%