"""Print every number the lesson quotes. Run: PYTHONDONTWRITEBYTECODE=1 python3 show_numbers.py"""
from decimal import ROUND_HALF_UP, Decimal

import gpu_split as g


def half_up(x: float, places: str = "0.01") -> str:
    """Round the way the widget's JavaScript toFixed does on an exact tie (0.625 -> 0.63); Python's format rounds ties to even."""
    return str(Decimal(x).quantize(Decimal(places), rounding=ROUND_HALF_UP))


def gib(b: int) -> str:
    return half_up(b / g.GIB)


def ms(s: float) -> str:
    return half_up(s * 1e3)


def halved_link_times(key: str, n: int, words: int) -> dict:
    """Word time in seconds on each link if the data sheet's link rates count both directions, so one direction runs at half."""
    saved = dict(g.LINKS)
    try:
        for link, (label, bandwidth, step_us) in saved.items():
            g.LINKS[link] = (label, bandwidth / 2, step_us)
        return {link: g.word_time(key, n, link, words=words)["total"] for link in g.LINKS}
    finally:
        g.LINKS.update(saved)


def main() -> None:
    print(f"GPU: {g.GPU_GIB} GiB, {g.USABLE_BYTES / g.GIB:.1f} GiB usable, {g.GPU_BANDWIDTH / 1e9:,.0f} GB/s, {g.GPU_FLOPS / 1e12:.0f} TFLOP/s")
    for key, (label, bandwidth, step_us) in g.LINKS.items():
        print(f"{label}: {bandwidth / 1e9:.0f} GB/s, {step_us} us a step, base {g.BASE_US} us")

    print("\n== Models (16 bits)")
    for key, (name, _) in g.MODELS.items():
        cfg = g.load_config(key)
        print(f"{name}: {g.param_count(cfg):,} weights, {gib(g.weight_bytes(cfg))} GiB, {cfg['num_hidden_layers']} layers, "
              f"hidden {cfg['hidden_size']}, message {g.message_bytes(cfg):,} bytes, "
              f"{g.ALL_REDUCES_PER_LAYER * cfg['num_hidden_layers']} swaps a word, fewest GPUs {g.least_n(key)}")

    print("\n== Memory per GPU in GiB: weights, cache, working, total, fits")
    for key, (name, _) in g.MODELS.items():
        for n in g.COUNTS:
            m = g.memory_per_gpu(key, n)
            print(f"{name} x{n}: {gib(m['weights'])} + {gib(m['cache'])} + {gib(m['working'])} = {gib(m['total'])}  fits={g.fits(key, n)}")

    print("\n== One word, one user, GPUs slice every layer: read ms | wait ms | total ms | wait share | words/s")
    for key, (name, _) in g.MODELS.items():
        for n in g.COUNTS:
            for link, (label, _, _) in g.LINKS.items():
                t = g.word_time(key, n, link)
                print(f"{name} x{n} {label}: {ms(t['read'])} | {ms(t['transfer'])} | {ms(t['total'])} | {t['share'] * 100:.1f}% | {t['words_per_s']:.1f}")

    print("\n== Best whole-number GPU count (any count that fits, up to 64), words at once 1/16/64")
    for key, (name, _) in g.MODELS.items():
        print(name, {w: {link: g.best_n(key, link, w) for link in g.LINKS} for w in g.WORDS})

    print("\n== Best GPU count that divides the model's attention heads (up to 64), one user")
    for key, (name, _) in g.MODELS.items():
        print(name, f"{g.load_config(key)['num_attention_heads']} heads", {link: g.best_n(key, link, divide_heads=True) for link in g.LINKS})

    print("\n== 8B: speed-up of each doubling, by link")
    for link, (label, _, _) in g.LINKS.items():
        times = [g.word_time("8b", n, link)["total"] for n in g.COUNTS]
        print(label, [f"x{a / b:.2f}" for a, b in zip(times, times[1:])])

    print("\n== 70B on 8 GPUs, words at once: total ms | wait share")
    for words in g.WORDS:
        row = [f"{label} {ms(g.word_time('70b', 8, link, words=words)['total'])} | {g.word_time('70b', 8, link, words=words)['share'] * 100:.1f}%"
               for link, (label, _, _) in g.LINKS.items()]
        print(words, *row)

    print("\n== Layer groups (pipeline), one user: total ms | wait share")
    for key, (name, _) in g.MODELS.items():
        for n in g.COUNTS:
            t = g.word_time(key, n, "fast", "groups")
            u = g.word_time(key, n, "slow", "groups")
            print(f"{name} x{n}: fast {ms(t['total'])} | {t['share'] * 100:.2f}%   slow {ms(u['total'])} | {u['share'] * 100:.2f}%   one GPU busy at a time: {100 / n:.1f}% of GPUs")

    print("\n== If the sheet's link rates count both directions (halve them), 70B on 8 GPUs, total ms: as written -> halved")
    for words in (1, 64):
        halved = halved_link_times("70b", 8, words)
        print(words, {link: f"{ms(g.word_time('70b', 8, link, words=words)['total'])} -> {ms(halved[link])}"
                      f" ({(halved[link] / g.word_time('70b', 8, link, words=words)['total'] - 1) * 100:.0f}% more)" for link in g.LINKS})


if __name__ == "__main__":
    main()