using StatsPlots
using Clustering
using Distances
using Distributions
using DataFrames

import Base: ∈
import Base: ∩

include("../src/utilities/utils.jl")

# a = rand(Uniform(0,.8), 2, 20)
# b = rand(Uniform(1.2,2), 2, 20)
# c = a + cat(1.2ones(1,20),zeros(1,20), dims=1)
# X = cat(a,b,c, dims=2)

# scatt = scatter(X[1,:], X[2,:], label=false)

# D_a = pairwise(Euclidean(), X, X)
# SL = hclust(D_a, linkage=:single)
# CL = hclust(D_a, linkage=:complete)
# AV = hclust(D_a, linkage=:average)
# plot(SL)
# plot(CL)
# plot(AV)

# propertynames(Hclust)

# propertynames(SL)
# SL.merges
# SL.heights
# SL.order
# SL.linkage
# SL.height
# SL.labels
# SL.merge
# SL.method

# plot(SL.heights)
# SL.heights[56]
# SL.heights[57]
# SL.heights[58]
# SL.heights[59]

# kmeans_ = kmeans(reshape(SL.heights, (1,59)), 2)
# kmeans_ = kmeans(reshape(CL.heights, (1,59)), 2)
# propertynames(kmeans_)
# @show kmeans_

# cutree(SL, k=2)

# D_SL = pairwise(Euclidean(), SL.heights)
# dbscan(reshape(SL.heights, (1,59)), 0.8maximum(SL.heights))
# dbscan(reshape(SL.heights, (1,59)), 0.05maximum(SL.heights))
# dbscan(reshape(SL.heights, (1,59)), 0.2maximum(SL.heights))

# dbscan_ = dbscan(reshape(CL.heights, (1,59)), 0.1maximum(CL.heights))
# dbscan_[1]
# dbscan_[2]
# dbscan(reshape(CL.heights, (1,59)), 0.08maximum(CL.heights))
# dbscan(reshape(CL.heights, (1,59)), 0.2maximum(CL.heights))

# cutree(CL, h=CL.heights[58])
# x = Interval
# x = Interval(1,2)

# struct OrderedPair
#     x::Real
#     y::Real
#     OrderedPair(x,y) = x > y ? error("out of order") : new(x,y)
# end
# OrderedPair(3,1)

# I = Interval(0.2, 4.0)
# -1 ∈ I
# 1 ∈ I
# J = Interval(-1., 3.6)

# # function intersection(I::Interval, J::Interval)
# #     min = maximum([I.a, J.a])
# #     max = minimum([I.b, J.b])

# #     return Interval(min, max)
# # end

# ∩(I,J)

begin
using StatsPlots
using Clustering
using Distances
using Distributions
using DataFrames

import Base: ∈
import Base: ∩

include("../src/utilities/utils.jl")
Data = noisy_circle_data(500, 0.03)

function proj_1d(X)
    min = minimum(X[1, :])
    output = X[1, :] .- min

    return output
end

struct Interval{T<:Real}
    a::T
    b::T
    # Interval(a,b) = a > b ? error("out of order") : new(a,b)
end

mutable struct MapperInformations
    data_number
    original_coordinate
    original_subset
    clustered_subset
    filtered_coordinate
    interval_subset
end

function ∈(x::T, interval::Interval) where T<:Real
    if interval.a ≤ x ≤ interval.b
        return true
    else
        return false
    end
end

function ∩(I::Interval, J::Interval)
    min = maximum([I.a, J.a])
    max = minimum([I.b, J.b])

    return Interval(min, max)
end

function filtered_coordinate(x, p)
    x.filtered_coordinate = p
end

#Y_ = MapperInformations(:None, :None, :None, :None, :None)
end

function mapper(X, filter, interval, overlab)

    original_X = []
    for i in 1:size(X,2)
        push!(original_X, X[:,i])
    end
    # findfirst(==(Info[500].original_coordinate), original_X)

    Info = [MapperInformations(i, X[:, i], Int64[], Tuple{Int64, Int64}[], :None, :None) for i ∈ 1:length(X[1,:])]

    filtered_X = filter(X)
    filtered_coordinate.(Info, filtered_X)
    Info

    min_X = minimum(filtered_X)
    max_X = maximum(filtered_X)
    filter_range = Interval(min_X, max_X)

    length_ = (max_X - min_X)/interval
    ϵ = overlab*length_

    Intervals = [∩(Interval(min_X +(i-1)*length_ - ϵ, min_X +i*length_ + ϵ), filter_range) for i ∈ 1:interval]

    X_subset = interval_test(Info, Intervals)

    Info

    X_subset_clustering = domain_clustering(Info, X_subset)

    # function myshowall(io, x, limit = false)
    #     println(io, summary(x), ":")
    #     Base.print_matrix(IOContext(io, :limit => limit), x)
    # end
    # myshowall(Base.stdout, Info, false)

    # X_subset_clustering[1]
    # X_subset_clustering[2]
    # X_subset_clustering[3]
    # X_subset_clustering[4]
    # X_subset_clustering[5]

    link = []

    for i ∈ 1:length(X_subset)-1
        for j ∈ 1:length(X_subset_clustering[i])
            for k ∈ 1:length(X_subset_clustering[i+1])
                if length(∩(X_subset_clustering[i][j], X_subset_clustering[i+1][k])) > 0
                    push!(link, [[i,j], [i+1,k]])
                end
            end
        end
    end

    center_ = coor_mean(X_subset_clustering, Info)

    center_number = 0
    for i ∈ 1:length(X_subset)
        for j ∈ 1:length(X_subset_clustering[i])
            center_number +=1
        end
    end

    center_matrix = zeros(2, center_number)
    k_temp = 1
    for i ∈ 1:length(X_subset)
        for j ∈ 1:length(X_subset_clustering[i])
            center_matrix[:,k_temp] = center_[i][j]
            k_temp += 1
        end
    end

    zcolor_temp = filter(cat(Data, center_matrix, dims=2))

    zcolor = [[] for i ∈ 1:length(X_subset)]
    k_temp = 1
    for i ∈ 1:length(X_subset)
        for j ∈ 1:length(X_subset_clustering[i])
            push!(zcolor[i], zcolor_temp[size(Data)[2]+k_temp])
            k_temp +=1
        end
    end

    return link, X_subset_clustering, Info, center_, zcolor
end

function mapper_plot(link, X_subset_clustering, center_, zcolor, xlim, ylim)

    linked_center = zeros(length(link), 4)
    for i ∈ 1:length(link)
        linked_center[i, 1:2] = center_[link[i][1][1]][link[i][1][2]]
        linked_center[i, 3:4] = center_[link[i][2][1]][link[i][2][2]]
    end
    linked_center

    p = plot([linked_center[1,1], linked_center[1,3]], [linked_center[1,2], linked_center[1,4]], color=:black, xlim=xlim, ylim=ylim, label="")
    for i ∈ 2:length(link)
        plot!([linked_center[i,1], linked_center[i,3]], [linked_center[i,2], linked_center[i,4]], color=:black, label="")
    end

    for i ∈ 1:length(X_subset_clustering)
        for j ∈ 1:length(X_subset_clustering[i])
            scatter!([center_[i][j][1]], [center_[i][j][2]], markersize=25*(452.106/181)*sqrt(length(X_subset_clustering[i][j])/500), zcolor=zcolor[i][j], label="", color=:jet, clim=(0,Inf))
            length_temp = length(X_subset_clustering[i][j])
            annotate!([center_[i][j][1]], [center_[i][j][2]], "$length_temp")
        end
    end

    return p
end

d1 = scatter(Data[1,:], Data[2,:], label="")
d2 = scatter(Data[1,:], Data[2,:], zcolor=proj_1d(Data), color=:jet, label="")

link, X_subset_clustering, Info, center_, zcolor = mapper(Data, proj_1d, 5, 0.2)
p = mapper_plot(link, X_subset_clustering, center_, zcolor, (-1.5, 1.5), (-1.5,1.5))

plot(d1, d2, p, size=(1400,800))

# plot(d, p, size=(1000,400))
# include("../src/utilities/tDATAsets.jl")

# Data = rand(Sphere(2), 500)
# Data = cat(rand(Infty(), 500), rand(Sphere(2), 500).+1, dims=2)

# link, X_subset_clustering, Info = mapper(Data, proj_1d, 5, 0.3)
# p = mapper_plot(link, X_subset_clustering, Info, (-1.5, 1.5), (-1.5,1.5))

# plot(d, p, size=(1000,400))

# d = scatter(Data[1,:], Data[2,:])
# link, X_subset_clustering, Info = mapper(Data, proj_1d, 10, 0.2)
# p = mapper_plot(link, X_subset_clustering, Info)

# plot(d, p, size=(1000,400))


using CSV

""" Game of Life iteration for the periodic boundary condition case"""
function game_it(A)
    M, N = size(A)
    j_l = j_r = 1
    i_u = i_d = 1
    B = zeros(Int64, M, N)
    for j in 1:N
        j == 1 ? j_l = N : j_l = j-1
        j == N ? j_r = 1 : j_r = j+1

        for i in 1:M
            i == 1 ? i_u = M : i_u = i-1
            i == M ? i_d = 1 : i_d = i+1

            sum_neigh = A[i_u,j_l] + A[i_u,j] + A[i_u,j_r] + A[i,j_l] + A[i,j_r] + A[i_d,j_l] + A[i_d,j] + A[i_d,j_r]
            if sum_neigh == 3 || (sum_neigh == 2 && A[i,j] == 1)
                B[i,j] = 1
            end
        end
    end
    return B
end

""" Game of Life iteration for the zero boundary condition case"""
function game_it_zero_bdry(A)
    M, N = size(A)
    B = zeros(Int64, M, N)
    for j in 2:N-1
        for i in 2:M-1
            sum_neigh = A[i-1,j-1] + A[i-1,j] + A[i-1,j+1] + A[i,j-1] + A[i,j+1] + A[i+1,j-1] + A[i+1,j] + A[i+1,j+1]
            if sum_neigh == 3 || (sum_neigh == 2 && A[i,j] == 1)
                B[i,j] = 1
            end
        end
    end
    return B
end

""" An auxiliary function for plotting the status of the game of life """
function plot_it(A)
    current_plot = heatmap(A,aspect_ratio=:equal,size = (400,400),grid=:none,axis = false,legend=nothing,color=:greys_r)

    # Plot a grid on top
    M, N = size(A)
    x1 = hcat([[0,N] for i in 0:M]...).+0.5
    y1 = hcat([[i,i] for i in 0:M]...).+0.5
    x2 = hcat([[i,i] for i in 0:N]...).+0.5
    y2 = hcat([[0,M] for i in 0:N]...).+0.5
    plot!(x1, y1,color=:lightgrey,linewidth=10/max(M,N))
    plot!(x2, y2,color=:lightgrey,linewidth=10/max(M,N))
    return current_plot
end

function gen_examples()

    # Check if the directory exists
    if !isdir("game_examples")
        mkdir("game_examples")
    end

    # Intro
    N = 35
    A = zeros(Int64, N, N)
    c = N ÷ 2+1
    A[[c-3,c+3],c-1:c+1] .= 1
    A[[c-2,c-1,c+1,c+2],[c-1,c+1]] .= 1
    CSV.write("game_examples/intro.csv", DataFrame(A))

    # Block
    A = zeros(Int64, 4, 4)
    A[2:3,2:3] .= 1
    CSV.write("game_examples/still_block.csv", DataFrame(A))

    # Hive
    A = zeros(Int64, 5, 6)
    A[[2,4],3:4] .= 1
    A[3,[2,5]] .= 1
    CSV.write("game_examples/still_hive.csv", DataFrame(A))

    # Tub
    A = zeros(Int64, 5, 5)
    A[[2,4],3] .= 1
    A[3,[2,4]] .= 1
    CSV.write("game_examples/still_tub.csv", DataFrame(A))

    # Blinker
    A = zeros(Int64, 5, 5)
    A[2:4,3] .= 1
    CSV.write("game_examples/osc_blinker.csv", DataFrame(A))

    # Beacon
    A = zeros(Int64, 6, 6)
    A[2:3,2:3] .= 1
    A[4:5,4:5] .= 1
    CSV.write("game_examples/osc_beacon.csv", DataFrame(A))

    # Ducks

    A = zeros(Int64, 7, 7)
    A[2,3:4] .= 1
    A[3,2:5] .= 1
    A[4,2:3] .= 1; A[4,5:6] .= 1;
    A[5,4:5] .= 1;

    function rotate(A)
        m, n = size(A)
        A = [A[n+1-i,j] for j in 1:m, i in 1:n]
        return A
    end

    N = 18
    B = zeros(Int64,N,N)
    B[1:7,1:7] .= A

    A = rotate(A); B[1:7,N-6:N] .= A[:,:];
    A = rotate(A); B[N-6:N,N-6:N] .= A[:,:]
    A = rotate(A); B[N-6:N,1:7] .= A[:,:]

    C = zeros(Int64, N+4, N+4)
    C[3:N+2,3:N+2] = B
    CSV.write("game_examples/final.csv", DataFrame(C))
end


#!/usr/bin/julia

# Copyright (C) 2021 Matheus Fernandes Bigolin <mfrdrbigolin@disroot.org>
# SPDX-License-Identifier: MIT

"Day Twelve, Passage Pathing."

#= The recursive solution works on a good time, therefore no need to use
 = dynamic programming, at least now.  Maybe I will return and refactor this
 = solution someday.
 =#

include("Utils.jl")

using DelimitedFiles
using .Utils

"Create adjacency matrix from `paths`."
function create_adjacency_matrix(paths :: Matrix{String})
    vertices = unique(paths)
    vert_len = length(vertices)

    vert_assoc = Dict(zip(vertices, 1:vert_len))
    path_assoc = zip(paths[:,1], paths[:,2])

    adj_mat = BitArray{2}(zeros(Bool, (vert_len, vert_len)))

    for path in path_assoc
        if "start" in path
            to = path[findfirst(path .!= "start")]
            adj_mat[vert_assoc["start"], vert_assoc[to]] = 1
        elseif "end" in path
            from = path[findfirst(path .!= "end")]
            adj_mat[vert_assoc[from], vert_assoc["end"]] = 1
        else
            from, to = path

            adj_mat[vert_assoc[from], vert_assoc[to]] = 1
            adj_mat[vert_assoc[to], vert_assoc[from]] = 1
        end
    end

    return adj_mat
end

"""Find the number of possible `paths` in the cavern system.  If `can_twice`,
allow at most two visits to a single small cave; if not, only allow one visit
per small cave."""
function solve(paths :: Matrix{String}, can_twice = false)
    vertices = unique(paths)

    vert_assoc = Dict(zip(vertices, 1:length(vertices)))
    adj_mat = create_adjacency_matrix(paths)

    small_cave = Array{Int}([])

    for v in vertices
        if all(islowercase.(collect(v))) && v ∉ ["start", "end"]
            push!(small_cave, vert_assoc[v])
        end
    end

    """Recurse through the paths, recording the small caves `visited` and the
    `start` point. For `can_twice`, vide supra."""
    function rec_path(can_twice = false, visited = [], start = vert_assoc["start"])
        if start == vert_assoc["end"]
            return 1
        elseif start in visited
            if !can_twice
                return 0
            else
                can_twice = false
            end
        end

        new_visited = visited

        if start in small_cave
            new_visited = copy(visited)
            push!(new_visited, start)
        end

        starts = findall(adj_mat[start,:])

        return sum([rec_path(can_twice, new_visited, st) for st in starts])
    end

    return rec_path(can_twice)
end

function main()
    usage_and_exit(length(ARGS) != 1)

    data = Matrix{String}(readdlm(first(ARGS), '-'))

    println(solve(data))
    println(solve(data, true))
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


using Catlab.CategoricalAlgebra, Catlab.Graphs, Catlab.Present, Catlab.Graphics, Catlab.Theories
using Catlab.CategoricalAlgebra.FinCats: FinCatGraphEq

using Distributions: Exponential, Geometric, cdf
using Random: randsubseq
using StatsBase: sample

using Plots

# include("utils.jl")

# Bernoulli sample of vector `m` (all same probability)
function sample_matches(m::AbstractVector{T}, r::AbstractFloat, Δt) where {T}
    p = cdf(Exponential(), r * Δt)
    randsubseq(m, p)
end

# update matches; return nothing if invalid
function postcompose_partial(kg::ACSetTransformation, kh::ACSetTransformation, m::ACSetTransformation)
    d = Dict()
    for (k,vs) in pairs(components(m))
      vs_ = Int[]
      for v in collect(vs)
        kv = findfirst(==(v), collect(kg[k]))
        if isnothing(kv)
          return nothing
        else
          push!(vs_, kh[k](kv))
        end
      end
      d[k] = vs_
    end
    return ACSetTransformation(dom(m), codom(kh); d...)
  end

  # stuff to store stuff
  struct Rule
      L::ACSetTransformation
      R::ACSetTransformation
  end

  mutable struct MatchedRule
      rule::Rule
      match::Union{Nothing, ACSetTransformation}
  end

# schema
@present ThSIR(FreeSchema) begin
    (S,I,R)::Ob
end

@acset_type SIR(ThSIR)

# infection rules
I = @acset SIR begin I=1 end
I2 = @acset SIR begin I=2 end
SI = @acset SIR begin S=1; I=1 end
L_infect =  ACSetTransformation(I, SI; I=[1]) # fn from dom I to codom SI, I is the injective function mapping stuff in dom to codom
R_infect = ACSetTransformation(I, I2; I=[1])

# recovery rules
R = @acset SIR begin R=1 end
L_recovery = ACSetTransformation(SIR(), I)
R_recovery = ACSetTransformation(SIR(), R)

# structs
rule_inf = Rule(L_infect, R_infect)
rule_rec = Rule(L_recovery, R_recovery)

N = 500
I0 = 10
S0 = N - I0

Δt = 0.1
tmax = 100
steps = Int(tmax/Δt)
γ = 1/10 # recovery rate
R0 = 2.5
β = R0 * γ # R0 for corresponding ODEs

state = @acset SIR begin S=S0; I=I0; R=N-S0-I0 end

# simulation loop
out = fill(-1, (steps, 3))

for t = 1:steps
    # possible events
    infections = homomorphisms(SI, state)
    recoveries = homomorphisms(I, state)
    # sample occurances
    infections = sample_matches(infections, β/N, Δt)
    recoveries = sample_matches(recoveries, γ, Δt)
    # apply events
    events = MatchedRule[[MatchedRule(rule_inf, m) for m in infections]..., [MatchedRule(rule_rec, m) for m in recoveries]...]
    while length(events) > 0
        ev = pop!(events)
        # apply event
        _, kg, _, kh = rewrite_match_maps(ev.rule.L, ev.rule.R, ev.match)
        state = codom(kh)
        # update the remaining matches post rewrite
        for j in 1:length(events)
            events[j].match = postcompose_partial(kg, kh, events[j].match)
        end
        events = filter((e) -> !isnothing(e.match), events)
    end
    # write output
    out[t, :] = [nparts(state, x) for x in [:S, :I, :R]]
end

plot(
    (1:steps) * Δt,
    out,
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number"
)

# # --------------------------------------------------------
# match = homomorphism(SI, state)
# match = homomorphisms(SI, state)

# match = homomorphism(I, state)
# match = homomorphisms(I, state)

# # rewrite_match(L_recovery, R_recovery, match[1])

# # sample_matches(match, γ, Δt)

# state

# rewrite(L_recovery, R_recovery, state)

# # not sure how to update morphisms after the first one was applied.
# nparts(state, :R)
# state = apply_matches(state,match,L_recovery,R_recovery)
# nparts(state, :R)


# ---
# title: 576. Out of Boundary Paths
# id: problem576
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming, Depth-first Search
# link: <https://leetcode.com/problems/out-of-boundary-paths/description/>
# hidden: true
# ---
#
# There is an **m** by **n** grid with a ball. Given the start coordinate
# **(i,j)** of the ball, you can move the ball to **adjacent** cell or cross the
# grid boundary in four directions (up, down, left, right). However, you can
# **at most** move **N** times. Find out the number of paths to move the ball
# out of grid boundary. The answer may be very large, return it after mod 109 \+
# 7.
#
#
#
# **Example 1:**
#
#
#
#     Input: m = 2, n = 2, N = 2, i = 0, j = 0
#     Output: 6
#     Explanation:
#     ![](https://assets.leetcode.com/uploads/2018/10/13/out_of_boundary_paths_1.png)
#
#
# **Example 2:**
#
#
#
#     Input: m = 1, n = 3, N = 3, i = 0, j = 1
#     Output: 12
#     Explanation:
#     ![](https://assets.leetcode.com/uploads/2018/10/12/out_of_boundary_paths_2.png)
#
#
#
#
# **Note:**
#
#   1. Once you move the ball out of boundary, you cannot move it back.
#   2. The length and height of the grid is in range [1,50].
#   3. N is in range [0,50].
#
#
## @lc code=start
using LeetCode

function find_paths(m, n, N, i, j)
    dp = zeros(Int, m, n)

    for _ in 1:N
        t = copy(dp)
        inds = CartesianIndices(t)
        for x in inds
            v = 0
            for dir in ((0, -1), (0, 1), (1, 0), (-1, 0))
                p = x + CartesianIndex(dir)
                v += p ∈ inds ? t[p] : 1
            end
            dp[x] = v
        end
    end
    return dp[i, j] % (10^9 + 7)
end
## @lc code=end


"""
    Cartpole w/ internal friction
"""
struct Friction end
struct Frictionless end

struct Cartpole{T,X} <: Model{T}
    # dimensions
    nq::Int # generalized coordinates
    nu::Int # controls
    nw::Int # parameters
	nc::Int # contact points

    mc::T     # mass of the cart in kg
    mp::T     # mass of the pole (point mass at the end) in kg
    l::T      # length of the pole in m
    g::T      # gravity m/s^2

    friction::Vector{T} # friction coefficients for slider and arm joints
end

function kinematics(model::Cartpole, q)
    [q[1] + model.l * sin(q[2]); -model.l * cos(q[2])]
end

lagrangian(model::Cartpole, q, q̇) = 0.0

function M_func(model::Cartpole, x)
    H = [model.mc + model.mp model.mp * model.l * cos(x[2]);
		 model.mp * model.l * cos(x[2]) model.mp * model.l^2.0]
    return H
end

function B_func(model::Cartpole, x)
    [1.0; 0.0]
end

function P_func(model::Cartpole, x)
    [1.0 0.0;
     0.0 1.0]
end

function C_func(model::Cartpole, q, q̇)
    C = [0.0 -1.0 * model.mp * q̇[2] * model.l * sin(q[2]);
	 	 0.0 0.0]
    G = [0.0,
		 model.mp * model.g * model.l * sin(q[2])]
    return -C * q̇ + G
end

function dynamics(model::Cartpole{T,Friction}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T
	# evalutate at midpoint
    qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

    D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1)
    D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2)

    d = 0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 # variational integrator (midpoint)

	return d + B_func(model, qm2) * u1[1] + transpose(P_func(model, q2)) * λ1
end

function dynamics(model::Cartpole{T,Frictionless}, mass_matrix, dynamics_bias, h, q0, q1, u1, w1, λ1, q2) where T
	# evalutate at midpoint
    qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

    D1L1, D2L1 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm1, vm1)
    D1L2, D2L2 = lagrangian_derivatives(mass_matrix, dynamics_bias, qm2, vm2)

    d = 0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2 # variational integrator (midpoint)

	return d + B_func(model, qm2) * u1[1]
end

function residual(model::Cartpole{T,Friction}, z, θ, κ) where T
    nq = model.nq
    nu = model.nu
    nc = model.nc

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    μ_slider = θ[2nq + nu .+ (1:1)]
    μ_angle = θ[2nq + nu + 1 .+ (1:1)]
    h = θ[2nq + nu + 1 + 1 .+ (1:1)]

    q2 = z[1:nq]
    ψ = z[nq .+ (1:nc)]
    b = z[nq + nc .+ (1:nc)]
    sψ = z[nq + nc + nc .+ (1:nc)]
    sb = z[nq + nc + nc + nc .+ (1:nc)]

    vT1 = (q2[1] - q1[1]) / h
    vT2 = (q2[2] - q1[2]) / h

    λ1 = b

    [
     dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b),
        h, q0, q1, u1, zeros(model.nw), λ1, q2);
     sb[1] - vT1[1];
     ψ[1] .- μ_slider[1] * (model.mp + model.mc) * model.g * h[1];
     sb[2] - vT2[1];
     ψ[2] .- μ_angle[1] * (model.mp * model.g * model.l) * h[1];
     cone_product([ψ[1]; b[1]], [sψ[1]; sb[1]]) - [κ[1]; 0.0];
     cone_product([ψ[2]; b[2]], [sψ[2]; sb[2]]) - [κ[1]; 0.0];
    ]
end

function residual(model::Cartpole{T,Frictionless}, z, θ, κ) where T
    nq = model.nq
    nu = model.nu

    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    h = θ[2nq + nu .+ (1:1)]

    q2 = z[1:nq]

    return dynamics(model, a -> M_func(model, a), (a, b) -> C_func(model, a, b),
        h, q0, q1, u1, zeros(model.nw), zeros(0), q2);
end

# models
cartpole_friction = Cartpole{Float64,Friction}(2, 1, 0, 2, 1.0, 0.2, 0.5, 9.81, [0.1; 0.1])
cartpole_frictionless = Cartpole{Float64,Frictionless}(2, 1, 0, 2, 1.0, 0.2, 0.5, 9.81, [0.0; 0.0])


using SparseDiffTools
using LightGraphs
using Test

using Random
Random.seed!(123)

#= Test data =#
test_graphs = Array{SimpleGraph, 1}(undef, 0)

for _ in 1:5
    nv = rand(5:20)
    ne = rand(1:100)
    graph = SimpleGraph(nv)
    for e in 1:ne
        v1 = rand(1:nv)
        v2 = rand(1:nv)
        while v1 == v2
            v2 = rand(1:nv)
        end
        add_edge!(graph, v1, v2)
    end
    push!(test_graphs, copy(graph))
end

#=
 Coloring needs to satisfy two conditions:

1. every pair of adjacent vertices receives distinct colors
(a distance-1 coloring)

2. For any vertex v, any color that leads to a two-colored path
involving v and three other vertices  is  impermissible  for  v.
In other words, every path on four vertices uses at least three
colors.
=#

#Sample graph from Gebremedhin AH, Manne F, Pothen A. **What color is your Jacobian? Graph coloring for computing derivatives.**

#=
     (2)
    /  \
   /    \
 (1)----(3)----(4)

=#

gx = SimpleGraph(4)

add_edge!(gx,1,2)
add_edge!(gx,1,3)
add_edge!(gx,2,3)
add_edge!(gx,3,4)

push!(test_graphs, gx)

#begin testing
for i in 1:6
    g = test_graphs[i]

    out_colors1 = SparseDiffTools.color_graph(g,SparseDiffTools.GreedyStar1Color())
    out_colors2 = SparseDiffTools.color_graph(g,SparseDiffTools.GreedyStar2Color())

    #test condition 1
    for v = vertices(g)
        color = out_colors1[v]
        for j in inneighbors(g, v)
            @test out_colors1[j] != color
        end
    end

    #test condition 2
    for j = vertices(g)
        walk = LightGraphs.self_avoiding_walk(g, j, 4)
        walk_colors = zeros(Int64, 0)
        if length(walk) >= 4
            for t in walk
                push!(walk_colors, out_colors1[t])
            end
            @test length(unique(walk_colors)) >= 3
        end
    end

    #test condition 1
    for v = vertices(g)
        color = out_colors2[v]
        for j in inneighbors(g, v)
            @test out_colors2[j] != color
        end
    end

    #test condition 2
    for j = vertices(g)
        walk = LightGraphs.self_avoiding_walk(g, j, 4)
        walk_colors = zeros(Int64, 0)
        if length(walk) >= 4
            for t in walk
                push!(walk_colors, out_colors2[t])
            end
            @test length(unique(walk_colors)) >= 3
        end
    end

end


StrikePrices = [13:0.25:19]

SensitivityMatrix = zeros(Float64, length(StrikePrices), 7)

SensitivityMatrix[:,1] = StrikePrices
SensitivityMatrix[:,2] = 0.25
for i in 1:length(StrikePrices) SensitivityMatrix[i,3] = Black76("c", 11, StrikePrices[i], 1/12, 0.05, SensitivityMatrix[i,2]) end

for i in 1:length(StrikePrices) SensitivityMatrix[i,4] = VannaVolga(10, 11, 12, StrikePrices[i], "c", 1, 0.05, 0, 0.26, 0.25, 0.35)[2] end

for i in 1:length(StrikePrices) SensitivityMatrix[i,5] = Black76("c", 11, StrikePrices[i], 1/12, 0.05, SensitivityMatrix[i,4]) end

for i in 1:length(StrikePrices) SensitivityMatrix[i,6] = Delta("c", 11, StrikePrices[i], 1/12, 0.05, 0, SensitivityMatrix[i,4]) end

for i in 1:length(StrikePrices) SensitivityMatrix[i,7] = Gamma("c", 11, StrikePrices[i], 1/12, 0.05, 0, SensitivityMatrix[i,4]) end

CallAugust = DataFrame(SensitivityMatrix)

impliedVol = zeros(Float64, length(StrikePrices))

for i in 1:length(StrikePrices) impliedVol[i] = VannaVolga(14, 15.75, 18, StrikePrices[i], "c", (84/365), 0.005, 0, 0.34, 0.3245, 0.3375)[2] end

DataFrame(Strikes=StrikePrices, Implied=impliedVol)

Dates.days(Date(2016, 9,29) - today())/365


using LinearAlgebra

"""
    filter_cube(snap_file::String, corner_lowerleft::Array{<:Real}, corner_upperright::Array{<:Real},
                          parttype::Integer)

Reads positions from `snap_file` and returns the indices of particles contained in a box defined by `corner_lowerleft` and `corner_upperright`.
"""
function filter_cube(snap_file::String, corner_lowerleft::Array{<:Real}, corner_upperright::Array{<:Real}; parttype::Integer=0)

    # read positions from file
    pos = read_snap(snap_file, "POS", parttype)

    # double check if the corners are set correctly
    cube_lower = Array{eltype(corner_lowerleft[1]),1}(undef, 3)
    cube_upper = Array{eltype(corner_lowerleft[1]),1}(undef, 3)
    @inbounds for dim = 1:3
        cube_lower[dim] = ((corner_lowerleft[dim] < corner_upperright[dim]) ? corner_lowerleft[dim] : corner_upperright[dim])
        cube_upper[dim] = ((corner_lowerleft[dim] > corner_upperright[dim]) ? corner_lowerleft[dim] : corner_upperright[dim])
    end

    return findall( @views ( ( cube_lower[1] .<= pos[1,:] .<= cube_upper[1] ) .&
                             ( cube_lower[2] .<= pos[2,:] .<= cube_upper[2] ) .&
                             ( cube_lower[3] .<= pos[3,:] .<= cube_upper[3] )) )
end

"""
    check_in_cylinder(x, pt1, pt2, r)

Checks if a 3D point `x` is in a cylinder defined by its endpoints `pt1` and `pt2` and radius `r`.
"""
function check_in_cylinder(x, pt1, pt2, r)
    # https://stackoverflow.com/questions/47932955/how-to-check-if-a-3d-point-is-inside-a-cylinder
    ( ( x .- pt1 ) ⋅ ( pt2 .- pt1 ) ≥ 0.0                             ) &
    ( ( x .- pt2 ) ⋅ ( pt2 .- pt1 ) ≤ 0.0                             ) &
    ( ( norm( (x .- pt1) × (pt2 .- pt1 ) ) / norm( pt2 .- pt1 ) ) ≤ r )
end

"""
    filter_cylinder(filename::String, pt1::Array{<:Real}, pt2::Array{<:Real}, r::Real)

Reads the positions contained in a file and returns the indices of particles contained in a cylinder defined by the endpoints `pt1` and `pt2` and radius `r`.
"""
function filter_cylinder(filename::String, pt1::Array{<:Real}, pt2::Array{<:Real}, r::Real; parttype::Integer=0)

    in_cube = filter_cube(filename, pt1, pt2, parttype=parttype )

    if size(in_cube,1) > 0
        check_in_cylinder_helper(x) = check_in_cylinder(x, pt1, pt2, r)
        # read positions from file
        pos = read_snap(filename, "POS", parttype)

        pos_cylinder = @views pos[:, in_cube]
        is_in_cylinder = mapslices(check_in_cylinder_helper, pos_cylinder, dims=1)
        return in_cube[findall( is_in_cylinder .== true )]
    else
        return in_cube
    end
end

"""
    filter_sphere(filename::String, center::Array{<:Real}, r::Real; parttype::Integer=0)

Reads the positions contained in a file and returns the indices of particles contained in a cylinder defined by the endpoints `pt1` and `pt2` and radius `r`.
"""
function filter_sphere(filename::String, center::Array{<:Real}, r::Real; parttype::Integer=0)

    in_cube = filter_cube(filename, center .- r, center .+ r, parttype=parttype)

    if size(in_cube,1) > 0

        distance = @views @. √( (pos[1, in_cube] - center[1])^2 +
                                (pos[2, in_cube] - center[2])^2 +
                                (pos[3, in_cube] - center[3])^2 )

        return in_cube[findall( distance .<= r )]
    else
        return in_cube
    end
end


# We thank QuTiP developers for sharing the notebook on Jaynes Cummings model,
# https://github.com/qutip/qutip-notebooks/blob/master/examples/example-JC-model-wigner-function.ipynb
# this is a direct translation of QuTiP version to solving in Julia using
# QuDynamics.

using QuBase
using QuDynamics

function jc_setup(N, wc, wa, g, kappa, gamma, use_rwa)
    # Hamiltonian
    idc = QuArray(eye(N))
    ida = QuArray(eye(2))
    a  = tensor(lowerop(N), ida)
    sm = tensor(idc, lowerop(2))
    if use_rwa
        # use the rotating wave approxiation
        H = wc * a' * a + wa * sm' * sm + g * (a' * sm + a * sm')
    else
        H = wc * a' * a + wa * sm' * sm + g * (a' + a) * (sm + sm')
    end
    # collapse operators
    c_op_list = Array(QuBase.AbstractQuMatrix, 0)
    n_th_a = 0.0
    rate = kappa * (1 + n_th_a)
    if rate > 0.0
       push!(c_op_list, full(sqrt(rate) * a))
    end
    rate = kappa * n_th_a
    if rate > 0.0
        push!(c_op_list, full(sqrt(rate) * a'))
    end
    rate = gamma
    if rate > 0.0
        push!(c_op_list, full(sqrt(rate) * sm))
    end
    return full(H), c_op_list
end

# parameters
wc = 1.0 * 2 * pi   # cavity frequency
wa = 1.0 * 2 * pi   # atom frequency
g  = 0.05 * 2 * pi  # coupling strength
kappa = 0.05        # cavity dissipation rate
gamma = 0.15        # atom dissipation rate
N = 2           # number of cavity fock states
use_rwa = true

# initial state
psi0 = complex(tensor(statevec(1, FiniteBasis(N)), statevec(2, FiniteBasis(2))))

tlist = 0.:0.1:2*pi

H, c_ops = jc_setup(N, wc, wa, g, kappa, gamma, use_rwa)

qprop = QuPropagator(H, c_ops, psi0*psi0', tlist, QuODE45())

# Calculating the expectation values for observable a'*a
a = tensor(lowerop(2), QuArray(eye(2)))
op = a'*a
for (t, rho) in qprop
    println("t = $t : <a'*a> = $(trace(rho*op))")
end


#-------------------------------------------------------------------------------
#general wrapper for interpolating profiles in log pressure coordinates

export AtmosphericProfile

struct AtmosphericProfile{Q}
    ϕ::LinearInterpolator{Q,NoBoundaries}
end

function Base.show(io::IO, ::AtmosphericProfile{Q}) where {Q}
    print(io, "AtmosphericProfile{$Q}")
end

Base.copy(x::AtmosphericProfile) = AtmosphericProfile(copy(x.ϕ))

function AtmosphericProfile(P::AbstractVector{<:Real}, y::AbstractVector{Q}) where {Q<:Real}
    @assert length(P) == length(y) "cannot form AtmosphericProfile with unequal numbers of points"
    #pressure in ascending order
    idx = sortperm(P)
    P = collect(Q, P[idx])
    y = y[idx]
    ϕ = LinearInterpolator(log.(P), y, NoBoundaries())
    AtmosphericProfile(ϕ)
end

(x::AtmosphericProfile)(P) = x.ϕ(log(P))

#-------------------------------------------------------------------------------
#constructing generalized hydrostatic pressure profiles and inverting for z

export scaleheight, hydrostatic, altitude
export Hydrostatic

"""
    scaleheight(g, μ, T)

Evaluate the [atmospheric scale height](https://en.wikipedia.org/wiki/Scale_height),

``\\frac{RT}{μg}``

where ``R`` is the [universial gas constant](https://en.wikipedia.org/wiki/Gas_constant).

# Arguments
* `g`: gravitational acceleration [m/s``^s``]
* `μ`: mean molar mass [kg/mole]
* `T`: temperature [K]
"""
scaleheight(g, μ, T) = 𝐑*T/(μ*g)

#parameterized hydrostatic relation in log coordinates
function dlnPdz(z, lnP, param::Tuple)
    #unpack parameters
    Pₛ, g, fT, fμ = param
    #evaluate temperature and mean molar mass [kg/mole]
    P = exp(lnP)
    P < 𝐏ₘᵢₙ && return zero(lnP)
    P = min(P, Pₛ) #don't allow tiny unphysical dips below Pₛ
    T = fT(P)
    μ = fμ(T, P)
    #evaluate derivative
    -μ*g/(𝐑*T)
end

"""
    hydrostatic(z, Pₛ, g, fT, fμ)

Compute the hydrostatic pressure [Pa] at a specific altitude using arbitrary atmospheric profiles of temperature and mean molar mass. This function integrates the hydrostatic relation,

``\\frac{dP}{dz} = \\frac{\\mu g}{R T}

from the surface to a height of ``z``, where ``R`` is the [universial gas constant](https://en.wikipedia.org/wiki/Gas_constant).

# Arguments

* `z`: altitude [m] to compute pressure at
* `Pₛ`: surface pressure [Pa]
* `g`: gravitational acceleration [m/s``^2``]
* `fT`: temperature [K] as a function of pressure, `fT(P)`
* `fμ`: mean molar mass [kg/mole] as a function of pressure and temperature `fμ(T,P)`
"""
function hydrostatic(z, Pₛ, g, fT::T, fμ::U) where {T,U}
    @assert z >= 0 "cannot compute pressure at negative altitude $z m"
    @assert Pₛ > 𝐏ₘᵢₙ "pressure cannot be less than $𝐏ₘᵢₙ Pa"
    #integration parameters
    param = (Pₛ, g, fT, fμ)
    #integrate in log coordinates and return
    exp(radau(dlnPdz, log(Pₛ), zero(z), z, param))
end

"""
    altitude(P, Pₛ, g, fT, fμ)

Compute the altitude [m] at which a specific hydrostatic pressure occurs using arbitrary atmospheric profiles of temperature and mean molar mass. This function applies a root finder to the [`hydrostatic`](@ref) function.

# Arguments

* `P`: pressure [Pa] to compute altitude at
* `Pₛ`: surface pressure [Pa]
* `g`: gravitational acceleration [m/s``^2``]
* `fT`: temperature [K] as a function of pressure, `fT(P)`
* `fμ`: mean molar mass [kg/mole] as a function of pressure and temperature `fμ(T,P)`
"""
function altitude(P, Pₛ, g, fT::T, fμ::U) where {T,U}
    @assert P < Pₛ "surface pressure must be greater than pressure aloft"
    #pressure decreases monotonically, find altitudes bracketing Pₜ
    z₁ = zero(P)
    z₂ = 1e2*one(P)
    P₁ = Pₛ
    P₂ = hydrostatic(z₂, Pₛ, g, fT, fμ)
    while P₂ > P
        z₁ = z₂
        z₂ *= 2
        P₁ = P₂
        P₂ = hydrostatic(z₂, Pₛ, g, fT, fμ)
    end
    #find precise altitude where P = Pₜ
    fₕ(z,::Any) = log(hydrostatic(z, Pₛ, g, fT, fμ)) - log(P)
    regulafalsi(fₕ, z₁, z₂)
end

"""
[Function-like type](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) for initializing and evaluating a hydrostatic pressure profile with arbitrary temperature and mean molar mass profiles. A `Hydrostatic` object maps altitude to pressure. Internally, a pressure vs altitude profile is generated and used for interpolation.

# Constructor

    Hydrostatic(Pₛ, Pₜ, g, fT, fμ, N=250)

* `Pₛ`: surface pressure [Pa]
* `Pₜ`: top of profile pressure [Pa]
* `g`: gravitational acceleration [m/s``^2``]
* `fT`: temperature [K] as a function of presssure, `fT(P)`
* `fμ`: mean molar mass [kg/mole] as a function of temperature and pressure, `fμ(T,P)`
* `N`: optional, number of interpolation nodes

For a constant molar mass or temperature, you can use [anonymous functions](https://docs.julialang.org/en/v1/manual/functions/#man-anonymous-functions) directly. For example, to construct a hydrostatic pressure profile for a crude Earth-like atmosphere:

```@example
#moist adiabatic temperature profile
M = MoistAdiabat(288, 1e5, 1040, 1996, 0.029, 0.018, 2.3e6, psatH2O, Ptropo=1e4);
#hydrostatic pressure profile with constant mean molar mass
H = Hydrostatic(1e5, 1, 9.8, M, (T,P)->0.029);
#evaluate pressures at a few different altitudes
H.([0, 1e3, 1e4])
```
"""
struct Hydrostatic{T}
    ϕ::LinearInterpolator{T,WeakBoundaries}
    zₜ::T
end

function Hydrostatic(Pₛ, Pₜ, g, fT::T, fμ::U, N::Int=100) where {T,U}
    Pₛ, Pₜ = promote(Pₛ, Pₜ)
    #find the altitude corresponding to Pₜ
    zₜ = altitude(Pₜ, Pₛ, g, fT, fμ)
    #interpolation knots and output array
    z = logrange(zero(zₜ), zₜ, N)
    lnP = zeros(typeof(Pₛ), N)
    #integration parameters
    param = (Pₛ, g, fT, fμ)
    #integrate to get a full pressure profile
    radau!(lnP, z, dlnPdz, log(Pₛ), zero(zₜ), zₜ, param)
    #construct and return
    Hydrostatic(LinearInterpolator(z, lnP, WeakBoundaries()), zₜ)
end

(H::Hydrostatic)(z) = exp(H.ϕ(z))

"""
    altitude(H::Hydrostatic, P)

Compute the altitude at which a specific pressure occurs in a [`Hydrostatic`](@ref) pressure profile. A root finder is applied to the object.
"""
function altitude(H::Hydrostatic, P::Real)::Float64
    regulafalsi((z,p) -> log(H(z)) - log(P), zero(H.zₜ), H.zₜ)
end

#-------------------------------------------------------------------------------

#general function for adiabat with one condensible in bulk non-condensible
function dTdP(P, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F) where {F}
    #molar mixing ratio of condensible
    α = psat(T)/P
    #specific gas constants
    Rₙ = 𝐑/μₙ
    Rᵥ = 𝐑/μᵥ
    #numerator
    𝑵 = 1.0 + α*L/(Rₙ*T)
    #denominator
    𝑫 = 1.0 + α*(cₚᵥ/cₚₙ + (L/(T*Rᵥ) - 1.0)*L/(cₚₙ*T))
    #final expression
    (T/P)*(Rₙ/cₚₙ)*(𝑵/𝑫)
end

#same function in ω coordinates
function dTdω(ω, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F) where {F}
    P = ω2P(ω)
    -dωfac(P)*dTdP(P, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F)
end

#slurp up the parameters for integration
dTdω(ω, T, param::Tuple) = dTdω(ω, T, param...)

#-------------------------------------------------------------------------------
#exported wrappers of dTdω

export lapserate, lapse!

#single-condensible moist lapse rate
function lapserate(T, P, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F) where {F}
    dTdP(P, T, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F)
end

#dry lapse rate
function lapserate(T, P, cₚ, μ)
    dTdP(P, T, cₚ, 1.0, μ, 1.0, 0.0, T->0.0)
end

function lapse!(T, P, cₚ, μ)
    @assert length(P) == length(T)
    idx = sortperm(P, rev=true) #pressure sorting in descending order
    for n ∈ 1:length(idx)-1
        i, j = idx[n], idx[n+1]
        #expected lapse rate
        Γₑ = lapserate(T[i], P[i], cₚ, μ)
        #lapse rate of profile
        Γₚ = (T[j] - T[i])/(P[j] - P[i])
        #heat the upper point if needed
        if Γₚ > Γₑ
            T[j] = T[i] + Γₑ*(P[j] - P[i])
        end
    end
end

#------------------------------------

abstract type AbstractAdiabat end

export MoistAdiabat, DryAdiabat
export tropopause

function checkadiabat(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    @assert Pₛ > Pₜ "Pₛ must be greater than Pₜ"
    @assert Pₜ > 0 "Pₜ must be greater than 0"
    @assert Tstrat >= 0 "stratosphere temperature cannot be negative"
    @assert Ptropo >= 0 "tropopause pressure cannot be negative"
    @assert smooth >= 0 "smoothing distance cannot be negative"
    if Tstrat > 0
        @assert Tstrat < Tₛ "Tstrat cannot be greater than Tₛ"
    end
    if (Tstrat != 0) & (Ptropo != 0)
        throw("Cannot have nonzero Tstrat and Ptropo, must use one or the other")
    end
end

#------------------------------------

"""
[Function-like type](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) for initializing and evaluating a dry adiabatic temperature profile. Optionally, a uniform upper atmospheric temperature can be set below a specified temperature or pressure.

# Constructor

    DryAdiabat(Tₛ, Pₛ, cₚ, μ; Tstrat=0.0, Ptropo=0.0, Pₜ=$𝐏ₘᵢₙ)

* `Tₛ`: surface temperature [K]
* `Pₛ`: surface pressure [K]
* `cₚ`: specific heat of the atmosphere [J/kg/K]
* `μ`: molar mass of the atmosphere [kg/mole]

If `Tstrat` is greater than zero, the temperature profile will never drop below that temperature. If `Ptropo` is greater than zero, the temperature profile at pressures lower than `Ptropo` will be equal to the temperature at exactly `Ptropo`. `Tstrat` and `Ptropo` cannot both greater than zero.

`Pₜ` defines the highest pressure [Pa] in the temperature profile. If should generally be small but cannot be zero.

# Example

Once constructed, use a `DryAdiabat` like a function to compute temperature at a given pressure.

```@example
Tₛ = 288; #surface temperature [K]
Pₛ = 1e5; #surface pressure [Pa]
cₚ = 1040; #specific heat of air [J/kg/K]
μ = 0.029; #molar mass of air [kg/mole]

#construct the dry adiabat with an upper atmosphere temperature of 190 K
D = DryAdiabat(Tₛ, Pₛ, cₚ, μ, Tstrat=190);

#temperatures at 40-10 kPa
D.([4e4, 3e4, 2e4, 1e4])
```
"""
struct DryAdiabat{U} <: AbstractAdiabat
    Tₛ::U
    Pₛ::U
    Pₜ::U
    cₚ::U
    μ::U
    Tstrat::U
    Ptropo::U
    #smoothing quantities
    smooth::U #smoothing distance
    T₂::U #temperature at beginning of smoothing
    h₂::U #hermite thing at T₂
end

function Base.show(io::IO, Γ::DryAdiabat{T}) where {T}
    print(io, "DryAdiabat{$T}:\n")
    print(io, "  Tₛ     = $(Γ.Tₛ) K\n")
    print(io, "  Pₛ     = $(Γ.Pₛ) Pa\n")
    print(io, "  cₚ     = $(Γ.cₚ) J/kg/K\n")
    print(io, "  μ      = $(Γ.μ) kg/mole\n")
    if (Γ.Tstrat != 0) & (Γ.Ptropo != 0)
        Tstrat = round(Γ.Tstrat, sigdigits=6)
        print(io, "  Tstrat = $Tstrat K\n")
        Ptropo = round(Γ.Ptropo, sigdigits=6)
        print(io, "  Ptropo = $Ptropo Pa\n")
        print(io, "  smoothing interval of $(Γ.smooth) Pa")
    end
end

function DryAdiabat(Tₛ, Pₛ, cₚ, μ;
                    Tstrat=0.0,
                    Ptropo=0.0,
                    smooth=1e2,
                    Pₜ=𝐏ₘᵢₙ)
    checkadiabat(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    #fill in Tstrat or Ptropo if needed
    if Tstrat != 0
        Ptropo = regulafalsi((P,_) -> temperature(P, Tₛ, Pₛ, cₚ, μ) - Tstrat, Pₛ, Pₜ)
    elseif Ptropo != 0
        Tstrat = temperature(Ptropo, Tₛ, Pₛ, cₚ, μ)
    end
    #get smoothing connection ready if Ptropo is nonzero
    h₂ = zero(Tₛ)
    T₂ = zero(Tₛ)
    if Ptropo != 0
        P₂ = Ptropo + smooth
        T₂ = temperature(P₂, Tₛ, Pₛ, cₚ, μ)
        T₂′ = lapserate(T₂, P₂, cₚ, μ)
        h₂ = smooth*T₂′
    end
    DryAdiabat(promote(Tₛ, Pₛ, Pₜ, cₚ, μ, Tstrat, Ptropo, smooth, T₂, h₂)...)
end

#direct calculation of raw temperature profile
temperature(P, Tₛ, Pₛ, cₚ, μ) = Tₛ*(P/Pₛ)^(𝐑/(μ*cₚ))

temperature(Γ::DryAdiabat, P) = temperature(P, Γ.Tₛ, Γ.Pₛ, Γ.cₚ, Γ.μ)

#------------------------------------

"""
[Function-like type](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) for initializing and evaluating a moist adiabatic temperature profile. Optional uniform upper atmospheric temperature below a specified temperature or pressure.

# Constructor

    MoistAdiabat(Tₛ, Pₛ, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat; Tstrat=0, Ptropo=0, N=1000, Pₜ=$𝐏ₘᵢₙ)

* `Tₛ`: surface temperature [K]
* `Pₛ`: surface pressure [K]
* `cₚₙ`: specific heat of the non-condensible atmospheric component (air) [J/kg/K]
* `cₚᵥ`: specific heat of the condensible atmospheric component [J/kg/K]
* `μₙ`: molar mass of the non-condensible atmospheric component (air) [kg/mole]
* `μᵥ`: molar mass of the condensible atmospheric component [kg/mole]
* `L`: condsible component's latent heat of vaporization [J/kg]
* `psat`: function defining the saturation vapor pressure for a given temperature, `psat(T)`

If `Tstrat` is greater than zero, the temperature profile will never drop below that temperature. If `Ptropo` is greater than zero, the temperature profile at pressures lower than `Ptropo` will be equal to the temperature at exactly `Ptropo`. `Tstrat` and `Ptropo` cannot both greater than zero.

`Pₜ` defines the highest pressure [Pa] in the temperature profile. If should generally be small but cannot be zero.

The profile is evaluated along a number of pressure values in the atmosphere set by `N`. Those points are then used to construct a cubic spline interpolator for efficient and accurate temperature calculation. Experience indicates that 1000 points is very accurate and also fast.

# Example

Once constructed, use a `MoistAdiabat` like a function to compute temperature at a given pressure.

```@example
Tₛ = 288; #surface temperature [K]
Pₛ = 1e5; #surface pressure [Pa]
cₚₙ = 1040; #specific heat of air [J/kg/K]
cₚᵥ = 1996; #specific heat of H2O [J/kg/K]
μₙ = 0.029; #molar mass of air [kg/mole]
μᵥ = 0.018; #molar mass of H2O [kg/mole]
L = 2.3e6; #H2O latent heat of vaporization [J/kg]

#a saturation vapor pressure function for H2O is built in
psat = psatH2O;

#construct the moist adiabat with a tropopause pressure of 1e4 Pa
M = MoistAdiabat(Tₛ, Pₛ, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat, Ptropo=1e4);

#temperatures at 30-5 kPa
M.([3e4, 2e4, 1e4, 5e3])
```
"""
struct MoistAdiabat{U} <: AbstractAdiabat
    ϕ::LinearInterpolator{U,WeakBoundaries}
    Pₛ::U
    Pₜ::U
    Tstrat::U
    Ptropo::U
    #smoothing quantities
    smooth::U #smoothing distance
    T₂::U #temperature at beginning of smoothing
    h₂::U #hermite thing at T₂
end

function Base.show(io::IO, Γ::MoistAdiabat{T}) where {T}
    print(io, "MoistAdiabat{$T}:\n")
    print(io, "  Tₛ     = $(Γ(Γ.Pₛ)) K\n")
    print(io, "  Pₛ     = $(Γ.Pₛ) Pa\n")
    if (Γ.Tstrat != 0) & (Γ.Ptropo != 0)
        Tstrat = round(Γ.Tstrat, sigdigits=6)
        print(io, "  Tstrat = $Tstrat K\n")
        Ptropo = round(Γ.Ptropo, sigdigits=6)
        print(io, "  Ptropo = $Ptropo Pa\n")
        print(io, "  smoothing interval of $(Γ.smooth) Pa")
    end
end

function Base.copy(Γ::MoistAdiabat)
    MoistAdiabat(copy(Γ.ϕ), Γ.Pₛ, Γ.Pₜ, Γ.Tstrat, Γ.Ptropo, Γ.smooth, Γ.T₂, Γ.h₂)
end

function MoistAdiabat(Tₛ, Pₛ, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat::F;
                      Tstrat=0.0,
                      Ptropo=0.0,
                      smooth=1e2,
                      N::Int=100,
                      Pₜ=𝐏ₘᵢₙ) where {F}
    #basic checks
    checkadiabat(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    #type uniformity
    Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth = promote(Tₛ, Pₛ, Pₜ, Tstrat, Ptropo, smooth)
    #interpolation knots and output vector
    ω₁, ω₂ = P2ω(Pₛ, Pₜ)
    ω = logrange(ω₁, ω₂, N)
    T = zeros(typeof(Tₛ), N)
    #pack the parameters
    param = (cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat)
    #integrate with in-place dense output
    radau!(T, ω, dTdω, Tₛ, ω[1], ω[end], param)
    #interpolator ω pressure coordinates
    ϕ = LinearInterpolator(ω, T, WeakBoundaries())
    #fill in Tstrat or Ptropo if needed
    if Tstrat != 0
        Ptropo = regulafalsi((P,_) -> temperature(ϕ, P) - Tstrat, Pₛ, Pₜ)
    elseif Ptropo != 0
        Tstrat = temperature(ϕ, Ptropo)
    end
    #get smoothing connection ready if Ptropo is nonzero
    h₂ = zero(Tₛ)
    T₂ = zero(Tₛ)
    if Ptropo != 0
        P₂ = Ptropo + smooth
        T₂ = temperature(ϕ, P₂)
        T₂′ = lapserate(T₂, P₂, cₚₙ, cₚᵥ, μₙ, μᵥ, L, psat)
        h₂ = smooth*T₂′
    end
    MoistAdiabat(ϕ, Pₛ, Pₜ, Tstrat, Ptropo, smooth, T₂, h₂)
end

#direct calculation of raw temperature profile
temperature(ϕ::LinearInterpolator, P) = ϕ(P2ω(P))

temperature(Γ::MoistAdiabat, P) = temperature(Γ.ϕ, P)

#------------------------------------
#general operations with an adiabat

#find the pressure corresponding to a temperature, ignoring Tstrat & Ptropo
function pressure(Γ::AbstractAdiabat, T)
    #surface temperature
    Tₛ = temperature(Γ, Γ.Pₛ)
    #TOA temperature (at very tiny pressure)
    Tₜ = temperature(Γ, Γ.Pₜ)
    #T profile should always decrease with P, so demand bracketing
    @assert Tₛ >= T >= Tₜ "temperature $T K out of adiabat range [$(Tₛ),$(Tₜ)] K"
    #find the root with false position
    regulafalsi((P,::Any) -> temperature(Γ, P) - T, Γ.Pₛ, Γ.Pₜ)
end

function (Γ::AbstractAdiabat)(P)
    #return tropopause temperature if P is below Ptropo
    P < Γ.Ptropo && return Γ.Tstrat*one(P)
    #check if P is inside the smoothing region
    if (Γ.Ptropo != 0) & (Γ.smooth != 0)
        P₁ = Γ.Ptropo
        if Γ.Ptropo < P < P₁ + Γ.smooth
            ψ = (P - Γ.Ptropo)/Γ.smooth
            T₁ = Γ.Tstrat
            T₂ = Γ.T₂
            h₂ = Γ.h₂
            #smooth cubic connection
            return ψ^3*(2*T₁ - 2*T₂ + h₂) + ψ^2*(-3*T₁ + 3*T₂ - h₂) + T₁
        end
    end
    #raw temperature from the profile, unadjusted
    T = temperature(Γ, P)
    #return stratosphere temperature if T is below Tstrat
    T < Γ.Tstrat && return Γ.Tstrat
    #ensure positive
    @assert T > 0 "non-positive temperature ($T K) encountered in adiabat at $P Pa"
    return T
end

"""
    tropopause(Γ::AbstractAdiabat)

Compute the temperature [K] and pressure [Pa] at which the tropopause occurs in an adiabatic temperature profile. This function can be called on a `DryAdiabat` or a `MoistAdiabat` if it was constructed with nonzero `Tstrat` or `Ptropo`. Returns the tuple `(T,P)`.
"""
function tropopause(Γ::AbstractAdiabat)
    Γ.Ptropo != 0 && Γ.Tstrat != 0 && return (Γ.Tstrat,Γ.Ptropo)
    error("no stratosphere temperature or pressure has been defined (Tstrat/Ptropo)")
end

#-------------------------------------------------------------------------------
export psatH2O, tsatCO2, ozonelayer

"""
    psatH2O(T)

Compute the saturation partial pressure of water vapor at a given temperature using expressions from

* [Murphy, D. M. & Koop, T. Review of the vapour pressures of ice and supercooled water for atmospheric applications. Q. J. R. Meteorol. Soc. 131, 1539–1565 (2005).](https://rmets.onlinelibrary.wiley.com/doi/10.1256/qj.04.94)

This function uses equation 10 in the paper above when ``T >= 273.15`` K and equation 7 otherwise.
"""
function psatH2O(T)
    a = log(T)
    b = 1/T
    if T >= 273.15
        #equation 10
        c = 53.878 - 1331.22*b - 9.44523*a + 0.014025*T
        d = c*tanh(0.0415*(T - 218.8))
        P = exp(54.842763 - 6763.22*b - 4.21*a + 3.67e-4*T + d)
    else
        #equation 7
        P = exp(9.550426 - 5723.265*b + 3.53068*a - 0.00728332*T)
    end
    return P
end

"""
    tsatCO2(P)

Compute the saturation pressure of carbon dioxide at a certain pressure using equation 19 from

* [Fanale, F. P., Salvail, J. R., Bruce Banerdt, W. & Steven Saunders, R. Mars: The regolith-atmosphere-cap system and climate change. Icarus 50, 381–407 (1982)](https://doi.org/10.1016/0019-1035(82)90131-2)

The equation is inverted to express temperature as a function of pressure.
"""
function tsatCO2(P)
    @assert P <= 518000.0 "Pressure cannot be above 518000 Pa for CO2 saturation temperature"
    A = 1.2264e12 #[Pa]
    B = -3167.8 #[K]
    B/log(P/A)
end

"""
    ozonelayer(P, Cmax=8e-6)

Approximate the molar concentration of ozone in Earth's ozone layer using an 8 ppm peak at 1600 Pa which falls to zero at 100 Pa and 25500 Pa. Peak concentration is defined by `Cmax`. This approximation is discussed in

* Jacob, D. Introduction to Atmospheric Chemistry. (Princeton University Press, 1999).

"""
function ozonelayer(P, Cmax=8e-6)
    P = log(P)
    P₁ = 10.146433731146518 #ln(25500)
    P₂ = 7.3777589082278725 #ln(1600)
    P₃ = 4.605170185988092  #ln(100)
    if P₂ <= P <= P₁
        return Cmax*(P₁ - P)/(P₁ - P₂)
    elseif P₃ <= P <= P₂
        return Cmax*(P - P₃)/(P₂ - P₃)
    end
    return 0.0*P
end

#-------------------------------------------------------------------------------
#miscellaneous stuff

export condensibleprofile, haircut!, rayleighCO2

"""
    condensibleprofile(Γ::AbstractAdiabat, fPₛ)

Create a function defining concentration vs pressure for a condensible with uniform upper-atmosphere (stratosphere) concentration. The new concentration profile is created with reference to an existing adiabatic profile ([`DryAdiabat`](@ref) or [`MoistAdiabat`](@ref)), which must have `Ptropo != 0` or `Tstrato != 0`. Lower atmospheric concentration is determined by the temperature dependent partial pressure function `fPₛ(T)`. The concentration is `P/(fPₛ + P)`, where `P` is the dry/non-condensible pressure.

"""
function condensibleprofile(Γ::AbstractAdiabat, fPₛ::F)::Function where {F}
    #insist on an isothermal stratosphere
    @assert ((Γ.Ptropo != 0) | (Γ.Tstrat != 0)) "adiabat must have isothermal stratosphere"
    #compute tropopause pressure and temperature
    Tₜ, Pₜ = tropopause(Γ)
    #compute saturation partial pressure at tropopause
    Pₛₜ = fPₛ(Tₜ)
    #create concentration function and return it
    let (Pₜ, Pₛₜ) = (Pₜ, Pₛₜ)
        function (T, P)
            if P >= Pₜ
                Pₛ = fPₛ(T)
                C = Pₛ/(Pₛ + P)
            else
                C = Pₛₜ/(Pₜ + Pₛₜ)
            end
            return C
        end
    end
end

#this is already exported in gases.jl
"""
    reconcentrate(G::Gas, Γ::AbstractAdiabat, fPₛ)

Create a new concentration function for a [`Gas`](@ref) and use it to [`reconcentrate`](@the gas). This does not automatically copy gas data.
"""
function reconcentrate(g::Gas, Γ::AbstractAdiabat, fPₛ)::Gas
    #assign the gas a new concentration profile
    reconcentrate(g, condensibleprofile(Γ, fPₛ))
end

export haircut!
"""
    haircut!(T, P, fTₛ)

Put a temperature floor on a temperature profile using the saturation temperature function `fTₛ(P)`.
"""
function haircut!(T, P, fTₛ::F)::Nothing where {F}
    @assert length(T) == length(P)
    for i ∈ eachindex(T)
        Tₛ = fTₛ(P[i])
        if Tₛ > T[i]
            T[i] = Tₛ
        end
    end
    nothing
end

function rayleighCO2(ν, Pₛ, g, θ)
    #convert to wavelength in micrometers
    λ = ν2λ(ν)*1e6
    #formula 2.32 from
    #  Hansen and Travis, Light Scattering in Planetary Atmospheres (1974)
    τ₀ = 1.527*(1/λ^4)*(1 + 0.013*(1/λ^2))
    #conversion from Venus conditions
    τ₀ *= 8.7/(93*𝐀)
    #final optical depth
    τ = (Pₛ/g)*τ₀

    γ = 0.75
    μ = cos(θ)
    β = 1 - exp(-τ/μ)
    f = γ*τ
    R⁻ = ((0.5 - γ*μ)*β + f)/(1 + f)
    R⁺ = f/(1 + f)
    R = 1 - (1 - R⁺)*(1 - R⁻)/((1 - R⁻))
end


using FluxExtensions, LinearAlgebra, MLDataPattern

FluxExtensions.log_normal(x, μ, σ2::T, d::Int) where {T<:Real}  = (d > 0) ? - sum((@. ((x - μ)^2) ./ σ2), dims=1) / 2 .- d * (log.(σ2) + log(2π)) / 2 : 0

"""
	LinearAlgebra.det(f, x::Vector)
	determinant of a jacobian of `f` at point `x`
"""
function LinearAlgebra.det(f, x::Vector)
	Σ = Flux.data(Flux.Tracker.jacobian(f, x))
	S = svd(Σ)
	2*sum(log.(S.S .+ 1f-6))
end

jacodeco(model, x::Vector, z::Vector) = sum(log_normal(z)) + sum(log_pxexpectedz(model, x, z)) - det(model.g, z)
function jacodeco(m, x, z)
	@assert nobs(x) == nobs(z)
	[jacodeco(m, getobs(x, i), getobs(z, i)) for i in 1:nobs(x)]
end
jacodeco(m, x) = jacodeco(m, x, m.g(x))

function correctedjacodeco(model, x::Vector, z::Vector)
	Σ = Flux.data(Flux.Tracker.jacobian(model.g, z))
	S = svd(Σ)
	J = inv(S)
	logd = 2*sum(log.(S.S .+ 1f-6))
	if (det(I + J * transpose(J)) < 0)
		# println("determinant = $(det(I + J * transpose(J))) was negative, returning jacodeco = -Inf")
		# println(J)
		# println(J * transpose(J))
		return -Inf32
	end
	logpz = log_normal(z, zeros(Float32, size(z)...), I + J * J')[1] # TODO I changed the transpose(J) to J' - might be an issue
	logpx = FluxExtensions.log_normal(x, x_from_z(model, z).data, 1, length(x) - length(z))[1]
	logpz + logpx - logd
end

function correctedjacodeco(m, x, z)
	@assert nobs(x) == nobs(z)
	[correctedjacodeco(m, getobs(x, i), getobs(z, i)) for i in 1:nobs(x)]
end

correctedjacodeco(m, x) = correctedjacodeco(m, x, manifoldz(model, x))

"""
	manifoldz(model, x, steps = 100, rand_init = false)
	find the point in the latent space `z` such that it minimizes the
	reconstruction error. This means that it finds the point on manifold
	defined by `model.g(z)` and x
"""
function manifoldz(model, x, steps = 100, z = Flux.data(zparams(model,x)[1]))
	z = param(z)
	ps = Flux.Tracker.Params([z])
	# println("Size z: $(size(z)), size x: $(size(x))")
	li = Flux.data(mean(log_pxexpectedz(model, x, z)))
	Flux.train!((i) -> -mean(log_pxexpectedz(model, x, z)), ps, 1:steps, ADAM())
	le = Flux.data(mean(log_pxexpectedz(model, x, z)))
	@info "initial = $(li) final = $(le)"
	Flux.data(z)
end


using StructuralCausalModels, Test

ProjDir = @__DIR__

amat_data = reshape([
  0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  1,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,
  0,1,0,0,1,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,1,1,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,
  0,0,0,0,1,0,1,0,1,1,0,0,0,0,0,0,
  1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  0,0,0,0,0,0,0,0,0,1,0,1,0,0,0,0,
  0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
  1,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,
  0,0,0,0,0,0,0,0,0,0,0,1,0,1,0,0
], (16,16));

vars = [Symbol("n$i") for i in 1:size(amat_data, 1)]
a = NamedArray(Int.(amat_data), (vars, vars), ("Rows", "Cols"));

dag = DAG("ag_example", a)

fn = joinpath(mktempdir(), "ag.dot")
to_graphviz(dag, fn)
#Sys.isapple() && run(`open -a GraphViz.app $(fn)`)

m = [:n3, :n5, :n6, :n15, :n16];
c = [:n4, :n7];

ag = ancestral_graph(dag; m=m, c=c)
#indx = indexin(names(ag, 1), vars)
indx = filter(x -> !isnothing(x), indexin(vars, names(ag, 1)))
ag = ag[indx, indx]

agmat = [
0 100 0  0  0  0   0  0 100;
100 0 0  0  0  0   0  0   0;
0   0 0  0  0  0   0  0   0;
0   0 1  0  0 10   0  0   0;
0   0 1  0  0  0   0  0   0;
0   1 0 10  1  0   0  0   0;
1   0 0  0  0  0   0  0 100;
0   0 0  0  1  0   1  0   0;
100 0 0  0  0  0 100  0   0;
]


function clean_array(a; threshold=1e-10)
   # cleanArray removes the negligible real and imaginary parts from an array.
   # varargin allows special key-value pairs to be used to set how cleanArray
   # operates
   # Optional parameters: 'threshold' (double); specifies the maximum ratio of
   # imag/real or real/imag before value is kept [default: 1e-10]

   cleanedArray = zeros(length(a),1); # init output

   realOverImag = abs(real(a)./imag(a));
   imagOverReal = abs(imag(a)./real(a));
   realTooSmall = map( x -> x - threshold < 0, realOverImag)
   imagTooSmall = map( x -> x - threshold < 0, imagOverReal)
   neitherTooSmall = ~(imagTooSmall | realTooSmall);
   cleanedArray = im*imag(a).*realTooSmall +
                  real(a).*imagTooSmall +
                  a.*neitherTooSmall;
   return cleanedArray
end


using SpecialFunctions

abstract type Integrator end
abstract type GaussianQuadrature end
abstract type GaussianOrthogonalPolynomial end
abstract type IntegrationFunction end

struct GaussLaguerrePolynomial <: GaussianOrthogonalPolynomial
  s::Float64
end

get_alpha(poly::GaussLaguerrePolynomial, i::Int) = 2 * i + 1 + poly.s
get_beta(poly::GaussLaguerrePolynomial, i::Int) = i * (i + poly.s)
get_mu_0(poly::GaussLaguerrePolynomial) = exp(lgamma(poly.s + 1.0))
get_w(poly::GaussLaguerrePolynomial, x::Float64) = ^(x, poly.s) * exp(-x)

mutable struct SegmentIntegral <: Integrator
  absoluteAccuracy::Float64
  absoluteError::Float64
  maxEvals::Int
  evals::Int
  intervals::Int
end

SegmentIntegral(intervals::Int) = SegmentIntegral(1.0, 0.0, 1, 0, intervals)

function operator(integrator::Integrator, f::Function, a::Float64, b::Float64)
  integrator.evals = 0
  if a == b
    return 0.0
  end

  if (b > a)
    return integrate(integrator, f, a, b)
  else
    return -integrate(integrator, f, b, a)
  end
end

function (integrator::SegmentIntegral)(f::Function, a::Float64, b::Float64)
  integrator.evals = 0
  if a == b
    return 0.0
  end

  if (b > a)
    return integrate(integrator, f, a, b)
  else
    return -integrate(integrator, f, b, a)
  end
end

function Math.integrate(integrator::SegmentIntegral, f::Union{Function, IntegrationFunction}, a::Float64, b::Float64)
  dx = (b - a) / (integrator.intervals * 1.0)
  sum_ = 0.5 * (f(a) + f(b))
  end_ = b - 0.5 * dx
  x = a + dx
  while x < end_
    sum_ += f(x)
    x += dx
  end

  return sum_ * dx
end

struct GaussLaguerreIntegration <: GaussianQuadrature
  x::Vector{Float64}
  w::Vector{Float64}
end

function (gli::GaussLaguerreIntegration)(f::IntegrationFunction)
  sum_ = 0.0
  @inbounds @simd for i = length(gli.x):-1:1
    sum_ += gli.w[i] * f(gli.x[i])
  end

  return sum_
end

function GaussLaguerreIntegration(n::Int, s::Float64 = 0.0)
  x, w = build_gaussian_quadrature(n, GaussLaguerrePolynomial(s))

  return GaussLaguerreIntegration(x, w)
end

get_order(integration::GaussianQuadrature) = length(integration.x)

function build_gaussian_quadrature(n::Int, poly::GaussianOrthogonalPolynomial)
  x = zeros(n)
  w = zeros(n)

  e = zeros(n - 1)

  @simd for i = 2:n
    @inbounds x[i] = get_alpha(poly, i - 1)
    @inbounds e[i - 1] = sqrt(get_beta(poly, i - 1))
  end

  x[1] = get_alpha(poly, 0)

  # sym tridiagonal
  # tridiag = SymTridiagonal(x, e)
  # eigtri = eigfact(tridiag)

  tqr = TqrEigenDecomposition(x, e, OnlyFirstRowEigenVector(), Overrelaxation())
  x_new = tqr.d
  ev = tqr.ev

  mu_0 = get_mu_0(poly)

  @simd for i = 1:n
    @inbounds w[i] = mu_0 * ev[i, 1] * ev[i, 1] / get_w(poly, x_new[i])
  end

  return x_new, w
end


using Test
using FFTW

using ClimateMachine.ConfigTypes
using ClimateMachine.Spectra
using ClimateMachine.Spectra:
    compute_gaussian!,
    compute_legendre!,
    SpectralSphericalMesh,
    trans_grid_to_spherical!,
    compute_wave_numbers

include("spherical_helper_test.jl")

@testset "power_spectrum_1d (GCM)" begin
    FT = Float64
    # -- TEST 1: power_spectrum_1d(AtmosGCMConfigType(), var_grid, z, lat, lon, weight)
    nlats = 32

    # Setup grid
    sinθ, wts = compute_gaussian!(nlats)
    yarray = asin.(sinθ) .* 180 / π
    xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0
    z = 1

    # Setup variable
    mass_weight = ones(Float64, length(z))
    var_grid =
        1.0 * reshape(
            sin.(xarray / xarray[end] * 5.0 * 2π) .* (yarray .* 0.0 .+ 1.0)',
            length(xarray),
            length(yarray),
            1,
        ) +
        1.0 * reshape(
            sin.(xarray / xarray[end] * 10.0 * 2π) .* (yarray .* 0.0 .+ 1.0)',
            length(xarray),
            length(yarray),
            1,
        )
    nm_spectrum, wave_numbers = power_spectrum_1d(
        AtmosGCMConfigType(),
        var_grid,
        z,
        yarray,
        xarray,
        mass_weight,
    )

    nm_spectrum_ = nm_spectrum[:, 10, 1]
    var_grid_ = var_grid[:, 10, 1]
    sum_spec = sum(nm_spectrum_)
    sum_grid = sum(var_grid_ .^ 2) / length(var_grid_)

    sum_res = (sum_spec - sum_grid) / sum_grid

    @test sum_res < 0.1
end

@testset "power_spectrum_2d (GCM)" begin
    # -- TEST 2: power_spectrum_2d
    # Setup grid
    FT = Float64
    nlats = 32
    sinθ, wts = compute_gaussian!(nlats)
    cosθ = sqrt.(1 .- sinθ .^ 2)
    yarray = asin.(sinθ) .* 180 / π
    xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0
    z = 1

    # Setup variable: use an example analytical P_nm function
    P_32 = sqrt(105 / 8) * (sinθ .- sinθ .^ 3)
    var_grid =
        1.0 * reshape(
            sin.(xarray / xarray[end] * 3.0 * π) .* P_32',
            length(xarray),
            length(yarray),
            1,
        )

    mass_weight = ones(Float64, z)
    spectrum, wave_numbers, spherical, mesh =
        power_spectrum_2d(AtmosGCMConfigType(), var_grid, mass_weight)

    # Grid to spherical to grid reconstruction
    reconstruction = trans_spherical_to_grid!(mesh, spherical)

    sum_spec = sum((0.5 * spectrum))
    dθ = π .* wts ./ sum(wts)
    dθ = π / length(wts)
    area_factor = reshape(cosθ .* dθ .^ 2 / 4π, (1, length(cosθ)))

    sum_grid = sum(0.5 .* var_grid[:, :, 1] .^ 2 .* area_factor) # scaled to average over Earth's area (units: m2/s2)
    sum_reco = sum(0.5 .* reconstruction[:, :, 1] .^ 2 .* area_factor)

    sum_res_1 = (sum_spec - sum_grid) / sum_grid
    sum_res_2 = (sum_reco - sum_grid) / sum_grid

    @test abs(sum_res_1) < 0.1
    @test abs(sum_res_2) < 0.1
end


export bezier, dbezier, bern_basis

"""
	bezier(npts, b[], dpts)

Calculate a Bezier curve over the `npts` control polygon points in `b`.

The method evaluates `dpts` points of the curve and gives back the tuple `(P, EV)`.
Here `P` is the vector of points memorized by columns while `EV` is the `1`-dimensional cellular complex.
This structure as it is formed could be directly imported in Plasm module by typing:
```julia
julia> P, EV = bezier(npts, b[], dpts);
julia> Plasm.view(P, EV)
```

---

# Arguments

- `npts::Int64`: the number of the control polygon vertices.
- `b::Array{Float64, 2}`: 2-dimensional Array containing the control polygon vertices in the x, y, z coordinates for columns.
- `dpts::Int64`: number of data points to be calculated on the curve.

---

# Examples

```jldoctest
julia> bezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[1]
3×7 Array{Float64,2}:
 1.0  1.56481  2.18519  2.75  3.14815  3.26852  3.0
 1.0  1.83333  2.33333  2.5   2.33333  1.83333  1.0
 0.0  0.0      0.0      0.0   0.0      0.0      0.0
```

```jldoctest
julia> bezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[2]
6-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
 [5, 6]
 [6, 7]
```

---

_By Elia Onofri_
"""

function bezier(npts::Int64, b::Array{Float64,2}, dpts::Int64)::Tuple{Array{Float64,2}, Array{Array{Int64,1},1}}

	@assert npts == length(b[1,:]) ("ERROR: there are not npts = $(npts) points of the polygon in b[]")

	step::Float64 = 1/(dpts-1)
	n::Int64 = npts-1

	C::Array{Float64,2} = zeros(npts, npts)
	for i = 0:n
		for j = 0:n-i
			if (n-j-i)%2==0
				C[i+1, j+1] = binomial(n-j, n-j-i)
			else
				C[i+1, j+1] = 0.0-binomial(n-j, n-j-i)
			end
		end
	end

	D::Array{Float64,2} = zeros(npts, npts)
	for i = 1:npts
		D[i,i] = abs(C[npts+1-i,1])
	end

	G::Array{Float64,2} = transpose(b)

	T::Array{Float64,2} = ones(dpts, npts)

	t::Float64 = 0.0
	for i = 1:dpts
		for j = 0:n-1  #Per j=0 il valore è 1 ed è già impostato
			T[i, n-j]=T[i, npts-j]*t
		end
		t = t + step
	end

	EV::Array{Array{Int64,1},1} = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

	for i = 1 : dpts-1
		app = Array{Int64}(2)
		app[1] = i
		app[2] = i+1
		push!(EV,app)
	end

	P::Array{Float64,2} = transpose(T*C*D*G)

	return (P, EV)
end

#-----------------------------------------------------------------------

"""
	dbezier(npts, b[], dpts)

Calculate Bezier curve and its first and second derivatives over the `npts` control polygon points in `b`.

The method evaluates `dpts` points of the curve and gives back the tuple `(P, D1, D2, EV, DEV)`.

Here `P` is the vector of points memorized by columns while `EV` is the `1`-dimensional cellular complex associated to it.
Moreover `D1` and `D2` are the derivatives vectors and `DEV` is the `1`-dimensional cellular complex associated.
Please note that only `dpts-2` points are evaluated for the derivatives as there are not enought informations about the first and the last points.

This structure as it is formed could be directly imported in Plasm module by typing:
```julia
julia> P, _, _, EV, _ = bezier(npts, b[], dpts);
julia> Plasm.view(P, EV)
```

---

# Arguments

- `npts::Int64`: the number of the control polygon vertices.
- `b::Array{Float64, 2}`: 2-dimensional Array containing the control polygon vertices in the x, y, z coordinates for columns.
- `dpts::Int64`: number of data points to be calculated on the curve.

---

# Examples

```jldoctest
julia> dbezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[1]
3×7 Array{Float64,2}:
 1.0  1.56481  2.18519  2.75  3.14815  3.26852  3.0
 1.0  1.83333  2.33333  2.5   2.33333  1.83333  1.0
 0.0  0.0      0.0      0.0   0.0      0.0      0.0
```

```jldoctest
julia> dbezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[2]
3×5 Array{Float64,2}:
 3.66667  3.66667  3.0   1.66667  3.26852
 4.0      2.0      0.0  -2.0      1.83333
 0.0      0.0      0.0   0.0      0.0
```

```jldoctest
julia> dbezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[3]
3×5 Array{Float64,2}:
   2.0   -2.0   -6.0  -10.0  3.26852
 -12.0  -12.0  -12.0  -12.0  1.83333
   0.0    0.0    0.0    0.0  0.0
```

```jldoctest
julia> dbezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[4]
6-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
 [5, 6]
 [6, 7]
```

```jldoctest
julia> dbezier(4,[1.0 2 4 3; 1 3 3 1; 0 0 0 0], 7)[5]
6-element Array{Array{Int64,1},1}:
 [1, 2]
 [2, 3]
 [3, 4]
 [4, 5]
```

---

_By Elia Onofri_
"""

function dbezier(npts::Int64, b::Array{Float64,2}, dpts::Int64)::Tuple{Array{Float64,2}, Array{Float64,2}, Array{Float64,2}, Array{Array{Int64,1},1}, Array{Array{Int64,1},1}}

	@assert npts == length(b[1,:]) ("ERROR: there are not npts = $(npts) points of the polygon in b[]")

	step::Float64 = 1/(dpts-1)
	n::Int64 = npts-1

	C::Array{Float64,2} = zeros(npts, npts)
	for i = 0:n
		for j = 0:n-i
			if (n-j-i)%2==0
				C[i+1, j+1] = binomial(n-j, n-j-i)
			else
				C[i+1, j+1] = 0.0-binomial(n-j, n-j-i)
			end
		end
	end

	D::Array{Float64,2} = zeros(npts, npts)
	for i = 1:npts
		D[i,i] = abs(C[npts+1-i,1])
	end

	G::Array{Float64,2} = transpose(b)

	T::Array{Float64,2} = ones(dpts, npts)

	t::Float64 = 0.0
	for i = 1:dpts
		for j = 0:n-1  #Per j=0 il valore è 1 ed è già impostato
			T[i, n-j] = T[i, npts-j]*t
		end
		t = t + step
	end

	Der1::Array{Float64,2} = ones(dpts-2, npts)
	Der2::Array{Float64,2} = ones(dpts-2, npts)

	t = step
	for i = 2:dpts-2
		den1::Float64 = t * (1 - t)
		den2::Float64 = den1^2
		for j = 0:n
			Der1[i-1, j+1] = (j - n * t) / den1
			Der2[i-1, j+1] = ((j - n * t)^2 - n * t^2 - j * (1 - 2 * t)) / den2
		end
		t = t + step
	end

	EV::Array{Array{Int64,1},1} = Array{Int64,1}[] #1-dimensional cellular complex used for plotting the curve with Plasm package

	for i = 1 : dpts-1
		app = Array{Int64}(2)
		app[1] = i
		app[2] = i+1
		push!(EV,app)
	end

	DEV::Array{Array{Int64,1},1} = EV[1:dpts-3] #1-dimensional cellular complex used for plotting the Derivatives with Plasm package

	F::Array{Float64,2} = T*C*D
	F1::Array{Float64,2} = Der1.*(F[range(2,dpts-2),:])
	F2::Array{Float64,2} = Der2.*(F[range(2,dpts-2),:])
	P::Array{Float64,2} = transpose(F*G)
	D1::Array{Float64,2} = transpose(F1*G)
	D2::Array{Float64,2} = transpose(F2*G)

	return (P, D1, D2, EV, DEV)
end

#-----------------------------------------------------------------------

"""
	bern_basis(n, i, t)

Calculate the Bernstein basis.

This method is meant to evaluate the `i`-th Bernstein Basis of order `n` in the parameter `t`.
Of course this method could be implemented in a more quick way, in order to evaluate all the basis of a specific order
but this is out of the scope of this library.

---

# Arguments

- `n::Int64`: the order of the basis (`n = npts-1`).
- `i::Int64`: the `i`-th basis of the given order. `i ∈ [0, …, n]`.
- `t::Float64`: the parameter of the basis `N_{n, i}(t)`.

---

# Examples

```jldoctest
julia> map(x -> bern_basis(3, 0, x), [0, 0.15, 0.35, 0.5, 0.65, 0.85, 1])
7-element Array{Float64,1}:
 1.0
 0.614125
 0.274625
 0.125
 0.042875
 0.003375
 0.0
```

```jldoctest
julia> map(x -> bern_basis(3, 3, x), [0, 0.15, 0.35, 0.5, 0.65, 0.85, 1])
7-element Array{Float64,1}:
 0.0
 0.003375
 0.042875
 0.125
 0.274625
 0.614125
 1.0
```

---

_By Elia Onofri_
"""

function bern_basis(n::Int64, i::Int64, t::Float64)::Float64
	@assert i >= 0 && i<= n ("ERROR: there are $(n+1) basis of order $(n). You cannot choose the $(i)-th one.")
	return binomial(n, i) * t^i * (1-t)^(n-i)
end


using Plots;
using Printf;
gr();
anim = Animation();
f(x) = 1/2*cos(8pi*(x-1/2));

function forward(u_old::Array, i::Int64)
	return -(2u_old[ir[i]] + 3u_old[i] - 6u_old[il[i]] + u_old[il2[i]])C / 6 + (u_old[ir[i]] - 2u_old[i] + u_old[il[i]])C^2 / 2 - (u_old[ir[i]] - 3u_old[i] + 3u_old[il[i]] - u_old[il2[i]])C^3 / 6;
end

function backward(u_old::Array, i::Int64)
	return (2u_old[il[i]] + 3u_old[i] - 6u_old[ir[i]] + u_old[ir2[i]])*(-C) / 6 + (u_old[il[i]] - 2u_old[i] + u_old[ir[i]])*(-C)^2 / 2 + (u_old[il[i]] - 3u_old[i] + 3u_old[ir[i]] - u_old[ir2[i]])*(-C)^3 / 6;
end

function main()
#=********Numerical Set Up*********=#
	global c = 1.0;
	global nx = 101;
	global lx = 100;
	global x = range(0.0, stop=lx, length=nx);
	global dx = lx /(nx-1);
	global u1 = zeros(Float64, nx);
	global u2 = zeros(Float64, nx);
	@inbounds for i = 1:nx
		if 40.0 < x[i] < 60.0;
			u1[i] = f(x[i]);
		end
	end
	u2 = copy(u1);
	u = zeros(Float64, nx);
	global t = 0.0;
	global dt = 5e-1;
	global tlims = 300;
	global C = dt * c / dx;
	global ir = zeros(Int64, nx);
	global ir2 = zeros(Int64, nx);
	global il = zeros(Int64, nx);
	global il2 = zeros(Int64, nx);
	@inbounds for i = 1:nx
		ir[i] = i + 1;
		ir2[i] = i + 2;
		il[i] = i - 1;
		il2[i] = i - 2;
	end
	ir[end] = 1;
	ir2[end - 1] = 1;
	ir2[end] = 2;
	il[begin] = nx;
	il2[begin + 1] = nx;
	il2[begin] = nx - 1;
#=******Finish Numerical Set Up*************=#

	while t < tlims
		u_old1 = copy(u1);
		u_old2 = copy(u2);
		@inbounds for i = 1:nx
			u1[i] = u_old1[i] + forward(u_old1, i);
			u2[i] = u_old2[i] + backward(u_old2, i);
		end
		u1[begin] = u1[end];
		u2[begin] = u2[end];
		u = (u1 + u2) ./ 2;

		plt = plot(x, u, xlims=(0, lx), xticks=0:lx/5:lx, ylims=(-0.2, 1.2), yticks=-0.2:0.2:1.2,
		xlabel="x", ylabel="u", color="blue", title="QUICKEST with euler", label=@sprintf("t=%.1f", t));
		frame(anim, plt);
		t += dt;
		@printf("t = %.4f\n", t);
	end
	gif(anim, "quickest.gif", fps=5);
end

main();


import Arpack

"""
    LpNormLMO{T, p}(right_hand_side)

LMO with feasible set being an L-p norm ball:
```
C = {x ∈ R^n, norm(x, p) ≤ right_hand_side}
```
"""
struct LpNormLMO{T,p} <: LinearMinimizationOracle
    right_hand_side::T
end

LpNormLMO{p}(right_hand_side::T) where {T,p} = LpNormLMO{T,p}(right_hand_side)

function compute_extreme_point(lmo::LpNormLMO{T,2}, direction; v = similar(direction), kwargs...) where {T}
    dir_norm = norm(direction, 2)
    n = length(direction)
    # if direction numerically 0
    if dir_norm <= 10eps(eltype(direction))
        @. v = lmo.right_hand_side / sqrt(n)
    else
        @. v = -lmo.right_hand_side * direction / dir_norm
    end
    return v
end

function compute_extreme_point(lmo::LpNormLMO{T,Inf}, direction; v = similar(direction), kwargs...) where {T}
    for idx in eachindex(direction)
        v[idx] = -lmo.right_hand_side * (1 - 2signbit(direction[idx]))
    end
    return v
end

function compute_extreme_point(lmo::LpNormLMO{T,1}, direction; v = nothing, kwargs...) where {T}
    idx = 0
    v = -one(eltype(direction))
    for i in eachindex(direction)
        if abs(direction[i]) > v
            v = abs(direction[i])
            idx = i
        end
    end
    sign_coeff = sign(direction[idx])
    if sign_coeff == 0.0
        sign_coeff -= 1
    end
    return ScaledHotVector(-lmo.right_hand_side * sign_coeff, idx, length(direction))
end

function compute_extreme_point(lmo::LpNormLMO{T,p}, direction; v = similar(direction), kwargs...) where {T,p}
    # covers the case where the Inf or 1 is of another type
    if p == Inf
        v = compute_extreme_point(LpNormLMO{T,Inf}(lmo.right_hand_side), direction, v = v)
        return v
    elseif p == 1
        v = compute_extreme_point(LpNormLMO{T,1}(lmo.right_hand_side), direction)
        return v
    end
    q = p / (p - 1)
    pow_ratio = q / p
    q_norm = norm(direction, q)^(pow_ratio)
    # handle zero_direction first
    # assuming the direction is a vector of 1
    if q_norm < eps()
        one_vec = trues(length(direction))
        @. v = -lmo.right_hand_side * one_vec^(pow_ratio) / oftype(q_norm, 1)
        return v
    end
    @. v = -lmo.right_hand_side * sign(direction) * abs(direction)^(pow_ratio) / q_norm
    return v
end

# temporary oracle for l_1 ball to

struct L1ballDense{T} <: LinearMinimizationOracle
    right_hand_side::T
end

function compute_extreme_point(lmo::L1ballDense{T}, direction; v = zeros(T, length(direction)), kwargs...) where {T}
    idx = 0
    val = -1.0
    for i in eachindex(direction)
        if abs(direction[i]) > val
            val = abs(direction[i])
            idx = i
        end
    end

    v .= 0
    v[idx] = T(-lmo.right_hand_side * sign(direction[idx]))
    return v
end

"""
    KNormBallLMO{T}(K::Int, right_hand_side::T)

LMO with feasible set being the K-norm ball in the sense of
[2010.07243](https://arxiv.org/abs/2010.07243),
i.e., the convex hull over the union of an
L_1-ball with radius τ and an L_∞-ball with radius τ/K:
```
C_{K,τ} = conv { B_1(τ) ∪ B_∞(τ / K) }
```
with `τ` the `right_hand_side` parameter. The K-norm is defined as
the sum of the largest `K` absolute entries in a vector.
"""
struct KNormBallLMO{T} <: LinearMinimizationOracle
    K::Int
    right_hand_side::T
end

function compute_extreme_point(lmo::KNormBallLMO{T}, direction; v = similar(direction), kwargs...) where {T}
    K = max(min(lmo.K, length(direction)), 1)

    oinf = zero(eltype(direction))
    idx_l1 = 0
    val_l1 = -one(eltype(direction))

    @inbounds for (i, dir_val) in enumerate(direction)
        temp = -lmo.right_hand_side / K * sign(dir_val)
        v[i] = temp
        oinf += dir_val * temp
        abs_v = abs(dir_val)
        if abs_v > val_l1
            idx_l1 = i
            val_l1 = abs_v
        end
    end

    v1 = ScaledHotVector(-lmo.right_hand_side * sign(direction[idx_l1]), idx_l1, length(direction))
    o1 = dot(v1, direction)
    if o1 < oinf
        @. v = v1
    end
    return v
end

"""
    NuclearNormLMO{T}(radius)

LMO over matrices that have a nuclear norm less than `radius`.
The LMO returns the rank-one matrix with singular value `radius`.
"""
struct NuclearNormLMO{T} <: LinearMinimizationOracle
    radius::T
end

NuclearNormLMO{T}() where {T} = NuclearNormLMO{T}(one(T))
NuclearNormLMO() = NuclearNormLMO(1.0)

"""
Best rank-one approximation using the greatest singular value computed with Arpack.

Warning: this does not work (yet) with all number types, BigFloat and Float16 fail.
"""
function compute_extreme_point(lmo::NuclearNormLMO{TL}, direction::AbstractMatrix{TD}; tol=1e-8, kwargs...) where {TL, TD}
    T = promote_type(TD, TL)
    Z = Arpack.svds(direction, nsv=1, tol=tol)[1]
    u = -lmo.radius * view(Z.U, :)
    return RankOneMatrix(u::Vector{T}, Z.V[:]::Vector{T})
end

function convert_mathopt(
    lmo::NuclearNormLMO,
    optimizer::OT;
    row_dimension::Integer,
    col_dimension::Integer,
    use_modify=true::Bool,
    kwargs...,
) where {OT}
    MOI.empty!(optimizer)
    x = MOI.add_variables(optimizer, row_dimension * col_dimension)
    (t, _) = MOI.add_constrained_variable(optimizer, MOI.LessThan(lmo.radius))
    MOI.add_constraint(optimizer, [t; x], MOI.NormNuclearCone(row_dimension, col_dimension))
    return MathOptLMO(optimizer, use_modify)
end

"""
    SpectraplexLMO{T,M}(radius::T,gradient_container::M,ensure_symmetry::Bool=true)

Feasible set
```
{X ∈ 𝕊_n^+, trace(X) == radius}
```
`gradient_container` is used to store the symmetrized negative direction.
`ensure_symmetry` indicates whether the linear function is made symmetric before computing the eigenvector.
"""
struct SpectraplexLMO{T,M} <: LinearMinimizationOracle
    radius::T
    gradient_container::M
    ensure_symmetry::Bool
end

function SpectraplexLMO(radius::T, side_dimension::Int, ensure_symmetry::Bool=true) where {T}
    return SpectraplexLMO(
        radius,
        Matrix{T}(undef, side_dimension, side_dimension),
        ensure_symmetry,
    )
end

function SpectraplexLMO(radius::Integer, side_dimension::Int, ensure_symmetry::Bool=true)
    return SpectraplexLMO(float(radius), side_dimension, ensure_symmetry)
end

function compute_extreme_point(lmo::SpectraplexLMO{T}, direction::M; v = nothing, maxiters=500, kwargs...) where {T,M <: AbstractMatrix}
    lmo.gradient_container .= direction
    if !(M <: Union{LinearAlgebra.Symmetric, LinearAlgebra.Diagonal, LinearAlgebra.UniformScaling}) && lmo.ensure_symmetry
        # make gradient symmetric
        @. lmo.gradient_container += direction'
    end
    lmo.gradient_container .*= -1

    _, evec = Arpack.eigs(lmo.gradient_container; nev=1, which=:LR, maxiter=maxiters)
    # type annotation because of Arpack instability
    unit_vec::Vector{T} = vec(evec)
    # scaling by sqrt(radius) so that x x^T has spectral norm radius while using a single vector
    unit_vec .*= sqrt(lmo.radius)
    return FrankWolfe.RankOneMatrix(unit_vec, unit_vec)
end

"""
    UnitSpectrahedronLMO{T,M}(radius::T, gradient_container::M)

Feasible set of PSD matrices with bounded trace:
```
{X ∈ 𝕊_n^+, trace(X) ≤ radius}
```
`gradient_container` is used to store the symmetrized negative direction.
`ensure_symmetry` indicates whether the linear function is made symmetric before computing the eigenvector.
"""
struct UnitSpectrahedronLMO{T,M} <: LinearMinimizationOracle
    radius::T
    gradient_container::M
    ensure_symmetry::Bool
end

function UnitSpectrahedronLMO(radius::T, side_dimension::Int, ensure_symmetry::Bool=true) where {T}
    return UnitSpectrahedronLMO(
        radius,
        Matrix{T}(undef, side_dimension, side_dimension),
        ensure_symmetry,
    )
end

UnitSpectrahedronLMO(radius::Integer, side_dimension::Int) = UnitSpectrahedronLMO(float(radius), side_dimension)

function compute_extreme_point(lmo::UnitSpectrahedronLMO{T}, direction::M; v = nothing, maxiters=500, kwargs...) where {T, M <: AbstractMatrix}
    lmo.gradient_container .= direction
    if !(M <: Union{LinearAlgebra.Symmetric, LinearAlgebra.Diagonal, LinearAlgebra.UniformScaling}) && lmo.ensure_symmetry
        # make gradient symmetric
        @. lmo.gradient_container += direction'
    end
    lmo.gradient_container .*= -1

    e_val::Vector{T}, evec::Matrix{T} = Arpack.eigs(lmo.gradient_container; nev=1, which=:LR, maxiter=maxiters)
    # type annotation because of Arpack instability
    unit_vec::Vector{T} = vec(evec)
    if e_val[1] < 0
        # return a zero rank-one matrix
        unit_vec .*= 0
    else
        # scaling by sqrt(radius) so that x x^T has spectral norm radius while using a single vector
        unit_vec .*= sqrt(lmo.radius)
    end
    return FrankWolfe.RankOneMatrix(unit_vec, unit_vec)
end

function convert_mathopt(
    lmo::Union{SpectraplexLMO{T}, UnitSpectrahedronLMO{T}},
    optimizer::OT;
    side_dimension::Integer,
    use_modify::Bool=true,
    kwargs...,
) where {T, OT}
    MOI.empty!(optimizer)
    X = MOI.add_variables(optimizer, side_dimension * side_dimension)
    MOI.add_constraint(optimizer, X, MOI.PositiveSemidefiniteConeSquare(side_dimension))
    sum_diag_terms = MOI.ScalarAffineFunction{T}([],zero(T))
    # collect diagonal terms of the matrix
    for i in 1:side_dimension
        push!(sum_diag_terms.terms, MOI.ScalarAffineTerm(one(T), X[i + side_dimension * (i-1)]))
    end
    constraint_set = if lmo isa SpectraplexLMO
        MOI.EqualTo(lmo.radius)
    else
        MOI.LessThan(lmo.radius)
    end
    MOI.add_constraint(optimizer, sum_diag_terms, constraint_set)
    return MathOptLMO(optimizer, use_modify)
end


using Plots
pyplot()
x = range(0, stop=2*pi, length=50)
println(x)
p = plot(x = x, y = sin(x), size =(900,300))

plot(x = x, y = sin(x))

plot(rand(100) / 3,reg=true,fill=(0,:green))
scatter!(rand(100),markersize=6,c=:orange)

plot(rand(100) / 3,reg=true,fill=(0,:green))
scatter!(rand(100),markersize=6,c=:orange)

plot(y = cumsum(randn(1000, 1)))


using Test
using SymbolicUtils: <ₑ
SymbolicUtils.show_simplified[] = false

@syms a b c

function istotal(x,y)
    #either
    if x <ₑ y
        return !(y <ₑ x)
    elseif y <ₑ x
        return !(x <ₑ y) # already tested
    else
        # neither, equal
        return true
    end
end

@test istotal(a,a)
@test istotal(a,b)
@test istotal(2,a)
@test 2 <ₑ a
@test a <ₑ b
@test istotal(a, 2a)
@test 2a <ₑ a
@test istotal(b*a, a)
@test a <ₑ b*a
@test !(b*a <ₑ b+a)

@testset "operator order" begin
    fs = (*, ^, /, \, -, +)
    for i in 1:length(fs)
        f = fs[i]
        @test f(a, b) <ₑ f(b, c)
        @test istotal(f(a, b), f(b, c))
        @test !(f(b, b) <ₑ f(b, b))
        @test istotal(f(b, b), f(b, b))
        @test !(f(b, c) <ₑ f(a, b))
        @test istotal(f(b, c), f(a, b))

        @test f(1, b) <ₑ f(2, b)
        @test !(f(2, b) <ₑ f(1, b))
        @test istotal(f(1, b), f(2, b))
        @test istotal(f(2, b), f(1, b))
        @test !(b <ₑ f(2,b))
        for j in i+1:length(fs)
            g = fs[j]
            @test !(g(a, b) <ₑ f(a, b)) && !(f(a, b) <ₑ g(a, b))
            @test istotal(f(a, b), g(a, b))
        end
    end
end


using Peacock, PyPlot, LinearAlgebra
orthonormalise = Peacock.orthonormalise
get_k = Peacock.get_k
sample_path = Peacock.sample_path

a = 1; a1 = [1, 0]; a2 = [0, 1]
w = 0.17 # 柱子边长
d = w/2 # w/2 中间， a/2 - w/2 四周 a/4 中心

eps0 = 1; mu0 = 1
mu1 = [14 12.4im 0; -12.4im 14 0; 0 0 1]
eps1 = 15

function epf(x, y)
    if abs(x - d) < w/2 && abs(y - d) < w/2 || abs(x + d) < w/2 && abs(y - d) < w/2 || abs(x - d) < w/2 && abs(y + d) < w/2 || abs(x + d) < w/2 && abs(y + d) < w/2
        return eps1
    else
        return eps0
    end
end

function muf(x, y)
    if abs(x - d) < w/2 && abs(y - d) < w/2 || abs(x + d) < w/2 && abs(y - d) < w/2 || abs(x - d) < w/2 && abs(y + d) < w/2 || abs(x + d) < w/2 && abs(y + d) < w/2
        return mu1
    else
        return mu0
    end
end

#epf(x, y) = x^2 + y^2 < 0.2^2 ? 15 : 1
#muf(x, y) = x^2 + y^2 < 0.2^2 ? 1 : 1

d1 = d2 = 0.005
geometry = Geometry(epf, muf, a1, a2, d1, d2, TM)
solver = Peacock.FDFD.Solver(geometry, 2*[d1, d2])

kx = -pi:0.1:pi
ky = -pi:0.1:pi
modes = Peacock.FDFD.solve(solver,[0,0],TM,bands=1:1)
weighting = modes[1].weighting
data = zeros(ComplexF64, length(kx),length(ky)+1,2,length(weighting))
i = 1
for kx in kx
    j = 1
    for ky in ky
        for n = 1:2
            modes = Peacock.FDFD.solve(solver,[kx,ky],TM,bands=1:2)
            data0 = get_field_FDFD(modes[n].data, modes[n].basis,k0=modes[n].k0)
            data0 = data0[:] / sqrt(abs(dot(data0[:], weighting.*data0[:])))
            data[i,j,n,:] = data0[:]
        end
        j = j + 1
    end
    data[i,end,1,:] = data[i,1,1,:]
    data[i,end,2,:] = data[i,1,2,:]
    i = i + 1
end

vy = zeros(ComplexF64, length(kx),length(ky),2)
W = zeros(ComplexF64, 2,2)
states = zeros(ComplexF64, length(kx),length(ky),2,2)
for i = 1:length(kx)
    for j1 = 1:length(ky)
        F = I
        for j2 = 0:length(ky)-1
            if j1 + j2 <= length(ky)
                j = j1+j2
            else
                j = j1 + j2 - length(ky)
            end
            for k = 1:2
                for l = 1:2
                    W[k,l] = dot(data[i,j,k,:],weighting.*data[i,j+1,l,:])
                end
            end
            F = F * W
        end
        vals,state = eigen(F,sortby=angle)
        vy[i,j1,:] = angle.(vals)
        states[i,j1,:,:] = state
    end
end

wx = zeros(ComplexF64, length(kx)+1,length(ky),length(weighting))
for j = 1:length(ky) #ky
    for i = 1:length(kx) #kx
        wx[i,j,:] = states[i,j,1,1] * data[i,j,1,:] + states[i,j,2,1] * data[i,j,2,:]
        wx[i,j,:] = wx[i,j,:]/sqrt(abs(dot(wx[i,j,:],weighting.*wx[i,j,:])))
    end
    wx[end,j,:] = wx[1,j,:]
end

px = zeros(length(ky))
for j = 1:length(ky) #ky
    W = 1
    for i = 1:length(kx) #kx
        W = W*dot(wx[i,j,:],weighting.*wx[i+1,j,:])
    end
#    vals = eigvals(W)
    px[j] = angle.(W)
#    px[j] = imag(log(W/2/pi))
end
plot(px,"*")
ylim(-pi,pi)
#figure()
#plot(vy,"*")


using Test
using PolygonArea
using PolygonArea: PolarHalfPlane, Point, corner_point, Intersection, Reunion
using PolygonArea: signed_distance, distance, outward_normal, angle
using PolygonArea: complement, rotate, translate, distance, exchange_x_and_y
using StaticArrays

@testset "Half-planes" begin

    @testset "creation and equality" begin
        @test HalfPlane(1.0, 0.0, 0.0) isa HalfPlane{Float64}
        @test PolarHalfPlane(0.0, π, center=(1.0, 1.0)) isa HalfPlane{Float64}

        @test HalfPlane(1, 0, 0) isa HalfPlane{Int}
        @test HalfPlane{Float64}(1, 0, 0) isa HalfPlane{Float64}

        @test HalfPlane((0.0, 0.0), (1.0, 0.0)) ≈ HalfPlane(0.0, 1.0, 0.0)

        @test HalfPlane(1.0, 0.0, 0.0) == HalfPlane(1.0, 0.0, 0.0)
        @test HalfPlane(1.0, 0.0, 0.0) != HalfPlane(1.0, 1e-15, 0.0)
        @test HalfPlane(1.0, 0.0, 0.0) ≈ HalfPlane(1.0, 0.0, 0.0)
    end

    @testset "contained points" begin
        @test Point(0, -1)   in HalfPlane(0, 1, 0)
        @test Point(0, -1)   in HalfPlane(0.0, 1.0, 0.0)
        @test Point(0, 0)    in HalfPlane(0, 1, 0)
        @test !(Point(0, 1)  in HalfPlane(0, 1, 0))
        @test Point(1, 0)    in HalfPlane(-1, 0, 0)
        @test !(Point(-1, 0) in HalfPlane(-1, 0, 0))

        @test (-1.0, -1.0) ∈ HalfPlane((0.0, 0.0), (1.0, 0.0), :right)
        @test (1.0, 1.0) ∈ HalfPlane((0.0, 0.0), (1.0, 0.0), :left)

        @test Point(0.5, 0.5) in PolarHalfPlane(-1.0, 0π)
        @test (0.5, 0.5) in PolarHalfPlane(-1.0, 0π)

        @test Point(0.5, 0.5) in PolarHalfPlane(-1.0, π/2)
        @test Point(0.5, 0.5) in PolarHalfPlane(-1.0, π)
        @test Point(0.5, 0.5) in PolarHalfPlane(-1.0, 3π/2)
        @test Point(1.0, 1.0) in PolarHalfPlane(-1.0, 0π)
        @test Point(0.5, 2.0) in PolarHalfPlane(-1.0, 0π)
        @test Point(2.0, 0.5) in PolarHalfPlane(-1.0, π/2)
        @test Point(0.0, 3.0) in PolarHalfPlane(-1.0, π)
    end

    @testset "properties" begin
        @test signed_distance((0.0, 0.0), PolarHalfPlane(-4.0, 2π/3)) ≈ -4.0
        @test signed_distance((0.0, 0.0), PolarHalfPlane(4.0, π/3)) ≈ 4.0

        # Works with Points (SVector) as well as tuples
        @test distance(Point(9.0, 0.0), PolarHalfPlane(0.0, 0.0)) ≈ 9.0
        @test distance((9.0, 0.0), PolarHalfPlane(0.0, 0.0)) ≈ 9.0

        # Scaling of the definition of the plane
        @test distance((1.0, 1.0), HalfPlane(-1.0, 0.0, 0.0)) == distance((1.0, 1.0), HalfPlane(-2.0, 0.0, 0.0))
        @test distance((1.0, 1.0), HalfPlane(-1.0, -1.0, 0.0)) == distance((1.0, 1.0), HalfPlane(-2.0, -2.0, 0.0))

        @test distance((9.0, 0.0), PolarHalfPlane(1.0, 0.0)) ≈ 10.0
        @test distance((9.0, -10.0), PolarHalfPlane(1.0, 0.0)) ≈ 10.0
        @test distance((9.0, 1.0), PolarHalfPlane(0.0, π/2)) ≈ 1.0
        @test distance((1.0, 1.0), PolarHalfPlane(0.0, π/4)) ≈ √2
        @test distance((2.0, 2.0), PolarHalfPlane(-√2, π/4)) ≈ √2
        @test distance((6.0, 1.0), PolarHalfPlane(0.0, π/4, center=Point(5.0, 0.0))) ≈ √2

        @test angle(PolarHalfPlane(0.0, π/8)) ≈ π/8
        @test angle(PolarHalfPlane(0.0, -π/8)) ≈ mod(-π/8, 2π)
        @test angle(PolarHalfPlane(4.0, π/3)) ≈ π/3

        @test outward_normal(HalfPlane(1.0, 0.0, 0.0)) == SVector(1.0, 0.0)
        @test outward_normal(HalfPlane(0.0, 1.0, rand(1)[1])) == SVector(0.0, 1.0)
    end

    @testset "transformations" begin
        # Rotate
        @test outward_normal(rotate(PolarHalfPlane(0π), π)) == outward_normal(PolarHalfPlane(π))
        @test outward_normal(rotate(PolarHalfPlane(0π, center=(1.0, 1.0)), π)) == outward_normal(PolarHalfPlane(π))
        @test outward_normal(rotate(PolarHalfPlane(π/2, center=(1.0, 1.0)), π)) == outward_normal(PolarHalfPlane(3π/2))

        @test (PolarHalfPlane(0.234π) |> rotate(π) |> rotate(-π)) ≈ PolarHalfPlane(0.234π)
        @test (PolarHalfPlane(0.234π) |> rotate(π) |> rotate(π)) ≈ PolarHalfPlane(0.234π)

        @test Point(1.0, 0.0) in rotate(HalfPlane(1.0, 0.0, 0.0) ∪ HalfPlane(1.0, 0.1, 0.0), π)
        @test Point(1.0, 0.0) in rotate(HalfPlane(1.0, 0.0, 0.0) ∩ HalfPlane(1.0, 0.1, 0.0), π)

        # Translate
        @test outward_normal(translate(PolarHalfPlane(0π), SVector{2}(1.0, 1.0))) == outward_normal(PolarHalfPlane(0π))
        @test distance(Point(2.0, 0.0), translate(PolarHalfPlane(0π), SVector{2}(1.0, 1.0))) == 1.0
        @test Point(1.0, 0.0) in translate(HalfPlane(1.0, 0.0, 0.0), SVector(2.0, 0.0))

        @test Point(1.0, 0.0) in translate(HalfPlane(1.0, 0.0, 0.0) ∪ HalfPlane(1.0, 0.1, 0.0), SVector(2.0, 0.0))
        @test Point(1.0, 0.0) in translate(HalfPlane(1.0, 0.0, 0.0) ∩ HalfPlane(1.0, 0.1, 0.0), SVector(2.0, 0.0))

        @test translate(HalfPlane(1.0, 0.0, 0.0), SVector(2.0, 0.0)) == (HalfPlane(1.0, 0.0, 0.0) |> translate(SVector(2.0, 0.0)))

        # Exchange x and y
        @test isapprox(exchange_x_and_y(PolarHalfPlane(0.0, 0.0)), PolarHalfPlane(0.0, π/2), atol=1e-14)
        @test isapprox(exchange_x_and_y(PolarHalfPlane(1.0, 0.0)), PolarHalfPlane(1.0, π/2), atol=1e-14)
        @test isapprox(exchange_x_and_y(PolarHalfPlane(0.0, 0.0, center=Point(1.0, 0.0))), PolarHalfPlane(0.0, π/2, center=Point(0.0, 1.0)), atol=1e-14)
    end

    @testset "unions, intersections and corners" begin
        inferior_hp = HalfPlane(0, 1, 0)
        right_hp = HalfPlane(-1, 0, 0)
        @test Point(1, -1) in intersect(inferior_hp, right_hp)
        @test Point(1, -1) in inferior_hp ∩ right_hp
        @test !(Point(1, 1) in intersect(inferior_hp, right_hp))
        @test Point(1, 1) in complement(inferior_hp ∩ right_hp)
        @test Point(1, 1) in complement(inferior_hp ∪ complement(right_hp))

        @test !(Point(-1, 1) in union(inferior_hp, right_hp))
        @test (Point(1, 1) in union(inferior_hp, right_hp))

        @test corner_point(HalfPlane(0.0, 1.0, 0.0), HalfPlane(-1.0, 0.0, 0.0)) == Point(0.0, 0.0)
        @test corner_point(HalfPlane(1.0, 0.0, 1.0), HalfPlane(0.0, -1.0, 0.5)) == Point(-1.0, 0.5)
        @test corner_point(PolarHalfPlane(0.0, 0.0), HalfPlane(0.0, 1.0, -1.0)) == Point(0.0, 1.0)
        @test corner_point(PolarHalfPlane(-1.0, π/2, center=Point(0.5, 0.5)), HalfPlane(1.0, 0.0, -1.0)) == Point(1.0, 1.5)

        # Union and intersection
        h1 = HalfPlane(1.0, 0.0, 0.0)
        @test ((h1 ∩ h1) ∪ (h1 ∩ h1)) isa Reunion{Intersection{HalfPlane{Float64}}}
        @test ((h1 ∪ h1) ∩ (h1 ∪ h1)) isa Reunion{Intersection{HalfPlane{Float64}}}

        HalfPlane(1.0, 0.0, 0.0) ∪ HalfPlane(0.0, 1.0, 0.0) ∩ (HalfPlane(1.0, 0.0, 0.0) ∪ HalfPlane(0.0, 1.0, 0.0))

        @test Point(0.0, 0.0) in ((h1 ∪ h1) ∩ (h1 ∪ h1) |> complement)
    end

    @testset "conversion" begin
        a = HalfPlane{Int}(1, 1, 1)
        b = HalfPlane{Int}(-1, 1, 1)
        @test convert(Reunion{Intersection{HalfPlane{Int}}}, a ∪ b) isa Reunion{Intersection{HalfPlane{Int}}}
    end

end


"""
    FVectorvariate

Structure that subtypes `VariateForm`, indicating that a single sample
is a vector from a fiber of a vector bundle.
"""
struct FVectorvariate <: VariateForm end

"""
    FVectorSupport(space::Manifold, VectorBundleFibers)

Value support for vector bundle fiber-valued distributions (values from a fiber of a vector
bundle at a `point` from the given manifold).
For example used for tangent vector-valued distributions.
"""
struct FVectorSupport{TSpace<:VectorBundleFibers,T} <: ValueSupport
    space::TSpace
    point::T
end

"""
    FVectorDistribution{TSpace<:VectorBundleFibers, T}

An abstract distribution for vector bundle fiber-valued distributions (values from a fiber
of a vector bundle at point `x` from the given manifold).
For example used for tangent vector-valued distributions.
"""
abstract type FVectorDistribution{TSpace<:VectorBundleFibers,T} <:
              Distribution{FVectorvariate,FVectorSupport{TSpace,T}} end

"""
    MPointvariate

Structure that subtypes `VariateForm`, indicating that a single sample
is a point on a manifold.
"""
struct MPointvariate <: VariateForm end

"""
    MPointSupport(M::Manifold)

Value support for manifold-valued distributions (values from given
[`Manifold`](@ref) `M`).
"""
struct MPointSupport{TM<:Manifold} <: ValueSupport
    manifold::TM
end

"""
    MPointDistribution{TM<:Manifold}

An abstract distribution for points on manifold of type `TM`.
"""
abstract type MPointDistribution{TM<:Manifold} <:
              Distribution{MPointvariate,MPointSupport{TM}} end

"""
    support(d::FVectorDistribution)

Get the object of type `FVectorSupport` for the distribution `d`.
"""
function Distributions.support(::T) where {T<:FVectorDistribution}
    return error("support not implemented for type $T")
end

@decorator_transparent_signature normal_tvector_distribution(
    M::AbstractDecoratorManifold,
    p,
    σ,
)

@decorator_transparent_signature projected_distribution(M::AbstractDecoratorManifold, d, p)
@decorator_transparent_signature projected_distribution(M::AbstractDecoratorManifold, d)


using PyPlot

AF = A - B1 * F1 - B2 * F2
n = 20
x = Array{Float64}(3, n)
x[:, 1] = [1 1 1]
for t in 1:n-1
    x[:, t+1] = AF * x[:, t]
end
q1 = x[2, :]
q2 = x[3, :]
q = q1 + q2         # Total output, MPE
p = a0 - a1 * q     # Price, MPE

fig, ax = subplots(figsize=(9, 5.8))
ax[:plot](q, "b-", lw=2, alpha=0.75, label="total output")
ax[:plot](p, "g-", lw=2, alpha=0.75, label="price")
ax[:set_title]("Output and prices, duopoly MPE")
ax[:legend](frameon=false)