#pragma once
#include "MpmSimulationBase.h"
#include <Ziran/Math/Geometry/SourceCollisionObject.h>

#include <math.h>

#include <math.h>

#include <tbb/concurrent_unordered_map.h>
#include <tbb/tbb.h>

#include <Eigen/Core>
#include <Eigen/Dense>
#undef B2

#include <Ziran/CS/Util/Timer.h>
#include <Ziran/Math/Geometry/Particles.h>
#include <Ziran/Math/Geometry/AnalyticLevelSet.h>
#include <Ziran/Math/Geometry/CollisionObject.h>
#include <Ziran/Math/Geometry/PartioIO.h>
#include <Ziran/Math/Splines/BSplines.h>
#include <Ziran/Physics/LagrangianForce/Inertia.h>
#include <Ziran/Physics/LagrangianForce/LagrangianForce.h>
#include <Ziran/Sim/DiffTest.h>
#include <Ziran/Sim/Scene.h>

#include <MPM/Force/MpmForceBase.h>
#include "MpmSimulationDataAnalysis.h"

namespace ZIRAN {

template <class T, int dim>
MpmSimulationBase<T, dim>::MpmSimulationBase()
    : transfer_scheme(APIC_blend_RPIC)
    , apic_rpic_ratio((T)1)
    , flip_pic_ratio((T)0)
    , flip_rpic_ratio((T)0)
    , print_stats(false)
    , write_partio(true)
    , write_meshes(true)
    , use_particle_collision(false)
    , autorestart(true)
    , gravity(TV::Unit(1) * (0))
    , cfl((T).6)
    , dx(1)
    , dt((T)1.0 / 24)
    , element_partitions(0)
    , particles(scene.particles)
    , element_measure(particles.add(element_measure_name<T>()))
    , objective(*this)
    , newton(objective, (T)1e-5, 3)
    , force((forces.emplace_back(std::make_unique<MpmForceBase<T, dim>>(mls_mpm, dx, dt, particles, scene, scratch_xp, scratch_vp, scratch_fp, scratch_gradV, scratch_stress, dv, vn, grid, particle_base_offset, particle_order, particle_group, block_offset, num_nodes, plasticity_appliers, full_implicit)), forces.back()))
{
    symplectic = true;
    quasistatic = false;
    openvdb::initialize();
    inertia = std::make_unique<MassLumpedInertia<T, dim>>(mass_matrix, dv, dt);
    explicit_velocity_field = nullptr;
    ignoreCollisionObject = false;
    mls_mpm = false;
    useTrialCollision = true; //default to true
}

template <class T, int dim>
MpmSimulationBase<T, dim>::~MpmSimulationBase()
{
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::setDx(const T new_dx)
{
    dx = new_dx;
}

template <class T, int dim>
auto MpmSimulationBase<T, dim>::getMpmForce() -> MpmForceBase<T, dim>*
{
    return force.get();
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::addCollisionObject(AnalyticCollisionObject<T, dim>&& object)
{
    collision_objects.emplace_back(std::make_unique<AnalyticCollisionObject<T, dim>>(std::move(object)));
}

template <class T, int dim>
int MpmSimulationBase<T, dim>::addSourceCollisionObject(SourceCollisionObject<T, dim>&& object)
{
    collision_objects.emplace_back(std::make_unique<SourceCollisionObject<T, dim>>(std::move(object)));
    return collision_objects.size() - 1;
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::initialize()
{
    ZIRAN_TIMER();
    Base::initialize();
    ZIRAN_INFO("Initialize called");
    ZIRAN_INFO("end_frame = ", end_frame);

    if (mls_mpm) {
        ZIRAN_ASSERT(transfer_scheme == APIC_blend_RPIC && interpolation_degree != 1 && symplectic == true,
            "mls_mpm requires apic quadratic/cubic and only supports explicit for now");
        ZIRAN_INFO("Using mls_mpm.");
    }

    if ((transfer_scheme == APIC_blend_RPIC)
        && interpolation_degree != 1) {

        // use datamanager to store C
        particles.add(C_name<TM>());
        // compute D inverse
        T dx2 = dx * dx;
        if constexpr (interpolation_degree == 2)
            D_inverse = 4 / dx2;
        else if (interpolation_degree == 3)
            D_inverse = 3 / dx2;
        else
            ZIRAN_ASSERT(false);

        ZIRAN_ASSERT(apic_rpic_ratio >= 0 && apic_rpic_ratio <= 1);

        // resize C array
        const DisjointRanges& x_range = particles.X.ranges;
        particles.DataManager::get(C_name<TM>()).lazyResize(x_range, TM::Zero());
    }

    if (!symplectic) {
        objective.initialize(
            [&](TVStack& dv) {
                for (auto iter = collision_nodes.begin(); iter != collision_nodes.end(); ++iter) {
                    int node_id = (*iter).node_id;
                    (*iter).project(dv.col(node_id));
                }
            });
    }

    for (auto& em : scene.element_managers) {
        em->registerParticleReorderingCallback(particles);
    }

    writeSimulationInformation();

    ZIRAN_INFO("Registering restart callback for collider state update.");
    restart_callbacks.emplace_back([this](int frame) {
        for (size_t k = 0; k < collision_objects.size(); ++k) {
            if (collision_objects[k]->updateState) {
                collision_objects[k]->updateState(step.time, *collision_objects[k]);
            }
        }
    });
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::reinitialize()
{
    ZIRAN_TIMER();
    Base::reinitialize();

    if (!restarting) {
        if (element_partitions) {
            int count = 0;
            bool repartitioned = false;
            for (auto& em : scene.element_managers) {
                if (em && em->needRepartitioning()) {
                    repartitioned = true;
                    em->partitionAndReindexElements(element_partitions);

                    if (count++ == 0) {

                        StdVector<int> old_to_new;
                        StdVector<int> new_to_old;
                        if (em->renumberParticles(particles.count, old_to_new, new_to_old)) {
                            particles.reorder(new_to_old);
                        }
                    }
                }
            }

            if (repartitioned) {

                scene.segmesh_to_write.indices.clear();
                AttributeName<int> no_write("no_write");

                for (auto& em : scene.element_managers) {

                    // TODO: rebuild boundary for tet mesh

                    if (SimplexElements<T, 2, dim>* tris = dynamic_cast<SimplexElements<T, 2, dim>*>(em.get()))
                        scene.trimesh_to_write.indices = tris->indices.array;

                    if (SimplexElements<T, 1, dim>* segs = dynamic_cast<SimplexElements<T, 1, dim>*>(em.get())) {
                        if (!segs->exist(no_write)) {
                            ZIRAN_INFO("segmesh_to_write did not find no_write tag");
                            scene.segmesh_to_write.indices = segs->indices.array;
                        }
                        else {
                            ZIRAN_INFO("segmesh_to_write found no_write tag");
                            DataArray<int>& no_write_array = segs->get(no_write);
                            for (size_t z = 0; z < segs->indices.array.size(); ++z) {
                                if (no_write_array.valueId(z) == -1) // that means z is not labled no write
                                    scene.segmesh_to_write.indices.push_back(segs->indices.array[z]);
                            }
                        }
                    }
                }
            }
        }
    }

    // TODO: Fix the restarting flag
    ZIRAN_INFO("calling mpmsimulationbase reinitialize");
    ZIRAN_INFO("collision object refresh to time = ", step.time);
    ZIRAN_INFO("restarting = ", restarting); // TODO: SimulationBase already sets this to false (see simulate()).
    for (size_t k = 0; k < collision_objects.size(); ++k) {
        if (step.time == 0 || restarting) {
            if (collision_objects[k]->updateState) {
                // ZIRAN_INFO("Time=0 or restarting, in reinitialize, updating collision object to time ", step.time);
                collision_objects[k]->updateState(step.time, *collision_objects[k]);
            }
        }
        else {
            if (collision_objects[k]->updateState) {
                // ZIRAN_INFO("Calling in reinitialize anyway, to update collision object to time ", step.time);
                collision_objects[k]->updateState(step.time, *collision_objects[k]);
            }
        }
    }

    collision_nodes.clear();

    force->reinitialize();
    inertia->reinitialize();

    sortParticlesAndPolluteGrid();

    // TV
    scratch_xp.resize(particles.count, TV::Constant(0));
    scratch_vp.resize(dim, particles.count);
    scratch_fp.resize(dim, particles.count);
    // TM
    scratch_gradV.resize(particles.count, TM::Constant(0));
    scratch_stress.resize(particles.count, TM::Constant(0));
    // resize C array
    if ((transfer_scheme == APIC_blend_RPIC)
        && interpolation_degree != 1) {
        const DisjointRanges& x_range = particles.X.ranges;
        particles.DataManager::get(C_name<TM>()).lazyResize(x_range, TM::Zero());
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::beginTimeStep(int frame, int substep, double time, double dt)
{
    this->dt = dt;
    for (auto f : begin_time_step_callbacks) {
        f(frame, substep, time, dt);
    }
}

template <class T, int dim>
T MpmSimulationBase<T, dim>::totalEnergy(std::string info)
{
    T pe = 0;
    for (auto& lf : forces) {
        pe += lf->totalEnergy();
    }

    MpmSimulationDataAnalysis<T, dim> mpm_data(*this);
    T ke = mpm_data.evalTotalKineticEnergyGrid();

    ZIRAN_INFO(info, pe, " ", ke, " ", ke + pe);
    return ke + pe;
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::advanceOneTimeStep(double dt)
{
    ZIRAN_TIMER();
    ZIRAN_INFO("Advance one time step from time ", std::setw(7), step.time, " with                     dt = ", dt);

    reinitialize();

    if (symplectic) {
        forcesUpdateParticleState();
        for (auto f : before_p2g_callbacks) {
            f();
        }
        particlesToGrid();
        for (auto f : before_euler_callbacks) {
            f();
        }
        gridVelocityExplicitUpdate(dt);
        for (size_t k = 0; k < collision_objects.size(); ++k)
            if (collision_objects[k]->updateState)
                collision_objects[k]->updateState(step.time + dt, *collision_objects[k]);
        if (useTrialCollision) {
            trialCollision();
        }
        gridToParticles(dt);
    }
    else {
        for (auto f : before_p2g_callbacks) {
            f();
        }
        particlesToGrid();
        for (auto f : before_euler_callbacks) {
            f();
        }
        backwardEulerStep();
        ZIRAN_INFO("After BE (but before G2P), updating collision object to time ", step.time + dt);
        for (size_t k = 0; k < collision_objects.size(); ++k)
            if (collision_objects[k]->updateState) {
                collision_objects[k]->updateState(step.time + dt, *collision_objects[k]);
            }
        gridToParticles(dt);
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::trialCollision()
{
    ZIRAN_TIMER();

    ZIRAN_ASSERT(symplectic);
    grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
        if (explicit_collision_nodes[g.idx])
            return;
        TV vi = g.new_v;
        TV xi = node.template cast<T>() * dx;
        AnalyticCollisionObject<T, dim>::trialCollision(collision_objects, xi, dt, vi);
        g.new_v = vi;
    });
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::buildAndPassGradVnToMpmForceHelpers()
{
    bool need_to_compute = false;
    for (auto& h : force->helpers) {
        if (h->needGradVn()) {
            need_to_compute = true;
            break;
        }
    }
    if (need_to_compute) {
        buildGradVn(scratch_gradV);
        for (auto& h : force->helpers)
            if (h->needGradVn())
                h->getGradVn(scratch_gradV);
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::buildGradVn(StdVector<TM>& grad_vn)
{
    // TODO: use real rasterized grad_vn instead of grad_v from previous time step!
    auto& Xarray = particles.X.array;
    for (int i = 0; i < particles.count; ++i) {
        TM& gradVp = grad_vn[i];
        if (gradVp != gradVp) //check for NAN/unitialized data
            continue;
        gradVp = TM::Zero();
        TV Xp = Xarray[i];
        BSplineWeights<T, dim> spline(Xp, dx);
        grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
            gradVp.noalias() += g.v * dw.transpose();
        });
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::performRpicVelocityFilteringAtBeginningOfTimeStep()
{
    ZIRAN_ASSERT(!mls_mpm);
    ZIRAN_ASSERT(FLIP_blend_PIC == transfer_scheme);
    ZIRAN_ASSERT(0 == flip_pic_ratio);
    ZIRAN_ASSERT(1 != ZIRAN_MPM_DEGREE);

    reinitialize();
    forcesUpdateParticleState();

    particlesToGridWithForceHelper<false, false, true>();
    num_nodes = grid.getNumNodes();
    grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
        if (g.m != 0) {
            g.v /= g.m;
        }
        else {
            g.v = TV::Zero();
        }
    });

    grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
        g.new_v = g.v;
    });

    auto* C = ((transfer_scheme == APIC_blend_RPIC) && interpolation_degree != 1)
        ? &particles.DataManager::get(C_name<TM>())
        : nullptr;
    static constexpr bool USE_APIC_BLEND_RPIC = false;
    static constexpr bool USE_MPM_DEGREE_ONE = false;
    static constexpr bool USE_MLS_MPM = false;
    {
        tbb::parallel_for(0, (int)particle_group.size(), [&](int group_idx) {
            for (int idx = particle_group[group_idx].first; idx <= particle_group[group_idx].second; ++idx) {

                int i = particle_order[idx];
                TV& Xp = particles.X[i];

                TV picV = TV::Zero();
                BSplineWeights<T, dim> spline(Xp, dx);
                if constexpr (!USE_APIC_BLEND_RPIC) {
                    TV oldV = TV::Zero();
                    TM& gradVp = scratch_gradV[i];
                    gradVp = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        picV += w * g.new_v;
                        oldV += w * g.v;
                        gradVp.noalias() += g.new_v * dw.transpose();
                    });
                    particles.V[i] *= flip_pic_ratio;
                    particles.V[i] += picV - flip_pic_ratio * oldV;
                }
                else if constexpr (!USE_MPM_DEGREE_ONE) {
                    TM Bp = TM::Zero();
                    TM& gradVp = scratch_gradV[i];
                    gradVp = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        picV += w * g.new_v;
                        TV xi_minus_xp = node.template cast<T>() * dx - Xp;
                        Bp.noalias() += w * g.new_v * xi_minus_xp.transpose();
                        // !mls_mpm
                        if constexpr (!USE_MLS_MPM) {
                            gradVp.noalias() += g.new_v * dw.transpose();
                        }
                    });
                    particles.V[i] = picV;
                    TM CC = Bp * D_inverse;
                    (*C)[i] = ((apic_rpic_ratio + 1) * (T)0.5) * CC + ((apic_rpic_ratio - 1) * (T)0.5) * CC.transpose();
                    // no check for NAN/unitialized data
                    if constexpr (USE_MLS_MPM) {
                        gradVp = (*C)[i];
                    }
                }
                else {
                    TM& gradVp = scratch_gradV[i];
                    gradVp = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        picV += w * g.new_v;
                        gradVp.noalias() += g.new_v * dw.transpose();
                    });
                    particles.V[i] = picV;
                }
            }
        });
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::particlesToGrid()
{
#if 0
    {
        ZIRAN_TIMER();
        ZIRAN_INFO("particle group size ", particle_group.size());
        for (uint64_t color = 0; color < (1 << dim); ++color) {
            tbb::parallel_for(0, (int)particle_group.size(), [&](int group_idx) {
                if ((block_offset[group_idx] & ((1 << dim) - 1)) != color)
                    return;
                for (int idx = particle_group[group_idx].first; idx <= particle_group[group_idx].second; ++idx) {
                    int i = particle_order[idx];
                    TV& Xp = particles.X.array[i];
                    T mass = particles.mass.array[i];
                    BSplineWeights<T, dim> spline(Xp, dx);
                    grid.iterateKernel(spline, particle_base_offset[i], [&](const IV& node, T w, const TV& dw, GridState<T, dim>& g) {
                        g.m += w * mass;
                    });
                }
            });
        }
        ZIRAN_WARN("Just finish normal grid mass p2g");
        //getchar();
    }
#endif
    if (symplectic) {
        ZIRAN_TIMER();
        if (mls_mpm) {
            if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE == 1)
                particlesToGridWithForceHelper<false, true, true>();
            else if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE != 1)
                particlesToGridWithForceHelper<false, false, true>();
            else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE == 1)
                particlesToGridWithForceHelper<true, true, true>();
            else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE != 1)
                particlesToGridWithForceHelper<true, false, true>();
        }
        else {
            if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE == 1)
                particlesToGridWithForceHelper<false, true, false>();
            else if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE != 1)
                particlesToGridWithForceHelper<false, false, false>();
            else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE == 1)
                particlesToGridWithForceHelper<true, true, false>();
            else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE != 1)
                particlesToGridWithForceHelper<true, false, false>();
        }
    }
    else {
        ZIRAN_TIMER();
        if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE == 1)
            particlesToGridHelper<false, true>();
        else if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE != 1)
            particlesToGridHelper<false, false>();
        else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE == 1)
            particlesToGridHelper<true, true>();
        else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE != 1)
            particlesToGridHelper<true, false>();
    }
    {
        ZIRAN_TIMER();
        num_nodes = grid.getNumNodes();
        grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
            if (g.m != 0) {
                g.v /= g.m;
            }
            else {
                g.v = TV::Zero();
            }
        });
    }
}

template <class T, int dim>
template <bool USE_APIC_BLEND_RPIC, bool USE_MPM_DEGREE_ONE, bool USE_MLS_MPM>
void MpmSimulationBase<T, dim>::particlesToGridWithForceHelper()
{
    auto& Xarray = particles.X.array;
    auto& Varray = particles.V.array;
    auto& marray = particles.mass.array;
    const StdVector<Matrix<T, dim, dim>>* Carray_pointer;
    Carray_pointer = ((transfer_scheme == APIC_blend_RPIC) && interpolation_degree != 1) ? (&(particles.DataManager::get(C_name<TM>()).array)) : NULL;

    {
        for (uint64_t color = 0; color < (1 << dim); ++color) {
            tbb::parallel_for(0, (int)particle_group.size(), [&](int group_idx) {
                if ((block_offset[group_idx] & ((1 << dim) - 1)) != color)
                    return;
                for (int idx = particle_group[group_idx].first; idx <= particle_group[group_idx].second; ++idx) {
                    int i = particle_order[idx];
                    TV& Xp = Xarray[i];
                    T mass = marray[i];
                    TV momentum = marray[i] * Varray[i];

                    TM C = TM::Zero();
                    if constexpr (USE_APIC_BLEND_RPIC) {
                        if constexpr (!USE_MPM_DEGREE_ONE) {
                            C = mass * (*Carray_pointer)[i];
                        }
                        else {
                            const TM& gradVp = scratch_gradV[i];
                            C = mass * (((apic_rpic_ratio + 1) * (T)0.5) * gradVp + ((apic_rpic_ratio - 1) * (T)0.5) * gradVp.transpose());
                        }
                    }
                    TM4 stress_density = TM4::Zero();
                    TM& stress = scratch_stress[i];
                    stress_density.template topLeftCorner<dim, dim>() = stress; // stress density
                    const TV& fp = scratch_fp.col(i);
                    stress_density.template topRightCorner<dim, 1>() = -fp; // fp (for meshed forces).
                    TM4 velocity_density = TM4::Zero();
                    velocity_density.template topLeftCorner<dim, dim>() = C;
                    velocity_density.template topRightCorner<dim, 1>() = momentum;
                    velocity_density(3, 3) = mass;
                    BSplineWeights<T, dim> spline(Xp, dx);
                    if constexpr (USE_MLS_MPM) {
                        grid.iterateKernel(spline, particle_base_offset[i],
                            [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                                TV4 xi_minus_xp = TV4::Zero();
                                xi_minus_xp.template topLeftCorner<dim, 1>() = node.template cast<T>() * dx - Xp; // top dim entries non-zero
                                xi_minus_xp(3) = 1;
                                TV4 force_delta = ((stress_density * xi_minus_xp) * (w * D_inverse));
                                TV4 velocity_delta = velocity_density * xi_minus_xp * w;
                                g.m += velocity_delta(3);
                                g.v += velocity_delta.template topLeftCorner<dim, 1>();
                                g.new_v -= force_delta.template topLeftCorner<dim, 1>(); // fi -= \sum_p (Ap (xi-xp)  - fp )w_ip Dp_inv
                            });
                    }
                    else {
                        grid.iterateKernel(spline, particle_base_offset[i],
                            [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                                TV4 weight = TV4::Zero();
                                weight.template topLeftCorner<dim, 1>() = dw;
                                weight(3) = w;
                                TV4 xi_minus_xp = TV4::Zero();
                                xi_minus_xp.template topLeftCorner<dim, 1>() = node.template cast<T>() * dx - Xp; // top dim entries non-zero
                                xi_minus_xp(3) = 1;
                                TV4 force_delta = (stress_density * weight);
                                TV4 velocity_delta = velocity_density * xi_minus_xp * w;
                                g.m += velocity_delta(3);
                                g.v += velocity_delta.template topLeftCorner<dim, 1>();
                                g.new_v -= force_delta.template topLeftCorner<dim, 1>();
                            });
                    }
                }
            });
        }
    }
}

template <class T, int dim>
template <bool USE_APIC_BLEND_RPIC, bool USE_MPM_DEGREE_ONE>
void MpmSimulationBase<T, dim>::particlesToGridHelper()
{
    auto& Xarray = particles.X.array;
    auto& Varray = particles.V.array;
    auto& marray = particles.mass.array;
    const StdVector<Matrix<T, dim, dim>>* Carray_pointer;
    Carray_pointer = ((transfer_scheme == APIC_blend_RPIC) && interpolation_degree != 1) ? (&(particles.DataManager::get(C_name<TM>()).array)) : NULL;

    for (uint64_t color = 0; color < (1 << dim); ++color) {
        tbb::parallel_for(0, (int)particle_group.size(), [&](int group_idx) {
            if ((block_offset[group_idx] & ((1 << dim) - 1)) != color)
                return;
            for (int idx = particle_group[group_idx].first; idx <= particle_group[group_idx].second; ++idx) {
                int i = particle_order[idx];
                TV& Xp = Xarray[i];
                T mass = marray[i];
                TV momentum = marray[i] * Varray[i];
                TM C = TM::Zero();
                if constexpr (USE_APIC_BLEND_RPIC) {
                    if constexpr (!USE_MPM_DEGREE_ONE) {
                        C = mass * (*Carray_pointer)[i];
                    }
                    else {
                        const TM& gradVp = scratch_gradV[i];
                        C = mass * (((apic_rpic_ratio + 1) * (T)0.5) * gradVp + ((apic_rpic_ratio - 1) * (T)0.5) * gradVp.transpose());
                    }
                }
                TM4 velocity_density = TM4::Zero();
                velocity_density.template topLeftCorner<dim, dim>() = C;
                velocity_density.template topRightCorner<dim, 1>() = momentum;
                velocity_density(3, 3) = mass;
                BSplineWeights<T, dim> spline(Xp, dx);
                grid.iterateKernel(spline, particle_base_offset[i], [&](const IV& node, T w, const TV& dw, GridState<T, dim>& g) {
                    TV4 xi_minus_xp = TV4::Zero();
                    xi_minus_xp.template topLeftCorner<dim, 1>() = node.template cast<T>() * dx - Xp; // top dim entries non-zero
                    xi_minus_xp(3) = 1;
                    TV4 velocity_delta = velocity_density * xi_minus_xp * w;
                    g.m += velocity_delta(3);
                    g.v += velocity_delta.template topLeftCorner<dim, 1>();
                });
            }
        });
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::startBackwardEuler()
{
    ZIRAN_TIMER();
    ZIRAN_ASSERT(objective.matrix_free, "MPM only supports matrix free");

    buildMassMatrix();
    objective.setPreconditioner(([&](const TVStack& in, TVStack& out) {
        for (int i = 0; i < num_nodes; i++) {
            for (int d = 0; d < dim; d++) {
                out(d, i) = in(d, i) / mass_matrix(i);
            }
        }
    }));

    dv.resize(dim, num_nodes);
    vn.resize(dim, num_nodes);

    buildInitialDvAndVnForNewton(); // This also builds collision_nodes
    // TODO: this should probably be done in objective reinitialize.
    // Which should be called at the beginning of newton.
    force->backupStrain();
    objective.reinitialize(); // Reinitialize matrix sparsity pattern
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::backwardEulerStep()
{
    ZIRAN_TIMER();
    startBackwardEuler();
    for (auto f : general_callbacks) {
        f(frame, BeforeNewtonSolve);
    }
    // run diff_test if we specify it
    if (diff_test)
        runDiffTest<T, dim, Objective>(num_nodes, dv, objective, diff_test_perturbation_scale);
    newton.solve(dv, verbose);
    for (auto f : general_callbacks) {
        f(frame, AfterNewtonSolve);
    }
    // run diff_test if we specify it
    if (diff_test)
        runDiffTest<T, dim, Objective>(num_nodes, dv, objective, diff_test_perturbation_scale);

    force->restoreStrain();

    constructNewVelocityFromNewtonResult();
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::forcesUpdateParticleState()
{
    ZIRAN_TIMER();

    // zero out scratch_fp and scratch_stress, prepare for adding lagrangian force
    tbb::parallel_for(tbb::blocked_range<size_t>(0, particles.count),
        [&](const tbb::blocked_range<size_t>& range) {
                          for (size_t b = range.begin(), b_end = range.end(); b < b_end; ++b) {
                              scratch_fp.col(b) = TV::Zero();
                              scratch_stress[b] = TM::Zero(); } });
    // Lagrangian
    for (auto& lf : scene.forces)
        lf->updatePositionBasedState();

    auto& vtau = scratch_stress;
    auto& fp = scratch_fp;
    auto ranges = particles.X.ranges;
    tbb::parallel_for(ranges,
        [&](DisjointRanges& subrange) {
            for (auto& h : force->helpers)
                h->updateState(subrange, vtau, fp); // this adds to vtau and fp
        });

    // Lagrangian
    scene.addScaledForces(1, fp); // add forces to fp
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::moveNodes(const TVStack& dv_in)
{
    ZIRAN_QUIET_TIMER();
    if (dv.data() == dv_in.data())
        return;
    tbb::parallel_for(tbb::blocked_range<size_t>(0, num_nodes),
        [&](const tbb::blocked_range<size_t>& range) {
            for (size_t i = range.begin(), i_end = range.end(); i < i_end; ++i) {
                dv.col(i) = dv_in.col(i);
            }
        });
}

/**
       Callback to write the simulation state

       Should be overridden to write the simulation state
    */
template <class T, int dim>
void MpmSimulationBase<T, dim>::writeState(std::ostream& out)
{
    ZIRAN_TIMER();
    ZIRAN_INFO("Write state called");

    if (write_partio) {
        std::string obj_filename = output_dir.absolutePath(outputFileName("partio", ".bgeo"));
        writePartio(obj_filename, particles);
    }

    scene.writeState(out, outputFileName("boundary", ".obj"), outputFileName("segmesh", ".poly"), output_dir, binary_ver, write_meshes);

    if (transfer_scheme == APIC_blend_RPIC && interpolation_degree == 1)
        writeSTDVector(out, scratch_gradV); // apic C matrix
}

/**
       Callback to read the simulation state

       Should be overridden to read the output of writeState
    */
template <class T, int dim>
void MpmSimulationBase<T, dim>::readState(std::istream& in)
{
    ZIRAN_QUIET_TIMER();
    ZIRAN_INFO("Read state called");
    scene.readState(in, binary_ver);

    if (transfer_scheme == APIC_blend_RPIC && interpolation_degree == 1)
        readSTDVector(in, scratch_gradV); // apic C matrix
}

/**
       Callback to calculate dt based on CFL
    */
template <class T, int dim>
double MpmSimulationBase<T, dim>::calculateDt()
{
    ZIRAN_QUIET_TIMER();
    TV p_min_corner, p_max_corner;
    Eigen::Array<T, 2, 1> max_speed = evalMaxParticleSpeed(p_min_corner, p_max_corner);
    ZIRAN_INFO("Particle Min Corner: ", p_min_corner.transpose());
    ZIRAN_INFO("Particle Max Corner: ", p_max_corner.transpose());
    p_min_corner.array() -= (interpolation_degree + 2) * dx; // expand by influenced region and buffer
    p_max_corner.array() += (interpolation_degree + 2) * dx;
    ZIRAN_INFO("Maximum particle speed: ", max_speed(0));
    ZIRAN_ASSERT(max_speed(0) == max_speed(1));
    T max_collision_object_speed = 0;
    for (size_t k = 0; k < collision_objects.size(); ++k)
        max_collision_object_speed = std::max(collision_objects[k]->evalMaxSpeed(p_min_corner, p_max_corner), max_collision_object_speed);
    ZIRAN_INFO("Maximum collision object speed: ", max_collision_object_speed);
    max_speed(1) = std::max(max_speed(1), max_collision_object_speed);
    double dt_compute = step.max_dt;
    if (max_speed(1))
        dt_compute = cfl * dx / max_speed(1);
    ZIRAN_INFO("Computed next dt (based on CFL): ", dt_compute);

    //if (step.time == 0) return std::min(dt_compute, (double)1e-4);
    return dt_compute;
}

// Build and mass_matrix (for implicit)
template <class T, int dim>
void MpmSimulationBase<T, dim>::buildMassMatrix()
{
    ZIRAN_QUIET_TIMER();
    mass_matrix.resize(num_nodes, 1);

    grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
        mass_matrix(g.idx) = g.m;
    });
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::addScaledForces(const T scale, TVStack& f) // called BE Newton objective (for computing residual)
{
    for (auto& lf : forces)
        lf->addScaledForces(scale, f);
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::addScaledForceDifferentials(const T scale, const TVStack& x, TVStack& f) // called BE Newton objective (for computing residual)
{
    for (auto& lf : forces)
        lf->addScaledForceDifferential(scale, x, f);
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::gridVelocityExplicitUpdate(double dt)
{
    ZIRAN_TIMER();
    assert(symplectic);

    explicit_collision_nodes.resize((int)num_nodes, 0);

    TV dv_gravity = gravity * dt;
    if (explicit_velocity_field) {
        grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
            T mi = g.m;
            TV xi = node.template cast<T>() * dx;
            TV vi = g.v;
            explicit_velocity_field(mi, xi, vi);
            g.new_v = vi;
        });
    }
    else {
        grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
            T mi = g.m;
            T dt_over_mi = (T)dt / mi;
            TV xi = node.template cast<T>() * dx;
            TV old_v = g.v;
            TV force_value = g.new_v;
            TV dvijk = force_value * dt_over_mi;
            TV vi = old_v + dv_gravity + dvijk;
            TM normal_basis;

            // apply external body force field: fext(x,t)
            for (auto ff : fext) {
                TV external_force = ff(xi, (T)(step.time));
                vi += dt * external_force;
            }
            auto old_vi = vi;

            bool collided = AnalyticCollisionObject<T, dim>::
                multiObjectCollision(collision_objects, xi, vi, normal_basis);

            explicit_collision_nodes[g.idx] = (int)collided;

            if (ignoreCollisionObject)
                vi = old_vi;

            g.new_v = vi;
        });
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::constructNewVelocityFromNewtonResult()
{
    ZIRAN_QUIET_TIMER();
    grid.iterateTouchedGrid([&](IV node, GridState<T, dim>& g) {
        g.new_v = TV::Zero();
    });
    grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
        g.new_v = g.v + dv.col(g.idx);
    });
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::gridToParticles(double dt)
{
    if (mls_mpm) {
        ZIRAN_TIMER();
        if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE == 1)
            gridToParticlesHelper<false, true, true>(dt);
        else if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE != 1)
            gridToParticlesHelper<false, false, true>(dt);
        else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE == 1)
            gridToParticlesHelper<true, true, true>(dt);
        else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE != 1)
            gridToParticlesHelper<true, false, true>(dt);
    }
    else {
        ZIRAN_TIMER();
        if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE == 1)
            gridToParticlesHelper<false, true, false>(dt);
        else if (transfer_scheme == FLIP_blend_PIC && ZIRAN_MPM_DEGREE != 1)
            gridToParticlesHelper<false, false, false>(dt);
        else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE == 1)
            gridToParticlesHelper<true, true, false>(dt);
        else if (transfer_scheme == APIC_blend_RPIC && ZIRAN_MPM_DEGREE != 1)
            gridToParticlesHelper<true, false, false>(dt);
    }
}

template <class T, int dim>
template <bool USE_APIC_BLEND_RPIC, bool USE_MPM_DEGREE_ONE, bool USE_MLS_MPM>
void MpmSimulationBase<T, dim>::gridToParticlesHelper(double dt)
{
    std::atomic<bool> faster_than_grid_cell(false);
    std::atomic<bool> faster_than_half_grid_cell(false);
    auto* C = ((transfer_scheme == APIC_blend_RPIC) && interpolation_degree != 1)
        ? &particles.DataManager::get(C_name<TM>())
        : nullptr;

    //    if constexpr (false) {
    {
        tbb::parallel_for(0, (int)particle_group.size(), [&](int group_idx) {
            for (int idx = particle_group[group_idx].first; idx <= particle_group[group_idx].second; ++idx) {
                bool local_faster_than_grid_cell = false;
                bool local_faster_than_half_grid_cell = false;

                int i = particle_order[idx];
                TV& Xp = particles.X[i];

                TV picV = TV::Zero();
                BSplineWeights<T, dim> spline(Xp, dx);
                if constexpr (!USE_APIC_BLEND_RPIC) {
                    TV oldV = TV::Zero();
                    TM& gradVp = scratch_gradV[i];
                    gradVp = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        picV += w * g.new_v;
                        oldV += w * g.v;
                        gradVp.noalias() += g.new_v * dw.transpose();
                    });
                    particles.V[i] *= flip_pic_ratio;
                    particles.V[i] += picV - flip_pic_ratio * oldV;
                }
                else if constexpr (!USE_MPM_DEGREE_ONE) {
                    TM Bp = TM::Zero();
                    TM& gradVp = scratch_gradV[i];
                    gradVp = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        picV += w * g.new_v;
                        TV xi_minus_xp = node.template cast<T>() * dx - Xp;
                        Bp.noalias() += w * g.new_v * xi_minus_xp.transpose();
                        // !mls_mpm
                        if constexpr (!USE_MLS_MPM) {
                            gradVp.noalias() += g.new_v * dw.transpose();
                        }
                    });
                    particles.V[i] = picV;
                    TM CC = Bp * D_inverse;
                    (*C)[i] = ((apic_rpic_ratio + 1) * (T)0.5) * CC + ((apic_rpic_ratio - 1) * (T)0.5) * CC.transpose();
                    // no check for NAN/unitialized data
                    if constexpr (USE_MLS_MPM) {
                        gradVp = (*C)[i];
                    }
                }
                else {
                    TM& gradVp = scratch_gradV[i];
                    gradVp = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        picV += w * g.new_v;
                        gradVp.noalias() += g.new_v * dw.transpose();
                    });
                    particles.V[i] = picV;
                }

                TV increment = dt * picV;
                particles.X[i] += increment; // particle position update is the same for all transfer schemes
                T inc = increment.squaredNorm();
                T dx2 = dx * dx;
                local_faster_than_grid_cell = local_faster_than_grid_cell + (inc > dx2);
                local_faster_than_half_grid_cell = local_faster_than_half_grid_cell + (inc > dx2 * (T)0.25 * (cfl * cfl));
                if (local_faster_than_half_grid_cell)
                    faster_than_half_grid_cell.store(true, std::memory_order_relaxed);
                if (local_faster_than_grid_cell)
                    faster_than_grid_cell.store(true, std::memory_order_relaxed);
            }
        });
    }

    static int latest_frame_restarted = 0;
    if (faster_than_grid_cell && autorestart) {
        ZIRAN_WARN("Particle traveling more than a grid cell detected");
        int frame_to_restart = frame - 1;
        if (step.max_dt <= step.min_dt) {
            ZIRAN_WARN("Unable to shrink dt further");
            frame_to_restart--;
            ZIRAN_ASSERT(frame_to_restart >= start_frame, "Unstable intial conditions detected, shrink min_dt or change other parameters");
        }
        else {
            double new_max_dt = std::max(step.max_dt / 2, step.min_dt);
            ZIRAN_WARN("Shrinking max dt to ", new_max_dt);

            restart_callbacks.emplace_back([new_max_dt, this](int frame) {
                step.max_dt = new_max_dt;
            });
        }
        latest_frame_restarted = std::max(frame_to_restart, latest_frame_restarted);
        throw RestartException(frame_to_restart);
    }

    if (!faster_than_half_grid_cell && frame > latest_frame_restarted + 1) {
        double new_max_dt = std::min(step.max_dt_original, step.max_dt * (double)1);
        if (new_max_dt != step.max_dt) {
            ZIRAN_WARN("All particles traveled less than half a grid cell");
            step.max_dt = new_max_dt;
            ZIRAN_WARN("Increasing max dt to ", step.max_dt);
        }
    }

    force->evolveStrain(dt);

    applyPlasticity();
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::applyPlasticity()
{
    // Split all particles to subranges.
    // In parallel, each subrange goes through all plasticity appliers in serial.
    ZIRAN_QUIET_TIMER();
    auto ranges = particles.X.ranges;
    tbb::parallel_for(ranges,
        [&](DisjointRanges& subrange) {
            for (auto& p : plasticity_appliers)
                p->applyPlasticity(subrange, particles);
        });

    if (use_particle_collision) {
        particles.parallel_for([&](const auto& X_local, auto& V_local) {
            for (size_t k = 0; k < collision_objects.size(); ++k)
                collision_objects[k]->particleCollision(dx, X_local, V_local);
        },
            particles.X_name(), particles.V_name());
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::sortParticlesAndPolluteGrid()
{
    auto& Xarray = particles.X.array;

    constexpr int index_bits = (32 - SparseMask::block_bits);
    ZIRAN_ASSERT(particles.count < (1 << index_bits));

    if (particles.count != (int)particle_sorter.size()) {
        particle_base_offset.resize(particles.count);
        particle_sorter.resize(particles.count);
        particle_order.resize(particles.count);
    }

    T one_over_dx = (T)1 / dx;
    tbb::parallel_for(0, particles.count, [&](int i) {
        uint64_t offset = SparseMask::Linear_Offset(to_std_array(
            baseNode<interpolation_degree, T, dim>(Xarray[i] * one_over_dx)));
        particle_sorter[i] = ((offset >> SparseMask::data_bits) << index_bits) + i;
    });

    tbb::parallel_sort(particle_sorter.begin(), particle_sorter.end());

    particle_group.clear();
    block_offset.clear();
    int last_index = 0;
    for (int i = 0; i < particles.count; ++i)
        if (i == particles.count - 1 || (particle_sorter[i] >> 32) != (particle_sorter[i + 1] >> 32)) {
            particle_group.push_back(std::make_pair(last_index, i));
            block_offset.push_back(particle_sorter[i] >> 32);
            last_index = i + 1;
        }

    grid.page_map->Clear();
    for (int i = 0; i < particles.count; ++i) {
        particle_order[i] = (int)(particle_sorter[i] & ((1ll << index_bits) - 1));
        uint64_t offset = (particle_sorter[i] >> index_bits) << SparseMask::data_bits;
        particle_base_offset[particle_order[i]] = offset;
        if (i == particles.count - 1 || (particle_sorter[i] >> 32) != (particle_sorter[i + 1] >> 32)) {
            grid.page_map->Set_Page(offset);
            if constexpr (dim == 2) {
                auto x = 1 << SparseMask::block_xbits;
                auto y = 1 << SparseMask::block_ybits;
                for (int i = 0; i < 2; ++i)
                    for (int j = 0; j < 2; ++j)
                        grid.page_map->Set_Page(SparseMask::Packed_Add(
                            offset, SparseMask::Linear_Offset(x * i, y * j)));
            }
            else {
                auto x = 1 << SparseMask::block_xbits;
                auto y = 1 << SparseMask::block_ybits;
                auto z = 1 << SparseMask::block_zbits;
                for (int i = 0; i < 2; ++i)
                    for (int j = 0; j < 2; ++j)
                        for (int k = 0; k < 2; ++k)
                            grid.page_map->Set_Page(SparseMask::Packed_Add(
                                offset, SparseMask::Linear_Offset(x * i, y * j, z * k)));
            }
        }
    }
    grid.page_map->Update_Block_Offsets();

    auto grid_array = grid.grid->Get_Array();
    auto blocks = grid.page_map->Get_Blocks();
    for (int b = 0; b < (int)blocks.second; ++b) {
        auto base_offset = blocks.first[b];
        std::memset(&grid_array(base_offset), 0, (size_t)(1 << MpmGrid<T, dim>::log2_page));
        GridState<T, dim>* g = reinterpret_cast<GridState<T, dim>*>(&grid_array(base_offset));
        for (int i = 0; i < (int)SparseMask::elements_per_block; ++i)
            g[i].idx = -1;
    }
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::buildInitialDvAndVnForNewton()
{
    ZIRAN_QUIET_TIMER();
    assert(!symplectic);

    grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
        TV old_v = g.v;
        TV vi = old_v;
        TV wn;

        const TV xi = node.template cast<T>() * dx;
        TM normal_basis, R;

        bool any_collision = AnalyticCollisionObject<T, dim>::
            multiObjectCollision(collision_objects, xi, vi, normal_basis, wn);

        int node_id = g.idx;
        assert(node_id < num_nodes);
        if (any_collision) {
#if 1
            bool isSlip = wn != TV::Zero();
            if (isSlip) {
                R = RotationExtractor<T, dim>::rotate(wn);
                //ZIRAN_INFO(isSlip ? "SLIP!\n" : "STICKY!\n");
                //printf("normal: %e\t%e\t%e\n\n%e\t%e\t%e\n%e\t%e\t%e\n%e\t%e\t%e\n\n", wn(0), wn(1), wn(2), R(0, 0), R(0, 1), R(0, 2), R(1, 0), R(1, 1), R(1, 2), R(2, 0), R(2, 1), R(2, 2));
                //getchar();
            }
            else
                R = TM::Identity();
            CollisionNode<T, dim> Z{ node_id, TM::Identity() - normal_basis * normal_basis.transpose(), /* wn, */ R, R.inverse(), isSlip };
#else
            CollisionNode<T, dim> Z{ node_id, TM::Identity() - normal_basis * normal_basis.transpose() };
#endif
            collision_nodes.push_back(Z);

            // Newton initial guess for collided nodes.
            // Setting it this way so that the new velocity constructed after the solve agrees with collision.
            dv.col(node_id) = vi - old_v;
        }
        else
            dv.col(node_id) = gravity * dt; // Newton initial guess for non-collided nodes

        vn.col(node_id) = old_v;
    });
}

/**
       result(0) is purely particle speed maximum.
       result(1) is APIC enhanced 'speed' maximum.
    */
template <class T, int dim>
Eigen::Array<T, 2, 1> MpmSimulationBase<T, dim>::evalMaxParticleSpeed(TV& min_corner, TV& max_corner)
{
    ZIRAN_TIMER();
    using TArray = Eigen::Array<T, 2 * dim + 2, 1>;
    TArray init = TArray::Zero();
    T Tmin = std::numeric_limits<float>::lowest();
    TV TVmin = TV::Zero();
    TVmin.array() = Tmin;
    init << (T)0, (T)0, TVmin, TVmin;
    TArray result;
    result = particles.map_reduce(
        init,
        [](const TV& position, const TV& velocity) {
            T velocity_norm = velocity.norm();
            TArray r;
            r << velocity_norm, velocity_norm, position, -position;
            return r;
        },
        [](const TArray& a, const TArray& b) {
            return a.max(b);
        },
        particles.X_name(), particles.V_name());

    max_corner = result.template segment<dim>(2);
    min_corner = -result.template tail<dim>();
    return result.template head<2>();
}

template <class T, int dim>
void MpmSimulationBase<T, dim>::writeSimulationInformation()
{
    ZIRAN_INFO("Particle count:      ", particles.count);
    ZIRAN_INFO("Grid dx:             ", dx);
    ZIRAN_INFO("Interp degree:       ", interpolation_degree);
    ZIRAN_INFO("Gravity:             ", gravity.transpose());
    ZIRAN_INFO("");
    ZIRAN_INFO("CFL:                 ", cfl);
    ZIRAN_INFO("Transfer scheme:     ", transfer_scheme);
    ZIRAN_INFO("Print stats:         ", print_stats);
    ZIRAN_INFO("Using mls_mpm:       ", mls_mpm);
    writeParticlePerCellHistogram();
}

// Called by writeSimulationInformation
template <class T, int dim>
void MpmSimulationBase<T, dim>::writeParticlePerCellHistogram()
{
    typedef tbb::concurrent_unordered_map<uint64_t, size_t> ConcurrentHash;
    ConcurrentHash cellParticleCount; // mapping IV to size_t
    AttributeName<Vector<T, dim>> x_name = Particles<T, dim>::X_name();
    for (auto iter = particles.iter(x_name); iter; ++iter) {
        Vector<T, dim>& X = iter.template get<0>();
        Vector<int, dim> IX = (X / dx).template cast<int>();
        uint64_t index = SparseMask::Linear_Offset(to_std_array(IX));
        auto search = cellParticleCount.find(index);
        if (search == cellParticleCount.end())
            cellParticleCount[index] = 1;
        else
            cellParticleCount[index]++;
    }
    StdVector<size_t> number_of_cells_with_particles(101, 0);
    for (auto& iter : cellParticleCount)
        number_of_cells_with_particles[std::min(iter.second, (size_t)100)]++;

    for (int i = 0; i < 101; i++) {
        int count = number_of_cells_with_particles[i];
        if (count > 0 && i != 100)
            ZIRAN_INFO("Cells with ", std::setw(2), i, " particles：", count);
        else if (count > 0 && i == 100)
            ZIRAN_INFO("Cells with 100 or more particles: ", count);
    }
}

template <class T, int dim>
constexpr int MpmSimulationBase<T, dim>::interpolation_degree;
} // namespace ZIRAN


// Copyright 2018, Intel Corporation

#include "tile/codegen/fuse.h"

#include <boost/format.hpp>

#include "base/util/stream_container.h"
#include "base/util/throw.h"
#include "tile/codegen/localize.h"
#include "tile/codegen/tile.h"

namespace vertexai {
namespace tile {
namespace codegen {

using namespace stripe;  // NOLINT

bool NoInnerBlock(Block *block) {
  for (const auto& stmt : block->stmts) {
    if (stmt->kind() == StmtKind::Block) {
      return false;
    }
  }
  return true;
}

static std::vector<Affine> TranslatedContraints(const AliasMap& map, std::map<std::string, std::string> remap,
                                                const Block& in) {
  std::vector<Affine> out;
  AliasMap inner = AliasMap(map, const_cast<Block*>(&in));
  std::map<std::string, stripe::Affine> remapped;
  for (const auto& kvp : inner.idx_sources()) {
    auto it = remap.find(kvp.first);
    if (it == remap.end()) {
      remapped.emplace(kvp);
    } else {
      remapped.emplace(kvp.first, Affine(it->second));
    }
  }
  for (const auto& aff : in.constraints) {
    IVLOG(4, "Remap = " << remap);
    IVLOG(4, "Translating " << aff << " to " << aff.sym_eval(remapped));
    out.push_back(aff.sym_eval(remapped));
  }
  std::sort(out.begin(), out.end());
  return out;
}

static Affine TranslateAffine(const Affine& src, const AliasMap& map, std::map<std::string, std::string> remap,
                              const Block* in) {
  Affine dest;
  AliasMap inner = AliasMap(map, const_cast<Block*>(in));
  std::map<std::string, stripe::Affine> remapped;
  for (const auto& kvp : inner.idx_sources()) {
    auto it = remap.find(kvp.first);
    if (it == remap.end()) {
      remapped.emplace(kvp);
    } else {
      remapped.emplace(kvp.first, Affine(it->second));
    }
  }
  return dest.sym_eval(remapped);
}

boost::optional<FusionPlan> ComputeFusionPlan(const AliasMap& scope, const Block& a, const Block& b,
                                              const std::string& buf_name) {
  IVLOG(3, "ComputeFusionPlan for " << buf_name << " between " << a.name << " and " << b.name);
  FusionPlan plan;
  plan.tile_a = TileShape(a.idxs.size(), 1);
  plan.a_interleave = false;
  plan.tile_b = TileShape(b.idxs.size(), 1);
  plan.b_interleave = false;
  // This is quite hueristic right now, but still beats our prior implementation
  auto it_a = a.ref_by_from(buf_name, false);
  if (it_a == a.refs.end()) {
    IVLOG(3, "ComputeFusionPlan: buffer name unknown in block a");
    return boost::none;
  }
  auto it_b = b.ref_by_from(buf_name, false);
  if (it_b == b.refs.end()) {
    IVLOG(3, "ComputeFusionPlan: buffer name unknown in block b");
    return boost::none;
  }
  assert(it_a->access.size() == it_b->access.size());
  for (size_t i = 0; i < it_a->access.size(); i++) {
    const Affine& poly_a = it_a->access[i];
    const Affine& poly_b = it_b->access[i];
    if (poly_a == 0 && poly_b == 0) {
      continue;
    }
    if (poly_a.getMap().size() != 1 || poly_a.getMap().begin()->first.empty()) {
      IVLOG(3, "ComputeFusionPlan: complex access in a: " << poly_a.toString());
      return boost::none;
    }
    if (poly_b.getMap().size() != 1 || poly_b.getMap().begin()->first.empty()) {
      IVLOG(3, "ComputeFusionPlan: complex access in b: " << poly_b.toString());
      return boost::none;
    }
    std::string idx_a = poly_a.getMap().begin()->first;
    std::string idx_b = poly_b.getMap().begin()->first;
    if (plan.remap_a.find(idx_a) != plan.remap_a.end()) {
      IVLOG(3, "ComputeFusionPlan: duplicate index");
      return boost::none;
    }
    int64_t mul_a = poly_a[idx_a];
    int64_t mul_b = poly_b[idx_b];
    if (mul_a % mul_b != 0) {
      IVLOG(3, "ComputeFusionPlan: uneven index division");
      return boost::none;
    }
    for (size_t i = 0; i < b.idxs.size(); i++) {
      if (b.idxs[i].name == idx_b) {
        int64_t mul = mul_a / mul_b;
        size_t a_range = a.idx_by_name(idx_a)->range;
        size_t b_range = b.idx_by_name(idx_b)->range;
        if (mul > 1) {
          plan.tile_b[i] = mul_a / mul_b;
        } else if (a_range != b_range) {
          plan.tile_b[i] = b_range / a_range;
          plan.b_interleave = true;
        }
      }
    }
    plan.remap_a.emplace(idx_a, idx_a);
    plan.remap_b.emplace(idx_b, idx_a);
  }
  // Process the index that are not in the ref access
  for (size_t i = 0; i < b.idxs.size(); ++i) {
    const auto& idx = b.idxs[i];
    if (plan.remap_b.find(idx.name) == plan.remap_b.end()) {
      plan.tile_b[i] = idx.range;
    }
  }

  // Translate the constraints
  if (TranslatedContraints(scope, plan.remap_a, a) != TranslatedContraints(scope, plan.remap_b, b)) {
    IVLOG(3, "Remap a: " << plan.remap_a);
    IVLOG(3, "Remap b: " << plan.remap_b);
    IVLOG(3, "ComputeFusionPlan: incompatible constraints");
    IVLOG(4, "    a: " << StreamContainer(a.constraints));
    IVLOG(4, "    b: " << StreamContainer(b.constraints));
    return boost::none;
  }
  // Compute induced remappings
  for (const auto& idx_b : b.idxs) {
    if (idx_b.affine != Affine()) {
      for (const auto& idx_a : a.idxs) {
        if (idx_b.affine == idx_a.affine) {
          plan.remap_b.emplace(idx_b.name, idx_a.name);
        }
      }
    }
  }
  // Translate constraints
  for (const auto& constraint : b.constraints) {
    for (const auto& term : constraint.getMap()) {
      auto it = plan.remap_b.find(term.first);
      if (it == plan.remap_b.end()) {
        plan.remap_b.emplace(term.first, term.first);
      }
    }
  }
  return plan;
}

void FlattenTrivial(stripe::Block* outer) {
  IVLOG(4, "FlattenTrivial before:\n" << *outer);
  auto it = outer->stmts.begin();
  while (it != outer->stmts.end()) {
    auto inner = Block::Downcast(*it);
    // Skip non blocks
    if (!inner) {
      IVLOG(4, "FlattenTrivial: skip> non-block");
      ++it;
      continue;
    }
    uint64_t range = 1;
    for (const auto& idx : inner->idxs) {
      range *= idx.range;
    }
    if (range != 1) {
      IVLOG(4, "FlattenTrivial: skip> range != 1");
      ++it;
      continue;
    }
    bool renames = false;
    for (const auto& ref : inner->refs) {
      if (ref.from != "" && ref.into() != ref.from) {
        renames = true;
      }
    }
    // TODO: renames technically can be applied to inner statements,
    // but it's really annoying!
    if (renames) {
      IVLOG(4, "FlattenTrivial: skip> renames");
      ++it;
      continue;
    }
    // Move out inner statements
    for (auto& stmt : inner->stmts) {
      auto deep = Block::Downcast(stmt);
      if (deep) {
        // Rewrite any copied down indexes
        for (auto& idx : deep->idxs) {
          std::vector<std::string> names;
          for (const auto& item : idx.affine.getMap()) {
            if (item.first != "") {
              names.push_back(item.first);
            }
          }
          for (const auto& name : names) {
            idx.affine.substitute(name, inner->idx_by_name(name)->affine);
          }
        }
      }
      outer->stmts.insert(it, stmt);
    }
    auto it_old = it;
    ++it;
    outer->stmts.erase(it_old);
  }

  IVLOG(4, "FlattenTrivial after:\n" << *outer);
}

std::shared_ptr<Block> FusionRefactor(const stripe::Block& orig,                          //
                                      const std::map<std::string, std::string>& mapping,  //
                                      const TileShape& tile,                              //
                                      bool interleave) {
  IVLOG(3, "FusionRefactor:\n" << orig);
  IVLOG(3, "mapping: " << StreamContainer(mapping) << ", tile: " << tile);
  // Possibly tile
  auto tiled = std::make_shared<Block>(orig);
  ApplyTile(tiled.get(), tile, true, true, interleave);
  // IVLOG(3, "Tiled:\n" << *tiled);
  // Make empty inner and outer blocks, and put inner into outer
  auto outer = std::make_shared<Block>();
  outer->name = tiled->name;
  auto inner = std::make_shared<Block>();
  inner->name = tiled->name;
  outer->set_attrs(*tiled);
  outer->stmts.push_back(inner);
  // Put constraints on outer block and rewrite
  for (const auto& aff : tiled->constraints) {
    Affine out;
    for (const auto& kvp : aff.getMap()) {
      auto it = mapping.find(kvp.first);
      if (it != mapping.end()) {
        out += Affine(it->second, kvp.second);
      } else {
        out += Affine(kvp.first, kvp.second);
      }
    }
    outer->constraints.push_back(out);
  }
  // Move / rename each index to the appropriate block
  for (const auto& idx : tiled->idxs) {
    auto it = mapping.find(idx.name);
    if (it == mapping.end()) {
      if (idx.affine != Affine()) {
        IVLOG(3, "Affine idx: " << idx.name);
        outer->idxs.push_back(idx);
        inner->idxs.emplace_back(Index{idx.name, 1, idx.name});
      } else {
        IVLOG(3, "New idx: " << idx.name);
        inner->idxs.push_back(idx);
      }
    } else {
      IVLOG(3, "Existing idx: " << idx.name);
      inner->idxs.emplace_back(Index{idx.name, 1, it->second});
      outer->idxs.push_back(idx);
      outer->idxs.back().name = it->second;
    }
  }
  outer->location = tiled->location;
  // Sort outer indexes by names
  std::sort(outer->idxs.begin(), outer->idxs.end(), [](const Index& a, const Index& b) { return a.name < b.name; });
  // Copy statements to the inner block
  inner->stmts = tiled->stmts;
  // Copy refinements to both blocks
  outer->refs = tiled->refs;
  inner->refs = tiled->refs;
  // Rename mapped, and remove unmapped access elements from outer refinements
  // Also expand sizes base on inner indexes that have been removed.
  for (auto& ref : outer->refs) {
    for (size_t i = 0; i < ref.access.size(); i++) {
      auto& acc = ref.mut().access[i];
      int64_t min_val = 0;
      int64_t max_val = ref.interior_shape.dims[i].size - 1;
      Affine affine = acc.constant();
      for (const auto& kvp : acc.getMap()) {
        auto it = mapping.find(kvp.first);
        if (it == mapping.end()) {
          if (kvp.first != "") {
            if (kvp.second < 0) {
              min_val += (tiled->idx_by_name(kvp.first)->range - 1) * kvp.second;
            } else {
              max_val += (tiled->idx_by_name(kvp.first)->range - 1) * kvp.second;
            }
          }
          continue;
        }
        affine += Affine(it->second, kvp.second);
      }
      ref.mut().interior_shape.dims[i].size = max_val - min_val + 1;
      acc = affine;
    }
  }
  // Remove mapped access elements from inner refinements
  for (auto& ref : inner->refs) {
    // Rename from to match outer into
    ref.mut().from = ref.into();
    // If original was an allocation, make R/W.
    if (ref.dir == RefDir::None) {
      ref.mut().dir = RefDir::InOut;
    }
    // Update accesses
    for (auto& acc : ref.mut().access) {
      Affine affine;
      for (const auto& kvp : acc.getMap()) {
        if (kvp.first != "" && !mapping.count(kvp.first)) {
          affine += Affine(kvp.first, kvp.second);
        }
      }
      acc = affine;
    }
  }
  // Remove any trivial loops remaining
  FlattenTrivial(outer.get());
  // Return final result
  IVLOG(3, "Refactor output:\n" << *outer);
  return outer;
}

bool FuseBlocks(const AliasMap& scope, Block* block_a, Block* block_b) {
  // If indexes don't match, fail
  if (block_a->idxs != block_b->idxs) {
    IVLOG(3, "Fuse failed due to mismatched indexes");
    IVLOG(3, "A: " << block_a->idxs);
    IVLOG(3, "B: " << block_b->idxs);
    return false;
  }
  // If constraints don't match, fail
  if (block_a->constraints != block_b->constraints) {
    IVLOG(3, "Fuse failed due to mismatched constraints");
    return false;
  }
  // If locations don't match, fail
  if (block_a->location != block_b->location) {
    IVLOG(3, "Fuse failed due to mismatched locations");
    return false;
  }
  // Make AliasMaps for the two blocks
  AliasMap a_map(scope, block_a);
  AliasMap b_map(scope, block_b);
  // Start by copying A's reference across
  auto tmp = std::make_shared<Block>();
  tmp->refs = block_a->refs;
  // Walk over refinements in B and move them across
  // Rename duplicate refinements in B to their name in A
  // Otherwise make a new unique name (keeping original if possible)
  std::map<std::string, std::string> remap_b;
  for (const auto& new_ref : block_b->refs) {
    // If it's a local, always safe to copy if across
    // Check if b matches something in the existing block
    bool merged = false;
    for (auto& old_ref : block_a->refs) {
      auto atype = AliasInfo::Compare(a_map.at(old_ref.into()), b_map.at(new_ref.into()));
      if (atype == AliasType::Partial) {
        // Conflict, if either do any writing, we have a problem
        if (IsWriteDir(new_ref.dir) || IsWriteDir(old_ref.dir)) {
          IVLOG(3, "Fuse failed due to mismatched aliases: " << old_ref.into() << " vs " << new_ref.into());
          return false;  // Fuse will not work, bail
        }
      } else if (atype == AliasType::Exact) {
        remap_b[new_ref.into()] = old_ref.into();
        old_ref.mut().dir = UnionDir(old_ref.dir, new_ref.dir);
        merged = true;
        break;
      }
    }
    if (!merged) {
      // Copy across as a new ref
      std::string new_name = tmp->unique_ref_name(new_ref.into());
      remap_b[new_ref.into()] = new_name;
      tmp->refs.emplace(new_ref.WithInto(std::move(new_name)));
    }
  }
  // We are now safe (cannot fail), move new reference over A's
  std::swap(block_a->refs, tmp->refs);
  if (!block_a->name.empty()) {
    block_a->name = str(boost::format("%s+%s") % block_a->name % block_b->name);
  } else if (!block_b->name.empty()) {
    block_a->name = block_b->name;
  }
  // Load all the scalars that exist as of block A
  std::set<std::string> all_scalars;
  std::map<std::string, std::string> scalar_rename;
  for (const auto& stmt : block_a->stmts) {
    for (const auto& name : stmt->scalar_defs()) {
      all_scalars.emplace(name);
    }
  }
  auto def_scalar = [&](const std::string& orig) -> std::string {
    if (all_scalars.count(orig) == 0) {
      all_scalars.emplace(orig);
      scalar_rename[orig] = orig;
      return orig;
    }
    for (size_t i = 0; true; i++) {
      std::string with_suffix = orig + "_" + std::to_string(i);
      if (all_scalars.count(with_suffix) == 0) {
        all_scalars.emplace(with_suffix);
        scalar_rename[orig] = with_suffix;
        return with_suffix;
      }
    }
    return "";
  };
  // Now move across statements, updating references/scalars as we do:
  for (const auto& stmt : block_b->stmts) {
    switch (stmt->kind()) {
      case StmtKind::Load: {
        auto op = Load::Downcast(stmt);
        op->into = def_scalar(op->into);
        op->from = remap_b.at(op->from);
      } break;
      case StmtKind::Store: {
        auto op = Store::Downcast(stmt);
        op->into = remap_b.at(op->into);
        op->from = scalar_rename.at(op->from);
      } break;
      case StmtKind::LoadIndex: {
        // Currently we can't reach here actually because LoadIndex
        // takes only index which must not be the output of block_a.
        // If we get here in the future, the following code need
        // to be tested.
        auto op = LoadIndex::Downcast(stmt);
        op->into = def_scalar(op->into);
        op->from = TranslateAffine(op->from, scope, remap_b, block_b);
      }
      case StmtKind::Special: {
        auto op = Special::Downcast(stmt);
        for (auto& in : op->inputs) {
          in = remap_b.at(in);
        }
        for (auto& out : op->outputs) {
          out = remap_b.at(out);
        }
      } break;
      case StmtKind::Block: {
        auto op = Block::Downcast(stmt);
        for (auto& ref : op->refs) {
          ref.mut().from = remap_b.at(ref.from);
        }
      } break;
      case StmtKind::Constant: {
        auto op = Constant::Downcast(stmt);
        op->name = def_scalar(op->name);
      } break;
      case StmtKind::Intrinsic: {
        auto op = Intrinsic::Downcast(stmt);
        for (auto& in : op->inputs) {
          in = scalar_rename.at(in);
        }
        for (auto& out : op->outputs) {
          out = def_scalar(out);
        }
      } break;
    }
    block_a->stmts.push_back(stmt);
  }
  // All is well
  return true;
}

void FusionInner(const AliasMap& scope, Block* block, FusionStrategy* strategy, bool no_inner) {
  // Start with the first statement, and keep tying to fuse until you can't anymore, then move to the next
  auto it = block->stmts.begin();
  while (it != block->stmts.end()) {
    // If it's not a block, forget it!
    if ((*it)->kind() != StmtKind::Block) {
      ++it;
      continue;
    }
    while (true) {
      // Get block everytime in case it's updated
      auto block1 = Block::Downcast(*it);
      IVLOG(3, "Attempting fusion on block:\n" << block1->name);
      // Get the next statement
      auto it_next = it;
      it_next++;
      // If there is no next statement, I'm done with this block
      if (it_next == block->stmts.end()) {
        break;
      }
      // Convert to block
      auto block2 = Block::Downcast(*it_next);
      // If it's not a block, forget it
      if (!block2) {
        break;
      }
      // Get the list of outputs for this block
      std::set<std::string> outs_for_fuse;
      for (const auto& ro : block1->ref_outs()) {
        IVLOG(3, "Considering output: " << ro->from);
        outs_for_fuse.emplace(ro->from);
      }
      IVLOG(3, "Outs for fuse size: " << outs_for_fuse.size());
      std::string fuse_on = "";
      // Check if it's a match to any of the inputs on the next block
      for (const auto& ri : block2->ref_ins()) {
        IVLOG(3, "Considering input: " << ri->from);
        if (outs_for_fuse.count(ri->from)) {
          fuse_on = ri->from;
          break;
        }
      }
      // Nothing to fuse on, done with this block
      if (fuse_on == "") {
        IVLOG(3, "Nothing to fuse on");
        break;
      }
      IVLOG(3, "Fuse on = " << fuse_on);
      // Compute a fusion plan for the two blocks, if fails, give up
      auto plan = ComputeFusionPlan(scope, *block1, *block2, fuse_on);
      if (!plan) {
        IVLOG(3, "Fusion plan failed");
        break;
      }
      // Now call the strategy to see if we should fuse
      if (!strategy->AttemptFuse(*block, *block1, *block2)) {
        IVLOG(3, "Fusion denied by strategy");
        break;
      }
      // Do the appropriate refactors
      auto refactor1 = FusionRefactor(*block1, plan->remap_a, plan->tile_a, plan->a_interleave);
      auto refactor2 = FusionRefactor(*block2, plan->remap_b, plan->tile_b, plan->b_interleave);
      if (no_inner) {
        // Check if there is any inner block in refactor1 and refactor2
        if (!NoInnerBlock(refactor1.get()) || !NoInnerBlock(refactor2.get())) {
          IVLOG(3, "Inner block exists");
          break;
        }
      }
      // Now copy computed indexes if it's safe
      for (const auto& idx : refactor1->idxs) {
        if (idx.affine != stripe::Affine() && !refactor2->idx_by_name(idx.name)) {
          refactor2->idxs.push_back(idx);
        }
      }
      for (const auto& idx : refactor2->idxs) {
        if (idx.affine != stripe::Affine() && !refactor1->idx_by_name(idx.name)) {
          refactor1->idxs.push_back(idx);
        }
      }
      // Sort one more time
      std::sort(refactor1->idxs.begin(), refactor1->idxs.end(),
                [](const Index& a, const Index& b) { return a.name < b.name; });
      std::sort(refactor2->idxs.begin(), refactor2->idxs.end(),
                [](const Index& a, const Index& b) { return a.name < b.name; });
      // IVLOG(3, "Fusion refactor 1:\n" << *refactor1);
      // IVLOG(3, "Fusion refactor 2:\n" << *refactor2);
      // Try the actual fusion
      if (!FuseBlocks(scope, refactor1.get(), refactor2.get())) {
        strategy->OnFailed();
        IVLOG(3, "Actual fusion failed");
        break;
      }
      IVLOG(3, "Fused block:\n" << *refactor1);
      // If it worked, update
      *it = refactor1;
      block->stmts.erase(it_next);
      strategy->OnFused(scope, refactor1.get(), *block1, *block2);
    }
    it++;
  }
}

struct FusionPassOptions {
  Tags parent_reqs;
  Tags a_block_reqs;
  Tags b_block_reqs;
  Tags fused_set;
  Tags exclude;
  bool perfect;
  bool output_match;
  bool no_inner;
};

class TagFusionStrategy : public FusionStrategy {
 public:
  explicit TagFusionStrategy(const FusionPassOptions& options) : options_(options) {}
  bool AttemptFuse(const stripe::Block& parent, const stripe::Block& a, const stripe::Block& b) {
    bool tag_match = parent.has_tags(options_.parent_reqs) &&  //
                     a.has_tags(options_.a_block_reqs) &&      //
                     b.has_tags(options_.b_block_reqs) &&      //
                     !a.has_any_tags(options_.exclude) &&      //
                     !b.has_any_tags(options_.exclude);
    if (!tag_match) {
      return false;
    }
    if (options_.output_match && a.idxs.size() != b.idxs.size()) {
      return false;
    }
    return true;
  }
  void OnFailed() {}
  void OnFused(const AliasMap& outer, stripe::Block* block, const stripe::Block& a, const stripe::Block& b) {
    block->add_tags(options_.fused_set);
  }
  bool NoInner() { return options_.no_inner; }

 private:
  const FusionPassOptions& options_;
};

static void FusionPassRecurse(const AliasMap& map, stripe::Block* block, TagFusionStrategy* strategy) {
  FusionInner(map, block, strategy, strategy->NoInner());
  for (const auto& stmt : block->stmts) {
    auto inner = Block::Downcast(stmt);
    if (inner) {
      AliasMap inner_map(map, inner.get());
      FusionPassRecurse(inner_map, inner.get(), strategy);
    }
  }
}

void FusionPass::Apply(CompilerState* state) const {
  stripe::Block* root = state->entry();
  FusionPassOptions fopts = {
      FromProto(options_.parent_reqs()),  // parent_reqs
      FromProto(options_.a_reqs()),       // a_block_reqs
      FromProto(options_.b_reqs()),       // b_block_reqs
      FromProto(options_.fused_set()),    // fused_set
      FromProto(options_.exclude()),      // exclude
      options_.perfect(),                 // perfect
      options_.output_match(),            // output_match
      options_.no_inner()                 // elide trivial plan
  };
  AliasMap base;
  AliasMap root_map(base, root);
  // Check if we should fuse this block
  TagFusionStrategy strategy(fopts);
  FusionPassRecurse(root_map, root, &strategy);
}

namespace {
[[gnu::unused]] char reg = []() -> char {
  CompilePassFactory<FusionPass, proto::FusionPass>::Register();
  return 0;
}();
}  // namespace
}  // namespace codegen
}  // namespace tile
}  // namespace vertexai


#include "examples/Cassie/input_supervisor.h"

#include "drake/multibody/joints/floating_base_types.h"
#include "drake/multibody/parsers/urdf_parser.h"
#include "drake/multibody/rigid_body_tree.h"
#include "drake/multibody/rigid_body_tree_construction.h"
#include "examples/Cassie/cassie_utils.h"
#include <Eigen/Dense>
#include <gtest/gtest.h>

namespace dairlib {
namespace systems {
namespace {

using Eigen::VectorXd;

// class InputSupervisorTest: public ::testing::Test {};

class InputSupervisorTest : public ::testing::Test {
protected:
  InputSupervisorTest() {
    buildCassieTree(tree_, "examples/Cassie/urdf/cassie_v2.urdf",
                    drake::multibody::joints::kQuaternion);
    drake::multibody::AddFlatTerrainToWorld(&tree_, 100, 0.2);
    supervisor_ = std::make_unique<InputSupervisor>(
        tree_, 10.0, 0.01, min_consecutive_failures, 20.0);
    context_ = supervisor_->CreateDefaultContext();
    status_output_ = std::make_unique<TimestampedVector<double>>(1);
    command_input_ =
        std::make_unique<TimestampedVector<double>>(tree_.get_num_actuators());
    state_input_ = std::make_unique<OutputVector<double>>(
        tree_.get_num_positions(), tree_.get_num_velocities(),
        tree_.get_num_actuators());
  }

  RigidBodyTree<double> tree_;
  const int min_consecutive_failures = 5;
  std::unique_ptr<InputSupervisor> supervisor_;
  std::unique_ptr<TimestampedVector<double>> status_output_;
  std::unique_ptr<TimestampedVector<double>> command_input_;
  std::unique_ptr<OutputVector<double>> state_input_;
  std::unique_ptr<drake::systems::Context<double>> context_;
};

TEST_F(InputSupervisorTest, StatusBitTest) {
  double output_bit;
  VectorXd zero_input = VectorXd::Zero(tree_.get_num_actuators());
  command_input_->get_mutable_value() = zero_input;
  context_->FixInputPort(0, *command_input_);

  supervisor_->SetStatus(*context_, status_output_.get());
  output_bit = status_output_->get_value()[0];
  EXPECT_EQ(output_bit, 0);

  VectorXd large_input = 100 * VectorXd::Ones(tree_.get_num_actuators());
  command_input_->get_mutable_value() = large_input;
  context_->FixInputPort(0, *command_input_);
  supervisor_->SetStatus(*context_, status_output_.get());
  output_bit = status_output_->get_value()[0];
  EXPECT_EQ(output_bit, 2);

  VectorXd high_velocities = 100 * VectorXd::Ones(tree_.get_num_velocities());
  state_input_->SetVelocities(high_velocities);
  context_->FixInputPort(1, *state_input_);
  supervisor_->UpdateErrorFlag(*context_,
                               &context_->get_mutable_discrete_state());
  supervisor_->SetStatus(*context_, status_output_.get());
  output_bit = status_output_->get_value()[0];
  EXPECT_EQ(output_bit, 3);

  // Trigger the min_consecutive_failures
  for (int i = 0; i < min_consecutive_failures; ++i) {
    supervisor_->UpdateErrorFlag(*context_,
                                 &context_->get_mutable_discrete_state());
  }
  supervisor_->SetStatus(*context_, status_output_.get());
  output_bit = status_output_->get_value()[0];
  EXPECT_EQ(output_bit, 7);
}

} // namespace
} // namespace systems
} // namespace dairlib

int main(int argc, char **argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


// Copyright PA Knowledge Ltd 2021
// MIT License. For licence terms see LICENCE.md file

#include "Client.hpp"
#include "DataSendPeriodController.hpp"
#include "TimerInterface.hpp"
#include "PacketGeneratorInterface.hpp"

#include "spdlog/spdlog.h"
#include <boost/algorithm/string.hpp>
#include <chrono>
#include <filesystem>
#include <istream>
#include <random>
#include <regex>
#include <utility>

Client::Client(
  std::shared_ptr<UdpClientInterface> udpClient,
  std::shared_ptr<TimerInterface> timer,
  std::uint16_t maxPayloadSize,
  std::unique_ptr<DataSendPeriodController> dataSendPeriodController,
  std::unique_ptr<PacketGeneratorInterface> packetGenerator)
    : udpClient(std::move(udpClient))
    , edTimer(std::move(timer))
    , maxPayloadSize(maxPayloadSize)
    , dataRateCalculator{std::move(dataSendPeriodController)}
    , packetGenerator {std::move(packetGenerator)}
{
}

Client::~Client() {} // NOLINT(modernize-use-equals-default)

void Client::send(std::istream& inputStream)
{
  inputStream.exceptions(std::istream::badbit);
  if (!inputStream)
  {
    throw std::runtime_error("file stream not found");
  }

  setSessionID();

  edTimer->runTimer([&]() {
    try
    {
      adjustDataRate();
      return sendFrame(inputStream);
    }
    catch (const std::exception& exception)
    {
      spdlog::error(std::string("exception in send frame ") + exception.what());
      return false;
    }
  });
}

void Client::adjustDataRate()
{
  if (!dataRateCalculator)
  {
    return;
  }

  edTimer->setTickInterval(dataRateCalculator->getPeriod());
}

bool Client::sendFrame(std::istream& inputStream)
{
  const auto packet_to_send = packetGenerator->next(inputStream, maxPayloadSize);
  if(!packet_to_send)
  {
    return false;
  }

  udpClient->send(*packet_to_send);
  return true;
}

void Client::setSessionID()
{
  const auto seed = std::chrono::system_clock::now().time_since_epoch().count();
  packetGenerator->setSessionID((std::uint32_t)std::default_random_engine(seed)());
}

boost::posix_time::microseconds calculateTimerPeriod(double dataRateMbps, std::uint32_t mtuSize)
{
  const auto period = std::round((static_cast<double>((mtuSize * 8)) * 1000000) / (dataRateMbps * 1024 * 1024));
  return boost::posix_time::microseconds {static_cast<std::int64_t>(period)};
}


#include <Rcpp.h>
#include <vector>
#include <string>
#include <unordered_set>
#include <algorithm>
#include <utility>
#include <random>
#include <thread>
#include <boost/format.hpp>
#include "defines.h"
#include "Segment.h"
#include "Region.h"
#include "CNVR.h"
#include "df_to_segments.h"
#include "get_regions.h"
#include "statistical_model.h"
#include "query_model.h"
#include "merge.h"

using namespace Rcpp;

// [[Rcpp::export]]
List convaqCpp(
    DataFrame df1,
    DataFrame df2,
    unsigned int model_num,
    bool qvalues,
    unsigned int qvalues_rep,
    bool merge,
    unsigned int merge_threshold,
    double cutoff,
    unsigned int comp1, double value1, unsigned int eq1, unsigned int type1,
    unsigned int comp2, double value2, unsigned int eq2, unsigned int type2,
    unsigned int nthreads
) {
  if(nthreads == 0) nthreads = std::thread::hardware_concurrency();

  MODEL model = (MODEL)model_num;

  std::vector<Segment> segments1, segments2;

  // convert data frames to vector of Segment objects
  df_to_segments(df1, segments1);
  df_to_segments(df2, segments2);

  // get number of patients in each group
  IntegerVector patients1 = df1["patient"];
  IntegerVector patients2 = df2["patient"];
  int npatients[2] = {
    Rcpp::max(patients1)+1,
    Rcpp::max(patients2)+1
  };

  // grab all chromosomes in data set
  std::unordered_set<std::string> chromosomes;
  StringVector chr1 = df1["chr"];
  StringVector chr2 = df2["chr"];
  for(size_t i = 0; i < chr1.length(); ++i) chromosomes.insert(std::string(chr1[i]));
  for(size_t i = 0; i < chr2.length(); ++i) chromosomes.insert(std::string(chr2[i]));

  std::vector<Region> regions;
  get_regions(segments1, segments2, npatients[0], npatients[1], chromosomes, regions);

  std::vector<CNVR> results;

  if(model == MODEL_STAT) {
    statistical_model(regions, npatients[0], npatients[1], cutoff, results);
  } else {
    query_model(
      regions, npatients[0], npatients[1],
      (COMPARISON)comp1, value1, (EQUALITY)eq1, (VARIATION_TYPE)type1,
      (COMPARISON)comp2, value2, (EQUALITY)eq2, (VARIATION_TYPE)type2,
      results
    );
  }

  if(results.size() > 0 && merge) merge_adjacent(results, merge_threshold);

  // sort by p-value
  std::sort(results.begin(), results.end(), [](const CNVR &a, const CNVR &b) { return a.pvalue < b.pvalue; });

  if(results.size() > 0 && qvalues) {
    std::vector<std::vector<int>> best(4);

    for(size_t i = 0; i < 4; ++i) best[i].resize(qvalues_rep, 0);

    std::vector<std::thread> threads;
    for(size_t tid = 0; tid < nthreads; ++tid) {
      threads.push_back(std::thread([&](size_t offset) {
        // collect all group-patient pairs
        std::vector<std::pair<int,int>> all_patients;
        for(size_t group = 0; group < 2; ++group) {
          for(int i = 0; i < npatients[group]; ++i) all_patients.emplace_back(group, i);
        }

        std::random_device rd;
        std::minstd_rand rand(rd());

        for(size_t rep = offset; rep < qvalues_rep; rep += nthreads) {
          std::shuffle(all_patients.begin(), all_patients.end(), rand);

          std::vector<std::set<int>> selected(2);
          for(size_t i = 0; i < npatients[0]; ++i) {
            selected[all_patients[i].first].insert(all_patients[i].second);
          }

          std::vector<Segment> q_segments1, q_segments2;
          for(const Segment &s : segments1) {
            if(selected[0].find(s.patient) != selected[0].end()) {
              q_segments1.push_back(s);
            } else {
              q_segments2.push_back(s);
            }
          }
          for(const Segment &s : segments2) {
            if(selected[1].find(s.patient) != selected[1].end()) {
              q_segments1.push_back(s);
            } else {
              q_segments2.push_back(s);
            }
          }

          std::vector<Region> q_regions;
          get_regions(q_segments1, q_segments2, npatients[0], npatients[1], chromosomes, q_regions);

          std::vector<CNVR> q_results;
          if(model == MODEL_STAT) {
            statistical_model(q_regions, npatients[0], npatients[1], cutoff, q_results);
          } else if(model == MODEL_QUERY) {
            query_model(
              q_regions, npatients[0], npatients[1],
              (COMPARISON)comp1, value1, (EQUALITY)eq1, (VARIATION_TYPE)type1,
              (COMPARISON)comp2, value2, (EQUALITY)eq2, (VARIATION_TYPE)type2,
              q_results
            );
          }

          if(q_results.size() > 0 && merge) merge_adjacent(q_results, merge_threshold);

          for(const CNVR &r : q_results) {
            best[r.type][rep] = std::max(best[r.type][rep], r.length);
          }
        }
      }, tid));
    }

    for(std::thread &th : threads) th.join();

    for(CNVR &r : results) {
      int better = 0;
      for(int l : best[r.type]) {
        if(l >= r.length) ++better;
      }
      r.qvalue = (double)better / qvalues_rep;
    }

    std::sort(results.begin(), results.end(), [](const CNVR &a, const CNVR &b) { return a.qvalue < b.qvalue; });
  }

  // prepare output
  std::vector<std::string> df_chr;
  std::vector<int> df_start, df_end, df_length, df_type;
  std::vector<double> df_pvalue, df_qvalue;

  std::transform(results.begin(), results.end(), std::back_inserter(df_chr),    [](const CNVR &r){ return r.chr; });
  std::transform(results.begin(), results.end(), std::back_inserter(df_start),  [](const CNVR &r){ return r.start; });
  std::transform(results.begin(), results.end(), std::back_inserter(df_end),    [](const CNVR &r){ return r.end; });
  std::transform(results.begin(), results.end(), std::back_inserter(df_length), [](const CNVR &r){ return r.length; });
  std::transform(results.begin(), results.end(), std::back_inserter(df_type),   [](const CNVR &r){ return r.type; });
  std::transform(results.begin(), results.end(), std::back_inserter(df_pvalue), [](const CNVR &r){ return r.pvalue; });
  std::transform(results.begin(), results.end(), std::back_inserter(df_qvalue), [](const CNVR &r){ return r.qvalue; });

  DataFrame out_regions = DataFrame::create(
    Named("chr") = df_chr,
    Named("start") = df_start,
    Named("end") = df_end,
    Named("length") = df_length,
    Named("type") = df_type,
    Named("pvalue") = df_pvalue,
    Named("qvalue") = df_qvalue
  );

  // get within-group frequencies
  List out_freq = List::create();
  for(const CNVR &r : results) {
    List region_freq = List::create();
    for(size_t group = 0; group < 2; ++group) {
      List group_freq = List::create();
      for(size_t type = 0; type < 3; ++type) {
        group_freq.push_back(r.get_freq(group, type));
      }
      region_freq.push_back(group_freq);
    }
    out_freq.push_back(region_freq);
  }

  List out_state = List::create();
  for(CNVR &r : results) {
    List r_state = List::create();
    for(size_t group = 0; group < 2; ++group) {
      r_state.push_back(r.get_state(group));
    }
    out_state.push_back(r_state);
  }

  return List::create(
    Named("regions") = out_regions,
    Named("freq") = out_freq,
    Named("state") = out_state
  );
}


#define BOOST_TEST_MODULE TestCore
#include <boost/test/included/unit_test.hpp>
#define BOOST_LOG_DYN_LINK 1
#include <boost/log/trivial.hpp>
#include <DepthCodecFactory.h>
#include <boost/serialization/export.hpp>
#include <boost/serialization/shared_ptr.hpp>
#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>
#include "opencv2/opencv.hpp"
#include "CodecEvalFramework.h"
#include "NodeAddress.h"
#include "QuadTreeTypes.h"
#include "QuadTreeCodecFactory.h"
#include "opencv2/core/version.hpp"
#include "ReadWrite.h"
#include "Range.h"
#include "BoostMarshaller.hpp"
#include "Core/Registration/Registration.h"

BOOST_CLASS_EXPORT(RollingQT32bitMinMaxAbsDiff)
BOOST_CLASS_EXPORT(RollingQT16bitMinMaxAbsDiff)
BOOST_CLASS_EXPORT_GUID(IDepthCodec, "IDepthCodec")
BOOST_CLASS_EXPORT_GUID(TiledCodec,"TiledCodec")

BOOST_AUTO_TEST_CASE(TestLoadRosBag){
        auto fs = LoadRosBag("/home/philip/Downloads/structured.bag");

        auto originalFrame = fs.grabFrame();

        boost::program_options::variables_map compressionOptions;
        compressionOptions.insert(std::make_pair("CodecType", po::variable_value(std::string("QuadTree"), false)));
        po::notify(compressionOptions);
        auto codecFactory = std::make_shared<DepthCodecFactory>(compressionOptions);
        codecFactory->registerSubFactory("QuadTree", std::make_shared<QuadTreeCodecFactory>(compressionOptions));

        auto codec = std::make_shared<TiledCodec>(codecFactory);

        cv::Mat depthImage = originalFrame.getDepthImage();

        /*CodecEvalFramework<cv::Mat, TiledCodec> testFramework(codec);
        auto compressionPerf = testFramework.evaluateCodecOnExample(depthImage, false);
        compressionPerf.printPerformance(std::cout);*/

        cv::Mat rslt;
        compressAndDecompress(codec, depthImage, rslt);
        cv::imshow("sdrf",rslt*10);
        cv::waitKey(0);
        Frame compressedFrame = originalFrame;
        compressedFrame.updateDepthImage(rslt);

        showPointCloudCompression(originalFrame, compressedFrame);

}

BOOST_AUTO_TEST_CASE(TestMapStream)
{
    std::map<int, std::string> test{ {2, "hello"}, {3,"world" } }, rslt;
	BoostMarshaller<std::map<int, std::string>> bm;
    auto data = bm.marshall(test);
    bm.unmarshall(data);

	BOOST_TEST(test == rslt);
}

BOOST_AUTO_TEST_CASE(TestTest){
    std::cout << "Using OpenCV "
              << CV_VERSION_MAJOR << "."  // major version
              << CV_VERSION_MINOR << "."  // minor version
              << CV_VERSION_REVISION    // patch level
              << std::endl;

    std::cout << "Using Boost "
              << BOOST_VERSION / 100000     << "."  // major version
              << BOOST_VERSION / 100 % 1000 << "."  // minor version
              << BOOST_VERSION % 100                // patch level
              << std::endl;
}

//return the depth in the tree given an address
BOOST_AUTO_TEST_CASE(TestQuadTreeGetDepth) {
    NodeAddress32bit address = 0b0000100101110100;
    BOOST_TEST(getDepth(address) == 4);
    address = 0b0000000101110100;
    BOOST_TEST(getDepth(address) == 3);
    address = 0b0000000111111111;
    BOOST_TEST(getDepth(address) == 3);
}

/// given an address, return the tl tr bl br of that address
BOOST_AUTO_TEST_CASE(TestQuadTreeSiblings) {
    NodeAddress32bit address = 0b0000101101110100;

    auto pa = parentAddress(address);
    auto tl = topLeft(pa);
    auto tr = topRight(pa);
    auto bl = bottomLeft(pa);
    auto br = bottomRight(pa);

    /// tl tr bl br in Z-order   https://en.wikipedia.org/wiki/Z-order_curve
    printAddress(tl);
    BOOST_TEST(tl == 0b0000100101110100);
    BOOST_TEST(tr == 0b0000101101110100);
    BOOST_TEST(bl == 0b0000110101110100);
    BOOST_TEST(br == 0b0000111101110100);
}

BOOST_AUTO_TEST_CASE(TestEncodeAddressRecurse) {
    auto a = encodeAddressRecurse<NodeAddress32bit>(3, 2, 4/2, 4/2, 2);
    printAddress(a);
    BOOST_TEST(a == 0b101111);

    auto b = encodeAddressRecurse<NodeAddress32bit>(199, 199, 200/2, 200/2, 6);
    printAddress(b);
    BOOST_TEST(b == 0b111111111111111111);

    auto c = encodeAddressRecurse<NodeAddress32bit>(0, 0, 200/2, 200/2, 6);
    printAddress(c);
    BOOST_TEST(c == 0b100100100100100100);

    auto d = encodeAddressRecurse<NodeAddress32bit>(1, 2, 4/2, 4/2, 2);
    printAddress(d);
    BOOST_TEST(d == 0b101110);
}

BOOST_AUTO_TEST_CASE(TestDecodeAddressRecurse) {
    auto a = decodeAddressRecurse(0b101111, 0, 0, 4, 4);
    BOOST_TEST(std::get<0>(a) == 3);
    BOOST_TEST(std::get<1>(a) == 2);
    BOOST_TEST(std::get<0>(a) == 3);
    BOOST_TEST(std::get<1>(a) == 2);

    auto b = decodeAddressRecurse(0b111111111111111111111111, 0, 0, 256, 256);
    BOOST_TEST(std::get<0>(b) == 255);
    BOOST_TEST(std::get<1>(b) == 255);

    auto c = decodeAddressRecurse(0b100100100100100100, 0, 0, 256, 256);
    BOOST_TEST(std::get<0>(c) == 0);
    BOOST_TEST(std::get<1>(c) == 0);

    auto d = decodeAddressRecurse(0b101110, 0, 0, 4, 4);
    BOOST_TEST(std::get<0>(d) == 1);
    BOOST_TEST(std::get<1>(d) == 2);

    auto e = decodeAddressRecurse(0b101110, 0, 0, 4, 4);
    BOOST_TEST(std::get<0>(e) == 1);
    BOOST_TEST(std::get<1>(e) == 2);

    auto f = decodeAddressRecurse(0b101110, 0, 0, 4, 4);
    BOOST_TEST(std::get<2>(f) == 1);
    BOOST_TEST(std::get<3>(f) == 1);

    auto g = decodeAddressRecurse(0b000110, 0, 0, 4, 4);
    BOOST_TEST(std::get<2>(g) == 2);
    BOOST_TEST(std::get<3>(g) == 2);

}

BOOST_AUTO_TEST_CASE(TestPerfectEncodeDecode){
    NodeAddress32bit address = 0b000100;
    auto pa = parentAddress(address);
    BOOST_TEST(pa == 0b000);

    cv::Mat example = cv::imread("../modules/core/resources/bgExampleDepth.tif");
    cv::resize(example, example, cv::Size(512,512), cv::INTER_NEAREST);

    cv::Mat gray_image;
    cv::cvtColor(example, gray_image, CV_BGR2GRAY);
    cv::Mat shortMat;
    gray_image.convertTo(shortMat, CV_16UC1, 65536/255);

    using Codec = RollingQuadTree<NodeAddress32bit, RangeLeaf<unsigned short>, PSNRPolicy>;
    Codec t(PSNRPolicy(30 * (65536/255)));

    BoostMarshaller<Codec> bm;
    t.compress(shortMat);
    auto data = bm.marshall(t);
    t = bm.unmarshall(data);
    cv::Mat decompressed(512, 512, CV_16UC1, cv::Scalar(0));
    t.decompress(decompressed);
    showCompressionArtifacts(shortMat, decompressed);
    std::vector<uchar> exampleVec(example.datastart,  example.dataend);
    std::vector<uchar> decompressedVec(decompressed.datastart,  decompressed.dataend);

    /*BOOST_TEST(exampleVec.size() == decompressedVec.size());
    for(size_t i = 0; i < exampleVec.size(); i++)
        BOOST_TEST(exampleVec[i] == decompressedVec[i]);*/

    cv::imshow("decompressed", decompressed);
    cv::waitKey(0);
}

class MockLossyCodec:public IDepthCodec{
    cv::Mat mImage;

public:
    MockLossyCodec(){};

    void compress(const cv::Mat& data){
        mImage = data.clone();
        BOOST_LOG_TRIVIAL(info) << "Compressing Data";
    };

    void decompress(cv::Mat& data){
        data = mImage;
    };

    template<class Archive>
    void serialize(Archive & ar, const unsigned int version)
    {
        // note, version is always the latest when saving
        boost::serialization::void_cast_register<MockLossyCodec, IDepthCodec>();
        ar & mImage;
    }

    friend class boost::serialization::access;
};

BOOST_AUTO_TEST_CASE(TestCodecFramework){
    auto codec = std::make_shared<MockLossyCodec>();
    CodecEvalFramework<cv::Mat, MockLossyCodec> testFramework(codec);
    auto rslt = testFramework.evaluateCodecOnExample(cv::Mat(640, 480, CV_16UC1, cv::Scalar(0)), false);
    rslt.printPerformance(std::cout);
}

BOOST_AUTO_TEST_CASE(TestPlayRosBag){
    FrameSource fs("/home/philip/Downloads/structured.bag");

    auto frame = fs.grabFrame();
    frame.getIntrin();
}

BOOST_AUTO_TEST_CASE(TestCodecFactory){
    cv::Mat example = cv::imread("../modules/core/resources/bgExampleDepth.tif");

    cv::Mat gray_image;
    cv::cvtColor(example, gray_image, CV_BGR2GRAY);
    cv::Mat shortMat;
    gray_image.convertTo(shortMat, CV_16UC1, 65536/255);

    cv::resize(shortMat, shortMat, cv::Size(1280,720));

    boost::program_options::variables_map compressionOptions;
    compressionOptions.insert(std::make_pair("CodecType",po::variable_value(std::string("QuadTree"), false)));
    po::notify(compressionOptions);
    auto codecFactory = std::make_shared<DepthCodecFactory>(compressionOptions);
    codecFactory->registerSubFactory("QuadTree", std::make_shared<QuadTreeCodecFactory>(compressionOptions));

    auto codec = std::make_shared<TiledCodec>(codecFactory);

    CodecEvalFramework<cv::Mat, IDepthCodec> testFramework(codec);
    auto rslt = testFramework.evaluateCodecOnExample(shortMat, true);
    rslt.printPerformance(std::cout);

    cv::waitKey(0);
}

//BOOST_CLASS_EXPORT(Tree)


#include "agency_rating.h"
#include "util.h"

#include <boost/algorithm/string.hpp>

#include <fstream>
#include <cmath>

void TAgencyRating::Load(const std::string& filePath, bool setMinAsUnk) {
    std::string line;
    std::ifstream rating(filePath);
    if (!rating.is_open()) {
        LOG_DEBUG("Rating file is not available");
        return;
    }
    while (std::getline(rating, line)) {
        std::vector<std::string> lineSplitted;
        boost::split(lineSplitted, line, boost::is_any_of("\t"));
        Records[lineSplitted[1]] = std::stod(lineSplitted[0]);
    }

    if (setMinAsUnk) {
        UnkRating = std::min_element(Records.begin(), Records.end(),
            [](const std::pair<std::string, double>& item1, const std::pair<std::string, double>& item2) {
                return item1.second < item2.second;
            }
        )->second;
    }
}

double TAgencyRating::ScoreUrl(const std::string& url) const {
    const auto iter = Records.find(GetHost(url));
    return (iter != Records.end()) ? iter->second : UnkRating;
}

void TAlexaAgencyRating::Load(const std::string& filePath) {
    std::ifstream fileStream(filePath);
    nlohmann::json json;
    fileStream >> json;
    for (const nlohmann::json& agency : json) {
        std::string host = agency.at("host").get<std::string>();
        double rating = agency.at("rating").get<double>();
        RawRating[host] = rating;

        for (auto& [key, value] : agency.at("country").items()) {
            CountryShare[host][key] = value;
        }
    }
}

double TAlexaAgencyRating::GetCountryShare(const std::string& host, const std::string& code) const {
    const auto iter = CountryShare.find(host);
    if (iter == CountryShare.end()) {
        return 0.;
    }
    const auto iter2 = iter->second.find(code);
    if (iter2 == iter->second.end()) {
        return 0.;
    }
    return iter2->second;
}

double TAlexaAgencyRating::GetRawRating(const std::string& host) const {
    const auto iter = RawRating.find(host);
    return (iter != RawRating.end()) ? iter->second : UnkRating;
}

double TAlexaAgencyRating::ScoreUrl(
    const std::string& host,
    tg::ELanguage language,
    ERatingType type,
    double shift
) const {
    if (type == RT_ONE) {
        return 1.;
    }
    double raw = GetRawRating(host);
    double coeff = 0;
    if (language == tg::LN_EN) {
        coeff = (100. - GetCountryShare(host, "US") - GetCountryShare(host, "GB"))/100.;
    } else {
        coeff = GetCountryShare(host, "RU");
    }
    if (type == RT_LOG) {
        return std::max(log(raw * coeff + shift), 0.3);
    } else if (type == RT_RAW) {
        return raw * coeff;
    }
    return 1.;
}