#include "viz/sankey.h"

#include <algorithm>
#include <cmath>
#include <queue>

namespace {

// BFS topologico: asigna a cada nodo el max(level(src)+1).
// Nodos sin in-edges arrancan en nivel 0.
std::vector<int> compute_levels(int n_nodes, const std::vector<SankeyLink>& links) {
    std::vector<int>              indeg(n_nodes, 0);
    std::vector<std::vector<int>> out(n_nodes);
    for (const auto& l : links) {
        if (l.src < 0 || l.src >= n_nodes || l.dst < 0 || l.dst >= n_nodes) continue;
        out[l.src].push_back(l.dst);
        indeg[l.dst]++;
    }
    std::vector<int>  level(n_nodes, 0);
    std::vector<bool> visited(n_nodes, false);
    std::queue<int>   q;
    for (int i = 0; i < n_nodes; i++) {
        if (indeg[i] == 0) {
            q.push(i);
            visited[i] = true;
        }
    }
    while (!q.empty()) {
        int u = q.front(); q.pop();
        for (int v : out[u]) {
            if (level[v] < level[u] + 1) level[v] = level[u] + 1;
            indeg[v]--;
            if (indeg[v] == 0 && !visited[v]) {
                visited[v] = true;
                q.push(v);
            }
        }
    }
    return level;
}

ImU32 node_color(int idx) {
    // Paleta indigo/teal/amber rotativa.
    static const ImU32 palette[] = {
        IM_COL32(120, 144, 252, 230),
        IM_COL32( 92, 200, 200, 230),
        IM_COL32(250, 176,  92, 230),
        IM_COL32(180, 120, 200, 230),
        IM_COL32( 92, 200, 130, 230),
        IM_COL32(250, 120, 130, 230),
        IM_COL32(180, 200,  92, 230),
        IM_COL32(120, 200, 230, 230),
    };
    constexpr int N = sizeof(palette) / sizeof(palette[0]);
    return palette[((idx % N) + N) % N];
}

} // namespace

void sankey(const char* id,
            const std::vector<SankeyNode>& nodes,
            const std::vector<SankeyLink>& links,
            ImVec2                          size) {
    ImGui::PushID(id);

    ImVec2 avail = ImGui::GetContentRegionAvail();
    float  W     = (size.x > 0.0f) ? size.x : avail.x;
    float  H     = (size.y > 0.0f) ? size.y : 300.0f;

    ImVec2 origin = ImGui::GetCursorScreenPos();
    ImGui::Dummy(ImVec2(W, H));

    int N = (int)nodes.size();
    if (N == 0) { ImGui::PopID(); return; }

    auto levels = compute_levels(N, links);
    int  L      = 0;
    for (int v : levels) L = std::max(L, v + 1);
    if (L < 1) L = 1;

    // Magnitud por nodo = max(in_total, out_total). Para sources es out_total,
    // para sinks es in_total, en intermedios usamos el max para que los
    // rectangulos encajen visualmente.
    std::vector<float> in_tot (N, 0.0f);
    std::vector<float> out_tot(N, 0.0f);
    for (const auto& l : links) {
        if (l.src < 0 || l.src >= N || l.dst < 0 || l.dst >= N) continue;
        in_tot [l.dst] += l.value;
        out_tot[l.src] += l.value;
    }
    std::vector<float> magnitude(N, 0.0f);
    for (int i = 0; i < N; i++) magnitude[i] = std::max(in_tot[i], out_tot[i]);

    // Por columna: total magnitud + lista de nodos.
    std::vector<std::vector<int>> by_col(L);
    for (int i = 0; i < N; i++) by_col[levels[i]].push_back(i);

    // Geometria
    const float pad_x      = 16.0f;
    const float pad_y      = 16.0f;
    const float node_w     = 14.0f;
    const float gap_y      = 6.0f;
    float       avail_w    = std::max(1.0f, W - 2.0f * pad_x);
    float       col_pitch  = (L > 1) ? (avail_w - node_w) / (float)(L - 1) : 0.0f;
    float       avail_h    = std::max(1.0f, H - 2.0f * pad_y);

    // Por columna calcula scale: pixel/value tal que (sum mag + gaps) cabe en avail_h.
    std::vector<float> col_scale(L, 1.0f);
    std::vector<float> col_yoff (L, 0.0f);
    for (int c = 0; c < L; c++) {
        float sum_m = 0.0f;
        for (int i : by_col[c]) sum_m += magnitude[i];
        int   cnt   = (int)by_col[c].size();
        float gaps  = (cnt > 1) ? (cnt - 1) * gap_y : 0.0f;
        float usable = std::max(1.0f, avail_h - gaps);
        col_scale[c] = (sum_m > 0.0f) ? (usable / sum_m) : 0.0f;
        // Centrar verticalmente
        float total_h = sum_m * col_scale[c] + gaps;
        col_yoff[c]   = pad_y + (avail_h - total_h) * 0.5f;
    }

    // Posiciones de nodos: para cada nodo, y_top y altura.
    struct NodeBox {
        float x_left, x_right;
        float y_top,  y_bot;
    };
    std::vector<NodeBox> boxes(N);
    // y_cursor por columna
    std::vector<float> y_cursor(L, 0.0f);
    for (int c = 0; c < L; c++) y_cursor[c] = col_yoff[c];
    for (int c = 0; c < L; c++) {
        for (int i : by_col[c]) {
            float h = magnitude[i] * col_scale[c];
            if (h < 1.0f) h = 1.0f;
            float xl = pad_x + c * col_pitch;
            boxes[i].x_left  = xl;
            boxes[i].x_right = xl + node_w;
            boxes[i].y_top   = y_cursor[c];
            boxes[i].y_bot   = y_cursor[c] + h;
            y_cursor[c]      = boxes[i].y_bot + gap_y;
        }
    }

    // Por nodo: cursores de salida (en src) y entrada (en dst), incrementan al consumir links.
    std::vector<float> src_cursor(N, 0.0f);
    std::vector<float> dst_cursor(N, 0.0f);
    for (int i = 0; i < N; i++) {
        src_cursor[i] = boxes[i].y_top;
        dst_cursor[i] = boxes[i].y_top;
    }

    ImDrawList* dl = ImGui::GetWindowDrawList();

    // Render links primero (por debajo de los nodos).
    // Sort por src para estetica (topdown).
    std::vector<int> link_order(links.size());
    for (size_t i = 0; i < links.size(); i++) link_order[i] = (int)i;
    std::sort(link_order.begin(), link_order.end(),
        [&](int a, int b) {
            const auto& la = links[a];
            const auto& lb = links[b];
            if (la.src != lb.src) return la.src < lb.src;
            return la.dst < lb.dst;
        });

    for (int li : link_order) {
        const auto& l = links[li];
        if (l.src < 0 || l.src >= N || l.dst < 0 || l.dst >= N) continue;
        if (l.value <= 0.0f) continue;

        float h_src = l.value * col_scale[levels[l.src]];
        float h_dst = l.value * col_scale[levels[l.dst]];
        float ya0 = src_cursor[l.src];
        float ya1 = ya0 + h_src;
        float yb0 = dst_cursor[l.dst];
        float yb1 = yb0 + h_dst;
        src_cursor[l.src] += h_src;
        dst_cursor[l.dst] += h_dst;

        float xa = boxes[l.src].x_right;
        float xb = boxes[l.dst].x_left;

        // Dos beziers cubicos formando una banda (top + bottom + cierre).
        ImVec2 p0(origin.x + xa, origin.y + (ya0 + ya1) * 0.5f);
        ImVec2 p1(origin.x + xb, origin.y + (yb0 + yb1) * 0.5f);
        float ctrl_dx = (xb - xa) * 0.5f;

        // Aproximar la banda con poligono de 24 segmentos top + 24 bottom.
        const int   STEPS = 24;
        std::vector<ImVec2> poly;
        poly.reserve(STEPS * 2 + 2);
        for (int s = 0; s <= STEPS; s++) {
            float t  = (float)s / (float)STEPS;
            float u  = 1.0f - t;
            float bx = u*u*u * (xa) + 3*u*u*t * (xa + ctrl_dx) + 3*u*t*t * (xb - ctrl_dx) + t*t*t * (xb);
            float by_top = u*u*u * (ya0) + 3*u*u*t * (ya0) + 3*u*t*t * (yb0) + t*t*t * (yb0);
            poly.push_back(ImVec2(origin.x + bx, origin.y + by_top));
            (void)p0; (void)p1;
        }
        for (int s = STEPS; s >= 0; s--) {
            float t  = (float)s / (float)STEPS;
            float u  = 1.0f - t;
            float bx = u*u*u * (xa) + 3*u*u*t * (xa + ctrl_dx) + 3*u*t*t * (xb - ctrl_dx) + t*t*t * (xb);
            float by_bot = u*u*u * (ya1) + 3*u*u*t * (ya1) + 3*u*t*t * (yb1) + t*t*t * (yb1);
            poly.push_back(ImVec2(origin.x + bx, origin.y + by_bot));
        }
        ImU32 src_c = node_color(l.src);
        // alpha bajo
        ImU32 band  = (src_c & 0x00FFFFFF) | (90u << 24);
        dl->AddConvexPolyFilled(poly.data(), (int)poly.size(), band);
    }

    // Render nodos + labels.
    for (int i = 0; i < N; i++) {
        ImVec2 a(origin.x + boxes[i].x_left,  origin.y + boxes[i].y_top);
        ImVec2 b(origin.x + boxes[i].x_right, origin.y + boxes[i].y_bot);
        dl->AddRectFilled(a, b, node_color(i));
        // Label: a la izquierda si es ultima columna, derecha si no.
        const char* lbl = nodes[i].label.c_str();
        ImVec2 ts = ImGui::CalcTextSize(lbl);
        float  ly = (a.y + b.y) * 0.5f - ts.y * 0.5f;
        if (levels[i] == L - 1) {
            // a la izquierda del rect
            dl->AddText(ImVec2(a.x - ts.x - 4.0f, ly),
                        IM_COL32(220, 222, 235, 230), lbl);
        } else {
            dl->AddText(ImVec2(b.x + 4.0f, ly),
                        IM_COL32(220, 222, 235, 230), lbl);
        }
    }

    ImGui::PopID();
}