#include "viz/treemap.h"

#include <algorithm>
#include <cmath>
#include <cstdio>

namespace {

// Worst aspect ratio en la fila si añadimos `next` (Bruls et al.).
// row: vector de areas (mismo orden que items), w: lado corto del rect actual.
float worst_ratio(const std::vector<float>& row, float w) {
    if (row.empty() || w <= 0.0f) return 1e30f;
    float sum = 0.0f, mn = row[0], mx = row[0];
    for (float v : row) {
        sum += v;
        if (v < mn) mn = v;
        if (v > mx) mx = v;
    }
    if (sum <= 0.0f) return 1e30f;
    float w2 = w * w;
    float s2 = sum * sum;
    float r1 = (w2 * mx) / s2;
    float r2 = s2 / (w2 * mn);
    return std::max(r1, r2);
}

// Coloca una fila en el rect actual, devuelve el rect remanente.
// Si el rect es mas ancho que alto, la fila ocupa el lado izquierdo (vertical strip).
struct RowOut {
    ImVec2 next_min;
    ImVec2 next_max;
};

RowOut place_row(const std::vector<float>& row,
                 const std::vector<int>&   row_idx,
                 ImVec2                    rect_min,
                 ImVec2                    rect_max,
                 std::vector<TreemapRect>& out) {
    float W   = rect_max.x - rect_min.x;
    float H   = rect_max.y - rect_min.y;
    float sum = 0.0f;
    for (float v : row) sum += v;
    if (sum <= 0.0f) return {rect_min, rect_max};

    bool horizontal_strip = (W >= H); // strip ocupa lado izquierdo (vertical column de cells horizontales)
    if (horizontal_strip) {
        float strip_w = sum / H;
        float y       = rect_min.y;
        for (size_t i = 0; i < row.size(); i++) {
            float h = row[i] / strip_w;
            TreemapRect r;
            r.min  = ImVec2(rect_min.x, y);
            r.max  = ImVec2(rect_min.x + strip_w, std::min(y + h, rect_max.y));
            r.item = nullptr;
            out[row_idx[i]] = r;
            y += h;
        }
        return {ImVec2(rect_min.x + strip_w, rect_min.y), rect_max};
    } else {
        float strip_h = sum / W;
        float x       = rect_min.x;
        for (size_t i = 0; i < row.size(); i++) {
            float w = row[i] / strip_h;
            TreemapRect r;
            r.min  = ImVec2(x, rect_min.y);
            r.max  = ImVec2(std::min(x + w, rect_max.x), rect_min.y + strip_h);
            r.item = nullptr;
            out[row_idx[i]] = r;
            x += w;
        }
        return {rect_min, ImVec2(rect_max.x, rect_min.y + strip_h)};
    }
}

} // namespace

std::vector<TreemapRect> treemap_layout(const std::vector<TreemapItem>& items,
                                        ImVec2                           region) {
    std::vector<TreemapRect> out(items.size(), {ImVec2(0,0), ImVec2(0,0), nullptr});
    if (items.empty() || region.x <= 0.0f || region.y <= 0.0f) return out;

    // Filtra y ordena indices por value desc.
    std::vector<int> idx;
    idx.reserve(items.size());
    float total = 0.0f;
    for (size_t i = 0; i < items.size(); i++) {
        if (items[i].value > 0.0f) {
            idx.push_back((int)i);
            total += items[i].value;
        }
    }
    if (idx.empty() || total <= 0.0f) return out;

    std::sort(idx.begin(), idx.end(),
        [&](int a, int b) { return items[a].value > items[b].value; });

    // Areas escaladas al area total del rect.
    float region_area = region.x * region.y;
    float scale       = region_area / total;
    std::vector<float> areas(idx.size());
    for (size_t i = 0; i < idx.size(); i++) areas[i] = items[idx[i]].value * scale;

    ImVec2 rmin(0.0f, 0.0f);
    ImVec2 rmax = region;

    std::vector<float> row;
    std::vector<int>   row_indices; // indices en `out` (= idx[i])
    size_t cursor = 0;

    while (cursor < idx.size()) {
        float w = std::min(rmax.x - rmin.x, rmax.y - rmin.y);
        if (w <= 0.0f) break;

        std::vector<float> row_try = row;
        row_try.push_back(areas[cursor]);
        float wr_now  = worst_ratio(row, w);
        float wr_next = worst_ratio(row_try, w);

        if (row.empty() || wr_next <= wr_now) {
            row     = row_try;
            row_indices.push_back(idx[cursor]);
            cursor++;
        } else {
            RowOut ro = place_row(row, row_indices, rmin, rmax, out);
            rmin = ro.next_min;
            rmax = ro.next_max;
            row.clear();
            row_indices.clear();
        }
    }
    if (!row.empty()) {
        place_row(row, row_indices, rmin, rmax, out);
    }

    // Asocia el item.
    for (size_t i = 0; i < items.size(); i++) {
        out[i].item = &items[i];
    }
    return out;
}

void treemap(const char* id,
             const std::vector<TreemapItem>& items,
             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 : 200.0f;

    ImVec2 origin = ImGui::GetCursorScreenPos();
    ImGui::Dummy(ImVec2(w, h));

    auto rects = treemap_layout(items, ImVec2(w, h));

    ImDrawList* dl = ImGui::GetWindowDrawList();

    // Borde exterior tenue
    dl->AddRect(origin, ImVec2(origin.x + w, origin.y + h),
                IM_COL32(80, 80, 90, 200), 0.0f, 0, 1.0f);

    for (const auto& r : rects) {
        if (!r.item) continue;
        ImVec2 a(origin.x + r.min.x, origin.y + r.min.y);
        ImVec2 b(origin.x + r.max.x, origin.y + r.max.y);
        if (b.x - a.x < 1.0f || b.y - a.y < 1.0f) continue;

        dl->AddRectFilled(a, b, r.item->color);
        dl->AddRect(a, b, IM_COL32(20, 22, 28, 220), 0.0f, 0, 1.0f);

        // Label si cabe
        const char* lbl     = r.item->label.c_str();
        ImVec2      ts      = ImGui::CalcTextSize(lbl);
        float       cell_w  = b.x - a.x;
        float       cell_h  = b.y - a.y;
        if (ts.x + 6.0f <= cell_w && ts.y + 4.0f <= cell_h) {
            dl->AddText(ImVec2(a.x + 4.0f, a.y + 3.0f),
                        IM_COL32(245, 246, 250, 240), lbl);
            // valor en segunda linea si tambien cabe
            char buf[32];
            std::snprintf(buf, sizeof(buf), "%.0f", r.item->value);
            ImVec2 vs = ImGui::CalcTextSize(buf);
            if (vs.x + 6.0f <= cell_w && ts.y + vs.y + 6.0f <= cell_h) {
                dl->AddText(ImVec2(a.x + 4.0f, a.y + 3.0f + ts.y + 1.0f),
                            IM_COL32(220, 222, 235, 200), buf);
            }
        }
    }

    ImGui::PopID();
}