//! Hydrology: where the water sits and where it runs.
//!
//! Three jobs, in order:
//!
//! 1. [`water_bodies`] floods the sea inward from the map border so that
//!    below-sea-level basins that *aren't* connected to the open ocean become
//!    inland **lakes**.
//! 2. [`flow_accumulation`] first **fills depressions** with a priority-flood
//!    (Barnes, Lehman & Mulla 2014) so every cell has a strictly downhill path
//!    to the border — without this, rivers dead-end in noise pits. It then routes
//!    each cell's rainfall downhill (steepest D8 descent) and sums the drainage
//!    passing through every cell.
//! 3. [`river_cells`] thresholds that drainage at a chosen density to pick out
//!    the channels that read as rivers.
//!
//! No randomness and no transcendental math: given the same elevation/moisture,
//! the water is identical on every machine. Float ordering for the priority
//! queue uses `f32::total_cmp` with the cell index as a deterministic tie-break.

use std::cmp::{Ordering, Reverse};
use std::collections::{BinaryHeap, VecDeque};

use serde::{Deserialize, Serialize};

use crate::geometry::Point;
use crate::grid::Grid;

use super::WorldConfig;

/// Which kind of standing water (if any) fills a cell, before rivers are traced.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum WaterBody {
    /// Dry land (elevation at or above sea level).
    Land,
    /// Below sea level and connected to the map edge — the open sea.
    Ocean,
    /// Below sea level but walled off from the open sea — an inland lake.
    Lake,
}

/// Classifies every cell as [`Land`](WaterBody::Land), [`Ocean`](WaterBody::Ocean),
/// or [`Lake`](WaterBody::Lake).
///
/// A cell is water when its elevation is below `cfg.sea_level`. The open ocean is
/// the set of water cells reachable (4-connected) from any below-sea-level cell
/// on the map border; water cells the flood can't reach are inland lakes. This is
/// what lets a high plateau hold a sea without it silently merging with the coast.
pub fn water_bodies(elevation: &Grid<f32>, cfg: &WorldConfig) -> Grid<WaterBody> {
    let (w, h) = (elevation.width(), elevation.height());
    let is_water = |p: Point| *elevation.get(p).unwrap_or(&1.0) < cfg.sea_level;

    let mut body = Grid::new(w, h, WaterBody::Land);
    // Mark all water cells as Lake first; the flood promotes the connected ones.
    for (p, _) in elevation.iter() {
        if is_water(p) {
            body.set(p, WaterBody::Lake);
        }
    }

    // BFS the open ocean inward from every below-sea border cell.
    let mut queue: VecDeque<Point> = VecDeque::new();
    let on_border = |p: Point| p.x == 0 || p.y == 0 || p.x == w as i32 - 1 || p.y == h as i32 - 1;
    for (p, _) in elevation.iter() {
        if on_border(p) && is_water(p) {
            body.set(p, WaterBody::Ocean);
            queue.push_back(p);
        }
    }
    while let Some(p) = queue.pop_front() {
        // Collect first: `neighbors4` borrows `body`, and we mutate `body` below.
        let neighbors: Vec<Point> = body.neighbors4(p).collect();
        for n in neighbors {
            if is_water(n) && body.get(n) == Some(&WaterBody::Lake) {
                body.set(n, WaterBody::Ocean);
                queue.push_back(n);
            }
        }
    }

    body
}

/// A cell waiting in the priority-flood queue, ordered by its spill elevation
/// (lowest first), with the flat cell index as a deterministic tie-break.
struct FloodCell {
    spill: f32,
    idx: u32,
}
impl PartialEq for FloodCell {
    fn eq(&self, o: &Self) -> bool {
        self.idx == o.idx && self.spill.to_bits() == o.spill.to_bits()
    }
}
impl Eq for FloodCell {}
impl Ord for FloodCell {
    fn cmp(&self, o: &Self) -> Ordering {
        // total_cmp gives a total order over f32 (no NaN surprises); the index
        // tie-break makes equal-elevation pops fully deterministic.
        self.spill.total_cmp(&o.spill).then(self.idx.cmp(&o.idx))
    }
}
impl PartialOrd for FloodCell {
    fn partial_cmp(&self, o: &Self) -> Option<Ordering> {
        Some(self.cmp(o))
    }
}

/// Fills depressions and returns the drainage **accumulation** field: each
/// cell's value is the total upstream rainfall that drains through it.
///
/// Rainfall per cell is `1 + river_moisture_weight * moisture`, so wetter
/// uplands feed larger rivers, but every land cell contributes at least its own
/// unit (drainage area). Water cells originate no flow but pass it through to the
/// sea. The returned grid shares the input's dimensions.
pub fn flow_accumulation(
    elevation: &Grid<f32>,
    moisture: &Grid<f32>,
    water: &Grid<WaterBody>,
    cfg: &WorldConfig,
) -> Grid<f32> {
    let (w, h) = (elevation.width(), elevation.height());
    let n = (w as usize) * (h as usize);
    if n == 0 {
        return Grid::new(w, h, 0.0);
    }

    let to_pt = |i: u32| Point::new((i % w) as i32, (i / w) as i32);
    let to_idx = |p: Point| (p.y as u32) * w + p.x as u32;

    // --- 1. Priority-flood (+epsilon) depression fill ----------------------
    //
    // `filled` is the elevation raised just enough that every cell has a
    // strictly downhill route to the border. We never display it — it only
    // defines flow direction — so the tiny epsilon creep is harmless.
    const EPS: f32 = 1e-6;
    let mut filled = vec![f32::INFINITY; n];
    let mut closed = vec![false; n];
    let mut heap: BinaryHeap<Reverse<FloodCell>> = BinaryHeap::new();

    let on_border = |p: Point| p.x == 0 || p.y == 0 || p.x == w as i32 - 1 || p.y == h as i32 - 1;
    for (p, &e) in elevation.iter() {
        if on_border(p) {
            let i = to_idx(p) as usize;
            filled[i] = e;
            closed[i] = true;
            heap.push(Reverse(FloodCell { spill: e, idx: to_idx(p) }));
        }
    }

    while let Some(Reverse(cell)) = heap.pop() {
        let cp = to_pt(cell.idx);
        for np in elevation.neighbors8(cp) {
            let ni = to_idx(np) as usize;
            if closed[ni] {
                continue;
            }
            closed[ni] = true;
            // Spill level: at least the neighbor's own ground, but never below
            // the cell we came from (plus epsilon to guarantee descent).
            let raised = elevation.get(np).copied().unwrap_or(0.0).max(cell.spill + EPS);
            filled[ni] = raised;
            heap.push(Reverse(FloodCell { spill: raised, idx: to_idx(np) }));
        }
    }

    // --- 2. Flow direction: steepest D8 descent on the filled surface ------
    //
    // `downstream[i]` is the flat index this cell drains into, or `i` itself for
    // a sink (a border outlet with no lower neighbor).
    let mut downstream = vec![0u32; n];
    for (p, _) in elevation.iter() {
        let i = to_idx(p) as usize;
        let here = filled[i];
        let mut best = here;
        let mut best_idx = to_idx(p);
        for np in elevation.neighbors8(p) {
            let ni = to_idx(np) as usize;
            if filled[ni] < best {
                best = filled[ni];
                best_idx = to_idx(np);
            }
        }
        downstream[i] = best_idx;
    }

    // --- 3. Accumulate rainfall from peaks down --------------------------
    //
    // Processing cells in descending filled-elevation order guarantees that when
    // we hand a cell's water downstream, every cell upstream of it has already
    // contributed.
    let precip = |p: Point| -> f32 {
        match water.get(p) {
            Some(WaterBody::Land) => 1.0 + cfg.river_moisture_weight * *moisture.get(p).unwrap_or(&0.0),
            _ => 0.0, // open water originates no river flow
        }
    };
    let mut acc = vec![0.0f32; n];
    for (p, _) in elevation.iter() {
        acc[to_idx(p) as usize] = precip(p);
    }

    let mut order: Vec<u32> = (0..n as u32).collect();
    order.sort_by(|&a, &b| {
        // Highest filled first; index tie-break keeps it deterministic.
        filled[b as usize]
            .total_cmp(&filled[a as usize])
            .then(a.cmp(&b))
    });
    for &i in &order {
        let d = downstream[i as usize];
        if d != i {
            acc[d as usize] += acc[i as usize];
        }
    }

    Grid::from_fn(w, h, |p| acc[to_idx(p) as usize])
}

/// Picks the river channels from an accumulation field at a given `density`.
///
/// `density` is the fraction of **land** cells that should read as river
/// (`0.05` ≈ the top 5% most-drained land cells). Because flow accumulation only
/// grows downstream, the chosen set is naturally contiguous — each river runs
/// unbroken from its headwater down to the sea. Returns a boolean mask the size
/// of the inputs; lakes and oceans are never rivers.
pub fn river_cells(acc: &Grid<f32>, water: &Grid<WaterBody>, density: f32) -> Grid<bool> {
    let is_land = |p: Point| matches!(water.get(p), Some(WaterBody::Land));

    // Sort land accumulations to find the density-th percentile threshold.
    let mut land_acc: Vec<f32> = acc
        .iter()
        .filter(|(p, _)| is_land(*p))
        .map(|(_, &v)| v)
        .collect();
    if land_acc.is_empty() {
        return Grid::new(acc.width(), acc.height(), false);
    }
    land_acc.sort_by(|a, b| a.total_cmp(b));

    let d = density.clamp(0.0, 1.0);
    // Index of the first cell that qualifies: (1 - d) of the way up the sorted
    // list. A density of 0 yields a threshold above every cell (no rivers).
    let cut = (((1.0 - d) * land_acc.len() as f32) as usize).min(land_acc.len() - 1);
    let threshold = land_acc[cut];

    Grid::from_fn(acc.width(), acc.height(), |p| {
        is_land(p) && *acc.get(p).unwrap_or(&0.0) >= threshold && d > 0.0
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    /// An inland basin below sea level becomes a Lake, while a below-sea cell
    /// touching the border (and its connected neighbors) becomes Ocean.
    #[test]
    fn lakes_are_disconnected_from_the_ocean() {
        let cfg = WorldConfig { sea_level: 0.5, ..WorldConfig::default() };
        // 5x5 of high land, with a single low cell in the dead center (a lake)
        // and the whole left column low (ocean reaching the border).
        let mut elev = Grid::new(5, 5, 0.8);
        elev.set(Point::new(2, 2), 0.1); // landlocked low cell
        for y in 0..5 {
            elev.set(Point::new(0, y), 0.1); // low border column → ocean
        }
        let body = water_bodies(&elev, &cfg);
        assert_eq!(body.get(Point::new(2, 2)), Some(&WaterBody::Lake));
        assert_eq!(body.get(Point::new(0, 2)), Some(&WaterBody::Ocean));
        assert_eq!(body.get(Point::new(3, 3)), Some(&WaterBody::Land));
    }

    /// Accumulation strictly increases downstream: a cell's drainage is at least
    /// the sum it received, so the coast carries more water than the headwaters.
    /// We build a simple tilted plane draining to the right edge and check the
    /// mouth out-drains the source.
    #[test]
    fn accumulation_grows_toward_the_outlet() {
        let cfg = WorldConfig { sea_level: 0.0, river_moisture_weight: 0.0, ..WorldConfig::default() };
        // Elevation decreasing left→right so everything drains east.
        let w = 10u32;
        let elev = Grid::from_fn(w, 3, |p| (w as i32 - p.x) as f32 / w as f32);
        let moist = Grid::new(w, 3, 0.0);
        let water = Grid::new(w, 3, WaterBody::Land);
        let acc = flow_accumulation(&elev, &moist, &water, &cfg);

        // Elevation decreases left→right, so everything drains east: the eastern
        // outlet out-drains a western headwater cell.
        let source = *acc.get(Point::new(1, 1)).unwrap();
        let mouth = *acc.get(Point::new(8, 1)).unwrap();
        assert!(mouth > source, "mouth {mouth} should out-drain source {source}");
    }

    /// A pit in the middle of a slope doesn't swallow the river: the priority
    /// flood fills it so flow still reaches the outlet (accumulation downstream
    /// of the pit exceeds the pit's own rainfall).
    #[test]
    fn depressions_are_filled_so_flow_passes_through() {
        let cfg = WorldConfig { sea_level: 0.0, river_moisture_weight: 0.0, ..WorldConfig::default() };
        let w = 10u32;
        let mut elev = Grid::from_fn(w, 3, |p| (w as i32 - p.x) as f32 / w as f32);
        elev.set(Point::new(5, 1), 0.0); // a deep pit mid-slope
        let moist = Grid::new(w, 3, 0.0);
        let water = Grid::new(w, 3, WaterBody::Land);
        let acc = flow_accumulation(&elev, &moist, &water, &cfg);
        // Cells east of (downstream of) the pit still accumulate a lot of water,
        // proving the river wasn't swallowed by the depression.
        let downstream = *acc.get(Point::new(8, 1)).unwrap();
        assert!(downstream > 5.0, "flow should pass the pit, got {downstream}");
    }

    /// `river_cells` honors density: a higher density marks at least as many
    /// river cells as a lower one, and density 0 marks none.
    #[test]
    fn river_density_is_monotonic() {
        let cfg = WorldConfig { sea_level: 0.0, river_moisture_weight: 0.0, ..WorldConfig::default() };
        let w = 16u32;
        let elev = Grid::from_fn(w, w, |p| (w as i32 - p.x) as f32 / w as f32);
        let moist = Grid::new(w, w, 0.5);
        let water = Grid::new(w, w, WaterBody::Land);
        let acc = flow_accumulation(&elev, &moist, &water, &cfg);

        let count = |d: f32| river_cells(&acc, &water, d).iter().filter(|(_, &b)| b).count();
        assert_eq!(count(0.0), 0, "density 0 marks no rivers");
        assert!(count(0.20) >= count(0.05), "more density ⇒ at least as many river cells");
        assert!(count(0.20) > 0, "a positive density should mark some rivers");
    }
}
