# # Specular Canopy Scattering

# Using packages:
using Plots, Distributions#, PyPlot
using CanopyOptics
pyplot()
theme(:ggplot2)

# Create a specular model with refractive index and "roughness" κ
specularMod = CanopyOptics.SpecularCanopyScattering(nᵣ=1.5, κ=0.2)

# Use standard planophile leaf distribution:
LD = CanopyOptics.planophile_leaves2()

# Create some discretized angles:
# Azimuth (in radians)
ϕ = range(0.0, 2π,  length=200);
# Polar angle as cos(Θ)
μ,w = CanopyOptics.gauleg(180,0.0,1.0);
# Create directional vectors:
dirs = [CanopyOptics.dirVector_μ(a,b) for a in μ, b in ϕ];

# Compute specular reflectance:
R = CanopyOptics.compute_reflection.([specularMod], [dirs[50,1]], dirs, [LD]);

# Polar plot of reflectance
contourf(ϕ, acos.(μ), R,proj=:polar, ylim=(0,π/2), alpha=0.8)

# ### Animation over different Beta leaf distributions
steps = 0.1:0.1:5
α = [collect(steps); 5*ones(length(steps))]
β = [ 5*ones(length(steps)); reverse(collect(steps));]
x = 0:0.01:1
l = @layout [a  b; c d; e f]

anim = @animate for i ∈ eachindex(α)
    LD = CanopyOptics.LeafDistribution(Beta(α[i],β[i]), 2/π)
    R = CanopyOptics.compute_reflection.([specularMod], [dirs[120,1]], dirs, [LD])
    pZ = []
    for m=0:3
        Z⁺⁺, Z⁻⁺ = CanopyOptics.compute_Z_matrices(specularMod, μ, LD, m)
        push!(pZ, contourf(μ, μ, Z⁻⁺, title="Z⁻⁺ Fourier moment $m", xlabel="μꜜ", aspect_ratio = :equal, ylabel="μꜛ"))
    end
    p0 = plot(rad2deg.(π * x/2), pdf.(LD.LD,x), legend=false, ylim=(0,3), title="Leaf angle distribution", xlabel="Θ (degrees)")
    p1 = contourf(ϕ, acos.(μ), R, proj=:polar, ylim=(0,π/2), clims=(0,0.02), alpha=0.8,legend = :none, title="Polar plot R")
    plot(p0, p1, pZ..., layout = l, margin=5Plots.mm)
    plot!(size=(700,700))
end

gif(anim, "anim_fps10.gif", fps = 10)


module Ex7_4
using DescriptorSystems
using LinearAlgebra
using Polynomials
using Test

# Example 7.4 - Nullspace-based synthesis
# Uses only DescriptorSystems.jl
println("Example 7.4")

# define the state-space realizations of Gu, Gd and Gf
n = 5; p = 3; mu = 3; md = 1; mf = 5;       # enter dimensions
# define matrices of the state space realization
A = [ 0 0 1.132 0 -1;
0 -0.0538 -0.1712 0 0.0705;
0 0 0 1 0;
0 0.0485 0 -0.8556 -1.013;
0 -0.2909 0 1.0532 -0.6859];
Bu = [0 0 0;-0.12 1 0;0 0 0;4.419 0 -1.665;1.575 0 -0.0732];
Bd = Bu[:,mu]; Bf = [zeros(n,p) Bu[:,1:mu-1]];
C=eye(p,n); Du=zeros(p,mu); Dd=zeros(p,md); Df=eye(p,mf);
sys = dss(A,[Bu Bd Bf],C,[Du Dd Df]);        # define system

# compute [Q Rf], where Q is a left nullspace basis of
# [Gu Gd;I 0] and Rf = Q*[Gf;0];
Q_Rf, info = glnull([sys;eye(mu,mu+md+mf)],mf);
info.degs

# determine 1-st and 2-nd order scalar output designs
Q1_Rf1 = glmcover1([0 1;eye(2)]*Q_Rf,1)[1];    #  [Q1 Rf1]
Q2_Rf2 = glmcover1([-1 1;eye(2)]*Q_Rf,1)[1];    #  [Q2 Rf2]

# compute stable left coprime factorizations
Q1_Rf1 = glcf(Q1_Rf1, sdeg=-1, smarg = -1)[1];
# choose poles to meet example
Q2_Rf2 = glcf(Q2_Rf2, evals = [-1.0000 + 0.3992im, -1.0000 - 0.3992im], sdeg = -1, smarg = -1)[1];

# compute Q1 and Rf1; check admissibility, i.e., finite reciprocal sensitivity condition
Q1 = Q1_Rf1[:,1:p+mu]; Rf1 = Q1_Rf1[:,p+mu+1:end];
g1 = abs.(evalfr(Rf1,im)); rscond1 = maximum(g1)/minimum(g1)
# compute Q2 and Rf2; check admissibility, i.e., finite reciprocal sensitivity condition
Q2 = Q2_Rf2[:,1:p+mu]; Rf2 = Q2_Rf2[:,p+mu+1:end];
g2 = abs.(evalfr(Rf2,im)); rscond2 = maximum(g2)/minimum(g2)
println("|Rf1(im)| = $g1")
println("rscond1 = $rscond1 ")
println("|Rf2(im)| = $g2")
println("rscond2 = $rscond2")
@test !isinf(rscond1) && !isinf(rscond1)

# display results
println("Q1 = $(dss2rm(Q1))")
println("Q2 = $(dss2rm(Q2))")

end # module


# This simply holds the feasibility checks
function isFeasiblePolyHedron(A, b, x)
    #=
    Checks if all polyhedron checks are satisfied
    Namely: Ax ≤ b
    =#

    state = A * x - b

    if maximum(state) > 0
        return false
    end

    return true
end


using DelimitedFiles
using FFTW

function derive(data, dim, dx)
    out = Array{Complex,2}(undef,size(data,1), size(data,2))
    dim2 = dim+1
    if dim2 == 3
        dim2 = 1
    end
    for i = 1:size(data,dim)
        for j = 1:size(data,dim2)
            #println(j,'\t',i)
            if (dim == 1)
                if (i == size(data,dim))
                    out[i,j] = (data[1, j] - data[i, j])/dx
                else
                    out[i,j] = (data[i + 1, j] - data[i, j])/dx
                end
            else
                if (i == size(data,dim))
                    out[j,i] = (data[j,1] - data[j,i])/dx
                else
                    out[j,i] = (data[j,i + 1] - data[j,i])/dx
                end
            end
        end
    end
    return out
end

# We are calculating the energy to check <Psi|H|Psi>
function calculate_energy(wfc, H_k, H_r, Ax, Ay, g, xDim, yDim, dx, dy)
    hbar = 1.05457148e-34
    omega = 0.6
    omegaX = 6.283

    # Creating momentum and conjugate wavefunctions
    wfc_k = fft(wfc)
    wfc_c = conj(wfc)

    # Finding the momentum and real-space energy terms
    energy_k = wfc_c.*ifft((H_k) .* wfc_k)
    energy_r = wfc_c.* (H_r).* wfc
    energy_i = wfc_c.* (0.5*g*abs2.(wfc)).* wfc

    energy_l = wfc_c.*(im*hbar*(Ax.*derive(wfc,1,dx) + Ay.*derive(wfc,2,dy)))

    # Integrating over all space
    energy_if = 0
    energy_lf = 0
    energy_kf = 0
    energy_rf = 0
    for i = 1:xDim*yDim
        energy_if += real(energy_i[i])*dx*dy
        energy_rf += real(energy_r[i])*dx*dy
        energy_lf += real(energy_l[i])*dx*dy
        energy_kf += real(energy_k[i])*dx*dy
    end

    println("Kinetic energy:", "\t\t\t", energy_kf)
    println("Potential energy:", "\t\t", energy_rf)
    println("Internal energy:", "\t\t", energy_if)
    println("Angular Momentum energy:", '\t', energy_lf)
    println("Total energy:", "\t\t\t", energy_kf+energy_if+energy_rf+energy_lf)

end

# Read in param.cfg file
function calculate(param_file::String, data_dir::String)
    parameters = Dict()

    for line in readlines(data_dir*param_file)
        if line != "[Params]"
            tmp = split(line,"=")
            parameters[tmp[1]] = tmp[2]
        end
    end

    xDim = parse(Int64, parameters["xDim"])
    yDim = parse(Int64, parameters["yDim"])
    dx = parse(Float64, parameters["dx"])
    dy = parse(Float64, parameters["dy"])

    omega = parse(Float64, parameters["omega"])
    omegaX = parse(Float64, parameters["omegaX"])

    g = parse(Float64, parameters["gDenConst"])

    Ax = readdlm(data_dir*"Ax_0")
    Ay = readdlm(data_dir*"Ay_0")
    K = readdlm(data_dir*"K_0")
    V = readdlm(data_dir*"V_0")

    Ax = reshape(Ax, xDim, yDim)
    Ay = reshape(Ay, xDim, yDim)
    K = reshape(K, xDim, yDim)
    V = reshape(V, xDim, yDim)

    start = 0
    ende = parse(Int64, parameters["esteps"])
    endg = parse(Int64, parameters["gsteps"])
    incr = parse(Int64, parameters["printSteps"])

    # Ground State Evolution
    println("Starting imaginary time energy calculation")
    for i = start:incr:endg
        wfc = readdlm(data_dir*"wfc_0_const_"*string(i)) +
              readdlm(data_dir*"wfc_0_consti_"*string(i))*im
        println(data_dir*"wfc_0_const_"*string(i), '\t',
                data_dir*"wfc_0_consti_"*string(i))
        wfc = reshape(wfc, xDim, yDim)
        calculate_energy(wfc, K, V, Ax, Ay, g, xDim, yDim, dx, dy)
        println()
    end

    println()

    # Ground State Evolution
    println("Starting real time energy calculation")
    for i = start:incr:ende
        wfc = readdlm(data_dir*"wfc_ev_"*string(i)) +
              readdlm(data_dir*"wfc_evi_"*string(i))*im
        println(data_dir*"wfc_0_const_"*string(i), '\t',
                data_dir*"wfc_0_consti_"*string(i))
        wfc = reshape(wfc, xDim, yDim)
        calculate_energy(wfc, K, V, Ax, Ay, g, xDim, yDim, dx, dy)
        println()
    end

end

calculate("Params.dat", "../data/")


using Knet: logsoftmax, softmax

function kl_div(input_logits, target_logits)
    logsoft_input = logsoftmax(input_logits; dims=1)
    soft_target = softmax(target_logits; dims=1)
    logsoft_target = logsoftmax(target_logits; dims=1)
    return mean(soft_target .* (logsoft_target .- logsoft_input))
end

function mse_loss(y_hat, y)
    return mean((y - y_hat)^2)
end

function dice_loss(score, target)
    smooth = 1e-5
    intersect = sum(score .* target)
    y_sum = sum(target .* target)
    z_sum = sum(score .* score)
    loss = (2 .* intersect .+ smooth) / (z_sum + y_sum .+ smooth)
    loss = 1 .- loss
    return loss
end


#===============================================================================

  Code for investment under uncertainty by Heterogenous firms

  Translated from matlab code from Ben Moll's website:
      www.princeton.edu/~moll/HACTproject.htm

================================================================================#

using LinearAlgebra, SparseArrays, Plots

# Set up model parameters
ρ = 0.05  #discount rate
α = 0.5 # Cobb-douglas production function
Θ = 2.7 # quadratic adjustment cost
δ = 0.025 # depreciation rate

# Parameters for the Ornstein-Uhlenbeck process for z
var = 0.026^2
zmean = exp(var/2)
corr = 0.859
θ = -log(corr)
σ_sq = 2*θ*var

# Set up the gridspace for z
J = 40
zmin = zmean*0.6
zmax=zmean*1.4
z = LinRange(zmin,zmax,J)
dz = (zmax-zmin)/(J-1)
dz_sq = dz^2

# Set up gridspace for capital
H =100
kmin=1
kmax = 100
k = LinRange(kmin,kmax, H)
dk = (kmax-kmin)/(H-1)

kk = k*ones(1,J)
zz = (z*ones(1,H))'

# Set up terms for drift and variance of the UO process that depend on zz
μ = (-θ*log.(z) .+ σ_sq/2).*z
Σ_sq = σ_sq*(z.^2)

# Production function
F = zz.*kk.^α

# Parameters for our loop
maxit = 20
crit=10^(-6)
Δ = 1000

# Empty spaces for saving things
Vkb, Vkf, c = [zeros(H,J) for i in 1:3]
dist =[]
x_all=[]

# Create the matrix that summarizes the evolution of our system in the z dimension
 yy = (-Σ_sq/dz_sq + min.(μ,0)/dz - max.(μ,0)/dz)
 χ = Σ_sq/(2*dz_sq)- min.(μ,0)/dz
 ζ = max.(μ,0)/dz + Σ_sq/(2*dz_sq)

 # Define the Diagonals of this matrix
 updiag_z = zeros(H,1)
 	for j = 1:J
		global updiag_z =[updiag_z; repeat([ζ[j]], H, 1)]
	end

 centerdiag_z=repeat([χ[1]+yy[1]],H,1)
	for j = 2:J-1
		global centerdiag_z = [centerdiag_z; repeat([yy[j]], H, 1)]
	end
 centerdiag_z=[centerdiag_z; repeat([yy[J]+ζ[J]], H, 1)]

lowdiag_z = repeat([χ[2]], H, 1)
	for j=3:J
		global lowdiag_z = [lowdiag_z; repeat([χ[j]],H,1)]
	end
# spdiags in Matlab allows for automatic trimming/adding of zeros
    # spdiagm does not do this

Bswitch =  spdiagm(0=>centerdiag_z[:], -H=>lowdiag_z[:], H=>updiag_z[H+1:end-H])

# Inital guess for v
v0 = (F-δ.*kk-0.5*Θ*δ^2 .*kk)/ρ
global v =v0

for n in 1:maxit
	global V=v
    #forward difference
    Vkf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/dk
    Vkf[H,:] = (1+Θ*δ)*ones(1,J)#will never be used, but impose state constraint a<=amax just in case
    #backward difference
    Vkb[2:H,:] = (V[2:H,:]-V[1:H-1,:])/dk
    Vkb[1,:] = (1+Θ*δ)*ones(1,J) #state constraint boundary condition

    #I_concave = Vab .> Vaf #indicator whether value function is concave (problems arise if this is not the case)

    # investment and savings with forward difference
    xf = (Vkf.-1)/Θ.*kk
    ssf = xf-δ.*kk
	Hf = F -xf -0.5*Θ*(xf./kk).^2 .*kk + Vkf.*ssf
    # consumption and savings with backward difference
	xb = (Vkb.-1)/Θ.*kk
    ssb = xb-δ.*kk
	Hb = F -xb -0.5*Θ*(xb./kk).^2 .*kk + Vkb.*ssb
    # investment at steady state
    x0 = δ.*kk

    #= dV_upwind makes a choice of forward or backward differences based on
     the sign of the drift    =#
	Ieither = (1 .- (ssf.>0)) .* (1 .- (ssb.<0))
	Iunique = (ssb.<0).*(1 .- (ssf.>0)) + (1 .- (ssb.<0)).*(ssf.>0)
	Iboth = (ssb.<0).*(ssf.>0)
    If= Iunique.*(ssf .> 0) + Iboth.*(Hf.>=Hb)  #positive drift → forward difference
    Ib = Iunique.*(ssb .< 0) + Iboth.*(Hb.>=Hf) #negative drift → backward difference
    I0 = Ieither  #at steady state

	x = xf.*If + xb.*Ib + x0.*I0
    profits = F-x-0.5*Θ*(x./kk).^2. .*kk

    # CONSTRUCT MATRIX
    X = -ssb.*Ib/dk
    Y = - ssf.*If/dk  + ssb.*Ib/dk
    Z = ssf.*If/dk

	# Define the Diagonals of this transition matrix
	updiag = 0
       for j = 1:J
           updiag =[updiag; Z[1:H-1,j]; 0]
       end
    updiag =(updiag[:])

    centerdiag=reshape(Y, H*J, 1)
    centerdiag = (centerdiag[:]) # for tuples

   lowdiag = X[2:H, 1]
       for j = 2:J
           lowdiag = [lowdiag; 0; X[2:H,j]]
       end
   lowdiag=(lowdiag)

   # spdiags in Matlab allows for automatic trimming/adding of zeros
       # spdiagm does not do this
  	AA = sparse(Diagonal(centerdiag))+ [zeros(1, H*J); sparse(Diagonal(lowdiag)) zeros(H*J-1,1)] + sparse(Diagonal(updiag))[2:end, 1:(H*J)] # trim first element
	A = AA + Bswitch

    B = (1/Δ + ρ)*sparse(I, H*J, H*J) - A

    profits_stacked = reshape(profits, H*J,1)
    V_stacked = reshape(V, H*J,1)

    b = profits_stacked + V_stacked/Δ
    V_stacked = B\b #SOLVE SYSTEM OF EQUATIONS

    global V = reshape(V_stacked, H,J)

    Vchange = V-v
	global v = V

	push!(dist, findmax(abs.(Vchange))[1])
	push!(x_all,x)

	if dist[n]<crit
		println("Value function converged iteration = $(n)")
		break
	end
end

plot(k,v, legend=false)

plot(dist, legend=false)

kdot=x_all[end]-δ.*kk
plot(kk,kdot,label="")
plot!(kk,zeros(H,1), line=:dash, color=:black, label="")


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("useful_routines.jl")

function bulk_modulus(l, m)
    return l + 2m / 3
end

function lame_lambda(k, m)
    return k - 2m / 3
end

function analytical_coefficient_matrix(inradius, outradius, ls, ms, lc, mc)
    a = zeros(3, 3)
    a[1, 1] = inradius
    a[1, 2] = -inradius
    a[1, 3] = -1.0 / inradius
    a[2, 1] = 2 * (lc + mc)
    a[2, 2] = -2 * (ls + ms)
    a[2, 3] = 2ms / inradius^2
    a[3, 2] = 2(ls + ms)
    a[3, 3] = -2ms / outradius^2
    return a
end

function analytical_coefficient_rhs(ls, ms, theta0)
    r = zeros(3)
    Ks = bulk_modulus(ls, ms)
    r[2] = -Ks * theta0
    r[3] = Ks * theta0
    return r
end

struct AnalyticalSolution
    inradius::Any
    outradius::Any
    center::Any
    A1c::Any
    A1s::Any
    A2s::Any
    ls::Any
    ms::Any
    lc::Any
    mc::Any
    theta0::Any
    function AnalyticalSolution(
        inradius,
        outradius,
        center,
        ls,
        ms,
        lc,
        mc,
        theta0,
    )
        a = analytical_coefficient_matrix(inradius, outradius, ls, ms, lc, mc)
        r = analytical_coefficient_rhs(ls, ms, theta0)
        coeffs = a \ r
        new(
            inradius,
            outradius,
            center,
            coeffs[1],
            coeffs[2],
            coeffs[3],
            ls,
            ms,
            lc,
            mc,
            theta0,
        )
    end
end

function radial_displacement(A::AnalyticalSolution, r)
    if r <= A.inradius
        return A.A1c * r
    else
        return A.A1s * r + A.A2s / r
    end
end

function (A::AnalyticalSolution)(x)
    relpos = x - A.center
    r = sqrt(relpos' * relpos)
    ur = radial_displacement(A, r)
    if ur ≈ 0.0
        [0.0, 0.0]
    else
        costheta = (x[1] - A.center[1]) / r
        sintheta = (x[2] - A.center[2]) / r
        u1 = ur * costheta
        u2 = ur * sintheta
        return [u1, u2]
    end
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function displacement_error(
    width,
    center,
    inradius,
    outradius,
    stiffness,
    theta0,
    nelmts,
    polyorder,
    numqp,
    penaltyfactor;
    eta = 1,
)
    L = W = width

    lambda1, mu1 = CutCellDG.lame_coefficients(stiffness, +1)
    lambda2, mu2 = CutCellDG.lame_coefficients(stiffness, -1)
    transfstress =
        CutCellDG.plane_strain_transformation_stress(lambda1, mu1, theta0)

    dx = width / nelmts
    meanmoduli = 0.5 * (lambda1 + lambda2 + mu1 + mu2)
    penalty = penaltyfactor / dx * meanmoduli

    analyticalsolution = AnalyticalSolution(
        inradius,
        outradius,
        center,
        lambda1,
        mu1,
        lambda2,
        mu2,
        theta0,
    )

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCellDG.DGMesh([0.0, 0.0], [width, width], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs = CutCellDG.levelset_coefficients(
        x -> -circle_distance_function(x, center, inradius),
        mesh,
    )

    cutmesh = CutCellDG.CutMesh(mesh, levelset, levelsetcoeffs)
    cellquads =
        CutCellDG.CellQuadratures(cutmesh, levelset, levelsetcoeffs, numqp)
    interfacequads =
        CutCellDG.InterfaceQuadratures(cutmesh, levelset, levelsetcoeffs, numqp)
    facequads =
        CutCellDG.FaceQuadratures(cutmesh, levelset, levelsetcoeffs, numqp)

    mergedwithcell, hasmergedcells = CutCellDG.merge_tiny_cells_in_mesh!(
        cutmesh,
        cellquads,
        facequads,
        interfacequads,
    )
    mergedmesh = CutCellDG.MergedMesh(cutmesh, mergedwithcell)

    sysmatrix = CutCellDG.SystemMatrix()
    sysrhs = CutCellDG.SystemRHS()

    CutCellDG.assemble_displacement_bilinear_forms!(
        sysmatrix,
        basis,
        cellquads,
        stiffness,
        mergedmesh,
    )
    CutCellDG.assemble_bulk_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        cellquads,
        mergedmesh,
    )
    CutCellDG.assemble_interelement_condition!(
        sysmatrix,
        basis,
        facequads,
        stiffness,
        mergedmesh,
        penalty,
        eta,
    )
    CutCellDG.assemble_interelement_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        mergedmesh,
    )

    CutCellDG.assemble_coherent_interface_condition!(
        sysmatrix,
        basis,
        interfacequads,
        stiffness,
        mergedmesh,
        penalty,
        eta,
    )
    CutCellDG.assemble_coherent_interface_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        interfacequads,
        mergedmesh,
    )

    CutCellDG.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        analyticalsolution,
        basis,
        facequads,
        stiffness,
        mergedmesh,
        x -> onboundary(x, L, W),
        penalty,
    )
    CutCellDG.assemble_penalty_displacement_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        mergedmesh,
        x -> onboundary(x, L, W),
    )

    matrix = CutCellDG.sparse_displacement_operator(sysmatrix, mergedmesh)
    rhs = CutCellDG.displacement_rhs_vector(sysrhs, mergedmesh)

    solution = matrix \ rhs
    nodaldisplacement = reshape(solution, 2, :)

    err1 = interface_maxnorm_error(
        nodaldisplacement,
        analyticalsolution,
        +1,
        basis,
        interfacequads,
        mergedmesh,
    )
    err2 = interface_maxnorm_error(
        nodaldisplacement,
        analyticalsolution,
        -1,
        basis,
        interfacequads,
        mergedmesh,
    )

    exactnorm = maxnorm_on_interface(
        analyticalsolution,
        interfacequads,
        +1,
        mergedmesh,
        2,
    )

    return err1 ./ exactnorm, err2 ./ exactnorm
end

lambda1, mu1 = 100.0, 80.0
lambda2, mu2 = 80.0, 60.0
theta0 = -0.067
stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)

width = 1.0
penaltyfactor = 1e2

polyorder = 2
numqp = required_quadrature_order(polyorder) + 2

center = [width / 2, width / 2]
inradius = width / 3
outradius = width

powers = [3, 4, 5]
nelmts = [2^p + 1 for p in powers]

err = [
    displacement_error(
        width,
        center,
        inradius,
        outradius,
        stiffness,
        theta0,
        ne,
        polyorder,
        numqp,
        penaltyfactor,
    ) for ne in nelmts
]

producterr = [er[1] for er in err]
parenterr = [er[2] for er in err]

productu1err = [p[1] for p in producterr]
productu2err = [p[2] for p in producterr]

parentu1err = [p[1] for p in parenterr]
parentu2err = [p[2] for p in parenterr]

dx = 1.0 ./ nelmts

productu1rate = convergence_rate(dx,productu1err)
productu2rate = convergence_rate(dx,productu2err)

parentu1rate = convergence_rate(dx,parentu1err)
parentu2rate = convergence_rate(dx,parentu2err)

@test all(productu1rate .> 2.6)
@test all(productu2rate .> 2.6)
@test all(parentu1rate .> 2.6)
@test all(parentu2rate .> 2.6)


mutable struct LogNormalFwdRateEuler{M <: AbstractMarketModel, G <: BrownianGenerator} <: AbstractMarketModelEvolver
  marketModel::M
  numeraires::Vector{Int}
  initialStep::Int
  generator::G
  fixedDrifts::Vector{Vector{Float64}}
  numberOfRates::Int
  numberOfFactors::Int
  curveState::LMMCurveState
  currentStep::Int
  forwards::Vector{Float64}
  displacement::Vector{Float64}
  logForwards::Vector{Float64}
  initialLogForwards::Vector{Float64}
  drifts1::Vector{Float64}
  initialDrifts::Vector{Float64}
  brownians::Vector{Float64}
  correlatedBrownians::Vector{Float64}
  alive::Vector{Int}
  calculators::Vector{LMMDriftCalculator}
end

function LogNormalFwdRateEuler(marketModel::AbstractMarketModel, factory::BrownianGeneratorFactory, numeraires::Vector{Int}, initialStep::Int = 1)
  numberOfRates = marketModel.numberOfRates
  numberOfFactors = marketModel.numberOfFactors
  curveState = LMMCurveState(marketModel.evolution.rateTimes)
  forwards = copy(marketModel.initialRates)
  displacements = copy(marketModel.displacements)
  logForwards = Vector{Float64}(undef, numberOfRates)
  initialLogForwards = Vector{Float64}(undef, numberOfRates)
  drifts1 = Vector{Float64}(undef, numberOfRates)
  initialDrifts = Vector{Float64}(undef, numberOfRates)
  brownians = Vector{Float64}(undef, numberOfFactors)
  correlatedBrownians = Vector{Float64}(undef, numberOfRates)
  alive = marketModel.evolution.firstAliveRate

  check_compatibility(marketModel.evolution, numeraires)

  steps = number_of_steps(marketModel.evolution)

  generator = create(factory, numberOfFactors, steps - (initialStep-1))

  currentStep = initialStep

  calculators = Vector{LMMDriftCalculator}(undef, steps)
  fixedDrifts = Vector{Vector{Float64}}(undef, steps)

  @inbounds @simd for j = 1:steps
    A = marketModel.pseudoRoots[j]
    calculators[j] = LMMDriftCalculator(A, displacements, marketModel.evolution.rateTaus, numeraires[j], alive[j])
    fixed = Vector{Float64}(undef, numberOfRates)

    for k in eachindex(fixed)
      # variance_ = dot(A[:, k], A[:, k])
      variance_ = dot(A[k, :], A[k, :])
      fixed[k] = -0.5 * variance_
    end

    fixedDrifts[j] = fixed
  end

  lognorm = LogNormalFwdRateEuler(marketModel, numeraires, initialStep, generator, fixedDrifts, numberOfRates, numberOfFactors, curveState, currentStep, forwards,
                              displacements, logForwards, initialLogForwards, drifts1, initialDrifts, brownians, correlatedBrownians, alive, calculators)

  set_forwards!(lognorm, marketModel.initialRates)
  return lognorm
end

function set_forwards!(lognorm::LogNormalFwdRateEuler, forwards::Vector{Float64})
  length(forwards) == lognorm.numberOfRates || error("mismatch between forwards and rate times")
  @inbounds @simd for i in eachindex(lognorm.initialLogForwards)
    lognorm.initialLogForwards[i] = log(forwards[i] + lognorm.displacement[i])
  end

  compute!(lognorm.calculators[lognorm.initialStep], forwards, lognorm.initialDrifts)

  return lognorm
end

function start_new_path!(lognorm::LogNormalFwdRateEuler)
  lognorm.currentStep = lognorm.initialStep
  lognorm.logForwards = copy(lognorm.initialLogForwards)

  return next_path!(lognorm.generator)
end

function advance_step!(lognorm::LogNormalFwdRateEuler)
  # we are going from T1 to T2

  # a) compute drifts D1 at T1
  if lognorm.currentStep > lognorm.initialStep
    compute!(lognorm.calculators[lognorm.currentStep], lognorm.forwards, lognorm.drifts1)
  else
    lognorm.drifts1 = copy(lognorm.initialDrifts)
  end

  # b) evolve forwards up to T2 using D1
  weight = next_step!(lognorm.generator, lognorm.brownians)
  A = copy(lognorm.marketModel.pseudoRoots[lognorm.currentStep])
  fixedDrift = copy(lognorm.fixedDrifts[lognorm.currentStep])

  alive = lognorm.alive[lognorm.currentStep]
  @inbounds @simd for i = alive:lognorm.numberOfRates
    lognorm.logForwards[i] += lognorm.drifts1[i] + fixedDrift[i]
    lognorm.logForwards[i] += dot(A[i, :], lognorm.brownians)
    lognorm.forwards[i] = exp(lognorm.logForwards[i]) - lognorm.displacement[i]
  end

  # same as PC evolver with two steps dropped

  # c) update curve state
  set_on_forward_rates!(lognorm.curveState, lognorm.forwards)
  lognorm.currentStep += 1

  return weight
end

current_state(lognorm::LogNormalFwdRateEuler) = lognorm.curveState
brownians_this_step(lognorm::LogNormalFwdRateEuler) = lognorm.brownians

# Clone #
clone(lognorm::LogNormalFwdRateEuler) = LogNormalFwdRateEuler(clone(lognorm.marketModel), copy(lognorm.numeraires), lognorm.initialStep, clone(lognorm.generator),
                                      deepcopy(lognorm.fixedDrifts), lognorm.numberOfRates, lognorm.numberOfFactors, clone(lognorm.curveState),
                                      lognorm.currentStep, copy(lognorm.forwards), copy(lognorm.displacement), copy(lognorm.logForwards), copy(lognorm.initialLogForwards),
                                      copy(lognorm.drifts1), copy(lognorm.initialDrifts), copy(lognorm.brownians), copy(lognorm.correlatedBrownians),
                                      copy(lognorm.alive), deepcopy(lognorm.calculators))


import Test: @test, @testset

@testset "Radial propagation" begin
import Luna
import Luna: Grid, Maths, PhysData, Nonlinear, Ionisation, NonlinearRHS, Output, Stats, LinearOps, Plotting
import Luna.PhysData: wlfreq
import FFTW
import Hankel
import LinearAlgebra: norm

gas = :Ar
pres = 1.2
τ = 20e-15
λ0 = 800e-9
w0 = 40e-6
energy = 1e-12
L = 0.6
R = 4e-3
N = 128

grid = Grid.RealGrid(L, 800e-9, (400e-9, 2000e-9), 0.2e-12)
q = Hankel.QDHT(R, N, dim=2)

energyfun, energyfun_ω = Fields.energyfuncs(grid, q)

function prop(E, z)
    Eω = FFTW.rfft(E, 1)
    Eωk = q * Eω
    kzsq = @. (grid.ω/PhysData.c)^2 - (q.k^2)'
    kzsq[kzsq .< 0] .= 0
    kz = sqrt.(kzsq)
    @. Eωk *= exp(-1im * z * (kz - grid.ω/PhysData.c))
    Eω = q \ Eωk
    E = FFTW.irfft(Eω, length(grid.t), 1)
    return E
end

dens0 = PhysData.density(gas, pres)
densityfun(z) = dens0
ionpot = PhysData.ionisation_potential(gas)
ionrate = Ionisation.ionrate_fun!_PPTcached(gas, λ0)
responses = (Nonlinear.Kerr_field(PhysData.γ3_gas(gas)),
             Nonlinear.PlasmaCumtrapz(grid.to, grid.to, ionrate, ionpot))
linop = LinearOps.make_const_linop(grid, q, PhysData.ref_index_fun(gas, pres))
normfun = NonlinearRHS.const_norm_radial(grid, q, PhysData.ref_index_fun(gas, pres))

inputs = Fields.GaussGaussField(λ0=λ0, τfwhm=τ, energy=energy, w0=w0, propz=-0.3)

Eω, transform, FT = Luna.setup(grid, q, densityfun, normfun, responses, inputs)
output = Output.MemoryOutput(0, grid.zmax, 201)
Luna.run(Eω, grid, linop, transform, FT, output)

Erout = (q \ output.data["Eω"])
Iωr = abs2.(Erout)

ω0idx = argmin(abs.(grid.ω .- 2π*PhysData.c/λ0))

Iλ0 = Iωr[ω0idx, :, :]
λ0 = 2π*PhysData.c/grid.ω[ω0idx]
w1 = w0*sqrt(1+(L/2*λ0/(π*w0^2))^2)
Iλ0_analytic = Maths.gauss.(q.r, w1/2)*(w0/w1)^2 # analytical solution (in paraxial approx)
Ir = Maths.normbymax(Iλ0[:, end])
Ira = Maths.normbymax(Iλ0_analytic)
@test maximum(abs.(Ir .- Ira)/norm(Ir)) < 1.5e-4
end


using ImplicitEquations
using Base.Test

f(x,y) = y-x
g(x,y) = y+x

ImplicitEquations.GRAPH(f == 0, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f <= 0, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f <  g, -5, 5, -5, 5, 2^4, 2^4)

ImplicitEquations.GRAPH(f ≶  g, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f ≷  g, -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(neq(f, g), -5, 5, -5, 5, 2^4, 2^4)

ImplicitEquations.GRAPH(eq(f, g), -5, 5, -5, 5, 2^4, 2^4)
ImplicitEquations.GRAPH(f ⩵ g, -5, 5, -5, 5, 2^4, 2^4)


using MultivariatePolynomials
using JuMP
using PolyJuMP
using SumOfSquares
using DynamicPolynomials
using Mosek

# Create symbolic variables (not JuMP decision variables)
@polyvar x1 x2

# Create a Sum of Squares JuMP model with the Mosek solver
m = SOSModel(solver = MosekSolver())

# Create a JuMP decision variable for the lower bound
@variable m γ

# f(x) is the Goldstein-Price function
f1 = x1+x2+1
f2 = 19-14*x1+3*x1^2-14*x2+6*x1*x2+3*x2^2
f3 = 2*x1-3*x2
f4 = 18-32*x1+12*x1^2+48*x2-36*x1*x2+27*x2^2

f = (1+f1^2*f2)*(30+f3^2*f4)

# Constraints f(x) - γ to be sum of squares
@constraint m f >= γ

@objective m Max γ

status = solve(m)

# The lower bound found is 3
println(getobjectivevalue(m))


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

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

                     @variable(m, -3.517322036224115 <= q <= 6.144134546328734)

                     add_NL_constraint(m, :(1/(1 + exp(-(0.9006950404367178 + -0.8785024386485549*1/(1 + exp(-(0.872987896263127 + 0.13249617278479375*$(x[1]) + 0.7878182386475117*$(x[2]) + -0.4203533662190604*$(x[3]) + -0.7649366056085984*$(x[4]) + -0.57871509222098*$(x[5])))) + 0.18543992065441683*1/(1 + exp(-(0.8256824299447256 + -0.6033988061517293*$(x[1]) + -0.15324994390362567*$(x[2]) + 0.6933635130851963*$(x[3]) + 0.1238690975426957*$(x[4]) + 0.2325901400223671*$(x[5]))))))) + 1/(1 + exp(-(0.10821024944347046 + 0.16585705930469619*1/(1 + exp(-(0.872987896263127 + 0.13249617278479375*$(x[1]) + 0.7878182386475117*$(x[2]) + -0.4203533662190604*$(x[3]) + -0.7649366056085984*$(x[4]) + -0.57871509222098*$(x[5])))) + 0.9368217650268278*1/(1 + exp(-(0.8256824299447256 + -0.6033988061517293*$(x[1]) + -0.15324994390362567*$(x[2]) + 0.6933635130851963*$(x[3]) + 0.1238690975426957*$(x[4]) + 0.2325901400223671*$(x[5]))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using ModelingToolkit
using NeuralPDE
using Flux
using DiffEqFlux
using GalacticOptim

# using CUDA
# CUDA.allowscalar(false)

function NeuralPDE.PhysicsInformedNN(dx, chain::Chain, initθ; autodiff=false, strategy = GridTraining(), kwargs...)
   indim = size(nnet.layers[1].W, 2)
   outdim = size(chain.layers[end].W, 1)
   isuinplace = outdim == 1
   φ = get_phi(chain, isuinplace)
   derivative = get_derivative(indim, φ, autodiff, isuinplace)
   PhysicsInformedNN(dx, chain, initθ, φ, autodiff, derivative, strategy, kwargs)
end

function fitPDESystem!(K, α, µ, ρ₀, ρ₀ˢ, ρ₀ˡ,
                       F₁, F₂,
                       tspan, x₁span, x₂span,
                       Δt, Δx₁, Δx₂;
                       maxiters=1000, strategy = GridTraining())

   @parameters t, x₁, x₂, θ

    u₁ = Variable{Float32}(:u₁)
    u₂ = Variable{Float32}(:u₂)
    v₁ = Variable{Float32}(:v₁)
    v₂ = Variable{Float32}(:v₂)
   σ₁₁ = Variable{Float32}(:σ₁₁)
   σ₁₂ = Variable{Float32}(:σ₁₂)
   σ₂₂ = Variable{Float32}(:σ₂₂)
     p = Variable{Float32}(:p)

   @derivatives Dt'~t
   @derivatives Dx₁'~x₁
   @derivatives Dx₂'~x₂

   # System of PDEs
   equations = [
      Dt(u₁(t,x₁,x₂,θ)) + ( Dx₁(σ₁₁(t,x₁,x₂,θ)) + Dx₂(σ₁₂(t,x₁,x₂,θ)) )/ρ₀ˢ + Dx₁(p(t,x₁,x₂,θ))/ρ₀ - F₁(t,x₁,x₂) ~ 0,
      Dt(u₂(t,x₁,x₂,θ)) + ( Dx₁(σ₁₂(t,x₁,x₂,θ)) + Dx₂(σ₂₂(t,x₁,x₂,θ)) )/ρ₀ˢ + Dx₂(p(t,x₁,x₂,θ))/ρ₀ - F₂(t,x₁,x₂) ~ 0,
      Dt(v₁(t,x₁,x₂,θ)) + Dx₁(p(t,x₁,x₂,θ))/ρ₀ - F₁(t,x₁,x₂) ~ 0,
      Dt(v₂(t,x₁,x₂,θ)) + Dx₂(p(t,x₁,x₂,θ))/ρ₀ - F₂(t,x₁,x₂) ~ 0,
      Dt(σ₁₁(t,x₁,x₂,θ)) + (ρ₀ˡ/ρ₀ * K + 4µ/3)Dx₁(u₁(t,x₁,x₂,θ)) + (ρ₀ˡ/ρ₀ * K - 2µ/3)Dx₂(u₂(t,x₁,x₂,θ)) - ρ₀ˢ/ρ₀ * K * ( Dx₁(v₁(t,x₁,x₂,θ)) + Dx₂(v₂(t,x₁,x₂,θ)) ) ~ 0,
      Dt(σ₁₂(t,x₁,x₂,θ)) + µ * ( Dx₂(u₁(t,x₁,x₂,θ)) + Dx₁(u₂(t,x₁,x₂,θ)) ) ~ 0,
      Dt(σ₂₂(t,x₁,x₂,θ)) + (ρ₀ˡ/ρ₀ * K - 2µ/3)Dx₁(u₁(t,x₁,x₂,θ)) + (ρ₀ˡ/ρ₀ * K + 4µ/3)Dx₂(u₂(t,x₁,x₂,θ)) - ρ₀ˢ/ρ₀ * K * ( Dx₁(v₁(t,x₁,x₂,θ)) + Dx₂(v₂(t,x₁,x₂,θ)) ) ~ 0,
      Dt(p(t,x₁,x₂,θ)) - (K - α*ρ₀*ρ₀ˢ)*( Dx₁(u₁(t,x₁,x₂,θ)) + Dx₂(u₂(t,x₁,x₂,θ)) ) + α*ρ₀*ρ₀ˡ*( Dx₁(v₁(t,x₁,x₂,θ)) + Dx₂(v₂(t,x₁,x₂,θ)) ) ~ 0
   ]

   # Initial and boundary conditions
   boundconds = [
        # initial conditions at t = 0
       u₁(tspan.lower, x₁, x₂, θ) ~ 0,
       u₂(tspan.lower, x₁, x₂, θ) ~ 0,
       v₁(tspan.lower, x₁, x₂, θ) ~ 0,
       v₂(tspan.lower, x₁, x₂, θ) ~ 0,
      σ₁₁(tspan.lower, x₁, x₂, θ) ~ 0,
      σ₁₂(tspan.lower, x₁, x₂, θ) ~ 0,
      σ₂₂(tspan.lower, x₁, x₂, θ) ~ 0,
        p(tspan.lower, x₁, x₂, θ) ~ 0,
      # boundary conditions at x₁ = x₁⁻
       u₁(t, x₁span.lower, x₂, θ) ~ 0,
       u₂(t, x₁span.lower, x₂, θ) ~ 0,
       v₁(t, x₁span.lower, x₂, θ) ~ 0,
       v₂(t, x₁span.lower, x₂, θ) ~ 0,
      σ₁₁(t, x₁span.lower, x₂, θ) ~ 0,
      σ₁₂(t, x₁span.lower, x₂, θ) ~ 0,
      σ₂₂(t, x₁span.lower, x₂, θ) ~ 0,
        p(t, x₁span.lower, x₂, θ) ~ 0,
      # boundary conditions at x₁ = x₁⁺
       u₁(t, x₁span.upper, x₂, θ) ~ 0,
       u₂(t, x₁span.upper, x₂, θ) ~ 0,
       v₁(t, x₁span.upper, x₂, θ) ~ 0,
       v₂(t, x₁span.upper, x₂, θ) ~ 0,
      σ₁₁(t, x₁span.upper, x₂, θ) ~ 0,
      σ₁₂(t, x₁span.upper, x₂, θ) ~ 0,
      σ₂₂(t, x₁span.upper, x₂, θ) ~ 0,
        p(t, x₁span.upper, x₂, θ) ~ 0,
      # boundary conditions at x₂ = x₂⁻
      σ₁₂(t, x₁, x₂span.lower, θ) ~ 0,
      # boundary conditions at x₂ = x₂⁺
       u₁(t, x₁, x₂span.upper, θ) ~ 0,
       u₂(t, x₁, x₂span.upper, θ) ~ 0,
       v₁(t, x₁, x₂span.upper, θ) ~ 0,
       v₂(t, x₁, x₂span.upper, θ) ~ 0,
      σ₁₁(t, x₁, x₂span.upper, θ) ~ 0,
      σ₁₂(t, x₁, x₂span.upper, θ) ~ 0,
      σ₂₂(t, x₁, x₂span.upper, θ) ~ 0,
        p(t, x₁, x₂span.upper, θ) ~ 0
   ]
   if ρ₀ˡ == 0
      push!(boundconds, σ₂₂(t, x₁, x₂span.lower, θ) ~ -p(t, x₁, x₂span.lower, θ))
   else
      push!(boundconds, σ₂₂(t, x₁, x₂span.lower, θ) ~ 0, p(t, x₁, x₂span.lower, θ) ~ 0)
   end

   # Time and space domains
   domain = [t ∈ tspan, x₁ ∈ x₁span, x₂ ∈ x₂span]

   # Neural network
   nnet = FastChain(FastDense(3,  16, Flux.selu),
                    FastDense(16, 16, Flux.σ),
                    FastDense(16, 16, Flux.σ),
                    FastDense(16, 8)) #|> gpu
   # nnet = Chain(Dense(3,  16, Flux.selu),
   #              Dense(16, 16, Flux.σ),
   #              Dense(16, 16, Flux.σ),
   #              Dense(16, 8)) #|> gpu

   # Initial neural network training parameters
   initθ = initial_params(nnet) #|> gpu

   discretization = PhysicsInformedNN([Δt, Δx₁, Δx₂], nnet, initθ; strategy)

   pde_system = PDESystem(equations, boundconds, domain, [t, x₁, x₂], [u₁, u₂, v₁, v₂, σ₁₁, σ₁₂, σ₂₂, p])

   problem = NeuralPDE.discretize(pde_system, discretization)

   optimizer = Flux.ADAM(0.1f0)
   callback = function (p, l)
      println("Current loss is: $l")
      return false
   end

   @time res = GalacticOptim.solve(problem, optimizer; progress=false, cb=callback, maxiters)

   ϕ = discretization.phi
   θopt = res.minimizer

   return function sol(t, x₁, x₂)
      ϕ([t, x₁, x₂], θopt)
   end
end

# Specify constants
const K, α, µ, ρ₀, ρ₀ˢ, ρ₀ˡ = let
   ρ₀ᶠˢ = 1400
   ρ₀ᶠˡ = 997
   cp₁ = 2000
   cp₂ = 1300
   cs = 1300
   d₀ = 0.3

   ρ₀ˢ = (1 - d₀)ρ₀ᶠˢ
   ρ₀ˡ = d₀ * ρ₀ᶠˡ
   ρ₀ = ρ₀ˢ + ρ₀ˡ
   µ  = ρ₀ˢ * cs^2
   K  = (cp₁^2 + cp₂^2 - 8/3*ρ₀ˡ/ρ₀*cs^2 - √((cp₁^2 - cp₂^2)^2 - 64/9*cs^4*ρ₀ˡ*ρ₀ˢ/ρ₀^2) ) * ρ₀ * ρ₀ˢ/(2ρ₀ˡ)
   α₃ = (cp₁^2 + cp₂^2 - 8/3*ρ₀ˢ/ρ₀*cs^2 - √((cp₁^2 - cp₂^2)^2 - 64/9*cs^4*ρ₀ˡ*ρ₀ˢ/ρ₀^2) ) / (2ρ₀^2)
   α  = ρ₀ * α₃ + K/ρ₀^2

   Float32.((K, α, µ, ρ₀, ρ₀ˢ, ρ₀ˡ))
end

# Specify force functions F₁ and F₂
F₁(t, x₁, x₂) = 100exp(-t^2 - x₁^2 - x₂^2) * sin(t)
F₂(t, x₁, x₂) = 1000exp(-t^2 - x₁^2 - x₂^2) * cos(t)

# Specify variable intervals
 tspan = IntervalDomain(0.f0, 60.f0)
x₁span = IntervalDomain(-100.f0, 100.f0)
x₂span = IntervalDomain(0.f0, 100.f0)

# Specify discretization
Δt  = 10.f0
Δx₁ = 20.f0
Δx₂ = 10.f0

# Fit the neural network to the PDE system
w = fitPDESystem!(K, α, µ, ρ₀, ρ₀ˢ, ρ₀ˡ,
                  F₁, F₂,
                  tspan, x₁span, x₂span,
                  Δt, Δx₁, Δx₂; maxiters=1000)

w(tspan.lower, x₁span.lower, x₂span.lower)

w(1, 2, 3)


function trailing_edge_wing(settings::PyObject, ac::PyObject, freq::Array{Float64, 1}, M_0, c_0, rho_0, mu_0, theta, phi)

    ### ---------------- Wing trailing-edge noise ----------------
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 5
    delta_w_star = 0.37 * (ac.af_S_w / ac.af_b_w^2) * (rho_0 * M_0 * c_0 * ac.af_S_w / (mu_0 * ac.af_b_w))^(-0.2)

    # Determine configuration constant and the sound power
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 7
    if ac.af_clean_w == 0
        K_w = 4.464e-5
    elseif ac.af_clean_w == 1
        K_w = 7.075e-6
    end

    Pi_star_w = K_w * M_0^5 * delta_w_star

    # Determine directivity function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 8
    D_w = 4. * cos(phi * π / 180.)^2 * cos(theta / 2 * π / 180.)^2

    # Determine spectral distribution function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 10-11-12
    S_w = freq * delta_w_star * ac.af_b_w / (M_0 * c_0) * (1 - M_0 * cos(theta * π / 180.))
    if ac.af_delta_wing == 1
        F_w = 0.613 * (10 * S_w).^4 .* ((10 * S_w).^1.35 .+ 0.5).^(-4)
    elseif ac.af_delta_wing == 0
        F_w = 0.485 * (10 * S_w).^4 .* ((10 * S_w).^1.5 .+ 0.5).^(-4)
    end

    # Determine msap
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 1
    r_s_star_af = settings.r_0 / ac.af_b_w
    msap_w = 1 / (4 * π * r_s_star_af^2) / (1 - M_0 * cos(theta * π / 180.))^4 * (Pi_star_w * D_w * F_w)

    #Write output
    return msap_w
end
function trailing_edge_horizontal_tail(settings::PyObject, ac::PyObject, freq::Array{Float64, 1}, M_0, c_0, rho_0, mu_0, theta, phi)

    # ---------------- Horizontal tail trailing-edge noise ----------------
    # Trailing edge noise of the horizontal tail
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 5
    delta_h_star = 0.37 * (ac.af_S_h / ac.af_b_h^2) * (rho_0 * M_0 * c_0 * ac.af_S_h / (mu_0 * ac.af_b_h))^(-0.2)

    # Determine configuration constant and the sound power
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 7
    if ac.af_clean_h == 0
        K_h = 4.464e-5
    elseif ac.af_clean_h == 1
        K_h = 7.075e-6
    end
    Pi_star_h = K_h * M_0^5 * delta_h_star * (ac.af_b_h / ac.af_b_w)^2

    # Determine the directivity function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 8
    D_h = 4 * cos(phi * π / 180.)^2 * cos(theta / 2 * π / 180.)^2

    # Determine spectral distribution function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 10-11-12
    S_h = freq * delta_h_star * ac.af_b_h / (M_0 * c_0) * (1 - M_0 * cos(theta * π / 180.))
    F_h = 0.485 * (10 * S_h).^4 .* ((10 * S_h).^1.5 .+ 0.5).^(-4)

    # Determine msap
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 1
    r_s_star_af = settings.r_0 / ac.af_b_w
    msap_h = 1. / (4. * π * r_s_star_af^2) / (1 - M_0 * cos(theta * π / 180.))^4 * (Pi_star_h * D_h * F_h)

    return msap_h
end
function trailing_edge_vertical_tail(settings::PyObject, ac::PyObject, freq::Array{Float64, 1}, M_0, c_0, rho_0, mu_0, theta, phi)

    ### ---------------- Vertical tail trailing-edge noise ----------------
    delta_v_star = 0.37 * (ac.af_S_v / ac.af_b_v^2) * (rho_0 * M_0 * c_0 * ac.af_S_v / (mu_0 * ac.af_b_v))^(-0.2)

    # Trailing edge noise of the vertical tail
    if ac.af_clean_v == 0
        K_v = 4.464e-5
    elseif ac.af_clean_v == 1
        K_v = 7.075e-6
    end

    Pi_star_v = K_v * M_0^5 * delta_v_star * (ac.af_b_v / ac.af_b_w)^2

    # Determine directivity function
    D_v = 4 * sin(phi * π / 180.)^2 * cos(theta / 2 * π / 180.)^2

    # Determine spectral distribution function
    S_v = freq * delta_v_star * ac.af_b_v / (M_0 * c_0) * (1 - M_0 * cos(theta * π / 180.))

    if ac.af_delta_wing == 1
        F_v = 0.613 * (10 * S_v).^4 .* ((10 * S_v).^1.35 .+ 0.5).^(-4)
    elseif ac.af_delta_wing == 0
        F_v = 0.485 * (10 * S_v).^4 .* ((10 * S_v).^1.35 .+ 0.5).^(-4)
    end

    # Determine msap
    r_s_star_af = settings.r_0 / ac.af_b_w
    msap_v = 1. / (4 * π * r_s_star_af^2) / (1 - M_0 * cos(theta * π / 180.))^4 * (Pi_star_v * D_v * F_v)

    return msap_v
end
function leading_edge_slat(settings::PyObject, ac::PyObject, freq::Array{Float64, 1}, M_0, c_0, rho_0, mu_0, theta, phi)

    ### ---------------- Slat noise ----------------
    delta_w_star = 0.37 * (ac.af_S_w / ac.af_b_w^2) * (rho_0 * M_0 * c_0 * ac.af_S_w / (mu_0 * ac.af_b_w))^(-0.2)

    # Noise power
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 4
    Pi_star_les1 = 4.464e-5 * M_0^5 * delta_w_star  # Slat noise
    Pi_star_les2 = 4.464e-5 * M_0^5 * delta_w_star  # Added trailing edge noise

    # Determine the directivity function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 8
    D_les = 4 * cos(phi * π / 180.)^2 * cos(theta / 2 * π / 180.)^2

    # Determine spectral distribution function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 10-12-13
    S_les = freq * delta_w_star * ac.af_b_w / (M_0 * c_0) * (1 - M_0 * cos(theta * π / 180.))

    F_les1 = 0.613 * (10 * S_les).^4 .* ((10. * S_les).^1.5 .+ 0.5).^(-4)
    F_les2 = 0.613 * (2.19 * S_les).^4 .* ((2.19 * S_les).^1.5 .+ 0.5).^(-4)

    # Calculate msap
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 1
    r_s_star_af = settings.r_0 / ac.af_b_w
    msap_les = 1 / (4 * π * r_s_star_af^2) / (1 - M_0 * cos(theta * π / 180.))^4 * (Pi_star_les1 * D_les * F_les1 + Pi_star_les2 * D_les * F_les2)

    return msap_les
end
function trailing_edge_flap(settings::PyObject, ac::PyObject, freq::Array{Float64, 1}, M_0, c_0, theta, phi, theta_flaps)

    ### ---------------- Flap noise ----------------
    # Calculate noise power
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 14-15
    if ac.af_s < 3
        Pi_star_tef = 2.787e-4 * M_0^6 * ac.af_S_f / ac.af_b_w^2 * sin(theta_flaps * π / 180.)^2
    elseif ac.af_s == 3
        Pi_star_tef = 3.509e-4 * M_0^6 * ac.af_S_f / ac.af_b_w^2 * sin(theta_flaps * π / 180.)^2
    end

    # Calculation of the directivity function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 16
    D_tef = 3 * (sin(theta_flaps * π / 180.) * cos(theta * π / 180.) + cos(theta_flaps * π / 180.) * sin(theta * π / 180.) * cos(phi * π / 180.))^2
    # Strouhal number
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 19
    S_tef = freq * ac.af_S_f / (M_0 * ac.af_b_f * c_0) * (1 - M_0 * cos(theta * π / 180.))

    # Calculation of the spectral function
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 17-18
    if ac.af_s < 3
        F_tef = zeros(eltype(S_tef), (settings.N_f, ))
        for i in range(1, settings.N_f, step=1)
            if S_tef[i] .< 2.
                F_tef[i] = 0.0480 * S_tef[i]
            elseif 2. .< S_tef[i] .< 20.
                F_tef[i] = 0.1406 * S_tef[i].^(-0.55)
            else
                F_tef[i] = 216.49 * S_tef[i].^(-3)
            end
        end
    elseif ac.af_s == 3
        F_tef = zeros(eltype(S_tef), (settings.N_f, ))
        for i in range(1, settings.N_f, step=1)
            if S_tef[i] .< 2.
                F_tef[i] = 0.0257 * S_tef[i]
            elseif 2. .< S_tef[i] .< 20.
                F_tef[i] = 0.0536 * S_tef[i].^(-0.0625)
            else
                F_tef[i] = 17078 * S_tef[i].^(-3)
            end
        end
    end

    # Calculate msap
    # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 1
    r_s_star_af = settings.r_0 / ac.af_b_w
    msap_tef = 1. / (4 * π * r_s_star_af^2) / (1 - M_0 * cos(theta * π / 180.))^4 * (Pi_star_tef * D_tef * F_tef)

    return msap_tef
end
function landing_gear(settings::PyObject, ac::PyObject, freq::Array{Float64, 1}, M_0, c_0, theta, phi, I_landing_gear)

    ### ---------------- Landing-gear noise ----------------
    if I_landing_gear == 1
        # Calculate nose-gear noise
        # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 29
        S_ng = freq * ac.af_d_ng / (M_0 * c_0) * (1 - M_0 * cos(theta * π / 180.))

        # Calculate noise power and spectral distribution function
        # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 20-21-22-25-26-27-28
        if ac.af_n_ng == 1 || ac.af_n_ng == 2
            Pi_star_ng_w = 4.349e-4 * M_0^6 * ac.af_n_ng * (ac.af_d_ng / ac.af_b_w)^2
            Pi_star_ng_s = 2.753e-4 * M_0^6 * (ac.af_d_ng / ac.af_b_w)^2 * (ac.af_l_ng / ac.af_d_ng)
            F_ng_w = 13.59 * S_ng.^2 .* (12.5 .+ S_ng.^2).^(-2.25)
            F_ng_s = 5.32 * S_ng.^2 .* (30 .+ S_ng.^8).^(-1)
        elseif ac.af_n_ng == 4
            Pi_star_ng_w = 3.414 - 4 * M_0^6 * ac.af_n_ng * (ac.af_d_ng / ac.af_b_w)^2
            Pi_star_ng_s = 2.753e-4 * M_0^6 * (ac.af_d_ng / ac.af_b_w)^2 * (ac.af_l_ng / ac.af_d_ng)
            F_ng_w = 0.0577 * S_ng.^2 .* (1 .+ 0.25 * S_ng.^2).^(-1.5)
            F_ng_s = 1.28 * S_ng.^3 .* (1.06 .+ S_ng.^2).^(-3)
        end

        # Calculate main-gear noise
        # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 29
        S_mg = freq * ac.af_d_mg / (M_0 * c_0) * (1 - M_0 * cos(theta * π / 180.))
        # Calculate noise power and spectral distribution function
        # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 20-21-22-25-26-27-28
        if ac.af_n_mg == 1 || ac.af_n_mg == 2
            Pi_star_mg_w = 4.349e-4 * M_0^6 * ac.af_n_mg * (ac.af_d_mg / ac.af_b_w)^2
            Pi_star_mg_s = 2.753e-4 * M_0^6 * (ac.af_d_mg / ac.af_b_w)^2 * (ac.af_l_ng / ac.af_d_mg)
            F_ng_w = 13.59 * S_mg.^2 .* (12.5 .+ S_mg.^2).^(-2.25)
            F_ng_s = 5.32 * S_mg.^2 .* (30 .+ S_mg.^8).^(-1)
        elseif ac.af_n_mg == 4
            Pi_star_mg_w = 3.414e-4 * M_0^6 * ac.af_n_mg * (ac.af_d_mg / ac.af_b_w)^2
            Pi_star_mg_s = 2.753e-4 * M_0^6 * (ac.af_d_mg / ac.af_b_w)^2 * (ac.af_l_ng / ac.af_d_mg)
            F_mg_w = 0.0577 * S_mg.^2 .* (1 .+ 0.25 * S_mg.^2).^(-1.5)
            F_mg_s = 1.28 * S_mg.^3 .* (1.06 .+ S_mg.^2).^(-3)
        end

        # Directivity function
        # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 23-24
        D_w = 1.5 * sin(theta * π / 180.)^2
        D_s = 3 * sin(theta * π / 180.)^2 * sin(phi * π / 180.)^2
        # Calculate msap
        # Source: Zorumski report 1982 part 2. Chapter 8.8 Equation 1
        # If landing gear is down
        r_s_star_af = settings.r_0 / ac.af_b_w
        msap_lg = 1 / (4 * π * r_s_star_af^2) / (1 - M_0 * cos(theta * π / 180.))^4 * (
                              ac.af_N_ng * (Pi_star_ng_w * F_ng_w * D_w + Pi_star_ng_s * F_ng_s * D_s) +
                              ac.af_N_mg * (Pi_star_mg_w * F_mg_w * D_w + Pi_star_mg_s * F_mg_s * D_s))

    # If landing gear is up
    else
        msap_lg = zeros(eltype(I_landing_gear), 1)
    end

    return msap_lg
end

function airframe(settings::PyObject, data::PyObject, ac::PyObject, n_t::Int64, M_0, mu_0, c_0, rho_0, theta, phi, theta_flaps, I_landing_gear)

    # Initialize solution
    T = eltype(theta_flaps)
    msap_af = zeros(T, (n_t, settings.N_f))

    for i in range(1, n_t, step=1)

        # Airframe noise only when aircraft is moving
        if M_0[i] != 0.

            # Apply HSR-era airframe suppression/calibration levels
            # Source: validation noise assessment data set of NASA STCA (Berton et al., 2019)
            if settings.hsr_calibration
                # Suppression data
                data_angles = data.supp_af_angles
                data_freq = data.supp_af_freq
                data_supp = data.supp_af
                f_supp = LinearInterpolation((data_freq, data_angles), data_supp)
                supp = f_supp.(data.f, theta[i]*ones(T, (settings.N_f, )))
            else
                supp = ones(T, (settings.N_f, ))
            end

            # Add airframe noise components
            msap_j = zeros(T, (settings.N_f, ))
            if "wing" in ac.comp_lst
                msap_w = trailing_edge_wing(settings, ac, data.f, M_0[i], c_0[i], rho_0[i], mu_0[i], theta[i], phi[i])
                msap_j = msap_j .+ msap_w .* supp
            end
            if "tail_v" in ac.comp_lst
                msap_v = trailing_edge_vertical_tail(settings, ac, data.f, M_0[i], c_0[i], rho_0[i], mu_0[i], theta[i], phi[i])
                msap_j = msap_j .+ msap_v .* supp
            end
            if "tail_h" in ac.comp_lst
                msap_h = trailing_edge_horizontal_tail(settings, ac, data.f, M_0[i], c_0[i], rho_0[i], mu_0[i], theta[i], phi[i])
                msap_j = msap_j .+ msap_h .* supp
            end
            if "les" in ac.comp_lst
                msap_les = leading_edge_slat(settings, ac, data.f, M_0[i], c_0[i], rho_0[i], mu_0[i], theta[i], phi[i])
                msap_j = msap_j .+ msap_les .* supp
            end
            if "tef" in ac.comp_lst
                msap_tef = trailing_edge_flap(settings, ac, data.f, M_0[i], c_0[i], theta[i], phi[i], theta_flaps[i])
                msap_j = msap_j .+ msap_tef .* supp
            end
            if "lg" in ac.comp_lst
                msap_lg = landing_gear(settings, ac, data.f, M_0[i], c_0[i], theta[i], phi[i], I_landing_gear[i])
                msap_j = msap_j .+ msap_lg .* supp
            end

            # Normalize msap by reference pressure
            msap_af[i,:] = msap_j/settings.p_ref^2

        else
            msap_af[i,:] =  1e-99 * ones(T, (settings.N_f, ))
        end
    end

    return msap_af
end


using Plots, VortexDistributions

@load "./examples/one_frame.jld2"

heatmap(x,y,abs2.(ψ1'))

psi = Torus(ψ1,x,y)
vortices = findvortices(psi)

# zoom window
x0,y0 = 70,200
Lx,Ly = 100,100

# simple cartesian mask
xm = x[@. abs(x-x0)<Lx/2]
ym = y[@. abs(y-y0)<Ly/2]
psim = @. ψ1[abs(x-x0)<Lx/2,abs(y-y0)<Ly/2]

heatmap(xm,ym,angle.(psim),transpose=true)
psi2 = Torus(psim,xm,ym)
vort = findvortices(psi2)

scatter!([vort[1].xv],[vort[1].yv],alpha=0.4)
scatter!([vort[2].xv],[vort[2].yv],alpha=0.4)


using Base.Test
using Clustering

srand(34568)

m = 3
n = 1000
k = 5

x = rand(m,n)

fuzziness = 2.0
srand(34568)
r = fuzzy_cmeans(x, k, fuzziness)
@test isa(r, FuzzyCMeansResult{Float64})
@test size(r.centers) == (m,k)
@test size(r.weights) == (n,k)
@test all([s ≈ 1.0 for s in sum(r.weights,2)])
@test all(0 .<= r.weights .<= 1)

fuzziness = 3.0
srand(34568)
r = fuzzy_cmeans(x, k, fuzziness)
@test isa(r, FuzzyCMeansResult{Float64})
@test size(r.centers) == (m,k)
@test size(r.weights) == (n,k)
@test all([s ≈ 1.0 for s in sum(r.weights,2)])
@test all(0 .<= r.weights .<= 1)


using TriMeshRender
using Test

Vertex = TriMeshRender.Vertex

@testset "TriMeshRender.jl" begin
	@test isometric(Vertex(0,0,0)) == Vertex(0,0,0)
	@test isometric(Vertex(0,1,0)) == Vertex(0,1,0)
	@test isometric(Vertex(1,0,0)) ≈ Vertex(cos(deg2rad(30)), sin(deg2rad(30)), 0)
	@test isometric(Vertex(2,0,0)) ≈ Vertex(2cos(deg2rad(30)), sin(2deg2rad(30)), 0)
	@test isometric(Vertex(1,1,0)) ≈ Vertex(cos(deg2rad(30)), 1+sin(deg2rad(30)), 0)
	@test isometric(Vertex(0,0,1)) ≈ Vertex(-cos(deg2rad(30)), sin(2eg2rad(30)), 0)
	@test isometric(Vertex(0,0,2)) ≈ Vertex(-2cos(deg2rad(30)), sin(2deg2rad(30)), 0)
end


include("../utils.jl")
include("group_utils.jl")

@testset "Product group" begin
    SOn = SpecialOrthogonal(3)
    Tn = TranslationGroup(2)
    Rn = Rotations(3)
    M = ProductManifold(SOn, Tn)
    G = ProductGroup(M)
    @test_throws ErrorException ProductGroup(ProductManifold(Rotations(3), Stiefel(3, 2)))
    @test G isa ProductGroup
    @test submanifold(G, 1) === SOn
    @test submanifold(G, 2) === Tn
    @test base_manifold(G) === M
    @test sprint(show, G) == "ProductGroup($(SOn), $(Tn))"
    @test sprint(show, "text/plain", G) == "ProductGroup with 2 subgroups:\n $(SOn)\n $(Tn)"
    x = Matrix{Float64}(I, 3, 3)
    t = Vector{Float64}.([1:2, 2:3, 3:4])
    ω = [[1.0, 2.0, 3.0], [3.0, 2.0, 1.0], [1.0, 3.0, 2.0]]
    tuple_pts = [(exp(Rn, x, hat(Rn, x, ωi)), ti) for (ωi, ti) in zip(ω, t)]
    tuple_v = (hat(Rn, x, [1.0, 0.5, -0.5]), [-1.0, 2.0])
    eA = [x, zeros(2)]

    @testset "Product Identity" begin
        i = Identity(G)
        @test submanifold_components(G, i) == (Identity(SOn), Identity(Tn))
        @test submanifold_component(G, i, 1) == Identity(SOn)
        @test submanifold_component(G, i, 2) == Identity(Tn)
    end

    @testset "product repr" begin
        pts = [ProductRepr(tp...) for tp in tuple_pts]
        X_pts = [ProductRepr(tuple_v...)]

        @test compose(G, pts[1], Identity(G)) == pts[1]
        @test compose(G, Identity(G), pts[1]) == pts[1]
        test_group(G, pts, X_pts, X_pts; test_diff=true)
        @test isapprox(
            G,
            Identity(G),
            exp_lie(G, X_pts[1]),
            ProductRepr(exp_lie(SOn, X_pts[1].parts[1]), exp_lie(Tn, X_pts[1].parts[2])),
        )
        @test isapprox(
            G,
            Identity(G),
            log_lie(G, pts[1]),
            ProductRepr(log_lie(SOn, pts[1].parts[1]), log_lie(Tn, pts[1].parts[2])),
        )
        X = log_lie(G, pts[1])
        Z = zero_vector(G, pts[1])
        log_lie!(G, Z, pts[1])
        @test isapprox(G, pts[1], X, Z)
        p = exp_lie(G, X)
        q = identity_element(G)
        @test is_identity(G, q)
        @test isapprox(G, q, Identity(G))
        @test isapprox(G, Identity(G), q)
        exp_lie!(G, q, X)
        @test isapprox(G, p, q)
        log_lie!(G, Z, Identity(G))
        @test isapprox(G, Identity(G), Z, zero_vector(G, identity_element(G)))
        @test isapprox(
            G,
            Identity(G),
            log_lie(G, Identity(G)),
            zero_vector(G, identity_element(G)),
        )

        @test compose(G, pts[1], Identity(G)) == pts[1]
        @test compose(G, Identity(G), pts[1]) == pts[1]
        test_group(G, pts, X_pts, X_pts; test_diff=true, test_mutating=false)
        @test isapprox(
            G,
            exp_lie(G, X_pts[1]),
            ProductRepr(exp_lie(SOn, X_pts[1].parts[1]), exp_lie(Tn, X_pts[1].parts[2])),
        )
        @test isapprox(
            G,
            log_lie(G, pts[1]),
            ProductRepr(log_lie(SOn, pts[1].parts[1]), log_lie(Tn, pts[1].parts[2])),
        )
    end
    @test sprint(show, "text/plain", G) === """
    ProductGroup with 2 subgroups:
     SpecialOrthogonal(3)
     TranslationGroup(2; field = ℝ)"""
end


function block_sizes(::Type{T}, _α, _β, R₁, R₂) where {T}
    W = pick_vector_width(T)
    α = _α * W
    β = _β * W
    L₁ₑ = first_cache_size(T) * R₁
    L₂ₑ = second_cache_size(T) * R₂
    block_sizes(W, α, β, L₁ₑ, L₂ₑ)
end
function block_sizes(W, α, β, L₁ₑ, L₂ₑ)
    mᵣ, nᵣ = matmul_params()
    MᵣW = mᵣ * W

    Mc = floortostaticint(√(L₁ₑ)*√(L₁ₑ*β + L₂ₑ*α)/√(L₂ₑ) / MᵣW) * MᵣW
    Kc = roundtostaticint(√(L₁ₑ)*√(L₂ₑ)/√(L₁ₑ*β + L₂ₑ*α))
    Nc = floortostaticint(√(L₂ₑ)*√(L₁ₑ*β + L₂ₑ*α)/√(L₁ₑ) / nᵣ) * nᵣ

    Mc, Kc, Nc
end
function block_sizes(::Type{T}) where {T}
    block_sizes(T, W₁Default(), W₂Default(), R₁Default(), R₂Default())
end

"""
    split_m(M, Miters_base, W)

Splits `M` into at most `Miters_base` iterations.
For example, if we wish to divide `517` iterations into roughly 7 blocks using multiples of `8`:

```julia
julia> split_m(517, 7, 8)
(72, 2, 69, 7)
```

This suggests we have base block sizes of size `72`, with two iterations requiring an extra remainder of `8 ( = W)`,
and a final block of `69` to handle the remainder. It also tells us that there are `7` total iterations, as requested.
```julia
julia> 80*2 + 72*(7-2-1) + 69
517
```
This is meant to specify roughly the requested amount of blocks, and return relatively even sizes.

This method is used fairly generally.
"""
@inline function split_m(M, _Mblocks, W)
    Miters = cld_fast(M, W)
    Mblocks = min(_Mblocks, Miters)
    Miter_per_block, Mrem = divrem_fast(Miters, Mblocks)
    Mbsize = Miter_per_block * W
    Mremfinal = M - Mbsize*(Mblocks-1) - Mrem * W
    Mbsize, Mrem, Mremfinal, Mblocks
end

"""
  solve_block_sizes(::Type{T}, M, K, N, α, β, R₂, R₃)

This function returns iteration/blocking descriptions `Mc`, `Kc`, and `Nc` for use when packing both `A` and `B`.

It tries to roughly minimize the cost
```julia
MKN/(Kc*W) + α * MKN/Mc + β * MKN/Nc
```
subject to constraints
```julia
Mc - M ≤ 0
Kc - K ≤ 0
Nc - N ≤ 0
Mc*Kc - L₁ₑ ≤ 0
Kc*Nc - L₂ₑ ≤ 0
```
That is, our constraints say that our block sizes shouldn't be bigger than the actual dimensions, and also that
our packed `A` (`Mc × Kc`) should fit into the first packing cache (generally, actually the `L₂`, and our packed
`B` (`Kc × Nc`) should fit into the second packing cache (generally the `L₃`).

Our cost model consists of three components:
1. Cost of moving data in and out of registers. This is done `(M/Mᵣ * K/Kc * N/Nᵣ)` times and the cost per is `(Mᵣ/W * Nᵣ)`.
2. Cost of moving strips from `B` pack from the low cache levels to the highest cache levels when multiplying `Aₚ * Bₚ`.
   This is done `(M / Mc * K / Kc * N / Nc)` times, and the cost per is proportional to `(Kc * Nᵣ)`.
   `α` is the proportionality-constant parameter.
3. Cost of packing `A`. This is done `(M / Mc * K / Kc * N / Nc)` times, and the cost per is proportional to
   `(Mc * Kc)`. `β` is the proportionality-constant parameter.

As `W` is a constant, we multiply the cost by `W` and absorb it into `α` and `β`. We drop it from the description
from  here on out.

In the full problem, we would have Lagrangian, with μ < 0:
f((Mc,Kc,Nc),(μ₁,μ₂,μ₃,μ₄,μ₅))
MKN/Kc + α * MKN/Mc + β * MKN/Nc - μ₁(Mc - M) - μ₂(Kc - K) - μ₃(Nc - N) - μ₄(Mc*Kc - L2) - μ₅(Kc*Nc - L3)
```julia
0 = ∂L/∂Mc = - α * MKN / Mc² - μ₁ - μ₄*Kc
0 = ∂L/∂Kc = - MKN / Kc² - μ₂ - μ₄*Mc - μ₅*Nc
0 = ∂L/∂Nc = - β * MKN / Nc² - μ₃ - μ₅*Kc
0 = ∂L/∂μ₁ = M - Mc
0 = ∂L/∂μ₂ = K - Kc
0 = ∂L/∂μ₃ = N - Nc
0 = ∂L/∂μ₄ = L₁ₑ - Mc*Kc
0 = ∂L/∂μ₅ = L₂ₑ - Kc*Nc
```
The first 3 constraints complicate things, because they're trivially solved by setting `M = Mc`, `K = Kc`, and `N = Nc`.
But this will violate the last two constraints in general; normally we will be on the interior of the inequalities,
meaning we'd be dropping those constraints. Doing so, this leaves us with:

First, lets just solve the cost w/o constraints 1-3
```julia
0 = ∂L/∂Mc = - α * MKN / Mc² - μ₄*Kc
0 = ∂L/∂Kc = - MKN / Kc² - μ₄*Mc - μ₅*Nc
0 = ∂L/∂Nc = - β * MKN / Nc² - μ₅*Kc
0 = ∂L/∂μ₄ = L₁ₑ - Mc*Kc
0 = ∂L/∂μ₅ = L₂ₑ - Kc*Nc
```
Solving:
```julia
Mc = √(L₁ₑ)*√(L₁ₑ*β + L₂ₑ*α)/√(L₂ₑ)
Kc = √(L₁ₑ)*√(L₂ₑ)/√(L₁ₑ*β + L₂ₑ*α)
Nc = √(L₂ₑ)*√(L₁ₑ*β + L₂ₑ*α)/√(L₁ₑ)
μ₄ = -K*√(L₂ₑ)*M*N*α/(L₁ₑ^(3/2)*√(L₁ₑ*β + L₂ₑ*α))
μ₅ = -K*√(L₁ₑ)*M*N*β/(L₂ₑ^(3/2)*√(L₁ₑ*β + L₂ₑ*α))
```
These solutions are indepedent of matrix size.
The approach we'll take here is solving for `Nc`, `Kc`, and then finally `Mc` one after the other, incorporating sizes.

Starting with `N`, we check how many iterations would be implied by `Nc`, and then choose the smallest value that would
yield that number of iterations. This also ensures that `Nc ≤ N`.
```julia
Niter = cld(N, √(L₂ₑ)*√(L₁ₑ*β + L₂ₑ*α)/√(L₁ₑ))
Nblock, Nrem = divrem(N, Niter)
Nblock_Nrem = Nblock + (Nrem > 0)
```
We have `Nrem` iterations of size `Nblock_Nrem`, and `Niter - Nrem` iterations of size `Nblock`.

We can now make `Nc = Nblock_Nrem` a constant, and solve the remaining three equations again:
```julia
0 = ∂L/∂Mc = - α * MKN / Mc² - μ₄*Kc
0 = ∂L/∂Kc = - MKN / Kc² - μ₄*Mc - μ₅*Ncm
0 = ∂L/∂μ₄ = L₂ₑ - Mc*Kc
```
yielding
```julia
Mc = √(L₁ₑ)*√(α)
Kc = √(L₁ₑ)/√(α)
μ₄ = -K*M*N*√(α)/L₁ₑ^(3/2)
```
We proceed in the same fashion as for `Nc`, being sure to reapply the `Kc * Nc ≤ L₂ₑ` constraint:
```julia
Kiter = cld(K, min(√(L₁ₑ)/√(α), L₂ₑ/Nc))
Kblock, Krem = divrem(K, Ki)
Kblock_Krem = Kblock + (Krem > 0)
```
This leaves `Mc` partitioning, for which, for which we use the constraint `Mc * Kc ≤ L₁ₑ` to set
the initial number of proposed iterations as `cld(M, L₁ₑ / Kcm)` for calling `split_m`.
```julia
Mbsize, Mrem, Mremfinal, Mblocks = split_m(M, cld(M, L₁ₑ / Kcm), StaticInt{W}())
```

Note that for synchronization on `B`, all threads must have the same values for `Kc` and `Nc`.
`K` and `N` will be equal between threads, but `M` may differ. By calculating `Kc` and `Nc`
independently of `M`, this algorithm guarantees all threads are on the same page.
"""
@inline function solve_block_sizes(::Type{T}, M, K, N, _α, _β, R₂, R₃, Wfactor) where {T}
    W = pick_vector_width(T)
    α = _α * W
    β = _β * W
    L₁ₑ =  first_cache_size(T) * R₂
    L₂ₑ = second_cache_size(T) * R₃

    # Nc_init = round(Int, √(L₂ₑ)*√(α * L₂ₑ + β * L₁ₑ)/√(L₁ₑ))
    Nc_init⁻¹ = √(L₁ₑ) / (√(L₂ₑ)*√(α * L₂ₑ + β * L₁ₑ))

    Niter = cldapproxi(N, Nc_init⁻¹) # approximate `ceil`
    Nblock, Nrem = divrem_fast(N, Niter)
    Nblock_Nrem = Nblock + One()#(Nrem > 0)

    ((Mblock, Mblock_Mrem, Mremfinal, Mrem, Miter), (Kblock, Kblock_Krem, Krem, Kiter)) = solve_McKc(T, M, K, Nblock_Nrem, _α, _β, R₂, R₃, Wfactor)

    (Mblock, Mblock_Mrem, Mremfinal, Mrem, Miter), (Kblock, Kblock_Krem, Krem, Kiter), promote(Nblock, Nblock_Nrem, Nrem, Niter)
end
# Takes Nc, calcs Mc and Kc
@inline function solve_McKc(::Type{T}, M, K, Nc, _α, _β, R₂, R₃, Wfactor) where {T}
    W = pick_vector_width(T)
    α = _α * W
    β = _β * W
    L₁ₑ =  first_cache_size(T) * R₂
    L₂ₑ = second_cache_size(T) * R₃

    Kc_init⁻¹ = Base.FastMath.max_fast(√(α/L₁ₑ), Nc*inv(L₂ₑ))
    Kiter = cldapproxi(K, Kc_init⁻¹) # approximate `ceil`
    Kblock, Krem = divrem_fast(K, Kiter)
    Kblock_Krem = Kblock + One()

    Miter_init = cldapproxi(M * inv(L₁ₑ), Kblock_Krem) # Miter = M * Kc / L₁ₑ
    Mbsize, Mrem, Mremfinal, Miter = split_m(M, Miter_init, W * Wfactor)
    Mblock_Mrem = Mbsize + W * Wfactor

    promote(Mbsize, Mblock_Mrem, Mremfinal, Mrem, Miter), promote(Kblock, Kblock_Krem, Krem, Kiter)
end

@inline cldapproxi(n, d⁻¹) = Base.fptosi(Int, Base.FastMath.add_fast(Base.FastMath.mul_fast(n, d⁻¹), 0.9999999999999432)) # approximate `ceil`

"""
  find_first_acceptable(M, W)

Finds first combination of `Miter` and `Niter` that doesn't make `M` too small while producing `Miter * Niter = num_cores()`.
This would be awkard if there are computers with prime numbers of cores. I should probably consider that possibility at some point.
"""
@inline function find_first_acceptable(M, W)
    Mᵣ, Nᵣ = matmul_params()
    factors = calc_factors()
    for (miter, niter) ∈ factors
        if miter * (StaticInt(2) * Mᵣ * W) ≤ M + (W + W)
            return miter, niter
        end
    end
    last(factors)
end
"""
  divide_blocks(M, Ntotal, _nspawn, W)

Splits both `M` and `N` into blocks when trying to spawn a large number of threads relative to the size of the matrices.
"""
@inline function divide_blocks(M, Ntotal, _nspawn, W)
    _nspawn == num_cores() && return find_first_acceptable(M, W)
    mᵣ, nᵣ = matmul_params()
    Miter = clamp(div_fast(M, W*mᵣ * MᵣW_mul_factor()), 1, _nspawn)
    nspawn = div_fast(_nspawn, Miter)
    if (nspawn ≤ 1) & (Miter < _nspawn)
        # rebalance Miter
        Miter = cld_fast(_nspawn, cld_fast(_nspawn, Miter))
        nspawn = div_fast(_nspawn, Miter)
    end
    Miter, cld_fast(Ntotal, max(2, cld_fast(Ntotal, nspawn)))
end


function solver_default(X,data::AbstractArray{T,2}) where {T<:Union{Missing, <:Number}}
    # likely much larger matrix, using lsqr
    modelinfo = []
    beta = Array{Float64}(undef,size(data,1),size(X,2))

    @showprogress .1 for ch in 1:size(data,1)
        ix = .!ismissing.(data[ch,:])
        beta[ch,:],h = lsmr(X[ix,:],data[ch,ix],log=true)
        push!(modelinfo,h)
    end

    return beta, modelinfo
end

function solver_default(X,data::AbstractArray{T,3}) where {T<:Union{Missing, <:Number}}
       beta = Array{Union{Missing,Number}}(undef,size(data,1),size(data,2),size(X,2))
       @showprogress .1 for ch in 1:size(data,1)
           for t in 1:size(data,2)
               @debug("$(ndims(data,)),$t,$ch")
               ix = .!ismissing.(data[ch,t,:])
               beta[ch,t,:] = X[ix,:] \ data[ch,t,ix]
           end
       end

       modelinfo = [undef] # no history implemented (yet?)
       return beta, modelinfo
end

function solver_b2b(X,data::AbstractArray{T,3},cross_val_reps = 10) where {T<:Union{Missing, <:Number}}

    X,data = dropMissingEpochs(X,data)

    E = zeros(size(data,2),size(X,2),size(X,2))
    W = Array{Float64}(undef,size(data,2),size(X,2),size(data,1))
    println("n = samples = $(size(X,1)) = $(size(data,3))")
    for t in 1:size(data,2)
        for m in 1:cross_val_reps
            k_ix = collect(Kfold(size(data,3),2))
            Y1 = data[:,t,k_ix[1]]
            Y2 = data[:,t,k_ix[2]]
            X1 = X[k_ix[1],:]
            X2 = X[k_ix[2],:]

            G = (Y1' \ X1)
            H = X2 \ (Y2'*G)

            E[t,:,:] = E[t,:,:]+Diagonal(H[diagind(H)])

        end
        E[t,:,:] = E[t,:,:] ./ cross_val_reps
        W[t,:,:] = (X*E[t,:,:])' / data[:,t,:]

    end

    # extract diagonal
    beta = mapslices(diag,E,dims=[2,3])
    # reshape to conform to ch x time x pred
    beta = permutedims(beta,[3 1 2])
    modelinfo = Dict("W"=>W,"E"=>E,"cross_val_reps"=>cross_val_reps) # no history implemented (yet?)
    return beta, modelinfo
end


# QCQP example
#    maximize x + y
#
#    s.t.  x, y >= 0
#          x^2 + y^2 <= 1
#
#    solution: (0.71, 0.71) objv = 1.414

using Gurobi, Compat.Test

@testset "QCQP1" begin

    env = Gurobi.Env()
    setparam!(env, "OutputFlag", 0)

    model = Gurobi.Model(env, "qcqp_01", :maximize)

    add_cvars!(model, [1., 1.], 0., Inf)
    update_model!(model)

     # add_qpterms!(model, linearindices, linearcoeffs, qrowinds, qcolinds, qcoeffs, sense, rhs)
    add_qconstr!(model, [], [], [1, 2], [1, 2], [1, 1.], '<', 1.0)
    update_model!(model)

    optimize(model)

    @test isapprox(get_solution(model), [0.707107, 0.707107], atol=1e-5)
    @test isapprox(get_objval(model), 1.414214, atol=1e-5)
end


@symbol_func function g_p(cur_reactor::AbstractReactor)
  cur_kappa = cur_reactor.kappa_95

  cur_delta = cur_reactor.delta_95

  cur_factor = 1.0

  cur_factor += 2 * cur_delta ^ 2

  cur_factor -= 1.2 * cur_delta ^ 3

  cur_factor *= cur_kappa ^ 2

  cur_factor += 1

  cur_factor /= 2

  cur_factor
end


using LinearAlgebra: I, rmul!, mul!
import Base.Iterators

# We the more specific OrdinaryDiffEq for a significantly lighter dependency tree and build time.
# However, this means we loose the automated solver picking. For now we use a recommended Tsit5
# solver but it is not clear whether it is the right choice for the stiffness of the
# Schrodinger/Lindblad equations.
# * `Rodas4/5` does not appear to support complex equations.
# * `DP5` used by QuantumOptics.jl also gave strange errors
# * the `reltol` and `abstol` were chosen somewhat arbitrarily to be "good enough"
using OrdinaryDiffEq: ODEProblem, solve, Tsit5

export unitary_propagator, unitary_state, me_propagator, me_state
export floquet_propagator, floquet_rise_fall_propagator, choose_times_floquet, decompose_times_floquet

"""
    unitary_propagator(cqs::CompositeQSystem, ts::AbstractVector{<:Real})

Compute the unitary propagator of a CompositeQSystem at given times by solving the differential
equation `dU/dt = -iHU`.

## args
* `cqs`: a CompositeQSystem.
* `ts`: a sorted array of times at which to compute the unitary propagator.

## returns
An array of unitary propagators at the specified times.
"""
function unitary_propagator(cqs::CompositeQSystem, ts::AbstractVector{<:Real})
    @assert issorted(ts)
    function ode(du, u, p, t)
        ham = p[3] # preallocated workspace array
        ham .= p[2] # start from fixed_ham
        add_parametric_hamiltonians!(ham, p[1], t)
        rmul!(ham, -2π * 1im)
        mul!(du, ham, u)
    end
    fixed_ham = hamiltonian(cqs)
    d = dimension(cqs)
    u0 = Matrix{ComplexF64}(I, d, d) # start with identity
    work_ham = similar(fixed_ham) # scratch space
    prob = ODEProblem(ode, u0, (float(ts[1]), float(ts[end])), (cqs, fixed_ham, work_ham))
    sol = solve(prob, Tsit5(); saveat=ts, save_start=true, reltol=1e-6, abstol=1e-8)
    return sol.u
end

unitary_propagator(cqs::CompositeQSystem, t::Real) = unitary_propagator(cqs, [0.0, t])[end]

"""
    unitary_state(cqs::CompositeQSystem, ts::AbstractVector{<:Real}, ψ0::Vector{<:Number})

Compute the unitary state evolution of a `CompositeQSystem` from a given initial state and at given
times by solving the differential equation `dψ/dt = -iHψ`.

## args
* `cqs`: a CompositeQSystem.
* `ts`: a sorted array of times at which to compute the state.
* `ψ0`: a vector indicating the initial state.

## returns
An array of state vectors for the system at the specified times.
"""
function unitary_state(cqs::CompositeQSystem, ts::AbstractVector{<:Real}, ψ0::Vector{<:Number})
    @assert issorted(ts)
    function ode(dψ, ψ, p, t)
        ham = p[3] # preallocated workspace array
        ham .= p[2] # start from fixed_ham
        add_parametric_hamiltonians!(ham, p[1], t)
        rmul!(ham, -2π * 1im)
        mul!(dψ, ham, ψ)
    end
    fixed_ham = hamiltonian(cqs)
    work_ham = similar(fixed_ham)
    prob = ODEProblem(ode, ψ0, (float(ts[1]), float(ts[end])), (cqs, fixed_ham, work_ham))
    sol = solve(prob, Tsit5(); saveat=ts, save_start=true, reltol=1e-6, abstol=1e-8)
    return sol.u
end

unitary_state(cqs::CompositeQSystem, t::Real, ψ0::Vector{<:Number}) = unitary_state(cqs, [0.0, t], ψ0)[end]

"""
    me_propagator(cqs::CompositeQSystem, ts::AbstractVector{<:Real})

Compute the master equation propagator evolution of a `CompositeQSystem` at given times by solving
the vectorized master equation `du/dt = (-i(I ⊗ H - transpose(H) ⊗ I) + Σ conj(L) ⊗ L - I ⊗ L^†L/2 -
transpose(L^†L/2) ⊗ I)u` which can be derived using the identity `vec(AXB) = (transpose(B) ⊗
A)vec(X)`.

## args
* `cqs`: a CompositeQSystem.
* `ts`: a sorted array of times at which to compute the density matrix.

## returns
An array of propagators for the system at the specified times.
"""
function me_propagator(cqs::CompositeQSystem, ts::AbstractVector{<:Real})
    @assert issorted(ts)
    function ode(du, u, p, t)
        ham = p[3] # preallocated workspace array
        ham .= p[2] # start from fixed_ham
        add_parametric_hamiltonians!(ham, p[1], t)
        d = size(ham, 1)
        I_mat = Matrix{ComplexF64}(I, d, d)
        mul!(du, -2π * 1im * (I_mat ⊗ ham - transpose(ham) ⊗ I_mat), u)
        lind_mat = p[5] # preallocated workspace array
        for (lind_op, idxs) in Iterators.flatten((p[1].fixed_Ls, ((l(t), idxs) for (l, idxs) in p[1].parametric_Ls)))
            lind_mat .= p[4] # start with empty array
            embed_add!(lind_mat, lind_op, idxs)
            du .+= 2π * (conj(lind_mat) ⊗ lind_mat .- 1//2 .* I_mat ⊗ (lind_mat'*lind_mat) .- 1//2 .* transpose(lind_mat'*lind_mat) ⊗ I_mat) * u
        end
    end
    fixed_ham = hamiltonian(cqs)
    work_ham = similar(fixed_ham)
    bare_lind = zeros(ComplexF64, size(fixed_ham))
    work_lind = similar(fixed_ham)
    d = dimension(cqs)^2
    u0 = Matrix{ComplexF64}(I, d, d) # start with identity
    prob = ODEProblem(ode, u0, (float(ts[1]), float(ts[end])), (cqs, fixed_ham, work_ham, bare_lind, work_lind))
    sol = solve(prob, Tsit5(); saveat=ts, save_start=true, reltol=1e-6, abstol=1e-8)
    return sol.u
end

me_propagator(cqs::CompositeQSystem, t::Real) = me_propagator(cqs, [0.0, t])[end]

"""
    me_state(cqs::CompositeQSystem, ts::AbstractVector{<:Real}, ρ0::Matrix{<:Number})

Compute the master equation state evolution of a `CompositeQSystem` from a given initial density
matrix and at given times by solving the differential equation `dρ/dt = -i[H, ρ] + Σ LρL^† - L^†Lρ/2
- ρL^†L/2`.

## args
* `cqs`: a CompositeQSystem.
* `ts`: a sorted array of times at which to compute the density matrix.
* `ρ0`: a matrix indicating the initial density matrix.

## returns
An array of density matrices for the system at the specified times.
"""
function me_state(cqs::CompositeQSystem, ts::AbstractVector{<:Real}, ρ0::Matrix{<:Number})
    @assert issorted(ts)
    function ode(dρ, ρ, p, t)
        ham = p[3] # preallocated workspace array
        ham .= p[2] # start from fixed_ham
        add_parametric_hamiltonians!(ham, p[1], t)
        dρ .= -2π * 1im * (ham*ρ - ρ*ham)
        lind_mat = p[5] # preallocated workspace array
        for (index, (lind_op, idxs)) in enumerate([p[1].fixed_Ls; p[1].parametric_Ls])
            lind_mat .= p[4] # start with empty array
            l = index <= length(p[1].fixed_Ls) ? lind_op : lind_op(t)
            embed_add!(lind_mat, l, idxs)
            dρ .+= 2π * (lind_mat*ρ*lind_mat' .- 1//2 .* lind_mat'*lind_mat*ρ .- 1//2 .* ρ*lind_mat'*lind_mat)
        end
    end
    fixed_ham = hamiltonian(cqs)
    work_ham = similar(fixed_ham)
    bare_lind = zeros(ComplexF64, size(fixed_ham))
    work_lind = similar(fixed_ham)
    prob = ODEProblem(ode, ρ0, (float(ts[1]), float(ts[end])), (cqs, fixed_ham, work_ham, bare_lind, work_lind))
    sol = solve(prob, Tsit5(); saveat=ts, save_start=true, reltol=1e-6, abstol=1e-8)
    return sol.u
end

me_state(cqs::CompositeQSystem, t::Real, ρ0::Matrix{<:Number}) = me_state(cqs, [0.0, t], ρ0)[end]

######################################################
# Define a `propagator_function` to be a function that takes a CompositeQSystem and an AbstractVector
# of times and returns an array of propagators (something that evolves a state vector or density
# matrix) at those times assuming an identity initial condition at the first given time.
######################################################

"""
    floquet_propagator(propagator_func::Function, t_period::Real, rtol::Real)

Given a propagator function and a time period, create a new propagator function that applies
correctly to CompositeQSystems that are periodic with the given time period. It makes use of the
identity `U(nτ + dt) = U(dt)U(τ)^n` where `U` is the propagator for a system with period `τ`.

## args
* `propagator_func`: a propagator function, e.g. `unitary_propagator`.
* `t_period`: the periodicity of the system equation of motion.
* `rtol`: a tolerance to use in `unique_tol` for the times mod the period expressed as a fraction
             of `t_period`.

## returns
A new propagator function.
"""
function floquet_propagator(propagator_func::Function, t_period::Real, rtol::Real)
    function p(cqs::CompositeQSystem, ts::AbstractVector{<:Real})
        @assert issorted(ts)
        quotients, unique_remainders, unique_inds = decompose_times_floquet(ts, t_period, rtol)
        # Perform time evolution for unique remainders as well as one full period. Because
        # unique_remainders[1] == ts[1] and all of unique_remainders .- unique_remainders[1] <=
        # t_period then we can calcuate the propagators for all the "remainder" times starting at
        # ts[1] and also calculate the period propagator by tacking on a final t = t_period + ts[1]
        us_remainders = propagator_func(cqs, [unique_remainders; t_period + ts[1]])
        u_period = us_remainders[end]
        # create desired unitaries
        us = Matrix{ComplexF64}[]
        quotient = quotients[1] # quotients[1] == 0
        u_floquet = u_period^quotient # effectively the identity map
        for i in 1:length(ts)
            # if we've advanced by at least a period update the base u_floquet
            if quotients[i] > quotient
                u_floquet *= u_period^(quotients[i] - quotient)
                quotient = quotients[i]
            end
            push!(us, us_remainders[unique_inds[i]] * u_floquet)
        end
        return us
    end
    return p
end

"""
    floquet_propagator(propagator_func::Function, t_period::Real)

Use a default value of `eps(t_period)` for `rtol`.
"""
floquet_propagator(propagator_func::Function, t_period::Real) = floquet_propagator(propagator_func, t_period, eps(t_period))

"""
    choose_times_floquet(center::Real, width::Real, t_period::Real, dt::Real)

Choose times for integration that maximize remainder overlap for Floquet integration.
`center` is guaranteed to be one of the times.

## args
* `center`: the center of the array of times.
* `width`: the desired approximate width of the array of times.
* `t_period`: the periodicity of the system.
* `dt`: the desired approximate time increment.

## returns
An array of times.
"""
function choose_times_floquet(center::Real, width::Real, t_period::Real, dt::Real)
    @assert width >= 0
    @assert dt >= 0
    if (width == 0) || (dt == 0)
        return [center]
    end

    dt = t_period / ceil(Int, t_period / dt) # reset dt so it divides time_period, making it smaller
    num_times = floor(Int, width / dt)
    num_times += iseven(num_times) # make sure num_times is odd for symmetry
    width = dt * num_times
    times = collect(range(center - width/2, stop=center + width/2, length=num_times))
    times[ceil(Int, num_times/2)] = center
    return times
end

"""
    decompose_times_floquet(ts::AbstractVector{<:Real}, t_period::Real, rtol::Real)

Decompose the given times as `ts = quotients * t_period + unique_remainders[unique_inds]` where
`unique_remainders` is as small of an array as possible and `unique_remainders[1] == ts[1]`.

## args
* `ts`: an array of times.
* `t_period`: the periodicity of the system equations of motion.
* `rtol`: a tolerance to use in `unique_tol` for the times mod the period expressed as a fraction
             of `t_period`.

## returns
`quotients`, `unique_remainders`, and `unique_inds`.

# Examples
```jldoctest
julia> quotients, unique_remainders, unique_inds = decompose_times_floquet(0.25:0.1:1, 0.25)
([0.0, 0.0, 0.0, 1.0, 1.0, 2.0, 2.0, 2.0], [0.25, 0.3, 0.35, 0.4, 0.45], [1, 3, 5, 2, 4, 1, 3, 5])
```
"""
function decompose_times_floquet(ts::AbstractVector{<:Real}, t_period::Real, rtol::Real)
    # find unique times mod a period
    quotients = fld.(ts .- ts[1], t_period)
    remainders = mod.(ts .- ts[1], t_period) .+ ts[1]
    # remove collisions in the times mod a period up to time_tol
    unique_remainders, unique_inds = unique_tol(remainders, rtol * t_period)
    # sort the unique remainders
    sort_inds = sortperm(unique_remainders)
    unique_remainders = unique_remainders[sort_inds]
    unique_inds = sortperm(sort_inds)[unique_inds]
    return quotients, unique_remainders, unique_inds
end

"""
    function decompose_times_floquet(ts::AbstractVector{<:Real}, t_period::Real)

Use a default value off `eps(t_period)` for `rtol`
"""
decompose_times_floquet(ts::AbstractVector{<:Real}, t_period::Real) = decompose_times_floquet(ts, t_period, eps(t_period))

"""
    floquet_rise_fall_propagator(propagator_func::Function, t_period::Real, rise_time::Real, fall_time::Real, rtol::Real)

Create a propagator function that uses the floquet evolution but also allows a rise time and a fall
time during which the pulse is not periodic. The rise time is assumed to be at the beginning of the
vector `ts` and the fall time is assumed to be at the end.

## args
* `propagator_func`: a propagator function, e.g. `unitary_propagator`.
* `t_period`: the periodicity of the system equations of motion.
* `rise_time`: the rise time of the pulse (the time at the beginning that is non-periodic).
* `fall_time`: the fall time of the pulse (the time at the end that is non-periodic).
* `rtol`: a tolerance to use in `unique_tol` for the times mod the period expressed as a fraction of `t_period`.

## returns
A propagator function.
"""
function floquet_rise_fall_propagator(propagator_func::Function, t_period::Real, rise_time::Real,
                                      fall_time::Real, rtol::Real)
    @assert all([t_period, rise_time, fall_time] .>= 0)
    floquet_prop = floquet_propagator(propagator_func, t_period, rtol)
    function p(cqs::CompositeQSystem, ts::Vector{<:Real})
        t0, t1 = ts[1] + rise_time, ts[end] - fall_time
        @assert t0 <= t1
        us_risetime = propagator_func(cqs, [ts[ts .< t0]; t0])
        u0 = us_risetime[end]
        us = us_risetime[1:end-1]
        us_floquet = floquet_prop(cqs, [t0; ts[t0 .<= ts .<= t1]; t1])
        u1 = us_floquet[end] * u0
        append!(us, [u * u0 for u in us_floquet[2:end-1]])
        us_falltime = propagator_func(cqs, [t1; ts[t1 .< ts]])
        append!(us, [u * u1 for u in us_falltime[2:end]])
        return us
    end
    return p
end

"""
    floquet_rise_fall_propagator(propagator_func::Function, t_period::Real, rise_time::Real, fall_time::Real)

Use a default value of `eps(t_period)` for `rtol`.
"""
floquet_rise_fall_propagator(propagator_func::Function, t_period::Real, rise_time::Real, fall_time::Real) =
    floquet_rise_fall_propagator(propagator_func, t_period, rise_time, fall_time, eps(t_period))

"""
    unique_tol(ts::AbstractVector{<:Real}, dt::Real)

Compute an array `a` of values where each is within `dt/2` of some value in `ts` and no two differ
by less than `dt`. Further return a vector of indices `inds` such that `a[inds]` is an approximation
of `ts` up to `dt`.

## args
* `ts`: an array of numbers.
* `dt`: a small number within which errors are unimportant.

## returns
An array of the unique values and an array of indices.
"""
function unique_tol(ts::AbstractVector{<:Real}, dt::Real)
    vals = round.(ts ./ dt) .* dt
    unique_vals = unique(vals)
    unique_inds = [findfirst(unique_vals .== v) for v in vals]
    return unique_vals, unique_inds
end


@doc raw"""
    AbstractAffineConnection

Abstract type for affine connections on a manifold.
"""
abstract type AbstractAffineConnection end

"""
    LeviCivitaConnection

The [Levi-Civita connection](https://en.wikipedia.org/wiki/Levi-Civita_connection) of a Riemannian manifold.
"""
struct LeviCivitaConnection <: AbstractAffineConnection end

struct MetricDecoratorType <: AbstractDecoratorType end

"""
    AbstractConnectionManifold{𝔽,M<:AbstractManifold{𝔽},G<:AbstractAffineConnection} <: AbstractDecoratorManifold{𝔽}

Equip an [`AbstractManifold`](@ref) explicitly with an [`AbstractAffineConnection`](@ref) `G`.

`AbstractConnectionManifold` is defined by values of [`christoffel_symbols_second`](@ref),
which is used for a default implementation of [`exp`](@ref), [`log`](@ref) and
[`vector_transport_to`](@ref). Closed-form formulae for particular connection manifolds may
be explicitly implemented when available.

An overview of basic properties of affine connection manifolds can be found in [^Pennec2020].

[^Pennec2020]:
    > X. Pennec and M. Lorenzi, “5 - Beyond Riemannian geometry: The affine connection
    > setting for transformation groups,” in Riemannian Geometric Statistics in Medical Image
    > Analysis, X. Pennec, S. Sommer, and T. Fletcher, Eds. Academic Press, 2020, pp. 169–229.
    > doi: 10.1016/B978-0-12-814725-2.00012-1.
"""
abstract type AbstractConnectionManifold{𝔽} <:
              AbstractDecoratorManifold{𝔽,MetricDecoratorType} end

"""
    connection(M::AbstractManifold)

Get the connection (an object of a subtype of [`AbstractAffineConnection`](@ref))
of [`AbstractManifold`](@ref) `M`.
"""
connection(::AbstractManifold)

"""
    ConnectionManifold(M, C)

Decorate the [`AbstractManifold`](@ref) `M` with [`AbstractAffineConnection`](@ref) `C`.
"""
struct ConnectionManifold{𝔽,M<:AbstractManifold{𝔽},C<:AbstractAffineConnection} <:
       AbstractConnectionManifold{𝔽}
    manifold::M
    connection::C
end

@doc raw"""
    christoffel_symbols_second(
        M::AbstractManifold,
        p,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
    )

Compute the Christoffel symbols of the second kind in local coordinates of basis `B`.
For affine connection manifold the Christoffel symbols need to be explicitly implemented
while, for a [`MetricManifold`](@ref) they are computed as (in Einstein summation convention)

````math
Γ^{l}_{ij} = g^{kl} Γ_{ijk},
````

where ``Γ_{ijk}`` are the Christoffel symbols of the first kind
(see [`christoffel_symbols_first`](@ref)), and ``g^{kl}`` is the inverse of the local
representation of the metric tensor. The dimensions of the resulting multi-dimensional array
are ordered ``(l,i,j)``.
"""
christoffel_symbols_second(::AbstractManifold, ::Any, ::AbstractBasis)

@decorator_transparent_signature christoffel_symbols_second(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)

@doc raw"""
    christoffel_symbols_second_jacobian(
        M::AbstractManifold,
        p,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
    )

Get partial derivatives of the Christoffel symbols of the second kind
for manifold `M` at `p` with respect to the coordinates of `B`, i.e.

```math
\frac{∂}{∂ p^l} Γ^{k}_{ij} = Γ^{k}_{ij,l}.
```

The dimensions of the resulting multi-dimensional array are ordered ``(i,j,k,l)``.
"""
christoffel_symbols_second_jacobian(::AbstractManifold, ::Any, B::AbstractBasis)
function christoffel_symbols_second_jacobian(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    n = size(p, 1)
    ∂Γ = reshape(
        _jacobian(q -> christoffel_symbols_second(M, q, B; backend=backend), p, backend),
        n,
        n,
        n,
        n,
    )
    return ∂Γ
end
@decorator_transparent_signature christoffel_symbols_second_jacobian(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)

"""
    connection(M::ConnectionManifold)

Return the connection associated with [`ConnectionManifold`](@ref) `M`.
"""
connection(M::ConnectionManifold) = M.connection

Base.copyto!(M::AbstractConnectionManifold, q, p) = copyto!(M.manifold, q, p)
Base.copyto!(M::AbstractConnectionManifold, Y, p, X) = copyto!(M.manifold, Y, p, X)

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

Compute the exponential map on the [`AbstractConnectionManifold`](@ref) `M` equipped with
corresponding affine connection.

If `M` is a [`MetricManifold`](@ref) with a metric that was declared the default metric
using [`is_default_metric`](@ref), this method falls back to `exp(M, p, X)`.

Otherwise it numerically integrates the underlying ODE, see [`solve_exp_ode`](@ref).
Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.
"""
exp(::AbstractConnectionManifold, ::Any...)

@decorator_transparent_fallback function exp!(M::AbstractConnectionManifold, q, p, X)
    tspan = (0.0, 1.0)
    A = get_default_atlas(M)
    i = get_chart_index(M, A, p)
    B = induced_basis(M, A, i, TangentSpace)
    sol = solve_exp_ode(M, p, X, tspan, B; dense=false, saveat=[1.0])
    n = length(p)
    return copyto!(q, sol.u[1][(n + 1):end])
end

"""
    gaussian_curvature(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Gaussian curvature of the manifold `M` at the point `p` using basis `B`.
This is equal to half of the scalar Ricci curvature, see [`ricci_curvature`](@ref).
"""
gaussian_curvature(::AbstractManifold, ::Any, ::AbstractBasis)
function gaussian_curvature(M::AbstractManifold, p, B::AbstractBasis; kwargs...)
    return ricci_curvature(M, p, B; kwargs...) / 2
end
@decorator_transparent_signature gaussian_curvature(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)

function injectivity_radius(M::AbstractConnectionManifold, p)
    return injectivity_radius(base_manifold(M), p)
end
function injectivity_radius(M::AbstractConnectionManifold, m::AbstractRetractionMethod)
    return injectivity_radius(base_manifold(M), m)
end
function injectivity_radius(M::AbstractConnectionManifold, m::ExponentialRetraction)
    return injectivity_radius(base_manifold(M), m)
end
function injectivity_radius(M::AbstractConnectionManifold, p, m::AbstractRetractionMethod)
    return injectivity_radius(base_manifold(M), p, m)
end
function injectivity_radius(M::AbstractConnectionManifold, p, m::ExponentialRetraction)
    return injectivity_radius(base_manifold(M), p, m)
end

"""
    ricci_tensor(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend = diff_backend())

Compute the Ricci tensor, also known as the Ricci curvature tensor,
of the manifold `M` at the point `p` using basis `B`,
see [`https://en.wikipedia.org/wiki/Ricci_curvature#Introduction_and_local_definition`](https://en.wikipedia.org/wiki/Ricci_curvature#Introduction_and_local_definition).
"""
ricci_tensor(::AbstractManifold, ::Any, ::AbstractBasis)
function ricci_tensor(M::AbstractManifold, p, B::AbstractBasis; kwargs...)
    R = riemann_tensor(M, p, B; kwargs...)
    n = size(R, 1)
    Ric = allocate(R, Size(n, n))
    @einsum Ric[i, j] = R[l, i, l, j]
    return Ric
end
@decorator_transparent_signature ricci_tensor(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)

@doc raw"""
    riemann_tensor(M::AbstractManifold, p, B::AbstractBasis; backend::AbstractDiffBackend=diff_backend())

Compute the Riemann tensor ``R^l_{ijk}``, also known as the Riemann curvature
tensor, at the point `p` in local coordinates defined by `B`. The dimensions of the
resulting multi-dimensional array are ordered ``(l,i,j,k)``.

The function uses the coordinate expression involving the second Christoffel symbol,
see [`https://en.wikipedia.org/wiki/Riemann_curvature_tensor#Coordinate_expression`](https://en.wikipedia.org/wiki/Riemann_curvature_tensor#Coordinate_expression)
for details.

# See also

[`christoffel_symbols_second`](@ref), [`christoffel_symbols_second_jacobian`](@ref)
"""
riemann_tensor(::AbstractManifold, ::Any, ::AbstractBasis)
function riemann_tensor(
    M::AbstractManifold,
    p,
    B::AbstractBasis;
    backend::AbstractDiffBackend=diff_backend(),
)
    n = size(p, 1)
    Γ = christoffel_symbols_second(M, p, B; backend=backend)
    ∂Γ = christoffel_symbols_second_jacobian(M, p, B; backend=backend) ./ n
    R = allocate(∂Γ, Size(n, n, n, n))
    @einsum R[l, i, j, k] =
        ∂Γ[l, i, k, j] - ∂Γ[l, i, j, k] + Γ[s, i, k] * Γ[l, s, j] - Γ[s, i, j] * Γ[l, s, k]
    return R
end
@decorator_transparent_signature riemann_tensor(
    M::AbstractDecoratorManifold,
    p,
    B::AbstractBasis;
    kwargs...,
)

@doc raw"""
    solve_exp_ode(
        M::AbstractConnectionManifold,
        p,
        X,
        tspan,
        B::AbstractBasis;
        backend::AbstractDiffBackend = diff_backend(),
        solver = AutoVern9(Rodas5()),
        kwargs...,
    )

Approximate the exponential map on the manifold over the provided timespan
assuming the default connection of the given manifold by solving the ordinary differential
equation

```math
\frac{d^2}{dt^2} p^k + Γ^k_{ij} \frac{d}{dt} p_i \frac{d}{dt} p_j = 0,
```

where ``Γ^k_{ij}`` are the Christoffel symbols of the second kind, and
the Einstein summation convention is assumed. The arguments `tspan` and
`solver` follow the `OrdinaryDiffEq` conventions. `kwargs...` specify keyword
arguments that will be passed to `OrdinaryDiffEq.solve`.

Currently, the numerical integration is only accurate when using a single
coordinate chart that covers the entire manifold. This excludes coordinates
in an embedded space.

!!! note
    This function only works when
    [OrdinaryDiffEq.jl](https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl) is loaded with
    ```julia
    using OrdinaryDiffEq
    ```
"""
function solve_exp_ode(M, p, X, tspan, B::AbstractBasis; kwargs...)
    return error(
        "solve_exp_ode not implemented on $(typeof(M)) for point $(typeof(p)), vector $(typeof(X)), and timespan $(typeof(tspan)). For a suitable default, enter `using OrdinaryDiffEq` on Julia 1.1 or greater.",
    )
end

#
# Introduce transparency
# (a) new functions & other parents
for f in [
    christoffel_symbols_second_jacobian,
    exp,
    gaussian_curvature,
    flat,
    get_coordinates,
    get_vector,
    log,
    mean,
    median,
    project,
    ricci_tensor,
    riemann_tensor,
    sharp,
    vector_transport_along,
    vector_transport_direction,
    vector_transport_direction!, #since it has a default using _to!
    vector_transport_to,
]
    eval(
        quote
            function decorator_transparent_dispatch(
                ::typeof($f),
                ::AbstractConnectionManifold,
                args...,
            )
                return Val(:parent)
            end
        end,
    )
end

# (b) changes / intransparencies.
for f in [
    christoffel_symbols_second, # this is basic for connection manifolds but not for metric manifolds
    exp!,
    flat!,
    get_coordinates!,
    get_vector!,
    get_basis,
    inner,
    inverse_retract!,
    log!,
    mean!,
    median!,
    norm,
    project!,
    sharp!,
    retract!,
    vector_transport_along!,
    vector_transport_to!,
]
    eval(
        quote
            function decorator_transparent_dispatch(
                ::typeof($f),
                ::AbstractConnectionManifold,
                args...,
            )
                return Val(:intransparent)
            end
        end,
    )
end
# (c) special cases
function decorator_transparent_dispatch(
    ::typeof(exp!),
    M::AbstractConnectionManifold,
    q,
    p,
    X,
    t,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(inverse_retract!),
    M::AbstractConnectionManifold,
    X,
    p,
    q,
    m::LogarithmicInverseRetraction,
)
    return Val(:parent)
end
function decorator_transparent_dispatch(
    ::typeof(retract!),
    M::AbstractConnectionManifold,
    q,
    p,
    X,
    m::ExponentialRetraction,
)
    return Val(:parent)
end


module Variance
using ValidatedNumerics;
using Distributions;
using Dierckx;

abstract BaseEstimator

function estimate(values, method::BaseEstimator)
    error("Unsupported method $method")
end

immutable byConfidenceProbability <: BaseEstimator
    confidence_probability :: Real
end

function estimate(values, method :: byConfidenceProbability)
    average = mean(values)
    count = length(values)
    gamma_down, gamma_up = getGammas(method.confidence_probability)
    quantile_down = getChiSquareQuantile(gamma_down, count)
    quantile_up = getChiSquareQuantile(gamma_up, count)

    reduced = mapreduce((x) -> (x - average)^2, +, values)

    @interval(reduced/quantile_down, reduced/quantile_up)
end

immutable byMeanAbsoluteDeviation <: BaseEstimator
    confidence_probability :: Real
end

function estimate(values, method :: byMeanAbsoluteDeviation)
    average = mean(values)
    count = length(values)
    quantile_down, quantile_up = getMeanAbsDeviationQuantiles(method.confidence_probability, count)

    mean_abs_deviation = mapreduce((x) -> abs(x - average), +, values)/count

    @interval(mean_abs_deviation/quantile_down, mean_abs_deviation/quantile_up)^2
end

immutable byPointVariance <: BaseEstimator
    confidence_probability :: Real
    variance :: Real
end

function estimate(values, method::byPointVariance)
    count = length(values)
    gamma_down, gamma_up = getGammas(method.confidence_probability)
    quantile_down = getChiSquareQuantile(gamma_down, count)
    quantile_up = getChiSquareQuantile(gamma_up, count)

    fixed_variance = method.variance * (1 + 0.254 / (count - 1))^2
    term = (count - 1) * fixed_variance

    @interval(term / quantile_down, term / quantile_up)
end

function getGammas(confidence_probability)
    (1 + confidence_probability) * 0.5, (1 - confidence_probability) * 0.5
end

function getChiSquareQuantile(value, count)
    d = Chisq(count - 1)
    quantile(d, value)
end

function getMeanAbsDeviationQuantiles(confidence_probability, count)
    if confidence_probability == 0.90
        const knots = [
        2 1.386 0.044;
        3 1.276 0.166;
        4 1.224 0.254;
        5 1.187 0.305;
        6 1.158 0.360;
        7 1.135 0.394;
        8 1.116 0.422;
        9 1.100 0.445;
        10 1.086 0.464;
        ]
    elseif confidence_probability == 0.95
        const knots = [
        2 1.585 0.022;
        3 1.417 0.116;
        4 1.344 0.199;
        5 1.292 0.260;
        6 1.253 0.306;
        7 1.222 0.342;
        8 1.196 0.372;
        9 1.175 0.396;
        10 1.156 0.417;
        ]
    elseif confidence_probability == 0.99
        const knots = [
        2 1.985 0.004;
        3 1.703 0.073;
        4 1.590 0.145;
        5 1.507 0.203;
        6 1.445 0.250;
        7 1.397 0.287;
        8 1.358 0.318;
        9 1.326 0.344;
        10 1.299 0.366;
        ]
    else
        error("Value=$confidence_probability is not allowed")
    end
    x_knots = knots[:, 1]
    if count > last(x_knots)
        error("Error $count too high")
    end

    y_down_knots = knots[:, 2]
    down_itp = Spline1D(x_knots, y_down_knots, k=1, bc="extrapolate")

    y_up_knots = knots[:, 3]
    up_itp = Spline1D(x_knots, y_up_knots, k=1, bc="extrapolate")
    evaluate(down_itp, count), evaluate(up_itp, count)
end

end


ChainRulesCore.@scalar_rule(
    betalogpdf(α::Real, β::Real, x::Number),
    @setup(z = digamma(α + β)),
    (
        log(x) + z - digamma(α),
        log1p(-x) + z - digamma(β),
        (α - 1) / x + (1 - β) / (1 - x),
    ),
)

ChainRulesCore.@scalar_rule(
    binomlogpdf(n::Real, p::Real, k::Real),
    @setup(z = digamma(n - k + 1)),
    (
        ChainRulesCore.NoTangent(),
        (k / p - n) / (1 - p),
        ChainRulesCore.NoTangent(),
    ),
)

ChainRulesCore.@scalar_rule(
    chisqlogpdf(k::Real, x::Number),
    @setup(hk = k / 2),
    (
        (log(x) - logtwo - digamma(hk)) / 2,
        (hk - 1) / x - one(hk) / 2,
    ),
)

ChainRulesCore.@scalar_rule(
    fdistlogpdf(ν1::Real, ν2::Real, x::Number),
    @setup(
        xν1 = x * ν1,
        temp1 = xν1 + ν2,
        a = (x - 1) / temp1,
        νsum = ν1 + ν2,
        di = digamma(νsum / 2),
    ),
    (
        (-log1p(ν2 / xν1) - ν2 * a + di - digamma(ν1 / 2)) / 2,
        (-log1p(xν1 / ν2) + ν1 * a + di - digamma(ν2 / 2)) / 2,
        ((ν1 - 2) / x - ν1 * νsum / temp1) / 2,
    ),
)

ChainRulesCore.@scalar_rule(
    gammalogpdf(k::Real, θ::Real, x::Number),
    @setup(
        invθ = inv(θ),
        xoθ = invθ * x,
        z = xoθ - k,
    ),
    (
        log(xoθ) - digamma(k),
        invθ * z,
        - (1 + z) / x,
    ),
)

ChainRulesCore.@scalar_rule(
    poislogpdf(λ::Number, x::Number),
    ((iszero(x) && iszero(λ) ? zero(x / λ) : x / λ) - 1, ChainRulesCore.NoTangent()),
)

ChainRulesCore.@scalar_rule(
    tdistlogpdf(ν::Real, x::Number),
    @setup(
        νp1 = ν + 1,
        xsq = x^2,
        invν = inv(ν),
        a = xsq * invν,
        b = νp1 / (ν + xsq),
    ),
    (
        (digamma(νp1 / 2) - digamma(ν / 2) + a * b - log1p(a) - invν) / 2,
        - x * b,
    ),
)