import java.util.*;

/**
 * Unit test for ollama-0001: kvcache.Causal.buildMask Except []int O(B*E) scan.
 *
 * Models the Ollama defect where buildMask() calls slices.Contains(opts.Except, i)
 * inside a loop over all batch tokens (curBatchSize = B). Except is populated by
 * Gemma3 for multimodal image token positions. For N images with 256 tokens each:
 *   |except| = N * 256, B = N * 256 + text_tokens ≈ N * 256
 *   Cost of Except check alone: O(B * E) = O((N*256)^2) — quadratic in image count.
 *
 * Fix: convert Except to a HashSet before the outer loop, giving O(1) per lookup
 * and O(B + E) total cost for the Except component.
 */
public class OllamaKvcacheBuildmaskExceptTest {

    // --- Defect: slices.Contains on []int Except, O(B*E) ---
    // Isolates just the enabled[] computation (the Except scan), not the full mask fill.
    static boolean[] computeEnabledLinear(int batchSize, int[] except) {
        boolean[] enabled = new boolean[batchSize];
        for (int i = 0; i < batchSize; i++) {
            enabled[i] = !listContains(except, i);
        }
        return enabled;
    }

    static boolean listContains(int[] arr, int val) {
        for (int v : arr) {
            if (v == val) return true;
        }
        return false;
    }

    // --- Fix: map[int]struct{} lookup, O(B+E) ---
    static boolean[] computeEnabledHashSet(int batchSize, int[] except) {
        Set<Integer> exceptSet = new HashSet<>(except.length * 2);
        for (int idx : except) exceptSet.add(idx);
        boolean[] enabled = new boolean[batchSize];
        for (int i = 0; i < batchSize; i++) {
            enabled[i] = !exceptSet.contains(i);
        }
        return enabled;
    }

    // Build an Except array simulating N images * tokensPerImage = 256 each
    static int[] buildExcept(int numImages, int tokensPerImage) {
        int[] except = new int[numImages * tokensPerImage];
        for (int i = 0; i < except.length; i++) {
            except[i] = i;
        }
        return except;
    }

    public static void main(String[] args) {
        // Correctness: small case
        int smallBatch = 16;
        int[] smallExcept = {0, 1, 3, 5};
        boolean[] enabledLinear = computeEnabledLinear(smallBatch, smallExcept);
        boolean[] enabledHash   = computeEnabledHashSet(smallBatch, smallExcept);
        assert Arrays.equals(enabledLinear, enabledHash) : "Correctness failed";
        System.out.println("Correctness OK: batch=" + smallBatch + " except=" + smallExcept.length);

        // Performance: simulate 10-image Gemma3 prefill
        // tokensPerImage=256, 10 images => batchSize=2560, except=2560
        int numImages = 10;
        int tokensPerImage = 256;
        int batchSize = numImages * tokensPerImage;
        int[] except = buildExcept(numImages, tokensPerImage);

        // Correctness at full scale
        boolean[] mL = computeEnabledLinear(batchSize, except);
        boolean[] mH = computeEnabledHashSet(batchSize, except);
        assert Arrays.equals(mL, mH) : "Correctness failed at full scale";

        // Warmup
        for (int w = 0; w < 3; w++) {
            computeEnabledLinear(batchSize, except);
            computeEnabledHashSet(batchSize, except);
        }

        int reps = 200;
        long t0 = System.nanoTime();
        for (int r = 0; r < reps; r++) computeEnabledLinear(batchSize, except);
        long linearNs = System.nanoTime() - t0;

        long t1 = System.nanoTime();
        for (int r = 0; r < reps; r++) computeEnabledHashSet(batchSize, except);
        long hashNs = System.nanoTime() - t1;

        double ratio = (double) linearNs / hashNs;
        System.out.printf("Gemma3 multi-image: %d images x %d tokens = B=%d |except|=%d%n",
                numImages, tokensPerImage, batchSize, except.length);
        System.out.printf("  Defect ([]int linear scan): %,d ns%n", linearNs);
        System.out.printf("  Fix    (HashSet O(1) lookup): %,d ns%n", hashNs);
        System.out.printf("  speedup: %.1fx%n", ratio);

        assert ratio > 3.0 : "Expected >3x speedup, got " + ratio;
        System.out.println("PASS");
    }
}