/******************************************************************************

 * (c) 2024 info@hobu.co

 *

 * SPDX-License-Identifier: MIT

 ****************************************************************************/

#include <algorithm>

#include <limits>

#include <thread>

#include "cpl_worker_thread_pool.h"

#include "memdataset.h"

#include "combiner.h"

#include "cumulative.h"

#include "notifyqueue.h"

#include "util.h"

#include "viewshed_executor.h"

namespace gdal

{

namespace viewshed

{

/// Constructor

///

/// @param opts  Options for viewshed generation.

Cumulative::Cumulative(const Options &opts) : m_opts(opts)

{

}

/// Destructor

///

Cumulative::~Cumulative() = default;

/// Compute the cumulative viewshed of a raster band.

///

/// @param srcFilename  Source filename.

/// @param pfnProgress  Pointer to the progress function. Can be null.

/// @param pProgressArg  Argument passed to the progress function

/// @return  True on success, false otherwise.

bool Cumulative::run(const std::string &srcFilename,

                     GDALProgressFunc pfnProgress, void *pProgressArg)

{

    // In cumulative mode, we run the executors in normal mode and want "1" where things

    // are visible.

    m_opts.outputMode = OutputMode::Normal;

    m_opts.visibleVal = 1;

    DatasetPtr srcDS(

        GDALDataset::FromHandle(GDALOpen(srcFilename.c_str(), GA_ReadOnly)));

    if (!srcDS)

    {

        CPLError(CE_Failure, CPLE_AppDefined, "Unable open source file.");

        return false;

    }

    GDALRasterBand *pSrcBand = srcDS->GetRasterBand(1);

    // In cumulative mode, the output extent is always the entire source raster.

    m_extent.xStop = GDALGetRasterBandXSize(pSrcBand);

    m_extent.yStop = GDALGetRasterBandYSize(pSrcBand);

    // Make a bunch of observer locations based on the spacing and stick them on a queue

    // to be handled by viewshed executors.

    for (int x = 0; x < m_extent.xStop; x += m_opts.observerSpacing)

        for (int y = 0; y < m_extent.yStop; y += m_opts.observerSpacing)

            m_observerQueue.push({x, y});

    m_observerQueue.done();

    // Run executors.

    const int numThreads = m_opts.numJobs;

    std::atomic<bool> err = false;

    std::atomic<int> running = numThreads;

    std::atomic<bool> hasFoundNoData = false;

    Progress progress(pfnProgress, pProgressArg,

                      m_observerQueue.size() * m_extent.ySize());

    CPLWorkerThreadPool executorPool(numThreads);

    for (int i = 0; i < numThreads; ++i)

        executorPool.SubmitJob(

            [this, &srcFilename, &progress, &err, &running, &hasFoundNoData]

            {

                runExecutor(srcFilename, progress, err, running,

                            hasFoundNoData);

            });

    // Run combiners that create 8-bit sums of executor jobs.

    CPLWorkerThreadPool combinerPool(numThreads);

    std::vector<Combiner> combiners(numThreads,

                                    Combiner(m_datasetQueue, m_rollupQueue));

    for (Combiner &c : combiners)

        combinerPool.SubmitJob([&c] { c.run(); });

    // Run 32-bit rollup job that combines the 8-bit results from the combiners.

    std::thread sum([this] { rollupRasters(); });

    // When the combiner jobs are done, all the data is in the rollup queue.

    combinerPool.WaitCompletion();

    if (m_datasetQueue.isStopped())

        return false;

    m_rollupQueue.done();

    // Wait for finalBuf to be fully filled.

    sum.join();

    // The executors should exit naturally, but we wait here so that we don't outrun their

    // completion and exit with outstanding threads.

    executorPool.WaitCompletion();

    if (hasFoundNoData)

    {

        CPLError(

            CE_Warning, CPLE_AppDefined,

            "Nodata value found in input DEM. Output will be likely incorrect");

    }

    // Scale the data so that we can write an 8-bit raster output.

    scaleOutput();

    if (!writeOutput(createOutputDataset(*pSrcBand, m_opts, m_extent)))

    {

        CPLError(CE_Failure, CPLE_AppDefined,

                 "Unable to write to output file.");

        return false;

    }

    progress.emit(1);

    return true;

}

/// Run an executor (single viewshed)

/// @param srcFilename  Source filename

/// @param progress  Progress supporting support.

/// @param err  Shared error flag.

/// @param running  Shared count of number of executors running.

/// @param hasFoundNoData Shared flag to indicate if a point at nodata has been encountered.

void Cumulative::runExecutor(const std::string &srcFilename, Progress &progress,

                             std::atomic<bool> &err, std::atomic<int> &running,

                             std::atomic<bool> &hasFoundNoData)

{

    DatasetPtr srcDs(GDALDataset::Open(srcFilename.c_str(), GA_ReadOnly));

    if (!srcDs)

    {

        err = true;

    }

    else

    {

        Location loc;

        while (!err && m_observerQueue.pop(loc))

        {

            DatasetPtr dstDs(MEMDataset::Create(

                "", m_extent.xSize(), m_extent.ySize(), 1, GDT_UInt8, nullptr));

            if (!dstDs)

            {

                err = true;

            }

            else

            {

                ViewshedExecutor executor(

                    *srcDs->GetRasterBand(1), *dstDs->GetRasterBand(1), loc.x,

                    loc.y, m_extent, m_extent, m_opts, progress,

                    /* emitWarningIfNoData = */ false);

                err = !executor.run();

                if (!err)

                    m_datasetQueue.push(std::move(dstDs));

                if (executor.hasFoundNoData())

                {

                    hasFoundNoData = true;

                }

            }

        }

    }

    // Job done. Set the output queue state.  If all the executor jobs have completed,

    // set the dataset output queue done.

    if (err)

        m_datasetQueue.stop();

    else

    {

        running--;

        if (!running)

            m_datasetQueue.done();

    }

}

// Add 8-bit rasters into the 32-bit raster buffer.

void Cumulative::rollupRasters()

{

    DatasetPtr pDS;

    m_finalBuf.resize(m_extent.size());

    while (m_rollupQueue.pop(pDS))

    {

        uint8_t *srcP =

            static_cast<uint8_t *>(pDS->GetInternalHandle("MEMORY1"));

        for (size_t i = 0; i < m_extent.size(); ++i)

            m_finalBuf[i] += srcP[i];

    }

}

/// Scale the output so that it's fully spread in 8 bits. Perhaps this shouldn't happen if

/// the max is less than 255?

void Cumulative::scaleOutput()

{

    uint32_t m = 0;  // This gathers all the bits set.

    for (uint32_t &val : m_finalBuf)

        m = std::max(val, m);

    if (m == 0)

        return;

    double factor =

        std::numeric_limits<uint8_t>::max() / static_cast<double>(m);

    for (uint32_t &val : m_finalBuf)

        val = static_cast<uint32_t>(std::floor(factor * val));

}

/// Write the output dataset.

/// @param pDstDS  Pointer to the destination dataset.

/// @return True if the write was successful, false otherwise.

bool Cumulative::writeOutput(DatasetPtr pDstDS)

{

    if (!pDstDS)

        return false;

    GDALRasterBand *pDstBand = pDstDS->GetRasterBand(1);

    return (pDstBand->RasterIO(GF_Write, 0, 0, m_extent.xSize(),

                               m_extent.ySize(), m_finalBuf.data(),

                               m_extent.xSize(), m_extent.ySize(), GDT_UInt32,

                               0, 0, nullptr) == 0);

}

}  // namespace viewshed

}  // namespace gdal