/****************************************************************************
 *
 *  MODULE:        r.terraflow
 *
 *  COPYRIGHT (C) 2007 Laura Toma
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *****************************************************************************/

#include <stdlib.h>
#include <ctype.h>
#include <time.h>

#include <string>
#include "fill.h"
#include "common.h"
#include "water.h"
#include "sortutils.h"
#include "streamutils.h"
#include "filldepr.h"
#include "grid.h"

/* globals in common.H

extern statsRecorder *stats;       stats file
extern userOptions *opt;           command-line options
extern struct  Cell_head *region;  header of the region
extern dimension_type nrows, ncols;
*/

#define FILL_DEBUG   if (0)
#define FILL_SAVEALL if (0)

/* defined in this module */
void recordWatersheds(AMI_STREAM<labelElevType> *labeledWater);
long assignDirections(AMI_STREAM<plateauStats> *statstr,
                      AMI_STREAM<plateauType> *platstr,
                      AMI_STREAM<waterType> *waterstr);
void assignFinalDirections(AMI_STREAM<plateauStats> *statstr,
                           AMI_STREAM<plateauType> *platstr,
                           AMI_STREAM<waterType> *waterstr);
template <class T1, class T2, class T3, class T4, class FUN>
void mergeStreamGridGrid(AMI_STREAM<T1> *grid1, AMI_STREAM<T2> *grid2,
                         dimension_type rows, dimension_type cols,
                         AMI_STREAM<T3> *str, FUN fo,
                         AMI_STREAM<T4> *outStream);
void merge2waterBase(AMI_STREAM<waterType> *unsortedWaterStr,
                     AMI_STREAM<direction_type> *dirStr,
                     AMI_STREAM<elevation_type> *elStr,
                     AMI_STREAM<waterWindowBaseType> *merge);
AMI_STREAM<waterGridType> *
merge2waterGrid(AMI_STREAM<waterType> *unsortedWaterStr,
                AMI_STREAM<direction_type> *dirStr,
                AMI_STREAM<elevation_type> *elStr);
AMI_STREAM<boundaryType> *
findBoundariesMain(AMI_STREAM<labelElevType> *labeledWater);

/* ********************************************************************** */
/* some helper classes */
/* ********************************************************************** */

class printElevation {
public:
    char *operator()(const elevation_type &p)
    {
        static char buf[20];
        snprintf(buf, sizeof(buf), "%.1f", (float)p);
        return buf;
    }
};

class printDirection {
public:
    char *operator()(const direction_type &p)
    {
        static char buf[20];
        snprintf(buf, sizeof(buf), "%3d", p);
        return buf;
    }
    char *operator()(const waterWindowBaseType &p)
    {
        static char buf[20];
        snprintf(buf, sizeof(buf), "%3d", p.dir);
        return buf;
    }
#if (0)
    char *operator()(const waterWindowBaseType &p)
    {
        static char buf[3];
        buf[0] = directionSymbol(p.dir);
        buf[1] = '\0';
        return buf;
    }
#endif
};

class printLabel {
public:
    char *operator()(const labelElevType &p)
    {
        static char buf[8];
        snprintf(buf, sizeof(buf), CCLABEL_FMT, p.getLabel());
        return buf;
    }
    char *operator()(const waterGridType &p)
    {
        static char buf[8];
        snprintf(buf, sizeof(buf), CCLABEL_FMT, p.getLabel());
        return buf;
    }
    char *operator()(const waterType &p)
    {
        static char buf[8];
        snprintf(buf, sizeof(buf), CCLABEL_FMT, p.getLabel());
        return buf;
    }
};

class printDepth {
public:
    char *operator()(const waterGridType &p)
    {
        static char buf[3];
        snprintf(buf, sizeof(buf), "%1u", p.depth);
        return buf;
    }
};

char *verbosedir(const std::string &s)
{
    static char buf[BUFSIZ];
    snprintf(buf, BUFSIZ, "dump/%s", s.c_str());
    return buf;
}

/* ---------------------------------------------------------------------- */
/* fill the terrain (if necessary) and compute flow direction stream;
   elstr must exist and contain grid data before call, filledstr and
   dirstr are created; elstr is deleted and replaced with the
   classified elstr, which has boundary nodata distinguished from
   inner nodata */
AMI_STREAM<waterWindowBaseType> *
computeFlowDirections(AMI_STREAM<elevation_type> *&elstr,
                      AMI_STREAM<elevation_type> *&filledstr,
                      AMI_STREAM<direction_type> *&dirstr,
                      AMI_STREAM<labelElevType> *&labeledWater)
{

    Rtimer rt, rtTotal;
    AMI_STREAM<elevation_type> *elstr_reclass = NULL;
    AMI_STREAM<plateauStats> *statstr = NULL;
    AMI_STREAM<plateauType> *platstr = NULL;
    AMI_STREAM<waterType> *waterstr = NULL;
    AMI_STREAM<waterGridType> *mergedWaterStr = NULL;
    AMI_STREAM<boundaryType> *boundaryStr = NULL;
    AMI_STREAM<waterWindowType> *waterWindows = NULL;

    rt_start(rtTotal);
    assert(elstr && filledstr == NULL && dirstr == NULL &&
           labeledWater == NULL);
    if (stats) {
        stats->comment("------------------------------");
        stats->comment("COMPUTING FLOW DIRECTIONS");

        /* adjust nodata -- boundary nodata distinguished from inner
           nodata */
        stats->comment("classifying nodata (inner & boundary)");
    }

    elstr_reclass = classifyNodata(elstr);
    delete elstr;
    elstr = elstr_reclass;

    /* ---------------------------------------------------------------------- */
    /* find the plateaus. */
    /* ---------------------------------------------------------------------- */
    if (stats) {
        stats->comment("----------", opt->verbose);
        stats->comment("assigning preliminary directions");
    }

    rt_start(rt);
    dirstr = new AMI_STREAM<direction_type>;
    statstr = new AMI_STREAM<plateauStats>;

    platstr = findPlateaus(elstr, nrows, ncols, nodataType::ELEVATION_NODATA,
                           dirstr, statstr);
    rt_stop(rt);

    if (stats) {
        stats->recordTime("findingPlateaus", rt);
        stats->recordLength("plateaus", platstr);
        stats->recordLength("plateau stats", statstr);
    }
    FILL_SAVEALL
    {
        /* printStream(*stats, statstr); */
        FILL_DEBUG cout << "sort plateauStr (by ij): ";
        AMI_STREAM<plateauType> *tmp = sort(platstr, ijCmpPlateauType());
        printStream2Grid(tmp, nrows, ncols, verbosedir("label1.asc"),
                         plateauType::printLabel);
        delete tmp;
    }

    /* ---------------------------------------------------------------------- */
    /* assign labels and directions & BFS ordering. depressions have
       labels, but no direction information.
    */
    /* ---------------------------------------------------------------------- */
    rt_start(rt);
    waterstr = new AMI_STREAM<waterType>();
    assignDirections(statstr, platstr, waterstr);
    delete platstr;
    delete statstr;
    rt_stop(rt);
    if (stats) {
        stats->recordTime("assigning directions", rt);
        *stats << "maxWatershedCount=" << labelFactory::getLabelCount() << endl;
    }

    rt_start(rt);
    mergedWaterStr = merge2waterGrid(waterstr, dirstr, elstr);
    delete dirstr;
    delete waterstr;
    rt_stop(rt);
    if (stats) {
        stats->recordTime("merging", rt);
        stats->recordLength("mergedWaterStr", mergedWaterStr);
    }

    /* ---------------------------------------------------------------------- */
    /* watershed analysis */
    /* IN: mergedWaterStr, labelFactory::... */
    /* ---------------------------------------------------------------------- */
    if (stats) {
        stats->comment("--------------", opt->verbose);
        stats->comment("generating watersheds and watershed graph");
    }

    rt_start(rt);
    waterWindows = new AMI_STREAM<waterWindowType>();
    createWaterWindows(mergedWaterStr, nrows, ncols, waterWindows);
    delete mergedWaterStr;
    rt_stop(rt);
    if (stats) {
        stats->recordTime("creating windows", rt);
        stats->recordLength("waterWindows", waterWindows);
    }

    /* ---------------------------------------------------------------------- */
    rt_start(rt);
    labeledWater = new AMI_STREAM<labelElevType>();
    boundaryStr = new AMI_STREAM<boundaryType>();
    generateWatersheds(&waterWindows, nrows, ncols, labeledWater, boundaryStr);

    /* do we need to make boundaries here?? */
    delete waterWindows;
    /* cout << "bogus boundary length = " << boundaryStr->stream_len() <<
     * endl;*/
    assert(boundaryStr->stream_len() == 0);
    delete boundaryStr;

    assert(labeledWater->stream_len() == (off_t)nrows * ncols);
    rt_stop(rt);
    if (stats)
        stats->recordTime("generating watersheds", rt);

    /* ---------------------------------------------------------------------- */
    /* find boundaries */
    /* ---------------------------------------------------------------------- */
    FILL_DEBUG cerr << "sort labeledWater (by ij):";
    sort(&labeledWater, ijCmpLabelElevType());

#ifdef SAVE_ASCII
    cerr << "saving WATERSHED grid as watershed_grid\n";
    printStream2Grid(labeledWater, nrows, ncols, "watershed.asc",
                     labelElevType::printLabel);
#endif
    boundaryStr = findBoundariesMain(labeledWater);

    /* ---------------------------------------------------------------------- */
    /* filling */
    /* valid streams are: boundaryStr, labeledWater */
    /* ---------------------------------------------------------------------- */
    rt_start(rt);
    elevation_type *raise;
    /*find the raise elevations */

    FILL_DEBUG cerr << "sort boundaryStr (by elev): ";
    sort(&boundaryStr, elevCmpBoundaryType());

    raise = fill_depression(boundaryStr, labelFactory::getLabelCount());
    delete boundaryStr;
    rt_stop(rt);
    if (stats)
        stats->recordTime("filling depressions", rt);

    /*fill the terrain*/
    rt_start(rt);
    filledstr = new AMI_STREAM<elevation_type>();
    commit_fill(labeledWater, raise, labelFactory::getLabelCount(), filledstr);
    assert(filledstr->stream_len() == (off_t)nrows * ncols);
    delete[] raise;
    rt_stop(rt);
    if (stats) {
        stats->recordTime("updating filled grid", rt);
        stats->recordLength("filledstr", filledstr);

        /* ----------------------------------------------------------------------
         */
        /* find plateaus again and reassign directions */
        /* ----------------------------------------------------------------------
         */

        stats->comment("------------------------------");
        stats->comment("REASSIGNING DIRECTIONS");
    }

    rt_start(rt);
    dirstr = new AMI_STREAM<direction_type>();
    statstr = new AMI_STREAM<plateauStats>();
    platstr = findPlateaus(filledstr, nrows, ncols,
                           nodataType::ELEVATION_NODATA, dirstr, statstr);
    rt_stop(rt);
    if (stats) {
        stats->recordTime("findingPlateaus2", rt);
        stats->recordLength("final plateaus", platstr);
        stats->recordLength("final plateau stats", statstr);
    }
    FILL_SAVEALL
    {
        FILL_DEBUG cout << "sort plateauStr (by ij): ";
        AMI_STREAM<plateauType> *tmp = sort(platstr, ijCmpPlateauType());
        printStream2Grid(tmp, nrows, ncols, verbosedir("plateaus.asc"),
                         plateauType::printLabel);
        delete tmp;
    }

    /* assign final directions */
    rt_start(rt);
    waterstr = new AMI_STREAM<waterType>();
    long dc = assignDirections(statstr, platstr, waterstr);
    if (dc && stats) {
        *stats << "WARNING: " << dc << " depressions (islands) detected\n";
    }
    delete platstr;
    delete statstr;
    rt_stop(rt);
    if (stats)
        stats->recordTime("final directions", rt);

    /* merge */
    rt_start(rt);
    AMI_STREAM<waterWindowBaseType> *flowStream;
    /*STREAM_TO_OPTION(flowStream, "flowStream");*/
    char path[BUFSIZ];
    char *base_dir = getenv(STREAM_TMPDIR);
    assert(base_dir);
    snprintf(path, BUFSIZ, "%s/flowStream", base_dir);
    flowStream = new AMI_STREAM<waterWindowBaseType>(path);
    /*flowStream->persist(PERSIST_PERSISTENT); */
    if (stats)
        stats->comment("creating flowStream: ", flowStream->sprint());

    merge2waterBase(waterstr, dirstr, filledstr, flowStream);
    delete waterstr;
    rt_stop(rt);
    if (stats)
        stats->recordTime("merge water,dir,elev to flow", rt);
    rt_stop(rtTotal);

#ifdef SAVE_ASCII
    /*write grids as ascii file */
    printGridStream(filledstr, nrows, ncols, "filled_elev.asc",
                    printElevation());
    printGridStream(flowStream, nrows, ncols, "direction.asc",
                    printDirection());
#endif

    if (stats) {
        stats->recordTime("Total compute flow direction running time", rtTotal);
        stats->comment("compute flow directions done.");
    }

    return flowStream;
}

/* ---------------------------------------------------------------------- */
void recordWatersheds(AMI_STREAM<labelElevType> *labeledWater)
{
    AMI_err ae;
    long labelCount = 0;
    AMI_STREAM<labelElevType> *tmp;

    if (stats)
        *stats << "--- watershed stats" << endl;
    FILL_DEBUG cout << "sort labeledWater (by wat label): ";
    tmp = sort(labeledWater, labelCmpLabelElevType());

    labelElevType *p;
    cclabel_type prev(LABEL_UNDEF);
    tmp->seek(0);
    while ((ae = tmp->read_item(&p)) == AMI_ERROR_NO_ERROR) {
        if (p->getLabel() != prev) {
            labelCount++;
            prev = p->getLabel();
        }
    }
    assert(ae == AMI_ERROR_END_OF_STREAM);

    if (stats) {
        *stats << "watershed count = " << labelCount << endl;
        *stats << "---" << endl;
        stats->flush();
    }

    delete tmp;
}

/* ********************************************************************** */
/* assign directions to plateaus that have sinks;
 * reassign labels to depressions (don't drain out).
 * all plateaus are written out to the waterstr. */
long assignDirections(AMI_STREAM<plateauStats> *statstr,
                      AMI_STREAM<plateauType> *platstr,
                      AMI_STREAM<waterType> *waterstr)
{
    size_t fmem;
    AMI_err ae;
    plateauStats *ps;

    if (stats) {
        stats->comment("----------", opt->verbose);
        stats->comment("assigning directions on plateaus");
    }

    labelFactory::reset(); /* we are relabeling now */

    statstr->seek(0);
    platstr->seek(0);
    fmem = getAvailableMemory();
    long depressionCount = 0;
    long spillCount = 0;
    while ((ae = statstr->read_item(&ps)) == AMI_ERROR_NO_ERROR) {
        if (ps->size * sizeof(gridElement) > fmem) {
            cerr << "WARNING: grid larger than memory (ignored)" << endl;
        }
        assert(ps->label != LABEL_NODATA);
        if (ps->hasSpill) {
            spillCount++;
            grid *platGrid = new grid(ps->iMin, ps->jMin, ps->iMax, ps->jMax,
                                      ps->size, ps->label);
            platGrid->load(*platstr);
            platGrid->assignDirections(opt->d8 ? 1 : 0);
            platGrid->save(*waterstr); /* this doesn't save labels */
            delete platGrid;
        }
        else {
            /* depression - just give contiguous labels only */
            depressionCount++;
            cclabel_type label = labelFactory::getNewLabel();
            for (int i = 0; i < ps->size; i++) {
                plateauType *pt;
                platstr->read_item(&pt);
                pt->cclabel = label; /* assign new label */
                waterType wt(*pt);   /* write it out */
                ae = waterstr->write_item(wt);
                assert(ae == AMI_ERROR_NO_ERROR);
            }
        }
    }
    if (stats) {
        *stats << "depression count = " << depressionCount << endl;
        *stats << "spill count = " << spillCount << endl;
    }

    return depressionCount;
}

/* ********************************************************************** */
/* assign directions to plateaus that have sinks;
 * check that there are no depressions.
 */
void assignFinalDirections(AMI_STREAM<plateauStats> *statstr,
                           AMI_STREAM<plateauType> *platstr,
                           AMI_STREAM<waterType> *waterstr)
{
    AMI_err ae;
    plateauStats *ps;

    if (stats)
        stats->comment("assigning final directions");

    statstr->seek(0);
    platstr->seek(0);
    while ((ae = statstr->read_item(&ps)) == AMI_ERROR_NO_ERROR) {

        if (ps->hasSpill) {
            grid *platGrid = new grid(ps->iMin, ps->jMin, ps->iMax, ps->jMax,
                                      ps->size, ps->label);
            platGrid->load(*platstr);
            platGrid->assignDirections(opt->d8 ? 1 : 0);
            platGrid->save(*waterstr); /* this doesn't save labels */
            delete platGrid;
        }
        else {
            /* could be legitimate */
            cerr << "WARNING: depression detected: " << *ps << endl;
            continue;
        }
    }
};

/* ********************************************************************** */
class directionElevationMerger {
public:
    waterGridType operator()(elevation_type el, direction_type dir UNUSED,
                             const waterType &p)
    {
        /* check that no (boundary) nodata values got in here */
        assert(el != nodataType::ELEVATION_BOUNDARY);
        assert(!is_nodata(el)); /* p should be a valid grid cell */
        return waterGridType(el, p.getDirection(), p.getLabel(), p.getDepth());
    }
    waterGridType operator()(elevation_type el, direction_type dir)
    {
        waterGridType wg(el, dir);
        if (el == nodataType::ELEVATION_BOUNDARY) { /* hack XXX (approved) */
            wg.setLabel(LABEL_BOUNDARY);
        }
        /* nodata => boundary or undef */
        assert(!is_nodata(el) || (wg.getLabel() == LABEL_BOUNDARY ||
                                  wg.getLabel() == LABEL_UNDEF));
        return wg;
    }
};

/* ---------------------------------------------------------------------- */
AMI_STREAM<waterGridType> *
merge2waterGrid(AMI_STREAM<waterType> *unsortedWaterStr,
                AMI_STREAM<direction_type> *dirStr,
                AMI_STREAM<elevation_type> *elStr)
{
    AMI_STREAM<waterType> *waterStr;

    FILL_DEBUG cout << "sort waterStr (by ij): ";
    waterStr = sort(unsortedWaterStr, ijCmpWaterType());

    FILL_SAVEALL printStream2Grid(waterStr, nrows, ncols,
                                  verbosedir("platlabels.asc"), printLabel());

    AMI_STREAM<waterGridType> *mergedWaterStr = new AMI_STREAM<waterGridType>();
    mergeStreamGridGrid(elStr, dirStr, nrows, ncols, waterStr,
                        directionElevationMerger(), mergedWaterStr);
    delete waterStr;
    FILL_SAVEALL printGridStream(mergedWaterStr, nrows, ncols,
                                 verbosedir("mergedlabels.asc"), printLabel());

    assert(mergedWaterStr->stream_len());
    return mergedWaterStr;
}

/* ---------------------------------------------------------------------- */
void merge2waterBase(AMI_STREAM<waterType> *unsortedWaterStr,
                     AMI_STREAM<direction_type> *dirStr,
                     AMI_STREAM<elevation_type> *elStr,
                     AMI_STREAM<waterWindowBaseType> *merge)
{
    AMI_STREAM<waterType> *waterStr;
    FILL_DEBUG cout << "sort waterStr (by ij): ";
    waterStr = sort(unsortedWaterStr, ijCmpWaterType());
    mergeStreamGridGrid(elStr, dirStr, nrows, ncols, waterStr,
                        directionElevationMerger(), merge);
    delete waterStr;
}

/* ---------------------------------------------------------------------- */
/*
 * merge 2 grids and a stream together to form a new grid.
 * str should be sorted in ij order
 */
template <class T1, class T2, class T3, class T4, class FUN>
void mergeStreamGridGrid(AMI_STREAM<T1> *grid1, AMI_STREAM<T2> *grid2,
                         dimension_type rows, dimension_type cols,
                         AMI_STREAM<T3> *str, FUN fo, AMI_STREAM<T4> *outStream)
{
    T1 *t1p;
    T2 *t2p;
    T3 *t3p;
    AMI_err aeUpd, ae;

    grid1->seek(0);
    grid2->seek(0);
    str->seek(0);
    aeUpd = str->read_item(&t3p);
    assert(aeUpd == AMI_ERROR_NO_ERROR || aeUpd == AMI_ERROR_END_OF_STREAM);

    for (dimension_type row = 0; row < rows; row++) {
        for (dimension_type col = 0; col < cols; col++) {
            ae = grid1->read_item(&t1p);
            assert(ae == AMI_ERROR_NO_ERROR);
            ae = grid2->read_item(&t2p);
            assert(ae == AMI_ERROR_NO_ERROR);

            T4 t4;
            if (aeUpd == AMI_ERROR_NO_ERROR && t3p->i == row && t3p->j == col) {
                /* cell present in stream */
                t4 = fo(*t1p, *t2p, *t3p);
                aeUpd = str->read_item(&t3p);
                assert(aeUpd == AMI_ERROR_NO_ERROR ||
                       aeUpd == AMI_ERROR_END_OF_STREAM);
            }
            else {
                /* not in stream */
                t4 = fo(*t1p, *t2p);
            }
            ae = outStream->write_item(t4);
            assert(ae == AMI_ERROR_NO_ERROR);
        }
        /*assert(outStream->stream_len() == (off_t)(row+1) * cols); */
    }
    assert(outStream->stream_len() == (off_t)rows * cols);
    return;
}

/* ---------------------------------------------------------------------- */
/* make boundaryStr from labeledWater */
AMI_STREAM<boundaryType> *
findBoundariesMain(AMI_STREAM<labelElevType> *labeledWater)
{
    AMI_STREAM<boundaryType> *boundaryStr;
    Rtimer rt;

    rt_start(rt);
    boundaryStr = new AMI_STREAM<boundaryType>();
    FILL_SAVEALL printGridStream(labeledWater, nrows, ncols,
                                 verbosedir("labels.asc"), printLabel());

    findBoundaries(labeledWater, nrows, ncols, boundaryStr);
    if (stats)
        stats->recordLength("all boundaries", boundaryStr);

    FILL_SAVEALL
    {
        FILL_DEBUG cout << "sort boundaryStr (by ij): ";
        sort(&boundaryStr, ijCmpBoundaryType());
        removeDuplicatesEx(&boundaryStr, ijCmpBoundaryType());
        printStream2Grid(boundaryStr, nrows, ncols, verbosedir("boundary.asc"),
                         boundaryType::print);
    }
    FILL_DEBUG cout << "sort boundaryStr (by wat label): ";
    sort(&boundaryStr, waterCmpBoundaryType());
    removeDuplicatesEx(&boundaryStr, boundaryCmpBoundaryType());

    rt_stop(rt);
    if (stats) {
        stats->recordTime("generating boundaries", rt);
        stats->recordLength("boundary stream", boundaryStr);
    }

    /*if(GETOPT("veryfillverbose")) printStream(cout, boundaryStr);
          SAVE_ON_OPTION(boundaryStr, "saveBoundaryStream");
    */
    return boundaryStr;
}
