#!/usr/bin/env python3
"""Original, offline teaching calculator. Every workload and rate is fictional."""
import argparse
import json
from decimal import Decimal, ROUND_HALF_UP

D = Decimal
REQUESTS = 100
INPUT_PER_ATTEMPT = 1_000
OUTPUT_PER_ATTEMPT = 400
# Illustrative currency units per million tokens, NOT vendor prices.
INPUT_RATE = D("40")
OUTPUT_RATE = D("200")
REVIEW_RATE_PER_MINUTE = D("0.25")
INFRASTRUCTURE = D("2")
SUPPORT = D("6")
MIN_ACCEPTED = 80  # A supplied comparison rule, not a production release gate.

# Counts describe final, distinct tasks accepted by a fictional reviewer.
ACCEPTED = {
    "manual": 90, "baseline": 80, "retries": 80,
    "cheaper": 60, "review-heavy": 80, "none-accepted": 0,
}
RETRY_TASKS = {
    "manual": 0, "baseline": 0, "retries": 20,
    "cheaper": 40, "review-heavy": 0, "none-accepted": 0,
}
REVIEW_MINUTES = {
    "manual": D("2.56"), "baseline": D("0.8"), "retries": D("0.8"),
    "cheaper": D("1.6"), "review-heavy": D("1.6"), "none-accepted": D("0.8"),
}
RATE_MULTIPLIER = {
    "manual": D("0"), "baseline": D("1"), "retries": D("1"),
    "cheaper": D("0.5"), "review-heavy": D("1"), "none-accepted": D("1"),
}


def task_ledger(case):
    """One row per requested task; retries add attempts, never accepted tasks."""
    if case not in ACCEPTED:
        raise ValueError("Unknown case")
    if not 0 <= ACCEPTED[case] <= REQUESTS or not 0 <= RETRY_TASKS[case] <= REQUESTS:
        raise ValueError("Task counts must fit the workload")
    if REVIEW_MINUTES[case] < 0 or RATE_MULTIPLIER[case] < 0:
        raise ValueError("Rates and effort cannot be negative")
    rows = []
    for index in range(REQUESTS):
        retried = index < RETRY_TASKS[case]
        attempts = 0 if case == "manual" else 1 + int(retried)
        # The retries case adds a minute of review to each of its 20 retried tasks.
        extra_review = D("1") if case == "retries" and retried else D("0")
        rows.append({
            "task_id": f"T-{index + 1:03}",
            "attempts": attempts,
            "input_tokens": attempts * INPUT_PER_ATTEMPT,
            "output_tokens": attempts * OUTPUT_PER_ATTEMPT,
            "review_minutes": REVIEW_MINUTES[case] + extra_review,
            "accepted": index < ACCEPTED[case],
        })
    return rows


def number(value, places="0.0001"):
    """Round only for display; Decimal arithmetic retains unrounded totals."""
    return str(value.quantize(D(places), rounding=ROUND_HALF_UP))


def summarize(case):
    rows = task_ledger(case)
    attempts = sum(row["attempts"] for row in rows)
    accepted = sum(row["accepted"] for row in rows)
    input_tokens = sum(row["input_tokens"] for row in rows)
    output_tokens = sum(row["output_tokens"] for row in rows)
    review_minutes = sum((row["review_minutes"] for row in rows), D("0"))
    input_cost = D(input_tokens) / 1_000_000 * INPUT_RATE * RATE_MULTIPLIER[case]
    output_cost = D(output_tokens) / 1_000_000 * OUTPUT_RATE * RATE_MULTIPLIER[case]
    model_cost = input_cost + output_cost
    review_cost = review_minutes * REVIEW_RATE_PER_MINUTE
    total = model_cost + review_cost + INFRASTRUCTURE + SUPPORT
    return {
        "case": case,
        "requests": len(rows),
        "model_attempts": attempts,
        "accepted_distinct_tasks": accepted,
        "not_accepted_tasks": len(rows) - accepted,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "review_minutes": number(review_minutes, "0.01"),
        "cost": {
            "model": number(model_cost, "0.01"),
            "review_or_manual_handling": number(review_cost, "0.01"),
            "infrastructure": number(INFRASTRUCTURE, "0.01"),
            "support": number(SUPPORT, "0.01"),
            "total": number(total, "0.01"),
        },
        "model_cost_per_attempt": number(model_cost / attempts) if attempts else None,
        "cost_per_requested_task": number(total / len(rows)),
        "cost_per_accepted_task": number(total / accepted) if accepted else None,
        "meets_supplied_acceptance_floor": accepted >= MIN_ACCEPTED,
        "decision": "No accepted denominator" if not accepted else
                    "Below the supplied acceptance floor" if accepted < MIN_ACCEPTED else
                    "Eligible for comparison; not a release approval",
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("case", choices=["compare", *ACCEPTED], nargs="?", default="compare")
    args = parser.parse_args()
    cases = list(ACCEPTED) if args.case == "compare" else [args.case]
    print(json.dumps({
        "synthetic": True,
        "currency": "illustrative units, not a vendor quote",
        "scope": "Same 100 fictional requests; supplied final acceptance labels",
        "limits": "No model call, quality assessment, billing integration or revenue estimate",
        "results": [summarize(case) for case in cases],
    }, indent=2))


if __name__ == "__main__":
    main()
