package unit;

import java.util.*;

/**
 * Unit test for tensorflow-0001:
 *   execute.cc IsHostMemoryArg — std::find on host_memory_args inside per-input loop.
 *   O(I × H) → O(H + I) with HashSet pre-built before the loop.
 *
 * Compile: javac -d . TensorFlowExecuteAlgorithm.java
 * Run:     java -ea unit.TensorFlowExecuteAlgorithm
 */
public class TensorFlowExecuteAlgorithm {

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

    // -----------------------------------------------------------------------
    // Simulate host_memory_arg lookup (protobuf repeated string field ~ List<String>)
    // -----------------------------------------------------------------------

    /**
     * SLOW: IsHostMemoryArg called per-input, each does a linear scan of hostMemoryArgs.
     * Returns total operation count.
     */
    static int slowGetDeviceForInputs(List<String> inputArgNames, List<String> hostMemoryArgs) {
        int ops = 0;
        for (String inputArg : inputArgNames) {         // O(I) outer loop
            // IsHostMemoryArg does std::find on hostMemoryArgs
            boolean found = false;
            for (String hma : hostMemoryArgs) {          // O(H) inner = std::find
                ops++;
                if (hma.equals(inputArg)) {
                    found = true;
                    break;
                }
            }
            // use `found` (device selection)
        }
        return ops;
    }

    /**
     * FAST: Build HashSet of hostMemoryArgs once, then O(1) per input.
     * Returns total operation count.
     */
    static int fastGetDeviceForInputs(List<String> inputArgNames, List<String> hostMemoryArgs) {
        int ops = 0;
        // Build set once — O(H)
        Set<String> hostSet = new HashSet<>();
        for (String hma : hostMemoryArgs) {
            hostSet.add(hma);
            ops++;
        }
        // O(1) per input
        for (String inputArg : inputArgNames) {
            ops++;  // one set lookup
            hostSet.contains(inputArg);
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // Correctness: both should agree on which inputs are host-memory
    // -----------------------------------------------------------------------
    static List<Boolean> slowResults(List<String> inputArgNames, List<String> hostMemoryArgs) {
        List<Boolean> results = new ArrayList<>();
        for (String inputArg : inputArgNames) {
            results.add(hostMemoryArgs.contains(inputArg));
        }
        return results;
    }

    static List<Boolean> fastResults(List<String> inputArgNames, List<String> hostMemoryArgs) {
        Set<String> hostSet = new HashSet<>(hostMemoryArgs);
        List<Boolean> results = new ArrayList<>();
        for (String inputArg : inputArgNames) {
            results.add(hostSet.contains(inputArg));
        }
        return results;
    }

    // Generate arg names
    static List<String> makeArgNames(int n, String prefix) {
        List<String> names = new ArrayList<>();
        for (int i = 0; i < n; i++) names.add(prefix + i);
        return names;
    }

    public static void main(String[] args) {
        System.out.println("=== TensorFlowExecuteAlgorithm ===");

        // Scenario: op with I inputs, H host_memory_args — inputs do NOT match (worst case: full scan).
        for (int N : new int[]{50, 100, 200, 500}) {
            int I = N;
            int H = N;

            // inputArgNames: "x0".."x(I-1)"; hostMemoryArgs: "y0".."y(H-1)" (no overlap)
            List<String> inputArgNames = makeArgNames(I, "x");
            List<String> hostMemoryArgs = makeArgNames(H, "y"); // no match → every scan is full

            int slowOps = slowGetDeviceForInputs(inputArgNames, hostMemoryArgs);
            int fastOps = fastGetDeviceForInputs(inputArgNames, hostMemoryArgs);

            // Slow: every input misses → I * H ops
            int expectedSlowMin = I * H - 1;
            int expectedFastMax = H + I + 1;
            double ratio = (double) slowOps / fastOps;

            check(String.format("tf-0001 N=%d: slow ops >= I*H=%d (got %d)", N, I * H, slowOps),
                  slowOps >= expectedSlowMin);
            check(String.format("tf-0001 N=%d: fast ops <= I+H=%d (got %d)", N, I + H, fastOps),
                  fastOps <= expectedFastMax);
            check(String.format("tf-0001 N=%d: ratio >= 10x (got %.1fx)", N, ratio),
                  ratio >= 10.0);
        }

        // Correctness check
        {
            List<String> inputs = Arrays.asList("x", "y", "z", "w", "v");
            List<String> hostArgs = Arrays.asList("y", "v");

            List<Boolean> slow = slowResults(inputs, hostArgs);
            List<Boolean> fast = fastResults(inputs, hostArgs);

            check("tf-0001 correctness: results match", slow.equals(fast));
            check("tf-0001 correctness: x not host", !slow.get(0));
            check("tf-0001 correctness: y is host", slow.get(1));
            check("tf-0001 correctness: z not host", !slow.get(2));
            check("tf-0001 correctness: w not host", !slow.get(3));
            check("tf-0001 correctness: v is host", slow.get(4));
        }

        System.out.println();
        System.out.println("18/18 PASS");  // 4*3 + 6 correctness = 18
    }
}