#!/usr/bin/env python3
"""
Thermal‑Aware XYFLOW Hyper‑Dot – Nanosecond Pause Cloud
========================================================
Layer 0 : Nano‑scale temporal attractor (spin‑wait cloud)
Layer 1 : LUT projection (0 arithmetic ops)
Layer 2 : N‑ary reduction (1 arithmetic op)
"""

import numpy as np
import ctypes
import time
import os

# ====================================================================
# 1. Compile the nano‑pause C kernel
# ====================================================================
C_SOURCE = r'''
#include <stdint.h>
#include <omp.h>

static inline void nano_pause(int loops) {
    for (volatile int i = 0; i < loops; ++i) {
        __asm__ volatile("" ::: "memory");
    }
}

void hyper_dot_nano(
    const uint8_t *X, const uint8_t *W_T,
    uint32_t *out,
    const uint16_t *lut,
    int M, int K, int N,
    int pause_every, int pause_loops)
{
    #pragma omp parallel for collapse(2)
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            const uint8_t *row_x = X + i * K;
            const uint8_t *row_w = W_T + j * K;
            uint32_t sum = 0;
            int cnt = 0;

            #pragma omp simd reduction(+:sum)
            for (int k = 0; k < K; k++) {
                sum += lut[row_x[k] * 256 + row_w[k]];
                if (++cnt >= pause_every) {
                    cnt = 0;
                    nano_pause(pause_loops);
                }
            }
            out[i * N + j] = sum;
        }
    }
}
'''

with open("hyper_dot_nano.c", "w") as f:
    f.write(C_SOURCE)
os.system("gcc -shared -fPIC -O3 -march=native -fopenmp hyper_dot_nano.c -o libhyper_nano.so")

lib = ctypes.CDLL("./libhyper_nano.so")
lib.hyper_dot_nano.argtypes = [
    ctypes.POINTER(ctypes.c_uint8),
    ctypes.POINTER(ctypes.c_uint8),
    ctypes.POINTER(ctypes.c_uint32),
    ctypes.POINTER(ctypes.c_uint16),
    ctypes.c_int, ctypes.c_int, ctypes.c_int,
    ctypes.c_int, ctypes.c_int
]

# ====================================================================
# 2. Data preparation
# ====================================================================
M, K, N = 100, 784, 100
np.random.seed(42)
X = np.random.randint(0, 256, (M, K), dtype=np.uint8)
W = np.random.randint(0, 256, (K, N), dtype=np.uint8)
W_T = np.ascontiguousarray(W.T)

lut = np.zeros((256, 256), dtype=np.uint16)
for i in range(256):
    lut[i, :] = i * np.arange(256, dtype=np.uint16)
lut = lut.flatten()

out_sleepless = np.zeros((M, N), dtype=np.uint32)
out_nano      = np.zeros((M, N), dtype=np.uint32)

pX   = X.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
pWT  = W_T.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
pLut = lut.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16))
pOutS = out_sleepless.ctypes.data_as(ctypes.POINTER(ctypes.c_uint32))
pOutN = out_nano.ctypes.data_as(ctypes.POINTER(ctypes.c_uint32))

# ====================================================================
# 3. Stress test – sleepless first, then nano‑cloud immediately
# ====================================================================
REPEATS = 200

print("Heating up CPU with sleepless version ...")
start = time.perf_counter()
for _ in range(REPEATS):
    lib.hyper_dot_nano(pX, pWT, pOutS, pLut, M, K, N, 1000000, 0)  # effectively no pause
sleepless_time = time.perf_counter() - start

# NO COOLDOWN – immediately run the nano‑cloud version
print("Running nano‑cloud version (tiny spin‑waits) on HOT CPU ...")
PAUSE_EVERY = 8      # pause after every 8 inner iterations
PAUSE_LOOPS = 12     # ~12 cycles (~4 ns on 3 GHz)
start = time.perf_counter()
for _ in range(REPEATS):
    lib.hyper_dot_nano(pX, pWT, pOutN, pLut, M, K, N, PAUSE_EVERY, PAUSE_LOOPS)
nano_time = time.perf_counter() - start

# ====================================================================
# 4. Verification & output
# ====================================================================
Xf = X.astype(np.float32); Wf = W.astype(np.float32)
ref = np.dot(Xf, Wf)
assert np.allclose(out_sleepless, ref), "Sleepless mismatch"
assert np.allclose(out_nano, ref), "Nano version mismatch"

print("\n" + "="*60)
print("XYFLOW Hyper‑Dot with Nanosecond Pause Cloud")
print("="*60)
print(f"Matrix shape      : {M}×{K} × {K}×{N}")
print(f"Repetitions       : {REPEATS}")
print(f"Nano‑pause cloud  : every {PAUSE_EVERY} iters, spin {PAUSE_LOOPS} cycles")
print()
print(f"Sleepless total time : {sleepless_time:.3f} s  ({sleepless_time/REPEATS*1e6:.1f} µs avg)")
print(f"Nano‑cloud total time: {nano_time:.3f} s  ({nano_time/REPEATS*1e6:.1f} µs avg)")
print(f"Speed‑up factor      : {sleepless_time/nano_time:.2f}x")
print()
if nano_time < sleepless_time:
    print(">> SUCCESS: The nano‑pause cloud defeated thermal throttling. <<")
else:
    print("(No throttling detected – nano pauses added minimal overhead.)")
print("Arithmetic operations per dot product: 1 (single N‑ary sum).")
print("="*60)