package unit;
import java.util.*;

/**
 * OmpiGroupTest — CWE-407 benchmark for ompi-0001
 *
 * ompi-0001: ompi_group O(N×M) nested process-name scan in group set operations
 *   SLOW: O(N×M) — nested for-loops comparing opal_process_name_t structs
 *   FAST: O(N+M) — HashMap keyed on (jobid<<32|vpid) built once from group2
 *
 * Affected functions in ompi/group/group.c:
 *   - ompi_group_translate_ranks (dense fallback at line 105)  O(n_ranks × M)
 *   - ompi_group_intersection (line 444)                       O(N × M)
 *   - ompi_group_overlap (line 629)                            O(N × M)
 *   - ompi_group_compare (line 491)                            O(N²)
 *
 * The sparse fast-path (#if OMPI_GROUP_SPARSE) only applies to parent-child
 * group relationships. Independent groups (the common case in MPI collective
 * communicator creation) always hit the O(N×M) path.
 */
public class OmpiGroupTest {

    // -------------------------------------------------------------------------
    // Benchmark harness
    // -------------------------------------------------------------------------

    static void bench(String label, Runnable slow, Runnable fast, long sOps, long fOps) {
        slow.run(); fast.run();
        long t0 = System.nanoTime(); slow.run(); long sMs = (System.nanoTime() - t0) / 1_000_000;
        long t1 = System.nanoTime(); fast.run(); long fMs = (System.nanoTime() - t1) / 1_000_000;
        System.out.printf("  %-60s  slow:%5dms (%,d ops)  fast:%5dms (%,d ops)  speedup:%.0fx%n",
                label, sMs, sOps, fMs, fOps, fOps > 0 ? (double) sOps / fOps : 0);
    }

    // =========================================================================
    // Model
    //
    // An MPI process is identified by (jobid, vpid) — an opal_process_name_t.
    // We model this as a long: (jobid << 32) | vpid.
    //
    // A group is an array of such process names.
    //
    // ompi_group_translate_ranks (dense path) does:
    //   for proc = 0 .. n_ranks-1:
    //       name1 = group1[ranks[proc]]
    //       for proc2 = 0 .. group2.size-1:
    //           name2 = group2[proc2]
    //           if opal_compare_proc(name1, name2) == 0:
    //               result[proc] = proc2; break
    //
    // ompi_group_intersection does:
    //   for proc1 = 0 .. group1.size-1:
    //       name1 = group1[proc1]
    //       for proc2 = 0 .. group2.size-1:
    //           name2 = group2[proc2]
    //           if opal_compare_proc(name1, name2) == 0:
    //               included[k++] = proc1; break
    //
    // ompi_group_overlap does:
    //   for i = 0 .. group1.size-1:
    //       for j = 0 .. group2.size-1:
    //           if same_proc: return true
    //   return false
    //
    // opal_compare_proc compares two 64-bit process names: one comparison each.
    // =========================================================================

    // --- SLOW: ompi_group_translate_ranks dense fallback O(n_ranks × M) ---

    static long translateRanksSlow(long[] group1, int[] ranks, long[] group2) {
        long ops = 0;
        int[] result = new int[ranks.length];
        for (int i = 0; i < ranks.length; i++) {
            long name1 = group1[ranks[i]];
            result[i] = -1;
            for (int proc2 = 0; proc2 < group2.length; proc2++) {
                ops++;
                if (group2[proc2] == name1) { result[i] = proc2; break; }
            }
        }
        return ops;
    }

    // --- FAST: O(n_ranks + M) with hash ---

    static long translateRanksFast(long[] group1, int[] ranks, long[] group2) {
        long ops = 0;
        // Build reverse map: process_name -> rank in group2
        Map<Long, Integer> revMap = new HashMap<>(group2.length * 2);
        for (int r = 0; r < group2.length; r++) {
            revMap.put(group2[r], r);
            ops++;
        }
        int[] result = new int[ranks.length];
        for (int i = 0; i < ranks.length; i++) {
            long name1 = group1[ranks[i]];
            ops++;
            result[i] = revMap.getOrDefault(name1, -1);
        }
        return ops;
    }

    // --- SLOW: ompi_group_intersection O(N × M) ---

    static long intersectionSlow(long[] group1, long[] group2) {
        long ops = 0;
        List<Integer> included = new ArrayList<>();
        for (int proc1 = 0; proc1 < group1.length; proc1++) {
            long name1 = group1[proc1];
            for (int proc2 = 0; proc2 < group2.length; proc2++) {
                ops++;
                if (group2[proc2] == name1) { included.add(proc1); break; }
            }
        }
        return ops;
    }

    // --- FAST: ompi_group_intersection O(N + M) with hash ---

    static long intersectionFast(long[] group1, long[] group2) {
        long ops = 0;
        Set<Long> g2Set = new HashSet<>(group2.length * 2);
        for (long name : group2) { g2Set.add(name); ops++; }
        List<Integer> included = new ArrayList<>();
        for (int proc1 = 0; proc1 < group1.length; proc1++) {
            ops++;
            if (g2Set.contains(group1[proc1])) included.add(proc1);
        }
        return ops;
    }

    // --- SLOW: ompi_group_overlap O(N × M) ---

    static long overlapSlow(long[] group1, long[] group2) {
        long ops = 0;
        for (long name1 : group1) {
            for (long name2 : group2) {
                ops++;
                if (name1 == name2) return ops;
            }
        }
        return ops;
    }

    // --- FAST: ompi_group_overlap O(N + M) ---

    static long overlapFast(long[] group1, long[] group2) {
        long ops = 0;
        // Hash the smaller group
        long[] small = group1.length <= group2.length ? group1 : group2;
        long[] large = small == group1 ? group2 : group1;
        Set<Long> smallSet = new HashSet<>(small.length * 2);
        for (long name : small) { smallSet.add(name); ops++; }
        for (long name : large) {
            ops++;
            if (smallSet.contains(name)) return ops;
        }
        return ops;
    }

    // =========================================================================
    // Correctness verification
    // =========================================================================

    static boolean PASS = true;

    static void assertIntersectionEqual(String test, long[] group1, long[] group2) {
        // Collect intersection ranks using both methods
        List<Integer> slowResult = new ArrayList<>();
        for (int p1 = 0; p1 < group1.length; p1++) {
            boolean found = false;
            for (long n2 : group2) if (group1[p1] == n2) { found = true; break; }
            if (found) slowResult.add(p1);
        }
        Set<Long> g2Set = new HashSet<>();
        for (long n : group2) g2Set.add(n);
        List<Integer> fastResult = new ArrayList<>();
        for (int p1 = 0; p1 < group1.length; p1++) {
            if (g2Set.contains(group1[p1])) fastResult.add(p1);
        }
        boolean ok = slowResult.equals(fastResult);
        System.out.printf("  CORRECTNESS %-40s %s%n", test, ok ? "PASS" : "FAIL");
        if (!ok) PASS = false;
    }

    static void assertOverlapEqual(String test, long[] group1, long[] group2, boolean expected) {
        boolean slowHas = overlapSlow(group1, group2) > 0 && hasOverlap(group1, group2);
        boolean fastHas = overlapFast(group1, group2) > 0 && hasOverlap(group1, group2);
        boolean ok = slowHas == fastHas && slowHas == expected;
        System.out.printf("  CORRECTNESS %-40s %s (expected=%b actual=%b)%n",
                test, ok ? "PASS" : "FAIL", expected, slowHas);
        if (!ok) PASS = false;
    }

    static boolean hasOverlap(long[] g1, long[] g2) {
        Set<Long> s = new HashSet<>();
        for (long x : g2) s.add(x);
        for (long x : g1) if (s.contains(x)) return true;
        return false;
    }

    static void assertTranslateEqual(String test, long[] group1, int[] ranks, long[] group2) {
        // slow
        int[] slowResult = new int[ranks.length];
        for (int i = 0; i < ranks.length; i++) {
            long name1 = group1[ranks[i]];
            slowResult[i] = -1;
            for (int r = 0; r < group2.length; r++) {
                if (group2[r] == name1) { slowResult[i] = r; break; }
            }
        }
        // fast
        Map<Long, Integer> revMap = new HashMap<>();
        for (int r = 0; r < group2.length; r++) revMap.put(group2[r], r);
        int[] fastResult = new int[ranks.length];
        for (int i = 0; i < ranks.length; i++) {
            fastResult[i] = revMap.getOrDefault(group1[ranks[i]], -1);
        }
        boolean ok = Arrays.equals(slowResult, fastResult);
        System.out.printf("  CORRECTNESS %-40s %s%n", test, ok ? "PASS" : "FAIL");
        if (!ok) PASS = false;
    }

    // =========================================================================
    // Main
    // =========================================================================

    public static void main(String[] args) {
        System.out.println("=== ompi-0001: ompi_group O(N×M) nested process-name scan ===");
        System.out.println();

        // Build process-name arrays: jobid=0, vpid=rank (single job, typical case)
        // Non-overlapping portion: g1 = procs 0..N-1, g2 = procs N/2..3N/2-1

        System.out.println("--- Correctness ---");

        // Intersection tests
        long[] ga = {10L, 20L, 30L, 40L, 50L};
        long[] gb = {20L, 40L, 60L, 80L};
        assertIntersectionEqual("intersection {10,20,30,40,50} ∩ {20,40,60,80}", ga, gb);

        long[] gc = {1L, 2L, 3L, 4L};
        long[] gd = {5L, 6L, 7L, 8L};
        assertIntersectionEqual("disjoint groups: intersection = empty", gc, gd);

        long[] ge = {1L, 2L, 3L};
        long[] gf = {1L, 2L, 3L};
        assertIntersectionEqual("identical groups: intersection = group1", ge, gf);

        // Overlap tests
        assertOverlapEqual("overlap: groups share proc 20", ga, gb, true);
        assertOverlapEqual("overlap: disjoint groups", gc, gd, false);

        // Translate ranks tests
        long[] grp1 = {100L, 200L, 300L, 400L, 500L};
        long[] grp2 = {200L, 400L, 600L, 800L};
        int[] ranks = {0, 1, 2, 3, 4};
        assertTranslateEqual("translate_ranks: partial overlap", grp1, ranks, grp2);

        System.out.println();
        System.out.println("--- Performance ---");

        int[] sizes = {100, 500, 1000};
        for (int N : sizes) {
            final long[] g1 = new long[N], g2 = new long[N];
            // jobid=1, vpid=i for group1; jobid=1, vpid=i+N/2 for group2 (partial overlap)
            for (int i = 0; i < N; i++) {
                g1[i] = (1L << 32) | i;
                g2[i] = (1L << 32) | (i + N / 2);
            }
            int[] allRanks = new int[N];
            for (int i = 0; i < N; i++) allRanks[i] = i;

            final long[] slowOps = {0}, fastOps = {0};

            bench(
                String.format("group_translate_ranks  N=%d", N),
                () -> slowOps[0] = translateRanksSlow(g1, allRanks, g2),
                () -> fastOps[0] = translateRanksFast(g1, allRanks, g2),
                slowOps[0], fastOps[0]
            );

            bench(
                String.format("group_intersection     N=%d", N),
                () -> slowOps[0] = intersectionSlow(g1, g2),
                () -> fastOps[0] = intersectionFast(g1, g2),
                slowOps[0], fastOps[0]
            );

            bench(
                String.format("group_overlap          N=%d (worst: no overlap)", N),
                () -> slowOps[0] = overlapSlow(g1, g2),
                () -> fastOps[0] = overlapFast(g1, g2),
                slowOps[0], fastOps[0]
            );
        }

        // Large scale: 4096 processes, two independent jobs (cross-job group ops
        // are common in MPI applications with spawned processes)
        {
            final int N = 4096;
            final long[] g1 = new long[N], g2 = new long[N];
            for (int i = 0; i < N; i++) {
                // jobid=1 for g1, jobid=2 for g2 — completely disjoint (worst case for overlap)
                g1[i] = (1L << 32) | i;
                g2[i] = (2L << 32) | i;
            }
            int[] allRanks = new int[N];
            for (int i = 0; i < N; i++) allRanks[i] = i;
            final long[] slowOps = {0}, fastOps = {0};
            bench(
                String.format("group_overlap disjoint N=%d (full scan worst case)", N),
                () -> slowOps[0] = overlapSlow(g1, g2),
                () -> fastOps[0] = overlapFast(g1, g2),
                slowOps[0], fastOps[0]
            );
            bench(
                String.format("group_intersection     N=%d (disjoint)", N),
                () -> slowOps[0] = intersectionSlow(g1, g2),
                () -> fastOps[0] = intersectionFast(g1, g2),
                slowOps[0], fastOps[0]
            );
        }

        System.out.println();
        System.out.println("Result: " + (PASS ? "PASS" : "FAIL"));
        if (!PASS) System.exit(1);
    }
}