using Test
using LinearAlgebra
using Nystrom
using StaticArrays
using Random
Random.seed!(1)

@testset "DIM" begin
    # construct interior solution
    rtol=5e-2
    @testset "Greens identity (interior) 2d" begin
        Geometry.clear_entities!()
        Ω   = ParametricSurfaces.Circle() |> Geometry.Domain
        Γ   = boundary(Ω)
        M   = ParametricSurfaces.meshgen(Γ,(5,))
        mesh = NystromMesh(view(M,Γ),order=3)
        xout = SVector(3,3)
        ops = (
            Laplace(;dim=2),
            Helmholtz(;dim=2,k=1.2),
            Elastostatic(;dim=2,μ=2,λ=3)
        )
        for pde in ops
            T    = Nystrom.default_density_eltype(pde)
            c    = rand(T)
            u    = (qnode) -> SingleLayerKernel(pde)(qnode,xout)*c
            dudn = (qnode) -> AdjointDoubleLayerKernel(pde)(qnode,xout)*c
            γ₀u   = Density(u,mesh)
            γ₁u   = Density(dudn,mesh)
            γ₀u_norm = norm(norm.(γ₀u,Inf),Inf)
            γ₁u_norm = norm(norm.(γ₁u,Inf),Inf)
            # single and double layer
            S     = SingleLayerOperator(pde,mesh)
            Smat  = S |> Matrix
            D     = DoubleLayerOperator(pde,mesh)
            Dmat = D |> Matrix
            e0   = Nystrom.error_interior_green_identity(Smat,Dmat,γ₀u,γ₁u)/γ₀u_norm
            Sdim  = Nystrom.assemble_dim(S)
            Ddim  = Nystrom.assemble_dim(D)
            e1    = Nystrom.error_interior_green_identity(Sdim,Ddim,γ₀u,γ₁u)/γ₀u_norm
            @testset "Single/double layer $(string(pde))" begin
                @test norm(e0,Inf) > norm(e1,Inf)
                @test norm(e1,Inf) < rtol
            end
            # adjoint double-layer and hypersingular
            K     = AdjointDoubleLayerOperator(pde,mesh)
            Kmat     = K |> Matrix
            H     = HyperSingularOperator(pde,mesh)
            Hmat = H |> Matrix
            e0    = Nystrom.error_interior_derivative_green_identity(Kmat,Hmat,γ₀u,γ₁u)/γ₁u_norm
            Kdim  = Nystrom.assemble_dim(K)
            Hdim  = Nystrom.assemble_dim(H)
            e1   = Nystrom.error_interior_derivative_green_identity(Kdim,Hdim,γ₀u,γ₁u)/γ₁u_norm
            @testset "Adjoint double-layer/hypersingular $(string(pde))" begin
                @test norm(e0,Inf) > rtol
                @test norm(e1,Inf) < rtol
            end
        end
    end

    @testset "Greens identity (interior) 3d" begin
        Geometry.clear_entities!()
        Ω   = ParametricSurfaces.Sphere(;radius=1) |> Geometry.Domain
        Γ   = boundary(Ω)
        M   = ParametricSurfaces.meshgen(Γ,(2,2))
        mesh = NystromMesh(view(M,Γ),order=3)
        xout = SVector(3,3,3)
        ops = (
            Laplace(;dim=3),
            Helmholtz(;dim=3,k=1.2),
            Elastostatic(;dim=3,μ=2,λ=3),
            Maxwell(;k=2.)
        )
        for pde in ops
            T    = Nystrom.default_density_eltype(pde)
            c    = rand(T)
            u    = (qnode) -> SingleLayerKernel(pde)(xout,qnode)*c
            dudn = (qnode) -> transpose(DoubleLayerKernel(pde)(xout,qnode))*c
            γ₀u   = Density(u,mesh)
            γ₁u   = Density(dudn,mesh)
            γ₀u_norm = norm(norm.(γ₀u,Inf),Inf)
            γ₁u_norm = norm(norm.(γ₁u,Inf),Inf)
            # single and double layer
            S     = SingleLayerOperator(pde,mesh)
            Smat  = S |> Matrix
            D     = DoubleLayerOperator(pde,mesh)
            Dmat = D |> Matrix
            e0   = Nystrom.error_interior_green_identity(Smat,Dmat,γ₀u,γ₁u)/γ₀u_norm
            Sdim  = Nystrom.assemble_dim(S)
            Ddim  = Nystrom.assemble_dim(D)
            e1    = Nystrom.error_interior_green_identity(Sdim,Ddim,γ₀u,γ₁u)/γ₀u_norm
            @testset "Single/double layer $(string(pde))" begin
                @test norm(e0,Inf) > norm(e1,Inf)
                @test norm(e1,Inf) < rtol
            end
            # adjoint double-layer and hypersingular
            pde isa Maxwell && continue
            K     = AdjointDoubleLayerOperator(pde,mesh)
            Kmat     = K |> Matrix
            H     = HyperSingularOperator(pde,mesh)
            Hmat = H |> Matrix
            e0   = Nystrom.error_interior_derivative_green_identity(Kmat,Hmat,γ₀u,γ₁u)/γ₁u_norm
            Kdim  = Nystrom.assemble_dim(K)
            Hdim  = Nystrom.assemble_dim(H)
            e1   = Nystrom.error_interior_derivative_green_identity(Kdim,Hdim,γ₀u,γ₁u)/γ₁u_norm
            @testset "Adjoint double-layer/hypersingular $(string(pde))" begin
                @test norm(e0,Inf) > norm(e1,Inf)
                @test norm(e1,Inf) < rtol
            end
        end
    end

    @testset "Greens identity (exterior) 2d" begin
        Geometry.clear_entities!()
        Ω   = ParametricSurfaces.Circle() |> Geometry.Domain
        Γ   = boundary(Ω)
        M   = ParametricSurfaces.meshgen(Γ,(7,))
        mesh = NystromMesh(view(M,Γ),order=3)
        xin = SVector(0.1,0.2)
        ops = (
            Laplace(;dim=2),
            Helmholtz(;dim=2,k=1.2),
            Elastostatic(;dim=2,μ=2,λ=3)
        )
        for pde in ops
            T    = Nystrom.default_density_eltype(pde)
            c    = rand(T)
            u    = (qnode) -> SingleLayerKernel(pde)(xin,qnode)*c
            dudn = (qnode) -> transpose(DoubleLayerKernel(pde)(xin,qnode))*c
            γ₀u   = Density(u,mesh)
            γ₁u   = Density(dudn,mesh)
            γ₀u_norm = norm(norm.(γ₀u,Inf),Inf)
            γ₁u_norm = norm(norm.(γ₁u,Inf),Inf)
            # single and double layer
            S     = SingleLayerOperator(pde,mesh)
            Smat  = S |> Matrix
            D     = DoubleLayerOperator(pde,mesh)
            Dmat = D |> Matrix
            e0   = Nystrom.error_exterior_green_identity(Smat,Dmat,γ₀u,γ₁u)/γ₀u_norm
            Sdim  = Nystrom.assemble_dim(S)
            Ddim  = Nystrom.assemble_dim(D)
            e1    = Nystrom.error_exterior_green_identity(Sdim,Ddim,γ₀u,γ₁u)/γ₀u_norm
            @testset "Single/double layer $(string(pde))" begin
                @test norm(e0,Inf) > norm(e1,Inf)
                @test norm(e1,Inf) < rtol
            end
            # adjoint double-layer and hypersingular
            K     = AdjointDoubleLayerOperator(pde,mesh)
            Kmat     = K |> Matrix
            H     = HyperSingularOperator(pde,mesh)
            Hmat = H |> Matrix
            e0   = Nystrom.error_exterior_derivative_green_identity(Kmat,Hmat,γ₀u,γ₁u)/γ₁u_norm
            Kdim  = Nystrom.assemble_dim(K)
            Hdim  = Nystrom.assemble_dim(H)
            e1   = Nystrom.error_exterior_derivative_green_identity(Kdim,Hdim,γ₀u,γ₁u)/γ₁u_norm
            @testset "Adjoint double-layer/hypersingular $(string(pde))" begin
                @test norm(e0,Inf) > rtol
                @test norm(e1,Inf) < rtol
            end
        end
    end

    @testset "Greens identity (exterior) 3d" begin
        Geometry.clear_entities!()
        Ω   = ParametricSurfaces.Sphere(;radius=3) |> Geometry.Domain
        Γ   = boundary(Ω)
        M   = ParametricSurfaces.meshgen(Γ,(4,4))
        mesh = NystromMesh(view(M,Γ),order=3)
        xs = SVector(0.1,-0.1,0.2)
        ops = (
            Laplace(;dim=3),
            Helmholtz(;dim=3,k=1.2),
            Elastostatic(;dim=3,μ=2,λ=3),
            Maxwell(;k=1)
        )
        for pde in ops
            T    = Nystrom.default_density_eltype(pde)
            c    = rand(T)
            u    = (qnode) -> SingleLayerKernel(pde)(xs,qnode)*c
            dudn = (qnode) -> transpose(DoubleLayerKernel(pde)(xs,qnode))*c
            γ₀u   = Density(u,mesh)
            γ₁u   = Density(dudn,mesh)
            γ₀u_norm = norm(norm.(γ₀u,Inf),Inf)
            γ₁u_norm = norm(norm.(γ₁u,Inf),Inf)
            # single and double layer
            S     = SingleLayerOperator(pde,mesh)
            Smat  = S |> Matrix
            D     = DoubleLayerOperator(pde,mesh)
            Dmat = D |> Matrix
            e0   = Nystrom.error_exterior_green_identity(Smat,Dmat,γ₀u,γ₁u)/γ₀u_norm
            Sdim  = Nystrom.assemble_dim(S)
            Ddim  = Nystrom.assemble_dim(D)
            e1    = Nystrom.error_exterior_green_identity(Sdim,Ddim,γ₀u,γ₁u)/γ₀u_norm
            @testset "Single/double layer $(string(pde))" begin
                @test norm(e0,Inf) > norm(e1,Inf)
                @test norm(e1,Inf) < rtol
            end
            # adjoint double-layer and hypersingular
            pde isa Maxwell && continue
            K     = AdjointDoubleLayerOperator(pde,mesh)
            Kmat     = K |> Matrix
            H     = HyperSingularOperator(pde,mesh)
            Hmat = H |> Matrix
            e0   = Nystrom.error_exterior_derivative_green_identity(Kmat,Hmat,γ₀u,γ₁u)/γ₁u_norm
            Kdim  = Nystrom.assemble_dim(K)
            Hdim  = Nystrom.assemble_dim(H)
            e1   = Nystrom.error_exterior_derivative_green_identity(Kdim,Hdim,γ₀u,γ₁u)/γ₁u_norm
            @testset "Adjoint double-layer/hypersingular $(string(pde))" begin
                @test norm(e0,Inf) > norm(e1,Inf)
                @test norm(e1,Inf) < rtol
            end
        end
    end
end


include("constants.jl")

function angulares(λ, D; arcs = true)
    if arcs == true
        θ = 2.5e5 * (λ / D) #Arcseconds
    else
        θ = 1.22 * (λ / D) #Radians
    end

    return θ
end

function rms(λ)
    return 0.03 * λ
end

export angulares, rms


# using Revise, Plots
using Test
using BifurcationKit, SparseArrays, Parameters, LinearAlgebra, Setfield, ForwardDiff
const BK = BifurcationKit
const FD = ForwardDiff

norminf = x -> norm(x, Inf)
function Laplacian1D(N, lx, bc = :none)
	hx = 2lx/N
	Δ = spdiagm(0 => -2ones(N), 1 => ones(N-1), -1 => ones(N-1) )
	Δ[1,end]=1; Δ[end,1]=1
	D = spdiagm(1 => ones(N-1), -1 => -ones(N-1) )
	D[1,end]=-1; D[end,1]=1
	D = D / (2hx)
	Δ = Δ / hx^2
	return Δ, D
end

# add to f the nonlinearity
@views function NL!(f, u, p)
	@unpack r, μ, ν, c3, c5, γ = p
	n = div(length(u), 2)
	u1 = u[1:n]
	u2 = u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1 = f[1:n]
	f2 = f[n+1:2n]

	f1 .+= @. r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1 + γ
	f2 .+= @. r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2

	return f
end

function Fcgl!(f, u, p, t = 0)
	mul!(f, p.Δ, u)
	NL!(f, u, p)
end

Fcgl(u, p, t = 0) = Fcgl!(similar(u), u, p)

# remark: I checked this against finite differences
@views function Jcgl(u, p)
	@unpack r, μ, ν, c3, c5, Δ = p

	n = div(length(u), 2)
	u1 = u[1:n]
	u2 = u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1u = zero(u1)
	f2u = zero(u1)
	f1v = zero(u1)
	f2v = zero(u1)

	@. f1u =  r - 2 * u1 * (c3 * u1 - μ * u2) - c3 * ua - 4 * c5 * ua * u1^2 - c5 * ua^2
	@. f1v = -ν - 2 * u2 * (c3 * u1 - μ * u2)  + μ * ua - 4 * c5 * ua * u1 * u2
	@. f2u =  ν - 2 * u1 * (c3 * u2 + μ * u1)  - μ * ua - 4 * c5 * ua * u1 * u2
	@. f2v =  r - 2 * u2 * (c3 * u2 + μ * u1) - c3 * ua - 4 * c5 * ua * u2 ^2 - c5 * ua^2

	jacdiag = vcat(f1u, f2v)

	Δ + spdiagm(0 => jacdiag, n => f1v, -n => f2u)
end

jet = BK.getJet(Fcgl, Jcgl)
####################################################################################################
n = 50
	l = pi

	Δ, D = Laplacian1D(n, l, :Periodic)
	par_cgl = (r = 0.0, μ = 0.5, ν = 1.0, c3 = -1.0, c5 = 1.0, Δ = blockdiag(Δ, Δ), Db = blockdiag(D, D), γ = 0.0, δ = 1.0, N = 2n)
	sol0 = zeros(par_cgl.N)

_sol0 = zeros(2n)
_J0 = Jcgl(_sol0, par_cgl)
_J1 = FD.jacobian(z->Fcgl(z, par_cgl), _sol0) |> sparse
@test _J0 ≈ _J1

eigls = EigArpack(1.0, :LM)
	eigls = DefaultEig()
	# eigls = eig_MF_KrylovKit(tol = 1e-8, dim = 60, x₀ = rand(ComplexF64, Nx*Ny), verbose = 1)
	opt_newton = NewtonPar(tol = 1e-9, verbose = true, eigsolver = eigls, maxIter = 20)
	out, hist, flag = @time newton(jet[1], jet[2], sol0, par_cgl, opt_newton, normN = norminf)

opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.15, ds = 0.001, pMax = 2.5, detectBifurcation = 3, nev = 9, plotEveryStep = 50, newtonOptions = (@set opt_newton.verbose = false), maxSteps = 1060, nInversion = 8, maxBisectionSteps=20)
br, = continuation(jet[1], jet[2], vec(sol0), par_cgl, (@lens _.r), opts_br, verbosity = 0)

####################################################################################################
# we test the jacobian
_J0 = jet[2](sol0, par_cgl)
_J1 = FD.jacobian(z->jet[1](z, par_cgl), sol0) |> sparse
@test _J0 == _J1
####################################################################################################
function guessFromHopfO2(branch, ind_hopf, eigsolver, M, z1, z2 = 0.; phase = 0, k = 1.)
	specialpoint = branch.specialpoint[ind_hopf]
	@show specialpoint.ind_ev

	# parameter value at the Hopf point
	p_hopf = specialpoint.param

	# frequency at the Hopf point
	# @show branch.eig[specialpoint.idx].eigenvals[specialpoint.ind_ev:specialpoint.ind_ev+4]
	ωH = imag(branch.eig[specialpoint.idx].eigenvals[specialpoint.ind_ev]) |> abs

	# vec_hopf is the eigenvector for the eigenvalues iω
	vec_hopf1 = geteigenvector(eigsolver, br.eig[specialpoint.idx][2], specialpoint.ind_ev)
	vec_hopf1 ./=  norm(vec_hopf1)

	vec_hopf2 = geteigenvector(eigsolver, br.eig[specialpoint.idx][2], specialpoint.ind_ev - 2)
	vec_hopf2 ./=  norm(vec_hopf2)

	 orbitguess = [real.(specialpoint.x .+
	 			z1 .* vec_hopf1 .* exp(2pi * complex(0, 1) .* (ii/(M-1) - phase)) .+
				z2 .* vec_hopf2 .* exp(2pi * complex(0, 1) .* (ii/(M-1) - phase))) for ii=0:M-1]

	 X = LinRange(-pi, pi, n) |> collect; X = vcat(X, X)
	 # orbitguess = [real.(specialpoint.x .+ z1 .* exp.(complex(0, 1) .* (2pi * ii/(M-1) .- k .* X))) for ii=0:M-1]

	return p_hopf, 2pi/ωH, orbitguess, specialpoint.x, vec_hopf1, vec_hopf2
end
####################################################################################################
# we test TWProblem: travelling wave problem
# number of time slices in the periodic orbit
M = 50

# TW ansatz
r_hopf, Th, orbitguess2, hopfpt, eigvec = guessFromHopfO2(br, 2, opt_newton.eigsolver, M, 1. + 0.0im, 1+0.0im; k = 2.) #TW

uold = copy(orbitguess2[1][1:2n])
# plot(uold[1:end-1]; linewidth = 5)

# we create a TW problem
probTW = BK.TWProblem(Fcgl, Jcgl, par_cgl.Db, copy(uold))
probTW(vcat(uold,.1), par_cgl)

# we test the sparse formulation of the problem jacobian
_sol0 = rand(2n+1)
_J1 = FD.jacobian(z->probTW(z, par_cgl), _sol0) |> sparse
_J0 = probTW(Val(:JacFullSparse), _sol0, par_cgl)
@test _J1 ≈ _J0

# we test the matrix-free formulation of the problem jacobian
_sol0 = rand(2n+1)
_dsol0 = rand(2n+1)
_out1 = FD.derivative(t -> probTW(_sol0 .+ t .* _dsol0, par_cgl), 0)
_out0 = probTW(_sol0, par_cgl, _dsol0)
@test _out0 ≈ _out1

BK.VFtw(probTW, uold, (user=par_cgl, s=Tuple(0.,)))

# we test the ∂
BK.applyD(probTW, rand(2n))

# we test update section
BK.updateSection!(probTW, probTW.u₀)
####################################################################################################
# test newton method, not meant to converge
newton(probTW, vcat(uold,.1), par_cgl, NewtonPar(verbose = true, maxIter = 5), jacobian = :AutoDiff)
newton(probTW, vcat(uold,.1), par_cgl, NewtonPar(verbose = true, maxIter = 5), jacobian = :FullLU)
####################################################################################################
# test continuation method witth different Generalised eigensolvers
optn = NewtonPar(tol = 1e-8)
opt_cont_br = ContinuationPar(pMin = -1., pMax = 1., newtonOptions = optn, maxSteps = 3, detectBifurcation = 2)
continuation(probTW, vcat(uold,.1), par_cgl, (@lens _.r), opt_cont_br; jacobian = :FullLU, verbosity = 0)

@set! opt_cont_br.newtonOptions.eigsolver = BK.DefaultGEig(B = diagm(0=>vcat(ones(2n),0)))
continuation(probTW, vcat(uold,.1), par_cgl, (@lens _.r), opt_cont_br; jacobian = :FullLU, verbosity = 0)

BK.GEigArpack(nothing, :LR)
@set! opt_cont_br.newtonOptions.eigsolver = EigArpack(nev = 5, which = :LM, sigma = 0.2, v0 = rand(2n+1))
continuation(probTW, vcat(uold,.1), par_cgl, (@lens _.r), opt_cont_br; jacobian = :AutoDiff, verbosity = 0)

@set! opt_cont_br.newtonOptions.linsolver = GMRESIterativeSolvers(N = 2n+1)
@set! opt_cont_br.newtonOptions.eigsolver = EigArpack(nev = 4, ncv = 2n+1, tol = 1e-3, v0 = rand(2n+1))
@set! opt_cont_br.detectBifurcation = 0
continuation(probTW, vcat(uold,.1), par_cgl, (@lens _.r), opt_cont_br; jacobian = :MatrixFreeAD, verbosity = 0)


# News Boy Problem solved by L-Shaped Benders Decomposition
# Author: Edward J. Xu, edxu96@outlook.com
# Date: April 5th, 2019

function get_data()
    demand = [12, 14, 16, 18, 20, 22, 24, 26, 28, 30] # Demand of newspapers in each scenario
    lengthS = length(demand)
    vec_prob = [0.05, 0.10, 0.10, 0.10, 0.15, 0.15, 0.10, 0.10, 0.10, 0.05 ] # probability of scenario
    c = 20 # purchase price
    p = 70 # selling price
    h = 10 # scrap value
    y = 20
    # 2. Data Transfer -----------------------------------------------------------------------------------------------------
    n_x = 10  # Number of all the x
    n_y = 1
    vec_min_y = hcat([0])
    vec_max_y = hcat([30])
    vec_f = hcat(c - 10)
    vec_pi = hcat(vec_prob)
    mat_c = zeros(10, 1, 1)
    for i = 1:10
        mat_c[i, :, :] = hcat([- 70 + 10])
    end
    mat_h = zeros(10, 2, 1)
    for i = 1: 10
        mat_h[i, 1, :] = hcat(- demand[i])
        mat_h[i, 2, :] = hcat(0)
    end
    mat3_t = zeros(10, 2, 1)
    for i = 1:10
        mat3_t[i, 1, :] = hcat(0)
        mat3_t[i, 2, :] = hcat(1)
    end
    mat3_w = - ones(10, 2, 1)
    epsilon = 0.001
    timesIterationMax = 500
    return
end

function main()

    BendersLshaped(; n_x = n_x, n_y = n_y, vec_min_y = vec_min_y, vec_max_y = vec_max_y, vec_f = vec_f,
        vec_pi = vec_pi, mat_c = mat_c, mat_h = mat_h, mat3_t = mat3_t, mat3_w = mat3_w,
        epsilon = epsilon, timesIterationMax = timesIterationMax)
    # answer: obj = - 976
end


"""
assembler is the functions in charge of the construction of the matrix and RHS of the main FEM system
"""
function assembler( nodalCoords, elemNodalConnec, MEBIValsMat, MEBIVec, materialsData, elementsData, BCsData, currentSolution )

    nElems = size( elemNodalConnec, 1)
    nNodes = size( nodalCoords,     1)

    # The global stiffness equations are computed for the unknown displacements (free dofs)
    KG = zeros( 6*nNodes, 6*nNodes )

    Kmatrices = [] ;

    for elem in (1:nElems);

        # extract the MEBI entry
        elemMEBIEntry = MEBIVec[elem]

        elemMat = MEBIValsMat[ elemMEBIEntry, 1 ]
        elemEle = MEBIValsMat[ elemMEBIEntry, 2 ]

        print( "elemMebi entry", elemMEBIEntry )

        if elemMat > 0  # if it has a material

            nodesOfThisElem = elemNodalConnec[ elem ] ;
            dofsElem = nodes2dofs( nodesOfThisElem , 6 ) ;

            if cmp( elementsData[elemEle].type, "truss" ) == 0

                dofsElem = dofsElem[ [1,5,7,11]]

                E = materialsData[ MEBIValsMat[ elemMEBIEntry, 1 ] ].youngModulus
                A = elementsData[  MEBIValsMat[ elemMEBIEntry, 2 ] ].crossSectionArea

                nodesOfElemCoords = nodalCoords[ nodesOfThisElem, : ] ;
                L = norm( nodalCoords[ nodesOfThisElem[1], : ] - nodalCoords[ nodesOfThisElem[2], : ] );

                KGelem = linearTrussStiffnessMatrix( E, A, nodesOfElemCoords );

            elseif cmp( elementsData[elemEle].type, "tetrahedron" ) == 0

                dofsElem = dofsElem[ Vector(1:2:6*4) ]

                E = materialsData[ MEBIValsMat[ elemMEBIEntry, 1 ] ].youngModulus

                KGelem = linearTetraSolid( E );

            end

            KG[ dofsElem, dofsElem ] = KG[ dofsElem, dofsElem ] .+ KGelem

            push!( Kmatrices, KGelem )
        end
    end

    FG, neumDofs = assembleFext( nodalCoords, elemNodalConnec, MEBIValsMat, MEBIVec, BCsData )

    return KG, FG, Kmatrices, neumDofs
end

function assembleFext( nodalCoords, elemNodalConnec, MEBIValsMat, MEBIVec, BCsData )

    nNodes     = size( nodalCoords,     1)

    neumDofs = computeNeumDofs( nNodes )

    print("neumdofs", neumDofs)
    FG = zeros( 6*nNodes ) ;
    elementsBC = computeElemsByMEBI( MEBIValsMat[:,3], MEBIVec )
    for indBC in (1:length(BCsData))
        print("indBC", indBC)
        for elem in elementsBC[indBC]
            print("elem:", elem)
            nodesElem = elemNodalConnec[ elem ]
            dofs      = nodes2dofs( nodesElem, 6)

            if length( BCsData[indBC].NeumannNodalDOFs) >0
                dofsNeu = dofs[ BCsData[indBC].NeumannNodalDOFs ]
                FG[ dofsNeu ] = FG[ dofsNeu ] + BCsData[indBC].NeumannNodalVals
            end

            if length( BCsData[indBC].DirichletNodalDOFs) >0
                @assert sum( BCsData[indBC].DirichletNodalVals )==0 "only homogeneous dirichlet BC are allowed by now!"
                print("dofs:", dofs)
                neumDofs[ dofs[ BCsData[indBC].DirichletNodalDOFs ] ] .= 0
            end
        end
    end

    print("neumdofs", neumDofs)

    neumDofs = sort( unique(neumDofs) )
    neumDofs[1] == 0 && popfirst!(neumDofs)

    print("neumdofs", neumDofs)
    print("listo\n")

    return FG, neumDofs
end


"""
TransformedSet

Set obtained by some transformation from a feasible set (`generatingset`).
Should support `forward!` and `backward!` methods.
"""
abstract type TransformedSet <: FeasibleSet end

function generatingset(ts::TransformedSet)
error("Cannot find the generating set for $ts.")
end

"""
forward!(ts::TransformedSet)

In-place forward transform assosiated with the set `ts`: `x ∈ ts ⇔ f(x) ∈ ts.set`
Use `forward!(ts)(q,p)` to obtain `q = f(p)`.
"""
function forward!(ts::TransformedSet)
error("The forward transform is not defined for $ts.")
end

function backward!(ts::TransformedSet)
error("The backward transform is not defined for $ts.")
end

# it's easy to get an element of a transformed set
function getelement(ts::TransformedSet)
    el = getelement(generatingset(ts))
    backward!(ts)(getelement(generatingset(ts)))
end

"""
LinearTransformedSet

Subtype of TransformedSet where `forward` and `backward` transformations are given by multiplication
by forward and backward "plans" (i.e. -- precomputed matrices).
"""
abstract type LinearTransformedSet <: TransformedSet end

generatingset(s::LinearTransformedSet) = s.set
forward!(s::LinearTransformedSet) = ((xf, x) -> mul!(xf, s.fplan, x))
backward!(s::LinearTransformedSet) = (x -> mul!(x, s.bplan, x))

project!(xp, x, s::LinearTransformedSet) = (forward!(s)(xp, x); project!(xp, s.set); backward!(s)(xp))

struct FourierTransformedSet{TS,PF,PB} <: LinearTransformedSet where {TS <: FeasibleSet,PF <: AbstractFFTs.Plan,PB <: AbstractFFTs.Plan}
set::TS
fplan::PF
bplan::PB
end


using LinearAlgebra

# Postcompose HuberLoss

f = HuberLoss(1.0, 1.0)
g = Postcompose(f, 2.5)
h = HuberLoss(1.0, 2.5)

x = randn(10)
x = 0.5*x/norm(x)

gx = call_test(g, x)
hx = call_test(h, x)

@test abs(gx-hx)/(1+abs(gx)) <= 1e-12

grad_gx, gx1 = gradient_test(g, x)
grad_hx, hx1 = gradient_test(h, x)

@test abs(gx1-gx)/(1+abs(gx)) <= 1e-12
@test abs(hx1-hx)/(1+abs(hx)) <= 1e-12
@test norm(grad_gx-grad_hx, Inf)/(1+norm(grad_gx, Inf)) <= 1e-12

yg, gyg = prox_test(g, x, 1.3)
yh, hyh = prox_test(h, x, 1.3)

@test abs(gyg-hyh)/(1+abs(gyg)) <= 1e-12
@test norm(yg-yh, Inf)/(1+norm(yg, Inf)) <= 1e-12

x = randn(10)
x = 1.2*x/norm(x)

gx = call_test(g, x)
hx = call_test(h, x)

@test abs(gx-hx)/(1+abs(gx)) <= 1e-12

grad_gx, gx1 = gradient_test(g, x)
grad_hx, hx1 = gradient_test(h, x)

@test abs(gx1-gx)/(1+abs(gx)) <= 1e-12
@test abs(hx1-hx)/(1+abs(hx)) <= 1e-12
@test norm(grad_gx-grad_hx, Inf)/(1+norm(grad_gx, Inf)) <= 1e-12

yg, gyg = prox_test(g, x, 1.3)
yh, hyh = prox_test(h, x, 1.3)

@test abs(gyg-hyh)/(1+abs(gyg)) <= 1e-12
@test norm(yg-yh, Inf)/(1+norm(yg, Inf)) <= 1e-12


using BenchmarkTools

data11 = """
8548335644
6576521782
1223677762
1284713113
6125654778
6435726842
5664175556
1445736556
2248473568
6451473526
"""

function day11()
    part = [0, 0]
    day11lines = filter(!isempty, strip.(split(data11, "\n")))
    nrows, ncols = length(day11lines), length(first(day11lines))
    mat = zeros(Int, nrows, ncols)
    for i in 1:nrows
        mat[i, :] .= parse.(Int, collect(day11lines[i]))
    end

    for step in 1:1000
        mat .+= 1
        while any(mat .> 9)
            for c in CartesianIndices(mat)
                if mat[c] > 9
                    for k in max(1, c[1] - 1):min(c[1] + 1, nrows),
                        l in max(c[2] - 1, 1):min(c[2] + 1, ncols)

                        if mat[k, l] > 0
                            mat[k, l] += 1
                        end
                    end
                    mat[c] = 0
                    step < 101 && (part[1] += 1)
                end
            end
        end
        if all(==(first(mat)), mat)
            part[2] == 0 && (part[2] = step)
            step >= 100 && break
        end
    end
    return part
end

part = day11()
println("Part 1:", part[1])
println("Part 2:", part[2])

@btime day11()

#=
Part 1:1601
Part 2:368
  7.539 ms (86240 allocations: 2.71 MiB)
=#


using LinearAlgebra
import ClimaCore:
    Fields,
    Domains,
    Topologies,
    Meshes,
    DataLayouts,
    Operators,
    Geometry,
    Spaces

using OrdinaryDiffEq: ODEProblem, solve, SSPRK33

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

const FT = Float64

a = Geometry.ZPoint(FT(-20.0))
b = Geometry.ZPoint(FT(20.0))
n = 64
α = FT(0.1)

function heaviside(pt)
    0.5 * (sign(pt.z) + 1)
end

domain = Domains.IntervalDomain(a, b, boundary_names = (:left, :right))
mesh = Meshes.IntervalMesh(domain, nelems = n)

cs = Spaces.CenterFiniteDifferenceSpace(mesh)
fs = Spaces.FaceFiniteDifferenceSpace(cs)

V = Geometry.WVector.(ones(FT, fs))
θ = heaviside.(Fields.coordinate_field(cs))

# Solve advection Equation: ∂θ/dt = -∂(vθ)

# upwinding
function tendency1!(dθ, θ, _, t)
    fcc = Operators.FluxCorrectionC2C(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    fcf = Operators.FluxCorrectionF2F(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    UB = Operators.UpwindBiasedProductC2F(
        left = Operators.SetValue(sin(a.z - t)),
        right = Operators.SetValue(sin(b.z - t)),
    )
    ∂ = Operators.DivergenceF2C()

    return @. dθ = -∂(UB(V, θ))
end
function tendency2!(dθ, θ, _, t)
    fcc = Operators.FluxCorrectionC2C(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    fcf = Operators.FluxCorrectionF2F(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    UB = Operators.UpwindBiasedProductC2F(
        left = Operators.SetValue(sin(a.z - t)),
        right = Operators.SetValue(sin(b.z - t)),
    )
    ∂ = Operators.DivergenceF2C()
    return @. dθ = -∂(UB(V, θ)) + fcc(V, θ)
end
# use the advection operator
function tendency3!(dθ, θ, _, t)

    fcc = Operators.FluxCorrectionC2C(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    fcf = Operators.FluxCorrectionF2F(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    A = Operators.AdvectionC2C(
        left = Operators.SetValue(sin(-t)),
        right = Operators.Extrapolate(),
    )
    return @. dθ = -A(V, θ)
end
# use the advection operator
function tendency4!(dθ, θ, _, t)

    fcc = Operators.FluxCorrectionC2C(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    fcf = Operators.FluxCorrectionF2F(
        left = Operators.Extrapolate(),
        right = Operators.Extrapolate(),
    )
    A = Operators.AdvectionC2C(
        left = Operators.SetValue(sin(-t)),
        right = Operators.Extrapolate(),
    )
    return @. dθ = -A(V, θ) + fcc(V, θ)
end

# use the advection operator

@show tendency1!(similar(θ), θ, nothing, 0.0)
# Solve the ODE operator
Δt = 0.001
prob1 = ODEProblem(tendency1!, θ, (0.0, 5.0))
prob2 = ODEProblem(tendency2!, θ, (0.0, 5.0))
prob3 = ODEProblem(tendency3!, θ, (0.0, 5.0))
prob4 = ODEProblem(tendency4!, θ, (0.0, 5.0))
sol1 = solve(
    prob1,
    SSPRK33(),
    dt = Δt,
    saveat = 10 * Δt,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);
sol2 = solve(
    prob2,
    SSPRK33(),
    dt = Δt,
    saveat = 10 * Δt,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);
sol3 = solve(
    prob3,
    SSPRK33(),
    dt = Δt,
    saveat = 10 * Δt,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);
sol4 = solve(
    prob4,
    SSPRK33(),
    dt = Δt,
    saveat = 10 * Δt,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);

ENV["GKSwstype"] = "nul"
import ClimaCorePlots, Plots
Plots.GRBackend()

dirname = "advect_step_function"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

anim = Plots.@animate for u in sol1.u
    Plots.plot(u, xlim = (-1, 1))
end
Plots.mp4(anim, joinpath(path, "UBP_advect_step_function.mp4"), fps = 10)
Plots.png(
    Plots.plot(sol1.u[end], xlim = (-1, 1)),
    joinpath(path, "sol1_advect_step_function_end.png"),
)

anim = Plots.@animate for u in sol2.u
    Plots.plot(u, xlim = (-1, 1))
end
Plots.mp4(anim, joinpath(path, "UBP_advect_step_function_fc.mp4"), fps = 10)
Plots.png(
    Plots.plot(sol2.u[end], xlim = (-1, 1)),
    joinpath(path, "sol2_advect_step_function_end.png"),
)

anim = Plots.@animate for u in sol3.u
    Plots.plot(u, xlim = (-1, 1))
end
Plots.mp4(anim, joinpath(path, "C2C_advect_step_function.mp4"), fps = 10)
Plots.png(
    Plots.plot(sol3.u[end], xlim = (-1, 1)),
    joinpath(path, "sol3_advect_step_function_end.png"),
)

anim = Plots.@animate for u in sol4.u
    Plots.plot(u, xlim = (-1, 1))
end
Plots.mp4(anim, joinpath(path, "C2C_advect_step_function_fc.mp4"), fps = 10)
Plots.png(
    Plots.plot(sol4.u[end], xlim = (-1, 1)),
    joinpath(path, "sol4_advect_step_function_end.png"),
)

p = Plots.plot(sol1.u[end], xlim = (-1, 1), label = "UBP")
p = Plots.plot!(sol2.u[end], xlim = (-1, 1), label = "UBP_FC")
p = Plots.plot!(sol3.u[end], xlim = (-1, 1), label = "C2C")
p = Plots.plot!(sol4.u[end], xlim = (-1, 1), label = "C2C_FC")
Plots.png(p, joinpath(path, "all_advect_step_function_end.png"))

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig(
    "examples/column/output/$(dirname)/advect_step_function_end.png",
    "Advect End Simulation",
)

linkfig(
    relpath(
        joinpath(path, "advect_step_function_end.png"),
        joinpath(@__DIR__, "../.."),
    ),
    "Advect End Simulation",
)


using Test
using PETSc, MPI, LinearAlgebra, SparseArrays

# This shows some examples of how SNES can be used.
# there are 2 ways to use it:
#    1) you will have julia vectors (but PETSc matrices) within your user-routine
#       That is fine for sequential runs.
#    2) you will receive pointers to PETSc vectors within your user routines.
#       That is important for parallel simulations, where global residual
#       and solution vectors are passed to the user-routines, but we work with
#       local vectors

# Case 1) Julia vectors

x = ones(2)
# this is an example of a case where julia vectors are passed to the routines
function fn_julia!(fx::Vector, x::Vector, args...)
  fx[1] = x[1]^2 + x[1]*x[2] - 3
  fx[2] = x[1]*x[2] + x[2]^2 - 6
end

J = zeros(2,2)
PJ = PETSc.MatSeqDense(J)
function update_jac_julia!(x::Vector, args...)
    J[1,1] = 2x[1] + x[2]
    J[1,2] = x[1]
    J[2,1] = x[2]
    J[2,2] = x[1] + 2x[2]
end

S_julia = PETSc.SNES{Float64}(MPI.COMM_SELF; ksp_rtol=1e-4, pc_type="none", ksp_monitor=false, ksp_converged_reason=true)
PETSc.setfunction!(S_julia, fn_julia!, PETSc.VecSeq(zeros(2)))
PETSc.setjacobian!(S_julia, update_jac_julia!, PJ, PJ)
x_julia = PETSc.VecSeq([2.0, 3.0]);
b_julia = PETSc.VecSeq([0.0, 0.0]);
PETSc.solve!(x_julia, S_julia, b_julia)
@test x_julia.array ≈ [1.0,2.0] rtol=1e-4

# Case 2) PETSc vectors
# this is an example of a case where pointers to PETSc vectors are passed to the routines
function fn!(cfx, cx, user_ctx)
  # We could do Global->Local here on cfx/cx, provided a pointer to the local
  #  vector is available in user_ctx
  x  = PETSc.unsafe_localarray(Float64, cx;  write=false)   # read array
  fx = PETSc.unsafe_localarray(Float64, cfx; read=false)    # write array

  fx[1] = x[1]^2 + x[1]*x[2] - 3
  fx[2] = x[1]*x[2] + x[2]^2 - 6

  Base.finalize(fx)
  Base.finalize(x)
end

function update_jac!(cx, J1::PETSc.AbstractMat{Float64}, args...)
    x  = PETSc.unsafe_localarray(Float64, cx;  write=false)
    @show J1, typeof(J1)
    J1[1,1] = 2x[1] + x[2]
    J1[1,2] = x[1]
    J1[2,1] = x[2]
    J1[2,2] = x[1] + 2x[2]

    Base.finalize(x)
    PETSc.assemble(J1)
end

# structure with which we can pass data to the user-routines above
mutable struct Data
    vec
    julia
end

julia_vec = 0;  # we want pointers to local vectors
S = PETSc.SNES{Float64}(MPI.COMM_SELF, julia_vec;
                ksp_rtol=1e-4,
                pc_type="none",
                ksp_monitor=true,
                ksp_converged_reason=true)

data        = Data([100;2], 1)
S.user_ctx  = data;      # we can pack anything we need in this struct

PETSc.setfunction!(S, fn!, PETSc.VecSeq(zeros(2)))
PJ = PETSc.MatSeqDense(zeros(2,2))
PETSc.setjacobian!(S, update_jac!, PJ, PJ)
x  = PETSc.VecSeq([2.0, 3.0]);
b  = PETSc.VecSeq([0.0, 0.0]);
PETSc.solve!(x, S, b)
@test x.array ≈ [1.0,2.0] rtol=1e-4

# cleanup
#finalize(x)
#finalize(J);
#PETSc.destroy(x0);
#PETSc.destroy(x);
#PETSc.destroy(b);
#PETSc.destroy(PJ);
#PETSc.destroy(S);

#include("test_dmstag.jl")

#PETSc.finalize()


# The code to generate Fig 6A
using PyPlot
using SurNoR_2020
using Statistics

PyPlot.svg(true)
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["svg.fonttype"] = "none"
rcParams["pdf.fonttype"] = 42

# ------------------------------------------------------------------------------
# Training Algorithms
# ------------------------------------------------------------------------------
Agents_SurNoR = SurNoR_2020.Func_SurNoR_Overall_train_opt_ind_i(3)

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Fig 6A1: MF learning rate as a function if Surprise
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
p_c = Agents_SurNoR[1,1,1].Param.p_change   # p_c = m/(m+1) in the paper
m = Agents_SurNoR[1,1,1].Param.p_change / (1 - Agents_SurNoR[1,1,1].Param.p_change)

Surp = 0:0.1:1e6
Gamma = (m .* Surp) ./ (1 .+ m .* Surp)

α = Agents_SurNoR[1,1,1].Param.α        # ρ_b in the paper
Δα = Agents_SurNoR[1,1,1].Param.Δα      # δρ in the paper

α_S = α .+ Δα .* Gamma

y = α_S
x = Surp
fig = figure(figsize=(7,6)); ax = gca()
ax.plot(x,y,color="k")
ax.plot(x,ones(size(x)) .* (α + Δα) ,color="#B71616",linewidth=1,linestyle="--")
ax.plot(x,ones(size(x)) .* (α) ,color="#167CB7",linewidth=1,linestyle="--")
ax.set_xscale("log")
ax.set_xlim(1e0,1e6)
ax.set_ylim(0,0.6)

ax.set_title("MF learning rate vs Surp")
ax.set_xlabel("Surprise")
ax.set_ylabel("Model free learning rate")

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Fig 6A3: Surprise histogram
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
Sub_Num = 12
Surp = []

for Sub = 1:Sub_Num
    for Epi = 1:5
        for Block = 1:2
            Surp_temp = Agents_SurNoR[Sub,Block,Epi].Surprise
            append!(Surp, Surp_temp[2:end])
        end
    end
end

fig = figure(figsize=(7,6)); ax = gca()
ax.hist(Surp,bins= 10 .^ (-1:0.5:6)  , log=true)
ax.set_yscale("log")
ax.set_xscale("log")
ax.set_xlim(1e0,1e6)
ax.set_ylim(0,1e4)

ax.set_title("Surprise histogram")
ax.set_xlabel("log Surprise")
ax.set_ylabel("Count")

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Fig 6A4: Relative importance of MF to MB
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Computing Q values for
Mean_diff_ratio_2Parts = zeros((Sub_Num,2,2)) # MF to MB

Epi_Phase_Sets = [[1],[2,3,4,5]]
for Sub = 1:Sub_Num
    for Epi_phase = 1:2
        for Block = 1:2
            DQ_R_MB = []
            DQ_R_MF = []
            DQ_N_MB = []
            DQ_N_MF = []

            for Epi = Epi_Phase_Sets[Epi_phase]
                T = length(Agents_SurNoR[Sub,Block,Epi].Input.obs)

                Q_R_MB = zeros((4,T))
                Q_R_MF = zeros((4,T))

                Q_N_MB = zeros((4,T))
                Q_N_MF = zeros((4,T))

                for t = 1:T
                    Q_R_MB[:,t] = Agents_SurNoR[Sub,Block,Epi].Q_R_MB[4*Agents_SurNoR[Sub,Block,Epi].Input.obs[t] .+ (1:4), t]
                    Q_R_MF[:,t] = Agents_SurNoR[Sub,Block,Epi].Q_R_MF[4*Agents_SurNoR[Sub,Block,Epi].Input.obs[t] .+ (1:4), t]
                    Q_N_MB[:,t] = Agents_SurNoR[Sub,Block,Epi].Q_N_MB[4*Agents_SurNoR[Sub,Block,Epi].Input.obs[t] .+ (1:4), t]
                    Q_N_MF[:,t] = Agents_SurNoR[Sub,Block,Epi].Q_N_MF[4*Agents_SurNoR[Sub,Block,Epi].Input.obs[t] .+ (1:4), t]
                end

                DQ_R_MB_temp = findmax(Q_R_MB,dims=1)[1] .- findmin(Q_R_MB,dims=1)[1]
                DQ_R_MF_temp = findmax(Q_R_MF,dims=1)[1] .- findmin(Q_R_MF,dims=1)[1]

                DQ_N_MB_temp = findmax(Q_N_MB,dims=1)[1] .- findmin(Q_N_MB,dims=1)[1]
                DQ_N_MF_temp = findmax(Q_N_MF,dims=1)[1] .- findmin(Q_N_MF,dims=1)[1]

                append!(DQ_R_MB, DQ_R_MB_temp)
                append!(DQ_R_MF, DQ_R_MF_temp)
                append!(DQ_N_MB, DQ_N_MB_temp)
                append!(DQ_N_MF, DQ_N_MF_temp)
            end

            if Epi_phase != 1
                Mean_diff_ratio_2Parts[Sub,Epi_phase,Block] = mean(DQ_R_MF) / mean(DQ_R_MB)
            else
                if Block == 1
                    β_ratio_N2R = Agents_SurNoR[1,1,1].Param.β_ratio_N2R
                    Mean_diff_ratio_2Parts[Sub,Epi_phase,Block] = (mean(DQ_R_MF) .+ (mean(DQ_N_MF) .* β_ratio_N2R)) /
                                                           (mean(DQ_R_MB) .+ (mean(DQ_N_MB) .* β_ratio_N2R))
                elseif Block == 2
                    β_ratio_N2R = Agents_SurNoR[1,2,1].Param.β_ratio_N2R
                    Mean_diff_ratio_2Parts[Sub,Epi_phase,Block] = (mean(DQ_R_MF) .+ (mean(DQ_N_MF) .* β_ratio_N2R)) /
                                                           (mean(DQ_R_MB) .+ (mean(DQ_N_MB) .* β_ratio_N2R))
                end
            end
        end
    end
end

# Effective weights
w_scale = Agents_SurNoR[1,1,1].Param.w_MF_to_MB_Scaler
w_11 = Agents_SurNoR[1,1,1].Param.w_MF_to_MB_1
w_12 = Agents_SurNoR[1,2,1].Param.w_MF_to_MB_1
w_0 = Agents_SurNoR[1,1,1].Param.w_MF_to_MB_2

Eff_MF_weight_2Parts = deepcopy(Mean_diff_ratio_2Parts)
Eff_MF_weight_2Parts[:,1,1] = Eff_MF_weight_2Parts[:,1,1] .* w_scale .* w_11 ./ (1 - w_11)
Eff_MF_weight_2Parts[:,1,2] = Eff_MF_weight_2Parts[:,1,2] .* w_scale .* w_12 ./ (1 - w_12)
Eff_MF_weight_2Parts[:,2,1:2] = Eff_MF_weight_2Parts[:,2,1:2].* w_scale .* w_0 ./ (1 - w_0)

fig = figure(figsize=(7,6)); ax = gca()

x = 1:4
y = mean(Eff_MF_weight_2Parts,dims=1)[:]
dy = std(Eff_MF_weight_2Parts,dims=1)[:] ./ sqrt(12)
ax.bar(x, y[x])
ax.errorbar(x ,y[x],yerr=dy[x],color="k",linewidth=1,drawstyle="steps",capsize=3,linestyle="")

ax.set_xticks(1:4)
ax.set_xticklabels(["E1","E2to5","E1","E2to5"])
title("Hybrid weights MF to MB")


my_func(x,y) = 2x + y

derivative_of_my_func(x,y) = ForwardDiff.derivative(x->my_func(x,y),x)


"""
    AbstractLattice

Contains information about a crystal lattice.
"""
abstract type AbstractLattice
end

"""
    RealLattice{N}

Describes an N-dimensional real space crystal lattice, with an associated primitive and
conventional lattice. These lattices may be the same depending on the lattice type.
"""
struct RealLattice{N} <: AbstractLattice
    # Primitive and conventional lattice vectors are stored together
    prim::SVector{N,SVector{N,Float64}}
    conv::SVector{N,SVector{N,Float64}}
end

# TODO: Implement lattice checking and conversion
# Some of this functionality might already be present in Crystalline.jl
# However, we need to be able to recognize primitive lattices with their own unique conditions

"""
    nuniqlen(basis::AbstractVector{<:AbstractVector{<:Real}}) -> Int

Returns the number of vectors that are the same length.
"""
function nuniqlen(basis::AbstractVector{<:AbstractVector{<:Real}})
    # unique!() does better than unique() here (per @btime)
    return length(unique!(norm.(basis)))
end

# Same thing for matrices
nuniqlen(basis::AbstractMatrix{<:Real}) = nuniqlen(_tovectors(basis))

"""
    northog(basis::AbstractVector{<:AbstractVector{<:Real}}) -> Int

Returns the number of pairs of orthogonal vectors.
"""
function northog(basis::AbstractVector{<:AbstractVector{<:Real}})
    # FIXME: This doesn't work for arbitrary dimensions
    # It needs to use the upper portion of the dot product matrix
    return count(iszero, dot.(basis, basis[circshift(1:length(basis), 1)]))
end

northog(basis::AbstractMatrix{<:Real}) = northog(_tovectors(basis))

# TODO: Write a number of tests for checking if lattices meet certain criteria

#=
Criteria for primitive unit cells in 3 dimensions:
    aP: 6 parameters (a, b, c, α, β, γ)
    mP: 4 parameters (a, b, c, β), α = γ = 90°
    mS: 4 parameters (a, c, α, γ), a = b, α = β
    oP: 3 parameters (a, b, c), α = β = γ = 90°
    oS: 3 parameters (a, c, γ), a = b, α = β = 90°
    oI: 3 parameters (a, α, γ), a = b = c, |cos(α)| + |cos(β)| + |cos(γ)| = 1
    oF: 3 parameters (a, b, c), sin(α)/a = sin(β)/b = sin(γ)/c, α + β + γ = 180°
        3 parameters (a, c, γ), [forgot the rest of this alternate test]
    tP: 2 parameters (a, c), a = b, α = β = γ = 90°
    tI: 2 parameters (a, α), a = b = c, |cos(α)| + |cos(β)| + |cos(γ)| = 1
    hP: 2 parameters (a, c), a = b, α = β = 90°, γ = 120°
    hR: 2 parameters (a, α), a = b = c, α = β = γ
    cP: 1 parameter (a), a = b = c, α = β = γ = 90°
    cI: 1 parameter (a), a = b = c, α = β = γ = arccos(-1/3)°
    cF: 1 parameter (a), a = b = c, α = β = γ = 60°

Rules that need to be checked:
    Number of unique lengths
    Number of orthogonal vectors
    Number of unique angles
    Cosine rule: |cos(α)| + |cos(β)| + |cos(γ)| = 1
    Sine rule: sin(α)/a = sin(β)/b = sin(γ)/c, α + β + γ = 180°

    The last three are probably going to be trickier to implement, or at the very least there
    might be some unintuitive linear algebra tricks to make those checks happen faster.
=#

# TODO: what are the crtieria for a "good" lattice basis?
# Potentially worth looking at the Niggli reduction algorithm to do this

# TODO: Tools for working with reciprocal space lattices


using Test
using DynamicPolynomials
# using Revise
using PolynomialBasis
PB = PolynomialBasis

function allapprox(v1,v2)
    return all(v1 .≈ v2)
end

basis = PB.HermiteTensorProductBasis(2)

testvals = zeros(16)
testvals[1] = 1.0
@test allapprox(basis(-1.,-1.),testvals)
@test allapprox(basis([-1.,-1.]),testvals)

fill!(testvals,0.0)
testvals[5] = 1.0
@test allapprox(basis(-1.,1.),testvals)
@test allapprox(basis([-1.,1.]),testvals)

fill!(testvals,0.0)
testvals[9] = 1.0
@test allapprox(basis(1,-1),testvals)
@test allapprox(basis([1,-1]),testvals)

fill!(testvals,0.0)
testvals[13] = 1.0
@test allapprox(basis(1,1),testvals)
@test allapprox(basis([1,1]),testvals)

testvals = zeros(16,2)
testvals[2,2] = 1.0
testvals[3,1] = 1.0
@test allapprox(gradient(basis,-1,-1),testvals)
@test allapprox(gradient(basis,[-1,-1]),testvals)

fill!(testvals,0.0)
testvals[6,2] = 1.0
testvals[7,1] = 1.0
@test allapprox(gradient(basis,-1,1),testvals)
@test allapprox(gradient(basis,[-1,1]),testvals)

fill!(testvals,0.0)
testvals[10,2] = 1.0
testvals[11,1] = 1.0
@test allapprox(gradient(basis,1,-1),testvals)
@test allapprox(gradient(basis,[1,-1]),testvals)

fill!(testvals,0.0)
testvals[14,2] = 1.0
testvals[15,1] = 1.0
@test allapprox(gradient(basis,1,1),testvals)
@test allapprox(gradient(basis,[1,1]),testvals)


#=
GasLaws.jl:
- Julia version: 1.5.1
- Author: tbald
- Date: 2020-12-20
=#

module GasLaws
using Unitful
using Unitful: 𝐋,𝐍

const R=.0831446261815324u"bar*L*K^-1*mol^-1"
@derived_dimension MolarVolume 𝐋^3/𝐍

#Returns amount in desired unit
function ideal(P::Unitful.Pressure,V::Unitful.Volume,
               n::Unitful.AmountUnits,T::Unitful.Temperature)::Unitful.Amount
    newT= T |> u"K"
    units=Unitful.unit(P)*Unitful.unit(V)/(u"K"*n)
    newR=R |> units
    newN=P*V/(newR*newT)
end

#Returns Temperature in desired unit
function ideal(P::Unitful.Pressure,V::Unitful.Volume,
               n::Unitful.Amount,T::Unitful.TemperatureUnits)::Unitful.Temperature
    units=Unitful.unit(P)*Unitful.unit(V)/(u"K"*Unitful.unit(n))
    newR=R |> units
    newT=P*V/(newR*n)
    newT |> T
end

#Returns Volume in desired unit
function ideal(P::Unitful.Pressure,V::Unitful.VolumeUnits,
               n::Unitful.Amount,T::Unitful.Temperature)::Unitful.Volume
    newT= T |> u"K"
    units=Unitful.unit(P)*V/(u"K"*Unitful.unit(n))
    newR=R |> units
    newV=(n*newR*newT)/P
end

#Returns Pressure in desired unit
function ideal(P::Unitful.PressureUnits,V::Unitful.Volume,
               n::Unitful.Amount,T::Unitful.Temperature)::Unitful.Pressure
    newT= T |> u"K"
    units=P*Unitful.unit(V)/(u"K"*Unitful.unit(n))
    newR=R |> units
    newP=(n*newR*newT)/V
end

#Solves for pressure using Van der Waals
#Default a and b for CO2
function Van_der_Waals(Vm::MolarVolume,T::Unitful.Temperature;a=3.658,b=0.04286)::typeof(1.0u"bar")
    full_a=a*1.0u"bar*L^2*mol^-2"
    full_b=b*1.0u"L*mol^-1"
    newT= T |> u"K"

    p=R*newT/(Vm-full_b)
    p+=-full_a/Vm^2
end

#Solves for pressure using Redlich-Kwong equation
##Default a and b for CO2
function Redlich_Kwong(Vm::MolarVolume,T::Unitful.Temperature;a=64.43,b=0.02963)::typeof(1.0u"bar")
    full_a=a*1.0u"bar*L^2*mol^-2*K^(1/2)"
    full_b=b*1.0u"L*mol^-1"
    newT= T |> u"K"

    p=R*newT/(Vm-full_b)
    p+=-full_a/(newT^(1/2)*Vm*(Vm+full_b))
end

##Could add reduced versions of the real gas laws

end


using AOPE, BusinessDays, Dates
using Test

@testset "settle" begin
    @test liqd(Date(2010, 8, 13), "T+3", :UKSettlement) == Date(2010, 8, 18)
    @test liqd(Date(2010, 8, 13),     3, :UKSettlement) == Date(2010, 8, 18)
    @test liqd(Date(2010, 8, 13),     0, :UKSettlement) == Date(2010, 8, 13)
end

# based on tests from OpenGamma Strata
# https://github.com/OpenGamma/Strata/blob/efaa0be0d08dc61c23692e7b86df101ea0fb7223/modules/basics/src/test/java/com/opengamma/strata/basics/date/DayCountTest.java
@testset "Actual365Fixed" begin
    dc = Actual365Fixed()
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == (4 + 58) / 365
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == (4 + 59) / 365
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == (4 + 60) / 365
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == (4 + 366 + 365 + 365 + 365 + 58) / 365
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == (4 + 366 + 365 + 365 + 365 + 59) / 365
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == (4 + 366 + 365 + 365 + 365 + 60) / 365
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == 29 / 365
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == 28 / 365
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == 27 / 365

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "Actual360" begin
    dc = Actual360()
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == (4 + 58) / 360
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == (4 + 59) / 360
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == (4 + 60) / 360
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == (4 + 366 + 365 + 365 + 365 + 58) / 360
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == (4 + 366 + 365 + 365 + 365 + 59) / 360
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == (4 + 366 + 365 + 365 + 365 + 60) / 360
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == 29 / 360
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == 28 / 360
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == 27 / 360

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "ActualActualISDA" begin
    dc = ActualActualISDA()
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == 4 / 365 + 58 / 366
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == 4 / 365 + 59 / 366
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == 4 / 365 + 60 / 366
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == 4 / 365 + 58 / 366 + 4
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == 4 / 365 + 59 / 366 + 4
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == 4 / 365 + 60 / 366 + 4
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == 29 / 366
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == 28 / 366
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == 27 / 366

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "Thirty360" begin
    dc = Thirty360()
    # usual rule: ((d2-d1) + (m2-m1)*30 + (y2-y1)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == ((28-28) + (2-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == ((29-28) + (2-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == ((1-28) + (3-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == ((28-28) + (2-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == ((29-28) + (2-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == ((1-28) + (3-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == ((28-28) + (3-2)*30 + (2012-2012)*360)/360
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == ((28-29) + (3-2)*30 + (2012-2012)*360)/360
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == ((28-1) + (3-3)*30 + (2012-2012)*360)/360

    @test frac(Date(2012, 5,29), Date(2013, 8,29), dc) == ((29-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,29), Date(2013, 8,30), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,29), Date(2013, 8,31), dc) == ((31-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "ThirtyE360" begin
    dc = ThirtyE360()
    # usual rule: ((d2-d1) + (m2-m1)*30 + (y2-y1)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == ((28-28) + (2-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == ((29-28) + (2-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == ((1-28) + (3-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == ((28-28) + (2-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == ((29-28) + (2-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == ((1-28) + (3-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == ((28-28) + (3-2)*30 + (2012-2012)*360)/360
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == ((28-29) + (3-2)*30 + (2012-2012)*360)/360
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == ((28-1) + (3-3)*30 + (2012-2012)*360)/360

    @test frac(Date(2012, 5,29), Date(2013, 8,29), dc) == ((29-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,29), Date(2013, 8,30), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,29), Date(2013, 8,31), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,30), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "ThirtyE360ISDA" begin
    @testset "maturity not end of Feb" begin
        dc = ThirtyE360ISDA(Date(2012, 8, 31))
        @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == ((28-28) + (2-12)*30 + (2012-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == ((30-28) + (2-12)*30 + (2012-2011)*360)/360 # exception
        @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == ((1-28) + (3-12)*30 + (2012-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == ((28-28) + (2-12)*30 + (2016-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == ((30-28) + (2-12)*30 + (2016-2011)*360)/360 # exception
        @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == ((1-28) + (3-12)*30 + (2016-2011)*360)/360
        @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == ((28-28) + (3-2)*30 + (2012-2012)*360)/360
        @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == ((28-30) + (3-2)*30 + (2012-2012)*360)/360 # exception
        @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == ((28-1) + (3-3)*30 + (2012-2012)*360)/360

        @test frac(Date(2012, 5,29), Date(2013, 8,29), dc) == ((29-29) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,29), Date(2013, 8,30), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,29), Date(2013, 8,31), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,30), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,30), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,30), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,31), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,31), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,31), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception

        # zeros
        @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
        @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
        @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
        @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
    end

    @testset "maturity end of Feb" begin
        dc = ThirtyE360ISDA(Date(2012, 2,29)) # leap year, end of month
        @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == ((28-28) + (2-12)*30 + (2012-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == ((29-28) + (2-12)*30 + (2012-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == ((1-28) + (3-12)*30 + (2012-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == ((28-28) + (2-12)*30 + (2016-2011)*360)/360
        @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == ((30-28) + (2-12)*30 + (2016-2011)*360)/360 # exception
        @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == ((1-28) + (3-12)*30 + (2016-2011)*360)/360
        @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == ((28-28) + (3-2)*30 + (2012-2012)*360)/360
        @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == ((28-30) + (3-2)*30 + (2012-2012)*360)/360 # exception
        @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == ((28-1) + (3-3)*30 + (2012-2012)*360)/360

        @test frac(Date(2012, 5,29), Date(2013, 8,29), dc) == ((29-29) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,29), Date(2013, 8,30), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,29), Date(2013, 8,31), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,30), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,30), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360
        @test frac(Date(2012, 5,30), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,31), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,31), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
        @test frac(Date(2012, 5,31), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception

        # zeros
        @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
        @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
        @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
        @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
    end
end

@testset "ActualActualICMA" begin
    dc = ActualActualICMA(Date(2011,7,1):Month(6):Date(2012,7,1))

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "ActuaActual ISDA Memo" begin
    # these test cases are all derived from
    # The Actual/Actual Day Count Fraction, ISDA
    # Paper for use with the ISDA Market Conventions Survey - 3rd June 1999
    # https://www.isda.org/a/pIJEE/The-Actual-Actual-Day-Count-Fraction-1999.pdf

    @testset "Semi-annual payment" begin
        dc_isda = ActualActualISDA()
        dc_icma = ActualActualICMA(Date(2003,11,1):Month(6):Date(2004,5,1))
    end

    @testset "short first" begin
        dc_isda = ActualActualISDA()
        dc_icma = ActualActualICMA(Date(1998,7,1):Month(12):Date(2000,7,1))

        @test frac(Date(1999,2,1), Date(1999,7,1), dc_isda) == 150/365
        @test frac(Date(1999,2,1), Date(1999,7,1), dc_icma) ≈ 150/365 atol=eps()

        @test frac(Date(1999,7,1), Date(2000,7,1), dc_isda) == 184/365 + 182/366
        @test frac(Date(1999,7,1), Date(2000,7,1), dc_icma) == 366/366

        @test frac(Date(1999,2,1), Date(2000,7,1), dc_isda) == (150 + 184)/365 + 182/366
        @test frac(Date(1999,2,1), Date(2000,7,1), dc_icma) ≈ 150/365 + 366/366 atol=eps()
    end

    @testset "long first" begin
        dc_isda = ActualActualISDA()
        dc_icma = ActualActualICMA(Date(2002,7,15):Month(6):Date(2004,1,15))

        @test frac(Date(2002,8,15), Date(2003,7,15), dc_isda) == 334/365
        @test frac(Date(2002,8,15), Date(2003,7,15), dc_icma) == 181/(181*2) + 153/(184*2)

        @test frac(Date(2003,7,15),Date(2004,1,15), dc_isda) == 170/365 + 14/366
        @test frac(Date(2003,7,15), Date(2004,1,15), dc_icma) == 184/(184*2)

        @test frac(Date(2002,8,15), Date(2004,1,15), dc_isda) ≈ (334+170)/365 + 14/366 rtol=eps()
        @test frac(Date(2002,8,15), Date(2004,1,15), dc_icma) == 181/(181*2) + (153+184)/(184*2)
    end

    @testset "short final" begin
        dc_isda = ActualActualISDA()
        dc_icma = ActualActualICMA(Date(1999,7,30):Month(6):Date(2000,7,30))

        @test frac(Date(1999,7,30), Date(2000,1,30), dc_isda) == 155/365 + 29/366
        @test frac(Date(1999,7,30), Date(2000,1,30), dc_icma) == 184/(184*2)

        @test frac(Date(2000,1,30), Date(2000,6,30), dc_isda) == 152/366
        @test frac(Date(2000,1,30), Date(2000,6,30), dc_icma) == 152/(182*2)

        @test frac(Date(1999,7,30), Date(2000,6,30), dc_isda) == 155/365 + (29+152)/366
        @test frac(Date(1999,7,30), Date(2000,6,30), dc_icma) == 184/(184*2) + 152/(182*2)
    end

    @testset "long final" begin
        dc_isda = ActualActualISDA()
        # TODO: need a way to express "last day of month" schedules.
        dc_icma = ActualActualICMA(Date(1999,5,31):Month(3):Date(2000,5,31))

        @test frac(Date(1999,11,30), Date(2000,4,30), dc_isda) == 32/365 + 120/366
        @test frac(Date(1999,11,30), Date(2000,4,30), dc_icma) == 91/(91*4) + 61/(92*4)
    end
end

@testset "Thirty360Excel" begin
    dc = Thirty360Excel()
    # usual rule: ((d2-d1) + (m2-m1)*30 + (y2-y1)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == ((28-28) + (2-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == ((29-28) + (2-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == ((1-28) + (3-12)*30 + (2012-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == ((28-28) + (2-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == ((29-28) + (2-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == ((1-28) + (3-12)*30 + (2016-2011)*360)/360
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == ((28-28) + (3-2)*30 + (2012-2012)*360)/360
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == ((28-30) + (3-2)*30 + (2012-2012)*360)/360 # exception
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == ((28-1) + (3-3)*30 + (2012-2012)*360)/360

    @test frac(Date(2012, 5,29), Date(2013, 8,29), dc) == ((29-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,29), Date(2013, 8,30), dc) == ((30-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,29), Date(2013, 8,31), dc) == ((31-29) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360
    @test frac(Date(2012, 5,30), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,29), dc) == ((29-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,30), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception
    @test frac(Date(2012, 5,31), Date(2013, 8,31), dc) == ((30-30) + (8-5)*30 + (2013-2012)*360)/360 # exception

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "ActualActualExcel" begin
    dc = ActualActualExcel()
    # calculated from Excel sheet
    @test frac(Date(2011,12,28), Date(2012, 2,28), dc) == 62/365
    @test frac(Date(2011,12,28), Date(2012, 2,29), dc) == 63/366
    @test frac(Date(2011,12,28), Date(2012, 3, 1), dc) == 64/366
    @test frac(Date(2011,12,28), Date(2016, 2,28), dc) == 1523/365.3333333333333
    @test frac(Date(2011,12,28), Date(2016, 2,29), dc) == 1524/365.3333333333333
    @test frac(Date(2011,12,28), Date(2016, 3, 1), dc) == 1525/365.3333333333333
    @test frac(Date(2012, 2,28), Date(2012, 3,28), dc) == 29/366
    @test frac(Date(2012, 2,29), Date(2012, 3,28), dc) == 28/366
    @test frac(Date(2012, 3, 1), Date(2012, 3,28), dc) == 27/366

    @test frac(Date(2012, 2,28), Date(2013, 2,28), dc) == 366/366
    @test frac(Date(2012, 2,29), Date(2013, 2,28), dc) == 365/366
    @test frac(Date(2012, 3, 1), Date(2013, 2,28), dc) == 364/365
    @test frac(Date(2012, 2,28), Date(2013, 3, 1), dc) == 367/365.5
    @test frac(Date(2012, 2,29), Date(2013, 3, 1), dc) == 366/365.5
    @test frac(Date(2012, 3, 1), Date(2013, 3, 1), dc) == 365/365

    @test frac(Date(2011, 2,28), Date(2012, 2,28), dc) == 365/365
    @test frac(Date(2011, 3, 1), Date(2012, 2,28), dc) == 364/365
    @test frac(Date(2011, 2,28), Date(2012, 2,29), dc) == 366/365.5
    @test frac(Date(2011, 3, 1), Date(2012, 2,29), dc) == 365/366
    @test frac(Date(2011, 2,28), Date(2012, 3, 1), dc) == 367/365.5
    @test frac(Date(2011, 3, 1), Date(2012, 3, 1), dc) == 366/366

    # zeros
    @test frac(Date(2011,12,28), Date(2011,12,28), dc) == 0
    @test frac(Date(2011,12,31), Date(2011,12,31), dc) == 0
    @test frac(Date(2012, 2,28), Date(2012, 2,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), dc) == 0
end

@testset "ActualActualISMA" begin
    dc = ActualActualISMA()
    # The following two test cases are derived from:
    # Accrued interest – a guide for private investors, January 2015, page 2.
    # Published by the London Stock Exchange Group:
    # https://www.londonstockexchange.com/traders-and-brokers/security-types/retail-bonds/accrued-interest-corp-supra.pdf
    @test frac(Date(2014,11,26), Date(2015, 4,1), Date(2015,11,26), dc) == 126/(365 * 1)  # annual coupon
    @test frac(Date(2015, 3, 7), Date(2015, 4,9), Date(2015, 9, 7), dc) ==  33/(184 * 2)  # semi-annual coupoun

    # zeros
    @test frac(Date(2012, 2,28), Date(2012, 2,28), Date(2012, 8,28), dc) == 0
    @test frac(Date(2012, 2,29), Date(2012, 2,29), Date(2012, 8,28), dc) == 0
end

# Dates for the swap example from http://www.derivativepricing.com/blogpage.asp?id=19:
# Effective Date (Start Date): Feb 28, 2008
# Maturity Date (End Date): Feb 28, 2011
# Frequency: Semi-Annual
@testset "Todos os dias são Dias Úteis: AllBD" begin
    @test utild(Date(2019, 12, 25),  AllBD, :USNYSE) == Date(2019,  12, 25)
    @test utild(Date(2019, 12, 25),  AllBD) == Date(2019,  12, 25)
end

@testset "Convenção Dias Úteis: Following" begin
    @test utild(Date(2019, 2, 23),    Following, :USNYSE) == Date(2019,  2, 25)
end
@testset "Following: Contrato de Swap" begin
    @test utild(Date(2008, 2, 28), Following, :USNYSE) == Date(2008,  2, 28)
    @test utild(Date(2008, 8, 28), Following, :USNYSE) == Date(2008,  8, 28)
    @test utild(Date(2009, 2, 28), Following, :USNYSE) == Date(2009,  3,  2)
    @test utild(Date(2009, 8, 28), Following, :USNYSE) == Date(2009,  8, 28)
    @test utild(Date(2010, 2, 28), Following, :USNYSE) == Date(2010,  3,  1)
    @test utild(Date(2010, 8, 28), Following, :USNYSE) == Date(2010,  8, 30)
    @test utild(Date(2011, 2, 28), Following, :USNYSE) == Date(2011,  2, 28)
end

@testset "Convenção Dias Úteis: Modified Following" begin
    @test utild(Date(2019, 2, 23),    ModifiedFollowing, :USNYSE) == Date(2019,  2, 25)
end
@testset "Modified Following: Contrato de Swap" begin
    @test utild(Date(2008, 2, 28), ModifiedFollowing, :USNYSE) == Date(2008,  2, 28)
    @test utild(Date(2008, 8, 28), ModifiedFollowing, :USNYSE) == Date(2008,  8, 28)
    @test utild(Date(2009, 2, 28), ModifiedFollowing, :USNYSE) == Date(2009,  2, 27)
    @test utild(Date(2009, 8, 28), ModifiedFollowing, :USNYSE) == Date(2009,  8, 28)
    @test utild(Date(2010, 2, 28), ModifiedFollowing, :USNYSE) == Date(2010,  2, 26)
    @test utild(Date(2010, 8, 28), ModifiedFollowing, :USNYSE) == Date(2010,  8, 30)
    @test utild(Date(2011, 2, 28), ModifiedFollowing, :USNYSE) == Date(2011,  2, 28)
end
# Note how in coupon #2 and #4 the dates are adjusted to the previous business date. While coupon #5 was adjusted to the next good business date.
@testset "Convenção Dias Úteis: Preceding" begin
    @test utild(Date(2019, 2, 23),    Preceding, :USNYSE) == Date(2019,  2, 22)
end
@testset "Convenção Dias Úteis: ModifiedPreceding" begin
    @test utild(Date(2019, 2, 23),    Preceding, :USNYSE) == Date(2019,  2, 22)
end
@testset "Convenção Dias Úteis: EndOfMonth (Unadjusted)" begin
    @test utild(Date(2019, 2, 23),     EndOfMonth) == Date(2019,  2, 28)
    @test utild(Date(2020, 2, 23),     EndOfMonth) == Date(2020,  2, 29)
    @test utild(Date(2020, 2, 29),     EndOfMonth) == Date(2020,  2, 29)
end
@testset "EndOfMonth (Unadjusted): Contrato de Swap" begin
    @test utild(Date(2008, 2, 28),  EndOfMonth) == Date(2008,  2, 29)
    @test utild(Date(2008, 8, 28),  EndOfMonth) == Date(2008,  8, 31)
    @test utild(Date(2009, 2, 28),  EndOfMonth) == Date(2009,  2, 28)
    @test utild(Date(2009, 8, 28),  EndOfMonth) == Date(2009,  8, 31)
    @test utild(Date(2010, 2, 28),  EndOfMonth) == Date(2010,  2, 28)
    @test utild(Date(2010, 8, 28),  EndOfMonth) == Date(2010,  8, 31)
    @test utild(Date(2011, 2, 28),  EndOfMonth) == Date(2011,  2, 28)
end
@testset "Convenção Dias Úteis: Aplicada a Matriz de Datas" begin
    A = [Date(2019,3,2) 0.05; Date(2019,3,4) 0.07; Date(2019,3,5) 0.02]
    B = [Date(2019,3,2) 0.05; Date(2019,3,4) 0.07; Date(2019,3,5) 0.02]
    @test bdarray(A, Following, :TARGET)[1,1] == Date(2019,3,4)
    @test bdarray(B, Following, :TARGET, [1])[1,1] == Date(2019,3,4)
end
@testset "Fluxos Caixa Obrigação: Com e Sem Amortização do Principal" begin
    a = fcaixa(Date(2019,3,3), Date(2021,3,12), 0.05, 2, principal = 1000)
    @test size(a) == (5,2)
    @test a[1,1] == Date(2021,3,12)
    @test a[3,1] == Date(2020,3,12)
    @test a[1,2] == 1025.0
    @test a[3,2] ==   25.0

    b = fcaixa(Date(2019,3,3), Date(2021,3,12), 0.05, 2, 100, principal = 1000)
    @test size(b) == (5,2)
    @test b[1,1] == Date(2021,3,12)
    @test b[3,1] == Date(2020,3,12)
    @test b[1,2] == 102.49999999999999
    @test b[3,2] == 107.5
end
@testset "z-spread de uma obrigação" begin
    price = 102.30
    cf = [3.50, 103.50]
    m = [0.589, 1.589]
    r = [0.0020, 0.0035]
    z = 0.025390646751015725
    compos = Continuous   # composição contínua
    @test spreço(cf,m,r, z, compos) == 102.30
    @test zspread(price, cf, m, r, compos) == 0.02539064675101572
end
@testset "dirty price <-> yield-to-maturity: Formula normal" begin
    preço = 108.185342
    fck = [2.20, 2.20, 2.20, 2.20, 102.20]
    tk = [0.665753, 1.665753, 2.67122, 3.66848, 4.66574]
    y = 0.005767516967531196
    compos = Discrete   # composição discreta
    @test spreço(fck,tk, y, compos) == 108.18534199999998
    @test ytm(preço, fck, tk, compos) == 0.005767516967531196
end


# Type definition
struct FourierSeries{T,S<:Real} <: FunctionSeries{T}
    coeffs::Vector{T}
    period::S
    function FourierSeries(coeffs::Vector{T}, period=1) where {T}
        isodd(length(coeffs)) ||
            throw(ArgumentError("expect odd length for coefficient vector"))
        S = typeof(period)
        new{T,S}(coeffs, period)
    end
end

# Basic properties
period(f::FourierSeries) = f.period
domain(f::FourierSeries) = (zero(period(f)), period(f))

nummodes(f::FourierSeries) = (length(f.coeffs)-1) >> 1
coefficients(f::FourierSeries) = f.coeffs

function Base.:(==)(f1::FourierSeries, f2::FourierSeries)
    period(f1) == period(f2) || return false
    K = max(nummodes(f1), nummodes(f2))
    for k = -K:K
        f1[k] == f2[k] || return false
    end
    return true
end

# Indexing, getting and setting coefficients
Base.eachindex(f::FourierSeries) = UnitRange(-nummodes(f), nummodes(f))
function Base.getindex(f::FourierSeries, k)
    @inbounds if abs(k) > nummodes(f)
        return zero(f.coeffs[1])
    else
        j = ifelse(k == 0, 1, ifelse(k < 0, -2*k+1, 2*k))
        return f.coeffs[j]
    end
end
function Base.setindex!(f::FourierSeries, v, k)
    if abs(k) > nummodes(f)
        setnummodes!(f, abs(k))
    end
    j = ifelse(k == 0, 1, ifelse(k < 0, -2*k+1, 2*k))
    @inbounds setindex!(f.coeffs, v, j)
end

# Use as function
function (f::FourierSeries)(x̃)
    xred = x̃/period(f)
    x = mod1(xred, one(xred))
    K = nummodes(f)
    factor = 2pi
    s, c = sincos(factor*0*x)
    v = zero(f[0]) + f[0]*(c+im*s) # for type stability
    for k = 1:K
        s, c = sincos(factor*k*x)
        v += (f[k]+f[-k])*c + im*(f[k]-f[-k])*s
    end
    return v
end

# Change number of coefficients
function truncate!(f::FourierSeries; Kmax::Integer = nummodes(f), tol::Real = 0)
    Ktol = findlast(x->(norm(x) >= tol), f.coeffs)
    Kmax = min(Kmax, Ktol === nothing ? 0 : Ktol>>1)
    if Kmax < nummodes(f)
        resize!(f.coeffs, 2*Kmax+1)
    end
    return f
end

function setnummodes!(f::FourierSeries, K::Int)
    while nummodes(f) < K
        push!(f.coeffs, zero(f.coeffs[1]), zero(f.coeffs[1]))
    end
    if nummodes(f) > K
        resize!(f.coeffs, 2*K+1)
    end
    return f
end

# Special purpose constructor
Base.similar(f::FourierSeries, ::Type{T} = scalartype(f)) where {T} =
    FourierSeries([zero(T)*f[0]], period(f))

Base.zero(f::FourierSeries) = FourierSeries([zero(f[0])], period(f))
Base.one(f::FourierSeries) = FourierSeries([one(f[0])], period(f))

# Arithmetic (out of place)
Base.copy(f::FourierSeries) = FourierSeries(copy.(coefficients(f)), period(f))

Base.:-(f::FourierSeries) = FourierSeries(.-coefficients(f), period(f))
Base.:+(f::FourierSeries) = FourierSeries(.+coefficients(f), period(f))

Base.:*(f::FourierSeries, a::Const) =
    FourierSeries([c*a for c in coefficients(f)], period(f))
Base.:*(a::Const, f::FourierSeries) =
    FourierSeries([a*c for c in coefficients(f)], period(f))
Base.:/(f::FourierSeries, a::Const) =
    FourierSeries([c/a for c in coefficients(f)], period(f))
Base.:\(a::Const, f::FourierSeries) =
    FourierSeries([a\c for c in coefficients(f)], period(f))

function Base.:+(f1::FourierSeries, f2::FourierSeries)
    domain(f1) == domain(f2) || throw(DomainMismatch())
    K1 = nummodes(f1)
    K2 = nummodes(f2)
    K = max(K1, K2)
    f = FourierSeries(sizehint!([f1[0]+f2[0]], 2*K+1), period(f1))
    for k = 1:K
        if k <= K1 && k <= K2
            f[k] = f1[k] + f2[k]
            f[-k] = f1[-k] + f2[-k]
        elseif k <= K1
            f[k] = copy(f1[k])
            f[-k] = copy(f1[-k])
        else
            f[k] = copy(f2[k])
            f[-k] = copy(f2[-k])
        end
    end
    return f
end

function Base.:-(f1::FourierSeries, f2::FourierSeries)
    domain(f1) == domain(f2) || throw(DomainMismatch())
    K1 = nummodes(f1)
    K2 = nummodes(f2)
    K = max(K1, K2)
    f = FourierSeries(sizehint!([f1[0] - f2[0]], 2*K+1), period(f1))
    for k = 1:K
        if k <= K1 && k <= K2
            f[k] = f1[k] - f2[k]
            f[-k] = f1[-k] - f2[-k]
        elseif k <= K1
            f[k] = copy(f1[k])
            f[-k] = copy(f1[-k])
        else
            f[k] = -(f2[k])
            f[-k] = -(f2[-k])
        end
    end
    return f
end

Base.:*(f1::FourierSeries, f2::FourierSeries) = truncmul(f1, f2)

function truncmul(f1::FourierSeries, f2::FourierSeries;
                    Kmax::Integer = nummodes(f1) + nummodes(f2), tol::Real = 0)
    period(f1) == period(f2) || throw(DomainMismatch())
    f = FourierSeries(sizehint!([zero(f1[0])*zero(f2[0])], 2*Kmax+1), period(f1))
    truncmul!(f, f1, f2, true, false; Kmax = Kmax, tol = tol)
    return f
end

function Base.conj(f::FourierSeries)
    fc = FourierSeries(map(conj, coefficients(f)), period(f))
    K = nummodes(f)
    for k = 1:K
        fc[k], fc[-k] = fc[-k], fc[k]
    end
    return fc
end
function Base.adjoint(f::FourierSeries)
    fc = FourierSeries(map(adjoint, coefficients(f)), period(f))
    K = nummodes(f)
    for k = 1:K
        fc[k], fc[-k] = fc[-k], fc[k]
    end
    return fc
end
Base.transpose(f::FourierSeries) = FourierSeries(map(transpose, coefficients(f)), period(f))

LinearAlgebra.tr(f::FourierSeries) = FourierSeries(map(tr, coefficients(f)), period(f))

Base.real(f::FourierSeries) = rmul!(f+conj(f), 1//2)
Base.imag(f::FourierSeries) = rmul!(f-conj(f), 1//(2*im))

# Arithmetic (in place / mutating methods)
function Base.copy!(fdst::FourierSeries, fsrc::FourierSeries)
    period(fdst) == period(fsrc) || throw(DomainMismatch())
    setnummodes!(fdst, nummodes(fsrc))
    if eltype(fdst) <: Number
        copy!(coefficients(fdst), coefficients(fsrc))
    else
        for k in eachindex(fdst)
            copy!(fdst[k], fsrc[k])
        end
    end
    return fdst
end
function LinearAlgebra.rmul!(f::FourierSeries, α::Number)
    if eltype(f) <: Number
        rmul!(coefficients(f), α)
    else
        for k in eachindex(f)
            rmul!(f[k], α)
        end
    end
    return f
end
function LinearAlgebra.lmul!(α::Number, f::FourierSeries)
    if eltype(f) <: Number
        lmul!(α, coefficients(f))
    else
        for k in eachindex(f)
            lmul!(α, f[k])
        end
    end
    return f
end

LinearAlgebra.axpy!(α::Number, fx::FourierSeries, fy::FourierSeries) =
    truncadd!(fy, fx, α)
LinearAlgebra.axpby!(α::Number, fx::FourierSeries, β::Number, fy::FourierSeries) =
    truncadd!(fy, fx, α, β)
LinearAlgebra.mul!(fy::FourierSeries, s::Number, fx::FourierSeries, α = true, β = false) =
    truncmul!(fy, s, fx, α, β)
LinearAlgebra.mul!(fy::FourierSeries, fx::FourierSeries, s::Number, α = true, β = false) =
    truncmul!(fy, fx, s, α, β)
LinearAlgebra.mul!(f::FourierSeries, f1::FourierSeries, f2::FourierSeries,
                    α = true, β = false) = truncmul!(f, f1, f2, α, β)

function truncadd!(fy::FourierSeries, fx::FourierSeries, α = true, β = true;
                    Kmax::Integer = max(iszero(β) ? 0 : nummodes(fy), nummodes(fx)),
                    tol::Real = 0)
    period(fx) == period(fy) || throw(DomainMismatch())
    Kx = nummodes(fx)
    Ky = min(Kmax, iszero(β) ? Kx : max(Kx, nummodes(fy)))
    setnummodes!(fy, Ky)
    if eltype(fy) <: Number
        if Ky > Kx
            LinearAlgebra.axpby!(α, coefficients(fx), β, view(coefficients(fy), 1:2Kx+1))
            lmul!(β, view(coefficients(fy), (2Kx+2):(2Ky+1)))
        else
            LinearAlgebra.axpby!(α, view(coefficients(fx), 1:2Ky+1), β, coefficients(fy))
        end
    else
        Threads.@threads for k in -Ky:Ky
            if abs(k) <= Kx
                axpby!(α, fx[k], β, fy[k])
            elseif !isone(β)
                lmul!(β, fy[k])
            end
        end
    end
    return iszero(tol) ? fy : truncate!(fy; tol = tol)
end

truncmul!(fdst::FourierSeries, α₁::Number, fsrc::FourierSeries, α₂ = true, β = false;
            kwargs...) = truncadd!(fdst, fsrc, α₁ * α₂, β; kwargs...)
truncmul!(fdst::FourierSeries, fsrc::FourierSeries, α₁::Number, α₂ = true, β = false;
            kwargs...) = truncadd!(fdst, fsrc, α₁ * α₂, β; kwargs...)

function truncmul!(f::FourierSeries, f1::FourierSeries, f2::FourierSeries,
                    α = true, β = false;
                    Kmax::Integer = max(iszero(β) ? 0 : nummodes(f),
                                                        nummodes(f1)+nummodes(f2)),
                    tol::Real = 0)
    period(f) == period(f1) == period(f2) || throw(DomainMismatch())
    K1 = nummodes(f1)
    K2 = nummodes(f2)
    K = min(Kmax, max(iszero(β) ? 0 : nummodes(f), K1+K2))
    setnummodes!(f, K)
    Threads.@threads for k = -K:K
        fk = f[k]
        if fk isa AbstractArray
            T = eltype(fk)
            if β == 0
                fill!(fk, T(β))
            elseif β != 1
                rmul!(fk, T(β))
            end
            for k1 = max(-K1, k-K2):min(K1, k+K2)
                k2 = k - k1
                fk1 = f1[k1]
                fk2 = f2[k2]
                if fk1 isa AbstractArray && fk2 isa AbstractArray
                    mul!(fk, fk1, fk2, T(α), T(true))
                elseif fk1 isa AbstractArray
                    axpy!(T(fk2*α), fk1, fk)
                elseif fk2 isa AbstractArray
                    axpy!(T(fk1*α), fk2, fk)
                else
                    @warn "unexpected branch"
                    fk .= fk + (fk1*fk2*α)
                end
            end
        else
            fk *= β
            for k1 = max(-K1, k-K2):min(K1, k+K2)
                k2 = k - k1
                fk1 = f1[k1]
                fk2 = f2[k2]
                fk += fk1*fk2*α
            end
            f[k] = fk
        end
    end
    return iszero(tol) ? f : truncate!(f; tol = tol)
end

# Inner product and norm
function localdot(f1::FourierSeries, f2::FourierSeries)
    period(f1) == period(f2) || throw(DomainMismatch())
    K1 = nummodes(f1)
    K2 = nummodes(f2)
    K = K1+K2
    f = FourierSeries(sizehint!([zero(dot(f1[0], f2[0]))], 2*K+1), period(f1))
    for k = -K:K
        fk = f[k]
        for k1 = max(-K1, -K2-k):min(K1, K2-k)
            k2 = k + k1
            fk += dot(f1[k1], f2[k2])
        end
        f[k] = fk
    end
    return f
end

function LinearAlgebra.dot(f1::FourierSeries, f2::FourierSeries)
    domain(f1) == domain(f2) || throw(DomainMismatch())
    K = min(nummodes(f1), nummodes(f2))
    return sum(k->dot(f1[k], f2[k]), -K:K)
end
LinearAlgebra.norm(f::FourierSeries) = norm(norm(fk) for fk in coefficients(f))

# Differentiate and integrate
function differentiate(f::FourierSeries)
    df = 0im * f
    ω = (2*pi)/period(f)
    K = nummodes(f)
    for k = 1:K
        ν = im*ω*k
        df[-k] = -ν*f[-k]
        df[+k] = +ν*f[+k]
    end
    return df
end
function integrate(f::FourierSeries, (a,b) = domain(F))
    p = period(f)
    if b-a == p
        return f[0]*p
    else
        error("not yet implemented")
    end
end

fit(f, ::Type{FourierSeries}, (a,b)::Tuple{Real,Real}; kwargs...) =
    fit(f, FourierSeries, b-a; kwargs...)

function fit(f, ::Type{FourierSeries}, period = 1; Kmax = 10, numpoints = 2*Kmax+1)
    K = Kmax
    x = (0:(numpoints-1)) * (period / numpoints)
    eox = exp.((im*2*pi/period) .* x)
    coeffs = similar(eox, (numpoints, 2*Kmax+1))
    coeffs[:, 1] .= 1
    for k = 1:Kmax
        coeffs[:, 2*k] .= view(coeffs, :, k == 1 ? 1 : 2*k-2) .* eox
        coeffs[:, 2*k+1] .= view(coeffs, :, 2*k-1) .* conj.(eox)
    end
    fx = f.(x)
    fs = FourierSeries(pinv(coeffs)*fx, period)
    if all(isreal, fx)
        return real(fs)
    else
        return fs
    end
end

# Inverse and square root
# ?


using DataDrivenControl
using FlightSims
using UnPack
using Plots
using LinearAlgebra
using Transducers
using LaTeXStrings
using DiffEqCallbacks

struct LinearSystem_ZOH_Gain
    linear::LinearSystem
    controller::LinearIRL
end

function FlightSims.State(env::LinearSystem_ZOH_Gain)
    @unpack linear = env
    State(linear)
end

function FlightSims.Dynamics!(env::LinearSystem_ZOH_Gain)
    @unpack linear, controller = env
    @Loggable function dynamics!(dx, x, w, t)
        u = optimal_input(controller, x, w)
        @onlylog param = w
        @onlylog i = controller.i
        @nested_log Dynamics!(linear)(dx, x, w, t; u=u)
    end
end

# See [1, Example 2. Continuous-Time Optimal Adaptive Control Using IRL]
# Refs
# [1] “Reinforcement Learning and Feedback Control: Using Natural Decision Methods to Design Optimal Adaptive Controllers,” IEEE Control Syst., vol. 32, no. 6, pp. 76–105, Dec. 2012, doi: 10.1109/MCS.2012.2214134.
function main()
        n, m = 2, 1
        A = [   -10  1;
             -0.002 -2]
        B = [0 2]'
        Q = Matrix(I, n, n)
        R = Matrix(I, m, m)
        cost = DataDrivenControl.QuadraticCost(Q, R)
        linear = LinearSystem(A, B)
        controller = DataDrivenControl.LinearIRL(Q, R, B; T=0.04)
        env = LinearSystem_ZOH_Gain(linear, controller)
        x0 = State(env)(rand(2))
        # TODO: add callback to update param `w`
        # simulation
        P_true = [0.0500 0.0039;
                  0.0039 0.2085]  # true solution
        w_true = [P_true[1, 1], P_true[2, 1]+P_true[1, 2], P_true[2, 2]]
        scale = 0.1
        w0 = w_true + scale*randn(3)  # perturbed initial guess
        tf = 5.0
        simulator = Simulator(x0, Dynamics!(env), w0;
                              tf=tf,
                             )
        Δt = 0.01
        # TODO: data buffer is a duplicate of saving callback (see `df`)
        # data buffer
        buffer = DataDrivenControl.DataBuffer()
        function update!(integrator)
            t = integrator.t
            x = integrator.u  # convention of DifferentialEquations.jl
            w = integrator.p  # convention of DifferentialEquations.jl
            u = optimal_input(controller, x, w)  # TODO: not to be duplicate of control input in dynamics for stable coding
            push!(buffer, controller, cost;
                  t=t, x=copy(x), u=copy(u), w=copy(w))
            eps = 1e-1
            sc = DistanceStopCondition(eps)
            # value_iteration!(controller, buffer, w; sc=sc)
            policy_iteration!(controller, buffer, w; sc=sc)
        end
        cb_irl = PeriodicCallback(update!, controller.T; initial_affect=true)  # stack initial data
        df = solve(simulator;
                   callback=cb_irl,
                   savestep=Δt,
                  )
        # plot
        ts = df.time
        xs = df.sol |> Map(datum -> datum.state) |> collect
        us = df.sol |> Map(datum -> datum.input) |> collect
        ws = df.sol |> Map(datum -> datum.param) |> collect
        is = df.sol |> Map(datum -> datum.i) |> collect
        fig_x = plot(ts, hcat(xs...)'; label=[L"x_{1}" L"x_{2}"], legend=:outerbottomright)
        fig_u = plot(ts, hcat(us...)'; label=L"u", legend=:outerbottomright)
        fig_w = plot(ts, hcat(ws...)'; label=[L"w_{1}" L"w_{2}" L"w_{3}"], legend=:outerbottomright)
        fig_iter = plot(ts, hcat(is...)'; label="iter")
        ws_true = ts |> Map(t -> w_true) |> collect
        plot!(fig_w, ts, hcat(ws_true...)'; color=:black, label=nothing)
        fig = plot(fig_x, fig_u, fig_w, fig_iter)
end


"""Script to compute the memory capacity of a AssemblyCalculus.BrainArea

Script projects random stimuli into the brain area and compares assembly that
form in the region. This is an attempt to answer the question, how many
suitably unique assemblies can a region generate.
"""

using AssemblyCalculus
using AssemblyCalculus: overlap, forget_parents!
using Graphs
using Plots

NUM_BRAINS = 4
MAX_ASSEMBLIES = 500
MAX_ITERS = 100
TOL = 3
N = 10_000
K = 100
EDGE_PROBS = [0.005f0, 0.01f0, 0.05f0, 0.1f0]
β = 0.1f0

assemblies = [Assembly{Float32}[] for i in 1:NUM_BRAINS]
max_distances = [Int[] for i in 1:NUM_BRAINS]
iters_taken = [Int[] for i in 1:NUM_BRAINS]

for brain_num in 1:NUM_BRAINS
    cutoff = (2EDGE_PROBS[brain_num] / π) ^ 0.5
    g, _ = euclidean_graph(N, 2, cutoff=cutoff)
    println("Mean degree: $(length(edges(g)) / N)")
    brain = Brain(DiGraph(g), [N], [K], [β])
    for assem_num in 1:MAX_ASSEMBLIES
        stim = rand_stim(brain[1])
        assems, _, dists = simulate!([stim], tol=TOL, max_iters=MAX_ITERS)
        a = assems[1]
        forget_parents!(a)
        iters = size(dists)[2]
        push!(iters_taken[brain_num], iters)
        if assem_num != 1
            if dists[iters] <= TOL
                common_neurons = map(x -> overlap(a, x), assemblies[brain_num])
                push!(max_distances[brain_num], maximum(common_neurons))
            else
                push!(max_distances[brain_num], -1)
            end
        end
        push!(assemblies[brain_num], a)
    end
end

p1 = plot(legend=:bottomright)
colors = [:teal, :pink, :green, :orange]
for (i, ep) in enumerate(EDGE_PROBS)
    max_ds = max_distances[i]
    mean_max_iters = cumsum(max_ds) ./ collect(1:length(max_ds))
    plot!(p1, mean_max_iters, label="p = $ep", linewidth=3, color=colors[i])
end
xlabel!(p1, "Number of Assemblies Created")
ylabel!(p1, "Average max overlap \n with previous assemblies")

p2 = plot(legend=:bottomright)
for (i, ep) in enumerate(EDGE_PROBS)
    plot!(p2, max_distances[i], label="p = $ep", linewidth=3, color=colors[i])
end
xlabel!(p2, "Number of Assemblies Created")
ylabel!(p2, "Max overlap with \n previous assemblies")

l = @layout [a b]
plot(p1, p2, layout=l, size=(1200, 400))
suptitle!("Assembly capacity for different edge probabilities \n(Random Geometric topology)")


import RCall  # slow

"""
Auto-Correlation function by RCall.
"""
function acf(X::AbstractVector{T}, lagmax::Int) where {T<:Real}
    res = RCall.rcopy(RCall.rcall(:acf, X, lagmax, plot=false, na_action=:na_pass))
    return res[:acf][2:end]
end

"""
Partial Auto-Correlation function by RCall.
"""
function pacf(X::AbstractVector{T}, lagmax::Int) where {T<:Real}
    res = RCall.rcopy(RCall.rcall(:pacf, X, lagmax, plot=false, na_action=:na_pass))
    return res[:acf][:]
end

"""
Auto-Correlation function of increment process.
"""
function acf_incr(X::AbstractVector{T}, dlags::Union{Int, AbstractVector{Int}}, lagmax::Int; method::Symbol=:acf) where {T<:Real}
    # for single value of dlag: convert to a list
    if typeof(dlags) <: Integer
        dlags = [dlags]
    end

    A = if method==:acf
        [acf((X[l+1:end]-X[1:end-l])[100:end-100], lagmax) for l in dlags]
    elseif method==:pacf
        [pacf((X[l+1:end]-X[1:end-l])[100:end-100], lagmax) for l in dlags]
    else
        error("Invalid method: $(method)")
    end

    # A = []  # output of acf
    # for l in dlags
    #     dX = X[l+1:end]-X[1:end-l]
    #     if method==:acf
    #         push!(A, acf(dX, lagmax))
    #         # push!(A, [cor(dX[t+1:end], dX[1:end-t]) for t in tidx])
    #     elseif method==:pacf
    #         push!(A, pacf(dX, lagmax))
    #     else
    #         error("Invalid method: $(method)")
    #     end
    # end
    return hcat(A...)
end