//
// Created by cflin on 4/20/23.
//

#include "Top3d.h"

namespace top {
    void Top3d::TopOptMainLoop() {
        Eigen::VectorXd xPhys_col(sp_mesh_->GetNumEles());
        xPhys_col.setConstant(sp_para_->volfrac);

        int loop = 0;
        double change = 1.0;
        double E0 = sp_material_->YM;
        double Emin = sp_material_->YM * sp_para_->E_factor;

        // precompute
        Eigen::VectorXd dv(sp_mesh_->GetNumEles());
        dv.setOnes();
        dv = H_ * (dv.array() / Hs_.array()).matrix().eval();

        Eigen::VectorXd ele_to_write = Eigen::VectorXd::Zero(sp_mesh_->GetLx() * sp_mesh_->GetLy() * sp_mesh_->GetLz());
        Eigen::VectorXi pixel_idx = sp_mesh_->GetPixelIdx();
        spdlog::info("end precompute");
        // clear output dir
        clear_dir(CMAKE_SOURCE_DIR "/output");
        // start iteration
        while (change > sp_para_->tol_x && loop < sp_para_->max_loop) {
            ++loop;
            Eigen::VectorXd sK = (sKe_ * (Emin + xPhys_col.array().pow(sp_para_->penal) *
                                                 (E0 - Emin)).matrix().transpose()).reshaped();
            auto v_tri = Vec2Triplet(iK_, jK_, sK);
            K_.setFromTriplets(v_tri.begin(), v_tri.end());
            IntroduceFixedDofs(K_, F_);
#ifdef USE_SUITESPARSE

            Eigen::CholmodSupernodalLLT<Eigen::SparseMatrix<double>> solver;
                solver.compute(K_);
                Eigen::VectorXd U = solver.solve(F_);

#else
            // dense direct method
            Eigen::MatrixXd denseK = Eigen::MatrixXd(K_);
            Eigen::VectorXd U = denseK.fullPivLu().solve(Eigen::VectorXd(F_));
#endif
            // compliance
            Eigen::VectorXd ce(sp_mesh_->GetNumEles());
            for (int i = 0; i < sp_mesh_->GetNumEles(); ++i) {
                Eigen::VectorXi dofs_in_ele_i = sp_mesh_->MapEleId2Dofs(i);
                Eigen::VectorXd Ue = U(dofs_in_ele_i);
                ce(i) = Ue.transpose() * Ke_ * Ue;
            }
            double c = ce.transpose() * (Emin + xPhys_col.array().pow(sp_para_->penal) * (E0 - Emin)).matrix();
            double v = xPhys_col.sum();

            Eigen::VectorXd dc =
                    -sp_para_->penal * (E0 - Emin) * xPhys_col.array().pow(sp_para_->penal - 1.0) * ce.array();

            // mma solver
            size_t num_constrants = 1;
            size_t num_variables = sp_mesh_->GetNumEles();
            auto mma = std::make_shared<MMASolver>(num_variables, num_constrants);
            Eigen::VectorXd variables_tmp = xPhys_col;
            double f0val = c;
            Eigen::VectorXd df0dx = dc;
            double fval = v - sp_mesh_->GetNumEles() * sp_para_->volfrac;
            Eigen::VectorXd dfdx = dv;
            static Eigen::VectorXd low_bounds = Eigen::VectorXd::Zero(sp_mesh_->GetNumEles());
            static Eigen::VectorXd up_bounds = Eigen::VectorXd::Ones(sp_mesh_->GetNumEles());

//                spdlog::info("mma update");
            mma->Update(variables_tmp.data(), df0dx.data(), &fval, dfdx.data(),
                        low_bounds.data(), up_bounds.data());

            change = (variables_tmp - xPhys_col).cwiseAbs().maxCoeff();
            xPhys_col = variables_tmp;
            spdlog::critical("Iter: {:3d}, Comp: {:.3f}, Vol: {:.2f}, Change: {:.3f}", loop, c, v, change);
#ifdef  WRITE_TENSOR
            // result
        ele_to_write(pixel_idx)=xPhys_col;
        Tensor3d ten_xPhys_to_write(sp_mesh_->GetLx()*sp_mesh_->GetLy()*sp_mesh_->GetLz(),1,1);
        for(int i=0;i<ele_to_write.size();++i){
            ten_xPhys_to_write(i,0,0)=ele_to_write(i);
        }
        ten_xPhys_to_write = ten_xPhys_to_write.reshape(
                Eigen::array<int, 3>{sp_mesh_->GetLx(), sp_mesh_->GetLy(), sp_mesh_->GetLz()});
        write_tensor3d(CMAKE_SOURCE_DIR "/output/txt/iter_" + std::to_string(loop) + ".txt", ten_xPhys_to_write);
        write_tensor3d_to_vtk(CMAKE_SOURCE_DIR "/output/vtk/iter_" + std::to_string(loop) + ".vtk",
                              ten_xPhys_to_write);
#endif
        }
        // result
        ele_to_write(pixel_idx) = xPhys_col;
        Tensor3d ten_xPhys_to_write(sp_mesh_->GetLx() * sp_mesh_->GetLy() * sp_mesh_->GetLz(), 1, 1);
        for (int i = 0; i < ele_to_write.size(); ++i) {
            ten_xPhys_to_write(i, 0, 0) = ele_to_write(i);
        }
        ten_xPhys_to_write = ten_xPhys_to_write.reshape(
                Eigen::array<int, 3>{sp_mesh_->GetLx(), sp_mesh_->GetLy(), sp_mesh_->GetLz()});
        write_tensor3d(CMAKE_SOURCE_DIR "/output/txt/iter_" + std::to_string(loop) + ".txt", ten_xPhys_to_write);
        write_tensor3d_to_vtk(CMAKE_SOURCE_DIR "/output/vtk/iter_" + std::to_string(loop) + ".vtk",
                              ten_xPhys_to_write);

    }

    void Top3d::Precompute() {
        Eigen::MatrixXi mat_ele2dofs = sp_mesh_->GetEleId2DofsMap();
        iK_ = Eigen::KroneckerProduct(mat_ele2dofs, Eigen::VectorXi::Ones(24)).transpose().reshaped();
        jK_ = Eigen::KroneckerProduct(mat_ele2dofs, Eigen::RowVectorXi::Ones(24)).transpose().reshaped();
        sp_material_->computeKe(0.5, 0.5, 0.5, sp_material_->D, Ke_);
        sKe_ = Ke_.reshaped();

        // precompute filter
        Eigen::VectorXi iH = Eigen::VectorXi::Ones(
                sp_mesh_->GetNumEles() * std::pow(2.0 * (std::ceil(sp_para_->r_min) - 1.0) + 1, 3));
        Eigen::VectorXi jH = iH;
        Eigen::VectorXd sH(iH.size());
        sH.setZero();
        int cnt = 0;
        Hs_ = Eigen::VectorXd(sp_mesh_->GetNumEles());

        int delta = std::ceil(sp_para_->r_min) - 1;
        for (int k = 0; k < sp_mesh_->GetLz(); ++k) {
            for (int j = 0; j < sp_mesh_->GetLy(); ++j) {
                for (int i = 0; i < sp_mesh_->GetLx(); ++i) {
                    int ele_id0 = sp_mesh_->MapEleCoord2Id((Eigen::MatrixXi(1, 3) << i, j, k).finished())(0);
                    if (ele_id0 == -1) {
                        continue;
                    }
                    for (int k2 = std::max(k - delta, 0); k2 <= std::min(k + delta, sp_mesh_->GetLz() - 1); ++k2) {
                        for (int j2 = std::max(j - delta, 0);
                             j2 <= std::min(j + delta, sp_mesh_->GetLy() - 1); ++j2) {
                            for (int i2 = std::max(i - delta, 0);
                                 i2 <= std::min(i + delta, sp_mesh_->GetLx() - 1); ++i2) {
                                int ele_id1 = sp_mesh_->MapEleCoord2Id(
                                        (Eigen::MatrixXi(1, 3) << i2, j2, k2).finished())(0);
                                if (ele_id1 == -1) {
                                    continue;
                                }
                                iH(cnt) = ele_id0;
                                jH(cnt) = ele_id1;
                                sH(cnt) = std::max(0.0,
                                                   sp_para_->r_min -
                                                   Eigen::Vector3d(i - i2, j - j2, k - k2).norm());
                                Hs_(ele_id0) += sH(cnt);
                                ++cnt;
                            }
                        }
                    }

                }

            }
        }
        std::vector<Eigen::Triplet<double>> v_tri = Vec2Triplet(iH, jH, sH);
        H_ = SpMat(sp_mesh_->GetNumEles(), sp_mesh_->GetNumEles());
        H_.setFromTriplets(v_tri.begin(), v_tri.end());
    }
} // top