visualizer.cpp

#include "visualizer.hpp"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <thread>
#include <chrono>
#include <sstream>
#include <iomanip>

Visualizer::Visualizer() = default;

Visualizer::~Visualizer() {
    if (initialized_) cleanup();
}

void Visualizer::init() {
    initscr();
    cbreak();
    noecho();
    keypad(stdscr, TRUE);
    nodelay(stdscr, TRUE);
    curs_set(0);             // hide cursor
    setup_colors();
    getmaxyx(stdscr, rows_, cols_);
    compute_layout();

    bubbles_.clear();
    int active_w = cols_ > 4 ? cols_ - 4 : 20;
    int active_h = (net_bottom_ - net_top_) > 4 ? (net_bottom_ - net_top_ - 4) : 10;
    for (int i = 0; i < 8; ++i) {
        bubbles_.push_back({
            static_cast<float>(2 + rand() % active_w),
            static_cast<float>(net_top_ + 2 + rand() % active_h),
            0.02f + 0.03f * (static_cast<float>(rand() % 100) / 100.f),
            0.3f + 0.2f * (static_cast<float>(rand() % 100) / 100.f)
        });
    }

    initialized_ = true;
}

void Visualizer::cleanup() {
    if (initialized_) {
        curs_set(1);
        endwin();
        initialized_ = false;
    }
}

void Visualizer::setup_colors() {
    start_color();
    use_default_colors();
    // init_pair(pair, fg, bg)
    init_pair(CP_DEFAULT, COLOR_WHITE,   -1);
    init_pair(CP_TITLE,   COLOR_BLACK,   COLOR_CYAN);
    init_pair(CP_ACTIVE,  COLOR_CYAN,    -1);
    init_pair(CP_DIM,     COLOR_BLUE,    -1);
    init_pair(CP_BALL,    COLOR_YELLOW,  -1);
    init_pair(CP_EDGE,    COLOR_WHITE,   -1);
    init_pair(CP_ANSWER,  COLOR_GREEN,   -1);
    init_pair(CP_WARN,    COLOR_YELLOW,  -1);
    init_pair(CP_BAR,     COLOR_BLACK,   COLOR_BLUE);
    init_pair(CP_PROMPT,  COLOR_MAGENTA, -1);
    init_pair(CP_LOSS,    COLOR_RED,     -1);
    init_pair(CP_LABEL,   COLOR_CYAN,    -1);
}

void Visualizer::compute_layout() {
    getmaxyx(stdscr, rows_, cols_);
    net_top_    = 1;
    net_bottom_ = rows_ - 9;
    info_top_   = rows_ - 8;
    bar_y_      = rows_ - 1;

    int usable = cols_ - 4;
    layer_col_x_.resize(4);
    for (int i = 0; i < 4; ++i)
        layer_col_x_[i] = 2 + (usable * i) / 3;
}


void Visualizer::attron_pair(ColorPair cp, bool bold) {
    attron(COLOR_PAIR(cp));
    if (bold) attron(A_BOLD);
}
void Visualizer::attroff_pair(ColorPair cp, bool bold) {
    attroff(COLOR_PAIR(cp));
    if (bold) attroff(A_BOLD);
}

std::string Visualizer::truncate(const std::string& s, int max_w) const {
    if (static_cast<int>(s.size()) <= max_w) return s;
    return s.substr(0, std::max(0, max_w - 3)) + "...";
}

int Visualizer::neuron_y(int center_y, int idx, int total) const {
    int spacing = (total <= 1) ? 0 : (net_bottom_ - net_top_ - 4) / (total - 1);
    spacing = std::max(1, spacing);
    int start = center_y - spacing * (total - 1) / 2;
    return std::max(net_top_ + 2, std::min(net_bottom_ - 1, start + idx * spacing));
}

int Visualizer::layer_center_y(int /*layer_idx*/, int /*n*/) const {
    return (net_top_ + net_bottom_) / 2;
}

int Visualizer::neuron_y_abs(int layer_idx, int neuron_idx, int total) const {
    int cy = (net_top_ + net_bottom_) / 2;
    return neuron_y(cy, neuron_idx, total);
}

void Visualizer::sleep_ms(int ms) {
    std::this_thread::sleep_for(std::chrono::milliseconds(ms));
}

int Visualizer::get_key() {
    return getch();
}

bool Visualizer::is_quit_key(int ch) {
    return ch == 'q' || ch == 'Q' || ch == 27 /*ESC*/ || ch == 3 /*Ctrl-C*/;
}


void Visualizer::draw_title_bar(const std::string& title) {
    attron_pair(CP_TITLE, true);
    // Fill entire row with spaces first
    std::string pad(cols_, ' ');
    mvaddstr(0, 0, pad.c_str());
    // Center the title
    int x = std::max(0, (cols_ - static_cast<int>(title.size())) / 2);
    mvaddstr(0, x, title.c_str());
    attroff_pair(CP_TITLE, true);
}

void Visualizer::draw_status_bar(const std::string& msg) {
    attron_pair(CP_BAR, true);
    std::string pad(cols_, ' ');
    mvaddstr(bar_y_, 0, pad.c_str());
    mvaddstr(bar_y_, 1, truncate(msg, cols_ - 2).c_str());
    attroff_pair(CP_BAR, true);
}


void Visualizer::draw_layer(int, int, const std::vector<float>&, const std::string&, bool) {}
void Visualizer::draw_edges(int, int, int, int, const std::vector<float>&, const std::vector<float>&) {}
void Visualizer::draw_ball(const BallPos&, bool) {}


void Visualizer::draw_network_panel(const NeuralNet& nn,
                                     int active_layer,
                                     const BallPos& ball,
                                     bool ball_visible) {
    animation_time_ += 0.2f;

    int active_w = cols_ > 4 ? cols_ - 4 : 20;
    int active_h = (net_bottom_ - net_top_) > 4 ? (net_bottom_ - net_top_ - 4) : 10;
    for (auto& b : bubbles_) {
        b.y -= b.vy;
        if (b.y < net_top_ + 2) {
            b.y = net_bottom_ - 1;
            b.x = 2 + rand() % active_w;
            b.vy = 0.02f + 0.03f * (static_cast<float>(rand() % 100) / 100.f);
            b.r = 0.3f + 0.2f * (static_cast<float>(rand() % 100) / 100.f);
        }
    }

    struct LayerInfo {
        std::vector<float> acts;
        std::string label;
        int col_x;
    };

    auto norm = [](std::vector<float> v) {
        float mx = 0.f;
        for (float x : v) mx = std::max(mx, std::abs(x));
        if (mx > 0.f) for (auto& x : v) x = std::abs(x) / mx;
        return v;
    };

    std::vector<LayerInfo> layers = {
        { norm(nn.a0),  "INPUT\n(emb)",  layer_col_x_[0] },
        { norm(nn.a1),  "HIDDEN1",       layer_col_x_[1] },
        { norm(nn.a2),  "HIDDEN2",       layer_col_x_[2] },
        { nn.ao,        "OUTPUT",        layer_col_x_[3] },
    };

    for (int i = 0; i < 4; ++i) {
        attron_pair(CP_LABEL, true);
        std::string lbl = layers[i].label;
        auto pos = lbl.find('\n');
        if (pos != std::string::npos) {
            std::string part1 = lbl.substr(0, pos);
            std::string part2 = lbl.substr(pos + 1);
            mvaddstr(net_top_, layers[i].col_x - part1.size() / 2, part1.c_str());
            mvaddstr(net_top_ + 1, layers[i].col_x - part2.size() / 2, part2.c_str());
        } else {
            mvaddstr(net_top_, layers[i].col_x - lbl.size() / 2, lbl.c_str());
        }
        attroff_pair(CP_LABEL, true);
    }

    for (int i = 0; i + 1 < 4; ++i) {
        int n1 = std::min(static_cast<int>(layers[i].acts.size()), MAX_VIS_NEURONS);
        int n2 = std::min(static_cast<int>(layers[i+1].acts.size()), MAX_VIS_NEURONS);
        
        for (int j = 0; j < n1; ++j) {
            int y1 = neuron_y_abs(i, j, n1);
            int x1 = layers[i].col_x;
            for (int k = 0; k < n2; ++k) {
                int y2 = neuron_y_abs(i+1, k, n2);
                int x2 = layers[i+1].col_x;

                attron(COLOR_PAIR(CP_DIM));
                for (int x = x1 + 1; x < x2; ++x) {
                    float t = static_cast<float>(x - x1) / (x2 - x1);
                    int y = static_cast<int>(y1 + t * (y2 - y1) + 0.5f);
                    if (y >= net_top_ + 2 && y < net_bottom_ && x >= 0 && x < cols_) {
                        mvaddch(y, x, '.');
                    }
                }
                attroff(COLOR_PAIR(CP_DIM));
            }
        }
    }

    struct Metaball {
        float x;
        float y;
        float r;
        int type;
    };

    std::vector<Metaball> metaballs;

    for (const auto& b : bubbles_) {
        metaballs.push_back({ b.x, b.y, b.r, 0 });
    }

    for (int i = 0; i < 4; ++i) {
        int total = std::min(static_cast<int>(layers[i].acts.size()), MAX_VIS_NEURONS);
        bool layer_is_active = (active_layer == -1 || active_layer == i);
        
        for (int j = 0; j < total; ++j) {
            float act = layers[i].acts[j];
            float r = layer_is_active ? (0.8f + 1.1f * act) : (0.5f + 0.3f * act);
            
            float pulse = std::sin(animation_time_ + i * 1.5f + j * 0.7f) * 0.12f;
            r += pulse;

            metaballs.push_back({
                static_cast<float>(layers[i].col_x),
                static_cast<float>(neuron_y_abs(i, j, total)),
                r,
                layer_is_active ? 1 : 0
            });
        }
    }

    for (int i = 0; i + 1 < 4; ++i) {
        int n1 = std::min(static_cast<int>(layers[i].acts.size()), MAX_VIS_NEURONS);
        int n2 = std::min(static_cast<int>(layers[i+1].acts.size()), MAX_VIS_NEURONS);
        bool edge_is_active = (active_layer == -1 || active_layer == i);

        for (int j = 0; j < n1; ++j) {
            float act1 = layers[i].acts[j];
            for (int k = 0; k < n2; ++k) {
                float act2 = layers[i+1].acts[k];
                float strength = (act1 + act2) * 0.5f;

                float weight = 0.0f;
                if (i == 0) {
                    if (k * nn.EMBED_DIM + j < static_cast<int>(nn.w1.size()))
                        weight = nn.w1[k * nn.EMBED_DIM + j];
                } else if (i == 1) {
                    if (k * nn.HIDDEN1 + j < static_cast<int>(nn.w2.size()))
                        weight = nn.w2[k * nn.HIDDEN1 + j];
                } else if (i == 2) {
                    if (k * nn.HIDDEN2 + j < static_cast<int>(nn.wo.size()))
                        weight = nn.wo[k * nn.HIDDEN2 + j];
                }

                if (strength > 0.15f) {
                    float x1 = layers[i].col_x;
                    float y1 = neuron_y_abs(i, j, n1);
                    float x2 = layers[i+1].col_x;
                    float y2 = neuron_y_abs(i+1, k, n2);

                    int bridge_type = 0;
                    if (edge_is_active) {
                        bridge_type = (weight > 0.03f) ? 1 : ((weight < -0.03f) ? 2 : 0);
                    }

                    for (int step = 1; step <= 4; ++step) {
                        float t = step * 0.2f;
                        float bx = x1 + t * (x2 - x1);
                        float by = y1 + t * (y2 - y1);
                        float br = 0.3f + 0.25f * strength + 0.15f * std::abs(weight);
                        metaballs.push_back({ bx, by, br, bridge_type });
                    }
                }
            }
        }
    }

    if (ball_visible) {
        auto get_acts = [&](int l) -> const std::vector<float>& {
            if (l == 0) return nn.a0;
            if (l == 1) return nn.a1;
            if (l == 2) return nn.a2;
            return nn.ao;
        };

        int vis_src = std::min(static_cast<int>(get_acts(ball.src_layer).size()), MAX_VIS_NEURONS);
        int vis_dst = std::min(static_cast<int>(get_acts(ball.dst_layer).size()), MAX_VIS_NEURONS);

        float x1 = layer_col_x_[ball.src_layer];
        float x2 = layer_col_x_[ball.dst_layer];
        float y1 = neuron_y_abs(ball.src_layer, ball.src_neuron % vis_src, vis_src);
        float y2 = neuron_y_abs(ball.dst_layer, ball.dst_neuron % vis_dst, vis_dst);

        for (int i = 0; i < 6; ++i) {
            float t = ball.x - i * 0.06f; // stagger distance
            if (t >= 0.0f && t <= 1.0f) {
                float bx = x1 + t * (x2 - x1);
                float by = y1 + t * (y2 - y1);
                float br = 1.3f * (1.0f - i * 0.14f); // size decay for tail
                if (br > 0.2f) {
                    metaballs.push_back({ bx, by, br, 3 });
                }
            }
        }
    }

    for (int y = net_top_ + 2; y < net_bottom_; ++y) {
        for (int x = 0; x < cols_; ++x) {
            float f_dim = 0.0f;
            float f_pos = 0.0f; // Excitatory
            float f_neg = 0.0f; // Inhibitory
            float f_ball = 0.0f;

            for (const auto& mb : metaballs) {
                float dx = x - mb.x;
                float dy = (y - mb.y) * 2.0f;
                float dist2 = dx * dx + dy * dy;
                float influence = (mb.r * mb.r) / (dist2 + 0.01f);
                
                if (mb.type == 3)      f_ball += influence;
                else if (mb.type == 1) f_pos  += influence;
                else if (mb.type == 2) f_neg  += influence;
                else                   f_dim  += influence;
            }

            float total_f = f_dim + f_pos + f_neg + f_ball;
            if (total_f > 0.35f) {
                int cp = CP_DIM;
                if (f_ball > 0.25f) {
                    cp = CP_BALL;
                } else if (f_pos > f_neg && f_pos > f_dim) {
                    cp = CP_ACTIVE;
                } else if (f_neg > f_pos && f_neg > f_dim) {
                    cp = CP_LOSS;
                }

                attron(COLOR_PAIR(cp));

                const char* ch = " ";
                if (total_f > 1.6f) {
                    ch = "█";
                } else if (total_f > 1.1f) {
                    ch = "▓";
                } else if (total_f > 0.7f) {
                    ch = "▒";
                } else if (total_f > 0.35f) {
                    ch = "░";
                }

                mvaddstr(y, x, ch);
                attroff(COLOR_PAIR(cp));
            }
        }
    }
}

void Visualizer::draw_info_panel(const std::string& prompt,
                                   const std::string& answer,
                                   float confidence,
                                   const std::vector<std::pair<float,std::string>>& top_preds,
                                   bool warn) {
    int y = info_top_;
    int maxw = cols_ - 4;

    attron(COLOR_PAIR(CP_DIM));
    std::string sep(cols_, '-');
    mvaddstr(y, 0, sep.c_str());
    attroff(COLOR_PAIR(CP_DIM));
    ++y;

    attron(COLOR_PAIR(CP_PROMPT) | A_BOLD);
    mvaddstr(y, 2, "Prompt: ");
    attroff(COLOR_PAIR(CP_PROMPT) | A_BOLD);
    attron(COLOR_PAIR(CP_DEFAULT));
    mvaddstr(y, 10, truncate(prompt, maxw - 10).c_str());
    attroff(COLOR_PAIR(CP_DEFAULT));
    ++y;

    attron(COLOR_PAIR(CP_ANSWER) | A_BOLD);
    mvaddstr(y, 2, "Answer: ");
    char conf_buf[64];
    snprintf(conf_buf, sizeof(conf_buf), "%s  (%.1f%% confidence)",
             answer.c_str(), confidence * 100.f);
    mvaddstr(y, 10, truncate(std::string(conf_buf), maxw - 10).c_str());
    attroff(COLOR_PAIR(CP_ANSWER) | A_BOLD);
    ++y;

    if (!top_preds.empty()) {
        attron(COLOR_PAIR(CP_DIM));
        mvaddstr(y, 2, "Top predictions: ");
        attroff(COLOR_PAIR(CP_DIM));
        ++y;
        int shown = std::min(static_cast<int>(top_preds.size()), 3);
        for (int i = 0; i < shown && y < bar_y_ - 1; ++i, ++y) {
            char buf[128];
            snprintf(buf, sizeof(buf), "  [%d] %-20s %.1f%%",
                     i + 1,
                     top_preds[i].second.substr(0, 20).c_str(),
                     top_preds[i].first * 100.f);
            attron(i == 0 ? COLOR_PAIR(CP_ANSWER) : COLOR_PAIR(CP_DIM));
            mvaddstr(y, 2, truncate(std::string(buf), maxw).c_str());
            attroff(i == 0 ? COLOR_PAIR(CP_ANSWER) : COLOR_PAIR(CP_DIM));
        }
    }

    if (warn && y < bar_y_ - 1) {
        attron(COLOR_PAIR(CP_WARN));
        mvaddstr(y, 2,
            "AI may not be accurate! Add to ~/.config/nvisual/dataset/dataset.jsonl and run nvisual --train to tune");
        attroff(COLOR_PAIR(CP_WARN));
    }
}

void Visualizer::draw_training_panel(int epoch, int total,
                                      float loss, float acc,
                                      int dataset_size,
                                      const std::string& msg) {
    int y = info_top_;
    int maxw = cols_ - 4;

    attron(COLOR_PAIR(CP_DIM));
    std::string sep(cols_, '=');
    mvaddstr(y, 0, sep.c_str());
    attroff(COLOR_PAIR(CP_DIM));
    ++y;

    attron(COLOR_PAIR(CP_LABEL) | A_BOLD);
    char ep_buf[64];
    snprintf(ep_buf, sizeof(ep_buf), "  Training: Epoch %d / %d", epoch, total);
    mvaddstr(y, 0, ep_buf);
    attroff(COLOR_PAIR(CP_LABEL) | A_BOLD);
    ++y;

    int bar_w = std::min(maxw - 4, 50);
    float pct = total > 0 ? static_cast<float>(epoch) / total : 0.f;
    int filled = static_cast<int>(pct * bar_w);
    std::string bar = "  [";
    bar += std::string(filled, '#');
    bar += std::string(bar_w - filled, '.');
    bar += "] ";
    char pct_buf[16];
    snprintf(pct_buf, sizeof(pct_buf), "%3.0f%%", pct * 100.f);
    bar += pct_buf;
    attron(COLOR_PAIR(CP_ANSWER) | A_BOLD);
    mvaddstr(y, 0, bar.c_str());
    attroff(COLOR_PAIR(CP_ANSWER) | A_BOLD);
    ++y;

    char stat_buf[128];
    snprintf(stat_buf, sizeof(stat_buf),
             "  Loss: %.4f   Accuracy: %.1f%%   Dataset: %d examples",
             loss, acc * 100.f, dataset_size);
    attron(COLOR_PAIR(CP_LOSS));
    mvaddstr(y, 0, stat_buf);
    attroff(COLOR_PAIR(CP_LOSS));
    ++y;

    if (!msg.empty() && y < bar_y_ - 1) {
        attron(COLOR_PAIR(CP_WARN));
        mvaddstr(y, 2, truncate(msg, maxw).c_str());
        attroff(COLOR_PAIR(CP_WARN));
    }
}


void Visualizer::draw_inference(const NeuralNet& nn,
                                  const std::string& prompt,
                                  const std::string& answer,
                                  float confidence,
                                  bool  warn,
                                  const BallPos& ball,
                                  int   active_layer) {
    erase();
    compute_layout();

    draw_title_bar("nvisual");

    draw_network_panel(nn, active_layer, ball, true);

    std::vector<std::pair<float,std::string>> top;
    for (int i = 0; i < static_cast<int>(nn.ao.size()); ++i)
        top.emplace_back(nn.ao[i], nn.labels[i]);
    std::sort(top.rbegin(), top.rend());

    draw_info_panel(prompt, answer, confidence, top, warn);

    char stat[256];
    snprintf(stat, sizeof(stat),
             "  %d params | Press Q/ESC to quit | --train to fine-tune",
             nn.param_count());
    draw_status_bar(stat);

    refresh();
}

void Visualizer::draw_training(const NeuralNet& nn,
                                 int epoch, int total_epochs,
                                 float loss, float accuracy,
                                 int dataset_size,
                                 const std::string& status_msg) {
    erase();
    compute_layout();

    draw_title_bar("  nvisual  ─  Training Mode  (CTRL+C to stop)  ");

    BallPos b; b.x = 0;
    draw_network_panel(nn, -1, b, false);

    draw_training_panel(epoch, total_epochs, loss, accuracy, dataset_size, status_msg);

    char stat[256];
    snprintf(stat, sizeof(stat),
             "  Training %d/%d epochs | %d params | %d examples",
             epoch, total_epochs, nn.param_count(), dataset_size);
    draw_status_bar(stat);

    refresh();
}

void Visualizer::draw_splash(const NeuralNet& nn, int dataset_size) {
    erase();
    compute_layout();

    draw_title_bar("  nvisual  ─  Initializing...  ");

    int cy = rows_ / 2;
    const char* logo[] = {
        "  ███╗   ██╗██╗   ██╗██╗███████╗██╗   ██╗ █████╗ ██╗     ",
        "  ████╗  ██║██║   ██║██║██╔════╝██║   ██║██╔══██╗██║     ",
        "  ██╔██╗ ██║██║   ██║██║███████╗██║   ██║███████║██║     ",
        "  ██║╚██╗██║╚██╗ ██╔╝██║╚════██║██║   ██║██╔══██║██║     ",
        "  ██║ ╚████║ ╚████╔╝ ██║███████║╚██████╔╝██║  ██║███████╗",
        "  ╚═╝  ╚═══╝  ╚═══╝  ╚═╝╚══════╝ ╚═════╝ ╚═╝  ╚═╝╚══════╝",
    };
    int n_logo = static_cast<int>(sizeof(logo) / sizeof(logo[0]));
    int start_y = cy - n_logo - 2;

    attron(COLOR_PAIR(CP_ACTIVE) | A_BOLD);
    for (int i = 0; i < n_logo; ++i) {
        int lx = std::max(0, (cols_ - static_cast<int>(strlen(logo[i]))) / 2);
        if (start_y + i >= 1 && start_y + i < rows_ - 1)
            mvaddstr(start_y + i, lx, logo[i]);
    }
    attroff(COLOR_PAIR(CP_ACTIVE) | A_BOLD);

    attron(COLOR_PAIR(CP_DEFAULT));
    char info[128];
    snprintf(info, sizeof(info),
             "CLI Neural Network Visualizer  •  %d params  •  %d training examples",
             nn.param_count(), dataset_size);
    int ix = std::max(0, (cols_ - static_cast<int>(strlen(info))) / 2);
    if (start_y + n_logo + 1 < rows_ - 1)
        mvaddstr(start_y + n_logo + 1, ix, info);
    attroff(COLOR_PAIR(CP_DEFAULT));

    attron(COLOR_PAIR(CP_WARN));
    const char* warn = "Small model — answers may be inaccurate. See ~/.config/nvisual/ to improve.";
    int wx = std::max(0, (cols_ - static_cast<int>(strlen(warn))) / 2);
    if (start_y + n_logo + 3 < rows_ - 1)
        mvaddstr(start_y + n_logo + 3, wx, warn);
    attroff(COLOR_PAIR(CP_WARN));

    draw_status_bar("  Loading model... Please wait.");
    refresh();
}

void Visualizer::animate_ball(const NeuralNet& nn,
                                const std::string& prompt,
                                const std::string& answer,
                                float confidence,
                                bool  warn,
                                int   src_layer,
                                int   dst_layer,
                                int   frames,
                                int   delay_ms) {
    int n_src = std::min(static_cast<int>(
        src_layer == 0 ? nn.a0.size() :
        src_layer == 1 ? nn.a1.size() :
        src_layer == 2 ? nn.a2.size() : nn.ao.size()), MAX_VIS_NEURONS);
    int n_dst = std::min(static_cast<int>(
        dst_layer == 0 ? nn.a0.size() :
        dst_layer == 1 ? nn.a1.size() :
        dst_layer == 2 ? nn.a2.size() : nn.ao.size()), MAX_VIS_NEURONS);

    auto get_acts = [&](int l) -> const std::vector<float>& {
        if (l == 0) return nn.a0;
        if (l == 1) return nn.a1;
        if (l == 2) return nn.a2;
        return nn.ao;
    };
    const auto& src_acts = get_acts(src_layer);
    const auto& dst_acts = get_acts(dst_layer);

    int best_src = 0;
    {
        float mx = -1e9f;
        for (int i = 0; i < std::min(static_cast<int>(src_acts.size()), n_src); ++i)
            if (src_acts[i] > mx) { mx = src_acts[i]; best_src = i; }
    }
    int best_dst = 0;
    {
        float mx = -1e9f;
        for (int i = 0; i < std::min(static_cast<int>(dst_acts.size()), n_dst); ++i)
            if (dst_acts[i] > mx) { mx = dst_acts[i]; best_dst = i; }
    }

    BallPos ball;
    ball.src_layer  = src_layer;
    ball.dst_layer  = dst_layer;
    ball.src_neuron = best_src;
    ball.dst_neuron = best_dst;

    for (int f = 0; f <= frames; ++f) {
        ball.x = static_cast<float>(f) / frames;
        ball.y = ball.x; // same interpolation for y
        draw_inference(nn, prompt, answer, confidence, warn, ball, src_layer);
        sleep_ms(delay_ms);

        int ch = get_key();
        if (is_quit_key(ch)) break;
    }
}