/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
#include "quantization.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <float.h>

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

const char* fq_dtype_name(fq_dtype dtype) {
    switch (dtype) {
        case FQ_INT8:    return "INT8";
        case FQ_INT4:    return "INT4";
        case FQ_UINT8:   return "UINT8";
        case FQ_UINT4:   return "UINT4";
        case FQ_TERNARY: return "TERNARY";
        default:         return "UNKNOWN";
    }
}

float fq_dtype_bytes(fq_dtype dtype) {
    switch (dtype) {
        case FQ_INT8:
        case FQ_UINT8:   return 1.0f;
        case FQ_INT4:
        case FQ_UINT4:   return 0.5f;
        case FQ_TERNARY: return 0.25f;
        default:         return 0.0f;
    }
}

static int64_t compute_numel(int64_t* shape, int ndim) {
    int64_t n = 1;
    for (int i = 0; i < ndim; i++) {
        n *= shape[i];
    }
    return n;
}

/* Find min/max of data */
static void find_min_max(const float* data, int64_t n, float* min_val, float* max_val) {
    *min_val = FLT_MAX;
    *max_val = -FLT_MAX;

    for (int64_t i = 0; i < n; i++) {
        if (data[i] < *min_val) *min_val = data[i];
        if (data[i] > *max_val) *max_val = data[i];
    }
}

void fq_compute_scale(const float* data, int64_t n, fq_mode mode,
                      float* scale, int32_t* zero_point) {
    float min_val, max_val;
    find_min_max(data, n, &min_val, &max_val);

    if (mode == FQ_SYMMETRIC) {
        /* Symmetric: scale = max(|min|, |max|) / 127 */
        float abs_max = fmaxf(fabsf(min_val), fabsf(max_val));
        *scale = abs_max / 127.0f;
        *zero_point = 0;
    } else {
        /* Asymmetric: map [min, max] to [0, 255] for uint8, [-128, 127] for int8 */
        float range = max_val - min_val;
        *scale = range / 255.0f;
        *zero_point = (int32_t)roundf(-min_val / *scale);
    }

    /* Avoid division by zero */
    if (*scale < 1e-10f) {
        *scale = 1e-10f;
    }
}

size_t fq_memory_usage(fq_tensor* qt) {
    if (!qt) return 0;

    size_t data_size = 0;
    if (qt->dtype == FQ_INT8 || qt->dtype == FQ_UINT8) {
        data_size = qt->numel;
    } else if (qt->dtype == FQ_TERNARY) {
        /* Ternary: 4 values packed per byte (2 bits each) */
        data_size = (qt->numel + 3) / 4;
    } else {
        /* INT4: 2 values packed per byte */
        data_size = (qt->numel + 1) / 2;
    }

    size_t scale_size = qt->num_scales * sizeof(float);
    size_t zp_size = (qt->mode == FQ_ASYMMETRIC) ? qt->num_scales * sizeof(int32_t) : 0;

    return data_size + scale_size + zp_size;
}

/* ============================================================= */
/* Tensor allocation                                             */
/* ============================================================= */

static fq_tensor* fq_alloc(int64_t* shape, int ndim, fq_dtype dtype,
                          fq_mode mode, fq_granularity granularity, int num_scales) {
    fq_tensor* qt = calloc(1, sizeof(fq_tensor));
    if (!qt) return NULL;

    qt->ndim = ndim;
    qt->numel = compute_numel(shape, ndim);
    qt->dtype = dtype;
    qt->mode = mode;
    qt->granularity = granularity;
    qt->num_scales = num_scales;

    qt->shape = malloc(ndim * sizeof(int64_t));
    memcpy(qt->shape, shape, ndim * sizeof(int64_t));

    /* Allocate data */
    if (dtype == FQ_INT8 || dtype == FQ_UINT8) {
        qt->data = calloc(qt->numel, 1);
    } else if (dtype == FQ_TERNARY) {
        /* Ternary: pack 4 values per byte (2 bits each) */
        qt->data = calloc((qt->numel + 3) / 4, 1);
    } else {
        /* INT4: pack 2 values per byte */
        qt->data = calloc((qt->numel + 1) / 2, 1);
    }

    /* Allocate scales */
    qt->scales = calloc(num_scales, sizeof(float));

    /* Allocate zero points for asymmetric */
    if (mode == FQ_ASYMMETRIC) {
        qt->zero_points = calloc(num_scales, sizeof(int32_t));
    }

    if (!qt->shape || !qt->data || !qt->scales) {
        fq_free(qt);
        return NULL;
    }

    return qt;
}

void fq_free(fq_tensor* qt) {
    if (!qt) return;

    free(qt->shape);
    free(qt->data);
    free(qt->scales);
    free(qt->zero_points);
    free(qt);
}

/* ============================================================= */
/* INT8 Quantization                                             */
/* ============================================================= */

fq_tensor* fq_quantize_int8(const float* data, int64_t* shape, int ndim) {
    if (!data || !shape || ndim <= 0) return NULL;

    fq_tensor* qt = fq_alloc(shape, ndim, FQ_INT8, FQ_SYMMETRIC, FQ_PER_TENSOR, 1);
    if (!qt) return NULL;

    /* Compute scale */
    float scale;
    int32_t zero_point;
    fq_compute_scale(data, qt->numel, FQ_SYMMETRIC, &scale, &zero_point);
    qt->scales[0] = scale;

    /* Quantize */
    int8_t* qdata = (int8_t*)qt->data;
    float inv_scale = 1.0f / scale;

    for (int64_t i = 0; i < qt->numel; i++) {
        float q = roundf(data[i] * inv_scale);
        if (q < -128.0f) q = -128.0f;
        if (q > 127.0f) q = 127.0f;
        qdata[i] = (int8_t)q;
    }

    return qt;
}

fq_tensor* fq_quantize_int8_per_channel(const float* data, int64_t* shape, int ndim, int channel_dim) {
    if (!data || !shape || ndim <= 0) return NULL;
    if (channel_dim < 0 || channel_dim >= ndim) return NULL;

    int num_channels = (int)shape[channel_dim];

    fq_tensor* qt = fq_alloc(shape, ndim, FQ_INT8, FQ_SYMMETRIC, FQ_PER_CHANNEL, num_channels);
    if (!qt) return NULL;

    /* Compute channel stride */
    int64_t channel_stride = 1;
    for (int i = channel_dim + 1; i < ndim; i++) {
        channel_stride *= shape[i];
    }

    /* First pass: compute scales per channel */
    for (int c = 0; c < num_channels; c++) {
        float abs_max = 0.0f;

        /* Find max in this channel */
        for (int64_t i = 0; i < qt->numel; i++) {
            /* Check if this element belongs to channel c */
            int64_t channel_idx = (i / channel_stride) % num_channels;
            if (channel_idx == c) {
                float abs_val = fabsf(data[i]);
                if (abs_val > abs_max) abs_max = abs_val;
            }
        }

        qt->scales[c] = abs_max / 127.0f;
        if (qt->scales[c] < 1e-10f) qt->scales[c] = 1e-10f;
    }

    /* Second pass: quantize */
    int8_t* qdata = (int8_t*)qt->data;

    for (int64_t i = 0; i < qt->numel; i++) {
        int64_t channel_idx = (i / channel_stride) % num_channels;
        float inv_scale = 1.0f / qt->scales[channel_idx];
        float q = roundf(data[i] * inv_scale);
        if (q < -128.0f) q = -128.0f;
        if (q > 127.0f) q = 127.0f;
        qdata[i] = (int8_t)q;
    }

    return qt;
}

/* ============================================================= */
/* INT4 Quantization                                             */
/* ============================================================= */

/* Pack two 4-bit values into one byte */
static void pack_int4(int8_t* packed, int64_t numel, const int8_t* values) {
    for (int64_t i = 0; i < numel; i += 2) {
        int8_t low = values[i] & 0x0F;
        int8_t high = (i + 1 < numel) ? (values[i + 1] & 0x0F) : 0;
        packed[i / 2] = (high << 4) | low;
    }
}

/* Unpack 4-bit values from packed byte */
static void unpack_int4(const int8_t* packed, int64_t numel, int8_t* values) {
    for (int64_t i = 0; i < numel; i += 2) {
        int8_t byte = packed[i / 2];
        values[i] = (byte & 0x0F);
        if (values[i] > 7) values[i] -= 16;  /* Sign extend */
        if (i + 1 < numel) {
            values[i + 1] = (byte >> 4);
            if (values[i + 1] > 7) values[i + 1] -= 16;  /* Sign extend */
        }
    }
}

/* ============================================================= */
/* Ternary Packing (2-bit: 4 values per byte)                    */
/* ============================================================= */

/* Encoding: 00 = -1, 01 = 0, 10 = +1 */
/* Byte layout: [val3:2bits][val2:2bits][val1:2bits][val0:2bits] */

static void pack_ternary(uint8_t *packed, int64_t numel, const int8_t *values) {
    int64_t n_bytes = (numel + 3) / 4;
    for (int64_t b = 0; b < n_bytes; b++) {
        uint8_t byte = 0;
        for (int j = 0; j < 4; j++) {
            int64_t idx = b * 4 + j;
            if (idx >= numel) break;
            uint8_t enc;
            if (values[idx] < 0)      enc = 0; /* -1 -> 00 */
            else if (values[idx] == 0) enc = 1; /* 0  -> 01 */
            else                       enc = 2; /* +1 -> 10 */
            byte |= (enc << (j * 2));
        }
        packed[b] = byte;
    }
}

static void unpack_ternary(const uint8_t *packed, int64_t numel, int8_t *values) {
    static const int8_t decode[4] = {-1, 0, 1, 0}; /* 00=-1, 01=0, 10=+1, 11=unused */
    for (int64_t i = 0; i < numel; i++) {
        uint8_t byte = packed[i / 4];
        int shift = (i % 4) * 2;
        uint8_t enc = (byte >> shift) & 0x03;
        values[i] = decode[enc];
    }
}

fq_tensor* fq_quantize_int4(const float* data, int64_t* shape, int ndim, int group_size) {
    if (!data || !shape || ndim <= 0) return NULL;
    if (group_size <= 0) group_size = 128;  /* Default group size */

    int64_t numel = compute_numel(shape, ndim);
    int num_groups = (int)((numel + group_size - 1) / group_size);

    fq_tensor* qt = fq_alloc(shape, ndim, FQ_INT4, FQ_SYMMETRIC, FQ_PER_GROUP, num_groups);
    if (!qt) return NULL;
    qt->group_size = group_size;

    /* Temporary buffer for unpacked values */
    int8_t* temp = malloc(numel);
    if (!temp) {
        fq_free(qt);
        return NULL;
    }

    /* Quantize per group */
    for (int g = 0; g < num_groups; g++) {
        int64_t start = g * group_size;
        int64_t end = start + group_size;
        if (end > numel) end = numel;
        /* Find max in group */
        float abs_max = 0.0f;
        for (int64_t i = start; i < end; i++) {
            float abs_val = fabsf(data[i]);
            if (abs_val > abs_max) abs_max = abs_val;
        }

        qt->scales[g] = abs_max / 7.0f;  /* INT4 range is -8 to 7 */
        if (qt->scales[g] < 1e-10f) qt->scales[g] = 1e-10f;

        float inv_scale = 1.0f / qt->scales[g];

        /* Quantize to INT4 range */
        for (int64_t i = start; i < end; i++) {
            float q = roundf(data[i] * inv_scale);
            if (q < -8.0f) q = -8.0f;
            if (q > 7.0f) q = 7.0f;
            temp[i] = (int8_t)q;
        }
    }

    /* Pack INT4 values */
    pack_int4((int8_t*)qt->data, numel, temp);

    free(temp);
    return qt;
}

/* ============================================================= */
/* Ternary Quantization (BitNet b1.58 absmean)                   */
/* ============================================================= */

fq_tensor* fq_quantize_ternary(const float* data, int64_t* shape, int ndim) {
    if (!data || !shape || ndim <= 0) return NULL;

    fq_tensor* qt = fq_alloc(shape, ndim, FQ_TERNARY, FQ_SYMMETRIC, FQ_PER_TENSOR, 1);
    if (!qt) return NULL;

    /* Compute absmean: γ = mean(|data|) */
    double abs_sum = 0.0;
    for (int64_t i = 0; i < qt->numel; i++) {
        abs_sum += fabsf(data[i]);
    }
    float gamma = (float)(abs_sum / qt->numel);
    if (gamma < 1e-10f) gamma = 1e-10f;
    qt->scales[0] = gamma;

    /* Quantize: round(data[i] / γ), clamp to {-1, 0, +1} */
    int8_t* temp = malloc(qt->numel);
    if (!temp) { fq_free(qt); return NULL; }

    float inv_gamma = 1.0f / gamma;
    for (int64_t i = 0; i < qt->numel; i++) {
        float q = roundf(data[i] * inv_gamma);
        if (q < -1.0f) q = -1.0f;
        if (q > 1.0f) q = 1.0f;
        temp[i] = (int8_t)q;
    }

    /* Pack into 2-bit format */
    pack_ternary((uint8_t*)qt->data, qt->numel, temp);

    free(temp);
    return qt;
}

fq_tensor* fq_quantize_ternary_threshold(const float* data, int64_t* shape, int ndim, float threshold) {
    if (!data || !shape || ndim <= 0) return NULL;
    if (threshold < 1e-10f) threshold = 1e-10f;

    fq_tensor* qt = fq_alloc(shape, ndim, FQ_TERNARY, FQ_SYMMETRIC, FQ_PER_TENSOR, 1);
    if (!qt) return NULL;

    qt->scales[0] = threshold; /* reconstruction: ternary_value * scale */

    /* Quantize to {-1, 0, +1} */
    int8_t* temp = malloc(qt->numel);
    if (!temp) { fq_free(qt); return NULL; }

    for (int64_t i = 0; i < qt->numel; i++) {
        if (data[i] > threshold) temp[i] = 1;
        else if (data[i] < -threshold) temp[i] = -1;
        else temp[i] = 0;
    }

    pack_ternary((uint8_t*)qt->data, qt->numel, temp);
    free(temp);

    return qt;
}

fq_tensor* fq_quantize_ternary_per_channel(const float* data, int64_t* shape, int ndim, int channel_dim) {
    if (!data || !shape || ndim <= 0) return NULL;
    if (channel_dim < 0 || channel_dim >= ndim) return NULL;

    int num_channels = (int)shape[channel_dim];

    fq_tensor* qt = fq_alloc(shape, ndim, FQ_TERNARY, FQ_SYMMETRIC, FQ_PER_CHANNEL, num_channels);
    if (!qt) return NULL;

    int64_t channel_stride = 1;
    for (int i = channel_dim + 1; i < ndim; i++) {
        channel_stride *= shape[i];
    }

    /* First pass: compute absmean threshold per channel */
    for (int c = 0; c < num_channels; c++) {
        double sum = 0.0;
        int64_t count = 0;
        for (int64_t i = 0; i < qt->numel; i++) {
            int64_t channel_idx = (i / channel_stride) % num_channels;
            if (channel_idx == c) {
                sum += fabsf(data[i]);
                count++;
            }
        }
        qt->scales[c] = (count > 0) ? (float)(sum / count) : 1e-10f;
        if (qt->scales[c] < 1e-10f) qt->scales[c] = 1e-10f;
    }

    /* Second pass: quantize */
    int8_t* temp = malloc(qt->numel);
    if (!temp) { fq_free(qt); return NULL; }

    for (int64_t i = 0; i < qt->numel; i++) {
        int64_t channel_idx = (i / channel_stride) % num_channels;
        float thresh = qt->scales[channel_idx];
        if (data[i] > thresh) temp[i] = 1;
        else if (data[i] < -thresh) temp[i] = -1;
        else temp[i] = 0;
    }

    pack_ternary((uint8_t*)qt->data, qt->numel, temp);
    free(temp);

    return qt;
}

/* ============================================================= */
/* Generic Quantization                                          */
/* ============================================================= */

fq_tensor* fq_quantize(const float* data, int64_t* shape, int ndim,
                       fq_dtype dtype, fq_mode mode, fq_granularity granularity,
                       int channel_dim, int group_size) {
    (void)mode;

    if (dtype == FQ_TERNARY) {
        if (granularity == FQ_PER_CHANNEL) {
            return fq_quantize_ternary_per_channel(data, shape, ndim, channel_dim);
        }
        return fq_quantize_ternary(data, shape, ndim);
    }

    if (dtype == FQ_INT4 || dtype == FQ_UINT4) {
        return fq_quantize_int4(data, shape, ndim, group_size);
    }

    if (granularity == FQ_PER_CHANNEL) {
        return fq_quantize_int8_per_channel(data, shape, ndim, channel_dim);
    }

    return fq_quantize_int8(data, shape, ndim);
}

/* ============================================================= */
/* Dequantization                                                */
/* ============================================================= */

float* fq_dequantize(fq_tensor* qt) {
    if (!qt) return NULL;

    float* data = malloc(qt->numel * sizeof(float));
    if (!data) return NULL;

    if (qt->dtype == FQ_TERNARY) {
        int8_t* temp = malloc(qt->numel);
        if (!temp) { free(data); return NULL; }

        unpack_ternary((const uint8_t*)qt->data, qt->numel, temp);

        if (qt->granularity == FQ_PER_CHANNEL) {
            int num_channels = qt->num_scales;
            int64_t elements_per_channel = qt->numel / num_channels;
            for (int64_t i = 0; i < qt->numel; i++) {
                int channel_idx = (int)(i / elements_per_channel);
                data[i] = temp[i] * qt->scales[channel_idx];
            }
        } else {
            float scale = qt->scales[0];
            for (int64_t i = 0; i < qt->numel; i++) {
                data[i] = temp[i] * scale;
            }
        }

        free(temp);
    } else if (qt->dtype == FQ_INT8 || qt->dtype == FQ_UINT8) {
        int8_t* qdata = (int8_t*)qt->data;

        if (qt->granularity == FQ_PER_TENSOR) {
            float scale = qt->scales[0];
            for (int64_t i = 0; i < qt->numel; i++) {
                data[i] = qdata[i] * scale;
            }
        } else if (qt->granularity == FQ_PER_CHANNEL) {
            /* Per-channel: elements are grouped by channel (first dim) */
            int num_channels = qt->num_scales;
            int64_t elements_per_channel = qt->numel / num_channels;

            for (int64_t i = 0; i < qt->numel; i++) {
                int channel_idx = (int)(i / elements_per_channel);
                data[i] = qdata[i] * qt->scales[channel_idx];
            }
        }
    } else {
        /* INT4 */
        int8_t* temp = malloc(qt->numel);
        if (!temp) {
            free(data);
            return NULL;
        }

        unpack_int4((int8_t*)qt->data, qt->numel, temp);

        if (qt->granularity == FQ_PER_GROUP) {
            int group_size = qt->group_size;
            for (int64_t i = 0; i < qt->numel; i++) {
                int group_idx = (int)(i / group_size);
                data[i] = temp[i] * qt->scales[group_idx];
            }
        } else {
            float scale = qt->scales[0];
            for (int64_t i = 0; i < qt->numel; i++) {
                data[i] = temp[i] * scale;
            }
        }

        free(temp);
    }

    return data;
}

float* fq_dequantize_range(fq_tensor* qt, int64_t start, int64_t end) {
    if (!qt || start < 0 || end > qt->numel || start >= end) return NULL;

    int64_t len = end - start;
    float* data = malloc(len * sizeof(float));
    if (!data) return NULL;

    /* For simplicity, dequantize all and copy range */
    float* full = fq_dequantize(qt);
    if (!full) {
        free(data);
        return NULL;
    }

    memcpy(data, full + start, len * sizeof(float));
    free(full);

    return data;
}

/* ============================================================= */
/* Quantized Operations                                          */
/* ============================================================= */

float* fq_matmul(const float* A, fq_tensor* B, int M, int K, int N) {
    if (!A || !B) return NULL;

    float* C = calloc(M * N, sizeof(float));
    if (!C) return NULL;

    fq_matmul_acc(A, B, C, M, K, N);
    return C;
}

void fq_matmul_acc(const float* A, fq_tensor* B, float* C, int M, int K, int N) {
    if (!A || !B || !C) return;

    /* Dispatch ternary to native integer-only kernel */
    if (B->dtype == FQ_TERNARY) {
        fq_matmul_ternary_acc(A, B, C, M, K, N);
        return;
    }

    /* Dequantize B for now (could be optimized with int8 accumulation) */
    float* B_f32 = fq_dequantize(B);
    if (!B_f32) return;

    /* Standard matmul: C[m,n] = sum_k A[m,k] * B[k,n] */
    for (int m = 0; m < M; m++) {
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                sum += A[m * K + k] * B_f32[k * N + n];
            }
            C[m * N + n] += sum;
        }
    }

    free(B_f32);
}

/* ============================================================= */
/* Native Ternary Matmul (integer-only, no dequant)              */
/* ============================================================= */

void fq_matmul_ternary_acc(const float* A, fq_tensor* B, float* C, int M, int K, int N) {
    if (!A || !B || !C) return;

    const uint8_t* packed = (const uint8_t*)B->data;
    float scale = B->scales[0];
    static const int8_t ternary_val[4] = {-1, 0, 1, 0};

    /* C[m,n] += scale * sum_k( A[m,k] * ternary(B[k,n]) ) */
    /* B is K×N row-major, element [k,n] is at linear index k*N+n */
    for (int m = 0; m < M; m++) {
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
            for (int k = 0; k < K; k++) {
                int64_t idx = (int64_t)k * N + n;
                uint8_t byte = packed[idx / 4];
                int shift = (idx % 4) * 2;
                int8_t w = ternary_val[(byte >> shift) & 0x03];
                /* Branch-free: w is -1, 0, or +1 so this is add/nop/sub */
                sum += A[m * K + k] * w;
            }
            C[m * N + n] += sum * scale;
        }
    }
}

float* fq_matmul_ternary(const float* A, fq_tensor* B, int M, int K, int N) {
    if (!A || !B) return NULL;

    float* C = calloc(M * N, sizeof(float));
    if (!C) return NULL;

    fq_matmul_ternary_acc(A, B, C, M, K, N);
    return C;
}

/* ============================================================= */
/* PRNG (caller-owned state for reproducibility)                 */
/* ============================================================= */

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

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

void fq_seed(uint64_t *rng_state, uint64_t seed) {
    *rng_state = seed ? seed : 1;
}

/* ============================================================= */
/* Ternary Backward Pass                                         */
/* ============================================================= */

void fq_matmul_ternary_backward_input(const float *d_output, fq_tensor *W,
                                       float *d_input, int M, int K, int N) {
    if (!d_output || !W || !d_input) return;

    const uint8_t *packed = (const uint8_t *)W->data;
    float scale = W->scales[0];
    static const int8_t ternary_val[4] = {-1, 0, 1, 0};

    /* d_input[m,k] += scale * sum_n d_output[m,n] * W[k,n] */
    for (int m = 0; m < M; m++) {
        for (int k = 0; k < K; k++) {
            float sum = 0.0f;
            for (int n = 0; n < N; n++) {
                int64_t idx = (int64_t)k * N + n;
                uint8_t byte = packed[idx / 4];
                int shift = (idx % 4) * 2;
                int8_t w = ternary_val[(byte >> shift) & 0x03];
                sum += d_output[m * N + n] * w;
            }
            d_input[m * K + k] += sum * scale;
        }
    }
}

void fq_matmul_ternary_backward_weight(const float *input, const float *d_output,
                                        float *d_weight, int M, int K, int N) {
    if (!input || !d_output || !d_weight) return;

    /* d_weight[k,n] += sum_m input[m,k] * d_output[m,n] */
    for (int k = 0; k < K; k++) {
        for (int n = 0; n < N; n++) {
            float sum = 0.0f;
            for (int m = 0; m < M; m++) {
                sum += input[m * K + k] * d_output[m * N + n];
            }
            d_weight[k * N + n] += sum;
        }
    }
}

/* ============================================================= */
/* Stochastic Ternary Flip                                       */
/* ============================================================= */

void fq_ternary_stochastic_flip(fq_tensor *W, const float *d_weight,
                                 float temperature, uint64_t *rng_state) {
    if (!W || !d_weight || !rng_state) return;

    uint8_t *packed = (uint8_t *)W->data;
    static const int8_t decode[4] = {-1, 0, 1, 0};
    /* Encoding: -1 -> 0b00, 0 -> 0b01, +1 -> 0b10 */

    for (int64_t i = 0; i < W->numel; i++) {
        float grad = d_weight[i];
        float p_flip = fabsf(grad) * temperature;
        if (p_flip > 1.0f) p_flip = 1.0f;

        if (p_flip <= 0.0f) continue;
        if (fq_rand_uniform(rng_state) >= p_flip) continue;

        /* Read current ternary value */
        int64_t byte_idx = i / 4;
        int bit_shift = (i % 4) * 2;
        uint8_t byte = packed[byte_idx];
        uint8_t enc = (byte >> bit_shift) & 0x03;
        int8_t cur = decode[enc];

        /* Flip toward negative gradient direction:
         * grad > 0 -> decrease weight (flip toward -1)
         * grad < 0 -> increase weight (flip toward +1) */
        int8_t next = cur;
        if (grad > 0.0f && cur > -1) next = cur - 1;
        if (grad < 0.0f && cur < 1)  next = cur + 1;

        if (next == cur) continue;

        /* Encode new value */
        uint8_t new_enc;
        if (next < 0)       new_enc = 0; /* -1 -> 00 */
        else if (next == 0) new_enc = 1; /*  0 -> 01 */
        else                new_enc = 2; /* +1 -> 10 */

        /* Clear old 2-bit field and write new */
        byte &= ~(0x03 << bit_shift);
        byte |= (new_enc << bit_shift);
        packed[byte_idx] = byte;
    }
}

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

void fq_print(fq_tensor* qt) {
    if (!qt) {
        printf("(null quantized tensor)\n");
        return;
    }

    printf("Quantized tensor: %s, [", fq_dtype_name(qt->dtype));
    for (int i = 0; i < qt->ndim; i++) {
        if (i > 0) printf(", ");
        printf("%lld", (long long)qt->shape[i]);
    }
    printf("]\n");

    printf("  Mode: %s\n", qt->mode == FQ_SYMMETRIC ? "symmetric" : "asymmetric");
    printf("  Granularity: %s (%d scales)\n",
           qt->granularity == FQ_PER_TENSOR ? "per-tensor" :
           qt->granularity == FQ_PER_CHANNEL ? "per-channel" : "per-group",
           qt->num_scales);
    printf("  Memory: %.2f KB (%.1fx compression)\n",
           fq_memory_usage(qt) / 1024.0f,
           (qt->numel * 4.0f) / fq_memory_usage(qt));
}

float fq_quantization_error(const float* original, fq_tensor* quantized) {
    if (!original || !quantized) return -1.0f;

    float* dequant = fq_dequantize(quantized);
    if (!dequant) return -1.0f;

    double mse = 0.0;
    for (int64_t i = 0; i < quantized->numel; i++) {
        double diff = original[i] - dequant[i];
        mse += diff * diff;
    }
    mse /= quantized->numel;

    free(dequant);
    return (float)mse;
}