package unit;

import java.util.*;

/**
 * kubeflow-0001 — O(T²×I) pipeline DAG compilation in pipeline_spec_builder.py
 *
 * Simulates the defective pattern from:
 *   sdk/python/kfp/compiler/pipeline_spec_builder.py  lines 1383–1540
 *
 * Defect:
 *   tasks_in_current_dag is a List[str] REBUILT inside the outer loop over T tasks,
 *   and then each task performs O(I) input lookups each doing O(T) list `in` check.
 *   Total: O(T) outer × O(I) inputs × O(T) in-check = O(T² × I).
 *
 * Fix:
 *   Build tasks_in_current_dag once as a Set[str] outside the loop.
 *   Each `in` check becomes O(1).
 *   Total: O(T × I).
 *
 * Op-count measures: number of string comparisons performed across all `in` checks.
 */
public class KubeflowTaskDagAlgorithm {

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

    // Simulates a PipelineTask with input channels that may reference other tasks
    static class Task {
        final String name;
        final List<String> inputProducers; // task names that produce each input (may be null)

        Task(String name, List<String> inputProducers) {
            this.name = name;
            this.inputProducers = inputProducers;
        }
    }

    // Simulates build result: which inputs are "task outputs" vs "component inputs"
    static class TaskSpec {
        final Map<String, String> taskOutputInputs = new LinkedHashMap<>();    // input → producer task
        final Map<String, String> componentInputs = new LinkedHashMap<>();     // input → component input
    }

    // Result carrying op-count
    static class Result {
        final List<TaskSpec> taskSpecs;
        final long inChecks; // total string comparisons
        Result(List<TaskSpec> specs, long checks) {
            this.taskSpecs = specs;
            this.inChecks = checks;
        }
    }

    // -----------------------------------------------------------------------
    // DEFECTIVE: List[str] rebuilt inside outer loop, O(T) in-check per input
    // -----------------------------------------------------------------------
    static long defectiveChecks;

    static boolean listContains(List<String> list, String item) {
        for (String s : list) {
            defectiveChecks++;
            if (s.equals(item)) return true;
        }
        return false;
    }

    static TaskSpec buildTaskSpecDefective(Task task, List<String> tasksInCurrentDag) {
        TaskSpec spec = new TaskSpec();
        for (String producer : task.inputProducers) {
            if (producer == null) {
                spec.componentInputs.put(task.name + "_null", "pipeline_input");
            } else if (listContains(tasksInCurrentDag, producer)) {
                spec.taskOutputInputs.put(task.name + "_from_" + producer, producer);
            } else {
                spec.componentInputs.put(task.name + "_from_" + producer, "outer_" + producer);
            }
        }
        return spec;
    }

    static Result runDefective(List<Task> tasks) {
        defectiveChecks = 0;
        List<TaskSpec> specs = new ArrayList<>();

        for (Task task : tasks) {
            // DEFECTIVE: rebuilt on every iteration of the outer loop
            List<String> tasksInCurrentDag = new ArrayList<>();
            for (Task t : tasks) {
                tasksInCurrentDag.add(t.name);
            }
            specs.add(buildTaskSpecDefective(task, tasksInCurrentDag));
        }
        return new Result(specs, defectiveChecks);
    }

    // -----------------------------------------------------------------------
    // FIXED: Set[str] built once outside the loop, O(1) in-check per input
    // -----------------------------------------------------------------------
    static long fixedChecks;

    static boolean setContains(Set<String> set, String item) {
        fixedChecks++;  // O(1) hash lookup counted as 1
        return set.contains(item);
    }

    static TaskSpec buildTaskSpecFixed(Task task, Set<String> tasksInCurrentDag) {
        TaskSpec spec = new TaskSpec();
        for (String producer : task.inputProducers) {
            if (producer == null) {
                spec.componentInputs.put(task.name + "_null", "pipeline_input");
            } else if (setContains(tasksInCurrentDag, producer)) {
                spec.taskOutputInputs.put(task.name + "_from_" + producer, producer);
            } else {
                spec.componentInputs.put(task.name + "_from_" + producer, "outer_" + producer);
            }
        }
        return spec;
    }

    static Result runFixed(List<Task> tasks) {
        fixedChecks = 0;
        List<TaskSpec> specs = new ArrayList<>();

        // FIXED: built once outside the loop
        Set<String> tasksInCurrentDag = new LinkedHashSet<>();
        for (Task t : tasks) {
            tasksInCurrentDag.add(t.name);
        }

        for (Task task : tasks) {
            specs.add(buildTaskSpecFixed(task, tasksInCurrentDag));
        }
        return new Result(specs, fixedChecks);
    }

    // -----------------------------------------------------------------------
    // Build test pipeline: T tasks, each with I inputs referencing earlier tasks
    // -----------------------------------------------------------------------
    static List<Task> buildPipeline(int T, int I) {
        List<Task> tasks = new ArrayList<>();
        for (int t = 0; t < T; t++) {
            List<String> inputs = new ArrayList<>();
            for (int i = 0; i < I; i++) {
                if (t > 0) {
                    // reference a task within the same DAG
                    inputs.add("task" + (t - 1));
                } else {
                    inputs.add(null); // pipeline-level input
                }
            }
            tasks.add(new Task("task" + t, inputs));
        }
        return tasks;
    }

    // -----------------------------------------------------------------------
    // Verify both produce identical specs
    // -----------------------------------------------------------------------
    static boolean specsEqual(List<TaskSpec> a, List<TaskSpec> b) {
        if (a.size() != b.size()) return false;
        for (int i = 0; i < a.size(); i++) {
            if (!a.get(i).taskOutputInputs.equals(b.get(i).taskOutputInputs)) return false;
            if (!a.get(i).componentInputs.equals(b.get(i).componentInputs)) return false;
        }
        return true;
    }

    public static void main(String[] args) {
        System.out.println("=== kubeflow-0001: O(T²×I) tasks_in_current_dag List in pipeline compiler ===");

        // ------ Test 1: correctness on small pipeline ------
        {
            List<Task> tasks = buildPipeline(5, 3);
            Result def = runDefective(tasks);
            Result fix = runFixed(tasks);
            check("pipeline-5x3: same task specs (defective == fixed)", specsEqual(def.taskSpecs, fix.taskSpecs));
            check("pipeline-5x3: fixed uses fewer checks", fix.inChecks <= def.inChecks);
        }

        // ------ Test 2: O(T²×I) vs O(T×I) scaling ------
        {
            int I = 5;
            int T = 200;
            List<Task> tasks = buildPipeline(T, I);

            Result def = runDefective(tasks);
            Result fix = runFixed(tasks);

            long ratio = def.inChecks / Math.max(fix.inChecks, 1);
            System.out.printf("  T=%d I=%d: defective checks=%d, fixed checks=%d, ratio=%dx%n",
                T, I, def.inChecks, fix.inChecks, ratio);

            // Defective: each of T tasks does I × O(T) checks = T² × I total
            // Expected ~200² × 5 / 2 avg = ~100,000
            check("T=200,I=5: defective checks > T*I*5 (quadratic evidence)",
                def.inChecks > (long) T * I * 5);
            // Fixed: T × I × 1 = 1000
            // task0 has null producers so its I inputs are skipped before the set check;
            // remaining T-1 tasks each do I checks → (T-1)*I total (bounded by T*I)
            check("T=200,I=5: fixed checks <= T*I (linear)",
                fix.inChecks <= (long) T * I);
            check("T=200,I=5: ratio >= 10x",
                ratio >= 10);

            check("T=200,I=5: same specs", specsEqual(def.taskSpecs, fix.taskSpecs));
        }

        // ------ Test 3: quadratic growth is clear ------
        {
            int I = 3;
            List<Task> tasks100 = buildPipeline(100, I);
            List<Task> tasks400 = buildPipeline(400, I);

            Result def100 = runDefective(tasks100);
            Result def400 = runDefective(tasks400);

            double growthRatio = (double) def400.inChecks / Math.max(def100.inChecks, 1);
            System.out.printf("  Defective: T=100 checks=%d, T=400 checks=%d, growth=%.1fx%n",
                def100.inChecks, def400.inChecks, growthRatio);
            check("quadratic growth: def400 > 4x def100 (O(T²) evidence)", growthRatio > 4.0);

            Result fix100 = runFixed(tasks100);
            Result fix400 = runFixed(tasks400);
            double fixGrowth = (double) fix400.inChecks / Math.max(fix100.inChecks, 1);
            System.out.printf("  Fixed:     T=100 checks=%d, T=400 checks=%d, growth=%.1fx%n",
                fix100.inChecks, fix400.inChecks, fixGrowth);
            check("linear growth: fix400 ≈ 4x fix100 (O(T) evidence)",
                fixGrowth >= 3.5 && fixGrowth <= 4.5);
        }

        // ------ Test 4: mixed producers (some inside, some outside DAG) ------
        {
            // Task t2 references task t0 (inside), task "external_task" (outside)
            List<String> t0inputs = Arrays.asList((String)null);
            List<String> t1inputs = Arrays.asList("task0");
            List<String> t2inputs = Arrays.asList("task0", "task1", "external_task");
            List<Task> tasks = Arrays.asList(
                new Task("task0", t0inputs),
                new Task("task1", t1inputs),
                new Task("task2", t2inputs)
            );

            Result def = runDefective(tasks);
            Result fix = runFixed(tasks);

            check("mixed: same specs", specsEqual(def.taskSpecs, fix.taskSpecs));

            // task2's "external_task" ref should be in componentInputs (not in DAG)
            TaskSpec t2defSpec = def.taskSpecs.get(2);
            check("mixed: external_task classified as component_input (defective)",
                t2defSpec.componentInputs.containsKey("task2_from_external_task"));
            check("mixed: external_task classified as component_input (fixed)",
                fix.taskSpecs.get(2).componentInputs.containsKey("task2_from_external_task"));
        }

        System.out.println("All tests PASS.");
    }
}