"""
GEMM (General Matrix Multiply) Memory Access Pattern Simulator

This module simulates different loop orderings for matrix multiplication
and tracks memory access patterns for visualization and cache analysis.
"""

import numpy as np
from typing import List, Tuple, Dict


class GEMMSimulator:
    """Simulates GEMM memory access patterns for different loop orderings."""

    # All possible loop orderings
    LOOP_ORDERS = ['ijk', 'ikj', 'jik', 'jki', 'kij', 'kji']

    def __init__(self, n: int, block_size: int = None):
        """
        Initialize GEMM simulator.

        Args:
            n: Matrix size (n x n)
            block_size: Block size for tiling (None for unblocked)
        """
        self.n = n
        self.block_size = block_size if block_size else n  # Unblocked if None
        self.tracks = []
        self.access_count = {'A': 0, 'B': 0, 'C': 0}

    def reset(self):
        """Reset simulation state."""
        self.tracks = []
        self.access_count = {'A': 0, 'B': 0, 'C': 0}

    def simulate(self, loop_order: str, blocked: bool = True) -> List[Tuple]:
        """
        Simulate GEMM with specified loop order.

        Args:
            loop_order: One of 'ijk', 'ikj', 'jik', 'jki', 'kij', 'kji'
            blocked: Whether to use blocking/tiling

        Returns:
            List of access patterns: [((i,k), (k,j), (i,j)), ...]
                                     (A_pos, B_pos, C_pos)
        """
        self.reset()

        if loop_order not in self.LOOP_ORDERS:
            raise ValueError(f"Invalid loop order. Must be one of {self.LOOP_ORDERS}")

        if blocked:
            self._simulate_blocked(loop_order)
        else:
            self._simulate_unblocked(loop_order)

        return self.tracks

    def _simulate_blocked(self, loop_order: str):
        """Simulate blocked/tiled GEMM."""
        n = self.n
        bs = self.block_size

        # Outer loops over blocks
        if loop_order == 'ijk':
            for i in range(0, n, bs):
                for j in range(0, n, bs):
                    for k in range(0, n, bs):
                        self._inner_loops(i, j, k, bs)

        elif loop_order == 'ikj':
            for i in range(0, n, bs):
                for k in range(0, n, bs):
                    for j in range(0, n, bs):
                        self._inner_loops(i, j, k, bs)

        elif loop_order == 'jik':
            for j in range(0, n, bs):
                for i in range(0, n, bs):
                    for k in range(0, n, bs):
                        self._inner_loops(i, j, k, bs)

        elif loop_order == 'jki':
            for j in range(0, n, bs):
                for k in range(0, n, bs):
                    for i in range(0, n, bs):
                        self._inner_loops(i, j, k, bs)

        elif loop_order == 'kij':
            for k in range(0, n, bs):
                for i in range(0, n, bs):
                    for j in range(0, n, bs):
                        self._inner_loops(i, j, k, bs)

        elif loop_order == 'kji':
            for k in range(0, n, bs):
                for j in range(0, n, bs):
                    for i in range(0, n, bs):
                        self._inner_loops(i, j, k, bs)

    def _inner_loops(self, i_base, j_base, k_base, bs):
        """Inner loops within a block."""
        n = self.n

        for ii in range(i_base, min(i_base + bs, n)):
            for jj in range(j_base, min(j_base + bs, n)):
                for kk in range(k_base, min(k_base + bs, n)):
                    # C[i,j] += A[i,k] * B[k,j]
                    a_pos = (ii, kk)
                    b_pos = (kk, jj)
                    c_pos = (ii, jj)
                    self.tracks.append((a_pos, b_pos, c_pos))
                    self.access_count['A'] += 1
                    self.access_count['B'] += 1
                    self.access_count['C'] += 1

    def _simulate_unblocked(self, loop_order: str):
        """Simulate unblocked GEMM."""
        n = self.n

        if loop_order == 'ijk':
            for i in range(n):
                for j in range(n):
                    for k in range(n):
                        self._record_access(i, j, k)

        elif loop_order == 'ikj':
            for i in range(n):
                for k in range(n):
                    for j in range(n):
                        self._record_access(i, j, k)

        elif loop_order == 'jik':
            for j in range(n):
                for i in range(n):
                    for k in range(n):
                        self._record_access(i, j, k)

        elif loop_order == 'jki':
            for j in range(n):
                for k in range(n):
                    for i in range(n):
                        self._record_access(i, j, k)

        elif loop_order == 'kij':
            for k in range(n):
                for i in range(n):
                    for j in range(n):
                        self._record_access(i, j, k)

        elif loop_order == 'kji':
            for k in range(n):
                for j in range(n):
                    for i in range(n):
                        self._record_access(i, j, k)

    def _record_access(self, i, j, k):
        """Record a single memory access."""
        a_pos = (i, k)
        b_pos = (k, j)
        c_pos = (i, j)
        self.tracks.append((a_pos, b_pos, c_pos))
        self.access_count['A'] += 1
        self.access_count['B'] += 1
        self.access_count['C'] += 1

    def get_heatmap_data(self) -> Dict[str, np.ndarray]:
        """
        Generate heatmap data showing access frequency.

        Returns:
            Dict with keys 'A', 'B', 'C' containing access frequency matrices
        """
        heatmaps = {
            'A': np.zeros((self.n, self.n)),
            'B': np.zeros((self.n, self.n)),
            'C': np.zeros((self.n, self.n))
        }

        for a_pos, b_pos, c_pos in self.tracks:
            heatmaps['A'][a_pos] += 1
            heatmaps['B'][b_pos] += 1
            heatmaps['C'][c_pos] += 1

        return heatmaps

    def get_statistics(self) -> Dict:
        """
        Get simulation statistics.

        Returns:
            Dictionary containing access counts and other metrics
        """
        return {
            'total_operations': len(self.tracks),
            'matrix_size': self.n,
            'block_size': self.block_size,
            'access_count': self.access_count.copy(),
            'theoretical_ops': self.n ** 3
        }


if __name__ == '__main__':
    # Test the simulator
    print("Testing GEMM Simulator...")

    # Test blocked version
    sim = GEMMSimulator(n=8, block_size=4)
    tracks = sim.simulate('kji', blocked=True)
    print(f"\nBlocked KJI (n=8, block=4):")
    print(f"  Total accesses: {len(tracks)}")
    print(f"  First 5 accesses: {tracks[:5]}")
    print(f"  Statistics: {sim.get_statistics()}")

    # Test unblocked version
    sim2 = GEMMSimulator(n=4, block_size=None)
    tracks2 = sim2.simulate('ijk', blocked=False)
    print(f"\nUnblocked IJK (n=4):")
    print(f"  Total accesses: {len(tracks2)}")
    print(f"  First 5 accesses: {tracks2[:5]}")

    print("\n✓ GEMM Simulator module working correctly!")