""" Simplicial sets in one, two, and three dimensions.

For the time being, this module provides data structures only for [delta
sets](https://en.wikipedia.org/wiki/Delta_set), also known as [semi-simplicial
sets](https://ncatlab.org/nlab/show/semi-simplicial+set). These include the face
maps but not the degeneracy maps of a simplicial set. In the future we may add
support for simplicial sets. The analogy to keep in mind is that graphs are to
semi-simpicial sets as reflexive graphs are to simplicial sets.

Also provided are the fundamental operators on simplicial sets used in nearly
all geometric applications, namely the boundary and coboundary (discrete
exterior derivative) operators. For additional operators, see the
`DiscreteExteriorCalculus` module.
"""
module SimplicialSets
export Simplex, V, E, Tri, SimplexChain, VChain, EChain, TriChain,
  SimplexForm, VForm, EForm, TriForm,
  HasDeltaSet, HasDeltaSet1D, HasDeltaSet2D,
  AbstractDeltaSet1D, DeltaSet1D, OrientedDeltaSet1D, EmbeddedDeltaSet1D,
  AbstractDeltaSet2D, DeltaSet2D, OrientedDeltaSet2D, EmbeddedDeltaSet2D,
  ∂, boundary, coface, d, coboundary, exterior_derivative,
  simplices, nsimplices, point, volume,
  orientation, set_orientation!, orient!, orient_component!,
  src, tgt, nv, ne, vertices, edges, has_vertex, has_edge, edge_vertices,
  add_vertex!, add_vertices!, add_edge!, add_edges!,
  add_sorted_edge!, add_sorted_edges!,
  triangle_edges, triangle_vertices, ntriangles, triangles,
  add_triangle!, glue_triangle!, glue_sorted_triangle!

using LinearAlgebra: det
using SparseArrays
using StaticArrays: @SVector, SVector, SMatrix

using Catlab, Catlab.CategoricalAlgebra, Catlab.Graphs
import Catlab.Graphs: src, tgt, nv, ne, vertices, edges, has_vertex, has_edge,
  add_vertex!, add_vertices!, add_edge!, add_edges!
using ..ArrayUtils

# 0-D simplicial sets
#####################

@present DeltaCategory0D(FreeSchema) begin
  V::Ob
end

""" Abstract type for C-sets that contain a delta set of some dimension.

This dimension could be zero, in which case the delta set consists only of
vertices (0-simplices).
"""
@abstract_acset_type HasDeltaSet

vertices(s::HasDeltaSet) = parts(s, :V)
nv(s::HasDeltaSet) = nparts(s, :V)
nsimplices(::Type{Val{0}}, s::HasDeltaSet) = nv(s)

has_vertex(s::HasDeltaSet, v) = has_part(s, :V, v)
add_vertex!(s::HasDeltaSet; kw...) = add_part!(s, :V; kw...)
add_vertices!(s::HasDeltaSet, n::Int; kw...) = add_parts!(s, :V, n; kw...)

# 1D simplicial sets
####################

@present DeltaCategory1D <: DeltaCategory0D begin
  E::Ob
  (∂v0, ∂v1)::Hom(E, V) # (∂₁(0), ∂₁(1))
end

""" Abstract type for C-sets that contain a one-dimensional delta set.
"""
@abstract_acset_type HasDeltaSet1D <: HasDeltaSet

""" Abstract type for one-dimensional delta sets, aka semi-simplicial sets.
"""
@abstract_acset_type AbstractDeltaSet1D <: HasDeltaSet1D

""" A one-dimensional delta set, aka semi-simplicial set.

Delta sets in 1D are isomorphic to graphs (in the category theorist's sense).
The source and target of an edge can be accessed using the face maps [`∂`](@ref)
(simplicial terminology) or `src` and `tgt` maps (graph-theoretic terminology).
More generally, this type implements the graphs interface in `Catlab.Graphs`.
"""
@acset_type DeltaSet1D(DeltaCategory1D, index=[:∂v0,:∂v1]) <: AbstractDeltaSet1D

edges(s::HasDeltaSet1D) = parts(s, :E)
edges(s::HasDeltaSet1D, src::Int, tgt::Int) =
  (e for e in coface(1,1,s,src) if ∂(1,0,s,e) == tgt)

ne(s::HasDeltaSet1D) = nparts(s, :E)
nsimplices(::Type{Val{1}}, s::HasDeltaSet1D) = ne(s)

has_edge(s::HasDeltaSet1D, e) = has_part(s, :E, e)
has_edge(s::HasDeltaSet1D, src::Int, tgt::Int) =
  has_vertex(s, src) && any(e -> ∂(1,0,s,e) == tgt, coface(1,1,s,src))

src(s::HasDeltaSet1D, args...) = subpart(s, args..., :∂v1)
tgt(s::HasDeltaSet1D, args...) = subpart(s, args..., :∂v0)
face(::Type{Val{(1,0)}}, s::HasDeltaSet1D, args...) = subpart(s, args..., :∂v0)
face(::Type{Val{(1,1)}}, s::HasDeltaSet1D, args...) = subpart(s, args..., :∂v1)

coface(::Type{Val{(1,0)}}, s::HasDeltaSet1D, args...) = incident(s, args..., :∂v0)
coface(::Type{Val{(1,1)}}, s::HasDeltaSet1D, args...) = incident(s, args..., :∂v1)

""" Boundary vertices of an edge.
"""
edge_vertices(s::HasDeltaSet1D, e...) = SVector(∂(1,0,s,e...), ∂(1,1,s,e...))

add_edge!(s::HasDeltaSet1D, src::Int, tgt::Int; kw...) =
  add_part!(s, :E; ∂v1=src, ∂v0=tgt, kw...)

function add_edges!(s::HasDeltaSet1D, srcs::AbstractVector{Int},
                    tgts::AbstractVector{Int}; kw...)
  @assert (n = length(srcs)) == length(tgts)
  add_parts!(s, :E, n; ∂v1=srcs, ∂v0=tgts, kw...)
end

""" Add edge to simplicial set, respecting the order of the vertex IDs.
"""
add_sorted_edge!(s::HasDeltaSet1D, v₀::Int, v₁::Int; kw...) =
  add_edge!(s, min(v₀, v₁), max(v₀, v₁); kw...)

""" Add edges to simplicial set, respecting the order of the vertex IDs.
"""
function add_sorted_edges!(s::HasDeltaSet1D, vs₀::AbstractVector{Int},
                           vs₁::AbstractVector{Int}; kw...)
  add_edges!(s, min.(vs₀, vs₁), max.(vs₀, vs₁); kw...)
end

# 1D oriented simplicial sets
#----------------------------

@present OrientedDeltaSchema1D <: DeltaCategory1D begin
  Orientation::AttrType
  edge_orientation::Attr(E,Orientation)
end

""" A one-dimensional oriented delta set.

Edges are oriented from source to target when `edge_orientation` is
true/positive and from target to source when it is false/negative.
"""
@acset_type OrientedDeltaSet1D(OrientedDeltaSchema1D,
                               index=[:∂v0,:∂v1]) <: AbstractDeltaSet1D

orientation(::Type{Val{1}}, s::HasDeltaSet1D, args...) =
  s[args..., :edge_orientation]
set_orientation!(::Type{Val{1}}, s::HasDeltaSet1D, e, orientation) =
  (s[e, :edge_orientation] = orientation)

function ∂_nz(::Type{Val{1}}, s::HasDeltaSet1D, e::Int)
  (edge_vertices(s, e), sign(1,s,e) * @SVector([1,-1]))
end

function d_nz(::Type{Val{0}}, s::HasDeltaSet1D, v::Int)
  e₀, e₁ = coface(1,0,s,v), coface(1,1,s,v)
  (lazy(vcat, e₀, e₁), lazy(vcat, sign(1,s,e₀), -sign(1,s,e₁)))
end

# 1D embedded simplicial sets
#----------------------------

@present EmbeddedDeltaSchema1D <: OrientedDeltaSchema1D begin
  Point::AttrType
  point::Attr(V, Point)
end

""" A one-dimensional, embedded, oriented delta set.
"""
@acset_type EmbeddedDeltaSet1D(EmbeddedDeltaSchema1D,
                               index=[:∂v0,:∂v1]) <: AbstractDeltaSet1D

""" Point associated with vertex of complex.
"""
point(s::HasDeltaSet, args...) = s[args..., :point]

struct CayleyMengerDet end

volume(::Type{Val{n}}, s::EmbeddedDeltaSet1D, x) where n =
  volume(Val{n}, s, x, CayleyMengerDet())
volume(::Type{Val{1}}, s::HasDeltaSet1D, e::Int, ::CayleyMengerDet) =
  volume(point(s, edge_vertices(s, e)))

# 2D simplicial sets
####################

@present DeltaCategory2D <: DeltaCategory1D begin
  Tri::Ob
  (∂e0, ∂e1, ∂e2)::Hom(Tri,E) # (∂₂(0), ∂₂(1), ∂₂(2))

  # Simplicial identities.
  ∂e1 ⋅ ∂v1 == ∂e2 ⋅ ∂v1 # ∂₂(1) ⋅ ∂₁(1) == ∂₂(2) ⋅ ∂₁(1) == v₀
  ∂e0 ⋅ ∂v1 == ∂e2 ⋅ ∂v0 # ∂₂(0) ⋅ ∂₁(1) == ∂₂(2) ⋅ ∂₁(0) == v₁
  ∂e0 ⋅ ∂v0 == ∂e1 ⋅ ∂v0 # ∂₂(0) ⋅ ∂₁(0) == ∂₂(1) ⋅ ∂₁(0) == v₂
end

""" Abstract type for C-sets containing a 2D delta set.
"""
@abstract_acset_type HasDeltaSet2D <: HasDeltaSet1D

""" Abstract type for 2D delta sets.
"""
@abstract_acset_type AbstractDeltaSet2D <: HasDeltaSet2D

""" A 2D delta set, aka semi-simplicial set.

The triangles in a semi-simpicial set can be interpreted in several ways.
Geometrically, they are triangles (2-simplices) whose three edges are directed
according to a specific pattern, determined by the ordering of the vertices or
equivalently by the simplicial identities. This geometric perspective is present
through the subpart names `∂e0`, `∂e1`, and `∂e2` and through the boundary map
[`∂`](@ref). Alternatively, the triangle can be interpreted as a
higher-dimensional link or morphism, going from two edges in sequence (which
might be called `src2_first` and `src2_last`) to a transitive edge (say `tgt2`).
This is the shape of the binary composition operation in a category.
"""
@acset_type DeltaSet2D(DeltaCategory2D,
                       index=[:∂v0,:∂v1,:∂e0,:∂e1,:∂e2]) <: AbstractDeltaSet2D

triangles(s::HasDeltaSet2D) = parts(s, :Tri)
ntriangles(s::HasDeltaSet2D) = nparts(s, :Tri)
nsimplices(::Type{Val{2}}, s::HasDeltaSet2D) = ntriangles(s)

face(::Type{Val{(2,0)}}, s::HasDeltaSet2D, args...) = subpart(s, args..., :∂e0)
face(::Type{Val{(2,1)}}, s::HasDeltaSet2D, args...) = subpart(s, args..., :∂e1)
face(::Type{Val{(2,2)}}, s::HasDeltaSet2D, args...) = subpart(s, args..., :∂e2)

coface(::Type{Val{(2,0)}}, s::HasDeltaSet2D, args...) = incident(s, args..., :∂e0)
coface(::Type{Val{(2,1)}}, s::HasDeltaSet2D, args...) = incident(s, args..., :∂e1)
coface(::Type{Val{(2,2)}}, s::HasDeltaSet2D, args...) = incident(s, args..., :∂e2)

""" Boundary edges of a triangle.
"""
function triangle_edges(s::HasDeltaSet2D, t...)
  SVector(∂(2,0,s,t...), ∂(2,1,s,t...), ∂(2,2,s,t...))
end

""" Boundary vertices of a triangle.

This accessor assumes that the simplicial identities hold.
"""
function triangle_vertices(s::HasDeltaSet2D, t...)
  SVector(s[s[t..., :∂e1], :∂v1],
          s[s[t..., :∂e2], :∂v0],
          s[s[t..., :∂e1], :∂v0])
end

""" Add a triangle (2-simplex) to a simplicial set, given its boundary edges.

In the arguments to this function, the boundary edges have the order ``0 → 1``,
``1 → 2``, ``0 → 2``.

!!! warning

    This low-level function does not check the simplicial identities. It is your
    responsibility to ensure they are satisfied. By contrast, triangles added
    using the function [`glue_triangle!`](@ref) always satisfy the simplicial
    identities, by construction. Thus it is often easier to use this function.
"""
add_triangle!(s::HasDeltaSet2D, src2_first::Int, src2_last::Int, tgt2::Int; kw...) =
  add_part!(s, :Tri; ∂e0=src2_last, ∂e1=tgt2, ∂e2=src2_first, kw...)

""" Glue a triangle onto a simplicial set, given its boundary vertices.

If a needed edge between two vertices exists, it is reused (hence the "gluing");
otherwise, it is created.
"""
function glue_triangle!(s::HasDeltaSet2D, v₀::Int, v₁::Int, v₂::Int; kw...)
  add_triangle!(s, get_edge!(s, v₀, v₁), get_edge!(s, v₁, v₂),
                get_edge!(s, v₀, v₂); kw...)
end

function get_edge!(s::HasDeltaSet1D, src::Int, tgt::Int)
  es = edges(s, src, tgt)
  isempty(es) ? add_edge!(s, src, tgt) : first(es)
end

""" Glue a triangle onto a simplicial set, respecting the order of the vertices.
"""
function glue_sorted_triangle!(s::HasDeltaSet2D, v₀::Int, v₁::Int, v₂::Int; kw...)
  v₀, v₁, v₂ = sort(SVector(v₀, v₁, v₂))
  glue_triangle!(s, v₀, v₁, v₂; kw...)
end

# 2D oriented simplicial sets
#----------------------------

@present OrientedDeltaSchema2D <: DeltaCategory2D begin
  Orientation::AttrType
  edge_orientation::Attr(E,Orientation)
  tri_orientation::Attr(Tri,Orientation)
end

""" A two-dimensional oriented delta set.

Triangles are ordered in the cyclic order ``(0,1,2)`` when `tri_orientation` is
true/positive and in the reverse order when it is false/negative.
"""
@acset_type OrientedDeltaSet2D(OrientedDeltaSchema2D,
                               index=[:∂v0,:∂v1,:∂e0,:∂e1,:∂e2]) <: AbstractDeltaSet2D

orientation(::Type{Val{2}}, s::HasDeltaSet2D, args...) =
  s[args..., :tri_orientation]
set_orientation!(::Type{Val{2}}, s::HasDeltaSet2D, t, orientation) =
  (s[t, :tri_orientation] = orientation)

function ∂_nz(::Type{Val{2}}, s::HasDeltaSet2D, t::Int)
  edges = triangle_edges(s,t)
  (edges, sign(2,s,t) * sign(1,s,edges) .* @SVector([1,-1,1]))
end

function d_nz(::Type{Val{1}}, s::HasDeltaSet2D, e::Int)
  sgn = sign(1, s, e)
  t₀, t₁, t₂ = coface(2,0,s,e), coface(2,1,s,e), coface(2,2,s,e)
  (lazy(vcat, t₀, t₁, t₂),
   lazy(vcat, sgn*sign(2,s,t₀), -sgn*sign(2,s,t₁), sgn*sign(2,s,t₂)))
end

# 2D embedded simplicial sets
#----------------------------

@present EmbeddedDeltaSchema2D <: OrientedDeltaSchema2D begin
  Point::AttrType
  point::Attr(V, Point)
end

""" A two-dimensional, embedded, oriented delta set.
"""
@acset_type EmbeddedDeltaSet2D(EmbeddedDeltaSchema2D,
                               index=[:∂v0,:∂v1,:∂e0,:∂e1,:∂e2]) <: AbstractDeltaSet2D

volume(::Type{Val{n}}, s::EmbeddedDeltaSet2D, x) where n =
  volume(Val{n}, s, x, CayleyMengerDet())
volume(::Type{Val{2}}, s::HasDeltaSet2D, t::Int, ::CayleyMengerDet) =
  volume(point(s, triangle_vertices(s,t)))

# General operators
###################

""" Wrapper for simplex or simplices of dimension `n`.

See also: [`V`](@ref), [`E`](@ref), [`Tri`](@ref).
"""
@parts_array_struct Simplex{n}

""" Vertex in simplicial set: alias for `Simplex{0}`.
"""
const V = Simplex{0}

""" Edge in simplicial set: alias for `Simplex{1}`.
"""
const E = Simplex{1}

""" Triangle in simplicial set: alias for `Simplex{2}`.
"""
const Tri = Simplex{2}

""" Wrapper for chain of oriented simplices of dimension `n`.
"""
@vector_struct SimplexChain{n}

const VChain = SimplexChain{0}
const EChain = SimplexChain{1}
const TriChain = SimplexChain{2}

""" Wrapper for discrete form, aka cochain, in simplicial set.
"""
@vector_struct SimplexForm{n}

const VForm = SimplexForm{0}
const EForm = SimplexForm{1}
const TriForm = SimplexForm{2}

""" Simplices of given dimension in a simplicial set.
"""
@inline simplices(n::Int, s::HasDeltaSet) = 1:nsimplices(Val{n}, s)

""" Number of simplices of given dimension in a simplicial set.
"""
@inline nsimplices(n::Int, s::HasDeltaSet) = nsimplices(Val{n}, s)

""" Face map and boundary operator on simplicial sets.

Given numbers `n` and `0 <= i <= n` and a simplicial set of dimension at least
`n`, the `i`th face map is implemented by the call

```julia
∂(n, i, s, ...)
```

The boundary operator on `n`-faces and `n`-chains is implemented by the call

```julia
∂(n, s, ...)
```

Note that the face map returns *simplices*, while the boundary operator returns
*chains* (vectors in the free vector space spanned by oriented simplices).
"""
@inline ∂(i::Int, s::HasDeltaSet, x::Simplex{n}) where n =
  Simplex{n-1}(face(Val{(n,i)}, s, x.data))
@inline ∂(n::Int, i::Int, s::HasDeltaSet, args...) =
  face(Val{(n,i)}, s, args...)

@inline coface(i::Int, s::HasDeltaSet, x::Simplex{n}) where n =
  Simplex{n+1}(coface(Val{(n+1,i)}, s, x.data))
@inline coface(n::Int, i::Int, s::HasDeltaSet, args...) =
  coface(Val{(n,i)}, s, args...)

∂(s::HasDeltaSet, x::SimplexChain{n}) where n =
  SimplexChain{n-1}(∂(Val{n}, s, x.data))
@inline ∂(n::Int, s::HasDeltaSet, args...) = ∂(Val{n}, s, args...)

function ∂(::Type{Val{n}}, s::HasDeltaSet, args...) where n
  operator_nz(Int, nsimplices(n-1,s), nsimplices(n,s), args...) do x
    ∂_nz(Val{n}, s, x)
  end
end

""" Alias for the face map and boundary operator [`∂`](@ref).
"""
const boundary = ∂

""" The discrete exterior derivative, aka the coboundary operator.
"""
d(s::HasDeltaSet, x::SimplexForm{n}) where n =
  SimplexForm{n+1}(d(Val{n}, s, x.data))
@inline d(n::Int, s::HasDeltaSet, args...) = d(Val{n}, s, args...)

function d(::Type{Val{n}}, s::HasDeltaSet, args...) where n
  operator_nz(Int, nsimplices(n+1,s), nsimplices(n,s), args...) do x
    d_nz(Val{n}, s, x)
  end
end

""" Alias for the coboundary operator [`d`](@ref).
"""
const coboundary = d

""" Alias for the discrete exterior derivative [`d`](@ref).
"""
const exterior_derivative = d

""" Orientation of simplex.
"""
orientation(s::HasDeltaSet, x::Simplex{n}) where n =
  orientation(Val{n}, s, x.data)
@inline orientation(n::Int, s::HasDeltaSet, args...) =
  orientation(Val{n}, s, args...)

@inline Base.sign(n::Int, s::HasDeltaSet, args...) = sign(Val{n}, s, args...)
Base.sign(::Type{Val{n}}, s::HasDeltaSet, args...) where n =
  numeric_sign.(orientation(Val{n}, s, args...))

numeric_sign(x) = sign(x)
numeric_sign(x::Bool) = x ? +1 : -1

""" Set orientation of simplex.
"""
@inline set_orientation!(n::Int, s::HasDeltaSet, args...) =
  set_orientation!(Val{n}, s, args...)

""" ``n``-dimensional volume of ``n``-simplex in an embedded simplicial set.
"""
volume(s::HasDeltaSet, x::Simplex{n}, args...) where n =
  volume(Val{n}, s, x.data, args...)
@inline volume(n::Int, s::HasDeltaSet, args...) = volume(Val{n}, s, args...)

""" Convenience function for linear operator based on structural nonzero values.
"""
operator_nz(f, ::Type{T}, m::Int, n::Int,
            x::Int, Vec::Type=SparseVector{T}) where T = fromnz(Vec, f(x)..., m)
operator_nz(f, ::Type{T}, m::Int, n::Int,
            vec::AbstractVector) where T = applynz(f, vec, m, n)
operator_nz(f, ::Type{T}, m::Int, n::Int,
            Mat::Type=SparseMatrixCSC{T}) where T = fromnz(f, Mat, m, n)

# Consistent orientation
########################

""" Consistently orient simplices in a simplicial set, if possible.

Two simplices with a common face are *consistently oriented* if they induce
opposite orientations on the shared face. This function attempts to consistently
orient all simplices of a given dimension and returns whether this has been
achieved. Each connected component is oriently independently using the helper
function [`orient_component!`](@ref).
"""
orient!(s::AbstractDeltaSet1D) = orient!(s, E)
orient!(s::AbstractDeltaSet2D) = orient!(s, Tri)

function orient!(s::HasDeltaSet, ::Type{Simplex{n}}) where n
  # Compute connected components as coequalizer of face maps.
  ndom, ncodom = nsimplices(n, s), nsimplices(n-1, s)
  face_maps = SVector{n+1}([ FinFunction(x -> ∂(n,i,s,x), ndom, ncodom)
                             for i in 0:n ])
  π = only(coequalizer(face_maps))

  # Choose an arbitrary representative of each component.
  reps = zeros(Int, length(codom(π)))
  for x in reverse(simplices(n, s))
    reps[π(∂(n,0,s,x))] = x
  end

  # Orient each component, starting at the chosen representative.
  init_orientation = one(eltype(orientation(n, s)))
  for x in reps
    orient_component!(s, Simplex{n}(x), init_orientation) || return false
  end
  true
end

""" Consistently orient simplices in the same connected component, if possible.

Given an ``n``-simplex and a choice of orientation for it, this function
attempts to consistently orient all ``n``-simplices that may be reached from it
by traversing ``(n-1)``-faces. The traversal is depth-first. If a consistent
orientation is possible, the function returns `true` and the orientations are
assigned; otherwise, it returns `false` and no orientations are changed.

If the simplicial set is not connected, the function [`orient!`](@ref) may be
more convenient.
"""
orient_component!(s::AbstractDeltaSet1D, e::Int, args...) =
  orient_component!(s, E(e), args...)
orient_component!(s::AbstractDeltaSet2D, t::Int, args...) =
  orient_component!(s, Tri(t), args...)

function orient_component!(s::HasDeltaSet, x::Simplex{n},
                           x_orientation::Orientation) where {n, Orientation}
  orientations = repeat(Union{Orientation,Nothing}[nothing], nsimplices(n, s))

  orient_stack = Vector{Pair{Int64, Orientation}}()

  push!(orient_stack, x[] => x_orientation)
  is_orientable = true
  while !isempty(orient_stack)
    x, target = pop!(orient_stack)
    current = orientations[x]
    if isnothing(current)
      # If not visited, set the orientation and add neighbors to stack.
      orientations[x] = target
      for i in 0:n, j in 0:n
        next = iseven(i+j) ? negate(target) : target
        for y in coface(n, j, s, ∂(n, i, s, x))
          y == x || push!(orient_stack, y=>next)
        end
      end
    elseif current != target
      is_orientable = false
      break
    end
  end

  if is_orientable
    component = findall(!isnothing, orientations)
    set_orientation!(n, s, component, orientations[component])
  end
  is_orientable
end

negate(x) = -x
negate(x::Bool) = !x

# Euclidean geometry
####################

""" ``n``-dimensional volume of ``n``-simplex spanned by given ``n+1`` points.
"""
function volume(points)
  CM = cayley_menger(points...)
  n = length(points) - 1
  sqrt(abs(det(CM)) / 2^n) / factorial(n)
end

""" Construct Cayley-Menger matrix for simplex spanned by given points.

For an ``n`-simplex, this is the ``(n+2)×(n+2)`` matrix that appears in the
[Cayley-Menger
determinant](https://en.wikipedia.org/wiki/Cayley-Menger_determinant).
"""
function cayley_menger(p0::V, p1::V) where V <: AbstractVector
  d01 = sqdistance(p0, p1)
  SMatrix{3,3}(0,  1,   1,
               1,  0,   d01,
               1,  d01, 0)
end
function cayley_menger(p0::V, p1::V, p2::V) where V <: AbstractVector
  d01, d12, d02 = sqdistance(p0, p1), sqdistance(p1, p2), sqdistance(p0, p2)
  SMatrix{4,4}(0,  1,   1,   1,
               1,  0,   d01, d02,
               1,  d01, 0,   d12,
               1,  d02, d12, 0)
end
function cayley_menger(p0::V, p1::V, p2::V, p3::V) where V <: AbstractVector
  d01, d12, d02 = sqdistance(p0, p1), sqdistance(p1, p2), sqdistance(p0, p2)
  d03, d13, d23 = sqdistance(p0, p3), sqdistance(p1, p3), sqdistance(p2, p3)
  SMatrix{5,5}(0,  1,   1,   1,   1,
               1,  0,   d01, d02, d03,
               1,  d01, 0,   d12, d13,
               1,  d02, d12, 0,   d23,
               1,  d03, d13, d23, 0)
end

""" Squared Euclidean distance between two points.
"""
sqdistance(x, y) = sum((x-y).^2)

end


#=
PSD cone
cuts are V*V^T for eigenvectors V corresponding to negative eigenvalues of matrix point X
=#
function get_init_cuts(cone::MOI.PositiveSemidefiniteConeTriangle)
    cuts = Vector{Float64}[]
    L = cone.side_dimension
    dim = MOI.dimension(cone)

    # diagonal variables are nonnegative
    for i in 1:L
        cut = zeros(dim)
        cut[mat_to_vec_idx(i, i)] = 1.0
        push!(cuts, cut)
    end

    # 2x2 principal minors satisfy linearizations of PSD/RSOC condition
    for i in 1:L, j in 1:(i - 1)
        cut = zeros(dim)
        cut[mat_to_vec_idx(i, i)] = 1.0
        cut[mat_to_vec_idx(j, j)] = 1.0
        ij_idx = mat_to_vec_idx(i, j)
        cut[ij_idx] = 2.0
        push!(cuts, cut)
        cut = copy(cut)
        cut[ij_idx] = -2.0
        push!(cuts, cut)
    end

    return cuts
end

function get_sep_cuts(
    x::Vector{Float64},
    cone::MOI.PositiveSemidefiniteConeTriangle,
    tol::Float64,
)
    cuts = Vector{Float64}[]
    L = cone.side_dimension
    X = Matrix{Float64}(undef, L, L)
    vec_to_mat_noscale!(X, x)

    # TODO try only getting negative eigenvalues (what lower bound?) vs smallest k eigenvalues
    F = eigen!(Symmetric(X, :L), 1:min(L, 5))

    for k in eachindex(F.values)
        if F.values[k] >= -0.1 * tol
            continue
        end

        V = F.vectors[:, k]
        W = V * V'
        cut = similar(x)
        mat_to_vec_scale!(cut, W)

        if dot(x, cut) < -tol
            push!(cuts, cut)
        end
    end

    return cuts
end


using RandomMatrices
using LinearAlgebra: I, tr
using Test

@testset "Haar" begin

N=5
A=randn(N,N)
B=randn(N,N)
Q=rand(Haar(1), N)

#Test case where there are no symbolic Haar matrices
@test_broken eval(expectation(:(A*B))) ≈ A*B
#Test case where there is one pair of symbolic Haar matrices
@test_broken tr(eval(expectation(:(A*Q*B*Q')))) ≈ tr(A)*tr(B)/N

println("Case 3")
@test_broken println("E(A*Q*B*Q'*A*Q*B*Q') = ", eval(expectation(N, :Q, :(A*Q*B*Q'*A*Q*B*Q'))))

for elty in (Float64, ComplexF64)
	A = Stewart(elty, N)
        @test A'A ≈ Matrix{elty}(I, N, N)
end

end # testset


"""
    Implements Thompson Sampling for Bandits with r ∈ [ 0,1]
"""

type TS <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    α0::Vector{Int64}
    β0::Vector{Int64}
    cummSuccess::Vector{Int64}
    cummFailure::Vector{Int64}

    samplingDist::Vector{Distributions.Beta}

    function TS( noOfArms::Integer )
        new( noOfArms,
             0,
             0,
             ones(Int64,noOfArms),
             ones(Int64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms),
             fill(Distributions.Beta(1,1),noOfArms)
        )
    end

    function TS( armParams::Array{Tuple{Int64,Int64},1} )
        _noOfArms   = length( armParams )
        _priorDist  = [ Distributions.Beta(armParams[idx][1],armParams[idx][2]) for idx=1:_noOfArms ]
        new( _noOfArms,
             0,
             0,
             [ armParams[idx][1] for idx=1:_noOfArms ],
             [ armParams[idx][2] for idx=1:_noOfArms ],
             zeros(Float64,_noOfArms),
             zeros(Float64,_noOfArms),
             _priorDist
        )
    end
end

function get_arm_index( agent::TS )
    agent.lastPlayedArm = findmax( [rand(armSample) for armSample in agent.samplingDist] )[2]
    return agent.lastPlayedArm
end

function update_reward!( agent::TS, r::Int64 )
    # Update S and F
    agent.cummSuccess[agent.lastPlayedArm] += (r==0?0:1)
    agent.cummFailure[agent.lastPlayedArm] += (r==0?1:0)

    # Update Distributions
    agent.samplingDist[agent.lastPlayedArm] = Distributions.Beta(
                                                agent.cummSuccess[agent.lastPlayedArm]+agent.α0[agent.lastPlayedArm],
                                                agent.cummFailure[agent.lastPlayedArm]+agent.β0[agent.lastPlayedArm]
                                            )

    # Update time steps
    agent.noOfSteps += 1

    nothing
end

function update_reward!( agent::TS, r::Float64 )
    rTilde = rand( Distributions.Bernoulli(r) )
    update_reward!( agent, rTilde )

    nothing
end

function reset!( agent::TS )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.cummSuccess   = zeros( Float64, agent.noOfArms )
    agent.cummFailure   = zeros( Float64, agent.noOfArms )
    agent.samplingDist  = [ Distributions.Beta(agent.α0[idx],agent.β0[idx]) for idx=1:agent.noOfArms ]

    nothing
end

function info_str( agent::TS, latex::Bool )
    return @sprintf( "Thompson Sampling" )
end

"""-----------------------------------------------------------------------------
    Dynamic Thompson Sampling
    Based on: Gupta, N., Granmo, O. C., & Agrawala, A. (2011). Thompson sampling for dynamic multi-armed bandits. Proceedings - 10th International Conference on Machine Learning and Applications, ICMLA 2011, 1, 484–489. http://doi.org/10.1109/ICMLA.2011.144
"""

type DynamicTS <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    α::Vector{Float64}
    β::Vector{Float64}
    C::Float64

    samplingDist::Vector{Distributions.Beta}

    function DynamicTS( noOfArms, _C )
        new( noOfArms,
             0,
             0,
             2*ones(Float64,noOfArms),
             2*ones(Float64,noOfArms),
             _C,
             fill(Distributions.Beta(2,2),noOfArms)
        )
    end
end

function get_arm_index( agent::DynamicTS )
    agent.lastPlayedArm = findmax( [rand(armSample) for armSample in agent.samplingDist] )[2]
    return agent.lastPlayedArm
end

function update_reward!( agent::DynamicTS, r::Integer )
    # Update reward to arm played
    if agent.α[agent.lastPlayedArm]+agent.β[agent.lastPlayedArm] < agent.C
        agent.α[agent.lastPlayedArm]    += (r==0?0:1)
        agent.β[agent.lastPlayedArm]    += (r==0?1:0)
    else
        agent.α[agent.lastPlayedArm]    = (agent.α[agent.lastPlayedArm]+(r==0?0:1)) * agent.C/(agent.C+1)
        agent.β[agent.lastPlayedArm]    = (agent.β[agent.lastPlayedArm]+(r==0?1:0)) * agent.C/(agent.C+1)
    end

    # Update arm's distribution
    agent.samplingDist[agent.lastPlayedArm]  = Distributions.Beta(
                                                agent.α[agent.lastPlayedArm],
                                                agent.β[agent.lastPlayedArm]
                                            )

    # Update time steps
    agent.noOfSteps += 1

    nothing
end

function update_reward!( agent::DynamicTS, r::AbstractFloat )
    update_reward!( agent, rand(Distributions.Bernoulli(r)) )   # Do a Bernoulli Trial to update the posterior

    nothing
end

function reset!( agent::DynamicTS )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.α             = 2 * ones( Float64, agent.noOfArms )
    agent.β             = 2 * ones( Float64, agent.noOfArms )
    agent.samplingDist  = fill( Distributions.Beta(2,2), agent.noOfArms )

    nothing
end

function info_str( agent::DynamicTS, latex::Bool )
    if latex
        return @sprintf( "Dynamic Thompson Sampling (\$C = %4.3f\$)", agent.C )
    else
        return @sprintf( "Dynamic Thompson Sampling (C = %4.3f)", agent.C )
    end
end

"""-----------------------------------------------------------------------------
    Optimistic Thompson Sampling
    Based on: May, B. C., Korda, N., Lee, A., & Leslie, D. S. (2012). Optimistic bayesian sampling in contextual-bandit problems. Journal of Machine Learning Research, 13, 2069–2106.
"""

type OTS <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    α0::Vector{Int64}
    β0::Vector{Int64}
    cummSuccess::Vector{Int64}
    cummFailure::Vector{Int64}

    samplingDist::Vector{Distributions.Beta}

    function OTS( noOfArms::Integer )
        new( noOfArms,
             0,
             0,
             ones(Int64,noOfArms),
             ones(Int64,noOfArms),
             zeros(Float64,noOfArms),
             zeros(Float64,noOfArms),
             fill(Distributions.Beta(1,1),noOfArms)
        )
    end

    function OTS( armParams::Array{Tuple{Int64,Int64},1} )
        _noOfArms   = length( armParams )
        _priorDist  = [ Distributions.Beta(armParams[idx][1],armParams[idx][2]) for idx=1:_noOfArms ]
        new( _noOfArms,
             0,
             0,
             [ armParams[idx][1] for idx=1:_noOfArms ],
             [ armParams[idx][2] for idx=1:_noOfArms ],
             zeros(Float64,_noOfArms),
             zeros(Float64,_noOfArms),
             _priorDist
        )
    end
end

function get_arm_index( agent::OTS )
    agent.lastPlayedArm = findmax( [ max(mean(armSample),rand(armSample)) for armSample in agent.samplingDist] )[2]
    return agent.lastPlayedArm
end

function update_reward!( agent::OTS, r::Int64 )
    # Update S and F
    agent.cummSuccess[agent.lastPlayedArm] += (r==1?1:0)
    agent.cummFailure[agent.lastPlayedArm] += (r==0?1:0)

    # Update Distributions
    agent.samplingDist[agent.lastPlayedArm] = Distributions.Beta(
                                                agent.cummSuccess[agent.lastPlayedArm]+agent.α0[agent.lastPlayedArm],
                                                agent.cummFailure[agent.lastPlayedArm]+agent.β0[agent.lastPlayedArm]
                                            )
    # Update time steps
    agent.noOfSteps += 1

    nothing
end

function update_reward!( agent::OTS, r::Float64 )
    rTilde = rand( Distributions.Bernoulli(r) )
    update_reward!( agent, rTilde )

    nothing
end

function reset!( agent::OTS )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.cummSuccess   = zeros( Float64, agent.noOfArms )
    agent.cummFailure   = zeros( Float64, agent.noOfArms )
    agent.samplingDist  = fill( Distributions.Beta(1,1), agent.noOfArms )

    nothing
end

function info_str( agent::OTS, latex::Bool )
    return @sprintf( "Optimistic Thompson Sampling" )
end

"""-----------------------------------------------------------------------------
    Thompson Sampling for Normal Rewards
    Based on: Sec 6.3, Russo, D., & Van Roy, B. (2014). Learning to Optimize Via Posterior Sampling. Mathematics of Operations Research. http://doi.org/10.1287/xxxx.0000.0000
"""

type TSNormal <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    η::Float64
    μ::Vector{Float64}
    σ::Vector{Float64}

    samplingDist::Vector{Distributions.Normal}

    function TSNormal( noOfArms::Int64, η::Real = 1 )
        new(
            noOfArms,
            0,
            0,
            η,
            zeros(Float64,noOfArms),
            100*ones(Float64,noOfArms),    # σ - Some high value
            fill(Distributions.Normal(0,100),noOfArms)
        )
    end

    function TSNormal( armParams::Array{Tuple{Real,Real},1}, η::Real = 1 )
        _K = length(armParams)
        new(
            _K,
            0,
            0,
            η,
            [ armParams[idx][1] for i = 1:_K ],
            [ armParams[idx][2] for i = 1:_K ],
            [ Distributions.Normal(armParams[idx][1],armParams[idx][2]) for i = 1:_K ]
        )
    end
end

function get_arm_index( agent::TSNormal )
    agent.lastPlayedArm = findmax( [ max(mean(armSample),rand(armSample)) for armSample in agent.samplingDist] )[2]
    return agent.lastPlayedArm
end

function update_reward!( agent::TSNormal, r::Float64 )

    agent.μ[agent.lastPlayedArm] = ( agent.μ[agent.lastPlayedArm]/agent.σ[agent.lastPlayedArm]^2 + r/agent.η^2 ) /
                                        ( 1/agent.σ[agent.lastPlayedArm]^2 + 1/agent.η^2 )
    agent.σ[agent.lastPlayedArm] = ( 1/agent.σ[agent.lastPlayedArm]^2 + 1/agent.η^2 )^(-1/2)

    # Update Distributions
    agent.samplingDist[agent.lastPlayedArm] = Distributions.Normal(
                                                agent.μ[agent.lastPlayedArm],
                                                agent.σ[agent.lastPlayedArm]
                                            )

    agent.noOfSteps += 1

    nothing
end

function reset!( agent::TSNormal )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.μ   = zeros( Float64, agent.noOfArms )
    agent.σ   = 100*ones(Float64,agent.noOfArms)
    agent.samplingDist  = [ Distributions.Normal(agent.μ[i],agent.σ[i]) for i = 1:agent.noOfArms ]

    nothing
end

function info_str( agent::TSNormal, latex::Bool )
    if latex
        return @sprintf( "TS-Normal(\$\\eta=%3.2f\$)", agent.η )
    else
        return @sprintf( "TS-Normal(η=%3.2f)", agent.η )
    end
end

"""-----------------------------------------------------------------------------
    Discounted Thompson Sampling
    Based on: Vishnu Raj, Sheetal Kalyani (2017). Taming Non-stationary Bandits: A Bayesian Approach. https://arxiv.org/abs/1707.09727
"""

type dTS <: BanditAlgorithmBase
    noOfArms::Int64
    noOfSteps::Int64
    lastPlayedArm::Int64

    γ::Float64                      # Discounting factor
    αₒ::Vector{Float64}             # Initial α value
    βₒ::Vector{Float64}             # Initial β value
    cummSuccess::Vector{Float64}    # Cummulative success
    cummFailure::Vector{Float64}    # Cummulative failure

    samplingDist::Vector{Distributions.Beta}    # Internal Distributions

    function dTS( noOfArms, γ::Float64 = 1.00,
                    αₒ::Float64 = 1.00, βₒ::Float64 = 1.00 )
        new( noOfArms,
             0,
             0,
             γ,
             αₒ * ones(Float64,noOfArms),
             βₒ * ones(Float64,noOfArms),
             αₒ * ones(Float64,noOfArms),
             βₒ * ones(Float64,noOfArms),
             fill(Distributions.Beta(αₒ,βₒ),noOfArms)
        )
    end

    function dTS( noOfArms, γ::Float64 ,
                    α₀::Vector{Float64}, β₀::Vector{Float64} )
        new( noOfArms,
             0,
             0,
             γ,
             α₀,
             β₀,
             α₀,
             β₀,
             Distributions.Beta.(α₀,β₀)
        )
    end
end

function get_arm_index( agent::dTS )
    agent.lastPlayedArm = findmax( [rand(armSample) for armSample in agent.samplingDist] )[2]
    return agent.lastPlayedArm
end

function update_reward!( agent::dTS, r::Int64 )
    # Update S and F
    agent.cummSuccess *= agent.γ
    agent.cummFailure *= agent.γ
    agent.cummSuccess[agent.lastPlayedArm] += (r==0?0:1)
    agent.cummFailure[agent.lastPlayedArm] += (r==0?1:0)

    # Update Distributions
    agent.samplingDist  = Distributions.Beta.(
                                                agent.cummSuccess + agent.αₒ,
                                                agent.cummFailure + agent.βₒ
                                            )

    # Update time steps
    agent.noOfSteps += 1
end

function update_reward!( agent::dTS, r::Float64 )
    rTilde = rand( Distributions.Bernoulli(r) )
    update_reward!( agent, rTilde )
end

function reset!( agent::dTS )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.cummSuccess   = copy( agent.αₒ )
    agent.cummFailure   = copy( agent.βₒ )
    agent.samplingDist  = Distributions.Beta.(agent.αₒ,agent.βₒ)
end

function info_str( agent::dTS, latex::Bool )
    if latex
        return @sprintf( "dTS(\$\\gamma = %4.3f\$)", agent.γ )
    else
        return @sprintf( "dTS(γ = %4.3f)", agent.γ )
    end
end

""" ----------------------------------------------------------------------------
    Discounted Optimistic Thompson Sampling
    Based on: Vishnu Raj, Sheetal Kalyani (2017). Taming Non-stationary Bandits: A Bayesian Approach. https://arxiv.org/abs/1707.09727
"""
type dOTS <: BanditAlgorithmBase
    _dTS::dTS

    function dOTS( noOfArms, γ::Float64 = 1.00,
                    αₒ::Float64 = 1.00, βₒ::Float64 = 1.00 )
        new(
             dTS( noOfArms, γ, αₒ, βₒ )
        )
    end
end

function get_arm_index( agent::dOTS )
    agent._dTS.lastPlayedArm = findmax( [max(rand(armSample),mean(armSample)) for armSample in agent._dTS.samplingDist] )[2]
    return agent._dTS.lastPlayedArm
end

function update_reward!( agent::dOTS, r::Int64 )
    update_reward!( agent._dTS, r )
    nothing
end

function update_reward!( agent::dOTS, r::Float64 )
    update_reward!( agent._dTS, r )
    nothing
end

function reset!( agent::dOTS )
    reset!( agent._dTS )
    nothing
end

function info_str( agent::dOTS, latex::Bool )
    if latex
        return @sprintf( "dOTS(\$\\gamma = %4.3f\$)", agent._dTS.γ )
    else
        return @sprintf( "dOTS(γ = %4.3f)", agent._dTS.γ )
    end
end

""" ----------------------------------------------------------------------------
    Restarting Thompson Sampling
    Based on:
"""

type RestartTS <: BanditAlgorithmBase
    _TS::TS
    Δ::Int64

    function RestartTS( noOfArms::Int, Δ )
        new( TS(noOfArms), Δ )
    end
end

function get_arm_index( agent::RestartTS )
    return get_arm_index( agent._TS )
end

function update_reward!( agent::RestartTS, r::Int64 )
    update_reward!( agent._TS, r )
    if agent._TS.noOfSteps >= agent.Δ
        reset!( agent._TS )
    end
    nothing
end

function update_reward!( agent::RestartTS, r::Float64 )
    update_reward!( agent._TS, r )
    if agent._TS.noOfSteps >= agent.Δ
        reset!( agent._TS )
    end
    nothing
end

function reset!( agent::RestartTS )
    reset!( agent._TS )
    nothing
end

function info_str( agent::RestartTS, latex::Bool )
    if latex
        return @sprintf( "Restarting TS(\$\\Delta = %d\$)", agent.Δ )
    else
        return @sprintf( "Restarting TS(Δ = %d)", agent.Δ )
    end
end

""" ----------------------------------------------------------------------------------------
    Thompson Sampling with Gaussian Priors
    Based on: Algorithm 2, S. Agrawal and N. Goyal, “Near-Optimal Regret Bounds for Thompson Sampling,” J. ACM, vol. 64, no. 5, 2017.
"""
type TSGaussPrior <: BanditAlgorithmBase
    noOfArms::Integer
    noOfSteps::Integer
    lastPlayedArm::Integer

    # Algorithm specific parameters
    μ::Vector{Float64}
    noOfPulls::Vector{Float64} # Even though an integer, keeping this as Float for performance

    samplingDist::Vector{Distributions.Normal}

    function TSGaussPrior( noOfArms::Integer )
        new(
            noOfArms,
            0,
            0,
            zeros( Float64, noOfArms),
            zeros( Float64, noOfArms ),
            fill( Distributions.Normal(0,1), noOfArms )
        )
    end
end

function get_arm_index( agent::TSGaussPrior )
    agent.lastPlayedArm = findmax( [rand(armSample) for armSample in agent.samplingDist] )[2]
    return agent.lastPlayedArm
end

function  update_reward!( agent::TSGaussPrior, r::Real )
    # Update observation
    agent.μ[agent.lastPlayedArm] = (agent.μ[agent.lastPlayedArm]*agent.noOfPulls[agent.lastPlayedArm]+r) /
                                        (agent.noOfPulls[agent.lastPlayedArm]+1)
    agent.noOfPulls[agent.lastPlayedArm] += 1

    # Update distribution
    agent.samplingDist[agent.lastPlayedArm] = Distributions.Normal(
                                                        agent.μ[agent.lastPlayedArm],
                                                        1.0/(agent.noOfPulls[agent.lastPlayedArm]+1)
                                                )

    # Update time steps
    agent.noOfSteps += 1

    nothing
end

function  reset!( agent::TSGaussPrior )
    agent.noOfSteps     = 0
    agent.lastPlayedArm = 0

    agent.μ             = zeros( Float64, agent.noOfArms )
    agent.noOfPulls     = zeros( Float64, agent.noOfArms )

    agent.samplingDist  = fill( Distributions.Normal(0,1), agent.noOfArms )
end

function  info_str( agent::TSGaussPrior, latex::Bool )
    return @sprintf( "TS Gauss Prior" )
end


#=

      Linear Layer

=#

export LinearLayer

type LinearLayer{D<:Device, F<:Float} <: Layer
  # Device
  device::Device

  # The parameters
  weight::matrixParameters(D,F)
  bias::matrixParameters(D,F)

  # State
  output::denseRealMatrix(D,F)
  input_gradient::denseRealMatrix(D,F)

  function LinearLayer(device, weight, bias)
    new(device, weight, bias, array(device, F, 0, 0), array(device, F, 0, 0))
  end
end

"Initialize a linear module with the input and output size"
function LinearLayer{D<:Device, F<:Float}(d::D, ::Type{F}, inputSize::Int64, outputSize::Int64)
  LinearLayer{D,F}(d, matrixParameters(d, F, inputSize, outputSize), matrixParameters(d, F, 1, outputSize))
end

function forward!{D<:Device, F<:Float}(linear::LinearLayer{D,F}, input::DenseMatrix)
  @ensuresize linear.output, size(input,1), size(linear.weight.values, 2)

  gemm!('N', 'N', 1., input, linear.weight.values, 0., linear.output)
  broadcast!(+, linear.output, linear.output, linear.bias.values)
end

function compute_inputgradient!{D<:Device, F<:Float}(linear::LinearLayer{D,F}, input::DenseMatrix, gradOutput::DenseMatrix)
  @ensuresize linear.input_gradient, size(gradOutput,1), size(linear.weight.values, 1)

  gemm!('N', 'T', 1., gradOutput, linear.weight.values, 0., linear.input_gradient)
end

function update_gradient!{D<:Device, F<:Float}(linear::LinearLayer{D,F}, input::DenseMatrix, gradOutput::DenseMatrix, scale::F=1.)
  if !isnull(linear.weight.gradient)
    gemm!('T', 'N', scale, input, gradOutput, 0., get(linear.weight.gradient))
  end

  if !isnull(linear.bias.gradient)
    axpy!(scale, sum(gradOutput, 1), get(linear.bias.gradient))
  end
end


using OrdinaryDiffEq, Zygote, DiffEqSensitivity, Test

p_model = [1f0]

u0 = Float32.([0.0])

function dudt(du, u, p, t)
    du[1] = p[1]
end

prob = ODEProblem(dudt,u0,(0f0,99.9f0))

function predict_neuralode(p)
    _prob = remake(prob,p=p)
    Array(solve(_prob,Tsit5(), saveat=0.1))
end

loss(p) = sum(abs2,predict_neuralode(p))/length(p)

p_model_ini = copy(p_model)

@test !iszero(Zygote.gradient(loss,p_model_ini)[1])


@doc raw"""
    βdifferential_geodesic_startpoint(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_x g(t;p,q)[X]$.
They are

```math
\beta(\kappa) = \begin{cases}
\frac{\sinh(d(1-t)\sqrt{-\kappa})}{\sinh(d\sqrt{-\kappa})}
&\text{ if }\kappa < 0,\\
1-t & \text{ if } \kappa = 0,\\
\frac{\sin((1-t)d\sqrt{\kappa})}{\sinh(d\sqrt{\kappa})}
&\text{ if }\kappa > 0.
\end{cases}
```

Due to a symmetry agrument, these are also used to compute $D_q g(t; p,q)[\eta]$

# See also

[`differential_geodesic_endpoint`](@ref), [`differential_geodesic_startpoint`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_geodesic_startpoint(κ,t,d)
    if (d==0) || (κ==0)
        return (1-t)
    else
        if κ < 0
            return sinh(sqrt(-κ)*(1-t)*d)/sinh(sqrt(-κ)*d)
        elseif κ > 0
            return sin(sqrt(κ)*(1-t)*d)/sin(sqrt(κ)*d)
        end
    end
end
@doc raw"""
    βdifferential_exp_basepoint(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_p \exp_p X [Y]$. They are

```math
\beta(\kappa) = \begin{cases}
\cosh(\sqrt{-\kappa})&\text{ if }\kappa < 0,\\
1 & \text{ if } \kappa = 0,\\
\cos(\sqrt{\kappa}) &\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_exp_basepoint`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_exp_basepoint(κ,t,d)
    if κ < 0
        return cosh(sqrt(-κ)*d)
    elseif κ > 0
        return cos(sqrt(κ)*d)
    else
        return 1.0;
    end
end
@doc raw"""
    βdifferential_exp_argument(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_X \exp_p X[Y]$. They are

```math
$\beta(\kappa) = \begin{cases}
\frac{\sinh(d\sqrt{-\kappa})}{d\sqrt{-\kappa}}&\text{ if }\kappa < 0,\\
1 & \text{ if } \kappa = 0,\\
\frac{\sin(d\sqrt{\kappa})}{\sqrt{d\kappa}}&\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_exp_argument`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_exp_argument(κ,t,d)
    if (d==0) || (κ==0)
        return 1.0
    else
        if κ < 0
            return sinh(sqrt(-κ)*d)/( d*sqrt((-κ)) )
        elseif κ > 0
            return sin( sqrt(κ)*d )/( d*sqrt(κ) )
        end
    end
end
@doc raw"""
    βdifferential_log_basepoint(κ,t,d)

weights for the [`jacobi_field`](@ref) corresponding to the differential of the geodesic
with respect to its start point $D_p \log_p q[X]$. They are

```math
\beta(\kappa) = \begin{cases}
-\sqrt{-\kappa}d\frac{\cosh(d\sqrt{-\kappa})}{\sinh(d\sqrt{-\kappa})}&\text{ if }\kappa < 0,\\
-1 & \text{ if } \kappa = 0,\\
-\sqrt{\kappa}d\frac{\cos(d\sqrt{\kappa})}{\sin(d\sqrt{\kappa})}&\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_log_argument`](@ref), [`differential_log_argument`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_log_basepoint(κ::Number,t::Number,d::Number)
    if (d==0) || (κ==0)
        return -1.0
    else
        if κ < 0
            return - sqrt(-κ)*d*cosh(sqrt(-κ)*d)/sinh(sqrt(-κ)*d)
        else #if κ > 0
            return - sqrt(κ)*d*cos(sqrt(κ)*d)/sin(sqrt(κ)*d)
        end
    end
end
@doc raw"""
    βdifferential_log_argument(κ,t,d)

weights for the JacobiField corresponding to the differential of the logarithmic
map with respect to its argument $D_q \log_p q[X]$. They are

```math
\beta(\kappa) = \begin{cases}
\frac{ d\sqrt{-\kappa} }{\sinh(d\sqrt{-\kappa})}&\text{ if }\kappa < 0,\\
1 & \text{ if } \kappa = 0,\\
\frac{ d\sqrt{\kappa} }{\sin(d\sqrt{\kappa})}&\text{ if }\kappa > 0.
\end{cases}
```

# See also

[`differential_log_basepoint`](@ref), [`jacobi_field`](@ref)
"""
function βdifferential_log_argument(κ::Number,t::Number,d::Number)
    if (d==0) || (κ==0)
        return 1.0
    else
        if κ < 0
            return sqrt(-κ)*d/sinh(sqrt(-κ)*d)
        else #if κ > 0
            return sqrt(κ)*d/sin(sqrt(κ)*d)
        end
    end
end

@doc raw"""
    Y = adjoint_Jacobi_field(M, p, q, t, X, β)

Compute the AdjointJacobiField $J$ along the geodesic $γ_{p,q}$ on the manifold
$\mathcal M$ with initial conditions (depending on the application)
$X ∈ T_{γ_{p,q}(t)}\mathcal M$ and weights $β$. The result is a vector
$Y ∈ T_p\mathcal M$. The main difference to [`jacobi_field`](@ref) is the,
that the input `X` and the output `Y` switched tangent spaces.
For detais see [`jacobi_field`](@ref)
"""
function adjoint_Jacobi_field(M::Manifold, p, q, t, X, β::Function=βdifferential_geodesic_startpoint)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M,p,q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M),Ref(p), V, Ref(x), Ref(ParallelTransport()))
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (
            inner.(Ref(M),Ref(x),Ref(X),Θ)
        ) .* (
            β.(B.data.eigenvalues,Ref(t),Ref(distance(M,p,q)))
        ) .* V
    )
    return Y
end
function adjoint_Jacobi_field(M::Circle{ℝ}, p::Real, q::Real, t::Real, X::Real, β::Function=βdifferential_geodesic_startpoint)
    x = shortest_geodesic(M, p, q, t)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M,p,q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.(Ref(M),Ref(p), V, Ref(x), Ref(ParallelTransport()))
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (
            inner.(Ref(M),Ref(x),Ref(X),Θ)
        ) .* (
            β.(B.data.eigenvalues,Ref(t),Ref(distance(M,p,q)))
        ) .* V
    )[1]
    return Y
end
function adjoint_Jacobi_field(M::AbstractPowerManifold, p, q, t, X, β::Function=βdifferential_geodesic_startpoint)
    rep_size = representation_size(M.manifold)
    Y = allocate_result(M, adjoint_Jacobi_field, p, X)
    for i in get_iterator(M)
        #lY = adjoint_Jacobi_field(M.manifold, p[i], q[i], t, X[i], β)
        lY = adjoint_Jacobi_field(
            M.manifold,
            get_component(M,p,i),
            get_component(M,q,i),
            t,
            get_component(M,X,i),
            β,
        )
        set_component!(M,Y,lY,i)
    end
    return Y
end

@doc raw"""
    Y = jacobi_field(M, p, q, t, X, β)
compute the Jacobi jield $J$ along the geodesic $γ_{p,q}$ on the manifold $\mathcal M$ with
initial conditions (depending on the application) $X ∈ T_p\mathcal M$ and weights $β$. The
result is a tangent vector `Y` from $T_{γ_{p,q}(t)}\mathcal M$.

# See also

[`adjoint_Jacobi_field`](@ref)
"""
function jacobi_field(M::Manifold, p, q, t, X, β::Function=βdifferential_geodesic_startpoint)
    x = shortest_geodesic(M, p, q, t);
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M,p,q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.( Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    Y = zero_tangent_vector(M,p)
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y .= sum(
        (
            inner.(Ref(M), Ref(p), Ref(X), V)
        ) .* (
            β.(B.data.eigenvalues,Ref(t),Ref(distance(M,p,q)))
        ) .* Θ
    )
    return Y
end
function jacobi_field(M::Circle{ℝ}, p::Real, q::Real, t::Real, X::Real, β::Function=βdifferential_geodesic_startpoint)
    x = shortest_geodesic(M, p, q, t);
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(log(M,p,q)))
    V = get_vectors(M, p, B)
    Θ = vector_transport_to.( Ref(M), Ref(p), V, Ref(x), Ref(ParallelTransport()))
    Y = zero_tangent_vector(M,p)
    # Decompose wrt. frame, multiply with the weights from w and recompose with Θ.
    Y = sum(
        (
            inner.(Ref(M), Ref(p), Ref(X), V)
        ) .* (
            β.(B.data.eigenvalues,Ref(t),Ref(distance(M,p,q)))
        ) .* Θ
    )[1]
    return Y
end
function jacobi_field(M::AbstractPowerManifold, p, q, t, X, β::Function=βdifferential_geodesic_startpoint)
    rep_size = representation_size(M.manifold)
    Y = allocate_result(M, adjoint_Jacobi_field, p, X)
    for i in get_iterator(M)
        #lY = jacobi_field(M.manifold, p[i], q[i], t, X[i], β)
        set_component!(
            M,
            Y,
            jacobi_field(
                M.manifold,
                get_component(M,p,i),
                get_component(M,q,i),
                t,
                get_component(M,X,i),
                β,
            ),
            i,
        )
    end
    return Y
end


function wsplintp!(n::Int64, x, w)
    ## arguments
    #integer, intent(in) :: n
    #real(8), intent(in) :: x(n)
    #real(8), intent(out) :: w(4,n)

    @assert n >= 4

    w[:,1] .= 0.0
    f = zeros(4)

    f[1] = 1.0
    w[1,2] = polynm(-1, 4, x, f, x[2])

    f[1] = 0.0
    f[2] = 1.0
    w[2,2] = polynm(-1, 4, x, f, x[2])
    f[2] = 0.0

    f[3] = 1.0
    w[3,2] = polynm(-1, 4, x, f, x[2])
    f[3] = 0.0

    f[4] = 1.0
    w[4,2] = polynm(-1, 4, x, f, x[2])

    for i in 3:n-1
        f[:] .= 0.0
        f[1] = 1.0
        @views t1 = polynm(-1, 4, x[i-2:end], f, x[i-1])
        @views t2 = polynm(-1, 4, x[i-2:end], f, x[i])
        w[1,i] = t2 - t1

        f[1] = 0.0
        f[2] = 1.0
        @views t1 = polynm(-1, 4, x[i-2:end], f, x[i-1])
        @views t2 = polynm(-1, 4, x[i-2:end], f, x[i])
        w[2,i] = t2 - t1

        f[2] = 0.0
        f[3] = 1.0
        @views t1 = polynm(-1, 4, x[i-2:end], f, x[i-1])
        @views t2 = polynm(-1, 4, x[i-2:end], f, x[i])
        w[3,i] = t2 - t1

        f[3] = 0.0
        f[4] = 1.0
        @views t1 = polynm(-1, 4, x[i-2:end], f, x[i-1])
        @views t2 = polynm(-1, 4, x[i-2:end], f, x[i])
        w[4,i] = t2 - t1
    end

    f[:] .= 0.0
    f[1] = 1.0
    @views t1 = polynm(-1, 4, x[n-3:end], f, x[n-1])
    @views t2 = polynm(-1, 4, x[n-3:end], f, x[n])
    w[1,n] = t2 - t1

    f[1] = 0.0
    f[2] = 1.0
    @views t1 = polynm(-1, 4, x[n-3:end], f, x[n-1])
    @views t2 = polynm(-1, 4, x[n-3:end], f, x[n])
    w[2,n] = t2 - t1

    f[2] = 0.0
    f[3] = 1.0
    @views t1 = polynm(-1, 4, x[n-3:end], f, x[n-1])
    @views t2 = polynm(-1, 4, x[n-3:end], f, x[n])
    w[3,n] = t2 - t1

    f[3] = 0.0
    f[4] = 1.0
    @views t1 = polynm(-1, 4, x[n-3:end], f, x[n-1])
    @views t2 = polynm(-1, 4, x[n-3:end], f, x[n])
    w[4,n] = t2 - t1

    return
end


## Foil type
#==========================================================================================#

"""
    Foil(coords)

Airfoil structure consisting of foil coordinates as an array of points. Should be in Selig format for compatibility with other AeroMDAO tools.
"""
struct Foil{T <: Real}
    coords :: Vector{SVector{2,T}}
    name :: String
    Foil(coords :: Vector{SVector{2,T}}, name = "Unnamed") where T <: Real = new{T}(coords, name)
end

"""
    scale_foil(foil :: Foil, chord)

Scale the coordinates of a Foil, usually to some chord length.
"""
scale_foil(foil :: Foil, chord) = chord * foil.coords

"""
    cosine_foil(foil :: Foil, num :: Integer)

Return a Foil with cosine spacing for a given number of points.
"""
cosine_foil(foil :: Foil, num :: Integer) = cosine_foil(foil.coords, num)

"""
    camber_thickness(foil :: Foil, num :: Integer)

Compute the camber-thickness distribution of a Foil with cosine spacing..
"""
camber_thickness(foil :: Foil, num :: Integer) = foil_camthick(cosine_foil(foil.coords), num + 1)

function max_thickness_to_chord_ratio_location(coords)
    xs, thiccs = getindex.(coords, 1), getindex.(coords, 3)
    max_thick_arg = argmax(thiccs)
    xs[max_thick_arg], thiccs[max_thick_arg]
end

## Foil processing
#==========================================================================================#

"""
    read_foil(path :: String; header = true)

Read a '.dat' file consisting of 2D coordinates, for an airfoil as an array of `SVector`s, with an optional argument to skip the header.
"""
function read_foil(path :: String; header = true)
    coords = readdlm(path, skipstart = header ? 1 : 0)
    @views SVector.(coords[:,1], coords[:,2])
end

function split_foil(coords)
    for (i, ((xp, yp), (x, y), (xn, yn))) ∈ (enumerate ∘ adj3)(coords)
        if x < xp && x < xn
            if slope(x, y, xp, yp) >= slope(x, y, xn, yn)
                return splitat(i, coords)
            else
                return splitat(i, coords[end:-1:1,:])
            end
        end
    end
    (coords, [])
end

function paneller(foil :: Foil, num_panels :: Integer)
    coords = cosine_foil(foil.coords, Int(ceil(num_panels / 2)))
    @views Panel2D.(coords[2:end], coords[1:end-1])[end:-1:1]
end

"""
Discretises a foil profile into panels by projecting the x-coordinates of a circle onto the geometry.
"""
function cosine_foil(coords, n :: Integer = 40)
    upper, lower = split_foil(coords)
    n_upper = [ upper    ;
               [lower[1]]] # Append leading edge point from lower to upper
    upper_cos, lower_cos = cosine_interp(n_upper[end:-1:1], n), cosine_interp(lower, n)

    @views [ upper_cos[end:-1:2] ;
             lower_cos           ]
end

## Class shape transformation method
#==========================================================================================#

# Basic shape function
function shape_function(x, basis_func, coeffs, coeff_LE = 0)
    n     = length(coeffs)
    terms = basis_func.(x, n - 1, 0:n-1)
    dot(coeffs, terms) + coeff_LE * (x^0.5) * (1 - x)^(n - 0.5)
end

# Computing coordinates
cst_coords(class_func, basis_func, x, alphas, dz, coeff_LE, args...) = class_func(x) * shape_function(x, basis_func, alphas, coeff_LE) + x * dz

## Bernstein basis
#==========================================================================================#

bernstein_class(x, N1, N2) = x^N1 * (1 - x)^N2
bernstein_basis(x, n, k)   = binomial(n, k) * bernstein_class(x, k, n - k)

"""
    kulfan_CST(alpha_u, alpha_l,
               (dz_u, dz_l) :: NTuple{2, Real},
               coeff_LE = 0,
               n :: Integer = 40)

Define a cosine-spaced airfoil with ``2n`` points using the Class Shape Transformation method on a Bernstein polynomial basis with arrays of coefficients ``(\\alpha_u,~ \\alpha_l)`` for the upper and lower surfaces, trailing-edge spacing values ``(\\Delta z_u,~ \\Delta z_l)``, and support for leading edge modifications.
"""
function kulfan_CST(alpha_u, alpha_l, (dz_u, dz_l), coeff_LE = 0, n :: Integer = 40, N1 = 0.5, N2 = 1.)
    # Cosine spacing for airfoil of unit chord length
    xs = cosine_spacing(0.5, 1, n)

    # λ-function for Bernstein polynomials
    bernie(x, alphas, dz) = cst_coords(y -> bernstein_class(y, N1, N2), bernstein_basis, x, alphas, dz, coeff_LE)

    # Upper and lower surface generation
    upper_surf = [ bernie(x, alpha_u, dz_u) for x ∈ xs ]
    lower_surf = [ bernie(x, alpha_l, dz_l) for x ∈ xs ]

    # Counter-clockwise ordering
    @views [ SVector.(xs, upper_surf)[end:-1:2] ;
             SVector.(xs, lower_surf)           ]
end

function camber_CST(alpha_cam, alpha_thicc, (dz_cam, dz_thicc), coeff_LE = 0, n :: Integer = 40, N1 = 0.5, N2 = 1.)
    # Cosine spacing for airfoil of unit chord length
    xs = cosine_spacing(0.5, 1, n)

    # λ-function for Bernstein polynomials
    bernie(x, alphas, dz) = cst_coords(y -> bernstein_class(y, N1, N2), bernstein_basis, x, alphas, dz, coeff_LE)

    # Upper and lower surface generation
    cam   = [ bernie(x, alpha_cam, dz_cam) for x ∈ xs ]
    thicc = [ bernie(x, alpha_thicc, dz_thicc) for x ∈ xs ]

    camthick_foil(xs, cam, thicc)
end

function coords_to_CST(coords, num_dvs)
    xs 		 = first.(coords)
    S_matrix = reduce(hcat, @. bernstein_class(xs, 0.5, 1.0) * bernstein_basis(xs, num_dvs - 1, i) for i in 0:num_dvs - 1)

    alphas 	 = S_matrix \ last.(coords)

    return alphas
end

function camthick_to_CST(coords, num_dvs)
    xs, camber, thickness = (columns ∘ foil_camthick)(coords)

    alpha_cam  	= coords_to_CST(SVector.(xs, camber), num_dvs)
    alpha_thick = coords_to_CST(SVector.(xs, thickness), num_dvs)

    alpha_cam, alpha_thick
end

## Camber-thickness representation
#==========================================================================================#

"""
Converts an airfoil to its camber-thickness representation in cosine spacing.
"""
function foil_camthick(coords, num :: Integer = 40)
    upper, lower = split_foil(cosine_foil(coords, num))

    xs, y_LE  	 = first.(lower), first(lower)[2]   # Getting abscissa and leading edge ordinate
    y_upper, y_lower = last.(upper[end:-1:1]), last.(lower[2:end]) # Excluding leading edge point

    camber    	 = [ y_LE; (y_upper + y_lower) / 2 ]
    thickness 	 = [ 0; y_upper - y_lower ]

    SVector.(xs, camber, thickness)
end

"""
    camthick_foil(xs, camber, thickness)

Converts the camber-thickness representation to coordinates given the ``x``-locations and their corresponding camber and thickness values.
"""
function camthick_foil(xs, camber, thickness)
    coords = [ [xs camber + thickness / 2][end:-1:2,:];
                xs camber - thickness / 2 			  ]
    @views SVector.(coords[:,1], coords[:,2])
end

camthick_foil(coords) = @views camthick_foil(getindex.(coords, 1), getindex.(coords, 2), getindex.(coords, 3))

"""
    camber_coordinates(coords :: Array{2, <: Real})

Provides the camber coordinates on the ``x``-``z`` plane at ``y = 0``.
"""
camber_coordinates(coords) = @views SVector.(getindex.(coords, 1), 0, getindex.(coords, 2))

"""
    thickness_coordinates(coords :: Array{2, <: Real})

Provides the thickness coordinates on the ``x``-``z`` plane at ``y = 0``.
"""
thickness_coordinates(coords) = @views SVector.(getindex.(coords, 1), 0, getindex.(coords, 3))

## NACA Parametrisation
#==========================================================================================#

# NACA 4-digit parameter functions
naca4_thickness(t_by_c, xc, sharp_trailing_edge :: Bool) = 5 * t_by_c * (0.2969 * √xc - 0.1260 * xc - 0.3516 * xc^2 + 0.2843 * xc^3 - (ifelse(sharp_trailing_edge, 0.1036, 0.1015) * xc^4))
naca4_camberline(pos, cam, xc) = ifelse(xc < pos, (cam / pos^2) * xc * (2 * pos - xc), cam / (1 - pos)^2 * ( (1 - 2 * pos) + 2 * pos * xc - xc^2) )
naca4_gradient(pos, cam, xc) = atan(2 * cam / (ifelse(xc < pos, pos^2, (1 - pos)^2)) * (pos - xc))

"""
    naca4(digits :: NTuple{4, <: Real}, n :: Integer; sharp_trailing_edge :: Bool)

Defines a NACA 4-digit series profile.
"""
function naca4(digits :: NTuple{4, <: Real}, n :: Integer = 40; sharp_trailing_edge = true)
    # Camber
    cam = digits[1] / 100
    # Position
    pos = digits[2] / 10
    # Thickness-to-chord ratio
    t_by_c = (10 * digits[3] + digits[4]) / 100

    # Cosine spacing
    xs = cosine_spacing(0.5, 1.0, n)

    # Thickness distribution
    thickness = naca4_thickness.(Ref(t_by_c), xs, Ref(sharp_trailing_edge))

    if pos == 0 || cam == 0
        x_upper = xs
        y_upper = thickness
        x_lower = xs
        y_lower = -thickness
    else
        # Compute camberline
        camber 	= naca4_camberline.(Ref(pos), Ref(cam), xs)
        # Compute gradients
        grads 	= naca4_gradient.(Ref(pos), Ref(cam), xs)
        # Upper surface
        x_upper = @. xs - thickness * sin(grads)
        y_upper = @. camber + thickness * cos(grads)
        # Lower surface
        x_lower = @. xs + thickness * sin(grads)
        y_lower = @. camber - thickness * cos(grads)
    end
    coords = [ [x_upper y_upper][end:-1:2,:];
                x_lower y_lower             ]

    @views SVector.(coords[:,1], coords[:,2])
end


my_f(x,y) = 2*x +3y
my_sigmoid(z) = 1/(1 + exp(-z))
function my_probability(z)
    if my_sigmoid(z)>= 0.5
        return 1
    else
        return 0
    end
end


module RayTrace # main model that calls on all the other modules

using Shapes, Vector3D

export trace! # nice, just one funnction to export

type Ray # the ray type
    pos::Vec # has a position, a location
    dir::Vec # and a direction, which has to be unitized, norm(dir) = 1
end

type Bull # this type is really just a container for the following fields. It can be thought of as a ray in the gradient lens.
    ray::Ray # the ray, note its origin is now the center of the lens
    N::Vec # the normal to the refractive surface
    dis::Float64 # the distance to the lens center
    n1::Float64 # the refractive index at the previous position
end

# functions:
# initiates the Bull type.
function Bull(ray::Ray,lens::Lens)
    ray.pos -= lens.c # set ray's origin on the lens' center
    return Bull(ray,Vec(0.,0.,0.),0.,0.)
end

# initiates the array of Rays. It starts them off from directly above the lens at distance L. Their directions are uniformly and randomly distributed across the lens.
function letBlight(lens::Lens,L::Float64)
    if isinf(L) # if the light source is infinitely far away
        # just ditribute the rays uniformly on the disk that the sphere makes:
        sqrtr = sqrt(rand()) # this is taken from http://mathworld.wolfram.com/DiskPointPicking.html
        θ = rand()*2π
        a = Vec([sqrtr*cos(θ),sqrtr*sin(θ),2lens.r]) # initiate the Ray with those starting points at distance 2*rad (could be any distance larger than rad really
        pos = lens.c + a # adjust in case the lens center is not at 0,0,0
        dir = Vec([0.,0.,-1.]) # just point them directly down
    else
        cosα = cos(asin(lens.r/L)) # theta is the angle between the optical axis and the ray that touches the surface of the sphere

        # z and α are randomly sampled from those ranges to generate the uniformly spaced points on the sphere
        z = rand()*(cosα - 1.) - cosα
        θ = rand()*2π

        # the x, y, and z that lie on the sphere's surface
        x = sqrt(1. - z*z)*cos(θ)
        y = sqrt(1. - z*z)*sin(θ)

        # this is the starting position of all the rays -- point source
        a = Vec([0.,0.,L])
        pos = lens.c + a
        dir = unitize(Vec([x,y,z]))
    end
    return Ray(pos,dir) # the ray type
end

# This function finds the intersection point of the light as it starts from its initial point. It's a basic vector-sphere intersection thing. It works. I copied it from the ray_sphere function in https://github.com/JuliaLang/julia/blob/master/test/perf/kernel/raytracer.jl
function intersect!(ray::Ray,s::Sphere)
    v = s.c - ray.pos
    b = dot(v, ray.dir)
    disc = b*b - dot(v, v) + s.r*s.r
    if disc >= 0.
        d = sqrt(disc)
        t2 = b + d
        if t2 >= 0.
            t1 = b - d
            ray.pos += (t1 > 0. ? t1 : t2)*ray.dir
        end
    end
end

# the refraction/reflection function, ray is ray, N is the normal to the refractive durface, n is the ratio between the refractive index before and afer the surface.
function refract!(ray::Ray,N::Vec,n::Float64)
    # this is the 'r' in page 74 in that pdf I just mentioned
    a = -dot(ray.dir,N)

    # this is the 'c' in page 74
    c = 1. - n*n*(1. - a*a)

    # and finally the 't' in page 74
    # added reflection as well... See page 54 (v is -ray.dir)
    t = c>0 ? n*ray.dir + (n*a - sqrt(c))*N : 2* dot(N,-ray.dir) *N + ray.dir

    # unitize the direction vector for the ray
    ray.dir = unitize(t)
end

# a step function. moves the ray forward one ε. Notice that the ray is now packed inside Bull, and is normalized to the lens center
function step!(bull::Bull,lens::Lens)
    bull.n1 = lens.ri(norm(bull.ray.pos)) # first I find the refractive index at the current position

    bull.ray.pos += lens.ε*bull.ray.dir # then I move the ray one ε forward with its original direction

    bull.dis = norm(bull.ray.pos) # get the distance to the lens center from the new location
    bull.N = bull.ray.pos/bull.dis # find the new location's normal
end

# bend the ray accordingly to the refraction/reflection.
function bend!(bull::Bull,lens::Lens)
    n2 = lens.ri(bull.dis) # new refractive index here

    # here I find out if the ray is on it's way in or way out. If its distance to the lens center is diminishing then the normal to the refractive surface is going outwards. But if the ray is leaving the lens the normal to the surface is pointing inwards.
    costheta = dot(bull.ray.dir,bull.ray.pos)

    out = -sign(costheta) # is it going out? If so make out equal to +1, otherwise make it -1

    bull.N = out*bull.ray.pos/bull.dis # this is the normal to the refractive surface

    refract!(bull.ray,bull.N,bull.n1/n2) # refraction

end

# main tracing function. p is a matrix where all the ray's locations are saved. lens is lens, retina is retina, L is the distance to the light source, n_medium you know, maxiterATION, diagnose is a Dict that contains some diagnostics.
function trace!(p::Array{Float64,2},lens::Lens,retina::Retina,L::Float64,n_medium::Float64,maxiter::Int,diagnose::Dict)
    ray = letBlight(lens,L) # initiate all the rays

    p[:,1] = myconvert(ray.pos) # populate with the start position

    intersect!(ray,lens) # find the intersection point of the ray with the lens and update ray's position

    p[:,2] = myconvert(ray.pos) # add said intersection to the position matrix

    refract!(ray,unitize(ray.pos - lens.c),n_medium/lens.ri(lens.r)) # refract at the surface of the lens due to the medium's refractive index (that might be different from that of the lens' periphery)
    bull = Bull(ray,lens) # make a bull instance
    for i = 3:maxiter-1 # go for maximum maxiterATIONS
        step!(bull,lens) # advance one ε
        if bull.dis > lens.r # is the ray exiting the lens? If so, intersect it with the lens (which is therefore closer)
            ray = bull.ray # extract the ray from bull
            ray.pos -= lens.ε*ray.dir-lens.c # correct for the mistake of exiting the lens AND the lens center
            intersect!(ray,lens) # intersect with the lens
            p[:,i] = myconvert(ray.pos) # add position to the position matrix
            diagnose[:Δ] = sqrt(sum((diagnose[:goal]-vec(p[:,i])).^2)) # find distance to the goal

            break # and break the loop, cause we're done with the lens
        end # if the ray didn't exit the lens:
        p[:,i] = myconvert(bull.ray.pos + lens.c) # save the new position in the positions matrix
        bend!(bull,lens) # and refract
    end
    refract!(ray,-unitize(ray.pos - lens.c),lens.ri(lens.r)/n_medium) # refract when exiting the lens and entering the surrounding medium
    intersect!(ray,retina) # intersection point with the retina
    i += 1 # just for the position matrix
    p[:,i] = myconvert(ray.pos) # add said position to the position matrix
    diagnose[:len] = i # and set the length of the ray refraction points in the positions matrix
end

end


########################################################
function lsststd(; zpwave=3670.7, z=0.5, lum=5e43) # Angstrom, unitless, erg/s
########################################################

    #Define QSO variability-luminosity-redshift relation
    #Morganson et al. (2014) ApJ 784:92
    #input luminosity in erg/s
    #See also Roberto Cid Fernandes, Jr, et al. (1996)
    angstonm = 0.1
    zpwave   = zpwave*angstonm
    sigmavar = 0.079*(1+z)^0.15 * (lum/1e46)^-0.2 * (zpwave/1000.)^-0.44
    return sigmavar
end

########################################################
function setlsstwavelengths(λ = [3670.7, 4826.9, 6223.2, 7546.0, 8690.0, 9710.3]) # Angstrom
########################################################

    lsstwaves() =  λ

end

########################################################
function unitgalaxyvector(T)
########################################################

    G = galaxyvector(T)

    G / norm(G)

end

########################################################
function galaxyvector(T)
########################################################

    G = map(λ -> B(; T=T, λ=λ), lsstwaves())

end

########################################################
function B(; λ=λ, T=T)   # λ is in Angstrom, T in Kelvin
########################################################

    h = 6.6261e-27      # Planck constant    [cgs]
    c = 2.9979245800e10 # Speed of light     [cm/s]
    k = 1.380648813e-16 # Boltzmann constant [cgs]

    wavcm = λ*1e-8      # convert Ang to cm
    nu    = c/wavcm     # convert wavelength to frequency [Hz]

    #--frequency units--
    # emittence [erg/sec/cm^2/Hz]

    In   = 2.0*pi*h*nu^3.0/c^2 * (1.0/(exp(h*nu/(k*T)) - 1.0))

    ImJy = In / 1.0e-26

    return ImJy # Flux in absolute units of milli jansky
                # Doesn't not contain observational effects
end


@block AnthonyWang ["implicits"] begin
    @cell "Implicit equation" ["2d", contour, implicit] begin
        r = range(-5, stop = 5, length = 400)
        (a, b) = -1, 2
        z = ((x,y) -> y.^4 - x.^4 + a .* y.^2 + b .* x.^2).(r, r')
        z2 =  z .* (abs.(z) .< 250)
        contour(r, r, z2)
    end

    @cell "Cube lattice" ["3d", contour, implicit, colormap] begin
        r = range(-3, stop = 3, length = 100)
        me = [((1 ./ x).^2 + (1 ./ y).^2 + (1 ./ z).^2) for x=r, y=r, z=r]
        me2 = me .* (abs.(me) .> 1.5)
        contour(me2, color = :Set2)
    end

    @cell "3D cube with sphere cutout, inside" ["3d", volume, implicit] begin
        scene = Scene()
        r = range(-1, stop = 1, length = 100)
        mat = [(x.^2 + y.^2 + z.^2) for x = r, y = r, z = r]
        mat2 = mat .* (mat .> 1.4)
        #plot the space inside
        volume(mat2, algorithm = :absorptionrgba)
    end

    @cell "3D cube with sphere cutout, outside" ["3d", volume, implicit] begin
        scene = Scene()
        r = range(-1, stop = 1, length = 100)
        mat = [(x.^2 + y.^2 + z.^2) for x=r, y=r, z=r]
        mat2 = mat .* (mat .< 1.4)
        #plot the space outside
        volume(50..100, 50..100, 50..100, mat2, algorithm = :absorptionrgba)
    end

    @cell "Biohazard" ["3d", volume, implicit, algorithm, absorption] begin
        (a, b) = -1, 2
        r = range(-5, stop = 5, length = 100)
        z = ((x,y) -> y.^4 - x.^4 + a.*y.^2 + b.*x.^2).(r, r')
        me = [cos.(2 .* pi .* sqrt.(x.^2 + y.^2)) .* (4 .* z) for x=r, y=r, z=r]
        me2 = me .* (abs.(me) .> z .* 3)
        volume(me2, algorithm = :absorptionrgba)
    end

    @cell "Twisty cube thing" ["3d", implicit, contour, colormap, colorrange] begin
        (a, b) = -1, 2
        r = range(-2, stop = 2, length = 100)
        z = ((x,y) -> x + y).(r, r') ./ 5
        me = [z .* sin.(3 .* (atan.(y ./ x) .+ z.^2 .+ pi .* (x .> 0))) for x=r, y=r, z=r]
        me2 = me .* (me .> z .* 0.25)
        contour(me2, levels = 6, colormap = :Spectral)
    end

    @cell "Spacecraft from a galaxy far, far away" ["3d", implicit, contour, surface, inequalities, colormap, colorrange] begin
        N = 100
        r = range(-1, stop = 1, length = N)

        # bunch of equations and inequalities
        f1(x,y,z) = x.^2 .+ y.^2 .+ z.^2 #center sphere
        f2(x,y,z) = y.^2 .+ z.^2 #command deck cylinder thing
        f3(x,y,z) = x.^2 .+ 4 .* y.^2 #controls the flattened cylinder connecting center pod to wings
        f4(x,y,z) = (y .* 0.7 .+ 0.05) #defines the diagonal spokes
        f5(x,y,z) = (y .* 0.7 .- 0.05) #defines the diagonal spokes
        f6(x,y,z) = abs.(x) + 0.3 .* abs.(y) #frame part of the wings

        e1(x,y,z) = 0.12 .* (1 .- abs.(z)) #limits of a hexagonal tube in the inside of the craft
        e2(x,y,z) = abs.(z) .* (abs.(z) .< 0.95) #outer limits of the wing plane
        e3(x,y,z) = abs.(z) .* (abs.(z) .> 0.9) #inner limits of the wing plane
        e4(x,y,z) = (abs.(x) + abs.(0.3 .*y)) .* ((abs.(x) + abs.(0.3 .* y)) .< 1) #frame of the wings
        e5(x,y,z) = abs.(z) .* (abs.(z) .< 1.05) #outside thickness of wing frames, including the spokes
        e6(x,y,z) = abs.(z) .* (abs.(z) .> 0.80) #inside thickness of wing frames, including the spokes
        e7(x,y,z) = abs.(x) .* (abs.(x) .< 0.7) #length of the straight bars part of frames
        e8(x,y,z) = abs.(y) .* (abs.(y) .> 0.9) #width of the straight bars part of frames
        e9(x,y,z) = abs.(y) .* (abs.(y) .< 0.035) #the thickness of the horizontal reinforcing bar on the wing planes

        amp = 15 #this just amplifies the "strength" of a volume, so that it shows up more clearly in the plot

        # spawn the tie fighter
        me = [(f1(x,y,z) .* f1(x,y,z).<0.2) .+ ((f2(x,y,z) .* f2(x,y,z).<0.02).*((x.<0.68).*(x.>0.50))) .+ amp .* (f3(x,y,z) .* (f3(x,y,z) .< e1(x,y,z))) .+ (e2(x,y,z).*e3(x,y,z).*e4(x,y,z)) .+ (e5(x,y,z).*e6(x,y,z)).*((e7(x,y,z)).*(e8(x,y,z)) .+ e9(x,y,z) .+ ((x.>f5(x,y,z)).*x).*((x.<f4(x,y,z)).*x) .+ (((-).(x).>f5(x,y,z)).*x).*(((-).(x).<f4(x,y,z)).*x) .+ ((f6(x,y,z).*(f6(x,y,z).<1.05)).*(f6(x,y,z).*(f6(x,y,z).>0.95)))) for x=r, y=r, z=r]

        me2 = me
        for i = 1:length(r)
            me2[:,:,i] = me2[:,:,i] .- min(me2[:,:,i]...)
            me2[:,:,i] = me2[:,:,i] ./ max(me2[:,:,i]...)
        end
        contour(me2, colormap = :Purples, colorrange = (0.001,0.6), alpha = 1.0, levels = 10)
    end

end


module ConvolutionTools

import Base.conv

export conv, conv_valid, pad_const

"""
    pad_const(a, n[, c])

Pads any multidimensional array with `n[i]` constant values `c` on both sides.
Returns an array of size `size(a) + 2*n`.
"""
function pad_const(a::Array, n::Array{Int, 1}, c=0)
    shape = collect(size(a)) .+ 2 .* n
    b = zeros(eltype(a), shape...) .+ c
    block = [(n[i] + 1):(size(a)[i] + n[i]) for i in 1:length(n)]
    b[block...] = a
    return b
end

function conv{T<:Base.LinAlg.BlasFloat, N}(u::Array{T, N}, v::Array{T, N})
    su = [size(u)...]
    sv = [size(v)...]
    sp = su .+ sv .- 1 # for zero-padding
    # pad for efficient FFT
    sp2 = [n > 1024 ? nextprod([2,3,5], n) : nextpow2(n) for n in sp]
    upad = zeros(T, (sp2...))
    upad[[1:n for n in su]...] = u
    vpad = zeros(T, (sp2...))
    vpad[[1:n for n in sv]...] = v
    if T <: Real
        p = plan_rfft(upad)
    else
        p = plan_fft(upad)
    end
    # perform convolution
    y = p \ ((p * upad) .* (p * vpad))
    # undo FFT padding
    full = [1:sp[i] for i in 1:N]
    return y[full...]
end
conv{T<:Integer, N}(u::Array{T, N}, v::Array{T, N}) = round(Int, conv(float(u), float(v)))
conv{T<:Integer, S<:Base.LinAlg.BlasFloat, N}(u::Array{T, N}, v::Array{S, N}) = conv(float(u), v)
conv{T<:Integer, S<:Base.LinAlg.BlasFloat, N}(u::Array{S, N}, v::Array{T, N}) = conv(u, float(v))

function strip_invalid{T, S, N}(output::Array{T, N}, stencil::Array{S, N})
    shape_padded = [size(output)...]
    shape_stencil = [size(stencil)...]
    shape_valid = shape_padded .- 2 .* (shape_stencil .- 1)
    invalid = shape_stencil .- 1
    lower = invalid .+ 1
    upper = invalid .+ shape_valid
    center = [lower[i]:upper[i] for i in 1:N]
    return output[center...]
end

function conv_valid(data, stencil)
    full = conv(data, stencil)
    return strip_invalid(full, stencil)
end

end # module


import COBRA

function lpproblem(S::AbstractMatrix, b::AbstractVector, c::AbstractVector,
                   lb::AbstractVector, ub::AbstractVector)
    @assert length(b) == size(S,1)
    @assert length(lb) == length(ub) == length(c) == size(S,2)
    osense = -1 # maximize
    csense = fill('=', length(b))
    mets = "M" .* string.(1 : size(S,1))
    rxns = "R" .* string.(1 : size(S,2))
    COBRA.LPproblem(S, b, c, lb, ub, osense, csense, rxns, mets)
end

"""
Converts a MetNet to COBRA.LPProblem.
The MetNet struct is defined in
https://github.com/anna-pa-m/Metabolic-EP
"""
function lpproblem(metnet)
    lpproblem(metnet.S, metnet.b, metnet.c, metnet.lb, metnet.ub)
end


"""
Kinetic energy: 1/2 sum_n f_n ∫ |∇ψn|^2.
"""
struct Kinetic
    scaling_factor::Real
end
Kinetic(; scaling_factor=1) = Kinetic(scaling_factor)
(kin::Kinetic)(basis) = TermKinetic(basis, kin.scaling_factor)
function Base.show(io::IO, kin::Kinetic)
    fac = isone(kin.scaling_factor) ? "" : ", scaling_factor=$scaling_factor"
    print(io, "Kinetic($fac)")
end

struct TermKinetic <: Term
    scaling_factor::Real  # scaling factor, absorbed into kinetic_energies
    kinetic_energies::Vector{<:AbstractVector}  # kinetic energy 1/2|G+k|^2 for every kpoint
end
function TermKinetic(basis::PlaneWaveBasis{T}, scaling_factor) where {T}
    kinetic_energies = [[T(scaling_factor) * sum(abs2, Gk) / 2
                         for Gk in Gplusk_vectors_cart(basis, kpt)]
                        for kpt in basis.kpoints]
    TermKinetic(T(scaling_factor), kinetic_energies)
end

@timing "ene_ops: kinetic" function ene_ops(term::TermKinetic, basis::PlaneWaveBasis{T},
                                            ψ, occ; kwargs...) where {T}
    ops = [FourierMultiplication(basis, kpoint, term.kinetic_energies[ik])
           for (ik, kpoint) in enumerate(basis.kpoints)]
    isnothing(ψ) && return (E=T(Inf), ops=ops)

    E = zero(T)
    for (ik, k) in enumerate(basis.kpoints)
        for iband = 1:size(ψ[1], 2)
            ψnk = @views ψ[ik][:, iband]
            E += (basis.kweights[ik] * occ[ik][iband]
                  * real(dot(ψnk, Diagonal(term.kinetic_energies[ik]), ψnk)))
        end
    end
    E = mpi_sum(E, basis.comm_kpts)

    (E=E, ops=ops)
end


#Stochastic collocation towards a polynomial and the lognormal density
using Polynomials
using Roots
import AQFED.Math: normcdf, normpdf, norminv
import AQFED.Black: blackScholesFormula, blackScholesVega, impliedVolatility
using AQFED.Bachelier
using LeastSquaresOptim
struct LogPolynomial
    p::AbstractPolynomial
    σ::Number
end
struct IsotonicLogCollocation
    p1::AbstractPolynomial
    p2::AbstractPolynomial
    σ::Number
    forward::Number
end

function LogPolynomial(iso::IsotonicLogCollocation; minSlope = 0.0)
    q = degree(iso.p1) + degree(iso.p2)
    p = integrate(iso.p1^2 + Polynomials.Polynomial([0.0, 1.0])*iso.p2^2 + minSlope * iso.forward)
    theoForward = hermiteIntegral(LogPolynomial(p,iso.σ))
    c =  iso.forward / theoForward
    p = p*c
    return LogPolynomial(p,iso.σ)
end

(p::LogPolynomial)(x::S) where {S} = p.p(x)

function solvePositiveStrike(p::LogPolynomial, strike::Number; useNumerical=true)::Number
    if useNumerical

        return find_zero(
            x -> p(x) - strike,
            (1e-7,1e5),
            Bisection()
        )
    else
        r = roots(pk)
        for ri in r
            if (abs(imag(ri)) < eps(strike)) && real(ri) >= 0
                return real(ri)
            end
        end
        throw(DomainError(strike, "no real roots at the given strike"))
    end
end

function priceEuropean(p::LogPolynomial, isCall::Bool, strike::Number, forward::Number, discountDf::Number)::Number
    ck = solvePositiveStrike(p, strike)
    if ck < 0
        throw(DomainError(ck, "expected a positive value corresponding to strike"))
    end
    logck = log(ck)/p.σ
    valuef = hermiteIntegralBounded(p, logck)
    valuek = normcdf(-logck)
    callPrice = valuef - strike * valuek
    putPrice = -(forward - strike) + callPrice
    if isCall
        return callPrice * discountDf
    else
        return putPrice * discountDf
    end
end

function density(p::LogPolynomial, strike::Number)::Number

    ck = solvePositiveStrike(p, strike)
    if ck < 0
        throw(DomainError(ck, "expected a positive value corresponding to strike"))
    end
    logck = log(ck)/p.σ
    #Phi(log(ck)/p.sigma) '
    dp = derivative(p.p)(ck)
    an = normpdf(logck) / ( dp * ck*p.σ )
    #num = (-2*priceEuropean(p, true, strike, 0.0, 1.0) +priceEuropean(p, true, strike+1e-4, 0.0, 1.0) +priceEuropean(p, true, strike-1e-4, 0.0, 1.0) )/(1e-8)
    return an
end

function rawMoment(p::LogPolynomial, moment::Int)::Number
    if moment == 0
        return 1.0
    end
    q = p.p
    for i = 2:moment
        q *= p.p
    end
    return hermiteIntegral(LogPolynomial(q,p.σ))
end

#return mean, standard dev, skew, kurtosis of the collocation
function stats(p::LogPolynomial)
    μ = hermiteIntegral(p)
    μ2 = hermiteIntegral(LogPolynomial((p.p - μ)^2,p.σ))
    μ3 = hermiteIntegral(LogPolynomial((p.p - μ)^3,p.σ))
    μ4 = hermiteIntegral(LogPolynomial((p.p - μ)^4,p.σ))

    skew = μ3 / μ2^1.5
    kurtosis = μ4 / μ2^2
    return μ, sqrt(μ2), skew, kurtosis
end

function hermiteIntegral(p::LogPolynomial)::Number
    c = coeffs(p.p)
    value = sum(c[i]*exp(((i-1)*p.σ)^2 / 2) for i=1:length(c))
    return value
end

function hermiteIntegralBounded(p::LogPolynomial, logck::Number)::Number
    c = coeffs(p.p)
    return sum(c[i]*exp(((i-1)*p.σ)^2 / 2)*normcdf(-logck+(i-1)*p.σ) for i=1:length(c))
end

function makeLogXFromUndiscountedPrices(strikesf::Vector{T}, pricesf::Vector{T}, σ::T) where {T}
    return makeXFromUndiscountedPrices(strikesf, pricesf, s -> exp(σ*norminv(1+s)))
end

function makeIsotonicLogCollocationGuess(
    strikes::Vector{T},
    callPrices::Vector{T},
    weights::Vector{T},
    τ::T,
    forward::T,
    discountDf::T;
    deg = 3, #1 = Black (trivial, smoother if least squares iterations small). otherwise 3 is good.
) where{T}
    b = fitLogBlack(strikes, callPrices, weights, τ, forward, discountDf)
    if deg >= 3
        strikesf, pricesf, weightsf = filterConvexPrices(strikes, callPrices ./ discountDf, weights, forward)
        strikesf, pif, xif = makeLogXFromUndiscountedPrices(strikesf, pricesf, b.σ)
        # cubic = Polynomials.fit(xif, strikesf, 3; weights = weightsf) #FIXME origin must be = 0
        # if !isCubicMonotone(cubic)
        #     cubic = LeastSquaresCubicMurrayForward(xif, strikesf, weightsf, forward)
        # end
        # if deg == 3
        #     return IsotonicLogCollocation(cubic, b.σ, forward)
        # else
        #     return fitLogMonotonic(xif, strikesf, weightsf, forward, cubic, deg = deg)
        # end
        return fitLogMonotonic(xif, strikesf, weightsf, forward, b, deg = deg)
    else
        return b
    end
end

function makeIsotonicLogCollocation(
    strikes::Vector{T},
    callPrices::Vector{T},
    weights::Vector{T},
    τ::T,
    forward::T,
    discountDf::T;
    deg = 3, #degree of collocation. 5 is usually best from a stability perspective.
    degGuess = 3, #1 = Bachelier (trivial, smoother if least squares iterations small). otherwise 3 is good.
    minSlope = 1e-4, penalty=0.0
)::Tuple{IsotonicLogCollocation,Number} where {T} #return collocation and error measure
    isoc = makeIsotonicLogCollocationGuess(strikes, callPrices, weights, τ, forward, discountDf, deg = degGuess)
    #optimize towards actual prices
    isoc, m = fit(isoc, strikes, callPrices, weights, forward, discountDf, deg = deg, minSlope=minSlope,penalty=penalty)
    return isoc, m
end

function fitLogBlack(strikes, prices, weights, τ, forward, discountDf)
    m = length(strikes)
    i = findfirst(x -> x > forward, strikes)
    if i == nothing
        i = m
    elseif i == 1
        i = 2
    end
    price = (prices[i] * (forward - strikes[i-1]) + prices[i-1] * (strikes[i] - forward)) / (strikes[i] - strikes[i-1])
    strike = forward
    vol = impliedVolatility(true, price, forward, strike, τ, discountDf)
    σ = vol*sqrt(τ)
    isoc = IsotonicLogCollocation(Polynomials.Polynomial([sqrt(forward*exp(-σ^2 / 2))]), Polynomials.Polynomial([0.0]), σ, forward)
    return isoc
end

Base.length(p::LogPolynomial) = Base.length(p.p)
Base.iterate(p::LogPolynomial) = Base.iterate(p.p)
Base.iterate(p::LogPolynomial, state) = Base.iterate(p.p, state)
Base.broadcastable(p::LogPolynomial) = Ref(p)
Polynomials.degree(p::LogPolynomial) = Polynomials.degree(p.p)

isotonicDegree(deg::Int) = isodd(deg) ? trunc(Int, (deg-1)/2) : trunc(Int, deg / 2)

function fit(isoc::IsotonicLogCollocation, strikes, prices, weights, forward, discountDf; deg = 3, minSlope = 1e-6, penalty=0.0)
    q = isotonicDegree(deg)
    c0length = 2*q
    if isodd(deg)
        q+=1
        c0length = 2*q-1
    end
    iter = 0
    function obj(ct)
        c = ct
        p1 = Polynomials.Polynomial(c[1:q])
        p2 = Polynomials.Polynomial(c[q+1:c0length])
        isoc = IsotonicLogCollocation(p1, p2, isoc.σ, forward)
        p = LogPolynomial(isoc)
        iter += 1
        if penalty > 0
            ip = hermiteIntegral(LogPolynomial(derivative(p.p,2)^2,p.σ)) #does not seem appropriate here
            pvalue = penalty * ip
            return @. weights * sqrt((priceEuropean(p, true, strikes, forward, discountDf) - prices)^2 + pvalue^2)
        else
            return @. weights * (priceEuropean(p, true, strikes, forward, discountDf) - prices)
        end
    end
    c0 = zeros(Float64, c0length)
    c1 = coeffs(isoc.p1)
    for i = 1:min(q, length(c1))
        c0[i] = c1[i]
    end
    c2 = coeffs(isoc.p2)
    for i = 1:min(c0length-q, length(c2))
        c0[q+i] = c2[i]
    end
    for i = 1:length(c0)
        if c0[i] == 0
            c0[i] = eps(forward)  #zero would break automatic differentiation
        end
    end
    fit = optimize(obj, c0, LevenbergMarquardt(); show_trace = false, autodiff = :forward,iterations = deg * 300) #autodiff breaks without Halley. Would need custom

    println(iter, " fit ", fit, obj(fit.minimizer))
    c0 = fit.minimizer
    measure = fit.ssr
    return IsotonicLogCollocation(Polynomials.Polynomial(c0[1:q]), Polynomials.Polynomial(c0[q+1:c0length]), isoc.σ, forward), measure
end

function IsotonicLogCollocation(cguess::LogPolynomial, forward::Number)
    cubic = cguess.p
    if degree(cubic) == 1
        return IsotonicLogCollocation(Polynomials.Polynomial([sqrt(cguess.p[1])]), Polynomials.Polynomial([0.0]), cguess.σ, forward)
    elseif degree(cubic) > 3
        throw(DomainError(degree(cubic), "expected a cubic"))
    end
    e = (cubic[1] - cubic[2]^2 / (3 * cubic[3]))
    e = max(e, forward * 1e-6)
    p2 = Polynomials.Polynomial([sqrt(e)])
    p1 = Polynomials.Polynomial([cubic[2] / sqrt(3 * cubic[3]), sqrt(3 * cubic[3])])
    return IsotonicLogCollocation(p1, p2, cguess.σ, forward)
end

function fitLogMonotonic(xif, strikesf, w1, forward, cubic; deg = 3)
    q = isotonicDegree(deg)
    c0length = 2*q
    if isodd(deg)
        q+=1
        c0length = 2*q-1
    end
    iter = 0
    function obj(c)
        #c = vcat(0.0, ct)
        p1 = Polynomials.Polynomial(c[1:q])
        p2 = Polynomials.Polynomial(c[q+1:c0length])
        isoc = IsotonicLogCollocation(p1, p2, cubic.σ, forward)
        p = LogPolynomial(isoc)
        iter += 1
        v = @. w1 * (p(xif) - strikesf)
        #println(iter, " ", c, " ", p.p, sum(v.^2))
        return v
    end

    isocubic = typeof(cubic) == IsotonicLogCollocation ? cubic : IsotonicLogCollocation(cubic, forward)
    c1 = coeffs(isocubic.p1)
    c2 = coeffs(isocubic.p2)
    c0 = zeros(Float64,c0length)
    c0[1:min(q, length(c1))] = c1
    c0[q+1:min(c0length, q + length(c2))] = c2
    # println("c0 ",c0," isocub ",isocubic)
    xamax = xif[end]
    for i = 1:length(c0)
        if c0[i] == 0
            c0[i] = eps(forward) / (xamax)^(i - 1)  #zero would break automatic differentiation
        end
    end
    fit = optimize(obj, c0, LevenbergMarquardt(); show_trace = false, autodiff = :forward)
    c0 = fit.minimizer
    isoc = IsotonicLogCollocation(Polynomials.Polynomial(c0[1:q]), Polynomials.Polynomial(c0[q+1:c0length]), cubic.σ, forward)
    println(iter, " fitMonotonic ", LogPolynomial(isoc), " ", fit)
    return isoc
end


using ZChop

"""
    QFI(ρ, dρ [, abstol])

Numerically evaluate the quantum Fisher information for the matrix ρ given its derivative dρ wrt the parameter

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `ρ`:  Density matrix
    * `dhro`: Derivative wrt the parameter to be estimated
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI(ρ, dρ; abstol = 1e-5)

    # Get the eigenvalues and eigenvectors of the density matrix
    # We enforce its Hermiticity so that the algorithm is more efficient and returns real values

    eigval, eigvec = eigen(Hermitian(Matrix(zchop(ρ, 1e-10))))

    dim = length(eigval)
    tmpvec = Array{eltype(ρ)}(undef, size(ρ, 1))
    res = 0.
    tmp = 0.
    for m = 1:dim
        mul!(tmpvec, dρ, view(eigvec, :, m))
        for n = 1:dim
            tmp = eigval[n] + eigval[m]
            if tmp > abstol
                res += 2 * (1. / tmp) * abs2(dot(view(eigvec, :, n), tmpvec))
            end
        end
    end
    return res
end

"""
    QFI!(ρ, dρ [, abstol])

Numerically evaluate the quantum Fisher information for the matrix ρ given its derivative dρ wrt the parameter.
Can overwrite ρ.

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `ρ`:  Density matrix (gets overwritten)
    * `dhro`: Derivative wrt the parameter to be estimated
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI!(ρ, dρ; abstol = 1e-5)
    # Get the eigenvalues and eigenvectors of the density matrix
    # We enforce its Hermiticity so that the algorithm is more efficient and returns real values
    eigval, eigvec = eigen!(Hermitian(ρ))
    dim = length(eigval)
    tmpvec = Array{eltype(ρ)}(undef, size(ρ, 1))
    res = 0.
    tmp = 0.
    for m = 1:dim
        mul!(tmpvec, dρ, view(eigvec, :, m))
        for n = 1:dim
            tmp = eigval[n] + eigval[m]
            if tmp > abstol
                res += 2 * (1. / tmp) * abs2(dot(view(eigvec, :, n), tmpvec))
            end
        end
    end
    return res
end

function QFI(ρ::BlockDiagonal, dρ::BlockDiagonal; abstol = 1e-5)
    qfi = 0.
    for i in 1:nblocks(ρ)
        qfi += QFI(ρ.blocks[i], dρ.blocks[i])
    end
    return qfi
end

function QFI!(ρ::BlockDiagonal, dρ::BlockDiagonal; abstol = 1e-5)
    qfi = 0.
    for i in 1:nblocks(ρ)
        qfi += QFI!(ρ.blocks[i], dρ.blocks[i])
    end
    return qfi
end

"""
    QFI(state::State, [, abstol])

Numerically evaluate the quantum Fisher information for the State state.

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `state`: A State object
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI(state::State, abstol = 1e-5)
    QFI(state.ρ, state.dρ)
end

"""
    QFI!(state::State, [, abstol])

Numerically evaluate the quantum Fisher information for the State state,
overwriting one of its temporary variables.

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `state`: A State object
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI!(state::State, abstol = 1e-5)
    copy!(state._tmp1, state.ρ)
    QFI!(state._tmp1, state.dρ)
end

"""
    FI(ρ, dρ, op, [, tol])

Evaluates the Fisher information for the measurement operator op, for density matrix ρ, with
derivative dρ with respect to the parameter.

tol gives the tolerance to treat the denominator as 0 (and avoid NaNs). Default is `tol=1e-10`.
"""
function FI(ρ::BlockDiagonal, dρ::BlockDiagonal, op::BlockDiagonal)
    ed = eigen(op)
    return FI(ρ, dρ, ed)
end

function FI(ρ::Matrix, dρ::Matrix, measurement::Eigen; tol=1e-10)
    # We use Eq. (4) of Paris, Int. J. Quantum Inf. 7, 125 (2009)
    # We need the projectors onto the subspaces of the operator op
    fi = 0.0
    for i in 1:size(measurement.vectors, 1)
        evec = measurement.vectors[:, i]
        denominator = real(evec' * ρ * evec)

        # Empty blocks would cause a division by 0
        # We use a tolerance
        if denominator > tol
            fi += real(evec' * dρ * evec)^2 / denominator
        end
    end
    return fi
end

function FI(ρ::BlockDiagonal, dρ::BlockDiagonal, measurement::Array{Eigen}; tol=1e-10)
    # We use Eq. (4) of Paris, Int. J. Quantum Inf. 7, 125 (2009)
    # We need the projectors onto the subspaces of the operator op
    fi = 0.0
    for i in 1:nblocks(ρ)
        fi += FI(ρ.blocks[i], dρ.blocks[i], measurement[i], tol=tol)
    end
    return fi
end

FI(state::State, measurement::Eigen) = FI(state.ρ, state.dρ, measurement)
FI(state::State, measurement::Array{Eigen}) = FI(state.ρ, state.dρ, measurement)