# # Van der Pol oscillator
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/VanDerPol.ipynb)
#
#md # !!! note "Overview"
#md #     System type: polynomial continuous system\
#md #     State dimension: 2\
#md #     Application domain: Nonlinear physics
#
# The Van der Pol oscillator was introduced by the Dutch physicist Balthasar van der Pol.
# This is a famous model, typically investigated in the study of nonlinear dynamics.
# The model presents non-conservative oscillations with non-linear damping.
# In the past, it has been of relevance in several practical problems of engineering such as
# circuits containing vacuum tubes. For more information on the model see the wikipedia article
# [Van der Pol oscillator](https://en.wikipedia.org/wiki/Van_der_Pol_oscillator).

# In this notebook we compute the evolution of the limit cycle in the phase plane
# for a set of initial conditions. For such set, we consider a *safety condition*
# originally from [^ARCHCOMP19LIN], that aims at verifying that there is no solution
# starting from the initial set that exceeds a prescribed upper bound on the
# velocity $y(t) := x'(t)$ at all times in the given time span. Moreover, we also
# illustrate the computation of an invariant set using the obtained flowipe.
# Finally, we study the limit cycle à la Poincaré, by defining a function that computes
# the cross section of the flowpipe at each revolution. Interestingly, this method
# gives an algorithmic proof that the safety condition obtained previously is actually
# verified at all times, i.e. over the unbounded time horizon $[0, \infty)$.

# ## Dynamics
#
# The dynamics of the Van der Pol oscillator are described by the following ODE with
# two variables:
#
# ```math
# \begin{aligned}
#   \dot{x} &= y \\
#   \dot{y} &= \mu (1 - x^2) y - x
# \end{aligned}
# ```
# The system has a stable limit cycle. Such limit cycle becomes increasingly sharp
# for higher values of ``μ``. Here we consider the parameter ``μ = 1``.

#!jl using ReachabilityAnalysis, Plots

@taylorize function vanderpol!(dx, x, params, t)
    local μ = 1.0
    dx[1] = x[2]
    dx[2] = (μ * x[2]) * (1 - x[1]^2) - x[1]
    return dx
end

#md # !!! tip "Performance tip"
#md #     Advanced users may want to consider the section [Some common gotchas](@ref)
#md #     to improve the code that is generated by `@taylorize` by re-writing
#md #     the right-hand side of `dx[2]` with the use of auxiliary variables.

#jl function vanderpol()
#jl    X0 = Hyperrectangle(low=[1.25, 2.35], high=[1.55, 2.45])
#jl    ivp = @ivp(x' = vanderpol!(x), dim: 2, x(0) ∈ X0)
#jl    tspan = (0.0, 5.0)
#jl    return ivp, tspan
#jl end

# ## Safety verification
#
# We set the initial condition ``x(0) ∈ [1.25, 1.55]``, ``y(0) ∈ [2.35,2.45]``.
# The *unsafe set* is given by ``y ≥ 2.75`` for a time span ``[0, 7]``.
# In other words, we would like to prove that there doesn't exist a solution of
# the model with a ``y`` value which is greater than 2.75, for any initial condition
# on the given domain. The time horizon of ``T = 7`` is chosen such that the oscillator
# can do at least one full cycle.

# We proceed by defining the initial conditions as a hyperrectangular set according
# to the problem's specifications. Then we build the initial-value problem and pass
# it to the `solve` function. We specify using `TMJets` algorithm with default options.

X0 = Hyperrectangle(low=[1.25, 2.35], high=[1.55, 2.45]) #!jl
prob = @ivp(x' = vanderpol!(x), dim=2, x(0) ∈ X0) #!jl
sol = solve(prob, T=7.0, alg=TMJets(abstol=1e-12)); #!jl

# For further computations, it is convenient to work with a [zonotopic](https://juliareach.github.io/ReachabilityAnalysis.jl/dev/man/basics/#Zonotopes-1)
# overapproximation of the flowpipe.

solz = overapproximate(sol, Zonotope); #!jl

# The maximum value of variable ``y`` is obtained by computing the [support function](https://juliareach.github.io/ReachabilityAnalysis.jl/dev/man/basics/#Representing-sets-with-support-functions-1)
# of the flowpipe along direction ``[0, 1]``:

ρ([0.0, 1.0], solz) #!jl

# That shows that the property is satisfied. Below we plot the flowpipe in the
# x-y plane, together with the horizontal line ``y = 2.75``.

plot(solz, vars=(1, 2), lw=0.2, xlims=(-2.5, 2.5), xlab="x", ylab="y") #!jl
plot!(x -> 2.75, color=:red, lab="y = 2.75", style=:dash, legend=:bottomright) #!jl

# We can also plot the state variables ``x(t)`` and ``y(t)`` as a function of time
# (recall that `0` in `vars` is used to denote the time variable):

plot(solz, vars=(0, 1), lw=0.2, xlab="t", ylab="x") #!jl

#-

plot(solz, vars=(0, 2), lw=0.2,  xlab="t", ylab="y") #!jl

# ## Invariant Set
#
# We can use the reachability result to examine an invariant of the system. In
# other words, we can algorithmically prove that the flowpipe re-enters from where
# it started after giving one loop, using inclusion checks.

plot(solz, vars=(1, 2), lw=0.2, xlims=(0.0, 2.5), ylims=(1.6, 2.8), xlab="x", ylab="y") #!jl
plot!(X0, color=:orange, lab="X0") #!jl
plot!(solz[1:13], vars=(1, 2), color=:green, lw=1.0, alpha=0.5, lab="F[1:13]") #!jl
plot!(solz[388], vars=(1, 2), color=:red, lw=1.0, alpha=0.6, lab="F[388]") #!jl

# It is seen that the reach-set corresponding to the time-span

tspan(solz[388])  #!jl

# is included in the set union ``F[1] \cup \cdots \cup F[13]`` of previously
# computed reach-sets. Notice that all future trajectories starting from
# the 388-th reach-set are already covered by the flowpipe. Therefore, we can claim
# that an *invariant set* was found.

# ## Limit cycle

# To examine the limit cycle we can intersect a line segment perpendicular to the
# flowpipe, that will allow us to get a cross-section of the sets in order to prove
# that after one cycle the intersection segment actually shrinks. This approach
# is similar to the method of [Poincaré sections](https://en.wikipedia.org/wiki/Poincar%C3%A9_map).

line = LineSegment([1, 2.], [2., 2.5]) #!jl
plot(solz, vars=(1, 2), lw=0.2, xlims=(0.0, 2.5), ylims=(1.6, 2.8), xlab="x", ylab="y") #!jl
plot!(X0, color=:orange, lab="X0") #!jl
plot!(solz[1:13], vars=(1, 2), color=:green, lw=1.0, alpha=0.5, lab="F[1:13]") #!jl
plot!(solz[388], vars=(1, 2), color=:red, lw=1.0, alpha=0.6, lab="F[388]") #!jl
plot!(line, lw=2.0) #!jl

# Then we can define a function to get the cross section of the flowpipe. The
# function needs the flowpipe, a line segment that cuts the flowpipe and the
# indices of the subsets to cut.

function cross_section(line::LineSegment, fp, idx) #!jl
    p = VPolygon() #!jl
    for i in idx #!jl
        x = intersection(line, set(fp[i])) #!jl
        if !isempty(x)  #!jl
            p = convex_hull(p, x) #!jl
        end  #!jl
    end #!jl
    vl = vertices_list(p) #!jl
    @assert length(vl) == 2 #!jl
    return LineSegment(vl[1], vl[2]) #!jl
end #!jl

# Then we can get the cross section of the first five sets and the last set,
# calling them `i1` and `i2` respectively.

ifirst = cross_section(line, solz, 1:13) #!jl
ilast = cross_section(line, solz, [388]) #!jl

# We can also calculate the length of each cross section, remember that the
# system is 2D, so the cross section will be a line segment.

lfirst = norm(ifirst.q - ifirst.p) #!jl
#-
llast = norm(ilast.q - ilast.p); #!jl

#-

plot(ifirst, lw=3.0, alpha=1.0, label="First subsets", legend=:bottomright) #!jl
plot!(ilast, lw=5.0, alpha=1.0, label="Last subset") #!jl

#-

# The inclusion check succeeds:

ilast ⊆ ifirst #!jl

# We can see, the cross section of the last subset is a subset of the first few
# sets, thus, the cycle will continue, presumably getting smaller each revolution.

# ## References

# [^ARCHCOMP19LIN]: Althoff, M., Bak, S., Forets, M., Frehse, G., Kochdumper, N., Ray, R., ... & Schupp, S. (2019, May). ARCH-COMP19 Category Report: Continuous and Hybrid Systems with Linear Continuous Dynamics. In ARCH@ CPSIoTWeek (pp. 14-40).


include("splinesParallel.jl")

Su0p = plarBezier(S1, [[0,0,0],[10,0,0]])
Su1p = plarBezier(S1, [[0,10,0],[2.5,10,3],[5,10,-3],[7.5,10,3],[10,10,0]])
Sv0p = plarBezier(S2, [[0,0,0],[0,0,3],[0,10,3],[0,10,0]])
Sv1p = plarBezier(S2, [[10,0,0],[10,5,3],[10,10,0]])

coons = plarCoonsPatch([Su0p,Su1p,Sv0p,Sv1p])
point = [0.2, 0.7]
coons(point)


# Load Julia packages (libraries)

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

df = CSV.read(sr_datadir("chimpanzees.csv"), DataFrame);

# Define the Stan language model

stan10_4sl = "
data {
    int N;
    int<lower=0, upper=1> L[N];
    vector[N] P;
    vector[N] C;

    int<lower=1, upper=N> N_chimps;
    int<lower=1, upper=N_chimps> chimp[N];
}
parameters {
    real a_chimp[N_chimps];
    real bp;
    real bpc;
}
model {
    vector[N] p;
    target += normal_lpdf(a_chimp | 0, 10);
    target += normal_lpdf(bp | 0, 10);
    target += normal_lpdf(bpc | 0, 10);
    for (i in 1:N) p[i] = a_chimp[chimp[i]] + (bp + bpc * C[i]) * P[i];
    target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
    vector[N] log_lik;
    {
        vector[N] p;
        for(n in 1:N) {
          p[n] = a_chimp[chimp[n]] + (bp + bpc * C[n]) * P[n];
          log_lik[n] = binomial_logit_lpmf(L[n] | 1, p[n]);
        }
    }
}
";

# Define the Stanmodel and set the output format to :mcmcchains.

m10_4sl = SampleModel("m10.4sl", stan10_4sl);

# Input data for cmdstan

m10_4_data = Dict("N" => size(df, 1), "N_chimps" => length(unique(df[!, :actor])),
"chimp" => df[!, :actor], "L" => df[!, :pulled_left],
"P" => df[!, :prosoc_left], "C" => df[!, :condition]);

# Sample using cmdstan

rc10_4sl = stan_sample(m10_4sl, data=m10_4_data);

# Result rethinking

rethinking = "
      mean   sd  5.5% 94.5% n_eff Rhat
bp    0.84 0.26  0.43  1.26  2271    1
bpC  -0.13 0.29 -0.59  0.34  2949    1

a[1] -0.74 0.27 -1.16 -0.31  3310    1
a[2] 10.88 5.20  4.57 20.73  1634    1
a[3] -1.05 0.28 -1.52 -0.59  4206    1
a[4] -1.05 0.28 -1.50 -0.60  4133    1
a[5] -0.75 0.27 -1.18 -0.32  4049    1
a[6]  0.22 0.27 -0.22  0.65  3877    1
a[7]  1.81 0.39  1.22  2.48  3807    1
";

# Update sections

if success(rc10_4sl)
    part10_4sl = read_samples(m10_4sl, :particles)
    part10_4sl |> display
end


@testset "kernelembeddings" begin
    rng = MersenneTwister(42)
    N = 3
    measure = MvNormal(ones(3), ones(3))
    k = SqExponentialKernel()
    l = 2.0
    σ = 0.5
    ke = KernelEmbedding(k, σ, l, measure)
    kernel = σ * with_lengthscale(k, l)
    @test BQ.scale(ke) == σ

    sample = [rand(rng, 3)]
    @test kernel_mean(ke, sample) ≈ [mean(kernel.(sample, eachcol(rand(rng, measure, 10000))))] atol=1e-2
    @test kernel_variance(ke) ≈ mean(kernel.(eachcol(rand(rng, measure, 10000)), eachcol(rand(rng, measure, 10000)))) atol=1e-2

    s = 2.0
    l = [1.0, 2.0]
    L = LowerTriangular(rand(2, 2))
    @test BQ.Λ(s) ≈ s^2 * I
    @test BQ.Λ(l) ≈ Diagonal(abs2.(l))
    @test Matrix(BQ.Λ(L)) ≈ L * L'
end


"""
    ZAdjoin

  Let α be an element satisfying a monic polynomial of degree n > 1 over the integers, Z. That is,

  ```
  α^n = p[1] + p[2]α + ... + p[n]α^(n-1),  (I)
  ```

  where the p[i]ϵZ. Let Z[α] be the ring formed by adjoining α to the Z. This module provides code for basic representation and manipulation of such rings.

  Elements of Z[α] can be represented as polynomials, q[1] + q[2]α + ... + q[n]α^(n-1), of degree n-1 in α, since higher powers of α can be reduced modulo the defining relationship, α^n=p(α). Thus, under addition, Z[α] is a module of dimension n over Z. Since, under multiplication, elements of Z[α] induce linear transformations of this module, elements may also be represented as matrices over Z. Specifically, α may be represented as the companion matrix of the defining polynomial, α^n-p(α), powers of α up to n-1 as corresponding powers of this companion matrix, and general elements as integer combinations of these powers.

  A defining relationship, essentially n integer coefficients p[1],...,p[n] as in equation (I) above, is represented by one of the types Modulus, KummerModulus, or QIModulus (QI for Quadratic Integer,) all concrete subtypes of AbstractModulus. Ring elements are represented as integer vectors and a modulus (instance of a concrete subtype of AbstractModulus) combined in the Element type. Addition, subtraction, multiplication, and exponentiation by integer powers of Elements are supported. Conversion of Elements to matrices is also supported.

  # Examples

  ```
  # A Kummer Ring of degree 4 is the Gaussian (complex) Integers
  julia> g = KummerModulus(4)
  ZAdjoin.KummerModulus{Int64}
  α^2 = -1

  # Function generator(modulus) will return the adjoined element.
  julia> α = generator(g)

  # Other elements can be formed from the adjoined element.
  julia> β = 3+4α # Another element
  ZAdjoin.Element{Int64}
  3+4α

  julia> α + β
  ZAdjoin.Element{Int64}
  3+5α

  julia> α*β
  ZAdjoin.Element{Int64}
  -4+3α

  # Compare with complex multiplication
  julia> im*(3+4im)
  -4 + 3im

  julia> α^2
  ZAdjoin.Element{Int64}
  -1

  # Matrix representation
  julia> Matrix(β)
  2×2 Array{Int64,2}:
   3  -4
   4   3

  # Representations are arithmetcally compatible (isomorphic)
  julia> Matrix(α+β^5)
  2×2 Array{Int64,2}:
    -237  3115
   -3115  -237

  julia> Matrix(α)+Matrix(β)^5
  2×2 Array{Int64,2}:
    -237  3115
   -3115  -237

  # The determinant of an element is the determinant of its Matrix
  # and acts as a multiplicative homomorphism into the integers.
  julia> det(β)
  25
  julia> det(1+α)
  2
  julia> det((1+α)*β)
  50

  ```

"""
module ZAdjoin

import Base: +, -, *, one, det, ==, string, show, Matrix
import Primes
export AbstractModulus, Modulus, KummerModulus, Element, generator

include("types.jl")
include("operations.jl")
include("utilities.jl")

end


# solution is x1=6, x2=2/3; objective = -6.666667
function nlp1(gm::Bool = false)
    m = JuMP.Model()
    @variable(m, 0 <= x[1:2])
    for i = 1:2
        JuMP.set_upper_bound(x[i], [6, 4][i])
    end
    @objective(m, Min, -x[1] - x[2])
    if !gm
        @NLconstraint(m, e1, x[1]*x[2] <= 4)
        set_optimizer(m, BARON_SILENT)
        return m
    else
        gm = GlobalModel(model = m, name = "nlp1")
        set_optimizer(gm, CPLEX_SILENT)
        add_nonlinear_constraint(gm, :(x ->  4 - x[1]*x[2]), name = "e1")
        return gm
    end
end


include("./make_column.jl")

function gen_impulse_response(Freq,F)
   # sanitize input and make a frequency interval variable
   if length(size(Freq)) > 1 || length(size(Freq)) > 2
      error("Freq and F both need to be vectors.")
   end
   deltaF = Freq[2]-Freq[1];

   # Make the time array.
   timePts = linspace(0,1/deltaF,length(Freq)); # generate time array

   # Generate the impulse response by inverting S11*e^{1i*2*pi*f*t)
   #
   #
   # Approach 1: No parallelization
   ###################################################################
   #exp_array = exp(1i*2*pi*Freq*times.'); # times indexes columns, doubleFreq indexes rows
   #impulse_response = deltaF * exp_array.' * S11 * 1;
   ##################################################################
   #
   #
   # Approach 2: Parallelization
   ##################################################################

   # make array of indices roughly corresponding to 5#, 10#,..., 95# of
   # timePts

   f = zeros(length(timePts),1);

   # parfor t = 1:length(timePts)
   #    f(t,1) = deltaF *  exp(1i*2*pi*Freq*timePts(t)).' * F;
   # end
   f = deltaF * length(timePts) * ifft(F);

   return (timePts, f)
end


using StanModels, CSV

df = filter(row -> row[:age] >= 18,
  CSV.read(joinpath(@__DIR__, "..", "..", "data", "Howell1.csv"), delim=';'))

m4_1s = "
// Inferring a Rate
data {
  int N;
  real<lower=0> h[N];
}
parameters {
  real<lower=0> sigma;
  real<lower=0,upper=250> mu;
}
model {
  // Priors for mu and sigma
  mu ~ normal(178, 20);
  sigma ~ uniform( 0 , 50 );

  // Observed heights
  h ~ normal(mu, sigma);
}
";

sm = SampleModel("m4.1s", m4_1s);

m4_1_data = Dict("N" => length(df[!, :height]), "h" => df[!, :height]);

(sample_file, log_file) = stan_sample(sm, data=m4_1_data);

if !(sample_file == nothing)
  chn = read_samples(sm)
  # Update parameter names
  chn = set_names(chn, Dict("mu" => "μ", "sigma" => "σ"))
  describe(chn)
end


"""ARMA process"""
function ARMA(X⃗::AbstractVector, ξ⃗::AbstractVector, ϕ⃗::AbstractVector, θ⃗::AbstractVector) # No constant, too boring
    # Just one update of an ARMA process
    if !all(length(X⃗)-1 .== (length(ξ⃗)-1, length(ϕ⃗), length(θ⃗)))
        p = max(length(X⃗)-1, length(ξ⃗)-1, length(ϕ⃗), length(θ⃗))
        X⃗ = vcat(X⃗, zeros(p+1-length(X⃗)))
        ξ⃗ = vcat(ξ⃗, randn(p+1-length(ξ⃗)))
        ϕ⃗ = vcat(ϕ⃗, zeros(p-length(ϕ⃗)))
        θ⃗ = vcat(θ⃗, zeros(p-length(θ⃗)))
    end
    X⃗[2:end] = X⃗[1:end-1]
    ξ⃗[2:end] = ξ⃗[1:end-1]
    ξ⃗[1] = randn()
    X⃗[1] = ξ⃗[1] + sum(ϕ⃗.*X⃗[2:end]) + sum(θ⃗.*ξ⃗[2:end])
    return (X⃗, ξ⃗)
end

function ARMA(X⃗::AbstractVector, ξ⃗::AbstractVector, ϕ₁::Float64=0.0, θ₁::Float64=0.0, args...)
    # Args are higher order parameters, alternating between ϕᵢ and θᵢ
    # E.g. ARMA(X⃗, ξ⃗, ϕ₁, θ₁, ϕ₂, θ₂)
    ϕ⃗ = vcat(ϕ₁, args[1:2:end])
    θ⃗ = vcat(θ₁, args[2:2:end])
    ARMA(X⃗, ξ⃗, ϕ⃗, θ⃗)
end

function ARMA(X0::AbstractVector, ξ0::AbstractVector, ϕ::Vector{Float64}, θ::Vector{Float64}, T::Int)
    X = zeros(T, length(X0))
    X[1, :] = X0
    for t ∈ 2:T
        X[t, :], ξ0 = ARMA(X[t-1, :], ξ0, ϕ, θ)
    end
    return X
end

function ARMA(X0::AbstractVector, ϕ::Vector{Float64}, θ::Vector{Float64}, T::Int)
    X = zeros(T, length(X0))
    X[1, :] = X0
    ξ = randn(length(X0))
    ARMA(X0, ξ, ϕ, θ, T)
end

"""AR process"""
function AR(X⃗::AbstractVector, ξ⃗::AbstractVector, ϕ⃗::AbstractVector)
    θ⃗ = zeros(length(ϕ⃗))
    ARMA(X⃗, ξ⃗, ϕ⃗, θ⃗)
end

function AR(X⃗::AbstractVector, ξ⃗::AbstractVector, ϕ₁::Float64=0.0, args...)
    ϕ⃗ = vcat(ϕ₁, args)
    AR(X⃗, ξ⃗, ϕ⃗)
end

function AR(X0::AbstractVector, ξ0::AbstractVector, ϕ::Vector{Float64}, T::Int)
    θ = zeros(length(ϕ))
    ARMA(X0, ξ0, ϕ, θ, T)
end

function AR(X0::AbstractVector, ϕ::Vector{Float64}, T::Int)
    θ = zeros(length(ϕ))
    ARMA(X0, ξ, ϕ, θ, T)
end

"""MA process"""
function MA(X⃗::AbstractVector, ξ⃗::AbstractVector, θ⃗::AbstractVector)
    ϕ⃗ = zeros(length(θ⃗))
    ARMA(X⃗, ξ⃗, ϕ⃗, θ⃗)
end

function MA(X⃗::AbstractVector, ξ⃗::AbstractVector, θ₁::Float64=0.0, args...)
    θ⃗ = vcat(θ₁, args)
    MA(X⃗, ξ⃗, θ⃗)
end

function MA(X0::AbstractVector, ξ0::AbstractVector, ϕ::Vector{Float64}, T::Int)
    ϕ = zeros(length(θ))
    ARMA(X0, ξ0, ϕ, θ, T)
end

function MA(X0::AbstractVector, ϕ::Vector{Float64}, T::Int)
    ϕ = zeros(length(θ))s
    ARMA(X0, ξ, ϕ, θ, T)
end

arma(args...) = ARMA(args...)
ar(args...) = AR(args...)
ma(args...) = MA(args...)


struct Prior{F <: AbstractFloat}
  K::Int
  p::Beta
  gamma::Beta
  eta::Dirichlet
  mu::Normal
  a_omega::F
  nu::LogNormal
  psi::Normal
  tau::Gamma
end

function compute_prior_mu(data::Data)
  yfinite = [data.yC_finite; data.yT_finite]
  return Normal(mean(yfinite), std(yfinite))
end

function Prior(K; p=Beta(100, 100), gamma=Beta(1,1), eta=nothing,
               mu=Normal(0, 3), a_omega=2.5, tau=Gamma(0.5, 1),
               nu=LogNormal(0.5, 0.1), # nu=LogNormal(1.6, 0.4),
               psi=Normal(-1, 1), data=nothing)
  eta == nothing && (eta = Dirichlet(K, 1/K))
  data == nothing || (mu = compute_prior_mu(data))
  return Prior(K, p, gamma, eta, mu, a_omega, nu, psi, tau)
end

#= TEST:
prior = Prior(3)
for fn in fieldnames(Prior)
  println("$(fn) => $(getfield(prior, fn))")
end
=#


#=
Sampling functions to generate random galaxy consistent with provided NFW parameters
=#

#=
Questions from while coding
1) Why more complicated PhiR (genphi()) function in previous functions
2) Why x^2*vt in fprob()
3) Why introduce dummy normalization? in f()
=#

using Optim, GSL

#=
Define gravitationla potential function,
results are in units of PhiS
=#
function GravPotential(x::Real)
    return 1 - (log(1 + x) / x)
end

#=
Calculate escape velocity for a radius scaled by the
characteristic radius, i.e. x = r / rs
=#
function EscapeVelocity(x, rlim, rs, PhiS)
    xlim = rlim / rs
    return sqrt(2 * PhiS * (GravPotential(xlim) - GravPotential(x)))
end

#=
Calculate the  Navarro-Frenk-White density based on x, the radius
scaled by the characteristic radius, vr the radial component of velocity,
vt the tangental component of velocity, and the NFW parameters.  The density
is not normalized, the normalization constant is set to 1000 for conveniance.
=#
function NFWDensity{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1})
    if x < 0
        return 0
    end
    #=
    a = parameters[1]
    d = parameters[2]
    e = parameters[3]
    Ec = parameters[4]
    rlim = parameters[5]
    b = parameters[6]
    q = parameters[7]
    Jb = parameters[8]
    vmax = parameters[9]
    rmax = parameters[10]
    =#

    a, d, e, Ec, rlim, b, q, Jb, vmax, rmax = parameters

    rs = rmax / 2.16
    PhiS = (vmax / 0.465)^2
    PhiR = GravPotential(x) * PhiS

    J = abs(x * rs * vt) #vt = v * sin(theta)
    E = (vt * vt + vr * vr) / 2 + PhiR

    Ec = Ec * PhiS
    xlim = rlim / rs
    Philim = PhiS * GravPotential(xlim)
    Jb = Jb * rs * sqrt(PhiS)

    if b <= 0
        gJ = 1 / (1 + (J / Jb)^(-b))
    else
        gJ = 1 + (J / Jb)^b
    end

    N = 1000
    if E < Philim && E >= 0
        hE = N * E^a * (E^q + Ec^q)^(d/q) * (Philim - E)^e
    else
        hE = 0
    end
    return hE * gJ
end

#=
Scale NFWDensity function
=#
function NFWProb{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1})
    return x * x * vt * NFWDensity(x, vr, vt, parameters)
end

#=
Randomly a generate with samplesize = # stars and NFW parameters parameters.
Basic rejection sampler with a factor of 1.1 is used.  Results are scaled radius (x),
radial velocity (vr), and tangential velocity (vt).
=#
function GenGalaxy{G<:Real}(parameters::Array{G,1}, samplesize::Integer)
    rlim = parameters[5]
    vmax = parameters[9]
    rmax = parameters[10]
    rs = rmax / 2.16
    PhiS = (vmax / 0.465)^2

    function NFWProbOpt(x)
        return -NFWProb(x[1], x[2], x[3], parameters)
    end

    minProb = optimize(NFWProbOpt, [0.1, 1.0, 1.0], method = :nelder_mead)
    fmax = max(0,  -1.1 * minProb.f_minimum)
    x0 = 10.0^-8
    vmax0 = EscapeVelocity(x0, rlim, rs, PhiS)
    #print("fmax: ", fmax, " vmax: ", vmax0, "\n")

    sampledValues = Array{Array{Float64,1}}(samplesize)
    tested = 0
    accepted = 0
    while accepted < samplesize
        x = (rlim  / rs) * rand()
        vr = vmax0 * rand()
        vt = vmax0 * rand()
        u = rand()
        sampleDensity = NFWProb(x, vr, vt, parameters)
        tested += 1
        if sampleDensity >= (u * fmax)
            accepted += 1
            sampledValues[accepted] = [x*rs , vr, vt]
        end
    end
    #print("acceptance rate:", accepted / tested, "\n")
    return sampledValues
end

#=
Transform the results of GenGalaxy to Euclidean coordinates.
=#
function converttoeuclidean{G<:Real}(coordinate::Array{G,1})
    radius = coordinate[1]
    theta = acos(1.0 - 2.0 * rand())
    phi = 2.0 * pi * rand()

    velocity = sqrt(coordinate[2]^2 + coordinate[3]^2)
    velSign = sign(rand() - 0.5)
    velPhi = 2.0 * pi * rand()

    x = radius * sin(theta) * cos(phi)
    y = radius * sin(theta) * sin(phi)
    z = radius * cos(theta)

    vz2 = velSign * coordinate[2]
    vx2 = coordinate[3] * cos(velPhi)
    vy2 = coordinate[3] * sin(velPhi)

    vx = cos(theta) * cos(phi) * vx2 - sin(phi) * vy2 + sin(theta) * cos(phi) * vz2
    vy = cos(theta) * sin(phi) * vx2 + cos(phi) * vy2 + sin(theta) * sin(phi) * vz2
    vz = -sin(theta) * vx2 + cos(theta) * vz2
    return x, y, z, vx, vy, vz
end

#=
Simplified posterior functions
=#
function NFWhEDensity{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1})
    a = parameters[1]
    d = parameters[2]
    e = parameters[3]
    Ec = parameters[4]
    rlim = parameters[5]
    b = parameters[6]
    q = parameters[7]
    Jb = parameters[8]
    vmax = parameters[9]
    rmax = parameters[10]

    rs = rmax / 2.16
    PhiS = (vmax / 0.465)^2
    PhiR = GravPotential(x) * PhiS

    J = abs(x * rs * vt) #vt = v * sin(theta)
    E = (vt * vt + vr * vr) / 2 + PhiR

    Ec = Ec * PhiS
    xlim = rlim / rs
    Philim = PhiS * GravPotential(xlim)
    Jb = Jb * rs * sqrt(PhiS)

    N = 1000
    if E < Philim && E >= 0
        hE = N * E^a * (E^q + Ec^q)^(d/q) * (Philim - E)^e
    else
        hE = 0
    end
    return hE
end

function NFWhEProb{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1})
    return x * x * vt * NFWhEDensity(x, vr, vt, parameters)
end

function GenGalaxySimplified{G<:Real}(parameters::Array{G,1}, samplesize::Integer)
    rlim = parameters[5]
    vmax = parameters[9]
    rmax = parameters[10]
    rs = rmax / 2.16
    PhiS = (vmax / 0.465)^2

    function NFWhEProbOpt(x)
        return -NFWhEProb(x[1], x[2], x[3], parameters)
    end

    minProb = optimize(NFWhEProbOpt, [0.1, 1.0, 1.0], method = :nelder_mead)
    fmax = max(0,  -1.1 * minProb.f_minimum)

    sampledValues = Array{Array{Float64,1}}(samplesize)
    tested = 0
    accepted = 0
    while accepted < samplesize
        x = (rlim  / rs) * rand()
        x0 = 10.0^-8
        vmax0 = EscapeVelocity(x0, rlim, rs, PhiS)
        vr = vmax0 * rand()
        vt = vmax0 * rand()
        u = rand()
        sampleDensity = NFWhEProb(x, vr, vt, parameters)
        tested += 1
        if sampleDensity >= (u * fmax)
            accepted += 1
            sampledValues[accepted] = [x*rs , vr, vt]
        end
    end
    #print("acceptance rate:", accepted / tested, "\n")
    return sampledValues
end

function convertgalaxy(galaxy::Array{Array{Float64, 1}, 1})
    n = length(galaxy)
    x = Array{Float64}(n)
    y = Array{Float64}(n)
    z = Array{Float64}(n)
    vx = Array{Float64}(n)
    vy = Array{Float64}(n)
    vz = Array{Float64}(n)
    for ii in 1:n
        x[ii], y[ii], z[ii], vx[ii], vy[ii], vz[ii] = converttoeuclidean(galaxy[ii])
    end
    return x, y, z, vx, vy, vz
end

function radiusvelocity(galaxy::Array{Array{Float64, 1}, 1})
    n = length(galaxy)
    r = Array{Float64}(n)
    v = Array{Float64}(n)
    for ii in 1:n
        r[ii] = galaxy[ii][1]
        v[ii] = sqrt(galaxy[ii][2]^2 + galaxy[ii][3]^2)
    end
    return r, v
end

function radiusvelocity(x::Array{Float64}, y::Array{Float64}, z::Array{Float64}, vx::Array{Float64}, vy::Array{Float64}, vz::Array{Float64})
    n = length(x)
    r = Array{Float64}(n)
    v = Array{Float64}(n)
    for ii in 1:n
        r[ii] = sqrt(x[ii]^2 + y[ii]^2 + z[ii]^2)
        v[ii] = sqrt(vx[ii]^2 + vy[ii]^2 + vz[ii]^2)
    end
    return r, v
end

function radius(x::Array{Float64}, y::Array{Float64})
    n = length(x)
    r = Array{Float64}(n)
    for ii in 1:n
        r[ii] = sqrt(x[ii]^2 + y[ii]^2)
    end
    return r
end

function rv2(r::Array{Float64, 1}, v::Array{Float64, 1})
    n = length(r)
    quantity = Array{Float64}(n)
    for ii in 1:n
        quantity[ii] = r[ii] * v[ii]^2
    end
    return quantity
end

#new functions

function genphi0_integral(x::Float64, alpha::Float64, beta::Float64, gamma::Float64, rhos::Float64, rs::Float64)
    PhiS = rhos * rs^2
    function integrand1(xi)
        return (xi ^ -gamma) * (1 + xi ^ alpha) ^ ((gamma - beta) / alpha) * xi * xi
    end
    function integrand2(xi)
        return (xi ^ -gamma) * (1 + xi ^ alpha) ^ ((gamma - beta) / alpha) * xi
    end
    phi1 = quadgk(integrand1, 0.0, x)[1]
    phi2 = quadgk(integrand2, x, Inf)[1]

    return PhiS * (1 - (phi1 / x + phi2))
end

function hypergeom_2F1_expanded(a::Float64, b::Float64, c::Float64, z::Float64)
    #if z >= -1.0 && z < 1.0
    #    return hypergeom([a, b], c, z)
    #end
    if z >= -1.0 && z < 0.5
        return hypergeom([a, b], c, z)
    end
    if z >= 0.5 && z < 1.0
        coef1 = gamma(c) * gamma(c - a - b) / (gamma(c - a) * gamma(c - b))
        hypergeom1 = hypergeom_2F1_expanded(a, b, a + b - c + 1, 1 - z)
        coef2 = (1 - z) ^ (c - a - b) * gamma(c) * gamma(a + b - c) / (gamma(a) * gamma(b))
        hypergeom2 = hypergeom_2F1_expanded(c - a, c - b, c - a - b + 1, 1 - z)
        return coef1 * hypergeom1 + coef2 * hypergeom2
    end
    if z >= 1.0
        coef1 = z ^ (-1.0 * a) * gamma(c) * gamma(c - a - b) / (gamma(c - a) * gamma(c - b))
        hypergeom1 = hypergeom_2F1_expanded(a, a - c + 1, a + b - c + 1, 1 - 1 / z)
        coef2 = z ^ (a - c) * (1 - z) ^ (c - a - b) * gamma(c) * gamma(a + b - c) / (gamma(a) * gamma(b))
        hypergeom2 = hypergeom_2F1_expanded(c - a, 1 - a, c - a - b + 1, 1 - 1 / z)
        return coef1 * hypergeom1 + coef2 * hypergeom2
    end
    if z < -1.0
        return (1 - z) ^ (-1.0 * a) * hypergeom_2F1_expanded(a, c - b, c, z / (z - 1))
    end
end

function genphi0(x::Float64, alpha::Float64, beta::Float64, gamma::Float64, rhos::Float64, rs::Float64)
    PhiS = rhos * rs^2
    x0 = 10.0 ^ -13
    p1a = 0
    p1b = hypergeom_2F1_expanded((3.0 - gamma) / alpha, (beta - gamma) / alpha, (3.0 + alpha - gamma) / alpha, -x ^ alpha)
    I1 = (x0 ^ (3.0 - gamma) * p1a - x ^ (3.0 - gamma) * p1b) / (x * (gamma - 3.0))
    p2 = hypergeom_2F1_expanded((beta - 2.0) / alpha, (beta - gamma) / alpha, (alpha + beta - 2.0) / alpha, -x ^ (-alpha))
    I2 = x ^ (2.0 - beta) * p2 / (beta - 2.0)
    return PhiS * (1 - (I1 + I2))
end

function genphi(x::Float64, alpha::Float64, beta::Float64, gamma::Float64, rhos::Float64, rs::Float64)
    x0 = 10.0 ^ -8
    return genphi0(x, alpha, beta, gamma, rhos, rs) -  genphi0(x0, alpha, beta, gamma, rhos, rs)
end

function genphi(x::Float64, alpha::Float64, beta::Float64, gamma::Float64, rhos::Float64, rs::Float64, Phi0::Float64)
    return genphi0(x, alpha, beta, gamma, rhos, rs) -  Phi0
end

#new vesc function line 385
function escape_velocity(x, rlim, rhos, rs, alpha, beta, gamma)
    xlim = rlim / rs
    PhiR = genphi(x, alpha, beta, gamma, rhos, rs)
    Philim = genphi(xlim, alpha, beta, gamma, rhos, rs)
    return sqrt(2 * (Philim - PhiR))
end

function escape_velocity(x, rlim, rhos, rs, alpha, beta, gamma, Philim)
    PhiR = genphi(x, alpha, beta, gamma, rhos, rs)
    return sqrt(2 * (Philim - PhiR))
end

function escape_velocity(x, rlim, rhos, rs, alpha, beta, gamma, Philim, Phi0)
    PhiR = genphi(x, alpha, beta, gamma, rhos, rs, Phi0)
    return sqrt(2 * (Philim - PhiR))
end

#new f function
function SFW_density{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1})
    if x < 0
        return 0
    end
    a, d, e, Ec, rlim, b, q, Jb, rhos, rs, alpha, beta, gamma = parameters

    PhiS = rhos * rs ^ 2
    PhiR = genphi(x, alpha, beta, gamma, rhos, rs)

    J = abs(x * rs * vt) #vt = v * sin(theta)
    E = (vt * vt + vr * vr) / 2.0 + PhiR

    Ec = Ec * PhiS
    xlim = rlim / rs
    Philim = genphi(xlim, alpha, beta, gamma, rhos, rs)
    Jb = Jb * rs * sqrt(PhiS)

    if b <= 0
        gJ = 1.0 / (1.0 + (J / Jb)^(-b))
    else
        gJ = 1.0 + (J / Jb)^b
    end

    N = 1000.0
    if E < Philim && E >= 0
        hE = N * E^a * (E^q + Ec^q)^(d/q) * (Philim - E)^e
    else
        hE = 0
    end
    return hE * gJ
end

function SFW_density{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1}, Philim::G)
    if x < 0
        return 0
    end
    a, d, e, Ec, rlim, b, q, Jb, rhos, rs, alpha, beta, gamma = parameters

    PhiS = rhos * rs ^ 2
    PhiR = genphi(x, alpha, beta, gamma, rhos, rs)

    J = abs(x * rs * vt) #vt = v * sin(theta)
    E = (vt * vt + vr * vr) / 2.0 + PhiR

    Ec = Ec * PhiS
    Jb = Jb * rs * sqrt(PhiS)

    if b <= 0
        gJ = 1.0 / (1.0 + (J / Jb)^(-b))
    else
        gJ = 1.0 + (J / Jb)^b
    end

    N = 1000.0
    if E < Philim && E >= 0
        hE = N * E^a * (E^q + Ec^q)^(d/q) * (Philim - E)^e
    else
        hE = 0
    end
    return hE * gJ
end

function SFW_density{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1}, Philim::G, Phi0::G)
    if x < 0
        return 0
    end
    a, d, e, Ec, rlim, b, q, Jb, rhos, rs, alpha, beta, gamma = parameters

    PhiS = rhos * rs ^ 2
    PhiR = genphi(x, alpha, beta, gamma, rhos, rs, Phi0)

    J = abs(x * rs * vt) #vt = v * sin(theta)
    E = (vt * vt + vr * vr) / 2.0 + PhiR

    Ec = Ec * PhiS
    Jb = Jb * rs * sqrt(PhiS)

    if b <= 0
        gJ = 1.0 / (1.0 + (J / Jb)^(-b))
    else
        gJ = 1.0 + (J / Jb)^b
    end

    N = 1000.0
    if E < Philim && E >= 0
        hE = N * E^a * (E^q + Ec^q)^(d/q) * (Philim - E)^e
    else
        hE = 0
    end
    return hE * gJ
end

#new fprob
function SFW_prob{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1})
    return x * x * vt * SFW_density(x, vr, vt, parameters)
end

function SFW_prob{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1}, Philim)
    return x * x * vt * SFW_density(x, vr, vt, parameters, Philim)
end

function SFW_prob{G<:Real}(x::G, vr::G, vt::G, parameters::Array{G,1}, Philim, Phi0)
    return x * x * vt * SFW_density(x, vr, vt, parameters, Philim, Phi0)
end

#new sample function
function general_gen_galaxy{G<:Real}(parameters::Array{G,1}, samplesize::Integer)
    a, d, e, Ec, rlim, b, q, Jb, rhos, rs, alpha, beta, gamma = parameters
    xlim = rlim / rs
    x0 = 10.0 ^ -8
    Phi0 = genphi0(x0, alpha, beta, gamma, rhos, rs)
    Philim = genphi(xlim, alpha, beta, gamma, rhos, rs, Phi0)
    function SFW_prob_opt(x)
        return -SFW_prob(x[1], x[2], x[3], parameters, Philim, Phi0)
    end

    minProb = optimize(SFW_prob_opt, [0.1, 1.0, 1.0], method = :nelder_mead)
    #minProb = optimize(SFW_prob_opt, [0.1, 1.0, 1.0], method = Optim.NelderMead())
    fmax = max(0,  -1.1 * minProb.f_minimum)
    x0 = 10.0^-8
    vmax0 = escape_velocity(x0, rlim, rhos, rs, alpha, beta, gamma, Philim, Phi0)
    #print("fmax: ", fmax, " vmax: ", vmax0, "\n")

    sampledValues = Array{Array{Float64,1}}(samplesize)
    tested = 0
    accepted = 0
    while accepted < samplesize
        x = (rlim  / rs) * rand()
        vr = vmax0 * rand()
        vt = vmax0 * rand()
        u = rand()
        sampleDensity = SFW_prob(x, vr, vt, parameters, Philim, Phi0)
        tested += 1
        if sampleDensity >= (u * fmax)
            accepted += 1
            sampledValues[accepted] = [x * rs , vr, vt]
        end
    end
    #print("acceptance rate:", accepted / tested, "\n")
    return sampledValues
end

#=
param=[2.0, -5.3, 2.5, 0.16, 1.5, -9.0, 6.9, 0.086, 4229.2, .694444, 1.01, 3.01, 1.01]
@time general_gen_galaxy(param, 1000); #7.5 seconds before Philim
@time general_gen_galaxy(param, 1000); #3.6 seconds before Phi0
@time general_gen_galaxy(param, 1000); #2.1 seconds with Philim and Phi0

#hypergeom([1.0, 2.0], 2.0, -0.98) #julia syntax
#ss.hyp2f1(1.0, 2.0, 2.0, -1.0) #python syntax
hypergeom()

@time genphi0(x, alpha, beta, gamma, rhos, rs)
@time genphi0_2(1.1, alpha, beta, gamma, rhos, rs)

@everywhere function f(n::Int64)
    x, vt, vr = general_gen_galaxy(param, 100)
    return mean(x)
end
=#


struct HuberAtom <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr}
    size::Tuple{Int,Int}
    M::Real

    function HuberAtom(x::AbstractExpr, M::Real)
        if sign(x) == ComplexSign()
            error("Arguemt must be real")
        elseif M <= 0
            error("Huber parameter must by a positive scalar")
        end
        children = (x,)
        return new(:huber, hash((children, M)), children, x.size, M)
    end
end

function sign(x::HuberAtom)
    return Positive()
end

function monotonicity(x::HuberAtom)
    return (Nondecreasing() * sign(x.children[1]),)
end

function curvature(x::HuberAtom)
    return ConvexVexity()
end

function evaluate(x::HuberAtom)
    c = evaluate(x.children[1])
    for i in 1:length(c)
        if c[i] <= x.M
            c[i] = c[i]^2
        else
            c[i] = 2 * x.M * c[i] - x.M^2
        end
    end
    return c
end

function conic_form!(x::HuberAtom, unique_conic_forms::UniqueConicForms)
    if !has_conic_form(unique_conic_forms, x)
        c = x.children[1]
        s = Variable(c.size)
        n = Variable(c.size)

        # objective given by s.^2 + 2 * M * |n|
        objective =
            conic_form!(square(s) + 2 * x.M * abs(n), unique_conic_forms)
        conic_form!(c == s + n, unique_conic_forms)

        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end

huber(x::AbstractExpr, M::Real = 1.0) = HuberAtom(x, M)


using LinearAlgebra
using MLDatasets.MNIST
using ImageCore
using Plots
using Random
using PyCall
@pyimport pickle

function h(v)
    map(x -> 1.0 / (1.0 + exp(-x)), v)
end

function σ(a)
    c = maximum(a)
    exp_a = map(x -> exp(x .- c), a)
    exp_a / sum(exp_a)
end

function init_network()::Matrix{Float64}
    rand(Float64, (2,3))
end

function loss(
    x::Matrix{Float64},
    t::Matrix{Float64},
    W::Matrix{Float64},
    )::Float64
    z = x * W
    y = σ(z)
    cross_entropy_error(y, t)
end

function forward(network, x)
    W1, W2, W3 = network["W1"], network["W2"], network["W3"]
    b1, b2, b3 = network["b1"], network["b2"], network["b3"]

    a1 = x * W1 + b1
    z1 = h(a1)
    a2 = z1 * W2 + b2
    z2 = h(a2)
    a3 = z2 * W3 + b3
end

function prepare_testdata(;batch=1)
    data, labels = MNIST.testdata()
    batch_num = div(length(labels), batch)
    data0 = eachslice(reshape(data, (28 * 28 * batch, batch_num)), dims=2)
    labels0 = eachslice(transpose(reshape(labels, (batch, batch_num))), dims=1)
    collect(zip(data0, labels0))
end

function cross_entropy_error(
    y::AbstractArray{T, N},
    t::AbstractArray{T, N},
    )::AbstractFloat where{T <: Number, N}
    - sum(float.(t) .* log.(float.(y)))
end

function numerical_gradient(
    f::Function,
    xs::AbstractArray{T, N};
    h=1e-4::T,
    )::AbstractArray{T, N} where{T <: Number, N}
    grad = zeros(size(xs))
    for (i, x) in enumerate(xs)
        xs[i] = x + h
        fxh1 = f(xs)
        xs[i] = x - h
        fxh2 = f(xs)
        grad[i] = (fxh1 - fxh2) / (2h)
        xs[i] = x
    end
    grad
end

function gradient_descent(
    f::Function,
    init_x::AbstractArray{T, N};
    learning_rate=0.01::T,
    steps=100::UInt128
    )::AbstractArray{T, N} where{T <: Number, N}
    x = init_x
    for _ = 1:steps
        x -= learning_rate * numerical_gradient(f, x)
    end
    x
end

function make_predict(fname="./sample_weight.pkl")
    f = pybuiltin("open")(fname, "rb")
    network = pickle.load(f)
    f.close()
    W₁ ,W₂, W₃ = network["W1"], network["W2"], network["W3"]
    b₁ ,b₂, b₃ = network["b1"], network["b2"], network["b3"]
    function (x::AbstractArray)
        a₁ = x * W₁ .+ transpose(b₁)
        z₁ = h(a₁)
        a₂ = z₁ * W₂ .+ transpose(b₂)
        z₂ = h(a₂)
        a₃ = z₂ * W₃ .+ transpose(b₃)
        σ(a₃)
    end
end

function bench01()
    predict = make_predict()
    data = prepare_testdata()
    count = 0
    for (x, l) in data
        if argmax(predict(x))[2] - 1 == l[1]
            count += 1
        end
    end
    count / length(data)
end
bench01()

function bench01_batch(;batch=100)
    predict = make_predict()
    data = prepare_testdata(batch=batch)
    count = 0
    for (xs, ls) in data
        for (p, l) in zip(eachslice(predict(transpose(reshape(xs, (28 * 28, batch)))), dims=1), ls)
            if argmax(p) - 1 == l
                count += 1
            end
        end
    end
    count / (length(data) * batch)
end


using PyPlot
using FFTW
using Bolt
using LinearAlgebra

N = 64
μ = 0
q = 0.0
r = 10 .^ range(-4., 4., length=N)

pl = Bolt.plan_fftlog(r, μ, q, 1.0; kropt=true)
aₙ = r .^ (μ + 1) .* exp.(-r.^2 / 2)
# aₙ = exp.(-(log.(r) ).^2 / 2)
y = similar(r, ComplexF64)
y2 = similar(r, ComplexF64)
mul!(y, pl, aₙ)

clf()
plt.plot(r, aₙ)
# plt.plot(r, aₙ .* r.^(-q))
# ylim(1e-5, 1.2)
plt.plot(pl.k, y)
yscale("symlog", linthreshy=1e-5)
# yscale("log")
xscale("log")
gcf()

##

# y .*= sin.(0.5 .* r)
ldiv!(y2,pl,y)

clf()
plt.plot(r, aₙ)
plt.plot(pl.r, y2)
# yscale("symlog", linthreshy=1e-5)
xscale("log")
gcf()

##
aₙ = sf[1200:end,1]
y = similar(r, ComplexF64)
r = x2a.(bg.x_grid[1200:end])
μ = 30
q = 0.9
pl = Bolt.plan_fftlog(r, μ, q, 1.0; kropt=true)
mul!(y, pl, aₙ)

clf()
plt.plot(r, aₙ)
plt.plot(r, y)

yscale("symlog", linthreshy=maximum(norm.(y)) * 1e-2)
xscale("log")
gcf()

##
using DelimitedFiles
fftdata = readdlm("test/fftlog64.txt", ' ', Float64, '\n')
k_ref = fftdata[:,1]
f_ref = fftdata[:,2]

##
clf()
plot(k_ref .- fftlogfreq(pl))
gcf()
##

isapprox(y, f_ref)

##
clf()
plot(y .- f_ref)
gcf()

##
clf()
plt.plot(fftlogfreq(pl), reverse(y))
plt.plot(k_ref, f_ref)
yscale("symlog", base=10, linthresh=1e-5)
xscale("log")
gcf()


"""
    infovalue(g::Vector{Integer}, b::Vector{Integer}))

Information value calculation of `g`, `b` vector of binned frequency counts

- weight of evidence = log(density g / density b), 0 adjusted
- infovalue = sum (density g - density b) * weight of evidence

Industry rule of thumb:
- iv <= 0.1         no significant change
- 0.1 < iv <= 0.25  minor change
- 0.25 < iv         major change
"""
function infovalue(
    g::Vector{TG} where {TG<:Integer},
    b::Vector{TB} where {TB<:Integer};
    woeadjust = 0.5,
)
    length(g) >= 2 || error("frequency vectors should have at least 2 elements")
    length(g) == length(b) || error("frequency vectors should be the same length")

    # infovalue is class invariant - swapping g/b doesn't matter
    # infovalue is order invariant - re-ordering rows doesn't matter
    # implementation modelled after SAS PROC BINNING where
    # weight of evidence is adjusted to avoid divide by 0

    pb = b ./ sum(b)
    pg = g ./ sum(g)
    pba = (b + (b .== 0) * woeadjust) ./ sum(b)     # adjust only pure bins
    pga = (g + (g .== 0) * woeadjust) ./ sum(g)
    woe = log.(pga ./ pba)
    sum((pg .- pb) .* woe)
end


using Ferrite, SparseArrays, LinearAlgebra, FerriteViz

struct J2Plasticity{T, S <: SymmetricTensor{4, 3, T}}
    G::T  # Shear modulus
    K::T  # Bulk modulus
    σ₀::T # Initial yield limit
    H::T  # Hardening modulus
    Dᵉ::S # Elastic stiffness tensor
end;

function J2Plasticity(E, ν, σ₀, H)
    δ(i,j) = i == j ? 1.0 : 0.0 # helper function
    G = E / 2(1 + ν)
    K = E / 3(1 - 2ν)

    Isymdev(i,j,k,l) = 0.5*(δ(i,k)*δ(j,l) + δ(i,l)*δ(j,k)) - 1.0/3.0*δ(i,j)*δ(k,l)
    temp(i,j,k,l) = 2.0G *( 0.5*(δ(i,k)*δ(j,l) + δ(i,l)*δ(j,k)) + ν/(1.0-2.0ν)*δ(i,j)*δ(k,l))
    Dᵉ = SymmetricTensor{4, 3}(temp)
    return J2Plasticity(G, K, σ₀, H, Dᵉ)
end;

mutable struct MaterialState{T, S <: SecondOrderTensor{3, T}}
    # Store "converged" values
    ϵᵖ::S # plastic strain
    σ::S # stress
    k::T # hardening variable

    # Store temporary values used during equilibrium iterations
    temp_ϵᵖ::S
    temp_σ::S
    temp_k::T
end

function MaterialState()
    return MaterialState(
                zero(SymmetricTensor{2, 3}),
                zero(SymmetricTensor{2, 3}),
                0.0,
                zero(SymmetricTensor{2, 3}),
                zero(SymmetricTensor{2, 3}),
                0.0)
end

function update_state!(state::MaterialState)
    state.ϵᵖ = state.temp_ϵᵖ
    state.σ = state.temp_σ
    state.k = state.temp_k
end;

function vonMises(σ)
    s = dev(σ)
    return sqrt(3.0/2.0 * s ⊡ s)
end;

function compute_stress_tangent(ϵ::SymmetricTensor{2, 3}, material::J2Plasticity, state::MaterialState)
    # unpack some material parameters
    G = material.G
    K = material.K
    H = material.H

    # We use (•)ᵗ to denote *trial*-values
    σᵗ = material.Dᵉ ⊡ (ϵ - state.ϵᵖ) # trial-stress
    sᵗ = dev(σᵗ)         # deviatoric part of trial-stress
    J₂ = 0.5 * sᵗ ⊡ sᵗ   # second invariant of sᵗ
    σᵗₑ = sqrt(3.0*J₂)   # effetive trial-stress (von Mises stress)
    σʸ = material.σ₀ + H * state.k # Previous yield limit

    φᵗ  = σᵗₑ - σʸ # Trial-value of the yield surface

    if φᵗ < 0.0 # elastic loading
        state.temp_σ = σᵗ
        return state.temp_σ, material.Dᵉ
    else # plastic loading
        h = H + 3G
        μ =  φᵗ / h   # plastic multiplier

        c1 = 1 - 3G * μ / σᵗₑ
        s = c1 * sᵗ           # updated deviatoric stress
        σ = s + vol(σᵗ)       # updated stress

        # Compute algorithmic tangent stiffness ``D = \frac{\Delta \sigma }{\Delta \epsilon}``
        κ = H * (state.k + μ) # drag stress
        σₑ = material.σ₀ + κ  # updated yield surface

        δ(i,j) = i == j ? 1.0 : 0.0
        Isymdev(i,j,k,l)  = 0.5*(δ(i,k)*δ(j,l) + δ(i,l)*δ(j,k)) - 1.0/3.0*δ(i,j)*δ(k,l)
        Q(i,j,k,l) = Isymdev(i,j,k,l) - 3.0 / (2.0*σₑ^2) * s[i,j]*s[k,l]
        b = (3G*μ/σₑ) / (1.0 + 3G*μ/σₑ)

        Dtemp(i,j,k,l) = -2G*b * Q(i,j,k,l) - 9G^2 / (h*σₑ^2) * s[i,j]*s[k,l]
        D = material.Dᵉ + SymmetricTensor{4, 3}(Dtemp)

        # Store outputs in the material state
        Δϵᵖ = 3/2 *μ / σₑ*s            # plastic strain
        state.temp_ϵᵖ = state.ϵᵖ + Δϵᵖ  # plastic strain
        state.temp_k = state.k + μ     # hardening variable
        state.temp_σ = σ               # updated stress
        return state.temp_σ, D
    end
end

function create_values(interpolation)
    # setup quadrature rules
    qr      = QuadratureRule{3,RefTetrahedron}(2)
    face_qr = QuadratureRule{2,RefTetrahedron}(3)

    # create geometric interpolation (use the same as for u)
    interpolation_geom = Lagrange{3,RefTetrahedron,1}()

    # cell and facevalues for u
    cellvalues_u = CellVectorValues(qr, interpolation, interpolation_geom)
    facevalues_u = FaceVectorValues(face_qr, interpolation, interpolation_geom)

    return cellvalues_u, facevalues_u
end;

function create_dofhandler(grid, interpolation)
    dh = DofHandler(grid)
    dim = 3
    push!(dh, :u, dim, interpolation) # add a displacement field with 3 components
    close!(dh)
    return dh
end

function create_bc(dh, grid)
    dbcs = ConstraintHandler(dh)
    # Clamped on the left side
    dofs = [1, 2, 3]
    dbc = Dirichlet(:u, getfaceset(grid, "left"), (x,t) -> [0.0, 0.0, 0.0], dofs)
    add!(dbcs, dbc)
    close!(dbcs)
    return dbcs
end;

function doassemble(cellvalues::CellVectorValues{dim},
                    facevalues::FaceVectorValues{dim}, K::SparseMatrixCSC, grid::Grid,
                    dh::DofHandler, material::J2Plasticity, u, states, t) where {dim}
    r = zeros(ndofs(dh))
    assembler = start_assemble(K, r)
    nu = getnbasefunctions(cellvalues)
    re = zeros(nu)     # element residual vector
    ke = zeros(nu, nu) # element tangent matrix

    for (cell, state) in zip(CellIterator(dh), states)
        fill!(ke, 0)
        fill!(re, 0)
        eldofs = celldofs(cell)
        ue = u[eldofs]
        assemble_cell!(ke, re, cell, cellvalues, facevalues, grid, material,
                       ue, state, t)
        assemble!(assembler, eldofs, re, ke)
    end
    return K, r
end

function assemble_cell!(Ke, re, cell, cellvalues, facevalues, grid, material,
                        ue, state, t)
    n_basefuncs = getnbasefunctions(cellvalues)
    reinit!(cellvalues, cell)

    for q_point in 1:getnquadpoints(cellvalues)
        # For each integration point, compute stress and material stiffness
        ∇u = function_gradient(cellvalues, q_point, ue)
        ϵ = symmetric(∇u) # Total strain
        σ, D = compute_stress_tangent(ϵ, material, state[q_point])

        dΩ = getdetJdV(cellvalues, q_point)
        for i in 1:n_basefuncs
            δϵ = symmetric(shape_gradient(cellvalues, q_point, i))

            re[i] += (δϵ ⊡ σ) * dΩ # add internal force to residual
            for j in 1:i
                Δϵ = symmetric(shape_gradient(cellvalues, q_point, j))
                Ke[i, j] += δϵ ⊡ D ⊡ Δϵ * dΩ
            end
        end
    end
    symmetrize_lower!(Ke)

    # Add traction as a negative contribution to the element residual `re`:
    for face in 1:nfaces(cell)
        if onboundary(cell, face) && (cellid(cell), face) ∈ getfaceset(grid, "right")
            reinit!(facevalues, cell, face)
            for q_point in 1:getnquadpoints(facevalues)
                dΓ = getdetJdV(facevalues, q_point)
                for i in 1:n_basefuncs
                    δu = shape_value(facevalues, q_point, i)
                    re[i] -= (δu ⋅ t) * dΓ
                end
            end
        end
    end

end

function symmetrize_lower!(K)
    for i in 1:size(K,1)
        for j in i+1:size(K,1)
            K[i,j] = K[j,i]
        end
    end
end;

function solve(liveplotting=false)
    # Define material parameters
    E = 200.0e9 # [Pa]
    H = E/20   # [Pa]
    ν = 0.3     # [-]
    σ₀ = 200e6  # [Pa]
    material = J2Plasticity(E, ν, σ₀, H)

    L = 10.0 # beam length [m]
    w = 1.0  # beam width [m]
    h = 1.0  # beam height[m]
    n_timesteps = 100
    u_max = zeros(n_timesteps)
    traction_magnitude = 1.e7 * range(0.5, 1.0, length=n_timesteps)

    # Create geometry, dofs and boundary conditions
    n = 2
    nels = (10n, n, 2n) # number of elements in each spatial direction
    P1 = Vec((0.0, 0.0, 0.0))  # start point for geometry
    P2 = Vec((L, w, h))        # end point for geometry
    grid = generate_grid(Tetrahedron, nels, P1, P2)
    interpolation = Lagrange{3, RefTetrahedron, 1}() # Linear tet with 3 unknowns/node

    dh = create_dofhandler(grid, interpolation) # JuaFEM helper function
    dbcs = create_bc(dh, grid) # create Dirichlet boundary-conditions

    cellvalues, facevalues = create_values(interpolation)

    # Pre-allocate solution vectors, etc.
    n_dofs = ndofs(dh)  # total number of dofs
    u  = zeros(n_dofs)  # solution vector
    u_history = Vector{Vector{Float64}}()
    if liveplotting
        plotter = MakiePlotter(dh,u)
        fig = ferriteviewer(plotter)
        display(fig)
    end
    Δu = zeros(n_dofs)  # displacement correction
    r = zeros(n_dofs)   # residual
    K = create_sparsity_pattern(dh); # tangent stiffness matrix

    # Create material states. One array for each cell, where each element is an array of material-
    # states - one for each integration point
    nqp = getnquadpoints(cellvalues)
    states = [[MaterialState() for _ in 1:nqp] for _ in 1:getncells(grid)]

    # states = [MaterialState() for _ in 1:nqp for _ in 1:getncells(grid)]
    # temp_states = [MaterialState() for _ in 1:nqp for _ in 1:getncells(grid)]

    # Newton-Raphson loop
    NEWTON_TOL = 1 # 1 N

    for timestep in 1:n_timesteps
        t = timestep # actual time (used for evaluating d-bndc)
        traction = Vec((0.0, 0.0, traction_magnitude[timestep]))
        newton_itr = -1
        update!(dbcs, t) # evaluates the D-bndc at time t
        apply!(u, dbcs)  # set the prescribed values in the solution vector

        while true; newton_itr += 1

            if newton_itr > 8
                error("Reached maximum Newton iterations, aborting")
                break
            end
            K, r = doassemble(cellvalues, facevalues, K, grid, dh, material, u,
                             states, traction);
            norm_r = norm(r[Ferrite.free_dofs(dbcs)])

            if norm_r < NEWTON_TOL
                break
            end

            apply_zero!(K, r, dbcs)
            Δu = Symmetric(K) \ r
            u -= Δu
        end
        if liveplotting
            FerriteViz.update!(plotter,u)
            sleep(0.1)
        end
        push!(u_history,u)

        # Update all the material states after we have reached equilibrium
        for cell_states in states
            foreach(update_state!, cell_states)
        end
        u_max[timestep] = max(abs.(u)...) # maximum displacement in current timestep
    end

    # ## Postprocessing
    # Only a vtu-file corrsponding to the last time-step is exported.
    #
    # The following is a quick (and dirty) way of extracting average cell data for export.
    mises_values = zeros(getncells(grid))
    κ_values = zeros(getncells(grid))
    for (el, cell_states) in enumerate(states)
        for state in cell_states
            mises_values[el] += vonMises(state.σ)
            κ_values[el] += state.k*material.H
        end
        mises_values[el] /= length(cell_states) # average von Mises stress
        κ_values[el] /= length(cell_states)     # average drag stress
    end
    return u, dh, u_history, mises_values, κ_values
end


function estimate(entity::FundamentalMatrix, method::DirectLinearTransform, 𝒟::Tuple{AbstractArray, Vararg{AbstractArray}})
    ℳ, ℳʹ =  𝒟
    N = length(ℳ)
    if (N != length(ℳʹ))
          throw(ArgumentError("There should be an equal number of points for each view."))
    end
    𝒪, 𝐓 = transform(HomogeneousCoordinates(),CanonicalToHartley(),ℳ)
    𝒪ʹ, 𝐓ʹ  = transform(HomogeneousCoordinates(),CanonicalToHartley(),ℳʹ)
    𝐀 = moments(FundamentalMatrix(), (𝒪, 𝒪ʹ))
    λ, f = smallest_eigenpair(Symmetric(𝐀))
    𝐅 = reshape(f,(3,3))
    𝐅 = enforce_ranktwo!(Array(𝐅))
    𝐅 = SMatrix{3,3,Float64,9}(𝐅 / norm(𝐅))
    # Transform estimate back to the original (unnormalised) coordinate system.
    𝐓ʹ'*𝐅*𝐓
end

# TODO fix numerical instability
# function estimate(entity::FundamentalMatrix, method::Taubin, matches...)
#     ℳ, ℳʹ = matches
#     N = length(ℳ)
#     if (N != length(ℳʹ))
#           throw(ArgumentError("There should be an equal number of points for each view."))
#     end
#     (ℳ,𝐓) = hartley_normalization(ℳ)
#     (ℳʹ,𝐓ʹ) = hartley_normalization(ℳʹ)
#     𝐀::Matrix{Float64} = moments(FundamentalMatrix(), ℳ, ℳʹ)
#     𝐁::Matrix{Float64} = mean_covariance(FundamentalMatrix(), ℳ, ℳʹ)
#     (λ::Float64, f::Vector{Float64}) = smallest_eigenpair(𝐀,𝐁)
#     𝐅::Matrix{Float64} = reshape(f,(3,3))
#     enforce_ranktwo!(𝐅)
#     # Transform estimate back to the original (unnormalised) coordinate system.
#     𝐅 = 𝐓ʹ'*𝐅*𝐓
# end

function estimate(entity::FundamentalMatrix, method::FundamentalNumericalScheme,  𝒞::Tuple{AbstractArray, Vararg{AbstractArray}}, 𝒟::Tuple{AbstractArray, Vararg{AbstractArray}})
    ℳ, ℳʹ = 𝒟
    Λ₀, Λ₀ʹ = 𝒞
    N = length(ℳ)
    if (N != length(ℳʹ))
          throw(ArgumentError("There should be an equal number of points for each view."))
    end
    if (N != length(Λ₀) || N != length(Λ₀ʹ) )
          throw(ArgumentError("There should be a covariance matrix for each point correspondence."))
    end
    # Map corresponding points to the normalized coordinate system.
    𝒪, 𝐓 = transform(HomogeneousCoordinates(),CanonicalToHartley(),ℳ)
    𝒪ʹ, 𝐓ʹ = transform(HomogeneousCoordinates(),CanonicalToHartley(),ℳʹ)
    # Map seed to the normalized coordinate system.
    𝛉 = (inv(𝐓') ⊗ inv(𝐓ʹ')) * method.𝛉₀
    # Map covariance matrices to the normalized coordinate system.
    Λ₁ = transform(CovarianceMatrices(), CanonicalToHartley(), Λ₀ , 𝐓)
    Λ₁ʹ = transform(CovarianceMatrices(), CanonicalToHartley(), Λ₀ʹ , 𝐓ʹ)
    for i = 1:method.max_iter
        𝐗 = X(AML(),FundamentalMatrix(), 𝛉, (Λ₁,Λ₁ʹ), (𝒪, 𝒪ʹ))
        λ, 𝛉⁺ = smallest_eigenpair(Symmetric(𝐗/N))
        𝛉 = reshape(𝛉⁺,length(𝛉⁺),1)
    end
    𝐅 = reshape(𝛉,(3,3))
    𝐅 = enforce_ranktwo!(Array(𝐅))
    𝐅 = SMatrix{3,3,Float64,9}(𝐅)
    # Transform estimate back to the original (unnormalised) coordinate system.
    𝐅 = 𝐓ʹ'*𝐅*𝐓
end

function estimate(entity::FundamentalMatrix, method::BundleAdjustment,  𝒟::Tuple{AbstractArray, Vararg{AbstractArray}})
    ℳ, ℳʹ = 𝒟
    N = length(ℳ)
    if (N != length(ℳʹ))
          throw(ArgumentError("There should be an equal number of points for each view."))
    end
    𝐅 = SMatrix{3,3,Float64,9}(reshape(method.𝛉₀,(3,3)))
    𝒳 = triangulate(DirectLinearTransform(),𝐅,(ℳ,ℳʹ))

    𝐏₁, 𝐏₂ = construct(ProjectionMatrix(),𝐅)

    # Construct a length-(12+3*N) vector consisting of the projection matrix associated
    # with the second view (the first twelve dimensions), as well as N three-dimensional points
    # (the remaining dimensions).
    𝛉 = pack(FundamentalMatrix(), 𝐏₂, 𝒳)

    index₁ = SVector(1,2)
    index₂ = SVector(3,4)
    pts = Matrix{Float64}(undef,4,N)
    for n = 1:N
        pts[index₁,n] = ℳ[n][index₁]
        pts[index₂,n] = ℳʹ[n][index₁]
    end
    fit = curve_fit(model_fundamental, jacobian_model,  𝐏₁, reshape(reinterpret(Float64,vec(pts)),(4*N,)) , 𝛉; show_trace = false)
    𝐏₂ = reshape(fit.param[1:12],(3,4))
    𝐅 = construct(FundamentalMatrix(), 𝐏₁, 𝐏₂)
    𝐅, fit
end

function model_fundamental(𝐏₁,𝛉)
    # Twelve parameters for the projection matrix, and 3 parameters per 3D point.
    N = Int((length(𝛉) - 12) / 3)
    index₁ = SVector(1,2)
    index₂ = SVector(3,4)
    reprojections = Matrix{Float64}(undef,4,N)
    𝛉v = @view 𝛉[1:12]
    𝐏₂ = SMatrix{3,4,Float64,12}(reshape(𝛉v,(3,4)))
    i = 13
    for n = 1:N
        # Extract 3D point and convert to homogeneous coordinates
        v = @view 𝛉[i:i+2]
        M = hom(SVector{3,Float64}(𝛉[i:i+2]))
        reprojections[index₁,n] = hom⁻¹(𝐏₁ * M)
        reprojections[index₂,n] = hom⁻¹(𝐏₂ * M)
        i = i + 3
    end
    reshape(reinterpret(Float64,vec(reprojections)),(4*N,))
end

function jacobian_model(𝐏₁,𝛉)
    # Twelve parameters for the projection matrix, and 3 parameters per 3D point.
    N = Int((length(𝛉) - 12) / 3)
    index₁ = SVector(1,2)
    index₂ = SVector(3,4)
    reprojections = Matrix{Float64}(undef,4,N)
    𝛉v = @view 𝛉[1:12]
    𝐏₂ = SMatrix{3,4,Float64,12}(reshape(𝛉v,(3,4)))
    𝐉 = zeros(4*N,12+3*N)
    # Create a view of the jacobian matrix 𝐉 and reshape it so that
    # it will be more convenient to index into the appropriate entires
    # whilst looping over all of the data points.
    𝐉v = reshape(reinterpret(Float64,𝐉), 4, N, 12+3*N)
    𝐀 = SMatrix{2,3,Float64,6}(1,0,0,1,0,0)
    𝐈₃ = SMatrix{3,3}(1.0I)
    i = 13
    for n = 1:N
        # Extract 3D point and convert to homogeneous coordinates.
        𝐌 = hom(SVector{3,Float64}(𝛉[i:i+2]))

        # Derivative of residual in first and second image w.r.t 3D point.
        ∂𝐫₁_d𝐌 = -𝐀 * ∂hom⁻¹(𝐏₁ * 𝐌) * 𝐏₁
        ∂𝐫₂_d𝐌 = -𝐀 * ∂hom⁻¹(𝐏₂ * 𝐌) * 𝐏₂

        # Derivative of residual in second image w.r.t projection martix
        # ∂𝐫₁_d𝐏₁ is the zero vector.
        ∂𝐫₂_d𝐏₂ = 𝐀 * ∂hom⁻¹(𝐏₂ * 𝐌) * (𝐌' ⊗ 𝐈₃)

        𝐉v[index₂,n,1:12] = ∂𝐫₂_d𝐏₂
        𝐉v[index₁,n,i:i+2] = ∂𝐫₁_d𝐌[:,1:3]
        𝐉v[index₂,n,i:i+2] = ∂𝐫₂_d𝐌[:,1:3]
        i = i + 3
    end
    𝐉
end

#𝛉 = reshape(𝛉₀,length(𝛉₀),1)
# function z(entity::FundamentalMatrix, 𝐦::Vector{Float64}, 𝐦ʹ::Vector{Float64})
# 𝐮 = 𝐦 ⊗ 𝐦ʹ
# 𝐮[1:end-1]
# end

function mean_covariance(entity::ProjectiveEntity, matches...)
    ℳ, ℳʹ = matches
    N = length(ℳ)
    if (N != length(ℳʹ))
          throw(ArgumentError("There should be an equal number of points for each view."))
    end
    𝐁 = fill(0.0,(9,9))
    𝚲 = eye(4)
    for correspondence in zip(ℳ, ℳʹ)
        m , mʹ = correspondence
        𝐦  = 𝑛(collect(Float64,m.coords))
        𝐦ʹ = 𝑛(collect(Float64,mʹ.coords))
        ∂ₓ𝐮 = ∂ₓu(entity, 𝐦 , 𝐦ʹ)
        𝐁 = 𝐁 + ∂ₓ𝐮 * 𝚲 * ∂ₓ𝐮'
    end
    𝐁/N
end

function enforce_ranktwo!(𝐅::AbstractArray)
    # Enforce the rank-2 constraint.
    U,S,V = svd(𝐅)
    S[end] = 0.0
    𝐅 = U*Matrix(Diagonal(S))*V'
end

# Construct a parameter vector consisting of a projection matrix and 3D points.
function pack(entity::FundamentalMatrix, 𝐏₂::AbstractArray, 𝒳::AbstractArray)
    N = length(𝒳)
    𝛉 = Vector{Float64}(undef,12+N*3)
    𝛉[1:12] = Array(𝐏₂[:])
    i = 13
    for n = 1:N
        𝛉[i:i+2] = 𝒳[n][1:3]
        i = i + 3
    end
    𝛉
end


using PastaQ
using ITensors
using Random

# Set the random seed so the results are the same
# each time it is run
Random.seed!(1234)

# Initialize the MPS state ψ = |0,0,0⟩
ψ = qubits(3)

# Apply the X gate on qubit 2
g = ("X", 2)
ψ = runcircuit(ψ, g)

# Show samples from P(x) = |⟨x|ψ⟩|²
println("Sample from |ψ⟩ = X₂|0,0,0⟩:")
display(getsamples(ψ, 3))
println()

# Custom gates
#
# There might be often problems that require quantum gates not included
# in the standard PastaQ gate set. If that is the case, a new gate can be
# added by overloading the gate(::GateName"...") function, using the format
# defined in gates.jl.
#

# Import gate so it can be overloaded
import PastaQ: gate

gate(::GateName"iX") = [
  0 im
  im 0
]

# Set the state back to |0,0,0⟩
resetqubits!(ψ)

# Show samples from P(x) = |⟨x|ψ⟩|²
println("Sample from |ψ⟩ = iX₁|0,0,0⟩:")
g = ("iX", 1)
ψ = runcircuit(ψ, g)
display(getsamples(ψ, 3))
println()


function getTableauSLRK(s, o, tsym, q, p, q̃, p̃, ω, d=Nothing)
    TableauSLRK(tsym, o, s, q, p, q̃, p̃, ω, d)
end

function TableauSLRKLobattoIII(s, q, p)
    o = 2s-2
    tsym = Symbol("SLRKLobattoIII")
    getTableauSLRK(s, o, tsym, q, p, q, p, get_lobatto_ω_matrix(s), get_lobatto_nullvector(s))
end

function TableauSLRKLobattoIIIAB(s)
    q = TableauLobattoIIIA(s)
    p = TableauLobattoIIIB(s)
    TableauSLRKLobattoIII(s, q, p)
end

function TableauSLRKLobattoIIIBA(s)
    q = TableauLobattoIIIB(s)
    p = TableauLobattoIIIA(s)
    TableauSLRKLobattoIII(s, q, p)
end

function TableauSLRKLobattoIIICC̄(s)
    q = TableauLobattoIIIC(s)
    p = TableauLobattoIIIC̄(s)
    TableauSLRKLobattoIII(s, q, p)
end

function TableauSLRKLobattoIIIC̄C(s)
    q = TableauLobattoIIIC̄(s)
    p = TableauLobattoIIIC(s)
    TableauSLRKLobattoIII(s, q, p)
end

function TableauSLRKLobattoIIID(s)
    l = TableauLobattoIIID(s)
    TableauSLRKLobattoIII(s, l, l)
end

function TableauSLRKLobattoIIIE(s)
    l = TableauLobattoIIIE(s)
    TableauSLRKLobattoIII(s, l, l)
end


using StatsModels, RDatasets, DataFrames, GLM, Random
Random.seed!(0)

n = 30
df = dataset("MASS", "cpus")[1:n,:]
df.Freq = map( x->10^9/x , df.CycT)
df = df[:, [:Perf, :MMax, :Cach, :ChMax, :Freq]]
df.Junk1 = rand(n)
df.Junk2 = rand(n)

function stepReg(df, reVar, pThresh)
    predVars = setdiff(propertynames(df), [reVar])
    numVars = length(predVars)
    model = nothing
    while numVars > 0
        fm = term(reVar) ~ term.((1, predVars...))
        model = lm(fm, df)
        pVals = coeftable(model).cols[4][2:end]
        println("Variables: ", predVars)
        println("P-values = ", round.(pVals,digits = 3))
        pVal, knockout = findmax(pVals)
        pVal < pThresh && break
        println("\tRemoving the variable ", predVars[knockout],
                " with p-value = ", round(pVal,digits=3))
        deleteat!(predVars,knockout)
        numVars = length(predVars)
    end
    model
end

model = stepReg(df, :Perf, 0.05)
println(model)


### CPO ###
harmonicMean(x) = 1 / mean(1 / x)

"""
CPO for one data point

likelihood: likelihood for that data point. Note that this should take one
sample of the parameters and return the likelihood, and NOT the log-likelihood.

http://webpages.math.luc.edu/~ebalderama/myfiles/modelchecking101_pres.pdf
"""

function cpo(params::Vector{T}, likelihood::Function) where T
  return harmonicMean(likelihood.(params))
end

mutable struct CPOstream{T <: AbstractFloat}
  invLikelihoodSum::Vector{T}
  counter::Int

  CPOstream{T}() where T = new(T[], 0)
end

function updateCPO(cs::CPOstream{T}, like::Vector{T}) where {T <: AbstractFloat}
  @assert cs.counter >= 0

  # Either initialize or Update inverse-likelihood sum
  if cs.counter == 0
    # Initialize CPO stream
    cs.invLikelihoodSum = 1 ./ like
  else
    # Update invLikelihoodSum
    cs.invLikelihoodSum += 1 ./ like
  end

  # Update counter
  cs.counter += 1
end

"""
Computes (elementwise) log CPO for each observation.
"""
function computeLogCPO(cs::CPOstream{T}) where {T <: AbstractFloat}
  # NOTE: Equivalent to `log(cs.counter ./ cs.invLikelihoodSum)`
  return log(cs.counter) .- log.(cs.invLikelihoodSum)
end

### LPML ###
function computeLPML(cs::CPOstream{T}; verbose::Int=1) where {T <: AbstractFloat}
  mean_log_cpo = mean(computeLogCPO(cs))

  if isinf(mean_log_cpo) || isnan(mean_log_cpo)
    # Reset LPML
    cs.counter = 0
  end

  return mean_log_cpo
end


using StatsBase
using Statistics

function fuel2pos(positions,finalposition)
    sum(abs.(positions .- finalposition))
end
function answer1()
    input = "input"
    crap_positions0 = parse.(Int,split(readline(input),","))
    test_final_position = median(crap_positions0)
    println("range: ", minimum(crap_positions0)," - ", maximum(crap_positions0))
    min = Inf
    minpos = -1
    for pos in minimum(crap_positions0):1:maximum(crap_positions0)
        distance = fuel2pos(crap_positions0,pos)
        if distance < min
            min = distance
            minpos = pos
        end
    end
    min
    # crap_positions0
end

@show answer1()

triangulardistance(n) = Int(n*(n+1)/2)
function fuel2pos2(positions,finalposition)
    sum(triangulardistance.(abs.(positions .- finalposition)))
end
function answer2()
    input = "input"
    crap_positions0 = parse.(Int,split(readline(input),","))
    println("range: ", minimum(crap_positions0)," - ", maximum(crap_positions0))
    min = Inf
    minpos = -1
    for pos in minimum(crap_positions0):1:maximum(crap_positions0)
        distance = fuel2pos2(crap_positions0,pos)
        if distance < min
            min = distance
            minpos = pos
        end
    end
    min
    # crap_positions0
end

@show answer2()


# https://github.com/grasingerm/lbxflow/blob/master/inc/lattice.jl
# http://mafija.fmf.uni-lj.si/seminar/files/2012_2013/Lattice_Boltzmann_method.pdf

# const original_order = [1, 2 ,3, 4]

# From: E3,1,5
struct D2Q5{
    T <: Real,
    Abscissaes <: AbstractMatrix{<:Integer},
    Weights <: AbstractVector{<:Real},
} <: Quadrature
    abscissae::Abscissaes
    weights::Weights
    speed_of_sound_squared::T
end

D2Q5() = D2Q5(
    [
        0 1 0 -1 0
        0 0 1 0 -1
    ],
    [
        4 / 6,
        1 / 12,
        1 / 12,
        1 / 12,
        1 / 12,
        # 1/3, 1/6, 1/6, 1/6, 1/6
    ],
    6.0,
)

order(q::D2Q5) = 3
function opposite(q::D2Q5, idx::Int64)
    if idx == 1
        return 1
    end
    if idx <= 3
        return idx + 2
    end
    return idx - 2
end

Base.show(io::IO, q::D2Q5) = show(io, "D2Q5")
Base.string(q::D2Q5) = "D2Q5"

# The D2Q5 lattice is unable to include temperature
pressure(q::D2Q5, f::VT, ρ::T, u::VT) where {T <: Real, VT <: AbstractVector{T}} = one(T)


module TestVariance
    using StreamStats
    using Distributions
    using Base.Test

    # Means of uniform draws
    for n in rand(1:1_000_000, 100)
        xs = rand(n)
        stat = StreamStats.Variance()
        for x in xs
            update!(stat, x)
        end
        online_v = var(stat)
        online_s = std(stat)
        online_vs = state(stat)
        online_n = nobs(stat)
        batch_v = var(xs)
        batch_s = std(xs)
        @test_approx_eq(online_v, batch_v)
        @test_approx_eq(online_vs, batch_v)
        @test_approx_eq(online_s, batch_s)
        @test online_n == n
    end
end


using MultivariateFunctions
tol = 10*eps()

# pe functions
function test_pe(a::Float64,b::Float64,base::Float64,d::Int,x::Float64)
    return a * exp(b*(x-base)) * (x-base)^d
end
f1 = PE_Function(1.0, 2.0,3.0, 4)
f1_test_result = test_pe(1.0, 2.0,3.0, 4,5.0)
abs(evaluate(f1, 5.0) - f1_test_result) < tol
f2 = PE_Function(0.0, 2.0,3.0, 4)
f2_test_result = 0.0
abs(evaluate(f2, 5.0) - f2_test_result) < tol
f3 = PE_Function(1.0, 8.0,6.0, 8)
f3_test_result = test_pe(1.0, 8.0,6.0, 8,5.0)
abs(evaluate(f3, 5.0) - f3_test_result) < tol

# Sum of functions
sum0 =  Sum_Of_Functions([])
typeof(sum0) == MultivariateFunctions.Sum_Of_Functions
sum1 =  Sum_Of_Functions([f1])
typeof(sum1) == MultivariateFunctions.Sum_Of_Functions
sum3 =  Sum_Of_Functions([f2,f3])
typeof(sum3) == MultivariateFunctions.Sum_Of_Functions
sum4 =  Sum_Of_Functions([f1,f2,f3])
typeof(sum4) == MultivariateFunctions.Sum_Of_Functions
length(sum4.functions_) == 2
sum5 =  Sum_Of_Functions([sum3, sum4])
typeof(sum5) == MultivariateFunctions.Sum_Of_Functions
length(sum5.functions_) == 2
abs(evaluate(sum1, 5.0) - f1_test_result) < tol
abs(evaluate(sum5, 5.0) - 2*f3_test_result - f1_test_result) < 100* tol
abs(evaluate(sum0, 5.0)) < tol

coordinates = Dict(Symbol.(["x", "y", "z", "w"]) .=> [0.5, 2.0, 3.5, 4.0])

f = PE_Function(2.4, Dict{Symbol,PE_Unit}([:x,:y,:z,:w] .=> [ PE_Unit(2.0,3.0,2), PE_Unit(2.0,3.0,2) , PE_Unit(2.0,3.0,2), PE_Unit(2.0,3.0,2)   ]))
new_bases = Dict{Symbol,Float64}([:x,:y,:z, :q] .=> [ 2.8,3.0,1.0,3.0  ])
converted_f = change_base(f,new_bases)
abs(sum(evaluate.(converted_f, Ref(coordinates))) - evaluate(f, coordinates)) < tol

fl = 8.9
integ = 7

function test_result(func, expected_type, eval_to, len = 1)
    val_test = abs(evaluate(func, 5.0) - eval_to) < 1e-09
    if (!val_test)
        print("Failed Val Test")
    end
    type_test = typeof(func) == expected_type
    if (!type_test)
        print("Failed Type Test")
    end
    if typeof(func) == MultivariateFunctions.Sum_Of_Functions
        len_test = length(func.functions_) == len
        if (!len_test)
            print("Failed Length Test")
        end
        return all([val_test,type_test,len_test])
    else
        return all([val_test,type_test])
    end
end

# pe and Float
test_result(f1 + fl, MultivariateFunctions.Sum_Of_Functions, f1_test_result + fl,2 )
test_result(f1 - fl, MultivariateFunctions.Sum_Of_Functions, f1_test_result - fl,2 )
test_result(f1 * fl, MultivariateFunctions.PE_Function, (f1_test_result) * fl )
test_result(f1 / fl, MultivariateFunctions.PE_Function, (f1_test_result) / fl )
test_result(fl + f1, MultivariateFunctions.Sum_Of_Functions, f1_test_result + fl,2 )
test_result(fl - f1, MultivariateFunctions.Sum_Of_Functions, fl - f1_test_result,2 )
test_result(fl * f1, MultivariateFunctions.PE_Function, (f1_test_result) * fl )
# pe and Int
test_result(f1 + integ, MultivariateFunctions.Sum_Of_Functions, f1_test_result + integ,2 )
test_result(f1 - integ, MultivariateFunctions.Sum_Of_Functions, f1_test_result - integ,2 )
test_result(f1 * integ, MultivariateFunctions.PE_Function, (f1_test_result) * integ )
test_result(f1 / integ, MultivariateFunctions.PE_Function, (f1_test_result) / integ )
test_result(integ + f1, MultivariateFunctions.Sum_Of_Functions, f1_test_result + integ,2 )
test_result(integ - f1, MultivariateFunctions.Sum_Of_Functions, -f1_test_result + integ,2 )
test_result(integ * f1, MultivariateFunctions.PE_Function, (f1_test_result) * integ )

# Sum and Float
test_result(sum4 + fl, MultivariateFunctions.Sum_Of_Functions, f1_test_result + f3_test_result + fl, 3 )
test_result(sum4 - fl, MultivariateFunctions.Sum_Of_Functions, f1_test_result + f3_test_result - fl, 3 )
test_result(sum4 * fl, MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result) * fl, 2 )
test_result(sum4 / fl, MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result) / fl, 2 )
test_result(fl + sum4, MultivariateFunctions.Sum_Of_Functions, f1_test_result + f3_test_result + fl, 3 )
test_result(fl - sum4, MultivariateFunctions.Sum_Of_Functions, -f1_test_result -f3_test_result + fl, 3 )
test_result(fl * sum4, MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result) * fl, 2 )
# Sum and Int
test_result(sum4 + integ, MultivariateFunctions.Sum_Of_Functions, f1_test_result + f3_test_result + integ, 3 )
test_result(sum4 - integ, MultivariateFunctions.Sum_Of_Functions, f1_test_result + f3_test_result - integ, 3 )
test_result(sum4 * integ, MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result) * integ, 2 )
test_result(sum4 / integ, MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result) / integ, 2 )
test_result(integ + sum4, MultivariateFunctions.Sum_Of_Functions, f1_test_result + f3_test_result + integ, 3 )
test_result(integ - sum4, MultivariateFunctions.Sum_Of_Functions, -f1_test_result -f3_test_result + integ, 3 )
test_result(integ * sum4, MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result) * integ, 2 )

### Making sums with addition and subtraction.
evaluate(sum4, 5.0) == evaluate(f1 + f2 + f3, 5.0)
abs(evaluate(sum4, 5.0) - 2*evaluate(f3, 5.0) - evaluate(f1 + f2 - f3, 5.0)) < 0.001

### Changing of base
f1_unchanged = change_base(f1, Dict{Symbol,Float64}(:default => 3.0))
abs(f1.units_[:default].base_ - f1_unchanged[1].units_[:default].base_) < tol

f1_changed = change_base(f1, Dict{Symbol,Float64}(:default => 4.0))
abs(f1_changed[1].units_[:default].base_ - 4.0) < tol
abs(f1_changed[3].units_[:default].base_ - 4.0) < tol
abs(sum(evaluate.(f1_changed, 5.0)) - f1_test_result) < 100*tol # Changing bases should not change this.

f1_changed_again = change_base(f1, Dict{Symbol,Float64}(:default => -1.0))
abs(f1_changed_again[1].units_[:default].base_ + 1.0) < tol
abs(f1_changed_again[3].units_[:default].base_ + 1.0) < tol
abs(sum(evaluate.(f1_changed_again, 5.0)) - f1_test_result) < 1000000*tol # Changing bases should not change this.

f3_changed = change_base(f3, Dict{Symbol,Float64}(:default => 11.0))
abs(f3_changed[1].units_[:default].base_ - 11.0) < tol
abs(f3_changed[3].units_[:default].base_ - 11.0) < tol
abs(sum(evaluate.(f3_changed, 5.0)) - f3_test_result) < 1000000*tol # Changing bases should not change this.

### multiplication of functions
test_result( f1 * f3 , MultivariateFunctions.Sum_Of_Functions, (f1_test_result * f3_test_result), 9 )
test_result( f1 * PE_Function(1.0,2.0,3.0,4) , MultivariateFunctions.PE_Function, (f1_test_result * f1_test_result), 1 )
test_result( f1 * sum4 , MultivariateFunctions.Sum_Of_Functions, f1_test_result * (f1_test_result + f3_test_result), 10 )
test_result( sum4 * f1 , MultivariateFunctions.Sum_Of_Functions, f1_test_result * (f1_test_result + f3_test_result), 10 )
test_result( sum4 * sum4 , MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result)^2, 11 )

### Powers pe
test_result( f1 ^ 0 , MultivariateFunctions.PE_Function, 1.0 )
test_result( f1 ^ 1 , MultivariateFunctions.PE_Function, (f1_test_result) )
test_result( f1 ^ 2 , MultivariateFunctions.PE_Function, (f1_test_result * f1_test_result) )
abs(evaluate(f1 ^ 4 ,5.0) - (f1_test_result * f1_test_result * f1_test_result * f1_test_result)) < 1e-03 # This fails normal tol.

### Powers sums
test_result( sum4 ^ 0 , MultivariateFunctions.PE_Function, 1.0 )
test_result( sum4 ^ 1 , MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result), 2 )
test_result( sum4 ^ 2 , MultivariateFunctions.Sum_Of_Functions, (f1_test_result + f3_test_result)^2, 11 )
# Higher powers give an underflow problem with base changes.

## Testing Calculus.
pe_const    = PE_Function(2.0,0.0,5.0,0)
pe_lin      = PE_Function(2.0,0.0,5.0,1)
pe_quad     = PE_Function(2.0,0.0,2.0,2)
pe_exp      = PE_Function(2.0,2.0,1.0,0)
pe_exp_quad = PE_Function(7.0,2.0,2.0,2)

# Linear gradient constant
abs(evaluate(derivative(pe_lin),5.0) - evaluate(derivative(pe_lin),1.0) ) < tol
typeof(derivative(pe_lin)) == MultivariateFunctions.Sum_Of_Functions
# This is also linear
abs(evaluate(derivative(derivative(pe_quad)),5.0) -evaluate(derivative(derivative(pe_quad)),9.0) ) < tol
# derivative into a sum of functions
typeof(derivative(pe_exp_quad)) == MultivariateFunctions.Sum_Of_Functions
abs(evaluate(derivative(pe_exp_quad),5.0) - ( 7.0*exp(2.0*(5.0-2.0))*(5.0-2.0)*(2 + 2.0*(5.0-2.0)) ) ) < tol
# Derivaitve of sum
abs(evaluate(derivative(sum4),5.0) - evaluate(derivative(f1),5.0) - evaluate(derivative(f3),5.0) ) < 1e-09

# Integral of quadratic
abs(integral(pe_quad,1.0,2.0) - ((2/3)*2^3 - 4*2^2 + 8*2) + ((2/3)*1^3 - 4*1^2 + 8*1)) < tol
# Integral of exponential
abs(integral(pe_exp,0.2,3.8) - (exp(2.0*(3.8-1.0)) - exp(2.0*(0.2-1.0))  )) < tol
# Integral of combined.
combined_analytical_integral = PE_Function(3.5,2.0,2.0,2) - PE_Function(14.0/4.0,2.0,2.0,1) + PE_Function(14.0/8.0,2.0,2.0,0)
abs(integral(pe_exp_quad,0.2,3.8) - ( evaluate(combined_analytical_integral, 3.8) - evaluate(combined_analytical_integral, 0.2))) < 1000*tol

# left and right integrals
result = integral(pe_exp_quad,0.2,3.8)
left_integral = integral(pe_exp_quad, Dict{Symbol,Tuple{Float64,Symbol}}(:default => (0.2, :default_right)))
abs(evaluate(left_integral, Dict{Symbol,Float64}(:default_right => 3.8)) - result) < 100 * tol
right_integral = integral(pe_exp_quad, Dict{Symbol,Tuple{Symbol,Float64}}(:default => (:default_left, 3.8)))
abs(evaluate(right_integral, Dict{Symbol,Float64}(:default_left => 0.2)) - result) <  100 * tol
both_integral = integral(pe_exp_quad, Dict{Symbol,Tuple{Symbol,Symbol}}(:default => (:default_left, :default_right)))
abs(evaluate(both_integral, Dict{Symbol,Float64}([:default_left, :default_right] .=> [0.2, 3.8])) - result) <  100 * tol

## Testing linear rescaling
# For PE_Functions
test_func = PE_Function(0.5,1.2,0.9,2)
inp = 4.0
alpha = 0.435
beta = -1.52
rescaled_input = alpha*inp + beta
converted_test_func = convert_to_linearly_rescale_inputs(test_func, alpha, beta)
abs(evaluate(test_func, inp) - evaluate(converted_test_func, rescaled_input)) < 1e-10
# SumOfFunctions
sumFunc = test_func + pe_exp_quad + pe_exp
converted_sumFunc = convert_to_linearly_rescale_inputs(sumFunc, alpha, beta)
abs(evaluate(sumFunc, inp) - evaluate(converted_sumFunc, rescaled_input)) < 1e-10

# Multiplications with no units.
n1 = PE_Function(5.0)
n2 = PE_Function(12.0)
abs((n1 * n2).multiplier_ - 60.0) < tol
abs((n1 * pe_lin).multiplier_ - 10.0) < tol


using Optim
using Devectorize

"""
Linear Regression Cost function with Regularisation

Parameters:
	params      = vector containing rolled up X and Θ
		X       = array of features for each movie
		Θ       = array of weights for each user
	Y           = array of ratings for each movie
	R           = array of the existance of ratings for each movie
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features
	λ           = regularisation constant

Returns:
	J = linear regression cost of the given weights with the given regularisation constant
"""
function cost(params, Y::Array{Int}, numUsers, numMovies, numFeatures, λ)
	# Unfold params -> X, Θ
	X, Θ = unroll(params, numUsers, numMovies, numFeatures)

	# Calculate cost
	J = 0
 	@simd for i = 1:size(Y)[2]
		@inbounds J += (dot(X[Y[2, i], :][:], Θ[Y[1, i], :][:]) - Y[3, i])^2
	end
	J = 1/2 * J
	# add regularisation
	J += λ/2 * (sum(Θ.^2) + sum(X.^2))

	return J
end

"""
Linear Regression Gradient with Regularisation

Parameters:
	params      = vector containing rolled up X and Θ
		X       = array of features for each movie
		Θ       = array of weights for each user
	gradient    = vector containing rolled up Xgrad and ΘGrad from previous iteration
	Y           = array of ratings for each movie
	R           = array of the existance of ratings for each movie
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features
	λ           = regularisation constant

Modifies:
	gradient = XGrad and ΘGrad are updated based on the given weights and regularisation costant
"""
function grad!(params, gradient, Y, numUsers, numMovies, numFeatures, λ)
	@printf "hi"
	# Unfold params -> X, Θ
	X, Θ = unroll(params, numUsers, numMovies, numFeatures)

	XGrad = zeros(size(X))
	ΘGrad = zeros(size(Θ))

	for n = 1:size(Y)[2]
		i = Y[2, n]
		j = Y[1, n]
		@devec x = X[i, :]
		@devec þ = Θ[j, :]
		temp = dot(x[:], þ[:]) - Y[3, n]
		@devec XGrad[i, :] += temp .* þ + λ.*x
		@devec ΘGrad[j, :] += temp .* x + λ.*þ
	end

	gradient[:] = rollup(XGrad, ΘGrad)
end

function costGrad!(params, gradient, Y, numUsers, numMovies, numFeatures, λ)
	# Unfold params -> X, Θ
	X, Θ = unroll(params, numUsers, numMovies, numFeatures)

	XGrad = zeros(size(X))
	ΘGrad = zeros(size(Θ))

	# Calculate cost
	J = 0
	count = Int(floor(size(Y)[2] * 0.1))
	start = rand(1:size(Y)[2] - count)
	for n = start:start+count
	#for n = 1:size(Y)[2]
		i = Y[2, n]
		j = Y[1, n]
		@devec x = X[i, :]
		@devec þ = Θ[j, :]
		temp = dot(x[:], þ[:]) - Y[3, n]
		J += temp^2
		@devec XGrad[i, :] += temp .* þ + λ.*x
		@devec ΘGrad[j, :] += temp .* x + λ.*þ
	end
	J = 1/2 * J
	# add regularisation
	J += λ/2 * (sum(Θ.^2) + sum(X.^2))

	gradient[:] = rollup(XGrad, ΘGrad)

	return J
end

"""
Unroll parameter vector into X and Θ

Parameters:
	params = vector containing rolled up X and Θ
	numUsers    = number of users
	numMovies   = number of movies
	numFeatures = number of features

Returns:
	X = array of features for each movie
	Θ = array of weights for each user
"""
function unroll(params, numUsers, numMovies, numFeatures)
	X = reshape(params[1 : numMovies * numFeatures], numMovies, numFeatures)
	Θ = reshape(params[numMovies * numFeatures + 1 : end], numUsers, numFeatures)
	return X, Θ
end

"""
Roll up X and Θ into single vector

Parameters:
	X = array of features for each movie
	Θ = array of weights for each user

Returns:
	params = vector containing rolled up X and Θ
"""
function rollup(X, Θ)
	return params = vcat(X..., Θ...)
end

"""
Entry point

Params:
	dataPath = path to() CSVs
"""
function main()
	dataPath = ""
	if ARGS[1] == "small"
		dataPath = "../ml-latest-small/fixed/"
	elseif ARGS[1] == "full"
		dataPath = "../ml-latest/fixed/"
	end

	# Load data from CSV - using wc to speed up loading by preallocating
	# N.B. requires cat and wc to be installed - a shell should not be nessecary though
	lines = parse(readall(pipeline(`cat $[dataPath]ratings.csv`, `wc -l`)))
	@time ratings = readdlm(dataPath * "ratings.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]train.csv`, `wc -l`)))
	@time train = readdlm(dataPath * "train.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]cross.csv`, `wc -l`)))
	@time cross = readdlm(dataPath * "cross.csv", ',', header=true, dims=(lines, 4))[1]

	lines = parse(readall(pipeline(`cat $[dataPath]movies.csv`, `wc -l`)))
	@time movies = readdlm(dataPath * "movies.csv", ',', header=true, dims=(lines, 3))
	moviesHeader = movies[2]
	movies = movies[1]

	@printf "Finished loading, begin processing\n"

	train[:, 3] *= 2
	train = Array{Int}(train[:, 1:3]')
	cross[:, 3] *= 2
	cross = Array{Int}(cross[:, 1:3]')

	# Count users, movies and features
	numUsers = length(union(ratings[:,1]))
	numMovies = size(movies)[1]

	numFeatures = 18
	λ = 0.05
	X = rand(numMovies, 18)
	Θ = rand(numUsers, 18)

	if numFeatures == 18
		for i = 1:numMovies
			X[i, 1] = contains(movies[i, 3], "Action") ? 2 : 0
			X[i, 2] = contains(movies[i, 3], "Adventure") ? 2 : 0
			X[i, 3] = contains(movies[i, 3], "Animation") ? 2 : 0
			X[i, 4] = contains(movies[i, 3], "Children") ? 2 : 0
			X[i, 5] = contains(movies[i, 3], "Comedy") ? 2 : 0
			X[i, 6] = contains(movies[i, 3], "Crime") ? 2 : 0
			X[i, 7] = contains(movies[i, 3], "Documentary") ? 2 : 0
			X[i, 8] = contains(movies[i, 3], "Drama") ? 2 : 0
			X[i, 9] = contains(movies[i, 3], "Fantasy") ? 2 : 0
			X[i, 10] = contains(movies[i, 3], "Film-Noir") ? 2 : 0
			X[i, 11] = contains(movies[i, 3], "Horror") ? 2 : 0
			X[i, 12] = contains(movies[i, 3], "Musical") ? 2 : 0
			X[i, 13] = contains(movies[i, 3], "Mystery") ? 2 : 0
			X[i, 14] = contains(movies[i, 3], "Romance") ? 2 : 0
			X[i, 15] = contains(movies[i, 3], "Sci-Fi") ? 2 : 0
			X[i, 16] = contains(movies[i, 3], "Thriller") ? 2 : 0
			X[i, 17] = contains(movies[i, 3], "War") ? 2 : 0
			X[i, 18] = contains(movies[i, 3], "Western") ? 2 : 0
		end
	end
	scoreboard = Array{Any}(0,3)
	for i = 0.1:0.10:1.00
		tempTrain = train[:, 1:Int(floor(i*size(train)[2]))]

		function costTemp(x)
			cost(x, train, numUsers, numMovies, numFeatures, λ)
		end

		function gradTemp!(x, y)
			grad!(x, y, train, numUsers, numMovies, numFeatures, λ)
		end

		function costGradTemp!(x, y)
			costGrad!(x, y, tempTrain, numUsers, numMovies, numFeatures, λ)
		end

		tic()
		@time @show results = optimize(DifferentiableFunction(costTemp, gradTemp!, costGradTemp!), rollup(X,Θ), iterations=10)
		time = toc()

		@show score = cost(results.minimum, tempTrain, numUsers, numMovies, numFeatures, 0) / size(tempTrain)[2]
		@show score = cost(results.minimum, cross, numUsers, numMovies, numFeatures, 0) / size(cross)[2]
		scoreboard = [scoreboard; λ score time]
	end

	#writecsv("lambba-scoreboard.csv", scoreboard)

	@printf "Done :D\n"
end

main()


using Plots

X = randn(100,10)

function pca(x::Array{Float64}, components::Int64)
    s = svdfact(x)
    estimate =  s[:U][:,1:components]*
                Diagonal(s[:S][1:components])*
                s[:Vt][1:components,:]
    residual = x-estimate
    loss = vecnorm(residual)
    return Dict(:eigenvalues    => (s[:S]/size(x,1)).^2,
                :estimate       => estimate,
                :residual       => residual,
                :loss           => loss)
end

function Plots.plot(x::Dict; kwargs...)
    if (:eigenvalues in keys(x))
        plot(x[:eigenvalues]/sum(x[:eigenvalues]), label = "",
             title = "Proportion of explained variance",
             xlabel = "Component number"; kwargs...)
    else
        plot(x; kwargs...)
    end
end

p1 = plot([pca(X,i)[:loss] for i in 1:10], labels = "", size = (500,500))
xlabel!("number of components")
ylabel!("loss")
title!("Principal Components loss")

p = pca(X, 3)
p2 = plot(p, size = (500,500))

p3 = plot(p1, p2, layout = (1,2), size = (1000,500))
savefig("pca")


function perform_dijkstra_fm(W, pstart; method="fm", svg_rate=10, niter=Inf, bound="sym")
    """
        # perform_fm_dijstra - *slow* (matlab) implementation of Dijstra and FM
        #
        #   [D,Dsvg,Ssvg] = perform_fm_dijstra(W, pstart, options);
        #
        #   W is an (n,n) metric matrix.
        #   pstart is a (2,k) starting points.
        #   options.method is either 'fm' or 'dijstra'
        #
        #   D is the final distance map to pstart
        #   options.svg_rate gives the rate at wich Dsvg and Ssvg is filled.
        #   options.niter can be used to limit the total number of steps (partial propagation).
        #
        #   Copyright (c) 2012 Gabriel Peyre
    """

##
# Size.

n = size(W,1)

##
# The four displacement vector to go to the four neightbors.

neigh = [[1;0] [-1;0] [0;1] [0;-1]]

##
# For simplicity of implementation, we use periodic boundary conditions.

function symmetrize(x,n)
    if (x<=0)
        x = 2-x;
    elseif (x>n)
        x = 2*n-x
    end
    return x
end

if bound == "per"
    boundary = x -> mod(x-1,n)+1
elseif bound == "sym"
    boundary = x -> [symmetrize(x[1],n), symmetrize(x[2],n)]
else
    error("Works only for per and sym.")
end
#     case 'sym'
#         boundary = @(x)x.*(x<=n & x>0) + (2-x).*(x<=0) + (2*n-x).*(x>n);
#     otherwise
#         error('Works only for per and sym.');
# end

##
# For a given grid index |k|, and a given neighboring index k in \({1,2,3,4}\),
# |Neigh(k,i)| gives the corresponding grid neigboring index.

ind2sub1 = k -> Int.([rem(k-1, n)+1; (k - rem(k-1, n) - 1)/n + 1])
sub2ind1 = u -> Int((u[2]-1)*n + u[1])
Neigh = (k,i) -> sub2ind1(boundary(ind2sub1(k) + neigh[:,i]));
extract   = (x,I) -> x[I]
extract1d = (x,I) -> extract(vec(x),I);

##
# Stack of starting points.
nstart = size(pstart,2)
I = Int.(zeros(nstart))
for i=1:nstart
    I[i] = sub2ind1(pstart[:,i])
end

##
# Initialize the distance to \(+\infty\), excepted for the boundary conditions.

D = zeros(n,n)+Inf; # current distance
for i in 1:nstart
    D[pstart[1,i], pstart[2,i]] = 0;
end

##
# Initialize the state to 0 (unexplored), excepted for the boundary point to \(1\)
# (front).

S = zeros(n,n);
for i in 1:nstart
    S[pstart[1,i],pstart[2,i]] = 1 # open
end

##
# Run!
compt=0
iter = 0
q = 100  # maximum number of saves
Dsvg = zeros(n,n,q)
Ssvg = zeros(n,n,q)
while !(isempty(I)) && iter<=niter
    iter = iter+1
    # pop from stack
    j = sortperm(extract1d(D,I))
    j=j[1]
    i = I[j]
    deleteat!(I,j)
    # declare dead
    u = ind2sub1(i);
    S[u[1],u[2]] = -1;
    # Make a list of neighbors that are not dead
    J = []
    for k in 1:4
        j = Neigh(i,k)
        if extract1d(S,j)!= -1
            # add to the list of point to update
            append!(J,j)
            if extract1d(S,j)==0
                # add to the front
                u = ind2sub1(j)
                S[u[1],u[2]] = 1
                append!(I,j)
            end
        end
    end
    # update neighbor values
    DNeigh = (D,k) -> extract1d(D,Neigh(j,k))
    for j in J
        dx = min(DNeigh(D,1), DNeigh(D,2))
        dy = min(DNeigh(D,3), DNeigh(D,4))
        u = ind2sub1(j)
        w = extract1d(W,j);
        if method=="dijkstra"
            D[u[1],u[2]] = min(dx + w, dy + w)
        else
            Delta = 2*w - (dx-dy).^2
            if (Delta>=0)
                D[u[1],u[2]] = (dx + dy + sqrt(Delta))/ 2
            else
                D[u[1],u[2]] = min(dx + w, dy + w)
            end
        end
    end
    # svd
    t = Base.div(iter,svg_rate)
    if mod(iter,svg_rate)==0 && (t<q)
        Dsvg[:,:,t] = D
        Ssvg[:,:,t] = S
    end
end
return (D,Dsvg,Ssvg)
end


using Plots
using MLFlowClient
using Random

# Parameters
αs = [0.0, 0.9, 0.98]
n = 100

"""Method that generates price paths of length `n` based on `α`"""
function getpricepath(α, n)
    x = zeros(n + 1)
    x[1] = 0.0
    for t in 1:n
        x[t+1] = α * x[t] + rand()
    end
    x
end
p = plot()

# Create MLFlow instance
mlf = MLFlow("http://localhost:5000")

# Initiate new experiment
experiment_id = createexperiment(mlf; name="price-paths")

# Create a run in the new experiment
exprun = createrun(mlf, experiment_id)

# Log parameters and their values
for (idx, α) in enumerate(αs)
    logparam(mlf, exprun, "alpha$(idx)", string(α)) # MLFlow only supports string parameter values
end

# Obtain pricepaths
pricepaths = [getpricepath(α, n) for α in αs]

# Log pricepaths in MLFlow
for (idx, pricepath) in enumerate(pricepaths)
    plot!(p,
        pricepath,
        title="Random price paths",
        label="alpha = $(αs[idx])",
        xlabel="Timestep",
        ylabel="Price"
    )

    logmetric(mlf, exprun, "pricepath$(idx)", pricepath)
end

# Save the price path plot as an image
plotfilename = "pricepaths-plot.png"
png(plotfilename)

# Upload the plot as an artifact associated with the MLFlow experiment's run
logartifact(mlf, exprun, plotfilename)

# remote temporary plot which was already uploaded in MLFlow
rm(plotfilename)

# complete the experiment
updaterun(mlf, exprun, "FINISHED")


myParse(x) = parse.(Int, x)
energyLevels = Matrix(hcat(myParse.(collect.(readlines("input.txt")))...)')

nRows, nCols = size(energyLevels)
nLevels = length(energyLevels)

function coordfromIndex(index)
    row = ((index - 1) % nRows) + 1
    col = ((index - 1) ÷ nRows) + 1
    return row, col
end

function indexfromCoord(row, col)
    return nRows * (col - 1) + row
end

function surroundingIndices(index)
    row, col = coordfromIndex(index)
    arr = [(row - 1, col), (row, col - 1), (row + 1, col), (row, col + 1), (row + 1, col +1), (row - 1, col - 1), (row - 1, col + 1), (row + 1, col - 1)]
    return [indexfromCoord(x, y) for (x, y) in arr if x in 1:nRows && y in 1:nCols]
end

function simulatePt1(steps, elevels)

    function flashNode(index, flashed)
        flashed[index] = true
        for i in surroundingIndices(index)
            elevels[i] += 1
            if elevels[i] > 9 && !flashed[i]
                flashNode(i, flashed)
            end
        end
    end

    function simulateFlashes()
        flashed = [false for i in 1:nLevels]
        for i in 1:nLevels
            if elevels[i] > 9 && !flashed[i]
                flashNode(i, flashed)
            end
        end

        for (i, hasFlashed) in enumerate(flashed)
            if hasFlashed
                elevels[i] = 0
            end
        end

        return count(flashed)
    end

    flashes = 0

    for _ in 1:steps
        elevels .= elevels .+ 1
        flashes += simulateFlashes()
    end

    return flashes
end

simulate(100, copy(energyLevels))

function simulatePt2(elevels)

    function flashNode(index, flashed)
        flashed[index] = true
        for i in surroundingIndices(index)
            elevels[i] += 1
            if elevels[i] > 9 && !flashed[i]
                flashNode(i, flashed)
            end
        end
    end

    function simulateFlashes()
        flashed = [false for i in 1:nLevels]
        for i in 1:nLevels
            if elevels[i] > 9 && !flashed[i]
                flashNode(i, flashed)
            end
        end

        for (i, hasFlashed) in enumerate(flashed)
            if hasFlashed
                elevels[i] = 0
            end
        end

        return count(flashed) == length(flashed)
    end

    step = 1

    while true
        elevels .= elevels .+ 1
        if simulateFlashes() == true
            return step
        end
        step += 1
    end
end

simulatePt2(energyLevels)