#include "synapse.h"

#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#ifdef _WIN32
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#endif

#include "ggml.h"
#include "gguf.h"

#define CACHE_SLOTS 2048   // vettori da n_embd in cache, per tipo (circa 40 MB ciascuna)

typedef struct {
    const uint8_t *data;   // inizio dei dati del tensore nel file mappato
    enum ggml_type type;
    size_t row_bytes;
} Mat;

typedef struct { int64_t key; float *v; } Slot;

struct Synapses {
    int n_layer, n_ff, n_embd;
    const uint8_t *map;
    size_t map_size;
#ifdef _WIN32
    HANDLE file, mapping;
#endif
    Mat   *down, *gate;    // [n_layer]
    float *norm;           // [n_layer][n_embd] post_attention_norm
    Slot  *cols, *rows;    // cache: colonne di down (uscite), righe di gate (ingressi)
};

static const uint8_t *map_file(Synapses *s, const char *path) {
#ifdef _WIN32
    s->file = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (s->file == INVALID_HANDLE_VALUE) return NULL;
    LARGE_INTEGER sz;
    GetFileSizeEx(s->file, &sz);
    s->map_size = (size_t)sz.QuadPart;
    s->mapping = CreateFileMappingA(s->file, NULL, PAGE_READONLY, 0, 0, NULL);
    if (!s->mapping) return NULL;
    return MapViewOfFile(s->mapping, FILE_MAP_READ, 0, 0, 0);
#else
    int fd = open(path, O_RDONLY);
    if (fd < 0) return NULL;
    struct stat st;
    fstat(fd, &st);
    s->map_size = (size_t)st.st_size;
    void *p = mmap(NULL, s->map_size, PROT_READ, MAP_SHARED, fd, 0);
    close(fd);
    return p == MAP_FAILED ? NULL : p;
#endif
}

static bool find_mat(struct gguf_context *g, const Synapses *s, const char *name, Mat *m, size_t row_elems) {
    int64_t id = gguf_find_tensor(g, name);
    if (id < 0) return false;
    m->type = gguf_get_tensor_type(g, id);
    m->data = s->map + gguf_get_data_offset(g) + gguf_get_tensor_offset(g, id);
    m->row_bytes = ggml_row_size(m->type, (int64_t)row_elems);
    return ggml_get_type_traits(m->type)->to_float != NULL || m->type == GGML_TYPE_F32;
}

static void row_to_float(const Mat *m, size_t r, float *out, int n) {
    const uint8_t *p = m->data + r * m->row_bytes;
    if (m->type == GGML_TYPE_F32) memcpy(out, p, sizeof(float) * (size_t)n);
    else ggml_get_type_traits(m->type)->to_float(p, out, n);
}

Synapses *syn_open(const char *model_path, int n_layer, int n_ff, int n_embd) {
    Synapses *s = calloc(1, sizeof *s);
    s->n_layer = n_layer; s->n_ff = n_ff; s->n_embd = n_embd;
    s->map = map_file(s, model_path);
    struct gguf_init_params gp = {.no_alloc = true, .ctx = NULL};
    struct gguf_context *g = s->map ? gguf_init_from_file(model_path, gp) : NULL;
    if (!g) { fprintf(stderr, "sinapsi: impossibile leggere %s\n", model_path); syn_close(s); return NULL; }

    s->down = calloc((size_t)n_layer, sizeof(Mat));
    s->gate = calloc((size_t)n_layer, sizeof(Mat));
    s->norm = calloc((size_t)n_layer * n_embd, sizeof(float));
    char name[96];
    bool ok = true;
    for (int L = 0; L < n_layer && ok; L++) {
        snprintf(name, sizeof name, "blk.%d.ffn_down.weight", L);
        ok &= find_mat(g, s, name, &s->down[L], (size_t)n_ff);     // riga = una dimensione del residuo
        snprintf(name, sizeof name, "blk.%d.ffn_gate.weight", L);
        ok &= find_mat(g, s, name, &s->gate[L], (size_t)n_embd);   // riga = un neurone
        Mat nm;
        snprintf(name, sizeof name, "blk.%d.post_attention_norm.weight", L);
        if (find_mat(g, s, name, &nm, (size_t)n_embd)) row_to_float(&nm, 0, s->norm + (size_t)L * n_embd, n_embd);
        else for (int d = 0; d < n_embd; d++) s->norm[(size_t)L * n_embd + d] = 1.0f;
    }
    gguf_free(g);
    if (!ok) { fprintf(stderr, "sinapsi: tensori FFN non trovati o formato non supportato\n"); syn_close(s); return NULL; }
    s->cols = calloc(CACHE_SLOTS, sizeof(Slot));
    s->rows = calloc(CACHE_SLOTS, sizeof(Slot));
    for (int i = 0; i < CACHE_SLOTS; i++) s->cols[i].key = s->rows[i].key = -1;
    return s;
}

void syn_close(Synapses *s) {
    if (!s) return;
    for (int i = 0; s->cols && i < CACHE_SLOTS; i++) { free(s->cols[i].v); free(s->rows[i].v); }
    free(s->cols); free(s->rows); free(s->down); free(s->gate); free(s->norm);
#ifdef _WIN32
    if (s->map) UnmapViewOfFile(s->map);
    if (s->mapping) CloseHandle(s->mapping);
    if (s->file && s->file != INVALID_HANDLE_VALUE) CloseHandle(s->file);
#else
    if (s->map) munmap((void *)s->map, s->map_size);
#endif
    free(s);
}

// Colonna i di ffn_down(L): cosa scrive il neurone nel residuo. Decodifica un blocco per riga.
static void down_column(const Synapses *s, int L, int i, float *out) {
    const Mat *m = &s->down[L];
    const struct ggml_type_traits *tt = ggml_get_type_traits(m->type);
    int64_t bs = m->type == GGML_TYPE_F32 ? 1 : tt->blck_size;
    size_t bbytes = m->type == GGML_TYPE_F32 ? sizeof(float) : tt->type_size;
    float blk[256];
    for (int d = 0; d < s->n_embd; d++) {
        const uint8_t *p = m->data + (size_t)d * m->row_bytes + (size_t)(i / bs) * bbytes;
        if (m->type == GGML_TYPE_F32) out[d] = *(const float *)p;
        else { tt->to_float(p, blk, bs); out[d] = blk[i % bs]; }
    }
}

static float *cached(Synapses *s, Slot *cache, int64_t key, bool column, int L, int idx) {
    Slot *sl = &cache[(uint64_t)(key * 2654435761u) % CACHE_SLOTS];
    if (sl->key == key) return sl->v;
    if (!sl->v) sl->v = malloc(sizeof(float) * (size_t)s->n_embd);
    if (column) down_column(s, L, idx, sl->v);
    else row_to_float(&s->gate[L], (size_t)idx, sl->v, s->n_embd);
    sl->key = key;
    return sl->v;
}

float syn_pair(Synapses *s, int L, int i, int j) {
    if (L < 0 || L + 1 >= s->n_layer) return 0;
    const float *c = cached(s, s->cols, (int64_t)L * s->n_ff + i, true, L, i);
    const float *r = cached(s, s->rows, (int64_t)(L + 1) * s->n_ff + j, false, L + 1, j);
    const float *g = s->norm + (size_t)(L + 1) * s->n_embd;
    float w = 0;
    for (int d = 0; d < s->n_embd; d++) w += c[d] * g[d] * r[d];
    return w;
}

// Mantiene i k legami più forti (in modulo), ordinati.
static void keep(SynLink *top, int *n, int k, SynLink l) {
    float a = fabsf(l.w);
    if (*n == k && a <= fabsf(top[k - 1].w)) return;
    int i = *n < k ? (*n)++ : k - 1;
    while (i > 0 && fabsf(top[i - 1].w) < a) { top[i] = top[i - 1]; i--; }
    top[i] = l;
}

// vecs (già moltiplicati per la norma dello strato bersaglio) contro tutte le righe di gate(L).
static void scan_gate(const Synapses *s, int L, const float *vecs, int n_vec, int k, SynLink *out, int *counts) {
    float *row = malloc(sizeof(float) * (size_t)s->n_embd);
    for (int v = 0; v < n_vec; v++) counts[v] = 0;
    for (int j = 0; j < s->n_ff; j++) {
        row_to_float(&s->gate[L], (size_t)j, row, s->n_embd);
        for (int v = 0; v < n_vec; v++) {
            const float *x = vecs + (size_t)v * s->n_embd;
            float w = 0;
            for (int d = 0; d < s->n_embd; d++) w += x[d] * row[d];
            keep(out + (size_t)v * k, &counts[v], k, (SynLink){L, j, w});
        }
    }
    free(row);
}

// vecs contro tutte le colonne di down(L): per ogni riga d accumula vec[d] * down[d, j].
static void scan_down(const Synapses *s, int L, const float *vecs, int n_vec, int k, SynLink *out, int *counts) {
    float *row = malloc(sizeof(float) * (size_t)s->n_ff);
    float *acc = calloc((size_t)n_vec * s->n_ff, sizeof(float));
    for (int d = 0; d < s->n_embd; d++) {
        row_to_float(&s->down[L], (size_t)d, row, s->n_ff);
        for (int v = 0; v < n_vec; v++) {
            float x = vecs[(size_t)v * s->n_embd + d];
            float *a = acc + (size_t)v * s->n_ff;
            for (int j = 0; j < s->n_ff; j++) a[j] += x * row[j];
        }
    }
    for (int v = 0; v < n_vec; v++) {
        counts[v] = 0;
        for (int j = 0; j < s->n_ff; j++) keep(out + (size_t)v * k, &counts[v], k, (SynLink){L, j, acc[(size_t)v * s->n_ff + j]});
    }
    free(row); free(acc);
}

void syn_vectors_out(Synapses *s, int L, const float *vecs, int n_vec, int k, SynLink *out) {
    int *counts = calloc((size_t)n_vec, sizeof(int));
    if (L + 1 < s->n_layer) {
        float *v = malloc(sizeof(float) * (size_t)n_vec * s->n_embd);
        const float *g = s->norm + (size_t)(L + 1) * s->n_embd;
        for (int i = 0; i < n_vec; i++)
            for (int d = 0; d < s->n_embd; d++) v[(size_t)i * s->n_embd + d] = vecs[(size_t)i * s->n_embd + d] * g[d];
        scan_gate(s, L + 1, v, n_vec, k, out, counts);
        free(v);
    }
    for (int i = 0; i < n_vec; i++)                       // posti vuoti: nessun legame
        for (int j = counts[i]; j < k; j++) out[(size_t)i * k + j] = (SynLink){-1, -1, 0};
    free(counts);
}

void syn_vectors_in(Synapses *s, int L, const float *vecs, int n_vec, int k, SynLink *out) {
    int *counts = calloc((size_t)n_vec, sizeof(int));
    if (L >= 1) {
        float *v = malloc(sizeof(float) * (size_t)n_vec * s->n_embd);
        const float *g = s->norm + (size_t)L * s->n_embd;
        for (int i = 0; i < n_vec; i++)
            for (int d = 0; d < s->n_embd; d++) v[(size_t)i * s->n_embd + d] = vecs[(size_t)i * s->n_embd + d] * g[d];
        scan_down(s, L - 1, v, n_vec, k, out, counts);
        free(v);
    }
    for (int i = 0; i < n_vec; i++)
        for (int j = counts[i]; j < k; j++) out[(size_t)i * k + j] = (SynLink){-1, -1, 0};
    free(counts);
}

bool syn_neuron(Synapses *s, int L, int i, int k, SynLink *out, int *n_out, SynLink *in, int *n_in) {
    if (L < 0 || L >= s->n_layer || i < 0 || i >= s->n_ff) return false;
    float *col = malloc(sizeof(float) * (size_t)s->n_embd);
    float *row = malloc(sizeof(float) * (size_t)s->n_embd);
    down_column(s, L, i, col);
    row_to_float(&s->gate[L], (size_t)i, row, s->n_embd);
    syn_vectors_out(s, L, col, 1, k, out);   // senza cache: sicuro da un thread qualsiasi
    syn_vectors_in(s, L, row, 1, k, in);
    *n_out = *n_in = 0;
    for (int j = 0; j < k; j++) { if (out[j].layer >= 0) (*n_out)++; if (in[j].layer >= 0) (*n_in)++; }
    free(col); free(row);
    return true;
}