/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.cassandra.index.sai.cql;

import java.util.ArrayList;
import java.util.List;

import org.junit.Test;

import org.apache.cassandra.cql3.UntypedResultSet;
import org.apache.cassandra.db.marshal.FloatType;
import org.apache.cassandra.db.marshal.VectorType;
import org.apache.cassandra.index.sai.disk.v1.segment.SegmentBuilder;

import static org.assertj.core.api.Assertions.assertThat;

public class VectorSegmentationTest extends VectorTester
{
    @Test
    public void testMultipleSegmentsForCreatingIndex() throws Throwable
    {
        createTable("CREATE TABLE %s (pk int, val vector<float, " + word2vec.dimension() + ">, PRIMARY KEY(pk))");

        int vectorCount = 100;
        List<float[]> vectors = new ArrayList<>();
        for (int row = 0; row < vectorCount; row++)
        {
            float[] vector = word2vec.vector(row);
            vectors.add(vector);
            execute("INSERT INTO %s (pk, val) VALUES (?, ?)", row, vector(vector));
        }

        flush();

        SegmentBuilder.updateLastValidSegmentRowId(17); // 17 rows per segment
        createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");

        int limit = 35;
        float[] queryVector = word2vec.vector(getRandom().nextIntBetween(0, vectorCount));
        UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
        assertThat(resultSet.size()).isEqualTo(limit);

        List<float[]> resultVectors = getVectorsFromResult(resultSet);
        double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
        assertThat(recall).isGreaterThanOrEqualTo(0.9);
    }

    @Test
    public void testMultipleSegmentsForCompaction() throws Throwable
    {
        createTable("CREATE TABLE %s (pk int, val vector<float, " + word2vec.dimension() + ">, PRIMARY KEY(pk))");
        createIndex("CREATE CUSTOM INDEX ON %s(val) USING 'StorageAttachedIndex'");

        List<float[]> vectors = new ArrayList<>();
        int vectorCount = 0;
        int rowsPerSSTable = 10;
        int sstables = 5;
        int pk = 0;
        for (int i = 0; i < sstables; i++)
        {
            for (int row = 0; row < rowsPerSSTable; row++)
            {
                float[] vector = word2vec.vector(vectorCount++);
                execute("INSERT INTO %s (pk, val) VALUES (?, ?)", pk++, vector(vector));
                vectors.add(vector);
            }

            flush();
        }

        int limit = 30;
        float[] queryVector = word2vec.vector(getRandom().nextIntBetween(0, vectorCount));
        UntypedResultSet resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
        assertThat(resultSet.size()).isEqualTo(limit);

        List<float[]> resultVectors = getVectorsFromResult(resultSet);
        double recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
        assertThat(recall).isGreaterThanOrEqualTo(0.99);


        SegmentBuilder.updateLastValidSegmentRowId(11); // 11 rows per segment
        compact();

        queryVector = word2vec.vector(getRandom().nextIntBetween(0, vectorCount));
        resultSet = execute("SELECT * FROM %s ORDER BY val ANN OF ? LIMIT " + limit, vector(queryVector));
        assertThat(resultSet.size()).isEqualTo(limit);

        resultVectors = getVectorsFromResult(resultSet);
        recall = rawIndexedRecall(vectors, queryVector, resultVectors, limit);
        assertThat(recall).isGreaterThanOrEqualTo(0.99);
    }

    private static List<float[]> getVectorsFromResult(UntypedResultSet result)
    {
        List<float[]> vectors = new ArrayList<>();
        VectorType<?> vectorType = VectorType.getInstance(FloatType.instance, word2vec.dimension());

        // verify results are part of inserted vectors
        for (UntypedResultSet.Row row: result)
        {
            vectors.add(vectorType.composeAsFloat(row.getBytes("val")));
        }

        return vectors;
    }
}
