/**
 * 使用 OpenAI text-embedding-3-large 的 MRL 特性
 * 支持直接指定维度，SDK会返回对应维度的向量
 */
public class MRLEmbeddingService {

    private final OpenAIClient openAIClient;

    // 支持的维度列表（MRL训练时指定的）
    private static final Set<Integer> SUPPORTED_DIMS = Set.of(
        64, 128, 256, 512, 1024, 1536, 3072
    );

    public MRLEmbeddingService(String apiKey) {
        this.openAIClient = OpenAIClient.builder()
            .apiKey(apiKey)
            .build();
    }

    /**
     * 获取指定维度的Embedding
     *
     * @param text      输入文本
     * @param dimension 目标维度（64/128/256/512/1024/1536/3072）
     */
    public float[] embed(String text, int dimension) {
        if (!SUPPORTED_DIMS.contains(dimension)) {
            throw new IllegalArgumentException(
                "不支持的维度: " + dimension +
                "，支持: " + SUPPORTED_DIMS);
        }

        var request = EmbeddingCreateParams.builder()
            .model("text-embedding-3-large")
            .input(text)
            .dimensions(dimension)   // MRL的关键参数
            .build();

        var response = openAIClient.embeddings().create(request);
        return toFloatArray(response.data().get(0).embedding());
    }

    /**
     * 批量获取Embedding（推荐，减少API调用次数）
     */
    public List<float[]> embedBatch(List<String> texts, int dimension) {
        // OpenAI支持单次最多2048个文本
        int batchSize = 512;
        List<float[]> results = new ArrayList<>();

        for (int i = 0; i < texts.size(); i += batchSize) {
            List<String> batch = texts.subList(i,
                Math.min(i + batchSize, texts.size()));

            var request = EmbeddingCreateParams.builder()
                .model("text-embedding-3-large")
                .input(batch)
                .dimensions(dimension)
                .build();

            var response = openAIClient.embeddings().create(request);
            response.data().forEach(item ->
                results.add(toFloatArray(item.embedding())));

            System.out.printf("批次 %d/%d 完成%n",
                Math.min(i + batchSize, texts.size()), texts.size());
        }
        return results;
    }

    private float[] toFloatArray(List<Double> doubles) {
        float[] arr = new float[doubles.size()];
        for (int i = 0; i < doubles.size(); i++) arr[i] = doubles.get(i).floatValue();
        return arr;
    }

    /**
     * 如果用的是非MRL模型（如BGE），需要PCA降维
     * 先离线训练PCA，再在线截取
     */
    public static float[] truncateForNonMRLModel(float[] fullVector,
                                                   PCAReducer pcaReducer) {
        // 非MRL模型：先做PCA投影，取前k个主成分
        double[][] input = new double[1][fullVector.length];
        for (int i = 0; i < fullVector.length; i++) input[0][i] = fullVector[i];
        double[][] output = pcaReducer.transform(input);

        float[] result = new float[output[0].length];
        for (int i = 0; i < output[0].length; i++) result[i] = (float) output[0][i];

        // 归一化
        return normalize(result);
    }

    private static float[] normalize(float[] v) {
        float norm = 0;
        for (float x : v) norm += x * x;
        norm = (float) Math.sqrt(norm);
        if (norm < 1e-12f) return v;
        float[] result = new float[v.length];
        for (int i = 0; i < v.length; i++) result[i] = v[i] / norm;
        return result;
    }
}

四、数值量化：float32 → int8

4.1 量化的数学基础

float32 用 32 位表示一个浮点数，范围大、精度高，但绝大多数情况下向量值集中在 [-1, 1] 的小范围内。

int8 用 8 位表示整数，范围是 [-128, 127]。

量化的思路：找到向量值的范围 [min, max]，线性映射到 [-128, 127]。

q = \text{round}\left(\frac{v - \text{zero\_point}}{\text{scale}}\right)

反量化：

v \approx q \times \text{scale} + \text{zero\_point}

其中：

\text{scale} = \frac{\max - \min}{255}

\text{zero\_point} = \text{round}\left(-\frac{\min}{\text{scale}}\right) - 128

4.2 两种量化策略

逐张量量化（Per-Tensor）：整个向量用同一个 scale 和 zero_point。

public class PerTensorQuantizer {

    /**
     * 逐张量量化：整个向量统一缩放
     * 优点：实现简单
     * 缺点：如果向量各维度量程差异大，部分维度精度损失严重
     */
    public static QuantizedVector quantize(float[] vector) {
        float min = Float.MAX_VALUE;
        float max = Float.MIN_VALUE;
        for (float v : vector) {
            if (v < min) min = v;
            if (v > max) max = v;
        }

        float scale = (max - min) / 255.0f;
        if (scale < 1e-8f) scale = 1e-8f; // 防止除以0

        byte[] quantized = new byte[vector.length];
        for (int i = 0; i < vector.length; i++) {
            int q = Math.round((vector[i] - min) / scale) - 128;
            quantized[i] = (byte) Math.max(-128, Math.min(127, q));
        }

        return new QuantizedVector(quantized, scale, min);
    }

    public record QuantizedVector(byte[] data, float scale, float minVal) {

        /**
         * 反量化还原
         */
        public float[] dequantize() {
            float[] result = new float[data.length];
            for (int i = 0; i < data.length; i++) {
                result[i] = (data[i] + 128) * scale + minVal;
            }
            return result;
        }

        /**
         * 直接在量化空间计算内积（近似）
         * 避免反量化，计算更快
         */
        public int dotProductInt(QuantizedVector other) {
            int sum = 0;
            for (int i = 0; i < data.length; i++) {
                sum += data[i] * other.data[i];
            }
            return sum;
        }
    }
}

逐通道量化（Per-Channel）：每个维度独立统计 scale 和 zero_point。

public class PerChannelQuantizer {

    private final float[] scales;
    private final float[] minVals;
    private final int dim;

    /**
     * 先用大量样本统计每个维度的分布
     * 这一步离线完成
     */
    public PerChannelQuantizer(List<float[]> calibrationData) {
        this.dim = calibrationData.get(0).length;
        this.scales = new float[dim];
        this.minVals = new float[dim];

        float[] mins = new float[dim];
        float[] maxs = new float[dim];
        Arrays.fill(mins, Float.MAX_VALUE);
        Arrays.fill(maxs, -Float.MAX_VALUE);

        for (float[] v : calibrationData) {
            for (int d = 0; d < dim; d++) {
                if (v[d] < mins[d]) mins[d] = v[d];
                if (v[d] > maxs[d]) maxs[d] = v[d];
            }
        }

        for (int d = 0; d < dim; d++) {
            minVals[d] = mins[d];
            scales[d] = Math.max((maxs[d] - mins[d]) / 255.0f, 1e-8f);
        }

        System.out.printf("逐通道量化校准完成，维度: %d%n", dim);
    }

    /**
     * 对单个向量做逐通道量化
     */
    public byte[] quantize(float[] vector) {
        byte[] result = new byte[dim];
        for (int d = 0; d < dim; d++) {
            int q = Math.round((vector[d] - minVals[d]) / scales[d]) - 128;
            result[d] = (byte) Math.max(-128, Math.min(127, q));
        }
        return result;
    }

    /**
     * 反量化
     */
    public float[] dequantize(byte[] quantized) {
        float[] result = new float[dim];
        for (int d = 0; d < dim; d++) {
            result[d] = (quantized[d] + 128) * scales[d] + minVals[d];
        }
        return result;
    }

    /**
     * 量化后的精度损失分析
     * 计算原始向量和量化→反量化向量之间的余弦相似度
     */
    public void analyzeQuantizationError(List<float[]> testVectors) {
        double totalCosineSim = 0;
        double minCosineSim = 1.0;

        for (float[] original : testVectors) {
            float[] reconstructed = dequantize(quantize(original));
            float cos = cosineSim(original, reconstructed);
            totalCosineSim += cos;
            if (cos < minCosineSim) minCosineSim = cos;
        }

        System.out.printf("量化精度分析 (n=%d):%n", testVectors.size());
        System.out.printf("  平均余弦相似度: %.6f%n",
            totalCosineSim / testVectors.size());
        System.out.printf("  最低余弦相似度: %.6f%n", minCosineSim);
        System.out.printf("  平均信息保留率: %.4f%%%n",
            totalCosineSim / testVectors.size() * 100);
    }

    private float cosineSim(float[] a, float[] b) {
        float dot = 0, na = 0, nb = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i];
        }
        if (na == 0 || nb == 0) return 0;
        return dot / (float)(Math.sqrt(na) * Math.sqrt(nb));
    }
}

4.3 量化感知的信息损失量化

量化会引入误差，误差的大小取决于向量值的分布。经验数据：

量化方案	平均余弦相似度（原始 vs 量化）	内存节省
float32（基准）	1.0000	1x
float16	0.9999	2x
int8 Per-Tensor	0.9987	4x
int8 Per-Channel	0.9994	4x
二值化（1bit）	0.9523	32x

int8 Per-Channel 量化的平均余弦相似度达到 0.9994，信息损失非常小，是工程上的最佳平衡点。

五、组合压缩的完整工程方案

把维度压缩和数值量化结合，构建完整的压缩流水线：

/**
 * 完整的 Embedding 压缩流水线
 * 原始: float32 / 768d → 目标: int8 / 256d
 * 存储压缩比: 3072字节 → 256字节，压缩约 12x
 */
public class EmbeddingCompressionPipeline {

    private final PCAReducer pcaReducer;
    private final PerChannelQuantizer quantizer;
    private final int targetDim;

    /**
     * 构建压缩流水线（离线，需要校准数据）
     *
     * @param calibrationData 用于拟合PCA和量化参数的样本数据
     * @param targetDim       目标维度（如256）
     */
    public static EmbeddingCompressionPipeline build(
            List<float[]> calibrationData, int targetDim) {

        System.out.println("=== 开始构建压缩流水线 ===");
        System.out.printf("校准数据量: %d，原始维度: %d，目标维度: %d%n",
            calibrationData.size(), calibrationData.get(0).length, targetDim);

        // Step 1: 拟合 PCA
        System.out.println("Step 1: 拟合 PCA 变换矩阵...");
        PCAReducer pca = new PCAReducer();
        double[][] pcaInput = calibrationData.stream()
            .map(v -> {
                double[] d = new double[v.length];
                for (int i = 0; i < v.length; i++) d[i] = v[i];
                return d;
            })
            .toArray(double[][]::new);
        pca.fit(pcaInput, targetDim);

        // Step 2: 用PCA降维后的数据校准量化器
        System.out.println("Step 2: 校准逐通道量化参数...");
        List<float[]> reducedData = calibrationData.stream()
            .map(v -> {
                double[][] inp = new double[1][v.length];
                for (int i = 0; i < v.length; i++) inp[0][i] = v[i];
                double[][] out = pca.transform(inp);
                float[] result = new float[targetDim];
                for (int i = 0; i < targetDim; i++) result[i] = (float) out[0][i];
                return normalize(result);
            })
            .collect(Collectors.toList());

        PerChannelQuantizer quant = new PerChannelQuantizer(reducedData);

        System.out.println("=== 压缩流水线构建完成 ===");
        return new EmbeddingCompressionPipeline(pca, quant, targetDim);
    }

    private EmbeddingCompressionPipeline(PCAReducer pca,
                                          PerChannelQuantizer quant,
                                          int targetDim) {
        this.pcaReducer = pca;
        this.quantizer = quant;
        this.targetDim = targetDim;
    }

    /**
     * 在线压缩：单个向量
     * 延迟目标：< 0.1ms
     */
    public byte[] compress(float[] original) {
        // Step 1: PCA 降维
        double[][] input = new double[1][original.length];
        for (int i = 0; i < original.length; i++) input[0][i] = original[i];
        double[][] reduced = pcaReducer.transform(input);

        float[] reducedFloat = new float[targetDim];
        for (int i = 0; i < targetDim; i++) reducedFloat[i] = (float) reduced[0][i];

        // Step 2: 归一化（量化前必须归一化）
        float[] normalized = normalize(reducedFloat);

        // Step 3: int8 量化
        return quantizer.quantize(normalized);
    }

    /**
     * 在压缩向量空间中计算近似相似度
     */
    public float approximateSimilarity(byte[] a, byte[] b) {
        // 在int8空间直接计算内积（不反量化），作为相似度近似
        int dotProduct = 0;
        int normA = 0, normB = 0;
        for (int i = 0; i < a.length; i++) {
            dotProduct += (int) a[i] * (int) b[i];
            normA += (int) a[i] * (int) a[i];
            normB += (int) b[i] * (int) b[i];
        }
        if (normA == 0 || normB == 0) return 0;
        return (float)(dotProduct / (Math.sqrt(normA) * Math.sqrt(normB)));
    }

    /**
     * 评估整体压缩方案的召回率影响
     */
    public void evaluateEndToEnd(List<float[]> queries,
                                  List<float[]> corpus,
                                  int k) {
        System.out.println("=== 端到端召回率评估 ===");

        // 原始空间的 ground truth
        List<List<Integer>> groundTruth = queries.stream()
            .map(q -> findTopKOriginal(q, corpus, k))
            .collect(Collectors.toList());

        // 压缩空间的检索结果
        List<byte[]> compressedCorpus = corpus.stream()
            .map(this::compress)
            .collect(Collectors.toList());

        int totalHits = 0;
        for (int q = 0; q < queries.size(); q++) {
            byte[] compressedQuery = compress(queries.get(q));
            List<Integer> topK = findTopKCompressed(
                compressedQuery, compressedCorpus, k);

            Set<Integer> gtSet = new HashSet<>(groundTruth.get(q));
            for (int idx : topK) {
                if (gtSet.contains(idx)) totalHits++;
            }
        }

        double recall = (double) totalHits / (queries.size() * k);
        System.out.printf("Recall@%d (压缩 vs 原始): %.4f%n", k, recall);
        System.out.printf("存储压缩比: %.1fx%n",
            (float)(corpus.get(0).length * 4) / targetDim);
    }

    private List<Integer> findTopKOriginal(float[] query, List<float[]> corpus, int k) {
        var pq = new PriorityQueue<int[]>(
            (a, b) -> Float.compare(Float.intBitsToFloat(a[1]),
                                     Float.intBitsToFloat(b[1])));
        for (int i = 0; i < corpus.size(); i++) {
            float sim = cosineSim(query, corpus.get(i));
            pq.offer(new int[]{i, Float.floatToIntBits(sim)});
            if (pq.size() > k) pq.poll();
        }
        return pq.stream().map(e -> e[0]).collect(Collectors.toList());
    }

    private List<Integer> findTopKCompressed(byte[] query,
                                              List<byte[]> corpus, int k) {
        var pq = new PriorityQueue<int[]>(
            (a, b) -> Float.compare(Float.intBitsToFloat(a[1]),
                                     Float.intBitsToFloat(b[1])));
        for (int i = 0; i < corpus.size(); i++) {
            float sim = approximateSimilarity(query, corpus.get(i));
            pq.offer(new int[]{i, Float.floatToIntBits(sim)});
            if (pq.size() > k) pq.poll();
        }
        return pq.stream().map(e -> e[0]).collect(Collectors.toList());
    }

    private static float[] normalize(float[] v) {
        float norm = 0;
        for (float x : v) norm += x * x;
        norm = (float) Math.sqrt(norm);
        if (norm < 1e-12f) return v;
        float[] r = new float[v.length];
        for (int i = 0; i < v.length; i++) r[i] = v[i] / norm;
        return r;
    }

    private float cosineSim(float[] a, float[] b) {
        float dot = 0, na = 0, nb = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i]; na += a[i] * a[i]; nb += b[i] * b[i];
        }
        if (na == 0 || nb == 0) return 0;
        return dot / (float)(Math.sqrt(na) * Math.sqrt(nb));
    }
}

六、与向量库集成（以 Milvus 为例）

import io.milvus.v2.client.*;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.response.*;

/**
 * 使用压缩向量的 Milvus 集成方案
 */
public class CompressedMilvusService {

    private final MilvusClientV2 milvusClient;
    private final EmbeddingCompressionPipeline pipeline;
    private final String collectionName;

    public CompressedMilvusService(MilvusClientV2 client,
                                    EmbeddingCompressionPipeline pipeline,
                                    String collectionName) {
        this.milvusClient = client;
        this.pipeline = pipeline;
        this.collectionName = collectionName;
    }

    /**
     * 创建使用 INT8 量化的 Collection
     * Milvus 2.4+ 支持 INT8 向量
     */
    public void createCollection(int vectorDim) {
        var schema = CreateCollectionReq.CollectionSchema.newBuilder()
            .addField(AddFieldReq.newBuilder()
                .fieldName("id").dataType(DataType.Int64).isPrimaryKey(true)
                .autoID(true).build())
            .addField(AddFieldReq.newBuilder()
                .fieldName("text").dataType(DataType.VarChar)
                .maxLength(2000).build())
            // 注意：使用 INT8 向量类型
            .addField(AddFieldReq.newBuilder()
                .fieldName("embedding")
                .dataType(DataType.Int8Vector)  // INT8向量
                .dimension(vectorDim).build())
            .build();

        var indexParam = IndexParam.builder()
            .fieldName("embedding")
            .indexType(IndexParam.IndexType.HNSW)
            .metricType(IndexParam.MetricType.COSINE)
            .extraParams(Map.of("M", 16, "efConstruction", 200))
            .build();

        milvusClient.createCollection(
            CreateCollectionReq.newBuilder()
                .collectionName(collectionName)
                .collectionSchema(schema)
                .indexParams(List.of(indexParam))
                .build()
        );
        System.out.println("Collection 创建成功，使用 INT8 压缩向量");
    }

    /**
     * 批量插入（在线压缩）
     */
    public void insertBatch(List<String> texts, List<float[]> embeddings) {
        List<JsonObject> rows = new ArrayList<>();

        for (int i = 0; i < texts.size(); i++) {
            byte[] compressed = pipeline.compress(embeddings.get(i));

            JsonObject row = new JsonObject();
            row.addProperty("text", texts.get(i));

            // byte[] 转 JSON array
            var arr = new com.google.gson.JsonArray();
            for (byte b : compressed) arr.add(b);
            row.add("embedding", arr);

            rows.add(row);
        }

        milvusClient.insert(InsertReq.newBuilder()
            .collectionName(collectionName)
            .data(rows)
            .build());
    }

    /**
     * 查询（在线压缩查询向量）
     */
    public List<SearchResp.SearchResult> search(float[] queryEmbedding, int k) {
        byte[] compressedQuery = pipeline.compress(queryEmbedding);

        // byte[] 转 List<Byte>（Milvus Java SDK 要求）
        List<Byte> queryVec = new ArrayList<>();
        for (byte b : compressedQuery) queryVec.add(b);

        var response = milvusClient.search(SearchReq.newBuilder()
            .collectionName(collectionName)
            .data(List.of(queryVec))
            .annsField("embedding")
            .topK(k)
            .outputFields(List.of("text"))
            .searchParams(Map.of("ef", 64))
            .build());

        return response.getSearchResults().get(0);
    }
}