package unit;

import java.util.*;

/**
 * CWE-407 unit test — spark-0002
 *
 * DAGScheduler BFS traversals use ListBuffer.remove(0) which is O(N) per call
 * (ArrayList/ListBuffer shift every remaining element left), making the overall
 * BFS O(N²) instead of O(N).
 *
 * This test simulates both the defective (ArrayList.remove(0)) and fixed
 * (ArrayDeque.removeFirst()) BFS traversal over fan-out DAGs, counting element
 * shift operations to prove the complexity difference.
 *
 * Topology that triggers O(N²): root node has N leaf children.
 * After visiting root, the queue holds N items. Each subsequent remove(0)
 * costs N-k shifts for the k-th dequeue → total (N-1) + (N-2) + ... + 0 = N(N-1)/2.
 *
 * No JUnit. Run: javac -d . SparkDAGSchedulerTest.java && java -ea unit.SparkDAGSchedulerTest
 */
public class SparkDAGSchedulerTest {

    // --- Node model: each node has a list of dependencies (parents in RDD lineage) ---

    static class RDDNode {
        final int id;
        final List<RDDNode> deps;

        RDDNode(int id, List<RDDNode> deps) {
            this.id = id;
            this.deps = deps;
        }
    }

    // -------------------------------------------------------------------------
    // Slow BFS: ArrayList simulating ListBuffer.remove(0)
    // remove(0) on ArrayList shifts all remaining elements → O(N) per call.
    // We track shift operations explicitly.
    // -------------------------------------------------------------------------

    /** Result of slow BFS: visited node IDs and total shift-op count. */
    static long[] bfsSlowListRemove(RDDNode start) {
        List<Integer> visited = new ArrayList<>();
        Set<Integer> seen = new HashSet<>();
        List<RDDNode> queue = new ArrayList<>();   // simulates ListBuffer
        queue.add(start);
        long shiftOps = 0;

        while (!queue.isEmpty()) {
            int sizeBefore = queue.size();
            RDDNode node = queue.remove(0);          // O(sizeBefore-1) shifts
            shiftOps += (sizeBefore - 1);

            if (!seen.contains(node.id)) {
                seen.add(node.id);
                visited.add(node.id);
                // Prepend deps — like waitingForVisit.prepend(dep) in Spark
                for (RDDNode dep : node.deps) {
                    shiftOps += queue.size();        // prepend shifts all current items right
                    queue.add(0, dep);
                }
            }
        }
        return new long[]{visited.size(), shiftOps};
    }

    // -------------------------------------------------------------------------
    // Fast BFS: ArrayDeque.removeFirst() is O(1) amortized.
    // -------------------------------------------------------------------------

    /** Result of fast BFS: visited node IDs and total dequeue-op count. */
    static long[] bfsFastArrayDeque(RDDNode start) {
        List<Integer> visited = new ArrayList<>();
        Set<Integer> seen = new HashSet<>();
        ArrayDeque<RDDNode> queue = new ArrayDeque<>();
        queue.add(start);
        long dequeueOps = 0;

        while (!queue.isEmpty()) {
            RDDNode node = queue.removeFirst();       // O(1)
            dequeueOps += 1;                          // constant work per dequeue

            if (!seen.contains(node.id)) {
                seen.add(node.id);
                visited.add(node.id);
                for (RDDNode dep : node.deps) {
                    queue.addFirst(dep);              // O(1) prepend
                }
            }
        }
        return new long[]{visited.size(), dequeueOps};
    }

    // -------------------------------------------------------------------------
    // DAG builders
    // -------------------------------------------------------------------------

    /** Root node with N leaf children — triggers O(N²) on ListBuffer.remove(0). */
    static RDDNode buildFanOutDAG(int n) {
        List<RDDNode> children = new ArrayList<>();
        for (int i = 1; i <= n; i++) {
            children.add(new RDDNode(i, Collections.emptyList()));
        }
        return new RDDNode(0, children);
    }

    /**
     * Two-level fan-out: root has K children, each child has K leaf grandchildren.
     * Total = 1 + K + K² nodes; queue grows to K then K² — maximises O(N²) gap.
     */
    static RDDNode buildTwoLevelFanOut(int k) {
        int id = 0;
        List<RDDNode> level1 = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            List<RDDNode> leaves = new ArrayList<>();
            for (int j = 0; j < k; j++) {
                leaves.add(new RDDNode(++id, Collections.emptyList()));
            }
            level1.add(new RDDNode(++id, leaves));
        }
        return new RDDNode(0, level1);
    }

    /** Binary tree of given depth. */
    static RDDNode buildBinaryTree(int depth, int[] counter) {
        int id = counter[0]++;
        if (depth == 0) return new RDDNode(id, Collections.emptyList());
        return new RDDNode(id, Arrays.asList(
            buildBinaryTree(depth - 1, counter),
            buildBinaryTree(depth - 1, counter)));
    }

    // -------------------------------------------------------------------------
    // Helpers
    // -------------------------------------------------------------------------

    static void test(String name, boolean condition) {
        if (!condition) throw new AssertionError("FAIL: " + name);
        System.out.println("PASS: " + name);
    }

    // -------------------------------------------------------------------------
    // Main
    // -------------------------------------------------------------------------

    public static void main(String[] args) {
        System.out.println("=== spark-0002: DAGScheduler ListBuffer.remove(0) O(N²) BFS ===");
        System.out.println();

        // --- T1: correctness — fan-out N=10 ---
        {
            RDDNode dag = buildFanOutDAG(10);
            long[] slow = bfsSlowListRemove(dag);
            long[] fast = bfsFastArrayDeque(dag);
            // 1 root + 10 leaves = 11 nodes visited
            test("T1: slow BFS visits all 11 nodes (N=10 fan-out)", slow[0] == 11);
            test("T1: fast BFS visits all 11 nodes (N=10 fan-out)", fast[0] == 11);
        }

        // --- T2: correctness — two-level K=5 fan-out ---
        {
            RDDNode dag = buildTwoLevelFanOut(5);
            long[] slow = bfsSlowListRemove(dag);
            long[] fast = bfsFastArrayDeque(dag);
            // 1 + 5 + 25 = 31 nodes
            test("T2: slow BFS visits all 31 nodes (two-level K=5)", slow[0] == 31);
            test("T2: fast BFS visits all 31 nodes (two-level K=5)", fast[0] == 31);
        }

        // --- T3: correctness — binary tree depth=5 ---
        {
            int[] counter = {0};
            RDDNode tree = buildBinaryTree(5, counter);
            int expected = counter[0]; // 63 nodes
            long[] slow = bfsSlowListRemove(tree);
            long[] fast = bfsFastArrayDeque(tree);
            test("T3: slow BFS visits all " + expected + " tree nodes", slow[0] == expected);
            test("T3: fast BFS visits all " + expected + " tree nodes", fast[0] == expected);
        }

        // --- T4: O(N²) proof — fan-out N=50 ---
        {
            int N = 50;
            RDDNode dag = buildFanOutDAG(N);
            long[] slow = bfsSlowListRemove(dag);
            long[] fast = bfsFastArrayDeque(dag);
            System.out.printf("T4: N=%d fan-out — slow shift-ops=%d, fast dequeue-ops=%d%n",
                N, slow[1], fast[1]);
            // Slow: prepend N children shifts 0..N-1 → N(N-1)/2 shifts; plus remove(0) costs.
            // Fast: exactly N+1 dequeue ops (root + N leaves).
            test("T4: slow shift-ops > fast dequeue-ops for N=50 fan-out",
                slow[1] > fast[1]);
            test("T4: slow shift-ops >= N*(N-1)/2 (O(N²) lower bound)",
                slow[1] >= (long) N * (N - 1) / 2);
            test("T4: fast dequeue-ops == N+1 (O(N) confirmed)",
                fast[1] == N + 1);
        }

        // --- T5: O(N²) scaling — doubling N should ~4x slow ops, ~2x fast ops ---
        {
            int N1 = 100;
            int N2 = 200;

            long[] slowN1 = bfsSlowListRemove(buildFanOutDAG(N1));
            long[] fastN1 = bfsFastArrayDeque(buildFanOutDAG(N1));
            long[] slowN2 = bfsSlowListRemove(buildFanOutDAG(N2));
            long[] fastN2 = bfsFastArrayDeque(buildFanOutDAG(N2));

            double slowRatio = (double) slowN2[1] / slowN1[1];
            double fastRatio = (double) fastN2[1] / fastN1[1];

            System.out.printf("T5: N=%d slow=%d fast=%d%n", N1, slowN1[1], fastN1[1]);
            System.out.printf("T5: N=%d slow=%d fast=%d%n", N2, slowN2[1], fastN2[1]);
            System.out.printf("T5: slow ops ratio %.2f (expect ~4.0 for O(N²))%n", slowRatio);
            System.out.printf("T5: fast ops ratio %.2f (expect ~2.0 for O(N))%n", fastRatio);

            test("T5: slow ratio >= 3.5 (super-linear, confirming O(N²))", slowRatio >= 3.5);
            test("T5: fast ratio <= 2.1 (linear, confirming O(N))", fastRatio <= 2.1);
            test("T5: fast ratio >= 1.9 (not sub-linear)", fastRatio >= 1.9);
        }

        // --- T6: two-level fan-out — larger queue, larger gap ---
        {
            int K1 = 20;  // 1 + 20 + 400 = 421 nodes
            int K2 = 40;  // 1 + 40 + 1600 = 1641 nodes
            long[] slowK1 = bfsSlowListRemove(buildTwoLevelFanOut(K1));
            long[] fastK1 = bfsFastArrayDeque(buildTwoLevelFanOut(K1));
            long[] slowK2 = bfsSlowListRemove(buildTwoLevelFanOut(K2));
            long[] fastK2 = bfsFastArrayDeque(buildTwoLevelFanOut(K2));

            System.out.printf("T6: K=%d (%d nodes) — slow=%d fast=%d%n",
                K1, (int)slowK1[0], slowK1[1], fastK1[1]);
            System.out.printf("T6: K=%d (%d nodes) — slow=%d fast=%d%n",
                K2, (int)slowK2[0], slowK2[1], fastK2[1]);

            test("T6: slow shift-ops >> fast dequeue-ops at K=20",
                slowK1[1] > fastK1[1] * 10);
            test("T6: slow shift-ops >> fast dequeue-ops at K=40",
                slowK2[1] > fastK2[1] * 10);
        }

        // --- T7: wall-clock — fast must be faster for large N ---
        {
            int N = 3000;
            RDDNode dag = buildFanOutDAG(N);

            long t0 = System.nanoTime();
            long[] slow = bfsSlowListRemove(dag);
            long slowNs = System.nanoTime() - t0;

            t0 = System.nanoTime();
            long[] fast = bfsFastArrayDeque(dag);
            long fastNs = System.nanoTime() - t0;

            double speedup = (double) slowNs / Math.max(fastNs, 1);
            System.out.printf("T7: N=%d fan-out — slow=%.3fms fast=%.3fms speedup=%.1fx%n",
                N, slowNs / 1e6, fastNs / 1e6, speedup);

            test("T7: slow BFS visits N+1=" + (N+1) + " nodes", slow[0] == N + 1);
            test("T7: fast BFS visits N+1=" + (N+1) + " nodes", fast[0] == N + 1);
            test("T7: slow shift-ops confirms O(N²) — >= N*(N-1)/2",
                slow[1] >= (long) N * (N - 1) / 2);
            test("T7: fast dequeue-ops == N+1 (pure O(N))",
                fast[1] == N + 1);
        }

        System.out.println();
        System.out.println("ALL PASS — spark-0002: ListBuffer.remove(0) is O(N²);" +
            " ArrayDeque.removeFirst() is O(N) — 6 BFS functions in DAGScheduler.scala");
    }
}