"""Offline packed-weight arithmetic, not a device-fit or runtime-memory estimator.

API inputs are int, float or Decimal. Parameter count must be a positive whole
number; bits per parameter may be fractional to model an ideal average precision.
No weights are downloaded, no devices are inspected and no inference is run.
"""

import argparse
from decimal import Decimal, InvalidOperation
from fractions import Fraction
import json
import math


DEFAULT_PARAMETERS = 24_000_000_000
GB_BYTES = 10**9
GIB_BYTES = 2**30
ESTIMATE_TYPE = "IDEAL_PACKED_WEIGHTS_ONLY_LOWER_BOUND"
EXCLUSIONS = (
    "quantization scales and other metadata",
    "higher-precision or mixed tensors",
    "checkpoint headers, alignment and auxiliary files",
    "KV cache and other runtime state",
    "activation buffers, workspaces and load-time copies",
    "batching, concurrent requests and non-model processes",
)


def _positive_number(value, name):
    if isinstance(value, bool) or not isinstance(value, (int, float, Decimal)):
        raise ValueError(f"{name} must be a positive finite number, not bool or text")
    number = Decimal(str(value))
    if not number.is_finite() or number <= 0:
        raise ValueError(f"{name} must be positive and finite")
    return number


def weight_floor(parameter_count, bits_per_parameter):
    """Return ideal P * bits / 8 in bytes, decimal GB and binary GiB.

    Sub-byte results are not rounded to allocation units. The caller must use
    the total parameter count, not the per-token active parameter count.
    """
    parameters = _positive_number(parameter_count, "parameter_count")
    bits = _positive_number(bits_per_parameter, "bits_per_parameter")
    if parameters != parameters.to_integral_value():
        raise ValueError("parameter_count must be a whole number")
    parameters = int(parameters)
    ideal_bytes = Fraction(parameters) * Fraction(bits) / 8
    try:
        numeric_bits = float(bits)
        numeric_bytes = float(ideal_bytes)
        gb = float(ideal_bytes / GB_BYTES)
        gib = float(ideal_bytes / GIB_BYTES)
    except OverflowError as error:
        raise ValueError("inputs exceed the supported numeric output range") from error
    if not all(math.isfinite(value) and value > 0
               for value in (numeric_bits, numeric_bytes, gb, gib)):
        raise ValueError("inputs exceed the supported numeric output range")
    return {
        "total_parameters": parameters,
        "bits_per_parameter": numeric_bits,
        "ideal_weight_bytes": (
            ideal_bytes.numerator if ideal_bytes.denominator == 1 else numeric_bytes
        ),
        "ideal_weight_gb": gb,
        "ideal_weight_gib": gib,
    }


def _cli_number(text):
    try:
        number = Decimal(text)
    except InvalidOperation as error:
        raise argparse.ArgumentTypeError("expected a numeric value") from error
    if not number.is_finite() or number <= 0:
        raise argparse.ArgumentTypeError("expected a positive finite numeric value")
    return number


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--parameters", type=_cli_number, default=DEFAULT_PARAMETERS,
        help="total parameters, not active parameters (default: 24000000000)",
    )
    parser.add_argument(
        "--bits", type=_cli_number, nargs="+", default=(16, 8, 4),
        help="ideal uniform bits per parameter (default: 16 8 4)",
    )
    args = parser.parse_args(argv)
    try:
        results = [weight_floor(args.parameters, bits) for bits in args.bits]
    except ValueError as error:
        parser.error(str(error))
    print(json.dumps({
        "estimate_type": ESTIMATE_TYPE,
        "formula": "total_parameters * bits_per_parameter / 8",
        "units": {"GB_bytes": GB_BYTES, "GiB_bytes": GIB_BYTES},
        "assumption": "all parameters packed at one ideal uniform bit width",
        "exclusions": list(EXCLUSIONS),
        "inference_status": "NOT_RUN",
        "device_fit": None,
        "measured_memory_bytes": None,
        "results": results,
    }, indent=2, allow_nan=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())