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

#include "Top3d.h"
#include <spdlog/spdlog.h>
#include <cassert>
#include "Eigen/src/Core/Matrix.h"
#include "Util.h"
#include "LinearSolver/Amgcl.h"

#ifdef USE_AMGCL_CUDA
#include "LinearSolver/AmgclCuda.h"
#endif
namespace da::sha {
    namespace top {
        Tensor3d
        Top3d::TopOptMainLoop(bool only_simulation, bool save_internal_rho) {
            Eigen::VectorXd xPhys_col(sp_mesh_->GetNumEles());
            Eigen::VectorXi chosen_ele_id(sp_mesh_->GetChosenEleIdx());
            bool flg_chosen = chosen_ele_id.size() != 0;
            if (!flg_chosen) {
                // no part chosen
                xPhys_col.setConstant(only_simulation ? 1 : sp_para_->volfrac);
            } else {
                // pick chosen part to sim
                xPhys_col = sp_mesh_->GetInitEleRho();
                xPhys_col(chosen_ele_id).setConstant(
                        only_simulation ? 1 : sp_para_->volfrac);
            }

            int loop = 0;
            double change = 1.0;
            double E0 = sp_fea_->sp_material_->E;
            double Emin = E0 * 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_->GetGlobalIdxOfEleInUse();
            spdlog::info("end precompute");
            //        // clear output dir
            //        clear_dir(CMAKE_SOURCE_DIR "/output");

            LOG_SOLVER

            // set 0 to rho of unchosen part
            assert(xPhys_col.size());
            Eigen::VectorXi continue_idx =
                    Eigen::VectorXi::LinSpaced(xPhys_col.size(), 0,
                                               xPhys_col.size() - 1);
            Eigen::VectorXi unchosen_idx = flg_chosen ? SetDifference(
                    continue_idx, chosen_ele_id)
                                                      : Eigen::VectorXi();

            auto RhoVec2Ten = [&](Eigen::VectorXd rho) -> Tensor3d {
                rho(unchosen_idx).setZero();
                Eigen::VectorXd ele_to_write =
                        Eigen::VectorXd::Zero(
                                sp_mesh_->GetLx() * sp_mesh_->GetLy() *
                                sp_mesh_->GetLz());
                ele_to_write(pixel_idx) = rho;
                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 < Eigen::DenseIndex, 3 > {
                                sp_mesh_->GetLx(), sp_mesh_->GetLy(),
                                sp_mesh_->GetLz()});
                return ten_xPhys_to_write;
            };

            if (save_internal_rho) {
                v_ten_rho_.emplace_back(RhoVec2Ten(xPhys_col));
            }

            // start iteration
            while (change > sp_para_->tol_x && loop < sp_para_->max_loop) {
                ++loop;
                // filter
                xPhys_col =
                        H_ * (xPhys_col.array() / Hs_.array()).matrix().eval();
                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_);
                INIT_SOLVER(solver, K_);
                U_ = solver.solve(F_);
                // 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 = flg_chosen ? xPhys_col(chosen_ele_id).sum()
                                      : xPhys_col.sum();
                // sensitivity
                Eigen::VectorXd dc_drho =
                        -sp_para_->penal * (E0 - Emin) *
                        xPhys_col.array().pow(sp_para_->penal - 1.0) *
                        ce.array();
                Eigen::VectorXd dc_dx = drho_dx_ * dc_drho;

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

                mma->Update(variables_tmp.data(), df0dx.data(), &fval,
                            dfdx.data(),
                            low_bounds.data(),
                            up_bounds.data());
                if (flg_chosen) {
                    change = (variables_tmp -
                              xPhys_col(chosen_ele_id)).cwiseAbs().maxCoeff();
                    xPhys_col(chosen_ele_id) = variables_tmp;
                } else {
                    change = (variables_tmp - xPhys_col).cwiseAbs().maxCoeff();
                    xPhys_col = variables_tmp;
                }

                spdlog::info(
                        "Iter: {:3d}, Comp: {:.3e}, Vol: {:.2f}, Change: {:f}",
                        loop, c, v,
                        change);

                v_compliance_.push_back(c);
                v_volume_.push_back(v);

                if (only_simulation) {
                    break;
                }
                if (save_internal_rho) {
                    v_ten_rho_.emplace_back(RhoVec2Ten(xPhys_col));
                }
            }

            // result
            rho_ = only_simulation ? sp_mesh_->GetInitEleRho() : xPhys_col;

            return RhoVec2Ten(rho_);
        }

        std::vector <Tensor3d>
        Top3d::GetTensorOfStress(const Eigen::VectorXd &which_col_of_stress) {
            Eigen::VectorXd ele_to_write =
                    Eigen::VectorXd::Zero(
                            sp_mesh_->GetLx() * sp_mesh_->GetLy() *
                            sp_mesh_->GetLz());
            Eigen::VectorXi pixel_idx = sp_mesh_->GetGlobalIdxOfEleInUse();
            // stress
            Eigen::MatrixXd mat_stress(sp_mesh_->GetNumEles(), 6);
            Eigen::MatrixXd B = sp_fea_->computeBe({0, 0, 0});
            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);
                mat_stress.row(i) = rho_(i) * sp_fea_->computeD() * B * Ue;
            }
            // fill
            std::vector <Tensor3d> vt;
            for (int i = 0; i < which_col_of_stress.size(); ++i) {
                ele_to_write(pixel_idx) = mat_stress.col(
                        which_col_of_stress(i));
                vt.push_back(GetTensorFromCol(ele_to_write));
            }
            return vt;
        }


        Eigen::VectorXd Top3d::GetVonStress() {
            Eigen::VectorXd ele_to_write =
                    Eigen::VectorXd::Zero(
                            sp_mesh_->GetLx() * sp_mesh_->GetLy() *
                            sp_mesh_->GetLz());
            Eigen::VectorXi pixel_idx = sp_mesh_->GetGlobalIdxOfEleInUse();
            // stress
            Eigen::VectorXd node_to_write = Eigen::VectorXd::Zero(
                    sp_mesh_->GetNumNodes());

            std::vector <Eigen::MatrixXd> v_B(8);
            static const Eigen::MatrixXd delta_coord = (Eigen::MatrixXd(8, 3) <<
                                                                              0, 0, 0,
                    1, 0, 0,
                    1, 1, 0,
                    0, 1, 0,
                    0, 0, 1,
                    1, 0, 1,
                    1, 1, 1,
                    0, 1, 1
            ).finished();
            for (int i = 0; i < delta_coord.rows(); ++i) {
                v_B[i] = sp_fea_->computeBe(delta_coord.row(i).array() - 0.5);
            }
            auto computeVonStress = [&](Eigen::VectorXd stress) -> double {
                double x = stress(0);
                double y = stress(1);
                double z = stress(2);
                double xy = stress(3);
                double yz = stress(4);
                double zx = stress(5);
                return sqrt(0.5 * (x * x + y * y + z * z) +
                            3 * (xy * xy + yz * yz + zx * zx));
            };
            for (int i = 0; i < sp_mesh_->GetNumEles(); ++i) {
                Eigen::VectorXi dofs_in_ele_i = sp_mesh_->MapEleId2Dofs(
                        i);// 1x 24
                Eigen::VectorXd Ue = U_(dofs_in_ele_i);
                Eigen::VectorXi node_id_in_ele_i = sp_mesh_->MapEleId2NodeIds(
                        i);// 1x8
                for (int j = 0; j < node_id_in_ele_i.size(); ++j) {
                    int i_node = node_id_in_ele_i(j);
                    if (node_to_write(i_node) != 0) {
                        continue;
                    }
                    Eigen::MatrixXd B = v_B[j];
                    Eigen::VectorXd s6 = rho_(i) * sp_fea_->computeD() * B * Ue;
                    node_to_write(i_node) = computeVonStress(s6);
                }
            }

            return node_to_write;
        }

        Tensor3d Top3d::GetTensorFromCol(const Eigen::VectorXd &proprty_col) {
            Tensor3d ten_prop_to_write(
                    sp_mesh_->GetLx() * sp_mesh_->GetLy() * sp_mesh_->GetLz(),
                    1,
                    1);
            assert(proprty_col.size() == ten_prop_to_write.size());
            for (int i = 0; i < proprty_col.size(); ++i) {
                ten_prop_to_write(i, 0, 0) = proprty_col(i);
            }
            ten_prop_to_write = ten_prop_to_write.reshape(
                    Eigen::array < Eigen::DenseIndex, 3 > {sp_mesh_->GetLx(),
                                                           sp_mesh_->GetLy(),
                                                           sp_mesh_->GetLz()});
            return ten_prop_to_write;
        }

        void Top3d::Precompute() {
            Eigen::MatrixXi mat_ele2dofs = sp_mesh_->GetEleId2DofsMap();
            int dofs_each_ele = sp_mesh_->Get_DOFS_EACH_ELE();// 24 for mathe; 8 for heat
            iK_ = Eigen::KroneckerProduct(mat_ele2dofs,
                                          Eigen::VectorXi::Ones(dofs_each_ele))
                    .transpose()
                    .reshaped();
            jK_ = Eigen::KroneckerProduct(mat_ele2dofs,
                                          Eigen::RowVectorXi::Ones(
                                                  dofs_each_ele))
                    .transpose()
                    .reshaped();
            Ke_ = sp_fea_->computeKe(1.0);
            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::Zero(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());

            Eigen::VectorXi i_Hs = Eigen::VectorXi::LinSpaced(Hs_.size(), 0,
                                                              Hs_.size());
            Eigen::SparseMatrix<double> sp_inv_Hs(i_Hs.size(), i_Hs.size());
            auto v_inv_Hs_tri = Vec2Triplet(i_Hs, i_Hs,
                                            Eigen::VectorXd(1.0 / Hs_.array()));
            sp_inv_Hs.setFromTriplets(v_inv_Hs_tri.begin(), v_inv_Hs_tri.end());
            drho_dx_ = sp_inv_Hs * H_.transpose();

        }
    }  // namespace top
}  // namespace da::sha