# Generate 2D vertices in a rectangular Lattice
function _generate_2d_vertices!(vertices, nx, ny, LL, LR, UR, UL)
    for i in 0:ny-1
      ratio_bounds = i / (ny-1)

      x0 = LL[1] * (1 - ratio_bounds) + ratio_bounds * UL[1]
      x1 = LR[1] * (1 - ratio_bounds) + ratio_bounds * UR[1]

      y0 = LL[2] * (1 - ratio_bounds) + ratio_bounds * UL[2]
      y1 = LR[2] * (1 - ratio_bounds) + ratio_bounds * UR[2]

      for j in 0:nx-1
          ratio = j / (nx-1)
          x = x0 * (1 - ratio) + ratio * x1
          y = y0 * (1 - ratio) + ratio * y1
          push!(vertices, Vertex(Tensors.Vec{2}((x, y))))
      end
  end
  end

  # Check edge orientation consistency
  function _check_vertex_data(vertices, n1,n2,n3)
  a = vertices[n2].x-vertices[n1].x
  b = vertices[n3].x-vertices[n1].x
  if (a[1]*b[2]-a[2]*b[1]) < 0
      #swap vertices 2 and 3
      return (n1,n3,n2)
  end
  return (n1,n2,n3)
  end

@inline _mapToGlobalIdx(cells,cellidx,localvertexidx) = cells[cellidx].vertices[localvertexidx]

function _get_vertexset_from_edges(cells,edgeset)
  vertices = Set{Int}()
  for edge in edgeset
      push!(vertices, _mapToGlobalIdx(cells, edge.cellidx, reference_edge_vertices(cells[edge.cellidx])[edge.idx][1]))
      push!(vertices, _mapToGlobalIdx(cells, edge.cellidx, reference_edge_vertices(cells[edge.cellidx])[edge.idx][2]))
  end
  return vertices
end

# Shortcuts
unitSquareMesh(cellType, nel::NTuple{2,Int}) = rectangle_mesh(cellType, nel, Tensors.Vec{2}((0.0,0.0)), Tensors.Vec{2}((1.0,1.0)))

#########################
# Triangle Cells 2D   #
#########################
"""
rectangle_mesh(::Type{RefTetrahedron}, ::Type{Val{2}}, nel::NTuple{2,Int}, LL::Vec{2,T}, UR::Vec{2,T})
Generate a rectangular mesh with triangular cells, where `LL` is the low left vertex
and `UR` is the upper right one. `nel` is a tuple with the number of partions to be
used in each dimension.
"""
function rectangle_mesh(::Type{TriangleCell}, nel::NTuple{2,Int}, LL::Tensors.Vec{2,T}, UR::Tensors.Vec{2,T}) where {T}
    LR = Tensors.Vec{2}((UR[1],LL[2]))
    UL = Tensors.Vec{2}((LL[1],UR[2]))
    nel_x = nel[1]; nel_y = nel[2]; nel_tot = 2*nel_x*nel_y
    n_vertices_x = nel_x + 1; n_vertices_y = nel_y + 1
    n_vertices = n_vertices_x * n_vertices_y

    # Generate vertices
    vertices = Vertex{2,T}[]
    _generate_2d_vertices!(vertices, n_vertices_x, n_vertices_y, LL, LR, UR, UL)

    # Generate cells
    vertex_array = reshape(collect(1:n_vertices), (n_vertices_x, n_vertices_y))
    cells = TriangleCell[]
    for j in 1:nel_y, i in 1:nel_x
        push!(cells, TriangleCell((vertex_array[i,j], vertex_array[i+1,j], vertex_array[i,j+1]))) # ◺
        push!(cells, TriangleCell((vertex_array[i+1,j], vertex_array[i+1,j+1], vertex_array[i,j+1]))) # ◹
    end

    # Cell edges
    cell_array = reshape(collect(1:nel_tot),(2, nel_x, nel_y))
    boundary = EdgeIndex[[EdgeIndex(cl, 3) for cl in cell_array[1,:,1]];
                           [EdgeIndex(cl, 3) for cl in cell_array[2,end,:]];
                           [EdgeIndex(cl, 1) for cl in cell_array[2,:,end]];
                           [EdgeIndex(cl, 2) for cl in cell_array[1,1,:]]]
    offset = 0
    edgesets = Dict{String,Set{EdgeIndex}}()
    edgesets["bottom"] = Set{EdgeIndex}(boundary[(1:length(cell_array[1,:,1]))   .+ offset]); offset += length(cell_array[1,:,1])
    edgesets["right"]  = Set{EdgeIndex}(boundary[(1:length(cell_array[2,end,:])) .+ offset]); offset += length(cell_array[2,end,:])
    edgesets["top"]    = Set{EdgeIndex}(boundary[(1:length(cell_array[2,:,end])) .+ offset]); offset += length(cell_array[2,:,end])
    edgesets["left"]   = Set{EdgeIndex}(boundary[(1:length(cell_array[1,1,:]))   .+ offset]); offset += length(cell_array[1,1,:])
    edgesets["boundary"] = union(edgesets["bottom"],edgesets["right"],edgesets["top"],edgesets["left"])
    vertexsets = Dict{String,Set{Int}}()
    for set in edgesets
        vertexsets[set.first] = _get_vertexset_from_edges(cells,set.second)
    end
    return PolytopalMesh(cells, vertices; edgesets = edgesets, vertexsets = vertexsets)
end

#########################
# Rectangle Cells 2D   #
#########################
@inline _build_cell(::Type{RectangleCell}, el_vertices, el_faces) = RectangleCell(el_vertices,(el_faces[1],el_faces[2],el_faces[3],el_faces[4]))
"""
rectangle_mesh(::Type{RectangleCell}, nel::NTuple{2,Int}, LL::Vec{2,T}, UR::Vec{2,T}) where {T}
Generate a rectangular mesh with triangular cells, where `LL` is the low left vertex
and `UR` is the upper right one. `nel` is a tuple with the number of partions to be
used in each dimension.
"""
function rectangle_mesh(::Type{RectangleCell}, nel::NTuple{2,Int}, LL::Tensors.Vec{2,T}, UR::Tensors.Vec{2,T}) where {T}
    LR = Tensors.Vec{2}((UR[1],LL[2]))
    UL = Tensors.Vec{2}((LL[1],UR[2]))
    nel_x = nel[1]; nel_y = nel[2]; nel_tot = nel_x*nel_y
    n_vertices_x = nel_x + 1; n_vertices_y = nel_y + 1
    n_vertices = n_vertices_x * n_vertices_y

    # Generate vertices
    vertices = Vertex{2,T}[]
    _generate_2d_vertices!(vertices, n_vertices_x, n_vertices_y, LL, LR, UR, UL)

    # Generate cells
    vertex_array = reshape(collect(1:n_vertices), (n_vertices_x, n_vertices_y))
    cells = RectangleCell[]
    for j in 1:nel_y, i in 1:nel_x
        push!(cells, RectangleCell((vertex_array[i,j], vertex_array[i+1,j], vertex_array[i+1,j+1], vertex_array[i,j+1])))
    end

    # Cell faces
    cell_array = reshape(collect(1:nel_tot),(nel_x, nel_y))
    boundary = EdgeIndex[[EdgeIndex(cl, 1) for cl in cell_array[:,1]];
                      [EdgeIndex(cl, 2) for cl in cell_array[end,:]];
                      [EdgeIndex(cl, 3) for cl in cell_array[:,end]];
                      [EdgeIndex(cl, 4) for cl in cell_array[1,:]]]
    # Cell face sets
    offset = 0
    edgesets = Dict{String, Set{EdgeIndex}}()
    edgesets["bottom"] = Set{EdgeIndex}(boundary[(1:length(cell_array[:,1]))   .+ offset]); offset += length(cell_array[:,1])
    edgesets["right"]  = Set{EdgeIndex}(boundary[(1:length(cell_array[end,:])) .+ offset]); offset += length(cell_array[end,:])
    edgesets["top"]    = Set{EdgeIndex}(boundary[(1:length(cell_array[:,end])) .+ offset]); offset += length(cell_array[:,end])
    edgesets["left"]   = Set{EdgeIndex}(boundary[(1:length(cell_array[1,:]))   .+ offset]); offset += length(cell_array[1,:])
    edgesets["boundary"] = union(edgesets["bottom"],edgesets["right"],edgesets["top"],edgesets["left"])
    vertexsets = Dict{String,Set{Int}}()
    for set in edgesets
        vertexsets[set.first] = _get_vertexset_from_edges(cells,set.second)
    end
    return PolytopalMesh(cells, vertices; edgesets = edgesets, vertexsets = vertexsets)
end

#########################
# Hexagonal Cells 2D
#############################

function _gen_hexagon(centroid,LL,UR,hex_width,hex_heigth)
    hex_verts = [
        Vertex(Tensors.Vec{2}((centroid[1]-hex_width/2,centroid[2]-hex_heigth/4))),
        Vertex(Tensors.Vec{2}((centroid[1],centroid[2]-hex_heigth/2))),
        Vertex(Tensors.Vec{2}((centroid[1]+hex_width/2,centroid[2]-hex_heigth/4))),
        Vertex(Tensors.Vec{2}((centroid[1]+hex_width/2,centroid[2]+hex_heigth/4))),
        Vertex(Tensors.Vec{2}((centroid[1],centroid[2]+hex_heigth/2))),
        Vertex(Tensors.Vec{2}((centroid[1]-hex_width/2,centroid[2]+hex_heigth/4))),
    ]
    filter_verts = filter(x->(x.x[1]>=LL[1] && x.x[1]<=UR[1]
                          && x.x[2]>=LL[2] && x.x[2] <= UR[2]),hex_verts)
    if size(filter_verts,1) == 2 #corners case
        x = centroid[1] + (centroid[1] == LL[1] ? hex_width/2 : -hex_width/2)
        new_vertex = Vertex(Tensors.Vec{2}((x, centroid[2])))
        if centroid == LL
            filter_verts = (Vertex(centroid),new_vertex,filter_verts...)
        elseif centroid[2] == LL[2] && centroid[1] == UR[1]
            filter_verts = (new_vertex,Vertex(centroid),filter_verts...)
        elseif centroid[2] == UR[2] && centroid[1] == LL[1]
            filter_verts = (filter_verts...,new_vertex,Vertex(centroid))
        elseif centroid == UR
            filter_verts = (filter_verts...,Vertex(centroid),new_vertex)
        else
            throw("error on hexagon at corner $centroid")
        end
    elseif size(filter_verts,1) == 3 #top and bottom pentagon/triangle case
        nv1 = Vertex(Tensors.Vec{2}((centroid[1]-hex_width/2, centroid[2])))
        nv2 = Vertex(Tensors.Vec{2}((centroid[1]+hex_width/2, centroid[2])))
        if centroid[2] == LL[2]
            filter_verts = (nv1,nv2,filter_verts...)
        elseif centroid[2] == UR[2]
            filter_verts = (nv1,filter_verts...,nv2)
        elseif centroid[2] > UR[2]
            filter_verts = (filter_verts...,)
        else
            throw("error at bottom centroid $centroid")
        end
    end
    return filter_verts
end

function rectangle_mesh(::Type{HexagonCell}, nel::NTuple{2,Int}, LL::Tensors.Vec{2,T}, UR::Tensors.Vec{2,T}) where {T}
    LR = Tensors.Vec{2}((UR[1],LL[2]))
    UL = Tensors.Vec{2}((LL[1],UR[2]))
    nel_x = nel[1]; nel_y = nel[2]
    nel_tot = 2*nel_x*nel_y +nel_y - nel_x

    # Generate vertices
    vertices = Vertex{2,T}[]
    hex_width = 1.0
    hex_heigth = 1.0
    w_scale = (LR[1] - LL[1])/nel_x
    h_scale = (UL[2] - LL[1])/nel_y
    LLn = Tensors.Vec{2}((0.0,0.0))
    URn = Tensors.Vec{2}((nel_x,nel_y))

    map_func(x::Vertex) = Vertex(Tensors.Vec{2}((x.x[1]*w_scale+LL[1],x.x[2]*h_scale+LL[2])))

    centroids = Tensors.Vec{2,T}[]
    _generate_2d_hex_centroids!(centroids,LL, nel_y, nel_x, hex_width, hex_heigth)
    cells = Cell[]
    used_vertices = Dict{Vertex,Int}()
    #Add cells
    nextvert = 1 # next free vertex to use
    for c_i in 1:size(centroids,1)
        cell_verts = Int[]
        for vert in _gen_hexagon(centroids[c_i],LLn,URn,hex_width,hex_heigth)
            nextvert = _push_cell_vertex!(vert,cell_verts, used_vertices,vertices,nextvert, map_func)
        end
        n = size(cell_verts,1)
        push!(cells,Cell{2,n,n,1}(Tuple(cell_verts)))
    end

    # Cell edges
    edgesets = Dict{String,Set{EdgeIndex}}()
    ncells = size(cells,1)
    t_edge = mod1(nel_y,3)==1 ? 4 : 3
    t_idxs = (ncells-nel_x + (mod1(nel_y,3)==3 ? 0 : 1)):ncells
    l_bounds = Tuple{Int,Int}[]
    r_bounds = Tuple{Int,Int}[]
    c_idx = 1
    for i in 1:nel_y
        if mod1(i,3) == 1
            push!(l_bounds,(c_idx,4))
            push!(r_bounds,(c_idx+nel_x,2))
            ed_idx = (i == nel_y ? 5 : 6)
            push!(l_bounds,(c_idx+nel_x+1,ed_idx))
            push!(r_bounds,(c_idx+2*nel_x,3))
            c_idx = c_idx + nel_x*2+1
        elseif mod1(i,3) == 2
            push!(l_bounds,(c_idx,4))
            push!(r_bounds,(c_idx+nel_x,2))
            c_idx = c_idx + nel_x + 1
        else
            push!(l_bounds,(c_idx,6))
            push!(r_bounds,(c_idx+nel_x-1,3))
            if i == nel_y
                push!(l_bounds,(c_idx+nel_x,4))
                push!(r_bounds,(c_idx+2*nel_x,2))
            end
            c_idx = c_idx + nel_x
        end
    end
    edgesets["bottom"] = Set{EdgeIndex}([EdgeIndex(c_i,1) for c_i in 1:(nel_x+1)])
    edgesets["top"]    = Set{EdgeIndex}([EdgeIndex(i,t_edge) for i in t_idxs])
    edgesets["right"]  = Set{EdgeIndex}([EdgeIndex(x[1],x[2]) for x in r_bounds])
    edgesets["left"]   = Set{EdgeIndex}([EdgeIndex(x[1],x[2]) for x in l_bounds])
    edgesets["boundary"] = union(edgesets["bottom"],edgesets["right"],edgesets["top"],edgesets["left"])

    vertexsets = Dict{String,Set{Int}}()
    for set in edgesets
        vertexsets[set.first] = _get_vertexset_from_edges(cells,set.second)
    end
    return PolytopalMesh(cells, vertices; edgesets = edgesets, vertexsets = vertexsets)
end


# HuberLoss
using LaTeXStrings, LossFunctions, Measures, Latexify

Plots.plot(HuberLoss(1.5), c=:olive, linewidth=2)
Plots.plot!(HuberLoss(1), c=:orangered4, linewidth=2.5)
Plots.plot!(HuberLoss(.5), c=:springgreen4, linewidth=2.5)


function minmod(a,b,c)
  if (a > 0 && b > 0 && c > 0)
    min(a,b,c)
  elseif (a < 0 && b < 0 && c < 0)
    max(a,b,c)
  else
    zero(a)
  end
end

function minmod(a,b)
  0.5*(sign(a)+sign(b))*min(abs(a),abs(b))
end

"Apply slope limiter πᴺ to u assuming u an Nth polynomial"
mutable struct DGLimiter{pType}
  problem::AbstractConservationLawProblem
  basis::PolynomialBasis
  limiter::AbstractDGLimiter
  params::pType
end

function DGLimiter(problem, basis, limiter)
    DGLimiter(problem, basis, limiter, params_init(limiter, basis, problem))
end

# Trivial parameter initialization
params_init(limiter::AbstractDGLimiter, basis::PolynomialBasis, problem::AbstractConservationLawProblem) = nothing

function (dgLimiter::DGLimiter)(u, f, t)
    @unpack problem, basis, limiter, params = dgLimiter
    apply_limiter!(u,f,t,problem, basis, params, limiter)
end

# Auxiliary functions
"""
get_cell_averages!(uavg, basis, uh, NN)
Compute cell averages of polynomial system u,
using modal values `uh`, a polynomial `basis` and number of
variables `NC`
"""
function get_cell_averages!(uavg::AbstractArray{T,2}, basis::PolynomialBasis, uh::AbstractArray{T,2}, NC::Int) where {T}
    NN = basis.order + 1
    for k = 1:NC
        uavg[k,:] = [basis.φ[1,1]*uh[(k-1)*NN+1,i] for i in 1:size(uh,2)]
    end
end

#MUSCL Limiter Van Leer
#Reference: Hesthaven, Warburton, Nodal Discontinuous Galerkin Methods Algorithms, Analysis and
#applications, pp 150.
immutable Linear_MUSCL_Limiter <: AbstractDGLimiter end

"""
apply_limiter!(u,f,t,mesh, basis, limiter::Linear_MUSCL_Limiter)
    apply Linear MUSCL limiter on function u
  """
function apply_limiter!(u,f,t,prob, basis, params, imiter::Linear_MUSCL_Limiter)
    # Compute modal coefficients
    mesh = prob.mesh; NC = prob.numvars
    NN = basis.order + 1
    uh = myblock(basis.invφ,NC)*u
    @assert basis.order > 1 "MUSCL limiter requires at least linear polynomials"
    #Extract linear polynomial
    ul = uh[:,:];
    for k = 1:NC;ul[((k-1)*NN+3):k*NN,:] = 0.0;end

    # Compute cell averages
    uavg = zeros(eltype(u), NC, numcells(mesh))
    get_cell_averages!(uavg, basis, uh, NC)

    #compute nodes in real coordinates
    x1 = zeros(basis.order+1,numcells(mesh))
    for k in 1:numcells(mesh)
        x1[:,k] = reference_to_interval(basis.nodes, (cell_faces(mesh)[k],cell_faces(mesh)[k+1]))
    end
    x0 = cell_centers(mesh)
    # Compute derivative
    h = cell_volumes(mesh)
    ux = 2*(1./h)'.*(myblock(basis.dφ,NC)*ul)

    ve = apply_boundary(uavg, mesh)
    uavgp1 = ve[:,3:end];  uavgm1 = ve[:,1:end-2]

    #limit nodal cell values
    for i in 1:size(u,2)
        for k = 1:NC
            for j in 1:NN
                u[(k-1)*NN+j,i] = uavg[k,i]+(x1[j,i]-x0[i])*
                minmod(ux[(k-1)*NN+j,i],(uavgp1[k,i]-uavg[k,i])/h[i],(uavg[k,i]-uavgm1[k,i])/h[i])
            end
      end
    end
end

#WENO limiter by Zhong-Shu (2013)
mutable struct WENO_Limiter <: AbstractDGLimiter end

function params_init(WENO_Limiter::WENO_Limiter, basis::PolynomialBasis, problem::AbstractConservationLawProblem)
    return WENOLimiterWeights(basis.order, basis.invφ)
end

"""
function WENOLimiterWeights(m, iV);
Compute operators to enable evaluation of WENO smoothness
indicator and WENO polynomial of order m.
"""
function WENOLimiterWeights(m,iV);
    Q = zeros(m+1,m+1)
    Pmat = zeros(m+1,m+1)
    Xm = Pmat; Xp = Pmat;

    # Compute quadrature points
    (x,w) = gausslegendre(m+1)
    Λ = diagm(w);

    # Initial matrices of Legendre polynomails
    for n = 0:m
      p = poly_legendre(n)
      # Eval nodes
      Pmat[n+1,:] = polyval(p, x)
      Xm[n+1,:] = polyval(p, x-2)
      Xp[n+1,:] = polyval(p, x+2)
    end

    # Compute matrices corresponding to increasing order of derivative
    for l=1:m
        # Set up operator to recover derivaties
        A = zeros(m+2-l,m+2-l); A[1,1] = 1/sqrt((2*l+1)*(2*l-1))
        A[m+2-l,m+2-l] = 1/(sqrt(2*(m+2)+1)*sqrt(2*(m+2)-1))
        for i=2:m-l+1
            Ah = 1/(sqrt(2*(l-1+i)+1)*sqrt(2*(l-1+i)-1))
            A[i,i] = Ah; A[i+1,i-1] = -Ah
        end

        # Recover derivatives at quadrature points
        Ph1 = A\Pmat[l:m+1,:]
        Pmat[1:l,:]=0; Pmat[l+1:m+1,:] = Ph1[1:m-l+1,:]

        # Compute smoothness operator for order l and update
        Qh = Pmat*Λ*Pmat'
        Q = Q + 2^(2*l-1)*Qh
    end

    # Initialize operator for smoothness indicator in nodal space
    Q = basis.invφ'*Q*basis.invφ;

    # Initialize interpolation matrices
    Xp = basis.invφ'*Xp; Xm = basis.invφ'*Xm
    return Q,Xm,Xp
end

"""
apply_limiter!(u,f,t,mesh, basis, limiter::WENO_Limiter)
    apply Zhong Shu WENO limiter on function u
  """
function apply_limiter!(u,f,t,prob, basis, params, limiter::WENO_Limiter)
    mesh = prob.mesh; NC = prob.numvars; NN = basis.order + 1
    (Q, Xp, Xm) = params
    Xp = myblock(Xp,NC)
    Xm = myblock(Xm,NC)
    iV = myblock(basis.invφ,NC)
    V = myblock(basis.φ,NC)
    eps0=1.0e-6;

    # Set constants for limiting
    eps1 = 1e-10; p=1;
    γm1 =0.001; γ0 = 0.998; γp1 = 0.001;

    # Compute cell averages and cell centers
    uh = iV*u
    uavg = zeros(eltype(u), NC, numcells(mesh))
    get_cell_averages!(uavg, basis, uh, NC)

    # Compute extended polynomials with zero cell averages
    ue = apply_boundary(u, mesh);
    Pm = Xp'*ue; Pp = Xm'*ue;
    Ph = iV*Pm; Ph[1:NN:end,:]=0; Pm = V*Ph;
    Ph = iV*Pp; Ph[1:NN:end,:]=0; Pp = V*Ph;

    # Extend cell averages
    ve = apply_boundary(uavg, mesh)

    # extract end values and cell averages for each element
    uel = u[1:NN:end,:]; uer = u[NN:NN:end,:]
    vj = uavg; vjm = ve[:,1:end-2]; vjp = ve[:,3:end]

    for k = 1:NC
        # Find elements that require limiting
        vel = vj[k,:] - minmod.(vj[k,:]-uel[k,:],vj[k,:]-vjm[k,:],vjp[k,:]-vj[k,:])
        ver = vj[k,:] + minmod.(uer[k,:]-vj[k,:],vj[k,:]-vjm[k,:],vjp[k,:]-vj[k,:])
        ids = union(find(x->abs(x)>eps0, vel-uel[k,:]), find(x->abs(x)>eps0, ver-uer[k,:]))

        # Apply limiting when needed
        if (!isempty(ids))
            # Extract local polynomials
            pm1 = Pm[(k-1)*NN+1:k*NN,ids] + ones(NN,1)*vj[k,ids]'
            p0  =  u[(k-1)*NN+1:k*NN,ids]
            pp1 = Pp[(k-1)*NN+1:k*NN,ids+2] + ones(NN,1)*vj[k,ids]'

            # Compute smoothness indicators and WENO weights
            βm1 = diag(pm1'*Q*pm1); αm1 = γm1./(eps1 + βm1).^(2*p);
            β0 = diag(p0'*Q*p0); α0  =  γ0./(eps1 + β0).^(2*p);
            βp1 = diag(pp1'*Q*pp1); αp1 = γp1./(eps1 + βp1).^(2*p);

            αs = αm1 + α0 + αp1;
            omm1 = αm1./αs; om0  = α0./αs; omp1 = αp1./αs;

            # Compute limited function
            u[(k-1)*NN+1:k*NN,ids] = pm1*diagm(omm1) + p0*diagm(om0) + pp1*diagm(omp1)
        end
    end
end

#TODO: Not Working
# immutable MUSCLV2Limiter <: AbstractDGLimiter end
#
# function apply_limiter(u,f,t,mesh, basis, limiter::MUSCLV2Limiter)
#     # Compute cell averages
#     ̄ū = [0.5*dot(basis.φ*u[:,i], basis.weights) for i in 1:size(u,2)]
#
#     # find end values of each element
#     ue1 = u[1,:];ue2=u[end,:]
#
#     #find cell averages
#     vk = ū; vkm1=[ū[1],ū[1:end-1]...];vkp1=[ū[2:end]...,ū[end]]
#
#     #apply reconstruction to find elements in need of limiting
#     ve1 = vk - minmod.(vk-ue1, vk -vkm1,vkp1-vk)
#     ve2 = vk + minmod.(ue2-vk,vk-vkm1,vkp1-vk)
#     tol = 1e-8
#     idx = (abs.(ve1-ue1).>tol)|(abs.(ve2-ue2).>tol)
#     if (!isempty(idx))
#       h2 = 2./diff(cell_faces(mesh))
#       x1 = zeros(basis.order+1,numcells(mesh))
#       for k in 1:numcells(mesh)
#         x1[:,k] = reference_to_interval(basis.nodes, (cell_faces(mesh)[k],cell_faces(mesh)[k+1]))
#       end
#       x0 = cell_centers(mesh)
#       ux = 0.5*h2.*((basis.dφ*u)'*basis.weights)
#       #println("idx: ", idx)
#       for j in 1:size(u,1)
#         u[j,idx] = vk[idx]+(x1[j,idx]-x0[idx]).*
#         minmod.(ux[idx],h2[idx].*(vkp1[idx]-vk[idx]),h2[idx].*(vk[idx]-vkm1[idx]))
#       end
#     end
# end


power_law_params = Dict(
    "A" => (α=0.108),
    "B" => (α=0.112),
    "C" => (α=0.120),
    "D" => (α=0.142),
    "E" => (α=0.203),
    "F" => (α=0.253)
)

"""
    windspeed(u0, z0, stability_class)
returns the windspeed function u(z) for a given Pasquill-Gifford
stability class, `z` is assumed to be in meters and `u` is in m/s

# Arguments
`u0` reference windspeed
`z0` height of reference point
`stability_class` Pasquill stability class (A, B, C, D, E, F)

# References

"""
function windspeed(u0::Number, z0::Number, stability_class::String)
    if stability_class ∈ Set(["A","B","C","D","E","F"])
        α, = power_law_params[stability_class]
    else
        err = "$stability_class is not a valid Pasquill-Gifford stability class"
        error(err)
    end

    return z -> u0*(z/z0)^α
end

"""
    windspeed(u, zR, λ, stability_class; k=0.35)
returns the windspeed function u(z) for a given Pasquill-Gifford
stability class, `z` is assumed to be in meters and `u` is in m/s

# Arguments
`u` friction velocity
`zR` surface roughness
`λ`  Monin-Obukhov length
`stability_class` Pasquill stability class (A, B, C, D, E, F)
`k`  von Karman's constant, 0.35

# References
    Businger et al. 1971
"""
function windspeed(u::Number, zR::Number, λ::Number, stability_class::String; k=0.35)
    if stability_class ∈ Set(["A","B","C"])
        a(z) = (1-15*(z/λ))^0.25
        Ψ(a) = 2*log((1+a)/2) + log((1+a^2)/2) - 2*atan(a) + π/2
        return z -> (u/k)*(log((z+zR)/zR) - Ψ(a(z)))
    elseif stability_class == "D"
        return z -> (u/k)*log((z+zR)/zR)
    elseif stability_class ∈ Set(["E", "F"])
        return z -> (u/k)*(log((z+zR)/zR) - 4.7*(z/λ))
    else
        err = "$stability_class is not a valid Pasquill-Gifford stability class"
        error(err)
    end
end


# Important packages
using JLD

# Loading data...
println("== Loading data...")
d = load("backup_matrices/eigenvalues.jld")["values"]
v = load("backup_matrices/eigenvectors.jld")["vectors"]

# Defining the size of the word vectors
println("== Defining word vectors size...")
m = 50

# Sorting eigenvalues by its absolute value (BubbleSort)
println("== BubbleSorting eigenvalues by its absolute values...")
for i = 1:length(d)
    swapped = false
    for j = 1:length(d) - i - 1
        if (abs(d[j]) < abs(d[j + 1]))
            aux = d[j]
            d[j] = d[j+1]
            d[j + 1] = aux
            swapped = true
        end
    end
    if swapped == false
        break
    end
end

# Filling a trunkated matrix (MxN) w/ most significant eigenvalues
println("== Filling the trunkated matrix (MxN)...")
d2 = zeros(m, length(d))
for i = 1:m
    d2[i, i] += sqrt(d[i])
end

# Reconstructing matrix
println("== Constructing W matrix with all the word vectors...")
word_vectors = d2 * v.'

# Saving reconstructed matrix
println("== Exporting word vectors matrix...")
save("backup_matrices/wordvectors.jld", "wordvectors", word_vectors)


"""Continuity equation's right-hand side with surface relaxation
-∂x(uh) - ∂y(vh) + γ*(η_ref - η)."""
function continuity_surf_relax!(η::Array{T,2},
                                Diag::DiagnosticVars{T,Tprog},
                                S::ModelSetup{T,Tprog},
                                t::Int) where {T,Tprog}

    @unpack dη = Diag.Tendencies
    @unpack dUdx,dVdy = Diag.VolumeFluxes
    @unpack η_ref = S.forcing
    @unpack γ = S.constants

    m,n = size(dη) .- (2,2)
    @boundscheck (m,n+2) == size(dUdx) || throw(BoundsError())
    @boundscheck (m+2,n) == size(dVdy) || throw(BoundsError())
    @boundscheck (m+2,n+2) == size(η) || throw(BoundsError())
    @boundscheck (m,n) == size(η_ref) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            dη[i+1,j+1] = -(Tprog(dUdx[i,j+1]) + Tprog(dVdy[i+1,j])) + Tprog(γ*(η_ref[i,j]-η[i+1,j+1]))
        end
    end
end

"""Continuity equation's right-hand side with time&space dependent forcing."""
function continuity_forcing!(   Diag::DiagnosticVars{T,Tprog},
                                S::ModelSetup{T,Tprog},
                                t::Int) where {T,Tprog}

    @unpack dη = Diag.Tendencies
    @unpack dUdx,dVdy = Diag.VolumeFluxes
    @unpack Fη = S.forcing

    m,n = size(dη) .- (2,2)         # cut off halo
    @boundscheck (m,n+2) == size(dUdx) || throw(BoundsError())
    @boundscheck (m+2,n) == size(dVdy) || throw(BoundsError())
    @boundscheck (m,n) == size(Fη) || throw(BoundsError())

    # avoid recomputation
    @unpack ωFη = S.constants
    Fηt = Ftime(T,t,ωFη)

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            dη[i+1,j+1] = -(Tprog(dUdx[i,j+1]) + Tprog(dVdy[i+1,j])) + Tprog(Fηt*Fη[i,j])
        end
    end
end

"""Continuity equation's right-hand side -∂x(uh) - ∂y(vh) without forcing."""
function continuity_itself!(Diag::DiagnosticVars{T,Tprog},
                            S::ModelSetup{T,Tprog},
                            t::Int) where {T,Tprog}

    @unpack dη = Diag.Tendencies
    @unpack dUdx,dVdy = Diag.VolumeFluxes

    m,n = size(dη) .- (2,2)     # cut off halo
    @boundscheck (m,n+2) == size(dUdx) || throw(BoundsError())
    @boundscheck (m+2,n) == size(dVdy) || throw(BoundsError())

    @inbounds for j ∈ 1:n
        for i ∈ 1:m
            dη[i+1,j+1] = -(Tprog(dUdx[i,j+1]) + Tprog(dVdy[i+1,j]))
        end
    end
end

"""Transit function to call the specified continuity function."""
function continuity!(   η::AbstractMatrix,
                        Diag::DiagnosticVars,
                        S::ModelSetup,
                        t::Int)

    if S.parameters.surface_relax
        continuity_surf_relax!(η,Diag,S,t)
    elseif S.parameters.surface_forcing
        continuity_forcing!(Diag,S,t)
    else
        continuity_itself!(Diag,S,t)
    end
end


function simulate(m::Translation, noise::Distribution;
    duration::Number = 6, Δt::Number = 0.01)

    start, stop = m.start, m.stop
    stop <= duration || error("Movement extends beyond simulation event.")
    m.D ∈ ["forwards", "backwards"] || error("Direction not specified.")

    D = m.D == "forwards" ? 1 : -1
    f = 1/(stop - start) # frequency: single sinusiodal head turn

    timesteps = range(0, stop=duration, step=Δt)
    T = length(timesteps)

    # κ1, κ2 = cupula_dynamics(4, Δt)
    α = Array{Float64}(undef, T) # acceleration
    ω = Array{Float64}(undef, T) # velocity
    θ = Array{Float64}(undef, T) # position
    # c = Array{Float64}(undef, T) # cupula dynamics
    y = Array{Float64}(undef, T) # sensory observations

    # α[1]
    ω[1] = 0
    # c[1] = 0
    θ[1] = 0
    y[1] = α[1] + rand(noise)

    for (i, t) ∈ enumerate(timesteps)
        α[i] = t > start && t < stop ? acceleration(D, m.A, f, t, start) : 0

        if i > 1
            ω[i] = ω[i-1] + Δt * α[i-1]
            # c[i] = κ1 * c[i-1] + κ2 * ω[i]
            θ[i] = θ[i-1] + Δt * ω[i-1] + 0.5*Δt^2 * α[i-1]
        end
        y[i] = α[i] + rand(noise)
    end
    return collect(timesteps), y, α, ω, θ, Δt
end

function simulate(m1::HeadTurn, m2::ExternalDisturbance,
    noise::Distribution; duration::Number = 6,
    Δt::Number = 0.01)

    # start, stop = m1.times
    m1.stop <= duration || error("Movement extends beyond simulation event.")
    m2.stop < duration || error("Movement extends beyond simulation event.")
    m1.D ∈ ["left", "right"] || error("Head turn direction not specified.")
    m2.D ∈ ["left", "right"] || error("Disturbance direction not specified.")

    Dh = m1.D == "left" ? 1 : -1
    Dd = m2.D == "left" ? 1 : -1

    f = 1/(m1.stop - m1.start) # frequency: single sinusiodal head turn

    timesteps = range(0, stop=duration, step=Δt)
    T = length(timesteps)

    κ1, κ2 = cupula_dynamics(4, Δt)

    ωᵤ = Array{Float64}(undef, T) # active (controlled)
    ωₑ = Array{Float64}(undef, T) # passive (external)
    ω = Array{Float64}(undef, T) # total velocity
    θ = Array{Float64}(undef, T) # position
    c = Array{Float64}(undef, T) # cupula dynamics
    y = Array{Float64}(undef, T) # sensory observations

    ωᵤ[1] = 0
    ωₑ[1] = 0
    ω[1] = ωᵤ[1] + ωₑ[1]
    θ[1] = 0
    y[1] = ω[1] - c[1] + rand(noise)

    for (i, t) ∈ enumerate(timesteps)
        ωᵤ[i] = t > m1.start && t < m1.stop ? acceleration(Dh, m1.A, f, t, m1.start) : 0
        ωₑ[i] = t > m2.start && t < m2.stop ? Dd * m2.A : 0
        ω[i] = ωᵤ[i] + ωₑ[i]
        if i > 1
            c[i] = κ1 * c[i-1] + κ2 * ω[i]
            θ[i] = θ[i-1] + Δt * ω[i-1]
        end
        y[i] = ω[i] - c[i] + rand(noise)
    end

    p = scatter(timesteps, y,
        # marker = (:circle, 1, 0.5, :black,
        #     stroke(0.5, 0.5, :black)))
            markershape = :circle,
            markersize = 2,
            markeralpha = 0.4,
            markercolor = :black)
    # plot!(p, time, c,
    #     linewidth = 4,
    #     linealpha = 0.5,
    #     linecolor = :steelblue)
    plot!(p, timesteps, ωₑ,
        linewidth = 6,
        linecolor = :orange,
        xlabel = "Time (s)",
        ylabel = "Angular displacement",
        legend = :none)
    plot!(p, timesteps, ωᵤ,
        linewidth = 6,
        linestyle = :dot,
        linecolor = :steelblue1)
    plot!(p, timesteps, ω,
        linewidth = 6,
        linecolor = :steelblue4)

    plot!(p, timesteps, θ,
        linewidth = 6,
        linestyle = :dot,
        linecolor = :purple)

    return collect(timesteps), y, ω, ωᵤ, ωₑ, θ, c, Δt
end


@testset "Operators" begin
    @testset "Coulomb operator" begin
        k = 7
        N = 31

        a,b = 0,70

        coulomb(r) = -1/r
        @testset "$name knot set" for (name, t, tol) in [("Linear", LinearKnotSet(k, a, b, N), 9e-3),
                                                         ("Exponential", ExpKnotSet(k, -1.0, log10(b), N), 2e-7)]
            B = BSpline(t)[:,2:end-1]
            S = B'B

            r = axes(B,1)

            f = B \ (r -> r^2 * exp(-r)).(r)
            g = B \ (r -> -r*exp(-r)).(r)

            V = B'*QuasiDiagonal(coulomb.(r))*B
            g̃ = S \ V*f

            @test g ≈ g̃ atol=tol
        end
    end
end


using Test, QSimulator

using LinearAlgebra: I
@testset "Primitives" begin
    r = Resonator("test", 3, QSimulator.ResonatorSpec(5.5))

    # test some trivial cases of the ladder operators
    @test raising(r) == [0 0 0;
                         1 0 0;
                         0 √2 0]
    @test lowering(r) == [0 1 0;
                          0 0 √2;
                          0 0 0]

    r = Resonator("test", 7, QSimulator.ResonatorSpec(5.5))

    # test harmonic oscillator operator identity a†a = N
    @test raising(r) * lowering(r) ≈ number(r)

    # test X = a + a†
    @test X(r) ≈ raising(r) + lowering(r)

    # test phase rotation of 0.25 on X → Y
    @test X(r, 0.25) ≈ Y(r)

    # test commutator relation [a, a†] = 1; because of truncation it doesn't hold for last row/columns
    commutator = lowering(r)*raising(r)- raising(r)*lowering(r)
    @test commutator[1:end-1,1:end-1] ≈ Matrix{eltype(commutator)}(I, dimension(r)-1, dimension(r)-1)

    # test trivial case of the `flip_flop` operator
    a = Resonator("a", 2, QSimulator.ResonatorSpec(5.5))
    b = Resonator("b", 2, QSimulator.ResonatorSpec(5.5))
    @test flip_flop(a,b) == [0 0 0 0;
                             0 0 1 0;
                             0 1 0 0;
                             0 0 0 0]

end


export AkersNicholls_fast,AkersNicholls

"""
    AkersNicholls_fast(param;kwargs)

Same as `AkersNicholls`, but faster.
"""
mutable struct AkersNicholls_fast <: AbstractModel

	label   :: String
	f!		:: Function
	mapto	:: Function
	mapfro	:: Function
	info    :: String

    function AkersNicholls_fast( param::NamedTuple;
									mesh = Mesh(param),
									IL	    = false,
									dealias=false,
									label="Akers-Nicholls" )

		# Set up
		μ 	= param.μ
		ϵ 	= param.ϵ
		if !in(:ν,keys(param))
			if μ > 1
				ν = 1/sqrt(μ)
				nu = "1/√μ (deep water case)"
			else
				ν = 1
				nu = "1 (shallow water case)"
			end
		else
			ν = param.ν
			nu = "$ν"
		end
		if μ == Inf || ν==0 || IL == true # infinite layer case
			IL = true;  # IL (=Infinite layer) is a flag to be used thereafter
			μ = 1; ν = 1; # Then we should set μ=ν=1 in subsequent formula.
		end

		# Print information
		info = "Akers-Nicholls model.\n"
		if IL == true
			info *= "├─Steepness parameter ϵ=$ϵ (infinite depth case).\n"
		else
			info *= "├─Shallowness parameter μ=$μ, nonlinearity parameter ϵ=$ϵ, \
					scaling parameter ν=$nu.\n"
		end
		if dealias == true || dealias == 1
			info *= "└─Dealiasing with Orszag’s 3/2 rule. "
		else
			info *= "└─No dealiasing. "
		end
		info *= "\nDiscretized with $(mesh.N) collocation points on [$(mesh.xmin), $(mesh.xmax)]."

		# Pre-allocate useful data
		x = mesh.x
		k = mesh.k
		if IL == true
			Tμ 	= -1im * sign.(k)
			Lμ 	= abs.(k);Lμ[1]=1;
		else
			Tμ 	= -1im* sign.(k) .* tanh.(sqrt(μ)*abs.(k))
			Lμ 	= sqrt(μ)*abs.(k)./tanh.(sqrt(μ)*abs.(k));Lμ[1]=1;
		end
		∂ₓ	=  1im * sqrt(μ)* k            # Differentiation
		Tμ∂ₓ= Tμ .* ∂ₓ
		if dealias == true || dealias == 1
			Π⅔    = abs.(k) .< (mesh.kmax-mesh.kmin)/3 	# Dealiasing low-pass filter
		else
			Π⅔    = zero(k) .+ 1     		# No dealiasing (Π⅔=Id)
		end

        ζ = zeros(Complex{Float64}, mesh.N)
        unew = zeros(Complex{Float64}, mesh.N)

        I₁ = zeros(Complex{Float64}, mesh.N)
        I₂ = zeros(Complex{Float64}, mesh.N)
        I₃ = zeros(Complex{Float64}, mesh.N)

        Px  = plan_fft(ζ; flags = FFTW.MEASURE)

		# Evolution equations are ∂t U = f(U)
		function f!(U)

		    ldiv!(ζ, Px , view(U,:,1))

		    I₁  .= view(U,:,2) .* Lμ
		    ldiv!(unew, Px , I₁)
		    unew .^= 2
		    mul!(I₁, Px , unew)
		    I₁ .*= Tμ

		    I₂  .= view(U,:,1) .* ∂ₓ
		    ldiv!(unew, Px , I₂)
		    unew .*= ζ
		    mul!(I₂, Px , unew)

		    I₃  .= view(U,:,1) .* Tμ∂ₓ
		    ldiv!(unew, Px, I₃)
		    unew .*= ζ
		    mul!(I₃ , Px , unew)
		    I₃ .*= Tμ

		    ζ  .= .- view(U,:,2) .* ∂ₓ

		    I₁ .*= ν
			I₁ .-= I₂
		    I₁ .-= I₃
		    I₁ .*= Π⅔
		    I₁ .*= ϵ

		    U[:,2] .= (view(U,:,1) .* Tμ .+ I₁)/sqrt(μ)/ν
		    U[:,1] .= ζ/sqrt(μ)

		end

		# Build raw data from physical data.
		function mapto(data::InitialData)
			ζ.=data.η(x);v=data.v(x);
			fftv=fft(v);
			[Π⅔ .* fft(ζ)  Π⅔ .* (fftv./Lμ+ϵ*fft(ζ.*v)+ ϵ*Tμ.*fft(ζ.*ifft(Tμ.*fftv)))/ν]
		end

		# Reconstruct physical variables from raw data
		# Return `(η,v,x)`, where
		# - `η` is the surface deformation;
		# - `v` is the derivative of the trace of the velocity potential;
		# - `x` is the vector of collocation points
		function mapfro(U;n=10)
			ζ.=ifft(view(U,:,1));I₁.=ν*view(U,:,2);I₂.=Lμ.*I₁;
			for j=1:n
				I₂.=Lμ.*(I₁-ϵ*Π⅔ .* ( fft(ζ.*ifft(I₂)) + Tμ.*fft(ζ.*ifft(Tμ.*I₂))) )
			end
			real(ζ),real(ifft(I₂)),mesh.x
		end

		new(label, f!, mapto, mapfro, info )

    end
end

"""
    AkersNicholls(param;kwargs)

Define an object of type `AbstractModel` in view of solving the initial-value problem for
the quadratic deep-water model proposed by [Akers and Nicholls](https://doi.org/10.1137/090771351)
and [Cheng, Granero-Belinchón, Shkoller and Milewski](https://doi.org/10.1007/s42286-019-00005-w)

# Argument
`param` is of type `NamedTuple` and must contain
- dimensionless parameters `ϵ` (nonlinearity) and `μ` (dispersion);
- optionally, `ν` the shallow/deep water scaling factor. By default, `ν=1` if `μ≦1` and `ν=1/√μ` otherwise. Set the infinite-layer case if `ν=0`, or `μ=Inf`.
- numerical parameters to construct the mesh of collocation points, if `mesh` is not provided as a keyword argument.

## Optional keyword arguments
- `mesh`: the mesh of collocation points. By default, `mesh = Mesh(param)`;
- `IL`: Set the infinite-layer case if `IL=true` (or `μ=Inf`, or `ν=0`), in which case `ϵ` is the steepness parameter. Default is `false`.
- `dealias`: dealiasing with `1/3` Orlicz rule if `true` or no dealiasing if `false` (by default);
- `label`: a label for future references (default is `"deep quadratic"`);

# Return values
Generate necessary ingredients for solving an initial-value problem via `solve!`:
1. a function `AkersNicholls.f!` to be called in explicit time-integration solvers;
2. a function `AkersNicholls.mapto` which from `(η,v)` of type `InitialData` provides the raw data matrix on which computations are to be executed;
3. a function `AkersNicholls.mapfro` which from such data matrix returns the Tuple of real vectors `(η,v,x)`, where
    - `η` is the values of surface deformation at collocation points `x`;
    - `v` is the derivative of the trace of the velocity potential at `x`.

"""
mutable struct AkersNicholls <: AbstractModel

	label   :: String
	f!		:: Function
	mapto	:: Function
	mapfro	:: Function
	info 	:: String

    function AkersNicholls( param::NamedTuple;
							mesh = Mesh(param),
							IL	    = false,
							dealias = false,
							label="Akers-Nicholls")

		# Set up
		μ 	= param.μ
		ϵ 	= param.ϵ
		if !in(:ν,keys(param))
			if μ > 1
				ν = 1/sqrt(μ)
				nu = "1/√μ (deep water case)"
			else
				ν = 1
				nu = "1 (shallow water case)"
			end
		else
			ν = param.ν
			nu = "$ν"
		end
		if μ == Inf || ν==0 || IL == true # infinite layer case
			IL = true;  # IL (=Infinite layer) is a flag to be used thereafter
			μ = 1; ν = 1; # Then we should set μ=ν=1 in subsequent formula.
		end

		# Print information
		info = "Akers-Nicholls model.\n"
		if IL == true
			info *= "├─Steepness parameter ϵ=$ϵ (infinite depth case).\n"
		else
			info *= "├─Shallowness parameter μ=$μ, nonlinearity parameter ϵ=$ϵ, \
					scaling parameter ν=$nu.\n"
		end
		if dealias == true || dealias == 1
			info *= "└─Dealiasing with Orszag’s 3/2 rule. "
		else
			info *= "└─No dealiasing. "
		end
		info *= "\nDiscretized with $(mesh.N) collocation points on [$(mesh.xmin), $(mesh.xmax)]."

		# Pre-allocate useful data
		x = mesh.x
		k = mesh.k
		if IL == true
			Tμ 	= -1im * sign.(k)
			Lμ 	= abs.(k);Lμ[1]=1;
		else
			Tμ 	= -1im* sign.(k) .* tanh.(sqrt(μ)*abs.(k))
			Lμ 	= sqrt(μ)*abs.(k)./tanh.(sqrt(μ)*abs.(k));Lμ[1]=1;
		end
		∂ₓ	=  1im * sqrt(μ)* k            # Differentiation
		Tμ∂ₓ= Tμ .* ∂ₓ

		if dealias == true || dealias == 1
			Π⅔    = abs.(k) .< (mesh.kmax-mesh.kmin)/3 	# Dealiasing low-pass filter
		else
			Π⅔    = zero(k) .+ 1     		# No dealiasing (Π⅔=Id)
		end

        ζ = zeros(Complex{Float64}, mesh.N)
        unew = zeros(Complex{Float64}, mesh.N)

        I₁ = zeros(Complex{Float64}, mesh.N)
        I₂ = zeros(Complex{Float64}, mesh.N)
        I₃ = zeros(Complex{Float64}, mesh.N)

		# Evolution equations are ∂t U = f(U)
		function f!(U)

			ζ .=ifft(U[:,1]);
			I₁ .=Tμ.*fft(ifft(Lμ.*U[:,2]).^2);
			# Tricomi identity : the above equals the two commented lines below
			# I₁ .=Tμ.*fft(ifft(Lμ.*U[:,2]).^2 .+ ifft(∂ₓ.*U[:,2]).^2)/2;
			# I₁ .-=fft(ifft(Lμ.*U[:,2]).*ifft(∂ₓ.*U[:,2]));

			I₂ .=fft(ζ.*ifft(∂ₓ.*U[:,1])) ;
			I₃ .=Tμ.*fft(ζ.*ifft(Tμ∂ₓ.*U[:,1]));
			ζ .= U[:,1] ;
			U[:,1] .= -∂ₓ.*U[:,2]/sqrt(μ) ;
			U[:,2] .= (Tμ.*ζ+ϵ*Π⅔.*(ν*I₁-I₂-I₃))/sqrt(μ)/ν ;

		end

		# Build raw data from physical data.
		function mapto(data::InitialData)
			ζ.=data.η(x);v=data.v(x);
			fftv=fft(v);
			[Π⅔ .* fft(ζ)  Π⅔ .* (fftv./Lμ+ϵ*fft(ζ.*v)+ ϵ*Tμ.*fft(ζ.*ifft(Tμ.*fftv)))/ν]
		end

		# Reconstruct physical variables from raw data
		# Return `(η,v,x)`, where
		# - `η` is the surface deformation;
		# - `v` is the derivative of the trace of the velocity potential;
		# - `x` is the vector of collocation points
		function mapfro(U;n=10)
			ζ.=ifft(U[:,1]);I₁.=ν*U[:,2];I₂.=Lμ.*I₁;
			for j=1:n
				I₂.=Lμ.*(I₁-ϵ*Π⅔ .* ( fft(ζ.*ifft(I₂)) + Tμ.*fft(ζ.*ifft(Tμ.*I₂))) )
			end
			real(ζ),real(ifft(I₂)),mesh.x
		end

		new(label, f!, mapto, mapfro, info )

    end
end


using Flux
using Plots
using DifferentialEquations

ω0 = 0.5
γ = 0.1
x₀ = 10.0
v₀ = 0.0
tspan = (0.0, 20.0)
acceleration(dx, x, p, t) = -ω0^2*x-2*γ*dx
prob = SecondOrderODEProblem(acceleration, v₀, x₀, tspan)
sol = solve(prob)

plot(sol,label=["Velocity" "Position"])

# generate the training dataset
plot_t = 0:0.01:10
data_plot = sol(plot_t)
positions_plot = [state[2] for state in data_plot]
force_plot = [force(state[1],state[2],k,t) for state in data_plot]

t = 0:3.3:10
dataset = sol(t)
position_data = [state[2] for state in sol(t)]
force_data = [force(state[1],state[2],k,t) for state in sol(t)]

plot(plot_t,force_plot,xlabel="t",label="True Force")
scatter!(t,force_data,label="Force Measurements")

# create NN
# create NN
NNForce = Chain(x -> [x],
                Dense(1,32,tanh),
                Dense(32,1),
                first)

loss() = sum(abs2,NNForce(position_data[i]) - force_data[i] for i in 1:length(position_data))
loss()

opt = Flux.Descent(0.01)
data = Iterators.repeated((), 5000)
iter = 0
cb = function () #callback function to observe training
    global iter += 1
    if iter % 500 == 0
        display(loss())
    end
end
display(loss())

# train the NN
Flux.train!(loss, Flux.params(NNForce), data, opt; cb=cb)

# visualize the result
learned_force_plot = NNForce.(positions_plot)
plot(plot_t,force_plot,xlabel="t",label="True Force")
plot!(plot_t,learned_force_plot,label="Predicted Force")
scatter!(t,force_data,label="Force Measurements")

# solve SHO problem to justify PINN
acceleration2(dx,x,k,t) = -ω0^2*x
prob_simplified = SecondOrderODEProblem(acceleration2, v₀, x₀, tspan)
sol_simplified = solve(prob_simplified)
p2 = plot(sol_simplified, label=["Velocity" "Position"])
plot!(p2, sol_simplified,label=["Velocity Simplified" "Position Simplified"])

# create Physics Loss Term to force solution to satisfy SHO ode as well
random_positions = [2rand()-1 for i in 1:100] # random values in [-1,1]
loss_ode() = sum(abs2,NNForce(x) - (-(ω0^2)*x) for x in random_positions)
loss_ode()

λ = 0.025  # factor controlling physics regularization
composed_loss() = loss() + λ*loss_ode()

opt = Flux.Descent(0.01)
data = Iterators.repeated((), 5000)
iter = 0
cb = function () #callback function to observe training
    global iter += 1
    if iter % 500 == 0
        display(composed_loss())
    end
end
display(composed_loss())

# train the new model
Flux.train!(composed_loss, Flux.params(NNForce), data, opt; cb=cb)

loss()
λ * loss_ode()

# visualize the final result!
learned_force_plot = NNForce.(positions_plot)
plot(plot_t,force_plot,xlabel="t",label="True Force")
plot!(plot_t,learned_force_plot,label="Predicted Force")
scatter!(t,force_data,label="Force Measurements")


using Plots
using ImageMagick
using FileIO
using Images
using ImageView
using FFTW

function convert_to_float(a::RGB)
    a = float(a)
    return (a.r + a.g + a.b)/3.0
end

function convert_to_float(a::RGBA)
    a = float(a)
    return (a.r + a.g + a.b + a.alpha)*0.25
end

function convert_to_rgba(r, i)
    return RGBA(r, 0, i, 1)
end

function normalize!(a)
    factor = 1/(sum(abs2.(a)))
    a .*= factor
end

function square_animation()
    x = zeros(100)
    x[1:5] .= 1
    x[end-5:end] .= 1
    for i = 1:length(x)-11
        x[i+5] = 1
        x[i] = 0
        if i < 6
            x[end-6+i] = 0
        end
        filename = "output/time"*lpad(i, 5, "0")*".png"
        time_plot = plot(x)
        savefig(time_plot, filename)

        y = fftshift(fft(x))
        filename = "output/frequency"*lpad(i, 5, "0")*".png"
        freq_plot = plot(real(y))
        freq_plot = plot!(freq_plot, imag(y))
        savefig(freq_plot, filename)

    end
end

function fourier_challenge(input_filename::String, output_filename::String)
    img = load(input_filename)
    img_real = float(red.(img))
    img_imag = float(blue.(img))

    array = img_real + im*img_imag
    array = fftshift(ifft(array))
    #array = ifft(array)
    array = abs2.(array)
    normalize!(array)
    img_final = convert_to_rgba.(real(array), imag(array))
    save(output_filename, map(clamp01nan, img_final))
end


# This file calculate the whole aerodynamic force of the rotor in the hub

function beforce(vall_r,chord,Cl,Cd,α_aero,θ,dr,beta=0.0)
  # vall_r = Array{Vector}(NR,Nbe)
  # for k in 1:NR # 将总速度由旋转坐标系转化到叶素坐标系
  #   for i in 1:Nbe
  #     vall_r[k,i] = rotobe(vall_r[k,i],beta,θ[k,i])
  #   end
  # end
  L = abs(1/2*ρ*norm(vall_r)^2*chord*Cl)
  D = abs(1/2*ρ*norm(vall_r)^2*chord*Cd)
  if α_aero>=0
    L_be = [0,cos(α_aero+π/2),-sin(α_aero+π/2)] # 叶素当地升力单位矢量
  else
    L_be = [0,cos(α_aero-π/2),-sin(α_aero-π/2)]
  end
  D_be = [0,cos(α_aero),-sin(α_aero)] # 叶素当地阻力单位矢量
  L_ro = betoro(L_be,beta,θ) # 旋转坐标系叶素升力单位矢量
  D_ro = betoro(D_be,beta,θ) # 旋转坐标系叶素阻力单位矢量
  Fbez = (L*L_ro[3]+D*D_ro[3])*dr # 旋转坐标系叶素z向力
  Fbey = (L*L_ro[2]+D*D_ro[2])*dr # 旋转坐标系叶素y向力
  return Fbez,Fbey
end

function rotorforce(ψ,vall_r,α_aero,θ,Cl,Cd,filename=pwd()*"\\force_record\\fhub.txt")
  # filebf = open(filename,"w")
  # write(filebf,"List  ZForce  YForce\n")

  Fz_r = 0.0 # 旋转坐标系轴向力
  Fy_r = 0.0 # 旋转坐标系y向力
  MQ = 0.0 # 扭矩
  Mβ_aero = zeros(Float64,NR) # 初始化气动力挥舞力矩
  for k in 1:NR
    # write(filebf,"$(k)\t\n")
    for i in 1:Nbe
      fbe = beforce(vall_r[k,i],ch[k,i],Cl[k,i],Cd[k,i],
                      α_aero[k,i],θ[k,i],dr[k,i])
      fbez = fbe[1]
      fbey = fbe[2]
      MQ += fbey*rb[k,i]
      Mβ_aero[k] += fbez*(rb[k,i]-ecut*R)
      Fz_r += fbez
      Fy_r += fbey
      # write(filebf,"$(i)  $(Fzr)  $(Fyr)\n")
    end
  end

  # write(filebf,"The Whole Force of Z is : $(Fzr)\n")
  # write(filebf,"The Whole Force of Y is : $(Fyr)\n")
  #
  # close(filebf)
  Fy_s = Fy_r*cos(ψ)
  Fx_s = -Fy_r*sin(ψ)
  Fz_s = Fz_r

  return Fx_s,Fy_s,Fz_s,MQ,Mβ_aero
end


function fib(n::Int)
  if n < 2
    return n
  end
  return fib(n - 2) + fib(n - 1)
end

println(fib(35))


function test_metric(metric, enc, cm, targets, predicts, scores, thres, val_true; kwargs...)
    @testset "$(metric)" begin
        @test metric(cm; kwargs...) == val_true
        @test metric([cm]; kwargs...) == [val_true]
        @test metric(targets, predicts; kwargs...) == val_true
        @test metric(enc, targets, predicts; kwargs...) == val_true
        @test metric(targets, scores, thres; kwargs...) == val_true
        @test metric(enc, targets, scores, thres; kwargs...) == val_true
        @test metric(targets, scores, [thres]; kwargs...) == [val_true]
        @test metric(enc, targets, scores, [thres]; kwargs...) == [val_true]
    end
end

n = 1000
targets = rand(0:1, n)
scores = rand(n)
thres = 0.7
predicts = scores .>= thres

βs = [1,2,3,4]

p  = sum(targets .== 1)
n  = sum(targets .== 0)
tp = sum(targets .* predicts)
tn = sum((1 .- targets) .* (1 .- predicts))
fp = sum((1 .- targets) .* predicts)
fn = sum(targets .* (1 .- predicts))

cm = ConfusionMatrix(p,n,tp,tn,fp,fn)

set_encoding(OneZero())

encs = [
    OneZero(),
    OneMinusOne(),
    OneTwo(),
    OneVsOne(3,4),
    OneVsOne(:three,:four),
    OneVsOne("three","four"),
    OneVsRest(1, [2,3,4]),
    OneVsRest(:one, [:two, :three]),
    OneVsRest("one", ["two", "three"]),
    RestVsOne([1,2,3], 4),
    RestVsOne([:one, :two], :three),
    RestVsOne(["one", "two"], "three")
]

@testset "Metrics for $enc encoding" for enc in encs
    global targets, predicts

    targets = recode.(current_encoding(), enc, targets)
    predicts = recode.(current_encoding(), enc, predicts)
    set_encoding(enc)

    # true positie
    test_metric(true_positive, enc, cm, targets, predicts, scores, thres, tp)

    # true negtives
    test_metric(true_negative, enc, cm, targets, predicts, scores, thres, tn)

    # false positive
    test_metric(false_positive, enc, cm, targets, predicts, scores, thres, fp)

    # false negative
    test_metric(false_negative, enc, cm, targets, predicts, scores, thres, fn)

    # true positie rate
    val_true = tp/p
    test_metric(true_positive_rate, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(sensitivity, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(recall, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(hit_rate, enc, cm, targets, predicts, scores, thres, val_true)

    # true negtives
    val_true = tn/n
    test_metric(true_negative_rate, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(specificity, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(selectivity, enc, cm, targets, predicts, scores, thres, val_true)

    # false positive
    val_true = fp/n
    test_metric(false_positive_rate, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(fall_out, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(type_I_error, enc, cm, targets, predicts, scores, thres, val_true)

    # false negative
    val_true = fn/p
    test_metric(false_negative_rate, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(miss_rate, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(type_II_error, enc, cm, targets, predicts, scores, thres, val_true)

    # precision
    val_true = tp/(tp + fp)
    if isnan(val_true)
        val_true = one(val_true)
    end
    test_metric(precision, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(positive_predictive_value, enc, cm, targets, predicts, scores, thres, val_true)

    # negative predictive value
    val_true = tn/(tn + fn)
    test_metric(negative_predictive_value, enc, cm, targets, predicts, scores, thres, val_true)

    # false discovery rate
    val_true = fp/(fp + tp)
    test_metric(false_discovery_rate, enc, cm, targets, predicts, scores, thres, val_true)

    # false omission rate
     val_true = fn/(fn + tn)
    test_metric(false_omission_rate, enc, cm, targets, predicts, scores, thres, val_true)

    # threat score
    val_true = tp/(tp + fn + fp)
    test_metric(threat_score, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(critical_success_index, enc, cm, targets, predicts, scores, thres, val_true)

    # accuracy
    val_true = (tp + tn)/(p + n)
    test_metric(accuracy, enc, cm, targets, predicts, scores, thres, val_true)

    # balanced accuracy
    val_true = (tp/p + tn/n)/2
    test_metric(balanced_accuracy, enc, cm, targets, predicts, scores, thres, val_true)

    # error rate
    val_true = 1 - (tp + tn)/(p + n)
    test_metric(error_rate, enc, cm, targets, predicts, scores, thres, val_true)

    # balanced accuracy
    val_true = 1 - (tp/p + tn/n)/2
    test_metric(balanced_error_rate, enc, cm, targets, predicts, scores, thres, val_true)

    # f1 score
    prec = tp/(tp + fp)
    if isnan(prec)
        prec = one(prec)
    end
    rec = tp/p
    val_true = 2*prec*rec/(prec + rec)
    test_metric(f1_score, enc, cm, targets, predicts, scores, thres, val_true)

    # fβ score
    @testset "fβ score: β = $(β)" for β in βs
        val_true = (1 + β^2)*prec*rec/(β^2*prec + rec)
        test_metric(fβ_score, enc, cm, targets, predicts, scores, thres, val_true; β = β)
    end

    # Matthews correlation coefficient
    val_true = (tp*tn - fp*fn)/sqrt((tp + fp)*(tp + fn)*(tn + fp)*(tn + fn))
    test_metric(matthews_correlation_coefficient, enc, cm, targets, predicts, scores, thres, val_true)
    test_metric(mcc, enc, cm, targets, predicts, scores, thres, val_true)

    # quant
    val_true = (fn + tn)/(p + n)
    test_metric(quant, enc, cm, targets, predicts, scores, thres, val_true)

    # topquant
    val_true = 1 - (fn + tn)/(p + n)
    test_metric(topquant, enc, cm, targets, predicts, scores, thres, val_true)

    # positive likelihood ratio
    val_true = tp/p/(fp/n)
    test_metric(positive_likelihood_ratio, enc, cm, targets, predicts, scores, thres, val_true)

    # negative likelihood ratio
    val_true = fn/p/(tn/n)
    test_metric(negative_likelihood_ratio, enc, cm, targets, predicts, scores, thres, val_true)

    # diagnostic odds ratio
    val_true = (tp/p)*(tn/n)/((fp/n)*(fn/p))
    test_metric(diagnostic_odds_ratio, enc, cm, targets, predicts, scores, thres, val_true)

    # prevalence
    val_true = p/(p + n)
    test_metric(prevalence, enc, cm, targets, predicts, scores, thres, val_true)
end


using Unitful
using DataStructures

t = range(0.0, stop=10.0, length=100)

result = OrderedDict{String,Any}()

result["time"] = t*u"s"
result["phi"]  = sin.(t)u"rad"
result["w"]    = cos.(t)u"rad/s"

plot(result, ("phi", "w"), prefix="Sim 1:", heading="Test reuse")

result["phi"]  = 1.2*sin.(t)u"rad"
result["w"]    = 0.8*cos.(t)u"rad/s"

plot(result, ("phi", "w"), prefix="Sim 2:", reuse=true)

println("\n... result info:")
printResultInfo(result)


import LazySets: dim, inner
using LazySets: isapproxzero

"""
    AbstractDirections{N, VN}

Abstract type for template direction representations.

### Notes

This type is parameterzed by `N` and `VN`, where:

- `N` stands for the numeric type
- `VN` stands for the vector type with coefficients of type `N`

Each subtype is an iterator over a set of prescribed directions.

All subtypes should implement the standard iterator methods from `Base`, namely
`Base.length` (returns the number of directions in the template), and `Base.iterate`.
Moreover, the following methods should be implemented:

- `dim`    -- return the ambient dimension of the template
- `eltype` -- return the type of each vector in the template

Optionally, subtypes may implement:

- `isbounding`   -- (defaults to `false`) return `true` if an overapproximation with
                    a list of template directions results in a bounded set, given a
                    bounded input set, and `false` otherwise
- `isnormalized` -- (defaults to `false`) returns `true` if each direction in the
                    given template has norm one w.r.t. the usual vector 2-norm
"""
abstract type AbstractDirections{N, VN} end

"""
    dim(ad::AbstractDirections)

Returns the dimension of the generated directions.

### Input

- `ad` -- template directions

### Output

The dimension of the generated directions.
"""
function dim(ad::AbstractDirections) end

"""
    isbounding(ad::Type{<:AbstractDirections})

Checks if an overapproximation with a list of template directions results in a
bounded set, given a bounded input set.

### Input

- `ad` -- template directions

### Output

Given a bounded set ``X``, we can construct an outer approximation of ``X`` by
using the template directions `ad` as normal vectors of the facets.
If this function returns `true`, then the result is again a bounded set (i.e., a
polytope).
Note that the result does not depend on the specific shape of ``X``, as long as
``X`` is bounded.

### Notes

By default, this function returns `false` in order to be conservative.
Custom subtypes of `AbstractDirections` should hence add a method for this
function.
"""
function isbounding(ad::Type{<:AbstractDirections})
    return false
end

isbounding(::AD) where {AD<:AbstractDirections} = isbounding(AD)

"""
    isnormalized(ad::Type{<:AbstractDirections})

Returns whether the given template directions is normalized with respect to the
2-norm.

### Input

- `ad` -- template directions

### Output

`true` if the 2-norm of each element in `ad` is one and `false` otherwise.
"""
function isnormalized(ad::Type{<:AbstractDirections})
    return false
end

isnormalized(::AD) where {AD<:AbstractDirections} = isnormalized(AD)

# ==================================================
# Box directions
# ==================================================

"""
    BoxDirections{N, VN} <: AbstractDirections{N, VN}

Box directions representation.

### Fields

- `n` -- dimension

### Notes

Box directions can be seen as the vectors where only one entry is ±1, and all
other entries are 0. In dimension ``n``, there are ``2n`` such directions.

The default vector representation used in this template is a
`LazySets.Arrays.SingleEntryVector`, although other implementations can be used
such as a regular `Vector` and a sparse vector, `SparseVector`.

### Examples

The template can be constructed by passing the dimension. For example, in dimension
two,

```jldoctest dirs_Box
julia> dirs = BoxDirections(2)
BoxDirections{Float64,LazySets.Arrays.SingleEntryVector{Float64}}(2)

julia> length(dirs)
4
```

By default, each direction is represented in this iterator as a `SingleEntryVector`,
i.e. a vector with only one non-zero element,

```jldoctest dirs_Box
julia> eltype(dirs)
LazySets.Arrays.SingleEntryVector{Float64}
```
In two dimensions, the directions defined by `BoxDirections` are normal to the
facets of a box.

```jldoctest dirs_Box
julia> collect(dirs)
4-element Array{LazySets.Arrays.SingleEntryVector{Float64},1}:
 [1.0, 0.0]
 [0.0, 1.0]
 [0.0, -1.0]
 [-1.0, 0.0]
```

The numeric type can be specified as well:

```jldoctest
julia> BoxDirections{Rational{Int}}(10)
BoxDirections{Rational{Int64},LazySets.Arrays.SingleEntryVector{Rational{Int64}}}(10)

julia> length(ans)
20
```
"""
struct BoxDirections{N, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}
    n::Int
end

# convenience constructor for type Float64
BoxDirections(n::Int) = BoxDirections{Float64, SingleEntryVector{Float64}}(n)

# constructor where only N is specified
BoxDirections{N}(n::Int) where {N} = BoxDirections{N, SingleEntryVector{N}}(n)

Base.eltype(::Type{BoxDirections{N, VN}}) where {N, VN} = VN
Base.length(bd::BoxDirections) = 2 * bd.n

# interface functions
dim(bd::BoxDirections) = bd.n
isbounding(::Type{<:BoxDirections}) = true
isnormalized(::Type{<:BoxDirections}) = true

# The idea is that positive states run through vectors with +1 entry,
# and negative states run through vectors with -1 entry
# (1, 0)   state = 1
# (0, 1)   state = 2
# (0, -1)  state = -2
# (-1, 0)  state = -1
function Base.iterate(bd::BoxDirections{N, SingleEntryVector{N}}, state::Int=1) where {N}
    if state == 0
        return nothing
    end
    vec = SingleEntryVector(abs(state), bd.n, convert(N, sign(state)))
    state = (state == bd.n) ? -bd.n : state + 1
    return (vec, state)
end

function Base.iterate(bd::BoxDirections{N, Vector{N}}, state::Int=1) where {N}
    if state == 0
        return nothing
    end
    vec = zeros(N, bd.n)
    vec[abs(state)] = convert(N, sign(state))
    state = (state == bd.n) ? -bd.n : state + 1
    return (vec, state)
end

function Base.iterate(bd::BoxDirections{N, SparseVector{N, Int}}, state::Int=1) where {N}
    if state == 0
        return nothing
    end
    vec = sparsevec([abs(state)], convert(N, sign(state)), bd.n)
    state = (state == bd.n) ? -bd.n : state + 1
    return (vec, state)
end

# ==================================================
# Octagonal directions
# ==================================================

"""
    OctDirections{N, VN} <: AbstractDirections{N, VN}

Octagon directions representation.

### Fields

- `n` -- dimension

### Notes

Octagon directions consist of all vectors that are zero almost everywhere except
in two dimensions ``i``, ``j`` (possibly ``i = j``) where it is ``±1``. In dimension
``n``, there are ``2n^2`` such directions.

### Examples

The template can be constructed by passing the dimension. For example, in dimension
two,

```jldoctest dirs_Oct
julia> dirs = OctDirections(2)
OctDirections{Float64,SparseArrays.SparseVector{Float64,Int64}}(2)

julia> length(dirs) # number of directions
8
```
By default, each direction is represented in this iterator as a sparse vector:

```jldoctest dirs_Oct
julia> eltype(dirs)
SparseArrays.SparseVector{Float64,Int64}
```
In two dimensions, the directions defined by `OctDirections` are normal to
the facets of an octagon.

```jldoctest dirs_Oct
julia> first(dirs)
2-element SparseArrays.SparseVector{Float64,Int64} with 2 stored entries:
  [1]  =  1.0
  [2]  =  1.0

julia> Vector.(collect(dirs))
8-element Array{Array{Float64,1},1}:
 [1.0, 1.0]
 [1.0, -1.0]
 [-1.0, 1.0]
 [-1.0, -1.0]
 [1.0, 0.0]
 [0.0, 1.0]
 [0.0, -1.0]
 [-1.0, 0.0]
```

The numeric type can be specified as well:

```jldoctest
julia> OctDirections{Rational{Int}}(10)
OctDirections{Rational{Int64},SparseArrays.SparseVector{Rational{Int64},Int64}}(10)

julia> length(ans)
200
```
"""
struct OctDirections{N, VN} <: AbstractDirections{N, VN}
    n::Int
end

# constructor for type Float64
OctDirections(n::Int) = OctDirections{Float64, SparseVector{Float64, Int}}(n)

# constructor where only N is specified
OctDirections{N}(n::Int) where {N} = OctDirections{N, SparseVector{N, Int}}(n)

Base.eltype(::Type{OctDirections{N, VN}}) where {N, VN} = VN
Base.length(od::OctDirections) = 2 * od.n^2

# interface functions
dim(od::OctDirections) = od.n
isbounding(::Type{<:OctDirections}) = true
isnormalized(::Type{<:OctDirections}) = false

function Base.iterate(od::OctDirections{N, SparseVector{N, Int}}) where {N}
    if od.n == 1
        # fall back to box directions in 1D case
        return iterate(od, 1)
    end
    vec = spzeros(N, od.n)
    vec[1] = one(N)
    vec[2] = one(N)
    return (copy(vec), (vec, 1, 2))
end

# basic idea: modify the vector from the previous iteration
# have two indices i = 1, j = i + 1
# - run j through all indices > i
# - i = i + 1, j = i + 1
# - for any pair (i, j) create four vectors [..., i: ±1, ..., j: ±1, ...]
# in the end continue with box directions
function Base.iterate(od::OctDirections{N, SparseVector{N, Int}}, state::Tuple) where {N}
    # continue with octagon directions
    vec = state[1]
    i = state[2]
    j = state[3]
    if vec[j] > 0
        # change sign at j
        vec[j] = -one(N)
    elseif vec[i] > 0
        # change sign at i and j
        vec[i] = -one(N)
        vec[j] = one(N)
    elseif j < od.n
        # advance j, reset i
        vec[i] = one(N)
        vec[j] = zero(N)
        j = j + 1
        vec[j] = one(N)
    elseif i < od.n - 1
        # advance i
        vec[i] = zero(N)
        vec[j] = zero(N)
        i = i + 1
        j = i + 1
        vec[i] = one(N)
        vec[j] = one(N)
    else
        # continue with box directions
        return iterate(od, 1)
    end
    return (copy(vec), (vec, i, j))
end

function Base.iterate(od::OctDirections{N, SparseVector{N, Int}}, state::Int) where {N}
    # continue with box directions
    return iterate(BoxDirections{N, SparseVector{N, Int}}(od.n), state)
end

# ==================================================
# Box-diagonal directions
# ==================================================

"""
    BoxDiagDirections{N, VN} <: AbstractDirections{N, VN}

Box-diagonal directions representation.

### Fields

- `n` -- dimension

### Notes

Box-diagonal directions can be seen as the union of diagonal directions (all
entries are ±1) and box directions (one entry is ±1, all other entries are 0).
The iterator first enumerates all diagonal directions, and then all box
directions. In dimension ``n``, there are in total ``2^n + 2n`` such directions.

## Examples

The template can be constructed by passing the dimension. For example, in dimension
two,

```jldoctest dirs_BoxDiag
julia> dirs = BoxDiagDirections(2)
BoxDiagDirections{Float64,Array{Float64,1}}(2)

julia> length(dirs) # number of directions
8
```
By default, each direction is represented in this iterator as a regular vector:

```jldoctest dirs_BoxDiag
julia> eltype(dirs)
Array{Float64,1}
```
In two dimensions, the directions defined by `BoxDiagDirections` are normal to
the facets of an octagon.

```jldoctest dirs_BoxDiag
julia> collect(dirs)
8-element Array{Array{Float64,1},1}:
 [1.0, 1.0]
 [-1.0, 1.0]
 [1.0, -1.0]
 [-1.0, -1.0]
 [1.0, 0.0]
 [0.0, 1.0]
 [0.0, -1.0]
 [-1.0, 0.0]
```

The numeric type can be specified as well:

```jldoctest
julia> BoxDiagDirections{Rational{Int}}(10)
BoxDiagDirections{Rational{Int64},Array{Rational{Int64},1}}(10)

julia> length(ans)
1044
```
"""
struct BoxDiagDirections{N, VN} <: AbstractDirections{N, VN}
    n::Int
end

# constructor for type Float64
BoxDiagDirections(n::Int) = BoxDiagDirections{Float64, Vector{Float64}}(n)

# constructor where only N is specified
BoxDiagDirections{N}(n::Int) where {N} = BoxDiagDirections{N, Vector{N}}(n)

Base.eltype(::Type{BoxDiagDirections{N, VN}}) where {N, VN} = VN
Base.length(bdd::BoxDiagDirections) = bdd.n == 1 ? 2 : 2^bdd.n + 2 * bdd.n

# interface function
dim(bdd::BoxDiagDirections) = bdd.n
isbounding(::Type{<:BoxDiagDirections}) = true
isnormalized(::Type{<:BoxDiagDirections}) = false

function Base.iterate(bdd::BoxDiagDirections{N, Vector{N}}) where {N}
    return (ones(N, bdd.n), ones(N, bdd.n))
end

function Base.iterate(bdd::BoxDiagDirections{N}, state::Vector{N}) where {N}
    # continue with diagonal directions
    i = 1
    while i <= bdd.n && state[i] < 0
        state[i] = -state[i]
        i = i+1
    end
    if i > bdd.n
        if bdd.n == 1
            # finish here to avoid duplicates
            return nothing
        else
            # continue with box directions
            return iterate(bdd, 1)
        end
    else
        state[i] = -state[i]
        return (copy(state), state)
    end
end

function Base.iterate(bdd::BoxDiagDirections{N, Vector{N}}, state::Int) where {N}
    # continue with box directions
    return iterate(BoxDirections{N, Vector{N}}(bdd.n), state)
end

# ==================================================
# Polar directions
# ==================================================

"""
    PolarDirections{N<:AbstractFloat, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}

Polar directions representation.

### Fields

- `Nφ` -- length of the partition of the polar angle

### Notes

The `PolarDirections` constructor provides a sample of the unit sphere
in ``\\mathbb{R}^2``, which is parameterized by the polar angle
``φ ∈ Dφ := [0, 2π]``; see the wikipedia entry
[Polar coordinate system](https://en.wikipedia.org/wiki/Polar_coordinate_system)
for details.

The integer argument ``Nφ`` defines how many samples of ``Dφ`` are taken. The
Cartesian components of each direction are obtained with

```math
[cos(φᵢ), sin(φᵢ)].
```

### Examples

The integer passed as an argument is used to discretize ``φ``:

```jldoctest; filter = r"2246[0-9]*e-16"
julia> pd = PolarDirections(2);

julia> pd.stack
2-element Array{Array{Float64,1},1}:
 [1.0, 0.0]
 [-1.0, 1.2246467991473532e-16]

julia> length(pd)
2
```
"""
struct PolarDirections{N<:AbstractFloat, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}
    Nφ::Int
    stack::Vector{VN} # stores the polar directions
end

# convenience constructors
PolarDirections(Nφ::Int) = PolarDirections{Float64, Vector{Float64}}(Nφ)
PolarDirections{N}(Nφ::Int) where {N} = PolarDirections{N, Vector{N}}(Nφ)

function PolarDirections{N, Vector{N}}(Nφ::Int) where {N}
    if Nφ <= 0
        throw(ArgumentError("Nφ = $Nφ is invalid; it shoud be at least 1"))
    end
    stack = Vector{Vector{N}}(undef, Nφ)
    φ = range(N(0), N(2*pi), length=Nφ+1)  # discretization of the polar angle

    @inbounds for i in 1:Nφ  # skip last (repeated) angle
        stack[i] = N[cos(φ[i]), sin(φ[i])]
    end
    return PolarDirections{N, Vector{N}}(Nφ, stack)
end

Base.eltype(::Type{PolarDirections{N, VN}}) where {N, VN} = VN
Base.length(pd::PolarDirections) = pd.Nφ

# interface functions
dim(pd::PolarDirections) = 2
isbounding(pd::PolarDirections) = pd.Nφ > 2
isnormalized(::Type{<:PolarDirections}) = true

function Base.iterate(pd::PolarDirections{N, Vector{N}}, state::Int=1) where {N}
    state == pd.Nφ + 1 && return nothing
    return (pd.stack[state], state + 1)
end

# ==================================================
# Spherical directions
# ==================================================

"""
    SphericalDirections{N<:AbstractFloat, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}

Spherical directions representation.

### Fields

- `Nθ`    -- length of the partition of the azimuthal angle
- `Nφ`    -- length of the partition of the polar angle
- `stack` -- list of computed directions

### Notes

The `SphericalDirections` constructor provides a sample of the unit sphere
in ``\\mathbb{R}^3``, which is parameterized by the azimuthal and polar angles
``θ ∈ Dθ := [0, π]`` and ``φ ∈ Dφ := [0, 2π]`` respectively, see the wikipedia
entry [Spherical coordinate system](https://en.wikipedia.org/wiki/Spherical_coordinate_system)
for details.

The integer arguments ``Nθ`` and ``Nφ`` define how many samples along the domains
``Dθ`` and ``Dφ`` respectively are taken. The Cartesian components of each direction
are obtained with

```math
[sin(θᵢ)*cos(φᵢ), sin(θᵢ)*sin(φᵢ), cos(θᵢ)].
```

The north and south poles are treated separately so that those points
are not considered more than once.

### Examples

A `SphericalDirections` template can be built in different ways. If you pass
only one integer, the same value is used to discretize both ``θ`` and ``φ``:

```jldoctest spherical_directions; filter = r"1232[0-9]*e-17.*2246[0-9]*e-16.*1232[0-9]*e-17"
julia> sd = SphericalDirections(3);

julia> sd.Nθ, sd.Nφ
(3, 3)

julia> length(sd)
4
```

Pass two integers to control the discretization in ``θ`` and in ``φ`` separately:

```jldoctest spherical_directions
julia> sd = SphericalDirections(4, 5);

julia> length(sd)
10

julia> sd = SphericalDirections(4, 8);

julia> length(sd)
16
```
"""
struct SphericalDirections{N<:AbstractFloat, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}
    Nθ::Int
    Nφ::Int
    stack::Vector{VN} # stores the spherical directions
end

# convenience constructors
SphericalDirections(Nθ::Int) = SphericalDirections(Nθ, Nθ)
SphericalDirections(Nθ::Int, Nφ::Int) = SphericalDirections{Float64, Vector{Float64}}(Nθ::Int, Nφ::Int)
SphericalDirections{N}(Nθ::Int, Nφ::Int) where {N} = SphericalDirections{N, Vector{N}}(Nθ::Int, Nφ::Int)

function SphericalDirections{N, Vector{N}}(Nθ::Int, Nφ::Int) where {N}
    if Nθ <= 1 || Nφ <= 1
        throw(ArgumentError("(Nθ, Nφ) = ($Nθ, $Nφ) is invalid; both shoud be at least 2"))
    end
    stack = Vector{Vector{N}}()
    θ = range(N(0), N(pi), length=Nθ)    # discretization of the azimuthal angle
    φ = range(N(0), N(2*pi), length=Nφ)  # discretization of the polar angle

    # add north pole (θ = 0)
    push!(stack, N[0, 0, 1])

    # add south pole (θ = pi)
    push!(stack, N[0, 0, -1])

    for φᵢ in φ[1:Nφ-1]  # delete repeated angle
        for θⱼ in θ[2:Nθ-1] # delete north and south poles
            d = N[sin(θⱼ)*cos(φᵢ), sin(θⱼ)*sin(φᵢ), cos(θⱼ)]
            push!(stack, d)
        end
    end
    return SphericalDirections{N, Vector{N}}(Nθ, Nφ, stack)
end

Base.eltype(::Type{SphericalDirections{N, VN}}) where {N, VN} = VN
Base.length(sd::SphericalDirections) = length(sd.stack)

# interface functions
dim(::SphericalDirections) = 3
isbounding(sd::SphericalDirections) = sd.Nθ > 2 && sd.Nφ > 2
isnormalized(::Type{<:SphericalDirections}) = true

function Base.iterate(sd::SphericalDirections, state::Int=1)
    state == length(sd.stack) + 1 && return nothing
    return (sd.stack[state], state + 1)
end

# ==================================================
# Custom template directions
# ==================================================

"""
    CustomDirections{N, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}

User-defined template directions.

### Fields

- `directions` -- list of template directions
- `n`          -- dimension
- `isbounding` -- boundedness status

### Notes

This struct is a wrapper type for a set of user-defined directions which are
iterated over. It has fields for the list of directions, the set dimension,
and (boolean) cache fields for the boundedness and normalization properties.
The latter are checked by default upon construction.

To check boundedness, we overapproximate the unit ball in the infinity
norm using the given directions and check if the resulting set is bounded.

The dimension will also be determined automatically, unless the empty vector is
passed (in which case the optional argument `n` needs to be specified).

## Examples

Creating a template with box directions in dimension two:

```jldoctest
julia> dirs = CustomDirections([[1.0, 0.0], [-1.0, 0.0], [0.0, 1.0], [0.0, -1.0]]);

julia> dirs.directions
4-element Array{Array{Float64,1},1}:
 [1.0, 0.0]
 [-1.0, 0.0]
 [0.0, 1.0]
 [0.0, -1.0]

julia> LazySets.Approximations.isbounding(dirs)
true

julia> LazySets.Approximations.isnormalized(dirs)
true
```
"""
struct CustomDirections{N, VN<:AbstractVector{N}} <: AbstractDirections{N, VN}
    directions::Vector{VN}
    n::Int
    bounded::Bool
    normalized::Bool
end

function CustomDirections(directions::Vector{VN},
                          n::Int=_determine_dimension(directions);
                          check_boundedness::Bool=true,
                          check_normalization::Bool=true) where {N, VN<:AbstractVector{N}}
    bounded = check_boundedness ? _isbounding(directions) : false
    normalized = check_normalization ? _isnormalized(directions) : false
    return CustomDirections(directions, n, bounded, normalized)
end

function _determine_dimension(directions)
    isempty(directions) && throw(ArgumentError("empty template directions " *
                                               "need a specified dimension"))
    return length(directions[1])
end

function _isbounding(directions::Vector{VN}) where {N, VN<:AbstractVector{N}}
    isempty(directions) && return false

    # check boundedness of the polyhedron `⋂_a ax <= 1` where `a` is a direction
    P = HPolyhedron([HalfSpace(dir, one(N)) for dir in directions])
    return isbounded(P)
end

function _isnormalized(directions::Vector{VN}) where {N, VN<:AbstractVector{N}}
    return all(x -> _isapprox(x, one(N)), norm.(directions, 2))
end

Base.eltype(::Type{CustomDirections{N, VN}}) where {N, VN} = VN
Base.length(cd::CustomDirections) = length(cd.directions)

# interface functions
dim(cd::CustomDirections) = cd.n
isbounding(cd::Type{<:CustomDirections}) = false # it is not a property of the type
isbounding(cd::CustomDirections) = cd.bounded
isnormalized(cd::Type{<:CustomDirections}) = false # it is not a property of the type
isnormalized(cd::CustomDirections) = cd.normalized

function Base.iterate(cd::CustomDirections{N}, state::Int=1) where {N}
    if state > length(cd.directions)
        return nothing
    end
    vec = cd.directions[state]
    state = state + 1
    return (vec, state)
end


@testset "363.max-sum-of-rectangle-no-larger-than-k.jl" begin
    @test max_sum_submatrix(hcat([[1, 0, 1], [0, -2, 3]]...), 2) == 2
    @test max_sum_submatrix(hcat([[2, 2, -1]]...), 3) == 3
end


using Printf

function variable_examples()
    println("--------------------------------------")
    println("variable_examples")
    println("Variables types can be dynamically assigned")
    s = 0
    println(s)
    s = true
    println(s)
    s = "Dog"
    println(s)

    println("We can also specify variable types")
    x::Int8 = 10
    println(x)
    y::Int16 = -10
    println(y)

    # If we try to assign a different type to a typed variable, we get an exception
    #y = 3.14
    #println(y)

    # Int8 : -128 - 127
    # Int16 : −32,768 - 32,767
    # Int32 : −2,147,483,648 - 2,147,483,647
    # Int64 : −2^63 - 2^63 - 1
    # Float32, Float64, UInt8, UInt16, etc.
    # UInt8, UInt16, etc.

    println("We can also use booleans")
    somebool = true
    println(somebool)

    println("Chars have to be surrounded by single quotes")
    c = 'c'
    println(c)

    println("Cast from Int to Char")
    c2 = Char(120)
    println(c2)

    println("Cast from Float to Int")
    i1 = UInt8(trunc(3.14))
    println(i1)

    println("Cast string to Float")
    f1 = parse(Float64,"1")
    println(f1)

    println("Cast string to Int")
    i2 = parse(Int8,"1")
    println(i2)

    println("We can also create BigInts")
    someBigInt::BigInt = 1234 * 1234 * 1234 * 1234
    println(someBigInt)

    println("And we can create BigFloats")
    someBigFloat::BigFloat = 1234 / 1234 / 1234 / 1234
    println(someBigFloat)

end