//
//  Created by Wei Chen on 3/31/22
//

#ifndef LinSysSolver_hpp
#define LinSysSolver_hpp

#include <Eigen/Eigen>
#include <Eigen/Sparse>

#include <set>
#include <map>
#include <unordered_map>
#include <iostream>

#ifdef USE_TBB
#include <tbb/tbb.h>
#endif

namespace ssim {

template <typename vectorTypeI, typename vectorTypeS>
class LinSysSolver {
protected:
    int numRows;
    Eigen::VectorXi ia, ja;
    std::vector<std::map<int, int>> IJ2aI;
    Eigen::VectorXd a;

public:
    virtual ~LinSysSolver(void){};

public:
    virtual void set_pattern(const std::vector<std::set<int>>& vNeighbor)
    {
        numRows = static_cast<int>(vNeighbor.size()) * DIM_;
        ia.resize(numRows + 1);
        ia[0] = 1; // 1 + nnz above row i
        ja.resize(0); // colI of each element
        IJ2aI.resize(0); // map from matrix index to ja index
        IJ2aI.resize(numRows);

        std::vector<Eigen::VectorXi> ja_v(vNeighbor.size());
        std::vector<int> rowNNZ(numRows);
#ifdef USE_TBB
        tbb::parallel_for(0, (int)vNeighbor.size(), 1, [&](int vI)
#else
        for (int vI = 0; vI < vNeighbor.size(); ++vI)
#endif
            {
                ja_v[vI].resize((vNeighbor[vI].size() + 1) * DIM_);

                ja_v[vI][0] = vI * DIM_;
                ja_v[vI][1] = ja_v[vI][0] + 1;
                IJ2aI[ja_v[vI][0]][ja_v[vI][0]] = 0;
                IJ2aI[ja_v[vI][0]][ja_v[vI][1]] = 1;
                if constexpr (DIM_ == 3) {
                    ja_v[vI][2] = ja_v[vI][0] + 2;
                    IJ2aI[ja_v[vI][0]][ja_v[vI][2]] = 2;
                }

                int nnz = DIM_;
                for (const auto& nbVI : vNeighbor[vI]) {
                    if (nbVI > vI) {
                        ja_v[vI][nnz] = nbVI * DIM_;
                        ja_v[vI][nnz + 1] = ja_v[vI][nnz] + 1;
                        IJ2aI[ja_v[vI][0]][ja_v[vI][nnz]] = nnz;
                        IJ2aI[ja_v[vI][0]][ja_v[vI][nnz + 1]] = nnz + 1;
                        if constexpr (DIM_ == 3) {
                            ja_v[vI][nnz + 2] = ja_v[vI][nnz] + 2;
                            IJ2aI[ja_v[vI][0]][ja_v[vI][nnz + 2]] = nnz + 2;
                        }
                        nnz += DIM_;
                    }
                }

                rowNNZ[ja_v[vI][0]] = nnz;
                if constexpr (DIM_ == 2) {
                    ja_v[vI].conservativeResize(nnz * DIM_ - 1);
                    ja_v[vI].tail(nnz - 1) = ja_v[vI].segment(1, nnz - 1);

                    IJ2aI[ja_v[vI][0] + 1] = IJ2aI[ja_v[vI][0]];
                    IJ2aI[ja_v[vI][0] + 1].erase(ja_v[vI][0]);

                    rowNNZ[ja_v[vI][0] + 1] = nnz - 1;
                }
                else {
                    ja_v[vI].conservativeResize(nnz * DIM_ - 3);
                    ja_v[vI].segment(nnz, nnz - 1) = ja_v[vI].segment(1, nnz - 1);
                    ja_v[vI].tail(nnz - 2) = ja_v[vI].segment(2, nnz - 2);

                    IJ2aI[ja_v[vI][0] + 1] = IJ2aI[ja_v[vI][0]];
                    IJ2aI[ja_v[vI][0] + 1].erase(ja_v[vI][0]);
                    IJ2aI[ja_v[vI][0] + 2] = IJ2aI[ja_v[vI][0] + 1];
                    IJ2aI[ja_v[vI][0] + 2].erase(ja_v[vI][0] + 1);

                    rowNNZ[ja_v[vI][0] + 1] = nnz - 1;
                    rowNNZ[ja_v[vI][0] + 2] = nnz - 2;
                }
            }
#ifdef USE_TBB
        );
#endif

        for (int rowI = 0; rowI < numRows; ++rowI) {
            ia[rowI + 1] = ia[rowI] + rowNNZ[rowI];
        }

        ja.resize(ia[numRows] - 1);
#ifdef USE_TBB
        tbb::parallel_for(0, (int)vNeighbor.size(), 1, [&](int vI)
#else
        for (int vI = 0; vI < vNeighbor.size(); ++vI)
#endif
            {
                int rowIStart = vI * DIM_;

                ja.segment(ia[rowIStart] - 1, ja_v[vI].size()) = ja_v[vI];

                for (auto& indexI : IJ2aI[rowIStart]) {
                    indexI.second += ia[rowIStart] - 1;
                }
                for (auto& indexI : IJ2aI[rowIStart + 1]) {
                    indexI.second += ia[rowIStart + 1] - 2;
                }
                if constexpr (DIM_ == 3) {
                    for (auto& indexI : IJ2aI[rowIStart + 2]) {
                        indexI.second += ia[rowIStart + 2] - 3;
                    }
                }
            }
#ifdef USE_TBB
        );
#endif
        ja.array() += 1;
        a.resize(ja.size());

        //NOTE: fixed verts nnz entries are not eliminated
    }

    virtual void load(const char* filePath, Eigen::VectorXd& rhs)
    {
        FILE* in = fopen(filePath, "rb");
        assert(in);

        size_t vecSize;
        fread(&vecSize, sizeof(size_t), 1, in);
        std::cout << "ia size " << vecSize << std::endl;
        ia.resize(vecSize);
        fread(ia.data(), sizeof(ia[0]), vecSize, in);

        fread(&vecSize, sizeof(size_t), 1, in);
        std::cout << "ja size " << vecSize << std::endl;
        ja.resize(vecSize);
        fread(ja.data(), sizeof(ja[0]), vecSize, in);

        if (ia[0] == 0) {
            ia.array() += 1;
            ja.array() += 1;
        }

        fread(&vecSize, sizeof(size_t), 1, in);
        std::cout << "a size " << vecSize << std::endl;
        a.resize(vecSize);
        fread(a.data(), sizeof(a[0]), vecSize, in);

        fread(&vecSize, sizeof(size_t), 1, in);
        std::cout << "rhs size " << vecSize << std::endl;
        rhs.resize(vecSize);
        fread(rhs.data(), sizeof(rhs[0]), vecSize, in);

        numRows = vecSize;

        fclose(in);
        std::cout << "load done" << std::endl;
    }
    virtual void write(const char* filePath, const Eigen::VectorXd& rhs)
    {
        FILE* out = fopen(filePath, "wb");

        size_t vecSize = ia.size();
        fwrite(&vecSize, sizeof(vecSize), 1, out);
        fwrite(ia.data(), sizeof(ia[0]), ia.size(), out);

        vecSize = ja.size();
        fwrite(&vecSize, sizeof(vecSize), 1, out);
        fwrite(ja.data(), sizeof(ja[0]), ja.size(), out);

        vecSize = a.size();
        fwrite(&vecSize, sizeof(vecSize), 1, out);
        fwrite(a.data(), sizeof(a[0]), a.size(), out);

        vecSize = rhs.size();
        fwrite(&vecSize, sizeof(vecSize), 1, out);
        fwrite(rhs.data(), sizeof(rhs[0]), rhs.size(), out);

        fclose(out);
    }

    virtual void set_pattern(const Eigen::SparseMatrix<double>& mtr)
    {
        //NOTE: mtr must be SPD

        numRows = static_cast<int>(mtr.rows());

        ja.conservativeResize(mtr.nonZeros());
        memcpy(ja.data(), mtr.innerIndexPtr(),
            mtr.nonZeros() * sizeof(mtr.innerIndexPtr()[0]));

        ia.conservativeResize(numRows + 1);
        memcpy(ia.data(), mtr.outerIndexPtr(),
            (numRows + 1) * sizeof(mtr.outerIndexPtr()[0]));

        a.conservativeResize(mtr.nonZeros());
        memcpy(a.data(), mtr.valuePtr(),
            mtr.nonZeros() * sizeof(mtr.valuePtr()[0]));
    }

    virtual void analyze_pattern(void) = 0;

    virtual bool factorize(void) = 0;

    virtual void solve(Eigen::VectorXd& rhs,
        Eigen::VectorXd& result)
        = 0;

    virtual void multiply(const Eigen::VectorXd& x,
        Eigen::VectorXd& Ax)
    {
        assert(x.size() == numRows);
        assert(IJ2aI.size() == numRows);

        Ax.setZero(numRows);
        for (int rowI = 0; rowI < numRows; ++rowI) {
            for (const auto& colI : IJ2aI[rowI]) {
                Ax[rowI] += a[colI.second] * x[colI.first];
                if (rowI != colI.first) {
                    Ax[colI.first] += a[colI.second] * x[rowI];
                }
            }
        }
    }

public:
    virtual void outputFactorization(const std::string& filePath)
    {
        assert(0 && "please implement!");
    }
    virtual double coeffMtr(int rowI, int colI) const
    {
        if (rowI > colI) {
            // return only upper right part for symmetric matrix
            int temp = rowI;
            rowI = colI;
            colI = temp;
        }
        assert(rowI < IJ2aI.size());
        const auto finder = IJ2aI[rowI].find(colI);
        if (finder != IJ2aI[rowI].end()) {
            return a[finder->second];
        }
        else {
            return 0.0;
        }
    }
    virtual void getCoeffMtr(Eigen::SparseMatrix<double>& mtr) const
    {
        mtr.resize(numRows, numRows);
        mtr.setZero();
        mtr.reserve(a.size() * 2 - numRows);
        for (int rowI = 0; rowI < numRows; rowI++) {
            for (const auto& colIter : IJ2aI[rowI]) {
                mtr.insert(rowI, colIter.first) = a[colIter.second];
                if (rowI != colIter.first) {
                    mtr.insert(colIter.first, rowI) = a[colIter.second];
                }
            }
        }
    }
    virtual void getCoeffMtr_lower(Eigen::SparseMatrix<double>& mtr) const
    {
        assert(numRows > 0);

        mtr.conservativeResize(numRows, numRows);
        mtr.reserve(a.size());

        memcpy(mtr.innerIndexPtr(), ja.data(), ja.size() * sizeof(ja[0]));
        memcpy(mtr.outerIndexPtr(), ia.data(), ia.size() * sizeof(ia[0]));
        memcpy(mtr.valuePtr(), a.data(), a.size() * sizeof(a[0]));
    }
    virtual void getTriplets(const Eigen::VectorXi& nodeList,
        std::vector<Eigen::Triplet<double>>& triplet) const
    {
        std::map<int, int> rowIMapper;
        for (int i = 0; i < nodeList.size(); ++i) {
            int startI = i * DIM_;
            int startRowI = nodeList[i] * DIM_;

            rowIMapper[startRowI] = startI;
            rowIMapper[startRowI + 1] = startI + 1;
            if constexpr (DIM_ == 3) {
                rowIMapper[startRowI + 2] = startI + 2;
            }
        }

        triplet.resize(0);
        for (int rowI = 0; rowI < numRows; rowI++) {
            auto rowIFinder = rowIMapper.find(rowI);
            for (const auto& colIter : IJ2aI[rowI]) {
                auto colIFinder = rowIMapper.find(colIter.first);
                if (rowIFinder != rowIMapper.end() && colIFinder != rowIMapper.end()) {
                    triplet.emplace_back(rowIFinder->second, colIFinder->second, a[colIter.second]);
                    if (rowIFinder->second != colIFinder->second) {
                        triplet.emplace_back(colIFinder->second, rowIFinder->second, a[colIter.second]);
                    }
                }
            }
        }
    }
    virtual void setCoeff(int rowI, int colI, double val)
    {
        if (rowI <= colI) {
            assert(rowI < IJ2aI.size());
            const auto finder = IJ2aI[rowI].find(colI);
            assert(finder != IJ2aI[rowI].end());
            a[finder->second] = val;
        }
    }
    virtual void setCoeff(const LinSysSolver<vectorTypeI, vectorTypeS>* other,
        double multiplier)
    {
        assert(numRows == other->numRows);
        assert(ja.size() == other->a.size());

        a = multiplier * other->a;
    }
    virtual void setZero(void)
    {
        a.setZero();
    }
    virtual void setUnit_row(int rowI)
    {
        assert(numRows == IJ2aI.size());
        assert(rowI < numRows);
        for (const auto& colIter : IJ2aI[rowI]) {
            a[colIter.second] = (colIter.first == rowI);
        }
    }
    virtual void setUnit_row(int rowI, std::unordered_map<int, double>& rowVec)
    {
        assert(numRows == IJ2aI.size());
        assert(rowI < numRows);
        rowVec.clear();
        for (const auto& colIter : IJ2aI[rowI]) {
            rowVec[colIter.first] = a[colIter.second];
            a[colIter.second] = (colIter.first == rowI);
        }
    }
    virtual void setZeroCol(int colI){
        assert(colI<numRows);
        for(int rowI=0; rowI<numRows; ++rowI){
            const auto finder = IJ2aI[rowI].find(colI);
            if(finder != IJ2aI[rowI].end()){
                a[finder->second] = 0.0;
            }
        }
    }
    virtual void setUnit_col(int colI, const std::set<int>& rowVIs)
    {
        assert(numRows == IJ2aI.size());
        assert(colI < numRows);
        for (const auto& rowVI : rowVIs) {
            for (int dimI = 0; dimI < DIM_; ++dimI) {
                int rowI = rowVI * DIM_ + dimI;
                assert(rowI < numRows);
                if (rowI <= colI) {
                    const auto finder = IJ2aI[rowI].find(colI);
                    if (finder != IJ2aI[rowI].end()) {
                        a[finder->second] = (rowI == colI);
                    }
                }
            }
        }
    }
    virtual void setUnit_col_dim1(int colI, const std::set<int>& rowVIs)
    {
        assert(numRows == IJ2aI.size());
        assert(colI < numRows);
        for (const auto& rowI : rowVIs) {
            assert(rowI < numRows);
            if (rowI <= colI) {
                const auto finder = IJ2aI[rowI].find(colI);
                if (finder != IJ2aI[rowI].end()) {
                    a[finder->second] = (rowI == colI);
                }
            }
        }
    }

    virtual void addCoeff(int rowI, int colI, double val)
    {
        if (rowI <= colI) {
            assert(rowI < IJ2aI.size());
            const auto finder = IJ2aI[rowI].find(colI);
            assert(finder != IJ2aI[rowI].end());
            a[finder->second] += val;
        }
    }
    virtual void precondition_diag(const Eigen::VectorXd& input, Eigen::VectorXd& output)
    {
        assert(numRows == input.size());
        output.resize(numRows);
        for (int rowI = 0; rowI < numRows; ++rowI) {
            const auto finder = IJ2aI[rowI].find(rowI);
            assert(finder != IJ2aI[rowI].end());
            output[rowI] = input[rowI] / a[finder->second];
        }
    }
    virtual void getMaxDiag(double& maxDiag)
    {
        maxDiag = -std::numeric_limits<double>::infinity();
        for (int rowI = 0; rowI < numRows; ++rowI) {
            const auto finder = IJ2aI[rowI].find(rowI);
            assert(finder != IJ2aI[rowI].end());
            if (maxDiag < a[finder->second]) {
                maxDiag = a[finder->second];
            }
        }
    }
    virtual void addCoeff(const LinSysSolver<vectorTypeI, vectorTypeS>* other,
        double multiplier)
    {
        assert(numRows == other->numRows);
        if (a.size() == other->a.size()) {
            a += multiplier * other->a;
        }
        else {
            for (int rowI = 0; rowI < numRows; ++rowI) {
                for (const auto& colIter : other->IJ2aI[rowI]) {
                    const auto finder = IJ2aI[rowI].find(colIter.first);
                    if (finder != IJ2aI[rowI].end()) {
                        a[finder->second] += multiplier * other->a[colIter.second];
                    }
                }
            }
        }
    }

    virtual int getNumRows(void) const
    {
        return numRows;
    }
    virtual int getNumNonzeros(void) const
    {
        return a.size();
    }
    virtual const std::vector<std::map<int, int>>& getIJ2aI(void) const
    {
        return IJ2aI;
    }
    virtual Eigen::VectorXi& get_ia(void) { return ia; }
    virtual Eigen::VectorXi& get_ja(void) { return ja; }
    virtual Eigen::VectorXd& get_a(void) { return a; }
    virtual const Eigen::VectorXd& get_a(void) const { return a; }
};

} // namespace SIM

#endif /* LinSysSolver_hpp */