"""Export final inference weights and reference logits; no converter dependency."""
import argparse
import hashlib
import gzip
import shutil
import json
from pathlib import Path
import numpy as np
import tensorflow as tf
from train import Config, make_model


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("run", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--quantize", action="store_true", help="Export affine uint8 download weights")
    args = parser.parse_args()
    tf.config.threading.set_intra_op_parallelism_threads(4)
    tf.config.threading.set_inter_op_parallelism_threads(1)
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        tf.config.set_visible_devices(gpus[0], 'GPU')
        tf.config.experimental.set_memory_growth(gpus[0], True)
    config = Config(**json.loads((args.run / "config.json").read_text())["model"])
    model = make_model(config)
    checkpoint = tf.train.latest_checkpoint(str(args.run / "checkpoints"))
    tf.train.Checkpoint(model=model).restore(checkpoint).expect_partial().assert_existing_objects_matched()
    args.output.mkdir(parents=True, exist_ok=True)
    specs = []
    digest = hashlib.sha256()
    original_digest = hashlib.sha256()
    binary = args.output / ("weights-int8.bin" if args.quantize else "weights.bin")
    with binary.open("wb") as stream:
        def weight(name, value, shape=None):
            array = value.numpy().astype("<f4")
            if shape:
                array = array.reshape(shape)
            original_digest.update(array.tobytes())
            spec = {"name": name, "shape": list(array.shape), "dtype": "float32"}
            if args.quantize:
                minimum = float(array.min())
                scale = (float(array.max()) - minimum) / 255 or 1.0
                codes = np.clip(np.rint((array - minimum) / scale), 0, 255).astype(np.uint8)
                spec["quantization"] = {"dtype":"uint8", "min":minimum, "scale":scale}
                # Reference logits must describe the decoded download, not FP32 originals.
                value.assign((codes.astype(np.float32) * scale + minimum).reshape(value.shape))
                raw = codes.tobytes()
            else:
                raw = array.tobytes()
            stream.write(raw)
            digest.update(raw)
            specs.append(spec)

        def norm(name, layer):
            weight(name + ".gamma", layer.gamma)
            weight(name + ".beta", layer.beta)

        def dense(name, layer, shape=None):
            weight(name + ".kernel", layer.kernel, shape)
            weight(name + ".bias", layer.bias, (int(np.prod(layer.bias.shape)),))

        weight("tokens", model.tokens.embeddings)
        weight("positions", model.positions.embeddings)
        for i, name in enumerate(model.block_names):
            block = getattr(model, name)
            prefix = f"b{i}"
            norm(prefix + ".norm1", block.norm1)
            norm(prefix + ".norm2", block.norm2)
            for short, projection in (("q", block.attention.query_dense),
                                      ("k", block.attention.key_dense),
                                      ("v", block.attention.value_dense),
                                      ("o", block.attention.output_dense)):
                dense(prefix + "." + short, projection, (config.width, config.width))
            dense(prefix + ".expand", block.expand)
            dense(prefix + ".contract", block.contract)
        norm("norm", model.norm)
        dense("head", model.output_head)
    if args.quantize:
        final = args.output / ("weights-int8-" + digest.hexdigest()[:12] + ".bin")
        binary.replace(final)
        binary = final
    with binary.open('rb') as source, (args.output / (binary.name + '.gz')).open('wb') as target:
        with gzip.GzipFile(filename='', mode='wb', compresslevel=6, fileobj=target, mtime=0) as compressed:
            shutil.copyfileobj(source, compressed)
    manifest_name = 'model-int8.json' if args.quantize else 'model.json'
    verification_name = 'verification-int8.json' if args.quantize else 'verification.json'
    manifest = {"context": config.context, "width": config.width, "layers": config.layers,
                "heads": config.heads, "parameters": model.count_params(),
                "checkpoint": Path(checkpoint).name, "weights": specs,
                "weight_file": binary.name, "weight_bytes": binary.stat().st_size,
                "gzip_bytes": (args.output / (binary.name + ".gz")).stat().st_size,
                "verification_file":verification_name, "source_sha256":original_digest.hexdigest(),
                "sha256": digest.hexdigest(), "epsilon": 1e-5}
    (args.output / manifest_name).write_text(json.dumps(manifest, indent=2) + "\n")
    cases = []
    for prompt in ("Far far away", "Mira found a key under the old tree.", "A little café"):
        ids = list(prompt.encode())
        steps = []
        for extra in ([], *[[byte] for byte in b" in the"]):
            ids.extend(extra)
            logits = model(tf.constant([ids], tf.int32), training=False)[0, -1].numpy()
            steps.append({"input": extra or list(prompt.encode()), "logits": logits.tolist()})
        cases.append({"prompt": prompt, "steps": steps})
    ids = list(("A robot found a bright star. " * 20).encode()[:config.context])
    cases.append({"prompt": "full context", "steps": [{"input": ids,
        "logits": model(tf.constant([ids], tf.int32), training=False)[0, -1].numpy().tolist()}]})
    (args.output / verification_name).write_text(json.dumps({"cases": cases}) + "\n")
    print(json.dumps({key: manifest[key] for key in ("checkpoint", "weight_bytes", "sha256")}), flush=True)


if __name__ == "__main__":
    main()