//
// 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 top::fs_path;
    using std::string;
    top::fs_path output_dir(OUTPUT_DIR);
    top::fs_path config_file(CONFIG_FILE);
    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("asserts 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);

    // 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!!! //TODO
    std::shared_ptr<top::Mesh> sp_mech_mesh;
    std::shared_ptr<top::HeatMesh> sp_thermal_mesh;
    std::string condition = j_config["model"]["predefined_condition"];
    if (condition == "L_shape") {
        // L-shape condition
        spdlog::info("Using {}!", condition);
        top::Tensor3d L_shape_model(len_x, len_y, len_z);
        L_shape_model.setConstant(1);
        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.5 && 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_thermal_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
    top::Boundary mech_bdr(sp_mech_mesh);
    top::Boundary thermal_bdr(sp_thermal_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());
        }
    }

    {
        //      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_thermal_top3d->AddDBC(mech_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 = mech_bdr.GetChosenCoordsByRelativeAlignedBox(t_neu.box);
            sp_thermal_top3d->AddNBC(coords, t_neu.val / coords.rows());
        }
    }


    // init thermoelastic top3d
    top::ThermoelasticTop3d ther_top3d(sp_mech_top3d, sp_thermal_top3d);

    // process topology optimization
    top::Tensor3d ten_rho = ther_top3d.TopOptMainLoop();

    // extract txt or vtk
    write_tensor3d(output_dir / "txt" / "top_mach_regular_field_matrix.txt", ten_rho,
                   sp_mech_mesh->GetOrigin(),
                   sp_mech_mesh->GetOrigin() + sp_mech_mesh->GetLenBox());
    WriteTensorToVtk(output_dir / "vtk" / "top_mach_regular_field_matrix.vtk", ten_rho,
                     sp_mech_mesh);
    WriteUToVtk(output_dir / "vtk" / "top_thermoelastic_regular_U.vtk", ther_top3d.GetU(),
                sp_mech_mesh);

}