
import math

def estimate_energy():
    # Model Constants (pJ per operation - approximate values based on 7nm/5nm technology)
    # Source: Horowitz / Stanford High Performance Computing
    E_FP32_MUL = 3.7   # pJ
    E_FP32_ADD = 0.9   # pJ
    E_INT8_ADD = 0.03  # pJ (Very cheap)
    E_INT32_ADD = 0.1  # pJ
    E_DRAM_ACCESS = 640 # pJ (Costly!)
    E_SRAM_ACCESS = 5   # pJ (L1 Cache)
    E_OPTICAL_OP = 0.001 # pJ (Theoretical passive optical propagation)

    sizes = [256, 512, 1024, 2048]

    print(f"{'N':<8} {'Method':<15} {'Ops (M)':<10} {'Energy (uJ)':<15} {'Efficiency (TOPS/W)':<20}")
    print("-" * 75)

    for N in sizes:
        # 1. Standard GEMM (FP32)
        # N^3 Mults + N^3 Adds + 2*N^2 Memory Accesses (Naive)
        # Ops = 2 * N^3
        ops_gemm = 2 * (N**3)
        # Assuming heavy cache reuse (tiling), memory cost is minimized but still present.
        # Simplification: 1 DRAM access per 100 FLOPS
        energy_gemm = (N**3) * (E_FP32_MUL + E_FP32_ADD) + ((N**3)/100) * E_DRAM_ACCESS

        # 2. RSR (Method of Four Russians)
        # Precompute: O(N * 2^(2*block_size)) -> Negligible for large N
        # Main Loop: O(N^3 / block_size) lookups + additions
        # Block size = 4
        # Ops: (N^3 / 4) Lookups + (N^3 / 4) INT32 Adds
        block_size = 4
        ops_rsr = (N**3) / block_size
        # Energy: SRAM Access (Lookup) + INT32 Add
        energy_rsr = ops_rsr * (E_SRAM_ACCESS + E_INT32_ADD)

        # 3. Optical (MZI Mesh)
        # Constant/Linear propagation energy
        # Power = Laser Power + DAC/ADC
        # Assume effective per-MAC energy is vastly lower
        ops_optical = 2 * (N**3) # Equivalent effective ops
        energy_optical = ops_optical * E_OPTICAL_OP

        # Convert pJ to uJ (1e-6 J = 1e6 pJ)
        print(f"{N:<8} {'FP32-GEMM':<15} {ops_gemm/1e6:<10.2f} {energy_gemm/1e6:<15.2f} { (ops_gemm/energy_gemm):<20.2f}")
        print(f"{N:<8} {'TRME-RSR':<15} {ops_rsr/1e6:<10.2f} {energy_rsr/1e6:<15.2f} { (ops_gemm/energy_rsr):<20.2f}") # Use equivalent FP32 Ops for TOPS/W comparison
        print(f"{N:<8} {'TRME-Optical':<15} {ops_optical/1e6:<10.2f} {energy_optical/1e6:<15.2f} { (ops_optical/energy_optical):<20.2f}")
        print("-" * 75)

if __name__ == "__main__":
    estimate_energy()
