//! D8 drainage routing — flow direction, flow accumulation, river network
//! extraction, and drainage basin delineation (D-208).
//!
//! **Determinism (D-010, D-208):** All flow-direction comparisons use integer
//! arithmetic on scaled elevation values (`(elev * 1_000_000.0) as i64`) to
//! avoid f32 comparison non-determinism. Tie-breaking uses a fixed D8 neighbor
//! priority order. The result is bit-identical across runs on the same inputs.
//!
//! **Algorithm:**
//! 1. Scale f32 elevation to i64 integers.
//! 2. Priority-flood depression fill (iterative, convergence in ≤10 passes).
//! 3. D8 flow direction: steepest descent, 8-neighbor, wraps horizontally.
//! 4. Flow accumulation via topological sort of the D8 DAG.
//! 5. River network extraction: cells with accumulation > RIVER_THRESHOLD.
//! 6. Basin labeling: flood-fill seeded at pour points.
//!
//! The grid is row-major. Row 0 is the north pole; row H-1 is the south pole.
//! Columns wrap horizontally (the globe is equirectangular).

use std::collections::VecDeque;

use crate::atlas::body_world_state::{DrainageBasin, RiverNetwork};

/// A cell is a river cell when its flow accumulation exceeds this threshold (D-208).
pub const RIVER_THRESHOLD: i32 = 200;

/// Scale factor for converting f32 elevation to integer for deterministic comparison.
const ELEV_SCALE: f64 = 1_000_000.0;

// D8 neighbor offsets (dr, dc) in fixed priority order for deterministic tie-breaking.
// Priority: cardinal directions first (N, S, E, W), then diagonals (NE, NW, SE, SW).
const D8: [(i32, i32); 8] = [
    (-1, 0),  // N
    (1, 0),   // S
    (0, 1),   // E
    (0, -1),  // W
    (-1, 1),  // NE
    (-1, -1), // NW
    (1, 1),   // SE
    (1, -1),  // SW
];

/// Result of the full D8 drainage analysis for one body.
#[derive(Debug, Clone)]
pub struct DrainageResult {
    pub river_network: RiverNetwork,
    pub drainage_basins: Vec<DrainageBasin>,
}

// ---------------------------------------------------------------------------
// Public entry point
// ---------------------------------------------------------------------------

/// Run the full D8 drainage analysis on an elevation grid.
///
/// `elevation` is a row-major float32 grid of shape `height × width`, values
/// in [0.0, 1.0]. `sea_level` is the fraction below which terrain is ocean.
///
/// Returns `DrainageResult` with the river network and drainage basins.
pub fn analyze(elevation: &[f32], width: u32, height: u32, sea_level: f32) -> DrainageResult {
    let w = width as usize;
    let h = height as usize;
    // 1. Scale to integers.
    let scaled: Vec<i64> = elevation
        .iter()
        .map(|&e| (e as f64 * ELEV_SCALE) as i64)
        .collect();

    // 2. Depression fill.
    let filled = depression_fill(&scaled, w, h);

    // 3. D8 flow direction. -1 = no outflow (edge or flat peak).
    let fdir = flow_direction(&filled, w, h);

    // 4. Flow accumulation.
    let accum = flow_accumulation(&fdir, w, h);

    // 5. River network.
    let river_network = extract_river_network(&accum, &fdir, w, h, sea_level, elevation);

    // 6. Basin labeling.
    let labels = label_basins(&fdir, &accum, w, h);

    // 7. Merge small basins + clamp count to [4, 12].
    let labels = merge_small_basins(labels, w, h, 4, 12);

    // 8. Build DrainageBasin structs.
    let drainage_basins = build_basins(&labels, w, h);

    DrainageResult {
        river_network,
        drainage_basins,
    }
}

// ---------------------------------------------------------------------------
// Step 2: Depression fill
// ---------------------------------------------------------------------------

fn depression_fill(scaled: &[i64], w: usize, h: usize) -> Vec<i64> {
    let mut filled = scaled.to_vec();
    for _ in 0..10 {
        let mut changed = false;
        for r in 1..h.saturating_sub(1) {
            for c in 0..w {
                let mut nbr_min = i64::MAX;
                for &(dr, dc) in &D8 {
                    let nr = r as i32 + dr;
                    let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
                    if nr >= 0 && nr < h as i32 {
                        let val = filled[nr as usize * w + nc];
                        if val < nbr_min {
                            nbr_min = val;
                        }
                    }
                }
                if filled[r * w + c] < nbr_min {
                    filled[r * w + c] = nbr_min + 1;
                    changed = true;
                }
            }
        }
        if !changed {
            break;
        }
    }
    filled
}

// ---------------------------------------------------------------------------
// Step 3: D8 flow direction
// ---------------------------------------------------------------------------

/// Returns per-cell flow direction index into D8 (0–7), or -1 for no outflow.
fn flow_direction(filled: &[i64], w: usize, h: usize) -> Vec<i8> {
    let mut fdir = vec![-1i8; w * h];
    for r in 0..h {
        for c in 0..w {
            let elev = filled[r * w + c];
            let mut best_drop = 0i64;
            let mut best_k: i8 = -1;
            for (k, &(dr, dc)) in D8.iter().enumerate() {
                let nr = r as i32 + dr;
                let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
                if nr < 0 || nr >= h as i32 {
                    continue;
                }
                let drop = elev - filled[nr as usize * w + nc];
                if drop > best_drop {
                    best_drop = drop;
                    best_k = k as i8;
                }
            }
            fdir[r * w + c] = best_k;
        }
    }
    fdir
}

// ---------------------------------------------------------------------------
// Step 4: Flow accumulation
// ---------------------------------------------------------------------------

fn flow_accumulation(fdir: &[i8], w: usize, h: usize) -> Vec<i32> {
    let n = w * h;
    let mut in_degree = vec![0i32; n];

    for r in 0..h {
        for c in 0..w {
            let k = fdir[r * w + c];
            if k < 0 {
                continue;
            }
            let (dr, dc) = D8[k as usize];
            let nr = r as i32 + dr;
            let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
            if nr >= 0 && nr < h as i32 {
                in_degree[nr as usize * w + nc] += 1;
            }
        }
    }

    let mut queue = VecDeque::new();
    for (i, &deg) in in_degree.iter().enumerate().take(n) {
        if deg == 0 {
            queue.push_back(i);
        }
    }

    let mut accum = vec![1i32; n];
    while let Some(idx) = queue.pop_front() {
        let r = idx / w;
        let c = idx % w;
        let k = fdir[idx];
        if k < 0 {
            continue;
        }
        let (dr, dc) = D8[k as usize];
        let nr = r as i32 + dr;
        let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
        if nr >= 0 && nr < h as i32 {
            let ni = nr as usize * w + nc;
            accum[ni] += accum[idx];
            in_degree[ni] -= 1;
            if in_degree[ni] == 0 {
                queue.push_back(ni);
            }
        }
    }

    accum
}

// ---------------------------------------------------------------------------
// Step 5: River network extraction
// ---------------------------------------------------------------------------

fn extract_river_network(
    accum: &[i32],
    fdir: &[i8],
    w: usize,
    h: usize,
    sea_level: f32,
    elevation: &[f32],
) -> RiverNetwork {
    let n = w * h;

    // River cells: above threshold AND above sea level.
    let is_river: Vec<bool> = (0..n)
        .map(|i| accum[i] > RIVER_THRESHOLD && elevation[i] >= sea_level)
        .collect();

    let river_cells: Vec<(u16, u16)> = (0..n)
        .filter(|&i| is_river[i])
        .map(|i| ((i / w) as u16, (i % w) as u16))
        .collect();

    // Confluences: river cells with 2+ river neighbors flowing into them.
    let mut inflow_count = vec![0u8; n];
    for r in 0..h {
        for c in 0..w {
            let i = r * w + c;
            if !is_river[i] {
                continue;
            }
            let k = fdir[i];
            if k < 0 {
                continue;
            }
            let (dr, dc) = D8[k as usize];
            let nr = r as i32 + dr;
            let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
            if nr >= 0 && nr < h as i32 {
                let ni = nr as usize * w + nc;
                if is_river[ni] {
                    inflow_count[ni] = inflow_count[ni].saturating_add(1);
                }
            }
        }
    }
    let confluences: Vec<(u16, u16)> = (0..n)
        .filter(|&i| is_river[i] && inflow_count[i] >= 2)
        .map(|i| ((i / w) as u16, (i % w) as u16))
        .collect();

    // Mouths: river cells that flow to a sea cell or to the polar edge.
    let mouths: Vec<(u16, u16)> = (0..n)
        .filter(|&i| {
            if !is_river[i] {
                return false;
            }
            let r = i / w;
            let c = i % w;
            let k = fdir[i];
            if k < 0 {
                return true; // no outflow — edge
            }
            let (dr, dc) = D8[k as usize];
            let nr = r as i32 + dr;
            let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
            if nr < 0 || nr >= h as i32 {
                return true; // polar edge
            }
            // Flows into a sub-sea-level cell = mouth
            elevation[nr as usize * w + nc] < sea_level
        })
        .map(|i| ((i / w) as u16, (i % w) as u16))
        .collect();

    RiverNetwork {
        river_cells,
        confluences,
        mouths,
    }
}

// ---------------------------------------------------------------------------
// Step 6: Basin labeling
// ---------------------------------------------------------------------------

fn label_basins(fdir: &[i8], accum: &[i32], w: usize, h: usize) -> Vec<i32> {
    let n = w * h;
    let mut labels = vec![-1i32; n];

    // Pour points: local accumulation maxima above river threshold.
    let mut pour_pts: Vec<usize> = Vec::new();
    for i in 0..n {
        if accum[i] <= RIVER_THRESHOLD {
            continue;
        }
        let r = i / w;
        let c = i % w;
        let mut is_max = true;
        for &(dr, dc) in &D8 {
            let nr = r as i32 + dr;
            let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
            if nr >= 0 && nr < h as i32 && accum[nr as usize * w + nc] > accum[i] {
                is_max = false;
                break;
            }
        }
        if is_max {
            pour_pts.push(i);
        }
    }

    if pour_pts.is_empty() {
        // Flat/ocean world — single basin.
        labels.iter_mut().for_each(|l| *l = 0);
        return labels;
    }

    for (basin_id, &idx) in pour_pts.iter().enumerate() {
        labels[idx] = basin_id as i32;
    }

    // Trace remaining cells: follow fdir until a labeled cell is reached.
    for start in 0..n {
        if labels[start] >= 0 {
            continue;
        }
        // Walk forward, accumulate path.
        let mut path: Vec<usize> = Vec::new();
        let mut cur = start;
        let label = loop {
            if labels[cur] >= 0 {
                break labels[cur];
            }
            path.push(cur);
            let k = fdir[cur];
            if k < 0 {
                break 0; // no outflow — assign to basin 0
            }
            let r = cur / w;
            let c = cur % w;
            let (dr, dc) = D8[k as usize];
            let nr = r as i32 + dr;
            let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
            if nr < 0 || nr >= h as i32 {
                break 0; // polar edge
            }
            let next = nr as usize * w + nc;
            // Cycle guard: if we're visiting a cell already in path, stop.
            if path.contains(&next) {
                break 0;
            }
            cur = next;
        };
        for idx in path {
            labels[idx] = label;
        }
    }

    labels
}

// ---------------------------------------------------------------------------
// Step 7: Merge small basins
// ---------------------------------------------------------------------------

fn merge_small_basins(
    mut labels: Vec<i32>,
    w: usize,
    h: usize,
    min_count: usize,
    max_count: usize,
) -> Vec<i32> {
    let n = w * h;
    let min_frac = 0.02f64; // 2% minimum basin area

    for _ in 0..200 {
        // Count basin sizes.
        let mut sizes: std::collections::BTreeMap<i32, usize> = std::collections::BTreeMap::new();
        for &l in &labels {
            *sizes.entry(l).or_insert(0) += 1;
        }
        let n_basins = sizes.len();

        // Stop if within target range and all basins are large enough.
        if n_basins <= max_count && sizes.values().all(|&s| s as f64 / n as f64 >= min_frac) {
            break;
        }
        if n_basins <= min_count {
            break;
        }

        // Find the smallest basin.
        let (&smallest_id, &smallest_size) = sizes.iter().min_by_key(|(_, &s)| s).unwrap();

        if n_basins <= max_count && smallest_size as f64 / n as f64 >= min_frac {
            break;
        }

        // Find its largest adjacent basin.
        let nbr_id = find_largest_neighbor(&labels, smallest_id, &sizes, w, h);
        let merge_into = nbr_id.unwrap_or(0);

        // Merge.
        for l in labels.iter_mut() {
            if *l == smallest_id {
                *l = merge_into;
            }
        }
    }

    // Renumber contiguously from 0.
    let unique: Vec<i32> = {
        let mut set: std::collections::BTreeSet<i32> = std::collections::BTreeSet::new();
        for &l in &labels {
            set.insert(l);
        }
        set.into_iter().collect()
    };
    let remap: std::collections::BTreeMap<i32, i32> = unique
        .iter()
        .enumerate()
        .map(|(new, &old)| (old, new as i32))
        .collect();
    for l in labels.iter_mut() {
        *l = remap[l];
    }

    labels
}

fn find_largest_neighbor(
    labels: &[i32],
    target_id: i32,
    sizes: &std::collections::BTreeMap<i32, usize>,
    w: usize,
    h: usize,
) -> Option<i32> {
    let n = w * h;
    let mut neighbor_sizes: std::collections::BTreeMap<i32, usize> =
        std::collections::BTreeMap::new();

    for i in 0..n {
        if labels[i] != target_id {
            continue;
        }
        let r = i / w;
        let c = i % w;
        for &(dr, dc) in &D8 {
            let nr = r as i32 + dr;
            let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
            if nr >= 0 && nr < h as i32 {
                let nbr_id = labels[nr as usize * w + nc];
                if nbr_id != target_id {
                    let size = sizes.get(&nbr_id).copied().unwrap_or(0);
                    let e = neighbor_sizes.entry(nbr_id).or_insert(0);
                    if size > *e {
                        *e = size;
                    }
                }
            }
        }
    }

    neighbor_sizes
        .into_iter()
        .max_by_key(|(_, s)| *s)
        .map(|(id, _)| id)
}

// ---------------------------------------------------------------------------
// Step 8: Build DrainageBasin structs
// ---------------------------------------------------------------------------

fn build_basins(labels: &[i32], w: usize, h: usize) -> Vec<DrainageBasin> {
    let n = w * h;
    let mut basin_map: std::collections::BTreeMap<i32, Vec<usize>> =
        std::collections::BTreeMap::new();

    for (i, &l) in labels.iter().enumerate() {
        basin_map.entry(l).or_default().push(i);
    }

    let mut basins: Vec<DrainageBasin> = Vec::with_capacity(basin_map.len());
    let mut ids: Vec<i32> = basin_map.keys().copied().collect();
    ids.sort();

    for basin_id in ids {
        let cells = &basin_map[&basin_id];
        let area_pct = cells.len() as f32 / n as f32;

        // Boundary cells: in this basin, adjacent to a different basin or edge.
        let mut boundary: Vec<(u16, u16)> = Vec::new();
        for &idx in cells {
            let r = idx / w;
            let c = idx % w;
            let mut on_boundary = false;
            for &(dr, dc) in &D8 {
                let nr = r as i32 + dr;
                let nc = (c as i32 + dc).rem_euclid(w as i32) as usize;
                if nr < 0 || nr >= h as i32 {
                    on_boundary = true;
                    break;
                }
                if labels[nr as usize * w + nc] != basin_id {
                    on_boundary = true;
                    break;
                }
            }
            if on_boundary {
                boundary.push((r as u16, c as u16));
            }
        }

        // Sort boundary by angle from centroid for a coherent polygon.
        if !boundary.is_empty() {
            let cr = boundary.iter().map(|&(r, _)| r as f32).sum::<f32>() / boundary.len() as f32;
            let cc = boundary.iter().map(|&(_, c)| c as f32).sum::<f32>() / boundary.len() as f32;
            boundary.sort_by(|&(r1, c1), &(r2, c2)| {
                let a1 = (r1 as f32 - cr).atan2(c1 as f32 - cc);
                let a2 = (r2 as f32 - cr).atan2(c2 as f32 - cc);
                a1.partial_cmp(&a2).unwrap_or(std::cmp::Ordering::Equal)
            });
            // Subsample to ≤500 points.
            if boundary.len() > 500 {
                let step = boundary.len() / 500;
                boundary = boundary.into_iter().step_by(step).collect();
            }
        }

        basins.push(DrainageBasin {
            basin_id: basin_id as u32,
            boundary,
            area_pct,
        });
    }

    basins
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn flat_grid(w: u32, h: u32, val: f32) -> Vec<f32> {
        vec![val; (w * h) as usize]
    }

    fn slope_grid(w: u32, h: u32) -> Vec<f32> {
        let n = (w * h) as usize;
        (0..n)
            .map(|i| {
                let r = i / w as usize;
                let c = i % w as usize;
                // Slope: higher in top-left, drains toward bottom-right.
                1.0 - (r as f32 / h as f32 * 0.5 + c as f32 / w as f32 * 0.5)
            })
            .collect()
    }

    #[test]
    fn flat_grid_produces_single_basin() {
        let elev = flat_grid(16, 8, 0.5);
        let result = analyze(&elev, 16, 8, 0.3);
        // Flat world → no pour points → single basin
        assert_eq!(result.drainage_basins.len(), 1);
        assert!((result.drainage_basins[0].area_pct - 1.0).abs() < 0.01);
    }

    #[test]
    fn slope_grid_has_no_river_cells_below_threshold_by_default() {
        // Small 8×4 grid: max flow_accum ≤ 32, below RIVER_THRESHOLD (200).
        let elev = slope_grid(8, 4);
        let result = analyze(&elev, 8, 4, 0.3);
        // River cells may be empty on this tiny grid — that is acceptable.
        // What matters: no panic and basin count ≥ 1.
        assert!(!result.drainage_basins.is_empty());
    }

    #[test]
    fn large_grid_river_cells_nonempty() {
        // 512×256: max flow accumulation ~131K >> RIVER_THRESHOLD.
        let elev = slope_grid(512, 256);
        let result = analyze(&elev, 512, 256, 0.3);
        assert!(
            !result.river_network.river_cells.is_empty(),
            "Expected river cells on a large sloped grid"
        );
    }

    #[test]
    fn basin_area_pcts_sum_to_one() {
        let elev = slope_grid(64, 32);
        let result = analyze(&elev, 64, 32, 0.3);
        let total: f32 = result.drainage_basins.iter().map(|b| b.area_pct).sum();
        assert!(
            (total - 1.0).abs() < 0.01,
            "Basin area fractions must sum to 1, got {}",
            total
        );
    }

    #[test]
    fn basin_count_within_target_range() {
        let elev = slope_grid(128, 64);
        let result = analyze(&elev, 128, 64, 0.3);
        let n = result.drainage_basins.len();
        assert!(
            n >= 1 && n <= 12,
            "Basin count {} out of expected range [1, 12]",
            n
        );
    }

    #[test]
    fn determinism() {
        // Running analyze twice on the same input must produce identical results.
        let elev = slope_grid(64, 32);
        let r1 = analyze(&elev, 64, 32, 0.3);
        let r2 = analyze(&elev, 64, 32, 0.3);
        assert_eq!(
            r1.river_network.river_cells, r2.river_network.river_cells,
            "River cells must be deterministic"
        );
        assert_eq!(
            r1.drainage_basins.len(),
            r2.drainage_basins.len(),
            "Basin count must be deterministic"
        );
    }
}