export compose_rotation

################################################################################
#                              Compose Rotations
################################################################################

"""
    @inline function compose_rotation(R1, [, R2, R3, R4, R5, ...])

Compute a composed rotation using the rotations `R1`, `R2`, `R3`, `R4`, ..., in
the following order:

     First rotation
     |
     |
    R1 => R2 => R3 => R4 => ...
           |
           |
           Second rotation

The rotations can be described by Direction Cosine Matrices or Quaternions.
Notice, however, that all rotations **must be** of the same type (DCM or
quaternion).

The output will have the same type as the inputs (DCM or quaternion).

# Example

```julia-repl
julia> D1 = angle2dcm(+pi/3,+pi/4,+pi/5,:ZYX);

julia> D2 = angle2dcm(-pi/5,-pi/4,-pi/3,:XYZ);

julia> compose_rotation(D1,D2)
3×3 StaticArrays.SArray{Tuple{3,3},Float64,2,9}:
 1.0          0.0          5.55112e-17
 0.0          1.0          5.55112e-17
 5.55112e-17  5.55112e-17  1.0

julia> q1 = angle2quat(+pi/3,+pi/4,+pi/5,:ZYX);

julia> q2 = angle2quat(-pi/5,-pi/4,-pi/3,:XYZ);

julia> compose_rotation(q1,q2)
Quaternion{Float64}:
  + 1.0 + 0.0.i + 2.0816681711721685e-17.j + 5.551115123125783e-17.k
```

"""
@inline compose_rotation(D::DCM) = D
@inline compose_rotation(D::DCM, Ds::DCM...) = compose_rotation(Ds...)*D

@inline compose_rotation(q::Quaternion) = q
@inline compose_rotation(q::Quaternion, qs::Quaternion...) =
    q*compose_rotation(qs...)

# This algorithm was proposed by @Per in
#
#   https://discourse.julialang.org/t/improve-the-performance-of-multiplication-of-an-arbitrary-number-of-matrices/10835/24


using SafeTestsets, Test

@safetestset "Values at origin" begin include("origin.jl") end
@safetestset "Orientation" begin include("orientation.jl") end
@safetestset "Field values compared with symbolic results" begin include("field_values.jl") end
@safetestset "Divergence in Polynomial Roots" begin include("roots.jl") end
@safetestset "LG(p = 0, m = 0, x, y, z, t) == G(x, y, z, t)" begin include("gauss.jl") end
@safetestset "Comparison to Symbolic Values" begin include("field_values.jl") end


using Powersense, GLPK
using JuMP

"""
The attributes to the Powersense and the external solver can be added using the following method with JuMP
"""
#optimizer_with_attributes( Powersense.Optimizer, "external_optimizer" => optimizer_with_attributes((GLPK.Optimizer,"msg_lev"=>GLPK.MSG_OFF)))

model = Model(Powersense.Optimizer);
set_optimizer_attribute(model, "external_optimizer", GLPK.Optimizer)

@variable(model, X);
@variable(model, Y);
@objective(model, Min, X^2 + X);
@NLconstraint(model, X^2 - X == 2);
@NLconstraint(model, X*Y == 1);
@NLconstraint(model, X*Y >= 0);
@constraint(model, X >= -2);

println("________________________________________");
print(model);
println("________________________________________");
JuMP.optimize!(model);

xsol = JuMP.value.(X)
ysol = JuMP.value.(Y)
status = termination_status(model)

println("Xsol = ", xsol);
println("Ysol = ", ysol);

println("Status: ", status);


using SparseArrays
using DataStructures: Accumulator, counter
using Turing
using Distributions

function extract_TDM(file = "epigrams.txt")
    epigram_bags = Vector{Accumulator{UInt64, Int8}}()

    for epigram in eachline(file)
        push!(epigram_bags, counter(hash.(split(epigram))))
    end

    D = length(epigram_bags)
    T = length(reduce(union, epigram_bags))
    W = sparse([], [], Int8[], D, T)

    for (epigram, bag) in enumerate(epigram_bags)
        for (word, count) in bag
            W[epigram, Int(rem(word, T)) + 1] += count
        end
    end

    return W
end

@model LatentDirichletAllocation(W, ntopics = 10, α = 0.5, β = 0.5) = begin
    ndocs, ntypes = size(W)
    ntokens = sum(W)

    # prior for topic distribution per document
    θ = Vector{Vector{Real}}(undef, ndocs)
    for d = 1:ndocs
        θ[d] ~ Dirichlet(ntopics, α)
    end

    # topic distribution per document
    Z = Vector{Vector{Int}}(undef, ndocs)
    for d = 1:ndocs
        z[d] ~ Categorical(θ[d])
    end

    # prior for topic distribution per word
    ϕ = Vector{Vector{Real}}(undef, ntopics)
    for t = 1:ntopics
        ϕ[t] ~ Dirichlet(ntypes, β)
    end

    # distribution of word per document
    for d = 1:ndocs, w = 1:ntypes
        W[d, w] ~ Categorical(ϕ[Z[d][w]])
    end
end

function test()
    W = extract_TDM()
    chain = sample(LatentDirichletAllocation(W), IS(), 10)
    describe(chain)
end

test()


using Zygote, LinearAlgebra
using KrylovKit: eigsolve

test_eigmin(A::AbstractMatrix) = eigsolve(A, 1, :SR)[1][1]

A = rand(4, 4); A = A + A';

@show eigmin(A)
@show test_eigmin(A)
@show Zygote.gradient(test_eigmin, A)


function solve()
    n, x, t = [parse(Int, x) for x in split(readline())]
    t * ((n ÷ x) + (n % x > 0 ? 1 : 0))
end

println(solve())


include("BaryTrees/BaryTrees.jl")
using Main.BaryTrees, Plots

# construct an empty tree from a Barycentric coordinate system defined
# by the triangle with vertices (0,sqrt(3)/2), (-0.5,0), (0.5,0)
T = BaryTree(
    BaryCentric(Point((0,sqrt(3)/2)),Point((-0.5,0)),Point((0.5,0.0)))
)

# random points in a square: (-1,-1), (-1,1), (1,1), (1,-1)
p = (rand(2^9,2) .- 0.5).*2;

pin = zeros(1,2)
pout = zeros(1,2)

for i in 1:size(p,1)
    b = Cartesian2BaryCentric(Point((p[i,:]...)),T.boundary.b)

    # test if in trig
    if b[1] >= 0 && b[2] >= 0 && b[1]+b[2] < 1
        pin = cat(pin,p[i:i,:],dims=1)
    else
        pout = cat(pout,p[i:i,:],dims=1)
    end
end

# insert only points inside trig
for i in 1:size(pin,1)
    insert!(Point((pin[i,:]...)),T)
end

scatter(pin[:,1],pin[:,2],label="",dpi=300,
    aspect_ratio=:equal,xaxis=nothing,yaxis=nothing,axis=:off
)

Draw!(T)
plot!()

savefig("barytree.png")


@testset "solve_diffusion.jl" begin
    domain = ClosedIntervalDomain(0.0, 1.0)
    a, b = domain.left, domain.right

    T = 1.0

    # Exact solutions (from Borgaard 2006)
    u(ξ, t) = t + sin(2π * ξ) + sin(8π * ξ)
    u_int(ξ, t) = t * ξ - 1 / 2π * cos(2π * ξ) - 1 / 8π * cos(8π * ξ)
    f(ξ, t) = 1 + 4π^2 * sin(2π * ξ) + 64π^2 * sin(8π * ξ)
    g_a(t) = t
    g_b(t) = t

    tols = (; abstol = 1e-6, reltol = 1e-4)

    # Number of mesh points
    NN = floor.(Int, 10 .^ LinRange(2.5, 4.0, 4))

    err_ref = 0
    err_bar_ref = 0
    err_allbar_ref = 0
    n_ref = 0
    for (i, N) ∈ enumerate(NN)
        M = N

        # Discretization
        x = discretize(domain, M)
        ξ = discretize(domain, N)
        Δx = (domain.right - domain.left) / N

        # Filter
        h(x) = Δx / 2

        # Exact filtered solution
        function ū(x, t)
            α = max(a, x - h(x))
            β = min(b, x + h(x))
            1 / (β - α) * (u_int(β, t) - u_int(α, t))
        end

        # Equations
        equation = DiffusionEquation(domain, IdentityFilter(), f, g_a, g_b)
        equation_filtered = DiffusionEquation(domain, TopHatFilter(h), f, g_a, g_b)

        # Solve discretized problem
        sol = solve(
            equation,
            ξ -> u(ξ, 0.0),
            (0.0, T),
            M,
            N;
            method = "discretizefirst",
            boundary_conditions = "derivative",
            tols...,
        )

        # Solve filtered-then-discretized problem
        @test_broken sol_bar = solve(
            equation_filtered,
            ξ -> u(ξ, 0.0),
            (0.0, T),
            M,
            N;
            method = "filterfirst",
            boundary_conditions = "derivative",
            tols...,
        )

        # Solve discretized-then-filtered problem
        sol_allbar = solve(
            equation_filtered,
            ξ -> u(ξ, 0.0),
            (0.0, T),
            M,
            N;
            method = "discretizefirst",
            boundary_conditions = "derivative",
            tols...,
        )

        ## Relative error
        u_exact = u.(ξ, T)
        ū_exact = ū.(x, T)
        err = norm(sol(T) - u_exact) / norm(u_exact)
        # err_bar = norm(sol_bar(T) - ū_exact) / norm(ū_exact)
        err_allbar = norm(sol_allbar(T) - ū_exact) / norm(ū_exact)

        if i == 1
            err_ref = err
            # err_bar_ref = err_bar
            err_allbar_ref = err_allbar
            n_ref = N
        end

        # Test for quadratic convergence
        @test err < 1.6 * err_ref * (n_ref / N)^2
        # @test err_bar < 1.6 * err_bar_ref * (n_ref / N)^2
        @test err_allbar < 1.6 * err_allbar_ref * (n_ref / N)^2
    end

    # Test exact boundary conditions
    N = 100
    M = N
    equation = DiffusionEquation(domain, IdentityFilter(), f, g_a, g_b)

    # Solve discretized problem with DAE formulation
    sol_exact = solve(
        equation,
        ξ -> u(ξ, 0.0),
        (0.0, T),
        M,
        N;
        method = "discretizefirst",
        boundary_conditions = "exact",
        tols...,
    )

    # Solve discretized problem with derivative BC
    sol_derivative = solve(
        equation,
        ξ -> u(ξ, 0.0),
        (0.0, T),
        M,
        N;
        method = "discretizefirst",
        boundary_conditions = "derivative",
        tols...,
    )

    @test sol_exact(T) ≈ sol_derivative(T)
end


module StreamStats
    import StatsBase, Distributions

    export update!, state, nobs, replicates

    include("streamstat.jl")
    include("mean.jl")
    include("var.jl")
    include("moments.jl")
    include("cov.jl")
    include("min.jl")
    include("max.jl")
    # include("approx_quantile.jl")
    # include("approx_cdf.jl")
    # include("approx_pmf.jl")
    include("hyper_log_log.jl")
    include("approx_distinct.jl")
    include("bootstrap.jl")
    include("sample.jl")
    include("approx_ols.jl")
    include("approx_logit.jl")
    include("approx_ridge.jl")
    include("approx_l2_logit.jl")
end


module IntegerIB

using Random
using DataFrames
import Base.merge!

"""
    Computes the KL divergence for two 1D probability distributions.
"""
function KL_single(P, Q)
    return sum(P.*log2.(P./Q))
end

"""
    Computes the KL divergence for a collection of distributions in matrix form.
Given P of shape MxN and Q of shape MxL, returns a matrix of shape NxL.
"""
function KL(P, Q)
    kl = zeros(size(P)[2], size(Q)[2])
    for i in 1:size(P)[2]
        for j in 1:size(Q)[2]
            kl[i,j] = KL_single(P[:,i],Q[:,j])
        end
    end
    return kl
end

"""
    Computes the entropy of a distribution. If distribution is of shape (MxN),
treats each column as an independant distribution and returns an 1xN entropy vector.
"""
function entropy(p)
    if any(size(p) .== 1)
        return -sum(p.*log2.(p))
    else
        return -sum(p.*log2.(p), dims = 1)
    end
end

"""
    get_y(x, algorithm = "nn")

Returns the corresponding context vector `y` to the data vector `x`.
Several types of context are possible, and are selected by the `type` argument.
type = "nn" is the next neighbor to every point of `x`.
type = "an" stands for "adjacent neighbor", or the previous and next element of every element of x.
"""
function get_y(x, type = "nn")
    if type == "nn"
        y = x[2:end]
    elseif type == "an"
        y = Any[]
        for idx in collect(2:length(x)-1)
            push!(y, (x[idx-1], x[idx+1]))
        end
    end
    return y
end

function findin(el, array)
    return findall(x -> x == el, array)[1]
end

"""
    get_pxy(x,y)
Returns joint probability of x and y from normalized 2D histogram.
Automatically replaces 0 by 10^-16 to prevent Inf and NaN values during KL computation.
"""
function get_pxy(x,y)
    X = x[1:end-1]
    if length(X) != length(y)
        X = X[2:end]
    end
    count = 0
    x_elts, y_elts = sort(unique(X)), sort(unique(y))
    pxy = zeros(size(x_elts)[1], size(y_elts)[1])
    for (idx, el) in enumerate(X)
        xpos, ypos = findin(el, x_elts), findin(y[idx], y_elts)
        pxy[xpos, ypos] += 1
        count += 1
    end
    pxy[pxy .== 0.0] .= 10^-16 #replacing zeros

    return pxy./count
end

"""
    returns the conditional probability of y given x.
"""
function get_py_x(x, y)
    pyx = get_pxy(x, y)
    inv_px = 1 ./ get_px(x, y)
    cond_pyx = inv_px.*pyx
    return cond_pyx' #transposition in order to get the y dimension in dims = 1.
end

function get_py_x(pxy)
    inv_px = 1 ./ get_px(pxy)
    cond_pyx = inv_px.*pxy
    return cond_pyx' #transposition in order to get the y dimension in dims = 1.
end

"""
    Computes marginal probability of 'x'.
"""
function get_px(x, y)
    return sum(get_pxy(x, y), dims = 2)
end

function get_px(pxy)
    return sum(pxy, dims = 2)
end

"""
    Computes marginal probability of 'y'.
"""
function get_py(x, y)
    return sum(get_pxy(x, y), dims = 1)
end

function get_py(pxy)
    return sum(pxy, dims = 1)
end

function get_qt_x_init(px, algorithm)
    if algorithm == "IB"
        qt_x_init = rand(size(px)[1], size(px)[1])  # randomly initialize qt_x_init
        qt_x_init = qt_x_init.*(1 ./ sum(qt_x_init, dims = 1)) #normalize qt_x_init
        return qt_x_init
    elseif algorithm == "DIB"
        index = shuffle(collect(1:size(px)[1]))
        qt_x_init = zeros(size(px)[1], size(px)[1])
        for (i,v) in enumerate(index)
            qt_x_init[v,i] = 1
        end
        return replace_zeros(qt_x_init)
    else
        throw("Algorithm must be \"IB\" pr \"DIB\"")
    end
end

"""
    init_values(x, y)

Initializes and return clustering porbability matrices `qt_x`, `q` and `qy_t`.
'qt_x' refers to the probability of a category in x to be mapped to an element of t.
'q' is the overall probability of an element of t to occur. 'qy_t' is the probability that an element of t
corresponds to a given element of `y`.
"""
function init_values(x, y, algorithm = "IB")
    pxy = get_pxy(x, y)
    px = get_px(x, y)
    qt_x_init = get_qt_x_init(px, algorithm)
    qt_init = qt_x_init*px
    qy_t_init = pxy'*qt_x_init'.*(1 ./ qt_init)'
    return (qt_x_init, qt_init, qy_t_init)
end

function init_values(pxy, algorithm = "IB")
    px = get_px(pxy)
    qt_x_init = get_qt_x_init(px, algorithm)
    qt_init = qt_x_init*px
    qy_t_init = pxy'*qt_x_init'.*(1 ./ qt_init)'
    return (qt_x_init, qt_init, qy_t_init)
end

"""
Interface for IB clustering computation.
Input data 'x' must be 1D. β controls the amount of clustering : the smaller the beta, the greater the clustering (and information loss).
Can be initalized directly with the data via IB(x), or if you already have the histograms/probability of co-occurences of 'x' and 'y' via IB(pxy).
To use the deterministic IB algorithm, set algorithm to "DIB" like IB(x, "DIB").
"""
mutable struct IB
    algorithm::String
    β::Real
    x
    px
    py
    pxy
    py_x
    qt_x
    qt
    qy_t
    IB(x, y, β = 100, algorithm = "IB") = new(algorithm, β, x, get_px(x,y), get_py(x,y), get_pxy(x,y), get_py_x(x,y), init_values(x, y, algorithm)...)
    IB(x::Array{Float64,1}, β = 100, algorithm = "IB") = new(algorithm, β, x, get_px(x, get_y(x)), get_py(x, get_y(x)), get_pxy(x, get_y(x)), get_py_x(x, get_y(x)), init_values(x,  get_y(x), algorithm)...)
    IB(x::Array{Int64,1}, β = 100, algorithm = "IB") = new(algorithm, β, x, get_px(x, get_y(x)), get_py(x, get_y(x)), get_pxy(x, get_y(x)), get_py_x(x, get_y(x)), init_values(x,  get_y(x), algorithm)...)
    IB(pxy::Array{Float64,2}, β = 100, algorithm = "IB") = new(algorithm, β, nothing, get_px(pxy), get_py(pxy), pxy, get_py_x(pxy), init_values(pxy, algorithm)...)
end

"""
    Re-initialies values of IB model
"""
function init!(m::IB)
    m.qt_x, m.qt, m.qy_t = init_values(m.pxy, m.algorithm)
end

"""
    Peforms q(t|x) update step.
"""
function qt_x_step!(m::IB)
    kl_part = log.(m.qt)' .- m.β*KL(m.py_x, m.qy_t)
    if m.algorithm == "DIB"
        m.qt_x = zeros(size(m.qt_x))
        for x_index in 1:size(kl_part)[1]
            m.qt_x[argmax(kl_part[x_index,:]), x_index] = 1
        end
        m.qt_x = replace_zeros(m.qt_x)
    elseif m.algorithm == "IB"
        q = replace_zeros(exp.(kl_part))
        m.qt_x = q'.*(1 ./ sum(q', dims = 1)) #transpose and normalize q.
    else
        throw(ArgumentError("'algorithm' unknown must be either \"IB\" or \"DIB\""))
    end
end

"""
    Peforms q(t) update step and drops unused clusters.
"""
function qt_step!(m::IB, clust_thres; report = true)
    m.qt = m.qt_x*m.px
    dropped = m.qt .< clust_thres
    if any(dropped) && size(m.qt)[1] > 1
        if report
            println("$(length(dropped[dropped])) clusters dropped, $(length(dropped[.~dropped])) ramaining.")
        end
        m.qt = m.qt[.~dropped]
        m.qt_x = drop_rows(m.qt_x, dropped)
        m.qt_x = m.qt_x.*(1 ./ sum(m.qt_x, dims = 1)) #normalize cluster distribution again
        m.qt = m.qt_x*m.px
    end
end

"""
    Remove rows of 2D array 'value_arr'. Boolean array 'drop_arr' indicates which rows to drop.
"""
function drop_rows(value_arr, drop_arr)
    new_size = length(drop_arr[.~drop_arr])
    dropped_arr = zeros(new_size, size(value_arr)[2])
    count = 1
    for (idx, drop) in enumerate(drop_arr)
        if ~drop
            dropped_arr[count,:] = value_arr[idx,:]
            count += 1
        end
    end
    return dropped_arr
end

"""
    Performs q(y|t) update step.
"""
function qy_t_step!(m::IB)
    qt_y = m.qt_x*m.pxy
    m.qy_t = qt_y'.*(1 ./ m.qt')
end

"""
    replace_zeros(arr::Array{Float64,2}, cut = 10^-305)

Replaces zeros (values below 'cut') by 'cut' to avoid exponential term causing NaN at high beta values.
"""
function replace_zeros(arr::Array{Float64,2}, cut = 10^-305)
    arr[arr .< cut] .= cut
    return arr
end

"""
    Performs one iteration of the IB algorithm.
"""
function IB_step!(m::IB, clust_thres; report = true)
    qt_x_step!(m)
    qt_step!(m, clust_thres; report = report)
    qy_t_step!(m)
end

"""
    Checks if absolute or relative convergence criterion on L function is satisfied.
"""
function check_conv(L_prev, L_current, abs_thres, rel_thres)
    abs_change = abs(L_prev - L_current)
    rel_change = abs_change/L_prev
    if abs_change <= abs_thres || rel_change <= rel_thres
        return true
    else
        return false
    end
end

function calc_metrics(m::IB)
    ht = entropy(m.qt)
    hy_t = entropy(m.qy_t)*m.qt
    iyt = entropy(m.py) - hy_t[1] # [1] cause of dumb julia rules focing dot product of vector to be 2 dimensional.
    ht_x = entropy(m.qt_x)*m.px
    ixt = entropy(m.qt) - ht_x[1] # same as iyt.
    L = ht - ht_x[1] - m.β*iyt    # same
    return (ht, ixt, iyt, L)
end

"""
    IB_optimize(model::IB; abs_thres = 10^-3, rel_thres = 0.01, ct_thres = 5, n_conv = 1000)

Optimizes IB model of the data.
Keywords arguments are :
    - abs_thres : threshold for convergence in absoulte change of the optimization function (lagrangian) L.
    - rel_thres : threshold for convergence in relative change of L.
    - ct_thres : number of steps where L has to be < to abs_thres or rel_thres to declare convergence.
    - n_conv : limit of the number of iterations for the optimization algorithm.
    - scan (bool): when using 'DIB' algorithm, checks and merges clusters if it leads to reduction of cost function.
    - report : wether or not to say when and how convergence was reached.
"""
function IB_optimize!(model::IB; merge_thres = 5*10^-2, abs_thres = 10^-10, rel_thres = 0, ct_thres = 25, n_conv = 1000, scan = true, report = true)
    consec_conv = 0
    total_steps = 0
    L = calc_metrics(model)[4]
    while consec_conv <= ct_thres && total_steps <= n_conv
        IB_step!(model, merge_thres; report = report)
        if model.algorithm == "DIB" && scan
            brute_optimize!(model, false)
        end
        L_current = calc_metrics(model)[4]
        if check_conv(L, L_current, abs_thres, rel_thres)
            consec_conv += 1
        else
            consec_conv = 0
        end
        L = L_current
        total_steps += 1
        if total_steps == n_conv && report
            println("Convergence not reached after $n_conv iterations, stopping optimization.")
        elseif consec_conv == ct_thres && report
            println("Convergence reached after $total_steps iterations, stopping optimization")
        end
    end
end

"""
    scans through the qt probabilities and merges clusters where values are too low.
"""
function merge!(m::IB, clust_thres = 10^-5)
    dropped = m.qt .< clust_thres
    if any(dropped) && size(m.qt)[1] > 1
        m.qt = m.qt[.~dropped]
        m.qt_x = drop_rows(m.qt_x, dropped)
        m.qt_x = m.qt_x.*(1 ./ sum(m.qt_x, dims = 1)) #normalize cluster distribution again
        m.qt = m.qt_x*m.px
        m.qy_t = drop_rows(m.qy_t', dropped)'
    end
end

"""
    brute_optimize!(m::IB, findbest = true)

Performs bute force optimization of IB model, looking at all possible cluster combinations
and choosing the merges that reduce the cost function.
"""
function brute_optimize!(m::IB, findbest = true)
    if m.algorithm == "IB"
        println("IB agorithm must be DIB.")
        return
    end
    if size(m.qt_x)[1] == 1
        return
    end
    anybetter = false
    best = deepcopy(m)
    for t1 in range(1, stop = size(m.qt_x)[1]-1)
        for t2 in range(2, stop = size(m.qt_x)[1])
            if ~anybetter || findbest
                tmp = deepcopy(best)
                tmp.qt_x[t1, findall(x -> x == 1, m.qt_x[t2,:])] .= 1
                tmp.qt_x[t2,:] .= 0
                IB_step!(tmp, 0)
                # check if cost function L reduced relative to best so far
                if calc_metrics(tmp)[4] < calc_metrics(best)[4]
                    best = deepcopy(tmp)
                    mergedt1 = t1
                    mergedt2 = t2
                    anybetter = true
                    print("merged clusters $mergedt1 and $mergedt2 \n")
                end
            end
        end
    end
    merge!(best)
    if anybetter
        m.qt_x = best.qt_x
        m.qt = best.qt
        m.qy_t = best.qy_t
        return true
    else
        return false
    end
end

"""
    Scans the IB plane with various vlaues of beta to get the optimal curve in the IB plane.
"""
function get_IB_curve(m::IB, start = 0.1, stop = 400, step = 0.05; glob = false)
    beta = collect(start:step:stop)
    ixt, iyt = [], []
    for (idx, b) in enumerate(beta)
        m.β = b
        init!(m)
        if glob
            search_optima!(m)
        else
            IB_optimize!(m; report = false)
        end
        x, y = calc_metrics(m)[2:3]
        if idx == 1
            push!(ixt, x)
            push!(iyt, y)
        elseif y >= iyt[end]
            push!(ixt, x)
            push!(iyt, y)
        end
    end
    return ixt, iyt
end

"""
    search_optima!(m::IB, n_iter = 10000)

The original IB algorithm can converge to a local optima. To find the global optima,
search_optima! performs 'n_iter' optimization of input model 'm', randomly re-initializing it at every iteration.
The best optimization is kept and 'm' is modified in-place.
"""
function search_optima!(m::IB, n_iter = 10000)
    if m.algorithm == "DIB"
        println("DIB algorithm is already giving optima.")
        return
    end
    best = deepcopy(m)
    for i in 1:n_iter
        init!(m)
        IB_optimize!(m; report = false)
        if calc_metrics(m)[4] < calc_metrics(best)[4]
            best.qt_x = m.qt_x
            best.qt = m.qt
            best.qy_t = m.qy_t
        end
    end
    m.qt_x = best.qt_x
    m.qt = best.qt
    m.qy_t = best.qy_t
end

"""
    Helper function that puts all similar clusters next to each other for ease of readability.
"""
function group_equivalent!(df)
  for i = 1:size(df,2)-1
    for j = i+1:size(df,2)
      if df[!,i] == df[!,j]
        idxs = collect(1:size(df,2))
        idxs[j], idxs[i+1] = idxs[i+1], idxs[j]
        permutecols!(df, idxs)
      end
    end
  end
end

"""
    print_results(m::IB, disp_thres = 0.1)

Displays the clustering of data for an optimized input model 'm'.

Since the IB algorithm is not deterministic, some probabilities in the clustering matrix 'm.qt_x' are non-zero.
For ease of readability, every probability under 'disp_thres' is displayed as 0, everything else as 1.
The result is a 2D matrix where the rows represent the clusters and the column the initial categories.
The 1 tell to which cluster which category belongs.
If several 1 are present for a single category, it means that the optimization wasn't able to assign
a single cluster to this category for the chosen β value. You might try a diffent β, or use the DIB variant of the algorithm.
"""
function print_results(m::IB, disp_thres = 0.1)
    if ~isnothing(m.x)
        df = DataFrame(model.qt_x .> disp_thres, Symbol.(sort(unique(m.x))))
        group_equivalent!(df)
        display(df)
    else
        @warn "Initial data is probability distribution, categories will be displayed as x1, x2 ..., xn."
        df = DataFrame(model.qt_x .> disp_thres)
        group_equivalent!(df)
        display(df)
    end
end

export IB, search_optima!, brute_optimize!, IB_optimize!, calc_metrics, get_IB_curve, get_y, print_results

end


abstract type Combinatroics{T} <: AbstractArray{T,1} end

"""
    Combination{T} <: Combinatroics{T}

Structure to generated all Combinations of the given moduels with given length
low_level structure that a user don't need to use normaly
mod = Array with modulels or where combination can be generated from
len = lenght of the made constructs
"""
struct Combination{T} <: Combinatroics{T}
    mod::Array{T}
    len::Int
end

# base types of
Base.length(c::Combination) = binomial(length(c.mod),c.len);
Base.size(c::Combination) = (length(c),1);
#Base.eltype(K::Combination{T}) where {T} = Unordered_Construct{N} where N <: eltype(T) ;
Base.eltype(K::Combination{T}) where T = Unordered_Construct{T} ;

"""
    Base.iterate(c::Combination{T} where T, state = [i for i in c.len:-1:1])

Iterater for Combinations, printed in lexicografical order.
"""
function Base.iterate(c::Combination{T} where T, state = [i for i in c.len:-1:1])
    max = length(c.mod)

    if c.len == 1 # special case len = 1 ( needed?)
        if state[1] <= max
            construct = [c.mod[state[1]]]
            state[1] +=1
            return (Unordered_Construct(construct),state)
        else
            return
        end
    end

    # check when done, last constructed is printed
    if state[1] > max
        return nothing
    end

    # make the current construct based on given state
    construct = sum(c.mod[state])

    # update the state
    # i start at the last index, length of the construct
    i = c.len
    #evaluated if the difference between the last and the next is higher than 1 (than their is still a construct that can be made)
    # if not than lower the index with one and recheck if the difference is 1 . i>1 prevents error
    while i > 1  && state[i-1] - state[i] == 1
        i -=1
    end
    # upate the state at position i and reset all higher index to lowest values
    if i < c.len
        state = _restate(state,i)
    else
        state[i] +=1
    end
    return (construct,state)
end;

"""
    _restate(state,i)

Update state at index i and reset all values behind to lowsted allow values.
function used in the Base.iterate(c::Combination{T} where T, state = [i for i in c.len:-1:1])
"""
function _restate(state,i)
    state[i] +=1;
    state[(i+1):end] = collect(length(state)-i:-1:1);
    return state
end

"""
    function getindex(c::Combination,pos::Int)

Get the index of `c` at position `pos`.

Index based on rank Combination.
https://en.wikipedia.org/wiki/Combinatorial_number_system
"""
function getindex(c::Combination,pos::Int)

    # spacial case of lenght is 1
    if c.len == 1
        return Unordered_Construct([c.mod[pos]])
    end
    # if not 1
    com = Array{Int}(undef, c.len, 1)
    for i in c.len:-1:1
        n = i-1
        b_val = binomial(n,i)
        while b_val < pos
            n += 1
            b_val = binomial(n,i)
        end
        pos -= binomial(n-1,i)
        com[c.len-i+1] = n-1
    end
    index =  com .+=1
    return (sum(c.mod[index]))
end;


function liu_fit(xdata, ydata)
    # each term in Liu's time-domain impedance is denoted by m# (i.e. m1, m2, m3)
    m1(L) = 2*L/xdata[end]
    # m2 is complicated because of the sqrt(t). The units don't strictly work out
    # because there's supposed to be a 1/sqrt(Hz) in the denominator. Now, L is to
    # be computed nH and t should be provided as ns. So, scaling everything would
    # result in something like (Rs*1e-6)*sqrt.(t*1e-9)/(pi*L*1e-9). But, we don't
    # want Julia to have to work with big numbers if she doesn't have to. So, we'll
    # combine all of these exponentials into one: 10^-6*10^-4.5*10^9 = 10^-1.5
    m2(t,Rs,L) = 10^-1.5*Rs*sqrt.(t)/(pi*L) # Rs - microOhms, t - ns, L - nH
    m3(t,Rdc,L) = .5*Rdc*t/L # Rdc- ohms, t - ns, L - nH
    m(t,p) = m1(p[1])*(1+m2(t,p[2],p[1])+m3(t,p[3],p[1])) # p[1:3] -> L,Rs,Rdc
    fit = LsqFit.curve_fit(m,xdata,ydata,rand(3))
    errors = LsqFit.estimate_errors(fit)
    L = fit.param[1]
    C = xdata[end]^2/(4*L)
    Rs = fit.param[2]
    Rdc = fit.param[3]
    #@printf("L [nH]: %2.2f +/- %2.2f (%2.2f,%2.2f)\n", L, errors[1],
            #fit.param[1]-errors[1],fit.param[1]+errors[1])
    #@printf("C [μF]: %2.2f\n", C)
    #@printf("Rs [μΩ]: %2.2f +/- %2.2f (%2.2f,%2.2f)\n", Rs, errors[2],
            #fit.param[2]-errors[2],fit.param[2]+errors[2])
    #@printf("Rdc [Ω]: %2.2f +/- %2.2f (%2.2f,%2.2f)\n", Rdc, errors[3],
            #fit.param[3]-errors[3],fit.param[3]+errors[3])
    return Dict(
                "L" => L,
                "C" => C,
                "Rs" => Rs,
                "Rdc" => Rdc,
               )
end


function visible_asteroids_at_coord(grid, coord)
    # Search outwards in "shells" from coord
    remaining = copy(grid)
    for r in 1:maximum(size(grid))
        shell = map(x -> coord + x, setdiff(CartesianIndices((-r:r, -r:r)), CartesianIndices((-(r-1):(r-1), -(r-1):(r-1)))))
        for i in filter(x -> x ∈ CartesianIndices(remaining) && remaining[x], shell)
            # We can see this asteroid. Find the minimum integer direction vector.
            delta = CartesianIndex(tuple((x -> x .÷ gcd(x))(collect((i - coord).I))...))
            # Delete all occluded locations by stepping along the direction vector
            occluded = i + delta
            while occluded ∈ CartesianIndices(grid)
                remaining[occluded] = false
                occluded += delta
            end
        end
    end
    # Remember to remove the asteroid we're sitting on
    remaining[coord] = false
    return remaining
end

# Part 1: read asteroid map as bool matrix and find out which one can detect most asteroids
grid = hcat(map(line -> '#' .== collect(line), open(f -> readlines(f), "day10_input.txt"))...)'
visible_counts = map(x -> grid[x] ? count(visible_asteroids_at_coord(grid, x)) : 0, CartesianIndices(grid))
println("Max detected asteroids: $(maximum(visible_counts))")

# Part 2: get coord of best station, then work out which is the 200th asteroid to be vaporized
# We actually have >200 detected, so we need less than one rotation. Sort detected by angle.
coord = argmax(visible_counts)
detected = visible_asteroids_at_coord(grid, coord)
item = sort(map(x -> (π - atan(reverse((x - coord).I)...), x), filter(x -> detected[x], CartesianIndices(detected))))[200]
println("200th to be vaporized: $item, checksum: $((item[2][2] - 1) * 100 + item[2][1] - 1))")


using LinearAlgebra
using SparseArrays
using CartesianCutCell
using Vofinit

bc = (periodic, periodic, periodic)
st = stencil(bc)
n = (64, 64, 64)

n̅, x̅, x̅ₘ, dx̅ = mesh(bc, st, n)

δ̅₋, δ̅₊ = backwarddiff(bc, n̅), forwarddiff(bc, n̅)
σ̅₋, σ̅₊ = backwardinterp(bc, n̅), forwardinterp(bc, n̅)

circle(x...; x₀=zero.(x), r=one(eltype(x))/2, f=identity) =
    f(norm(r)) - f(norm(x .- x₀))

V̅ = map(tuple.(x̅...), tuple.(dx̅...)) do x₀, h₀
    getcelltype(circle, x₀, h₀)
end

V = restrict(V̅, n̅, n)

# -1 : =0 (interface)
# 0 : >0 (inside)
# 1 : <0 (outside)
cnt = count.((==).((-1, 0, 1)), Ref(V))
cnt[2] / cnt[3] * 48 / π

#V̅, A̅ = integrate(circle, x̅, dx̅)


using SchumakerSpline
using LinearAlgebra

x_row = [0.1,0.3]
x_col = [0.2,0.3,0.5]
ygrid = [1.0 2.0 3.0;
         18.0 -2.0 3.3]
sp = Schumaker2d(x_row,x_col,ygrid; extrapolation = (Linear, Linear))
sp(0.4, 0.3)


using Flux, DiffEqFlux, StochasticDiffEq, Plots

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end
function lotka_volterra_noise(du,u,p,t)
  du[1] = 0.1u[1]
  du[2] = 0.1u[2]
end
prob = SDEProblem(lotka_volterra,lotka_volterra_noise,[1.0,1.0],(0.0,10.0))
p = param([2.2, 1.0, 2.0, 0.4])
function predict_fd_sde()
  diffeq_fd(p,sol->sol[1,:],101,prob,SOSRI(),saveat=0.1)
end
loss_fd_sde() = sum(abs2,x-1 for x in predict_fd_sde())
loss_fd_sde()
Flux.back!(loss_fd_sde())

function predict_rd_sde()
  Array(diffeq_rd(p,prob,SOSRI(),saveat=0.1))
end
loss_rd_sde() = sum(abs2,x-1 for x in predict_rd_sde())
loss_rd_sde()
@test_broken Flux.back!(loss_rd_sde())


# Set random seed:
srand(1234)

#uniform.jl

"""
Create a NxM uniformly distributed matrix using the method rand()*(b-a)-a where a,b are the Unif Parameters.

Arguments
---------
a
    Lower Bound
b
    Upper Bound
N
    Row Dimension
M
    Column Dimension

"""
function uniform(a,b,N,M)
    #Returns a NxM uniformly distributed matrix
    #for random number in range [a,b] we note that it can be generated by: rand() * (b-a) - a
    #where rand() prints real numbers:
    rand(N,M)*(b-a)+a
end

#sigmoid.jl

"""
Return the sigmoid of x:

Arguments
---------
x
    input
h
    hyperparemeter

"""

function sigmoid(x,h)
    return 1/(1+exp(-h*x))
end

#sigmoid_deriv.jl

"""
Return the sigmoid of x:

Arguments
---------
x
    input
h
    hyperparemeter

"""

function sigmoid_deriv(x,h)
    return (1/(1+exp(-h*x)))*(1-1/(1+exp(-h*x)))
end

#Train_Neural_Net.jl

"""
Train the XOR neural net for a given number of epochs:

Arguments
---------
epochs
    number of times the network will foward and back propagate
Layer_1
    weighting for hidden layer
Layer_2
    weighting for output layer
learning_rate
    Learning Rate for Gradient Descent
"""

function Train_Neural_Net_Loop(epochs,Layer_1,Layer_2,learning_rate,node_function,node_deriv)

    #Initialise XOR truth values:
    X=[ 0 0 ; 0 1 ; 1 0 ; 1 1] #When selected X[i,:] we have column vectors so need to transpose for matmult
    Y=[ 0;1;1;0]

    MSE=0

    for i=1:epochs #For the weight matrix wij is the weight for input i going to Neuron j:

        s_1=map(node_function,X*Layer_1) #Applies Sigmoid to X*Layer_1 the weighting matrix
        s_2=map(node_function, s_1*Layer_2) #Applies Sigmoid to a2 * second weighting matrix
        direct_error=s_2-Y
        MSE=0.5*sum(direct_error.*direct_error) #element wise squares direct_error then summs up and divides by length

    #Begin updating:
    delta_outer=-1.0*direct_error.*map(node_deriv,s_1*Layer_2) #this is the delta (which is the raw loss element-wise
                                                        #multiplied by z(3) see explanation in readme

    delta_inner=delta_outer*transpose(Layer_2).*map(node_deriv,X*Layer_1)
    # map(sigmoid_deriv,X*Layer_1) is f'(z2)

    #Update the weights:

    Layer_2 +=learning_rate*transpose(s_1)*delta_outer #Outer layer has been updated.
    Layer_1 += learning_rate*transpose(X)*delta_inner

    print(string(MSE,"\r"))

    end
    return Layer_1, Layer_2, MSE
    end

#Train_Neural_Net.jl

"""
Train the XOR neural net for a given mean square error it needs to reach:

Arguments
---------
MSE_Min
    Error the network trains till
Layer_1
    weighting for hidden layer
Layer_2
    weighting for output layer
learning_rate
    Learning Rate for Gradient Descent
"""

function Train_Neural_Net_MSE(MSE_Min,Layer_1,Layer_2,learning_rate,node_function,node_deriv)

    #Initialise XOR truth values:
    X=[ 0 0 ; 0 1 ; 1 0 ; 1 1] #When selected X[i,:] we have column vectors so need to transpose for matmult
    Y=[ 0;1;1;0]

    MSE=MSE_Min+1

    while MSE>=MSE_Min #For the weight matrix wij is the weight for input i going to Neuron j:

        s_1=map(node_function,X*Layer_1) #Applies Sigmoid to X*Layer_1 the weighting matrix
        s_2=map(node_function, s_1*Layer_2) #Applies Sigmoid to a2 * second weighting matrix
        direct_error=s_2-Y
        MSE=0.5*sum(direct_error.*direct_error) #element wise squares direct_error then summs up and divides by length

    #Begin updating:
    delta_outer=-1.0*direct_error.*map(node_deriv,s_1*Layer_2) #this is the delta (which is the raw loss element-wise
                                                        #multiplied by z(3) see explanation in readme

    delta_inner=delta_outer*transpose(Layer_2).*map(node_deriv,X*Layer_1)
    # map(sigmoid_deriv,X*Layer_1) is f'(z2)

    #Update the weights:

    Layer_2 +=learning_rate*transpose(s_1)*delta_outer #Outer layer has been updated.
    Layer_1 += learning_rate*transpose(X)*delta_inner

    print(string(MSE,"\r"))

    end
    return Layer_1, Layer_2, MSE
    end

#XOR_Net.jl

"""
Output the values of a given XOR_Net

Arguments
---------
M
    Input for the XOR
w1
    Layer 1
w2
    Layer 2

"""

function XOR_Net(M,w1,w2,node_function)
    #Where M is a matrix that
    s_1=map(node_function,M*w1) #Applies Sigmoid to X*Layer_1 the weighting matrix
    s_2=map(node_function,s_1*w2)
    return s_2

end

#-----------------------------Training Example--------------------------------#

# #Initialise Layers of Neuron Weights:

# Layer_1=uniform(0,1,2,2) #(In video above W1)
# Layer_2=uniform(0,1,2,1) #Column vector for the outer layer (in video above W2)

# function hyper_curry(h)
#     return (x->sigmoid(x,h))
# end

# function hyper_curry_deriv(h)
#     return (x->sigmoid_deriv(x,h))
# end

# node_function=hyper_curry(1)
# node_deriv=hyper_curry_deriv(1)

# w1, w2, MSE1 =Train_Neural_Net_Loop(1000000,Layer_1,Layer_2,0.01,node_function,node_deriv)

# print(XOR_Net([0 1],w1,w2,node_function))


function mre(predictions, actualdata)
	#Julia: simple version of mre( without checking the third parameter)
	#%if the third parameter is provided as true, then we will first convert
	#%everything into integers!
	#if nargin == 3 && varargin{1}
	#    predictions = ceil(predictions);
	#    actualdata = ceil(actualdata);
	#end;

	#% mean absolute error

	out = zeros(1,size(predictions,2));
	for i=1:size(predictions,2)
		if i > size(actualdata,2)
			continue
		end
		act = actualdata[:,i];
		pre = predictions[:,i];
		nonz = find(act);

		#%My own definition
		out[i] = mean(abs(pre[nonz]-act[nonz])./abs(act[nonz]));
		#%Weka!!
		#%avgAct = mean(act);
		#%out(i) = sum(abs(pre-act)) / sum(abs(avgAct-act));
	end
	out
end


export check_dΠψ, _dψ1dθρ, _ρ, _ψ2, _ψ1

@inline _ρ(θρ::Real) = logistic(θρ)
@inline _dρdθρ(θρ::Real) = (z = logistic(θρ); z*(1-z) )

@inline _ρ2(θρ::Real) = _ρ(θρ)^2
_dρdσ = _dρdθρ # alias

@inline _ψ1(u::Real,v::Real,ρ::Real) = ρ*u + sqrt(1-ρ^2)*v
@inline _ψ2(u::Real,v::Real, ρ::Real) = _ψ2(u,v)
@inline _ψ2(u::Real,v::Real) = u

@inline _dψ1dρ(u::Real,v::Real,ρ::Real) = u - ρ/sqrt(1-ρ^2)*v

@inline _dψ1dθρ(u::Real, v::Real, ρ::Real, θρ::Real) = _dψ1dρ(u,v,ρ)*_dρdθρ(θρ::Real)

@inline _z(x2::Real, x1::Real, Δ::Real, ρ::Real) = (x2 - ρ*x1 + Δ)/sqrt(1-ρ^2)

# derivatives
@inline _dzdρ(x2::Real, x1::Real, ρ::Real, z::Real) = -x1/sqrt(1-ρ^2) + ρ*z/(1-ρ^2)
@inline _dπdρ(x2::Real, x1::Real, Δ::Real, ρ::Real) = (z = _z(x2,x1,Δ,ρ);  normpdf(z) * _dzdρ(x2,x1,ρ,z))

# finite difference versions
@inline _ρ(x::Real, h::Real) = _ρ(x+h)
@inline _z(x2::Real, x1::Real, Δ::Real, ρ::Real, h::Real) = _z(x2, x1+h, Δ, ρ)

# ------------------------------ matrix updates -------------------------

function _βΠψdθρ!(P::AbstractMatrix, y::StepRangeLen, θρ::Real, β::Real)
    n = length(y)
    (n,n) == size(P) || throw(DimensionMismatch())
    ρ = _ρ(θρ)
    dρdθρ = _dρdθρ(θρ)
    Δ = 0.5 * step(y)

    @inbounds for (j,yj) in enumerate(y)
        if j == 1
            @. P[ :,j] = β * _dπdρ(yj,y,Δ,ρ) * dρdθρ
        elseif j == n
            @. P[ :,j] = -β * _dπdρ(yj,y,-Δ,ρ) * dρdθρ
        else
            @. P[ :,j] = β * ( _dπdρ(yj,y,Δ,ρ) - _dπdρ(yj,y,-Δ,ρ) ) * dρdθρ
        end
    end
end

function _βΠψ!(P::AbstractMatrix, y1::StepRangeLen, y2::StepRangeLen, θp::Real, β::Real)
    n = length(y2)
    n == length(y1) || throw(DimensionMismatch())
    (n,n) == size(P) || throw(DimensionMismatch())
    ρ = _ρ(θp)
    Δ = 0.5 * step(y2)

    @inbounds for (j,yj) in enumerate(y2)
        if j == 1
            @. P[ :,j] = β * normcdf(_z(yj, y1, Δ, ρ))
        elseif j == n
            @. P[ :,j] = β * normccdf(_z(yj, y1, -Δ, ρ))
        else
            @. P[ :,j] = β * ( normcdf(_z(yj, y1, Δ, ρ)) - normcdf(_z(yj, y1, -Δ, ρ) ))
        end
    end
end

_βΠψ!(P::AbstractMatrix, y1::StepRangeLen, θp::Real, β::Real) = _βΠψ!(P, y1, y1, θp, β)

function _βΠψ(y1::StepRangeLen{T}, θp::Real, β::Real) where {T}
    P = Matrix{T}(undef, length(y1), length(y1))
    _βΠψ!(P, y1, y1, θp, β)
    return P
end

# ------------------------------ derivative check -------------------------

function check_dΠψ(θρ::Real, ψspace::StepRangeLen)

    Δ = 0.5 * step(ψspace)

    for yj in ψspace
        for y in ψspace
            fdσ = Calculus.derivative((sig) -> normcdf(_z(yj, y, Δ, _ρ(sig))), θρ)
            dσ  = _dρdθρ(θρ) * _dπdρ(yj, y, Δ, _ρ(θρ))
            abs(fdσ - dσ ) < 1e-7 || throw(error("bad σ grad at σ = $σ, ψ2 = $yj, ψ1 = $y"))
        end
    end
    return true
end


using NGWP, Plots, LightGraphs, Random
using PrettyTables

nsim = 500

function generate_ROT_TSD_ratio(nsim, 𝚽, ∇𝚽, 𝛌, Q; edge_length = 1)
    Random.seed!(1234)
    ρ = zeros(nsim)
    for i = 1:nsim
        p = rand(N); p ./= norm(p, 1)
        q = rand(N); q ./= norm(q, 1)
        W1 = ROT_Distance(p, q, Q; edge_length = edge_length)
        K = K_functional(p, q, 𝚽, ∇𝚽, 𝛌; length = edge_length)[1]
        ρ[i] = K / W1
    end
    return ρ
end

function display_basic_stats(ρs)
    header = ["min" "max" "mean" "std"]
    basic_stats = zeros(0, 4)
    for ρ in ρs
        basic_stats = vcat(
            basic_stats,
            round.([minimum(ρ) maximum(ρ) mean(ρ) std(ρ)]; digits = 4)
            )
    end
    pretty_table(basic_stats, header)
end

function ROT_TSD_ratio_histogram(ρ)
    plt = histogram(ρ, grid = false, legend = false, c = :teal,
              xlims = [minimum(ρ) - std(ρ), maximum(ρ) + std(ρ)],
              xlab = "ρ")
    return plt
end


struct Cache{TF<:AbstractFloat}
    Ax::Vector{TF}
    resid::Vector{TF}
    dresid::Vector{TF}
    g::Vector{TF}
    dx::Vector{TF}
end

function Cache(TF::Type, M::Integer, N::Integer)
    Ax = Vector{TF}(undef, M)
    resid = Vector{TF}(undef, M)
    dresid = Vector{TF}(undef, M)
    g = Vector{TF}(undef, N)
    dx = Vector{TF}(undef, N)
    return Cache{TF}(Ax, resid, dresid, g, dx)
end

"""
    Cache(A::AbstractMatrix)

Allocate arrays for the solver based on the size and element type of `A`.
"""
Cache(A::AbstractMatrix{TF}) where TF<:AbstractFloat = Cache(TF, size(A)...)

function updaterss!(A, b, x, ca)
    mul!(ca.Ax, A, x)
    @inbounds for i in eachindex(ca.resid)
        ca.resid[i] = ca.Ax[i] - b[i]
    end
    return sum(abs2, ca.resid)
end

function direction!(A, x, λ, ca)
    copyto!(ca.g, x)
    # Compute the half-gradient g = A'(Ax-b) + λx
    mul!(ca.g, A', ca.resid, true, λ)
    gmin, ix = findmin(ca.g)
    mul!(ca.dx, -one(eltype(A)), x)
    ca.dx[ix] += one(eltype(x))
    return ix
end

function stepsize!(ix, A, λ, ca)
    @inbounds for i in eachindex(ca.dresid)
        ca.dresid[i] = A[i,ix] - ca.Ax[i]
    end
    ss = -dot(ca.g, ca.dx) / (sum(abs2,ca.dresid) + λ*sum(abs2,ca.dx))
    TF = eltype(A)
    return min(one(TF), max(zero(TF), ss))
end

function updatex!(ss, x, ca)
    @inbounds for i in eachindex(x)
        x[i] += ss*ca.dx[i]
    end
end

"""
    convexfit(A::Matrix, b::Vector, λ::Real=0; kwargs...)

Fit vector `b` with a convex combination of the columns in matrix `A`
and an optional regularization parameter `λ`.

The associated constrained optimization problem can be expressed as follows:
```math
\\min_x \\|Ax - b\\|_2^2 + λ\\|x\\|_2^2 \\\\
st. x_i ≥ 0 \\forall i, \\, \\sum_i x_i = 1
```
where `x` is the vector of coefficients to be solved.

# Keywords
- `x0::AbstractVector{<:Real}=fill(1/size(A,2),1/size(A,2))`: initial value of `x`.
- `ftol::Real=1e-6`: convergence criterion based on the change in the objective function.
- `xtol::Real=1e-6`: convergence criterion based on the inf-norm change in `x`.
- `maxiter::Integer=1000`: maximum number of iterations.
- `store_trace::Bool=false`: save the solver state in each iteration.
- `show_trace::Bool=false`: print the solver state in each iteration to `stdout`.
- `show_thread::Bool=false`: show the thread ID when printing solver states.
- `cache::Cache=Cache(A)`: `Cache` that stores intermediate results.
"""
function convexfit(A::Matrix{TF}, b::Vector{TF}, λ::Real=zero(TF);
        x0::AbstractVector{<:Real}=fill(convert(TF,1/size(A,2)), size(A,2)),
        ftol::Real=convert(TF, 1e-6),
        xtol::Real=convert(TF, 1e-6),
        maxiter::Integer=1000,
        store_trace::Bool=false,
        show_trace::Bool=false,
        show_thread::Bool=false,
        cache::Cache=Cache(A)) where TF<:AbstractFloat

    sx = sum(x0)
    if !(sx ≈ 1)
        @warn "sum of elements in x0 ($sx) is normalized to 1"
        x0 = x0./sx
    end
    λ = convert(TF, λ)
    λ < 0 && throw(ArgumentError("regularization parameter cannot be negative"))

    f_converged = false
    x_converged = false
    rss = -one(TF)
    xnorm = -one(TF)
    flast = -one(TF)
    f = -one(TF)
    x = Vector{TF}(undef, length(x0))
    copyto!(x, x0)
    dfnorm = convert(TF, NaN)
    dxnorm = convert(TF, NaN)
    iter = 0
    if store_trace || show_trace
        tr = SolverTrace(TF)
    else
        tr = nothing
    end

    for _ in 1:maxiter
        flast = f
        rss = updaterss!(A, b, x, cache)
        xnorm = sum(abs2, x)
        f = rss + λ * xnorm
        dfnorm = abs(f-flast)
        if dfnorm < ftol
            f_converged = true
            break
        elseif x_converged
            break
        elseif iszero(f)
            f_converged = true
            x_converged = true
            dfnorm = zero(TF)
            dxnorm = zero(TF)
            break
        end
        ix = direction!(A, x, λ, cache)
        if iszero(cache.dx[ix])
            f_converged = true
            x_converged = true
            dfnorm = zero(TF)
            dxnorm = zero(TF)
            break
        end
        ss = stepsize!(ix, A, λ, cache)
        dxnorm = ss*maximum(abs, cache.dx)
        if dxnorm < xtol
            x_converged = true
            # Do not immediately break to allow checking f_converged with the new x
        end
        updatex!(ss, x, cache)
        # Ensure that iter always represents the number of steps from x0r
        iter += 1
        # xnorm should not be updated until the next iteration
        update!(tr, iter, f, rss, xnorm, ix, ss, store_trace, show_trace, show_thread)
    end
    return SolverResult{TF}(x0, x, copy(cache.Ax), iter, f, rss, λ, xnorm,
        dfnorm, ftol, f_converged, dxnorm, xtol, x_converged, tr)
end

function convexfit(A::Matrix{TF}, B::Matrix{TF}, λ=0.0;
        multithreads::Bool=false, show_thread::Bool=multithreads,
        kwargs...) where TF<:AbstractFloat
    N = size(B, 2)
    rs = Vector{Any}(undef, N)
    if multithreads
        haskey(kwargs, :cache) && @warn "keyword cache is ignored with multithreads=true"
        Threads.@threads for i in axes(B, 2)
            # Create new Cache for each call no matter whether cache is specified
            rs[i] = convexfit(A, view(B,:,i), λ;
                kwargs..., cache=Cache(A), show_thread=show_thread)
        end
    else
        ca = get(kwargs, :cache, nothing)
        ca === nothing && (ca = Cache(A))
        for i in axes(B, 2)
            get(kwargs, :show_trace, false) && println("target $i of $N:")
            rs[i] = convexfit(A, view(B,:,i), λ;
                kwargs..., cache=ca, show_thread=show_thread)
        end
    end
    return rs
end

"""
    convexfit(A::AbstractMatrix, B::AbstractVecOrMat, λ=0; multithreads=false, kwargs...)

A wrapper method of `convexfit` that accepts multiple `b`s in a matrix `B` and
fits each column in `B` with a convex combination of the columns in matrix `A`.
The actual problems to be solved depend on the type of `λ`.
Results are collected in an array in the order of the columns of `B`.

Problems across different columns can be solved simultaneously
by setting `multithreads=true`.
See also [`convexfit(::Matrix,::Vector,::Real)`](@ref).

!!! compat "Julia 1.2"
    Setting `multithreads=true` requires at least Julia 1.2.

"""
function convexfit(A::AbstractMatrix, b::AbstractVecOrMat, λ=0.0; kwargs...)
    TF = promote_type(eltype(A), eltype(b))
    TF <: Integer && (TF = Float64)
    A = convert(Matrix{TF}, A)
    b = convert(Array{TF}, b)
    # Convert x0 at this stage to avoid making multiple copies of x0 later
    x0 = get(kwargs, :x0, nothing)
    if x0 !== nothing
        x0 = convert(Vector{TF}, x0)
        return convexfit(A, b, λ; kwargs..., x0=x0)
    else
        return convexfit(A, b, λ; kwargs...)
    end
end

function convexfit(A::Matrix{TF}, λ=0.0;
        loo::Bool=true, multithreads::Bool=false, show_thread::Bool=multithreads,
        kwargs...) where TF<:AbstractFloat
    if loo
        N = size(A, 2)
        N > 1 || throw(ArgumentError("matrix A must contain at least two columns if loo=true"))
        rs = Vector{Any}(undef, N)
        if multithreads
            haskey(kwargs, :cache) && @warn "keyword cache is ignored with multithreads=true"
            Threads.@threads for i in axes(A, 2)
                Aloo = view(A,:,axes(A,2).!=i)
                # Create new Cache for each call no matter whether cache is specified
                rs[i] = convexfit(Aloo, view(A,:,i), λ;
                    kwargs..., cache=Cache(Aloo), show_thread=show_thread)
            end
        else
            ca = get(kwargs, :cache, nothing)
            ca === nothing && (ca = Cache(view(A,:,2:N)))
            for i in axes(A, 2)
                get(kwargs, :show_trace, false) && println("target $i of $N:")
                rs[i] = convexfit(view(A,:,axes(A,2).!=i), view(A,:,i), λ;
                    kwargs..., cache=ca, show_thread=show_thread)
            end
        end
        return rs
    else
        ca = get(kwargs, :cache, nothing)
        ca === nothing && (ca = Cache(A))
        return convexfit(A, A, λ; kwargs..., cache=ca)
    end
end

"""
    convexfit(A::AbstractMatrix, λ=0; loo=true, multithreads=false, kwargs...)

A wrapper method of `convexfit` for fitting
each column in matrix `A` with a convex combination of the other columns in matrix `A`.
The actual problems to be solved depend on the type of `λ`.
Results are collected in an array in the order of the columns of `B`.

By default, with `loo=true` the fitted column is not involved in the convex combination.
This can be altered by setting `loo=false`.
Problems across different columns can be solved simultaneously
by setting `multithreads=true`.
See also [`convexfit(::Matrix,::Vector,::Real)`](@ref).

!!! compat "Julia 1.2"
    Setting `multithreads=true` requires at least Julia 1.2.

"""
function convexfit(A::AbstractMatrix, λ=0.0; kwargs...)
    TF = eltype(A)
    TF <: Integer && (TF = Float64)
    A = convert(Matrix{TF}, A)
    # Convert x0 at this stage to avoid making multiple copies of x0 later
    x0 = get(kwargs, :x0, nothing)
    if x0 !== nothing
        x0 = convert(Vector{TF}, x0)
        return convexfit(A, λ; kwargs..., x0=x0)
    else
        return convexfit(A, λ; kwargs...)
    end
end


using LinearAlgebra
using DataStructures

struct P{T<:Real}
    x::T
    y::T
end

function Base.:+(A::P,B::P)
 return P(A.x+B.x,A.y+B.y)
end
function Base.:-(A::P,B::P)
 return P(A.x-B.x,A.y-B.y)
end

function LinearAlgebra.norm(P::P)
    sqrt(P.x^2+P.y^2)
end

struct Polygon{T}
    name::Symbol
    edges::T
end

function Base.show(io::IO,  P::Polygon)
    print(io, P.name)
end
universe = [Polygon(:box,(P(0.0,0.0),P(0.0,5.0),P(10.0,5.0),P(10.0,0.0))),
            Polygon(:square,(P(0.0,0.0),P(0.0,10.0),P(10.0,10.0),P(10.0,0.0))),
            Polygon(:isosceles,(P(0.0,0.0),P(5.0,5.0),P(10.0,0.0))),
            Polygon(:right_triangle,(P(0.0,0.0),P(4.0,3.0),P(4.0,0.0))),
            Polygon(:trapezoid,(P(0.0,0.0),P(5.0,5.0),P(20.0,5.0),P(25.0,0.0))),
            Polygon(:parallelogram,(P(0.0, 0.0), P(5.0, 5.0), P(15.0, 5.0), P(10.0, 0.0))),
            Polygon(:rhombus,(P(0.0, 0.0), P(4., 3.), P(9., 3.), P(5., 0.))),
            Polygon(:multi,(P(0., 0.), P(0., 10.), P(4., 15.), P(10., 15.), P(15., 10.), P(15., 4.), P(10., 0.) )),
            Polygon(:line,(P(0.0,0.0),P(10.0,0.0))),
            Polygon(:dot,(P(0.0,0.0),))]

"""
    many_sides(P::Polygon)

Returns T if P has more than 6 sides.
"""
function many_sides(P::Polygon)
    return length(P.edges) > 6;
end

"""
    nonzero_area(P::Polygon)

Returns T if P has nonzero area.
"""
function nonzero_area(P::Polygon)
    return area(P)!=0
end # function

"""
    area(P::Polygon)

    Computes the area enclosed by polygon P.
    Taken from https://rosettacode.org/wiki/Shoelace_formula_for_polygonal_area#Kotlin
"""
function area(P::Polygon)
    n = length(P.edges)
    a = 0.0
    for i in range(1,stop=n-1)
        a+= P.edges[i].x * P.edges[i+1].y - P.edges[i+1].x * P.edges[i].y
    end
    return abs(a + P.edges[n].x * P.edges[1].y - P.edges[1].x * P.edges[n].y) / 2.0
end # function

"""
    equal_sides(P::Polygon)

EQUAL-SIDES is true if all sides of P have equal length
"""
function equal_sides(P::Polygon)
    n = length(P.edges)
    len = norm(P.edges[1]-P.edges[n])
    for i=1:n-1
        if norm(P.edges[i]-P.edges[i+1]) != len
            return false
        end
    end
    return true
end # function

predicates = [many_sides,nonzero_area,equal_sides]

reporting = true

agenda = PriorityQueue{Any,Float64}(Base.Order.Reverse);

Singleton = World(universe,predicates,reporting,agenda)

Object = Concept(Singleton, "Abstract Object", (x)->true, 50, [], [], missing, [])

Singleton.agenda[(find_examples_of, Object)] = 50.0;
Singleton.agenda[(make_specialization, Object)] = 25.0;


immutable Beta <: ContinuousUnivariateDistribution
    alpha::Float64
    beta::Float64
    function Beta(a::Real, b::Real)
        if a > 0.0 && b > 0.0
            new(float64(a), float64(b))
        else
            error("Both alpha and beta must be positive")
        end
    end
end

Beta(a::Real) = Beta(a, a) # symmetric in [0, 1]
Beta() = Beta(1.0) # uniform

@_jl_dist_2p Beta beta

function entropy(d::Beta)
    o = lbeta(d.alpha, d.beta)
    o -= (d.alpha - 1.0) * digamma(d.alpha)
    o -= (d.beta - 1.0) * digamma(d.beta)
    o += (d.alpha + d.beta - 2.0) * digamma(d.alpha + d.beta)
    return o
end

insupport(d::Beta, x::Number) = isreal(x) && 0.0 < x < 1.0

function kurtosis(d::Beta)
    a, b = d.alpha, d.beta
    d = 6.0 * ((a - b)^2 * (a + b + 1.0) - a * b * (a + b + 2.0))
    n = a * b * (a + b + 2.0) * (a + b + 3.0)
    return d / n
end

mean(d::Beta) = d.alpha / (d.alpha + d.beta)

function modes(d::Beta)
    if d.alpha > 1.0 && d.beta > 1.0
        return [(d.alpha - 1.0) / (d.alpha + d.beta - 2.0)]
    else
        error("Beta distribution with a <= 1 || b <= 1 has no modes")
    end
end

function rand(d::Beta)
    u = rand(Gamma(d.alpha))
    return u / (u + rand(Gamma(d.beta)))
end

# TODO: Don't create temporaries here
function rand(d::Beta, dims::Dims)
    u = rand(Gamma(d.alpha), dims)
    return u ./ (u + rand(Gamma(d.beta), dims))
end

function rand!(d::Beta, A::Array{Float64})
    for i in 1:length(A)
        A[i] = rand(d)
    end
    return A
end

function skewness(d::Beta)
    d = 2.0 * (d.beta - d.alpha) * sqrt(d.alpha + d.beta + 1.0)
    n = (d.alpha + d.beta + 2.0) * sqrt(d.alpha * d.beta)
    return d / n
end

function var(d::Beta)
    ab = d.alpha + d.beta
    return d.alpha * d.beta / (ab * ab * (ab + 1.0))
end


using EmpiricalRiskMinimization
using CSV, Dataframes

# data from  https://web.stanford.edu/~hastie/Papers/LARS/diabetes.data
# There are 10 explanatory variables, including age (age),
# sex (sex), body mass index (bmi) and mean arterial blood pressure (map)
# of 442 patients as well as six blood serum measurements (tc, ldl, hdl,
# tch, ltg and glu). The response (y) is a continuous measure of disease
# progression one year after the baseline measurements.

DATA_URL = "https://web.stanford.edu/~hastie/Papers/LARS/diabetes.data"

function load_data()
    data = download(DATA_URL)
    D, header = CSV.read(data; delim="\t")
    V = floatarray(D[:,11])
    Vnames = Any[header[11]]
    U = floatarray(D[:,1:10])
    Unames = header[1:10]
    return U, V, Unames, Vnames
end

function individual_plots(U, V, colnames, theta, loss, yhat)
    cols = [1 2 3 4 5; 6 7 8 9 10]
    nrows, ncols = size(cols)

    # just plot data
    function f1(ax, i, j)
        col = cols[i,j]
        plot(ax, U[:,col], V, markersize=6)
             #xlabel=string(colnames[col])
    end
    arrayplot(f1, nrows, ncols, name="alldata.pdf")

    # plot regression lines
    function f2(ax, i, j)
        col = cols[i,j]
        plot(ax, U[:,col], V, xlabel=string(colnames[col]), markersize=6)
        regline(ax, theta[col])
        annotate(ax, 0.95, 0.01, @sprintf("MSE %.0f\n", loss[col]), horizontalalignment="right")
    end
    arrayplot(f2, nrows, ncols, name="alldata_reg.pdf")

    # plot performance of BMI, X column 3
    col = 3
    perfplot(yhat[col], V, xmin=0, xmax=350, name="bmi_perf.pdf")
    println("BMI loss = ", loss[col])
    showmatrix(theta[col], "BMI theta")

    # plot just bmi regression
    ax = plot(U[:,col], V, xlabel=string(colnames[col]))
    regline(ax, theta[col])
    plt.savefig(plotname("bmi_reg.pdf"))

end

function pairwise_table()
    srand(1)
    U, V, Unames, Vnames = loaddata()
    M = Model(U, V, Unames, Vnames, SquareLoss(), L2Reg();
              embedall=true, stand=false)

    d = size(U, 2)
    # all regressors, validated
    tests = Any[]
    train(M; trainfrac=0.8)
    push!(tests, ["all", trainloss(M), testloss(M)])

    # no regressors
    train(M, features=[1])
    push!(tests, ["NONE", trainloss(M), testloss(M)])

    # one regressor
    for i=1:d
        train(M, features=[1, i+1])
        push!(tests, [Unames[i], trainloss(M), testloss(M)])
    end

    # losses for pairs of regressors
    for i=2:d
        for j=1:i-1
            train(M, features=[1, i+1, j+1])
            push!(tests, [ [Unames[i], Unames[j]], trainloss(M), testloss(M)])
        end
    end
    tests = sort(tests, by = x->x[3])
    for a in tests
        if isa(a[1], Array)
            print(a[1][1], " and ", a[1][2])
        else
            print(a[1])
        end
        @printf("  &  %.4g   &  %.4g \\\\\n", a[2], a[3])
    end
end

function full_regression()
    U, V, Unames, Vnames = loaddata()
    M = Model(U, V, Unames, Vnames, SquareLoss(), L2Reg();
              embedall=true, stand=false)

    train(M, trainfrac=1)
    fulltheta = thetaopt(M)
    fullloss = trainloss(M)
    fullyhat = predict_y_from_train(M)

    ###################################################
    # plotting and display
    @printf("Loss based on all data = %.0f\n", fullloss)
    perfplot(fullyhat, V,  xmin=0, xmax=350, name="all_perf.pdf")
    showmatrix(fulltheta, "all theta")

end

function mean_predictor()
    U, V, Unames, Vnames = loaddata()

    # constant predictor
    n = length(V)
    Vmean = sum(V)/n
    meanloss = (1/n)*(norm(V-fill(Vmean,n))^2)

    println("ymean = ", Vmean)
    @printf("Loss of constant predictor = %.0f\n", meanloss)

end

function individual_regressions()
    U, V, Unames, Vnames = loaddata()
    M = Model(U, V, Unames, Vnames, SquareLoss(), L2Reg();
              embedall=true, stand=false)

    d = size(U,2)
    # compute all single variable regressions
    theta = Array{Any}(d)
    losses = zeros(d)
    yhat = Array{Any}(d)
    for i=1:10
        setfeatures(M, [1,i+1])
        train(M; trainfrac=1)
        theta[i] = thetaopt(M)
        losses[i] = trainloss(M)
        yhat[i] = predict_y_from_train(M)
    end
    individual_plots(U, V, Unames, theta, losses, yhat)
end

function repeated_train()
    U, V, Unames, Vnames = loaddata()
    M = Model(U, V, Unames, Vnames, SquareLoss(), L2Reg();
              embedall=true, stand=false)

    features = [3, 9]  # BMI and S5 in colnames
    features = features + 1 # shift to account for constant feature
    unshift!(features, 1)   # include constant feature in list
    setfeatures(M, features)

    trainlosses = Float64[]
    testlosses = Float64[]
    for i=1:1000
        train(M; resplit=true, trainfrac=0.8)
        push!(trainlosses, trainloss(M))
        push!(testlosses, testloss(M))
    end
    plot(testlosses; bins=20, plottype=:hist, xlabel="test loss",
         ylabel="no. of outcomes", name="repeated_test_loss.pdf")
    println("mean test loss = ", mean(testlosses))

    plot(trainlosses; bins=20, plottype=:hist, ylabel="training loss",
         xlim=[2000,5000])
    println("mean train loss = ", mean(trainlosses))

end

function main()
    close("all")

    mean_predictor()
    individual_regressions()
    full_regression()
    pairwise_table()

    srand(1)
    repeated_train()
end


using FlightMechanicsUtils
using Test

@testset "Rigid body vel and accel fields" begin
    # TEST rigid_body_velocity
    # Test without angular velocity
    vel_P = [10.0, 0.0, 0.0]
    ω = [0.0, 0.0, 0.0]
    r_PQ = [5.0, 5.0, 5.0]

    exp_vel_Q = [10.0, 0.0, 0.0]
    vel_Q = rigid_body_velocity(vel_P, ω, r_PQ)
    @test isapprox(vel_Q, exp_vel_Q)

    # Only angular velocity in Z axis
    vel_P = [0.0, 0.0, 0.0]
    ω = [0.0, 0.0, 1.0]
    r_PQ = [5.0, 5.0, 5.0]

    exp_vel_Q = [-sqrt(50) * sin(pi/4), sqrt(50) * cos(pi/4), 0.0]
    vel_Q = rigid_body_velocity(vel_P, ω, r_PQ)
    @test isapprox(vel_Q, exp_vel_Q)

    # X linear velocity and Z angular velocity
    vel_P = [10.0, 0.0, 0.0]
    ω = [0.0, 0.0, 1.0]
    r_PQ = [5.0, 5.0, 5.0]

    exp_vel_Q = [-sqrt(50) * sin(pi/4) + 10.0, sqrt(50) * cos(pi/4), 0.0]
    vel_Q = rigid_body_velocity(vel_P, ω, r_PQ)
    @test isapprox(vel_Q, exp_vel_Q)

    # TEST rigid_body_acceleration
    # Only X linear acceleration
    acc_P = [10.0, 0.0, 0.0]
    ω = [0.0, 0.0, 0.0]
    ω_dot = [0.0, 0.0, 0.0]
    r_PQ = [5.0, 5.0, 5.0]

    exp_accel_Q = [10.0, 0.0, 0.0]
    accel_Q = rigid_body_acceleration(acc_P, ω, ω_dot, r_PQ)
    @test isapprox(accel_Q, exp_accel_Q)

    # Only Z angular acceleration
    acc_P = [0.0, 0.0, 0.0]
    ω = [0.0, 0.0, 0.0]
    ω_dot = [0.0, 0.0, 1.0]
    r_PQ = [5.0, 5.0, 5.0]

    exp_accel_Q = [-sqrt(50) * sin(pi/4), sqrt(50) * cos(pi/4), 0.0]
    accel_Q = rigid_body_acceleration(acc_P, ω, ω_dot, r_PQ)
    @test isapprox(accel_Q, exp_accel_Q)

    # Only Z angular velocity
    acc_P = [0.0, 0.0, 0.0]
    ω = [0.0, 0.0, 2.0]
    ω_dot = [0.0, 0.0, 0.0]
    r_PQ = [5.0, 5.0, 5.0]

    exp_accel_Q = [-4.0 * sqrt(50) * sin(pi/4), -4.0 * sqrt(50) * cos(pi/4), 0.0]
    accel_Q = rigid_body_acceleration(acc_P, ω, ω_dot, r_PQ)
    @test isapprox(accel_Q, exp_accel_Q)
end

@testset "Angular kinematic equations" begin
    # Null Euler angles
    pqr = [0.5, 0.3, 0.7]
    euler_angles_rates = pqr_2_ψθϕ_dot(pqr..., 0, 0)
    @test isapprox(euler_angles_rates, pqr[end:-1:1])

    ψθϕ_dot = [0.5, 0.3, 0.7]
    pqr_ = ψθϕ_dot_2_pqr(ψθϕ_dot..., 0, 0)
    @test isapprox(pqr_, ψθϕ_dot[end:-1:1])

    # Reciprocal relationships
    θ = 0.15  # rad
    ϕ = 0.6  # rad
    pqr = [0.5, 0.3, 0.7]
    euler_angles_rates = pqr_2_ψθϕ_dot(pqr..., θ, ϕ)
    pqr_ = ψθϕ_dot_2_pqr(euler_angles_rates..., θ, ϕ)
    @test isapprox(pqr, pqr_)

    # TODO: test body_angular_velocity_to_quaternion_rates
end

@testset "Aero <--> Body velocity" begin
    uvw = [30, 3, 5]
    tas, α, β = uvw_to_tasαβ(30, 3, 5)
    u, v, w = wind2body(tas, 0, 0, α, β)
    @test isapprox([u, v, w], uvw)

    uvw = [100, 0, 0]  # m/s
    rv = uvw_to_tasαβ(uvw...)
    @test isapprox(rv, [100, 0, 0])

    uvw = 10, 0, 10  # m/s
    rv = uvw_to_tasαβ(uvw...)
    @test isapprox(rv, [sqrt(200), atan(1), 0])

    uvw = [10, 10, 0]  # m/s
    rv = uvw_to_tasαβ(uvw...)
    @test isapprox(rv, [sqrt(200), 0, asin(10/sqrt(200))])

    uvw_dot = [5, 1, 10]
    tasαβ_dot = uvw_dot_to_tasαβ_dot(uvw..., uvw_dot...)
    uvw_dot_ = tasαβ_dot_to_uvw_dot(uvw_to_tasαβ(uvw...)..., tasαβ_dot...)
    @test isapprox(uvw_dot_, uvw_dot)
end

@testset "Rate of climb constrain" begin
    # wings-level, no sideslip: θ = α + γ
    θ = rate_of_climb_constrain_no_wind(0.1, 0.05, 0, 0)
    @test isapprox(θ, 0.1 + 0.05)

    # Check case:  Stevens, B. L., Lewis, F. L., (1992).
    # Aircraft control and simulation: dynamics, controls design, and autonomous systems.
    # John Wiley & Sons. (Section 3.6, figure 3.6-2, page 192)
    θ =rate_of_climb_constrain_no_wind(0.0, deg2rad(13.7), deg2rad(0.0292), deg2rad(78.3))
    # Take into account that this check case is a full trim not just the kinematic realtionship
    @test isapprox(rad2deg(θ), 2.87, rtol=0.01)
end


"""
    null1(N::BinaryNetwork)

Given a matrix `A`, `null1(A)` returns a matrix with the same dimensions,
where every interaction happens with a probability equal to the connectance of
`A`.
"""
function null1(N::BinaryNetwork)
    return linearfilter(N; α=[0.0, 0.0, 0.0, 1.0])
end

"""
    null3out(N::BinaryNetwork)

Given a matrix `A`, `null3out(A)` returns a matrix with the same dimensions,
where every interaction happens with a probability equal to the out-degree
(number of successors) of each species, divided by the total number of possible
successors.
"""
function null3out(N::BinaryNetwork)
  return linearfilter(N; α=[0.0, 1.0, 0.0, 0.0])
end

"""
    null3in(N::BinaryNetwork)

Given a matrix `A`, `null3in(A)` returns a matrix with the same dimensions,
where every interaction happens with a probability equal to the in-degree
(number of predecessors) of each species, divided by the total number of
possible predecessors.
"""
function null3in(N::BinaryNetwork)
  return linearfilter(N; α=[0.0, 0.0, 1.0, 0.0])
end

"""
    null2(N::BinaryNetwork)

Given a matrix `A`, `null2(A)` returns a matrix with the same dimensions, where
every interaction happens with a probability equal to the degree of each
species.
"""
function null2(N::BinaryNetwork)
  return linearfilter(N; α=[0.0, 0.5, 0.5, 0.0])
end


# The "pound" symbol or "hash" symbol or "octothorp?" is the comment character
# in a julia file. So all of these commands are ignored!

##
grades = [
  30
  40
  30
  45
  50
  20
  35
  40
  50
  49
  55
]
avggrade = sum(grades)/length(grades)
##
println("The average grade is ", avggrade)

##
# Let's compute the standard deviation too!
# The formula for the standard deviation is
# sqrt(mean((x_i - mean(x)))^2)
# Julia has a built in stdev function, we'll use that in a second.

cgrades = grades .- avggrade
stdgrades = sqrt(sum(cgrades.^2)/length(cgrades))

##
using Statistics
##
mean(grades)
##
std(grades; corrected=false)


function cardinal_rotations()
    six_dirs = [let o=zeros(3);o[i] = sign*1.0;o end for sign in [-1,1] for i in 1:3]
    cardinal_rots = []
    for dir in six_dirs
        for dir2 in six_dirs
            if dot(dir, dir2) == 0.0
                dir3 = cross(dir, dir2)
                rot = vcat(dir', dir2', dir3')
                push!(cardinal_rots, RotMatrix3{Float64}(rot))
            end
        end
    end
    cardinal_rots
end

function rotations_90()
    [
        let
            rot = zeros(3)
            rot[i] = sign * pi/2
            RotXYZ(rot...)
        end
        for i in 1:3 for sign in [-1, 1]
    ]
end

export cardinal_rotations, rotations_90


function HDirProdPs(x,y)
#HDirProdFn    Calculates horizontal direct product of two matrices with equal number of rows.
#              z[i,:] is the Kronecker product of x[i:i,:] and y[i:i,:]
#
#
#  Usage:    z = HDirProdFn(x,y)
#
#  Input:    x      T x Kx matrix
#            y      T x Ky matrix
#
#  Output:   z      T x (Kx*Ky)
#
#
#  Example:   x = [ 1 2;
#                   3 4]
#             y = [5 6 1;
#                  7 8 1]
#
#             the z = HDirProdFn(x,y) gives
#
#             z = [ 5  6  1  10  12  2;
#                  21 24  3  28  32  4]
#
# Paul.Soderlind@unisg.ch, Oct 2015
#----------------------------------------------------------------------------

  #T = size(x,1)       #rows and columns in x
  Kx = size(x,2)       #columns in x
  Ky = size(y,2)       #columns in y

  z = repmat(y,1,Kx) .* kron(x,ones(1,Ky))

  return z

end
#----------------------------------------------------------------------------


cubic(x::Real) = x * x * x

function apy2(x::Real, y::Real)
    isnan(x) && return x
    isnan(y) && return y
    # general case
    xabs = abs(x)
    yabs = abs(y)
    w = max(xabs, yabs)
    z = min(xabs, yabs)
    iszero(z) && return w
    return w * sqrt(1.0 + abs2(z / w))
end

function sqr2(x::Real, y::Real)
    isnan(x) && return x
    isnan(y) && return y
    # general case
    xabs = abs(x)
    yabs = abs(y)
    w = max(xabs, yabs)
    z = min(xabs, yabs)
    iszero(z) && return abs2(w)
    return abs2(w) * (1.0 + abs2(z / w))
end

function unsafe_copy!(des::AbstractArray, src::AbstractArray)
    @simd for i in eachindex(des)
        @inbounds des[i] = src[i]
    end
    return nothing
end

nrm2(x::VecI{Tx}) where Tx<:Real = sqrt(dot(x, x))

function nrm2(x::VecI{Tx}, y::VecI{Ty}, b::VecB{Tb}) where {Tx<:Real,Ty<:Real,Tb<:Real} # @code_warntype ✓
    @simd for i in eachindex(b)
        @inbounds b[i] = abs2(x[i] - y[i])
    end
    return sqrt(sum(b))
end

function swap!(v::VecIO, i::Int, j::Int) # @code_warntype ✓
    @inbounds temp = v[i]
    @inbounds v[i] = v[j]
    @inbounds v[j] = temp
    return nothing
end

function swap!(m::MatI, i1::Int, j1::Int, i2::Int, j2::Int) # @code_warntype ✓
    @inbounds temp     = m[i1,j1]
    @inbounds m[i1,j1] = m[i2,j2]
    @inbounds m[i2,j2] = temp
    return nothing
end

function checksquare(A::MatI)
    m, n = size(A)
    m == n || error("tr: A should be a square matrix")
    return n
end

function diagmax(A::MatI, n::Int)
    r = -Inf
    @inbounds for i in eachindex(1:n)
        r = max(r, A[i,i])
    end
    return r
end

diagmax(A::MatI) = (n = checksquare(A); return diagmax(A, n))

function tr(A::MatI, n::Int)
    ret = 0.0
    @inbounds for i in eachindex(1:n)
        ret += A[i,i]
    end
    return ret
end

"""
    tr(A)
Matrix trace. Sums the diagonal elements of `A`.
`A` must be a square matrix.
"""
tr(A::MatI) = (n = checksquare(A); return tr(A, n))

function dot(x::VecI, y::VecI, n::Int)
    r = 0.0
    m = mod(n, 5)
    if m ≠ 0
        @inbounds for i in 1:m
            r += x[i] * y[i]
        end
        n < 5 && return r
    end
    m += 1
    @inbounds for i in m:5:n
        r += x[i] * y[i] + x[i+1] * y[i+1] + x[i+2] * y[i+2] + x[i+3] * y[i+3] + x[i+4] * y[i+4]
    end
    return r
end

function dot(x::VecI, m::Int, A::MatI, y::VecI, n::Int)
    ret = 0.0
    for j in eachindex(1:n)
        @inbounds yj = y[j]
        if !iszero(yj)
            tmp = 0.0
            @inbounds for i in eachindex(1:m)
                tmp += A[i,j] * x[i]
            end
            ret += tmp * yj
        end
    end
    return ret
end

"""
    dot(x, y)
Compute the dot product `x' * y` between two vectors `x` and `y`.
"""
function dot(x::VecI, y::VecI)
    n = length(x)
    n == length(y) || error("dot: length(y) ≠ length(x) = $n.")
    return dot(x, y, n)
end

"""
    dot(x, A, y)
Compute the dot product `x' * A * y` between two vectors `x` and `y`.
"""
function dot(x::VecI, A::MatI, y::VecI)
    m, n = size(A)
    m == length(x) || error("dot: length(x) ≠ size(A, 1) = $m.")
    n == length(y) || error("dot: length(y) ≠ size(A, 2) = $n.")
    return dot(x, m, A, y, n)
end

function logdet(A::MatI, n::Int)
    r = 0.0
    @inbounds for i in eachindex(1:n)
        r += log(abs(A[i,i]))
    end
    return r
end

"""
    logdet(A)
Compute the log-determinant `log(det(A))` where `A` is a triangular matrix.
"""
logdet(A::MatI) = (n = checksquare(A); return logdet(A, n))


function binarize_proportion(g::Union{SimpleWeightedGraph,SimpleWeightedDiGraph}, prop)
    b = nonzeros(g.weights)
    N = round(Int, length(b) * prop)
    return binarize_absolute(g, sort(b)[N])
end

function compute_threshold_from_proportion(
    g::Union{SimpleWeightedGraph,SimpleWeightedDiGraph},
    prop::AbstractVector
)
    b = nonzeros(g.weights)
    N = length(b)
    sb = sort(b)
    return [sb[round(Int, N * p)] for p in prop]
end

function binarize_absolute(g, thresh)
    out = SimpleGraph(nv(g))
    for e in edges(g)
        if weight(e) > thresh
            add_edge!(out, src(e), dst(e))
        end
    end
    return out
end

each_threshold(op, g, thresh) = [op(binarize_absolute(g, t)) for t in thresh]
each_proportion(op, g, prop) = each_threshold(op, g, compute_threshold_from_proportion(g, prop))


using FileIO

"""
    get_data(args; is_save=true)

Get data for the specified `dataset`.

For synthetic data, it will be randomly generated and
if `is_save` is `true` it will be saved in the first time and for the next time it will be loaded.
"""
function get_data(args; is_save=true)
    tr_sz = args[:tr_sz]
    te_sz = args[:te_sz]
    dataset = args[:dataset]
    local features = nothing
    if dataset == "synth" || dataset == "synth_large" || dataset == "ibp_gauss_gen" || dataset == "sbc_gauss_gen"
        data_file = "../data/$(args[:dataset]).jld2"
        # Load saved generated data
        if ispath(data_file)
           return load(data_file, "x_tr", "x_te", "features")
        end
        local x
        n = tr_sz + te_sz
        gauss_dim = 256
        k_max = 25
        if dataset == "ibp_gauss_gen"
            x, features = gen_synth(args[:alpha], gauss_dim, n)
        elseif dataset == "sbc_gauss_gen"
            x, features = gen_synth(args[:alpha], gauss_dim, n, k_max)
        else
            if dataset == "synth"
                features = get_features_griffiths2011indian()
            elseif dataset == "synth_large"
                features = get_features_large()
            end
            x = mix_features(features, n; alpha=args[:alpha])
        end
        x_tr, x_te = Matrix{FT}(x[:,1:tr_sz]), Matrix{FT}(x[:,tr_sz+1:end])
        # Save data
        if is_save
            data_folder = "../data"
            if !ispath(data_folder)
                mkpath(data_folder)
            end
            save(data_file, Dict("x_tr" => x_tr, "x_te" => x_te, "features" => features))
            @info "Synthetic data ($(args[:dataset])) is saved to $data_file"
        end
    elseif dataset == "mnist"
        x_tr, _, x_te, _ = load_mnist(:mnist, tr_sz, te_sz)
    elseif dataset == "fmnist"
        x_tr, _, x_te, _ = load_mnist(:fmnist, tr_sz, te_sz)
    else throw("[main] Unkown dataset: $dataset") end

    feature_size = features == nothing ? "nothing" : size(features)
    @info "Dataset size" size(x_tr) size(x_te) feature_size
    return x_tr, x_te, features
end

function gen_synth(alpha, D, N, K...)
    Z = rand(IBP(alpha), N, K...)
    K = size(Z, 2)
    A = randn(D, K)
    X = A * Z'
    return X, A
end

"""
    mix_features(features, N; alpha=nothing, clip01=true, noise_level=0.1)

Generate data by randomly mixing features.

By default it uses each feature uniformly, but if `alpha` is provided, it generates the activations following a corresponding stick-breaking construction of IBP.
`features` is assumed to be in a `D` by `K` matrix, where `K` is the number of features and `D` is the dimension of each feature.
"""
function mix_features(features, N; alpha=nothing, clip01=true, noise_level=0.1)
    D, K = size(features)

    # Get activation matrix
    Z = if alpha != nothing
        # Mix via stick-breaking construction
        rand(IBP(alpha), N, K)
    else
        # Mix by uniformly activation
        z_list = []
        for n = 1:N
            z_raw = Vector{Int}(map(x->parse(Int, string(x)), collect(Base.bin(unsigned(rand(0:2^K-1)), 4, false))))
            z = zeros(Int, K)
            j = K
            for i = length(z_raw):-1:1
                z[j] = z_raw[i]
                j = j - 1
            end
            push!(z_list, z)
        end
        hcat(z_list...)'
    end
    @info "Maximum number of features" sum(sum(Z; dims=1) .> 0)

    # Mixing feature and activations
    X = features * Z'

    # Add Gaussian noise
    X = X + randn(D, N) * noise_level

    # Clip to be between 0 and 1
    if clip01
        X[X .> 1.0] .= 1.0
        X[X .< 0.0] .= 0.0
    end

    return X
end


function next_gen(a::BitArray{1}, isperiodic=false)
    b = copy(a)
    if isperiodic
        ncnt = prepend!(a[1:end-1], [a[end]]) + append!(a[2:end], [a[1]])
    else
        ncnt = prepend!(a[1:end-1], [false]) + append!(a[2:end], [false])
    end
    b[ncnt .== 0] = false
    b[ncnt .== 2] = ~b[ncnt .== 2]
    return b
end

function show_gen(a::BitArray{1})
    s = join([i ? "\u2588" : " " for i in a], "")
    s = "\u25ba"*s*"\u25c4"
end

hi = 70
a = bitrand(hi)
b = falses(hi)
println("A 1D Cellular Atomaton with ", hi, " cells and empty bounds.")
while any(a) && any(a .!= b)
    println("    ", show_gen(a))
    b = copy(a)
    a = next_gen(a)
end
a = bitrand(hi)
b = falses(hi)
println()
println("A 1D Cellular Atomaton with ", hi, " cells and periodic bounds.")
while any(a) && any(a .!= b)
    println("    ", show_gen(a))
    b = copy(a)
    a = next_gen(a, true)
end


using DelimitedFiles

function ReadData(fname::AbstractString)
    return readdlm(fname, '\n', String)
end

function Part1(data::Array)
    nDigits = length(data[1])
    nLines = length(data)

    # convert raw data into binary
    D = map(x -> parse(Int,x,base=2), data)

    # sums of the digits = number of 1s in each column
    S = [sum((i >> j) & 1 for i in D) for j=0:nDigits-1]

    # compute gamma
    gamma = sum(Int(S[i]>nLines/2) << (i-1) for i=1:nDigits)

    # compute epsilon
    epsilon = ~gamma & sum(1<<i for i=0:nDigits-1)

    return gamma*epsilon
end

function Part2(data::Array)
    nDigits = length(data[1])
    nLines = length(data)

    # convert raw data into binary, do this twice for O2 and CO2 ratings
    D_O2  = map(x -> parse(Int,x,base=2), data)
    D_CO2 = map(x -> parse(Int,x,base=2), data)

    j = 1 # the bit we're currently working on
    while length(D_O2) > 1
        # determine most common bit
        cBit = Int(sum((i >> (nDigits-j)) & 1 for i in D_O2) >= length(D_O2)/2)
        # filter out numbers which don't have same common first bit
        D_O2 = filter(x -> ((x >> (nDigits-j)) & 1) == cBit, D_O2)

        j += 1
    end

    j = 1 # the bit we're currently working on
    while length(D_CO2) > 1
        # determine least common bit
        cBit = Int(sum((i >> (nDigits-j)) & 1 for i in D_CO2) < length(D_CO2)/2)
        # filter out numbers which don't have same common first bit
        D_CO2 = filter(x -> ((x >> (nDigits-j)) & 1) == cBit, D_CO2)

        j += 1
    end

    return D_O2[1]*D_CO2[1]
end

# test data
#data = ["00100", "11110", "10110", "10111", "10101", "01111", "00111", "11100",
#        "10000", "11001", "00010", "01010"]

data = ReadData("../data/day3.txt")

println(Part1(data))

println(Part2(data))


include("../include.jl")

using JuMP, Ipopt

function create_1(n_v::Int64)
  m = Model()
  @variable(m, 0.0 <= theta[i=1:n_v] <= pi)
  @variable(m, 0.0 <= r[i=1:n_v] <= 1.0)

  @NLobjective(m, Min, -0.5 * sum(r[i+1] * r[i] * sin(theta[i+1] - theta[i]) for i = 1:(n_v - 1)))
  for i = 1:(n_v-1)
    for j = (i+1):n_v
      @NLconstraint(m, r[i]^2 + r[j]^2 - 2 * r[i] * r[j] * cos(theta[i] - theta[j]) <= 1)
    end
    @NLconstraint(m, theta[i] <= theta[i+1])
  end
  @constraint(m, r[n_v] == 0.0)
  @constraint(m, theta[n_v] == pi)

  return m;
end

m = create_1(20)
nlp_raw = MathProgNLPModel(m);
my_pars = Class_parameters();
it = one_phase_solve(nlp_raw,my_pars);
it.point.x
#=
setsolver(m,IpoptSolver())
solve(m)
=#


#=
Created on Monday 13 Jan 2020
Last update: Monday 20 Jan 2020

@author: Michiel Stock
michielfmstock@gmail.com

Just some plots of the testing functions
=#

using STMO, Plots, LaTeXStrings

using STMO.TestFuns: branin, rosenbrock, rastrigine, flower, booth, fquadr, fnonquadr

cmap = :speed

# NOTE: some contours are log-scaled to visualize the minimum better

# branin
contour(-5:0.02:20, -3:0.02:23, (x1, x2) -> log(branin([x1, x2])), xlabel="\$ x_1 \$",
                                    ylabel="\$ x_2 \$", color=cmap)
savefig("figures/brainin.png")

# rosenbrock
contour(-2:0.01:2, -2:0.01:2, (x1, x2) -> log(rosenbrock([x1, x2])), xlabel="\$ x_1 \$",
                                    ylabel="\$ x_2 \$", color=cmap)
savefig("figures/rosenbrock.png")

# rastrigine
contourf(-5.2:0.1:5.2, -5.2:0.1:5.2, (x1, x2) -> rastrigine([x1, x2]),
                    xlabel="\$ x_1 \$", ylabel="\$ x_2 \$", color=cmap)
savefig("figures/rastrigine.png")

# flower
contourf(-3:0.1:3, -3:0.1:3, (x1, x2) -> flower([x1, x2]),
                        xlabel="\$ x_1 \$", ylabel="\$ x_2 \$", color=cmap)
savefig("figures/flower.png")

# booth
contour(-10:0.1:10, -10:0.1:10, (x1, x2) -> booth([x1, x2]) |> log,
                        xlabel="\$ x_1 \$", ylabel="\$ x_2 \$", color=cmap)
savefig("figures/booth.png")

# quadratic
contour(-10:0.1:10, -5:0.1:5, (x1, x2) -> fquadr((x1, x2)), xlabel="\$ x_1 \$",
                ylabel="\$ x_2 \$", title="quadratic", fill=true)
savefig("figures/quadratic.png")

# non-quadratic
contour(-2:0.1:2, -1:0.1:1, (x1, x2) -> fnonquadr((x1, x2)), xlabel="\$ x_1 \$",
                ylabel="\$ x_2 \$", title="non-quadratic", fill=true)
savefig("figures/nonquadratic.png")


function maxconcat(arr::Vector{<:Integer})
    b = sort(dec.(arr); lt=(x, y) -> x * y < y * x, rev=true) |> join
    return try parse(Int, b) catch parse(BigInt, b) end
end

tests = ([1, 34, 3, 98, 9, 76, 45, 4],
         [54, 546, 548, 60],
         [1, 34, 3, 98, 9, 76, 45, 4, 54, 546, 548, 60])

for arr in tests
    println("Max concatenating in $arr:\n -> ", maxconcat(arr))
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.6.2
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown]
# https://cometscome.github.io/JuliaFromFortran/build/chapter2/01/

# %% [markdown]
# https://docs.julialang.org/en/v1/manual/variables-and-scoping/#Loops-and-Comprehensions

# %%
function f()
    i = -99
    for i in 1:3 end
    i
end

@show f();

# %%
function g()
    local i
    for outer i in 1:3 end
    i
end

@show g();

# %%
function mandel()
    nx = 61; ny = 31; maxiter = 90
    x0 = -2.0; x1 = 2.0; y0 = -2.0; y1 = 2.0
    mandelbrot = zeros(Int64,nx, ny)
    mandelbrot .= maxiter
    for iy=1:ny
        y = y0 + (y1 - y0) * (iy - 1) / real(ny - 1)
        for ix=1:nx
            x = x0 + (x1 - x0) * (ix - 1) / real(nx - 1)
            c = x + im*y
            z = 0im
            for iter=0:maxiter
                z = z * z + c
                if abs(z) > 2
                    mandelbrot[ix,iy] = iter
                    break
                end
            end

        end
    end
    for iy=1:ny
        for ix=1:nx
            if mandelbrot[ix,iy] == maxiter
                print("*")
            else
                print((mandelbrot[ix,iy]+9) ÷ 10)
            end
        end
        println("\t")
    end
end
mandel()

# %%
function mandel_outer()
    nx = 61; ny = 31; maxiter = 90
    x0 = -2.0; x1 = 2.0; y0 = -2.0; y1 = 2.0
    mandelbrot = zeros(Int64,nx, ny)
    for iy=1:ny
        y = y0 + (y1 - y0) * (iy - 1) / real(ny - 1)
        for ix=1:nx
            x = x0 + (x1 - x0) * (ix - 1) / real(nx - 1)
            c = x + im*y
            z = 0im
            local iter; for outer iter=0:maxiter
                z = z * z + c
                abs(z) > 2 && break
            end
            mandelbrot[ix,iy] = iter
        end
    end
    for iy=1:ny
        for ix=1:nx
            if mandelbrot[ix,iy] == maxiter
                print("*")
            else
                print((mandelbrot[ix,iy]+9) ÷ 10)
            end
        end
        println("\t")
    end
end
mandel_outer()

# %%


# ---
# title: 947. Most Stones Removed with Same Row or Column
# id: problem947
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Depth-first Search, Union Find
# link: <https://leetcode.com/problems/most-stones-removed-with-same-row-or-column/description/>
# hidden: true
# ---
#
# On a 2D plane, we place stones at some integer coordinate points.  Each
# coordinate point may have at most one stone.
#
# Now, a _move_ consists of removing a stone that shares a column or row with
# another stone on the grid.
#
# What is the largest possible number of moves we can make?
#
#
#
# **Example 1:**
#
#
#
#     Input: stones = [[0,0],[0,1],[1,0],[1,2],[2,1],[2,2]]
#     Output: 5
#
#
# **Example 2:**
#
#
#
#     Input: stones = [[0,0],[0,2],[1,1],[2,0],[2,2]]
#     Output: 3
#
#
# **Example 3:**
#
#
#
#     Input: stones = [[0,0]]
#     Output: 0
#
#
#
#
# **Note:**
#
#   1. `1 <= stones.length <= 1000`
#   2. `0 <= stones[i][j] < 10000`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


import Base.*

# Group action is x ↦ a*x+b
struct SpaceGroupElement{N,T<:Integer} <: GroupElement
    a::SMatrix{N,N,T}
    b::SVector{N, Rational{T}}
    function SpaceGroupElement{N,T}(a::SMatrix{N,N,T}, b::SVector{N,Rational{T}}) where {N, T<:Integer}
        new(a, b-floor.(b))
    end
end

function SpaceGroupElement(a::SMatrix{N,N,T}, b::SVector{N,Rational{T}}) where {N,T<:Integer}
    SpaceGroupElement{N,T}(a, b)
end

SpaceGroupElement{N,T}() where {N,T<:Integer} =
    SpaceGroupElement(one(SMatrix{N,N,T}), zeros(SVector{N, Rational{T}}))

function *(e1::SpaceGroupElement{N,T}, e2::SpaceGroupElement{N,T}) where {N,T<:Integer}
    SpaceGroupElement(e1.a*e2.a, e1.a*e2.b+e1.b)
end

const SpaceGroup{N,T} = FiniteGroup{SpaceGroupElement{N,T}}


te = [-0.112639 0.124715 0.124715 0.268717 0.124715 0.268717 0.268717 0.046154]
st = SymmetricTensor(reshape(te, (2,2,2)))
Random.seed!(43)
c = cumulants(randn(1000, 20), 4)

@testset "axiliary functions" begin
    @test rep((1,2,3)) == 6
    @test rep((1,2,2)) == 3
    @test rep((1,1,1)) == 1
end

@testset "norm" begin
    @test norm(st) ≈ 0.5273572868359742
    @test norm(st, 1) ≈ 1.339089
    @test norm(st, 2.5) ≈ norm(te, 2.5)
    n = norm(c[2])
    @test n ≈ norm(Array(c[2]))
    @test cnorms(c) ≈ [norm(c[3])/(n^(3/2)), norm(c[4])/(n^2)]
end


# Intersection methods
# TODO: Change intersection methods -- This breaks with infinite slope
function intersect(l1::LineSegment, l2::LineSegment)
    # Check for all cases in which they don't intersect
    # Will check for: 1) Whether l1.p1 and l1.p2 are on opposite sides of
    # l2.p1->l2.p2 2) Whether l2.p1 and l2.p2 are on opposite sides of
    # l1.p1->l2.p2 3) Whether they have same slope
    # See: http://jeffe.cs.illinois.edu/teaching/373/notes/x06-sweepline.pdf
    p1, p2 = l1.p1, l1.p2
    p3, p4 = l2.p1, l2.p2
    if ccw(p1, p3, p4) == ccw(p2, p3, p4)
        return (false, Point(Inf, Inf))
    elseif ccw(p3, p1, p2) == ccw(p4, p1, p2)
        return (false, Point(Inf, Inf))
    elseif l1.slope == l2.slope
        return (false, Point(Inf, Inf))
    end

    # If we make it this far, then line segments should intersect
    # Pull out points
    x1 = l1.p1[1]
    y1 = l1.p1[2]
    x2 = l1.p2[1]
    y2 = l1.p2[2]
    x3 = l2.p1[1]
    y3 = l2.p1[2]
    x4 = l2.p2[1]
    y4 = l2.p2[2]

    # Get intersection
    denom = (x1-x2)*(y3-y4) - (y1-y2)*(x3-x4)
    x_intersect = (x1*y2 - y1*x2)*(x3-x4) - (x1-x2)*(x3*y4 - y3*x4)/denom
    y_intersect = (x1*y2 - y1*x2)*(y3-y4) - (y1-y2)*(x3*y4 - y3*x4)/denom
    intersection = Point(x_intersect, y_intersect)

    # Check if intersection is on both intervals
    if in(x_intersect, Interval(x1, x2)) && in(x_intersect, Interval(x3, x4))
        out = (true, intersection)
    # If not, return intersection but negative status
    else
        out = (false, intersection)
    end
in
    return out
end

function intersect(ch1::ConvexHull, ch2::ConvexHull)
    error("Intersection between convex sets is not currently implemented")
end

function intersect(ch::ConvexHull, q::Quadrant)

    # Get information about the set of extreme points
    ext_pts = ch.extremepoints
    xmin, _, xmax, _ = extremax(ext_pts)
    ymin, _, ymax, _ = extremay(ext_pts)
    npts = length(ch)

    # Get information about quadrant
    o = q.origin
    lx = LineSegment(Point(clamp(q.lx.p1[1], xmin, xmax), clamp(q.lx.p1[2], ymin, ymax)),
                     Point(clamp(q.lx.p2[1], xmin, xmax), clamp(q.lx.p2[2], ymin, ymax)))
    ly = LineSegment(Point(clamp(q.ly.p1[1], xmin, xmax), clamp(lx.p1[2], ymin, ymax)),
                     Point(clamp(q.ly.p2[1], xmin, xmax), clamp(lx.p2[2], ymin, ymax)))

    # Start a convex hull to store outpoints in
    prev_pt = ext_pts[end]
    outhull = in(prev_pt, q) ? [prev_pt] : []

    # Will check between each point individually
    # Points need to be ordered either clockwise
    # or counter-clockwise
    for i=1:npts
        # Current extreme point and corresponding line
        c_ep = ext_pts[i]
        c_ls = LineSegment(prev_pt, c_ep)

        # Compute intersections with lines
        stat_x, c_inter_x = intersect(lx, c_ls)
        stat_y, c_inter_y = intersect(ly, c_ls)

        # If there is intersection, push to hull
        stat_x ? push!(outhull, c_inter_x) : nothing
        stat_y ? push!(outhull, c_inter_y) : nothing

        # Add current point if it is in quadrant
        in(c_ep, q) ? push!(outhull, c_ep) : nothing
        prev_pt = c_ep

    end

    return outhull
end

# function intersect(ch::ConvexHull, q::Quadrant)
#     # TODO: Need to consider special cases where it is not binding on one side
#     # but is on the other. It should be as simple as returning all extrema to
#     # the right side of the binding extreme. Right thing to do here is break
#     # this function into cases and write the cases out, that will make it
#     # look much less messy

#     # Get information about the set of extreme points
#     ext_pts = ch.extremepoints
#     _, xmin, _, xmax = extremax(ext_pts)
#     _, ymin, _, ymax = extremay(ext_pts)
#     npts = length(ep)

#     # Get information about quadrant
#     o = q.origin

#     # Counters to determine whether points belong in the
#     # set or not. vcount and hcount are for number of intersections
#     # the line makes with the convex hull each way. If either are
#     # even then origin is not in intersection
#     checknext = true
#     storecurr = false
#     vcount = 0
#     hcount = 0
#     # Start a convex hull to store outpoints in
#     outhull = []

#     # Will check between each point individually
#     # Points need to be ordered either clockwise
#     # or counter-clockwise
#     for i=2:npts+1
#         # Pull out points
#         if i<=npts
#             x1 = ep[i-1]
#             x2 = ep[i]
#         else
#             x1 = ep[npts]
#             x2 = ep[1]
#         end

#         # First add point if we are supposed to store
#         if storecurr
#             push!(outhull, x1)
#         end

#         # Check whether intersects
#         w1intersect = evalliney(x1, x2, o.x)
#         w2intersect = evallinex(x1, x2, o.y)
#         cond1 = ((x2.x - o.x)*(x1.x - o.x) < 0) && (w1intersect >= o.y)
#         cond2 = ((x2.y - o.y)*(x1.y - o.y) < 0) && (w2intersect >= o.x)

#         # Check first condition
#         if cond1
#             storecurr = !storecurr
#             push!(outhull, Point(w1, w1intersect))
#             vcount += 1
#         # Account for if point is right on top of value
#         elseif isapprox(x2.x, o.x) && isapprox(x2.y, w1intersect)
#             if storecurr
#                 push!(outhull, x2)
#             end
#             storecurr = !storecurr
#             vcount += 1
#         end

#         # Check second condition
#         if cond2
#             storecurr = !storecurr
#             # Make sure we don't double count if both conditions hold
#             push!(outhull, Point(w2intersect, w2))
#             hcount += 1
#         elseif isapprox(x2.y, o.y) && isapprox(x2.x, w2intersect)
#             if storecurr
#                 push!(outhull, x2)
#             end
#             storecurr = !storecurr
#             hcount += 1
#         end
#     end

#     if (hcount*vcount)%2!=0 || (o.x<minx(ep)[1].x && o.y<miny(ep)[1].y)
#         push!(outhull, o)
#     end

#     return outhull
# end


alpha = 0.5
accuracy = 0.0000001
edgelist = "test/fixtures/simple.adj"

rank = Rank{Float64,Int64}(alpha, accuracy, edgelist)
r = stationary_distribution(rank)
count = length(r)

expected = [float(5 / 18), float(4 / 9), float(5 / 18)]
precision = 5

@assert isequal(floor(expected, precision), floor(r, precision))


"""
COSINEKNN
---------
    compute the k-nearest neighbors similarity metric between the
    vertices of A or the upper half of a bipartite graph A.

Functions
---------
- S = cosineknn{T}(A::SparseMatrixCSC{T,Int64},K::Int64)
- S = cosineknn(A::MatrixNetwork,K::Int64)

Example
-------
~~~
A = load_matrix_network("bfs_example")
S = cosineknn(A,2)
~~~
"""

## CSC support
function cosineknn{T}(A::SparseMatrixCSC{T,Int64},K::Int64)
    (rp,ci,ai) = sparse_to_csr(A)
    (rpt,cit,ait) = sparse_to_csr(A')
    (m,n) = size(A)
    return cosineknn_internal(rp,ci,ai,rpt,cit,ait,m,K)
end

##  MatrixNetwork support
function cosineknn(A::MatrixNetwork,K::Int64)
    # for the original matrix
    (rp,ci,ai) = (A.rp,A.ci,A.vals)
    # for the transposed matrix:
    M = sparse_transpose(A)
    At = MatrixNetwork(M)
    (rpt,cit,ait) = (At.rp,At.ci,At.vals)
    m = A.n
    return cosineknn_internal(rp,ci,ai,rpt,cit,ait,m,K)

end

function cosineknn_internal{T}(rp::Vector{Int64},ci::Vector{Int64},ai::Vector{T},
                rpt::Vector{Int64},cit::Vector{Int64},ait::Vector{T},m::Int64,K::Int64)

    # accumarray
    rn = zeros(Float64,maximum(cit))
    for i = 1:length(ait)
        rn[cit[i]] += ait[i]^2
    end
    rn = sqrt(rn)

    rn[rn.>eps()] = 1./rn[rn.>eps()] # do division once

    si = zeros(Int64,m*K)
    sj = zeros(Int64,m*K)
    sv = zeros(Float64,m*K)
    nused = 0

    dwork = zeros(Float64,m)
    iwork = zeros(Int64,m)
    iused = zeros(Int64,m)

    for i = 1:m

        # for each row of A, compute an inner product against all rows of A.
        # to do this efficiently, we know that the inner-product will
        # only be non-zero if two rows of A overlap in a column.  Thus, in
        # row i of A, we look at all non-zero columns, and then look at all
        # rows touched by those columns in the A' matrix.  We  compute
        # the similarity metric, then sort and only store the top k.

        curused = 0 # track how many non-zeros we compute during this operation
        for ri = rp[i]:rp[i+1]-1
            # find all columns j in row i
            j = ci[ri]
            aij = ai[ri]*rn[i]
            for rit = rpt[j]:rpt[j+1]-1
                # find all rows k in column j
                k = cit[rit]
                if k == i
                    continue
                end # skip over the standard entries
                akj = ait[rit]*rn[k]
                if iwork[k]>0
                    # we already have a non-zero between row i and row k
                    dwork[iwork[k]] += aij*akj
                else
                    # we need to add a non-zero betwen row i and row k
                    curused = curused + 1
                    iwork[k] = curused
                    dwork[curused] = aij*akj
                    iused[curused] = k
                end
            end
        end

        # don't sort if fewer than K elements
        if curused < K
            sperm = 1:curused
        else
            sperm = sortperm(dwork[1:curused],rev=true)
        end

        for k = 1:min(K,length(sperm))
            nused = nused + 1
            si[nused] = i
            sj[nused] = iused[sperm[k]]
            sv[nused] = dwork[sperm[k]]
        end

        # reset the iwork array, no need to reset dwork, as it's just
        # overwritten
        for k = 1:curused
            iwork[iused[k]] = 0
        end
    end

    ei = si[1:nused]
    ej = sj[1:nused]
    ev = sv[1:nused]

    S = sparse(ei,ej,ev,m,m)

    return S
end