#########################################################
# Hinge Loss
#
# L(\hat{y},y) = 1/2N \sum_{i=1}^N (1-y_i \hat{y}_i)
#########################################################
@defstruct HingeLossLayer Layer (
    name :: AbstractString = "hinge-loss",
    (weight :: AbstractFloat = 1.0, weight >= 0),
    (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == 2),
)

@characterize_layer(HingeLossLayer,
    has_loss  => true,
    can_do_bp => true,
    is_sink   => true,
    has_stats => true,
)

type HingeLossLayerState{T} <: LayerState
    layer :: HingeLossLayer
    loss  :: T

    loss_accum :: T
    n_accum    :: Int

#     # a helper blob used to compute the loss without destroying
#     # the prediction results passed up
#     pred_copy :: Blob
#     bp_mask :: Blob

    # a helper blob used to compute the loss
    loss_blob :: Blob
end

function setup(backend::Backend, layer::HingeLossLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
    data_type = eltype(inputs[1])
#     pred_copy = make_blob(backend, data_type, size(inputs[1])...)
#     bp_mask = make_blob(backend, Bool, size(inputs[1])...)
    loss_blob = make_blob(backend, data_type, 1)

    state = HingeLossLayerState(layer, zero(data_type), zero(data_type), 0, loss_blob)#, pred_copy, bp_mask)
    return state
end
function shutdown(backend::Backend, state::HingeLossLayerState)
#     destroy(state.pred_copy)
#     destroy(state.bp_mask)
  destroy(state.loss_blob)
  nothing
end
function reset_statistics(state::HingeLossLayerState)
    state.n_accum = 0
    state.loss_accum = zero(typeof(state.loss_accum))
end
function dump_statistics(storage, state::HingeLossLayerState, show::Bool)
    update_statistics(storage, "$(state.layer.name)-hinge-loss", state.loss_accum)

    if show
        loss = @sprintf("%.4f", state.loss_accum)
        @info("  Hinge-loss (avg over $(state.n_accum)) = $loss")
    end
end

function forward(backend::CPUBackend, state::HingeLossLayerState, inputs::Vector{Blob})
  pred  = inputs[1]
  label = inputs[2]

  data_type = eltype(pred)
  n = length(pred)

  state.loss = zero(data_type)
  for i=1:n
    loss = one(data_type) - pred.data[i]*label.data[i]
    if loss > zero(data_type)
      state.loss += loss
    end
  end
  state.loss /= convert(data_type, get_num(pred))

  # accumulate statistics
  state.loss_accum = (state.loss_accum*state.n_accum + state.loss*get_num(pred)) / (state.n_accum+get_num(pred))
  state.n_accum += get_num(pred)
end

function backward(backend::CPUBackend, state::HingeLossLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  const pred = inputs[1]
  const label = inputs[2]

  const data_type = eltype(pred)
  const n = length(pred)
  const num = get_num(pred)

  if isa(diffs[1], CPUBlob)
    erase!(diffs[1])
    for i = 1:n
      if pred.data[i]*label.data[i] < one(data_type)
        diffs[1].data[i] = convert(data_type, -state.layer.weight/num) * label.data[i]
        end
    end
  end

  if isa(diffs[2], CPUBlob)
    erase!(diffs[2])
    for i = 1:n
      if pred.data[i]*label.data[i] < one(data_type)
        diffs[2].data[i] = convert(data_type, -state.layer.weight/num) * pred.data[i]
        end
    end
  end

end


const BasisSet = Array{CGBF,1}

include("read_basis_g94.jl")
import .g94basis

function build_basis_new( atoms::Atoms )

    data_dict = g94basis.read_basis_set("g94_sto-3g")

    basis = BasisSet()

    Natoms = atoms.Natoms
    atsymbs = atoms.atsymbs

    Nbasis = 0
    for ia = 1:Natoms
        #
        x = atoms.positions[1,ia]
        y = atoms.positions[2,ia]
        z = atoms.positions[3,ia]
        #
        data = data_dict[atsymbs[ia]]
        Nsyms = data.Nsyms
        #
        for isym = 1:Nsyms
            sym = data.syms[isym]
            for (I,J,K) in sym2power[sym]
                cbf = init_CGBF( (x,y,z), (I,J,K) )
                ncontr = data.Ncontr[isym]
                for i = 1:ncontr
                    expn = data.expns[isym][i]
                    coef = data.coeffs[isym][i]
                    push!(cbf,expn,coef)
                end
                push!(basis,cbf)
                Nbasis = Nbasis + 1
            end
        end
    end

    println("Nbasis = ", Nbasis)

    return basis

end

function build_basis( atoms::Atoms, name="sto3g", verbose::Bool=false )

    data = basis_set_data[name]
    basis = BasisSet()

    Natoms = atoms.Natoms
    atsymbs = atoms.atsymbs

    Nbasis = 0
    for ia = 1:Natoms
        #
        atno = ZATOMS[atsymbs[ia]]
        x = atoms.positions[1,ia]
        y = atoms.positions[2,ia]
        z = atoms.positions[3,ia]
        #
        for btuple in data[atno]
            sym, primlist = btuple
            for (I,J,K) in sym2power[sym]
                cbf = init_CGBF( (x,y,z), (I,J,K) )
                for (expn,coef) in primlist
                    push!(cbf,expn,coef)
                end
                push!(basis,cbf)
                Nbasis = Nbasis + 1
            end
        end
    end

    println("Nbasis = ", Nbasis)

    return basis
end

const sym2power = Dict{Any,Any}(
  "S" => [(0,0,0)],
  "P" => [(1,0,0),(0,1,0),(0,0,1)],
  "D" => [(2,0,0),(0,2,0),(0,0,2),(1,1,0),(1,0,1),(0,1,1)],
  "F" => [(3,0,0),(0,3,0),(0,0,3),(2,1,0),(2,0,1),(1,2,0),
          (1,1,1),(1,0,2),(0,2,1),(0,1,2)]
  )


# ```math
# x1[a_, b_, c_, t_] :=  Sin[.5 t] - a Sin[b t]*Cos[t] - .1 c Sin[10 b t];
# y1[a_, b_, c_, t_] :=  Cos[.5 t] - a Sin[b t]*Sin[t] - .1 c Cos[10 b t];
#
# GraphicsGrid[Partition[ParallelTable[
#    With[{
#      a = RandomReal[{-1.5, 1.5}],
#      b = RandomInteger[{3, 15}],
#      c = RandomReal[{0, 1.5}],
#      clr1 = Black,
#      clr2 = RGBColor @@ RandomReal[1, 3],
#      clr3 = RGBColor @@ RandomReal[1, 3],
#      thick = RandomReal[{.04, .5}],
#      tm = 1},
#     ParametricPlot[
#      Evaluate[{{x1[a, b, c, t], y1[a, b, c, t]}, {x1[a, b, c, t],
#         y1[a, b, c, t]}}], {t, 0, tm 4 \[Pi]},
#      PlotStyle -> {{clr2, Thickness[0.001` + 0.05` thick]}, {clr3,
#         Thickness[0.001` + 0.01` thick]}}, Axes -> False,
#      PlotPoints -> 200, PlotRange -> All, Background -> clr1]]
#    , {n, 32}], 4], ImageSize -> 600]
# ```
using Distributions
using OdinSon
using PyPlot

function colorflake()
    clr1 = RGB(rand(3)...)
    clr2 = RGB(rand(3)...)
    thick = rand(Uniform(2, 7)) # need to get this better, MM linewidth is based on size of figure
    tm = 1 # this seems needless in the current code

    a = rand(Uniform(0, 1.5))
    b = rand(3:15)
    c = rand(Uniform(0, 1.5))
    # this creates a closure over the parameters, might not be the best way
    # this is also vectorized since sin/cos are, can't wait for f.() syntax!
    x1(t) = sin(0.5*t) - a*sin(b*t).*cos(t) - 0.1*c*sin(10*b*t)
    y1(t) = cos(0.5*t) - a*sin(b*t).*sin(t) - 0.1*c*cos(10*b*t)

    t = linspace(0, tm*4.0*π, 5000)
    plot(x1(t), y1(t), c=clr1, linewidth=thick)
    plot(x1(t), y1(t), c=clr2, linewidth=0.2*thick)
    gca()[:set_aspect]("equal")
    gca()[:set_axis_bgcolor](NC"black")
end

colorflake()


using EffectiveWaves
using LaTeXStrings
using Plots
height=500
# unicodeplots()
 pyplot(linewidth=3, size=(2.6*height,height))

 Plots.scalefontsizes(1.7)

## choose material

  # concrete
  medium = LimeStone
  inclusion1 = AirDry
  inclusion2 = Brick
  ωs = LinRange(real(medium.c/10000),real(medium.c),100) # k from 0 to 1
  # ωs = LinRange(real(medium.c/10000),real(medium.c/1000),100) # k from 0 to 1

  volfrac = 0.16
  r1 = 0.1; vol1 = 0.06
  r2 = 1.0; vol2 = volfrac - vol1

  sp1 = Specie(ρ=inclusion1.ρ ,r=r1, c=inclusion1.c, volfrac = vol1)
  sp2 = Specie(ρ=inclusion2.ρ ,r=r2, c=inclusion2.c, volfrac = volfrac-vol1)

  # True wavenumber
  kTs  = wavenumber_low_volumefraction(ωs, medium, [sp1,sp2]);
  # Approximate wavenumber
  kTLSs = two_species_approx_wavenumber(ωs, medium, [sp1,sp2]);

  speed_arr = [ ωs./real(kTs), ωs./real(kTLSs), 0 .*ωs + real(medium.c)]
  atten_arr = imag([kTs,kTLSs])

  styles = reshape([:solid,:solid,:dot],1,3)
  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$", "Lime Stone"] ,1, 3);
  ys_arr = speed_arr;
  xs = r1.*(ωs./real(medium.c));
  m =5;
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  plot(xs, ys_arr, xlabel=L"a_S k", ylabel="sound speed", labels=labs, line = styles);
  p1 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  ys_arr = atten_arr;
  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$"] ,1, 2);
  plot(xs, ys_arr, labels=labs, xlabel=L"a_S k", ylabel="attenuation");
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  p2 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  plot(p1,p2)
  savefig("media/compare_concrete.png")
  savefig("media/compare_concrete.pdf")
  gui()

## Zoomed in version
  ωs = LinRange(real(medium.c/10000),ωs[m],100) # k from 0 to 1
  m = length(ωs);
  kTs  = wavenumber_low_volumefraction(ωs, medium, [sp1,sp2]);
  kTLSs = two_species_approx_wavenumber(ωs, medium, [sp1,sp2]);

  speed_arr = [ ωs./real(kTs), ωs./real(kTLSs)]
  atten_arr = imag([kTs,kTLSs])

  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$"] ,1, 2);
  ys_arr = speed_arr;
  xs = r1.*(ωs./real(medium.c));
  plot(xs, ys_arr, xlabel=L"a_S k", ylabel="sound speed", labels=labs
                 , border = false);
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  p1 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  ys_arr = atten_arr;
  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$"] ,1, 2);
  plot(xs, ys_arr, xlabel=L"a_S k", ylabel="attenuation", labels=labs
                 , border=false);
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  dy = abs(y2- y1)/120.0;
  p2 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  plot(p1,p2)
  mkdir("media")
  savefig("media/compare_concrete_zoom.png")
  savefig("media/compare_concrete_zoom.pdf")

Plots.scalefontsizes(1/1.7)


using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

#Cubohemioctahedron prototype DEBUG
"""
Return the Lar.Struct of cubohemioctahedron prototype
"""
function cubohemioctahedron()
	a=0.135
	b=0.268
	c=0.5
	d=1
	h=0.709219858
	V=[0 b d c -a;0 d b -a c; 0 0 0 h h]
	VV=[[1],[2],[3],[4],[5]]
	EV=[[1,2],[1,3],[1,4],[1,5],[2,3],[3,4],[4,5],[2,5]]
	FV=[[1,2,3],[1,3,4],[1,2,5],[1,4,5]]#,[2,3,4,5]
	CV=[[1, 2, 3, 4, 5]]
	tetra = V,EV,FV,CV
	twotetra = Lar.Struct([ tetra ,Lar.r(0,0,2π/3),tetra])
	threetetra = Lar.Struct([ twotetra ,Lar.r(0,0,4π/3),tetra])
	sixtetra=Lar.Struct([ threetetra ,Lar.r(0,π,0)*Lar.r(0,0,853π/1024), threetetra])
	V,EV,FV,CV = Lar.struct2lar(Lar.Struct([ sixtetra ]))
	return Lar.Struct([(V,collect(Set(EV)),FV,CV)])
end


# =======================================================================
# 2KP

using vOptSpecific

# ==========================================
# instance

n  = 3
p1 = [ 3,  9 , 7]
p2 = [16, 15 , 6]
w  = [ 2,  3 , 4]
c  = 8

# ==========================================

myInstance = set2UKP( n , p1 , p2 , w , c )
algorithm = UKP_Jorge2010()
z1, z2, r, x = vSolve( myInstance , algorithm )

for i = 1:length(z1)
    println("(" , z1[i], " | ", z2[i], ") : ", x[i,:])
end


# # Batched Generalized Minimal Residual
# In this tutorial we describe the basics of using the batched gmres iterative solver.
# At the end you should be able to
# 1. Use BatchedGeneralizedMinimalResidual to solve batches of linear systems
# 2. Construct a columnwise linear solver with BatchedGeneralizedMinimalResidual

# ## What is the Generalized Minimal Residual Method?
# The  Generalized Minimal Residual Method (GMRES) is a [Krylov subspace](https://en.wikipedia.org/wiki/Krylov_subspace) method for solving linear systems:
# ```math
#  Ax = b
# ```
# See the [wikipedia](https://en.wikipedia.org/wiki/Generalized_minimal_residual_method) for more details.

# ## What is the Batched Generalized Minimal Residual Method?
# As the name suggests it solves a whole bunch of independent GMRES problems

# ## Basic Example
# First we must load a few things
using ClimateMachine,
    ClimateMachine.LinearSolvers,
    ClimateMachine.BatchedGeneralizedMinimalResidualSolver
using LinearAlgebra, Random, Plots

# Next we define two linear systems that we would like to solve simultaneously.
# The matrix for the first linear system is
A1 = [
    2.0 -1.0 0.0
    -1.0 2.0 -1.0
    0.0 -1.0 2.0
];
# And the right hand side is
b1 = ones(typeof(1.0), 3);
# The exact solution to the first linear system is
x1_exact = [1.5, 2.0, 1.5];
# The matrix for the first linear system is
A2 = [
    2.0 -1.0 0.0
    0.0 2.0 -1.0
    0.0 0.0 2.0
];
# And the right hand side is
b2 = ones(typeof(1.0), 3);
# The exact solution to second linear system is
x2_exact = [0.875, 0.75, 0.5];

# We now define a function that performs the action of each linear operator independently.
function closure_linear_operator(A1, A2)
    function linear_operator!(x, y)
        mul!(view(x, :, 1), A1, view(y, :, 1))
        mul!(view(x, :, 2), A2, view(y, :, 2))
        return nothing
    end
    return linear_operator!
end;

# To understand how this works let us construct an instance
# of the linear operator and apply it to a vector
linear_operator! = closure_linear_operator(A1, A2);
# Let us see what the action of this linear operator is
y1 = ones(typeof(1.0), 3);
y2 = ones(typeof(1.0), 3) * 2.0;
y = [y1 y2];
x = copy(y);
linear_operator!(x, y);
x
# We see that the first column is `A1 * [1 1 1]'`
# and the second column is `A2 * [2 2 2]'`
# that is,
[A1 * y1 A2 * y2]

# We are now ready to set up our Batched Generalized Minimal Residual solver
# We must now set up the right hand side of the linear system
b = [b1 b2];
# as well as the exact solution, (to verify convergence)
x_exact = [x1_exact x2_exact];
# !!! warning
#     For BatchedGeneralizedMinimalResidual the assumption is that each column of b is independent and corresponds to a batch. This will come back later.

# We now use an instance of the solver
linearsolver = BatchedGeneralizedMinimalResidual(b);
# As well as an initial guess, denoted by the variable x
x1 = ones(typeof(1.0), 3);
x2 = ones(typeof(1.0), 3);
x = [x1 x2];
# To solve the linear system, we just need to pass to the linearsolve! function
iters = linearsolve!(linear_operator!, linearsolver, x, b)
# which is guaranteed to converge in 3 iterations since `length(b1)=length(b2)=3`
# We can now check that the solution that we computed, x
x
# has converged to the exact solution
x_exact
# Which indeed it has.
# ## Advanced Example

# We now go through a more advanced application of the Batched Generalized Minimal Residual solver
# !!! warning
#     Iterative methods should be used with preconditioners!
# The first thing we do is define a linear operator that mimics
# the behavior of a columnwise operator in ClimateMachine
function closure_linear_operator!(A, tup)
    function linear_operator!(y, x)
        alias_x = reshape(x, tup)
        alias_y = reshape(y, tup)
        for i6 in 1:tup[6]
            for i4 in 1:tup[4]
                for i2 in 1:tup[2]
                    for i1 in 1:tup[1]
                        tmp = alias_x[i1, i2, :, i4, :, i6][:]
                        tmp2 = A[i1, i2, i4, i6] * tmp
                        alias_y[i1, i2, :, i4, :, i6] .=
                            reshape(tmp2, (tup[3], tup[5]))
                    end
                end
            end
        end
    end
end;
# Next we define the array structure of an MPIStateArray
# in its true high dimensional form
tup = (2, 2, 5, 2, 10, 2);
# We define our linear operator as a random matrix
Random.seed!(1234);
B = [
    randn(tup[3] * tup[5], tup[3] * tup[5])
    for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
];
columnwise_A = [
    B[i1, i2, i4, i6] + 3 * (i1 + i2 + i4 + i6) * I
    for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
];
# as well as its inverse
columnwise_inv_A = [
    inv(columnwise_A[i1, i2, i4, i6])
    for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]
];
columnwise_linear_operator! = closure_linear_operator!(columnwise_A, tup);
columnwise_inverse_linear_operator! =
    closure_linear_operator!(columnwise_inv_A, tup);
# The structure of an MPIStateArray is related to its true
# higher dimensional form as follows:
mpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6]);
# We now define the right hand side of our Linear system
b = randn(mpi_tup);
# As well as the initial guess
x = copy(b);
x += randn(mpi_tup) * 0.1;
# In the previous tutorial we mentioned that it is assumed that
# the right hand side is an array whose column vectors all independent linear
# systems. But right now the array structure of ``x`` and ``b`` do not follow
# this requirement.
# To handle this case we must pass in additional arguments that tell the
# linear solver how to reconcile these differences.
# The first thing that the linear solver must know of is the higher tensor
# form of the MPIStateArray, which is just the `tup` from before
reshape_tuple_f = tup;
# The second thing it needs to know is which indices correspond to a column
# and we want to make sure that these are the first set of indices that appear
# in the permutation tuple (which can be thought of as enacting
# a Tensor Transpose).
permute_tuple_f = (5, 3, 4, 6, 1, 2);
# It has this format since the 3 and 5 index slots
# are the ones associated with traversing a column. And the 4 index
# slot corresponds to a state.
# We also need to tell our solver which kind of Array struct to use
ArrayType = Array;

# We are now ready to finally define our linear solver, which uses a number
# of keyword arguments
gmres = BatchedGeneralizedMinimalResidual(
    b,
    ArrayType = ArrayType,
    m = tup[3] * tup[5] * tup[4],
    n = tup[1] * tup[2] * tup[6],
    reshape_tuple_f = reshape_tuple_f,
    permute_tuple_f = permute_tuple_f,
    atol = eps(Float64) * 10^2,
    rtol = eps(Float64) * 10^2,
);
# ```m``` is the number of gridpoints along a column. As mentioned previously,
# this is `tup[3]*tup[5]*tup[4]`. The ```n``` term corresponds to the batch size
# or the number of columns in this case. ```atol``` and ```rtol``` are relative and
# absolute tolerances

# All the hard work is done, now we just call our linear solver
iters = linearsolve!(
    columnwise_linear_operator!,
    gmres,
    x,
    b,
    max_iters = tup[3] * tup[5] * tup[4],
)
# We see that it converged in less than `tup[3]*tup[5] = 50` iterations.
# Let us verify that it is indeed correct by computing the exact answer
# numerically and comparing it against the iterative solver.
x_exact = copy(x);
columnwise_inverse_linear_operator!(x_exact, b);
# Now we can compare with some norms
norm(x - x_exact) / norm(x_exact)
columnwise_linear_operator!(x_exact, x);
norm(x_exact - b) / norm(b)
# Which we see are small, given our choice of ```atol``` and ```rtol```.
# The struct also keeps a record of its convergence rate
# in the residual member. The convergence rate of each column
# can be visualized via
plot(log.(gmres.residual[1:iters, :]) / log(10));
plot!(legend = false, xlims = (1, iters), ylims = (-15, 2));
plot!(ylabel = "log10 residual", xlabel = "iterations")


#=
Samuel J Palmer
Implementation is based on:
[1]: Akima, Hiroshi (1970). "A new method of interpolation and smooth curve fitting
    based on local procedures". Journal of the ACM. 17: 589–602.

Original module by Samuel J. Palmer, see https://github.com/sp94/CubicSplines.jl
Modified by Stefan Mathis to allow for customized extrapolation and to use a
binary search for interval identification.
=#

module CubicSplines

export CubicSpline, gradient

# As described in section 2.3 of 1
function extrapolate(xs, ys, extrap)

    # Get first / last three values
    x1, x2, x3 = xs
    y1, y2, y3 = ys

    # Extrapolate the x-coordinates of two additional datapoints with (8) from [1]
    x4 = x3-x1+x2
    x5 = x3-x1+x3

    # Quadratic extrapolation of (x4,y4), (x5,y5)
    if extrap === nothing
        y4 = y3 + 2*(x4-x3)*(y3-y2)/(x3-x2) - (x4-x3)*(y2-y1)/(x2-x1)
        y5 = y4 + 2*(x5-x4)*(y4-y3)/(x4-x3) - (x5-x4)*(y3-y2)/(x3-x2)
    else
        # Use the evalpoly macro to calculate the extrapolation values
        y4 = @evalpoly(x4-x3, y3, extrap...)
        y5 = @evalpoly(x5-x3, y3, extrap...)
    end
    return x4, x5, y4, y5
end

"""
Construct an Akima spline according to [1]. In addition to the original method,
this implementation also allows for customized extrapolation at both spline
sides. The extrapolation behaviour is defined by the coefficients of a 3rd degree
polynomial which can be given as additional keyword arguments extrapl and extrapr
to the constructor. extrapl and/or extrapr should be either "nothing" or a
set of coefficients.

Case 1: Extrapolation is set to "nothing"
================================================================================
In this case, the default quadratic extrapolation according to section 2.3 of [1]
is used to construct the spline and the spline throws an error when evaluated
outside its boundaries.

Case 2: Extrapolation is set to a vector
================================================================================
The container is given as a n-element collection of values:
    extrapl = [p1, p2, p3, ..., pn]
which are used to evaluate the following function:
    y = p0 + p1*(x-x3) + p2*(x-x3)^2 + p3*(x-x3)^3 + ... + pn*(x-x3)^n
The constant p0 is equal to y3 by definition (continuous spline). x3 and y3
are the first/last datapoint values respectively.
"""
struct CubicSpline{T <: Real}
    xs::Vector{T}
    ys::Vector{T}
    ts::Vector{T}
    ps::Matrix{T} # Coefficients of the polynom for each spline section
    logx::Bool
    extrapl::Union{Nothing,Vector{T}}
    extrapr::Union{Nothing,Vector{T}}
    function CubicSpline(xs, ys; logx=false, extrapl=nothing, extrapr=nothing)

        # Check input data
        @assert length(xs) == length(ys)
        @assert issorted(xs)
        if !issorted(xs)
            idx = sortperm(xs)
            xs = xs[idx]
            ys = ys[idx]
        end
        if logx
            xs = log.(xs)
        end
        if length(xs) < 5
            throw("At least five datapoints are necessary to construct a cubic Akima spline, however only $(length(xs)) datapoints were given.")
        end

        xb, xa, yb, ya = extrapolate(xs[3:-1:1], ys[3:-1:1], extrapl)
        xy, xz, yy, yz = extrapolate(xs[end-2:end], ys[end-2:end], extrapr)
        xs = [xa; xb; xs; xy; xz]
        ys = [ya; yb; ys; yy; yz]
        ms = zeros(eltype(ys), length(xs)-1)
        for i in 1:length(ms)
            ms[i] = (ys[i+1]-ys[i]) ./ (xs[i+1]-xs[i])
        end
        ts = zeros(eltype(ys), length(xs))
        for i in 3:length(xs)-2

            # Equals (1) in [1]
            m1, m2, m3, m4 = ms[i-2:i+1]
            # As described in [1] p.591 (parentheses block), this is an arbitrary
            # convention to guarantee uniqueness of the solution
            if m1 == m2 && m3 == m4
                ts[i] = (m2+m3)/2
            else
                numer = abs(m4-m3)*m2 + abs(m2-m1)*m3
                denom = abs(m4-m3)    + abs(m2-m1)
                ts[i] = numer ./ denom
            end
        end

        # Spline coefficients
        ps = zeros(eltype(ys), length(xs)-1, 4)
        for i in 3:length(xs)-2
            x1, x2 = xs[i:i+1]
            y1, y2 = ys[i:i+1]
            t1, t2 = ts[i:i+1]
            p0 = y1
            p1 = t1
            p2 = (3(y2-y1)/(x2-x1)-2t1-t2)/(x2-x1)
            p3 = (t1+t2-2(y2-y1)/(x2-x1))/(x2-x1)^2
            ps[i,:] = [p0, p1, p2, p3]
        end

        # Create extrapolation coefficients vector
        if extrapl !== nothing
            extrapl = vcat([ys[3]],[c for c in extrapl]) # 3rd value equals first value of the original dataset
        end
        if extrapr !== nothing
            extrapr = vcat([ys[end-2]],[c for c in extrapr]) # 3rd last value equals last value of the original dataset
        end
        return new{eltype(xs)}(xs[3:end-2], ys[3:end-2], ts[3:end-2], ps[3:end-2,:], logx, extrapl, extrapr)
    end
end
# Functor to make the CubicSpline class callable
function (obj::CubicSpline)(x)
    return obj[x]
end
# Non-allocating version for vectors
function (obj::CubicSpline)(y::AbstractVector,x::AbstractVector)
    for i in 1:length(x)
        y[i] = obj[x[i]]
    end
    return y
end

# As described in section 2.2 of [1]
function Base.getindex(spline::CubicSpline, x::Real)

    if spline.logx
        x = log(x)
    end
    if isnan(x)
        return NaN
    end

    # Value sits on upper interval border
    if x == spline.xs[end]
        idx = length(spline.xs)-1

    # Find the corresponding interval
    else

        # Find the interval index
        idx = _binary_search_interval(spline.xs, x)

        # Extrapolation to the "left"
        if idx == 0
            if spline.extrapl === nothing
                throw("x too small ($x < $(spline.xs[1]))")
            else
                return @inbounds @evalpoly(x-spline.xs[1], spline.extrapl...)
            end
        end

        # Extrapolation to the "right"
        if idx == length(spline.xs)
            if spline.extrapr === nothing
                throw("x too big ($x > $(spline.xs[end]))")
            else
                return @inbounds @evalpoly(x-spline.xs[end], spline.extrapr...)
            end
        end
    end

    # Evaluate the spline segment
    return @inbounds @evalpoly(x-spline.xs[idx], spline.ps[idx,:]...)
end

function Base.getindex(spline::CubicSpline, xs::AbstractArray{<:Real,1})
    return [spline[x] for x in xs]
end

function Base.firstindex(spline::CubicSpline)
    return spline.xs[1]
end

function Base.lastindex(spline::CubicSpline)
    return spline.xs[end]
end

"""
Evaluate the spline gradient of degree one at a given point x
"""
function gradient(spline::CubicSpline, x::Real)
    return gradient(spline, x, 1)
end

"""
Generalized gradient of degree n at a given point x
"""
function gradient(spline::CubicSpline, x::Real, n::Integer)

    # Value sits on upper interval border
    if x == spline.xs[end]
        idx = length(spline.xs)-1

    # Find the corresponding interval
    else

        # Find the interval index
        idx = _binary_search_interval(spline.xs, x)

        # Extrapolation to the "left"
        if idx == 0
            if spline.extrapl === nothing
                throw("x too small ($x < $(spline.xs[1]))")
            else
                coeff = spline.extrapl
                x_ref = spline.xs[1]
            end

        # Extrapolation to the "right"
        elseif idx == length(spline.xs)
            if spline.extrapr === nothing
                throw("x too big ($x > $(spline.xs[end]))")
            else
                coeff = spline.extrapr
                x_ref = spline.xs[end]
            end
        else
            coeff = spline.ps[idx,:]
            x_ref = spline.xs[idx]

        end
    end

    # Calculate the coefficients of degree n
    coeff = _diffpoly(coeff, n, 1)

    # If the coefficients are empty, the derivative is 0.
    if isempty(coeff)
        return zero(x)
    else
        return @evalpoly(x-x_ref, coeff...)
    end
end

"""
Binary search to find the interval which contains a given number. The interval
is given as a sorted array of real numbers. Each neighbouring pair is interpreted
as an interval. The function returns the lower index of the interval containing
the given number.

If the returned index is 0, the number is located in the interval (-Inf, array[1])
If the returned index is length(array), the number is located in the interval [array[end], Inf)
"""
@inline function _binary_search_interval(array::Vector{T}, number::Real) where T <: Real

    left = 1
    right = length(array)

    @inbounds while left <= right
        center = (left+right)÷2
        if array[center] > number
            right = center - 1
        else # array[center] <= number
            left = center + 1
        end
    end
    return left - 1
end

"""
Calculate the coefficients for the differentiated polygon
"""
function _diffpoly(coeff::AbstractArray{<:Real,1}, n::Integer=1, counter::Integer=1)

    # Discard the first coefficient, since it vanishes during the differentiation process
    coeff = coeff[2:end] .* collect(1:(length(coeff)-1))

    # If the differentiation degree is larger than the counter, differentiate the
    # differentiated coefficients again.
    if n == counter
        return coeff
    else
        return _diffpoly(coeff, n, counter+1)
    end
end

end # module


# NOTE BLAS.set_num_threads(1) may increase performance significantly

using LinearAlgebra

"""
Calculate a unitary tensor with size (chi1, chi2, chi1*chi2)
that approximately disentangles the input tensor 'A'.
'A' must have the size (chi1*chi2, chi3, chi4) where either
chi2 <= ceil(chi4 / ceil(chi1/chi3)) or chi1 <= ceil(chi3 / ceil(chi2/chi4)).
If these ratios are integers, then chi1*chi2 <= chi3*chi4 is sufficient.
chi1 <= chi3 and chi2 <= chi4 is also sufficient.

example: fastDisentangle(2, 3, randn(6,5,7))
"""
function fastDisentangle(chi1::Int, chi2::Int, A::AbstractArray{T,3},
                         transposeQ::Union{Bool,Nothing}=nothing) where T
    n,chi3,chi4 = size(A)
    if n != chi1*chi2
        throw(ArgumentError("fastDisentangle: The input array must have the shape (chi1*chi2, chi3, chi4)"))
    end

    # implementing Appendix B in https://arxiv.org/pdf/2104.08283
    if chi1 > chi3
        chi4to3 = divUp(chi1, chi3)
        chi4new = divUp(chi4, chi4to3)
        if chi2 > chi4new
            throw(ArgumentError("""
                fastDisentangle: The input array must have the shape (chi1*chi2, chi3, chi4) where either
                chi2 <= ceil(chi4 / ceil(chi1/chi3)) or
                chi1 <= ceil(chi3 / ceil(chi2/chi4))."""))
        end
        V1      = svd(flatten(A, 2), full=true).V # 1
        V       = permutedims(reshape([V1 zeros(chi4, chi4to3*chi4new - chi4)], (chi4, chi4new, chi4to3)), [1,3,2]) # 2
        Anew    = reshape(tensordot(A, V), (n, chi3*chi4to3, chi4new)) # 3
        return fastDisentangle(chi1, chi2, Anew, false)
    end

    if chi2 > chi4
        return permutedims(fastDisentangle(chi2, chi1, permutedims(A, [1,3,2])), [2,1,3])
    end

    # implementing Algorithm 1 in https://arxiv.org/pdf/2104.08283
    r = randn(T, n) # 1
    svd_A = svd(tensordot(r, A)) # 2
    alpha3, alpha4 = conj(svd_A.U[:,1]), conj(svd_A.Vt[1,:])
    V3  = svd(tensordot(A, alpha4)).Vt[1:chi1,:]' # 3
    V4  = svd(tensordot(permutedims(A,[1,3,2]), alpha3)).Vt[1:chi2,:]' # 4
    B   = permutedims(tensordot(permutedims(tensordot(A, V4), [1,3,2]), V3), [1,3,2]) # 5 TODO optimize order
    @assert size(B) == (n, chi1, chi2)
    if transposeQ == nothing
        transposeQ = chi1 > chi2
    end
    Bt  = transposeQ ? permutedims(B, [1,3,2]) : B
    Ut  = reshape(orthogonalize!(flatten(conj(Bt), 1)), size(Bt)) # 6
    permutedims(Ut, transposeQ ? [3,2,1] : [2,3,1])
end

"""Compute the entanglement entropy of UA = tensordot(U, A)"""
# defined in equation (10) of https://arxiv.org/pdf/2104.08283
function entanglement(UA::AbstractArray{T,4}) where T
    lambdas  = svdvals(flatten(permutedims(UA, [1,3,2,4]), 2))
    ps   = lambdas .* lambdas
    ps ./= sum(ps)
    max(0., -dot(ps, log.(max.(ps, floatmin(real(T))))))
end

divUp(x,y) = div(x,y,RoundUp)

"""Convert an array to a matrix (or vector if n==0 or ndims(A)) by merging the first n and last ndims(A)-n indices."""
function flatten(A::AbstractArray{T,N}, n) where {T,N}
    n==0 || n==ndims(A) ? vec(A) : reshape(A, (prod(size(A)[1:n]), prod(size(A)[n+1:end])))
end

"""
Contract the last n indices of A with the first n indices of B.
This is similar to Python's numpy.tensordot, except n defaults to 1.
"""
function tensordot(A::AbstractArray{T,M}, B::AbstractArray{T,N}, n::Int=1) where {T,M,N}
    if n == ndims(A) # needed since 'vector * matrix' returns an error
        return dropdims(tensordot(reshape(A, (1, size(A)...)), B, n); dims=1)
    elseif n == 0
        return tensordot(reshape(A, (size(A)..., 1)), reshape(B, (1, size(B)...)), 1)
    end
    reshape(flatten(A, ndims(A)-n) * flatten(B, n), (size(A)[1:ndims(A)-n]..., size(B)[n+1:end]...))
end

"""
Gram-Schmidt orthonormalization of the rows of M.
Inserts random vectors in the case of linearly dependent rows.
"""
function orthogonalize!(M::Matrix{T}, reorthogonalizeQ::Bool=true) where T
    epsMin = sqrt(eps(real(T))) # once eps0<epsMin, we add random vectors if needed
    eps0 = 0.5sqrt(epsMin) # only accept new orthogonal vectors if their relative norm is at last eps0 after orthogonalization
    n,m = size(M)
    @assert n >= m
    norms = [norm(M[:,i]) for i in 1:m]
    maxNorm = maximum(norms)
    orthogQ = zeros(Bool, m)
    allOrthog = false
    while !allOrthog
        allOrthog = true
        eps1 = max(eps0, epsMin)
        for i in 1:m
            if !orthogQ[i]
                if norms[i] > eps1 * maxNorm
                    Mi = M[:,i]
                    for j in 1:m
                        if orthogQ[j]
                            Mi = Mi - M[:,j] * dot(M[:,j], Mi)
                        end
                    end
                    normMi = norm(Mi)
                    if normMi > eps1 * norms[i]
                        M[:,i] = Mi / normMi
                        orthogQ[i] = true
                        continue
                    end
                end
                # M[:,i] was a linear combination of the other vectors
                if eps0 < epsMin
                    M[:,i]   = maxNorm * normalize(randn(T, n))
                    norms[i] = maxNorm
                end
                allOrthog = false
            end
        end
        eps0 = eps0*eps0
    end
    if reorthogonalizeQ && norm(M' * M - I) > eps(real(T)) ^ 0.75
        return orthogonalize!(M, false)
    end
    M
end

# verification code

"""Check that the ansatz in equation (1) of https://arxiv.org/pdf/2104.08283 results in the minimal entanglement entropy."""
function checkAnsatz(chi1::Int, chi2::Int, chi3a::Int, chi4b::Int, chi3c::Int, chi4c::Int,
                     eps0::T=ComplexF64(0)) where T
    M1,M2,M3 = randn(T,chi1,chi3a), randn(T,chi2,chi4b), randn(T,chi3c,chi4c)
    A = reshape(permutedims(tensordot(tensordot(M1, M2, 0), M3, 0), (1,3,2,5,4,6)),
                (chi1*chi2, chi3a*chi3c, chi4b*chi4c))
    A = A + eps0 * norm(A) * normalize(randn(T,size(A)...))
    U = fastDisentangle(chi1,chi2,A)
    return entanglement(tensordot(U,A)) - entanglement(reshape(M3,(1,1,chi3c,chi4c)))
end

"""repeatedly check the ansatz"""
function checkAnsatzRepeated(maxChi::Int=5)
    c = 0
    BLAS.set_num_threads(1) # this seems to increase performance significantly
    while true
        c += 1
        if c%1000 == 0
            println(c)
        end
        chis = rand(1:maxChi, 6)
        chi1,chi2,chi3a,chi4b,chi3c,chi4c = chis
        chi3 = chi3a*chi3c
        chi4 = chi4b*chi4c
        eps0 = 10^(-6 - 14*rand(Float64))
        complexQ = rand(Bool)
        if (chi1 <= chi3 && chi2 <= chi4) ||
           ((chi3c==1 || chi4c==1) && (chi2 <= divUp(chi4, divUp(chi1,chi3)) ||
                                       chi1 <= divUp(chi3, divUp(chi2,chi4))))
            args = (chis..., complexQ ? ComplexF64(eps0) : eps0)
            S = checkAnsatz(args...)
            if S > 0.1sqrt(max(eps0, eps(Float64)))
                println(args, " -> ", S)
                return
            end
        end
    end
end


get_input() = open("./inputs/day15.txt") do f
	readlines(f)[1] |> l -> split(l, ",") .|> string .|> String |> l -> parse.(Int, l)
end

function run_game(upto)
	nums = get_input()

	prev = pop!(nums)

	memo = Dict((num, i) for (i, num) in enumerate(nums))

	turn = length(nums) + 1

	while turn != upto
		if haskey(memo, prev)
			memo[prev], prev = turn, turn - memo[prev]
		else
			memo[prev] = turn
			prev = 0
		end

		turn += 1
	end

	return prev
end

part1() = run_game(2020)
part2() = run_game(30000000)

println("day 15:\n\tpart 1: ", part1(), "\n\tpart 2: ", part2())


using Smolyak
import PyPlot

# function to plot a 2d or 3d grid
function plot(sg::SmolyakGrid)
    g = sg.grid
    mu  = sg.mu
    d = sg.d
    mu_str = isa(mu, Array) ? replace(strip(string(mu)),'\n'," \\times ") : mu
    if size(g, 2) == 2
        PyPlot.scatter(g[:, 1], g[:, 2])
        PyPlot.title("Smolyak grid: \$d=$(d), \\; \\mu=$(mu_str)\$")
    elseif size(g, 2) == 3
        PyPlot.scatter3D(g[:, 1], g[:, 2], g[:, 3], "b.")
        PyPlot.title("Smolyak grid: \$d=$(d), \\; \\mu=$(mu_str)\$")
    else
        error("ERROR: can only plot 2d or 3d grids")
    end
end


using NumericalContinuation
using NumericalContinuationShootingProblems
using StaticArrays
using Test

NC = NumericalContinuation

@testset "Shooting problems" begin
    hopf = (out, u, p, t) -> (out[1] = p[1]*u[1] - u[2] + p[2]*u[1]*(u[1]^2 + u[2]^2);
                              out[2] = u[1] + p[1]*u[2] + p[2]*u[2]*(u[1]^2 + u[2]^2))
    prob = NC.ProblemStructure()
    @test_throws ArgumentError add_shootingproblem!(prob, "hopf", hopf, [1.0, 0.0], [1.0, -1.0], [0, 2π], abstol=1e-8, pnames=["α", "β", "γ"])
    add_shootingproblem!(prob, "hopf", hopf, [1.0, 0.0], [1.0, -1.0], [0, 2π], abstol=1e-8)
    NC.initialize!(prob)
    out = ones(6)
    u0 = NC.get_u0(Float64, NC.get_vars(prob))
    d0 = NC.get_data(NC.get_data(prob))
    get_funcs(prob)[:embedded](out, u0, data=d0, prob=prob)
    @test isapprox(out, zero(out), atol=1e-5)

    hopf2 = (u, p, t) -> SVector(p[1]*u[1] - u[2] + p[2]*u[1]*(u[1]^2 + u[2]^2),
                                 u[1] + p[1]*u[2] + p[2]*u[2]*(u[1]^2 + u[2]^2))
    prob = NC.ProblemStructure()
    add_shootingproblem!(prob, "hopf", hopf2, SVector(1.0, 0.0), SVector(1.0, -1.0), [0, 2π], abstol=1e-8)
    NC.initialize!(prob)
    out = ones(6)
    u0 = NC.get_u0(Float64, NC.get_vars(prob))
    d0 = NC.get_data(NC.get_data(prob))
    get_funcs(prob)[:embedded](out, u0, data=d0, prob=prob)
    @test isapprox(out, zero(out), atol=1e-5)
end


module FastUnitfulOperations

using LinearAlgebra
using Unitful

export *ᵘ, uproduct, uproduct!, usort, usort!, usortperm, usortperm!

"""
    uproduct(a::AbstractArray{T1}, b::AbstractArray{T2}) where {T1 and/or T2 <:Quantity}
    *ᵘ(a::AbstractArray{T1}, b::AbstractArray{T2}) where {T1 and/or T2 <:Quantity}

Fast matrix multiplication for Unitful arrays.

This method reinterprets the unitful arrays as arrays with the number's type to make use of
optimized matrix multiplication for integers or floats.
The method will throw an error if the array contains Unitful elements that do not have a fixed
value or unit type.

Type `*\\^u` followed by the tab key to input the operator `*ᵘ`.
"""
function uproduct(a::AbstractArray{T1}, b::AbstractArray{T2}) where {T1<:Quantity,T2<:Quantity}
    c = ustrip(a) * ustrip(b)
    u = unit(T1) * unit(T2)
    return reinterpret(Quantity{eltype(c),dimension(u),typeof(u)}, c)
end

function uproduct(a::AbstractArray{T1}, b::AbstractArray) where {T1<:Quantity}
    c = ustrip(a) * b
    return reinterpret(T1, c)
end

function uproduct(a::AbstractArray, b::AbstractArray{T2}) where {T2<:Quantity}
    c = a * ustrip(b)
    return reinterpret(T2, c)
end

"""
    uproduct!(c::AbstractArray, a::AbstractArray{T1}, b::AbstractArray{T2}) where {T1 and/or T2 <:Quantity}

Fast in-place matrix multiplication for Unitful arrays.

This method reinterprets the unitful arrays as arrays with the number's type to make use of
optimized matrix multiplication for integers or floats.
The method will throw an error if the array contains Unitful elements that do not have a fixed
value or unit type.
"""
function uproduct!(c::AbstractArray{T1}, a::AbstractArray{T2}, b::AbstractArray{T3}) where {T1<:Quantity,T2<:Quantity,T3<:Quantity}
    mul!(ustrip(c), ustrip(a), ustrip(b))
    return c
end

function uproduct!(c::AbstractArray, a::AbstractArray{T2}, b::AbstractArray{T3}) where {T2<:Quantity,T3<:Quantity}
    mul!(c, ustrip(a), ustrip(b))
    return c
end

function uproduct!(c::AbstractArray{T1}, a::AbstractArray, b::AbstractArray{T3}) where {T1<:Quantity,T3<:Quantity}
    mul!(ustrip(c), a, ustrip(b))
    return c
end

function uproduct!(c::AbstractArray{T1}, a::AbstractArray{T2}, b::AbstractArray) where {T1<:Quantity,T2<:Quantity}
    mul!(ustrip(c), ustrip(a), b)
    return c
end

"""
    uproduct(a, b)
    *ᵘ(a, b)

Fallback to regular multiplication: `uproduct(a, b) = a * b`
"""
uproduct(a, b) = a * b

"""
    uproduct!(c, a, b)

Fallback to regular multiplication: `uproduct!(c, a, b) = mul!(c, a, b)`
"""
uproduct!(c, a, b) = mul!(c, a, b)

*ᵘ = uproduct

"""
    usort(a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    usort(a; kwargs...)

Sort a unitful array as a reinterpreted array according to the number's type.

Calls [`Base.sort`](https://docs.julialang.org/en/v1/base/sort/#Base.sort) on the reinterpreted
array. This enables the use of sorting algorithms that only work for integers or floats
(e.g. RadixSort from
[SortingAlgorithms](https://github.com/JuliaCollections/SortingAlgorithms.jl).

The method falls back `Base.sort` for other array types.
"""
function usort(a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    a_sorted = sort(ustrip(a); kwargs...)
    return reinterpret(T, a_sorted)
end

"""
    usort!(a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    usort!(a; kwargs...)

Sort a unitful array in-place as a reinterpreted array according to the number's type.

Calls [`Base.sort!`](https://docs.julialang.org/en/v1/base/sort/#Base.sort!) on the reinterpreted
array. This enables the use of sorting algorithms that only work for integers or floats
(e.g. RadixSort from
[SortingAlgorithms](https://github.com/JuliaCollections/SortingAlgorithms.jl).

The method falls back `Base.sort!` for other array types.
"""
function usort!(a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    sort!(ustrip(a); kwargs...)
    return a
end

"""
    usortperm(a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    usortperm(a; kwargs...)

Return the permutation vector to sort a unitful array as a reinterpreted array according to the
number's type.

Calls [`Base.sortperm`](https://docs.julialang.org/en/v1/base/sort/#Base.sortperm) on the reinterpreted
array. This enables the use of sorting algorithms that only work for integers or floats
(e.g. RadixSort from
[SortingAlgorithms](https://github.com/JuliaCollections/SortingAlgorithms.jl).

The method falls back `Base.sortperm` for other array types.
"""
function usortperm(a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    return sortperm(ustrip(a); kwargs...)
end

"""
    usortperm!(ix, a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    usortperm!(ix, a; kwargs...)

Return the permutation vector in the preallocated vector `ix` to sort a unitful array as a
reinterpreted array according to the number's type.

Calls [`Base.sortperm!`](https://docs.julialang.org/en/v1/base/sort/#Base.sortperm!) on the reinterpreted
array. This enables the use of sorting algorithms that only work for integers or floats
(e.g. RadixSort from
[SortingAlgorithms](https://github.com/JuliaCollections/SortingAlgorithms.jl).

The method falls back `Base.sortperm!` for other array types.
"""
function usortperm!(ix, a::AbstractArray{T}; kwargs...) where {T<:Quantity}
    return sortperm!(ix, ustrip(a); kwargs...)
end

usort(a; kwargs...) = sort(a; kwargs...)
usort!(a; kwargs...) = sort!(a; kwargs...)
usortperm(a; kwargs...) = sortperm(a; kwargs...)
usortperm!(ix, a; kwargs...) = sortperm!(ix, a; kwargs...)

end


function mass_matrix_machine_entries!(
    mass_matrix,
    machine::M,
    global_index::Base.ImmutableDict{Symbol, Int64},
) where {M <: PSY.Machine}
    @debug "Using default mass matrix entries $M"
end

"""
Model of 0-state synchronous machine in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""
function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{PSY.BaseMachine, S, A, TG, P}},
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.BaseMachine)
    δ = device_states[external_ix[1]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    Xd_p = PSY.get_Xd_p(machine)
    eq_p = PSY.get_eq_p(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric variables
    i_d = (1.0 / (R^2 + Xd_p^2)) * (Xd_p * (eq_p - V_dq[2]) - R * V_dq[1])  #15.36
    i_q = (1.0 / (R^2 + Xd_p^2)) * (Xd_p * V_dq[1] + R * (eq_p - V_dq[2])) #15.36
    Pe = (V_dq[1] + R * i_d) * i_d + (V_dq[2] + R * i_q) * i_q      #15.35

    #Update inner_vars
    inner_vars[τe_var] = Pe #Model assume ω approx 1.0

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [i_d; i_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

"""
Model of 2-state (One d- and One q-) synchronous machine in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""
function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{PSY.OneDOneQMachine, S, A, TG, P}},
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.OneDOneQMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    eq_p = internal_states[1]
    ed_p = internal_states[2]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.OneDOneQMachine)
    δ = device_states[external_ix[1]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xq_p = PSY.get_Xq_p(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Tq0_p = PSY.get_Tq0_p(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric variables
    i_d = (1.0 / (R^2 + Xd_p * Xq_p)) * (Xq_p * (eq_p - V_dq[2]) + R * (ed_p - V_dq[1]))  #15.32
    i_q = (1.0 / (R^2 + Xd_p * Xq_p)) * (-Xd_p * (ed_p - V_dq[1]) + R * (eq_p - V_dq[2]))  #15.32
    Pe = (V_dq[1] + R * i_d) * i_d + (V_dq[2] + R * i_q) * i_q      #15.35

    #Compute ODEs
    output_ode[local_ix[1]] = (1.0 / Td0_p) * (-eq_p - (Xd - Xd_p) * i_d + Vf)     #15.29 eq_p
    output_ode[local_ix[2]] = (1.0 / Tq0_p) * (-ed_p + (Xq - Xq_p) * i_q)          #15.30 ed_p

    #Update inner_vars
    inner_vars[τe_var] = Pe #Model assume ω approx 1.0

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [i_d; i_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

"""
Model of 6-state (MarconatoMachine) synchronous machine in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""
function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{PSY.MarconatoMachine, S, A, TG, P}},
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.MarconatoMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    ψq = internal_states[1]
    ψd = internal_states[2]
    eq_p = internal_states[3]
    ed_p = internal_states[4]
    eq_pp = internal_states[5]
    ed_pp = internal_states[6]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.MarconatoMachine)
    δ = device_states[external_ix[1]]
    ω = device_states[external_ix[2]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xq_p = PSY.get_Xq_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = PSY.get_Xq_pp(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Tq0_p = PSY.get_Tq0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    T_AA = PSY.get_T_AA(machine)
    γd = PSY.get_γd(machine)
    γq = PSY.get_γq(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric variables
    i_d = (1.0 / Xd_pp) * (eq_pp - ψd)      #15.18
    i_q = (1.0 / Xq_pp) * (-ed_pp - ψq)     #15.18
    τ_e = ψd * i_q - ψq * i_d               #15.6

    #Compute ODEs
    output_ode[local_ix[1]] = 2 * π * f0 * (R * i_q - ω * ψd + V_dq[2])                        #15.9 ψq
    output_ode[local_ix[2]] = 2 * π * f0 * (R * i_d + ω * ψq + V_dq[1])                        #15.9 ψd
    output_ode[local_ix[3]] =
        ((1.0 / Td0_p) * (-eq_p - (Xd - Xd_p - γd) * i_d + (1 - (T_AA / Td0_p)) * Vf))                             #15.16 eq_p
    output_ode[local_ix[4]] = (1.0 / Tq0_p) * (-ed_p + (Xq - Xq_p - γq) * i_q)             #15.16 ed_p
    output_ode[local_ix[5]] =
        ((1.0 / Td0_pp) * (-eq_pp + eq_p - (Xd_p - Xd_pp + γd) * i_d + (T_AA / Td0_p) * Vf))       #15.16 eq_pp
    output_ode[local_ix[6]] = (1.0 / Tq0_pp) * (-ed_pp + ed_p + (Xq_p - Xq_pp + γq) * i_q) #15.16 ed_pp

    #Update inner_vars
    inner_vars[τe_var] = τ_e

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [i_d; i_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

"""
Model of 4-state (SimpleMarconatoMachine) synchronous machine in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""
function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{
        PSY.DynamicGenerator{PSY.SimpleMarconatoMachine, S, A, TG, P},
    },
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)
    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.SimpleMarconatoMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    eq_p = internal_states[1]
    ed_p = internal_states[2]
    eq_pp = internal_states[3]
    ed_pp = internal_states[4]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.SimpleMarconatoMachine)
    δ = device_states[external_ix[1]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xq_p = PSY.get_Xq_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = PSY.get_Xq_pp(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Tq0_p = PSY.get_Tq0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    T_AA = PSY.get_T_AA(machine)
    γd = PSY.get_γd(machine)
    γq = PSY.get_γq(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric variables
    #i_dq = inv([[R -Xq_pp]; [Xq_pp R]])*[ed_pp - V_dq[1]; eq_pp - V_dq[2]]
    i_d =
        (1.0 / (R^2 + Xd_pp * Xq_pp)) * (Xq_pp * (eq_pp - V_dq[2]) + R * (ed_pp - V_dq[1]))      #15.25
    i_q =
        (1.0 / (R^2 + Xd_pp * Xq_pp)) * (-Xd_pp * (ed_pp - V_dq[1]) + R * (eq_pp - V_dq[2]))      #15.25
    τ_e = (V_dq[1] + R * i_d) * i_d + (V_dq[2] + R * i_q) * i_q               #15.6

    #Compute ODEs
    output_ode[local_ix[1]] =
        ((1.0 / Td0_p) * (-eq_p - (Xd - Xd_p - γd) * i_d + (1 - (T_AA / Td0_p)) * Vf))                             #15.16 eq_p
    output_ode[local_ix[2]] = (1.0 / Tq0_p) * (-ed_p + (Xq - Xq_p - γq) * i_q)             #15.16 ed_p
    output_ode[local_ix[3]] =
        ((1.0 / Td0_pp) * (-eq_pp + eq_p - (Xd_p - Xd_pp + γd) * i_d + (T_AA / Td0_p) * Vf))       #15.16 eq_pp
    output_ode[local_ix[4]] = (1.0 / Tq0_pp) * (-ed_pp + ed_p + (Xq_p - Xq_pp + γq) * i_q) #15.16 ed_pp

    #Update inner_vars
    inner_vars[τe_var] = τ_e

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [i_d; i_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

"""
Model of 6-state (AndersonFouadMachine) synchronous machine in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""
function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{
        PSY.DynamicGenerator{PSY.AndersonFouadMachine, S, A, TG, P},
    },
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.AndersonFouadMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    ψq = internal_states[1]
    ψd = internal_states[2]
    eq_p = internal_states[3]
    ed_p = internal_states[4]
    eq_pp = internal_states[5]
    ed_pp = internal_states[6]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.AndersonFouadMachine)
    δ = device_states[external_ix[1]]
    ω = device_states[external_ix[2]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xq_p = PSY.get_Xq_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = PSY.get_Xq_pp(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Tq0_p = PSY.get_Tq0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric variables
    i_d = (1.0 / Xd_pp) * (eq_pp - ψd)      #15.18
    i_q = (1.0 / Xq_pp) * (-ed_pp - ψq)     #15.18
    τ_e = ψd * i_q - ψq * i_d               #15.6

    #Compute ODEs
    output_ode[local_ix[1]] = 2 * π * f0 * (R * i_q - ω * ψd + V_dq[2])                        #15.9 ψq
    output_ode[local_ix[2]] = 2 * π * f0 * (R * i_d + ω * ψq + V_dq[1])                        #15.9 ψd
    output_ode[local_ix[3]] = (1.0 / Td0_p) * (-eq_p + (Xd - Xd_p) * i_d + Vf)             #15.19 eq_p
    output_ode[local_ix[4]] = (1.0 / Tq0_p) * (-ed_p + (Xq - Xq_p) * i_q)                  #15.19 ed_p
    output_ode[local_ix[5]] = (1.0 / Td0_pp) * (-eq_pp + eq_p - (Xd_p - Xd_pp) * i_d)      #15.19 eq_pp
    output_ode[local_ix[6]] = (1.0 / Tq0_pp) * (-ed_pp + ed_p + (Xq_p - Xq_pp) * i_q)      #15.19 ed_pp

    #Update inner_vars
    inner_vars[τe_var] = τ_e

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [i_d; i_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

"""
Model of 4-state (SimpleAFMachine) synchronous machine in Julia.
Refer to Power System Modelling and Scripting by F. Milano for the equations
"""
function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{PSY.SimpleAFMachine, S, A, TG, P}},
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.SimpleAFMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    eq_p = internal_states[1]
    ed_p = internal_states[2]
    eq_pp = internal_states[3]
    ed_pp = internal_states[4]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.SimpleAFMachine)
    δ = device_states[external_ix[1]]
    ω = device_states[external_ix[2]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xq_p = PSY.get_Xq_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = PSY.get_Xq_pp(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Tq0_p = PSY.get_Tq0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric variables
    i_d =
        (1.0 / (R^2 + Xd_pp * Xq_pp)) * (Xq_pp * (eq_pp - V_dq[2]) + R * (ed_pp - V_dq[1]))      #15.25
    i_q =
        (1.0 / (R^2 + Xd_pp * Xq_pp)) * (-Xd_pp * (ed_pp - V_dq[1]) + R * (eq_pp - V_dq[2]))      #15.25
    τ_e = (V_dq[1] + R * i_d) * i_d + (V_dq[2] + R * i_q) * i_q               #15.6

    #Compute ODEs
    output_ode[local_ix[1]] = (1.0 / Td0_p) * (-eq_p + (Xd - Xd_p) * i_d + Vf)            #15.19 eq_p
    output_ode[local_ix[2]] = (1.0 / Tq0_p) * (-ed_p + (Xq - Xq_p) * i_q)                 #15.19 ed_p
    output_ode[local_ix[3]] = (1.0 / Td0_pp) * (-eq_pp + eq_p - (Xd_p - Xd_pp) * i_d)     #15.19 eq_pp
    output_ode[local_ix[4]] = (1.0 / Tq0_pp) * (-ed_pp + ed_p + (Xq_p - Xq_pp) * i_q)     #15.19 ed_pp

    #Update inner_vars
    inner_vars[τe_var] = τ_e

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [i_d; i_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{M, S, A, TG, P}},
) where {
    M <: Union{PSY.RoundRotorQuadratic, PSY.RoundRotorExponential},
    S <: PSY.Shaft,
    A <: PSY.AVR,
    TG <: PSY.TurbineGov,
    P <: PSY.PSS,
}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)

    #Obtain indices for component w/r to device
    machine = PSY.get_machine(dynamic_device)
    local_ix = get_local_state_ix(dynamic_device, M)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    eq_p = internal_states[1]
    ed_p = internal_states[2]
    ψ_kd = internal_states[3]
    ψ_kq = internal_states[4]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, M)
    δ = device_states[external_ix[1]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var] #E_fd: Field voltage

    #Get parameters
    R = PSY.get_R(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_p = PSY.get_Tq0_p(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xq_p = PSY.get_Xq_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = Xd_pp
    Xl = PSY.get_Xl(machine)
    γ_d1 = PSY.get_γ_d1(machine)
    γ_q1 = PSY.get_γ_q1(machine)
    γ_d2 = PSY.get_γ_d2(machine)
    γ_q2 = PSY.get_γ_q2(machine)
    γ_qd = PSY.get_γ_qd(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Additional Fluxes
    ψq_pp = γ_q1 * ed_p + ψ_kq * (1 - γ_q1)
    ψd_pp = γ_d1 * eq_p + γ_d2 * (Xd_p - Xl) * ψ_kd
    ψ_pp = sqrt(ψd_pp^2 + ψq_pp^2)
    #Currents
    I_d =
        (1.0 / (R^2 + Xq_pp * Xd_pp)) *
        (-R * (V_dq[1] - ψq_pp) + Xq_pp * (-V_dq[2] + ψd_pp))
    I_q =
        (1.0 / (R^2 + Xq_pp * Xd_pp)) * (Xd_pp * (V_dq[1] - ψq_pp) + R * (-V_dq[2] + ψd_pp))
    Se = saturation_function(machine, ψ_pp)
    Xad_Ifd = eq_p + (Xd - Xd_p) * (γ_d1 * I_d - γ_d2 * ψ_kd + γ_d2 * eq_p) + Se * ψd_pp
    Xaq_I1q =
        ed_p + (Xq - Xq_p) * (γ_q2 * ed_p - γ_q2 * ψ_kq - γ_q1 * I_q) + Se * ψq_pp * γ_qd
    τ_e = I_d * (V_dq[1] + I_d * R) + I_q * (V_dq[2] + I_q * R)

    #Compute ODEs
    output_ode[local_ix[1]] = (1.0 / Td0_p) * (Vf - Xad_Ifd)                        #2.20b eq_p
    output_ode[local_ix[2]] = (1.0 / Tq0_p) * (-Xaq_I1q)                            #2.21b ed_p
    output_ode[local_ix[3]] = (1.0 / Td0_pp) * (-ψ_kd + eq_p - (Xd_p - Xl) * I_d)   #2.20c ψ_kd
    output_ode[local_ix[4]] = (1.0 / Tq0_pp) * (-ψ_kq + ed_p + (Xq_p - Xl) * I_q)   #2.21c ψ_kq

    #Update inner_vars
    inner_vars[τe_var] = τ_e
    inner_vars[Xad_Ifd_var] = Xad_Ifd

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [I_d; I_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{
        PSY.DynamicGenerator{PSY.SalientPoleQuadratic, S, A, TG, P},
    },
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)

    #Obtain indices for component w/r to device
    machine = PSY.get_machine(dynamic_device)
    local_ix = get_local_state_ix(dynamic_device, typeof(machine))

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    eq_p = internal_states[1]
    ψ_kd = internal_states[2]
    ψq_pp = internal_states[3]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, typeof(machine))
    δ = device_states[external_ix[1]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var] #E_fd: Field voltage

    #Get parameters
    R = PSY.get_R(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = Xd_pp
    Xl = PSY.get_Xl(machine)
    γ_d1 = PSY.get_γ_d1(machine)
    γ_q1 = PSY.get_γ_q1(machine)
    γ_d2 = PSY.get_γ_d2(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_d, V_q = ri_dq(δ) * [V_tR; V_tI]

    #Additional Fluxes
    ψd_pp = γ_d1 * eq_p + γ_q1 * ψ_kd

    #Currents
    I_d = (1.0 / (R^2 + Xd_pp^2)) * (-R * (V_d + ψq_pp) + Xd_pp * (ψd_pp - V_q))
    I_q = (1.0 / (R^2 + Xd_pp^2)) * (Xd_pp * (V_d + ψq_pp) + R * (ψd_pp - V_q))
    Se = saturation_function(machine, eq_p)
    Xad_Ifd =
        eq_p + Se * eq_p + (Xd - Xd_p) * (I_d + γ_d2 * (eq_p - ψ_kd - (Xd_p - Xl) * I_d))
    τ_e = I_d * (V_d + I_d * R) + I_q * (V_q + I_q * R)

    #Compute ODEs
    output_ode[local_ix[1]] = (1.0 / Td0_p) * (Vf - Xad_Ifd)                        #2.34b eq_p
    output_ode[local_ix[2]] = (1.0 / Td0_pp) * (-ψ_kd + eq_p - (Xd_p - Xl) * I_d)   #2.34c ψ_kd
    output_ode[local_ix[3]] = (1.0 / Tq0_pp) * (-ψq_pp - (Xq - Xq_pp) * I_q)        #2.35b ψq_pp

    #Update inner_vars
    inner_vars[τe_var] = τ_e
    inner_vars[Xad_Ifd_var] = Xad_Ifd

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [I_d; I_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

function mdl_machine_ode!(
    device_states::AbstractArray{<:ACCEPTED_REAL_TYPES},
    output_ode::AbstractArray{<:ACCEPTED_REAL_TYPES},
    inner_vars::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{
        PSY.DynamicGenerator{PSY.SalientPoleExponential, S, A, TG, P},
    },
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}
    Sbase = get_system_base_power(dynamic_device)
    f0 = get_system_base_frequency(dynamic_device)

    #Obtain indices for component w/r to device
    machine = PSY.get_machine(dynamic_device)
    local_ix = get_local_state_ix(dynamic_device, typeof(machine))

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    eq_p = internal_states[1]
    ψ_kd = internal_states[2]
    ψq_pp = internal_states[3]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, typeof(machine))
    δ = device_states[external_ix[1]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var] #E_fd: Field voltage

    #Get parameters
    R = PSY.get_R(machine)
    Td0_p = PSY.get_Td0_p(machine)
    Td0_pp = PSY.get_Td0_pp(machine)
    Tq0_pp = PSY.get_Tq0_pp(machine)
    Xd = PSY.get_Xd(machine)
    Xq = PSY.get_Xq(machine)
    Xd_p = PSY.get_Xd_p(machine)
    Xd_pp = PSY.get_Xd_pp(machine)
    Xq_pp = Xd_pp
    Xl = PSY.get_Xl(machine)
    γ_d1 = PSY.get_γ_d1(machine)
    γ_q1 = PSY.get_γ_q1(machine)
    γ_d2 = PSY.get_γ_d2(machine)
    γ_qd = (Xq - Xl) / (Xd - Xl)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_d, V_q = ri_dq(δ) * [V_tR; V_tI]

    #Additional Fluxes
    ψd_pp = γ_d1 * eq_p + γ_q1 * ψ_kd
    ψ_pp = sqrt(ψd_pp^2 + ψq_pp^2)

    #Currents
    I_d = (1.0 / (R^2 + Xd_pp^2)) * (-R * (V_d - ψq_pp) + Xq_pp * (-V_q + ψd_pp))
    I_q = (1.0 / (R^2 + Xd_pp^2)) * (Xd_pp * (V_d - ψq_pp) + R * (-V_q + ψd_pp))
    Se = saturation_function(machine, ψ_pp)
    Xad_Ifd =
        eq_p + Se * ψd_pp + (Xd - Xd_p) * (I_d + γ_d2 * (eq_p - ψ_kd - (Xd_p - Xl) * I_d))
    τ_e = I_d * (V_d + I_d * R) + I_q * (V_q + I_q * R)

    #Compute ODEs
    output_ode[local_ix[1]] = (1.0 / Td0_p) * (Vf - Xad_Ifd)                        #2.34b eq_p
    output_ode[local_ix[2]] = (1.0 / Td0_pp) * (-ψ_kd + eq_p - (Xd_p - Xl) * I_d)   #2.34c ψ_kd
    output_ode[local_ix[3]] =
        (1.0 / Tq0_pp) * (-ψq_pp + (Xq - Xq_pp) * I_q - Se * γ_qd * ψq_pp)        #2.35b ψq_pp

    #Update inner_vars
    inner_vars[τe_var] = τ_e
    inner_vars[Xad_Ifd_var] = Xad_Ifd

    #Compute current from the generator to the grid
    I_RI = (basepower / Sbase) * dq_ri(δ) * [I_d; I_q]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

#Not already implemented:
#=
"""
Model of 5-state (FullOrderMachine) synchronous machine in Julia.
Refer to Power System Dynamics: Stability and Control, by J. Machowski, J. Bialek and J. Bumby,
or Power System Stability and Control by P. Kundur, for the equations
"""

function mdl_machine_ode!(
    device_states,
    output_ode,
inner_vars,
current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{PSY.FullMachine, S, A, TG, P}},
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.FullMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    ψd = internal_states[1]
    ψq = internal_states[2]
    ψf = internal_states[3]
    ψ1d = internal_states[4]
    ψ1q = internal_states[5]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.FullMachine)
    δ = device_states[external_ix[1]]
    ω = device_states[external_ix[2]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]
    #ψd = inner_vars[ψd_var]
    #ψq = inner_vars[ψq_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    R_f = PSY.get_R_f(machine)
    R_1d = PSY.get_R_1d(machine)
    R_1q = PSY.get_R_1q(machine)
    inv_d_fluxlink = PSY.get_inv_d_fluxlink(machine)
    inv_q_fluxlink = PSY.get_inv_q_fluxlink(machine)
    basepower = PSY.get_base_power(dynamic_device)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric currents] using Flux Linkage equations
    i_df1d = inv_d_fluxlink * [ψd; ψf; ψ1d]  #11.18 in Machowski (3.127, 3.130 and 3.131 in Kundur)
    i_q1q = inv_q_fluxlink * [ψq; ψ1q]        #11.19 in Machowski (3.128 and 3.132 in Kundur)

    #Compute ODEs
    output_ode[local_ix[1]] = 2 * pi * f0 * (ψq * ω + R * i_df1d[1] + V_dq[1])  #11.32 in Machowski or 3.120 in Kundur
    output_ode[local_ix[2]] = 2 * pi * f0 * (-ψd * ω + R * i_q1q[1] + V_dq[2])
    output_ode[local_ix[3]] = -R_f * i_df1d[2] + Vf      #11.33 in Machowski or 3.123 in Kundur
    output_ode[local_ix[4]] = -R_1d * i_df1d[3]           #11.33 in Machowski or 3.124 in Kundur
    output_ode[local_ix[5]] = -R_1q * i_q1q[2]            #11.33 in Machowski or 3.125 in Kundur

    #Update stator fluxes and torque
    #ψd = (1.0/ω)*(R*i_q1q[1] + V_dq[2])         #11.32 in Machowski or 3.121 in Kundur
    #ψq = (1.0/ω)*(-R*i_df1d[1] - V_dq[1])       #11.32 in Machowski or 3.120 in Kundur
    τ_e = ψd * i_q1q[1] - ψq * i_df1d[1]            #15.6 in Milano or 3.117 in Kundur

    #Update inner_vars
    inner_vars[τe_var] = τ_e
    inner_vars[ψd_var] = ψd
    inner_vars[ψq_var] = ψq

    #Compute current from the generator to the grid
    I_RI = basepower * dq_ri(δ) * [i_df1d[1]; i_q1q[1]]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end

"""
Model of 3-state (FullOrderMachine) synchronous machine in Julia.
Refer to Power System Dynamics: Stability and Control, by J. Machowski, J. Bialek and J. Bumby,
or Power System Stability and Control by P. Kundur, for the equations
"""

function mdl_machine_ode!(
    device_states,
    output_ode,
inner_vars,
current_r::AbstractArray{<:ACCEPTED_REAL_TYPES},
    current_i::AbstractArray{<:ACCEPTED_REAL_TYPES},
    dynamic_device::DynamicWrapper{PSY.DynamicGenerator{PSY.SimpleFullMachine, S, A, TG, P}},
) where {S <: PSY.Shaft, A <: PSY.AVR, TG <: PSY.TurbineGov, P <: PSY.PSS}

    #Obtain indices for component w/r to device
    local_ix = get_local_state_ix(dynamic_device, PSY.SimpleFullMachine)

    #Define internal states for component
    internal_states = @view device_states[local_ix]
    ψf = internal_states[1]
    ψ1d = internal_states[2]
    ψ1q = internal_states[3]

    #Obtain external states inputs for component
    external_ix = get_input_port_ix(dynamic_device, PSY.SimpleFullMachine)
    δ = device_states[external_ix[1]]
    ω = device_states[external_ix[2]]

    #Obtain inner variables for component
    V_tR = inner_vars[VR_gen_var]
    V_tI = inner_vars[VI_gen_var]
    Vf = inner_vars[Vf_var]
    ψd = inner_vars[ψd_var]
    ψq = inner_vars[ψq_var]

    #Get parameters
    machine = PSY.get_machine(dynamic_device)
    R = PSY.get_R(machine)
    R_f = PSY.get_R_f(machine)
    R_1d = PSY.get_R_1d(machine)
    R_1q = PSY.get_R_1q(machine)
    inv_d_fluxlink = PSY.get_inv_d_fluxlink(machine)
    inv_q_fluxlink = PSY.get_inv_q_fluxlink(machine)
    BaseMVA = PSY.get_MVABase(machine)

    #RI to dq transformation
    V_dq = ri_dq(δ) * [V_tR; V_tI]

    #Obtain electric currents] using Flux Linkage equations
    i_df1d = inv_d_fluxlink * [ψd; ψf; ψ1d]  #11.18 in Machowski (3.127, 3.130 and 3.131 in Kundur)
    i_q1q = inv_q_fluxlink * [ψq; ψ1q]        #11.19 in Machowski (3.128 and 3.132 in Kundur)

    #Compute ODEs
    output_ode[local_ix[1]] = -R_f * i_df1d[2] + Vf      #11.33 in Machowski or 3.123 in Kundur
    output_ode[local_ix[2]] = -R_1d * i_df1d[3]           #11.33 in Machowski or 3.124 in Kundur
    output_ode[local_ix[3]] = -R_1q * i_q1q[2]            #11.33 in Machowski or 3.125 in Kundur

    #Update stator fluxes and torque
    ψd = (1.0 / ω) * (R * i_q1q[1] + V_dq[2])         #11.32 in Machowski or 3.121 in Kundur
    ψq = (1.0 / ω) * (-R * i_df1d[1] - V_dq[1])       #11.32 in Machowski or 3.120 in Kundur
    τ_e = ψd * i_q1q[1] - ψq * i_df1d[1]            #15.6 in Milano or 3.117 in Kundur

    #Update inner_vars
    inner_vars[τe_var] = τ_e
    inner_vars[ψd_var] = ψd
    inner_vars[ψq_var] = ψq

    #Compute current from the generator to the grid
    I_RI = (BaseMVA / Sbase) * dq_ri(δ) * [i_df1d[1]; i_q1q[1]]

    #Update current
    current_r[1] += I_RI[1]
    current_i[1] += I_RI[2]

    return
end
=#


using Test

@testset "Hypothesis tests" begin include("test_hypothesis_tests.jl") end
@testset "FSSD" begin include("fssd_tests.jl") end
@testset "KSD" begin include("ksd_tests.jl") end
@testset "Univariate distributions" begin include("test_univariate.jl") end


# An implementation of CRLS for the solution of the
# over-determined linear least-squares problem
#
#  minimize ‖Ax - b‖
#
# equivalently, of the linear system
#
#  A'Ax = A'b.
#
# This implementation follows the formulation given in
#
# D. C.-L. Fong, Minimum-Residual Methods for Sparse
# Least-Squares using Golubg-Kahan Bidiagonalization,
# Ph.D. Thesis, Stanford University, 2011
#
# with the difference that it also recurs r = b - Ax.
#
# Dominique Orban, <dominique.orban@gerad.ca>
# Princeton, NJ, March 2015.

export crls

"""Solve the linear least-squares problem

  minimize ‖b - Ax‖₂² + λ ‖x‖₂²

using the Conjugate Residuals (CR) method. This method is equivalent to
applying MINRES to the normal equations

  (A'A + λI) x = A'b.

This implementation recurs the residual r := b - Ax.

CRLS produces monotonic residuals ‖r‖₂ and optimality residuals ‖A'r‖₂.
It is formally equivalent to LSMR, though can be slightly less accurate,
but simpler to implement.
"""
function crls(A :: AbstractLinearOperator, b :: AbstractVector{T};
              M :: AbstractLinearOperator=opEye(size(b,1)), λ :: Float64=0.0,
              atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, radius :: Float64=0.0,
              itmax :: Int=0, verbose :: Bool=false) where T <: Number

  m, n = size(A);
  size(b, 1) == m || error("Inconsistent problem size");
  verbose && @printf("CRLS: system of %d equations in %d variables\n", m, n);

  x = zeros(T, n)
  r  = copy(b)
  bNorm = @knrm2(m, r)  # norm(b - A * x0) if x0 ≠ 0.
  bNorm == 0 && return x, SimpleStats(true, false, [0.0], [0.0], "x = 0 is a zero-residual solution");
  Mr = M * r;

  # The following vector copy takes care of the case where A is a LinearOperator
  # with preallocation, so as to avoid overwriting vectors used later. In other
  # cases, this should only add minimum overhead.
  Ar = copy(A' * Mr);  # - λ * x0 if x0 ≠ 0.

  s  = A * Ar;
  Ms = M * s;

  p  = copy(Ar);
  Ap = copy(s);
  q  = A' * Ms;  # Ap;
  λ > 0 && @kaxpy!(n, λ, p, q)  # q = q + λ * p;
  γ  = @kdot(m, s, Ms)  # Faster than γ = dot(s, Ms);
  iter = 0;
  itmax == 0 && (itmax = m + n);

  rNorm = bNorm;  # + λ * ‖x0‖ if x0 ≠ 0 and λ > 0.
  ArNorm = @knrm2(n, Ar)  # Marginally faster than norm(Ar);
  λ > 0 && (γ += λ * ArNorm * ArNorm);
  rNorms = [rNorm;];
  ArNorms = [ArNorm;];
  ε = atol + rtol * ArNorm;
  verbose && @printf("%5s  %8s  %8s\n", "Aprod", "‖A'r‖", "‖r‖")
  verbose && @printf("%5d  %8.2e  %8.2e\n", 3, ArNorm, rNorm);

  status = "unknown";
  on_boundary = false
  solved = ArNorm <= ε;
  tired = iter >= itmax;
  psd = false

  while ! (solved || tired)
    qNorm² = @kdot(n, q, q) # dot(q, q)
    α = γ / qNorm²

    # if a trust-region constraint is give, compute step to the boundary
    # (note that α > 0 in CRLS)
    if radius > 0.0
      pNorm = @knrm2(n, p)
      if @knrm2(m, Ap)^2 ≤ ε * sqrt(qNorm²) * pNorm # the quadratic is constant in the direction p
        psd = true # det(AᵀA) = 0
        p = Ar # p = Aᵀr
        pNorm² = ArNorm * ArNorm
        q = A' * s
        α = min(ArNorm^2 / γ, maximum(to_boundary(x, p, radius, flip = false, dNorm2 = pNorm²))) # the quadratic is minimal in the direction Aᵀr for α = ‖Ar‖²/γ
      else
        pNorm² = pNorm * pNorm
        σ = maximum(to_boundary(x, p, radius, flip = false, dNorm2 = pNorm²))
        if α ≥ σ
          α = σ
          on_boundary = true
        end
      end
    end

    @kaxpy!(n,  α, p,   x)     # Faster than  x =  x + α *  p;
    @kaxpy!(n, -α, q,  Ar)     # Faster than Ar = Ar - α *  q;
    ArNorm = @knrm2(n, Ar)
    solved = psd || on_boundary
    solved && continue
    @kaxpy!(m, -α, Ap,  r)     # Faster than  r =  r - α * Ap;
    s = A * Ar;
    Ms = M * s;
    γ_next = @kdot(m, s, Ms)   # Faster than γ_next = dot(s, s);
    λ > 0 && (γ_next += λ * ArNorm * ArNorm);
    β = γ_next / γ;

    @kscal!(n, β, p)
    @kaxpy!(n, 1.0, Ar, p)    # Faster than  p = Ar + β *  p;
    # The combined call uses less memory but tends to trigger more gc.
    #     BLAS.axpy!(n, 1.0, Ar, 1, BLAS.scal!(n, β, p, 1), 1);

    @kscal!(m, β, Ap)
    @kaxpy!(m, 1.0, s, Ap)    # Faster than Ap =  s + β * Ap;
    q = A' * M * Ap;
    λ > 0 && @kaxpy!(n, λ, p, q)  # q = q + λ * p;

    γ = γ_next;
    if λ > 0
      rNorm = sqrt(@kdot(m, r, r) + λ * @kdot(n, x, x))
    else
      rNorm = @knrm2(m, r)  # norm(r);
    end
    #     ArNorm = norm(Ar);
    push!(rNorms, rNorm);
    push!(ArNorms, ArNorm);
    iter = iter + 1;
    verbose && @printf("%5d  %8.2e  %8.2e\n", 3 + 2 * iter, ArNorm, rNorm);
    solved = (ArNorm <= ε) || on_boundary
    tired = iter >= itmax;
  end

  status = psd ? "zero-curvature encountered" : (on_boundary ? "on trust-region boundary" : (tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"))
  stats = SimpleStats(solved, false, rNorms, ArNorms, status);
  return (x, stats);
end


@testset "Coordinate constructor" begin
    coord = Coordinate([1.0, 2.0, 3.0])
    @test isapprox(coord.x, 1.0, atol = 1e-5)
    @test isapprox(coord.y, 2.0, atol = 1e-5)
    @test isapprox(coord.z, 3.0, atol = 1e-5)
end

@testset "operator + for Coordinate" begin
    coord1 = Coordinate([1.0, 2.0, 3.0])
    coord2 = Coordinate([2.0, 3.0, 4.0])
    res = coord1 + coord2
    @test isapprox(res.x, 3.0, atol = 1e-5)
    @test isapprox(res.z, 7.0, atol = 1e-5)
end

@testset "operator - for Coordinate" begin
    coord1 = Coordinate([3.0, 2.0, 1.0])
    coord2 = Coordinate([2.0, 3.0, 4.0])
    res = coord2 - coord1
    @test isapprox(res.x, -1.0, atol = 1e-5)
    @test isapprox(res.z, 3.0, atol = 1e-5)
end

@testset "operator * for Coordinate" begin
    coord1 = Coordinate([1.0, 2.0, 3.0])
    coord2 = Coordinate([2.0, 3.0, 4.0])
    coord_coord_case = coord1 * coord2
    @test isapprox(coord_coord_case, 20.0, atol = 1e-5)
    coord_scholar_case = coord1 * 2.0
    @test isapprox(Array(coord_scholar_case), [2.0, 4.0, 6.0], atol = 1e-5)

end

@testset "operator / for Coordinate" begin
    coord1 = Coordinate([1.0, 2.0, 3.0])
    res = coord1 / 2.0
    @test isapprox(res.x, 0.5, atol = 1e-5)
    @test isapprox(res.z, 1.5, atol = 1e-5)
end

@testset "norm for Coordinate" begin
    coord = Coordinate([1.0, 2.0, 3.0])
    resnorm = norm(coord)
    predicted_val    = sqrt(1.0^2 + 2.0^2 + 3.0^2)
    @test isapprox(resnorm, predicted_val, atol = 1e-5)
end

@testset "broadcast for Coordinate" begin
    coord_a = Coordinate([1.0, 2.0, 3.0])
    coord_b = Coordinate([2.0, 3.0, 4.0])
    coord_c = Coordinate([3.0, 4.0, 5.0])
    coord_d = Coordinate([3.0, 2.0, 1.0])
    coords_vec = [coord_a, coord_b, coord_c]
    @test isapprox(Array((coords_vec .+ coord_d)[3]), [6.0, 6.0, 6.0],    atol = 1e-5)
    @test isapprox(Array((coords_vec .- coord_d)[2]), [-1.0, 1.0, 3.0],   atol = 1e-5)
    @test isapprox(       coords_vec .* coord_d,      [10.0, 16.0, 22.0], atol = 1e-5)
    @test isapprox(Array((coords_vec .* 2.0)[2]),     [4.0, 6.0, 8.0],    atol = 1e-5)
end


"""copy a FEFunction on a FE Space"""
clone_fe_function(space,f)=FEFunction(space,copy(get_free_dof_values(f)))

"""
 Given an arbitrary number of FEFunction arguments,
 returns a tuple with their corresponding free DoF Values
"""
function Gridap.get_free_dof_values(functions...)
  map(get_free_dof_values,functions)
end

function setup_mixed_spaces(model, order)
  reffe_rt  = ReferenceFE(raviart_thomas, Float64, order)
  V = FESpace(model, reffe_rt ; conformity=:HDiv)
  U = TrialFESpace(V)
  reffe_lgn = ReferenceFE(lagrangian, Float64, order)
  Q = FESpace(model, reffe_lgn; conformity=:L2)
  P = TrialFESpace(Q)
  reffe_lgn = ReferenceFE(lagrangian, Float64, order+1)
  S = FESpace(model, reffe_lgn; conformity=:H1)
  R = TrialFESpace(S)

  R, S, U, V, P, Q
end

function setup_and_factorize_mass_matrices(dΩ,
                                           R, S, U, V, P, Q;
                                           mass_matrix_solver=BackslashSolver())
  amm(a,b) = ∫(a⋅b)dΩ
  H1MM   = assemble_matrix(amm, R, S)
  RTMM   = assemble_matrix(amm, U, V)
  L2MM   = assemble_matrix(amm, P, Q)

  ssH1MM = symbolic_setup(mass_matrix_solver,H1MM)
  nsH1MM = numerical_setup(ssH1MM,H1MM)

  ssRTMM = symbolic_setup(mass_matrix_solver,RTMM)
  nsRTMM = numerical_setup(ssRTMM,RTMM)

  ssL2MM = symbolic_setup(mass_matrix_solver,L2MM)
  nsL2MM = numerical_setup(ssL2MM,L2MM)

  H1MM, RTMM, L2MM, nsH1MM, nsRTMM, nsL2MM
end

function new_vtk_step(Ω,file,_cellfields)
  n = size(_cellfields)[1]
  createvtk(Ω,
            file,
            cellfields=_cellfields,
            nsubcells=n)
end


### A Pluto.jl notebook ###
# v0.14.3

using Markdown
using InteractiveUtils

# ╔═╡ dba34531-0a3d-4598-a4ba-52cf81203caf
using Pkg; Pkg.add(["Plots","LinearAlgebra"])

# ╔═╡ ed657f60-a5f3-11eb-2974-8714af0f16dc
using LinearAlgebra,Plots

# ╔═╡ 97916012-f5e5-4957-99c3-4f91d7b86b9d
function ansatz(r1,r2,α)
    norm_r1 = norm(r1)
    norm_r2 = norm(r2)
    r12 = norm(r1-r2)
    wf = exp(-2*norm_r1)*exp(-2*norm_r2)*exp(r12/(2*(1+α*r12)))
    return wf
end;

# ╔═╡ 9a7c2034-7b47-4c70-b469-60d0e1bcf159
function prob_density(r1,r2,α)
    return ansatz(r1,r2,α)^2
end;

# ╔═╡ 9f80362c-5e59-4e73-b878-ddba6429974a
function E_local(r1,r2,α)
    norm_r1 = norm(r1)
    norm_r2 = norm(r2)
    r12 = norm(r1-r2)
    dot_product = dot(r1/norm_r1-r2/norm_r2,r1-r2)
    energy = -4+dot_product/(r12*(1+α*r12)^2)-1/(r12*(1+α*r12)^3)-1/(4*(1+α*r12)^4)+1/r12
    return energy
end;

# ╔═╡ 3c677d47-d02c-4328-9e12-7ad2ce540ade
function metropolis(N, α)

    L = 1
    r1 = rand(3)*2*L .- L
    r2 = rand(3)*2*L .- L #random number from -L to L
    E = 0
    E2 = 0
    Eln_average = 0
    ln_average = 0
    rejection_ratio = 0
    step = 0
    max_steps = -1

    for i in 1:N
        chose = rand()
        step = step + 1
        if chose < 0.5
            r1_trial = r1 + 0.5*(rand(3)*2*L .- L)
            r2_trial = r2
        else
            r2_trial = r2 + 0.5*(rand(3)*2*L .- L)
            r1_trial = r1
        end
        if prob_density(r1_trial,r2_trial,α) >= prob_density(r1,r2,α)
            r1 = r1_trial
            r2 = r2_trial
        else
            dummy = rand()
            if dummy < prob_density(r1_trial,r2_trial,α)/prob_density(r1,r2,α)
                r1 = r1_trial
                r2 = r2_trial
            else
                rejection_ratio += 1/N
            end
        end

        if step > max_steps
            E += E_local(r1,r2,α)/(N-max_steps)
            E2 += E_local(r1,r2,α)^2/(N-max_steps)
            r12 = norm(r1-r2)
            Eln_average += (E_local(r1,r2,α)*-r12^2/(2*(1+α*r12)^2))/(N-max_steps)
            ln_average += -r12^2/(2*(1+α*r12)^2)/(N-max_steps)
        end
    end
    return E, E2, Eln_average, ln_average, rejection_ratio
end;

# ╔═╡ 4ad41744-94c3-4624-b5b3-f5bf679a7957

function iterate_alpha()
	α = 0

	alpha_iterations = 15
	N_metropolis = 5000
	random_walkers = 200
	γ = 0.5

	energy_plot = []
	alpha_plot = []
	variance_plot = []

	for i in 1:alpha_iterations
	    E = 0
	    E2 = 0
	    dE_dalpha = 0
	    Eln = 0
	    ln = 0
	    rejection_ratio = 0

	    for j in 1:random_walkers
	        E_met, E2_met, Eln_met, ln_met, rejections_met = metropolis(N_metropolis,α)
	        E += E_met/random_walkers
	        E2 += E2_met/random_walkers
	        Eln += Eln_met/random_walkers
	        ln += ln_met/random_walkers
	        rejection_ratio += rejections_met/random_walkers
	    end

		α = α + 0.05

	    energy_plot = append!(energy_plot, E)
	    alpha_plot = append!(alpha_plot, α)
	    variance_plot = append!(variance_plot, E2-E^2)
	end
	return energy_plot,alpha_plot,variance_plot
end;

# ╔═╡ 82e68fc3-f949-43f8-8704-f05e19dec3aa
energy,alpha,variance = iterate_alpha();

# ╔═╡ 6f50f51d-1ec8-4031-af1a-0d6cd69171a0
plot(energy,
	label = false,
	title = "Energy expenctation Plot",
	linecolor = :blue,
	ylabel = "<E>",
	xlabel = "timestep"
)

# ╔═╡ d797e924-293c-496b-8ac8-f06b59b9ad89
plot(variance,
	title = "Variance in energy",
	label = false,
	linecolor = :blue,
	ylabel = "Var(E)",
	xlabel = "timestep"
)

# ╔═╡ 03294208-6397-4602-9157-58ca24154003
plot(alpha,
	title = "Variation Parameter",
	label = false,
	linecolor = :blue,
	ylabel = "α",
	xlabel = "timestep"
)

# ╔═╡ 53b76a30-942b-4191-982a-f06fb38bd52f
scatter(alpha,
	energy,
	title = "Energy vs α",
	xlabel = "α",
	ylabel = "<E>",
	label = false,
	markercolor = :blue
)

# ╔═╡ Cell order:
# ╠═dba34531-0a3d-4598-a4ba-52cf81203caf
# ╠═ed657f60-a5f3-11eb-2974-8714af0f16dc
# ╠═97916012-f5e5-4957-99c3-4f91d7b86b9d
# ╠═9a7c2034-7b47-4c70-b469-60d0e1bcf159
# ╠═9f80362c-5e59-4e73-b878-ddba6429974a
# ╠═3c677d47-d02c-4328-9e12-7ad2ce540ade
# ╠═4ad41744-94c3-4624-b5b3-f5bf679a7957
# ╠═82e68fc3-f949-43f8-8704-f05e19dec3aa
# ╠═6f50f51d-1ec8-4031-af1a-0d6cd69171a0
# ╠═d797e924-293c-496b-8ac8-f06b59b9ad89
# ╠═03294208-6397-4602-9157-58ca24154003
# ╠═53b76a30-942b-4191-982a-f06fb38bd52f