"""Kernel build and ctypes bindings for the Mach Boom runtime."""
import ctypes, os, subprocess, sys
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))


def tiles_from_qidxs(Q, m, n, K):
    Q = np.asarray(Q).view(np.uint16).reshape(m, -1)
    nb, K = n // 16, int(K)
    assert Q.shape[1] == K * nb, (Q.shape, K, nb)
    t = Q.reshape(m // 16, 16, nb, K).transpose(0, 2, 1, 3)
    return t.astype(">u2").tobytes()


def signbits(s):
    return np.packbits(np.asarray(s).reshape(-1) > 0, bitorder="little").tobytes()


def tlut10(tl):
    tl = np.asarray(tl, np.float32).reshape(512, 2)
    t = np.concatenate([tl, tl]).copy()
    t[512:, 0] *= -1
    return np.ascontiguousarray(t)


def build(out=None, isa=None):
    src = [os.path.join(HERE, f) for f in ("tmk.c", "tmk_act.c", "tmk_model.c", "tmk_amx.c")]
    if out is None:
        out = os.path.join(HERE, f"tmk_{isa or 'native'}.so")
        if not os.access(HERE, os.W_OK):
            import hashlib
            h = hashlib.sha256(b"".join(open(f, "rb").read() for f in src)).hexdigest()[:12]
            d = os.path.join(os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")), "mach_boom")
            os.makedirs(d, exist_ok=True)
            out = os.path.join(d, f"tmk_{isa or 'native'}_{h}.so")
    if os.path.exists(out) and os.path.getmtime(out) >= max(os.path.getmtime(f) for f in src):
        return out
    cc = os.environ.get("CC", "cc")
    flags = {"native": "-march=native", "avx512": "-march=skylake-avx512", "avx2": "-march=haswell", "scalar": "",
             None: "-march=native"}[isa]
    tmp = out + f".{os.getpid()}.tmp"
    objs = [tmp + ".act.o", tmp + ".model.o", tmp + ".amx.o"]

    def run(flags, libs, amx):
        for f, o in zip(src[1:3], objs):
            subprocess.check_call([cc, "-O3", "-ffast-math", *flags.split(), "-fPIC", "-c", f, "-o", o])
        subprocess.check_call([cc, "-O3", *flags.split(), *([] if amx else ["-DTMK_NO_AMX"]), "-fPIC", "-c", src[3],
                               "-o", objs[2]], stderr=subprocess.DEVNULL if amx else None)
        subprocess.check_call([cc, "-O3", *flags.split(), "-shared", "-fPIC", src[0], *objs, "-o", tmp, *libs.split()])
    mv, nm = "-lmvec -lm -pthread", "-lm -pthread"
    tries = [(flags, mv, 1), (flags, mv, 0), (flags, nm, 1), (flags, nm, 0), ("", nm, 0)]
    try:
        for i, (fl, libs, amx) in enumerate(tries):
            try:
                run(fl, libs, amx)
                break
            except subprocess.CalledProcessError:
                if i == len(tries) - 1:
                    raise
        os.replace(tmp, out)
    finally:
        for o in objs:
            if os.path.exists(o):
                os.remove(o)
    return out


P_, I_, F_ = ctypes.c_void_p, ctypes.c_int, ctypes.c_float
MAXL = 8


class Chunk(ctypes.Structure):
    _fields_ = [("codes", P_), ("su", P_), ("sv", P_), ("ws", F_), ("m", I_), ("n", I_)]


class Lin(ctypes.Structure):
    _fields_ = [("ch", ctypes.POINTER(Chunk)), ("nch", I_), ("step", I_), ("tlut", P_), ("bias", P_), ("dense", P_),
                ("m", I_), ("n", I_)]


class LMS(ctypes.Structure):
    _fields_ = [("D", I_), ("NL", I_), ("NH", I_), ("HD", I_), ("FF", I_), ("Pmax", I_), ("eps", F_), ("inv_freq", P_),
                ("ln", (P_ * 4) * MAXL), ("lin", (Lin * 6) * MAXL), ("lnf_w", P_), ("lnf_b", P_), ("in_proj", Lin),
                ("in_add", P_)]


class FlowS(ctypes.Structure):
    _fields_ = [("D", I_), ("NL", I_), ("LD", I_), ("eps", F_), ("inp", Lin), ("cond", Lin), ("ada", Lin * MAXL),
                ("fc1", Lin * MAXL), ("fc2", Lin * MAXL), ("fada", Lin), ("fproj", Lin), ("norm", (P_ * 2) * MAXL),
                ("time", P_)]


class K:
    def __init__(self, so=None):
        lib = ctypes.CDLL(so or build())
        P, I, F, Lg = ctypes.c_void_p, ctypes.c_int, ctypes.c_float, ctypes.c_long
        lib.tmk_isa.restype = ctypes.c_char_p
        lib.tmk_qgemv.argtypes = [P, I, I, I, P, P, P]
        lib.tmk_qgemv_rows.argtypes = [P, I, I, P, P, P, I, I]
        lib.tmk_tlut10.argtypes = [P, P]
        lib.tmk_threads.argtypes = [I]; lib.tmk_threads.restype = I
        lib.tmk_qdense.argtypes = [P, I, I, I, P, F, P]
        lib.tmk_fwht.argtypes = [P, I, I]
        lib.tmk_signs.argtypes = [P, I, I, P]
        lib.tmk_rot.argtypes = [P, I, I, P, P, F]
        lib.tmk_linear_vec.argtypes = [P, I, I, I, P, P, P, F, P, P, P]
        lib.tmk_lm_work.argtypes = [ctypes.POINTER(LMS)]; lib.tmk_lm_work.restype = ctypes.c_size_t
        lib.tmk_lm_step.argtypes = [ctypes.POINTER(LMS), P, P, I, P, P, P, P]
        lib.tmk_flow_step.argtypes = [ctypes.POINTER(FlowS), P, P, P, P]
        lib.tmk_softmax_band.argtypes = [P, Lg, I, I, I, I]
        lib.tmk_bf16_pack_k.argtypes = [P, Lg, I, I, P, I, I]
        lib.tmk_bf16_pack_j.argtypes = [P, Lg, I, I, P, I, I]
        lib.tmk_bf16_x.argtypes = [P, Lg, I, I, I, I, P, I]
        lib.tmk_colrot.argtypes = [P, I, I, P, F]
        lib.tmk_dtf_work.argtypes = [I, I, I, I]; lib.tmk_dtf_work.restype = Lg
        lib.tmk_dtf_layer.argtypes = [P, I, I, I, I, F, P, P, P, P, P, P, P, P, P, I, I, P, Lg, P]
        lib.tmk_wpack.argtypes = [P, I, I, I, P, F, P, P, P, P, P, I, I]
        lib.tmk_amx_gemm.argtypes = [P, I, I, P, I, P, Lg, I]
        for f in ("tmk_gelu", "tmk_silu", "tmk_elu"):
            getattr(lib, f).argtypes = [P, Lg]
        self.lib, self.isa = lib, lib.tmk_isa().decode()


def ptr(a):
    return a.ctypes.data if isinstance(a, np.ndarray) else ctypes.cast(a, ctypes.c_void_p).value