package unit;

import java.util.*;

/**
 * flink-0003: AggregateReduceGroupingRule — List<Integer>.contains() in for loop.
 *
 * newGroupingList is a List<Integer> from ImmutableBitSet.toList().
 * Each .contains(column) is O(G). With G grouping columns, total is O(G^2).
 * Fix: use Set<Integer> = new HashSet<>(newGrouping.toList()) for O(G).
 */
public class FlinkAggregateGroupingRuleTest {

    // Slow path: List<Integer>.contains() — simulates AggregateReduceGroupingRule.onMatch()
    // Returns (indexOldToNewMap, opCount)
    static Object[] slowBuildIndexMap(List<Integer> originalGrouping, List<Integer> newGroupingList) {
        Map<Integer, Integer> indexOldToNewMap = new HashMap<>();
        int idxOfNewGrouping = 0;
        int idxOfAggCallsForDropped = newGroupingList.size();
        int index = 0;
        long ops = 0;
        for (int column : originalGrouping) {
            // List.contains: linear scan
            boolean found = false;
            for (Integer g : newGroupingList) {
                ops++;
                if (g == column) { found = true; break; }
            }
            if (found) {
                indexOldToNewMap.put(index, idxOfNewGrouping++);
            } else {
                indexOldToNewMap.put(index, idxOfAggCallsForDropped++);
            }
            index++;
        }
        return new Object[]{indexOldToNewMap, ops};
    }

    // Fast path: Set<Integer>.contains() — O(1) lookup
    static Object[] fastBuildIndexMap(List<Integer> originalGrouping, Set<Integer> newGroupingSet, int newGroupingSize) {
        Map<Integer, Integer> indexOldToNewMap = new HashMap<>();
        int idxOfNewGrouping = 0;
        int idxOfAggCallsForDropped = newGroupingSize;
        int index = 0;
        long ops = 0;
        for (int column : originalGrouping) {
            ops++; // O(1) HashSet lookup
            if (newGroupingSet.contains(column)) {
                indexOldToNewMap.put(index, idxOfNewGrouping++);
            } else {
                indexOldToNewMap.put(index, idxOfAggCallsForDropped++);
            }
            index++;
        }
        return new Object[]{indexOldToNewMap, ops};
    }

    public static void main(String[] args) {
        int passed = 0;
        int total = 0;

        // Test 1: correctness — all columns retained
        {
            total++;
            List<Integer> original = Arrays.asList(0, 1, 2, 3, 4);
            List<Integer> newGrouping = Arrays.asList(0, 1, 2, 3, 4);
            Set<Integer> newGroupingSet = new HashSet<>(newGrouping);

            @SuppressWarnings("unchecked")
            Map<Integer,Integer> slowMap = (Map<Integer,Integer>) slowBuildIndexMap(original, newGrouping)[0];
            @SuppressWarnings("unchecked")
            Map<Integer,Integer> fastMap = (Map<Integer,Integer>) fastBuildIndexMap(original, newGroupingSet, newGrouping.size())[0];

            assert slowMap.equals(fastMap)
                : "All retained: slow and fast must agree. slow=" + slowMap + " fast=" + fastMap;
            // All mapped to 0..4
            for (int i = 0; i < 5; i++) {
                assert slowMap.get(i) == i : "column " + i + " should map to " + i;
            }
            System.out.println("  Test 1 PASS: all-retained case, map=" + slowMap);
            passed++;
        }

        // Test 2: correctness — some columns dropped
        {
            total++;
            List<Integer> original = Arrays.asList(0, 1, 2, 3, 4);
            List<Integer> newGrouping = Arrays.asList(1, 3); // only cols 1 and 3 are unique
            Set<Integer> newGroupingSet = new HashSet<>(newGrouping);

            @SuppressWarnings("unchecked")
            Map<Integer,Integer> slowMap = (Map<Integer,Integer>) slowBuildIndexMap(original, newGrouping)[0];
            @SuppressWarnings("unchecked")
            Map<Integer,Integer> fastMap = (Map<Integer,Integer>) fastBuildIndexMap(original, newGroupingSet, newGrouping.size())[0];

            assert slowMap.equals(fastMap)
                : "Partial drop: slow and fast must agree. slow=" + slowMap + " fast=" + fastMap;
            // col 0 not in newGrouping → agg call index (starts at 2)
            assert slowMap.get(0) == 2 : "col 0 not in newGrouping → idx 2, got " + slowMap.get(0);
            assert slowMap.get(1) == 0 : "col 1 in newGrouping → idx 0, got " + slowMap.get(1);
            assert slowMap.get(2) == 3 : "col 2 not in newGrouping → idx 3, got " + slowMap.get(2);
            assert slowMap.get(3) == 1 : "col 3 in newGrouping → idx 1, got " + slowMap.get(3);
            assert slowMap.get(4) == 4 : "col 4 not in newGrouping → idx 4, got " + slowMap.get(4);
            System.out.println("  Test 2 PASS: partial-drop case, map=" + slowMap);
            passed++;
        }

        // Test 3: op count — O(G^2) vs O(G)
        {
            total++;
            int G = 60; // grouping columns
            List<Integer> original = new ArrayList<>();
            List<Integer> newGrouping = new ArrayList<>();
            for (int i = 0; i < G; i++) {
                original.add(i);
                newGrouping.add(i);
            }
            Set<Integer> newGroupingSet = new HashSet<>(newGrouping);

            long slowOps = (Long) slowBuildIndexMap(original, newGrouping)[1];
            long fastOps = (Long) fastBuildIndexMap(original, newGroupingSet, G)[1];

            // Slow: best case (all found on first scan) = G * 1 = G ops
            // Worst case: not found = G * G ops
            // In our worst-case-measuring sim above, found on first match (column == 0..G-1 sequential)
            // so slowOps = sum(pos+1) for each element. Let's just verify slowOps > fastOps
            // and fastOps == G.
            assert fastOps == G : "Fast path must make exactly G=" + G + " ops, got " + fastOps;
            assert slowOps >= G : "Slow path must make at least G ops";

            long speedup = slowOps / fastOps;
            System.out.println("  Test 3 PASS: slowOps=" + slowOps + " fastOps=" + fastOps + " speedup=" + speedup + "x");
            passed++;
        }

        // Test 4: worst-case O(G^2) when target always at end of list
        {
            total++;
            int G = 50;
            // original grouping has elements NOT in newGrouping → every contains() scans all G
            List<Integer> original = new ArrayList<>();
            List<Integer> newGrouping = new ArrayList<>();
            // newGrouping = [0..G-1], original = [G..2G-1] (no overlap → all scan full list)
            for (int i = 0; i < G; i++) newGrouping.add(i);
            for (int i = G; i < 2 * G; i++) original.add(i);

            Set<Integer> newGroupingSet = new HashSet<>(newGrouping);

            // slow: each original element scans all G elements of newGrouping (no match)
            long expectedSlowOps = (long) G * G; // G * G scans
            long actualSlowOps = 0;
            for (int col : original) {
                for (int g : newGrouping) {
                    actualSlowOps++;
                    if (g == col) break; // never matches
                }
            }

            assert actualSlowOps == expectedSlowOps
                : "Expected " + expectedSlowOps + " slow ops, got " + actualSlowOps;

            long fastOps = G; // G × O(1)
            long speedup = actualSlowOps / fastOps;

            assert speedup == G : "Speedup should equal G=" + G + " in worst case, got " + speedup;
            System.out.println("  Test 4 PASS: O(G^2)=" + actualSlowOps + " vs O(G)=" + fastOps + " speedup=" + speedup + "x (G=" + G + ")");
            passed++;
        }

        System.out.println(passed + "/" + total + " PASS");
        if (passed != total) System.exit(1);
    }
}