"""Reader for CTranslate2 `model.bin` files.

Format (see ctranslate2/specs/model_spec.py::ModelSpec._serialize):
  uint32   binary_version
  string   spec name          (uint16 len+1, bytes, NUL)
  uint32   spec revision
  uint32   num_variables
  per variable:
      string  name
      uint8   rank
      uint32  dim  * rank
      uint8   dtype id  (0 float32, 1 int8, 2 int16, 3 int32, 4 float16, 5 bfloat16)
      uint32  num_bytes
      raw data
  uint32   num_aliases
  per alias: string alias, string target_variable_name
"""

import struct
import numpy as np

DTYPES = {
    0: np.float32,
    1: np.int8,
    2: np.int16,
    3: np.int32,
    4: np.float16,
    5: None,  # bfloat16, handled specially
}


def read_ct2_model(path, load_data=True):
    with open(path, "rb") as f:
        buf = f.read()
    off = 0

    def u32():
        nonlocal off
        v = struct.unpack_from("I", buf, off)[0]
        off += 4
        return v

    def u16():
        nonlocal off
        v = struct.unpack_from("H", buf, off)[0]
        off += 2
        return v

    def u8():
        nonlocal off
        v = buf[off]
        off += 1
        return v

    def s():
        nonlocal off
        n = u16()
        v = buf[off:off + n - 1].decode("utf-8")
        off += n
        return v

    binary_version = u32()
    spec_name = s()
    revision = u32()
    nvars = u32()

    variables = {}
    for _ in range(nvars):
        name = s()
        rank = u8()
        shape = tuple(u32() for _ in range(rank))
        dtype_id = u8()
        nbytes = u32()
        if load_data:
            if dtype_id == 5:  # bfloat16 -> float32
                raw = np.frombuffer(buf, dtype=np.uint16, count=nbytes // 2, offset=off)
                arr = (raw.astype(np.uint32) << 16).view(np.float32).reshape(shape)
            else:
                dt = DTYPES[dtype_id]
                arr = np.frombuffer(
                    buf, dtype=dt, count=nbytes // np.dtype(dt).itemsize, offset=off
                ).reshape(shape)
            variables[name] = arr
        else:
            variables[name] = (shape, dtype_id, nbytes)
        off += nbytes

    naliases = u32()
    aliases = {}
    for _ in range(naliases):
        a = s()
        t = s()
        aliases[a] = t

    return dict(
        binary_version=binary_version,
        spec=spec_name,
        revision=revision,
        variables=variables,
        aliases=aliases,
    )


def dequantize(variables, name):
    """Return a float32 array for `name`, dequantizing with its scale if present."""
    w = variables[name]
    scale_name = name + "_scale"
    if scale_name in variables:
        scale = variables[scale_name]
        w = w.astype(np.float32)
        scale = np.asarray(scale, dtype=np.float32)
        if scale.ndim == 0:
            return w / scale
        return w / scale.reshape(-1, 1)
    return np.asarray(w, dtype=np.float32)


if __name__ == "__main__":
    import sys
    m = read_ct2_model(sys.argv[1], load_data=False)
    print("binary_version", m["binary_version"], "spec", m["spec"], "rev", m["revision"])
    print("n vars", len(m["variables"]), "n aliases", len(m["aliases"]))
    for k, v in m["variables"].items():
        print(k, v)
    for k, v in m["aliases"].items():
        print("ALIAS", k, "->", v)