taxa-core 0.1.0

//! The focus-bounded treemap — Polars edition. Same contract as the SQL port:
//! deepen one level at a time, keep top-`top_k` children per parent by
//! `size_by`, fold the rest into one "Other" node (sufficient-statistic combine
//! → correct for additive/mean metrics, NULL for non-decomposable ones). Result
//! is `≤ Σ(top_k+1)^level` nodes regardless of dataset size.

use std::collections::HashMap;

use indexmap::IndexMap;
use polars::prelude::*;
use serde_json::{Map, Value as Json};

use crate::error::{Error, Result};
use crate::filters::{axis_row_filter, filter_exprs, focus_exprs};
use crate::manifest::FrameDataset;
use crate::metric::{combines, metric_plans, stat_exprs, MetricPlan};
use crate::output::{fingerprint, TreeNode};
use crate::source::Source;

pub const OTHER: &str = "__other__";
const MAX_TOP_K: i64 = 100;
const MAX_DEPTH: i64 = 12;

pub struct TreemapArgs {
    pub axis: String,
    pub focus: Vec<Json>, // canonical focus keys WITHOUT the leading "root"
    pub filters: Map<String, Json>,
    pub size_by: Option<String>,
    pub depth: i64,
    pub top_k: i64,
}

impl TreemapArgs {
    pub fn new(axis: impl Into<String>) -> Self {
        TreemapArgs {
            axis: axis.into(),
            focus: vec![],
            filters: Map::new(),
            size_by: None,
            depth: 2,
            top_k: 12,
        }
    }
}

struct Built {
    name: String,
    key: Vec<AnyValue<'static>>,
    measures: IndexMap<String, Json>,
    is_other: bool,
    n_folded: i64,
    has_more: bool,
}

pub fn treemap(ds: &FrameDataset, source: &dyn Source, args: &TreemapArgs) -> Result<TreeNode> {
    let columns = source.columns()?;
    let axis = ds
        .axis(&args.axis)
        .ok_or_else(|| Error::Schema(format!("unknown axis {:?}", args.axis)))?;
    let is_path = axis.path.is_some();
    let size_by = ds
        .resolve_size_by(args.size_by.as_deref())
        .ok_or_else(|| Error::Schema("no size_by / metrics".into()))?;

    // A ratio metric isn't an area — sizing a treemap by it is the `share_pct`
    // trap (a ratio doesn't roll up by summing, so rectangles wouldn't nest).
    if let Some(m) = ds.metrics.iter().find(|m| m.id == size_by) {
        if m.is_ratio() {
            return Err(Error::Schema(format!(
                "cannot size a treemap by ratio metric {size_by:?} (ratios aren't additive areas)"
            )));
        }
    }
    let plans: Vec<MetricPlan> = metric_plans(&ds.metrics, &columns)?;
    let size_plan = plans
        .iter()
        .find(|p| p.id == size_by)
        .ok_or_else(|| Error::Schema(format!("size_by {size_by:?} not a metric")))?;
    let stats = stat_exprs(&plans);
    let combos = combines(&plans);

    // Base frame: filtered, then (for a path axis) augmented with the derived
    // `__lvl*` level columns — after which the engine treats it like a fixed
    // axis. `levels` are the authored columns for a fixed axis, the derived
    // names for a path axis.
    let mut base = source.frame()?;
    for e in filter_exprs(ds, &args.filters) {
        base = base.filter(e);
    }
    // Per-axis row universe: applied ONLY for the selected axis, after the base
    // filters and before level derivation/grouping. Other axes (no `row_filter`)
    // keep every row; only this axis drops the excluded rows from its rollup.
    if let Some(pred) = axis_row_filter(axis, &columns)? {
        base = base.filter(pred);
    }
    // Entity-grain metrics: mask each to its grain group's picked row over the
    // now-filtered universe, so the `rollup` (e.g. SUM) gives the per-entity total
    // (no double-count; the picked row respects active filters incl. any time window).
    let masks = crate::metric::entity_mask_exprs(&ds.metrics, &ds.id_column, &columns)?;
    let base = if masks.is_empty() {
        base
    } else {
        base.with_columns(masks)
    };
    let (levels, base) = crate::path::resolved_levels(axis, base)?;
    let levels = &levels;

    let top_k = args.top_k.clamp(1, MAX_TOP_K) as usize;
    let mut depth = args.depth.clamp(1, MAX_DEPTH);
    let f = args.focus.len();
    depth = depth.min(levels.len() as i64 - f as i64);

    // Focus-narrow on the (now resolved) level columns.
    let mut base = base;
    for e in focus_exprs(levels, &args.focus) {
        base = base.filter(e);
    }

    // Totals (root measures): whole-frame aggregation → 1 row.
    let totals = base.clone().select(stats.clone()).collect()?;
    let mut root = TreeNode {
        name: "root".into(),
        measures: finalize_row(&plans, &totals, 0),
        children: vec![],
        is_other: false,
        n_folded: 0,
        has_more: false,
    };
    if depth <= 0 {
        return Ok(root);
    }

    let mut by_parent: HashMap<String, Vec<Built>> = HashMap::new();
    let mut kept_parents: Option<DataFrame> = None;

    for l in (f + 1)..=(f + depth as usize) {
        let level_cols: Vec<String> = levels[..l].to_vec();
        let parent_cols: Vec<String> = levels[..l - 1].to_vec();
        // Kept nodes at the deepest fetched level have unmaterialized children
        // iff there are deeper axis levels — mark them zoomable for windowed fetches.
        // For a path axis this blunt positional flag is replaced by a per-node
        // `_has_more` computed from the next component (see below).
        let boundary_has_more = l == (f + depth as usize) && l < levels.len();

        let mut lf = base.clone();
        // Path-axis correctness fix #1: a shorter path is null-padded at deeper
        // levels. Drop rows whose NEW component at this level is null, so they
        // don't spawn a spurious null node and don't descend — a path that ends
        // here is already counted in its prefix's aggregate.
        if is_path {
            lf = lf.filter(col(level_cols[l - 1].as_str()).is_not_null());
        }
        if let Some(kp) = &kept_parents {
            // descend only into kept parents (semi-join on the parent key cols)
            let keys: Vec<Expr> = parent_cols.iter().map(|c| col(c.as_str())).collect();
            lf = lf.join(
                kp.clone().lazy(),
                keys.clone(),
                keys,
                JoinArgs::new(JoinType::Semi),
            );
        }

        let group: Vec<Expr> = level_cols.iter().map(|c| col(c.as_str())).collect();
        // Path-axis correctness fix #2: per-node `has_more` = does any row under
        // this node carry a non-null NEXT component? Aggregated alongside stats.
        let mut agg_exprs = stats.clone();
        let has_more_col = is_path && l < levels.len();
        if has_more_col {
            agg_exprs.push(
                col(levels[l].as_str())
                    .is_not_null()
                    .max()
                    .alias("_has_more"),
            );
        }
        let agg = lf
            .group_by(group)
            .agg(agg_exprs)
            .with_column(size_plan.rank_expr.clone().alias("_rankval"));

        // Sort by parent asc, rank desc (nulls last); assign within-parent ordinal.
        // Final key: the child's own column ascending — a DETERMINISTIC tiebreaker so
        // equal-`_rankval` siblings always land in the same order (otherwise polars'
        // unstable sort lets equal-sized branches swap positions between requests as
        // you change unrelated settings). (parent, child) is unique, so this fully
        // orders the rows.
        let mut sort_by: Vec<Expr> = parent_cols.iter().map(|c| col(c.as_str())).collect();
        sort_by.push(col("_rankval"));
        sort_by.push(col(level_cols[l - 1].as_str()));
        let mut desc = vec![false; parent_cols.len()];
        desc.push(true); // _rankval desc
        desc.push(false); // child label asc (tiebreaker)
        let sorted = agg
            .sort_by_exprs(
                sort_by,
                SortMultipleOptions::new()
                    .with_order_descending_multi(desc)
                    .with_nulls_last(true),
            )
            .collect()?;

        // _rn within parent (Rust scan over contiguous parent groups).
        let h = sorted.height();
        let mut rn = Vec::with_capacity(h);
        let mut last_fp: Option<String> = None;
        let mut counter: u32 = 0;
        for i in 0..h {
            let pfp = fingerprint(&row_key(&sorted, &parent_cols, i)?);
            if last_fp.as_ref() != Some(&pfp) {
                counter = 0;
                last_fp = Some(pfp);
            }
            counter += 1;
            rn.push(counter);
        }
        let mut sorted = sorted;
        let rn_i64: Vec<i64> = rn.iter().map(|v| *v as i64).collect();
        sorted.with_column(Series::new("_rn".into(), rn_i64))?;

        // kept rows → nodes; tail rows → one "Other" per parent.
        let kept_mask: BooleanChunked = sorted.column("_rn")?.i64()?.lt_eq(top_k as i64);
        let tail_mask = !&kept_mask;
        let kept = sorted.filter(&kept_mask)?;
        let tail = sorted.filter(&tail_mask)?;

        for i in 0..kept.height() {
            let key = row_key(&kept, &level_cols, i)?;
            let parent_fp = fingerprint(&key[..key.len() - 1]);
            let name = av_to_label_static(&key[key.len() - 1]);
            // Path axis: per-node flag from the next-component aggregate (a file
            // leaf at the boundary has no deeper component → not zoomable). At the
            // last derived level there is no next column → never has_more.
            let has_more = if is_path {
                has_more_col
                    && kept
                        .column("_has_more")
                        .ok()
                        .and_then(|c| c.get(i).ok())
                        .map(|av| matches!(av, AnyValue::Boolean(true)))
                        .unwrap_or(false)
            } else {
                boundary_has_more
            };
            by_parent.entry(parent_fp).or_default().push(Built {
                name,
                key: key.clone(),
                measures: finalize_row(&plans, &kept, i),
                is_other: false,
                n_folded: 0,
                has_more,
            });
        }

        if tail.height() > 0 {
            let other_exprs: Vec<Expr> = combos
                .iter()
                .map(|(a, c)| MetricPlan::other_expr(a, *c))
                .collect();
            let pgroup: Vec<Expr> = parent_cols.iter().map(|c| col(c.as_str())).collect();
            // `__nf` = number of sibling branches folded into each Other node, so the
            // frontend can show "+N others" instead of a hardcoded "+1".
            let other = if parent_cols.is_empty() {
                tail.clone().lazy().select(other_exprs).collect()?
            } else {
                let mut agg = other_exprs;
                agg.push(len().alias("__nf"));
                tail.clone().lazy().group_by(pgroup).agg(agg).collect()?
            };
            for i in 0..other.height() {
                let parent_key = row_key(&other, &parent_cols, i)?;
                let parent_fp = fingerprint(&parent_key);
                let mut key = parent_key.clone();
                key.push(AnyValue::Null);
                let n_folded = if parent_cols.is_empty() {
                    tail.height() as i64
                } else {
                    other
                        .column("__nf")
                        .ok()
                        .and_then(|c| c.get(i).ok())
                        .map(|av| crate::output::av_to_json(&av))
                        .and_then(|j| j.as_i64())
                        .unwrap_or(1)
                };
                by_parent.entry(parent_fp).or_default().push(Built {
                    name: OTHER.to_string(),
                    key,
                    measures: finalize_row(&plans, &other, i),
                    is_other: true,
                    n_folded,
                    has_more: false,
                });
            }
        }

        // Next level descends into kept parents only.
        kept_parents =
            Some(kept.select(level_cols.iter().map(|s| s.as_str()).collect::<Vec<_>>())?);
    }

    attach(
        &mut root,
        &fingerprint(&focus_key(&args.focus)),
        &mut by_parent,
    );
    Ok(root)
}

/// The treemap's kept branch set ONE LEVEL below the focus, for the same
/// `axis`/`focus`/`size_by`/`top_k` — the snapshot ranking the series Line tab
/// follows when a view sets `branch_set: "treemap"`. Returns `(keep_keys,
/// has_other)`: the kept (non-Other) child names in treemap order, plus whether
/// an "Other" node exists at that level. Runs the standard rollup at depth 1.
pub struct BranchSet {
    pub keep: Vec<String>,
    pub has_other: bool,
}

pub fn branch_set(ds: &FrameDataset, source: &dyn Source, args: &TreemapArgs) -> Result<BranchSet> {
    // Resolve exactly one level below the focus: the immediate children are the
    // branch set. Reuse the requested axis/focus/size_by/top_k.
    let mut a = TreemapArgs::new(args.axis.clone());
    a.focus = args.focus.clone();
    a.filters = args.filters.clone();
    a.size_by = args.size_by.clone();
    a.top_k = args.top_k;
    a.depth = 1;
    let tree = treemap(ds, source, &a)?;
    let mut keep = Vec::new();
    let mut has_other = false;
    for c in &tree.children {
        if c.is_other {
            has_other = true;
        } else {
            keep.push(c.name.clone());
        }
    }
    Ok(BranchSet { keep, has_other })
}

fn attach(node: &mut TreeNode, key_fp: &str, by_parent: &mut HashMap<String, Vec<Built>>) {
    let built = by_parent.remove(key_fp).unwrap_or_default();
    for b in built {
        let child_fp = fingerprint(&b.key);
        let mut child = TreeNode {
            name: b.name,
            measures: b.measures,
            children: vec![],
            is_other: b.is_other,
            n_folded: b.n_folded,
            has_more: b.has_more,
        };
        if !b.is_other {
            attach(&mut child, &child_fp, by_parent);
        }
        node.children.push(child);
    }
}

/// Measures for one row of a stat DataFrame (every metric finalized).
fn finalize_row(plans: &[MetricPlan], df: &DataFrame, i: usize) -> IndexMap<String, Json> {
    let get = |alias: &str| -> AnyValue<'static> {
        df.column(alias)
            .ok()
            .and_then(|c| c.get(i).ok())
            .map(|av| av.into_static())
            .unwrap_or(AnyValue::Null)
    };
    plans
        .iter()
        .map(|p| (p.id.clone(), p.finalize_row(&get)))
        .collect()
}

/// The AnyValue tuple of the given columns at row `i` (owned).
fn row_key(df: &DataFrame, cols: &[String], i: usize) -> Result<Vec<AnyValue<'static>>> {
    cols.iter()
        .map(|c| Ok(df.column(c.as_str())?.get(i)?.into_static()))
        .collect()
}

fn focus_key(focus: &[Json]) -> Vec<AnyValue<'static>> {
    focus.iter().map(json_to_av).collect()
}

fn json_to_av(v: &Json) -> AnyValue<'static> {
    match v {
        Json::Null => AnyValue::Null,
        Json::Bool(b) => AnyValue::Boolean(*b),
        Json::Number(n) => {
            if let Some(i) = n.as_i64() {
                AnyValue::Int64(i)
            } else {
                AnyValue::Float64(n.as_f64().unwrap_or(0.0))
            }
        }
        Json::String(s) => AnyValue::StringOwned(s.as_str().into()),
        other => AnyValue::StringOwned(other.to_string().into()),
    }
}

fn av_to_label_static(av: &AnyValue<'static>) -> String {
    crate::output::av_to_label(av)
}