from utils import Configuration
from pynq import Overlay, PL, allocate
from pathlib import Path
import numpy as np
import time

def get_bitstream_file(data_dir):
    """
    Ensure that exactly one .bit file and one .hwh file exist in dir_path.
    Raises FileNotFoundError if missing, or RuntimeError if there are duplicates.
    """
    p = Path(data_dir)

    # Gather files
    bit_files = list(p.glob('*.bit'))
    hwh_files = list(p.glob('*.hwh'))

    # Check counts
    if not bit_files:
        raise FileNotFoundError(f"No .bit file found in {p!r}")
    if len(bit_files) > 1:
        raise RuntimeError(f"Expected exactly one .bit file in {p!r}, but found {len(bit_files)}: {bit_files}")

    if not hwh_files:
        raise FileNotFoundError(f"No .hwh file found in {p!r}")
    if len(hwh_files) > 1:
        raise RuntimeError(f"Expected exactly one .hwh file in {p!r}, but found {len(hwh_files)}: {hwh_files}")

    # If we reach here, exactly one of each was found
    bit_file = bit_files[0]
    hwh_file = hwh_files[0]

    return bit_file

def get_overlay(bitstream_filename):

    try:
        bitstream = str(bitstream_filename)
    except (FileNotFoundError, RuntimeError) as e:
        print(f"Error: {e}")
        exit(1)

    PL.reset()
    ol = Overlay(bitstream)
    # print(ol.ip_dict.keys())

    return ol

def safe_set_register(ip, reg_name, value):
    """Set an IP register if it exists in the register map."""

    if hasattr(ip.register_map, reg_name):
        setattr(ip.register_map, reg_name, value)
        # print(f"✅ Set {reg_name} to {hex(value)}")
    else:
        print(f"Register '{reg_name}' not found in IP '{ip}'")

class DualEncoderKernel:
    def __init__(self, data_dir, config: Configuration):
        """Prepare the PYNQ overlay and persistent DMA buffers for one run."""
        self.data_dir = data_dir
        self.config = config
        self.bit_file = get_bitstream_file(data_dir)
        self.overlay = get_overlay(self.bit_file)

        self.ip = self.overlay.siso_0
        self.dma = self.overlay.axi_dma_0

        # print(self.ip.register_map)

        self.SEQ_LEN = self.config.model.sequence_length
        self.D_MODEL = self.config.model.d_model


        self.store_type = np.float16 if self.config.kernel.store_fp16 else np.float32

        # persistent buffers
        self.buf_in = allocate(shape=(self.SEQ_LEN, self.D_MODEL), dtype=np.float32)
        self.buf_out = allocate(shape=(self.SEQ_LEN, self.D_MODEL), dtype=np.float32)

        # Weight buffers (for model parameters)
        self.buf_weights = {}

    def load_parameters(self, weights: dict):
        """Copy flattened encoder weights into contiguous FPGA buffers."""

        def alloc_and_set(name, arr):
            buf = allocate(shape=arr.shape, dtype=self.store_type)
            buf[:] = arr
            buf.flush()
            safe_set_register(self.ip, name, buf.physical_address)
            self.buf_weights[name] = buf

        alloc_and_set("W_QKV_1", weights["W_QKV"])
        alloc_and_set("B_QKV_1", weights["B_QKV"])
        alloc_and_set("W_O_1", weights["W_O"])
        alloc_and_set("B_O_1", weights["B_O"])
        alloc_and_set("norm1_gamma_1", weights["norm1_gamma"])
        alloc_and_set("norm1_beta_1", weights["norm1_beta"])
        alloc_and_set("norm2_gamma_1", weights["norm2_gamma"])
        alloc_and_set("norm2_beta_1", weights["norm2_beta"])
        alloc_and_set("ff_w1_1", weights["ff_w1"])
        alloc_and_set("ff_b1_1", weights["ff_b1"])
        alloc_and_set("ff_w2_1", weights["ff_w2"])
        alloc_and_set("ff_b2_1", weights["ff_b2"])

        start = time.time()
        self.ip.register_map.load_parameters = 1
        self.ip.register_map.seq_id = 0
        self.ip.register_map.CTRL.AP_START = 1

        while int(self.ip.register_map.CTRL.AP_DONE) != 1:
            if (time.time() - start) > 10:
                raise TimeoutError("Timeout while loading parameters into kernel.")
            time.sleep(0.1)
            pass

        finished = time.time()
        print(f"Loaded parameters into kernel in {finished - start:.3f} seconds.")
        self.ip.register_map.load_parameters = 0

    def run(self, input_array, mask_count):
        """Run one encoder input through the accelerator."""

        self.buf_in[:] = input_array.astype(np.float32, copy=False)
        self.buf_in.flush()

        self.ip.register_map.mask_count = mask_count
        self.ip.register_map.load_parameters = 0

        start = time.time()
        self.dma.recvchannel.transfer(self.buf_out)
        self.ip.register_map.CTRL.AP_START = 1
        self.dma.sendchannel.transfer(self.buf_in)

        self.dma.sendchannel.wait()
        self.dma.recvchannel.wait()

        while int(self.ip.register_map.CTRL.AP_DONE) != 1:
            if (time.time() - start) > 10:
                raise TimeoutError("Timeout while running kernel.")
            time.sleep(0.01)
            pass

        self.buf_out.invalidate()
        out32 = np.array(self.buf_out, dtype=np.float32)
        return out32.astype(np.float16)

    def push_zeros(self):
        """Advance the hardware pipeline during warm-up/cool-down cycles."""
        print("Pushing zeros through the kernel")
        in_zeros = np.zeros((self.SEQ_LEN, self.D_MODEL), dtype=self.store_type)
        return self.run(in_zeros, 10)


        
