# Script to replicate figure 2 and 3

using StatsBase, Statistics, LinearAlgebra, StatsPlots, XLSX, PrettyTables

# -------------- Previous scripts -----------------------------------------------------------------------------------------------------------------------

include("src/data_wrangling.jl") # formats the raw input data obtained from WIOD

# -------------- Sectoral aggregation for country level I/O-Table ---------------------------------------------------------------------------------------

# the function "country_level_GVC_indicators" uses the output from the function "obtain_matrices"
# it prepares the data to replicate figure 2 and 3 by aggregating from a country-sector level to a country level I/O-Table
# outputs:
# t_country_wide ... yearly statistics of final demand share, value added share and measures of U and D at country level

function country_level_GVC_indicators(df::DataFrame, year::Int64)

    ID, FD, GO, VA, IV, other = obtain_matrices(df, year) # N*S×N*S, N*S×N, N*S×1, N*S×1, N*S×N, N*S×1

    # aggregate to country level
    ID_ctry = [sum(ID[i:i+n_sectors-1, j:j+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_sectors:n_sectors*n_ctrys] # N×N
    FD_ctry = [sum(FD[i:i+n_sectors-1, j]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_ctrys] # N×N, sum over origin sector
    IV_ctry = [sum(IV[i:i+n_sectors-1, j]) for i in 1:n_sectors:n_sectors*n_ctrys, j in 1:n_ctrys] # N×N, sum over origin sector
    GO_ctry = [sum(GO[i:i+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys] # N×1, sum over destination countries
    VA_ctry = [sum(VA[i:i+n_sectors-1]) for i in 1:n_sectors:n_sectors*n_ctrys] # N×1

    # net inventory correction
    iv_correction = [GO_ctry[i]/(GO_ctry[i]-IV_ctry[i, j]) for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, to match country wise inventory changes

    ID_ctry = ID_ctry .* iv_correction # N×N, element wise multiplication
    FD_ctry = FD_ctry .* iv_correction # N×N

    FD_ctry = [sum(FD_ctry[i, :]) for i in 1:n_ctrys] # N×1, sum over destination countries

    # input/output coefficient matrices
    A = [ID_ctry[i, j] / GO_ctry[j] for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, iterate output over destination country
    B = [ID_ctry[i, j] / GO_ctry[i] for i in 1:n_ctrys, j in 1:n_ctrys] # N×N, iterate output over reporting country

    # country level statistics
    FD_GO = FD_ctry ./ GO_ctry # N×1
    VA_GO = VA_ctry ./ GO_ctry # N×1
    U_num = inv(I - A) * GO_ctry # N×1, numerator of equation (5)
    D_num = transpose(GO_ctry) * inv(I - B) # N×1, numerator of equation (7), notice difference in matrix algebra

    U = U_num ./ GO_ctry # N×1
    D = transpose(D_num) ./ GO_ctry # N×1, transpose to have column vector

    # weighted average of all countries to create a world GVC position
    FD_GO_world = sum(FD_GO .* GO_ctry) / sum(GO_ctry)
    VA_GO_world = sum(VA_GO .* GO_ctry) / sum(GO_ctry)
    U_world = sum(U .* GO_ctry) / sum(GO_ctry)
    D_world = sum(D .* GO_ctry) / sum(GO_ctry)

    # collect data in a nice dataframe
    summary_stats = hcat(fill(year, size(ctrys)), ctrys, FD_GO, VA_GO, U, D)
    t_country_wide = DataFrame(summary_stats, :auto)
    push!(t_country_wide, [year "world" FD_GO_world VA_GO_world U_world D_world]) # add data for world
    rename!(t_country_wide, [:year, :country, :FD_GO, :VA_GO, :U, :D])
    transform!(t_country_wide, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(x -> round(x, digits=3)) .=> [:FD_GO, :VA_GO, :U, :D]) # rescale U, D for better comparison

    return t_country_wide

end

# gather data for all years in a single data frame
t_fig_2_3 = DataFrame(year=Int[], country=String[], FD_GO=Float64[], VA_GO=Float64[], U=Float64[], D=Float64[])

for i in years
    append!(t_fig_2_3, country_level_GVC_indicators(df, i))
end

XLSX.writetable("clean/t_fig_2_3.xlsx", t_fig_2_3, overwrite=true) # export table to folder as input for plotting figure 2 and 3

# -------------- Figure 2 --------------------------------------------------------------------------------------------------------------------------------

t_fig_2_3 = DataFrame(XLSX.readtable("clean/t_fig_2_3.xlsx", "Sheet1")...) # load directly from file
transform!(t_fig_2_3, :country .=> ByRow(string) .=> :country, renamecols=false) # change type to String
transform!(t_fig_2_3, :year .=> ByRow(Int) .=> :year, renamecols=false) # change type to Int64
transform!(t_fig_2_3, [:FD_GO, :VA_GO, :U, :D] .=> ByRow(Float64) .=> [:FD_GO, :VA_GO, :U, :D], renamecols=false) # change type to Float64

t_fig_2_3_long = stack(t_fig_2_3, Not([:year, :country])) # long format for plotting

function plot_fig2(df::DataFrame, country::Vector, series::Vector) # use vectors as input so we can plot multiple countries for comparison

    df = subset(df, :country => ByRow(x -> x in country), :variable => ByRow(x -> x in series))

    p = @df df plot(:year, :value, group=(:country, :variable), lw=2, legend=:topleft,
        ylims=(round(minimum(:value), digits=1)-0.15,round(maximum(:value), digits=1)+0.15))

    return p
end

l = @layout [a ; b] # 2×1 layout
p1 = plot_fig2(t_fig_2_3_long, ["world"], ["FD_GO", "VA_GO"]) # add both measures (should be the same in theory) - slight difference (should adjust for "others"?)
p2 = plot_fig2(t_fig_2_3_long, ["world"], ["U", "D"]) # coincide very well
title = "Figure 2: GVC Positioning over Time (World Average)"
p_fig2 = plot(p1, p2, layout=l, plot_title=title, plot_titlefontsize=10)
savefig(p_fig2, "images/figure2.png") # export image to folder

# -------------- Figure 3 ------------------------------------------------------------------------------------------------------------------------------

# computing percentiles per country-series group
t_fig3 = subset(t_fig_2_3_long, :country => ByRow(x -> x != "world")) # take out world to calculate percentiles
gdf = groupby(t_fig3, [:year, :variable])
begin
    t_fig3 = combine(gdf,
     :value => (x -> percentile(x, 25)) => :quantile_25,
     :value => (x -> percentile(x, 50)) => :quantile_50,
     :value => (x -> percentile(x, 75)) => :quantile_75)
end
rename!(t_fig3, :variable => :series) # name conflict in formatting
df_fig3_long = stack(t_fig3, Not([:year, :series]))

function plot_fig3(df::DataFrame, series::String) # input vectors in so we can plot multiple countries for comparison

    df = subset(df, :series => ByRow(x -> x == series))

    p = @df df plot(:year, :value, group=:variable, lw=2, legend=:none,
    ylims=(round(minimum(:value)-0.05, digits=1),round(maximum(:value)+0.05, digits=1)),
    ylabel="$(series) over time")

    return p
end

l = @layout [a b ; c d] # 2×2 layout
p1 = plot_fig3(df_fig3_long, "FD_GO")
p2 = plot_fig3(df_fig3_long, "VA_GO")
p3 = plot_fig3(df_fig3_long, "U")
p4 = plot_fig3(df_fig3_long, "D")
title = "Figure 3: GVC Positioning over Time (25th, 50th, 75th country percentiles)"
p_fig3 = plot(p1, p2, p3, p4, layout=l, plot_title=title, plot_titlefontsize=10)
savefig(p_fig3, "images/figure3.png") # export image to folder

# -------------- Table 1 ------------------------------------------------------------------------------------------------------------------------------

t_tab1 = subset(t_fig_2_3_long, :country => ByRow(x -> x != "world")) # take out world
gdf = groupby(t_tab1, [:year, :variable])
#t_tab1 = combine(gdf, :value => (x -> ordinalrank(x; lt=isless)) => :rank) # ordinal rank per year and series
t_tab1 = transform(gdf, :value => (x -> ordinalrank(x; lt=isless)) => :rank) # ordinal rank per year and series
# notice the difference between combine and transform! => transform returns all other columns as well (very useful!)
sort!(t_tab1, [:year, :variable, :rank]) # sorting is still respecting groups

# ------------ Spearman coefficient of correlation (page 15)
# need to sort correctly, i.e. sort according to countries!
# function produced a table computing the spearman correlation coefficient between two specified years
function spearman_coeff(df::DataFrame, year_base::Int64, year_comp::Int64)
    df = sort(df, :country)
    spearman_coefficients = DataFrame(year_base=Int[], year_comp=Int[], measure=String[], value=Float64[])

    for series in unique(df.variable)
        base = subset(df, :year => ByRow(x -> x == year_base), :variable => ByRow(x -> x == series))
        comp = subset(df, :year => ByRow(x -> x == year_comp), :variable => ByRow(x -> x == series))

        spearman_row = [year_base year_comp series round(StatsBase.corspearman(base.rank, comp.rank), digits=3)]

        push!(spearman_coefficients, spearman_row)
    end

    return spearman_coefficients
end

spear_coeff = spearman_coeff(t_tab1, 1995, 2011) # coincide well with paper

# export table to include in README.md
open("images/spearman_coefficients.txt", "w") do f
    pretty_table(f,spear_coeff; backend = Val(:html), nosubheader=true, standalone=false) # only prints html code for table itself
end

# ------------ Table entries: returns top5/bottom5
function table_1(df::DataFrame, year::Int64, series::String, selection::String)
    df = subset(df, :year => ByRow(x -> x == year), :variable => ByRow(x -> x == series))
    view = selection == "top5" ? first(df, 5) : last(df, 5)
    return view
end

# to produce output in one go
for i in [1995, 2011]
    for j in unique(t_tab1.variable)
        for k in ["top5", "bottom5"]
            print(table_1(t_tab1, i, j, k))
        end
    end
end

# final demand share of output
table_1(t_tab1, 1995, "FD_GO", "top5")
table_1(t_tab1, 1995, "FD_GO", "bottom5")
table_1(t_tab1, 2011, "FD_GO", "top5")
table_1(t_tab1, 2011, "FD_GO", "bottom5")

# value added share of output
table_1(t_tab1, 1995, "VA_GO", "top5")
table_1(t_tab1, 1995, "VA_GO", "bottom5")
table_1(t_tab1, 2011, "VA_GO", "top5")
table_1(t_tab1, 2011, "VA_GO", "bottom5")

# upstreamness
table_1(t_tab1, 1995, "U", "top5")
table_1(t_tab1, 1995, "U", "bottom5")
table_1(t_tab1, 2011, "VA_GO", "top5")
table_1(t_tab1, 2011, "VA_GO", "bottom5")

# downstreamness
table_1(t_tab1, 1995, "U", "top5")
table_1(t_tab1, 1995, "U", "bottom5")
table_1(t_tab1, 2011, "VA_GO", "top5")
table_1(t_tab1, 2011, "VA_GO", "bottom5")


struct CauchyPoint end

function (cp::CauchyPoint)(state, B, Δ)
    p = similar(state.x)
    ∇f = state.∇f
    norm_∇f = norm(∇f)
    prod = ∇f'*B*∇f
    _cauchypoint!(p, Δ, ∇f, norm_∇f, prod)
    return p
end

function _cauchypoint!(p, Δ, ∇f, norm_∇f, prod)
    p .= -Δ/norm_∇f*∇f
    if prod > 0
        p .*= min(1, norm_∇f^3/(Δ*prod))
    end
    return p
end


# STO-nG basis definition

abstract type BasisFunction end
abstract type GTO <: BasisFunction end
abstract type GaussianFunction <: GTO end

"Gaussian s wave function"
struct Gaussian_s <: GaussianFunction
    α::Float64       # gussian function exponent
    R::Vec3{Float64} # center (or actually nuclues' position)
end

"Gaussian p wave function"
struct Gaussian_p <:GaussianFunction
    α::Float64       # gussian function exponent
    R::Vec3{Float64} # center
    n::Vec3{Float64} # orbital direction
end

"STO-nG function"
struct STONG{N} <: GTO
    d::NTuple{N, Float64}           # coefficients
    g::NTuple{N, GaussianFunction}  # gaussian functions
    STONG{N}(d::NTuple{N, Float64}, g::NTuple{N, GaussianFunction}) where N = new(d, g)
end

STONG(d::NTuple{N, Float64}, g::NTuple{N, GaussianFunction}) where N = STONG{N}(d, g)

struct STONG_coeff{N}
    coefficients::NTuple{N, Float64}
    exponents::NTuple{N,Float64}
    STONG_coeff{N}(coeff::NTuple{N, Float64}, _exp::NTuple{N, Float64}) where N = new(coeff, _exp)
end

"construct STONG from STONG_coeff"
STONG(coeff::STONG_coeff{N}, center::Vec3{Float64}; type::Symbol=:s) where N = begin
    if type === :s
        return STONG{N}(
            coeff.coefficients,
            ntuple(i->Gaussian_s(coeff.exponents[i], center), N)
        )
    elseif type === :px
        return STONG{N}(
            coeff.coefficients,
            ntuple(i->Gaussian_p(coeff.exponents[i], center, vec3(1., 0., 0.)), N)
        )
    elseif type === :py
        return STONG{N}(
            coeff.coefficients,
            ntuple(i->Gaussian_p(coeff.exponents[i], center, vec3(0., 1., 0.)), N)
        )
    elseif type === :pz
        return STONG{N}(
            coeff.coefficients,
            ntuple(i->Gaussian_p(coeff.exponents[i], center, vec3(0., 0., 1.)), N)
        )
    end
    error("orbital type $type is not support yet!")
end

const STO3G = STONG{3}
const STO3G_coeff = STONG_coeff{3}

# the data comes from basis_set_exchange
# https://github.com/MolSSI-BSE/basis_set_exchange
const coefficient_data = Dict(
    "H" => Dict(
        "1s" => STO3G_coeff(
            (0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00), # coefficients
            (0.3425250914E+01, 0.6239137298E+00, 0.1688554040E+00)  # exponents
        )
    ),
    "He" => Dict(
        "1s" => STO3G_coeff(
            (0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00), # coefficients
            (0.6362421394E+01, 0.1158922999E+01, 0.3136497915E+00)  # exponents
        )
    ),
    "C" => Dict(
        "1s" => STO3G_coeff(
            (0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00), # coefficients
            (0.7161683735E+02, 0.1304509632E+02, 0.3530512160E+01)  # exponents
        ),
        "2s" => STO3G_coeff(
            (-0.9996722919E-01, 0.3995128261E+00, 0.7001154689E+00),# coefficients
            (0.2941249355E+01, 0.6834830964E+00, 0.2222899159E+00)  # exponents
        ),
        "2p" => STO3G_coeff(
            (0.1559162750E+00, 0.6076837186E+00, 0.3919573931E+00), # coefficients
            (0.2941249355E+01, 0.6834830964E+00, 0.2222899159E+00)  # exponents
        )
    ),
    "N" => Dict(
        "1s" => STO3G_coeff(
            (0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00), # coefficients
            (0.9910616896E+02, 0.1805231239E+02, 0.4885660238E+01)  # exponents
        ),
        "2s" => STO3G_coeff(
            (-0.9996722919E-01, 0.3995128261E+00, 0.7001154689E+00),# coefficients
            (0.3780455879E+01, 0.8784966449E+00, 0.2857143744E+00)  # exponents
        ),
        "2p" => STO3G_coeff(
            (0.1559162750E+00, 0.6076837186E+00, 0.3919573931E+00), # coefficients
            (0.3780455879E+01, 0.8784966449E+00, 0.2857143744E+00)  # exponents
        )
    ),
    "O" => Dict(
        "1s" => STO3G_coeff(
            (0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00), # coefficients
            (0.1307093214E+03, 0.2380886605E+02, 0.6443608313E+01)  # exponents
        ),
        "2s" => STO3G_coeff(
            (-0.9996722919E-01, 0.3995128261E+00, 0.7001154689E+00),# coefficients
            (0.5033151319E+01, 0.1169596125E+01, 0.3803889600E+00)  # exponents
        ),
        "2p" => STO3G_coeff(
            (0.1559162750E+00, 0.6076837186E+00, 0.3919573931E+00), # coefficients
            (0.5033151319E+01, 0.1169596125E+01, 0.3803889600E+00)  # exponents
        )
    ),
    "F" => Dict(
        "1s" => STO3G_coeff(
            (0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00), # coefficients
            (0.1666791340E+03, 0.3036081233E+02, 0.8216820672E+01)  # exponents
        ),
        "2s" => STO3G_coeff(
            (-0.9996722919E-01, 0.3995128261E+00, 0.7001154689E+00),# coefficients
            (0.6464803249E+01, 0.1502281245E+01, 0.4885884864E+00)  # exponents
        ),
        "2p" => STO3G_coeff(
            (0.1559162750E+00, 0.6076837186E+00, 0.3919573931E+00), # coefficients
            (0.6464803249E+01, 0.1502281245E+01, 0.4885884864E+00)  # exponents
        )
    )
)

Base.length(::STONG{N}) where N = N

Base.show(io::IO, gs::Gaussian_s) = begin
    print(io, "gs(α=$(gs.α), R=$(gs.R))")
end

Base.show(io::IO, gp::Gaussian_p) = begin
    print(io, "gp(α=$(gp.α), R=$(gp.R), n=$(gp.n))")
end

Base.show(io::IO, stong::STONG{N}) where N = begin
    println(io, "STO-$(N)G")
    for i = 1:N
        println(io, "  $(stong.d[i]) × $(stong.g[i])")
    end
end


include("../utils.jl")

struct TestVectorSpaceType <: VectorSpaceType end

@testset "Tangent bundle" begin
    M = Sphere(2)

    @testset "Nice access to vector bundle components" begin
        TB = TangentBundle(M)
        p = ProductRepr([1.0, 0.0, 0.0], [0.0, 2.0, 4.0])
        @test p[TB, :point] === p.parts[1]
        @test p[TB, :vector] === p.parts[2]
        p[TB, :vector] = [0.0, 3.0, 1.0]
        @test p.parts[2] == [0.0, 3.0, 1.0]
        p[TB, :point] = [0.0, 1.0, 0.0]
        @test p.parts[1] == [0.0, 1.0, 0.0]
        @test_throws DomainError p[TB, :error]
        @test_throws DomainError p[TB, :error] = [1, 2, 3]
    end

    types = [Vector{Float64}]
    TEST_FLOAT32 && push!(types, Vector{Float32})
    TEST_STATIC_SIZED && push!(types, MVector{3,Float64})

    for T in types
        p = convert(T, [1.0, 0.0, 0.0])
        TB = TangentBundle(M)
        TpM = TangentSpaceAtPoint(M, p)
        @test sprint(show, TB) == "TangentBundle(Sphere(2, ℝ))"
        @test base_manifold(TB) == M
        @test manifold_dimension(TB) == 2 * manifold_dimension(M)
        @test representation_size(TB) == (6,)
        CTB = CotangentBundle(M)
        @test sprint(show, CTB) == "CotangentBundle(Sphere(2, ℝ))"
        @test sprint(show, VectorBundle(TestVectorSpaceType(), M)) ==
              "VectorBundle(TestVectorSpaceType(), Sphere(2, ℝ))"
        @testset "Type $T" begin
            pts_tb = [
                ProductRepr(convert(T, [1.0, 0.0, 0.0]), convert(T, [0.0, -1.0, -1.0])),
                ProductRepr(convert(T, [0.0, 1.0, 0.0]), convert(T, [2.0, 0.0, 1.0])),
                ProductRepr(convert(T, [1.0, 0.0, 0.0]), convert(T, [0.0, 2.0, -1.0])),
            ]
            @inferred ProductRepr(
                convert(T, [1.0, 0.0, 0.0]),
                convert(T, [0.0, -1.0, -1.0]),
            )
            for pt in pts_tb
                @test bundle_projection(TB, pt) ≈ pt.parts[1]
            end
            diag_basis = DiagonalizingOrthonormalBasis(log(TB, pts_tb[1], pts_tb[2]))
            basis_types = (
                DefaultOrthonormalBasis(),
                get_basis(TB, pts_tb[1], DefaultOrthonormalBasis()),
                diag_basis,
                get_basis(TB, pts_tb[1], diag_basis),
            )
            test_manifold(
                TB,
                pts_tb,
                test_injectivity_radius=false,
                test_reverse_diff=isa(T, Vector),
                test_forward_diff=isa(T, Vector),
                test_tangent_vector_broadcasting=false,
                test_vee_hat=true,
                test_project_tangent=true,
                test_project_point=true,
                test_default_vector_transport=true,
                vector_transport_methods=[ParallelTransport()],
                basis_types_vecs=basis_types,
                projection_atol_multiplier=4,
                test_inplace=true,
            )

            # tangent space at point
            pts_TpM = map(
                p -> convert(T, p),
                [[0.0, 0.0, 1.0], [0.0, 2.0, 0.0], [0.0, -1.0, 1.0]],
            )
            basis_types = (
                DefaultOrthonormalBasis(),
                get_basis(TpM, pts_TpM[1], DefaultOrthonormalBasis()),
            )
            @test get_basis(TpM, pts_TpM[1], basis_types[2]) === basis_types[2]
            test_manifold(
                TpM,
                pts_TpM,
                test_injectivity_radius=true,
                test_reverse_diff=isa(T, Vector),
                test_forward_diff=isa(T, Vector),
                test_tangent_vector_broadcasting=true,
                test_vee_hat=false,
                test_project_tangent=true,
                test_project_point=true,
                test_default_vector_transport=true,
                basis_types_vecs=basis_types,
                projection_atol_multiplier=4,
                test_inplace=true,
            )
        end
    end

    @test TangentBundle{ℝ,Sphere{2,ℝ}} ==
          VectorBundle{ℝ,Manifolds.TangentSpaceType,Sphere{2,ℝ}}
    @test CotangentBundle{ℝ,Sphere{2,ℝ}} ==
          VectorBundle{ℝ,Manifolds.CotangentSpaceType,Sphere{2,ℝ}}

    @test base_manifold(TangentBundle(M)) == M
    @testset "spaces at point" begin
        p = [1.0, 0.0, 0.0]
        t_p = TangentSpaceAtPoint(M, p)
        t_p2 = TangentSpace(M, p)
        @test t_p == t_p2
        ct_p = CotangentSpaceAtPoint(M, p)
        t_ps = sprint(show, "text/plain", t_p)
        sp = sprint(show, "text/plain", p)
        sp = replace(sp, '\n' => "\n ")
        t_ps_test = "Tangent space to the manifold $(M) at point:\n $(sp)"
        @test t_ps == t_ps_test
        @test base_manifold(t_p) == M
        @test base_manifold(ct_p) == M
        @test t_p.fiber.manifold == M
        @test ct_p.fiber.manifold == M
        @test t_p.fiber.fiber == TangentSpace
        @test ct_p.fiber.fiber == CotangentSpace
        @test t_p.point == p
        @test ct_p.point == p
        @test injectivity_radius(t_p) == Inf
        @test representation_size(t_p) == representation_size(M)
        v = [0.0, 0.0, 1.0]
        @test embed(t_p, v) == v
        @test embed(t_p, v, v) == v
        # generic vector space at
        fiber = VectorBundleFibers(TestVectorSpaceType(), M)
        v_p = VectorSpaceAtPoint(fiber, p)
        v_ps = sprint(show, "text/plain", v_p)
        fiber_s = sprint(show, "text/plain", fiber)
        v_ps_test = "$(typeof(v_p))\nFiber:\n $(fiber_s)\nBase point:\n $(sp)"
        @test v_ps == v_ps_test
    end

    @testset "tensor product" begin
        TT = Manifolds.TensorProductType(TangentSpace, TangentSpace)
        @test sprint(show, TT) == "TensorProductType(TangentSpace, TangentSpace)"
        @test vector_space_dimension(VectorBundleFibers(TT, Sphere(2))) == 4
        @test vector_space_dimension(VectorBundleFibers(TT, Sphere(3))) == 9
        @test base_manifold(VectorBundleFibers(TT, Sphere(2))) == M
        @test sprint(show, VectorBundleFibers(TT, Sphere(2))) ==
              "VectorBundleFibers(TensorProductType(TangentSpace, TangentSpace), Sphere(2, ℝ))"
    end

    @testset "Error messages" begin
        vbf = VectorBundleFibers(TestVectorSpaceType(), Euclidean(3))
        @test_throws ErrorException inner(vbf, [1, 2, 3], [1, 2, 3], [1, 2, 3])
        @test_throws ErrorException Manifolds.project!(vbf, [1, 2, 3], [1, 2, 3], [1, 2, 3])
        @test_throws ErrorException zero_vector!(vbf, [1, 2, 3], [1, 2, 3])
        @test_throws MethodError vector_space_dimension(vbf)
        a = fill(0.0, 6)
        @test_throws ErrorException get_coordinates!(
            TangentBundle(M),
            a,
            ProductRepr([1.0, 0.0, 0.0], [0.0, 0.0, 0.0]),
            ProductRepr([1.0, 0.0, 0.0], [0.0, 0.0, 1.0]),
            CachedBasis(DefaultOrthonormalBasis(), []),
        )
    end

    @testset "log and exp on tangent bundle for power and product manifolds" begin
        M = PowerManifold(Circle(ℝ), 2)
        N = TangentBundle(M)
        p1 = ProductRepr([0.0, 0.0], [0.0, 0.0])
        p2 = ProductRepr([-1.047, -1.047], [0.0, 0.0])
        X1 = log(N, p1, p2)
        @test isapprox(N, p2, exp(N, p1, X1))
        @test is_vector(N, p2, vector_transport_to(N, p1, X1, p2))

        M2 = ProductManifold(Circle(ℝ), Euclidean(2))
        N2 = TangentBundle(M2)
        p1_2 = ProductRepr(ProductRepr([0.0], [0.0, 0.0]), ProductRepr([0.0], [0.0, 0.0]))
        p2_2 = ProductRepr(
            ProductRepr([-1.047], [1.0, 0.0]),
            ProductRepr([-1.047], [0.0, 1.0]),
        )
        @test isapprox(N2, p2_2, exp(N2, p1_2, log(N2, p1_2, p2_2)))

        ppt = PowerVectorTransport(ParallelTransport())
        tbvt = Manifolds.VectorBundleVectorTransport(ppt, ppt)
        @test TangentBundle(M, tbvt).vector_transport === tbvt
        @test CotangentBundle(M, tbvt).vector_transport === tbvt
        @test VectorBundle(TangentSpace, M, tbvt).vector_transport === tbvt
    end
end


# This file contains methods to solve an instance (heuristically or with CPLEX)
using CPLEX
using Printf
include("generation.jl")

TOL = 0.00001
directions=[1 0;-1 0;0 1;0 -1]

"""
Solve an instance with CPL

Argument
-t: array which stores the puzzle data

Return

"""

function cplexSolve(t::Array{Int,2})

    rowNum=size(t,1)
    colNum=size(t,2)
    if rowNum!=colNum
        println("not square matrix")
        return
    end
    @printf("%d-%d\n",rowNum,colNum)
    # Create the model
    m = Model(CPLEX.Optimizer)

    # x[i,j]=1 if the cell is white, 0 if the cell is black
    @variable(m,x[1:rowNum,1:colNum],Bin)

    ss=0
    # set the cell white if it has a number
    for l in 1:rowNum
        for c in 1:colNum
            if t[l,c]!=0
                @constraint(m,x[l,c]==1)
                ss=ss+1
            end
        end
    end

    # No two black squares are orthogonally adjacent
    @constraint(m,[r in 1:rowNum-1, c in 1:colNum],x[r,c]+x[r+1,c]>=1)
    @constraint(m,[r in 1:rowNum, c in 1:colNum-1],x[r,c]+x[r,c+1]>=1)

    # For any two white squares, there is a path between them using only white squares.
    @constraint(m,[c in 2:colNum-1],sum(x[1,c-1:c+1])+x[2,c]>=2) # the first line except the corners
    @constraint(m,[c in 2:colNum-1],sum(x[rowNum,c-1:c+1])+x[rowNum-1,c]>=2) # the last line except the corners
    @constraint(m,[r in 2:rowNum-1],sum(x[r-1:r+1,1])+x[r,2]>=2) # the first column except the corners
    @constraint(m,[r in 2:rowNum-1],sum(x[r-1:r+1,colNum])+x[r,colNum-1]>=2) # the last colNum except the corners
    @constraint(m,sum(x[1:2,1])+x[1,2]>=2) # the left top corner
    @constraint(m,sum(x[1:2,colNum])+x[1,colNum-1]>=2) # the right top corner
    @constraint(m,sum(x[rowNum-1:rowNum,1])+x[rowNum,2]>=2) # the left bottom corner
    @constraint(m,sum(x[rowNum-1:rowNum,colNum])+x[rowNum,colNum-1]>=2) # the right bottom corner

    @constraint(m,[r in 2:rowNum-1, c in 2:colNum-1],sum(x[r-1:r+1,c-1:c+1])>=2)

    #case when the black is in diagonal line
    @constraint(m,sum(x[r,r] for r=1:rowNum)>=1)
    @constraint(m, sum(x[rowNum-r+1,r] for r=1:rowNum)>=1)

    function my_callback_function(cb_data)
        x_vals = callback_value.(Ref(cb_data), x)
        zz=0
        para=Array{Int64,2}(undef,0,6)
        for l in 1:rowNum
            for c in 1:colNum
                if t[l,c]!=0
                    up=0 # the number of white cells in top direction until meet black cell
                    right=0 # the number of white cells in right direction until meet black cell
                    down=0 # the number of white cells in down direction until meet black cell
                    left=0 # the number of white cells in left direction until meet black cell
                    for i1 in 1:l-1
                        if  x_vals[l-i1,c]== 0  #meet the black cell
                            down=i1-1
                            break
                        else
                            down=l-1
                        end
                    end
                    for j1 in 1:rowNum-l
                        if  x_vals[l+j1,c]== 0
                            up=j1-1
                            break
                        else
                            up=rowNum-l
                        end
                    end
                    for i2 in 1:c-1
                        if  x_vals[l,c-i2]== 0
                            left=i2-1
                            break
                        else
                            left=c-1
                        end
                    end
                    for j2 in 1:colNum-c
                        if  x_vals[l,c+j2]== 0
                            right=j2-1
                            break
                        else
                            right=colNum-c
                        end
                    end
                    if up+down+left+right+1>t[l,c]
                        # println(up+down+left+right+1)
                        con1 = @build_constraint((sum(x[l:l+up,c])+sum(x[l-down:l,c])+sum(x[l,c-left:c])+sum(x[l,c:c+right])-3)>=t[l,c]) #counting the white cells from four orthogonal directions
                        # # # con1=@build_constraint(up+down+left+right+1==t[l,c])
                        MOI.submit(m, MOI.LazyConstraint(cb_data), con1)
                    elseif up+down+left+right+1<t[l,c]
                        con = @build_constraint((sum(x[l:l+up,c])+sum(x[l-down:l,c])+sum(x[l,c-left:c])+sum(x[l,c:c+right])-3)<=t[l,c])
                        MOI.submit(m, MOI.LazyConstraint(cb_data), con)
                    elseif up+down+left+right+1==t[l,c]
                        zz=zz+1
                        para=vcat(para,[l c up down left right])
                    end
                end
            end
        end
        # println(zz,"/",ss)
        if(zz>=ss)

            for i=1:ss
                l,c,up,down,left,right=para[i,:]
                # @printf("func:%d+%d+%d+%d+1=%d\n",up,down,left,right,t[l,c]) #up,down,left,right,target
                # con = @build_constraint((sum(x[l:l+up,c])+sum(x[l-down:l,c])+sum(x[l,c-left:c])+sum(x[l,c:c+right])-3)==t[l,c])

                # MOI.submit(m, MOI.LazyConstraint(cb_data), con)

            end
        end
    end

    MOI.set(m, MOI.LazyConstraintCallback(), my_callback_function)
    # Maximize the number of white cells
    @objective(m,Max, sum(x))

    # Start a chronometer
    start = time()

    # Solve the model
    optimize!(m)

    # println(num_constraints(m, VariableRef, MOI.ZeroOne))
    # println("-----",termination_status(m))
    # Return:
    # 1 - true if an optimum is found(type: Bool)
    # 2 - the resolution time(type Float64)
    # 3 - the value of each cell(type:Array{undef,2})
    return JuMP.primal_status(m) == JuMP.MathOptInterface.FEASIBLE_POINT, time() - start,x

end

"""
Heuristically solve an instance
"""
function visiable(x_vals::Array{Bool, 2},l::Int,c::Int)
    rowNum=size(x_vals,1)
    colNum=size(x_vals,2)

    up=0 # the number of white cells in top direction until meet black cell
    right=0 # the number of white cells in right direction until meet black cell
    down=0 # the number of white cells in down direction until meet black cell
    left=0 # the number of white cells in left direction until meet black cell
    for i1 in 1:l-1
        if  x_vals[l-i1,c]== 0  #meet the black cell
            down=i1-1
            break
        else
            down=l-1
        end
    end
    for j1 in 1:rowNum-l
        if  x_vals[l+j1,c]== 0
            up=j1-1
            break
        else
            up=rowNum-l
        end
    end
    for i2 in 1:c-1
        if  x_vals[l,c-i2]== 0
            left=i2-1
            break
        else
            left=c-1
        end
    end
    for j2 in 1:colNum-c
        if  x_vals[l,c+j2]== 0
            right=j2-1
            break
        else
            right=colNum-c
        end
    end
    return up+down+left+right+x_vals[l,c]
end

function dfs(l::Int,c::Int,x_vals::Array{Bool, 2},walked::Array{Bool, 2})

    rowNum=size(x_vals,1)
    colNum=size(x_vals,2)
    ans=1
    for i in 1:4
        r1=l+directions[i,1]
        c1=c+directions[i,2]
        if r1>=1 && r1<=rowNum && c1>=1 && c1<=colNum && walked[r1,c1]==0 && x_vals[r1,c1]!=0
            walked[r1,c1]=1
            ans+=dfs(r1,c1,x_vals,walked)
        end
    end
    walked=zeros(Bool,rowNum,colNum)
    return ans
end

"""
check if it is possible to place a black cell in position(r,c) and if possible, place black cell
return true if possible, false if impossible
"""
function placeBlack(r::Int, c::Int,x_vals::Array{Bool,2})
    rowNum=size(x_vals,1)
    colNum=size(x_vals,2)
    # if the cell is already a black one
    if x_vals[r,c]!=1
        return false
    end

    # there is black cells in its 4 orthogonal neighbors
    for i=1:size(directions,1)
        if r+directions[i,1]>=1 && r+directions[i,1]<=rowNum && c+directions[i,2]>=1 &&c+directions[i,2]<=colNum && x_vals[r+directions[i,1],c+directions[i,2]]==0
            return false
        end
    end
    x_vals[r,c]=0
    return true
end

"""
unplace the black cell, change it to white
"""
function unplaceBlack(r::Int, c::Int,x_vals::Array{Bool,2})
    x_vals[r,c]=1
end

"""
iterativly solve the puzzle
"""
function solve(r::Int,c::Int,rw::Int,cw::Int,x_vals::Array{Bool,2},walked::Array{Bool,2},t::Array{Int,2})
    rowNum=size(x_vals,1)
    colNum=size(x_vals,2)

    if c > colNum
        c = 1
        r+=1
    end
    # check if white cells are continuous
    reach=0
    walked[rw,cw]=1
    if x_vals[rw,cw]!=1
        x_vals[rw,cw]=1
        reach+=dfs(rw,cw,x_vals,walked)
        x_vals=0
    else
        reach+=dfs(rw,cw,x_vals,walked)
    end
    walked=zeros(Bool,rowNum,colNum)
    if reach!=sum(x_vals)
        # println("white cells are not contiguous,reach=",reach,"sum=",sum(x_vals))
        return 0
    end
    # check after all iteration whether the puzzle is finally solve
    if r>rowNum
        solved=true
        for i in 1:rowNum
            for j in 1:colNum
                if t[i,j]>0 # for the cell who has a number
                    if visiable(x_vals,i,j)!=t[i,j] # not solved
                        # println("no solution")
                        solved=false
                        return 0
                    end
                end
            end
        end
        if solved==true
            # finally solved
            # println("solve successfully")
            # println(x_vals)
            return x_vals
        end
        return 0
    end
    # check in each iteration
    for i in 1:rowNum
        for j in 1:colNum
            if t[i,j]>0 # for the cell who has a number
                if visiable(x_vals,i,j)<t[i,j] # not solved
                    # println(visiable(x_vals,i,j)," less than ",t[i,j])
                    return 0
                end
                if i<r && visiable(x_vals[1:r-1,:],i,j)>t[i,j]
                    # println("before ",r)
                    return 0
                end
                # println("nothing:",visiable(x_vals,i,j),"->",t[i,j])
            end
        end
    end
    if placeBlack(r,c,x_vals)==true
        rw1=rw
        cw1=cw
        # update the position of white cell
        if rw1==r && cw1==c
            cw1+=1 # move right
            if cw1>colNum
                cw1=1
                rw1+=1
            end
        end
        # println("new->",r,c+1,rw1,cw1,x_vals,walked,t)
        ret_val=solve(r,c+1,rw1,cw1,x_vals,walked,t)
        # println(ret_val)
        if ret_val!=0 && ret_val!=nothing
            return ret_val
        end
        unplaceBlack(r,c,x_vals)
    end
    walked=zeros(Bool,rowNum,colNum)
    # println("new start",r,c+1,rw,cw,x_vals,walked,t)
    ret_val=solve(r,c+1,rw,cw,x_vals,walked,t)
    if ret_val!=0 && ret_val!=nothing
        return ret_val
    end
    # println(ret_val)
end

function heuristicSolve(t::Array{Int, 2})
    rowNum=size(t,1)
    colNum=size(t,2)
    walked=zeros(Bool,rowNum,colNum)
    x_vals=ones(Bool,rowNum,colNum)
    ret_val=solve(1,1,1,1,x_vals,walked,t)
    isOptimal=false
    if ret_val!=0 && ret_val!=nothing
        isOptimal=true
        println("solve successfully")
        println(ret_val)
    else
        isOptimal=false
        println("no solution")
        println(ret_val)
    end
    return isOptimal,ret_val
end

"""
Solve all the instances contained in "../data" through CPLEX and heuristics

The results are written in "../res/cplex" and "../res/heuristic"

Remark: If an instance has previously been solved (either by cplex or the heuristic) it will not be solved again
"""
function solveDataSet()

    dataFolder = "../data/"
    resFolder = "../res/"

    # Array which contains the name of the resolution methods
    # resolutionMethod = ["cplex"]
    resolutionMethod = ["cplex", "heuristique"]

    # Array which contains the result folder of each resolution method
    resolutionFolder = resFolder .* resolutionMethod

    # Create each result folder if it does not exist
    for folder in resolutionFolder
        if !isdir(folder)
            mkdir(folder)
        end
    end

    global isOptimal = false
    global solveTime = -1

    # For each instance
    # (for each file in folder dataFolder which ends by ".txt")
    for file in filter(x->occursin(".txt", x), readdir(dataFolder))

        println("-- Resolution of ", file)
        t=readInputFile(dataFolder * file)

        # For each resolution method
        for methodId = 1:size(resolutionMethod,1)

            outputFile = resolutionFolder[methodId] * "/" * file

            # If the instance has not already been solved by this method
            if !isfile(outputFile)

                fout = open(outputFile, "w")

                resolutionTime = -1
                isOptimal = false

                # If the method is cplex
                if resolutionMethod[methodId] == "cplex"

                    # Solve it and get the results
                    isOptimal, resolutionTime,x = cplexSolve(t)

                    # If a solution is found, write it
                    if isOptimal
                        writeSolution(fout, x)
                    end

                # If the method is one of the heuristics
                else

                    isSolved = false

                    # Start a chronometer
                    startingTime = time()

                    solution=nothing

                    # While the grid is not solved and less than 100 seconds are elapsed

                    # Solve it and get the results
                    isOptimal, solution = heuristicSolve(t)

                    # Stop the chronometer
                    resolutionTime = time() - startingTime

                    # Write the solution (if any)
                    if isOptimal

                        writeSolution(fout, solution)

                    end
                end

                println(fout, "solveTime = ", resolutionTime)
                println(fout, "isOptimal = ", isOptimal)

                close(fout)
            end

            # Display the results obtained with the method on the current instance
            include(outputFile)
            println(resolutionMethod[methodId], " optimal: ", isOptimal)
            println(resolutionMethod[methodId], " time: " * string(round(solveTime, sigdigits=2)) * "s\n")
        end
    end
end


## Load the MNIST Graph and draw it
##
using DelimitedFiles
using SparseArrays
using LinearAlgebra
ei,ej = copy.(eachcol(Int.(readdlm("mnist-train-4.edges"))))
n = maximum(ei)
A = sparse(ei,ej,1,n,n)
@assert issymmetric(A)
##
using PyCall
using Conda
pyimport_conda("igraph","python-igraph","conda-forge")
##
using PyCall
igraph = pyimport("igraph")

function igraph_layout(A::SparseMatrixCSC{T}, layoutname::AbstractString="drl") where T
    ei,ej,ew = findnz(A)
    edgelist = [(ei[i]-1,ej[i]-1) for i = 1:length(ei)]
    nverts = size(A)
    G = igraph.Graph(nverts, edges=edgelist, directed=false)
    xy = G.layout(layoutname)
    xy = [Float64(xy[i][j]) for i in 1:length(xy),  j in 1:length(xy[1])]
end

xy = igraph_layout(A, "drl")
##
writedlm("mnist-train-4.xy", xy)
##
xy = readdlm("mnist-train-4.xy")
##
using GraphRecipes
using LinearAlgebra
using Plots
graphplot(ei, ej, x =xy[:,1], y=xy[:,2],
  markercolor=:black, markerstrokecolor=:white,
  size=(1200,1200),dpi=400,
  markersize=2, linecolor=1, linealpha=0.2, linewidth=0.5,
  markeralpha=0.3,
  #colorbar=true,
  marker_z=train_y,
  axis_buffer=0.02, background=nothing)
for i=1:100
  annotate!(xy[i,1],xy[i,2],"$(train_y[i])",color=:white)
end
plot!()
##
savefig("mnist-train-4.png")


#=
void acb_hypgeom_u(acb_t res, const acb_t a, const acb_t b, const acb_t z, slong prec)¶
Computes U(a,b,z) using an automatic algorithm choice.

void acb_hypgeom_0f1(acb_t res, const acb_t a, const acb_t z, int regularized, slong prec)¶
Computes the confluent hypergeometric function 0F1(a,z), or inv(Γ(a))0F1(a,z) if regularized is set,

void acb_hypgeom_1f1(acb_t res, const acb_t a, const acb_t b, const acb_t z, int regularized, slong prec)¶
Computes the confluent hypergeometric function M(a,b,z)=1F1(a,b,z), or M(a,b,z)=inv(Γ(b))1F1(a,b,z) if regularized is set.

acb_hypgeom_2f1(acb_t res, const acb_t a, const acb_t b,
        const acb_t c, const acb_t z, int flags, slong prec)
=#

"""
    hypergeometric_0F1(a, z)

confluent hypergeometric function ``₀F₁``
"""
function hypergeometric_0F1(a::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    regularized = Cint(0)
    ccall(@libarb(acb_hypgeom_0f1), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong),
         result, a, z, regularized, P)
    return result
end

const F₀₁ = hypergeometric_0F1

"""
    regular_hypergeometric_0F1(a, z)

regularized confluent hypergeometric function, ``₀F₁ / gamma(a)``
"""
function regular_hypergeometric_0F1(a::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    regularized = Cint(1)
    ccall(@libarb(acb_hypgeom_0f1), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong),
         result, a, z, regularized, P)
    return result
end

const regularF₀₁ = regular_hypergeometric_0F1

"""
    hypergeometric_U(a, b, z)

confluent hypergeometric function of the second kind
"""
function hypergeometric_U(a::ArbComplex{P}, b::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbReal{P}()
    ccall(@libarb(acb_hypgeom_u), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Clong),
         result, a, b, z, P)
    return result
end

"""
    hypergeometric_1F1(a, b, z)

confluent hypergeometric function ``₁F₁``
"""
function hypergeometric_1F1(a::ArbComplex{P}, b::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    regularized = Cint(0)
    ccall(@libarb(acb_hypgeom_1f1), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong),
         result, a, b, z, regularized, P)
    return result
end

const F₁₁ = hypergeometric_1F1

"""
    regular_hypergeometric_1F1(a, b, z)

regularized confluent hypergeometric function, ``₁F₁ / gamma(a)``
"""
function regular_hypergeometric_1F1(a::ArbComplex{P}, b::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    regularized = Cint(1)
    ccall(@libarb(acb_hypgeom_1f1), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong),
         result, a, b, z, regularized, P)
    return result
end

const regularF₁₁ = regular_hypergeometric_1F1

"""
    hypergeometric_2F1(a, b, c, z)

confluent hypergeometric function ``₂F₁``
"""
function hypergeometric_2F1(a::ArbComplex{P}, b::ArbComplex{P}, c::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    regularized = Cint(0)
    ccall(@libarb(acb_hypgeom_2f1), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong),
         result, a, b, c, z, regularized, P)
    return result
end

const F₂₁ = hypergeometric_2F1

"""
    regular_hypergeometric_2F1_regularized(a, b, c, z)

regularized confluent hypergeometric function, ``₂F₁ / gamma(a)``
"""
function regular_hypergeometric_2F1(a::ArbComplex{P}, b::ArbComplex{P}, c::ArbComplex{P}, z::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    regularized = Cint(1)
    ccall(@libarb(acb_hypgeom_2f1), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint, Clong),
         result, a, b, c, z, regularized, P)
    return result
end

const regularF₂₁ = regular_hypergeometric_2F1

function hypergeometric_0F1(a::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    regularized = Cint(0)
    ccall(@libarb(arb_hypgeom_0f1), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Cint, Clong),
         result, a, z, regularized, P)
    return result
end

function regular_hypergeometric_0F1(a::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    regularized = Cint(1)
    ccall(@libarb(arb_hypgeom_0f1), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Cint, Clong),
         result, a, z, regularized, P)
    return result
end

function hypergeometric_U(a::ArbReal{P}, b::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    ccall(@libarb(arb_hypgeom_u), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Clong),
         result, a, b, z, P)
    return result
end

function hypergeometric_1F1(a::ArbReal{P}, b::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    regularized = Cint(0)
    ccall(@libarb(arb_hypgeom_1f1), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Cint, Clong),
         result, a, b, z, regularized, P)
    return result
end

function regular_hypergeometric_1F1(a::ArbReal{P}, b::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    regularized = Cint(1)
    ccall(@libarb(arb_hypgeom_1f1), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Cint, Clong),
         result, a, b, z, regularized, P)
    return result
end

function hypergeometric_2F1(a::ArbReal{P}, b::ArbReal{P}, c::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    regularized = Cint(0)
    ccall(@libarb(arb_hypgeom_2f1), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Cint, Clong),
         result, a, b, c, z, regularized, P)
    return result
end

function regular_hypergeometric_2F1(a::ArbReal{P}, b::ArbReal{P}, c::ArbReal{P}, z::ArbReal{P}) where {P}
    result = ArbReal{P}()
    regularized = Cint(1)
    ccall(@libarb(arb_hypgeom_2f1), Cvoid, (Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Ref{ArbReal}, Cint, Clong),
         result, a, b, c, z, regularized, P)
    return result
end

# ArbFloat

hypergeometric_0F1(a::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeom0f1(ArbReal{P}(a), ArbReal{P}(z)))

hypergeometric_0F1_regularized(a::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeom0f1reg(ArbReal{P}(a), ArbReal{P}(z)))

hypergeometric_U(a::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeomu(ArbReal{P}(a), ArbReal{P}(z)))

hypergeometric_1F1(a::ArbFloat{P}, b::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeom1f1(ArbReal{P}(a), ArbReal{P}(b), ArbReal{P}(z)))

hypergeometric_1F1_regularized(a::ArbFloat{P}, b::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeom1f1reg(ArbReal{P}(a), ArbReal{P}(b), ArbReal{P}(z)))

hypergeometric_1F2(a::ArbFloat{P}, b::ArbFloat{P}, c::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeom1f2(ArbReal{P}(a), ArbReal{P}(b), ArbReal{P}(c), ArbReal{P}(z)))

hypergeometric_1F2_regularized(a::ArbFloat{P}, b::ArbFloat{P}, c::ArbFloat{P}, z::ArbFloat{P}) where {P} =
    ArbFloat{P}(hypgeom1f2reg(ArbReal{P}(a), ArbReal{P}(b), ArbReal{P}(c), ArbReal{P}(z)))


#
#  Hénon-Heiles
#

function HenonHam(u,params)

    @inbounds begin
    q = view(u,1:2)
    p = view(u,3:4)

    T=zero(eltype(u))
    V=zero(eltype(u))

    T=1/2*(p[1]^2+p[2]^2)
    V=1/2*(q[1]^2+q[2]^2)+q[1]^2*q[2]-1/3*q[2]^3

    return(T+V)
    end

end

#
#  Dynamical Pro
#

function dotq(dq,q,p,params,t)
    @inbounds begin
    dq[1] = p[1]
    dq[2] = p[2]
    end
end

function dotp(dp,q,p,params,t)
    @inbounds begin
    dp[1] = -q[1]-2*q[1]*q[2]
    dp[2] = -q[2]-q[1]^2+q[2]^2
    end
end

#
#  Second Order Problem
#

function f2nd!(ddu,du,u,p,t)

  @inbounds begin
  q = view(u,1:2)   # x

  ddu[1] = -q[1]-2*q[1]*q[2]
  ddu[2] = -q[2]-q[1]^2+q[2]^2

  end
end


# strangle_stu.py
#
# This is a MOS model: Mudd Open Source only.
# This is a deep strangle model based upon the Taboga concept (student version).
# Much of the documentation can be found in the core Taboga model or finutil.py
# Designed in Python 3.6 on for Economics 136
# Original version developed on May 30, 2017 by Prof Evans
# This is version 2.1stu, dated April 11, 2017
#

# This model draws upon utility model finutil_stu. That model must either be in
# the same directory as this model or the path command below must be used to
# point to its location.
# import sys
# sys.path.append('c:/Users/Prof Gary Evans/Dropbox/PyGo/PyFi')
import finutil_stu
fu = finutil_stu
# This version calculates IDV for the put and call, but also requires that you
# supply historical volatility (252 and 30 day). You can include an earnings
# date if date is relevant.
#
# Load data
aname = "Prof Evans"
stosym = "MSFT"
tnow = Dates.today()
# Set earnings to True if this is an earnings-based strangle.
earnings = true
# If you don't need to calculate the days to expiry and earnings (optional)
# then calcdays is False (default).
calcdays = false
# BEGIN of optional time data. Load below ONLY if calcdays = True
# Expiration date:
exyear = 2018
exmonth = 2
exday = 16
# Earnings date (load only if relevant):
eayear = 2018
eamonth = 1
eaday = 31
# END of optional time data load
#
# This model does not use the alpha.
# stock = 15.15
# callstrike = 16
# callbid = 0.90
# callask = 0.93
# putstrike = 15
# putbid = 1.150
# putask = 1.18
# sigma252 = 0.043534
# sigma30 = 0.025
###########################
stock = 166.86
callstrike = 170
callbid = 3.20
callask = 3.25
putstrike = 165
putbid = 4.50
putask = 4.60
sigma252 = 0.011358
sigma30 = 0.016387
if calcdays
    days = fu.daysto(exyear,exmonth,exday)
    if earnings
        daystoearn = fu.daysto(eayear,eamonth,eaday)
    end
else
    # Here you have to supply the days if you don't calculate them.
    days = 14.0
    if earnings
        daystoearn = 11.0
    end
end
#
moname = fu.monthname(exmonth)
# Set effective options prices (PEG)
peg = 0.6
cspread = callask - callbid
callpr = callbid + (peg*cspread)
pspread = putask - putbid
putpr = putbid + (peg*pspread)
cost = callpr + putpr
# Note: callbe and putbe should be calced based upon the prices they would
# have to be the day after earnings, especially for more distant options.
callbe = stock + cost
putbe = stock - cost
callbegr = log(callbe/stock)*100
putbegr = log(putbe/stock)*100
# Consider adding log percent breakeven to current value
#
# Calculate call and put ideal prices based upon historical volatilities.
hv252 = sigma252*fu.durvol(days)
cop252 = fu.otranche(stock,callstrike,hv252,true)
pop252 = fu.otranche(stock,putstrike,hv252,false)
hv30 = sigma30*fu.durvol(days)
cop30 = fu.otranche(stock,callstrike,hv30,true)
pop30 = fu.otranche(stock,putstrike,hv30,false)
#
# Calculate call and put IDV based upon actual prices
# See the documentation for finutil and tabogacidv to see what these are doing.
# Start with the call:
#
seedsigma = 1e-6
cutoff = 1e-4
tempop = 0.00
price2 = 0.00
derive = 0.00
cdurseed = seedsigma*fu.durvol(days)
cdaysigma = 0.06
cdursigma = cdaysigma*fu.durvol(days)
#
# The convergence shown here (Newton's method) was designed by Alec Griffith '17
#
while abs(tempop - callpr) > cutoff
    tempop = fu.otranche(stock,callstrike,cdursigma,true)
    price2 = fu.otranche(stock,callstrike,cdursigma+cdurseed,true)
    deriv = (price2-tempop)/seedsigma
    cdaysigma -= (tempop-callpr)/deriv
    cdursigma = cdaysigma*fu.durvol(days)
end
# End of Loop!
#
# Repeat for the put:
#
seedsigma = 1e-6
cutoff = 1e-4
tempop = 0.00
price2 = 0.00
derive = 0.00
pdurseed = seedsigma*fu.durvol(days)
pdaysigma = 0.01
pdursigma = pdaysigma*fu.durvol(days)
call = false
#
#
while abs(tempop - putpr) > cutoff
    tempop = fu.otranche(stock,putstrike,pdursigma,false)
    price2 = fu.otranche(stock,putstrike,pdursigma+pdurseed,false)
    deriv = (price2-tempop)/seedsigma
    pdaysigma -= (tempop-putpr)/deriv
    pdursigma = pdaysigma*fu.durvol(days)
end
# End of Loop!
callsigmaratio = cdaysigma/sigma252
putsigmaratio = pdaysigma/sigma252
#
dateString = Dates.format(now(), "e u d y")
println()
println("Strangle Model")
println("Analyst's name: $(aname)")
println( "Date: $(dateString)")
println( "Stock symbol: $(stosym)")
println("Option chain: $(moname), $(exday)")
println("Days to expiry: $(days)")
println( "Stock price: $(stock)")
println( "Call strike price: $(callstrike)")
println( "Call Bid: $(callbid)")
println( "Call Ask: $(callask)")
println( "Call price: $(callpr)")
println( "Put strike price: $(putstrike)")
println( "Put Bid: $(putbid)")
println( "Put Ask: $(putask)")
println( "Put price: $(putpr)")
println( "Position cost: $(cost)")
if earnings
    println("Days to earnings: $(daystoearn)")
println( "Call breakeven price: $(callbe)")
println( "Call breakeven ln percent: $(callbegr)")
println( "Put breakeven price: $(putbe)")
println( "Put breakeven ln percent: $(putbegr)")
println( "Call value based upon 252-day HV ($(sigma252)) : $(cop252)")
println( "Call value based upon 30-day HV ($(sigma30)) : $(cop30)")
println( "Put value based upon 252-day HV ($(sigma252)) : $(pop252)")
println( "Put value based upon 30-day HV ($(sigma30)) : $(pop30)")
println( "Call implied duration volatility: $(cdursigma)")
println( "Put implied duration volatility: $(pdursigma)")
println("CRITICAL")
println( "Call implied daily volatility: $(cdaysigma)")
println( "Put implied daily volatility: $(pdaysigma)")
println( "Call sigma ratio: $(callsigmaratio)")
println( "Put sigma ratio: $(putsigmaratio)")
#
# Calculate one day time decay using our calculated IDV.
#
if days >= 2.0
    days = days - 1.0
    cdursigma = cdaysigma*fu.durvol(days)
    cprice1d = fu.otranche(stock,callstrike,cdursigma,true)
    pdursigma = pdaysigma*fu.durvol(days)
    pprice1d = fu.otranche(stock,putstrike,pdursigma,false)
    timedecay = (callpr - cprice1d) + (putpr - pprice1d)
    println( "One day time decay: $(timedecay)")
else
    println( "No time decay with one day remaining.")
end
end
#


using LowRank
using Plots

function plot_changing_n(method, method_name, d, k, std, num_times, ns)
    mean_errors = []
    mean_indices_correctnesses = []
    mean_times = []
    baseline_errors = []
    for n in ns
        N = 10*n
        A = LowRank.matrix_with_normal_noise(N, d, n, k, std)
        errors, indicess, times = LowRank.run_method(B -> method(B, n, k), A, num_times)
        baseline_error, _ = LowRank.truncated_svd(A[:,1:n], k)
        baseline_indices = Set(1:n)
        mean_error = sum(errors)/num_times
        mean_indices_correctness = sum([length(intersect(inds,baseline_indices))/n for inds in indicess])/num_times
        mean_time = sum(times)/num_times
        push!(mean_errors, mean_error)
        push!(mean_indices_correctnesses, mean_indices_correctness)
        push!(mean_times, mean_time)
        push!(baseline_errors, baseline_error)
    end

    fontsize = 10
    texts = ["method: $method_name", "number of points N: 10n", "dimension d: $d", "target rank k: $k", "noise std: $std", "number of runs: $num_times"]
    y_locations = 0.9:-0.15:0
    annotations = [(0.2,y_locations[i], text(t, :left, fontsize)) for (i,t) in enumerate(texts)]
    p_text = plot([],[], ann=annotations, border =:none, label=:none)
    p_error = plot()
    plot!(p_error, ns, baseline_errors, label = "baseline errors")
    plot!(p_error, ns, mean_errors, label = "mean errors")
    p_indices = plot(ns, mean_indices_correctnesses, label = "mean indices correctness")
    ylims!(p_indices, (0,1))
    p_time = plot(ns, mean_times, label = "mean time")
    xaxis!(p_error, "size of target set, n")
    xaxis!(p_indices, "size of target set, n")
    xaxis!(p_time, "size of target set, n")
    p = plot(p_text,p_error, p_indices, p_time)
end

method_name = "Global FastLowRank"
method = FastLowRank
d = 10
k = 4
std = 0.1
num_times = 10
ns = 5:5:50
p = plot_changing_n(method, method_name, d, k, std, num_times, ns)
savefig(p, "GlobalFastLowRank.png")
method_name = "Local FastLowRank"
method = (B, n, k) -> FastLowRank(B, n, k, step = LowRank.L2_step_local)
p = plot_changing_n(method, method_name, d, k, std, num_times, ns)
savefig(p, "LocalFastLowRank.png")


using GLVisualize, Colors, GeometryTypes, GLAbstraction, Reactive

if !isdefined(:runtests)
	window = glscreen()
    timesignal = loop(linspace(0f0, 1f0, 360))
end
const N = 87
const range = linspace(-5f0, 5f0, N)

function contour_inner(i, x, y)
    Float32(sin(1.3*x*i)*cos(0.9*y)+cos(.8*x)*sin(1.9*y)+cos(y*.2*x))
end
const data = zeros(Float32, N, N)

function contourdata(t)
    for i=1:size(data, 1)
        @simd for j=1:size(data, 2)
            @inbounds data[i,j] = contour_inner(t, range[i], range[j])
        end
    end
    data
end

heightfield = map(contourdata, timesignal)
mini = Vec3f0(first(range), first(range), minimum(value(heightfield)))
maxi = Vec3f0(last(range), last(range), maximum(value(heightfield)))
barsvis = visualize(
    heightfield,
    scale_x = 0.07,
    scale_y = 0.07,
    color_map=map(RGBA{U8}, colormap("Blues")),
    color_norm=Vec2f0(0,1),
    ranges=(range, range),
    boundingbox=Signal(AABB{Float32}(mini, maxi))
)
_view(barsvis, window)

if !isdefined(:runtests)
	renderloop(window)
end


abstract type VariableType end
abstract type Quantitative <: VariableType end
struct Categorical <: VariableType end
struct Continuous <: Quantitative end
struct Discrete <: Quantitative end

# Variable type of a single value
function vartype(v)
    if is_quantitative(v)
        if is_continuous(v)
            return Continuous
        elseif is_discrete(v)
            return Discrete
        else
            return Quantitative
        end
    elseif is_categorical(v)
        # Not currently making a distinction between nominal/dichotomous or ordinal
        return Categorical
    else
        error("Could not classify variable")
    end
end

# Variable type for a vector
# will check the lesser of the vector length or `searchlen`
function vartype(v::T where T <: AbstractVector, searchlen = 100)
    subarr = view(v, 1:min(length(v), searchlen))

    # Assess
    if any(is_categorical.(subarr))
        return Categorical
    else
        if any(is_continuous.(subarr))
            return Continuous
        elseif any(is_discrete.(subarr))
            return Discrete
        else
            return Quantitative
        end
    end
end

# Variable types for a matrix
function vartype(v::T where T <: AbstractMatrix, dim=1, searchlen = 100)
    if dim !== 1 && dim !== 2
        error("dim must equal 1 or 2")
    end

    len = size(v, dim == 1 ? 2 : 1)
    vartypes = Array{Any, 1}(len)

    @inbounds for i = 1:len
        if dim == 1
            vartypes[i] = vartype(v[:, i], searchlen)
        else
            vartypes[i] = vartype(v[i, :], searchlen)
        end
    end

    return vartypes
end

# Type => Variable Type logic
is_quantitative(v) = typeof(v) <: Number && typeof(v) !== Bool
is_categorical(v) = !is_quantitative(v)
is_continuous(v) = typeof(v) <: AbstractFloat || typeof(v) <: Complex || typeof(v) <: Irrational
is_discrete(v) = typeof(v) <: Int || typeof(v) <: Rational


# K-means algorithm

#### Interface

type KmeansResult{T<:AbstractFloat} <: ClusteringResult
    centers::Matrix{T}         # cluster centers (d x k)
    assignments::Vector{Int}   # assignments (n)
    costs::Vector{T}           # costs of the resultant assignments (n)
    counts::Vector{Int}        # number of samples assigned to each cluster (k)
    cweights::Vector{Float64}  # cluster weights (k)
    totalcost::Float64         # total cost (i.e. objective) (k)
    iterations::Int            # number of elapsed iterations
    converged::Bool            # whether the procedure converged
end

const _kmeans_default_init = :kmpp
const _kmeans_default_maxiter = 100
const _kmeans_default_tol = 1.0e-6
const _kmeans_default_display = :none

function kmeans!{T<:AbstractFloat}(X::Matrix{T}, centers::Matrix{T};
                                   weights=nothing,
                                   maxiter::Integer=_kmeans_default_maxiter,
                                   tol::Real=_kmeans_default_tol,
                                   display::Symbol=_kmeans_default_display,
                                   distance::SemiMetric=SqEuclidean())

    m, n = size(X)
    m2, k = size(centers)
    m == m2 || throw(DimensionMismatch("Inconsistent array dimensions."))
    (2 <= k < n) || error("k must have 2 <= k < n.")

    assignments = zeros(Int, n)
    costs = zeros(T, n)
    counts = Vector{Int}(k)
    cweights = Vector{Float64}(k)

    _kmeans!(X, conv_weights(T, n, weights), centers,
             assignments, costs, counts, cweights,
             round(Int, maxiter), tol, display_level(display), distance)
end

function kmeans{T<:AbstractFloat}(X::Matrix{T}, k::Int;
                weights=nothing,
                init=_kmeans_default_init,
                maxiter::Integer=_kmeans_default_maxiter,
                tol::Real=_kmeans_default_tol,
                display::Symbol=_kmeans_default_display,
                distance::SemiMetric=SqEuclidean())

    m, n = size(X)
    (2 <= k < n) || error("k must have 2 <= k < n.")
    iseeds = initseeds(init, X, k)
    centers = copyseeds(X, iseeds)
    kmeans!(X, centers;
            weights=weights,
            maxiter=maxiter,
            tol=tol,
            display=display,
            distance=distance)
end

#### Core implementation

# core k-means skeleton
function _kmeans!{T<:AbstractFloat}(
    x::Matrix{T},                   # in: sample matrix (d x n)
    w::Union{Void, Vector{T}},      # in: sample weights (n)
    centers::Matrix{T},             # in/out: matrix of centers (d x k)
    assignments::Vector{Int},       # out: vector of assignments (n)
    costs::Vector{T},               # out: costs of the resultant assignments (n)
    counts::Vector{Int},            # out: the number of samples assigned to each cluster (k)
    cweights::Vector{Float64},      # out: the weights of each cluster
    maxiter::Int,                   # in: maximum number of iterations
    tol::Real,                      # in: tolerance of change at convergence
    displevel::Int,                 # in: the level of display
    distance::SemiMetric)             # in: function to calculate the distance with

    # initialize

    k = size(centers, 2)
    to_update = Vector{Bool}(k) # indicators of whether a center needs to be updated
    unused = Int[]
    num_affected::Int = k # number of centers, to which the distances need to be recomputed

    dmat = pairwise(distance, centers, x)
    dmat = convert(Array{T}, dmat) #Can be removed if one day Distance.result_type(SqEuclidean(), T, T) == T
    update_assignments!(dmat, true, assignments, costs, counts, to_update, unused)
    objv = w == nothing ? sum(costs) : dot(w, costs)

    # main loop
    t = 0
    converged = false
    if displevel >= 2
        @printf "%7s %18s %18s | %8s \n" "Iters" "objv" "objv-change" "affected"
        println("-------------------------------------------------------------")
        @printf("%7d %18.6e\n", t, objv)
    end

    while !converged && t < maxiter
        t = t + 1

        # update (affected) centers

        update_centers!(x, w, assignments, to_update, centers, cweights)

        if !isempty(unused)
            repick_unused_centers(x, costs, centers, unused)
        end

        # update pairwise distance matrix

        if !isempty(unused)
            to_update[unused] = true
        end

        if t == 1 || num_affected > 0.75 * k
            pairwise!(dmat, distance, centers, x)
        else
            # if only a small subset is affected, only compute for that subset
            affected_inds = find(to_update)
            dmat_p = pairwise(distance, centers[:, affected_inds], x)
            dmat[affected_inds, :] = dmat_p
        end

        # update assignments

        update_assignments!(dmat, false, assignments, costs, counts, to_update, unused)
        num_affected = sum(to_update) + length(unused)

        # compute change of objective and determine convergence

        prev_objv = objv
        objv = w == nothing ? sum(costs) : dot(w, costs)
        objv_change = objv - prev_objv

        if objv_change > tol
            warn("The objective value changes towards an opposite direction")
        end

        if abs(objv_change) < tol
            converged = true
        end

        # display iteration information (if asked)

        if displevel >= 2
            @printf("%7d %18.6e %18.6e | %8d\n", t, objv, objv_change, num_affected)
        end
    end

    if displevel >= 1
        if converged
            println("K-means converged with $t iterations (objv = $objv)")
        else
            println("K-means terminated without convergence after $t iterations (objv = $objv)")
        end
    end

    return KmeansResult(centers, assignments, costs, counts, cweights,
                        Float64(objv), t, converged)
end

#
#  Updates assignments, costs, and counts based on
#  an updated (squared) distance matrix
#
function update_assignments!{T<:AbstractFloat}(
    dmat::Matrix{T},            # in:  distance matrix (k x n)
    is_init::Bool,              # in:  whether it is the initial run
    assignments::Vector{Int},   # out: assignment vector (n)
    costs::Vector{T},           # out: costs of the resultant assignment (n)
    counts::Vector{Int},        # out: number of samples assigned to each cluster (k)
    to_update::Vector{Bool},    # out: whether a center needs update (k)
    unused::Vector{Int})        # out: the list of centers get no samples assigned to it

    k::Int, n::Int = size(dmat)

    # re-initialize the counting vector
    fill!(counts, 0)

    if is_init
        fill!(to_update, true)
    else
        fill!(to_update, false)
        if !isempty(unused)
            empty!(unused)
        end
    end

    # process each sample
    @inbounds for j = 1 : n

        # find the closest cluster to the i-th sample
        a::Int = 1
        c::T = dmat[1, j]
        for i = 2 : k
            ci = dmat[i, j]
            if ci < c
                a = i
                c = ci
            end
        end

        # set/update the assignment
        if is_init
            assignments[j] = a
        else  # update
            pa = assignments[j]
            if pa != a
                # if assignment changes,
                # both old and new centers need to be updated
                assignments[j] = a
                to_update[a] = true
                to_update[pa] = true
            end
        end

        # set costs and counts accordingly
        costs[j] = c
        counts[a] += 1
    end

    # look for centers that have no associated samples

    for i = 1 : k
        if counts[i] == 0
            push!(unused, i)
            to_update[i] = false # this is handled using different mechanism
        end
    end
end

#
#  Update centers based on updated assignments
#
#  (specific to the case where samples are not weighted)
#
function update_centers!{T<:AbstractFloat}(
    x::Matrix{T},                   # in: sample matrix (d x n)
    w::Void,                        # in: sample weights
    assignments::Vector{Int},       # in: assignments (n)
    to_update::Vector{Bool},        # in: whether a center needs update (k)
    centers::Matrix{T},             # out: updated centers (d x k)
    cweights::Vector)               # out: updated cluster weights (k)

    d::Int = size(x, 1)
    n::Int = size(x, 2)
    k::Int = size(centers, 2)

    # initialize center weights
    for i = 1 : k
        if to_update[i]
            cweights[i] = 0.
        end
    end

    # accumulate columns
    @inbounds for j = 1 : n
        cj = assignments[j]
        1 <= cj <= k || error("assignment out of boundary.")
        if to_update[cj]
            if cweights[cj] > 0
                for i = 1:d
                    centers[i, cj] += x[i, j]
                end
            else
                for i = 1:d
                    centers[i, cj] = x[i, j]
                end
            end
            cweights[cj] += 1
        end
    end

    # sum ==> mean
    for j = 1:k
        if to_update[j]
            @inbounds cj::T = 1 / cweights[j]
            vj = view(centers,:,j)
            for i = 1:d
                @inbounds vj[i] *= cj
            end
        end
    end
end

#
#  Update centers based on updated assignments
#
#  (specific to the case where samples are weighted)
#
function update_centers!{T<:AbstractFloat}(
    x::Matrix{T},                   # in: sample matrix (d x n)
    weights::Vector{T},             # in: sample weights (n)
    assignments::Vector{Int},       # in: assignments (n)
    to_update::Vector{Bool},        # in: whether a center needs update (k)
    centers::Matrix{T},             # out: updated centers (d x k)
    cweights::Vector)               # out: updated cluster weights (k)

    d::Int = size(x, 1)
    n::Int = size(x, 2)
    k::Int = size(centers, 2)

    # initialize center weights
    for i = 1 : k
        if to_update[i]
            cweights[i] = 0.
        end
    end

    # accumulate columns
    # accumulate_cols_u!(centers, cweights, x, assignments, weights, to_update)
    for j = 1 : n
        @inbounds wj = weights[j]

        if wj > 0
            @inbounds cj = assignments[j]
            1 <= cj <= k || error("assignment out of boundary.")

            if to_update[cj]
                rj = view(centers, :, cj)
                xj = view(x, :, j)
                if cweights[cj] > 0
                    for i = 1:d
                        @inbounds rj[i] += xj[i] * wj
                    end
                else
                    for i = 1:d
                        @inbounds rj[i] = xj[i] * wj
                    end
                end
                cweights[cj] += wj
            end
        end
    end

    # sum ==> mean
    for j = 1:k
        if to_update[j]
            @inbounds cj::T = 1 / cweights[j]
            vj = view(centers,:,j)
            for i = 1:d
                @inbounds vj[i] *= cj
            end
        end
    end
end

#
#  Re-picks centers that get no samples assigned to them.
#
function repick_unused_centers{T<:AbstractFloat}(
    x::Matrix{T},           # in: the sample set (d x n)
    costs::Vector{T},       # in: the current assignment costs (n)
    centers::Matrix{T},     # to be updated: the centers (d x k)
    unused::Vector{Int})    # in: the set of indices of centers to be updated

    # pick new centers using a scheme like kmeans++
    ds = similar(costs)
    tcosts = copy(costs)
    n = size(x, 2)

    for i in unused
        j = wsample(1:n, tcosts)
        tcosts[j] = 0
        v = x[:,j]
        centers[:,i] = v

        colwise!(ds, SqEuclidean(), v, x)
        tcosts = min(tcosts, ds)
    end
end


# In the various algorithms we store the tolerance as a generic Real in order to
# construct these objects using keyword arguments in a type-stable manner. At
# the beginning of the corresponding algorithm, the actual tolerance will be
# converted to the concrete numeric type appropriate for the problem at hand

# Orthogonalization and orthonormalization
"""
    abstract type Orthogonalizer

Supertype of a hierarchy for representing different orthogonalization strategies or
algorithms.

See also: [`ClassicalGramSchmidt`](@ref), [`ModifiedGramSchmidt`](@ref),
[`ClassicalGramSchmidt2`](@ref), [`ModifiedGramSchmidt2`](@ref),
[`ClassicalGramSchmidtIR`](@ref), [`ModifiedGramSchmidtIR`](@ref).
"""
abstract type Orthogonalizer end
abstract type Reorthogonalizer <: Orthogonalizer end

# Simple
"""
    ClassicalGramSchmidt()

Represents the classical Gram Schmidt algorithm for orthogonalizing different vectors,
typically not an optimal choice.
"""
struct ClassicalGramSchmidt <: Orthogonalizer
end

"""
    ModifiedGramSchmidt()

Represents the modified Gram Schmidt algorithm for orthogonalizing different vectors,
typically a reasonable choice for linear systems but not for eigenvalue solvers with a
large Krylov dimension.
"""
struct ModifiedGramSchmidt <: Orthogonalizer
end

# A single reorthogonalization always
"""
    ClassicalGramSchmidt2()

Represents the classical Gram Schmidt algorithm with a second reorthogonalization step
always taking place.
"""
struct ClassicalGramSchmidt2 <: Reorthogonalizer
end

"""
    ModifiedGramSchmidt2()

Represents the modified Gram Schmidt algorithm with a second reorthogonalization step
always taking place.
"""
struct ModifiedGramSchmidt2 <: Reorthogonalizer
end

# Iterative reorthogonalization
"""
    ClassicalGramSchmidtIR(η::Real = 1/sqrt(2))

Represents the classical Gram Schmidt algorithm with iterative (i.e. zero or more)
reorthogonalization until the norm of the vector after an orthogonalization step has not
decreased by a factor smaller than `η` with respect to the norm before the step. The
default value corresponds to the Daniel-Gragg-Kaufman-Stewart condition.
"""
struct ClassicalGramSchmidtIR{S<:Real} <: Reorthogonalizer
    η::S
end
ClassicalGramSchmidtIR() = ClassicalGramSchmidtIR(1/sqrt(2)) # Daniel-Gragg-Kaufman-Stewart

"""
    ModifiedGramSchmidtIR(η::Real = 1/sqrt(2))

Represents the modified Gram Schmidt algorithm with iterative (i.e. zero or more)
reorthogonalization until the norm of the vector after an orthogonalization step has not
decreased by a factor smaller than `η` with respect to the norm before the step. The
default value corresponds to the Daniel-Gragg-Kaufman-Stewart condition.
"""
struct ModifiedGramSchmidtIR{S<:Real} <: Reorthogonalizer
    η::S
end
ModifiedGramSchmidtIR() = ModifiedGramSchmidtIR(1/sqrt(2)) # Daniel-Gragg-Kaufman-Stewart

# Solving eigenvalue problems
abstract type KrylovAlgorithm end

# General purpose; good for linear systems, eigensystems and matrix functions
"""
    Lanczos(; krylovdim = KrylovDefaults.krylovdim, maxiter = KrylovDefaults.maxiter,
        tol = KrylovDefaults.tol, orth = KrylovDefaults.orth, eager = false, verbosity = 0)

Represents the Lanczos algorithm for building the Krylov subspace; assumes the linear
operator is real symmetric or complex Hermitian. Can be used in `eigsolve` and
`exponentiate`. The corresponding algorithms will build a Krylov subspace of size at most
`krylovdim`, which will be repeated at most `maxiter` times and will stop when the norm of
the residual of the Lanczos factorization is smaller than `tol`. The orthogonalizer `orth`
will be used to orthogonalize the different Krylov vectors. Eager mode, as selected by
`eager = true`, means that the algorithm that uses this Lanczos process (e.g. `eigsolve`)
can try to finish its computation before the total Krylov subspace of dimension `krylovdim`
is constructed. Default verbosity level `verbosity` is zero, meaning that no output will be
printed.

Use `Arnoldi` for non-symmetric or non-Hermitian linear operators.

See also: `factorize`, `eigsolve`, `exponentiate`, `Arnoldi`, `Orthogonalizer`
"""
struct Lanczos{O<:Orthogonalizer, S<:Real} <: KrylovAlgorithm
    orth::O
    krylovdim::Int
    maxiter::Int
    tol::S
    eager::Bool
    verbosity::Int
end
Lanczos(; krylovdim::Int = KrylovDefaults.krylovdim,
            maxiter::Int = KrylovDefaults.maxiter,
            tol::Real = KrylovDefaults.tol,
            orth::Orthogonalizer = KrylovDefaults.orth,
            eager::Bool = false,
            verbosity::Int = 0) =
    Lanczos(orth, krylovdim, maxiter, tol, eager, verbosity)

"""
    GKL(; krylovdim = KrylovDefaults.krylovdim, maxiter = KrylovDefaults.maxiter,
        tol = KrylovDefaults.tol, orth = KrylovDefaults.orth, verbosity = 0)

Represents the Golub-Kahan-Lanczos bidiagonalization algorithm for sequentially building a
Krylov-like factorization of a general matrix or linear operator with a bidiagonal reduced
matrix. Can be used in `svdsolve`. The corresponding algorithm builds a Krylov subspace of
size at most `krylovdim`, which will be repeated at most `maxiter` times and will stop when
the norm of the residual of the Arnoldi factorization is smaller than `tol`. The
orthogonalizer `orth` will be used to orthogonalize the different Krylov vectors. Default
verbosity level `verbosity` is zero, meaning that no output will be printed.

See also: `svdsolve`, `Orthogonalizer`
"""
struct GKL{O<:Orthogonalizer, S<:Real} <: KrylovAlgorithm
    orth::O
    krylovdim::Int
    maxiter::Int
    tol::S
    eager::Bool
    verbosity::Int
end
GKL(; krylovdim::Int = KrylovDefaults.krylovdim,
        maxiter::Int = KrylovDefaults.maxiter,
        tol::Real = KrylovDefaults.tol,
        orth::Orthogonalizer = KrylovDefaults.orth,
        eager::Bool = false,
        verbosity::Int = 0) =
    GKL(orth, krylovdim, maxiter, tol, eager, verbosity)

"""
    Arnoldi(; krylovdim = KrylovDefaults.krylovdim, maxiter = KrylovDefaults.maxiter,
        tol = KrylovDefaults.tol, orth = KrylovDefaults.orth, verbosity = 0)

Represents the Arnoldi algorithm for building the Krylov subspace for a general matrix or
linear operator. Can be used in `eigsolve` and `exponentiate`. The corresponding algorithms
will build a Krylov subspace of size at most `krylovdim`, which will be repeated at most
`maxiter` times and will stop when the norm of the residual of the Arnoldi factorization is
smaller than `tol`. The orthogonalizer `orth` will be used to orthogonalize the different
Krylov vectors. Eager mode, as selected by `eager = true`, means that the algorithm that
uses this Arnoldi process (e.g. `eigsolve`) can try to finish its computation before the
total Krylov subspace of dimension `krylovdim` is constructed. Default verbosity level
`verbosity` is zero, meaning that no output will be printed.

Use `Lanczos` for real symmetric or complex Hermitian linear operators.

See also: [`eigsolve`](@ref), [`exponentiate`](@ref), [`Lanczos`](@ref),
[`Orthogonalizer`](@ref)
"""
struct Arnoldi{O<:Orthogonalizer, S<:Real} <: KrylovAlgorithm
    orth::O
    krylovdim::Int
    maxiter::Int
    tol::S
    eager::Bool
    verbosity::Int
end
Arnoldi(; krylovdim::Int = KrylovDefaults.krylovdim,
            maxiter::Int = KrylovDefaults.maxiter,
            tol::Real = KrylovDefaults.tol,
            orth::Orthogonalizer = KrylovDefaults.orth,
            eager::Bool = false,
            verbosity::Int = 0) =
    Arnoldi(orth, krylovdim, maxiter, tol, eager, verbosity)

"""
    GolubYe(; krylovdim = KrylovDefaults.krylovdim, maxiter = KrylovDefaults.maxiter,
        tol = KrylovDefaults.tol, orth = KrylovDefaults.orth, verbosity = 0)

Represents the Golub-Ye algorithm for solving hermitian (symmetric) generalized eigenvalue
problems `A x = λ B x` with positive definite `B`, without the need for inverting `B`.
Builds a Krylov subspace of size `krylovdim` starting from an estimate `x` by acting with
`(A - ρ(x) B)`, where `ρ(x) = dot(x, A*x)/dot(x, B*x)`, and employing the Lanczos
algorithm. This process is repeated at most `maxiter` times. In every iteration `k>1`, the
subspace will also be expanded to size `krylovdim+1` by adding ``x_k - x_{k-1}``, which is
known as the LOPCG correction and was suggested by Money and Ye. With `krylovdim = 2`, this
algorithm becomes equivalent to `LOPCG`.
"""
struct GolubYe{O<:Orthogonalizer, S<:Real} <: KrylovAlgorithm
    orth::O
    krylovdim::Int
    maxiter::Int
    tol::S
    verbosity::Int
end
GolubYe(; krylovdim::Int = KrylovDefaults.krylovdim,
            maxiter::Int = KrylovDefaults.maxiter,
            tol::Real = KrylovDefaults.tol,
            orth::Orthogonalizer = KrylovDefaults.orth,
            verbosity::Int = 0) =
    GolubYe(orth, krylovdim, maxiter, tol, verbosity)

# Solving linear systems specifically
abstract type LinearSolver <: KrylovAlgorithm end

"""
    CG(; maxiter = KrylovDefaults.maxiter, tol = KrylovDefaults.tol)

Construct an instance of the conjugate gradient algorithm with specified parameters, which
can be passed to `linsolve` in order to iteratively solve a linear system with a positive
definite (and thus symmetric or hermitian) coefficient matrix or operator. The `CG` method
will search for the optimal `x` in a Krylov subspace of maximal size `maxiter`, or stop when
`norm(A*x - b) < tol`. Default verbosity level `verbosity` is zero, meaning that no output
will be printed.

See also: [`linsolve`](@ref), [`MINRES`](@ref), [`GMRES`](@ref), [`BiCG`](@ref),
[`BiCGStab`](@ref)
"""
struct CG{S<:Real} <: LinearSolver
    maxiter::Int
    tol::S
    verbosity::Int
end
CG(; maxiter::Integer = KrylovDefaults.maxiter,
        tol::Real = KrylovDefaults.tol,
        verbosity::Int = 0) =
    CG(maxiter, tol, verbosity)

"""
    GMRES(; krylovdim = KrylovDefaults.krylovdim, maxiter = KrylovDefaults.maxiter,
        tol = KrylovDefaults.tol, orth::Orthogonalizer = KrylovDefaults.orth)

Construct an instance of the GMRES algorithm with specified parameters, which can be passed
to `linsolve` in order to iteratively solve a linear system. The `GMRES` method will search
for the optimal `x` in a Krylov subspace of maximal size `krylovdim`, and repeat this
process for at most `maxiter` times, or stop when `norm(A*x - b) < tol`. In building the
Krylov subspace, `GMRES` will use the orthogonalizer `orth`. Default verbosity level
`verbosity` is zero, meaning that no output will be printed.

Note that in the traditional nomenclature of `GMRES`, the parameter `krylovdim` is referred
to as the restart parameter, and `maxiter` is the number of outer iterations, i.e. restart
cycles. The total iteration count, i.e. the number of expansion steps, is roughly
`krylovdim` times the number of iterations.

See also: [`linsolve`](@ref), [`BiCG`](@ref), [`BiCGStab`](@ref), [`CG`](@ref),
[`MINRES`](@ref)
"""
struct GMRES{O<:Orthogonalizer,S<:Real} <: LinearSolver
    orth::O
    maxiter::Int
    krylovdim::Int
    tol::S
    verbosity::Int
end
GMRES(; krylovdim::Integer = KrylovDefaults.krylovdim,
        maxiter::Integer = KrylovDefaults.maxiter,
        tol::Real = KrylovDefaults.tol,
        orth::Orthogonalizer = KrylovDefaults.orth,
        verbosity::Int = 0) =
    GMRES(orth, maxiter, krylovdim, tol, verbosity)

# TODO
"""
    MINRES(; maxiter = KrylovDefaults.maxiter, tol = KrylovDefaults.tol)

Construct an instance of the conjugate gradient algorithm with specified parameters, which
can be passed to `linsolve` in order to iteratively solve a linear system with a real
symmetric or complex hermitian coefficient matrix or operator. The `MINRES` method will
search for the optimal `x` in a Krylov subspace of maximal size `maxiter`, or stop when
`norm(A*x - b) < tol`. In building the Krylov subspace, `MINRES` will use the
orthogonalizer `orth`. Default verbosity level `verbosity` is zero, meaning that no output
will be printed.

!!! warning "Not implemented yet"

See also: [`linsolve`](@ref), [`CG`](@ref), [`GMRES`](@ref), [`BiCG`](@ref), [`BiCGStab`](@ref)
"""
struct MINRES{S<:Real} <: LinearSolver
    maxiter::Int
    tol::S
    verbosity::Int
end
MINRES(; maxiter::Integer = KrylovDefaults.maxiter,
            tol::Real = KrylovDefaults.tol,
            verbosity::Int = 0) =
    MINRES(maxiter, tol, verbosity)

"""
    BiCG(; maxiter = KrylovDefaults.maxiter, tol = KrylovDefaults.tol)

Construct an instance of the Biconjugate gradient algorithm with specified parameters, which
can be passed to `linsolve` in order to iteratively solve a linear system general linear
map, of which the adjoint can also be applied. The `BiCG` method will search for the
optimal `x` in a Krylov subspace of maximal size `maxiter`, or stop when `norm(A*x - b) <
tol`. Default verbosity level `verbosity` is zero, meaning that no output will be printed.

!!! warning "Not implemented yet"

See also: [`linsolve`](@ref), [`BiCGStab`](@ref), [`GMRES`](@ref), [`CG`](@ref),
[`MINRES`](@ref)
"""
struct BiCG{S<:Real} <: LinearSolver
    maxiter::Int
    tol::S
    verbosity::Int
end
BiCG(; maxiter::Integer = KrylovDefaults.maxiter,
        tol::Real = KrylovDefaults.tol,
        verbosity::Int = 0) =
    BiCG(maxiter, tol, verbosity)

"""
    BiCGStab(; maxiter = KrylovDefaults.maxiter, tol = KrylovDefaults.tol)

Construct an instance of the Biconjugate gradient algorithm with specified parameters, which
can be passed to `linsolve` in order to iteratively solve a linear system general linear
map. The `BiCGStab` method will search for the optimal `x` in a Krylov subspace of maximal
size `maxiter`, or stop when `norm(A*x - b) < tol`. Default verbosity level `verbosity` is
zero, meaning that no output will be printed.

!!! warning "Not implemented yet"

See also: [`linsolve`](@ref), [`BiCG`](@ref), [`GMRES`](@ref), [`CG`](@ref),
[`MINRES`](@ref)
"""
struct BiCGStab{S<:Real} <: LinearSolver
    maxiter::Int
    tol::S
    verbosity::Int
end
BiCGStab(; maxiter::Integer = KrylovDefaults.maxiter,
            tol::Real = KrylovDefaults.tol,
            verbosity::Int = 0) =
    BiCGStab(maxiter, tol, verbosity)

# Solving eigenvalue systems specifically
abstract type EigenSolver <: KrylovAlgorithm end

struct JacobiDavidson <: EigenSolver
end

# Default values
"""
    module KrylovDefaults
        const orth = KrylovKit.ModifiedGramSchmidtIR()
        const krylovdim = 30
        const maxiter = 100
        const tol = 1e-12
    end

A module listing the default values for the typical parameters in Krylov based algorithms:
*   `orth`: the orthogonalization routine used to orthogonalize the Krylov basis in the
    `Lanczos` or `Arnoldi` iteration
*   `krylovdim`: the maximal dimension of the Krylov subspace that will be constructed
*   `maxiter`: the maximal number of outer iterations, i.e. the maximum number of times the
    Krylov subspace may be rebuilt
*   `tol`: the tolerance to which the problem must be solved, based on a suitable error
    measure, e.g. the norm of some residual.
!!! warning
    The default value of `tol` is a `Float64` value, if you solve problems in `Float32` or
    `ComplexF32` arithmetic, you should always specify a new `tol` as the default value
    will not be attainable.
"""
module KrylovDefaults
    using ..KrylovKit
    const orth = KrylovKit.ModifiedGramSchmidt2() # conservative choice
    const krylovdim = 30
    const maxiter = 100
    const tol = 1e-12
end


"""
Colorscheme: phase plot with conformal polar grid
"""
function cs_c(; colormap=hsv_colors(), phaseres=20)
    stepfct = generate_stepfct(length(colormap))
    sawfct1 = generate_sawfct(1.0/phaseres, 0.75, 1.0)
    sawfct2 = generate_sawfct(2*π/phaseres, 0.75, 1.0)

    let cm = RGB.(colormap)
      return function(z, fz)
          anglefz = angle_zero_one(fz)
          absfz = abs(fz)
          factor = sawfct1(anglefz) * (
            absfz == 0.0 ? 0.75 : absfz == Inf ? 1.0 : sawfct2(log(absfz)))
          color = cm[stepfct(anglefz)]
          return brightenRGB(
            RGB(factor*color.r, factor*color.g, factor*color.b), 0.1)
      end
    end
end

"""
Colorscheme: colored phase plot with module jumps
"""
function cs_m(; colormap=hsv_colors(), logabsres=20)
    sawfct = generate_sawfct(2*π/logabsres, 0.7, 1.0)
    stepfct = generate_stepfct(length(colormap))
    let cm = RGB.(colormap)
      return function(z, fz)
        absfz = abs(fz)
        factor = absfz == 0.0 ? 0.7 :
                 absfz == Inf ? 1.0 : sawfct(log(absfz))
        color = cm[stepfct(angle_zero_one(fz))]
        return RGB(factor*color.r, factor*color.g, factor*color.b)
      end
    end
end

"""
Colorscheme: enhanced domain coloring
"""
cs_e(; colormap=hsv_colors(), logabsres=20) = cs_mod_abs_brightness(
    cs_m(colormap=colormap, logabsres=logabsres))

"""
Colorscheme: colored phase plot with specific phase jumps
"""
function cs_j(; colormap=hsv_colors(),
                jumps=collect(LinRange(-π, π, 21))[1:20], a=0.8, b=1.0)
    sawfct = generate_sawfctx(
      angle_zero_one.(exp.(1im .* jumps)); a=a, b=b)
    stepfct = generate_stepfct(length(colormap))
    let cm = RGB.(colormap)
      return function(z, fz)
          angle_fz = angle_zero_one(fz)
          factor = sawfct(angle_fz)
          color = cm[stepfct(angle_fz)]
          return RGB(factor*color.r, factor*color.g, factor*color.b)
      end
    end
end

"""
Colorscheme: proper phase plot
"""
function cs_p(; colormap=hsv_colors())
    stepfct = generate_stepfct(length(colormap))
    let cm = copy(colormap)
      return (z, fz) -> cm[stepfct(angle_zero_one(fz))]
    end
end

"""
Colorscheme: phase plot colored in steps
"""
cs_q(; colormap=hsv_colors(20)) = cs_p(; colormap=colormap)

"""
Colorscheme: standard domain coloring
"""
cs_d(; colormap=hsv_colors()) = cs_mod_abs_brightness(
  cs_p(colormap=colormap))

"""
Colorscheme: for exams; cs_j with NIST colors and more pronounced shadow jumps.
"""
cs_exam(; colormap=nist_colors(), jumps=collect(LinRange(-π, π, 21))[1:20],
          a=0.6, b=1.0) = cs_j(colormap=colormap, jumps=jumps, a=a, b=b)

"""
Colorscheme: Stripes in (real(z), imag(z)) plane

Directions is a vector of complex numbers.
For each such number d in directions the stripes are perpendicular
to (real(d), imag(d)). And abs(d) is the number of stripes per unit
length.
This devides the complex plane in two sets Sd and ℂ\\Sd.

Each z ∊ ℂ has n=length(directions) flags (z ∊ S1, z ∊ S1, ..., z ∊ Sn).
This make 2^n possible flag-combinations. For every combination there
need to be a color in colors.
"""
function cs_stripes(directions, colors)
  let no_dirs = length(directions)
    if 2^no_dirs != length(colors)
      error("for $(no_dirs) directions $(2^no_dirs) colors are needed.")
    end
    rectfct = generate_rectfct(1.0)
    return function(z, fz)
      color_index = 0
      color_pos = 1
      for dir in directions
        color_index += color_pos*Int(rectfct(real(fz*dir)))
        color_pos *= 2
      end
      return colors[1 + color_index]
    end
  end
end

"""
Colorscheme: Stripes in (angle(z), abs(z)) plane

Use `cs_stripes` for (angle(z), abs(z)).
"""
function cs_angle_abs_stripes(directions, colors)
  col_fct = cs_stripes(directions, colors)
  return (z, fz) -> col_fct(z,angle(fz) + 1im*abs(fz))
end

"""
Colorscheme: uniform low-brightness grid lines in real- and imag-part and
             low saturation grid lines in lowsat(abs(fz)).
"""
function cs_gridReIm_lowsat(; reim_a=-9.0, reim_b=0.5, abs_a=-9.0, abs_b=0.8,
                              lowsat=absfz -> log(exp(0.5)*absfz),
                              vcorr=(s,v) -> max(1.0 - s, v))
  sp_reim = generate_spike(a=reim_a, b=reim_b)
  sp_abs = generate_spike(a=abs_a, b=abs_b)
  return function(z, fz)
    absfz = abs(fz)
    h = nist_trafo(4*angle_zero_one(fz))
    s = absfz == 0.0 ? 1.0 : absfz == Inf ? 0.0 : sp_abs(lowsat(absfz))
    v = absfz == 0.0 ? 0.0 :
        absfz == Inf ? 1.0 : min(sp_reim(real(fz)-0.5), sp_reim(imag(fz)-0.5))
    return HSV(h, s, vcorr(s,v))
  end
end

"""
Colorscheme: uniform low-brightness grid lines in real- and imag-part and
             low saturation grid lines in log(abs(fz)).
"""
cs_gridReIm_logabs(; reim_a=-9.0, reim_b=0.5, abs_a=-9.0, abs_b=0.8,
                     vcorr=(s,v) -> max(1.0 - s, v)) = cs_gridReIm_lowsat(
  reim_a=reim_a, reim_b=reim_b, abs_a=abs_a, abs_b=abs_b, vcorr=vcorr,
  lowsat=absfz -> log(exp(0.5)*absfz))

"""
Colorscheme: uniform low-brightness grid lines in real- and imag-part and
             low saturation grid lines in abs(fz).
"""
cs_gridReIm_abs(; reim_a=-9.0, reim_b=0.5, abs_a=-9.0, abs_b=0.8,
                  vcorr=(s,v) -> max(1.0 - s, v)) = cs_gridReIm_lowsat(
  reim_a=reim_a, reim_b=reim_b, abs_a=abs_a, abs_b=abs_b, vcorr=vcorr,
  lowsat=absfz -> absfz+0.5)

"""
Colorscheme: uniform low-brightness grid lines in real- and imag-part.
"""
cs_gridReIm(; reim_a=-9.0, reim_b=0.5) = cs_gridReIm_abs(
  reim_a=reim_a, reim_b=reim_b, abs_a=-Inf, vcorr=(s,v) -> v)

"""
Colorscheme: use the pixels of an image as color source.
"""
function cs_useImage(image;
                      img_upperleft=0.0 + 1.0im, img_lowerright=1.0 + 0.0im,
                      is_hv_periodic=(true, true))
  let (img_height, img_width) = size(image),
      x_left = real(img_upperleft),
      x_width = real(img_lowerright) - x_left,
      y_top = imag(img_upperleft),
      y_height = y_top - imag(img_lowerright),
      epsilon = eps()

      width_fct = generate_stepfct(img_width)
      height_fct = generate_stepfct(img_height)

      return function (z,fz)
        x = (real(fz) - x_left)/x_width
        y = (y_top - imag(fz))/y_height
        if !is_hv_periodic[1]
          x = min(max(0.0, x), 1.0-epsilon)
        end
        if !is_hv_periodic[2]
          y = min(max(0.0, y), 1.0-epsilon)
        end
        return image[height_fct(y), width_fct(x)]
      end
  end
end

# vim:syn=julia:cc=79:fdm=indent:sw=2:ts=2:


module PosteriorSummary

export effective_sample_size,
       hpd_interval

import MCMCDiagnostics.effective_sample_size

function hpd_interval(x::AbstractVector{Float64}; alpha::Float64 = 0.05, sorted::Bool = false)
    n = length(x)
    span = Int(round((1.0 - alpha) * n))
    smallest_rng = Inf
    rng_upper = Inf
    rng_lower = -Inf
    sx::AbstractVector{Float64} = sorted ? x : sort(x)
    lower = 0
    for upper = span:n
        lower = upper - span + 1
        rng = sx[upper] - sx[lower]
        if rng < smallest_rng
            smallest_rng = rng
            rng_lower = sx[lower]
            rng_upper = sx[upper]
        end
    end

    return (lower = rng_lower, upper = rng_upper)
end

end