#include "getAscanBlock.h"

#include "Function.h"
#include "Matrix.h"
#include "Parser.h"
#include "Data/OneTasAScanData.h"
#include "Data/AScanData.h"
#include "Data/TasElementIndex.h"
#include "Data/ElementIndex.h"
#include "Data/GeometryIndex.h"
#include "Data/ElectricIndex.h"
#include "Data/MetaData.h"
#include "ShotList/ShotList.h"
#include <algorithm>
#include <cstddef>
#include <vector>

using namespace Recon;
using namespace Aurora;

namespace
{
    std::vector<int> findTasAndElementIndex(TasIndicesPointer aTasIndices, ReceiverIndicesPointer aReceiverIndices, const Matrix& aRl, const Matrix& aRn)
    {
        std::vector<int> result;
        Matrix sortRl = Matrix::copyFromRawData(aRl.getData(), aRl.getDimSize(0), aRl.getDimSize(1), aRl.getDimSize(2));
        Matrix sortRn = Matrix::copyFromRawData(aRn.getData(), aRn.getDimSize(0), aRn.getDimSize(1), aRn.getDimSize(2));
        std::sort(sortRl.getData(), sortRl.getData() + sortRl.getDataSize());
        std::sort(sortRn.getData(), sortRn.getData() + sortRn.getDataSize());
        for(int i=0; i<aTasIndices.getLength(); ++i)
        {
            if (std::binary_search(sortRl.getData(), sortRl.getData() + sortRl.getDataSize(), aTasIndices.get()[i]) &&
                std::binary_search(sortRn.getData(), sortRn.getData() + sortRn.getDataSize(), aReceiverIndices.get()[i]) )
            {
                result.push_back(i);
            }
        }

        return result;
    }
}

AscanBlock Recon::getAscanBlock(Parser* aParser, const Aurora::Matrix& aMp, const Aurora::Matrix& aSl, const Aurora::Matrix& aSn, 
                                const Aurora::Matrix& aRl, const Aurora::Matrix& aRn)
{
    AscanBlock result;
    size_t numScans =  aMp.getDataSize() * aSl.getDataSize() * aSn.getDataSize() * aRl.getDataSize() * aRn.getDataSize();
    //size_t numScansIndex = 0;
    int ascanBlockSize = numScans * aParser->getMetaData().getSampleNumber(); 
    float* ascanBlockData = new float[ascanBlockSize];
    float* mpBlockData = new float[numScans];
    float* slBlockData = new float[numScans];
    float* snBlockData = new float[numScans];
    float* rlBlockData = new float[numScans];
    float* rnBlockData = new float[numScans];
    float* gainBlockData = new float[numScans];
    result.ascanBlock = Matrix::fromRawData(ascanBlockData, aParser->getMetaData().getSampleNumber(), numScans);
    result.mpBlock = Matrix::fromRawData(mpBlockData, 1, numScans);
    result.slBlock = Matrix::fromRawData(slBlockData, 1, numScans);
    result.snBlock = Matrix::fromRawData(snBlockData, 1, numScans);
    result.rlBlock = Matrix::fromRawData(rlBlockData, 1, numScans);
    result.rnBlock = Matrix::fromRawData(rnBlockData, 1, numScans);
    result.gainBlock = Matrix::fromRawData(gainBlockData, 1, numScans);
    auto tasIndices = aParser->getMetaData().getTasIndices();
    auto receiverIndices = aParser->getMetaData().getReceiverIndices();
    auto tasElementMapper = findTasAndElementIndex(tasIndices, receiverIndices, aRl, aRn);

    for(int mpIndex=0; mpIndex<aMp.getDataSize();++mpIndex)
    {
        for(int slIndex=0; slIndex<aSl.getDataSize();++slIndex)
        {
            OneTasAScanData oneTasData = aParser->getOneTasAscanDataOfMotorPosition(aSl[slIndex]);
            #pragma omp parallel for
            for(int snIndex=0; snIndex<aSn.getDataSize();++snIndex)
            {
                //int mapperIndex = 0;           
                //#pragma omp parallel for
                for(int rlIndex=0; rlIndex<aRl.getDataSize();++rlIndex)
                {                            
                    for(int rnIndex=0; rnIndex<aRn.getDataSize(); ++rnIndex)
                    {
                        size_t mapperIndex = rnIndex + rlIndex*aRn.getDataSize();
                        size_t numScansIndex = rnIndex  + rlIndex*aRn.getDataSize() + snIndex * aRl.getDataSize() * aRn.getDataSize() + 
                                                slIndex * aRl.getDataSize() * aRn.getDataSize() * aSn.getDataSize() + 
                                                mpIndex * aRl.getDataSize() * aRn.getDataSize() * aSn.getDataSize() * aSl.getDataSize();
                        MotorPosition mp = aMp[mpIndex] == 1 ? MotorPosition::MotorPosition1 : MotorPosition::MotorPosition2;
                        AScanData ascan = aParser->searchAscanDataFromOneTasAscanDataOfMP(oneTasData, ElementIndex(GeometryIndex(aSn[snIndex])), TasElementIndex(tasIndices.get()[tasElementMapper[mapperIndex]], ElementIndex(GeometryIndex(receiverIndices.get()[tasElementMapper[mapperIndex]]))), mp);
                        float* startPointer = ascanBlockData + numScansIndex *  ascan.getAscanDataLength();
                        std::copy(ascan.get() ,ascan.get() + ascan.getAscanDataLength(), startPointer);
                        //ascanBlockData += ascan.getAscanDataLength();
                        mpBlockData[numScansIndex] = aMp[mpIndex];
                        slBlockData[numScansIndex] = aSl[slIndex];
                        snBlockData[numScansIndex] = aSn[snIndex];
                        rlBlockData[numScansIndex] = tasIndices.get()[tasElementMapper[mapperIndex]];
                        rnBlockData[numScansIndex] = receiverIndices.get()[tasElementMapper[mapperIndex]];
                        gainBlockData[numScansIndex] = ascan.getAmplification()[0];
                        //++numScansIndex;
                        //++mapperIndex;
                    }
                }
            }
        }
    }    

    // 改分支不会进，暂不实现 todo
    // % initialized and calculated numbers agree?, otherwise reduce block sizes
    // if num ~= numScans
    // warning([mfilename, ':MissingSenderReceiverInfos'],'Not all defined data available.')
    // writeReconstructionLog('Not all defined sender/receiver available', 3);
    // usedDataIdxs = 1:num;
    // [slBlock, snBlock, rlBlock, rnBlock, mpBlock, gainBlock] = removeDataFromArrays(usedDataIdxs, slBlock, snBlock, rlBlock, rnBlock, mpBlock, gainBlock);
    // AscanBlock = AscanBlock(:,usedDataIdxs);
    // end
    

    return result;
}