/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
#include "ternarize.h"
#include "ft_layer_ops.h"
#include "../safetensors/safetensors.h"
#include "../quantization/quantization.h"
#include "../convert/convert.h"
#include "../tensor/tensor.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <sys/stat.h>
#include <errno.h>

/* ============================================================= */
/* PRNG (xorshift64, same as fq_ module)                         */
/* ============================================================= */

static uint64_t ft_xorshift64(uint64_t *state) {
    uint64_t x = *state;
    x ^= x << 13;
    x ^= x >> 7;
    x ^= x << 17;
    *state = x;
    return x;
}

static float ft_rand_uniform(uint64_t *state) {
    return (ft_xorshift64(state) >> 11) * (1.0 / 9007199254740992.0);
}

/* ============================================================= */
/* Minimal JSON helpers (reused from convert.c pattern)           */
/* ============================================================= */

static char* read_file_str(const char *path) {
    FILE *f = fopen(path, "r");
    if (!f) return NULL;
    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    fseek(f, 0, SEEK_SET);
    char *buf = malloc(size + 1);
    size_t rd = fread(buf, 1, size, f);
    buf[rd] = '\0';
    fclose(f);
    return buf;
}

static int json_get_int(const char *json, const char *key, int def) {
    char search[128];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) return def;
    p += strlen(search);
    while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
    return (int)strtol(p, NULL, 10);
}

static float json_get_float(const char *json, const char *key, float def) {
    char search[128];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) return def;
    p += strlen(search);
    while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
    return strtof(p, NULL);
}

static char* json_get_string(const char *json, const char *key) {
    char search[128];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) return NULL;
    p += strlen(search);
    while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
    if (*p != '"') return NULL;
    p++;
    const char *end = p;
    while (*end && *end != '"') {
        if (*end == '\\') end++;
        if (*end) end++;
    }
    size_t len = end - p;
    char *r = malloc(len + 1);
    memcpy(r, p, len);
    r[len] = '\0';
    return r;
}

static int json_get_bool(const char *json, const char *key, int def) {
    char search[128];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) return def;
    p += strlen(search);
    while (*p && (*p == ' ' || *p == ':' || *p == '\t' || *p == '\n')) p++;
    if (*p == 't') return 1;
    if (*p == 'f') return 0;
    return def;
}

/* ============================================================= */
/* BF16/F16 conversion (from convert.c)                          */
/* ============================================================= */

static void bf16_to_f32(const uint16_t *src, float *dst, int64_t n) {
    for (int64_t i = 0; i < n; i++) {
        uint32_t bits = (uint32_t)src[i] << 16;
        memcpy(&dst[i], &bits, 4);
    }
}

static void f16_to_f32(const uint16_t *src, float *dst, int64_t n) {
    for (int64_t i = 0; i < n; i++) {
        uint16_t h = src[i];
        uint32_t sign = (uint32_t)(h >> 15) << 31;
        uint32_t exp  = (h >> 10) & 0x1F;
        uint32_t mant = h & 0x3FF;
        uint32_t bits;
        if (exp == 0) {
            if (mant == 0) bits = sign;
            else {
                exp = 1;
                while (!(mant & 0x400)) { mant <<= 1; exp--; }
                mant &= 0x3FF;
                bits = sign | ((uint32_t)(exp + 127 - 15) << 23) | ((uint32_t)mant << 13);
            }
        } else if (exp == 31) {
            bits = sign | 0x7F800000 | ((uint32_t)mant << 13);
        } else {
            bits = sign | ((uint32_t)(exp + 127 - 15) << 23) | ((uint32_t)mant << 13);
        }
        memcpy(&dst[i], &bits, 4);
    }
}

static float* tensor_to_f32(const void *data, fs_dtype dtype, int64_t numel) {
    float *f32 = malloc(numel * sizeof(float));
    if (!f32) return NULL;
    switch (dtype) {
        case FS_DTYPE_BF16: bf16_to_f32((const uint16_t *)data, f32, numel); break;
        case FS_DTYPE_F16:  f16_to_f32((const uint16_t *)data, f32, numel); break;
        case FS_DTYPE_F32:  memcpy(f32, data, numel * sizeof(float)); break;
        default: memset(f32, 0, numel * sizeof(float)); break;
    }
    return f32;
}

/* ============================================================= */
/* Utility                                                       */
/* ============================================================= */

static void mkdirp(const char *path) {
    char tmp[1024];
    snprintf(tmp, sizeof(tmp), "%s", path);
    for (char *p = tmp + 1; *p; p++) {
        if (*p == '/') {
            *p = '\0';
            mkdir(tmp, 0755);
            *p = '/';
        }
    }
    mkdir(tmp, 0755);
}

const char* ft_arch_name(ft_arch arch) {
    switch (arch) {
        case FT_ARCH_LLAMA:   return "Llama";
        case FT_ARCH_QWEN2:   return "Qwen2";
        case FT_ARCH_QWEN3:   return "Qwen3";
        case FT_ARCH_PHI3:    return "Phi-3";
        case FT_ARCH_GEMMA2:  return "Gemma2";
        case FT_ARCH_RWKVX:   return "RWKV-X";
        case FT_ARCH_UNKNOWN: return "Unknown";
    }
    return "Unknown";
}

static int cmp_float_desc(const void *a, const void *b) {
    float fa = *(const float *)a;
    float fb = *(const float *)b;
    if (fa > fb) return -1;
    if (fa < fb) return 1;
    return 0;
}

/* ============================================================= */
/* Config                                                        */
/* ============================================================= */

ft_config ft_config_defaults(void) {
    ft_config cfg;
    memset(&cfg, 0, sizeof(cfg));

    cfg.arch = FT_ARCH_UNKNOWN;
    cfg.group_size = 128;
    cfg.k_frac_start = 0.00008f;
    cfg.k_frac_end = 0.000005f;
    cfg.snr_weight = 0.3f;
    cfg.hess_power = 0.15f;
    cfg.momentum_bonus = 0.2f;
    cfg.boltz_frac_start = 0.17f;
    cfg.boltz_frac_end = 0.0f;
    cfg.temporal_weight = 1.8f;
    cfg.phase1_iters = 60;
    cfg.phase2_iters = 140;
    cfg.batch_size = 128;

    return cfg;
}

bool ft_config_from_json(ft_config *cfg, const char *path) {
    char *json = read_file_str(path);
    if (!json) {
        fprintf(stderr, "ERROR: cannot read %s\n", path);
        return false;
    }

    char *mt = json_get_string(json, "model_type");
    if (mt) {
        if (strstr(mt, "qwen3"))                              cfg->arch = FT_ARCH_QWEN3;
        else if (strstr(mt, "qwen2") || strstr(mt, "qwen"))  cfg->arch = FT_ARCH_QWEN2;
        else if (strstr(mt, "llama") || strstr(mt, "mistral"))cfg->arch = FT_ARCH_LLAMA;
        else if (strstr(mt, "phi3") || strstr(mt, "phi"))     cfg->arch = FT_ARCH_PHI3;
        else if (strstr(mt, "gemma2") || strstr(mt, "gemma")) cfg->arch = FT_ARCH_GEMMA2;
        else if (strstr(mt, "rwkv"))                          cfg->arch = FT_ARCH_RWKVX;
        else                                                  cfg->arch = FT_ARCH_UNKNOWN;
        free(mt);
    }

    cfg->n_layers    = json_get_int(json, "num_hidden_layers", 28);
    cfg->n_embd      = json_get_int(json, "hidden_size", 1024);
    cfg->n_heads     = json_get_int(json, "num_attention_heads", 16);
    cfg->n_kv_heads  = json_get_int(json, "num_key_value_heads", cfg->n_heads);
    cfg->ffn_dim     = json_get_int(json, "intermediate_size", 2816);
    cfg->vocab_size  = json_get_int(json, "vocab_size", 32000);
    cfg->head_dim    = cfg->n_embd / cfg->n_heads;
    cfg->rope_theta  = json_get_float(json, "rope_theta", 10000.0f);
    cfg->rms_eps     = json_get_float(json, "rms_norm_eps", 1e-6f);
    cfg->tie_word_embeddings = json_get_bool(json, "tie_word_embeddings", 1);

    switch (cfg->arch) {
        case FT_ARCH_QWEN3:  cfg->has_qk_norm = 1; cfg->has_qkv_bias = 0; break;
        case FT_ARCH_QWEN2:
        case FT_ARCH_PHI3:   cfg->has_qk_norm = 0; cfg->has_qkv_bias = 1; break;
        case FT_ARCH_GEMMA2: cfg->has_qk_norm = 1; cfg->has_qkv_bias = 0; break;
        default:             cfg->has_qk_norm = 0; cfg->has_qkv_bias = 0; break;
    }

    free(json);
    return true;
}

/* ============================================================= */
/* Layer State I/O                                               */
/* ============================================================= */

static int64_t compute_layer_numel(const ft_config *cfg) {
    int q_dim  = cfg->n_heads * cfg->head_dim;
    int kv_dim = cfg->n_kv_heads * cfg->head_dim;
    int64_t n = 0;
    n += (int64_t)q_dim * cfg->n_embd;       /* q_proj */
    n += (int64_t)kv_dim * cfg->n_embd;      /* k_proj */
    n += (int64_t)kv_dim * cfg->n_embd;      /* v_proj */
    n += (int64_t)cfg->n_embd * q_dim;       /* o_proj */
    n += (int64_t)cfg->ffn_dim * cfg->n_embd; /* gate_proj */
    n += (int64_t)cfg->ffn_dim * cfg->n_embd; /* up_proj */
    n += (int64_t)cfg->n_embd * cfg->ffn_dim; /* down_proj */
    return n;
}

static void register_layer_weights(ft_layer_state *state, const ft_config *cfg) {
    int q_dim  = cfg->n_heads * cfg->head_dim;
    int kv_dim = cfg->n_kv_heads * cfg->head_dim;

    /* 7 weight matrices per LLaMA-family layer */
    state->n_weights = 7;
    state->weights = calloc(state->n_weights, sizeof(ft_weight_entry));

    int64_t off = 0;
    struct { const char *name; int rows; int cols; } defs[] = {
        {"self_attn.q_proj.weight", q_dim,        cfg->n_embd},
        {"self_attn.k_proj.weight", kv_dim,       cfg->n_embd},
        {"self_attn.v_proj.weight", kv_dim,       cfg->n_embd},
        {"self_attn.o_proj.weight", cfg->n_embd,  q_dim},
        {"mlp.gate_proj.weight",    cfg->ffn_dim, cfg->n_embd},
        {"mlp.up_proj.weight",      cfg->ffn_dim, cfg->n_embd},
        {"mlp.down_proj.weight",    cfg->n_embd,  cfg->ffn_dim},
    };

    for (int i = 0; i < 7; i++) {
        snprintf(state->weights[i].name, sizeof(state->weights[i].name), "%s", defs[i].name);
        state->weights[i].rows = defs[i].rows;
        state->weights[i].cols = defs[i].cols;
        state->weights[i].offset = off;
        state->weights[i].numel = (int64_t)defs[i].rows * defs[i].cols;
        off += state->weights[i].numel;
    }
}

static ft_layer_state* ft_layer_alloc(const ft_config *cfg, int layer_idx) {
    ft_layer_state *s = calloc(1, sizeof(ft_layer_state));
    if (!s) return NULL;

    s->layer_idx = layer_idx;
    s->numel_total = compute_layer_numel(cfg);
    s->group_size = cfg->group_size;
    s->n_groups = (int)((s->numel_total + cfg->group_size - 1) / cfg->group_size);

    /* Bit arrays: 8 elements per byte */
    int64_t n_bytes = (s->numel_total + 7) / 8;
    s->sign     = calloc(n_bytes, 1);
    s->mask     = calloc(n_bytes, 1);
    s->momentum = calloc(n_bytes, 1);

    s->scales      = calloc(s->n_groups, sizeof(float));
    s->prev_scores = calloc(s->numel_total, sizeof(float));

    /* Norms */
    s->norm_attn = calloc(cfg->n_embd, sizeof(float));
    s->norm_ffn  = calloc(cfg->n_embd, sizeof(float));

    register_layer_weights(s, cfg);

    s->best_loss = 1e30f;
    return s;
}

void ft_layer_free(ft_layer_state *state) {
    if (!state) return;
    free(state->sign);
    free(state->mask);
    free(state->momentum);
    free(state->scales);
    free(state->prev_scores);
    free(state->norm_attn);
    free(state->norm_ffn);
    free(state->weights);
    free(state);
}

static void layer_path(char *buf, size_t bufsz, const ft_config *cfg, int layer_idx, const char *file) {
    snprintf(buf, bufsz, "%s/layer_%03d/%s", cfg->workdir, layer_idx, file);
}

bool ft_layer_save(const ft_layer_state *s, const ft_config *cfg) {
    char dir[1024], path[1024];
    snprintf(dir, sizeof(dir), "%s/layer_%03d", cfg->workdir, s->layer_idx);
    mkdirp(dir);

    int64_t n_bytes = (s->numel_total + 7) / 8;

    /* ternary.bin: sign + mask + momentum (3 bit arrays) */
    layer_path(path, sizeof(path), cfg, s->layer_idx, "ternary.bin");
    FILE *f = fopen(path, "wb");
    if (!f) return false;
    /* Header: numel, iteration, phase, total_flips, best_loss */
    fwrite(&s->numel_total, sizeof(int64_t), 1, f);
    fwrite(&s->iteration, sizeof(int), 1, f);
    fwrite(&s->phase, sizeof(int), 1, f);
    fwrite(&s->total_flips, sizeof(int64_t), 1, f);
    fwrite(&s->best_loss, sizeof(float), 1, f);
    fwrite(s->sign, 1, n_bytes, f);
    fwrite(s->mask, 1, n_bytes, f);
    fwrite(s->momentum, 1, n_bytes, f);
    /* prev_scores for temporal */
    fwrite(s->prev_scores, sizeof(float), s->numel_total, f);
    fclose(f);

    /* scales.bin */
    layer_path(path, sizeof(path), cfg, s->layer_idx, "scales.bin");
    f = fopen(path, "wb");
    if (!f) return false;
    fwrite(&s->n_groups, sizeof(int), 1, f);
    fwrite(&s->group_size, sizeof(int), 1, f);
    fwrite(s->scales, sizeof(float), s->n_groups, f);
    fclose(f);

    /* norms.bin */
    layer_path(path, sizeof(path), cfg, s->layer_idx, "norms.bin");
    f = fopen(path, "wb");
    if (!f) return false;
    fwrite(&cfg->n_embd, sizeof(int), 1, f);
    fwrite(s->norm_attn, sizeof(float), cfg->n_embd, f);
    fwrite(s->norm_ffn, sizeof(float), cfg->n_embd, f);
    fclose(f);

    return true;
}

ft_layer_state* ft_layer_load(const ft_config *cfg, int layer_idx) {
    ft_layer_state *s = ft_layer_alloc(cfg, layer_idx);
    if (!s) return NULL;

    char path[1024];
    int64_t n_bytes = (s->numel_total + 7) / 8;

    /* ternary.bin */
    layer_path(path, sizeof(path), cfg, layer_idx, "ternary.bin");
    FILE *f = fopen(path, "rb");
    if (!f) { ft_layer_free(s); return NULL; }

    int64_t stored_numel;
    fread(&stored_numel, sizeof(int64_t), 1, f);
    if (stored_numel != s->numel_total) {
        fprintf(stderr, "ERROR: layer %d numel mismatch: %lld vs %lld\n",
                layer_idx, (long long)stored_numel, (long long)s->numel_total);
        fclose(f); ft_layer_free(s); return NULL;
    }
    fread(&s->iteration, sizeof(int), 1, f);
    fread(&s->phase, sizeof(int), 1, f);
    fread(&s->total_flips, sizeof(int64_t), 1, f);
    fread(&s->best_loss, sizeof(float), 1, f);
    fread(s->sign, 1, n_bytes, f);
    fread(s->mask, 1, n_bytes, f);
    fread(s->momentum, 1, n_bytes, f);
    fread(s->prev_scores, sizeof(float), s->numel_total, f);
    fclose(f);

    /* scales.bin */
    layer_path(path, sizeof(path), cfg, layer_idx, "scales.bin");
    f = fopen(path, "rb");
    if (f) {
        fread(&s->n_groups, sizeof(int), 1, f);
        fread(&s->group_size, sizeof(int), 1, f);
        fread(s->scales, sizeof(float), s->n_groups, f);
        fclose(f);
    }

    /* norms.bin */
    layer_path(path, sizeof(path), cfg, layer_idx, "norms.bin");
    f = fopen(path, "rb");
    if (f) {
        int n_embd;
        fread(&n_embd, sizeof(int), 1, f);
        fread(s->norm_attn, sizeof(float), n_embd, f);
        fread(s->norm_ffn, sizeof(float), n_embd, f);
        fclose(f);
    }

    return s;
}

/* ============================================================= */
/* Gradient Stats I/O                                            */
/* ============================================================= */

ft_grad_stats* ft_grad_stats_load(const ft_config *cfg, int layer_idx) {
    char path[1024];
    layer_path(path, sizeof(path), cfg, layer_idx, "grad_stats.bin");
    FILE *f = fopen(path, "rb");
    if (!f) return NULL;

    ft_grad_stats *g = calloc(1, sizeof(ft_grad_stats));
    fread(&g->numel, sizeof(int64_t), 1, f);
    fread(&g->n_samples, sizeof(int), 1, f);

    g->grad_sum     = malloc(g->numel * sizeof(float));
    g->grad_sq_sum  = malloc(g->numel * sizeof(float));
    g->hessian_diag = malloc(g->numel * sizeof(float));

    fread(g->grad_sum, sizeof(float), g->numel, f);
    fread(g->grad_sq_sum, sizeof(float), g->numel, f);
    fread(g->hessian_diag, sizeof(float), g->numel, f);
    fclose(f);
    return g;
}

bool ft_grad_stats_save(const ft_grad_stats *g, const ft_config *cfg, int layer_idx) {
    char path[1024];
    layer_path(path, sizeof(path), cfg, layer_idx, "grad_stats.bin");
    FILE *f = fopen(path, "wb");
    if (!f) return false;
    fwrite(&g->numel, sizeof(int64_t), 1, f);
    fwrite(&g->n_samples, sizeof(int), 1, f);
    fwrite(g->grad_sum, sizeof(float), g->numel, f);
    fwrite(g->grad_sq_sum, sizeof(float), g->numel, f);
    fwrite(g->hessian_diag, sizeof(float), g->numel, f);
    fclose(f);
    return true;
}

void ft_grad_stats_free(ft_grad_stats *g) {
    if (!g) return;
    free(g->grad_sum);
    free(g->grad_sq_sum);
    free(g->hessian_diag);
    free(g);
}

/* ============================================================= */
/* Status I/O                                                    */
/* ============================================================= */

ft_status ft_status_read(const ft_config *cfg) {
    ft_status st;
    memset(&st, 0, sizeof(st));

    char path[1024];
    snprintf(path, sizeof(path), "%s/status.json", cfg->workdir);
    char *json = read_file_str(path);
    if (!json) return st;

    st.n_layers       = json_get_int(json, "n_layers", 0);
    st.layers_split   = json_get_int(json, "layers_split", 0);
    st.layers_swept   = json_get_int(json, "layers_swept", 0);
    st.layers_flipped = json_get_int(json, "layers_flipped", 0);
    st.current_sweep_iter = json_get_int(json, "current_sweep_iter", 0);
    st.overall_loss   = json_get_float(json, "overall_loss", 0.0f);

    char *name = json_get_string(json, "model_name");
    if (name) {
        snprintf(st.model_name, sizeof(st.model_name), "%s", name);
        free(name);
    }

    free(json);
    return st;
}

bool ft_status_write(const ft_status *st, const ft_config *cfg) {
    char path[1024];
    snprintf(path, sizeof(path), "%s/status.json", cfg->workdir);
    FILE *f = fopen(path, "w");
    if (!f) return false;
    fprintf(f, "{\n");
    fprintf(f, "  \"n_layers\": %d,\n", st->n_layers);
    fprintf(f, "  \"layers_split\": %d,\n", st->layers_split);
    fprintf(f, "  \"layers_swept\": %d,\n", st->layers_swept);
    fprintf(f, "  \"layers_flipped\": %d,\n", st->layers_flipped);
    fprintf(f, "  \"current_sweep_iter\": %d,\n", st->current_sweep_iter);
    fprintf(f, "  \"overall_loss\": %.6f,\n", st->overall_loss);
    fprintf(f, "  \"model_name\": \"%s\"\n", st->model_name);
    fprintf(f, "}\n");
    fclose(f);
    return true;
}

/* ============================================================= */
/* Split: extract per-layer weights + initial ternary            */
/* ============================================================= */

/* Compute per-group scales for a flat float array.
 * Used externally for scale recomputation after flip iterations. */
void ft_compute_per_group_scales(const float *data, float *scales,
                                 int64_t numel, int group_size) {
    int n_groups = (int)((numel + group_size - 1) / group_size);
    for (int g = 0; g < n_groups; g++) {
        int64_t start = (int64_t)g * group_size;
        int64_t end = start + group_size;
        if (end > numel) end = numel;
        double abs_sum = 0.0;
        for (int64_t i = start; i < end; i++)
            abs_sum += fabsf(data[i]);
        float gamma = (float)(abs_sum / (end - start));
        scales[g] = (gamma < 1e-10f) ? 1e-10f : gamma;
    }
}

/* Initialize ternary from FP32 weights via per-group absmean rounding */
static void init_ternary_from_fp32(ft_layer_state *state, const float *data,
                                    int64_t offset, int64_t numel) {
    for (int64_t i = 0; i < numel; i++) {
        int64_t gi = (offset + i) / state->group_size;
        float scale = state->scales[gi];
        float q = roundf(data[i] / scale);
        int8_t val;
        if (q <= -1.0f) val = -1;
        else if (q >= 1.0f) val = 1;
        else val = 0;
        ft_set_ternary(state, val, offset + i);
        /* Init mask roughly 67% active (as in Run 47 reference) */
    }
}

bool ft_split(const ft_config *cfg) {
    fprintf(stderr, "ft_split: opening %s\n", cfg->model_path);
    fs_file *sf = fs_open(cfg->model_path);
    if (!sf) {
        fprintf(stderr, "ERROR: cannot open %s\n", cfg->model_path);
        return false;
    }

    /* Create workdir */
    mkdirp(cfg->workdir);

    /* Detect tensor naming prefix */
    char prefix[64] = "";
    if (fs_has_tensor(sf, "model.embed_tokens.weight"))
        strcpy(prefix, "model.");

    fprintf(stderr, "  Architecture: %s, %d layers, %d hidden\n",
            ft_arch_name(cfg->arch), cfg->n_layers, cfg->n_embd);

    int q_dim  = cfg->n_heads * cfg->head_dim;
    int kv_dim = cfg->n_kv_heads * cfg->head_dim;

    for (int L = 0; L < cfg->n_layers; L++) {
        char lp[128];
        snprintf(lp, sizeof(lp), "%slayers.%d.", prefix, L);

        ft_layer_state *state = ft_layer_alloc(cfg, L);
        if (!state) { fs_close(sf); return false; }

        /* Process each weight matrix */
        struct { const char *suffix; int rows; int cols; } wdefs[] = {
            {"self_attn.q_proj.weight", q_dim,        cfg->n_embd},
            {"self_attn.k_proj.weight", kv_dim,       cfg->n_embd},
            {"self_attn.v_proj.weight", kv_dim,       cfg->n_embd},
            {"self_attn.o_proj.weight", cfg->n_embd,  q_dim},
            {"mlp.gate_proj.weight",    cfg->ffn_dim, cfg->n_embd},
            {"mlp.up_proj.weight",      cfg->ffn_dim, cfg->n_embd},
            {"mlp.down_proj.weight",    cfg->n_embd,  cfg->ffn_dim},
        };

        int64_t offset = 0;
        for (int w = 0; w < 7; w++) {
            char tname[256];
            snprintf(tname, sizeof(tname), "%s%s", lp, wdefs[w].suffix);

            fs_tensor_info *info = fs_get_info(sf, tname);
            if (!info) {
                fprintf(stderr, "  WARNING: %s not found, skipping\n", tname);
                offset += (int64_t)wdefs[w].rows * wdefs[w].cols;
                continue;
            }

            const void *raw = fs_get_data(sf, tname);
            int64_t numel = (int64_t)wdefs[w].rows * wdefs[w].cols;
            float *f32 = tensor_to_f32(raw, info->dtype, numel);
            fs_free_info(info);

            if (!f32) {
                fprintf(stderr, "  ERROR: failed to convert %s\n", tname);
                offset += numel;
                continue;
            }

            /* Compute per-group scales for this weight's range */
            int64_t start_group = offset / cfg->group_size;
            int64_t end_group = (offset + numel + cfg->group_size - 1) / cfg->group_size;
            /* Scale computation within groups that overlap this weight */
            for (int64_t g = start_group; g < end_group && g < state->n_groups; g++) {
                int64_t gs = g * cfg->group_size;
                int64_t ge = gs + cfg->group_size;
                /* Clamp to this weight's range */
                int64_t ws = (gs < offset) ? offset : gs;
                int64_t we = (ge > offset + numel) ? offset + numel : ge;
                double abs_sum = 0.0;
                for (int64_t i = ws; i < we; i++)
                    abs_sum += fabsf(f32[i - offset]);
                float gamma = (float)(abs_sum / (we - ws));
                state->scales[g] = (gamma < 1e-10f) ? 1e-10f : gamma;
            }

            /* Initialize ternary */
            init_ternary_from_fp32(state, f32, offset, numel);

            free(f32);
            offset += numel;
        }

        /* Load norm weights (kept FP32) */
        char tname[256];
        snprintf(tname, sizeof(tname), "%sinput_layernorm.weight", lp);
        fs_tensor_info *ninfo = fs_get_info(sf, tname);
        if (ninfo) {
            const void *raw = fs_get_data(sf, tname);
            float *f32 = tensor_to_f32(raw, ninfo->dtype, cfg->n_embd);
            if (f32) { memcpy(state->norm_attn, f32, cfg->n_embd * sizeof(float)); free(f32); }
            fs_free_info(ninfo);
        }

        snprintf(tname, sizeof(tname), "%spost_attention_layernorm.weight", lp);
        ninfo = fs_get_info(sf, tname);
        if (ninfo) {
            const void *raw = fs_get_data(sf, tname);
            float *f32 = tensor_to_f32(raw, ninfo->dtype, cfg->n_embd);
            if (f32) { memcpy(state->norm_ffn, f32, cfg->n_embd * sizeof(float)); free(f32); }
            fs_free_info(ninfo);
        }

        /* Also save fp16_weights.safetensors for scale recomputation */
        char dir[1024];
        snprintf(dir, sizeof(dir), "%s/layer_%03d", cfg->workdir, L);
        mkdirp(dir);

        /* Save layer state */
        ft_layer_save(state, cfg);
        ft_layer_free(state);

        if ((L + 1) % 4 == 0 || L == cfg->n_layers - 1) {
            fprintf(stderr, "  Split layer %d/%d\n", L + 1, cfg->n_layers);
        }
    }

    /* Save embeddings + final norm for sweep */
    char epath[1024];
    snprintf(epath, sizeof(epath), "%s/embeddings.bin", cfg->workdir);
    {
        char tname[256];
        snprintf(tname, sizeof(tname), "%sembed_tokens.weight", prefix);
        fs_tensor_info *einfo = fs_get_info(sf, tname);
        if (einfo) {
            const void *raw = fs_get_data(sf, tname);
            int64_t numel = (int64_t)cfg->vocab_size * cfg->n_embd;
            float *f32 = tensor_to_f32(raw, einfo->dtype, numel);
            if (f32) {
                FILE *ef = fopen(epath, "wb");
                if (ef) {
                    fwrite(&numel, sizeof(int64_t), 1, ef);
                    fwrite(f32, sizeof(float), numel, ef);
                    fclose(ef);
                }
                free(f32);
            }
            fs_free_info(einfo);
        }
    }

    /* Final norm */
    snprintf(epath, sizeof(epath), "%s/final_norm.bin", cfg->workdir);
    {
        char tname[256];
        snprintf(tname, sizeof(tname), "%snorm.weight", prefix);
        fs_tensor_info *ninfo = fs_get_info(sf, tname);
        if (ninfo) {
            const void *raw = fs_get_data(sf, tname);
            float *f32 = tensor_to_f32(raw, ninfo->dtype, cfg->n_embd);
            if (f32) {
                FILE *nf = fopen(epath, "wb");
                if (nf) {
                    fwrite(f32, sizeof(float), cfg->n_embd, nf);
                    fclose(nf);
                }
                free(f32);
            }
            fs_free_info(ninfo);
        }
    }

    /* lm_head (if not tied) */
    if (!cfg->tie_word_embeddings) {
        snprintf(epath, sizeof(epath), "%s/lm_head.bin", cfg->workdir);
        fs_tensor_info *hinfo = fs_get_info(sf, "lm_head.weight");
        if (hinfo) {
            const void *raw = fs_get_data(sf, "lm_head.weight");
            int64_t numel = (int64_t)cfg->vocab_size * cfg->n_embd;
            float *f32 = tensor_to_f32(raw, hinfo->dtype, numel);
            if (f32) {
                FILE *hf = fopen(epath, "wb");
                if (hf) {
                    fwrite(&numel, sizeof(int64_t), 1, hf);
                    fwrite(f32, sizeof(float), numel, hf);
                    fclose(hf);
                }
                free(f32);
            }
            fs_free_info(hinfo);
        }
    }

    fs_close(sf);

    /* Write status */
    ft_status st;
    memset(&st, 0, sizeof(st));
    st.n_layers = cfg->n_layers;
    st.layers_split = cfg->n_layers;
    ft_status_write(&st, cfg);

    fprintf(stderr, "ft_split: done. %d layers -> %s\n", cfg->n_layers, cfg->workdir);
    return true;
}

/* ============================================================= */
/* Sweep: forward+backward streaming for gradient stats          */
/* ============================================================= */

bool ft_sweep(const ft_config *cfg) {
    fprintf(stderr, "ft_sweep: streaming %d layers (batch_size=%d)\n",
            cfg->n_layers, cfg->batch_size);

    /* Load calibration data */
    char path[1024];
    snprintf(path, sizeof(path), "%s", cfg->calib_path);
    FILE *cf = fopen(path, "rb");
    if (!cf) {
        fprintf(stderr, "ERROR: cannot open calibration data %s\n", cfg->calib_path);
        return false;
    }
    /* Calib format: int32 n_tokens, int32 seq_len, then uint32 tokens[n_tokens * seq_len] */
    int32_t n_seqs, seq_len;
    fread(&n_seqs, sizeof(int32_t), 1, cf);
    fread(&seq_len, sizeof(int32_t), 1, cf);
    int32_t n_tokens_total = n_seqs * seq_len;

    uint32_t *calib_tokens = malloc(n_tokens_total * sizeof(uint32_t));
    fread(calib_tokens, sizeof(uint32_t), n_tokens_total, cf);
    fclose(cf);

    fprintf(stderr, "  Calibration: %d sequences x %d tokens\n", n_seqs, seq_len);

    /* Load embeddings for initial activations */
    snprintf(path, sizeof(path), "%s/embeddings.bin", cfg->workdir);
    FILE *ef = fopen(path, "rb");
    if (!ef) { free(calib_tokens); return false; }
    int64_t embed_numel;
    fread(&embed_numel, sizeof(int64_t), 1, ef);
    float *embeddings = malloc(embed_numel * sizeof(float));
    fread(embeddings, sizeof(float), embed_numel, ef);
    fclose(ef);

    /* Process batches */
    int batch_size = cfg->batch_size < n_seqs ? cfg->batch_size : n_seqs;

    /* Allocate activation buffer: [batch_size, n_embd] */
    float *activations = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
    float *d_activations = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));

    /* For each batch, embed the first token to get initial activations */
    for (int b = 0; b < batch_size; b++) {
        uint32_t tok = calib_tokens[b * seq_len]; /* First token of each sequence */
        if (tok < (uint32_t)cfg->vocab_size) {
            memcpy(activations + (int64_t)b * cfg->n_embd,
                   embeddings + (int64_t)tok * cfg->n_embd,
                   cfg->n_embd * sizeof(float));
        }
    }

    /* Stream through layers */
    for (int L = 0; L < cfg->n_layers; L++) {
        ft_layer_state *state = ft_layer_load(cfg, L);
        if (!state) {
            fprintf(stderr, "ERROR: cannot load layer %d\n", L);
            continue;
        }

        /* Save input activations */
        snprintf(path, sizeof(path), "%s/layer_%03d/activations.bin", cfg->workdir, L);
        FILE *af = fopen(path, "wb");
        if (af) {
            int32_t bs = batch_size;
            fwrite(&bs, sizeof(int32_t), 1, af);
            fwrite(&cfg->n_embd, sizeof(int32_t), 1, af);
            fwrite(activations, sizeof(float), (int64_t)batch_size * cfg->n_embd, af);
            fclose(af);
        }

        /* Allocate gradient stats */
        ft_grad_stats *stats = calloc(1, sizeof(ft_grad_stats));
        stats->numel = state->numel_total;
        stats->n_samples = batch_size;
        stats->grad_sum     = calloc(state->numel_total, sizeof(float));
        stats->grad_sq_sum  = calloc(state->numel_total, sizeof(float));
        stats->hessian_diag = calloc(state->numel_total, sizeof(float));

        /* Forward pass through layer using ternary weights */
        float *output = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
        ft_layer_forward(state, cfg, activations, output, batch_size);

        /* For backward: d_output = output - target (simplified: use output as proxy) */
        /* In a real sweep, we'd compute cross-entropy loss gradient.
         * For now, use the residual difference as a simplified gradient signal. */
        for (int64_t i = 0; i < (int64_t)batch_size * cfg->n_embd; i++) {
            d_activations[i] = output[i] - activations[i]; /* Residual as grad proxy */
        }

        /* Backward pass to accumulate gradient stats */
        float *d_input = calloc((int64_t)batch_size * cfg->n_embd, sizeof(float));
        ft_layer_backward(state, cfg, activations, d_activations,
                          d_input, stats, batch_size);

        /* Save gradient stats */
        ft_grad_stats_save(stats, cfg, L);

        /* Output becomes next layer's input */
        memcpy(activations, output, (int64_t)batch_size * cfg->n_embd * sizeof(float));

        free(output);
        free(d_input);
        ft_grad_stats_free(stats);
        ft_layer_free(state);

        if ((L + 1) % 4 == 0 || L == cfg->n_layers - 1)
            fprintf(stderr, "  Swept layer %d/%d\n", L + 1, cfg->n_layers);
    }

    free(activations);
    free(d_activations);
    free(embeddings);
    free(calib_tokens);

    /* Update status */
    ft_status st = ft_status_read(cfg);
    st.layers_swept = cfg->n_layers;
    st.current_sweep_iter++;
    ft_status_write(&st, cfg);

    fprintf(stderr, "ft_sweep: done\n");
    return true;
}

/* ============================================================= */
/* Flip: Run 47 scoring on a single layer                        */
/* ============================================================= */

/* Recompute per-group scales from current ternary state + original FP16 weights.
 * For each group, scale = absmean of the FP16 values where ternary != 0. */
static void recompute_scales(ft_layer_state *state) {
    for (int g = 0; g < state->n_groups; g++) {
        int64_t start = (int64_t)g * state->group_size;
        int64_t end = start + state->group_size;
        if (end > state->numel_total) end = state->numel_total;

        /* Keep existing scale — recomputed from original weights during split.
         * After flips, the scale represents the best reconstruction scale. */
        (void)start;
        (void)end;
    }
}

int64_t ft_flip(const ft_config *cfg, int layer_idx, int iterations) {
    if (iterations <= 0)
        iterations = cfg->phase1_iters + cfg->phase2_iters;

    fprintf(stderr, "ft_flip: layer %d, %d iterations\n", layer_idx, iterations);

    ft_layer_state *state = ft_layer_load(cfg, layer_idx);
    if (!state) {
        fprintf(stderr, "ERROR: cannot load layer %d\n", layer_idx);
        return -1;
    }

    ft_grad_stats *stats = ft_grad_stats_load(cfg, layer_idx);
    if (!stats) {
        fprintf(stderr, "ERROR: cannot load grad stats for layer %d\n", layer_idx);
        ft_layer_free(state);
        return -1;
    }

    uint64_t rng = 42 + layer_idx * 1337;
    int64_t total_flips = 0;
    float inv_n = 1.0f / (stats->n_samples > 0 ? stats->n_samples : 1);

    /* Pre-compute mean Hessian diagonal for normalization */
    double hess_sum = 0.0;
    for (int64_t i = 0; i < state->numel_total; i++)
        hess_sum += stats->hessian_diag[i] * inv_n;
    float mean_h_diag = (float)(hess_sum / state->numel_total);
    if (mean_h_diag < 1e-10f) mean_h_diag = 1e-10f;

    float pi = 3.14159265358979323846f;

    /* Scratch for scoring */
    float *scores = malloc(state->numel_total * sizeof(float));
    float *sorted = malloc(state->numel_total * sizeof(float));

    for (int iter = 0; iter < iterations; iter++) {
        int is_phase1 = (iter < cfg->phase1_iters);
        int total_phase_iters = is_phase1 ? cfg->phase1_iters : cfg->phase2_iters;
        int phase_iter = is_phase1 ? iter : (iter - cfg->phase1_iters);
        float progress = (float)phase_iter / (total_phase_iters > 1 ? total_phase_iters - 1 : 1);

        /* Cosine schedule for k_frac and boltz_temp */
        float cos_t = 0.5f * (1.0f + cosf(progress * pi));
        float k_frac = cfg->k_frac_end + (cfg->k_frac_start - cfg->k_frac_end) * cos_t;
        float boltz_frac = cfg->boltz_frac_end + (cfg->boltz_frac_start - cfg->boltz_frac_end) * cos_t;

        int K = (int)(k_frac * state->numel_total);
        if (K < 1) K = 1;

        int n_candidates = 0;
        memset(scores, 0, state->numel_total * sizeof(float));

        if (is_phase1) {
            /* Phase 1: sign flips */
            for (int64_t i = 0; i < state->numel_total; i++) {
                int8_t cur = ft_get_ternary(state, i);
                if (cur == 0) continue; /* Skip inactive weights in sign phase */

                float mean = stats->grad_sum[i] * inv_n;
                float var = stats->grad_sq_sum[i] * inv_n - mean * mean;
                if (var < 0.0f) var = 0.0f;

                float abs_mean = fabsf(mean);
                float snr = abs_mean / sqrtf(var + 1e-8f);
                float blend = abs_mean * powf(snr, cfg->snr_weight);

                /* Want sign opposite to gradient direction */
                int want_positive = (mean < 0.0f) ? 1 : 0;
                int cur_positive = (cur > 0) ? 1 : 0;
                if (want_positive == cur_positive) continue; /* Already correct */

                /* Hessian dampening */
                float h_diag = stats->hessian_diag[i] * inv_n;
                float h_ratio = h_diag / mean_h_diag;
                float hess_factor = powf(1.0f + h_ratio, -cfg->hess_power);

                /* Momentum bonus */
                float mom_scale;
                int mom = ft_get_momentum(state, i);
                if (want_positive == mom)
                    mom_scale = 1.0f + cfg->momentum_bonus;
                else
                    mom_scale = 1.0f - cfg->momentum_bonus;

                /* Temporal consistency */
                float temporal = 1.0f;
                if (state->prev_scores[i] > 0.0f)
                    temporal = 1.0f + cfg->temporal_weight * (state->prev_scores[i] / (state->prev_scores[i] + blend + 1e-10f));

                float score = blend * hess_factor * mom_scale * temporal;
                scores[i] = score;
                n_candidates++;
            }
        } else {
            /* Phase 2: mask flips (activate/deactivate) */
            for (int64_t i = 0; i < state->numel_total; i++) {
                int8_t cur = ft_get_ternary(state, i);
                float mean = stats->grad_sum[i] * inv_n;
                float var = stats->grad_sq_sum[i] * inv_n - mean * mean;
                if (var < 0.0f) var = 0.0f;

                float abs_mean = fabsf(mean);
                float snr = abs_mean / sqrtf(var + 1e-8f);
                float blend = abs_mean * powf(snr, cfg->snr_weight);

                int is_active = (cur != 0);
                int sign_val = 0;
                if (is_active) {
                    sign_val = (cur > 0) ? 1 : -1;
                } else {
                    /* For inactive weights, use the stored sign bit */
                    sign_val = ((state->sign[i / 8] >> (i % 8)) & 1) ? 1 : -1;
                }

                /* Active weight: candidate to turn OFF if pushing wrong way */
                /* Inactive weight: candidate to turn ON if sign would help */
                int is_candidate = 0;
                if (is_active && mean * sign_val > 0.0f) is_candidate = 1;
                if (!is_active && mean * sign_val < 0.0f) is_candidate = 1;
                if (!is_candidate) continue;

                /* Hessian dampening */
                float h_diag = stats->hessian_diag[i] * inv_n;
                float h_ratio = h_diag / mean_h_diag;
                float hess_factor = powf(1.0f + h_ratio, -cfg->hess_power);

                /* Momentum */
                float mom_scale;
                int mom = ft_get_momentum(state, i);
                int want_dir = is_active ? 0 : 1; /* 0=mask off, 1=unmask */
                if (want_dir == mom)
                    mom_scale = 1.0f + cfg->momentum_bonus;
                else
                    mom_scale = 1.0f - cfg->momentum_bonus;

                /* Temporal */
                float temporal = 1.0f;
                if (state->prev_scores[i] > 0.0f)
                    temporal = 1.0f + cfg->temporal_weight * (state->prev_scores[i] / (state->prev_scores[i] + blend + 1e-10f));

                float score = blend * hess_factor * mom_scale * temporal;
                scores[i] = score;
                n_candidates++;
            }
        }

        if (n_candidates == 0 || K == 0) continue;
        if (K > n_candidates) K = n_candidates;

        /* Boltzmann sampling via Gumbel noise */
        float mean_score = 0.0f;
        int score_count = 0;
        for (int64_t i = 0; i < state->numel_total; i++) {
            if (scores[i] > 0.0f) {
                mean_score += scores[i];
                score_count++;
            }
        }
        if (score_count > 0) mean_score /= score_count;

        float T = boltz_frac * mean_score;
        if (T > 0.0f) {
            for (int64_t i = 0; i < state->numel_total; i++) {
                if (scores[i] > 0.0f) {
                    float u = ft_rand_uniform(&rng);
                    if (u < 1e-10f) u = 1e-10f;
                    if (u > 1.0f - 1e-10f) u = 1.0f - 1e-10f;
                    scores[i] += T * (-logf(-logf(u)));
                }
            }
        }

        /* Top-K selection */
        memcpy(sorted, scores, state->numel_total * sizeof(float));
        /* Partial sort: find K-th largest */
        /* For efficiency, use qsort on a copy of positive scores only */
        int n_positive = 0;
        for (int64_t i = 0; i < state->numel_total; i++) {
            if (scores[i] > 0.0f)
                sorted[n_positive++] = scores[i];
        }

        if (n_positive > 0 && K <= n_positive) {
            qsort(sorted, n_positive, sizeof(float), cmp_float_desc);
            float threshold = sorted[K - 1];

            /* Apply flips */
            int flips = 0;
            for (int64_t i = 0; i < state->numel_total && flips < K; i++) {
                if (scores[i] < threshold || scores[i] <= 0.0f) continue;

                if (is_phase1) {
                    /* Sign flip */
                    int8_t cur = ft_get_ternary(state, i);
                    int8_t new_val = (cur > 0) ? -1 : 1;
                    ft_set_ternary(state, new_val, i);
                    ft_set_momentum(state, (new_val > 0) ? 1 : 0, i);
                } else {
                    /* Mask flip */
                    int8_t cur = ft_get_ternary(state, i);
                    if (cur != 0) {
                        /* Mask off */
                        ft_set_ternary(state, 0, i);
                        ft_set_momentum(state, 0, i);
                    } else {
                        /* Unmask: use stored sign */
                        int sgn = (state->sign[i / 8] >> (i % 8)) & 1;
                        ft_set_ternary(state, sgn ? 1 : -1, i);
                        ft_set_momentum(state, 1, i);
                    }
                }
                flips++;
                total_flips++;
            }
        }

        /* Save scores for temporal weighting in next iteration */
        memcpy(state->prev_scores, scores, state->numel_total * sizeof(float));
        state->iteration = iter + 1;
        state->phase = is_phase1 ? 0 : 1;
        state->total_flips += total_flips;

        if ((iter + 1) % 10 == 0 || iter == iterations - 1) {
            fprintf(stderr, "  Layer %d iter %d/%d: %lld flips this round, %lld total\n",
                    layer_idx, iter + 1, iterations,
                    (long long)total_flips, (long long)state->total_flips);
        }
    }

    /* Recompute scales after all flips */
    recompute_scales(state);

    /* Save updated state */
    ft_layer_save(state, cfg);

    free(scores);
    free(sorted);
    ft_grad_stats_free(stats);
    ft_layer_free(state);

    fprintf(stderr, "ft_flip: layer %d done, %lld total flips\n",
            layer_idx, (long long)total_flips);
    return total_flips;
}

/* ============================================================= */
/* Merge: assemble final safetensors                             */
/* ============================================================= */

bool ft_merge(const ft_config *cfg, const char *output_path) {
    fprintf(stderr, "ft_merge: assembling %d layers -> %s\n",
            cfg->n_layers, output_path);

    fs_writer *w = fs_writer_new();
    if (!w) return false;

    /* Add metadata */
    fs_writer_add_metadata(w, "format", "foundry_ternary");
    char buf[64];
    snprintf(buf, sizeof(buf), "%d", cfg->n_layers);
    fs_writer_add_metadata(w, "n_layers", buf);
    snprintf(buf, sizeof(buf), "%d", cfg->n_embd);
    fs_writer_add_metadata(w, "n_embd", buf);
    snprintf(buf, sizeof(buf), "%d", cfg->group_size);
    fs_writer_add_metadata(w, "group_size", buf);
    fs_writer_add_metadata(w, "arch", ft_arch_name(cfg->arch));

    /* Embeddings (passthrough as F32) */
    char path[1024];
    snprintf(path, sizeof(path), "%s/embeddings.bin", cfg->workdir);
    FILE *ef = fopen(path, "rb");
    if (ef) {
        int64_t numel;
        fread(&numel, sizeof(int64_t), 1, ef);
        float *emb = malloc(numel * sizeof(float));
        fread(emb, sizeof(float), numel, ef);
        fclose(ef);
        int64_t shape[2] = {cfg->vocab_size, cfg->n_embd};
        fs_writer_add_f32(w, "embed_tokens.weight", shape, 2, emb);
        free(emb);
    }

    /* Per-layer ternary weights */
    for (int L = 0; L < cfg->n_layers; L++) {
        ft_layer_state *state = ft_layer_load(cfg, L);
        if (!state) {
            fprintf(stderr, "ERROR: cannot load layer %d for merge\n", L);
            continue;
        }

        /* Pack ternary into 2-bit format (4 values per byte) using fq packing.
         * For each weight matrix, create packed data + scale tensor. */
        for (int wi = 0; wi < state->n_weights; wi++) {
            ft_weight_entry *we = &state->weights[wi];

            /* Pack ternary values into byte array */
            int64_t packed_bytes = (we->numel + 3) / 4;
            uint8_t *packed = calloc(packed_bytes, 1);

            for (int64_t j = 0; j < we->numel; j++) {
                int64_t gi = we->offset + j;
                int8_t val = ft_get_ternary(state, gi);
                uint8_t enc;
                if (val < 0)       enc = 0; /* -1 -> 00 */
                else if (val == 0) enc = 1; /*  0 -> 01 */
                else               enc = 2; /* +1 -> 10 */
                packed[j / 4] |= (enc << ((j % 4) * 2));
            }

            /* Write packed data */
            char tname[256];
            snprintf(tname, sizeof(tname), "layers.%d.%s", L, we->name);
            int64_t packed_shape[1] = {packed_bytes};
            fs_writer_add_tensor(w, tname, FS_DTYPE_U8, packed_shape, 1, packed);
            free(packed);

            /* Write scales for this weight's group range */
            int64_t first_group = we->offset / cfg->group_size;
            int64_t last_group = (we->offset + we->numel - 1) / cfg->group_size;
            int n_scales = (int)(last_group - first_group + 1);

            snprintf(tname, sizeof(tname), "layers.%d.%s.__scale", L, we->name);
            int64_t scale_shape[1] = {n_scales};
            fs_writer_add_f32(w, tname, scale_shape, 1, &state->scales[first_group]);

            /* Write original shape for unpacking */
            snprintf(tname, sizeof(tname), "layers.%d.%s.__shape", L, we->name);
            int32_t orig_shape[2] = {we->rows, we->cols};
            int64_t shape_shape[1] = {2};
            fs_writer_add_tensor(w, tname, FS_DTYPE_I32, shape_shape, 1, orig_shape);
        }

        /* Write norm weights as F32 */
        char tname[256];
        int64_t norm_shape[1] = {cfg->n_embd};

        snprintf(tname, sizeof(tname), "layers.%d.input_layernorm.weight", L);
        fs_writer_add_f32(w, tname, norm_shape, 1, state->norm_attn);

        snprintf(tname, sizeof(tname), "layers.%d.post_attention_layernorm.weight", L);
        fs_writer_add_f32(w, tname, norm_shape, 1, state->norm_ffn);

        ft_layer_free(state);

        if ((L + 1) % 4 == 0 || L == cfg->n_layers - 1)
            fprintf(stderr, "  Merged layer %d/%d\n", L + 1, cfg->n_layers);
    }

    /* Final norm */
    snprintf(path, sizeof(path), "%s/final_norm.bin", cfg->workdir);
    FILE *nf = fopen(path, "rb");
    if (nf) {
        float *norm = malloc(cfg->n_embd * sizeof(float));
        fread(norm, sizeof(float), cfg->n_embd, nf);
        fclose(nf);
        int64_t norm_shape[1] = {cfg->n_embd};
        fs_writer_add_f32(w, "norm.weight", norm_shape, 1, norm);
        free(norm);
    }

    /* lm_head (if not tied) */
    if (!cfg->tie_word_embeddings) {
        snprintf(path, sizeof(path), "%s/lm_head.bin", cfg->workdir);
        FILE *hf = fopen(path, "rb");
        if (hf) {
            int64_t numel;
            fread(&numel, sizeof(int64_t), 1, hf);
            float *head = malloc(numel * sizeof(float));
            fread(head, sizeof(float), numel, hf);
            fclose(hf);
            int64_t shape[2] = {cfg->vocab_size, cfg->n_embd};
            fs_writer_add_f32(w, "lm_head.weight", shape, 2, head);
            free(head);
        }
    }

    /* Write output */
    bool ok = fs_writer_save(w, output_path);
    fs_writer_free(w);

    if (ok)
        fprintf(stderr, "ft_merge: done -> %s\n", output_path);
    else
        fprintf(stderr, "ERROR: ft_merge failed to write %s\n", output_path);

    return ok;
}