#include "detection.h"

#include <algorithm>
#include <cstddef>
#include <sys/types.h>

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

#include "Matrix.h"
#include "common/getMeasurementMetaData.h"
#include "config/config.h"
#include "transmissionReconstruction/detection/detection.cuh"

//remove detectTofAndAttMex function, use detectTofAndAtt function instead of
//保留函数和函数申明备忘，以防以后再次使用，以下头文件从动态库中映入了detectTofAndAttMex的核心函数
// #include "calculateBankDetectAndHilbertTransformation.hpp"

using namespace Aurora;
namespace Recon {
    Matrix  analyzeDetections(const Matrix &slBlockTotal, const Matrix &snBlockTotal, const Matrix &rlBlockTotal, const Matrix &rnBlockTotal, const Matrix &tofDataTotal){
        //TODO: 待完成，暂时用不到
        return Matrix();
    }


    Matrix calculateSOSOffset(const Matrix &sosBlock, float referenceSOS, const Matrix &distBlock, float sampleRate){
        auto offset = (distBlock / sosBlock - distBlock / referenceSOS) * sampleRate;
        return offset;
    }


    Matrix calculateAttenuation(const Matrix &ascans, const Matrix &startPos, const Matrix &endPos, const Matrix &ascansRef, const Matrix &startPosRef, const Matrix &endPosRef)
    {
        auto ascans2 = zeros(ascans.getDimSize(0), ascans.getDimSize(1));
        auto ascansRef2 = zeros(ascansRef.getDimSize(0), ascansRef.getDimSize(1));
        #pragma omp parallel for 
        for (size_t i = 0; i < ascans.getDimSize(1); i++)
        {
            for (size_t j = startPos[i]-1; j < endPos[i]; j++)
            {
                ascans2(j, i) = ascans(j, i);
            }
            for (size_t j = startPosRef[i]-1; j < endPosRef[i]; j++)
            {
                ascansRef2(j, i) = ascansRef(j, i);
            }
        }

        auto pulseEnergy = sum(ascans2^2);
        auto pulseEnergyEmpty = sum(ascansRef2^2);

        return log(pulseEnergyEmpty/pulseEnergy);
    }

    SearchPosition calculateStarEndSearchPosition(const Matrix &aVDistBlock,
                                float minSpeedOfSound, float maxSpeedOfSound,
                                float sampleRate, float maxSample,
                                const Matrix &aVSosOffsetBlock,
                                float startOffset, float segmentLenOffset)
    {
        auto startSearch = floor((aVDistBlock / maxSpeedOfSound)*sampleRate+aVSosOffsetBlock+startOffset);
        for (size_t i = 0; i < startSearch.getDataSize(); i++)
        {
            startSearch[i] = (uint)(startSearch[i]);
            if(startSearch[i]<1) startSearch[i]=1;
            else if(startSearch[i]>maxSample) startSearch[i] = maxSample;
        }

        auto endSearch = ceil(aVDistBlock/minSpeedOfSound*sampleRate+aVSosOffsetBlock+startOffset+segmentLenOffset);
        for (size_t i = 0; i < endSearch.getDataSize(); i++)
        {
            endSearch[i] = (uint)(endSearch[i]);
            if(endSearch[i]<1) endSearch[i]=1;
            else if(endSearch[i]>maxSample) endSearch[i] = maxSample;
        }
        SearchPosition result;
        result.startSearch = startSearch;
        result.endSearch = endSearch;
        return result;
    }
    
    TimeWindowResult applyTimeWindowing(const Aurora::Matrix &AscanBlock, float sampleRate,
        const Aurora::Matrix &distBlock, float sosWater,
        float startOffset, float segmentLenOffset,
        float minSpeedOfSound, float maxSpeedOfSound, bool gaussWindow)
    {
        Aurora::Matrix sosOffset = Aurora::Matrix::fromRawData(new float[1]{0}, 1);
        auto calcResult = calculateStarEndSearchPosition(distBlock, minSpeedOfSound, maxSpeedOfSound, sampleRate,AscanBlock.getDimSize(0), sosOffset, startOffset, segmentLenOffset);

        auto AscanBlockProcessed = zeros(AscanBlock.getDimSize(0),AscanBlock.getDimSize(1));

        if(gaussWindow)
        {
            auto expectedPosWater = (distBlock / sosWater) * sampleRate + startOffset;
            auto windowWidth = calcResult.endSearch-calcResult.startSearch;
            #pragma omp parallel for
            for (size_t i = 0; i < AscanBlock.getDimSize(1); i++)
            {
                auto t = linspace(-5,5,windowWidth[i]);
                auto gauss = exp( -0.1 * (transpose(t) ^ 2));
                
                for (size_t j = calcResult.startSearch[i]-1; j < calcResult.endSearch[i]; j++)
                {
                    AscanBlockProcessed(j, i) = AscanBlock(j, i);
                }
                auto temp = zeros(AscanBlockProcessed.getDimSize(0), 1);
                size_t end = std::round(expectedPosWater[i])-std::round(windowWidth[i]/2)+gauss.getDataSize()-1;
                size_t gIdx = 0;
                for (size_t k = std::round(expectedPosWater[i])- std::round(windowWidth[i]/2)-1; k < end; k++)
                {
                    temp[k] = gauss[gIdx++];
                }
                AscanBlockProcessed($,i) = AscanBlockProcessed($,i).toMatrix() * temp;
            }
        }
        else{
            #pragma omp parallel for
            for (size_t i = 0; i < AscanBlock.getDimSize(1); i++) {
                for (size_t j = calcResult.startSearch[i] - 1;
                     j < calcResult.endSearch[i]; j++) {
                    AscanBlockProcessed(j, i) = AscanBlock(j, i);
                }
            }
        }
        TimeWindowResult result;
        result.startSearch = calcResult.startSearch;
        result.AscanBlockProcessed = AscanBlockProcessed;
        return result;
    }
    
    Aurora::Matrix detectAttVectorized(const Aurora::Matrix &Ascan, const Aurora::Matrix &AscanRef,
                        const Aurora::Matrix &distRef,
                        float sosWaterRef,
                        const Aurora::Matrix &tof, int aScanReconstructionFrequency,
                        float offsetElectronic, int detectionWindowATT)
    {
        auto sizeAscan = size(Ascan);

        auto sampleRate = aScanReconstructionFrequency;
        float offsetElectronicSamples = offsetElectronic * sampleRate;

        auto envelopeOfAScan = abs(hilbert(Ascan));
        auto envelopeOfReferenceAScan = abs(hilbert(AscanRef));
        auto tof2 = (distRef / sosWaterRef) * sampleRate + offsetElectronicSamples;

        auto startPos = zeros(Ascan.getDimSize(1), 1);
        auto endPos = zeros(Ascan.getDimSize(1), 1);
        auto startPosRef = zeros(Ascan.getDimSize(1), 1);
        auto endPosRef = zeros(Ascan.getDimSize(1), 1);
        #pragma omp parallel for
        for (size_t i = 0; i < Ascan.getDimSize(1); i++)
        {
            startPos[i] = std::floor(std::max(tof[i]*sampleRate+offsetElectronicSamples,(float)1.0));
            endPos[i] = std::ceil(std::min(sizeAscan[0], tof[i]*sampleRate+offsetElectronicSamples+detectionWindowATT));
            startPosRef[i] =  std::floor(std::max( tof2[i],(float)1.0));
            endPosRef[i] = std::ceil(std::min(sizeAscan[1], tof2[i]+detectionWindowATT));
        }
        return calculateAttenuation(envelopeOfAScan,startPos,endPos,envelopeOfReferenceAScan,startPosRef,endPosRef);
    }

    int nextpow2(unsigned int value){
        unsigned int compareValue = 2;
        int result = 1;    
        while(compareValue<=value){
            compareValue=compareValue<<1;
            result++;
        }
        return result;
    }

    DetectResult detectTofVectorized(
        const Aurora::Matrix &AscanBlock, const Aurora::Matrix &AscanRefBlock,
        const Aurora::Matrix &distBlock, const Aurora::Matrix &distBlockRef,
        float aSOSWater, float aSOSWaterRef,
        int useTimeWindowing, int aScanReconstructionFrequency,
        float offsetElectronic, int detectionWindowSOS, float minSpeedOfSound,
        float maxSpeedOfSound, bool gaussWindow) 
    {
        auto sampleRate = aScanReconstructionFrequency;
        float offsetElectronicSamples = offsetElectronic * sampleRate;
        Matrix diffStartSearch;
        TimeWindowResult timeResult1;
        timeResult1.AscanBlockProcessed = AscanBlock;
        TimeWindowResult timeResult2;
        timeResult2.AscanBlockProcessed = AscanRefBlock;
        if (useTimeWindowing == 1) {
            timeResult1 = applyTimeWindowing(
                AscanBlock, sampleRate, distBlock, 
                aSOSWater, offsetElectronicSamples, detectionWindowSOS,
                minSpeedOfSound, maxSpeedOfSound, gaussWindow);
            timeResult2 = applyTimeWindowing(
                AscanRefBlock, sampleRate, distBlockRef, 
                aSOSWaterRef, offsetElectronicSamples, detectionWindowSOS,
                minSpeedOfSound, maxSpeedOfSound, gaussWindow);

            diffStartSearch = timeResult1.startSearch - timeResult2.startSearch;
        }
        auto _AscanBlock = timeResult1.AscanBlockProcessed;
        auto _AscanRefBlock = timeResult2.AscanBlockProcessed;

        auto m = std::max(size(_AscanBlock, 1), size(_AscanRefBlock, 1));
        auto maxlag =
            std::max(size(_AscanBlock, 1), size(_AscanRefBlock, 1)) - 1;
        auto mxl = std::min(maxlag, m - 1);
        auto ceilLog2 = nextpow2(2 * m - 1);
        auto m2 = pow(2, ceilLog2);

        auto x = fft(_AscanBlock, m2);
        auto y = fft(_AscanRefBlock, m2);
        auto c_1_1 = x * conj(y);
        auto c_1_2 = ifft(c_1_1);

        auto c1 = real(c_1_2);
        auto c = zeros(mxl + mxl + 1, c1.getDimSize(1));
        c.setBlock(0, 0, mxl-1, c1.block(0, m2-mxl,m2-1));
        c.setBlock(0, mxl, mxl*2, c1.block(0, 0, mxl));
        // #pragma omp parallel for
        // for (size_t i = 0; i < mxl; i++) {
        //     c(i, $) = c1(m2 - mxl + i, $);
        //     c(i + mxl, $) = c1(i, $);
        // }
        // c(mxl + mxl, $) = c1(mxl, $);

        auto shiftInSamples = zeros(1, c1.getDimSize(1));
        #pragma omp parallel for
        for (size_t i = 0; i < c1.getDimSize(1); i++) {
            long rowID = 0, colID = 0;
            max(c($, i).toMatrix(), All, rowID, colID);
            shiftInSamples[i] = -maxlag + colID + rowID;
        }
        if (useTimeWindowing) {
            shiftInSamples = shiftInSamples - diffStartSearch;
        }
        auto tof = shiftInSamples / sampleRate + distBlock / aSOSWater;
        auto sosValue = distBlock / tof;
        DetectResult result;
        result.tof = tof;
        result.sosValue = sosValue;
        return result;
    }

    DetectResult detectTofAndAtt(
        const Aurora::Matrix &AscanBlock, const Aurora::Matrix &AscanRefBlock,
        const Aurora::Matrix &distBlock, const Aurora::Matrix &distRefBlock,
        int resampleFactor,int nthreads, float aSOSWater, float aSOSWaterRef,
        int useTimeWindowing, int aScanReconstructionFrequency,int detectionWindowATT,
        float offsetElectronic, int detectionWindowSOS, float minSpeedOfSound,
        float maxSpeedOfSound, bool gaussWindow)
    {
        auto sampleRate = aScanReconstructionFrequency;
        float offsetElectronicSamples = offsetElectronic * sampleRate;
        Matrix diffStartSearch;
        TimeWindowResult timeResult1;
        timeResult1.AscanBlockProcessed = AscanBlock;
        TimeWindowResult timeResult2;
        timeResult2.AscanBlockProcessed = AscanRefBlock;
        if (useTimeWindowing == 1) {
            timeResult1 = applyTimeWindowing(
                AscanBlock, sampleRate, distBlock, 
                aSOSWater, offsetElectronicSamples, detectionWindowSOS,
                minSpeedOfSound, maxSpeedOfSound, gaussWindow);
            timeResult2 = applyTimeWindowing(
                AscanRefBlock, sampleRate, distRefBlock, 
                aSOSWaterRef, offsetElectronicSamples, detectionWindowSOS,
                minSpeedOfSound, maxSpeedOfSound, gaussWindow);

            diffStartSearch = timeResult1.startSearch - timeResult2.startSearch;
        }
        auto _AscanBlock = timeResult1.AscanBlockProcessed;
        auto _AscanRefBlock = timeResult2.AscanBlockProcessed;

        auto m = std::max(size(_AscanBlock, 1), size(_AscanRefBlock, 1));
        auto maxlag =
            std::max(size(_AscanBlock, 1), size(_AscanRefBlock, 1)) - 1;
        auto mxl = std::min(maxlag, m - 1);
        auto ceilLog2 = nextpow2(2 * m - 1);
        auto m2 = pow(2, ceilLog2);

        auto x = fft(_AscanBlock, m2);
        auto y = fft(_AscanRefBlock, m2);
        auto c_1_1 = x * conj(y);
        auto c_1_2 = ifft(c_1_1);

        auto c1 = real(c_1_2);
        auto c = zeros(mxl + mxl + 1, c1.getDimSize(1));
        c.setBlock(0, 0, mxl-1, c1.block(0, m2-mxl,m2-1));
        c.setBlock(0, mxl, mxl*2, c1.block(0, 0, mxl));
        // #pragma omp parallel for
        // for (size_t i = 0; i < mxl; i++) {
        //     c(i, $) = c1(m2 - mxl + i, $);
        //     c(i + mxl, $) = c1(i, $);
        // }
        // c(mxl + mxl, $) = c1(mxl, $);

        auto shiftInSamples = zeros(1, c1.getDimSize(1));
        #pragma omp parallel for
        for (size_t i = 0; i < c1.getDimSize(1); i++) {
            long rowID = 0, colID = 0;
            max(c($, i).toMatrix(), All, rowID, colID);
            shiftInSamples[i] = -maxlag + colID + rowID;
        }
        if (useTimeWindowing) {
            shiftInSamples = shiftInSamples - diffStartSearch;
        }
        auto tof = shiftInSamples / sampleRate + distBlock / aSOSWater;
        auto sosValue = distBlock / tof;
        DetectResult result;
        result.sosValue = sosValue;
        auto tofRel = tof - distBlock / aSOSWater;
        result.att = detectAttVectorized(
            _AscanBlock, _AscanRefBlock, distRefBlock, aSOSWaterRef,
            tof, aScanReconstructionFrequency, offsetElectronic,
            detectionWindowATT);
        result.tof = tofRel;
        return result;
    }

    //remove detectTofAndAttMex function, use detectTofAndAtt function instead of
    //保留函数和函数申明备忘，以防以后再次使用 
    // DetectResult detectTofAndAttMex(
    //     const Aurora::Matrix &AscanBlock, const Aurora::Matrix &AscanRefBlock,
    //     const Aurora::Matrix &distBlock, const Aurora::Matrix &distRefBlock,
    //     const Aurora::Matrix &sosWaterBlock,
    //     const Aurora::Matrix &sosWaterRefBlock,
    //     int resampleFactor,int nthreads, float expectedSOSWater,
    //     int useTimeWindowing, int aScanReconstructionFrequency,int detectionWindowATT,
    //     float offsetElectronic, int detectionWindowSOS, float minSpeedOfSound,
    //     float maxSpeedOfSound, bool gaussWindow) 
    // {

    //     auto sizeAscan = size(AscanBlock);

    //     auto sampleRate = aScanReconstructionFrequency;
    //     float offsetElectronicSamples = offsetElectronic * sampleRate;
    //     Matrix diffStartSearch;
    //     TimeWindowResult timeResult1;
    //     timeResult1.AscanBlockProcessed = AscanBlock;
    //     TimeWindowResult timeResult2;
    //     timeResult2.AscanBlockProcessed = AscanRefBlock;
    //     if (useTimeWindowing == 1) {
    //         timeResult1 = applyTimeWindowing(
    //             AscanBlock, sampleRate, distBlock, sosWaterBlock,
    //             expectedSOSWater, offsetElectronicSamples, detectionWindowSOS,
    //             minSpeedOfSound, maxSpeedOfSound, gaussWindow);
    //         timeResult2 = applyTimeWindowing(
    //             AscanRefBlock, sampleRate, distRefBlock, sosWaterRefBlock,
    //             expectedSOSWater, offsetElectronicSamples, detectionWindowSOS,
    //             minSpeedOfSound, maxSpeedOfSound, gaussWindow);

    //         diffStartSearch = timeResult1.startSearch - timeResult2.startSearch;
    //     }
    //     auto _AscanBlock = timeResult1.AscanBlockProcessed;
    //     auto _AscanRefBlock = timeResult2.AscanBlockProcessed;
        
    //     int M =std::min(AscanBlock.getDimSize(0),AscanRefBlock.getDimSize(0));
        
    //     Matrix resDetect;
    //     float * resEnvelopeD = nullptr;
    //     float * resEnvelopeRefD = nullptr;

    //     size_t N = _AscanBlock.getDimSize(1);
    //     size_t totalSize = N*M;
    //     Matrix _AscanBlock_trim = _AscanBlock.getDimSize(0)!=M?_AscanBlock.block(0, 0, M-1):_AscanBlock;
    //     Matrix _AscanRefBlock_trim = _AscanRefBlock.getDimSize(0)!=M?_AscanRefBlock.block(0, 0, M-1):_AscanRefBlock;
    //     resDetect = Aurora::zeros(1,N);
    //     resEnvelopeD = Aurora::malloc(totalSize);
    //     resEnvelopeRefD = Aurora::malloc(totalSize);
    //     calculateBankDetectAndHilbertTransformation(
    //             _AscanBlock_trim.getData(), _AscanRefBlock_trim.getData(), N, M, resampleFactor, nthreads,
    //             resDetect.getData(), resEnvelopeD, resEnvelopeRefD);
    //     auto resEnvelope =Matrix::New(resEnvelopeD,M,N);
    //     auto resEnvelopeRef =Matrix::New(resEnvelopeRefD,M,N);
    //     //floor(size(AscanBlock, 1)*inits.resampleFactor / 2 - 1),
    //     int end_1 = std::floor(_AscanBlock.getDimSize(0)*resampleFactor/2-1);
    //     //-ceil(size(AscanBlock, 1)*inits.resampleFactor / 2)
    //     int begin_2 = -std::ceil(_AscanBlock.getDimSize(0)*resampleFactor/2);
    //     int *lags2 = new int[_AscanBlock.getDimSize(0)];
    //     for (size_t i = 0; i <= end_1; i++)
    //     {
    //         lags2[i] = i;
    //         lags2[i+end_1+1] = begin_2+i;
    //     }
    //     auto resDetectLags = zeros(1,_AscanBlock.getDimSize(1));
    //     for (size_t i = 0; i < _AscanBlock.getDimSize(1); i++)
    //     {
    //         resDetectLags[i] = lags2[(int)resDetect[i]];
    //     }
    //     delete [] lags2;
    //     resDetectLags =resDetectLags/resampleFactor;
    //     if (useTimeWindowing == 1) {
    //         resDetectLags = resDetectLags - diffStartSearch;
    //     }
    //     auto tofRel = (resDetectLags / sampleRate);
    //     auto tofAbs = tofRel + (distBlock / sosWaterBlock);
    //     auto sosValue = distBlock / tofAbs;
    //     auto tof2 = (distRefBlock / sosWaterRefBlock) * sampleRate + offsetElectronicSamples;
    //     auto startPos = zeros(_AscanBlock.getDimSize(1), 1);
    //     auto endPos = zeros(_AscanBlock.getDimSize(1), 1);
    //     auto startPosRef = zeros(_AscanBlock.getDimSize(1), 1);
    //     auto endPosRef = zeros(_AscanBlock.getDimSize(1), 1);
    //     #pragma omp parallel for
    //     for (size_t i = 0; i < _AscanBlock.getDimSize(1); i++)
    //     {
    //         startPos[i] = std::floor(std::max(tofAbs[i]*sampleRate+offsetElectronicSamples,(float)1.0));
    //         endPos[i] = std::ceil(std::min(sizeAscan[0], tofAbs[i]*sampleRate+offsetElectronicSamples+detectionWindowATT));
    //         startPosRef[i] =  std::floor(std::max( tof2[i],(float)1.0));
    //         endPosRef[i] = std::ceil(std::min(sizeAscan[0], tof2[i]+detectionWindowATT));
    //     }
    //     DetectResult result;
    //     result.att =  calculateAttenuation(resEnvelope,startPos,endPos,resEnvelopeRef,startPosRef,endPosRef);
    //     result.tof = tofRel;
    //     result.sosValue = sosValue;
    //     return result;
    // }

    
    DetectResult transmissionDetection(const Aurora::CudaMatrix &AscanBlock,
                                       const Aurora::CudaMatrix &AscanRefBlock,
                                       const Aurora::CudaMatrix &distBlock,
                                       const Aurora::CudaMatrix &distRefBlock,
                                       float aSOSWater, float aSOSWaterRef)
    {
        switch (Recon::transParams::version) {
        // case 1: {
        //     return detectTofAndAttMex(
        //         AscanBlock, AscanRefBlock, distBlock, distRefBlock,
        //         _sosWaterBlock, _sosWaterRefBlock, Recon::transParams::resampleFactor, Recon::transParams::nThreads,
        //         expectedSOSWater, Recon::transParams::useTimeWindowing,
        //         Recon::transParams::aScanReconstructionFrequency, Recon::transParams::detectionWindowATT,
        //         Recon::transParams::offsetElectronic, Recon::transParams::detectionWindowSOS, Recon::transParams::minSpeedOfSound,
        //         Recon::transParams::maxSpeedOfSound, Recon::transParams::gaussWindow);
        // }
        // case 2:
        default:
            auto r =  detectTofAndAtt(
                AscanBlock, AscanRefBlock, distBlock, distRefBlock,
                Recon::transParams::resampleFactor, Recon::transParams::nThreads,
                aSOSWater, aSOSWaterRef, Recon::transParams::useTimeWindowing,
                Recon::transParams::aScanReconstructionFrequency, Recon::transParams::detectionWindowATT,
                Recon::transParams::offsetElectronic, Recon::transParams::detectionWindowSOS, Recon::transParams::minSpeedOfSound,
                Recon::transParams::maxSpeedOfSound, Recon::transParams::gaussWindow);
            DetectResult ret;
            ret.att = r.att.toHostMatrix();
            ret.sosValue = r.sosValue.toHostMatrix();
            ret.tof = r.tof.toHostMatrix();
            return ret;
        }
    }
}