#=
# FV solver
=#
module FVSolver

export solve, nonlinear_solve

# See finch_import_symbols.jl for a list of all imported symbols.
import ..Finch: @import_finch_symbols
@import_finch_symbols()

using LinearAlgebra, SparseArrays

include("fv_boundary.jl");

# Things to do before and after each step(or stage) ########
function default_pre_step() end
function default_post_step() end

pre_step_function = default_pre_step;
post_step_function = default_post_step;

function set_pre_step(fun) global pre_step_function = fun; end
function set_post_step(fun) global post_step_function = fun; end
############################################################

function init_solver()
    global pre_step_function = default_pre_step;
    global post_step_function = default_post_step;
end

function linear_solve(var, source_lhs, source_rhs, flux_lhs, flux_rhs, stepper=nothing, assemble_func=nothing)
    # Treat explicit and implicit solvers differently
    if stepper === nothing
        printerr("FV assumes a time dependent problem. Set time stepper and initial conditions.", fatal=true)
    elseif stepper.implicit
        return linear_solve_implicit(var, source_lhs, source_rhs, flux_lhs, flux_rhs, stepper, assemble_func)
    else # explicit
        return linear_solve_explicit(var, source_lhs, source_rhs, flux_lhs, flux_rhs, stepper, assemble_func)
    end
end

function linear_solve_explicit(var, source_lhs, source_rhs, flux_lhs, flux_rhs, stepper=nothing, assemble_func=nothing)
    # If more than one variable
    if typeof(var) <: Array
        # multiple variables being solved for simultaneously
        dofs_per_node = 0;
        dofs_per_loop = 0;
        for vi=1:length(var)
            dofs_per_loop += length(var[vi].symvar);
            dofs_per_node += var[vi].total_components;
        end
    else
        # one variable
        dofs_per_loop = length(var.symvar);
        dofs_per_node = var.total_components;
    end
    if fv_grid === nothing
        grid = grid_data;
        nel = grid_data.nel_owned;
        nfaces = size(grid_data.face2element, 2);
    else
        grid = fv_grid;
        nel = fv_grid.nel_owned;
        nfaces = size(fv_grid.face2element, 2);
    end

    Nn = dofs_per_node * nel;
    Nf = dofs_per_node * nfaces

    # Allocate arrays that will be used by assemble
    # These vectors will hold the integrated values(one per cell).
    # They will later be combined.
    sourcevec = zeros(Nn);
    fluxvec = zeros(Nn);
    facefluxvec = zeros(Nf);
    face_done = zeros(nfaces); # Increment when the corresponding flux value is computed.
    allocated_vecs = (sourcevec, fluxvec, facefluxvec, face_done);

    if prob.time_dependent && !(stepper === nothing)
        log_entry("Beginning "*string(stepper.Nsteps)*" time steps.");
        t = 0;
        if config.num_partitions > 1 exchange_ghosts(var, grid, 0); end
        sol = get_var_vals(var);

        # allocate storage used by steppers
        if stepper.type == LSRK4
            tmppi = zeros(size(sol));
            tmpki = zeros(size(sol));
        elseif stepper.type == RK4
            tmpvals = zeros(size(sol));
            tmpki = zeros(length(sol), stepper.stages);
        end

        start_t = Base.Libc.time();
        last2update = 0;
        last10update = 0;
        if config.proc_rank == 0 print("Time stepping progress(%): 0"); end
        for i=1:stepper.Nsteps
            if stepper.stages > 1
                # LSRK4 is a special case, low storage
                if stepper.type == LSRK4
                    # Low storage RK4:
                    # p0 = u
                    #   ki = ai*k(i-1) + dt*f(p(i-1), t+ci*dt)
                    #   pi = p(i-1) + bi*ki
                    # u = p5

                    tmppi = get_var_vals(var, tmppi);
                    tmpki = zeros(size(sol));
                    for rki=1:stepper.stages
                        rktime = t + stepper.c[rki]*stepper.dt;
                        # p(i-1) is currently in u

                        if config.num_partitions > 1 exchange_ghosts(var, grid, i); end
                        pre_step_function();

                        sol = assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, rktime, stepper.dt, assemble_loops=assemble_func);

                        if rki == 1 # because a1 == 0
                            tmpki = stepper.dt .* sol;
                        else
                            tmpki = stepper.a[rki].*tmpki + stepper.dt.*sol;
                        end
                        tmppi = tmppi + stepper.b[rki].*tmpki

                        FV_copy_bdry_vals_to_vector(var, tmppi, grid, dofs_per_node);
                        place_sol_in_vars(var, tmppi, stepper);

                        post_step_function();
                    end

                else
                    # Explicit multi-stage methods:
                    # x = x + dt*sum(bi*ki)
                    # ki = rhs(t+ci*dt, x+dt*sum(aij*kj)))   j < i

                    # will hold the final result
                    # sol = get_var_vals(var, sol);
                    # will be placed in var.values for each stage
                    for stage=1:stepper.stages
                        stime = t + stepper.c[stage]*stepper.dt;

                        # Update the values in vars to be used in this stage
                        if stage > 1
                            initialized_tmpvals = false;
                            for j=1:stage
                                if stepper.a[stage, j] > 0
                                    if !initialized_tmpvals
                                        initialized_tmpvals = true;
                                        for k=1:length(sol)
                                            tmpvals[k] = sol[k] + stepper.dt * stepper.a[stage, j] * tmpki[k,j];
                                        end
                                    else
                                        for k=1:length(sol)
                                            tmpvals[k] += stepper.dt * stepper.a[stage, j] * tmpki[k,j];
                                        end
                                    end
                                end
                            end

                            if initialized_tmpvals
                                FV_copy_bdry_vals_to_vector(var, tmpvals, grid, dofs_per_node);
                                place_sol_in_vars(var, tmpvals, stepper);
                            end
                            post_step_function(); # seems weird, but imagine this is happening after stage-1
                        end

                        if config.num_partitions > 1 exchange_ghosts(var, grid, i); end
                        pre_step_function();

                        tmpki[:,stage] = assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, stime, stepper.dt, assemble_loops=assemble_func);

                    end
                    # Stages are done. Assemble the final result for this step
                    for stage=1:stepper.stages
                        sol += stepper.dt * stepper.b[stage] .* tmpki[:, stage];
                    end
                    FV_copy_bdry_vals_to_vector(var, sol, grid, dofs_per_node);
                    place_sol_in_vars(var, sol, stepper);

                    post_step_function();
                end

            elseif stepper.type == EULER_EXPLICIT
                if config.num_partitions > 1 exchange_ghosts(var, grid, i); end
                pre_step_function();

                sol = sol .+ stepper.dt .* assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, stepper.dt, assemble_loops=assemble_func);

                FV_copy_bdry_vals_to_vector(var, sol, grid, dofs_per_node);
                place_sol_in_vars(var, sol, stepper);

                post_step_function();

            elseif stepper.type == PECE
                # Predictor (explicit Euler)
                if config.num_partitions > 1 exchange_ghosts(var, grid, i); end
                pre_step_function();
                tmpsol = sol .+ stepper.dt .* assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, stepper.dt, assemble_loops=assemble_func);

                FV_copy_bdry_vals_to_vector(var, tmpsol, grid, dofs_per_node);
                place_sol_in_vars(var, tmpsol, stepper);
                post_step_function();

                # Corrector (implicit Euler)
                if config.num_partitions > 1 exchange_ghosts(var, grid, i); end
                pre_step_function();
                sol = sol .+ stepper.dt .* assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t+stepper.dt, stepper.dt, assemble_loops=assemble_func);

                FV_copy_bdry_vals_to_vector(var, sol, grid, dofs_per_node);
                place_sol_in_vars(var, sol, stepper);
                post_step_function();

            else
                printerr("Only explicit time steppers for FV are ready. TODO")
                return sol;
            end

            # ########## uncomment to return after one time step
            # return sol
            # ############

            t += stepper.dt;

            if config.proc_rank == 0
                progressPercent = Int(floor(i*100.0/stepper.Nsteps));
                if progressPercent - last2update >= 2
                    last2update = progressPercent;
                    if progressPercent - last10update >= 10
                        print(progressPercent);
                        last10update = progressPercent;
                    else
                        print(".");
                    end
                end
            end
        end
        if config.proc_rank == 0 println(""); end
        if config.num_partitions > 1 exchange_ghosts(var, grid, 0); end
        end_t = Base.Libc.time();

        log_entry("Stepping took "*string(end_t-start_t)*" seconds.");
        return sol;

    else
        # Does it make any sense to do this for time-independent problems?
        return sol;
    end
end

function linear_solve_implicit(var, source_lhs, source_rhs, flux_lhs, flux_rhs, stepper=nothing, assemble_func=nothing)
    # If more than one variable
    if typeof(var) <: Array
        # multiple variables being solved for simultaneously
        dofs_per_node = 0;
        dofs_per_loop = 0;
        for vi=1:length(var)
            dofs_per_loop += length(var[vi].symvar);
            dofs_per_node += var[vi].total_components;
        end
    else
        # one variable
        dofs_per_loop = length(var.symvar);
        dofs_per_node = var.total_components;
    end
    if fv_grid === nothing
        grid = grid_data;
        nel = grid_data.nel_owned;
        nfaces = size(grid_data.face2element, 2);
    else
        grid = fv_grid;
        nel = fv_grid.nel_owned;
        nfaces = size(fv_grid.face2element, 2);
    end

    dofs_squared = dofs_per_node*dofs_per_node;
    Nn = dofs_per_node * nel; # dofs values for each cell
    Nf = dofs_per_node * nfaces; # dofs values per face
    Nn2 = dofs_squared * nel; # dofs values per dof per cell
    Nf2 = dofs_squared * nfaces * 4; # dofs values per dof per face * 4

    # Allocate arrays that will be used by assemble
    sourcevec = zeros(Nn);
    fluxvec = zeros(Nn);
    face_done = zeros(nfaces); # true when the corresponding flux value is computed.
    facefluxvec = zeros(Nf);
    lhsmatI = zeros(Int, Nn2+Nf2);
    lhsmatJ = zeros(Int, Nn2+Nf2);
    lhsmatV = zeros(Nn2+Nf2);
    allocated_vecs = (lhsmatI, lhsmatJ, lhsmatV, sourcevec, fluxvec, facefluxvec, face_done);

    # The spare matrix indices only need to be set once
    set_matrix_indices!(lhsmatI, lhsmatJ, dofs_per_node);

    sol = get_var_vals(var);

    if prob.time_dependent && !(stepper === nothing)
        log_entry("Beginning "*string(stepper.Nsteps)*" time steps.");
        t = 0;
        if config.num_partitions > 1 exchange_ghosts(var, grid, 0); end

        start_t = Base.Libc.time();
        last2update = 0;
        last10update = 0;
        if config.proc_rank == 0 print("Time stepping progress(%): 0"); end
        for i=1:stepper.Nsteps
            if stepper.type == EULER_IMPLICIT || stepper.type == CRANK_NICHOLSON
                if config.num_partitions > 1 exchange_ghosts(var, grid, i); end
                pre_step_function();

                b = assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, stepper.dt, assemble_loops=assemble_func, is_explicit=false);
                # Need to add var values to b for u(-) + dt*rhs
                # Need to add identity to A for u(+) + dt*lhs
                sol = get_var_vals(var, sol);
                for j=1:length(b)
                    b[j] += sol[j];
                end
                if dofs_per_node == 1
                    for j=1:nel
                        lhsmatV[j] += 1;
                    end
                else
                    for j=1:nel
                        for k=1:dofs_per_node
                            lhsmatV[(j-1)*dofs_squared + (k-1)*dofs_per_node + k] += 1;
                        end
                    end
                end

                # The matrix IJV vectors are now set.
                # The dt part should already be included.
                A = sparse(lhsmatI, lhsmatJ, lhsmatV);
                # scoper = sol;
                sol = A\b;

                FV_copy_bdry_vals_to_vector(var, sol, grid, dofs_per_node);
                place_sol_in_vars(var, sol, stepper);

                post_step_function();

            else
                printerr("Unknown time stepper: "*string(stepper), fatal=true)
            end

            # ########## uncomment to return after one time step
            # return sol
            # ############

            t += stepper.dt;

            if config.proc_rank == 0
                progressPercent = Int(floor(i*100.0/stepper.Nsteps));
                if progressPercent - last2update >= 2
                    last2update = progressPercent;
                    if progressPercent - last10update >= 10
                        print(progressPercent);
                        last10update = progressPercent;
                    else
                        print(".");
                    end
                end
            end
        end
        if config.proc_rank == 0 println(""); end
        if config.num_partitions > 1 exchange_ghosts(var, grid, 0); end
        end_t = Base.Libc.time();

        log_entry("Stepping took "*string(end_t-start_t)*" seconds.");
        return sol;

    else
        # Does it make any sense to do this for time-independent problems?
        return sol;
    end
end

#
function assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node=1, dofs_per_loop=1, t=0, dt=0; assemble_loops=nothing, is_explicit=true)
    # If an assembly loop function was provided, use it
    if !(assemble_loops === nothing)
        return assemble_loops.func(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, dt, is_explicit=is_explicit);
    end
    # If parent maps were created for high order flux, use a slightly different function
    if !(parent_maps === nothing)
        return assemble_using_parent_child(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, dt, is_explicit=is_explicit);
    end

    nel = fv_grid.nel_owned;
    dofs_squared = dofs_per_node*dofs_per_node;

    # Label things that were allocated externally
    if is_explicit
        sourcevec = allocated_vecs[1];
        fluxvec = allocated_vecs[2];
        facefluxvec = allocated_vecs[3];
        face_done = allocated_vecs[4];
    else
        # lhsmatI = allocated_vecs[1]; # Don't need these here. Already set
        # lhsmatJ = allocated_vecs[2];
        lhsmatV = allocated_vecs[3];
        sourcevec = allocated_vecs[4]; # In the implicit case, these vectors are for the RHS
        fluxvec = allocated_vecs[5];
        facefluxvec = allocated_vecs[6];
        face_done = allocated_vecs[7];

        for i=1:length(lhsmatV)
            lhsmatV[i] = 0;
        end
    end

    for i=1:length(face_done)
        face_done[i] = 0;
    end

    # Elemental loop
    for ei=1:nel
        eid = elemental_order[ei]; # The index of this element
        first_ind = (eid-1)*dofs_per_node + 1; # First global index for this element
        last_ind = eid*dofs_per_node; # last global index

        # Zero the result vectors for this element
        sourcevec[first_ind:last_ind] .= 0;
        fluxvec[first_ind:last_ind] .= 0;

        ##### Source integrated over the cell #####
        # Compute RHS volume integral
        if !(source_rhs === nothing)
            sourceargs = (var, eid, 0, fv_grid, fv_geo_factors, fv_info, refel, t, dt);
            source = source_rhs.func(sourceargs);
            # Add to global source vector
            sourcevec[first_ind:last_ind] = source;
        end

        # Compute LHS volume integral = (block)diagonal components
        if !(source_lhs === nothing) && !is_explicit
            sourceargs = (var, eid, 0, fv_grid, fv_geo_factors, fv_info, refel, t, dt);
            source = source_lhs.func(sourceargs);
            # Add to global matrix
            for di = 1:dofs_per_node
                first = (eid-1)*dofs_squared + (di-1)*dofs_per_node + 1;
                last = first + dofs_per_node - 1;
                lhsmatV[first:last] = source[((di-1)*dofs_per_node + 1):(di*dofs_per_node)];
                # lhsmatI[first:last] .= first_ind + di - 1; # previously set
                # lhsmatJ[first:last] = first_ind:last_ind;
            end
        end

        ##### Flux integrated over the faces #####
        # Loop over this element's faces.
        for i=1:refel.Nfaces
            fid = fv_grid.element2face[i, eid];
            # Only one element on either side is available here. For more use parent/child version.
            (leftel, rightel) = fv_grid.face2element[:,fid];
            if rightel == 0
                neighborhood = [[leftel],[]];
            else
                neighborhood = [[leftel],[rightel]];
            end
            do_face_here = (face_done[fid] == 0);

            # RHS surface integral
            if !(flux_rhs === nothing)
                if do_face_here
                    face_done[fid] = 1; # Possible race condition, but in the worst case it will be computed twice.

                    fluxargs = (var, eid, fid, neighborhood, fv_grid, fv_geo_factors, fv_info, refel, t, dt);
                    flux = flux_rhs.func(fluxargs) .* fv_geo_factors.area[fid];
                    # Add to global flux vector for faces
                    facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] = flux;
                    # Combine all flux for this element
                    fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] += flux ./ fv_geo_factors.volume[eid];

                else
                    # This flux has either been computed or is being computed by another thread.
                    # The state will need to be known before paralellizing, but for now assume it's complete.
                    fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] -= facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] ./ fv_geo_factors.volume[eid];
                end
            end

            # LHS surface integral
            if !(flux_lhs === nothing) && !is_explicit
                # These update the (block)diagonal for eid and one (block of)off diagonal for neighborID
                if do_face_here
                    face_done[fid] = 1; # Possible race condition.
                    neighborID = (rightel==eid ? leftel : rightel);
                    eid_flux_index = (leftel==eid ? 1 : 2);
                    neighbor_flux_index = (leftel==eid ? 2 : 1);
                    nfirst_ind = (neighborID-1)*dofs_per_node + 1;
                    nlast_ind = neighborID*dofs_per_node;

                    fluxargs = (var, eid, fid, neighborhood, fv_grid, fv_geo_factors, fv_info, refel, t, dt);
                    flux = flux_lhs.func(fluxargs) .* fv_geo_factors.area[fid];
                    # insert the matrix elements
                    for di = 1:dofs_per_node
                        # The eid components
                        first = nel * dofs_squared + (fid-1)*dofs_squared*4 + (di-1)*dofs_per_node + 1;
                        last = first + dofs_per_node - 1;
                        lhsmatV[first:last] = flux[((di-1)*dofs_per_node + 1):(di*dofs_per_node), eid_flux_index] ./ fv_geo_factors.volume[eid];
                        # lhsmatI[first:last] .= first_ind + di - 1; # previously set
                        # lhsmatJ[first:last] = first_ind:last_ind;

                        # The neighborID components
                        if neighborID > 0
                            # first = nel * dofs_squared + (fid-1)*dofs_squared*4 + dofs_squared + (di-1)*dofs_per_node + 1;
                            # last = first + dofs_per_node - 1;
                            first += dofs_squared;
                            last += dofs_squared;
                            lhsmatV[first:last] = flux[((di-1)*dofs_per_node + 1):(di*dofs_per_node), neighbor_flux_index] ./ fv_geo_factors.volume[eid];
                            # lhsmatI[first:last] .= first_ind + di - 1; # previously set
                            # lhsmatJ[first:last] = nfirst_ind:nlast_ind;
                        end
                    end

                    # While we're here, set the components for the neighborID rows as well.
                    if neighborID > 0
                        for di = 1:dofs_per_node
                            # The eid components
                            first = nel * dofs_squared + (fid-1)*dofs_squared*4 + dofs_squared*2 + (di-1)*dofs_per_node + 1;
                            last = first + dofs_per_node - 1;
                            lhsmatV[first:last] = -flux[((di-1)*dofs_per_node + 1):(di*dofs_per_node), eid_flux_index] ./ fv_geo_factors.volume[neighborID];
                            # lhsmatI[first:last] .= nfirst_ind + di - 1; # previously set
                            # lhsmatJ[first:last] = first_ind:last_ind;

                            # The neighborID components
                            # first = nel * dofs_squared + (fid-1)*dofs_squared*4 + dofs_squared*3 + (di-1)*dofs_per_node + 1;
                            # last = first + dofs_per_node - 1;
                            first += dofs_squared;
                            last += dofs_squared;

                            lhsmatV[first:last] = -flux[((di-1)*dofs_per_node + 1):(di*dofs_per_node), neighbor_flux_index] ./ fv_geo_factors.volume[neighborID];
                            # lhsmatI[first:last] .= nfirst_ind + di - 1; # previously set
                            # lhsmatJ[first:last] = nfirst_ind:nlast_ind;
                        end
                    end

                else
                    # This flux has either been computed or is being computed by another thread.
                    # Everything was set there.
                end
            end

            # Boundary conditions are applied to flux
            fbid = fv_grid.facebid[fid]; # BID of this face
            if fbid > 0
                facex = fv_grid.allnodes[:, fv_grid.face2glb[:,1,fid]];  # face node coordinates

                if typeof(var) <: Array
                    dofind = 0;
                    for vi=1:length(var)
                        for compo=1:length(var[vi].symvar)
                            dofind = dofind + 1;
                            if prob.bc_type[var[vi].index, fbid] == NO_BC
                                # do nothing
                            elseif prob.bc_type[var[vi].index, fbid] == FLUX
                                # compute the value and add it to the flux directly
                                # Qvec = (refel.surf_wg[fv_grid.faceRefelInd[1,fid]] .* fv_geo_factors.face_detJ[fid])' * (refel.surf_Q[fv_grid.faceRefelInd[1,fid]])[:, refel.face2local[fv_grid.faceRefelInd[1,fid]]]
                                # Qvec = Qvec ./ fv_geo_factors.area[fid];
                                # bflux = FV_flux_bc_rhs_only(prob.bc_func[var[vi].index, fbid][compo], facex, Qvec, t, dofind, dofs_per_node) .* fv_geo_factors.area[fid];
                                bflux = FV_flux_bc_rhs_only_simple(prob.bc_func[var[vi].index, fbid][compo], fid, t) .* fv_geo_factors.area[fid];
                                if !is_explicit
                                    bflux = bflux .* dt;
                                end

                                fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ fv_geo_factors.volume[eid];
                                facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                            elseif prob.bc_type[var[vi].index, fbid] == DIRICHLET
                                # Set variable array and handle after the face loop
                                var[vi].values[compo,eid] = FV_evaluate_bc(prob.bc_func[var[vi].index, fbid][compo], eid, fid, t);
                                # If implicit, this needs to be handled before solving
                                if !is_explicit
                                    #TODO
                                end
                            else
                                printerr("Unsupported boundary condition type: "*prob.bc_type[var[vi].index, fbid]);
                            end
                        end
                    end
                else
                    for d=1:dofs_per_node
                        dofind = d;
                        if prob.bc_type[var.index, fbid] == NO_BC
                            # do nothing
                        elseif prob.bc_type[var.index, fbid] == FLUX
                            # compute the value and add it to the flux directly
                            # Qvec = (refel.surf_wg[fv_grid.faceRefelInd[1,fid]] .* fv_geo_factors.face_detJ[fid])' * (refel.surf_Q[fv_grid.faceRefelInd[1,fid]])[:, refel.face2local[fv_grid.faceRefelInd[1,fid]]]
                            # Qvec = Qvec ./ fv_geo_factors.area[fid];
                            # bflux = FV_flux_bc_rhs_only(prob.bc_func[var.index, fbid][d], facex, Qvec, t, dofind, dofs_per_node) .* fv_geo_factors.area[fid];
                            bflux = FV_flux_bc_rhs_only_simple(prob.bc_func[var.index, fbid][d], fid, t) .* fv_geo_factors.area[fid];
                            if !is_explicit
                                bflux = bflux .* dt;
                            end

                            fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ fv_geo_factors.volume[eid];
                            facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                        elseif prob.bc_type[var.index, fbid] == DIRICHLET
                            # Set variable array and handle after the face loop
                            var.values[d,eid] = FV_evaluate_bc(prob.bc_func[var.index, fbid][d], eid, fid, t);
                            # If implicit, this needs to be handled before solving
                            if !is_explicit
                                #TODO
                            end
                        else
                            printerr("Unsupported boundary condition type: "*prob.bc_type[var.index, fbid]);
                        end
                    end
                end
            end# BCs

        end# face loop

    end# element loop

    for i=1:length(sourcevec)
        sourcevec[i] += fluxvec[i];
    end
    return sourcevec;
end

function assemble_using_parent_child(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node=1, dofs_per_loop=1, t=0, dt=0; is_explicit=true)
    debug = false;
    if !is_explicit
        printerr("Higher order FV is not ready for implicit time stepping. Sorry.", fatal=true);
    end

    # Label things that were allocated externally
    sourcevec = allocated_vecs[1];
    fluxvec = allocated_vecs[2];
    facefluxvec = allocated_vecs[3];
    face_done = allocated_vecs[4];

    face_done .= 0;
    fluxvec .= 0; # need to do this here

    # Loop over parent elements
    # Since flux depends only on this and nearest neighbor parent cells, it is safe
    # to partition between parents using only the nearest parents as ghosts.
    Nparents = size(parent_maps.parent2child,2);
    nchildren = size(parent_maps.parent2child,1); # (assumes one element type)
    npfaces = size(fv_grid.element2face,1); # Outer faces of parent. Parents are elements in fv_grid (assumes one element type)
    ncfaces = size(parent_maps.parent2face,1); # All faces of children in a parent (assumes one element type)
    dim = size(fv_grid.allnodes,1);

    for parentid=1:Nparents
        # Only work on owned parents
        if !fv_grid.is_subgrid || fv_grid.element_owner[parent_maps.parent2child[1,parentid]] < 0
            # Get the cells in the local patch
            patch_cells = parent_maps.patches[:, parentid];

            # Split the work into two parts:
            # - source terms that don't depend on neighbors and can be done independently (loop over elements)
            # - flux terms that depend on neighbors (loop over faces)

            # The loop over elements for source terms.
            for childid = 1:nchildren
                eid = parent_maps.parent2child[childid, parentid];
                # Zero the result vectors for this element
                sourcevec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] .= 0;

                # Use source_rhs for RHS volume integral
                if !(source_rhs === nothing)
                    sourceargs = (var, eid, 0, fv_grid, fv_geo_factors, fv_info, fv_refel, t, dt);
                    source = source_rhs.func(sourceargs);
                    # source = source_rhs.func(sourceargs) ./ fv_geo_factors.volume[eid];
                    # Add to global source vector
                    sourcevec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] = source;
                end
            end

            # Loop over this parent's faces for flux terms.
            for pfaceid = 1:ncfaces
                fid = parent_maps.parent2face[pfaceid, parentid];
                els = fv_grid.face2element[:,fid];
                eid = els[1];
                neighbor = els[2];

                if debug && parentid < 3 && t<=dt
                    println("parent: "*string(parentid)*", fid: "*string(fid)*"("*string(pfaceid)*"), eid/neigh: "*string(eid)*"/"*string(neighbor));
                    println("patch: "*string(patch_cells));
                end

                if !(flux_rhs === nothing)
                    if face_done[fid] == 0
                        face_done[fid] = 1; # Possible race condition, but in the worst case it will be computed twice.\

                        (left_cells, right_cells) = get_left_right_cells(patch_cells, pfaceid, parent_maps, dim, fv_info.fluxOrder);

                        # leftness and rightness depends on the normal vector direction
                        # normal points left to right.
                        if !(eid == left_cells[1])
                            #oops, swap them
                            tmp = left_cells;
                            left_cells = right_cells;
                            right_cells = tmp;
                        end

                        # remove zero values
                        left_cells = remove_zero_cells(left_cells);
                        right_cells = remove_zero_cells(right_cells);
                        neighborhood = [left_cells, right_cells];

                        fluxargs = (var, eid, fid, neighborhood, fv_grid, fv_geo_factors, fv_info, fv_refel, t, dt);
                        flux = flux_rhs.func(fluxargs) .* fv_geo_factors.area[fid];

                        # Add to global flux vector for faces
                        facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] = flux;
                        # Contribute to flux for elements on both sides
                        if fv_grid.is_subgrid
                            if fv_grid.element_owner[eid] < 0
                                fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] += flux ./ fv_geo_factors.volume[eid];
                            end
                            if neighbor > 0 && fv_grid.element_owner[neighbor] < 0
                                fluxvec[((neighbor-1)*dofs_per_node + 1):(neighbor*dofs_per_node)] -= flux ./ fv_geo_factors.volume[neighbor];
                            end

                        else
                            fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] += flux ./ fv_geo_factors.volume[eid];
                            if neighbor > 0 # not a boundary, so the other side exists
                                fluxvec[((neighbor-1)*dofs_per_node + 1):(neighbor*dofs_per_node)] -= flux ./ fv_geo_factors.volume[neighbor];
                            end
                        end

                        if debug && parentid < 3 && t<=dt
                            println("cells left: "*string(left_cells)*", right: "*string(right_cells));
                            println("flux: "*string(flux));
                            println("contribute "*string(flux ./ fv_geo_factors.volume[eid])*" to "*string(((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)));
                            if neighbor > 0
                                println("contribute "*string(-flux ./ fv_geo_factors.volume[neighbor])*" to "*string(((neighbor-1)*dofs_per_node + 1):(neighbor*dofs_per_node)));
                            end
                            println("fluxvec "*string(fluxvec[1:8]));
                            println("");
                        end

                    else
                        # This flux has either been computed or is being computed by another thread.
                        # The state will need to be known before paralellizing, but for now assume it's complete.
                        # fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] .-= facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] ./ fv_geo_factors.volume[eid];
                    end
                end

                # Boundary conditions are applied to flux
                fbid = fv_grid.facebid[fid]; # BID of this face
                if fbid > 0
                    if debug && parentid < 3 && t==0
                        println("applying bdry cond "*string(fbid)*" to face "*string(fid));
                    end
                    # facex = fv_grid.allnodes[:, fv_grid.face2glb[:,1,fid]];  # face node coordinates

                    if typeof(var) <: Array
                        dofind = 0;
                        for vi=1:length(var)
                            for compo=1:length(var[vi].symvar)
                                dofind = dofind + 1;
                                if prob.bc_type[var[vi].index, fbid] == NO_BC
                                    # do nothing
                                elseif prob.bc_type[var[vi].index, fbid] == FLUX
                                    # compute the value and add it to the flux directly
                                    # Qvec = (fv_refel.surf_wg[fv_grid.faceRefelInd[1,fid]] .* fv_geo_factors.face_detJ[fid])' * (fv_refel.surf_Q[fv_grid.faceRefelInd[1,fid]])[:, fv_refel.face2local[fv_grid.faceRefelInd[1,fid]]]
                                    # Qvec = Qvec ./ fv_geo_factors.area[fid];
                                    # bflux = FV_flux_bc_rhs_only(prob.bc_func[var[vi].index, fbid][compo], facex, Qvec, t, dofind, dofs_per_node) .* fv_geo_factors.area[fid];
                                    bflux = FV_flux_bc_rhs_only_simple(prob.bc_func[var[vi].index, fbid][compo], fid, t) .* fv_geo_factors.area[fid];

                                    fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ fv_geo_factors.volume[eid];
                                    facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                                elseif prob.bc_type[var[vi].index, fbid] == DIRICHLET
                                    # Set variable array and handle after the face loop
                                    var[vi].values[compo,eid] = FV_evaluate_bc(prob.bc_func[var[vi].index, fbid][compo], eid, fid, t);
                                else
                                    printerr("Unsupported boundary condition type: "*prob.bc_type[var[vi].index, fbid]);
                                end
                            end
                        end
                    else
                        for d=1:dofs_per_node
                            dofind = d;
                            if prob.bc_type[var.index, fbid] == NO_BC
                                # do nothing
                            elseif prob.bc_type[var.index, fbid] == FLUX
                                # compute the value and add it to the flux directly
                                # Qvec = (fv_refel.surf_wg[fv_grid.faceRefelInd[1,fid]] .* fv_geo_factors.face_detJ[fid])' * (fv_refel.surf_Q[fv_grid.faceRefelInd[1,fid]])[:, fv_refel.face2local[fv_grid.faceRefelInd[1,fid]]]
                                # Qvec = Qvec ./ fv_geo_factors.area[fid];
                                # bflux = FV_flux_bc_rhs_only(prob.bc_func[var.index, fbid][d], facex, Qvec, t, dofind, dofs_per_node) .* fv_geo_factors.area[fid];
                                bflux = FV_flux_bc_rhs_only_simple(prob.bc_func[var.index, fbid][d], fid, t) .* fv_geo_factors.area[fid];

                                fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ fv_geo_factors.volume[eid];
                                facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                            elseif prob.bc_type[var.index, fbid] == DIRICHLET
                                # Set variable array and handle after the face loop
                                var.values[d,eid] = FV_evaluate_bc(prob.bc_func[var.index, fbid][d], eid, fid, t);
                            else
                                printerr("Unsupported boundary condition type: "*prob.bc_type[var.index, fbid]);
                            end
                        end
                    end
                end# BCs
            end# face loop
        end # is owned?

    end# element loop
    if debug && t<=dt
        println("fluxvec: "*string(fluxvec[1:20]));
    end
    return sourcevec + fluxvec;
end

# Returns a copy of a with zeros removed.
function remove_zero_cells(a)
    b = similar(a);
    nnz = 0;
    for ai in a
        if !(ai == 0)
            nnz += 1;
            b[nnz] = ai;
        end
    end
    b = b[1:nnz];

    return b;
end

# Returns arrays of left and right cells
function get_left_right_cells(patch, face, maps, dim, order)
    # Provide all of the cells that can be used for this face in a particular order.
    # Make two arrays, one for left one for right, starting from the nearest neighbor.
    left_cells = [];
    right_cells = [];
    if dim == 1
        left_cell_table_1d = [
            [ # children=1
                [2],
                [1, 2]
            ],
            [ # children=2
                [4, 3],
                [1, 4, 3],
                [2, 1, 4, 3]
            ],
            [ # children=3
                [6, 5, 4],
                [1, 6, 5, 4],
                [2, 1, 6, 5, 4],
                [3, 2, 1, 6, 5, 4]
            ],
            [ # children=4
                [8, 7, 6, 5],
                [1, 8, 7, 6, 5],
                [2, 1, 8, 7, 6, 5],
                [3, 2, 1, 8, 7, 6, 5],
                [4, 3, 2, 1, 8, 7, 6, 5],
            ]
        ]
        right_cell_table_1d = [
            [ # children=1
                [1, 3],
                [3]
            ],
            [ # children=2
                [1, 2, 5, 6],
                [2, 5, 6],
                [5, 6]
            ],
            [ # children=3
                [1, 2, 3, 7, 8, 9],
                [2, 3, 7, 8, 9],
                [3, 7, 8, 9],
                [7, 8, 9]
            ],
            [ # children=4
                [1, 2, 3, 4, 9, 10, 11, 12],
                [2, 3, 4, 9, 10, 11, 12],
                [3, 4, 9, 10, 11, 12],
                [4, 9, 10, 11, 12],
                [9, 10, 11, 12]
            ]
        ]
        nchildren = size(maps.parent2child,1);
        left_cells = left_cell_table_1d[nchildren][face];
        right_cells = right_cell_table_1d[nchildren][face];

    elseif dim == 2
        if size(maps.parent2neighbor,1) == 3 # triangles
            # Triangle parents have 9 faces, patches have 16 cells
            # Here Left means toward the center of the central parent
            left_cell_table_triangle = [
                [1, 4, 3, 15, 16, 13, 2],
                [2, 4, 3, 9, 12, 11, 1],
                [2, 4, 1, 7, 8, 5, 3],
                [3, 4, 1, 13, 16, 15, 2],
                [3, 4, 2, 11, 12, 9, 1],
                [1, 4, 2, 5, 8, 7, 3],
                [4, 2, 3, 9, 11, 12],
                [4, 1, 3, 15, 13, 16],
                [4, 1, 2, 5, 7, 8]
            ]
            right_cell_table_triangle = [
                [5, 8, 6, 7],
                [7, 8, 6, 5],
                [9, 12, 10, 11],
                [11, 12, 10, 9],
                [13, 16, 14, 15],
                [15, 16, 14, 13],
                [1, 5, 15, 8, 16],
                [2, 7, 9, 8, 12],
                [3, 11, 13, 12, 16]
            ]
            left_cells = left_cell_table_triangle[face];
            right_cells = right_cell_table_triangle[face];
        else # quads
            # Quad parents have 12 faces, patches have 20 cells
            # Here Left means toward the center of the central parent
            left_cell_table_quad = [
                [1, 4, 2, 16, 15, 3, 13, 14],
                [2, 3, 1, 13, 14, 4, 16, 15],
                [2, 1, 3, 20, 19, 4, 17, 18],
                [3, 4, 2, 17, 18, 1, 20, 19],
                [3, 2, 4, 8, 7, 1, 5, 6],
                [4, 1, 3, 5, 6, 2, 8, 7],
                [4, 3, 1, 12, 11, 2, 9, 10],
                [1, 2, 4, 9, 10, 3, 12, 11],
                [1, 20, 4, 19, 17, 18, 5],
                [2, 8, 1, 7, 5, 6, 9],
                [3, 12, 2, 11, 9, 10, 13],
                [4, 16, 3, 15, 13, 14, 17]
            ]
            right_cell_table_quad = [
                [5, 6, 8, 7],
                [8, 7, 5, 6],
                [9, 10, 12, 11],
                [12, 11, 9, 10],
                [13, 14, 16, 15],
                [16, 15, 13, 14],
                [17, 18, 20, 19],
                [20, 19, 17, 18],
                [2, 9, 3, 10, 12, 11],
                [3, 13, 4, 14, 16, 15],
                [4, 17, 1, 18, 20, 19],
                [1, 5, 2, 6, 8, 7]
            ]
            left_cells = left_cell_table_quad[face];
            right_cells = right_cell_table_quad[face];
        end

    elseif dim == 3
        #TODO
    end

    if length(left_cells) > order
        left_cells = left_cells[1:order];
    end
    if length(right_cells) > order
        right_cells = right_cells[1:order];
    end

    return (patch[left_cells], patch[right_cells]);
end

# Insert the single dof into the greater vector
function insert_linear!(b, bel, glb, dof, Ndofs)
    # group nodal dofs
    for d=1:length(dof)
        ind = glb.*Ndofs .- (Ndofs-dof[d]);
        ind2 = ((d-1)*length(glb)+1):(d*length(glb));

        b[ind] = b[ind] + bel[ind2];
    end
end

# place the values in the variable value arrays
function place_sol_in_vars(var, sol, stepper)
    nel = fv_grid.nel_owned;
    if typeof(var) <: Array
        tmp = 0;
        totalcomponents = 0;
        for vi=1:length(var)
            totalcomponents = totalcomponents + var[vi].total_components;
        end
        for vi=1:length(var)
            components = var[vi].total_components;
            for compi=1:components
                var[vi].values[compi,1:nel] = sol[(compi+tmp):totalcomponents:end];
                tmp = tmp + 1;
            end
        end
    else
        components = var.total_components;
        for compi=1:components
            var.values[compi,1:nel] = sol[compi:components:end];
        end
    end
end

# place the variable values in a vector
function get_var_vals(var, vect=nothing)
    nel = fv_grid.nel_owned;
    if typeof(var) <: Array
        tmp = 0;
        totalcomponents = 0;
        for vi=1:length(var)
            totalcomponents = totalcomponents + var[vi].total_components;
        end
        if vect === nothing
            vect = zeros(totalcomponents * nel);
        end

        for vi=1:length(var)
            components = var[vi].total_components;
            for compi=1:components
                vect[(compi+tmp):totalcomponents:end] = var[vi].values[compi,1:nel];
                tmp = tmp + 1;
            end
        end
    else
        components = var.total_components;
        if vect === nothing
            vect = zeros(components * nel);
        end
        for compi=1:components
            vect[compi:components:end] = var.values[compi,1:nel];
        end
    end

    return vect;
end

# Set the i,j indices in the ai and aj vectors for building the sparse matrix
function set_matrix_indices!(ai, aj, dofs_per_node)
    dofs_squared = dofs_per_node*dofs_per_node;
    nel = fv_grid.nel_owned;
    nfaces = size(fv_grid.face2element, 2);
    face_done = fill(false, nfaces);

    # Elemental loop
    for ei=1:nel
        eid = elemental_order[ei]; # The index of this element
        first_ind = (eid-1)*dofs_per_node + 1; # First global index for this element
        last_ind = eid*dofs_per_node; # last global index

        # Diagonal blocks for each cell
        for di = 1:dofs_per_node
            first = (eid-1)*dofs_squared + (di-1)*dofs_per_node + 1;
            last = first + dofs_per_node - 1;
            ai[first:last] .= first_ind + di - 1;
            aj[first:last] = first_ind:last_ind;
        end

        # Diagonal and off diagonal blocks for each face
        # Loop over this element's faces.
        for i=1:refel.Nfaces
            fid = fv_grid.element2face[i, eid];

            if !face_done[fid]
                face_done[fid] = true;
                (leftel, rightel) = fv_grid.face2element[:,fid];
                neighborID = (rightel==eid ? leftel : rightel);
                nfirst_ind = (neighborID-1)*dofs_per_node + 1;
                nlast_ind = neighborID*dofs_per_node;

                for di = 1:dofs_per_node
                    # The eid components
                    first = nel * dofs_squared + (fid-1)*dofs_squared*4 + (di-1)*dofs_per_node + 1;
                    last = first + dofs_per_node - 1;
                    ai[first:last] .= first_ind + di - 1;
                    aj[first:last] = first_ind:last_ind;

                    # The neighborID components
                    if neighborID > 0
                        first += dofs_squared;
                        last += dofs_squared;
                        ai[first:last] .= first_ind + di - 1;
                        aj[first:last] = nfirst_ind:nlast_ind;
                    end
                end

                # While we're here, set the components for the neighborID rows as well.
                if neighborID > 0
                    for di = 1:dofs_per_node
                        # The eid components
                        first = nel * dofs_squared + (fid-1)*dofs_squared*4 + dofs_squared*2 + (di-1)*dofs_per_node + 1;
                        last = first + dofs_per_node - 1;
                        ai[first:last] .= nfirst_ind + di - 1;
                        aj[first:last] = first_ind:last_ind;

                        # The neighborID components
                        first += dofs_squared;
                        last += dofs_squared;
                        ai[first:last] .= nfirst_ind + di - 1;
                        aj[first:last] = nfirst_ind:nlast_ind;
                    end
                end

            else
                # already done
            end
        end# face loop
    end# element loop

    # Some indices may be zero due to boundaries. set those all to [1,1] (not the most efficient, but it's consistent)
    # Since the av part will be zero, this won't change the matrix.
    for i=1:length(ai)
        if ai[i]<1 || aj[i]<1
            ai[i] = 1;
            aj[i] = 1;
        end
    end
end

end #module


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

@testset "Time arithmetic" begin
    a = Dates.Time(23, 59, 59)
    b = Dates.Time(11, 59, 59)
    @test Dates.CompoundPeriod(a - b) == Dates.Hour(12)
end
@testset "Wrapping arithmetic for Months" begin
    # This ends up being trickier than expected because
    # the user might do 2014-01-01 + Month(-14)
    # monthwrap figures out the resulting month
    # when adding/subtracting months from a date
    @test Dates.monthwrap(1, -14) == 11
    @test Dates.monthwrap(1, -13) == 12
    @test Dates.monthwrap(1, -12) == 1
    @test Dates.monthwrap(1, -11) == 2
    @test Dates.monthwrap(1, -10) == 3
    @test Dates.monthwrap(1, -9) == 4
    @test Dates.monthwrap(1, -8) == 5
    @test Dates.monthwrap(1, -7) == 6
    @test Dates.monthwrap(1, -6) == 7
    @test Dates.monthwrap(1, -5) == 8
    @test Dates.monthwrap(1, -4) == 9
    @test Dates.monthwrap(1, -3) == 10
    @test Dates.monthwrap(1, -2) == 11
    @test Dates.monthwrap(1, -1) == 12
    @test Dates.monthwrap(1, 0) == 1
    @test Dates.monthwrap(1, 1) == 2
    @test Dates.monthwrap(1, 2) == 3
    @test Dates.monthwrap(1, 3) == 4
    @test Dates.monthwrap(1, 4) == 5
    @test Dates.monthwrap(1, 5) == 6
    @test Dates.monthwrap(1, 6) == 7
    @test Dates.monthwrap(1, 7) == 8
    @test Dates.monthwrap(1, 8) == 9
    @test Dates.monthwrap(1, 9) == 10
    @test Dates.monthwrap(1, 10) == 11
    @test Dates.monthwrap(1, 11) == 12
    @test Dates.monthwrap(1, 12) == 1
    @test Dates.monthwrap(1, 13) == 2
    @test Dates.monthwrap(1, 24) == 1
    @test Dates.monthwrap(12, -14) == 10
    @test Dates.monthwrap(12, -13) == 11
    @test Dates.monthwrap(12, -12) == 12
    @test Dates.monthwrap(12, -11) == 1
    @test Dates.monthwrap(12, -2) == 10
    @test Dates.monthwrap(12, -1) == 11
    @test Dates.monthwrap(12, 0) == 12
    @test Dates.monthwrap(12, 1) == 1
    @test Dates.monthwrap(12, 2) == 2
    @test Dates.monthwrap(12, 11) == 11
    @test Dates.monthwrap(12, 12) == 12
    @test Dates.monthwrap(12, 13) == 1
end
@testset "Wrapping arithmetic for years" begin
    # yearwrap figures out the resulting year
    # when adding/subtracting months from a date
    @test Dates.yearwrap(2000, 1, -3600) == 1700
    @test Dates.yearwrap(2000, 1, -37) == 1996
    @test Dates.yearwrap(2000, 1, -36) == 1997
    @test Dates.yearwrap(2000, 1, -35) == 1997
    @test Dates.yearwrap(2000, 1, -25) == 1997
    @test Dates.yearwrap(2000, 1, -24) == 1998
    @test Dates.yearwrap(2000, 1, -23) == 1998
    @test Dates.yearwrap(2000, 1, -14) == 1998
    @test Dates.yearwrap(2000, 1, -13) == 1998
    @test Dates.yearwrap(2000, 1, -12) == 1999
    @test Dates.yearwrap(2000, 1, -11) == 1999
    @test Dates.yearwrap(2000, 1, -2) == 1999
    @test Dates.yearwrap(2000, 1, -1) == 1999
    @test Dates.yearwrap(2000, 1, 0) == 2000
    @test Dates.yearwrap(2000, 1, 1) == 2000
    @test Dates.yearwrap(2000, 1, 11) == 2000
    @test Dates.yearwrap(2000, 1, 12) == 2001
    @test Dates.yearwrap(2000, 1, 13) == 2001
    @test Dates.yearwrap(2000, 1, 23) == 2001
    @test Dates.yearwrap(2000, 1, 24) == 2002
    @test Dates.yearwrap(2000, 1, 25) == 2002
    @test Dates.yearwrap(2000, 1, 36) == 2003
    @test Dates.yearwrap(2000, 1, 3600) == 2300
    @test Dates.yearwrap(2000, 2, -2) == 1999
    @test Dates.yearwrap(2000, 3, 10) == 2001
    @test Dates.yearwrap(2000, 4, -4) == 1999
    @test Dates.yearwrap(2000, 5, 8) == 2001
    @test Dates.yearwrap(2000, 6, -18) == 1998
    @test Dates.yearwrap(2000, 6, -6) == 1999
    @test Dates.yearwrap(2000, 6, 6) == 2000
    @test Dates.yearwrap(2000, 6, 7) == 2001
    @test Dates.yearwrap(2000, 6, 19) == 2002
    @test Dates.yearwrap(2000, 12, -3600) == 1700
    @test Dates.yearwrap(2000, 12, -36) == 1997
    @test Dates.yearwrap(2000, 12, -35) == 1998
    @test Dates.yearwrap(2000, 12, -24) == 1998
    @test Dates.yearwrap(2000, 12, -23) == 1999
    @test Dates.yearwrap(2000, 12, -14) == 1999
    @test Dates.yearwrap(2000, 12, -13) == 1999
    @test Dates.yearwrap(2000, 12, -12) == 1999
    @test Dates.yearwrap(2000, 12, -11) == 2000
    @test Dates.yearwrap(2000, 12, -2) == 2000
    @test Dates.yearwrap(2000, 12, -1) == 2000
    @test Dates.yearwrap(2000, 12, 0) == 2000
    @test Dates.yearwrap(2000, 12, 1) == 2001
    @test Dates.yearwrap(2000, 12, 11) == 2001
    @test Dates.yearwrap(2000, 12, 12) == 2001
    @test Dates.yearwrap(2000, 12, 13) == 2002
    @test Dates.yearwrap(2000, 12, 24) == 2002
    @test Dates.yearwrap(2000, 12, 25) == 2003
    @test Dates.yearwrap(2000, 12, 36) == 2003
    @test Dates.yearwrap(2000, 12, 37) == 2004
    @test Dates.yearwrap(2000, 12, 3600) == 2300
end
@testset "DateTime arithmetic" begin
    # Dates.DateTime arithmetic
    a = Dates.DateTime(2013, 1, 1, 0, 0, 0, 1)
    b = Dates.DateTime(2013, 1, 1, 0, 0, 0, 0)
    @test a - b == Dates.Millisecond(1)
    @test Dates.DateTime(2013, 1, 2) - b == Dates.Millisecond(86400000)

    @testset "Dates.DateTime-Year arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Year(1) == Dates.DateTime(2000, 12, 27)
        @test dt + Dates.Year(100) == Dates.DateTime(2099, 12, 27)
        @test dt + Dates.Year(1000) == Dates.DateTime(2999, 12, 27)
        @test dt - Dates.Year(1) == Dates.DateTime(1998, 12, 27)
        @test dt - Dates.Year(100) == Dates.DateTime(1899, 12, 27)
        @test dt - Dates.Year(1000) == Dates.DateTime(999, 12, 27)
        dt = Dates.DateTime(2000, 2, 29)
        @test dt + Dates.Year(1) == Dates.DateTime(2001, 2, 28)
        @test dt - Dates.Year(1) == Dates.DateTime(1999, 2, 28)
        @test dt + Dates.Year(4) == Dates.DateTime(2004, 2, 29)
        @test dt - Dates.Year(4) == Dates.DateTime(1996, 2, 29)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Year(1) == Dates.DateTime(1973, 6, 30, 23, 59, 59)
        @test dt - Dates.Year(1) == Dates.DateTime(1971, 6, 30, 23, 59, 59)
        @test dt + Dates.Year(-1) == Dates.DateTime(1971, 6, 30, 23, 59, 59)
        @test dt - Dates.Year(-1) == Dates.DateTime(1973, 6, 30, 23, 59, 59)

        dt = Dates.DateTime(2000, 1, 1, 12, 30, 45, 500)
        dt2 = dt + Dates.Year(1)
        @test Dates.year(dt2) == 2001
        @test Dates.month(dt2) == 1
        @test Dates.day(dt2) == 1
        @test Dates.hour(dt2) == 12
        @test Dates.minute(dt2) == 30
        @test Dates.second(dt2) == 45
        @test Dates.millisecond(dt2) == 500
    end
    @testset "DateTime-Month arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Month(1) == Dates.DateTime(2000, 1, 27)
        @test dt + Dates.Month(-1) == Dates.DateTime(1999, 11, 27)
        @test dt + Dates.Month(-11) == Dates.DateTime(1999, 1, 27)
        @test dt + Dates.Month(11) == Dates.DateTime(2000, 11, 27)
        @test dt + Dates.Month(-12) == Dates.DateTime(1998, 12, 27)
        @test dt + Dates.Month(12) == Dates.DateTime(2000, 12, 27)
        @test dt + Dates.Month(13) == Dates.DateTime(2001, 1, 27)
        @test dt + Dates.Month(100) == Dates.DateTime(2008, 4, 27)
        @test dt + Dates.Month(1000) == Dates.DateTime(2083, 4, 27)
        @test dt - Dates.Month(1) == Dates.DateTime(1999, 11, 27)
        @test dt - Dates.Month(-1) == Dates.DateTime(2000, 1, 27)
        @test dt - Dates.Month(100) == Dates.DateTime(1991, 8, 27)
        @test dt - Dates.Month(1000) == Dates.DateTime(1916, 8, 27)
        dt = Dates.DateTime(2000, 2, 29)
        @test dt + Dates.Month(1) == Dates.DateTime(2000, 3, 29)
        @test dt - Dates.Month(1) == Dates.DateTime(2000, 1, 29)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Month(1) == Dates.DateTime(1972, 7, 30, 23, 59, 59)
        @test dt - Dates.Month(1) == Dates.DateTime(1972, 5, 30, 23, 59, 59)
        @test dt + Dates.Month(-1) == Dates.DateTime(1972, 5, 30, 23, 59, 59)
    end
    @testset "DateTime-Week arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Week(1) == Dates.DateTime(2000, 1, 3)
        @test dt + Dates.Week(100) == Dates.DateTime(2001, 11, 26)
        @test dt + Dates.Week(1000) == Dates.DateTime(2019, 2, 25)
        @test dt - Dates.Week(1) == Dates.DateTime(1999, 12, 20)
        @test dt - Dates.Week(100) == Dates.DateTime(1998, 1, 26)
        @test dt - Dates.Week(1000) == Dates.DateTime(1980, 10, 27)
        dt = Dates.DateTime(2000, 2, 29)
        @test dt + Dates.Week(1) == Dates.DateTime(2000, 3, 7)
        @test dt - Dates.Week(1) == Dates.DateTime(2000, 2, 22)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Week(1) == Dates.DateTime(1972, 7, 7, 23, 59, 59)
        @test dt - Dates.Week(1) == Dates.DateTime(1972, 6, 23, 23, 59, 59)
        @test dt + Dates.Week(-1) == Dates.DateTime(1972, 6, 23, 23, 59, 59)
    end
    @testset "DateTime-Day arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Day(1) == Dates.DateTime(1999, 12, 28)
        @test dt + Dates.Day(100) == Dates.DateTime(2000, 4, 5)
        @test dt + Dates.Day(1000) == Dates.DateTime(2002, 9, 22)
        @test dt - Dates.Day(1) == Dates.DateTime(1999, 12, 26)
        @test dt - Dates.Day(100) == Dates.DateTime(1999, 9, 18)
        @test dt - Dates.Day(1000) == Dates.DateTime(1997, 4, 1)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Day(1) == Dates.DateTime(1972, 7, 1, 23, 59, 59)
        @test dt - Dates.Day(1) == Dates.DateTime(1972, 6, 29, 23, 59, 59)
        @test dt + Dates.Day(-1) == Dates.DateTime(1972, 6, 29, 23, 59, 59)
    end
    @testset "DateTime-Hour arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Hour(1) == Dates.DateTime(1999, 12, 27, 1)
        @test dt + Dates.Hour(100) == Dates.DateTime(1999, 12, 31, 4)
        @test dt + Dates.Hour(1000) == Dates.DateTime(2000, 2, 6, 16)
        @test dt - Dates.Hour(1) == Dates.DateTime(1999, 12, 26, 23)
        @test dt - Dates.Hour(100) == Dates.DateTime(1999, 12, 22, 20)
        @test dt - Dates.Hour(1000) == Dates.DateTime(1999, 11, 15, 8)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Hour(1) == Dates.DateTime(1972, 7, 1, 0, 59, 59)
        @test dt - Dates.Hour(1) == Dates.DateTime(1972, 6, 30, 22, 59, 59)
        @test dt + Dates.Hour(-1) == Dates.DateTime(1972, 6, 30, 22, 59, 59)
    end
    @testset "DateTime-Minute arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Minute(1) == Dates.DateTime(1999, 12, 27, 0, 1)
        @test dt + Dates.Minute(100) == Dates.DateTime(1999, 12, 27, 1, 40)
        @test dt + Dates.Minute(1000) == Dates.DateTime(1999, 12, 27, 16, 40)
        @test dt - Dates.Minute(1) == Dates.DateTime(1999, 12, 26, 23, 59)
        @test dt - Dates.Minute(100) == Dates.DateTime(1999, 12, 26, 22, 20)
        @test dt - Dates.Minute(1000) == Dates.DateTime(1999, 12, 26, 7, 20)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Minute(1) == Dates.DateTime(1972, 7, 1, 0, 0, 59)
        @test dt - Dates.Minute(1) == Dates.DateTime(1972, 6, 30, 23, 58, 59)
        @test dt + Dates.Minute(-1) == Dates.DateTime(1972, 6, 30, 23, 58, 59)
    end
    @testset "DateTime-Second arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Second(1) == Dates.DateTime(1999, 12, 27, 0, 0, 1)
        @test dt + Dates.Second(100) == Dates.DateTime(1999, 12, 27, 0, 1, 40)
        @test dt + Dates.Second(1000) == Dates.DateTime(1999, 12, 27, 0, 16, 40)
        @test dt - Dates.Second(1) == Dates.DateTime(1999, 12, 26, 23, 59, 59)
        @test dt - Dates.Second(100) == Dates.DateTime(1999, 12, 26, 23, 58, 20)
        @test dt - Dates.Second(1000) == Dates.DateTime(1999, 12, 26, 23, 43, 20)
    end
    @testset "DateTime-Millisecond arithmetic" begin
        dt = Dates.DateTime(1999, 12, 27)
        @test dt + Dates.Millisecond(1) == Dates.DateTime(1999, 12, 27, 0, 0, 0, 1)
        @test dt + Dates.Millisecond(100) == Dates.DateTime(1999, 12, 27, 0, 0, 0, 100)
        @test dt + Dates.Millisecond(1000) == Dates.DateTime(1999, 12, 27, 0, 0, 1)
        @test dt - Dates.Millisecond(1) == Dates.DateTime(1999, 12, 26, 23, 59, 59, 999)
        @test dt - Dates.Millisecond(100) == Dates.DateTime(1999, 12, 26, 23, 59, 59, 900)
        @test dt - Dates.Millisecond(1000) == Dates.DateTime(1999, 12, 26, 23, 59, 59)
        dt = Dates.DateTime(1972, 6, 30, 23, 59, 59)
        @test dt + Dates.Millisecond(1) == Dates.DateTime(1972, 6, 30, 23, 59, 59, 1)
        @test dt - Dates.Millisecond(1) == Dates.DateTime(1972, 6, 30, 23, 59, 58, 999)
        @test dt + Dates.Millisecond(-1) == Dates.DateTime(1972, 6, 30, 23, 59, 58, 999)
    end
end
@testset "Date arithmetic" begin
    @testset "Date-Year arithmetic" begin
        dt = Dates.Date(1999, 12, 27)
        @test dt + Dates.Year(1) == Dates.Date(2000, 12, 27)
        @test dt + Dates.Year(100) == Dates.Date(2099, 12, 27)
        @test dt + Dates.Year(1000) == Dates.Date(2999, 12, 27)
        @test dt - Dates.Year(1) == Dates.Date(1998, 12, 27)
        @test dt - Dates.Year(100) == Dates.Date(1899, 12, 27)
        @test dt - Dates.Year(1000) == Dates.Date(999, 12, 27)
        dt = Dates.Date(2000, 2, 29)
        @test dt + Dates.Year(1) == Dates.Date(2001, 2, 28)
        @test dt - Dates.Year(1) == Dates.Date(1999, 2, 28)
        @test dt + Dates.Year(4) == Dates.Date(2004, 2, 29)
        @test dt - Dates.Year(4) == Dates.Date(1996, 2, 29)
    end
    @testset "Date-Month arithmetic" begin
        dt = Dates.Date(1999, 12, 27)
        @test dt + Dates.Month(1) == Dates.Date(2000, 1, 27)
        @test dt + Dates.Month(100) == Dates.Date(2008, 4, 27)
        @test dt + Dates.Month(1000) == Dates.Date(2083, 4, 27)
        @test dt - Dates.Month(1) == Dates.Date(1999, 11, 27)
        @test dt - Dates.Month(100) == Dates.Date(1991, 8, 27)
        @test dt - Dates.Month(1000) == Dates.Date(1916, 8, 27)
        dt = Dates.Date(2000, 2, 29)
        @test dt + Dates.Month(1) == Dates.Date(2000, 3, 29)
        @test dt - Dates.Month(1) == Dates.Date(2000, 1, 29)
    end
    @testset "Date-Week arithmetic" begin
        dt = Dates.Date(1999, 12, 27)
        @test dt + Dates.Week(1) == Dates.Date(2000, 1, 3)
        @test dt + Dates.Week(100) == Dates.Date(2001, 11, 26)
        @test dt + Dates.Week(1000) == Dates.Date(2019, 2, 25)
        @test dt - Dates.Week(1) == Dates.Date(1999, 12, 20)
        @test dt - Dates.Week(100) == Dates.Date(1998, 1, 26)
        @test dt - Dates.Week(1000) == Dates.Date(1980, 10, 27)
        dt = Dates.Date(2000, 2, 29)
        @test dt + Dates.Week(1) == Dates.Date(2000, 3, 7)
        @test dt - Dates.Week(1) == Dates.Date(2000, 2, 22)
    end
    @testset "Date-Day arithmetic" begin
        dt = Dates.Date(1999, 12, 27)
        @test dt + Dates.Day(1) == Dates.Date(1999, 12, 28)
        @test dt + Dates.Day(100) == Dates.Date(2000, 4, 5)
        @test dt + Dates.Day(1000) == Dates.Date(2002, 9, 22)
        @test dt - Dates.Day(1) == Dates.Date(1999, 12, 26)
        @test dt - Dates.Day(100) == Dates.Date(1999, 9, 18)
        @test dt - Dates.Day(1000) == Dates.Date(1997, 4, 1)
    end
end
@testset "Time-TimePeriod arithmetic" begin
    t = Dates.Time(0)
    @test t + Dates.Hour(1) == Dates.Time(1)
    @test t - Dates.Hour(1) == Dates.Time(23)
    @test t - Dates.Nanosecond(1) == Dates.Time(23, 59, 59, 999, 999, 999)
    @test t + Dates.Nanosecond(-1) == Dates.Time(23, 59, 59, 999, 999, 999)
    @test t + Dates.Hour(24) == t
    @test t + Dates.Nanosecond(86400000000000) == t
    @test t - Dates.Nanosecond(86400000000000) == t
    @test t + Dates.Minute(1) == Dates.Time(0, 1)
    @test t + Dates.Second(1) == Dates.Time(0, 0, 1)
    @test t + Dates.Millisecond(1) == Dates.Time(0, 0, 0, 1)
    @test t + Dates.Microsecond(1) == Dates.Time(0, 0, 0, 0, 1)
    @test_throws MethodError t + Dates.Day(1)
end
@testset "Month arithmetic and non-associativity" begin
    # Month arithmetic minimizes "edit distance", or number of changes
    # needed to get a correct answer
    # This approach results in a few cases of non-associativity
    a = Dates.Date(2012, 1, 29)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 1, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 2, 29)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 3, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 4, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 5, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 6, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 8, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 9, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 10, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
    a = Dates.Date(2012, 11, 30)
    @test (a + Dates.Day(1)) + Dates.Month(1) != (a + Dates.Month(1)) + Dates.Day(1)
end
@testset "Vectorized arithmetic" begin
    # Vectorized Time arithmetic
    a = Dates.Time(1, 1, 1)
    dr = [a, a, a, a, a, a, a, a, a, a]
    b = a + Dates.Hour(1)
    @test dr .+ Dates.Hour(1) == repmat([b], 10)
    b = a + Dates.Second(1)
    @test dr .+ Dates.Second(1) == repmat([b], 10)
    b = a + Dates.Millisecond(1)
    @test dr .+ Dates.Millisecond(1) == repmat([b], 10)
    b = a + Dates.Microsecond(1)
    @test dr .+ Dates.Microsecond(1) == repmat([b], 10)
    b = a + Dates.Nanosecond(1)
    @test dr .+ Dates.Nanosecond(1) == repmat([b], 10)

    b = a - Dates.Hour(1)
    @test dr .- Dates.Hour(1) == repmat([b], 10)
    b = a - Dates.Second(1)
    @test dr .- Dates.Second(1) == repmat([b], 10)
    b = a - Dates.Millisecond(1)
    @test dr .- Dates.Millisecond(1) == repmat([b], 10)
    b = a - Dates.Microsecond(1)
    @test dr .- Dates.Microsecond(1) == repmat([b], 10)
    b = a - Dates.Nanosecond(1)
    @test dr .- Dates.Nanosecond(1) == repmat([b], 10)

    # Vectorized arithmetic
    a = Dates.Date(2014, 1, 1)
    dr = [a, a, a, a, a, a, a, a, a, a]
    b = a + Dates.Year(1)
    @test dr .+ Dates.Year(1) == repmat([b], 10)
    b = a + Dates.Month(1)
    @test dr .+ Dates.Month(1) == repmat([b], 10)
    b = a + Dates.Day(1)
    @test dr .+ Dates.Day(1) == repmat([b], 10)
    b = a - Dates.Year(1)
    @test dr .- Dates.Year(1) == repmat([b], 10)
    b = a - Dates.Month(1)
    @test dr .- Dates.Month(1) == repmat([b], 10)
    b = a - Dates.Day(1)
    @test dr .- Dates.Day(1) == repmat([b], 10)

    # Vectorized arithmetic
    b = a + Dates.Year(1)
    @test dr .+ Dates.Year(1) == repmat([b], 10)
    b = a + Dates.Month(1)
    @test dr .+ Dates.Month(1) == repmat([b], 10)
    b = a + Dates.Day(1)
    @test dr .+ Dates.Day(1) == repmat([b], 10)
    b = a - Dates.Year(1)
    @test dr .- Dates.Year(1) == repmat([b], 10)
    b = a - Dates.Month(1)
    @test dr .- Dates.Month(1) == repmat([b], 10)
    b = a - Dates.Day(1)
    @test dr .- Dates.Day(1) == repmat([b], 10)
    t1 = [Dates.Date(2009, 1, 1) Dates.Date(2009, 1, 2) Dates.Date(2009, 1, 3); Dates.Date(2009, 2, 1) Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3)]
    t2 = [Dates.Date(2009, 1, 2) Dates.Date(2009, 2, 2) Dates.Date(2010, 1, 3); Dates.Date(2010, 2, 1) Dates.Date(2009, 3, 2) Dates.Date(2009, 2, 4)]
    t3 = [Dates.DateTime(2009, 1, 1), Dates.DateTime(2009, 1, 2), Dates.DateTime(2009, 1, 3)]
    t4 = [Dates.DateTime(2009, 1, 1, 0, 0, 1), Dates.DateTime(2009, 1, 2, 0, 1), Dates.DateTime(2009, 1, 3, 1)]
    t5 = [Dates.Time(0, 0, 0) Dates.Time(0, 0, 1) Dates.Time(0, 0, 2); Dates.Time(0, 1, 0) Dates.Time(0, 2, 0) Dates.Time(0, 3, 0)]

    @testset "TimeType, Array{TimeType}" begin
        @test Dates.Date(2010, 1, 1) .- t1 == [Dates.Day(365) Dates.Day(364) Dates.Day(363); Dates.Day(334) Dates.Day(333) Dates.Day(332)]
        @test t1 .- Dates.Date(2010, 1, 1) == [Dates.Day(-365) Dates.Day(-364) Dates.Day(-363); Dates.Day(-334) Dates.Day(-333) Dates.Day(-332)]
        @test Dates.DateTime(2009, 1, 1) .- t3 == [Dates.Millisecond(0), Dates.Millisecond(-86400000), Dates.Millisecond(-172800000)]
        @test t3 .- Dates.DateTime(2009, 1, 1) == [Dates.Millisecond(0), Dates.Millisecond(86400000), Dates.Millisecond(172800000)]
        @test Dates.Date(2010, 1, 1) - t1 == [Dates.Day(365) Dates.Day(364) Dates.Day(363); Dates.Day(334) Dates.Day(333) Dates.Day(332)]
        @test t1 - Dates.Date(2010, 1, 1) == [Dates.Day(-365) Dates.Day(-364) Dates.Day(-363); Dates.Day(-334) Dates.Day(-333) Dates.Day(-332)]
        @test Dates.DateTime(2009, 1, 1) - t3 == [Dates.Millisecond(0), Dates.Millisecond(-86400000), Dates.Millisecond(-172800000)]
        @test t3 - Dates.DateTime(2009, 1, 1) == [Dates.Millisecond(0), Dates.Millisecond(86400000), Dates.Millisecond(172800000)]
        @test Dates.Time(2) .- t5 == [Dates.Nanosecond(7200000000000) Dates.Nanosecond(7199000000000) Dates.Nanosecond(7198000000000);
                                      Dates.Nanosecond(7140000000000) Dates.Nanosecond(7080000000000) Dates.Nanosecond(7020000000000)]
    end
    @testset "GeneralPeriod, Array{TimeType}" begin
        @test Dates.Day(1) .+ t1 == [Dates.Date(2009, 1, 2) Dates.Date(2009, 1, 3) Dates.Date(2009, 1, 4); Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4)]
        @test Dates.Hour(1) .+ t3 == [Dates.DateTime(2009, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1), Dates.DateTime(2009, 1, 3, 1)]
        @test t1 .+ Dates.Day(1) == [Dates.Date(2009, 1, 2) Dates.Date(2009, 1, 3) Dates.Date(2009, 1, 4); Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4)]
        @test t3 .+ Dates.Hour(1) == [Dates.DateTime(2009, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1), Dates.DateTime(2009, 1, 3, 1)]
        @test Dates.Day(1) + t1 == [Dates.Date(2009, 1, 2) Dates.Date(2009, 1, 3) Dates.Date(2009, 1, 4); Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4)]
        @test Dates.Hour(1) + t3 == [Dates.DateTime(2009, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1), Dates.DateTime(2009, 1, 3, 1)]
        @test t1 + Dates.Day(1) == [Dates.Date(2009, 1, 2) Dates.Date(2009, 1, 3) Dates.Date(2009, 1, 4); Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4)]
        @test t3 + Dates.Hour(1) == [Dates.DateTime(2009, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1), Dates.DateTime(2009, 1, 3, 1)]
        @test t5 + Dates.Hour(1) == [Dates.Time(1, 0, 0) Dates.Time(1, 0, 1) Dates.Time(1, 0, 2); Dates.Time(1, 1, 0) Dates.Time(1, 2, 0) Dates.Time(1, 3, 0)]

        @test (Dates.Month(1) + Dates.Day(1)) .+ t1 == [Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4); Dates.Date(2009, 3, 2) Dates.Date(2009, 3, 3) Dates.Date(2009, 3, 4)]
        @test (Dates.Hour(1) + Dates.Minute(1)) .+ t3 == [Dates.DateTime(2009, 1, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1, 1), Dates.DateTime(2009, 1, 3, 1, 1)]
        @test t1 .+ (Dates.Month(1) + Dates.Day(1)) == [Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4); Dates.Date(2009, 3, 2) Dates.Date(2009, 3, 3) Dates.Date(2009, 3, 4)]
        @test t3 .+ (Dates.Hour(1) + Dates.Minute(1)) == [Dates.DateTime(2009, 1, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1, 1), Dates.DateTime(2009, 1, 3, 1, 1)]
        @test (Dates.Month(1) + Dates.Day(1)) + t1 == [Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4); Dates.Date(2009, 3, 2) Dates.Date(2009, 3, 3) Dates.Date(2009, 3, 4)]
        @test (Dates.Hour(1) + Dates.Minute(1)) + t3 == [Dates.DateTime(2009, 1, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1, 1), Dates.DateTime(2009, 1, 3, 1, 1)]
        @test t1 + (Dates.Month(1) + Dates.Day(1)) == [Dates.Date(2009, 2, 2) Dates.Date(2009, 2, 3) Dates.Date(2009, 2, 4); Dates.Date(2009, 3, 2) Dates.Date(2009, 3, 3) Dates.Date(2009, 3, 4)]
        @test t3 + (Dates.Hour(1) + Dates.Minute(1)) == [Dates.DateTime(2009, 1, 1, 1, 1), Dates.DateTime(2009, 1, 2, 1, 1), Dates.DateTime(2009, 1, 3, 1, 1)]

        @test t1 .- Dates.Day(1) == [Dates.Date(2008, 12, 31) Dates.Date(2009, 1, 1) Dates.Date(2009, 1, 2); Dates.Date(2009, 1, 31) Dates.Date(2009, 2, 1) Dates.Date(2009, 2, 2)]
        @test t3 .- Dates.Hour(1) == [Dates.DateTime(2008, 12, 31, 23), Dates.DateTime(2009, 1, 1, 23), Dates.DateTime(2009, 1, 2, 23)]
        @test t1 - Dates.Day(1) == [Dates.Date(2008, 12, 31) Dates.Date(2009, 1, 1) Dates.Date(2009, 1, 2); Dates.Date(2009, 1, 31) Dates.Date(2009, 2, 1) Dates.Date(2009, 2, 2)]
        @test t3 - Dates.Hour(1) == [Dates.DateTime(2008, 12, 31, 23), Dates.DateTime(2009, 1, 1, 23), Dates.DateTime(2009, 1, 2, 23)]

        @test t1 .- (Dates.Month(1) + Dates.Day(1)) == [Dates.Date(2008, 11, 30) Dates.Date(2008, 12, 1) Dates.Date(2008, 12, 2); Dates.Date(2008, 12, 31) Dates.Date(2009, 1, 1) Dates.Date(2009, 1, 2)]
        @test t3 .- (Dates.Hour(1) + Dates.Minute(1)) == [Dates.DateTime(2008, 12, 31, 22, 59), Dates.DateTime(2009, 1, 1, 22, 59), Dates.DateTime(2009, 1, 2, 22, 59)]
        @test t1 - (Dates.Month(1) + Dates.Day(1)) == [Dates.Date(2008, 11, 30) Dates.Date(2008, 12, 1) Dates.Date(2008, 12, 2); Dates.Date(2008, 12, 31) Dates.Date(2009, 1, 1) Dates.Date(2009, 1, 2)]
        @test t3 - (Dates.Hour(1) + Dates.Minute(1)) == [Dates.DateTime(2008, 12, 31, 22, 59), Dates.DateTime(2009, 1, 1, 22, 59), Dates.DateTime(2009, 1, 2, 22, 59)]
    end
    @testset "Array{TimeType}, Array{TimeType}" begin
        @test t2 - t1 == [Dates.Day(1) Dates.Day(31) Dates.Day(365); Dates.Day(365) Dates.Day(28) Dates.Day(1)]
        @test t4 - t3 == [Dates.Millisecond(1000), Dates.Millisecond(60000), Dates.Millisecond(3600000)]
        @test (Dates.Date(2009, 1, 1):Dates.Week(1):Dates.Date(2009, 1, 21)) - (Dates.Date(2009, 1, 1):Dates.Day(1):Dates.Date(2009, 1, 3)) == [Dates.Day(0), Dates.Day(6), Dates.Day(12)]
        @test (Dates.DateTime(2009, 1, 1, 1, 1, 1):Dates.Second(1):Dates.DateTime(2009, 1, 1, 1, 1, 3)) - (Dates.DateTime(2009, 1, 1, 1, 1):Dates.Second(1):Dates.DateTime(2009, 1, 1, 1, 1, 2)) == [Dates.Second(1), Dates.Second(1), Dates.Second(1)]
    end
    @testset "#20205" begin
        # ensure commutative subtraction methods are not defined
        @test_throws MethodError Dates.Day(1) .- t1
        @test_throws MethodError Dates.Hour(1) .- t3
        @test_throws MethodError Dates.Day(1) - t1
        @test_throws MethodError Dates.Hour(1) - t3
        @test_throws MethodError (Dates.Month(1) + Dates.Day(1)) .- t1
        @test_throws MethodError (Dates.Hour(1) + Dates.Minute(1)) .- t3
        @test_throws MethodError (Dates.Month(1) + Dates.Day(1)) - t1
        @test_throws MethodError (Dates.Hour(1) + Dates.Minute(1)) - t3
    end
end


using AutoLandmarking
using Flux
using CuArrays # to run models on the GPU for massive performance increase
import ImageView

push!(LOAD_PATH, Base.@__DIR__)

# Import data with custom import functions
image_data, imname = load_imgs("C:/Users/immanueldiener/Desktop/Master/scripts/no_enamel", (128,128,128), false)
landmark_data = read_landmarks("C:/Users/immanueldiener/Desktop/Master/scripts/no_enamel", 10, "@1")

# X and Y coordinates are swapped in avizo files, wo I'll swap them back here.
#this step is essential for all the data augmentation to work.
landmark_data = swap_xy(landmark_data)

#some images have a higher value in the fill voxels than in the actual object voxels. This will fix the issue.
# The highest value must always be the filler value!
image_data[:,:,:,[109,130]] = image_data[:,:,:,[109,130]] .- 1
image_data[:,:,:,[109,130]] = image_data[:,:,:,[109,130]] .* -1

image_data2, imname2 = load_imgs("C:/Users/immanueldiener/Desktop/Master/master_data\\additional_teeth_volumes\\", (128,128,128), true)
landmark_data2 = read_landmarks("C:/Users/immanueldiener/Desktop/Master/master_data\\additional_teeth_volumes\\", 10, "@1")

#Swap X and Y again
landmark_data2 = swap_xy(landmark_data2)

#all these images have reversed fill and object values. Tha highest value must always be the filler value!
image_data2 = (image_data2 .- 1) .* -1

names = vcat(imname, imname2)
images = cat(image_data, image_data2, dims=4)
lms = cat(landmark_data, landmark_data2, dims=2)

# align along principal axes
aligned, aligned_lms, retro = align_principal(images, lms, 192)
resized, resized_lms, scales = resize_relevant(aligned, aligned_lms, 144)
ImageView.imshow(depth_map_all_sides(resized))
#define training and testing sets
X_train, X_test, y_train, y_test = regular_train_test_split_3d(resized, resized_lms)

# perform data augmentation. In case of RAM limitations the depth-map generation can be
# done right after every augmentation step, after which the augmented volumes can be deleted.
X_train, y_train = (mirror_vol(X_train, y_train))

X_train = Float32.(X_train)

flip1, lm_flip1 = flip_3D(X_train, y_train)
flip2, lm_flip2 = flip_volume_front(flip1, lm_flip1)
flip3, lm_flip3 = flip_volume_side(flip2, lm_flip2)

rot, lm_rot = rotate_volumes(flip2, lm_flip2, 10)
rot2, lm_rot2 = rotate_volumes(X_train, y_train, 20)
rot3, lm_rot3 = rotate_volumes(flip3, lm_flip3, -15)

jit, lm_jit = jitter_3D(flip1, lm_flip1, 10)
jit2, lm_jit2 = jitter_3D(X_train, y_train, 10)
jit3, lm_jit3 = jitter_3D(flip3, lm_flip3, 10)
jit4, lm_jit4 = jitter_3D(rot2, lm_rot2, 10)

complete1 = depth_map_all_sides(X_train)
X_train=nothing
complete2 = depth_map_all_sides(flip1)
flip1=nothing
complete3 = depth_map_all_sides(flip2)
flip2=nothing
complete4 = depth_map_all_sides(flip3)
flip3=nothing
complete5 = depth_map_all_sides(rot)
rot=nothing
complete6 = depth_map_all_sides(rot2)
rot2=nothing
complete7 = depth_map_all_sides(rot3)
rot3=nothing
complete8 = depth_map_all_sides(jit)
jit=nothing
complete9 = depth_map_all_sides(jit2)
jit2=nothing
complete10 = depth_map_all_sides(jit3)
jit3=nothing
complete11 = depth_map_all_sides(jit4)

X_train = cat(complete1, complete2, complete3, complete4, complete5, complete6,
    complete7, complete8, complete9, complete10, complete11, dims=4)
y_train = cat(y_train, lm_flip1, lm_flip2, lm_flip3, lm_rot, lm_rot2, lm_rot3, lm_jit, lm_jit2, lm_jit3, lm_jit4, dims=2)

X_train = image_gradients(X_train)

# make depthmaps of the testing data
X_test = depth_map_all_sides(X_test)
X_test = image_gradients(X_test)

# define the cost function
cost(x, y) = sum((model(x)-y).^2)|>gpu

#define the model
model = Flux.mapleaves(cu, AutoLandmarking.vgg19)

# define the trainingrate and optimiser
opt = Flux.ADAM(0.00003)

# redefinition of the dropout function to work with testmode!() since this
# functionality is not working at the moment. This step only needs to be done
# if you wish to perform uncertainty estimation using the dropout method (AutoLM.response_distribution())
# after training.
using Random
function Flux.dropout(x, p; dims = :)
    q = 1 - p
    y = rand!(similar(x, Flux._dropout_shape(x, dims)))
    y .= Flux._dropout_kernel.(y, p, q)
    x .* y
end

import Zygote
Zygote.@adjoint function Flux.dropout(x, p; dims = :)
   q = 1 - p
   y = rand!(similar(x, Flux._dropout_shape(x, dims)))
   y .= Flux._dropout_kernel.(y, p, q)
   return x .* y, Δ -> (Δ .* y, nothing)
end
Zygote.refresh()

function run_model(modell, X, y)
  train_data = Flux.mapleaves(cu, X)
  train_labels = Flux.mapleaves(cu, y)
  dataset = Flux.Data.DataLoader(train_data, train_labels, batchsize = 4, shuffle=true)
  Flux.train!(cost, params(model), dataset, opt)
  testmode!(model)
  cosima = cost_whole_data_2D(train_data, train_labels, cost)
  testmode!(model, false)
  return cosima
end

# create lists to store metrics
accs = []
costs = []

# run the training for 300 epochs. batches of 128 will be loaded onto the GPU,
# which will be further subdevided into minibatches of 4 (as defined in AutoML.run_model).
for i in 1:400
  costr = 0
  for j in 1:128:size(X_train, 4)
    if j+127 > size(X_train, 4)
      train_data = X_train[:,:,:,j:end]
      train_labels = y_train[:,j:end]
    else
      train_data = X_train[:,:,:,j:j+127]
      train_labels = y_train[:,j:j+127]
    end
    cosima = run_model(model, train_data, train_labels)
    costr = costr + cosima
    train_data = nothing
    train_labels = nothing
    # GC.gc()
  end
  push!(costs, costr)
  println("epoch ", i, " finished")
  println(costr)
  testmode!(model)
  acc2, max1, min1 = AutoLandmarking.avg_accuracy_per_point(model, gpu(X_test[:,:,:,:]), gpu(y_test[:,:]), 3)
  testmode!(model, false)
  println("median deviation per point on testing dataset: ", acc2, "with maximum", max1)
  println("median sum is ", sum(acc2))
  push!(accs, sum(acc2))
  if sum(acc2)[1]<7.6
    break
  end
end

#save model
using BSON
testmode!(model) #save the model in testmode so that the dropout layers are inactive!
model_cpu = cpu(model)
BSON.@save "C:/Users/immanueldiener/Desktop/Master/master_data/model_teeth.bson" model_cpu

#load model back in
BSON.@load "C:/Users/immanueldiener/Desktop/Master/master_data/model_teeth.bson" model_cpu
model = gpu(model_cpu)

# ================== After training ===================

#solve issue with arrays inside vector
aks = []
for i in accs
  push!(aks, i[1])
end

# check what the best accuracy is
minimum(aks)

#plot the development of cost and accuracy
import Plots
Plots.plot(aks, legend=:topright, label="sum of deviations per point", color= :red, xlabel="epochs", ylabel = "acuracy")
plt = Plots.plot!(Plots.twinx(), costs, label="training loss", legend=:topleft, ylabel = "loss")

# check for outliers (greatest distance from the procrustes mean).
# teh idea is to eliminate out all the predicted landmark sets that
# deviate a lot from the procrustes mean (reference), since these are
# most likely bad predictions. do get the procrustes mean, we can use training
# testing data together to get a more accureate mean. This is an ideal approch
# because otherwise we would not realize if all the predictions are biased in some
# way that is congruent over all predictions but still wrong.

function print_costs(y, resp)
  costs=[]
  for i in 1:size(resp,2)
    dev = sum((y[:,i] .- resp[:,i]) .^2)
    push!(costs, dev)
    println(i, "    ", dev)
  end
  return costs
end

X_train2, X_test2, y_train2, y_test2 = regular_train_test_split_3d(images, lms) #to get the volumes again

response = cpu(predict_set(gpu(X_test), model)) #predictions

# translate predictions back to the original volumes and get the names for training
# and testing set. This is necessary since we performed PC-alignment and resizing before training.
retro_test = zeros(3,5,48)
scales_test = zeros(1,48)
names_test = []
names_train = []
for i in 1:48
  push!(names_test, names[i*5])
  retro_test[:,:,i] = retro[:,:,i*5]
  scales_test[1,i] = scales[i*5]
end
for i in 1:243
  if i%5!=0
    push!(names_train, names[i])
  end
end

response_scaled = deepcopy(response)

for i in 1:48
  response_scaled[:,i] .= response[:,i]./scales_test[i]
end

accuracies = []
for i in 1:48
  acc2, max1, min1 = AutoLandmarking.avg_accuracy_per_point(model, gpu(X_test[:,:,:,i:i]), gpu(y_test[:,i:i]), 3)
  push!(accuracies, sum(acc2))
  println(i, "   ", sum(acc2))
end

# set negative values to 0 and values outside the volume to the maximum along each dimension
response_unaligned = relu.(translate_lms_back(response_scaled, retro_test))
AutoLandmarking.change_values!(response_unaligned, 12.8, 12.8, >)

# put landmarks near the surface onto the surface and look at the cost change
response_on_vol = landmark_to_surface(X_test2, response_unaligned, 10)
costs = print_costs(y_test2, response_on_vol) # using the landmark_to_surface function is useful

#outlier detection based on procrustes distances

three_d_train = AutoLandmarking.to_3d_array(y_train2)
three_d_test = AutoLandmarking.to_3d_array(response_on_vol)

aligned_train = AutoLandmarking.align_all(three_d_train)
aligned_test = AutoLandmarking.align_all(three_d_test)

mean_train, stdevs = AutoLandmarking.mean_shape(aligned_train)

distances = AutoLandmarking.proc_distance(mean_train, aligned_test, stdevs)
dists = []
for i in 1:48
  push!(dists, distances[i])
end

# look at distances from mean and respective costs
for i in 1:48
  println("ind $i cost is ", costs[i], "  and dist is ", dists[i])
end

# outlier detection based on model uncertainty using the response_distribution function
uncertainties = []
for i in 1:48
  means, stddev = AutoLandmarking.response_distribution(model, gpu(X_test[:,:,:,i:i]), gpu(y_test[:,i:i]), 200)
  println("$i has uncevertaty", sum(stddev), "  and cost  ", costs[i])
  push!(uncertainties, sum(stddev))
end


function POMDPs.reward(mdp::PedCarMDP, s::PedCarMDPState, a::PedCarMDPAction, sp::PedCarMDPState)
    r = mdp.action_cost
    ego = sp.ego
    if sp.crash
        r += mdp.collision_cost
    end
    if ego.posF.s >= get_end(mdp.env.roadway[mdp.ego_goal]) &&
       get_lane(mdp.env.roadway, ego).tag == mdp.ego_goal
        r += mdp.goal_reward
    end
    return r
end


# common setup for Svalbard runs
#
# Todo

println("Time to load BITEModel.jl")
@time import BITEModel
const BM=BITEModel

update_cache = false
verbose = false
plotyes = false

# RGI data
all_svalbard = readdlm("data/rgi-6.0-huss/outlines/xyzn/svalbard/svalbard.dat")
cutoff_km2 = (5,10^10) # km^2
westfonna_ll = [18.5, 79.7] # long lat
austfonna_ll = [20, 79.2] # long lat
svalbard_nrs = Int[all_svalbard[i,1] for i=1:size(all_svalbard,1) if
                   cutoff_km2[1]<=all_svalbard[i,end]<=cutoff_km2[2] && # sizes
                   # remove Austfonna and Westfonna:
                   !(all_svalbard[i,2]>westfonna_ll[1] && all_svalbard[i,3]>westfonna_ll[2]) &&
                   !(all_svalbard[i,2]>austfonna_ll[1] && all_svalbard[i,3]>austfonna_ll[2])]
#with_thick = ["00027", "00408", "00409", "00427", "00428", "00429", "00434"][2:end] # 27 is on west or austfonna
with_thick = ["00409", "00428"] # the others are useless

brucebreen = BM.RGIGlacier("RGI60-07.01455", "Brucebreen") # area 6km^2, land, long lat: 17.291 78.485
rabotbreen = BM.RGIGlacier("RGI60-07.01478", "Rabotbreen") # area 70km^2, land
tunabreen = BM.RGIGlacier("RGI60-07.01458", "Tunabreen") # area 160km^2, tide-water
region_svalbard = 7

init_svalbard_gl(nr::Int; kws...) = init_svalbard_gl(BM.RGIGlacier(nr, region_svalbard); kws...)
init_svalbard_gl(rgi::Union{String,Symbol}; kws...) = init_svalbard_gl(BM.RGIGlacier(rgi); kws...)
function init_svalbard_gl(gid::BM.GlacierID; kws...)
    dt, pl = init_svalbard_gl_dt(gid; kws...)
    init_svalbard_gl(dt, pl)
end

function init_svalbard_gl(dt::BM.DataTable, pl)
    gl,pp,pm,pn,pl = BM.load_glacier(dt,pl)
    gb, pm = BM.make_bands(gl, pp, pm, pn);
    return gl,gb,pp,pm,pn,pl,dt
end
function init_svalbard_gl_dt(gid; kws...)
    pl = BM.LoadPara(gid; kws...)
    dt, pl = BM.make_datatable(gid,pl)

    # set temperature profile:
    dt.temp.opts[:temp] = [-10, -1.0] # colder at tongue than at top
    dt.temp.opts[:ele] = linspace(-30.0, 1717, 2) # a bit below sea to height of highest peak on Svalbard
    return dt, pl
end


const OCTAVIAN_NUM_TASKS = Ref(1)
_nthreads() = OCTAVIAN_NUM_TASKS[]

@generated function calc_factors(::Union{Val{nc},StaticInt{nc}} = num_cores()) where {nc}
    t = Expr(:tuple)
    for i ∈ nc:-1:1
        d, r = divrem(nc, i)
        iszero(r) && push!(t.args, (i, d))
    end
    t
end
# const CORE_FACTORS = calc_factors()

MᵣW_mul_factor(::True) = StaticInt{4}()
MᵣW_mul_factor(::False) = StaticInt{9}()
MᵣW_mul_factor() = MᵣW_mul_factor(has_feature(Val(:x86_64_avx512f)))

W₁Default(::True) = StaticFloat64{0.006089395198610773}()
W₂Default(::True) = StaticFloat64{0.7979822724696168}()
R₁Default(::True) = StaticFloat64{0.5900561503730485}()
R₂Default(::True) = StaticFloat64{0.762152930709678}()

W₁Default_arch(::Val{:znver1}) = StaticFloat64{0.053918949422353986}()
W₂Default_arch(::Val{:znver1}) = StaticFloat64{0.3013238122374886}()
R₁Default_arch(::Val{:znver1}) = StaticFloat64{0.6077103834481342}()
R₂Default_arch(::Val{:znver1}) = StaticFloat64{0.8775382433240162}()

W₁Default_arch(::Union{Val{:znver2},Val{:znver3}}) = StaticFloat64{0.1}()
W₂Default_arch(::Union{Val{:znver2},Val{:znver3}}) = StaticFloat64{0.993489411720157}()
R₁Default_arch(::Union{Val{:znver2},Val{:znver3}}) = StaticFloat64{0.6052218809954467}()
R₂Default_arch(::Union{Val{:znver2},Val{:znver3}}) = StaticFloat64{0.7594052633561165}()

W₁Default_arch(_) = StaticFloat64{0.1}()
W₂Default_arch(_) = StaticFloat64{0.15989396641218157}()
R₁Default_arch(_) = StaticFloat64{0.4203583148344484}()
R₂Default_arch(_) = StaticFloat64{0.8775382433240162}()

W₁Default(::False) = W₁Default_arch(VectorizationBase.cpu_name())
W₂Default(::False) = W₂Default_arch(VectorizationBase.cpu_name())
R₁Default(::False) = R₁Default_arch(VectorizationBase.cpu_name())
R₂Default(::False) = R₂Default_arch(VectorizationBase.cpu_name())

W₁Default() = W₁Default(has_feature(Val(:x86_64_avx512f)))
W₂Default() = W₂Default(has_feature(Val(:x86_64_avx512f)))
R₁Default() = R₁Default(has_feature(Val(:x86_64_avx512f)))
R₂Default() = R₂Default(has_feature(Val(:x86_64_avx512f)))

_first_cache(::StaticInt{1}) = StaticInt{1}()
_first_cache(::StaticInt) = StaticInt{2}()
first_cache() = _first_cache(VectorizationBase.num_l2cache())
second_cache() = first_cache() + One()

_first_cache_size(fcs::StaticInt) = ifelse(eq(first_cache(), StaticInt(2)) & cache_inclusive(StaticInt(2)), fcs - cache_size(One()), fcs)
_first_cache_size(::Nothing) = StaticInt(262144)
first_cache_size() = _first_cache_size(cache_size(first_cache()))

_second_cache_size(scs::StaticInt, ::True) = scs - cache_size(first_cache())
_second_cache_size(scs::StaticInt, ::False) = scs
_second_cache_size(::StaticInt{0}, ::Nothing) = StaticInt(3145728)
function second_cache_size()
    sc = second_cache()
    _second_cache_size(cache_size(sc), cache_inclusive(sc))
end

first_cache_size(::Type{T}) where {T} = first_cache_size() ÷ static_sizeof(T)
second_cache_size(::Type{T}) where {T} = second_cache_size() ÷ static_sizeof(T)

bcache_count() = VectorizationBase.num_cache(second_cache())

const BCACHEPTR = Ref{Ptr{Cvoid}}(C_NULL)
const BCACHE_LOCK = Threads.Atomic{UInt}(zero(UInt))

if Sys.WORD_SIZE ≤ 32
    const ACACHEPTR = Ref{Ptr{Cvoid}}(C_NULL)
end


using HDF5
using JSON
using Flux

"""
Loads the weights from the weights.h5 file.
"""
function weights(file="weights.h5")
    f = h5open(file, "r")
#    for ele in keys(read(f))
#        if contains(ele ,"/")
#            temp = split(ele, "/")
#            read(f)[temp[1]][temp[1]] = read(f)[temp[1]][temp[2]][temp[1]][temp[2]]
#        end
#    end
    weight = read(f)
    for ele in keys(weight)
        if !isempty(weight[ele])
            if haskey(weight[ele], ele*"_1")
                for ele2 in keys(weight[ele][ele*"_1"])
                    weight[ele][ele*"_1"][ele2] = convert(Array{Float64, N} where N, weight[ele][ele*"_1"][ele2])
                end
            elseif haskey(weight[ele], ele)
                for ele2 in keys(weight[ele][ele])
                    weight[ele][ele][ele2] = convert(Array{Float64, N} where N, weight[ele][ele][ele2])
                end
            elseif ele == "conv1"
                #weight[ele]["conv"][ele]["conv"] = convert(Array{Float64, N} where N, weight[ele]["conv"][ele]["conv"]["kernel:0"])
                weight[ele*"/conv"] = Dict{Any, Any}()
                weight[ele*"/conv"][ele*"/conv"] = Dict{Any, Any}()
                weight[ele*"/conv"][ele*"/conv"]["kernel:0"] = convert(Array{Float64, N} where N, weight[ele]["conv"][ele]["conv"]["kernel:0"])
            end
        end
    end
    return weight
end

"""
Load the structure of the model from the .JSON
file.
"""
function load_structure(file="structure.json")
    res = JSON.parse(String(read(open(file, "r"))))
    return res["config"]
end

"""
Check if model uses sequential/functional API.
"""
function check_modeltype(file)
    res = JSON.parse(String(read(open(file, "r"))))
    return res["class_name"]
end

struct new_type
    layer_type::Symbol
    input_nodes::Any
    fields::Any
end

"""
Get structure information from class_name
"""
function layer_type(a)
    if (a["class_name"] == "Conv2D" || a["class_name"] == "Conv3D")
        return :Conv
    elseif (a["class_name"] == "MaxPooling2D" || a["class_name"] == "MaxPooling3D")
        return :MaxPool
    else
        return Symbol(a["class_name"])
    end
    #    if (a["class_name"] == "Dropout")
    #    return :Dropout
    #elseif (a["class_name"] == "Flatten")
    #    return :Flatten
    #elseif (a["class_name"] == "Dense")
    #    return :Dense
    #elseif (a["class_name"] == "Activation")
    #    return :Activation
    #elseif (a["class_name"] == "Reshape")
    #    return :Reshape
    #elseif (a["class_name"] == "BatchNormalization")
    #    return :BatchNormalization
    #elseif (a["class_name"] == "InputLayer")
    #    return :InputLayer
    #elseif (a["class_name"] == "ZeroPadding2D")
    #    return :ZeroPadding2D
    #elseif (a["class_name"] == "Add")
    #    return :Add
    #elseif (a["class_name"] == "AveragePooling2D")
    #    return :AveragePooling2D
    #elseif (a["class_name"] == "Concatenate")
    #    return :Concatenate
    #elseif (a["class_name"] == "GlobalAveragePooling2D")
    #    return :GlobalAveragePooling2D
    #elseif (a["class_name"] == "Embedding")
    #    return :Embedding
    #elseif (a["class_name"] == "LSTM")
    #    return :LSTM
    #elseif (a["class_name"] == "MaxPooling1D")
    #    return :MaxPooling1D
    #elseif (a["class_name"] == "AveragePooling1D")
    #    return :AveragePooling1D
    #elseif (a["class_name"] == "Conv1D")
    #    return :Conv1D
    #end
end

"""
Extract necessary fields only from the type of layer.
"""
function fields(a)
    if layer_type(a) == :Conv
        return ["name", "strides", "activation", "kernel_size", "padding"]
    elseif layer_type(a) == :MaxPool
        return ["name", "strides", "padding", "pool_size"]
    elseif layer_type(a) == :Dropout
        return ["name", "rate"]
    elseif layer_type(a) == :Flatten
        return ["name"]
    elseif layer_type(a) == :Dense
        return ["name", "activation"]
    elseif layer_type(a) == :relu
        return ["name", "activation"]
    elseif layer_type(a) == :softmax
        return ["name", "activation"]
    elseif layer_type(a) == :Reshape
        return ["name", "target_shape"]
    elseif layer_type(a) == :BatchNormalization
        return ["name", "momentum", "epsilon"]
    elseif layer_type(a) == :InputLayer
        return ["name"]
    elseif layer_type(a) == :ZeroPadding2D
        return ["name", "padding"]
    elseif layer_type(a) == :Add
        return ["name"]
    elseif layer_type(a) == :AveragePooling2D
        return ["name","strides","padding","pool_size"]
    elseif layer_type(a) == :Concatenate
        return ["name","axis"]
    elseif layer_type(a) == :GlobalAveragePooling2D
        return ["name"]
    elseif layer_type(a) == :Embedding
        return ["name", "output_dim", "input_dim", "input_length"]
    elseif layer_type(a) == :LSTM
        return ["name", "activation", "units", "return_sequences"]
    elseif layer_type(a) == :Activation
        return ["name", "activation"]
    elseif layer_type(a) == :MaxPooling1D
        return ["name","strides", "pool_size", "padding"]
    elseif layer_type(a) == :AveragePooling1D
        return ["name","strides", "pool_size", "padding"]
    elseif layer_type(a) == :Conv1D
        return ["name", "activation", "strides", "dilation_rate", "kernel_size", "filters", "padding"]
    end
end

"""
Extract data layerwise from the structure["layers"]
array.
"""
function load_layers(a::Array{Any, 1})
    res = Array{Any, 1}()
    for ele in a
        d = Dict{Any, Any}()
        for ele2=1:length(fields(ele))
            d[fields(ele)[ele2]] = ele["config"][fields(ele)[ele2]]
        end
        if !haskey(ele, "inbound_nodes")
            ele["inbound_nodes"] = nothing
        end
        inputs = ele["inbound_nodes"]
    push!(res, new_type(layer_type(ele), inputs, d))
    end
    return res
end


using SCS
const SCS_PRIMAL_OPT = SCS.Optimizer(verbose = false)
const SCS_DUAL_OPT = DualOptimizer(SCS.Optimizer(verbose = false))
const SCS_PRIMAL_FACTORY = with_optimizer(SCS.Optimizer, verbose = false)
const SCS_DUAL_FACTORY = with_optimizer(DualOptimizer, SCS.Optimizer(verbose = false))

push!(primal_power_cone_factory, SCS_PRIMAL_FACTORY)
push!(dual_power_cone_factory, SCS_DUAL_FACTORY)
push!(dual_power_cone_optimizer, SCS_DUAL_OPT)
push!(primal_power_cone_optimizer, SCS_PRIMAL_OPT)

@testset "SCS Exponential Cone Problems" begin
    list_of_exp_problems = [
        exp1_test,
        exp2_test
    ]
    test_strong_duality(list_of_exp_problems, SCS_PRIMAL_FACTORY)
end

@testset "SCS Power Cone Problems" begin
    list_of_pow_problems = [
        pow1_test,
        pow2_test
    ]
    test_strong_duality(list_of_pow_problems, SCS_PRIMAL_FACTORY; atol = 1e-3, rtol = 1e-3)
end