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

#ifndef TOP3D_UTIL_H
#define TOP3D_UTIL_H
#include <unsupported/Eigen/CXX11/Tensor>
#include <Eigen/Eigen>
#include "TensorWrapper.h"
#include <fstream>
#include <memory>
#include <vtkSmartPointer.h>
#include <vtkStructuredGrid.h>
#include <vtkXMLStructuredGridWriter.h>
#include <vtkPoints.h>
#include <vtkPointData.h>
#include <vtkFloatArray.h>
#include <vtkStructuredGridWriter.h>
#include <Eigen/Dense>

#include <iostream>
namespace top {
    using Tensor3d=TensorWrapper<double,3>;
    using Tensor3i=TensorWrapper<int,3>;
    using Eigen::all;
    using SpMat=Eigen::SparseMatrix<double>;

    template<typename Scalar>
    inline std::vector<Eigen::Triplet<Scalar>> Vec2Triplet(const Eigen::VectorXi& I,const Eigen::VectorXi& J,const Eigen::Matrix<Scalar,-1,1>& V){
        std::vector<Eigen::Triplet<Scalar>> v_tri;
        for(int i=0;i<I.size();++i){
            v_tri.push_back({I(i),J(i),V(i)});
        }
        return v_tri;
    }

    inline void write_tensor3d(const std::string &save_path,  const Tensor3d& t3){
        std::ofstream of(save_path);
        of<<t3.dimension(0)<<'\t'<<t3.dimension(1)<<'\t'<<t3.dimension(2)<<std::endl;
        for(int k=0;k<t3.dimension(2);++k){
            for(int j=0;j<t3.dimension(1);++j){
                for(int i=0;i<t3.dimension(0);++i){
                    of<<t3(i,j,k)<<'\t';
                }
                of<<'\n';
            }
            of<<'\n';
        }
    }

    inline Tensor3d read_tensor3d(const std::string &read_path){
        std::ifstream iif(read_path);
        int lx,ly,lz;
        iif>>lx>>ly>>lz;
        Tensor3d t3(lx,ly,lz);
        for(int k=0;k<t3.dimension(2);++k){
            for(int j=0;j<t3.dimension(1);++j){
                for(int i=0;i<t3.dimension(0);++i){
                    iif>>t3(i,j,k);
                }
            }
        }
        return t3;
    }

    inline void write_tensor3d_to_vtk(const std::string &vtk_path, const Tensor3d &t3){
        // 将Eigen::Tensor变量转换为vtkStructuredGrid对象
        vtkSmartPointer<vtkStructuredGrid> grid = vtkSmartPointer<vtkStructuredGrid>::New();
        grid->SetDimensions(t3.dimension(0), t3.dimension(1), t3.dimension(2));
        vtkSmartPointer<vtkPoints> points = vtkSmartPointer<vtkPoints>::New();
        for (int k = 0; k < t3.dimension(2); k++) {
            for (int j = 0; j < t3.dimension(1); j++) {
                for (int i = 0; i < t3.dimension(0); i++) {
                    points->InsertNextPoint(i, j, k);
                }
            }
        }
        grid->SetPoints(points);

        // 将Eigen::Tensor变量的值添加到网格中
        vtkSmartPointer<vtkFloatArray> values = vtkSmartPointer<vtkFloatArray>::New();
        values->SetNumberOfTuples(t3.size());
        values->SetNumberOfComponents(1);
        for (int k = 0; k < t3.dimension(2); k++) {
            for (int j = 0; j < t3.dimension(1); j++) {
                for (int i = 0; i < t3.dimension(0); i++) {
                    int index = i + j * t3.dimension(0) + k * t3.dimension(0) * t3.dimension(1);
                    values->SetTuple1(index, t3(i,j,k));
                }
            }
        }
        grid->GetPointData()->SetScalars(values);


        // 将网格写入VTK文件
        vtkSmartPointer<vtkStructuredGridWriter> writer = vtkSmartPointer<vtkStructuredGridWriter>::New();
        writer->SetFileName(vtk_path.c_str());
        writer->SetInputData(grid);
        writer->Write();

    }

} // top


#endif //TOP3D_UTIL_H