"""Calculate an average and record the inputs, result, and report in MLflow."""

import argparse
import json
import os

import mlflow

from get_token import get_token


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--numbers", nargs="+", type=int, default=[10, 20, 30])
    parser.add_argument("--experiment", default="my-experiments")
    parser.add_argument("--run-name", default="calculate-average")
    args = parser.parse_args()

    tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "").strip()
    if not tracking_uri:
        raise SystemExit("Set MLFLOW_TRACKING_URI to your instance URL first.")

    # Select the instance and authenticate before making tracking requests.
    mlflow.set_tracking_uri(tracking_uri)
    os.environ["MLFLOW_TRACKING_TOKEN"] = get_token()
    mlflow.set_experiment(args.experiment)

    with mlflow.start_run(run_name=args.run_name) as run:
        # Replace this calculation with your own processing code.
        average = sum(args.numbers) / len(args.numbers)

        # Choose the inputs, numerical results, and files to record.
        mlflow.log_param("numbers", json.dumps(args.numbers))
        mlflow.log_metric("average", average)
        mlflow.log_text(
            json.dumps({"numbers": args.numbers, "average": average}, indent=2),
            "result.json",
        )

    print(f"Average: {average}")
    print(f"Run ID: {run.info.run_id}")


if __name__ == "__main__":
    main()