#!/usr/bin/env python3
"""Original LockedIn FDE teaching lab. No GPU, packages, network or credentials.

Inputs are hypothetical unless the learner replaces them with measured values.
This is budgeting arithmetic, not a Nemotron benchmark or a supplier quotation.
Memory is for one device; aggregate multi-GPU placement is deliberately not modeled.
"""
import argparse
import json
import math
import unittest


def positive(value):
    number = float(value)
    if not math.isfinite(number) or number <= 0:
        raise argparse.ArgumentTypeError("Use a finite number greater than zero.")
    return number


def count(value):
    number = int(value)
    if number < 1:
        raise argparse.ArgumentTypeError("Use a positive whole number.")
    return number


def memory(parameters_b=8, bits=16, capacity_gb=48, overhead_gb=4,
           request_gb=2, concurrency=4, reserve_pct=20):
    values = [parameters_b, bits, capacity_gb, overhead_gb, request_gb, concurrency]
    if any(not math.isfinite(x) or x <= 0 for x in values):
        raise ValueError("Memory inputs must be finite and positive.")
    if concurrency != int(concurrency) or bits not in (4, 8, 16, 32):
        raise ValueError("Use whole concurrency and a teaching precision of 4, 8, 16 or 32 bits.")
    if not math.isfinite(reserve_pct) or not 0 <= reserve_pct < 100:
        raise ValueError("Reserve percentage must be at least zero and below 100.")
    weights = parameters_b * bits / 8
    reserve = capacity_gb * reserve_pct / 100
    total = weights + overhead_gb + request_gb * concurrency + reserve
    return {
        "scenario": "synthetic single-device budget; not measured GPU performance",
        "units": "decimal GB (1 GB = 1,000,000,000 bytes)",
        "weights_gb": round(weights, 3),
        "weights_gib": round(weights * 1e9 / (1024 ** 3), 3),
        "request_state_gb": round(request_gb * concurrency, 3),
        "runtime_overhead_gb": overhead_gb,
        "safety_reserve_gb": round(reserve, 3),
        "total_budget_gb": round(total, 3),
        "device_capacity_gb": capacity_gb,
        "fits_estimate": total <= capacity_gb,
        "next": "Measure the exact model/runtime before setting a real admission limit.",
    }


def compare(hours=160, hourly=8, purchase=80000, months=36, operating=800):
    if any(not math.isfinite(x) or x <= 0 for x in [hours, hourly, purchase, months, operating]):
        raise ValueError("Cost inputs must be finite and positive.")
    owned = purchase / months + operating
    rented = hours * hourly
    return {
        "scenario": "hypothetical prices for equivalent capacity; not supplier quotes",
        "scope": "same capacity bundle; hours per month; all amounts in one chosen currency",
        "rental_per_month": round(rented, 2),
        "ownership_per_month": round(owned, 2),
        "arithmetic_break_even_hours_per_month": round(owned / hourly, 2),
        "lower_estimated_cost": "rent" if rented < owned else "own" if owned < rented else "equal",
        "excluded": ["tax", "financing", "residual value", "migration", "unmodeled downtime"],
        "next": "Obtain dated comparable quotes and validate workload fit, availability and operating readiness.",
    }


class BudgetTests(unittest.TestCase):
    def test_weights_and_units(self):
        self.assertEqual(memory()["weights_gb"], 16)
        self.assertAlmostEqual(memory()["weights_gib"], 14.901, places=3)

    def test_admission_changes_with_concurrency(self):
        self.assertTrue(memory(concurrency=4)["fits_estimate"])
        self.assertFalse(memory(concurrency=10)["fits_estimate"])

    def test_exact_capacity_and_overhead(self):
        self.assertTrue(memory(capacity_gb=35, reserve_pct=20, request_gb=2, concurrency=4)["fits_estimate"])
        self.assertFalse(memory(capacity_gb=34, reserve_pct=20, request_gb=2, concurrency=4)["fits_estimate"])

    def test_precision_is_not_a_quality_claim(self):
        self.assertEqual(memory(bits=8)["weights_gb"], 8)
        self.assertIn("not measured", memory(bits=8)["scenario"])

    def test_invalid_values_rejected(self):
        for arguments in [{"capacity_gb": 0}, {"request_gb": float("nan")}, {"reserve_pct": 100}, {"concurrency": 1.2}]:
            with self.assertRaises(ValueError):
                memory(**arguments)

    def test_usage_can_reverse_cost_comparison(self):
        self.assertEqual(compare(hours=160)["lower_estimated_cost"], "rent")
        self.assertEqual(compare(hours=600)["lower_estimated_cost"], "own")
        self.assertEqual(compare()["arithmetic_break_even_hours_per_month"], 377.78)

    def test_missing_cost_is_not_zero(self):
        with self.assertRaises(ValueError):
            compare(months=0)
        with self.assertRaises(ValueError):
            compare(operating=float("inf"))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    actions = parser.add_subparsers(dest="action", required=True)
    budget = actions.add_parser("memory", help="Synthetic single-device memory budget")
    for name, default in [("parameters-b", 8), ("capacity-gb", 48), ("overhead-gb", 4), ("request-gb", 2)]:
        budget.add_argument("--" + name, type=positive, default=default)
    budget.add_argument("--bits", type=int, choices=[4, 8, 16, 32], default=16)
    budget.add_argument("--concurrency", type=count, default=4)
    budget.add_argument("--reserve-pct", type=float, default=20)
    cost = actions.add_parser("compare", help="Hypothetical equivalent-capacity rent/own comparison")
    for name, default in [("hours", 160), ("hourly", 8), ("purchase", 80000), ("months", 36), ("operating", 800)]:
        cost.add_argument("--" + name, type=positive, default=default)
    actions.add_parser("test", help="Check the arithmetic and boundary cases")
    args = vars(parser.parse_args())
    action = args.pop("action")
    if action == "test":
        result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(BudgetTests))
        raise SystemExit(0 if result.wasSuccessful() else 1)
    try:
        print(json.dumps(memory(**args) if action == "memory" else compare(**args), indent=2))
    except ValueError as error:
        parser.error(str(error))


if __name__ == "__main__":
    main()