package unit;
import java.util.*;

/**
 * MpichGroupTest — CWE-407 benchmark for mpich-0001
 *
 * mpich-0001: pmap_lpid_to_rank O(N) linear scan in group set operations
 *   SLOW: O(N×M) — pmap_lpid_to_rank linear scan inside outer loop over group1
 *   FAST: O(N+M) — HashMap reverse lookup built once from group2, O(1) per query
 *
 * Affected functions (all call pmap_lpid_to_rank inside a loop):
 *   - MPIR_Group_difference_impl      group_impl.c:330
 *   - MPIR_Group_intersection_impl    group_impl.c:368
 *   - MPIR_Group_union_impl           group_impl.c:414
 *   - MPIR_Group_translate_ranks_impl group_impl.c:93
 *   - MPIR_Group_compare_impl         group_impl.c:55
 *
 * The defect is acknowledged by a code comment in grouputil.c:475:
 *   "Use linear search for now.
 *    Optimization: build hash map in MPIR_Group_create_map and do O(1) hash lookup"
 * and by a TODO in grouputil.c:226:
 *   "TODO: build hash to accelerate MPIR_Group_lpid_to_rank"
 */
public class MpichGroupTest {

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

    static void bench(String label, Runnable slow, Runnable fast, long sOps, long fOps) {
        // warmup
        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 group is modeled as an array of lpids (logical process IDs).
    // map-format groups arise when ranks are not a simple arithmetic progression
    // (MPI_Group_incl with arbitrary rank lists, MPI_Comm_split with gaps, etc.)
    //
    // MPIR_Group_lpid_to_rank scans the map[] array linearly:
    //   for (rank = 0; rank < size; rank++)
    //       if (map[rank] == lpid) return rank;
    //   return MPI_UNDEFINED;
    //
    // This is called from MPIR_Group_difference_impl inside:
    //   for (i = 0; i < group1->size; i++)
    //       if (MPI_UNDEFINED == MPIR_Group_lpid_to_rank(group2, lpid)) ...
    //
    // Total: O(N × M) comparisons.
    // =========================================================================

    /** Simulate the current slow MPIR_Group_lpid_to_rank: linear scan */
    static int lpid_to_rank_slow(long[] map, long lpid) {
        for (int rank = 0; rank < map.length; rank++) {
            if (map[rank] == lpid) return rank;
        }
        return -1; // MPI_UNDEFINED
    }

    /** Simulate the fast version: O(1) HashMap lookup */
    static int lpid_to_rank_fast(Map<Long, Integer> reverseMap, long lpid) {
        return reverseMap.getOrDefault(lpid, -1);
    }

    // -------------------------------------------------------------------------
    // SLOW: MPIR_Group_difference_impl current behavior — O(N × M)
    //   For each lpid in group1, scan all of group2 linearly.
    // -------------------------------------------------------------------------

    static long groupDifferenceSlow(long[] group1, long[] group2) {
        long ops = 0;
        long[] result = new long[group1.length];
        int nnew = 0;
        for (long lpid : group1) {
            // pmap_lpid_to_rank: linear scan
            boolean found = false;
            for (int r = 0; r < group2.length; r++) {
                ops++;
                if (group2[r] == lpid) { found = true; break; }
            }
            if (!found) result[nnew++] = lpid;
        }
        return ops;
    }

    // -------------------------------------------------------------------------
    // FAST: MPIR_Group_difference_impl with hash — O(N + M)
    //   Build reverse hash of group2 once, then O(1) per group1 element.
    // -------------------------------------------------------------------------

    static long groupDifferenceFast(long[] group1, long[] group2) {
        long ops = 0;
        // Build hash: O(M)
        Set<Long> group2Set = new HashSet<>(group2.length * 2);
        for (long lpid : group2) {
            group2Set.add(lpid);
            ops++;
        }
        // O(N) lookups
        long[] result = new long[group1.length];
        int nnew = 0;
        for (long lpid : group1) {
            ops++;
            if (!group2Set.contains(lpid)) result[nnew++] = lpid;
        }
        return ops;
    }

    // -------------------------------------------------------------------------
    // SLOW: MPIR_Group_translate_ranks_impl — O(n_ranks × M)
    //   For each of n_ranks ranks in group1, scan all of group2 for the lpid.
    // -------------------------------------------------------------------------

    static long translateRanksSlow(long[] group1, int[] ranksToTranslate, long[] group2) {
        long ops = 0;
        int[] result = new int[ranksToTranslate.length];
        for (int i = 0; i < ranksToTranslate.length; i++) {
            long lpid = group1[ranksToTranslate[i]];
            // linear scan of group2
            result[i] = -1;
            for (int r = 0; r < group2.length; r++) {
                ops++;
                if (group2[r] == lpid) { result[i] = r; break; }
            }
        }
        return ops;
    }

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

    static long translateRanksFast(long[] group1, int[] ranksToTranslate, long[] group2) {
        long ops = 0;
        // Build reverse map of group2: lpid -> rank
        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[ranksToTranslate.length];
        for (int i = 0; i < ranksToTranslate.length; i++) {
            long lpid = group1[ranksToTranslate[i]];
            ops++;
            result[i] = revMap.getOrDefault(lpid, -1);
        }
        return ops;
    }

    // =========================================================================
    // Test cases
    // =========================================================================

    static boolean PASS = true;

    static void assertCorrect(String test, long[] group1, long[] group2) {
        // Verify slow and fast produce the same result set for group difference
        Set<Long> slowResult = new HashSet<>();
        Set<Long> g2Set = new HashSet<>();
        for (long x : group2) g2Set.add(x);
        for (long lpid : group1) {
            boolean found = false;
            for (long x : group2) if (x == lpid) { found = true; break; }
            if (!found) slowResult.add(lpid);
        }
        Set<Long> fastResult = new HashSet<>();
        for (long lpid : group1) {
            if (!g2Set.contains(lpid)) fastResult.add(lpid);
        }
        boolean ok = slowResult.equals(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("=== mpich-0001: pmap_lpid_to_rank O(N) linear scan in group ops ===");
        System.out.println();

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

        // Basic difference: group1 = {0..9}, group2 = {5..14}, diff = {0..4}
        long[] g1 = new long[10], g2 = new long[10];
        for (int i = 0; i < 10; i++) { g1[i] = i; g2[i] = i + 5; }
        assertCorrect("difference {0..9} \\ {5..14} = {0..4}", g1, g2);

        // Intersection via complement: group1 = {0,2,4,6,8}, group2 = {1,2,3,4,5}
        long[] ga = {0, 2, 4, 6, 8}, gb = {1, 2, 3, 4, 5};
        assertCorrect("difference {even 0-8} \\ {1-5}", ga, gb);

        // Disjoint groups
        long[] gc = {0, 1, 2}, gd = {10, 11, 12};
        assertCorrect("disjoint groups: diff = group1", gc, gd);

        // Identical groups
        long[] ge = {5, 3, 7, 1}, gf = {5, 3, 7, 1};
        assertCorrect("identical groups: diff = empty", ge, gf);

        System.out.println();

        // --- Performance benchmarks ---
        System.out.println("--- Performance (group_difference, group_translate_ranks) ---");

        int[] sizes = {100, 500, 1000};
        for (int N : sizes) {
            // Create two groups with partial overlap: group1 = 0..N-1, group2 = N/2..3N/2-1
            final long[] grp1 = new long[N], grp2 = new long[N];
            for (int i = 0; i < N; i++) { grp1[i] = i; grp2[i] = i + N / 2; }

            final long[] slowOps = {0}, fastOps = {0};
            bench(
                String.format("group_difference  N=%d", N),
                () -> slowOps[0] = groupDifferenceSlow(grp1, grp2),
                () -> fastOps[0] = groupDifferenceFast(grp1, grp2),
                slowOps[0], fastOps[0]
            );

            // translate_ranks: translate all N ranks
            int[] allRanks = new int[N];
            for (int i = 0; i < N; i++) allRanks[i] = i;
            bench(
                String.format("translate_ranks   N=%d", N),
                () -> slowOps[0] = translateRanksSlow(grp1, allRanks, grp2),
                () -> fastOps[0] = translateRanksFast(grp1, allRanks, grp2),
                slowOps[0], fastOps[0]
            );
        }

        // Large scale
        {
            final int N = 4096;
            final long[] grp1 = new long[N], grp2 = new long[N];
            for (int i = 0; i < N; i++) {
                // Non-strided (map format) with partial overlap
                grp1[i] = (long) i * 3;
                grp2[i] = (long) i * 3 + 1;
            }
            final long[] slowOps = {0}, fastOps = {0};
            bench(
                String.format("group_difference  N=%d (non-strided)", N),
                () -> slowOps[0] = groupDifferenceSlow(grp1, grp2),
                () -> fastOps[0] = groupDifferenceFast(grp1, grp2),
                slowOps[0], fastOps[0]
            );
        }

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