#pragma once

#include <richdem/common/Array2D.hpp>
#include <richdem/common/math.hpp>
#include <richdem/common/logger.hpp>
#include <richdem/common/ProgressBar.hpp>
#include <richdem/common/timer.hpp>
#include <richdem/depressions/depression_hierarchy.hpp>
#include <richdem/depressions/depressions.hpp>

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <queue>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>

namespace richdem::dephier {

constexpr double FP_ERROR = 1e-4;

///////////////////////////////////
//Function prototypes
///////////////////////////////////

template<class elev_t>
void ResetDH(DepressionHierarchy<elev_t> &deps);

template<class elev_t, class wtd_t>
void FillSpillMerge(
  const Array2D<elev_t>       &topo,
  const Array2D<dh_label_t>   &label,
  const Array2D<flowdir_t>    &flowdirs,
  DepressionHierarchy<elev_t> &deps,
  Array2D<wtd_t>              &wtd
);

template<class elev_t, class wtd_t>
static void MoveWaterIntoPits(
  const Array2D<elev_t>        &topo,
  const Array2D<dh_label_t>    &label,
  const Array2D<flowdir_t>     &flowdirs,
  DepressionHierarchy<elev_t>  &deps,
  Array2D<wtd_t>               &wtd
);

template<class elev_t>
static void MoveWaterInDepHier(
  int                                         current_depression,
  DepressionHierarchy<elev_t>                &deps,
  std::unordered_map<dh_label_t, dh_label_t> &jump_table
);

template<class elev_t>
static dh_label_t OverflowInto(
  const dh_label_t                            root,
  const dh_label_t                            stop_node,
  DepressionHierarchy<elev_t>                &deps,
  std::unordered_map<dh_label_t, dh_label_t> &jump_table,
  double                                      extra_water
);

class SubtreeDepressionInfo;

template<class elev_t, class wtd_t>
static SubtreeDepressionInfo FindDepressionsToFill(
  const int                          current_depression,
  const DepressionHierarchy<elev_t> &deps,
  const Array2D<elev_t>             &topo,
  const Array2D<dh_label_t>         &label,
  Array2D<wtd_t>                    &wtd
);

template<class elev_t, class wtd_t>
void FillDepressions(
  const flat_c_idx                      pit_cell,
  const flat_c_idx                      out_cell,
  const std::unordered_set<dh_label_t> &dep_labels,
  double                                water_vol,
  const Array2D<elev_t>                &topo,
  const Array2D<dh_label_t>            &label,
  Array2D<wtd_t>                       &wtd
);

template<class elev_t, class wtd_t>
void BackfillDepression(
  const double water_level,
  const Array2D<elev_t>   &topo,
  Array2D<wtd_t>          &wtd,
  std::vector<flat_c_idx> &cells_affected
);

template<class wtd_t>
double DetermineWaterLevel(
  wtd_t &sill_wtd,
  double water_vol,
  const double sill_elevation,
  const size_t cells_to_spread_across,
  const double total_elevation
);

template<class elev_t>
double DepressionVolume(
  const elev_t sill_elevation,
  const size_t cells_in_depression,
  const double total_elevation
);



///////////////////////////////////
//Implementations
///////////////////////////////////



///This function routes surface water from into pit cells and then distributes
///it so that it fills the bottoms of depressions, taking account of overflows.
///
///@param topo     Topography used to generate the DepressionHierarchy
///@param label    Labels from GetDepressionHierarchy indicate which depression
///                each cell belongs to.
///@param flowdirs Flowdirs generated by GetDepressionHierarchy
///@param deps     The DepressionHierarchy generated by GetDepressionHierarchy
///@param wtd      Water table depth. Values of 0 indicate saturation.
///                Negative values indicate additional water can be added to the
///                cell. Positive values indicate standing surface water.
///
///Note that from GetDepressionHierarchy we already know the number of cells
///within each depression as well as its volume.
///
///@return Modifies the depression hierarchy `deps` to indicate the amount of
///        water contained in each depression. Modifies `wtd` to indicate how
///        saturated a cell is or how much standing surface water it has.
template<class elev_t, class wtd_t>
void FillSpillMerge(
  const Array2D<elev_t>       &topo,
  const Array2D<dh_label_t>   &label,
  const Array2D<flowdir_t>    &flowdirs,
  DepressionHierarchy<elev_t> &deps,
  Array2D<wtd_t>              &wtd
){
  Timer timer_overall;
  timer_overall.start();

  RDLOG_ALG_NAME<<"FillSpillMerge";

  //If we're using the DH more than once we need to reset its water variables
  //between uses.
  ResetDH(deps);

  //We move standing water downhill to the pit cells of each depression
  MoveWaterIntoPits(topo, label, flowdirs, deps, wtd);

  {
    //Scope to limit `timer_overflow` and `jump_table`. Also ensures
    //`jump_table` frees its memory
    Timer timer_overflow;
    timer_overflow.start();
    std::unordered_map<dh_label_t, dh_label_t> jump_table;
    //Now that the water is in the pit cells, we move it around so that
    //depressions which contain too much water overflow into depressions that
    //have less water. If enough overflow happens, then the water is ultimately
    //routed to the ocean.
    MoveWaterInDepHier(OCEAN, deps, jump_table);
    RDLOG_TIME_USE<<"FlowInDepressionHierarchy: Overflow time = "<<timer_overflow.stop();
  }

  //Sanity checks
  #ifndef NDEBUG
  for(int d=1;d<(int)deps.size();d++){
      const auto &dep = deps.at(d);
      assert(dep.water_vol==0 || dep.water_vol<=dep.dep_vol);
      assert(dep.water_vol==0 || (dep.lchild==NO_VALUE && dep.rchild==NO_VALUE) || (dep.lchild!=NO_VALUE && deps.at(dep.lchild).water_vol<dep.water_vol));
      assert(dep.water_vol==0 || (dep.lchild==NO_VALUE && dep.rchild==NO_VALUE) || (dep.rchild!=NO_VALUE && deps.at(dep.rchild).water_vol<dep.water_vol));
  }
  #endif

  RDLOG_PROGRESS<<"p Finding filled...";
  Timer timer_filled;
  timer_filled.start();
  //We start at the ocean, crawl to the bottom of the depression hierarchy and
  //determine which depressions or metadepressions contain standing water. We
  //then modify `wtd` in order to distribute this water across the cells of the
  //depression which will lie below its surface.
  FindDepressionsToFill(OCEAN, deps, topo, label, wtd);
  RDLOG_TIME_USE<<"t FlowInDepressionHierarchy: Fill time = "<<timer_filled.stop()<<" s";
  RDLOG_TIME_USE<<"t FlowInDepressionHierarchy = "<<timer_overall.stop()<<" s";
}



///This function works much like a standard flow accumulation algorithm except
///that as the water moves downhill it contributes to saturating the water table
///and, when it reaches the pit cell of a depression, all the remaining water is
///moved into the DepressionHierarchy data structure for rapid flood-spill-merge
///calculations.
///
///One might think that MoveWaterIntoPits could be tested by adding a lot of
///water to a DEM, running FSM (including MoveWaterIntoPits), taking the
///resulting WTD, running MoveWaterIntoPits on it, and checking that the results
///of MoveWaterIntoPits are the same after both runs. But this isn't so: water
///from the first run comes from a broadly-distributed area. Once it is
///redistributed within a depression more water might get assigned to one pit or
///the other depending on the internal flow field.
///
///@param topo     Topography used to generate the DepressionHierarchy
///@param label    Labels from GetDepressionHierarchy indicate which depression
///                each cell belongs to.
///@param flowdirs Flowdirs generated by GetDepressionHierarchy
///@param deps     The DepressionHierarchy generated by GetDepressionHierarchy
///@param wtd      Water table depth. Values of 0 indicate saturation.
///                Negative values indicate additional water can be added to the
///                cell. Positive values indicate standing surface water.
///
///@return Modifies the depression hierarchy `deps` to indicate the amount of
///        water contained in each LEAF depression. Modifies `wtd` to indicate
///        how saturated a cell is. All values in wtd will be <=0 following this
///        operation.
template<class elev_t, class wtd_t>
static void MoveWaterIntoPits(
  const Array2D<elev_t>        &topo,
  const Array2D<dh_label_t>    &label,
  const Array2D<flowdir_t>     &flowdirs,
  DepressionHierarchy<elev_t>  &deps,
  Array2D<wtd_t>               &wtd
){
  Timer timer;
  ProgressBar progress;
  timer.start();

  //Our first step is to move all of the water downstream into pit cells. To do
  //so, we use the steepest-descent flow directions provided by the depression
  //hierarchy code

  RDLOG_PROGRESS<<"Moving surface water downstream...";

  //Calculate how many upstream cells flow into each cell
  Array2D<char>  dependencies(topo.width(),topo.height(),0);
  #pragma omp parallel for collapse(2)
  for(int y=0;y<topo.height();y++)
  for(int x=0;x<topo.width(); x++)
  for(int n=1;n<=8;n++){               //Loop through neighbours
    const int nx = x+d8x[n];           //Identify coordinates of neighbour
    const int ny = y+d8y[n];
    if(!topo.inGrid(nx,ny))
      continue;
    if(flowdirs(nx,ny)==d8_inverse[n])  //Does my neighbour flow into me?
      dependencies(x,y)++;              //Increment my dependencies
  }

  //Find the peaks. These are the cells into which no other cells pass flow (i.e. 0 dependencies). We
  //know the flow accumulation of the peaks without having to perform any
  //recursive calculations; they just pass flow downstream. From the peaks, we
  //can begin a breadth-first traversal in the downstream direction by adding
  //each cell to the frontier/queue as its dependency count drops to zero.
  std::queue<flat_c_idx> q;
  for(unsigned int i=0;i<topo.size();i++){
    if(dependencies(i)==0)// && flowdirs(i)!=NO_FLOW)  //Is it a peak?
      q.emplace(i);
  }  //Yes.

  //Starting with the peaks, pass flow downstream
  progress.start(topo.size());
  while(!q.empty()){
    ++progress;

    const auto c = q.front();          //Copy focal cell from queue
    q.pop();                           //Clear focal cell from queue


    //TODO: It would be nice to abstract the following so that we have a simple
    //way of converting the neighbour direction to the neighbour address

    //Coordinates of downstream neighbour, if any
    const auto ndir = flowdirs(c); //Neighbour direction
    int n = NO_FLOW;               //Neighbour address
    if(ndir!=NO_FLOW){
      const auto [x, y] = topo.iToxy(c);
      const int nx = x+d8x[ndir];
      const int ny = y+d8y[ndir];
      n            = topo.xyToI(nx,ny);
      assert(n>=0);
    }

    if(n == NO_FLOW){    //If this is a pit cell, move the water to the appropriate depression's water_vol.
      //Note that OCEAN cells have NO_FLOW so this will route water to the OCEAN
      //depression. This can be used to find out how much total run-off made it
      //into the ocean.
      if(wtd(c)>0){
        deps[label(c)].water_vol += wtd(c);
        wtd(c) = 0; //Clean up as we go
      }
    } else {                               //not a pit cell
      //If we have water, pass it downstream.
      if(wtd(c)>0){ //Groundwater can go negative, so it's important to make sure that we are only passing positive water around
        wtd(n) += wtd(c);  //Add water to downstream neighbour. This might result in filling up the groundwater table, which could have been negative
        wtd(c)  = 0;       //Clean up as we go
      }

      //Decrement the neighbour's dependencies. If there are no more dependencies,
      //we can process the neighbour.
      if(--dependencies(n)==0){                            //CHECK this is a hack! Something is wrong with the dependencies matrix, should never go below 0 but it sometimes does.
        assert(dependencies(n)>=0);
        q.emplace(n);                   //Add neighbour to the queue
      }
    }
  }
  progress.stop();

  RDLOG_TIME_USE<<"t FlowInDepressionHierarchy: Surface water = "<<timer.stop()<<" s";
}













///At this point all the values of the water table `wtd`, which is not used by
///this function, are <=0. The excess water, which will eventually become
///standing surface water, is stored in the DepressionHierarchy itself in the
///leaf depressions.
///
///In this function we will perform a depth-first post-order traversal of the
///depression hierarchy starting with the OCEAN. When we reach the leaf
///depressions we check if `water_vol>dep_vol`. If so, we try to overflow into
///the geographically proximal leaf depression indicated by our outlet. If there
///is not sufficient room in the depression linked to by our outlet to hold all
///the water, then the excess is passed to the depression's parent.
///
///Thus, by the time we exit this function, water will have been redistributed
///from leaf depressions as far up the depression hierarchy as necessary to
///ensure that there is sufficient volume to hold it. Any excess water is
///redistributed to the OCEAN, which is unaffected.
///
///The modified depression hierarchy's values for depression and water volume
///indicate total volumes: the volume in just the metadepression as well as its
///children.
///
///@param current_depression  We're doing a tree traversal; this is the id of
///                           the depression we're currently considering.
///@param deps                The DepressionHierarchy generated by
///                           GetDepressionHierarchy
///@param jump_table          A data structure that persists throughout the
///                           traversal and is used to skip from leaf nodes to
///                           the highest known meta-depression which still has
///                           unfilled volume. Ensures the traversal happens in
///                           O(N) time.
///
///@return Modifies the depression hierarchy `deps` to indicate the amount of
///        water in each depression. This information can be used to add
///        standing surface water to the cells within a depression.
template<class elev_t>
static void MoveWaterInDepHier(
  dh_label_t                                  current_depression,
  DepressionHierarchy<elev_t>                &deps,
  std::unordered_map<dh_label_t, dh_label_t> &jump_table
){
  if(current_depression==NO_VALUE)
    return;

  auto &this_dep = deps.at(current_depression);

  //Catch depressions that link to the ocean through this one. These are special
  //cases because we will never spread water across the union of these
  //depressions and the current depressions: they only flow into the current
  //depression. We also process the ocean-linked depressions before the children
  //of this depression because the ocean-linked depressions can flow into the
  //children (but the children do not flow into the ocean-linked).
  for(const auto c: this_dep.ocean_linked)
    MoveWaterInDepHier(c, deps, jump_table);

  //Visit child depressions. When these both overflow, then we spread water
  //across them by spreading water across their common metadepression
  MoveWaterInDepHier(this_dep.lchild, deps, jump_table);
  MoveWaterInDepHier(this_dep.rchild, deps, jump_table);

  //If the current depression is the ocean then at this point we've visited all
  //of its ocean-linked depressions (the ocean has no children). Since we do not
  //otherwise want to modify the ocean we quit here.
  if(current_depression==OCEAN)
    return;

  {
    const auto lchild = this_dep.lchild;
    const auto rchild = this_dep.rchild;

    //Only if both children are full should their water make its way to this
    //parent and, then, only if water hasn't been passed to the parent by
    //OverflowInto.
    if(lchild!=NO_VALUE //If I am a leaf, then I can't get water from my children
      && deps.at(lchild).water_vol==deps.at(lchild).dep_vol
      && deps.at(rchild).water_vol==deps.at(rchild).dep_vol
      && this_dep.water_vol==0 //If water_vol>0 then children have already overflowed into parent
    ){
      this_dep.water_vol += deps.at(lchild).water_vol + deps.at(rchild).water_vol;
    }
  }

  //Each depression has an associated dep_vol. This is the TOTAL volume of the
  //meta-depression including all of its children. This property answers the
  //question, "How much water can this meta-depression hold?"

  //Each depression also has an associated water volume called `water_vol`. This
  //stores the TOTAL volume of the water in the depression. That is, for each
  //depression this indicates how much water is in this depression and its
  //children. However, if a depression has sufficient volume to contain its
  //water, then its water_vol will not propagate up to its parent. In this way
  //we can distinguish between depressions whose water needs to be spread versus
  //metadepressions whose children might need water spread, but which will not
  //receive any spreading themselves.

  //Is the depression overflowing?
  if(this_dep.water_vol>this_dep.dep_vol) {
    //The neighbouring depression is not the ocean and this depression is
    //overflowing (therefore, all of its children are full)
    assert(this_dep.lchild==NO_VALUE || deps.at(this_dep.lchild).water_vol==deps.at(this_dep.lchild).dep_vol);
    assert(this_dep.rchild==NO_VALUE || deps.at(this_dep.rchild).water_vol==deps.at(this_dep.rchild).dep_vol);

    //OverflowInto will start at this depression and then, since there is
    //insufficient room, move to its neighbour via geolinks. If everything fills
    //up, the water will end up in this depression's parent. Afterwards, this
    //depression won't have extra water here any more. Any excess passed to the
    //parent will be dealt with it as we backtrack up the DH.
    OverflowInto(current_depression, this_dep.parent, deps, jump_table, 0);

    assert(
         this_dep.water_vol==0
      || this_dep.water_vol<=this_dep.dep_vol
      || (this_dep.lchild==NO_VALUE && this_dep.rchild==NO_VALUE)
      || (
              this_dep.lchild!=NO_VALUE && this_dep.rchild!=NO_VALUE
           && deps.at(this_dep.lchild).water_vol<this_dep.water_vol
           && deps.at(this_dep.rchild).water_vol<this_dep.water_vol
         )
    );
  }

  //All overflowing depressions should by now have overflowed all the way down
  //to the ocean. We must now spread the water in the depressions by setting
  //appropriate values for wtd.
}



//When water overflows from one depression into another, this function ensures
//that chained overflows and infiltration take place; that is, that water first
//fills the water table on the path between the two depressions and then moves
//from leaf depressions to metadepressions.
//
//A depression has three places it can put the water it's received.
//The depression will try to stash water in these three places sequentially.
//  1. It can store the water in itself
//  2. It can overflow into its neighbouring depression (by following a geolink to that depression's leaf)
//  3. It can overflow into its parent
//
//Options (2) and (3) result in a recursive call. If there's enough water,
//eventually repeated calls to (3) will bring the function to the parent of the
//depression that originally called it (through its neighbour). At this point we
//stash the water in the parent and exit.
//
//Since we might end up calling the function again and again if there's a
//complex series of overflows, the `jump_table` argument holds the location of
//where the water ultimately ended up. Everything between the original node and
//this destination is then full which means that we only traverse each part of a
//filled hierarchy once.
//
//NOTE: This function is written recursively, but the jump table means it could
//follow a long string of depressions thereby hitting the stack/recursion limit.
//A better way of writing this would be to use an iterative stack technique, but
//this would also be slightly harder to describe, so we leave it for the future.
//
//@param root         The depression we're currently considering
//@param stop_node    When we reach this depression we dump all the excess water
//                    we're carrying into it. This depression is the parent of
//                    the depression that first called this function. Reaching
//                    it means that both the original metadepression and its
//                    neighbour are both full.
///@param deps        The DepressionHierarchy generated by
///                   GetDepressionHierarchy
///@param jump_table  A data structure that persists throughout the traversal
///                   and is used to skip from leaf nodes to the highest known
///                   meta-depression which still has unfilled volume. Ensures
///                   the traversal happens in O(N) time.
///@param extra_water The amount of water left to distribute. We'll try to stash
///                   it in root. If we fail, we'll pass it to root's neighbour
///                   or, if the neighbour's full, to root's parent.
//@return The depression where the water ultimately ended up. This is used to
//        update the jump table
template<class elev_t>
static dh_label_t OverflowInto(
  const dh_label_t                            root,
  const dh_label_t                            stop_node,
  DepressionHierarchy<elev_t>                &deps,
  std::unordered_map<dh_label_t, dh_label_t> &jump_table,  //Shortcut from one depression to its next known empty neighbour
  double                                      extra_water
){
  auto &this_dep = deps.at(root);

  //If this depression is too full, then pick up its water. We'll find somewhere
  //to put it.
  if(this_dep.water_vol>this_dep.dep_vol){
    //The marginal volume of this depression is larger than what it can hold, so
    //we determine the amount that overflows, the "extra water".
    extra_water += this_dep.water_vol - this_dep.dep_vol;

    //Now that we've figured out how much excess water there is, we fill this
    //depression to its brim. Note that we don't use addition or subtraction
    //here to ensure floating-point equality.
    this_dep.water_vol = this_dep.dep_vol;
  }

  //TODO: Could simulate water running down flowpath into depression so that wtd
  //fills up more realistically

  //Determine if we've completed the loop. That'll happen either when we get to
  //the stop_node or the OCEAN. We can reach the OCEAN without reaching the
  //stop_node because the jump table can allow us to skip the stop_node.
  if(root==stop_node || root==OCEAN){
    //Otherwise, if we've reached the stop_node, that means we're at the
    //original node's parent. We give it our extra water.
    this_dep.water_vol += extra_water;
    return root;
  }

  //FIRST PLACE TO STASH WATER: IN THIS DEPRESSION

  if(this_dep.water_vol<this_dep.dep_vol){                                              //Can this depression hold any water?
    const double capacity = this_dep.dep_vol - this_dep.water_vol;                      //Yes. How much can it hold?
    if(extra_water<capacity){                                                           //Is it enough to hold all the extra water?
      this_dep.water_vol  = std::min(this_dep.water_vol+extra_water,this_dep.dep_vol);  //Yup. But let's be careful about floating-point stuff
      extra_water         = 0;                                                          //No more extra water
    } else {                                                                            //It wasn't enough to hold all the water
      this_dep.water_vol = this_dep.dep_vol;                                            //So we fill it all the way.
      extra_water       -= capacity;                                                    //And have that much less extra water to worry about
    }
  }

  //If there's no more extra water then call it quits
  if(fp_eq(extra_water,0)){
    return root;
  }

  //Jump table ensures that the two traversals collectively take no more than
  //O(N) time by ensuring that we never revisit a filled depression twice.
  if(jump_table.count(root)!=0){
    return jump_table[root] = OverflowInto(jump_table.at(root), stop_node, deps, jump_table, extra_water);
  }

  //Okay, so there's extra water and we can't fit it into this depression

  //SECOND PLACE TO STASH WATER: IN THIS DEPRESSION'S NEIGHBOUR

  if(this_dep.odep!=NO_VALUE){  //Does the depression even have such a neighbour?
    auto &odep = deps.at(this_dep.odep);
    if(odep.water_vol<odep.dep_vol){  //Can overflow depression hold more water?
      //Yes. Move the water geographically into that depression's leaf.
      return jump_table[root] = OverflowInto(this_dep.geolink, stop_node, deps, jump_table, extra_water);
    } else if(odep.water_vol>odep.dep_vol) {
      //Neighbour is overfull. Since our next stop is the parent, let's take the
      //neighbour's water with us when we go.
      extra_water += odep.water_vol-odep.dep_vol;
      odep.water_vol = odep.dep_vol;
    }
  }

  //Okay, so the extra water didn't fit into this depression or its overflow
  //depression. That means we pass it to this depression's parent. Since we're
  //going to add water to that parent, we need it to know about all the water
  //beneath it, so we add our water and our overflow neighbour's water before
  //climbing up to the parent. Since it's possible our overflow neighbour has
  //already done this, we only add the water we and our neighbour own if the
  //parent currently has no water in it. If we're ocean-linked to our parent we
  //don't want to do this since we are not contained within our parent.
  auto &pdep = deps.at(this_dep.parent);
  if(pdep.water_vol==0 && !this_dep.ocean_parent){
    pdep.water_vol += this_dep.water_vol;
    if(this_dep.odep!=NO_VALUE){
      pdep.water_vol += deps.at(this_dep.odep).water_vol;
    }
  }

  //THIRD PLACE TO STASH WATER: IN THIS DEPRESSION'S PARENT
  return jump_table[root] = OverflowInto(this_dep.parent, stop_node, deps, jump_table, extra_water);
}



//When we're determine which depressions to spread standing surface water
//across, we need to know a leaf depression to start filling from and the
//metadepression that contains all of the water. We also need to know all the
//depressions across which we're allowed to spread water. This data structure
//keeps track of this information.
class SubtreeDepressionInfo {
 public:
  //One of the depressions at the bottom of the meta-depression. We use this to
  //identify a pit cell from which to start flooding.
  dh_label_t leaf_label = NO_VALUE;
  //The metadepression containing all of the children. This metadepression is
  //guaranteed to be large enough to hold all of the water of its children plus
  //whatever exists only in the metadepression itself. We use this to determine
  //the water and depression volumes.
  dh_label_t top_label = NO_VALUE;
  //Here we keep track of which depressions are contained within the
  //metadepression. This allows us to limit the spreading function to cells
  //within the metadepression.
  std::unordered_set<dh_label_t> my_labels;
};



///This function recursively traverses the depression hierarchy until it reaches
///a leaf depression. It then starts climbing back up. If a leaf depression is
///full, it notes the leaf depression's id. This a potential place to start
///trying to flood a metadepression.
///
///At higher levels if both of a metadepression's children are full, then the
///metadepression makes a note of the ids of all the depressions contained
///within their subtrees. It also chooses arbitrarily one of the leaf
///depressions to start flooding from (since both children are full, the choice
///of a starting point for flooding doesn't matter).
///
///Eventually, a metadepression is reached which has more volume than water. At
///this point a helper function `FillDepressions()` is called. This is passed
///the ids of the arbitrarily chosen leaf depression, the partially-filled
///metadepression, and all the depressions in between. It then modifies the
///water table so that standing water rises to its natural level within the
///partially-filled metadepression.
///
///@param current_depression  The depression we're currently considering
///@param deps     The DepressionHierarchy generated by GetDepressionHierarchy
///@param topo     Topography used to generate the DepressionHierarchy
///@param label    Labels from GetDepressionHierarchy indicate which depression
///                each cell belongs to.
///@param wtd      Water table depth. Values of 0 indicate saturation.
///                Negative values indicate additional water can be added to the
///                cell. Positive values indicate standing surface water.
///@return Information about the subtree: its leaf node, depressions it
///        contains, and its root node.
template<class elev_t, class wtd_t>
static SubtreeDepressionInfo FindDepressionsToFill(
  const dh_label_t                   current_depression,    //Depression we are currently in
  const DepressionHierarchy<elev_t> &deps,                  //Depression hierarchy
  const Array2D<elev_t>             &topo,                  //Topographic data (used for determinining volumes as we're spreading stuff)
  const Array2D<dh_label_t>         &label,                 //Array indicating which leaf depressions each cell belongs to
  Array2D<wtd_t>                    &wtd                    //Water table depth
){
  //Stop when we reach one level below the leaves
  if(current_depression==NO_VALUE)
    return SubtreeDepressionInfo();

  const auto& this_dep = deps.at(current_depression);

  //We start by visiting all of the ocean-linked depressions. They don't need to
  //pass us anything because their water has already been transferred to this
  //metadepression tree by MoveWaterInDepHier(). Similar, it doesn't mater what their leaf
  //labels are since we will never spread water into them.
  for(const auto c: this_dep.ocean_linked)
    FindDepressionsToFill(c, deps, topo, label, wtd);

  //At this point we've visited all of the ocean-linked depressions. Since all
  //depressions link to the ocean and the ocean has no children, this means we
  //have visited all the depressions and spread their water. Since we don't wish
  //to modify the ocean, we are done.
  if(current_depression==OCEAN)
    return SubtreeDepressionInfo();

  //We visit both of the children. We need to keep track of info from these
  //because we may spread water across them.
  SubtreeDepressionInfo left_info  = FindDepressionsToFill(this_dep.lchild, deps, topo, label, wtd);
  SubtreeDepressionInfo right_info = FindDepressionsToFill(this_dep.rchild, deps, topo, label, wtd);

  SubtreeDepressionInfo combined;
  combined.my_labels.emplace(current_depression);
  combined.my_labels.merge(left_info.my_labels);
  combined.my_labels.merge(right_info.my_labels);

  combined.leaf_label = left_info.leaf_label;  //Choose left because right is not guaranteed to exist
  if(combined.leaf_label==NO_VALUE)            //If there's no label, then there was no child
    combined.leaf_label = current_depression;  //Therefore, this is a leaf depression

  combined.top_label = current_depression;

  //The water volume should never be greater than the depression volume because
  //otherwise we would have overflowed the water into the neighbouring
  //depression and moved the excess to the parent.
  assert(this_dep.water_vol<=this_dep.dep_vol);

  //If the depression holds less water than its volume, we can fill it now.
  //Alternatively, if its parent is ocean-linked then we need to fill the
  //depression now because the excess will have already overflowed through the
  //oceanlink and we won't get another chance (oceanlinked depressions may be at
  //the bottom of a cliff with respect to the focal depression). Alternatively,
  //if the depression's parent contains no water then we know the depression is
  //large enough to contain its water.
  if(this_dep.water_vol<this_dep.dep_vol || this_dep.ocean_parent || (this_dep.water_vol==this_dep.dep_vol && deps.at(this_dep.parent).water_vol==0)){
    assert(this_dep.water_vol<=this_dep.dep_vol);

    //If both of a depression's children have already spread their water, we do
    //not want to attempt to do so again in an empty parent depression. We check
    //to see if both children have finished spreading water.

    FillDepressions(
      deps.at(combined.leaf_label).pit_cell,
      deps.at(combined.top_label).out_cell,
      combined.my_labels,
      this_dep.water_vol,
      topo,
      label,
      wtd
    );

    //At this point there should be no more water all the way up the tree until
    //we pass through an ocean link, so we pass this up as a kind of null value.
    return SubtreeDepressionInfo();
  } else {
    return combined;
  }
}



///This function adjusts the water table to reflect standing surface water that
///has pooled at the bottom of depressions.
///
///At this point we know the id of an arbitrarily-chosen leaf depression and the
///id of the metadepression across which we will spread the water.
///
///Starting at the pit cell of the leaf depression we use a priority queue to
///climb from the lowest cells to the highest cell. For each cell we determine
///how much volume is contained within the depression if it were flooded to the
///level of this cell. If the volume is sufficent, we do the flooding.
///
///Therefore, we need an efficient way to determine the volume of a depression.
///To do so, we note that if the elevation of the current cell is $o$ and a
///depression contains cells of elevations $\{a,b,c,d\}$, then the volume of the
///depression is $(o-a)+(o-b)+(o-c)+(o-d)=4o-a-b-c-d=4o-sum(elevations)$. Thus,
///if we keep track of the number of cells in a depression and their total
///elevation, it is possible to calculate the volume of a depression at any time
///based on a hypothetical outlet level. We call this the Water Level Equation.
///
///@param pit_cell   Cell to start trying to fill from
///@param dep_labels Labels contained within the metadepression we are trying to
///                  fill
///@param water_vol How much water needs to be spread throughout the depression
///@param topo     Topography used to generate the DepressionHierarchy. Used
///                for, e.g. calculating marginal volumes as we attempt to
///                spread water
///@param label    Labels from GetDepressionHierarchy indicate which leaf
///                depression each cell belongs to.
///@param wtd      Water table depth. Values of 0 indicate saturation.
///                Negative values indicate additional water can be added to the
///                cell. Positive values indicate standing surface water. We may
///                add water to it.
///@return         N/A
template<class elev_t, class wtd_t>
void FillDepressions(
  const flat_c_idx                      pit_cell,
  const flat_c_idx                      out_cell,
  const std::unordered_set<dh_label_t> &dep_labels,
  double                                water_vol,
  const Array2D<elev_t>                &topo,
  const Array2D<dh_label_t>            &label,
  Array2D<wtd_t>                       &wtd
){
  //Nothing to do if we have no water
  if(water_vol==0)
    return;

  //Hashset stores the ids of the cells we've visited. We don't want to use a 2D
  //array because the size of the DEM as a whole could be massive and that's a
  //lot of memory to allocate/deallocate each time this function is called. We
  //arbitrarily reserve enough space in the hashset for a few thousand items.
  //This should be large than most depressions while still being small by the
  //computer's standards.
  std::unordered_set<flat_c_idx> visited(2048);

  //Priority queue that sorts cells by lowest elevation first. If two cells are
  //of equal elevation the one added most recently is popped first. The ordering
  //of the cells processed by this priority queue need not match the ordering of
  //the cells processed by the depression hierarchy.
  GridCellZk_high_pq<elev_t> flood_q;

  //Cell from which we can begin flooding the meta-depression. Which one we
  //choose is arbitrary, since we will fill all of the leaf depressions and
  //intermediate metadepressions until and including when we reach the
  //depression identified by stdi.top_label
  assert(pit_cell>=0);

  //We start flooding at the pit cell of the depression and work our way
  //upwards
  flood_q.emplace(
    pit_cell % topo.width(),
    pit_cell / topo.width(),
    topo(pit_cell)
  );

  visited.emplace(pit_cell);

  //Cells whose wtd will be affected as we spread water around
  std::vector<flat_c_idx> cells_affected;

  //Stores the sum of the elevations of all of the cells in cells_affected. Used
  //for calculating the volume we've seen so far. (See explanation above or in
  //dephier.hpp TODO)
  double total_elevation = 0;

  while(!flood_q.empty()){
    const auto c = flood_q.top();
    flood_q.pop();

    //We keep track of the current volume of the depression by noting the total
    //elevation of the cells we've seen as well as the number of cells we've
    //seen.

    //Note that the current cell's above ground volume and wtd do not contribute
    //at all. It is as though there is a virtual water line coincident with the
    //edge of the current cell. No water infiltrates into this cell or is stored
    //above it - only cells previously visited are considered when doing volume
    //calculations.

    //Current volume of this subset of the metadepression. Since we might climb
    //over a saddle point, this value can occasionally be negative. It will be
    //positive by the time we need to spread the water around.
    const double current_volume = DepressionVolume(topo(c.x,c.y), cells_affected.size(), total_elevation);

    //NOTE: If this is false by a small margin, then it's a floating point issue
    //and this should be adjusted to be >=-1e-6 and water_vol should be made 0
    assert(water_vol>=0);

    //All the cells within this depression should have water table depths less
    //than or equal to zero because we have moved all of their water down slope
    //into the pit cell. Since we may already have filled other depressions
    //their cells are allowed to have wtd>0. Thus, we raise a warning if we are
    //looking at a cell in this unfilled depression with wtd>0.
    if(dep_labels.count(label(c.x,c.y))==1 && wtd(c.x,c.y)>0)
      throw std::runtime_error("A cell was discovered in an unfilled depression with wtd>0!");

    //There are two possibilities:
    //1. The virtual water level exceeds slightly the height of the cell. The
    //   cell's water table then fills up as much as it can. The water surface
    //   is then level with the height of the cell.
    //
    //2. The water surface is below the height of the cell because there is
    //   sufficient topographic volume to hold all the water. In this case, the
    //   cell's water table is left unaffected.

    if(fp_le(water_vol,current_volume-wtd(c.x,c.y))){
      //The current scope of the depression plus the water storage capacity of
      //this cell is sufficient to store all of the water. We'll stop adding
      //cells and fill things now.

      auto water_level = DetermineWaterLevel(wtd(c.x,c.y), water_vol, topo(c.x,c.y), cells_affected.size(), total_elevation);
      if(fp_eq(water_level, topo(out_cell))){
        water_level = topo(out_cell);
      }

      //Water level must be higher than (or equal to) the previous cell we looked at, but lower than (or equal to) the current cell
      assert(cells_affected.size()==0 || fp_le(topo(cells_affected.back()),water_level));
      assert(topo(c.x,c.y)>water_level || fp_eq(topo(c.x,c.y),water_level));

      BackfillDepression(water_level, topo, wtd, cells_affected);

      //We've spread the water, so we're done
      return;
    }

    //We haven't found enough volume for the water yet.

    //The outlet cell never impacts the amount of water the depression can hold
    //and its water table is adjusted in the if-clause above.
    if(topo.xyToI(c.x,c.y)!=out_cell){
      //Add this cell to those affected so that its volume is available for
      //filling.
      cells_affected.emplace_back(topo.xyToI(c.x,c.y));

      //Fill in cells' water tables as we go
      assert(wtd(c.x,c.y)<=0);
      water_vol += wtd(c.x,c.y);  //We use += because wtd is less than or equal to zero
      wtd(c.x,c.y) = 0;           //Now we are sure that wtd is 0, since we've just filled it //TODO: Why are we sure?

      //Add the current cell's information to the running total
      total_elevation += topo(c.x,c.y);
    }

    //Add focal cell's neighbours to expand the search for more volume.
    for(int n=1;n<=8;n++){
      const auto nx = c.x + d8x[n];      // Get neighbour's x-coordinate using an offset
      const auto ny = c.y + d8y[n];      // Get neighbour's y-coordinate using an offset
      if(!topo.inGrid(nx,ny))            // Is this cell in the grid?
        continue;                        // Nope: out of bounds.
      const auto ni = topo.xyToI(nx,ny); // Get neighbour's flat index

      //Don't add cells which are not part of the depression unless the cell in
      //question is the outlet cell.
      if(dep_labels.count(label(ni))==0 && ni!=out_cell)
        continue;

      if(topo(nx,ny) > topo(out_cell))  //We may not add any cells higher than the out_cell. Without this check, we would sometimes look at higher cells before checking out_cell for the second time.
        continue;
      //Ocean cells may be found at the edge of a depression. They might get
      //added to this list even if there are other, lower, cells within the
      //depression which have not yet been explored. This happens when a flat
      //abutts an ocean. The side of the flat near the ocean will see the
      //ocean and try to add it. The ocean would then be called instead of
      //more cells within the depression. Therefore, we do not add ocean
      //cells.
      if(visited.count(ni)==0){
        flood_q.emplace(nx,ny,topo(nx,ny));
        visited.emplace(ni);
      }
    }

    //The queue is empty, so we add the outlet cell. We may have already visited
    //the cell while moving from one side of a depression to the other, but
    //since it doesn't affect depression volume it's safe to visit it twice.
    //NOTE: if our logic is wrong somewhere adding the cell like this could
    //result in an infinite loop.
    if(flood_q.empty()){
      const auto [x, y] = topo.iToxy(out_cell);
      flood_q.emplace(x, y, topo(out_cell));
      visited.emplace(out_cell);
    }
  }

  //Since we're in this function we are supposed to be guaranteed to be able to
  //fill our depression, since we only enter this function once that is true.
  //Therefore, if we've reached this point, something has gone horribly wrong
  //somewhere. :-(

  throw std::runtime_error("PQ loop exited without filling a depression!");
}



///Calculates the volume of a depression contained by a sill
///@param sill_elevation  Elevation of the sill
///@param cells_in_depression Number of cells contained by the sill
///@param total_elevation Sum of the elevations of cells contained by the sill
///
///@return The volume of the depression dammed by a sill of the given elevation
template<class elev_t>
double DepressionVolume(
  const elev_t sill_elevation,
  const size_t cells_in_depression,
  const double total_elevation
){
  return cells_in_depression * static_cast<double>(sill_elevation) - total_elevation;
}


///Once a set of cells sufficient to hold the water in a depression has been
///found, we use this function to spread the water across those cells.
///@param water_level    Elevation of the water's surface
///@param topo           Elevations of cells
///@param wtd            Water table depth
///@param cells_affected Queue of cells to add water to
template<class elev_t, class wtd_t>
void BackfillDepression(
  const double water_level,
  const Array2D<elev_t>   &topo,
  Array2D<wtd_t>          &wtd,
  std::vector<flat_c_idx> &cells_affected
){
  for(const auto c: cells_affected){
    //While searching for cells to fill, we simulated infiltration. The result
    //is that surface water should now be covering cells whose water table is
    //flush with the surface.
    assert(wtd(c)==0);
    //Water level should be equal to or greater than the elevation of each cell
    //to be filled
    if(water_level<topo(c)){
      assert(fp_ge(water_level, topo(c)));
    }
    //Raise the cell's water table so that it is flush with the water's surface
    //(Note that the elevation of the surface is topo+wtd)
    wtd(c) = water_level - topo(c);
    //The water table must have been zero when we added the water, but there's a
    //chance floating point effects will make it negative. Therefore, we make
    //sure that it's >=0. This is most likely to be necessary when considering
    //saddle cells within a metadepression.
    wtd(c) = std::max(static_cast<wtd_t>(0), wtd(c));
    assert(wtd(c)>=0);
  }
}



///The Lake-Level Equation.
///@param sill_wtd               Water table depth of the sill cell
///@param water_vol              Volume of water to spread across the depression
///@param sill_elevation         Elevation of the cell that dams the water into
///                              the depression.
///@param cells_to_spread_across How many cells the water is to be spread across
///@param total_elevation        The sum of those cells' elevations
///
///@return The elevation of the water level
template<class wtd_t>
double DetermineWaterLevel(
  wtd_t &sill_wtd,
  double water_vol,
  const double sill_elevation,
  const size_t cells_to_spread_across,
  const double total_elevation
){
  const double current_dep_volume = DepressionVolume(sill_elevation, cells_to_spread_across, total_elevation);

  if(water_vol>current_dep_volume){
    if(fp_eq(water_vol,current_dep_volume))
      water_vol = current_dep_volume;

    //The volume of water exceeds what we can hold above ground, so we will
    //stash as much as we can in this cell's water table. This is okay
    //because the above ground volume plus this cell's water table IS enough
    //volume or we wouldn't be in this function. But we assert that it is so.
    assert(fp_le(water_vol,current_dep_volume-sill_wtd));

    //Fill in as much of this cell's water table as we can
    const double fill_amount = water_vol - current_dep_volume;
    assert(fill_amount>=0);
    assert(fill_amount<=-sill_wtd);
    sill_wtd += fill_amount;
    //At this point we would subtract `fill_amount` from `water_vol` because
    //we've just put that into the sill's water table. But we're not using
    //`water_vol` any more, so we just forget about it.
    //water_vol   -= fill_amount;
    return sill_elevation;
  } else if (fp_eq(water_vol,current_dep_volume)){
    //The volume of water is exactly equal to the above ground volume so we
    //set the water level equal to this cell's elevation
    return sill_elevation;
  } else {
    //The water volume is less than this cell's elevation, so we calculate
    //what the water level should be.

    //We have that Volume = (Water Level)*(Cell Count)-(Total Elevation)
    //rearranging this gives us:
    const auto nominal_level = (water_vol+total_elevation)/cells_to_spread_across;
    if(fp_eq(nominal_level, sill_elevation))
      return sill_elevation;
    else
      return nominal_level;
  }
}



///Reset water volumes in the Depression Hierarchy to zero
///
///@param deps Depression Hierarchy to reset
template<class elev_t>
void ResetDH(DepressionHierarchy<elev_t> &deps){
  for(auto &dep: deps){
    dep.water_vol = 0;
  }
}

}
