import numpy as np
import os

# -----------------------------
# HDGL Channel
# -----------------------------
class HDGLChannel:
    def __init__(self, id):
        self.id = id
        self.pages = []  # list of memory pages (bytearrays)
        self.tension = 0.0

    def assign_page(self, page_data: bytes):
        self.pages.append(bytearray(page_data))
        self.update_tension(len(page_data))

    def operate(self):
        # Safe increment of all bytes in all pages
        for i, page in enumerate(self.pages):
            arr = np.frombuffer(page, dtype=np.uint8).copy()
            arr = ((arr.astype(np.uint16) + 1) % 256).astype(np.uint8)
            self.pages[i] = bytearray(arr)
        self.update_tension(sum(len(p) for p in self.pages))

    def update_tension(self, size):
        self.tension = 0.8 * self.tension + 0.2 * (size / 4096.0)  # normalize by page size

# -----------------------------
# HDGL Lattice
# -----------------------------
class HDGLLattice:
    def __init__(self, num_channels=16):
        self.channels = [HDGLChannel(i) for i in range(num_channels)]

    def select_channel(self):
        # Pick channel with lowest tension
        tensions = np.array([ch.tension for ch in self.channels])
        min_idx = np.argmin(tensions)
        return self.channels[min_idx]

    def map_page(self, page_data: bytes):
        # Assign page dynamically to lowest-tension channel
        ch = self.select_channel()
        ch.assign_page(page_data)

    def operate_all(self):
        # Operate on all channels (dynamic flip-chip)
        for ch in self.channels:
            ch.operate()

# -----------------------------
# Kernel / Initramfs Loader
# -----------------------------
def load_pages(path, page_size=4096):
    """Load a file and split into pages"""
    pages = []
    with open(path, "rb") as f:
        while True:
            data = f.read(page_size)
            if not data:
                break
            if len(data) < page_size:
                data += b'\x00' * (page_size - len(data))
            pages.append(data)
    return pages

# -----------------------------
# Main Execution
# -----------------------------
if __name__ == "__main__":
    lattice = HDGLLattice(num_channels=16)

    # Demo: load kernel and initramfs pages
    kernel_pages = load_pages("vmlinuz.bin")
    initramfs_pages = load_pages("initrd.img")

    for page in kernel_pages + initramfs_pages:
        lattice.map_page(page)

    # Run lattice operations (simulate HDGL superposition processing)
    lattice.operate_all()

    # Print channel stats
    for ch in lattice.channels:
        print(f"Channel {ch.id}: {len(ch.pages)} pages, tension {ch.tension:.4f}")
