main.cpp

#include "nn.hpp"
#include "tokenizer.hpp"
#include "dataset.hpp"
#include "config.hpp"
#include "visualizer.hpp"

#include <iostream>
#include <string>
#include <vector>
#include <filesystem>
#include <fstream>
#include <cstdlib>
#include <chrono>

namespace fs = std::filesystem;

// Default dataset JSONL content
const std::string DEFAULT_DATASET_JSONL = R"({"prompt": "What is 2 + 2?", "answer": "4"}
{"prompt": "What is the capital of France?", "answer": "Paris"}
{"prompt": "What color is the sky?", "answer": "Blue"}
{"prompt": "Is water wet?", "answer": "Yes"}
{"prompt": "What is the speed of light?", "answer": "299,792,458 m/s"}
{"prompt": "Who invented the telephone?", "answer": "Alexander Graham Bell"}
{"prompt": "What is the largest planet?", "answer": "Jupiter"}
{"prompt": "How many days in a week?", "answer": "7"}
{"prompt": "What gas do plants absorb?", "answer": "Carbon dioxide"}
{"prompt": "What is the boiling point of water?", "answer": "100 degrees Celsius"}
{"prompt": "What is the capital of Japan?", "answer": "Tokyo"}
{"prompt": "What is the capital of Germany?", "answer": "Berlin"}
{"prompt": "What color is grass?", "answer": "Green"}
{"prompt": "How many continents are there?", "answer": "7"}
{"prompt": "Who wrote Hamlet?", "answer": "William Shakespeare"}
{"prompt": "What is the square root of 9?", "answer": "3"}
{"prompt": "What is the chemical formula for water?", "answer": "H2O"}
{"prompt": "Who was the first president of the United States?", "answer": "George Washington"}
{"prompt": "Which country is home to the kangaroo?", "answer": "Australia"}
{"prompt": "How many legs does a spider have?", "answer": "8"}
)";

std::string get_home_dir() {
    const char* h = std::getenv("HOME");
    return h ? std::string(h) : std::string("/root");
}

void ensure_directory_exists(const std::string& path) {
    try {
        fs::create_directories(path);
    } catch (...) {}
}

void setup_default_files(const std::string& config_dir) {
    ensure_directory_exists(config_dir);
    ensure_directory_exists(config_dir + "/dataset");

    std::string config_path = config_dir + "/config.conf";
    if (!fs::exists(config_path)) {
        Config::write_default(config_path);
    }

    std::string dataset_path = config_dir + "/dataset/dataset.jsonl";
    if (!fs::exists(dataset_path)) {
        std::ofstream f(dataset_path);
        if (f) {
            f << DEFAULT_DATASET_JSONL;
        }
    }
}

// Train
bool run_training(NeuralNet& nn, Tokenizer& tok, const std::string& dataset_path,
                 const std::string& model_path, const std::string& tokenizer_path,
                 int epochs, float lr, Visualizer* vis) {
    Dataset ds;
    if (ds.load(dataset_path) == 0) {
        if (vis) {
            vis->draw_training(nn, 0, epochs, 0.0f, 0.0f, 0, "Error: Dataset is empty or not found!");
            vis->sleep_ms(3000);
        } else {
            std::cerr << "Error: Dataset is empty or not found at " << dataset_path << std::endl;
        }
        return false;
    }

    // Vocab
    tok.build_vocab(ds.all_prompts());
    nn.init(ds.labels());

    if (vis) {
        vis->draw_training(nn, 0, epochs, 0.0f, 0.0f, ds.size(), "Starting training...");
        vis->sleep_ms(500);
    } else {
        std::cout << "Starting training on " << ds.size() << " examples..." << std::endl;
    }

    // Prep Train input
    struct TrainItem {
        std::vector<int> tokens;
        int label_idx;
    };
    std::vector<TrainItem> train_data;
    train_data.reserve(ds.size());
    for (const auto& ex : ds.examples()) {
        train_data.push_back({
            tok.tokenize(ex.prompt),
            ds.label_index(ex.answer)
        });
    }

    // Train (Main)
    for (int epoch = 1; epoch <= epochs; ++epoch) {
        float total_loss = 0.0f;
        int correct = 0;

        for (const auto& item : train_data) {
            float loss = nn.train_step(item.tokens, item.label_idx, lr);
            total_loss += loss;

            int pred = nn.predict(item.tokens);
            if (pred == item.label_idx) {
                correct++;
            }
        }

        float avg_loss = total_loss / train_data.size();
        float accuracy = static_cast<float>(correct) / train_data.size();

        if (vis) {
            vis->draw_training(nn, epoch, epochs, avg_loss, accuracy, ds.size(), "Optimizing parameters...");
            vis->sleep_ms(15); // visual smooth updates
            
            // Check for exit
            int ch = vis->get_key();
            if (ch == 3 || ch == 'q') {
                break;
            }
        } else {
            if (epoch % 10 == 0 || epoch == epochs) {
                std::cout << "Epoch " << epoch << "/" << epochs
                          << " - Loss: " << avg_loss
                          << " - Accuracy: " << (accuracy * 100.0f) << "%" << std::endl;
            }
        }
    }

    nn.save(model_path);
    tok.save(tokenizer_path);

    if (vis) {
        vis->draw_training(nn, epochs, epochs, 0.0f, 1.0f, ds.size(), "Training complete! Model saved.");
        vis->sleep_ms(1500);
    } else {
        std::cout << "Training complete! Saved model to " << model_path << std::endl;
    }

    return true;
}

int main(int argc, char* argv[]) {
    bool train_only = false;
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--train" || arg == "-t") {
            train_only = true;
        } else if (arg == "--help" || arg == "-h") {
            std::cout << "nvisual - CLI Neural Network Visualizer\n\n"
                      << "Usage:\n"
                      << "  nvisual            Run the CLI neural network visualizer\n"
                      << "  nvisual --train    Train/Fine-tune the neural network on the local dataset\n\n"
                      << "Configuration:\n"
                      << "  Config file:    ~/.config/nvisual/config.conf\n"
                      << "  Dataset:        ~/.config/nvisual/dataset/dataset.jsonl\n"
                      << "  Model weights:  ~/.config/nvisual/model.bin\n";
            return 0;
        }
    }

    std::string home_dir = get_home_dir();
    std::string config_dir = home_dir + "/.config/nvisual";
    setup_default_files(config_dir);

    std::string config_path = config_dir + "/config.conf";
    std::string dataset_path = config_dir + "/dataset/dataset.jsonl";
    std::string model_path = config_dir + "/model.bin";
    std::string tokenizer_path = config_dir + "/tokenizer.txt";

    Config cfg;
    cfg.load(config_path);

    NeuralNet nn;
    Tokenizer tok;

    if (train_only) {
        Visualizer vis;
        vis.init();
        run_training(nn, tok, dataset_path, model_path, tokenizer_path,
                     cfg.get().train_epochs, cfg.get().learning_rate, &vis);
        vis.cleanup();
        return 0;
    }

    if (!fs::exists(model_path) || !fs::exists(tokenizer_path)) {
        Visualizer vis;
        vis.init();
        vis.draw_splash(nn, 0);
        vis.sleep_ms(1500);

        run_training(nn, tok, dataset_path, model_path, tokenizer_path,
                     cfg.get().train_epochs, cfg.get().learning_rate, &vis);
        vis.cleanup();
    }

    // Load model
    if (!nn.load(model_path) || !tok.load(tokenizer_path)) {
        std::cerr << "Failed to load model or tokenizer. Please run: nvisual --train" << std::endl;
        return 1;
    }

    Visualizer vis;
    vis.init();

    std::vector<std::string> prompts = cfg.get().prompts;
    if (prompts.empty()) {
        prompts = {
            "What is 2 + 2?",
            "What is the capital of France?",
            "What color is the sky?",
            "Is water wet?",
            "What is the speed of light?"
        };
    }

    int prompt_idx = 0;
    bool running = true;

    while (running) {
        std::string current_prompt = prompts[prompt_idx];
        auto tokens = tok.tokenize(current_prompt);
        
        auto probs = nn.forward(tokens);

        int pred_idx = nn.predict(tokens);
        std::string pred_label = nn.labels[pred_idx];
        float confidence = probs[pred_idx];

        vis.animate_ball(nn, current_prompt, "Processing...", confidence, cfg.get().warn_accuracy, 0, 1, 15, 30);
        vis.animate_ball(nn, current_prompt, "Processing...", confidence, cfg.get().warn_accuracy, 1, 2, 15, 30);
        vis.animate_ball(nn, current_prompt, pred_label, confidence, cfg.get().warn_accuracy, 2, 3, 15, 30);

        BallPos static_ball;
        static_ball.x = 1.0f;
        static_ball.src_layer = 2;
        static_ball.dst_layer = 3;
        static_ball.src_neuron = 0;
        static_ball.dst_neuron = pred_idx;

        vis.draw_inference(nn, current_prompt, pred_label, confidence, cfg.get().warn_accuracy, static_ball, -1);

        int delay = cfg.get().cycle_delay_ms;
        auto start_wait = std::chrono::steady_clock::now();
        while (std::chrono::duration_cast<std::chrono::milliseconds>(
                   std::chrono::steady_clock::now() - start_wait).count() < delay) {
            int ch = vis.get_key();
            if (vis.is_quit_key(ch)) {
                running = false;
                break;
            } else if (ch == ' ' || ch == '\n' || ch == KEY_RIGHT) {
                break;
            }
            vis.sleep_ms(20);
        }

        prompt_idx = (prompt_idx + 1) % prompts.size();
    }

    vis.cleanup();
    return 0;
}