Calculate and show caffeine level in body with time.

caffeines_level_calculator.py

"""Calculate and show caffeine level in body with time."""

from math import exp
from typing import Iterable
import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = [24, 12]

# caffeine half life in hours
HALF_LIFE_HOURS = 5


def calculate(caffeines: int, time: int) -> Iterable[float]:
    return [
        round(caffeines, 2) * exp(-0.693 * hour / HALF_LIFE_HOURS)
        for hour
        in range(time + 1)
    ]


def plot(data: Iterable[float]) -> None:
    plt.title("Caffeine level in body with time")
    plt.xlabel("Hours)")
    plt.ylabel("mg")
    plt.plot(data, color="red")
    plt.grid()
    plt.show()


def main() -> None:
    while True:
        try:
            level = int(input("Enter initial caffeine level in mg: "))
            time = int(input("Enter time to calculate caffeine in body: "))
            break
        except ValueError:
            print("Invalid value. Try again")

    plot(calculate(level, time))


if __name__ == "__main__":
    main()