#include "viz.h"
#include "implot.h"

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <string>
#include <unordered_map>
#include <vector>

namespace viz {

using data_table::StageOutput;
using data_table::ColumnType;
using data_table::ViewMode;
using data_table::ViewConfig;
using data_table::parse_number;
using data_table::nearest_index_2d;
using data_table::pie_angle;
using data_table::pie_slice_at_angle;
using data_table::heatmap_cell_at;

static int find_header(const StageOutput& out, const std::string& name) {
    if (name.empty()) return -1;
    for (size_t c = 0; c < out.headers.size(); ++c)
        if (out.headers[c] == name) return (int)c;
    return -1;
}

static int resolve_x(const StageOutput& out, const ViewConfig& cfg, int fallback) {
    int c = find_header(out, cfg.x_col);
    return (c >= 0) ? c : fallback;
}
static int resolve_cat(const StageOutput& out, const ViewConfig& cfg, int fallback) {
    int c = find_header(out, cfg.cat_col);
    return (c >= 0) ? c : fallback;
}
static int resolve_size(const StageOutput& out, const ViewConfig& cfg, int fallback) {
    int c = find_header(out, cfg.size_col);
    return (c >= 0) ? c : fallback;
}

int first_numeric_col(const StageOutput& out) {
    for (size_t c = 0; c < out.types.size(); ++c) {
        if (out.types[c] == ColumnType::Int || out.types[c] == ColumnType::Float) return (int)c;
    }
    return -1;
}

int first_category_col(const StageOutput& out) {
    for (size_t c = 0; c < out.types.size(); ++c) {
        ColumnType t = out.types[c];
        if (t == ColumnType::String || t == ColumnType::Date || t == ColumnType::Bool ||
            t == ColumnType::Json) return (int)c;
    }
    return -1;
}

std::vector<double> extract_numeric(const StageOutput& out, int col) {
    std::vector<double> v;
    if (col < 0 || col >= out.cols) return v;
    v.reserve(out.rows);
    for (int r = 0; r < out.rows; ++r) {
        const char* s = out.cells[(size_t)r * out.cols + col];
        double d = 0;
        if (s && *s && parse_number(s, d)) v.push_back(d);
        else                                v.push_back(std::nan(""));
    }
    return v;
}

std::vector<std::string> extract_category(const StageOutput& out, int col) {
    std::vector<std::string> v;
    if (col < 0 || col >= out.cols) return v;
    v.reserve(out.rows);
    for (int r = 0; r < out.rows; ++r) {
        const char* s = out.cells[(size_t)r * out.cols + col];
        v.emplace_back(s ? s : "");
    }
    return v;
}

namespace {

struct NumCol { int idx; std::string name; std::vector<double> vals; };

std::vector<NumCol> collect_numeric(const StageOutput& out, int max_n = 16) {
    std::vector<NumCol> r;
    for (size_t c = 0; c < out.types.size() && (int)r.size() < max_n; ++c) {
        if (out.types[c] == ColumnType::Int || out.types[c] == ColumnType::Float) {
            NumCol nc;
            nc.idx  = (int)c;
            nc.name = out.headers[c];
            nc.vals = extract_numeric(out, (int)c);
            r.push_back(std::move(nc));
        }
    }
    return r;
}

std::vector<NumCol> collect_numeric_filtered(const StageOutput& out,
                                              const ViewConfig& cfg,
                                              int max_n = 16) {
    if (cfg.y_cols.empty()) return collect_numeric(out, max_n);
    std::vector<NumCol> r;
    for (const auto& name : cfg.y_cols) {
        if ((int)r.size() >= max_n) break;
        int c = find_header(out, name);
        if (c < 0) continue;
        if (out.types[c] != ColumnType::Int && out.types[c] != ColumnType::Float) continue;
        NumCol nc;
        nc.idx  = c;
        nc.name = out.headers[c];
        nc.vals = extract_numeric(out, c);
        r.push_back(std::move(nc));
    }
    if (r.empty()) return collect_numeric(out, max_n);
    return r;
}

ImPlotSpec spec_with_color(unsigned int rgba_color) {
    if (rgba_color == 0) return ImPlotSpec();
    ImU32 c = (ImU32)rgba_color;
    return ImPlotSpec(ImPlotProp_LineColor, c, ImPlotProp_FillColor, c);
}

// Axis flags: locked = no pan/zoom; unlocked = 0 (sin AutoFit, para preservar
// pan/zoom del user). Re-fit explicito via SetNextAxesToFit cuando fit_request.
ImPlotAxisFlags axflag(const ViewConfig& cfg, ImPlotAxisFlags base = 0) {
    if (cfg.locked) return base | ImPlotAxisFlags_Lock;
    return base;
}

// Llamar antes de BeginPlot. Si cfg.fit_request -> fuerza re-fit y limpia el flag.
void maybe_fit(const ViewConfig& cfg) {
    if (cfg.fit_request) {
        ImPlot::SetNextAxesToFit();
        cfg.fit_request = false;
    }
}

void info_text(const char* msg) {
    ImVec2 avail = ImGui::GetContentRegionAvail();
    ImVec2 sz    = ImGui::CalcTextSize(msg);
    ImGui::SetCursorPos(ImVec2(ImGui::GetCursorPosX() + (avail.x - sz.x) * 0.5f,
                                ImGui::GetCursorPosY() + (avail.y - sz.y) * 0.5f));
    ImGui::TextDisabled("%s", msg);
}

// Drop NaN and pair with optional labels.
std::vector<double> finite(const std::vector<double>& v) {
    std::vector<double> r; r.reserve(v.size());
    for (double d : v) if (!std::isnan(d)) r.push_back(d);
    return r;
}

bool render_bar_like(const StageOutput& out, ViewMode mode,
                     const ViewConfig& cfg, ImVec2 size,
                     int* clicked_row_out = nullptr) {
    int cat_col = resolve_cat(out, cfg, first_category_col(out));
    auto nums = collect_numeric_filtered(out, cfg, 8);
    if (cat_col < 0 || nums.empty()) {
        info_text("Need 1 category + 1+ numeric columns");
        return false;
    }
    auto cats = extract_category(out, cat_col);
    int n = (int)cats.size();
    if (n == 0) { info_text("Empty data"); return false; }

    // Ticks
    std::vector<double> ticks(n);
    std::vector<const char*> labels(n);
    for (int i = 0; i < n; ++i) { ticks[i] = i; labels[i] = cats[i].c_str(); }

    bool horiz = (mode == ViewMode::Bar);
    ImPlotFlags pflags = cfg.show_legend ? 0 : ImPlotFlags_NoLegend;
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##bar", size, pflags)) return false;

    ImPlotAxisFlags ax_cat = axflag(cfg);
    ImPlotAxisFlags ax_num = axflag(cfg);

    if (horiz) {
        ImPlot::SetupAxes(out.headers[nums[0].idx].c_str(), out.headers[cat_col].c_str(),
                          ax_num, ax_cat);
        ImPlot::SetupAxisTicks(ImAxis_Y1, ticks.data(), n, labels.data(), false);
    } else {
        ImPlot::SetupAxes(out.headers[cat_col].c_str(), out.headers[nums[0].idx].c_str(),
                          ax_cat, ax_num);
        ImPlot::SetupAxisTicks(ImAxis_X1, ticks.data(), n, labels.data(), false);
    }

    if (mode == ViewMode::StackedBar || mode == ViewMode::GroupedBar) {
        // Build flat matrix items x groups
        int items  = (int)nums.size();
        std::vector<double> mat((size_t)items * n, 0.0);
        std::vector<const char*> series_labels(items);
        for (int it = 0; it < items; ++it) {
            series_labels[it] = nums[it].name.c_str();
            for (int g = 0; g < n; ++g) {
                double d = nums[it].vals[g];
                mat[(size_t)it * n + g] = std::isnan(d) ? 0.0 : d;
            }
        }
        int flags = (mode == ViewMode::StackedBar) ? ImPlotBarGroupsFlags_Stacked : 0;
        if (horiz) flags |= ImPlotBarGroupsFlags_Horizontal;
        ImPlot::PlotBarGroups(series_labels.data(), mat.data(), items, n, 0.67, 0,
                              ImPlotSpec(ImPlotProp_Flags, flags));
    } else {
        // Single series (first numeric col).
        std::vector<double> ys(n);
        for (int i = 0; i < n; ++i) {
            double d = nums[0].vals[i];
            ys[i] = std::isnan(d) ? 0.0 : d;
        }
        ImPlotSpec spc = spec_with_color(cfg.primary_color);
        if (horiz) {
            if (cfg.primary_color != 0) {
                ImU32 col = (ImU32)cfg.primary_color;
                ImPlot::PlotBars(nums[0].name.c_str(), ys.data(), ticks.data(), n, 0.67,
                                 ImPlotSpec(ImPlotProp_Flags, ImPlotBarsFlags_Horizontal,
                                            ImPlotProp_FillColor, col,
                                            ImPlotProp_LineColor, col));
            } else {
                ImPlot::PlotBars(nums[0].name.c_str(), ys.data(), ticks.data(), n, 0.67,
                                 ImPlotSpec(ImPlotProp_Flags, ImPlotBarsFlags_Horizontal));
            }
        } else {
            ImPlot::PlotBars(nums[0].name.c_str(), ticks.data(), ys.data(), n, 0.67, spc);
        }
    }
    // Hit-test fase 10: idx = round(plot.{x|y}) en single-series mode.
    if (clicked_row_out &&
        mode != ViewMode::GroupedBar && mode != ViewMode::StackedBar &&
        ImPlot::IsPlotHovered() && ImGui::IsMouseClicked(ImGuiMouseButton_Left)) {
        ImPlotPoint p = ImPlot::GetPlotMousePos();
        double target = horiz ? p.y : p.x;
        int idx = (int)(target + 0.5);
        if (idx >= 0 && idx < n) *clicked_row_out = idx;
    }
    ImPlot::EndPlot();
    return true;
}

bool render_line_like(const StageOutput& out, ViewMode mode,
                      const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 8);
    if (nums.empty()) { info_text("Need at least 1 numeric column"); return false; }

    ImPlotFlags pflags = cfg.show_legend ? 0 : ImPlotFlags_NoLegend;
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##line", size, pflags)) return false;
    ImPlot::SetupAxes(nullptr, nullptr, axflag(cfg), axflag(cfg));

    int n = nums.empty() ? 0 : (int)nums[0].vals.size();
    if (n == 0) { ImPlot::EndPlot(); return false; }

    // X column: cfg.x_col override; sino primer numeric si hay >=2; sino indices.
    int x_idx = -1;
    if (!cfg.x_col.empty()) {
        int xc = find_header(out, cfg.x_col);
        if (xc >= 0 && (out.types[xc] == ColumnType::Int || out.types[xc] == ColumnType::Float)) {
            x_idx = xc;
        }
    }
    std::vector<double> idx_xs;
    const double* xs = nullptr;
    int start_y = 0;
    std::vector<double> x_data_external;
    if (x_idx >= 0) {
        x_data_external = extract_numeric(out, x_idx);
        xs = x_data_external.data();
    } else if (nums.size() >= 2 && cfg.y_cols.empty()) {
        xs = nums[0].vals.data();
        start_y = 1;
    } else {
        idx_xs.resize(n);
        for (int i = 0; i < n; ++i) idx_xs[i] = i;
        xs = idx_xs.data();
    }

    bool only_one = (cfg.primary_color != 0) && (nums.size() - start_y == 1);
    for (size_t i = (size_t)start_y; i < nums.size(); ++i) {
        const auto& nc = nums[i];
        ImU32 col = only_one ? (ImU32)cfg.primary_color : 0;
        int marker = cfg.show_markers ? ImPlotMarker_Circle : ImPlotMarker_None;
        if (mode == ViewMode::Area) {
            if (col) {
                ImPlot::PlotShaded(nc.name.c_str(), xs, nc.vals.data(), (int)nc.vals.size(), 0.0,
                                    ImPlotSpec(ImPlotProp_FillColor, col, ImPlotProp_LineColor, col));
            } else {
                ImPlot::PlotShaded(nc.name.c_str(), xs, nc.vals.data(), (int)nc.vals.size(), 0.0);
            }
        } else if (mode == ViewMode::Stairs) {
            if (col) {
                ImPlot::PlotStairs(nc.name.c_str(), xs, nc.vals.data(), (int)nc.vals.size(),
                                    ImPlotSpec(ImPlotProp_LineColor, col));
            } else {
                ImPlot::PlotStairs(nc.name.c_str(), xs, nc.vals.data(), (int)nc.vals.size());
            }
        } else {
            if (col) {
                ImPlot::PlotLine(nc.name.c_str(), xs, nc.vals.data(), (int)nc.vals.size(),
                                  ImPlotSpec(ImPlotProp_Flags, ImPlotLineFlags_SkipNaN,
                                             ImPlotProp_LineColor, col,
                                             ImPlotProp_Marker, marker));
            } else {
                ImPlot::PlotLine(nc.name.c_str(), xs, nc.vals.data(), (int)nc.vals.size(),
                                  ImPlotSpec(ImPlotProp_Flags, ImPlotLineFlags_SkipNaN,
                                             ImPlotProp_Marker, marker));
            }
        }
    }
    ImPlot::EndPlot();
    return true;
}

bool render_scatter(const StageOutput& out, const ViewConfig& cfg, ImVec2 size,
                     int* clicked_row_out = nullptr) {
    // Soporte cfg.x_col + cfg.y_cols[0]
    int xc = find_header(out, cfg.x_col);
    int yc = !cfg.y_cols.empty() ? find_header(out, cfg.y_cols[0]) : -1;
    std::vector<NumCol> nums;
    if (xc >= 0 && yc >= 0) {
        NumCol a{xc, out.headers[xc], extract_numeric(out, xc)};
        NumCol b{yc, out.headers[yc], extract_numeric(out, yc)};
        nums = {a, b};
    } else {
        nums = collect_numeric(out, 4);
    }
    if (nums.size() < 2) { info_text("Need 2 numeric columns"); return false; }
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##scatter", size, cfg.show_legend ? 0 : ImPlotFlags_NoLegend)) return false;
    ImPlot::SetupAxes(nums[0].name.c_str(), nums[1].name.c_str(),
                      axflag(cfg), axflag(cfg));
    if (cfg.primary_color) {
        ImU32 col = (ImU32)cfg.primary_color;
        ImPlot::PlotScatter("##s", nums[0].vals.data(), nums[1].vals.data(),
                             (int)nums[0].vals.size(),
                             ImPlotSpec(ImPlotProp_MarkerFillColor, col,
                                        ImPlotProp_MarkerLineColor, col));
    } else {
        ImPlot::PlotScatter("##s", nums[0].vals.data(), nums[1].vals.data(),
                             (int)nums[0].vals.size());
    }
    if (clicked_row_out &&
        ImPlot::IsPlotHovered() && ImGui::IsMouseClicked(ImGuiMouseButton_Left)) {
        ImPlotPoint p = ImPlot::GetPlotMousePos();
        int idx = nearest_index_2d(p.x, p.y,
                                    nums[0].vals.data(), nums[1].vals.data(),
                                    (int)nums[0].vals.size());
        if (idx >= 0) *clicked_row_out = idx;
    }
    ImPlot::EndPlot();
    return true;
}

bool render_bubble(const StageOutput& out, const ViewConfig& cfg, ImVec2 size,
                    int* clicked_row_out = nullptr) {
    int xc = find_header(out, cfg.x_col);
    int yc = !cfg.y_cols.empty() ? find_header(out, cfg.y_cols[0]) : -1;
    int sc = resolve_size(out, cfg, -1);
    std::vector<NumCol> nums;
    if (xc >= 0 && yc >= 0 && sc >= 0) {
        nums = {
            {xc, out.headers[xc], extract_numeric(out, xc)},
            {yc, out.headers[yc], extract_numeric(out, yc)},
            {sc, out.headers[sc], extract_numeric(out, sc)},
        };
    } else {
        nums = collect_numeric(out, 4);
    }
    if (nums.size() < 3) { info_text("Need 3 numeric columns (x, y, size)"); return false; }
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##bubble", size, cfg.show_legend ? 0 : ImPlotFlags_NoLegend)) return false;
    ImPlot::SetupAxes(nums[0].name.c_str(), nums[1].name.c_str(),
                      axflag(cfg), axflag(cfg));
    ImPlot::PlotBubbles("##b", nums[0].vals.data(), nums[1].vals.data(),
                         nums[2].vals.data(), (int)nums[0].vals.size());
    if (clicked_row_out &&
        ImPlot::IsPlotHovered() && ImGui::IsMouseClicked(ImGuiMouseButton_Left)) {
        ImPlotPoint p = ImPlot::GetPlotMousePos();
        int idx = nearest_index_2d(p.x, p.y,
                                    nums[0].vals.data(), nums[1].vals.data(),
                                    (int)nums[0].vals.size());
        if (idx >= 0) *clicked_row_out = idx;
    }
    ImPlot::EndPlot();
    return true;
}

bool render_histogram(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 4);
    if (nums.empty()) { info_text("Need 1 numeric column"); return false; }
    auto vals = finite(nums[0].vals);
    if (vals.empty()) { info_text("No finite values"); return false; }
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##hist", size, cfg.show_legend ? 0 : ImPlotFlags_NoLegend)) return false;
    ImPlot::SetupAxes(nums[0].name.c_str(), "count",
                      axflag(cfg), axflag(cfg));
    int bins = (cfg.hist_bins > 0) ? cfg.hist_bins : ImPlotBin_Sturges;
    if (cfg.primary_color) {
        ImU32 col = (ImU32)cfg.primary_color;
        ImPlot::PlotHistogram("##h", vals.data(), (int)vals.size(), bins, 1.0,
                               ImPlotRange(),
                               ImPlotSpec(ImPlotProp_FillColor, col,
                                          ImPlotProp_LineColor, col));
    } else {
        ImPlot::PlotHistogram("##h", vals.data(), (int)vals.size(), bins);
    }
    ImPlot::EndPlot();
    return true;
}

bool render_hist2d(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    int xc = find_header(out, cfg.x_col);
    int yc = !cfg.y_cols.empty() ? find_header(out, cfg.y_cols[0]) : -1;
    std::vector<NumCol> nums;
    if (xc >= 0 && yc >= 0) {
        nums = {
            {xc, out.headers[xc], extract_numeric(out, xc)},
            {yc, out.headers[yc], extract_numeric(out, yc)},
        };
    } else {
        nums = collect_numeric(out, 2);
    }
    if (nums.size() < 2) { info_text("Need 2 numeric columns"); return false; }
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##hist2d", size, cfg.show_legend ? 0 : ImPlotFlags_NoLegend)) return false;
    ImPlot::SetupAxes(nums[0].name.c_str(), nums[1].name.c_str());
    int bins = (cfg.hist_bins > 0) ? cfg.hist_bins : ImPlotBin_Sturges;
    ImPlot::PlotHistogram2D("##h2", nums[0].vals.data(), nums[1].vals.data(),
                             (int)nums[0].vals.size(), bins, bins);
    ImPlot::EndPlot();
    return true;
}

bool render_heatmap(const StageOutput& out, const ViewConfig& cfg, ImVec2 size,
                     int* clicked_row_out = nullptr) {
    auto nums = collect_numeric_filtered(out, cfg, 64);
    if (nums.empty()) { info_text("Need numeric columns"); return false; }
    int cols = (int)nums.size();
    int rows = (int)nums[0].vals.size();
    if (rows == 0) { info_text("No rows"); return false; }
    std::vector<double> mat((size_t)rows * cols, 0.0);
    double mn = +1e300, mx = -1e300;
    for (int c = 0; c < cols; ++c) {
        for (int r = 0; r < rows; ++r) {
            double d = nums[c].vals[r];
            if (std::isnan(d)) d = 0;
            mat[(size_t)r * cols + c] = d;
            if (d < mn) mn = d; if (d > mx) mx = d;
        }
    }
    if (mn == mx) { mx = mn + 1; }
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##heatmap", size, 0)) return false;
    ImPlot::PlotHeatmap("##hm", mat.data(), rows, cols, mn, mx, nullptr);
    if (clicked_row_out &&
        ImPlot::IsPlotHovered() && ImGui::IsMouseClicked(ImGuiMouseButton_Left)) {
        ImPlotPoint p = ImPlot::GetPlotMousePos();
        // ImPlot heatmap Y se pinta de top a bottom; plot mouse_y va igual
        // (default scale 0..rows). Mapeo directo.
        int rr, cc;
        heatmap_cell_at(p.x, p.y, rows, cols, rr, cc);
        if (rr >= 0) *clicked_row_out = rr;
        (void)cc;
    }
    ImPlot::EndPlot();
    return true;
}

bool render_pie(const StageOutput& out, const ViewConfig& cfg, bool donut, ImVec2 size,
                 int* clicked_row_out = nullptr) {
    int cat = resolve_cat(out, cfg, first_category_col(out));
    auto nums = collect_numeric_filtered(out, cfg, 1);
    if (cat < 0 || nums.empty()) { info_text("Need 1 category + 1 numeric"); return false; }
    auto cats = extract_category(out, cat);
    int n = std::min((int)cats.size(), (int)nums[0].vals.size());
    if (n == 0) return false;
    std::vector<double> values(n);
    std::vector<const char*> labels(n);
    for (int i = 0; i < n; ++i) {
        double d = nums[0].vals[i];
        values[i] = std::isnan(d) ? 0.0 : std::abs(d);
        labels[i] = cats[i].c_str();
    }
    ImPlotFlags pf = ImPlotFlags_Equal;
    if (!cfg.show_legend) pf |= ImPlotFlags_NoLegend;
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##pie", size, pf)) return false;
    ImPlot::SetupAxes(nullptr, nullptr, ImPlotAxisFlags_NoDecorations,
                      ImPlotAxisFlags_NoDecorations);
    ImPlot::SetupAxesLimits(0, 1, 0, 1, ImPlotCond_Always);
    double radius = (cfg.pie_radius > 0) ? cfg.pie_radius : (donut ? 0.4 : 0.45);
    ImPlot::PlotPieChart(labels.data(), values.data(), n, 0.5, 0.5, radius, "%.1f");
    if (donut) {
        // Draw inner hole as solid circle by overlaying a smaller pie of one slice transparent.
        // Simpler: just visually it's a circle with text. Use no extra primitive for now.
    }
    if (clicked_row_out &&
        ImPlot::IsPlotHovered() && ImGui::IsMouseClicked(ImGuiMouseButton_Left)) {
        ImPlotPoint p = ImPlot::GetPlotMousePos();
        double dx = p.x - 0.5, dy = p.y - 0.5;
        double dist2 = dx*dx + dy*dy;
        double inner = donut ? (radius * 0.5) : 0.0;
        if (dist2 <= radius * radius && dist2 >= inner * inner) {
            double ang = pie_angle(0.5, 0.5, p.x, p.y);
            int idx = pie_slice_at_angle(ang, values.data(), n);
            if (idx >= 0) *clicked_row_out = idx;
        }
    }
    ImPlot::EndPlot();
    return true;
}

bool render_funnel(const StageOutput& out, const ViewConfig& cfg, ImVec2 size,
                    int* clicked_row_out = nullptr) {
    int cat = resolve_cat(out, cfg, first_category_col(out));
    auto nums = collect_numeric_filtered(out, cfg, 1);
    if (cat < 0 || nums.empty()) { info_text("Need 1 category + 1 numeric"); return false; }
    auto cats = extract_category(out, cat);
    int n = std::min((int)cats.size(), (int)nums[0].vals.size());
    if (n == 0) return false;
    // Sort desc by value
    std::vector<int> idx(n);
    for (int i = 0; i < n; ++i) idx[i] = i;
    std::sort(idx.begin(), idx.end(), [&](int a, int b) {
        double da = std::isnan(nums[0].vals[a]) ? -1e300 : nums[0].vals[a];
        double db = std::isnan(nums[0].vals[b]) ? -1e300 : nums[0].vals[b];
        return da > db;
    });
    std::vector<double> ys(n);
    std::vector<double> ticks(n);
    std::vector<const char*> labels(n);
    std::vector<std::string> labels_store(n);
    for (int i = 0; i < n; ++i) {
        double d = nums[0].vals[idx[i]];
        ys[i]      = std::isnan(d) ? 0 : d;
        ticks[i]   = n - 1 - i;       // descending order
        labels_store[i] = cats[idx[i]];
        labels[i]  = labels_store[i].c_str();
    }
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##funnel", size, 0)) return false;
    ImPlot::SetupAxes(nums[0].name.c_str(), out.headers[cat].c_str(),
                      axflag(cfg), axflag(cfg));
    ImPlot::SetupAxisTicks(ImAxis_Y1, ticks.data(), n, labels.data(), false);
    ImPlot::PlotBars(nums[0].name.c_str(), ys.data(), ticks.data(), n, 0.85,
                     ImPlotSpec(ImPlotProp_Flags, ImPlotBarsFlags_Horizontal));
    if (clicked_row_out &&
        ImPlot::IsPlotHovered() && ImGui::IsMouseClicked(ImGuiMouseButton_Left)) {
        ImPlotPoint p = ImPlot::GetPlotMousePos();
        int tick_idx = (int)(p.y + 0.5);
        // ticks[i] = n-1-i. Invertir para idx en orden sorted descendiente.
        int sorted_pos = (n - 1) - tick_idx;
        if (sorted_pos >= 0 && sorted_pos < n) {
            // idx[sorted_pos] da indice de row original en out.
            *clicked_row_out = idx[sorted_pos];
        }
    }
    ImPlot::EndPlot();
    return true;
}

bool render_waterfall(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 1);
    if (nums.empty()) { info_text("Need 1 numeric column"); return false; }
    int n = (int)nums[0].vals.size();
    if (n == 0) return false;
    int cat = resolve_cat(out, cfg, first_category_col(out));
    auto cats = (cat >= 0) ? extract_category(out, cat) : std::vector<std::string>();

    std::vector<double> running(n + 1, 0);
    for (int i = 0; i < n; ++i) {
        double d = std::isnan(nums[0].vals[i]) ? 0 : nums[0].vals[i];
        running[i + 1] = running[i] + d;
    }
    std::vector<double> ticks(n);
    for (int i = 0; i < n; ++i) ticks[i] = i;
    std::vector<const char*> labels(n);
    for (int i = 0; i < n; ++i) labels[i] = (i < (int)cats.size()) ? cats[i].c_str() : "";

    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##waterfall", size, 0)) return false;
    ImPlot::SetupAxes(nullptr, nums[0].name.c_str(),
                      axflag(cfg), axflag(cfg));
    if (cat >= 0) ImPlot::SetupAxisTicks(ImAxis_X1, ticks.data(), n, labels.data(), false);
    // Draw stems with rectangles via error-bars trick: low=cum_prev, high=cum_curr.
    std::vector<double> mid(n), err(n);
    for (int i = 0; i < n; ++i) {
        mid[i] = (running[i] + running[i + 1]) * 0.5;
        err[i] = std::abs((running[i + 1] - running[i]) * 0.5);
    }
    ImPlot::PlotErrorBars("##wf", ticks.data(), mid.data(), err.data(), n);
    ImPlot::PlotLine("cum", running.data() + 1, n);
    ImPlot::EndPlot();
    return true;
}

bool render_kpi_single(const StageOutput& out, const ViewConfig& cfg) {
    int nc = !cfg.y_cols.empty() ? find_header(out, cfg.y_cols[0]) : -1;
    if (nc < 0) nc = first_numeric_col(out);
    if (nc < 0) { info_text("Need 1 numeric column"); return false; }
    auto vals = extract_numeric(out, nc);
    if (vals.empty()) { info_text("Empty"); return false; }
    double last = std::nan("");
    for (auto v : vals) if (!std::isnan(v)) last = v;
    if (std::isnan(last)) { info_text("No finite values"); return false; }

    char buf[64];
    if (std::abs(last) >= 1e6)       std::snprintf(buf, sizeof(buf), "%.2fM", last / 1e6);
    else if (std::abs(last) >= 1e3)  std::snprintf(buf, sizeof(buf), "%.2fK", last / 1e3);
    else                              std::snprintf(buf, sizeof(buf), "%.3g", last);

    ImVec2 avail = ImGui::GetContentRegionAvail();
    ImGui::SetWindowFontScale(4.0f);
    ImVec2 sz = ImGui::CalcTextSize(buf);
    ImGui::SetCursorPos(ImVec2(ImGui::GetCursorPosX() + (avail.x - sz.x) * 0.5f,
                                ImGui::GetCursorPosY() + (avail.y - sz.y) * 0.5f - 20));
    ImGui::TextUnformatted(buf);
    ImGui::SetWindowFontScale(1.0f);

    sz = ImGui::CalcTextSize(out.headers[nc].c_str());
    ImGui::SetCursorPos(ImVec2(ImGui::GetCursorPosX() + (avail.x - sz.x) * 0.5f,
                                ImGui::GetCursorPosY() + (avail.y - sz.y) * 0.5f - 10));
    ImGui::TextDisabled("%s", out.headers[nc].c_str());
    return true;
}

bool render_kpi_grid(const StageOutput& out, const ViewConfig& cfg) {
    auto nums = collect_numeric_filtered(out, cfg, 12);
    if (nums.empty()) { info_text("Need numeric columns"); return false; }
    ImVec2 avail = ImGui::GetContentRegionAvail();
    int per_row = std::max(1, (int)(avail.x / 220));
    int idx = 0;
    for (auto& nc : nums) {
        double last = std::nan("");
        for (auto v : nc.vals) if (!std::isnan(v)) last = v;
        if (std::isnan(last)) last = 0;
        char buf[64];
        if (std::abs(last) >= 1e6)       std::snprintf(buf, sizeof(buf), "%.2fM", last / 1e6);
        else if (std::abs(last) >= 1e3)  std::snprintf(buf, sizeof(buf), "%.2fK", last / 1e3);
        else                              std::snprintf(buf, sizeof(buf), "%.4g", last);

        ImGui::BeginChild((ImGuiID)(0x1000 + idx), ImVec2(210, 100), true);
        ImGui::TextDisabled("%s", nc.name.c_str());
        ImGui::SetWindowFontScale(2.4f);
        ImGui::TextUnformatted(buf);
        ImGui::SetWindowFontScale(1.0f);
        ImGui::EndChild();

        if ((idx % per_row) != (per_row - 1)) ImGui::SameLine();
        idx++;
    }
    return true;
}

bool render_stem(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 1);
    if (nums.empty()) { info_text("Need 1 numeric column"); return false; }
    int n = (int)nums[0].vals.size();
    std::vector<double> xs(n); for (int i = 0; i < n; ++i) xs[i] = i;
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##stem", size, 0)) return false;
    ImPlot::SetupAxes(nullptr, nums[0].name.c_str(),
                      axflag(cfg), axflag(cfg));
    ImPlot::PlotStems(nums[0].name.c_str(), xs.data(), nums[0].vals.data(), n);
    ImPlot::EndPlot();
    return true;
}

bool render_errorbars(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 4);
    if (nums.size() < 2) { info_text("Need 2 numeric columns (value, err)"); return false; }
    int n = (int)nums[0].vals.size();
    std::vector<double> xs(n); for (int i = 0; i < n; ++i) xs[i] = i;
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##eb", size, 0)) return false;
    ImPlot::SetupAxes(nullptr, nums[0].name.c_str(),
                      axflag(cfg), axflag(cfg));
    ImPlot::PlotErrorBars(nums[0].name.c_str(), xs.data(),
                          nums[0].vals.data(), nums[1].vals.data(), n);
    ImPlot::PlotScatter("##s", xs.data(), nums[0].vals.data(), n);
    ImPlot::EndPlot();
    return true;
}

// BoxPlot: agrupar por categoria, calcular min/p25/p50/p75/max y dibujar
// rectangulos manuales via PlotShaded + lineas.
bool render_boxplot(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    int cat = resolve_cat(out, cfg, first_category_col(out));
    auto nums = collect_numeric_filtered(out, cfg, 1);
    if (cat < 0 || nums.empty()) { info_text("Need 1 category + 1 numeric"); return false; }
    auto cats = extract_category(out, cat);
    int n = std::min((int)cats.size(), (int)nums[0].vals.size());
    if (n == 0) return false;

    // Group values by category
    std::unordered_map<std::string, std::vector<double>> groups;
    std::vector<std::string> order;
    for (int i = 0; i < n; ++i) {
        if (groups.find(cats[i]) == groups.end()) order.push_back(cats[i]);
        double d = nums[0].vals[i];
        if (!std::isnan(d)) groups[cats[i]].push_back(d);
    }
    int G = (int)order.size();
    if (G == 0) return false;

    std::vector<double> mn(G), p25(G), p50(G), p75(G), mx(G), xs(G);
    std::vector<const char*> labels(G);
    for (int g = 0; g < G; ++g) {
        auto& v = groups[order[g]];
        std::sort(v.begin(), v.end());
        int N = (int)v.size();
        xs[g]    = g;
        labels[g]= order[g].c_str();
        if (N == 0) { mn[g]=p25[g]=p50[g]=p75[g]=mx[g]=0; continue; }
        mn[g]  = v.front();
        mx[g]  = v.back();
        p25[g] = v[std::min(N - 1, (int)(N * 0.25))];
        p50[g] = v[std::min(N - 1, (int)(N * 0.50))];
        p75[g] = v[std::min(N - 1, (int)(N * 0.75))];
    }

    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##box", size, 0)) return false;
    ImPlot::SetupAxes(out.headers[cat].c_str(), nums[0].name.c_str(),
                      axflag(cfg), axflag(cfg));
    ImPlot::SetupAxisTicks(ImAxis_X1, xs.data(), G, labels.data(), false);

    // Whiskers: stems from min to max
    for (int g = 0; g < G; ++g) {
        double lo[2] = { mn[g], mx[g] };
        double xx[2] = { xs[g], xs[g] };
        ImPlot::PlotLine("##wh", xx, lo, 2);
    }
    // Box: p25..p75 as bars centered on p50
    std::vector<double> mid(G), half(G);
    for (int g = 0; g < G; ++g) {
        mid[g]  = (p25[g] + p75[g]) * 0.5;
        half[g] = (p75[g] - p25[g]) * 0.5;
    }
    ImPlot::PlotErrorBars("box", xs.data(), mid.data(), half.data(), G);
    ImPlot::PlotScatter("median", xs.data(), p50.data(), G);
    ImPlot::EndPlot();
    return true;
}

// Candlestick: tiempo + O/H/L/C. Asume 4 primeras cols numericas en ese orden.
bool render_candlestick(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 8);
    if (nums.size() < 4) { info_text("Need 4 numeric columns: O/H/L/C"); return false; }
    int n = (int)nums[0].vals.size();
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##candle", size, 0)) return false;
    ImPlot::SetupAxes("t", "price", axflag(cfg), axflag(cfg));
    std::vector<double> xs(n); for (int i = 0; i < n; ++i) xs[i] = i;
    const auto& O = nums[0].vals;
    const auto& H = nums[1].vals;
    const auto& L = nums[2].vals;
    const auto& C = nums[3].vals;
    // Wicks
    for (int i = 0; i < n; ++i) {
        double xx[2] = { xs[i], xs[i] };
        double yy[2] = { L[i], H[i] };
        ImPlot::PlotLine("##wick", xx, yy, 2);
    }
    // Body via PlotBars(mid, |C-O|)? Simpler: separate lines.
    std::vector<double> body_low(n), body_high(n), body_mid(n), body_err(n);
    for (int i = 0; i < n; ++i) {
        body_low[i]  = std::min(O[i], C[i]);
        body_high[i] = std::max(O[i], C[i]);
        body_mid[i]  = (body_low[i] + body_high[i]) * 0.5;
        body_err[i]  = (body_high[i] - body_low[i]) * 0.5;
    }
    ImPlot::PlotErrorBars("OHLC", xs.data(), body_mid.data(), body_err.data(), n);
    ImPlot::EndPlot();
    return true;
}

bool render_radar(const StageOutput& out, const ViewConfig& cfg, ImVec2 size) {
    auto nums = collect_numeric_filtered(out, cfg, 12);
    if (nums.size() < 3) { info_text("Need 3+ numeric columns"); return false; }
    int K = (int)nums.size();
    int n = (int)nums[0].vals.size();
    if (n == 0) return false;
    // Take first row as the polygon.
    std::vector<double> xs(K + 1), ys(K + 1);
    double radius_norm = 0;
    for (int k = 0; k < K; ++k) {
        double d = nums[k].vals[0];
        if (std::isnan(d)) d = 0;
        radius_norm = std::max(radius_norm, std::abs(d));
    }
    if (radius_norm == 0) radius_norm = 1;
    for (int k = 0; k < K; ++k) {
        double v = nums[k].vals[0]; if (std::isnan(v)) v = 0;
        double angle = 2 * 3.14159265358979 * k / K - 3.14159265358979 / 2;
        double r = v / radius_norm;
        xs[k] = std::cos(angle) * r;
        ys[k] = std::sin(angle) * r;
    }
    xs[K] = xs[0]; ys[K] = ys[0];
    maybe_fit(cfg);
    if (!ImPlot::BeginPlot("##radar", size,
                           ImPlotFlags_Equal | ImPlotFlags_NoLegend)) return false;
    ImPlot::SetupAxesLimits(-1.2, 1.2, -1.2, 1.2, ImPlotCond_Always);
    ImPlot::SetupAxes(nullptr, nullptr, ImPlotAxisFlags_NoDecorations,
                      ImPlotAxisFlags_NoDecorations);
    // Grid rings
    for (double rr : {0.25, 0.5, 0.75, 1.0}) {
        double gx[64], gy[64];
        for (int i = 0; i < 64; ++i) {
            double a = 2 * 3.14159265358979 * i / 63;
            gx[i] = std::cos(a) * rr; gy[i] = std::sin(a) * rr;
        }
        ImPlot::PlotLine("##grid", gx, gy, 64);
    }
    ImPlot::PlotLine("radar", xs.data(), ys.data(), K + 1);
    // Axis labels
    for (int k = 0; k < K; ++k) {
        double a = 2 * 3.14159265358979 * k / K - 3.14159265358979 / 2;
        ImPlot::PlotText(nums[k].name.c_str(), std::cos(a) * 1.1, std::sin(a) * 1.1);
    }
    ImPlot::EndPlot();
    return true;
}

} // anon

bool render(const StageOutput& out, ViewMode mode,
            const ViewConfig& cfg, ImVec2 size,
            int* clicked_row_out) {
    if (clicked_row_out) *clicked_row_out = -1;
    if (out.rows == 0 || out.cols == 0) {
        info_text("No data");
        return false;
    }
    switch (mode) {
        case ViewMode::Table:        return false;
        case ViewMode::Bar:
        case ViewMode::Column:
        case ViewMode::GroupedBar:
        case ViewMode::StackedBar:   return render_bar_like(out, mode, cfg, size, clicked_row_out);
        case ViewMode::Line:
        case ViewMode::Area:
        case ViewMode::Stairs:       return render_line_like(out, mode, cfg, size);
        case ViewMode::Scatter:      return render_scatter(out, cfg, size, clicked_row_out);
        case ViewMode::Bubble:       return render_bubble(out, cfg, size, clicked_row_out);
        case ViewMode::Histogram:    return render_histogram(out, cfg, size);
        case ViewMode::Histogram2D:  return render_hist2d(out, cfg, size);
        case ViewMode::Heatmap:      return render_heatmap(out, cfg, size, clicked_row_out);
        case ViewMode::BoxPlot:      return render_boxplot(out, cfg, size);
        case ViewMode::Stem:         return render_stem(out, cfg, size);
        case ViewMode::ErrorBars:    return render_errorbars(out, cfg, size);
        case ViewMode::Pie:          return render_pie(out, cfg, false, size, clicked_row_out);
        case ViewMode::Donut:        return render_pie(out, cfg, true,  size, clicked_row_out);
        case ViewMode::Funnel:       return render_funnel(out, cfg, size, clicked_row_out);
        case ViewMode::Waterfall:    return render_waterfall(out, cfg, size);
        case ViewMode::KPI:          return render_kpi_single(out, cfg);
        case ViewMode::KPIGrid:      return render_kpi_grid(out, cfg);
        case ViewMode::Candlestick:  return render_candlestick(out, cfg, size);
        case ViewMode::Radar:        return render_radar(out, cfg, size);
    }
    return false;
}

} // namespace viz