package unit;

import java.util.*;

/**
 * pytorch-geometric-0001: smiles.py list.index() O(L×A) → O(A) with dict
 *
 * Simulates the defective pattern from torch_geometric/utils/smiles.py:
 *   from_rdmol() calls x_map[key].index(value) for each atom feature,
 *   where x_map[key] is a List. This is O(L) per call, called A×9 times
 *   per molecule, M molecules total → O(M × A × L).
 *
 * Fixed: pre-build x_idx[key] = {value → index} HashMap → O(1) lookup.
 *
 * Run: javac -d . PyGSmilesIndexAlgorithm.java && java unit.PyGSmilesIndexAlgorithm
 */
public class PyGSmilesIndexAlgorithm {

    // ---- Result ----
    static class Result {
        final int slowOps;
        final int fastOps;
        Result(int slowOps, int fastOps) {
            this.slowOps = slowOps;
            this.fastOps = fastOps;
        }
    }

    // ---- Defective: list.index() O(L) per lookup ----
    static class DefectiveMolProcessor {
        // Simulates x_map['atomic_num'] = list(range(0, 119))
        final List<Integer> atomicNumList;
        // Simulates x_map['chirality'] = 9 entries
        final List<String> chiralityList;
        // Simulates e_map['bond_type'] = 22 entries
        final List<String> bondTypeList;

        int opCount = 0;

        DefectiveMolProcessor() {
            atomicNumList = new ArrayList<>();
            for (int i = 0; i < 119; i++) atomicNumList.add(i);

            chiralityList = new ArrayList<>();
            String[] chiralities = {"CHI_UNSPECIFIED","CHI_TETRAHEDRAL_CW","CHI_TETRAHEDRAL_CCW",
                "CHI_OTHER","CHI_TETRAHEDRAL","CHI_ALLENE","CHI_SQUAREPLANAR",
                "CHI_TRIGONALBIPYRAMIDAL","CHI_OCTAHEDRAL"};
            chiralityList.addAll(Arrays.asList(chiralities));

            bondTypeList = new ArrayList<>();
            for (int i = 0; i < 22; i++) bondTypeList.add("BOND_TYPE_" + i);
        }

        // O(L) list scan — mirrors list.index()
        int listIndex(List<?> list, Object value) {
            for (int i = 0; i < list.size(); i++) {
                opCount++;
                if (list.get(i).equals(value)) return i;
            }
            throw new NoSuchElementException("Value not found: " + value);
        }

        // Process one molecule: A atoms, B bonds
        // Each atom: 3 list.index() calls (atomic_num, chirality, dummy)
        // Each bond: 1 list.index() call (bond_type)
        void processMolecule(int[] atomNums, String[] chiralTags, String[] bondTypes) {
            for (int i = 0; i < atomNums.length; i++) {
                listIndex(atomicNumList, atomNums[i]);
                listIndex(chiralityList, chiralTags[i % chiralTags.length]);
                // Simulate remaining 7 atom feature lookups (simplified):
                // degree, formal_charge, num_hs, num_radical, hybridization,
                // is_aromatic, is_in_ring — all small lists, count as 7 ops each
                opCount += 7;
            }
            for (String bt : bondTypes) {
                listIndex(bondTypeList, bt);
            }
        }

        Result run(int[][] atomNumsPerMol, String[] chiralTags, String[] bondTypes) {
            opCount = 0;
            for (int[] atomNums : atomNumsPerMol) {
                processMolecule(atomNums, chiralTags, bondTypes);
            }
            return new Result(opCount, -1);
        }
    }

    // ---- Fixed: HashMap O(1) lookup ----
    static class FixedMolProcessor {
        final Map<Integer, Integer> atomicNumIdx;
        final Map<String, Integer> chiralityIdx;
        final Map<String, Integer> bondTypeIdx;

        int opCount = 0;

        FixedMolProcessor() {
            // Build index maps once at module load
            atomicNumIdx = new HashMap<>();
            for (int i = 0; i < 119; i++) atomicNumIdx.put(i, i);

            chiralityIdx = new HashMap<>();
            String[] chiralities = {"CHI_UNSPECIFIED","CHI_TETRAHEDRAL_CW","CHI_TETRAHEDRAL_CCW",
                "CHI_OTHER","CHI_TETRAHEDRAL","CHI_ALLENE","CHI_SQUAREPLANAR",
                "CHI_TRIGONALBIPYRAMIDAL","CHI_OCTAHEDRAL"};
            for (int i = 0; i < chiralities.length; i++) chiralityIdx.put(chiralities[i], i);

            bondTypeIdx = new HashMap<>();
            for (int i = 0; i < 22; i++) bondTypeIdx.put("BOND_TYPE_" + i, i);
        }

        int mapGet(Map<?, Integer> map, Object key) {
            opCount++;  // one op per lookup (hash + compare)
            return map.get(key);
        }

        void processMolecule(int[] atomNums, String[] chiralTags, String[] bondTypes) {
            for (int i = 0; i < atomNums.length; i++) {
                mapGet(atomicNumIdx, atomNums[i]);
                mapGet(chiralityIdx, chiralTags[i % chiralTags.length]);
                opCount += 7;  // 7 other O(1) dict lookups
            }
            for (String bt : bondTypes) {
                mapGet(bondTypeIdx, bt);
            }
        }

        Result run(int[][] atomNumsPerMol, String[] chiralTags, String[] bondTypes) {
            opCount = 0;
            for (int[] atomNums : atomNumsPerMol) {
                processMolecule(atomNums, chiralTags, bondTypes);
            }
            return new Result(-1, opCount);
        }
    }

    // ---- check helper ----
    static int passed = 0;
    static int total = 0;

    static void check(String name, boolean condition) {
        total++;
        if (condition) {
            passed++;
            System.out.println("  PASS: " + name);
        } else {
            System.out.println("  FAIL: " + name);
        }
    }

    public static void main(String[] args) {
        System.out.println("pytorch-geometric-0001: smiles.py list.index() O(L*A) -> O(A)");
        System.out.println();

        // Build test data: simulate molecules with varying atom counts
        // Atoms use atomic nums spread across the list to trigger long scans
        String[] chiralTags = {"CHI_UNSPECIFIED", "CHI_TETRAHEDRAL_CW", "CHI_OTHER"};
        String[] bondTypes = {"BOND_TYPE_0", "BOND_TYPE_10", "BOND_TYPE_21"};

        // ---- Test 1: Small molecule (10 atoms, 9 bonds) ----
        {
            int M = 1;
            int[][] atomNums = new int[M][];
            // Use high atomic numbers to stress the list scan
            atomNums[0] = new int[]{100, 110, 90, 50, 118, 60, 80, 30, 5, 1};

            DefectiveMolProcessor slow = new DefectiveMolProcessor();
            Result slowR = slow.run(atomNums, chiralTags, bondTypes);

            FixedMolProcessor fast = new FixedMolProcessor();
            Result fastR = fast.run(atomNums, chiralTags, bondTypes);

            System.out.println("Test 1: 1 molecule, 10 atoms");
            System.out.println("  Slow ops: " + slowR.slowOps);
            System.out.println("  Fast ops: " + fastR.fastOps);
            check("slow > fast (list scan > dict lookup)", slowR.slowOps > fastR.fastOps);
            check("slow ops > 300 (O(L) scans dominate)", slowR.slowOps > 300);
        }

        // ---- Test 2: Large dataset (1000 molecules, 18 atoms each) ----
        {
            int M = 1000;
            int A = 18;
            int[][] atomNums = new int[M][A];
            Random rng = new Random(42);
            for (int m = 0; m < M; m++) {
                for (int a = 0; a < A; a++) {
                    atomNums[m][a] = rng.nextInt(119);
                }
            }

            DefectiveMolProcessor slow = new DefectiveMolProcessor();
            Result slowR = slow.run(atomNums, chiralTags, bondTypes);

            FixedMolProcessor fast = new FixedMolProcessor();
            Result fastR = fast.run(atomNums, chiralTags, bondTypes);

            System.out.println("Test 2: 1000 molecules, 18 atoms each");
            System.out.println("  Slow ops: " + slowR.slowOps);
            System.out.println("  Fast ops: " + fastR.fastOps);
            double ratio = (double) slowR.slowOps / fastR.fastOps;
            System.out.printf("  Ratio slow/fast: %.1fx%n", ratio);
            check("slow > 5x fast ops (O(L) overhead confirmed)", ratio > 5.0);
            check("slow ops scale with L=119 (>500k for M=1000,A=18)", slowR.slowOps > 500_000);
        }

        // ---- Test 3: Correctness — same results ----
        {
            int[][] singleMol = {{6, 7, 8, 1, 16}};  // C, N, O, H, S

            DefectiveMolProcessor slow = new DefectiveMolProcessor();
            FixedMolProcessor fast = new FixedMolProcessor();

            // Verify same index returned for known values
            int slowIdx = slow.listIndex(slow.atomicNumList, 6);   // Carbon = index 6
            int fastIdx = fast.atomicNumIdx.get(6);                // Carbon = index 6
            System.out.println("Test 3: correctness check for atomic_num=6 (Carbon)");
            System.out.println("  Slow index: " + slowIdx + ", Fast index: " + fastIdx);
            check("slow and fast return same index (6)", slowIdx == fastIdx && fastIdx == 6);

            int slowChiral = slow.listIndex(slow.chiralityList, "CHI_TETRAHEDRAL_CW");
            int fastChiral = fast.chiralityIdx.get("CHI_TETRAHEDRAL_CW");
            System.out.println("Test 4: correctness check for chirality=CHI_TETRAHEDRAL_CW");
            check("slow and fast return same chirality index (1)", slowChiral == fastChiral && fastChiral == 1);
        }

        // ---- Test 4: O(n^2) scaling — double molecule count doubles slow more ----
        {
            int A = 10;
            String[] chiralTagsT = {"CHI_UNSPECIFIED"};
            String[] bondTypesT = {"BOND_TYPE_21"};  // index 21 = worst case scan

            int M1 = 100;
            int M2 = 200;

            int[][] atomNums1 = new int[M1][A];
            int[][] atomNums2 = new int[M2][A];
            for (int m = 0; m < M2; m++) {
                for (int a = 0; a < A; a++) {
                    int v = 100 + (m % 19);  // high atomic nums to stress scan
                    if (m < M1) atomNums1[m][a] = v;
                    atomNums2[m][a] = v;
                }
            }

            DefectiveMolProcessor slow1 = new DefectiveMolProcessor();
            DefectiveMolProcessor slow2 = new DefectiveMolProcessor();
            slow1.run(atomNums1, chiralTagsT, bondTypesT);
            slow2.run(atomNums2, chiralTagsT, bondTypesT);

            System.out.println("Test 5: O(n) scaling — 2x molecules → ~2x ops");
            System.out.println("  M=" + M1 + " ops: " + slow1.opCount);
            System.out.println("  M=" + M2 + " ops: " + slow2.opCount);
            double scalingRatio = (double) slow2.opCount / slow1.opCount;
            System.out.printf("  Scaling: %.2fx (should be ~2.0)%n", scalingRatio);
            check("slow ops scale ~2x when molecules doubled", scalingRatio > 1.9 && scalingRatio < 2.1);
        }

        System.out.println();
        System.out.println(passed + "/" + total + " PASS");
    }
}