/*******************************************************************************
* Copyright 2012-2020 Intel Corporation.
*
* This software and the related documents are Intel copyrighted  materials,  and
* your use of  them is  governed by the  express license  under which  they were
* provided to you (License).  Unless the License provides otherwise, you may not
* use, modify, copy, publish, distribute,  disclose or transmit this software or
* the related documents without Intel's prior written permission.
*
* This software and the related documents  are provided as  is,  with no express
* or implied  warranties,  other  than those  that are  expressly stated  in the
* License.
*******************************************************************************/

#include <math.h>
#include <memory>

#include "base.h"
#include "base_image.h"
#include "base_ipp.h"
#include "base_renderer.h"

#include "ippcore.h"
#include "ipps.h"
#include "ippi.h"
#include "ippcv.h"


static void printVersion()
{
    const IppLibraryVersion *pVersion;
    printf("\nIntel(R) IPP:\n");
    PRINT_LIB_VERSION(  , pVersion)
    PRINT_LIB_VERSION(s,  pVersion)
    PRINT_LIB_VERSION(i,  pVersion)
    PRINT_LIB_VERSION(cv, pVersion)
}

static void printHelp(const cmd::OptDef pOptions[], char* argv[])
{
    printf("\nUsage: %s [-i] InputFile [[-o] OutputFile] [Options]\n", GetProgName(argv));
    printf("Options:\n");
    cmd::OptUsage(pOptions);
}

static const Ipp8u g_mask[] =
{
    0, 0, 1, 0, 0,
    0, 1, 1, 1, 0,
    1, 1, 1, 1, 1,
    0, 1, 1, 1, 0,
    0, 0, 1, 0, 0,
};
static const IppiSizeL g_maskSize = {5, 5};


class Morphology
{
public:
    Morphology()
    {
    }

    virtual ~Morphology()
    {
        Close();
    }

    void Close()
    {
    }

    virtual Status Init(Image *pSrcImage, Image *pDstImage)
    {
        IppStatus       ippSts;
        IppSizeL        specSize = 0;

        if(!pSrcImage || !pSrcImage->ptr() || !pDstImage)
            return STS_ERR_NULL_PTR;

        Close();

        IppiSizeL ippSize = {(IppSizeL)pSrcImage->m_size.width, (IppSizeL)pSrcImage->m_size.height};

        ippSts = ippiErodeGetSpecSize_L(ippSize, g_maskSize, &specSize);
        CHECK_STATUS_PRINT_AC(ippSts, "ippiErodeGetSpecSize_L()", ippGetStatusString(ippSts), return STS_ERR_FAILED);

        m_specErode.Alloc(1, specSize);
        if(!m_specErode)
        {
            PRINT_MESSAGE("Cannot allocate memory for morphology spec");
            return STS_ERR_ALLOC;
        }

        ippSts = ippiDilateGetSpecSize_L(ippSize, g_maskSize, &specSize);
        CHECK_STATUS_PRINT_AC(ippSts, "ippiDilateGetSpecSize_L()", ippGetStatusString(ippSts), return STS_ERR_FAILED);

        m_specDilate.Alloc(1, specSize);
        if(!m_specDilate)
        {
            PRINT_MESSAGE("Cannot allocate memory for morphology spec");
            return STS_ERR_ALLOC;
        }

        ippSts = ippiErodeInit_L (ippSize, g_mask, g_maskSize, m_specErode);
        CHECK_STATUS_PRINT_AC(ippSts, "ippiErodeInit_L()", ippGetStatusString(ippSts), return STS_ERR_FAILED);

        ippSts = ippiDilateInit_L(ippSize, g_mask, g_maskSize, m_specDilate);
        CHECK_STATUS_PRINT_AC(ippSts, "ippiDilateInit_L()", ippGetStatusString(ippSts), return STS_ERR_FAILED);

        m_inter = *pSrcImage;
        m_inter.Alloc();

        m_templ = *pSrcImage;

        return STS_OK;
    }

    Status Dilate(Image *pSrcImage, Image *pDstImage, int iterations)
    {
        Status      status;
        IppStatus   ippSts;

        AutoBuffer<Ipp8u> tmpBuffer;
        IppSizeL          tmpBufferSize = 0;

        if(!pSrcImage || !pSrcImage->ptr() || !pDstImage || !pDstImage->ptr())
            return STS_ERR_NULL_PTR;

        if(m_templ != *pSrcImage || !m_specDilate)
        {
            status = Init(pSrcImage, pDstImage);
            CHECK_STATUS_PRINT_RS(status, "Resize::Init()", GetBaseStatusString(status));
        }

        if(!iterations)
            pSrcImage->CopyTo(*pDstImage);

        IppiSizeL   ippSize = {(IppSizeL)pSrcImage->m_size.width, (IppSizeL)pSrcImage->m_size.height};
        IppDataType ippType = ImageFormatToIpp(pSrcImage->m_sampleFormat);

        ippSts = ippiDilateGetBufferSize_L(ippSize, g_maskSize, ippType, pSrcImage->m_samples, &tmpBufferSize);
        CHECK_STATUS_PRINT_AC(ippSts, "ippiDilateGetBufferSize_L()", ippGetStatusString(ippSts), return STS_ERR_FAILED);

        tmpBuffer.Alloc(tmpBufferSize);
        if(!tmpBuffer)
        {
            PRINT_MESSAGE("Cannot allocate memory for morphology buffer");
            return STS_ERR_ALLOC;
        }

        Image srcSwitch = *pSrcImage;
        Image dstSwitch = (iterations > 1)?m_inter:*pDstImage;
        for(int i = 0; i < iterations; i++)
        {
            if(i)
            {
                srcSwitch = (i%2)?m_inter:*pDstImage;
                dstSwitch = (i%2)?*pDstImage:m_inter;
            }

            // Apply morphology filter
            switch(ippType)
            {
            case ipp8u:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiDilate_8u_C1R_L((Ipp8u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp8u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                case 3:  ippSts = ippiDilate_8u_C3R_L((Ipp8u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp8u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                case 4:  ippSts = ippiDilate_8u_C4R_L((Ipp8u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp8u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            case ipp16u:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiDilate_16u_C1R_L((Ipp16u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp16u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            case ipp16s:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiDilate_16s_C1R_L((Ipp16s*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp16s*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            case ipp32f:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiDilate_32f_C1R_L((Ipp32f*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp32f*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                case 3:  ippSts = ippiDilate_32f_C3R_L((Ipp32f*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp32f*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                case 4:  ippSts = ippiDilate_32f_C4R_L((Ipp32f*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp32f*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specDilate, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            default: ippSts = ippStsDataTypeErr; break;
            }
            CHECK_STATUS_PRINT_AC(ippSts, "ippiDilate()", ippGetStatusString(ippSts), return STS_ERR_FAILED);
        }
        if(iterations > 1 && (iterations%2))
            m_inter.CopyTo(*pDstImage);

        return STS_OK;
    }

    Status Erode(Image *pSrcImage, Image *pDstImage, int iterations)
    {
        Status      status;
        IppStatus   ippSts;

        AutoBuffer<Ipp8u> tmpBuffer;
        IppSizeL          tmpBufferSize = 0;

        if(!pSrcImage || !pSrcImage->ptr() || !pDstImage || !pDstImage->ptr())
            return STS_ERR_NULL_PTR;

        if(m_templ != *pSrcImage || !m_specErode)
        {
            status = Init(pSrcImage, pDstImage);
            CHECK_STATUS_PRINT_RS(status, "Resize::Init()", GetBaseStatusString(status));
        }

        if(!iterations)
            pSrcImage->CopyTo(*pDstImage);

        IppiSizeL   ippSize = {(IppSizeL)pSrcImage->m_size.width, (IppSizeL)pSrcImage->m_size.height};
        IppDataType ippType = ImageFormatToIpp(pSrcImage->m_sampleFormat);

        ippSts = ippiErodeGetBufferSize_L(ippSize, g_maskSize, ippType, pSrcImage->m_samples, &tmpBufferSize);
        CHECK_STATUS_PRINT_AC(ippSts, "ippiErodeGetBufferSize_L()", ippGetStatusString(ippSts), return STS_ERR_FAILED);

        tmpBuffer.Alloc(tmpBufferSize);
        if(!tmpBuffer)
        {
            PRINT_MESSAGE("Cannot allocate memory for morphology buffer");
            return STS_ERR_ALLOC;
        }

        Image srcSwitch = *pSrcImage;
        Image dstSwitch = (iterations > 1)?m_inter:*pDstImage;
        for(int i = 0; i < iterations; i++)
        {
            if(i)
            {
                srcSwitch = (i%2)?m_inter:*pDstImage;
                dstSwitch = (i%2)?*pDstImage:m_inter;
            }

            // Apply morphology filter
            switch(ippType)
            {
            case ipp8u:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiErode_8u_C1R_L((Ipp8u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp8u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                case 3:  ippSts = ippiErode_8u_C3R_L((Ipp8u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp8u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                case 4:  ippSts = ippiErode_8u_C4R_L((Ipp8u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp8u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            case ipp16u:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiErode_16u_C1R_L((Ipp16u*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp16u*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            case ipp16s:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiErode_16s_C1R_L((Ipp16s*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp16s*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            case ipp32f:
                switch(srcSwitch.m_samples)
                {
                case 1:  ippSts = ippiErode_32f_C1R_L((Ipp32f*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp32f*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                case 3:  ippSts = ippiErode_32f_C3R_L((Ipp32f*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp32f*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                case 4:  ippSts = ippiErode_32f_C4R_L((Ipp32f*)srcSwitch.ptr(), srcSwitch.m_step, (Ipp32f*)dstSwitch.ptr(), dstSwitch.m_step, ippSize, ippBorderRepl, NULL, m_specErode, tmpBuffer); break;
                default: ippSts = ippStsNumChannelsErr; break;
                }
                break;
            default: ippSts = ippStsDataTypeErr; break;
            }
            CHECK_STATUS_PRINT_AC(ippSts, "ippiErode()", ippGetStatusString(ippSts), return STS_ERR_FAILED);
        }
        if(iterations > 1 && (iterations%2))
            m_inter.CopyTo(*pDstImage);

        return STS_OK;
    }

protected:
    Image m_templ;

    Image m_inter;

    AutoBuffer<IppiMorphStateL> m_specDilate;
    AutoBuffer<IppiMorphStateL> m_specErode;
};

int main(int argc, char *argv[])
{
    /*
    // Variables initialization
    */
    Status       status         = STS_OK;
    DString      sInputFile     = CheckTestDirs( BMP_GRAYSCALE_FILE );
    DString      sOutputFile;
    char*        sIppCpu        = 0;
    bool         bNoWindow      = false;
    bool         bPrintHelp     = false;


    Image      srcData;
    Image      dstData;
    Morphology morph;
    int        morphIter = 1;

    // General timing
    vm_tick      tickStart   = 0;
    vm_tick      tickAcc     = 0;
    vm_tick      tickFreq    = vm_time_get_frequency()/1000;
    double       fTime       = 0;
    unsigned int iTimeLimit  = 0;
    unsigned int iLoops      = 0;
    unsigned int iLoopsLimit = 0;

    /*
    // Cmd parsing
    */
    const cmd::OptDef cmdOpts[] = {
        { 'i', "", 1, cmd::KT_DSTRING,   cmd::KF_OPTIONAL,  &sInputFile,      "input file name" },
        { 'o', "", 1, cmd::KT_DSTRING,   cmd::KF_OPTIONAL,  &sOutputFile,     "output file name" },
        { 'd', "", 1, cmd::KT_INTEGER,   0,                 &morphIter,       "initial number of dilate iterations (1 by default). Negative values apply erode" },
#if defined(ENABLE_RENDERING)
        { 's', "", 1, cmd::KT_BOOL,      0,                 &bNoWindow,       "suppress window output" },
#endif
        { 'w', "", 1, cmd::KT_POSITIVE,  0,                 &iTimeLimit,      "minimum test time in milliseconds" },
        { 'l', "", 1, cmd::KT_POSITIVE,  0,                 &iLoopsLimit,     "number of loops (overrides test time)" },
        { 'T', "", 1, cmd::KT_STRING,    0,                 &sIppCpu,         "target Intel(R) IPP optimization (" IPP_OPT_LIST ")" },
        { 'h', "", 1, cmd::KT_BOOL,      0,                 &bPrintHelp,      "print help and exit" },
        {0}
    };

    if(cmd::OptParse(argc, argv, cmdOpts))
    {
        printHelp(cmdOpts, argv);
        PRINT_MESSAGE("invalid input parameters");
        return 1;
    }

    InitPreferredCpu(sIppCpu);

    printVersion();
    
    // Check default image availability
    if ( !strcmp(sInputFile.c_str(), BMP_GRAYSCALE_FILE) ) {
        bPrintHelp = ( -1 == vm_file_access(sInputFile.c_str(), 0) );
    }

    if(bPrintHelp)
    {
        printHelp(cmdOpts, argv);
        return 0;
    }

    if(!sInputFile.Size())
    {
        printHelp(cmdOpts, argv);
        PRINT_MESSAGE("Cannot open input file");
        return 1;
    }

    for(;;)
    {
        // Read from file
        printf("\nInput file: %s\n", sInputFile.c_str());
        status = srcData.Read(sInputFile.c_str());
        CHECK_STATUS_PRINT_BR(status, "Image::Read()", GetBaseStatusString(status));
        printf("Input info: %dx%d %s\n", (int)srcData.m_size.width, (int)srcData.m_size.height, colorFormatName[srcData.m_color]);

        dstData = srcData;
        status  = dstData.Alloc();
        CHECK_STATUS_PRINT_BR(status, "Image::Alloc()", GetBaseStatusString(status));

        printf("\nOutput file: %s\n", (sOutputFile.Size())?sOutputFile.c_str():"-");
        printf("Output info: %dx%d %s\n", (int)dstData.m_size.width, (int)dstData.m_size.height, colorFormatName[dstData.m_color]);

        status = morph.Init(&srcData, &dstData);
        CHECK_STATUS_PRINT_BR(status, "Resize::Init()", GetBaseStatusString(status));

        // Rendering
        if(!bNoWindow)
        {
            WindowDraw draw("Intel(R) IPP Morphology example");
            if(draw.IsInitialized())
            {
                printf("\nPress Up key to increase dilation, press Down key to decrease dilation and increase erosion\n");
                printf("\nClose window to exit.\n");
                printf("\nIterations: %d", morphIter);

                bool bRedraw = true;
                bool bRecalc = true;
                while(!draw.IsClosed())
                {
                    vm_time_sleep(10);
                    int key = draw.CheckKey();
                    if(key == KK_UP)
                    {
                        morphIter++;
                        bRecalc = true;
                        printf(", %d", morphIter);
                    }
                    else if(key == KK_DOWN)
                    {
                        morphIter--;
                        bRecalc = true;
                        printf(", %d", morphIter);
                    }
                    if(bRecalc)
                    {
                        if(morphIter > 0)
                            morph.Dilate(&srcData, &dstData, morphIter);
                        else
                            morph.Erode(&srcData, &dstData, -morphIter);

                        bRedraw = true;
                    }
                    if(draw.IsInvalidated())
                        bRedraw = true;

                    if(bRedraw)
                    {
                        draw.DrawImage(&dstData, bRecalc);
                        bRedraw = false;
                        bRecalc = false;
                    }
                }
                printf("\n");
            }
        }
        else
            printf("\nIterations: %d\n", morphIter);

        for(iLoops = 1, tickAcc = 0;; iLoops++)
        {
            if(morphIter > 0)
            {
                tickStart = vm_time_get_tick();
                status = morph.Dilate(&srcData, &dstData, morphIter);
                tickAcc += (vm_time_get_tick() - tickStart);
                CHECK_STATUS_PRINT_BR(status, "Morphology::Dilate()", GetBaseStatusString(status));
            }
            else
            {
                tickStart = vm_time_get_tick();
                status = morph.Erode(&srcData, &dstData, -morphIter);
                tickAcc += (vm_time_get_tick() - tickStart);
                CHECK_STATUS_PRINT_BR(status, "Morphology::Erode()", GetBaseStatusString(status));
            }

            fTime = (double)tickAcc/tickFreq;
            if(iLoopsLimit)
            {
                if(iLoops >= iLoopsLimit)
                    break;
            }
            else
            {
                if(fTime >= (double)iTimeLimit)
                    break;
            }
        }
        if(status < 0) break;

        /*
        // Results output
        */
        printf("\nLoops:      %d\n", iLoops);
        printf("Time total: %0.3fms\n",  fTime);
        printf("Loop avg:   %0.3fms\n",  fTime/iLoops);

        if(sOutputFile.Size())
        {
            status = dstData.Write(sOutputFile.c_str());
            CHECK_STATUS_PRINT_BR(status, "Image::Write()", GetBaseStatusString(status));
        }

        break;
    }
    if(status < 0)
        return status;
    return 0;
}
