using Distributions
using LinearAlgebra
using LinearManifoldModels
using Statistics
import StatsBase
import KernelDensity
import LMCLUS
import Random
using Test
# using Logging; global_logger(ConsoleLogger(stderr, LMCLUS.Base.CoreLogging.Debug))

const SEED = 986163731

# generate dataset
Random.seed!(SEED)
X1 = rand(Normal(0., 0.7), 2, 100)
X2 = rand(MvNormal([2., 2.], diagm(0=>[0.02, 0.6])), 100)
X = [X1 X2]
Xn = X .- mean(X, dims=2)

# generate clusters
p = LMCLUS.Parameters(1)
p.random_seed = SEED
p.sampling_heuristic = 1
p.sampling_factor = 1.5
p.max_bin_portion = 0.2
p.best_bound = 1.0
p.basis_alignment = true
p.dim_adjustment = true
p.dim_adjustment_ratio = 0.1
res = LMCLUS.lmclus(Xn, p)

# test distribution properties
@testset for E in [StatsBase.Histogram, KernelDensity.UnivariateKDE]
    for C in LMCLUS.manifolds(res)
        L = LMCLUS.points(C)
        lmcm = estimate(NonparametricEstimator(EmpiricalLinearManifold, E), C, Xn[:,L])
        @test length(lmcm) == 2
        @test params(lmcm) == (mean(C), LMCLUS.projection(C))
        @test partype(lmcm) == typeof(Xn[1])
        @test all(pdf(lmcm, Xn[:,L]) .> 0.)
        @test all(.!isinf.(logpdf(lmcm, Xn[:,L])))
        @test entropy(lmcm) > 0.
    end
end


struct BackwardReach end

import LazySets:HPoly

# This is the main function
function solve(solver::BackwardReach, problem::Problem)
    safe_zone = backward_network(solver, problem.network, problem.output)
    if isempty(safe_zone)
        return BasicResult(:unknown)
    end
    print("safe_zone $(dim(safe_zone))\n")
    return issubset(problem.input, safe_zone) ? BasicResult(:holds) : BasicResult(:violated)
end

function backward_network(solver::BackwardReach, nnet::Network, input::HPolyhedron{Float64})
    reach = input
    for i in length(nnet.layers):-1:1
        reach = backward_layer(solver, nnet.layers[i], nnet.layers[i].activation, reach)
    end
    return reach
end

function backward_layer(solver::BackwardReach, layer::Layer, act::Id, input::HPolyhedron{Float64})
    lc = get_constraints(input)
    (W, b) = (layer.weights, layer.bias)
    lco = backward_affine_map(lc, W, b)
    remove_redundant_constraints!(lco)
    return HPolyhedron(lco)
end

function backward_layer(solver::BackwardReach, layer::Layer, act::ReLU, input::HPolyhedron{Float64})
    n_dim = dim(input)
    lc = get_constraints(input)
    (W, b) = (layer.weights, layer.bias)
    lcs = Vector{LinearConstraint{Float64, Vector{Float64}}}(undef, 0)
    diag = zeros(Float64, (n_dim, n_dim))
    b = zeros(Float64, n_dim)
    act_status = zeros(Int64, n_dim)
    n_part = 2 ^ n_dim
    for i in 1:n_part
        n = i
        lci = Vector{LinearConstraint{Float64, Vector{Float64}}}(undef, 0)
        for j in 1:n_dim
            n, act_status[j] = fldmod1(n, 2)
            diag[j, j] = 1.0
            if act_status[j] == 1
                push!(lci, LinearConstraint(diag[j,:], 0.0))
            end
            push!(lci, LinearConstraint(-diag[j,:], 0.0))
        end #setfield! immutable struct of type HalfSpace cannot be changed
        append!(lci, lc)
        if isempty(HPolyhedron(lci))
            continue
        end
        remove_redundant_constraints!(lci) #LinearConstraint
        lci_to_delete = Vector{Int64}(undef, 0)
        for k in 1:length(lci)
            if !in(lci[k], lc)
                push!(lci_to_delete, k)
            end
        end
        deleteat!(lci, lci_to_delete)
        Wi = diag
        bi = b
        for l in 1:n_dim
            Wi[l, l] = act_status[l] == 1 ? 0.0 : 1.0
            bi[l] = 0.0
        end
        append!(lcs, backward_affine_map(lci, Wi, bi))
    end
    if length(lcs) == 0
        ni = size(W, 2)
        diagi = zeros(Float64, (ni, ni))
        for m in 1:ni
            diagi[m, m] = 1.0
            push!(lcs, LinearConstraint(diagi[m,:], 0.0))
        end
        return HPolyhedron(lcs)
    end
    lco = backward_affine_map(lcs, W, b)
    remove_redundant_constraints!(lco)
    return HPolyhedron(lco)
end

function backward_layer(solver::BackwardReach, layer::Layer, act::PiecewiseLinearActivation, input::HPolyhedron{Float64})
    n_dim = dim(input)
    lc = get_constraints(input)
    (W, b) = (layer.weights, layer.bias)
    inflections = get_inflection(act)
    lcs = Vector{LinearConstraint{Float64, Vector{Float64}}}(undef, 0)
    diag = zeros(Float64, (n_dim, n_dim))
    b = zeros(Float64, n_dim)
    if inflections == []
        W1 = diag
        b1 = b
        for l in 1:n_dim
            W1[l, l] = act.pieces[1].slope
            b1[l] = act.pieces[1].intercept
        end
        append!(lcs, backward_affine_map(lc, W1, b1))
    else
        out_bounds = [-Inf; act.(inflections); Inf]
        n_interval = length(inflections) + 1
        act_status = zeros(Int64, n_dim)
        n_part = n_interval ^ n_dim
        for i in 1:n_part
            n = i
            lci = Vector{LinearConstraint{Float64, Vector{Float64}}}(undef, 0)
            for j in 1:n_dim
                n, act_status[j] = fldmod1(n, n_interval)
                diag[j, j] = 1.0
                act_status[j] == n_interval || push!(lci, LinearConstraint(diag[j,:], out_bounds[act_status[j] + 1])) #right bound
                act_status[j] == 1 || push!(lci, LinearConstraint(-diag[j,:], -out_bounds[act_status[j]])) #left bound
            end #setfield! immutable struct of type HalfSpace cannot be changed
            append!(lci, lc)
            if isempty(HPolyhedron(lci))
                continue
            end
            remove_redundant_constraints!(lci) #LinearConstraint
            lci_to_delete = Vector{Int64}(undef, 0)
            for k in 1:length(lci)
                if !in(lci[k], lc)
                    push!(lci_to_delete, k)
                end
            end
            deleteat!(lci, lci_to_delete)
            Wi = diag
            bi = b
            for l in 1:n_dim
                Wi[l, l] = act.pieces[act_status[l]].slope
                bi[l] = act.pieces[act_status[l]].intercept
            end
            append!(lcs, backward_affine_map(lci, Wi, bi))
        end
    end
    lco = backward_affine_map(lcs, W, b)
    remove_redundant_constraints!(lco)
    return HPolyhedron(lco)
end

#ensure LinearConstraint{Float64, Vector{Float64}}
function get_constraints(input::HPolyhedron{Float64})
    C, d = tosimplehrep(input) #Cannot `convert` an object of type HalfSpace{Float64,LazySets.Arrays.SingleEntryVector{Float64}} to an object of type HalfSpace{Float64,Array{Float64,1}}
    lc = Vector{LinearConstraint{Float64, Vector{Float64}}}(undef, length(d))
    for i in 1:length(d)
        lc[i] = LinearConstraint(C[i,:], d[i])
    end
    return lc
end

function backward_affine_map(input::Vector{LinearConstraint{Float64, Vector{Float64}}}, M::Matrix{Float64}, v::Vector{Float64})
    output = Vector{LinearConstraint{Float64, Vector{Float64}}}(undef, 0)
    for lc in input
        vp = vprod(lc.a, M)
        if !iszero(vp)
            push!(output, LinearConstraint(vp, lc.b - lc.a' * v))
        end
    end
    return output
end

function vprod(x::Vector{Float64}, y::Matrix{Float64})
    @assert length(x) == size(y, 1) "Vector length $(length(x)) doesn't match Matrix size $(size(y))"
    n = size(y, 2)
    prod = zeros(Float64, n)
    for i in 1:n
        prod[i] = x' * y[:,i]
    end
    return prod
end


struct Gaussian{T}
    scaling::T
    width::T
    delay::T
end

Gaussian(;scaling=1.0, width, delay) = Gaussian(scaling, width, delay)

(g::Gaussian)(s::Real) = 4*g.scaling/(g.width*√π) * exp(-(4*(s-g.delay)/g.width)^2)

function creategaussian(width,s0,scaling=one(typeof(width)))
    #f(s) = 4*scaling/(width*sqrt(π)) * exp(-(4*(s-s0)/width)^2)
    Gaussian(scaling, width, s0)
    #f(s) = scaling * exp(-(4*(s-s0)/width)^2)
end

function  fouriertransform(g::Gaussian)
    scaling = g.scaling
    width = g.width
    s0 = g.delay
    ft(w) = scaling * exp(-im*w*s0 - (width*w/8)^2) / sqrt(2π)
end

function fouriertransform(a::Array, dt, t0, dim=1)
    n = size(a,dim)
    dω = 2π / (n*dt)
    b = fftshift(fft(a, dim), dim) * dt / sqrt(2π)
    ω0 = -dω * div(n,2)
    b, dω, ω0
end

fouriertransform(a::Array; stepsize, offset, dim=1) = fouriertransform(a, stepsize, offset, dim)

derive(g::Gaussian) =  s -> g(s) * (-8 * (s-g.delay)/g.width) * (4/g.width)

struct ErrorFunction{T}
    scaling::T
    width::T
    delay::T
end

function (f::ErrorFunction)(s)
    f.scaling * 0.5 * (1 + erf(4*(s-f.delay)/f.width))
end

function integrate(f::Gaussian)
    return ErrorFunction(f.scaling, f.width, f.delay)
end


####### Mass and temperature dependent biological rates

##### Parameters
# Create a NamedTuple containing the parameters required to calculate the mass/ temperature dependent biological rates

# define parameters
param = (
    # Intercepts
    I_r = -15.68,
    I_y = -9.66,
    I_B0 = 3.44,
    I_x = -16.54,
    I_αm = -1.81,
    I_thm = 1.92,
    I_thT = 0.5,

    # Resource slopes
    s_res_K = 0.28,
    s_res_r = -0.25,
    s_res_y = 0.45,
    s_res_B0 = 0.2,
    s_res_x = -0.31,

    # Consumer slopes (consumption)
    s_con_y = -0.47,
    s_con_B0 = 0.33,

    # Ratio slopes (attack and handling)
    s1_αm = 0.39,
    s1_thm = -0.48,
    s1_thT = -0.055,

    # Quadratic ratio slopes (attack and handling)
    s2_αm = -0.017,
    s2_thm = 0.0256,
    s2_thT = 0.0013,

    # Activation energies
    Ea_K = 0.71,
    Ea_r = -0.84,
    Ea_y = -0.26,
    Ea_B0 = 0.12,
    Ea_x = -0.69,

    # Assimilation efficiencies
    e = 0.85
)
#typeof(param)

#### Mass and temperature dependence
# write a function to calculate the Boltzmann term from the activation energy
function boltzmann(Ea, T)
    T0 = 293.15
    k = 8.617e-5
    normalised_T = T0 - T
    denom = k * T0 * T
    return exp(Ea *(normalised_T / denom))
end

# the following parameters are required for growth rate, carrying capacity and metabolism, which are only specific to one species at a time
    # I (exponent) - intercept
    # si - slope
    # Ea - activation energy

## Growth rate - r
# function to calculate the growth rate for the basal species
function growth_BA(param, M, T)
    # extract parameter values from param
    I = param.I_r
    s_res = param.s_res_r
    Ea = param.Ea_r
    # calculate terms in growth equation
    intercept = exp(I)
    mass = M .^ s_res
    boltz = boltzmann(Ea, T)
    growth = intercept .* mass .* boltz
    growth[2:3] .= 0  # only basal species has growth rate
    return growth
end

## Carrying capacity - K
# the intercept I is varied to investigate the effects of enrichment
function carryingcapacity_BA(I_K, param, M, T)
    # extract parameter values from param
    I = I_K
    s_res = param.s_res_K
    Ea = param.Ea_K
    # calculate terms in carrying capacity equation
    intercept = I
    mass = M .^ s_res
    boltz = boltzmann(Ea, T)
    carrycap = intercept .* mass .* boltz
    #carrycap[2:3] .= 0   # only basal species has carrying capacity
    return carrycap
end

## Metabolism - x
function metabolism_BA(param, M, T)
    # extract parameter values from param
    I = param.I_x
    s_res = param.s_res_x
    Ea = param.Ea_x
    # calculate terms in metabolism equation
    intercept = exp(I)
    mass = M .^ s_res
    boltz = boltzmann(Ea, T)
    metab = intercept .* mass .* boltz
    metab[1] = 0 # basal species has no metabolic rate
    return metab
end

# the following parameters are required for feeding terms maximum ingestion and half saturation density, which are resource and consumer specific
    # I (exponent)
    # si - resource slope
    # sj - consumer slope
    # Ea - activation energy

# Max ingestion and half sat density also require mass dependent attack rate and mass and temperature dependent handling time, which is consumer-resource body mass ratio dependent and needs:
    # I - intercept
    # s1 - simple Z coefficient
    # s2 - quadratic Z coefficient
    # Z

## Attack rate (mass dependence) αm
function attack_mass(param, Z)
    # extract parameter values from param
    I = param.I_αm
    s1 = param.s1_αm
    s2 = param.s2_αm
    # Z scaling
    Z_scale1 = s1 * log(Z)
    Z_scale2 = s2 * (log(Z))^2
    return exp(I + Z_scale1 + Z_scale2)
end

## Handling time (mass dependence) thm
function handling_mass(param, Z)
    # extract parameter values from param
    I = param.I_thm
    s1 = param.s1_thm
    s2 = param.s2_thm
    # Z scaling
    Z_scale1 = s1 * log(Z)
    Z_scale2 = s2 * (log(Z))^2
    return exp(I + Z_scale1 + Z_scale2)
end

## Handling time (temperature dependence) thT
function handling_temp(param, T)
    # Temperature conversion from K to C
    TC = T - 273.15
    # extract parameter values from param
    I = param.I_thT
    s1 = param.s1_thT
    s2 = param.s2_thT
    # temperature scaling
    T_scale1 = s1 * TC
    T_scale2 = s2 * TC^2
    return exp(I + T_scale1 + T_scale2)
end

## Maximum ingestion - y
function max_ingestion_BA(param, M, Z, T)
    # extract parameter values from param
    I = param.I_y
    s_res = param.s_res_y
    s_con = param.s_con_y
    Ea = param.Ea_y
    # calculate terms in max ingestion equation
    intercept = exp(I)
    mass_res = M .^ s_res
    mass_con = M .^ s_con
    boltz = boltzmann(Ea, T)
    th_m = handling_mass(param, Z)
    th_T = handling_temp(param, T)

    # create empty matrix for values
    mat = zeros(3,3)

    # intermediate eating basal
    mat[2,1] = intercept * mass_res[1] * mass_con[2] * boltz * 1/th_m * 1/th_T
    # top eating intermediate
    mat[3,2] = intercept * mass_res[2] * mass_con[3] * boltz * 1/th_m * 1/th_T

    return mat
end

## Half saturation density - β
function half_saturation_BA(param, M, Z, T)
    # extract parameter values from param
    I = param.I_B0
    s_res = param.s_res_B0
    s_con = param.s_con_B0
    Ea = param.Ea_B0
    # calculate terms in half saturation equation
    intercept = exp(I)
    mass_res = M .^ s_res
    mass_con = M .^ s_con
    boltz = boltzmann(Ea, T)
    α_m = attack_mass(param, Z)
    th_m = handling_mass(param, Z)
    th_T = handling_temp(param, T)

    # create empty matrix for values
    mat = zeros(3,3)

    # intermediate eating basal
    mat[2,1] = intercept * mass_res[1] * mass_con[2] * boltz * 1/(α_m * th_m * th_T)
    # top eating intermediate
    mat[3,2] = intercept * mass_res[2] * mass_con[3] * boltz * 1/(α_m * th_m * th_T)
    return mat
end


# p8.m - eigenvalues of harmonic oscillator -u"+x^2 u on R

L = 8;                             # domain is [-L L], periodic
for N = 6:6:36
    h = 2*pi/N; x = h*(1:N); x = L*(x-pi)/pi;
    column = [-pi^2/(3*h^2)-1/6; @. -.5*(-1)^(1:N-1)/sin(h*(1:N-1)/2)^2 ];
    D2 = (pi/L)^2*toeplitz(column);  # 2nd-order differentiation
    eigenvalues = sort(eigvals(-D2 + diagm(x.^2)));
    @show N
    [ println(eigenvalues[i]) for i=1:4 ];
    println("");
end


@defcomp PAGE09SeaLevelRise begin

  # Parameters

    rt_g_globaltemperature = Parameter(index=[time], unit="degreeC")
    sltemp_SLtemprise = Parameter(unit="m-degreeC", default=1.7333333333333334)
    sla_SLbaselinerise = Parameter(unit="m", default=1.00)
    sltau_SLresponsetime = Parameter(unit="years", default=1000.)
    s0_initialSL = Parameter(unit="m", default=0.15)
    y_year = Parameter(index=[time], unit="year")
    y_year_0 = Parameter(unit="year")

  # Variables

    es_equilibriumSL = Variable(index=[time], unit="m")
    s_sealevel = Variable(index=[time], unit="m")
    expfs_exponential = Variable(index=[time], unit="unitless")
    yp_timestep = Variable(index=[time], unit="years")

    function run_timestep(p, v, d, t)
        s0_initialSL_2009 = quantile(TriangularDist(0.1, 0.2, 0.15), cdf(TriangularDist(0.17, 0.21, 0.19), p.s0_initialSL))
        sltau_SLresponsetime_2009 = quantile(TriangularDist(500, 1500, 1000), cdf(Gamma(16.0833333333333333, 24.), p.sltau_SLresponsetime))

        if is_first(t)
            v.yp_timestep[t] = p.y_year[1] - p.y_year_0
            v.es_equilibriumSL[t] = p.sltemp_SLtemprise * p.rt_g_globaltemperature[t] + p.sla_SLbaselinerise
            v.expfs_exponential[t] = exp(-v.yp_timestep[t] / sltau_SLresponsetime_2009)
            v.s_sealevel[t] = s0_initialSL_2009 + (v.es_equilibriumSL[t] - s0_initialSL_2009) * (1 - v.expfs_exponential[t])

        else
            v.yp_timestep[t] = p.y_year[t] - p.y_year[t - 1]
            v.es_equilibriumSL[t] = p.sltemp_SLtemprise * p.rt_g_globaltemperature[t] + p.sla_SLbaselinerise
            v.expfs_exponential[t] = exp(-v.yp_timestep[t] / sltau_SLresponsetime_2009)
            v.s_sealevel[t] = v.s_sealevel[t - 1] + (v.es_equilibriumSL[t] - v.s_sealevel[t - 1]) * (1 - v.expfs_exponential[t])

        end
    end
end


"""
    batch_matvec(A::Tensor, X::Matrix)

for A ∈ ℝ^{m × d × n}, X ∈ ℝ^{d × n}. Performs ``n`` matrix-vec multiplications

    A[:,:,i] * X[:,i]

for ``i \\in 1,\\ldots,n``. This is performed via expansion of `X` and can be
efficiently evaluated using BLAS operations, which (I have assumed!) will often
be faster than a loop, especially in relation to AD and GPUs. Cf. `multitask.jl`
"""
batch_matvec(A::AbstractArray{T,3}, X::AbstractArray{T,2}) where T =
    dropdims(sum(A .* unsqueeze(X, 1), dims=2), dims=2)

"""
    batch_matmul(A::MatrixOrArray, X::Array)

for A ∈ ℝ^{n × m × d}, X ∈ ℝ^{m × r × d}. Performs ``n`` matrix multiplications

    A[:,:,i] * X[:,:,i]

for ``i \\in 1,\\ldots,n``. This is performed via expansion of `X` and can be
efficiently evaluated using BLAS operations, which (I have assumed!) will often
be faster than a loop in relation to AD and GPUs.

We also permit A ∈ ℝ^{n × m}, X ∈ ℝ^{m × r × d} => ℝ^{m × r × d}.
"""
batch_matmul(A::AbstractArray{T,3}, X::AbstractArray{T,3}) where T =
    dropdims(sum(unsqueeze(A,3) .* unsqueeze(X, 1), dims=2), dims=2)

batch_matmul(A::AbstractArray{T,2}, X::AbstractArray{T,3}) where T =
    dropdims(sum(unsqueeze(unsqueeze(A,3),4) .* unsqueeze(X, 1), dims=2), dims=2)

ensure_matrix(x::AbstractMatrix) = x
ensure_matrix(x::AbstractVector) = unsqueeze(x, 2)

ensure_tensor(x::AbstractMatrix) = unsqueeze(x, 3)
ensure_tensor(x::AbstractVector) = reshape(x, length(x), 1, 1)
ensure_tensor(x::AbstractArray{T,3}) where T = x

eye(d) = Matrix(I, d, d)

################################################################################
##                                                                            ##
##                         Special Matrix Constructors                        ##
##          *** THIS IS TAKEN DIRECTLY FROM MY JULIA UTILS PKG,               ##
##          BUT COPIED HERE FOR COMPLETENESS / TO SATISFY ANONYMITY ***       ##
##                                                                            ##
################################################################################

#= have a CuArray version in CUDA file. Cannot make more generally available
   because CuArrays.jl will not compile (e.g. on my machine) if no GPU. See
   @requires in [REDACTED].jl and CUDA.jl files in this project. =#

"""
    make_lt(x::AbstractVector{T}, d::Int)

Make lower triangular matrix (``d \\times d``) from vector `x`.
"""
function make_lt(x::AbstractVector{T}, d::Int)::Array{T,2} where T <: Real
    @argcheck (length(x) == Int(d*(d+1)/2))
    make_lt!(zeros(T, d, d), x, d)
end

function make_lt!(out::AbstractMatrix{T}, x::AbstractArray{T,1}, d::Int)::Array{T,2} where T <: Real
    x_i = 1
    for j=1:d, i=j:d
        out[i,j] = x[x_i]
        x_i += 1
    end
    return out
end

# Gradients
function unmake_lt(M::AbstractMatrix{T}, d)::Array{T,1} where T <: Real
    return M[tril!(trues(d,d))]
end

make_lt(x::TrackedArray, d::Int) = Tracker.track(make_lt, x, d)

@grad function make_lt(x, d::Int)
    return make_lt(Tracker.data(x), d), Δ -> (unmake_lt(Δ, d), nothing)
end

"""
    make_lt_strict(x::AbstractVector{T}, d::Int)

Make *strictly* lower triangular matrix (``d \\times d``) from vector `x`.
"""
function make_lt_strict(x::AbstractVector{T}, d::Int)::Array{T,2} where T <: Real
    @argcheck (length(x) == Int(d*(d-1)/2))
    return make_lt_strict!(zeros(T, d, d), x, d)
end

function make_lt_strict!(out::AbstractMatrix{T}, x::AbstractVector, d::Int)::Array{T,2} where T <: Real
    x_i = 1
    for j=1:d-1, i=j+1:d
        out[i,j] = x[x_i]
        x_i += 1
    end
    return out
end

make_lt_strict(x::TrackedArray, d::Int) = Tracker.track(make_lt_strict, x, d)

@grad function make_lt_strict(x, d::Int)
    return make_lt_strict(Tracker.data(x), d), Δ -> (unmake_lt_strict(Δ, d), nothing)
end

function unmake_lt_strict(M::AbstractMatrix{T}, d::Int)::Array{T,1} where T <: Real
    return M[tril!(trues(d,d), -1)]
end

"""
    make_skew(x::AbstractArray{T,1}, d::Int)

Make ``d \\times d`` skew-symmetric matrix from vector `x`.
"""
function make_skew(x::AbstractArray{T,1}, d::Int)::AbstractArray{T,2} where T <: Real
    S = make_lt_strict(x, d)
    return S - S'
end

function cayley_orthog(x::AbstractArray{T,1}, d)::AbstractArray{T,2} where T <: Real
    S = make_skew(x, d)
    I = eye(d)
    return (I - S) / (I + S)  # (I - S)^{-1}(I + S). Always nonsingular.
end

function inverse_cayley_orthog(Q::AbstractMatrix{T}, d::Int=size(Q,1)) where T <: Real
    unmake_lt_strict((I - Q ) / (I + Q), d)
end

"""
    diag0(x::Array{T,1})
Make diagonal matrix from vector `x`. This may have been superseded by `Diag`,
but is here for legacy reasons since Flux's version used to result in Tracked{Tracked}.
"""
function diag0(x::Array{T,1})::Array{T,2} where T <: Real
    d = length(x)
    M = zeros(T, d,d)
    M[diagind(M)] = x
    return M
end

diag0(x::TrackedArray) = Tracker.track(diag0, x)

@grad function diag0(x)
    return diag0(Tracker.data(x)), Δ -> (Δ[diagind(Δ)],)
end

"""
    LDSCell_diag_u{U,V,W}
Simple LDS struct containing three fields: a (transition coeffs), B (input
matrix) and h, the initial state.
"""
mutable struct LDSCell_diag_u{U,V,W}
    a::U
    B::V
    h::W
end

# Operation
function (m::LDSCell_diag_u)(h, x)
    h = m.a .* h + m.B * x
    return h, h
end

Flux.hidden(m::LDSCell_diag_u) = m.h
Flux.@treelike LDSCell_diag_u

"""
    LDSCell_diag_batch_u{U,V,W}
Simple LDS struct containing three fields: a (transition coeffs), B (input
matrix) and h, the initial state. Permits 3D arrays for a, B and performs
batch updates on the inputs.
"""
mutable struct LDSCell_diag_batch_u{U,V,W}
    a::U
    B::V
    h::W
end

# Operation
function (m::LDSCell_diag_batch_u)(h, x)
    h = m.a .* h + batch_matvec(m.B, x)
    return h, h
end

Flux.hidden(m::LDSCell_diag_batch_u) = m.h
Flux.@treelike LDSCell_diag_batch_u

"""
    LDSCell_simple_u{U,V,W}
Simple LDS struct containing three fields: A, B (for transition and input
matrices) and h, the initial state.
"""
mutable struct LDSCell_simple_u{U,V,W}
    A::U
    B::V
    h::W
end

# Operation
function (m::LDSCell_simple_u)(h, x)
    h = m.A * h + m.B * x
    return h, h
end

Flux.hidden(m::LDSCell_simple_u) = m.h
Flux.@treelike LDSCell_simple_u

"""
    LDSCell_batch_u{U,V,W}
Simple LDS struct containing three fields: A, B (for transition and input
matrices) and h, the initial state. Permits 3D arrays for A, B and performs
batch updates on the inputs.
"""
mutable struct LDSCell_batch_u{U,V,W}
    A::U
    B::V
    h::W
end

# Operation
function (m::LDSCell_batch_u)(h, x)
    h = batch_matvec(m.A, h) + batch_matvec(m.B, x)
    return h, h
end

Flux.hidden(m::LDSCell_batch_u) = m.h
Flux.@treelike LDSCell_batch_u

################################################################################
##                                                                            ##
##                         Model specific utilities                           ##
##                                                                            ##
################################################################################

"""
    A(ψ, d)
Construct transition matrix `A` using reduced Cayley form from parameter vector
ψ (extracting relevant elements).
"""
function A(ψ, d)
    n_skew = Int(d*(d-1)/2)
    x_S, x_V = ψ[1:d], ψ[d+1:d+n_skew]
    V = cayley_orthog(x_V/10, d)
    S = diag0(σ.(x_S))
    return S * V
end

"""
    B(ψ, d)
Construct transition matrix `B` using sparse construction from parameter vector
ψ (extracting relevant elements).
"""
function B(ψ::AbstractVector, d, n_u)
    n_skew = Int(d*(d-1)/2)
    i_init = d+n_skew
    ψ_b1 = reshape(ψ[(i_init+1):(i_init+n_u*d)], d, n_u)
    ψ_b2 = reshape(ψ[(i_init + n_u*d +1):(i_init + 2*n_u*d)], d, n_u)
    return σ.(ψ_b1) .* tanh.(ψ_b2)
end

function B(ψ::AbstractMatrix, d, n_u)
    n_skew = Int(d*(d-1)/2)
    n_b = size(ψ, 2)
    i_init = d+n_skew
    ψ_b1 = reshape(ψ[(i_init+1):(i_init+n_u*d), :], d, n_u, n_b)
    ψ_b2 = reshape(ψ[(i_init + n_u*d +1):(i_init + 2*n_u*d), :], d, n_u, n_b)
    return σ.(ψ_b1) .* tanh.(ψ_b2)
end

"""
    C(ψ, d)
Construct emission matrix `C` with no special constructor from parameter vector
ψ (extracting relevant elements).
"""
function C(ψ::AbstractVector, d, n_u, n_out)
    n_skew = Int(d*(d-1)/2)
    i_init = d + n_skew + 2*n_u*d
    ψ_c= ψ[(i_init+1):(i_init+n_out*d)]
    return reshape(ψ_c, n_out, d)
end

function C(ψ::AbstractMatrix, d, n_u, n_out)
    n_skew = Int(d*(d-1)/2)
    n_b = size(ψ, 2)
    i_init = d + n_skew + 2*n_u*d
    ψ_c= ψ[(i_init+1):(i_init+n_out*d), :]
    return reshape(ψ_c, n_out, d, n_b)
end

"""
    dmt(ψ, d)
Construct emission bias `d` with no special constructor from parameter vector
ψ (extracting relevant elements).
"""
function dmt(ψ::AbstractVector, d, n_u, n_out)
    n_skew = Int(d*(d-1)/2)
    i_init = d + n_skew + 2*n_u*d +n_out*d
    return ψ[(i_init+1):(i_init+n_out)]
end

function dmt(ψ::AbstractMatrix, d, n_u, n_out)
    n_skew = Int(d*(d-1)/2)
    i_init = d + n_skew + 2*n_u*d +n_out*d
    return ψ[(i_init+1):(i_init+n_out), :]
end


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigInt()
    l = Exo1Binion(Exo1Real(a, b))
    r = Exo1Binion(a, 0, b, 0)
    l == r
end

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(Exo1Binion(a))
end

@test begin
    a = 1
    !isreal(Exo1Binion(a, a, a, a))
end

@test begin
    z = random(Exo1Binion{Int})
    z == +(z)
end

@test begin
    a = 1
    real(Exo1Binion(a)) == a
end

@test begin
    a = 1
    b = 2
    c = 3
    l = unreal(Exo1Binion(0, a, b, c))
    r = [a, b, c]
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo1Binion(1, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1, i: 2, a: 3, ia: 4]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    l = Exo1Binion(Binion(a, b))
    r = Exo1Binion(a, b)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    l = Exo1Binion(a, b, c)
    r = Exo1Binion(a, b, c, 0)
    l == r
end

@test begin
    x = random(Exo1Binion{BigFloat})
    y = random(Exo1Binion{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo1Binion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo1Binion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = associator(x, y, z)
    r = zero(Exo1Binion{BigInt})
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(Exo1Binion{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

@test begin
    x = random(Exo1Binion{Rational{BigInt}})
    y = random(Exo1Binion{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Exo1Binion{Rational{BigInt}})
    y = random(Exo1Binion{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Exo1Binion{Rational{BigInt}})
    y = random(Exo1Binion{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo1Binion{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(Binion{BigInt})
    x = random(Exo1Binion{BigInt})
    l = a * x
    r = x * a
    l != r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    z = random(Exo1Binion{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Exo1Binion{BigInt})
    quadrance(x) > big(0) || quadrance(x) == big(0)
end

@test begin
    x = random(Exo1Binion{BigInt})
    y = random(Exo1Binion{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(Exo1Binion(0, 0.0))
end

@test_throws ErrorException begin
    inv(Exo1Binion(0, 0, 1, 1))
end

@test_throws ErrorException begin
    inv(Exo1Binion(0, 0, 2, 1))
end

@test_throws ErrorException begin
    random(Exo1Binion{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Exo1Binion{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo1Binion{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo1Binion{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == z + zero(Exo1Binion{BigInt})
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == z * one(Exo1Binion{BigInt})
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == z * one(z)
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == cloak(cloak(z))
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == dagger(dagger(z))
end

@test begin
    z = random(Exo1Binion{BigInt})
    z == star(star(z))
end

@test begin
    z = random(Exo1Binion{BigInt})
    selfstar(z) == selfstar(selfstar(z))
end

@test begin
    z = random(Exo1Binion{BigInt})
    anti_selfstar(z) == anti_selfstar(anti_selfstar(z))
end


show(io::IO, surr::T) where {T<:AbstractSurrogate} =
    print(io, "$(nameof(T))($(surr.ϑ))")

@with_kw_noshow struct Hinge{T<:Real} <: AbstractSurrogate
    ϑ::T = 1.0
    value::Function       = (s) -> max(0, 1 + ϑ*s)
    value_exact::Function = (s) -> Convex.max(0, 1 + ϑ*s)
    gradient::Function    = (s) -> 1 + ϑ*s <= 0 ? zero(ϑ) : ϑ
end

Hinge(ϑ) = Hinge(ϑ = ϑ)

@with_kw_noshow struct Quadratic{T<:Real} <: AbstractSurrogate
    ϑ::T = 1.0
    value::Function       = (s) -> max(0, 1 + ϑ*s)^2
    value_exact::Function = (s) -> Convex.square(Convex.max(0, 1 + ϑ*s))
    gradient::Function    = (s) -> 2*ϑ*max(0, 1 + ϑ*s)
end

Quadratic(ϑ) = Quadratic(ϑ = ϑ)


mutable struct AugmentedLagrangianCosts <: Objective
    costs::StageCosts
    cons::Constraints
    ρ::Vector  # penalty
    λ::Vector  # dual estimates
    a::Vector  # active set
end

function augmented_lagrangian(costs::StageCosts, cons::Constraints)
    λ = [zeros(cons.con[t].p) for t = 1:cons.T]
    a = [ones(cons.con[t].p) for t = 1:cons.T]
    ρ = [ones(cons.con[t].p) for t = 1:cons.T]
    AugmentedLagrangianCosts(costs, cons, ρ, λ, a)
end

function objective(obj::AugmentedLagrangianCosts, x, u)
    # costs
    J = objective(obj.costs, x, u)

    # constraints
    T = obj.cons.T
    c = obj.cons.data.c
    ρ = obj.ρ
    λ = obj.λ
    a = obj.a

    constraints!(obj.cons, x, u)
    active_set!(a, obj.cons, λ)

    for t = 1:T
        J += λ[t]' * c[t]
        J += 0.5 * c[t]' * Diagonal(ρ[t] .* a[t]) * c[t]
    end

    return J
end

function active_set!(a, cons::StageConstraints, λ)
    T = cons.T
    c = cons.data.c

    for t = 1:T
        # set all constraints active
        fill!(a[t], 1.0)

        # find inequality constraints
        if haskey(cons.con[t].info, :inequality)
            for i in cons.con[t].info[:inequality]
                # check active-set criteria
                (c[t][i] < 0.0 && λ[t][i] == 0.0) && (a[t][i] = 0.0)
            end
        end
    end
end

function augmented_lagrangian_update!(obj::AugmentedLagrangianCosts;
        s = 10.0, max_penalty = 1.0e12)
    # constraints
    T = obj.cons.T
    c = obj.cons.data.c
    ρ = obj.ρ
    λ = obj.λ
    a = obj.a

    for t = 1:T
        # dual estimate update
        λ[t] .+= ρ[t] .* c[t]

        # inequality projection
        if haskey(obj.cons.con[t].info, :inequality)
            idx = obj.cons.con[t].info[:inequality]
            λ[t][idx] = max.(0.0, view(λ[t], idx))
        end
        ρ[t] .= min.(s .* ρ[t], max_penalty)
    end
end


module TestSpatialDiscretization

using Modia
using ModiaMath:plot

const n = 5
@model SpatialDiscretization begin
    u = Var(start=zeros(n))
    c = 1
    @equations begin
        # der(u[2:n]) = ( u[2:n] - u[1:n-1] ) / c
        der(u) = [1; ( u[2:n] - u[1:n - 1] ) / c]
    end
end

result = simulate(SpatialDiscretization, 1)
plot(result, "u")

splice(uodd, ueven) = [if isodd(i); uodd[div(i,2)+1] else ueven[div(i,2)] end for i in 1:length(ueven)+length(uodd)]

@model SpatialDiscretization2 begin
    u = Var(size=(5,))
    ueven = Var(start=[1,2])
    uodd = Var(size=(3,))
    @equations begin
      u = splice(uodd, ueven)
      der(ueven) = - u[[2,4]] + u[[1,3]]
      uodd = ones(3)*time
    end
end

result = simulate(SpatialDiscretization2, 1, logTranslation=true, storeEliminated=false)
plot(result, "ueven")

end


using SparseArrays
using LinearAlgebra
using SparseMatricesCSR
using Test

include("SparseMatrixCSC.jl")
include("SparseMatrixCSR.jl")
include("SymSparseMatrixCSR.jl")


# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools

function main()
    # Define beam
    λ = LinRange(200,5000,1000) # wavelength range [nm]
    θ = LinRange(0,90,90)
    beam = PlaneWave(λ,θ)
    # Define layers
    l0 = LayerTMMO(RIdb.air(beam.λ))
    l1 = LayerTMMO(RIdb.dummy(beam.λ,1.45,0.0); type=:OT, d=1/4.)
    l2 = LayerTMMO(RIdb.dummy(beam.λ,3.45,0.0); type=:OT, d=1/4.)
    layers = [l0, l1, l2, l1, l2, l1, l2, l1, l2, l0]
    # Reference wavelenth
    λ0 = 700.0
    # call main script
    sol = tmm_optics(beam, layers; λ0=λ0, emfflag=true, h=10, pbgflag=true)
    return sol
end

function plotEMF(sol, ϕ)
    ϕ_ = findmin(abs.(sol.beam.θ .- ϕ))[2][1]
    plot(EMF2D(),
              sol.beam.λ, sol.Misc.ℓ, log10.(sol.Field.emfp[:,ϕ_,:]),
              title="Log EMF intesnsity, p-TM",
    )
    gui()
    plot(EMF2D(),
              sol.beam.λ, sol.Misc.ℓ, log10.(sol.Field.emfs[:,ϕ_,:]),
              title="Log EMF intesnsity, s-TE",
    )
    gui()
    return nothing
end

sol = main()

# plot the R, T and A spectra
plot(Spectrum2D(),
        sol.beam.λ, sol.beam.θ, sol.Spectra.Rp,
        title="Reflectance, p-TM", clims=(0.0, 1.0),
)
gui()
plot(Spectrum2D(),
        sol.beam.λ, sol.beam.θ, sol.Spectra.Rs,
        title="Reflectance, s-TE", clims=(0.0, 1.0),
)
gui()

# Plot the EMF at normal incidence
plotEMF(sol, 0.0)

# Plot the EMF at 15 degrees
plotEMF(sol, 15.0)

# plot the refractive index profile
plot(RIprofile(), sol; plotemf=true)
gui()
### Note that if you choose θ outside the range it will show the EMF for one extrema. Same for λ.

# plot the photonic dispersion with custom function
plot(PBGDispersion2D(), sol.Bloch; wave=:p)
gui()

plot(PBGDispersion2D(), sol.Bloch; wave=:s)
gui()

plot(PBGDispersion2Dalt(), sol.Bloch)
gui()


# This file contains algorithms needed for the solution of
# the Risch differential equation from chapter 6 of the book
#
#    Manuel Bronstein, Symbolic Integration I: Transcendental Functions,
#    2nd ed., Springer 2005.
#

"""
    WeakNormalizer(f, D) -> q

Weak normalization.

Given a field `k`, a derivation `D` on `k[t]` and `f` in `k(t)`, return
`q` in `k[t]` such that `f-D(q)/q` is weakly normalized with respect to D.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.1, p. 183.
"""
function WeakNormalizer(f::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(f, D) || error("rational function f must be in the domain of derivation D")
    dn, ds = SplitFactor(denominator(f), D)
    g = gcd(dn, derivative(dn))
    dstar = divexact(dn, g)
    d1 = divexact(dstar, gcd(dstar, g))
    a, b = gcdx(divexact(denominator(f), d1), d1, numerator(f))
    kz, z = PolynomialRing(base_ring(a), :ζ)
    kzt, t = PolynomialRing(kz, var(parent(a)))
    dd1 = D(d1)
    r = resultant(d1(t), a(t)-z*dd1(t))
    if iszero(r)
        return one(parent(a))
    end
    D1 = CoefficientLiftingDerivation(kz, BaseDerivation(D)) # dummy derivation compatible with r
    ns = [numerator(rationalize_over_Int(x)) for x in constant_roots(r, D1) if
            isrational(x) && isone(denominator(x)) && x>0]
    if isempty(ns)
        return one(numerator(f))
    end
    prod([gcd(a - n*dd1, d1)^n for n in ns])
end

"""
    RdeNormalDenominator(f, g, D) -> (a, b, c, h, ρ)

Normal part of the denomninator.

Given a field `k`, a derivation `D` on `k[t]`, `f`, `g` in `k(t)`
with `f` weakly normalized with respect to `t`, return
either `ρ=0`, in which case the equation `D(y)+f*y=g` has no solution in `k⟨t⟩`,
or `ρ=1`, `a`, `h` in `k[t]` and `b`, `c` in `k⟨t⟩`, such that for any  solution `y`
in `k(t)` of `D(y)+f*y=g`, `q=y*h` in `k⟨t⟩` satisfies `a*D(q)+b*q=c`.

(Here, `k⟨t⟩` denotes the elements of `k(t)` which are reduced w.r.t. `D`.)

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.1, p. 185.
"""
function RdeNormalDenominator(f::F, g::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(f, D) && iscompatible(g, D) ||
        error("rational functions f and g must be in the domain of derivation D")
    # Note: f must be weakly normalized which we do not check. It is recommended
    # to use this function only for rational functions f which were computed by WeakNormalizer.
    (dn, ds) = SplitFactor(denominator(f), D)
    (en, es) = SplitFactor(denominator(g), D)
    p = gcd(dn, en)
    h = divexact(gcd(en, derivative(en)), gcd(p, derivative(p)))
    if !iszero(rem(dn*h^2, en))
        return zero(h), zero(f), zero(f), zero(h), 0 # no solution
    end
    dn*h, dn*h*f-dn*D(h), dn*h^2*g, h, 1
end

function one_times_n_solve(A::Vector{T}, B::Vector{T}) where T<:FieldElement
    if length(A)!=length(B) || length(A)==0
        return T[] # no solution
    end
    eqfound = false
    c = 0
    for i=1:length(A)
        a = A[i]
        b = B[i]
        if iszero(b)
            if !iszero(a)
                return T[] # no solution
            end
        elseif eqfound
            if c!=a//b
                return T[] # no solution
            end
        else
            c = a//b
            eqfound = true
        end
    end
    eqfound ? [c] : T[]
end

"""
    ParametricLogarithmicDerivative(f, w, D) -> (n, m, v, ρ)

Parametric logarithmic derivative heuristic.

Given a field `k`, a derivation `D` on `k[t]`, `f` in `k(t)` and `w=D(θ)/θ` in `k` such that `θ`
is a hyperexponential monomial over `k(t)`, return either `ρ=0`, in which case
`n*f=D(v)/v+m*D(θ)/θ` has no solution `v≠0` in `k(t)` and
integers `n≠0` and `m`, or `ρ=1` and a solution `n`, `m`, `v` of that equation.

This function may throw exception `AlgorithmFailedError`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 7.3, p. 253.
"""
function ParametricLogarithmicDerivative(f::F, w::F, D::Derivation) where F<:FieldElement
    # base case f,w \in constant field, D must be the null derivation
    iszero(D) || error("base case only for null derivations")
    v = one(parent(f))
    q = f//w
    if !isrational(q) # could be complex for example
        return 0, 0, zero(parent(f)), 0
    end
    q = rationalize_over_Int(f//w)
    m = numerator(q)
    n = denominator(q)
    n, m, v, 1
end

function ParametricLogarithmicDerivative(f::F, w::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    # See Bronstein's book, Section 7.3, p. 253
    Z = zero(f)
    no_solution = (0, 0, Z, 0)
    if degree(numerator(f))<=0 && degree(denominator(f))<=0 &&
        degree(numerator(w))<=0 && degree(denominator(w))<=0
        #f, w \in k(t) are actually in k, so search solution not in k(t) but in k
        n, m, v, ρ = ParametricLogarithmicDerivative(
            constant_coefficient(numerator(f)),
            constant_coefficient(numerator(w)),
            BaseDerivation(D))
        return n, m, v + Z, ρ # make v \in k an element of k(t)
    end
    d = denominator(f)
    e = denominator(w)
    p, a = divrem(numerator(f), d)
    q, b = divrem(numerator(w), e)
    t = gen(parent(d))
    B = max(0, degree(D(t))-1)
    C = max(degree(q), degree(p))
    if degree(q)>B
        ss = one_times_n_solve([coeff(p, i) for i=B+1:C], [coeff(q, i) for i=B+1:C])
        if isempty(ss)
            return no_solution
        end
        if !isrational(ss[1])
            return no_solution
        end
        s = rationalize_over_Int(ss[1])
        M = numerator(s)  # In Bronstein's book, M and N are switched, this is the first and so far only bug I found there
        N = denominator(s)
        Q, v, ρ = InFieldLogarithmicDerivativeOfRadical(N*f - M*w, D)
        if ρ>=1 && !iszero(Q) && !iszero(v)
            return Q*N, Q*M, v, 1
        else
            return no_solution
        end
    end
    if degree(p)>B
        return no_solution
    end
    l = lcm(d, e)
    ln, ls = SplitFactor(l, D)
    z = ls*gcd(ln, derivative(ln))
    if degree(z)<=0 # z \in k
        throw(AlgorithmFailedError("ParametricLogarithmicDerivative\n@ $(@__FILE__):$(@__LINE__)"))
    end
    u1, r1 = divrem(numerator(l*f), z)
    u2, r2 = divrem(numerator(l*w), z)
    ss = one_times_n_solve([coeff(r1, i) for i=0:(degree(z)-1)], [coeff(r2, i) for i=0:(degree(z)-1)])
    if isempty(ss)
        return no_solution
    end
    if !isrational(ss[1])
        return no_solution
    end
    s = rationalize_over_Int(ss[1])
    M = numerator(s)
    N = denominator(s)
    Q, v, ρ = InFieldLogarithmicDerivativeOfRadical(N*f - M*w, D)
    if ρ>=1 && !iszero(Q) && !iszero(v)
        return Q*N, Q*M, v, 1
    else
        return no_solution
    end
end

"""
    RdeSpecialDenomExp(a, b, c, D) -> (A, B, C, h)

Special part of the denominator - hyperexponential case.

Given a field `k`, a derivation `D` on `k[t]` with `D(t)/t` in `k`,
`a≠0` in `k[t]` with `gcd(a,t)=1`, and `b`, `c` in `k⟨t⟩`
return  `A`, `B`, `C`, `h` in `k[t]` such that
for any solution `q` in `k⟨t⟩` of `a*D(q)+b*q=c`, `r=q*h` in `k[t]` satisfies `A*D(r)+B*r=C`.

(Here, `k⟨t⟩` denotes the elements of `k(t)` which are reduced w.r.t. `D`.)

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.2, p. 190.
"""
function RdeSpecialDenomExp(a::P, b::F, c::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    !iszero(a) || error("a must be != 0")
    ishyperexponential(D) ||
        error("monomial of derivation D must be hyperexponential")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomial a and rational functions b and c must be in the domain of derivation D")
    isreduced(b, D) && isreduced(c, D) ||
        error("rational functions b and c must be reduced with respect to derivation D")
    t = gen(parent(a))
    degree(gcd(t, a))==0 || error("gcd(a, t) must be == 1")
    p = t
    nb = iszero(b) ? typemax(Int) : valuation(b, p)
    nc = iszero(c) ? typemax(Int) : valuation(c, p)
    n = min(0, nc - min(0, nb))
    if nb==0
        α = constant_coefficient(Remainder(-b//a, p))
        w = coeff(MonomialDerivative(D), 1)
        n0, m, z, ρ = ParametricLogarithmicDerivative(α, w, BaseDerivation(D))
        if  ρ>0 && n0==1 && !iszero(z)
            n = min(n, m)
        end
    end
    N = max(0, -nb, n-nc)
    p_power_N = p^N
    b1 = (b+n*a*divexact(D(p), p))*p_power_N
    @assert isone(denominator(b1))
    c1 = c*p^(N-n)
    @assert isone(denominator(c1))
    a*p_power_N, numerator(b1), numerator(c1), p^(-n)
end

"""
    RdeSpecialDenomTan(a, b, c, D) -> (A, B, C, h)

Special part of the denominator - hypertangent case.

Given a field `k` not containing `√-1`, a derivation `D` on `k[t]` with `D(t)/(t^2+1)` in `k`,
`a≠0` in `k[t]` with `gcd(a,t^2+1)=1`, and `b`, `c` in `k⟨t⟩`,
return  `A`, `B`, `C`, `h` in `k[t]` such that
for any solution `q` in `k⟨t⟩` of `a*D(q)+b*q=c`, `r=q*h` in `k[t]` satisfies `A*D(r)+B*r=C`.

(Here, `k⟨t⟩` denotes the elements of `k(t)` which are reduced w.r.t. `D`.)

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.2, p. 192.
"""
function RdeSpecialDenomTan(a::P, b::F, c::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    !iszero(a) || error("a must be != 0")
    ishypertangent(D) ||
        error("monomial of derivation D must be hypertangent")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomial a and rational functions b and c must be in the domain of derivation D")
    isreduced(b, D) && isreduced(c, D) ||
        error("rational functions b and c must be reduced with respect to derivation D")
    t = gen(parent(a))
    p = t^2+1
    degree(gcd(a, p))==0 || error("gcd(a, t^2+1) must be == 1")
    nb = iszero(b) ? typemax(Int) : valuation(b, p)
    nc = iszero(c) ? typemax(Int) : valuation(c, p)
    n = min(0, nc - min(0, nb))
    if nb==0
        αI_plus_β = Remainder(-b//a, p)
        α = coeff(αI_plus_β, 1)
        β = coeff(αI_plus_β, 0)
        η = constant_coefficient(divexact(MonomialDerivative(D), p))
        D0 = BaseDerivation(D)
        v, ρ = InFieldDerivative(2*β, D0)
        if ρ>0 && !iszero(v)
            _, I, D0I = Complexify(parent(η), D0)
            n0, m, v, ρ = ParametricLogarithmicDerivative(α*I+β, 2*η*I, D0I)
            if  ρ>0 && n0==1 && !iszero(v)
                n = min(n, m)
            end
        end
    end
    N = max(0, -nb, n-nc)
    p_power_N = p^N
    b1 = (b+n*a*divexact(D(p), p))*p_power_N
    @assert isone(denominator(b1))
    c1 = c*p^(N-n)
    @assert isone(denominator(c1))
    a*p_power_N, numerator(b1), numerator(c1), p^(-n)
end

"""
    RdeSpecialDenomTanI(a, b, c, D) -> (A, B, C, h)

Special part of the denominator - hypertangent case, field contains I=√-1.

Given a field `k` containing `√-1`, a derivation `D` on `k[t]` with `D(t)/(t^2+1)` in `k`,
`a≠0` in `k[t]` with `gcd(a,t^2+1)=1`, and `b`, `c` in `k⟨t⟩`,
return  `A`, `B`, `C`, `h` in `k[t]` such that
for any solution `q` in `k⟨t⟩` of `a*D(q)+b*q=c`, `r=q*h` in `k[t]` satisfies `A*D(r)+B*r=C`.

(Here, `k⟨t⟩` denotes the elements of `k(t)` which are reduced w.r.t. `D`.)

The implementation is  analogous to `RdeSpecialDenomPrim` and `RdeSpecialDenomTan`,
but now with `p=t-√-1` instead of `p=t` resp. `p=t^2+1`,
see [Bronstein's book](https://link.springer.com/book/10.1007/b138171), last sentence
before algorithm `RdeSpecialDenomTan` in Section 6.2, p. 192 and Exercise 6.1, p. 216.
"""
function RdeSpecialDenomTanI(a::P, b::F, c::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    !iszero(a) || error("a must be != 0")
    ishypertangent(D) ||
        error("monomial of derivation D must be hypertangent")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomial a and rational functions b and c must be in the domain of derivation D")
    isreduced(b, D) && isreduced(c, D) ||
        error("rational functions b and c must be reduced with respect to derivation D")
    t = gen(parent(a))
    degree(gcd(a, t^2 + 1))==0 || error("gcd(a, t^2+1) must be == 1")
    contains_I(parent(a)) || error("field k must contain I=sqrt(-1)")
    I = get_I(base_ring(parent(a)))
    p = t-I
    nb = iszero(b) ? typemax(Int) : valuation(b, p)
    nc = iszero(c) ? typemax(Int) : valuation(c, p)
    n = min(0, nc - min(0, nb))
    if nb==0
        α = constant_coefficient(Remainder(-b//a, p))
        η = constant_coefficient(divexact(MonomialDerivative(D), t^2 + 1))
        n0, m, v, ρ = ParametricLogarithmicDerivative(α, 2*η*I, BaseDerivation(D))
        if  ρ>0 && n0==1 && !iszero(v)
            n = min(n, m)
        end
    end
    N = max(0, -nb, n-nc)
    pp = t^2 + 1
    pp_power_N = pp^N
    b1 = (b+n*a*divexact(D(pp), pp))*pp_power_N
    @assert isone(denominator(b1))
    c1 = c*pp^(N-n)
    @assert isone(denominator(c1))
    a*pp_power_N, numerator(b1), numerator(c1), pp^(-n)
end

"""
    RdeSpecialDenominator(a, b, c, D) -> (A, B, C, h)

Special part of the denominator.

Given a field `k`, a derivation `D` on `k[t]`, `a≠0` in `k[t]`, and `b`, `c` in `k⟨t⟩`,
return  `A`, `B`, `C`, `h` in `k[t]` such that
for any solution `q` in `k⟨t⟩` of `a*D(q)+b*q=c`, `r=q*h` in `k[t]` satisfies `A*D(r)+B*r=C`.

(Here, `k⟨t⟩` denotes the elements of `k(t)` which are reduced w.r.t. `D`.)

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.2, p. 186.
"""
function RdeSpecialDenominator(a::P, b::F, c::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    !iszero(a) || error("a must be != 0")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomial a and rational functions b and c must be in the domain of derivation D")
    isreduced(b, D) && isreduced(c, D) ||
        error("rational functions b and c must be reduced with respect to derivation D")
    if isprimitive(D)
        @assert isone(denominator(b))
        @assert isone(denominator(c))
        return a, numerator(b), numerator(c), one(b)
    elseif ishyperexponential(D)
        t = gen(parent(a))
        d = gcd(a, t)
        a = divexact(a, d)
        b = b//d
        c = c//d
        return  RdeSpecialDenomExp(a, b, c, D)
    elseif ishypertangent(D)
        t = gen(parent(a))
        p = t^2 + 1
        d = gcd(a, p)
        a = divexact(a, d)
        b = b//d
        c = c//d
        if contains_I(parent(a))
            return RdeSpecialDenomTanI(a, b, c, D) #case √-1 in k
        else
            return RdeSpecialDenomTan(a, b, c, D)
        end
    else
        H = MonomialDerivative(D)
        throw(NotImplementedError("RdeSpecialDenominator: monomial derivative $H\n@ $(@__FILE__):$(@__LINE__)"))
    end
end

"""
    RdeBoundDegreePrim(a, b, c, D) -> n

Bound on polynomial solutions - primitive case.

Given a field `k`, a derivation `D` on `k[t]` with `D(t)` in `k`, and `a`, `b`, `c` in `k[t]`
with `a≠0`, return integer `n` such that `degree(q)≤n` for any solution `q` in `k[t]` of `a*D(q)+b*q=c`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.3, p. 198.
"""
function RdeBoundDegreePrim(a::P, b::P, c::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    isprimitive(D) ||
        error("monomial of derivation D must be primitive")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials a, b, and c must be in the domain of derivation D")
    !iszero(a) || error("polynomial a must be nonzero")
    da = degree(a)
    db = degree(b)
    dc = degree(c)
    if db>da
        n = max(0, dc-db)
    else
        n = max(0, dc-da+1)
    end
    if db==da-1
        α = -leading_coefficient(b)//leading_coefficient(a)
        z, m0, ρ = LimitedIntegrate(α, leading_coefficient(D), BaseDerivative(D))
        if ρ>0 && is_rational(m0)
            m = rationalize_over_Int(m0)
            if denominator(m)==1
                n = max(n, numerator(m))
            end
        end
    end
    D0 = BaseDerivation(D)
    if db==da
        α = -leading_coefficient(b)//leading_coefficient(a)
        z, ρ = InFieldDerivative(α, D0)
        if ρ>0 && !iszero(z)
            β = -leading_coefficient(a*D0(z)+b*z)//(z*leading_coefficient(a))
            w, m0, ρ = LimitedIntegrate(β, leading_coefficient(D), D0)
            if ρ>0 && isrational(m0)
                m = rationalize_over_Int(m0)
                if denominator(m)==1
                    n = max(n, numerator(m))
                end
            end
        end
    end
    n
end

"""
    RdeBoundDegreeBase(a, b, c) -> n

Bound on polynomial solutions - base case.

Given a field `k` and  `a`, `b`, `c` in `k[t]` with `a≠0`, return integer `n`
such that `degree(q)≤n` for any solution `q` in `k[t]` of `a*(d/dt)(q)+b*q=c`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.3, p. 199.
"""
function RdeBoundDegreeBase(a::P, b::P, c::P) where
    {T<:FieldElement, P<:PolyElem{T}}
    !iszero(a) || error("polynomial a must be nonzero")
    da = degree(a)
    db = degree(b)
    dc = degree(c)
    n = max(0, dc - max(db, da - 1))
    if db==da-1
        m0 = -leading_coefficient(b)//leading_coefficient(a)
        if isrational(m0)
            m = rationalize_over_Int(m0)
            if isone(denominator(m))
                n = max(0, numerator(m), dc - db)
            end
        end
    end
    return n
end

"""
    RdeBoundDegreeExp(a, b, c, D) -> n

Bound on polynomial solutions - hyperexponential case.

Given a field `k`, a derivation `D` on `k[t]` with `D(t)/t` in `k`, and `a`, `b`, `c` in `k[t]`
with`a≠0`, return integer `n` such that `degree(q)≤n` for any solution `q` in `k[t]` of `a*D(q)+b*q=c`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.3, p. 200.
"""
function RdeBoundDegreeExp(a::P, b::P, c::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    ishyperexponential(D) ||
        error("monomial of derivation D must be hyperexponential")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials a, b, and c must be in the domain of derivation D")
    !iszero(a) || error("polynomial a must be nonzero")
    da = degree(a)
    db = degree(b)
    dc = degree(c)
    n = max(0, dc - max(db, da))
    if da==db
        α = -leading_coefficient(b)//leading_coefficient(a)
        w = coeff(MonomialDerivative(D), 1)
        n0, m, z, ρ = ParametricLogarithmicDerivative(α, w, BaseDerivation(D))
        if  ρ>0 && n0==1 && !iszero(z)
            n = max(n, m)
        end
    end
    n
end

"""
    RdeBoundDegreeNonLinear(a, b, c, D) -> n

Bound on polynomial solutions - nonlinear case.

Given a field `k`, a derivation `D` on `k[t]` with `degree(D(t))≥2`, and `a`, `b`, `c` in `k[t]`
with `a≠0`, return integer `n` such that `degree(q)≤n` for any solution `q` in `k[t]` of `a*D(q)+b*q=c`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.3, p. 201.
"""
function RdeBoundDegreeNonLinear(a::P, b::P, c::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    isnonlinear(D) ||
        error("monomial of derivation D must be nonlinear")
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials a, b, and c must be in the domain of derivation D")
    !iszero(a) || error("polynomial a must be nonzero")
    da = degree(a)
    db = degree(b)
    dc = degree(c)
    δ = degree(D)
    λ = leading_coefficient(D)
    n = max(0, dc - max(da + δ - 1, db))
    if db==da+δ-1
        m0 = -leading_coefficient(b)//(λ*leading_coefficient(a))
        if isrational(m0)
            m = rationalize_over_Int(m0)
            if isone(denominator(m))
                n = max(0, numerator(m), dc - db)
            end
        end
    end
    n
end

"""
    RdeBoundDegree(a, b, c, D) -> n

Bound on polynomial solutions.

Given a field `k`, a derivation `D` on `k[t]` and `a`, `b`, `c` in `k[t]`
with `a≠0`, return integer `n` such that `degree(q)≤n`
for any solution `q` in `k[t]` of `a*D(q)+b*q=c`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.3, p. 193.
"""
function RdeBoundDegree(a::P, b::P, c::P, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}}
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials a, b, and c must be in the domain of derivation D")
    !iszero(a) || error("polynomial a must be nonzero")
    if isbasic(D) || (isprimitive(D) && isone(MonomialDerivative(D)))
        return RdeBoundDegreeBase(a, b, c)
    elseif isprimitive(D)
        return RdeBoundDegreePrim(a, b, c, D)
    elseif ishyperexponential(D)
        return RdeBoundDegreeExp(a, b, c, D)
    elseif isnonlinear(D)
        return RdeBoundDegreeNonLinear(a, b, c, D)
    else
        H = MonomialDerivative(D)
        throw(NotImplementedError("RdeBoundDegree: monomial derivative $H\n@ $(@__FILE__):$(@__LINE__)"))
    end
end

"""
    SPDE(a, b, c, D, n) -> (B, C, m, α, β, ρ)

Rothstein's Special Polynomial Differential Equation algorithm.

Given a field `k`, a derivation `D` on `k[t]`, an integer `n`
and `a`, `b`, `c` in `k[t]`  with `a≠0`, return either
`ρ=0`, in which case the equation `a*D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1`, `B`, `C`, `α`, `β` in `k[t]` and an integer `m`, such that any solution `q` in `k[t]`
of `a*D(q)+b*q=c` must be of the form `q=α*h+β`, where `h` is in `k[t]`, `degree(h)≤m` and
`D(h)+B*h=C`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.4, p. 204.
"""
function SPDE(a::P, b::P, c::P, D::Derivation, n::Int) where
    {T<:FieldElement, P<:PolyElem{T}}
    iscompatible(a, D) && iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials a, b, and c must be in the domain of derivation D")
    !iszero(a) || error("polynomial a must be nonzero")
    Z = zero(a)
    no_solution = (Z, Z, 0, Z, Z, 0)
    if n<0
        if iszero(c)
            return Z, Z, 0, Z, Z, 1
        else
            return no_solution
        end
    end
    g = gcd(a, b)
    if !iszero(rem(c, g))
        return no_solution
    end
    a = divexact(a, g)
    b = divexact(b, g)
    c = divexact(c, g)
    if degree(a)==0
        return divexact(b, a), divexact(c, a), n, one(a), Z, 1
    end
    r, z = gcdx(b, a, c)
    u = SPDE(a, b + D(a), z - D(r), D, n - degree(a))
    if u[6]==0 # no solution of recursive call of SPDE
        return no_solution
    end
    b1, c1, m, α, β, _ = u
    b1, c1, m, a*α, a*β+r, 1
end

"""
    PolyRischDENoCancel1(b, c, D[, n=∞]) -> (q, ρ)

Polynomial Risch differential equation - no cancellation.

Given a field `k`, a derivation `D` on `k[t]`, an integer `n`
and `b`, `c` in `k[t]`  with `b≠0` such that either `D=d/dt`or
`degree(b)>max(0, degree(D(t))-1)`, return either
`ρ=0`, in which case the equation `D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.5, p. 208.
"""
function PolyRischDENoCancel1(b::P, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}}      # here typemax(Int) represents +infinity
    iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials b and c must be in the domain of derivation D")
    !iszero(b) || error("polynomial b must be nonzero")
    isbasic(D) || degree(b)>max(0, degree(D)-1) ||
        error("either derivation D must be basic or degree(b)>max(0, degree(D)-1)")
    Z  = zero(b)
    t = gen(parent(b))
    q = Z
    while !iszero(c)
        m = degree(c)-degree(b)
        if n<0 || m<0 || m>n
            return Z, 0 # no solution
        end
        p = (leading_coefficient(c)//leading_coefficient(b))*t^m
        q += p
        n = m - 1
        c -= D(p)+b*p
    end
    q, 1
end

"""
    PolyRischDENoCancel2(b, c, D[, n=∞]) -> (q, B, C, ρ)

Polynomial Risch differential equation - no cancellation.

Given a field `k`, a derivation `D` on `k[t]`, an integer `n`
and `b`, `c` in `k[t]` such that `degree(b)<degree(D(t))-1` and either
`D=d/dt` or `degree(D(t))≥2`, return either
`ρ=0`, in which case the equation `D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`,
or `ρ=2` and `q`, `A`, `B` in `k[t]` such that for any solution `y` in `k[t]`
of degree at most `n` of `D(y)+b*y=c`, `z=y-q` is a solution of `D(z)+B*z=C`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.5, p. 209.
"""
function PolyRischDENoCancel2(b::P, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}}      # here typemax(Int) represents +infinity
    iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials b and c must be in the domain of derivation D")
    δ = degree(D)
    iszero(b) || degree(b)<δ-1 || error("degree(b)<degree(D)-1 must hold")
    # Note: In Bronstein's book the convention degree(0)=-infinity is used,
    # but in AbstractAlgebra we have degree(0)=-1, so we have to explicitly
    # include the clause iszero(b) in the above condition.
    isbasic(D) || δ>=2 ||
        error("either derivation D must be basic or degree(D)>=2")
    t = gen(parent(b))
    λ = leading_coefficient(D)
    Z  = zero(b)
    Z0 = zero(λ)
    no_solution = (Z, Z0, Z0, 0)
    q = Z
    while !iszero(c)
        if n==0
            m=0
        else
            m = degree(c) - δ + 1
        end
        if n<0 || m<0 || m>n
            return no_solution
        end
        if m>0
            p = (leading_coefficient(c)//(m*λ))*t^m
        else
            if degree(b)!=degree(c)
                return no_solution
            end
            if degree(b)==0
                return q, constant_coefficient(b), constant_coefficient(c), 2
            end
            p = (leading_coefficient(c)//leading_coefficient(b)) + Z #make p\in k an element of k(t)
        end
        q += p
        n = m - 1
        c -= D(p)+b*p
    end
    q, Z0, Z0, 1
end

"""
    PolyRischDENoCancel3(b, c, D[, n=∞]) -> (q, m, C, ρ)

Polynomial Risch differential equation - no cancellation.

Given a field `k`, a derivation `D` on `k[t]`, an integer `n`
and `b`, `c` in `k[t]` such that `degree(D(t))≥2` and `degree(b)=degree(D(t))-1`,
return either `ρ=0`, in which case the equation `D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`,
or `ρ=2`, `q`, `C` in `k[t]` and an integer `m`such that for any solution `y` in `k[t]`
of degree at most `n` of `D(y)+b*y=c`, `z=y-q` is a solution in `k[t]`
of degree at most `m` of `D(z)+b*z=C`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.5, p. 210.
"""
function PolyRischDENoCancel3(b::P, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}}      # here typemax(Int) represents +infinity
    iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials b and c must be in the domain of derivation D")
    δ = degree(D)
    δ>=2 || error("degree(D) must be >= 2")
    degree(b)==δ-1 || error("degree(b)==degree(D)-1 must hold")
    Z = zero(b)
    no_solution = (Z, 0, Z, 0)
    t = gen(parent(b))
    λ = leading_coefficient(D)
    M = -1
    h = -leading_coefficient(b)//λ
    if isrational(h)
        h = rationalize_over_Int(h)
        if isone(denominator(h)) &&  h>0
            M = numerator(h)
        end
    end
    q = Z
    while !iszero(c)
        m = max(M, degree(c)-δ+1)
        if n<0 || m<0 || m>n
            return no_solution
        end
        u = m*λ + leading_coefficient(b)
        if iszero(u)
            return q, m, c, 2
        end
        if m>0
            p = (leading_coefficient(c)//u)*t^m
        else
            if degree(c)!=δ-1
                return no_solution
            end
            p = (leading_coefficient(c)//leading_coefficient(b)) + Z #make p\in k an element of k[t]
        end
        q += p
        n = m - 1
        c -= D(p)+b*p
    end
    q, 0, Z, 1
end

"""
    PolyRischDECancelPrim(b, c, D[, n=∞]) -> (q, ρ)

Polynomial Risch differential equation, cancellation - primitive case.

Given a field `k`, a derivation `D` on `k[t]` with `D(t)` in `k`, an integer `n`,
`b≠0` in `k` and `c` in `k[t]`, return either
`ρ=0`, in which case the equation `D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.6, p. 212.
"""
function PolyRischDECancelPrim(b::T, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}}       # here typemax(Int) represents +infinity
    isprimitive(D) ||
        error("monomial of derivation D must be primitive")
    D0 = BaseDerivation(D)
    iscompatible(b, D0) ||
        error("coefficient b must be in the domain of the base derivation of D")
    iscompatible(c, D) ||
        error("polynomial c must be in the domain of derivation D")
    Z = zero(c)
    no_solution = (Z, 0)
    if b==0  #Note: case b==0 allowed, TODO: test this and update docstring
        q0, ρ = InFieldDerivative(c//1, D) # make poly c a rational function
        q = numerator(q0)
        if ρ<=0 || !isone(denominator(q0)) || degree(q)>n
            return no_solution
        end
        return q, 1
    end
    t = gen(parent(c))
    z, ρ = InFieldLogarithmicDerivative(b, D0)
    if ρ>0
        p0, ρ = InFieldDerivative(z*c, D)
        p = numerator(p0)
        if ρ<=0 || !isone(denominator(p0)) || degree(p)>n
            return no_solution
        end
        return p//z, 1
    end
    if iszero(c)
        return Z, 1 # zero is solution
    end
    if n<degree(c)
        return no_solution
    end
    q = Z
    while !iszero(c)
        m = degree(c)
        if n<m
            return no_solution
        end
        s, ρ = RischDE(b, leading_coefficient(c), D0)
        if ρ<=0
            return Z, 0
        end
        q += s*t^m
        n = m - 1
        c -= b*s*t^m + D(s*t^m)
    end
    q, 1
end

"""
    PolyRischDECancelExp(b, c, D[, n=∞]) -> (q, ρ)

Polynomial Risch differential equation, cancellation - hyperexponential case.

Given a field `k`, a derivation `D` on `k[t]` with `D(t)/t` in `k`, an integer `n`,
`b≠0` in `k` and `c` in `k[t]`, return either
`ρ=0`, in which case the equation `D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.6, p. 213.
"""
function PolyRischDECancelExp(b::T, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}}      # here typemax(Int) represents +infinity
    ishyperexponential(D) ||
        error("monomial of derivation D must be hyperexponential")
    D0 = BaseDerivation(D)
    iscompatible(b, D0) ||
        error("coefficient b must be in the domain of the base derivation of D")
    iscompatible(c, D) ||
        error("polynomial c must be in the domain of derivation D")
    Z = zero(c)
    no_solution = (Z, 0)
    if b==0
        q0, ρ = InFieldDerivative(c//1, D) # make poly c a rational function
        q = numerator(q0)
        if ρ<=0 || !isone(denominator(q0)) || degree(q)>n
            return no_solution
        end
        return q, 1
    end
    t = gen(parent(c))
    H = MonomialDerivative(D)
    w = coeff(H,1) # = Dt/t for hyperexponentialm case
    n1, m, z, ρ = ParametricLogarithmicDerivative(b, w, D0)
    if  ρ>0 && n1==1
        p, ρ = InFieldDerivative(c*z*(t//1)^m, D) # t//1 in k(t)
        if ρ<=0 || !isreduced(p, D)
            return no_solution
        end
        q0 = p//(z*(t//1)^m) # t((1 in k(t)))
        q = numerator(q0)
        if isone(denominator(q0)) && degree(q)<=n
            return q, 1
        else
            return no_solution
        end
    end
    if iszero(c)
        return Z, 1 # zero is solution
    end
    if n<degree(c)
        return no_solution
    end
    q = Z
    while !iszero(c)
        m = degree(c)
        if n<m
            return no_solution
        end
        s, ρ = RischDE(b+m*w, leading_coefficient(c), D0)
        if ρ<=0
            return no_solution
        end
        q += s*t^m
        n = m - 1
        c -= b*s*t^m + D(s*t^m)
    end
    q, 1
end

"""
    PolyRischDECancelTan(b₀, c, D[, n=∞]) -> (q, ρ)

Polynomial Risch differential equation, cancellation - hypertangent case.

Given a field `k` not containing `√-1`, a derivation `D` on `k[t]` with `D(t)/(t^2+1)` in `k`, an integer `n`,
`b₀` in `k` and `c` in `k[t]`, return either
`ρ=0`, in which case the equation `D(q)+(b₀-n*t*D(t)/(t^2+1))*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Section 6.6, p. 215.
"""
function PolyRischDECancelTan(b0::T, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}}      # here typemax(Int) represents +infinity
    ishypertangent(D) ||
        error("monomial of derivation D must be hypertangent")
    D0 = BaseDerivation(D)
    iscompatible(b0, D0) ||
        error("coefficient b0 must be in the domain of the base derivation of D")
    iscompatible(c, D) ||
        error("polynomial c must be in the domain of derivation D")
    Z = zero(c)
    no_solution = (Z, 0)
    if n==0
        if degree(c)<=0
            if !iszero(b0)
                q, ρ = RischDE(b0, constant_coefficient(c), BaseDerivation(D))
                if ρ<=0
                    return no_solution
                else
                    return q + Z, 1
                end
            end
            q, ρ = InFieldDerivative(constant_coefficient(c), BaseDerivation(D))
            if ρ<=0
                return no_solution
            else
                return q + Z, 1
            end
        end
    end
    t = gen(parent(c))
    p = t^2 + 1
    η = constant_coefficient(divexact(MonomialDerivative(D), p))
    r = rem(c, p)
    c0 = coeff(r, 0)
    c1 = coeff(r, 1)
    u, v, ρ   = CoupledDESystem(b0, -n*η, c0, c1, BaseDerivation(D))
    if ρ<=0
        return no_solution
    end
    r = u + v*t
    if n==1
        return r, 1
    end
    c = divexact(c - D(r) - (b0 - n*η*t)*r, p)
    h, ρ = PolyRischDECancelTan(b0, c, D, n-2)
    if ρ<=0
        return no_solution
    end
    p*h + r, 1
end

"""
    PolyRischDE(b, c, D[, n=∞]) -> (q, ρ)

Polynomial Risch differential equation.

Given a field `k`, a derivation `D` on `k[t]`, an integer `n` and
`b`, `c` in `k[t]`, return either
`ρ=0`, in which case the equation `D(q)+b*q=c` has no solution of degree at most `n` in `k[t]`,
or `ρ=1` and a solution `q` in `k[t]` of this equation with `degree(q)≤n`.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Sections 6.6 and 6.7, p_power_N. 206-216.
"""
function PolyRischDE(b::P, c::P, D::Derivation, n::Int=typemax(Int)) where
    {T<:FieldElement, P<:PolyElem{T}} # here typemax(Int) represents +infinity
    iscompatible(b, D) && iscompatible(c, D) ||
        error("polynomials b and c must be in the domain of derivation D")
    δ = degree(D)
    Z = zero(parent(b))
    if !iszero(b) && isbasic(D) || degree(b)>max(0, δ-1)
        return PolyRischDENoCancel1(b, c, D, n)
    elseif (iszero(b) || degree(b)<δ-1) && (isbasic(D) || δ>=2)
        q, b, c, ρ = PolyRischDENoCancel2(b, c, D, n)
        if ρ==2
            q1, ρ = RischDE(b, c, BaseDerivation(D))
            ρ>=1 || return Z, ρ
            q = q1 - q
        end
        return q, 1
    elseif δ>=2 && degree(b)==δ-1
        q, m, c, ρ = PolyRischDENoCancel3(b, c, D, n)
        if ρ<=0
            return Z, 0
        elseif ρ==1
            return q, 1
        elseif ρ==2
            if ishypertangent(D)
                if contains_I(parent(b))
                    # Seems that PolyRischDE was called by CoupledDESystem,
                    # this case should be handled by CoupledDESystem.
                    return c, 101 + max(-1, m) # return new c, m
                else
                    t = gen(parent(b))
                    η = divexact(MonomialDerivative(D), t^2+1)
                    b0 = b + n*t*η
                    @assert degree(b0)<=0
                    return PolyRischDECancelTan(constant_coefficient(b0), c, D, n)
                end
            else
                H = MonomialDerivative(D)
                throw(NotImplementedError("PolyRischDE: no cancellation, degree(b)==δ-1, monomial derivative $H\n@ $(@__FILE__):$(@__LINE__)"))
            end
        else
            @assert false # never reach this point
        end
    # At this point only the cancellation case δ<=1, D!=d/dt is possible;
    # this is only compatible with primitive or hyperexponential.
    elseif contains_I(parent(b))
        # Seems that PolyRischDE was called by CoupledDESystem,
        # this case should be handled by CoupledDESystem.
        return c, 101 + max(-1, n) # return original c, n
    elseif isprimitive(D)
        return PolyRischDECancelPrim(constant_coefficient(b), c, D, n)
    elseif ishyperexponential(D)
        return PolyRischDECancelExp(constant_coefficient(b), c, D, n)
    else
        H = MonomialDerivative(D)
        throw(NotImplementedError("PolyRischDE: cancellation case, monomial derivative $H\n@ $(@__FILE__):$(@__LINE__)"))
    end
end

"""
    RischDE(f, g, D) -> (y, ρ)

Risch differential equation.

Given a field `K`, a derivation `D` on `K` and  `f`, `g` in `K`, return
either `ρ=0`, in which case the equation `D(y)+f*q=g` has no solution in `K`,
or `ρ=1` and a solution `y` in `K` of this equation.

See [Bronstein's book](https://link.springer.com/book/10.1007/b138171), Chapter 6, p. 181.
"""
function RischDE(f::F, g::F, D::Derivation) where F<:FieldElement
    iszero(D) || error("base case only for null derivations")
    #base case => pure algebra problem ...
    iscompatible(f, D) && iscompatible(g,D) ||
        error("coefficients f and g must be in the domain of the base derivation of D")
    if iszero(f)
        if iszero(g)
            return zero(parent(f)), 1
        else
            return zero(parent(f)), 0 # no solution
        end
    else
        return g//f, 1
    end
 end

function RischDE(f::F, g::F, D::Derivation) where
    {T<:FieldElement, P<:PolyElem{T}, F<:FracElem{P}}
    iscompatible(f, D) && iscompatible(g, D) ||
        error("rational functions f and g must be in the domain of derivation D")
    if iszero(f)
        return InFieldDerivative(g::F, D::Derivation)
    end

    Z = zero(f)
    h0 = WeakNormalizer(f, D)

    f = f - D(h0)//h0
    g = h0*g
    a, b, c, h1, ρ = RdeNormalDenominator(f, g, D)
    ρ>=1 || return Z, ρ

    a, b, c, h2 = RdeSpecialDenominator(a, b, c, D)

    n = RdeBoundDegree(a, b, c, D)

    b, c, n, α, β, ρ = SPDE(a, b, c, D, n)
    ρ>=1 || return Z, ρ

    z, ρ = PolyRischDE(b, c, D, n)
    ρ>=1 || return Z, ρ

    (α*z+β)//(h0*h1*h2), 1
end


default_metric_dispatch(::Stiefel, ::EuclideanMetric) = Val(true)

function decorator_transparent_dispatch(::typeof(distance), ::Stiefel, args...)
    return Val(:intransparent)
end

@doc raw"""
    exp(M::Stiefel, p, X)

Compute the exponential map on the [`Stiefel`](@ref)`{n,k,𝔽}`() manifold `M`
emanating from `p` in tangent direction `X`.

````math
\exp_p X = \begin{pmatrix}
   p\\X
 \end{pmatrix}
 \operatorname{Exp}
 \left(
 \begin{pmatrix} p^{\mathrm{H}}X & - X^{\mathrm{H}}X\\
 I_n & p^{\mathrm{H}}X\end{pmatrix}
 \right)
\begin{pmatrix}  \exp( -p^{\mathrm{H}}X) \\ 0_n\end{pmatrix},
````

where $\operatorname{Exp}$ denotes matrix exponential,
$\cdot^{\mathrm{H}}$ denotes the complex conjugate transpose or Hermitian, and $I_k$ and
$0_k$ are the identity matrix and the zero matrix of dimension $k × k$, respectively.
"""
exp(::Stiefel, ::Any...)

function exp!(::Stiefel{n,k}, q, p, X) where {n,k}
    A = p' * X
    B = exp([A -X'*X; I A])
    @views begin
        r = p * B[1:k, 1:k]
        mul!(r, X, B[(k + 1):(2 * k), 1:k], true, true)
    end
    mul!(q, r, exp(-A))
    return q
end

@doc raw"""
    get_basis(M::Stiefel{n,k,ℝ}, p, B::DefaultOrthonormalBasis) where {n,k}

Create the default basis using the parametrization for any $X ∈ T_p\mathcal M$.
Set $p_\bot \in ℝ^{n\times(n-k)}$ the matrix such that the $n\times n$ matrix of the common
columns $[p\ p_\bot]$ is an ONB.
For any skew symmetric matrix $a ∈ ℝ^{k\times k}$ and any $b ∈ ℝ^{(n-k)\times k}$ the matrix

````math
X = pa + p_\bot b ∈ T_p\mathcal M
````

and we can use the $\frac{1}{2}k(k-1) + (n-k)k = nk-\frac{1}{2}k(k+1)$ entries
of $a$ and $b$ to specify a basis for the tangent space.
using unit vectors for constructing both
the upper matrix of $a$ to build a skew symmetric matrix and the matrix b, the default
basis is constructed.

Since $[p\ p_\bot]$ is an automorphism on $ℝ^{n\times p}$ the elements of $a$ and $b$ are
orthonormal coordinates for the tangent space. To be precise exactly one element in the upper
trangular entries of $a$ is set to $1$ its symmetric entry to $-1$ and we normalize with
the factor $\frac{1}{\sqrt{2}}$ and for $b$ one can just use unit vectors reshaped to a matrix
to obtain orthonormal set of parameters.
"""
function get_basis(M::Stiefel{n,k,ℝ}, p, B::DefaultOrthonormalBasis{ℝ}) where {n,k}
    V = get_vectors(M, p, B)
    return CachedBasis(B, V)
end

function get_coordinates!(
    M::Stiefel{n,k,ℝ},
    c,
    p,
    X,
    B::DefaultOrthonormalBasis{ℝ},
) where {n,k}
    V = get_vectors(M, p, B)
    c .= inner.(Ref(M), Ref(p), V, Ref(X))
    return c
end

function get_vector!(M::Stiefel{n,k,ℝ}, X, p, c, B::DefaultOrthonormalBasis{ℝ}) where {n,k}
    V = get_vectors(M, p, B)
    zero_tangent_vector!(M, X, p)
    length(c) < length(V) && error(
        "Coordinate vector too short. Excpected $(length(V)), but only got $(length(c)) entries.",
    )
    @inbounds for i in 1:length(V)
        X .+= c[i] .* V[i]
    end
    return X
end

function get_vectors(::Stiefel{n,k,ℝ}, p, ::DefaultOrthonormalBasis{ℝ}) where {n,k}
    p⊥ = nullspace([p zeros(n, n - k)])
    an = div(k * (k - 1), 2)
    bn = (n - k) * k
    V = vcat(
        [p * vec2skew(1 / sqrt(2) .* _euclidean_unit_vector(an, i), k) for i in 1:an],
        [p⊥ * reshape(_euclidean_unit_vector(bn, j), (n - k, k)) for j in 1:bn],
    )
    return V
end

_euclidean_unit_vector(n, i) = [k == i ? 1.0 : 0.0 for k in 1:n]

@doc raw"""
    project(M::Stiefel,p)

Projects `p` from the embedding onto the [`Stiefel`](@ref) `M`, i.e. compute `q`
as the polar decomposition of $p$ such that $q^{\mathrm{H}q$ is the identity,
where $\cdot^{\mathrm{H}}$ denotes the hermitian, i.e. complex conjugate transposed.
"""
project(::Stiefel, ::Any, ::Any)

function project!(::Stiefel, q, p)
    s = svd(p)
    mul!(q, s.U, s.Vt)
    return q
end

@doc raw"""
    project(M::Stiefel, p, X)

Project `X` onto the tangent space of `p` to the [`Stiefel`](@ref) manifold `M`.
The formula reads

````math
\operatorname{proj}_{\mathcal M}(p, X) = X - p \operatorname{Sym}(p^{\mathrm{H}}X),
````

where $\operatorname{Sym}(q)$ is the symmetrization of $q$, e.g. by
$\operatorname{Sym}(q) = \frac{q^{\mathrm{H}}+q}{2}$.
"""
project(::Stiefel, ::Any...)

function project!(::Stiefel, Y, p, X)
    A = p' * X
    T = eltype(Y)
    copyto!(Y, X)
    mul!(Y, p, A + A', T(-0.5), true)
    return Y
end

function vector_transport_to!(M::Stiefel, Y, ::Any, X, q, ::ProjectionTransport)
    return project!(M, Y, q, X)
end


# module ExponentialFamilies
using Distributions

# A parameterized exponential family
struct ExponentialFamily{P,X}
    logh :: Function # :: X -> Real
    logg :: Function # :: P -> Real
    η    :: Function # :: P -> StaticVector{N, Real}
    t    :: Function # :: X -> StaticVector{N, Real}
end

# A specific instance of a distribution belonging to an exponential family
struct ExpFamDist{P,X}
    fam :: ExponentialFamily{P,X}
    θ :: P
end

# Find the exponential family distribution instance for a given distribution.
# This needs to be reworked, since many distributions are not from exponential families
efam(d::ExpFamDist{P,X} where {P,X}) = d

Distributions.logdensity(d::ExpFamDist{P,X}, x :: X) where {P,X} = d.fam.logh(x) + d.fam.logg(d.θ) + sum(d.fam.η(d.θ) .* d.fam.t(x))

# Exponential families are closed under iid
function iid(n::Integer, d::ExpFamDist{P,X} ) where {P,X}
    fam = d.fam
    logh(xs) = sum(fam.logh.(xs))
    logg(θ) = n * fam.logg(θ)
    η(θ) = fam.η(θ)
    t(xs) = sum(fam.t.(xs))

    newfam = ExponentialFamily{P,Vector{X}}(logh, logg, η, t)
    ExpFamDist{P,Vector{X}}(newfam,d.θ)
end

iid(n, d ) = iid(n,efam(d))

################################
# Bernoulli example

logit(p) = log(p/(1-p))

function efam(d::Bernoulli{P}) where {P <: Real}
    logh(x) = 0
    logg(p::P) = -log(1 - p)
    η(p::P) = logit(p)
    t(x) = Float64(x)

    fam = ExponentialFamily{P,Bool}(logh, logg, η, t)

    ExpFamDist(fam, d.p)
end

efam(Bernoulli(0.2))

binom = iid(1000,Bernoulli(0.2))

using BenchmarkTools

@btime logdensity(binom, repeat([true],1000))

@btime sum(logdensity.(Bernoulli(0.2),repeat([true],1000)))

################################
# Normal example

const sqrt2π = sqrt(2*π)

const invsqrt2π = 1/sqrt(2*π)

function efam(d::Normal{P}) where {P <: Real}
    logh(x) = invsqrt2π
    logg(θ) = 0.5 * (θ.μ/θ.σ)^2 + log(θ.σ)
    η(θ) = (θ.μ/θ.σ^2, -1/(2*θ.σ^2))
    t(x) = (x,x^2)

    fam = ExponentialFamily{NamedTuple{(:μ, :σ),Tuple{P,P}}, Float64}(logh, logg, η, t)

    ExpFamDist(fam, (μ=d.μ,σ=d.σ))
end

const xs = 3 .+ 4 .* randn(1000);

@btime logdensity(iid(1000,Normal(3.1,4.3)), xs)

@btime sum(logdensity.(Normal(3,4),xs))

# end # module


using LinearAlgebra, Random, Test
using SafeTestsets

@time @safetestset "metrics" begin include("metrics.jl") end


module SummaryTest

using OnlineStats, StatsBase, Base.Test

x = randn(500)
x1 = randn(500)
x2 = randn(501)
xs = hcat(x1, x)

@testset "Summary" begin
@testset "Mean / Means" begin
    o = Mean()
    fit!(o, x1, 5)
    @test mean(o) ≈ mean(x1)
    @test nobs(o) == 500
    o = Mean(x2)
    @test mean(o) ≈ mean(x2)
    @test mean(o) == value(o)
    @test center(o, x[1]) == x[1] - mean(o)

    o = Means(2)
    fit!(o, xs, 5)
    @test mean(o) ≈ vec(mean(xs, 1))
    fit!(o, xs, 1)
    o = Means(xs)
    o2 = Means(2)
    fit!(o2, xs)
    @test mean(o) ≈ mean(o2)
    @test center(o, vec(xs[1, :])) == vec(xs[1, :]) - mean(o)

    o = Mean(x1)
    o2 = Mean(x2)
    merge!(o, o2)
    @test mean(o) ≈ mean(vcat(x1, x2))
    merge!(o, o2, :mean)
    merge!(o, o2, :singleton)
end
@testset "Variance / Variances" begin
    o = Variance(x1)
    @test var(o) ≈ var(x1)
    o = Variance()
    fit!(o, x1)
    @test var(o)            ≈ var(x1)
    @test std(o)            ≈ std(x1)
    @test mean(o)           ≈ mean(x1)
    @test var(o)            == value(o)
    @test center(o, x[1])   == x[1] - mean(o)
    @test zscore(o, x[1])   == (x[1] - mean(o)) / std(o)
    o = Variance(x1)
    o2 = Variance(x2)
    merge!(o, o2)
    @test var(o) ≈ var(vcat(x1, x2))

    o = Variances(xs)
    @test var(o) ≈ vec(var(xs, 1))
    @test mean(o) ≈ vec(mean(xs, 1))
    @test std(o) ≈ vec(std(xs, 1))
    @test var(o) == value(o)
    o2 = Variances(2)
    fit!(o2, xs)
    @test var(o) ≈ var(o2)
    @test center(o, vec(xs[1, :])) == vec(xs[1, :]) - mean(o)
    @test zscore(o, vec(xs[1, :])) == (vec(xs[1, :]) - mean(o)) ./ std(o)
    @test zscore(Variances(5), ones(5)) == ones(5)

    o = Variances(xs)
    o2 = Variances(xs)
    merge!(o, o2)
    @test var(o) ≈ vec(var(vcat(xs, xs), 1))
end
@testset "CovMatrix" begin
    o = CovMatrix(xs)
    @test cov(o) ≈ cov(xs)
    @test cor(o) ≈ cor(xs)

    o2 = CovMatrix(2)
    fit!(o2, xs)
    @test cov(o) ≈ cov(o2)
    @test cor(o) ≈ cor(o2)
    @test value(o) ≈ value(o2)
    @test mean(o) ≈ mean(o2)
    @test mean(o2) ≈ vec(mean(xs, 1))
    @test var(o) ≈ vec(var(xs, 1))
    @test std(o) ≈ vec(std(xs, 1))
    fit!(o2, xs, 1)

    o = CovMatrix(xs)
    o2 = CovMatrix(xs)
    merge!(o, o2)
    @test cov(o) ≈ cov(vcat(xs, xs))
end
@testset "Extrema" begin
    o = Extrema(x1)
    @test extrema(o) == extrema(x1)
    o2 = Extrema()
    fit!(o2, x1)
    @test extrema(o) == extrema(o2)

    o = Extrema(x1)
    o2 = Extrema(x2)
    merge!(o, o2)
    @test extrema(o) == extrema(vcat(x1, x2))
end
@testset "QuantileSGD / QuantileMM" begin
    o = QuantileSGD(x1)
    fit!(o, x2, 2)

    o = QuantileMM(x1)
    fit!(o, x2, 2)
end
@testset "Moments" begin
    o = Moments(x1)
    @test mean(o) ≈ mean(x1)
    @test var(o) ≈ var(x1)
    @test std(o) ≈ std(x1)
    @test_approx_eq_eps skewness(o) skewness(x1) .01
    @test_approx_eq_eps kurtosis(o) kurtosis(x1) .1
end
@testset "Diff / Diffs" begin
    Diff()
    Diff(Float64)
    Diff(Float32)
    Diff(Int64)
    Diff(Int32)
    y = randn(100)
    o = Diff(y)
    @test typeof(o) == Diff{Float64}
    @test last(o) == y[end]
    @test diff(o) == y[end] - y[end-1]
    @test value(o) == diff(o)
    y = rand(Int, 100)
    o = Diff(y)
    @test typeof(o) == Diff{Int64}
    @test last(o) == y[end]
    @test diff(o) == y[end] - y[end-1]

    Diffs(10)
    Diffs(Int32, 10)
    y = randn(100, 10)
    o = Diffs(y)
    @test last(o) == vec(y[end, :])
    @test diff(o) == vec(y[end, :] - y[end-1, :])
    @test value(o) == diff(o)
end
@testset "Sum / Sums" begin
    Sum()
    Sum(Float64)
    Sum(Float32)
    Sum(Int64)
    Sum(Int32)
    y = randn(100)
    o = Sum(y)
    @test typeof(o) == Sum{Float64}
    @test sum(o) ≈ sum(y)
    @test value(o) == sum(o)
    y = rand(Int, 100)
    o = Sum(y)
    @test typeof(o) == Sum{Int64}
    @test sum(o) ≈ sum(y)

    Sums(10)
    Sums(Int32, 10)
    y = randn(100, 10)
    o = Sums(y)
    @test sum(o) ≈ vec(sum(y,1))
    @test value(o) == sum(o)
end
end # summary
end #module


#
# This file is the aerodynamic drag calculations using "Zero Lift Drag"
# Individual components are calculated and the drags combined for an approximation
# This calculation is currently using the wing as well as the empennage.
# Wing lift induced drag will also be included so adds a velocity variable to the solution.
# Interference & excrescence drag are approximated using a factor based on the aircraft type.
# In this case = 1.15, which is slightly better than the general aviation factor of 1.2 - 1.3.
#
# This code is licensed under MIT license (see LICENSE.txt for details)
#
# Dec 12, 2021 		Initial code
#

using Plots

#						Mach	A		B	  R^2
Skin_friction_coeff =  [0.0	0.0391	-0.157	0.9967
						0.3	0.0399	-0.159	0.9976
						0.7	0.0392	-0.16	0.9971
						0.9	0.0376	-0.159	0.9965
						1.0	0.0381	-0.161	0.9970
						1.5	0.0371	-0.164	0.9966
						2.0	0.0329	-0.162	0.9941
						2.5	0.0286	-0.161	0.9944
						3.0	0.0261	-0.161	0.9936]

						Mach_indexes = [0.0, 0.3, 0.7, 0.9, 1.0, 1.5, 2.0, 2.5, 3.0]

# dataframe for keeping track of (book-keeping)drag variables
drag_vars = DataFrame(
		Velocity			=	Float32[],
		Cd_fuselage 		=	Float32[],
		Cd_wing 			=	Float32[],
		Cd_horizontal_tail	=	Float32[],
		Cd_vertical_tail	=	Float32[],
		Cd0_total			=	Float32[]
	)

#  interference + excrescence drag coefficient
interference_drag = 1.15

	#------------------------------ Find the Nearest Value in an Array) ------------------------------#
function findnearest(a,x)
   idx = searchsortedfirst(a,x)
   if (idx==1); return idx; end
   if (idx>length(a)); return length(a); end
   if (a[idx]==x); return idx; end
   if (abs(a[idx]-x) < abs(a[idx-1]-x))
      return idx
   else
      return idx-1
   end
end

	#------------------------------ Drag Simulation Wrapper (drag_sim) ------------------------------#
	#  These calculations assume "zero lift drag method"
	#  Itetrates through various temperatures and fills out a matrix of constituient drags
	#  of various components

function drag_sim(df)

	# define structures to hold definitions for the call: (name, value, units, stored)
	SVstart  	= req_vars("SVstart",  	0, "m/s", 	0)	# Lowest simulation velocity
	SVend	   	= req_vars("SVend",		0, "m/s",	0)	# Highest simulation velocity
	SVincr	   	= req_vars("SVincr",	0, "m/s",	0)	# Loop increment

	var_array = [SVstart, SVend, SVincr]				# Array of structures to enable easy iteration

	Extract_data(df, var_array)							# Loads the array of structures with data, checks units and presence

	# loop through the various velocities
	for V in SVstart.value:SVincr.value:SVend.value
            f_drag = fuselage_drag(df, V)
			w_drag = wing_drag(df, V)
			h_tail = 3.0
			v_tail = 4.0
			Cd0_total = (f_drag + w_drag + h_tail + v_tail) * interference_drag

		# write the results to the storage matrix
		push!(drag_vars,[V*2.237, f_drag, w_drag, h_tail, v_tail, Cd0_total])  ### temporarily changing velocity to MPH
       end

###----------------- some temporary display values to check results
# print("\n\n", drag_vars)
display(plot(drag_vars[:,1], drag_vars[:,2])); sleep(30)

end

	#---------------------------- Fuselage drag component (fuselage_drag) --------------------------#
	#	lf = fuselage length
	#	dfuse = maximum fuselage diameter
	#	Swetf = total fuselage surface area exposed to air
	#	Sw = total wing area

function fuselage_drag(df, V)

	# define structures to hold definitions for the call: (name, value, units, stored)
	Sw  	= req_vars("Sw",  		0, 	"m^2", 	0)			# wing area
	Swetf 	= req_vars("Swetf",  	0, 	"m^2", 	0)			# fuselage wet area
	dfuse	= req_vars("dfuse",		0, 	"m", 	0)			# fuselage maximum diameter
	lf		= req_vars("lf",  		0, 	"m", 	0)			# fuselage maximum length

	var_array = [Sw, Swetf, dfuse, lf]						# Array of structures to enable easy iteration
	Extract_data(df, var_array)								# Loads the array of structures with data, checks units and presence

	Cf_fuse = Cf(df, V)										# calculate the skin friction coefficient
	df_equiv = sqrt((4/π)*dfuse.value)						# factor assumes a non-circular fuselage cross section

	Cd0_fueslage = Cf_fuse * ((1 + 60/(lf.value/df_equiv)^3) + (0.0025 * (lf.value/df_equiv))) * (Swetf.value/Sw.value)

	return (Cd0_fueslage)
end

	#--------------------------------  Wing drag component (wing_drag) -----------------------------#

function wing_drag(df, V)
			return (2)
end

	#------------------------------------- Speed of Sound (c) --------------------------------------#

function c(df)

	# define structures to hold definitions for the call: (name, value, units, stored)
	T  	   = req_vars("T",  0, 	"C", 	0)			# T = Temperature

	var_array = [T]									# Array of structures to enable easy iteration

	Extract_data(df, var_array)						# Loads the array of structures with data, checks units and presence

	return ((331.3 * sqrt(1 + T.value/273.15)))

end

	#--------------------------------- Skin Friction Coefficient (Cf) --------------------------------#
	#									Cf = A * Re^B

function Cf(df, V)

	lf		= req_vars("lf",  		0, 	"m", 	0)			# fuselage maximum length
	var_array = [lf]										# Array of structures to enable easy iteration
	Extract_data(df, var_array)								# Loads the array of structures with data, checks units and presence

	Re_number = Re(V, lf.value)								# calculatealculate the fuselage reynolds number for velocity
	speed_of_sound = c(df)									# calculate the speed of sound for the temperature
	Mach_no = V / speed_of_sound							# calculate the mach number
	Mach_nearest = findnearest(Mach_indexes, Mach_no)		# find the index of the closest fit in the mach table

	Cf = Skin_friction_coeff[Mach_nearest,2] * Re_number^Skin_friction_coeff[Mach_nearest,3]
	return (Cf)

end


# Simple kernel for matrix multiplication
@kernel function matmul_kernel!(a, b, c)
    i, j = @index(Global, NTuple)

    # creating a temporary sum variable for matrix multiplication
    tmp_sum = zero(eltype(c))
    for k = 1:size(a)[2]
        @inbounds tmp_sum += a[i,k] * b[k, j]
    end

    c[i,j] = tmp_sum
end

function matmul_testsuite(backend, ArrayT)

    matmul = matmul_kernel!(backend(), (32, 32))
    a = ArrayT(rand(128, 256))
    b = ArrayT(rand(256, 128))
    c = ArrayT(zeros(128, 128))
    wait(matmul(a, b, c, ndrange=size(c)))

    @test c ≈ a*b

    dmatmul = Enzyme.autodiff(matmul)
    da = similar(a)
    da .= 0
    db = similar(b)
    db .= 0
    dc = similar(c)
    dc .= 1
    c .= 0

    compare_dc = copy(dc)
    wait(dmatmul(
        Duplicated(a, da),
        Duplicated(b, db),
        Duplicated(c, dc), ndrange=size(c)))

    @test da ≈ compare_dc * b'
    @test db ≈ a' * compare_dc
end


"""
    struct LinearCombinationDofVector{T} <: AbstractVector{Dof}
      change_of_basis::Matrix{T}
      dof_basis::AbstractVector{<:Dof}
    end

Type that implements a dof basis (a) as the linear combination of a dof basis
(b). The dofs are first evaluated at dof basis (b) (field `dof_basis`) and the
dof values are next mapped to dof basis (a) applying a change of basis (field
`change_of_basis`).

Fields:

- `change_of_basis::Matrix{T}` the matrix of the change from dof basis (b) to (a)
- `dof_basis::AbstractVector{<:Dof}` A type representing dof basis (b)
"""
struct LinearCombinationDofVector{T} <: AbstractVector{Dof}
  change_of_basis::Matrix{T}
  dof_basis::AbstractVector{<:Dof}
end

@inline Base.size(a::LinearCombinationDofVector) = size(a.dof_basis)
@inline Base.axes(a::LinearCombinationDofVector) = axes(a.dof_basis)
@inline Base.getindex(a::LinearCombinationDofVector,i::Integer) = getindex(a.dof_basis,i)
@inline Base.IndexStyle(::LinearCombinationDofVector) = IndexLinear()

function linear_combination(a::AbstractMatrix{<:Number},
                            b::AbstractVector{<:Dof})
  LinearCombinationDofVector(a,b)
end

function linear_combination(a::LinearCombinationDofVector{T},
                            b::AbstractVector{<:Dof}) where T
  linear_combination(a.change_of_basis,b)
end

function return_cache(b::LinearCombinationDofVector,field)
  c, cf = return_cache(b.dof_basis,field)
  c, cf, return_cache(*,b.change_of_basis,c)
end

@inline function evaluate!(cache,b::LinearCombinationDofVector,field)
  c, cf, cc = cache
  vals = evaluate!(cache,b.dof_basis,field)
  evaluate!(cc,*,b.change_of_basis,vals)
end