//
// Created by cflin on 6/9/23.
//
#include <memory>
#include <nlohmann/json.hpp>
#include "Boundary.h"
#include "Mesh/HeatMesh.h"
#include "Util.h"
#include "ThermoelasticTop3d.h"
#include "FEA/MechanicalLinearFEA.h"
#include "FEA/ThermalLinearFEA.h"
#include "TensorWrapper.h"

int main() {
    using namespace da::sha;
    using namespace da::sha::top;
    using top::fs_path;
    using std::string;
    top::fs_path output_dir(OUTPUT_DIR);
    top::fs_path config_file(
            CMAKE_SOURCE_DIR "/examples/top-thermoelastic-compare-3d/config_Lshape.json");
    top::fs_path assets_dir(ASSETS_DIR);
    spdlog::info("Algo read from '{}'", config_file.string());
    spdlog::info("Algo output to '{}'", output_dir.string());
    spdlog::info("assets dir: '{}'", assets_dir.string());

    // read json
    std::ifstream f(config_file.c_str());
    if (!f) {
        spdlog::critical("f open fail!");
        exit(-7);
    }
    nlohmann::json j_config = nlohmann::json::parse(f);
    // read title
    std::string ex_name = j_config["TopologyOptimizationExample"];
    spdlog::critical("TopologyOptimizationExample: {}", ex_name);

    // set topology parameters
    auto para = std::make_shared<top::CtrlPara>();
    para->max_loop = j_config["topology"]["max_loop"];
    para->volfrac = j_config["topology"]["volfrac"];
    para->r_min = j_config["topology"]["r_min"];
    para->penal = j_config["topology"]["penal"];
    para->T_ref = j_config["topology"]["T_ref"];
    para->T_limit = j_config["topology"]["T_limit"];
    para->R_E = j_config["topology"]["R_E"];
    para->R_lambda = j_config["topology"]["R_lambda"];
    para->R_beta = j_config["topology"]["R_beta"];

    // set material parameters
    double E = j_config["material"]["E"];
    double Poisson_ratio = j_config["material"]["poisson_ratio"];
    double thermal_conductivity = j_config["material"]["thermal_conductivity"];
    double thermal_expansion_coefficient = j_config["material"]["thermal_expansion_coefficient"];
    auto material = std::make_shared<top::Material>(E, Poisson_ratio,
                                                    thermal_conductivity,
                                                    thermal_expansion_coefficient);
    // set fea
    auto sp_mech_fea = std::make_shared<top::MechanicalLinearFEA>(
            material);//  for mechanical
    auto sp_thermal_fea = std::make_shared<top::ThermalLinearFEA>(
            material);// for thermal

    // set mesh(regular)
    int len_x = j_config["model"]["regular_model"]["lx"];
    int len_y = j_config["model"]["regular_model"]["ly"];
    int len_z = j_config["model"]["regular_model"]["lz"];
    // NOTE: USER DEFINE GRID HERE!!!
    std::shared_ptr<top::Mesh> sp_mech_mesh;
    std::shared_ptr<top::HeatMesh> sp_thermal_mesh;
    if (ex_name == "Lshape") {
        // L-shape condition
        spdlog::critical("Using User Defined density!");
        top::Tensor3d L_shape_model(len_x, len_y, len_z);
        L_shape_model.setConstant(1);
        // set the initial voxel model
        for (int k = 0; k < len_z; ++k) {
            for (int j = 0; j < len_y; ++j) {
                for (int i = 0; i < len_x; ++i) {
                    if (j > len_y * 0.6 & k > len_z * 0.5) {
                        L_shape_model(i, j, k) = 0;
                    }
                }
            }
        }
        auto &model = L_shape_model;
        sp_mech_mesh = std::make_shared<top::Mesh>(len_x, len_y, len_z, model);
        sp_thermal_mesh = std::make_shared<top::HeatMesh>(len_x, len_y, len_z,
                                                          model);
    } else {
        sp_mech_mesh = std::make_shared<top::Mesh>(len_x, len_y, len_z);
        sp_thermal_mesh = std::make_shared<top::HeatMesh>(len_x, len_y, len_z);
    }
    // initialize Top3d
    auto sp_mech_top3d = std::make_shared<top::Top3d>(para, sp_mech_fea,
                                                      sp_mech_mesh);
    auto sp_ther_top3d = std::make_shared<top::Top3d>(para, sp_thermal_fea,
                                                      sp_thermal_mesh);

    // auxiliary class(Boundary) help to get boundary coordinates
    //      see the comments in Boundary.h for more information
    auto AddBoundaryMechanicalCondition = [&j_config](
            std::shared_ptr<top::Top3d> sp_mech_top3d,
            std::shared_ptr<top::Mesh> sp_mech_mesh) {
        top::Boundary mech_bdr(sp_mech_mesh);

        {
            //      add Dirichlet boundary, see the comments in Top3d::AddDBC for more information
            assert(j_config.count("mechanical_boundary_condition"));
            assert(j_config["mechanical_boundary_condition"].count("DBC"));
            assert(j_config["mechanical_boundary_condition"].count("NBC"));
            int DBCNum = j_config["mechanical_boundary_condition"]["DBC"].size();
            for (int _i = 0; _i < DBCNum; ++_i) {
                const auto &DBCI = j_config["mechanical_boundary_condition"]["DBC"][_i];
                Eigen::Vector3d minBBox(DBCI["min"][0], DBCI["min"][1],
                                        DBCI["min"][2]);
                Eigen::Vector3d maxBBox(DBCI["max"][0], DBCI["max"][1],
                                        DBCI["max"][2]);
                Eigen::Vector3i dir(DBCI["dir"][0], DBCI["dir"][1],
                                    DBCI["dir"][2]);
                top::Dir t_dir(minBBox, maxBBox, dir);

                sp_mech_top3d->AddDBC(
                        mech_bdr.GetChosenCoordsByRelativeAlignedBox(t_dir.box),
                        t_dir.dir);
            }

            //      add Neumann boundary, see the comments in Top3d::AddNBC for more information
            int NBCNum = j_config["mechanical_boundary_condition"]["NBC"].size();
            for (int _i = 0; _i < NBCNum; ++_i) {
                const auto &NBCI = j_config["mechanical_boundary_condition"]["NBC"][_i];
                Eigen::Vector3d minBBox(NBCI["min"][0], NBCI["min"][1],
                                        NBCI["min"][2]);
                Eigen::Vector3d maxBBox(NBCI["max"][0], NBCI["max"][1],
                                        NBCI["max"][2]);
                Eigen::Vector3d val(NBCI["val"][0], NBCI["val"][1],
                                    NBCI["val"][2]);
                top::Neu t_neu(minBBox, maxBBox, val);

                Eigen::MatrixXi coords = mech_bdr.GetChosenCoordsByRelativeAlignedBox(
                        t_neu.box);
                sp_mech_top3d->AddNBC(coords, t_neu.val / coords.rows());
            }
        }

    };
    auto AddBoundaryThermalCondition = [&j_config](
            std::shared_ptr<top::Top3d> sp_ther_top3d,
            std::shared_ptr<top::HeatMesh> sp_ther_mesh) {
        top::Boundary thermal_bdr(sp_ther_mesh);
        {
            //      add Dirichlet boundary, see the comments in Top3d::AddDBC for more information
            assert(j_config.count("thermal_boundary_condition"));
            assert(j_config["thermal_boundary_condition"].count("DBC"));
            assert(j_config["thermal_boundary_condition"].count("NBC"));
            int DBCNum = j_config["thermal_boundary_condition"]["DBC"].size();
            for (int _i = 0; _i < DBCNum; ++_i) {
                const auto &DBCI = j_config["thermal_boundary_condition"]["DBC"][_i];
                Eigen::Vector3d minBBox(DBCI["min"][0], DBCI["min"][1],
                                        DBCI["min"][2]);
                Eigen::Vector3d maxBBox(DBCI["max"][0], DBCI["max"][1],
                                        DBCI["max"][2]);
                top::Dir t_dir(minBBox, maxBBox, Eigen::Vector3i(1, 0, 0));

                sp_ther_top3d->AddDBC(
                        thermal_bdr.GetChosenCoordsByRelativeAlignedBox(
                                t_dir.box),
                        t_dir.dir, DBCI["temperature"]);
            }

            //      add Neumann boundary, see the comments in Top3d::AddNBC for more information
            int NBCNum = j_config["thermal_boundary_condition"]["NBC"].size();
            for (int _i = 0; _i < NBCNum; ++_i) {
                const auto &NBCI = j_config["thermal_boundary_condition"]["NBC"][_i];
                Eigen::Vector3d minBBox(NBCI["min"][0], NBCI["min"][1],
                                        NBCI["min"][2]);
                Eigen::Vector3d maxBBox(NBCI["max"][0], NBCI["max"][1],
                                        NBCI["max"][2]);
                Eigen::Vector3d val(NBCI["heat_flux"], 0, 0);
                top::Neu t_neu(minBBox, maxBBox, val);

                Eigen::MatrixXi coords = thermal_bdr.GetChosenCoordsByRelativeAlignedBox(
                        t_neu.box);
                sp_ther_top3d->AddNBC(coords, t_neu.val / coords.rows());
            }
        }
    };

    AddBoundaryMechanicalCondition(sp_mech_top3d, sp_mech_mesh);
    AddBoundaryThermalCondition(sp_ther_top3d, sp_thermal_mesh);

    // process topology optimization
    spdlog::critical("start to mechanical top opt ...");
    top::Tensor3d t_me_rho = sp_mech_top3d->TopOptMainLoop();
    {
        spdlog::critical("extract compliance and volume each iteration ...");
        // extract compliance and volume each iteration
        fs_path compliance_path =
                output_dir / "txt" / ex_name /
                (ex_name + "_MeTop" + "_compliance.txt");
        WriteStdVector(compliance_path, sp_mech_top3d->v_compliance_);
        spdlog::info("write compliance txt to: {}", compliance_path.c_str());

        fs_path volume_path =
                output_dir / "txt" / ex_name /
                (ex_name + "_MeTop" + "_volume.txt");
        WriteStdVector(volume_path, sp_mech_top3d->v_volume_);
        spdlog::info("write volume txt to: {}", volume_path.c_str());

        // extract rho (txt and vtk)
        fs_path rho_txt_path =
                output_dir / "txt" / ex_name /
                (ex_name + "_MeTop" + "_rho.txt");
        write_tensor3d(rho_txt_path, t_me_rho, sp_mech_mesh->GetOrigin(),
                       sp_mech_mesh->GetOrigin() + sp_mech_mesh->GetLenBox());
        spdlog::info("write density txt to: {}", rho_txt_path.c_str());

        fs_path rho_vtk_path =
                output_dir / "vtk" / ex_name /
                (ex_name + "_MeTop" + "_rho.vtk");
        WriteTensorToVtk(rho_vtk_path, t_me_rho, sp_mech_mesh);
        spdlog::info("write density vtk to: {}", rho_vtk_path.c_str());
    }

    spdlog::critical("start to mechanical thermal top opt ...");
    // init thermoelastic top3d
    top::ThermoelasticTop3d mech_ther_top3d(sp_mech_top3d, sp_ther_top3d);
    top::Tensor3d t_meth_rho = mech_ther_top3d.TopOptMainLoop();
    {
        spdlog::critical("extract compliance and volume each iteration ...");
        // extract compliance and volume each iteration
        fs_path compliance_path =
                output_dir / "txt" / ex_name /
                (ex_name + "_MeThTop" + "_compliance.txt");
        WriteStdVector(compliance_path, mech_ther_top3d.v_compliance_);
        spdlog::info("write compliance txt to: {}", compliance_path.c_str());

        fs_path volume_path =
                output_dir / "txt" / ex_name /
                (ex_name + "_MeThTop" + "_volume.txt");
        WriteStdVector(volume_path, mech_ther_top3d.v_volume_);
        spdlog::info("write volume txt to: {}", volume_path.c_str());

        // extract rho (txt and vtk)
        fs_path rho_txt_path =
                output_dir / "txt" / ex_name /
                (ex_name + "_MeThTop" + "_rho.txt");
        write_tensor3d(rho_txt_path, t_meth_rho, sp_mech_mesh->GetOrigin(),
                       sp_mech_mesh->GetOrigin() + sp_mech_mesh->GetLenBox());
        spdlog::info("write density txt to: {}", rho_txt_path.c_str());

        fs_path rho_vtk_path =
                output_dir / "vtk" / ex_name /
                (ex_name + "_MeThTop" + "_rho.vtk");
        WriteTensorToVtk(rho_vtk_path, t_meth_rho, sp_mech_mesh);
        spdlog::info("write density vtk to: {}", rho_vtk_path.c_str());
    }

//    // extract rho (txt and vtk)
//    fs_path rho_txt_path = output_dir / "txt" / (ex_name + "_rho.txt");
//    write_tensor3d(rho_txt_path, ten_rho, sp_mech_mesh->GetOrigin(),
//                   sp_mech_mesh->GetOrigin() + sp_mech_mesh->GetLenBox());
//    spdlog::info("write density txt to: {}", rho_txt_path.c_str());
//
//    fs_path rho_vtk_path = output_dir / "vtk" / (ex_name + "_rho.vtk");
//    WriteTensorToVtk(rho_vtk_path, ten_rho, sp_mech_mesh);
//    spdlog::info("write density vtk to: {}", rho_vtk_path.c_str());

//    // extract temperature(vtk)
//    fs_path T_vtk_path = output_dir / "vtk" / (ex_name + "_T.vtk");
//    WriteUToVtk(T_vtk_path, ther_top3d.GetTemperature(), sp_thermal_mesh);
//    spdlog::info("write temperature vtk to: {}", T_vtk_path.c_str());
//
//    // extract displacement(vtk)
//    fs_path U_vtk_path = output_dir / "vtk" / (ex_name + "_U.vtk");
//    WriteUToVtk(U_vtk_path, ther_top3d.GetNormedDisplacement(), sp_mech_mesh);
//    spdlog::info("write displacement norm vtk to: {}", U_vtk_path.c_str());
//
//    // extract stress field(txt or vtk)
//    auto t_von_stress = ther_top3d.GetVonStress();
//    fs_path von_stress_txt_path =
//            output_dir / "txt" / (ex_name + "_von_stress.txt");
//    write_tensor3d(von_stress_txt_path, t_von_stress, sp_mech_mesh->GetOrigin(),
//                   sp_mech_mesh->GetOrigin() + sp_mech_mesh->GetLenBox());
//    spdlog::info("write von stress txt to: {}", von_stress_txt_path.c_str());
//
//    fs_path von_stress_vtk_path =
//            output_dir / "vtk" / (ex_name + "_von_stress.vtk");
//    WriteTensorToVtk(von_stress_vtk_path, t_von_stress, sp_mech_mesh);
//    spdlog::info("write von stress vtk to: {}", von_stress_vtk_path.c_str());


    // postprocess--------------------------------------------------------------
    auto MechanicalSimulation = [&](const top::Tensor3d &t_rho,
                                    const string &suffix) {
        sp_mech_mesh = std::make_shared<top::Mesh>(t_rho.dimension(0),
                                                   t_rho.dimension(1),
                                                   t_rho.dimension(2), t_rho);
        // postprocess: simulation
        // MeSim
        {
            auto sp_mech_top3d_sim = std::make_shared<top::Top3d>(para,
                                                                  sp_mech_fea,
                                                                  sp_mech_mesh);
            sp_mech_top3d_sim->TopOptMainLoop(true);

            // extract displacement(vtk)
            fs_path U_vtk_path = output_dir / "vtk" /
                                 (ex_name + "_MeSim" + "_U" + suffix +
                                  ".vtk");
            WriteUToVtk(U_vtk_path, sp_mech_top3d_sim->GetNormedDisplacement(),
                        sp_mech_mesh);
            spdlog::info("write displacement norm vtk to: {}",
                         U_vtk_path.c_str());

            // extract stress field(txt and vtk)
            auto t_von_stress = sp_mech_top3d_sim->GetVonStress();
            fs_path von_stress_vtk_path =
                    output_dir / "vtk" /
                    (ex_name + "_MeSim" + "_von_stress" + suffix + ".vtk");
            WriteNodeToVtk(von_stress_vtk_path, t_von_stress, sp_mech_mesh);
            spdlog::info("write von stress vtk to: {}",
                         von_stress_vtk_path.c_str());

            // extract compliance
            fs_path compliance_path =
                    output_dir / "txt" /
                    (ex_name + "_MeSim" + "_compliance" + suffix + ".txt");
            WriteStdVector(compliance_path, mech_ther_top3d.v_compliance_);
            spdlog::info("write compliance txt to: {}",
                         compliance_path.c_str());

        }
    };
    auto MechanicalThermalSimulation = [&output_dir, &ex_name, &para, &sp_mech_fea, &sp_thermal_fea, &AddBoundaryThermalCondition, &AddBoundaryMechanicalCondition](
            const top::Tensor3d &t_rho,
            const string &suffix) {
        auto sp_mech_mesh_sim = std::make_shared<top::Mesh>(t_rho.dimension(0),
                                                            t_rho.dimension(1),
                                                            t_rho.dimension(2),
                                                            t_rho);
        auto sp_thermal_mesh_sim = std::make_shared<top::HeatMesh>(
                t_rho.dimension(0),
                t_rho.dimension(1),
                t_rho.dimension(2),
                t_rho);
        auto sp_mech_top3d_sim = std::make_shared<top::Top3d>(para, sp_mech_fea,
                                                              sp_mech_mesh_sim);
        auto sp_thermal_top3d_sim = std::make_shared<top::Top3d>(para,
                                                                 sp_thermal_fea,
                                                                 sp_thermal_mesh_sim);
        AddBoundaryMechanicalCondition(sp_mech_top3d_sim, sp_mech_mesh_sim);
        AddBoundaryThermalCondition(sp_thermal_top3d_sim, sp_thermal_mesh_sim);
        top::ThermoelasticTop3d ther_top3d_sim(sp_mech_top3d_sim,
                                               sp_thermal_top3d_sim);
        ther_top3d_sim.TopOptMainLoop(true);
        {
            // extract clamped rho (vtk)
            fs_path rho_vtk_path = output_dir / "vtk" / ex_name /
                                   (ex_name + suffix + "_rho" +
                                    ".vtk");
            auto t_temp = t_rho;
            t_temp.SetConst(1);
            WriteTensorToVtk(rho_vtk_path, t_temp, sp_mech_mesh_sim);
            spdlog::info("write clamped density vtk to: {}",
                         rho_vtk_path.c_str());

            // extract temperature (vtk)
            fs_path T_vtk_path = output_dir / "vtk" / ex_name /
                                 (ex_name + suffix + "_T" +
                                  ".vtk");
            WriteNodeToVtk(T_vtk_path, ther_top3d_sim.GetTemperature(),
                           sp_thermal_mesh_sim);
            spdlog::info("write temperature vtk to: {}", T_vtk_path.c_str());

            // extract displacement (vtk)
            fs_path U_vtk_path = output_dir / "vtk" / ex_name /
                                 (ex_name + suffix + "_U" +
                                  ".vtk");
            WriteNodeToVtk(U_vtk_path, ther_top3d_sim.GetNormedDisplacement(),
                           sp_mech_mesh_sim);
            spdlog::info("write displacement norm vtk to: {}",
                         U_vtk_path.c_str());

            // extract stress field(txt or vtk)
            auto t_von_stress = ther_top3d_sim.GetVonStress();
            fs_path von_stress_vtk_path =
                    output_dir / "vtk" / ex_name /
                    (ex_name + suffix + "_von_stress" +
                     ".vtk");
            WriteNodeToVtk(von_stress_vtk_path, t_von_stress, sp_mech_mesh_sim);
            spdlog::info("write von stress vtk to: {}",
                         von_stress_vtk_path.c_str());

            // extract compliance
            fs_path compliance_path =
                    output_dir / "txt" / ex_name /
                    (ex_name + suffix + "_compliance" + ".txt");
            WriteStdVector(compliance_path, ther_top3d_sim.v_compliance_);
            spdlog::info("write compliance txt to: {}",
                         compliance_path.c_str());
        }
    };
    auto ClampAndSimulation = [&](const top::Tensor3d &res_rho,
                                  const std::string &suffix) {
        // clamp density to 0 or 1
        // set different thresholds to pick a suitable density result
        for (double threshold = 0.2;
             threshold < 0.5 + 0.0001; threshold += 0.05) {
            std::string str_thresh =
                    "_thresh" + std::to_string((int) (threshold * 100));
            top::Tensor3d t_rho = res_rho;
            for (int k = 0; k < t_rho.dimension(2); ++k) {
                for (int j = 0; j < t_rho.dimension(1); ++j) {
                    for (int i = 0; i < t_rho.dimension(0); ++i) {
                        t_rho(i, j, k) = t_rho(i, j, k) >= threshold ? 1 : 0;
                    }
                }
            }

            try {
                // mech ther sim
                MechanicalThermalSimulation(t_rho, suffix + str_thresh);
            } catch (std::exception &e) {
                spdlog::error(e.what() + std::string(" skip") + str_thresh);
            }
        }
    };

    {
        spdlog::critical("postprocess: mechanical result simulation ...");
        // SIM: Mechanical result -> clamped -> MeThSim
        ClampAndSimulation(t_me_rho, "_MeSim");
    }

    {
        spdlog::critical(
                "postprocess: mechanical thermal result simulation ...");
        // SIM: Mechanical thermal result -> clamped -> MeThSim
        ClampAndSimulation(t_meth_rho, "_MeThSim");
    }


}