package unit;

import java.util.*;

/**
 * airflow-0001 — O(N²) topological sort in TaskGroup.topological_sort()
 *
 * Simulates the defective "modified Kahn's" algorithm from:
 *   task-sdk/src/airflow/sdk/definitions/taskgroup.py:536-568
 *
 * Defect: Each round of the outer while loop rescans ALL remaining unsorted nodes
 * (inner for loop), doing O(N) work per round. Worst case (linear chain): O(N²).
 *
 * Fix: Pre-compute in-degrees and maintain a ready-queue (true Kahn's), O(N + E).
 *
 * The op-count measures "node examination" — how many times a node is inspected
 * in the scan loop. In the defective version every remaining node is examined
 * each round; in the fixed version each node is processed exactly once.
 */
public class AirflowTopoSortAlgorithm {

    static void check(String desc, boolean cond) {
        System.out.println((cond ? "PASS" : "FAIL") + ": " + desc);
        if (!cond) throw new AssertionError("FAIL: " + desc);
    }

    // Node: simulates TaskGroup child with upstream_list
    static class Node {
        final String id;
        final List<String> upstreamIds = new ArrayList<>();

        Node(String id) { this.id = id; }
        void addUpstream(String uid) { upstreamIds.add(uid); }
    }

    // Result carrying an op-count alongside the sorted list
    static class SortResult {
        final List<String> sorted;
        final long nodeExaminations; // how many times a node was pulled from the scan

        SortResult(List<String> sorted, long nodeExaminations) {
            this.sorted = sorted;
            this.nodeExaminations = nodeExaminations;
        }
    }

    // -----------------------------------------------------------------------
    // DEFECTIVE: O(N²) "modified Kahn's" — direct port of Airflow Python
    // graph_unsorted is Map<id, Node>, mimics Python dict
    // Each pass through while scans ALL remaining nodes
    // -----------------------------------------------------------------------
    static SortResult defectiveTopoSort(Map<String, Node> nodes) {
        Map<String, Node> graphUnsorted = new LinkedHashMap<>(nodes);
        List<String> graphSorted = new ArrayList<>();
        long examinations = 0;

        while (!graphUnsorted.isEmpty()) {
            boolean acyclic = false;
            for (Node node : new ArrayList<>(graphUnsorted.values())) {
                examinations++; // Each node examined once per pass
                boolean blocked = false;
                for (String upId : node.upstreamIds) {
                    if (graphUnsorted.containsKey(upId)) {
                        blocked = true;
                        break;
                    }
                }
                if (!blocked) {
                    acyclic = true;
                    graphUnsorted.remove(node.id);
                    graphSorted.add(node.id);
                }
            }
            if (!acyclic) throw new RuntimeException("Cycle detected");
        }
        return new SortResult(graphSorted, examinations);
    }

    // -----------------------------------------------------------------------
    // FIXED: True Kahn's with ready-queue — O(N + E)
    // Each node is dequeued exactly once → examinations == N
    // -----------------------------------------------------------------------
    static SortResult fixedTopoSort(Map<String, Node> nodes) {
        Map<String, Integer> inDegree = new HashMap<>();
        Map<String, List<String>> dependents = new HashMap<>();

        for (String id : nodes.keySet()) {
            inDegree.put(id, 0);
            dependents.put(id, new ArrayList<>());
        }

        for (Node node : nodes.values()) {
            for (String upId : node.upstreamIds) {
                if (nodes.containsKey(upId)) {
                    inDegree.merge(node.id, 1, Integer::sum);
                    dependents.get(upId).add(node.id);
                }
            }
        }

        Deque<String> ready = new ArrayDeque<>();
        for (Map.Entry<String, Integer> e : inDegree.entrySet()) {
            if (e.getValue() == 0) ready.add(e.getKey());
        }

        List<String> result = new ArrayList<>();
        long examinations = 0;
        while (!ready.isEmpty()) {
            String id = ready.poll();
            examinations++; // Each node dequeued exactly once
            result.add(id);
            for (String dep : dependents.get(id)) {
                inDegree.merge(dep, -1, Integer::sum);
                if (inDegree.get(dep) == 0) ready.add(dep);
            }
        }

        if (result.size() != nodes.size()) throw new RuntimeException("Cycle detected");
        return new SortResult(result, examinations);
    }

    // -----------------------------------------------------------------------
    // Build a reversed linear chain: nodes inserted in dependency-last order.
    //
    // Insertion order: t(N-1), t(N-2), ..., t1, t0
    // Dependency:      t(k) depends on t(k-1)
    //
    // When graph_unsorted preserves insertion order (LinkedHashMap, like Python dict),
    // each round of the while loop scans the full snapshot but can only free the LAST
    // node (t0) on the first pass, then (t1) on the second, etc.
    //
    // Round 1: examine all N nodes, free t0  (it has no upstreams)
    // Round 2: examine remaining N-1 nodes, free t1 (t0 now gone)
    // ...
    // Total: N + (N-1) + ... + 1 = N(N+1)/2 = O(N²) node examinations
    // -----------------------------------------------------------------------
    static Map<String, Node> buildLinearChain(int n) {
        Map<String, Node> nodes = new LinkedHashMap<>();
        // Insert in REVERSE order (t(N-1) first) — worst case for the defective algo
        for (int i = n - 1; i >= 0; i--) {
            nodes.put("t" + i, new Node("t" + i));
        }
        // t(k) depends on t(k-1): dependency is at the END of the insertion order
        for (int i = 1; i < n; i++) {
            nodes.get("t" + i).addUpstream("t" + (i - 1));
        }
        return nodes;
    }

    // Build a wide DAG: single source, all others depend on it
    static Map<String, Node> buildWideDag(int n) {
        Map<String, Node> nodes = new LinkedHashMap<>();
        nodes.put("root", new Node("root"));
        for (int i = 1; i < n; i++) {
            Node node = new Node("t" + i);
            node.addUpstream("root");
            nodes.put("t" + i, node);
        }
        return nodes;
    }

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

        // ------------------------------------------------------------------
        // Test 1: Correctness on small chain (3 nodes)
        // ------------------------------------------------------------------
        total++;
        {
            Map<String, Node> dag = buildLinearChain(3);
            List<String> defResult = defectiveTopoSort(new LinkedHashMap<>(dag)).sorted;
            List<String> fixResult = fixedTopoSort(new LinkedHashMap<>(dag)).sorted;
            // Both must produce a valid topological order: t0 before t1 before t2
            boolean defOk = defResult.indexOf("t0") < defResult.indexOf("t1")
                         && defResult.indexOf("t1") < defResult.indexOf("t2");
            boolean fixOk = fixResult.indexOf("t0") < fixResult.indexOf("t1")
                         && fixResult.indexOf("t1") < fixResult.indexOf("t2");
            check("Correctness: defective produces valid topo order for chain-3", defOk);
            check("Correctness: fixed produces valid topo order for chain-3", fixOk);
            passed += 2;
            total++;
        }

        // ------------------------------------------------------------------
        // Test 2: Correctness on wide dag
        // ------------------------------------------------------------------
        total++;
        {
            Map<String, Node> dag = buildWideDag(11);
            List<String> defResult = defectiveTopoSort(new LinkedHashMap<>(dag)).sorted;
            List<String> fixResult = fixedTopoSort(new LinkedHashMap<>(dag)).sorted;
            boolean rootFirst = defResult.get(0).equals("root") && fixResult.get(0).equals("root");
            check("Correctness: wide dag — root is first in both", rootFirst);
            passed++;
        }

        // ------------------------------------------------------------------
        // Test 3: Defective shows O(N²) node examinations on linear chain
        // For N=200 chain: expected = 200 + 199 + ... + 1 = 200*201/2 = 20100
        // Fixed: expected = N = 200
        // ------------------------------------------------------------------
        total++;
        {
            int n = 200;
            long expected = (long) n * (n + 1) / 2;
            SortResult defR = defectiveTopoSort(new LinkedHashMap<>(buildLinearChain(n)));
            SortResult fixR = fixedTopoSort(new LinkedHashMap<>(buildLinearChain(n)));

            System.out.printf("  N=%d chain: defective=%d exams (expected %d), fixed=%d exams (expected %d)%n",
                              n, defR.nodeExaminations, expected, fixR.nodeExaminations, n);

            check("Defective: N=200 chain examinations == N(N+1)/2=" + expected,
                  defR.nodeExaminations == expected);
            check("Fixed: N=200 chain examinations == N=" + n,
                  fixR.nodeExaminations == n);
            passed += 2;
            total++;
        }

        // ------------------------------------------------------------------
        // Test 4: Ratio confirms O(N²) vs O(N)
        // ------------------------------------------------------------------
        total++;
        {
            int n = 500;
            SortResult defR = defectiveTopoSort(new LinkedHashMap<>(buildLinearChain(n)));
            SortResult fixR = fixedTopoSort(new LinkedHashMap<>(buildLinearChain(n)));

            double ratio = (double) defR.nodeExaminations / fixR.nodeExaminations;
            System.out.printf("  N=%d chain: ratio = %.1fx%n", n, ratio);
            // For N=500: defective = 500*501/2 = 125250, fixed = 500 → ratio = 250.5
            check("Op-count ratio >= 100x for linear chain N=" + n, ratio >= 100.0);
            passed++;
        }

        // ------------------------------------------------------------------
        // Test 5: Quadratic growth — defective at N=200 should be ~4x N=100
        // N=100: 100*101/2 = 5050; N=200: 200*201/2 = 20100; ratio = 20100/5050 = 3.98
        // ------------------------------------------------------------------
        total++;
        {
            long ops100 = defectiveTopoSort(new LinkedHashMap<>(buildLinearChain(100))).nodeExaminations;
            long ops200 = defectiveTopoSort(new LinkedHashMap<>(buildLinearChain(200))).nodeExaminations;
            double growthRatio = (double) ops200 / ops100;
            System.out.printf("  Defective growth ratio (N=200 vs N=100): %.2fx (expected ~3.98x for O(N²))%n",
                              growthRatio);
            check("Defective shows quadratic growth (ratio >= 3.5x)", growthRatio >= 3.5);
            passed++;
        }

        // ------------------------------------------------------------------
        // Test 6: Linear growth — fixed at N=200 should be ~2x N=100
        // ------------------------------------------------------------------
        total++;
        {
            long ops100 = fixedTopoSort(new LinkedHashMap<>(buildLinearChain(100))).nodeExaminations;
            long ops200 = fixedTopoSort(new LinkedHashMap<>(buildLinearChain(200))).nodeExaminations;
            double growthRatio = (double) ops200 / ops100;
            System.out.printf("  Fixed growth ratio (N=200 vs N=100): %.2fx (expected ~2x for O(N))%n",
                              growthRatio);
            check("Fixed shows linear growth (ratio < 3.0x)", growthRatio < 3.0);
            passed++;
        }

        System.out.printf("%n%d/%d PASS%n", passed, total);
    }
}