#include <array>
#include <bitset>
#include <cassert>
#include <iostream>
#include <booluarray.hpp>

#include <cstddef>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include "bernstein.hpp"
#include "multiloop.hpp"
#include "organizer/blobtree.hpp"
#include "quadrature_multipoly.hpp"
#include "binomial.hpp"

#include "real.hpp"
#include "sparkstack.hpp"
#include "uvector.hpp"
#include "vector"
#include "xarray.hpp"
#include <chrono>
#include <cmath>
#include <memory>

#include "organizer/primitive.hpp"

namespace algoim::Organizer
{

namespace detail
{
void bernstein2PowerTensor(const tensor3& phiBernstein, tensor3& phiPower) {}

void restrictToInnerFace(const tensor3& phi, int k, real facePos, tensor2& res)
{
    assert(0 <= k && k < 3);
    assert(all(res.ext() == remove_component(phi.ext(), k)));
    assert(0 < facePos && facePos < 1);
    real* basisAlongK;
    int   P = phi.ext(k);
    algoim_spark_alloc(real, &basisAlongK, P);
    bernstein::evalBernsteinBasis(facePos, P, basisAlongK);

    for (auto i = res.loop(); ~i; ++i) {
        real evalAlongKAtFacePos = 0;
        for (int j = 0; j < P; ++j) { evalAlongKAtFacePos += basisAlongK[j] * phi.m(add_component(i(), k, j)); }
        res.l(i) = evalAlongKAtFacePos;
    }
}

template <int N>
int numOfZero(const uvector<int, N>& v)
{
    int res = 0;
    for (int i = 0; i < N; ++i) {
        if (v(i) == 0) { res++; }
    }
    return res;
}

template <int N>
int binaryToDecimal(const uvector<int, N>& v, int zeroRep = 0)
{
    int res  = 0;
    int base = 1;
    for (int i = N - 1; i >= 0; --i) {
        if (v(i) != zeroRep) { res += base; }
        base *= 2;
    }
    return res;
}

template <int N>
int replaceFirst(uvector<int, N>& v, int oldVal, int newVal, int startIdx = 0)
{
    for (int i = startIdx; i < N; ++i) {
        if (v(i) == oldVal) {
            v(i) = newVal;
            return i;
        }
    }
    return -1;
}

template <int N>
int findFirst(const uvector<int, N>& v, int val, int startIdx = 0)
{
    for (int i = startIdx; i < N; ++i) {
        if (v(i) == val) return i;
    }
    return -1;
}

std::pair<uvector<int, 3>, uvector<int, 3>> getOneEightCellAABB(const uvector3&       min,
                                                                const uvector3&       max,
                                                                const uvector<int, 3> side,
                                                                int                   negativeRep = -1)
{
    std::pair<uvector<int, 3>, uvector<int, 3>> res = {min, max};
    uvector3                                    mid = (min + max) / 2;
    for (int i = 0; i < 3; ++i) {
        if (side(i) == negativeRep) {
            res.second(i) = mid(i);
        } else {
            res.first(i) = mid(i);
        }
    }
    return res;
}
} // namespace detail

// TODO: delete it
bool isPointInside(const std::vector<tensor3>& originTensors, const uvector3& originPt)
{
    for (auto& t : originTensors) {
        if (evalPower(t, originPt) >= 0) { return false; }
    }
    return true;
}

// bool isPointInside(const std::vector<CompleteTensorRep>& completeTensorReps, const uvector3& originPt)
// {
//     for (auto& completeTensorRep : completeTensorReps) {
//         if (isPointInside(completeTensorRep.originalPower, originPt)) { return false; }
//     }
//     return true;
// }

// 这里blobTree是拷贝传参
bool keepQuadraturePoint(const Scene&        scene,
                         organizer::BlobTree blobTree,
                         const OcTreeNode&   ocTreeNode,
                         const uvector3&     originPt)
{
    // 只需要考虑intersect polys，不用考虑fully contained polys
    const auto&       polyIntersectIndices = ocTreeNode.polyIntersectIndices;
    std::vector<bool> primitiveInOuts(polyIntersectIndices.size());
    for (int i = 0; i < ocTreeNode.polyIntersectIndices.size(); ++i) {
        primitiveInOuts[i] = isPointInside(scene.polys[polyIntersectIndices[i]].originalPower, originPt);
    }
    int res = organizer::traverse(blobTree, ocTreeNode.polyIntersectIndices, primitiveInOuts);
    assert(res == organizer::NODE_IN || res == organizer::NODE_OUT);
    if (res == organizer::NODE_OUT) {
        return false;
    } else {
        return true;
    }
}

// std::vector<std::shared_ptr<PrimitiveDesc>> primitives;

// BasicTask(std::vector<std::shared_ptr<PrimitiveDesc>> ps) {};


const std::array<int, 2> sides = {-1, 1};

struct BasicTaskRes {
    real volume;
    real area;
};

// 对于mark含2个及以上0的情况，尝试对每个0填1或-1的所有组合
// TODO: 参数太多了，考虑换用std::function + lambda
void visitSubcellOnBothSidesOfDir(const uvector3& nodeMid,
                                  const int       polyIntersectIndex,
                                  const tensor3&  poly,

                                  std::array<OcTreeNode, CHILD_NUM>& subNodes,
                                  int                                startIdxToCheck,
                                  uvector<int, 3>&                   mark)
{
    int zeroIdx = detail::findFirst(mark, 0, startIdxToCheck);
    if (zeroIdx == -1) {
        tensor3 halfCellPoly(nullptr, poly.ext());
        algoim_spark_alloc(real, halfCellPoly);
        int   subIdx  = detail::binaryToDecimal(mark);
        auto& subNode = subNodes[subIdx];
        bernstein::deCasteljau(poly, subNode.min, subNode.max, halfCellPoly); // 求1/8空间下的表达
        if (bernstein::uniformSign(halfCellPoly) != 1) {
            subNode.polyIntersectIndices.emplace_back(polyIntersectIndex);
        } else {
            organizer::traverse(subNode.blobTree, polyIntersectIndex, organizer::NODE_OUT);
        }
    } else {
        for (auto side : sides) {
            mark(zeroIdx) = side;
            visitSubcellOnBothSidesOfDir(nodeMid, polyIntersectIndex, poly, subNodes, zeroIdx + 1, mark);
        }
        mark(zeroIdx) = 0;
    }
}

// 叶节点单独存在leaves中，是因为最后只需要叶节点信息，没有自顶向下/自底向上遍历
// 所以树结构是不需要的记录的
// void buildOcTree(const Scene& scene, std::vector<Node>& intermediateNodes, std::vector<Node> leaves, int nowNodeIdx)
void buildOcTree(const Scene& scene, const OcTreeNode& node, std::vector<OcTreeNode>& leaves)
{
    const std::vector<int>& polyIntersectIndices = node.polyIntersectIndices;
    if (polyIntersectIndices.size() <= 4) {
        leaves.emplace_back(node);
        return;
    }
    const uvector3&                   nowNodeMin = node.min;
    const uvector3&                   nowNodeMax = node.max;
    std::array<OcTreeNode, CHILD_NUM> subNodes;
    // intermediateNodes.resize(lastIdx + 8);
    int                               subIdx = 0;
    for (MultiLoop<3> j(0, 2); ~j; ++j, ++subIdx) {
        auto nodeAABB             = detail::getOneEightCellAABB(node.min, node.max, j(), 0);
        subNodes[subIdx].min      = nodeAABB.first;
        subNodes[subIdx].max      = nodeAABB.second;
        subNodes[subIdx].blobTree = node.blobTree;
    }
    uvector3 nodeMid = (nowNodeMin + nowNodeMax) / 2;
    for (int i = 0; i < polyIntersectIndices.size(); ++i) {
        const int       polyIntersectIndex = polyIntersectIndices[i];
        const auto&     poly               = scene.polys[polyIntersectIndex];
        uvector<int, 3> mark(0, 0, 0);
        for (int faceAxis = 0; faceAxis < 3; ++faceAxis) {
            real    centerPlane = nodeMid(faceAxis);
            tensor2 restrictToCenterPlane(nullptr, remove_component(poly.compositedBernstein.ext(), faceAxis));
            algoim_spark_alloc(real, restrictToCenterPlane);
            // NOTE: 这里restrict To Face用的是合成的tensor，这可能导致产生很多实际不需要考虑的交线
            detail::restrictToInnerFace(poly.compositedBernstein, faceAxis, centerPlane, restrictToCenterPlane);
            int signOnHalfPlane = bernstein::uniformSign(restrictToCenterPlane);
            if (signOnHalfPlane == -1) {
                // primitive contain the centerface
                mark(faceAxis) = 2;
            } else if (signOnHalfPlane == 1) {
                // primitive intersects either side or both sides of the centerface
                // deCasteljau to transformation
                uvector3 halfCellMin = nowNodeMin, halfCellMax = nowNodeMax;
                halfCellMax(faceAxis) = nodeMid(faceAxis);
                tensor3 halfCellPoly(nullptr, poly.compositedBernstein.ext());
                algoim_spark_alloc(real, halfCellPoly);
                bernstein::deCasteljau(poly.compositedBernstein, halfCellMin, halfCellMax, halfCellPoly);
                if (bernstein::uniformSign(halfCellPoly) != 1) {
                    // 负空间有
                    mark(faceAxis) = -1;
                }
                halfCellMax(faceAxis) = nowNodeMax(faceAxis);
                halfCellMin(faceAxis) = nodeMid(faceAxis);
                bernstein::deCasteljau(poly.compositedBernstein, halfCellMin, halfCellMax, halfCellPoly);
                if (bernstein::uniformSign(halfCellPoly) != 1) {
                    // 正空间有
                    mark(faceAxis) += 1; // 当正负空间都有，记0
                }
            }
        }

        if (any(mark == uvector3(2, 2, 2))) {
            if (all(mark == uvector3(2, 2, 2))) {
                // fully containing cases
                for (int i = 0; i < CHILD_NUM; ++i) {
                    tensor3 subPoly(nullptr, poly.compositedBernstein.ext());
                    bernstein::deCasteljau(poly.compositedBernstein, subNodes[i].min, subNodes[i].max, subPoly);
                    if (bernstein::uniformSign(subPoly) == -1) {
                        subNodes[i].polyFullyContainedIndices.emplace_back(polyIntersectIndex);
                    } else {
                        subNodes[i].polyIntersectIndices.emplace_back(polyIntersectIndex);
                    }
                }
                continue;
            }
            for (int subIdx = 0; subIdx < CHILD_NUM; ++subIdx) {
                // intermediateNodes[lastIdx + subIdx].polyIntersectIndices.emplace_back(poly);
                // intermediateNodes[lastIdx + subIdx].min = nowNodeMin;
                // intermediateNodes[lastIdx + subIdx].max = nowNodeMax;
                subNodes[subIdx].polyIntersectIndices.emplace_back(polyIntersectIndex);
            }
            continue;
        }
        int zeroCount = detail::numOfZero<3>(mark);
        if (zeroCount == 0) {
            // mark has -1 or 1 only
            real nodeMidVal = bernstein::evalBernsteinPoly(poly.compositedBernstein, nodeMid);
            int  markIdx    = detail::binaryToDecimal(mark);
            for (int subIdx = 0; subIdx < CHILD_NUM; ++subIdx) {
                if (subIdx == markIdx)
                    subNodes[subIdx].polyIntersectIndices.emplace_back(polyIntersectIndex);
                else
                    organizer::traverse(subNodes[subIdx].blobTree, polyIntersectIndex, organizer::NODE_OUT);
            }
        } else if (zeroCount == 1) {
            // poly related to 2 subcells
            uvector<int, 3> sideMark = mark;
            int             zeroIdx  = detail::replaceFirst(sideMark, 0, 1);
            int             sideIdx1 = detail::binaryToDecimal(sideMark, -1);
            sideMark(zeroIdx)        = -1;
            int sideIdx2             = detail::binaryToDecimal(sideMark, -1);
            subNodes[sideIdx1].polyIntersectIndices.emplace_back(polyIntersectIndex);
            subNodes[sideIdx2].polyIntersectIndices.emplace_back(polyIntersectIndex);
            for (int subIdx = 0; subIdx < CHILD_NUM; ++subIdx) {
                if (subIdx != sideIdx1 && subIdx != sideIdx2) {
                    organizer::traverse(subNodes[subIdx].blobTree, polyIntersectIndex, organizer::NODE_OUT);
                }
            }
        } else {
            // zeroCount == 2 or 3
            visitSubcellOnBothSidesOfDir(nodeMid, polyIntersectIndex, poly.compositedBernstein, subNodes, 0, mark);
        }
    }
    // launch subdivision in subcells
    for (subIdx = 0; subIdx < CHILD_NUM; ++subIdx) {
        subNodes[subIdx].polyFullyContainedIndices.resize(subNodes[subIdx].polyFullyContainedIndices.size()
                                                          + node.polyFullyContainedIndices.size());
        subNodes[subIdx].polyFullyContainedIndices.insert(subNodes[subIdx].polyFullyContainedIndices.end(),
                                                          node.polyFullyContainedIndices.begin(),
                                                          node.polyFullyContainedIndices.end());
        buildOcTree(scene, subNodes[subIdx], leaves);
    }
}

/** single object quadrature */
void basicTask(const std::shared_ptr<PrimitiveDesc>& p, int q = 20, real xmin = -1, real xmax = 1)
{
    real     volume    = 0;
    auto     integrand = [](const uvector<real, 3>& x) { return 1.0; };
    uvector3 range     = xmax - xmin;

    if (auto pt = std::dynamic_pointer_cast<SphereDesc>(p)) {
        tensor3 tensor(nullptr, 3);
        algoim_spark_alloc(real, tensor);
        makeSphere(*pt, tensor);
        detail::powerTransformation(range, xmin, tensor);

        tensor3 phi(nullptr, 3);
        algoim_spark_alloc(real, phi);
        detail::power2BernsteinTensor(tensor, phi);

        uvector<real, 3> testX(0., 0., 0.25);
        real             testEvalBernstein = bernstein::evalBernsteinPoly(phi, testX);
        std::cout << "eval bernstein without interpolation:" << testEvalBernstein << std::endl;

        ImplicitPolyQuadrature<3> ipquad(phi);
        ipquad.integrate(AutoMixed, q, [&](const uvector<real, 3>& x, real w) {
            if (isInsideBernstein(phi, x)) volume += w * integrand(xmin + x * (xmax - xmin));
        });
    } else if (auto pt = std::dynamic_pointer_cast<MeshDesc>(p)) {
        const int faceCount = pt->indexInclusiveScan.size();
        assert(faceCount > 1);
        std::vector<tensor3> planeTensors(faceCount, tensor3(nullptr, 2));
        algoim_spark_alloc(real, planeTensors);

        tensor3 compositeTensor(nullptr, 1 + faceCount);
        algoim_spark_alloc(real, compositeTensor);

        makeMesh(*pt, compositeTensor, planeTensors);
        detail::powerTransformation(range, xmin, compositeTensor);

        auto planeStdVector1          = xarray2StdVector(planeTensors[0]);
        auto planeStdVector2          = xarray2StdVector(planeTensors[1]);
        auto compositeTensorStdVector = xarray2StdVector(compositeTensor);

        uvector<real, 3> testX(0., 0.75, 0.2);
        real             textEvalPower = evalPower(compositeTensor, testX);

        tensor3 phi(nullptr, 1 + faceCount);
        algoim_spark_alloc(real, phi);
        detail::power2BernsteinTensor(compositeTensor, phi);

        int                       quadraturePointCount = 0;
        ImplicitPolyQuadrature<3> ipquad(phi);

        real testEvalBernstein = bernstein::evalBernsteinPoly(phi, testX);
        ipquad.integrate(AutoMixed, q, [&](const uvector<real, 3>& x, real w) {
            quadraturePointCount++;
            auto realX = x * range + xmin;
            if (isInsidePowers(planeTensors, realX)) volume += w * integrand(realX);
        });
        std::cout << "textEvalPower: " << textEvalPower << std::endl;
        std::cout << "quadraturePointCount: " << quadraturePointCount << std::endl;
    }
    volume *= pow(xmax - xmin, 3);
    std::cout << "Volume xxx: " << volume << std::endl;
};

void basicTask(const std::vector<std::shared_ptr<PrimitiveDesc>>& primitives, int q = 20, real xmin = -1, real xmax = 1)
{
    std::vector<SparkStack<real>*> phiStacks;
    std::vector<tensor3>           phis;

    std::vector<SparkStack<real>*> originTensorStacks;
    std::vector<tensor3>           originTensors;
    real                           volume;
    auto                           integrand = [](const uvector<real, 3>& x) { return 1.0; };
    uvector3                       range     = xmax - xmin;
    uvector<real, 3>               testX(0., 0.75, 0.2);
    for (int i = 0; i < primitives.size(); i++) {
        if (auto pt = std::dynamic_pointer_cast<SphereDesc>(primitives[i])) {
            tensor3 originTensor(nullptr, 3), transformedTensor(nullptr, 3);
            originTensorStacks.emplace_back(algoim_spark_alloc_heap(real, originTensor)); // 记录，用以最后bool

            tensor3 phi(nullptr, 3);
            phiStacks.emplace_back(
                algoim_spark_alloc_heap(real, phi)); // 必须先于algoim_spark_alloc，使transformedTensor的内存以栈形式释放

            algoim_spark_alloc(real, transformedTensor);
            makeSphere(*pt, originTensor);
            originTensors.emplace_back(originTensor);
            detail::powerTransformation(range, xmin, originTensor, transformedTensor);


            detail::power2BernsteinTensor(transformedTensor, phi);
            phis.emplace_back(phi);
        } else if (auto pt = std::dynamic_pointer_cast<MeshDesc>(primitives[i])) {
            const int faceCount = pt->indexInclusiveScan.size();
            assert(faceCount > 1);
            std::vector<tensor3> planeTensors(faceCount, tensor3(nullptr, 2));
            algoimSparkAllocHeapVector(originTensorStacks, planeTensors);

            tensor3 phi(nullptr, 1 + faceCount);
            phiStacks.emplace_back(
                algoim_spark_alloc_heap(real, phi)); // 必须先于algoim_spark_alloc，使compositeTensor的内存以栈形式释放

            tensor3 compositeTensor(nullptr, 1 + faceCount);
            algoim_spark_alloc(real, compositeTensor);

            makeMesh(*pt, compositeTensor, planeTensors);
            detail::powerTransformation(range, xmin, compositeTensor);

            real testEvalPower = evalPower(compositeTensor, testX);

            originTensors.insert(originTensors.end(), planeTensors.begin(), planeTensors.end());

            detail::power2BernsteinTensor(compositeTensor, phi);
            real testEvalBernstein = bernstein::evalBernsteinPoly(phi, testX);
            phis.emplace_back(phi);
        }
    }
    real                      testEvalBernstein = bernstein::evalBernsteinPoly(phis[0], testX);
    ImplicitPolyQuadrature<3> ipquad(phis);

    ipquad.integrate(AutoMixed, q, [&](const uvector<real, 3>& x, real w) {
        auto realX = x * range + xmin;
        if (isPointInside(originTensors, realX)) volume += w * integrand(realX);
    });
    volume *= pow(xmax - xmin, 3);
    std::cout << "Volume xxx: " << volume << std::endl;

    // free memory, further deallocating memory of xarray
    for (auto& p : phiStacks) delete p;
    for (auto& p : originTensorStacks) delete p;
};

BasicTaskRes basicTask(const Scene& scene, const organizer::BlobTree& blobTree, const OcTreeNode& node, int q = 20)
{
    auto                      integrand = [](const uvector<real, 3>& x) { return 1.0; };
    real                      volume = 0., surf = 0.;
    auto                      range = node.max - node.min;
    ImplicitPolyQuadrature<3> ipquad(scene.polys);
    ipquad.integrate(AutoMixed, q, [&](const uvector<real, 3>& x, real w) {
        auto realX = x * range + node.min;
        if (keepQuadraturePoint(scene, blobTree, node, realX)) volume += w * integrand(realX);
    });
}

void quadratureScene(const std::vector<std::shared_ptr<PrimitiveDesc>>& primitives,
                     const uvector3&                                    xmin,
                     const uvector3&                                    xmax,
                     const organizer::BlobTree&                         blobTree)
{
    OcTreeNode              rootNode(xmin, xmax, blobTree);
    std::vector<OcTreeNode> leaves;

    std::vector<SparkStack<real>*> tensorStacks;
    std::vector<CompleteTensorRep> completeTensorReps;

    real             volume;
    auto             integrand = [](const uvector<real, 3>& x) { return 1.0; };
    uvector3         range     = xmax - xmin;
    uvector<real, 3> testX(0., 0.75, 0.2);
    for (int i = 0; i < primitives.size(); i++) {
        if (auto pt = std::dynamic_pointer_cast<SphereDesc>(primitives[i])) {
            tensor3 originTensor(nullptr, 3);
            tensor3 transformedTensor(nullptr, 3);
            tensorStacks.emplace_back(algoim_spark_alloc_heap(real, originTensor));

            tensor3 phi(nullptr, 3);
            tensorStacks.emplace_back(
                algoim_spark_alloc_heap(real, phi)); // 必须先于algoim_spark_alloc，使transformedTensor的内存以栈形式释放

            algoim_spark_alloc(real, transformedTensor);
            makeSphere(*pt, originTensor);

            detail::powerTransformation(range, xmin, originTensor, transformedTensor);

            detail::power2BernsteinTensor(transformedTensor, phi);
            completeTensorReps.emplace_back(CompleteTensorRep{phi, {originTensor}});
        } else if (auto pt = std::dynamic_pointer_cast<MeshDesc>(primitives[i])) {
            const int faceCount = pt->indexInclusiveScan.size();
            assert(faceCount > 1);
            std::vector<tensor3> planeTensors(faceCount, tensor3(nullptr, 2));
            algoimSparkAllocHeapVector(tensorStacks, planeTensors);

            tensor3 phi(nullptr, 1 + faceCount);
            tensorStacks.emplace_back(
                algoim_spark_alloc_heap(real, phi)); // 必须先于algoim_spark_alloc，使compositeTensor的内存以栈形式释放

            tensor3 compositeTensor(nullptr, 1 + faceCount);
            algoim_spark_alloc(real, compositeTensor);

            makeMesh(*pt, compositeTensor, planeTensors);
            detail::powerTransformation(range, xmin, compositeTensor);

            real testEvalPower = evalPower(compositeTensor, testX);

            detail::power2BernsteinTensor(compositeTensor, phi);
            real testEvalBernstein = bernstein::evalBernsteinPoly(phi, testX);

            completeTensorReps.emplace_back(CompleteTensorRep{phi, planeTensors});
        }
    }

    Scene scene{completeTensorReps, blobTree};
    buildOcTree(scene, rootNode, leaves);

    for (const auto& leaf : leaves) {
        auto basicRes  = basicTask(blobTree, scene, leaf, xmin, xmax, 10);
        volume        += basicRes.volume;
    }

    volume *= prod(xmax - xmin);
    // TODO: surface area
    std::cout << "Volume xxx: " << volume << std::endl;


    // free memory, further deallocating memory of xarray
    for (auto& p : tensorStacks) delete p;
}
}; // namespace algoim::Organizer

// clang-format off
/** 
	最一开始的cell应当包含所有的primitive
	func(given cell) {
		得到8个subcell
		for each primitive {
			mark = [0,0,0]
			for k in 0,1,2 {
				测试primitive和k轴向的centerface有无交
				if 无交 {0
					if primitive包裹centerface {
						mark[k] = 2
					} else {
						判断primitive位于哪一侧
						if k正方向 mark[k] = 1
						else mark[k] = -1
					}
				}
			}
			if mark 含2 {
				// TODO: 可能需要优化：在极端情况下，
				// 例如primitive的边刚好与centerface重合，
				// 此时不是所有subcell都关联primitive
				将primitive放入所有subcell
			} elif mark 没有0 {
				将primitive放入对应的subcell
			} elif mark 有1个0 {
				将primitive放入两个对应的subcell
			} else { // >= 2个0
				测试每个可能有交subcell是否与primitive相交
				相交的放入对于subcell
			}
		}
		for each subcell {
			func(subcell)
		}
	}
	
 */

// clang-format on