package unit;

import java.util.*;

/**
 * Unit test for pytorch-0001 and pytorch-0002:
 *   graph_fuser.cpp fuseChunkByReusingExistingFusedChunk and tryToMoveChunk
 *   O(C×I) std::find on group inputs inside chunk output loop → O(C+I) with HashMap.
 *
 * Compile: javac -d . PyTorchGraphFuserAlgorithm.java
 * Run:     java -ea unit.PyTorchGraphFuserAlgorithm
 */
public class PyTorchGraphFuserAlgorithm {

    static int opCount = 0;

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

    // -----------------------------------------------------------------------
    // Simulated Value pointer — identity by object reference
    // -----------------------------------------------------------------------
    static class Value {
        final int id;
        Value(int id) { this.id = id; }
        @Override public String toString() { return "v" + id; }
    }

    // -----------------------------------------------------------------------
    // SLOW: fuseChunkByReusingExistingFusedChunk — std::find equivalent
    // -----------------------------------------------------------------------
    static int slowFuseChunk(List<Value> groupInputs, List<Value> chunkOutputs) {
        int ops = 0;
        for (Value val : chunkOutputs) {                     // O(C) outer
            for (Value gi : groupInputs) {                   // O(I) inner = std::find
                ops++;
                if (gi == val) break;
            }
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // FAST: fuseChunkByReusingExistingFusedChunk — HashMap lookup
    // -----------------------------------------------------------------------
    static int fastFuseChunk(List<Value> groupInputs, List<Value> chunkOutputs) {
        int ops = 0;
        // Build reverse index once: O(I)
        Map<Value, Integer> indexMap = new HashMap<>();
        for (int k = 0; k < groupInputs.size(); k++) {
            indexMap.put(groupInputs.get(k), k);
            ops++;
        }
        // Lookup per chunk output: O(1) each
        for (Value val : chunkOutputs) {
            ops++;  // one lookup
            indexMap.get(val); // O(1)
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // SLOW: tryToMoveChunk — std::find on bchunk_inputs per producer input
    // -----------------------------------------------------------------------
    static int slowTryMoveChunk(List<Value> producerInputs, List<Value> bchunkInputs) {
        int ops = 0;
        for (Value input : producerInputs) {                 // O(I) outer
            for (Value bi : bchunkInputs) {                  // O(B) inner = std::find
                ops++;
                if (bi == input) break;
            }
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // FAST: tryToMoveChunk — HashMap on bchunk_inputs
    // -----------------------------------------------------------------------
    static int fastTryMoveChunk(List<Value> producerInputs, List<Value> bchunkInputs) {
        int ops = 0;
        // Build index once: O(B)
        Map<Value, Integer> indexMap = new HashMap<>();
        for (int k = 0; k < bchunkInputs.size(); k++) {
            indexMap.put(bchunkInputs.get(k), k);
            ops++;
        }
        // Lookup per producer input: O(1) each
        for (Value input : producerInputs) {
            ops++;
            indexMap.get(input); // O(1)
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // Build test data — some chunk outputs are in groupInputs, some are not
    // -----------------------------------------------------------------------
    static List<Value> makeValues(int n) {
        List<Value> vs = new ArrayList<>();
        for (int i = 0; i < n; i++) vs.add(new Value(i));
        return vs;
    }

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

        // --- pytorch-0001: fuseChunkByReusingExistingFusedChunk ---
        // Worst case: chunk outputs are NOT in groupInputs → slow scans full I list each time.
        for (int N : new int[]{50, 100, 200, 500}) {
            int C = N;          // chunk outputs
            int I = N;          // group inputs (same size)

            // groupInputs: values 0..I-1
            List<Value> groupInputs = makeValues(I);
            // chunkOutputs: fresh values NOT in groupInputs (ids I..I+C-1)
            List<Value> chunkOutputs = new ArrayList<>();
            List<Value> allVals = makeValues(I + C);
            for (int i = I; i < I + C; i++) chunkOutputs.add(allVals.get(i));

            int slowOps = slowFuseChunk(groupInputs, chunkOutputs);
            int fastOps = fastFuseChunk(groupInputs, chunkOutputs);

            // Slow: every chunk output misses, scans all I entries → C * I
            int expectedSlowMin = C * I - 1;
            // Fast: I + C (build index + lookups)
            int expectedFastMax = I + C + 1;

            double ratio = (double) slowOps / fastOps;

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

        // --- pytorch-0002: tryToMoveChunk ---
        // Worst case: producer inputs are NOT in bchunkInputs → each scan misses (full B scan).
        for (int N : new int[]{50, 100, 200, 500}) {
            int I = N;      // producer inputs
            int B = N;      // bchunk inputs

            // producerInputs: ids 0..I-1; bchunkInputs: ids I..I+B-1 (no overlap → full miss)
            List<Value> allVals = makeValues(I + B);
            List<Value> producerInputs = new ArrayList<>(allVals.subList(0, I));
            List<Value> bchunkInputs   = new ArrayList<>(allVals.subList(I, I + B));

            int slowOps = slowTryMoveChunk(producerInputs, bchunkInputs);
            int fastOps = fastTryMoveChunk(producerInputs, bchunkInputs);

            // All I producer inputs miss fully → I * B ops
            int expectedSlowMin = I * B - 1;
            int expectedFastMax = I + B + 1;
            double ratio = (double) slowOps / fastOps;

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

        // --- pytorch-0001 correctness check ---
        {
            List<Value> vals = makeValues(10);
            List<Value> groupInputs = new ArrayList<>(vals);
            // chunk outputs = vals[2], vals[5], vals[8]
            List<Value> chunkOutputs = Arrays.asList(vals.get(2), vals.get(5), vals.get(8));

            // Slow and fast should agree on which indices are found
            Map<Value, Integer> fastMap = new HashMap<>();
            for (int k = 0; k < groupInputs.size(); k++) fastMap.put(groupInputs.get(k), k);

            for (Value v : chunkOutputs) {
                int slowIdx = groupInputs.indexOf(v);
                int fastIdx = fastMap.getOrDefault(v, -1);
                check(String.format("pytorch-0001 correctness: val %s slow=%d fast=%d", v, slowIdx, fastIdx),
                      slowIdx == fastIdx);
            }
        }

        System.out.println();
        System.out.println("27/27 PASS");
    }
}