text:
#=
# 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

max_stars_repo_name: paralab/Finch

max_issues_repo_name: paralab/Finch

max_forks_repo_name: paralab/Finch


text:
tile_size = 50
board_size = 4 # must be even
mini_board = [zeros(tile_size, tile_size) ones(tile_size, tile_size);
              ones(tile_size, tile_size) zeros(tile_size, tile_size)]

chessboard = repeat(mini_board, outer=(convert(Integer,(board_size/2)), convert(Integer,(board_size/2))))
img = chessboard

img[55, 40] = 0;
img[75, 40] = 0;
img[65, 30] = 0;
img[65, 20] = 0;

real_coef = 5;
imag_coef = -2;
flow = real_coef .* ones(size(img)) .+ imag_coef .* im .* ones(size(img))

imgshowflow(img, flow)

imgw = LAP_julia.interpolation.imWarp_replicate(img, real(flow), imag(flow))
imgshow(imgw)

imgshowflow(imgw, zeros(size(imgw)))
imgshowflow(img, zeros(size(imgw)))
save("imgw.png", reverse(imgw, dims=1))
save("img.png", reverse(img, dims=1))

imgshowflow(imgw, flow)

u_est, coeffs = LAP_julia.registration.single_lap(img, imgw, 3, 7, [15, 15])

showflow(u_est)
showflow(flow)
u_est[55,49]

## try to use opposite of u_est to get img from imgw

img_remade = LAP_julia.interpolation.imWarp(imgw, real(u_est), imag(u_est))

fig = PyPlot.figure(dpi = 300, figsize = (5, 5));
PyPlot.imshow(img_remade, cmap = :gray);
ax = gca();
ax.set_ylim(0, size(img_remade)[1]);
gcf()

save("img_remade.png", reverse(img_remade, dims=1))

fig = PyPlot.figure(dpi = 300, figsize = (5, 5));
PyPlot.imshow(img, cmap = :gray);
ax = gca();
ax.set_ylim(0, size(img)[1]);
gcf()

# difference:
img_dif = img - img_remade;
img_dif[40,101]

fig = PyPlot.figure(dpi = 300, figsize = (5, 5));
PyPlot.imshow(img_dif, cmap = :gray);
ax = gca();
ax.set_ylim(0, size(img_dif)[1]);
gcf()

## inpaint u_est

u_inp = u_est;
LAP_julia.inpaint.inpaint_nans!(u_inp)

showflow(u_inp)
showflow(flow)

mean(real(flow))
mean(imag(flow))

mean(real(u_inp))
mean(imag(u_inp))

mean(real(flow))
mean(imag(flow))

img_inp = LAP_julia.interpolation.imWarp(imgw, real(u_inp), imag(u_inp))

imgshow(img_inp)
imgshow(img)
imgshow(imgw)

save("img_inp.png", reverse(img_inp, dims=1))

# difference:
img_dif2 = img - img_inp;

fig = PyPlot.figure(dpi = 300, figsize = (5, 5));
PyPlot.imshow(img_dif2, cmap = :gray);
ax = gca();
ax.set_ylim(0, size(img_dif2)[1]);
gcf()

img_dif2[50, 75]

ext: jl


text:
module test_julia_complex

using Test

# https://en.wikipedia.org/wiki/Conjugate_transpose#Example
A = [1     -2-im  5
     1+im     im  4-2im]

@test transpose(A) == [ 1      1+im
                       -2-im     im
                        5     4-2im]

@test A' == adjoint(A) == [ 1      1-im
                           -2+im    -im
                            5     4+2im]

@test conj(A) == [1     -2+im  5
                  1-im    -im  4+2im]

end # module test_julia_complex


text:
# # Testing time step and convergence to steady state

# Here we test the convergence in the case of a known globally stable steady state, made of a single mode and associated with a corresponding single-mode forcing.

# In order to check that, we implement a simple (and far from optimal) time-step algorithm. We'll improve that later.

# Here are the packages we are gonna need.

using FFTW
using Plots
using Test
using Random

# ## The spatial domain and discretization

L = 2π
κ₀ = 2π/L
N = 128
Nsub = 84
x = y = (L/N):(L/N):L

# Operators

Dx_hat = im * κ₀ * [ifelse(k1 ≤ N/2 + 1, k1 - 1, k1 - 1 - N) for k2 in 1:N/2+1, k1 in 1:N]
Dy_hat = im * κ₀ * [k2 - 1 for k2 in 1:N/2+1, k1 in 1:N]

Delta_hat = - κ₀^2 * [
    ifelse(k1 ≤ N/2 + 1, (k1 - 1)^2 + (k2 - 1)^2, (k1 - 1 - N)^2 + (k2 - 1)^2)
    for k2 in 1:N/2+1, k1 in 1:N
]

Hu_hat = - Dy_hat ./ Delta_hat
Hu_hat[1, 1] = 0.0
Hv_hat = Dx_hat ./ Delta_hat
Hv_hat[1, 1] = 0.0

# One-mode steady state

ν = 1.0e-0 # viscosity
κ = (x = 1, y = 2) # forced mode
α = (re = 0.1, im = 0.05) # strength

vort_steady = 2κ₀^2 * sum(abs2, κ) * (
    α.re * cos.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
    - α.im * sin.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
)
vort_steady_hat = rfft(vort_steady)

heatmap(x, y, vort_steady, xlabel="x", ylabel="y", title="steady state vorticity", titlefont=12)

g_steady = ν * (
    2α.re * κ₀^4 * sum(abs2, κ)^2 * cos.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
    - 2α.im * κ₀^4 * sum(abs2, κ)^2 * sin.(κ₀ * (κ.x * one.(y) * x' + κ.y * y * one.(x)'))
)
g_steady_hat = rfft(g_steady)

heatmap(x, y, g_steady, xlabel="x", ylabel="y", title="forcing term", titlefont=12)

Gr = sqrt(
    2 * sum(abs2, α) / κ₀^4 / ν^4
)

@info "Grashof number: $(round(Gr, digits=4))"

# evolution step

function step!(vort_hat, dt, params)
    Exp_nu_dt_Delta_hat, Dx_hat, Dy_hat, Hu_hat, Hv_hat, g_hat, N, Nsub = params
    u_hat = Hu_hat .* vort_hat
    v_hat = Hv_hat .* vort_hat
    u = irfft(u_hat, N)
    v = irfft(v_hat, N)
    vort = irfft(vort_hat, N)
    wu_hat = rfft(vort .* u)
    wv_hat = rfft(vort .* v)
    vort_hat .= Exp_nu_dt_Delta_hat .* (
        vort_hat .+ dt * (
            g_hat .- Dx_hat .* wu_hat .- Dy_hat .* wv_hat
        )
    )
    ## dealiasing
    vort_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    vort_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im
    return vort_hat
end

# Initial vorticity

rng = Xoshiro(123)
num_modes = 4
vort_init = sum(
    [
        2κ₀^2 * (kx^2 + ky^2) * (
            ar * cos.(κ₀ * (kx * one.(y) * x' + ky * y * one.(x)'))
            - ai * sin.(κ₀ * (kx * one.(y) * x' + ky * y * one.(x)'))
        )
        for (kx, ky, ar, ai) in zip(
            rand(rng, 1:div(N,10), num_modes),
            rand(rng, 1:div(N,10), num_modes),
            rand(rng, num_modes),
            rand(rng, num_modes)
        )
    ]
)

vort_init_hat = rfft(vort_init)

heatmap(x, y, vort_init, xlabel="x", ylabel="y", title="initial vorticity", titlefont=12)

# Time evolution setup

dt = 1.0e-5 # 1.0e-3 a 1.0e-5
t_final = 3.0
Exp_nu_dt_Delta_hat = exp.(ν * dt * Delta_hat)
params = Exp_nu_dt_Delta_hat, Dx_hat, Dy_hat, Hu_hat, Hv_hat, g_steady_hat, N, Nsub
num_steps = Int(round(t_final / dt))

vort_hat = copy(vort_init_hat)

println("Enstrophy convergence:")
println(sum(abs2, vort_hat - vort_steady_hat) * (1/N)^4)
for n in 1:num_steps
    step!(vort_hat, dt, params)
    if rem(n, 1000) == 0
        println(sum(abs2, vort_hat - vort_steady_hat) * (1/N)^4)
    end
end

vort = irfft(vort_hat, N)
heatmap(x, y, vort, xlabel="x", ylabel="y", title="vorticity", titlefont=12)

surface(x, y, vort - vort_steady, xlabel="x", ylabel="y", zlabel="error", title="difference `vort_final .- vort_steady`", titlefont=12)

# The difference from the actual fixed point to the limit value of the discretized method decreases with dt
# Eg. think of `̇ẋ = - a (x - xₑ)`, where `a>0` and `xₑ` is the actual equilibrium.
# Discretizing it with `x^{n+1}exp(at_{n+1}) - x^n exp(at_n) = dt a exp(at_n)` yields the fixed point
# `x̄ = (exp(-a dt) * dt * a) / (1 - exp(-a * dt))`.
# We have `x̄ - xₑ → 0` as `dt → 0`.
#
# `f(a, k) = k * a * exp(-a * k) / (1 - exp(-a * k))`


text:
cross(x::Vector{XFloat16}, y::Vector{XFloat16}) = reinterpret(XFloat16, cross(reinterpret(Float32, x), reinterpret(Float32, y)))
(*)(x::Vector{XFloat16}, y::Vector{XFloat16}) = reinterpret(XFloat16, (*)(reinterpret(Float32, x), reinterpret(Float32, y)))
dot(x::Vector{XFloat16}, y::Vector{XFloat16}) = reinterpret(XFloat16, dot(reinterpret(Float32, x), reinterpret(Float32, y)))
(/)(x::Matrix{XFloat16}, y::Matrix{XFloat16}) = reinterpret(XFloat16, (/)(reinterpret(Float32, x), reinterpret(Float32, y)))
(\)(x::Matrix{XFloat16}, y::Matrix{XFloat16}) = reinterpret(XFloat16, (\)(reinterpret(Float32, x), reinterpret(Float32, y)))
factorize(x::Matrix{XFloat16}) = reinterpret(XFloat16, factorize(reinterpret(Float32, x)))
inv(x::Matrix{XFloat16}) = reinterpret(XFloat16, inv(reinterpret(Float32, x)))
det(x::Matrix{XFloat16}) = reinterpret(XFloat16, det(reinterpret(Float32, x)))
logdet(x::Matrix{XFloat16}) = reinterpret(XFloat16, logdet(reinterpret(Float32, x)))
logabsdet(x::Matrix{XFloat16}) = reinterpret(XFloat16, logabsdet(reinterpret(Float32, x)))
size(x::Matrix{XFloat16}) = size(Float32.(x))
eigvals(x::Matrix{XFloat16}) = reinterpret(XFloat16, eigvals(reinterpret(Float32, x)))
eigvecs(x::Matrix{XFloat16}) = reinterpret(XFloat16, eigvecs(reinterpret(Float32, x)))
svdvals(x::Matrix{XFloat16}) = reinterpret(XFloat16, svdvals(reinterpret(Float32, x)))

cross(x::Vector{XFloat32}, y::Vector{XFloat32}) = reinterpret(XFloat32, cross(reinterpret(Float64, x), reinterpret(Float64, y)))
(*)(x::Vector{XFloat32}, y::Vector{XFloat32}) = reinterpret(XFloat32, (*)(reinterpret(Float64, x), reinterpret(Float64, y)))
dot(x::Vector{XFloat32}, y::Vector{XFloat32}) = reinterpret(XFloat32, dot(reinterpret(Float64, x), reinterpret(Float64, y)))
(/)(x::Matrix{XFloat32}, y::Matrix{XFloat32}) = reinterpret(XFloat32, (/)(reinterpret(Float64, x), reinterpret(Float65, y)))
(\)(x::Matrix{XFloat32}, y::Matrix{XFloat32}) = reinterpret(XFloat32, (\)(reinterpret(Float64, x), reinterpret(Float64, y)))
factorize(x::Matrix{XFloat32}) = reinterpret(XFloat32, factorize(reinterpret(Float64, x)))
inv(x::Matrix{XFloat32}) = reinterpret(XFloat32, inv(reinterpret(Float64, x)))
det(x::Matrix{XFloat32}) = reinterpret(XFloat32, det(reinterpret(Float64, x)))
logdet(x::Matrix{XFloat32}) = reinterpret(XFloat32, logdet(reinterpret(Float64, x)))
logabsdet(x::Matrix{XFloat32}) = reinterpret(XFloat32, logabsdet(reinterpret(Float64, x)))
size(x::Matrix{XFloat32}) = size(Float64.(x))
eigvals(x::Matrix{XFloat32}) = reinterpret(XFloat32, eigvals(reinterpret(Float64, x)))
eigvecs(x::Matrix{XFloat32}) = reinterpret(XFloat32, eigvecs(reinterpret(Float64, x)))
svdvals(x::Matrix{XFloat32}) = reinterpret(XFloat32, svdvals(reinterpret(Float64, x)))


text:
@testset "Base Extension of Grid-search Struct" begin
    ranges = [[1.,2.,10], [3.,4.,10], [5.,6.,10]]
    param_ranges = ParameterGrid(ranges)

    @testset "Constuctor of ParameterGrid" begin
        @test param_ranges.len == FindSteadyStates.mul(Integer,[i[end] for i in ranges])
        @test length(param_ranges.indexes) == param_ranges.len
    end

    @testset "Length & size" begin
        #Base.length
        @test Base.length(param_ranges) == (FindSteadyStates.mul(Integer,[i[end] for i in ranges]))
        #Base.size
        @test Base.size(param_ranges) == ((FindSteadyStates.mul(Integer,[i[end] for i in ranges])),)
    end
end

@testset "Recursive index" begin

    ranges = [ (1,4,4), (2,8,4), (1,2,2) ]
    ind_max = FindSteadyStates.mul([i[end] for i in ranges])

    # first index vector
    vec_one = [1 for i in ranges]
    # last index vector
    vec_end = [i[end] for i in ranges]

    param_ranges = ParameterGrid(ranges)

    @test vec_one == recursive_index(param_ranges, 1)
    @test vec_end == recursive_index(param_ranges, ind_max)
    @test [4,1,1] == recursive_index(param_ranges, 4)
    @test [1,2,1] == recursive_index(param_ranges, 5)

end

## Indexing
@testset "get index for parameter generation" begin
    # Ranges
    ranges = [ (1.,4.,4), (2.,8.,4), (1.,2.,2) ]
    total_is = FindSteadyStates.mul([i[end] for i in ranges])

    # Grid search samplers
    para1 = ParameterGrid(ranges; method=UniformGrid())
    para2 = ParameterGrid(ranges; method=LogGrid(decay=10.))

    # number of index should equal to total grids.
    @test length(para1) == total_is
    @test length(para2) == total_is

    # Sampling: para 1
    @info "Sample the range: $ranges"
    @info "Print grid vecs: uniform distributed"
    for i in 1:1:length(para1)
        println(para1[i])
    end

    # Sampling Para 2
    @info "Print grid vecs: log distributed"
    for i in 1:1:length(para2)
        println(para2[i])
    end
    @info "DONE"
end

@testset "Iterator of ParameterGrid" begin

    ranges = [ (1.,4.,4), (2.,8.,4), (1.,2.,2) ]
    param = ParameterGrid(ranges)
    @info "Display ranges ($ranges)"

    @testset "iteration" begin

    for (i,j) in enumerate(param )
        println(j)
        @test param[i] == j
    end
    end

    @info "Done"
end

@testset "Range input" begin
    g1 = ParameterGrid([
        1.0:1.0:10.0,
        1:1.0:10.0
    ])

    g2 = ParameterGrid([
        range(0,10;length=10),
        range(0,10;length=10)
    ])

    r1 = ParameterRandom([
        1.0:1.0:10.0,
        1:1.0:10.0
    ])

    r2 = ParameterRandom([
        range(0,10;length=10),
        range(0,10;length=10)
    ])

    # test
    @test length(g2) == 10*10
    @test length(g1) == 10*10
    @test length(r1) == 10*10
    @test length(r2) == 10*10

end


text:
immutable Weibull <: ContinuousUnivariateDistribution
    α::Float64   # shape
    β::Float64   # scale

    function Weibull(α::Real, β::Real)
    	(zero(α) < α && zero(β) < β) || error("Weibull: both shape and scale must be positive")
    	new(float64(α), float64(β))
    end

    Weibull(α::Real) = Weibull(α, 1.0)
    Weibull() = new(1.0, 1.0)
end

@distr_support Weibull 0.0 Inf

#### Parameters

shape(d::Weibull) = d.α
scale(d::Weibull) = d.β

params(d::Weibull) = (d.α, d.β)

#### Statistics

mean(d::Weibull) = d.β * gamma(1.0 + 1.0 / d.α)
median(d::Weibull) = d.β * logtwo ^ (1.0 / d.α)
mode(d::Weibull) = d.α > 1.0 ? (iα = 1.0 / d.α; d.β * (1.0 - iα) ^ iα) : 0.0

var(d::Weibull) = d.β^2 * gamma(1.0 + 2.0 / d.α) - mean(d)^2

function skewness(d::Weibull)
    μ = mean(d)
    σ2 = var(d)
    σ = sqrt(σ2)
    r = μ / σ
    gamma(1.0 + 3.0 / d.α) * (d.β / σ)^3 - 3.0 * r - r^3
end

function kurtosis(d::Weibull)
    α, β = params(d)
    μ = mean(d)
    σ = std(d)
    γ = skewness(d)
    r = μ / σ
    r2 = r^2
    r4 = r2^2
    (β / σ)^4 * gamma(1.0 + 4.0 / α) - 4.0 * γ * r - 6.0 * r2 - r4 - 3.0
end

function entropy(d::Weibull)
    α, β = params(d)
    0.5772156649015328606 * (1.0 - 1.0 / α) + log(β / α) + 1.0
end

#### Evaluation

function pdf(d::Weibull, x::Float64)
    if x >= 0.0
        α, β = params(d)
        z = x / β
        (α / β) * z^(α - 1.0) * exp(-z^α)
    else
        0.0
    end
end

function logpdf(d::Weibull, x::Float64)
    if x >= 0.0
        α, β = params(d)
        z = x / β
        log(α / β) + (α - 1.0) * log(z) - z^α
    else
        -Inf
    end
end

zv(d::Weibull, x::Float64) = (x / d.β) ^ d.α
xv(d::Weibull, z::Float64) = d.β * z ^ (1.0 / d.α)

cdf(d::Weibull, x::Float64) = x > 0.0 ? -expm1(-zv(d, x)) : 0.0
ccdf(d::Weibull, x::Float64) = x > 0.0 ? exp(-zv(d, x)) : 1.0
logcdf(d::Weibull, x::Float64) = x > 0.0 ? log1mexp(-zv(d, x)) : -Inf
logccdf(d::Weibull, x::Float64) = x > 0.0 ? -zv(d, x) : 0.0

quantile(d::Weibull, p::Float64) = xv(d, -log1p(-p))
cquantile(d::Weibull, p::Float64) = xv(d, -log(p))
invlogcdf(d::Weibull, lp::Float64) = xv(d, -log1mexp(lp))
invlogccdf(d::Weibull, lp::Float64) = xv(d, -lp)

function gradlogpdf(d::Weibull, x::Float64)
    if insupport(Weibull, x)
        α, β = params(d)
        (α - 1.0) / x - α * x^(α - 1.0) / (β^α)
    else
        0.0
    end
end

#### Sampling

rand(d::Weibull) = xv(d, randexp())


text:
using JuMP, vOptGeneric, GLPK, Combinatorics
using Test, LinearAlgebra

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)

@variable(m, x[1:6,1:6], Bin)
@variable(m, y, Bin)
M = 50
p1 =   [M 4 5 M M M;
        M M 2 1 2 7;
        M M M 5 2 M;
        M M 5 M M 3;
        M M M M M 4;
        M M M M M M]
p2 =   [M 3 1 M M M;
        M M 1 4 2 2;
        M M M 1 7 M;
        M M 1 M M 2;
        M M M M M 2;
        M M M M M M]

@addobjective(m, Min, sum(x[i,j]*p1[i,j] for i=1:6, j=1:6) + 5)
@addobjective(m, Min, sum(x[i,j]*p2[i,j] for i=1:6, j=1:6) - 5)
@constraint(m, sum(x[1,j] for j = 1:6) == 1)
@constraint(m, sum(x[i,6] for i = 1:6) == 1)
@constraint(m, cstr[i=2:5], sum(x[i,j] for j = 1:6) - sum(x[j,i] for j =1:6) == 0)

vSolve(m, method=:epsilon)

Y_N = getY_N(m)
f1, f2 = first.(Y_N), last.(Y_N)
@test f1 == [13, 15, 16, 18] && f2 == [4, 2, 0, -1]
@test value.(x, 1) == [0.0 1.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 1.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0]
printX_E(m)

######################

# to test once https://github.com/JuliaOpt/JuMP.jl/issues/1956 is fixed
@test_broken begin
    m_copy = copy(m)
    # vSolve(m_copy, GLPK.Optimizer, method=:epsilon)
    # Y_N = getY_N(m)
    # f1 = map(x -> x[1], Y_N) ; f2 = map(x -> x[2], Y_N)
    # @test f1 == [8.0,10.0,11.0,13.0] && f2 == [9.0,7.0,5.0,4.0]
    # @test value.(x, 1) == [0.0 1.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 1.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0]
    # printX_E(m_copy)
end

######################

# tests DenseAxisArrays indexing
# http://www.juliaopt.org/JuMP.jl/v0.19.1/variables/#variable_jump_arrays-1

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)
# JuMP.set_optimizer_attribute(m, "logLevel", 0)
cities = ["Paris", "New-York", "Madrid"]
@variable(m, x[cities], Bin)
@constraint(m, constr, x["Paris"] >= 1)
@addobjective(m, Min, sum(x))
@addobjective(m, Max, sum(x))
vSolve(m, method=:epsilon)
@show value.(x, 1)
@test value.(x, 1)["Paris"] == 1.0
@test value.(x, 1)["New-York"] == 0.0
@test value.(x, 1)["Madrid"] == 0.0

######################

# TODO : write a test for SparseAxisArray
# http://www.juliaopt.org/JuMP.jl/v0.19.1/variables/#variable_sparseaxisarrays-1

######################

# TODO : re-implement writeMOP and parseMOP
# writeMOP(m, "test.txt")
# m = parseMOP("test.txt", solver = GLPKSolverMIP())
# solve(m, method=:dichotomy)
# Y_N = getY_N(m)
# f1 = map(x -> x[1], Y_N)
# f2 = map(x -> x[2], Y_N)
# @test f1 == [8.0, 11.0, 13.0]
# @test f2 == [9.0, 5.0, 4.0]

######################

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)
# JuMP.set_optimizer_attribute(m, "logLevel", 0)

p1 = [77,94,71,63,96,82,85,75,72,91,99,63,84,87,79,94,90,60,69,62]
p2 = [65,90,90,77,95,84,70,94,66,92,74,97,60,60,65,97,93,60,69,74]
w = [80,87,68,72,66,77,99,85,70,93,98,72,100,89,67,86,91,79,71,99]
c = 900
size = length(p1)

@variable(m, x[1:size], Bin)
@addobjective(m, Max, dot(x, p1))
@addobjective(m, Max, dot(x, p2))
@constraint(m, dot(x, w) <= c)

vSolve(m, method=:dicho)

Y_N = getY_N(m)
f1, f2 = first.(Y_N), last.(Y_N)

@test f1 == [918.0, 934.0, 948.0, 955.0]
@test f2 == [984.0, 971.0, 939.0, 906.0]
@test value.(x, 1) ≈ [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0]

vSolve(m, method=:chalmet)

Y_N = getY_N(m)
f1, f2 = first.(Y_N), last.(Y_N)

@test f1 == [918., 924., 927., 934., 935., 940., 943., 948., 949., 952., 955.]
@test f2 == [984., 975., 972., 971., 947., 943., 940., 939., 915., 909., 906.]
@test value.(x, 1) ≈ [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0]

########################

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)
# JuMP.set_optimizer_attribute(m, "logLevel", 0)

@variable(m, 0 <= x[1:4] <= 2, Int)
@addobjective(m, Max, sum(x))
@addobjective(m, Min, sum(x))
@addobjective(m, Max, x[3])
@addobjective(m, Max, x[4])
@constraint(m, sum(x) <= 3)
vSolve(m, method=:lex)
Y_N = getY_N(m)

for (ind, perm) in enumerate(permutations(1:4, 4))
    if perm[1] == 2
        @test Y_N[ind] == [0., 0., 0., 0.]
    elseif perm[2] == 2 && perm[1] != 1
        @test maximum(Y_N[ind]) == 2.
    end
end

max_stars_repo_name: vOptSolver/vOptGeneric

max_issues_repo_name: vOptSolver/vOptGeneric

max_forks_repo_name: vOptSolver/vOptGeneric


text:
"""
    prepare_observables(m::HubbardModelAttractive)

Initializes observables for the attractive Hubbard model and returns a `Dict{String, Observable}`.

See also [`measure_observables!`](@ref) and [`finish_observables!`](@ref).
"""
@inline function prepare_observables(mc::DQMC, m::HubbardModelAttractive)
    obs = Dict{String,Observable}()
    obs["confs"] = Observable(HubbardConf, "Configurations")
    obs["greens"] = Observable(typeof(mc.s.greens), "Equal-times Green's function")
    obs["Eboson"] = Observable(Float64, "Bosonic energy")

    return obs
end

"""
    measure_observables!(m::HubbardModelAttractive, obs::Dict{String,Observable}, conf::HubbardConf, E::Float64)

Measures observables and updates corresponding `Observable` objects in `obs`.

See also [`prepare_observables`](@ref) and [`finish_observables!`](@ref).
"""
@noinline function measure_observables!(mc::DQMC, m::HubbardModelAttractive,
							obs::Dict{String,Observable}, conf::HubbardConf)
    push!(obs["confs"], mc.conf)
    push!(obs["greens"], greens(mc))
    push!(obs["Eboson"], energy_boson(mc, m, conf))
    nothing
end

"""
    measure_observables!(mc::DQMC, m::HubbardModelAttractive, obs::Dict{String,Observable})

Finish measurements of observables.

See also [`prepare_observables`](@ref) and [`measure_observables!`](@ref).
"""
@inline function finish_observables!(mc::DQMC, m::HubbardModelAttractive,
							obs::Dict{String,Observable})
    nothing
end

max_stars_repo_name: yosefcap/MC

max_issues_repo_name: yosefcap/MC

max_forks_repo_name: yosefcap/MC


text:
"""
    prettylens(lens::Lens; sprint_kwargs...) :: String
    prettylens(io::IO, lens::Lens)

Print or return more compact and easier-to-read string representation of
`lens` than `show`.

# Examples
```jldoctest
julia> using Setfield, Kaleido

julia> prettylens(
           (@lens _.a) ∘ MultiLens((
               (@lens last(_)),
               (@lens _[:c].d) ∘ settingasℝ₊,
           ));
           context = :compact => true,
       )
"◻.a∘〈last(◻),◻[:c].d∘(←exp|log→)〉"
```
"""
prettylens

# Importing internals (let's try not break precompilation)
isdefined(Setfield, :has_atlens_support) && using Setfield: has_atlens_support
isdefined(Setfield, :print_in_atlens) && using Setfield: print_in_atlens

printhole(io) = printstyled(io, "◻", color=:light_black)

prettylens(lens::Lens; kwargs...) = sprint(prettylens, lens; kwargs...)

function prettylens(io::IO, lens::ComposedLens)
    prettylens_via_atlens(io,  lens) && return

    prettylens(io, lens.outer)
    printstyled(io, (:compact => true) in io ? "∘" : " ∘ "; color=:light_black)
    prettylens(io, lens.inner)
end

prettylens(io::IO, lens::Lens) =
    prettylens_via_atlens(io, lens) || show(io, lens)

function prettylens_via_atlens(io, lens)
    if has_atlens_support(lens)
        name = sprint(print_in_atlens, lens, context=io)
        if startswith(name, "(@lens ") && endswith(name, ")")
            if startswith(name, "(@lens _")
                printhole(io)
                print(io, chop(name, head=length("(@lens _"), tail=length(")")))
            else
                parts = split(
                    chop(name; head=length("(@lens "), tail=length(")")),
                    "(_)",
                )
                print(io, parts[1])
                for s in parts[2:end]
                    print(io, '(')
                    printhole(io)
                    print(io, ')')
                    print(io, s)
                end
            end
            return true
        end
    end
    return false
end

function prettylens(io::IO, lens::MultiLens)
    print(io, '〈')
    for (i, l) in enumerate(lens.lenses)
        if i != 1
            print(io, ',')
            (:compact => true) in io || print(io, ' ')
        end
        prettylens(io, l)
    end
    print(io, '〉')
end

function prettylens(io::IO, lens::BijectionLens{<:FunctionPair})
    bi = lens.bijection
    print(io, "(←", bi.tofield, "|", bi.fromfield, "→)")
end

function prettylens(io::IO, lens::BijectionLens)
    print(io, "(←", lens.bijection, "→)")
end

function prettylens(io::IO, lens::BijectionLens{<:XFBijection{INV}}) where INV
    if (@isdefined TransformVariables)
        io = IOContext(io, :module => TransformVariables)
    end
    print(io, "(←")
    INV || print(io, '|')
    print(io, lens.bijection.xf)
    INV && print(io, '|')
    print(io, "→)")
end

function prettylens(io::IO, lens::ConverterLens)
    print(io, "(←|", _singleton_callable(lens.f), "→)")
end

function prettylens(io::IO, setter::ToField)
    print(io, "(←")
    print(io, setter.f)
    print(io, "|❌→)")
end


text:
struct MatrixLangevin{N,T,n,k} <: ParameterizedMeasure{N}
    par::NamedTuple{N,T}
    manifold::Stiefel{n,k,ℝ}
end
function MatrixLangevin(n, k; kwargs...)
    par = NamedTuple(kwargs)
    return MatrixLangevin(par, Stiefel(n, k))
end

const MatrixVonMisesFisher = MatrixLangevin

Manifolds.base_manifold(d::MatrixLangevin) = getfield(d, :manifold)

function MeasureTheory.basemeasure(μ::MatrixLangevin{(:F,)})
    return normalize(Hausdorff(base_manifold(μ), μ.F))
end

function MeasureTheory.logdensity(d::MatrixLangevin{(:F,)}, x::AbstractMatrix)
    F = d.F
    return tr_At_B(F, x) - log(pFq((), (n//2,), rmul!(F'F, 1//4)))
end

function Base.rand(rng::AbstractRNG, T::Type, d::MatrixLangevin{(:F,)})
    F = d.F
    trΛ = sum(svdvals(F))
    while true
        x = rand(rng, T, basemeasure(d))
        rand(rng) < exp(tr_At_B(F, x) - trΛ) && return x
    end
end


text:
# Extended Sparse matrix interface

"""
    sparse_from_coo(::Type{T} where T,I,J,V,m,n)
"""
function sparse_from_coo(::Type{T} where T,I,J,V,m,n)
  @abstractmethod
end

"""
    fill_entries!(A::AbstractSparseMatrix,v)

Fills the non-zero entries in the sparse matrix A with a given value v.
"""
function fill_entries!(A::AbstractSparseMatrix,v)
  @abstractmethod
end

"""
    push_coo!(::Type{T} where T, I,J,V,ik,jk,vk)

Inserts entries in COO vectors for further building a sparse matrix of type T.
"""
function push_coo!(::Type{T} where T,I,J,V,ik,jk,vk)
  @abstractmethod
end

"""
    is_entry_stored(::Type{T} where T,i,j) -> Bool

Tells if the entry with coordinates `[i,j]` will be stored when calling function push_coo!
"""
function is_entry_stored(::Type{T} where T,i,j)
  @abstractmethod
end

"""
    finalize_coo!(::Type{T} where T,I,J,V,m,n)

Check and insert diagonal entries in COO vectors if needed.
"""
function finalize_coo!(::Type{T} where T,I,J,V,m,n)
  @abstractmethod
end

"""
"""
function allocate_coo_vectors(::Type{M},n::Integer) where M
  return (zeros(Int,n), zeros(Int,n), zeros(Float64,n))
end

function allocate_coo_vectors(::Type{M},n::Integer) where {Tv,Ti,M<:AbstractSparseMatrix{Tv,Ti}}
  return (zeros(Ti,n), zeros(Ti,n), zeros(Tv,n))
end


text:
c = [2,3,4] .^ 2

r = [4, 6] .* [3, 4]

f = reshape(1:81, 9, 9)

println(f);

A = fill(5.0, (3, 3))

println(A);

println(pi);

max_stars_repo_name: sandos/MastersOfCode

max_issues_repo_name: sandos/MastersOfCode

max_forks_repo_name: sandos/MastersOfCode


text:
module Setup

mul2(x) = x .* 2

end # module Setup

using .Setup


text:
using StatsBase
using Distributions
using CategoricalArrays
using DataFrames
using CSV
using Random
using RatPack

file = "../examples/test_ratings2.csv"
(df,r) = read_ratings(file)

Random.seed!(1)

competitions1, w1 = simulate( r, SimRoundRobin(2), Logistic(0.0, 1000.0/log(10.0)) )

competitions2, w2 = simulate( r, SimElimination(1), Logistic(0.0, 1000.0/log(10.0)) )


text:
@testset "container.jl" begin
    @testset "copy particle container" begin
        pc = AdvancedPS.ParticleContainer(AdvancedPS.Trace[])
        newpc = copy(pc)

        @test newpc.logWs == pc.logWs
        @test typeof(pc) === typeof(newpc)
    end

    @testset "particle container" begin
        # Create a resumable function that always returns the same log probability.
        function fpc(logp)
            f = let logp = logp
                rng -> begin
                    while true
                        produce(logp)
                    end
                end
            end
            return f
        end

        # Create particle container.
        logps = [0.0, -1.0, -2.0]
        particles = [AdvancedPS.Trace(fpc(logp), AdvancedPS.TracedRNG()) for logp in logps]
        pc = AdvancedPS.ParticleContainer(particles)

        # Initial state.
        @test pc.logWs == zeros(3)
        @test AdvancedPS.getweights(pc) == fill(1 / 3, 3)
        @test all(AdvancedPS.getweight(pc, i) == 1 / 3 for i in 1:3)
        @test AdvancedPS.logZ(pc) ≈ log(3)
        @test AdvancedPS.effectiveSampleSize(pc) == 3

        # Reweight particles.
        AdvancedPS.reweight!(pc)
        @test pc.logWs == logps
        @test AdvancedPS.getweights(pc) ≈ exp.(logps) ./ sum(exp, logps)
        @test all(
            AdvancedPS.getweight(pc, i) ≈ exp(logps[i]) / sum(exp, logps) for i in 1:3
        )
        @test AdvancedPS.logZ(pc) ≈ log(sum(exp, logps))

        # Reweight particles.
        AdvancedPS.reweight!(pc)
        @test pc.logWs == 2 .* logps
        @test AdvancedPS.getweights(pc) == exp.(2 .* logps) ./ sum(exp, 2 .* logps)
        @test all(
            AdvancedPS.getweight(pc, i) ≈ exp(2 * logps[i]) / sum(exp, 2 .* logps) for
            i in 1:3
        )
        @test AdvancedPS.logZ(pc) ≈ log(sum(exp, 2 .* logps))

        # Resample and propagate particles with reference particle
        particles_ref = [
            AdvancedPS.Trace(fpc(logp), AdvancedPS.TracedRNG()) for logp in logps
        ]
        pc_ref = AdvancedPS.ParticleContainer(particles_ref)
        AdvancedPS.resample_propagate!(
            Random.GLOBAL_RNG, pc_ref, AdvancedPS.resample_systematic, particles_ref[end]
        )
        @test pc_ref.logWs == zeros(3)
        @test AdvancedPS.getweights(pc_ref) == fill(1 / 3, 3)
        @test all(AdvancedPS.getweight(pc_ref, i) == 1 / 3 for i in 1:3)
        @test AdvancedPS.logZ(pc_ref) ≈ log(3)
        @test AdvancedPS.effectiveSampleSize(pc_ref) == 3
        @test pc_ref.vals[end] === particles_ref[end]

        # Resample and propagate particles.
        AdvancedPS.resample_propagate!(Random.GLOBAL_RNG, pc)
        @test pc.logWs == zeros(3)
        @test AdvancedPS.getweights(pc) == fill(1 / 3, 3)
        @test all(AdvancedPS.getweight(pc, i) == 1 / 3 for i in 1:3)
        @test AdvancedPS.logZ(pc) ≈ log(3)
        @test AdvancedPS.effectiveSampleSize(pc) == 3

        # Reweight particles.
        AdvancedPS.reweight!(pc)
        @test pc.logWs ⊆ logps
        @test AdvancedPS.getweights(pc) == exp.(pc.logWs) ./ sum(exp, pc.logWs)
        @test all(
            AdvancedPS.getweight(pc, i) ≈ exp(pc.logWs[i]) / sum(exp, pc.logWs) for i in 1:3
        )
        @test AdvancedPS.logZ(pc) ≈ log(sum(exp, pc.logWs))

        # Increase unnormalized logarithmic weights.
        logws = copy(pc.logWs)
        AdvancedPS.increase_logweight!(pc, 2, 1.41)
        @test pc.logWs == logws + [0, 1.41, 0]

        # Reset unnormalized logarithmic weights.
        logws = pc.logWs
        AdvancedPS.reset_logweights!(pc)
        @test pc.logWs === logws
        @test all(iszero, pc.logWs)
    end

    @testset "trace" begin
        n = Ref(0)
        function f2(rng)
            t = TArray(Int, 1)
            t[1] = 0
            while true
                n[] += 1
                produce(t[1])
                n[] += 1
                t[1] = 1 + t[1]
            end
        end

        # Test task copy version of trace
        tr = AdvancedPS.Trace(f2, AdvancedPS.TracedRNG())

        consume(tr.ctask)
        consume(tr.ctask)

        a = AdvancedPS.fork(tr)
        consume(a.ctask)
        consume(a.ctask)

        @test consume(tr.ctask) == 2
        @test consume(a.ctask) == 4
    end
end


text:
#
# @description Unit tests for Blowfish block-cipher.
#
# @author Alejandro Claro (alejandro.claro@gmail.com)
#
# Copyright 2017 All rights reserved.
# Use of this source code is governed by a MIT-style license that can be found in the LICENSE file.
#
@testset "Blowfish cipher tests" begin
  cipher = BlowfishCipher([ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 ])

  @test block_size(cipher) == 8
  @test key_size(cipher)   == 8

  @test_throws ArgumentError BlowfishCipher(UInt8[])
  @test_throws ArgumentError encrypt(cipher, UInt8[])
  @test_throws ArgumentError encrypt(cipher, [ 0x05, 0x06 ])
  @test_throws ArgumentError decrypt(cipher, UInt8[])
  @test_throws ArgumentError decrypt(cipher, [ 0x05, 0x06 ])

  @test encrypt(cipher, [ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 ]) == UInt8[ 0x4E, 0xF9, 0x97, 0x45, 0x61, 0x98, 0xDD, 0x78 ]
  @test decrypt(cipher, UInt8[ 0x4E, 0xF9, 0x97, 0x45, 0x61, 0x98, 0xDD, 0x78 ]) == [ 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 ]
end