using LinearAlgebra
using Distributions, Random

#softmax(x) = exp.(x) / sum(exp.(x))
function softmax(x)
    x = clamp.(x, -32., 32.)
    return exp.(x) / sum(exp.(x))
end

function sample_discrete(p, rng)
    n = length(p)
    i = 1
    c = p[1]
    u = rand(rng)
    while c < u && i < n
        c += p[i += 1]
    end

    return i
end

mutable struct LinearSoftmaxPolicy{TP, TB} <: AbstractPolicy where {TP<:Real,TB<:AbstractFuncApprox}
    f::LinearFunction{TP, TB}
    action::Int
    probs::Array{TP,1}
    function LinearSoftmaxPolicy(state_dim::Int, num_actions::Int)
        f = LinearFunction(state_dim, num_actions)
        new{eltype(f.θ),typeof(f.ϕ)}(f, 0, zeros(eltype(f.θ), num_actions))
    end
    function LinearSoftmaxPolicy(::Type{T}, state_dim::Int, num_actions::Int) where {T<:Real}
        f = LinearFunction(T, state_dim, num_actions)
        new{T,typeof(f.ϕ)}(f, 0, zeros(T, num_actions))
    end
    function LinearSoftmaxPolicy(f::LinearFunction{T}) where {T<:Real}
        new{T,typeof(f.ϕ)}(clone(f), 0, zeros(T, get_num_outputs(f)))
    end
end

function get_num_params(π::LinearSoftmaxPolicy)::Int
    return get_num_params(π.f)
end

function get_action_probabilities!(π::LinearSoftmaxPolicy, x)
    call!(π.probs,π.f,x)
    π.probs .= softmax(π.probs)
end

function get_action_probabilities(π::LinearSoftmaxPolicy{T}, x)::Array{T,1} where {T<:Real}
    get_action_probabilities!(π, x)
    return π.probs
end

function get_action!(π::LinearSoftmaxPolicy{T}, x, rng::AbstractRNG)::T where {T<:Real}
    get_action_probabilities!(π, x)
    π.action = sample_discrete(π.probs, rng)
    logp = log(π.probs[π.action])
    return logp
end

function gradient_logp!(grad::Array{T}, π::LinearSoftmaxPolicy{T, IdentityBasis}, x::Int, action::Int)::T where {T<:Real}
    grad = reshape(grad, size(get_params(π.f)))
    fill!(grad, 0.)

    get_action_probabilities!(π, x)
    # probs = -π.probs
    logp = log(π.probs[action])
    # probs[action] += 1

    grad[x, :] .-= π.probs
    grad[x, action] += 1.

    return logp
end

function gradient_logp!(grad::Array{T}, π::LinearSoftmaxPolicy{T, IdentityBasis}, x::Array{T, 1}, action::Int)::T where {T<:Real}
    grad = reshape(grad, size(get_params(π.f)))
    fill!(grad, 0.)

    get_action_probabilities!(π, x)
    probs = -π.probs
    logp = log(π.probs[action])
    probs[action] += 1.

    for i in 1:length(probs)
        @. grad[:, i] = x * probs[i]
    end
    # grad .= x*probs'

    return logp
end

function gradient_logp!(grad::Array{T}, π::LinearSoftmaxPolicy{T}, x, action::Int)::T where {T<:Real}
    grad = reshape(grad, size(get_params(π.f)))
    fill!(grad, 0.)

    get_action_probabilities!(π, x)
    probs = -π.probs
    logp = log(π.probs[action])
    probs[action] += 1.
    for i in 1:length(probs)
        @. grad[:, i] = π.f.feats * probs[i]
    end
    # grad .= π.f.feats*probs'

    return logp
end

function gradient_entropy!(grad, π::LinearSoftmaxPolicy{T, IdentityBasis}, x::Int; reuse::Bool=false) where {T}
    grad = reshape(grad, size(get_params(π.f)))
    fill!(grad, 0.)
    if !reuse
        get_action_probabilities!(π, x)
    end

    p1 = zeros(length(π.probs))
    @. p1 = -π.probs * log(π.probs)
    H = sum(p1)  # entropy (p1 has negative in it)
    @. p1 += π.probs * H
    grad[x, :] .= p1
    return H
end

function gradient_entropy!(grad, π::LinearSoftmaxPolicy{T, IdentityBasis}, x::Array{T, 1}; reuse::Bool=false) where {T}
    grad = reshape(grad, size(get_params(π.f)))
    fill!(grad, 0.)
    if !reuse
        get_action_probabilities!(π, x)
    end

    p1 = zeros(length(π.probs))
    @. p1 = -π.probs * log(π.probs)
    H = sum(p1)  # entropy (p1 has negative in it)
    @. p1 += π.probs * H

    for i in 1:length(p1)
        @. grad[:, i] = x * p1[i]
    end
    # grad .= x*p1'
    return H
end

function gradient_entropy!(grad, π::LinearSoftmaxPolicy{T}, x; reuse::Bool=false) where {T}
    grad = reshape(grad, size(get_params(π.f)))
    fill!(grad, 0.)
    if !reuse
        get_action_probabilities!(π, x)
    end

    p1 = zeros(length(π.probs))
    @. p1 = -π.probs * log(π.probs)
    H = sum(p1)  # entropy (p1 has negative in it)
    @. p1 += π.probs * H

    for i in 1:length(p1)
        @. grad[:, i] = π.f.feats * p1[i]
    end

    # grad .= π.f.feats*p1'
    return H
end

function get_action_gradient_logp!(grad::Array{T}, π::LinearSoftmaxPolicy{T, IdentityBasis}, x::Int, rng::AbstractRNG)::T where {T<:Real}
    get_action_probabilities!(π, x)
    π.action = sample_discrete(π.probs, rng)
    logp = log(π.probs[π.action])
    # probs = -π.probs
    # probs[π.action] += 1.
    fill!(grad, 0.)
    grad[x, :] .-= π.probs
    grad[x, π.action] += 1.
    # grad[x, :] .= probs
    return logp
end

function get_action_gradient_logp!(grad::Array{T}, π::LinearSoftmaxPolicy{T, IdentityBasis}, x::Array{T, 1}, rng::AbstractRNG)::T where {T<:Real}
    get_action_probabilities!(π, x)
    π.action = sample_discrete(π.probs, rng)
    logp = log(π.probs[π.action])
    probs = -π.probs
    probs[π.action] += 1.
    grad .= x*probs'
    return logp
end

function get_action_gradient_logp!(grad::Array{T}, π::LinearSoftmaxPolicy{T}, x, rng::AbstractRNG)::T where {T<:Real}
    get_action_probabilities!(π, x)
    π.action = sample_discrete(π.probs, rng)
    logp = log(π.probs[π.action])
    probs = deepcopy(-π.probs)
    probs[π.action] += 1.
    grad .= π.f.feats*probs'
    return logp
end

function get_action_gradient_logp(grad::Array{T}, π::LinearSoftmaxPolicy{T}, x, rng::AbstractRNG) where {T<:Real}
    grad = zeros(T, size(get_params(π)))
    logp = get_action_gradient_logp!(grad, π, x)
    action = π.action
    return action, grad, logp
end

function set_params!(π::LinearSoftmaxPolicy{T}, θ) where {T<:Real}
    set_params!(π.f, θ)
end

function get_params(π::LinearSoftmaxPolicy)
    get_params(π.f)
end

function copy_params!(params::Array{T}, π::LinearSoftmaxPolicy{T}) where {T<:Real}
    copy_params!(params, π.f)
end

function copy_params(π::LinearSoftmaxPolicy{T})::Array{T} where {T<:Real}
    return copy_params(π.f)
end

function add_to_params!(π::LinearSoftmaxPolicy{T}, Δθ) where {T<:Real}
    add_to_params!(π.f, Δθ)
end

function clone(π::LinearSoftmaxPolicy)::LinearSoftmaxPolicy
    π₂ = LinearSoftmaxPolicy(π.f)
    π₂.action = deepcopy(π.action)
    π₂.probs = deepcopy(π.probs)
    return π₂
end


#=
author: mingchuno
smart permutation inspired by python-high-perf
=#
array = [0 ,2 ,3, 4, 5, 6, 7 ,8, 9]
sarray = Set(array)

for last4_arr in combinations(array, 4), efgh_arr in permutations(last4_arr)
  e = efgh_arr[1]
  if e == 0 continue end # fail fast

  f = efgh_arr[2]

  g = efgh_arr[3]
  if g == 0 continue end # fail fast

  h = efgh_arr[4]

  ef = e*10 + f
  gh = g*10 + h

  if ef + gh == 111
    remain_arr = collect(setdiff(sarray, Set(last4_arr)))
    for first4_arr in combinations(remain_arr, 4), abcd_arr in permutations(first4_arr)
      a = abcd_arr[1]
      if a == 0 continue end # fail fast

      b = abcd_arr[2]

      c = abcd_arr[3]
      if c == 0 continue end # fail fast

      d = abcd_arr[4]
      ab = a*10 + b
      cd = c*10 + d
      if ab - cd == ef
        println("($(ab) - $(cd) = $(ef)) + $(gh) = 111")
      end
    end
  end
end


"""
For converting the concentrations into molecules

Arguments:  cell_radius = the radius of the modelled cell (m)
            baseline,
            indicator,
            endogeneous,
            immobile,
            peak,
            f_i,
            f_e,
            f_im
Returns:    m_baseline, m_indicator, m_endogeneous, m_immobile, m_peak, m_f_i,
            m_f_e, m_f_im
"""
function convertUnits(cell_radius::Float64, baseline::Float64,  indicator::Float64, endogeneous::Float64,
                      immobile::Float64,    peak::Float64,      f_i::Float64,       f_e::Float64,
                      f_im::Float64)
  cell_volume = sphereVolume(cell_radius)
  m_baseline = molarsToMolecules(baseline, cell_volume)
  m_indicator = molarsToMolecules(indicator, cell_volume)
  m_endogeneous = molarsToMolecules(endogeneous, cell_volume)
  m_immobile = molarsToMolecules(immobile, cell_volume)
  m_peak = molarsToMolecules(peak, cell_volume)
  m_f_i = perMolarToPerMolecule(f_i, cell_volume)
  m_f_e = perMolarToPerMolecule(f_e, cell_volume)
  m_f_im = perMolarToPerMolecule(f_im, cell_volume)
  return cell_volume, m_baseline, m_indicator, m_endogeneous, m_immobile, m_peak, m_f_i, m_f_e, m_f_im
end


using Plots,FFTW,JLD2

function sqaure_windowing_signal(signal::Array{Float64})
    windowed_signal = signal
    window_correction_factor = 1
    return (windowed_signal,window_correction_factor)
end

function hanning_windowing_signal(signal::Array{Float64})
    windowed_signal::Array{Float64} = []
    window::Array{Float64} = []
    total_window_power = 0
    j = 1
    for w = signal
        window_val = 0.5*(1 - cos(2*pi*j/length(signal)))
        push!(window,window_val)
        windowed = window_val*w
        push!(windowed_signal, windowed)
        total_window_power += window_val
        j += 1
    end
    window_correction_factor = length(signal)/total_window_power
    return (windowed_signal,window_correction_factor)
end

function main()
    sampling_rate = 0
    window_size = 64*1024*1024
    w_fm_vals::Array{Float64} = []

    jldopen(ARGS[1],"r") do f
        sampling_rate = f["samplingrate"]
        w_fm_vals = f["w_fm"]
    end

    (windowed_signal, window_correction_factor) = sqaure_windowing_signal(w_fm_vals[1:window_size])

    powers = map(x->20*log10(abs(x)*window_correction_factor),fft(windowed_signal)|>fftshift)
    freqs = fftfreq( window_size, sampling_rate )|>fftshift

    println(freqs)

    plot_graph=plot(freqs, powers, size=(2000,700), fmt =:png)
    savefig(plot_graph, "fm_fft1.png")

    plot_graph=plot(freqs, powers, xlims=(7.6*10^6, 8.0*10^6), size=(2000,700), fmt =:png)
    savefig(plot_graph, "fm_fft2.png")

end

main()


struct FurutaParams{T}
    α::T
    β::T
    γ::T
    δ::T
    Ku::T
    fϕ::T
    fθ::T
    h::T
    max_speed::T
    max_torque::T
    noise::T
    friction::Symbol
end

mutable struct SimulatedFurutaPendulum{T, R<:AbstractRNG} <: SimulatedProcess
    params::FurutaParams{T}
    x::Vector{T} # θ, θdot, ϕ, ϕdot = base angle & speed, arm angle & speed
    u::T
    rng::R
end

"""
    SimulatedFurutaPendulum(;
        J = 1.54e-4,
        M = 0,
        ma = 0,
        mp = 5.44e-3,
        la = 4.3e-2,
        lp = 6.46e-2,
        friction = :none,
        fϕ = 0.0238,
        fθ = 0.000008,
        R = 0.13,
        Ka = 0.03,
        h = 0.002,
        x0 = zeros(4),
        max_speed = 100,
        max_torque = 1,
        noise = 0.01,
        friction = :viscous,
        rng = Random.GLOBAL_RNG,
    )

Creates an instance of a simulator for the furutapendulum.
It is made up of two inertial bodies, a center pillar with
intertia `J` connected to a horizontal arm with length `la`
and homogenuously distributed mass `ma`. The other end of the
arm is connected to the pendulum which has length `lp`, homogenuously
distributed mass `mp` and point distributed mass `M` at the end.
The `ϕ`-joint is in the center pillar while the `θ`-joint is for the
pendulum arm. There is significant `friction` in the `ϕ`-joint which
is modeled using either coulomb and viscous friction `:viscous`,
or coulomb friction and stiction `:stiction` (for more info see Gäfvert
https://portal.research.lu.se/portal/files/4453844/8727127.pdf).
The simulation takes steps of `h` seconds, and the initial value `x0`
is for `[ϕ, ϕdot, θ, θdot]`.
There is `max_speed` and `max_torque` limiting the speed and torque, and
`noise` that adds white noise with this scaling factor on top of the
measurements.
"""
function SimulatedFurutaPendulum(;
        J = 1.54e-4,
        M = 0,
        ma = 0,
        mp = 5.44e-3,
        la = 4.3e-2,
        lp = 6.46e-2,
        fϕ = 0.0238,
        fθ = 0.000008,
        R = 0.13,
        Ka = 0.03,
        h = 0.002, # Simulation time, TODO find what this actually is
        max_speed = 100,
        max_torque = 0.04,
        x0 = [0.0, 0.0, π, 0.0], # Start down as default
        noise = 0.01,
        friction::Symbol = :viscous,
        rng = Random.GLOBAL_RNG,
    )
    g = 9.81
    α = J+(M+ma/3+mp)*la^2
    β = (M+mp/3)*lp^2
    γ = (M+mp/2)*la*lp
    δ = (M+mp/2)*g*lp
    Ku = Ka / R

    params = FurutaParams{Float64}(α, β, γ, δ, Ku, fϕ, fθ, h, max_speed, max_torque, noise, friction)
    SimulatedFurutaPendulum{Float64, typeof(rng)}(params, x0, 0.0, rng)
end

function step!(p::SimulatedFurutaPendulum, dt)
    x = p.x
    u = p.u
    h = p.params.h

    steps = round(Int, dt / h)
    steps ≈ dt / h || throw(ArgumentError("`dt`=$(dt) need to be a multiple of the simulation step `h`=$(h)."))
    for _ in 1:steps
        #RK3/8
        k1 = f(p, x, u)
        k2 = f(p, x + h * k1 / 3, u)
        k3 = f(p, x + h * (-k1 / 3 + k2), u)
        k4 = f(p, x + h * (k1 - k2 + k3), u)
        x .+= h .* (k1 .+ 3 .* k2 .+ 3 .* k3 .+ k4) ./ 8
    end

    x[1] = mod2pi(x[1])
    x[3] = mod2pi(x[3])

    # Should this be here?
    x[2] = clamp(x[2] , -p.params.max_speed, p.params.max_speed)
    x[4] = clamp(x[4] , -p.params.max_speed, p.params.max_speed)

    nothing
end

# See https://portal.research.lu.se/portal/files/4453844/8727127.pdf
# for source of dynamic equations
function f(p::SimulatedFurutaPendulum, x, u)
    # Base angle/speed, arm angle/speed
    ϕ, ϕdot, θ, θdot = x

    @unpack α, β, γ, δ, Ku, fϕ, fθ = p.params

    # clamp(Ku * p.u, -p.params.max_torque, p.params.max_torque)
    τϕm = Ku * u # Torque from motor (approximating that transients between voltage and current are fast)
    τϕf = -fϕ * ϕdot
    τϕ = τϕm + τϕf
    # theta joint
    τθ = -fθ * θdot

    ψ = 1 / (α*β - γ^2 + (β^2 + γ^2)*sin(θ)^2)
    dϕ    = ϕdot
    dϕdot = ψ*(β*γ*(sin(θ)^2-1)*sin(θ)*ϕdot^2 - 2*β^2*cos(θ)*sin(θ)*ϕdot*θdot
        + β*γ*sin(θ)*θdot^2 - γ*δ*cos(θ)*sin(θ) + β*τϕ - γ*cos(θ)*τθ)
    dθ    = θdot
    dθdot = ψ * (β*(α+β*sin(θ)^2)*cos(θ)*sin(θ)*ϕdot^2
        + 2*β*γ*(1-sin(θ)^2)*sin(θ)*ϕdot*θdot - γ^2*cos(θ)*sin(θ)*θdot^2
        + δ*(α+β*sin(θ)^2)*sin(θ) - γ*cos(θ)*τϕ + (α+β*sin(θ)^2)*τθ)

    return [dϕ, dϕdot, dθ, dθdot]
end

# To allow reading values on the same low level interface as for the hardware
# using read(furuta.port)
# Currently only implemented for energy readings
struct ValueContainer
    v::Float64
end

function Base.getproperty(p::SimulatedFurutaPendulum, s::Symbol)
    K0 = 0.0000652 # kinetic energy constant base: 0.5*moment of inertia, kgm^2
    K1 = 0.00000387 # kinetic energy constant pendulum

    if s === :arm_energy
        ϕ, ϕdot, θ, θdot = p.x
        arm_energy = (K1 * θdot^2 + 0.0054 * 9.8 * 0.07/2 * (1 - cos(θ)))
        return ValueContainer(arm_energy)
    elseif s === :total_energy
        ϕ, ϕdot, θ, θdot = p.x
        arm_energy = (K1 * θdot^2 + 0.0054 * 9.8 * 0.07/2 * (1 - cos(θ)))
        base_energy = K0 * ϕdot^2
        return ValueContainer(base_energy + arm_energy)
    else
        return getfield(p, s)
    end
end

Base.read(c::ValueContainer) = c.v


"""
    FiniteDifferenceGrids

A simple 1-dimensional uniform grid
for finite difference method.
"""
module FiniteDifferenceGrids

export Grid, over_elems, over_elems_real, over_elems_ghost
export ZBoundary, Zmin, Zmax, n_hat, binary, ghost_vec, ghost_dual
export first_interior, boundary, second_interior, first_ghost, boundary_points

"""
    Grid{T, I}

A simple 1-dimensional uniform grid of
type `T` for finite difference method.
"""
struct Grid{T, I}
  """
  Number of physical elements
  """
  n_elem :: I

  """
  Number of real elements
  """
  n_elem_real :: I

  """
  Number of ghost points per side
  """
  n_ghost :: I

  """
  Element size
  """
  Δz :: T

  """
  Inverse of element size squared
  """
  Δzi :: T

  """
  Inverse of element size squared
  """
  Δzi2 :: T

  """
  z-coordinate at cell edges
  """
  ze :: Vector{T}

  """
  z-coordinate at cell centers
  """
  zc :: Vector{T}

  """
  z-coordinate at bottom of boundary node
  """
  zn_min :: T

  """
  z-coordinate at top of boundary node
  """
  zn_max :: T
end

"""
    Grid(z_min::T, z_max::T, n_elem_real::Int, n_ghost::Int = 1)

A simple grid implementation that accepts the domain interval
`z_min` and `z_max`, the number of elements `n_elem_real` and
the number of ghost points per side `n_ghost`.
"""
function Grid(z_min::T, z_max::T, n_elem_real::I, n_ghost::I = 1) where {T, I}
  n_elem = n_elem_real+2*n_ghost
  Δz = (z_max-z_min)/n_elem_real
  Δzi = inv(Δz)
  Δzi2 = Δzi^2
  zc = [(i-n_ghost)*Δz-Δz/2 for i in 1:n_elem]
  ze = [(i-n_ghost)*Δz-Δz for i in 1:n_elem+1]
  zn_min = (zc[n_ghost]+zc[1+n_ghost])/2
  zn_max = (zc[end-n_ghost]+zc[end-n_ghost+1])/2
  return Grid{T, I}(n_elem, n_elem_real, n_ghost, Δz, Δzi, Δzi2, ze, zc, zn_min, zn_max)
end

function Base.show(io::IO, grid::Grid)
  println(io, "----------------- Grid")
  println(io, "zc      = ",grid.zc)
  println(io, "ze      = ",grid.ze)
  println(io, "n_elem  = ",grid.n_elem)
  println(io, "n_ghost = ",grid.n_ghost)
  println(io, "zn_min  = ",grid.zn_min)
  println(io, "zn_max  = ",grid.zn_max)
  println(io, "-----------------")
end

abstract type ZBoundary end

"""
    Zmax
Type to dispatch function to particular location
"""
struct Zmax <: ZBoundary end

"""
    Zmin
Type to dispatch function to particular location
"""
struct Zmin <: ZBoundary end

"""
    n_hat(::ZBoundary)

The outward normal vector to the boundary
"""
n_hat(::Zmin) = -1
n_hat(::Zmax) = 1

"""
    binary(::ZBoundary)

Returns 0 for Zmin and 1 for Zmax
"""
binary(::Zmin) = 0
binary(::Zmax) = 1

"""
    ghost_vec(::ZBoundary)

A 3-element vector near the boundary with 1's on ghost cells and 0's on interior cells
"""
ghost_vec(::Zmin) = [1, 0, 0]
ghost_vec(::Zmax) = [0, 0, 1]

"""
    ghost_dual(::ZBoundary)

A 2-element vector near the boundary with 1's on ghost cells and 0's on interior cells
"""
ghost_dual(::Zmin) = [1, 0]
ghost_dual(::Zmax) = [0, 1]

"""
    over_elems(grid::Grid)

Get the range of indexes to traverse real and ghost grid elements
"""
@inline over_elems(grid::Grid) = 1:grid.n_elem

"""
    over_elems_real(grid::Grid)

Get the range of indexes to traverse only real grid elements
"""
@inline over_elems_real(grid::Grid) = 1+grid.n_ghost : grid.n_elem-grid.n_ghost

"""
    first_interior(grid::Grid, ::ZBoundary)

Get the first element index in the physical domain
"""
@inline first_interior(grid::Grid, ::Zmax) = grid.n_elem-grid.n_ghost
@inline first_interior(grid::Grid, ::Zmin) = 1+grid.n_ghost

"""
    boundary(grid::Grid, ::ZBoundary)

Get the element index for edge data at the boundary
"""
@inline boundary(grid::Grid, ::Zmax) = grid.n_elem-grid.n_ghost
@inline boundary(grid::Grid, ::Zmin) = 1+grid.n_ghost

"""
    second_interior(grid::Grid, ::ZBoundary)

Get the second element index in the physical domain
"""
@inline second_interior(grid::Grid, ::Zmax) = first_interior(grid, Zmax())-1
@inline second_interior(grid::Grid, ::Zmin) = first_interior(grid, Zmin())+1

"""
    first_ghost(grid::Grid, ::ZBoundary)

Get the first element index in the ghost domain
"""
@inline first_ghost(grid::Grid, ::Zmax) = first_interior(grid, Zmax())+1
@inline first_ghost(grid::Grid, ::Zmin) = first_interior(grid, Zmin())-1

"""
    boundary_points(grid::Grid, ::ZBoundary)

Get the first element index in the ghost domain
"""
@inline boundary_points(grid::Grid, b::ZBoundary) = first_ghost(grid, b), first_interior(grid, b), second_interior(grid, b)

"""
    over_elems_ghost(grid::Grid)

Get the range of indexes to traverse only ghost grid elements
"""
@inline over_elems_ghost(grid::Grid) = setdiff(over_elems(grid), over_elems_real(grid))

"""
    over_elems_ghost(grid::Grid, ::ZBoundary)

Get the range of indexes to traverse only ghost grid elements below ``z_{min}`` and above ``z_{max}``.
"""
@inline over_elems_ghost(grid::Grid, ::Zmin) = 1 : grid.n_ghost
@inline over_elems_ghost(grid::Grid, ::Zmax) = grid.n_elem-grid.n_ghost+1 : grid.n_elem

Base.eltype(::Grid{T}) where T = T

include("GridOperators.jl")

end


# # Random particles in a circle

# The code [particles_in_circle.jl](particles_in_circle.jl) compares the scattered wave from one big circle, with the scattered wave from a circle filled with small particles.

# ```julia
using MultipleScattering

#You can also pick your own shape, an generate random particles inside it
#with a certain radius ands volume fraction
radius = 0.3
volfrac = 0.45
centre = [0.,0.]
big_radius = 3.0

particle_medium = Acoustic(2; ρ=0.0, c=0.0) # 2D particle with density ρ = 0.0 and soundspeed c = 0.0
particle_shape = Circle(radius)

circle = Sphere(centre, big_radius)

particles = random_particles(particle_medium, particle_shape; region_shape = circle, volume_fraction = volfrac, seed=1)

x = [-10.,0.] # position to receive the reflected wave
host_medium = Acoustic(2; ρ=1.0, c=1.0)
source =  plane_source(host_medium; position = x, direction = [1.0,0.])
simulation = FrequencySimulation(particles, source)
# ```
# The particles chosen are impenetrable, i.e. the wave is 100\% reflected. So this circle filled with scatterers should act like one big particle.
# ```julia
big_particle = Particle(particle_medium, circle)
big_particle_simulation = FrequencySimulation([big_particle], source)

#define a bounding box for plot
    bottomleft = [-10, -2*big_radius]
    topright = [big_radius, 2*big_radius]
    box = Box([bottomleft, topright])

using Plots
height = 300
#gr(size=(1.4*height,height))
# pyplot(leg=false, size=(1.4*height,height))

ω = 0.5
#plot(big_particle_simulation, ω; res=15, bounds = box);
#plot!(big_particle)

#savefig("plot_field_big.png")

#plot(simulation, ω; res=15, bounds = box);
#plot!(particles, linecolor = :green)

#savefig("plot_field.png")
# ```
# Resulting in the figures:

# ![The field with big particle](plot_field_big.png)
# ![The field with particles](plot_field.png)

# If we compare the response measured at the listener `[-10., 0.]`, they should be very similar:
# ```julia
#define angular frequency range
ωs = collect(LinRange(0.1,1.0,10))
result = run(simulation, x, ωs)
big_result = run(big_particle_simulation, x, ωs)

#plot(result, lab = "scattering from particles")
#plot!(big_result,
#    lab = "scattering from big particle",
#    title="Compare scattered wave from one big particle, \n and a circle filled with small particles")
# ```
# ![The response comparison](plot_response_compare.png)


function bfs(G, source)
    visited = Set(source);
    q = Deque{typeof(source)}()
    push!(q, source);
    while ! isempty(q)
        v = popfirst!(q);
        for n in neighbors(G, v)
            if !(n in visited)
                push!(q, n);
                push!(visited, n);
            end
        end
    end
    return visited
end

function connected_components(G)
    seen = []
    components = []
    for v in vertices(G)
        if !(v in seen)
            c = bfs(G, v)
            [push!(seen, elem) for elem in c]
            push!(components, c)
        end
    end
    return components
end


#module ResNetModule
export ResNet

import Knet
using Knet.Layers21: Conv, BatchNorm, Linear, Block, Add
using Knet.Ops21: relu, pool, mean
using Artifacts

"""
    ResNet(; nblocks=(2,2,2,2), block=ResNetBottleneck, groups=1, bottleneck=1, classes=1000)
    ResNet(name::String; pretrained=true)

Return a ResNet model based on keyword arguments or a name specifying a predefined
structure. Load pretrained weights if `pretrained=true`, randomly initialize otherwise. The
models satisfy the indexing interface, e.g. model[1] is the preprocessing layer and
model[1:end-1] omits the imagenet classification top.

A ResNet model `r` applied to an input of size `(width,height,channels=3,images)` returns a
`(classes,1)` vector of class scores. The model is a sequence of 6 functions: `r[1]` is the
input block, `r[2:5]` are four stages of residual blocks, and `r[6]` is the output
block. One can manipulate a ResNet model using the array interface, e.g. `r[1:5](img)` would
give the output of the last residual block, `pushfirst!(r, preprocess)` would add the
`preprocess` function to the beginning, etc.

The input and output blocks are the same for every ResNet model:

    Input:
        Conv(7×7, 3=>64, padding=3, stride=2, BatchNorm(), relu)
        Op(pool; window=3, stride=2, padding=1)

    Output:
        Op(pool; op=mean, window=typemax(Int))
        reshape2d
        Linear(classes, bias)

Stage i in 1:4 consists of `nblocks[i]` blocks of type `block` which can be
`ResNetBottleneck` (default) or `ResNetBasic`. `ResNetBasic` is a residual block with two
3×3 convolutions. `ResNetBottleneck` is a residual block with a 1×1 convolution that brings
the number of channels down, a 3×3 convolution (possibly grouped if `groups > 1`), and
another 1×1 convolution that brings the number of channels up. The keyword argument
`bottleneck` controls the ratio of the output channels to the intermediate channels. The
number of channels is doubled at each stage: blocks in the i'th stage have `2^(i+5)` output
channels for `ResNetBasic` and `2^(i+7)` output channels for `ResNetBottleneck`.  The
spatial resolution is halved at each stage: the first 3×3 convolution at each stage except
the first has `stride=2`.

The pretrained models come from
[torchvision](https://pytorch.org/vision/stable/models.html). A generic ResNet model can
handle any size input and any number of classes, but for pretrained models the training
settings should be replicated for best results: The number of classes is 1000 and the
expected size of the input is (224,224,3,N). The input images are resized to 256 pixels at
the short edge and center cropped to size (224,224). The pixel values are transformed from
the (0,1) range to:

    (rgb_pixel_0_1 .- [0.485, 0.456, 0.406]) ./ [0.229, 0.224, 0.225]

The class labels, as well as training, validation and test sets can be found at
[Kaggle](https://www.kaggle.com/c/imagenet-object-localization-challenge). Here are all the
pretrained models with name, settings, size in bytes, inference time compared to resnet50,
and top-1 validation accuracy.

    name              size   time  top1   ref1    settings
    ----              ----   ----  ----   ----    --------
    resnet18           45M   0.41  68.94  69.758  (nblocks=(2,2,2,2), block=ResNetBasic)
    resnet34           84M   0.57  72.76  73.314  (nblocks=(3,4,6,3), block=ResNetBasic)
    resnet50           98M   1.00  75.28  76.130  (nblocks=(3,4,6,3), bottleneck=4)
    resnet101         171M   1.58  76.83  77.374  (nblocks=(3,4,23,3), bottleneck=4)
    resnet152         231M   2.19  77.80  78.312  (nblocks=(3,8,36,3), bottleneck=4)
    wide_resnet50_2   264M   2.11  77.75  78.468  (nblocks=(3,4,6,3), bottleneck=2)
    wide_resnet101_2  485M   3.00  78.19  78.848  (nblocks=(3,4,23,3), bottleneck=2)
    resnext50_32x4d    96M   1.34  77.06  77.618  (nblocks=(3,4,6,3), groups=32, bottleneck=2)
    resnext101_32x8d  340M   3.45  78.69  79.312  (nblocks=(3,4,23,3), groups=32)

Note: The top1 accuracy figures are slightly different from the ones given by torchvision
(ref1), which is probably due to the differences in preprocessing, in particular `imresize`
in Julia gives different results compared to `Resize` in `torchvision.transforms`.

References:
* He, Kaiming et al. "Deep Residual Learning for Image Recognition." 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2016): 770-778.
* Zagoruyko, Sergey and Nikos Komodakis. "Wide Residual Networks." ArXiv abs/1605.07146 (2016)
* Xie, Saining et al. "Aggregated Residual Transformations for Deep Neural Networks." 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2017): 5987-5995.
* https://pytorch.org/vision/master/models.html

"""
function ResNet(; nblocks = (2,2,2,2), block = ResNetBottleneck, groups = 1, bottleneck = 1, classes = 1000)
    s = Block(; name="$block$nblocks")
    push!(s, Op(imagenet_preprocess; normalize=resnet_pt_normalize))
    push!(s, ResNetInput())
    x, y = 64, (block === ResNetBasic ? 64 : 256)
    for (stage, nblock) in enumerate(nblocks)
        if stage > 1; y *= 2; end
        b = y ÷ bottleneck
        blocks = Block(; name="Stage$stage")
        for iblock in 1:nblock
            stride = (stage > 1 && iblock == 1) ? 2 : 1
            push!(blocks, block(x, b, y; stride, groups))
            x = y
        end
        push!(s, blocks)
    end
    push!(s, ResNetOutput(y, classes))
    resnet_init(s)
end

function ResNetBottleneck(x, b, y; activation=relu, groups=1, padding=1, stride=1, o...)
    Add(
        Block(
            ConvBN(1, 1, x, b; activation),
            ConvBN(3, 3, b÷groups, b; activation, groups, padding, stride),
            ConvBN(1, 1, b, y),
        ),
        (x != y ? ConvBN(1, 1, x, y; stride) : identity);
        activation)
end

function ResNetBasic(x, b, y; stride=1, padding=1, activation=relu, o...)
    Add(
        Block(
            ConvBN(3, 3, x, b; activation, padding, stride),
            ConvBN(3, 3, b, y; padding),
        ),
        (x != y ? ConvBN(1, 1, x, y; stride) : identity);
        activation)
end

function ResNetInput()
    Block(
        ConvBN(7, 7, 3, 64; stride=2, padding=3, activation=relu),
        Op(pool; window=3, stride=2, padding=1);
        name = "Input"
    )
end

function ResNetOutput(xchannels, classes)
    Block(
        Op(pool; op=mean, window=typemax(Int)),
        reshape2d,
        Linear(xchannels, classes; binit=zeros);
        name = "Output"
    )
end

# Frequently used combo
ConvBN(x...; o...) = Conv(x...; o..., normalization=BatchNorm())

# Torchvision normalization
resnet_pt_normalize(x) = (x .- [0.485, 0.456, 0.406]) ./ [0.229, 0.224, 0.225]

# Run a single image so weights get initialized
resnet_init(m) = (m(Knet.atype(zeros(Float32,224,224,3,1))); m)

# Pretrained models from torchvision
resnetmodels = Dict{String,NamedTuple}(
    "resnet18" => (nblocks=(2,2,2,2), block=ResNetBasic),
    "resnet34" => (nblocks=(3,4,6,3), block=ResNetBasic),
    "resnet50" => (nblocks=(3,4,6,3), bottleneck=4),
    "resnet101" => (nblocks=(3,4,23,3), bottleneck=4),
    "resnet152" => (nblocks=(3,8,36,3), bottleneck=4),
    "wide_resnet50_2" => (nblocks=(3,4,6,3), bottleneck=2),
    "wide_resnet101_2" => (nblocks=(3,4,23,3), bottleneck=2),
    "resnext50_32x4d" => (nblocks=(3,4,6,3), groups=32, bottleneck=2),
    "resnext101_32x8d" => (nblocks=(3,4,23,3), groups=32),
)

function ResNet(s::String; pretrained=true)
    @assert haskey(resnetmodels, s)  "Please choose from known ResNet models:\n$(collect(keys(resnetmodels)))"
    kwargs = resnetmodels[s]
    model = ResNet(; kwargs...)
    pretrained && setweights!(model, joinpath(@artifact_str(s), "$s.jld2"))
    return model
end

#end # ResNetModule


using AccuracyAtTop: find_threshold, find_kth, find_quantile

s = collect(1:10)
y = s .>= 6

@testset "auxiliary threshold functions" begin
    @testset "maximum" begin
        @test find_threshold(All, findmax, y, s) == (10, 10)
        @test find_threshold(Neg, findmax, y, s) == (5, 5)
        @test find_threshold(Pos, findmax, y, s) == (10, 10)

        @test find_threshold(Maximum(All), y, s) == (10, 10)
        @test find_threshold(Maximum(Neg), y, s) == (5, 5)
        @test find_threshold(Maximum(Pos), y, s) == (10, 10)
    end

    @testset "minimum" begin
        @test find_threshold(All, findmin, y, s) == (1, 1)
        @test find_threshold(Neg, findmin, y, s) == (1, 1)
        @test find_threshold(Pos, findmin, y, s) == (6, 6)

        @test find_threshold(Minimum(All), y, s) == (1, 1)
        @test find_threshold(Minimum(Neg), y, s) == (1, 1)
        @test find_threshold(Minimum(Pos), y, s) == (6, 6)
    end

    @testset "$(k) all samples" for k in 1:10
        @test find_kth(s, k, false) == (k, k)
        @test find_kth(s, k, true) == (11 - k, 11 - k)

        @test find_threshold(All, find_kth, y, s, k, false) == (k, k)
        @test find_threshold(All, find_kth, y, s, k, true) == (11 - k, 11 - k)

        @test find_threshold(Kth(k, All; rev = false), y, s) == (k, k)
        @test find_threshold(Kth(k, All; rev = true), y, s) == (11 - k, 11 - k)
    end

    @testset "$(k) pos/neg samples" for k in 1:5
        @test find_threshold(Neg, find_kth, y, s, k, false) == (k, k)
        @test find_threshold(Neg, find_kth, y, s, k, true) == (6 - k, 6 - k)
        @test find_threshold(Kth(k, Neg; rev = false), y, s) == (k, k)
        @test find_threshold(Kth(k, Neg; rev = true), y, s) == (6 - k, 6 - k)

        @test find_threshold(Pos, find_kth, y, s, k, false) == (5 + k, 5 + k)
        @test find_threshold(Pos, find_kth, y, s, k, true) == (11 - k, 11 - k)
        @test find_threshold(Kth(k, Pos; rev = false), y, s) == (5 + k, 5 + k)
        @test find_threshold(Kth(k, Pos; rev = true), y, s) == (11 - k, 11 - k)
    end

    @testset "$(τ)-quantile all samples" for (k, τ) in zip(vcat(1, 1:10), 0:0.1:1)
        @test find_quantile(s, τ, false) == (k, k)
        @test find_threshold(All, find_quantile, y, s, τ, false) == (k, k)
        @test find_threshold(Quantile(τ, All; rev = false), y, s) == (k, k)
    end

    @testset "$(τ)-quantile all samples reversed" for (k, τ) in zip(vcat(10:-1:1, 1), 0:0.1:1)
        @test find_quantile(s, τ, true) == (k, k)
        @test find_threshold(All, find_quantile, y, s, τ, true) == (k, k)
        @test find_threshold(Quantile(τ, All; rev = true), y, s) == (k, k)
    end

    @testset "$(τ)-quantile pos/neg samples" for (k, τ) in zip(vcat(1, 1:5), 0:0.2:1)
        @test find_threshold(Neg, find_quantile, y, s, τ, false) == (k, k)
        @test find_threshold(Quantile(τ, Neg; rev = false), y, s) == (k, k)

        @test find_threshold(Pos, find_quantile, y, s, τ, false) == (k + 5, k + 5)
        @test find_threshold(Quantile(τ, Pos; rev = false), y, s) == (k + 5, k + 5)
    end

    @testset "$(τ)-quantile pos/neg samples reversed" for (k, τ) in zip(vcat(5:-1:1, 1), 0:0.2:1)
        @test find_threshold(Neg, find_quantile, y, s, τ, true) == (k, k)
        @test find_threshold(Quantile(τ, Neg; rev = true), y, s) == (k, k)

        @test find_threshold(Pos, find_quantile, y, s, τ, true) == (k + 5, k + 5)
        @test find_threshold(Quantile(τ, Pos; rev = true), y, s) == (k + 5, k + 5)
    end
end

s = rand(4, 100)
y = rand(Bool, 4, 100)

function test_indices(tp, y, s)
    t, ind = find_threshold(tp, y, s)

    @testset "$tp for $i-th row" for i in 1:size(s, 1)
        @test t[i] == s[i, ind[i]]
    end
end

@testset "thresholds rrule" begin
    @testset "$tp" for tp in Maximum.([All, Neg, Pos])
        test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
        test_indices(tp, y, s)
    end

    @testset "$tp" for tp in Minimum.([All, Neg, Pos])
        test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
        test_indices(tp, y, s)
    end

    @testset "$K" for K in [1, 2, 5, 10, 15]
        @testset "$rev" for rev in [true, false]
            @testset "$tp" for tp in Kth.(K, [All, Neg, Pos]; rev)
                test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
                test_indices(tp, y, s)
            end
        end
    end

    @testset "$τ" for τ in [0.05, 0.1, 0.15, 0.2, 0.5]
        @testset "$rev" for rev in [true, false]
            @testset "$tp" for tp in Quantile.(τ, [All, Neg, Pos]; rev)
                test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
                test_indices(tp, y, s)
            end
        end
    end

    @testset "$τ" for τ in [0.05, 0.1, 0.15, 0.2, 0.5]
        sampler() = τ
        @testset "$rev" for rev in [true, false]
            @testset "$tp" for tp in SampledQuantile.(sampler, [All, Neg, Pos]; rev)
                test_rrule(threshold, tp ⊢ NoTangent(), y ⊢ NoTangent(), s)
                test_indices(tp, y, s)
            end
        end
    end
end


struct MatSpread
    mat::Vector{Vector{Float64}}
    spread::Float64
end

struct MorrisSensitivity
    means
    variances
    elementary_effects
end

function generate_design_matrix(p_range,p_steps;k = 10)
    ps = [range(p_range[i][1], stop=p_range[i][2], length=p_steps[i]) for i in 1:length(p_range)]
    indices = [rand(1:i) for i in p_steps]
    all_idxs = Vector{typeof(indices)}(undef,k)

    for i in 1:k
        flag = 0
        j = rand(1:length(p_range))
        indices[j] += (rand() < 0.5 ? -1 : 1)
        if indices[j] > p_steps[j]
            indices[j] -= 2
        elseif indices[j] < 1.0
            indices[j] += 2
        end
        all_idxs[i] = copy(indices)
    end

    B = Array{Array{Float64}}(undef,k)
    for j in 1:k
        cur_p = [ps[u][(all_idxs[j][u])] for u in 1:length(p_range)]
        B[j] = cur_p
    end
    B
end

function calculate_spread(matrix)
    spread = 0.0
    for i in 2:length(matrix)
        spread += sqrt(sum(abs2.(matrix[i] - matrix[i-1])))
    end
    spread
end

function sample_matrices(p_range,p_steps;len_trajectory=10,num_trajectory=10,total_num_trajectory=5*num_trajectory)
    matrix_array = []
    if total_num_trajectory<num_trajectory
        error("total_num_trajectory should be greater than num_trajectory preferably atleast 3-4 times higher")
    end
    for i in 1:total_num_trajectory
        mat = generate_design_matrix(p_range,p_steps;k=len_trajectory)
        spread = calculate_spread(mat)
        push!(matrix_array,MatSpread(mat,spread))
    end
    sort!(matrix_array,by = x -> x.spread,rev=true)
    matrices = [i.mat for i in matrix_array[1:num_trajectory]]
    matrices
end

function morris_sensitivity(prob::DiffEqBase.DEProblem,alg,t,p_range,p_steps;kwargs...)
    f = function (p)
      prob1 = remake(prob;p=p)
      Array(solve(prob1,alg;saveat=t))
    end
    morris_sensitivity(f,p_range,p_steps;kwargs...)
end

function morris_sensitivity(f,p_range,p_steps;relative_scale=false,kwargs...)
    design_matrices = sample_matrices(p_range,p_steps;kwargs...)
    y1 = f(design_matrices[1][1])
    effects = [[]]
    for i in design_matrices
        y1 = f(i[1])
        for j in 1:length(i)-1
            y2 = y1
            del = i[j+1] - i[j]
            change_index = 0
            for k in 1:length(del)
                if abs(del[k]) > 0
                    change_index = k
                    break
                end
            end
            del = sum(del)
            y1 = f(i[j+1])
            if relative_scale == false
                elem_effect = @. abs((y1-y2)/(del))
            else
                if del > 0
                    elem_effect = @. abs((y1-y2)/(y2*del))
                else
                    elem_effect = @. abs((y1-y2)/(y1*del))
                end
            end
            if length(effects) >= change_index && change_index > 0
                push!(effects[change_index],elem_effect)
            elseif change_index > 0
                push!(effects,[elem_effect])
            end
        end
    end
    means = []
    variances = []
    for k in effects
        push!(means,mean(convert(Array, VectorOfArray(k)),dims=2))
        push!(variances,var(convert(Array, VectorOfArray(k)),dims=2))
    end
    MorrisSensitivity(means,variances,effects)
end


"""
    Mlm(B, varB, sigma, data, weights, targetType, lambda)

Type for storing the results of an mlm model fit.

"""
mutable struct Mlm

    # Coefficient estimates
    B::Array{Float64,2}
    # Coefficient variance estimates
    varB::Array{Float64,2}
    # Estimated variance of errors
    sigma::Array{Float64,2}

    # Response and predictor matrices
    data::RawData

    # Column weights for `Y`, or `nothing`
    weights
    # String indicating target type to shrink the variance toward, or `nothing`
    targetType
    # Estimated shrinkage coefficient
    lambda::Float64
end

"""
    mlm_fit(data, weights, targetType)

Matrix linear model using least squares method. Optionally incorporates
shrinkage of the variance of the errors.

# Arguments

- data = RawData object
- weights = `nothing`
- targetType = string indicating the target type toward which to shrink the
  error variance, or `nothing`. If the former, acceptable inputs are "A", "B",
  "C", and "D".
    - "A": Target is identity matrix
    - "B": Target is diagonal matrix with constant diagonal
    - "C": Target is has same diagonal element, and same off-diagonal element
    - "D": Target is diagonal matrix with unequal entries

# Value

An Mlm object

"""
function mlm_fit(data::RawData, weights::Nothing, targetType)

    # Calculate and store transpose(X)*X
    XTX = transpose(get_X(data))*get_X(data)
    # Calculate and store transpose(Z)*Z
    ZTZ = transpose(get_Z(data))*get_Z(data)

    # Estimate MLM coefficients
    B = calc_coeffs(get_X(data), get_Y(data), get_Z(data), XTX, ZTZ)

    # Calculate residuals
    resid = calc_resid(get_X(data), get_Y(data), get_Z(data), B)

    # Estimate variance of errors, optionally with variance shrinkage
    sigma, lambda = calc_sigma(resid, targetType)

    # Estimate variance of coefficient estimates
    varB = calc_var(get_X(data), get_Z(data), XTX, ZTZ, sigma)

    # Return Mlm object with estimates for coefficients, variances, and sigma
    return Mlm(B, varB, sigma, data, weights, targetType, lambda)
end

"""
    mlm_fit(data, weights, targetType)

Matrix linear model using column weighted least squares method. Optionally
incorporates shrinkage of the variance of the errors.

# Arguments

- data = RawData object
- weights = 1d array of floats to use as column weights for `Y`. Must be the
  same length as the number of columns of `Y`.
- targetType = string indicating the target type toward which to shrink the
  error variance, or `nothing`. If the former, acceptable inputs are "A", "B",
  "C", and "D".
    - "A": Target is identity matrix
    - "B": Target is diagonal matrix with constant diagonal
    - "C": Target is has same diagonal element, and same off-diagonal element
    - "D": Target is diagonal matrix with unequal entries

# Value

An Mlm object

"""
function mlm_fit(data::RawData, weights::Array{Float64,1}, targetType)

    # Throw an error if the length of weights is incorrect
    if (length(weights) != size(get_Y(data),2))
        error("Weights must be same length as number of columns in Y.")
    end

    # Diagonal matrix with weights on diagonal.
    W = diagm(weights)

    # Calculate and store WZ product
    WZ = W*get_Z(data)

    # Calculate and store transpose(X)*X
    XTX = transpose(get_X(data))*get_X(data)
    # Calculate and store transpose(Z)*W*Z
    ZTWZ = transpose(get_Z(data))*WZ

    # Estimate MLM coefficients
    B = calc_coeffs(get_X(data), get_Y(data), WZ, XTX, ZTWZ)

    # Calculate residuals
    resid = calc_resid(get_X(data), get_Y(data), get_Z(data), B)

    # Estimate variance of errors, optionally with variance shrinkage
    sigma, lambda = calc_sigma(resid, targetType)

    # Estimate variance of coefficient estimates
    varB = calc_var(get_X(data), WZ, XTX, ZTWZ, sigma)

    # Return Mlm object with estimates for coefficients, variances, and sigma
    return Mlm(B, varB, sigma, data, weights, targetType, lambda)
end

"""
    mlm(data; hasXIntercept, hasZIntercept, weights, targetType)

Matrix linear model using least squares method. Column weighted least squares
and shrinkage of the variance of the errors are options.

# Arguments

- data = RawData object

# Keyword arguments

- hasXIntercept = boolean flag indicating whether or not to include an `X`
  intercept (row main effects). Defaults to `true`.
- hasZIntercept = boolean flag indicating whether or not to include a `Z`
  intercept (column main effects). Defaults to `true`.
- weights = 1d array of floats to use as column weights for `Y`, or `nothing`.
  If the former, must be the same length as the number of columns of `Y`.
  Defaults to `nothing`.
- targetType = string indicating the target type toward which to shrink the
  error variance, or `nothing`. If the former, acceptable inputs are "A", "B",
  "C", and "D". Defaults to `nothing`.
    - "A": Target is identity matrix
    - "B": Target is diagonal matrix with constant diagonal
    - "C": Target is has same diagonal element, and same off-diagonal element
    - "D": Target is diagonal matrix with unequal entries

# Value

An Mlm object

"""
function mlm(data::RawData; hasXIntercept::Bool=true, hasZIntercept::Bool=true,
             weights=nothing, targetType=nothing, kwargs...)

    #= Deprecation:check for previous version keyword arguments
    ---------------------------------------------------------------------------------------------=#
    if haskey(kwargs, :isXIntercept)
        @warn "Keyword arguments `isXIntercept` and `isZIntercept` are deprecated, use
                    `hasXIntercept` and `hasZIntercept` instead."

        hasXIntercept = values(kwargs).isXIntercept
    end

    if haskey(kwargs, :isZIntercept)
        @warn "Keyword arguments `isXIntercept` and `isZIntercept` are deprecated, use
                    `hasXIntercept` and `hasZIntercept` instead."

        hasZIntercept = values(kwargs).isZIntercept
    end
    #-----------------------------------------------------------------------------------------------

    # Add X and Z intercepts if necessary
    if hasXIntercept==true && data.predictors.hasXIntercept==false
        data.predictors.X = add_intercept(data.predictors.X)
        data.predictors.hasXIntercept = true
        data.p = data.p + 1
    end
    if hasZIntercept==true && data.predictors.hasZIntercept==false
        data.predictors.Z = add_intercept(data.predictors.Z)
        data.predictors.hasZIntercept = true
        data.q = data.q + 1
    end

    # Remove X and Z intercepts in new predictors if necessary
    if hasXIntercept==false && data.predictors.hasXIntercept==true
        data.predictors.X = remove_intercept(data.predictors.X)
        data.predictors.hasXIntercept = false
        data.p = data.p - 1
    end
    if hasZIntercept==false && data.predictors.hasZIntercept==true
        data.predictors.Z = remove_intercept(data.predictors.Z)
        data.predictors.hasZIntercept = false
        data.q = data.q - 1
    end

    if (typeof(targetType) != Nothing) & !(targetType in ["A", "B", "C", "D"])
        println("Unrecognizable targetType will be ignored and no variance shrinkage will be performed.")
        targetType = nothing
    end

    # Run matrix linear models
    return mlm_fit(data, weights, targetType)
end

"""
    t_stat(MLM, isMainEff)

Calculates t-statistics of an Mlm object

# Arguments

- MLM = Mlm object
- isMainEff = boolean flag indicating whether or not to include t-statistics
  for the main effects

# Value

2d array of floats

"""
function t_stat(MLM::Mlm, isMainEff::Bool=false)

    # Cases when not including main effects
    if isMainEff== false
        # If X and Z intercepts were both included
        if (MLM.data.predictors.hasXIntercept==true) &&
           (MLM.data.predictors.hasZIntercept==true)
            return MLM.B[2:end, 2:end]./sqrt.(MLM.varB[2:end, 2:end])

        # If only X intercept was included
        elseif MLM.data.predictors.hasXIntercept==true
            return MLM.B[2:end, :]./sqrt.(MLM.varB[2:end, :])

        # If only Z intercept was included
        elseif  MLM.data.predictors.hasZIntercept==true
            return MLM.B[:, 2:end]./sqrt.(MLM.varB[:, 2:end])

        end
    end

    # Case when including main effects
    return MLM.B ./ sqrt.(MLM.varB)
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

module Math

export sin, cos, sincos, tan, sinh, cosh, tanh, asin, acos, atan,
       asinh, acosh, atanh, sec, csc, cot, asec, acsc, acot,
       sech, csch, coth, asech, acsch, acoth,
       sinpi, cospi, sinc, cosc,
       cosd, cotd, cscd, secd, sind, tand, sincosd,
       acosd, acotd, acscd, asecd, asind, atand,
       rad2deg, deg2rad,
       log, log2, log10, log1p, exponent, exp, exp2, exp10, expm1,
       cbrt, sqrt, significand,
       hypot, max, min, minmax, ldexp, frexp,
       clamp, clamp!, modf, ^, mod2pi, rem2pi,
       @evalpoly, evalpoly

import .Base: log, exp, sin, cos, tan, sinh, cosh, tanh, asin,
             acos, atan, asinh, acosh, atanh, sqrt, log2, log10,
             max, min, minmax, ^, exp2, muladd, rem,
             exp10, expm1, log1p

using .Base: sign_mask, exponent_mask, exponent_one,
            exponent_half, uinttype, significand_mask,
            significand_bits, exponent_bits, exponent_bias,
            exponent_max, exponent_raw_max

using Core.Intrinsics: sqrt_llvm

using .Base: IEEEFloat

@noinline function throw_complex_domainerror(f::Symbol, x)
    throw(DomainError(x, string("$f will only return a complex result if called with a ",
                                "complex argument. Try $f(Complex(x)).")))
end
@noinline function throw_exp_domainerror(x)
    throw(DomainError(x, string("Exponentiation yielding a complex result requires a ",
                                "complex argument.\nReplace x^y with (x+0im)^y, ",
                                "Complex(x)^y, or similar.")))
end

# non-type specific math functions

"""
    clamp(x, lo, hi)

Return `x` if `lo <= x <= hi`. If `x > hi`, return `hi`. If `x < lo`, return `lo`. Arguments
are promoted to a common type.

# Examples
```jldoctest
julia> clamp.([pi, 1.0, big(10.)], 2., 9.)
3-element Array{BigFloat,1}:
 3.141592653589793238462643383279502884197169399375105820974944592307816406286198
 2.0
 9.0

julia> clamp.([11,8,5],10,6) # an example where lo > hi
3-element Array{Int64,1}:
  6
  6
 10
```
"""
clamp(x::X, lo::L, hi::H) where {X,L,H} =
    ifelse(x > hi, convert(promote_type(X,L,H), hi),
           ifelse(x < lo,
                  convert(promote_type(X,L,H), lo),
                  convert(promote_type(X,L,H), x)))

"""
    clamp(x, T)::T

Clamp `x` between `typemin(T)` and `typemax(T)` and convert the result to type `T`.

# Examples
```jldoctest
julia> clamp(200, Int8)
127
julia> clamp(-200, Int8)
-128
```
"""
clamp(x, ::Type{T}) where {T<:Integer} = clamp(x, typemin(T), typemax(T)) % T

"""
    clamp!(array::AbstractArray, lo, hi)

Restrict values in `array` to the specified range, in-place.
See also [`clamp`](@ref).
"""
function clamp!(x::AbstractArray, lo, hi)
    @inbounds for i in eachindex(x)
        x[i] = clamp(x[i], lo, hi)
    end
    x
end

"""
    evalpoly(x, p)

Evaluate the polynomial ``\\sum_k x^{k-1} p[k]`` for the coefficients `p[1]`, `p[2]`, ...;
that is, the coefficients are given in ascending order by power of `x`.
Loops are unrolled at compile time if the number of coefficients is statically known, i.e.
when `p` is a `Tuple`.
This function generates efficient code using Horner's method if `x` is real, or using
a Goertzel-like [^DK62] algorithm if `x` is complex.

[^DK62]: Donald Knuth, Art of Computer Programming, Volume 2: Seminumerical Algorithms, Sec. 4.6.4.

!!! compat "Julia 1.4"
    This function requires Julia 1.4 or later.

# Example
```jldoctest
julia> evalpoly(2, (1, 2, 3))
17
```
"""
function evalpoly(x, p::Tuple)
    if @generated
        N = length(p.parameters)
        ex = :(p[end])
        for i in N-1:-1:1
            ex = :(muladd(x, $ex, p[$i]))
        end
        ex
    else
        _evalpoly(x, p)
    end
end

evalpoly(x, p::AbstractVector) = _evalpoly(x, p)

function _evalpoly(x, p)
    N = length(p)
    ex = p[end]
    for i in N-1:-1:1
        ex = muladd(x, ex, p[i])
    end
    ex
end

function evalpoly(z::Complex, p::Tuple)
    if @generated
        N = length(p.parameters)
        a = :(p[end])
        b = :(p[end-1])
        as = []
        for i in N-2:-1:1
            ai = Symbol("a", i)
            push!(as, :($ai = $a))
            a = :(muladd(r, $ai, $b))
            b = :(muladd(-s, $ai, p[$i]))
        end
        ai = :a0
        push!(as, :($ai = $a))
        C = Expr(:block,
                 :(x = real(z)),
                 :(y = imag(z)),
                 :(r = x + x),
                 :(s = muladd(x, x, y*y)),
                 as...,
                 :(muladd($ai, z, $b)))
    else
        _evalpoly(z, p)
    end
end
evalpoly(z::Complex, p::Tuple{<:Any}) = p[1]

evalpoly(z::Complex, p::AbstractVector) = _evalpoly(z, p)

function _evalpoly(z::Complex, p)
    length(p) == 1 && return p[1]
    N = length(p)
    a = p[end]
    b = p[end-1]

    x = real(z)
    y = imag(z)
    r = 2x
    s = muladd(x, x, y*y)
    for i in N-2:-1:1
        ai = a
        a = muladd(r, ai, b)
        b = muladd(-s, ai, p[i])
    end
    ai = a
    muladd(ai, z, b)
end

"""
    @horner(x, p...)

Evaluate `p[1] + x * (p[2] + x * (....))`, i.e. a polynomial via Horner's rule.
"""
macro horner(x, p...)
     xesc, pesc = esc(x), esc.(p)
    :(invoke(evalpoly, Tuple{Any, Tuple}, $xesc, ($(pesc...),)))
end

# Evaluate p[1] + z*p[2] + z^2*p[3] + ... + z^(n-1)*p[n].  This uses
# Horner's method if z is real, but for complex z it uses a more
# efficient algorithm described in Knuth, TAOCP vol. 2, section 4.6.4,
# equation (3).

"""
    @evalpoly(z, c...)

Evaluate the polynomial ``\\sum_k z^{k-1} c[k]`` for the coefficients `c[1]`, `c[2]`, ...;
that is, the coefficients are given in ascending order by power of `z`.  This macro expands
to efficient inline code that uses either Horner's method or, for complex `z`, a more
efficient Goertzel-like algorithm.

# Examples
```jldoctest
julia> @evalpoly(3, 1, 0, 1)
10

julia> @evalpoly(2, 1, 0, 1)
5

julia> @evalpoly(2, 1, 1, 1)
7
```
"""
macro evalpoly(z, p...)
    zesc, pesc = esc(z), esc.(p)
    :(evalpoly($zesc, ($(pesc...),)))
end

"""
    rad2deg(x)

Convert `x` from radians to degrees.

# Examples
```jldoctest
julia> rad2deg(pi)
180.0
```
"""
rad2deg(z::AbstractFloat) = z * (180 / oftype(z, pi))

"""
    deg2rad(x)

Convert `x` from degrees to radians.

# Examples
```jldoctest
julia> deg2rad(90)
1.5707963267948966
```
"""
deg2rad(z::AbstractFloat) = z * (oftype(z, pi) / 180)
rad2deg(z::Real) = rad2deg(float(z))
deg2rad(z::Real) = deg2rad(float(z))
rad2deg(z::Number) = (z/pi)*180
deg2rad(z::Number) = (z*pi)/180

log(b::T, x::T) where {T<:Number} = log(x)/log(b)

"""
    log(b,x)

Compute the base `b` logarithm of `x`. Throws [`DomainError`](@ref) for negative
[`Real`](@ref) arguments.

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> log(4,8)
1.5

julia> log(4,2)
0.5

julia> log(-2, 3)
ERROR: DomainError with -2.0:
log will only return a complex result if called with a complex argument. Try log(Complex(x)).
Stacktrace:
 [1] throw_complex_domainerror(::Symbol, ::Float64) at ./math.jl:31
[...]

julia> log(2, -3)
ERROR: DomainError with -3.0:
log will only return a complex result if called with a complex argument. Try log(Complex(x)).
Stacktrace:
 [1] throw_complex_domainerror(::Symbol, ::Float64) at ./math.jl:31
[...]
```

!!! note
    If `b` is a power of 2 or 10, [`log2`](@ref) or [`log10`](@ref) should be used, as these will
    typically be faster and more accurate. For example,

    ```jldoctest
    julia> log(100,1000000)
    2.9999999999999996

    julia> log10(1000000)/2
    3.0
    ```
"""
log(b::Number, x::Number) = log(promote(b,x)...)

# type specific math functions

const libm = Base.libm_name

# functions with no domain error
"""
    sinh(x)

Compute hyperbolic sine of `x`.
"""
sinh(x::Number)

"""
    cosh(x)

Compute hyperbolic cosine of `x`.
"""
cosh(x::Number)

"""
    tanh(x)

Compute hyperbolic tangent of `x`.
"""
tanh(x::Number)

"""
    atan(y)
    atan(y, x)

Compute the inverse tangent of `y` or `y/x`, respectively.

For one argument, this is the angle in radians between the positive *x*-axis and the point
(1, *y*), returning a value in the interval ``[-\\pi/2, \\pi/2]``.

For two arguments, this is the angle in radians between the positive *x*-axis and the
point (*x*, *y*), returning a value in the interval ``[-\\pi, \\pi]``. This corresponds to a
standard [`atan2`](https://en.wikipedia.org/wiki/Atan2) function.
"""
atan(x::Number)

"""
    asinh(x)

Compute the inverse hyperbolic sine of `x`.
"""
asinh(x::Number)

"""
    expm1(x)

Accurately compute ``e^x-1``.
"""
expm1(x)
for f in (:exp2, :expm1)
    @eval begin
        ($f)(x::Float64) = ccall(($(string(f)),libm), Float64, (Float64,), x)
        ($f)(x::Float32) = ccall(($(string(f,"f")),libm), Float32, (Float32,), x)
        ($f)(x::Real) = ($f)(float(x))
    end
end

"""
    exp2(x)

Compute the base 2 exponential of `x`, in other words ``2^x``.

# Examples
```jldoctest
julia> exp2(5)
32.0
```
"""
exp2(x::AbstractFloat) = 2^x

"""
    exp10(x)

Compute the base 10 exponential of `x`, in other words ``10^x``.

# Examples
```jldoctest
julia> exp10(2)
100.0
```
"""
exp10(x::AbstractFloat) = 10^x

for f in (:sinh, :cosh, :tanh, :atan, :asinh, :exp, :expm1)
    @eval ($f)(x::AbstractFloat) = error("not implemented for ", typeof(x))
end

# functions with special cases for integer arguments
@inline function exp2(x::Base.BitInteger)
    if x > 1023
        Inf64
    elseif x <= -1023
        # if -1073 < x <= -1023 then Result will be a subnormal number
        # Hex literal with padding must be used to work on 32bit machine
        reinterpret(Float64, 0x0000_0000_0000_0001 << ((x + 1074) % UInt))
    else
        # We will cast everything to Int64 to avoid errors in case of Int128
        # If x is a Int128, and is outside the range of Int64, then it is not -1023<x<=1023
        reinterpret(Float64, (exponent_bias(Float64) + (x % Int64)) << (significand_bits(Float64) % UInt))
    end
end

# utility for converting NaN return to DomainError
# the branch in nan_dom_err prevents its callers from inlining, so be sure to force it
# until the heuristics can be improved
@inline nan_dom_err(out, x) = isnan(out) & !isnan(x) ? throw(DomainError(x, "NaN result for non-NaN input.")) : out

# functions that return NaN on non-NaN argument for domain error
"""
    sin(x)

Compute sine of `x`, where `x` is in radians.
"""
sin(x::Number)

"""
    cos(x)

Compute cosine of `x`, where `x` is in radians.
"""
cos(x::Number)

"""
    tan(x)

Compute tangent of `x`, where `x` is in radians.
"""
tan(x::Number)

"""
    asin(x)

Compute the inverse sine of `x`, where the output is in radians.
"""
asin(x::Number)

"""
    acos(x)

Compute the inverse cosine of `x`, where the output is in radians
"""
acos(x::Number)

"""
    acosh(x)

Compute the inverse hyperbolic cosine of `x`.
"""
acosh(x::Number)

"""
    atanh(x)

Compute the inverse hyperbolic tangent of `x`.
"""
atanh(x::Number)

"""
    log(x)

Compute the natural logarithm of `x`. Throws [`DomainError`](@ref) for negative
[`Real`](@ref) arguments. Use complex negative arguments to obtain complex results.

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> log(2)
0.6931471805599453

julia> log(-3)
ERROR: DomainError with -3.0:
log will only return a complex result if called with a complex argument. Try log(Complex(x)).
Stacktrace:
 [1] throw_complex_domainerror(::Symbol, ::Float64) at ./math.jl:31
[...]
```
"""
log(x::Number)

"""
    log2(x)

Compute the logarithm of `x` to base 2. Throws [`DomainError`](@ref) for negative
[`Real`](@ref) arguments.

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> log2(4)
2.0

julia> log2(10)
3.321928094887362

julia> log2(-2)
ERROR: DomainError with -2.0:
NaN result for non-NaN input.
Stacktrace:
 [1] nan_dom_err at ./math.jl:325 [inlined]
[...]
```
"""
log2(x)

"""
    log10(x)

Compute the logarithm of `x` to base 10.
Throws [`DomainError`](@ref) for negative [`Real`](@ref) arguments.

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> log10(100)
2.0

julia> log10(2)
0.3010299956639812

julia> log10(-2)
ERROR: DomainError with -2.0:
NaN result for non-NaN input.
Stacktrace:
 [1] nan_dom_err at ./math.jl:325 [inlined]
[...]
```
"""
log10(x)

"""
    log1p(x)

Accurate natural logarithm of `1+x`. Throws [`DomainError`](@ref) for [`Real`](@ref)
arguments less than -1.

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> log1p(-0.5)
-0.6931471805599453

julia> log1p(0)
0.0

julia> log1p(-2)
ERROR: DomainError with -2.0:
log1p will only return a complex result if called with a complex argument. Try log1p(Complex(x)).
Stacktrace:
 [1] throw_complex_domainerror(::Symbol, ::Float64) at ./math.jl:31
[...]
```
"""
log1p(x)
for f in (:log2, :log10)
    @eval begin
        @inline ($f)(x::Float64) = nan_dom_err(ccall(($(string(f)), libm), Float64, (Float64,), x), x)
        @inline ($f)(x::Float32) = nan_dom_err(ccall(($(string(f, "f")), libm), Float32, (Float32,), x), x)
        @inline ($f)(x::Real) = ($f)(float(x))
    end
end

@inline function sqrt(x::Union{Float32,Float64})
    x < zero(x) && throw_complex_domainerror(:sqrt, x)
    sqrt_llvm(x)
end

"""
    sqrt(x)

Return ``\\sqrt{x}``. Throws [`DomainError`](@ref) for negative [`Real`](@ref) arguments.
Use complex negative arguments instead. The prefix operator `√` is equivalent to `sqrt`.

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> sqrt(big(81))
9.0

julia> sqrt(big(-81))
ERROR: DomainError with -81.0:
NaN result for non-NaN input.
Stacktrace:
 [1] sqrt(::BigFloat) at ./mpfr.jl:501
[...]

julia> sqrt(big(complex(-81)))
0.0 + 9.0im
```
"""
sqrt(x::Real) = sqrt(float(x))

"""
    hypot(x, y)

Compute the hypotenuse ``\\sqrt{|x|^2+|y|^2}`` avoiding overflow and underflow.

This code is an implementation of the algorithm described in:
An Improved Algorithm for `hypot(a,b)`
by Carlos F. Borges
The article is available online at ArXiv at the link
  https://arxiv.org/abs/1904.09481

# Examples
```jldoctest; filter = r"Stacktrace:(\\n \\[[0-9]+\\].*)*"
julia> a = Int64(10)^10;

julia> hypot(a, a)
1.4142135623730951e10

julia> √(a^2 + a^2) # a^2 overflows
ERROR: DomainError with -2.914184810805068e18:
sqrt will only return a complex result if called with a complex argument. Try sqrt(Complex(x)).
Stacktrace:
[...]

julia> hypot(3, 4im)
5.0
```
"""
hypot(x::Number, y::Number) = hypot(promote(x, y)...)
hypot(x::Complex, y::Complex) = hypot(abs(x), abs(y))
hypot(x::T, y::T) where {T<:Real} = hypot(float(x), float(y))
function hypot(x::T, y::T) where {T<:Number}
    if !iszero(x)
        z = y/x
        z2 = z*z

        abs(x) * sqrt(oneunit(z2) + z2)
    else
        abs(y)
    end
end

function hypot(x::T, y::T) where T<:AbstractFloat
    # Return Inf if either or both inputs is Inf (Compliance with IEEE754)
    if isinf(x) || isinf(y)
        return T(Inf)
    end

    # Order the operands
    ax,ay = abs(x), abs(y)
    if ay > ax
        ax,ay = ay,ax
    end

    # Widely varying operands
    if ay <= ax*sqrt(eps(T)/2)  #Note: This also gets ay == 0
        return ax
    end

    # Operands do not vary widely
    scale = eps(sqrt(floatmin(T)))  #Rescaling constant
    if ax > sqrt(floatmax(T)/2)
        ax = ax*scale
        ay = ay*scale
        scale = inv(scale)
    elseif ay < sqrt(floatmin(T))
        ax = ax/scale
        ay = ay/scale
    else
        scale = one(scale)
    end
    h = sqrt(muladd(ax,ax,ay*ay))
    # This branch is correctly rounded but requires a native hardware fma.
    if Base.Math.FMA_NATIVE
        hsquared = h*h
        axsquared = ax*ax
        h -= (fma(-ay,ay,hsquared-axsquared) + fma(h,h,-hsquared) - fma(ax,ax,-axsquared))/(2*h)
    # This branch is within one ulp of correctly rounded.
    else
        if h <= 2*ay
            delta = h-ay
            h -= muladd(delta,delta-2*(ax-ay),ax*(2*delta - ax))/(2*h)
        else
            delta = h-ax
            h -= muladd(delta,delta,muladd(ay,(4*delta-ay),2*delta*(ax-2*ay)))/(2*h)
        end
    end
    return h*scale
end

"""
    hypot(x...)

Compute the hypotenuse ``\\sqrt{\\sum |x_i|^2}`` avoiding overflow and underflow.

# Examples
```jldoctest
julia> hypot(-5.7)
5.7

julia> hypot(3, 4im, 12.0)
13.0
```
"""
hypot(x::Number...) = sqrt(sum(abs2(y) for y in x))

atan(y::Real, x::Real) = atan(promote(float(y),float(x))...)
atan(y::T, x::T) where {T<:AbstractFloat} = Base.no_op_err("atan", T)

max(x::T, y::T) where {T<:AbstractFloat} = ifelse((y > x) | (signbit(y) < signbit(x)),
                                    ifelse(isnan(x), x, y), ifelse(isnan(y), y, x))

min(x::T, y::T) where {T<:AbstractFloat} = ifelse((y < x) | (signbit(y) > signbit(x)),
                                    ifelse(isnan(x), x, y), ifelse(isnan(y), y, x))

minmax(x::T, y::T) where {T<:AbstractFloat} =
    ifelse(isnan(x) | isnan(y), ifelse(isnan(x), (x,x), (y,y)),
           ifelse((y > x) | (signbit(x) > signbit(y)), (x,y), (y,x)))

"""
    ldexp(x, n)

Compute ``x \\times 2^n``.

# Examples
```jldoctest
julia> ldexp(5., 2)
20.0
```
"""
function ldexp(x::T, e::Integer) where T<:IEEEFloat
    xu = reinterpret(Unsigned, x)
    xs = xu & ~sign_mask(T)
    xs >= exponent_mask(T) && return x # NaN or Inf
    k = Int(xs >> significand_bits(T))
    if k == 0 # x is subnormal
        xs == 0 && return x # +-0
        m = leading_zeros(xs) - exponent_bits(T)
        ys = xs << unsigned(m)
        xu = ys | (xu & sign_mask(T))
        k = 1 - m
        # underflow, otherwise may have integer underflow in the following n + k
        e < -50000 && return flipsign(T(0.0), x)
    end
    # For cases where e of an Integer larger than Int make sure we properly
    # overflow/underflow; this is optimized away otherwise.
    if e > typemax(Int)
        return flipsign(T(Inf), x)
    elseif e < typemin(Int)
        return flipsign(T(0.0), x)
    end
    n = e % Int
    k += n
    # overflow, if k is larger than maximum possible exponent
    if k >= exponent_raw_max(T)
        return flipsign(T(Inf), x)
    end
    if k > 0 # normal case
        xu = (xu & ~exponent_mask(T)) | (rem(k, uinttype(T)) << significand_bits(T))
        return reinterpret(T, xu)
    else # subnormal case
        if k <= -significand_bits(T) # underflow
            # overflow, for the case of integer overflow in n + k
            e > 50000 && return flipsign(T(Inf), x)
            return flipsign(T(0.0), x)
        end
        k += significand_bits(T)
        z = T(2.0)^-significand_bits(T)
        xu = (xu & ~exponent_mask(T)) | (rem(k, uinttype(T)) << significand_bits(T))
        return z*reinterpret(T, xu)
    end
end
ldexp(x::Float16, q::Integer) = Float16(ldexp(Float32(x), q))

"""
    exponent(x) -> Int

Get the exponent of a normalized floating-point number.
"""
function exponent(x::T) where T<:IEEEFloat
    @noinline throw1(x) = throw(DomainError(x, "Cannot be NaN or Inf."))
    @noinline throw2(x) = throw(DomainError(x, "Cannot be subnormal converted to 0."))
    xs = reinterpret(Unsigned, x) & ~sign_mask(T)
    xs >= exponent_mask(T) && throw1(x)
    k = Int(xs >> significand_bits(T))
    if k == 0 # x is subnormal
        xs == 0 && throw2(x)
        m = leading_zeros(xs) - exponent_bits(T)
        k = 1 - m
    end
    return k - exponent_bias(T)
end

"""
    significand(x)

Extract the `significand(s)` (a.k.a. mantissa), in binary representation, of a
floating-point number. If `x` is a non-zero finite number, then the result will be
a number of the same type on the interval ``[1,2)``. Otherwise `x` is returned.

# Examples
```jldoctest
julia> significand(15.2)/15.2
0.125

julia> significand(15.2)*8
15.2
```
"""
function significand(x::T) where T<:IEEEFloat
    xu = reinterpret(Unsigned, x)
    xs = xu & ~sign_mask(T)
    xs >= exponent_mask(T) && return x # NaN or Inf
    if xs <= (~exponent_mask(T) & ~sign_mask(T)) # x is subnormal
        xs == 0 && return x # +-0
        m = unsigned(leading_zeros(xs) - exponent_bits(T))
        xs <<= m
        xu = xs | (xu & sign_mask(T))
    end
    xu = (xu & ~exponent_mask(T)) | exponent_one(T)
    return reinterpret(T, xu)
end

"""
    frexp(val)

Return `(x,exp)` such that `x` has a magnitude in the interval ``[1/2, 1)`` or 0,
and `val` is equal to ``x \\times 2^{exp}``.
"""
function frexp(x::T) where T<:IEEEFloat
    xu = reinterpret(Unsigned, x)
    xs = xu & ~sign_mask(T)
    xs >= exponent_mask(T) && return x, 0 # NaN or Inf
    k = Int(xs >> significand_bits(T))
    if k == 0 # x is subnormal
        xs == 0 && return x, 0 # +-0
        m = leading_zeros(xs) - exponent_bits(T)
        xs <<= unsigned(m)
        xu = xs | (xu & sign_mask(T))
        k = 1 - m
    end
    k -= (exponent_bias(T) - 1)
    xu = (xu & ~exponent_mask(T)) | exponent_half(T)
    return reinterpret(T, xu), k
end

rem(x::Float64, y::Float64, ::RoundingMode{:Nearest}) =
    ccall((:remainder, libm),Float64,(Float64,Float64),x,y)
rem(x::Float32, y::Float32, ::RoundingMode{:Nearest}) =
    ccall((:remainderf, libm),Float32,(Float32,Float32),x,y)
rem(x::Float16, y::Float16, r::RoundingMode{:Nearest}) = Float16(rem(Float32(x), Float32(y), r))

"""
    modf(x)

Return a tuple `(fpart, ipart)` of the fractional and integral parts of a number. Both parts
have the same sign as the argument.

# Examples
```jldoctest
julia> modf(3.5)
(0.5, 3.0)

julia> modf(-3.5)
(-0.5, -3.0)
```
"""
modf(x) = isinf(x) ? (flipsign(zero(x), x), x) : (rem(x, one(x)), trunc(x))

function modf(x::Float32)
    temp = Ref{Float32}()
    f = ccall((:modff, libm), Float32, (Float32, Ptr{Float32}), x, temp)
    f, temp[]
end

function modf(x::Float64)
    temp = Ref{Float64}()
    f = ccall((:modf, libm), Float64, (Float64, Ptr{Float64}), x, temp)
    f, temp[]
end

@inline function ^(x::Float64, y::Float64)
    z = ccall("llvm.pow.f64", llvmcall, Float64, (Float64, Float64), x, y)
    if isnan(z) & !isnan(x+y)
        throw_exp_domainerror(x)
    end
    z
end
@inline function ^(x::Float32, y::Float32)
    z = ccall("llvm.pow.f32", llvmcall, Float32, (Float32, Float32), x, y)
    if isnan(z) & !isnan(x+y)
        throw_exp_domainerror(x)
    end
    z
end
@inline function ^(x::Float64, y::Integer)
    y == -1 && return inv(x)
    y == 0 && return one(x)
    y == 1 && return x
    y == 2 && return x*x
    y == 3 && return x*x*x
    ccall("llvm.pow.f64", llvmcall, Float64, (Float64, Float64), x, Float64(y))
end
@inline function ^(x::Float32, y::Integer)
    y == -1 && return inv(x)
    y == 0 && return one(x)
    y == 1 && return x
    y == 2 && return x*x
    y == 3 && return x*x*x
    ccall("llvm.pow.f32", llvmcall, Float32, (Float32, Float32), x, Float32(y))
end
@inline ^(x::Float16, y::Integer) = Float16(Float32(x) ^ y)
@inline literal_pow(::typeof(^), x::Float16, ::Val{p}) where {p} = Float16(literal_pow(^,Float32(x),Val(p)))

## rem2pi-related calculations ##

function add22condh(xh::Float64, xl::Float64, yh::Float64, yl::Float64)
    # This algorithm, due to Dekker, computes the sum of two
    # double-double numbers and returns the high double. References:
    # [1] http://www.digizeitschriften.de/en/dms/img/?PID=GDZPPN001170007
    # [2] https://doi.org/10.1007/BF01397083
    r = xh+yh
    s = (abs(xh) > abs(yh)) ? (xh-r+yh+yl+xl) : (yh-r+xh+xl+yl)
    zh = r+s
    return zh
end

# multiples of pi/2, as double-double (ie with "tail")
const pi1o2_h  = 1.5707963267948966     # convert(Float64, pi * BigFloat(1/2))
const pi1o2_l  = 6.123233995736766e-17  # convert(Float64, pi * BigFloat(1/2) - pi1o2_h)

const pi2o2_h  = 3.141592653589793      # convert(Float64, pi * BigFloat(1))
const pi2o2_l  = 1.2246467991473532e-16 # convert(Float64, pi * BigFloat(1) - pi2o2_h)

const pi3o2_h  = 4.71238898038469       # convert(Float64, pi * BigFloat(3/2))
const pi3o2_l  = 1.8369701987210297e-16 # convert(Float64, pi * BigFloat(3/2) - pi3o2_h)

const pi4o2_h  = 6.283185307179586      # convert(Float64, pi * BigFloat(2))
const pi4o2_l  = 2.4492935982947064e-16 # convert(Float64, pi * BigFloat(2) - pi4o2_h)

"""
    rem2pi(x, r::RoundingMode)

Compute the remainder of `x` after integer division by `2π`, with the quotient rounded
according to the rounding mode `r`. In other words, the quantity

    x - 2π*round(x/(2π),r)

without any intermediate rounding. This internally uses a high precision approximation of
2π, and so will give a more accurate result than `rem(x,2π,r)`

- if `r == RoundNearest`, then the result is in the interval ``[-π, π]``. This will generally
  be the most accurate result. See also [`RoundNearest`](@ref).

- if `r == RoundToZero`, then the result is in the interval ``[0, 2π]`` if `x` is positive,.
  or ``[-2π, 0]`` otherwise. See also [`RoundToZero`](@ref).

- if `r == RoundDown`, then the result is in the interval ``[0, 2π]``.
  See also [`RoundDown`](@ref).
- if `r == RoundUp`, then the result is in the interval ``[-2π, 0]``.
  See also [`RoundUp`](@ref).

# Examples
```jldoctest
julia> rem2pi(7pi/4, RoundNearest)
-0.7853981633974485

julia> rem2pi(7pi/4, RoundDown)
5.497787143782138
```
"""
function rem2pi end
function rem2pi(x::Float64, ::RoundingMode{:Nearest})
    abs(x) < pi && return x

    n,y = rem_pio2_kernel(x)

    if iseven(n)
        if n & 2 == 2 # n % 4 == 2: add/subtract pi
            if y.hi <= 0
                return add22condh(y.hi,y.lo,pi2o2_h,pi2o2_l)
            else
                return add22condh(y.hi,y.lo,-pi2o2_h,-pi2o2_l)
            end
        else          # n % 4 == 0: add 0
            return y.hi+y.lo
        end
    else
        if n & 2 == 2 # n % 4 == 3: subtract pi/2
            return add22condh(y.hi,y.lo,-pi1o2_h,-pi1o2_l)
        else          # n % 4 == 1: add pi/2
            return add22condh(y.hi,y.lo,pi1o2_h,pi1o2_l)
        end
    end
end
function rem2pi(x::Float64, ::RoundingMode{:ToZero})
    ax = abs(x)
    ax <= 2*Float64(pi,RoundDown) && return x

    n,y = rem_pio2_kernel(x)

    if iseven(n)
        if n & 2 == 2 # n % 4 == 2: add pi
            z = add22condh(y.hi,y.lo,pi2o2_h,pi2o2_l)
        else          # n % 4 == 0: add 0 or 2pi
            if y.hi > 0
                z = y.hi+y.lo
            else      # negative: add 2pi
                z = add22condh(y.hi,y.lo,pi4o2_h,pi4o2_l)
            end
        end
    else
        if n & 2 == 2 # n % 4 == 3: add 3pi/2
            z = add22condh(y.hi,y.lo,pi3o2_h,pi3o2_l)
        else          # n % 4 == 1: add pi/2
            z = add22condh(y.hi,y.lo,pi1o2_h,pi1o2_l)
        end
    end
    copysign(z,x)
end
function rem2pi(x::Float64, ::RoundingMode{:Down})
    if x < pi4o2_h
        if x >= 0
            return x
        elseif x > -pi4o2_h
            return add22condh(x,0.0,pi4o2_h,pi4o2_l)
        end
    end

    n,y = rem_pio2_kernel(x)

    if iseven(n)
        if n & 2 == 2 # n % 4 == 2: add pi
            return add22condh(y.hi,y.lo,pi2o2_h,pi2o2_l)
        else          # n % 4 == 0: add 0 or 2pi
            if y.hi > 0
                return y.hi+y.lo
            else      # negative: add 2pi
                return add22condh(y.hi,y.lo,pi4o2_h,pi4o2_l)
            end
        end
    else
        if n & 2 == 2 # n % 4 == 3: add 3pi/2
            return add22condh(y.hi,y.lo,pi3o2_h,pi3o2_l)
        else          # n % 4 == 1: add pi/2
            return add22condh(y.hi,y.lo,pi1o2_h,pi1o2_l)
        end
    end
end
function rem2pi(x::Float64, ::RoundingMode{:Up})
    if x > -pi4o2_h
        if x <= 0
            return x
        elseif x < pi4o2_h
            return add22condh(x,0.0,-pi4o2_h,-pi4o2_l)
        end
    end

    n,y = rem_pio2_kernel(x)

    if iseven(n)
        if n & 2 == 2 # n % 4 == 2: sub pi
            return add22condh(y.hi,y.lo,-pi2o2_h,-pi2o2_l)
        else          # n % 4 == 0: sub 0 or 2pi
            if y.hi < 0
                return y.hi+y.lo
            else      # positive: sub 2pi
                return add22condh(y.hi,y.lo,-pi4o2_h,-pi4o2_l)
            end
        end
    else
        if n & 2 == 2 # n % 4 == 3: sub pi/2
            return add22condh(y.hi,y.lo,-pi1o2_h,-pi1o2_l)
        else          # n % 4 == 1: sub 3pi/2
            return add22condh(y.hi,y.lo,-pi3o2_h,-pi3o2_l)
        end
    end
end

rem2pi(x::Float32, r::RoundingMode) = Float32(rem2pi(Float64(x), r))
rem2pi(x::Float16, r::RoundingMode) = Float16(rem2pi(Float64(x), r))
rem2pi(x::Int32, r::RoundingMode) = rem2pi(Float64(x), r)
function rem2pi(x::Int64, r::RoundingMode)
    fx = Float64(x)
    fx == x || throw(ArgumentError("Int64 argument to rem2pi is too large: $x"))
    rem2pi(fx, r)
end

"""
    mod2pi(x)

Modulus after division by `2π`, returning in the range ``[0,2π)``.

This function computes a floating point representation of the modulus after division by
numerically exact `2π`, and is therefore not exactly the same as `mod(x,2π)`, which would
compute the modulus of `x` relative to division by the floating-point number `2π`.

!!! note
    Depending on the format of the input value, the closest representable value to 2π may
    be less than 2π. For example, the expression `mod2pi(2π)` will not return `0`, because
    the intermediate value of `2*π` is a `Float64` and `2*Float64(π) < 2*big(π)`. See
    [`rem2pi`](@ref) for more refined control of this behavior.

# Examples
```jldoctest
julia> mod2pi(9*pi/4)
0.7853981633974481
```
"""
mod2pi(x) = rem2pi(x,RoundDown)

# generic fallback; for number types, promotion.jl does promotion

"""
    muladd(x, y, z)

Combined multiply-add: computes `x*y+z`, but allowing the add and multiply to be merged
with each other or with surrounding operations for performance.
For example, this may be implemented as an [`fma`](@ref) if the hardware supports it
efficiently.
The result can be different on different machines and can also be different on the same machine
due to constant propagation or other optimizations.
See [`fma`](@ref).

# Examples
```jldoctest
julia> muladd(3, 2, 1)
7

julia> 3 * 2 + 1
7
```
"""
muladd(x,y,z) = x*y+z

# Float16 definitions

for func in (:sin,:cos,:tan,:asin,:acos,:atan,:sinh,:cosh,:tanh,:asinh,:acosh,
             :atanh,:exp,:exp2,:exp10,:log,:log2,:log10,:sqrt,:lgamma,:log1p)
    @eval begin
        $func(a::Float16) = Float16($func(Float32(a)))
        $func(a::ComplexF16) = ComplexF16($func(ComplexF32(a)))
    end
end

for func in (:atan,:hypot)
    @eval begin
        $func(a::Float16,b::Float16) = Float16($func(Float32(a),Float32(b)))
    end
end

cbrt(a::Float16) = Float16(cbrt(Float32(a)))
sincos(a::Float16) = Float16.(sincos(Float32(a)))

# helper functions for Libm functionality

"""
    highword(x)

Return the high word of `x` as a `UInt32`.
"""
@inline highword(x::Float64) = highword(reinterpret(UInt64, x))
@inline highword(x::UInt64)  = (x >>> 32) % UInt32
@inline highword(x::Float32) = reinterpret(UInt32, x)

@inline fromhighword(::Type{Float64}, u::UInt32) = reinterpret(Float64, UInt64(u) << 32)
@inline fromhighword(::Type{Float32}, u::UInt32) = reinterpret(Float32, u)

"""
    poshighword(x)

Return positive part of the high word of `x` as a `UInt32`.
"""
@inline poshighword(x::Float64) = poshighword(reinterpret(UInt64, x))
@inline poshighword(x::UInt64)  = highword(x) & 0x7fffffff
@inline poshighword(x::Float32) = highword(x) & 0x7fffffff

# More special functions
include("special/cbrt.jl")
include("special/exp.jl")
include("special/exp10.jl")
include("special/ldexp_exp.jl")
include("special/hyperbolic.jl")
include("special/trig.jl")
include("special/rem_pio2.jl")
include("special/log.jl")

# `missing` definitions for functions in this module
for f in (:(acos), :(acosh), :(asin), :(asinh), :(atan), :(atanh),
          :(sin), :(sinh), :(cos), :(cosh), :(tan), :(tanh),
          :(exp), :(exp2), :(expm1), :(log), :(log10), :(log1p),
          :(log2), :(exponent), :(sqrt))
    @eval $(f)(::Missing) = missing
end
clamp(::Missing, lo, hi) = missing

end # module


using FastGaussQuadrature: gausslegendre

"""
	cache = POOrthogonalCollocationCache(Ntst::Int, m::Int, Ty = Float64)

Structure to hold the cache for the collocation method.

# Arguments
- `Ntst` number of time steps
- `m` degree of the polynomials
- `Ty` type of the time variable
"""
struct POOrthogonalCollocationCache{T}
	Ntst::Int
	degree::Int
    lagrange_vals::Matrix{T}
    lagrange_driv::Matrix{T}
	gauss_nodes::Vector{T}
    gauss_weight::Vector{T}
	# TODO how do we ensure that eltype(mesh) is what we expect?
	mesh::Vector{T} 		# τs, we need a vector here for mesh adaptation
	mesh_coll::LinRange{T} 	# σs
	full_mesh::Vector{T}
end

function POOrthogonalCollocationCache(Ntst::Int, m::Int, Ty = Float64)
	τs = LinRange{Ty}( 0, 1, Ntst + 1) |> collect
	σs = LinRange{Ty}(-1, 1, m + 1)
	L, ∂L = getL(σs)
	zg, wg = gausslegendre(m)
	prob = POOrthogonalCollocationCache(Ntst, m, L, ∂L, zg, wg, τs, σs, zeros(Ty, 1 + m * Ntst))
	# put the mesh where we removed redundant timing
	prob.full_mesh .= getTimes(prob)
	return prob
end

@inline Base.eltype(pb::POOrthogonalCollocationCache) = eltype(pb.lagrange_vals)
@inline Base.size(pb::POOrthogonalCollocationCache) = (pb.degree, pb.Ntst)
@inline getLs(pb::POOrthogonalCollocationCache) = (pb.lagrange_vals, pb.lagrange_driv)
@inline getMesh(pb::POOrthogonalCollocationCache) = pb.mesh
@inline getMeshColl(pb::POOrthogonalCollocationCache) = pb.mesh_coll
getMaxTimeStep(pb::POOrthogonalCollocationCache) = maximum(diff(getMesh(pb)))
τj(σ, τs, j) = τs[j] + (1 + σ)/2 * (τs[j+1] - τs[j])
# get the sigma corresponding to τ in the interval (𝜏s[j], 𝜏s[j+1])
σj(τ, τs, j) = -(2*τ - τs[j] - τs[j + 1])/(-τs[j + 1] + τs[j])

# code from Jacobi.lagrange
function lagrange(i, x, z)
    nz = length(z)
    l = one(z[1])
	for k in 1:(i-1)
        l = l * (x - z[k]) / (z[i] - z[k])
    end
    for k in (i+1):nz
        l = l * (x - z[k]) / (z[i] - z[k])
    end
    return l
end

dlagrange(i, x, z) = ForwardDiff.derivative(x -> lagrange(i, x, z), x)

# accept a range, ie σs = LinRange(-1, 1, m + 1)
function getL(σs::AbstractVector)
	m = length(σs) - 1
	zs, = gausslegendre(m)
	L = zeros(m, m+1); ∂L = zeros(m, m + 1)
	for j in 1:m+1
		for i in 1:m
			L[i, j]  =  lagrange(j, zs[i], σs)
			∂L[i, j] = dlagrange(j, zs[i], σs)
		end
	end
	return (;L, ∂L)
end

"""
$(SIGNATURES)

Return the times at which is evaluated the collocation problem.
"""
function getTimes(pb::POOrthogonalCollocationCache)
	m, Ntst = size(pb)
	Ty = eltype(pb)
	ts = Matrix{Ty}(undef, Ntst, m + 1)
	tsvec = Ty[0]
	τs = pb.mesh
	σs = pb.mesh_coll
	for j in 1:Ntst
		for l in 1:m+1
			ts[j, l] = τj(σs[l], τs, j)
			l>1 && push!(tsvec, τj(σs[l], τs, j))
		end
	end
	return vec(tsvec)
end

function updateMesh!(pb::POOrthogonalCollocationCache, mesh)
	pb.mesh .= mesh
	pb.full_mesh .= getTimes(pb)
end
####################################################################################################

"""
	pb = PeriodicOrbitOCollProblem(kwargs...)

This composite type implements an orthogonal collocation (at Gauss points) method of piecewise polynomials to locate periodic orbits. More details (maths, notations, linear systems) can be found [here](https://bifurcationkit.github.io/BifurcationKitDocs.jl/dev/periodicOrbitCollocation/).

## Arguments
- `F` vector field specified as a function of two arguments `F(x,p)`
- `J` is the jacobian of `F` at `(x, p)`. It can assume three forms.
    1. Either `J` is a function and `J(x, p)` returns a `::AbstractMatrix`. In this case, the default arguments of `contParams::ContinuationPar` will make `continuation` work.
    2. Or `J` is a function and `J(x, p)` returns a function taking one argument `dx` and returning `dr` of the same type as `dx`. In our notation, `dr = J * dx`. In this case, the default parameters of `contParams::ContinuationPar` will not work and you have to use a Matrix Free linear solver, for example `GMRESIterativeSolvers`,
    3. Or `J` is a function and `J(x, p)` returns a variable `j` which can assume any type. Then, you must implement a linear solver `ls` as a composite type, subtype of `AbstractLinearSolver` which is called like `ls(j, rhs)` and which returns the solution of the jacobian linear system. See for example `examples/SH2d-fronts-cuda.jl`. This linear solver is passed to `NewtonPar(linsolver = ls)` which itself passed to `ContinuationPar`. Similarly, you have to implement an eigensolver `eig` as a composite type, subtype of `AbstractEigenSolver`.
- `ϕ::AbstractVector` used to set a section for the phase constraint equation
- `xπ::AbstractVector` used in the section for the phase constraint equation
- `N::Int` dimension of the state space
- `coll_cache::POOrthogonalCollocationCache` cache for collocation. See docs of `POOrthogonalCollocationCache` .

## Methods

Here are some useful methods you can apply to `pb`

- `length(pb)` gives the total number of unknowns
- `size(pb)` returns the triplet `(N, m, Ntst)`
- `getMesh(pb)` returns the mesh `0 = τ0 < ... < τNtst+1 = 1`. This is useful because this mesh is born to vary by automatic mesh adaptation
- `getMeshColl(pb)` returns the (static) mesh `0 = σ0 < ... < σm+1 = 1`
- `getTimes(pb)` returns the vector of times (length `1 + m * Ntst`) at the which the collocation is applied.
- `generateSolution(pb, orbit, period)` generate a guess from a function `t -> orbit(t)` which approximates the periodic orbit.
- `POOcollSolution(pb, x)` return a function interpolating the solution `x` using a piecewise polynomials function

# Orbit guess
You will see below that you can evaluate the residual of the functional (and other things) by calling `pb(orbitguess, p)` on an orbit guess `orbitguess`. Note that `orbitguess` must be of size 1 + N * (1 + m * Ntst) where N is the number of unknowns in the state space and `orbitguess[end]` is an estimate of the period ``T`` of the limit cycle.

Note that you can generate this guess from a function using `generateSolution`.

# Functional
 A functional, hereby called `G`, encodes this problem. The following methods are available

- `pb(orbitguess, p)` evaluates the functional G on `orbitguess`

"""
@with_kw_noshow struct PeriodicOrbitOCollProblem{TF, TJ, vectype, Tmass, Tcache <: POOrthogonalCollocationCache} <: AbstractPeriodicOrbitProblem
	# Function F(x, par)
	F::TF = nothing

	# Jacobian of F w.r.t. x
	J::TJ = nothing

	# variables to define a Section for the phase constraint equation
	ϕ::vectype = nothing
	xπ::vectype = nothing

	# dimension of the problem in case of an AbstractVector
	N::Int = 0

	# whether the time discretisation is adaptive
	adaptmesh::Bool = false

	# whether the problem is nonautonomous
	isautonomous::Bool = true

	# mass matrix
	massmatrix::Tmass = nothing

	# collocation cache
	coll_cache::Tcache = nothing
end

# trivial constructor
function PeriodicOrbitOCollProblem(Ntst, m, N = 0)
	cache = POOrthogonalCollocationCache(Ntst, m)
	PeriodicOrbitOCollProblem(; N = N, coll_cache = cache)
end

# TODO rename this in num_mesh? or meshSize
@inline getMeshSize(pb::PeriodicOrbitOCollProblem) = pb.coll_cache.Ntst

@inline length(pb::PeriodicOrbitOCollProblem) = ( (n, m, Ntst) = size(pb); return n * (1 + m * Ntst) )
# the size is (n, m, Ntst)
@inline Base.size(pb::PeriodicOrbitOCollProblem) = (pb.N, size(pb.coll_cache)...)
@inline Base.eltype(pb::PeriodicOrbitOCollProblem) = eltype(pb.coll_cache)
getLs(pb::PeriodicOrbitOCollProblem) = getLs(pb.coll_cache)

# these functions extract the time slices components
getTimeSlices(x::AbstractVector, N, degree, Ntst) = @views reshape(x[1:end-1], N, (degree) * Ntst + 1)
# array of size Ntst ⋅ (m+1) ⋅ n
getTimeSlices(pb::PeriodicOrbitOCollProblem, x) = getTimeSlices(x, size(pb)...)
getTimes(pb::PeriodicOrbitOCollProblem) = getTimes(pb.coll_cache)
getMesh(pb::PeriodicOrbitOCollProblem) = getMesh(pb.coll_cache)
getMeshColl(pb::PeriodicOrbitOCollProblem) = getMeshColl(pb.coll_cache)
getMaxTimeStep(pb::PeriodicOrbitOCollProblem) = getMaxTimeStep(pb.coll_cache)
updateMesh!(pb::PeriodicOrbitOCollProblem, mesh) = updateMesh!(pb.coll_cache, mesh)

function Base.show(io::IO, pb::PeriodicOrbitOCollProblem)
	N, m, Ntst = size(pb)
	println(io, "┌─ Collocation problem for periodic orbits")
	println(io, "├─ type            : Vector{", eltype(pb), "}")
	println(io, "├─ time slices     : ", Ntst)
	println(io, "├─ degree          : ", m)
	println(io, "├─ space dimension : ", pb.N)
	println(io, "└─ # unknowns      : ", pb.N * (1 + m * Ntst))
end

"""
$(SIGNATURES)

This function generates an initial guess for the solution of the problem `pb` based on the orbit `t -> orbit(t)` for t ∈ [0,1] and the period `period`.
"""
function generateSolution(pb::PeriodicOrbitOCollProblem, orbit, period)
	n, _m, Ntst = size(pb)
	ts = getTimes(pb)
	Nt = length(ts)
	ci = zeros(eltype(pb), n, Nt)
	for (l, t) in pairs(ts)
		ci[:, l] .= orbit(t * period)
	end
	return vcat(vec(ci), period)
end

@views function (prob::PeriodicOrbitOCollProblem)(u::AbstractVector, pars)
	uc = getTimeSlices(prob, u)
	T = getPeriod(prob, u, pars)
	result = zero(u)
	resultc = getTimeSlices(prob, result)
	functionalColl!(prob, resultc, uc, T, getLs(prob.coll_cache), pars)
	# add  the phase condition
	result[end] = dot(u[1:end-1], prob.ϕ) - dot(prob.xπ, prob.ϕ)
	return result
end

# function or collocation problem
@views function functionalColl!(pb, out, u, period, (L, ∂L), pars)
	Ty = eltype(u)
	n, ntimes = size(u)
	m = pb.coll_cache.degree
	Ntst = pb.coll_cache.Ntst
	# on veut faire des slices a t fixes, donc des gj[:, j], c'est le plus rapide
	gj = zeros(Ty, n, m)
	∂gj = zeros(Ty, n, m)
	uj = zeros(Ty, n, m+1)
	mesh = getMesh(pb)
	rg = 1:m+1
	for j in 1:Ntst
		uj .= u[:, rg]
		gj 	.=  uj * L'
		∂gj .=  uj * ∂L'
		# mul!(gj, uj, L')
		# mul!(∂gj, uj, ∂L')
		# compute the collocation residual
		for l in 1:m
			out[:, rg[l]] .= ∂gj[:, l] .- (period * (mesh[j+1]-mesh[j]) / (2)) .* pb.F(gj[:, l], pars)
		end
		rg = rg .+ m
	end
	# add the periodicity condition
	out[:, end] .= u[:, end] .- u[:, 1]
end

"""
$(SIGNATURES)

Compute the full periodic orbit associated to `x`. Mainly for plotting purposes.
"""
@views function getPeriodicOrbit(prob::PeriodicOrbitOCollProblem, u::AbstractVector, p)
	T = getPeriod(prob, u, p)
	ts = getTimes(prob)
	uc = getTimeSlices(prob, u)
	n, m, Ntst = size(prob)
	return SolPeriodicOrbit(t = ts .* T, u = uc)
end

# function needed for automatic Branch switching from Hopf bifurcation point
function reMake(prob::PeriodicOrbitOCollProblem, F, dF, par, hopfpt, ζr::AbstractVector, orbitguess_a, period; orbit = t->t)
	M = length(orbitguess_a)
	N = length(ζr)

	n, m, Ntst = size(prob)
	nunknows = N * (1 + m*Ntst)

	# update the problem
	probPO = setproperties(prob, N = N, F = F, J = dF, ϕ = zeros(nunknows), xπ = zeros(nunknows))

	probPO.ϕ[1:N] .= ζr
	probPO.xπ[1:N] .= hopfpt.x0

	# append period at the end of the initial guess
	orbitguess = generateSolution(probPO, t -> orbit(2pi*t/period), period)

	return probPO, orbitguess
end

####################################################################################################
const DocStrjacobianPOColl = """
- `jacobianPO` Specify the choice of the linear algorithm, which must belong to `(:autodiffDense, :FiniteDifferences)`. This is used to select a way of inverting the jacobian dG
    - For `:autodiffDense`. Same as for `:autodiffMF` but the jacobian is formed as a dense Matrix. You can use a direct solver or an iterative one using `options`.
    - For `:FiniteDifferencesDense`, same as for `:autodiffDense` but we use Finite Differences to compute the jacobian of `x -> prob(x, p)` using the `δ = 1e-8` which can be passed as an argument.
"""

"""
$(SIGNATURES)

This is the Newton Solver for computing a periodic orbit using (Standard / Poincaré) Shooting method.
Note that the linear solver has to be apropriately set up in `options`.

# Arguments

Similar to [`newton`](@ref) except that `prob` is either a [`ShootingProblem`](@ref) or a [`PoincareShootingProblem`](@ref). These two problems have specific options to be tuned, we refer to their link for more information and to the tutorials.

- `prob` a problem of type `<: AbstractShootingProblem` encoding the shooting functional G.
- `orbitguess` a guess for the periodic orbit. See [`ShootingProblem`](@ref) and See [`PoincareShootingProblem`](@ref) for information regarding the shape of `orbitguess`.
- `par` parameters to be passed to the functional
- `options` same as for the regular [`newton`](@ref) method.

# Optional argument
$DocStrjacobianPOColl

# Output:
- solution
- history of residuals
- flag of convergence
- number of iterations
"""
function newton(prob::PeriodicOrbitOCollProblem, orbitguess, par, options::NewtonPar;
		jacobianPO = :autodiffDense, kwargs...)
	@assert jacobianPO in
			(:autodiffDense, ) "This jacobian is not defined. Please chose another one."

	if jacobianPO == :autodiffDense
		jac = (x, p) -> ForwardDiff.jacobian(z -> prob(z, p), x)
	end

	return newton(prob, jac, orbitguess, par, options; kwargs...)
end

"""
$(SIGNATURES)

This is the continuation method for computing a periodic orbit using a (Standard / Poincaré) Shooting method.

# Arguments

Similar to [`continuation`](@ref) except that `prob` is either a [`ShootingProblem`](@ref) or a [`PoincareShootingProblem`](@ref). By default, it prints the period of the periodic orbit.

# Optional argument
- `δ = 1e-8` used for finite differences
- `jacobianPO` Specify the choice of the linear algorithm, which must belong to `[:autodiffMF, :MatrixFree, :autodiffDense, :autodiffDenseAnalytical, :FiniteDifferences]`. This is used to select a way of inverting the jacobian dG
- `updateSectionEveryStep = 0` updates the section every `updateSectionEveryStep` step during continuation

# Choices for `jacobianPO`
$DocStrjacobianPOColl
"""
function continuation(
	prob::PeriodicOrbitOCollProblem,
	orbitguess, par, lens::Lens, _contParams::ContinuationPar,
	_linearAlgo::AbstractBorderedLinearSolver;
	jacobianPO = :autodiffDense,
	updateSectionEveryStep = 0, kwargs...)
	@assert jacobianPO in
			(:autodiffDense,) "This jacobian is not defined. Please chose another one."

	jac = (x, p) -> FloquetWrapper(prob, ForwardDiff.jacobian(z -> prob(z, p), x), x, p)

	# we have to change the Bordered linearsolver to cope with our type FloquetWrapper
	linearAlgo = @set _linearAlgo.solver = FloquetWrapperLS(_linearAlgo.solver)

	options = _contParams.newtonOptions
	contParams = @set _contParams.newtonOptions.linsolver = FloquetWrapperLS(options.linsolver)

	if computeEigenElements(contParams)
		contParams = @set contParams.newtonOptions.eigsolver = FloquetLUColl(contParams.newtonOptions.eigsolver, length(prob), prob.N)
	end

	# change the user provided finalise function by passing prob in its parameters
	_finsol = modifyPOFinalise(prob, kwargs, updateSectionEveryStep)
	_recordsol = modifyPORecord(prob, kwargs, par, lens)
	_plotsol = modifyPOPlot(prob, kwargs)

	branch, u, τ = continuation(
					prob, jac,
					orbitguess, par, lens,
					contParams, linearAlgo;
					kwargs...,
					recordFromSolution = _recordsol,
					finaliseSolution = _finsol,
					plotSolution = _plotsol)
	return setproperties(branch; type = :PeriodicOrbit, functional = prob), u, τ
end

"""
$(SIGNATURES)

Compute the maximum of the periodic orbit associated to `x`.
"""
function getMaximum(prob::PeriodicOrbitOCollProblem, x::AbstractVector, p)
	sol = getPeriodicOrbit(prob, x, p).u
	return maximum(sol)
end

# this function updates the section during the continuation run
@views function updateSection!(prob::PeriodicOrbitOCollProblem, x, par; stride = 0)
	n, m, Ntst = size(prob)
	xc = getTimeSlices(prob, x)
	T = getPeriod(prob, x, par)

	# update the reference point
	prob.xπ .= x[1:end-1]

	# update the normals
	ϕc = @views reshape(prob.ϕ, n, (m) * Ntst + 1)
	for ii in 1: (m) * Ntst + 1
		ϕc[:, ii] .= prob.F(xc[:, ii], par) ./ (m*Ntst)
	end
	return true
end
####################################################################################################
# mesh adaptation routine

# iterated derivatives
∂(f) = x -> ForwardDiff.derivative(f, x)
∂(f, n) = n == 0 ? f : ∂(∂(f), n-1)

struct POOcollSolution{Tpb, Tx}
	pb::Tpb
	x::Tx
end

@views function (sol::POOcollSolution)(t0)
	n,m,Ntst = size(sol.pb)
	xc = getTimeSlices(sol.pb, sol.x)

	T = getPeriod(sol.pb, sol.x, 0)
	t = t0 / T

	mesh = getMesh(sol.pb)
	indτ = searchsortedfirst(mesh, t) - 1
	if indτ <= 0
		return sol.x[1:n]
	elseif indτ > Ntst
		return xc[:, end]
	end
	# println("--> ", t, " belongs to ", (mesh[indτ], mesh[indτ+1])) # waste lots of ressources
	@assert mesh[indτ] <= t <= mesh[indτ+1] "Please open an issue on the website of BifurcationKit.jl"
	σ = σj(t, mesh, indτ)
	# @assert -1 <= σ <= 1 "Strange value of $σ"
	σs = getMeshColl(sol.pb)
	out = zeros(typeof(t), sol.pb.N)
	rg = (1:m+1) .+ (indτ-1)*m
	for l in 1:m+1
		out .+= xc[:, rg[l]] .* lagrange(l, σ, σs)
	end
	out
end


function reallocation!(state)
    len = length(state)
    index = argmax(state)
    value = state[index]
    state[index] = 0
    for i in 1:value
        target_index = index + i
        if target_index > len
            target_index = target_index % len
        end
        state[target_index] += 1
    end
end

function solve(state)
    # part a
    count = 1
    know_states = Set([copy(state)])

    while true
        reallocation!(state)
        if state ∈ know_states
            break
        end
        push!(know_states, copy(state))
        count += 1
    end
    sol_a = count
    # part b
    count = 1
    target_state = copy(state)
    while true
        reallocation!(state)
        if state == target_state
            break
        end
        count += 1
    end

    sol_b = count
    sol_a, sol_b
end

data = parse.(Int, split(readline("input.txt")));
sol_a, sol_b = solve(data)
println("sol a: $(sol_a)")
println("sol b: $(sol_b)")


using PtFEM

data = Dict(
  # Plane(ndim, nst, nxe, nye, nip, direction, finite_element(nod, nodof), axisymmetric)
  :struc_el => Plane(2, 3, 20, 2, 4, :y, Quadrilateral(8, 2), false),
  :properties => [1.0e5 0.0 1.0e-5 1.0e-5;],
  :x_coords => collect(range(0.0, stop=1.0, length=21)),
  :y_coords => collect(range(0.05, stop=-0.05, length=3)),
  :temp => [],
  :support => [
    (1, [0 0]),
    (2, [0 0]),
    (3, [0 0]),
    (4, [0 0]),
    (5, [0 0])
  ],
  :dtemp => [
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50, 0.50, 0.00, -0.50,
    0.50, 0.25, 0.00, -0.25, -0.50
  ],
  :nspr => [
    (163, 2, 50.0)
  ]
)

data |> display
println()

@time m, dis_df, fm_df = p55(data)
println()

dis_df |> display
println()

println("\nThe integration point (nip = $(data[:struc_el].nip)) stresses are:\n")

fm_df |> display
println()