text:
include("bernstein.jl")
using Plots
using LaTeXStrings
using WriteVTK
plotly()

mutable struct Mesh
    elemVertex::Array{Float64,2}
    elemNode::Array{Array{Int64,1},1}
    degP::Array{Int64,1}
    C::Array{Array{Float64,2},1}
    numBasis::Int64
    numElem::Int64
    #localKnots::Array{Array{Float64,1},1}
    controlPoints::Array{Float64, 2}
    weights::Array{Float64, 1}
end

mutable struct Neighbors
    left::Array{Int64,1}
    right::Array{Int64,1}
end

"""
create the IEN (node index to element) array for a given knot vector and p and
elementVertex array

INPUT: knotVector - the given knot vector
        numElem - number of non-zero knot-spans
        p - polynomial degree
OUTPUT: IEN - array with nb rows and p+1 columns, where each row indicates the
                global node indices for a knot-span
        elemVertex - array with nb rows and 2 columns, where each row indicates
        the left and right knot of a non-empty knot-span
"""
function makeIEN(knotVector, numElem, p)
    IEN  = zeros(Int64, numElem, p+1)
    tolEq = 1e-10
    elemVertex = zeros(numElem, 2)
    elementCounter = 0
    for indexKnot = 1:length(knotVector)-1
        if knotVector[indexKnot+1]>knotVector[indexKnot]+tolEq
            elementCounter += 1
            IEN[elementCounter,:] = indexKnot-p:indexKnot
            elemVertex[elementCounter, :] = [knotVector[indexKnot], knotVector[indexKnot+1]]
        end
    end
    @assert numElem==elementCounter "Wrong number of elements passed"
    return IEN, elemVertex
end

function makeIEN(knotVectors::Array{Array{Float64,1},1}, numElem::Integer, p::Array{Int64, 1})
    numDim = length(knotVectors)
    numEntries = prod(p.+1)
    IEN = zeros(Int64, numElem, numEntries)
    tolEq = 1e-10
    elementVertex = zeros(numElem, 2^numDim)
    elementCounter = 0
    if numDim==2
        knotU = knotVectors[1]
        knotV = knotVectors[2]
        lenU = length(knotU)-p[1]-1
        for j=1:length(knotV)-1
            for i=1:length(knotU)-1
                if (knotU[i+1]>knotU[i]) && (knotV[j+1]>knotV[j])
                    elementCounter += 1
                    elementVertex[elementCounter, :] = [knotU[i], knotV[j], knotU[i+1], knotV[j+1]]

                    # now we add the nodes from i-p,..., i in the u direction
                    # j-q,..., j in the v direction
                    tcount = 0
                    for t2=j-degP[2]:j
                        for t1=i-degP[1]:i
                            tcount += 1
                            IEN[elementCounter, tcount] = t1+(t2-1)*lenU
                        end
                    end
                end
            end
        end
    end
    @assert numElem==elementCounter "Wrong number of elements passed"
    return IEN, elementVertex
end

"""
Initializes the elementNode, and elementChildren arrays
INPUT: knotVector - the given knot vector
        numElem - number of non-zero knot-spans
OUTPUT:
        elementNeighbors - array with nb elements of type Neighbors
        elementChildren - array with nb elements of type Children
"""
function initNeighborsChildren(knotVector, numElem)
    tolEq = 1e-10
    elementNeighbors = Array{Neighbors,1}(undef, numElem)
    elementChildren = Array{Array{Int64,1},1}(undef, numElem)
    elementCounter = 0
    for indexKnot = 1:length(knotVector)-1
        if knotVector[indexKnot+1]>knotVector[indexKnot]+tolEq
            elementCounter = elementCounter + 1
            elementNeighbors[elementCounter] = Neighbors(Array{Int64,1}(undef,1),
                                                         Array{Int64,1}(undef,1))

            if elementCounter>1
                elementNeighbors[elementCounter].left = [elementCounter-1]
            else
                elementNeighbors[elementCounter].left = []
            end
            if elementCounter<numElem
                elementNeighbors[elementCounter].right = [elementCounter+1]
            else
                elementNeighbors[elementCounter].right = []
            end
            elementChildren[elementCounter] = []
        end
    end
    @assert numElem==elementCounter "Wrong number of elements passed"
    return elementNeighbors, elementChildren
end

"""
Initialize the local knots corresponding the basis functions determined by the
knotVector with degree degP.
INPUT:
    knotVector - the given knot vector
    degP - polynomial degree
OUTPUT:
    localKnots: array of local knots vectors
"""
function initLocalKnots(knotVector, degP)
    numBasis = length(knotVector) - degP - 1
    localKnots = Array{Array{Float64,1}}(undef, numBasis)
    splineChildren = Array{Array{Float64,1}}(undef, numBasis)
    indexKnot = 1
    for i = 1:numBasis
        localKnots[i]= knotVector[i:i+degP+1]
    end
    return localKnots
end

"""
Initializes an IGA mesh from a NURBS object
"""
function genMesh(nurbs::NURBS)
    if length(nurbs.knots)==1
        #1D mesh
        knotU = nurbs.knots[1]
        degP = nurbs.order.-1
        Cmat, numElem = bezierExtraction(knotU, degP[1])
        C = Array{Array{Float64,2},1}(undef, numElem)
        for i=1:numElem
            C[i] = Cmat[:,:,i]
        end
        numBasis = length(knotU) - degP[1] - 1
        IEN, elemVertex = makeIEN(knotU, numElem, degP[1])

        #localKnots = initLocalKnots(knotU, degP)

    elseif length(nurbs.knots)==2
        #2D mesh
        knotU = nurbs.knots[1]
        knotV = nurbs.knots[2]
        degP = nurbs.order.-1
        CmatU, numElemU = bezierExtraction(knotU, degP[1])
        CmatV, numElemV = bezierExtraction(knotV, degP[2])
        numElem = numElemU * numElemV
        C = Array{Array{Float64,2},1}(undef, numElem)
        indexMatrix = permutedims(reshape(1:numElem, numElemU, numElemV), [2,1])
        for j=1:numElemV
            for i=1:numElemU
                elementIndex = indexMatrix[j,i]
                C[elementIndex] = kron(CmatV[:,:,j], CmatU[:,:,i])
            end
        end
        IEN, elemVertex = makeIEN([knotU, knotV], numElem, degP)
        numBasis = (length(knotU) - degP[1] - 1)*(length(knotV) - degP[2] - 1)
    end
    cpts = reshape(nrb.coefs[1:3,:,:,:], 3, numBasis)
    wgts = reshape(nrb.coefs[4,:,:,:], numBasis)
    for i=1:3
        cpts[i,:] ./= wgts
    end
    elemNode = Array{Array{Int64,1},1}(undef, numElem)
    for i=1:numElem
        elemNode[i]=IEN[i,:]
    end
    IGAmesh = Mesh(elemVertex, elemNode, degP, C, numBasis, numElem, cpts, wgts)
    return IGAmesh
end
"""
Plots the basis functions of a mesh in parameter space
"""
function plotBasisParam(mesh::Mesh)
    if length(degP)==1
        #1D plot
        colorList = ["blue", "red", "green", "black", "magenta"]
        graph=Plots.plot(title="B-Splines of degree $(mesh.degP[1])")
        numPtsElem = 11
        evalPts = LinRange(-1, 1, numPtsElem)
        B, dB = bernsteinBasis(evalPts, degP[1])

        for iBasis = 1:mesh.numBasis
            colorIndex = ((iBasis-1) % length(colorList))+1
            for iElem = 1:mesh.numElem
                localIndex = findall(isequal(iBasis), mesh.elemNode[iElem])
                if length(localIndex)>0
                    uMin = mesh.elemVertex[iElem, 1]
                    uMax = mesh.elemVertex[iElem, 2]
                    plotPts = LinRange(uMin, uMax, numPtsElem)
                    plotVal = B*(mesh.C[iElem][localIndex,:])'
                    graph = plot!(plotPts, plotVal, color=colorList[colorIndex], leg=false, line=2)
                    graph = scatter!([uMin], [0], color="red", leg=false, markersize=5)
                    graph = scatter!([uMax], [0], color="red", leg=false, markersize=5)
                end
            end
        end
    elseif length(degP)==2
        #2D plot
        graph=plot(title="B-Splines of degree $(mesh.degP)")
        numPtsElem = 11
        evalPtsU = LinRange(-1, 1, numPtsElem)
        evalPtsV = LinRange(-1, 1, numPtsElem)
        Bu, dBu = bernsteinBasis(evalPtsU, degP[1])
        Bv, dBv = bernsteinBasis(evalPtsV, degP[2])
        basisCounter = 0
        Buv = zeros(numPtsElem, numPtsElem, (degP[1]+1)*(degP[2]+1))
        for j=1:degP[2]+1
            for i=1:degP[1]+1
                basisCounter += 1
                Buv[:,:,basisCounter] = Bu[:,i]*Bv[:,j]'
            end
        end

        for iBasis = 1:mesh.numBasis
            for iElem = 1:mesh.numElem
                localIndex = findall(isequal(iBasis), mesh.elemNode[iElem])
                if length(localIndex)>0
                    uMin = mesh.elemVertex[iElem, 1]
                    uMax = mesh.elemVertex[iElem, 3]
                    vMin = mesh.elemVertex[iElem, 2]
                    vMax = mesh.elemVertex[iElem, 4]
                    plotPtsU = LinRange(uMin, uMax, numPtsElem)
                    plotPtsV = LinRange(vMin, vMax, numPtsElem)
                    cR = zeros(numPtsElem, numPtsElem)
                    for j=1:numPtsElem
                        for i=1:numPtsElem
                            cR[i,j] = (mesh.C[iElem][localIndex,:]*Buv[i,j,:])[1]
                        end
                    end

                end
            end
        end
    end

    display(graph)
end

"""
Plots the computed solution
"""
function plotSol(mesh::Mesh, sol0, fileName::String)

    if length(degP)==1
        graph=plot(title="Approximate solution")
        numPtsElem = 11
        evalPts = LinRange(-1, 1, numPtsElem)
        B, dB = bernsteinBasis(evalPts, degP[1])

        for iElem in 1:mesh.numElem
            curNodes = mesh.elemNode[iElem]
            uMin = mesh.elemVertex[iElem, 1]
            uMax = mesh.elemVertex[iElem, 2]
            plotPts = LinRange(uMin, uMax, numPtsElem)
            splineVal = B*(mesh.C[iElem])'
            cpts = mesh.controlPoints[1, curNodes]
            wgts = mesh.weights[curNodes]
            basisVal = zero(splineVal)
            for iPlotPt = 1:numPtsElem
                RR = splineVal[iPlotPt,:].* wgts
                w_sum = sum(RR)
                RR /= w_sum
                basisVal[iPlotPt,:] = RR
            end
            solVal = basisVal*sol0[curNodes]
            graph = plot!(plotPts, solVal, color="blue", leg=false, line=2)
            graph = plot!(plotPts, zeros(length(plotPts)), color="black", leg=false, line=1)
            graph = scatter!([uMin], [0], color="red", leg=false, markersize=5)
            graph = scatter!([uMax], [0], color="red", leg=false, markersize=5)
        end
        display(graph)
    elseif length(degP)==2
        #Set the order of points for VTK_QUADRATIC_QUAD  w.r.t. a 3x3 grid from Figure 3 of
        # https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf
        VTKOrder = [1, 3, 9, 7, 2, 6, 8, 4]
        numPtsElem = 3
        evalPtsU = LinRange(-1, 1, numPtsElem)
        evalPtsV = LinRange(-1, 1, numPtsElem)
        Bu, dBu = bernsteinBasis(evalPtsU, degP[1])
        Bv, dBv = bernsteinBasis(evalPtsV, degP[2])
        basisCounter = 0
        Buv = zeros(numPtsElem, numPtsElem, (degP[1]+1)*(degP[2]+1))
        for j=1:degP[2]+1
            for i=1:degP[1]+1
                basisCounter += 1
                Buv[:,:,basisCounter] = Bu[:,i]*Bv[:,j]'
            end
        end
        pointsVTK = zeros(2, mesh.numElem*8)
        pointValVTK = zeros(1, mesh.numElem*8)
        localPoints = zeros(2, 9)
        localPointVal = zeros(9)
        meshCells = Array{MeshCell, 1}(undef, mesh.numElem)
        for iElem in 1:mesh.numElem
            curNodes = mesh.elemNode[iElem]
            uMin = mesh.elemVertex[iElem, 1]
            uMax = mesh.elemVertex[iElem, 3]
            vMin = mesh.elemVertex[iElem, 2]
            vMax = mesh.elemVertex[iElem, 4]
            cpts = mesh.controlPoints[1:2, curNodes]
            wgts = mesh.weights[curNodes]
            localPtCounter = 0
            for jPt = 1:numPtsElem
                for iPt = 1:numPtsElem
                    #compute the (B-)spline basis functions and derivatives with Bezier extraction
                    N_mat = mesh.C[iElem] * Buv[iPt, jPt, :]
                    #compute the rational basis
                    RR = N_mat.* wgts
                    w_sum = sum(RR)
                    RR /= w_sum
                    phys_pt = cpts*RR
                    localPtCounter += 1
                    localPoints[:, localPtCounter ]=phys_pt
                    solVal = RR'*sol0[curNodes]
                    localPointVal[localPtCounter]=solVal
                end
            end
            ptRange = (iElem-1)*8+1:iElem*8
            pointsVTK[:, ptRange] = localPoints[:, VTKOrder]
            pointValVTK[:, ptRange] = localPointVal[VTKOrder]
            meshCells[iElem] = MeshCell(VTKCellTypes.VTK_QUADRATIC_QUAD, collect(ptRange))
        end
        #Output data to VTK
        vtkfile = vtk_grid(fileName, pointsVTK, meshCells)
        vtkfile["U", VTKPointData()] = pointValVTK
        vtk_save(vtkfile)
        println("Output written to $fileName.vtk")
    end

end

"""
Plots the computed solution for elasticity
"""
function plotSolElast(mesh::Mesh, Cmat::Array{Float64}, sol0, fileName::String)
    if length(degP)==2
        #Set the order of points for VTK_QUADRATIC_QUAD  w.r.t. a 3x3 grid from Figure 3 of
        # https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf
        VTKOrder = [1, 3, 9, 7, 2, 6, 8, 4]
        numPtsElem = 3
        evalPtsU = LinRange(-1, 1, numPtsElem)
        evalPtsV = LinRange(-1, 1, numPtsElem)
        Buv, dBdu, dBdv = bernsteinBasis2D(collect(evalPtsU), collect(evalPtsV), mesh.degP)
        pointsVTK = zeros(2, mesh.numElem*8)
        pointDispVTK = zeros(2, mesh.numElem*8)
        pointStressVTK = zeros(3, mesh.numElem*8)
        pointStressVMVTK = zeros(mesh.numElem*8)
        localPoints = zeros(2, 9)
        localPointDisp = zeros(2, 9)
        localPointStress = zeros(3, 9)
        localPointStressVM = zeros(1, 9)
        meshCells = Array{MeshCell, 1}(undef, mesh.numElem)
        for iElem in 1:mesh.numElem
            curNodes = mesh.elemNode[iElem]
            numNodes = length(curNodes)
            curNodesXY = reshape(hcat(2*curNodes.-1, 2*curNodes)', 2*numNodes)
            uMin = mesh.elemVertex[iElem, 1]
            uMax = mesh.elemVertex[iElem, 3]
            vMin = mesh.elemVertex[iElem, 2]
            vMax = mesh.elemVertex[iElem, 4]
            cpts = mesh.controlPoints[1:2, curNodes]
            wgts = mesh.weights[curNodes]
            localPtCounter = 0
            B = zeros(2*numNodes,3)
            for jPt = 1:numPtsElem
                for iPt = 1:numPtsElem
                    #compute the (B-)spline basis functions and derivatives with Bezier extraction
                    N_mat = mesh.C[iElem] * Buv[iPt, jPt, :]
                    dN_du = mesh.C[iElem] * dBdu[iPt, jPt, :] * 2/(uMax-uMin)
                    dN_dv = mesh.C[iElem] * dBdv[iPt, jPt, :] * 2/(vMax-vMin)

                    #compute the rational basis
                    RR = N_mat.* wgts
                    dRdu = dN_du.* wgts
                    dRdv = dN_dv.* wgts
                    w_sum = sum(RR)
                    dw_xi = sum(dRdu)
                    dw_eta = sum(dRdv)
                    dRdu = dRdu/w_sum - RR*dw_xi/w_sum^2
                    dRdv = dRdv/w_sum - RR*dw_eta/w_sum^2

                    #compute the derivatives w.r.t to the physical space
                    dR = [dRdu'; dRdv']
                    dxdxi = dR * cpts'
                    RR /= w_sum
                    phys_pt = cpts*RR
                    if abs(det(dxdxi))<1e-12
                        @warn "Singularity in mapping at $phys_pt"
                        dR = pinv(dxdxi)*dR
                    else
                        dR = dxdxi\dR
                    end
                    Jac_par_phys = det(dxdxi)

                    B[1:2:2*numNodes-1,1] = dR[1,:]
                    B[2:2:2*numNodes,2] = dR[2,:]
                    B[1:2:2*numNodes-1,3] = dR[2,:]
                    B[2:2:2*numNodes,3] = dR[1,:]
                    localPtCounter += 1
                    localPoints[:, localPtCounter ]=phys_pt
                    solValX = RR'*sol0[2*curNodes.-1]
                    solValY = RR'*sol0[2*curNodes]
                    localPointDisp[:, localPtCounter]=[solValX, solValY]
                    stressVect = Cmat*B'*sol0[curNodesXY]
                    localPointStress[:, localPtCounter] = stressVect
                    stressVM = sqrt(stressVect[1]^2 - stressVect[1]*stressVect[2] + stressVect[2]^2
                                    +3*stressVect[3]^2)
                    localPointStressVM[localPtCounter] = stressVM
                end
            end
            ptRange = (iElem-1)*8+1:iElem*8
            pointsVTK[:, ptRange] = localPoints[:, VTKOrder]
            pointDispVTK[:, ptRange] = localPointDisp[:, VTKOrder]
            pointStressVTK[:, ptRange] = localPointStress[:, VTKOrder]
            pointStressVMVTK[ptRange] = localPointStressVM[VTKOrder]
            meshCells[iElem] = MeshCell(VTKCellTypes.VTK_QUADRATIC_QUAD, collect(ptRange))
        end
        #Output data to VTK
        vtkfile = vtk_grid(fileName, pointsVTK, meshCells)
        vtkfile["U", VTKPointData()] = pointDispVTK
        vtkfile["Stress", VTKPointData()] = pointStressVTK
        vtkfile["StressVM"] = pointStressVMVTK
        vtk_save(vtkfile)
        println("Output written to $fileName.vtk")
    end

end

"""
Plots the error in the approximate solution
"""
function plotSolError(mesh::Mesh, sol0, exactSol::Function, fileName::String)
    if length(degP)==1
        graph=plot(title="Error \$u-u_h\$")
        numPtsElem = 11
        evalPts = LinRange(-1, 1, numPtsElem)
        B, dB = bernsteinBasis(evalPts, degP[1])

        for iElem in 1:mesh.numElem
            curNodes = mesh.elemNode[iElem]
            uMin = mesh.elemVertex[iElem, 1]
            uMax = mesh.elemVertex[iElem, 2]
            plotPts = LinRange(uMin, uMax, numPtsElem)
            physPts = zeros(numPtsElem)
            splineVal = B*(mesh.C[iElem])'
            cpts = mesh.controlPoints[1, curNodes]
            wgts = mesh.weights[curNodes]
            basisVal = zero(splineVal)
            for iPlotPt = 1:numPtsElem
                RR = splineVal[iPlotPt,:].* wgts
                w_sum = sum(RR)
                RR /= w_sum
                physPts[iPlotPt] = RR'*cpts
                basisVal[iPlotPt,:] = RR
            end
            exSolVal = real(exactSol.(physPts))
            solVal = basisVal*sol0[curNodes]
            graph = plot!(plotPts, exSolVal-solVal, color="blue", leg=false, line=2)
            graph = plot!(plotPts, zeros(length(plotPts)), color="black", leg=false, line=1)
            graph = scatter!([uMin], [0], color="red", leg=false, markersize=5)
            graph = scatter!([uMax], [0], color="red", leg=false, markersize=5)
        end
        display(graph)
    elseif length(degP)==2
        #Set the order of points for VTK_QUADRATIC_QUAD  w.r.t. a 3x3 grid from Figure 3 of
        # https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf
        VTKOrder = [1, 3, 9, 7, 2, 6, 8, 4]
        numPtsElem = 3
        evalPtsU = LinRange(-1, 1, numPtsElem)
        evalPtsV = LinRange(-1, 1, numPtsElem)
        Bu, dBu = bernsteinBasis(evalPtsU, degP[1])
        Bv, dBv = bernsteinBasis(evalPtsV, degP[2])
        basisCounter = 0
        Buv = zeros(numPtsElem, numPtsElem, (degP[1]+1)*(degP[2]+1))
        for j=1:degP[2]+1
            for i=1:degP[1]+1
                basisCounter += 1
                Buv[:,:,basisCounter] = Bu[:,i]*Bv[:,j]'
            end
        end
        pointsVTK = zeros(2, mesh.numElem*8)
        pointValVTK = zeros(1, mesh.numElem*8)
        localPoints = zeros(2, 9)
        localPointVal = zeros(9)
        meshCells = Array{MeshCell, 1}(undef, mesh.numElem)
        for iElem in 1:mesh.numElem
            curNodes = mesh.elemNode[iElem]
            uMin = mesh.elemVertex[iElem, 1]
            uMax = mesh.elemVertex[iElem, 3]
            vMin = mesh.elemVertex[iElem, 2]stressvect
            vMax = mesh.elemVertex[iElem, 4]
            cpts = mesh.controlPoints[1:2, curNodes]
            wgts = mesh.weights[curNodes]
            localPtCounter = 0
            for jPt = 1:numPtsElem
                for iPt = 1:numPtsElem
                    #compute the (B-)spline basis functions and derivatives with Bezier extraction
                    N_mat = mesh.C[iElem] * Buv[iPt, jPt, :]
                    #compute the rational basis
                    RR = N_mat.* wgts
                    w_sum = sum(RR)
                    RR /= w_sum
                    phys_pt = cpts*RR
                    localPtCounter += 1
                    localPoints[:, localPtCounter ]=phys_pt

                    solVal = RR'*sol0[curNodes]
                    exSolVal = real(exactSol(phys_pt[1], phys_pt[2]))
                    localPointVal[localPtCounter]=exSolVal-solVal
                end
            end
            ptRange = (iElem-1)*8+1:iElem*8
            pointsVTK[:, ptRange] = localPoints[:, VTKOrder]
            pointValVTK[:, ptRange] = localPointVal[VTKOrder]
            meshCells[iElem] = MeshCell(VTKCellTypes.VTK_QUADRATIC_QUAD, collect(ptRange))
        end
        #Output data to VTK
        vtkfile = vtk_grid(fileName, pointsVTK, meshCells)
        vtkfile["U-U_h", VTKPointData()] = pointValVTK
        vtk_save(vtkfile)
        println("Output written to $fileName.vtk")
    end
end

"""
Plots the error in the approximate solution
"""
function plotSolErrorElast(mesh::Mesh, Cmat::Array{Float64}, sol0, exactSolDisp::Function, exactSolStress::Function, fileName::String)
    if length(degP)==2
         #Set the order of points for VTK_QUADRATIC_QUAD  w.r.t. a 3x3 grid from Figure 3 of
        # https://vtk.org/wp-content/uploads/2015/04/file-formats.pdf
        VTKOrder = [1, 3, 9, 7, 2, 6, 8, 4]
        numPtsElem = 3
        evalPtsU = LinRange(-1, 1, numPtsElem)
        evalPtsV = LinRange(-1, 1, numPtsElem)
        Buv, dBdu, dBdv = bernsteinBasis2D(collect(evalPtsU), collect(evalPtsV), mesh.degP)
        pointsVTK = zeros(2, mesh.numElem*8)
        pointDispVTK = zeros(2, mesh.numElem*8)
        pointStressVTK = zeros(3, mesh.numElem*8)
        pointStressVMVTK = zeros(mesh.numElem*8)
        localPoints = zeros(2, 9)
        localPointDisp = zeros(2, 9)
        localPointStress = zeros(3, 9)
        localPointStressVM = zeros(1, 9)
        meshCells = Array{MeshCell, 1}(undef, mesh.numElem)
        for iElem in 1:mesh.numElem
            curNodes = mesh.elemNode[iElem]
            numNodes = length(curNodes)
            curNodesXY = reshape(hcat(2*curNodes.-1, 2*curNodes)', 2*numNodes)
            uMin = mesh.elemVertex[iElem, 1]
            uMax = mesh.elemVertex[iElem, 3]
            vMin = mesh.elemVertex[iElem, 2]
            vMax = mesh.elemVertex[iElem, 4]
            cpts = mesh.controlPoints[1:2, curNodes]
            wgts = mesh.weights[curNodes]
            localPtCounter = 0
            for jPt = 1:numPtsElem
                for iPt = 1:numPtsElem
                    #compute the (B-)spline basis functions and derivatives with Bezier extraction
                    N_mat = mesh.C[iElem] * Buv[iPt, jPt, :]
                    dN_du = mesh.C[iElem] * dBdu[iPt, jPt, :] * 2/(uMax-uMin)
                    dN_dv = mesh.C[iElem] * dBdv[iPt, jPt, :] * 2/(vMax-vMin)

                    #compute the rational basis
                    RR = N_mat.* wgts
                    dRdu = dN_du.* wgts
                    dRdv = dN_dv.* wgts
                    w_sum = sum(RR)
                    dw_xi = sum(dRdu)
                    dw_eta = sum(dRdv)
                    dRdu = dRdu/w_sum - RR*dw_xi/w_sum^2
                    dRdv = dRdv/w_sum - RR*dw_eta/w_sum^2

                    #compute the derivatives w.r.t to the physical space
                    dR = [dRdu'; dRdv']
                    dxdxi = dR * cpts'
                    RR /= w_sum
                    phys_pt = cpts*RR
                    if abs(det(dxdxi))<1e-12
                        @warn "Singularity in mapping at $phys_pt"
                        dR = pinv(dxdxi)*dR
                    else
                        dR = dxdxi\dR
                    end
                    Jac_par_phys = det(dxdxi)

                    B = zeros(2*numNodes,3);
                    B[1:2:2*numNodes-1,1] = dR[1,:]
                    B[2:2:2*numNodes,2] = dR[2,:]
                    B[1:2:2*numNodes-1,3] = dR[2,:]
                    B[2:2:2*numNodes,3] = dR[1,:]
                    localPtCounter += 1
                    localPoints[:, localPtCounter ]=phys_pt
                    solValX = RR'*sol0[2*curNodes.-1]
                    solValY = RR'*sol0[2*curNodes]

                    #evaluate the exact displacement and stresses
                    exSolVal = exactSolDisp(phys_pt[1], phys_pt[2])
                    exSolStressVal = exactSolStress(phys_pt[1], phys_pt[2])
                    exStressVM = sqrt(exSolStressVal[1]^2 - exSolStressVal[1]*exSolStressVal[2] + exSolStressVal[2]^2
                                    +3*exSolStressVal[3]^2)

                    #compute the errors (exact - approximate)
                    localPointDisp[:, localPtCounter]=[exSolVal[1]-solValX, exSolVal[2]-solValY]
                    stressVect = Cmat*B'*sol0[curNodesXY]
                    localPointStress[:, localPtCounter] = exSolStressVal - stressVect
                    stressVM = sqrt(stressVect[1]^2 - stressVect[1]*stressVect[2] + stressVect[2]^2
                                    +3*stressVect[3]^2)
                    localPointStressVM[localPtCounter] = exStressVM - stressVM
                end
            end
            ptRange = (iElem-1)*8+1:iElem*8
            pointsVTK[:, ptRange] = localPoints[:, VTKOrder]
            pointDispVTK[:, ptRange] = localPointDisp[:, VTKOrder]
            pointStressVTK[:, ptRange] = localPointStress[:, VTKOrder]
            pointStressVMVTK[ptRange] = localPointStressVM[VTKOrder]
            meshCells[iElem] = MeshCell(VTKCellTypes.VTK_QUADRATIC_QUAD, collect(ptRange))
        end
        #Output data to VTK
        vtkfile = vtk_grid(fileName, pointsVTK, meshCells)
        vtkfile["error U", VTKPointData()] = pointDispVTK
        vtkfile["error Stress", VTKPointData()] = pointStressVTK
        vtkfile["error StressVM"] = pointStressVMVTK
        vtk_save(vtkfile)
        println("Output written to $fileName.vtk")
    end
end


text:
using JET
using Test

using ClimaCore.RecursiveApply

@static if @isdefined(var"@test_opt") # v1.7 and higher
    @testset "RecursiveApply optimization test" begin
        for x in [
            1.0,
            1.0f0,
            (1.0, 2.0),
            (1.0f0, 2.0f0),
            (a = 1.0, b = (x1 = 2.0, x2 = 3.0)),
            (a = 1.0f0, b = (x1 = 2.0f0, x2 = 3.0f0)),
        ]
            @test_opt 2 ⊠ x
            @test_opt x ⊞ x
            @test_opt RecursiveApply.rdiv(x, 3)
        end
    end
end

@testset "RecursiveApply nary ops" begin
    for x in [
        1.0,
        1.0f0,
        (1.0, 2.0),
        (1.0f0, 2.0f0),
        (a = 1.0, b = (x1 = 2.0, x2 = 3.0)),
        (a = 1.0f0, b = (x1 = 2.0f0, x2 = 3.0f0)),
    ]
        FT = eltype(x[1])
        @test RecursiveApply.rmul(x, one(FT), one(FT), one(FT)) == x
        @test RecursiveApply.rmul(x, one(FT), x, one(FT)) ==
              RecursiveApply.rmul(x, x)
        @test RecursiveApply.radd(x, zero(FT), zero(FT), zero(FT)) == x
        @test RecursiveApply.radd(x, zero(FT), x, zero(FT)) ==
              RecursiveApply.rmul(x, FT(2))
    end
end

ext: jl


text:
using Destruct
using BenchmarkTools

unpack_bc(w::Array{<:Tuple}) = Tuple((v->v[i]).(w) for i=1:length(w[1]))

f(a, b) = a+b, a*b, a-b
println("getindex-broadcast / destruct")
println("  Homogeneous types")
packed = f.(rand(2,2,1), rand(1,1,2))
println("$(typeof(packed))")
for sz in [10 50 100]
    packed = f.(rand(sz,sz,1), rand(1,1,sz))
    r = (@belapsed unpack_bc($packed))/(@belapsed destruct($packed))
    println("    $sz^3 : $r")
end
println("")

g(a, b) = a+b+1im*(a-b), a*b, convert(Int, round(a-b))

println("  Heterogeneous types")
packed = g.(rand(2,2,1), rand(1,1,2))
println("$(typeof(packed))")
for sz in [10 50 100]
    packed = g.(rand(sz,sz,1), rand(1,1,sz))
    r = (@belapsed unpack_bc($packed))/(@belapsed destruct($packed))
    println("    $sz^3 : $r")
end


text:
const cufftNumber = Union{cufftDoubleReal,cufftReal,cufftDoubleComplex,cufftComplex}
const cufftReals = Union{cufftDoubleReal,cufftReal}
const cufftComplexes = Union{cufftDoubleComplex,cufftComplex}
const cufftDouble = Union{cufftDoubleReal,cufftDoubleComplex}
const cufftSingle = Union{cufftReal,cufftComplex}
const cufftTypeDouble = Union{Type{cufftDoubleReal},Type{cufftDoubleComplex}}
const cufftTypeSingle = Union{Type{cufftReal},Type{cufftComplex}}

cufftfloat(x) = _cufftfloat(float(x))
_cufftfloat(::Type{T}) where {T<:cufftReals} = T
_cufftfloat(::Type{Float16}) = Float32
_cufftfloat(::Type{Complex{T}}) where {T} = Complex{_cufftfloat(T)}
_cufftfloat(::Type{T}) where {T} = error("type $T not supported")
_cufftfloat(x::T) where {T} = _cufftfloat(T)(x)

complexfloat(x::DenseCuArray{Complex{<:cufftReals}}) = x
realfloat(x::DenseCuArray{<:cufftReals}) = x

complexfloat(x::DenseCuArray{T}) where {T<:Complex} = copy1(typeof(cufftfloat(zero(T))), x)
complexfloat(x::DenseCuArray{T}) where {T<:Real} = copy1(typeof(complex(cufftfloat(zero(T)))), x)

realfloat(x::DenseCuArray{T}) where {T<:Real} = copy1(typeof(cufftfloat(zero(T))), x)

function copy1(::Type{T}, x) where T
    y = CuArray{T}(undef, map(length, axes(x)))
    #copy!(y, x)
    y .= broadcast(xi->convert(T,xi),x)
end


text:
using OptionalData
using Test

struct TestType
    a::Float64
    b::Float64
    c::Float64
end

@OptionalData opt1 Symbol
@OptionalData opt2 TestType

@testset "OptionalData" begin
    @test string(opt1) == "OptData{Symbol}()"
    @test_throws ErrorException get(opt1)
    push!(opt1, :Test)
    @test string(opt1) == "OptData{Symbol}(Test)"
    @test get(opt1) == :Test
    push!(opt1, "Test")
    @test get(opt1) == :Test

    @test string(opt2) == "OptData{TestType}()"
    @test_throws ErrorException get(opt2)
    push!(opt2, TestType(1, 2, 3))
    @test string(opt2) == "OptData{TestType}(TestType(1.0, 2.0, 3.0))"
    @test get(opt2) == TestType(1, 2, 3)
    push!(opt2, 1, 2, 3)
    @test get(opt2) == TestType(1, 2, 3)
end


text:
@testset "Quaternions" begin

    qx = qrotation(Vec(1, 0, 0), pi / 4)
    @test qx * qx ≈ qrotation(SVector(1.0, 0.0, 0.0), pi / 2)
    @test Base.power_by_squaring(qx, 2) ≈ qrotation(SVector(1.0, 0.0, 0.0), pi / 2)
    theta = pi / 8
    qx = qrotation(SVector(1.0, 0.0, 0.0), theta)
    c = cos(theta); s = sin(theta)
    Rx = [1 0 0; 0 c -s; 0 s c]

    @test Mat3f0(qx) ≈ Rx
    theta = pi / 6
    qy = qrotation(SVector(0.0, 1.0, 0.0), theta)
    c = cos(theta); s = sin(theta)
    Ry = [c 0 s; 0 1 0; -s 0 c]
    @test Mat3f0(qy) ≈ Ry
    theta = 4pi / 3
    qz = qrotation(SVector(0.0, 0.0, 1.0), theta)
    c = cos(theta); s = sin(theta)
    Rz = [c -s 0; s c 0; 0 0 1]
    @test Mat3f0(qz) ≈ Rz

    @test Mat3f0(qx * qy * qz) ≈ Rx * Ry * Rz
    @test Mat3f0(qy * qx * qz) ≈ Ry * Rx * Rz
    @test Mat3f0(qz * qx * qy) ≈ Rz * Rx * Ry

    a, b = qrotation(SVector(0.0, 0.0, 1.0), deg2rad(0)), qrotation(SVector(0.0, 0.0, 1.0), deg2rad(180))
    # @test slerp(a, b, 0.0) ≈ a
    # @test slerp(a, b, 1.0) ≈ b
    # @test slerp(a, b, 0.5) ≈ qrotation([0, 0, 1], deg2rad(90))
    #
    # @test angle(qrotation(SVector(1.0, 0.0, 0.0), 0)) ≈ 0
    # @test angle(qrotation([0, 1, 0], pi / 4)) ≈ pi / 4
    # @test angle(qrotation([0, 0, 1], pi / 2)) ≈ pi / 2
    #
    # let # test numerical stability of angle
    #     ax = randn(3)
    #     for θ in [1e-9, π - 1e-9]
    #         q = qrotation(ax, θ)
    #         @test angle(q) ≈ θ
    #     end
    # end
end


text:
using LazyJSON
using Test

const JSON = LazyJSON

@testset "IOString" begin

@testset "IOString Error" begin
    log = [] ; io = Base.BufferStream()

    @sync begin
        @async begin
            try
                v = LazyJSON.value(io)
                push!(log, Int64(v))
            catch e
                push!(log, e)
            end
        end

        @test isempty(log)
        write(io, "a")
        while isempty(log) sleep(0.01) end; @test log[1] isa LazyJSON.ParseError
        close(io)
    end
end

@testset "IOString Number" begin
    log = [] ; io = Base.BufferStream()

    @sync begin
        @async begin
            v = LazyJSON.value(io)
            push!(log, Int64(v))
        end

        @test isempty(log)
        write(io, "1") ; sleep(0.01)  ;@test isempty(log)
        write(io, "2") ; sleep(0.01)  ;@test isempty(log)
        write(io, "3") ; sleep(0.01)  ;@test isempty(log)
        close(io)      ; while isempty(log) sleep(0.01) end
                                       @test log[1] == 123
    end
end

@testset "IOString String" begin
    log = [] ; io = Base.BufferStream()

    @sync begin
        @async begin
            v = LazyJSON.value(io)
            push!(log, String(v))
        end

        @test isempty(log)
        write(io, "\"1") ; sleep(0.01)  ;@test isempty(log)
        write(io, "2")   ; sleep(0.01)  ;@test isempty(log)
        write(io, "3\"") ; while isempty(log) sleep(0.01) end
                                         @test log[1] == "123"
        close(io)
    end
end

@testset "IOString Lazy String" begin
    log = [] ; io = Base.BufferStream()

    @sync begin
        @async begin
            v = LazyJSON.value(io)
            push!(log, v)
        end

        etb = LazyJSON.IOStrings.ASCII_ETB
        @test isempty(log)
        write(io, "\"") ; sleep(0.01) ;@test log[1] isa JSON.String
                                      ;@test log[1].s.buf.data[1] == UInt8('"')
                                       @test log[1].s.buf.data[2] == etb
        write(io, "12") ; sleep(0.01) ;@test log[1].s.buf.data[1] == UInt8('"')
                                       @test log[1].s.buf.data[2] == etb
        write(io, "3\""); sleep(0.01) ;@test log[1] == "123"
        close(io)
    end
end

@testset "IOString Array" begin
    log = [] ; io = Base.BufferStream()

    @sync begin
        @async begin
            v = LazyJSON.value(io)
            for x in v
                push!(log, Int64(x))
            end
        end

        @test isempty(log)
        write(io, "[")    ; sleep(0.01)  ;@test isempty(log)
        write(io, "1")    ; sleep(0.01)  ;@test isempty(log)
        write(io, ",")    ; sleep(0.01)  ;@test log[1] == 1
        write(io, "2,")   ; sleep(0.01)  ;@test log[2] == 2
        write(io, "3")    ; sleep(0.01)  ;@test length(log) == 2
        write(io, "4")    ; sleep(0.01)  ;@test length(log) == 2
        write(io, "5")    ; sleep(0.01)  ;@test length(log) == 2
        write(io, "]")    ; sleep(0.01)  ;@test log[3] == 345
        close(io)
    end
end

@testset "IOString Object" begin
    log = [] ; io = Base.BufferStream()

    @sync begin
        t = @async begin
            v = LazyJSON.value(io)
            for (k, x) in v
                push!(log, (String(k), String(x)))
            end
        end

        @test isempty(log)
        write(io, "{")            ; sleep(0.01)  ;@test isempty(log)
        write(io, "\"f")          ; sleep(0.01)  ;@test isempty(log)
        write(io, "oo\":")        ; sleep(0.01)  ;@test isempty(log)
        write(io, "\"A")          ; sleep(0.01)  ;@test isempty(log)
        write(io, "\"")           ; sleep(0.01)  ;@test log[1] == ("foo", "A")
        write(io, ",")            ; sleep(0.01)
        write(io, "\"bar\":")
        write(io, "   \"123\"  ,"); sleep(0.01)  ;@test log[2] == ("bar", "123")
        write(io, "}")            ; sleep(0.01)
        close(io)
    end
end

@testset "IOString split" begin
    log = [] ; io = Base.BufferStream()

    gist_json = """{
      "id": 1296269,
      "owner": {
        "login": "octocat"
      },
      "parent": {
        "name": "test-parent"
      },
      "full_name": "octocat/Hello-World",
      "private": false,
      "url": "https://api.github.com/repos/octocat/Hello-World",
      "language": null,
      "pushed_at": "2011-01-26T19:06:43Z",
      "permissions": {
        "admin": false,
        "push": false,
        "pull": true
      }
    }"""
    bytes = Vector{UInt8}(codeunits(gist_json))

    @sync begin
        t = @async begin
            v = LazyJSON.value(io, getproperty=true)
            push!(log, v.id)
            push!(log, v.owner.login)
            push!(log, v.private)
            push!(log, v.permissions.admin)
        end

        @test isempty(log)
        write(io, bytes[1:80])    ; sleep(0.01)  ;@test log[1] == 1296269
                                                  @test log[2] == "octocat"
                                                  @test length(log) == 2
        write(io, bytes[81:200])  ; sleep(0.01)  ;@test log[3] == false
                                                  @test length(log) == 3
        write(io, bytes[201:end]) ; sleep(0.01)  ;@test log[4] == false
                                                  @test length(log) == 4
        close(io)
    end
end

end #testset


text:
using IterativeSolvers
using FactCheck

facts("svdvals_gkl") do

context("Small diagonal matrix") do
    A = full(Diagonal([10.0, 9, 8, 6, 1]))
    @fact norm(svdvals_gkl(A)[1] - svdvals(A)) --> less_than(1e-10)
end

for T in (Float32, Float64)
    context("Medium random UpperTriangular{$T}") do
        srand(1)
        n = 150
        nvals = 6

        Base.Ac_mul_B!(p::Vector, A::UpperTriangular, u::Vector) = Ac_mul_B!(A, copy!(p, u))
        Base.At_mul_B!(p::Vector, A::UpperTriangular, u::Vector) = At_mul_B!(A, copy!(p, u))

        A = UpperTriangular(convert(Matrix{T}, rand(n, n)))
        σth = 0.1*√eps(T)

        svals, = svdvals_gkl(A, nvals, σth=σth)
        @fact length(svals) --> greater_than_or_equal(nvals)

        svals = svals[1:nvals]
        svals2 = svdvals(A)[1:nvals]
        @fact norm(svals - svals2) --> less_than(nvals*σth) "Disagreement in top $nvals singular values:\nsvdvals_gkl\tsvdvals\n"*repr([svals svals2])
    end
end

for T in (Float32, Float64)
    context("Medium wide rectangular Matrix{$T}") do
        srand(1)
        n = 500
        m = 200
        nvals = 10

        A = convert(Matrix{T}, rand(n, m))
        σth = 0.1*√eps(T)

        svals, = svdvals_gkl(A, nvals, σth=σth)
        @fact length(svals) --> greater_than_or_equal(nvals)

        nvals = 4
        svals = svals[1:nvals]
        svals2 = svdvals(A)[1:nvals]
        @fact norm(svals - svals2) --> less_than(nvals*σth) "Disagreement in top $nvals singular values:\nsvdvals_gkl\tsvdvals\n"*repr([svals svals2])
    end

    context("Medium tall rectangular Matrix{$T}") do
        srand(1)
        n = 200
        m = 500
        nvals = 10

        A = convert(Matrix{T}, rand(n, m))
        σth = 0.1*√eps(T)

        svals, = svdvals_gkl(A, nvals, σth=σth)
        @fact length(svals) --> greater_than_or_equal(nvals)

        nvals = 4
        svals = svals[1:nvals]
        svals2 = svdvals(A)[1:nvals]
        @fact norm(svals - svals2) --> less_than(nvals*σth) "Disagreement in top $nvals singular values:\nsvdvals_gkl\tsvdvals\n"*repr([svals svals2])
    end
e
end

end


text:
using DataFrames
using Test

m = test_page_model()

for gas in [:CO2, :CH4, :N2O, :Lin]
    MimiPAGE2009.addabatementcostparameters(m, gas)
    MimiPAGE2009.addabatementcosts(m, gas)

    comp_name1 = Symbol("AbatementCostParameters$gas")
    comp_name2 = Symbol("AbatementCosts$gas")

    connect_param!(m, comp_name1 => :cbe_absoluteemissionreductions, comp_name2 => :cbe_absoluteemissionreductions)

    connect_param!(m, comp_name2 => :zc_zerocostemissions, comp_name1 => :zc_zerocostemissions)
    connect_param!(m, comp_name2 => :q0_absolutecutbacksatnegativecost, comp_name1 => :q0_absolutecutbacksatnegativecost)
    connect_param!(m, comp_name2 => :blo, comp_name1 => :blo)
    connect_param!(m, comp_name2 => :alo, comp_name1 => :alo)
    connect_param!(m, comp_name2 => :bhi, comp_name1 => :bhi)
    connect_param!(m, comp_name2 => :ahi, comp_name1 => :ahi)

end

set_param!(m, :yagg, readpagedata(m,"test/validationdata/yagg_periodspan.csv"))

p = load_parameters(m)
p["y_year_0"] = 2008.
p["y_year"] = Mimi.dim_keys(m.md, :time)
set_leftover_params!(m, p)

run(m)

@test !isnan(m[:AbatementCostsCO2, :tc_totalcost][10, 5])
@test !isnan(m[:AbatementCostsCH4, :tc_totalcost][10, 5])
@test !isnan(m[:AbatementCostsN2O, :tc_totalcost][10, 5])
@test !isnan(m[:AbatementCostsLin, :tc_totalcost][10, 5])

#compare output to validation data
tc_compare_co2=readpagedata(m, "test/validationdata/tc_totalcosts_co2.csv")
tc_compare_ch4=readpagedata(m, "test/validationdata/tc_totalcosts_ch4.csv")
tc_compare_n2o=readpagedata(m, "test/validationdata/tc_totalcosts_n2o.csv")
tc_compare_lin=readpagedata(m, "test/validationdata/tc_totalcosts_linear.csv")

zc_compare_co2=readpagedata(m, "test/validationdata/zc_zerocostemissionsCO2.csv")
zc_compare_ch4=readpagedata(m, "test/validationdata/zc_zerocostemissionsCH4.csv")
zc_compare_n2o=readpagedata(m, "test/validationdata/zc_zerocostemissionsN2O.csv")
zc_compare_lin=readpagedata(m, "test/validationdata/zc_zerocostemissionsLG.csv")

@test m[:AbatementCostsCO2, :tc_totalcost] ≈ tc_compare_co2 rtol=1e-2
@test m[:AbatementCostsCH4, :tc_totalcost] ≈ tc_compare_ch4 rtol=1e-2
@test m[:AbatementCostsN2O, :tc_totalcost] ≈ tc_compare_n2o rtol=1e-2
@test m[:AbatementCostsLin, :tc_totalcost] ≈ tc_compare_lin rtol=1e-2

@test m[:AbatementCostParametersCO2, :zc_zerocostemissions] ≈ zc_compare_co2 rtol=1e-2
@test m[:AbatementCostParametersCH4, :zc_zerocostemissions] ≈ zc_compare_ch4 rtol=1e-3
@test m[:AbatementCostParametersN2O, :zc_zerocostemissions] ≈ zc_compare_n2o rtol=1e-3
@test m[:AbatementCostParametersLin, :zc_zerocostemissions] ≈ zc_compare_lin rtol=1e-3


text:
# TODO: allow passing a rng as first parameter?

"""
    randobs(data, [n], [obsdim])

Pick a random observation or a batch of `n` random observations
from `data`.

The optional (keyword) parameter `obsdim` allows one to specify
which dimension denotes the observations. see `LearnBase.ObsDim`
for more detail.

For this function to work, the type of `data` must implement
[`nobs`](@ref) and [`getobs`](@ref).
"""
randobs(data; obsdim = default_obsdim(data)) =
    randobs(data, convert(LearnBase.ObsDimension,obsdim))

randobs(data, n; obsdim = default_obsdim(data)) =
    randobs(data, n, convert(LearnBase.ObsDimension,obsdim))

randobs(data, obsdim::Union{Tuple,ObsDimension}) =
    getobs(data, rand(1:nobs(data, obsdim)), obsdim)

randobs(data, n, obsdim::Union{Tuple,ObsDimension}) =
    getobs(data, rand(1:nobs(data, obsdim), n), obsdim)


text:
"""
    $(SIGNATURES)
Get the data for the specified entry, returns the data or Nothing.
"""
function getBigData(store::D, entry::E)::Union{Nothing, T} where {T, D <: AbstractDataStore{T}, E <: AbstractBigDataEntry}
    error("$(typeof(store)) doesn't override 'getData'.")
end

"""
    $(SIGNATURES)
Adds the data to the store with the given entry. The function will warn if the entry already
exists and will overwrite it.
"""
function addBigData!(store::D, entry::E, data::T)::Union{Nothing, T} where {T, D <: AbstractDataStore{T}, E <: AbstractBigDataEntry}
    error("$(typeof(store)) doesn't override 'addData!'.")
end

"""
    $(SIGNATURES)
Update the data in the store. The function will error and return nothing if
the entry does not exist.
"""
function updateBigData!(store::D, entry::E, data::T)::Union{Nothing, T} where {T, D <: AbstractDataStore{T}, E <: AbstractBigDataEntry}
    error("$(typeof(store)) doesn't override 'updateData!'.")
end

"""
    $(SIGNATURES)
Delete the data in the store for the given entry. The function will error and return nothing if
the entry does not exist.
"""
function deleteBigData!(store::D, entry::E)::Union{Nothing, T} where {T, D <: AbstractDataStore{T}, E <: AbstractBigDataEntry}
    error("$(typeof(store)) doesn't override 'deleteData!'.")
end

"""
    $(SIGNATURES)
List all entries in the data store.
"""
function listStoreEntries(store::D)::Vector{E} where {D <: AbstractDataStore, E <: AbstractBigDataEntry}
    error("$(typeof(store)) doesn't override 'listEntries'.")
end

"""
    $(SIGNATURES)
Copies all the entries from the source into the destination.
Can specify which entries to copy with the `sourceEntries` parameter.
Returns the list of copied entries.
"""
function copyStore(sourceStore::D1, destStore::D2; sourceEntries=listEntries(sourceStore))::Vector{E} where {T, D1 <: AbstractDataStore{T}, D2 <: AbstractDataStore{T}, E <: AbstractBigDataEntry}
    # Quick check
    destEntries = listStoreEntries(destStore)
    typeof(sourceEntries) != typeof(destEntries) && error("Can't copy stores, source has entries of type $(typeof(sourceEntries)), destination has entries of type $(typeof(destEntries)).")
    # Same source/destination check
    sourceStore == destStore && error("Can't specify same store for source and destination.")
    # Otherwise, continue
    for sourceEntry in sourceEntries
        addBigData!(destStore, deepcopy(sourceEntry), getBigData(sourceStore, sourceEntry))
    end
    return sourceEntries
end


text:
module How2Mix

# Unassuming code in an unassuming module...
module SubFoo

g() = rand()

function h()
    return g()
end

function f()
    x = rand()
    y = rand()
    return x + y + h()
end

end

using Mixtape
import Mixtape: CompilationContext, transform, optimize!, allow, show_after_inference,
                show_after_optimization, debug
using MacroTools

# 101: How2Mix
struct MyMix <: CompilationContext end

# A few little utility functions for working with Expr instances.
swap(e) = e
function swap(e::Expr)
    new = MacroTools.postwalk(e) do s
        isexpr(s, :call) || return s
        s.args[1] == Base.rand || return s
        return 4
    end
    return new
end

# This is pre-inference - you get to see a CodeInfoTools.Builder instance.
function transform(::MyMix, b)
    for (v, st) in b
        replace!(b, v, swap(st))
    end
    return b
end

function optimize!(::MyMix, ir)
    return ir
end

# MyMix will only transform functions which you explicitly allow.
# You can also greenlight modules.
allow(ctx::MyMix, m::Module) = m == SubFoo
show_after_inference(ctx::MyMix) = false
show_after_optimization(ctx::MyMix) = false
debug(ctx::MyMix) = false

# This loads up a call interface which will cache the result of the pipeline.
Mixtape.@load_call_interface()
@assert(call(MyMix(), SubFoo.f) == 12)
@assert(call(MyMix(), SubFoo.f) == 12)
@assert(SubFoo.f() != 12)

end # module


text:
# Tests for functions listed by file
#

# solverInterface.jl

a = randn(10,10)
a = Laplacians.forceLap(a)

Laplacians.forceLap(a)
Laplacians.forceLap(abs.(a))

a = wted_chimera(100,101);
la = lap(a);
la[1,1] += 1;
f = Laplacians.sddmWrapLap(approxchol_lap, la, verbose=true);
b = randn(100);
x = f(b);

a = grid2(20)
a2 = subsampleEdges(grid2(20),0.45)
f = approxchol_lap(a2)
its = [0]
f = approxchol_lap(a2, pcgIts=its, verbose=true)

# conditionNumber

support(a,a2)
support(a2,a)

conditionNumber(a, akpw(a), verbose=true)

f = approxchol_lap(a, maxits=1)
conditionNumber(a, f, verbose=true)

println("End of testByFile")


text:
using PtFEM

data = Dict(
  # GenericSolid(ndim, nst, nels, nn, nip, finite_element(nod, nodof), axisymmentrix)
  :struc_el => GenericSolid(2, 3, 6, 35, 9, Quadrilateral(9, 2), false),
  :properties => [1.0e6 0.3 0.0;],
  :g_coord => [
    0.0  0.0; 1.5  0.0; 3.0  0.0; 4.5  0.0; 6.0  0.0; 0.0 -1.5;
    1.5 -1.5; 3.0 -1.5; 4.5 -1.5; 6.0 -1.5; 0.0 -3.0; 1.5 -3.0;
    3.0 -3.0; 4.5 -3.0; 6.0 -3.0; 0.0 -4.5; 1.5 -4.5; 3.0 -4.5;
    4.5 -4.5; 6.0 -4.5; 0.0 -6.0; 1.5 -6.0; 3.0 -6.0; 4.5 -6.0;
    6.0 -6.0; 0.0 -7.5; 1.5 -7.5; 3.0 -7.5; 4.5 -7.5; 6.0 -7.5;
    0.0 -9.0; 1.5 -9.0; 3.0 -9.0; 4.5 -9.0; 6.0 -9.0],
  :g_num => [
    11,  6,  1,  2,  3,  8, 13, 12,  7,
    21, 16, 11, 12, 13, 18, 23, 22, 17,
    31, 26, 21, 22, 23, 28, 33, 32, 27,
    13,  8,  3,  4,  5, 10, 15, 14,  9,
    23, 18, 13, 14, 15, 20, 25, 24, 19,
    33, 28, 23, 24, 25, 30, 35, 34, 29
  ],
  :support => [
    ( 1, [0 1]), ( 5, [0 1]), ( 6, [0 1]), (10, [0 1]), (11, [0 1]), (15, [0 1]),
    (16, [0 1]), (20, [0 1]), (21, [0 1]), (25, [0 1]), (26, [0 1]), (30, [0 1]),
    (31, [0 0]), (32, [0 0]), (33, [0 0]), (34, [0 0]), (35, [0 0])
  ],
  :loaded_nodes => [
    ( 1, [0.0 -0.5]),
    ( 2, [0.0 -2.0]),
    ( 3, [0.0 -0.5])
  ]
)

data |> display
println()

@time m, dis_df, fm_df = p54(data)
println()

dis_df |> display
println()

println("\nThe integration point (nip = $(data[:struc_el].nip)) stresses are:\n")

fm_df |> display
println()


text:
function allocate_nzs(mat, bv)
	ptm = sparse(mat)
	ptm.nzval .= 1.0 # set all the nonzeros to 1
	res = sparse(ptm * bv)
	return res
end

struct SqrScaling{dim} <: AbstractScaling
	iWiW::Diagonal{Float64, Vector{Float64}}
	iW::Diagonal{Float64, Vector{Float64}}
	l::Vector{Float64}

	us :: Vector{Vector{Float64}}
	vs :: Vector{Vector{Float64}}
	#per cone
	mu::Vector{Float64}
	wbs::Vector{Float64} # sqrt(s/z) for POC, wb for SOC

	# workspaces for scaling functions
	ws_1::Vector{Vector{Float64}}
	ws_2::Vector{Vector{Float64}}
	ws_3::Vector{Vector{Float64}}
	ws_4::Vector{Vector{Float64}}
	ws_5::Vector{SparseMatrixCSC{Float64, Int}}
	ws_6::Vector{SuiteSparse.CHOLMOD.Sparse{Float64}}
	function SqrScaling{T}(p::Problem) where T
		return SqrScaling(Diagonal(zeros(p.k)), Diagonal(zeros(p.k)), zeros(p.k), p.cones, p.n, p.G)
	end
	function SqrScaling(iWiW, iW,l,cones, id, G)
		psz = cones[end].offs + conedim(cones[end])
		ws5 = [allocate_nzs(G', collect(if cone.offs <= i && i <= cone.offs+conedim(cone) 1.0 else 0.0 end for i in 1:psz)) for cone in cones]
		ws6 = [SuiteSparse.CHOLMOD.Sparse(exv) for exv in ws5]
		return new{id}(iWiW, iW, l,
			[zeros(psz) for cone in cones],
			[zeros(psz) for cone in cones],
			zeros(length(cones)), zeros(psz),
			[alloc_ws1(cone, psz) for cone in cones],
			[alloc_ws2(cone, psz) for cone in cones],
			[zeros(id) for cone in cones],
			[zeros(conedim(cone)) for cone in cones],
			ws5, ws6)
	end
end

function alloc_ws1(cone::POC{dim}, total) where dim return zeros(0) end
function alloc_ws2(cone::POC{dim}, total) where dim return zeros(0) end
function alloc_ws1(cone::SOC{dim}, total) where dim return zeros(total) end
function alloc_ws2(cone::SOC{dim}, total) where dim return zeros(total) end

function compute_scaling(cone::POC{dim}, cind, scaling::SqrScaling, s, z) where dim
	for i = 1:dim
		ii = cti(cone, i)
		scaling.iWiW[ii,ii] = z[ii] / s[ii]
		scaling.iW[ii,ii] = sqrt(z[ii] / s[ii])
		scaling.l[ii] = sqrt(s[ii] * z[ii])
		scaling.wbs[ii] = sqrt(s[ii] / z[ii])
	end
end

function modify_factors!(cone::POC{dim}, cind, factor::SuiteSparse.CHOLMOD.Factor,
						 scaling::SqrScaling, g::SparseMatrixCSC) where dim
	# we don't need to do anything here
end

function compute_scaling(c::SOC{dim}, cind, scaling::SqrScaling, s, z) where dim
	function cone_prod(vect)
		bsum = vect[1] * vect[1]
		for i=2:dim
			bsum -= vect[i] * vect[i]
		end
		return bsum
	end
	sbk = scaling.ws_1[cind]
	zbk = scaling.ws_2[cind]
	sbk[1:dim] .= @view s[c.offs+1:c.offs+dim]
	zbk[1:dim] .= @view z[c.offs+1:c.offs+dim]
	sprod, zprod = cone_prod(sbk), cone_prod(zbk)
	sbk .*= 1/sqrt(sprod)
	zbk .*= 1/sqrt(zprod)

	nsum = 0.0
	for i=1:dim
		nsum += zbk[i] * sbk[i]
	end
	gamma = sqrt((1 + nsum)/2)

	# compute scaling point
	wb = scaling.ws_4[cind]
	wb[1] = (sbk[1] + zbk[1]) / (2*gamma)
	for i=2:dim
		wb[i] = (sbk[i] - zbk[i]) / (2*gamma)
	end

	mu = sqrt(sqrt(sprod/zprod))
	scaling.mu[cind] = mu

	# compute scaling matrix
	inusq = 1/sqrt(sprod/zprod)
	inu = 1/sqrt(sqrt(sprod/zprod))
	wb0 = wb[1]
	wb1 = @view wb[2:dim]
	wb1sq = 0.0 #wb1sq = sum(wb1 .* wb1)
	for i=1:dim-1
		wb1sq += wb1[i]*wb1[i]
	end
	cv = -(1 + wb0 + wb1sq/(1+wb0))
	d = 1 + 2/(1+wb0) + wb1sq/((1+wb0)*(1+wb0))
	a = (wb0*wb0 + wb1sq - cv*cv*wb1sq/(1+d*wb1sq))/2
	u0 = sqrt(wb0*wb0 + wb1sq - a)
	u1 = cv/u0
	v1 = sqrt(cv*cv/(u0*u0) - d)

	scaling.iWiW.diag[cti(c, 1)] = a*inusq
	scaling.iW.diag[cti(c, 1)] = sqrt(abs(a*inusq))
	for i=2:dim
		scaling.iWiW.diag[cti(c, i)] = inusq
		scaling.iW.diag[cti(c, i)] = sqrt(inusq)
	end
	scu, scv = scaling.us[cind], scaling.vs[cind]
	# u = inu .* vcat(u0, u1 .* wb1); v = inu .* vcat(0, v1 .* wb1)
	scu[c.offs+1] = inu * u0
	scv[c.offs+1] = 0.0
	for i=2:dim
		wbv = inu * wb1[i-1]
		scu[c.offs+i] = u1 * wbv
		scv[c.offs+i] = v1 * wbv
	end

	# compute scaling variable
	ziv,siv = zbk[1], sbk[1]
	tmv1 = sqrt(sqrt(sprod)*sqrt(zprod))
	mult = tmv1/(ziv + siv + 2*gamma)
	for i=2:dim
		scaling.l[cti(c, i)] = (sbk[i]*(gamma + ziv) + zbk[i]*(gamma + siv))*mult
	end
	(@view scaling.wbs[cti(c,1):cti(c,conedim(c))]) .= wb
	scaling.l[cti(c,1)] = gamma*tmv1
end

# this removes every nonzero in the sparse vector dest first, then re-fills it. This retains the # nonzeros
function clear_reset!(dest::SparseMatrixCSC, src::Vector)
	nnz = 0
	cnz = 1
	dest.nzval .= 0.0
	# this is all 1d so we don't have to touch colptr
	for i=1:length(src)
		if src[i] != 0.0
			if cnz > length(dest.nzval)
				throw(ArgumentError("too many nonzeros in source; limit $(length(dest.nzval))"))
			end
			dest.rowval[cnz] = i
			dest.nzval[cnz] = src[i]
			cnz += 1
			nnz += 1
		end
	end
end

function modify_factors!(cone::SOC{dim}, cind, factor::SuiteSparse.CHOLMOD.Factor,
						 scaling::SqrScaling{odim}, g::SparseMatrixCSC) where {dim, odim}
	perm = SuiteSparse.CHOLMOD.get_perm(factor) # this is allocating but... meh. The permutation might be updated.

	temp = scaling.ws_3[cind]
	mul!(temp, g', scaling.us[cind])
	permute!(temp, perm)
	clear_reset!(scaling.ws_5[cind], temp)
	copyto!(scaling.ws_6[cind], scaling.ws_5[cind])
	SuiteSparse.CHOLMOD.lowrankupdowndate!(factor, scaling.ws_6[cind], Cint(1))
	mul!(temp, g', scaling.vs[cind])
	permute!(temp, perm)
	clear_reset!(scaling.ws_5[cind], temp)
	copyto!(scaling.ws_6[cind], scaling.ws_5[cind])
	SuiteSparse.CHOLMOD.lowrankupdowndate!(factor, scaling.ws_6[cind], Cint(0))
end

@unroll function compute_scaling(cones::Tuple{Vararg{Cone}}, scaling::SqrScaling, s, z)
	# assume that cones is in the order POC ... POC SOC ... SOC
	i=1
	@unroll for cone in cones
		compute_scaling(cone, i, scaling, s, z)
		i += 1
	end
	return scaling
end

@unroll function modify_factors!(cones::Tuple{Vararg{Cone}}, factor::SuiteSparse.CHOLMOD.Factor,
						 scaling::SqrScaling, g::SparseMatrixCSC) where dim
	i=1
	@unroll for cone in cones
		modify_factors!(cone, i, factor, scaling, g)
		i += 1
	end
end

@unroll function compute_full_scaling(cones::Tuple{Vararg{Cone}}, scaling::SqrScaling)
	i=1
	oiWiW = zeros(MMatrix{length(scaling.l), length(scaling.l)})
	@unroll for cone in cones
		compute_full_scaling(cone, i, scaling, oiWiW)
		i += 1
	end
	return oiWiW
end

function compute_full_scaling(cone::SOC{dim}, cind, scaling::SqrScaling, output) where dim
	output[cone.offs+1:cone.offs+dim, cone.offs+1:cone.offs+dim] .= scaling.iWiW[cone.offs+1:cone.offs+dim, cone.offs+1:cone.offs+dim]
	BLAS.gemm!('N', 'T', 1.0, scaling.us[cind], scaling.us[cind], 1.0, output)
	BLAS.gemm!('N', 'T', -1.0, scaling.vs[cind], scaling.vs[cind], 1.0, output)
end

function compute_full_scaling(cone::POC{dim}, cind, scaling::SqrScaling, output) where dim
	output[cone.offs+1:cone.offs+dim, cone.offs+1:cone.offs+dim] .= scaling.iWiW[cone.offs+1:cone.offs+dim, cone.offs+1:cone.offs+dim]
end


text:
include("precompile_funcs.jl")

using OptimDash

using .PrecompileFuncs
PrecompileFuncs.precompile_optim()

app = OptimDash.make_app()

port = haskey(ENV, "PORT") ? parse(Int64, ENV["PORT"]) : 8050

OptimDash.Dash.run_server(app, "0.0.0.0", port)


text:
using .PyCall

function GymEnv(name::String; seed::Union{Int,Nothing} = nothing)
    if !PyCall.pyexists("gym")
        error(
            "Cannot import module 'gym'.\n\nIf you did not yet install it, try running\n`ReinforcementLearningEnvironments.install_gym()`\n",
        )
    end
    gym = pyimport_conda("gym", "gym")
    if PyCall.pyexists("d4rl")
        pyimport("d4rl")
    end
    pyenv = try
        gym.make(name)
    catch e
        error(
            "Gym environment $name not found.\n\nRun `ReinforcementLearningEnvironments.list_gym_env_names()` to find supported environments.\n",
        )
    end
    if seed !== nothing
        pyenv.seed(seed)
    end
    obs_space = space_transform(pyenv.observation_space)
    act_space = space_transform(pyenv.action_space)
    obs_type = if obs_space isa Space{<:Union{Array{<:Interval},Array{<:ZeroTo}}}
        PyArray
    elseif obs_space isa Interval
        Float64
    elseif obs_space isa ZeroTo
        Int
    elseif obs_space isa Space{<:Tuple}
        PyVector
    elseif obs_space isa Space{<:Dict}
        PyDict
    else
        error("don't know how to get the observation type from observation space of $obs_space")
    end
    env = GymEnv{obs_type,typeof(act_space),typeof(obs_space),typeof(pyenv)}(
        pyenv,
        obs_space,
        act_space,
        PyNULL(),
    )
    reset!(env) # reset immediately to init env.state
    env
end

Base.nameof(env::GymEnv) = env.pyenv.__class__.__name__

function Base.copy(env::GymEnv)
    @warn "clone method is not exposed in GymEnv"
    env
end

function (env::GymEnv{T})(action) where {T}
    if env.action_space isa Tuple
        action = Tuple(action)
    end
    pycall!(env.state, env.pyenv.step, PyObject, action)
    nothing
end

function RLBase.reset!(env::GymEnv)
    pycall!(env.state, env.pyenv.reset, PyObject)
    nothing
end

RLBase.action_space(env::GymEnv) = env.action_space
RLBase.state_space(env::GymEnv) = env.observation_space

function RLBase.reward(env::GymEnv{T}) where {T}
    if pyisinstance(env.state, PyCall.@pyglobalobj :PyTuple_Type) && length(env.state) == 4
        obs, reward, isdone, info = convert(Tuple{T,Float64,Bool,PyDict}, env.state)
        reward
    else
        0.0
    end
end

function RLBase.is_terminated(env::GymEnv{T}) where {T}
    if pyisinstance(env.state, PyCall.@pyglobalobj :PyTuple_Type) && length(env.state) == 4
        obs, reward, isdone, info = convert(Tuple{T,Float64,Bool,PyDict}, env.state)
        isdone
    else
        false
    end
end

function RLBase.state(env::GymEnv{T}) where {T}
    if pyisinstance(env.state, PyCall.@pyglobalobj :PyTuple_Type) && length(env.state) == 4
        obs, reward, isdone, info = convert(Tuple{T,Float64,Bool,PyDict}, env.state)
        obs
    else
        convert(T, env.state)
    end
end

Random.seed!(env::GymEnv, s) = env.pyenv.seed(s)

# Base.display(env::GymEnv) = env.pyenv.render()

###
### utils
###
function space_transform(s::PyObject)
    spacetype = s.__class__.__name__
    if spacetype == "Box"
        Space(ClosedInterval.(s.low, s.high))
    elseif spacetype == "Discrete"  # for GymEnv("CliffWalking-v0"), `s.n` is of type PyObject (numpy.int64)
        Space(0:(py"int($s.n)"-1))
    elseif spacetype == "MultiBinary"
        Space(Bool, s.n)
    elseif spacetype == "MultiDiscrete"
        Space(map(x -> 0:x-one(typeof(x)), s.nvec))
    elseif spacetype == "Tuple"
        Tuple(space_transform(x) for x in s.spaces)
    elseif spacetype == "Dict"
        Dict((k => space_transform(v) for (k, v) in s.spaces)...)
    else
        error("Don't know how to convert Gym Space of class [$(spacetype)]")
    end
end

function list_gym_env_names(;
    modules = [
        "gym.envs.algorithmic",
        "gym.envs.box2d",
        "gym.envs.classic_control",
        "gym.envs.mujoco",
        "gym.envs.mujoco.ant_v3",
        "gym.envs.mujoco.half_cheetah_v3",
        "gym.envs.mujoco.hopper_v3",
        "gym.envs.mujoco.humanoid_v3",
        "gym.envs.mujoco.swimmer_v3",
        "gym.envs.mujoco.walker2d_v3",
        "gym.envs.robotics",
        "gym.envs.toy_text",
        "gym.envs.unittest",
        "d4rl.pointmaze",
        "d4rl.hand_manipulation_suite",
        "d4rl.gym_mujoco.gym_envs",
        "d4rl.locomotion.ant",
        "d4rl.gym_bullet.gym_envs",
        "d4rl.pointmaze_bullet.bullet_maze", # yet to include flow and carla
    ],
)
    if PyCall.pyexists("d4rl")
        pyimport("d4rl")
    end
    gym = pyimport("gym")
    [x.id for x in values(gym.envs.registry) if split(x.entry_point, ':')[1] in modules]
end

"""
    install_gym(; packages = ["gym", "pybullet"])
"""
function install_gym(; packages = ["gym", "pybullet"])
    # Use eventual proxy info
    proxy_arg = String[]
    if haskey(ENV, "http_proxy")
        push!(proxy_arg, "--proxy")
        push!(proxy_arg, ENV["http_proxy"])
    end
    # Import pip
    if !PyCall.pyexists("pip")
        # If it is not found, install it
        println("Pip not found on your system. Downloading it.")
        get_pip = joinpath(dirname(@__FILE__), "get-pip.py")
        download("https://bootstrap.pypa.io/get-pip.py", get_pip)
        run(`$(PyCall.python) $(proxy_arg) $get_pip --user`)
    end
    println("Installing required python packages using pip")
    run(`$(PyCall.python) $(proxy_arg) -m pip install --user --upgrade pip setuptools`)
    run(`$(PyCall.python) $(proxy_arg) -m pip install --user $(packages)`)
end


text:
using Test
using STAMP
using BioSequences
import BioGenerics

record_str = """
DE gps_test_1_m0 7.09 1 k7_c1526
1 1297 0 53 176 A
2 13 2 0 1511 T
3 5 0 1394 126 G
4 0 1171 73 282 C
5 1055 6 2 463 A
6 1281 22 201 22 A
7 1503 4 7 11 A
XX
"""

@testset "STAMP" begin
    @testset "Record" begin

        # Initialise empty record.
        record = STAMP.Record()
        @test !BioGenerics.isfilled(record)
        @test STAMP.header(record) === missing

        record = STAMP.Record(record_str)

        # Initialise from string.
        @test BioGenerics.isfilled(record)
        @test STAMP.header(record) == "gps_test_1_m0 7.09 1 k7_c1526"
        @test STAMP.frequencies(record) == [
                [1297 13 5 0 1055 1281 1503];
                [0 2 0 1171 6 22 4];
                [53 0 1394 73 2 201 7];
                [176 1511 126 282 463 22 11]
            ]
        @test STAMP.sequence(record) == dna"ATGCAAA"

        # Check negative.
        record = STAMP.Record("""
            DE test 9.59 -19 k8_c740
            1 7 76 657 0 G
            XX
            """)
        @test STAMP.header(record) == "test 9.59 -19 k8_c740"

    end # testset

    @testset "Reader" begin

        # Read from buffer.
        reader = STAMP.Reader(IOBuffer(record_str))

        record = STAMP.Record()
        @test read!(reader, record) === record
        @test STAMP.sequence(record) == dna"ATGCAAA"

        # Read from file.
        reader = STAMP.Reader(open("gps_test_1.all.PFM.txt", "r"))

        record = STAMP.Record()
        @test read!(reader, record) === record
        @test STAMP.sequence(record) == dna"ATGCAAA"
        @test STAMP.sequence(read!(reader, record)) == dna"TTGTTATG"
        @test STAMP.sequence(read!(reader, record)) == dna"CCTTTT"
        @test STAMP.sequence(read!(reader, record)) == dna"TGTTATGCA"
        @test STAMP.sequence(read!(reader, record)) == dna"AACAAAG"
        @test STAMP.sequence(read!(reader, record)) == dna"TTCTTTG"
        @test STAMP.sequence(read!(reader, record)) == dna"ACAAAAGA"
        @test STAMP.sequence(read!(reader, record)) == dna"CCCACCC"
        @test STAMP.sequence(read!(reader, record)) == dna"CCTGCT"
        @test STAMP.sequence(read!(reader, record)) == dna"CCTTCC"
        @test STAMP.sequence(read!(reader, record)) == dna"AGCAGGTGGC"
        @test STAMP.sequence(read!(reader, record)) == dna"CATTAGAATGGA"
        @test STAMP.sequence(read!(reader, record)) == dna"CCTGCTGG"
        @test STAMP.sequence(read!(reader, record)) == dna"CTGCCCCCTGCTG"

        @test_throws EOFError read!(reader, record)
    end # testset

end


text:
# """
#     struct Ellipse{T,TR} <: AbstractLinePrimitive{T}
#
# * `r::TR`:
#     * TR = Real -> Circle (a = b = r)
#     * TR = (Real, Real) -> Circular Annulus (r_in = r[1], r_out = r[2])
#     * TR = ((Real,), (Real,)) -> Ellipse (a = r[1][1], b = r[2][1])
#     * TR = ((Real, Real),(Real, Real)) -> Elliptical Annulus \n(a_in = r[1][1], a_out = r[1][2], b_in = r[2][1], b_out = r[2][2])
#     * Not all are implemented yet
#
# * `φ::TP`:
#     * TP = Nothing <-> Full in φ
#     * ...
# """
@with_kw struct Ellipse{T,TR,TP} <: AbstractLinePrimitive{T}
    r::TR = 1
    φ::TP = nothing

    origin::CartesianPoint{T} = zero(CartesianPoint{T})
    rotation::SMatrix{3,3,T,9} = one(SMatrix{3, 3, T, 9})
end

const Circle{T} = Ellipse{T,T,Nothing}
const PartialCircle{T} = Ellipse{T,T,Tuple{T,T}}

extremum(e::Ellipse{T,T}) where {T} = e.r
extremum(e::Ellipse{T,Tuple{T,T}}) where {T} = max(e.r...)

function sample(e::Circle{T}; n = 4)::Vector{CartesianPoint{T}} where {T}
    φs = range(T(0), step = T(2π) / n, length = n)
    pts = Vector{CartesianPoint{T}}(undef, n)
    for i in eachindex(pts)
        pts[i] = _transform_into_global_coordinate_system(CartesianPoint(CylindricalPoint{T}(e.r, φs[i], zero(T))), e)
    end
    pts
end
function sample(e::PartialCircle{T}; n = 2)::Vector{CartesianPoint{T}} where {T}
    φs = range(e.φ[1], stop = e.φ[1], length = n)
    pts = Vector{CartesianPoint{T}}(undef, n)
    for i in eachindex(pts)
        pts[i] = _transform_into_global_coordinate_system(CartesianPoint(CylindricalPoint{T}(e.r, φs[i], zero(T))), e)
    end
    pts
end

# function Arc(; r = 0, center = PlanarPoint(0,0), αMin = 0, αMax = 2π)
#     T = float(promote_type(typeof.((r, αMin, αMax))..., eltype(center)))
#     α = mod(T(αMax    ) - T(αMin), T(2π)) == 0 ? nothing : T(αMin)..T(αMax)
#     Arc(T, T(r), PlanarPoint{T}(center), α)
# end

# Arc(r, center, αMin, αMax) = Arc(; r = r, center = center, αMin = αMin, αMax = αMax)

# Circle(r::T, center::PlanarPoint{T}) where {T} = Arc(T, r, center, nothing)
# Circle(a::Arc{T}) where {T} = Arc(T, a.r, a.center, nothing)

# get_α_at_u_v(a::Arc{T}, u::Real, v::Real) where {T} = mod(atan(v - a.center.v, u - a.center.u), 2π) #u,v are planar coordinates

# get_α_limits(a::Arc{T, Nothing}) where {T} = (T(0), T(2π), true)
# get_α_limits(a::Arc{T, <:AbstractInterval}) where {T} = (a.α.left, a.α.right, false)

# distance_to_line(point::PlanarPoint{T}, a::Arc{T, Nothing}) where {T} = abs(norm(point - a.center) - a.r)

# function distance_to_line(point::PlanarPoint{T}, a::Arc{T, <:AbstractInterval})::T where {T}
#     αMin::T, αMax::T, _ = get_α_limits(a)
#     p_t = point - a.center
#     α = atan(p_t.v, p_t.u)
#     if _in_angular_interval_closed(α, a.α)
#         return abs(norm(p_t) - a.r)
#     else
#         sαNear, cαNear = sincos(_φNear(α, αMin, αMax))
#         return norm(p_t - PlanarPoint{T}(a.r*cαNear, a.r*sαNear))
#     end
# end


text:
pop1 = UncertainValue(
        [3.0, UncertainValue(Normal, 0, 1),
        UncertainValue(Gamma, 2, 3),
        UncertainValue(Uniform, rand(1000))],
        [0.5, 0.5, 0.5, 0.5]
    )

pop2 = UncertainValue(1:10 |> collect, rand(10))
pop3 = UncertainValue(1.0:10.0 |> collect, Weights(rand(10)))

# Uncertain population consisting of uncertain populations and other stuff
pop4 = UncertainValue([pop1, pop2], [0.1, 0.5])
pop5 = UncertainValue([pop1, pop2, 2, UncertainValue(Normal, -2, 3)], Weights(rand(4)));

uncertain_scalar_populations = [pop1, pop2, pop3, pop4, pop5];

T = Union{Nothing, ConstrainedUncertainScalarPopulation}

@test constrain(pop1, TruncateMinimum(0.2)) isa T
@test constrain(pop1, TruncateMaximum(3.0)) isa T
@test constrain(pop1, TruncateRange(0.0, 3.0)) isa T
@test constrain(pop1, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop1, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop1, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop1, TruncateStd(1)) isa T
@test constrain(pop1, TruncateStd(0.5)) isa T

# pop2.values = [1, 2, 3]
@test constrain(pop2, TruncateMinimum(0.2)) isa T
@test constrain(pop2, TruncateMaximum(2.5)) isa T
@test constrain(pop2, TruncateRange(0.5, 2.5)) isa T
@test constrain(pop2, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop2, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop2, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop2, TruncateStd(1)) isa T
@test constrain(pop2, TruncateStd(0.5)) isa T

# pop3.values = [1.0, 2.0, 3.0]
@test constrain(pop3, TruncateMinimum(0.2)) isa T
@test constrain(pop3, TruncateMaximum(2.5)) isa T
@test constrain(pop3, TruncateRange(0.5, 2.5)) isa T
@test constrain(pop3, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop3, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop3, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop3, TruncateStd(1)) isa T
@test constrain(pop3, TruncateStd(0.5)) isa T

@test constrain(pop4, TruncateMinimum(-10)) isa T
@test constrain(pop4, TruncateMaximum(10)) isa T
@test constrain(pop4, TruncateRange(-10, 10)) isa T
@test constrain(pop4, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop4, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop4, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop4, TruncateStd(1)) isa T
@test constrain(pop4, TruncateStd(0.5)) isa T

@test constrain(pop5, TruncateMinimum(-10)) isa T
@test constrain(pop5, TruncateMaximum(10)) isa T
@test constrain(pop5, TruncateRange(-10, 10)) isa T
@test constrain(pop5, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop5, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop5, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop5, TruncateStd(1)) isa T
@test constrain(pop5, TruncateStd(0.5)) isa T

# Some subpopulations consisting of both scalar values and distributions
subpop1_members = [UncertainValue(Normal, 0, 1), UncertainValue(Uniform, -2, 2), -5]
subpop2_members = [
    UncertainValue(Normal, -2, 1),
    UncertainValue(Uniform, -6, -1),
    -3,
    UncertainValue(Gamma, 1, 0.4)]

# Define the probabilities of sampling the different population members within the
# subpopulations. Weights are normalised, so we can input any numbers here indicating
# relative importance
subpop1_probs = [1, 2, 1]
subpop2_probs = [0.1, 0.2, 0.3, 0.1]

pop1 = UncertainValue(subpop1_members, subpop1_probs)
pop2 = UncertainValue(subpop2_members, subpop2_probs)

# Define the probabilities of sampling the two subpopulations in the overall population.
pop_probs = [0.3, 0.7]

# Construct overall population
pop_mixed = UncertainValue([pop1, pop2], pop_probs)

@test constrain(pop_mixed, TruncateMinimum(1)) isa T
@test constrain(pop_mixed, TruncateMaximum(2)) isa T
@test constrain(pop_mixed, TruncateRange(1, 2)) isa T
@test constrain(pop_mixed, TruncateQuantiles(0.1, 0.9)) isa T
@test constrain(pop_mixed, TruncateLowerQuantile(0.1)) isa T
@test constrain(pop_mixed, TruncateUpperQuantile(0.9)) isa T
@test constrain(pop_mixed, TruncateStd(1)) isa T
@test constrain(pop_mixed, TruncateStd(0.5)) isa T

# Truncation
#------------------------------------------
pop1 = UncertainValue(
        [3.0, UncertainValue(Normal, 0, 1),
        UncertainValue(Gamma, 2, 3),
        UncertainValue(Uniform, rand(1000))],
        [0.5, 0.5, 0.5, 0.5]
    )

pop2 = UncertainValue([1, 2, 3], rand(3))
pop3 = UncertainValue([1.0, 2.0, 3.0], Weights(rand(3)))

# Uncertain population consisting of uncertain populations and other stuff
pop4 = UncertainValue([pop1, pop2], [0.1, 0.5])
pop5 = UncertainValue([pop1, pop2, 2, UncertainValue(Normal, -2, 3)], Weights(rand(4)));

n = 2000
p = pop1
sb = resample(p, n)

v_min, v_max = -1, 1
r_min, r_max = -1, 1
c_max = TruncateMaximum(v_max)
c_min = TruncateMinimum(v_min)
c_range = TruncateRange(r_min, r_max)
ql, qh = 0.2, 0.8
c_ql = TruncateLowerQuantile(ql)
c_qh = TruncateUpperQuantile(qh)
c_qs = TruncateQuantiles(ql, qh)
c_std = TruncateStd(1.0)

@test truncate(p, c_max) isa T
@test truncate(p, c_min) isa T
@test truncate(p, c_range) isa T
@test truncate(p, c_ql) isa T
@test truncate(p, c_qh) isa T
@test truncate(p, c_qs) isa T
@test truncate(p, c_std) isa T

n = 2000
p = pop2
sb = resample(p, n)

v_min, v_max = -1, 1
r_min, r_max = -1, 1
c_max = TruncateMaximum(v_max)
c_min = TruncateMinimum(v_min)
c_range = TruncateRange(r_min, r_max)
ql, qh = 0.2, 0.8
c_ql = TruncateLowerQuantile(ql)
c_qh = TruncateUpperQuantile(qh)
c_qs = TruncateQuantiles(ql, qh)
c_std = TruncateStd(1.0)

@test truncate(p, c_max) isa T
@test truncate(p, c_min) isa T
@test truncate(p, c_range) isa T
@test truncate(p, c_ql) isa T
@test truncate(p, c_qh) isa T
@test truncate(p, c_qs) isa T
@test truncate(p, c_std) isa T

n = 2000
p = pop3
sb = resample(p, n)

v_min, v_max = -1, 1
r_min, r_max = -1, 1
c_max = TruncateMaximum(v_max)
c_min = TruncateMinimum(v_min)
c_range = TruncateRange(r_min, r_max)
ql, qh = 0.2, 0.8
c_ql = TruncateLowerQuantile(ql)
c_qh = TruncateUpperQuantile(qh)
c_qs = TruncateQuantiles(ql, qh)
c_std = TruncateStd(1.0)

@test truncate(p, c_max) isa T
@test truncate(p, c_min) isa T
@test truncate(p, c_range) isa T
@test truncate(p, c_ql) isa T
@test truncate(p, c_qh) isa T
@test truncate(p, c_qs) isa T
@test truncate(p, c_std) isa T

n = 2000
p = pop4
sb = resample(p, n)

v_min, v_max = -2, 2
r_min, r_max = -2, 2
c_max = TruncateMaximum(v_max)
c_min = TruncateMinimum(v_min)
c_range = TruncateRange(r_min, r_max)
ql, qh = 0.2, 0.8
c_ql = TruncateLowerQuantile(ql)
c_qh = TruncateUpperQuantile(qh)
c_qs = TruncateQuantiles(ql, qh)
c_std = TruncateStd(1.0)

@test truncate(p, c_max) isa T
@test truncate(p, c_min) isa T
@test truncate(p, c_range) isa T
@test truncate(p, c_ql) isa T
@test truncate(p, c_qh) isa T
@test truncate(p, c_qs) isa T
@test truncate(p, c_std) isa T

n = 2000
p = pop5
sb = resample(p, n)

v_min, v_max = -2, 2
r_min, r_max = -2, 2
c_max = TruncateMaximum(v_max)
c_min = TruncateMinimum(v_min)
c_range = TruncateRange(r_min, r_max)
ql, qh = 0.2, 0.8
c_ql = TruncateLowerQuantile(ql)
c_qh = TruncateUpperQuantile(qh)
c_qs = TruncateQuantiles(ql, qh)
c_std = TruncateStd(1.0)

@test truncate(p, c_max) isa T
@test truncate(p, c_min) isa T
@test truncate(p, c_range) isa T
@test truncate(p, c_ql) isa T
@test truncate(p, c_qh) isa T
@test truncate(p, c_qs) isa T
@test truncate(p, c_std) isa T

pop1 = UncertainValue(
        [3.0, UncertainValue(Normal, 0, 1),
        UncertainValue(Gamma, 2, 3),
        UncertainValue(Uniform, rand(10000))],
        [0.5, 0.5, 0.5, 0.5]
    )

pop2 = UncertainValue([1, 2, 3], rand(3))
pop3 = UncertainValue([1.0, 2.0, 3.0], Weights(rand(3)))

# Uncertain population consisting of uncertain populations and other stuff
pop4 = UncertainValue([pop1, pop2], [0.1, 0.5])
pop5 = UncertainValue([pop1, pop2, 2, UncertainValue(Normal, -2, 3)], Weights(rand(4)));

uncertain_scalar_populations = [pop1, pop2, pop3, pop4, pop5];

@test constrain(pop1, TruncateMinimum(0.2)) isa T
@test constrain(pop1, TruncateMaximum(3.0)) isa T
@test constrain(pop1, TruncateRange(0.0, 3.0)) isa T
@test constrain(pop1, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop1, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop1, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop1, TruncateStd(1)) isa T

# pop2.values = [1, 2, 3]
@test constrain(pop2, TruncateMinimum(0.2)) isa T
@test constrain(pop2, TruncateMaximum(2.5)) isa T
@test constrain(pop2, TruncateRange(0.5, 2.5)) isa T
@test constrain(pop2, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop2, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop2, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop2, TruncateStd(1)) isa T

# pop3.values = [1.0, 2.0, 3.0]
@test constrain(pop3, TruncateMinimum(0.2)) isa T
@test constrain(pop3, TruncateMaximum(2.5)) isa T
@test constrain(pop3, TruncateRange(0.5, 2.5)) isa T
@test constrain(pop3, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop3, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop3, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop3, TruncateStd(1)) isa T

@test constrain(pop4, TruncateMinimum(-10)) isa T
@test constrain(pop4, TruncateMaximum(10)) isa T
@test constrain(pop4, TruncateRange(-10, 10)) isa T
@test constrain(pop4, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop4, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop4, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop4, TruncateStd(1)) isa T

@test constrain(pop5, TruncateMinimum(-10)) isa T
@test constrain(pop5, TruncateMaximum(10)) isa T
@test constrain(pop5, TruncateRange(-10, 10)) isa T
@test constrain(pop5, TruncateLowerQuantile(0.2)) isa T
@test constrain(pop5, TruncateUpperQuantile(0.8)) isa T
@test constrain(pop5, TruncateQuantiles(0.2, 0.8)) isa T
@test constrain(pop5, TruncateStd(1)) isa T

# Some subpopulations consisting of both scalar values and distributions
subpop1_members = [UncertainValue(Normal, 0, 1), UncertainValue(Uniform, -2, 2), -5]
subpop2_members = [
    UncertainValue(Normal, -2, 1),
    UncertainValue(Uniform, -6, -1),
    -3,
    UncertainValue(Gamma, 1, 0.4)]

# Define the probabilities of sampling the different population members within the
# subpopulations. Weights are normalised, so we can input any numbers here indicating
# relative importance
subpop1_probs = [1, 2, 1]
subpop2_probs = [0.1, 0.2, 0.3, 0.1]
pop1 = UncertainValue(subpop1_members, subpop1_probs)
pop2 = UncertainValue(subpop2_members, subpop2_probs)

# Define the probabilities of sampling the two subpopulations in the overall population.
pop_probs = [0.3, 0.7]

# Construct overall population
pop_mixed = UncertainValue([pop1, pop2], pop_probs)

@test constrain(pop_mixed, TruncateMinimum(1)) isa T
@test constrain(pop_mixed, TruncateMaximum(2)) isa T
@test constrain(pop_mixed, TruncateRange(1, 2)) isa T
@test constrain(pop_mixed, TruncateQuantiles(0.1, 0.9)) isa T
@test constrain(pop_mixed, TruncateLowerQuantile(0.1)) isa T
@test constrain(pop_mixed, TruncateUpperQuantile(0.9)) isa T
@test constrain(pop_mixed, TruncateStd(1)) isa T


text:
# XXX things to possibly speed this up
# replace recursion with while loop
# cache simulation results

function create_policy{S}(solver::AbstractPOMCPSolver, pomdp::POMDPs.POMDP{S})
    if isa(solver.node_belief_updater, DefaultReinvigoratorStub)
        node_belief_updater = DeadReinvigorator{S}()
    else
        node_belief_updater = solver.node_belief_updater
    end
    return POMCPPlanner(pomdp, solver, node_belief_updater,
                        convert_estimator(solver.estimate_value, solver, pomdp),
                        Nullable{Any}())
end

function action(policy::POMCPPlanner, belief, a=nothing)
    try
        a = search(policy, make_node(policy, belief), policy.solver.tree_queries)
    catch ex
        a = default_action(policy.solver.default_action, belief, ex)
    end
    return a
end

"""
    solve(solver::POMCPSolver, pomdp::POMDPs.POMDP)

Simply return a properly constructed POMCPPlanner object.
"""
function solve{S}(solver::AbstractPOMCPSolver, pomdp::POMDPs.POMDP{S})
    create_policy(solver, pomdp)
end

solve(solver::AbstractPOMCPSolver, pomdp::POMDPs.POMDP, dummy_policy::Any) = solve(solver, pomdp)

"""
    function search(pomcp::POMCPPlanner, b::BeliefNode, tree_queries)

Search the tree for the next best move.
"""
function search(pomcp::POMCPPlanner, b::BeliefNode, tree_queries)
    # XXX hack
    pomcp._tree_ref = b
    # end hack

    all_terminal = true
    for i in 1:tree_queries
        s = rand(pomcp.solver.rng, b)
        if !POMDPs.isterminal(pomcp.problem, s)
            simulate(pomcp, b, s, 0)
            b.N += 1
            all_terminal = false
        end
    end

    if all_terminal
        throw(AllSamplesTerminal(b.B))
    end

    best_node = first(values(b.children))
    best_V = best_node.V
    @assert !isnan(best_V)
    for node in collect(values(b.children))[2:end]
        if node.V >= best_V
            best_V = node.V
            best_node = node
        end
    end

    return best_node.label
end

"""
    simulate{S}(pomcp::POMCPPlanner, h::BeliefNode, s::S, depth)

Move the simulation forward a single step and update the BeliefNode h accordingly.
"""
function simulate{S,A,O,B}(pomcp::POMCPPlanner{S,A,O,B,POMCPSolver}, h::BeliefNode, s::S, depth)

    sol = pomcp.solver

    if POMDPs.discount(pomcp.problem)^depth < sol.eps ||
            POMDPs.isterminal(pomcp.problem, s) ||
            depth >= sol.max_depth
        return 0
    end
	if isempty(h.children)
        action_space_iter = sparse_actions(pomcp, pomcp.problem, h, sol.num_sparse_actions)
		for a in action_space_iter
			h.children[a] = ActNode(a,
                                    init_N(sol.init_N, pomcp.problem, h, a),
                                    init_V(sol.init_V, pomcp.problem, h, a),
                                    Dict{O,ObsNode{B,A,O}}())
		end

        if depth > 0 # no need for a rollout if this is the root node
            if POMDPs.discount(pomcp.problem) == 1.0
                step_to_eps = typemax(Int)
            else
                steps_to_eps = ceil(Int, log(sol.eps)/log(POMDPs.discount(pomcp.problem))-depth)
            end
            steps = min(sol.max_depth-depth, steps_to_eps)
            return POMDPs.discount(pomcp.problem)^depth * estimate_value(pomcp.solved_estimate, pomcp.problem, s, h, steps)
        else
            return 0.
        end
	end

    best_criterion_val = -Inf
    local best_node
    for node in values(h.children)
        if node.N == 0 && h.N == 1
            criterion_value = node.V
        elseif node.N == 0 && node.V == -Inf
            criterion_value = Inf
        else
            criterion_value = node.V + sol.c*sqrt(log(h.N)/node.N)
        end
        if criterion_value >= best_criterion_val
            best_criterion_val = criterion_value
            best_node = node
        end
    end
    a = best_node.label

    (sp, o, r) = POMDPs.generate_sor(pomcp.problem, s, a, sol.rng)

    if r == Inf
        warn("POMCP: +Inf reward. This is not recommended and may cause future errors.")
    end

    if haskey(best_node.children, o)
        hao = best_node.children[o]
    else
        if isa(pomcp.node_belief_updater, ParticleReinvigorator)
            hao = ObsNode(o, 0, ParticleCollection{S}(), Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
        else
            new_belief = update(pomcp.node_belief_updater, h.B, a, o) # this relies on h.B not being modified
            hao = ObsNode(o, 0, new_belief, Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
        end
        best_node.children[o]=hao
    end

    R = r + POMDPs.discount(pomcp.problem)*simulate(pomcp, hao, sp, depth+1)

    if uses_states_from_planner(hao.B)
        push!(hao.B, sp) # insert this state into the particle collection
    end
    hao.N += 1

    best_node.N += 1
    if best_node.V != -Inf
        best_node.V += (R-best_node.V)/best_node.N
    end

    return R
end

function simulate{S,A,O,B}(pomcp::POMCPPlanner{S,A,O,B,POMCPDPWSolver}, h::BeliefNode, s::S, depth)

    sol = pomcp.solver

    if POMDPs.discount(pomcp.problem)^depth < sol.eps ||
            POMDPs.isterminal(pomcp.problem, s) ||
            depth >= sol.max_depth
        return 0
    end

    total_N = reduce(add_N, 0, values(h.children))
    if sol.enable_action_pw
        if length(h.children) <= sol.k_action*total_N^sol.alpha_action
            a = next_action(sol.next_action, pomcp.problem, h.B, h) # XXX should this be a function of s or h.B?
            if !(a in keys(h.children))
                h.children[a] = ActNode(a,
                                        init_N(sol.init_N, pomcp.problem, h, a),
                                        init_V(sol.init_V, pomcp.problem, h, a),
                                        Dict{O,ObsNode{B,A,O}}())
            end
            if length(h.children) <= 1
                if depth > 0
                    return POMDPs.discount(pomcp.problem)^depth * estimate_value(pomcp.solved_estimate, pomcp.problem, s, h, depth)
                else
                    return 0.
                end
            end
        end
    else # run through all the actions
        if isempty(h.children)
            action_space_iter = POMDPs.iterator(POMDPs.actions(pomcp.problem))
            for a in action_space_iter
                h.children[a] = ActNode(a,
                                        init_N(sol.init_N, pomcp.problem, h, a),
                                        init_V(sol.init_V, pomcp.problem, h, a),
                                        Dict{O,ObsNode{B,A,O}}())
            end

            if depth > 0 # no need for a rollout if this is the root node
                return POMDPs.discount(pomcp.problem)^depth * estimate_value(pomcp.solved_estimate, pomcp.problem, s, h, depth)
            else
                return 0.
            end
        end
        total_N = h.N
    end

    # Calculate UCT
    best_criterion_val = -Inf
    local best_node
    for (a,node) in h.children
        if node.N == 0 && total_N <= 1
            criterion_value = node.V
        elseif node.N == 0 && node.V == -Inf
            criterion_value = Inf
        else
            criterion_value = node.V + sol.c*sqrt(log(total_N)/node.N)
        end
        if criterion_value >= best_criterion_val
            best_criterion_val = criterion_value
            best_node = node
        end
    end
    a = best_node.label

    if length(best_node.children) <= sol.k_observation*(best_node.N^sol.alpha_observation)
        state_was_generated = true

        # observation progressive widening
        (sp, o, r) = POMDPs.generate_sor(pomcp.problem, s, a, sol.rng)

        if r == Inf
            warn("POMCP: +Inf reward. This is not recommended and may cause future errors.")
        end

        if haskey(best_node.children, o)
            hao = best_node.children[o]
        else
            if isa(pomcp.node_belief_updater, ParticleReinvigorator)
                hao = ObsNode(o, 0, ParticleCollection{S}(), Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
            else
                new_belief = update(pomcp.node_belief_updater, h.B, a, o)
                hao = ObsNode(o, 0, new_belief, Dict{A,ActNode{A,O,ObsNode{B,A,O}}}())
            end
            best_node.children[o] = hao
        end

    else
        state_was_generated = false
        # otherwise sample nodes
        os = values(best_node.children)
        wv = WeightVec(Int[node.N for node in os]) # allocation
        hao = sample(sol.rng, os, wv)
        sp = rand(sol.rng, hao.B)
        r = POMDPs.reward(pomcp.problem, s, a, sp)
    end

    if state_was_generated
        if uses_states_from_planner(hao.B)
            push!(hao.B, sp)
        end
        hao.N += 1
    end

    R = r + POMDPs.discount(pomcp.problem)*simulate(pomcp, hao, sp, depth+1)

    best_node.N += 1
    if best_node.V != -Inf
        best_node.V += (R-best_node.V)/best_node.N
    end

    return R
end

make_node{P<:POMCPPlanner, N<:BeliefNode}(::P, b::N) = b

function make_node{S,A,O,B}(pomcp::POMCPPlanner{S,A,O,B}, belief)
    ANodeType = ActNode{A, O, ObsNode{B,A,O}}
    return RootNode{typeof(belief), A, ANodeType}(0, belief, Dict{A,ANodeType}())
end

"""
Add the N's of two nodes - for use in reduce
"""
add_N(a::AbstractActNode, b::AbstractActNode) = a.N + b.N
add_N(a::Int, b::AbstractActNode) = a + b.N