import java.util.*;

/**
 * transformers-0003: convert_tokens_to_string O(T×S) list scan vs set lookup
 *
 * Models marian/m2m100/speech_to_text convert_tokens_to_string():
 *   DEFECTIVE: for each token, call all_special_tokens property (list) → O(T×S)
 *   PATCHED:   cache set(all_special_tokens) before loop → O(T)
 *
 * Mirrors the pattern at:
 *   src/transformers/models/marian/tokenization_marian.py ~line 283
 *   src/transformers/models/m2m_100/tokenization_m2m_100.py ~line 216
 *   src/transformers/models/speech_to_text/tokenization_speech_to_text.py ~line 193
 */
public class Transformers0003ConvertTokensSpecialListTest {

    // Simulate `all_special_tokens` property — returns a new list each call
    // (in Python, the property calls convert_tokens_to_ids which reconstructs)
    static List<String> allSpecialTokensList(List<String> specialTokens) {
        return new ArrayList<>(specialTokens); // new list each time (mirrors Python property)
    }

    // --- Defective implementation ---
    static String convertTokensDefective(List<String> tokens, List<String> specialTokens) {
        StringBuilder current = new StringBuilder();
        StringBuilder out = new StringBuilder();
        for (String token : tokens) {
            // O(S) list scan per token — mirrors `if token in self.all_special_tokens`
            if (allSpecialTokensList(specialTokens).contains(token)) {
                out.append(current).append(token).append(" ");
                current.setLength(0);
            } else {
                current.append(token);
            }
        }
        out.append(current);
        return out.toString().trim();
    }

    // --- Patched implementation ---
    static String convertTokensPatched(List<String> tokens, List<String> specialTokens) {
        Set<String> specialSet = new HashSet<>(specialTokens); // cache once
        StringBuilder current = new StringBuilder();
        StringBuilder out = new StringBuilder();
        for (String token : tokens) {
            if (specialSet.contains(token)) { // O(1)
                out.append(current).append(token).append(" ");
                current.setLength(0);
            } else {
                current.append(token);
            }
        }
        out.append(current);
        return out.toString().trim();
    }

    public static void main(String[] args) {
        // Build special tokens list: 5 base tokens + 100 language codes (m2m100 style)
        int S = 108;
        List<String> specialTokens = new ArrayList<>();
        specialTokens.add("<unk>");
        specialTokens.add("<s>");
        specialTokens.add("</s>");
        specialTokens.add("<pad>");
        specialTokens.add("<mask>");
        for (int i = 0; i < 103; i++) {
            specialTokens.add("__lang" + i + "__");
        }

        // Build output token sequence: T=512, with a few special tokens mixed in
        int T = 512;
        List<String> tokens = new ArrayList<>(T);
        for (int i = 0; i < T; i++) {
            if (i % 50 == 0) {
                tokens.add("__lang0__"); // insert a special token
            } else {
                tokens.add("word" + (i % 1000));
            }
        }

        // Correctness check
        String defectOut = convertTokensDefective(tokens, specialTokens);
        String patchOut  = convertTokensPatched(tokens, specialTokens);
        assert defectOut.equals(patchOut)
            : "FAIL: defect and patched outputs differ:\n  defect=" + defectOut.substring(0, 50)
              + "\n  patch=" + patchOut.substring(0, 50);

        // Timing comparison
        int REPS = 2000;

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

        long t1 = System.nanoTime();
        for (int r = 0; r < REPS; r++) {
            convertTokensPatched(tokens, specialTokens);
        }
        long patchNs = System.nanoTime() - t1;

        double ratio = (double) defectNs / patchNs;
        System.out.printf("transformers-0003 MOAD-0001 convert_tokens_to_string O(T×S)%n");
        System.out.printf("  T=%d tokens, S=%d special tokens, %d reps%n", T, S, REPS);
        System.out.printf("  Defective (list scan + property call): %,d ms%n", defectNs / 1_000_000);
        System.out.printf("  Patched   (set cache before loop):      %,d ms%n", patchNs  / 1_000_000);
        System.out.printf("  Speedup: %.1fx%n", ratio);

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