using Test
using PETSc, MPI
using ForwardDiff, SparseArrays
MPI.Initialized() || MPI.Init()

@testset "DMStagCreate1d" begin

    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)
    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        # Create 1D DMStag
        dm = PETSc.DMStagCreate1d(
                petsclib,
                comm,
                PETSc.DM_BOUNDARY_PERIODIC,
                20,
                2,
                2,
                PETSc.DMSTAG_STENCIL_BOX,
                2)
        @test PETSc.DMStagGetBoundaryTypes(dm)==PETSc.DM_BOUNDARY_PERIODIC
        PETSc.destroy(dm)

        # Create 1D DMStag with array of local @ of points
        dm = PETSc.DMStagCreate1d(petsclib,comm,PETSc.DM_BOUNDARY_NONE,20,2,2,PETSc.DMSTAG_STENCIL_BOX,2,[20])

        # Test get size
        @test PETSc.DMStagGetGlobalSizes(dm) == 20
        @test PETSc.DMStagGetLocalSizes(dm) == 20

        # Test
        @test PETSc.gettype(dm) == "stag"
        @test PETSc.getdimension(dm) == 1

        # Info about ranks
        @test PETSc.DMStagGetIsFirstRank(dm) == (true,false,false)
        @test PETSc.DMStagGetIsLastRank(dm) == (true,false,false)

        # Boundary
        @test PETSc.DMStagGetBoundaryTypes(dm)==PETSc.DM_BOUNDARY_NONE

        # Corners
        corners         = PETSc.getcorners(dm)
        ghost_corners   = PETSc.getghostcorners(dm)

        @test corners.lower[1] == 1
        @test corners.upper[1] == 20
        @test corners.size[1]  == 20
        @test corners.extra[1] == 1

        @test ghost_corners.lower[1] == 1
        @test ghost_corners.upper[1] == 21
        @test ghost_corners.size[1]  == 21
        @test ghost_corners.extra[1]  == 0

        # DOF
        @test PETSc.DMStagGetDOF(dm) == (2,2)
        PETSc.destroy(dm)

        # Create new struct and pass keyword arguments
        dm_1D = PETSc.DMStagCreate1d(petsclib,comm,PETSc.DM_BOUNDARY_NONE,200,2,2; stag_grid_x=10);
        @test PETSc.DMStagGetGlobalSizes(dm_1D) == 10
        @test PETSc.DMStagGetEntriesPerElement(dm_1D)==4

        # Stencil width & type
        @test  PETSc.DMStagGetStencilWidth(dm_1D)==2
        @test  PETSc.DMStagGetBoundaryTypes(dm_1D) == PETSc.DM_BOUNDARY_NONE

        PETSc.destroy(dm_1D)

        # test ghosted array set using keywords
        dm_ghosted = PETSc.DMStagCreate1d(petsclib,comm,PETSc.DM_BOUNDARY_GHOSTED,200,2,2; stag_grid_x=10);
        @test  PETSc.DMStagGetStencilWidth(dm_ghosted)==2
        corners         = PETSc.getcorners(dm_ghosted)

        @test corners.size[1]==10       # keyword overrides the specified value
        @test  PETSc.DMStagGetBoundaryTypes(dm_ghosted) == PETSc.DM_BOUNDARY_GHOSTED

        ind = PETSc.DMStagGetIndices(dm_ghosted);
        @test ind.center.x[3] == 5

        # simple test to retrieve the KSP object
        # NOTE: need to implement a similar SNES routine
        ksp = PETSc.KSP(dm_ghosted)
        @test PETSc.gettype(ksp)=="gmres"

        PETSc.destroy(dm_ghosted)

        PETSc.finalize(petsclib)
    end
end

@testset "DMStagCreate2d" begin

    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)
    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        # Create 2D DMStag
        dm_2D = PETSc.DMStagCreate2d(petsclib, comm,
                PETSc.DM_BOUNDARY_NONE,
                PETSc.DM_BOUNDARY_NONE,
                20,
                21,
                1,
                1,
                1,
                1,
                1,
                PETSc.DMSTAG_STENCIL_BOX,
                2)

        @test PETSc.DMStagGetGlobalSizes(dm_2D) == (20,21)
        corners = PETSc.getcorners(dm_2D)
        @test corners.size  == [20,21,0]
        @test corners.extra == [1, 1, 0]
        @test corners.lower == [1, 1, 1]
        @test corners.upper == [20,21,0]

        PETSc.finalize(petsclib)
    end
end

@testset "DMStagCreate3d" begin

    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)
    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        # Create 3D DMStag
        dm_3D = PETSc.DMStagCreate3d(petsclib,comm,
                PETSc.DM_BOUNDARY_NONE,
                PETSc.DM_BOUNDARY_NONE,
                PETSc.DM_BOUNDARY_NONE,
                20,
                21,
                22,
                1,
                1,
                1,
                2,
                2,
                2,
                2,
                PETSc.DMSTAG_STENCIL_BOX,
                1,
                [],
                [],
                [])
        @test PETSc.DMStagGetGlobalSizes(dm_3D) == (20, 21, 22)

        # copy struct
        dmnew = PETSc.DMStagCreateCompatibleDMStag(dm_3D,1,1,2,2)
        @test PETSc.DMStagGetGlobalSizes(dmnew) == (20, 21, 22)

        @test PETSc.DMStagGetGlobalSizes(dm_3D) == (20,21,22)
        corners = PETSc.getcorners(dm_3D)
        @test corners.size  == [20,21,22]
        @test corners.extra == [1, 1, 1]
        @test corners.lower == [1, 1, 1]
        @test corners.upper == [20,21,22]

        PETSc.finalize(petsclib)
    end
end

@testset "DMStag Vectors and Coordinates" begin
    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)
    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        # Create 1D DMStag
        dm_1D = PETSc.DMStagCreate1d(petsclib,comm,PETSc.DM_BOUNDARY_NONE,200,2,2; stag_grid_x=10);
        @test PETSc.DMStagGetGlobalSizes(dm_1D) == 10

        # Set coordinates
        PETSc.setuniformcoordinates!(dm_1D, (0,), (10,))

        DMcoord = PETSc.getcoordinateDM(dm_1D)
        @test PETSc.gettype(DMcoord)=="stag"

        # Retrieve array with staggered coordinates
        coord_vec   = PETSc.getcoordinateslocal(dm_1D)
        X_coord     = PETSc.DMStagVecGetArray(DMcoord, coord_vec);
        @test  X_coord[1,2] == 0.5

        # Set coordinates using product (1D) arrays
        dm_1D = PETSc.DMStagCreate1d(petsclib,comm,PETSc.DM_BOUNDARY_NONE,200,2,2; stag_grid_x=10);
        PETSc.setuniformcoordinatesproduct!(dm_1D, (0,), (10,))

        # retrieve DM with coordinates
        DMcoord = PETSc.getcoordinateDM(dm_1D)
        @test PETSc.gettype(DMcoord)=="product"
        coord_vec   = PETSc.getcoordinateslocal(dm_1D)

        # Note: retrieving 1D coordinate vectors using the "product" type appears broken
        #  This is something to be looked at later
        #x_coord,y_coord,z_coord = PETSc.DMStagGetProductCoordinateArraysRead(dm_1D); # BROKEN

        PETSc.DMStagGetLocationSlot(dm_1D, PETSc.DMSTAG_RIGHT, 0) ==4
        @test PETSc.DMStagGetProductCoordinateLocationSlot(dm_1D, PETSc.DMSTAG_RIGHT) == 2

        global_vec      = PETSc.createglobalvector(dm_1D)
        local_vec       = PETSc.createlocalvector(dm_1D)

        # Fill everything with some data
        fill!(local_vec, mpisize)
        fill!(global_vec, mpisize)
        @test global_vec[3] == 1.0

        # Add the local values to the global values
        PETSc.update!(global_vec, local_vec, PETSc.ADD_VALUES)
        @test global_vec[3] == 2.0

        # PETSc.destroy(DMcoord);

        # Do 2D tests
        dm_2D = PETSc.DMStagCreate2d(petsclib, comm,
            PETSc.DM_BOUNDARY_NONE,
            PETSc.DM_BOUNDARY_NONE,
            3,
            4,
            1,
            1,
            1,
            1,
            1,
            PETSc.DMSTAG_STENCIL_BOX,
            2)

        PETSc.setuniformcoordinates!(dm_2D, (1,3), (10,11))
        coord_vec   = PETSc.getcoordinateslocal(dm_2D)

        # Retrieve array with staggered coordinates
        DMcoord_2D = PETSc.getcoordinateDM(dm_2D)
        X_coord_2D = PETSc.DMStagVecGetArray(DMcoord_2D, coord_vec);

        @test X_coord_2D[3,3,1] ≈ 7.0
        @test X_coord_2D[3,4,2] ≈ 9.0
        @test X_coord_2D[3,3,3] ≈ 8.5
        @test X_coord_2D[3,3,4] ≈ 7.0

        vec_test_2D     = PETSc.createlocalvector(dm_2D)
        X               = PETSc.DMStagVecGetArray(dm_2D,vec_test_2D);
        X[end,end,end] = 111;                   # modify 3D array @ some point and DOF
        @test vec_test_2D[end]==111.0           # verify that this modified the vector as well
        Base.finalize(X)                        # release from memory

        #test stencil locations
        pos1 = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_LEFT,1,0,0,1)
        @test pos1.c == 1
        pos2 = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_RIGHT,4,0,0,0)
        pos  = [pos1, pos2]
        @test pos2.loc == PETSc.DMSTAG_RIGHT
        @test pos2.i == 4

        # Retrieve value from stencil
        vec_test       = PETSc.createlocalvector(dm_1D)
        vec_test      .= 1:length(vec_test)                 # point wise copy of data to PetscVec
        val             = PETSc.DMStagVecGetValuesStencil(dm_1D, vec_test, pos1) # this gets a single value
        @test val ==6
        vals            = PETSc.DMStagVecGetValuesStencil(dm_1D, vec_test, 2, pos)     # this gets an array of values
        @test vals[1] == 6

        X_1D            = PETSc.DMStagVecGetArray(dm_1D,vec_test);
        @test X_1D[2,3] == 7.0

        # Set values using stencils
        vec_test_global = PETSc.createglobalvector(dm_1D)
        val1 = PetscScalar.([2222.2, 3.2]);
        PETSc.DMStagVecSetValuesStencil(dm_1D, vec_test_global,    pos1, val1[1],   PETSc.INSERT_VALUES)
        @test vec_test_global[6] ≈ 2222.2
        PETSc.DMStagVecSetValuesStencil(dm_1D, vec_test_global, 2, pos,  val1,      PETSc.INSERT_VALUES)
        @test vec_test_global[21] ≈ 3.2

        pos3 = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_LEFT,1,0,0,1)
        val = PETSc.DMStagVecGetValuesStencil(dm_1D, vec_test, 2, [pos3; pos3])
        @test val[2] == 6.0
        PETSc.destroy(dm_1D);

        PETSc.finalize(petsclib)
    end
end

@testset "DMStag create matrixes" begin
    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)
    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        dm_1D = PETSc.DMStagCreate1d(petsclib,comm,PETSc.DM_BOUNDARY_NONE,200,2,2; stag_grid_x=10);
        PETSc.setuniformcoordinatesproduct!(dm_1D, (0,), (10,))

        A = PETSc.creatematrix(dm_1D);  #
        @test size(A) == (42,42)
        PETSc.assembled(A)

        # set some values using normal indices:
        A[1,1]  = 1.0
        A[1,10] = 1.0

        pos1 = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_LEFT,1,0,0,1)
        pos2 = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_RIGHT,4,0,0,0)
        pos  = [pos1, pos2]
        val1 = PetscScalar.([2222.2, 3.2]);
        PETSc.DMStagMatSetValuesStencil(dm_1D, A, pos1, pos1, 11.1, PETSc.INSERT_VALUES)
        PETSc.DMStagMatSetValuesStencil(dm_1D, A, 1, [pos2], 2, pos, val1, PETSc.INSERT_VALUES)

        PETSc.assemble(A)
        @test A[1,10] == 1.0

        # Reads a value from the matrix, using the stencil structure
        @test PETSc.DMStagMatGetValuesStencil(dm_1D, A, pos1, pos1)== PetscScalar(11.1)
        @test PETSc.DMStagMatGetValuesStencil(dm_1D, A, 1, [pos2], 2, pos)==val1

        PETSc.destroy(dm_1D);

        dofCenter       =   1;
        dofEdge         =   1;
        dofVertex       =   0
        stencilWidth    =   1;
        dm_2D = PETSc.DMStagCreate2d(petsclib,comm,
                                        PETSc.DM_BOUNDARY_GHOSTED,
                                        PETSc.DM_BOUNDARY_GHOSTED,
                                        10,11,
                                        PETSc.PETSC_DECIDE,PETSc.PETSC_DECIDE,
                                        dofVertex,dofEdge,dofCenter,
                                        PETSc.DMSTAG_STENCIL_BOX,stencilWidth)

        vec_test_2D_global      =   PETSc.createglobalvector(dm_2D)
        vec_test_2D_local       =   PETSc.createlocalvector(dm_2D)

        corners                 =   PETSc.getcorners(dm_2D)
        ghost_corners           =   PETSc.getghostcorners(dm_2D)

        for ix=corners.lower[1]:corners.upper[1]
            for iy=corners.lower[2]:corners.upper[2]
                local dof
                # DOF at the center point
                dof     = 0;
                posA    = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_DOWN,ix,iy,0,dof)
                value   = PetscScalar(ix+10);
                PETSc.DMStagVecSetValuesStencil(dm_2D, vec_test_2D_global, posA, value, PETSc.INSERT_VALUES)

                dof     = 0;
                posB    = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_LEFT,ix,iy,0,dof)
                value   = PetscScalar(33);
                PETSc.DMStagVecSetValuesStencil(dm_2D, vec_test_2D_global, posB, value, PETSc.INSERT_VALUES)

                dof     = 0;
                posC    = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_ELEMENT,ix,iy,0,dof)
                value   = PetscScalar(44);
                PETSc.DMStagVecSetValuesStencil(dm_2D, vec_test_2D_global, posC, value, PETSc.INSERT_VALUES)

            end
        end
        PETSc.assemble(vec_test_2D_global) # assemble global vector

        # Add the global values to the local values
        PETSc.update!(vec_test_2D_local, vec_test_2D_global,PETSc.INSERT_VALUES)

        # retrieve value back from the local array and check that it agrees with global one
        dof     = 0;
        pos     = PETSc.DMStagStencil{PetscInt}(PETSc.DMSTAG_DOWN,2,2,0,dof)
        @test PETSc.DMStagVecGetValuesStencil(dm_2D, vec_test_2D_local, pos) == 12.0

        # Extract an array that holds all DOF's
        X2D_dofs  = PETSc.DMStagVecGetArray(dm_2D,vec_test_2D_local)           # extract arrays with all DOF (mostly for visualizing)
        @test X2D_dofs[4,4,1] ≈ PetscScalar(12.0)
        @test X2D_dofs[4,4,2] ≈ PetscScalar(33.0)
        @test X2D_dofs[4,4,3] ≈ PetscScalar(44.0)

        # Extract an array of a specific DOF (here a face velocity @ the left)
        Xarray = PETSc.DMStagGetGhostArrayLocationSlot(dm_2D,vec_test_2D_local, PETSc.DMSTAG_LEFT, 0)
        @test sum(X2D_dofs[:,:,2]-Xarray)==0                # check if the local array is identical to the full array

        Xarray .= 111.                                      # Set a value @ a specific location
        @test vec_test_2D_local[2] ≈ PetscScalar(111)       # verify that this is changed in the PETSc Vec

        PETSc.finalize(petsclib)
    end
end

# -----------------
# Example of DMStag & SNES with AD jacobian
@testset "DMStag: 1D SNES AD" begin

    # Tell AD that it can handle Complex as scalars
    ForwardDiff.can_dual(::Type{ComplexF64}) = true
    ForwardDiff.can_dual(::Type{ComplexF32}) = true

    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)

    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        # Define a struct that holds data we need in the local SNES routines below
        mutable struct Data_1{PetscScalar,PetscInt}
            dm
            x_l
            f_l
        end

        user_ctx = Data_1{PetscScalar,PetscInt}(nothing, nothing, nothing);  # holds data we need in the local

        function FormRes!(ptr_fx_g, ptr_x_g, user_ctx)

            # Note that in PETSc, ptr_x_g and ptr_fx_g are pointers to global vectors.
            # Copy global to local vectors that are stored in user_ctx
            PETSc.update!(user_ctx.x_l, ptr_x_g,   PETSc.INSERT_VALUES)
            PETSc.update!(user_ctx.f_l, ptr_fx_g,  PETSc.INSERT_VALUES)

            # Retrieve arrays from the local vectors
            ArrayLocal_x     =   PETSc.DMStagVecGetArrayRead(user_ctx.dm, user_ctx.x_l);  # array with all local x-data
            ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, user_ctx.f_l);      # array with all local residual

            # Compute local residual
            ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx)

            # Finalize local arrays
            Base.finalize(ArrayLocal_x)
            Base.finalize(ArrayLocal_f)

            # Copy local into global residual vector
            PETSc.update!(ptr_fx_g, user_ctx.f_l,   PETSc.INSERT_VALUES)

        end

        function ComputeLocalResidual(dm, ArrayLocal_x, ArrayLocal_f, user_ctx)
            # Compute the local residual. The vectors include ghost points

            T              =   PETSc.DMStagGetGhostArrayLocationSlot(dm,ArrayLocal_x, PETSc.DMSTAG_LEFT,    0);
            fT             =   PETSc.DMStagGetGhostArrayLocationSlot(dm,ArrayLocal_f, PETSc.DMSTAG_LEFT,    0);

            P              =   PETSc.DMStagGetGhostArrayLocationSlot(dm,ArrayLocal_x, PETSc.DMSTAG_ELEMENT, 0);
            fP             =   PETSc.DMStagGetGhostArrayLocationSlot(dm,ArrayLocal_f, PETSc.DMSTAG_ELEMENT, 0);

            # compute the FD stencil
            indices         =     PETSc.DMStagGetIndices(dm);      # indices of (center/element) points, not including ghost values.
            gc              =     PETSc.getghostcorners(user_ctx.dm);   # start and end of loop including ghost points
            c               =     PETSc.getcorners(user_ctx.dm);        # start and end of loop including ghost points

            nT             =     length(T);                             # array length
            dx             =     1.0/(c.size[1]-1);
            xp             =     (gc.lower[1]:gc.upper[1]).*dx;         # coordinates including ghost points (to define source term)
            F              =     6.0.*xp .+ (xp .+1.e-12).^6.0;         # define source term function

            # Nonlinear equation @ nodal points
            ind            =     indices.vertex.x;                         #  There is one more "vertex" point
            i              =     ind[2:end-1]
            fT[ind[1]]     =     T[ind[1]  ]-0.5;                       # left BC
            fT[ind[end]]   =     T[ind[end]]-2.0;                       # right BC
            fT[i]          =     (T[i .+ 1] - 2*T[i] + T[i .- 1])/dx^2  + T[i].*T[i] - F[i] # NL diffusion with source term

            # second, non-coupled, equation @ center points
            ind            =     indices.center.x;                             #  There is one more "vertex" point
            i              =     ind[2:end-1];
            fP[ind[1]]     =     P[ind[1]]-30.;                             # left BC
            fP[ind[end]]   =     P[ind[end]]-20.;                           # right BC
            fP[i]          =     (P[i .+ 1] - 2*P[i] + P[i .- 1])/dx^2      # steady state diffusion

        end

        function  ForwardDiff_res(x, user_ctx)

            f   = zero(x)               # vector of zeros, of same type as x (local vector)

            ArrayLocal_x     =   PETSc.DMStagVecGetArray(user_ctx.dm, x);        # array with all local x-data
            ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, f);        # array with all local residual

            ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx);
            # As the residual vector f is linked with ArrayLocal_f, we don't need to pass ArrayLocal_f back

            return f;
        end

        function FormJacobian!(ptr_x_g, J, P, user_ctx)

            # This requires several steps:
            #
            #   1) Extract local vector from global solution (x) vector
            #   2) Compute local jacobian from the residual routine (note that
            #       this routine requires julia vectors as input)

            # Extract the local vector
            PETSc.update!(user_ctx.x_l, ptr_x_g,  PETSc.INSERT_VALUES)
            x               =   PETSc.unsafe_localarray(PetscScalar, user_ctx.x_l.ptr;  write=false, read=true)
            f_Residual      =   (x -> ForwardDiff_res(x, user_ctx));        # pass additional arguments into the routine
            J_julia         =   ForwardDiff.jacobian(f_Residual,x);

            # Note: since x is the LOCAL vector, J_julia also ends up having the same size.
            ind             =   PETSc.LocalInGlobalIndices(user_ctx.dm);
            if PETSc.assembled(J) == false
                J           =   PETSc.MatSeqAIJ(sparse(J_julia[ind,ind]));
            else
                J           .=   sparse(J_julia[ind,ind]);
            end

           Base.finalize(x_g)
           return sparse(J_julia[ind,ind]), ind
        end

        # Main part

        # Construct a 1D test case for a coupled P-T diffusion solver, with 1 DOF @ the center & 1 DOF @ faces
        nx              =   21;
        user_ctx.dm     =   PETSc.DMStagCreate1d(petsclib,comm,
                                PETSc.DM_BOUNDARY_GHOSTED,
                                nx,
                                1,                              # DOF @ vertex
                                1,                              # DOF @ center
                                PETSc.DMSTAG_STENCIL_BOX,
                                1);                             # Stencil width

        x_g             =   PETSc.createglobalvector(user_ctx.dm)
        f_g             =   PETSc.createglobalvector(user_ctx.dm)
        user_ctx.x_l    =   PETSc.createlocalvector(user_ctx.dm)
        user_ctx.f_l    =   PETSc.createlocalvector(user_ctx.dm)

        PJ           =      PETSc.creatematrix(user_ctx.dm);                  # extract (global) matrix from DMStag
        J_julia, ind =      FormJacobian!(x_g, PJ, PJ, user_ctx)
        PJ           =      PETSc.MatSeqAIJ(J_julia)                # assemble non-zero structure

        S = PETSc.SNES{PetscScalar}(petsclib, comm;
                snes_rtol=1e-12,
                snes_view=false,
                snes_monitor=true,
                ksp_view=true,
                # pc_type="none",
                snes_monitor_true_residual=false,
                snes_converged_reason=false);
        S.user_ctx  =       user_ctx;

        PETSc.setfunction!(S, FormRes!, f_g)
        PETSc.setjacobian!(S, FormJacobian!, PJ, PJ)

        # Solve
        PETSc.solve!(x_g, S);
       # @show x_g

        # check
        @test x_g[4] ≈ 29.5
        @test x_g[11] ≈ 0.63797 rtol=1e-4

        PETSc.finalize(petsclib)

    end
end

@testset "DMStag: 2D SNES AD"  begin

    comm = MPI.COMM_WORLD
    mpirank = MPI.Comm_rank(comm)
    mpisize = MPI.Comm_size(comm)

    # Tell AD that it can handle Complex as scalars
    ForwardDiff.can_dual(::Type{ComplexF64}) = true
    ForwardDiff.can_dual(::Type{ComplexF32}) = true

    for petsclib in PETSc.petsclibs
        PETSc.initialize(petsclib)
        PetscScalar = PETSc.scalartype(petsclib)
        PetscInt    = PETSc.inttype(petsclib)

        mutable struct Data_2D{PetscScalar,PetscInt}
            dm
            x_l
            f_l
        end
        user_ctx = Data_2D{PetscScalar,PetscInt}(nothing, nothing, nothing);  # holds data we need in the local

        function FormRes!(ptr_fx_g, ptr_x_g, user_ctx)
            # Note that in PETSc, ptr_x_g and ptr_fx_g are pointers to global vectors.

            # Copy global to local vectors
            PETSc.update!(user_ctx.x_l, ptr_x_g,  PETSc.INSERT_VALUES)
            PETSc.update!(user_ctx.f_l, ptr_fx_g, PETSc.INSERT_VALUES)

            # Retrieve arrays from the local vectors
            ArrayLocal_x     =   PETSc.DMStagVecGetArrayRead(user_ctx.dm, user_ctx.x_l);  # array with all local x-data
            ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, user_ctx.f_l);      # array with all local residual

            # Compute local residual
            ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx)

            # Finalize local arrays
            Base.finalize(ArrayLocal_x)
            Base.finalize(ArrayLocal_f)

            # Copy local into global residual vector
            PETSc.update!(ptr_fx_g, user_ctx.f_l, PETSc.INSERT_VALUES)

        end

        function  ForwardDiff_res(x, user_ctx)
            f   = zero(x)               # vector of zeros, of same type as x (local vector)

            ArrayLocal_x     =   PETSc.DMStagVecGetArray(user_ctx.dm, x);        # array with all ocal x-data
            ArrayLocal_f     =   PETSc.DMStagVecGetArray(user_ctx.dm, f);        # array with all ocal residual

            ComputeLocalResidual(user_ctx.dm, ArrayLocal_x, ArrayLocal_f, user_ctx);

            # As the residual vector f is linked with ArrayLocal_f, we don't need to pass ArrayLocal_f back

            return f;
        end

        function ComputeLocalResidual(dm, ArrayLocal_x, ArrayLocal_f, user_ctx)
            # Compute the local residual. The vectors include ghost points

            # Important! Make sure you retrieve the values from the correct locations. In this example we have a
            T              =   PETSc.DMStagGetGhostArrayLocationSlot(dm,ArrayLocal_x, PETSc.DMSTAG_ELEMENT, 0);
            fT             =   PETSc.DMStagGetGhostArrayLocationSlot(dm,ArrayLocal_f, PETSc.DMSTAG_ELEMENT, 0);

            # compute the FD stencil
            indices         =     PETSc.DMStagGetIndices(dm);      # indices of (center/element) points, not including ghost values.

            sz             =     size(user_ctx.dm);                                 # array length
            dx             =     1.0/(sz[1]-1);
            dz             =     1.0/(sz[2]-1);

            # set ghost points for BC'S
            bnd            =    PETSc.DMStagGetBoundaryTypes(user_ctx.dm)
            if bnd[1] == PETSc.DM_BOUNDARY_GHOSTED
                T[1,:]     =    T[2,:];        # zero flux; dT/dx=0
                T[end,:]   =    T[end-1,:];    # zero flux
            end

            # Diffusion @ center points
            indx        = indices.center.x;
            indz        = indices.center.y;

            ix          =     indx[1:end]                             # use ghost points in x  (required GHOSTED x-boundary)
            iz          =     indz[2:end-1]                           # center points

            # upper and lower BC (including corners)
            fT[indx,indz[1]]   =    T[indx,indz[1]] .- 0.5;                             # bottom BC
            fT[indx,indz[end]] =    T[indx,indz[end]] .- 2.0;                             # top BC

            fT[ix,iz]       =    (T[ix .+ 1,iz] - 2*T[ix,iz] + T[ix .- 1,iz])/dx^2   +
                                 (T[ix,iz .+ 1] - 2*T[ix,iz] + T[ix,iz .- 1])/dz^2

        end

        function FormJacobian!(ptr_x_g, J, P, user_ctx)
            # This requires several steps:
            #
            #   1) Extract local vector from global solution (x) vector
            #   2) Compute local jacobian from the residual routine (note that
            #       this routine requires julia vectors as input)

            # Extract the local vector
            PETSc.update!(user_ctx.x_l, ptr_x_g,  PETSc.INSERT_VALUES)
            x               =   PETSc.unsafe_localarray(PetscScalar, user_ctx.x_l.ptr;  write=false, read=true)

            f_Residual      =   (x -> ForwardDiff_res(x, user_ctx));        # pass additional rguments into the routine
            J_julia         =   ForwardDiff.jacobian(f_Residual,x);

            # Note: since x is the LOCAL vector, J_julia also ends up having the same size.
            ind             =   PETSc.LocalInGlobalIndices(user_ctx.dm);
            if PETSc.assembled(J) == false
                J           =   PETSc.MatSeqAIJ(sparse(J_julia[ind,ind]));
            else
                J           .=   sparse(J_julia[ind,ind]);
            end

            return sparse(J_julia[ind,ind]), ind, sparse(J_julia)
        end

        # Main routine starts here ----

        dofVertex   =   0
        dofEdge     =   0
        dofCenter   =   1
        nx,nz       =   6,25
        user_ctx.dm =   PETSc.DMStagCreate2d(petsclib,comm,
                    PETSc.DM_BOUNDARY_GHOSTED,
                    PETSc.DM_BOUNDARY_NONE,
                    nx,
                    nz,
                    1,
                    1,
                    dofVertex,
                    dofEdge,
                    dofCenter,
                    PETSc.DMSTAG_STENCIL_BOX,
                    1)
        PJ           =      PETSc.creatematrix(user_ctx.dm);                  # extract global) matrix from DMStag
        x_g             =   PETSc.createglobalvector(user_ctx.dm)
        f_g             =   PETSc.createglobalvector(user_ctx.dm)
        user_ctx.x_l    =   PETSc.createlocalvector(user_ctx.dm)
        user_ctx.f_l    =   PETSc.createlocalvector(user_ctx.dm)

        S = PETSc.SNES{PetscScalar}(petsclib, comm;
                snes_rtol=1e-12,
                snes_monitor=true,
                pc_type="none",
                snes_monitor_true_residual=true,
                snes_converged_reason=false);
        S.user_ctx  =       user_ctx;

        J_julia, ind, J_full =      FormJacobian!(x_g, PJ, PJ, user_ctx)
        PJ           =      PETSc.MatSeqAIJ(J_julia)                # assemble non-zero structure

        PETSc.setfunction!(S, FormRes!, f_g)
        PETSc.setjacobian!(S, FormJacobian!, PJ, PJ)

        # Solve 2D system
        sol = PETSc.solve!(x_g, S);

        PETSc.update!(user_ctx.x_l,sol, PETSc.INSERT_VALUES);   # copy global solution -> local vector
        T2d =   PETSc.DMStagGetGhostArrayLocationSlot(user_ctx.dm,user_ctx.x_l, PETSc.DMSTAG_ELEMENT,    0);

        @test T2d[5,5] ≈ 0.75 rtol=1e-3
        #
        # -----------------

        PETSc.finalize(petsclib)

    end
end


using CmdStan

N1 = 10
data1 = Dict(
  "Scalar" => rand(),
  "Single_element_vector" => [rand()],
  "Vector" => rand(N1),
  "Array" => rand(N1,3)
)

N = 100000
data = Dict(
  "Scalar" => rand(),
  "Single_element_vector" => [rand()],
  "Vector" => rand(N),
  "Array" => rand(N,3)
)

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  Stan.update_R_file("testdata0", data1)
  Stan.update_R_file_jon("testdata1", data1)

  @time Stan.update_R_file("testdata0", data)
  @time Stan.update_R_file_jon("testdata1", data)

  isfile("testdata0") && rm("testdata0");
  isfile("testdata1") && rm("testdata1");

end # cd


abstract type NamedRange{N,T,R} <: StaticRangeWrapper{T,R} end

Base.names(r::NamedRange{N,T,R}) where {N,T,R} = N
Base.names(r::Type{<:NamedRange{N,T,R}}) where {N,T,R} = N

NamedRange(n::NTuple{N,Symbol}, r::SRange) where N = NamedSRange{n}(r)
NamedRange(n::NTuple{N,Symbol}, r::MRange) where N = NamedMRange{n}(r)

struct NamedSRange{N,T,R} <: NamedRange{N,T,R}

    function NamedSRange{N,T,R}() where {N,T,R}
        isstatic(R) || error("NamedSRange only supports static ranges.")
        new{N,T,R}()
    end
end

NamedSRange{N}(r::SRange{T}) where {N,T} = NamedSRange{N,T,typeof(r)}()

"""
    NamedMRange

A mutable range that may be indexed by
"""
mutable struct NamedMRange{N,T,R} <: NamedRange{N,T,R}
    names::N
    parent::R

    NamedMRange{N}(r::AbstractRange{T}) where {N,T} = new{N,T,typeof(r)}(names, r)
end

Base.parent(r::NamedMRange) = getfield(r, :parent)

@inline getindex(r::NamedRange, i::Integer) = getindex(parent(r), i)

@inline getindex(r::NamedRange, s::Symbol) = getindex(r, SSymbol(s))

@inline @generated function getindex(r::NamedRange, ::SSymbol{S}) where S
    idx = findfirst(y->y==S::Symbol, names(r))
    :(parent(r)[$idx])
end

#@inline Base.setindex!(x::LArray,y,i...) = getfield(x,:__x)[i...] = y

@inline function getindex(r::NamedRange, inds::AbstractVector{<:Number})
    _namedrange(names(r)[inds], getindex(parent(r), inds))
end

_namedrange(n::NTuple{N,Symbol}, r::Vector) where N = NamedTuple{n}(r...)
_namedrange(n::NTuple{N,Symbol}, r::SRange) where N = NamedSRange{n}(r)
_namedrange(n::NTuple{N,Symbol}, r::AbstractRange) where N = NamedMRange{n}(r)

function getindex(x::NamedRange, s::AbstractArray{Symbol,1})
    NamedTuple{(s...,)}([getindex(x, si) for si in s]...)
end


mutable struct force_torque_t <: LCMType
    timestamp::Int64
    fx::Float64
    fy::Float64
    fz::Float64
    tx::Float64
    ty::Float64
    tz::Float64
end

@lcmtypesetup(force_torque_t)


using IRTools, MacroTools, InteractiveUtils, Test
using IRTools: @dynamo, IR, meta, isexpr, xcall, self, insertafter!, recurse!,
  argument!, return!, func, var

@dynamo roundtrip(a...) = IR(a...)

@dynamo function passthrough(a...)
  ir = IR(a...)
  ir == nothing && return
  recurse!(ir)
  return ir
end

add(a, b) = a+b
@test roundtrip(add, 2, 3) == 5
@test passthrough(add, 2, 3) == 5

@test @code_ir(passthrough, add(2, 3)) isa IR

IRTools.refresh(roundtrip)

add(a, b) = a*b
@test roundtrip(add, 2, 3) == 6

relu(x) = x > 0 ? x : 0
@test roundtrip(relu, 1) == 1
@test passthrough(relu, 1) == 1

foo(x) = y = x > 0 ? x + 1 : x - 1
@test roundtrip(foo, 1) == 2
@test roundtrip(foo, -1) == -2
@test passthrough(foo, 1) == 2

@test passthrough(() -> [1, 2, 3]) == [1, 2, 3]

@dynamo function mullify(a...)
  ir = IR(a...)
  ir == nothing && return
  ir = MacroTools.prewalk(ir) do x
    x isa GlobalRef && x.name == :(*) && return GlobalRef(Base, :+)
    return x
  end
  for (x, st) in ir
    isexpr(st.expr, :call) || continue
    ir[x] = xcall(Main, :mullify, st.expr.args...)
  end
  return ir
end

@test mullify(prod, [5, 10]) == 15

@dynamo err(a...) = 0//0

@test_throws IRTools.Inner.CompileError err(+, 2, 3)

mutable struct Context
  calls::Int
end

@dynamo function (cx::Context)(a...)
  ir = IR(a...)
  ir == nothing && return
  recurse!(ir)
  calls = pushfirst!(ir, xcall(:getfield, self, QuoteNode(:calls)))
  calls = insertafter!(ir, calls, xcall(:+, calls, 1))
  insertafter!(ir, calls, xcall(:setfield!, self, QuoteNode(:calls), calls))
  return ir
end

@code_lowered Context(0)(add, 2, 3)

cx = Context(0)
@test cx(add, 2, 3.0) == 6
@test cx.calls > 5

ir = IR()
x = argument!(ir)
y = push!(ir, xcall(:*, x, x))
return!(ir, y)

@test IRTools.evalir(ir, 5) == 25

ir = IR()
x = argument!(ir)
y = argument!(ir)
z = push!(ir, xcall(:*, x, y))
return!(ir, z)

@test IRTools.evalir(ir, 5, 3) == 15

function pow(x, n)
  r = 1
  while n > 0
    n -= 1
    r *= x
  end
  return r
end

ir = @code_ir pow(2, 3)

ir[var(8)] = xcall(:+, var(5), var(2))

mul = func(ir)

@test mul(nothing, 10, 3) == 31


"""
    permute_loci!(data::PopData)
Edits `PopData` in place with loci permuted across populations within
the `.loci` dataframe.
"""
@inline function permute_loci!(data::PopData)
    @inbounds for locus in groupby(data.loci, :locus)
        shuffle!(locus.population)
    end
    data
end

"""
    permute_samples!(data::PopData; meta::Bool = false)
Edits `PopData` in place with samples permuted across populations within
the `.loci` dataframe. Since performance is important for many permutations,
the default is to only edit the `.loci` table in place; use `meta = true`
if you also require the `.meta` dataframe edited in place.
"""
@inline function permute_samples!(data::PopData; meta::Bool = false)
    pops = shuffle(data.meta.population)

    if meta == true
        meta_pops = deepcopy(pops)
        @inbounds for name in groupby(data.meta, :name)
            @inbounds name.population .= pop!(meta_pops)
        end
    end

    @inbounds for name in groupby(data.loci, :name)
        @inbounds name.population .= pop!(pops)
    end
    data
end

@inline function permute_samples!(data::AbstractDataFrame, popnames::Vector{String})
    pops = shuffle(popnames)

    @inbounds for name in groupby(data, :name)
        @inbounds name.population .= pop!(pops)
    end
    data
end

"""
    permute_genotypes!(data::PopData; by::String = "locus")
Edits `PopData` in place with genotypes permuted across individuals within
the `.loci` dataframe. Use `by = "population"` (or `"pop"`) to permute genotypes
within populations.
"""
@inline function permute_genotypes!(data::PopData; by::String = "locus")
    #establish mode of operation
    if by in ["locus", "loci"]
        groupings = :locus
    elseif by in ["population", "pop"]
        groupings = [:locus, :population]
    else
        error("Please choose from either \"locus\" or \"population\" run methods.")
    end
    @inbounds for grp in groupby(data.loci, groupings)
        grp.genotype .= strict_shuffle!(grp.genotype)
    end
    data
end

"""
    permute_alleles!(data::PopData; ploidy::Union{Nothing, Int} = nothing, by::String = "locus")
Edits `PopData` in place with alleles permuted and reconstructed into genotypes
for each locus within the `.loci` dataframe. Use `by = "population"` (or `"pop"`)
to permute alleles within populations. If `ploidy` is not provided (default `ploidy = nothing`),
then ploidy will be identified from the PopData. If performance is important,
it would be best to identify ploidy in advance and set it to a specific integer.
"""
@inline function permute_alleles!(data::PopData; ploidy::Union{Nothing, Int} = nothing, by::String = "locus")
    #establish mode of operation
    if by in ["locus", "loci"]
        groupings = :locus
    elseif by in ["population", "pop"]
        groupings = [:locus, :population]
    else
        error("Please choose from either \"locus\" or \"population\" run methods.")
    end

    if ploidy == nothing
        tmp = unique(data.meta.ploidy)
        length(tmp) > 1 && error("This permutation method is not appropriate for mixed-ploidy data")
        ploidy = tmp[1]
    end

    @inbounds for grp in groupby(data.loci, groupings)
        alle = shuffle(alleles(grp.genotype))
        new_genos = Tuple.(Base.Iterators.partition(alle, ploidy))
        (@view grp.genotype[.!ismissing.(grp.genotype)]) .= new_genos
    end
    data
end


#=
MathOptInterface wrapper of Hypatia solver
=#

"""
$(TYPEDEF)

A MathOptInterface optimizer type for Hypatia.
"""
mutable struct Optimizer{T <: Real} <: MOI.AbstractOptimizer
    solver::Solvers.Solver{T} # Hypatia solver object

    # data for transforming certificates
    obj_sense::MOI.OptimizationSense
    zeros_idxs::Vector{UnitRange{Int}}
    moi_cones::Vector{MOI.AbstractVectorSet}
    moi_cone_idxs::Vector{UnitRange{Int}}

    function Optimizer{T}(; options...) where {T <: Real}
        opt = new{T}()
        opt.solver = Solvers.Solver{T}(; options...)
        return opt
    end
end

Optimizer(; options...) = Optimizer{Float64}(; options...) # default to Float64

MOI.is_empty(opt::Optimizer) = (opt.solver.status == Solvers.NotLoaded)

MOI.empty!(opt::Optimizer) = (opt.solver.status = Solvers.NotLoaded)

MOI.get(::Optimizer, ::MOI.SolverName) = "Hypatia"

MOI.get(opt::Optimizer, ::MOI.RawSolver) = opt.solver

MOI.supports(::Optimizer, ::MOI.Silent) = true

function MOI.set(opt::Optimizer, ::MOI.Silent, value::Bool)
    opt.solver.verbose = !value
    return
end

MOI.get(opt::Optimizer, ::MOI.Silent) = !opt.solver.verbose

MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true

function MOI.set(opt::Optimizer, ::MOI.TimeLimitSec, value::Union{Real, Nothing})
    opt.solver.time_limit = something(value, Inf)
    return
end

function MOI.get(opt::Optimizer, ::MOI.TimeLimitSec)
    if isfinite(opt.solver.time_limit)
        return opt.solver.time_limit
    end
    return
end

function MOI.get(opt::Optimizer, ::MOI.SolveTimeSec)
    if opt.solver.status == Solvers.NotLoaded
        error("solve has not been called")
    end
    return opt.solver.solve_time
end

MOI.get(opt::Optimizer, ::MOI.RawStatusString) = string(opt.solver.status)

MOI.get(opt::Optimizer, ::MOI.BarrierIterations) = opt.solver.num_iters

function MOI.set(opt::Optimizer, param::MOI.RawOptimizerAttribute, value)
    return setproperty!(opt.solver, Symbol(param.name), value)
end

function MOI.get(opt::Optimizer, param::MOI.RawOptimizerAttribute)
    return getproperty(opt.solver, Symbol(param.name))
end

MOI.supports(
    ::Optimizer{T},
    ::Union{MOI.ObjectiveSense, MOI.ObjectiveFunction{<:Union{VI, SAF{T}}}},
    ) where {T <: Real} = true

MOI.supports_constraint(
    ::Optimizer{T},
    ::Type{<:Union{VV, VAF{T}}},
    ::Type{<:Union{MOI.Zeros, SupportedCone{T}}},
    ) where {T <: Real} = true

function MOI.copy_to(
    opt::Optimizer{T},
    src::MOI.ModelLike,
    ) where {T <: Real}
    idx_map = MOI.Utilities.IndexMap()

    # variables
    n = MOI.get(src, MOI.NumberOfVariables()) # columns of A
    j = 0
    for vj in MOI.get(src, MOI.ListOfVariableIndices())
        j += 1
        idx_map[vj] = VI(j)
    end
    @assert j == n
    for attr in MOI.get(src, MOI.ListOfVariableAttributesSet())
        if attr == MOI.VariableName() || attr == MOI.VariablePrimalStart()
            continue
        end
        throw(MOI.UnsupportedAttribute(attr))
    end

    # objective
    opt.obj_sense = MOI.MIN_SENSE
    model_c = zeros(T, n)
    obj_offset = zero(T)
    for attr in MOI.get(src, MOI.ListOfModelAttributesSet())
        if attr == MOI.Name()
            continue
        elseif attr == MOI.ObjectiveSense()
            opt.obj_sense = MOI.get(src, MOI.ObjectiveSense())
        elseif attr isa MOI.ObjectiveFunction
            F = MOI.get(src, MOI.ObjectiveFunctionType())
            if !(F <: Union{VI, SAF{T}})
                error("objective function type $F not supported")
            end
            obj = convert(SAF{T}, MOI.get(src, MOI.ObjectiveFunction{F}()))
            for t in obj.terms
                model_c[idx_map[t.variable].value] += t.coefficient
            end
            obj_offset = obj.constant
        else
            throw(MOI.UnsupportedAttribute(attr))
        end
    end
    if opt.obj_sense == MOI.MAX_SENSE
        model_c .*= -1
        obj_offset *= -1
    end

    # constraints
    get_src_cons(F, S) = MOI.get(src, MOI.ListOfConstraintIndices{F, S}())
    get_con_fun(con_idx) = MOI.get(src, MOI.ConstraintFunction(), con_idx)
    get_con_set(con_idx) = MOI.get(src, MOI.ConstraintSet(), con_idx)

    # equality constraints
    (IA, JA, VA) = (Int[], Int[], T[])
    model_b = T[]
    opt.zeros_idxs = zeros_idxs = Vector{UnitRange{Int}}()
    for F in (VV, VAF{T}), ci in get_src_cons(F, MOI.Zeros)
        fi = get_con_fun(ci)
        si = get_con_set(ci)
        _con_IJV(IA, JA, VA, model_b, zeros_idxs, fi, si, idx_map)
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.Zeros}(length(zeros_idxs))
    end
    model_A = dropzeros!(sparse(IA, JA, VA, length(model_b), n))

    # conic constraints
    (IG, JG, VG) = (Int[], Int[], T[])
    model_h = T[]
    moi_cones = MOI.AbstractVectorSet[]
    moi_cone_idxs = Vector{UnitRange{Int}}()
    cones = Cones.Cone{T}[]

    # build up one nonnegative cone
    for F in (VV, VAF{T}), ci in get_src_cons(F, MOI.Nonnegatives)
        fi = get_con_fun(ci)
        si = get_con_set(ci)
        _con_IJV(IG, JG, VG, model_h, moi_cone_idxs, fi, si, idx_map)
        push!(moi_cones, si)
        idx_map[ci] = MOI.ConstraintIndex{F, MOI.Nonnegatives}(length(moi_cones))
    end
    if !isempty(moi_cones)
        push!(cones, cone_from_moi(T, MOI.Nonnegatives(length(model_h))))
    end

    # other conic constraints
    for (F, S) in MOI.get(src, MOI.ListOfConstraintTypesPresent())
        if !MOI.supports_constraint(opt, F, S)
            throw(MOI.UnsupportedConstraint{F, S}())
        end
        for attr in MOI.get(src, MOI.ListOfConstraintAttributesSet{F, S}())
            if attr == MOI.ConstraintName() ||
               attr == MOI.ConstraintPrimalStart() ||
               attr == MOI.ConstraintDualStart()
                continue
            end
            throw(MOI.UnsupportedAttribute(attr))
        end
        if S == MOI.Zeros || S == MOI.Nonnegatives
            continue # already copied these constraints
        end

        for ci in get_src_cons(F, S)
            fi = get_con_fun(ci)
            si = get_con_set(ci)
            _con_IJV(IG, JG, VG, model_h, moi_cone_idxs, fi, si, idx_map)
            push!(cones, cone_from_moi(T, si))
            push!(moi_cones, si)
            idx_map[ci] = MOI.ConstraintIndex{F, S}(length(moi_cones))
        end
    end

    model_G = dropzeros!(sparse(IG, JG, VG, length(model_h), n))
    opt.moi_cone_idxs = moi_cone_idxs
    opt.moi_cones = moi_cones

    # finalize model and load into solver
    model = Models.Model{T}(model_c, model_A, model_b, model_G, model_h,
        cones; obj_offset = obj_offset)
    Solvers.load(opt.solver, model)

    return idx_map
end

MOI.optimize!(opt::Optimizer) = Solvers.solve(opt.solver)

function MOI.modify(
    opt::Optimizer{T},
    ::MOI.ObjectiveFunction{SAF{T}},
    chg::MOI.ScalarConstantChange{T},
    ) where {T}
    obj_offset = chg.new_constant
    if opt.obj_sense == MOI.MAX_SENSE
        obj_offset = -obj_offset
    end
    Solvers.modify_obj_offset(opt.solver, obj_offset)
    return
end

function MOI.modify(
    opt::Optimizer{T},
    ::MOI.ObjectiveFunction{SAF{T}},
    chg::MOI.ScalarCoefficientChange{T},
    ) where {T}
    new_c = chg.new_coefficient
    if opt.obj_sense == MOI.MAX_SENSE
        new_c = -new_c
    end
    Solvers.modify_c(opt.solver, [chg.variable.value], [new_c])
    return
end

function MOI.modify(
    opt::Optimizer{T},
    ci::MOI.ConstraintIndex{VAF{T}, MOI.Zeros},
    chg::MOI.VectorConstantChange{T},
    ) where {T}
    idxs = opt.zeros_idxs[ci.value]
    Solvers.modify_b(opt.solver, idxs, chg.new_constant)
    return
end

function MOI.modify(
    opt::Optimizer{T},
    ci::MOI.ConstraintIndex{VAF{T}, <:SupportedCone{T}},
    chg::MOI.VectorConstantChange{T},
    ) where {T}
    i = ci.value
    idxs = opt.moi_cone_idxs[i]
    set = opt.moi_cones[i]
    new_h = chg.new_constant
    if needs_permute(set)
        @assert !needs_rescale(set)
        new_h = permute_affine(set, new_h)
    end
    if needs_rescale(set)
        rescale_affine(set, new_h)
    end
    Solvers.modify_h(opt.solver, idxs, new_h)
    return
end

function MOI.get(opt::Optimizer, ::MOI.TerminationStatus)
    status = opt.solver.status
    if status in (Solvers.NotLoaded, Solvers.Loaded)
        return MOI.OPTIMIZE_NOT_CALLED
    elseif status == Solvers.Optimal
        return MOI.OPTIMAL
    elseif status == Solvers.NearOptimal
        return MOI.ALMOST_OPTIMAL
    elseif status == Solvers.PrimalInfeasible || status == Solvers.PrimalInconsistent
        return MOI.INFEASIBLE
    elseif status == Solvers.NearPrimalInfeasible
        return MOI.ALMOST_INFEASIBLE
    elseif status == Solvers.DualInfeasible || status == Solvers.DualInconsistent
        return MOI.DUAL_INFEASIBLE
    elseif status == Solvers.NearDualInfeasible
        return MOI.ALMOST_DUAL_INFEASIBLE
    elseif status == Solvers.SlowProgress
        return MOI.SLOW_PROGRESS
    elseif status == Solvers.IterationLimit
        return MOI.ITERATION_LIMIT
    elseif status == Solvers.TimeLimit
        return MOI.TIME_LIMIT
    elseif status == Solvers.NumericalFailure
        return MOI.NUMERICAL_ERROR
    elseif status in (Solvers.IllPosed, Solvers.NearIllPosed)
        return MOI.OTHER_LIMIT
    else
        @warn("Hypatia status $(opt.solver.status) not handled")
        return MOI.OTHER_ERROR
    end
end

function MOI.get(opt::Optimizer, attr::MOI.PrimalStatus)
    if attr.result_index != 1
        return MOI.NO_SOLUTION
    end
    status = opt.solver.status
    if status == Solvers.Optimal
        return MOI.FEASIBLE_POINT
    elseif status == Solvers.NearOptimal
        return MOI.NEARLY_FEASIBLE_POINT
    elseif status == Solvers.PrimalInfeasible
        return MOI.INFEASIBLE_POINT
    elseif status == Solvers.DualInfeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == Solvers.NearDualInfeasible
        return MOI.NEARLY_INFEASIBILITY_CERTIFICATE
    elseif status in (Solvers.IllPosed, Solvers.NearIllPosed)
        return MOI.OTHER_RESULT_STATUS
    else
        return MOI.UNKNOWN_RESULT_STATUS
    end
end

function MOI.get(opt::Optimizer, attr::MOI.DualStatus)
    if attr.result_index != 1
        return MOI.NO_SOLUTION
    end
    status = opt.solver.status
    if status == Solvers.Optimal
        return MOI.FEASIBLE_POINT
    elseif status == Solvers.NearOptimal
        return MOI.NEARLY_FEASIBLE_POINT
    elseif status == Solvers.PrimalInfeasible
        return MOI.INFEASIBILITY_CERTIFICATE
    elseif status == Solvers.NearPrimalInfeasible
        return MOI.NEARLY_INFEASIBILITY_CERTIFICATE
    elseif status == Solvers.DualInfeasible
        return MOI.INFEASIBLE_POINT
    elseif status in (Solvers.IllPosed, Solvers.NearIllPosed)
        return MOI.OTHER_RESULT_STATUS
    else
        return MOI.UNKNOWN_RESULT_STATUS
    end
end

_sense_val(sense::MOI.OptimizationSense) = (sense == MOI.MAX_SENSE ? -1 : 1)

function MOI.get(opt::Optimizer, attr::MOI.ObjectiveValue)
    MOI.check_result_index_bounds(opt, attr)
    return _sense_val(opt.obj_sense) * opt.solver.primal_obj
end

function MOI.get(opt::Optimizer, attr::MOI.DualObjectiveValue)
    MOI.check_result_index_bounds(opt, attr)
    return _sense_val(opt.obj_sense) * opt.solver.dual_obj
end

MOI.get(opt::Optimizer, ::MOI.ResultCount) = 1

function MOI.get(opt::Optimizer, attr::MOI.VariablePrimal, vi::VI)
    MOI.check_result_index_bounds(opt, attr)
    return opt.solver.result.x[vi.value]
end

function MOI.get(
    opt::Optimizer{T},
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{<:Union{VV, VAF{T}}, MOI.Zeros},
    ) where {T}
    MOI.check_result_index_bounds(opt, attr)
    return opt.solver.result.y[opt.zeros_idxs[ci.value]]
end

function MOI.get(
    opt::Optimizer{T},
    attr::MOI.ConstraintDual,
    ci::MOI.ConstraintIndex{<:Union{VV, VAF{T}}, <:SupportedCone{T}},
    ) where {T}
    MOI.check_result_index_bounds(opt, attr)
    i = ci.value
    z_i = opt.solver.result.z[opt.moi_cone_idxs[i]]
    return untransform_affine(opt.moi_cones[i], z_i)
end

function MOI.get(
    opt::Optimizer{T},
    attr::MOI.ConstraintPrimal,
    ci::MOI.ConstraintIndex{<:Union{VV, VAF{T}}, <:SupportedCone{T}},
    ) where {T}
    MOI.check_result_index_bounds(opt, attr)
    i = ci.value
    s_i = opt.solver.result.s[opt.moi_cone_idxs[i]]
    return untransform_affine(opt.moi_cones[i], s_i)
end

function _con_IJV(
    IM::Vector{Int},
    JM::Vector{Int},
    VM::Vector{T},
    vect::Vector{T},
    idxs_vect::Vector{UnitRange{Int}},
    func::VV,
    set::MOI.AbstractVectorSet,
    idx_map::MOI.IndexMap,
    ) where {T <: Real}
    dim = MOI.output_dimension(func)
    start = length(vect)
    idxs = start .+ (1:dim)
    push!(idxs_vect, idxs)
    append!(vect, zero(T) for _ in 1:dim)
    if needs_permute(set)
        perm_idxs = permute_affine(set, 1:dim)
        perm_idxs .+= start
        append!(IM, perm_idxs)
    else
        append!(IM, idxs)
    end
    append!(JM, idx_map[vi].value for vi in func.variables)
    append!(VM, -one(T) for _ in 1:dim)
    if needs_rescale(set)
        @views rescale_affine(set, VM[(end - dim + 1):end])
    end
    return
end

function _con_IJV(
    IM::Vector{Int},
    JM::Vector{Int},
    VM::Vector{T},
    vect::Vector{T},
    idxs_vect::Vector{UnitRange{Int}},
    func::VAF{T},
    set::MOI.AbstractVectorSet,
    idx_map::MOI.IndexMap,
    ) where {T <: Real}
    dim = MOI.output_dimension(func)
    start = length(vect)
    idxs = start .+ (1:dim)
    push!(idxs_vect, idxs)
    if needs_permute(set)
        @assert !needs_rescale(set)
        append!(vect, permute_affine(set, func.constants))
        perm_idxs = permute_affine(set, func)
        perm_idxs .+= start
        append!(IM, perm_idxs)
    else
        append!(vect, func.constants)
        append!(IM, start + vt.output_index for vt in func.terms)
    end
    append!(JM, idx_map[vt.scalar_term.variable].value for vt in func.terms)
    append!(VM, -vt.scalar_term.coefficient for vt in func.terms)
    if needs_rescale(set)
        @views vm = VM[(end - length(func.terms) + 1):end]
        rescale_affine(set, func, vm)
        @views rescale_affine(set, vect[idxs])
    end
    return
end


module _Finance

using Brainstorm._Finance
using Base.Test

include("daycount.jl")
include("cashflow.jl")

function test_all()
    println("")
    test_daycount_all()
    test_cashflow_all()
end

end


_hbox(args...) = Widgets.div(args...; style = Dict("display" => "flex", "flex-direction"=>"row"))
_vbox(args...) = Widgets.div(args...; style = Dict("display" => "flex", "flex-direction"=>"column"))

get_kwargs(; kwargs...) = kwargs
string2kwargs(s::AbstractString) = eval(Meta.parse("get_kwargs($s)"))

const analysis_options = OrderedDict(
    "" => nothing,
    "Cumulative" => Recombinase.cumulative,
    "Density" => Recombinase.density,
    "Hazard" => Recombinase.hazard,
    "Prediction" => Recombinase.prediction,
)

function is_small(col, n = 100)
    for (ind, _) in enumerate(IterTools.distinct(col))
        ind > n && return false
    end
    return true
end

function take_small(t, vec::AbstractVector{Symbol}, n = 100)
    filter(vec) do sym
        col = get(columns(t), sym, ())
        return is_small(col, n)
    end
end

"""
`gui(data, plotters; postprocess = NamedTuple())`

Create a gui around `data::IndexedTable` given a list of plotting
functions plotters.

## Examples

```julia
using StatsPlots, Recombinase, JuliaDB, Interact
school = loadtable(joinpath(Recombinase.datafolder, "school.csv"))
plotters = [plot, scatter, groupedbar]
Recombinase.gui(school, plotters)
```
"""
function gui(data′, plotters; postprocess = NamedTuple())
    (data′ isa AbstractObservable) || (data′ = Observable{Any}(data′))
    data = @map table(&data′, copy = false, presorted = true)
    ns = @map collect(colnames(&data))
    maybens = @map vcat(Symbol(), &ns)
    xaxis = dropdown(ns,label = "X")
    yaxis = dropdown(maybens,label = "Y")
    an_opt = dropdown(analysis_options, label = "Analysis")
    axis_type = dropdown([:automatic, :continuous, :discrete, :vectorial], label = "Axis type")
    error = dropdown(@map(vcat(automatic, &ns)), label="Error")
    styles = collect(keys(style_dict))
    sort!(styles)
    smallns = map(take_small, data, maybens)
    splitters = [dropdown(smallns, label = string(style)) for style in styles]
    plotter = dropdown(plotters, label = "Plotter")
    ribbon = toggle("Ribbon", value = false)
    btn = button("Plot")
    output = Observable{Any}("Set the dropdown menus and press plot to get started.")
    plot_kwargs = Widgets.textbox("Insert optional plot attributes")
    @map! output begin
        &btn
        select = yaxis[] == Symbol() ? xaxis[] : (xaxis[], yaxis[])
        grps = Dict(key => val[] for (key, val) in zip(styles, splitters) if val[] != Symbol())
        an = an_opt[]
        an_inf = isnothing(an) ? nothing : Analysis{axis_type[]}(an)
        args, kwargs = series2D(
                                an_inf,
                                &data,
                                Group(; grps...);
                                postprocess = postprocess,
                                select = select,
                                error = error[],
                                ribbon = ribbon[]
                               )
        plotter[](args...; kwargs..., string2kwargs(plot_kwargs[])...)
    end
    ui = Widget(
        OrderedDict(
            :xaxis => xaxis,
            :yaxis => yaxis,
            :analysis => an_opt,
            :axis_type => axis_type,
            :error => error,
            :plotter => plotter,
            :plot_button => btn,
            :plot_kwargs => plot_kwargs,
            :ribbon => ribbon,
            :splitters => splitters,
        ),
        output = output
    )
    Widgets.@layout! ui Widgets.div(
                                    _hbox(
                                          :xaxis,
                                          :yaxis,
                                          :analysis,
                                          :axis_type,
                                          :error,
                                          :plotter
                                         ),
                                    :ribbon,
                                    :plot_button,
                                    _hbox(
                                          _vbox(:splitters...),
                                          _vbox(output, :plot_kwargs)
                                         )
                                   )
end


include("plot_measurements.jl")
include("plot_observables.jl")
include("plot_ranking.jl")
include("plot_intervals.jl")


#=
type RobustNoCrashMDP <: RobustMDP{MLPhysicalState, MLAction}
    base::MLMDP
end
RobustMCTS.representative_mdp(rmdp::RobustNoCrashMDP) = StochasticBehaviorNoCrashMDP(rmdp.base)

function RobustMCTS.next_model(gen::RandomModelGenerator, rmdp::RobustNoCrashMDP, s::MLPhysicalState, a::MLAction)
    behaviors = Dict{Int, Nullable{BehaviorModel}}(1=>nothing)
    for c in s.cars
        # TODO instead of just using sample, try to pick models that will be bad
        behaviors[c.id] = rand(gen.rng, rmdp.base.dmodel.behaviors)
    end
    return FixedBehaviorNoCrashMDP(behaviors, rmdp.base)
end
=#

abstract type EmbeddedBehaviorMDP <: MDP{MLPhysicalState, MLAction} end

actions(p::EmbeddedBehaviorMDP) = actions(p.base)
actions(p::EmbeddedBehaviorMDP, s::MLPhysicalState, as::NoCrashActionSpace) = actions(p.base, s, as)
# create_action(p::EmbeddedBehaviorMDP) = create_action(p.base)
create_state(p::EmbeddedBehaviorMDP) = MLPhysicalState(false, 0.0, 0.0, [])
discount(p::EmbeddedBehaviorMDP) = discount(p.base)
# reward(p::EmbeddedBehaviorMDP, s::MLPhysicalState, a::MLAction, sp::MLPhysicalState)

mutable struct FixedBehaviorNoCrashMDP <: EmbeddedBehaviorMDP
    behaviors::Dict{Int,Nullable{BehaviorModel}}
    base::MLMDP
end

function generate_sr(mdp::FixedBehaviorNoCrashMDP, s::MLPhysicalState, a::MLAction, rng::AbstractRNG, sp::MLPhysicalState=create_state(mdp))
    full_s = MLState(s, Vector{CarState}(length(s.cars)))
    for (i,c) in enumerate(s.cars)
        full_s.cars[i] = CarState(c, mdp.behaviors[c.id])
    end
    full_sp = generate_s(mdp.base, full_s, a, rng)
    for c in sp.cars
        if !haskey(mdp.behaviors, c.id)
            mdp.behaviors[c.id] = c.behavior
        end
    end
    return MLPhysicalState(full_sp), reward(mdp.base, full_s, a, full_sp)
end

mutable struct StochasticBehaviorNoCrashMDP <: EmbeddedBehaviorMDP
    base::MLMDP
end

function generate_sr(mdp::StochasticBehaviorNoCrashMDP, s::MLPhysicalState, a::MLAction, rng::AbstractRNG, sp::MLPhysicalState=create_state(mdp))
    full_s = MLState(s, Vector{CarState}(length(s.cars)))
    for i in 1:length(s.cars)
        full_s.cars[i] = CarState(s.cars[i], rand(rng, mdp.base.dmodel.behaviors))
    end
    full_sp = generate_s(mdp.base, full_s, a, rng)
    return MLPhysicalState(full_sp), reward(mdp.base, full_s, a, full_sp)
end

function initial_state(mdp::Union{FixedBehaviorNoCrashMDP,StochasticBehaviorNoCrashMDP}, rng::AbstractRNG, s=nothing)
    full_state = initial_state(mdp.base, rng::AbstractRNG)
    return MLPhysicalState(full_state)
end

#=
type RobustMLSolver <: Solver
    rsolver
end

type RobustMLPolicy <: Policy{MLState}
    rpolicy
end

function action(p::RobustMLPolicy, s::MLState, a::MLAction=MLAction(0,0))
    mdp = representative_mdp(p.rpolicy.rmdp)
    as = actions(mdp, MLPhysicalState(s), actions(mdp))
    if length(as) == 1
        return collect(as)[1]
    end
    action(p.rpolicy, MLPhysicalState(s))
end

function solve(s::RobustMLSolver, mdp::MLMDP)
    RobustMLPolicy(solve(s.rsolver, RobustNoCrashMDP(mdp::MLMDP)))
end
=#


# function that handles the dispatch of cars
function incident_dispatch!(sim_data::DataFrame,
                            incidents::DataFrame,
                            ressource_flow::Array{Int64,3},
                            drivingtime::Array{Float64,2},
                            traffic::Array{Float64,2},
                            locations::Vector{Int64},
                            current_case::Int64,
                            now::Int64,
                            district::Int64,
                            k::Int64,
                            dispatch::String,
                            fastest_time::Int64)
#  create the neccessary variables
    inc::Int64           = incidents[current_case,:location] # incident_location
    dc::Int64            = locations[district]               # location district center

    # create a vector with all necessary timestamps
    # timestamps[1]: minute of arrival at incident
    # timestamps[2]: minute of departure at incident
    # timestamps[3]: minute of arrival at district center
    timestamps = Vector{Int64}(undef,3)

# calculate the minute the car will arrive at the incident location
    if fastest_time == 0
        timestamps[1] = now + dispatch_drivingtime(drivingtime[dc,inc],traffic,now)
    else
        timestamps[1] = now + fastest_time
    end

# calculate the minute of arrival at the incident and the minute of departure from the incident
    if sim_data[current_case,:cars_dispatched] == 0
        timestamps[2] = timestamps[1] + incidents[current_case,:length]
        timestamps[3] = timestamps[2] + dispatch_drivingtime(drivingtime[inc,dc],traffic,timestamps[2])
        ressource_flow[timestamps[3],district,7] += incidents[current_case,:backlog]
    else
        timestamps[2] = min(timestamps[1] + incidents[current_case,:length], timestamps[1] +
                        incidents[current_case,:length] - sim_data[current_case,:dispatch_minute_first] + now)
        timestamps[3] = timestamps[2] + dispatch_drivingtime(drivingtime[inc,dc],traffic,timestamps[2])
    end

# calculate the number of cars to dispatch to the incident
    # cars[1]:  number of cars dispatched from location currently free
    # cars[2]:  number of cars dispatched from location currently working on backlog
    # cars[3]:  number of cars dispatched currently on patrol
    cars = Vector{Int64}(undef,3) .= 0

    # k = 1: only one car will be dispatched
    # k = 2: the number of dispatched cars dependends on the availability
    if dispatch == "own_department"
        if ressource_flow[now,district,1] > 0
            if k == 1
                cars[1] = 1
            else
                cars[1] = min(sim_data[current_case,:cars_missing], ressource_flow[now,district,1])
            end
        else
            cars[2] = 1
        end
    elseif dispatch == "own_patrol"
        cars[3] = 1
    elseif dispatch == "exchange_department"
        cars[1] = 1
    end

    # calculate the numer of dispatched cars
    total = sum(cars)

    # assign the calculated capacities into the ressource flow matrix
    ressource_flow[now,district,1]             -= cars[1]
    ressource_flow[now,district,5]             -= cars[2]
    ressource_flow[now,district,6]             -= cars[3]
    ressource_flow[now,district,2]             += total
    ressource_flow[timestamps[1],district,2]   -= total
    ressource_flow[timestamps[1],district,3]   += total
    ressource_flow[timestamps[2],district,3]   -= total
    ressource_flow[timestamps[2],district,4]   += total
    ressource_flow[timestamps[3],district,4]   -= total
    ressource_flow[timestamps[3],district,1]   += total

    # write the resulting times to the sim_data DataFrame
    if sim_data[current_case,:cars_dispatched] == 0
        sim_data[current_case,:location_dispatched_first] = locations[district]
        sim_data[current_case,:dispatch_minute_first]     = now
        sim_data[current_case,:arrival_minute_first]      = timestamps[1]
    end
    sim_data[current_case,:arrival_minute_first]    = min(sim_data[current_case,:arrival_minute_first],timestamps[1])
    sim_data[current_case,:cars_dispatched]         += total
    sim_data[current_case,:cars_missing]            -= total
    if sim_data[current_case,:cars_missing] == 0
        sim_data[current_case,:dispatch_minute_all] = now
        sim_data[current_case,:arrival_minute_all]  = timestamps[1]
    end
    if  dispatch == "own_department" || dispatch == "own_patrol"
        sim_data[current_case,:cars_location_responsible] += total
    end
end


#=
    OLD OUT DATED CODE
=#

struct Disk_old
    identifier::String
    color::String
    location::Tuple{Float64, Float64}
    radius::Float64
    parameters::Vector{Disk_old}
end

"""
    add_disk(rt::Disk_old, parent, child = [])

The basic tree creation for disk operads. Specifying the new Disk_old, the parent identifier, and any potential children.
"""
function add_disk(rt::Disk_old, parent, child = [])
    found = false

    if parent == rt.identifier
        push!(rt.parameters, child)
        return (rt, true)
    end

    for i in rt.parameters
        found_m = false
        (i, found_m) = add_disk(i, parent, child)
        if found_m == true
            found = true
        end
    end

    return rt, found
end

"""
    draw_disk(to_draw, origin, dash)
Basics for drawing a disk. THe origin is what the disk will be relative to, and dash will be used for circles with outlines
"""
function draw_disk(to_draw, origin, dash)
    center = (origin[1] + (to_draw.location[1] * cosd(to_draw.location[2])), origin[2] + (to_draw.location[1] * sind(to_draw.location[2])))
    return compose(context(), circle(center[1], center[2], to_draw.radius), fill(to_draw.color))
    #if dash == 0
    #    circle = compose(context(), circle(center[1], center[2], to_draw.radius), fill(to_draw.color))
    #elseif dash == 1
    #    circle = compse(context(), arc(center[1], center[2], to_draw.radius, range(0, step=pi/4, length=8), range(pi/8, step=pi/4, length = 8)), fill(to_draw.color))
    #else
    #    circle = 0;
    #end
end


using POMDPs
using POMDPPolicies
using POMDPSimulators
using POMDPModels
using BenchmarkTools
using Random

mdps = [LegacyGridWorld(terminals = Set()),
        SimpleGridWorld(terminate_in = Set()),
       ]

for m in mdps
    @show typeof(m)
    policy = RandomPolicy(m, rng=MersenneTwister(7))
    rosim = RolloutSimulator(max_steps=10_000, rng=MersenneTwister(2))
    @btime simulate($rosim, $m, $policy)
end


"""
ULMFiT - Text Classifier

This is wrapper around the LanguageMode struct. It has three fields:

vocab           : contains the same vocabulary from the LanguageModel
rnn_layers      : contains same DroppedEmebeddings, LSTM (AWD_LSTM) and VarDrop layers of LanguageModel except for last softmax layer
linear_layers   : contains Chain of two Dense layers [PooledDense and Dense] with softmax layer

To train create and instance and give it as first argument to 'train_classifier!' function
"""
mutable struct TextClassifier
    vocab::Vector
    rnn_layers::Flux.Chain
    linear_layers::Flux.Chain
end

function TextClassifier(lm::LanguageModel=LanguageModel(), clsfr_out_sz::Integer=1, clsfr_hidden_sz::Integer=50, clsfr_hidden_drop::Float64=0.4)
    return TextClassifier(
        lm.vocab,
        lm.layers[1:8],
        Chain(
            gpu(PooledDense(length(lm.layers[7].layer.cell.h), clsfr_hidden_sz)),
            gpu(BatchNorm(clsfr_hidden_sz, relu)),
            Dropout(clsfr_hidden_drop),
            gpu(Dense(clsfr_hidden_sz, clsfr_out_sz)),
            gpu(BatchNorm(clsfr_out_sz)),
            softmax
        )
    )
end

Flux.@functor TextClassifier

"""
Cross Validate

This function will be used to cross-validate the classifier

Arguments:

tc              : Instance of TextClassfier
gen             : 'Channel' to get a mini-batch from validation set
num_of_batches  : specifies the number of batches the validation will be done

If num_of_batches is not specified then all the batches which can be given by the
gen will be used for validation
"""
function validate(tc::TextClassifier, gen::Channel, num_of_batches::Union{Colon, Integer})
    n_classes = size(tc.linear_layers[end-2].W, 1)
    classifier = tc
    Flux.testmode!(classifier)
    loss = 0
    iters = take!(gen)
    ((num_of_batches != :) & (num_of_batches < iters)) && (iters = num_of_batches)
    TP, TN = gpu(zeros(Float32, n_classes, 1)), gpu(zeros(Float32, n_classes, 1))
    FP, FN = gpu(zeros(Float32, n_classes, 1)), gpu(zeros(Float32, n_classes, 1))
    for i=1:num_of_batches
        X = take!(gen)
        Y = gpu(take!(gen))
        X = map(x -> indices(x, classifier.vocab, "_unk_"), X)
        H = classifier.rnn_layers.(X)
        H = classifier.linear_layers(H)
        l = crossentropy(H, Y)
        Flux.reset!(classifier.rnn_layers)
        TP .+= sum(H .* Y, dims=2)
        FN .+= sum(((-1 .* H) .+ 1) .* Y, dims=2)
        FP .+= sum(H .* ((-1 .* Y) .+ 1), dims=2)
        TN .+= sum(((-1 .* H) .+ 1) .* ((-1 .* Y) .+ 1), dims=2)
        loss += l
    end
    precisions = TP ./ (TP .+ FP)
    recalls = TP ./ (TP .+ FN)
    F1 = (2 .* (precisions .* recalls)) ./ (precisions .+ recalls)
    accuracy = (TP[1] + TN[1])/(TP[1] + TN[1] + FP[1] + FN[1])
    return (loss, accuracy, precisions, recalls, F1)
end

"""
Forward pass

This funciton does the main computation of a mini-batch.
It computes the output of the all the layers [RNN and DENSE layers] and returns the predicted output for that pass.
It uses Truncated Backprop through time to compute the output.

Arguments:
tc              : Instance of TextClassifier
gen             : data loader, which will give 'X' of the mini-batch in one call
tracked_steps   : This is the number of tracked time-steps for Truncated Backprop thorugh time,
                  these will be last time-steps for which gradients will be calculated.
"""
function forward(tc::TextClassifier, gen::Channel, tracked_steps::Integer=32)
  	# swiching off tracking
    classifier = tc
    X = take!(gen)
    l = length(X)
    # Truncated Backprop through time
    Zygote.ignore() do
	for i=1:ceil(l/tracked_steps)-1   # Tracking is swiched off inside this loop
	    (i == 1 && l%tracked_steps != 0) ? (last_idx = l%tracked_steps) : (last_idx = tracked_steps)
	    H = broadcast(x -> indices(x, classifier.vocab, "_unk_"), X[1:last_idx])
	    H = classifier.rnn_layers.(H)
	    X = X[last_idx+1:end]
	end
    end
    # set the lated hidden states to original model
    for (t_layer, unt_layer) in zip(tc.rnn_layers[2:end], classifier.rnn_layers[2:end])
        if t_layer isa AWD_LSTM
            t_layer.layer.state = unt_layer.layer.state
            continue
        end
        if !unt_layer.reset
            t_layer.mask = unt_layer.mask
            t_layer.reset = false
        end
    end
    # last part of the sequecnes in X - Tracking is swiched on
    H = broadcast(x -> tc.rnn_layers[1](indices(x, classifier.vocab, "_unk_")), X)
    H = tc.rnn_layers[2:end].(H)
    H = tc.linear_layers(H)
    return H
end

"""
    loss(classifier::TextClassifier, gen::Channel, tracked_steps::Integer=32)

LOSS function

It takes the output of the forward funciton and returns crossentropy loss.

Arguments:

classifier    : Instance of TextClassifier
gen           : 'Channel' [data loader], to give a mini-batch
tracked_steps : specifies the number of time-steps for which tracking is on
"""
function loss(classifier::TextClassifier, gen::Channel, tracked_steps::Integer=32)
    H = forward(classifier, gen, tracked_steps)
    Y = gpu(take!(gen))
    l = crossentropy(H, Y)
    Flux.reset!(classifier.rnn_layers)
    return l
end

function discriminative_step!(layers, classifier::TextClassifier, gen::Channel, tracked_steps::Integer, ηL::Float64, opts::Vector)
    @assert length(opts) == length(layers)
    # Gradient calculation
    grads = Zygote.gradient(() -> loss(classifier, gen, tracked_steps = tracked_steps), get_trainable_params(layers))

    # discriminative step
    ηl = ηL/(2.6^(length(layers)-1))
    for (layer, opt) in zip(layers, opts)
        opt.eta = ηl
        for ps in get_trainable_params([layer])
            Flux.Optimise.update!(opt, ps, grads[ps])
        end
        ηl *= 2.6
    end
    return
end

"""
    train_classifier!(classifier::TextClassifier=TextClassifier(), classes::Integer=1,
            data_loader::Channel=imdb_classifier_data, hidden_layer_size::Integer=50;kw...)

It contains main training loops for training a defined classifer for specified classes and data.
Usage is discussed in the docs.
"""
function train_classifier!(classifier::TextClassifier=TextClassifier(), classes::Integer=1,
    data_loader::Channel=imdb_classifier_data, hidden_layer_size::Integer=50;
    stlr_cut_frac::Float64=0.1, stlr_ratio::Number=32, stlr_η_max::Float64=0.01,
    val_loader::Channel=nothing, cross_val_batches::Union{Colon, Integer}=:,
    epochs::Integer=1, checkpoint_itvl=5000, tracked_steps::Integer=32)

    trainable = []
    append!(trainable, [classifier.rnn_layers[[1, 3, 5, 7]]...])
    push!(trainable, [classifier.linear_layers[1:2]...])
    push!(trainable, [classifier.linear_layers[4:5]...])
    opts = [ADAM(0.001, (0.7, 0.99)) for i=1:length(trainable)]
    gpu!.(classifier.rnn_layers)

    for epoch=1:epochs
        println("Epoch: $epoch")
        gen = data_loader()
        num_of_iters = take!(gen)
        cut = num_of_iters * epochs * stlr_cut_frac
        for iter=1:num_of_iters

            # Slanted triangular learning rates
            t = iter + (epoch-1)*num_of_iters
            p_frac = (iter < cut) ? iter/cut : (1 - ((iter-cut)/(cut*(1/stlr_cut_frac-1))))
            ηL = stlr_η_max*((1+p_frac*(stlr_ratio-1))/stlr_ratio)

            # Gradual-unfreezing Step with discriminative fine-tuning
            unfreezed_layers, cur_opts = (epoch < length(trainable)) ? (trainable[end-epoch+1:end], opts[end-epoch+1:end]) : (trainable, opts)
            discriminative_step!(unfreezed_layers, classifier, gen, tracked_steps,ηL, cur_opts)

            reset_masks!.(classifier.rnn_layers)    # reset all dropout masks
        end
        println("Train set accuracy: $trn_accu , Training loss: $trn_loss")
        if val_loader != nothing
            val_loss, val_acc, val_precisions, val_reacalls, val_F1_scores = validate(classifer, val_loader)
        else
            continue
        end
        #!(val_loader isa nothing) ? (val_loss, val_acc, val_precisions, val_reacalls, val_F1_scores = validate(classifer, val_loader)) : continue
        println("Cross validation loss: $val_loss")
        println("Cross validation accuracy:\n $val_acc")
        println("Cross validation class wise Precisions:\n $val_precisions")
        println("Cross validation class wise Recalls:\n $val_recalls")
        println("Cross validation class wise F1 scores:\n $val_F1_scores")
    end
end

"""
    predict(tc::TextClassifier, text_sents::Corpus)

This function can be used to test the model after training.
It returns the predictions done by the model for given `Corpus` of `Documents`
All the preprocessing related to the used vocabulary should be done before using this function.
Use `prepare!` function to do preprocessing
"""
function predict(tc::TextClassifier, text_sents::Corpus)
    classifier = tc
    Flux.testmode!(classifier)
    predictions = []
    expr(x) = indices(x, classifier.vocab, "_unk_")
    for text in text_sents
        tokens_ = tokens(text)
        h = classifier.rnn_layers.(expr.(tokens_))
        probability_dist = classifier.linear_layers(h)
        class = argmax(probaility_dist)
        push!(predictions, class)
    end
    return predictions
end


using BenchmarkTools
using Meshing
using GeometryTypes

# Define a parent BenchmarkGroup to contain our suite
const suite = BenchmarkGroup()
suite["SDF Construction"] = BenchmarkGroup()
suite["SDF Mesh"] = BenchmarkGroup()
suite["Function Mesh"] = BenchmarkGroup()

println("Benchmarking Meshing.jl...")

#
# Algorithms to benchmark
#

algos_sdf = [MarchingCubes(), MarchingTetrahedra(), NaiveSurfaceNets()]
algos_fn = [MarchingCubes(), NaiveSurfaceNets()]

#
# Benchmark SDF constructon for SDF and Direct sample comparisons
#

suite["SDF Construction"]["Torus"] = @benchmarkable SignedDistanceField(HyperRectangle(Vec(-2,-2,-2.),Vec(4,4,4.)),0.05) do v
    (sqrt(v[1]^2+v[2]^2)-0.5)^2 + v[3]^2 - 0.25 # torus
end

#
# Torus Constructions
#

sdf_torus = SignedDistanceField(HyperRectangle(Vec(-2,-2,-2.),Vec(4,4,4.)),0.05) do v
    (sqrt(v[1]^2+v[2]^2)-0.5)^2 + v[3]^2 - 0.25 # torus
end

fn_torus(v) = (sqrt(v[1]^2+v[2]^2)-0.5)^2 + v[3]^2 - 0.25 # torus

#
# Benchmark algorithms
#

for algo in algos_sdf
    suite["SDF Mesh"][string(typeof(algo))] = @benchmarkable HomogenousMesh(sdf_torus, $algo)
end

for algo in algos_fn
    suite["Function Mesh"][string(typeof(algo))] = @benchmarkable HomogenousMesh(fn_torus, HyperRectangle(Vec(-2,-2,-2.),Vec(4,4,4.)), (81,81,81), $algo)
end

# If a cache of tuned parameters already exists, use it, otherwise, tune and cache
# the benchmark parameters. Reusing cached parameters is faster and more reliable
# than re-tuning `suite` every time the file is included.
paramspath = joinpath(dirname(@__FILE__), "params.json")

if isfile(paramspath)
    loadparams!(suite, BenchmarkTools.load(paramspath)[1], :evals);
else
    tune!(suite)
    BenchmarkTools.save(paramspath, params(suite));
end

#
# Perform benchmarks and print results
#

results = run(suite)

for trial in results
    ctx = IOContext(stdout, :verbose => true, :compact => false)
    println(ctx)
    println(ctx, trial.first)
    println(ctx, trial.second)
end


"""
    inbounds(x, max, overflow)

Check grid boundaries for a single coordinate and max value or a tuple
of coorinates and max values.

Returns a tuple containing the coordinate(s) followed by a boolean `true`
if the cell is in bounds, `false` if not.

Overflow of type [`RemoveOverflow`](@ref) returns the coordinate and `false` to skip
coordinates that overflow outside of the grid.
[`WrapOverflow`](@ref) returns a tuple with the current position or it's
wrapped equivalent, and `true` as it is allways in-bounds.
"""
@inline inbounds(xs::Tuple, maxs::Tuple, overflow) = begin
    a, inbounds_a = inbounds(xs[1], maxs[1], overflow)
    b, inbounds_b = inbounds(xs[2], maxs[2], overflow)
    (a, b), inbounds_a & inbounds_b
end
@inline inbounds(x::Number, max::Number, overflow::RemoveOverflow) =
    x, isinbounds(x, max, overflow)
@inline inbounds(x::Number, max::Number, overflow::WrapOverflow) =
    if x < oneunit(x)
        max + rem(x, max), true
    elseif x > max
        rem(x, max), true
    else
        x, true
    end

@inline isinbounds(xs::Tuple, maxs::Tuple, overflow) = all(isinbounds.(xs, maxs, Ref(overflow)))
@inline isinbounds(x::Number, max::Number, overflow::RemoveOverflow) =
    x > zero(x) && x <= max

"""
    sizefromradius(radius)

Get the size of a neighborhood dimension from its radius,
which is always 2r + 1.
"""
hoodsize(radius::Integer) = 2radius + 1

"""
Return a tuple of the base types of the rules in the ruleset
"""
ruletypes(ruleset::Ruleset) = ruletypes(typeof(Ruleset.rules))
ruletypes(t::Type) = t.name.wrapper
ruletypes(ts::Type{<:Tuple}) = (ruletypes.(ts.parameters)...,)

@inline ismasked(data::AbstractSimData, i...) = ismasked(mask(data), i...)
@inline ismasked(mask::Nothing, i...) = false
@inline ismasked(mask::AbstractArray, i...) = @inbounds return !mask[i...]


"""
    convert_units!(p::Particles, h, units::Symbol=:full)

Converts the units of the particles properties (see [`Particles`](@ref)) from a snapshot or
subfind header in-place.
Use `:full` for units as `Unitful` quantities, `:physical` for values converted
to physical units (kpc, km/s, solar metallicities etc.), or `:sim` for values in simulation units.
The individual metallicities in "Zs" are converted to a one-dimensional `Vector` of values in solar
metallicities in any case.
"""
function convert_units!(p::Particles, h, units::Symbol=:full)
    if units === :physical
        for (key, vals) in pairs(p)
            convert_units_physical!(vals, key, h)
        end
    elseif units === :full
        for (key, vals) in pairs(p)
            p[key] = convert_units_full(vals, key, h)
        end
    end

    props = keys(p)
    if :zs in props && length(p.zs) > 0 && (:im in props || :mass in props)
        p.zs = convert_units_solar_metallicity(p.zs, :im in props ? p.iM : p.mass)
    end

    return p
end

function convert_units_subfind_prop(
    val::Union{Real,AbstractArray{<:Real}},
    prop::AbstractString,
    h,
    units::Symbol=:full;
    verbose::Bool=false,
)
    # TODO: SPIN (units: angular momentum), DSUB (vel. dispersion units?),
    # SLUM, SLAT, SLOB, DUST, SZ (units?), SSFR (units really M⊙/yr?)
    if prop[1] === 'R' || prop in ["GPOS", "BGPO", "BGRA", "SPOS", "SCM", "SHMR"]
        return convert_units_pos(val, h, units)
    elseif prop[1] === 'V' || prop in ["SVEL"]
        return convert_units_vel(val, h, units)
    elseif (prop[1] === 'M' && prop != "MBID") || prop in ["BGMA", "SMST"]
        return convert_units_mass(val, h, units)
    elseif prop in ["SAGE"]
        return convert_units_age(val, h, units)
        # elseif prop in ["TGAS"] # TODO: unclear what KeV unit is
        # return convert_units_temp(val, h, units)
        # elseif prop in ["LGAS"] # TODO: energy unit conversion (10^44 erg/s)
        # return convert_units_temp(val, h, units)
    elseif prop in ["SSFR"] && units === :full
        return val * u"Msun" / Unitful.yr
    else
        verbose && @warn "The quantity $prop is returned in simulation units"
        return val
    end
end

# define generated functions for static linting
function convert_units_full end
function convert_units_physical end
function convert_units_physical! end

function simulation_units_pos end
function simulation_units_pos! end
function convert_units_pos end
function convert_units_full_pos end
function convert_units_physical_pos end
function convert_units_physical_pos! end
function simulation_units_vel end
function simulation_units_vel! end
function convert_units_vel end
function convert_units_full_vel end
function convert_units_physical_vel end
function convert_units_physical_vel! end
function simulation_units_temp end
function simulation_units_temp! end
function convert_units_temp end
function convert_units_full_temp end
function convert_units_physical_temp end
function convert_units_physical_temp! end
function simulation_units_mass end
function simulation_units_mass! end
function convert_units_mass end
function convert_units_full_mass end
function convert_units_physical_mass end
function convert_units_physical_mass! end

for (type, excl) in [("physical", ""), ("physical", "!"), ("full", "")]
    quote
        function $(Symbol("convert_units_", type, excl))(vals::AbstractArray{<:Real}, prop::Symbol, h)
            if prop === :pos
                $(Symbol("convert_units_", type, "_pos", excl))(vals, h)
            elseif prop === :vel
                $(Symbol("convert_units_", type, "_vel", excl))(vals, h)
            elseif prop === :age
                $(Symbol("convert_units_", type, "_age", excl))(vals, h)
            elseif prop === :temp
                $(Symbol("convert_units_", type, "_temp", excl))(vals, h)
            elseif prop in [:mass, :zs, :im]
                $(Symbol("convert_units_", type, "_mass", excl))(vals, h)
            else
                vals
            end
        end
    end |> eval
end
@doc """
    convert_units_physical(vals::AbstractArray{<:Real}, prop::Symbol, h)
    convert_units_physical!(vals::AbstractArray{<:Real}, prop::Symbol, h)
    convert_units_full(vals::AbstractArray{<:Real}, prop::Symbol, h)

Converts simulation values to the respective physical values, depending on `prop`,
which can take any of the following values: `:pos`, `:vel`, `:temp`, `:mass`, `:age`.

Full returns values in `Unitful` quantities, whereas physical returns the physical value without unit.
"""
convert_units_physical,
#convert_units_physical!,
convert_units_full

for (key, factor, unit, plural, eq, eqsim) in [
    (
        "pos",
        :(1 / (h.h0 * (h.z + 1))),
        :(u"kpc"),
        "positions",
        raw"x \to x/(h_0 (z+1)))",
        raw"x \to x \times h_0 (z+1)",
    ),
    (
        "vel",
        :(1 / sqrt(h.z + 1)),
        :(u"km/s"),
        "velocities",
        raw"v \to v/\sqrt{z+1})",
        raw"v \to v \times \sqrt{z+1}",
    ),
    ("temp", :(1), :(u"K"), "temperatures", raw"T \to T", raw"T \to T"),
    (
        "mass",
        :(1e10 / h.h0),
        :(u"Msun"),
        "masses",
        raw"m \to m \times 10^{10} / h_0",
        raw"m \to m \times h_0 / 10^{10}",
    ),
]
    quote
        function $(Symbol("convert_units_physical_", key))(val::T, h) where {T<:Real}
            val * convert(T, $factor)
        end
        function $(Symbol("convert_units_full_", key))(val::Real, h)
            $(Symbol("convert_units_physical_", key))(val, h) * $unit
        end
        function $(Symbol("convert_units_physical_", key))(vals::AbstractArray{T}, h) where {T<:Real}
            vals .* convert(T, $factor)
        end
        function $(Symbol("convert_units_physical_", key, "!"))(vals::AbstractArray{T}, h) where {T<:Real}
            vals .*= convert(T, $factor)
        end
        function $(Symbol("convert_units_full_", key))(vals::AbstractArray{T}, h) where {T<:Real}
            vals .* (convert(T, $factor) * $unit)
        end
        function $(Symbol("convert_units_", key))(
            val::Union{Real,AbstractArray{<:Real}},
            h,
            units::Symbol=:full,
        )
            if units === :physical
                $(Symbol("convert_units_physical_", key))(val, h)
            elseif units === :full
                $(Symbol("convert_units_full_", key))(val, h)
            else
                val
            end
        end
        function $(Symbol("simulation_units_", key))(val::Quantity{T}, h) where {T}
            ustrip($unit, val) / convert(T, $factor)
        end
        function $(Symbol("simulation_units_", key))(val::T, h) where {T<:Real}
            val / convert(T, $factor)
        end
        function $(Symbol("simulation_units_", key))(vals::AbstractArray{<:Quantity{T}}, h) where {T}
            ustrip.($unit, vals) ./ convert(T, $factor)
        end
        function $(Symbol("simulation_units_", key))(vals::AbstractArray{T}, h) where {T<:Real}
            vals ./ convert(T, $factor)
        end
        function $(Symbol("simulation_units_", key, "!"))(vals::AbstractArray{T}, h) where {T<:Real}
            vals ./= convert(T, $factor)
        end
    end |> eval

    quote
        @doc """
             convert_units_$($key)(val::Union{Real,AbstractArray{<:Real}}, h, units::Symbol=:full)
             convert_units_physical_$($key)(val::Real, h)
             convert_units_physical_$($key)(vals::AbstractArray{<:Real}, h)
             convert_units_physical_$($key)!(vals::AbstractArray{<:Real}, h)
             convert_units_full_$($key)(val::Real, h)
             convert_units_full_$($key)(vals::AbstractArray{<:Real}, h)

         Converts $($plural) via ``$($eq)`` according to the cosmology defined by the snapshot header.
         `:full` returns values in `Unitful` quantities, whereas `:physical` returns the physical value
         in $(eval($unit)). `:sim` simply returns `val`.
         """
        $(Symbol("convert_units_", key))#,
        # $(Symbol("convert_units_physical_", key)),
        # $(Symbol("convert_units_physical_", key, "!")),
        # $(Symbol("convert_units_full_", key))

        @doc """
             simulation_units_$($key)(val::Number, h)
             simulation_units_$($key)(vals::AbstractArray{<:Number}, h)
             simulation_units_$($key)!(vals::AbstractArray{<:Real}, h)

         Converts $($plural) via ``$($eqsim)`` according to the cosmology defined by the snapshot
         header from physical to simulation units.
         Unitless parameters need to be given in $(eval($unit)).
         """
        $(Symbol("simulation_units_", key))#,
        # $(Symbol("simulation_units_", key, "!"))
    end |> eval
end

@doc raw"""
    convert_units_age(val::Union{Real,AbstractArray{<:Real}}, h, units::Symbol=:full)
    convert_units_physical_age(val::Real, h)
    convert_units_physical_age(vals::AbstractArray{<:Real}, h)
    convert_units_physical_age!(vals::AbstractArray{<:Real}, h)
    convert_units_full_age(val::Real, h)
    convert_units_full_age(vals::AbstractArray{<:Real}, h)

Converts ages from the scale factor ``a`` via the lookback time according to the cosmology
defined by the snapshot header.
Full returns values in `Unitful` quantities, whereas physical returns the physical value in Gyr.
"""
function convert_units_age(val::Union{Real,AbstractArray{<:Real}}, h, units::Symbol=:full)
    if units === :physical
        convert_units_physical_age(val, h)
    elseif units === :full
        convert_units_full_age(val, h)
    else
        val
    end
end
function convert_units_physical_age(val::Real, h)
    convert_units_full_age(val, h) |> ustrip
end
function convert_units_full_age(val::T, h) where {T<:Real}
    c = cosmology(h=h.h0, OmegaM=h.omega_0)
    lookback_time(c, 1 / val - 1) - lookback_time(c, h.z) # in Gyr
end
function convert_units_physical_age(vals::AbstractArray{<:Real}, h)
    convert_units_full_age(vals, h) .|> ustrip
end
function convert_units_physical_age!(vals::AbstractArray{<:Real}, h)
    vals .= convert_units_physical_age(vals, h)
end
function convert_units_full_age(vals::AbstractArray{T}, h) where {T<:Real}
    c = cosmology(h=h.h0, OmegaM=h.omega_0)
    lookback_time.((c,), 1 ./ vals .- 1) .- lookback_time(c, h.z) # in Gyr
end

"""
    convert_units_solar_metallicity(zs::AbstractMatrix{<:Number}, mass::AbstractVector{<:Number})

Converts an `11×n` mass matrix and an `n`-length mass vector to a metallicity vector of length `n`,
in solar metallicities.
"""
function convert_units_solar_metallicity(zs::AbstractMatrix{<:Number}, mass::AbstractVector{<:Number})
    dropdims(sum(zs[2:11, :]; dims=1); dims=1) ./ (mass .- dropdims(sum(zs; dims=1); dims=1)) ./ 0.0134 # in solar metallicities
end