# diversity / dissimilarity index calculation and plots:

using JLD2, DataFrames, CSV, Dates, Query
using LinearAlgebra, Clustering
using MAT
#using StatsPlots
using Gadfly

include("/home/kristen/Documents/coda_utilities/coda_ops_transforms.jl")

#load in the data:
@load "/home/kristen/Documents/V6V8_analysis/analysis_products/oligo_df_zero_plus.jld2"

## compute distance of each sample to all others, and plot by season:

y=closure(convert(Array,oligo_df_plus[!,Between(:csO1,:csO6)])) #rows are samples, columns are components
m=size(y,1)
oligo_dist = zeros(Float64,m,m)#Array{Float64,2}(undef,m,m)

for i=1:m
    for j=i:m
       oligo_dist[i,j] = ait_dist(y[i,:],y[j,:])
       oligo_dist[j,i] = ait_dist(y[i,:],y[j,:])
    end
end

## hmmm...perhaps, easier to organize by yearday?

ind=sortperm(dayofyear.(oligo_df_plus[!,:date])) #rev=true) #returns indicies...
# dayofyear.(oligo_df_plus[ind,:date]) #to check....
y_yrdy = y[ind,:]
oligo_byday = zeros(Float64,m,m)#Array{Float64,2}(undef,m,m)

for i=1:m
    for j=i:m
       oligo_byday[i,j] = ait_dist(y_yrdy[i,:],y_yrdy[j,:])
       oligo_byday[j,i] = ait_dist(y_yrdy[i,:],y_yrdy[j,:])
    end
end

## NEED TO WORK ON COLOR CODING, ETC...
file = matopen("/home/kristen/Documents/V6V8_analysis/paper_V6V8/matlab_figure_scripts/exports_from_julia/dissimilarity_matrix.mat", "w")
write(file, "oligo_dissim", oligo_dist)
write(file, "oligo_dissim_yrdy", oligo_byday)
write(file,"ind", ind)
write(file,"sample_yearday", dayofyear.(oligo_df_plus[!,:date]))
close(file)

## easier in a dataframe for plotting?
y=closure(convert(Array,oligo_df_plus[!,Between(:csO1,:csO6)])) #rows are samples, columns are components
m=size(y,1)
df_dist = DataFrame(x=Int64[], y=Int64[], dist=Float64[], date=DateTime[])

for i=1:m
    for j=i:m
        push!(df_dist,[i,j,ait_dist(y[i,:],y[j,:]),oligo_df_plus[i,:date]])
    end
end

## #yes - definitely easier for plotting! Very cool!

df_dist.yearday=dayofyear.(df_dist[!,:date])
plot(sort(df_dist,:yearday),x="x",y="y",color="dist", Geom.rectbin) #oops, no - this gives the same thing!

#CSV.write("/home/kristen/Documents/V6V8_analysis/paper_V6V8/matlab_figure_scripts/exports_from_julia/etc...", clade_array, delim = ',')

#what about a dendrogram?
#visualize! but must be with Plots.jl!
hc = hclust(Symmetric(oligo_dist), linkage=:single, uplo=:U)
#-> a few other linkage options, check hclust help

w=plot(hc,xticks=false,linewidth=2)
#node_labels=["O1-I","O2-CB5","O3-II","O4-III/IV","O5-I","O6-I*"] #oligotypes 1-6
#node_labels[hc.order]
#xticks!(w,([1,2,3,4,5,6],node_labels[hc.order]),fontsize=16) #maps from hc and matching up oligos to clades
#ylabel!(w,"Aitchision distance")
#savefig(w, "/home/kristen/Documents/V6V8_analysis/paper_V6V8/V6V8_Syn_CoDa_paper/figures/dendrogram_single.pdf")

## very cool! now, let's check out the aitchsion square norm for inqueality:

y=closure(convert(Array,oligo_df_plus[!,Between(:csO1,:csO6)])) #rows are samples, columns are components
m=size(y,1)
oligo_ineq = zeros(Float64,m,1)#Array{Float64,2}(undef,m,m)

for i=1:m
    oligo_ineq[i] = square_ait_norm(y[i,:])
end

## export it!
file = matopen("/home/kristen/Documents/V6V8_analysis/paper_V6V8/matlab_figure_scripts/exports_from_julia/ineq_metric.mat", "w")
write(file, "oligo_ineq", oligo_ineq)
rr=1 .- exp.(-oligo_ineq)
write(file,"scaled_ineq", rr) #should already ahve the smaples dates somewhere else...
close(file)

## plot it up!

Gadfly.set_default_plot_size(30cm,10cm)

xticks=DateTime("2010-01-01"):Year(1):DateTime("2019-01-01")
xyearticks = Dates.value.(xticks)
xyearticks_label=map(x-> Dates.format(x,"yyyy"),xticks)
labels = Dict(zip(xyearticks, xyearticks_label))

fontsize=Theme(minor_label_font_size=16pt,major_label_font_size=16pt,plot_padding=[5mm, 5mm, 2mm, 2mm])

l1=layer(x=Dates.value.(oligo_df_plus[!,:date]),y=oligo_ineq,Geom.line,Geom.point)
p2=plot(l1,Guide.xticks(ticks=xyearticks),Guide.xlabel(""),Scale.x_continuous(labels = x -> labels[x]),fontsize)

## can also scale it from 0-1:

rr=1 .- exp.(-oligo_ineq)
l2=layer(x=Dates.value.(oligo_df_plus[!,:date]),y=rr,Geom.line,Geom.point)
p2=plot(l2,Guide.xticks(ticks=xyearticks),Guide.xlabel(""),Scale.x_continuous(labels = x -> labels[x]),fontsize)

        #Guide.yticks(ticks=[0 0.25 0.5 0.75 1])
        # Coord.Cartesian(ymin=0, ymax=1),
        # Guide.ylabel("Relative abundance"),fontsize,
        # Guide.annotation(compose(context(), text(Dates.value.(DateTime(2010,1,10)), 0.9, "B"))))


"""
    RobertsBL1D(delta::FT, h::FT) where FT

The β for a given boundary layer
as described in section 5.6 (page 333) of
Computational Fluid Mechanics and Heat Transfer,
2nd edition, J. Tannehill et al.

# Input
```
hmin     = wall boundary (minimum value)
hmax     = wall boundary (maximum value)
delta    = thickness of boundary layer
```

The `if` statement protects against the case when
β = infinity.
"""
function RobertsBL1D(delta::FT, h::FT) where FT
  if delta < h*FT(99/100)
    β = (FT(1) - delta/h)^(-FT(1/2))
  else
    β = FT(1000)
  end
  return β
end

RobertsBL(delta::FT, hmin::FT, hmax::FT) where FT = RobertsBL1D(delta, hmax-hmin)

function Hartmann_BL_1D(Ha::FT, hmin::FT, hmax::FT) where FT
  @assert Ha>eps(FT)
  return RobertsBL((hmax-hmin)/Ha, hmin, hmax)
end

function Reynolds_BL_1D(Re::FT, hmin::FT, hmax::FT) where FT
  @assert Re>eps(FT)
  β = RobertsBL((hmax-hmin)/sqrt(Re), hmin, hmax)
end

function HartmannBL(Ha::FT, hmin::Vector{FT}, hmax::Vector{FT}) where FT
  @assert Ha>eps(FT)
  @assert length(hmin)==length(hmax)
  return [RobertsBL((hmax[i]-hmin[i])/Ha, hmin[i],hmax[i]) for i in 1:3]
end

function Re_Ha_BL_1D(Re::FT, Ha::FT, hmin::FT, hmax::FT) where FT
 temp1 = Reynolds_BL_1D(Re, hmin, hmax)
 temp2 = Hartmann_BL_1D(Ha, hmin, hmax)
 return min(temp1, temp2)
end

function ReynoldsBL(Re::FT, hmin::Vector{FT}, hmax::Vector{FT}) where FT
  @assert Re>eps(FT)
  β = [RobertsBL((hmax[i] - hmin[i])/sqrt(Re), hmin[i], hmax[i]) for i in 1:3]
end

function Re_Ha_BL(Re::FT, Ha::FT, hmin::Vector{FT}, hmax::Vector{FT}) where FT
  @assert length(hmin)==length(hmax)
  temp1 = ReynoldsBL(Re, hmin, hmax)
  temp2 = HartmannBL(Ha, hmin, hmax)
  β = [min(temp1[i],temp2[i]) for i in 1:length(hmin)]
end


# # Stan Economics Tutorial
# http://nbviewer.jupyter.org/github/QuantEcon/QuantEcon.notebooks/blob/master/IntroToStan_basics_workflow.ipynb

using Distributions
using Parameters
using DynamicHMC
using ContinuousTransformations
using MCMCDiagnostics
using DiffWrappers
using PyPlot

# Generate the data from a more complex model
function dgp_rng(X::Matrix, β::Vector, ν, σ)
    y = fill(0.0, size(X, 1))
    for n in 1:size(X, 1)
        #TODO: do a pull request to add the non-centered version
        y[n] = σ * (rand(TDist(ν))) + sum(X[n, :] .* β)
    end
    return y
end

#TODO: add random data
# Number of observations
N = 1000
# Number of covariates
P = 10
X = rand(Normal(), N, P)

# Degrees of freedom for Cauchy noise (\nu)
ν = 5.0
# scale parameter
σ_true = 5.0
# generate some random coefficients that we'll try to recover
β_true = rand(Normal(), 10)
# Make sure the first element of  eta is positive as in our chosen DGP
β_true[1] = abs(β_true[1])
# fixed values
β_true = [1.16, -0.19, -0.29, -0.39, 0.7, -1.6, 0.77, -1.29, 0.89, -0.17]

y_sim = dgp_rng(X, β_true, ν, σ_true)
begin
    figure()
    plt[:hist](y_sim, bins = 30)
end

struct EcoRegProblem
    X::Matrix{Float64}
    data::Vector{Float64}
end

function (prob::EcoRegProblem)(θ)
    β, σ = θ

    # Parameter Priors
    βp = Normal(0, 5)
    σp = Cauchy(0, 2.5)
    #β1, βrest, σ = θ
    #β = vcat(β1, βrest)
    llike = logpdf(σp, σ)
    for i = 1:length(β)
        llike += logpdf(βp, β[i])
    end

    # Likelihood
    # y ~ normal(x*beta, sigma)
    for (i, μ) in enumerate(prob.X * β)
        llike += logpdf(Normal(μ, σ), prob.data[i])
    end

    return llike
end

# # "Incorrect" Model
prob = EcoRegProblem(X, y_sim)
θ_transform = TransformationTuple(ArrayTransformation(IDENTITY, P), bridge(ℝ, ℝ⁺))
#TODO: the first element of β needs to be positive
tprob = TransformLogLikelihood(prob, θ_transform)
tprob∇ = ForwardGradientWrapper(tprob, fill(0.0, length(tprob)))

@time tprob∇(randn(P + 1))

#TODO: this seems to much slower than stan
iter = 2000
@time chain1, NUTS_tuned = NUTS_init_tune_mcmc(tprob∇, fill(1.1, P + 1), iter)
@time chain2, NUTS_tuned = NUTS_init_tune_mcmc(tprob∇, fill(0.1, P + 1), iter)
# Thin
chain1 = chain1[floor(Int, iter / 2):end]
chain2 = chain2[floor(Int, iter/ 2):end]
samp = vcat(chain1, chain2)

# # # "Correct" Model
# θ_transform = TransformationTuple(bridge(ℝ, ℝ⁺),                     # β₁
# ArrayTransformation(IDENTITY, P - 1),  # β rest
# bridge(ℝ, ℝ⁺))                     # σ

NUTS_statistics(samp)
ESS = squeeze(mapslices(effective_sample_size, ungrouping_map(Array, get_position, samp), 1), 1)

# Note the use of ungrouping_map, a utility function that maps the collects posterior results,
# and groups then in tuples of vectors or arrays. You could do this manually, but you have to
# use get_position to extract the posterior position for each point. Also, in order to do inference,
# you would want to transform the "raw" values in $\mathbb{R}^n$ using the parameter transformation.
#β1, βrest, σ = ungrouping_map(Array, θ_transform ∘ get_position, samp)
#β = hcat(β1, βrest)
β, σ = ungrouping_map(Array, θ_transform ∘ get_position, samp)

# Posterior means are fairly close to the parameters (note that they will not be the same,
# because of sampling variation):
for icol = 1:size(β, 2)
    @show mean(β[:, icol]), β_true[icol]
end
mean(σ), σ_true

for icol = 1:size(β, 2)
    @show abs(mean(β[:, icol]) - β_true[icol])
end
mean(σ), σ_true

# # Plot some results
# ## Posterior
function plot_posterior(known_value, posterior_draws, varname; nbins = 20)
    #postr = fit(Histogram, posterior_draws; closed = :right, args...)
    plt[:hist](posterior_draws, bins = nbins, label = "posterior")
    xlabel(varname)
    axvline([known_value], label = "known value", c = "red")
    #legend()
end

begin
    figure()
    for i = 1:size(β, 2)
        subplot(3, 4, i)
        plot_posterior(β_true[i], β[:, i], "β$i"; nbins = 20)
    end
    subplot(3, 4, P + 1)
    plot_posterior(σ_true, σ, "σ"; nbins = 20)
    tight_layout()
end

# ## MCMC sampling
begin
    figure()
    for i = 1:size(β, 2)
        subplot(3, 4, i)
        plot(β[:, i], label = "β$i")
        axhline([β_true[i]], c = "red")
    end
    tight_layout()
end


# Gumbel distribution

export Gumbel

@parameterized Gumbel()

basemeasure(::Gumbel{()}) = Lebesgue(ℝ)

@kwstruct Gumbel()

@kwstruct Gumbel(β)

@kwstruct Gumbel(μ,β)

# map affine names to those more common for Gumbel
for N in [(:μ,), (:σ,), (:μ,:σ)]
    G = tuple(replace(collect(N), :σ => :β)...)
    @eval begin
        proxy(d::Gumbel{$G}) = affine(NamedTuple{$N}(values(params(d))), Gumbel())
        logdensity(d::Gumbel{$G}, x) = logdensity(proxy(d), x)
        basemeasure(d::Gumbel{$G}) = basemeasure(proxy(d))
    end
end

function logdensity(d::Gumbel{()} , x)
    return -exp(-x) - x
end

import Base

function Base.rand(rng::AbstractRNG, d::Gumbel{()})
    u = rand(rng)
    -log(-log(u))
end

TV.as(::Gumbel) = asℝ

≪(::Gumbel, ::Lebesgue{X}) where X <: Real = true

distproxy(::Gumbel{()}) = Dists.Gumbel()


################################ marginal ######################################
function marginal end
marginal(G::Univariate, i) = i == 1 ? G : throw(DimensionMismatch("Univariate Gaussian does not have index $i"))
marginal(G::Multivariate, i) = Gaussian(G.μ[i], G.Σ[i, i])
marginal(G::Multivariate, i::AbstractVector) = Gaussian(G.μ[i], G.Σ[i, i])
# IDEA: could pass view::Val{true} in marginal
marginalview(G::Multivariate, i) = @views Gaussian(G.μ[i], G.Σ[i, i])
marginalview(G::Multivariate, i::AbstractVector) = @views Gaussian(G.μ[i], G.Σ[i, i])

# WARNING: expects x to be vector of data
function marginal(G::Gaussian, x::AbstractVector)
	μ = marginal(mean(G), x) # evaluation of the mean could be lazy
	Σ = gramian(cov(G), x)
	return Gaussian(μ, Σ)
end
marginal(k::MercerKernel, x::AbstractVector) = gramian(k, x)
# fallback for GP mean functions
# marginal(f, x::AbstractVector) = f.(x)
# this works for vector functions
# IDEA: could this be made better by just working with static arrays as input
# and then vectorizing them? seems feasible
# marginal(f, x::StaticVector) = vec(f.(x)) # output needs to be static array
function marginal(f, x::AbstractVector)
	fx = f(x[1])
	d, T = length(fx), eltype(fx)
	marginal(f, x, d, T)
end
function marginal(f, x::AbstractVector, d::Int, T::DataType = eltype(f(x[1])))
	fx = zeros(T, (d, length(x)))
	for i in eachindex(x)
		fx[:, i] .= f(x[i])
	end
	return vec(fx)
end


# LSAC JUNE 2007 Form 8LSN75
# https://www.lsac.org/sites/default/files/legacy/docs/default-source/jd-docs/sampleptjune.pdf

# Questions 6-10
using AbstractLogic

"""
Exactly three films—Greed, Harvest, and Limelight—are
shown during a film club’s festival held on Thursday, Friday,
and Saturday. Each film is shown at least once during the
festival but never more than once on a given day. On each day
at least one film is shown. Films are shown one at a time. The
following conditions apply:
"""

# Three films Greed, Harvest, and Limelight as well as 0 no film
# t1 is the first film, t2 the second, and t3 are the last film on Thursday
# f1, f2, f3 and s1, s2, s3 are similarly coded for friday and saturday
logicset = logicalparse("t.1, t.2, t.3, f.1, f.2, f.3, s.1, s.2, s.3  ∈  _, Greed, Harvest, Limelight")

# On each day at least one film is shown. Since 3 is last slot then .3 can't be empty.
logicset = logicalparse("{{j}}.3 != '_'", logicset)

# Let's impose the non-binding contrainst that the highest number 3 is the
# last film shown. Thus if t2|t1 != 0 => t3 != 0
# logicset = logicalparse("t1 != 0 ==> t2 !=0", logicset)
logicset = logicalparse("{{j}}.1 != '_' ==> {{j}}.2 != '_'", logicset)

# Each film is shown at least once during the festival
logicset = logicalparse("{{i}} == 'Greed' {{1,}}", logicset)
logicset = logicalparse("{{i}} == 'Harvest' {{1,}}", logicset)
logicset = logicalparse("{{i}} == 'Limelight' {{1,}}", logicset)

# Never more than once on a given day. If t1 does not equal '_' then t1 != t2.
# Since we know that t3 can't be '_' then we can assert that t2 != t3
logicset = logicalparse("{{j}}.1 != '_' ==> {{j}}.1 != {{j}}.2 &&& {{j}}.1 != {{j}}.3 &&& {{j}}.2 != {{j}}.3", logicset)

# On Thursday Harvest is shown, and no film is shown after it on that day
logicset = logicalparse("t.3 = 'Harvest'", logicset)

# On Friday either Greed or Limelight, but not both, and no film is
# shown after it on that day.
logicset = logicalparse(["f.3 = 'Greed'|'Limelight'; f.2,f.1 != 'Greed','Limelight'"], logicset);

# On Saturday either Greed or Harvest, but not both, is shown, and no film is
# shown after it on that day.
logicset = logicalparse(["s.3 = 'Greed'|'Harvest'; s.1,s.2 != 'Greed','Harvest'"], logicset)

showfeasible(logicset)

range(logicset)

#6. Which one of the following could be a complete and accurate description of
# the order in which the films are shown at the festival?
#(A) Thursday: Limelight, then Harvest; Friday: Limelight; Saturday: Harvest
checkfeasible("t.2 = 'Limelight' ; t.3 = 'Harvest' ; f.3 = 'Limelight' ; s.3 = 'Harvest' ; t.1,f.2,f.1,s.2,s.1='_'", logicset)
# false 0 out of 72 possible combinations true

#(B) Thursday: Harvest; Friday: Greed, then Limelight; Saturday: Limelight, then Greed
checkfeasible("t.3 = 'Harvest' ; f.2 = 'Greed' ; f.3 = 'Limelight' ; s.2 = 'Limelight' ; s.3 = 'Greed'", logicset)
# false 0 out of 72 possible combinations true

#(C) Thursday: Harvest; Friday: Limelight; Saturday: Limelight, then Greed
checkfeasible("t.3 = 'Harvest' ; f.3 = 'Limelight' ; s.2 = 'Limelight' ; s.3 = 'Greed'; t.1, t.2, f.1, f.2, s.1='_' ", logicset)
# true 1 out of 72 possible combinations true

#(D) Thursday: Greed, then Harvest, then Limelight; Friday: Limelight; Saturday: Greed
checkfeasible("t.1 = 'Greed' ; t.2 = 'Harvest' ; t.3 = 'Limelight' ; f.3 = 'Limelight' ; s.3 = 'Greed'", logicset)
# false, 0 out of 72 possible combinations 'true'.

#(E) Thursday: Greed, then Harvest; Friday: Limelight, then Harvest; Saturday: Harvest
checkfeasible("t.2 = 'Greed' ; t.3 = 'Harvest' ; f.2 = 'Limelight' ; f.3 = 'Harvest' ; s.3 = 'Harvest' ", logicset)
# false, 0 out of 72 possible combinations 'true'.

#7. Which one of the following CANNOT be true?
#(A) Harvest is the last film shown on each day of the festival.
checkfeasible("t.3,f.3,s.3 = 'Harvest'", logicset)
# false, 0 out of 64 possible combinations 'true'

# (B) Limelight is shown on each day of the festival.
checkfeasible("{{i}} == 'Limelight' {{3,}}", logicset)
# possible,  10 out of 72 possible combinations 'true'.

#(C) Greed is shown second on each day of the festival.
checkfeasible("f.1,f.2,f.3 |= 'Greed'", logicset)
# possible,  32 out of 72 possible combinations 'true'

# (D) A different film is shown first on each day of the festival.
## This will be a little trickier. Define .0 as first film
logicset2 = logicalparse("t.0, f.0, s.0 ∈ Greed, Harvest, Limelight", logicset)
logicset2 = logicalparse("{{j}}.0 == {{j}}.1 if {{j}}.1 != '_'", logicset2)
logicset2 = logicalparse("{{j}}.1 == '_' &&& {{j}}.2 != '_' ==> {{j}}.0 == {{j}}.2", logicset2)
logicset2 = logicalparse("{{j}}.0 == {{j}}.3 if {{j}}.2 == '_'", logicset2)
checkfeasible("t.0 != f.0, s.0 ; f.0 != s.0", logicset2)
# possible,  8 out of 18 possible combinations 'true'.

# (E) A different film is shown last on each day of the festival.
checkfeasible("t.3, f.3 != s.3; t.3 != f.3", logicset)
# possible,  24 out of 74 possible combinations 'true'.

# If Limelight is never shown again during the festival once Greed is shown,
logicsetreduced = logicalparse("{{i}} = 'Greed' ==> {{>i}} != 'Limelight'", logicset)

# which one of the following is the maximum number of film showings that could
# occur during the festival?
#(A) three
checkfeasible("{{i}} != '_' {{3}}", logicsetreduced, countany=true)
# true,  1 out of 20 possible combinations 'true'.

#(B) four
checkfeasible("{{i}} != '_' {{4}}", logicsetreduced, countany=true)
# true,  6 out of 20 possible combinations 'true'.

#(C) five
checkfeasible("{{i}} != '_' {{5}}", logicsetreduced, countany=true)
# true,  9 out of 20 possible combinations 'true'.

# (D) six
checkfeasible("{{i}} != '_' {{6}}", logicsetreduced, countany=true)
# true,  4 out of 20 possible combinations 'true'.
# * Answer

# (E) seven
checkfeasible("{{i}} != '_' {{7}}", logicsetreduced, countany=true)
# true, 0 out of 20 possible combinations 'true'.

# If Greed is shown exactly three times, Harvest is shown exactly twice, and
# Limelight is shown exactly once, then which one of the following must be true?
logicsetreduced = logicalparse(
  "{{i}} = 'Greed' {{3}}; {{i}} = 'Harvest' {{2}}; {{i}} = 'Limelight' {{1}}", logicset)

#(A) All three films are shown on Thursday.
checkfeasible("t.1 != '_'", logicsetreduced, force=true)
# false, 2 out of 3 possible combinations 'true'.

#(B) Exactly two films are shown on Saturday.
checkfeasible("s.1 = '_'", logicsetreduced, force=true)
# true, 3 out of 3 possible combinations 'true'.

#(C) Limelight and Harvest are both shown on Thursday.
checkfeasible("t.1,t.2,t.3 |= 'Limelight'; t.1,t.2,t.3 |= 'Harvest'", logicsetreduced, force=true)
# false, 2 out of 3 possible combinations 'true'.

#(D) Greed is the only film shown on Saturday.
checkfeasible("s.1,s.2 != '_'; s.3 = 'Greed'", logicsetreduced, force=true)
# false, 0 out of 3 possible combinations 'true'.

#(E) Harvest and Greed are both shown on Friday
checkfeasible("s.1,s.2,s.3 |= 'Harvest'; s.1,s.2,s.3 |= 'Greed'", logicsetreduced, force=true)
# false, 0 out of 3 possible combinations 'true'.

# If Limelight is shown exactly three times, Harvest is shown exactly twice, and
# Greed is shown exactly once...
logicsetreduced = logicalparse(
  "{{i}} = 'Limelight' {{3}}; {{i}} = 'Harvest' {{2}}; {{i}} = 'Greed' {{1}}", logicset)

# Reintroduce first film
logicsetreduced2 = logicalparse("t.0, f.0, s.0 ∈ Greed, Harvest, Limelight", logicsetreduced)
logicsetreduced2 = logicalparse("{{j}}.0 == {{j}}.1 if {{j}}.1 != '_'", logicsetreduced2)
logicsetreduced2 = logicalparse("{{j}}.2 == '_' ==> {{j}}.0 == {{j}}.3", logicsetreduced2)
# Order of operations helps here as the &&& block is evaluated before the if operator
logicsetreduced2 = logicalparse("{{j}}.2 != '_' &&& {{j}}.1 == '_' ==> {{j}}.0 == {{j}}.2", logicsetreduced2)

# then which one of the following is a complete and accurate list of the films
# that could be the first film shown on Thursday?
range(logicsetreduced2)[Symbol("t.0")]
# Greed, Limelight

# (A) Harvest
# (B) Limelight
# (C) Greed, Harvest
# (D) Greed, Limelight
# (E) Greed, Harvest, Limelight


# SteinBanana
# The banana posterior distribution p(theta | y_1,...,y_L) induced by the model
#
#   y_l | theta ~iid N(theta1 + theta2^2, sigma2y)
#   theta_1, theta_2 ~iid N(0, sigma2theta)
#
# This distribution is analyzed in the papers
#
#   http://arxiv.org/pdf/1211.3759v2.pdf
#   http://jmlr.org/proceedings/papers/v33/dubois14.pdf

using Distributions: Normal, logpdf
using SteinDistributions

type SteinBanana <: SteinPosterior
    y::Array{Float64, 1}
    sigma2y::Float64
    sigma2theta::Float64
    c1::Array{Float64, 1}
    c2::Array{Float64, 1}
    c3::Array{Float64, 1}
end

# Constructor with uniform setting of (c1,c2,c3) and
# default setting of sigma2y and sigma2theta
defaultsigma2y = 4.0
SteinBanana(y; sigma2y = defaultsigma2y, sigma2theta = 1.0) =
SteinBanana(y, sigma2y, sigma2theta,
            [1.0 for i=1:2],[1.0 for i=1:2],[1.0 for i=1:2])

# Returns lower bound of support
function supportlowerbound(d::SteinBanana, j::Int64)
    -Inf
end

# Returns upper bound of support
function supportupperbound(d::SteinBanana, j::Int64)
    Inf
end

# Draws n datapoints from the likelihood distribution
#   N(theta1 + theta2^2, sigma2y)
#
# Args:
#   mu = theta1 + theta2^2
function randbanana(mu::Float64, n::Int64)
    datanormal = Normal(mu, sqrt(defaultsigma2y))
    rand(datanormal, n)
end

# Computes the log prior for N(theta1 + theta2^2, sigma2y)
function logprior(d::SteinBanana,
                  theta::Array{Float64, 1})
    priornormal = Normal(0.0, sqrt(d.sigma2theta))
    sum(logpdf(priornormal, theta))
end

# Returns log p(theta | d.y[idx]), the log density of the posterior distribution
# evaluated at theta, conditioned only on those datapoints indexed by idx
function loglikelihood(d::SteinBanana,
                       theta::Array{Float64, 1};
                       idx=1:length(d.y))
    y = d.y[idx]
    theta0 = theta[1] + theta[2]^2

    datanormal = Normal(theta0, sqrt(d.sigma2y))
    sum(logpdf(datanormal, y))
end

# Returns grad_{theta} log p(d.y[idx] | theta)
function gradloglikelihood(d::SteinBanana,
                           theta::Array{Float64, 1};
                           idx=1:length(d.y))
    y = d.y[idx]
    theta0 = theta[1] + theta[2]^2

    gradtheta1 = sum(y - theta0) / d.sigma2y
    gradtheta2 = 2 * theta[2] * gradtheta1

    [gradtheta1, gradtheta2]
end

# Returns grad_{theta} log p(theta)
function gradlogprior(d::SteinBanana,
                      theta::Array{Float64, 1})
    -(1/d.sigma2theta) .* theta
end

# diaghessian(x) = (grad_1 b1(x), grad_2 b2(x), ....) where b = grad log p.
function diaghessian(d::SteinBanana, theta::Array{Float64,1}; idx=1:length(d.y))
    m = length(idx)
    priorterm = -1.0/d.sigma2theta
    ysum = sum(d.y[idx])
    likelihoodterm = [-m, -2*ysum + 2*m*theta[1] + 6*m*theta[2]^2] ./ d.sigma2y
    priorterm .+ likelihoodterm
end

function numdatapoints(d::SteinBanana)
    length(d.y)
end


function incremental(f,min::Int64,max::Int64,N::Int64)
    t = range(min,stop = max,length = N) #makes the domain
    intervalo = Array{Float64}(undef,N-1,2) #prepares all intervals
    intervalo[:,1], intervalo[:,2] = t[1:end-1], t[2:end]
    return intervalo[(f.(t[1:end-1])) .* (f.(t[2:end])) .< 0 ,:] #chooses all intervals that the product of the edges is negative
end

function incremental(f,min::Float64,max::Float64,N::Int64)
    t = range(min,stop = max,length = N) #makes the domain
    intervalo = Array{Float64}(undef,N-1,2) #prepares all intervals
    intervalo[:,1], intervalo[:,2] = t[1:end-1], t[2:end]
    return intervalo[(f.(t[1:end-1])) .* (f.(t[2:end])) .< 0 ,:]
end


#=
# Moment-matching conditions

The stationarity conditions for ML learning of the RBM imply certain moment-matching
conditions between the data and the RBM distribution.
However, due to biased nature of the PCD approximation, these conditions might not hold
exactly in practice.
=#

import Makie, CairoMakie
import MLDatasets
import Flux
using Statistics: mean, cor
using LinearAlgebra: norm
using RestrictedBoltzmannMachines: RBM, Binary, sample_from_inputs, free_energy, log_pseudolikelihood, training_epochs
using RestrictedBoltzmannMachines: sample_v_from_v, sample_h_from_v, initialize!, pcd!, minibatch_count, moving_average

# load MNIST dataset

train_x = MLDatasets.MNIST(split=:train).features .> 0.5; # binarize
train_y = MLDatasets.MNIST(split=:train).targets;
train_x = train_x[:, :, train_y .== 0]; # only zeros for speed

# floating type we will use

Float = Float32

# initialize RBM

nhidden = 64
batchsize = 64
nupdates = 20000
epochs = training_epochs(; nsamples = size(train_x, 3), nupdates, batchsize)
rbm = RBM(Binary(Float,28,28), Binary(Float,nhidden), randn(Float,28,28,nhidden)/28)
initialize!(rbm, train_x);
lpls = Float64[]
Fm = zeros(batchsize, epochs * minibatch_count(train_x; batchsize))
Fd = zeros(batchsize, epochs * minibatch_count(train_x; batchsize))
t = lastepoch = 0
function mycb(; epoch, vm, vd, kw...)
    global lastepoch
    global t += 1
    Fm[:, t] .= free_energy(rbm, vm)
    Fd[:, t] .= free_energy(rbm, vd)
    if epoch > lastepoch
        push!(lpls, mean(log_pseudolikelihood(rbm, train_x)))
        lastepoch = epoch
    end
end
@time pcd!(
    rbm, train_x;
    epochs, batchsize, callback=mycb, l2_weights=1e-4,
    optim=Flux.ADAM(5e-4, (0.9, 0.999))
);

# Plot log_pseudolikelihood during training

Makie.lines(lpls)

# Average energies of data and model fantasy particles used during training.

fig = Makie.Figure()
ax = Makie.Axis(fig[1,1])
Makie.lines!(ax, vec(mean(Fm; dims=1)), color=(:red, 0.25))
Makie.lines!(ax, vec(mean(Fd; dims=1)), color=(:blue, 0.25))
Makie.lines!(ax, moving_average(vec(mean(Fm; dims=1)), 20), color=:red, label="Fm")
Makie.lines!(ax, moving_average(vec(mean(Fd; dims=1)), 20), color=:blue, label="Fd")
Makie.axislegend(ax)
fig

# Difference in average energies of data and model fantasy particles used during training.

fig = Makie.Figure()
ax = Makie.Axis(fig[1,1])
Makie.lines!(ax, vec(mean(Fm .- Fd; dims=1)), color=(:blue, 0.25))
Makie.lines!(ax, moving_average(vec(mean(Fm .- Fd; dims=1)), 20), color=:red)
fig

# Generate samples, in v-space and h-space.

nsteps = 2000
nsamples = 1000
v_model = falses(28, 28, nsamples, nsteps);
v_model[:,:,:,1] .= sample_from_inputs(rbm.visible, falses(28,28,nsamples));
for t in 2:nsteps
    v_model[:, :, :, t] .= sample_v_from_v(rbm, v_model[:, :, :, t - 1])
end
h_data = sample_h_from_v(rbm, train_x);
h_model = sample_h_from_v(rbm, v_model[:,:,:,end]);

v_data_center = train_x .- mean(train_x; dims=3);
h_data_center = h_data .- mean(h_data; dims=2);
vh_data = reshape(v_data_center, 28*28, :) * h_data_center' / size(v_data_center, 3);

v_model_center = v_model[:,:,:,end] .- mean(v_model[:,:,:,end]; dims=3);
h_model_center = h_model .- mean(h_model; dims=2);
vh_model = reshape(v_model_center, 28*28, :) * h_model_center' / size(v_model_center, 3);

# Average energy of generated and data points

fig = Makie.Figure()
ax = Makie.Axis(fig[1,1])
Makie.hist!(ax, free_energy(rbm, train_x), color=:purple, label="data", normalization=:pdf)
Makie.hist!(ax, free_energy(rbm, v_model[:,:,:,end]), color=:orange, label="model", normalization=:pdf)
Makie.axislegend(ax)
fig

# Plot moment-matching conditions

# <v>_data vs. <v>_model

Makie.scatter(
    vec(mean(train_x; dims=3)),
    vec(mean(v_model[:,:,:,end]; dims=3))
)
Makie.xlims!(0,1)
Makie.ylims!(0,1)
Makie.current_figure()

# <h>_data vs. <h>_model

Makie.scatter(
    vec(mean(h_data; dims=2)),
    vec(mean(h_model; dims=2)),
)
Makie.xlims!(0,1)
Makie.ylims!(0,1)
Makie.current_figure()

# <vh>_data vs. <vh>_model

Makie.scatter(vec(vh_data), vec(vh_model))
Makie.xlims!(0,1)
Makie.ylims!(0,1)
Makie.current_figure()

# Correlation

cor(vec(vh_data), vec(vh_model))

# Verify convergence of sampling

Makie.lines(vec(mean(free_energy(rbm, v_model[:,:,1:100,:]); dims=1)))

# Plot average sampled digit

Makie.heatmap(mean(v_model[:,:,:,end]; dims=3)[:,:,1])

# Plot average data digit

Makie.heatmap(mean(train_x; dims=3)[:,:,1])

# Difference

Makie.heatmap(mean(v_model[:,:,:,end]; dims=3)[:,:,1] - mean(train_x; dims=3)[:,:,1])

# Plot weights  of maximum intensity.

𝒫 = sortperm([norm(rbm.w[:,:,m]) for m in 1:nhidden]; rev=true);

fig = Makie.Figure()
for i = 1:7
    ax = Makie.Axis(fig[1,i], width=50, height=50)
    Makie.heatmap!(ax, rbm.w[:, :, 𝒫[i]])
    Makie.hidedecorations!(ax)
    Makie.hidespines!(ax)
end
Makie.resize_to_layout!(fig)
fig


module PulsarSearch
using StatsBase
include("stats.jl")

"""
`z_n_binned(profile, n) --> zsq`

``Z^2_n`` statistic for pulse profiles from binned events

See Bachetti+2021, arXiv:2012.11397

Parameters
----------
`profile` : array of floats
    The folded pulse profile (containing the number of
    photons falling in each pulse bin)

`n` : int
    Number of harmonics, including the fundamental

Returns
-------
`zsq` : float
    The value of the statistic

# Examples
```jldoctest
julia> using PulsarSearch

julia> z_n_binned(zeros(10), 2)
0.0

julia> z_n_binned(ones(10), 2) < 1e-30
true

julia> z_n_binned([10., 0., 0., 0., 0.], 2)
40.0
```
"""
function z_n_binned(profile::AbstractVector{}, n::Integer)::Float64
    total = sum(profile)
    N = length(profile)
    phase = range(0, stop = N - 1)  * (2pi/ N)

    if iszero(total)
        return 0.0
    end

    z = zero(Float64)
    for k in range(1, stop = n)
        s = 0.0
        c = 0.0
        for i in eachindex(profile)
            sk, ck = sincos(k * phase[i])
            s += profile[i] * sk
            c += profile[i] * ck
        end
        s = s^2
        c = c^2
        z += c + s
    end
    return z * 2 / total
end

"""
`z_n(phases, n) --> zsq`

``Z^2_n`` statistics, à la Buccheri+83, A&A, 128, 245, eq. 2.

Parameters
----------
`phase` : array of floats
    The phases of the events

`n` : int, default 2
    Number of harmonics, including the fundamental

Returns
-------
`zsq` : float
    The ``Z^2_n`` statistic

# Examples
```jldoctest
julia> using PulsarSearch

julia> z_n([10., 0., 0., 0., 0.], 2)
20.0
julia> z_n(ones(10), 2)
40.0
julia> z_n(Array([0.5]), 2)
0.0
```
"""
function z_n(phases::AbstractVector{T}, n::Integer)::T where {T<:AbstractFloat}
    N = length(phases)
    if N < 2
        return zero(T)
    end
    z = zero(T)
    twopiphase = 2 * pi * phases
    for k in range(1, stop = n)
        s = zero(T)
        c = zero(T)
        for ph in twopiphase
            sk, ck = sincos(k * ph)
            s += sk
            c += ck
        end
        s = s^2
        c = c^2
        z += c + s
    end
    return z * 2 / N
end

"""
`z_n_search(times, n, fmin, fmax [,oversample]) --> freqs, zsq_stat`

Calculate the ``Z^2_n`` statistics at trial frequencies in photon data.

Parameters
----------
`times` : array-like
    the event arrival times

`n` : int
    the number of harmonics in ``Z^2_n``

Other Parameters
----------------
`fmin` : float
    minimum pulse frequency to search

`fmax` : float
    maximum pulse frequency to search

`oversample` : float
    Oversampling factor with respect to the usual 1/T/n rule

Returns
-------
`fgrid` : array-like
    frequency grid of the epoch folding periodogram

`zsq_stat` : array-like
    the Z^2_n statistics corresponding to each frequency bin.
"""
function z_n_search(
    times::AbstractVector{T},
    n::Integer,
    fmin::Number,
    fmax::Number;
    oversample::Number = 2,
) where {T <: AbstractFloat}
    t0 = first(times)
    t1 = last(times)
    df = 1 / (t1 - t0) / oversample
    freqs = fmin:df:fmax
    N = length(freqs)
    stats = Vector{T}(undef, N)
    for i in eachindex(freqs)
        phases = times * freqs[i]
        phases .-= floor.(phases)
        stats[i] = z_n(phases, n)
    end
    return freqs, stats
end

"""
`z_n_search_hist(times, n, fmin, fmax [,oversample]) --> freqs, zsq_stat`

Calculate the ``Z^2_n`` statistics at trial frequencies in photon data.
Pre-bins the data using a histogram.
At the moment, it is _not_ faster. It will be after I add the 2-d hist
+ shift-and-add. But it allowed to test that `z_n_binned works` with
binned data.

Parameters
----------
`times` : array-like
    the event arrival times

`n` : int
    the number of harmonics in ``Z^2_n``

Other Parameters
----------------
`fmin` : float
    minimum pulse frequency to search

`fmax` : float
    maximum pulse frequency to search

`oversample` : float
    Oversampling factor with respect to the usual 1/T/n rule

Returns
-------
`fgrid` : array-like
    frequency grid of the epoch folding periodogram

`zsq_stat` : array-like
    the Z^2_n statistics corresponding to each frequency bin.
"""
function z_n_search_hist(
    times::AbstractVector{T},
    n::Integer,
    fmin::Number,
    fmax::Number;
    oversample::Number = 2,
    nbin::Integer = 16
) where {T <: AbstractFloat}
    t0 = first(times)
    t1 = last(times)
    df = 1 / (t1 - t0) / oversample
    freqs = fmin:df:fmax
    N = length(freqs)
    stats = Vector{T}(undef, N)
    for i in eachindex(freqs)
        phases = times * freqs[i]
        phases .-= floor.(phases)
        hist = fit(Histogram, phases, 0.0:(1/nbin):1.0)
        stats[i] = z_n_binned(hist.weights, n)
    end
    return freqs, stats
end

export z_n
export z_n_binned
export z_n_search
export z_n_search_hist
# Write your package code here.

end


using FuzzySets, Test

empty = Interval() # empty interval
@test isnan(empty.left)
@test isnan(empty.right)

a = Interval(2, 4)
@test mid(a) == 3
@test rad(a) == 1

b = Interval(3, 5)
@test b == Interval(3, 5)

@test a + b == Interval(5, 9)
@test a - b == Interval(-3, 1)
@test a * b == Interval(6, 20)
@test a / b == Interval(2/5, 4/3)

c = 1 .. 2
@test c * (c - c) == Interval(-2, 2)
@test c * c - c * c == Interval(-3, 3)

a⃗ = [a, b, c]
for (idx, x) in enumerate(zip(a, b))
    if idx == 1
        @test x == (2, 3)
    elseif idx == 2
        @test x == (4, 5)
    end
end

# Fuzzy distances
a = Interval(1.5, 2.5)
b = Interval(1.5, 2.5)
@test a - b == Interval(-1.0, 1.0)
@test (a - b) ^ 2 == Interval(0, 1.0)

a = Interval(0.5, 1.5)
b = Interval(0.6, 1.6)
@test a - b == Interval(-1.1, 0.9)
@test (a - b) ^ 2 ≈ Interval(0, 1.21)

a = Interval(0, 2.21)
b = Interval(0, 1.48)
@test a - b == Interval(-1.48, 2.21)
@test (a - b) ^ 2 == Interval(0, 4.8841)

@test 1 / Interval(0.4, 1.6) == Interval(0.625, 2.5)
@test 1 / Interval(-0.4, 1.6) == Interval(-2.5, 0.625)
@test 1 / Interval(0, 25) == Interval(0, 0.04)
@test 1 / (1 + Interval(0.4, 4)) ≈ Interval(0.2, 0.7142857142857143)

a = Interval(0.3, 0.6)
b = Interval(0.4, 0.7)
@test a - b == Interval(0.3-0.7, 0.6-0.4)
@test (a - b)^2 == Interval(0, max((0.3-0.7)^2, (0.6-0.4)^2))


"""
    polynomial_mutation!(vector, bounds, η=20, prob = 1 / length(vector))

Polynomial Mutation applied to a vector of real numbers.
"""
function polynomial_mutation!(vector, bounds, η=20, prob = 1 / length(vector))
    do_mutation = rand(length(vector)) .< prob

    xu = view(bounds, 2,do_mutation)
    xl = view(bounds, 1,do_mutation)
    x = view(vector, do_mutation)

    δ1 = (x - xl) ./ (xu - xl)
    δ2 = (xu - x) ./ (xu - xl)

    D = length(xu)
    R = rand(D)
    mask = R .< 0.5
    s = η+1.0
    mut_pow = 1.0 / (η + 1.0)
    δq = [ mask[i] ?
            ^(2.0R[i] + (1. - 2.0R[i]) * ^(1.0 - δ1[i], s), mut_pow) - 1.0 :
            1.0 - (2.0 * (1.0 - R[i]) + 2.0 * (R[i] - 0.5) * ^(1.0 - δ2[i], s))^mut_pow
            for i in 1:D
        ]

    vector[do_mutation] = x + δq .* ( xu - xl)
    # correct using reset to bound
    #
    vector

end

"""
    DE_mutation(population, F = 1.0, strategy = :rand1)

Generate a `Vector` computed from population used in Differential Evolution.
Parameters: `F` is the stepsize, `strategy` can be one the following `:best1`, `:rand2`,
`:randToBest1` or `:best2`.
"""
function DE_mutation(population,
                    F::Float64 = 1.0,
                    strategy::Symbol=:rand1,
                    best_ind=0
    )

    N = length(population)

    i = rand(1:N)
    # select participats
    r1 = rand(1:N)
    while r1 == i
        r1 = rand(1:N)
    end

    r2 = rand(1:N)
    while r2 == i || r1 == r2
        r2 = rand(1:N)
    end

    x = get_position(population[i] )
    a = get_position(population[r1])
    b = get_position(population[r2])

    # strategy is selected here
    if strategy == :rand1
        # DE/rand/1
        return  x + F * (a - b)
    end

    r3 = rand(1:N)
    while r3 == i || r3 == r1 || r3 == r2
        r3 = rand(1:N)
    end

    c = get_position(population[r3])

    if strategy == :rand2
        # DE/rand/2

        r4 = rand(1:N)
        while r4 == i || r4 == r1 || r4 == r2 || r4 == r3
            r4 = rand(1:N)
        end

        r5 = rand(1:N)
        while r5 == i || r5 == r1 || r5 == r2 || r5 == r3 || r5 == r4
            r5 = rand(1:N)
        end

        d = get_position(population[r4])
        ee = get_position(population[r5])

        return ee + F * (a - b + c - d)
    end

    best_ind = best_ind == 0 ? argbest(population) : best_ind
    x_best = get_position(population[best_ind])
    if strategy == :best1
        # DE/best/1
        u = x_best + F * (b - c)
    elseif strategy == :randToBest1
        # DE/rand-to-best/1
        u = x + F * (x_best - x + a - b)
    elseif strategy == :best2
        # DE/best/2
        r4 = rand(1:N)
        while r4 == i || r4 == r1 || r4 == r2 || r4 == r3 || r4 == best_ind
            r4 = rand(1:N)
        end
        d = population[r4].x
        u = x_best + F * (a - b + c - d)
    else
        error("Unknown strategy " * string(strategy))
    end

    return u
end

"""
    MOEAD_DE_reproduction(a, b, c, F, CR, p_m, η, bounds)

Perform Differential Evolution operators and polynomial mutation using three vectors
`a, b, c` and parameters `F, CR, p_m, η`, i.e., stepsize, crossover and
mutation probability.
"""
function MOEAD_DE_reproduction(a, b, c, F, CR, p_m, η, bounds)
    D = length(a)
    # binomial crossover
    v = zeros(length(a))

    la = view(bounds, 1, :)
    lb = view(bounds, 2, :)

    # binomial crossover
    for j in 1:D
        # binomial crossover
        if rand() < CR
            v[j] = a[j] + F * (b[j] - c[j])
        else
            v[j] = a[j]
        end
        # polynomial mutation

        if rand() < p_m
            r = rand()
            if r < 0.5
                σ_k = (2.0 * r)^(1.0 / (η + 1)) - 1
            else
                σ_k = 1 - (2.0 - 2.0 * r)^(1.0 / (η + 1))
            end
            v[j] = v[j] + σ_k * (lb[j] - la[j])
        end
    end

    v
end


function traverse(posx::Float64,posy::Float64,velx::Float64,vely::Float64,
                    x0::Float64,x1::Float64,y0::Float64,y1::Float64)
    # Given a particle in a cell with position (posx,posy)
    # and velocity (velx,vely) where the cell is defined by
    # [x0,x1] x [y0,y1],
    # this function returns nexti, nextj = the direction of the cell
    # that the particle will, as well as nextx,nexty = the intersection point
    # and dist = the travelled distance.
    # move into after having passed the current cell.
    #
    # The layout is as follows:
    #
    #                     j-1            j            j+1
    #     \  y
    #      \
    #     x \
    #
    #   i+1
    #
    #                 ---------------x1---------------
    #                 |                              |
    #                 |                              |
    #     i           y0                             y1
    #                 |                              |
    #                 |                              |
    #                 ---------------x0---------------
    #
    #   i-1
    #
    # The coordinate system is
    #
    #  ^  x
    #  |
    #  |
    #  |
    #    --------> y
    #

    # assert that we are in the domain
    #eps = 1e-14
    #@assert posx>= x0-eps
    #@assert posx<= x1+eps
    #@assert posy>= y0-eps
    #@assert posy<= y1+eps

    # assert a non zero velocity
    normv = sqrt(vely^2+velx^2)
    #@assert normv>0

    # are we moving north east, south east, south west or north west
    vertical   = ifelse(velx>=0,+1,-1)
    horizontal = ifelse(vely>=0,+1,-1)

    # which face (left right top bottom) will we hit
    # i.e. to which y or x do we need to calculate the distance
    x = ifelse(vertical==+1,x1,x0)
    y = ifelse(horizontal==+1,y1,y0)

    # calculate distance to that point considering the velocity
    distx = abs((x-posx)/velx)
    disty = abs((y-posy)/vely)

    # check which face we hit first and save the corresponding i/j shift
    # as well as the distance travelled
    if disty<distx
        nexti = 0
        nextj = horizontal
        dist = disty
    elseif disty>distx
        nexti = vertical
        nextj = 0
        dist = distx
    else
        nexti = vertical
        nextj = horizontal
        dist = distx # == disty
    end

    # at what point to we cross the interface
    nextx = posx + dist*velx
    nexty = posy + dist*vely
    return nexti,nextj, dist*normv, nextx,nexty,velx,vely
end


function numerical_gradient(feval, x, dx=0.00001)

  g = zeros(size(x))

  for ii = 1:length(x)

    x[ii] = x[ii] + dx

    D1 = feval(x)

    x[ii] = x[ii] - 2*dx

    D2 = feval(x)

    x[ii] = x[ii] + dx

    g[ii] = (D1-D2)/(2*dx)

  end

  return g

end


doc"""
    VonMises(μ, κ)

The *von Mises distribution* with mean `μ` and concentration `κ` has probability density function

$f(x; \mu, \kappa) = \frac{1}{2 \pi I_0(\kappa)} \exp \left( \kappa \cos (x - \mu) \right)$

```julia
VonMises()       # von Mises distribution with zero mean and unit concentration
VonMises(κ)      # von Mises distribution with zero mean and concentration κ
VonMises(μ, κ)   # von Mises distribution with mean μ and concentration κ
```

External links

* [von Mises distribution on Wikipedia](http://en.wikipedia.org/wiki/Von_Mises_distribution)

"""
immutable VonMises <: ContinuousUnivariateDistribution
    μ::Float64      # mean
    κ::Float64      # concentration
    I0κ::Float64    # I0(κ), where I0 is the modified Bessel function of order 0

    function VonMises(μ::Real, κ::Real)
        @check_args(VonMises, κ > zero(κ))
        new(μ, κ, besseli(0, κ))
    end
    VonMises(κ::Real) = VonMises(0.0, κ)
    VonMises() = new(0.0, 1.0)
end

show(io::IO, d::VonMises) = show(io, d, (:μ, :κ))

@distr_support VonMises d.μ - π d.μ + π

#### Parameters

params(d::VonMises) = (d.μ, d.κ)

#### Statistics

mean(d::VonMises) = d.μ
median(d::VonMises) = d.μ
mode(d::VonMises) = d.μ
circvar(d::VonMises) = 1.0 - besseli(1, d.κ) / d.I0κ
entropy(d::VonMises) = log(twoπ * d.I0κ) - d.κ * (besseli(1, d.κ) / d.I0κ)

cf(d::VonMises, t::Real) = (besseli(abs(t), d.κ) / d.I0κ) * cis(t * d.μ)

#### Evaluation

pdf(d::VonMises, x::Float64) = exp(d.κ * cos(x - d.μ)) / (twoπ * d.I0κ)
logpdf(d::VonMises, x::Float64) = d.κ * cos(x - d.μ) - log(d.I0κ) - log2π

cdf(d::VonMises, x::Float64) = _vmcdf(d.κ, d.I0κ, x - d.μ, 1.0e-15)

function _vmcdf(κ::Float64, I0κ::Float64, x::Float64, tol::Float64)
    j = 1
    cj = besseli(j, κ) / j
    s = cj * sin(j * x)
    while abs(cj) > tol
        j += 1
        cj = besseli(j, κ) / j
        s += cj * sin(j * x)
    end
    return (x + 2.0 * s / I0κ) / twoπ + 0.5
end

#### Sampling

sampler(d::VonMises) = VonMisesSampler(d.μ, d.κ)


@testset "test_matrix_comps" begin
A = [-0.21 0.2; 0.2 -0.21]
B = 0.01*[1 0; 0 1]
C = [1 0; 0 1]
D = 0
sys = ss(A,B,C,D)
sysr, G = balreal(sys)

@test gram(sysr, :c) ≈ G
@test gram(sysr, :o) ≈ G
@test sort(pole(sysr)) ≈ sort(pole(sys))

sysb,T = ControlSystems.balance_statespace(sys)
Ab,Bb,Cb,T = ControlSystems.balance_statespace(A,B,C)

@test Ab*T ≈ T*A
@test Bb ≈ T*B
@test Cb*T ≈ C

@test sysb.A ≈ Ab

@test ControlSystems.balance_transform(A,B,C) ≈ ControlSystems.balance_transform(sys)

W = [1 0; 0 1]
@test covar(sys, W) ≈ [0.002560975609756 0.002439024390244; 0.002439024390244 0.002560975609756]
D2 = [1 0; 0 1]
@test covar(ss(A,B,C,D2, 1), W) ≈ [1.000110108378310 -0.000010098377310; -0.000010098377310 1.000110108378310]
# Direct term means infinite covariance
@test covar(ss(A,B,C,D2), W) ≈ [Inf Inf; Inf Inf]

# No noise on second output should give finite variance
@test covar(ss(A,B,C,[1 0; 0 0]), W) ≈ [Inf Inf; Inf 0.002560975609756]

# Unstable system has inf covar
@test covar(ss([1 0; 0 1],B,C,0), W) == [Inf Inf; Inf Inf]

# Discrete system can have direct term
@test covar(ss(A,B,C,D2,0.1),W) ≈ [1.00011010837831 -1.0098377309782909e-5; -1.0098377309782909e-5 1.00011010837831]

# TODO test in Julia 0.7 to see if supported
# # Test special matrices
As = sparse(A)
Bs = sparse(B)
Cs = sparse(C)
Asb,Bsb,Csb,Ts = ControlSystems.balance_statespace(As,Bs,Cs) #Error no LAPACK function
#
# @test Abs*Ts ≈ Ts*As
# @test Bbs ≈ Ts*Bs
# @test Cbs*Ts ≈ Cs

# Test special values
Ar = rationalize.(A)
Br = rationalize.(B)
Cr = rationalize.(Float64.(C))    # When did ever rationalize work on Int?
Arb,Brb,Crb,Tr = ControlSystems.balance_statespace(Ar,Br,Cr)

@test Arb*Tr ≈ Tr*Ar
@test Brb ≈ Tr*Br
@test Crb*Tr ≈ Cr

end


using SphericalFourierBesselDecompositions
SFB = SphericalFourierBesselDecompositions

using Test

@testset "top-level" begin
    rmin = 500.0
    rmax = 1000.0
    kmax = 0.02
    amodes = SFB.AnlmModes(kmax, rmin, rmax)
    cmodes = SFB.ClnnModes(amodes, Δnmax=0)
    pixwin = SFB.pixwin(cmodes)
end

# vim: set sw=4 et sts=4 :


function dot_simd(a::AbstractVector, b::AbstractVector)
    s = zero(eltype(a))
    @fastmath @inbounds @simd for i ∈ eachindex(a)
        s += a[i]' * b[i]
    end
    s
end
function cdot_mat(ca::AbstractVector{Complex{T}}, cb::AbstractVector{Complex{T}}) where {T}
    a = reinterpret(reshape, T, ca)
    b = reinterpret(reshape, T, cb)
    re = zero(T); im = zero(T)
    @turbo for i ∈ axes(a,2)
        re += a[1,i] * b[1,i] + a[2,i] * b[2,i]
        im += a[1,i] * b[2,i] - a[2,i] * b[1,i]
    end
    Complex(re, im)
end
function cdot_swizzle(ca::AbstractVector{Complex{T}}, cb::AbstractVector{Complex{T}}) where {T}
  a = reinterpret(T, ca)
  b = reinterpret(T, cb)
  reim = Vec(zero(T),zero(T))
  @turbo for i ∈ eachindex(a)
    reim = vfmsubadd(vmovsldup(a[i]), b[i], vfmsubadd(vmovshdup(a[i]), vpermilps177(b[i]), reim))
  end
  Complex(reim(1), reim(2))
end
function cdot_affine(ca::AbstractVector{Complex{T}}, cb::AbstractVector{Complex{T}}) where {T}
    a = reinterpret(T, ca);
    b = reinterpret(T, cb);
    re = zero(T); im = zero(T)
    # with a multiplier, we go from `i = 1 -> 2i = 2` to `i = 0 -> 2i = 0
    # 2(i+1-1) = 2i + 2 - 2, so....
    @turbo for i ∈ 1:length(a)>>>1
        re += a[2i-1] * b[2i-1] + a[2i] * b[2i  ]
        im += a[2i-1] * b[2i  ] - a[2i] * b[2i-1]
    end
    Complex(re, im)
end
function cdot_stride(ca::AbstractVector{Complex{T}}, cb::AbstractVector{Complex{T}}) where {T}
    a = reinterpret(T, ca);
    b = reinterpret(T, cb);
    re = zero(T); im = zero(T)
    @turbo for i ∈ 1:2:length(a)
        re += a[i] * b[i  ] + a[i+1] * b[i+1]
        im += a[i] * b[i+1] - a[i+1] * b[i  ]
    end
    Complex(re, im)
end
function qdot_simd(x::AbstractVector{NTuple{4,T}}, y::AbstractVector{NTuple{4,T}}) where {T}
    a = zero(T)
    b = zero(T)
    c = zero(T)
    d = zero(T)
    @fastmath @inbounds @simd for i ∈ eachindex(x)
        a₁, b₁, c₁, d₁ = x[i]
        a₂, b₂, c₂, d₂ = y[i]
        a += a₁*a₂ + b₁*b₂ + c₁*c₂ + d₁*d₂
        b += a₁*b₂ - b₁*a₂ - c₁*d₂ + d₁*c₂
        c += a₁*c₂ + b₁*d₂ - c₁*a₂ - d₁*b₂
        d += a₁*d₂ - b₁*c₂ + c₁*b₂ - d₁*a₂
    end
    (a,b,c,d)
end
function qdot_mat(x::AbstractMatrix, y::AbstractMatrix)
    a = zero(eltype(x))
    b = zero(eltype(x))
    c = zero(eltype(x))
    d = zero(eltype(x))
    @turbo for i ∈ axes(x,2)
        a₁ = x[1,i]
        b₁ = x[2,i]
        c₁ = x[3,i]
        d₁ = x[4,i]
        a₂ = y[1,i]
        b₂ = y[2,i]
        c₂ = y[3,i]
        d₂ = y[4,i]
        a += a₁*a₂ + b₁*b₂ + c₁*c₂ + d₁*d₂
        b += a₁*b₂ - b₁*a₂ - c₁*d₂ + d₁*c₂
        c += a₁*c₂ + b₁*d₂ - c₁*a₂ - d₁*b₂
        d += a₁*d₂ - b₁*c₂ + c₁*b₂ - d₁*a₂
    end
    (a,b,c,d)
end
function qdot_affine(x::AbstractVector, y::AbstractVector)
    a = zero(eltype(x))
    b = zero(eltype(x))
    c = zero(eltype(x))
    d = zero(eltype(x))
    @turbo for i ∈ 1:length(x)>>2
        a₁ = x[4i-3]
        b₁ = x[4i-2]
        c₁ = x[4i-1]
        d₁ = x[4i]
        a₂ = y[4i-3]
        b₂ = y[4i-2]
        c₂ = y[4i-1]
        d₂ = y[4i]
        a += a₁*a₂ + b₁*b₂ + c₁*c₂ + d₁*d₂
        b += a₁*b₂ - b₁*a₂ - c₁*d₂ + d₁*c₂
        c += a₁*c₂ + b₁*d₂ - c₁*a₂ - d₁*b₂
        d += a₁*d₂ - b₁*c₂ + c₁*b₂ - d₁*a₂
    end
    (a,b,c,d)
end
function qdot_stride(x::AbstractVector, y::AbstractVector)
    a = zero(eltype(x))
    b = zero(eltype(x))
    c = zero(eltype(x))
    d = zero(eltype(x))
    @turbo for i ∈ 1:4:length(x)
        a₁ = x[i]
        b₁ = x[i+1]
        c₁ = x[i+2]
        d₁ = x[i+3]
        a₂ = y[i]
        b₂ = y[i+1]
        c₂ = y[i+2]
        d₂ = y[i+3]
        a += a₁*a₂ + b₁*b₂ + c₁*c₂ + d₁*d₂
        b += a₁*b₂ - b₁*a₂ - c₁*d₂ + d₁*c₂
        c += a₁*c₂ + b₁*d₂ - c₁*a₂ - d₁*b₂
        d += a₁*d₂ - b₁*c₂ + c₁*b₂ - d₁*a₂
    end
    (a,b,c,d)
end
function cmatmul_array!(Cc::AbstractMatrix{Complex{T}}, Ac::AbstractMatrix{Complex{T}}, Bc::AbstractMatrix{Complex{T}}) where {T}
  C = reinterpret(reshape, Float64, Cc);
  A = reinterpret(reshape, Float64, Ac);
  B = reinterpret(reshape, Float64, Bc);
  @turbo for n ∈ indices((C,B),3), m ∈ indices((C,A),2)
    Cre = zero(T)
    Cim = zero(T)
    for k ∈ indices((A,B),(3,2))
      Cre += A[1,m,k] * B[1,k,n] - A[2,m,k] * B[2,k,n]
      Cim += A[1,m,k] * B[2,k,n] + A[2,m,k] * B[1,k,n]
    end
    C[1,m,n] = Cre
    C[2,m,n] = Cim
  end
  return Cc
end
function cmatmul_array_v2!(Cc::AbstractMatrix{Complex{T}}, Ac::AbstractMatrix{Complex{T}}, Bc::AbstractMatrix{Complex{T}}) where {T}
  C = reinterpret(Float64, Cc);
  A = reinterpret(Float64, Ac);
  B = reinterpret(reshape, Float64, Bc);
  @turbo vectorize=2 for n ∈ indices((C,B),(2,3)), m ∈ indices((C,A),1)
    Cmn = zero(T)
    for k ∈ indices((A,B),(2,2))
      Amk = A[m,k]
      Aperm = vpermilps177(Amk)
      Cmn = vfmaddsub(Amk, B[1,k,n], vfmaddsub(Aperm, B[2,k,n], Cmn))
    end
    C[m,n] = Cmn
  end
  return Cc
end

function issue209(M, G, J, H, B, ϕ)
    # tmp = similar(ϕ, G-1, (2*J+1)*(H + 1));
    tmp = view(fill(eltype(ϕ)(123456789), G+15, (2*J+1)*(H + 1)+16), 9:G+7, 9:(2*J+1)*(H + 1)+8);
    Bf = reinterpret(reshape, Float64, B);
    ϕf = reinterpret(reshape, Float64, ϕ);
    tmpf = reinterpret(reshape, Float64, tmp);
    jmax = 2*J + 1
    # B is being indexed via ptr offsetting
    # thus B's initial `gesp`ing must set it up for this
    # currently, it isn't because  `jj` and `hh` are loop induct vars
    for mm = 1:M
        m_idx = M + 2 - mm
        @turbo for hh = 1:H+1
            h_idx = (hh - 1)*jmax
            for jj = 1:jmax, gg = 1:G-1
                tmpf[1, gg, jj + h_idx] = ϕf[1, jj, gg+1, hh, m_idx] +
                                          Bf[1, jj, gg,   hh, m_idx]
                tmpf[2, gg, jj + h_idx] = ϕf[2, jj, gg+1, hh, m_idx] +
                                          Bf[2, jj, gg,   hh, m_idx]
            end
        end
    end
    parent(tmp)
end
function issue209_noavx(M, G, J, H, B, ϕ)
    tmp = view(fill(eltype(ϕ)(123456789), G+15, (2*J+1)*(H + 1)+16), 9:G+7, 9:(2*J+1)*(H + 1)+8);
    Bf = reinterpret(reshape, Float64, B);
    ϕf = reinterpret(reshape, Float64, ϕ);
    tmpf = reinterpret(reshape, Float64, tmp);
    jmax = 2*J + 1
    for mm = 1:M
        m_idx = M + 2 - mm
        for hh = 1:H+1
            h_idx = (hh - 1)*jmax
            for jj = 1:jmax, gg = 1:G-1
                tmpf[1, gg, jj + h_idx] = ϕf[1, jj, gg+1, hh, m_idx] +
                                          Bf[1, jj, gg,   hh, m_idx]
                tmpf[2, gg, jj + h_idx] = ϕf[2, jj, gg+1, hh, m_idx] +
                                          Bf[2, jj, gg,   hh, m_idx]
            end
        end
    end
    parent(tmp)
end
using LoopVectorization

function sumdim2_turbo!(r1, r2)
  @turbo thread=true for j = indices((r1,r2),(3,4)), i ∈ indices((r1,r2),(2,3))
    r1[1,i,j] = r2[1,1,i,j] + r2[1,2,i,j]
    r1[2,i,j] = r2[2,1,i,j] - r2[2,2,i,j]
    r1[3,i,j] = r2[3,1,i,j] * r2[3,2,i,j]
    r1[4,i,j] = r2[4,1,i,j] / r2[4,2,i,j]
  end
  r1
end
function sumdim2!(r1, r2)
  @inbounds @fastmath for j = indices((r1,r2),(3,4)), i ∈ indices((r1,r2),(2,3))
    r1[1,i,j] = r2[1,1,i,j] + r2[1,2,i,j]
    r1[2,i,j] = r2[2,1,i,j] - r2[2,2,i,j]
    r1[3,i,j] = r2[3,1,i,j] * r2[3,2,i,j]
    r1[4,i,j] = r2[4,1,i,j] / r2[4,2,i,j]
  end
  r1
end

# Issue 287
 function my_gemm_noturbo!(out, s::Matrix{UInt8}, V)
           Vcols = size(V, 2)
           srows = size(s, 1)
           scols = size(s, 2)
           k = srows >> 2
           rem = srows & 3
           @inbounds @fastmath for c in 1:Vcols
               for j in 1:scols
                   for l in 1:k
                       block = s[l, j]
                       for p in 1:4
                           Aij = (block >> (2 * (p - 1))) & 3
                           out[4*(l - 1) + p, c] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
                       end
                   end
               end
           end
           # TODO handle rem
       end
function my_gemm_unroll(out, s::Matrix{UInt8}, V)
  Vcols = size(V, 2)
  srows = size(s, 1)
  scols = size(s, 2)
  k = srows >> 2
  rem = srows & 3
  @avx for c in 1:Vcols
    for j in 1:scols
      for l in 1:k
        block = s[l, j]
        for p in 1:4
          Aij = (block >> (2 * (p - 1))) & 3
          out[4*(l - 1) + p, c] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
        end
      end
    end
  end
  # TODO handle rem
end
function my_gemm_manual_unroll(out, s::Matrix{UInt8}, V)
  Vcols = size(V, 2)
  srows = size(s, 1)
  scols = size(s, 2)
  k = srows >> 2
  rem = srows & 3
  @avx for c in 1:Vcols
    for j in 1:scols
      for l in 1:k
        block = s[l, j]
        # unrolled loop
        thisiszero = 0
        p = 1
        Aij = (block >> (2 * (p - 1))) & 3
        out[4*(l - 1) + p, c + thisiszero] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
        p = 2
        Aij = (block >> (2 * (p - 1))) & 3
        out[4*(l - 1) + p, c] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
        p = 3
        Aij = (block >> (2 * (p - 1))) & 3
        out[4*(l - 1) + p + thisiszero, c] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
        p = 4
        Aij = (block >> (2 * (p - 1))) & 3
        out[4*(l - 1) + p, c] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
      end
    end
  end
  # TODO handle rem
end
function my_gemm_nexpr_unroll(out, s::Matrix{UInt8}, V)
  Vcols = size(V, 2)
  srows = size(s, 1)
  scols = size(s, 2)
  k = srows >> 2
  rem = srows & 3
  @turbo for c in 1:Vcols
    for j in 1:scols
      for l in 1:k
        block = s[l, j]
        # unrolled loop
        Base.Cartesian.@nexprs 4 p -> begin
          Aij = (block >> (2 * (p - 1))) & 3
          out[4*(l - 1) + p, c] += ((Aij >= 2) + (Aij == 3)) * V[j, c]
        end
      end
    end
  end
  # TODO handle rem
end

function readraw_turbo!(img, raw)
  npack = length(raw) ÷ 3
  @turbo for i in 0:npack-1
    img[1+4i] = raw[2+3i] << 4
    img[2+4i] = raw[1+3i]
    img[3+4i] = raw[2+3i]
    img[4+4i] = raw[3+3i]
  end
  img
end
function readraw!(img, raw)
  npack = length(raw) ÷ 3
  @inbounds @simd for i in 0:npack-1
    img[1+4i] = raw[2+3i] << 4
    img[2+4i] = raw[1+3i]
    img[3+4i] = raw[2+3i]
    img[4+4i] = raw[3+3i]
  end
  img
end

@testset "shuffles load/stores" begin
  @show @__LINE__
  for i ∈ 1:128
    ac = rand(Complex{Float64}, i);
    bc = rand(Complex{Float64}, i);
    dsimd = dot_simd(ac, bc)
    if VERSION ≥ v"1.6.0-rc1"
      @test dsimd ≈ cdot_mat(ac, bc)
    end
    @test dsimd ≈ cdot_affine(ac, bc) ≈ cdot_stride(ac, bc)

    xq = [ntuple(_ -> rand(), Val(4)) for _ ∈ 1:i];
    yq = [ntuple(_ -> rand(), Val(4)) for _ ∈ 1:i];
    xqv = reinterpret(Float64, xq);
    yqv = reinterpret(Float64, yq);
    qsimd = Base.vect(qdot_simd(xq, yq)...);
    if VERSION ≥ v"1.6.0-rc1"
      xqm = reinterpret(reshape, Float64, xq);
      yqm = reinterpret(reshape, Float64, yq);
      @test qsimd ≈ Base.vect(qdot_mat(xqm, yqm)...)
    end
    @test qsimd ≈ Base.vect(qdot_affine(xqv, yqv)...) ≈ Base.vect(qdot_stride(xqv, yqv)...)

    if VERSION ≥ v"1.6.0-rc1"
      for j ∈ max(1,i-5):i+5, k ∈ max(1,i-5,i+5)
        Ac = rand(Complex{Float64}, j, i);
        Bc = rand(Complex{Float64}, i, k);
        Cc1 = Ac*Bc;
        Cc2 = similar(Cc1);
        Cc3 = similar(Cc1);
        @test Cc1 ≈ cmatmul_array!(Cc2, Ac, Bc)
        Cc2 .= NaN;
        @test Cc1 ≈ cmatmul_array_v2!(Cc2, Ac, Bc)
      end
    end
  end
  @show @__LINE__
  if VERSION ≥ v"1.6.0-rc1"
    M = 10
    G = 50
    J = 50
    H = 30

    # B = rand(Complex{Float64}, 2*J+1, G-1, H+1, M+1);
    # ϕ = rand(Complex{Float64}, 2*J+1, G+1, H+1, M+1);
    rbc = let rb = 1.0:((2*J+17) * (G+15) * (H+17) * (M+17)), rbr = reverse(rb)
      Complex{Float64}[rb[i] + im * rbr[i] for i ∈ eachindex(rb)];
    end
    B = view(reshape(rbc, (2*J+17, G+15, H+17, M+17)), 9:2*J+9, 9:G+9, 9:H+9, 9:M+9) .= rand.() .+ rand.().*im;
    ϕ = view(fill(1e5+1e7im, 2*J+17, G+17, H+17, M+17), 9:2*J+9, 9:G+9, 9:H+9, 9:M+9) .= rand.() .+ rand.().*im;
    @test issue209(M, G, J, H, B, ϕ) ≈ issue209_noavx(M, G, J, H, B, ϕ)
  end

  s = Array{Float64}(undef, 4, 128, 128);
  s2 = rand(4, 2, 128, 128);
  @test sumdim2_turbo!(s, s2) ≈ sumdim2!(similar(s), s2)

  # issue 287
  out_test = zeros(100, 10);
  out_test1 = zeros(100, 10);
  s = rand(UInt8, 25, 100);
  V = rand(100, 10);
  my_gemm_noturbo!(out_test, s, V);
  my_gemm_unroll(out_test1, s, V);
  @test out_test ≈ out_test1
  my_gemm_manual_unroll(fill!(out_test1, 0), s, V);
  @test out_test ≈ out_test1
  my_gemm_nexpr_unroll(fill!(out_test1, 0), s, V);
  @test out_test ≈ out_test1

  w = 2048; raw = rand(UInt8, (3w*w) ÷ 4);
  img1 = Matrix{UInt8}(undef, w, w);
  img2 = Matrix{UInt8}(undef, w, w);
  @test readraw!(img1, raw) == readraw_turbo!(img2, raw)
end


module PolyUtils

using ForwardDiff

# taken from Jacobi.lagrange
function lagrange(i, x, z)
    nz = length(z)

    l = one(z[1])

    for k = 1:(i-1)
        l = l * (x-z[k]) / (z[i]-z[k])
    end

    for k = (i+1):nz
        l = l * (x-z[k]) / (z[i]-z[k])
    end

    return l
end

dlagrange(i, x, z) = ForwardDiff.derivative(x -> lagrange(i, x, z), x)

end  # module


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Compute the lighting condition of a satellite:
#       1) Sunlight,
#       2) penumbra, or
#       3) umbra.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Longo, C. R. O., Rickman, S. L (1995). Method for the Calculation of
#       Spacecraft Umbra and Penumbra Shadow Terminator Points. NASA Technical
#       Paper 3547.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export SAT_LIGHTING_SUNLIGHT, SAT_LIGHTING_UMBRA, SAT_LIGHTING_PENUMBRA
export satellite_lighting_condition

# Constants.
const SAT_LIGHTING_SUNLIGHT = 0
const SAT_LIGHTING_PENUMBRA = 1
const SAT_LIGHTING_UMBRA = 2

"""
    satellite_lighting_condition(r_i::AbstractVector, s_i::AbstractVector)

Compute the satellite lighting condition given the Sun unitary vector `s_i` [m]
and the satellite position vector `r_i` [m].

# Returns

* `SAT_LIGHTING_SUNLIGHT`: Satellite is under sunlight.
* `SAT_LIGHTING_PENUMBRA`: Satellite is at penumbra region.
* `SAT_LIGHTING_UMBRA`: Satellite is at umbra region.

"""
function satellite_lighting_condition(r_i::AbstractVector, s_i::AbstractVector)
    # Norme of the sun position vector.
    norm_s_i = norm(s_i)

    # Umbra section [1].
    Xu = R0*norm_s_i/(Rs-R0)
    Au = asin(R0/Xu)

    # Penumbra section [1].
    Xp = R0*norm_s_i/(Rs+R0)
    Ap = asin(R0/Xp)

    # Projection of the satellite position vector on the sun direction [1].
    rs_i = dot(r_i, s_i)*s_i/(norm_s_i^2)

    # Distance of the umbral cone and the spacecraft [1].
    delta_i = r_i - rs_i

    # Location of the umbral cone terminator at the projected spacecraft
    # location [1].
    ep_i = (Xu - norm(rs_i))*tan(Au)

    # Location of the penumbral cone terminator at the projected spacecraft
    # location [1].
    Kp_i = (Xp + norm(rs_i))*tan(Ap)

    # Check if the satellite is under sunlight or umbra / penumbra.
    if dot(rs_i, s_i/norm_s_i) < 0
        if norm(delta_i) > Kp_i
            return SAT_LIGHTING_SUNLIGHT
        else
            if norm(delta_i) < ep_i
                return SAT_LIGHTING_UMBRA
            else
                return SAT_LIGHTING_PENUMBRA
            end
        end
    else
        return SAT_LIGHTING_SUNLIGHT
    end
end


@testset "537.complex-number-multiplication.jl" begin
    @test complex_number_multiply("1+1i", "1+1i") == "0+2i"
    @test complex_number_multiply("1+-1i", "1+-1i") == "0+-2i"
end


function real_to_parametric(p::Point2D, poly::QuadraticPolygon{N, 2, T}) where {N, T}
    return real_to_parametric(p, poly, 30)
end

function real_to_parametric(p::Point2D{T}, poly::QuadraticPolygon{N, 2, T},
                            max_iters::Int64) where {N, T}
    # Convert from real coordinates to the polygon's local parametric coordinates using
    # Newton-Raphson.
    # If a conversion doesn't exist, the minimizer is returned.
    # Initial guess at polygon centroid
    if N === 6 # Triangle
        rs = SVector{2,T}(1//3, 1//3)
    else # Quadrilateral
        rs = SVector{2,T}(1//2, 1//2)
    end
    for i ∈ 1:max_iters
        Δrs = inv(𝗝(poly, rs[1], rs[2]))*(p - poly(rs[1], rs[2]))
        if Δrs ⋅ Δrs < T((1e-8)^2)
            break
        end
        rs += Δrs
    end
    return Point2D{T}(rs[1], rs[2])
end


import NumericalIntegration

"""
    trapz(x, y)

Compute the area under the curve of 2-dimensional points (x, y) using
the trapezoidal method.
"""
function trapz end

trapz(x, y) = trapz(promote(x, y)...)

function trapz(
        x::AbstractVector{T},
        y::AbstractVector{T},
        )::T where T
    if length(x) != length(y)
        error("length(x) != length(y)")
    end
    if length(x) == 0
        error("length(x) == 0")
    end
    if !all(x .== sort(x; rev = false))
        error("x needs to be sorted in ascending order")
    end
    twoI::T = zero(T)
    for k = 2:length(x)
        twoI += ( y[k] + y[k-1] ) * ( x[k] - x[k-1] )
    end
    I_verify::T = NumericalIntegration.integrate(
        x,
        y,
        NumericalIntegration.Trapezoidal(),
        )
    @assert isapprox(twoI/2, I_verify; atol=0.00000001)
    return I_verify
end


# Tests for Cost DEA Models
@testset "CostDEAModel" begin

    ## Test Cost DEA Model with Cooper et al. (2007)
    # Test agains book results
    X = [3 2; 1 3; 4 6]
    Y = [3; 5; 6]
    W = [4 2; 4 2; 4 2]

    deacostcooper = deacost(X, Y, W, rts = :CRS)
    @test efficiency(deacostcooper, :Economic)   ≈ [0.375; 1; 0.429] atol = 1e-3
    @test efficiency(deacostcooper, :Technical)  ≈ [0.9  ; 1; 0.6  ] atol = 1e-3
    @test efficiency(deacostcooper, :Allocative) ≈ [0.417; 1; 0.714] atol = 1e-3

    ## Test Cost DEA Model with Zofío and Prieto (2006) data.
    # Test agains results with R
    X = [5 3; 2 4; 4 2; 4 8; 7 9]
    Y = [7 4; 10 8; 8 10; 5 4; 3 6]
    W = [2 1; 2 1; 2 1; 2 1; 2 1.0]

    # Cost CRS
    deacostcrs = deacost(X, Y, W, rts = :CRS)
    @test efficiency(deacostcrs) ≈ [0.4307692308;
                                    1.000;
                                    1.000;
                                    0.250;
                                    0.2608695652]
    @test efficiency(deacostcrs, :Technical) ≈ [0.6364;
                                1.000;
                                1.000;
                                0.250;
                                0.2609]  atol = 1e-3
    @test efficiency(deacostcrs, :Allocative) ≈ [0.6769230769;
                                1.000;
                                1.000;
                                1.000;
                                1.000]

    # Cost VRS
    deacostvrs = deacost(X, Y, W, rts = :VRS)
    @test efficiency(deacostvrs) ≈ [0.6153846154;
                                    1.000;
                                    1.000;
                                    0.500;
                                    0.3478260870]
    @test efficiency(deacostvrs, :Technical) ≈ [0.750;
                                1.000;
                                1.000;
                                0.500 ;
                                0.375]  atol = 1e-3
    @test efficiency(deacostvrs, :Allocative) ≈ [0.8205128205;
                                1.000;
                                1.000;
                                1.000;
                                0.9275362319]

    # Check defaults
    @test efficiency(deacost(X, Y, W)) == efficiency(deacostvrs)
    @test efficiency(deacostvrs, :Economic) == efficiency(deacostvrs)

    # Print
    show(IOBuffer(), deacostcooper)

    # Test errors
    @test_throws ErrorException deacost([1; 2 ; 3], [4 ; 5], [1; 1; 1]) #  Different number of observations
    @test_throws ErrorException deacost([1; 2; 3], [4; 5; 6], [1; 2; 3], rts = :Error) # Invalid returns to scale
    @test_throws ErrorException deacost([1; 2; 3], [4; 5; 6], [1; 2; 3; 4]) # Different number of observation in prices
    @test_throws ErrorException deacost([1 1; 2 2; 3 3 ], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3]) # Different number of input prices and inputs

end


module TotalVariation

using  DSP, SparseArrays, LinearAlgebra
export gstv, tv
#See ``Total Variation Denoising With Overlapping Group Sparsity'' by
# Ivan Selesnick and Po-Yu Chen (2013)

#Group Sparse Total Variation Denoising
function gstv(y::Vector{Float64}, k::Int, λ::Float64;
    show_cost::Bool=false, iter::Int=100)
    #Initialize Solution
    N  = length(y)
    x  = copy(y)

    #Differential of input
    b = diff(y)

    #Precalculate D D' where D is the first-order difference matrix
    DD::SparseMatrixCSC{Float64,Int} = spdiagm(-1=>-ones(N-2), 0=>2*ones(N-1), 1=>-ones(N-2))

    #Value To Prevent Singular Matrices
    epsilon = 1e-15

    #Convolution Mask - spreads D x over a larger area
    #This regularizes the problem with applying a gradient to a larger area.
    #at k=1 the normal total variational sparse solution (for D x) is found.
    h = ones(k)

    for i=1:iter
        u::Vector{Float64}              = diff(x)
        r::Vector{Float64}              = sqrt.(max.(epsilon, DSP.conv(u.^2,h)))
        Λ::Vector{Float64}              = DSP.conv(1 ./ r, h)[k:end-(k-1)]
        F::SparseMatrixCSC{Float64,Int} = sparse(Diagonal(1 ./ Λ))/λ + DD
        if(show_cost)
            #1/2||y-x||_2^2 + λΦ(Dx)
            #Where Φ(.) is the group sparse regularizer
            println("Cost at iter ",i," is ", 0.5*sum(abs2.(x.-y)) + λ*sum(r))
        end

        tmp::Vector{Float64} = F\b
        dfb::Vector{Float64} = diff(tmp)

        x[1]       = y[1]       + tmp[1]
        x[2:end-1] = y[2:end-1] + dfb[:]
        x[end]     = y[end]     - tmp[end]
    end

    return x
end

#Normal Total Variation Problem
function tv(y::Vector{Float64}, λ::Float64;
    show_cost::Bool=false, iter::Int=100)
    gstv(y, 1, λ, show_cost=show_cost, iter=iter)
end

# package code goes here
end # module


SameXYZLength = Union{Cube,Sphere}
SameXYLength  = Union{Cylinder,SphericalCap,SquarePyramid}
NoSameLength  = Union{Ellipsoid,EllipticCylinder,Parallelepiped,
                      RectangularPyramid}

halflengths(s::SameXYZLength   ) = Vec(s[1], s[1], s[1])
halflengths(s::SameXYLength    ) = Vec(s[1], s[1], s[2])
halflengths(s::NoSameLength    ) = Vec(s[1], s[2], s[3])
halflengths(s::HollowCylinder  ) = Vec(s[1], s[1], s[3])
halflengths(s::TSP             ) = Vec(s[1], s[1], s[2] * s[3])
halflengths(s::Ring            ) = Vec(s[1] + s[2], s[1] + s[2], s[3])
halflengths(s::Torus           ) = Vec(s[1] + s[2], s[1] + s[2], s[2])
halflengths(s::TriangularToroid) = Vec(s[1] + s[2], s[1] + s[2], s[3])

"""    in(p::Point, s::Shape)

Returns `true` if the point `p` lies inside or in the boundary of `s`,
or `false` otherwise.
"""
in(p::Point, s::Sphere            ) = p.x^2 + p.y^2 + p.z^2 ≤ s.r^2

in(p::Point, s::Ellipsoid         ) = (p.x / s.a)^2 + (p.y / s.b)^2 +
                                      (p.z / s.c)^2 ≤ 1

in(p::Point, s::Cylinder          ) = (p.x^2 + p.y^2 ≤ s.r^2) &&
                                      (abs(p.z) ≤ s.c)

in(p::Point, s::EllipticCylinder  ) = (p.x / s.a)^2 + (p.y / s.b)^2 ≤ 1 &&
                                      (abs(p.z) ≤ s.c)

in(p::Point, s::Cube              ) = (abs(p.x) ≤ s.a) && (abs(p.y) ≤ s.a) &&
                                      (abs(p.z) ≤ s.a)

in(p::Point, s::Parallelepiped    ) = (abs(p.x) ≤ s.a) && (abs(p.y) ≤ s.b) &&
                                      (abs(p.z) ≤ s.c)

in(p::Point, s::SquarePyramid     ) = (p.z + s.m * abs(p.x) ≤ s.b) &&
                                      (p.z + s.m * abs(p.y) ≤ s.b) &&
                                      (abs(p.z) ≤ s.b)

in(p::Point, s::RectangularPyramid) = (p.z + s.ma * abs(p.x) ≤ s.c) &&
                                      (p.z + s.mb * abs(p.y) ≤ s.c) &&
                                      (abs(p.z) ≤ s.c)

function in(p::Point, s::TSP)
    z = s.r * s.b
    b = 2 * s.b - z
    return (p.z + s.m * abs(p.x) ≤ b) &&
           (p.z + s.m * abs(p.y) ≤ b) &&
           (abs(p.z) ≤ z)
end


# FitzHugh_Nagumo_2D_compute_spiral_wave_PS_v2.jl
# ================================================   u.p. 4.9.21

# Computes and saves spiral wave solution of the Fitzhugh-Nagumo model.
# Used for figures 11.4 and 11.6, but also in the animation spiralwave in animations/11.

using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, OrdinaryDiffEq
using Random
using Statistics

using Logging: global_logger
using TerminalLoggers: TerminalLogger
global_logger(TerminalLogger())

function FitzHugh_Nagumo_ODE(du,u,p,t)
    a, b, d, ε, hsq6 = p
    uu = @view u[1,:,:]
    vv = @view u[2,:,:]

    # Create padded u matrix to incorporate Newman boundary conditions - 9-point stencil
    uuu=[ [uu[2,2] uu[2,:]' uu[2,end-1] ] ; [ uu[:,2] uu uu[:,end-1] ] ; [ uu[end-1,2] uu[end-1,:]' uu[end-1,end-1] ] ]
    diff_term = d .* (
                4 .* (uuu[2:end-1,1:end-2] .+ uuu[2:end-1,3:end] .+
                        uuu[3:end,2:end-1] .+ uuu[1:end-2,2:end-1] ) .+
                        uuu[3:end,3:end] .+ uuu[3:end,1:end-2] .+ uuu[1:end-2,3:end] .+
                        uuu[1:end-2,1:end-2] .- 20 .*uu  )  ./ hsq6

    du[1,:,:] = @. a*uu*(1 - uu)*(uu - b) - vv + diff_term
    du[2,:,:] = @. ε*(uu - vv)
end

L = 300.0  #  size of domain
N = 300   #  no. of grid points
h = L/N # spatial steps size
hsq6 = 6*h*h

# parameters
a = 3.
b = 0.2
d = 1.    #  diffusion
ε = 0.01

#  initial conditions
uu = 0.5 .* ones(N,N)
vv = 0.5 .* ones(N,N)
# This initial condition sets the spiral wave
uu[1:end, 1:(2N÷3)] .= 0
vv[1:(2N÷3 - N÷12), 1:end] .= 0

u0 = zeros(2,N,N)
u0[1,:,:] = uu
u0[2,:,:] = vv

# solve ODE system
p = a, b, d, ε, hsq6

saveat = 0.0:600.0
tspan = (0.0, saveat[end])

prob = ODEProblem(FitzHugh_Nagumo_ODE, u0, tspan, p)
@time sol = solve(prob, Tsit5();
    reltol=1e-6, abstol=1e-6, maxiters = Inf, saveat = saveat,
    progress = true, progress_steps = 10,
)
tvec = sol.t
uout = sol.u

params = @ntuple a b d ε hsq6 N L
data = @strdict tvec uout params

wsave(datadir("FitzHugh", "spiralwave.jld2"), data)


include("../src/EDKit.jl")
include("deprecate/ExactDiagonalization.jl")
using .EDKit
using LinearAlgebra
using Test

#-------------------------------------------------------------------------------------------------------------------------
# Tesnsor Basis
#-------------------------------------------------------------------------------------------------------------------------
println("--------------------------------------------------")
println("                   Tensor Basis                   ")
println("--------------------------------------------------")

const X = [0 1; 1 0]
const Y = [0 -1; 1 0] * 1im
const Z = [1 0; 0 -1]

directkron(mat, i, j, L, base=2) = kron(I(base^(i-1)), mat, I(base^(L-j)))
boundarykron(m1, m2, i, j, base=2) = kron(m1, I(base^(j-i+1)), m2)

@testset "Base-2 Random Kron" begin
    L = 10
    for i=1:10
        mats = [rand(2^i, 2^i) for j=1:L-i+1]
        inds = [j:j+i-1 for j=1:L-i+1]
        M_test = operator(mats, inds, L) |> Array
        M_targ = sum(directkron(mats[j], j, j+i-1, L) for j = 1:L-i+1)
        @test M_test ≈ M_targ
    end
end

@testset "Base-3 Random Kron" begin
    L = 6
    for i=1:6
        mats = [rand(3^i, 3^i) for j=1:L-i+1]
        inds = [j:j+i-1 for j=1:L-i+1]
        M_test = operator(mats, inds, L) |> Array
        M_targ = sum(directkron(mats[j], j, j+i-1, L, 3) for j = 1:L-i+1)
        @test M_test ≈ M_targ
    end
end

@testset "Base-3 Boundary Kron" begin
    L = 6
    for i = 1:3, j = 1:3
        m1, m2 = rand(3^i, 3^i), rand(3^j, 3^j)
        mats = kron(m1, m2)
        inds = vcat(Vector(1:i), Vector(L-j+1:L))
        M_test = operator(mats, inds, L) |> Array
        M_targ = boundarykron(m1, m2, i+1, L-j, 3)
        @test M_test ≈ M_targ
    end
end

@testset "Operator Multiplication" begin
    L = 10
    mat = kron(I(2), X, Y, Z)
    v = rand(2^L)
    opt = trans_inv_operator(mat, 4, L)
    v2 = opt * v
    optm = Array(opt)
    v3 = optm * v
    @test v2 ≈ v3
end

#-------------------------------------------------------------------------------------------------------------------------
# Projected Basis
#-------------------------------------------------------------------------------------------------------------------------
println("--------------------------------------------------")
println("                 Projected Basis                  ")
println("--------------------------------------------------")

@testset "Spin-1/2 XY" begin
    L = 10
    mat = spin((1, "xx"), (1, "yy"))
    E = zeros(2^L)
    P = 0
    for n = 0:L
        basis = ProjectedBasis(x->sum(x)==n, L)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 2, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    @test P == 2^L
    vals = trans_inv_operator(mat, 2, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

@testset "Spin-1/2 Random" begin
    L = 10
    mat = begin
        M = zeros(ComplexF64, 8, 8)
        M[1,1] = rand()
        M[8,8] = rand()
        M[[2,3,5], [2,3,5]] = rand(ComplexF64, 3, 3) |> Hermitian
        M[[4,6,7], [4,6,7]] = rand(ComplexF64, 3, 3) |> Hermitian
        M
    end
    E = zeros(2^L)
    P = 0
    for n = 0:L
        basis = ProjectedBasis(x->sum(x)==n, L)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 3, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    @test P == 2^L
    vals = trans_inv_operator(mat, 3, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

@testset "PXP_OBC" begin
    function fib(n::Integer)
        if n == 0
            return 1
        elseif n == 1
            return 2
        end
        a, b = 1, 2
        for m = 2:n
            a, b = b, a+b
        end
        b
    end
    mat = begin
        P = Diagonal([1, 1, 1, 0, 1, 1, 0, 0])
        P * kron(I(2), X, I(2)) * P
    end
    pxpf(v::Vector{<:Integer}) = all(v[i]==0 || v[i+1]==0 for i=1:length(v)-1)
    for L = 3:20
        basis = ProjectedBasis(pxpf, L)
        @test size(basis, 1) == fib(L)
    end
    for L = 3:10
        basis = ProjectedBasis(pxpf, L)
        mats = fill(mat, L-2)
        inds = [[i,i+1, i+2] for i=1:L-2]
        E = operator(mats, inds, basis) |> Array |> Hermitian |> eigvals
        H = operator(mats, inds, L) |> Array
        vals = H[basis.I, basis.I] |> Hermitian |> eigvals
        @test norm(vals - E) ≈ 0.0
    end
end

@testset "PXP_PBC Projected Basis" begin
    mat = begin
        P = Diagonal([1, 1, 1, 0, 1, 1, 0, 0])
        P * kron(I(2), X, I(2)) * P
    end
    pxpf(v::Vector{<:Integer}) = all(v[i]==0 || v[mod(i, length(v))+1]==0 for i=1:length(v))
    for L = 3:14
        basis = ProjectedBasis(pxpf, L)
        E = trans_inv_operator(mat, 3, basis) |> Array |> Hermitian |> eigvals
        H = trans_inv_operator(mat, 3, L) |> Array
        vals = H[basis.I, basis.I] |> Hermitian |> eigvals
        @test norm(vals - E) ≈ 0.0
    end
end

@testset "Spin-1 Random" begin
    L = 6
    R3 = begin
        m2 = [2, 4, 10]
        m1 = [3, 5, 7, 11, 13, 19]
        m0 = [6, 8, 12, 14, 16, 20, 22]
        p1 = [9, 15, 17, 21, 23, 25]
        p2 = [18, 24, 26]
        mat = zeros(ComplexF64, 27, 27)
        mat[1, 1]    = rand()
        mat[27, 27]  = rand()
        mat[m2, m2] .= rand(ComplexF64, 3, 3) |> Hermitian
        mat[m1, m1] .= rand(ComplexF64, 6, 6) |> Hermitian
        mat[m0, m0] .= rand(ComplexF64, 7, 7) |> Hermitian
        mat[p1, p1] .= rand(ComplexF64, 6, 6) |> Hermitian
        mat[p2, p2] .= rand(ComplexF64, 3, 3) |> Hermitian
        mat
    end
    E = zeros(3^L)
    P = 0
    for n = 0:2L
        basis = ProjectedBasis(x->sum(x)==n, L, base=3)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 3, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    @test P == 3^L
    vals = trans_inv_operator(mat, 3, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

#-------------------------------------------------------------------------------------------------------------------------
# Translation Basis
#-------------------------------------------------------------------------------------------------------------------------
println("--------------------------------------------------")
println("               Translational Basis                ")
println("--------------------------------------------------")

@testset "Spin-1/2 XY" begin
    L = 10
    θ = 0.34
    expθ = exp(-1im*θ)
    mat = spin((expθ, "+-"), (1/expθ, "-+"), (1, "z1"), (1, "1z"))
    E = zeros(2^L)
    P = 0
    for n = 0:L, k = 0:L-1
        basis = TranslationalBasis(x->sum(x)==n, k, L)
        #println(size(basis,1))
        if (l = size(basis,1)) > 0
            vals = trans_inv_operator(mat, 2, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    @test P == 2^L
    vals = trans_inv_operator(mat, 2, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

@testset "Spin-1/2 Random" begin
    L = 10
    mat = begin
        M = zeros(ComplexF64, 8, 8)
        M[1,1] = rand()
        M[8,8] = rand()
        M[[2,3,5], [2,3,5]] = rand(ComplexF64, 3, 3) |> Hermitian
        M[[4,6,7], [4,6,7]] = rand(ComplexF64, 3, 3) |> Hermitian
        M
    end
    E = zeros(2^L)
    P = 0
    for n = 0:L, k = 0:L-1
        basis = TranslationalBasis(x->sum(x)==n, k, L)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 3, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    @test P == 2^L
    vals = trans_inv_operator(mat, 3, L) |> Array |> Hermitian |> eigvals
    @test norm(vals - sort(E)) ≈ 0.0 atol = 1e-10
end

@testset "Spin-1 XY" begin
    L = 6
    h = 1
    mat = spin((1, "xx"), (1, "yy"), (h/2, "1z"), (h/2, "z1"), D=3)
    E = zeros(3^L)
    P = 0
    for n = 0:2L, k = 0:L-1
        basis = TranslationalBasis(x->sum(x)==n, k, L, base=3)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 2, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    @test P == 3^L
    vals = trans_inv_operator(mat, 2, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

@testset "AKLT" begin
    L = 6
    mat = begin
        ss = spin((1, "xx"), (1, "yy"), (1, "zz"), D=3)
        1/2 * ss + 1/6 * ss^2 + 1/3 * I
    end
    E = zeros(3^L)
    P = 0
    for n = 0:2L, k = 0:L-1
        basis = TranslationalBasis(x->sum(x)==n, k, L, base=3)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 2, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    vals = trans_inv_operator(mat, 2, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

@testset "Spin-1 Random" begin
    L = 6
    mat = begin
        m2 = [2, 4, 10]
        m1 = [3, 5, 7, 11, 13, 19]
        m0 = [6, 8, 12, 14, 16, 20, 22]
        p1 = [9, 15, 17, 21, 23, 25]
        p2 = [18, 24, 26]
        M = zeros(ComplexF64, 27, 27)
        M[1, 1]    = rand()
        M[27, 27]  = rand()
        M[m2, m2] .= rand(ComplexF64, 3, 3) |> Hermitian
        M[m1, m1] .= rand(ComplexF64, 6, 6) |> Hermitian
        M[m0, m0] .= rand(ComplexF64, 7, 7) |> Hermitian
        M[p1, p1] .= rand(ComplexF64, 6, 6) |> Hermitian
        M[p2, p2] .= rand(ComplexF64, 3, 3) |> Hermitian
        M
    end
    E = zeros(3^L)
    P = 0
    for n = 0:2L, k = 0:L-1
        basis = TranslationalBasis(x->sum(x)==n, k, L, base=3)
        if (l = size(basis, 1)) > 0
            vals = trans_inv_operator(mat, 3, basis) |> Array |> Hermitian |> eigvals
            E[P+1:P+l] = vals
            P += l
        end
    end
    vals = trans_inv_operator(mat, 3, L) |> Array |> Hermitian |> eigvals
    @test vals ≈ sort!(E)
end

#-------------------------------------------------------------------------------------------------------------------------
# Translation + Parity
#-------------------------------------------------------------------------------------------------------------------------
println("--------------------------------------------------")
println("             Translation Parity Basis             ")
println("--------------------------------------------------")

@testset "Spin-1/2 XY" begin
    mat = spin((1, "+-"), (1, "-+"), (1, "z1"), (1, "1z"))
    for L = 2:2:10
        for n = 0:L, k in [0, L ÷ 2]
            be = TranslationParityBasis(x->sum(x)==n, k, +1, L)
            bo = TranslationParityBasis(x->sum(x)==n, k, -1, L)
            ba = TranslationalBasis(x->sum(x)==n, k, L)
            ve = trans_inv_operator(mat, 2, be) |> Array |> Hermitian |> eigvals
            vo = trans_inv_operator(mat, 2, bo) |> Array |> Hermitian |> eigvals
            va = trans_inv_operator(mat, 2, ba) |> Array |> Hermitian |> eigvals
            E = sort(vcat(ve, vo))
            @test norm(E-va) ≈ 0.0 atol = 1e-12
        end
    end
end

@testset "Spin-1 XY" begin
    mat = spin((1, "+-"), (1, "-+"), (1, "z1"), (1, "1z"), D=3)
    for L = 2:2:6
        for n = 0:2L, k in [0, L ÷ 2]
            be = TranslationParityBasis(x->sum(x)==n, k, +1, L, base=3)
            bo = TranslationParityBasis(x->sum(x)==n, k, -1, L, base=3)
            ba = TranslationalBasis(x->sum(x)==n, k, L, base=3)
            ve = trans_inv_operator(mat, 2, be) |> Array |> Hermitian |> eigvals
            vo = trans_inv_operator(mat, 2, bo) |> Array |> Hermitian |> eigvals
            va = trans_inv_operator(mat, 2, ba) |> Array |> Hermitian |> eigvals
            E = sort(vcat(ve, vo))
            @test norm(E-va) ≈ 0.0 atol = 1e-12
        end
    end
end

@testset "PXP" begin
    mat = begin
        P = Diagonal([1, 1, 1, 0, 1, 1, 0, 0])
        P * kron(I(2), X, I(2)) * P
    end
    pxpf(v::Vector{<:Integer}) = all(v[i]==0 || v[mod(i, length(v))+1]==0 for i=1:length(v))
    for L = 4:2:20
        for k in [0, L ÷ 2]
            be = TranslationParityBasis(pxpf, k, +1, L)
            bo = TranslationParityBasis(pxpf, k, -1, L)
            ba = TranslationalBasis(pxpf, k, L)
            ve = trans_inv_operator(mat, 2, be) |> Array |> Hermitian |> eigvals
            vo = trans_inv_operator(mat, 2, bo) |> Array |> Hermitian |> eigvals
            va = trans_inv_operator(mat, 2, ba) |> Array |> Hermitian |> eigvals
            E = sort(vcat(ve, vo))
            @test norm(E-va) ≈ 0.0 atol = 1e-12
        end
    end
end

@testset "Ising Spin-flip" begin
    zz = spin((1, "zz"))
    x = spin((0.3,"x"))
    for L = 4:2:12
        E = zeros(2^L)
        for k in 0:L-1
            ba = TranslationalBasis(k, L)
            be = TranslationFlipBasis(k, 1, L)
            bo = TranslationFlipBasis(k, -1, L)
            va = trans_inv_operator(zz, 2, ba) + trans_inv_operator(x, 1, ba) |> Array |> Hermitian |> eigvals
            ve = trans_inv_operator(zz, 2, be) + trans_inv_operator(x, 1, be) |> Array |> Hermitian |> eigvals
            vo = trans_inv_operator(zz, 2, bo) + trans_inv_operator(x, 1, bo) |> Array |> Hermitian |> eigvals
            veo = sort(vcat(ve, vo))
            @test norm(va-veo) ≈ 0.0 atol = 1e-12
        end
    end
end


using Compat: cispi
# using DoubleFloats
using FourierTransforms
using Random
using Test

const FT = FourierTransforms

# `DoubleFloat` does not work: <https://github.com/JuliaMath/DoubleFloats.jl/issues/139>
const Types = [Complex{Float32}, Complex{Float64}, Complex{BigFloat}]

# # Random.rand(rng::AbstractRNG, ::Random.Random.Random.CloseOpen01{DoubleFloat}) = ???
# Random.randn(::Type{DoubleFloat}, dims::Integer...) = DoubleFloat.(randn(Float64, dims...))
# Random.randn(::Type{Complex{DoubleFloat}}, dims::Integer...) = Complex{DoubleFloat}.(randn(Complex{Float64}, dims...))
# Base.eps(::Type{DoubleFloat}) = eps(Float64)^2

# Random.rand(rng::AbstractRNG, ::Random.Random.Random.CloseOpen01{BigFloat}) = ???
Random.randn(::Type{BigFloat}, dims::Integer...) = BigFloat.(randn(Float64, dims...))
Random.randn(::Type{Complex{BigFloat}}, dims::Integer...) = Complex{BigFloat}.(randn(Complex{Float64}, dims...))

Random.seed!(100)
@testset "Direct Fourier transforms (T=$T)" for T in Types
    maxbits = T === Complex{BigFloat} ? 8 : 10
    for bits in 1:maxbits
        N = 2^bits
        x = randn(T, N)
        X = FT.direct_ft(x)
        x′ = conj(FT.direct_ft(conj(X)) / N)
        @test x′ ≈ x
    end
end

Random.seed!(100)
@testset "Radix-2 Cooley–Tukey Fast Fourier transforms (T=$T)" for T in Types
    maxbits = T === Complex{BigFloat} ? 15 : 20
    for bits in 1:maxbits
        N = 2^bits
        x = randn(T, N)
        X = FT.ditfft2(x)
        x′ = conj(FT.ditfft2(conj(X)) / N)
        @test x′ ≈ x
    end
end

Random.seed!(100)
@testset "General radix Cooley–Tukey Fast Fourier transforms (T=$T)" for T in Types
    maxbits = T === Complex{BigFloat} ? 15 : 20
    for bits in 1:maxbits
        N = 2^bits
        x = randn(T, N)
        X = FT.radix_fft(x)
        x′ = conj(FT.radix_fft(conj(X)) / N)
        @test x′ ≈ x
    end
end

Random.seed!(100)
@testset "Comparing algorithms for Fourier transforms (T=$T)" for T in Types
    maxbits = T === Complex{BigFloat} ? 8 : 10
    for bits in 1:maxbits
        N = 2^bits
        x = randn(T, N)
        X = FT.direct_ft(x)
        if N == 1
            X1 = FT.direct_ft_1(x)
            @test X1 ≈ X
        end
        if N == 2
            X2 = FT.direct_ft_2(x)
            @test X2 ≈ X
        end
        if N == 4
            X4 = FT.direct_ft_4(x)
            @test X4 ≈ X
        end
        Xct2 = FT.ditfft2(x)
        @test Xct2 ≈ X
        Xr = FT.radix_fft(x)
        @test Xr ≈ X
        for radix in 2:8
            if N % radix == 0
                Xr = FT.radix_fft(x; choose_radix=N -> radix)
                @test Xr ≈ X
            end
        end
    end
end

Random.seed!(100)
@testset "Simple Fourier transforms (T=$T)" for T in Types
    RT = typeof(real(zero(T)))
    maxbits = T === Complex{BigFloat} ? 15 : 20
    for bits in 1:maxbits
        for iter in 1:10
            N = 2^bits
            k = rand(1:N)
            x = [cispi(2 * ((k - 1) * (n - 1) % N) / RT(N)) for n in 1:N]
            X₀ = [(k′ == k) * N for k′ in 1:N]
            X = fft(x)
            @test isapprox(X, X₀; atol=N^2 * sqrt(eps()))
        end
    end
end

Random.seed!(100)
@testset "Inverse Fourier transforms (T=$T)" for T in Types
    maxbits = T === Complex{BigFloat} ? 15 : 20
    for bits in 1:maxbits
        for iter in 1:10
            N = 2^bits
            x = randn(T, N)
            X = fft(x)
            x′ = inv_fft(X)
            @test x′ ≈ x
        end
    end
end

Random.seed!(100)
@testset "Linearity of Fourier transforms (T=$T)" for T in Types
    maxbits = T === Complex{BigFloat} ? 10 : 15
    for bits in 1:maxbits
        for iter in 1:10
            N = 2^bits
            x = randn(T, N)
            y = randn(T, N)
            α = randn(T)
            z = α * x + y
            X = fft(x)
            Y = fft(y)
            Z = fft(z)
            @test α * X + Y ≈ Z
        end
    end
end


# ---
# title: 963. Minimum Area Rectangle II
# id: problem963
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Geometry
# link: <https://leetcode.com/problems/minimum-area-rectangle-ii/description/>
# hidden: true
# ---
#
# Given a set of points in the xy-plane, determine the minimum area of **any**
# rectangle formed from these points, with sides **not necessarily parallel** to
# the x and y axes.
#
# If there isn't any rectangle, return 0.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2018/12/21/1a.png)
#
#
#
#     Input: [[1,2],[2,1],[1,0],[0,1]]
#     Output: 2.00000
#     Explanation: The minimum area rectangle occurs at [1,2],[2,1],[1,0],[0,1], with an area of 2.
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2018/12/22/2.png)
#
#
#
#     Input: [[0,1],[2,1],[1,1],[1,0],[2,0]]
#     Output: 1.00000
#     Explanation: The minimum area rectangle occurs at [1,0],[1,1],[2,1],[2,0], with an area of 1.
#
#
# **Example 3:**
#
# ![](https://assets.leetcode.com/uploads/2018/12/22/3.png)
#
#
#
#     Input: [[0,3],[1,2],[3,1],[1,3],[2,1]]
#     Output: 0
#     Explanation: There is no possible rectangle to form from these points.
#
#
# **Example 4:**
#
# ![](https://assets.leetcode.com/uploads/2018/12/21/4c.png)
#
#
#
#     Input: [[3,1],[1,1],[0,1],[2,1],[3,3],[3,2],[0,2],[2,3]]
#     Output: 2.00000
#     Explanation: The minimum area rectangle occurs at [2,1],[2,3],[3,3],[3,1], with an area of 2.
#
#
#
#
# **Note:**
#
#   1. `1 <= points.length <= 50`
#   2. `0 <= points[i][0] <= 40000`
#   3. `0 <= points[i][1] <= 40000`
#   4. All points are distinct.
#   5. Answers within `10^-5` of the actual value will be accepted as correct.
#
#
## @lc code=start
using LeetCode

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


# NegativeBinomial is the distribution of the number of failures
# before the r-th success in a sequence of Bernoulli trials.
# We do not enforce integer size, as the distribution is well defined
# for non-integers, and this can be useful for e.g. overdispersed
# discrete survival times.

immutable NegativeBinomial <: DiscreteUnivariateDistribution
    r::Float64
    prob::Float64

    function NegativeBinomial(r::Real, p::Real)
        zero(p) < p <= one(p) || error("prob must be in (0, 1].")
        zero(r) < r || error("r must be positive.")
        new(float64(r), float64(p))
    end

    NegativeBinomial() = new(1.0, 0.5)
end

@_jl_dist_2p NegativeBinomial nbinom

insupport(::NegativeBinomial, x::Real) = isinteger(x) && zero(x) <= x
insupport(::Type{NegativeBinomial}, x::Real) = isinteger(x) && zero(x) <= x

function mgf(d::NegativeBinomial, t::Real)
    r, p = d.size, d.prob
    return ((1.0 - p) * exp(t))^r / (1.0 - p * exp(t))^r
end

function cf(d::NegativeBinomial, t::Real)
    r, p = d.size, d.prob
    return ((1.0 - p) * exp(im * t))^r / (1.0 - p * exp(im * t))^r
end

function mean(d::NegativeBinomial)
    p = d.prob
    (1.0 - p) * d.r / p
end

function var(d::NegativeBinomial)
    p = d.prob
    (1.0 - p) * d.r / (p * p)
end

function std(d::NegativeBinomial)
    p = d.prob
    sqrt((1.0 - p) * d.r) / p
end

function skewness(d::NegativeBinomial)
    p = d.prob
    (2.0 - p) / sqrt((1.0 - p) * r)
end

function kurtosis(d::NegativeBinomial)
    p = d.prob
    6.0 / r + (p * p) / ((1.0 - p) * r)
end

function mode(d::NegativeBinomial)
    p = d.prob
    ifloor((1.0 - p) * (d.r - 1.) / p)
end

modes(d::NegativeBinomial) = [mode(d)]


# Noise radius illustration using e.g. Poincare section
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot, Random

ro = Systems.roessler([0.1,0.2,0.1])
fig, axs = subplots(2, 1; figsize = (0.4*figx, 1.5figy))
plane = (2, 0.0)

err = (1e-4, 1e-12)
εs = 10 .^ (-5:0.5:1)

C = ["C0", "C2"]

for (i, e) ∈ enumerate(err)
    p = poincaresos(
        ro, plane, 10000; Ttr = 100.0,
        rootkw = (xrtol = e, atol = e)
    )

    axs[1].scatter(p[:, 1], p[:, 3]; s = 20, alpha = 0.75/i^2,
    label = "tol = \$10^{$(round(Int, log10(e)))}\$", color = C[i])
    Cs = correlationsum(p, εs)
    y = log.(Cs ./ maximum(Cs))
    axs[2].plot(log.(εs), y; color = C[i])

    # add noise line
    if i == 1
        z = log.(εs)[end-3]
        axs[2].plot([z, z], [-14, y[end-3]]; color = "C1", lw = 2, ls = ":")
        axs[2].text(z+0.2, -9, "\$\\sigma\$"; color = "C1")
    end
end

axs[1].set_xlabel("\$x\$";labelpad = -15)
axs[1].set_ylabel("\$z\$")
axs[1].legend(markerscale = 5, fontsize=26)
axs[1].set_yticks(0:6:18)

axs[2].plot([-9, 0], 0.9 .* [-9, 0]; ls = "--", color = C[2])
axs[2].text(-7, -4, "0.9"; color = C[2])
axs[2].plot([-6, -3], 2.0 .* [-6, -3] .- 2; ls = "--", color = C[1])
axs[2].text(-3, -11, "2"; color = C[1])
axs[2].set_ylabel("\$\\log(C)\$"; labelpad = -10)
axs[2].set_xlabel("\$\\log(\\varepsilon)\$"; labelpad = -12)
axs[2].set_xticks(-12:5:2)
axs[2].set_yticks(-14:5:2)

fig.tight_layout(pad = 0.5)
add_identifiers!(fig)
wsave(plotsdir("5", "fractaldim_noise"), fig)