#!/usr/bin/env python3
"""Apodex scaling-rule hyperparameter calculator.

This is a standalone calculator distilled from:
  /mnt/pre_train/sean/Megatron-LM/examples/apodex/scaling/design_matrix.py
  /mnt/pre_train/sean/Megatron-LM/examples/apodex/scaling/redesign_lr_bs.py

It supports two common entry points:
  1. Provide compute budget C and ratio r = D / N_active.
  2. Provide N_active and D_tokens directly.

By default it uses the newer v2 training-hparam rule:
  B*  = 0.0249 * D^0.4111
  lr  = 2.6561 * N^(-0.5256) * B^0.4897

Use --rule v1 to reproduce the older design_matrix.py batch/lr rule.

These are project-specific heuristics, not universal scaling laws. B is the
number of sequences per effective batch; N counts active non-embedding matrix
parameters in this architecture. C = 6ND omits attention's sequence-dependent
cost, communication, embeddings and other overhead. Node counts are planning
heuristics, not memory-fit guarantees. Inspect realized values after rounding.
"""

from __future__ import annotations

import argparse
import json
import math
from dataclasses import asdict, dataclass
from typing import Any


# Architecture invariants inherited from a1b_base.yaml.
VOCAB_SIZE = 151936
QK_HEAD_DIM = 128
QK_POS_EMB_HEAD_DIM = 64
V_HEAD_DIM = 128
NUM_EXPERTS = 128
TOP_K = 8
SEQ_LEN = 8192

H_PER_LAYER = 1536 / 28
H_PER_HEAD = 1536 / 16
Q_LORA_RATIO = 1024 / 1536
KV_LORA_RATIO = 512 / 1536
MOE_FFN_RATIO = 640 / 1536
DENSE_FFN_RATIO = 4096 / 1536

H_MULT = 64
LORA_MULT = 8
FFN_MULT = 64
HIDDEN_MIN = 192
HIDDEN_MAX = 4096

WARMUP_TOKENS = 1.5e9
WARMUP_FRAC_CAP = 0.05


@dataclass(frozen=True)
class Arch:
    hidden_size: int
    num_layers: int
    num_heads: int
    q_lora_rank: int
    kv_lora_rank: int
    moe_ffn: int
    dense_ffn: int
    num_experts: int = NUM_EXPERTS
    top_k: int = TOP_K


def round_to_multiple(x: float, mult: int) -> int:
    return max(mult, int(round(x / mult) * mult))


def nearest_pow2(x: float, lo: int = 1) -> int:
    """Nearest power of two, ties rounded up."""
    if x <= lo:
        return lo
    lo_pow = 1 << int(math.floor(math.log2(x)))
    hi_pow = lo_pow * 2
    return hi_pow if (hi_pow - x) <= (x - lo_pow) else lo_pow


def arch_from_hidden(hidden_size: int) -> Arch:
    return Arch(
        hidden_size=hidden_size,
        num_layers=max(4, round(hidden_size / H_PER_LAYER)),
        num_heads=nearest_pow2(hidden_size / H_PER_HEAD, lo=2),
        q_lora_rank=round_to_multiple(hidden_size * Q_LORA_RATIO, LORA_MULT),
        kv_lora_rank=round_to_multiple(hidden_size * KV_LORA_RATIO, LORA_MULT),
        moe_ffn=round_to_multiple(hidden_size * MOE_FFN_RATIO, FFN_MULT),
        dense_ffn=round_to_multiple(hidden_size * DENSE_FFN_RATIO, FFN_MULT),
    )


def attn_params_per_layer(arch: Arch) -> int:
    h, heads = arch.hidden_size, arch.num_heads
    return (
        h * arch.q_lora_rank
        + arch.q_lora_rank * heads * (QK_HEAD_DIM + QK_POS_EMB_HEAD_DIM)
        + h * (arch.kv_lora_rank + QK_POS_EMB_HEAD_DIM)
        + arch.kv_lora_rank * heads * (QK_HEAD_DIM + V_HEAD_DIM)
        + heads * V_HEAD_DIM * h
    )


def moe_active_ffn_per_layer(arch: Arch) -> int:
    return (arch.top_k + 1) * 3 * arch.hidden_size * arch.moe_ffn


def dense_ffn_params(arch: Arch) -> int:
    return 3 * arch.hidden_size * arch.dense_ffn


def n_active_nonembed(arch: Arch) -> int:
    layers = arch.num_layers
    return (
        layers * attn_params_per_layer(arch)
        + (layers - 1) * moe_active_ffn_per_layer(arch)
        + dense_ffn_params(arch)
    )


def find_hidden_for_target_n(target_n: float) -> int:
    lo, hi = HIDDEN_MIN, HIDDEN_MAX
    while hi - lo > H_MULT:
        mid = round_to_multiple((lo + hi) / 2, H_MULT)
        if n_active_nonembed(arch_from_hidden(mid)) < target_n:
            lo = mid
        else:
            hi = mid
    return min(
        [lo, hi],
        key=lambda hidden_size: abs(
            n_active_nonembed(arch_from_hidden(hidden_size)) - target_n
        ),
    )


def global_bs_v1(n_active: float) -> int:
    """design_matrix.py: doubles every 2x in N, anchored at 64 @ 40M, capped at 2048."""
    k = max(0, int(math.floor(math.log2(n_active / 40e6))))
    return 64 * (2 ** min(k, 5))


def peak_lr_v1(n_active: float) -> float:
    """design_matrix.py Chinchilla-inspired LR fit."""
    return min(6e-4, 6e-4 * (75e6 / n_active) ** 0.35)


def global_bs_v2(d_tokens: float) -> int:
    """redesign_lr_bs.py: B* = 0.0249 * D^0.4111, snapped to nearest power of 2."""
    return nearest_pow2(0.0249 * d_tokens**0.4111)


def peak_lr_v2(n_active: float, global_bs: int) -> float:
    """redesign_lr_bs.py: eta = 2.6561 * N^-0.5256 * B^0.4897."""
    return 2.6561 * n_active**-0.5256 * global_bs**0.4897


def nodes_for(n_active: float) -> int:
    if n_active <= 2.0e8:
        return 1
    if n_active <= 7.0e8:
        return 2
    if n_active <= 1.2e9:
        return 4
    return 8


def warmup_iters_for(d_tokens: float, global_bs: int, seq_len: int, train_iters: int) -> int:
    warmup_tokens = min(WARMUP_TOKENS, WARMUP_FRAC_CAP * d_tokens)
    raw = round(warmup_tokens / (global_bs * seq_len))
    iter_cap = max(0, int(train_iters * WARMUP_FRAC_CAP))
    return min(iter_cap, max(50, raw))


def parse_number(text: str) -> float:
    """Accept values like 3e19, 10B, 500M, 1.2T."""
    value = text.strip().replace("_", "")
    suffix_multipliers = {
        "k": 1e3,
        "m": 1e6,
        "b": 1e9,
        "t": 1e12,
        "p": 1e15,
    }
    if value[-1:].lower() in suffix_multipliers:
        return float(value[:-1]) * suffix_multipliers[value[-1].lower()]
    return float(value)


def prompt_number(name: str, default: float | None = None) -> float | None:
    suffix = f" [{default:g}]" if default is not None else ""
    raw = input(f"{name}{suffix}: ").strip()
    if not raw and default is not None:
        return default
    if not raw:
        return None
    return parse_number(raw)


def calculate(
    *,
    rule: str,
    compute_budget: float | None,
    ratio: float | None,
    n_active_input: float | None,
    d_tokens_input: float | None,
    seq_len: int,
) -> dict[str, Any]:
    if not isinstance(seq_len, int) or isinstance(seq_len, bool) or seq_len <= 0:
        raise ValueError("--seq-len must be a positive integer")
    for name, value in {
        "compute-budget": compute_budget, "ratio": ratio,
        "n-active": n_active_input, "d-tokens": d_tokens_input,
    }.items():
        if value is not None and (not math.isfinite(value) or value <= 0):
            raise ValueError(f"--{name} must be finite and positive")
    if compute_budget is not None and (n_active_input is not None or d_tokens_input is not None):
        raise ValueError("choose compute-budget/ratio OR n-active/d-tokens, not both")
    if compute_budget is None and ratio is not None:
        raise ValueError("--ratio requires --compute-budget")
    if compute_budget is not None:
        if ratio is None:
            raise ValueError("--ratio is required when --compute-budget is provided")
        target_n = math.sqrt(compute_budget / (6 * ratio))
        target_d = ratio * target_n
    else:
        target_n = n_active_input
        target_d = d_tokens_input

    if target_n is None and target_d is None:
        raise ValueError("provide either --compute-budget/--ratio or --n-active/--d-tokens")

    if target_n is None:
        if compute_budget is None:
            raise ValueError("--n-active is required when only --d-tokens is provided")
        target_n = compute_budget / (6 * target_d)

    requested_d = target_d
    min_n = n_active_nonembed(arch_from_hidden(HIDDEN_MIN))
    max_n = n_active_nonembed(arch_from_hidden(HIDDEN_MAX))
    if not min_n <= target_n <= max_n:
        raise ValueError(f"active parameter target must lie in architecture range [{min_n}, {max_n}]")
    hidden_size = find_hidden_for_target_n(target_n)
    arch = arch_from_hidden(hidden_size)
    n_actual = n_active_nonembed(arch)

    if target_d is None:
        if compute_budget is None:
            raise ValueError("--d-tokens is required when no --compute-budget is provided")
        target_d = compute_budget / (6 * n_actual)
    elif compute_budget is not None:
        # Match design_matrix.py behavior: after architecture rounding, hold C close.
        target_d = compute_budget / (6 * n_actual)

    if rule == "v1":
        global_bs = global_bs_v1(n_actual)
        peak_lr = peak_lr_v1(n_actual)
    elif rule == "v2":
        global_bs = global_bs_v2(target_d)
        peak_lr = peak_lr_v2(n_actual, global_bs)
    else:
        raise ValueError(f"unknown rule: {rule}")

    tokens_per_step = global_bs * seq_len
    train_iters = max(500, round(target_d / tokens_per_step / 100) * 100)
    d_final = train_iters * tokens_per_step
    c_final = 6 * n_actual * d_final
    r_final = d_final / n_actual
    min_lr = peak_lr * 0.1
    warmup_iters = warmup_iters_for(d_final, global_bs, seq_len, train_iters)
    nodes = nodes_for(n_actual)
    warnings = []
    if train_iters == 500 and target_d < 500 * tokens_per_step:
        warnings.append("The legacy 500-step minimum exceeds the requested token budget.")
    if abs(d_final / target_d - 1) > 0.05:
        warnings.append("Rounded training tokens differ by more than 5% from the planned token budget.")
    if global_bs % 2:
        warnings.append("Global batch is not divisible by the suggested micro batch of 2.")
    warnings.append("Choose TP/PP/CP/EP and data parallel size before validating batch divisibility and memory fit.")

    return {
        "rule": rule,
        "targets": {
            "C_target": compute_budget,
            "r_target_D_over_N": ratio,
            "N_active_target": target_n,
            "D_tokens_target": requested_d,
            "D_tokens_after_arch_rounding": target_d,
        },
        "realized": {
            "C_final": c_final,
            "r_final_D_over_N": r_final,
            "N_active": n_actual,
            "D_tokens": d_final,
        },
        "architecture": asdict(arch),
        "warnings": warnings,
        "training": {
            "seq_len": seq_len,
            "micro_batch_size": 2,
            "global_batch_size": global_bs,
            "tokens_per_step": tokens_per_step,
            "train_iters": train_iters,
            "lr": peak_lr,
            "min_lr": min_lr,
            "lr_warmup_iters": warmup_iters,
            "lr_decay_style": "cosine",
            "lr_decay_iters": train_iters,
            "weight_decay": 0.01,
        },
        "resources": {
            "nodes": nodes,
            "gpus": nodes * 8,
        },
    }


def fmt_sci(x: Any) -> str:
    if x is None:
        return "None"
    if isinstance(x, float):
        return f"{x:.6e}"
    return str(x)


def print_text(result: dict[str, Any]) -> None:
    print("\n=== Apodex Hyperparameter Calculator ===")
    print(f"rule: {result['rule']}")
    for warning in result["warnings"]:
        print(f"warning: {warning}")

    print("\n[targets]")
    for key, value in result["targets"].items():
        print(f"{key:24s}: {fmt_sci(value)}")

    print("\n[realized]")
    for key, value in result["realized"].items():
        print(f"{key:24s}: {fmt_sci(value)}")

    print("\n[architecture]")
    for key, value in result["architecture"].items():
        print(f"{key:24s}: {value}")

    print("\n[training]")
    for key, value in result["training"].items():
        print(f"{key:24s}: {fmt_sci(value)}")

    print("\n[resources]")
    for key, value in result["resources"].items():
        print(f"{key:24s}: {value}")

    arch = result["architecture"]
    train = result["training"]
    print("\n[megatron args snippet]")
    print(f"--num-layers {arch['num_layers']}")
    print(f"--hidden-size {arch['hidden_size']}")
    print(f"--ffn-hidden-size {arch['dense_ffn']}")
    print(f"--num-attention-heads {arch['num_heads']}")
    print(f"--q-lora-rank {arch['q_lora_rank']}")
    print(f"--kv-lora-rank {arch['kv_lora_rank']}")
    print(f"--num-experts {arch['num_experts']}")
    print(f"--moe-router-topk {arch['top_k']}")
    print(f"--moe-ffn-hidden-size {arch['moe_ffn']}")
    print(f"--moe-shared-expert-intermediate-size {arch['moe_ffn']}")
    print(f"--moe-layer-freq ([0]+[1]*{arch['num_layers'] - 1})")
    print(f"--seq-length {train['seq_len']}")
    print(f"--max-position-embeddings {train['seq_len']}")
    print(f"--micro-batch-size {train['micro_batch_size']}")
    print(f"--global-batch-size {train['global_batch_size']}")
    print(f"--train-iters {train['train_iters']}")
    print(f"--lr {train['lr']:.3e}")
    print(f"--min-lr {train['min_lr']:.3e}")
    print(f"--lr-warmup-iters {train['lr_warmup_iters']}")
    print(f"--lr-decay-style {train['lr_decay_style']}")
    print(f"--lr-decay-iters {train['lr_decay_iters']}")
    print(f"--weight-decay {train['weight_decay']}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Calculate Apodex model and training hyperparameters.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--rule", choices=["v1", "v2"], default="v2")
    parser.add_argument("--compute-budget", type=parse_number, help="FLOPs, e.g. 3e19")
    parser.add_argument("--ratio", type=parse_number, help="r = D_tokens / N_active")
    parser.add_argument("--n-active", type=parse_number, help="target active params, e.g. 300M")
    parser.add_argument("--d-tokens", type=parse_number, help="target train tokens, e.g. 16B")
    parser.add_argument("--seq-len", type=int, default=SEQ_LEN)
    parser.add_argument("--json", action="store_true", help="print JSON instead of text")
    parser.add_argument(
        "--interactive",
        action="store_true",
        help="prompt for missing inputs interactively",
    )
    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    if args.interactive or (
        args.compute_budget is None and args.n_active is None and args.d_tokens is None
    ):
        print("Choose input mode:")
        print("  1) compute budget C + ratio r = D / N_active")
        print("  2) N_active + D_tokens")
        mode = input("mode [1]: ").strip() or "1"
        if mode == "1":
            args.compute_budget = prompt_number("compute budget C, e.g. 3e19")
            args.ratio = prompt_number("ratio r = D / N_active, e.g. 50")
            args.n_active = None
            args.d_tokens = None
        elif mode == "2":
            args.n_active = prompt_number("N_active, e.g. 300M")
            args.d_tokens = prompt_number("D_tokens, e.g. 16B")
            args.compute_budget = None
            args.ratio = None
        else:
            raise SystemExit(f"unknown mode: {mode}")

        rule = input(f"rule v1/v2 [{args.rule}]: ").strip()
        if rule:
            args.rule = rule

    try:
        result = calculate(
            rule=args.rule,
            compute_budget=args.compute_budget,
            ratio=args.ratio,
            n_active_input=args.n_active,
            d_tokens_input=args.d_tokens,
            seq_len=args.seq_len,
        )
    except ValueError as error:
        parser.error(str(error))

    if args.json:
        print(json.dumps(result, indent=2, sort_keys=False))
    else:
        print_text(result)


if __name__ == "__main__":
    main()