#include "blockingGeometryInfo.h"
#include "Function.h"
#include "Function2D.h"
#include "Function3D.h"

using namespace Aurora;
using namespace Recon;

GeometryBlock Recon::blockingGeometryInfos(const GeometryInfo& aGeom, const Aurora::Matrix& aRnBlock, const Aurora::Matrix& aRlBlock,
                                           const Aurora::Matrix& aSnBlock, const Aurora::Matrix& aSlBlock, const Aurora::Matrix& aMpBlock)
{
    GeometryBlock result;
    int numData = aRnBlock.getDimSize(1);
    Matrix arrayOfDataSize = zeros(1, numData, 1) + 1;
    int dim = aGeom.receiverPositions[0].getDimSize(0);
    double* receiverPositionBlockData = Aurora::malloc(dim * numData);
    double* receiverNormalBlockData = Aurora::malloc(dim * numData);
    result.receiverPositionBlock = Matrix::New(receiverPositionBlockData, dim, numData);
    result.receiverNormalBlock = Matrix::New(receiverNormalBlockData, dim, numData);

    double* sizeData = Aurora::malloc(4);
    sizeData[0] = aGeom.receiverPositions[0].getDimSize(0);
    sizeData[1] = aGeom.receiverPositions[0].getDimSize(1);
    sizeData[2] = aGeom.receiverPositions[0].getDimSize(2);
    sizeData[3] = aGeom.receiverPositions.size();
    size_t matrixSize = aGeom.receiverPositions[0].getDataSize();
    Matrix receiverPositionsSize = Matrix::New(sizeData, 1, 4);
    #pragma omp parallel for
    for(int i=0; i<dim; ++i)
    {
        Matrix idx = sub2ind(receiverPositionsSize, {arrayOfDataSize * (i+1), aRnBlock, aRlBlock, aMpBlock});
        for(int j=0; j<idx.getDataSize(); ++j)
        {

            unsigned long long index4D = (idx[j] - 1) / matrixSize;
            unsigned long long index3D = idx[j] - 1;
            if(index4D > 0)
            {
                index3D -= matrixSize;
            }
            receiverPositionBlockData[dim*j + i] = aGeom.receiverPositions[index4D][index3D];
            receiverNormalBlockData[dim*j + i] = aGeom.receiverNormals[index4D][index3D];
        }
    }

    double* senderPositionBlockData = Aurora::malloc(dim * numData);
    double* senderNormalBlockData = Aurora::malloc(dim * numData);
    result.senderPositionBlock = Matrix::New(senderPositionBlockData, dim, numData);
    result.senderNormalBlock = Matrix::New(senderNormalBlockData, dim, numData);
    sizeData = Aurora::malloc(4);
    sizeData[0] = aGeom.senderPositions[0].getDimSize(0);
    sizeData[1] = aGeom.senderPositions[0].getDimSize(1);
    sizeData[2] = aGeom.senderPositions[0].getDimSize(2);
    sizeData[3] = aGeom.senderPositions.size();
    matrixSize = aGeom.receiverPositions[0].getDataSize();
    Matrix senderPositionsSize = Matrix::New(sizeData, 1, 4);
    dim = aGeom.senderPositions[0].getDimSize(0);
    #pragma omp parallel for
    for(int i=0; i<dim; ++i)
    {
        Matrix idx = sub2ind(senderPositionsSize, {arrayOfDataSize * (i+1), aSnBlock, aSlBlock, aMpBlock});
        for(int j=0; j<idx.getDataSize(); ++j)
        {

            unsigned long long index4D = (idx[j] - 1) / matrixSize;
            unsigned long long index3D = idx[j] - 1;
            if(index4D > 0)
            {
                index3D -= matrixSize;
            }
            senderPositionBlockData[dim*j + i] = aGeom.senderPositions[index4D][index3D];
            senderNormalBlockData[dim*j + i] = aGeom.senderNormals[index4D][index3D];
        }
    }
    return result;
}