#ifndef ZIRAN_ADMMSKETCH_H
#define ZIRAN_ADMMSKETCH_H

#include <Ziran/CS/Util/Forward.h>
#include <Eigen/Eigenvalues>

namespace ZIRAN {
/*

            cg_objective.setMultiplier([&](const TVStack& x, TVStack& b) {
                b.setZero();

                grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
                    b.col(g.idx) = g.m * x.col(g.idx).cwiseProduct(Q[g.idx]).cwiseProduct(Q[g.idx]);
                });

                // use new_v to store x
                // G2P
                StdVector<TM> WTWD;
                WTWD.resize((int)particles.count);
                grid.iterateTouchedGrid([&](IV node, GridState<T, dim>& g) {
                    g.new_v = TV::Zero();
                });
                grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
                    g.new_v = x.col(g.idx);
                });
                tbb::parallel_for(0, (int)particles.count, [&](int i) {
                    TV& Xp = Xarray[i];
                    TM& Fn = (*Fn_pointer)[i];
                    BSplineWeights<T, dim> spline(Xp, dx);
                    WTWD[i] = TM::Zero();
                    grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                        TM tmp = dt * rho_scale * (g.new_v.cwiseProduct(Q[g.idx])) * dw.transpose() * Fn;
                        WTWD[i] += tmp.cwiseProduct(omega[i]).cwiseProduct(omega[i]).cwiseProduct(R[i]).cwiseProduct(R[i]);
                    });
                });
                // P2G
                grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
                    g.new_v = TV::Zero();
                });
                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];
                            TM& Fn = (*Fn_pointer)[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) {
                                TV tmp = dt * WTWD[i] * Fn.transpose() * dw;
                                g.new_v += tmp.cwiseProduct(Q[g.idx]);
                            });
                        }
                    });
                }
                grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
                    b.col(g.idx) += g.new_v;
                });
            });

            cg_objective.setPreconditioner([&](const TVStack& in, TVStack& out) {
                grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
                    g.new_v = g.m * (Q[g.idx].cwiseProduct(Q[g.idx]));
                });
                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];
                            TM& Fn = (*Fn_pointer)[i];
                            BSplineWeights<T, dim> spline(Xp, dx);
                            grid.iterateKernel(spline, particle_base_offset[i], [&](IV node, T w, TV dw, GridState<T, dim>& g) {
                                TM tmp = dt * dt * rho_scale * Q[g.idx] * dw.transpose() * Fn;
                                TM tmpooRR = tmp.cwiseProduct(omega[i]).cwiseProduct(R[i]).cwiseProduct(R[i]).cwiseProduct(omega[i]);
                                TV opt = tmpooRR * Fn.transpose() * dw;
                                g.new_v += opt.cwiseProduct(Q[g.idx]);
                            });
                        }
                    });
                }
                grid.iterateGrid([&](IV node, GridState<T, dim>& g) {
                    out.col(g.idx) = (in.col(g.idx)).cwiseQuotient(g.new_v);
                });
            });
 */
} // namespace ZIRAN

#endif