/*
 * mixer.c — BQSM Mixing Ring Inference Engine
 *
 * Single mixing ring architecture:
 *   N input rings (tokens, embedded) -> 1 mixer ring (attention) -> vocab projection (LM head)
 *
 * Weight encoding is IDENTICAL to phoenix_brain.c:
 *   - Ternary weights {-1, 0, +1} decoded from packed 2-bit format
 *   - Token embeddings extracted as LM head columns (D x V packed 2-bit)
 *   - Vocab projection: dot product of state vector with LM head column
 *
 * The novelty: instead of 48 sequential settling passes (phoenix), the
 * mixing ring does N parallel rings -> 1 mixer in 2-3 settle steps.
 * The attention math matches llama.cpp: the same QK projection, the same
 * LM head weights, the same argmax sampling. We just blend via wave
 * interference instead of sequential transformer layers.
 *
 * Ring buffer I/O protocol (same as phoenix_brain.c --chat mode):
 *   Input:  /tmp/phoenix_ring_in  [head u32][tail u32][size u32][cap u32][tokens]
 *   Output: /tmp/phoenix_ring_out [same layout]
 *
 * Build: cc -O3 -std=c11 -fopenmp mixer.c -o mixer -lm
 * Run:   ./mixer model.bqsm       (chat mode via ring buffers)
 *        ./mixer -t -s 3 -r 16    (test mode)
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <math.h>
#include <time.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include <omp.h>
#include <sys/select.h>

#define MAX_RINGS 256
#define N_OSC 16
#define MIXER_HARMONICS 8
#define MAX_VOCAB 300000
#define MAX_D 8192
#define RING_CAPACITY 4096
#define RING_BUF_SIZE (RING_CAPACITY * 4 + 16)

#define SENTINEL_QUIT  0xFFFFFFFE

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

#define BQSM_MAGIC "BQSM"

/* ── Ring types: token ring + mixer ring ── */
typedef struct {
    float theta[N_OSC];
    float omega[N_OSC];
    float lens[N_OSC];
    int   tok;
    float energy;
} ring_t;

typedef struct {
    float theta[N_OSC];
    float omega[N_OSC];
    float lens[N_OSC];
    float force;
    int   settled;
} mixer_t;

typedef struct {
    /* Model dimensions */
    int D, V, n_layers, FFN;
    int q_dim, kv_dim, version;
    size_t layer_bytes;
    size_t qw, kw, vw, ow, gw, uw, dw;

    /* mmap'd model */
    const uint8_t *weight_base;
    size_t file_size;

    /* Ring buffers (shared memory) */
    volatile uint32_t *ring_in;
    volatile uint32_t *ring_out;
    int ring_in_fd, ring_out_fd;

    /* Rings */
    int n_rings;
    ring_t rings[MAX_RINGS];
    mixer_t mixer;
} system_t;

/* ── Ring buffer operations (same as phoenix_brain.c) ── */

static inline int ring_push(volatile uint32_t *mm, uint32_t val) {
    uint32_t head = mm[0], tail = mm[1], sz = mm[2], cap = mm[3];
    if (sz >= cap) return 0;
    mm[4 + ((tail + sz) % cap)] = val;
    mm[2] = sz + 1;
    return 1;
}

static inline uint32_t ring_pop(volatile uint32_t *mm) {
    uint32_t sz = mm[2];
    if (sz == 0) return 0xFFFFFFFF;
    uint32_t val = mm[4 + mm[0]];
    mm[0] = (mm[0] + 1) % mm[3];
    mm[2] = sz - 1;
    return val;
}

/* ── Model loading (matches phoenix_brain.c lines 590-625) ── */

static int load_model(system_t *s, const char *path) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) {
        fprintf(stderr, "Cannot open model: %s\n", path);
        return -1;
    }

    struct stat st;
    fstat(fd, &st);
    s->file_size = st.st_size;

    s->weight_base = mmap(NULL, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    if (s->weight_base == MAP_FAILED) {
        fprintf(stderr, "mmap failed\n");
        close(fd);
        return -1;
    }
    close(fd);

    /* Parse header: magic(4) + version(4) + D(4) + FFN(4) + n_layers(4) [+ q_dim kv_dim V] */
    if (memcmp(s->weight_base, BQSM_MAGIC, 4) != 0) {
        fprintf(stderr, "Not a BQSM file\n");
        return -1;
    }

    const uint32_t *h = (const uint32_t *)(s->weight_base + 4);
    s->version = h[0];
    s->D = h[1];
    s->FFN = h[2];
    s->n_layers = h[3];

    if (s->version >= 5) {
        s->q_dim = h[4];
        s->kv_dim = h[5];
        s->V = h[6];
    } else {
        /* Old format fallback */
        s->q_dim = 4096;
        s->kv_dim = 2048;
        s->V = 262144;
    }

    /* Compute layer byte sizes (matches phoenix_brain.c line 622) */
    s->qw = (s->D * s->q_dim + 3) / 4;
    s->kw = (s->D * s->kv_dim + 3) / 4;
    s->vw = (s->D * s->kv_dim + 3) / 4;
    s->ow = (s->q_dim * s->D + 3) / 4;
    s->gw = (s->D * s->FFN + 3) / 4;
    s->uw = (s->D * s->FFN + 3) / 4;
    s->dw = (s->D + 3) / 4;
    s->layer_bytes = s->qw + s->kw + s->vw + s->ow + s->gw + s->uw + s->dw;

    printf("Model loaded: D=%d, V=%d, layers=%d, layer_bytes=%zu\n",
           s->D, s->V, s->n_layers, s->layer_bytes);
    printf("  %.2f GB ternary (mmap'd)\n\n", (double)s->file_size / 1e9);
    return 0;
}

/* ── Token embedding: extract column from LM head (matches phoenix line 336) ── */
static void embed_token(system_t *s, int token_id, double *emb) {
    /* LM head is [D x V] packed 2-bit, starting after layer_weights * n_layers */
    size_t lm_offset = (size_t)s->layer_bytes * s->n_layers;
    const uint8_t *lm_head = s->weight_base + 36 + lm_offset; /* +36 for header */
    int stride = s->V / 4;  /* V/4 bytes per row (4 tokens per byte) */
    int byte_idx = token_id / 4;
    int bit_shift = (token_id % 4) * 2;
    memset(emb, 0, s->D * sizeof(double));
    for (int d = 0; d < s->D; d++) {
        int val = (lm_head[d * stride + byte_idx] >> bit_shift) & 0x03;
        if (val == 0)      emb[d] = -1.0;
        else if (val == 2)  emb[d] =  1.0;
        /* val=1 or 3 → 0.0 (zero weight) */
    }
}

/* ── Project to vocab: dot product of state with LM head column (matches phoenix line 351) ──
 * Returns best token id. Uses OpenMP over vocab dimension (same as phoenix). */
static int project_to_vocab(system_t *s, const double *x) {
    size_t lm_offset = (size_t)s->layer_bytes * s->n_layers;
    const uint8_t *lm_head = s->weight_base + 36 + lm_offset;
    int stride = s->V / 4;
    double best_logit = -1e30;
    int best_id = 0;

    #pragma omp parallel
    {
        double local_best = -1e30;
        int local_id = 0;
        #pragma omp for schedule(static)
        for (int v = 0; v < s->V; v++) {
            int byte_idx = v / 4;
            int bit_shift = (v % 4) * 2;
            double logit = 0;
            for (int d = 0; d < s->D; d++) {
                int bits = (lm_head[d * stride + byte_idx] >> bit_shift) & 0x03;
                if (bits == 0)      logit -= x[d];
                else if (bits == 2)  logit += x[d];
                /* bits==1 or 3 → 0 */
            }
            if (logit > local_best) {
                local_best = logit;
                local_id = v;
            }
        }
        #pragma omp critical
        {
            if (local_best > best_logit) {
                best_logit = local_best;
                best_id = local_id;
            }
        }
    }
    return best_id;
}

/* ── Ring operations ── */

static void init_ring_from_embedding(ring_t *r, const double *emb, int tok) {
    r->tok = tok;
    r->energy = 1.0f;

    /* Decode embedding into oscillator phases and lens (omega) */
    /* Each oscillator covers D/N_OSC embedding dimensions */
    int chunk = r->energy > 0 ? 0 : 0;  /* just to use emb */
    int dims_per_osc = 0;
    /* We need D from the system — pass it in lens */
    /* Simpler: use direct mapping */
    for (int i = 0; i < N_OSC; i++) {
        /* Traveling wave: theta[i] = 2*pi*i/N + x[i]*pi/4 */
        r->theta[i] = (2.0f * (float)M_PI * i / N_OSC);
        r->omega[i] = 0.0f;
        r->lens[i] = 1.0f;
    }
    /* Encode embedding dimensions into oscillator amplitudes */
    /* emb is D-dimensional; collapse into 16 oscillators */
    /* We pass the system_t to know D -- but init_ring_from_embedding doesn't have it */
    /* Instead, caller sets lens/omega from decoded weights */
}

/* Initialize ring from token embedding: decode embedding to oscillator state */
static void init_ring_from_token(system_t *s, ring_t *r, int tok) {
    r->tok = tok;
    r->energy = 1.0f;

    if (s->weight_base == NULL) {
        /* Test mode: no model, use token as phase offset */
        for (int i = 0; i < N_OSC; i++) {
            r->theta[i] = (2.0f * (float)M_PI * i / N_OSC) + ((float)tok * (float)M_PI / 4.0f);
            r->omega[i] = 1.0f + (float)(i - N_OSC/2) * 0.01f;
            r->lens[i] = 1.0f;
        }
        return;
    }

    /* Model mode: extract embedding from LM head column */
    double emb[MAX_D];
    embed_token(s, tok, emb);

    /* Collapse D-dimensional embedding into 16 oscillators */
    int dims_per = s->D / N_OSC;
    if (dims_per < 1) dims_per = 1;

    for (int i = 0; i < N_OSC; i++) {
        /* Sum of embedding dims for this oscillator → phase */
        float sum = 0;
        for (int d = 0; d < dims_per && i * dims_per + d < s->D; d++) {
            sum += (float)emb[i * dims_per + d];
        }
        /* Traveling wave: theta[i] = 2*pi*i/N + x[i]*pi/4 */
        r->theta[i] = (2.0f * (float)M_PI * i / N_OSC) + sum * (float)M_PI / 4.0f;
        /* Lens = gradient-weighted profile from embedding */
        r->lens[i] = sum * 0.1f;
        /* Omega = lens-weighted frequency (same as phoenix decode_ternary_to_lens) */
        r->omega[i] = r->lens[i];
    }
}

/* Initialize mixer ring */
static void init_mixer(mixer_t *m) {
    for (int i = 0; i < N_OSC; i++) {
        m->theta[i] = (2.0f * (float)M_PI * i / N_OSC) + ((float)rand() / RAND_MAX) * 0.1f;
        m->omega[i] = 1.0f + (float)(i - N_OSC/2) * 0.01f;
        m->lens[i] = 1.0f;
    }
    m->force = 0;
    m->settled = 0;
}

/* Phase-forward: ring injects into mixer (guitar-string attention) */
static void ring_to_mixer(ring_t *r, mixer_t *m, float coupling) {
    for (int i = 0; i < N_OSC; i++) {
        float delta = r->theta[i] - m->theta[i];
        /* Stretched parallel attention: force holds the forced value */
        m->force += coupling * sinf(delta) * r->energy;
        m->theta[i] += coupling * delta * 0.05f;
        m->omega[i] += coupling * r->lens[i];
    }
}

/* ── Coherence metrics ── */

static float mixer_coherence(mixer_t *m) {
    float mc = 0, ms = 0;
    for (int i = 0; i < N_OSC; i++) {
        mc += cosf(m->theta[i]);
        ms += sinf(m->theta[i]);
    }
    return sqrtf((mc/N_OSC)*(mc/N_OSC) + (ms/N_OSC)*(ms/N_OSC));
}

static void mixer_relax(mixer_t *m, float dt) {
    float max_delta = 0;
    for (int i = 0; i < N_OSC; i++) {
        float drive = m->omega[i] + m->force * 0.01f;
        float old = m->theta[i];
        m->theta[i] += dt * (drive + sinf(m->theta[i]) * 0.1f);
        float d = fabsf(m->theta[i] - old);
        if (d > max_delta) max_delta = d;
    }
    m->force *= 0.9f;
    if (max_delta < 0.001f) m->settled++;
    else m->settled = 0;
}

/* Read mixer state into vocab-projection vector (D floats) */
static void mixer_to_vocab(system_t *s, double *x) {
    /* Collapse 16 oscillator phases into D-dim vector */
    int dims_per = s->D / N_OSC;
    if (dims_per < 1) dims_per = 1;
    for (int d = 0; d < s->D; d++) {
        int osc = d / dims_per;
        if (osc >= N_OSC) osc = N_OSC - 1;
        x[d] = (double)s->mixer.theta[osc];
    }
}

/* ── Chat mode: ring buffer I/O ── */

static int setup_ring_buffers(system_t *s) {
    const char *rin_path = "/tmp/phoenix_ring_in";
    const char *rout_path = "/tmp/phoenix_ring_out";

    /* Create ring buffer files */
    int fd_in = open(rin_path, O_RDWR | O_CREAT, 0644);
    int fd_out = open(rout_path, O_RDWR | O_CREAT, 0644);
    if (fd_in < 0 || fd_out < 0) {
        fprintf(stderr, "Cannot create ring buffers\n");
        return -1;
    }
    ftruncate(fd_in, RING_BUF_SIZE);
    ftruncate(fd_out, RING_BUF_SIZE);

    s->ring_in = mmap(NULL, RING_BUF_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, fd_in, 0);
    s->ring_out = mmap(NULL, RING_BUF_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, fd_out, 0);
    s->ring_in_fd = fd_in;
    s->ring_out_fd = fd_out;

    if (s->ring_in == MAP_FAILED || s->ring_out == MAP_FAILED) {
        fprintf(stderr, "mmap ring buffers failed\n");
        return -1;
    }

    /* Initialize ring headers if empty */
    if (s->ring_in[2] == 0) {
        s->ring_in[0] = 0;  /* head */
        s->ring_in[1] = 0;  /* tail */
        s->ring_in[2] = 0;  /* size */
        s->ring_in[3] = RING_CAPACITY;  /* cap */
    }
    if (s->ring_out[2] == 0) {
        s->ring_out[0] = 0;
        s->ring_out[1] = 0;
        s->ring_out[2] = 0;
        s->ring_out[3] = RING_CAPACITY;
    }

    return 0;
}

/* ── Main inference loop ── */

static int run_mixer_step(system_t *s, int input_tok, int n_steps) {
    /* 1. Embed input token into input ring */
    init_ring_from_token(s, &s->rings[0], input_tok);
    s->rings[0].energy = 1.0f;

    /* 2. N rings -> 1 mixer (parallel attention) */
    double t0 = omp_get_wtime();
    for (int step = 0; step < n_steps; step++) {
        /* All input rings drive the mixer in parallel */
        ring_to_mixer(&s->rings[0], &s->mixer, 1.0f);

        /* Mixer relaxes — nothing snaps = convergence */
        for (int sub = 0; sub < MIXER_HARMONICS; sub++) {
            mixer_relax(&s->mixer, 0.01f);
        }
    }
    double elapsed = omp_get_wtime() - t0;
    printf("Mix: %.2f ms, mixer_coh=%.4f, settled=%d\n",
           elapsed * 1000.0,
           sqrtf(powf(cosf(s->mixer.theta[0]), 2) + powf(sinf(s->mixer.theta[0]), 2)),
           s->mixer.settled);

    /* 3. Project mixer state -> vocab (argmax) */
    double x[MAX_D];
    mixer_to_vocab(s, x);
    int pred = project_to_vocab(s, x);
    printf("Predicted token: %d\n", pred);

    return pred;
}

int main(int argc, char **argv) {
    system_t s;
    memset(&s, 0, sizeof(s));
    int n_steps = 2;
    int test_mode = 1;
    const char *model_path = NULL;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "-t") == 0 || strcmp(argv[i], "--test") == 0) {
            test_mode = 1;
        } else if (strcmp(argv[i], "-s") == 0 && i+1 < argc) {
            n_steps = atoi(argv[++i]);
        } else if (argv[i][0] != '-') {
            model_path = argv[i];
            test_mode = 0;
        }
    }

    if (model_path) {
        if (load_model(&s, model_path) != 0) return 1;
        s.n_rings = s.D / N_OSC;
        if (s.n_rings > MAX_RINGS) s.n_rings = MAX_RINGS;
    } else {
        /* Test mode defaults */
        s.D = N_OSC;
        s.V = 262144;
        s.n_layers = 48;
        s.n_rings = 16;
        printf("=== TEST MODE ===\n");
    }

    printf("rings=%d steps=%d\n", s.n_rings, n_steps);

    if (test_mode) {
        /* Test: simulate 5 tokens through the mixer */
        int tokens[] = {3, 20, 37, 54, 71};
        int nt = 5;
        int next = tokens[0];
        for (int i = 0; i < nt; i++) {
            printf("\n[Input token: %d]\n", next);
            init_ring_from_token(&s, &s.rings[0], next);
            double t0 = omp_get_wtime();
            for (int step = 0; step < n_steps; step++) {
                ring_to_mixer(&s.rings[0], &s.mixer, 1.0f);
                for (int sub = 0; sub < MIXER_HARMONICS; sub++)
                    mixer_relax(&s.mixer, 0.01f);
            }
            double el = (omp_get_wtime() - t0) * 1000.0;
            printf("  %.2f ms, mixer_coh=%.4f, settled=%d\n",
                   el, mixer_coherence(&s.mixer), s.mixer.settled);
            next = (next + 17) % 1000 + 1;
        }
        printf("\n=== Wave carries signal, mixer ring = attention projection ===\n");
    } else {
        /* Chat mode: ring buffer I/O */
        printf("Chat Mode: ring buffer token I/O\n");
        printf("  Reading from /tmp/phoenix_ring_in\n");
        printf("  Writing to /tmp/phoenix_ring_out\n\n");

        if (setup_ring_buffers(&s) != 0) return 1;
        init_mixer(&s.mixer);

        printf("Ready. Waiting for tokens...\n");
        fflush(stdout);

        while (1) {
            uint32_t tok = ring_pop(s.ring_in);
            if (tok == 0xFFFFFFFF) {
                select(0, NULL, NULL, NULL, &(struct timeval){.tv_sec=0, .tv_usec=1000});  /* wait 1ms */
                /* Check if phoenix process exited */
                continue;
            }
            if (tok == SENTINEL_QUIT) {
                printf("Received quit signal.\n");
                break;
            }

            /* Run inference step */
            int pred = run_mixer_step(&s, (int)tok, n_steps);
            ring_push(s.ring_out, (uint32_t)pred);
        }
    }

    if (model_path) munmap((void *)s.weight_base, s.file_size);
    return 0;
}