import java.util.*;

/**
 * vllm-0002: Grok2Tokenizer decode O(N×S) dict.values() linear scan
 *
 * Models the two decode implementations:
 *   DEFECTIVE:  for each token, scan all V dict values  → O(N×V)
 *   PATCHED:    use a frozenset (HashSet) of special IDs → O(N×1) = O(N)
 *
 * Mirrors vllm/vllm/tokenizers/grok2.py lines 354-363 and 376-382.
 */
public class Vllm0002Grok2SpecialTokenTest {

    // --- Defective implementation ---
    static List<Integer> decodeDefective(List<Integer> ids,
                                         Map<String, Integer> specialTokens,
                                         boolean skipSpecial) {
        if (!skipSpecial) return new ArrayList<>(ids);
        List<Integer> result = new ArrayList<>();
        for (int tokenId : ids) {
            // O(V) dict.values() scan per token — mirrors Python dict.values() membership
            if (!specialTokens.containsValue(tokenId)) {
                result.add(tokenId);
            }
        }
        return result;
    }

    // --- Patched implementation ---
    static List<Integer> decodePatched(List<Integer> ids,
                                        Set<Integer> specialTokenIds,
                                        boolean skipSpecial) {
        if (!skipSpecial) return new ArrayList<>(ids);
        List<Integer> result = new ArrayList<>();
        for (int tokenId : ids) {
            // O(1) frozenset/HashSet lookup
            if (!specialTokenIds.contains(tokenId)) {
                result.add(tokenId);
            }
        }
        return result;
    }

    public static void main(String[] args) {
        // Build a special token dict simulating Grok-2 (200 special tokens)
        int V = 200;
        Map<String, Integer> specialTokens = new LinkedHashMap<>();
        // Regular vocab: 0..99999; special tokens start at 100000
        for (int i = 0; i < V; i++) {
            specialTokens.put("<special_" + i + ">", 100000 + i);
        }
        Set<Integer> specialTokenIds = new HashSet<>(specialTokens.values());

        // Build a realistic output sequence: N=2048 tokens, 10 special, rest regular
        int N = 2048;
        List<Integer> ids = new ArrayList<>(N);
        for (int i = 0; i < N; i++) {
            if (i % 200 == 0) {
                ids.add(100000); // special token <special_0>
            } else {
                ids.add(i % 50000); // regular token
            }
        }

        // Correctness check
        List<Integer> defectResult = decodeDefective(ids, specialTokens, true);
        List<Integer> patchResult  = decodePatched(ids, specialTokenIds, true);
        assert defectResult.equals(patchResult)
            : "FAIL: defect and patched outputs differ";

        // Timing comparison
        int REPS = 500;

        long t0 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            decodeDefective(ids, specialTokens, true);
        }
        long defectNs = System.nanoTime() - t0;

        long t1 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            decodePatched(ids, specialTokenIds, true);
        }
        long patchNs = System.nanoTime() - t1;

        double ratio = (double) defectNs / patchNs;
        System.out.printf("vllm-0002 MOAD-0001 Grok2 decode O(N×S) dict.values() scan%n");
        System.out.printf("  N=%d tokens, V=%d special tokens, %d reps%n", N, V, REPS);
        System.out.printf("  Defective (dict.values scan): %,d ms%n", defectNs / 1_000_000);
        System.out.printf("  Patched   (frozenset lookup):  %,d ms%n", patchNs  / 1_000_000);
        System.out.printf("  Speedup: %.1fx%n", ratio);

        assert ratio > 5.0 : "FAIL: expected >5x speedup, got " + ratio;
        System.out.println("PASS");
    }
}