function tomography(probabilities::Dict{Tuple,<:Dict}, sites::Vector{<:Index};
                    method::String="linear_inversion",
                    fillzeros::Bool=true,
                    process::Bool = false,
                    trρ::Number = 1.0,
                    max_iters::Int=10000,
                    kwargs...)

  # Generate the projector matrix corresponding to the probabilities.
  A, p = design_matrix(probabilities; return_probs = true, process = process)

  if (method == "LI"  || method == "linear_inversion")
    # Invert the Born rule and reshape
    ρ_vec = pinv(A) * p
    d = Int(sqrt(size(ρ_vec,1)))
    ρ̂ = reshape(ρ_vec,(d,d))

    ρ̂ .= ρ̂ * (trρ / tr(ρ̂))
    # Make PSD
    ρ̂ = make_PSD(ρ̂)
  else
    d = Int(sqrt(size(A,2)))
    N = Int(sqrt(d))
    n = N ÷ 2

    # Variational density matrix
    ρ = Convex.ComplexVariable(d,d)

    if (method == "LS"  || method == "least_squares")
      # Minimize the cost function C = ||A ρ⃗ - p̂||²
      cost_function = Convex.norm(A * vec(ρ) - p)
    elseif (method == "MLE" || method == "maximum_likelihood")
      # Minimize the negative log likelihood:
      cost_function = - p' * Convex.log(real(A * vec(ρ)) + 1e-10)
    else
      error("Tomography method not recognized
       Currently available methods: - LI  : linear inversion
                                      LS  : least squares
                                      MLS : maximum likelihood")
    end

    # Contrained the trace and enforce positivity and hermitianity
    if process
      function tracepreserving(ρ)
        for j in 1:n
          subsystem_dims = [2 for _ in 1:(N+1-j)]
          ρ = Convex.partialtrace(ρ,j+1,subsystem_dims)
        end
        return ρ
      end

      constraints = [Convex.tr(ρ) == (1<<n) * trρ
                    Convex.isposdef(ρ)
                    tracepreserving(ρ) == Matrix{Float64}(I,1<<n,1<<n)
                    ρ == ρ']
    else
      constraints = [Convex.tr(ρ) == trρ
                     Convex.isposdef(ρ)
                     ρ == ρ']
    end
    # Use Convex.jl to solve the optimization
    problem = Convex.minimize(cost_function,constraints)
    Convex.solve!(problem, () -> SCS.Optimizer(verbose=false,max_iters=max_iters),verbose=false)
    ρ̂ = ρ.value
  end
  return PastaQ.itensor(ρ̂, sites)
end

tomography(data::Matrix{Pair{String, Int}}; method::String="linear_inversion", fillzeros::Bool=true, kwargs...) =
  tomography(empirical_probabilities(data; fillzeros=fillzeros), siteinds("Qubit", size(data,2)); method = method, kwargs...)

tomography(data::Matrix{Pair{String, Int}}, sites::Vector{<:Index}; method::String="linear_inversion", fillzeros::Bool=true, kwargs...) =
  tomography(empirical_probabilities(data; fillzeros=fillzeros), sites; method = method, kwargs...)

tomography(data::Matrix{Pair{String,Pair{String, Int}}}; kwargs...) =
  tomography(data, siteinds("Qubit", size(data,2)); kwargs...)

function tomography(data::Matrix{Pair{String,Pair{String, Int}}}, sites::Vector{<:Index}; method::String="linear_inversion", fillzeros::Bool=true, kwargs...)
  sites_in  = addtags.(sites, "Input")
  sites_out = addtags.(sites, "Output")
  process_sites = Index[]
  for j in 1:size(data,2)
    push!(process_sites, sites_in[j])
    push!(process_sites, sites_out[j])
  end
  tomography(empirical_probabilities(data; fillzeros=fillzeros), process_sites; method = method, process = true, kwargs...)
end

"""
    tomography(train_data::Matrix{Pair{String,Pair{String, Int}}}, L::LPDO;
               optimizer::Optimizer,
               observer! = nothing,
               batchsize::Int64 = 100,
               epochs::Int64 = 1000,
               kwargs...)

Run quantum process tomography using a variational model `L` to fit `train_data`.
The model can be either a unitary circuit (MPO) or a Choi matrix (LPDO).

# Arguments:
- `train_data`: pairs of preparation/ (basis/outcome): `("X+"=>"X"=>0, "Z-"=>"Y"=>1, "Y+"=>"Z"=>0, …)`.
 - `L`: variational model (MPO/LPDO).
 - `optimizer`: algorithm used to update the model parameters.
 - `observer!`: if provided, keep track of training metrics.
 - `batch_size`: number of samples used for one gradient update.
 - `epochs`: number of training iterations.
 - `target`: target quantum process (if provided, compute fidelities).
 - `test_data`: data for computing cross-validation.
 - `outputpath`: if provided, save metrics on file.
"""
function tomography(
  train_data::AbstractMatrix, L::LPDO; (observer!)=nothing, kwargs...
)
  # Read arguments
  opt = get(kwargs, :optimizer, Optimisers.Descent(0.01))
  batchsize::Int64 = get(kwargs, :batchsize, 100)
  epochs::Int64 = get(kwargs, :epochs, 1000)
  trace_preserving_regularizer = get(kwargs, :trace_preserving_regularizer, 0.0)
  observe_step::Int64 = get(kwargs, :observe_step, 1)
  test_data = get(kwargs, :test_data, nothing)
  outputpath = get(kwargs, :fout, nothing)
  print_metrics = get(kwargs, :print_metrics, [])
  outputpath = get(kwargs, :outputpath, nothing)
  outputlevel = get(kwargs, :outputlevel, 1)
  savestate = get(kwargs, :savestate, false)
  earlystop = get(kwargs, :earlystop, false)

  model = copy(L)
  isqpt = train_data isa Matrix{Pair{String,Pair{String,Int}}}
  localnorm = isqpt ? 2.0 : 1.0

  # observer is not passed but earlystop is called
  observer! = (isnothing(observer!) && earlystop) ? Observer() : observer!

  # observer is defined
  if !isnothing(observer!)
    observer!["train_loss"] = nothing
    if !isnothing(test_data)
      observer!["test_loss"] = nothing
    end
    # add the standard early stop function to the observer
    if earlystop
      stop_if(; loss::Vector) = stopif_loss(; loss = loss, ϵ = 1e-3, min_iter = 10)
      observer!["earlystop"] = stopif
    end
  end

  # initialize optimizer
  st = PastaQ.state(opt, model)
  optimizer = (opt, st)

  @assert size(train_data, 2) == length(model)
  !isnothing(test_data) && @assert size(test_data)[2] == length(model)

  batchsize = min(size(train_data)[1], batchsize)
  num_batches = Int(floor(size(train_data)[1] / batchsize))

  tot_time = 0.0
  best_model = nothing
  best_testloss = 1000.0
  test_loss = nothing

  # Training iterations
  for ep in 1:epochs
    ep_time = @elapsed begin
      train_data = train_data[shuffle(1:end), :]
      train_loss = 0.0

      # Sweep over the data set
      for b in 1:num_batches
        batch = train_data[((b - 1) * batchsize + 1):(b * batchsize), :]

        normalized_model = copy(model)
        sqrt_localnorms = []
        normalize!(normalized_model;
                   (sqrt_localnorms!)=sqrt_localnorms,
                   localnorm=localnorm)

        grads, loss = gradients(
          normalized_model,
          batch;
          sqrt_localnorms=sqrt_localnorms,
          trace_preserving_regularizer=trace_preserving_regularizer,
        )

        nupdate = ep * num_batches + b
        train_loss += loss / Float64(num_batches)
        update!(model, grads, optimizer)
      end
    end # end @elapsed
    !isnothing(observer!) && push!(last(observer!["train_loss"]), train_loss)
    tot_time += ep_time

    # measurement stage
    if ep % observe_step == 0
      normalized_model = copy(model)
      sqrt_localnorms = []
      normalize!(normalized_model;
                 (sqrt_localnorms!)=sqrt_localnorms,
                 localnorm=localnorm)

      if !isnothing(test_data)
        test_loss = nll(normalized_model, test_data)
        !isnothing(observer!) && push!(last(observer!["test_loss"]), test_loss)
        if test_loss ≤ best_testloss
          best_testloss = test_loss
          best_model = copy(normalized_model)
        end
      else
        best_model = copy(normalized_model)
      end

      # update observer
      if !isnothing(observer!)
        loss = (!isnothing(test_data) ? results(observer!, "test_loss") :
                                        results(observer!, "train_loss"))

        model_to_observe = (isqpt && (normalized_model isa LPDO{MPS}) ? choi_mps_to_unitary_mpo(normalized_model) :
                                  !isqpt && (normalized_model isa LPDO{MPS}) ? normalized_model.X :
                                                                               normalized_model)
        update!(observer!, model_to_observe; train_loss = train_loss,
                                             test_loss  = test_loss)
      end

      # printing
      if outputlevel ≥ 1
        @printf("%-4d  ", ep)
        @printf("⟨logP⟩ = %-4.4f  ", train_loss)
        if !isnothing(test_data)
          @printf("(%.4f)  ", test_loss)
        end
        # TODO: add the trace preserving cost function here for QPT
        !isnothing(observer!) && printobserver(observer!, print_metrics)
        @printf("elapsed = %-4.3fs", ep_time)
        println()
      end
      # saving
      if !isnothing(outputpath)
        observerpath = outputpath * "_observer.jld2"
        save(observerpath, observer!)
        if savestate
          if isqpt
            model_to_be_saved = model isa LPDO{MPS} ? choi_mps_to_unitary_mpo(best_model) : best_model
          else
            model_to_be_saved = model isa LPDO{MPS} ? best_model.X : best_model
          end
          statepath = outputpath * "_state.h5"
          h5rewrite(statepath) do fout
            write(fout, "state", model_to_be_saved)
          end
        end
      end
    end
    !isnothing(observer!) && haskey(observer!.data,"earlystop") && results(observer!, "earlystop")[end] && break
  end
  return best_model
end

# QST
tomography(data::Matrix{Pair{String,Int}}, ψ::MPS; kwargs...) =
  tomography(data, LPDO(ψ); kwargs...).X

tomography(train_data::Vector{<:Vector{Pair{String,Int}}}, args...; kwargs...) =
  tomography(permutedims(hcat(train_data...)), args...; kwargs...)

# QPT
tomography(data::Matrix{Pair{String,Pair{String,Int}}}, U::MPO; kwargs...) =
  choi_mps_to_unitary_mpo(tomography(data, LPDO(unitary_mpo_to_choi_mps(U)); kwargs...))

tomography(train_data::Vector{<:Vector{Pair{String,Pair{String,Int}}}}, args...; kwargs...) =
  tomography(permutedims(hcat(train_data...)), args...; kwargs...)

"""
EARLY STOPPING FUNCTIONS
"""

#stopif_fidelity(M1, M2; ϵ::Number, kwargs...)
#  fidelity(M1,M2) ≤ ϵ

function stopif_loss(; loss::Vector, ϵ::Number, min_iter::Number)
  length(loss) < min_iter+1 && return false
  avgloss = StatsBase.mean(loss[end-size:end])
  Δ = StatsBase.sem(historyloss[end-size:end])
  return Δ/avgloss < ϵ
end


@symbol_func function cross_sectional_area(cur_reactor::AbstractReactor)
  cur_volume = cur_reactor.pi

  cur_volume *= a(cur_reactor) ^ 2

  cur_volume *= cur_reactor.kappa_95

  cur_volume *= g(cur_reactor)

  cur_volume
end


function seager_mr(mp,comp)
# computes the radius of a planet
# given the mass & composition of the planet
# according to the formulae in Seager et al. (2007)
k1=-0.20945
k2=0.0804
k3=0.394
if comp == "Fe"
  m1=4.34
  r1=2.23
end
if comp == "Earth"
  m1=6.41
  r1=3.19
#  r1=2.84
end
if comp == "MgSiO3"
  m1=7.38
  r1=3.58
end
if comp == "H2O"
  m1=8.16
  r1=4.73
end
if comp == "FeMgH2O"
#  m1 = 6.88
#  r1 = 4.02
  m1 = 6.41
  r1 = 3.63
end
ms=mp/m1
lgrp=k1 .+ log10.(ms)/3.0 .-k2*ms.^k3
return 10.0.^lgrp*r1
end


function Euler6(N::Int64)

	println( div(N*(N+1)*(3*N+2)*(N-1),12))

end

@time Euler6(100)


using StructJuMP

numScens = 2
m = StructuredModel(num_scenarios=numScens)

@variable(m, x >= 1)
@variable(m, y <= 2)

@constraint(m,  x + y == 1)
@constraint(m, -x + y <= 1)
@objective(m, :Min, x*x + 0.5x*y + 0.25y*y - y)

rhs  = [5,4]
coef = [2,3]
qc   = [1,0.5]
ac   = [1,0.75]

for i = 1:numScens
    bl = StructuredModel(parent=m, id=i)
    @variable(bl, 0 <= w <= 1)
    @constraint(bl, coef[i]w - x - y <= rhs[i])
    @constraint(bl, coef[i]w + x     == rhs[i])
    @objective(bl, :Min, qc[i]*w*w + ac[i]*w)
end


using CSV, GLM, Distributions, Plots; pyplot()
data = CSV.read("../data/weightHeight.csv")
n = size(data)[1]
model = fit(LinearModel, @formula(Height ~ Weight), data)

alpha = 0.1
tVal = quantile(TDist(n-2),1-alpha/2)

xbar = mean(data.Weight)
Sxx = std(data.Weight)*(n-1)
MSE = deviance(model)/(n-2)
pred(x) = coef(model)'*[1, x]

interval(x,sign,prediction = 0) = sign(pred(x),
			tVal * sqrt(MSE*(prediction+1/n+(x-xbar)^2/Sxx)) )

xGrid = 40:1:140
scatter(data.Weight, data.Height, c=:black, ms=2, label="")
plot!(xGrid,pred.(xGrid), c=:red, label="Linear model")
plot!(xGrid,interval.(xGrid,+),c=:green, label="Confidence bands")
plot!(xGrid,interval.(xGrid,-),c=:green, label="")
plot!(xGrid,interval.(xGrid,+,1),c=:blue,label="Prediction bands")
plot!(xGrid,interval.(xGrid,-,1),
	c=:blue, label="", xlims=(40, 120), ylims=(145, 200), legend=:topleft,
	xlabel = "Weight (kg)", ylabel = "Height (cm)")


@testset "Initial scenario tests" begin
    @test isapprox(initialScenario([0], [10], 5), [5.00061], atol = 1e-5)

    @test isapprox(initialScenario([2, 3], [8, 9], 10), [7.50073, 7.50073], atol = 1e-5)

    @test isapprox(initialScenario([2, 3, 4, 8, 1], [5, 9, 7, 10, 10], 13), [5.74969, 5.74969, 5.74969, 8.0, 5.74969], atol = 1e-5)
end


include("../../src/Transform/dct.jl")

@testset "Transform/dct" begin
	a=rand(10, 20, 30)
	b=copy(a)
	dct!(b)
	idct!(b)
	@test isapprox(a, b)
	a=rand(12, 22, 32)
	b=copy(a)
	dct!(b)
	idct!(b)
	@test isapprox(a, b)
end


module TestPermutation
  using Base.Test
  using Brim

  # Test utility functions

  c1 = [1 0; 0 1]
  c2 = [0 1; 1 0]
  c3 = [0 0; 1 1]
  c4 = [1 1; 1 1]
  c5 = [0 0; 0 0]

  @test Brim.free_swapable(c1) == true
  @test Brim.free_swapable(c2) == true
  @test Brim.free_swapable(c3) == true
  @test Brim.free_swapable(c4) == false
  @test Brim.free_swapable(c5) == false

  @test Brim.constrained_swapable(c1) == true
  @test Brim.constrained_swapable(c2) == true
  @test Brim.constrained_swapable(c3) == false
  @test Brim.constrained_swapable(c4) == false
  @test Brim.constrained_swapable(c5) == false

  @test Brim.constrained_swap(c1) == c2
  @test Brim.constrained_swap(c2) == c1

  @test Brim.no_empty_rows(c3) == false
  @test Brim.no_empty_rows(c1) == true
  @test Brim.no_empty_columns(c3') == false
  @test Brim.no_empty_columns(c1') == true

  @test Brim.same_row_marginals(c1, c2) == true
  @test Brim.same_row_marginals(c1, c1) == true
  @test Brim.same_row_marginals(c1, c3) == false

  A = [1 1 1 0 0 0; 1 1 1 0 0 0; 1 1 1 0 0 0; 0 0 0 1 1 1; 0 0 0 1 1 1; 0 0 0 1 1 1]

  # null_preserve_marginals
  @test sum(A) == sum(null_preserve_marginals(A))
  @test sum(A) == sum(null_preserve_marginals(A))
  @test sum(A) == sum(null_preserve_marginals(A))

  # null_preserve_rows_marginals
  @test sum(A, 2) == sum(null_preserve_rows_marginals(A), 2)
  @test sum(A, 2) == sum(null_preserve_rows_marginals(A), 2)
  @test sum(A, 2) == sum(null_preserve_rows_marginals(A), 2)

  # null_preserve_columns_marginals
  @test sum(A, 1) == sum(null_preserve_columns_marginals(A), 1)
  @test sum(A, 1) == sum(null_preserve_columns_marginals(A), 1)
  @test sum(A, 1) == sum(null_preserve_columns_marginals(A), 1)

  # null_preserve_fill
  X = null_preserve_fill(A)
  @test sum(A) == sum(X)
  @test Brim.no_empty_rows(X)
  @test Brim.no_empty_columns(X)

end


export MySumSpace,MatrixOperator

struct MySumSpace{D,R} <: Space{D,R}
    spaces::Array{<:Space{D,R},1}
    function MySumSpace(S::Array{<:Space{D,R},1}) where {D,R}
        if !foldl(domainscompatible,domain.(S))
            error("Components of MySumSpace do not have the same domain.")
        else
            new{D,R}(S)
        end
    end
end
struct MatrixOperator{T} <: Operator{T}
    matrix::Array{<:Operator{T},2}
    function MatrixOperator(C::Array{<:Operator{T},2}) where T
        for m in 1:size(C)[1]
            if !foldl(spacescompatible,rangespace.(C[m,:]))
                error("Row $(m) do not have the same rangespace.")
            end
        end
        for n in 1:size(C)[2]
            if !foldl(spacescompatible,domainspace.(C[:,n]))
                error("Column $(n) do not have the same domainspace.")
            end
        end
        new{T}(C)
    end
end

# Basic
Base.getindex(S::MySumSpace,k::Integer)=S.spaces[k]
Base.length(S::MySumSpace)=Base.length(S.spaces)
Base.size(S::MySumSpace)=Base.size(S.spaces)
Base.getindex(D::MatrixOperator,k::Integer,m::Integer)=D.matrix[k,m]
Base.getindex(D::MatrixOperator,k::Integer)=D.matrix[k]
Base.size(D::MatrixOperator)=size(D.matrix)
Base.show(io::IO,C::MatrixOperator)=print(io,typeof(C))

domainspace(S::MatrixOperator)=MySumSpace([domainspace(S.matrix[1,n]) for n in 1:size(S.matrix)[2]])
rangespace(S::MatrixOperator)=MySumSpace([rangespace(S.matrix[m,1]) for m in 1:size(S.matrix)[1]])
domain(S::MySumSpace)=domain(S.spaces[1])

spacescompatible(S1::MySumSpace,S2::MySumSpace)=all(spacescompatible.(S1.spaces,S2.spaces))

# Interact with coefficients
*(D::MatrixOperator,f::Array{<:AbstractArray,1})=D.matrix*f

evaluate(f::AbstractArray{<:AbstractArray,1}, S::MySumSpace, x) = sum(evaluate.(f,S.spaces,x))

# MatrixOperator Constructor
function Conversion(S1::MySumSpace, S2::MySumSpace)
    @assert length(S1) == length(S2)
    MatrixOperator([m==n ? Conversion(S1[m],S2[n]) : ZeroOperator(S1[m],S2[n]) for n in 1:length(S2), m in 1:length(S1)])
end

function Conversion(S1::MySumSpace, S2::MySumSpace, plan::Array{<:Integer,1})
    @assert length(S1) == length(S2)
    MatrixOperator([plan[m]==n ? Conversion(S1[m],S2[n]) : ZeroOperator(S1[m],S2[n]) for n in 1:length(S2), m in 1:length(S1)])
end


## We'll learn about control flow!

## Simple stuff
x = Inf # we can use Infinity as a number !
        # this makes a Float64
typeof(x)
##
if x < 1
    println("x is smaller than 1")
end
##
if x < 1
    println("x is smaller than 1")
else
    println("x is larger than or equal to 1")
end

##
x = -5
if x < 1
    println("x is smaller than 1")
elseif x < 0
    println("x is smaller than 0")
else
    println("x is larger than or equal to 1")
end

##
if x < 0
    println("x is smaller than 0")
elseif x < 1
    println("x is smaller than 1")
else
    println("x is larger than or equal to 1")
end

##
x = [1.0, 2.0]
if all(x >= 0) && all( x <= 5)
    println("all elements between 0 and 5")
end

##
x = [1.0, 2.0]
if all(x .>= 0) && all(x .<= 5)
    println("all elements between 0 and 5")
end
##
x = 0
y = 5
if x+y > 10 || x < 0
    y = 0
else
    y = 10
end

## That block returned something,  you can "save"
# the result of a comparison
z = if x+y > 10 || x < 0
    y = 0
else
    y = 10
end

## That's it for "if, elseif, else"

## On to for floops

for i=1:10 # for i = 1 to 10, starting a 1 , ending at 10
    @show i
    println("end loop iteration")
end

##
v = -5:5
for i=v # for i = 1 to 10, starting a 1 , ending at 10
    @show i
    println("end loop iteration")
end

##
for i in v
    @show i
    println("end loop iteration")
end
##
v = [1.0,Inf,-Inf,π]
for i in v
    @show i
    println("end loop iteration")
end
##
v = [1 2 3 4]
for i in v
    @show i
    println("end loop iteration")
end
## Julia iterates over matrices by columns
@show [v' v'.+1]
for i=[v' v'.+1]
    @show i
end
##
x = 1:5
x[3] = -3
##
x = collect(1:5)
x[3] = -3

##
x = copy(1:5)
x[3] = -3
## We want to turn this into a real array
x = collect(1:5)
x[3] = -3

##
@show A = [x' x'.+1]
for i in A
    @show i
end

##
##
@show A = [x x.+1]
for i in A
    @show i
end

##While loops
i = 0
while i < 10
    global i
    if i == 5
        break # this works in for loops too!
    end
    if i==4
        i = i+5
    end
    @show i
    i = i+1
end
@show i

##
function example_while()
    i = 0
    while i < 10
        if i == 5
            break # this works in for loops too!
        end
        if i==4
            i = i+5
        end
        @show i
        i = i+1
    end
    @show i
end
example_while()

## The place I usually run into that issue.
# Let's say I want to compute the number of random variables
# that are bigger than pi.
numsuccess = 0
numtrials = 10000
for i=1:numtrials
    numsuccess += randn() > pi
end

# This case hits that scoping issue :(
# In IJulia/Jupyter, they include a patch
# And working out something else is an active area of
# development for Julia, so this may change in the
# future.

##
numsuccess = 0
numtrials = 10000
for i=1:numtrials
    global numsuccess
    numsuccess += randn() > pi
end


immutable Pareto <: ContinuousUnivariateDistribution
    scale::Float64
    shape::Float64
    function Pareto(sc::Real, sh::Real)
        if sc > 0.0 && sh > 0.0
            new(float64(sc), float64(sh))
        else
            error("shape and scale must be positive")
        end
    end
end

Pareto(scale::Real) = Pareto(scale, 1.0)
Pareto() = Pareto(1.0, 1.0)

cdf(d::Pareto, q::Real) = q >= d.scale ? 1.0 - (d.scale / q)^d.shape : 0.0

entropy(d::Pareto) = log(d.shape / d.scale) + 1. / d.scale + 1.

insupport(d::Pareto, x::Number) = isreal(x) && isfinite(x) && x > d.scale

function kurtosis(d::Pareto)
    a = d.shape
    if a > 4.0
        return (6.0 * (a^3 + a^2 - 6.0 * a - 2.0)) / (a * (a - 3.0) * (a - 4.0))
    else
        error("Kurtosis undefined for Pareto w/ shape <= 4")
    end
end

mean(d::Pareto) = d.shape <= 1.0 ? Inf : (d.scale * d.shape) / (d.scale - 1.0)

median(d::Pareto) = d.scale * 2.0^d.shape

function pdf(d::Pareto, q::Real)
    if q >= d.scale
        return (d.shape * d.scale^d.shape) / (q^(d.shape + 1.0))
    else
        return 0.0
    end
end

quantile(d::Pareto, p::Real) = d.scale / (1.0 - p)^(1.0 / d.shape)

rand(d::Pareto) = d.shape / (rand()^(1.0 / d.scale))

function skewness(d::Pareto)
    a = d.shape
    if a > 3.0
        return ((2.0 * (1.0 + a)) / (a - 3.0)) * sqrt((a - 2.0) / a)
    else
        error("Skewness undefined for Pareto w/ shape <= 3")
    end
end

function var(d::Pareto)
    if d.scale < 2.0
        return Inf
    else
        return (d.shape^2 * d.scale) / ((d.scale - 1.0)^2 * (d.scale - 2.0))
    end
end


using Koma, Plots, LaTeXStrings
pgfplotsx()
rf_wf = [
    #Waveform am_sg_100_100_0
    0, 413, 859, 1337, 1847, 2389, 2962, 3568, 4204, 4870,
    5565, 6288, 7038, 7814, 8613, 9435, 10277, 11138, 12015, 12906,
    13809, 14721, 15640, 16563, 17487, 18409, 19327, 20237, 21136, 22023,
    22892, 23742, 24570, 25372, 26146, 26889, 27598, 28271, 28905, 29497,
    30046, 30549, 31005, 31411, 31767, 32070, 32320, 32515, 32655, 32739,
    32767, 32739, 32655, 32515, 32320, 32070, 31767, 31411, 31005, 30549,
    30046, 29497, 28905, 28271, 27598, 26889, 26146, 25372, 24570, 23742,
    22892, 22023, 21136, 20237, 19327, 18409, 17487, 16563, 15640, 14721,
    13809, 12906, 12015, 11138, 10277, 9435, 8613, 7814, 7038, 6288,
    5565, 4870, 4204, 3568, 2962, 2389, 1847, 1337, 859, 413,
    0
] ./ (2^16/2 - 1) #signed int 16 bits
Trf = 0.486e-3
Arf = 21e-6 * rf_wf
Gx1 = 12.89e-3
ζ11 = 0.29e-3
ζ21 = 0.37e-3
Tg1 = Trf
delay = Tg1 + ζ11 - Trf
ex1 = Sequence([Grad(0,0); Grad(0,0); Grad(Gx1,Tg1,ζ11,ζ21,0);;],[RF(Arf,Trf,-5000,delay);;],[ADC(0,0)])
ex2 = Sequence([Grad(0,0); Grad(0,0); Grad(Gx1,Tg1,ζ11,ζ21,0);;],[RF(Arf,Trf,    0,delay);;],[ADC(0,0)])
ex3 = Sequence([Grad(0,0); Grad(0,0); Grad(Gx1,Tg1,ζ11,ζ21,0);;],[RF(Arf,Trf, 5000,delay);;],[ADC(0,0)])
Gx2 = -.83e-3
ζ12 = 0.026e-3
ζ22 = 0.28e-3
Tg2 = 5.75e-3
ref = Sequence([Grad(0,0); Grad(0,0); Grad(Gx2,Tg2,ζ12,ζ22,0);;])
#Different slices
seq1 = ex1 + ref
seq2 = ex2 + ref
seq3 = ex3 + ref
#Focus Time
t_focus_aux = (abs(Gx2)*(2Tg2+ζ12+ζ22)-ζ21*Gx1)/(2Gx1)
t_focus = (Tg1+ζ11-t_focus_aux)
#SIM
z = range(-2e-2,2e-2,200)
M1 = simulate_slice_profile(seq1;z)
M2 = simulate_slice_profile(seq2;z)
M3 = simulate_slice_profile(seq3;z)
#Plots
plot_seq(seq2; width=480, height=350, slider=false)
f = γ*Gx1*z*1e-3
plot(f,  abs.(M1.xy),label=L"|M_{xy}| @ \Delta f = -5\, \mathrm{kHz}",
        xlabel=L"\mathrm{kHz}",line=3,title="α=$(Koma.get_flip_angle(seq2.RF[1]))")
plot!(f, abs.(M2.xy),label=L"|M_{xy}| @ \Delta f = +0\, \mathrm{kHz}",line=3)
plot!(f, abs.(M3.xy),label=L"|M_{xy}| @ \Delta f = +5\, \mathrm{kHz}",line=3)

#GIF
# anim = @animate for tf = range(0,7.5e-3,60)
#     simParams = Dict("return_Mag"=>true, "end_sim_at"=>tf, "gpu"=>false)
#     M = simulate(phantom, seq2, sys; simParams)
#     #Plots
#     plot(xx*1e2, abs.(M.xy),label="Mxy",line=4,ylim=(-1,1),title="Simulation ended at $(round(simParams["end_sim_at"]*1e3,digits=3)) ms")
#     plot!(xx*1e2, abs.(M.z),label="Mz",line=:dash)
#     plot!(xx*1e2, real.(M.xy),label="Mx",line=:dash)
#     plot!(xx*1e2, imag.(M.xy),label="My",line=:dash)
# end
# gif(anim, "AnimSliceExc.gif", fps = 10)


using Statistics
using DataFrames
using CSV
using Plots; pyplot(fmt="png")
PyPlot.matplotlib.rc("mathtext", fontset="cm");
PyPlot.matplotlib.rc("text", usetex=true);
PyPlot.matplotlib.rc("font", family="serif", size=12);
using LaTeXStrings
using LsqFit

path = "/home/user/thesis/data/scaling/";
graph_type = "Network";
evolution_method = "stochastic_edge_acceptance";
data_file = string(path, graph_type, "_", evolution_method, ".csv");

savepath = "/home/user/thesis/latex/images"
save_bool = true
dpi = 144
if save_bool
	dpi = 300
end

# %% loading and computing df

df_raw = CSV.read(data_file) |> DataFrame;

Int.(df_raw[:, :N] .* df_raw[:, :r_1])
maximum(df_raw[:, :r_1])
cols = [
	:N,
	:r_0,
	:r_1,
	:m_0,
	:m_1,
	:Δr,
	:Δm,
];

df = aggregate(df_raw[:, cols], :N, [mean, std, median]);

# %% plotting Δr

x_min  = 15;
Δ_x    = 1;
x_max  = 24;
border = 0.25;
y_min  = 0.68;
Δ_y    = 0.01;
y_max  = 0.74;

plot_ = plot(legend=:right, dpi=dpi);
plot!(log2.(df[:, :N]), df[:, :r_0_mean], yerror=df[:, :r_0_std],
	label=L"t_0/N",
	line=(1, :black, 0.8),
	marker=(5, :dodgerblue, :circle, stroke(:dodgerblue)),
	xaxis=(latexstring("\\log_2 (N)"), (x_min-border, x_max+border), x_min:Δ_x:x_max),
	yaxis=(latexstring("t/N"), (y_min, y_max), y_min:Δ_y:y_max)
);
plot!(log2.(df[:, :N]), df[:, :r_1_mean], yerror=df[:, :r_1_std],
	label=L"t_1/N",
	line=(1, :black, 0.8),
	marker=(5, :firebrick1, :circle, stroke(:firebrick1)),
);

if save_bool
	savefig(plot_, joinpath(savepath, "r_scaling.png"))
end

df[:, :r_0_mean]
err = round((df[end, :r_1_mean] .- df[end, :r_0_mean]) / 2, digits=3)
round(df[end, :r_0_mean] + err, digits=3)

plot_

# %% plotting and fitting Δr

@. function f_linear(x, α)
	return α[1] * x + α[2]
end

α     = [-0.1, -1.5];
α_exp = [1., -1e6, 0.];
x     = log2.(df[:, :N]);
y     = log2.(df[:, :Δr_mean]);
σ     = df[:, :Δr_std];
fit   = curve_fit(f_linear, x, y, α);
y_fit = f_linear(x, fit.param);
α     = fit.param
σ     = stderror(fit)

plot_ = plot(dpi=dpi);
plot!(x, y_fit,
	label=latexstring("\\mathrm{Fit} \\sim $(round(α[1], digits=3)) ($(Int(round(σ[1] * 1e3, digits=0)))) \\cdot \\log_2(N)"),
	line=(1, :black, 0.8)
);
scatter!(x, y,
	label="Ensemble Average",
	markerstrokecolor=:black,
	marker=(6, :dodgerblue, :circle, stroke(:dodgerblue)),
	xaxis=(L"\log_2 (N)", 15:1:24),
	yaxis=(L"\log_2 (\Delta r)", (-8, -4.5))
);

if save_bool
	savefig(plot_, joinpath(savepath, "delta_r_scaling.png"))
end

plot_

# %%

@. function f_linear(x, α)
	return α[1] * x + α[2]
end

α     = [-0.1, -1.5];
α_exp = [1., -1e6, 0.];
x     = log2.(df[:, :N]);
y     = log2.(df[:, :Δm_mean]);
σ     = df[:, :Δm_std];
fit   = curve_fit(f_linear, x, y, α);
y_fit = f_linear(x, fit.param);
α     = fit.param
σ     = stderror(fit)

plot_ = plot(dpi=dpi);
plot!(x, y_fit,
	label=latexstring("\\mathrm{Fit} \\sim $(round(α[1], digits=3)) ($(Int(round(σ[1] * 1e3, digits=0)))) \\cdot \\log_2(N)"),
	line=(1, :black, 0.8)
);

δ = 0.302
τ = 2.213
exp_upper_bound = -δ * (τ - 2) / (τ - 1);
α_upper_bound = [exp_upper_bound, -exp_upper_bound * x[1] + y[1]]
plot!(x, f_linear(x, α_upper_bound),
	label=latexstring("\\mathrm{Upper Bound} \\sim $(round(exp_upper_bound, digits=3)) \\cdot \\log_2(N)"),
	line=(1, :black, :dash)
);

scatter!(x, y,
	label="Ensemble Average",
	markerstrokecolor=:black,
	marker=(6, :dodgerblue, :circle, stroke(:dodgerblue)),
	xaxis=(L"\log_2 (N)", 15:1:24),
	yaxis=(L"\log_2 (\Delta m)", (-3.8, -2), -3.8:0.2:-2)
);

if save_bool
	savefig(plot_, joinpath(savepath, "delta_m_scaling.png"))
end

plot_


cd(@__DIR__) # changes the directory to the current directory, the default I guess is the HOME
using Pkg; Pkg.activate("."); Pkg.instantiate()

Pkg.add("NeuralPDE")
Pkg.add("Flux")
Pkg.add("ModelingToolkit")
Pkg.add("GalacticOptim")
Pkg.add("Optim")
Pkg.add("DiffEqFlux")
Pkg.add("Plots")
Pkg.add("OrdinaryDiffEq")
Pkg.add("DiffEqOperators")

using OrdinaryDiffEq, DiffEqOperators
using NeuralPDE, Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
import ModelingToolkit: Interval, infimum, supremum

@parameters t, x
@variables u(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2

# PDE
eq  = Dt(u(t,x)) + u(t,x)*Dx(u(t,x)) - (0.01/pi)*Dxx(u(t,x)) ~ 0

# Initial and boundary conditions
bcs = [u(0,x) ~ -sin(pi*x),
       u(t,-1) ~ 0.,
       u(t,1) ~ 0.,
       u(t,-1) ~ u(t,1)]

# Space and time domains
domains = [t ∈ Interval(0.0,1.0),
           x ∈ Interval(-1.0,1.0)]
# Discretization
dx = 0.05

# finite difference
# Method of lines discretization
# PDE system - ModelingToolkit
@named pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])

discretization_fd = MOLFiniteDifference([x=>dx],t)

# Convert the PDE problem into an ODE problem
prob = discretize(pdesys,discretization_fd)

# Solve ODE problem
using OrdinaryDiffEq
sol = solve(prob,Tsit5(),saveat=0.2)

println("antes de chain")

# Neural network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
eltypeθ = eltype(initθ)
parameterless_type_θ = DiffEqBase.parameterless_type(initθ)
strategy = NeuralPDE.GridTraining(dx)

phi = NeuralPDE.get_phi(chain,parameterless_type_θ)
derivative = NeuralPDE.get_numeric_derivative()

indvars = [t,x]
depvars = [u]

_pde_loss_function = NeuralPDE.build_loss_function(eq,indvars,depvars,phi,derivative,nothing,
                                                   chain,initθ,strategy)

bc_indvars = NeuralPDE.get_variables(bcs,indvars,depvars)
_bc_loss_functions = [NeuralPDE.build_loss_function(bc,indvars,depvars,
                                                    phi,derivative,nothing,chain,initθ,strategy,
                                                    bc_indvars = bc_indvar) for (bc,bc_indvar) in zip(bcs,bc_indvars)]

train_sets = NeuralPDE.generate_training_sets(domains,dx,[eq],bcs,eltypeθ,indvars,depvars)
train_domain_set, train_bound_set = train_sets

pde_loss_function = NeuralPDE.get_loss_function(_pde_loss_function,
                                                train_domain_set[1],
                                                eltypeθ,parameterless_type_θ,
                                                strategy)

bc_loss_functions = [NeuralPDE.get_loss_function(loss,set,
                                                 eltypeθ, parameterless_type_θ,
                                                 strategy) for (loss, set) in zip(_bc_loss_functions,train_bound_set)]

loss_functions = [pde_loss_function; bc_loss_functions]
loss_function__ = θ -> sum(map(l->l(θ) ,loss_functions))

function loss_function_(θ,p)
    return loss_function__(θ)
end

f = OptimizationFunction(loss_function_, GalacticOptim.AutoZygote())
prob = GalacticOptim.OptimizationProblem(f, initθ)

cb_ = function (p,l)
    println("loss: ", l , "losses: ", map(l -> l(p), loss_functions))
    return false
end

# optimizer
opt = BFGS()
res = GalacticOptim.solve(prob, opt; cb = cb_, maxiters=2000)

using Plots

ts,xs = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
u_predict_contourf = reshape([first(phi([t,x],res.minimizer)) for t in ts for x in xs] ,length(xs),length(ts))
plot(ts, xs, u_predict_contourf, linetype=:contourf,title = "predict")

u_predict = [[first(phi([t,x],res.minimizer)) for x in xs] for t in ts ]
p1= plot(xs, u_predict[3],title = "t = 0.1");
p2= plot(xs, u_predict[11],title = "t = 0.5");
p3= plot(xs, u_predict[end],title = "t = 1");
p4= plot(xs, sol.u[end],label="Numerical, t=1")
plot(p1,p2,p3,p4)


# This code tests kep_drift_ell_hyp!

@testset "kepler_drift" begin

function test_kepler_drift(dlnq::BigFloat,drift_first::Bool)
# Call as: save,jac_num,jacobian=test_kepler_drift(1e-4,true)
# This routine runs a test of the kep_elliptic_jacobian function in kepler_solver.jl
# Define the central force constant in terms of AU and days:
k = (2.0*pi/365.25)^2
# Initial position at 1 AU:
x0 = [0.01,1.0,0.01]
# Circular velocity:
r0 = sqrt(x0[1]*x0[1]+x0[2]*x0[2]+x0[3]*x0[3])
vcirc = sqrt(k/r0)
# Define initial velocity at apastron:
v0 = [.9*vcirc,0.01*vcirc,0.01*vcirc]  # The eccentricity is about ~2(1-v0/vcirc).
h = 18.0 # 18-day timesteps
hbig = big(h)

s0::Float64 = 0.0
x = zeros(Float64,3)
v = zeros(Float64,3)
s::Float64 = 0.0
t = 0.0
nsteps = 1
xsave=zeros(Float64,12,nsteps)  # First is time (1,:); next three are position (2-4,:); next three velocity (5-7,:); then r (8,:), drdt (9,:); then beta (10,:); finally s & ds (11:12,:)
# Create a variable to store the state at the end of a step:
state=zeros(Float64,12)
state_diffbig=zeros(BigFloat,12)

xsave[1,1]=0.0
for jj=1:3
  xsave[jj+1,1]=x0[jj]
  xsave[jj+4,1]=v0[jj]
end
# Save beta:
xsave[8,1]=r0
beta0 = 2.0*k/r0-dot(v0,v0)
xsave[10,1]=beta0
jacobian=zeros(Float64,7,7)
iter = kep_drift_ell_hyp!(x0,v0,k,h,s0,state,jacobian,drift_first)

# Now, do finite differences at higher precision:
kbig = big(k); s0big = big(0.0);  statebig = big.(state)
x0big = big.(x0); v0big = big.(v0)
r0big = sqrt(x0big[1]*x0big[1]+x0big[2]*x0big[2]+x0big[3]*x0big[3])
beta0big = 2*kbig/r0big-dot(v0big,v0big)
jacobian_big =zeros(BigFloat,7,7)
iter = kep_drift_ell_hyp!(x0big,v0big,kbig,hbig,s0big,statebig,jacobian_big,drift_first)
jac_frac = jacobian./convert(Array{Float64,2},jacobian_big).-1.0
println("Fractional Jacobian difference: ",maxabs(jac_frac[.~isnan.(jac_frac)]))

# Now, compute the Jacobian numerically:
jac_num = zeros(BigFloat,7,7)
# jac_num[i,j]: derivative of (x_i,v_i,k) with respect to (x_{0,j},v_{0,j},k):
for j=1:3
  x0save = copy(x0big)
  dq = dlnq * x0big[j]
  if x0big[j] != 0.0
    x0big[j] +=  dq
  else
    dq = dlnq
    x0big[j] = dq
  end
  # Recompute quantities:
  r0big = sqrt(x0big[1]*x0big[1]+x0big[2]*x0big[2]+x0big[3]*x0big[3])
  beta0big = 2*kbig/r0big-dot(v0big,v0big)
  iter = kep_drift_ell_hyp!(x0big,v0big,kbig,hbig,s0big,state_diffbig,drift_first)
  x0big = copy(x0save)
  for i=1:3
    jac_num[  i,  j] = (state_diffbig[1+i]-statebig[1+i])/dq
    jac_num[3+i,  j] = (state_diffbig[4+i]-statebig[4+i])/dq
  end
  v0save = copy(v0big)
  dq = dlnq * v0big[j]
  if v0big[j] != 0.0
    v0big[j] +=  dq
  else
    dq = dlnq
    v0big[j] = dq
  end
  r0big = sqrt(x0big[1]*x0big[1]+x0big[2]*x0big[2]+x0big[3]*x0big[3])
  beta0big = 2*kbig/r0big-dot(v0big,v0big)
  iter = kep_drift_ell_hyp!(x0big,v0big,kbig,hbig,s0big,state_diffbig,drift_first)
  v0big = copy(v0save)
  for i=1:3
    jac_num[  i,3+j] = (state_diffbig[1+i]-statebig[1+i])/dq
    jac_num[3+i,3+j] = (state_diffbig[4+i]-statebig[4+i])/dq
  end
  # Now vary mass:
  ksave = copy(kbig)
  dq = kbig*dlnq
  kbig += dq
  r0big = sqrt(x0big[1]*x0big[1]+x0big[2]*x0big[2]+x0big[3]*x0big[3])
  beta0big = 2*kbig/r0big-dot(v0big,v0big)
  iter = kep_drift_ell_hyp!(x0big,v0big,kbig,hbig,s0big,state_diffbig,drift_first)
  for i=1:3
    jac_num[  i,7] = (state_diffbig[1+i]-statebig[1+i])/dq
    jac_num[3+i,7] = (state_diffbig[4+i]-statebig[4+i])/dq
  end
  kbig = copy(ksave)
  jac_num[7,  7] = 1.0
end
println("Maximum jac_big-jac_num: ",maxabs(convert(Array{Float64,2},jacobian_big-jac_num)))
#println(convert(Array{Float64,2},jacobian_big))
#println(convert(Array{Float64,2},jacobian_big./jac_num.-1.0))
return xsave,jac_num,jacobian
end

# First try:
xsave,jac_num1,jacobian=test_kepler_drift(big(1e-20),true)
# Second try:
xsave,jac_num1,jacobian=test_kepler_drift(big(1e-20),false)

emax = 0.0; imax = 0; jmax = 0
for i=1:7, j=1:7
  if jacobian[i,j] != 0.0
    diff = abs(convert(Float64,jac_num1[i,j])/jacobian[i,j]-1.0)
    if  diff > emax
      emax = diff; imax = i; jmax = j
    end
  end
end
println("Maximum fractional error: ",emax," ",imax," ",jmax)
println("Maximum jacobian-jac_num: ",maxabs(jacobian-convert(Array{Float64,2},jac_num1)))

@test isapprox(jacobian,jac_num1;norm=maxabs)
end


using SimpleSolvers
using SimpleSolvers: LinesearchState, StaticState
using Test

include("optimizers_problems.jl")

f(x) = x^2-1
g(x) = 2x

function F(x)
    # sum(x.^2) - 1
    y = - one(eltype(x))
    for _x in x
        y += _x^2
    end
    return y
end

@testset "$(rpad("Bracketing",80))" begin
    @test bracket_minimum(x -> x^2) == (-SimpleSolvers.DEFAULT_BRACKETING_s, +SimpleSolvers.DEFAULT_BRACKETING_s)
    @test bracket_minimum(x -> (x-1)^2) == (0.64, 2.56)
end

@testset "$(rpad("Static",80))" begin
    # xᵤ = 3.
    x₀ = [-3.]
    x₁ = [+1.]
    δx = x₁ - x₀
    x  = copy(x₀)

    o  = MultivariateObjective(f, x)
    ls = StaticState()

    @test ls == LinesearchState(Static(), o)
    @test ls == LinesearchState(Static(), f, x)
    @test ls == LinesearchState(Static(), f, x; D = g)
    @test ls == LinesearchState(Static(1.0), o)
    @test ls == LinesearchState(Static(1.0), f, x)
    @test ls == LinesearchState(Static(1.0), f, x; D = g)

    @test ls == StaticState(o)
    @test ls == StaticState(f, x)
    @test ls == StaticState(f, x; D = g)

    x1 = copy(x₀); x2 = copy(x₀); @test solve!(x1, δx, ls) == ls(x2, δx) == x₁

    o1  = MultivariateObjective(f, x)
    o2  = MultivariateObjective(f, g, x)
    ls1 = Linesearch(x, o1; algorithm = Static())
    ls2 = Linesearch(x, o2; algorithm = Static())
    ls3 = Linesearch(x, o1; algorithm = Static(1.0))

    # @test ls1 == Linesearch(x, F; algorithm = Static())
    # @test ls2 == Linesearch(x, F; algorithm = Static(), D = D)

    x1 = copy(x₀); x2 = copy(x₀); @test solve!(x1, δx, ls1) == ls1(x2, δx) == x₁
    x1 = copy(x₀); x2 = copy(x₀); @test solve!(x1, δx, ls2) == ls2(x2, δx) == x₁
    x1 = copy(x₀); x2 = copy(x₀); @test solve!(x1, δx, ls3) == ls3(x2, δx) == x₁

end

@testset "$(rpad("Bisection",80))" begin

    x₀ = -3.0
    x₁ = +0.5
    δx = x₁ .- x₀
    x  = copy(x₀)

    x1 = bisection(f, x₀, x₁; config = Options(x_abstol = zero(x)))
    x2 = bisection(f, x; config = Options(x_abstol = zero(x)))

    @test x1 ≈ -1  atol=∛(2eps())
    @test x2 ≈ -1  atol=∛(2eps())

    x1 = bisection(f, x₀, x₁; config = Options(f_abstol = zero(f(x))))
    x2 = bisection(f, x; config = Options(f_abstol = zero(f(x))))

    @test x1 ≈ -1  atol=2eps()
    @test x2 ≈ -1  atol=2eps()

    # x₀ = [-3.0]
    # x₁ = [+0.5]
    # xₛ = [-1.0]
    # δx = x₁ .- x₀
    # x  = copy(x₀)

    # x_abstol = zero(eltype(x))
    # f_abstol = zero(F(x))

    # ls = Linesearch(x, F; algorithm = Bisection(), config = Options(x_abstol = x_abstol))

    # x1 = copy(x₀)
    # x2 = copy(x₀)
    # x3 = copy(x₀)

    # ls(x1, δx)
    # solve!(x2, δx, ls)
    # solve!(x3, δx, ls.state)

    # @test x1 ≈ xₛ  atol=∛(2eps())
    # @test x2 ≈ xₛ  atol=∛(2eps())
    # @test x3 ≈ xₛ  atol=∛(2eps())

    # ls = Linesearch(x, F; algorithm = Bisection(), config = Options(f_abstol = f_abstol))

    # x1 = copy(x₀)
    # x2 = copy(x₀)
    # x3 = copy(x₀)

    # ls(x1, δx)
    # solve!(x2, δx, ls)
    # solve!(x3, δx, ls.state)

    # @test x1 ≈ xₛ  atol=2eps()
    # @test x2 ≈ xₛ  atol=2eps()
    # @test x3 ≈ xₛ  atol=2eps()

end

@testset "$(rpad("Backtracking",80))" begin

    x₀ = -3.0
    x₁ = +3.0
    xₛ =  0.0
    δx = x₁ .- x₀

    _f = α -> f(x₀ + α * δx)
    _d = α -> g(x₀ + α * δx)

    ls = Linesearch(x₀, _f; D = _d, algorithm = Backtracking(), config = Options(x_abstol = zero(x₀)))

    α1 = ls()
    # solve!(x2, δx, ls)
    α2 = backtracking(_f, x₀; config = Options(x_abstol = zero(x₀)))
    α3 = backtracking(_f, x₀; D = _d, config = Options(x_abstol = zero(x₀)))

    x1 = x₀ + α1 * δx
    x2 = x₀ + α2 * δx
    x3 = x₀ + α3 * δx

    @test x1 ≈ xₛ  atol=∛(2eps())
    @test x2 ≈ xₛ  atol=∛(2eps())
    @test x3 ≈ xₛ  atol=∛(2eps())

    # x₀ = [-3.0]
    # x₁ = [+3.0]
    # xₛ = [ 0.0]
    # δx = x₁ .- x₀
    # x  = copy(x₀)

    # x1 = copy(x₀)
    # x2 = copy(x₀)
    # x3 = copy(x₀)

    # ls = Linesearch(x, F; algorithm = Backtracking(), config = Options(x_abstol = zero(eltype(x))))

    # ls(x1, δx)
    # solve!(x2, δx, ls)
    # backtracking(F, x3, δx; config = Options(x_abstol = zero(eltype(x))))

    # @test x1 ≈ xₛ  atol=∛(2eps())
    # @test x2 ≈ xₛ  atol=∛(2eps())
    # @test x3 ≈ xₛ  atol=∛(2eps())

    # n = 1
    # x₀ = -0.5*ones(n)
    # x₁ = +1.0*ones(n)
    # x  = copy(x₀)
    # f  = zeros(n)
    # g  = zeros(n,n)
    # ls = Solver(F!, x, f)

    # F!(f,x)
    # J!(g,x)

    # solve!(x, f, g, x₀, x₁, ls)

    # F!(f,x)

    # @test x ≈ zero(x) atol=4E-1
    # @test f ≈ zero(f) atol=8E-2

    # @test solve!(x, f, g, x₀, x₁, ls) == ls(x, f, g, x₀, x₁) == SFunc(F!, x, f, g, x₀, x₁)

    # n = 3
    # x₀ = -ones(n)
    # x₁ = +ones(n)
    # x  = copy(x₀)
    # f  = zeros(n)
    # g  = zeros(n,n)
    # ls = Solver(F!, x, f)

    # F!(f,x)
    # J!(g,x)

    # solve!(x, f, g, x₀, x₁, ls)

    # F!(f,x)

    # @test x == zero(x)
    # @test f == zero(f)

    # @test solve!(x, f, g, x₀, x₁, ls) == ls(x, f, g, x₀, x₁) == SFunc(F!, x, f, g, x₀, x₁)

end

@testset "$(rpad("Polynomial Linesearch",80))" begin

    x₀ = -3.0
    x₁ = +3.0
    xₛ =  0.0
    δx = x₁ .- x₀

    _f = α -> f(x₀ + α * δx)
    _d = α -> g(x₀ + α * δx)

    ls = Linesearch(x₀, _f; D = _d, algorithm = Quadratic(), config = Options(x_abstol = zero(x₀)), σ₀=0.5, σ₁=0.8)

    α1 = ls()
    # solve!(x2, δx, ls)
    α2 = quadratic(_f, x₀; config = Options(x_abstol = zero(x₀)), σ₀=0.5, σ₁=0.8)
    α3 = quadratic(_f, x₀; D = _d, config = Options(x_abstol = zero(x₀)), σ₀=0.5, σ₁=0.8)

    x1 = x₀ + α1 * δx
    x2 = x₀ + α2 * δx
    x3 = x₀ + α3 * δx

    @test x1 ≈ xₛ  atol=∛(2eps())
    @test x2 ≈ xₛ  atol=∛(2eps())
    @test x3 ≈ xₛ  atol=∛(2eps())

    # x₀ = [-3.0]
    # x₁ = [+3.0]
    # xₛ = [ 0.0]
    # δx = x₁ .- x₀
    # x  = copy(x₀)

    # x1 = copy(x₀)
    # x2 = copy(x₀)
    # x3 = copy(x₀)

    # ls = Linesearch(x, F; algorithm = Quadratic(), config = Options(x_abstol = zero(eltype(x))))

    # ls(x1, δx)
    # solve!(x2, δx, ls)
    # quadratic(F, x3, δx; config = Options(x_abstol = zero(eltype(x))))

    # @test x1 ≈ xₛ  atol=∛(2eps())
    # @test x2 ≈ xₛ  atol=∛(2eps())
    # @test x3 ≈ xₛ  atol=∛(2eps())

    # n = 1
    # x₀ = -0.5*ones(n)
    # x₁ = +1.0*ones(n)
    # x  = copy(x₀)
    # f  = zeros(n)
    # g  = zeros(n,n)
    # ls = Solver(F!, x, f)

    # F!(f,x)
    # J!(g,x)

    # solve!(x, f, g, x₀, x₁, ls)

    # F!(f,x)

    # @test x ≈ zero(x) atol=4E-1
    # @test f ≈ zero(f) atol=8E-2

    # @test solve!(x, f, g, x₀, x₁, ls) == ls(x, f, g, x₀, x₁) == SFunc(F!, x, f, g, x₀, x₁)

    # n = 3
    # x₀ = -ones(n)
    # x₁ = +ones(n)
    # x  = copy(x₀)
    # f  = zeros(n)
    # g  = zeros(n,n)
    # ls = Solver(F!, x, f)

    # F!(f,x)
    # J!(g,x)

    # solve!(x, f, g, x₀, x₁, ls)

    # F!(f,x)

    # @test x == zero(x)
    # @test f == zero(f)

    # @test solve!(x, f, g, x₀, x₁, ls) == ls(x, f, g, x₀, x₁) == SFunc(F!, x, f, g, x₀, x₁)

end


using Test
using ForestBiometrics
using DelimitedFiles

datapath = joinpath(@__DIR__, "data")
dat=readdlm(joinpath(datapath,"IEsubset_Data_CSV.csv"),',')

#select columns from example data
dbh=dat[2:35,10] #dbh; ignore header
spp=dat[2:35,9] #species; ignore header

# 2 parameter Wyckoff coefficients. Default values for IE variant of FVS
FVS_IE=Dict{String,Array{Float64}}(
"WP"=>[5.19988	-9.26718],
"WL"=>[4.97407	-6.78347],
"DF"=>[4.81519	-7.29306],
"GF"=>[5.00233	-8.19365],
"WH"=>[4.97331	-8.19730],
"RC"=>[4.89564	-8.39057],
"LP"=>[4.62171	-5.32481],
"ES"=>[4.9219	-8.30289],
"AF"=>[4.76537	-7.61062],
"PP"=>[4.9288	-9.32795],
"MH"=>[4.77951	-9.31743],
"WB"=>[4.97407	-6.78347],
"LM"=>[4.19200	-5.16510],
"LL"=>[4.76537	-7.61062],
"PI"=>[3.20000	-5.00000],
"RM"=>[3.20000	-5.00000],
"PY"=>[4.19200	-5.16510],
"AS"=>[4.44210	-6.54050],
"CO"=>[4.44210	-6.54050],
"MM"=>[4.44210	-6.54050],
"PB"=>[4.44210	-6.54050],
"OH"=>[4.44210	-6.54050],
"OS"=>[4.77951	-9.31743] )

wyckoff_test=[
100.3704728
61.35043145
11.31948442
11.99184897
71.56773647
75.48849752
4.662859967
38.1318738
86.57926895
89.09331657
83.29702964
82.02991682
59.36359424
92.24130804
80.79924609
95.99589413
89.52703552
73.04225223
95.88342114
59.38502782
4.662859967
4.662859967
97.69440543
96.76176795
106.4984871
75.82474076
100.3704728
105.4603241
92.79168745
78.84582366
79.6182448
77.22661356
92.79168745
86.208065]

#Wyckoff=(x,b)->4.5+exp(b[1]+(b[2]/(x+1)))
p=HeightDiameter(Wyckoff,FVS_IE)
wyckoff_out= calculate_height(p,dbh,spp)

@test isapprox(wyckoff_out, wyckoff_test)

#end test 1

user_eq_test=[
5.50000110437066
5.50000077727861
5.50000025616043
5.50000027222188
5.50000085504705
5.50000088557174
5.49997579696978
5.50000059480051
5.50000097615539
5.50000099792679
5.50000094851992
5.50002825335774
5.50002185801763
5.50003169790446
5.50000092800344
5.50003314012292
5.50000100174185
5.50000086646161
5.50003309510479
5.50000076242647
5.49997579696978
5.49997579696978
5.50003383536612
5.50003344996446
5.50000117082406
5.50000088821741
5.50000110437066
5.50000115899896
5.50003190202163
5.50002728005371
5.50002751279364
5.50002679851972
5.50003190202163
5.50002959423972]

p2=HeightDiameter((x,b)->4.5+x^b[1]^b[2],FVS_IE)

user_eq_out = calculate_height(p2,dbh,spp)

@test isapprox(user_eq_out, user_eq_test)
#end test 2


function test_zeros(T=Float64, dims=(13, 45))
  a1, b1 = zeros(T, dims), zeros(T, dims)
  OceanTurb.@zeros T dims a2 b2
  eltype(a1) == T && a1 == a2 && b1 == b2
end

function test_run_until(stepper=:ForwardEuler)
    model = Diffusion.Model(N=3, H=1, K=1.0, stepper=stepper)

    dt, tfinal = 0.5, 1.3
    run_until!(model, dt, tfinal)

    return time(model) == tfinal
end

function test_diffusive_flux()
    model = Diffusion.Model(N=4, H=2, K=0.1)
    c0(z) = z
    model.solution.c = c0
    flux = OceanTurb.diffusive_flux(:c, model)
    return flux[2] == -0.1
end

@testset "Utils" begin
    @test test_zeros(Float64)
    @test test_zeros(Float32)
    @test test_diffusive_flux()
    for stepper in steppers
        @test test_run_until(stepper)
    end
end


# using Base.Test
function testdata(T_)
    # srand(1)

    n,m      = 2,1
    At_      = [0.95 0.1; 0 0.95]
    Bt_      = reshape([0.2; 1],2,1)
    u        = randn(1,T_)
    x        = zeros(n,T_)
    for t = 1:T_-1
        if t == T_÷2
            At_ = [0.5 0.05; 0 0.5]
        end
        x[:,t+1] = At_*x[:,t] + Bt_*u[:,t] + 0.2randn(n)
    end
    xm = x + 0.2randn(size(x));
    x',xm',u',n,m
end

T = 1000
x,xm,u,n,m = testdata(T)

nc = 2

model = LPVSS(x, u, nc; normalize=true, λ = 1e-3)
xh = model(x,u)

eRMS = √(mean((xh[1:end-1,:]-x[2:end,:]).^2))
println("RMS error: ", eRMS)
@test eRMS <= 0.37

# using Plots
# plotlyjs(show=true)
# plot(xh[1:end-1,:], lab="Prediction", c=:red, layout=(2,1))
# plot!(x[2:end,:], lab="True", c=:blue)

nc = 4
v = 1:T
model = LPVSS(x, u, v, nc; normalize=true, λ = 1e-3)
xh = model(x,u,v)

eRMS = √(mean((xh[1:end-1,:]-x[2:end,:]).^2))
println("RMS error: ", eRMS)
@test eRMS <= 0.28

# using Plots
# plotlyjs(show=true)
# plot(xh[1:end-1,:], lab="Prediction", c=:red, layout=(2,1))
# plot!(x[2:end,:], lab="True", c=:blue)


function LShape(domain::LShape, parent::ArbField)
    LShape(parent)
end

function Base.show(io::IO, domain::LShape)
    print(io, "L-shaped domain")
end

function area(domain::LShape)
    domain.parent(3)
end

function boundary_parameterization(t, domain::LShape, i::Integer)
    1 <= i <= 4 || throw(ErrorException("give 1 ≤ i ≤ 4 not i = $i"))

    if i == 1
        v = SVector(1, 0)
        w = SVector(1, 1)
    elseif i == 2
        v = SVector(1, 1)
        w = SVector(-1, 1)
    elseif i == 3
        v = SVector(-1, 1)
        w = SVector(-1, -1)
    elseif i == 4
        v = SVector(-1, -1)
        w = SVector(0, -1)
    end

    xy = v .+ t .* (w .- v)

    θ = atan(xy[2], xy[1])
    if typeof(t) == arb_series
        θ = ifelse(θ[0] < 0, θ + 2domain.parent(π), θ)
    else
        θ = ifelse(θ < 0, θ + 2domain.parent(π), θ)
    end
    (sqrt(xy[1]^2 + xy[2]^2), θ)
end

function boundary_points(
    domain::LShape,
    u::AbstractEigenfunction,
    ::Integer,
    n::Integer;
    distribution = :linear,
)
    points = Vector{Tuple{arb,arb}}(undef, n)

    for i = 1:n
        θ = 1.5 * domain.parent(π) * i / (n + 1)
        r = 1 / max(abs(sin(θ)), abs(cos(θ)))

        points[i] = (r, θ)
    end

    return points, fill(nothing, n)
end

function interior_points(domain::LShape, n::Integer; rng = MersenneTwister(42))
    # The points are generated by randomly picking a point in the
    # square [-1, 1]×[-1, 1] and keeping it if it is inside the
    # domain.
    points = Vector{Tuple{arb,arb}}(undef, n)

    for i = 1:n
        x = 2rand(rng) - 1
        y = 2rand(rng) - 1

        while x >= 0 && y <= 0
            x = 2rand(rng) - 1
            y = 2rand(rng) - 1
        end

        x, y = domain.parent(x), domain.parent(y)
        θ = atan(y, x)
        θ = ifelse(θ < 0, θ + 2domain.parent(π), θ)
        points[i] = (hypot(x, y), θ)
    end

    points
end


@testset "Takacs B&B vertical single large segment" begin

pressurecorrelation = BeggsAndBrill
p_initial = 346.6 - 14.65
dp_est = 100
t_avg = 107.2
md_initial = 0
md_end = 3700
tvd_initial = 0
tvd_end = 3700
inclination = 0
id = 2.441
roughness = 0.0003
q_o = 375
q_w = 0
GLR = 480
APIoil = 41.06
sg_water = 1
sg_gas = 0.916
molFracCO2 = 0
molFracH2S = 0
pseudocrit_pressure_correlation = HankinsonWithWichertPseudoCriticalPressure
pseudocrit_temp_correlation = HankinsonWithWichertPseudoCriticalTemp
Z_correlation = PapayZFactor
gas_viscosity_correlation = LeeGasViscosity
solutionGORcorrelation = StandingSolutionGOR
bubblepoint = 3000
oilVolumeFactor_correlation = StandingOilVolumeFactor
waterVolumeFactor_correlation = GouldWaterVolumeFactor
dead_oil_viscosity_correlation = GlasoDeadOilViscosity
live_oil_viscosity_correlation = ChewAndConnallySaturatedOilViscosity
frictionfactor = SerghideFrictionFactor
error_tolerance = 1.0 #psi

P_pc = pseudocrit_pressure_correlation(sg_gas, molFracCO2, molFracH2S)
    _, T_pc, _ = pseudocrit_temp_correlation(sg_gas, molFracCO2, molFracH2S)

ΔP_est = PressureDrop.calculate_pressuresegment(pressurecorrelation, p_initial, dp_est, t_avg,
                                md_end - md_initial, tvd_end - tvd_initial, inclination, true, id, roughness,
                                q_o, q_w, GLR, GLR, APIoil, sg_water, sg_gas,
                                Z_correlation, P_pc, T_pc,
                                gas_viscosity_correlation, solutionGORcorrelation, bubblepoint, oilVolumeFactor_correlation, waterVolumeFactor_correlation,
                                dead_oil_viscosity_correlation, live_oil_viscosity_correlation, frictionfactor, error_tolerance)

@test ΔP_est ≈ 3770 * (0.170 + 0.312) / 2 atol = 100 #order of magnitude test based on average between example points in Takacs (50).

end #testset Takacs B&B vertical

@testset "IHS Cleveland 6 - B&B full wellbore" begin

#%% end to end test

testpath = joinpath(dirname(dirname(pathof(PressureDrop))), "test/testdata/Cleveland_6/Test_survey_Cleveland_6.csv")

testwell = read_survey(path = testpath, id_included = false, maxdepth = 10000, id = 2.441)
test_temp = collect(range(85, stop = 160, length = length(testwell.md)))

pressure_values = traverse_topdown(wellbore = testwell, roughness = 0.0006, temperatureprofile = test_temp,
                                    pressurecorrelation = BeggsAndBrill, dp_est = 10, error_tolerance = 0.1,
                                    q_o = 400, q_w = 500, GLR = 2000, APIoil = 36, sg_water = 1.05, sg_gas = 0.75,
                                    WHP = 150 - 14.65)

ihs_data = joinpath(dirname(dirname(pathof(PressureDrop))), "test/testdata/Cleveland_6/Perform_results_cleveland_6_long.csv")
ihs_pressures = [parse.(Float64, split(line, ',', keepempty = false)) for line in readlines(ihs_data)[2:end]] |>
                        x -> hcat(x...)'

@test (ihs_pressures[end, 2] - 14.65) ≈ pressure_values[end] atol = 25

model = WellModel(wellbore = testwell, roughness = 0.0006, temperatureprofile = test_temp,
                    pressurecorrelation = BeggsAndBrill, dp_est = 10, error_tolerance = 0.1,
                    q_o = 400, q_w = 500, GLR = 2000, APIoil = 36, sg_water = 1.05, sg_gas = 0.75,
                    WHP = 150 - 14.65)

@test (ihs_pressures[end, 2] - 14.65) ≈ traverse_topdown(model)[end] atol = 25

#= view results
ihs_temps = "C:/pressuredrop.git/test/testdata/Cleveland_6/Perform_temps_cleveland_6_long.csv"
ihs_temps = [parse.(Float64, split(line, ',', keepempty = false)) for line in readlines(ihs_temps)[2:end]] |>
                        x -> hcat(x...)' ;

using Gadfly
set_default_plot_size(8.5inch, 11inch)

plot(   layer(x = pressure_values, y = testwell.md, Geom.line),
        layer(x = test_temp, y = testwell.md, Geom.line, Theme(default_color = "red")),
        layer(x = ihs_pressures[:,2], y = ihs_pressures[:,1], Geom.line, Theme(default_color = "green")),
        layer(x = ihs_temps[:,2], y = ihs_temps[:,1], Geom.line, Theme(default_color = "orange")),
        Coord.cartesian(yflip = true)) #matches

[pressure_values[i] - pressure_values[i-1] for i in 2:length(pressure_values)] #negative drops only occur where inclination is negative
=#

end #testset IHS Cleveland 6 - B&B


# Wilcoxon.jl
# Wilcoxon rank sum (Mann-Whitney U) tests
#
# Copyright (C) 2012   Simon Kornblith
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

export MannWhitneyUTest, ExactMannWhitneyUTest, ApproximateMannWhitneyUTest

## COMMON MANN-WHITNEY U

# Automatic exact/normal selection
"""
    MannWhitneyUTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform a Mann-Whitney U test of the null hypothesis that the probability that an
observation drawn from the same population as `x` is greater than an observation drawn
from the same population as `y` is equal to the probability that an observation drawn
from the same population as `y` is greater than an observation drawn from the same
population as `x` against the alternative hypothesis that these probabilities are not
equal.

The Mann-Whitney U test is sometimes known as the Wilcoxon rank-sum test.

When there are no tied ranks and ≤50 samples, or tied ranks and ≤10 samples,
`MannWhitneyUTest` performs an exact Mann-Whitney U test. In all other cases,
`MannWhitneyUTest` performs an approximate Mann-Whitney U test. Behavior may be further
controlled by using [`ExactMannWhitneyUTest`](@ref) or [`ApproximateMannWhitneyUTest`](@ref)
directly.

Implements: [`pvalue`](@ref)
"""
function MannWhitneyUTest(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real}
    (U, ranks, tieadj, nx, ny, median) = mwustats(x, y)
    if nx + ny <= 10 || (nx + ny <= 50 && tieadj == 0)
        ExactMannWhitneyUTest(U, ranks, tieadj, nx, ny, median)
    else
        ApproximateMannWhitneyUTest(U, ranks, tieadj, nx, ny, median)
    end
end

# Get U, ranks, and tie adjustment for Mann-Whitney U test
function mwustats(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real}
    nx = length(x)
    ny = length(y)
    if nx <= ny
        (ranks, tieadj) = tiedrank_adj([x; y])
        U = sum(ranks[1:nx]) - nx*(nx+1)/2
    else
        (ranks, tieadj) = tiedrank_adj([y; x])
        U = nx*ny - sum(ranks[1:ny]) + ny*(ny+1)/2
    end
    (U, ranks, tieadj, nx, ny, median(x)-median(y))
end

## EXACT MANN-WHITNEY U TEST
struct ExactMannWhitneyUTest{T<:Real} <: HypothesisTest
    U::Float64              # test statistic: Mann-Whitney-U statistic
    ranks::Vector{Float64}  # ranks
    tie_adjustment::Float64 # adjustment for ties
    nx::Int                 # number of observations
    ny::Int
    median::T               # sample median
end

"""
    ExactMannWhitneyUTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform an exact Mann-Whitney U test of the null hypothesis that the probability that an
observation drawn from the same population as `x` is greater than an observation drawn
from the same population as `y` is equal to the probability that an observation drawn
from the same population as `y` is greater than an observation drawn from the same
population as `x` against the alternative hypothesis that these probabilities are not
equal.

When there are no tied ranks, the exact p-value is computed using the `pwilcox` function
from the `Rmath` package. In the presence of tied ranks, a p-value is computed by exhaustive
enumeration of permutations, which can be very slow for even moderately sized data sets.

Implements: [`pvalue`](@ref)
"""
ExactMannWhitneyUTest(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real} =
    ExactMannWhitneyUTest(mwustats(x, y)...)

testname(::ExactMannWhitneyUTest) = "Exact Mann-Whitney U test"
population_param_of_interest(x::ExactMannWhitneyUTest) = ("Location parameter (pseudomedian)", 0, x.median) # parameter of interest: name, value under h0, point estimate
default_tail(test::ExactMannWhitneyUTest) = :both

function show_params(io::IO, x::ExactMannWhitneyUTest, ident)
    println(io, ident, "number of observations in each group: ", [x.nx, x.ny])
    println(io, ident, "Mann-Whitney-U statistic:             ", x.U)
    println(io, ident, "rank sums:                            ", [sum(x.ranks[1:x.nx]), sum(x.ranks[x.nx+1:end])])
    println(io, ident, "adjustment for ties:                  ", x.tie_adjustment)
end

# Enumerate all possible Mann-Whitney U results for a given vector,
# determining left-and right-tailed p values
function mwuenumerate(x::ExactMannWhitneyUTest)
    # Get the other U if inverted by mwu_stats
    n = min(x.nx, x.ny)
    U = x.U
    if x.ny > x.nx
        U = x.nx*x.ny -x.U
    end
    le = 0
    gr = 0
    tot = 0
    k = n*(n+1)/2
    for comb in combinations(x.ranks, n)
        Up = sum(comb) - k
        tot += 1
        le += Up <= U
        gr += Up >= U
    end
    (le/tot, gr/tot)
end

function pvalue(x::ExactMannWhitneyUTest; tail=:both)
    if x.tie_adjustment == 0
        # Compute exact p-value using method from Rmath, which is fast but
        # cannot account for ties
        if tail == :both
            if x.U < x.nx * x.ny / 2
                2 * pwilcox(x.U, x.nx, x.ny, true)
            else
                2 * pwilcox(x.U - 1, x.nx, x.ny, false)
            end
        elseif tail == :left
            pwilcox(x.U, x.nx, x.ny, true)
        elseif tail == :right
            pwilcox(x.U - 1, x.nx, x.ny, false)
        else
            throw(ArgumentError("tail=$(tail) is invalid"))
        end
    else
        # Compute exact p-value by enumerating possible ranks in the tied data
        if tail == :both
            min(1, 2 * minimum(mwuenumerate(x)))
        elseif tail == :left
            mwuenumerate(x)[1]
        elseif tail == :right
            mwuenumerate(x)[2]
        else
            throw(ArgumentError("tail=$(tail) is invalid"))
        end
    end
end

struct ApproximateMannWhitneyUTest{T<:Real} <: HypothesisTest
    U::Float64              # test statistic: Mann-Whitney-U statistic
    ranks::Vector{T}        # ranks
    tie_adjustment::Float64 # adjustment for ties
    nx::Int                 # number of observations
    ny::Int
    median::Float64         # sample median
    mu::Float64             # normal approximation: mean
    sigma::Float64          # normal approximation: std
end

## APPROXIMATE MANN-WHITNEY U TEST
"""
    ApproximateMannWhitneyUTest(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Perform an approximate Mann-Whitney U test of the null hypothesis that the probability that
an observation drawn from the same population as `x` is greater than an observation drawn
from the same population as `y` is equal to the probability that an observation drawn
from the same population as `y` is greater than an observation drawn from the same
population as `x` against the alternative hypothesis that these probabilities are not
equal.

The p-value is computed using a normal approximation to the distribution of the
Mann-Whitney U statistic:
```math
    \\begin{align*}
        μ & = \\frac{n_x n_y}{2}\\\\
        σ & = \\frac{n_x n_y}{12}\\left(n_x + n_y + 1 - \\frac{a}{(n_x + n_y)(n_x +
            n_y - 1)}\\right)\\\\
        a & = \\sum_{t \\in \\mathcal{T}} t^3 - t
    \\end{align*}
```
where ``\\mathcal{T}`` is the set of the counts of tied values at each tied position.

Implements: [`pvalue`](@ref)
"""
function ApproximateMannWhitneyUTest(U::Real, ranks::AbstractVector{T},
    tie_adjustment::Real, nx::Int, ny::Int, median::Float64) where T<:Real
    mu = U - nx * ny / 2
    sigma = sqrt((nx * ny * (nx + ny + 1 - tie_adjustment /
        ((nx + ny) * (nx + ny - 1)))) / 12)
    ApproximateMannWhitneyUTest(U, ranks, tie_adjustment, nx, ny, median, mu, sigma)
end
ApproximateMannWhitneyUTest(x::AbstractVector{S}, y::AbstractVector{T}) where {S<:Real,T<:Real} =
    ApproximateMannWhitneyUTest(mwustats(x, y)...)

testname(::ApproximateMannWhitneyUTest) = "Approximate Mann-Whitney U test"
population_param_of_interest(x::ApproximateMannWhitneyUTest) = ("Location parameter (pseudomedian)", 0, x.median) # parameter of interest: name, value under h0, point estimate
default_tail(test::ApproximateMannWhitneyUTest) = :both

function show_params(io::IO, x::ApproximateMannWhitneyUTest, ident)
    println(io, ident, "number of observations in each group: ", [x.nx, x.ny])
    println(io, ident, "Mann-Whitney-U statistic:             ", x.U)
    println(io, ident, "rank sums:                            ", [sum(x.ranks[1:x.nx]), sum(x.ranks[x.nx+1:end])])
    println(io, ident, "adjustment for ties:                  ", x.tie_adjustment)
    println(io, ident, "normal approximation (μ, σ):          ", (x.mu, x.sigma))
end

function pvalue(x::ApproximateMannWhitneyUTest; tail=:both)
    if x.mu == x.sigma == 0
        1
    else
        if tail == :both
            2 * ccdf(Normal(), abs(x.mu - 0.5 * sign(x.mu))/x.sigma)
        elseif tail == :left
            cdf(Normal(), (x.mu + 0.5)/x.sigma)
        elseif tail == :right
            ccdf(Normal(), (x.mu - 0.5)/x.sigma)
        end
    end
end


module Test_dlyap

using LinearAlgebra
using MatrixEquations
using Test

@testset "Testing discrete Lyapunov equation solvers" begin

n = 10
Ty = Float64

@testset "Discrete Lyapunov equations" begin

reltol = eps(float(1000))
a = -2+im; b = 2; @time x = lyapd(a,b)
@test abs(a*x*a'-x+b) < reltol

reltol = eps(float(1000f0))
a = 1f0-2f0im; b = 2f0; @time x = lyapd(a,b)
@test abs(a*x*a'-x+b)  < reltol

try
   x = lyapd(ones(1,1),ones(1,1))
   @test false
 catch
   @test true
 end

 try
   x = lyapd([0 1;  -1 0 ],ones(2,2))
   @test false
 catch
   @test true
 end

  try
   x = lyapd([0.5 1;  0 2],ones(2,2))
   @test false
 catch
   @test true
 end

 try
   x = lyapd([1 1 1 1;  -1 1 0 1; 0 0 0.5 0.5; 0 0 -0.5 0.5],ones(4,4))
   @test false
 catch
   @test true
 end

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
c = rand(Ty,n,n)+im*rand(Ty,n,n)
qc = c'*c
Qr = real(qc)
Ty == Float64 ? reltol = eps(float(10000)) : reltol = eps(10000*n*one(Ty))

@time x = lyapd(ac,qc);
@test norm(ac*x*ac'-x+qc)/norm(x)/max(1.,norm(ac)^2) < reltol

α = 3+im; # α = 1; #SingularException
@time x = lyapd(α*I,qc);
@test norm(α*x*α'-x+qc)/norm(x) < reltol

try
   α = 1; #SingularException
   @time x = lyapd(α,qc);
catch
   @test true
end

α = 3. +im;
@time x = lyapd(α,qc);
@test norm(α*x*α'-x+qc)/norm(x) < reltol

@time x = lyapd(ac',qc);
@test norm(ac'*x*ac-x+qc)/norm(x)/max(1.,norm(ac)^2) < reltol

@time x = lyapd(ar,Qr)
@test norm(ar*x*ar'-x+Qr)/norm(x)/max(1.,norm(ar)^2) < reltol

@time x = lyapd(ar',Qr)
@test norm(ar'*x*ar-x+Qr)/norm(x)/max(1.,norm(ar)^2) < reltol

@time x = lyapd(ac,Qr);
@test norm(ac*x*ac'-x+Qr)/norm(x)/max(1.,norm(ac)^2)  < reltol

@time x = lyapd(ac',Qr);
@test norm(ac'*x*ac-x+Qr)/norm(x)/max(1.,norm(ac)^2)  < reltol

@time x = lyapd(ar,qc)
@test norm(ar*x*ar'-x+qc)/norm(x)/max(1.,norm(ar)^2) < reltol

@time x = lyapd(ar',qc)
@test norm(ar'*x*ar-x+qc)/norm(x)/max(1.,norm(ar)^2) < reltol
end
end

@testset "Discrete generalized Lyapunov equations" begin

reltol = eps(float(10000))
ac = -2+im; ec = 4+im; b = 2; @time x = lyapd(ac,ec,b)
@test abs(ac*x*ac'-ec*x*ec'+b) < reltol

try
   x = lyapd(ones(1,1),ones(1,1),ones(1,1))
   @test false
catch
   @test true
end

try
   x = lyapd([0 2;  -2 0 ],[2 0; 0 2],ones(2,2))
   @test false
catch
   @test true
end

try
   x = lyapd([2 1;  0 2],[2 0; 0 2],ones(2,2))
   @test false
catch
   @test true
end

try
   x = lyapd(2*[1 1 1 1;  -1 1 0 1; 0 0 0.5 0.5; 0 0 -0.5 0.5],2*Matrix{Float64}(I,4,4),ones(4,4))
   @test false
catch
   @test true
end

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
er = rand(Ty,n,n)
ec = er+im*rand(Ty,n,n)

c = rand(Ty,n,n)+im*rand(Ty,n,n)
qc = c'*c
Qr = real(qc)
Ty == Float64 ? reltol = eps(float(1000)) : reltol = eps(1000*n*one(Ty))

@time x = lyapd(ac,ec,qc);
@test norm(ac*x*ac'-ec*x*ec'+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

β = 3
@time x = lyapd(ac,β*I,qc);
@test norm(ac*x*ac'-β*x*β'+qc)/norm(x)/norm(ac) < reltol

β = (1+im);
@time x = lyapd(ac,β,qc);
@test norm(ac*x*ac'-β*x*β'+qc)/norm(x)/norm(ac) < reltol

α = 2+3im
@time x = lyapd(α*I,ec,qc);
@test norm(α*x*α'-ec*x*ec'+qc)/norm(x)/norm(ec) < reltol

α = 2+3im
@time x = lyapd(α,ec,qc);
@test norm(α*x*α'-ec*x*ec'+qc)/norm(x)/norm(ec) < reltol

try
   @time x = lyapd(0*ac,0*ec,qc);
   @test false
catch
   @test true
end

α = 2+3im; β = (1+im);
@time x = lyapd(α*I,β*I,qc);
@test norm(α*x*α'-β*x*β'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

try
   α = 1; β = 1; #SingularException
   @time x = lyapd(α,β,qc);
   @test false
catch
   @test true
end

α = 2+3im; β = (1+im);
@time x = lyapd(α,β,qc);
@test norm(α*x*α'-β*x*β'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

@time x = lyapd(ac',ec',qc);
@test norm(ac'*x*ac-ec'*x*ec+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

β = (1+im);
@time x = lyapd(ac',β*I,qc);
@test norm(ac'*x*ac-β*x*β'+qc)/norm(x)/norm(ac) < reltol

α = 2+3im
@time x = lyapd(α*I,ec',qc);
@test norm(α*x*α'-ec'*x*ec+qc)/norm(x)/norm(ec) < reltol

@time x = lyapd(ac',ec,qc);
@test norm(ac'*x*ac-ec*x*ec'+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac,ec',qc);
@test norm(ac*x*ac'-ec'*x*ec+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ar,er,Qr);
@test norm(ar*x*ar'-er*x*er'+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er',Qr);
@test norm(ar'*x*ar-er'*x*er+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er,Qr);
@test norm(ar'*x*ar-er*x*er'+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar,er',Qr);
@test norm(ar*x*ar'-er'*x*er+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ac,ec,Qr);
@test norm(ac*x*ac'-ec*x*ec'+Qr)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac',ec',Qr);
@test norm(ac'*x*ac-ec'*x*ec+Qr)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac',ec,Qr);
@test norm(ac'*x*ac-ec*x*ec'+Qr)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac,ec',Qr);
@test norm(ac*x*ac'-ec'*x*ec+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar,er,qc);
@test norm(ar*x*ar'-er*x*er'+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er',qc);
@test norm(ar'*x*ar-er'*x*er+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er,qc);
@test norm(ar'*x*ar-er*x*er'+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar,er',qc);
@test norm(ar*x*ar'-er'*x*er+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol
end
end

@testset "Discrete Lyapunov equations - Schur form" begin

for Ty in (Float64, Float32)

ar = rand(Ty,n,n);
ac = rand(Ty,n,n)+im*rand(Ty,n,n);
er = rand(Ty,n,n);
ec = er+im*rand(Ty,n,n);
as, es = schur(ar,er);
acs, ecs = schur(ac,ec);

c = rand(Ty,n,n)+im*rand(Ty,n,n);
qc = c'*c;
Qr = real(qc);
Ty == Float64 ? reltol = eps(float(1000)) : reltol = eps(1000*n*one(Ty))

x = copy(Qr)
@time lyapds!(as,x);
@test norm(as*x*as'+Qr-x)/norm(x)/max(norm(as)^2,1.) < reltol

x = copy(Qr)
@time lyapds!(as,x,adj=true);
@test norm(as'*x*as+Qr-x)/norm(x)/max(norm(as)^2,1.) < reltol

x = copy(qc)
@time lyapds!(acs,ecs,x);
@test norm(acs*x*acs'+qc-ecs*x*ecs')/norm(x)/max(norm(acs)^2,norm(ecs)^2) < reltol

x = copy(qc)
@time lyapds!(acs,I,x);
@test norm(acs*x*acs'+qc-x)/norm(x)/max(norm(acs)^2,norm(ecs)^2) < reltol

x = copy(qc)
@time lyapds!(acs,ecs,x,adj=true);
@test norm(acs'*x*acs+qc-ecs'*x*ecs)/norm(x)/max(norm(acs)^2,norm(ecs)^2) < reltol

x = copy(Qr)
@time lyapds!(as,es,x);
@test norm(as*x*as'+Qr-es*x*es')/norm(x)/max(norm(as)^2,norm(es)^2) < reltol

x = copy(Qr)
@time lyapds!(as,es,x,adj=true);
@test norm(as'*x*as+Qr-es'*x*es)/norm(x)/max(norm(as)^2,norm(es)^2) < reltol

end
end

end

end


n = 10000;
a = rand(Float32,(n,n))
b = rand(Float32,(n,n))
c = zeros(Float32,(n,n));
t1 = time();

for i = 1:n
  for j = 1:n
    c[i,j] = a[i,j] + b[i,j]
  end
end

t2 = time();

for i = 1:10
c = a+ b;
end

t3 = time();

dt1 = t2 - t1
dt2 = t3 - t2
n2 = n*n

mflops1 = n2 / dt1 / 1e6;
mflops2 = n2 / dt2 / 1e5;

println("times $dt1,$dt2, mflops $mflops1,$mflops2")


#  MIT License
#
#  Copyright (c) 2020 Michael J Simms. All rights reserved.
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.

module Signals

import Statistics
export smooth

# Smooths the data, which should be a list, by averaging with the given window size.
function smooth(indata, windowsize)
    outdata = []
    outdatalen = length(indata) - windowsize
    if outdatalen <= 0
        return outdata
    end

    for i in 1:outdatalen
        tempdata = indata[i:i + windowsize]
        val = Statistics.mean(tempdata)
        push!(outdata, val)
    end

    outdata
end

end


"""

```
remove_drift!(v)
```

Removes possible drift from velocities.

"""
function remove_drift!(v)
  vmean = mean(v)
  for i in 1:length(v)
    v[i] -= vmean
  end
end


using Test
using MappedArrays
using HyperbolicDiffEq
@static if Int==Int64
    # NOTE: Only on 64 bit machines due to
    # - https://github.com/JuliaPlots/Plots.jl/issues/968
    # - https://github.com/JuliaPlots/Plots.jl/issues/963
    # Fixed for julia v0.7 (https://github.com/JuliaLang/julia/pull/22644)
    using Plots; unicodeplots()
end

const godunov = GodunovFlux()

# Burgers' equation
model = Burgers()

# Rarefaction wave
prob1 = RiemannProblem(model, 0., 1.)
print(devnull, prob1)
@test prob1(0-eps()) ≈ 0
@test prob1(0+eps()) ≈ 1
sol = solve(prob1)
print(devnull, sol)
@test sol(-1) ≈ 0
@test sol( 2) ≈ 1
ξ = range(0, 1, length=10)
@test all(sol.(ξ) .== ξ)
@test sol(2., -0.5) ≈ 0
@test sol(2., 1.) ≈ 0.5
@test sol(2., 2.) ≈ 1
@static if Int==Int64
    plot(sol)
end
@test flux(sol(0), sol.prob.model) ≈ godunov(sol.prob.uₗ, sol.prob.uᵣ, sol.prob.model)

# Stationary shock wave
prob2 = RiemannProblem(model, 1., -1., 1.)
print(devnull, prob2)
@test prob2(1-eps()) ≈ 1
@test prob2(1+eps()) ≈ -1
sol = solve(prob2)
print(devnull, sol)
@test sol(-1) ≈  1
@test sol( 1) ≈ -1
@static if Int==Int64
    plot(sol)
end
@test flux(sol(0), sol.prob.model) ≈ godunov(sol.prob.uₗ, sol.prob.uᵣ, sol.prob.model)

# Tuples of Riemann problems
prob3 = RiemannProblem(model, -1., -2., 2.f0)
prob4 = RiemannProblem(model, -2., 5., 7//2)

@test_throws ErrorException prob3 * prob1
prob1 * prob2
(prob1 * prob2) * prob3
prob1 * (prob2 * prob3)
(prob1 * prob2) * (prob3 * prob4)
prob = prob1 * prob2
print(devnull, prob)

sol = solve(prob)
print(devnull, sol)
@test sol(1., -1.) ≈ 0
@test sol(1., 0.) ≈ 0
@test sol(1., 0.5) ≈ 0.5
@test sol(1., 1-eps()) ≈ 1
@test sol(1., 1+eps()) ≈ -1
@test sol(1., 2.) ≈ -1
@static if Int==Int64
    plot(sol)
end

us = range(-1, 1, length=10)
for (uₗ,uᵣ) in Iterators.product(us,us)
  prob = RiemannProblem(model, uₗ, uᵣ)
  sol = solve(prob)
  @test isapprox(godunov(uₗ, uᵣ, model), flux(sol(0), model), atol=1.e-30)
end

################################################################################

# Buckley-Leverette equation
model = BuckleyLeverette()
us = range(0, 1, length=10)
for (uₗ,uᵣ) in Iterators.product(us,us)
  prob = RiemannProblem(model, uₗ, uᵣ)
  sol = solve(prob)
  @test godunov(uₗ, uᵣ, model) ≈ flux(sol(0), model)
end

################################################################################

# Cubic conservation law
model = Cubic()
us = range(-5, 5, length=100)
for (uₗ,uᵣ) in Iterators.product(us,us)
  prob = RiemannProblem(model, uₗ, uᵣ)
  sol = solve(prob)
  @test godunov(uₗ, uᵣ, model) ≈ flux(sol(0), model)
end

################################################################################

# Shallow water equations
model = ShallowWater()

# Dam Break
# Example 5.20 of Holden & Risebro, "Front Tracking for Hyperbolic Conservation Laws", 2002
uₗ = ShallowWaterVar1D(1., 0.)
uᵣ = ShallowWaterVar1D(0., 0.)
prob = RiemannProblem(model, uₗ, uᵣ)
sol = solve(prob)
x = range(-3, 3, length=50)
u = sol.(1., x)
h = mappedarray(u->u.h, u)
v = mappedarray(u->(u.h ≈ 0 ? zero(u.hv) : u.hv/u.h), u)
@test all(h[x .< -1] .≈ 1)
@test all(h[-1 .< x .< 2] ≈ (2 .- x[-1 .< x .< 2]).^2 ./ 9)
@test all(h[x .> 2] .≈ 0)
@test all(v[x .< -1] .≈ 0)
@test all(v[-1 .< x .< 2] ≈ 2 .* (1 .+ x[-1 .< x .< 2]) ./ 3)
@test all(v[x .> 2] .≈ 0)
@static if Int==Int64
    plot(sol)
end

uₗ, uᵣ = uᵣ, uₗ
prob = RiemannProblem(model, uₗ, uᵣ)
sol = solve(prob)
x = range(-3, 3, length=50)
u = sol.(1., x)
h = mappedarray(u->u.h, u)
v = mappedarray(u->(u.h ≈ 0 ? zero(u.hv) : u.hv/u.h), u)
@test all(h[x .< -2] .≈ 0)
@test all(h[-2 .< x .< 1] ≈ (2 .+ x[-2 .< x .< 1]).^2 ./ 9)
@test all(h[x .> 1] .≈ 1)
@test all(v[x .< -2] .≈ 0)
@test all(v[-2 .< x .< 1] ≈ -2 .* (1 .- x[-2 .< x .< 1]) ./ 3)
@test all(v[x .> 1] .≈ 0)
@static if Int==Int64
    plot(sol)
end

uₗ = ShallowWaterVar1D(1., 0.)
uᵣ = ShallowWaterVar1D(eps(), 0.)
prob = RiemannProblem(model, uₗ, uᵣ)
sol = solve(prob)
x = range(3, 3, length=50)
u = sol.(1., x)
h = mappedarray(u->u.h, u)
v = mappedarray(u->(u.h ≈ 0 ? zero(u.hv) : u.hv/u.h), u)
@test all(h[x .< -1] .≈ 1)
@test all(h[-1 .< x .< 2] ≈ (2 .- x[-1 .< x .< 2]).^2 ./ 9)
@test all(h[x .> 2] .≈ eps())
@test all(v[x .< -1] .≈ 0)
@test all(v[-1 .< x .< 2] ≈ 2 .* (1 .+ x[-1 .< x .< 2]) ./ 3)
@test all(v[x .> 2] .≈ 0)
@static if Int==Int64
    plot(sol)
end

uₗ, uᵣ = uᵣ, uₗ
prob = RiemannProblem(model, uₗ, uᵣ)
sol = solve(prob)
x = range(3, 3, length=50)
u = sol.(1., x)
h = mappedarray(u->u.h, u)
v = mappedarray(u->(u.h ≈ 0 ? zero(u.hv) : u.hv/u.h), u)
@test all(h[x .< -2] .≈ eps())
@test all(h[-2 .< x .< 1] ≈ (2 .+ x[-2 .< x .< 1]).^2 ./ 9)
@test all(h[x .> 1] .≈ 1)
@test all(v[x .< -2] .≈ 0)
@test all(v[-2 .< x .< 1] ≈ -2 .* (1 .- x[-2 .< x .< 1]) ./ 3)
@test all(v[x .> 1] .≈ 0)
@static if Int==Int64
    plot(sol)
end

# Moses' First Problem
# Example 5.21 of Holden & Risebro, "Front Tracking for Hyperbolic Conservation Laws", 2002
h₀ = 1.
v₀ = 2.5
uₗ = ShallowWaterVar1D(h₀, -v₀)
uᵣ = ShallowWaterVar1D(h₀,  v₀)
prob = RiemannProblem(model, uₗ, uᵣ)
sol = solve(prob)
x = range(-1-(v₀+sqrt(h₀)), 1+v₀+sqrt(h₀), length=50)
u = sol.(1., x)
h = mappedarray(u->u.h, u)
v = mappedarray(u->(u.h ≈ 0 ? zero(u.hv) : u.hv/u.h), u)
idx1 = x .< -(v₀+sqrt(h₀))
idx2 = -(v₀+sqrt(h₀)) .< x .< (2*sqrt(h₀)-v₀)
idx3 = (2*sqrt(h₀)-v₀) .< x .< (v₀-2*sqrt(h₀))
idx4 = (v₀-2*sqrt(h₀)) .< x .< (v₀+sqrt(h₀))
idx5 = (v₀+sqrt(h₀)) .< x
@test all(h[idx1] .≈ h₀)
@test all(h[idx2] .≈ (-v₀ .+ 2 .* sqrt(h₀) .- x[idx2]).^2 ./ 9)
@test all(h[idx3] .≈ 0)
@test all(h[idx4] .≈ (v₀ .- 2 .* sqrt(h₀) .- x[idx4]).^2 ./ 9)
@test all(h[idx5] .≈ h₀)
@test all(v[idx1] .≈ -v₀)
@test all(v[idx2] .≈ (-v₀ .+ 2 .* sqrt(h₀) .+ 2 .* x[idx2]) ./ 3)
@test all(v[idx3] .≈ 0)
@test all(v[idx4] .≈ (v₀ .- 2 .* sqrt(h₀) .+ 2 .* x[idx4]) ./ 3)
@test all(v[idx5] .≈ v₀)
@static if Int==Int64
    plot(sol)
end

################################################################################

# Euler equations
model = HyperbolicDiffEq.Euler()

# intermediate values
# Toro (2009), Riemann Solvers and Numerical Methods for Fluid Dynamics,
# Table 4.3
compute_pₘ_vₘ = (ϱₗ, vₗ, pₗ, ϱᵣ, vᵣ, pᵣ) -> begin
    γ = 1.4
    HyperbolicDiffEq.compute_pₘ_vₘ(ϱₗ, vₗ, pₗ, sqrt(γ*pₗ/ϱₗ), ϱᵣ, vᵣ, pᵣ, sqrt(γ*pᵣ/ϱᵣ), γ)
end
pₘ, vₘ = compute_pₘ_vₘ(1.0, 0.0, 1.0, 0.125, 0.0, 0.1)
@test isapprox(pₘ, 0.30313, atol=1.e-5)
@test isapprox(vₘ, 0.92745, atol=1.e-5)
pₘ, vₘ = compute_pₘ_vₘ(1.0, -2.0, 0.4, 1.0, 2.0, 0.4)
@test isapprox(pₘ, 0.00189, atol=1.e-5)
@test isapprox(vₘ, 0.00000, atol=1.e-5)
pₘ, vₘ = compute_pₘ_vₘ(1.0, 0.0, 1000.0, 1.0, 0.0, 0.01)
@test isapprox(pₘ, 460.894, atol=1.e-3)
@test isapprox(vₘ, 19.5975, atol=1.e-4)
pₘ, vₘ = compute_pₘ_vₘ(1.0, 0.0, 0.01, 1.0, 0.0, 100.0)
@test isapprox(pₘ, 46.0950, atol=1.e-4)
@test isapprox(vₘ, -6.19633, atol=1.e-5)
pₘ, vₘ = compute_pₘ_vₘ(5.99924, 19.5975, 460.894, 5.99242, -6.19633, 46.0950)
@test isapprox(pₘ, 1691.64, atol=1.e-2)
@test isapprox(vₘ, 8.68975, atol=3.e-5)

# Toro (2009), Riemann Solvers and Numerical Methods for Fluid Dynamics:
# Test 1 on page 129
uₗ = conserved_variables(1.0, 0.0, 1.0, model)
uᵣ = conserved_variables(0.125, 0.0, 0.1, model)
sol = solve(RiemannProblem(model, uₗ, uᵣ))
@static if Int==Int64
    plot(sol)
end

# Test 2 on page 129
uₗ = conserved_variables(1.0, -2.0, 0.4, model)
uᵣ = conserved_variables(1.0,  2.0, 0.4, model)
sol = solve(RiemannProblem(model, uₗ, uᵣ))
@static if Int==Int64
    plot(sol)
end

# Test 3 on page 129
uₗ = conserved_variables(1.0, 0.0, 1000.0, model)
uᵣ = conserved_variables(1.0, 0.0, 0.01, model)
sol = solve(RiemannProblem(model, uₗ, uᵣ))
@static if Int==Int64
    plot(sol)
end

# Test 4 on page 129
uₗ = conserved_variables(1.0, 0.0, 0.01, model)
uᵣ = conserved_variables(1.0, 0.0, 100.0, model)
sol = solve(RiemannProblem(model, uₗ, uᵣ))
@static if Int==Int64
    plot(sol)
end

# Test 5 on page 129
uₗ = conserved_variables(5.99924, 19.5975, 460.894, model)
uᵣ = conserved_variables(5.99242, -6.19633, 46.0950, model)
sol = solve(RiemannProblem(model, uₗ, uᵣ))
@static if Int==Int64
    plot(sol)
end


# ## Two y axis plot

#md # ![FILE_NAME.png](images/FILE_NAME.png)

## pyplot()
using Measures, Random

Random.seed!(2018)
plot(randn(100), ylabel="y1", leg=:topright)
plot!(twinx(), randn(100)*10,
    c=:red,
    ylabel="y2",
    leg=:bottomright,
    size=(600, 400))
plot!(right_bottom=10mm)


using GR
x = 0:0.1:100
y = sin.(x)
xlim([0, 120])
ylim([-1.5, 1.5])
xlabel("time")
ylabel("sin-value")
title("sin-plot")
plot(x, y)


function logistic_test()
    @testset "Logistic" begin
        n = 15;
        xV = range(-2.0, 3.0, length = n);

        lb = -2.5;
        ub = 1.5;
        lrShifter = 3.0;
        yV = logistic(xV; lb = lb, ub = ub, lrShifter = lrShifter);
        @test isa(yV, Vector{Float64})
        @test size(yV) == size(xV)
        @test all(yV .> lb)  &&  all(yV .< ub)

        glf = GeneralizedLogistic(lb = lb, ub = ub, lrShifter = lrShifter);
        y2V = logistic(glf, xV);
        @test yV ≈ y2V
    end
end

@testset "Logistic" begin
    logistic_test();
end

# ----------


#=
  "Cats and rats and elephants"
  https://www.allendowney.com/blog/2018/12/11/cats-and-rats-and-elephants/
  """
  A few weeks ago I posted 'Lions and Tigers and Bears',
  [https://www.allendowney.com/blog/2018/12/03/lions-and-tigers-and-bears/]
  which poses a Bayesian problem related to the Dirichlet distribution.  If you have
  not read it, you might want to start there.

  Now here’s the follow-up question:

    Suppose there are six species that might be in a zoo: lions and tigers and bears,
    and cats and rats and elephants. Every zoo has a subset of these species, and every
    subset is equally likely.

    One day we visit a zoo and see 3 lions, 2 tigers, and one bear. Assuming that every
    animal in the zoo has an equal chance to be seen, what is the probability that the
    next animal we see is an elephant?
  """

  Note: The approach in this model is the same as in the
  "Lions and Tigers and Bears" problem, just added
    - the three new animals
    - queries which calculates subsets and number of different animals

Distributions of variable animal7
lion       =>    3755  (0.375500)
tiger      =>    3677  (0.367700)
elephant   =>    2205  (0.220500)
bear       =>     323  (0.032300)
cat        =>      21  (0.002100)
rat        =>      19  (0.001900)

Distribution of unique combinations (unsorted)
Distributions of variable (num:0)
[2, 1, 0]	=>	4362 (0.4362)
[2, 0]	=>	3645 (0.3645)
[3, 2, 1, 0]	=>	1540 (0.154)
[1, 2, 0]	=>	324 (0.0324)
[3, 1, 0]	=>	29 (0.0029)
[2, 0, 1]	=>	25 (0.0025)
[1, 0, 2]	=>	18 (0.0018)
[3, 0, 1, 2]	=>	12 (0.0012)
[3, 1, 0, 2]	=>	11 (0.0011)
[4, 0, 2]	=>	5 (0.0005)
[3, 0, 2, 1]	=>	5 (0.0005)
[4, 0, 1]	=>	4 (0.0004)
[1, 0, 3]	=>	3 (0.0003)
[6, 0]	=>	3 (0.0003)
[1, 3, 0, 2]	=>	2 (0.0002)
[4, 1, 0]	=>	2 (0.0002)
[2, 0, 3, 1]	=>	2 (0.0002)
[3, 2, 0, 1]	=>	2 (0.0002)
[1, 3, 0]	=>	2 (0.0002)
[3, 0]	=>	1 (0.0001)
[1]	=>	1 (0.0001)
[2, 0, 1, 3]	=>	1 (0.0001)
[3, 0, 1]	=>	1 (0.0001)

  Cf lions__tiger_and_bears*.jl
     ~/blog/cats_rats_and_elephants.blog
     ~/webppl/cats_rats_and_elephants.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function cats_rats_and_elephants()
    #  The animals.
    lion     = 1
    tiger    = 2
    bear     = 3
    cat      = 4
    rat      = 5
    elephant = 6

    #  Prior
    #  We draw 6 times with the multinomial distribution.
    #  What is the probability of different combinations of the number of each animal?
    alphas = [1/6, 1/6, 1/6, 1/6, 1/6, 1/6]; #  must sum to 1
    #  alphas = [2/12, 2/12, 2/12, 2/12, 2/12, 2/12];

    #  Draw 6 animals
    n = 6
    x ~ Multinomial(n,alphas)

    #  The probabilities to calculate ("aliased" for simplicity)
    probLion      = x[1]
    probTiger     = x[2]
    probBear      = x[3]
    probCat       = x[4]
    probRat       = x[5]
    probElephant  = x[6]

    #  Posterior: What is the probability of lion, tiger, and bear given the observations?
    n2 = n+1
    o = tzeros(n2)
    for i in 1:n2
        o[i] ~ Categorical(simplex([probLion,probTiger,probBear,probCat,probRat,probElephant]))
    end

    #  Observations
    true ~ Dirac(o[1] == lion)
    true ~ Dirac(o[2] == lion)
    true ~ Dirac(o[3] == lion)
    true ~ Dirac(o[4] == tiger)
    true ~ Dirac(o[5] == tiger)
    true ~ Dirac(o[6] == bear)

    animal7 ~ Dirac(o[7])

    # probbear7     ~ Dirac(o[7] == bear)
    # problion7     ~ Dirac(o[7] == lion)
    # probtiger7    ~ Dirac(o[7] == tiger)
    # probcat7      ~ Dirac(o[7] == cat)
    # probrat7      ~ Dirac(o[7] == rat)
    # probelephant7 ~ Dirac(o[7] == elephant)

    #  Which subset of the animals is the most probable?
    # xunique_ = unique(x)
    # xunique = sort(reshape(xunique_,length(xunique_),1),dims=1)
    xunique = unique(x) # unsorted

    #  Number of different animals in the zoo
    numAnimals ~ Dirac(length(xunique))

    return xunique

end

model = cats_rats_and_elephants()
num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(1000), 10_000)
# chns = sample(model, IS(), 10_000)

#
display(chns)
show_var_dist_pct(chns,:animal7,["lion","tiger","bear","cat","rat","elephant"])

println()
show_var_dist_pct(chns,:probbear7)

# println()
# show_var_dist_pct(chns,:problion7)

# println()
# show_var_dist_pct(chns,:probtiger7)

# println()
# show_var_dist_pct(chns,:probcat7)

# println()
# show_var_dist_pct(chns,:probrat7)

# println()
# show_var_dist_pct(chns,:probelephant7)

println()
show_var_dist_pct(chns,:numAnimals)

println("Distribution of unique combinations (unsorted)")
chains_params = Turing.MCMCChains.get_sections(chns, :parameters)
genq = generated_quantities(model, chains_params)
show_var_dist_pct(genq)


@info("-------[Testing Complementarity/MPEC]------------------------------------------")

include("../examples/mpec/bard1.jl")
include("../examples/mpec/dempe.jl")
include("../examples/mpec/gnash1m.jl")
include("bard1_test.jl")


"Load data from file."
function loaddata(file::String)
    return readlines(file)
end

"Take toboggan down the slope on map."
function toboggan(map::Array, go::Array)::Int
    pos = [1, 1]
    n_trees = 0
    while pos[1] < length(map)
        pos += go
        pos[2] -= div(pos[2], length(map[1]) + 1)*length(map[1])
        n_trees += (map[pos[1]][pos[2]] == '#')
    end
    return n_trees
end

function main()
    map = loaddata("../assets/data/03.txt")

    # part one
    n_trees = toboggan(map, [1, 3])
    println("Encountered ", n_trees, " trees while tobogganing.")

    # part two
    n_trees = 1
    for go in [[1, 1], [1, 3], [1, 5], [1, 7], [2, 1]]
        n_trees *= toboggan(map, go)
    end
    println("Product of trees while tobogganing is ", n_trees, ".")
end

main()


cd(@__DIR__); include("setups/grid7x3.jl");

## TSD (T = 0.1)
D = natural_eigdist(𝚽, 𝛌, Q; T = 0.1, distance = :TSD)
E = transform(fit(MDS, D, maxoutdim=2, distances=true))
# E[1, :] .*= sign(E[1, 2]); E[2, :] .*= -sign(E[1, 2])
plt = grid7x3_mds_heatmaps(E, 𝚽; backend = :pyplot)
savefig(plt, "../figs/Grid7x3_MDS_TSD_T01.png")
display(plt)


#using ClustForOpt_priv
using JuMP
using Clp
using DataFrames
 # parameters

acf = 0.1 # annualized cost of capital factor

c_cap_s = (11.94e6 + acf*2.58e9)*20/100 # $/GW   # 20/100 ~ 0.5$/W, 40/100~ 1$/W
c_cap_w = 30.98e6 + acf*1.74e9 # $/GW
c_cap_gas = 11.96e6 + acf*0.968e9 # $/GW
c_cap_nuc = 92.04e6 + acf*3.82e9 # $/GW
c_cap_coal = 28.15e6 + acf*2.22e9 # $/GW
c_cap = [c_cap_coal,c_cap_nuc,c_cap_gas,c_cap_w,c_cap_s]

c_var_s = 0.5 # $/GWh
c_var_w = 0.5 # $/GWh
c_var_gas = 41500 # $/GWh
c_var_nuc = 7025 # $/GWh
c_var_coal = 23000 # $/GWh
c_var = [c_var_s,c_var_w,c_var_gas,c_var_nuc,c_var_coal]
c_var = [c_var_coal,c_var_nuc,c_var_gas,c_var_w,c_var_s]

c_lost = 1000e6 # $/GW # VOLL

n_years = 10  # depreciation period for capital
r = 0.07 # real discount rate

 # load input data

de= readtable("/home/hteich/.julia/v0.6/ClustForOpt_priv/data/texas_merrick/demand.txt",separator=' ')[:DEM] # MW
d=reshape(de,(size(de)[1],1))
 # load growth
d=1.486*d
d=d/1000 # GW

 #d = rand(8760,1)  # demand
 # load growth *1.5 (footnote 6 merrick)
a_s= readtable("/home/hteich/.julia/v0.6/ClustForOpt_priv/data/texas_merrick/TexInsolationFactorV1.txt",separator=' ')[:solar_61]
a_s=reshape(a_s,(size(a_s)[1],1))
  a_s = a_s/1000
 #a_s = rand(8760,1) # availability solar
 #a_w = rand(8760,1) # availability wind
a_w= readtable("/home/hteich/.julia/v0.6/ClustForOpt_priv/data/texas_merrick/windfactor2.txt",separator=' ')[:Wind_61]
a_w=reshape(a_w,(size(a_w)[1],1))
a_gas = ones(8760,1)
a_nuc = ones(8760,1)
a_coal =ones(8760,1)
a= [a_coal,a_nuc,a_gas,a_w,a_s]

N=ones(8760,1)

del_t = 1 # hour

  ### SETS ####
K = size(d)[1]
T = size(d)[2]
G_names = ["coal","nuclear","gas","wind","solar"] # generators
n_G = length(G_names)

  ##### MODEL #####
m = Model(optimizer=ClpSolver())

  ### VARIABLES ##
  ### DESIGN VARIABLES ###
# generation capacity build
@variable(m,G_max[1:n_G] >= 0) # GWh

 ### OPERATIONAL VARIABLES ####
# generation of individual generator per time step
@variable(m,G[1:n_G,1:K,1:T] >=0)

 ##### OBJECTIVE FUNCTION #####
 @objective(m,Min,sum(c_cap[g]*G_max[g] for g=1:n_G)  +
     sum(sum(sum( c_var[g]*G[g,k,t]
           for k=1:K )
         for t=1:T )
       for g=1:n_G )
     )

 #### CONSTRAINTS ######
 # energy balance constraint
@constraint(m,[k=1:K,t=1:T], sum(G[g,k,t] for g=1:n_G) >= d[k,t])
 # maximum power production constraint
@constraint(m,[g=1:n_G,k=1:K,t=1:T], G[g,k,t] <= G_max[g] * a[g][k,t] )

  #### SOLVE ####
  status = solve(m)

 #### Output results ####
 G_max = getvalue(G_max)
 G = getvalue(G) # installed capacity
 G_gen =sum(G[:,:,1],2) # generated power

 if true
    using PyPlot

    figure()
    bt = zeros(length(G_max))
    for i=2:length(bt)
      bt[i]=cumsum(G_max)[i-1]
    end
    for i=1:length(G_max)
      bar(1, G_max[i],bottom=bt[i],color=cols[i],label=G_names[i])
    end
    ylabel("GW")
    legend()
    title("Capacity built")

    figure()
    bt = zeros(length(G_gen))
    for i=2:length(bt)
      bt[i]=cumsum(G_gen)[i-1]
    end
    for i=1:length(G_gen)
      bar(1, G_gen[i],bottom=bt[i],color=cols[i],label=G_names[i])
    end
    ylabel("GWh")
    legend()
    title("Power generated")

 end

 # TODO - add clustering, check on paper what to do next


srand(0)

facts("qrdelcol") do
    m = 100
    A = randn(m,m)
    Q, R = qr(A)
    for i in 100:-1:1
        k = rand(1:i)
        A = A[:,1:i .!= k]
        R = qrdelcol(R, k)
        @fact vecnorm( R'*R - A'*A ) --> less_than(1e-5)
    end
end


function mats(p)
    p.A, p.B, p.K
end
function model_from_params(p, h, ny)
    A, B, K = mats(p)
    x0      = copy(p.x0)
    sys     = StateSpaceNoise(A, B, K, h)
    sysf    = SysFilter(sys, x0, similar(x0, ny))
end

function pem_costfun(p, y, u, h, metric::F) where {F} # To ensure specialization on metric
    nu, ny = obslength(u), obslength(y)
    model = model_from_params(p, h, ny)
    return mean(metric(e) for e in prediction_errors(model, y, u))
end
function sem_costfun(p, y, u, h, metric::F) where {F} # To ensure specialization on metric
    nu, ny = obslength(u), obslength(y)
    model = model_from_params(p, h, ny)
    return mean(metric(e) for e in simulation_errors(model, y, u))
end

"""
    sys, x0, opt = pem(data; nx, kwargs...)

System identification using the prediction-error method.

# Arguments:
- `data`: iddata object containing `y` and `u`.
    - `y`: Measurements, either a matrix with time along dim 2, or a vector of vectors
    - `u`: Control signals, same structure as `y`
- `nx`: Number of poles in the estimated system. Thus number should be chosen as number of system poles plus number of poles in noise models for measurement noise and load disturbances.
- `focus`: Either `:prediction` or `:simulation`. If `:simulation` is chosen, a two stage problem is solved with prediction focus first, followed by a refinement for simulation focus.
- `metric`: A Function determining how the size of the residuals is measured, default `sse` (e'e), but any Function such as `norm`, `e->sum(abs,e)` or `e -> e'Q*e` could be used.
- `regularizer(p)=0`: function for regularization. The structure of `p` is detailed below
- `solver` Defaults to `Optim.BFGS()`
- `stabilize_predictor=true`: Modifies the estimated Kalman gain `K` in case `A-KC` is not stable by moving all unstable eigenvalues to the unit circle.
- `difficult=false`: If the identification problem appears to be difficult and ends up in a local minimum, set this flag to true to solve an initial global optimization problem to supply a good initial guess. This is expected to take some time.
- `kwargs`: additional keyword arguments are sent to `Optim.Options`.

# Return values
- `sys::StateSpaceNoise`: identified system. Can be converted to `StateSpace` by `convert(StateSpace, sys)`, but this will discard the Kalman gain matrix.
- `x0`: Estimated initial state
- `opt`: Optimization problem structure. Contains info of the result of the optimization problem

## Structure of parameter vector `p`
The parameter vector is of type [`ComponentVector`](https://github.com/jonniedie/ComponentArrays.jl) and the fields `A,B,K,x0` can be accessed as `p.A` etc. The internal storage is according to
```julia
A = size(nx,nx)
B = size(nx,nu)
K = size(nx,ny)
x0 = size(nx)
p = [A[:]; B[:]; K[:]; x0]
```
"""
function pem(
    d;
    nx,
    solver              = BFGS(),
    focus               = :prediction,
    metric              = sse,
    regularizer         = p -> 0,
    iterations          = 1000,
    stabilize_predictor = true,
    difficult           = false,
    A                   = 0.0001randn(nx, nx),
    B                   = 0.001randn(nx, obslength(input(d))),
    # C                   = 0.001randn(obslength(output(d)), nx),
    K                   = 0.001randn(nx, obslength(output(d))),
    x0                  = 0.001randn(nx),
    kwargs...,
)

    y, u = output(d), input(d)
    nu, ny = obslength(u), obslength(y)
    if size(A,1) != size(A,2) # Old API
        A = [A zeros(nx, nx-ny)]
    end
    p = ComponentVector((; A, B, K, x0))
    options = Options(; iterations = iterations, kwargs...)
    cost_pred = p -> pem_costfun(p, y, u, d.Ts, metric) + regularizer(p)
    if difficult
        stabilizer = p -> 10000 * !stabfun(d.Ts, ny)(p)
        options0 = Options(; iterations = iterations, kwargs...)
        opt = optimize(
            p -> cost_pred(p) + stabilizer(p),
            1000p,
            # ParticleSwarm(n_particles = 100length(p)),
            NelderMead(),
            options0,
        )
        p = minimizer(opt)
    end
    opt = optimize(cost_pred, p, solver, options; autodiff = :forward)
    println(opt)
    if focus == :simulation
        @info "Focusing on simulation"
        cost_sim = p -> sem_costfun(p, y, u, d.Ts, metric) + regularizer(p)
        opt = optimize(
            cost_sim,
            minimizer(opt),
            NewtonTrustRegion(),
            Options(; iterations = iterations ÷ 2, kwargs...);
            autodiff = :forward,
        )
        println(opt)
    end
    model = model_from_params(minimizer(opt), d.Ts, ny)
    if !isstable(model.sys)
        @warn("Estimated system does not have a stable prediction filter (A-KC)")
        if stabilize_predictor
            @info("Stabilizing predictor")
            # model = stabilize(model)
            model = stabilize(model, solver, options, cost_pred)
        end
    end
    isstable(ss(model.sys)) || @warn("Estimated system is not stable")

    model.sys, copy(model.state), opt
end

function stabilize(model)
    s            = model.sys
    @unpack A, K = s
    C            = s.C
    poles        = eigvals(A - K * C)
    newpoles     = map(poles) do p
        ap = abs(p)
        ap <= 1 && (return p)
        p / (ap + sqrt(eps()))
    end
    K2           = ControlSystems.acker(A', C', newpoles)' .|> real
    all(abs(p) <= 1 for p in eigvals(A - K * C)) || @warn("Failed to stabilize predictor")
    s.K .= K2
    model
end

function stabilize(model, solver, options, cfi)
    s = model.sys
    cost = function (p)
        maximum(abs.(eigvals(s.A - p * s.K * s.C))) - 0.9999999
    end
    p = fzero(cost, 1e-9, 1 - 1e-9)
    s.K .*= p
    model
end

function stabfun(h, ny)
    function (p)
        model = model_from_params(p, h, ny)
        isstable(model.sys)
    end
end


input = joinpath(@__DIR__, "input")
eqs = readlines(input)

EQ = r"(\d+)|([+])|([*])|([(])|([)])"

function parse_eqs(eqs)
    acc = []
    for eq in eqs
        matches = eachmatch(EQ, eq)
        r = map(m -> only(filter(!isnothing, m.captures)), matches)
        push!(acc, r)
    end
    return acc
end

function convert(eqs, prec)::Vector{Vector{String}}
    acc = []
    for eq in eqs
        res = postfix(eq, prec)
        push!(acc, res)
    end
    return acc
end

function postfix(eq, prec)
    postfix = []
    operators = []
    for symbol in eq
        if symbol in ["+", "*"]
            if !isempty(operators) && last(operators) in ["+", "*"] && prec[symbol] <= prec[last(operators)]
                push!(postfix, pop!(operators))
                push!(operators, symbol)
            else
                push!(operators, symbol)
            end
        elseif symbol == "("
            push!(operators, symbol)
        elseif symbol == ")"
            while (op = pop!(operators); op != "(")
                push!(postfix, op)
            end
        else
            push!(postfix, symbol)
        end
    end
    while !isempty(operators)
        op = pop!(operators)
        push!(postfix, op)
    end
    return postfix
end

function evaluate(eqs::Vector{Vector{String}})
    return map(eq -> evaluate(eq), eqs)
end

function evaluate(eq::Vector{String})
    stack = []
    for symbol in eq
        if symbol == "+"
            left = pop!(stack)
            right = pop!(stack)
            push!(stack, left + right)
        elseif symbol == "*"
            left = pop!(stack)
            right = pop!(stack)
            push!(stack, left * right)
        else
            push!(stack, parse(Int, symbol))
        end
    end
    return only(stack)
end

eqs = filter(!isempty, eqs)
eqs = parse_eqs(eqs)

p1 = sum(evaluate(convert(eqs, Dict("+" => 1, "*" => 1))))
p2 = sum(evaluate(convert(eqs, Dict("+" => 2, "*" => 1))))

println("-----------------------------------------------------------------------")
println("operation order -- part one :: $p1")
println("operation order -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 21993583522852)
@assert(p2 == 122438593522757)


# calculating error between two poses

using Plots
pyplot()

function calc_lon_lat_error(pose::Array, pose_ref::Array)
  error_x = pose[1] - pose_ref[1]
  error_y = pose[2] - pose_ref[2]
  error_theta = pose[3] - pose_ref[3]
  error_lon = cos(pose_ref[3])*error_x + sin(pose_ref[3])*error_y
  error_lat = -sin(pose_ref[3])*error_x + cos(pose_ref[3])*error_y
  return [round(error_lon, digits=2), round(error_lat, digits=2),
          round(error_theta, digits=2)]
end

function draw_lon_lat_error!(pose::Array, pose_ref::Array)
  error = calc_lon_lat_error(pose, pose_ref)
  annotate!(pose_ref[1], pose_ref[2]+0.1,
            text("($(error[1]), $(error[2]), $(error[3]))",
            :left, 8))
end


struct FuzzyMatrix
    numbers::Matrix{FuzzyNumber}
end

function zeros(levels::Vector{<:Real}, dim1::Int64, dim2::Int64)
    M = Matrix{FuzzyNumber}(undef, dim1, dim2)
    Z = SingletonFuzzyNumber(levels, number=0)
    for i = 1:dim1
        for j = 1:dim2
            M[i, j] = Z
        end
    end
    M
end