#include "common.h"
#include <assert.h>
#include <errno.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define seq_len 5

typedef struct {
    char ch;
    uint16_t idx;
} TokMatchOpt;

typedef struct {
    TokMatchOpt *opts;
    uint16_t opts_cap;
    uint16_t opts_size;
    uint16_t id;
} TokMatchNode;

typedef struct {
    TokMatchNode *nodes;
    size_t nodes_cap;
    size_t nodes_size;
} TokMatcher;

void tokmatcher_init(TokMatcher *m)
{
    *m = (TokMatcher) {
        .nodes = NULL,
        .nodes_cap = 0,
        .nodes_size = 0,
    };

    TokMatchNode head = { NULL, 0, 0, 0 };

    array_push(
        (void **)&m->nodes, &m->nodes_cap, &m->nodes_size, &head, sizeof(head));
}

void tokmatcher_deinit(TokMatcher *m)
{
    for (size_t i = 0; i < m->nodes_size; ++i)
        if (m->nodes[i].opts)
            free(m->nodes[i].opts);
    if (m->nodes)
        free(m->nodes);
}

static TokMatchNode *tokmatcher_make_path(
    TokMatcher *m, size_t node_idx, const char *key)
{
    TokMatchNode *node = &m->nodes[node_idx];
    if (!*key)
        return node;

    for (size_t i = 0; i < node->opts_size; ++i)
        if (node->opts[i].ch == *key)
            return tokmatcher_make_path(m, node->opts[i].idx, key + 1);

    size_t idx = m->nodes_size;

    TokMatchNode new_node = { NULL, 0, 0, 0 };
    array_push((void **)&m->nodes,
        &m->nodes_cap,
        &m->nodes_size,
        &new_node,
        sizeof(new_node));

    // WARNING: m->nodes might be relocated here,
    // hence we need to re-index node.
    node = &m->nodes[node_idx];

    assert(idx <= UINT16_MAX);
    TokMatchOpt new_opt = {
        .ch = *key,
        .idx = (uint16_t)idx,
    };

    size_t opts_cap = node->opts_cap;
    size_t opts_size = node->opts_size;
    array_push(
        (void **)&node->opts, &opts_cap, &opts_size, &new_opt, sizeof(new_opt));

    assert(opts_cap <= UINT16_MAX && opts_size <= UINT16_MAX);
    node->opts_cap = opts_cap;
    node->opts_size = opts_size;

    return tokmatcher_make_path(m, idx, key + 1);
}

void tokmatcher_add(TokMatcher *m, const char *key, uint16_t id)
{
    TokMatchNode *node = tokmatcher_make_path(m, 0, key);
    node->id = id;
}

typedef struct {
    uint32_t len;
    uint16_t id;
} TokRecMatch;

static TokRecMatch tokmatcher_match_rec(TokMatcher *m,
    TokMatchNode *node,
    TokRecMatch cand,
    size_t len,
    const char *text,
    size_t text_len)
{
    if (node->id != 0) {
        cand = (TokRecMatch) { len, node->id };
    }

    if (!*text || text_len == 0)
        return cand;

    for (size_t i = 0; i < node->opts_size; ++i) {
        if (node->opts[i].ch != *text)
            continue;

        TokMatchNode *next = &m->nodes[node->opts[i].idx];

        return tokmatcher_match_rec(
            m, next, cand, len + 1, text + 1, text_len - 1);
    }

    return cand;
}

uint16_t tokmatcher_match(
    TokMatcher *m, const char *text, size_t text_len, size_t *chars_matched)
{
    TokRecMatch cand = { .len = 1, .id = 0 };

    TokRecMatch match
        = tokmatcher_match_rec(m, &m->nodes[0], cand, 0, text, text_len);

    *chars_matched += match.len;
    return match.id;
}

static void tokmatcher_dump_rec(
    TokMatcher *m, TokMatchNode *node, FILE *file, int depth)
{
    fprintf(file, " -> %d", node->id);

    for (size_t i = 0; i < node->opts_size; ++i) {
        fprintf(file, "\n");
        for (int j = 0; j < depth + 1; ++j)
            fprintf(file, "       | ");

        fprintf(file, "'%c'", node->opts[i].ch);

        tokmatcher_dump_rec(m, &m->nodes[node->opts[i].idx], file, depth + 1);
    }
}

void tokmatcher_dump(TokMatcher *m, FILE *file)
{
    printf(" * ");
    tokmatcher_dump_rec(m, &m->nodes[0], file, 0);
}

typedef enum {
    SpecialTok_Unk,
    SpecialTok_Begin,
    SpecialTok_End,
    SpecialTokCount,
} SpecialTok;

typedef struct {
    char *toks_strs;
    size_t toks_strs_cap;
    size_t toks_strs_size;

    size_t *toks;
    size_t toks_cap;
    size_t toks_size;

    TokMatcher tok_matcher;
    uint16_t special_toks[SpecialTokCount];

    int next_word_ch;

    FILE *file;
    char *word;
    size_t word_cap;
    size_t word_len;
    size_t chars_matched;
} Tokenizer;

static uint16_t t_add_token_str(Tokenizer *t, const char *str)
{
    size_t tok_str_idx = t->toks_strs_size;

    array_push_n((void **)&t->toks_strs,
        &t->toks_strs_cap,
        &t->toks_strs_size,
        (void *)str,
        1,
        strlen(str) + 1);

    array_push((void **)&t->toks,
        &t->toks_cap,
        &t->toks_size,
        &tok_str_idx,
        sizeof(tok_str_idx));

    assert(tok_str_idx <= UINT16_MAX);
    uint16_t id = tok_str_idx;
    return id;
}

void tokenizer_init(Tokenizer *t)
{
    *t = (Tokenizer) {
        .toks_strs = NULL,
        .toks_strs_cap = 0,
        .toks_strs_size = 0,

        .toks = NULL,
        .toks_cap = 0,
        .toks_size = 0,

        .tok_matcher = { 0 },
        .special_toks = { 0 },

        .next_word_ch = 0,

        .file = NULL,
        .word = NULL,
        .word_cap = 0,
        .word_len = 0,
        .chars_matched = 0,
    };

    tokmatcher_init(&t->tok_matcher);

    t->special_toks[SpecialTok_Unk] = t_add_token_str(t, "<unk>");
    t->special_toks[SpecialTok_Begin] = t_add_token_str(t, "<s>");
    t->special_toks[SpecialTok_End] = t_add_token_str(t, "</s>");
}

void tokenizer_deinit(Tokenizer *t)
{
    if (t->toks_strs)
        free(t->toks_strs);
    if (t->toks)
        free(t->toks);
    tokmatcher_deinit(&t->tok_matcher);
    if (t->file)
        fclose(t->file);
    if (t->word)
        free(t->word);
}

int tokenizer_load_tokens_file(Tokenizer *t, const char *filename)
{
    FILE *file = fopen(filename, "r");
    if (!file) {
        fprintf(stderr,
            "error: could not open file '%s': %s\n",
            filename,
            strerror(errno));
        return 1;
    }

    char chars_line[34] = "";
    fgets(chars_line, 34, file);
    size_t chars_line_len = strlen(chars_line);

    // This is just a sanity check that 1) we have read a
    // correct file, and 2) that the chars are what we expect.
    assert(strcmp(chars_line, " !',.?abcdefghijklmnopqrstuvwxyz"));
    assert(chars_line_len == 33);

    for (size_t i = 0; i < chars_line_len; ++i) {
        char str[2] = { chars_line[i], '\0' };

        uint16_t id = t_add_token_str(t, str);
        tokmatcher_add(&t->tok_matcher, str, id);
    }

    while (!feof(file)) {
        char line[41] = "";
        fgets(line, 41, file);
        size_t line_len = strlen(line);
        if (line_len == 0)
            break;
        assert(line[line_len - 1] == '\n');
        line[line_len - 1] = '\0';

        uint16_t id = t_add_token_str(t, line);
        tokmatcher_add(&t->tok_matcher, line, id);
    }

    fclose(file);

    return 0;
}

int tokenizer_open(Tokenizer *t, const char *filename)
{
    FILE *file = fopen(filename, "r");
    if (!file) {
        fprintf(stderr,
            "error: could not open file '%s': %s\n",
            filename,
            strerror(errno));
        return 1;
    }

    t->file = file;

    t->word_cap = 64;
    t->word = malloc(t->word_cap);
    t->word_len = 0;
    t->chars_matched = 0;

    return 0;
}

int32_t tokenizer_next(Tokenizer *t)
{
    if (t->chars_matched >= t->word_len) {
        size_t len = next_word(t->word, t->word_cap, t->file, &t->next_word_ch);

        if (len == 0)
            return -1;
        if (t->word[0] == '\n')
            return -2;

        t->word[len] = '\0';
        t->word_len = len;
        t->chars_matched = 0;
    }

    uint16_t id = tokmatcher_match(&t->tok_matcher,
        t->word + t->chars_matched,
        t->word_len - t->chars_matched,
        &t->chars_matched);

    return id;
}

typedef struct {
    uint16_t tok;
    uint32_t node_idx;
} SeqIdxMapOpt;

int seqidxmapopt_cmp(const SeqIdxMapOpt *a, const SeqIdxMapOpt *b)
{
    if (a->tok != b->tok)
        return (int)a->tok - (int)b->tok;
    if (a->node_idx != b->node_idx)
        return a->node_idx < b->node_idx ? 1 : -1;
    return 0;
}

typedef struct {
    uint32_t val;
    uint32_t opts_cap_exp : 8, opts_size : 24;
    SeqIdxMapOpt opts[];
} SeqIdxMapNode;

typedef struct {
    SeqIdxMapNode **nodes;
    size_t nodes_size;
    size_t nodes_cap;
} SeqIdxMap;

void seqidxmap_init(SeqIdxMap *m)
{
    *m = (SeqIdxMap) {
        .nodes = NULL,
        .nodes_cap = 0,
        .nodes_size = 0,
    };

    SeqIdxMapNode *head
        = malloc(sizeof(SeqIdxMapNode) + sizeof(SeqIdxMapOpt) * 1);
    *head = (SeqIdxMapNode) { .val = 0, .opts_cap_exp = 0, .opts_size = 0 };
    array_push((void **)&m->nodes,
        &m->nodes_cap,
        &m->nodes_size,
        &head,
        sizeof(SeqIdxMapNode *));
}

void seqidxmap_deinit(SeqIdxMap *m)
{
    for (size_t i = 0; i < m->nodes_size; ++i)
        free(m->nodes[i]);
    free(m->nodes);
}

uint32_t *seqidxmap_get(SeqIdxMap *m, uint16_t seq[seq_len])
{
    size_t node_idx = 0;

    for (size_t seq_idx = 0; seq_idx < seq_len; ++seq_idx) {
        assert(node_idx < m->nodes_size);
        SeqIdxMapNode *node = m->nodes[node_idx];

        size_t opt_idx;
        SeqIdxMapOpt *opt = binary_search(&node->opts,
            node->opts_size,
            &seq[seq_idx],
            sizeof(uint16_t),
            &opt_idx,
            (CmpFn)seqidxmapopt_cmp);

        if (opt) {
            node_idx = opt->node_idx;
            continue;
        }

        if (node->opts_size >= (1 << node->opts_cap_exp)) {
            assert(node->opts_cap_exp >= 0 && node->opts_cap_exp <= 255);
            node->opts_cap_exp += 1;

            SeqIdxMapNode *new_ptr = realloc(m->nodes[node_idx],
                sizeof(SeqIdxMapNode)
                    + sizeof(SeqIdxMapOpt) * (1ul << node->opts_cap_exp));
            assert(new_ptr);
            m->nodes[node_idx] = new_ptr;
        }

        size_t new_node_idx = m->nodes_size;
        SeqIdxMapNode *new_node
            = malloc(sizeof(SeqIdxMapNode) + sizeof(SeqIdxMapOpt) * 1);
        *new_node
            = (SeqIdxMapNode) { .val = 0, .opts_cap_exp = 0, .opts_size = 0 };
        array_push((void **)&m->nodes,
            &m->nodes_cap,
            &m->nodes_size,
            &new_node,
            sizeof(SeqIdxMapNode *));

        assert(node_idx < m->nodes_size);
        node = m->nodes[node_idx];

        assert(node->opts_size < (1 << 24) - 1);
        node->opts[node->opts_size] = (SeqIdxMapOpt) {
            .node_idx = new_node_idx,
            .tok = seq[seq_idx],
        };
        node->opts_size += 1;

        node_idx = new_node_idx;
    }

    return &m->nodes[node_idx]->val;
}

int seq_cmp(const uint16_t a[seq_len], const uint16_t b[seq_len])
{
    for (size_t i = 0; i < seq_len; ++i)
        if (a[i] != b[i])
            return (int)a[i] - (int)b[i];
    return 0;
}

typedef struct {
    uint32_t seq_idx;
    uint16_t id;
    uint16_t count;
} SeqPred;

int seqpred_cmp(const SeqPred *a, const SeqPred *b)
{
    if (a->seq_idx != b->seq_idx)
        return a->seq_idx < b->seq_idx ? 1 : -1;
    if (a->id != b->id)
        return (int)a->id - (int)b->id;
    return 0;
}

typedef uint64_t SeqPredKey;

SeqPredKey seqpredkey_from_seqpred(const SeqPred *pred)
{
    SeqPredKey key = (uint64_t)pred->seq_idx | (uint64_t)pred->id << 32ul;
    return key;
}

int seqpred_key_cmp(const SeqPredKey *a, const SeqPredKey *b)
{
    if (*a != *b)
        return *a < *b ? 1 : -1;
    return 0;
}

typedef struct {
    uint16_t *seqs;
    size_t seqs_cap;
    size_t seqs_size;

    SeqIdxMap seqidxmap;

    SeqPred *preds;
    size_t preds_cap;
    size_t preds_size;

    RbTree *pred_map;
} Model;

void model_init(Model *m)
{
    *m = (Model) {
        .seqs = NULL,
        .seqs_cap = 0,
        .seqs_size = 0,

        .seqidxmap = { 0 },

        .preds = NULL,
        .preds_cap = 0,
        .preds_size = 0,

        .pred_map = rbtree_new(
            sizeof(SeqPredKey), sizeof(SeqPred *), (CmpFn)seqpred_key_cmp),
    };

    seqidxmap_init(&m->seqidxmap);
}

void model_deinit(Model *m)
{
    if (m->seqs)
        free(m->seqs);
    seqidxmap_deinit(&m->seqidxmap);

    if (m->preds)
        free(m->preds);

    rbtree_free(m->pred_map);
}

size_t model_get_seq_or_insert(Model *m, uint16_t seq[seq_len])
{
    uint32_t *idx = seqidxmap_get(&m->seqidxmap, seq);

    if (*idx)
        return *idx;

    assert(m->seqs_size <= UINT32_MAX);
    uint32_t new_idx = m->seqs_size;

    array_push((void **)&m->seqs,
        &m->seqs_cap,
        &m->seqs_size,
        seq,
        sizeof(uint16_t) * seq_len);

    *idx = new_idx;

    return new_idx;
}

uint32_t *model_get_seq(Model *m, uint16_t seq[seq_len])
{
    return seqidxmap_get(&m->seqidxmap, seq);
}

SeqPred *model_find_seqpred(Model *m, SeqPred *pred)
{
    SeqPredKey key = (uint64_t)pred->seq_idx | (uint64_t)pred->id << 32ul;
    return rbtree_get(m->pred_map, &key);
}

SeqPred *model_insert_seqpred(Model *m, SeqPred *pred)
{
    array_push(
        (void **)&m->preds, &m->preds_cap, &m->preds_size, pred, sizeof(*pred));

    SeqPred *ptr = &m->preds[m->preds_size - 1];

    SeqPredKey key = seqpredkey_from_seqpred(pred);
    rbtree_insert(m->pred_map, &key, &ptr);

    return ptr;
}

int main(int argc, char *argv[])
{
    test_common();
    if (argc < 2) {
        fprintf(stderr, "error: invalid arguments\n");
        fprintf(stderr, "usage: ./train_model <training text> <tokens>\n");
        return 1;
    }

    const char *training_text_filename = argv[1];
    const char *tokens_filename = argv[2];

    Tokenizer tokenizer;
    tokenizer_init(&tokenizer);

    printf("reading tokens...\n");
    int result = tokenizer_load_tokens_file(&tokenizer, tokens_filename);
    if (result != 0)
        return 1;

    result = tokenizer_open(&tokenizer, training_text_filename);
    if (result != 0)
        return 1;

    printf("training...\n");

    Model model;
    model_init(&model);

    const uint16_t begin_tok = tokenizer.special_toks[SpecialTok_Begin];
    const uint16_t end_tok = tokenizer.special_toks[SpecialTok_End];

    uint16_t current_seq[seq_len] = {
        begin_tok,
        begin_tok,
        begin_tok,
        begin_tok,
        begin_tok,
    };

    int toks_counter = 0;
    int iters = 0;

    int32_t id;
    while ((id = tokenizer_next(&tokenizer)) != -1) {
        id = id == -2 ? end_tok : id;

        size_t seq_idx = model_get_seq_or_insert(&model, current_seq);

        SeqPred cand_seq = { .seq_idx = seq_idx, .id = id, .count = 0 };

        SeqPred *pred = model_find_seqpred(&model, &cand_seq);

        if (!pred) {
            pred = model_insert_seqpred(&model, &cand_seq);
        }

        pred->count += 1;

        if (id == end_tok) {

            for (size_t i = 0; i < seq_len; ++i)
                current_seq[i] = begin_tok;
        } else {
            for (size_t i = 0; i < seq_len - 1; ++i)
                current_seq[i] = current_seq[i + 1];
            current_seq[seq_len - 1] = id;
        }

        toks_counter += 1;
        iters += 1;
        if (iters >= 1000000) {
            iters = 0;

            printf(
                "  toks_counter: %d, seqs: %ld, seq nodes: %ld, preds: %ld\n",
                toks_counter,
                model.seqs_size,
                model.seqidxmap.nodes_size,
                model.preds_size);
            break;
        }
    }

    printf("done!\n\n");

    for (size_t j = 0; j < model.preds_size; j += 100) {
        for (size_t i = j; i < j + 100; ++i) {
            SeqPred *pred = &model.preds[i];
            uint16_t *seq = &model.seqs[pred->seq_idx];
            printf(" '");
            for (size_t i = 0; i < seq_len; ++i) {
                printf("%s", &tokenizer.toks_strs[seq[i]]);
            }
            printf("' -> %d x '%s'\n",
                pred->count,
                &tokenizer.toks_strs[pred->id]);
        }
        getc(stdin);
    }

    //
    // printf("> ");
    //
    // for (size_t i = 0; i < seq_len; ++i)
    //     current_seq[i] = begin_tok;
    //
    // {
    //     char word[64] = "";
    //     size_t word_cap = sizeof(word);
    //     size_t word_len = 0;
    //     size_t chars_matched = 0;
    //
    //     int ch = 0;
    //     while (true) {
    //         if (chars_matched >= word_len) {
    //             size_t len = next_word(word, word_cap, stdin, &ch);
    //
    //             if (len == 0)
    //                 break;
    //             if (word[0] == '\n')
    //                 break;
    //
    //             word[len] = '\0';
    //             word_len = len;
    //             chars_matched = 0;
    //         }
    //
    //         uint16_t id = tokmatcher_match(&tokenizer.tok_matcher,
    //             word + chars_matched,
    //             word_len - chars_matched,
    //             &chars_matched);
    //
    //         for (size_t i = 0; i < seq_len - 1; ++i)
    //             current_seq[i] = current_seq[i + 1];
    //         current_seq[seq_len - 1] = id;
    //     }
    // }
    //
    // printf("generating prompt...\n");
    //
    // printf("\n");
    //
    // for (size_t i = 0; i < seq_len; ++i) {
    //     printf("%s", &tokenizer.toks_strs[current_seq[i]]);
    // }
    //
    // for (int iters = 0; iters < 1000; ++iters) {
    //
    //     uint32_t *seq_idx = model_get_seq(&model, current_seq);
    //     if (!seq_idx) {
    //         printf("no sequence\n");
    //         break;
    //     }
    //
    //     SeqPred *pred = NULL;
    //
    //     for (size_t i = 0; i < model.preds_size; ++i) {
    //         if (model.preds[i].seq_idx != *seq_idx)
    //             continue;
    //         if (!pred || model.preds[i].count > pred->count)
    //             pred = &model.preds[i];
    //     }
    //
    //     if (!pred) {
    //         printf("no prediction\n");
    //         break;
    //     }
    //
    //     printf("%s", &tokenizer.toks_strs[pred->id]);
    //
    //     for (size_t i = 0; i < seq_len - 1; ++i)
    //         current_seq[i] = current_seq[i + 1];
    //     current_seq[seq_len - 1] = pred->id;
    // }

    tokenizer_deinit(&tokenizer);
    model_deinit(&model);
}