## Constants

const OPEN = -1

struct Polyhedron{Dp,Tp,Td}
  vertices::Vector{Point{Dp,Tp}}
  edge_vertex_graph::Vector{Vector{Int32}}
  isopen::Bool
  data::Td
end

struct PolyhedronData
  vertex_to_planes::Vector{Vector{Int32}}
  vertex_to_original_faces::Vector{Vector{Int32}}
  vertex_to_parent_vertex::Vector{Int32}
  vertex_to_parent_edge::Vector{Tuple{Int32,Int32}}
  plane_to_vertex_to_distances::Vector{Vector{Float64}}
  plane_to_ref_plane::Vector{Int32}
  plane_to_ids::Vector{Int32}
end

# Constructors

function Polyhedron(stl::STL)
  𝓖 = compute_graph(stl)
  X = get_vertex_coordinates(stl)
  isopen = is_open_surface(stl)
  p = Polyhedron(X,𝓖;isopen)
  set_original_faces!(p,stl)
  p
end

function Polyhedron(
  vertices::AbstractVector{<:Point},
  graph::Vector{<:Vector},
  isopen,
  data)

  Polyhedron(collect(vertices),graph,isopen,data)
end

function Polyhedron(
  vertices::AbstractVector{<:Point},
  graph::Vector{<:Vector}
  ;isopen=false::Bool)

  Polyhedron(vertices,graph,isopen,polyhedron_data(length(vertices)))
end

function Polyhedron(
  vertices::AbstractVector{<:Point},
  graph::Vector{<:Vector},
  data)

  isopen = false
  Polyhedron(vertices,graph,isopen,data)
end

function Polyhedron(p::Polytope{2},vertices::AbstractVector{<:Point})
  if p == TRI
    e_v_graph = [[2,3],[3,1],[1,2]]
  elseif p == QUAD
    e_v_graph = [[2, 3],[4, 1],[1, 4],[3, 2]]
  else
    @unreachable
  end
  e_v_graph = map(i->Int32.(i),e_v_graph)
  data = polyhedron_data(p)
  Polyhedron(vertices,e_v_graph,data)
end

function Polyhedron(p::Polytope{3},vertices::AbstractVector{<:Point})
  if p == TET
    e_v_graph = [[2,4,3],[3,4,1],[1,4,2],[1,2,3]]
  elseif p == HEX
    e_v_graph = [
      [5, 2, 3],
      [6, 4, 1],
      [7, 1, 4],
      [8, 3, 2],
      [1, 7, 6],
      [2, 5, 8],
      [3, 8, 5],
      [4, 6, 7] ]
  else
    @unreachable
  end
  e_v_graph = map(i->Int32.(i),e_v_graph)
  data = polyhedron_data(p)
  Polyhedron(vertices,e_v_graph,data)
end

function Polyhedron(p::Polytope)
  Polyhedron(p,get_vertex_coordinates(p))
end

# Operations

function clip(poly::Polyhedron,Π;inside=true,inout=trues(length(Π)),boundary=nothing)
  p = poly
  for (i,Πi) in enumerate(Π)
    side = inside == inout[i] ? :left : :right
    p = clip(p,Πi,side;boundary)
    !isnothing(p) || break
  end
  p
end

function split_overlapping(p::Polyhedron,Π)
  p⁻ = clip(p,Π,:left,boundary=true)
  p⁺ = clip(p,Π,:right,boundary=true)
  p⁻,p⁺
end

function split_gapping(p::Polyhedron,Π)
  p⁻ = clip(p,Π,:left,boundary=false)
  p⁺ = clip(p,Π,:right,boundary=false)
  p⁻,p⁺
end

function clip(p::Polyhedron,Π,side;boundary=nothing::Union{Nothing,Bool},invert=false)
  @assert side ∈ (:right,:left)
  if boundary === nothing
    p⁻,p⁺ = split(p,Π,side;invert)
    p  = side == :left ? p⁻ : p⁺
  else
    invert = false
    if side == :right && !boundary
      invert = true
      side = :left
    elseif side == :left && boundary
      invert = true
      side = :right
    end
    p = clip(p,Π,side;invert)
  end
  p
end

function split(p::Polyhedron,Π,side=:both;invert=false)
  @assert side ∈ (:both,:left,:right)
  distances = get_plane_distances(get_data(p),Π)
  _sign = invert ? (-) : (+)
  distances = lazy_map(_sign,distances)
  smin = Inf
  smax = -Inf
  for v in 1:num_vertices(p)
    isactive(p,v) || continue
    smin = min(smin,distances[v])
    smax = max(smax,distances[v])
  end
  if smin ≥ 0
    return nothing,p
  end
  if smax < 0
    if has_coplanars(get_data(p),Π)
      add_plane!(get_data(p),Π)
    end
    return p,nothing
  end
  new_vertices = empty( get_vertex_coordinates(p) )
  in_graph = _copy(get_graph(p))
  if side == :both
    out_graph = _copy(get_graph(p))
  end
  ≶ = side ∈ (:both,:left) ? (≥) : (<)
  data = copy( get_data(p) )
  D = num_dims(p)
  for v in 1:num_vertices(p)
    isactive(p,v) || continue
    distances[v] ≶ 0 && continue
    for (i,vneig) in enumerate( get_graph(p)[v] )
      vneig ∉ (UNSET,OPEN) || continue
      distances[vneig] ≶ 0 || continue
      vertex = compute_intersection(p[v],distances[v],p[vneig],distances[vneig])
      push!( new_vertices, vertex )
      add_vertex!(data,v,vneig,Π)
      push!( in_graph, fill(UNSET,D) )
      in_graph[v][i] = num_vertices(p) + length(new_vertices)
      in_graph[end][1] = v
      if side == :both
        push!( out_graph, fill(UNSET,D) )
        ineig = findfirst( isequal(v), get_graph(p)[vneig] )
        out_graph[vneig][ineig] = num_vertices(p) + length(new_vertices)
        out_graph[end][1] = vneig
      end
    end
  end
  if side == :both
    p_out = split_postprocess!(out_graph,copy(data),new_vertices,p,distances,(<))
  else
    p_out = nothing
  end
  p_in = split_postprocess!(in_graph,data,new_vertices,p,distances,(≶))
  p⁻ = side ≠ :right ? p_in : p_out
  p⁺ = side ≠ :right ? p_out : p_in
  p⁻,p⁺
end

function simplexify(poly::Polyhedron{3})
  !isopen(poly) || return simplexify_surface(poly)
  vstart = fill(UNSET,num_vertices(poly))
  stack = Int32[]
  for v in 1:num_vertices(poly)
    isactive(poly,v) || continue
    vstart[v] == UNSET || continue
    vstart[v] = v
    empty!(stack)
    push!(stack,v)
    while !isempty(stack)
      vcurrent = pop!(stack)
      for vneig in get_graph(poly)[vcurrent]
        if vstart[vneig] == UNSET
          vstart[vneig] = v
          push!(stack,vneig)
        end
      end
    end
  end
  v_to_pv = get_data(poly).vertex_to_parent_vertex
  istouch = map( i -> falses(length(i)), get_graph(poly) )
  vertex_coordinates = get_vertex_coordinates(poly)
  T = Vector{Int32}[]
  X = vertex_coordinates
  for v in 1:num_vertices(poly)
    isactive(poly,v) || continue
    for i in 1:length(get_graph(poly)[v])
      !istouch[v][i] || continue
      istouch[v][i] = true
      vcurrent = v
      vnext = get_graph(poly)[v][i]
      while vnext != v
        inext = findfirst( isequal(vcurrent), get_graph(poly)[vnext] )
        !isnothing(inext) || break
        inext = ( inext % length( get_graph(poly)[vnext] ) ) + 1
        istouch[vnext][inext] = true
        vcurrent = vnext
        vnext = get_graph(poly)[vnext][inext]
        @assert vcurrent ≠ vnext
        vcurrent ≠ vnext || break
        if v ∉ (vstart[v],vcurrent,vnext)
          k = [vstart[v],v,vcurrent,vnext]
          push!(T,k)
        end
      end
    end
  end
  T,X
end

function simplexify_surface(poly::Polyhedron{3})
  stack = Int32[]
  v_to_pv = get_data(poly).vertex_to_parent_vertex
  istouch = map( i -> falses(length(i)), get_graph(poly) )
  T = Vector{Int32}[]
  for v in 1:num_vertices(poly)
    isactive(poly,v) || continue
    for i in 1:length(get_graph(poly)[v])
      !istouch[v][i] || continue
      istouch[v][i] = true
      vcurrent = v
      vnext = get_graph(poly)[v][i]
      vnext ∉ (OPEN,UNSET) || continue
      while vnext != v
        inext = findfirst( isequal(vcurrent), get_graph(poly)[vnext] )
        inext = ( inext % length( get_graph(poly)[vnext] ) ) + 1
        istouch[vnext][inext] = true
        vcurrent = vnext
        vnext = get_graph(poly)[vnext][inext]
        vnext ∉ (OPEN,UNSET) || break
        if v ∉ (vcurrent,vnext)
          k = [v,vcurrent,vnext]
          push!(T,k)
        end
      end
    end
  end
  T,get_vertex_coordinates(poly)
end

function simplexify_boundary(poly::Polyhedron{3},stl::STL)
  D = 3
  stack = Int32[]
  v_to_pv = get_data(poly).vertex_to_parent_vertex
  if isopen(poly)
    v_to_Π = get_data(poly).vertex_to_original_faces
  else
    v_to_Π = get_data(poly).vertex_to_planes
  end
  facedims = get_facedims(stl)
  facet_list = Int32[]
  istouch = map( i -> falses(length(i)), get_graph(poly) )
  vertex_coordinates = get_vertex_coordinates(poly)
  T = Vector{Int32}[]
  X = empty(vertex_coordinates)
  f_to_stlf = Int32[]
  for v in 1:num_vertices(poly)
    isactive(poly,v) || continue
    for i in 1:length(get_graph(poly)[v])
      !istouch[v][i] || continue
      istouch[v][i] = true
      vcurrent = v
      vnext = get_graph(poly)[v][i]
      vnext ∉ (OPEN,UNSET) || continue
      empty!(facet_list)
      append!( facet_list, v_to_Π[v] )
      filter!( i -> i > 0, facet_list)
      filter!( i -> facedims[i] == num_dims(stl), facet_list)
      filter!( i -> i ∈ v_to_Π[vnext], facet_list )
      while vnext != v
        inext = findfirst( isequal(vcurrent), get_graph(poly)[vnext] )
        inext = ( inext % length( get_graph(poly)[vnext] ) ) + 1
        istouch[vnext][inext] = true
        vcurrent = vnext
        vnext = get_graph(poly)[vnext][inext]
        vnext ∉ (OPEN,UNSET) || break
        filter!( i -> i ∈ v_to_Π[vnext], facet_list )
        !isempty(facet_list) || break
        if v ∉ (vcurrent,vnext)
          kx = (v,vcurrent,vnext)
          k = (1:D) .+ length(X)
          for x in kx
            push!(X,vertex_coordinates[x])
          end
          push!(T,k)
          push!(f_to_stlf,first(facet_list))
        end
      end
    end
  end
  T,X,f_to_stlf
end

function simplexify_boundary(::Nothing,args...)
  nothing,nothing,nothing
end

function simplexify(polys::AbstractVector{<:Polyhedron{Dp,Tp}}) where {Dp,Tp}
  T = Vector{Int32}[]
  X = Point{Dp,Tp}[]
  for poly in polys
    Ti,Xi = simplexify(poly)
    append!(T, map(i->i.+length(X),Ti) )
    append!(X,Xi)
  end
  T,X
end

function simplexify_boundary(
  polys::AbstractVector{<:Polyhedron{Dp,Tp}},
  stl::STL) where {Dp,Tp}

  T = Vector{Int32}[]
  X = Point{Dp,Tp}[]
  f_to_stlf = Int32[]
  for poly in polys
    Ti,Xi,f_to_f_i = simplexify_boundary(poly,stl)
    append!(T, map(i->i.+length(X),Ti) )
    append!(X,Xi)
    append!(f_to_stlf,f_to_f_i)
  end
  T,X,f_to_stlf
end

function simplexify_boundary(::NTuple{N,Nothing},args...) where N
  nothing,nothing,nothing,nothing
end

function simplexify_boundary(p::Tuple,labels::Tuple,stl::STL)
  P = eltype(p[1])
  Dp = num_dims(P)
  Tp = point_eltype(P)
  T = Vector{Int32}[]
  X = Point{Dp,Tp}[]
  f_to_stlf = Int32[]
  f_to_ios = Int8[]
  for (p_i,ios) in zip( p, labels )
    Ti, Xi, f_to_f_i = simplexify_boundary(p_i,stl)
    if !isnothing(Ti)
      append!(T, map(i->i.+length(X),Ti) )
      append!(X,Xi)
      append!(f_to_stlf,f_to_f_i)
      append!(f_to_ios, fill(Int8(ios),length(Ti)) )
    end
  end
  T,X,f_to_stlf,f_to_ios
end

function surface(poly::Polyhedron{3})
  T,X = simplexify_surface(poly)
  p = TRI
  c = array_cache(T)
  !isempty(T) || return 0.0
  sum( i -> measure(get_cell!(c,T,X,p,i)), 1:length(T) )
end

function surface(poly::Polyhedron{3},stl::STL)
  T,X, = simplexify_boundary(poly,stl)
  p = TRI
  c = array_cache(T)
  !isempty(T) || return 0.0
  sum( i -> measure(get_cell!(c,T,X,p,i)), 1:length(T) )
end

function surface(polys::AbstractVector{<:Polyhedron},args...)
  s = 0.0
  for p in polys
    s += surface(p,args...)
  end
  s
end

function volume(poly::Polyhedron{3})
  T,X = simplexify(poly)
  p = TET
  c = array_cache(T)
  sum( i -> measure(get_cell!(c,T,X,p,i)), 1:length(T) )
end

function volume(polys::AbstractVector{<:Polyhedron},args...)
  v = 0.0
  for p in polys
    v += volume(p,args...)
  end
  v
end

function restrict(poly::Polyhedron,stl::STL,stl_facets)
  cell_to_nodes = get_cell_vertices(stl)
  c = array_cache(cell_to_nodes)
  nodes = Int32[]
  for f in stl_facets
    for n in getindex!(c,cell_to_nodes,f)
      if n ∉ nodes
        push!(nodes,n)
      end
    end
  end
  sort!(nodes)
  restrict(poly,nodes)
end

function restrict(p::Polyhedron,nodes)
  graph = get_graph(p)[nodes]
  f = i -> Int32( i ∈ nodes ? findfirst(isequal(i),nodes) : OPEN )
  graph = map(i->map(f,i),graph)
  data = restrict(get_data(p),nodes)
  vertices = get_vertex_coordinates(p)[nodes]
  isopen = true
  Polyhedron(vertices,graph,isopen,data)
end

function restrict(data::PolyhedronData,nodes)
  v_to_Π = data.vertex_to_planes[nodes]
  v_to_of = data.vertex_to_original_faces[nodes]
  v_to_v = collect(1:length(nodes))
  v_to_e = fill((UNSET,UNSET),length(nodes))
  Π_to_rΠ = copy(data.plane_to_ref_plane)
  Π_to_id = copy(data.plane_to_ids)
  _Π_to_v_to_d = data.plane_to_vertex_to_distances
  Π_to_v_to_d = [ _Π_to_v_to_d[i][nodes] for i in 1:length(Π_to_id) ]
  args = v_to_Π, v_to_of, v_to_v, v_to_e, Π_to_v_to_d, Π_to_rΠ, Π_to_id
  PolyhedronData(args...)
end

# Printers

function plot(p::Polyhedron,filename=nothing)
  vertices = Int32[]
  for i in 1:num_vertices(p)
    if isactive(p,i)
      push!(vertices,i)
    end
  end
  vertex_to_node = Dict( zip( vertices, 1:length(vertices) ) )
  g = empty( get_graph(p) )
  for (i,vneigs) in enumerate( get_graph(p) )
    isactive(p,i) || continue
    nodes = Int32[]
    for (j,vneig) in enumerate(vneigs)
      if vneig ∈ (UNSET,OPEN)
        push!(vertices,UNSET)
        n = length(vertices)
      else
        n = vertex_to_node[vneig]
      end
      push!(nodes,n)
    end
    push!(g,nodes)
  end
  v_to_pv = p.data.vertex_to_parent_vertex
  names = String[]
  for v in vertices
    if v == UNSET
      name = ""
    else
      name =  v_to_pv[v] == v ? "$v" : "$v($(v_to_pv[v]))"
    end
    push!(names,name)
  end
  kwargs = (fontsize=10,node_size=0,line=(:dot,0.5,1),thickness_scaling=0.5)
  GraphRecipes.graphplot(g,names=names,curves=false;kwargs...)
  isnothing(filename) || Plots.savefig(filename)
  Plots.plot!()
end

function writevtk(p::Polyhedron,filename;kwargs...)
  writevtk(edge_mesh(p),filename;kwargs...)
end

function writevtk(ps::Array{<:Polyhedron},filename)
  for (i,p) in enumerate(ps)
    writevtk(p,filename*"$i")
  end
end

# Getters

num_dims(::T) where T<:Polyhedron = num_dims(T)

num_dims(::Type{<:Polyhedron{D}}) where D = D

point_eltype(::T) where T<:Polyhedron = point_eltype(T)

point_eltype(::Type{<:Polyhedron{D,T}}) where {D,T} = T

num_vertices(a::Polyhedron) = length(a.vertices)

get_vertex_coordinates(a::Polyhedron) = a.vertices

Base.getindex(a::Polyhedron,i::Integer) = a.vertices[i]

@inline get_graph(a::Polyhedron) = a.edge_vertex_graph

get_data(a::Polyhedron) = a.data

Base.isopen(a::Polyhedron) = a.isopen

get_vertex_to_planes(a::PolyhedronData) = a.vertex_to_planes

function isactive(p::Polyhedron,vertex::Integer)
  !isempty( get_graph(p)[vertex] )
end

function get_original_reflex_faces(p::Polyhedron{D},stl::STL;empty=true) where D
  get_original_faces(p,stl,Val{D-2}();empty)
end

function get_original_facets(p::Polyhedron{D},stl::STL;empty=false) where D
  get_original_faces(p,stl,Val{D-1}();empty)
end

function get_original_faces(p::Polyhedron,stl::STL,::Val{d};empty) where d
  faces = collect_original_faces(p,stl,d)
  filter!(f->has_original_face(p,f,Val{d}();empty),faces)
  faces
end

function collect_original_faces(p::Polyhedron,stl::STL,d::Integer)
  v_to_f = get_data(p).vertex_to_original_faces
  facedims = get_facedims(stl)
  faces = Int32[]
  for v in 1:num_vertices(p)
    if isactive(p,v)
      for f ∈ v_to_f[v]
        if facedims[f] == d && f ∉ faces
          push!(faces,f)
        end
      end
    end
  end
  faces
end

function has_original_reflex_face(p::Polyhedron{D},face::Integer;empty=true) where D
  has_original_face(p,face,Val{D-2}();empty)
end

function has_original_facet(p::Polyhedron{D},face::Integer;empty=false) where D
  has_original_face(p,face,Val{D-1}();empty)
end

function has_original_face(p::Polyhedron,face::Integer,::Val{0};empty=false)
  v_to_f = get_data(p).vertex_to_original_faces
  for v in 1:num_vertices(p)
    if isactive(p,v) && face ∈ v_to_f[v]
      return true
    end
  end
  false
end

function has_original_face(p::Polyhedron,face::Integer,::Val{1};empty=true)
  v_to_f = get_data(p).vertex_to_original_faces
  v_to_v = get_data(p).vertex_to_parent_vertex
  for v in 1:num_vertices(p)
    if isactive(p,v) && face ∈ v_to_f[v]
      for vneig in get_graph(p)[v]
        vneig ∉ (OPEN,UNSET) || continue
        if face ∈ v_to_f[vneig] && ( empty || v_to_v[v] ≠ v_to_v[vneig] )
          return true
        end
      end
    end
  end
  false
end

function has_original_face(p::Polyhedron,face::Integer,::Val{2};empty=false)
  d = 2
  v_to_f = get_data(p).vertex_to_original_faces
  v_to_v = get_data(p).vertex_to_parent_vertex
  for v in 1:num_vertices(p)
    if isactive(p,v) && face ∈ v_to_f[v]
      for vneig in get_graph(p)[v]
        vneig ∉ (OPEN,UNSET) || continue
        if face ∈ v_to_f[vneig]
          num_v = 1
          vcurrent = v
          vnext = vneig
          while vnext ≠ v
            if v_to_v[vnext] ∉ (v_to_v[v],v_to_v[vcurrent]) || empty
              num_v += 1
            end
            vcurrent,vnext = vnext,next_vertex(p,vcurrent,vnext)
            vnext ∉ (UNSET,OPEN) || break
            face ∈ v_to_f[vnext] || break
          end
          if vnext == v && num_v ≥ d+1
            return true
          end
        end
      end
    end
  end
  false
end

function has_faces(p::Polyhedron,::Val{0})
  for v in 1:num_vertices(p)
    isactive(p,v) || continue
    return true
  end
  false
end

function has_faces(p::Polyhedron,::Val{1})
  for v in 1:num_vertices(p)
    isactive(p,v) || continue
    for vneig in get_graph(p)[v]
      @assert isactive(p,vneig)
      vneig ∉ (OPEN,UNSET) || continue
      v_to_v[v] ≠ v_to_v[vneig] || continue
      return true
    end
    @unreachable
  end
  false
end

function has_faces(p::Polyhedron,::Val{2})
  d = 2
  v_to_v = get_data(p).vertex_to_parent_vertex
  for v in 1:num_vertices(p)
    isactive(p,v) || continue
    for vneig in get_graph(p)[v]
      vneig ∉ (OPEN,UNSET) || continue
      num_v = 1
      vcurrent = v
      vnext = vneig
      while vnext ≠ v
        if v_to_v[vnext] ∉ (v_to_v[v],v_to_v[vcurrent])
          num_v += 1
        end
        vcurrent,vnext = vnext,next_vertex(p,vcurrent,vnext)
        vnext ∉ (UNSET,OPEN) || break
        if vnext == v && num_v ≥ d+1
          return true
        end
      end
    end
  end
  false
end

has_vertices(p::Polyhedron) = has_faces(p,Val{0}())

has_edges(p::Polyhedron) = has_faces(p,Val{1}())

has_facets(p::Polyhedron{D}) where D = has_faces(p,Val{D-1}())

function get_plane_distances(data::PolyhedronData)
  data.plane_to_vertex_to_distances
end

function get_plane_distances(data::PolyhedronData,Π)
  i = findfirst(isequal(Π),get_plane_ids(data))
  get_plane_distances(data)[i]
end

function get_plane_ids(data::PolyhedronData)
  data.plane_to_ids
end

function has_plane(data::PolyhedronData,Π)
  Π ∈ get_plane_ids(data)
end

## Helpers

function check_graph(p::Polyhedron)
  check_graph(get_graph(p))
end

function check_graph(graph::AbstractVector{<:AbstractVector})
  for v in 1:length(graph)
    !isempty(graph[v]) || continue
    for vneig in graph[v]
      vneig ∉ (OPEN,UNSET) || continue
      v ∈ graph[vneig] || return false
    end
  end
  true
end

function compact!(p::Polyhedron)
  ids = findall(i->!isactive(p,i),1:num_vertices(p))
  old_to_new = fill(UNSET,num_vertices(p))
  new = 0
  for i in 1:num_vertices(p)
    isactive(p,i) || continue
    new += 1
    old = i
    old_to_new[old] = new
  end
  vertices = get_vertex_coordinates(p)
  graph = get_graph(p)
  deleteat!(vertices,ids)
  deleteat!(graph,ids)
  f = i -> i ∈ (OPEN,UNSET) ? i : old_to_new[i]
  map!(i->map!(f,i,i),graph,graph)
  compact!(p.data,ids,old_to_new)
  p
end

function compact!(data::PolyhedronData,ids,old_to_new)
  v_to_Π = data.vertex_to_planes
  v_to_of = data.vertex_to_original_faces
  v_to_pv = data.vertex_to_parent_vertex
  v_to_pe = data.vertex_to_parent_edge
  Π_to_v_to_dist = data.plane_to_vertex_to_distances
  Π_to_ref_Π = data.plane_to_ref_plane
  Π_to_id = data.plane_to_ids

  deleteat!(v_to_Π,ids)
  deleteat!(v_to_of,ids)
  deleteat!(v_to_pv,ids)
  for (i,pv) in enumerate(v_to_pv)
    if old_to_new[pv] == UNSET
      old_to_new[pv] = i
    end
    v_to_pv[i] = old_to_new[pv]
  end
  deleteat!(v_to_pe,ids)
  f = i -> i == UNSET ? i : old_to_new[i]
  for (i,pe) in enumerate(v_to_pe)
    v_to_pe[i] = (f(pe[1]),f(pe[2]))
  end
  for v_to_dist in Π_to_v_to_dist
    deleteat!(v_to_dist,ids)
  end
  data
end

function Base.copy(poly::Polyhedron)
  vertices = get_vertex_coordinates(poly)
  graph = get_graph(poly)
  open = isopen(poly)
  data = copy(get_data(poly))
  Polyhedron(vertices,graph,open,data)
end

function Base.copy(data::PolyhedronData)
  v_to_Π = copy(data.vertex_to_planes)
  v_to_of = copy(data.vertex_to_original_faces)
  v_to_pv = copy(data.vertex_to_parent_vertex)
  v_to_pe = copy(data.vertex_to_parent_edge)
  Π_to_v_to_dist = _copy(data.plane_to_vertex_to_distances)
  Π_to_ref_Π = copy(data.plane_to_ref_plane)
  Π_to_id = copy(data.plane_to_ids)
  args = v_to_Π, v_to_of, v_to_pv, v_to_pe, Π_to_v_to_dist, Π_to_ref_Π, Π_to_id
  PolyhedronData(args...)
end

function compute_distances!(p::Polyhedron,Π,faces;atol=0)
  data = get_data(p)
  v_to_f = get_data(p).vertex_to_original_faces
  Π_to_v_to_d = get_plane_distances(data)
  Π_ids = get_plane_ids(data)
  Π_to_rΠ = data.plane_to_ref_plane
  num_Π = length(Π_ids)
  append!(Π_ids,faces)
  append!(Π_to_rΠ,fill(UNSET,length(faces)))
  append!(Π_to_v_to_d, [ zeros(num_vertices(p)) for _ in faces ] )
  for (i,(Πi,fi)) in enumerate(zip(Π,faces))
    v_to_d = Π_to_v_to_d[i+num_Π]
    for v in 1:num_vertices(p)
      if fi ∈ v_to_f[v]
        dist = 0.0
      else
        dist = signed_distance(p[v],Πi)
        if abs(dist) < atol
          dist = 0.0
        end
      end
      v_to_d[v] = dist
    end
  end
end

function edge_mesh(p::Polyhedron)
  edge_to_vertices = Vector{Int32}[]
  for v in 1:num_vertices(p)
    for vneig in get_graph(p)[v]
      if vneig > v
        push!(edge_to_vertices,[v,vneig])
      end
    end
  end
  X = get_vertex_coordinates(p)
  T = Table(edge_to_vertices)
  UnstructuredGrid(X,T,[SEG2],fill(1,length(T)))
end

function compute_graph(stl::STL{2})
  @notimplementedif get_polytope(stl) ≠ TRI
  D = 2
  v_to_e = get_faces(stl,0,1)
  v_to_f = get_faces(stl,0,D)
  f_to_v = get_faces(stl,D,0)
  ec = array_cache(v_to_e)
  fc = array_cache(v_to_f)
  vc = array_cache(f_to_v)
  graph = [ Int32[] for _ in 1:num_vertices(stl) ]
  for v in 1:num_vertices(stl)
    vfacets = getindex!(fc,v_to_f,v)
    !isempty(vfacets) || continue
    f0 = first(vfacets)
    fnext = f0
    while true
      i = findfirst(isequal(v), getindex!(vc,f_to_v,fnext) )
      inext = i == D+1 ? 1 : i+1
      vnext = getindex!(vc,f_to_v,fnext)[inext]
      push!(graph[v],vnext)
      faces = getindex!(fc,v_to_f,vnext)
      i = findfirst( f -> f ≠ fnext && v ∈ getindex!(vc,f_to_v,f), faces )
      fnext = isnothing(i) ? UNSET : faces[i]
      fnext ≠ UNSET || break
      fnext ≠ f0 || break
    end
    if fnext == UNSET
      n_edges = length( getindex!(ec,v_to_e,v) )
      if length(graph[v]) < n_edges
        v0 = first(graph[v])
        fnext = f0
        while fnext ≠ UNSET
          i = findfirst(isequal(v), getindex!(vc,f_to_v,fnext) )
          inext = i == 1 ? D+1 : i-1
          vnext = getindex!(vc,f_to_v,fnext)[inext]
          pushfirst!(graph[v],vnext)
          faces = getindex!(fc,v_to_f,vnext)
          i = findfirst( f -> f ≠ fnext && v ∈ getindex!(vc,f_to_v,f), faces )
          fnext = isnothing(i) ? UNSET : faces[i]
        end
        @assert length(graph[v]) == n_edges
      end
      push!(graph[v],OPEN)
    end
  end
  check_graph(graph) || error("Unable to build vertex-edge graph")
  graph
end

function set_original_faces!(p::Polyhedron,a...)
  set_original_faces!(get_data(p),a...)
  p
end

function polyhedron_data(num_vertices::Integer)
  v_to_Π = [ Int32[] for _ in 1:num_vertices ]
  v_to_of = [ Int32[] for _ in 1:num_vertices ]
  v_to_v = collect(1:num_vertices)
  v_to_e = fill((UNSET,UNSET),num_vertices)
  Π_to_v_to_d = Vector{Int32}[]
  Π_to_rΠ = Int32[]
  Π_to_id = Int32[]
  args = v_to_Π, v_to_of, v_to_v, v_to_e, Π_to_v_to_d, Π_to_rΠ, Π_to_id
  PolyhedronData(args...)
end

function polyhedron_data(p::Polytope)
  v_to_Π = map( i-> Int32.(i), -get_faces(p,0,num_dims(p)-1) )
  v_to_of = [ Int32[] for _ in 1:num_vertices(p) ]
  v_to_v = Int32.(collect(1:num_vertices(p)))
  v_to_e = fill((Int32(UNSET),Int32(UNSET)),num_vertices(p))
  Π_to_v_to_d = Vector{Int32}[]
  Π_to_rΠ = Int32[]
  Π_to_id = Int32[]
  args = v_to_Π, v_to_of, v_to_v, v_to_e, Π_to_v_to_d, Π_to_rΠ, Π_to_id
  PolyhedronData(args...)
end

@inline function split_postprocess!(graph,data,new_vertices,input_poly,distances,(≶))
  complete_graph!(graph,num_vertices(input_poly))
  disconnect_graph!(graph,num_vertices(input_poly),distances,(≶))
  add_open_vertices!(graph,input_poly)
  vertices = [input_poly.vertices;new_vertices]
  update_data!(data,graph,num_vertices(input_poly))
  poly = Polyhedron(vertices,graph,isopen(input_poly),data)
  compact!(poly)
end

function complete_graph!(edge_graph,num_vertices::Integer)
  for v in num_vertices+1:length(edge_graph)
    vnext = next_vertex(edge_graph,num_vertices,v)
    vnext ∉ (UNSET,OPEN) || continue
    edge_graph[v][end] = vnext
    edge_graph[vnext][2] = v
  end
end

function add_open_vertices!(graph,poly::Polyhedron)
  if isopen(poly)
    for v in num_vertices(poly)+1:length(graph)
      if UNSET ∉ graph[v]
        push!(graph[v],graph[v][end])
        graph[v][end-1] = OPEN
      end
    end
  end
end

function disconnect_graph!(edge_graph,num_vertices,distances,(≶)::Function)
  for i in 1:num_vertices
    if distances[i] ≶ 0
      edge_graph[i] = empty!( edge_graph[i] )
    end
  end
end

function update_data!(data,graph,num_vertices)
  v_to_pv = data.vertex_to_parent_vertex
  v_to_pe = data.vertex_to_parent_edge
  Π_to_v_to_d = get_plane_distances(data)
  for v in num_vertices+1:length(graph)
    v_to_pv[v] == v || continue
    for vnext in num_vertices+1:length(graph)
      v_to_pv[vnext] == vnext || continue
      if v_to_pe[v] == v_to_pe[vnext] && v_to_pe[v] ≠ (UNSET,UNSET)
        v_to_pv[vnext] = v
        for v_to_d in Π_to_v_to_d
          v_to_d[v] = v_to_d[v_to_pv[v]]
        end
      end
    end
  end
end

function next_vertex(edge_graph::Vector,num_vertices::Integer,vstart::Integer)
  vcurrent = vstart
  vnext = first( edge_graph[vcurrent] )
  while vnext ≤ num_vertices && vnext ∉ (UNSET,OPEN)
    i = findfirst( isequal(vcurrent), edge_graph[vnext] )
    vcurrent = vnext
    inext = ( i % length( edge_graph[vcurrent] ) ) + 1
    vnext = edge_graph[vcurrent][ inext ]
  end
  vnext
end

function compute_intersection(p1::Point,d1::Real,p2::Point,d2::Real)
  ( p1*(abs(d2))+p2*(abs(d1)) ) / ( abs(d1) + abs(d2) )
end

add_vertex!(data::Nothing,a...) = data

function _intersect(a::AbstractVector{T},b::AbstractVector{T}) where T
  n = count(in(a),b)
  r = Vector{T}(undef,n)
  n = 0
  for i in a
    if i ∈ b
      n += 1
      r[n] = i
    end
  end
  r
end

function _copy(a::AbstractVector{<:AbstractVector})
  [ copy(i) for i in a ]
end

function add_vertex!(data::PolyhedronData,v1::Integer,v2::Integer,Πid::Integer)
  v_to_Π = data.vertex_to_planes
  v_to_f = data.vertex_to_original_faces
  v_to_pv = data.vertex_to_parent_vertex
  v_to_pe = data.vertex_to_parent_edge
  Πs = _intersect( v_to_Π[v1], v_to_Π[v2] )
  fs = _intersect( v_to_f[v1], v_to_f[v2] )
  pe = (v_to_pv[v1],v_to_pv[v2])
  push!( Πs, Πid )
  push!( v_to_Π, Πs )
  push!( v_to_f, fs )
  push!( v_to_pe, pe)
  Π_to_v_to_d = get_plane_distances(data)
  Π_to_id = get_plane_ids(data)
  i = findfirst(isequal(Πid),Π_to_id)
  d1 = Π_to_v_to_d[i][v1]
  d2 = Π_to_v_to_d[i][v2]
  pv = iszero(d1) ? v_to_pv[v1] : length(v_to_pv)+1
  pv = iszero(d2) ? v_to_pv[v2] : pv
  push!(v_to_pv,pv)
  @assert d1 ≠ d2
  ref_i = data.plane_to_ref_plane[i]
  for i in 1:length(Π_to_v_to_d)
    v_to_d = Π_to_v_to_d[i]
    if pv ≠ length(v_to_pv)
      dist = v_to_d[pv]
    elseif ref_i == UNSET || data.plane_to_ref_plane[i] ∈ (UNSET,i)
      if Π_to_id[i] == Πid
        dist = 0.0
      else
        dist = (d1*v_to_d[v2]-d2*v_to_d[v1])/(d1-d2)
      end
    else
      j = data.plane_to_ref_plane[i]
      @assert abs(j) < i
      dist = sign(j) * Π_to_v_to_d[abs(j)][end]
    end
    push!(v_to_d,dist)
  end
  data
end

function set_original_faces!(data::PolyhedronData,stl::STL)
  Dc = num_dims(stl)
  v_to_f = data.vertex_to_original_faces
  for d in 0:Dc
    offset = get_offset(stl,d)
    _v_to_f = get_faces(stl,0,d)
    c = array_cache(_v_to_f)
    for v in 1:length(v_to_f)
      append!( v_to_f[v], map( i -> i.+offset, getindex!(c,_v_to_f,v) ) )
    end
  end
  data
end

function next_vertex(p::Polyhedron,vprevious::Integer,vcurrent::Integer)
  i = findfirst( isequal(vprevious), get_graph(p)[vcurrent] )
  i = ( i % length( get_graph(p)[vcurrent] ) ) + 1
  get_graph(p)[vcurrent][ i ]
end

## Kd Tree stuff

function refine_by_vertices(Γ::Polyhedron,K::Polyhedron,atol=0)
  refine_by_vertices(Γ,K,get_vertex_coordinates(Γ),atol)
end

function refine_by_vertices(Γ,K,vertices,atol)
  i = findfirst( v -> has_intersection(K,v), vertices )
  if isnothing(i)
    R = [Γ],[K]
  else
    v = vertices[i]
    d = get_direction(v,K)
    Π,Π⁻,Π⁺ = get_planes(v,d,atol)
    Πid,Πid⁻,Πid⁺ = get_new_plane_ids(Γ)
    compute_distances!(Γ,[Π⁻,Π⁺],[Πid⁻,Πid⁺])
    compute_distances!(K,[Π],[Πid])
    _,γ⁻ = split(Γ,Πid⁺)
    _,γ⁺ = split(Γ,Πid⁻)
    k⁻,k⁺ = split(K,Πid)
    round_distances!(k⁻,atol/10)
    round_distances!(k⁺,atol/10)
    Γn = [γ⁻,γ⁺]
    Kn = [k⁻,k⁺]
    if length(vertices) == i
      R = Γn,Kn
    else
      Γr,Kr = empty(Γn),empty(Kn)
      vertices = view(vertices,i+1:length(vertices))
      for (γ,k) in zip(Γn,Kn)
        Γi,Ki = refine_by_vertices(γ,k,vertices,atol)
        append!(Γr,Γi)
        append!(Kr,Ki)
      end
      R = Γr,Kr
    end
  end
  R
end

function get_bounding_box(cell::Polyhedron)
  i = findfirst(i->isactive(cell,i),1:num_vertices(cell))
  pmin,pmax = cell[i],cell[i]
  for i in 1:num_vertices(cell)
    if isactive(cell,i)
      pmin = Point( min.(Tuple(cell[i]),Tuple(pmin)) )
      pmax = Point( max.(Tuple(cell[i]),Tuple(pmax)) )
    end
  end
  pmin,pmax
end

function has_intersection(cell::Polyhedron,v::Point)
  pmin,pmax = get_bounding_box(cell)
  all( Tuple(pmin) .< Tuple(v) ) || return false
  all( Tuple(pmax) .> Tuple(v) ) || return false
  true
end

function get_direction(v,cell)
  pmin,pmax = get_bounding_box(cell)
  d0 = Tuple(v) .- Tuple(pmin)
  d1 = Tuple(pmax) .- Tuple(v)
  d = min.(d0,d1)
  d_max = maximum(d)
  findfirst(isequal(d_max),d)
end

function get_planes(v,d,atol)
  δ = VectorValue(Base.setindex(Tuple(zero(v)),atol,d))
  Π = CartesianPlane(v,d,+1)
  Π⁻ = CartesianPlane(v-δ,d,+1)
  Π⁺ = CartesianPlane(v+δ,d,-1)
  Π,Π⁻,Π⁺
end

function get_new_plane_ids(poly)
  id = minimum( get_plane_ids(poly.data) )
  id-1,id-2,id-3
end

function round_distances!(poly,atol)
  Π_to_v_to_dist = poly.data.plane_to_vertex_to_distances
  for i in 1:length(Π_to_v_to_dist)
    v_to_dist = Π_to_v_to_dist[i]
    for v in 1:length(v_to_dist)
      if abs(v_to_dist[v]) < atol
        v_to_dist[v] = 0.0
      end
    end
  end
end

function delete_inactive_planes!(Γ::Polyhedron,K::Polyhedron,stl::STL)
  @assert isopen(Γ)
  planes = get_active_planes(Γ,K,stl)
  restrict_planes!(Γ,planes)
  restrict_planes!(K,planes)
end

function get_active_planes(Γ::Polyhedron,K::Polyhedron,stl::STL)
  used_planes_Γ = _get_used_planes(K)
  used_planes_K = _get_used_planes(Γ)
  used_planes = lazy_append( used_planes_Γ, used_planes_K )
  face_planes = _get_face_planes(Γ,stl)
  planes = lazy_append( used_planes, face_planes )
  _append_ref_planes( Γ, planes )
end

function restrict_planes!(poly::Polyhedron,planes)
  restrict_planes!(get_data(poly),planes)
end

function restrict_planes!(data::PolyhedronData,planes)
  Π_to_id = data.plane_to_ids
  Π_to_v_to_dist = data.plane_to_vertex_to_distances
  Π_to_ref_Π = data.plane_to_ref_plane
  is_inactive = trues(length(Π_to_id))
  is_touch = falses(length(Π_to_id))
  for Π in planes
    i = findfirst(isequal(Π),Π_to_id)
    !isnothing(i) || continue
    is_inactive[i] = false
  end
  id = 0
  for i in 1:length(Π_to_id)
    if !is_inactive[i]
      id += 1
      if !is_touch[i]
        ref_plane = abs( Π_to_ref_Π[i] )
        if ref_plane != UNSET
          @assert i == ref_plane
          for j in i:length(Π_to_id)
            if !is_inactive[j] &&
               !is_touch[j] &&
               abs( Π_to_ref_Π[j] ) == ref_plane

              _sign = sign( Π_to_ref_Π[j] )
              Π_to_ref_Π[j] = _sign * id
              is_touch[j] = true
            end
          end
        end
      end
    end
  end
  deleteat!( Π_to_v_to_dist, is_inactive)
  deleteat!( Π_to_ref_Π, is_inactive)
  deleteat!( Π_to_id, is_inactive)
end

function _get_used_planes(poly::Polyhedron)
  v_to_Π = get_data(poly).vertex_to_planes
  planes = Int32[]
  for v in 1:length(v_to_Π)
    for Π in v_to_Π[v]
      if Π ∉ planes
        push!(planes,Π)
      end
    end
  end
  planes
end

function _get_face_planes(poly::Polyhedron,stl::STL)
  rfaces = get_original_reflex_faces(poly,stl,empty=true)
  facets = get_original_facets(poly,stl,empty=true)
  lazy_append(rfaces,facets)
end

function _append_ref_planes(poly::Polyhedron,planes)
  ref_planes = Int32[]
  Π_to_ref_Π = get_data(poly).plane_to_ref_plane
  Π_to_id = get_plane_ids( get_data(poly) )
  for Π in planes
    i = findfirst( isequal(Π), Π_to_id )
    i !== nothing || continue
    iref = abs(Π_to_ref_Π[i])
    if iref ≠ UNSET
      Π_ref = Π_to_id[iref]
      if Π_ref ∉ ref_planes && Π_ref ∉ planes
        push!(ref_planes,Π_ref)
      end
    end
  end
  lazy_append(planes,ref_planes)
end


module Lasso
module Util
    # Extract fields from object into function locals
    # See https://github.com/JuliaLang/julia/issues/9755
    macro extractfields(from, fields...)
        esc(Expr(:block, [:($(fields[i]) = $(from).$(fields[i])) for i = 1:length(fields)]...))
    end
    export @extractfields
end

include("FusedLasso.jl")
include("TrendFiltering.jl")

using Reexport, LinearAlgebra, SparseArrays, Random, .Util, MLBase
import Random: Sampler
@reexport using GLM, Distributions, .FusedLassoMod, .TrendFiltering
using GLM: FPVector, LinPred, Link, LmResp, GlmResp, cholpred, updateμ!,
    linkfun, linkinv, LinPredModel
export RegularizationPath, LassoPath, GammaLassoPath, NaiveCoordinateDescent,
       CovarianceCoordinateDescent, fit, fit!, coef, predict,
       minAICc, hasintercept, dof, aicc, distfun, linkfun, cross_validate_path,
       SegSelect, segselect, selectmodel,
       AllSeg, MinAIC, MinAICc, MinBIC, CVSegSelect, MinCVmse, MinCV1se,
       RegularizedModel, LassoModel, GammaLassoModel

## HELPERS FOR SPARSE COEFFICIENTS

struct SparseCoefficients{T} <: AbstractVector{T}
    coef::Vector{T}              # Individual coefficient values
    coef2predictor::Vector{Int}  # Mapping from indices in coef to indices in original X
    predictor2coef::Vector{Int}  # Mapping from indices in original X to indices in coef

    SparseCoefficients{T}(n::Int) where {T} = new(T[], Int[], zeros(Int, n))
end

function LinearAlgebra.mul!(out::Vector, X::Matrix, coef::SparseCoefficients{T}) where T
    fill!(out, zero(eltype(out)))
    @inbounds for icoef = 1:nnz(coef)
        ipred = coef.coef2predictor[icoef]
        c = coef.coef[icoef]
        @simd for i = 1:size(X, 1)
            out[i] += c*X[i, ipred]
        end
    end
    out
end

function LinearAlgebra.mul!(out::Vector, X::SparseMatrixCSC, coef::SparseCoefficients{T}) where T
    @extractfields X colptr rowval nzval
    fill!(out, zero(eltype(out)))
    @inbounds for icoef = 1:nnz(coef)
        ipred = coef.coef2predictor[icoef]
        c = coef.coef[icoef]
        @simd for i = colptr[ipred]:colptr[ipred+1]-1
            out[rowval[i]] += c*nzval[i]
        end
    end
    out
end

function LinearAlgebra.dot(x::Vector{T}, coef::SparseCoefficients{T}) where T
    v = 0.0
    @inbounds @simd for icoef = 1:nnz(coef)
        v += x[coef.coef2predictor[icoef]]*coef.coef[icoef]
    end
    v
end

Base.size(x::SparseCoefficients) = (length(x.predictor2coef),)
SparseArrays.nnz(x::SparseCoefficients) = length(x.coef)
Base.getindex(x::SparseCoefficients{T}, ipred::Int) where {T} =
    x.predictor2coef[ipred] == 0 ? zero(T) : x.coef[x.predictor2coef[ipred]]

function Base.setindex!(A::Matrix{T}, coef::SparseCoefficients, rg::UnitRange{Int}, i::Int) where T
    A[:, i] = zero(T)
    for icoef = 1:nnz(coef)
        A[rg[coef.coef2predictor[icoef]], i] = coef.coef[icoef]
    end
    A
end

function Base.copyto!(x::SparseCoefficients, y::SparseCoefficients)
    length(x) == length(y) || throw(DimensionMismatch())
    n = length(y.coef)
    resize!(x.coef, n)
    resize!(x.coef2predictor, n)
    copyto!(x.coef, y.coef)
    copyto!(x.coef2predictor, y.coef2predictor)
    copyto!(x.predictor2coef, y.predictor2coef)
    x
end

# Add a new coefficient to x, returning its index in x.coef
function addcoef!(x::SparseCoefficients{T}, ipred::Int) where T
    push!(x.coef, zero(T))
    push!(x.coef2predictor, ipred)
    coefindex = nnz(x)
    x.predictor2coef[ipred] = coefindex
end

# Add newcoef to column i of coefs
function addcoefs!(coefs::SparseMatrixCSC, newcoef::SparseCoefficients, i::Int)
    n = nnz(coefs)
    nzval = coefs.nzval
    rowval = coefs.rowval
    resize!(nzval, n+nnz(newcoef))
    resize!(rowval, n+nnz(newcoef))
    @inbounds for ipred = 1:length(newcoef.predictor2coef)
        icoef = newcoef.predictor2coef[ipred]
        if icoef != 0
            cval = newcoef.coef[icoef]
            if cval != 0
                n += 1
                nzval[n] = cval
                rowval[n] = ipred
            end
        end
    end
    resize!(nzval, n)
    resize!(rowval, n)
    coefs.colptr[i+1:end] .= n+1
end

## COEFFICIENT ITERATION IN SEQUENTIAL OR RANDOM ORDER
struct RandomCoefficientIterator
    rng::MersenneTwister
    rg::Sampler
    coeforder::Vector{Int}
end
const RANDOMIZE_DEFAULT = true

function RandomCoefficientIterator()
    rng = MersenneTwister(1337)
    RandomCoefficientIterator(rng, Sampler(rng, 1:2), Int[])
end

const CoefficientIterator = Union{UnitRange{Int},RandomCoefficientIterator}

# Iterate over coefficients in random order
function Base.iterate(x::RandomCoefficientIterator)
    if !isempty(x.coeforder)
        @inbounds for i = length(x.coeforder):-1:2
            j = rand(x.rng, x.rg)
            x.coeforder[i], x.coeforder[j] = x.coeforder[j], x.coeforder[i]
        end
        x.coeforder[1], 2
    else
        nothing
    end
end

Base.iterate(x::RandomCoefficientIterator, i) = (i > length(x.coeforder)) ? nothing : (x.coeforder[i], i += 1)

# Add an additional coefficient and return a new CoefficientIterator
function addcoef(x::RandomCoefficientIterator, icoef::Int)
    push!(x.coeforder, icoef)
    RandomCoefficientIterator(x.rng, Sampler(x.rng, 1:length(x.coeforder)), x.coeforder)
end
addcoef(x::UnitRange{Int}, icoef::Int) = 1:length(x)+1

abstract type RegularizationPath{S<:Union{LinearModel,GeneralizedLinearModel},T} <: RegressionModel end

# don't add an intercept when using a @formula because we use the intercept keyword arg to add an intercept
StatsModels.drop_intercept(::Type{R}) where R<:RegularizationPath = true

## LASSO PATH

mutable struct LassoPath{S<:Union{LinearModel,GeneralizedLinearModel},T} <: RegularizationPath{S,T}
    m::S
    nulldev::T                    # null deviance
    nullb0::T                     # intercept of null model, if one was fit
    λ::Vector{T}                  # shrinkage parameters
    autoλ::Bool                   # whether λ is automatically determined
    Xnorm::Vector{T}              # original squared norms of columns of X before standardization
    pct_dev::Vector{T}            # percent deviance explained by each model
    coefs::SparseMatrixCSC{T,Int} # model coefficients
    b0::Vector{T}                 # model intercepts
    niter::Int                    # number of coordinate descent iterations

    LassoPath{S,T}(m, nulldev::T, nullb0::T, λ::Vector{T}, autoλ::Bool, Xnorm::Vector{T}) where {S,T} =
        new(m, nulldev, nullb0, λ, autoλ, Xnorm)
end

function Base.show(io::IO, path::RegularizationPath)
    prefix = isa(path.m, GeneralizedLinearModel) ? string(typeof(distfun(path)).name.name, " ") : ""
    pathsize = size(path)
    println(io, "$(prefix)$(typeof(path).name.name) ($(pathsize[2])) solutions for $(pathsize[1]) predictors in $(path.niter) iterations):")

    if isdefined(path, :coefs)
        coefs = path.coefs
        ncoefs = zeros(Int, size(coefs, 2))
        for i = 1:size(coefs, 2)-1
            ncoefs[i] = coefs.colptr[i+1] - coefs.colptr[i]
        end
        ncoefs[end] = nnz(coefs) - coefs.colptr[size(coefs, 2)] + 1
        show(io, CoefTable(Union{Vector{Int},Vector{Float64}}[path.λ, path.pct_dev, ncoefs], ["λ", "pct_dev", "ncoefs"], []))
    else
        print(io, "    (not fit)")
    end
end

## MODEL CONSTRUCTION

# Controls early stopping criteria with automatic λ
const MIN_DEV_FRAC_DIFF = 1e-5
const MAX_DEV_FRAC = 0.999

# Compute automatic λ values based on λmax and λminratio
function computeλ(λmax, λminratio, α, nλ)
    λmax /= α
    logλmax = log(λmax)
    exp.(range(logλmax, stop=logλmax + log(λminratio), length=nλ))
end

# Compute automatic λ values based on X'y and λminratio without weights ω
function computeλ(Xy::Vector{T}, λminratio, α, nλ, ω::Nothing) where T
    λmax = zero(T)
    for i = eachindex(Xy)
        x = abs(Xy[i])
        if x > λmax
            λmax = x
        end
    end
    computeλ(λmax, λminratio, α, nλ)
end

# Compute automatic λ values based on X'y and λminratio with specific weights ω
function computeλ(Xy::Vector{T}, λminratio, α, nλ, ω::Vector) where T
    λmax = zero(T)
    for i = eachindex(Xy)
        if ω[i] > 0
            x = abs(Xy[i]) / ω[i]
            if x > λmax
                λmax = x
            end
        end
    end
    computeλ(λmax, λminratio, α, nλ)
end

function weightedresiduals(model::LinearModel)
    r = model.rr
    y = r.y
    mu = r.mu
    if isempty(r.wts)
        y - mu
    else
        wts = r.wts
        resid = similar(y)
        @simd for i = eachindex(resid,y,mu,wts)
            @inbounds resid[i] = (y[i] - mu[i]) * wts[i]
        end
        resid
    end
end

# NOTE: mutates nullmodel
maxλ!(nullmodel, λ::Vector, X::AbstractMatrix{T}, λminratio, α, nλ, ω) where T =
    convert(Vector{T}, λ)

function maxλ!(nullmodel::LinearModel, λ::Nothing,
    X::AbstractMatrix{T}, λminratio, α, nλ, ω) where T

    # muscratch = nullmodel.rr.mu.*nullmodel.rr.wts
    Xy = X'*weightedresiduals(nullmodel)

    computeλ(Xy, λminratio, α, nλ, ω)
end

function maxλ!(nullmodel::GeneralizedLinearModel, λ::Nothing,
    X::AbstractMatrix{T}, λminratio, α, nλ, ω) where T

    Xy = X'*broadcast!(*, nullmodel.rr.wrkresid, nullmodel.rr.wrkresid, nullmodel.rr.wrkwt)
    computeλ(Xy, λminratio, α, nλ, ω)
end

# rescales A so that it sums to base
rescale(A, base) = A * (base / sum(A))

# null model's X includes unpenalized (ω_j==0) variables
nullX(X::AbstractMatrix{T}, intercept::Bool, ω::Nothing) where T =
    ones(T, size(X,1), ifelse(intercept, 1, 0))

function nullX(X::AbstractMatrix{T}, intercept::Bool, ω::Vector) where T
    ixunpenalized = findall(iszero, ω)
    [ones(T, size(X,1), ifelse(intercept, 1, 0)) view(X, :, ixunpenalized)]
end

# version for Linear models
function build_model(X::AbstractMatrix{T}, y::FPVector, d::Normal, l::IdentityLink,
                     lp::LinPred, λminratio::Real, λ::Union{Vector,Nothing},
                     wts::Union{FPVector,Nothing}, offset::Vector, α::Real, nλ::Int,
                     ω::Union{Vector,Nothing}, intercept::Bool, irls_tol::Real, dofit::Bool) where T
    mu = isempty(offset) ? copy(y) : y + offset
    nullmodel = LinearModel(LmResp{typeof(y)}(fill!(similar(y), 0), offset, wts, y),
                            GLM.cholpred(nullX(X, intercept, ω), false))
    fit!(nullmodel)
    nulldev = deviance(nullmodel)
    nullb0 = intercept ? coef(nullmodel)[1] : zero(T)

    λ = maxλ!(nullmodel, λ, X, λminratio, α, nλ, ω)

    # mu is just a placeholder here, setting it to nullmodel.mu causes problems
    model = LinearModel(LmResp{typeof(y)}(mu, offset, wts, y), lp)

    (model, nulldev, nullb0, λ)
end

# version for GLMs
function build_model(X::AbstractMatrix{T}, y::FPVector, d::UnivariateDistribution, l::Link,
                     lp::LinPred, λminratio::Real, λ::Union{Vector,Nothing},
                     wts::Union{FPVector,Nothing}, offset::Vector, α::Real, nλ::Int,
                     ω::Union{Vector, Nothing}, intercept::Bool, irls_tol::Real, dofit::Bool) where T
    # Fit to find null deviance
    # Maybe we should reuse this GlmResp object?
    nullmodel = fit(GeneralizedLinearModel, nullX(X, intercept, ω), y, d, l;
                    wts=wts, offset=offset, rtol=irls_tol, dofit=dofit)
    nulldev = deviance(nullmodel)
    nullb0 = intercept ? coef(nullmodel)[1] : zero(T)

    λ = maxλ!(nullmodel, λ, X, λminratio, α, nλ, ω)

    rr = GlmResp(y, d, l, offset, wts)
    model = GeneralizedLinearModel(rr, lp, false)

    (model, nulldev, nullb0, λ)
end

defaultalgorithm(d::Normal, l::IdentityLink, n::Int, p::Int) = p > 5n ? NaiveCoordinateDescent : CovarianceCoordinateDescent
defaultalgorithm(d::UnivariateDistribution, l::Link, n::Int, p::Int) = NaiveCoordinateDescent

# following glmnet rescale penalty factors to sum to the number of coefficients
initpenaltyfactor(penalty_factor::Nothing,p::Int,::Bool) = nothing
initpenaltyfactor(penalty_factor::Vector,p::Int,standardizeω::Bool) =
    standardizeω ? rescale(penalty_factor, p) : penalty_factor

# Standardize predictors if requested
function standardizeX(X::AbstractMatrix{T}, standardize::Bool) where T
    if standardize
        Xnorm = vec(convert(Matrix{T},std(X; dims=1, corrected=false)))
        if any(x -> x == zero(T), Xnorm)
            @warn("""One of the predicators (columns of X) is a constant, so it can not be standardized.
                  To include a constant predicator set standardize = false and intercept = false""")
        end
        for i = 1:length(Xnorm)
            @inbounds Xnorm[i] = 1/Xnorm[i]
        end
        X = X .* transpose(Xnorm)
    else
        Xnorm = T[]
    end

    X, Xnorm
end

@doc raw"""
    fit(LassoPath, X, y, d=Normal(), l=canonicallink(d); ...)

fits a linear or generalized linear Lasso path given the design
matrix `X` and response `y`:

``\underset{\beta}{\operatorname{argmin}} -\frac{1}{N} \mathcal{L}(y|X,\beta) + \lambda\left[(1-\alpha)\frac{1}{2}\|\beta\|_2^2 + \alpha\|\beta\|_1\right]``

The optional argument `d` specifies the conditional distribution of
response, while `l` specifies the link function. Lasso.jl inherits
supported distributions and link functions from GLM.jl. The default
is to fit an linear Lasso path, i.e., `d=Normal(), l=IdentityLink()`,
or ``\mathcal{L}(y|X,\beta) = -\frac{1}{2}\|y - X\beta\|_2^2 + C``

# Examples
```julia
fit(LassoPath, X, y)    # L1-regularized linear regression
fit(LassoPath, X, y, Binomial(), Logit();
    α=0.5) # Binomial logit regression with an Elastic net combination of
           # 0.5 L1 and 0.5 L2 regularization penalties
```
# Arguments
- `wts=ones(length(y))`: Weights for each observation
- `offset=zeros(length(y))`: Offset of each observation
- `λ`: can be used to specify a specific set of λ values at which models are fit.
    If λ is unspecified, Lasso.jl selects nλ logarithmically spaced λ values from
    `λmax`, the smallest λ value yielding a null model, to
    `λminratio * λmax`.
- `nλ=100` number of λ values to use
- `λminratio=1e-4` if more observations than predictors otherwise 0.001.
- `stopearly=true`: When `true`, if the proportion of deviance explained
    exceeds 0.999 or the difference between the deviance explained by successive λ
    values falls below `1e-5`, the path stops early.
- `standardize=true`: Whether to standardize predictors to unit standard deviation
    before fitting.
- `intercept=true`: Whether to fit an (unpenalized) model intercept.
- `algorithm`: Algorithm to use.
    `NaiveCoordinateDescent` iteratively computes the dot product of the
    predictors with the  residuals, as opposed to the
    `CovarianceCoordinateDescent` algorithm, which uses a precomputed Gram matrix.
    `NaiveCoordinateDescent` is typically faster when there are many
    predictors that will not enter the model or when fitting
    generalized linear models.
    By default uses `NaiveCoordinateDescent` if more than 5x as many predictors
    as observations or model is a GLM. `CovarianceCoordinateDescent` otherwise.
- `randomize=true`: Whether to randomize the order in which coefficients are
    updated by coordinate descent. This can drastically speed
    convergence if coefficients are highly correlated.
- `maxncoef=min(size(X, 2), 2*size(X, 1))`: maximum number of coefficients
    allowed in the model. If exceeded, an error will be thrown.
- `dofit=true`: Whether to fit the model upon construction. If `false`, the
    model can be fit later by calling `fit!(model)`.
- `cd_tol=1e-7`: The tolerance for coordinate descent iterations iterations in
    the inner loop.
- `irls_tol=1e-7`: The tolerance for outer iteratively reweighted least squares
    iterations. This is ignored unless the model is a generalized linear model.
- `criterion=:coef` Convergence criterion. Controls how ``cd_tol`` and ``irls_tol``
    are to be interpreted. Possible values are:
    - ``:coef``: The model is considered to have converged if the
    the maximum absolute squared difference in coefficients
    between successive iterations drops below the specified
    tolerance. This is the criterion used by glmnet.
    - ``:obj``: The model is considered to have converged if the
    the relative change in the Lasso/Elastic Net objective
    between successive iterations drops below the specified
    tolerance. This is the criterion used by GLM.jl.
- `minStepFac=0.001`: The minimum step fraction for backtracking line search.
- `penalty_factor=ones(size(X, 2))`: Separate penalty factor ``\omega_j``
    for each coefficient ``j``, i.e. instead of ``\lambda`` penalties become
    ``\lambda\omega_j``.
    Note the penalty factors are internally rescaled to sum to
    the number of variables (`glmnet.R` convention).
- `standardizeω=true`: Whether to scale penalty factors to sum to the number of
    variables (glmnet.R convention).
"""
function StatsBase.fit(::Type{LassoPath},
                       X::AbstractMatrix{T}, y::V, d::UnivariateDistribution=Normal(),
                       l::Link=canonicallink(d);
                       wts::Union{FPVector,Nothing}=ones(T, length(y)),
                       offset::V=similar(y, 0),
                       α::Number=one(eltype(y)), nλ::Int=100,
                       λminratio::Number=ifelse(size(X, 1) < size(X, 2), 0.01, 1e-4),
                       λ::Union{Vector,Nothing}=nothing, standardize::Bool=true,
                       intercept::Bool=true,
                       algorithm::Type=defaultalgorithm(d, l, size(X, 1), size(X, 2)),
                       dofit::Bool=true,
                       irls_tol::Real=1e-7, randomize::Bool=RANDOMIZE_DEFAULT,
                       maxncoef::Int=min(size(X, 2), 2*size(X, 1)),
                       penalty_factor::Union{Vector,Nothing}=nothing,
                       standardizeω::Bool=true,
                       fitargs...) where {T<:AbstractFloat,V<:FPVector}
    size(X, 1) == size(y, 1) || DimensionMismatch("number of rows in X and y must match")
    n = length(y)
    length(wts) == n || error("length(wts) = $(length(wts)) should be 0 or $n")

    X, Xnorm = standardizeX(X, standardize)

    # Lasso initialization
    α = convert(T, α)
    λminratio = convert(T, λminratio)
    coefitr = randomize ? RandomCoefficientIterator() : (1:0)

    # penalty_factor (ω) defaults to a vector of ones
    ω = initpenaltyfactor(penalty_factor, size(X, 2), standardizeω)

    cd = algorithm{T,intercept,typeof(X),typeof(coefitr),typeof(ω)}(X, α, maxncoef, 1e-7, coefitr, ω)

    # GLM response initialization
    autoλ = λ == nothing
    model, nulldev, nullb0, λ = build_model(X, y, d, l, cd, λminratio, λ, wts .* T(1/sum(wts)),
                                            Vector{T}(offset), α, nλ, ω, intercept, irls_tol, dofit)

    # Fit path
    path = LassoPath{typeof(model),T}(model, nulldev, nullb0, λ, autoλ, Xnorm)
    if dofit
        fit!(path; irls_tol=irls_tol, fitargs...)
    end
    path
end

StatsBase.nobs(path::RegularizationPath) = length(path.m.rr.y)

dispersion_parameter(path::RegularizationPath) = GLM.dispersion_parameter(distfun(path))

function StatsBase.loglikelihood(path::RegularizationPath{S,T}) where {S<:LinearModel,T}
    n = nobs(path)
    -0.5.*n.*log.(deviance(path))
end

function StatsBase.loglikelihood(path::RegularizationPath{S,T}) where {S,T}
    n = nobs(path)
    -0.5.*n.*deviance(path)
end

"""
    dof(path::RegularizationPath)

Approximates the degrees-of-freedom in each segment of the path as the number of non zero coefficients
plus a dispersion parameter when appropriate.
Note that for GammaLassoPath this may be a crude approximation, as gamlr does this differently.
"""
function StatsBase.dof(path::RegularizationPath)
    nλ = length(path.λ)
    βs = coef(path)
    df = zeros(Int,nλ)
    for s=1:nλ
        df[s] = sum(βs[:,s].!=0)
    end

    if dispersion_parameter(path)
        # add one for dispersion_parameter
        df.+=1
    end

    df
end

function infocrit(d::T,l::F,n,k) where {T,F}
    if d + one(T) > n
        floatmax(F)
    else
        -2l + (k*d + k*d*(d+1)/(n-d-1))
    end
end

function StatsBase.aicc(path::RegularizationPath;k=2)
    dfs = dof(path)
    ls = loglikelihood(path)
    n = nobs(path)
    broadcast((d,l)->infocrit(d,l,n,k), dfs, ls)
end

minAIC(path::RegularizationPath)=argmin(aic(path))
minAICc(path::RegularizationPath;k=2)=argmin(aicc(path;k=k))
minBIC(path::RegularizationPath)=argmin(bic(path))

hasintercept(path::RegularizationPath) = hasintercept(path.m.pp)

"""
size(path) returns (p,nλ) where p is the number of coefficients (including
any intercept) and nλ is the number of path segments.
If model was only initialized but not fit, returns (p,1).
"""
function Base.size(path::RegularizationPath)
  if isdefined(path,:coefs)
    p,nλ = size(path.coefs)
  else
    X = path.m.pp.X
    p = size(X,2)
    nλ = 1
  end

  if hasintercept(path)
    p += 1
  end

  p,nλ
end

"link of underlying GLM"
GLM.linkfun(path::RegularizationPath{M}) where {M<:LinearModel} = IdentityLink()
GLM.linkfun(path::RegularizationPath{GeneralizedLinearModel{GlmResp{V,D,L},L2}}) where {V<:FPVector,D<:UnivariateDistribution,L<:Link,L2<:GLM.LinPred} = L()

"""
predict(path::RegularizationPath, newX::AbstractMatrix; kwargs...)

Predicted values for a selected segment of a regularization path.

# Examples
```julia
predict(path, newX; select=MinBIC())     # predict using BIC minimizing segment
"""
function StatsBase.predict(path::RegularizationPath, newX::AbstractMatrix{T}; offset::FPVector=T[], select=AllSeg()) where {T<:AbstractFloat}
    # add an interecept to newX if the model has one
    if hasintercept(path)
        newX = [ones(eltype(newX),size(newX,1),1) newX]
    end

    # calculate etas for each obs x segment
    eta = newX * coef(path, select)

    # get model
    mm = path.m

    # adjust for any offset
    if length(mm.rr.offset) > 0
        length(offset) == size(newX, 1) ||
            throw(ArgumentError("fit with offset, so `offset` kw arg must be an offset of length `size(newX, 1)`"))
        broadcast!(+, eta, eta, offset)
    else
        length(offset) > 0 && throw(ArgumentError("fit without offset, so value of `offset` kw arg does not make sense"))
    end

    if typeof(mm) <: LinearModel
        eta
    else
        # invert all etas to mus
        μ(η) = linkinv(linkfun(path), η)
        map(μ, eta)
    end
end

"predicted values for data used to estimate `path`"
function StatsBase.predict(path::RegularizationPath; select=AllSeg())
    X = path.m.pp.X
    offset = path.m.rr.offset

    # destandardize X if needed
    if !isempty(path.Xnorm)
        X = X ./ transpose(path.Xnorm)
    end

    predict(path, X; offset=offset, select=select)
end

"distribution of underlying GLM"
distfun(path::RegularizationPath{M}) where {M<:LinearModel} = Normal()
distfun(path::RegularizationPath) = path.m.rr.d

"deviance at each segment of the path for the fitted model and data"
StatsBase.deviance(path::RegularizationPath) = (1 .- path.pct_dev) .* path.nulldev # * nobs(path)

"""
deviance at each segement of the path for (potentially new) data X and y
select=AllSeg() or MinAICc() like in coef()
"""
function StatsBase.deviance(path::RegularizationPath, X::AbstractMatrix{T}, y::V;
                    offset::FPVector=T[], select=AllSeg(),
                    wts::FPVector=ones(T, length(y))) where {T<:AbstractFloat,V<:FPVector}
    μ = predict(path, X; offset=offset, select=select)
    deviance(path, y, μ, wts)
end

"""
deviance at each segment of the path for (potentially new) y and predicted values μ
"""
function StatsBase.deviance(path::RegularizationPath, y::V, μ::AbstractArray{T},
                wts::FPVector=ones(T, length(y))) where {T<:AbstractFloat,V<:FPVector}
    # get model specs from path
    d = distfun(path)

    # rescale weights
    wts = wts .* convert(T, 1/sum(wts))

    # closure for deviance of a single observation
    dev(ys,μs,ws) = ws * devresid(d, ys, μs)

    # deviances of all obs x segment
    devresidv = broadcast(dev,y,μ,wts)

    # deviance is just their sum
    if size(μ,2) > 1
        vec(sum(devresidv,dims=1))
    else
        sum(devresidv)
    end
end

include("coordinate_descent.jl")
include("gammalasso.jl")
include("segselect.jl")
include("cross_validation.jl")
include("deprecated.jl")

end


"""
spin17.jl - derivative robustness for the δa problem
"""

WDIR = get(ENV, "ROBUST_QOC_PATH", "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin17"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
const ACONTROL_SIZE = CONTROL_COUNT
# state indices
const STATE1_IDX = SVector{HDIM_ISO}(1:HDIM_ISO)
const STATE2_IDX = SVector{HDIM_ISO}(STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO)
const INTCONTROLS_IDX = SVector{CONTROL_COUNT}(STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT)
const CONTROLS_IDX = SVector{CONTROL_COUNT}(INTCONTROLS_IDX[end]
                                            + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT)
const DCONTROLS_IDX = SVector{CONTROL_COUNT}(CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT)
const DSTATE1_IDX = SVector{HDIM_ISO}(DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO)
const D2STATE1_IDX = SVector{HDIM_ISO}(DSTATE1_IDX[end] + 1:DSTATE1_IDX[end] + HDIM_ISO)
# const STATE3_IDX = DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO
# const STATE4_IDX = STATE3_IDX[end] + 1:STATE3_IDX[end] + HDIM_ISO
# const DSTATE1_IDX = STATE4_IDX[end] + 1:STATE4_IDX[end] + HDIM_ISO
# const DSTATE2_IDX = DSTATE1_IDX[end] + 1:DSTATE1_IDX[end] + HDIM_ISO
# const DSTATE3_IDX = DSTATE2_IDX[end] + 1:DSTATE2_IDX[end] + HDIM_ISO
# const DSTATE4_IDX = DSTATE3_IDX[end] + 1:DSTATE3_IDX[end] + HDIM_ISO
# const D2STATE1_IDX = DSTATE4_IDX[end] + 1:DSTATE4_IDX[end] + HDIM_ISO
# const D2STATE2_IDX = D2STATE1_IDX[end] + 1:D2STATE1_IDX[end] + HDIM_ISO
# const D2STATE3_IDX = D2STATE2_IDX[end] + 1:D2STATE2_IDX[end] + HDIM_ISO
# const D2STATE4_IDX = D2STATE3_IDX[end] + 1:D2STATE3_IDX[end] + HDIM_ISO
# control indices
const D2CONTROLS_IDX = SVector{CONTROL_COUNT}(1:CONTROL_COUNT)

# model
struct Model{DO} <: AbstractModel
end
@inline RD.state_dim(::Model{DO}) where {DO} = ASTATE_SIZE_BASE + DO * HDIM_ISO
@inline RD.control_dim(::Model) = ACONTROL_SIZE
# @inline RD.state_dim(::Model{DO}) where {DO} = (
#     ASTATE_SIZE_BASE + (SAMPLE_STATE_COUNT - STATE_COUNT) * HDIM_ISO + DO * SAMPLE_STATE_COUNT * HDIM_ISO
# )

# dynamics
const NEGI2_H1_ISO = 2 * NEGI_H1_ISO
function RD.discrete_dynamics(::Type{RK3}, model::Model{DO}, astate::SVector,
                              acontrol::SVector, time::Real, dt::Real) where {DO}
    negi_h = (
        FQ_NEGI_H0_ISO
        + astate[CONTROLS_IDX[1]] * NEGI_H1_ISO
    )
    h_prop = exp(negi_h * dt)
    state1_ = astate[STATE1_IDX]
    state2_ = astate[STATE2_IDX]
    state1 =  h_prop * state1_
    state2 = h_prop * state2_
    intcontrols = astate[INTCONTROLS_IDX] + astate[CONTROLS_IDX] * dt
    controls = astate[CONTROLS_IDX] + astate[DCONTROLS_IDX] * dt
    dcontrols = astate[DCONTROLS_IDX] + acontrol[D2CONTROLS_IDX] * dt

    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
    ]

    if DO >= 1
        dstate1_ = astate[DSTATE1_IDX]
        dstate1 = h_prop * (dstate1_ + dt * NEGI_H1_ISO * state1_)
        astate_ = [astate_; dstate1]
        # state3_ = astate[STATE3_IDX]
        # state3 = h_prop * state3_
        # state4_ = astate[STATE4_IDX]
        # state4 = h_prop * state4_
        # dstate1_ = astate[DSTATE1_IDX]
        # dstate1 = h_prop * (dstate1_ + dt * NEGI_H1_ISO * state1_)
        # dstate2_ = astate[DSTATE2_IDX]
        # dstate2 = h_prop * (dstate2_ + dt * NEGI_H1_ISO * state2_)
        # dstate3_ = astate[DSTATE3_IDX]
        # dstate3 = h_prop * (dstate3_ + dt * NEGI_H1_ISO * state3_)
        # dstate4_ = astate[DSTATE4_IDX]
        # dstate4 = h_prop * (dstate4_ + dt * NEGI_H1_ISO * state4_)
        # append!(astate_, [state3; state4; dstate1; dstate2; dstate3; dstate4])
    end
    if DO >= 2
        d2state1_ = astate[D2STATE1_IDX]
        d2state1 = h_prop * (d2state1_ + dt * NEGI2_H1_ISO * dstate1_)
        astate_ = [astate_; d2state1]
        # d2state1_ = astate[D2STATE1_IDX]
        # d2state1 = h_prop * (d2state1_ + dt * NEGI2_H1_ISO * dstate1_)
        # d2state2_ = astate[D2STATE2_IDX]
        # d2state2 = h_prop * (d2state2_ + dt * NEGI2_H1_ISO * dstate2_)
        # d2state3_ = astate[D2STATE3_IDX]
        # d2state3 = h_prop * (d2state3_ + dt * NEGI2_H1_ISO * dstate3_)
        # d2state4_ = astate[D2STATE4_IDX]
        # d2state4 = h_prop * (d2state4_ + dt * NEGI2_H1_ISO * dstate4_)
        # append!(astate_, [d2state1; d2state2; d2state3; d2state4])
    end

    return astate_
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=56.8, solver_type=altro,
                  sqrtbp=false, derivative_order=0, integrator_type=rk3, qs=ones(7),
                  smoke_test=false, dt_inv=Int64(1e1), constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true, max_iterations=Int64(2e5))
    # model configuration
    model = Model{derivative_order}()
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.

    # initial state, target state
    x0 = zeros(n)
    xf = zeros(n)
    gate = GT_GATE_ISO[gate_type]
    x0[STATE1_IDX] = IS1_ISO_
    x0[STATE2_IDX] = IS2_ISO_
    # x0[STATE3_IDX] = IS3_ISO_
    # x0[STATE4_IDX] = IS4_ISO_
    xf[STATE1_IDX] = gate * IS1_ISO_
    xf[STATE2_IDX] = gate * IS2_ISO_
    # xf[STATE3_IDX] = gate * IS3_ISO_
    # xf[STATE4_IDX] = gate * IS4_ISO_
    x0 = SVector{n}(x0)
    xf = SVector{n}(xf)

    # control amplitude constraint
    x_max = fill(Inf, n)
    x_max[CONTROLS_IDX] .= MAX_CONTROL_NORM_0
    x_max = SVector{n}(x_max)
    x_min = fill(-Inf, n)
    x_min[CONTROLS_IDX] .= -MAX_CONTROL_NORM_0
    x_min = SVector{n}(x_min)

    # control amplitude constraint at boundary
    x_max_boundary = fill(Inf, n)
    x_max_boundary[CONTROLS_IDX] .= 0
    x_max_boundary = SVector{n}(x_max_boundary)
    x_min_boundary = fill(-Inf, n)
    x_min_boundary[CONTROLS_IDX] .= 0
    x_min_boundary = SVector{n}(x_min_boundary)

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}(
        fill(1e-4, CONTROL_COUNT)
    ) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], CONTROL_COUNT); # ∫a
        fill(qs[3], CONTROL_COUNT); # a
        fill(qs[4], CONTROL_COUNT); # ∂a
        fill(qs[5], eval(:($derivative_order >=1 ? HDIM_ISO : 0))); #∂ψ1
        fill(qs[6], eval(:($derivative_order >=2 ? HDIM_ISO : 0))); #∂2ψ1
        # fill(0, eval(:($derivative_order >= 1 ?
        # ($SAMPLE_COUNT - $STATE_COUNT) * $HDIM_ISO : 0))) # ψ3, ψ4
        # fill(qs[5], eval(:($derivative_order >= 1 ? $SAMPLE_COUNT * $HDIM_ISO : 0))); # ∂ψ
        # fill(qs[6], eval(:($derivative_order >= 2 ? $SAMPLE_COUNT * $HDIM_ISO : 0))); # ∂2ψ
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[7], CONTROL_COUNT);
    ]))
    objective = LQRObjective(Q, R, Qf, xf, N)

    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX]) # ; INTCONTROLS_IDX])
    # must obey unit norm.
    norm_idxs = [STATE1_IDX, STATE2_IDX]
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idx) for idx in norm_idxs]

    constraints = ConstraintList(n, m, N)
    # add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N);
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints,
                                            x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
                 projected_newton_tolerance=al_tol, n_steps=n_steps,
                 penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
                 projected_newton=projected_newton, iterations_inner=iterations_inner,
                 iterations_outer=iterations_outer, iterations=max_iterations)
    Altro.solve!(solver)

    # post-process
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TrajectoryOptimization.max_violation(solver)
    cmax_info = TrajectoryOptimization.findmax_violation(TO.get_constraints(solver))
    iterations_ = Altro.iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "Q" => Q_arr,
        "Qf" => Qf_arr,
        "R" => R_arr,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "derivative_order" => derivative_order,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "integrator_type" => Integer(integrator_type),
        "iterations" => iterations_,
        "max_iterations" => max_iterations,
        "pn_steps" => pn_steps,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    return result
end


# This file is a part of Julia. License is MIT: https://julialang.org/licenOAse

import LinearAlgebra: BlasInt
using Logging

# A convenient shortcut to show unexpected behavior from libarpack
const ERR_UNEXPECTED_BEHAVIOR = -999

struct XYAUPD_Exception <: Exception
    info::BlasInt
end

const AUPD_ERRORS = [
    (3, "No shifts could be applied during a cycle of the Implicitly restarted Arnoldi iteration. One possibility is to increase the size of NCV relative to NEV. "),
    (2, "No longer an informational error. Deprecated starting with release 2 of ARPACK."),
    (1, """Maximum number of iterations taken. All possible eigenvalues of OP has been found.
          IPARAM(5) returns the number of wanted converged Ritz values."""),
    (0, "Normal exit."),
    (-1, "N must be positive."),
    (-2, "NEV must be positive."),
    (-3, "NCV-NEV >= 2 and less than or equal to N."),
    (-4, "The maximum number of Arnoldi update iterations allowed must be greater than zero."),
    (-5, " WHICH must be one of 'LM', 'SM', 'LR', 'SR', 'LI', 'SI'"),
    (-6, "BMAT must be one of 'I' or 'G'."),
    (-7, "Length of private work array WORKL is not sufficient."),
    (-8, "Error return from LAPACK eigenvalue calculation."),
    (-9, "Starting vector is zero."),
    (-10, "IPARAM(7) must be 1,2,3,4."),
    (-11, "IPARAM(7) = 1 and BMAT = 'G' are incompatible."),
    (-12, "IPARAM(1) must be equal to 0 or 1."),
    (-13, "NEV and WHICH = 'BE' are incompatible."),
    (-9999, """Could not build an Arnoldi factorization.
           IPARAM(5) returns the size of the current Arnoldi factorization.
           The user is advised to check that enough workspace and array storage has been allocated.""")
]

function Base.showerror(io::IO, ex::XYAUPD_Exception)
    info = ex.info

    if info == ERR_UNEXPECTED_BEHAVIOR
        @error "XYAUPD_Exception: Undefined error"
    else
      idx = searchsorted(AUPD_ERRORS, info, by=first, rev=true)
      if isempty(idx)
          @error "XYAUPD_Exception: Please check XYAUPD error codes in the ARPACK manual." info
      else
          @error "XYAUPD_Exception: $(last(AUPD_ERRORS[first(idx)]))" info
      end
    end
end

struct XYEUPD_Exception <: Exception
    info::BlasInt
end

const EUPD_ERRORS = [
    (1, """The Schur form computed by LAPACK routine lahqr could not be reordered by LAPACK routine trsen.
       Re-enter subroutine neupd with IPARAM(5)NCV and increase the size of the arrays DR and DI to have dimension at least dimension NCV and allocate at least NCV columns for Z.
       NOTE: Not necessary if Z and V share the same space. Please notify the authors if this error occurs."""),
    (0, "Normal exit."),
    (-1, "N must be positive."),
    (-2, "NEV must be positive."),
    (-3, "NCV-NEV >= 2 and less than or equal to N."),
    (-5, "WHICH must be one of 'LM', 'SM', 'LR', 'SR', 'LI', 'SI'"),
    (-6, "BMAT must be one of 'I' or 'G'."),
    (-7, "Length of private work WORKL array is not sufficient."),
    (-8, """Error return from calculation of a real Schur form.
        "Informational error from LAPACK routine lahqr ."""),
    (-9, """Error return from calculation of eigenvectors.
        "Informational error from LAPACK routine dtrevc."""),
    (-10, "IPARAM(7) must be 1,2,3,4."),
    (-11, "IPARAM(7) = 1 and BMAT = 'G' are incompatible."),
    (-12, "HOWMNY = 'S' not yet implemented"),
    (-13, "HOWMNY must be one of 'A' or 'P' if RVEC = .true."),
    (-14, "DNAUPD  did not find any eigenvalues to sufficient accuracy."),
    (-15, """DNEUPD got a different count of the number of converged Ritz values than NAUPD got.
          This indicates the user probably made an error in passing data from NAUPD to NEUPD or that the data was modified before entering NEUPD""")
]

function Base.showerror(io::IO, ex::XYEUPD_Exception)
    info = ex.info
    if info == ERR_UNEXPECTED_BEHAVIOR
        @error "XYEUPD_Exception: Undefined error"
    else
      idx = searchsorted(EUPD_ERRORS, info, by=first, rev=true)
      if isempty(idx)
          @error "XYEUPD_Exception: Please check XYEUPD error codes in the ARPACK manual." info
      else
          @error "XYEUPD_Exception: $(last(EUPD_ERRORS[first(idx)]))" info
      end
    end
end

## aupd and eupd wrappers

function aupd_wrapper(T, matvecA!::Function, matvecB::Function, solveSI::Function, n::Integer,
                      sym::Bool, cmplx::Bool, bmat,
                      nev::Integer, ncv::Integer, which,
                      tol::Real, maxiter::Integer, mode::Integer, v0::Vector, check::Integer)
    lworkl = cmplx ? ncv * (3*ncv + 5) : (sym ? ncv * (ncv + 8) :  ncv * (3*ncv + 6) )
    TR = cmplx ? T.types[1] : T
    TOL = Ref{TR}(tol)

    v     = Matrix{T}(undef, n, ncv)
    workd = Vector{T}(undef, 3*n)
    workl = Vector{T}(undef, lworkl)
    rwork = cmplx ? Vector{TR}(undef, ncv) : Vector{TR}()

    if isempty(v0)
        resid = Vector{T}(undef, n)
        info  = Ref{BlasInt}(0)
    else
        resid = deepcopy(v0)
        info  = Ref{BlasInt}(1)
    end
    iparam = zeros(BlasInt, 11)
    ipntr  = zeros(BlasInt, (sym && !cmplx) ? 11 : 14)
    ido    = Ref{BlasInt}(0)

    iparam[1] = BlasInt(1)       # ishifts
    iparam[3] = BlasInt(maxiter) # maxiter
    iparam[7] = BlasInt(mode)    # mode

    zernm1 = 0:(n-1)

    while true
        if cmplx
            naupd(ido, bmat, n, which, nev, TOL, resid, ncv, v, n,
                  iparam, ipntr, workd, workl, lworkl, rwork, info)
        elseif sym
            saupd(ido, bmat, n, which, nev, TOL, resid, ncv, v, n,
                  iparam, ipntr, workd, workl, lworkl, info)
        else
            naupd(ido, bmat, n, which, nev, TOL, resid, ncv, v, n,
                  iparam, ipntr, workd, workl, lworkl, info)
        end
        if info[] != 0
            if info[] == 1 && check != 0
                return (resid, v, n, iparam, ipntr, workd, workl, lworkl, rwork, TOL)
            end
            throw(XYAUPD_Exception(info[]))
        end

        x = view(workd, ipntr[1] .+ zernm1)
        y = view(workd, ipntr[2] .+ zernm1)
        if mode == 1  # corresponds to dsdrv1, dndrv1 or zndrv1
            if ido[] == -1 || ido[] == 1
                matvecA!(y, x)
            elseif ido[] == 99
                break
            else
                throw(XYAUPD_Exception(ERR_UNEXPECTED_BEHAVIOR))
            end
        elseif mode == 3 && bmat == "I" # corresponds to dsdrv2, dndrv2 or zndrv2
            if ido[] == -1 || ido[] == 1
                y[:] = solveSI(x)
            elseif ido[] == 99
                break
            else
                throw(XYAUPD_Exception(ERR_UNEXPECTED_BEHAVIOR))
            end
        elseif mode == 2 # corresponds to dsdrv3, dndrv3 or zndrv3
            if ido[] == -1 || ido[] == 1
                matvecA!(y, x)
                if sym
                    x[:] = y    # overwrite as per Remark 5 in dsaupd.f
                end
                y[:] = solveSI(y)
            elseif ido[] == 2
                y[:] = matvecB(x)
            elseif ido[] == 99
                break
            else
                throw(XYAUPD_Exception(ERR_UNEXPECTED_BEHAVIOR))
            end
        elseif mode == 3 && bmat == "G" # corresponds to dsdrv4, dndrv4 or zndrv4
            if ido[] == -1
                y[:] = solveSI(matvecB(x))
            elseif  ido[] == 1
                y[:] = solveSI(view(workd,ipntr[3] .+ zernm1))
            elseif ido[] == 2
                y[:] = matvecB(x)
            elseif ido[] == 99
                break
            else
                throw(XYAUPD_Exception(ERR_UNEXPECTED_BEHAVIOR))
            end
        else
            throw(ArgumentError("ARPACK mode ($mode) not yet supported"))
        end
    end

    return (resid, v, n, iparam, ipntr, workd, workl, lworkl, rwork, TOL)
end

function eupd_wrapper(T, n::Integer, sym::Bool, cmplx::Bool, bmat,
                      nev::Integer, which, ritzvec::Bool,
                      TOL::Ref, resid, ncv::Integer, v, ldv, sigma, iparam, ipntr,
                      workd, workl, lworkl, rwork)
    howmny = "A"
    select = Vector{BlasInt}(undef, ncv)
    info   = Ref{BlasInt}(0)

    dmap = x -> abs.(x)
    if iparam[7] == 3 # shift-and-invert
        dmap = x -> abs.(1 ./ (x .- sigma))
    elseif which == "LR" || which == "LA" || which == "BE"
        dmap = real
    elseif which == "SR" || which == "SA"
        dmap = x -> -real(x)
    elseif which == "LI"
        dmap = imag
    elseif which == "SI"
        dmap = x -> -imag(x)
    end

    if cmplx
        d = Vector{T}(undef, nev+1)
        sigmar = Ref{T}(sigma)
        workev = Vector{T}(undef, 2ncv)
        neupd(ritzvec, howmny, select, d, v, ldv, sigmar, workev,
              bmat, n, which, nev, TOL, resid, ncv, v, ldv,
              iparam, ipntr, workd, workl, lworkl, rwork, info)
        if info[] != 0
            throw(XYEUPD_Exception(info[]))
        end

        p = sortperm(dmap(d[1:nev]), rev=true)
        return ritzvec ? (d[p], v[1:n, p],iparam[5],iparam[3],iparam[9],resid) : (d[p],iparam[5],iparam[3],iparam[9],resid)
    elseif sym
        d = Vector{T}(undef, nev)
        sigmar = Ref{T}(sigma)
        seupd(ritzvec, howmny, select, d, v, ldv, sigmar,
              bmat, n, which, nev, TOL, resid, ncv, v, ldv,
              iparam, ipntr, workd, workl, lworkl, info)
        if info[] != 0
            throw(XYEUPD_Exception(info[]))
        end

        p = sortperm(dmap(d), rev=true)
        return ritzvec ? (d[p], v[1:n, p],iparam[5],iparam[3],iparam[9],resid) : (d[p],iparam[5],iparam[3],iparam[9],resid)
    else
        dr = Vector{T}(undef, nev+1)
        di = Vector{T}(undef, nev+1)
        fill!(dr,NaN)
        fill!(di,NaN)
        sigmar = Ref{T}(real(sigma))
        sigmai = Ref{T}(imag(sigma))
        workev = Vector{T}(undef, 3*ncv)
        neupd(ritzvec, howmny, select, dr, di, v, ldv, sigmar, sigmai,
              workev, bmat, n, which, nev, TOL, resid, ncv, v, ldv,
              iparam, ipntr, workd, workl, lworkl, info)
        if info[] != 0
            throw(XYEUPD_Exception(info[]))
        end
        evec = complex.(Matrix{T}(undef, n, nev+1), Matrix{T}(undef, n, nev+1))

        j = 1
        while j <= nev
            if di[j] == 0
                evec[:,j] = v[:,j]
            else # For complex conjugate pairs
                evec[:,j]   = v[:,j] + im*v[:,j+1]
                evec[:,j+1] = v[:,j] - im*v[:,j+1]
                j += 1
            end
            j += 1
        end
        if j == nev+1 && !isnan(di[j])
            if di[j] == 0
                evec[:,j] = v[:,j]
                j += 1
            else
                throw(XYEUPD_Exception(ERR_UNEXPECTED_BEHAVIOR))
            end
        end

        d = complex.(dr, di)

        if j == nev+1
            p = sortperm(dmap(d[1:nev]), rev=true)
        else
            p = sortperm(dmap(d), rev=true)
            p = p[1:nev]
        end

        return ritzvec ? (d[p], evec[1:n, p],iparam[5],iparam[3],iparam[9],resid) : (d[p],iparam[5],iparam[3],iparam[9],resid)
    end
end

for (T, saupd_name, seupd_name, naupd_name, neupd_name) in
    ((:Float64, :dsaupd_, :dseupd_, :dnaupd_, :dneupd_),
     (:Float32, :ssaupd_, :sseupd_, :snaupd_, :sneupd_))
    @eval begin
        function naupd(ido, bmat, n, evtype, nev, TOL::Ref{$T}, resid::Vector{$T}, ncv, v::Matrix{$T}, ldv,
                       iparam, ipntr, workd::Vector{$T}, workl::Vector{$T}, lworkl, info)
            ccall(($(string(naupd_name)), libarpack), Cvoid,
                  (Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt},
                   Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ptr{$T}, Ref{BlasInt},
                   Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ref{BlasInt}, Clong, Clong),
                  ido, bmat, n, evtype, nev,
                  TOL, resid, ncv, v, ldv,
                  iparam, ipntr, workd, workl, lworkl, info, 1, 2)
        end

        function neupd(rvec, howmny, select, dr, di, z, ldz, sigmar, sigmai,
                  workev::Vector{$T}, bmat, n, evtype, nev, TOL::Ref{$T}, resid::Vector{$T}, ncv, v, ldv,
                  iparam, ipntr, workd::Vector{$T}, workl::Vector{$T}, lworkl, info)
            ccall(($(string(neupd_name)), libarpack), Cvoid,
                  (Ref{BlasInt}, Ptr{UInt8}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ptr{$T}, Ref{BlasInt},
                   Ref{$T}, Ref{$T}, Ptr{$T}, Ptr{UInt8}, Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt},
                   Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ptr{$T}, Ref{BlasInt},
                   Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ref{BlasInt}, Clong, Clong, Clong),
                  rvec, howmny, select, dr, di, z, ldz,
                  sigmar, sigmai, workev, bmat, n, evtype, nev,
                  TOL, resid, ncv, v, ldv,
                  iparam, ipntr, workd, workl, lworkl, info, 1, 1, 2)
        end

        function saupd(ido, bmat, n, which, nev, TOL::Ref{$T}, resid::Vector{$T}, ncv, v::Matrix{$T}, ldv,
                       iparam, ipntr, workd::Vector{$T}, workl::Vector{$T}, lworkl, info)
            ccall(($(string(saupd_name)), libarpack), Cvoid,
                  (Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt},
                   Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ptr{$T}, Ref{BlasInt},
                   Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ref{BlasInt}, Clong, Clong),
                  ido, bmat, n, which, nev,
                  TOL, resid, ncv, v, ldv,
                  iparam, ipntr, workd, workl, lworkl, info, 1, 2)
        end

        function seupd(rvec, howmny, select, d, z, ldz, sigma,
                       bmat, n, evtype, nev, TOL::Ref{$T}, resid::Vector{$T}, ncv, v::Matrix{$T}, ldv,
                       iparam, ipntr, workd::Vector{$T}, workl::Vector{$T}, lworkl, info)
            ccall(($(string(seupd_name)), libarpack), Cvoid,
                  (Ref{BlasInt}, Ptr{UInt8}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt},
                   Ptr{$T}, Ptr{UInt8}, Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt},
                   Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ptr{$T}, Ref{BlasInt},
                   Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ref{BlasInt}, Clong, Clong, Clong),
                  rvec, howmny, select, d, z, ldz,
                  sigma, bmat, n, evtype, nev,
                  TOL, resid, ncv, v, ldv,
                  iparam, ipntr, workd, workl, lworkl, info, 1, 1, 2)
        end
    end
end

for (T, TR, naupd_name, neupd_name) in
    ((:ComplexF64, :Float64, :znaupd_, :zneupd_),
     (:ComplexF32, :Float32, :cnaupd_, :cneupd_))
    @eval begin
        function naupd(ido, bmat, n, evtype, nev, TOL::Ref{$TR}, resid::Vector{$T}, ncv, v::Matrix{$T}, ldv,
                       iparam, ipntr, workd::Vector{$T}, workl::Vector{$T}, lworkl,
                       rwork::Vector{$TR}, info)
            ccall(($(string(naupd_name)), libarpack), Cvoid,
                  (Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt},
                   Ptr{$TR}, Ptr{$T}, Ref{BlasInt}, Ptr{$T}, Ref{BlasInt},
                   Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ptr{$TR}, Ref{BlasInt}, Clong, Clong),
                  ido, bmat, n, evtype, nev,
                  TOL, resid, ncv, v, ldv,
                  iparam, ipntr, workd, workl, lworkl, rwork, info, 1, 2)
        end

        function neupd(rvec, howmny, select, d, z, ldz, sigma, workev::Vector{$T},
                       bmat, n, evtype, nev, TOL::Ref{$TR}, resid::Vector{$T}, ncv, v::Matrix{$T}, ldv,
                       iparam, ipntr, workd::Vector{$T}, workl::Vector{$T}, lworkl,
                       rwork::Vector{$TR}, info)
            ccall(($(string(neupd_name)), libarpack), Cvoid,
                  (Ref{BlasInt}, Ptr{UInt8}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt},
                   Ptr{$T}, Ptr{$T}, Ptr{UInt8}, Ref{BlasInt}, Ptr{UInt8}, Ref{BlasInt},
                   Ptr{$TR}, Ptr{$T}, Ref{BlasInt}, Ptr{$T}, Ref{BlasInt},
                   Ptr{BlasInt}, Ptr{BlasInt}, Ptr{$T}, Ptr{$T}, Ref{BlasInt}, Ptr{$TR}, Ref{BlasInt}, Clong, Clong, Clong),
                  rvec, howmny, select, d, z, ldz,
                  sigma, workev, bmat, n, evtype, nev,
                  TOL, resid, ncv, v, ldv,
                  iparam, ipntr, workd, workl, lworkl, rwork, info, 1, 1, 2)
        end
    end
end


module ElboMaximize

using Optim
using Optim: Options, NewtonTrustRegion
using ..Model
using ..SensitiveFloats
using ..DeterministicVI
using ..DeterministicVI: init_thread_pool!, ElboIntermediateVariables
using ..DeterministicVI.ConstraintTransforms: TransformDerivatives,
                                              VariationalParams,
                                              ConstraintBatch,
                                              ParameterConstraint,
                                              BoxConstraint,
                                              SimplexConstraint,
                                              allocate_free_params,
                                              to_free!,
                                              to_bound!,
                                              enforce!,
                                              propagate_derivatives!

##################################################
# defaults for optional arguments to `maximize!` #
##################################################

struct ElboConfig{N,T}
    bound_params::VariationalParams{T}
    free_params::VariationalParams{T}
    free_initial_input::Vector{T}
    free_previous_input::Vector{T}
    free_result::SensitiveFloat{T}
    bvn_bundle::Model.BvnBundle{T}
    constraints::ConstraintBatch
    derivs::TransformDerivatives{N,T}
    optim_options::Options
    trust_region::NewtonTrustRegion{T}
end

function ElboConfig{T}(ea::ElboArgs,
                   vp::VariationalParams{T},
                   bound_params::VariationalParams{T} = vp[ea.active_sources];
                   termination_callback = nothing,
                   loc_width::Float64 = 1e-4,
                   loc_scale::Float64 = 1.0,
                   max_iters::Int = 50,
                   constraints::ConstraintBatch = elbo_constraints(bound_params, loc_width, loc_scale),
                   optim_options::Options = elbo_optim_options(max_iters=max_iters,
                                                               termination_callback=termination_callback),
                   trust_region::NewtonTrustRegion = elbo_trust_region())
    free_params = allocate_free_params(bound_params, constraints)
    free_initial_input = to_flat_vector(free_params)
    free_previous_input = similar(free_initial_input)
    free_result = SensitiveFloat{Float64}(length(free_params[1]), length(bound_params), true, true)
    bvn_bundle = Model.BvnBundle{T}(ea.psf_K, ea.S)
    derivs = TransformDerivatives(bound_params, free_params)
    return ElboConfig(bound_params, free_params, free_initial_input,
                  free_previous_input, free_result, bvn_bundle,
                  constraints, derivs, optim_options, trust_region)
end

function elbo_constraints{T}(bound::VariationalParams{T},
                             loc_width::Real = 1.0e-4,
                             loc_scale::Real = 1.0)
    n_sources = length(bound)
    boxes = Vector{Vector{ParameterConstraint{BoxConstraint}}}(n_sources)
    simplexes = Vector{Vector{ParameterConstraint{SimplexConstraint}}}(n_sources)
    for src in 1:n_sources
        i1, i2 = ids.pos[1], ids.pos[2]
        u1, u2 = bound[src][i1], bound[src][i2]
        boxes[src] = [
            ParameterConstraint(BoxConstraint(u1 - loc_width, u1 + loc_width, loc_scale), i1),
            ParameterConstraint(BoxConstraint(u2 - loc_width, u2 + loc_width, loc_scale), i2),
            ParameterConstraint(BoxConstraint(1e-2, 0.99, 1.0), ids.gal_fracdev),
            ParameterConstraint(BoxConstraint(1e-2, 0.99, 1.0), ids.gal_ab),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.gal_angle),
            ParameterConstraint(BoxConstraint(0.10, 70.0, 1.0), ids.gal_scale),
            ParameterConstraint(BoxConstraint(-1.0, 10.0, 1.0), ids.flux_loc),
            ParameterConstraint(BoxConstraint(1e-4, 0.10, 1.0), ids.flux_scale),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.color_mean[:, 1]),
            ParameterConstraint(BoxConstraint(-10.0, 10.0, 1.0), ids.color_mean[:, 2]),
            ParameterConstraint(BoxConstraint(1e-4, 1.0, 1.0), ids.color_var[:, 1]),
            ParameterConstraint(BoxConstraint(1e-4, 1.0, 1.0), ids.color_var[:, 2])
        ]
        simplexes[src] = [
            ParameterConstraint(SimplexConstraint(0.005, 1.0, 2), ids.is_star),
            ParameterConstraint(SimplexConstraint(0.01/NUM_COLOR_COMPONENTS, 1.0, NUM_COLOR_COMPONENTS), ids.k[:, 1]),
            ParameterConstraint(SimplexConstraint(0.01/NUM_COLOR_COMPONENTS, 1.0, NUM_COLOR_COMPONENTS), ids.k[:, 2])
        ]
    end
    return ConstraintBatch(boxes, simplexes)
end

function elbo_optim_options(args...; xtol_abs = 1e-7, ftol_rel = 1e-6, max_iters = 50,
                            termination_callback = nothing, kwargs...)
    return Optim.Options(args...; x_tol = xtol_abs, f_tol = ftol_rel,
                         g_tol = 1e-8, iterations = max_iters,
                         store_trace = false, show_trace = false,
                         extended_trace = false,
                         callback = termination_callback,
                         kwargs...)
end

function elbo_trust_region(; initial_delta = 1.0, delta_hat = 1e9)
    return Optim.NewtonTrustRegion(initial_delta = initial_delta,
                                   delta_hat = delta_hat)
end

function enforce_references!{T}(ea::ElboArgs, vp::VariationalParams{T}, cfg::ElboConfig)
    @assert length(cfg.bound_params) == length(ea.active_sources)
    for i in 1:length(ea.active_sources)
        cfg.bound_params[i] = vp[ea.active_sources[i]]
    end
    return nothing
end

to_flat_vector{T}(sources::VariationalParams{T}) = vcat(sources...)::Vector{T}

function to_flat_vector!{T}(x::Vector{T}, sources::VariationalParams{T})
    i = 1
    for src in sources
        for j in eachindex(src)
            x[i] = src[j]
            i += 1
        end
    end
    return sources
end

function to_variational_params!{T}(sources::VariationalParams{T}, x::Vector{T})
    i = 1
    for src in sources
        for j in eachindex(src)
            src[j] = x[i]
            i += 1
        end
    end
    return sources
end

##########################################
# Preallocated ElboIntermediateVariables #
##########################################

const ELBO_VARS_POOL = Vector{ElboIntermediateVariables{Float64}}()

get_elbo_vars() = ELBO_VARS_POOL[Base.Threads.threadid()]

function __init__()
    init_thread_pool!(ELBO_VARS_POOL, () -> ElboIntermediateVariables(Float64, 1, true, true))
end

# explicitly call this for use with compiled system image
__init__()

##################################
# Callable Types Passed to Optim #
##################################

function evaluate!{T}(ea::ElboArgs, vp::VariationalParams{T}, cfg::ElboConfig, x::Vector)
    if x != cfg.free_previous_input
        copy!(cfg.free_previous_input, x)
        to_variational_params!(cfg.free_params, x)
        to_bound!(cfg.bound_params, cfg.free_params, cfg.constraints)
        bound_result = elbo(ea, vp, get_elbo_vars(), cfg.bvn_bundle)
        propagate_derivatives!(to_bound!, bound_result,
                               cfg.free_result, cfg.free_params,
                               cfg.constraints, cfg.derivs)
    end
    return nothing
end

# Objective #
#-----------#

struct Objective{N,T} <: Function
    ea::ElboArgs
    vp::VariationalParams{T}
    cfg::ElboConfig{N,T}
end

function (f::Objective)(x::Vector)
    evaluate!(f.ea, f.vp, f.cfg, x)
    return -(f.cfg.free_result.v[])
end

# Gradient #
#----------#

struct Gradient{N,T} <: Function
    ea::ElboArgs
    vp::VariationalParams{T}
    cfg::ElboConfig{N,T}
end

function (f::Gradient)(x::Vector, result::Vector)
    evaluate!(f.ea, f.vp, f.cfg, x)
    free_gradient = f.cfg.free_result.d
    for i in eachindex(result)
        result[i] = -(free_gradient[i])
    end
    return result
end

# Hessian #
#---------#

struct Hessian{N,T} <: Function
    ea::ElboArgs
    vp::VariationalParams{T}
    cfg::ElboConfig{N,T}
end

function (f::Hessian)(x::Vector, result::Matrix)
    evaluate!(f.ea, f.vp, f.cfg, x)
    free_hessian = f.cfg.free_result.h
    for i in eachindex(result)
        result[i] = -(free_hessian[i])
    end
    return result
end

#############
# maximize! #
#############

function maximize!(ea::ElboArgs, vp::VariationalParams{Float64}, cfg::ElboConfig = ElboConfig(ea, vp))
    enforce_references!(ea, vp, cfg)
    enforce!(cfg.bound_params, cfg.constraints)
    to_free!(cfg.free_params, cfg.bound_params, cfg.constraints)
    to_flat_vector!(cfg.free_initial_input, cfg.free_params)
    fill!(cfg.free_previous_input, NaN)
    R = Optim.MultivariateOptimizationResults{Float64,1,Optim.NewtonTrustRegion{Float64}}
    result::R = Optim.optimize(Objective(ea, vp, cfg), Gradient(ea, vp, cfg), Hessian(ea, vp, cfg),
                               cfg.free_initial_input, cfg.trust_region, cfg.optim_options)
    min_value::Float64 = -(Optim.minimum(result))
    min_solution::Vector{Float64} = Optim.minimizer(result)
    to_variational_params!(cfg.free_params, min_solution)
    to_bound!(cfg.bound_params, cfg.free_params, cfg.constraints)
    return result.f_calls, min_value, min_solution, result
end

end # module