// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "train/dataloader.h"

#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

#include <cstdio>
#include <cstdlib>
#include <cstring>

namespace forge::train {

namespace {
constexpr int32_t kMagic = 20240520;
constexpr size_t kHeaderInts = 256;

[[noreturn]] void die(const std::string& msg) {
    std::fprintf(stderr, "forge/data: %s\n", msg.c_str());
    std::abort();
}
} // namespace

DataLoader::DataLoader(const std::string& bin_path, int64_t context_length, uint64_t seed)
    : context_(context_length), rng_(seed) {
    const int fd = ::open(bin_path.c_str(), O_RDONLY);
    if (fd < 0) die("cannot open " + bin_path);
    struct stat st{};
    if (fstat(fd, &st) != 0) die("fstat failed on " + bin_path);
    map_len_ = size_t(st.st_size);
    map_ = mmap(nullptr, map_len_, PROT_READ, MAP_PRIVATE, fd, 0);
    ::close(fd);
    if (map_ == MAP_FAILED) die("mmap failed on " + bin_path);
    madvise(map_, map_len_, MADV_RANDOM);

    const int32_t* header = static_cast<const int32_t*>(map_);
    if (map_len_ >= kHeaderInts * 4 && header[0] == kMagic && header[1] == 1) {
        num_tokens_ = header[2];
        tokens_ = reinterpret_cast<const uint16_t*>(header + kHeaderInts);
        if (size_t(num_tokens_) * 2 + kHeaderInts * 4 > map_len_)
            die("header token count exceeds file size: " + bin_path);
    } else {
        // headerless: the whole file is uint16 tokens
        num_tokens_ = int64_t(map_len_ / 2);
        tokens_ = static_cast<const uint16_t*>(map_);
    }
    if (num_tokens_ < context_ + 1)
        die("dataset smaller than one context window: " + bin_path);
}

DataLoader::~DataLoader() {
    if (map_ && map_ != MAP_FAILED) munmap(map_, map_len_);
}

void DataLoader::fill(int64_t start, int64_t T, int32_t* ids_row, int32_t* tgt_row) const {
    for (int64_t t = 0; t < T; ++t) {
        ids_row[t] = int32_t(tokens_[start + t]);
        tgt_row[t] = int32_t(tokens_[start + t + 1]);
    }
}

void DataLoader::next_batch(Tensor& ids, Tensor& targets) {
    const int64_t B = ids.size(0), T = ids.size(1);
    std::uniform_int_distribution<int64_t> dist(0, num_tokens_ - T - 1);
    for (int64_t b = 0; b < B; ++b) {
        fill(dist(rng_), T, ids.data<int32_t>() + b * T, targets.data<int32_t>() + b * T);
    }
}

void DataLoader::seq_batch(int64_t index, Tensor& ids, Tensor& targets) const {
    const int64_t B = ids.size(0), T = ids.size(1);
    const int64_t span = num_tokens_ - T - 1;
    for (int64_t b = 0; b < B; ++b) {
        const int64_t start = ((index * B + b) * T) % span;
        fill(start, T, ids.data<int32_t>() + b * T, targets.data<int32_t>() + b * T);
    }
}

} // namespace forge::train