using LinearAlgebra

include("gaussnum.jl")
include("ode-solver.jl")

# Define & Solve ODE

function lotkavolterra(x,θ)
    α, β, γ, δ = θ
    x₁, x₂ = x

    dx₁ = α*x₁ - β*x₁*x₂
    dx₂ = δ*x₁*x₂ - γ*x₂

    [dx₁, dx₂]
end

struct GaussODEProblem{F,T<:Tuple{Number,Number},U<:AbstractMatrix,P<:AbstractVector,I} <: AbstractODEProblem
    f::F
    tspan::T
    u0::U
    θ::P
    θ_idx::I
    size::Int
end
Base.size(prob::GaussODEProblem) = size(prob.u0,1)
statesize(prob::GaussODEProblem) = prob.size

function GaussODEProblem(f,tspan,u0::V,θ::V) where V<:AbstractVector{<:GaussNum}
    θ_idx = findall(isuncertain, θ)

    μ = vcat(mean.(u0), mean.(θ[θ_idx]))
    d = length(μ)
    Σ = vcat(var.(u0), var.(θ[θ_idx])) |> Diagonal |> Matrix
    Σ½ = cholesky(Σ).L
    Qp = sqrt(d)*[I(d) -I(d)]
    Xp = μ .+ Σ½*Qp

    GaussODEProblem(f,tspan,Xp,mean.(θ),θ_idx,length(u0))
end

function setmean!(prob::GaussODEProblem,xi)
    prob.θ[prob.θ_idx] .= xi[(statesize(prob)+1):end]
    xi[1:statesize(prob)]
end

struct GaussODESolver{S<:ODESolver} <: ODESolver
    solver::S
end

function (s::GaussODESolver)(prob::GaussODEProblem, Xp, t)
    Xp = reduce(hcat, map(eachcol(Xp)) do xi
        ui = setmean!(prob, xi)
        ui = s.solver(prob, ui, t)[1]
        xi[1:statesize(prob)] .= ui
        xi
    end)
    Xp, t+s.solver.dt
end

function filter(prob::GaussODEProblem, solver::GaussODESolver, data, σy)
    t = prob.tspan[1]; u = prob.u0
    us = [u]; ts = [t]

    for (ty, y) in data
        while t <= ty
            (u,t) = solver(prob, u, t)
            push!(us,u)
            push!(ts,t)
        end

        yₚ = u[1,:] # prediction. we measure only the first variable
        μᵧ = mean(yₚ)
        μᵤ = mean(u, dims=2)

        Σᵧᵧ = cov(yₚ, corrected=false) + σy
        Σᵤᵤ = cov(u, dims=2, corrected=false)

        Σᵤᵧ = cov(u, yₚ, dims=2, corrected=false)
        K   = Σᵤᵧ*inv(Σᵧᵧ)

        μ = vec(μᵤ + K*(y .- μᵧ))
        Σ = Hermitian(Σᵤᵤ - K*Σᵤᵧ')

        d = size(prob)
        Qp = sqrt(d)*[I(d) -I(d)]
        u = μ .+ cholesky(Σ).L * Qp
    end
    ts, us
end

function lotkavolterra(x,θ)
    α, β, γ, δ = θ
    x₁, x₂ = x

    dx₁ = α*x₁ - β*x₁*x₂
    dx₂ = δ*x₁*x₂ - γ*x₂

    [dx₁, dx₂]
end

θ = [0.1±0.01, 0.2, 0.3, 0.2]
u0 = [1.0±0.2, 1.0±0.2]
tspan = (0., 100.)
prob = GaussODEProblem(lotkavolterra,tspan,u0,θ)
solver = GaussODESolver(RK2(0.1))

using JLD2
raw_data = load("../../lecture_12/lotkadata.jld2")
ts = Float32.(raw_data["t"][1:100:end])
ys = Float32.(raw_data["u"][:,1:100:end])
data = [(t,y[1]) for (t,y) in zip(ts,eachcol(ys))]
t, us = solve(prob,solver)
t, us = filter(prob, solver, data, 0.01)

gus = map(us) do u
    u = u[1:statesize(prob),:]
    μ = mean(u,dims=2) |> vec
    σ = std(u,dims=2) |> vec
    GaussNum.(μ,σ)
end
gus = reduce(hcat,gus)
p1 = plot(t, gus[1,:], lw=3)
plot!(p1, t, gus[2,:], lw=3)
plot!(p1, raw_data["t"], raw_data["u"][1,:], lw=3, alpha=0.5) |> display

solver = RK2(0.1)
prob = ODEProblem(lotkavolterra,tspan,u0,θ)

function Base.rand(x::AbstractVector{<:GaussNum{T}}) where T
    mean.(x) .+ std.(x) .* randn(T,length(x))
end
Base.rand(prob::ODEProblem) = ODEProblem(prob.f, prob.tspan, rand(prob.u0), rand(prob.θ))

p2 = plot()
Us = map(1:200) do i
    t, us = solve(rand(prob),solver.solver)
    us = reduce(hcat,us)
end

plot!(p2, t, GaussNum.(mean(Us)[1,:],std(Us)[1,:]), lw=3)
plot!(p2, t, GaussNum.(mean(Us)[2,:],std(Us)[2,:]), lw=3)
display(plot(p1,p2))


using Ipopt
using PowerModelsPrivacyPreserving
using PowerModels
using Distributions
using Random

const PMs = PowerModels
const PMPP = PowerModelsPrivacyPreserving

ipopt = Ipopt.Optimizer
file =  "test/data/matpower/case5.m"
data_unpert = parse_file(file)
data_min_loss = deepcopy(data_unpert)
data_min_cost = deepcopy(data_unpert)

"this is the canonical OPF problem with generation cost minimization"
result_unpert_cost = PMPP.run_ac_opf_cost(data_unpert, ipopt)
PMs.print_summary(result_unpert_cost["solution"])
PMPP.calculate_losses!(result_unpert_cost, data_unpert)
"store faithfulness info"
data_min_cost["cost"] = Dict()
data_min_cost["cost"]["value"] = result_unpert_cost["objective"]
data_min_cost["cost"]["beta"] = 1

"this variant of the OPF problem minimizes grid losses instead of generation cost"
result_unpert_loss = PMPP.run_ac_opf_loss(data_unpert, ipopt)
PMs.print_summary(result_unpert_loss["solution"])
PMPP.calculate_losses!(result_unpert_loss, data_unpert)
"store faithfulness info"
data_min_loss["loss"] = Dict()
data_min_loss["loss"]["value"] = result_unpert_loss["totalloss"]
data_min_loss["loss"]["beta"] = 1

# Set constant parameters
α = 0.01
ϵ = 1
λ = 50

# Add impedance perturbation to both data dictionaries.
data_pert_min_loss = PMPP.create_impedance_perturbation(data_min_loss, α, ϵ, λ)
data_pert_min_cost = PMPP.create_impedance_perturbation(data_min_cost, α, ϵ, λ)

# Calculate the result
result_pert_loss = PMPP.run_opf_variable_impedance_loss(data_pert_min_loss, ipopt)
PMs.print_summary(result_pert_loss["solution"])
PMPP.calculate_losses!(result_pert_loss, data_pert_min_loss)

result_pert_cost = PMPP.run_opf_variable_impedance_cost(data_pert_min_cost, ipopt)
PMs.print_summary(result_pert_cost["solution"])
PMPP.calculate_losses!(result_pert_cost, data_pert_min_cost)


function ispalindrome(a)
  s = string(a)
  mid = Int(floor(length(s)/ 2))
  r = true
  for j = 1:mid
    if s[j] != s[end-(j-1)]
      r = false
      break
    end
  end
  return r
end

function istwo3digitprod(a)
  r = false
  for i = 100:999
    if a%i == 0 && 100 <= a/i <= 999
      r = true
      break
    end
  end
  return r
end

for i = 999*999:-1:100*100
  if ispalindrome(i) && istwo3digitprod(i)
    println(i)
    break
  end
end


using Symbolics
using Symbolics: symbol_to_poly, poly_to_symbol

using Test

@variables x y z

syms = [
    [1], [1, BigInt(2)], [x], [x^2 + 2], [x + (5//8)y],
    [1 - x*y*z], [(x + y + z)^5], [0, BigInt(0)y^2, Rational(1)z^3],
    [x, sin((44//31)y)*z]
]
for sym in syms
    polynoms, pvar2sym, sym2term = Symbolics.symbol_to_poly(sym)
    sym2 = Symbolics.poly_to_symbol(polynoms, pvar2sym, sym2term, Real)
    @test isequal(expand.(sym2), expand.(sym))
end

@test isequal(expand.(groebner_basis([x, y])), [y, x])
@test isequal(expand.(groebner_basis([y, x])), [y, x])
@test isequal(expand.(groebner_basis([x])), [x])
@test isequal(expand.(groebner_basis([x, x^2])), [x])
@test isequal(expand.(groebner_basis([BigInt(1)x + 2//3])), [x  + 2//3])

@variables x1 x2 x3 x4
@test isequal(expand.(groebner_basis([x1, x, y])), [y, x1, x])

# input unchanged
f1 = [x2, x1, x4, x3]
f2 = [x2, x1, x4, x3]
groebner_basis(f1)
@test isequal(expand.(f1), expand.(f2))

@variables x1 x2 x3 x4 x5
system = [
   x1 + x2 + x3 + x4 + x5,
   x1*x2 + x1*x3 + x1*x4 + x1*x5 + x2*x3 + x2*x4 + x2*x5 + x3*x4 + x3*x5 + x4*x5,
   x1*x2*x3 + x1*x2*x4 + x1*x2*x5 + x1*x3*x4 + x1*x3*x5 + x1*x4*x5 + x2*x3*x4 + x2*x3*x5 + x2*x4*x5 + x3*x4*x5,
   x1*x2*x3*x4 + x1*x2*x3*x5 + x1*x2*x4*x5 + x1*x3*x4*x5 + x2*x3*x4*x5,
   x1*x2*x3*x4*x5 - 1
]
truebasis = [
        x1 + x2 + x3 + x4 + x5,
        x2^2 + x2*x3 + x2*x4 + x2*x5 + x3^2 + x3*x4 + x3*x5 + x4^2 + x4*x5 + x5^2,
        x3^3 + x3*(x4^2) + x3*(x5^2) + x4^3 + x4*(x3^2) + x5*(x4^2) + x4*(x5^2) + x5^3 + x5*(x3^2) + x3*x4*x5,
        x4^4 + x4*(x5^3) + x5^4 + x5*(x4^3) + (x4^2)*(x5^2),
        x5^5 - 1
]
basis = expand.(groebner_basis(system))
@test isequal(basis, truebasis)

N = 45671930739135174346839766056203605080877915151
system = [
    x1 + x2 + x3 + x4,
    x1*x2 + x1*x3 + x1*x4 + x2*x3 + x2*x4 + x3*x4,
    x1*x2*x3 + x1*x2*x4 + x1*x3*x4 + x2*x3*x4,
    x1*x2*x3*x4 + N
]
truebasis = [
    x1 + x2 + x3 + x4,
    x2^2 + x2*x3 + x3^2 + x2*x4 + x3*x4 + x4^2,
    x3^3 + x3^2*x4 + x3*x4^2 + x4^3,
    x4^4 - N
]
basis = groebner_basis(system)
@test isequal(expand.(basis), truebasis)

# Groebner does not yet work with constant ideals
@test_broken groebner_basis([1])


module NonlinearPDE

# Write your package code here.
    import LinearAlgebra: Diagonal, I
    include("truth_solution.jl")
    export u, f, g
    include("kernel_functions.jl")
    export kappa, Delta_x_kappa, Delta_x_y_kappa
    include("generate_sampled_points.jl")
    export sample_points
    include("assembly_Theta.jl")
    export assembly_Theta
    include("pde_solver.jl")
    export PDE_solver
    include("extend_solution.jl")
    export construct_inpTheta, extend_solution, get_extended_solution
    include("error_calculation.jl")
    export solution_error
end


#==========================================================
This routine initializes an optimization problem on a
single qubit with 2 energy levels (and no guard states)
where the analytical solution is a constant control
function, i.e. a Rabi oscillator. The drift Hamiltonian in
the rotating frame is
              H_0 = - 0.5*ξ_a(a^†a^†aa),
where a is the annihilation operator for the qubit. Here
the control Hamiltonian includes the usual symmetric and
anti-symmetric terms
    H_{sym} = p(t)(a + a^†),    H_{asym} = q(t)(a - a^†),
which come from the rotating frame approximation and hence
we refer to these as "coupled" controls. For this
example we evolve the state forward a full period T=2π. The
parameters for this example are:
                ω_a =  2π × 0.0       Grad/s,
                ξ_a =  2π × 2(0.1099) Grad/s.
We use the usual Bsplines (no carrier waves) in this
example.
==========================================================#
using LinearAlgebra
using Plots
#pyplot()
using FFTW
using DelimitedFiles
using Printf
using Ipopt
using Random

Base.show(io::IO, f::Float64) = @printf(io, "%20.13e", f)

using Juqbox

N = 2

Nguard = 0
Ntot = N + Nguard

samplerate = 32 # default number of time steps per unit time
casename = "rabi" # base file name (used in optimize-once.jl)

# frequencies (in GHz, will be multiplied by 2*pi to get angular frequencies in the Hamiltonian matrix)
fa = 5.0
xa = 2* 0.1099
rot_freq = [fa] # Rotational frequencies

# period of oscillation
Tperiod = 100.0 # ns # 2*pi

# duration
T = Tperiod # Tperiod/4

utarget = Matrix{ComplexF64}(I, Ntot, N)
vtarget = Matrix{ComplexF64}(I, Ntot, N)

# Rabi target (one period)

theta = pi/4 # phase angle
aOmega = pi/Tperiod
Omega = (cos(theta) + 1im*sin(theta)) * aOmega
println("Amplitude |Omega| = ", aOmega, " phase angle theta = ", theta)

# unitary target matrix
utarget[1,1] = cos(aOmega*T)
utarget[2,1] = -(sin(theta) + 1im*cos(theta))*sin(aOmega*T)
utarget[1,2] = (sin(theta) - 1im*cos(theta))*sin(aOmega*T)
utarget[2,2] = cos(aOmega*T)

omega1 = Juqbox.setup_rotmatrices([N], [Nguard], rot_freq)

# Compute Ra*utarget
rot1 = Diagonal(exp.(im*omega1*T))

# target in the rotating frame
# vtarget = rot1*utarget

# target in lab frame
vtarget = utarget

# setup ansatz for control functions
use_bcarrier = true # new Bcarrier allows a constant control function

Nctrl = 1
Nosc = Nctrl
Nfreq = 1 # number of carrier frequencies

Random.seed!(2456)
# initial parameter guess

# setup carrier frequencies
om = zeros(Nctrl,Nfreq)
# Note: same frequencies for each p(t) (x-drive) and q(t) (y-drive)
println("Carrier frequencies [GHz]: ", om[:,:]./(2*pi))

# setup drift Hamiltonian
number = Diagonal(collect(0:Ntot-1))

H0 = 2*pi* (fa*number - 0.5*xa* (number*number - number) )
println("Drift Hamiltonian/2*pi: ", H0)

# lowering matrix
amat = Bidiagonal(zeros(Ntot),sqrt.(collect(1:Ntot-1)),:U) # standard lowering operator matrix
# raising matrix
adag = transpose(amat) # raising operator matrix

# max parameter amplitude
maxpar = 1.0*aOmega/Nfreq

# dense matrices
Hunc_ops=[Array(amat + adag)]
H0 = Array(H0)

# Estimate time step
Pmin = 100
nsteps = calculate_timestep(T, H0, Hunc_ops, [maxpar], Pmin)
println("Duration = ", T, " # time steps per min-period, P = ", Pmin, " # time steps: ", nsteps)

# Initial conditions
Ident = Matrix{Float64}(I, Ntot, Ntot)
U0 = Ident[1:Ntot,1:N]

# setup the simulation parameters
params = Juqbox.objparams([N], [Nguard], T, nsteps, Uinit=U0, Utarget=vtarget, Cfreq=om, Rfreq=rot_freq,
                          Hconst=H0, Hunc_ops=Hunc_ops)

# setup the initial parameter vector, either randomized or from file
startFromScratch = false # true
startFile = "drives/rabi-pcof-opt-t100.jld2"

if startFromScratch
    # D1 smaller than 3 does not work
    D1 = 3 # Number of B-spline coefficients per frequency, sin/cos and real/imag
    nCoeff = 2*Nosc*Nfreq*D1 # factor '2' is for sin/cos
    pcof0  = zeros(nCoeff)

    if Nfreq == 1 # analytical solution
        pcof0[1:D1]  .= aOmega*cos(theta) # real coefficients
        pcof0[D1+1:2*D1]  .= aOmega*sin(theta) # imag coefficients
        println("*** Starting from constant pcof with (Re/Im) amplitudes: ", aOmega*cos(theta), " ", aOmega*sin(theta) )
    else
        pcof0 = (rand(nCoeff) .- 0.5).*maxpar*0.1
        println("*** Starting from random pcof with amplitude ", maxpar*0.1)
    end
else
    # the data on the startfile must be consistent with the setup!
    # use if you want to have initial coefficients read from file
    pcof0 = read_pcof(startFile)
    nCoeff = length(pcof0)
    D1 = div(nCoeff, 2*Nosc*Nfreq) # factor '2' is for sin/cos
    nCoeff = 2*Nosc*Nfreq*D1 # just to be safe if the file doesn't contain the right number of elements
    println("*** Starting from B-spline coefficients in file: ", startFile)
end

# min and max coefficient values
minCoeff = -maxpar*ones(nCoeff);
maxCoeff = maxpar*ones(nCoeff);

# For ipopt
maxIter = 150 # optional argument
lbfgsMax = 250 # optional argument

println("*** Settings ***")
println("System Hamiltonian coefficients: (fa, xa) =  ", fa, xa)
println("Total number of states, Ntot = ", Ntot, " Total number of guard states, Nguard = ", Nguard)
if use_bcarrier
  println("Using B-spline basis functions with carrier wave, # freq = ", Nfreq)
else
  println("Using regular B-spline basis functions")
end
println("Number of coefficients per spline = ", D1, " Total number of parameters = ", nCoeff)
println("Max parameter amplitudes: maxpar = ", maxpar)
println("Tikhonov coefficients: tik0 = ", params.tik0)

# Allocate all working arrays
wa = Juqbox.Working_Arrays(params, nCoeff)
prob = Juqbox.setup_ipopt_problem(params, wa, nCoeff, minCoeff, maxCoeff, maxIter=maxIter, lbfgsMax=lbfgsMax, startFromScratch=startFromScratch)

# uncomment to run the gradient checker for the initial pcof
# addOption( prob, "derivative_test", "first-order"); # for testing the gradient

println("Initial coefficient vector stored in 'pcof0'")

# evaluate objective function
objf, uhist, trfid = traceobjgrad(pcof0, params, wa, true, false);
println("Trace fidelity: ", trfid);


export alphaFilter, alphaSimplex, delaunayTriangulation
#===============================================================================
#
#	src/alpha_complex.jl
#
#	This File Contains:
#
#	 - alphaFilter(
#			V::Matrix,
#			DT = Array{Int64,1}[];
#			digits=64
#		)::DataStructures.SortedMultiDict{}
#
#	 - alphaSimplex(
#			V::Matrix,
#			filtration::DataStructures.SortedMultiDict{},
#			α_threshold::Float64
#		)::Array{Array{Array{Int64,1},1},1}
#
#	 - delaunayTriangulation(
#			V::Matrix
#		)::Array{Array{Int64,1},1}
#
===============================================================================#

"""
	alphaFilter(
		V::Matrix, DT = Array{Int64,1}[];
		digits=64
	)::DataStructures.SortedMultiDict{}

Return ordered collection of pairs `(alpha charatteristic, complex)`.

This method evaluates the `α`-filter over the sites `S`.
If a Delaunay Triangulation `DT` is not specified than it is evaluated
via `AlphaStructures.delaunayTriangulation()`.

# Examples
```jldoctest
julia> V = [1. 2. 1. 2.; 0. 0. 1. 2. ];

julia> AlphaStructures.alphaFilter(V)
SortedMultiDict(Base.Order.ForwardOrdering(),
	0.0 => [1],
	0.0 => [2],
	0.0 => [3],
	0.0 => [4],
	0.5 => [1, 2],
	0.5 => [1, 3],
	0.70710678118655 => [2, 3],
	0.70710678118655 => [3, 4],
	0.70710678118655 => [1, 2, 3],
	1.0 => [2, 4],
	1.0 => [2, 3, 4]
)

```
"""
function alphaFilter(
		V::Matrix,
		DT = Array{Int64,1}[];
		digits=64
	)::DataStructures.SortedDict{}

	dim = size(V, 1)
	filtration = DataStructures.SortedDict{Array{Int64,1},Float64}()

	# 1 - Each point => alpha_char = 0.
	for i = 1 : size(V, 2)
		insert!(filtration, [i], 0.)
	end

	# 2 - Delaunay triangulation of ``V``
	if isempty(DT)
		DT = AlphaStructures.delaunayTriangulation(V)
	end

	n_upsimplex = length(DT)

	# 3 - process all upper simplex
	ind = 1
	for upper_simplex in DT
		if ind % 500000 == 0
			println(ind," simplices processed of ", n_upsimplex)
		end
		AlphaStructures.processuppersimplex(V,upper_simplex,filtration; digits = digits)
		ind = ind + 1
	end

	return filtration
end

"""
	processuppersimplex(
		V::Matrix,
		up_simplex::Array{Int64,1},
		filtration::DataStructures.SortedDict{};
		digits=64
		)

Process the upper simplex.
"""
function processuppersimplex(
		V::Matrix,
		up_simplex::Array{Int64,1},
		filtration::DataStructures.SortedDict{};
		digits=64
		)

	α_char = AlphaStructures.findRadius(V[:, up_simplex], digits=digits);
	insert!(filtration, up_simplex, α_char)

	d = length(up_simplex)-1
	if d > 1
		# It gives back combinations in natural order
		newsimplex = collect(Combinatorics.combinations(up_simplex,d))
		for lowsimplex in newsimplex
			AlphaStructures.processlowsimplex(V, up_simplex, lowsimplex, filtration; digits=digits)
		end
	end
end

"""
	processlowsimplex(
		V::Matrix,
		up_simplex::Array{Int64,1},
		lowsimplex::Array{Int64,1},
		filtration::DataStructures.SortedDict{};
		digits=64
		)

Process the lower simplex knowing the upper.
"""
function processlowsimplex(
	V::Matrix,
	up_simplex::Array{Int64,1},
	lowsimplex::Array{Int64,1},
	filtration::DataStructures.SortedDict{};
	digits=64)

	α_char = AlphaStructures.findRadius(V[:, lowsimplex], digits=digits)
	point = V[:, setdiff(up_simplex, lowsimplex)]

	if AlphaStructures.vertexInCircumball(V[:, lowsimplex], α_char, point)
		filtration[lowsimplex] = filtration[up_simplex]

	elseif !haskey(filtration, lowsimplex)
		filtration[lowsimplex] = α_char

	end

	d = length(lowsimplex)-1
	if d > 1
		# It gives back combinations in natural order
		newsimplex = collect(Combinatorics.combinations(lowsimplex,d))
		for simplex in newsimplex
			 AlphaStructures.processlowsimplex(V, lowsimplex, simplex, filtration, digits=digits)
		end
	end
end

#-------------------------------------------------------------------------------

"""
	alphaSimplex(
		V::Matrix,
		filtration::DataStructures.SortedMultiDict{},
		α_threshold::Float64
	)::Array{Array{Array{Int64,1},1},1}

Return collection of all `d`-simplex, for `d ∈ [0,dimension]`,
	with characteristic α less than a given value `α_threshold`.
"""
function alphaSimplex(
		V::Matrix,
		filtration::DataStructures.SortedDict{},
		α_threshold::Float64
	)::Array{Array{Array{Int64,1},1},1}

	dim = size(V, 1)
	# [VV, EV, FV, ...]
	simplexCollection = [ Array{Array{Int64,1},1}() for i = 1 : dim+1 ]

	for (k, v) in filtration
        if v <= α_threshold
        	push!(simplexCollection[length(k)], k)
    	end
    end

	sort!.(simplexCollection)

	return simplexCollection
end


using OdinSon
using Distributions

# Graphics[Table[Circle[RandomReal[10, 2]], {100}]]
f1 = Canvas([Circle(rand(Uniform(0, 10), 2), 1, style=Style(fill=nothing)) for i = 1:100])
render(f1) # the render in the Mathematica notebook is implicit
f2 =Canvas(mapslices(p->Circle(p, 1, style=Style(fill=nothing)), rand(Uniform(0, 10), 100, 2), 2))
render(f2)
# So very verbose in comparison. Thow if instead we have like Mathematica implicit defaults
# which is very close, but I don't like the defaults as they don't feel natural to me,
# save for maybe having the style not fill by default
#=
```
Graphics[Table[Circle[RandomReal[10, 2]], {100}]]
Canvas([Circle(rand(Uniform(0, 10), 2)) for i = 1:100])
```
=#


using Dates

using Pumas, GlobalSensitivity, CairoMakie, PumasPlots

peg_inf_model = @model begin

  @param begin
    logKa   ∈  RealDomain()
    logKe   ∈  RealDomain()
    logVd   ∈  RealDomain()
    logn    ∈  RealDomain()
    logd    ∈  RealDomain()
    logc    ∈  RealDomain()
    logEC50 ∈  RealDomain()
    ω²Ka   ∈ RealDomain(lower=0.0)
    ω²Ke   ∈ RealDomain(lower=0.0)
    ω²Vd   ∈ RealDomain(lower=0.0)
    ω²n    ∈ RealDomain(lower=0.0)
    ω²δ    ∈ RealDomain(lower=0.0)
    ω²c    ∈ RealDomain(lower=0.0)
    ω²EC50 ∈ RealDomain(lower=0.0)
    σ²PK ∈ RealDomain(lower=0.0)
    σ²PD ∈ RealDomain(lower=0.0)
  end

  @random begin
    ηKa   ~ Normal(0.0, sqrt(ω²Ka))
    ηKe   ~ Normal(0.0, sqrt(ω²Ke))
    ηVd   ~ Normal(0.0, sqrt(ω²Vd))
    ηn    ~ Normal(0.0, sqrt(ω²n))
    ηδ    ~ Normal(0.0, sqrt(ω²δ))
    ηc    ~ Normal(0.0, sqrt(ω²c))
    ηEC50 ~ Normal(0.0, sqrt(ω²EC50))
  end

  @pre begin
    # constants
    p = 100.0
    d = 0.001
    e = 1e-7
    s = 20000.0
    log_Ka   = logKa   + ηKa
    log_Ke   = logKe   + ηKe
    log_Vd   = logVd   + ηVd
    log_n    = logn    + ηn
    log_d    = logd    + ηδ
    log_c    = logc    + ηc
    log_EC50 = logEC50 + ηEC50
  end

  @init begin
    T = exp(log_c + log_d)/(p*e)
    I = (s*e*p - d*exp(log_c + log_d))/(p*exp(log_d)*e)
    W = (s*e*p - d*exp(log_c + log_d))/(exp(log_c + log_d)*e)
  end

  @dynamics begin
    X' = -exp(log_Ka)*X
    A' = exp(log_Ka)*X - exp(log_Ke)*A
    T' = s - T*(e*W + d)
    I' = e*W*T - exp(log_d)*I
    W' = p/((A/exp(log_Vd)/exp(log_EC50) + 1e-100)^exp(log_n) + 1)*I - exp(log_c)*W
  end

  @derived begin
    conc   = @. A/exp(log_Vd)
    log10W = @. log10(max.(W, 1e-5))
    yPK ~ @. Normal(A/exp(log_Vd), sqrt(σ²PK))
    yPD ~ @. Normal(log10W, sqrt(σ²PD))
    nca1 := @nca conc
    cmax_pk = NCA.cmax(nca1)
    nca2 := @nca log10W
    cmax_pd = NCA.cmax(nca2)
  end

end

peg_inf_dr = DosageRegimen(180.0)
t = collect(0.0:1.0:28.0)
_pop = map(i -> Subject(id=i, observations=(yPK=[], yPD=[]), events=peg_inf_dr, time=t), 1:3)

param_PKPD = (
  logKa   = log(0.80),
  logKe   = log(0.15),
  logVd   = log(100.0),
  logn    = log(2.0),
  logd    = log(0.20),
  logc    = log(7.0),
  logEC50 = log(0.12),
  # random effects variance parameters, set to 0 to neutralize effect of random effects
  ω²Ka   = 0.0,
  ω²Ke   = 0.0,
  ω²Vd   = 0.0,
  ω²n    = 0.0,
  ω²δ    = 0.0,
  ω²c    = 0.0,
  ω²EC50 = 0.0,
  # variance parameter in proportional error model
  σ²PK = 0.04,
  σ²PD = 0.04)

simdata = simobs(peg_inf_model, _pop, param_PKPD, ensemblealg=EnsembleThreads())
sim_plot(peg_inf_model ,simdata, observations=[:yPD,:yPK], all = true)

p_range_low =  (logKa   = log(0.80)/2,
                logKe   = log(0.15)/2,
                logVd   = log(100.0)/2,
                logn    = log(2.0)/2,
                logd    = log(0.20)/2,
                logc    = log(7.0)/2,
                logEC50 = log(0.12)/2, )

p_range_high = (logKa   = log(0.80)*2,
                logKe   = log(0.15)*2,
                logVd   = log(100.0)*2,
                logn    = log(2.0)*2,
                logd    = log(0.20)*2,
                logc    = log(7.0)*2,
                logEC50 = log(0.12)*2, )

morris_ = Pumas.gsa(
  peg_inf_model,
  _pop,param_PKPD,
  GlobalSensitivity.Morris(),
  [:cmax_pk,:cmax_pd],
  p_range_low,
  p_range_high,
  ensemblealg=EnsembleThreads())

keys_ = [keys(p_range_low)...]
cmax_pk_meansstar = [morris_.means_star[1,:][key] for key in keys_]
cmax_pd_meansstar = [morris_.means_star[2,:][key] for key in keys_]

fig = Figure(resolution = (1200, 400))
plot_pk = barplot(fig[1,1], string.(keys_), cmax_pk_meansstar, axis = (xlabel ="Parameters", title="Cmax PK Morris means"))
plot_pd = barplot(fig[1,2], string.(keys_), cmax_pd_meansstar, axis = (xlabel ="Parameters", title="Cmax PD Morris means"))
display(fig)

sobol_ = Pumas.gsa(
  peg_inf_model,
  _pop,
  param_PKPD,
  GlobalSensitivity.Sobol(),
  [:cmax_pk,:cmax_pd],
  p_range_low,
  p_range_high,
  N=4000,
  ensemblealg=EnsembleThreads())

cmax_pk_s1 = [sobol_.first_order[1,:][key] for key in keys_]
cmax_pd_s1 = [sobol_.first_order[2,:][key] for key in keys_]
cmax_s1 = hcat(cmax_pk_s1, cmax_pd_s1)

fig = Figure(resolution = (1200, 400))
plot_s1 = heatmap(fig[1,1], 1:2, 1:7, cmax_s1, axis = (xticks = (1:2, ["Cmax_PK","Cmax_PD"]), yticks = (1:7, string.(keys_)), title="First Order Indices"), colormap = "darkrainbow")

cmax_pk_st = [sobol_.total_order[1,:][key] for key in keys_]
cmax_pd_st = [sobol_.total_order[2,:][key] for key in keys_]
cmax_st = hcat(cmax_pk_st, cmax_pd_st)

plot_st = heatmap(fig[1,2], 1:2, 1:7, cmax_st,  axis = (xticks = (1:2, ["Cmax_PK","Cmax_PD"]), yticks = (1:7, string.(keys_)), title="First Order Indices"), colormap = "darkrainbow")
display(fig)


#==
Functions related to the Booker quartic equation and the four waves associated with the
quartic roots.

Besides calculating the quartic roots themselves, this file also includes functions
related to calculating wavefields and vacuum reflection coefficients from the quartic roots.
==#

"""
    bookerquartic(ea::EigenAngle, M)
    bookerquartic(T::TMatrix)

Compute roots `q` and the coefficients `B` of the Booker quartic described by the
susceptibility tensor `M` or `T` matrix.

# References

[Budden1988]: K. G. Budden, “The propagation of radio waves: the theory of radio
    waves of low power in the ionosphere and magnetosphere,” First paperback edition.
    New York: Cambridge University Press, 1988.
"""
bookerquartic

function bookerquartic(ea::EigenAngle, M)
    S, C, C² = ea.sinθ, ea.cosθ, ea.cos²θ

    # Precompute
    M11p1 = 1 + M[1,1]
    M33p1 = 1 + M[3,3]

    C²pM22 = C² + M[2,2]
    C²pM33 = C² + M[3,3]

    M13pM31 = M[1,3] + M[3,1]

    M12M23 = M[1,2]*M[2,3]
    M13M31 = M[1,3]*M[3,1]
    M21M32 = M[2,1]*M[3,2]
    M23M32 = M[2,3]*M[3,2]

    # Booker quartic coefficients
    B4 = M33p1
    B3 = S*M13pM31
    B2 = -C²pM33*M11p1 + M13M31 - M33p1*C²pM22 + M23M32
    B1 = S*(M12M23 + M21M32 - C²pM22*M13pM31)
    B0 = M11p1*C²pM22*C²pM33 +
         M12M23*M[3,1] + M[1,3]*M21M32 -
         M13M31*C²pM22 - M11p1*M23M32 -
         M[1,2]*M[2,1]*C²pM33

    booker_quartic_coeffs = MVector{5}(B0, B1, B2, B3, B4)
    booker_quartic_roots = MVector{4,eltype(booker_quartic_coeffs)}(0, 0, 0, 0)

    # `roots!` dominates this functions runtime
    roots!(booker_quartic_roots, booker_quartic_coeffs, NaN, 4, false)

    return booker_quartic_roots, SVector(booker_quartic_coeffs)
end

function bookerquartic(T::TMatrix)
    # Precompute
    T34T42 = T[3,4]*T[4,2]

    # This is the depressed form of the quartic
    b3 = -T[1,1] - T[4,4]
    b2 = T[1,1]*T[4,4] - T[1,4]*T[4,1] - T[3,2]
    b1 = T[3,2]*(T[1,1] + T[4,4]) - T[1,2]*T[3,1] - T34T42
    b0 = T[1,1]*(T34T42 - T[3,2]*T[4,4]) +
        T[1,2]*(T[3,1]*T[4,4] - T[3,4]*T[4,1]) +
        T[1,4]*(T[3,2]*T[4,1] - T[3,1]*T[4,2])

    booker_quartic_coeffs = MVector{5}(b0, b1, b2, b3, 1)
    booker_quartic_roots = MVector{4,eltype(booker_quartic_coeffs)}(0, 0, 0, 0)

    roots!(booker_quartic_roots, booker_quartic_coeffs, NaN, 4, false)

    return booker_quartic_roots, SVector(booker_quartic_coeffs)
end

"""
    dbookerquartic(ea::EigenAngle, M, q, B)
    dbookerquartic(T::TMatrix, dT, q, B)

Compute derivative `dq` of the Booker quartic roots `q` with respect to ``θ`` for the
ionosphere described by susceptibility tensor `M` or `T` matrix.
"""
dbookerquartic

function dbookerquartic(ea::EigenAngle, M, q, B)
    S, C, C² = ea.sinθ, ea.cosθ, ea.cos²θ

    dS = C
    dC = -S
    dC² = -2*S*C

    dB3 = dS*(M[1,3] + M[3,1])
    dB2 = -dC²*(2 + M[1,1] + M[3,3])
    dB1 = dS/S*B[2] - S*dC²*(M[1,3] + M[3,1])
    dB0 = dC²*(2*C²*(1 + M[1,1]) + M[3,3] + M[2,2] + M[1,1]*(M[3,3] + M[2,2]) -
            M[1,3]*M[3,1] - M[1,2]*M[2,1])

    dq = similar(q)
    @inbounds for i in eachindex(dq)
        dq[i] = -(((dB3*q[i] + dB2)*q[i] + dB1)*q[i] + dB0) /
                (((4*B[5]*q[i] + 3*B[4])*q[i] + 2*B[3])*q[i] + B[2])
    end

    return SVector(dq)
end

function dbookerquartic(T::TMatrix, dT, q, B)
    dB3 = -dT[1,1] - dT[4,4]
    dB2 = T[1,1]*dT[4,4] + dT[1,1]*T[4,4] - dT[1,4]*T[4,1] - dT[3,2]
    dB1 = -(-T[3,2]*(dT[1,1] + dT[4,4]) - dT[3,2]*(T[1,1] + T[4,4]) + dT[1,2]*T[3,1] +
        dT[3,4]*T[4,2])
    dB0 = -T[1,1]*(T[3,2]*dT[4,4] + dT[3,2]*T[4,4] - dT[3,4]*T[4,2]) -
        dT[1,1]*(T[3,2]*T[4,4] - T[3,4]*T[4,2]) +
        T[1,2]*(T[3,1]*dT[4,4] - dT[3,4]*T[4,1]) + dT[1,2]*(T[3,1]*T[4,4] - T[3,4]*T[4,1]) -
        T[1,4]*(-dT[3,2]*T[4,1]) - dT[1,4]*(T[3,1]*T[4,2] - T[3,2]*T[4,1])

    dq = similar(q)
    @inbounds for i in eachindex(dq)
        dq[i] = -(((dB3*q[i] + dB2)*q[i] + dB1)*q[i] + dB0) /
                (((4*B[5]*q[i] + 3*B[4])*q[i] + 2*B[3])*q[i] + B[2])
    end

    return SVector(dq)
end

"""
    upgoing(q)

Calculate the absolute angle of `q` in radians from 315°×π/180 on the complex plane. Smaller
values indicate upgoing waves.
"""
function upgoing(q::Complex)
    # -π/4 is 315°
    a = -π/4 - angle(q)

    # angle() is +/-π
    # `a` can never be > 3π/4, but there is a region on the plane where `a` can be < -π
    a < -π && (a += 2π)
    return abs(a)
end

upgoing(q::Real) = oftype(q, π/4)

"""
    sortquarticroots!(q)

Sort array of quartic roots `q` in place such that the first two correspond to upgoing waves
and the latter two correspond to downgoing waves.

General locations of the four roots of the Booker quartic on the complex plane corresponding
to the:

    1) upgoing evanescent wave
    2) upgoing travelling wave
    3) downgoing evanescent wave
    4) downgoing travelling wave

```
          Im
      3 . |
          |
          |
  . 4     |
 ------------------Re
          |      . 2
          |
          |
          | . 1
```

Based on [Pitteway1965] fig. 5.

# References

[Pitteway1965]: M. L. V. Pitteway, “The numerical calculation of wave-fields, reflexion
    coefficients and polarizations for long radio waves in the lower ionosphere. I.,” Phil.
    Trans. R. Soc. Lond. A, vol. 257, no. 1079, pp. 219–241, Mar. 1965,
    doi: 10.1098/rsta.1965.0004.
"""
function sortquarticroots!(q)
    # It looks like we could just `sort!(q, by=real)`, but the quadrants of each root are
    # not fixed and the positions in the Argand diagram are just approximate.

    length(q) == 4 || @warn "length of `q` is not 4"

    # Calculate and sort by distance from 315°
    # The two closest are upgoing and the two others are downgoing
    # This is faster than `sort!(q, by=upgoing)`
    dist = MVector{4,real(eltype(q))}(undef)
    @inbounds for i in eachindex(q)
        dist[i] = upgoing(q[i])
    end

    i = 4
    @inbounds while i > 1
        if dist[i] < dist[i-1]
            dist[i], dist[i-1] = dist[i-1], dist[i]
            q[i], q[i-1] = q[i-1], q[i]
            i = 4  # Need to restart at the end
        else
            i -= 1
        end
    end

    # Now arrange the roots so that they correspond to:
    # 1) upgoing evanescent, 2) upgoing travelling
    # 3) downgoing evanescent, and 4) downgoing travelling
    if imag(q[1]) > imag(q[2])
        q[2], q[1] = q[1], q[2]
    end
    if imag(q[3]) < imag(q[4])
        q[4], q[3] = q[3], q[4]
    end

    # NOTE: Nagano et al 1975 says that of q1 and q2, the one with the largest
    # absolute value corresponds to e1 and the other to e2, although this method
    # does not guarantee that. In fact, I find that this criteria is often not met.

    return q
end

@doc raw"""
    bookerwavefields(ea::EigenAngle, M)
    bookerwavefields(T::TMatrix)

Compute the two-column wavefields matrix `e` from the ionosphere with susceptibility tensor
`M` or `T` matrix for the two upgoing wavefields.

The first column of `e` is the evanescent wave and the second is the travelling wave.

```math
e =
    \begin{pmatrix}
         Ex₁  &  Ex₂ \\
        -Ey₁  & -Ey₂ \\
         Hx₁  &  Hx₂ \\
         Hy₁  &  Hy₂
    \end{pmatrix}
```

This function solves the eigenvalue problem ``Te = qe``. First, the Booker quartic is solved
for the roots `q`. Then they are sorted so that the roots associated with the two upgoing
waves can be selected. `e` is solved as the eigenvectors for the two `q`s. An
analytical solution is used where `e[2,:] = 1`.
"""
bookerwavefields

function bookerwavefields(T::TMatrix)
    q, B = bookerquartic(T)
    sortquarticroots!(q)
    return bookerwavefields(T, q)
end

function bookerwavefields(ea::EigenAngle, M)
    q, B = bookerquartic(ea, M)
    sortquarticroots!(q)
    T = tmatrix(ea, M)
    return bookerwavefields(T, q)
end

function bookerwavefields(T::TMatrix, q)
    # Precompute
    T14T41 = T[1,4]*T[4,1]
    T14T42 = T[1,4]*T[4,2]
    T12T41 = T[1,2]*T[4,1]

    e = MArray{Tuple{4,2},ComplexF64,2,8}(undef)
    @inbounds for i = 1:2
        d = T14T41 - (T[1,1] - q[i])*(T[4,4] - q[i])

        e[1,i] = (T[1,2]*(T[4,4] - q[i]) - T14T42)/d
        e[2,i] = 1
        e[3,i] = q[i]
        e[4,i] = (-T12T41 + T[4,2]*(T[1,1] - q[i]))/d
    end

    # By returning as SArray instead of MArray, the MArray doesn't get hit by GC
    return SArray(e)
end

"""
    bookerwavefields(ea::EigenAngle, M, ::Dθ)
    bookerwavefields(T::TMatrix, dT, ::Dθ)

Compute the two-column wavefields matrix `e` as well as its derivative with respect to
``θ``, returned as a tuple `(e, de)` for the ionosphere with susceptibility tensor `M` or
`T` matrix and its derivative with respect to ``θ``, `dT`.
"""
function bookerwavefields(T::TMatrix, dT, ::Dθ)
    q, B = bookerquartic(T)
    sortquarticroots!(q)
    dq = dbookerquartic(T, dT, q, B)
    return bookerwavefields(T, dT, q, dq)
end

function bookerwavefields(ea::EigenAngle, M, ::Dθ)
    q, B = bookerquartic(ea, M)
    sortquarticroots!(q)
    dq = dbookerquartic(ea, M, q, B)
    T = tmatrix(ea, M)
    dT = dtmatrix(ea, M)

    return bookerwavefields(T, dT, q, dq)
end

function bookerwavefields(T::TMatrix, dT, q, dq)
    # Precompute
    T14T41 = T[1,4]*T[4,1]
    T14T42 = T[1,4]*T[4,2]
    T12T41 = T[1,2]*T[4,1]

    e = MArray{Tuple{4,2},ComplexF64,2,8}(undef)
    de = similar(e)
    @inbounds for i = 1:2
        den = T14T41 - (T[1,1] - q[i])*(T[4,4] - q[i])
        den² = abs2(den)
        dden = dT[1,4]*T[4,1] - (T[1,1]*dT[4,4] + dT[1,1]*T[4,4]) +
            q[i]*(dT[1,1] + dT[4,4] - 2*dq[i]) + dq[i]*(T[1,1] + T[4,4])

        e1num = T[1,2]*(T[4,4] - q[i]) - T14T42
        e4num = -T12T41 + T[4,2]*(T[1,1] - q[i])
        e[1,i] = e1num/den
        e[2,i] = 1
        e[3,i] = q[i]
        e[4,i] = e4num/den

        # some dT terms == 0
        de1num = T[1,2]*(dT[4,4] - dq[i]) + dT[1,2]*(T[4,4] - q[i]) - dT[1,4]*T[4,2]
        de4num = -dT[1,2]*T[4,1] + T[4,2]*(dT[1,1] - dq[i])
        de[1,i] = (de1num*den - e1num*dden)/den²
        de[2,i] = 0
        de[3,i] = dq[i]
        de[4,i] = (de4num*den - e4num*dden)/den²
    end

    # By returning as SArray instead of MArray, the MArray doesn't get hit by GC
    return SArray(e), SArray(de)
end

@doc raw"""
    bookerreflection(ea::EigenAngle, M::SMatrix{3,3})
    bookerreflection(ea::EigenAngle, e)

Compute the ionosphere reflection coefficient matrix for a sharply bounded ionosphere from
4×2 wavefields matrix `e` or the susceptibility matrix `M`.

The ionosphere reflection coefficient matrix is computed from a ratio of the downgoing to
upgoing plane waves in the free space beneath the ionosphere [Budden1988] pg. 307.
These are obtained from the two upgoing characteristic waves found from the Booker quartic.
Each make up a column of `e`.

```math
R =
    \begin{pmatrix}
        Ce₁[4] - e₁[1] &  Ce₂[4] - e₂[1] \\
       -Ce₁[2] + e₁[3] & -Ce₂[2] + e₂[3]
    \end{pmatrix}
    \begin{pmatrix}
        Ce₁[4] + e₁[1] &  Ce₂[4] + e₂[1] \\
       -Ce₁[2] - e₁[3] & -Ce₂[2] - e₂[3]
    \end{pmatrix}^{-1}
```

The reflection coefficient matrix for the sharply bounded case is commonly used as a
starting solution for integration of the reflection coefficient matrix through the
ionosphere.

# References

[Budden1988]: K. G. Budden, “The propagation of radio waves: the theory of radio
    waves of low power in the ionosphere and magnetosphere,” First paperback edition.
    New York: Cambridge University Press, 1988.

# Extended help

The set of horizontal field components ``e = (Ex, -Ey, Z₀Hx, Z₀Hy)ᵀ`` can be separated into
an upgoing and downgoing wave, each of which is generally elliptically polarized. A ratio of
the amplitudes of these two waves give a reflection coefficient, except it would only apply
for an incident wave of that particular elliptical polarization. However, the first set of
fields can be linearly combined with a second independent solution for the fields, which
will generally have a different elliptical polarization than the first. Two linear
combinations of the two sets of fields are formed with unit amplitude, linearly polarized
incident waves. The reflected waves then give the components ``R₁₁``, ``R₂₁`` or ``R₁₂``,
``R₂₂`` for the incident wave in the plane of incidence and perpendicular to it,
respectively [Budden1988] pg 552.

The process for determining the reflection coefficient requires resolving the two sets of
fields ``e₁`` and ``e₂`` into the four linearly polarized vacuum modes. The layer of vacuum
can be assumed to be so thin that it does not affect the fields. There will be two upgoing
waves and two downgoing waves, each which has one ``E`` and one ``H`` in the plane of
incidence. If ``f₁, f₂, f₃, f₄`` are the complex amplitudes of the four component waves,
then in matrix notation ``e = Lf`` where ``L`` is the appropriate transformation matrix.

For ``e₁`` and ``e₂``, we can find the corresponding vectors ``f1`` and ``f2`` by
``f1 = L⁻¹e₁``, ``f2 = L⁻¹e₂`` where the two column vectors are partitioned such that
``f1 = (u1, d1)ᵀ`` and ``f2 = (u2, d2)ᵀ`` for upgoing and downgoing 2-element vectors ``u``
and ``d``. From the definition of the reflection coefficient ``R``, ``d = Ru``. Letting
``U = (u1, u2)``, ``D = (d1, d2)``, then ``D = RU`` and the reflection coefficient is
``R = DU⁻¹``. Because the reflection coefficient matrix is a ratio of fields, either ``e₁``
and/or ``e₂`` can be independently multiplied by an arbitrary constant and the value of
``R`` is unaffected.

This function directly computes ``D`` and ``U`` and solves for ``R`` using the right
division operator `R = D/U`.

For additional details, see [Budden1988], chapter 18, section 7.
"""
bookerreflection

function bookerreflection(ea::EigenAngle, e)
    C = ea.cosθ

    # The order of the two upgoing waves doesn't matter. Swapping the first and second
    # columns of `e` will result in the same reflection coefficient matrix.

    D = SMatrix{2,2}(C*e[4,1]-e[1,1], -C*e[2,1]+e[3,1], C*e[4,2]-e[1,2], -C*e[2,2]+e[3,2])
    U = SMatrix{2,2}(C*e[4,1]+e[1,1], -C*e[2,1]-e[3,1], C*e[4,2]+e[1,2], -C*e[2,2]-e[3,2])

    R = D/U

    return R
end

function bookerreflection(ea::EigenAngle, M::SMatrix{3,3})
    e = bookerwavefields(ea, M)
    return bookerreflection(ea, e)
end

"""
    bookerreflection(ea::EigenAngle, M, ::Dθ)

Compute the ionosphere reflection coefficient matrix ``R`` for a sharply bounded
ionosphere with susceptibility tensor `M`, as well as its derivative ``dR/dθ`` returned as
the tuple `(R, dR)`.
"""
function bookerreflection(ea::EigenAngle, M, ::Dθ)
    S, C = ea.sinθ, ea.cosθ

    e, de = bookerwavefields(ea, M, Dθ())

    D = SMatrix{2,2}(C*e[4,1]-e[1,1], -C*e[2,1]+e[3,1], C*e[4,2]-e[1,2], -C*e[2,2]+e[3,2])
    dD = SMatrix{2,2}(-S*e[4,1] + C*de[4,1] - de[1,1], S*e[2,1] - C*de[2,1] + de[3,1],
                      -S*e[4,2] + C*de[4,2] - de[1,2], S*e[2,2] - C*de[2,2] + de[3,2])

    U = SMatrix{2,2}(C*e[4,1]+e[1,1], -C*e[2,1]-e[3,1], C*e[4,2]+e[1,2], -C*e[2,2]-e[3,2])
    dU = SMatrix{2,2}(-S*e[4,1] + C*de[4,1] + de[1,1], S*e[2,1] - C*de[2,1] - de[3,1],
                      -S*e[4,2] + C*de[4,2] + de[1,2], S*e[2,2] - C*de[2,2] - de[3,2])

    R = D/U
    dR = dD/U + D*(-U\dU/U)

    return R, dR
end


Nor(x;m=0.0,s=1.0) = 1.0/sqrt(2*pi)/s * exp( -0.5/(s^2)*(x-m)^2 )
Nor2(x;m1=0.0,m2=0.0,s=1.0) = 0.7/sqrt(2*pi)/s * exp( -0.5/(s^2)*(x-m1)^2 ) + 0.3/sqrt(2*pi)/s * exp( -0.5/(s^2)*(x-m2)^2 )
function simdata01_true(;N1=100,N2=100)
  # common struct to hold all simulated data
  data = DataGroundTruth()
  samples = SampleData(zeros(0,0), 1, N1+N2)

  temp1 = (randn(round(Int,0.7*N1)).-1.0)
  cluster1 = ([temp1; randn(N1-length(temp1)).-5.0]')
  cluster2 = randn(1,N2).+1

  data.clusters[1] = cluster1
  data.clustersizes[1] = N1
  data.clusters[2] = cluster2
  data.clustersizes[2] = N2
  # one dimensional, Array{Float64,1}, with N1+N2 elements [vector alias]
  sd = [cluster1';cluster2']'
  permu = randperm(N1+N2)
  samples.samples = sd[:,permu]

y = randn(10)

  data.assignment = ([ones(Int,N1);2*ones(Int,N2)])[permu]

  # [OPTIONAL] going to carry plot layers for convenience also
  drawN = 200
  pts1 = StatsBase.sample(cluster1[1,:][:], ProbabilityWeights(1.0/N1*ones(N1)), drawN, replace=false)
  yy1 = zeros(drawN)
  for i in 1:drawN  yy1[i] = Nor2(pts1[i],m1=-1.0,m2=-5.0) end
  data.plt_lyr_cluster[1] = layer(x=pts1, y=yy1, Geom.line, Theme(default_color=colorant"blue"))[1] #parse(Colorant,"red")
  data.plt_lyr_cluster_nocolor[1] = layer(x=pts1, y=yy1, Geom.line, Theme(default_color=colorant"gray"))[1] #parse(Colorant,"red")

  pts2 = StatsBase.sample(cluster2[1,:][:], ProbabilityWeights(1.0/N2*ones(N2)), drawN, replace=false)
  yy2 = zeros(drawN)
  for i in 1:drawN  yy2[i] = Nor(pts2[i],m=+1.0) end
  data.plt_lyr_cluster[2] = layer(x=pts2, y=yy2, Geom.line, Theme(default_color=colorant"red"))[1] #parse(Colorant,"red")
  data.plt_lyr_cluster_nocolor[2] = layer(x=pts2, y=yy2, Geom.line, Theme(default_color=colorant"gray"))[1] #parse(Colorant,"red")

  return samples, data
end

# two dimensional example
function simdata02_true(;N1=100,N2=100)
  # common struct to hold all simulated data
  data = DataGroundTruth()
  samples = SampleData(zeros(0,0), 1, N1+N2)

  cluster1 = rand(MvNormal([3.0;3.0],[2.5 -1.5; -1.5 2.5]),N1)
  cluster2 = rand(MvNormal([-2.0;-2.0],[4.0 3.0; 3.0 4.0]),N2)

  data.clusters[1] = cluster1
  data.clustersizes[1] = N1
  data.clusters[2] = cluster2
  data.clustersizes[2] = N2
  # two rows, N1+N2 columns (Array{Float64,2})
  # data.samples = [cluster1';cluster2']'
  # TODO: not sure if sd should be a vector or matrix
  sd = [cluster1';cluster2']'
  permu = randperm(N1+N2)
  samples.samples = sd[:,permu]

  data.assignment = ([ones(Int,N1);2*ones(Int,N2)])[permu]

  drawN = 200

  pts1 = StatsBase.sample(cluster1[1,:][:], StatsBase.ProbabilityWeights(1.0/N1*ones(N1)), drawN, replace=false)
  data.plt_lyr_cluster[1] = layer(x=cluster1[1,:], y=cluster1[2,:], Geom.point, Theme(default_color=colorant"blue"))
  data.plt_lyr_cluster_nocolor[1] = layer(x=cluster1[1,:], y=cluster1[2,:], Geom.point, Theme(default_color=colorant"gray"))
  # data.plt_lyr_cluster[1] = layer(x=pts1, y=yy1, Geom.point, Theme(default_color=colorant"blue"))[1] #parse(Colorant,"red")
  # data.plt_lyr_cluster_nocolor[1] = layer(x=pts1, y=yy1, Geom.point, Theme(default_color=colorant"gray"))[1] #parse(Colorant,"red")

  pts2 = StatsBase.sample(cluster2[1,:][:], StatsBase.ProbabilityWeights(1.0/N2*ones(N2)), drawN, replace=false)
  data.plt_lyr_cluster[2] = layer(x=cluster2[1,:], y=cluster2[2,:], Geom.point, Theme(default_color=colorant"red"))
  data.plt_lyr_cluster_nocolor[2] = layer(x=cluster2[1,:], y=cluster2[2,:], Geom.point, Theme(default_color=colorant"gray"))
  # data.plt_lyr_cluster[2] = layer(x=pts2, y=yy2, Geom.point, Theme(default_color=colorant"red"))[1] #parse(Colorant,"red")
  # data.plt_lyr_cluster_nocolor[2] = layer(x=pts2, y=yy2, Geom.point, Theme(default_color=colorant"gray"))[1] #parse(Colorant,"red")

  return samples, data
end


import ForwardDiff: valtype
const ∇ = Tensors.gradient
const Δ = Tensors.hessian

function Ψ(C, μ, Kb)
    T = ForwardDiff.valtype(eltype(C))
    detC = det(C)
    J = sqrt(detC)
    Ĉ = detC^(convert(T, -1 / 3)) * C
    return (convert(T, μ) * (tr(Ĉ) - 3) / 2 + convert(T, Kb) * (J - 1)^2)
end

function S(C, μ, Kb)
    T = ForwardDiff.valtype(eltype(C))
    I = one(C)
    J = sqrt(det(C))
    invC = inv(C)
    return  convert(T, μ) * det(C)^(convert(T, -1/3)) * (I - tr(C) * invC / 3) + convert(T, Kb) * (J - 1) * J * invC
end

const μ = 1e10;
const Kb = 1.66e11;

Ψ(C) = Ψ(C, μ, Kb)
S(C) = S(C, μ, Kb)

SUITE["gradient"] = BenchmarkGroup(["ad"])
SUITE["hessian"] = BenchmarkGroup(["ad"])

for dim in (ALL_DIMENSIONS ? (1,2,3) : (3,))
    for T in (Float64, Float32)
        V2 = tensor_dict[(dim, 2, T)]
        V2 = V2' ⋅ V2
        V2sym = symtensor_dict[(dim, 2, T)]
        V2sym = V2sym' ⋅ V2sym

        @assert eltype(Ψ(V2)) == T
        @assert eltype(S(V2)) == T
        @assert eltype(∇(Ψ, V2)) == T
        @assert eltype(Δ(Ψ, V2)) == T

        SUITE["gradient"]["∇Ψ(SymmetricTensor{2, $dim, $T})"]       = @benchmarkable ∇(Ψ, $V2sym)
        SUITE["gradient"]["∇Ψ(Tensor{2, $dim, $T})"]                = @benchmarkable ∇(Ψ, $V2)
        SUITE["gradient"]["∇Ψ(SymmetricTensor{2, $dim, $T}) - ana"] = @benchmarkable S($V2sym)
        SUITE["gradient"]["∇Ψ(Tensor{2, $dim, $T}) - ana"]          = @benchmarkable S($V2)
        SUITE["gradient"]["∇S(SymmetricTensor{2, $dim, $T})"]       = @benchmarkable ∇(S, $V2sym)
        SUITE["gradient"]["∇S(Tensor{2, $dim, $T})"]                = @benchmarkable ∇(S, $V2)

        SUITE["hessian"]["ΔΨ(SymmetricTensor{2, $dim, $T})"] = @benchmarkable Δ(Ψ, $V2sym)
        SUITE["hessian"]["ΔΨ(Tensor{2, $dim, $T})"]          = @benchmarkable Δ(Ψ, $V2)
    end
end


using Pkg, DelimitedFiles, MPIReco

# Install required packages
for P in [:HTTP,:TensorDecompositions, :PyPlot]
 try
   @eval using $P
 catch
   Pkg.add(String(P))
   @eval using $P
 end
end

# Download data
include("downloadData.jl")

include("./utils/subsampling.jl")
include("./utils/FRTruncation.jl")

# load system matrix
datadirSM = "./data/"
bSF = MPIFile(datadirSM*"SM.mdf")
# backup for local testing
# bSF = MPIFile("/opt/mpidata/20180228_180048_OpenMPIData_1_1/9")
nx,ny,nz = calibSize(bSF);
f = [2145,8629,60903]
S = ComplexF64.( getSystemMatrix(bSF,f,bgCorrection=true, tfCorrection=false) )

# sampling locations (linear index)
samplingIdx = vec( readdlm(joinpath(@__DIR__,"data/pdSamplingIdx.txt"), '\t', Int64) )

# get undersampled data
y = sfMeas(S, samplingIdx)

# CS-recovery
par_cs = Dict{Symbol,Any}()
par_cs[:shape] = (nx,ny,nz)   # size of the spatial measurement grid
par_cs[:λ_l1] = 0.02          # l1-regularization parameter
par_cs[:ρ_l1] = 2.0	      # parameter ρ associated with the l1-regularization
par_cs[:iterationsInner] = 50 # max. number of inner Split Bregman iterations
par_cs[:iterations] = 10      # number of outer Split Bregman iterations
par_cs[:relTol] = 1.e-2       # stopping criterion for the inner Split Bregman loop
par_cs[:absTol] = 1.e-4       # stopping criterion for the inner Split Bregman loop

S_cs = smRecovery(y, samplingIdx, par_cs)

# CSLR-recovery
par_cslr = Dict{Symbol,Any}()
par_cslr[:shape] = (nx,ny,nz)
par_cslr[:λ_l1] = 0.02
par_cslr[:ρ_l1] = 2.0
par_cslr[:λ_lr] = 0.02	        # TNN regularization parameter (for CSLR and CSFR)
par_cslr[:ρ_lr] = 0.2           # parameter ρ associated with the low rank regularization
par_cslr[:iterationsInner] = 50
par_cslr[:iterations] = 10
par_cslr[:relTol] = 1.e-2
par_cslr[:absTol] = 1.e-4

S_cslr = smRecovery(y, samplingIdx, par_cslr)

# reshape into tensor format
S = reshape(S,nx,ny,nz,:)
S_cs = reshape(S_cs,nx,ny,nz,:)
S_cslr = reshape(S_cslr,nx,ny,nz,:)

# CSFR
r = (6,6,6) # maximum HOSVD rank
S_csfr = FRTruncation3d(S_cslr,r)

# plot frequency components
figure()
subplot(4,3,1)
title("k=$(f[1])")
ylabel("measured")
imshow(abs.(S[:,:,17,1]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,2)
title("k=$(f[2])")
imshow(abs.(S[:,:,17,2]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,3)
title("k=$(f[3])")
imshow(abs.(S[:,:,17,3]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,4)
ylabel("CS")
imshow(abs.(S_cs[:,:,17,1]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,5)
imshow(abs.(S_cs[:,:,17,2]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,6)
imshow(abs.(S_cs[:,:,17,3]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,7)
ylabel("CSLR")
imshow(abs.(S_cslr[:,:,17,1]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,8)
imshow(abs.(S_cslr[:,:,17,2]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,9)
imshow(abs.(S_cslr[:,:,17,3]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,10)
ylabel("CSFR")
imshow(abs.(S_csfr[:,:,17,1]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,11)
imshow(abs.(S_csfr[:,:,17,2]),cmap="gray"); xticks([]); yticks([])
subplot(4,3,12)
imshow(abs.(S_csfr[:,:,17,3]),cmap="gray"); xticks([]); yticks([])


#=

  Curious numbers in ConstraintSolver.jl

  """
  Curious Numbers from "Amusements in Mathematics, Dudeney", number 114.

  The number 48 has this peculiarity, that if you add 1 to it the result
  is a square number, and if you add 1 to its half, you also get a
  square number. Now, there is no limit to the numbers that have this
  peculiarity, and it is an interesting puzzle to find three more of
  them---the smallest possible numbers. What are they?
  """

  The least such numbers are:
  [
   [48,49,7,24,25,5],
   [1680,1681,41,840,841,29],
   [57120,57121,239,28560,28561,169],
   [1940448,1940449,1393,970224,970225,985]
  ]

  Note: This model use the experimental (and very slow)
  non-linear mult constraint (mult(model,A,B,C) -> C = A*B)
  It thus only handle the first two results shows above.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function cur_num(max_val=100,print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = 6
    @variable(model, 1 <= x[1:n] <= max_val, Int)
    X,A,B,C,D,E = x

    @constraint(model,X + 1 == A) # if you add 1 to it
    # @constraint(model,A == B * B) # the result is a square number
    mult(model,B,B,A)

    # @cons)raint(model,X == 2 * C) # if you to its half
    mult(model,2,C,X)
    @constraint(model,C + 1 == D) # add 1
    # @constraint(model,D == E * E) # you also get a square number
    mult(model,E,E,D)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$x_val")

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

# for max_val in [100,2000]
for max_val in [100,2000]
    println("\nmax_val:$max_val")
    @time cur_num(max_val)
end


lines = readlines()

function solve(use_aim=false)
    depth, position, aim = 0, 0, 0
    for line in lines
        word, num = split(line)
        num = parse(Int, num)
        if word == "forward"
            position += num
            depth += use_aim ? aim * num : 0
        else
            if use_aim
                aim += word == "down" ? num : -num
            else
                depth += word == "down" ? num : -num
            end
        end
    end
    return depth * position
end
println(solve(false))
println(solve(true))


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 6316f37d-24d7-488a-9916-47a75512ea54
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ 9fa2be99-1c4a-4e53-8a01-546a3670f19f
begin
    using FFTW
    using Flux
    using Plots
    plotly()
end

# ╔═╡ b682a76e-75c6-4e0f-b542-210725501e5b
html"""<style>
main {
    max-width: 700pt;
}
"""

# ╔═╡ b5ba5d3a-eb07-4d5f-a802-6d85075ed528
md"
# Fourier Net

*“Can neural networks learn the Fourier transform?”*

This notebook is a implementation of [The Fourier transform is a neural network](https://sidsite.com/posts/fourier-nets/).
"

# ╔═╡ 21b015ea-aa4b-11eb-182b-a7590312f5ba
md"
**By the definition of DFT:**

A single amplitude for a singal duration of frequency can be formulated as follow:

$x̂_k = \sum_{n=0}^{N-1} x_n \cdot exp(\frac{-i2\pi k}{N}n)$

$x̂_k =
\begin{bmatrix} x_0 & x_1 & ... & x_{N-1} \end{bmatrix}
\begin{bmatrix}
    exp(0) \\
    exp(\frac{-i 2 \pi k}{N}1) \\
    exp(\frac{-i 2 \pi k}{N}2) \\
    exp(\frac{-i 2 \pi k}{N}3) \\
    exp(\frac{-i 2 \pi k}{N}4) \\
    ... \\
    exp(\frac{-i 2 \pi k}{N}(N-1))
\end{bmatrix}$

Therefore, N of amplitide for N durations of frequency can be illustrated like:

$\begin{bmatrix} x̂_0 \\ ... \\ x̂_{N-1} \end{bmatrix} =
\begin{bmatrix}
    exp(0) & exp(0) & exp(0) & ... & exp(0) \\
    exp(0) & exp(\frac{-i2\pi}{N}1) & exp(\frac{-i2\pi 2}{N}1) & ... & exp(\frac{-i2\pi (N-1)}{N}1) \\
    exp(0) & exp(\frac{-i2\pi}{N}2) & exp(\frac{-i2\pi 2}{N}2) & ... & exp(\frac{-i2\pi (N-1)}{N}2) \\
    exp(0) & exp(\frac{-i2\pi}{N}3) & exp(\frac{-i2\pi 2}{N}3) & ... & exp(\frac{-i2\pi (N-1)}{N}3) \\
    exp(0) & exp(\frac{-i2\pi}{N}4) & exp(\frac{-i2\pi 2}{N}4) & ... & exp(\frac{-i2\pi (N-1)}{N}4) \\
    ... \\
    exp(0) & exp(\frac{-i2\pi}{N}(N-1)) & exp(\frac{-i2\pi 2}{N}(N-1)) & ... & exp(\frac{-i2\pi (N-1)}{N}(N-1))
\end{bmatrix}_{N \times N}^T
\begin{bmatrix} x_0 \\ ... \\ x_{N-1} \end{bmatrix}$

Or in short:

$x̂⃗ = 𝐖_{ComplexFourier}^T x⃗$
"

# ╔═╡ 212ac582-1270-49da-9cd4-078ab695be66
md"
## Generate random signal

$x⃗ = \begin{bmatrix} x_0 \\ ... \\ x_{N-1} \end{bmatrix}$
"

# ╔═╡ 7efa004c-57c4-4b9c-bc74-90252de25218
begin
    len = 128
    x = rand(len)
end;

# ╔═╡ 0d6607c2-0881-40f7-81fd-494d49faa697
md"
## Computing the Fourier transform using Fourier weights
"

# ╔═╡ eeded1aa-3edd-4aa7-b147-9c69e2293267
md"
**Create Fourier weighrts**

$𝐖_{ComplexFourier}^T = \begin{bmatrix}
    exp(0) & exp(0) & exp(0) & ... & exp(0) \\
    exp(0) & exp(\frac{-i2\pi}{N}1) & exp(\frac{-i2\pi 2}{N}1) & ... & exp(\frac{-i2\pi (N-1)}{N}1) \\
    exp(0) & exp(\frac{-i2\pi}{N}2) & exp(\frac{-i2\pi 2}{N}2) & ... & exp(\frac{-i2\pi (N-1)}{N}2) \\
    exp(0) & exp(\frac{-i2\pi}{N}3) & exp(\frac{-i2\pi 2}{N}3) & ... & exp(\frac{-i2\pi (N-1)}{N}3) \\
    exp(0) & exp(\frac{-i2\pi}{N}4) & exp(\frac{-i2\pi 2}{N}4) & ... & exp(\frac{-i2\pi (N-1)}{N}4) \\
    ... \\
    exp(0) & exp(\frac{-i2\pi}{N}(N-1)) & exp(\frac{-i2\pi 2}{N}(N-1)) & ... & exp(\frac{-i2\pi (N-1)}{N}(N-1))
\end{bmatrix}_{N \times N}^T$
"

# ╔═╡ 1f66b59e-bdcd-4cdb-9f28-3bbf0f7dd08e
function create_𝐰_fourier(len::Integer)
    𝐰 = Matrix{ComplexF64}(undef, len, len)
    for n in 0:len-1, k in 0:len-1
        𝐰[n+1, k+1] = exp(-im*2π*k*n/len)
    end

    return 𝐰'
end;

# ╔═╡ 036016cd-e4f5-40f7-96d9-cac8fdbebd01
md"
**Do the Fourier transform by Fourier weights and FFT**

$x̂⃗ = 𝐖_{ComplexFourier}^T x⃗$

$x̂⃗_{fft} = 𝐹x⃗$
"

# ╔═╡ d5933fdb-b9ed-4475-af43-58dac24a47d5
begin
    # Fourier transform via Fourier weights
    𝐰 = create_𝐰_fourier(len)
    x̂ = 𝐰 * x

    # Fourier transform via fft
    x̂_fft = fft(x)

    # root-mean-sqrt error between those two
    ϵ_x̂ = sqrt(sum(abs.(x̂_fft - x̂).^2)/len)
end;

# ╔═╡ 9d937bef-f325-4299-8c4b-7ae358aeb1c9
md"
Spectrum's RMSE: $ϵ_x̂
"

# ╔═╡ f65c7059-9c2e-4993-a7e6-32cb9a848f23
begin
    plot(title="Spectrum", xlabel="Freq (arb. unit)", ylabel="Amp (arb. unit)")
    plot!(real(x̂_fft), label="FFT")
    plot!(real(x̂), label="Fourier weights")
end

# ╔═╡ 6e200b44-97fb-43e0-8026-aabef11f4f17
md"
**Reconstruct signal by inverse Fourier weights and iFFT**

$x̂_k = \sum_{n=0}^{N-1} x_n \cdot exp(\frac{-i2\pi k}{N}n)$

$x_t = \frac{1}{N} \sum_{k=0}^{N-1} x̂_k \cdot exp(\frac{i2\pi k}{N}n)$
"

# ╔═╡ 45844ffc-9f5c-4952-abde-4ce770ab77f6
function create_𝐰¯¹_fourier(len::Integer)
    𝐰¯¹ = Matrix{ComplexF64}(undef, len, len)
    for n in 0:len-1, k in 0:len-1
        𝐰¯¹[n+1, k+1] = exp(im*2π*k*n/len)
    end

    return (𝐰¯¹/len)'
end;

# ╔═╡ 856e62d3-53c8-49a9-9061-d223dbb92ef1
md"
**Do the inverse Fourier transform by inverse Fourier weights and iFFT**

$x⃗ = 𝐖¯¹_{ComplexFourier}^T x̂⃗$

$x⃗_{fft} = 𝐹¯¹x⃗$
"

# ╔═╡ ac36e1c5-dade-4b0f-974e-2d261ea260f8
begin
    # reconstruct signal via Fourier weights
    𝐰¯¹ = create_𝐰¯¹_fourier(len)
    𝑓¯¹x̂ = real(𝐰¯¹ * x̂)

    # reconstruct signal via ifft
    𝑓¯¹x̂_fft = real(ifft(x̂_fft))

    # root-mean-sqrt error between those two
    ϵ_𝑓¯¹x̂_fft = sqrt(sum(abs.(𝑓¯¹x̂_fft - x).^2)/len)
    ϵ_𝑓¯¹x̂ = sqrt(sum(abs.(𝑓¯¹x̂ - x).^2)/len)
end;

# ╔═╡ 7339df20-651e-42ac-8286-e9696fb72458
md"
Inverse Fourier transform's RMSE:

* Via IFFT: $ϵ_𝑓¯¹x̂_fft
* Via Fourier waights: $ϵ_𝑓¯¹x̂
"

# ╔═╡ e9c246d8-ae77-4236-8ca0-15de7cced221
begin
    plot(title="Signal", xlabel="Time (arb. unit)", ylabel="Amp (arb. unit)", ticks=[], ylims=(0, 3.5))
    plot!(x .+2 , label="Original signal")
    plot!(𝑓¯¹x̂ .+ 1, label="Reconstructed signal (Fourier weights)")
    plot!(𝑓¯¹x̂_fft, label="Reconstructed signal (ifft)")
end

# ╔═╡ 3c03ec6d-8994-4886-9006-3cddbf0ac027
md"
## Learning the Fourier transform via gradient-descent

We can consider the discrete Fourier transform (DFT) to be an artificial neural network: it is a single layer network, with no bias, no activation function, and particular values for the weights. The number of output nodes is equal to the number of frequencies we evaluate.
"

# ╔═╡ 0b4853ce-8dd0-4fe5-881d-cdf5d0289cc8
md"
**Build Fourier Net**
"

# ╔═╡ 1cbd900c-b74b-41c8-a57c-6a6ae644b88c
begin
    mutable struct FourierNet
        𝐰
    end

    Flux.@functor FourierNet

    (m::FourierNet)(x) = m.𝐰 * x

    loss(m::FourierNet, x, x̂) = sum(abs2, x̂ .- m(x)) / len
end;

# ╔═╡ 0abb581e-6835-4dfc-bd01-5b2b6306d242
md"
**Initiating the training process**
"

# ╔═╡ 895ee250-e3ac-4879-b8ee-f1ba8729adaa
begin
    # define the net and loss function
    m = FourierNet(rand(ComplexF64, len, len))
    loss(x, x̂) = loss(m, x, x̂)

    # prepare data
    xs = rand(ComplexF64, len, 50000)
    x̂s = 𝐰 * xs
    data = Flux.Data.DataLoader((xs, x̂s), batchsize=15, shuffle=true)

    # define call-back function for plotting
    training_loss = []
    function cb()
        lossₜ = loss(xs, x̂s)
        push!(training_loss, lossₜ)
    end

    # training process
    train_𝐰!(n, η) = Flux.@epochs n Flux.Optimise.train!(
        loss,
        params(m),
        data,
        Descent(η),
        # cb=Flux.throttle(@show(loss(xs, x̂s)), 1, leading=false, trailing=true)
        cb=Flux.throttle(cb, 1, leading=false, trailing=true)
    )

    # start training
    train_𝐰!(15, 1e-1)
    train_𝐰!(30, 1e-2)
    train_𝐰!(20, 1e-3)
end;

# ╔═╡ 3c1157bd-9780-429a-9cbf-cd5a18d676da
plot(training_loss, title="Loss", xlabel="time (sec)", ylabel="loss", label="training")

# ╔═╡ 9ef928f0-5577-4507-8fbe-193ce6925716
md"
Weights' residual: $(sum(abs, 𝐰 - m.𝐰)/length(𝐰))

Weights' RMSE: $(sqrt(sum(abs2, 𝐰 - m.𝐰)/length(𝐰)))
"

# ╔═╡ 04189afa-6cdd-45f4-8eaa-12140bde3ada
begin
    x̂_grad = m(x)

    plot(title="Spectrum", xlabel="Freq (arb. unit)", ylabel="Amp (arb. unit)")
    plot!(real(x̂_grad), label="Fourier weights by GD")
    plot!(real(x̂_fft), label="FFT")
end

# ╔═╡ Cell order:
# ╟─b682a76e-75c6-4e0f-b542-210725501e5b
# ╟─b5ba5d3a-eb07-4d5f-a802-6d85075ed528
# ╠═6316f37d-24d7-488a-9916-47a75512ea54
# ╠═9fa2be99-1c4a-4e53-8a01-546a3670f19f
# ╟─21b015ea-aa4b-11eb-182b-a7590312f5ba
# ╟─212ac582-1270-49da-9cd4-078ab695be66
# ╠═7efa004c-57c4-4b9c-bc74-90252de25218
# ╟─0d6607c2-0881-40f7-81fd-494d49faa697
# ╟─eeded1aa-3edd-4aa7-b147-9c69e2293267
# ╠═1f66b59e-bdcd-4cdb-9f28-3bbf0f7dd08e
# ╟─036016cd-e4f5-40f7-96d9-cac8fdbebd01
# ╠═d5933fdb-b9ed-4475-af43-58dac24a47d5
# ╟─9d937bef-f325-4299-8c4b-7ae358aeb1c9
# ╠═f65c7059-9c2e-4993-a7e6-32cb9a848f23
# ╟─6e200b44-97fb-43e0-8026-aabef11f4f17
# ╠═45844ffc-9f5c-4952-abde-4ce770ab77f6
# ╟─856e62d3-53c8-49a9-9061-d223dbb92ef1
# ╠═ac36e1c5-dade-4b0f-974e-2d261ea260f8
# ╟─7339df20-651e-42ac-8286-e9696fb72458
# ╠═e9c246d8-ae77-4236-8ca0-15de7cced221
# ╟─3c03ec6d-8994-4886-9006-3cddbf0ac027
# ╟─0b4853ce-8dd0-4fe5-881d-cdf5d0289cc8
# ╠═1cbd900c-b74b-41c8-a57c-6a6ae644b88c
# ╟─0abb581e-6835-4dfc-bd01-5b2b6306d242
# ╠═895ee250-e3ac-4879-b8ee-f1ba8729adaa
# ╠═3c1157bd-9780-429a-9cbf-cd5a18d676da
# ╟─9ef928f0-5577-4507-8fbe-193ce6925716
# ╠═04189afa-6cdd-45f4-8eaa-12140bde3ada


function logsumexp_dim1(x_km)
  max_m = maximum(x_km, dims=1)
  exp_km = exp.(x_km .- max_m)
  s_m = log.(sum(exp_km, dims=1)) .+ max_m
end

function logsumexp_dim2(x_km)
  max_k = maximum(x_km, dims=2)
  exp_km = exp.(x_km .- max_k)
  s_k = log.(sum(exp_km, dims=2)) .+ max_k
end

function wham_equation!(F, x, K, M, bias_km, log_Neff_km, log_heff_m)
    log_prob_m = reshape(x[1:M], 1, M)
    f_k = reshape(x[(M+1):(M+K)], K, 1)

    # 1st equation
    log_denominator_km = f_k .- bias_km
    log_denominator_km = log_denominator_km .+ log_Neff_km
    log_denominator_m  = logsumexp_dim1(log_denominator_km)
    log_numerator_m = log_heff_m
    F[1:M] = - log_prob_m .+ log_numerator_m .- log_denominator_m

    # 2nd equation
    F[(M+1):(M+K)] = f_k .+ logsumexp_dim2(log_prob_m .- bias_km)
    F[(M+1):(M+K)] = F[(M+1):(M+K)] .- F[M+1]
end

"""
wham

estimates (reduced) free energies of umbrella-windows and potential of mean force in data-bins by using the WHAM equations.

References
[1] S. Kumar, D. Bouzida, R. H. Swendsen, P. A. Kollman, and
    J. M. Rosenberg, J. Comput. Chem. 13, 1011 (1992).
"""
function wham(h_km, bias_km; ftol=1e-10, iterations=10^5)
    # K: number of umbrella-windows
    K = size(h_km, 1)

    # M: number of bins
    M = size(h_km, 2)

    # conversion from Int64 to Float64
    h_km = map(Float64, h_km)

    # calculate counts (N_k)
    # N_k: number of samples counted in the histogram from umbrella-window (k)
    N_k = sum(h_km, dims=2)
    log_N_k = zeros(Float64, K, 1)
    log_N_k[N_k .> 0.0] .= log.(N_k[N_k .> 0.0])

    # calculate effective histogram (heff_m)
    # heff_m: effective number of independent samples in data-bin (m)
    g_km = ones(Float64, K, M)
    heff_m = sum(h_km./g_km, dims=1)
    log_heff_m = zeros(Float64, 1, M)
    log_heff_m[:, :] .= -10.0 #double(log(eps(Float32)))
    log_heff_m[heff_m .> 0.0] .= log.(heff_m[heff_m .> 0.0])

    # calculate effective counts (Neff_km)
    # Neff_km: effective number of independent samples for umbrella-window (k) in data-bin (m)
    Neff_km = N_k ./ g_km
    log_Neff_km = zeros(Float64, K, M)
    log_Neff_km[:, :] .= -10.0 #double(log(eps(Float32)))
    log_Neff_km[Neff_km .> 0] .= log.(Neff_km[Neff_km .> 0])

    #F_tmp = zeros(Float64, M+K)
    #initial_x = zeros(Float64, M+K)
    #wham_equation!(F_tmp, initial_x, K, M, bias_km, log_Neff_km, log_heff_m)

    f!(F, x) = wham_equation!(F, x, K, M, bias_km, log_Neff_km, log_heff_m)
    initial_x = zeros(Float64, M+K)
    sol = nlsolve(f!, initial_x, autodiff=:forward, ftol=ftol, iterations=iterations)

    log_prob_m = reshape(sol.zero[1:M], 1, M)
    f_k = reshape(sol.zero[(M+1):(M+K)], K, 1)
    pmf_m = - log_prob_m
    f_k .- f_k[1, 1], pmf_m .- pmf_m[1, 1]
end

"""
wham_iteration

Old version of wham(). this function iteratively solve wham equations.

References
[1] S. Kumar, D. Bouzida, R. H. Swendsen, P. A. Kollman, and
    J. M. Rosenberg, J. Comput. Chem. 13, 1011 (1992).
"""
function wham_iteration(h_km, bias_km)
    TOLERANCE = 10^(-8)

    # K: number of umbrella-windows
    K = size(h_km, 1)

    # M: number of bins
    M = size(h_km, 2)

    # conversion from Int64 to Float64
    h_km = map(Float64, h_km)

    # calculate counts (N_k)
    # N_k: number of samples counted in the histogram from umbrella-window (k)
    N_k = sum(h_km, dims=2)
    log_N_k = zeros(Float64, K, 1)
    log_N_k[N_k .> 0.0] .= log.(N_k[N_k .> 0.0])

    # calculate effective histogram (heff_m)
    # heff_m: effective number of independent samples in data-bin (m)
    g_km = ones(Float64, K, M)
    heff_m = sum(h_km./g_km, dims=1)
    log_heff_m = zeros(Float64, 1, M)
    log_heff_m[:, :] .= -10.0 #double(log(eps(Float32)))
    log_heff_m[heff_m .> 0.0] .= log.(heff_m[heff_m .> 0.0])

    # calculate effective counts (Neff_km)
    # Neff_km: effective number of independent samples for umbrella-window (k) in data-bin (m)
    Neff_km = N_k ./ g_km
    log_Neff_km = zeros(Float64, K, M)
    log_Neff_km[:, :] .= -10.0 #double(log(eps(Float32)))
    log_Neff_km[Neff_km .> 0] .= log.(Neff_km[Neff_km .> 0])

    # solve the WHAM equations by self-consistent iteration
    f_k = zeros(Float64, K, 1)
    check_convergence = Inf

    log_prob_m = zeros(Float64, 1, M)
    count_iteration = 0
    while check_convergence > TOLERANCE

      # 1st equation
      log_denominator_km = f_k .- bias_km
      log_denominator_km = log_denominator_km .+ log_Neff_km
      log_denominator_m  = logsumexp_dim1(log_denominator_km)

      log_numerator_m = log_heff_m

      log_prob_m = log_numerator_m .- log_denominator_m

      # 2nd equation
      f_k_new = - logsumexp_dim2(log_prob_m .- bias_km)
      f_k_new = f_k_new .- f_k_new[1]

      # check convergence
      check_convergence = maximum(abs.(f_k_new - f_k))./std(f_k_new)
      f_k = f_k_new;

      count_iteration = count_iteration + 1
      if mod(count_iteration, 100) == 0
        @printf("%dth iteration  delta = %e  tolerance = %e\n", count_iteration, check_convergence, TOLERANCE)
        @printf("free energies = ")
        for k = 1:K
          @printf("%f ", f_k[k])
        end
        @printf("\n")
        @printf("\n")
      end
    end

    pmf_m = - log_prob_m
    f_k .- f_k[1, 1], pmf_m .- pmf_m[1, 1]
end


module AngleBetweenVectors

import Base: angle

import LinearAlgebra: norm

Floats = Union{AbstractFloat, Complex{T} where T<:AbstractFloat}

@inline unitize(p) = p ./ norm(p)

"""
    angle(point1::T, point2::T) where {T}

Accurately ascertains the undirected angle (0 <= radians < pi)
between two points given in N-dimensional Cartesian coordinates.

Prefer this to `acos` alternatives
- more reliably accurate
- more consistently stable

Suggested when any |coordinate| of either point may be outside 2^±20 or [1/1_000_000, 1_000_000].
Strongly recommended when any |coordinate| is outside 2^±24 or [1/16_000_000, 16_000_000].

If either of the points is at the origin, the result is `NaN` because angle is undefined in that case.

You *must* define a tuple constructor `Tuple(x::YourPointType) = ...` if one does not already exist.
"""
angle(tuple1::NTuple{N,T}, tuple2::NTuple{N,T}) where {N, T<:Floats} = angle(real(T), tuple1, tuple2)

# because of the broadcasts .- and .+ below, it is essential to check size compatibility
# for arrays, whereas for tuples the consistency ensured by the "N" in the type
function angle(a1::AbstractArray{T}, a2::AbstractArray{T}) where {T<:Floats}
	size(a1) == size(a2) || throw(DimensionMismatch())
    return angle(real(T), a1, a2)
end

function angle(::Type{T}, point1, point2) where {T<:AbstractFloat}
    unitpoint1 = unitize(point1)
    unitpoint2 = unitize(point2)

    y = unitpoint1 .- unitpoint2
    x = unitpoint1 .+ unitpoint2

    a = 2 * atan(norm(y) / norm(x))

    !(signbit(a) || signbit(T(pi) - a)) ? a : (signbit(a) ? zero(T) : T(pi))
end

# this method allows the arrays to have different types, even non-float types
function angle(a1::AbstractArray{T1}, a2::AbstractArray{T2}) where {T1 <: Number, T2 <: Number}
    T = float(promote_type(T1, T2)) # the "T" for T(pi) must be a float type
    return angle(convert(AbstractArray{T}, a1), convert(AbstractArray{T}, a2))
end

@inline angle(point1::T, point2::T) where {T} = angle(Tuple(point1), Tuple(point2))

end # AngleBetweenVectors


using Flux
using CuArrays

bn=BatchNorm(3,ϵ=Float32(1e-7),momentum=Float32(0.1)) |> gpu

@show typeof(bn.β)
@show typeof(bn.γ)
@show typeof(bn.λ)
@show typeof(bn.μ)
@show typeof(bn.σ)
@show typeof(bn.ϵ)
@show typeof(bn.momentum)

data = rand(28,28,3,2) |> gpu
bn(data)


using DiscretisedFluidQueues
using LinearAlgebra, SparseArrays, StableRNGs
using Test
include("SFM_operators.jl")

T = [-2.5 2 0.5; 1 -2 1; 1 2 -3]
T_nz = T - [0.01 0 0;0 0 0;0 0 0]
T_warn = T - [0.000001 0 0;0 0 0;0 0 0]

C = [0.0; 2.0; -3.0]
m = -1 .+ 2*Int.(DiscretisedFluidQueues._strictly_pos.(C))
S = PhaseSet(C)

model = BoundedFluidQueue(T,S,12.0)

T_aug = [-2.5 0 2 0.5; 0 -2.5 2 0.5; 1 0 -2 1; 0 1 2 -3]
C_aug = [0.0;-0.0;C[2:3]]
m_aug = -1 .+ 2*Int.(DiscretisedFluidQueues._strictly_pos.(C_aug))
S_aug = PhaseSet(C_aug)
am = BoundedFluidQueue(T_aug,S_aug,12.0)

nds_vec = Array{Any,1}(undef,2)
nds_vec[1] = [0.0;3.0;4.0;12.0]
nds_vec[2] = 0.0:4.0:12.0
for nds in 1:2 # probably overkill to test everything twice but oh well!
    global nodes=nds_vec[nds]

    global nbases=3
    global dgmesh=DGMesh(nodes,nbases)

    global am=augment_model(model)

    global fv_order=3
    global fvmesh=FVMesh(nodes,fv_order)

    global order=3
    global frapmesh=FRAPMesh(nodes,order)

    @testset begin

        include("models.jl")

        include("mesh.jl")

        include("generators.jl")

    # test 5_SFM_operators

        include("me_tools.jl")
        include("polynomials.jl")

        include("distributions.jl")

        include("time_integration.jl")

        include("sims.jl")

        include("numerics.jl")

        include("limiters.jl")

    # more testing for lazy_generators now with new modularised code
    # testing for DiscretisedFluidQueue
    # etc...
    end
end


file = open(f->read(f, String), "day_13.txt")

ids = [x.captures[1] for x in eachmatch(r"(\d+)",file)]
ids = parse.(Int64,ids)

(start,ids) = (ids[1], ids[2:end])

calulate_delay(x) = x -((start) % x )

index_min = argmin(calulate_delay.(ids))

print("awnser1 : $(ids[index_min] * calulate_delay(ids[index_min])) ")

#we can solve question 2 as a linear algebra problem,
#for example: The input 2,3 can be solve as:
#start_time + 2X = c
#start_time + 3X = c -1

#There is problably a way to find the basis of the problem and resolve it;
#I dont know now, lets just learn more about linear algebra, RIP

a = [2 1;1 0]
b = [4;3]
a\b


using Test
using NeXLCore

@testset "Electron" begin

    @test isapprox(λₑ(1.0), 1.227e-7, atol = 0.001e-7)
    @test isapprox(λₑ(100.0), 1.227e-8, atol = 0.001e-8)
    @test isapprox(λₑ(1.0e4), 1.227e-9, atol = 0.001e-9)
    @test isapprox(λₑ(5.0e3), 0.0173 * 1.0e-7, atol = 0.0001 * 1.0e-7)
    @test isapprox(λₑ(10.0e3), 0.0122 * 1.0e-7, atol = 0.0001 * 1.0e-7)
    @test isapprox(λₑ(200.0e3), 0.0027 * 1.0e-7, atol = 0.0001 * 1.0e-7)

    @test isapprox(
        kₑ(1.0),
        0.01 * sqrt(2.0 * 9.109e-31 * 1.602e-19 * 1.0) / 1.05457e-34,
        rtol = 1.0e-3,
    )
    @test isapprox(
        kₑ(10.0),
        0.01 * sqrt(2.0 * 9.109e-31 * 1.602e-19 * 10.0) / 1.05457e-34,
        rtol = 1.0e-3,
    )
    @test isapprox(
        kₑ(20.0),
        0.01 * sqrt(2.0 * 9.109e-31 * 1.602e-19 * 20.0) / 1.05457e-34,
        rtol = 1.0e-3,
    )
    @test isapprox(
        kₑ(200.0),
        0.01 * sqrt(2.0 * 9.109e-31 * 1.602e-19 * 200.0) / 1.05457e-34,
        rtol = 1.0e-3,
    )
    @test isapprox(
        kₑ(2000.0),
        0.01 * sqrt(2.0 * 9.109e-31 * 1.602e-19 * 2000.0) / 1.05457e-34,
        rtol = 1.0e-3,
    )

    @test isapprox(mₑ, 0.511e6, atol = 0.001e6)  # Must be in eV!!!!
end


function _besselk_ser(v, x, maxit, tol, modify)
  T    = promote_type(typeof(x), typeof(v))
  out  = zero(T)
  oneT = one(T)
  twoT = oneT+oneT
  # precompute a handful of things:
  xd2      = x/twoT
  xd22     = xd2*xd2
  half     = oneT/twoT
  if modify
    e2v      = exp2(v)
    xd2_v    = (x^(2*v))/e2v
    xd2_nv   = e2v
  else
    lxd2     = log(xd2)
    xd2_v    = exp(v*lxd2)
    xd2_nv   = exp(-v*lxd2)
  end
  gam_v    = gamma(v)
  gam_nv   = gamma(-v)
  gam_1mv  = -gam_nv*v # == gamma(one(T)-v)
  gam_1mnv = gam_v*v   # == gamma(one(T)+v)
  xd2_pow  = oneT
  fact_k   = oneT
  floatk   = convert(T, 0)
  (gpv, gmv) = (gam_1mnv, gam_1mv)
  # One final re-compression of a few things:
  _t1 = gam_v*xd2_nv*gam_1mv
  _t2 = gam_nv*xd2_v*gam_1mnv
  # now the loop using Oana's series expansion, with term function manually
  # inlined for max speed:
  for _j in 0:maxit
    t1   = half*xd2_pow
    tmp  = _t1/(gmv*fact_k)
    tmp += _t2/(gpv*fact_k)
    term = t1*tmp
    out += term
    ((abs(term) < tol) && _j>5) && return out
    # Use the trick that gamma(1+k+1+v) == gamma(1+k+v)*(1+k+v) to skip gamma calls:
    (gpv, gmv) = (gpv*(oneT+v+floatk), gmv*(oneT-v+floatk))
    xd2_pow *= xd22
    fact_k  *= (floatk+1)
    floatk  += 1.0
  end
  throw(error("$maxit iterations reached without achieving atol $tol."))
end


"""
    SeisAddNoise(d, snr; <keyword arguments>)

Add noise to an N-dimensional input
data `d`. Should specify the signal-to-noise ratio level `snr`.
Noise can be band limited using kewyord `L`.

# Arguments
- `d::Array{Real, N}`: N-dimensional data.
- `snr::Real`: signal-to-noise ratio.
- `db::Bool=false`: Flag is false if snr is given by amplitude. Flag is true if
snr is given in dB.
- `pdf::AbstractString="gaussian"`: random noise probability distribution:
`"gaussian"` or `"uniform"`.
- `L::Int=1`: averaging operator length to band-limit the random noise.

# Examples
```julia
julia> using PyPlot
julia> w = Ricker(); wn = SeisAddNoise(w, 2); plot(w); plot(wn);
MeasureSNR(w, wn)
julia> d = SeisHypEvents(); dn = SeisAddNoise(d, 1.0, db=true, L=9);
SeisPlotTX([d dn]); MeasureSNR(d, dn, db=true)
```
Credits: Juan I. Sabbione, 2016
"""
function SeisAddNoise(d::Array{T, N}, snr::Real; db::Bool=false,
                      pdf::AbstractString="gaussian", L::Int=1) where {T <: Real, N}

    noise = GenNoise(size(d), pdf, L=L)
    if db==false
        noise = noise/norm(noise) * norm(d)/snr
    elseif db==true
        noise = noise/norm(noise) * norm(d)/10.0^(0.05*snr)
    end
    noisy = d + noise
    @assert(abs(MeasureSNR(d, noisy, db=db) - snr) < 1e10*eps(AbstractFloat(snr)))
    return noisy

end

# Generate the noise
function GenNoise(dims::Tuple, pdf::AbstractString; L::Int=1)
    N = length(dims)
    n1 = dims[1]
    N == 1 ? nx = 1 : nx = *(dims[2:end]...)
    hamm = Hamming(L)
    L2 = floor(Int,L/2)
    if pdf=="gaussian"
        noise = randn(n1, nx)
    elseif pdf=="uniform"
        noise = -1 + 2*rand(n1, nx)
    else
        error("pdf must be gussian or uniform")
    end
    for j = 1:nx
         noise[:,j] = conv(noise[:,j], hamm)[L2+1:n1+L2]
    end
    reshape(noise, dims...)
end


@testset "IDW" begin
  # basic problem
  geodata = georef((variable=[1.,0.,1.],), [25. 50. 75.;  25. 75. 50.])
  domain  = CartesianGrid(100,100)
  problem = EstimationProblem(geodata, domain, :variable)

  solver = IDW(:variable => (neighbors=3,))

  solution = solve(problem, solver)

  if visualtests
    @test_reference "data/idw.png" contourf(solution,size=(800,400))
  end

  # haversine distance
  geodata = georef((variable=[4.0,-1.0,3.0],), [50.0 100.0 200.0; -30.0 30.0 10.0])
  domain  = CartesianGrid((1.0, -89.0), (359.0, 89.0), dims=(200, 100))
  problem = EstimationProblem(geodata, domain, :variable)

  solver = IDW(:variable => (neighbors=3, distance=Haversine(1.0)))

  solution = solve(problem, solver)

  if visualtests
    @test_reference "data/idw-haversine.png" contourf(solution,size=(800,200))
  end
end


using LightGraphs
using MonteCarloX
using Random
using Distributions
import Distributions.pdf
import Distributions.cdf
using StatsBase
import StatsBase.kldivergence

include("utils.jl")

""" Testing reweighting on 2D Ising model"""
function test_ising_reweighting(;verbose = false)
    log_P(E, beta) = -beta * E
    P(E, beta) = exp(-beta * E)

    list_beta = [0.0,0.3,0.7,1.0,1.5,2.0]
    pass = true

    rng = MersenneTwister(1000)
    system = constructIsing([8,8], rng)
    nearest_neighbors(index) = outneighbors(system.lattice, index)
    N = length(system.spins)

    for beta in list_beta
        samples = 1000
        list_energy = zeros(samples)
        # thermalization 100 sweeps
        for sweep in 1:samples + 100
            # Metropolis:
            for i in 1:N
                update_spin_flip(system, beta, rng)
            end
            if sweep > 100
                list_energy[sweep - 100] = E(system)
            end
        end
        H_meas = fit(Histogram, list_energy, minimum(list_energy):4:maximum(list_energy) + 4, closed = :left)
        P_meas = normalize(H_meas, mode = :probability)
        # reweighting both with timeseries (list_energy) and histrogram
        beta_source = beta
        beta_target = beta + 0.05
        log_P_source(E) = log_P(E, beta_source)
        log_P_target(E) = log_P(E, beta_target)
        P_source(E) = P(E, beta_source)
        P_target(E) = P(E, beta_target)
        E_ref = analytic_expectation_value_E(beta_target, log_dos_beale_8x8)
        P_ref_source = BoltzmannDistribution(beta_source, log_dos_beale_8x8).pdf
        P_ref_target = BoltzmannDistribution(beta_target, log_dos_beale_8x8).pdf
        kld_ref = kldivergence(P_meas, x->P_ref_source[x])
        if verbose
            println("result analytic = $(E_ref)")
        end
        ###########################################################################
        E_reweight_list1 = MonteCarloX.expectation_value_from_timeseries(log_P_target, log_P_source, list_energy, list_energy)
        E_reweight_list1_error = abs((E_reweight_list1 - E_ref) / E_ref)
        if verbose
            println("result expectation_value_from_timeseries_log: <E> = $(E_reweight_list1); difference from exact = $(E_reweight_list1_error)")
        end
        pass &= E_reweight_list1_error < 0.1
        ###########################################################################
        P_reweight_list = MonteCarloX.distribution_from_timeseries(log_P_target, log_P_source, list_energy, minimum(list_energy):4:maximum(list_energy) + 4)
        kld_source = kldivergence(P_reweight_list, x->P_ref_source[x])
        kld_target = kldivergence(P_reweight_list, x->P_ref_target[x])
        if verbose
            println("result distribution_from_timeseries_log: kld_target ($(kld_target)) !< kld_source ($(kld_source))")
        end
        # only compare reweighted distribution to target and source. There result is comparable.
        pass &= kld_target < kld_source
        ###########################################################################
        E_reweight_hist1 = MonteCarloX.expectation_value_from_histogram(E->E, log_P_target, log_P_source, H_meas)
        E_reweight_hist1_error = abs((E_reweight_hist1 - E_ref) / E_ref)
        if verbose
            println("result expectation_value_from_histogram_log-1: <E> = $(E_reweight_hist1); difference from exact = $(E_reweight_hist1_error)")
            println("result expectation_value_from_histogram_log-1: <E> = $(E_reweight_hist1); difference from timeseries = $(E_reweight_hist1 - E_reweight_list1)")
        end
        pass &= E_reweight_hist1_error < 0.1
        ###########################################################################
        hist_obs = zero(H_meas)
        for E in list_energy
            index = StatsBase.binindex(hist_obs, E)
            hist_obs.weights[index] += E
        end
        E_reweight_hist2 = MonteCarloX.expectation_value_from_histogram(log_P_target, log_P_source, H_meas, hist_obs)
        E_reweight_hist2_error = abs(E_reweight_hist2 - E_reweight_list1)
        if verbose
            println("result expectation_value_from_histogram_log: <E> = $(E_reweight_hist2); difference from timeseries = $(E_reweight_hist2_error)")
        end
        pass &= E_reweight_hist2_error < 0.1
        ###########################################################################
        # TODO: canonical reweighting
    end

    # TODO: multi-histogram reweighting

    return pass
end

""" Testing Metropolis accept on 2D Ising model"""
function test_ising_metropolis(;verbose = false)
    list_beta = [0.0,0.3,0.7,1.0,1.5,2.0]
    pass = true

    rng = MersenneTwister(1000)
    system = constructIsing([8,8], rng)
    nearest_neighbors(index) = outneighbors(system.lattice, index)
    N = length(system.spins)

    for beta in list_beta
        samples = 1000
        list_energy = zeros(samples)
        # thermalization 100 sweeps
        for sweep in 1:samples + 100
            # Metropolis:
            for i in 1:N
                update_spin_flip(system, beta, rng)
            end
            if sweep > 100
                list_energy[sweep - 100] = E(system)
            end
        end
        P_meas = fit(Histogram, list_energy, minimum(list_energy):4:maximum(list_energy) + 4, closed = :left)
        P_meas = normalize(P_meas, mode = :probability)
        P_true = BoltzmannDistribution(beta, log_dos_beale_8x8).pdf
        # KL divergence sum P(E)logP(E)/Q(E)
        # P=P_meas, Q=P_true s.t P(E)=0 simply ignored
        kld = abs(kldivergence(P_meas, x->P_true[x]))
        if verbose
            println(beta, " ", kld, " ", length(keys(P_meas)))
        end
        pass &= kld < 0.1
    end

    return pass
end

""" Testing Cluster update on 2D Ising model"""
function test_ising_cluster(;verbose = false)
    # avoid beta_c = 0.44
    # why does it consistently fail around beta=0.6
    list_beta = [0,0.3,0.5,1,1.5,2.0]
    pass = true

    rng = MersenneTwister(1000)
    system = constructIsing([8,8], rng)
    nearest_neighbors(index) = outneighbors(system.lattice, index)
    N = length(system.spins)

    cluster_algorithm = MonteCarloX.ClusterWolff()

    for beta in list_beta
        samples = Int(1e4)
        list_energy = zeros(samples)

        # thermalization 100 sweeps
        for sweep in 1:samples + 100
            for i in 1:5
                MonteCarloX.update(cluster_algorithm, rng, system.spins, nearest_neighbors, beta)
            end
            if sweep > 100
                list_energy[sweep - 100] = E(system)
            end
        end

        #
        P_meas = fit(Histogram, list_energy, minimum(list_energy):4:maximum(list_energy) + 4, closed = :left)
        P_meas = normalize(P_meas, mode = :probability)
        P_true = BoltzmannDistribution(beta, log_dos_beale_8x8).pdf

        kld = abs(kldivergence(P_meas, x->P_true[x]))
        if verbose
            println("result kldivergence P_meas to P_true = $(kld)")
        end
        pass &= kld < 0.02
    end

    return pass
end

###############################################################################
###############################################################################
### Details functions

struct IsingSystem{F}
    dims::Vector{Int}
    lattice::SimpleGraph
    spins::Vector{Int}
    nearest_neighbors::F
end

function constructIsing(dims::Vector{Int}, rng::AbstractRNG)::IsingSystem
    lattice = LightGraphs.SimpleGraphs.grid(dims, periodic = true)
    spins   = rand(rng, [-1,1], nv(lattice))
    nearest_neighbors(i) = outneighbors(lattice, i)
    return IsingSystem(dims, lattice, spins, nearest_neighbors)
end

function M(system::IsingSystem)::Int
    return abs(sum(system.spins))
end

function analytic_expectation_value_E(beta::Float64, log_dos::Vector{Tuple{Int,Float64}})::Float64
    mean_E  = 0.0
    norm    = 0.0
    for (E, log_d) in log_dos
        mean_E  +=     E * exp(log_d - beta * E)
        norm    +=         exp(log_d - beta * E)
    end
    mean_E /= norm
    return mean_E
end

# J=1
# E=-J*sum s_i s_j
function E(system::IsingSystem)::Int
    E = 0
    for i in 1:length(system.spins)
        E += E_local(system, i)
    end
    return E / 2
end

function E_local(system::IsingSystem, index::Int)::Int
    e = 0
    for j in system.nearest_neighbors(index)
        e -= system.spins[index] * system.spins[j]
    end
    return e
end

function update_spin_flip(system::IsingSystem, beta::Float64, rng::AbstractRNG)::Int
    index = rand(rng, 1:length(system.spins))
    dE        = -2 * E_local(system, index)
    if accept(Metropolis(), rng, beta, dE)
        system.spins[index] *= -1
    else
        dE = 0
    end
    return dE
end

# """ Definition of the Boltzmann distribution for hypothesis testing."""
struct BoltzmannDistribution <: ContinuousUnivariateDistribution
    pdf::Dict{Any,Float64}
    #cdf::Dict{Any,Float64}

    BoltzmannDistribution(beta, log_dos) = new(initialize_BoltzmannDistribution(beta, log_dos))
    # BoltzmannDistribution(beta, log_dos, log_Z) = new(Float64(beta), Dict{Any,Float64}(log_dos), Float64(log_Z))
end

# Probability distribution (for both precompute in constructor...)
pdf(d::BoltzmannDistribution, E::Real) = d.pdf[E]

# cumulated probability distribution
#cdf(d::BoltzmannDistribution, E::Real) = d.cdf[E]

function initialize_BoltzmannDistribution(beta, log_dos)
    # partition sum
    log_Z = -Inf
    for (E, log_d) in log_dos
        log_Z = log_sum(log_Z, log_d - beta * E)
    end

    pdf = Dict{Int64,Float64}()
    cdf = Dict{Int64,Float64}()

    log_cdf = 0.0
    for (E, log_d) in log_dos
        log_pdf = log_d - beta * E - log_Z
        log_cdf = log_sum(log_cdf, log_pdf)
        pdf[E] = exp(log_pdf)
        cdf[E] = exp(log_cdf)
    end
    #return pdf, cdf
    return pdf
end

log_dos_beale_8x8 = [ (-128, 0.6931471805599453), (-124, 0.0), (-120, 4.852030263919617), (-116, 5.545177444479562), (-112, 8.449342524508063), (-108, 9.793672686528922), (-104, 11.887298863200714), (-100, 13.477180596840947), (-96, 15.268195474147658), (-92, 16.912371686315282), (-88, 18.59085846191256), (-84, 20.230089202801466), (-80, 21.870810400320693), (-76, 23.498562234123614), (-72, 25.114602234581373), (-68, 26.70699035290573), (-64, 28.266152815389898), (-60, 29.780704423363996), (-56, 31.241053997806176), (-52, 32.63856452513369), (-48, 33.96613536105969), (-44, 35.217576663643314), (-40, 36.3873411250109), (-36, 37.47007844691906), (-32, 38.46041522581422), (-28, 39.35282710786369), (-24, 40.141667825183845), (-20, 40.82130289691285), (-16, 41.38631975325592), (-12, 41.831753810069756), (-8, 42.153328313883975), (-4, 42.34770636939425), (0, 42.41274640460084), (4, 42.34770636939425), (8, 42.153328313883975), (12, 41.831753810069756), (16, 41.38631975325592), (20, 40.82130289691285), (24, 40.141667825183845), (28, 39.35282710786369), (32, 38.46041522581422), (36, 37.47007844691906), (40, 36.3873411250109), (44, 35.217576663643314), (48, 33.96613536105969), (52, 32.63856452513369), (56, 31.241053997806176), (60, 29.780704423363996), (64, 28.266152815389898), (68, 26.70699035290573), (72, 25.114602234581373), (76, 23.498562234123614), (80, 21.870810400320693), (84, 20.230089202801466), (88, 18.59085846191256), (92, 16.912371686315282), (96, 15.268195474147658), (100, 13.477180596840947), (104, 11.887298863200714), (108, 9.793672686528922), (112, 8.449342524508063), (116, 5.545177444479562), (120, 4.852030263919617), (124, 0.0), (128, 0.6931471805599453) ];

"""
Kullback-Leibler divergence between two distributions

distributions are here considered to be dictionaries

valid for n-dimensional dictionaries (args=Tuple or larger?)
"""
function kldivergence(P::Dict, Q::Dict)
        ##KL divergence sum P(args)logP(args)/Q(args)
        ## P=P_meas, Q=P_true s.t P(args)=0 simply ignored
        #
        kld = 0.0
        for (args, p) in P
            q = Q[args]
            kld += p * log(p / q)
        end
        return kld
end


using LinearAlgebra
using MinimumVolumeEllipsoids
using PDMats
using Random
using Statistics
using Test

include("minimumvolumeellipsoid.jl")
include("minvol.jl")
include("rand.jl")
include("volume.jl")


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -8.186501529651974 <= q <= 9.622483160840265)

                     add_NL_constraint(m, :(1/(1 + exp(-(-0.9342689023247241 + -0.06643065035242302*1/(1 + exp(-(-0.6967397920955185 + 0.4612305646899797*1/(1 + exp(-(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])))) + 0.5917923558177942*1/(1 + exp(-(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))))))) + 0.13022424784864484*1/(1 + exp(-(-0.8092681495291991 + -0.6577545463201151*1/(1 + exp(-(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])))) + 0.816846389391447*1/(1 + exp(-(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4])))))))))) + 1/(1 + exp(-(0.9502905692637253 + 0.2731932353707758*1/(1 + exp(-(-0.6967397920955185 + 0.4612305646899797*1/(1 + exp(-(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])))) + 0.5917923558177942*1/(1 + exp(-(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))))))) + 0.7525923835832526*1/(1 + exp(-(-0.8092681495291991 + -0.6577545463201151*1/(1 + exp(-(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])))) + 0.816846389391447*1/(1 + exp(-(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4])))))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using ADCME
using LinearAlgebra
using AdFem
using DelimitedFiles

p = readdlm("fenics/x.txt")[:]
f = readdlm("fenics/f.txt")[:]

mesh = Mesh(8, 8, 1. /8)
f0 = compute_interaction_term(p, mesh)
sess = Session(); init(sess)
f1 = run(sess, f0)

@show norm(f - f1)

p = readdlm("fenics/x2.txt")[:]
f = readdlm("fenics/f2.txt")[:]

mesh = Mesh(8, 8, 1. /8, degree=2)
f0 = compute_interaction_term(p, mesh)
sess = Session(); init(sess)
f1 = run(sess, f0)

E = Int64.(readdlm("fenics/edges.txt"))
Edof = get_edge_dof(E, mesh)
DOF = [1:mesh.nnode; Edof .+ mesh.nnode]
f1 = f1[[DOF; DOF .+ mesh.ndof]]
@show norm(f - f1)


using LinearAlgebra, Images, FileIO, Colors

function sinkhorn( a, b, precision, costMatrix, n_iter )
	n, m = length( a ), length( b )
	C = costMatrix
	costMatrix = costMatrix ./ max( costMatrix... ) # renormalizing to avoid underflows in division
	K = exp.( -costMatrix/precision )
	u, v = ones( n ), ones( m )
	for i in 1:n_iter
		# u = a ./ K*v
		u = exp.( log.( a ) - log.( K*v ) )
		v = exp.( log.( b ) - log.( K'*u ) )
	end
	# P = [ u[i]*K[i,j]*v[j] for i in 1:n, j in 1:m ]
	P = u.*K*diagm( v )
	entropy = -sum( P.*log.( P ) )
	return P, sum( P.*C ) - precision*entropy
end

function getImage( path::String, quantization::Int )
	execstring =`convert $path -dither None -colors $quantization $path.temp.png`
	run(execstring)
	img = load( "$path.temp.png" )
	rm( "$path.temp.png" )
	return img
end

function swapPalette( paletteSourceImage::String, paletteTargetImage::String, maxColors=10_000, mode="mean" )
	source = RGB.( getImage( paletteSourceImage, maxColors ) )
	target = RGB.( getImage( paletteTargetImage, maxColors ) )

	scols = unique( source )
	tcols = unique( target )

	# we don't account for frequency of colors, so the distributions are uniform
	scdist = ones( length( scols ) )/ length( scols )
	tcdist = ones( length( tcols ) )/ length( tcols )

	costMatrix = [ colordiff( t, s ) for t in tcols, s in scols ]

	P, d = sinkhorn( tcdist, scdist, 0.01, costMatrix, 100 )

	if mode=="matching"
		perm = [ argmax( P[i,:] ) for i in 1:length( tcols ) ]
		mp = [ tcols[i] => scols[perm][i] for i in 1:length( tcols ) ]
	elseif mode=="mean"
		scols = LCHab.( scols )
		scols = [ [a.l, a.c, a.h] for a in scols ]
		mp = [tcols[i] => LCHab( sum( scols.*P[i,:]/sum( P[i,:] ) )... ) for i in 1:length( tcols )]
	else
		@error "Unknown mode:\t$mode."
	end
	for i in 1:length( tcols )  replace!( target, mp[i] )  end
	return target
end

nargs = length( ARGS )
if 3 <= nargs<= 5
	maxColors = nargs>=4 ? ARGS[4] : 10_000
	mode      = nargs==5 ? ARGS[5] : "mean"
	@time img = swapPalette( ARGS[1], ARGS[2], maxColors, mode )
	save( ARGS[3], img )
else
	println( "Usage: julia paletteSwap.jl paletteSourceImage paletteTargetImage outputFile [maxColors=10_000 mode='mean']" )
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -12.503038214935698 <= q <= 11.138318577023822)

                     add_NL_constraint(m, :(swish(-0.2825986234777189 + -0.06477303738755547*swish(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + 0.23435836886254213*swish(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + 0.7516401137686608*swish(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) + swish(0.12692347233027856 + 0.060984281016326136*swish(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + 0.4767325000240725*swish(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + -0.9515049437955176*swish(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) + swish(-0.022625917828582143 + -0.9740734403326825*swish(0.91111659897616 + 0.32194568310871174*$(x[1]) + -0.04079743189253193*$(x[2]) + -0.7834407613653229*$(x[3]) + -0.5105294854228863*$(x[4])) + -0.9306328050206538*swish(-0.27698098400696525 + 0.765733032903245*$(x[1]) + 0.6374561618281618*$(x[2]) + -0.006430667300674209*$(x[3]) + -0.6595935184751247*$(x[4])) + 0.61848968844355*swish(0.7789283911146407 + -0.8853253179157212*$(x[1]) + -0.6613182381252614*$(x[2]) + 0.5007525575333553*$(x[3]) + 0.7960733427996214*$(x[4]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using Plots
gr()

v = rand(10)
w = rand(10)

print(v + w)

using LinearAlgebra
# soma dos vetores v e w
print(dot(v, w))


using PQEcolaPoint
using Dates

# build a binary tree with three levels -
SF = 82.00
model = build(CRRLatticeModel; Sₒ = SF, number_of_levels=80, σ=0.4510, T=(37.0 / 365), μ=0.0045)

# what is the contract?
short_put_contract = PutContractModel()
short_put_contract.ticker = "XYZ"
short_put_contract.expiration_date = Date(2022, 05, 06)
short_put_contract.strike_price = (75.0)
short_put_contract.premium = 0.0
short_put_contract.number_of_contracts = 1
short_put_contract.direction = 1

# what is the premimum?
parray = premium(short_put_contract, model)


module Utils_states
using LinearAlgebra

export eᵢ,
       maketraceone,
       take_Pᵀ,
       isPPᵀ,
       ψ,
       sq,
       psep

"""The standard basis vector eᵢ in dimension n"""
eᵢ(n::Int,i::Int) = [Int(j==i) for j in 1:n]

""" Returns the partial transpose of the sysemt. """
function take_Pᵀ(ρ,sys,d)
    #d₁,d₂ = d
    d₂,d₁ = d
    D = [(i₁,i₂,j₁,j₂)   for i₁ in 1:d₁, i₂ in 1:d₁, j₁ in 1:d₂, j₂ in 1:d₂]
    rcs =  [(i,j)   for j in 1:d₂  for i in 1:d₁]
    DB₀ = [D[ij[1],kl[1],ij[2],kl[2]] for ij in rcs , kl in rcs ]
    DB₁ = [D[kl[1],ij[1],ij[2],kl[2]] for ij in rcs , kl in rcs ]
    DB₂ = [D[ij[1],kl[1],kl[2],ij[2]] for ij in rcs , kl in rcs ]
    Zar = Dict(zip(DB₀, ρ))
    if sys == 1
        return map(a -> Zar[a],DB₁)
    elseif sys == 2
        return map(a -> Zar[a],DB₂)
    else
        return ρ
    end
end
"""
Checks if a sysetem is a PPT via checking if the smallest eigenvalues > -1.e-16
"""
function isPPᵀ(ρ,d)
    for sys ∈ 0:2
        ρᵀ = take_Pᵀ(ρ, sys, d)
        isPSDr = real(LinearAlgebra.eigvals(ρᵀ)[1]) > -1.e-15
        isPSDi = imag(LinearAlgebra.eigvals(ρᵀ)[1]) <  1.e-10
        ~(isPSDr && isPSDi) ? (return false) : 0
    end
    return true
end
"""Divides ρ by its trace"""
maketraceone(ρ) = ρ/tr(ρ)

""" diagM(1/√diag(ρ))*ρ*diagM(1/√diag(ρ))"""
# function makediagone(ρ)
#     mpli = 1 ./ map(x -> iszero(x) ? 1 : x,  sqrt.([ρ[i,i] for i in 1:size(ρ)[1]]))
#     return diagm(mpli)*ρ*diagm(mpli)
# end

eᵢ(n::Int,i::Int) = [j==i for j in 1:n] #"""The standard basis vector eᵢ in dimension n"""
ψ(i₀::Int,i₁::Int,n::Int) = kron(eᵢ(n,i₀),eᵢ(n,i₁))
ψ(n::Int) = kron(eᵢ(n,1) + eᵢ(n,2),eᵢ(n,1) + eᵢ(n,2))
sq(ϕ) =  ϕ*transpose(conj.(ϕ))
psep(i₀::Int,i₁::Int,n::Int) = sq(ψ(i₀,i₁,n))
# all(isreal.(ϕ)) ? ϕ*transpose(ϕ) :
end