import java.util.*;

/**
 * Unit test for trino-0002: SkewedPartitionRebalancer scaledPartitions ArrayList.contains O(P²)
 *
 * Simulates the defective and fixed versions of the scaledPartitions dedup check.
 */
public class TestSkewedPartitionRebalancer {

    // --- Defective: ArrayList.contains per partition ---
    static int simulateDefective(int numPartitions, int numTaskBuckets) {
        int totalCost = 0;
        List<Integer> scaledPartitions = new ArrayList<>();

        for (int bucket = 0; bucket < numTaskBuckets; bucket++) {
            for (int partition = 0; partition < numPartitions; partition++) {
                totalCost += scaledPartitions.size();  // cost of ArrayList.contains
                if (scaledPartitions.contains(partition)) {
                    continue;
                }
                scaledPartitions.add(partition);
            }
        }
        return totalCost;
    }

    // --- Fixed: HashSet.contains per partition ---
    static int simulateFixed(int numPartitions, int numTaskBuckets) {
        int totalCost = 0;
        Set<Integer> scaledPartitions = new HashSet<>();

        for (int bucket = 0; bucket < numTaskBuckets; bucket++) {
            for (int partition = 0; partition < numPartitions; partition++) {
                totalCost += 1;  // O(1) HashSet.contains
                if (scaledPartitions.contains(partition)) {
                    continue;
                }
                scaledPartitions.add(partition);
            }
        }
        return totalCost;
    }

    public static void main(String[] args) {
        int pass = 0, fail = 0;

        // Test 1: correctness - same partitions get deduplicated
        {
            int P = 20, B = 3;
            List<Integer> seenDef = new ArrayList<>();
            Set<Integer> seenFix = new HashSet<>();

            for (int b = 0; b < B; b++) {
                for (int p = 0; p < P; p++) {
                    if (!seenDef.contains(p)) seenDef.add(p);
                    seenFix.add(p);
                }
            }
            boolean ok = (seenDef.size() == seenFix.size()) &&
                          new HashSet<>(seenDef).equals(seenFix);
            System.out.printf("[%s] Correctness P=%d B=%d: defective=%d unique, fixed=%d unique%n",
                ok ? "PASS" : "FAIL", P, B, seenDef.size(), seenFix.size());
            if (ok) pass++; else fail++;
        }

        // Test 2: cost at P=100, B=5
        {
            int P = 100, B = 5;
            int defCost = simulateDefective(P, B);
            int fixCost = simulateFixed(P, B);
            double ratio = (double) defCost / fixCost;
            boolean ok = ratio >= 30.0;  // expect significant quadratic vs linear gap
            System.out.printf("[%s] Cost ratio P=%d B=%d: defective=%d, fixed=%d, ratio=%.1fx%n",
                ok ? "PASS" : "FAIL", P, B, defCost, fixCost, ratio);
            if (ok) pass++; else fail++;
        }

        // Test 3: cost at P=1000, B=10
        {
            int P = 1000, B = 10;
            int defCost = simulateDefective(P, B);
            int fixCost = simulateFixed(P, B);
            double ratio = (double) defCost / fixCost;
            boolean ok = ratio >= 200.0;
            System.out.printf("[%s] Cost ratio P=%d B=%d: defective=%d, fixed=%d, ratio=%.1fx%n",
                ok ? "PASS" : "FAIL", P, B, defCost, fixCost, ratio);
            if (ok) pass++; else fail++;
        }

        // Test 4: empty partitions case
        {
            int P = 0, B = 5;
            int defCost = simulateDefective(P, B);
            int fixCost = simulateFixed(P, B);
            boolean ok = (defCost == 0 && fixCost == 0);
            System.out.printf("[%s] Empty partitions P=%d B=%d: defective=%d, fixed=%d%n",
                ok ? "PASS" : "FAIL", P, B, defCost, fixCost);
            if (ok) pass++; else fail++;
        }

        // Test 5: single bucket - verify O(P²) vs O(P) scaling
        {
            int[] sizes = {10, 50, 100, 200};
            boolean allOk = true;
            for (int P : sizes) {
                int defCost = simulateDefective(P, 1);
                int fixCost = simulateFixed(P, 1);
                double ratio = (double) defCost / fixCost;
                System.out.printf("  P=%d: defective=%d, fixed=%d, ratio=%.1fx%n", P, defCost, fixCost, ratio);
                // ratio should grow roughly proportional to P
            }
            // Check that ratio at P=200 is >> ratio at P=10 (demonstrates quadratic growth)
            double ratioSmall = (double) simulateDefective(10, 1) / simulateFixed(10, 1);
            double ratioLarge = (double) simulateDefective(200, 1) / simulateFixed(200, 1);
            allOk = ratioLarge > ratioSmall * 5;  // quadratic: 200/10=20x more ratio
            System.out.printf("[%s] Quadratic scaling check: small ratio=%.1fx, large ratio=%.1fx%n",
                allOk ? "PASS" : "FAIL", ratioSmall, ratioLarge);
            if (allOk) pass++; else fail++;
        }

        System.out.printf("%nResult: %d PASS, %d FAIL%n", pass, fail);
        System.exit(fail > 0 ? 1 : 0);
    }
}