// requant_trunk.c — Surgical GGUF rewriter for Kimi-K3 UD-Q4_K_XL.
//
// Requantizes ONLY the Q8_0 "trunk" tensors (attention, shared experts, output,
// token embedding) from Q8_0 -> Q4_K. Byte-copies everything else unchanged:
//   - MXFP4 routed-expert tensors (ffn_*_exps)  : QAT-native, must NOT requant
//   - F32/BF16 norm & bias tensors              : tiny, keep full precision
//
// Why not llama-quantize? In --allow-requantize mode it forces EVERY non-overridden
// tensor to the positional type, which would dequant->requant the MXFP4 experts and
// destroy their QAT calibration. We must byte-preserve the experts.
//
// Processes one shard at a time (split-in = split-out, like --keep-split).
//
// Build (on box):
//   gcc -O2 -o requant_trunk requant_trunk.c \
//       -I/root/llama.cpp/ggml/include -I/root/llama.cpp/ggml/src \
//       /root/llama.cpp/build/ggml/src/libggml-base.a \
//       /root/llama.cpp/build/ggml/src/libggml-cpu.a \
//       /root/llama.cpp/build/ggml/src/libggml.a -lm -lpthread
//
// Usage: requant_trunk <in_shard.gguf> <out_shard.gguf>
//
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

#include "ggml.h"
#include "ggml-quants.h"

#define GGUF_MAGIC 0x46554747  // "GGUF"
#define ALIGNMENT 32

// ---- little-endian read helpers ----
static uint32_t rd_u32(FILE *f){ uint32_t v; if(fread(&v,4,1,f)!=1){fprintf(stderr,"EOF u32\n");exit(1);} return v; }
static uint64_t rd_u64(FILE *f){ uint64_t v; if(fread(&v,8,1,f)!=1){fprintf(stderr,"EOF u64\n");exit(1);} return v; }

// ---- write helpers (dynamic buffer for header) ----
typedef struct { uint8_t *data; size_t len, cap; } Buf;
static void buf_put(Buf *b, const void *p, size_t n){
    if (b->len + n > b->cap){ b->cap = (b->len + n)*2 + 1024; b->data = realloc(b->data, b->cap); if(!b->data){fprintf(stderr,"oom\n");exit(1);} }
    memcpy(b->data + b->len, p, n); b->len += n;
}
static void buf_u32(Buf *b, uint32_t v){ buf_put(b, &v, 4); }
static void buf_u64(Buf *b, uint64_t v){ buf_put(b, &v, 8); }

// skip a KV value of given vtype in input file
static void skip_kv_value(FILE *f, uint32_t vtype){
    switch(vtype){
        case 0: case 1: case 7: fseek(f,1,SEEK_CUR); break;
        case 2: case 3: fseek(f,2,SEEK_CUR); break;
        case 4: case 5: case 6: fseek(f,4,SEEK_CUR); break;
        case 10: case 11: case 12: fseek(f,8,SEEK_CUR); break;
        case 8: { uint64_t l=rd_u64(f); fseek(f,(long)l,SEEK_CUR); } break;
        case 9: {
            uint32_t at=rd_u32(f); uint64_t al=rd_u64(f);
            if (at==8){ for(uint64_t i=0;i<al;i++){ uint64_t sl=rd_u64(f); fseek(f,(long)sl,SEEK_CUR);} }
            else {
                int sz; switch(at){case 0:case 1:case 7:sz=1;break;case 2:case 3:sz=2;break;case 4:case 5:case 6:sz=4;break;default:sz=8;}
                fseek(f,(long)(al*sz),SEEK_CUR);
            }
        } break;
        default: fprintf(stderr,"bad vtype %u\n",vtype); exit(1);
    }
}

// copy a KV value verbatim from input file into buffer
static void copy_kv_value(FILE *f, Buf *out, uint32_t vtype){
    buf_u32(out, vtype);
    switch(vtype){
        case 0: case 1: case 7: { uint8_t b[1]; if(fread(b,1,1,f)!=1){exit(1);} buf_put(out,b,1);} break;
        case 2: case 3: { uint8_t b[2]; if(fread(b,2,1,f)!=1){exit(1);} buf_put(out,b,2);} break;
        case 4: case 5: case 6: { uint8_t b[4]; if(fread(b,4,1,f)!=1){exit(1);} buf_put(out,b,4);} break;
        case 10: case 11: case 12: { uint8_t b[8]; if(fread(b,8,1,f)!=1){exit(1);} buf_put(out,b,8);} break;
        case 8: { uint64_t l=rd_u64(f); buf_u64(out,l); uint8_t *tmp=malloc(l); if(fread(tmp,l,1,f)!=1){exit(1);} buf_put(out,tmp,l); free(tmp);} break;
        case 9: {
            uint32_t at=rd_u32(f); uint64_t al=rd_u64(f);
            buf_u32(out,at); buf_u64(out,al);
            if (at==8){ for(uint64_t i=0;i<al;i++){ uint64_t sl=rd_u64(f); buf_u64(out,sl); uint8_t *tmp=malloc(sl); if(fread(tmp,sl,1,f)!=1){exit(1);} buf_put(out,tmp,sl); free(tmp);} }
            else { int sz; switch(at){case 0:case 1:case 7:sz=1;break;case 2:case 3:sz=2;break;case 4:case 5:case 6:sz=4;break;default:sz=8;}
                   uint64_t nb=al*sz; uint8_t *tmp=malloc(nb); if(fread(tmp,nb,1,f)!=1){exit(1);} buf_put(out,tmp,nb); free(tmp); }
        } break;
        default: fprintf(stderr,"bad vtype %u\n",vtype); exit(1);
    }
}

// decide whether a tensor name is a routed expert (MXFP4, byte-copy)
static bool is_routed_expert(const char *name){
    return strstr(name, "_exps.") != NULL;   // ffn_gate_exps / ffn_up_exps / ffn_down_exps
}

int main(int argc, char **argv){
    if (argc < 3){ fprintf(stderr,"usage: %s <in> <out>\n", argv[0]); return 1; }
    const char *fin_name = argv[1], *fout_name = argv[2];
    FILE *fin = fopen(fin_name, "rb"); if(!fin){ perror("open in"); return 1; }
    FILE *fout = fopen(fout_name, "wb"); if(!fout){ perror("open out"); return 1; }

    // ---- header ----
    uint32_t magic = rd_u32(fin);
    if (magic != GGUF_MAGIC){ fprintf(stderr,"not GGUF\n"); return 1; }
    uint32_t version = rd_u32(fin);
    uint64_t n_tensors = rd_u64(fin);
    uint64_t n_kv = rd_u64(fin);
    fprintf(stderr,"[requant] %s: ver=%u tensors=%llu kv=%llu\n", fin_name, version,
            (unsigned long long)n_tensors, (unsigned long long)n_kv);

    Buf hdr = {0};
    buf_u32(&hdr, magic); buf_u32(&hdr, version);
    buf_u64(&hdr, n_tensors); buf_u64(&hdr, n_kv);

    // ---- copy KV verbatim ----
    for (uint64_t i=0;i<n_kv;i++){
        uint64_t klen = rd_u64(fin);
        uint8_t *key = malloc(klen); if(fread(key,klen,1,fin)!=1){exit(1);}
        uint32_t vtype = rd_u32(fin);
        buf_u64(&hdr, klen); buf_put(&hdr, key, klen);
        copy_kv_value(fin, &hdr, vtype);
        free(key);
    }

    // ---- tensor infos: read all, decide new type, record ----
    typedef struct {
        char *name; uint32_t n_dims; uint64_t dims[8]; uint32_t orig_type; uint32_t new_type;
        uint64_t orig_offset;   // offset within data section
        uint64_t orig_size;     // bytes in source
        uint64_t new_size;      // bytes in output
        bool requant;           // true => q8_0 -> q4_K
    } TInfo;
    TInfo *tis = calloc(n_tensors, sizeof(TInfo));

    uint64_t n_requant=0, n_copy=0, bytes_in=0, bytes_out=0;
    for (uint64_t i=0;i<n_tensors;i++){
        TInfo *ti = &tis[i];
        uint64_t nl = rd_u64(fin);
        ti->name = malloc(nl+1); if(fread(ti->name,nl,1,fin)!=1){exit(1);} ti->name[nl]=0;
        ti->n_dims = rd_u32(fin);
        uint64_t nelem = 1;
        for (uint32_t d=0; d<ti->n_dims; d++){ ti->dims[d]=rd_u64(fin); nelem *= ti->dims[d]; }
        ti->orig_type = rd_u32(fin);
        ti->orig_offset = rd_u64(fin);

        enum ggml_type ot = (enum ggml_type)ti->orig_type;
        // total tensor bytes = row_size(type, ne[0]) * ne[1]*ne[2]*ne[3]
        int64_t nrow_mult = 1;
        for (uint32_t d=1; d<ti->n_dims; d++) nrow_mult *= (int64_t)ti->dims[d];
        ti->orig_size = ggml_row_size(ot, (int64_t)ti->dims[0]) * nrow_mult;

        // decision: requant Q8_0 non-expert tensors to Q4_K; copy everything else
        if (ot == GGML_TYPE_Q8_0 && !is_routed_expert(ti->name)){
            ti->requant = true;
            ti->new_type = GGML_TYPE_Q4_K;
            ti->new_size = ggml_row_size(GGML_TYPE_Q4_K, (int64_t)ti->dims[0]) * nrow_mult;
            n_requant++;
        } else {
            ti->requant = false;
            ti->new_type = ti->orig_type;
            ti->new_size = ti->orig_size;
            n_copy++;
        }
        // sanity: Q4_K requires the row dim divisible by QK_K(256). If not, copy instead.
        if (ti->requant && (ti->dims[0] % 256) != 0){
            fprintf(stderr,"  [warn] %s dims[0]=%llu not mult of 256, copying instead\n", ti->name,(unsigned long long)ti->dims[0]);
            ti->requant=false; ti->new_type=ti->orig_type; ti->new_size=ti->orig_size; n_requant--; n_copy++;
        }
        bytes_in += ti->orig_size; bytes_out += ti->new_size;
    }
    fprintf(stderr,"[requant] tensors: %llu requant(q8_0->q4_K), %llu copy | data %.1f GB -> %.1f GB\n",
            (unsigned long long)n_requant,(unsigned long long)n_copy, bytes_in/1e9, bytes_out/1e9);

    // ---- write tensor infos with new types & recomputed offsets ----
    // first compute new data offsets (sequential, aligned per-tensor to ALIGNMENT within data section)
    uint64_t *new_off = calloc(n_tensors, sizeof(uint64_t));
    uint64_t cur = 0;
    for (uint64_t i=0;i<n_tensors;i++){
        // each tensor's data starts aligned relative to data section start
        new_off[i] = cur;
        uint64_t sz = tis[i].new_size;
        // pad to ALIGNMENT after each tensor
        cur += (sz + ALIGNMENT - 1)/ALIGNMENT*ALIGNMENT;
    }

    Buf tinfos = {0};
    for (uint64_t i=0;i<n_tensors;i++){
        TInfo *ti=&tis[i];
        uint64_t nl=strlen(ti->name);
        buf_u64(&tinfos, nl); buf_put(&tinfos, ti->name, nl);
        buf_u32(&tinfos, ti->n_dims);
        for (uint32_t d=0; d<ti->n_dims; d++) buf_u64(&tinfos, ti->dims[d]);
        buf_u32(&tinfos, ti->new_type);
        buf_u64(&tinfos, new_off[i]);
    }

    // ---- emit header + tinfos, pad to data start ----
    uint64_t head_len = hdr.len + tinfos.len;
    uint64_t data_start = (head_len + ALIGNMENT - 1)/ALIGNMENT*ALIGNMENT;
    fwrite(hdr.data, hdr.len, 1, fout);
    fwrite(tinfos.data, tinfos.len, 1, fout);
    for (uint64_t p=head_len; p<data_start; p++) fputc(0, fout);

    // ---- source data section start ----
    // we must know where source tensor data begins to seek by orig_offset.
    // Recompute: after reading all tensor infos from fin, current position = end of tinfos.
    uint64_t src_tinfos_end = (uint64_t)ftell(fin);
    uint64_t src_data_start = (src_tinfos_end + ALIGNMENT - 1)/ALIGNMENT*ALIGNMENT;

    // ---- process tensors: copy or requant ----
    float *fbuf = NULL; void *qbuf = NULL; size_t fbuf_n=0, qbuf_n=0;
    uint64_t done=0;
    for (uint64_t i=0;i<n_tensors;i++){
        TInfo *ti=&tis[i];
        // read source tensor data
        fseek(fin, (long)(src_data_start + ti->orig_offset), SEEK_SET);
        if (!ti->requant){
            // byte-copy
            uint8_t *tmp = malloc(ti->orig_size);
            if (fread(tmp, ti->orig_size, 1, fin)!=1){ fprintf(stderr,"read tensor %s fail\n",ti->name); return 1; }
            // write at new offset
            uint64_t pos = data_start + new_off[i];
            fseek(fout, (long)pos, SEEK_SET);
            fwrite(tmp, ti->orig_size, 1, fout);
            free(tmp);
        } else {
            int64_t nelem = 1; for (uint32_t d=0; d<ti->n_dims; d++) nelem *= (int64_t)ti->dims[d];
            if ((size_t)nelem > fbuf_n){ fbuf = realloc(fbuf, nelem*sizeof(float)); fbuf_n=nelem; }
            if (ti->orig_size > qbuf_n){ qbuf = realloc(qbuf, ti->orig_size); qbuf_n=ti->orig_size; }
            // read q8_0
            if (fread(qbuf, ti->orig_size, 1, fin)!=1){ fprintf(stderr,"read q8 %s fail\n",ti->name); return 1; }
            // dequant q8_0 -> f32
            dequantize_row_q8_0((const block_q8_0*)qbuf, fbuf, nelem);
            // requant f32 -> q4_K (into a second region; reuse qbuf after? need separate)
            void *out = malloc(ti->new_size);
            quantize_row_q4_K_ref(fbuf, (block_q4_K*)out, nelem);
            uint64_t pos = data_start + new_off[i];
            fseek(fout, (long)pos, SEEK_SET);
            fwrite(out, ti->new_size, 1, fout);
            free(out);
        }
        done++;
        if (done % 50 == 0 || done==n_tensors){
            fprintf(stderr,"\r[requant] %llu/%llu tensors", (unsigned long long)done,(unsigned long long)n_tensors);
            fflush(stderr);
        }
    }
    fprintf(stderr,"\n[requant] wrote %s\n", fout_name);

    fclose(fin); fclose(fout);
    return 0;
}