module Yields

import BSplineKit
import ForwardDiff
using LinearAlgebra
using UnicodePlots
using Roots

# don't export type, as the API of Yields.Zero is nicer and
# less polluting than Zero and less/equally verbose as ZeroYieldCurve or ZeroCurve
export rate, discount, accumulation, forward,
LinearSpline, QuadraticSpline, Periodic, Continuous

include("AbstractYield.jl")
include("Rate.jl")
const DEFAULT_COMPOUNDING = Yields.Continuous()

include("utils.jl")
include("bootstrap.jl")
include("SmithWilson.jl")
include("generics.jl")
include("RateCombination.jl")

end


module random

export randstate, randoperator

using ..bases, ..states, ..operators_dense

"""
    randstate(basis)

Calculate a random normalized ket state.
"""
function randstate(b::Basis)
    psi = Ket(b, rand(Complex128, length(b)))
    normalize!(psi)
    psi
end

"""
    randoperator(b1[, b2])

Calculate a random unnormalized dense operator.
"""
randoperator(b1::Basis, b2::Basis) = DenseOperator(b1, b2, rand(Complex128, length(b1), length(b2)))
randoperator(b::Basis) = randoperator(b, b)

end #module


# import libraries
using GeometricIntegrators
using ProgressMeter

# import ParticleMethods package
using ParticleMethods

# import HDF5 package for storing solution
using HDF5

# parameters
Δt = 0.1    # time step size
nt = 200    # number of time steps
np = 10000  # number of particles
nx = 16     # number of grid points

# output file
h5file = "vlasov_poisson.hdf5"

# random initial conditions
x₀ = randn(np)
v₀ = randn(np)

# shift x₀ to the interval [0,1]
xmax = ceil(maximum(abs.(x₀)))
x₀ .+= xmax
x₀ ./= 2*xmax

# concatenate x₀ and v₀
z₀ = collect(hcat(x₀, v₀)')

# vector field
function lorentz_force!(t, z, ż, p::PoissonSolver)
    Particles.solve!(p, z[1,:])
    for i in axes(ż, 2)
        ż[1,i] = z[2,i]
        ż[2,i] = eval_field(p, z[1,i])
    end
end

# splitting fields
function v_advection!(t, z, ż)
    for i in axes(ż, 2)
        ż[1,i] = z[2,i]
        ż[2,i] = 0
    end
end

function s_advection!(t, z, s, h)
    for i in axes(s, 2)
        s[1,i] = z[1,i] + h * z[2,i]
        s[2,i] = z[2,i]
    end
end

function v_lorentz_force!(t, z, ż, p::PoissonSolver)
    Particles.solve!(p, z[1,:])
    for i in axes(ż, 2)
        ż[1,i] = 0
        ż[2,i] = eval_field(p, z[1,i])
    end
end

function s_lorentz_force!(t, z, s, h, p::PoissonSolver)
    Particles.solve!(p, z[1,:])
    for i in axes(s, 2)
        s[1,i] = z[1,i]
        s[2,i] = z[2,i] + h * eval_field(p, z[1,i])
    end
end

# solution storage
function copy_to_hdf5(h5z, z, n)
    h5z[:,:,n+1] = z
end

# create Poisson solver
p = PoissonSolverFFT{eltype(z₀)}(nx)
Particles.solve!(p, x₀)

# create an ODE instance
# ode = ODE((t, z, ż) -> lorentz_force!(t, z, ż, p), z₀)

# create a splitting ODE instance
sode = SODE((v_advection!, (t, z, ż) -> v_lorentz_force!(t, z, ż, p)),
            (s_advection!, (t, z, s, h) -> s_lorentz_force!(t, z, s, h, p)),
            z₀)

# create HDF5 file and copy initial conditions
h5  = h5open(h5file, "w")
h5z = create_dataset(h5, "z", eltype(z₀), ((2, np, nt+1), (2, np, -1)), chunk=(2,np,1))
copy_to_hdf5(h5z, z₀, 0)

# create integrator
# int = IntegratorExplicitEuler(ode, Δt)
int = Integrator(sode, TableauStrang(), Δt)

# create atomic solution
# asol = AtomicSolution(ode)
asol = AtomicSolution(sode)

# copy initial conditons to atomic solution
# set_initial_conditions!(asol, ode)
set_initial_conditions!(asol, sode)

# initilize integrator
initialize!(int, asol)

# loop over time steps showing progress bar
try
    @showprogress 5 for n in 1:nt
        integrate_step!(int, asol)
        copy_to_hdf5(h5z, asol.q, n)
    end
finally
    # close HDF5 file
    close(h5)
end

# load Plots package
using Plots

# read array from HDF5 file
z = h5read(h5file, "z")

# compute plot ranges
vmax = ceil(maximum(abs.(v₀)))
xlim = (0, 1)
vlim = (-vmax, +vmax)

# plot initial condition
scatter(mod.(z[1,:,begin], 1), z[2,:,begin],
        marker = 3,
        xlim = xlim,
        ylim = vlim,
        title = "Vlasov-Poisson",
        legend = false,
        size = (800, 600)
)

# save figure to file
savefig("vlasov_poisson_z₀.png")

# create animation
anim = @animate for n in axes(z,3)
    scatter(mod.(z[1,:,n], 1), z[2,:,n],
        marker = 3,
        xlim = xlim,
        ylim = vlim,
        title = "Vlasov-Poisson",
        legend = false,
        size = (800, 600)
    )
end

# save animation to file
gif(anim, "vlasov_poisson_anim.gif", fps=10)


raw = read(stdin, String)
raw = convert(Array{Any}, split(raw, ".\n", keepempty=false))

# raw[i] = [left, [[count 1, color 1], ...]]
for (i, r) in enumerate(raw)
    raw[i] = convert(Array{Any}, map(strip, split(r, "contain")))

    # Left
    raw[i][begin] = raw[i][begin][begin:end - length(" bags")]

    # Right
    raw[i][2] = convert(Array{Any}, map(strip, split(raw[i][2], ',')))

    if raw[i][2][begin] == "no other bags"
        raw[i][2] = []
        continue
    end

    for (j, rule) in enumerate(raw[i][2])
        first = findfirst(' ', rule)
        last = findlast(' ', rule)

        count, color = parse(Int, strip(rule[1:first - 1])), strip(rule[first + 1:last - 1])
        raw[i][2][j] = [count, color]
    end
end

# Every color containing a shiny gold
shiny = Set(["shiny gold"])
lastcount = nothing
while lastcount != length(shiny)
    global lastcount = length(shiny)

    for (left, right) in raw
        # If an item in right contains a shiny gold bag, left
        # contains a shiny gold bag
        any((x) -> x[2] ∈ shiny, right) && (push!(shiny, left))
    end
end

# Without shiny gold
println("Part 1 : $(length(shiny) - 1)")

name2prod = Dict(left => right for (left, right) in raw)
name2count = Dict()

# DP approach (memoization)
function count(color)
    global name2count

    if haskey(name2count, color)
        return name2count[color]
    end

    c = 1
    # For each right item
    for (number, name) in name2prod[color]
        c += number * count(name)
    end

    name2count[color] = c

    return c
end

println("Part 2 : $(count("shiny gold") - 1)")


import SparseArrays.sparse

function lar2plasm(EV, FE)
    EV = sparse(EV)
    faces = Array{Array{Int64, 1}, 1}()

    for f in 1:size(FE, 1)

        verts = Array{Int64, 1}()
        done = Array{Int64, 1}()
        face = FE[f, :]
        edges = EV[face.nzind, :]

        push!(done, 1)
        vs = edges[1, :].nzind
        if face.nzval[1] < 0
            vs = vs[end:-1:1]
        end
        startv, nextv = vs

        # println(startv)
        push!(verts, startv)
        while nextv != startv
            e = setdiff(edges[:, nextv].nzind, done)[1]
            push!(done, e)
            vs = edges[e, :].nzind
            if face.nzval[e] < 0
                vs = vs[end:-1:1]
            end
            curv, nextv = vs
            push!(verts, curv)
        end

        push!(faces, verts)
    end

    "FV = " * string(faces)[15:end]
end

function lar2py(V, EV)
    text = "V = [[0,0]"
    for i in 1:size(V, 1)
        text = text * "," * string(vec(V[i, :]))
    end
    text = text * "]\n"

    text = text * "EV = ["
    for i in 1:size(EV, 1)
        if i > 1 text = text * "," end
        text = text * string(EV[i, :].nzind)
    end
    text = text * "]"
    return text
end

function visualize_numbers(set, font_size)

    f = open("tmp.py", "w")

    write(f, "from larlib import *\n" *
        lar2py(set[1], set[2]) * "\n" *
        lar2plasm(set[2], set[3]) * "\n" *
        "VV = AA(LIST)(range(len(V)))\n" *
        "submodel = STRUCT(MKPOLS((V,EV)))\n" *
        "VIEW(larModelNumbering(1,1,1)(V,[VV,EV,FV],submodel,$font_size))\n")

    close(f)
    run(`python tmp.py`)
end

function visualize_boundary(set, boundary, font_size)
    f = open("tmp.py", "w")

    write(f, "from larlib import *\n" *
        lar2py(set[1], set[2]) * "\n" *
        lar2plasm(set[2], set[3]) * "\n" *
        "VV = AA(LIST)(range(len(V)))\n" *
        "submodel = STRUCT(MKPOLS((V,EV)))\n" *
        "model = larModelNumbering(1,1,1)(V,[VV,EV,FV],submodel,$font_size)\n" *
        lar2py(set[1], set[2][boundary.nzind, :]) * "\n" *
        "overmodel = STRUCT(MKPOLS((V,EV)))\n" *
        "VIEW(STRUCT([model, COLOR(RED)(overmodel)]))")

    close(f)
    run(`python tmp.py`)
end


#=

  Test of uniform draw in Turing.jl

  What I can see Turing.jl don't have a function to randomly draw
  an element from a list, e.g.
     uniformDraw([1,10,32,100])
  It use have DiscreteUniform() and Categorical()

  So I'm trying to simulating uniformDraw....

  Asked about this at Discourse:
  # https://discourse.julialang.org/t/turing-jl-unformdraw-distribution/65123
  https://discourse.julialang.org/t/turing-jl-uniformdraw-distribution/65123

  Ah, it's there as DiscreteNonParametric()
  https://juliastats.org/Distributions.jl/stable/univariate/#Discrete-Distributions

  I was looking in an old documentation of Distributions.jl!

=#

using Turing
include("jl_utils.jl")

# Defined in jl_utils.jl
# UniformDraw(xs) = DiscreteNonParametric(xs, ones(length(xs)) ./ length(xs))

@model function uniformDrawTest(x,pcts)
    #=
    function uniformDraw(x)
        n = length(x)
        ix1 ~ DiscreteUniform(1,n)
        return x[ix1]
    end
    function uniformDraw(x,pcts)
        n = length(x)
        ix2 ~ Categorical(simplex(pcts))
        return x[ix2]
    end
    val1 ~ Dirac(uniformDraw(x))
    val2 ~ Dirac(uniformDraw(x,pcts))
    val3 ~ Dirac(uniformDraw(x))
    =#
    val1 ~ UniformDraw(x)
    val2 ~ DiscreteNonParametric(x,simplex(pcts))
    val3 ~ UniformDraw(x)
    println([val1,val2,val3])
end

x = [1,10,32,100]
pcts = [1,2,3,4] # ones(length(x))
model = uniformDrawTest(x,pcts)

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

display(chns)
show_var_dist_pct(chns, :ix)
show_var_dist_pct(chns, :val1)
show_var_dist_pct(chns, :val2)
show_var_dist_pct(chns, :val3)


# run: ~/.julia/bin/mpiexecjl -n 4 julia --project diffusion_1D_mpi.jl
using Plots, Printf, MAT
import MPI

# enable plotting by default
if !@isdefined do_save; do_save = true end

# MPI functions
@views function update_halo(A, neighbors_x, comm)
    if neighbors_x[1] >= 0 # MPI_PROC_NULL?
        sendbuf = A[2]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[1], 0, comm)
        MPI.Recv!(recvbuf, neighbors_x[1], 1, comm)
        A[1] = recvbuf[1]
    end
    if neighbors_x[2] >= 0 # MPI_PROC_NULL?
        sendbuf = A[end-1]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[2], 1, comm)
        MPI.Recv!(recvbuf, neighbors_x[2], 0, comm)
        A[end] = recvbuf[1]
    end
    return
end

@views function diffusion_1D_mpi(; do_save=false)
    # MPI
    MPI.Init()
    dims        = [0]
    comm        = MPI.COMM_WORLD
    nprocs      = MPI.Comm_size(comm)
    MPI.Dims_create!(nprocs, dims)
    comm_cart   = MPI.Cart_create(comm, dims, [0], 1)
    me          = MPI.Comm_rank(comm_cart)
    coords      = MPI.Cart_coords(comm_cart)
    neighbors_x = MPI.Cart_shift(comm_cart, 0, 1)
    if (me==0) println("nprocs=$(nprocs), dims[1]=$(dims[1])") end
    # Physics
    lx    = 10.0
    λ     = 1.0
    nt    = 100
    # Numerics
    nx    = 32                 # local
    nx_g  = dims[1]*(nx-2) + 2 # global
    # Derived numerics
    dx    = lx/nx_g            # global
    dt    = dx^2/λ/2.1
    # Array allocation
    qHx   = zeros(nx-1)
    # Initial condition
    x0    = coords[1]*(nx-2)*dx
    xc    = [x0 + ix*dx - dx/2 - 0.5*lx  for ix=1:nx]
    H     = exp.(.-xc.^2)
    t_tic = 0.0
    # Time loop
    for it = 1:nt
        if (it==11) t_tic = Base.time() end
        qHx        .= -λ*diff(H)/dx
        H[2:end-1] .= H[2:end-1] - dt*diff(qHx)/dx
        update_halo(H, neighbors_x, comm_cart)
    end
    t_toc = Base.time()-t_tic
    if (me==0) @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", t_toc, round((2/1e9*nx*sizeof(lx))/(t_toc/(nt-10)), sigdigits=2)) end
    if do_save file = matopen("$(@__DIR__)/H_$(me).mat", "w"); write(file, "H", Array(H)); close(file) end
    MPI.Finalize()
    return
end

diffusion_1D_mpi(; do_save=do_save)


"""
    flts(A::AbstractArray{<:Real}, y::Vector{<:Real})

Fast least trimmed squares by an algorithm from Rousseeuw et al. - Computing LTS Regression for Large Data Sets (1999)
The basic idea is to find an outlier-free subset `H` of regressors that minimizes the sum of squared residuals `Q`, resulting in a robust version of least squares up to 50% outliers

# Args

* A::AbstractArray{<:Real}: Vector or Matrix ∈ R(n x p) containing regressors
* y::Vector{<:Real}: Output Vector ∈ R(n)

# Keyword arguments

* h::Int: The size of the subset H, has to be between 0.5(n + p + 1)) <= h <= n and defaults to 0.5(n + p + 1)
* outliers::Real: Amount of outliers, between 0 <= outliers <= 0.5. A valid definition of h overwrites an outlier definition!
* N::Int: Number of initial p-subsets, defaults to 500 (minimum is 10). This is the main parameter regarding speed, to estimate a reasonable N, the probability of having at least one outlier free p-subset is given by: 1-(1-(1-outliers)^p)^N > 0. If verbose = true information about this gets printed.
* maxiter::Int: Maximum number of C-steps applied in the final optimization, defaults to 100
* ΔQmin::Real: Minimum change of the objective function (squared residuals) that defines convergence
* return_set::Bool: If true, the subset H and the regarding value Q are also returned in the form of (H, θ, Q)
* verbose::Bool: Print some more information

# Return
The parameters θ that minimize Q and optionally a tuple (H, θ, Q), which also contains the selected subset H and the minimal Q

# Example
Taken from the original paper
```jldoctest
julia> x = rand(Normal(0,10), 1000)
julia> y = x .+ 1 + randn(1000)
julia> y[801:end] = rand(Normal(0,5), 200)
julia> x[801:end] = rand(Normal(50,5), 200)
julia> A = [x ones(1000)]
julia> flts(xb,y, return_set = true)

([656, 336, 293, 763, 281, 248, 269, 231, 678, 372  …  734, 220, 434, 658, 174, 465, 745, 371, 759, 732], [1.0003968687209952, 1.0058265326498548], 750.0607838759548)
```
"""
function flts(A::AbstractArray{<:Real}, y::AbstractVector{<:Real};
    h::Integer = 0, outliers::Real = -1, N::Integer = 500, maxiter::Integer = 100, ΔQmin::Real = 0.0001, return_set::Bool = false, verbose::Bool = false)
    # Input checking
    n = length(y)
    size(A, 1) == n || throw(DimensionMismatch("Both inputs A and y should have the same number of rows"))
    N >= 10 || throw(DomainError("N needs to be >= 10"))
    if ndims(A) == 2
        p = size(A, 2)
    else
        # make sure matrix multiplication works
        p = 1
        A = reshape(A, :, 1)
    end

    # Get the subset length h
    if (round(0.5(n + p + 1)) <= h <= n)
        verbose && @info "h was set to: h = $h. Breakdown point is at ~$(round(100 - h / n * 100))% outliers"
    elseif (0.0 <= outliers <= 0.5)
        h = Int(round((1-outliers)*n))
        verbose && @info "h was set to: h = $h. Breakdown point is at ~$(round(100 - h / n * 100))% outliers"
    else
        h = Int(round(0.5(n + p + 1)))
        verbose && @info "h was set to default: h = $h. Breakdown point is at ~$(round(100 - h / n * 100))% outliers"
    end
    # Information about possible success
    verbose && @info "Chance to find an outlier-free p subset is at least ~$((1-(1-(h/n)^p)^N) * 100) %"

    # Initial subset H1 using p-subsets (Option b from paper)
    hs = fill(h, N)
    f(h) = get_initial_H(A, y, p, n, h)
    # Array of Tuples (H, θ, Q)
    initials = map(f, hs)

    # Carry out two C-steps, take top 10 candidates for further optimization
    g1(x) = optimize_H(A, y, h, x, 2, 0)
    opts = map(g1, initials)
    sort!(opts, by = x -> x[3])
    candidates = opts[1:10]

    # Optimze until convergence and return the best
    g2(x) = optimize_H(A, y, h, x, maxiter, ΔQmin)
    results = map(g2, candidates)
    sort!(results, by = x -> x[3])

    # Return the winner
    winner = results[1]
    if return_set
        return winner
    else
        return winner[2]
    end
end

# Helper function that applies a certain number of C-steps to a subset H and checks for convergence
function optimize_H(A, y, h::Int, initial, maxiter::Int, ΔQmin)
    Q_old = initial[3]
    # Create object in outer scope
    opt = nothing
    # Apply C-steps until convergence / maxiter
    for i in 1:maxiter
        opt = C_step(A, y, initial[2], h)
        ΔQ = Q_old - opt[3]
        if ΔQ < ΔQmin
            break
        end
        Q_old = opt[3]
    end
    return opt
end

# Helper function that gets an initial set H from a random p subset
function get_initial_H(A, y, p, n, h)
    inds = 1:n
    # sample a random p subset
    J = sample(inds, p, replace = false)
    # if subset does not specify a unique hyperplane add random observation
    i = 1
    while ((p+i+1) < n) && (rank(A[J, :]) < p)
        J = sample(inds, p+i, replace = false)
        i += 1
    end

    # Estimate first parameters
    θ_J = \(A[J, :], y[J])
    # Use a C-step for an initial H, θ, Q
    H_initial, θ_initial, Q_initial = C_step(A, y, θ_J, h)
    return H_initial, θ_initial, Q_initial
end

# The C-step, workinghorse that guarentees converging parameters
function C_step(A, y, θ_old, h)
    # Use the old parameters to get the residuals
    residuals = y .- A * θ_old
    # Sort the residuals and return a new subset with the smallest residuals
    # H_new = sortperm(abs.(residuals))[1:h] # should be more efficient, but apperantly is not
    H_new = sortperm(abs.(residuals))[1:h]
    # Get new parameters using oridinary least squares
    θ_new = \(A[H_new, :], y[H_new])
    Q_new = get_Q(A, y, H_new, θ_new)
    return H_new, θ_new, Q_new
end

# get the objective function Q
function get_Q(A, y, H, θ)
    residuals = y .- A * θ
    return sum(abs2, residuals[H])
end


function div_thread(N :: Int64, nt :: Int64)

	npt = zeros(Int64, nt, 2)

	n = div(N, nt)

	for i in 1:nt
		npt[i, 1] = (i-1)*n + 1
		npt[i, 2] = npt[i, 1]-1+n
	end

	npt[nt, 2] += mod(N, nt)

	return npt
end


@model binormal() = begin
  y ~ MvNormal(zeros(2), [1.0 0.1; 0.1 1.0])
end


function trainfunction2d(func,x,y)
    if func == 1
        a1 = 3 #rand(Uniform(-10,10))
        b1 = -4#rand(Uniform(-10,10))
        u = a1*x + b1*y
    elseif func == 2
        a = [0.9, 0.1, -0.4]#rand(Uniform(-1,1),3)
        b = [0.2, -0.6, -0.3]#rand(Uniform(-1,1),3)
        u = a[2]*sin(1*pi*y)+cos(1*pi*x)
    elseif func == 3
        a = [0.9464328087730216; 0.17162310348664844; 0.07857856828523646]    #rand(Uniform(-1,1),3)
        b = [0.9355991072384002, 0.5353806651162376, -0.0613925812939633]
        u = a[3]*sin(2*pi*y)+b[3]*cos(2*pi*x)
    elseif func == 4
        a = [-0.8752648836596824, -0.2937424545901659, 0.5352037919227479]#rand(Uniform(-1,1),3)
        b = [-0.9137179533411737, -0.4640307310161864, -0.8663071119500332]#rand(Uniform(-1,1),3)
        u1 = a[3]*sin(pi*x)+b[3]*cos(pi*y)
        u = u1 + a[1]*sin(2*pi*x)+b[1]*cos(2*pi*y)
    elseif func == 5
        a = [-0.6867267253726719, 0.21059347965868014, -0.728510882976293]#rand(Uniform(-1,1),3)
        b = [0.676235507814718, 0.8294240477938528, -0.3998500839840293]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[2]*cos(pi*y)
        u = u1 + a[3]*sin(2*pi*x)+b[3]*cos(2*pi*y)
    elseif func == 6
        a = [0.7983982959430316, 0.903381400725598, 0.6022370326216002]#rand(Uniform(-1,1),3)
        b = [-0.7513545425595294, -0.7714624763400932, -0.8408910576188569]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 7
        a = [0.4456995188427575, -0.760694865555275, 0.5341392645364422]#rand(Uniform(-1,1),3)
        b = [0.6038471708245794, -0.9293109017054246, -0.03067895573440227]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 8
        a1 = 8 #rand(Uniform(-10,10))
        b1 = -3 #rand(Uniform(-10,10))
        u = a1*x + b1*y
    elseif func == 9
        a =[0.2653910349699058, 0.4700076088116263, 0.09668422924668407, -0.046234710248600486, 0.20474920422233067, 0.5826924980829493]
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 10
        a =[-0.9017561219736625, -0.03447706434100306, 0.3305030118808854, 0.684511518915635, -0.661785410962842, 0.877713860698909]
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 11
        a =[-0.4, 1.0, 0.2, 0.6, -0.9, 0.8]   # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 13
        a = [0.8977676302413324, 0.5503759548450637, -0.6451820733536873]#rand(Uniform(-1,1),3)
        b = [0.8, -0.6, -0.3]#rand(Uniform(-1,1),3)
        u = a[2]*sin(1*pi*x)+cos(1*pi*y)
    elseif func == 12
        a =[-0.21587992491755914, -0.041287018187041724, -0.8635531658151585, -0.040103752643044555, -0.0632059673237162, -0.09508661768127391]
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 14
        a = [-0.9, 0.1, -0.4]#rand(Uniform(-1,1),3)
        b = [0.8, -0.6, -0.3]#rand(Uniform(-1,1),3)
        u = a[2]*sin(pi*x)+b[2]*cos(pi*y)
    elseif func == 15
        a1 = -1 #rand(Uniform(-10,10))
        b1 = 3 #rand(Uniform(-10,10))
        u = a1*x + b1*y
    elseif func == 16
        a = [0.2, 0.9, -0.6]#rand(Uniform(-1,1),3)
        b = [-0.1, 0.7, 0.1]#rand(Uniform(-1,1),3)
        u = a[1]*sin(2*pi*x)+b[1]*cos(2*pi*y)
    elseif func == 17
        a = [-0.6, 0.2, -0.9]#rand(Uniform(-1,1),3)
        b = [-1.0, -0.8, 0.3]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 18
        a =[-0.4, 1.0, 0.2, 0.6, -0.9, 0.8]   # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 19 #10
        a = [0.2, 0.1, -0.4]#rand(Uniform(-1,1),3)
        b = [0.1, -0.6, -0.3]#rand(Uniform(-1,1),3)
        u = a[2]*sin(pi*x)+b[2]cos(pi*y)
    elseif func == 20
        a = [-0.6, 0.2, -0.9]#rand(Uniform(-1,1),3)
        b = [-1.0, -0.8, 0.3]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 21
        a = [0.3, 0.4, -1.0]#rand(Uniform(-1,1),3)
        b = [1.0, -0.4, 0.6]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 22
        a = [0.5, -0.9, -0.2]#rand(Uniform(-1,1),3)
        b = [0.1, 0.2, 0.8]#rand(Uniform(-1,1),3)
        u = a[1]*sin(2*pi*x)+b[1]*cos(2*pi*y)
    elseif func == 23
        a = [-0.2, -0.4, 0.9]#rand(Uniform(-1,1),3)
        b = [0.8, 0.6, -0.1]#rand(Uniform(-1,1),3)
        u = a[1]*sin(pi*x)+b[1]cos(pi*y)
    elseif func == 24
        a = [-0.2, -0.4, 0.9]#rand(Uniform(-1,1),3)
        b = [0.8, 0.6, -0.1]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 25
        a = [-0.3, -0.2, 1.0]#rand(Uniform(-1,1),3)
        b = [1.0, -0.2, -0.6]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 26
        a = [-0.1, -0.4, 1.0]#rand(Uniform(-1,1),3)
        b = [-1.0, 0.1, 0.8]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
    elseif func == 27
        a = [0.8, 0.9, -0.2]#rand(Uniform(-1,1),3)
        b = [0.1, 0.2, -0.1]#rand(Uniform(-1,1),3)
        u = a[1]*sin(2*pi*x)+b[1]*cos(2*pi*y)
    elseif func == 28
        a = [0.2, -0.4, 0.9]#rand(Uniform(-1,1),3)
        b = [-0.8, 0.6, -0.1]#rand(Uniform(-1,1),3)
        u = a[1]*sin(pi*x)+b[1]cos(pi*y)
    elseif func == 29  #20
        a = [0.9, 0.4, -0.7]#rand(Uniform(-1,1),3)
        b = [0.9, -0.5, -0.3]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 30
        a = [0.7, 0.0, -1.0]#rand(Uniform(-1,1),3)
        b = [0.0, -0.2, -0.6]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 31
        a = [0.1, 1.0, 1.0]#rand(Uniform(-1,1),3)
        b = [-1.0, -0.3, 0.8]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
    elseif func == 32
        a = [0.1, 0.9, -0.2]#rand(Uniform(-1,1),3)
        b = [0.2, 0.2, -0.1]#rand(Uniform(-1,1),3)
        u = a[1]*sin(2*pi*x)+b[1]*cos(2*pi*y)
    elseif func == 33
        a = [-0.3, -0.4, 0.9]#rand(Uniform(-1,1),3)
        b = [0.8, 0.6, -0.1]#rand(Uniform(-1,1),3)
        u = a[1]*sin(pi*x)+b[1]cos(pi*y)
    elseif func == 34
        a = [0.9, 0.1, -0.8]#rand(Uniform(-1,1),3)
        b = [0.9, 0.3, -0.3]#rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]*cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]*cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]*cos(3*pi*y)
    elseif func == 35
        a =[1.0, 0.3, -0.7, -0.2, 0.1, -0.9]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 36
        a =[-0.6, 1.0, 0.6, -0.8, -0.1, 0.7]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 37
        a =[0.9, -0.1, 0.1, -0.8, 0.4, -0.9]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 38
        a =[-1.0, 0.5, 0.2, -0.8, 0.0, -0.9]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 39
        a =[-0.4, 1.0, 0.2, 0.6, -0.9, 0.8]   # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 40
        a =[0.9, -1.0, 0.4, -0.2, 0.6, -0.2]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 41
        a =[0.4, -0.6, 1.0, 0.5, -0.5, 0.2]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 42
        a =[-0.2, 0.6, 0.3, 0.9, -0.7, 0.9]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 43
        a =[-1.0, -0.1, 0.1, 1.0, 0.1, 0.9]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 44
        a =[0.4, -0.2, 0.8, 0.6, -0.9, 0.8]   # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 45
        a =[1.0, 0.5, 0.9, -1.0, -0.1, 0.3]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 46
        a =[-0.4, -0.4, 0.9, 0.6, -0.9, -0.8]   # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    end
    return u
end

function troubledcellfunctionabs2d(x, y, a, m, x0, y0)
    u =  a*abs((y-y0)-m*(x-x0))
    return u
end

function troubledcellfunctionstep2d(x, y, ui, m, x0, y0)
    h1 = y0+m*(x-x0)
    h2 = y0-(1/m)*(x-x0)

    if y <= h1 && y < h2
        u=ui[1]
    elseif y >= h1 && y > h2
        u=ui[2]
    elseif y < h1 && y >= h2
        u=ui[3]
    elseif y > h1 && y <= h2
        u=ui[4]
    elseif y == h1 && y == h2
        u=ui[1]
    end
    return u
end

function roundstep2d(x, y, ui, r0, x0, y0)
    inicenter = SVector(x0, y0)
    x_norm = x - inicenter[1]
    y_norm = y - inicenter[2]
    r = sqrt(x_norm^2 + y_norm^2)

    # Calculate primitive variables
    u = r > r0 ? ui[1] : ui[2]
    return u
end

function good_cell2d(node_coord, length, func, m, x0, y0)
    x1 = node_coord[1]
    y1 = node_coord[2]
    x2 = node_coord[1] + length
    y2 = node_coord[2] + length

    if func == 1
        y1_abs = y0 + m*(x1-x0)
        y2_abs = y0 + m*(x2-x0)
        x1_abs = x0 + (1/m)*(y1-y0)
        x2_abs = x0 + (1/m)*(y2-y0)
        if  y1_abs <=y2 && y1_abs >= y1
            false
        elseif y2_abs <=y2 && y2_abs >= y1
            false
        elseif x1_abs <= x2 && x1_abs >= x1
            false
        elseif x2_abs <= x2 && x2_abs >= x1
            false
        else
            true
        end
    elseif func == 2 || func == 3
        y1_step = y0 + m*(x1-x0)
        y2_step = y0 + m*(x2-x0)
        x1_step = x0 + (1/m)*(y1-y0)
        x2_step = x0 + (1/m)*(y2-y0)
        y1_step2 = y0 -(1/m)*(x1-x0)
        y2_step2 = y0 -(1/m)*(x2-x0)
        x1_step2 = x0 - m*(y1-y0)
        x2_step2 = x0 - m*(y2-y0)
        if  (y1_step <=y2 && y1_step >= y1) ||  (y1_step2 <=y2 && y1_step2 >= y1 )
            false
        elseif (y2_step <=y2 && y2_step >= y1 ) || (y2_step2 <=y2 && y2_step2 >= y1  )
            false
        elseif (x1_step <= x2 && x1_step >= x1 ) ||  (x1_step2 <= x2 && x1_step2 >= x1)
            false
        elseif (x2_step <= x2 && x2_step >= x1)  || (x2_step2 <= x2 && x2_step2 >= x1)
            false
        else
            true
        end
    elseif func == 4
        inicenter = SVector(x0, y0)
        x_norm1 = x1 - inicenter[1]
        x_norm2 = x2 - inicenter[1]
        y_norm1 = y1 - inicenter[2]
        y_norm2 = y2 - inicenter[2]

        r1 = sqrt(x_norm1^2 + y_norm1^2)
        r2 = sqrt(x_norm2^2 + y_norm1^2)
        r3 = sqrt(x_norm1^2 + y_norm2^2)
        r4 = sqrt(x_norm2^2 + y_norm2^2)
        if (r1 <= m && r2 >=m) || (r1>= m && r2 <=m)
            false
        elseif (r1 <= m && r3 >=m) || (r1>= m && r3 <=m)
            false
        elseif (r2 <= m && r4 >=m) || (r2>= m && r4 <=m)
            false
        elseif (r3 <= m && r4 >=m) || (r3>= m && r4 <=m)
            false
        else
            true
        end
    end
end

function validfunction2d(func,x,y)
    if func == 1
        a = [-0.2, 0.8, -0.7    ]#rand(Uniform(-1,1),3)
        b = [-1.0, 0.4, -0.1]   #rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]cos(3*pi*y)
    elseif func == 2
        a = [0.4, -0.8, 0.1]    #rand(Uniform(-1,1),3)
        b = [0.1, -0.6, 0.9]    #rand(Uniform(-1,1),3)
        u1 = a[1]*sin(pi*x)+b[1]cos(pi*y)
        u2 = u1 + a[2]*sin(2*pi*x)+b[2]cos(2*pi*y)
        u = u2 + a[3]*sin(3*pi*x)+b[3]cos(3*pi*y)
    elseif func == 3
        a = [0.4]
        b = [0.1]
        u = a[1]*sin(pi*x)+b[1]cos(pi*y)
    elseif func == 4
        a = [-0.4]
        b = [0.9]
        u = a[1]*sin(pi*x)+b[1]sin(pi*y)
    elseif func == 5
        a = [0.4]
        b = [-0.7]
        u = a[1]*sin(2*pi*x)+b[1]sin(2*pi*y)
    elseif func == 6
        a =[0.3, -0.1]# rand(Uniform(-1,1),6)
        u = a[1]*x + a[2]*x^2
    elseif func == 7
        a =[-0.1, 0.6, 0.8]# rand(Uniform(-1,1),6)
        u = a[1]*x + a[2]*x^2 + a[3]*x^3
    elseif func == 8
        #a =[-0.5, -0.1, -0.8, 0.8]# rand(Uniform(-1,1),6)
        #u = a[1]*x + a[2]*x^2 + a[3]*x^3 + a[4]*x^4
        a =[1.0, 0.5, 0.9, -1.0, -0.1, 0.3]    # rand(Uniform(-1,1),6)
        u = exp(a[1]*((x-a[2])^2+(y-a[3])^2)) + exp(a[4]*((x-a[5])^2+(y-a[6])^2))
    elseif func == 9
        a =[0.8, -0.4, 0.1]# rand(Uniform(-1,1),6)
        u = a[1]*x + a[2]*x^2 + a[3]*x^3
    end
    return u
end


#!/usr/bin/env julia
# Author: Joseph Sahyoun
# This node uses MPC algorithms to calculate the optimal inputs that minimizes
# the distance between the BARC and the reference trajectory. For this demonstration
# a dummy reference trajectory is fed into this node which then publishes the optimal
# state trajectory.

println("Initializing Packages...")

using RobotOS
@rosimport std_msgs.msg: Float64MultiArray, Float64
@rosimport barc.msg: Input, barc_state
@rosimport data_service.msg: TimeData
@rosimport geometry_msgs.msg: Vector3
rostypegen()
using barc.msg
using data_service.msg
using geometry_msgs.msg
using std_msgs.msg
using JuMP
using Ipopt

println("DONE")

# -------------------------------------------------------------------
#                           MPC PARAMETERS
# -------------------------------------------------------------------

println("Defining Parameters...")

# Number of states / Number of inputs
nz = get_param("/numStates")
nu = get_param("/numInputs")

#### MAKE MPC TERMS PARAMETERS
# Horizons
n = get_param("/numStepsToLookAhead")           # MPC Horizon (CFTOC Horizon)

# Cost matrices
Q = get_param("/Q")
Q = reshape(Q,(nz,nz))
R = get_param("/R")
R = reshape(R,(nu,nu))

# Hardware parameters
lr = get_param("/lr")
lf = get_param("/lf")

#Box constraints on states
x_min = get_param("/x_min")
x_max = get_param("/x_max")
y_min = get_param("/y_min")
y_max = get_param("/y_max")

# Box constraints on input, A*u <= B
v_max = get_param("/v_max")
v_min = get_param("/v_min")
steer_min = get_param("/steer_min")
steer_max = get_param("/steer_max")

vref = get_param("/reference_velocity")
numMovingAveragePoints = get_param("/numMovingAveragePoints")
velSum = []
deltaSum = []
count = 0
version = 1

println("DONE")

# -------------------------------------------------------------------
#                         SET UP CFTOC MODEL
# -------------------------------------------------------------------

println("Creating CFTOC Model...")

m = Model(solver = IpoptSolver(print_level=0))

# Variables
@variable(m,z[1:nz,1:n+1])              # State vectors (nz x n+1)
@variable(m,u[1:nu,1:n])                # Input vectors (nu x n)

# Reference Trajectory
@NLparameter(m, xref[i=1:n] == 0)
@NLparameter(m, yref[i=1:n] == 0)
@NLparameter(m, pref[i=1:n] == 0)

# Objective Function

if version == 1
    @NLobjective(m, :Min, sum(Q[1,1]*(z[1,i+1]-xref[i])^2 + Q[2,2]*(z[2,i+1]-yref[i])^2 for i in 1:n) + sum(R[1,1]*(u[1,i]-vref)^2 for i in 1:n))
end

# State Constraints

for i in 1:n+1
    @constraint(m, x_min <= z[1,i] <= x_max)
    @constraint(m, y_min <= z[2,i] <= y_max)
    @constraint(m, -3.1415926/2 <= z[3,i] <= 3.1415926/2)
end

for i in 1:n
    @constraint(m, v_min <= u[1,i] <= v_max)
    @constraint(m, steer_min <= u[2,i] <= steer_max)
end

# Initial Condition and Constraints
@NLparameter(m, x0 == 0); @NLconstraint(m, z[1,1] == x0);
@NLparameter(m, y0 == 0); @NLconstraint(m, z[2,1] == y0);
@NLparameter(m, p0 == 0); @NLconstraint(m, z[3,1] == p0);
@NLparameter(m, dt == 0.05);

# Dynamics Constraints
for i in 1:n
    @NLconstraint(m, z[1,i+1] == z[1,i] + dt*u[1,i]*cos(z[3,i]+atan(lr*tan(u[2,i])/(lf+lr))))
    @NLconstraint(m, z[2,i+1] == z[2,i] + dt*u[1,i]*sin(z[3,i]+atan(lr*tan(u[2,i])/(lf+lr))))
    @NLconstraint(m, z[3,i+1] == z[3,i] + dt*(u[1,i]/lr)*sin(atan(lr*tan(u[2,i])/(lf+lr))))
end

println("DONE")

# -------------------------------------------------------------------
#                             INITIAL SOLVE
# -------------------------------------------------------------------

println("Initial solve...")
solve(m)
println("DONE")

# -------------------------------------------------------------------
#                             ROS FUNCTIONS
# -------------------------------------------------------------------

println("Running MPC...")

function dt_callback(msg::Float64Msg)
    setvalue(dt, msg.data)
end

function callback2(msg::barc_state)
    # Assign new reference trajectory
    for i in 1:n
        setvalue(xref[i], msg.x[i])
        setvalue(yref[i], msg.y[i])
        #setvalue(pref[i], msg.psi[i])
    end
end

function loop(pub_one, pub_two)
    global count,velSum,deltaSum
    while ! is_shutdown()
	loop_rate = Rate(get_param("/loop_rate"))
	count = count+1
        # -----------------------------------------------------------
        #                      RUNNING MPC
        # -----------------------------------------------------------
		if count>1
	    toc()
		end
        tic()
        solve(m)                                # Solve CFTOC
        zOpt = getvalue(z)
        uOpt = getvalue(u)
        JOpt = getobjectivevalue(m)
		if count<=numMovingAveragePoints
			velSum = [velSum;uOpt[1,1]]
			deltaSum = [deltaSum;uOpt[2,1]]
		else
			velSum = [velSum[2:numMovingAveragePoints];uOpt[1,1]]
			deltaSum = [deltaSum[2:numMovingAveragePoints];uOpt[2,1]]
		end
        # Publish optimal inputs and optimal state trajectory
        # Only need to publish state traj for this demonstration

        cmdInt = Input(mean(velSum), mean(deltaSum))    # Command to publish
        optTraj = barc_state(zOpt[1,1:end],zOpt[2,1:end],zOpt[3,1:end])
        publish(pub_one,cmdInt)
        publish(pub_two,optTraj)

        rossleep(loop_rate)

    end
end

function main()
    init_node("MPC_node")
	global count,velSum,deltaSum
    pub = Publisher("uOpt", Input, queue_size=10)
    pub2 = Publisher("optimal_state_trajectory", barc_state, queue_size=1)
    sub1 = Subscriber("dt", Float64Msg, dt_callback, queue_size=1)
    sub2 = Subscriber("reference_trajectory", barc_state, callback2, queue_size=1)
    loop(pub, pub2)
end

if ! isinteractive()
    main()
end


__precompile__()

module IScatterSpectrum

using SpecialFunctions,Polynomials

import PhysicalConstants.MKS: Boltzmann,ElectronCharge,MassElectron,SpeedOfLight,
                              UnifiedAtomicMass,VacuumPermittivity

friedconte(z) = im*sqrt(pi)*erfcx(-im*z)

struct Faddeeva
    n::Int
    l::Float64
    a::Poly{Float64}
    function Faddeeva(n=16)
        m  = 2*n
        l  = sqrt(n/sqrt(2))

        theta = (-m+1:m-1)*pi/m
        t = l*tan.(theta/2)
        f = [0; exp.(-t.^2).*(l*l + t.^2)]
        a = real.(fft(fftshift(f)))[2:n-1]/(2*m)   # no flip because of convention in polyval
        new(n, l, Poly(a))
    end
end

const deff = Faddeeva()

"""
    faddeeva(z, f=deff)

   The Faddeeva function is defined as:

   w(z) = exp(-z^2) * erfc(-im*z)

   where erfc(x) is the complex complementary error function. Also,

   friedconte(z) = im * sqrt(pi) * faddeeva(z)

   where friedconte(z) is the plasma dispersion function as defined in [2].

   w = faddeeva(z, f) reuses previously calculated coefficients in struct f
   where f is initialized like f=Faddeeva(n) with n specifying the number of terms
   to truncate the expansion (see (13) in [1]). n = 16 is used as default, i.e. for f=Faddeeva().

   Example:
       x = linspace(-10,10,1001); [X,Y] = meshgrid(x,x);
       W = faddeeva(complex(X,Y));

   Reference:
   [1] J.A.C. Weideman, "Computation of the Complex Error Function," SIAM
       J. Numerical Analysis, pp. 1497-1518, No. 5, Vol. 31, Oct., 1994
       Available Online: http://www.jstor.org/stable/2158232

   [2] B.D. Fried, S.D. Conte, The plasma dispersion function. New York Acadenic Press, 1961.
"""
function faddeeva(z, f::Faddeeva=deff)
    if real(z)==0      # purely imaginary input
        exp.(-z.^2).*erfc(imag(z))
    else               # complex valued input
        # make sure input in the upper half-plane (positive imag. values)
        zz = imag(z)<0? conj(z):z
        yy = (f.l + im*zz)/(f.l - im*zz)
        w  = 2*polyval(f.a, yy)/(f.l - im*zz)^2 + (1/sqrt(pi))/(f.l - im*zz)
        # convert the upper half-plane results to the lower half-plane if necessary
        imag(z)<0? conj(2*exp(-zz^2) - w):w
    end
end

const fnormf = SpeedOfLight*sqrt(MassElectron/(2*Boltzmann))/2
const debyef = (4*π/ElectronCharge/SpeedOfLight)^2*VacuumPermittivity*Boltzmann

"gfe(b in T) -> electron gyrofrequency in Hz"
gfe(b) = ElectronCharge*b/MassElectron

"""
    ScatterVolume(tfeff, system, b, alpha)

has parameters of the scattering volume that are constant in a fit.
"""
struct ScatterVolume{T<:Real}
    tfeff::T    # effective transmitter frequency in Hz
    fnormc::T   # factor for frequency normalization f -> theta
    debyec::T   # Debye length correction factor

    system::T   # spectrum scale factor (depends on tx power, volume size, ... )

    b::T        # magnetic field strength in Tesla
    gfe::T      # electron gyrofrequency
    phiel::T    # normalized electron gyrofrequency
    alpha::T    # angle between b and scattering vector in rad
    sinfac::T   # correction factor for magnetic field
    cosal::T    #
    cosal2::T

    function ScatterVolume{T}(tfeff::T, system::T, b::T, alpha::T) where T<:Real
        fnormc = fnormf/tfeff
        debyec = debyef*tfeff^2
        gf     = gfe(b)
        phiel  = fnormc*gf
        sinfac = 0.5*(sin(alpha)/phiel)^2
        cosal  = cos(alpha)
        new(tfeff, fnormc, debyec, system, b, gf, phiel, alpha, sinfac, cosal, cosal*cosal)
    end
end
ScatterVolume(tfeff::T, system::T, b::T, alpha::T) where {T<:Real} =
    ScatterVolume{T}(tfeff, system, b, alpha)

struct Ion{T<:Real}
    mass::T   # ion mass in amu
    srmr::T   # sqrt of mass ratio ion/electron
end
Ion(m=15.999) = Ion(m, sqrt(m*UnifiedAtomicMass/MassElectron))

"""
Plasma(ne, te, ti, ::ScatterVolume, [ions, fractn]

has parameters of the plasma that are constant over all frequencies of the spectrum.
"""
struct Plasma{T<:Real}
    ne::T
    te::T
    ti::T
    fractn1::T
    fractn::NTuple{2, T}    # up to two more ions
    ions::NTuple{3, Ion{T}}
    nion::Int

    spectrsf::T
    debyen::T   # normalized Debye length correction factor
    sqte::T
    sqti::T
    thetc::T
    tex::T
    tefac::T
    dpsidc::T

    # parameter for spectrum at zero frequency:
    cdtine::T
    ym::T

    function Plasma{T}(ne::T, te::T, ti::T, s::ScatterVolume{T}, fractn, ions, nion) where T<:Real
        sqti    = sqrt(ti)
        thetc   = s.fnormc/sqti
        tex     = exp(-te*s.sinfac)
        tefac   = 0.5/te + s.sinfac/tex
        fractn1 = one(ne) - fractn[1] - fractn[2]
        dpsidc  = thetc*ions[1].srmr
        cdtine  = ti*s.debyec + ne
        ym      = ti*ne + te*cdtine
        new(ne, te, ti, fractn1, fractn, ions, nion,
            s.system*ne, s.debyec/ne, sqrt(te), sqti, thetc,
            tex, tefac, dpsidc, cdtine, ym)
    end
end
Plasma(ne::T, te::T, ti::T,
       s::ScatterVolume{T}, nion=1, ions=(Ion(16*one(ne)), Ion(31*one(ne)), Ion(one(ne))),
       fractn=(zero(ne),zero(ne))) where {T<:Real} =
           Plasma{T}(ne, te, ti, s, fractn, ions, nion)

struct CollisionalPlasma{T<:Real}
    p::Plasma{T}
    cf::T
end

"""
    ionadm(theta, z=friedconte(theta))
-> unmagnetized ion admittance for real normalized frequency
"""
ionadm{T<:Real}(theta::T, z::Complex{T}=friedconte(theta)) = Complex(theta*imag(z), 1 + theta*real(z))

"""
    ionadm(freq, ::Plasma)
-> unmagnetized ion admittance for real unnormalized frequency and plasma parameters
"""
function ionadm{T<:Real}(freq::T, p::Plasma{T}, s::ScatterVolume{T})
    fnormt = s.fnormc/p.sqti
    thetai = fnormt*p.ions[1].srmr*freq
    yi     = ionadm(thetai)
    for k in 2:p.nion
        if k==2
            yi *= p.fractn1
        end
        thetai = fnormt*p.ions[k].srmr*freq
        yi    += p.fractn[k-1]*ionadm(thetai)
    end
    yi
end

"""
    ionadm(theta, z=friedconte(theta))
-> unmagnetized ion admittance for complex normalized frequency.
"""
function ionadm{T<:Complex}(theta::T, z::T=friedconte(theta))
    psizm = imag(theta)*(real(z)^2 + imag(z)^2)
    yd    = 1 - imag(theta)*(2*imag(z) - psizm)
    Complex(real(theta)*(imag(z) - psizm)/yd, 1 + real(theta)*real(z)/yd)
end

"""
    ionadm(freq, ::CollisionalPlasma)
-> unmagnetized ion admittance for unnormalized frequency and collisional plasma parameters
"""
function ionadm{T<:Real}(freq::T, coll::CollisionalPlasma{T}, s::ScatterVolume{T})
    fnormt = s.fnormc/coll.p.sqti
    thetai = fnormt*coll.p.ions[1].srmr*Complex(freq, coll.cf)
    yi     = ionadm(thetai)
    for k in 2:coll.p.nion
        if k==2
            yi *= coll.p.fractn1
        end
        thetai = fnormt*coll.p.ions[k].srmr*Complex(freq, coll.cf)
        yi    += coll.p.fractn[k-1]*ionadm(thetai)
    end
    yi
end

"""
    elecadm(theta, tex, z)
-> magnetized electron admittance for real normalized frequency
"""
function elecadm{T<:Real}(theta::T, tex::T, z::Complex{T}=friedconte(theta))
    tt = tex*theta
    Complex(tt*imag(z), 1 + tt*real(z)), imag(z)
end

# helper function, this part of the spectrum code is the same for non-collisional and collisional
"""
    espec(freq, ::Plasma, ::ScatterVolume, yi::Complex)
-> power spectral density for non-zero real frequency
"""
function espec{T<:Real}(freq::T, p::Plasma{T}, s::ScatterVolume{T}, yi::Complex{T})
    yiti   = yi/p.ti + im*p.debyen

    # electron admittance:
    thetae = s.fnormc/(p.sqte*s.cosal)*freq
    ye,    = elecadm(thetae, p.tex)
    yete   = ye/p.te

    the2   = thetae^2
    texthe = p.tex*thetae

    ym     = abs2(yete+yiti)
    yn     = abs2(yete)*real(yi) + abs2(yiti)*real(ye)

    (p.spectrsf/freq)*(yn/ym)
end

"""
    pwrspec(freq, ::Plasma, ::ScatterVolume)
-> power spectral density for non-zero real frequency
"""
function pwrspec{T<:Real}(freq::T, p::Plasma{T}, s::ScatterVolume{T})
    yi = ionadm(freq, p, s)
    espec(freq, p, s, yi)
end

"""
    pwrspec(freq, ::CollisionalPlasma, ::ScatterVolume)
-> power spectral density for non-zero frequency and collisional plasma
"""
function pwrspec{T<:Real}(freq::T, coll::CollisionalPlasma{T}, s::ScatterVolume{T})
    yi = ionadm(freq, coll, s)
    espec(freq, coll.p, s, yi)
end

# helper function
function yn0{T<:Real}(p::Plasma{T}, s::ScatterVolume{T}, yir::T)
    tineq  = p.sqti*p.ne^2*yir*p.ti
    sqtec  = p.cdtine*p.sqte*sqrt(pi)/s.cosal2 # zi(1) in fortran code = sqrt(π) for real zero frequency
    stecx  = sqtec*p.tex*p.te
    tineq + stecx*p.cdtine
end

"""
    pwrspec(::Plasma)
-> power spectral density for zero real frequency
"""
function pwrspec(p::Plasma, s::ScatterVolume)
    yir = sqrt(π)*p.ions[1].srmr
    for k in 2:p.nion
        if k==2
            yir *= p.fractn1
        end
        yir += p.fractn[k-1]*sqrt(π)*p.ions[k].srmr
    end
    yn = yn0(p, s, yir)
    s.fnormc*p.spectrsf*yn/p.ym^2
end

function yirc(coll::CollisionalPlasma, kion::Int)
    psiion = coll.p.dpsidc*coll.cf
    z      = friedconte(Complex(zero(coll.p.ne), psiion))
    psiz   = 1 - psiion*imag(z)
    psizm  = coll.p.ions[kion].srmr/psiz
    psizm*imag(z)
end

"""
    pwrspec(::CollisionalPlasma)
-> power spectral density for zero frequency and collisional plasma
"""
function pwrspec(coll::CollisionalPlasma, s::ScatterVolume)
    yir = yirc(coll, 1)
    for k in 2:p.nion
        if k==2
            yir *= p.fractn1
        end
        z    = friedconte(Complex(zero(coll.p.ne), p.thetc*coll.p.ions[k].srmr*coll.cf))
        yir += p.fractn[k-1]*yirc(coll, k)
    end

    yn = yn0(p, s, yir)
    s.fnormc*coll.p.spectrsf*yn/p.ym^2
end

end   # module IScatterSpectrum


### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 1eb559b3-dbce-4caf-800f-96a2495dee1b
using Revise

# ╔═╡ 77d14c3f-bcaf-432c-9005-f5ec184d84b7
using QuadGK

# ╔═╡ ba2dadc8-d9a2-452e-9f94-406ccfc7404b
using Optim

# ╔═╡ c05cd489-68c8-418e-bbc2-a33bc42781f0
using Plots

# ╔═╡ 1d300e7d-5efb-4ad2-a826-c3f8a86083d8
using PolaronMobility

# ╔═╡ 6a99cd13-68cb-474c-96af-7d78ee4e7095
using DataFrames

# ╔═╡ 49663b33-1604-430b-a815-9eb985802987
using CSV

# ╔═╡ 314af942-8922-40d4-a92b-f75b715c276c
using LaTeXStrings

# ╔═╡ 6ff68ba5-df0c-42a0-8cf1-dedeced77e8d
using ColorSchemes

# ╔═╡ 29d62408-0629-4f39-909e-a6ba8696e955
begin
	# Physical constants
	const ħ = 1.05457162825e-34; # Reduced Planck's constant (kg m^2 s^{-1})
	const eV = 1.602176487e-19; # Electron Volt (kg m^2 s^{-2})
	const m_e = 9.10938188e-31; # Electron Mass (kg)
	const k_B = 1.3806504e-23; # Boltzmann's constant (kg m^2 K^{-1} s^2)
	const ϵ_0 = 8.854e-12 # Dielectric constant (C^2 N^{-1} m^{-2})
	const c = 2.99792458e8 # Speed of light (m s^{-1})
	const amu = 1.66053906660e-27 # Atomic Mass Unit (kg)
end

# ╔═╡ 35390bdf-e260-4566-8a97-7a0b5a3ef07a
MAPI= [
# 96.20813558773261 0.4996300522819191
# 93.13630357703363 1.7139631746083817
# 92.87834578121567 0.60108592692181
# 92.4847918585963 0.0058228799414729
# 92.26701437594754 0.100590086574602
# 89.43972834606603 0.006278895133832249
# 46.89209141511332 0.2460894564364346
# 46.420949316788 0.14174282581124137
# 44.0380222871706 0.1987196948553428
# 42.89702947649343 0.011159939465770681
# 42.67180170168193 0.02557751102757614
# 41.46971205834201 0.012555230726601503
# 37.08982543385215 0.00107488277468418
# 36.53555265689563 0.02126940080871224
# 30.20608114002676 0.009019481779712388
# 27.374810898415028 0.03994453721421388
# 26.363055017011728 0.05011922682554448
# 9.522966890022039 0.00075631870522737
4.016471586720514 0.08168931020200264
3.887605410774121 0.006311654262282101
3.5313112232401513 0.05353548710183397
2.755392921480459 0.021303020776321225
2.4380741812443247 0.23162784335484837
2.2490917637719408 0.2622203718355982
2.079632190634424 0.23382298607799906
2.0336707697261187 0.0623239656843172
1.5673011873879714 0.0367465760261409
1.0188379384951798 0.0126328938653956
1.0022960504442775 0.006817361620021601
0.9970130778462072 0.0103757951973341
0.9201781906386209 0.01095811116040592
0.800604081794174 0.0016830270365341532
0.5738689505255512 0.00646428491253749
# 0.022939578929507105 8.355742795827834e-06   # Acoustic modes!
# 0.04882611767873102 8.309858592685e-06
# 0.07575149723846182 2.778248540373041e-05
]

# ╔═╡ e0ce47a5-368a-4da7-83fc-9f921af0f747
function ϵ_ionic_mode(phonon_mode_freq, ir_activity, volume) # single ionic mode
    ω_j = 2π * phonon_mode_freq * 1e12 # angular phonon freq in Hz
    ϵ_mode = eV^2 * ir_activity / (3 * volume * ω_j^2 * amu) # single ionic mode
    return ϵ_mode / ϵ_0 # normalise with 1 / (4π ϵ_0)
end

# ╔═╡ 053d717a-2b63-40e7-9dc0-487e3892de98
function ϵ_total(freqs_and_ir_activity, volume) # total ionic contribution to dielectric
    phonon_freqs = freqs_and_ir_activity[:, 1]
    ir_activity = freqs_and_ir_activity[:, 2]
    result = 0.0
    for (f, r) in zip(phonon_freqs, ir_activity)
        result += ϵ_ionic_mode(f, r, volume) # sum over all ionic contributions
    end
    return result
end

# ╔═╡ 86fe8a15-4d9e-48fa-8387-3cfaff7aa182
function frohlich_α_j(ϵ_optic, ϵ_ionic, ϵ_total, phonon_mode_freq, m_eff) # Frohlich alpha decomposed into phonon branch contributions
    Ry = eV^4 * m_e / (2 * ħ^2) # Rydberg energy
    ω = 2π * 1e12 * phonon_mode_freq # angular phonon freq (Hz)
    ϵ_static = ϵ_total + ϵ_optic # static dielectric. Calculate here instead of input so that ionic modes properly normalised.
    return (m_eff * Ry / (ħ * ω))^(1 / 2) * ϵ_ionic / (4π * ϵ_0) / (ϵ_optic * ϵ_static) # 1 / (4π ϵ_0) dielectric normalisation
end

# ╔═╡ 16e7d781-bc61-44b3-84a2-e8c8892bf773
function κ_i(i, v, w) # fictitious spring constant, multiple variational params
    κ = v[i]^2 - w[i]^2
    if length(v) > 1
        for j in 1:length(v)
            if j != i
                κ *= (v[j]^2 - w[i]^2) / (w[j]^2 - w[i]^2)
            end
        end
    end
    return κ
end

# ╔═╡ a78017d7-8835-4a8f-8dd3-f3a4a50eb37b
function h_i(i, v, w) # some vector relating to harmonic eigenmodes
    h = v[i]^2 - w[i]^2
    if length(v) > 1
        for j in 1:length(v)
            if j != i
                h *= (w[j]^2 - v[i]^2) / (v[j]^2 - v[i]^2)
            end
        end
    end
    return h
end

# ╔═╡ d7346684-6b68-4dfc-81da-12c435c4acf0
function C_ij(i, j, v, w) # generalised Feynman C variational parameter (inclusive of multiple v and w params)
    C = w[i] * κ_i(i, v, w) * h_i(j, v, w) / (4 * (v[j]^2 - w[i]^2))
    return C
end

# ╔═╡ 798a4f7a-a2a0-48de-82ad-f85b4199496a
function D_j(τ, β, v, w) # log of dynamic structure factor for polaron
    D = τ * (1 - τ / β)
    for i in 1:length(v)
        if v[i] != w[i]
        D += (h_i(i, v, w) / v[i]^2) * (2 * sinh(v[i] * τ / 2) * sinh(v[i] * (β - τ) / 2) / (v[i] * sinh(v[i] * β / 2)) - τ * (1 - τ / β))
        end
    end
    return D
end

# ╔═╡ ec0744c0-2530-11ec-04c3-1b5d83413b78
function multi_free_energy(v_params, w_params, T, ϵ_optic, m_eff, volume, freqs_and_ir_activity)

    setprecision(BigFloat, 32) # Speed up. Stops potential overflows.

    # Extract phonon frequencies and ir activities.
    phonon_freqs = freqs_and_ir_activity[:, 1]
    ir_activity = freqs_and_ir_activity[:, 2]

    num_of_branches = length(phonon_freqs)

    # total dielectric contribution from all phonon branches (used as a normalisation)
    ϵ_tot = ϵ_total(freqs_and_ir_activity, volume)

    # Generalisation of B i.e. Equation 62c in Hellwarth.
    B_j_integrand(τ, β, v, w) = cosh(β / 2 - abs(τ)) / (sinh(β / 2) * sqrt(D_j(abs(τ), β, v, w)))
    B_j(β, α, v, w) = α / √π * quadgk(τ -> B_j_integrand(τ, β, v, w), 0.0, β / 2)[1]

    # Generalisation of C i.e. Equation 62e in Hellwarth.
    function C_j(β, v, w)
        s = 0.0
        for i in 1:length(v)
            for j in 1:length(v)
                s += C_ij(i, j, v, w) / (v[j] * w[i]) * (coth(β * v[j] / 2)  - 2 / (β * v[j]))
            end
        end
        3 * s / num_of_branches
    end

    # Generalisation of A i.e. Equation 62b in Hellwarth.
    function A_j(β, v, w)
        s = -log(2π * β) / 2
        for i in 1:length(v)
            if v[i] != w[i]
                s += log(v[i] / w[i]) - log(sinh(v[i] * β / 2) / sinh(w[i] * β / 2))
            end
        end
        3 / β * s / num_of_branches
    end

	F = 0.0
	for j in 1:num_of_branches
		ω_j = 2π * 1e12 * phonon_freqs[j] # angular phonon freq im 2π Hz
		β_j = BigFloat(ħ * ω_j / (k_B * T)) # reduced thermodynamic beta
		ϵ_ionic_j = ϵ_ionic_mode(phonon_freqs[j], ir_activity[j], volume) # ionic dielectric contribution for current phonon branch
		α_j = frohlich_α_j(ϵ_optic, ϵ_ionic_j, ϵ_tot, phonon_freqs[j], m_eff) # decomposed alpha for current phonon branch

		# F = -(A + B + C) in Hellwarth.
		F += -(B_j(β_j, α_j, v_params, w_params) + C_j(β_j, v_params, w_params) + A_j(β_j, v_params, w_params)) * ω_j # × ħω branch phonon energy
	end

    return F * ħ / eV * 1e3 # change to meV
end

# ╔═╡ f3b91e16-26d3-45c8-b28f-193afac15d7f
function multi_variation(T, ϵ_optic, m_eff, volume, freqs_and_ir_activity; initial_vw = false, N = 1) # N number of v and w params

    setprecision(BigFloat, 32) # Speed up. Stops potential overflows.

    if initial_vw isa Bool
		# Intial guess for v and w.
    	initial = sort(rand(2 * N)) .* 4.0 .+ 1.0 # initial guess around 4 and ≥ 1.

		# Limits of the optimisation.
		lower = repeat([0.1], 2 * N)
		upper = repeat([60.0], 2 * N)
	else
		# Intial guess for v and w.
		initial = sort(vcat(initial_vw...))

		# Limits of the optimisation.
		lower = repeat([0.1], 2 * N)
		upper = repeat([60.0], 2 * N)
	end

	println("Initial guess: ", initial)

	# Osaka Free Energy function to minimise.
	f(x) = multi_free_energy([x[2 * n] for n in 1:Int(N)], [x[2 * n - 1] for n in 1:Int(N)], T, ϵ_optic, m_eff, volume, freqs_and_ir_activity)

	# Use Optim to optimise the free energy function w.r.t v and w.
	solution = Optim.optimize(
		Optim.OnceDifferentiable(f, initial; autodiff = :forward),
		lower,
		upper,
		initial,
		Fminbox(BFGS()),
		Optim.Options(time_limit = 20.0),
	)

	# Get v and w params that minimised free energy.
	var_params = Optim.minimizer(solution)

	# Update matrices for v and w parameters.
	v_params = [var_params[2 * n] for n in 1:N]
	w_params = [var_params[2 * n - 1] for n in 1:N]

	# Show current v and w that minimise jth phonon branch.
	println("Variational parameters: ", var_params)

    # Return variational parameters that minimise the free energy.
    return v_params, w_params
end

# ╔═╡ a941ccc6-90cb-448a-91bb-f80233b630f8
function HellwarthAScheme(phonon_modes; T = 295)

	phonon_mode_freqs = phonon_modes[:, 1]
	ir_activities = phonon_modes[:, 2]

	condition(f) = coth(π * f * 1e12 * ħ / (k_B * T)) / f - sum(ir_activities .* coth.(π .* phonon_mode_freqs .* 1e12 .* ħ ./ (k_B * T)) ./ phonon_mode_freqs) / sum(ir_activities)

	minimum_frequency = minimum(phonon_mode_freqs)
	maximum_frequency = maximum(phonon_mode_freqs)
	middle_frequency = (maximum_frequency + minimum_frequency) / 2
	print("\n")

	while (maximum_frequency - minimum_frequency) / 2 > 1e-6

		if sign(condition(middle_frequency)) == sign(condition(minimum_frequency))
			minimum_frequency = middle_frequency
			middle_frequency = (maximum_frequency + minimum_frequency) / 2
		else
			maximum_frequency = middle_frequency
			middle_frequency = (maximum_frequency + minimum_frequency) / 2
		end
	end

	return middle_frequency
end

# ╔═╡ 1372f247-4622-4388-b458-89aa6e695388
function HellwarthBScheme(LO)
    println("Hellwarth B Scheme... (athermal)")
    H58 = sum( LO[:,2] ./ LO[:,1].^2 )
    println("Hellwarth (58) summation: ",H58)

    H59 = sum( LO[:,2] ) # sum of total ir activity squarred
    println("Hellwarth (59) summation (total ir activity ^2): ", H59)
    println("Hellwarth (59) W_e (total ir activity ): ", sqrt(H59))

    omega = sqrt(H59 / H58)
    println("Hellwarth (61) Omega (freq): ",omega)

	return(omega)
end

# ╔═╡ c28f6d35-c3ed-4984-9c36-b282fb50001b
function ridgeline(x, y, z; shift = 1, jump = 10, ylabel = :none, xlabel = :none, palette = :twilight, linewidth = 1.5, color = :black, size = (595, 842), ymirror = false, fillalpha = 1, step = 1)
	if !ymirror
		p = plot(x, z[:, 1], fillrange = z[:, jump] .- shift, fillalpha = fillalpha, ylabel = ylabel, xlabel = xlabel, xticks = (0:maximum(x)/10:maximum(x), string.(x[1:Int(ceil((length(x)-1)/10)):end])), yticks = (z[1, 1:jump * step:end] .- shift .* range(0, stop = length(y[1:jump:end]) - 1, step = step), string.(y[1:jump * step:end])), legend = false, size = size, grid = false, ymirror = ymirror, palette = palette, tickfontsize = 12, labelfontsize = 12)
	else
		p = plot(x, z[:, 1], fillrange = z[:, jump] .- shift, fillalpha = fillalpha, ylabel = ylabel, xlabel = xlabel, xticks = (0:maximum(x)/10:maximum(x), string.(x[1:Int((length(x))/10):end])), yticks = (z[end, 1:jump * step:end] .- shift .* range(0, stop = length(y[1:jump:end]) - 1, step = step), string.(y[1:jump * step:end])), legend = false, size = size, grid = false, ymirror = ymirror, palette = palette, tickfontsize = 12, labelfontsize = 12)
	end

	plot!(x, z[:, 1], color = color, linewidth = linewidth)
	plot!(x, z[:, jump] .- shift, color = color, linewidth = linewidth)

	for i in 1:Int(floor(length(y)/jump) - 1)
	plot!(x, z[:, i * jump] .- shift * i, fillrange = z[:, (i + 1) * jump] .- shift * (i + 1), fillalpha = fillalpha, palette = palette)
	plot!(x, z[:, i * jump] .- shift * i, color = color, linewidth = linewidth)
	plot!(x, z[:, (i + 1) * jump] .- shift * (i + 1), color = color, linewidth = linewidth)
	end
	return p
end

# ╔═╡ 49480256-2dcc-4d10-a5a1-4202b243cfd2
T_range = 100:100:400

# ╔═╡ 4966ecf1-ef39-4da3-a8a5-9e47a6c6b0ec
phonon_mode_freqs = MAPI[:, 1]

# ╔═╡ 3321246b-76bd-4678-a5ae-332c3543a1ef
ir_activities = MAPI[:, 2]

# ╔═╡ b2944c8a-158e-461b-a83c-842925587179
A_scheme = [HellwarthAScheme(MAPI, T = i) for i in T_range]

# ╔═╡ fe4e5ef5-833a-4b61-9ea9-92fa52d7a21f
B_scheme = HellwarthBScheme(hcat(phonon_mode_freqs, ir_activities))

# ╔═╡ 8350c7a2-e9fd-4312-9561-44e2524734af
ϵ_ionic = [ϵ_ionic_mode(f, r, (6.29e-10)^3) for (f, r) in zip(phonon_mode_freqs, ir_activities)]

# ╔═╡ db2b4169-0737-44c1-9af5-aa2f2403f5db
ϵ_tot = sum(ϵ_ionic)

# ╔═╡ 2b6d00e1-317b-41cc-9489-14a1450012f4
α_j = [frohlich_α_j(4.5, ϵ_i, ϵ_tot, f, 0.12) for (ϵ_i, f) in zip(ϵ_ionic, phonon_mode_freqs)]

# ╔═╡ 3e1418c2-01a1-4651-8029-59f53264193d
begin
	scatter(phonon_mode_freqs, α_j, yaxis = :log, xticks = (phonon_mode_freqs, hcat(["$(round(i, digits = 2))" for i in phonon_mode_freqs]...)), label = L"\textbf{\alpha_j}", size = (500, 500), xlabel = "Phonon Frequencies (THz)", tickfontsize = 9, xrotation = 90)
	scatter!(phonon_mode_freqs, ir_activities, markershape = :diamond, label = "IR")
	plot!(phonon_mode_freqs, α_j ./ ir_activities, linewidth = 2, linestyle = :dash, label = "ratio")
end

# ╔═╡ d4f392ef-968a-439c-ba53-5707464446e7
begin
	multi_data = DataFrames.DataFrame(
			alpha = α_j,
			phonon_freqs = phonon_mode_freqs,
			ir_activities = ir_activities,
			ionic = ϵ_ionic
	    )
	CSV.write("multi_data.csv", multi_data)
end

# ╔═╡ b3c94ea5-d152-4705-9a30-f8d3e4da0874
α_eff = sum(α_j)

# ╔═╡ 3608af5e-7b44-4d50-9770-5e8afd4fca28
α_hellwarth_A = PolaronMobility.frohlichalpha.(4.5, 24.1, A_scheme .* 1e12, 0.12)

# ╔═╡ 1716f263-f8ad-40a1-a805-d8b7ea6f38df
α_hellwarth_B = PolaronMobility.frohlichalpha(4.5, 24.1, B_scheme * 1e12, 0.12)

# ╔═╡ da629bec-b53c-4acf-a0ab-6057c3b973a9
β_j = [ħ * 2π * 1e12 * f / k_B / T for f in phonon_mode_freqs, T in T_range]

# ╔═╡ 0641f475-0f35-4f5c-abf0-85ec56f4e30c
begin
	multi_beta = DataFrames.DataFrame([[0.0, [i for i in T_range]...]'; [phonon_mode_freqs β_j]], :auto)
	CSV.write("multi_beta.csv", multi_beta)
end

# ╔═╡ c49a4dc3-2546-472f-9fde-b59ec3246237
β_j_avg = [sum(β_j[:, T]) / length(phonon_mode_freqs) for T in 1:length(T_range)]

# ╔═╡ 0e3cf449-6e2a-4353-853e-58d15693743d
β_hellwarth_A = [ħ * 2π / k_B / T_range[T] * A_scheme[T] * 1e12 for T in 1:length(T_range)]

# ╔═╡ 16040051-57d5-4c6c-aa13-8280ae964261
β_hellwarth_B = [ħ * 2π / k_B / T * B_scheme * 1e12 for T in T_range]

# ╔═╡ 32f32534-2a6b-4675-8508-890579058e07
begin
	var_params = []
	push!(var_params, multi_variation(T_range[1], 4.5, 0.12, (6.29e-10)^3, MAPI; N = 1))
	for i in 2:length(T_range)
		push!(var_params, multi_variation(T_range[i], 4.5, 0.12, (6.29e-10)^3, MAPI; initial_vw = var_params[i-1], N = 1))
	end
	var_params
end

# ╔═╡ de89f0aa-21fb-42d2-9a75-5b45938b1e65
var_A = PolaronMobility.feynmanvw.(α_hellwarth_A, β_hellwarth_A)

# ╔═╡ a74836cf-f9ab-41d6-9f01-5066f1afac34
var_B = PolaronMobility.feynmanvw.(α_hellwarth_B, β_hellwarth_B)

# ╔═╡ 2ea4fc88-5285-4c8e-baf9-f0b2f75b33e0
v_j = [var_params[T][1] for T in 1:length(T_range)]

# ╔═╡ e5b21223-f7a5-4560-b03d-f015cdc99c41
sum.(v_j) ./ 2

# ╔═╡ ff8beae7-50e2-46cb-a38d-eb65a8cedaa4
w_j = [var_params[T][2] for T in 1:length(T_range)]

# ╔═╡ 757c41b7-e32b-4a67-b02d-c3f25acc266f
v_A = [i[1] for i in var_A]

# ╔═╡ 368ee8d7-84a8-4711-a9f9-f08d3c03c0ec
w_A = [i[2] for i in var_A]

# ╔═╡ b032bda5-4920-4546-a80f-5491b4044f45
v_B = [i[1] for i in var_B]

# ╔═╡ 5dd88a96-75d2-409a-abf8-960c09c74188
w_B = [i[2] for i in var_B]

# ╔═╡ 83ec77a6-fcf8-4025-ac6d-cd7819312427
begin
	vw_temp = plot(T_range, v_A, label = "H-A v", legend = :bottomright, linewidth = 2, tickfontsize = 12, labelfontsize = 12, legendfontsize = 12, xlabel = "T(K)", ylabel = "v & w (THz)", size = (550, 550), minorgrid = true, linestyle = :dashdot, color = theme_palette(:default)[1])
	plot!(T_range, w_A, label = "H-A w", linewidth = 2, linestyle = :dashdot, color = theme_palette(:default)[2])
	plot!(T_range, v_B, label = "H-B v", linewidth = 2, linestyle = :dash, color = theme_palette(:default)[3])
	plot!(T_range, w_B, label = "H-B w", linewidth = 2, linestyle = :dash, color = theme_palette(:default)[4])
	plot!(T_range, sum.(v_j) ./ length(v_j[1]), label = "multi v", linewidth = 2, linestyle = :solid, color = theme_palette(:default)[1])
	plot!(T_range, sum.(w_j) ./ length(w_j[1]), label = "multi w", linewidth = 2, linestyle = :solid, color = theme_palette(:default)[2])
end

# ╔═╡ 3a52ff01-110b-4755-b43a-acc5b3b870aa
savefig(vw_temp, "vw_temp.pdf")

# ╔═╡ 8fafb718-d95a-4ca7-921f-4abe1ab846ac
F_j = [multi_free_energy(v_j[T], w_j[T], T_range[T], 4.5, 0.12, (6.29e-10)^3, MAPI) for T in 1:length(T_range)]

# ╔═╡ 552b01e3-fd72-489a-aa2b-ae00b6d37c05
Hellwarth_energy_A = F.(v_A, w_A, β_hellwarth_A, α_hellwarth_A) .* 1e3 .* ħ .* 2π .* A_scheme .* 1e12 ./ eV

# ╔═╡ a465eeb5-5b8f-4202-bd58-651867c9d423
Hellwarth_energy_B = F.(v_B, w_B, β_hellwarth_B, α_hellwarth_B) .* 1e3 .* ħ .* 2π .* B_scheme .* 1e12 ./ eV

# ╔═╡ 6b60ab62-6706-4c33-ad7a-efb11803eb6b
begin
	free_energy_temp = plot(T_range, Hellwarth_energy_A, linewidth = 2, linestyle = :dashdot, label = "H-A", tickfontsize = 12, labelfontsize = 12, legendfontsize = 12, xlabel = "T(K)", ylabel = "Free energy (meV)", minorgrid = true, size = (550, 550), legend = :topright)
	plot!(T_range, Hellwarth_energy_B, linewidth = 2, linestyle = :dash, label = "H-B")
	plot!(T_range, F_j, linewidth = 2, linestyle = :solid, label = "multi")
end

# ╔═╡ 7ae9767d-d18f-4188-8e59-dea17c9d5ac4
savefig(free_energy_temp, "free_energy_temp.pdf")

# ╔═╡ 523af3a6-7cc9-41cf-9acb-9a1d25086969
begin
	multi_vw = DataFrames.DataFrame(
			temperature = T_range,
			beta = β_j_avg,
			v = v_j,
			w = w_j,
			F = F_j
	    )
	CSV.write("multi_vw.csv", multi_vw)
end

# ╔═╡ 8db8a4fd-fc04-4d45-8d05-8f69eff611df
begin
	A_data = DataFrames.DataFrame(
		alpha = α_hellwarth_A,
		efffreq = A_scheme,
		temp = T_range,
		beta = β_hellwarth_A,
		v = v_A,
		w = w_A,
		F = Hellwarth_energy_A
    )
	CSV.write("A_data.csv", A_data)
end

# ╔═╡ 0e921a7a-a07d-4afc-a991-e06d9be96e3f
begin
	B_data = DataFrames.DataFrame(
			alpha = [α_hellwarth_B for i in 1:length(T_range)],
			efffreq = [B_scheme for i in 1:length(T_range)],
			temp = T_range,
			beta = β_hellwarth_B,
			v = v_B,
			w = w_B,
			F = Hellwarth_energy_B
		)
	CSV.write("B_data.csv", B_data)
end

# ╔═╡ 2d4759ab-18d8-4368-a568-0da25474366c
function multi_conductivity(ν, β, α, v, w, ω, m_eff)
	z_integrand(t, β, ν) = (1 - exp(1im * 2π * ν * t)) * imag(cos(t - 1im * β / 2) / (  sinh(β / 2) * D_j(-1im * t, β, v, w)^(3/2)))
	z = 0.0
	for j in 1:length(ω)
		println("Photon frequency = $ν, Phonon mode frequency = $(ω[j] / 2π)")
		z += -1im * 2π * ν / length(ω) + 1im * 2 * α[j] * ω[j]^2 * quadgk(t -> z_integrand(t, β[j], ν / ω[j]), 0.0, Inf)[1] / (3 * √π * 2π * ν)
	end
	1 / z * eV * 100^2 / (m_eff * m_e * 1e12)
end

# ╔═╡ 2285006a-a115-414a-b788-b05fd6720bcd
ν_range = 0.01:0.01:3.51

# ╔═╡ 0fdcd049-8f06-4651-9864-a76647154680
begin
	σ_j = Array{ComplexF64}(undef, length(ν_range), length(T_range))
	for j in 1:length(T_range), i in 1:length(ν_range)
		println("ν: $(ν_range[i]) THz, T: $(T_range[j]) K")
		σ_j[i, j] = multi_conductivity(ν_range[i], β_j[:, j], α_j, v_j[1], w_j[1], phonon_mode_freqs .* 2π, 0.12)
		CSV.write("multi_conductivity_data.csv", DataFrames.DataFrame([[0.0, [i for i in T_range]...]'; [ν_range σ_j]], :auto))
	end
	σ_j
end

# ╔═╡ f09c4429-9415-49c1-b22d-a8867cbd4249
begin
	plot(ν_range, real.(σ_j), linewidth = 2, label = "Reσ", legend = :outerright, minorgrid = true)
	plot!(ν_range, imag.(σ_j), linewidth = 2, linestyle = :dash, label = "Imσ")
	plot!(ν_range, abs.(σ_j), linewidth = 2, linestyle = :dot, label = "|σ|")
end

# ╔═╡ 0178c326-eeca-4d7a-a111-58b589cba561
multi_contour_real = Plots.contourf(ν_range, T_range, real.(σ_j)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ d2bbb4c9-98e5-4959-a36c-b3e48c37a03c
multi_contour_imag = Plots.contourf(ν_range, T_range, imag.(σ_j)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ 0d78d4ef-6360-43b3-9952-3467410d8908
multi_contour_abs = Plots.contourf(ν_range, T_range, abs.(σ_j)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ 8822d1c3-5bda-4e4e-9b0b-703c75d10816
begin
	savefig(multi_contour_real, "multi_contour_real.pdf")
	savefig(multi_contour_imag, "multi_contour_imag.pdf")
	savefig(multi_contour_abs, "multi_contour_abs.pdf")
end

# ╔═╡ 72d02e9b-d59e-498b-8947-3a86634b7eb4
multi_plot_temp_real = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in real.(σ_j)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ 0916b2de-ff08-4c60-9bb1-70efbea6f6f6
multi_plot_temp_imag = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in imag.(σ_j)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ dfd78449-6db6-4f6b-9ecf-00a202ad685d
multi_plot_temp_abs = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in abs.(σ_j)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ 7d101426-86d9-4925-b4f8-52c2d5415491
begin
	savefig(multi_plot_temp_real, "multi_plot_temp_real.pdf")
	savefig(multi_plot_temp_imag, "multi_plot_temp_imag.pdf")
	savefig(multi_plot_temp_abs, "multi_plot_temp_abs.pdf")
end

# ╔═╡ ad91591f-1ccf-4347-8156-1d0e2e090fdb
begin
	σ_hellwarth_A = Array{ComplexF64}(undef, length(ν_range), length(T_range))
	for j in 1:length(T_range)
		for i in 1:length(ν_range)
			println("ν: $(ν_range[i]) THz, T: $(T_range[j]) K")
			σ_hellwarth_A[i, j] = multi_conductivity(ν_range[i], β_hellwarth_A[j], α_hellwarth_A[j], v_A[1], w_A[1], A_scheme[j] * 2π, 0.12)
		end
		CSV.write("A_conductivity_data.csv", DataFrames.DataFrame([[0.0, [i for i in T_range]...]'; [ν_range σ_hellwarth_A]], :auto))
	end
	σ_hellwarth_A
end

# ╔═╡ 2c4a84f5-6839-4ea4-bc9c-4eb2f8e6281b
begin
	term = 40
	plot(ν_range, real.(σ_hellwarth_A)[:, term], linewidth = 2, label = "Reσ", legend = :outerright, minorgrid = true, ylims = (0, 500))
	plot!(ν_range, imag.(σ_hellwarth_A)[:, term], linewidth = 2, linestyle = :dash, label = "Imσ")
	plot!(ν_range, abs.(σ_hellwarth_A)[:, term], linewidth = 2, linestyle = :dot, label = "|σ|")
end

# ╔═╡ ee677bc3-2e9e-4e35-a66c-4c6052c4e32e
A_contour_real = Plots.contourf(ν_range, T_range, real.(σ_hellwarth_A)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ 87154df5-6996-4b4b-9f63-45606205b945
A_contour_imag = Plots.contourf(ν_range, T_range, imag.(σ_hellwarth_A)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (-6, 600))

# ╔═╡ f4f84543-d7b8-43b2-93cb-0de8d7fa0a30
A_contour_abs = Plots.contourf(ν_range, T_range, abs.(σ_hellwarth_A)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true, scale = :log), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ 005705cd-a857-431a-8956-8ba0035c7c05
begin
	savefig(A_contour_real, "A_contour_real.pdf")
	savefig(A_contour_imag, "A_contour_imag.pdf")
	savefig(A_contour_abs, "A_contour_abs.pdf")
end

# ╔═╡ 77a6559f-6c40-4673-9a8d-9e0c8fe3476f
A_plot_temp_real = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in abs.(σ_hellwarth_A)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ 498c610b-0646-4a84-ac40-ab8c6efc5278
A_plot_temp_imag = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in imag.(σ_hellwarth_A)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ bc47027d-21a0-4a02-914d-aba2381d2e01
A_plot_temp_abs = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in abs.(σ_hellwarth_A)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ 1e902322-d319-42c5-889c-c3048515674c
begin
	savefig(A_plot_temp_real, "A_plot_temp_real.pdf")
	savefig(A_plot_temp_imag, "A_plot_temp_imag.pdf")
	savefig(A_plot_temp_abs, "A_plot_temp_abs.pdf")
end

# ╔═╡ 905efcbf-b3f5-4da0-a45c-3204cd69c0ff
begin
	σ_hellwarth_B = Array{ComplexF64}(undef, length(ν_range), length(T_range))
	for j in 1:length(T_range)
		for i in 1:length(ν_range)
			println("ν: $(ν_range[i]) THz, T: $(T_range[j]) K")
			σ_hellwarth_B[i, j] = multi_conductivity(ν_range[i], β_hellwarth_B[j], α_hellwarth_B, v_B[1], w_B[1], B_scheme * 2π, 0.12)
		end
		CSV.write("B_conductivity_data.csv", DataFrames.DataFrame([[0.0, [i for i in T_range]...]'; [ν_range σ_hellwarth_B]], :auto))
	end
	σ_hellwarth_B
end

# ╔═╡ 38a23177-633a-4b2d-8083-5e8afd74796a
begin
	term_B = 2
	plot(ν_range, real.(σ_hellwarth_B)[:, term_B], linewidth = 2, label = "Reσ", legend = :outerright, minorgrid = true, ylims = (0, 600))
	plot!(ν_range, imag.(σ_hellwarth_B)[:, term_B], linewidth = 2, linestyle = :dash, label = "Imσ")
	plot!(ν_range, abs.(σ_hellwarth_B)[:, term_B], linewidth = 2, linestyle = :dot, label = "|σ|")
end

# ╔═╡ 64821290-9bbe-4a7b-af4c-79825f6d4015
B_contour_real = Plots.contourf(ν_range, T_range, real.(σ_hellwarth_B)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ d292b963-3fd1-4d05-b72b-d71b66c5a919
B_contour_imag = Plots.contourf(ν_range, T_range, imag.(σ_hellwarth_B)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (-6, 600))

# ╔═╡ bfdc8a9b-9522-4f48-8b17-da79447994a8
B_contour_abs = Plots.contourf(ν_range, T_range, abs.(σ_hellwarth_B)', xlabel = "ν (THz)", ylabel = "T (K)", fill = cgrad(:thermal, rev = false, categorical = true), linewidth = 1.3, color = :grey10, size = (550, 500), tickfontsize = 12, labelfontsize = 12, right_margin = 7Plots.mm, clims = (0, 600))

# ╔═╡ 273c8ea1-0f7c-42b4-8954-67671d157181
begin
	savefig(B_contour_real, "B_contour_real.pdf")
	savefig(B_contour_imag, "B_contour_imag.pdf")
	savefig(B_contour_abs, "B_contour_abs.pdf")
end

# ╔═╡ b96343f4-bb61-4f4b-9ec9-ad2f37a96cfa
B_plot_temp_real = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in real.(σ_hellwarth_A)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ b4ee14e9-a89f-4a29-b6ed-d777fe477b9a
B_plot_temp_imag = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in imag.(σ_hellwarth_B)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ f711553f-e567-437f-9c38-79e644bd18b9
B_plot_temp_abs = ridgeline(round.(ν_range, digits = 1), Int.(T_range), [i < 600 ? i : 600 for i in abs.(σ_hellwarth_B)], jump = 5, shift = 10, xlabel = "ν (THz)", ylabel = "T (K)", size = (550, 550), linewidth = 0.5, palette = :thermal, fillalpha = 0.75, step = 5)

# ╔═╡ 26e1a5da-6bdb-49b9-973b-768d901a2a60
begin
	savefig(B_plot_temp_real, "B_plot_temp_real.pdf")
	savefig(B_plot_temp_imag, "B_plot_temp_imag.pdf")
	savefig(B_plot_temp_abs, "B_plot_temp_abs.pdf")
end

# ╔═╡ Cell order:
# ╠═29d62408-0629-4f39-909e-a6ba8696e955
# ╠═1eb559b3-dbce-4caf-800f-96a2495dee1b
# ╠═77d14c3f-bcaf-432c-9005-f5ec184d84b7
# ╠═ba2dadc8-d9a2-452e-9f94-406ccfc7404b
# ╠═c05cd489-68c8-418e-bbc2-a33bc42781f0
# ╠═1d300e7d-5efb-4ad2-a826-c3f8a86083d8
# ╠═6a99cd13-68cb-474c-96af-7d78ee4e7095
# ╠═49663b33-1604-430b-a815-9eb985802987
# ╠═314af942-8922-40d4-a92b-f75b715c276c
# ╠═6ff68ba5-df0c-42a0-8cf1-dedeced77e8d
# ╠═35390bdf-e260-4566-8a97-7a0b5a3ef07a
# ╠═e0ce47a5-368a-4da7-83fc-9f921af0f747
# ╠═053d717a-2b63-40e7-9dc0-487e3892de98
# ╠═86fe8a15-4d9e-48fa-8387-3cfaff7aa182
# ╠═16e7d781-bc61-44b3-84a2-e8c8892bf773
# ╠═a78017d7-8835-4a8f-8dd3-f3a4a50eb37b
# ╠═d7346684-6b68-4dfc-81da-12c435c4acf0
# ╠═798a4f7a-a2a0-48de-82ad-f85b4199496a
# ╠═ec0744c0-2530-11ec-04c3-1b5d83413b78
# ╠═f3b91e16-26d3-45c8-b28f-193afac15d7f
# ╠═a941ccc6-90cb-448a-91bb-f80233b630f8
# ╠═1372f247-4622-4388-b458-89aa6e695388
# ╠═c28f6d35-c3ed-4984-9c36-b282fb50001b
# ╠═49480256-2dcc-4d10-a5a1-4202b243cfd2
# ╠═4966ecf1-ef39-4da3-a8a5-9e47a6c6b0ec
# ╠═3321246b-76bd-4678-a5ae-332c3543a1ef
# ╠═b2944c8a-158e-461b-a83c-842925587179
# ╠═fe4e5ef5-833a-4b61-9ea9-92fa52d7a21f
# ╠═8350c7a2-e9fd-4312-9561-44e2524734af
# ╠═db2b4169-0737-44c1-9af5-aa2f2403f5db
# ╠═2b6d00e1-317b-41cc-9489-14a1450012f4
# ╠═3e1418c2-01a1-4651-8029-59f53264193d
# ╠═d4f392ef-968a-439c-ba53-5707464446e7
# ╠═b3c94ea5-d152-4705-9a30-f8d3e4da0874
# ╠═3608af5e-7b44-4d50-9770-5e8afd4fca28
# ╠═1716f263-f8ad-40a1-a805-d8b7ea6f38df
# ╠═da629bec-b53c-4acf-a0ab-6057c3b973a9
# ╠═0641f475-0f35-4f5c-abf0-85ec56f4e30c
# ╠═c49a4dc3-2546-472f-9fde-b59ec3246237
# ╠═0e3cf449-6e2a-4353-853e-58d15693743d
# ╠═16040051-57d5-4c6c-aa13-8280ae964261
# ╠═32f32534-2a6b-4675-8508-890579058e07
# ╠═de89f0aa-21fb-42d2-9a75-5b45938b1e65
# ╠═a74836cf-f9ab-41d6-9f01-5066f1afac34
# ╠═2ea4fc88-5285-4c8e-baf9-f0b2f75b33e0
# ╠═e5b21223-f7a5-4560-b03d-f015cdc99c41
# ╠═ff8beae7-50e2-46cb-a38d-eb65a8cedaa4
# ╠═757c41b7-e32b-4a67-b02d-c3f25acc266f
# ╠═368ee8d7-84a8-4711-a9f9-f08d3c03c0ec
# ╠═b032bda5-4920-4546-a80f-5491b4044f45
# ╠═5dd88a96-75d2-409a-abf8-960c09c74188
# ╠═83ec77a6-fcf8-4025-ac6d-cd7819312427
# ╠═3a52ff01-110b-4755-b43a-acc5b3b870aa
# ╠═8fafb718-d95a-4ca7-921f-4abe1ab846ac
# ╠═552b01e3-fd72-489a-aa2b-ae00b6d37c05
# ╠═a465eeb5-5b8f-4202-bd58-651867c9d423
# ╠═6b60ab62-6706-4c33-ad7a-efb11803eb6b
# ╠═7ae9767d-d18f-4188-8e59-dea17c9d5ac4
# ╠═523af3a6-7cc9-41cf-9acb-9a1d25086969
# ╠═8db8a4fd-fc04-4d45-8d05-8f69eff611df
# ╠═0e921a7a-a07d-4afc-a991-e06d9be96e3f
# ╠═2d4759ab-18d8-4368-a568-0da25474366c
# ╠═2285006a-a115-414a-b788-b05fd6720bcd
# ╠═0fdcd049-8f06-4651-9864-a76647154680
# ╠═f09c4429-9415-49c1-b22d-a8867cbd4249
# ╠═0178c326-eeca-4d7a-a111-58b589cba561
# ╠═d2bbb4c9-98e5-4959-a36c-b3e48c37a03c
# ╠═0d78d4ef-6360-43b3-9952-3467410d8908
# ╠═8822d1c3-5bda-4e4e-9b0b-703c75d10816
# ╠═72d02e9b-d59e-498b-8947-3a86634b7eb4
# ╠═0916b2de-ff08-4c60-9bb1-70efbea6f6f6
# ╠═dfd78449-6db6-4f6b-9ecf-00a202ad685d
# ╠═7d101426-86d9-4925-b4f8-52c2d5415491
# ╠═ad91591f-1ccf-4347-8156-1d0e2e090fdb
# ╠═2c4a84f5-6839-4ea4-bc9c-4eb2f8e6281b
# ╠═ee677bc3-2e9e-4e35-a66c-4c6052c4e32e
# ╠═87154df5-6996-4b4b-9f63-45606205b945
# ╠═f4f84543-d7b8-43b2-93cb-0de8d7fa0a30
# ╠═005705cd-a857-431a-8956-8ba0035c7c05
# ╠═77a6559f-6c40-4673-9a8d-9e0c8fe3476f
# ╠═498c610b-0646-4a84-ac40-ab8c6efc5278
# ╠═bc47027d-21a0-4a02-914d-aba2381d2e01
# ╠═1e902322-d319-42c5-889c-c3048515674c
# ╠═905efcbf-b3f5-4da0-a45c-3204cd69c0ff
# ╠═38a23177-633a-4b2d-8083-5e8afd74796a
# ╠═64821290-9bbe-4a7b-af4c-79825f6d4015
# ╠═d292b963-3fd1-4d05-b72b-d71b66c5a919
# ╠═bfdc8a9b-9522-4f48-8b17-da79447994a8
# ╠═273c8ea1-0f7c-42b4-8954-67671d157181
# ╠═b96343f4-bb61-4f4b-9ec9-ad2f37a96cfa
# ╠═b4ee14e9-a89f-4a29-b6ed-d777fe477b9a
# ╠═f711553f-e567-437f-9c38-79e644bd18b9
# ╠═26e1a5da-6bdb-49b9-973b-768d901a2a60


"""
    copras(decisionMat, weights, fs)
Apply COPRAS (COmplex PRoportional ASsesment) method for a given matrix, weights and, type of criteria.

# Arguments:
 - `decisionMat::DataFrame`: n × m matrix of objective values for n alternatives and m criteria
 - `weights::Array{Float64, 1}`: m-vector of weights that sum up to 1.0. If the sum of weights is not 1.0, it is automatically normalized.
 - `fs::Array{Function,1}`: m-vector of type of criteria. The benefit criteria shown with "maximum", and the cost criteria shown with "minimum".

 # Description
copras() applies the COPRAS method to rank n alternatives subject to m criteria and criteria type vector.

# Output
- `::COPRASResult`: COPRASResult object that holds multiple outputs including scores and best index.

# Examples
```julia-repl
julia> decmat = [2.50 240 57 45 1.10 0.333333;
       2.50 285 60 75 4.00 0.428571;
       4.50 320 100 65 7.50 1.111111;
       4.50 365 100 90 7.50 1.111111;
       5.00 400 100 90 11.00 1.111111;
       2.50 225 60 45 1.10 0.333333;
       2.50 270 57 60 4.00 0.428571;
       4.50 330 100 70 7.50 1.111111;
       4.50 365 100 80 7.50 1.111111;
       5.00 380 110 65 8.00 1.111111;
       2.50 285 65 80 4.00 0.400000;
       4.00 280 75 65 4.00 0.400000;
       4.50 365 102 95 7.50 1.111111;
       4.50 400 102 95 7.50 1.111111;
       6.00 450 110 95 11.00 1.176471;
       6.00 510 110 105 11.00 1.176471;
       6.00 330 140 110 18.50 1.395349;
       2.50 240 65 80 4.00 0.400000;
       4.00 280 75 75 4.00 0.400000;
       4.50 355 102 95 7.50 1.111111;
       4.50 385 102 90 7.50 1.111111;
       5.00 385 114 95 7.50 1.000000;
       6.00 400 110 90 11.00 1.000000;
       6.00 480 110 95 15.00 1.000000;
       6.00 440 140 100 18.50 1.200000;
       6.00 500 140 100 18.50 1.200000;
       5.00 450 125 100 15.00 1.714286;
       6.00 500 150 125 18.50 1.714286;
       6.00 515 180 140 22.00 2.307692;
       7.00 550 200 150 30.00 2.307692;
       6.00 500 180 140 15.00 2.307692;
       6.00 500 180 140 18.50 2.307692;
       6.00 500 180 140 22.00 2.307692;
       7.00 500 180 140 30.00 2.307692;
       7.00 500 200 140 37.00 2.307692;
       7.00 500 200 140 45.00 2.307692;
       7.00 500 200 140 55.00 2.307692;
       7.00 500 200 140 75.00 2.307692]
38×6 Array{Float64,2}:
 2.5  240.0   57.0   45.0   1.1  0.333333
 2.5  285.0   60.0   75.0   4.0  0.428571
 4.5  320.0  100.0   65.0   7.5  1.11111
 4.5  365.0  100.0   90.0   7.5  1.11111
 5.0  400.0  100.0   90.0  11.0  1.11111
 2.5  225.0   60.0   45.0   1.1  0.333333
 2.5  270.0   57.0   60.0   4.0  0.428571
 4.5  330.0  100.0   70.0   7.5  1.11111
 4.5  365.0  100.0   80.0   7.5  1.11111
 5.0  380.0  110.0   65.0   8.0  1.11111
 2.5  285.0   65.0   80.0   4.0  0.4
 ⋮                               ⋮
 6.0  500.0  150.0  125.0  18.5  1.71429
 6.0  515.0  180.0  140.0  22.0  2.30769
 7.0  550.0  200.0  150.0  30.0  2.30769
 6.0  500.0  180.0  140.0  15.0  2.30769
 6.0  500.0  180.0  140.0  18.5  2.30769
 6.0  500.0  180.0  140.0  22.0  2.30769
 7.0  500.0  180.0  140.0  30.0  2.30769
 7.0  500.0  200.0  140.0  37.0  2.30769
 7.0  500.0  200.0  140.0  45.0  2.30769
 7.0  500.0  200.0  140.0  55.0  2.30769
 7.0  500.0  200.0  140.0  75.0  2.30769

julia> df = makeDecisionMatrix(decmat)
38×6 DataFrame
 Row │ Crt1     Crt2     Crt3     Crt4     Crt5     Crt6
     │ Float64  Float64  Float64  Float64  Float64  Float64
─────┼───────────────────────────────────────────────────────
   1 │     2.5    240.0     57.0     45.0      1.1  0.333333
   2 │     2.5    285.0     60.0     75.0      4.0  0.428571
   3 │     4.5    320.0    100.0     65.0      7.5  1.11111
   4 │     4.5    365.0    100.0     90.0      7.5  1.11111
   5 │     5.0    400.0    100.0     90.0     11.0  1.11111
   6 │     2.5    225.0     60.0     45.0      1.1  0.333333
   7 │     2.5    270.0     57.0     60.0      4.0  0.428571
   8 │     4.5    330.0    100.0     70.0      7.5  1.11111
   9 │     4.5    365.0    100.0     80.0      7.5  1.11111
  ⋮  │    ⋮        ⋮        ⋮        ⋮        ⋮        ⋮
  30 │     7.0    550.0    200.0    150.0     30.0  2.30769
  31 │     6.0    500.0    180.0    140.0     15.0  2.30769
  32 │     6.0    500.0    180.0    140.0     18.5  2.30769
  33 │     6.0    500.0    180.0    140.0     22.0  2.30769
  34 │     7.0    500.0    180.0    140.0     30.0  2.30769
  35 │     7.0    500.0    200.0    140.0     37.0  2.30769
  36 │     7.0    500.0    200.0    140.0     45.0  2.30769
  37 │     7.0    500.0    200.0    140.0     55.0  2.30769
  38 │     7.0    500.0    200.0    140.0     75.0  2.30769
                                              20 rows omitted

julia> weights = [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667];

julia> fns = [maximum, maximum, maximum, maximum, maximum, minimum];

julia> result = copras(df, w, fs);

julia> result.bestIndex
38

julia> result.scores
38-element Array{Float64,1}:
 0.021727395411605937
 0.019814414550092637
 0.01789214190869233
 0.01624057709923278
 0.01507318798582843
 0.021837811311495522
 0.020735423365838293
 0.01754833367014725
 0.016658288043259514
 0.016258710042371068
 0.019388734250223458
 ⋮
 0.01064826476628102
 0.01000964101170639
 0.007056714858865703
 0.010627520153194962
 0.010419795157349462
 0.010212070161503961
 0.00876261130160425
 0.007733739643860005
 0.007258939653356005
 0.006665439665226005
 0.005478439688966004
```
# References
Kaklauskas, A., Zavadskas, E. K., Raslanas, S., Ginevicius, R., Komka, A., & Malinauskas, P. (2006). Selection of low-e windows in retrofit of public buildings by applying multiple criteria method COPRAS: A Lithuanian case. Energy and buildings, 38(5), 454-462.
Özdağoğlu, A. (2013). İmalat işletmeleri için eksantrik pres alternatiflerinin COPRAS yöntemi ile karşılaştırılması. Gümüşhane Üniversitesi Sosyal Bilimler Enstitüsü Elektronik Dergisi, 4(8), 1-22.
Yıldırım, B. F., Timor, M. (2019). "Bulanık ve Gri COPRAS Yöntemleri Kullanılarak Tedarikçi Seçim Modeli Geliştirilmesi". Optimum Ekonomi ve Yönetim Bilimleri Dergisi, 6 (2), 283-310.
"""
function copras(decisionMat::DataFrame, weights::Array{Float64,1}, fns::Array{Function,1})::COPRASResult

    #mat = convert(Matrix, decisionMat)
    mat = Matrix(decisionMat)

    nrows, ncols = size(mat)
    w = unitize(weights)
    normalizedMat = copy(mat)
    for col in 1:ncols
        for row in 1:nrows
            normalizedMat[row, col] = w[col] .* (mat[row, col] ./ sum(mat[:, col]))
        end
    end

    sPlus = zeros(Float64, nrows)
    sMinus = zeros(Float64, nrows)

    for row in 1:nrows
        for col in 1:ncols
            if fns[col] == maximum
                @inbounds sPlus[row] = sPlus[row] + normalizedMat[row, col]
            elseif fns[col] == minimum
                @inbounds sMinus[row] = sMinus[row] + normalizedMat[row, col]
            end
        end
    end

    Q = zeros(Float64, nrows)
    Z = sum(1 ./ sMinus)

    for row in 1:nrows
        Q[row] = sPlus[row] + (sum(sMinus) / (sMinus[row] * Z))
    end
    scores = zeros(Float64, nrows)
    scores = Q ./ maximum(Q)

    rankings = sortperm(scores)
    bestIndex = rankings |> last

    result = COPRASResult(
        decisionMat,
        w,
        scores,
        rankings,
        bestIndex
    )

    return result
end

"""
    copras(setting)
Apply COPRAS (COmplex PRoportional ASsesment) method for a given matrix, weights and, type of criteria.

# Arguments:
 - `setting::MCDMSetting`: MCDMSetting object.

 # Description
copras() applies the COPRAS method to rank n alternatives subject to m criteria and criteria type vector.

# Output
- `::COPRASResult`: COPRASResult object that holds multiple outputs including scores and best index.
"""
function copras(setting::MCDMSetting)::COPRASResult
    copras(
        setting.df,
        setting.weights,
        setting.fns
    )
end


@testset "Problems" begin
    F = equations(katsura(5))
    G = equations(cyclic(6))

    P1 = StartTargetInput(G, F)
    (PP1, start1) = problem_startsolutions(P1, [rand(ComplexF64, 6), rand(ComplexF64, 6)])
    @test PP1 isa Problem{AffineTracking}
    @test length(start1) == 2

    P2 = TargetSystemInput(G)
    (PP2, start2) = problem_startsolutions(P2)
    @test PP2 isa Problem{AffineTracking}
    @test length(start2) == 720

    @test homotopy(PP2) isa AbstractHomotopy
    @test homvars(PP2) == nothing

    @polyvar x y z
    F = [x^2+y^2+2z^2, x+y+3z]
    P, _ = problem_startsolutions(TargetSystemInput(F), nothing, homvar=y)
    @test homvars(P) == (2,)

    P, _ = problem_startsolutions(TargetSystemInput(F), nothing)
    @test homvars(P) === nothing

    @test_throws ArgumentError problem_startsolutions(StartTargetInput(
        [x^2+y^2+z^2, x^4+y^4+z^3], [x^3+z^3,y^3-z^3]), [rand(ComplexF64, 3)])

    P, _ = problem_startsolutions(StartTargetInput(
        [x^2+y^2+z^2, x^4+y^4+z^4], [x^3+z^3,y^3-z^3]), [rand(ComplexF64, 3)])
    @test P isa Problem{ProjectiveTracking}
    @test homvars(P) === nothing

    P, _ = problem_startsolutions(StartTargetInput(
        [x^2+y^2+z^2, x^4+y^4+z^2], [x^3+z^3,y^3-z^2]), [rand(ComplexF64, 3)])
    @test P isa Problem{AffineTracking}

    # not supported right now
    @test_throws MethodError problem_startsolutions(
        StartTargetInput(SPSystem([x^2+y^2+z^2, x^4+y^4+z^2]), SPSystem([x^3+z^3,y^3-z^2])),
        [rand(ComplexF64, 3)])

    P, starts = problem_startsolutions(HC.input_startsolutions(
        [x^2+y^2+1, x^4+y^4+1], [x^3+1,y^3-1], rand(ComplexF64, 3))...)
    @test P isa Problem{AffineTracking}
    @test length(starts) == 1

    P, _ = problem_startsolutions(StartTargetInput(
        [x^2+y^2+z^2, x^4+y^4+z^4], [x^3+z^3,y^3-z^3]), [rand(ComplexF64, 3)], homvar=z)
    @test homvars(P) === (3,)

    F = FPSystem(homogenize(equations(cyclic(6))))
    P, startvals = problem_startsolutions(TargetSystemInput(F))
    @test P isa Problem{ProjectiveTracking}
    @test homvars(P) == nothing

    F = homogenize(equations(cyclic(6)))
    P, _ = problem_startsolutions(TargetSystemInput(F))
    @test P isa Problem{ProjectiveTracking}

    F = FPSystem(homogenize(equations(cyclic(6))))
    P, startvals = problem_startsolutions(TargetSystemInput(F), homvar=5)
    @test P isa Problem{ProjectiveTracking}
    @test homvars(P) == (5,)
    @test length(startvals) == 720
end


module DirichletProcess

using Distributions

export DP, dp_sample!
export HDP, hdp_sample!

# Write your package code here.
mutable struct DP
    base_measure::Function
    alpha::Float64
    cache
    weights
    total_stick_used::Float64
    DP(base_measure, alpha) = new(base_measure, alpha, [], [], 0.0)
end

mutable struct HDP
    base_measure::Function
    alpha1::Float64
    alpha2::Float64
    dp1::DP
    dp2::DP
    function HDP(base_measure, alpha1, alpha2)
        dp1 = DP(base_measure, alpha1)
        dp2 = DP(() -> dp_sample!(dp1), alpha2)
        return new(base_measure, alpha1, alpha2, dp1, dp2)
    end
end

function dp_sample!(dp::DP)
    remaining = 1.0 - dp.total_stick_used
    i = rand(Categorical([dp.weights..., remaining]))
    if i <= length(dp.weights)
        return dp.cache[i]
    else
        stick_piece = rand(Beta(1, dp.alpha)) * remaining
        dp.total_stick_used += stick_piece
        push!(dp.weights, stick_piece)
        new_value = dp.base_measure()
        push!(dp.cache, new_value)
        return new_value
    end
end

function hdp_sample!(hdp::HDP)
    return dp_sample!(hdp.dp2)
end

end


export dubins, dubins_length, dubins_waypoints

dubins_length(q0::StaticVector{3}, qf::StaticVector{3}; r=1) = dubins(q0, qf, r=r).cost
function dubins_waypoints(q0::StaticVector{3}, qf::StaticVector{3}, dt_or_N; v=1, r=1)
    waypoints(SimpleCarDynamics{0,0}(), SE2State(q0), dubins(q0, qf, v=v, r=r).controls, dt_or_N)
end

function dubins((x0, y0, θ0)::StaticVector{3,T0}, (xf, yf, θf)::StaticVector{3,Tf}; v=1, r::R=1) where {T0,Tf,R}
    T = promote_type(T0, Tf, R)
    dx = (xf - x0)/r
    dy = (yf - y0)/r
    d = hypot(dx, dy)
    θ = atan(dy, dx)
    a = θ0 - θ
    b = θf - θ
    sa, ca = sincos(a)
    sb, cb = sincos(b)

    cmin = T(Inf)
    # LSL
    tmp = 2 + d*d - 2*(ca*cb + sa*sb - d*(sa - sb))
    if tmp >= 0
        θ = atan(cb - ca, d + sa - sb)
        t = mod2piF(-a + θ)
        p = sqrt(max(tmp, 0))
        q = mod2piF(b - θ)
        c = t + p + q

        cmin = c
        ctrl = SVector(
            carsegment2stepcontrol(1, t),
            carsegment2stepcontrol(0, p),
            carsegment2stepcontrol(1, q)
        )
    end

    # RSR
    tmp = 2 + d*d - 2*(ca*cb + sa*sb - d*(sb - sa))
    if tmp >= 0
        θ = atan(ca - cb, d - sa + sb)
        t = mod2piF(a - θ)
        p = sqrt(max(tmp, 0))
        q = mod2piF(-b + θ)
        c = t + p + q
        if c < cmin
            cmin = c
            ctrl = SVector(
                carsegment2stepcontrol(-1, t),
                carsegment2stepcontrol( 0, p),
                carsegment2stepcontrol(-1, q)
            )
        end
    end

    # RSL
    tmp = d*d - 2 + 2*(ca*cb + sa*sb - d*(sa + sb))
    if tmp >= 0
        p = sqrt(max(tmp, 0))
        θ = atan(ca + cb, d - sa - sb) - atan(T(2), p)
        t = mod2piF(a - θ)
        q = mod2piF(b - θ)
        c = t + p + q
        if c < cmin
            cmin = c
            ctrl = SVector(
                carsegment2stepcontrol(-1, t),
                carsegment2stepcontrol( 0, p),
                carsegment2stepcontrol( 1, q)
            )
        end
    end

    # LSR
    tmp = -2 + d*d + 2*(ca*cb + sa*sb + d*(sa + sb))
    if tmp >= 0
        p = sqrt(max(tmp, 0))
        θ = atan(-ca - cb, d + sa + sb) - atan(-T(2), p)
        t = mod2piF(-a + θ)
        q = mod2piF(-b + θ)
        c = t + p + q
        if c < cmin
            cmin = c
            ctrl = SVector(
                carsegment2stepcontrol( 1, t),
                carsegment2stepcontrol( 0, p),
                carsegment2stepcontrol(-1, q)
            )
        end
    end

    # RLR
    tmp = (6 - d*d  + 2*(ca*cb + sa*sb + d*(sa - sb)))/8
    if abs(tmp) < 1
        p = 2*T(pi) - acos(tmp)
        θ = atan(ca - cb, d - sa + sb)
        t = mod2piF(a - θ + p/2)
        q = mod2piF(a - b - t + p)
        c = t + p + q
        if c < cmin
            cmin = c
            ctrl = SVector(
                carsegment2stepcontrol(-1, t),
                carsegment2stepcontrol( 1, p),
                carsegment2stepcontrol(-1, q)
            )
        end
    end

    # LRL
    tmp = (6 - d*d  + 2*(ca*cb + sa*sb - d*(sa - sb)))/8
    if abs(tmp) < 1
        p = 2*T(pi) - acos(tmp)
        θ = atan(-ca + cb, d + sa - sb)
        t = mod2piF(-a + θ + p/2)
        q = mod2piF(b - a - t + p)
        c = t + p + q
        if c < cmin
            cmin = c
            ctrl = SVector(
                carsegment2stepcontrol( 1, t),
                carsegment2stepcontrol(-1, p),
                carsegment2stepcontrol( 1, q)
            )
        end
    end

    ctrl = scalespeed.(scaleradius.(ctrl, r), v)
    (cost=cmin*r/v, controls=ctrl)
end


"""
    $SIGNATURES

Calculate the ellipse from covariance matrix and return as lambda function.

Notes
- https://cookierobotics.com/007/

Related

plotCovEllipseLayer
"""
function covEllipseParameterized(distr::MvNormal; meanOffset::Bool=true)
  # println(round.(cov(distr), digits=3) )
  a = cov(distr)[1,1]
  b = cov(distr)[1,2] + cov(distr)[2,1]
  b *= 0.5
  c = cov(distr)[2,2]

  λ1 = 0.5*(a+c) + sqrt(0.25*(a-c)^2 + b^2)
  λ2 = 0.5*(a+c) - sqrt(0.25*(a-c)^2 + b^2)
  θ = if b ≈ 0 && a >= c
    0
  elseif b ≈ 0 && a < c
    pi/2
  else
    atan(λ1-a, b)
  end
  sλ1 = sqrt(λ1)
  sλ2 = sqrt(λ2)
  sθ =  sin(θ)
  cθ =  cos(θ)
  ox = meanOffset ? distr.μ[1] : 0
  oy = meanOffset ? distr.μ[2] : 0
  (t) -> [sλ1*cθ*cos(t)-sλ2*sθ*sin(t)+ox; sλ1*sθ*cos(t)+sλ2*cθ*sin(t)+oy]
end

covEllipseParameterized(pts::Array{Float64,2}; meanOffset::Bool=true) = covEllipseParameterized( fit(MvNormal, pts), meanOffset=meanOffset )
covEllipseParameterized(X::Union{<:BallTreeDensity,<:ManifoldKernelDensity}; meanOffset::Bool=true) = covEllipseParameterized( getPoints(X), meanOffset=meanOffset )
covEllipseParameterized(dfg::AbstractDFG, sym::Symbol; meanOffset::Bool=true, solveKey::Symbol=:default) = covEllipseParameterized( getBelief(dfg, sym, solveKey), meanOffset=meanOffset, solveKey=solveKey )

"""
    $SIGNATURES

Plotting tool to draw Gadfly layers of ellipses of 2D covariance fitted to the belief of factor graph variable nonparametric points.

Related

covEllipseParameterized, plotSLAM2DPosesLandms
"""
function plotCovEllipseLayer( dfg::AbstractDFG,
                              vsym::Symbol;
                              solveKey::Symbol=:default,
                              drawPoints::Bool=true,
                              ellipseColor::AbstractString="gray30",
                              pointsColor::AbstractString="gray30",
                              drawEllipse::Bool=true  )
  #
  PL = []

  if !(drawEllipse || drawPoints)
    return PL
  end

  # points to work from
  bel = getBelief(dfg, vsym, solveKey)
  pp__ = getPoints(bel)
  pp_ = if bel.manifold isa SpecialEuclidean
    # assume Pose2
    (x->x.parts[1]).(pp__)
  else
    # assume Point2
    pp__
  end

  @cast pp[i,j] := pp_[j][i]

  if drawEllipse
    # get ellipse function
    eX = covEllipseParameterized(pp[1:2,:], meanOffset=false)
    vEl = eX.(0:0.02:2pi)
    el = [(x->x[1]).(vEl) (x->x[2]).(vEl)]
    # add suggested PPE mean offset
    el[:,1] .+= getVariablePPE(dfg, vsym, solveKey).suggested[1]
    el[:,2] .+= getVariablePPE(dfg, vsym, solveKey).suggested[2]

    # add the ellipse layers
    plelX2 = Gadfly.layer(x=el[:,1], y=el[:,2], Geom.path, Theme(default_color=parse(Colorant, ellipseColor)))
    push!(PL, plelX2[1])
  end

  # add the points layer if needed
  if drawPoints
    plelX1 = Gadfly.layer(x=pp[1,:],
                          y=pp[2,:],
                          Geom.point,
                          Theme(default_color=parse(Colorant, pointsColor),
                          point_size=1pt))
    #
    push!(PL, plelX1[1])
  end

  return PL
end


@testset "Matrix algebras" begin

  @testset "Construction" begin

    M1 = matrix(FlintQQ, [ 1 1; 0 1])
    M2 = matrix(FlintQQ, [ 1 0; 1 1])
    A = matrix_algebra(FlintQQ, [ M1, M2 ])
    @test dim(A) == 4

    QG = group_algebra(FlintQQ, small_group(2, 1))
    M1 = matrix(QG, [ 1 1; 0 1])
    M2 = matrix(QG, [ 1 0; 1 1])
    A = matrix_algebra(FlintQQ, QG, [ M1, M2 ])
    @test dim(A) == 4
  end

end


using Makie
function img(m, n)
	a = zeros(Int, m, n)
	for (i, z) in enumerate(GoldenCartesianSequence((m, n)))
		a[z] = i
		i > m*n && break
	end
	!any(a .== 0) && @warn("Not spacefilling")
	a
end

m, n = 359, 451
m, n = @. denominator(rationalize(GoldenSequences.phis[2]^(-[1, 2]), tol=0.0000001))
@assert gcd(m, n) == 1
a = img(m, n)

pl = heatmap( -a[1:200, 1:200] )
save("cartesian1.png", pl)

pt(c) = Point2f0(c[1], c[2])

pl = scatter(pt.(take(GoldenCartesianSequence((m, n)), 20*20)), color=:red, markersize=20.0)
scatter!(pl, pt.(take(GoldenCartesianSequence((m, n)), m*n ÷ 525)), markersize=7.0)
#scatter!(pt.(take(GoldenCartesianSequence((m, n)), m*n ÷ 2)), grid=false, markersize=1.0)

save("cartesian2.png", pl)