#include "reconstructionSAFT.h"

#include "Function1D.h"
#include "Function2D.h"
#include "Function3D.h"
#include "Matrix.h"
#include "SAFTStructs.h"
#include "SAFT_ATT.h"
#include "SAFT_TOFI.h"
#include "common/dataBlockCreation/removeDataFromArrays.h"
#include "config/config.h"

#include <algorithm>
#include <cstddef>
#include <cstdlib>
#include <iostream>
#include <iterator>
#include <sys/types.h>
#include <vector>

namespace Recon {

    int compareColumn(const Aurora::Matrix &aMatrix, int colCount,
                    const Aurora::Matrix &aOther, int col2) {

    for (size_t j = 0; j < colCount; j++) {
        size_t i = 0;
        for (; i < aMatrix.getDimSize(0); i++) {
        if (aMatrix[j * aMatrix.getDimSize(0) + 1] !=
            aOther[col2 * aOther.getDimSize(0) + i])
            break;
        }
        if (i == aMatrix.getDimSize(0))
        return j;
    }
    return -1;
    }


    Aurora::Matrix callSAFT(const Aurora::Matrix &AScans,
                            const Aurora::Matrix &receiverIndex,
                            const Aurora::Matrix &senderIndex,
                            const Aurora::Matrix &receiverPosGeom,
                            const Aurora::Matrix &senderPosGeom,
                            int SAFT_mode, double TimeInterval,
                            const TransRecos& transRecos,
                            const Aurora::Matrix &Env) 
    {
        if (reflectParams::useAscanIndex == 0) {
            if (reflectParams::soundSpeedCorrection == 0) {
                std::cerr
                    << "reconstruction without sound speed correction not allowed!!!!"
                    << std::endl;
                return Aurora::Matrix();
            }
        }
        std::vector<Matrix_t> params;
        Matrix_t AScan{nullptr,
                        (size_t)AScans.getDims(),
                        {(size_t)AScans.getDimSize(0), (size_t)AScans.getDimSize(1),
                            (size_t)AScans.getDimSize(2)},
                        AScans.getDataSize()};
        float *ASdata = new float[AScans.getDataSize()]{0};
        std::copy(AScans.getData(), AScans.getData() + AScans.getDataSize(),
                ASdata);
        AScan.Data = ASdata;
        params.push_back(AScan);

        Matrix_t imageStartpoint{nullptr, 2, {1, 3}, 3};
        imageStartpoint.Data =
            new float[3]{(float)reflectParams::imageStartpoint[0],
                        (float)reflectParams::imageStartpoint[1],
                        (float)reflectParams::imageStartpoint[2]};
        params.push_back(imageStartpoint);

        Matrix_t receiverIdx{nullptr,
                            2,
                            { 1,receiverIndex.getDataSize(), 1},
                            receiverIndex.getDataSize()};
        unsigned short *temp1 = new unsigned short[receiverIndex.getDataSize()];
        std::copy(receiverIndex.getData(),
                receiverIndex.getData() + receiverIndex.getDataSize(), temp1);
        receiverIdx.Data = temp1;
        params.push_back(receiverIdx);

        Matrix_t senderIdx{nullptr,
                        2,
                        {senderIndex.getDataSize(),1, 1},
                        senderIndex.getDataSize()};
        temp1 = new unsigned short[senderIndex.getDataSize()];
        std::copy(senderIndex.getData(),
                senderIndex.getData() + senderIndex.getDataSize(), temp1);
        senderIdx.Data = temp1;
        params.push_back(senderIdx);

        Matrix_t receiverGeom{nullptr,
                            (size_t)receiverPosGeom.getDims(),
                            {(size_t)receiverPosGeom.getDimSize(0),
                            (size_t)receiverPosGeom.getDimSize(1),
                            (size_t)receiverPosGeom.getDimSize(2)},
                            receiverPosGeom.getDataSize()};
        float *fdata = new float[receiverPosGeom.getDataSize()]{0};
        std::copy(receiverPosGeom.getData(),
                receiverPosGeom.getData() + receiverPosGeom.getDataSize(), fdata);
        receiverGeom.Data = fdata;
        params.push_back(receiverGeom);

        Matrix_t senderGeom{nullptr,
                            (size_t)senderPosGeom.getDims(),
                            {(size_t)senderPosGeom.getDimSize(0),
                            (size_t)senderPosGeom.getDimSize(1),
                            (size_t)senderPosGeom.getDimSize(2)},
                            senderPosGeom.getDataSize()};
        fdata = new float[senderPosGeom.getDataSize()]{0};
        std::copy(senderPosGeom.getData(),
                senderPosGeom.getData() + senderPosGeom.getDataSize(), fdata);
        senderGeom.Data = fdata;
        params.push_back(senderGeom);

        Matrix_t _SAFT_mode{&SAFT_mode, 1, {1, 1, 1}, 1};
        params.push_back(_SAFT_mode);

        int saftVariant_v[6]{(int)reflectParams::saftVariant[0],(int)reflectParams::saftVariant[1],(int)reflectParams::saftVariant[2],
                            (int)reflectParams::saftVariant[3],(int)reflectParams::saftVariant[4],(int)reflectParams::saftVariant[5]};
        Matrix_t saftVariant{saftVariant_v,2,{1,6,1},3};
        params.push_back(saftVariant);

        Matrix_t SpeedMap3D{nullptr,
                            (size_t)transRecos.speedMap3D.getDims(),
                            {(size_t)transRecos.speedMap3D.getDimSize(0),
                            (size_t)transRecos.speedMap3D.getDimSize(1),
                            (size_t)transRecos.speedMap3D.getDimSize(2)},
                            transRecos.speedMap3D.getDataSize()};
        fdata = new float[transRecos.speedMap3D.getDataSize()]{0};
        std::copy(transRecos.speedMap3D.getData(),
                transRecos.speedMap3D.getData() + transRecos.speedMap3D.getDataSize(), fdata);
        SpeedMap3D.Data = fdata;
        params.push_back(SpeedMap3D);

        Matrix_t BeginTransMap{nullptr,
                            (size_t)transRecos.beginTransMap.getDims(),
                            {(size_t)transRecos.beginTransMap.getDimSize(0),
                            (size_t)transRecos.beginTransMap.getDimSize(1),
                            (size_t)transRecos.beginTransMap.getDimSize(2)},
                            transRecos.beginTransMap.getDataSize()};
        fdata = new float[transRecos.beginTransMap.getDataSize()]{0};
        std::copy(transRecos.beginTransMap.getData(),
                transRecos.beginTransMap.getData() + transRecos.beginTransMap.getDataSize(), fdata);
        BeginTransMap.Data = fdata;
        params.push_back(BeginTransMap);

        Matrix_t DeltaTransMap{nullptr,
                            1,
                            1,
                            1,
                            1,
                            1};
        fdata = new float[1]{0};
        fdata[0] = transRecos.deltaTransMap;
        DeltaTransMap.Data = fdata;
        params.push_back(DeltaTransMap);

        Matrix_t AttenuationMap3D{nullptr,
                            (size_t)transRecos.attenuationMap3D.getDims(),
                            {(size_t)transRecos.attenuationMap3D.getDimSize(0),
                            (size_t)transRecos.attenuationMap3D.getDimSize(1),
                            (size_t)transRecos.attenuationMap3D.getDimSize(2)},
                            transRecos.attenuationMap3D.getDataSize()};
        fdata = new float[transRecos.attenuationMap3D.getDataSize()]{0};
        std::copy(transRecos.attenuationMap3D.getData(),
                transRecos.attenuationMap3D.getData() + transRecos.attenuationMap3D.getDataSize(), fdata);
        AttenuationMap3D.Data = fdata;
        params.push_back(AttenuationMap3D);

        float imageResolution_v = (float)reflectParams::imageResolution;
        Matrix_t imageResolution{&imageResolution_v, 1, {1, 1, 1}, 1};
        params.push_back(imageResolution);
        
        float TimeInterval_v = (float)TimeInterval;
        Matrix_t _TimeInterval{&TimeInterval_v, 1, {1, 1, 1}, 1};
        params.push_back(_TimeInterval);

        int imageXYZ_v[3]{(int)reflectParams::imageXYZ[0],(int)reflectParams::imageXYZ[1],(int)reflectParams::imageXYZ[2]};
        Matrix_t imageXYZ{imageXYZ_v,2,{1,3,1},3};
        params.push_back(imageXYZ);

        Matrix_t _Env{nullptr,
                    (size_t)Env.getDims(),
                    {(size_t)Env.getDimSize(0), (size_t)Env.getDimSize(1),
                        (size_t)Env.getDimSize(2)},
                    Env.getDataSize()};
        float *Envdata = new float[Env.getDataSize()]{0};
        std::copy(Env.getData(), Env.getData() + Env.getDataSize(),
                Envdata);
        _Env.Data = Envdata;
        params.push_back(_Env);

        int blockXYZ_v[3]{(int)reflectParams::blockDimXYZ[0],(int)reflectParams::blockDimXYZ[1],(int)reflectParams::blockDimXYZ[2]};
        Matrix_t blockXYZ{blockXYZ_v,2,{1,3,1},3};
        params.push_back(blockXYZ);

        Matrix_t gpus{nullptr,
                    2,
                    {1, (size_t)reconParams::gpuSelectionList.getDataSize(),
                        1},
                    reconParams::gpuSelectionList.getDataSize()};
        int *gdata = new int[reconParams::gpuSelectionList.getDataSize()]{0};
        std::copy(reconParams::gpuSelectionList.getData(), reconParams::gpuSelectionList.getData() + reconParams::gpuSelectionList.getDataSize(),
                gdata);
        gpus.Data = gdata;
        params.push_back(gpus);

        if (reflectParams::useAscanIndex == 0){
            Matrix_t medianWindowSize{&reflectParams::medianWindowSize,1,{1,1,1},1};
        }

        int other[2]={reflectParams::debugMode, reflectParams::attenuationCorrectionLimit} ;
        Matrix_t otherP{other,2,{1,2,1},2};
        params.push_back(otherP);
        if (reflectParams::useAscanIndex == 0) {
            auto result =  SAFT_ATT(params);
            return Aurora::Matrix::fromRawData((double *)result.Data, result.Dims[0],
                                            result.Dims[1], result.Dims[2]);
        } else {
            auto result = SAFT_TOFI(params);
            return Aurora::Matrix::fromRawData((double *)result.Data, result.Dims[0],
                                            result.Dims[1], result.Dims[2]);
        }
    }

    Aurora::Matrix
    recontructSAFT(const Aurora::Matrix &AScans, const Aurora::Matrix &senderPos,
                const Aurora::Matrix &receiverPos, const Aurora::Matrix &mpIndex,
                int SAFT_mode, TransRecos &transRecos, Aurora::Matrix &Env) {
        double TimeInterval = 1.0 / (reflectParams::aScanReconstructionFrequency);
        std::vector<int> motorPosAvailable;
        std::unique_copy(mpIndex.getData(), mpIndex.getData() + mpIndex.getDataSize(),
                         std::back_inserter(motorPosAvailable));
        for (auto mp : motorPosAvailable) {
            auto mpIdxs = mpIndex == mp;
            size_t countMp = Aurora::sum(mpIdxs).getScalar();

            auto senderPos_s = Aurora::transpose(removeDataFromArrays(senderPos, mpIdxs));
            auto receiverPos_s = Aurora::transpose(removeDataFromArrays(receiverPos, mpIdxs));
            Aurora::Matrix idsSender ;
            auto senderPosGeom = Aurora::uniqueByRows(receiverPos_s, idsSender);
            idsSender = Aurora::transpose(idsSender);
            Aurora::Matrix idsReceiver ;
            auto receiverPosGeom = Aurora::uniqueByRows(receiverPos_s, idsReceiver);
            idsReceiver = Aurora::transpose(idsReceiver);

            auto numUsedData = countMp;
            int count = numUsedData / reflectParams::blockSizeReco;
            auto blockIdxs = Aurora::zeros(1, count + 2);
            for (size_t i = 0; i < count + 1; i++) {
                blockIdxs[i] = i * reflectParams::blockSizeReco;
            }
            blockIdxs[count + 1] = numUsedData;
            for (size_t i = 0; i < count + 1; i++) {
                size_t length = blockIdxs[i + 1] - blockIdxs[i];
                size_t begin = blockIdxs[i];
                size_t end = blockIdxs[i + 1];
                auto _AScans = Aurora::zeros(AScans.getDimSize(0), length);
                auto receiverIndex = Aurora::zeros(1, length);
                auto senderIndex = Aurora::zeros(1, length);
                for (size_t j = begin, k = 0; j < end; j++, k++) {
                    _AScans(Aurora::$, k) = AScans(Aurora::$, mpIdxs[j]);
                    receiverIndex[k] = idsReceiver[j];
                    senderIndex[k] = idsSender[j];
                }
                Env = callSAFT(_AScans, receiverIndex, senderIndex, receiverPosGeom,
                            senderPosGeom, SAFT_mode, TimeInterval, transRecos, Env);
            }
        }
        return Env;
    }

    //TODO:untested
    Aurora::Matrix polyvaln(PolyModel polymodel, const Aurora::Matrix &indepvar)
    {
        auto np = Aurora::size(indepvar);
        int n = (int)np[0];
        int p = (int)np[1];
        Aurora::Matrix _indepvar = indepvar;
        if (n == 1 && polymodel.ModelTerms.getDimSize(1) == 1){
            _indepvar = Aurora::transpose(indepvar);
            np = Aurora::size(indepvar);
            n = (int)np[0];
            p = (int)np[1];
        }
        else if(polymodel.ModelTerms.getDimSize(1) != p){
            std::cerr<<"Size of indepvar array and this model are inconsistent."<<std::endl;
            return Aurora::Matrix();
        }
        size_t nt = Aurora::size(polymodel.ModelTerms,1);
        auto ypred = Aurora::zeros(n,1);
        for (size_t i = 0; i < nt; i++)
        {
            auto t = Aurora::ones(n,1);
            for (size_t j = 0; j < p; j++)
            {
                t = (t*(_indepvar(Aurora::$,j).toMatrix()^polymodel.Coefficients(i,j).toMatrix().getScalar()));
            }
            ypred = ypred+t*polymodel.Coefficients[i];
        }
        return ypred;
    }
} // namespace Recon